Compare commits

..
Author SHA1 Message Date
Nicolò Boschi 0022d427d3 feat(integrations): add hindsight-opencode-coding plugin
Reflect-only long-term memory for coding agents in OpenCode, with a git+chat
backfill and (opt-in) live session write-back.

- reflect + INJECT: on a task, reflect() the symptom and push the root-cause
  answer into the system prompt (no tools/recall).
- backfill: every commit (full message + full diff, commit timestamp + git
  metadata) under a 'git' retain strategy; each chat as a JSON user/assistant
  transcript with custom extraction (<=2 coherent facts) under a 'chat' strategy;
  observations on; optional codebase knowledge pages.
- live write-back (opt-in HINDSIGHT_RETAIN_SESSIONS): every N turns upsert the
  tool-filtered transcript under a stable conversation:<sessionID> document_id.
2026-07-02 13:55:40 +02:00
Nicolò Boschi 1f9bad0858 feat(knowledge-base): default pages to living-document trigger + 4096 tokens
Client-created pages had no server curation applying a trigger, so they fell back
to the plain mental-model default (no refresh, full mode, all fact types). Make a
knowledge page a living document by default: when the client omits `trigger`, use
observation-only + delta + exclude_mental_models + refresh_after_consolidation;
when it omits `max_tokens`, default to 4096 (vs the mental-model 2048). Clients
can still override either.
2026-07-02 10:30:40 +02:00
Nicolò Boschi 138bf02f29 refactor(knowledge-base): drop server-side curation + folder missions
The knowledge base is now purely client-managed (CRUD over folders/pages); the
server no longer auto-curates. Removes the folder curator entirely and the
folder `mission` concept, and leads the sidebar with Knowledge Base.

- Remove engine/knowledge_curator.py, the curate_folder task (handler + dispatch
  + submit_async_curate_folder / _bank_folders), the post-consolidation curation
  hook, and the folder-create / mission-update curation triggers.
- Remove folder `mission` and `last_curated_at` (columns + engine + API + UI);
  keep `managed` as a client-set flag. Migration a5b6 now adds `managed` only;
  the last_curated_at migration is dropped and the unique-index migration
  repointed. Single alembic head preserved.
- API: KnowledgeNode/CreateFolderRequest/UpdateNodeRequest lose `mission`;
  PATCH node handles name/parent_id only.
- Control plane: sidebar leads with Knowledge Base (before Memories); remove the
  mission field, edit-mission dialog, and mission display from the KB view.
- Delete the curator tests; regenerate OpenAPI + SDK clients.
2026-07-02 10:30:39 +02:00
Nicolò Boschi df178aae8a refactor(hindsight-fs): mirror the knowledge-base tree, not mental models
Re-point hindsight-fs at the knowledge base so it projects a bank's folder/page
hierarchy as nested directories + .md files, instead of a flat list of mental
models.

- client: fetch GET /knowledge-base/tree + /export (two calls, any bank size)
  and join by page id; replaces the paginated mental-models list.
- format: planMirror() walks the tree into folder dirs + page files at nested
  paths (slug per segment, collision-safe); pages render the page's OKF doc.
- sync: create folder dirs, write pages at nested paths, prune removed pages and
  emptied folders; state keyed by relative path + tracked dirs.
- config/cli: drop the mental-model `detail` flag; `list` prints folders+pages;
  help/README updated. Tests rewritten for the tree/export model.

Verified live against a bank's knowledge base: the `people` folder mirrors to
people/anna.md + people/marco.md with OKF frontmatter.
2026-07-02 10:30:39 +02:00
Nicolò Boschi a43026b8f4 feat(hindsight-fs): mirror a bank's mental models as a live local folder
Add @vectorize-io/hindsight-fs, a CLI under hindsight-tools/ that mirrors a
Hindsight bank's mental models as real markdown files (YAML frontmatter + body)
in a local directory, refreshed from the API on an interval. Once mounted,
ordinary shell tools (ls, cat, grep, find, ...) work against current memory.

- Pull-based sync engine: full list each tick, write changed/new/tampered
  files, skip unchanged (content-hashed), prune deleted models. Atomic writes;
  a transient API error never wipes the mirror.
- One-way mirror enforced two ways: files are read-only (0444) so agent edits
  fail with EACCES, plus a tamper-revert backstop that compares on-disk bytes
  and overwrites drift on the next pass. --writable opts out.
- Commands: mount/start/stop/restart/sync/status/list/logs/unmount. Background
  daemon via detached process + pidfile; per-mount config is remembered.
- status doubles as a healthcheck: --json report and a non-zero exit when the
  mount is dead/failed/stale (--stale-after overrides the threshold).
- Tests: unit (sync engine, frontmatter, health) + e2e that spawns the real
  CLI against a mock API and exercises real bash commands. 26 tests.
2026-07-02 10:30:39 +02:00
Nicolò Boschi 5c425e276e feat(knowledge-base): self-curating knowledge base (OKF pages + folder missions)
Server-side knowledge base: a hierarchy of folders and pages over mental
models, projected to the Open Knowledge Format, with a mission-driven curator
that maintains pages automatically after each consolidation.

- knowledge_pages table (PG + Oracle): parent_id tree, kind folder/page,
  mission, managed, last_curated_at; partial unique index on (folder, name)
  for concurrency-safe dedup; added to BACKUP_TABLES.
- api/okf.py: OKF serializer (frontmatter + body, index/log, constellation graph).
- engine/knowledge_curator.py: folder curator (LLM op plan + safe apply); reads
  new memories since last curation (delta, not recall); ops create/merge/delete
  page + spawn sub-folder (bounded depth<=3, <=8). Runs as an async curate_folder
  task on folder/mission create and after consolidation. Curator pages use an
  observation-only delta trigger with exclude_mental_models.
- MemoryEngine: folder/page CRUD, tree, curate, async submit + worker handler.
- /v1/default/banks/{bank}/knowledge-base/* endpoints.
- Control plane: knowledge-base tree view + constellation toggle, missions,
  OKF page panel + bundle export; proxies, client, sidebar, i18n.
- Tests: okf unit, knowledge-base HTTP, curator apply + dedup guard, hs_llm_core e2e.
- Regenerated OpenAPI + SDK clients + docs-skill.
2026-07-02 10:30:39 +02:00
1936 changed files with 28841 additions and 225289 deletions
+1 -6
View File
@@ -1,7 +1,7 @@
{
"$schema": "https://anthropic.com/claude-code/marketplace.schema.json",
"name": "hindsight",
"version": "0.7.5",
"version": "0.7.2",
"description": "Official Hindsight integrations for Claude Code",
"owner": {
"name": "vectorize-io"
@@ -11,11 +11,6 @@
"name": "hindsight-memory",
"description": "Automatic long-term memory for Claude Code via Hindsight",
"source": "./hindsight-integrations/claude-code"
},
{
"name": "hindsight-zcode",
"description": "No-MCP long-term memory for ZCode via Hindsight hooks",
"source": "./hindsight-integrations/zcode"
}
]
}
-92
View File
@@ -78,22 +78,6 @@ results = await asyncio.gather(*tasks, return_exceptions=True)
- **Authentication/tenancy is enforced inside each engine method, not assumed by the handler.** Every engine method that touches bank-scoped data must authenticate via `request_context` — typically `await self._authenticate_tenant(request_context)` (often indirectly through `get_bank_profile(...)`) — so the correct tenant schema is resolved before any query runs. Handlers must thread `request_context` through to the engine method; never query a tenant-scoped table assuming the schema is already set.
- Engine methods return typed models (Pydantic/dataclass), not raw dicts (see Type Safety).
### 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.
- The pre-existing usage in `hindsight_api/migrations.py` is grandfathered, not a precedent — it is tracked for removal. Don't copy it.
- Design the concurrency out instead of locking around it: give each process its own object to write (e.g. per-schema DDL rather than a shared `public.` object), make the operation idempotent, or use a real row/table constraint (`INSERT ... ON CONFLICT`, `SELECT ... FOR UPDATE` in a fixed order). See #2690 for a migration that reached for `pg_advisory_xact_lock` and had to be reverted.
### Branch Hygiene
- **Always start new feature branches from `origin/main`** — rebase to ensure a clean base.
- **Only include commits relevant to the PR/branch/feature** — no unrelated changes. If the branch contains commits that don't belong, they must be removed before merging.
@@ -170,35 +154,12 @@ If any files in `hindsight-api-slim/hindsight_api/api/` were changed:
- Were the client SDKs regenerated? (`./scripts/generate-clients.sh`)
- Were the control plane proxy routes updated? (`hindsight-control-plane/src/app/api/`)
### 7a. Check TS/Python wrapper-client parity
Two of the generated SDKs ship a **hand-written, maintained convenience wrapper** on top of the auto-generated low-level client — and *only* these two:
- **TypeScript**: `hindsight-clients/typescript/src/index.ts` (`HindsightClient`)
- **Python**: `hindsight-clients/python/hindsight_client/hindsight_client.py` (`Hindsight`)
(The Rust/Go/etc. clients are generated-only — no wrapper to keep in sync.)
These wrappers are what most third-party consumers actually call, and they must expose the same surface. **If a change touches one wrapper's method — adds/removes a parameter, changes a default, forwards a new query/body field — the equivalent method in the *other* wrapper must get the same change in the same (or an immediately-following) PR.** A parameter that exists in the generated SDK but is dropped by one wrapper silently strips it for every consumer of that language (this is exactly what #2975 / #3042 fixed for `detail`/`tags_match`/`limit`/`offset` on `listMentalModels`/`getMentalModel`). **Should fix** — flag any wrapper method that gains capabilities in one language but not the other, and add a matching mapping regression test on both sides.
Note: the `client-coverage-check` CI tool only validates **request-body** fields, not GET **query** parameters — so query-param parity gaps are *not* caught automatically and must be checked by hand here.
### 7b. Check API-layer data-access boundary
For each changed handler in `hindsight-api-slim/hindsight_api/api/` (e.g. `http.py`, `mcp.py`):
- **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.
### 8. Check code comments
For each non-trivial change:
@@ -215,48 +176,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`:
@@ -285,14 +204,6 @@ in `hindsight-api-slim/hindsight_api/config.py`):
The `test_bundled_template_matches_repo_root` sync test fails on drift; if the
root file changed without re-copying, flag it as a **must fix**.
### 11c. Check for advisory locks
Grep the diff for `advisory` (`git diff main...HEAD | grep -in advisory`). Any new
`pg_advisory_lock` / `pg_try_advisory_lock` / `pg_advisory_xact_lock` /
`pg_advisory_unlock` call is a **must fix** — see Database Locking above. Point the
author at the alternatives (per-process objects, idempotent DDL, row-level
constraints) rather than just asking them to drop the lock.
### 12. Review against other coding standards
Check the diff for violations of the standards listed above:
@@ -318,12 +229,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
+2 -209
View File
@@ -2,15 +2,12 @@
# Copy this file to .env and fill in your values
# LLM Configuration (Required)
# Supported providers: openai, openai-responses, groq, ollama, gemini, anthropic, lmstudio, vertexai, minimax, deepseek, zai, atlas, volcano
# Supported providers: openai, groq, ollama, gemini, anthropic, lmstudio, vertexai, minimax, deepseek, zai, atlas, volcano
HINDSIGHT_API_LLM_PROVIDER=openai
HINDSIGHT_API_LLM_API_KEY=your-api-key-here
HINDSIGHT_API_LLM_MODEL=gpt-4o-mini
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`
@@ -24,51 +21,6 @@ HINDSIGHT_API_LLM_BASE_URL=https://api.openai.com/v1
# HINDSIGHT_API_LLM_TEMPERATURE_REFLECT=0.9
# HINDSIGHT_API_LLM_TEMPERATURE_CONSOLIDATION=0.0
# Grammar-enforce structured output (json_schema strict) instead of the soft
# schema-in-prompt path. Helps weaker self-hosted models that emit prose preambles
# or invalid JSON. The global override below applies to every operation;
# per-operation overrides take precedence, in both directions -- set one to false
# to opt that operation out while the global flag is on.
# HINDSIGHT_API_LLM_STRICT_SCHEMA=false
# HINDSIGHT_API_LLM_STRICT_SCHEMA_RETAIN=true
# HINDSIGHT_API_LLM_STRICT_SCHEMA_REFLECT=true
# HINDSIGHT_API_LLM_STRICT_SCHEMA_CONSOLIDATION=true
# Some backends, including Bedrock Converse, reject JSON Schema maxItems.
# Disable it only for those backends; consolidation still enforces the cap.
# HINDSIGHT_API_LLM_SUPPORTS_MAX_ITEMS=true
# Pin a conversation to one backend prompt cache (OpenAI-compatible providers only).
# Server-side prompt caches are per backend server, so the same conversation has to
# reach the same one to hit. Values: auto (default), xai_conv_id (sends xAI's
# x-grok-conv-id header), openai_prompt_cache_key (sends OpenAI's prompt_cache_key
# field), none (sends nothing). "auto" picks from the base URL host and is an
# allowlist: x.ai / grok.com get the header, native OpenAI / openai.com / Azure
# OpenAI get the field, and every other backend gets nothing. Per-operation
# overrides take precedence. Set to none to disable entirely.
# HINDSIGHT_API_LLM_CACHE_AFFINITY=auto
# HINDSIGHT_API_RETAIN_LLM_CACHE_AFFINITY=none
# HINDSIGHT_API_REFLECT_LLM_CACHE_AFFINITY=xai_conv_id
# HINDSIGHT_API_CONSOLIDATION_LLM_CACHE_AFFINITY=none
# Ask litellm/litellmrouter/bedrock for structured output via a forced tool call
# instead of response_format. Enable it for backends that reject response_format
# outright -- e.g. Bedrock Claude in ap-southeast-2 ("Extra inputs are not permitted");
# 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.
# HINDSIGHT_API_LLM_DEBUG_DUMP_4XX=false
# Example: Anthropic Claude configuration
# HINDSIGHT_API_LLM_PROVIDER=anthropic
# HINDSIGHT_API_LLM_API_KEY=your-anthropic-api-key
@@ -86,12 +38,6 @@ HINDSIGHT_API_LLM_BASE_URL=https://api.openai.com/v1
# HINDSIGHT_API_LLM_API_KEY=your-minimax-api-key
# HINDSIGHT_API_LLM_MODEL=MiniMax-M3 # or MiniMax-M2.7 for the previous generation
# Example: OpenAI Responses API (/v1/responses) — reasoning + function tools together
# HINDSIGHT_API_LLM_PROVIDER=openai-responses
# HINDSIGHT_API_LLM_API_KEY=your-openai-api-key
# HINDSIGHT_API_LLM_MODEL=gpt-5.6 # reasoning model (gpt-5.x / o-series); e.g. gpt-5.6-terra
# HINDSIGHT_API_LLM_REASONING_EFFORT=high # sent alongside tools, unlike chat/completions
# Example: DeepSeek configuration (https://api.deepseek.com)
# HINDSIGHT_API_LLM_PROVIDER=deepseek
# HINDSIGHT_API_LLM_API_KEY=your-deepseek-api-key
@@ -113,15 +59,6 @@ HINDSIGHT_API_LLM_BASE_URL=https://api.openai.com/v1
# HINDSIGHT_API_LLM_BASE_URL=http://localhost:1234/v1
# HINDSIGHT_API_LLM_MODEL=qwen2.5-32b-instruct
# Example: Ollama local configuration (native provider)
# HINDSIGHT_API_LLM_PROVIDER=ollama
# HINDSIGHT_API_LLM_BASE_URL=http://localhost:11434/v1
# HINDSIGHT_API_LLM_MODEL=gemma3:12b
# Native Ollama context-window override (num_ctx). Leave unset to let Ollama use
# the model Modelfile / server default; set a positive integer only to force a
# specific context size (e.g. 16384 to keep the previous request behavior).
# HINDSIGHT_API_LLM_OLLAMA_NUM_CTX=16384
# Multi-LLM strategies: configure extra LLMs by index alongside the primary above,
# then pick a routing strategy. Unset = single primary LLM (default). Members are
# numbered from 1; indices must be contiguous. Each operation can override with a
@@ -143,16 +80,6 @@ HINDSIGHT_API_LOG_LEVEL=info
# Unset uses HINDSIGHT_API_RETAIN_CHUNK_SIZE as the structured-chunk limit.
# HINDSIGHT_API_RETAIN_STRUCTURED_CHUNK_SIZE=
# When true, a retain operation that hit any fact-extraction errors is marked
# 'failed' (not 'completed'), surfacing silently-dropped facts. Default false.
# HINDSIGHT_API_FAIL_ON_EXTRACTION_ERRORS=false
# Wall-clock ceiling (seconds) for one retain task in the worker. A retain that
# blocks indefinitely is cancelled and marked 'failed' — and so becomes
# retryable — instead of holding its worker slot until the process restarts.
# Set well above your slowest healthy retain; 0 disables. Default 3600.
# HINDSIGHT_API_RETAIN_WALL_TIMEOUT=3600
# Dry-run extraction preview endpoint (POST /memories/dry-run-extract). Enabled by default; it makes
# a real LLM call but stores nothing. Set to false to remove the endpoint (returns 404).
# HINDSIGHT_API_ENABLE_DRY_RUN_EXTRACT=true
@@ -168,22 +95,7 @@ HINDSIGHT_API_LOG_LEVEL=info
# HINDSIGHT_API_READ_DATABASE_URL= # Optional read-replica URL. When set, recall queries (semantic, BM25, graph, temporal) flow through a separate pool against this URL, offloading the primary. Typically points to a read-only endpoint (CNPG's <cluster>-ro service or Aurora reader endpoint).
# HINDSIGHT_API_MIGRATION_DATABASE_URL= # Direct PostgreSQL URL for migrations (bypasses PgBouncer). Falls back to DATABASE_URL.
# HINDSIGHT_API_DATABASE_SCHEMA=public # PostgreSQL schema name (default: public)
# HINDSIGHT_API_DB_MAX_PARALLEL_WORKERS_PER_GATHER= # Optional cap on Postgres planner parallelism for this process's pool connections. Unset leaves the server default; 0 makes background/bulk queries run serially (useful on worker processes sharing a primary with latency-sensitive traffic).
# HINDSIGHT_API_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.
# HINDSIGHT_API_MIGRATION_CONCURRENCY=1 # Tenant schemas to migrate concurrently (PG only, each in its own process; per-schema work stays sequential). Each worker has ~1-2s startup cost + uses ~3 DB connections, so it only pays off with many schemas (tens+) or slow migrations; keep concurrency*3 <= spare max_connections. Default: 1 (sequential).
# HINDSIGHT_API_OPERATION_RETENTION_DAYS=30 # Prune terminal operation rows, payloads, and metadata after this many days; 0 (the default) keeps them forever.
# HINDSIGHT_API_OPERATION_CLEANUP_BATCH_SIZE=1000 # Maximum expired terminal rows deleted per tenant schema in each cleanup cycle; must be positive.
# 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)
@@ -203,18 +115,6 @@ 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
# File Parser (Optional - uses markitdown by default)
# HINDSIGHT_API_FILE_PARSER=markitdown
@@ -227,19 +127,12 @@ HINDSIGHT_API_LOG_LEVEL=info
# HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_BASE_URL=
# HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_MODEL=
# HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_PROMPT=
# Optional JSON dict of custom headers for the OCR OpenAI client (e.g. proxies / request tracing).
# HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_DEFAULT_HEADERS=
# Embeddings Configuration (Optional - uses local by default)
# Provider: "local" (default), "onnx", "tei", "openai", "cohere", "google", "openrouter", "zeroentropy", "litellm", or "litellm-sdk"
# HINDSIGHT_API_EMBEDDINGS_PROVIDER=local
# For local provider:
# HINDSIGHT_API_EMBEDDINGS_LOCAL_MODEL=BAAI/bge-small-en-v1.5
# Force CPU if local embeddings hit MPS/XPC instability on macOS:
# HINDSIGHT_API_EMBEDDINGS_LOCAL_FORCE_CPU=false
# Opt in to the Apple Silicon MPS GPU (off by default: MPS leaks memory under
# variable-length workloads). CUDA/XPU still auto-select regardless:
# HINDSIGHT_API_EMBEDDINGS_LOCAL_ALLOW_MPS=false
# For ONNX provider (local CPU embeddings without an Ollama/TEI sidecar):
# HINDSIGHT_API_EMBEDDINGS_PROVIDER=onnx
# HINDSIGHT_API_EMBEDDINGS_ONNX_MODEL_ID=intfloat/multilingual-e5-small
@@ -255,17 +148,6 @@ HINDSIGHT_API_LOG_LEVEL=info
# HINDSIGHT_API_EMBEDDINGS_ONNX_TOKENIZER_NAME_OR_PATH=/models/multilingual-e5-small
# Optional for China network / restricted HF access:
# HF_ENDPOINT=https://hf-mirror.com
# Applies to any provider: cap each input at this many tiktoken tokens before
# embedding, so oversized content is truncated instead of failing the embed call
# permanently (e.g. Bedrock Titan V2's 8192, or a llama.cpp server's context). Off
# by default. (Deprecated alias: HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_MAX_INPUT_TOKENS)
# HINDSIGHT_API_EMBEDDINGS_MAX_INPUT_TOKENS=8192
# 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:
@@ -287,62 +169,13 @@ HINDSIGHT_API_LOG_LEVEL=info
# DeepSeek note: DeepSeek is supported for LLM calls, but not for embeddings.
# If using DeepSeek as LLM provider, keep embeddings on local/openai/cohere/google/etc.
# Embedding similarity thresholds. These defaults preserve the behavior calibrated
# for BAAI/bge-small-en-v1.5. Recalibrate each threshold independently when changing
# embedding models because cosine-similarity distributions are model-dependent.
# HINDSIGHT_API_SEMANTIC_MIN_SIMILARITY=0.3
# HINDSIGHT_API_GRAPH_SEED_MIN_SIMILARITY=0.3
# HINDSIGHT_API_TEMPORAL_SEMANTIC_MIN_SIMILARITY=0.1
# HINDSIGHT_API_SEMANTIC_LINK_MIN_SIMILARITY=0.7
# HINDSIGHT_API_CONSOLIDATION_DEDUP_THRESHOLD=0.97
# Recall pipeline stages (all on by default). Each is hierarchical, so a single
# bank can switch a stage off via the config API without changing the server
# default. Turning all three off leaves semantic + BM25 fused by RRF, the
# lowest-latency recall path.
# Temporal retrieval arm, plus the date-aware query analysis that feeds it:
# HINDSIGHT_API_ENABLE_TEMPORAL_RETRIEVAL=true
# Entity/link graph traversal arm:
# HINDSIGHT_API_ENABLE_GRAPH_RETRIEVAL=true
# Cross-encoder rerank of the fused candidates (false = use the RRF order):
# HINDSIGHT_API_ENABLE_RERANKING=true
# Reranker Configuration (Optional - uses local by default)
# Provider: "local" (default) or "tei" (HuggingFace Text Embeddings Inference)
# HINDSIGHT_API_RERANKER_PROVIDER=local
# Trusted gateway attribution (disabled by default). When enabled, remote
# reranker requests include X-Hindsight-Bank-Id with the current bank ID.
# HINDSIGHT_API_RERANKER_SEND_BANK_AS_HEADER=false
# For local provider:
# HINDSIGHT_API_RERANKER_LOCAL_MODEL=cross-encoder/ms-marco-MiniLM-L-6-v2
# Force CPU if the local reranker hits MPS/XPC instability on macOS:
# HINDSIGHT_API_RERANKER_LOCAL_FORCE_CPU=false
# Opt in to the Apple Silicon MPS GPU (off by default: MPS leaks memory under
# variable-length workloads). CUDA/XPU still auto-select regardless:
# HINDSIGHT_API_RERANKER_LOCAL_ALLOW_MPS=false
# For TEI provider:
# HINDSIGHT_API_RERANKER_TEI_URL=http://localhost:8081
# 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
# a large recall's latency). 0 = fall back to the flat cap above; fully backwards-compatible.
# HINDSIGHT_API_RERANKER_MAX_CANDIDATES_LOW=0
# HINDSIGHT_API_RERANKER_MAX_CANDIDATES_MID=0
# HINDSIGHT_API_RERANKER_MAX_CANDIDATES_HIGH=0
# Reranker failover chain: extra rerankers tried, in order, when the one above
# fails. Members are numbered from 1 (indices must be contiguous) and every
# setting of member n carries the same index. A member inherits nothing from the
# primary, so spell out everything it needs. Unset = no fallback (default): a
# failing reranker fails the recall. End the chain with "rrf" to fail open and
# keep the retrieval order instead.
# HINDSIGHT_API_RERANKER_1_PROVIDER=cohere
# HINDSIGHT_API_RERANKER_1_COHERE_API_KEY=your-cohere-api-key
# HINDSIGHT_API_RERANKER_2_PROVIDER=rrf
# Observability & Tracing (Optional - disabled by default)
# Enable OpenTelemetry tracing for LLM calls (GenAI semantic conventions)
@@ -362,46 +195,6 @@ HINDSIGHT_API_LOG_LEVEL=info
# Expose async-operation queue + consolidation-backlog gauges on /metrics.
# Runs periodic per-schema COUNT queries on a background task (disabled by default).
# HINDSIGHT_API_METRICS_BACKLOG_ENABLED=true
#
# Runtime-stall observability (enabled by default). When a liveness probe fails,
# these tell you WHY: a blocked event loop vs DB connection-pool exhaustion.
# The loop watchdog logs the offending stack when the loop is unresponsive; the
# DB-pool acquire timing logs (and exposes hindsight.db.pool.waiting) when
# callers queue for a connection. Both are cheap; tune or disable if needed.
# HINDSIGHT_API_LOOP_WATCHDOG_ENABLED=false
# HINDSIGHT_API_LOOP_WATCHDOG_STALL_THRESHOLD_MS=1000
# HINDSIGHT_API_LOOP_WATCHDOG_POLL_INTERVAL_MS=250
# HINDSIGHT_API_DB_ACQUIRE_WARN_THRESHOLD_MS=1000
# -----------------------------------------------------------------------------
# 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)
# -----------------------------------------------------------------------------
# Outbound webhook delivery targets caller-supplied URLs. To prevent SSRF, the
# delivery worker blocks private, loopback, and link-local destinations
# (including the cloud metadata address 169.254.169.254) by default. List hosts
# or IP/CIDR ranges here (comma-separated) to re-permit specific internal
# destinations — e.g. 127.0.0.1 for local testing, or an internal receiver.
# HINDSIGHT_API_WEBHOOK_ALLOWED_HOSTS=127.0.0.1,internal-receiver.svc,10.0.0.0/8
# Whether the webhook delivery-history API returns the raw upstream response
# body. Off by default: returning arbitrary response bodies to callers is an
# information-exfiltration primitive. The delivery status code is always
# returned regardless. Enable only if you trust your webhook destinations.
# HINDSIGHT_API_WEBHOOK_EXPOSE_RESPONSE_BODY=false
# -----------------------------------------------------------------------------
# Control Plane (Optional)
File diff suppressed because one or more lines are too long

Before

Width:  |  Height:  |  Size: 41 KiB

File diff suppressed because it is too large Load Diff
+1 -71
View File
@@ -25,20 +25,6 @@ jobs:
with:
python-version-file: ".python-version"
# Each package is built from its own directory, so stage the repository's
# canonical license inside each isolated build context.
- name: Stage Python package licenses
run: |
for package in \
hindsight-clients/python \
hindsight-api-slim \
hindsight-api \
hindsight-all \
hindsight-all-slim \
hindsight-embed; do
cp LICENSE "$package/LICENSE"
done
# Build all packages
- name: Build hindsight-client
working-directory: ./hindsight-clients/python
@@ -64,24 +50,6 @@ jobs:
working-directory: ./hindsight-embed
run: uv build --out-dir dist
- name: Verify Python package licenses
run: |
for package in \
hindsight-clients/python \
hindsight-api-slim \
hindsight-api \
hindsight-all \
hindsight-all-slim \
hindsight-embed; do
wheel=$(find "$package/dist" -maxdepth 1 -name '*.whl' -print -quit)
sdist=$(find "$package/dist" -maxdepth 1 -name '*.tar.gz' -print -quit)
unzip -Z1 "$wheel" | grep -Eq '\.dist-info/licenses/LICENSE$'
unzip -p "$wheel" '*/METADATA' | grep -Fxq 'License-Expression: MIT'
unzip -p "$wheel" '*/METADATA' | grep -Fxq 'License-File: LICENSE'
tar -tzf "$sdist" | grep -Eq '/LICENSE$'
done
# Publish in order (client and api-slim first, then api/all wrappers which depend on them)
- name: Publish hindsight-client to PyPI
uses: pypa/gh-action-pypi-publish@release/v1
@@ -327,50 +295,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 +569,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/
-29
View File
@@ -1,29 +0,0 @@
name: Update star history
on:
schedule:
- cron: '17 3 * * *'
workflow_dispatch:
permissions:
contents: write
jobs:
update:
concurrency:
group: star-history
cancel-in-progress: false
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v4
- uses: nicoloboschi/gh-stars@v1
with:
github-token: ${{ secrets.GITHUB_TOKEN }}
line-color: '#14b8a6'
- name: Commit chart
run: |
git config user.name 'github-actions[bot]'
git config user.email '41898282+github-actions[bot]@users.noreply.github.com'
git add .github/star-history/data.json .github/star-history/chart.svg
git diff --cached --quiet || git commit -m 'chore: update star history'
git push
+29 -427
View File
@@ -25,27 +25,22 @@ jobs:
cli: ${{ steps.filter.outputs.cli }}
docker: ${{ steps.filter.outputs.docker }}
helm: ${{ steps.filter.outputs.helm }}
doc-examples: ${{ steps.filter.outputs.doc-examples }}
docs: ${{ steps.filter.outputs.docs }}
embed: ${{ steps.filter.outputs.embed }}
all-npm: ${{ steps.filter.outputs.all-npm }}
hindsight-all: ${{ steps.filter.outputs.hindsight-all }}
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 }}
integrations-claude-code: ${{ steps.filter.outputs.integrations-claude-code }}
integrations-coding-agents: ${{ steps.filter.outputs.integrations-coding-agents }}
integrations-cline: ${{ steps.filter.outputs.integrations-cline }}
integrations-codex: ${{ steps.filter.outputs.integrations-codex }}
integrations-github-copilot: ${{ steps.filter.outputs.integrations-github-copilot }}
integrations-continue: ${{ steps.filter.outputs.integrations-continue }}
integrations-cursor-cli: ${{ steps.filter.outputs.integrations-cursor-cli }}
integrations-zcode: ${{ steps.filter.outputs.integrations-zcode }}
integrations-agent-plugin: ${{ steps.filter.outputs.integrations-agent-plugin }}
integrations-copilot-cli: ${{ steps.filter.outputs.integrations-copilot-cli }}
integrations-crewai: ${{ steps.filter.outputs.integrations-crewai }}
integrations-litellm: ${{ steps.filter.outputs.integrations-litellm }}
integrations-pydantic-ai: ${{ steps.filter.outputs.integrations-pydantic-ai }}
@@ -119,17 +114,12 @@ jobs:
- 'docker/**'
helm:
- 'helm/**'
# The RUNNABLE samples only. This replaces a broad `docs` filter that also
# matched 'hindsight-docs/**', '*.md' and 'hindsight-integrations/**' — the
# last of those so an integration rename would be caught by the docs build's
# integrations check, except that build (build-docs) has no `if:` and runs
# unconditionally anyway. test-doc-examples was the filter's only consumer,
# and it executes every sample against a live LLM-backed server, so every
# integration and prose-only PR paid for four provider-credentialed runs that
# none of those files can affect.
doc-examples:
- 'hindsight-docs/examples/**'
- 'scripts/test-doc-examples.sh'
docs:
- 'hindsight-docs/**'
- '*.md'
# Integration changes can add/rename integrations, which the docs
# build's integrations check validates against integrations.json.
- 'hindsight-integrations/**'
embed:
- 'hindsight-embed/**'
all-npm:
@@ -144,8 +134,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:
@@ -154,8 +142,6 @@ jobs:
- 'hindsight-integrations/chat/**'
integrations-claude-code:
- 'hindsight-integrations/claude-code/**'
integrations-coding-agents:
- 'hindsight-integrations/coding-agents/**'
integrations-cline:
- 'hindsight-integrations/cline/**'
integrations-codex:
@@ -166,8 +152,6 @@ jobs:
- 'hindsight-integrations/continue/**'
integrations-cursor-cli:
- 'hindsight-integrations/cursor-cli/**'
integrations-copilot-cli:
- 'hindsight-integrations/copilot-cli/**'
integrations-crewai:
- 'hindsight-integrations/crewai/**'
integrations-litellm:
@@ -196,10 +180,6 @@ jobs:
- 'hindsight-integrations/cursor/**'
integrations-zed:
- 'hindsight-integrations/zed/**'
integrations-zcode:
- 'hindsight-integrations/zcode/**'
integrations-agent-plugin:
- 'hindsight-integrations/agent-plugin/**'
integrations-n8n:
- 'hindsight-integrations/n8n/**'
integrations-zapier:
@@ -303,18 +283,6 @@ jobs:
working-directory: ./hindsight-api-slim
run: uv build
# `uv build` only packages the source; it does not prove the dependency set
# resolves or that the code imports on this interpreter. Install into a fresh
# env and run a byte-compile + import smoke test so the matrix actually
# exercises each Python version (notably 3.14).
- name: Install and smoke-test on Python ${{ matrix.python-version }}
working-directory: ./hindsight-api-slim
run: |
uv venv --python ${{ matrix.python-version }} .venv-smoke
VIRTUAL_ENV=.venv-smoke uv pip install .
.venv-smoke/bin/python -m compileall -q hindsight_api
.venv-smoke/bin/python -c "import hindsight_api, hindsight_api.main, hindsight_api.config; from hindsight_api.engine import memory_engine, llm_wrapper; print('import OK')"
build-typescript-client:
needs: [detect-changes]
if: >-
@@ -486,42 +454,6 @@ jobs:
working-directory: ./hindsight-integrations/openclaw
run: ./scripts/smoke-test.sh
test-coding-agents:
needs: [detect-changes]
if: >-
(github.event_name == 'workflow_dispatch' ||
needs.detect-changes.outputs.integrations-coding-agents == 'true' ||
needs.detect-changes.outputs.ci == 'true')
runs-on: ubuntu-latest
timeout-minutes: 15
steps:
- uses: actions/checkout@v6
with:
ref: ${{ github.event.pull_request.head.sha || '' }}
- uses: actions/setup-node@v6
with:
node-version: '22'
cache: npm
cache-dependency-path: hindsight-integrations/coding-agents/package-lock.json
- name: Install dependencies
working-directory: ./hindsight-integrations/coding-agents
run: npm ci
- name: Typecheck
working-directory: ./hindsight-integrations/coding-agents
run: npx tsc --noEmit
- name: Unit tests
working-directory: ./hindsight-integrations/coding-agents
run: npm test
- name: Build
working-directory: ./hindsight-integrations/coding-agents
run: npm run build
test-claude-code-integration:
needs: [detect-changes]
if: >-
@@ -588,17 +520,22 @@ jobs:
with:
ref: ${{ github.event.pull_request.head.sha || '' }}
- name: Set up Node.js
uses: actions/setup-node@v6
- name: Set up Python
uses: actions/setup-python@v6
with:
node-version: '22'
python-version: '3.11'
- name: Install package and pytest
working-directory: ./hindsight-integrations/zed
# Installs the package (incl. the zstandard runtime dep) so the threads.db
# reader tests can decompress Zed's zstd blobs.
run: pip install -e . pytest
- name: Run tests
working-directory: ./hindsight-integrations/zed
# Config-only integration with no dependencies — it uses Node's built-in
# test runner. The runtime MCP bridge is `npx mcp-remote` (Node), so this
# integration requires only Node.js (no Python).
run: npm test
# PR CI runs only the deterministic bucket; the real-LLM E2E bucket
# (requires_real_llm) needs a live Hindsight server and runs separately.
run: python -m pytest tests/ -v -m "not requires_real_llm"
test-omo-integration:
needs: [detect-changes]
@@ -765,66 +702,6 @@ jobs:
working-directory: ./hindsight-integrations/cursor-cli
run: uv run pytest tests -v
test-zcode-integration:
needs: [detect-changes]
if: >-
(github.event_name == 'workflow_dispatch' ||
needs.detect-changes.outputs.integrations-zcode == 'true' ||
needs.detect-changes.outputs.ci == 'true')
runs-on: ubuntu-latest
timeout-minutes: 30
steps:
- uses: actions/checkout@v6
with:
ref: ${{ github.event.pull_request.head.sha || '' }}
- name: Install uv
uses: astral-sh/setup-uv@v7
with:
enable-cache: true
prune-cache: false
- name: Set up Python
uses: actions/setup-python@v6
with:
python-version-file: ".python-version"
- name: Build zcode integration
working-directory: ./hindsight-integrations/zcode
run: uv build
- name: Install dependencies
working-directory: ./hindsight-integrations/zcode
run: uv sync --frozen
- name: Run tests
working-directory: ./hindsight-integrations/zcode
run: uv run pytest tests -v
test-agent-plugin-integration:
needs: [detect-changes]
if: >-
(github.event_name == 'workflow_dispatch' ||
needs.detect-changes.outputs.integrations-agent-plugin == 'true' ||
needs.detect-changes.outputs.ci == 'true')
runs-on: ubuntu-latest
timeout-minutes: 15
steps:
- uses: actions/checkout@v6
with:
ref: ${{ github.event.pull_request.head.sha || '' }}
- name: Set up Python
uses: actions/setup-python@v6
with:
python-version-file: ".python-version"
- name: Validate Agent Plugin manifests
working-directory: ./hindsight-integrations/agent-plugin
run: python3 validate.py
build-ai-sdk-integration:
needs: [detect-changes]
if: >-
@@ -888,37 +765,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: >-
@@ -1383,10 +1229,10 @@ jobs:
build-docs:
needs: [detect-changes]
# Keep the production docs build as an unconditional PR check. OpenAPI
# generation used to build the site again inside verify-generated-files;
# running the existing job for every PR preserves that coverage without
# serializing two full Docusaurus builds in the generated-files check.
if: >-
(github.event_name == 'workflow_dispatch' ||
needs.detect-changes.outputs.docs == 'true' ||
needs.detect-changes.outputs.ci == 'true')
runs-on: ubuntu-latest
timeout-minutes: 30
@@ -1455,30 +1301,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
@@ -2014,27 +1836,6 @@ jobs:
with:
ref: ${{ github.event.pull_request.head.sha || '' }}
# The Oracle 23ai `free` service image is large; combined with the Python ML
# deps (torch) it exhausts the runner's ~14 GB root disk, so uv fails to
# extract a wheel with "No space left on device (os error 28)".
# Only the cheap, high-yield reclaims are enabled — the Android/.NET/Haskell
# dirs plus swap are a few `rm -rf`s worth ~16-21 GB, which is ample headroom:
# - large-packages runs apt-get remove and costs minutes for little gain;
# - tool-cache would delete the preinstalled Python that actions/setup-python
# then has to re-download, making the job slower, not faster;
# - docker-images is pointless here (the Oracle service container is already
# running, so its image is in use and cannot be pruned anyway).
- name: Free Disk Space
uses: jlumbroso/free-disk-space@main
with:
tool-cache: false
android: true
dotnet: true
haskell: true
large-packages: false
docker-images: false
swap-storage: true
- name: Setup Oracle test user
# The SYSTEM tablespace uses manual segment space management which
# doesn't support VECTOR types. Create an ASSM tablespace and a
@@ -2395,27 +2196,6 @@ jobs:
with:
ref: ${{ github.event.pull_request.head.sha || '' }}
# The Oracle 23ai `free` service image is large; combined with the Python ML
# deps (torch) it exhausts the runner's ~14 GB root disk, so uv fails to
# extract a wheel with "No space left on device (os error 28)".
# Only the cheap, high-yield reclaims are enabled — the Android/.NET/Haskell
# dirs plus swap are a few `rm -rf`s worth ~16-21 GB, which is ample headroom:
# - large-packages runs apt-get remove and costs minutes for little gain;
# - tool-cache would delete the preinstalled Python that actions/setup-python
# then has to re-download, making the job slower, not faster;
# - docker-images is pointless here (the Oracle service container is already
# running, so its image is in use and cannot be pruned anyway).
- name: Free Disk Space
uses: jlumbroso/free-disk-space@main
with:
tool-cache: false
android: true
dotnet: true
haskell: true
large-packages: false
docker-images: false
swap-storage: true
- name: Setup Oracle test user
run: |
pip install oracledb
@@ -2576,27 +2356,6 @@ jobs:
with:
ref: ${{ github.event.pull_request.head.sha || '' }}
# The Oracle 23ai `free` service image is large; combined with the Python ML
# deps (torch) it exhausts the runner's ~14 GB root disk, so uv fails to
# extract a wheel with "No space left on device (os error 28)".
# Only the cheap, high-yield reclaims are enabled — the Android/.NET/Haskell
# dirs plus swap are a few `rm -rf`s worth ~16-21 GB, which is ample headroom:
# - large-packages runs apt-get remove and costs minutes for little gain;
# - tool-cache would delete the preinstalled Python that actions/setup-python
# then has to re-download, making the job slower, not faster;
# - docker-images is pointless here (the Oracle service container is already
# running, so its image is in use and cannot be pruned anyway).
- name: Free Disk Space
uses: jlumbroso/free-disk-space@main
with:
tool-cache: false
android: true
dotnet: true
haskell: true
large-packages: false
docker-images: false
swap-storage: true
- name: Setup Oracle test user
run: |
pip install oracledb
@@ -3722,43 +3481,6 @@ jobs:
# (requires_real_llm) needs a live Hindsight server and runs separately.
run: uv run pytest tests -v -m "not requires_real_llm"
test-copilot-cli-integration:
needs: [detect-changes]
if: >-
(github.event_name == 'workflow_dispatch' ||
needs.detect-changes.outputs.integrations-copilot-cli == 'true' ||
needs.detect-changes.outputs.ci == 'true')
runs-on: ubuntu-latest
timeout-minutes: 30
steps:
- uses: actions/checkout@v6
with:
ref: ${{ github.event.pull_request.head.sha || '' }}
- name: Install uv
uses: astral-sh/setup-uv@v7
with:
enable-cache: true
prune-cache: false
- name: Set up Python
uses: actions/setup-python@v6
with:
python-version-file: ".python-version"
- name: Build copilot-cli integration
working-directory: ./hindsight-integrations/copilot-cli
run: uv build
- name: Install dependencies
working-directory: ./hindsight-integrations/copilot-cli
run: uv sync --frozen
- name: Run tests
working-directory: ./hindsight-integrations/copilot-cli
run: uv run pytest tests -v
test-crewai-integration:
needs: [detect-changes]
if: >-
@@ -4617,73 +4339,6 @@ jobs:
fi
done || true
# Compatibility gate against hermes-agent's *main* branch.
#
# `hermes memory setup` installs `hindsight-all` into Hermes' own venv to run
# memory in local_embedded mode, and Hermes exact-pins every direct dependency
# (`==X.Y.Z`) as a deliberate supply-chain policy — they will not loosen a pin
# for us. So any version range Hindsight declares that excludes one of their
# pins makes the two impossible to co-install for every Hermes user on
# embedded memory. That was #3251: our `cryptography>=48.0.1` / `pillow>=12.3.0`
# against their `==46.0.7` / `==12.2.0`, which left `pip check` permanently
# broken. Both sides bump on their own schedule, so this needs a standing gate
# rather than a one-off fix; tracking main surfaces the next collision while
# it is still cheap to fix on either side.
#
# Deliberately not gated on has_secrets — the resolution, wiring, runtime and
# daemon-boot checks need no credentials, so this runs on fork PRs too. Only
# retain/recall need an LLM and the script skips them when no key is present.
test-hermes-compat:
needs: [detect-changes]
if: >-
(github.event_name == 'workflow_dispatch' ||
needs.detect-changes.outputs.core == 'true' ||
needs.detect-changes.outputs.embed == 'true' ||
needs.detect-changes.outputs.hindsight-all == 'true' ||
needs.detect-changes.outputs.clients-python == 'true' ||
needs.detect-changes.outputs.ci == 'true')
runs-on: ubuntu-latest
timeout-minutes: 45
permissions:
contents: read
steps:
- uses: actions/checkout@v6
with:
ref: ${{ github.event.pull_request.head.sha || '' }}
- name: Install uv
uses: astral-sh/setup-uv@v7
with:
enable-cache: true
prune-cache: false
# hindsight-all pulls the local-ml extra, so the embedded daemon can load
# sentence-transformers models. Cache them like the test-embed job does.
- name: Cache HuggingFace models
uses: actions/cache@v5
with:
path: ~/.cache/huggingface
key: ${{ runner.os }}-huggingface-hermes-compat-${{ hashFiles('hindsight-api-slim/pyproject.toml') }}
restore-keys: |
${{ runner.os }}-huggingface-hermes-compat-
${{ runner.os }}-huggingface-
- name: Run Hermes compatibility test
env:
GITHUB_TOKEN: ${{ github.token }}
run: ./scripts/test-hermes-compat.sh
- name: Collect embedded daemon logs on failure
if: failure()
run: |
for f in ~/.hindsight/profiles/hermes-ci*.log ~/.hindsight/profiles/hermes-ci*.stderr.log; do
if [ -f "$f" ]; then
echo "=== $f ==="
cat "$f"
fi
done || true
test-hindsight-all:
needs: [detect-changes]
if: >-
@@ -4785,7 +4440,7 @@ jobs:
needs.detect-changes.outputs.clients-python == 'true' ||
needs.detect-changes.outputs.clients-go == 'true' ||
needs.detect-changes.outputs.cli == 'true' ||
needs.detect-changes.outputs.doc-examples == 'true' ||
needs.detect-changes.outputs.docs == 'true' ||
needs.detect-changes.outputs.ci == 'true')
runs-on: ubuntu-latest
timeout-minutes: 30
@@ -5116,8 +4771,7 @@ jobs:
cd ../hindsight-embed && uv sync --frozen --index-strategy unsafe-best-match
- name: Run generate-openapi
working-directory: hindsight-dev
run: uv run generate-openapi
run: ./scripts/generate-openapi.sh
- name: Run generate-bank-template-schema
run: ./scripts/generate-bank-template-schema.sh
@@ -5189,26 +4843,13 @@ jobs:
exit 0
fi
# Compare against the point this branch was cut from, NOT the live tip of
# the base branch. Using the tip reports every endpoint main has gained
# since the branch was cut as "removed by this PR" — a false positive that
# fails PRs which touch no spec at all, and whose only cure is an unrelated
# rebase. The merge-base answers the question the check actually asks:
# did *this branch* remove something?
MERGE_BASE="$(git merge-base "origin/$BASE_BRANCH" HEAD)"
echo "Checking OpenAPI compatibility against base branch: $BASE_BRANCH"
if [ -z "$MERGE_BASE" ]; then
echo "⚠️ Warning: Could not determine merge-base with $BASE_BRANCH. Skipping compatibility check."
exit 0
fi
echo "Checking OpenAPI compatibility against $BASE_BRANCH merge-base: $MERGE_BASE"
# Extract the old OpenAPI spec from the merge-base
git show "$MERGE_BASE:hindsight-docs/static/openapi.json" > /tmp/old-openapi.json
# Extract the old OpenAPI spec from base branch
git show "origin/$BASE_BRANCH:hindsight-docs/static/openapi.json" > /tmp/old-openapi.json
if [ ! -s /tmp/old-openapi.json ]; then
echo "⚠️ Warning: Could not find OpenAPI spec at the merge-base. Skipping compatibility check."
echo "⚠️ Warning: Could not find OpenAPI spec in base branch. Skipping compatibility check."
exit 0
fi
@@ -5250,42 +4891,6 @@ jobs:
cd hindsight-dev
uv run cli-coverage-check
# hindsight-dev/tests had no job of its own, so nothing ran it: the benchmark
# harness was only exercised by the nightly Performance Tests workflow, where a
# plain construction bug in the answer/judge LLM config surfaced as a red
# benchmark hours later instead of on the PR that introduced it.
test-dev:
needs: [detect-changes]
if: >-
(github.event_name == 'workflow_dispatch' ||
needs.detect-changes.outputs.core == 'true' ||
needs.detect-changes.outputs.dev == 'true' ||
needs.detect-changes.outputs.ci == 'true')
runs-on: ubuntu-latest
timeout-minutes: 30
steps:
- uses: actions/checkout@v6
with:
ref: ${{ github.event.pull_request.head.sha || '' }}
- name: Install uv
uses: astral-sh/setup-uv@v7
with:
enable-cache: true
- name: Set up Python
uses: actions/setup-python@v6
with:
python-version-file: ".python-version"
- name: Install hindsight-dev dependencies
run: |
cd hindsight-dev && uv sync --frozen --extra test --index-strategy unsafe-best-match
- name: Run hindsight-dev tests
working-directory: hindsight-dev
run: uv run pytest tests/ -v
# Report CI status back to the PR for pull_request_review events.
# GitHub does not automatically link pull_request_review check runs to the PR,
# so we create a commit status on the PR head SHA and post a comment.
@@ -5304,8 +4909,6 @@ jobs:
- test-github-copilot-integration
- test-codex-integration
- test-cursor-cli-integration
- test-zcode-integration
- test-agent-plugin-integration
- build-ai-sdk-integration
- test-ai-sdk-integration-deno
- test-opencode-integration
@@ -5361,7 +4964,6 @@ jobs:
- test-embed
- test-embed-windows
- verify-embed-control-center-bundle
- test-hermes-compat
- test-hindsight-all
- test-hindsight-agent-sdk
- test-claude-agent-sdk-integration
-11
View File
@@ -5,13 +5,6 @@ build/
dist/
wheels/
*.egg-info
# Release builds stage the canonical root license in each package context.
/hindsight-clients/python/LICENSE
/hindsight-api-slim/LICENSE
/hindsight-api/LICENSE
/hindsight-all/LICENSE
/hindsight-all-slim/LICENSE
/hindsight-embed/LICENSE
.mcp.json
.playwright-mcp/
.osgrep
@@ -20,9 +13,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
@@ -51,7 +41,6 @@ nltk_data/
logs/
.DS_Store
.sesskey
# Generated docs files
hindsight-docs/static/llms-full.txt
+1
View File
@@ -0,0 +1 @@
fcac2839-1db5-432f-91e1-c5dac07d7290
-29
View File
@@ -286,35 +286,6 @@ When adding or modifying parameters in the dataplane API (hindsight-api), you mu
- Update the client type definition in `lib/api.ts`
- Update any UI components that need to use the new parameter
### Harness Attribution (which coding agent wrote a document)
`hindsight-integrations/hindsight-coding-agents/` stamps the coding agent on every
document it retains, so the control plane can show its logo instead of another
`key=value` chip:
- `metadata.harness = "<id>"` — the authoritative field
- tag `harness:<id>` — the same value, so the documents list can filter on it
The ids are defined by that integration's HookSpecs
(`src/harness/hook-lifecycle.ts`) plus the persistent-plugin entrypoints
registered in `src/harness/registry.ts`, whose id is their
`createPluginEntry(...)` argument — currently `antigravity-cli`, `claude-code`,
`cline-cli`, `codex`, `copilot-cli`, `cursor-cli`, `devin-cli`, `grok-build`,
`kilo`, `opencode`.
The control plane resolves the value in
`hindsight-control-plane/src/lib/harness-logo.ts` (metadata wins over the tag) and
renders it with `components/ui/harness-logo.tsx` in the documents table and the
document detail dialog. **Adding a harness to the integration means adding it to
that registry in the same change**: copy its icon from
`hindsight-docs/static/img/icons/` (or take it from the agent's own brand assets
when the docs site carries none) into
`hindsight-control-plane/public/img/harness/` and add one entry. Don't register
ids nothing writes — a test asserts the registry matches the emitted set, plus an
explicit list of retired ids kept so already-retained documents keep their logo.
An unregistered harness is not an error: it renders no logo and still shows as
ordinary metadata.
### Adding New Integrations
Every new integration in `hindsight-integrations/` must satisfy all of the following before it can be merged:
+1 -1
View File
@@ -298,7 +298,7 @@ client.reflect(bank_id="my-bank", query="What should I know about Alice?")
---
## Star History
[![Star history](https://raw.githubusercontent.com/vectorize-io/hindsight/main/.github/star-history/chart.svg)](https://github.com/vectorize-io/hindsight/stargazers)
[![Star History Chart](https://api.star-history.com/svg?repos=vectorize-io/hindsight&type=date&legend=top-left)](https://www.star-history.com/#vectorize-io/hindsight&type=date&legend=top-left)
---
## Supported Platforms
Binary file not shown.

Before

Width:  |  Height:  |  Size: 303 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 152 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 153 KiB

@@ -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}
@@ -17,7 +17,7 @@ FROM ghcr.io/vectorize-io/hindsight:latest-slim
# `pip install` would fall back to user site-packages and not be visible
# to the runtime python.
RUN uv pip install --python /app/api/.venv/bin/python --no-cache \
'sentence-transformers>=5.0.0' \
'sentence-transformers>=3.3.0' \
'transformers>=4.53.0' \
'torch>=2.6.0'
@@ -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}
+14 -4
View File
@@ -3,11 +3,21 @@
# pgroonga is a multilingual full-text search extension built on Groonga.
# It works out of the box for CJK (Chinese, Japanese, Korean) and other
# non-whitespace-segmented languages via the TokenBigram tokenizer.
FROM groonga/pgroonga:4.0.8-debian-17
FROM groonga/pgroonga:latest-debian-pg17
# Install pgvector on top of the pgroonga base image (which already provides
# pgroonga, the Groonga library, and the PostgreSQL PGDG package repository).
# pgroonga and the Groonga library).
RUN apt-get update && apt-get install -y --no-install-recommends \
postgresql-17-pgvector=0.8.6-1.pgdg13+1 \
&& apt-get clean \
build-essential \
git \
postgresql-server-dev-17 \
&& rm -rf /var/lib/apt/lists/*
RUN cd /tmp && \
git clone --branch v0.8.0 https://github.com/pgvector/pgvector.git && \
cd pgvector && \
make && \
make install
RUN rm -rf /tmp/pgvector && \
apt-get purge -y --auto-remove build-essential git postgresql-server-dev-17
@@ -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
-107
View File
@@ -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:
+7 -2
View File
@@ -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
+3 -3
View File
@@ -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}
+9 -30
View File
@@ -41,43 +41,28 @@ RUN apt-get update && apt-get install -y \
&& rm -rf /var/lib/apt/lists/* \
&& pip install --no-cache-dir uv
# Copy the workspace lock and member metadata before source code so dependency
# installation stays cacheable while matching the versions tested in CI.
COPY pyproject.toml uv.lock ./
COPY hindsight-all/pyproject.toml ./hindsight-all/
COPY hindsight-api/pyproject.toml ./hindsight-api/
# Copy dependency files and README (required by pyproject.toml)
COPY hindsight-api-slim/pyproject.toml ./api/
COPY hindsight-api-slim/README.md ./api/
COPY hindsight-all-slim/pyproject.toml ./hindsight-all-slim/
COPY hindsight-dev/pyproject.toml ./hindsight-dev/
COPY hindsight-clients/python/pyproject.toml ./hindsight-clients/python/
COPY hindsight-embed/pyproject.toml ./hindsight-embed/
RUN ln -s api hindsight-api-slim
WORKDIR /app/api
# Sync dependencies using appropriate extras based on INCLUDE_LOCAL_MODELS
# local-ml: torch, sentence-transformers, transformers, einops, flashrank, mlx (optional)
# embedded-db: pg0-embedded (always included for embedded PostgreSQL support)
# ONNX Runtime embeddings are intentionally not bundled into the official
# standalone image; install the local-onnx extra in custom images when needed.
ENV UV_PROJECT_ENVIRONMENT=/app/api/.venv
RUN if [ "$INCLUDE_LOCAL_MODELS" = "true" ]; then \
uv sync --locked --package hindsight-api-slim --no-install-package hindsight-api-slim --extra local-ml --extra embedded-db; \
uv sync --extra local-ml --extra embedded-db; \
else \
uv sync --locked --package hindsight-api-slim --no-install-package hindsight-api-slim --extra embedded-db; \
uv sync --extra embedded-db; \
fi
# Copy source code (alembic migrations are inside hindsight_api/)
WORKDIR /app/api
COPY hindsight-api-slim/hindsight_api ./hindsight_api
# Install the local package from the same validated lock after source is present.
WORKDIR /app
RUN if [ "$INCLUDE_LOCAL_MODELS" = "true" ]; then \
uv sync --locked --package hindsight-api-slim --extra local-ml --extra embedded-db; \
else \
uv sync --locked --package hindsight-api-slim --extra embedded-db; \
fi \
&& uv pip check --python /app/api/.venv/bin/python
# Install the local package (uv sync only installed dependencies, not the package itself)
RUN uv pip install -e .
# =============================================================================
# Stage: SDK Builder (needed for Control Plane)
@@ -160,8 +145,6 @@ FROM python:3.11-slim AS api-only
WORKDIR /app
# Note: libicu version varies by Debian version - try common versions in order
# Runtime images use uv directly; remove pip build tooling after installation so
# vulnerable setuptools-vendored packages and wheel are not shipped in production.
RUN apt-get update && apt-get install -y \
curl \
procps \
@@ -171,8 +154,7 @@ RUN apt-get update && apt-get install -y \
libossp-uuid16 \
&& (apt-get install -y libicu72 2>/dev/null || apt-get install -y libicu74 2>/dev/null || apt-get install -y libicu76 2>/dev/null || true) \
&& rm -rf /var/lib/apt/lists/* \
&& pip install --no-cache-dir uv \
&& pip uninstall --yes setuptools wheel
&& pip install --no-cache-dir uv
RUN useradd -m -s /bin/bash hindsight
@@ -310,8 +292,6 @@ WORKDIR /app
# Install Node.js, curl, uv, and system dependencies
# Note: libicu version varies by Debian version - try common versions in order
# Runtime images use uv directly; remove pip build tooling after installation so
# vulnerable setuptools-vendored packages and wheel are not shipped in production.
RUN apt-get update && apt-get install -y \
curl \
procps \
@@ -323,8 +303,7 @@ RUN apt-get update && apt-get install -y \
&& curl -fsSL https://deb.nodesource.com/setup_20.x | bash - \
&& apt-get install -y nodejs \
&& rm -rf /var/lib/apt/lists/* \
&& pip install --no-cache-dir uv \
&& pip uninstall --yes setuptools wheel
&& pip install --no-cache-dir uv
RUN useradd -m -s /bin/bash hindsight
+2 -2
View File
@@ -2,8 +2,8 @@ apiVersion: v2
name: hindsight
description: Hindsight helm chart
type: application
version: 0.9.1
appVersion: "0.9.1"
version: 0.8.4
appVersion: "0.8.4"
keywords:
- ai
- memory
@@ -60,13 +60,13 @@ spec:
valueFrom:
fieldRef:
fieldPath: metadata.name
{{- /* Explicitly set port to override K8s service discovery env var (HINDSIGHT_API_PORT) */}}
- name: HINDSIGHT_API_PORT
value: {{ .Values.worker.service.targetPort | quote }}
{{- /* Inherit LLM config from api.env, then apply worker-specific env.
Merge (worker.env wins) so a key set in both does not emit a
duplicate env entry, which server-side apply rejects. */}}
{{- range $key, $value := merge (deepCopy (.Values.worker.env | default dict)) (.Values.api.env | default dict) }}
{{- /* Inherit LLM config from api.env */}}
{{- range $key, $value := .Values.api.env }}
- name: {{ $key }}
value: {{ $value | quote }}
{{- end }}
{{- /* Worker-specific env vars */}}
{{- range $key, $value := .Values.worker.env }}
- name: {{ $key }}
value: {{ $value | quote }}
{{- end }}
+4 -15
View File
@@ -36,22 +36,16 @@ api:
cpu: 500m
memory: 1Gi
# Liveness and readiness probes.
# Liveness uses /health/live, which performs no database access: a slow or
# unreachable database must gate traffic (readiness), never restart pods.
# Needs an image from this chart's appVersion or newer — older ones serve
# /health only, and would fail this probe with a 404.
# Liveness and readiness probes
livenessProbe:
httpGet:
path: /health/live
path: /health
port: 8888
initialDelaySeconds: 30
periodSeconds: 10
timeoutSeconds: 5
failureThreshold: 3
# Readiness checks the database, so a pod that cannot reach it is pulled out
# of the Service and put back once the database recovers.
readinessProbe:
httpGet:
path: /health
@@ -137,15 +131,10 @@ worker:
cpu: 500m
memory: 1Gi
# Liveness and readiness probes.
# Liveness uses /health/live, which performs no database access. Restarting a
# worker whose database is merely slow requeues its claimed operations with
# retry_count incremented, so DB checks must stay out of liveness.
# Needs an image from this chart's appVersion or newer — older ones serve
# /health only, and would fail this probe with a 404.
# Liveness and readiness probes
livenessProbe:
httpGet:
path: /health/live
path: /health
port: 8889
initialDelaySeconds: 30
periodSeconds: 10
+1 -1
View File
@@ -1,6 +1,6 @@
{
"name": "@vectorize-io/hindsight-all",
"version": "0.9.1",
"version": "0.8.4",
"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 -5
View File
@@ -1,18 +1,17 @@
[build-system]
requires = ["setuptools>=77"]
requires = ["setuptools>=61"]
build-backend = "setuptools.build_meta"
[project]
name = "hindsight-all-slim"
version = "0.9.1"
version = "0.8.4"
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.8.4",
"hindsight-client>=0.0.7",
"hindsight-embed==0.9.1",
"hindsight-embed>=0.1.0",
]
[tool.uv.sources]
+5 -6
View File
@@ -1,18 +1,17 @@
[build-system]
requires = ["hatchling>=1.27"]
requires = ["hatchling"]
build-backend = "hatchling.build"
[project]
name = "hindsight-all"
version = "0.9.1"
version = "0.8.4"
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.8.4",
"hindsight-client>=0.0.7",
"hindsight-embed==0.9.1",
"hindsight-embed>=0.1.0",
]
[tool.uv.sources]
@@ -22,7 +21,7 @@ hindsight-embed = { workspace = true }
[project.optional-dependencies]
local-llm = [
"hindsight-api-slim[local-llm]==0.9.1",
"hindsight-api-slim[local-llm]==0.8.4",
]
test = [
"pytest>=7.0.0",
+1 -1
View File
@@ -53,4 +53,4 @@ __all__ = [
"RemoteTEICrossEncoder",
"LLMConfig",
]
__version__ = "0.9.1"
__version__ = "0.8.4"
@@ -1,23 +0,0 @@
"""Text-search SQL shapes shared by index DDL and the queries that must hit it.
A PostgreSQL expression index is only selectable when the query repeats the
indexed expression verbatim, so the DDL (``migrations.py`` and the Alembic
versions) and the read arms (``engine/sql/postgresql.py``) cannot be allowed to
drift. Both sides call the helpers here — same idea as
``_pg_search.pg_search_bm25_columns``.
"""
def mental_models_text_document(alias: str | None = None) -> str:
"""The ``mental_models`` full-text document: model/page name + content.
Mirrors the generating expression of the native tsvector column created by
the ``n9i0j1k2l3m4`` (learnings / pinned_reflections) migration, so every
backend indexes and queries the exact same document. ``content`` is NOT NULL,
hence the deliberate lack of a ``COALESCE`` around it.
``alias`` qualifies the columns for queries that join the table (``mm``);
leave it unset for DDL, where the expression is already table-scoped.
"""
prefix = f"{alias}." if alias else ""
return f"(COALESCE({prefix}name, '') || ' ' || {prefix}content)"
+32 -419
View File
@@ -8,9 +8,7 @@ import asyncio
import io
import json
import logging
import struct
import zipfile
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Any
@@ -18,12 +16,10 @@ from typing import Any
import asyncpg
import typer
from ..config import DEFAULT_DATABASE_SCHEMA, HindsightConfig, load_dotenv_for_entrypoint
from ..config import DEFAULT_DATABASE_SCHEMA, HindsightConfig
from ..engine.memory_engine import _current_schema
from ..engine.retain.bank_utils import _vector_index_clause
from ..engine.schema import fq_table_explicit as _fq_table
from ..engine.transfer import export_bank
from ..engine.vector_index_health import SchemaVectorIndexResult, repair_vector_indexes
from ..extensions import TenantExtension, load_extension
from ..pg0 import parse_pg0_url, resolve_database_url
@@ -68,176 +64,9 @@ BACKUP_TABLES = [
"audit_log",
"llm_requests",
"graph_maintenance_queue",
"entity_maintenance_queue",
]
MANIFEST_VERSION = "2"
@dataclass(frozen=True)
class BackupColumn:
"""A PostgreSQL column shape required to decode a binary COPY stream."""
name: str
type_name: str
@dataclass(frozen=True)
class TableRestorePlan:
"""How one table's backed-up binary COPY stream is replayed onto the target.
``columns`` is the target column list handed to ``copy_to_table``, in stream
order. When the target no longer has a backed-up column, its field is stripped
from every tuple (``dropped_field_indices``) before the stream is replayed —
binary COPY is positional, so the column list and the tuple fields must agree.
"""
columns: list[str]
dropped_field_indices: tuple[int, ...]
source_field_count: int
# Header of a PostgreSQL binary COPY stream: an 11-byte signature, an int32 flags
# field, and an int32 header-extension length followed by that many bytes.
_COPY_BINARY_SIGNATURE = b"PGCOPY\n\xff\r\n\x00"
_COPY_BINARY_HEADER_LEN = len(_COPY_BINARY_SIGNATURE) + 8
def _strip_binary_copy_fields(data: bytes, plan: TableRestorePlan) -> bytes:
"""Drop `plan.dropped_field_indices` from every tuple of a binary COPY stream.
Restore used to reject a backup whose columns the target no longer had — the
preflight raised "target is missing backup columns …", which made any backup
taken before a column-dropping migration unrestorable afterwards. Those columns
are now ignored instead, but they cannot simply be left out of the
``copy_to_table`` column list: binary COPY carries no column identities, so each
tuple's fields are matched to the column list purely by position and an unedited
stream would desynchronise (or, worse, land values in the wrong columns). So the
stream itself is rewritten here.
Tuple format: int16 field count, then per field an int32 length (-1 for NULL)
followed by that many bytes. An int16 of -1 is the end-of-data trailer.
"""
if not plan.dropped_field_indices:
return data
if not data.startswith(_COPY_BINARY_SIGNATURE):
raise ValueError("Backup stream is not in PostgreSQL binary COPY format")
(extension_len,) = struct.unpack_from("!i", data, len(_COPY_BINARY_SIGNATURE) + 4)
pos = _COPY_BINARY_HEADER_LEN + extension_len
out = bytearray(data[:pos])
dropped = set(plan.dropped_field_indices)
kept_count = plan.source_field_count - len(dropped)
while True:
(field_count,) = struct.unpack_from("!h", data, pos)
pos += 2
if field_count == -1: # end-of-data trailer
out += struct.pack("!h", -1)
break
if field_count != plan.source_field_count:
raise ValueError(
f"Backup stream tuple has {field_count} fields, manifest declares {plan.source_field_count}"
)
out += struct.pack("!h", kept_count)
for index in range(field_count):
(length,) = struct.unpack_from("!i", data, pos)
pos += 4
payload = b"" if length == -1 else data[pos : pos + length]
pos += max(length, 0)
if index in dropped:
continue
out += struct.pack("!i", length)
out += payload
return bytes(out)
async def _table_columns(conn: asyncpg.Connection, schema: str, table: str) -> list[BackupColumn]:
rows = await conn.fetch(
"""
SELECT a.attname AS name, pg_catalog.format_type(a.atttypid, a.atttypmod) AS type_name
FROM pg_catalog.pg_attribute AS a
JOIN pg_catalog.pg_class AS c ON c.oid = a.attrelid
JOIN pg_catalog.pg_namespace AS n ON n.oid = c.relnamespace
WHERE n.nspname = $1 AND c.relname = $2 AND a.attnum > 0 AND NOT a.attisdropped
AND a.attgenerated = ''
ORDER BY a.attnum
""",
schema,
table,
)
return [BackupColumn(name=row["name"], type_name=row["type_name"]) for row in rows]
async def _validate_restore_schema(
conn: asyncpg.Connection, manifest: dict[str, Any], schema: str
) -> dict[str, TableRestorePlan]:
"""Validate every COPY stream against the target before destructive work starts.
A column the target no longer has is **not** an error: a migration that drops a
column would otherwise make every backup taken before it permanently
unrestorable. Such columns are skipped (their fields are stripped from the
stream by ``_strip_binary_copy_fields``) and reported, so the operator sees what
was discarded instead of the restore failing outright.
Type mismatches remain fatal. Type equality is an exact ``format_type`` string
match. This is deliberately stricter than binary-COPY wire compatibility (e.g.
``varchar`` and ``text`` share a binary format yet compare unequal here): we
would rather fail a genuinely-restorable backup with a clear, actionable error
than silently risk a subtle binary mismatch. Restores blocked this way can be
recovered by aligning the target schema.
"""
plans: dict[str, TableRestorePlan] = {}
errors: list[str] = []
for table, table_manifest in manifest["tables"].items():
source_columns = [BackupColumn(**column) for column in table_manifest["columns"]]
target_by_name = {column.name: column for column in await _table_columns(conn, schema, table)}
unknown = [
(index, column.name) for index, column in enumerate(source_columns) if column.name not in target_by_name
]
mismatched = [
f"{column.name} ({column.type_name} in backup, {target_by_name[column.name].type_name} in target)"
for column in source_columns
if column.name in target_by_name and target_by_name[column.name].type_name != column.type_name
]
if mismatched:
errors.append(f"{table}: incompatible column types: {', '.join(mismatched)}")
if unknown:
typer.echo(
f" {table}: ignoring {len(unknown)} backup column(s) absent from the target schema: "
f"{', '.join(name for _, name in unknown)}"
)
plans[table] = TableRestorePlan(
columns=[column.name for column in source_columns if column.name in target_by_name],
dropped_field_indices=tuple(index for index, _ in unknown),
source_field_count=len(source_columns),
)
if errors:
details = "; ".join(errors)
raise ValueError(f"Backup schema is incompatible with target schema '{schema}': {details}")
return plans
def _effective_backup_tables() -> list[str]:
"""Core backup tables plus any bank-scoped tables a loaded extension declares.
``BACKUP_TABLES`` covers only the tables core owns. An extension that
provisions its own bank-scoped tables (via ``TenantExtension``) declares
them through ``extra_bank_tables()`` so they aren't dropped on restore.
Extension tables are appended *after* the core set so restore's forward
COPY inserts them after their FK parents (e.g. ``banks``) and the reversed
TRUNCATE clears them before those parents.
"""
tables = list(BACKUP_TABLES)
tenant_extension = load_extension("TENANT", TenantExtension)
if tenant_extension is not None:
seen = set(tables)
for spec in tenant_extension.extra_bank_tables():
if spec.include_in_backup and spec.name not in seen:
tables.append(spec.name)
seen.add(spec.name)
return tables
MANIFEST_VERSION = "1"
async def _admin_connect(db_url: str) -> asyncpg.Connection:
@@ -248,8 +77,7 @@ async def _admin_connect(db_url: str) -> asyncpg.Connection:
is the only step needed to connect. JSON codecs are registered so ``jsonb``
columns decode to Python objects (used by the export row dumps).
"""
_pg0 = parse_pg0_url(db_url)
is_pg0, instance_name = _pg0.is_pg0, _pg0.instance_name
is_pg0, instance_name, _ = parse_pg0_url(db_url)
if is_pg0:
typer.echo(f"Starting embedded PostgreSQL (instance: {instance_name})...")
conn = await asyncpg.connect(await resolve_database_url(db_url))
@@ -258,18 +86,8 @@ async def _admin_connect(db_url: str) -> asyncpg.Connection:
return conn
async def _backup(
database_url: str,
output_path: Path,
schema: str = "public",
backup_tables: list[str] | None = None,
) -> dict[str, Any]:
"""Backup all tables to a zip file using binary COPY protocol.
``backup_tables`` defaults to the core ``BACKUP_TABLES``; callers pass the
extension-augmented list from ``_effective_backup_tables()``.
"""
backup_tables = backup_tables if backup_tables is not None else BACKUP_TABLES
async def _backup(database_url: str, output_path: Path, schema: str = "public") -> dict[str, Any]:
"""Backup all tables to a zip file using binary COPY protocol."""
conn = await asyncpg.connect(database_url)
try:
tables: dict[str, Any] = {}
@@ -286,24 +104,14 @@ async def _backup(
# entities table was backed up.
async with conn.transaction(isolation="repeatable_read"):
with zipfile.ZipFile(output_path, "w", zipfile.ZIP_DEFLATED) as zf:
for i, table in enumerate(backup_tables, 1):
typer.echo(f" [{i}/{len(backup_tables)}] Backing up {table}...", nl=False)
for i, table in enumerate(BACKUP_TABLES, 1):
typer.echo(f" [{i}/{len(BACKUP_TABLES)}] Backing up {table}...", nl=False)
buffer = io.BytesIO()
columns = await _table_columns(conn, schema, table)
# Pin the ordered columns into both the stream and manifest.
# PostgreSQL binary COPY does not encode column identities, so
# restore must validate this shape before truncating any data.
# Use binary COPY for exact type preservation
# asyncpg requires schema_name as separate parameter
await conn.copy_from_table(
table,
schema_name=schema,
columns=[column.name for column in columns],
output=buffer,
format="binary",
)
await conn.copy_from_table(table, schema_name=schema, output=buffer, format="binary")
data = buffer.getvalue()
zf.writestr(f"{table}.bin", data)
@@ -314,7 +122,6 @@ async def _backup(
tables[table] = {
"rows": row_count,
"size_bytes": len(data),
"columns": [{"name": column.name, "type_name": column.type_name} for column in columns],
}
typer.echo(f" {row_count} rows")
@@ -326,20 +133,8 @@ async def _backup(
await conn.close()
async def _restore(
database_url: str,
input_path: Path,
schema: str = "public",
backup_tables: list[str] | None = None,
) -> dict[str, Any]:
"""Restore all tables from a zip file using binary COPY protocol.
``backup_tables`` defaults to the core ``BACKUP_TABLES``; callers pass the
extension-augmented list from ``_effective_backup_tables()``. Tables named
here but absent from the archive are truncated then skipped for restore, so
a stale extension registration never leaves pre-restore rows behind.
"""
backup_tables = backup_tables if backup_tables is not None else BACKUP_TABLES
async def _restore(database_url: str, input_path: Path, schema: str = "public") -> dict[str, Any]:
"""Restore all tables from a zip file using binary COPY protocol."""
conn = await asyncpg.connect(database_url)
try:
with zipfile.ZipFile(input_path, "r") as zf:
@@ -348,42 +143,29 @@ async def _restore(
if manifest.get("version") != MANIFEST_VERSION:
raise ValueError(f"Unsupported backup version: {manifest.get('version')}")
# Complete the compatibility check before entering the transaction
# that truncates tables. This turns historical schema drift into an
# actionable error without risking the target's existing data.
restore_plans = await _validate_restore_schema(conn, manifest, schema)
# Use a transaction for atomic restore - either all tables are
# restored or none are, preventing partial/inconsistent state.
async with conn.transaction():
typer.echo(" Clearing existing data...")
# Truncate tables in reverse order (respects FK constraints)
for table in reversed(backup_tables):
for table in reversed(BACKUP_TABLES):
qualified_table = _fq_table(table, schema)
await conn.execute(f"TRUNCATE TABLE {qualified_table} CASCADE")
# Restore tables in forward order
for i, table in enumerate(backup_tables, 1):
for i, table in enumerate(BACKUP_TABLES, 1):
filename = f"{table}.bin"
if filename not in zf.namelist():
typer.echo(f" [{i}/{len(backup_tables)}] {table}: skipped (not in backup)")
typer.echo(f" [{i}/{len(BACKUP_TABLES)}] {table}: skipped (not in backup)")
continue
expected_rows = manifest["tables"].get(table, {}).get("rows", "?")
typer.echo(f" [{i}/{len(backup_tables)}] Restoring {table}... {expected_rows} rows")
typer.echo(f" [{i}/{len(BACKUP_TABLES)}] Restoring {table}... {expected_rows} rows")
plan = restore_plans[table]
# Strips the fields of any column the target no longer has;
# a no-op when the schemas still line up.
buffer = io.BytesIO(_strip_binary_copy_fields(zf.read(filename), plan))
data = zf.read(filename)
buffer = io.BytesIO(data)
# asyncpg requires schema_name as separate parameter
await conn.copy_to_table(
table,
schema_name=schema,
columns=plan.columns,
source=buffer,
format="binary",
)
await conn.copy_to_table(table, schema_name=schema, source=buffer, format="binary")
# Refresh materialized view
typer.echo(" Refreshing materialized views...")
@@ -396,22 +178,20 @@ async def _restore(
async def _run_backup(db_url: str, output: Path, schema: str = "public") -> dict[str, Any]:
"""Resolve database URL and run backup."""
_pg0 = parse_pg0_url(db_url)
is_pg0, instance_name = _pg0.is_pg0, _pg0.instance_name
is_pg0, instance_name, _ = parse_pg0_url(db_url)
if is_pg0:
typer.echo(f"Starting embedded PostgreSQL (instance: {instance_name})...")
resolved_url = await resolve_database_url(db_url)
return await _backup(resolved_url, output, schema, backup_tables=_effective_backup_tables())
return await _backup(resolved_url, output, schema)
async def _run_restore(db_url: str, input_file: Path, schema: str = "public") -> dict[str, Any]:
"""Resolve database URL and run restore."""
_pg0 = parse_pg0_url(db_url)
is_pg0, instance_name = _pg0.is_pg0, _pg0.instance_name
is_pg0, instance_name, _ = parse_pg0_url(db_url)
if is_pg0:
typer.echo(f"Starting embedded PostgreSQL (instance: {instance_name})...")
resolved_url = await resolve_database_url(db_url)
return await _restore(resolved_url, input_file, schema, backup_tables=_effective_backup_tables())
return await _restore(resolved_url, input_file, schema)
@app.command()
@@ -435,7 +215,7 @@ def backup(
manifest = asyncio.run(_run_backup(config.database_url, output, schema))
total_rows = sum(t["rows"] for t in manifest["tables"].values())
typer.echo(f"Backed up {total_rows} rows across {len(manifest['tables'])} tables")
typer.echo(f"Backed up {total_rows} rows across {len(BACKUP_TABLES)} tables")
typer.echo(f"Backup saved to {output}")
@@ -468,7 +248,7 @@ def restore(
manifest = asyncio.run(_run_restore(config.database_url, input_file, schema))
total_rows = sum(t["rows"] for t in manifest["tables"].values())
typer.echo(f"Restored {total_rows} rows across {len(manifest['tables'])} tables")
typer.echo(f"Restored {total_rows} rows across {len(BACKUP_TABLES)} tables")
typer.echo("Restore complete")
@@ -482,17 +262,17 @@ async def _run_migration(
"""Resolve database URL and run migrations for one schema or all discovered schemas."""
from ..migrations import run_migrations_for_schemas
_pg0 = parse_pg0_url(db_url)
is_pg0, instance_name = _pg0.is_pg0, _pg0.instance_name
is_pg0, instance_name, _ = parse_pg0_url(db_url)
if is_pg0:
typer.echo(f"Starting embedded PostgreSQL (instance: {instance_name})...")
resolved_url = await resolve_database_url(db_url)
config = HindsightConfig.from_env()
tenant_extension = load_extension("TENANT", TenantExtension)
if schema:
schemas = [schema]
else:
tenant_extension = load_extension("TENANT", TenantExtension)
schemas = [base_schema or DEFAULT_DATABASE_SCHEMA]
if tenant_extension:
tenants = await tenant_extension.list_tenants()
@@ -517,36 +297,9 @@ async def _run_migration(
ensure_extensions=ensure_extensions,
)
# After core migrations, provision any extension-owned bank-scoped tables
# per schema so extension schema evolves on the same lifecycle as core
# schema (rather than via a lazy first-request path).
if tenant_extension is not None:
await _provision_extra_bank_tables(resolved_url, schemas, tenant_extension)
return schemas
async def _provision_extra_bank_tables(
resolved_url: str, schemas: list[str], tenant_extension: TenantExtension
) -> None:
"""Run the tenant extension's table provisioner for each migrated schema.
Fires after core migrations complete so extension-owned bank tables are
created/evolved on the same lifecycle as core schema. A failure aborts the
migration command (and names the offending schema) rather than being
swallowed — provisioning is idempotent, so the operator can fix and re-run.
"""
for schema in schemas:
conn = await asyncpg.connect(resolved_url)
try:
await tenant_extension.provision_bank_tables(conn, schema)
except Exception as e:
typer.echo(f" Failed to provision extension tables for schema '{schema}': {e}", err=True)
raise
finally:
await conn.close()
@app.command(name="run-db-migration")
def run_db_migration(
schema: str | None = typer.Option(
@@ -601,134 +354,6 @@ def run_db_migration(
typer.echo(f"Database migrations completed successfully for {len(schemas)} schema(s)")
async def _resolve_schemas(base_schema: str | None) -> list[str]:
"""Base schema plus every discovered tenant schema, de-duplicated in order."""
schemas = [base_schema or DEFAULT_DATABASE_SCHEMA]
tenant_extension = load_extension("TENANT", TenantExtension)
if tenant_extension:
tenants = await tenant_extension.list_tenants()
schemas.extend(tenant.schema for tenant in tenants if tenant.schema)
return list(dict.fromkeys(schemas))
async def _run_repair_bank(
db_url: str,
*,
base_schema: str,
schema: str | None,
bank_id: str | None,
dry_run: bool,
) -> list[SchemaVectorIndexResult]:
"""Reconcile per-(bank, fact_type) vector index coverage over a raw connection.
A single autocommit connection is used because ``CREATE INDEX CONCURRENTLY``
(used by ``repair_vector_indexes``) cannot run inside a transaction block.
"""
schemas = [schema] if schema else await _resolve_schemas(base_schema)
index_clause = _vector_index_clause()
# Guarded by the command, but assert so this helper is never called for a
# backend without per-bank indexes.
assert index_clause is not None
conn = await _admin_connect(db_url)
try:
results = await repair_vector_indexes(conn, schemas, index_clause, dry_run=dry_run, bank_id=bank_id)
for result in results:
typer.echo(
f" schema '{result.schema}': {result.banks_scanned} bank(s) scanned, "
f"{result.already_present} present, {result.created} created, "
f"{result.skipped} to-create (dry-run), {result.failed} failed"
)
return results
finally:
await conn.close()
@app.command(name="repair-bank")
def repair_bank(
bank_id: str | None = typer.Option(
None,
"--bank",
"-b",
help="Bank id to repair. Mutually exclusive with --all.",
),
all_banks: bool = typer.Option(
False,
"--all",
help="Repair every bank in the base schema and all discovered tenant schemas.",
),
schema: str | None = typer.Option(
None,
"--schema",
"-s",
help="Limit to a single schema. Defaults to the base schema plus discovered tenant schemas.",
),
dry_run: bool = typer.Option(
False,
"--dry-run",
help="Report what would be repaired without creating or dropping any index.",
),
):
"""Verify and repair a bank's per-(bank, fact_type) vector index coverage.
Per-bank partial vector indexes are created when a bank is first created
(instant on an empty bank). Banks that arrive populated — via logical
restore, a cross-version upgrade, or a vector-extension switch — never hit
that path, so their recall silently falls back to a global index +
post-filter (slower, under-returning). This command detects missing OR
invalid coverage (an INVALID leftover or an index whose access method
drifted counts as missing) and rebuilds it with CREATE INDEX CONCURRENTLY,
so it never blocks the live fleet. Idempotent and safe to re-run — the
escape hatch after a restore, upgrade, or backend switch.
"""
if bool(bank_id) == all_banks:
typer.echo("Error: pass exactly one of --bank <id> or --all.", err=True)
raise typer.Exit(2)
config = HindsightConfig.from_env()
if not config.database_url:
typer.echo("Error: Database URL not configured.", err=True)
typer.echo("Set HINDSIGHT_API_DATABASE_URL environment variable.", err=True)
raise typer.Exit(1)
# Backend guard: backends with a single global vector index (AlloyDB ScaNN,
# Oracle) have no per-bank indexes to repair.
if _vector_index_clause() is None:
typer.echo("Configured vector backend does not use per-bank vector indexes — nothing to repair.")
return
target = f"bank '{bank_id}'" if bank_id else "all banks"
scope = f"schema '{schema}'" if schema else "base schema and all discovered tenant schemas"
typer.echo(f"Repairing per-bank vector indexes for {target} across {scope}...")
if dry_run:
typer.echo("Dry run: no indexes will be created or dropped.")
results = asyncio.run(
_run_repair_bank(
config.database_url,
base_schema=config.database_schema,
schema=schema,
bank_id=bank_id,
dry_run=dry_run,
)
)
total_banks = sum(r.banks_scanned for r in results)
total_present = sum(r.already_present for r in results)
total_created = sum(r.created for r in results)
total_skipped = sum(r.skipped for r in results)
total_failed = sum(r.failed for r in results)
typer.echo(
f"Done: {len(results)} schema(s), {total_banks} bank(s) scanned, "
f"{total_present} already present, {total_created} created, "
f"{total_skipped} to-create (dry-run), {total_failed} failed"
)
if total_failed:
failed_names = [name for r in results for name in r.failed_indexes]
typer.echo(f"Failed indexes (dropped, retry with a re-run): {', '.join(failed_names)}", err=True)
raise typer.Exit(1)
async def _run_export_bank(db_url: str, bank_id: str, output: Path, schema: str, include_history: bool) -> int:
"""Export a whole bank to a ZIP archive."""
conn = await _admin_connect(db_url)
@@ -736,14 +361,7 @@ async def _run_export_bank(db_url: str, bank_id: str, output: Path, schema: str,
# export_bank resolves table names via fq_table (the _current_schema
# contextvar); set it so the raw connection targets the right schema.
_current_schema.set(schema)
# _admin_connect registers JSON codecs, so row dumps already contain
# decoded Python values (including JSON scalar strings).
data = await export_bank(
conn,
bank_id,
include_history=include_history,
bank_rows_json_encoding="decoded",
)
data = await export_bank(conn, bank_id, include_history=include_history)
finally:
await conn.close()
@@ -848,16 +466,14 @@ def import_bank_command(
f"Imported bank '{result.bank_id}': {result.documents_imported} doc(s), "
f"{result.facts_imported} fact(s), {result.observations_imported} observation(s), "
f"{result.mental_models_imported} mental model(s), "
f"{result.mental_model_history_imported} mm-history row(s), "
f"{result.knowledge_pages_imported} knowledge page(s), {result.directives_imported} directive(s), "
f"{result.mental_model_history_imported} mm-history row(s), {result.directives_imported} directive(s), "
f"{result.webhooks_imported} webhook(s), {result.history_rows_imported} history row(s)"
)
async def _decommission_worker(db_url: str, worker_id: str, schema: str = "public") -> int:
"""Release all tasks owned by a worker, setting them back to pending status."""
_pg0 = parse_pg0_url(db_url)
is_pg0, instance_name = _pg0.is_pg0, _pg0.instance_name
is_pg0, instance_name, _ = parse_pg0_url(db_url)
if is_pg0:
typer.echo(f"Starting embedded PostgreSQL (instance: {instance_name})...")
resolved_url = await resolve_database_url(db_url)
@@ -916,8 +532,7 @@ def decommission_worker(
async def _decommission_all_workers(db_url: str, schema: str = "public") -> list[dict[str, Any]]:
"""Release all processing tasks from all workers, setting them back to pending status."""
_pg0 = parse_pg0_url(db_url)
is_pg0, instance_name = _pg0.is_pg0, _pg0.instance_name
is_pg0, instance_name, _ = parse_pg0_url(db_url)
if is_pg0:
typer.echo(f"Starting embedded PostgreSQL (instance: {instance_name})...")
resolved_url = await resolve_database_url(db_url)
@@ -982,8 +597,7 @@ def decommission_workers(
async def _worker_status(db_url: str, schema: str = "public") -> list[dict[str, Any]]:
"""Get all processing tasks grouped by worker with their last update time."""
_pg0 = parse_pg0_url(db_url)
is_pg0, instance_name = _pg0.is_pg0, _pg0.instance_name
is_pg0, instance_name, _ = parse_pg0_url(db_url)
if is_pg0:
typer.echo(f"Starting embedded PostgreSQL (instance: {instance_name})...")
resolved_url = await resolve_database_url(db_url)
@@ -1051,7 +665,6 @@ def worker_status(
def main():
load_dotenv_for_entrypoint()
app()
@@ -96,8 +96,7 @@ def get_database_url() -> str:
# for the sync engine used during migrations.
database_url = to_libpq_url(database_url)
# Alembic stores options through ConfigParser, where '%' is interpolation.
config.set_main_option("sqlalchemy.url", database_url.replace("%", "%%"))
config.set_main_option("sqlalchemy.url", database_url)
return database_url
@@ -1,67 +0,0 @@
"""Drop observation_history's FK to memory_units.
The history table records one snapshot per observation change, keyed by
``(bank_id, observation_id)``. Its foreign key to ``memory_units`` existed only to
cascade-delete history when the observation row went away.
That assumes every observation *is* a ``memory_units`` row, which is true only
while Postgres is the memories store. When another store owns the memories the
observation lives there and Postgres holds no row for it, so every history insert
raises a foreign-key violation — swallowed by the writer as "a race with parallel
consolidation" and logged at warning level. The audit trail goes silently empty.
Dropping the constraint lets history be recorded wherever the observation is
stored. The cleanup the cascade used to do is now explicit, in the paths that
delete observations (``_execute_delete_action``, ``clear_observations``,
``delete_bank``). Rows orphaned by a path that misses — a document delete
cascading through ``memory_units``, for instance — are invisible to readers,
which always filter by ``(bank_id, observation_id)``, and are reclaimed when the
bank is deleted.
Oracle builds this schema through its own DDL runner and never had the
constraint, so the Oracle slot is a deliberate no-op.
Revision ID: a1c9e7f3b2d8
Revises: c7d1e9a4b3f2
"""
from collections.abc import Sequence
from alembic import op
from hindsight_api.alembic._dialect import run_for_dialect
revision: str = "a1c9e7f3b2d8"
down_revision: str | Sequence[str] | None = "c7d1e9a4b3f2"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
_CONSTRAINT = "observation_history_observation_id_fkey"
def _pg_upgrade() -> None:
op.execute(f"ALTER TABLE observation_history DROP CONSTRAINT IF EXISTS {_CONSTRAINT}")
def _pg_downgrade() -> None:
# Re-adding the FK requires every row to reference a live memory_unit, so
# clear any history whose observation is not a Postgres row first — those are
# exactly the rows this migration made possible.
op.execute(
"DELETE FROM observation_history h "
"WHERE NOT EXISTS (SELECT 1 FROM memory_units m WHERE m.id = h.observation_id)"
)
op.execute(
f"ALTER TABLE observation_history ADD CONSTRAINT {_CONSTRAINT} "
"FOREIGN KEY (observation_id) REFERENCES memory_units(id) ON DELETE CASCADE"
)
def upgrade() -> None:
# Oracle never had the constraint (its schema is built by a separate DDL
# runner), so only Postgres has anything to drop.
run_for_dialect(pg=_pg_upgrade, oracle=None)
def downgrade() -> None:
run_for_dialect(pg=_pg_downgrade, oracle=None)
@@ -0,0 +1,52 @@
"""Add managed flag to knowledge_pages.
The knowledge base is managed by clients (CRUD over folders/pages). ``managed``
lets a client tag a node as system-owned vs. hand-authored; it carries no
server-side behaviour.
Revision ID: a5b6c7d8e9f0
Revises: a9b8c7d6e5f4
Create Date: 2026-06-26
"""
from collections.abc import Sequence
from alembic import context, op
from hindsight_api.alembic._dialect import run_for_dialect
revision: str = "a5b6c7d8e9f0"
down_revision: str | Sequence[str] | None = "a9b8c7d6e5f4"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
def _pg_schema_prefix() -> str:
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}knowledge_pages ADD COLUMN IF NOT EXISTS managed BOOLEAN NOT NULL DEFAULT false")
def _pg_downgrade() -> None:
schema = _pg_schema_prefix()
op.execute(f"ALTER TABLE {schema}knowledge_pages DROP COLUMN IF EXISTS managed")
def _oracle_upgrade() -> None:
op.execute("ALTER TABLE knowledge_pages ADD (managed NUMBER(1) DEFAULT 0 NOT NULL)")
def _oracle_downgrade() -> None:
op.execute("ALTER TABLE knowledge_pages DROP COLUMN managed")
def upgrade() -> None:
run_for_dialect(pg=_pg_upgrade, oracle=_oracle_upgrade)
def downgrade() -> None:
run_for_dialect(pg=_pg_downgrade, oracle=_oracle_downgrade)
@@ -1,82 +0,0 @@
"""Add indexes for terminal cleanup and newest-first operation listing.
Revision ID: a8c1e4f7b0d3
Revises: e7c3a9f1b2d5
Create Date: 2026-07-14
"""
from collections.abc import Sequence
from alembic import context, op
from hindsight_api.alembic._dialect import run_for_dialect
revision: str = "a8c1e4f7b0d3"
down_revision: str | Sequence[str] | None = "e7c3a9f1b2d5"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
def _pg_schema_prefix() -> str:
"""Schema-qualifier for PostgreSQL multi-tenant migration runs."""
schema = context.config.get_main_option("target_schema")
return f'"{schema}".' if schema else ""
def _pg_upgrade() -> None:
schema = _pg_schema_prefix()
# These can be large tables in long-running installations. Concurrent DDL
# keeps operation submission, polling, and status reads available.
with op.get_context().autocommit_block():
op.execute(
"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_async_operations_terminal_cleanup "
f"ON {schema}async_operations (updated_at, operation_id) "
"WHERE status IN ('completed', 'failed', 'cancelled')"
)
op.execute(
"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_async_operations_bank_created_desc "
f"ON {schema}async_operations (bank_id, created_at DESC)"
)
def _pg_downgrade() -> None:
schema = _pg_schema_prefix()
with op.get_context().autocommit_block():
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_async_operations_bank_created_desc")
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_async_operations_terminal_cleanup")
def _oracle_create_index(sql: str) -> None:
"""Create an index idempotently for rerun-safe Oracle migrations."""
block = (
"BEGIN "
"EXECUTE IMMEDIATE :stmt; "
"EXCEPTION WHEN OTHERS THEN "
"IF SQLCODE = -955 THEN NULL; ELSE RAISE; END IF; "
"END;"
)
op.get_bind().exec_driver_sql(block, {"stmt": sql})
def _oracle_upgrade() -> None:
# Oracle migrations run with CURRENT_SCHEMA set to each tenant, so table
# and index names intentionally remain unqualified here.
_oracle_create_index(
"CREATE INDEX idx_async_operations_terminal_cleanup ON async_operations (updated_at, operation_id, status)"
)
_oracle_create_index(
"CREATE INDEX idx_async_operations_bank_created_desc ON async_operations (bank_id, created_at DESC)"
)
def _oracle_downgrade() -> None:
op.execute("DROP INDEX idx_async_operations_bank_created_desc")
op.execute("DROP INDEX idx_async_operations_terminal_cleanup")
def upgrade() -> None:
run_for_dialect(pg=_pg_upgrade, oracle=_oracle_upgrade)
def downgrade() -> None:
run_for_dialect(pg=_pg_downgrade, oracle=_oracle_downgrade)
@@ -7,12 +7,8 @@ NULL). Hierarchy is a single self-referential ``parent_id`` so folders can nest
arbitrarily. Content stays in ``mental_models`` — this table is metadata + tree
structure only.
``managed`` lets a client tag a node as system-owned vs. hand-authored; it
carries no server-side behaviour. A partial unique index keeps page names unique
within a folder (case-insensitive; root pages compared under an empty parent).
Revision ID: a9b8c7d6e5f4
Revises: a1c9e7f3b2d8
Revises: b57a7c9e0d13
Create Date: 2026-06-25
"""
@@ -23,7 +19,7 @@ from alembic import context, op
from hindsight_api.alembic._dialect import run_for_dialect
revision: str = "a9b8c7d6e5f4"
down_revision: str | Sequence[str] | None = "a1c9e7f3b2d8"
down_revision: str | Sequence[str] | None = "b57a7c9e0d13"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
@@ -51,7 +47,6 @@ def _pg_upgrade() -> None:
name TEXT NOT NULL,
mental_model_id VARCHAR(64),
sort_order INTEGER NOT NULL DEFAULT 0,
managed BOOLEAN NOT NULL DEFAULT false,
created_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT now(),
updated_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT now(),
CONSTRAINT pk_knowledge_pages PRIMARY KEY (id),
@@ -68,25 +63,15 @@ def _pg_upgrade() -> None:
op.execute(
f"CREATE INDEX IF NOT EXISTS idx_kp_bank_parent ON {schema}knowledge_pages (bank_id, parent_id, sort_order)"
)
# COALESCE(parent_id, '') so root-level pages (NULL parent) are also unique by
# name — NULLs would otherwise compare distinct and allow duplicates.
op.execute(
"CREATE UNIQUE INDEX IF NOT EXISTS uq_kp_folder_pagename "
f"ON {schema}knowledge_pages (bank_id, COALESCE(parent_id, ''), lower(name)) "
"WHERE kind = 'page'"
)
def _pg_downgrade() -> None:
schema = _pg_schema_prefix()
op.execute(f"DROP INDEX IF EXISTS {schema}uq_kp_folder_pagename")
op.execute(f"DROP INDEX IF EXISTS {schema}idx_kp_bank_parent")
op.execute(f"DROP TABLE IF EXISTS {schema}knowledge_pages")
def _oracle_upgrade() -> None:
# No case-insensitive unique index on Oracle: `name` is a CLOB and cannot be
# indexed with lower(); page-name uniqueness is enforced on PG only.
op.execute(
"""
CREATE TABLE IF NOT EXISTS knowledge_pages (
@@ -97,7 +82,6 @@ def _oracle_upgrade() -> None:
name CLOB NOT NULL,
mental_model_id VARCHAR2(64),
sort_order NUMBER DEFAULT 0 NOT NULL,
managed NUMBER(1) DEFAULT 0 NOT NULL,
created_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,
updated_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,
CONSTRAINT pk_knowledge_pages PRIMARY KEY (id),
@@ -1,213 +0,0 @@
"""Add entities.entity_kind and exclude label entities from the trigram index.
Label entities (values of ``entity_labels`` config groups, stored as
``key:value`` canonical names) resolve by exact match only — fuzzy resolution
must never merge distinct label values (#1558), and since #3187 they are looked
up via the exact-match unique index rather than probed through pg_trgm. Their
rows were still covered by the shared trigram index, so every fuzzy probe for a
*regular* entity name pulled them into its candidate set only to discard them
in the bitmap recheck. On banks where a free-text label group accumulated tens
of thousands of mutually-similar values this recheck-discard overhead dominated
database CPU under ingest bursts (#3208).
"Is this row a label" was previously derived at runtime from the bank's
``entity_labels`` config, which an index predicate cannot reference — so the
classification is now materialised on the row:
1. Add ``entity_kind`` ("regular"/"label", CHECK-constrained) on both dialects.
A kind column rather than a boolean so future entity kinds don't need
another column.
2. Backfill per bank by classifying ``canonical_name`` against the bank's
``entity_labels`` config with the same ``is_label_entity()`` the resolver
uses at insert time — a SQL reimplementation would be a second source of
truth (and the map-group recursion doesn't translate). Banks hold at most
tens of thousands of entities, so the synchronous per-bank backfill is fine.
Label configs supplied only by a tenant extension (not stored in
``banks.config``) can't be seen here; their rows stay "regular", which
costs index size but never correctness — label *texts* still resolve via
the exact-match unique index.
3. Rebuild the PG trigram index as a partial index excluding label rows.
Built CONCURRENTLY (autocommit block, invalid-leftover sweep, IF NOT
EXISTS — same shape as 2071c7518f88) and only then drop the old full
index, so fuzzy probes never lose index coverage. Skipped entirely when
pg_trgm is absent (the resolver falls back to the "full" strategy, #626).
Oracle has no trigram index — it fuzzy-matches with a UTL_MATCH scan — so it
only gets the column + backfill; the resolver adds the matching
``entity_kind != 'label'`` filter to that scan.
Revision ID: b3e8d1c6f4a9
Revises: f2a6d8c4b1e9
Create Date: 2026-08-06
"""
import json
from collections.abc import Sequence
import sqlalchemy as sa
from alembic import context, op
from hindsight_api.alembic._dialect import run_for_dialect
revision: str = "b3e8d1c6f4a9"
down_revision: str | Sequence[str] | None = "f2a6d8c4b1e9"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
_OLD_INDEX = "entities_canonical_name_lower_trgm_idx"
_NEW_INDEX = "entities_canonical_name_lower_trgm_nonlabel_idx"
def _pg_schema_prefix() -> str:
"""Schema-qualifier for raw SQL on PG (multi-tenant search_path)."""
schema = context.config.get_main_option("target_schema")
return f'"{schema}".' if schema else ""
def _backfill_entity_kind(schema: str) -> None:
"""Set entity_kind='label' on rows matching their bank's entity_labels config.
Runs the resolver's own classification (``is_label_entity``) per bank in
Python rather than reimplementing the enum/text/map prefix rules in SQL.
Shared by both dialects: plain SELECT/UPDATE with expanding IN binds.
"""
from hindsight_api.engine.retain.entity_labels import (
build_labels_lookup,
is_label_entity,
parse_entity_labels,
)
bind = op.get_bind()
banks = bind.execute(sa.text(f"SELECT bank_id, config FROM {schema}banks")).fetchall()
for bank_id, raw_config in banks:
# PG JSONB arrives as a dict; Oracle CLOB arrives as a LOB object on
# raw text() fetches (oracledb's fetch_lobs default) — read it into a
# JSON string first.
if raw_config is not None and not isinstance(raw_config, (str, dict)):
raw_config = raw_config.read()
config = json.loads(raw_config) if isinstance(raw_config, str) else (raw_config or {})
labels_cfg = parse_entity_labels(config.get("entity_labels"))
if labels_cfg is None:
continue
lookup = build_labels_lookup(labels_cfg)
rows = bind.execute(
sa.text(f"SELECT id, canonical_name FROM {schema}entities WHERE bank_id = :bank_id"),
{"bank_id": bank_id},
).fetchall()
label_ids = [entity_id for entity_id, name in rows if is_label_entity(name, labels_cfg, lookup)]
# Chunked to stay under Oracle's 1000-element IN limit; also keeps PG
# bind arrays bounded.
for start in range(0, len(label_ids), 500):
chunk = label_ids[start : start + 500]
stmt = sa.text(f"UPDATE {schema}entities SET entity_kind = 'label' WHERE id IN :ids").bindparams(
sa.bindparam("ids", expanding=True)
)
bind.execute(stmt, {"ids": chunk})
def _pg_upgrade() -> None:
bind = op.get_bind()
schema = _pg_schema_prefix()
# `or None` collapses an unset option and an explicit empty string into NULL
# so the COALESCE below falls back to current_schema() in both cases.
target_schema = context.config.get_main_option("target_schema") or None
# IF NOT EXISTS: the transactional part below commits when the autocommit
# block is entered, so a failure during the CONCURRENTLY build leaves the
# revision unstamped with the column already added — the retry must not
# trip over it. The constant default is a metadata-only change on PG 11+.
op.execute(
f"ALTER TABLE {schema}entities ADD COLUMN IF NOT EXISTS entity_kind TEXT DEFAULT 'regular' NOT NULL "
f"CONSTRAINT chk_entities_entity_kind CHECK (entity_kind IN ('regular', 'label'))"
)
_backfill_entity_kind(schema)
# Without pg_trgm neither the old index nor the extension's opclass exists;
# the resolver already runs the "full" strategy there (#626).
has_trgm = bind.execute(sa.text("SELECT EXISTS(SELECT 1 FROM pg_extension WHERE extname = 'pg_trgm')")).scalar()
if not has_trgm:
return
# CREATE INDEX CONCURRENTLY cannot run inside a transaction block; the
# autocommit_block runs each statement outside Alembic's migration
# transaction. Build the partial index first and drop the old full index
# only afterwards, so fuzzy probes never lose index coverage.
with op.get_context().autocommit_block():
# A CONCURRENTLY build that errored on a previous run leaves an INVALID
# index of this name behind, which IF NOT EXISTS would skip forever.
leftover_invalid = bind.execute(
sa.text(
"SELECT NOT i.indisvalid "
"FROM pg_class c "
"JOIN pg_index i ON c.oid = i.indexrelid "
"JOIN pg_namespace n ON c.relnamespace = n.oid "
"WHERE c.relname = :index_name "
" AND n.nspname = COALESCE(:target_schema, current_schema())"
),
{"index_name": _NEW_INDEX, "target_schema": target_schema},
).scalar()
if leftover_invalid:
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}{_NEW_INDEX}")
# The predicate must textually match the resolver's candidate query
# (`entity_kind != 'label'`) for the planner to choose this index.
op.execute(
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS {_NEW_INDEX} "
f"ON {schema}entities USING GIN (LOWER(canonical_name) gin_trgm_ops) "
f"WHERE entity_kind != 'label'"
)
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}{_OLD_INDEX}")
def _pg_downgrade() -> None:
bind = op.get_bind()
schema = _pg_schema_prefix()
has_trgm = bind.execute(sa.text("SELECT EXISTS(SELECT 1 FROM pg_extension WHERE extname = 'pg_trgm')")).scalar()
if has_trgm:
# Restore the full index before dropping the partial one so fuzzy
# probes keep index coverage throughout.
with op.get_context().autocommit_block():
op.execute(
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS {_OLD_INDEX} "
f"ON {schema}entities USING GIN (LOWER(canonical_name) gin_trgm_ops)"
)
# Dropping the column also drops the partial index and CHECK constraint.
op.execute(f"ALTER TABLE {schema}entities DROP COLUMN IF EXISTS entity_kind")
def _oracle_upgrade() -> None:
# Swallow ORA-01430 (column already exists) so a retry after a mid-run
# failure is idempotent — Oracle DDL auto-commits statement by statement.
op.execute(
"""
BEGIN
EXECUTE IMMEDIATE 'ALTER TABLE entities ADD (entity_kind VARCHAR2(16) DEFAULT ''regular'' NOT NULL
CONSTRAINT chk_entities_entity_kind CHECK (entity_kind IN (''regular'', ''label'')))';
EXCEPTION WHEN OTHERS THEN
IF SQLCODE != -1430 THEN RAISE; END IF;
END;
"""
)
_backfill_entity_kind("")
def _oracle_downgrade() -> None:
# Swallow ORA-00904 (column does not exist).
op.execute(
"""
BEGIN
EXECUTE IMMEDIATE 'ALTER TABLE entities DROP COLUMN entity_kind';
EXCEPTION WHEN OTHERS THEN
IF SQLCODE != -904 THEN RAISE; END IF;
END;
"""
)
def upgrade() -> None:
run_for_dialect(pg=_pg_upgrade, oracle=_oracle_upgrade)
def downgrade() -> None:
run_for_dialect(pg=_pg_downgrade, oracle=_oracle_downgrade)
@@ -1,259 +0,0 @@
"""Install the maintenance discovery routines into the configured schema.
The three discovery routines driving the background maintenance loop —
``banks_needing_consolidation()``, ``schemas_with_expired_rows(...)`` and
``mental_models_with_cron()`` — were installed into ``public`` and gated on the
run being the base run (no ``target_schema``) or an explicit
``target_schema='public'`` run (``e5f6a7b8c9d0`` → ``b2d4f6a8c1e3`` →
``c7e9f1a3b5d2``, ``f4d1c2b3a5e6``).
That leaves a **single-tenant deployment migrated into a dedicated, non-**
``public`` **schema** (``HINDSIGHT_API_DATABASE_SCHEMA=<non-public>``) with no
routines at all: the runtime migrates only that one schema, so ``target_schema``
is never falsy or ``public``, the gate never opens, and the maintenance loop
logs, forever::
function public.banks_needing_consolidation() does not exist
function public.schemas_with_expired_rows(...) does not exist
The revision is stamped applied, so redeploying the same version does not help
(issue #2638; #2056 only fixed the ``public``/base-run case).
**The bug was the hardcoded literal, not the gating.** These routines are
database-global — each enumerates ``pg_class`` across every schema and dispatches
per schema — so exactly one copy should exist, and the maintenance loop calls the
one in ``get_config().database_schema`` (see ``fq_routine``). The old gate
installed into whichever schema was named ``public`` instead of whichever schema
the deployment is actually configured to use. Comparing ``target_schema`` against
the configured schema instead of the literal fixes #2638 at the source.
That also keeps the property the gate existed for: exactly one migration run
satisfies the predicate, so concurrent per-schema runs never issue competing
``CREATE OR REPLACE`` against the same ``pg_proc`` row and cannot hit
``tuple concurrently updated``. No cross-process coordination is required — in
particular no advisory lock, which is unusable here because Hindsight runs behind
connection poolers and managed PG services (see #2817).
Runs targeting any *other* schema drop the routines from that schema rather than
merely skipping. An earlier revision of this migration installed a copy into
every schema it touched, which left one dead duplicate per tenant on any database
that ran it; the drop makes the next migration pass clean those up instead of
leaving them behind forever.
PostgreSQL only: the maintenance loop and worker poller are PG-only, so the
Oracle slot is intentionally absent (mirrors ``e5f6a7b8c9d0``).
Revision ID: b6d2f8a4c1e7
Revises: a8c1e4f7b0d3
Create Date: 2026-07-20
"""
from collections.abc import Sequence
from alembic import context, op
from hindsight_api.alembic._dialect import run_for_dialect
from hindsight_api.config import get_config
revision: str = "b6d2f8a4c1e7"
down_revision: str | Sequence[str] | None = "a8c1e4f7b0d3"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
def _configured_schema() -> str:
"""The one schema this deployment's routines live in and are called from."""
return get_config().database_schema or "public"
def _target_schema() -> str | None:
return context.config.get_main_option("target_schema")
def _is_install_run() -> bool:
"""True for the single run that owns the routines.
The base run (no ``target_schema``) and the run targeting the configured
schema are the same deployment-level run; every other target is a tenant
schema that must not carry its own copy.
"""
target = _target_schema()
return not target or target == _configured_schema()
def _prefix(schema: str | None) -> str:
"""Qualifier for ``schema``, or ``""`` to fall back to ``search_path``."""
return f'"{schema}".' if schema else ""
def _pg_upgrade() -> None:
if not _is_install_run():
_drop_stray_copies()
return
schema = _prefix(_target_schema())
op.execute(
f"""
CREATE OR REPLACE FUNCTION {schema}banks_needing_consolidation()
RETURNS TABLE(schema_name text, bank_id text)
LANGUAGE plpgsql STABLE
AS $fn$
DECLARE
sch text;
BEGIN
FOR sch IN
SELECT n.nspname
FROM pg_class c
JOIN pg_namespace n ON n.oid = c.relnamespace
WHERE c.relname = 'memory_units' AND c.relkind = 'r'
LOOP
BEGIN
RETURN QUERY EXECUTE format($q$
SELECT %1$L::text, m.bank_id
FROM %1$I.memory_units m
JOIN %1$I.banks b ON b.bank_id = m.bank_id
WHERE m.consolidated_at IS NULL
AND m.consolidation_failed_at IS NULL
AND m.fact_type IN ('experience', 'world')
AND COALESCE(b.config -> 'enable_auto_consolidation', 'true'::jsonb) <> 'false'::jsonb
AND NOT EXISTS (
SELECT 1 FROM %1$I.async_operations o
WHERE o.bank_id = m.bank_id
AND o.operation_type = 'consolidation'
AND o.status IN ('pending', 'processing')
)
GROUP BY m.bank_id
$q$, sch);
EXCEPTION
-- Schema or its tables vanished between the pg_class
-- snapshot and this query (tenant dropped or migrating).
WHEN undefined_table OR invalid_schema_name OR undefined_column THEN
CONTINUE;
END;
END LOOP;
END;
$fn$;
"""
)
op.execute(
f"""
CREATE OR REPLACE FUNCTION {schema}schemas_with_expired_rows(
p_table text, p_ts_col text, p_days int
)
RETURNS SETOF text
LANGUAGE plpgsql STABLE
AS $fn$
DECLARE
sch text;
has_expired boolean;
BEGIN
IF p_days IS NULL OR p_days <= 0 THEN
RETURN;
END IF;
FOR sch IN
SELECT n.nspname
FROM pg_class c
JOIN pg_namespace n ON n.oid = c.relnamespace
WHERE c.relname = p_table AND c.relkind = 'r'
LOOP
BEGIN
EXECUTE format(
'SELECT EXISTS (SELECT 1 FROM %I.%I WHERE %I < NOW() - make_interval(days => $1))',
sch, p_table, p_ts_col
) INTO has_expired USING p_days;
EXCEPTION
-- Schema or its table vanished mid-scan; skip it.
WHEN undefined_table OR invalid_schema_name OR undefined_column THEN
CONTINUE;
END;
IF has_expired THEN
RETURN NEXT sch;
END IF;
END LOOP;
END;
$fn$;
"""
)
op.execute(
f"""
CREATE OR REPLACE FUNCTION {schema}mental_models_with_cron()
RETURNS TABLE(schema_name text, bank_id text, mental_model_id text,
refresh_cron text, last_refreshed_at timestamptz)
LANGUAGE plpgsql STABLE
AS $fn$
DECLARE
sch text;
BEGIN
FOR sch IN
SELECT n.nspname
FROM pg_class c
JOIN pg_namespace n ON n.oid = c.relnamespace
WHERE c.relname = 'mental_models' AND c.relkind = 'r'
LOOP
BEGIN
RETURN QUERY EXECUTE format($q$
SELECT %1$L::text, mm.bank_id::text, mm.id::text,
mm.trigger->>'refresh_cron', mm.last_refreshed_at
FROM %1$I.mental_models mm
WHERE COALESCE(mm.trigger->>'refresh_cron', '') <> ''
AND NOT EXISTS (
SELECT 1 FROM %1$I.async_operations o
WHERE o.bank_id = mm.bank_id
AND o.operation_type = 'refresh_mental_model'
AND o.status IN ('pending', 'processing')
AND o.task_payload->>'mental_model_id' = mm.id::text
)
$q$, sch);
EXCEPTION
-- Schema or its tables vanished between the pg_class
-- snapshot and this query (tenant dropped or migrating).
WHEN undefined_table OR invalid_schema_name OR undefined_column THEN
CONTINUE;
END;
END LOOP;
END;
$fn$;
"""
)
def _drop_routines(schema: str | None) -> None:
prefix = _prefix(schema)
op.execute(f"DROP FUNCTION IF EXISTS {prefix}mental_models_with_cron()")
op.execute(f"DROP FUNCTION IF EXISTS {prefix}schemas_with_expired_rows(text, text, int)")
op.execute(f"DROP FUNCTION IF EXISTS {prefix}banks_needing_consolidation()")
def _drop_stray_copies() -> None:
"""Remove per-tenant duplicates left by the first cut of this migration.
That version installed a copy into every schema it touched, so a database
that ran it carries one dead duplicate per tenant — only the copy in the
configured schema is ever called. Dropping here means the next migration pass
cleans them up; without it they would persist for the life of the database.
Safe on a database that never had them: ``DROP FUNCTION IF EXISTS`` is a
no-op, and this branch never runs for the configured schema.
"""
_drop_routines(_target_schema())
def _pg_downgrade() -> None:
# Only drop what this migration uniquely owns. When the configured schema is
# ``public`` the copies there belong to e5f6a7b8c9d0 / f4d1c2b3a5e6, which are
# still applied at this point and drop them on their own downgrade — removing
# them here would strand those migrations without the functions they claim to
# have installed.
if not _is_install_run() or _configured_schema() == "public":
return
_drop_routines(_target_schema())
def upgrade() -> None:
run_for_dialect(pg=_pg_upgrade)
def downgrade() -> None:
run_for_dialect(pg=_pg_downgrade)
@@ -0,0 +1,71 @@
"""Unique page name per folder in knowledge_pages.
The folder curator can fire concurrently (folder-create trigger + the
post-consolidation sweep), and an in-process lock can't serialize runs that
execute in different threads/loops. A partial unique index on
(bank_id, parent, lower(name)) for pages makes duplicate-named pages in the same
folder impossible at the DB level — the second concurrent insert fails and the
curator treats it as "already exists".
PostgreSQL only: the Oracle ``name`` column is a CLOB and cannot back a
functional unique index; Oracle relies on the in-process serialization instead.
Revision ID: c3d4e5f6a7b8
Revises: a5b6c7d8e9f0
Create Date: 2026-06-26
"""
from collections.abc import Sequence
from alembic import context, op
from hindsight_api.alembic._dialect import run_for_dialect
revision: str = "c3d4e5f6a7b8"
down_revision: str | Sequence[str] | None = "a5b6c7d8e9f0"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
def _pg_schema_prefix() -> str:
schema = context.config.get_main_option("target_schema")
return f'"{schema}".' if schema else ""
def _pg_upgrade() -> None:
schema = _pg_schema_prefix()
# First drop any pre-existing duplicate pages (created by the racy curator
# before this guard existed), keeping the earliest row of each duplicate set,
# so the unique index can be built. Their backing mental models are left in
# place (harmless orphans).
op.execute(
f"""
DELETE FROM {schema}knowledge_pages a
USING {schema}knowledge_pages b
WHERE a.kind = 'page' AND b.kind = 'page'
AND a.bank_id = b.bank_id
AND COALESCE(a.parent_id, '') = COALESCE(b.parent_id, '')
AND lower(a.name) = lower(b.name)
AND a.ctid > b.ctid
"""
)
# COALESCE(parent_id, '') so root-level pages (NULL parent) are also unique by
# name — NULLs would otherwise compare distinct and allow duplicates.
op.execute(
"CREATE UNIQUE INDEX IF NOT EXISTS uq_kp_folder_pagename "
f"ON {schema}knowledge_pages (bank_id, COALESCE(parent_id, ''), lower(name)) "
"WHERE kind = 'page'"
)
def _pg_downgrade() -> None:
schema = _pg_schema_prefix()
op.execute(f"DROP INDEX IF EXISTS {schema}uq_kp_folder_pagename")
def upgrade() -> None:
run_for_dialect(pg=_pg_upgrade) # oracle slot intentionally absent (CLOB name)
def downgrade() -> None:
run_for_dialect(pg=_pg_downgrade)
@@ -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)
@@ -1,90 +0,0 @@
"""Add ``causal_links`` to the curation archive (invalidated_memory_units).
Causal edges (``caused_by`` and the historical ``causes``/``enables``/
``prevents``) are retain-time extraction output: unlike temporal and semantic
links they cannot be recomputed from dates or embeddings, and graph maintenance
never rebuilds them. Invalidation MOVES a fact out of ``memory_units``, so the
``memory_links → memory_units`` FK cascade deletes every incident edge — and
revert had no way to bring the causal ones back (#2864).
This column parks the descriptors of the causal edges incident to an archived
fact — ``[{"from_unit_id", "to_unit_id", "link_type", "weight"}, ...]`` — so
revert can rematerialize them. It is deliberately unindexed and lives only on
the archive: live facts keep their causal edges in ``memory_links`` (curation
edits no longer delete them), and the archive is small, cold, and only read by
low-frequency curation operations.
Revision ID: c7d1e9a4b3f2
Revises: d7b2f8a1c934
Create Date: 2026-07-24
"""
from collections.abc import Sequence
from alembic import context, op
from hindsight_api.alembic._dialect import run_for_dialect
revision: str = "c7d1e9a4b3f2"
down_revision: str | Sequence[str] | None = "d7b2f8a1c934"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
def _pg_schema_prefix() -> str:
"""Schema-qualifier for raw SQL on PG (multi-tenant search_path)."""
schema = context.config.get_main_option("target_schema")
return f'"{schema}".' if schema else ""
def _pg_upgrade() -> None:
schema = _pg_schema_prefix()
# NOT NULL DEFAULT is metadata-only on PG 11+, so this is cheap even on a
# large archive. Existing rows read as "no causal edges captured" — edges
# lost before this migration cannot be reconstructed and are not guessed.
op.execute(
f"ALTER TABLE {schema}invalidated_memory_units "
f"ADD COLUMN IF NOT EXISTS causal_links JSONB NOT NULL DEFAULT '[]'::jsonb"
)
def _pg_downgrade() -> None:
schema = _pg_schema_prefix()
op.execute(f"ALTER TABLE {schema}invalidated_memory_units DROP COLUMN IF EXISTS causal_links")
def _oracle_upgrade() -> None:
# Kept in sync with PG for schema parity (curation itself is PostgreSQL-only
# today — it introspects pg_attribute to move rows between the two tables).
# Swallow ORA-01430 (column already exists) so the migration is idempotent.
op.execute(
"""
BEGIN
EXECUTE IMMEDIATE 'ALTER TABLE invalidated_memory_units ADD (causal_links CLOB DEFAULT ''[]''
CONSTRAINT imu_causal_links_json CHECK (causal_links IS JSON))';
EXCEPTION WHEN OTHERS THEN
IF SQLCODE != -1430 THEN RAISE; END IF;
END;
"""
)
def _oracle_downgrade() -> None:
# Swallow ORA-00904 (column does not exist).
op.execute(
"""
BEGIN
EXECUTE IMMEDIATE 'ALTER TABLE invalidated_memory_units DROP COLUMN causal_links';
EXCEPTION WHEN OTHERS THEN
IF SQLCODE != -904 THEN RAISE; END IF;
END;
"""
)
def upgrade() -> None:
run_for_dialect(pg=_pg_upgrade, oracle=_oracle_upgrade)
def downgrade() -> None:
run_for_dialect(pg=_pg_downgrade, oracle=_oracle_downgrade)
@@ -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)
@@ -1,150 +0,0 @@
"""Add the ``schemas_with_expired_operations`` cross-tenant discovery routine.
The worker's terminal-operation cleanup (``a8c1e4f7b0d3``) opens a connection
and a prune transaction against *every* tenant schema on every cleanup cycle,
whether or not that tenant has anything to prune. At thousands of tenants that
is a per-cycle query storm whose cost is paid entirely by idle schemas.
This is the same problem ``public.schemas_with_expired_rows`` already solves for
the ``audit_log`` / ``llm_requests`` retention sweeps (``e5f6a7b8c9d0``): one
round-trip returns just the schemas that actually hold expired rows, and the
caller then does real work only there. ``async_operations`` needs its own
routine rather than reusing that one because eligibility is not "row older than
N days" — pending and processing rows are never prunable, so the status filter
has to be part of the predicate.
Install policy mirrors ``b6d2f8a4c1e7`` (#2638/#2824), the current behaviour for
the sibling routines: the routine is database-global — it enumerates ``pg_class``
across every schema and dispatches per schema — so exactly one copy should exist,
installed into the schema this deployment is *configured* to use and called from
there via ``fq_routine``. Gating on the literal ``"public"`` instead of the
configured schema is what left single-tenant deployments in a dedicated
non-``public`` schema without the routine (#2638).
Exactly one migration run satisfies that predicate, so concurrent per-schema runs
never issue competing ``CREATE OR REPLACE`` against the same ``pg_proc`` row and
cannot hit ``tuple concurrently updated``. No cross-process coordination is
required — in particular no advisory lock, which is unusable here because
Hindsight runs behind connection poolers and managed PG services (see #2817).
Each per-schema probe runs in its own ``BEGIN ... EXCEPTION`` block so a tenant
dropped mid-scan is skipped instead of aborting the sweep (see ``c7e9f1a3b5d2``).
Revision ID: d7b2f8a1c934
Revises: b6d2f8a4c1e7
Create Date: 2026-07-20
"""
from collections.abc import Sequence
from alembic import context, op
from hindsight_api.alembic._dialect import run_for_dialect
from hindsight_api.config import get_config
revision: str = "d7b2f8a1c934"
down_revision: str | Sequence[str] | None = "b6d2f8a4c1e7"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
def _configured_schema() -> str:
"""The one schema this deployment's routines live in and are called from."""
return get_config().database_schema or "public"
def _target_schema() -> str | None:
return context.config.get_main_option("target_schema")
def _is_install_run() -> bool:
"""True for the single run that owns the routine (mirrors b6d2f8a4c1e7)."""
target = _target_schema()
return not target or target == _configured_schema()
def _prefix(schema: str | None) -> str:
"""Qualifier for ``schema``, or ``""`` to fall back to ``search_path``."""
return f'"{schema}".' if schema else ""
def _drop_routine(schema: str | None) -> None:
op.execute(f"DROP FUNCTION IF EXISTS {_prefix(schema)}schemas_with_expired_operations(int)")
def _pg_upgrade() -> None:
if not _is_install_run():
# Tenant schemas must not carry their own copy: the routine is
# database-global and only the configured schema's copy is ever called.
# Dropping (rather than skipping) also cleans up after any interim build
# of this branch that installed per-schema copies.
_drop_routine(_target_schema())
return
schema = _prefix(_target_schema())
op.execute(
f"""
CREATE OR REPLACE FUNCTION {schema}schemas_with_expired_operations(p_days int)
RETURNS SETOF text
LANGUAGE plpgsql STABLE
AS $fn$
DECLARE
sch text;
has_expired boolean;
BEGIN
-- Zero (or negative) retention means "keep forever": report nothing
-- so the caller skips the sweep entirely.
IF p_days IS NULL OR p_days <= 0 THEN
RETURN;
END IF;
FOR sch IN
SELECT n.nspname
FROM pg_class c
JOIN pg_namespace n ON n.oid = c.relnamespace
WHERE c.relname = 'async_operations' AND c.relkind = 'r'
LOOP
BEGIN
-- Matches the worker's prune predicate: only terminal rows
-- are eligible, so a schema holding nothing but pending or
-- processing work is correctly reported as having nothing
-- to prune. Uses idx_async_operations_terminal_cleanup.
EXECUTE format(
'SELECT EXISTS ('
' SELECT 1 FROM %I.async_operations'
' WHERE status IN (''completed'', ''failed'', ''cancelled'')'
' AND updated_at < NOW() - make_interval(days => $1)'
')',
sch
) INTO has_expired USING p_days;
EXCEPTION
-- Schema or its table vanished between the pg_class
-- snapshot and this probe (tenant dropped or migrating).
WHEN undefined_table OR invalid_schema_name OR undefined_column THEN
CONTINUE;
END;
IF has_expired THEN
RETURN NEXT sch;
END IF;
END LOOP;
END;
$fn$;
"""
)
def _pg_downgrade() -> None:
# This migration is the sole creator of this routine — no older migration
# owns a copy the way e5f6a7b8c9d0 owns the public sibling routines — so the
# install run's own copy is always ours to drop.
if not _is_install_run():
return
_drop_routine(_target_schema())
def upgrade() -> None:
# Oracle slot intentionally absent: this mirrors the PostgreSQL-only
# maintenance routines, and the Oracle worker keeps its per-schema sweep.
run_for_dialect(pg=_pg_upgrade)
def downgrade() -> None:
run_for_dialect(pg=_pg_downgrade)
@@ -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)
@@ -1,114 +0,0 @@
"""Drop the never-written `access_count` column from memory_units (and its archive).
``memory_units.access_count`` has been dead since the initial schema
(5a366d414dce): no code path anywhere in the repo ever writes it, and — despite
the ``access_count DESC`` index created alongside it — no query ever reads or
orders by it either. It is 0 on every row of every install. The lone remaining
mentions were an index, a stale comment naming an ``access_count_update`` task
type that was never implemented, and the column's name in the Oracle backend's
numeric-RETURNING list; all three go away with this change.
The column is dropped from the curation archive too. ``invalidated_memory_units``
was cloned ``LIKE memory_units`` (c9a1b2d3e4f5), so it inherited the column, and
curation's INSERT…SELECT round-trip builds its column list from the catalog
(``writes.py::_memory_unit_columns``) — the two tables must stay in lockstep or
the round-trip breaks on a column-count mismatch.
Dropping the column implicitly drops its index on both dialects
(``idx_memory_units_access_count`` on PG, ``idx_mu_access_count`` on Oracle), so
PostgreSQL also stops maintaining a btree that nothing ever probed.
Cost: on PostgreSQL ``DROP COLUMN`` is metadata-only (the attribute is marked
dropped, no table rewrite). On Oracle it does delete the column data row by row,
so on a large ``memory_units`` this migration is not free — it is still bounded
work on a single small integer column, and Oracle installs of that size can run
it during a maintenance window ahead of the upgrade if they prefer.
Revision ID: e4a7c1b9d2f6
Revises: a9b8c7d6e5f4
Create Date: 2026-08-03
"""
from collections.abc import Sequence
from alembic import context, op
from hindsight_api.alembic._dialect import run_for_dialect
revision: str = "e4a7c1b9d2f6"
down_revision: str | Sequence[str] | None = "a9b8c7d6e5f4"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
_TABLES = ("memory_units", "invalidated_memory_units")
def _pg_schema_prefix() -> str:
"""Schema-qualifier for raw SQL on PG (multi-tenant search_path)."""
schema = context.config.get_main_option("target_schema")
return f'"{schema}".' if schema else ""
def _pg_upgrade() -> None:
schema = _pg_schema_prefix()
for table in _TABLES:
# Drops idx_memory_units_access_count along with the column.
op.execute(f"ALTER TABLE {schema}{table} DROP COLUMN IF EXISTS access_count")
def _pg_downgrade() -> None:
schema = _pg_schema_prefix()
for table in _TABLES:
op.execute(f"ALTER TABLE {schema}{table} ADD COLUMN IF NOT EXISTS access_count integer NOT NULL DEFAULT 0")
# The archive was cloned without indexes; only the live table carried one.
op.execute(f"CREATE INDEX IF NOT EXISTS idx_memory_units_access_count ON {schema}memory_units (access_count DESC)")
def _oracle_upgrade() -> None:
# Oracle has no `DROP COLUMN IF EXISTS`; swallow ORA-00904 (column does not
# exist) so the migration is idempotent and safe on a schema that already
# lacks the column. Dropping the column also drops idx_mu_access_count.
for table in _TABLES:
op.execute(
f"""
BEGIN
EXECUTE IMMEDIATE 'ALTER TABLE {table} DROP COLUMN access_count';
EXCEPTION WHEN OTHERS THEN
IF SQLCODE != -904 THEN RAISE; END IF;
END;
"""
)
def _oracle_downgrade() -> None:
# Swallow ORA-01430 (column already exists) for idempotency. Matches the
# Oracle baseline's declaration: NUMBER(10) DEFAULT 0 NOT NULL.
for table in _TABLES:
op.execute(
f"""
BEGIN
EXECUTE IMMEDIATE
'ALTER TABLE {table} ADD (access_count NUMBER(10) DEFAULT 0 NOT NULL)';
EXCEPTION WHEN OTHERS THEN
IF SQLCODE != -1430 THEN RAISE; END IF;
END;
"""
)
# ORA-00955: index name already in use.
op.execute(
"""
BEGIN
EXECUTE IMMEDIATE 'CREATE INDEX idx_mu_access_count ON memory_units(access_count DESC)';
EXCEPTION WHEN OTHERS THEN
IF SQLCODE != -955 THEN RAISE; END IF;
END;
"""
)
def upgrade() -> None:
run_for_dialect(pg=_pg_upgrade, oracle=_oracle_upgrade)
def downgrade() -> None:
run_for_dialect(pg=_pg_downgrade, oracle=_oracle_downgrade)
@@ -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)
@@ -1,96 +0,0 @@
"""Drop the search_vector column from the curation archive (invalidated_memory_units).
The archive is cold storage, never a recall surface, and carries no text-search
index. Like ``embedding`` (dropped in d4f6a8c2e1b3), ``search_vector`` is a
recall-surface column whose type follows the configured text-search backend, so
it has no business living on the archive. Earlier curation code copied the live
row's ``search_vector`` into ``invalidated_memory_units`` on invalidate; the
engine now leaves it out on invalidate and recomputes it on revert, so the
column is dead weight.
Dropping it removes a latent failure mode (#2503): under a non-native backend
(pgroonga / pg_textsearch / pg_search / vchord) ``ensure_text_search_extension``
reconciles ``memory_units.search_vector`` to ``text`` / ``bm25vector`` but never
touched the archive, which the ``LIKE memory_units`` clone (c9a1b2d3e4f5) created
as ``tsvector``. The type mismatch then broke the curation INSERT … SELECT
round-trip:
column "search_vector" is of type tsvector but expression is of type text
With no column at all, there is nothing to mismatch. Unlike ``embedding`` (whose
creation sites already omit it), the ``LIKE`` clone still adds ``search_vector``,
so this migration does real work on both fresh and existing PostgreSQL databases.
DROP COLUMN is a metadata-only operation on both PostgreSQL and Oracle 23ai (no
table rewrite), so it is cheap even across many tenant schemas. The downgrade
re-adds an empty ``tsvector`` column (its original creation type).
Revision ID: e7c3a9f1b2d5
Revises: b57a7c9e0d13
Create Date: 2026-07-02
"""
from collections.abc import Sequence
from alembic import context, op
from hindsight_api.alembic._dialect import run_for_dialect
revision: str = "e7c3a9f1b2d5"
down_revision: str | Sequence[str] | None = "b57a7c9e0d13"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
def _pg_schema_prefix() -> str:
"""Schema-qualifier for raw SQL on PG (multi-tenant search_path)."""
schema = context.config.get_main_option("target_schema")
return f'"{schema}".' if schema else ""
def _pg_upgrade() -> None:
schema = _pg_schema_prefix()
op.execute(f"ALTER TABLE {schema}invalidated_memory_units DROP COLUMN IF EXISTS search_vector")
def _pg_downgrade() -> None:
schema = _pg_schema_prefix()
# Re-add as the original tsvector creation type; comes back empty regardless.
op.execute(f"ALTER TABLE {schema}invalidated_memory_units ADD COLUMN IF NOT EXISTS search_vector tsvector")
def _oracle_upgrade() -> None:
# Oracle has no `DROP COLUMN IF EXISTS`; swallow ORA-00904 (column does not
# exist) so the migration is idempotent and safe on a schema whose baseline
# may already omit the column.
op.execute(
"""
BEGIN
EXECUTE IMMEDIATE 'ALTER TABLE invalidated_memory_units DROP COLUMN search_vector';
EXCEPTION WHEN OTHERS THEN
IF SQLCODE != -904 THEN RAISE; END IF;
END;
"""
)
def _oracle_downgrade() -> None:
# Swallow ORA-01430 (column already exists) for idempotency. Oracle stores
# search_vector as CLOB (see the Oracle baseline), so re-add it as CLOB.
op.execute(
"""
BEGIN
EXECUTE IMMEDIATE 'ALTER TABLE invalidated_memory_units ADD (search_vector CLOB)';
EXCEPTION WHEN OTHERS THEN
IF SQLCODE != -1430 THEN RAISE; END IF;
END;
"""
)
def upgrade() -> None:
run_for_dialect(pg=_pg_upgrade, oracle=_oracle_upgrade)
def downgrade() -> None:
run_for_dialect(pg=_pg_downgrade, oracle=_oracle_downgrade)
@@ -1,85 +0,0 @@
"""Repair: drop the stale global memory_units vector index on per-bank backends.
Revision ID: f2a6d8c4b1e9
Revises: e4a7c1b9d2f6
Create Date: 2026-08-06
Migration d5e6f7a8b9c0 dropped the global ``idx_memory_units_embedding`` for
per-bank backends (every vector search is bank + fact_type scoped and served
by the ``idx_mu_emb_*`` partial indexes; the global index is never chosen by
the planner). However, older versions of the post-migration reconcile
(``ensure_vector_extension``) recreated the index when they found none, so
schemas that were provisioned or reconciled in that window carry it to this
day — paying a second vector graph insertion on every ``memory_units`` write
for an index no query uses.
This repair drops the leftover index. It is intentionally a migration, not
runtime reconcile behavior: ``DROP INDEX`` takes an ACCESS EXCLUSIVE lock on
``memory_units``, which belongs in the versioned, once-per-schema migration
path — not in code that runs at unpredictable times during startup or tenant
provisioning. The reconcile now leaves memory_units vector-index DDL to
migrations entirely on per-bank backends.
ScaNN deployments keep the global index by design (filtered vector search over
a global index; per-bank partial indexes cannot be built safely there), so the
migration is a no-op for them.
"""
import os
from collections.abc import Sequence
from alembic import context, op
from hindsight_api.alembic._dialect import run_for_dialect
revision: str = "f2a6d8c4b1e9"
down_revision: str | Sequence[str] | None = "e4a7c1b9d2f6"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
def _pg_schema_prefix() -> str:
"""Schema-qualifier for raw SQL on PG (multi-tenant search_path)."""
schema = context.config.get_main_option("target_schema")
return f'"{schema}".' if schema else ""
def _configured_vector_extension() -> str:
ext = os.getenv("HINDSIGHT_API_VECTOR_EXTENSION", "pgvector").lower()
if ext not in {"pgvector", "pgvectorscale", "vchord", "scann"}:
raise ValueError(
f"Invalid HINDSIGHT_API_VECTOR_EXTENSION: {ext}. Must be 'pgvector', 'vchord', 'pgvectorscale', or 'scann'"
)
return ext
def _pg_upgrade() -> None:
# ScaNN uses a global vector index by design — nothing stale to repair.
if _configured_vector_extension() == "scann":
return
schema = _pg_schema_prefix()
# DROP INDEX needs ACCESS EXCLUSIVE on memory_units. While it waits for
# in-flight transactions, every new query on the table queues behind it,
# so on a write-busy schema an unbounded wait can pile up traffic. Fail
# fast instead: the migration errors, the schema stays below head, and
# the next migration pass retries — preferable to freezing the table.
# SET LOCAL scopes the timeout to this migration's transaction.
op.execute("SET LOCAL lock_timeout = '10s'")
op.execute(f"DROP INDEX IF EXISTS {schema}idx_memory_units_embedding")
def _pg_downgrade() -> None:
# Intentional no-op: recreating a potentially multi-GB vector index that no
# query uses is not a safe downgrade action. Downgrading past d5e6f7a8b9c0
# restores the global index for deployments that genuinely need it.
pass
def upgrade() -> None:
# PG-only repair: the stale index is a PostgreSQL artifact of the old
# reconcile; Oracle deployments never had a reconcile that created it.
run_for_dialect(pg=_pg_upgrade)
def downgrade() -> None:
run_for_dialect(pg=_pg_downgrade)
@@ -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)
File diff suppressed because it is too large Load Diff
+1 -35
View File
@@ -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):
@@ -174,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,
@@ -394,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)
@@ -438,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
@@ -462,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
@@ -514,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
@@ -561,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,40 +1,76 @@
"""Markdown rendering for knowledge pages.
"""Open Knowledge Format (OKF) projection for knowledge pages.
Knowledge pages render as *read-only* markdown documents over the existing mental
models: each mental model becomes a markdown body with a YAML frontmatter block
(``type`` required; ``title``/``description``/``tags``/``timestamp`` optional).
Knowledge pages are a *read-only* OKF view over the existing mental models: each
mental model is projected into an OKF document a markdown body with YAML
frontmatter (``type`` required; ``title``/``description``/``tags``/``timestamp``
optional) and pages are linked into a constellation graph via shared tags.
See the Open Knowledge Format spec:
https://github.com/GoogleCloudPlatform/knowledge-catalog/tree/main/okf
This module is intentionally pure: every function transforms the mental-model
dicts returned by ``MemoryEngine.list_mental_models`` / ``get_mental_model`` and
never touches the database. That keeps rendering unit-testable without a DB or
LLM and lets the HTTP layer stay a thin wrapper.
never touches the database. That keeps the OKF contract unit-testable without a
DB or LLM and lets the HTTP layer stay a thin wrapper.
"""
from __future__ import annotations
from dataclasses import dataclass
from collections.abc import Callable
from dataclasses import dataclass, field
from typing import Any
# Every page carries exactly one ``type`` frontmatter field. We default to this
# when a page does not declare one via a ``type:<x>`` tag.
# OKF requires exactly one frontmatter field — ``type``. We default to this when
# a page does not declare one via a ``type:<x>`` tag.
DEFAULT_PAGE_TYPE = "knowledge-page"
# A page declares its ``type`` through a tag of the form ``type:runbook``.
# This keeps rendering schema-free (no new mental_models column): the type is
# lifted from the existing tags array.
# A page declares its OKF ``type`` through a tag of the form ``type:runbook``.
# This keeps the projection schema-free (no new mental_models column): the type
# is lifted from the existing tags array.
TYPE_TAG_PREFIX = "type:"
INDEX_FILENAME = "index.md"
# Deterministic, colour-blind-friendly palette. Type → colour is stable across
# requests so the constellation keeps the same colours between reloads.
_PALETTE = (
"#0074d9", # blue
"#2ecc40", # green
"#b10dc9", # purple
"#ff851b", # orange
"#39cccc", # teal
"#f012be", # magenta
"#3d9970", # olive
"#ff4136", # red
)
_EDGE_COLOR = "#9aa5b1"
@dataclass(frozen=True)
class PageType:
"""A page's ``type`` and the tags that remain after the type tag is split off."""
"""A page's OKF ``type`` and the tags that remain after the type tag is split off."""
type: str
display_tags: list[str]
@dataclass(frozen=True)
class KnowledgeGraph:
"""Cytoscape-style node/edge graph of knowledge pages linked by shared tags."""
nodes: list[dict[str, Any]] = field(default_factory=list)
edges: list[dict[str, Any]] = field(default_factory=list)
def _color_for(key: str) -> str:
"""Stable colour for a string key (FNV-ish hash into the fixed palette)."""
h = 0
for ch in key:
h = (h * 31 + ord(ch)) & 0xFFFFFFFF
return _PALETTE[h % len(_PALETTE)]
def _scalar(value: Any) -> str:
"""Emit a YAML-safe double-quoted scalar.
@@ -47,11 +83,11 @@ def _scalar(value: Any) -> str:
def page_type(tags: list[str] | None) -> PageType:
"""Split a ``type`` out of the tag list.
"""Split an OKF ``type`` out of the tag list.
The first ``type:<x>`` tag wins; all ``type:`` tags are removed from the
returned ``display_tags`` so they don't leak into the page's displayed tags.
Falls back to :data:`DEFAULT_PAGE_TYPE`.
returned ``display_tags`` so they don't pollute the constellation's
shared-tag edges. Falls back to :data:`DEFAULT_PAGE_TYPE`.
"""
resolved = DEFAULT_PAGE_TYPE
display: list[str] = []
@@ -70,7 +106,7 @@ def _timestamp(mm: dict[str, Any]) -> str | None:
def frontmatter(mm: dict[str, Any]) -> dict[str, Any]:
"""Build the ordered frontmatter mapping for a mental model.
"""Build the ordered OKF frontmatter mapping for a mental model.
``None``/empty values are dropped by :func:`render_frontmatter`.
"""
@@ -103,23 +139,23 @@ def render_frontmatter(fm: dict[str, Any]) -> str:
def render_document(mm: dict[str, Any]) -> str:
"""Render a full markdown document: frontmatter block + markdown body."""
"""Render a full OKF document: frontmatter block + markdown body."""
body = (mm.get("content") or "").strip()
return f"{render_frontmatter(frontmatter(mm))}\n\n{body}\n" if body else f"{render_frontmatter(frontmatter(mm))}\n"
def page_filename(page_id: str) -> str:
"""Bundle filename for a page id."""
"""OKF bundle filename for a page id."""
return f"{page_id}.md"
def log_filename(page_id: str) -> str:
"""Reserved per-page history filename."""
"""OKF reserved per-page history filename."""
return f"{page_id}.log.md"
def render_index(nodes: list[dict[str, Any]]) -> str:
"""Render the reserved ``index.md`` — nested markdown navigation over the tree.
"""Render the reserved ``index.md`` — nested OKF navigation over the tree.
``nodes`` is the flat folder/page list (each with ``id``, ``kind``, ``name``,
``parent_id``); folders nest their children, pages link to their ``.md``.
@@ -169,3 +205,59 @@ def render_log(mm: dict[str, Any], history: list[dict[str, Any]]) -> str:
lines.append(previous if previous else "_(empty)_")
lines.append("")
return "\n".join(lines).rstrip() + "\n"
def knowledge_graph(
pages: list[dict[str, Any]],
cluster_for: "Callable[[dict[str, Any]], str] | None" = None,
) -> KnowledgeGraph:
"""Derive the constellation graph: pages as nodes, shared tags as edges.
Two pages are linked when they share at least one (non-``type:``) tag; the
edge weight is the number of shared tags. Each node's cluster (``type`` field
+ colour) comes from ``cluster_for(page)`` the knowledge base groups by
parent folder; the default groups by OKF ``type``.
"""
nodes: list[dict[str, Any]] = []
tag_sets: list[tuple[str, frozenset[str]]] = []
for mm in pages:
page_id = mm["id"]
pt = page_type(mm.get("tags"))
cluster = cluster_for(mm) if cluster_for else pt.type
tag_sets.append((page_id, frozenset(pt.display_tags)))
nodes.append(
{
"data": {
"id": page_id,
"label": mm.get("name") or page_id,
"type": cluster,
"tagCount": len(pt.display_tags),
"color": _color_for(cluster),
}
}
)
edges: list[dict[str, Any]] = []
for i in range(len(tag_sets)):
source_id, source_tags = tag_sets[i]
if not source_tags:
continue
for j in range(i + 1, len(tag_sets)):
target_id, target_tags = tag_sets[j]
shared = source_tags & target_tags
if not shared:
continue
edges.append(
{
"data": {
"id": f"{source_id}--{target_id}",
"source": source_id,
"target": target_id,
"sharedTags": sorted(shared),
"weight": len(shared),
"color": _EDGE_COLOR,
}
}
)
return KnowledgeGraph(nodes=nodes, edges=edges)
@@ -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
@@ -66,6 +66,11 @@ def color_end(text: str) -> str:
return color(text, 1.0)
def color_mid(text: str) -> str:
"""Color text with gradient middle color."""
return color(text, 0.5)
def dim(text: str) -> str:
"""Dim/gray text."""
return f"\033[38;2;128;128;128m{text}\033[0m"
File diff suppressed because it is too large Load Diff
@@ -11,10 +11,8 @@ multiple API servers.
import asyncio
import json
import logging
from dataclasses import asdict, fields, replace
from functools import lru_cache
from types import UnionType
from typing import TYPE_CHECKING, Any, Union, get_args, get_origin
from dataclasses import asdict, replace
from typing import TYPE_CHECKING, Any
from hindsight_api.config import (
RECALL_BUDGET_FUNCTIONS,
@@ -34,14 +32,6 @@ if TYPE_CHECKING:
logger = logging.getLogger(__name__)
class BankConfigPersistenceConflictError(ValueError):
"""Raised when a validated bank config update can no longer be persisted."""
def __init__(self, bank_id: str):
self.bank_id = bank_id
super().__init__(f"Cannot update config for bank '{bank_id}': the bank does not exist")
def _validate_retain_strategy_chunking(base_config: HindsightConfig, strategies: Any) -> None:
"""Validate retain strategy chunking with the same semantics as apply_strategy()."""
if not isinstance(strategies, dict):
@@ -138,13 +128,12 @@ class ConfigResolver:
# Return full config object (dataclass doesn't have __init__ that accepts kwargs, so we update the object)
# Create a new config instance by copying the global config and updating fields
resolved_config = HindsightConfig(**config_dict)
# Multi-LLM chains and the reranker failover chain are static credential fields
# (never tenant/bank-overridable), but asdict() above flattened their member
# dataclasses into plain dicts. Restore the original typed objects from the global
# config so the resolved object stays well-typed for any consumer that reads them.
# Multi-LLM chains are static credential fields (never tenant/bank-overridable),
# but asdict() above flattened their member dataclasses into plain dicts. Restore
# the original typed objects from the global config so the resolved object stays
# well-typed for any consumer that reads them.
resolved_config = replace(
resolved_config,
reranker_members=self._global_config.reranker_members,
llm_members=self._global_config.llm_members,
llm_strategy=self._global_config.llm_strategy,
retain_llm_members=self._global_config.retain_llm_members,
@@ -297,8 +286,7 @@ class ConfigResolver:
# Only return active overrides for configurable fields. JSON null is a tombstone
# for "Server Default" in the bank-config UI and should not override defaults.
active = {k: v for k, v in normalized.items() if k in self._configurable_fields and v is not None}
return _coerce_stored_bank_overrides(bank_id, active)
return {k: v for k, v in normalized.items() if k in self._configurable_fields and v is not None}
except Exception as e:
logger.error(f"Failed to load bank config for {bank_id}: {e}")
@@ -338,22 +326,16 @@ class ConfigResolver:
k: v for k, v in normalized.items() if k in self._configurable_fields and v is not None
}
if overrides:
result[row["bank_id"]] = _coerce_stored_bank_overrides(row["bank_id"], overrides)
result[row["bank_id"]] = overrides
except Exception as e:
logger.error(f"Failed to bulk-load bank configs: {e}")
return result
async def validate_bank_config_updates(
self,
bank_id: str,
updates: dict[str, Any],
context: RequestContext | None = None,
*,
projected_bank_overrides: dict[str, Any] | None = None,
check_permissions: bool = True,
) -> dict[str, Any]:
async def update_bank_config(
self, bank_id: str, updates: dict[str, Any], context: RequestContext | None = None
) -> None:
"""
Normalize and validate bank configuration overrides.
Update bank configuration overrides (with permission checking).
Args:
bank_id: Bank identifier
@@ -362,16 +344,9 @@ class ConfigResolver:
or Python field format (llm_provider).
Only configurable fields are allowed.
context: Request context for permission checking
projected_bank_overrides: Bank overrides to use as the validation
base instead of loading the current bank row.
check_permissions: Whether client field permissions apply to these
updates. Server-owned projected values set this to false.
Returns:
Normalized updates ready to persist.
Raises:
ValueError: If attempting to override invalid/disallowed fields.
ValueError: If attempting to override invalid/disallowed fields
"""
# Normalize keys
normalized_updates = normalize_config_dict(updates)
@@ -403,7 +378,7 @@ class ConfigResolver:
)
# PERMISSIONS: Check tenant/bank permissions
if check_permissions and self.tenant_extension and context:
if self.tenant_extension and context:
try:
allowed_fields = await self.tenant_extension.get_allowed_config_fields(context, bank_id)
if allowed_fields is not None: # None means "allow all"
@@ -413,7 +388,7 @@ class ConfigResolver:
f"Not allowed to modify fields: {sorted(disallowed)}. "
f"Your permissions allow: {sorted(list(allowed_fields)[:10])}..."
if allowed_fields
else f"Not allowed to modify fields: {sorted(disallowed)}. "
else "Not allowed to modify fields: {sorted(disallowed)}. "
"Your permissions do not allow any config modifications."
)
except ValueError:
@@ -422,11 +397,6 @@ class ConfigResolver:
logger.warning(f"Failed to check permissions for bank {bank_id}: {e}")
# Continue without permission check (fail open for backward compatibility)
# Validate every value against its declared field type before the
# field-specific checks below, so a wrong-shaped value is reported as such
# instead of tripping a structural validator with a confusing message.
_validate_config_value_types(normalized_updates)
# Validate entity_labels structure
if "entity_labels" in normalized_updates and normalized_updates["entity_labels"] is not None:
from .engine.retain.entity_labels import parse_entity_labels
@@ -443,16 +413,6 @@ class ConfigResolver:
raise ValueError(
"Strategy names must not be empty strings. Remove entries with empty names before saving."
)
# A strategy's overrides are applied with dataclasses.replace() at retain
# time, so a wrong-shaped value there wedges the bank exactly as a
# top-level one would. Same contract, same door.
for strategy_name, strategy_overrides in normalized_updates["retain_strategies"].items():
if not isinstance(strategy_overrides, dict):
raise ValueError(f"Invalid retain strategy {strategy_name!r}: must be an object")
try:
_validate_config_value_types(normalize_config_dict(strategy_overrides))
except ValueError as e:
raise ValueError(f"Invalid retain strategy {strategy_name!r}: {e}") from e
# Validate recall budget fields
_validate_recall_budget_updates(normalized_updates)
@@ -467,11 +427,7 @@ class ConfigResolver:
)
if chunking_fields_updated:
config_dict = await self._resolve_parent_config_dict(bank_id, context)
active_bank_overrides = (
await self._load_bank_config(bank_id)
if projected_bank_overrides is None
else dict(projected_bank_overrides)
)
active_bank_overrides = await self._load_bank_config(bank_id)
for key, value in normalized_updates.items():
if key not in self._configurable_fields:
continue
@@ -487,26 +443,17 @@ class ConfigResolver:
)
_validate_retain_strategy_chunking(base_config, base_config.retain_strategies)
return normalized_updates
# Persist the override. Banks are created lazily (on first retain), so a
# PATCH that precedes any ingestion would otherwise UPDATE zero rows and
# silently no-op while returning 200. Ensure the bank row exists first
# (this also creates its per-bank vector indexes), then merge defensively:
# COALESCE guards against a NULL config column (NULL || jsonb is NULL),
# which would drop the override even when a row is updated.
from .engine.retain.fact_storage import ensure_bank_exists
async def update_bank_config(
self, bank_id: str, updates: dict[str, Any], context: RequestContext | None = None
) -> None:
"""Validate and persist bank configuration overrides for an existing bank.
Bank creation belongs to ``MemoryEngine``; this raises ``ValueError`` if
the bank does not exist rather than silently discarding the overrides.
"""
normalized_updates = await self.validate_bank_config_updates(bank_id, updates, context)
await self._persist_bank_config(bank_id, normalized_updates)
async def _persist_bank_config(self, bank_id: str, normalized_updates: dict[str, Any]) -> None:
"""Persist already-validated overrides without changing bank lifecycle state."""
# Bank lifecycle belongs to MemoryEngine. Callers must create the row
# before reaching this persistence step. COALESCE guards against a NULL
# config column (NULL || jsonb is NULL), which would drop the override.
async with self._backend.acquire() as conn:
result = await conn.execute(
await ensure_bank_exists(conn, bank_id, ops=self._backend.ops)
await conn.execute(
f"""
UPDATE {fq_table("banks")}
SET config = COALESCE(config, '{{}}'::jsonb) || $1::jsonb,
@@ -517,14 +464,6 @@ class ConfigResolver:
bank_id,
)
# A missing bank row matches zero rows, which would otherwise persist
# nothing while reporting success. Fail loudly instead: reaching here
# without the row means a caller skipped the engine's provisioning step.
# (The Oracle wrapper reshapes rowcount into the same "UPDATE <n>" form.)
updated = int(result.split()[-1]) if isinstance(result, str) and result.startswith("UPDATE") else 0
if updated == 0:
raise BankConfigPersistenceConflictError(bank_id)
logger.info(f"Updated bank config for {bank_id}: {list(normalized_updates.keys())}")
async def reset_bank_config(self, bank_id: str) -> None:
@@ -548,147 +487,6 @@ class ConfigResolver:
logger.info(f"Reset bank config for {bank_id} to defaults")
# Fields whose accepted input shape is deliberately wider than the dataclass
# annotation, because a dedicated structural validator normalizes them later.
_WIDENED_FIELD_TYPES: dict[str, tuple[type, ...]] = {
# parse_entity_labels() accepts both the bare list of label groups and the
# {"attributes": [...]} envelope, though the field is annotated `list | None`.
"entity_labels": (list, dict),
}
def _runtime_types(declared: Any) -> tuple[type, ...]:
"""Runtime-checkable base classes for a dataclass field annotation.
Unwraps unions (``str | None``) and generic aliases (``list[str]`` -> ``list``);
``None`` is dropped because callers handle the tombstone separately. Returns an
empty tuple for anything not reducible to concrete classes, which the callers
read as "no type contract to enforce".
"""
if declared is type(None):
return ()
origin = get_origin(declared)
if origin in (Union, UnionType):
return tuple(t for arg in get_args(declared) for t in _runtime_types(arg))
if origin is not None:
return (origin,) if isinstance(origin, type) else ()
return (declared,) if isinstance(declared, type) else ()
@lru_cache(maxsize=1)
def _configurable_field_types() -> dict[str, tuple[type, ...]]:
"""Map each configurable field to the value types it accepts."""
configurable = HindsightConfig.get_configurable_fields()
field_types: dict[str, tuple[type, ...]] = {}
for field in fields(HindsightConfig):
if field.name not in configurable:
continue
allowed = _WIDENED_FIELD_TYPES.get(field.name) or _runtime_types(field.type)
if allowed:
field_types[field.name] = allowed
return field_types
def _value_matches_type(value: Any, allowed: tuple[type, ...]) -> bool:
"""Whether ``value`` satisfies a field's declared type contract."""
if isinstance(value, bool):
# bool is an int subclass; it must not slip into a numeric field.
return bool in allowed
if isinstance(value, int) and float in allowed:
# JSON draws no int/float distinction: 1 is a valid ratio.
return True
return isinstance(value, allowed)
# Field types are reported to API clients, so name them the way the JSON payload
# reads rather than by their Python class.
_TYPE_DESCRIPTIONS: dict[type, str] = {
bool: "a boolean",
int: "an integer",
float: "a number",
str: "a string",
list: "a list",
dict: "an object",
}
def _describe_types(allowed: tuple[type, ...]) -> str:
return " or ".join(dict.fromkeys(_TYPE_DESCRIPTIONS.get(t, t.__name__) for t in allowed))
def _validate_config_value_types(updates: dict[str, Any]) -> None:
"""Reject values whose type contradicts the declared HindsightConfig type.
Without this, the bank-config API happily stores e.g. a JSON object in
``observations_mission``; the write succeeds and the bank then fails every
consolidation with ``expected string or bytes-like object, got 'dict'`` from
deep inside prompt assembly (issue #3218). Reject at the door instead, naming
the field and the expected type.
"""
field_types = _configurable_field_types()
for key, value in updates.items():
allowed = field_types.get(key)
# None is the "clear this override" tombstone; unknown keys are rejected
# elsewhere as non-configurable.
if allowed is None or value is None:
continue
if not _value_matches_type(value, allowed):
raise ValueError(f"{key} must be {_describe_types(allowed)}, got {type(value).__name__}")
def _coerce_stored_bank_overrides(bank_id: str, overrides: dict[str, Any], where: str = "") -> dict[str, Any]:
"""Make stored bank overrides safe to consume, tolerating pre-validation shapes.
``_validate_config_value_types`` rejects bad types at write time, but banks
configured before that landed can still hold e.g. a JSON object in a
string-typed field. Every consumer that treats such a value as text blows up
identically on every run (``escape_for_prompt`` -> ``re.sub`` ->
"expected string or bytes-like object, got 'dict'"), so the bank's
consolidation never recovers on its own (issue #3218).
String fields are JSON-encoded, which preserves the author's intent — the
structure still reaches the prompt, as text. Anything else is dropped so the
bank falls back to the tenant/global value rather than wedging.
``where`` labels the location in warnings; it is set when recursing into a
retain strategy, whose overrides reach the same fields via ``apply_strategy``.
"""
field_types = _configurable_field_types()
coerced: dict[str, Any] = {}
for key, value in overrides.items():
allowed = field_types.get(key)
# None passes through: the caller has already dropped top-level tombstones,
# and inside a retain strategy a null is a deliberate override to None.
if allowed is None or value is None or _value_matches_type(value, allowed):
coerced[key] = value
continue
if str in allowed:
coerced[key] = json.dumps(value, ensure_ascii=False)
logger.warning(
f"Bank {bank_id} config field '{key}'{where} holds a {type(value).__name__} but is a string field; "
f"using its JSON encoding. Re-save this field as a string to silence this warning."
)
else:
logger.warning(
f"Bank {bank_id} config field '{key}'{where} holds a {type(value).__name__} but must be "
f"{_describe_types(allowed)}; ignoring the override and falling back to the server default."
)
# Strategy overrides are spliced onto the resolved config by apply_strategy(),
# so a bad value nested there wedges the bank just as a top-level one does.
strategies = coerced.get("retain_strategies")
if isinstance(strategies, dict):
coerced["retain_strategies"] = {
name: (
_coerce_stored_bank_overrides(bank_id, strategy, where=f" in retain strategy {name!r}")
if isinstance(strategy, dict)
else strategy
)
for name, strategy in strategies.items()
}
return coerced
_RECALL_BUDGET_FIXED_KEYS = (
"recall_budget_fixed_low",
"recall_budget_fixed_mid",
+30 -4
View File
@@ -14,7 +14,10 @@ import subprocess
import sys
import time
from pathlib import Path
from typing import IO
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from typing import IO
logger = logging.getLogger(__name__)
@@ -39,28 +42,37 @@ class IdleTimeoutMiddleware:
self.app = app
self.idle_timeout = idle_timeout
self.last_activity = time.time()
self._checker_task = None
async def __call__(self, scope, receive, send):
# Update activity timestamp on each request
self.last_activity = time.time()
await self.app(scope, receive, send)
def start_idle_checker(self):
"""Start the background task that checks for idle timeout."""
self._checker_task = asyncio.create_task(self._check_idle())
async def _check_idle(self):
"""Exit the daemon after the configured period without requests."""
"""Background task that exits the process after idle timeout."""
# If idle_timeout is 0, don't auto-exit
if self.idle_timeout <= 0:
return
while True:
await asyncio.sleep(30)
await asyncio.sleep(30) # Check every 30 seconds
idle_time = time.time() - self.last_activity
if idle_time > self.idle_timeout:
logger.info(f"Idle timeout reached ({self.idle_timeout}s), shutting down daemon")
# Give a moment for any in-flight requests
await asyncio.sleep(1)
# Send SIGTERM to ourselves to trigger graceful shutdown
import signal
os.kill(os.getpid(), signal.SIGTERM)
def _detach_popen_kwargs(log_handle: IO[bytes]) -> dict:
def _detach_popen_kwargs(log_handle: "IO[bytes]") -> dict:
"""Cross-platform kwargs to spawn a subprocess detached from the caller.
On POSIX, ``start_new_session=True`` calls ``setsid(2)`` so the child
@@ -157,3 +169,17 @@ def daemonize():
subprocess.Popen(cmd, env=env, **_detach_popen_kwargs(log_handle))
sys.exit(0)
def check_daemon_running(port: int = DEFAULT_DAEMON_PORT) -> bool:
"""Check if a daemon is running and responsive on the given port."""
import socket
try:
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
sock.settimeout(1)
result = sock.connect_ex(("127.0.0.1", port))
sock.close()
return result == 0
except Exception:
return False
@@ -3,7 +3,7 @@ Memory Engine - Core implementation of the memory system.
This package contains all the implementation details of the memory engine:
- MemoryEngine: Main class for memory operations
- Utility modules: embedding_utils, link_utils, bank_utils
- Utility modules: embedding_utils, link_utils, think_utils, bank_utils
- Supporting modules: embeddings, cross_encoder, entity_resolver, etc.
"""
@@ -10,7 +10,7 @@ import asyncio
import json
import logging
import uuid
from collections.abc import Awaitable, Callable
from collections.abc import Callable
from contextlib import asynccontextmanager
from dataclasses import dataclass, field
from datetime import datetime, timezone
@@ -19,8 +19,6 @@ from typing import Any
from pydantic import BaseModel, Field
from ..engine.db_utils import acquire_with_retry
from ..models import RequestContext
from .schema import fq_table_explicit
logger = logging.getLogger(__name__)
@@ -121,60 +119,23 @@ class AuditLogger:
schema_getter: Callable[[], str],
enabled: bool,
allowed_actions: list[str],
bank_enabled_resolver: Callable[[str, RequestContext | None], Awaitable[bool]] | None = None,
) -> None:
self._pool_getter = pool_getter
self._schema_getter = schema_getter
self._enabled = enabled
self._allowed_actions: frozenset[str] | None = frozenset(allowed_actions) if allowed_actions else None
# Resolves the hierarchical ``audit_log_enabled`` for one bank
# (env -> tenant -> bank). None means "no per-bank resolution wired",
# in which case the global value alone decides.
self._bank_enabled_resolver = bank_enabled_resolver
def action_allowed(self, action: str) -> bool:
"""Global action-allowlist check. Cheap, synchronous, bank-independent.
The allowlist is deployment-wide, so this is a valid pre-filter to skip
work for actions that can never be audited. It deliberately does NOT
consult the enabled flag: that is per-bank overridable, so a bank may
turn auditing ON even when the deployment default is off.
"""
if self._allowed_actions is None:
return True
return action in self._allowed_actions
async def should_log(self, action: str, bank_id: str | None, context: RequestContext | None = None) -> bool:
"""Full audit decision: action allowlist AND the bank's resolved switch.
``audit_log_enabled`` is hierarchical (env -> tenant -> bank), so the
effective value depends on which bank the action targets. Falls back to
the global value when there is no bank in scope or no resolver wired.
"""
if not self.action_allowed(action):
def is_enabled(self, action: str) -> bool:
"""Check if audit logging is enabled for this action."""
if not self._enabled:
return False
if bank_id is None or self._bank_enabled_resolver is None:
return self._enabled
try:
return await self._bank_enabled_resolver(bank_id, context)
except Exception as e:
# Never let a config-resolution failure break the request. Fall back
# to the deployment default: a transient DB blip must not silently
# create an audit gap for a bank meant to be audited. The tradeoff is
# the opt-out direction — a bank that overrode to false under a
# default-on deployment will be audited during the outage. We accept
# that: a few extra audit rows during a DB blip is the safer failure
# than dropping records that compliance may require.
logger.warning(f"Audit config resolution failed for bank={bank_id}: {e}; using global default")
return self._enabled
if self._allowed_actions is not None:
return action in self._allowed_actions
return True
def log_fire_and_forget(self, entry: AuditEntry) -> None:
"""Schedule an audit write as a background task.
Assumes the caller already made the audit decision via ``should_log``;
only the bank-independent allowlist is re-checked here.
"""
if not self.action_allowed(entry.action):
"""Schedule an audit write as a background task."""
if not self.is_enabled(entry.action):
return
try:
asyncio.create_task(self._safe_log(entry))
@@ -189,12 +150,8 @@ class AuditLogger:
logger.debug("Audit log skipped: pool not available")
return
try:
# fq_table_explicit qualifies per dialect: "schema".audit_log on
# PostgreSQL, bare audit_log on Oracle (where the schema is set at the
# session level). A raw f"{schema}.audit_log" produced public.audit_log
# on Oracle, where "public" is a reserved word — every write failed
# with ORA-00903 even though the table exists.
table = fq_table_explicit("audit_log", self._schema_getter())
schema = self._schema_getter()
table = f"{schema}.audit_log"
async with acquire_with_retry(pool, max_retries=1) as conn:
await conn.execute(
f"""
@@ -225,7 +182,6 @@ async def audit_context(
bank_id: str | None = None,
request: dict[str, Any] | None = None,
metadata: dict[str, Any] | None = None,
context: RequestContext | None = None,
):
"""Async context manager that times the operation and writes audit on exit.
@@ -234,7 +190,7 @@ async def audit_context(
result = await do_work()
entry.response = result_dict
"""
if audit_logger is None or not await audit_logger.should_log(action, bank_id, context):
if audit_logger is None or not audit_logger.is_enabled(action):
entry = AuditEntry(action=action, transport=transport, bank_id=bank_id)
yield entry
return
@@ -13,8 +13,6 @@ but operators should opt in with that in mind.
from typing import Any
RERANKER_BANK_ID_HEADER = "X-Hindsight-Bank-Id"
def apply_bank_attribution(request: dict[str, Any]) -> None:
"""Tag ``request`` with ``user=<bank_id>`` for per-bank cost attribution.
@@ -34,14 +32,3 @@ def apply_bank_attribution(request: dict[str, Any]) -> None:
bank_id = get_current_bank_id()
if bank_id:
request["user"] = bank_id
def reranker_bank_attribution_headers() -> dict[str, str]:
"""Return the fixed per-bank header for trusted remote reranker endpoints."""
from ..config import get_config
from .memory_engine import get_current_bank_id
if not get_config().reranker_send_bank_as_header:
return {}
bank_id = get_current_bank_id()
return {RERANKER_BANK_ID_HEADER: bank_id} if bank_id else {}
@@ -1,191 +0,0 @@
"""Server-side prompt-cache affinity hints for OpenAI-compatible providers.
Prompt caching only pays off when the same conversation reaches the same backend
cache, and providers expose different mechanisms for that:
- xAI stores prompt-cache entries **per backend server** and routes requests
carrying the same ``x-grok-conv-id`` to one server (docs.x.ai, "Maximizing
Cache Hits"). Without it, consecutive calls of one agentic loop can each land
on a cache-cold replica.
- OpenAI accepts a ``prompt_cache_key`` request field that improves its own
cache routing.
Hindsight already does provider-specific cache work for its first-class
providers (``anthropic_llm`` sets ``cache_control`` breakpoints; ``gemini_llm``
runs an explicit ``CachedContent`` manager). This module is the equivalent for
the OpenAI-compatible family ``OpenAICompatibleLLM`` and its ``fireworks``
and ``nous`` subclasses which sent no affinity hint at all.
Default ``auto`` per member (``cache_affinity``). ``auto`` is an allowlist, not a
best-effort probe: it emits a hint only for hosts documented to accept one and
resolves to ``none`` for everything else, so an unknown OpenAI-compatible backend
never receives an unfamiliar field. Every helper here is fail-open when no id
can be derived the request goes out byte-identical to before. Set ``none`` to
disable entirely.
"""
from __future__ import annotations
import hashlib
import json
import logging
from enum import StrEnum
from typing import Any
from urllib.parse import urlparse
logger = logging.getLogger(__name__)
# xAI's documented cache-pinning header, and OpenAI's cache-routing field.
XAI_CONV_ID_HEADER = "x-grok-conv-id"
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",)
class CacheAffinityMode(StrEnum):
"""How (and whether) to pin a request to a backend prompt cache."""
NONE = "none"
XAI_CONV_ID = "xai_conv_id"
OPENAI_PROMPT_CACHE_KEY = "openai_prompt_cache_key"
AUTO = "auto"
def parse_cache_affinity(value: str | None) -> CacheAffinityMode:
"""Validate a configured cache-affinity mode, defaulting to ``none``.
Raises ``ValueError`` on an unrecognized value so a typo fails loudly at
provider construction rather than silently disabling the feature the whole
point of the setting is that its effect is invisible in the response.
"""
if not value:
return CacheAffinityMode.NONE
try:
return CacheAffinityMode(value.strip().lower())
except ValueError as e:
valid = ", ".join(mode.value for mode in CacheAffinityMode)
raise ValueError(f"Invalid cache_affinity {value!r}. Must be one of: {valid}.") from e
def _host_matches(hostname: str, domain: str) -> bool:
"""True when ``hostname`` is ``domain`` itself or a subdomain of it.
Parsed-host suffix matching, never a substring test: a bare
``"x.ai" in base_url`` also matches ``vertex.ai`` and
``https://x.ai.evil.example``. The in-tree Azure check
(``".openai.azure.com" in self.base_url``) gets away with a substring only
because its needle is long and dotted; ``x.ai`` is four characters.
"""
return hostname == domain or hostname.endswith(f".{domain}")
def resolve_cache_affinity(mode: CacheAffinityMode, provider: str, base_url: str | None) -> CacheAffinityMode:
"""Resolve ``auto`` to a concrete mode from the provider and base-URL host.
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
``none`` for everything else an unknown backend gets no unfamiliar field.
The xAI check is host-only and deliberately provider-independent: the
documented setup for an xAI endpoint is ``provider=openai`` plus an x.ai base
URL, exactly like Azure OpenAI, so keying on the provider name would miss it.
"""
if mode is not CacheAffinityMode.AUTO:
return mode
hostname = (urlparse(base_url).hostname or "") if base_url else ""
if hostname and any(_host_matches(hostname, domain) for domain in _XAI_DOMAINS):
return CacheAffinityMode.XAI_CONV_ID
if provider.lower() == "openai":
if not hostname:
return CacheAffinityMode.OPENAI_PROMPT_CACHE_KEY
if any(_host_matches(hostname, domain) for domain in _OPENAI_DOMAINS):
return CacheAffinityMode.OPENAI_PROMPT_CACHE_KEY
return CacheAffinityMode.NONE
def _first_message_fingerprint(messages: Any) -> str | None:
"""Hash the first message into a 32-hex id, or None if the shape is wrong.
Used only when no trace context is bound (direct provider use, tests). The
first message is the system prompt, so the id is stable as the message list
grows through an agent loop which is the property cache pinning needs
while differing across conversations whose first messages differ.
Shape-checked rather than truthiness-checked: a bare string ``messages``
would index to its first character and mint an id from garbage. Anything
unexpected returns None and the request goes out with no affinity hint.
"""
if not isinstance(messages, list) or not messages or not isinstance(messages[0], dict):
return None
try:
canonical = json.dumps(messages[0], sort_keys=True, ensure_ascii=False, default=str)
except (TypeError, ValueError):
logger.debug("Cache affinity: first message not serializable; sending no hint", exc_info=True)
return None
return hashlib.sha256(canonical.encode("utf-8")).hexdigest()[:32]
def cache_affinity_id(messages: Any) -> str | None:
"""Return the affinity id for the in-flight call, or None to send nothing.
Primary source is the operation's ``trace_id`` — one uuid per
retain/reflect/consolidation run, generated in ``LLMProvider.with_config``
and bound around every underlying provider call, so every LLM call of one
run shares it. That is engine identity rather than payload hashing: it stays
constant across a run even when the first message changes mid-run.
The value is always 32 lowercase hex characters, including for the trace_id
path (hashed rather than passed through) so the wire format is uniform and
carries no uuid semantics.
"""
from .llm_trace import current_trace_context
trace_ctx = current_trace_context()
if trace_ctx is not None and trace_ctx.trace_id:
return hashlib.sha256(str(trace_ctx.trace_id).encode("utf-8")).hexdigest()[:32]
return _first_message_fingerprint(messages)
def apply_cache_affinity(request: dict[str, Any], mode: CacheAffinityMode) -> None:
"""Add this request's cache-affinity hint to ``request`` in place.
``mode`` must already be resolved (see :func:`resolve_cache_affinity`);
``none`` and an unresolved ``auto`` add nothing.
User-wins semantics throughout, matching the file's ``setdefault`` precedent
in ``_apply_provider_extra_body_defaults``: an ``x-grok-conv-id`` the caller
already placed in ``extra_headers`` is kept, and a ``prompt_cache_key`` in
the operator's configured ``extra_body`` (the escape hatch for a backend
that wants its own value) suppresses ours entirely.
Never raises: when no id can be derived the request is left byte-identical
to a pre-affinity one.
"""
affinity_id = cache_affinity_id(request.get("messages"))
if affinity_id is None:
return
if mode is CacheAffinityMode.XAI_CONV_ID:
extra_headers = request.setdefault("extra_headers", {})
extra_headers.setdefault(XAI_CONV_ID_HEADER, affinity_id)
elif mode is CacheAffinityMode.OPENAI_PROMPT_CACHE_KEY:
# prompt_cache_key is a first-class named parameter on
# chat.completions.create() in the resolved openai SDK, so it goes at the
# top level rather than through extra_body. An operator value in
# extra_body would still reach the same wire field, so honour it and
# send nothing rather than sending both.
extra_body = request.get("extra_body")
if isinstance(extra_body, dict) and OPENAI_PROMPT_CACHE_KEY_PARAM in extra_body:
return
request.setdefault(OPENAI_PROMPT_CACHE_KEY_PARAM, affinity_id)
@@ -1,70 +0,0 @@
"""Shared causal-link taxonomy.
Retain writes only the canonical relationship. Transfer import/export also
preserves historical relationship types so existing banks keep their graph
semantics without allowing new retain output to create those types.
"""
from dataclasses import dataclass
from typing import Any
CANONICAL_CAUSAL_LINK_TYPE = "caused_by"
LEGACY_CAUSAL_LINK_TYPE_NAMES = ("causes", "enables", "prevents")
CANONICAL_CAUSAL_LINK_TYPES = frozenset({CANONICAL_CAUSAL_LINK_TYPE})
LEGACY_CAUSAL_LINK_TYPES = frozenset(LEGACY_CAUSAL_LINK_TYPE_NAMES)
CAUSAL_LINK_TYPES = (CANONICAL_CAUSAL_LINK_TYPE, *LEGACY_CAUSAL_LINK_TYPE_NAMES)
DEFAULT_CAUSAL_LINK_WEIGHT = 1.0
@dataclass(frozen=True)
class CausalLinkDescriptor:
"""One causal edge, parked on the curation archive while an endpoint is invalidated.
Invalidation moves a fact out of ``memory_units``, so the FK cascade deletes
its ``memory_links`` rows and nothing could recreate a causal edge, which
is extraction output rather than derived data. The descriptor is what the
archive row stores so revert can rematerialize the edge (#2864).
"""
from_unit_id: str
to_unit_id: str
link_type: str
weight: float = DEFAULT_CAUSAL_LINK_WEIGHT
def as_json_dict(self) -> dict[str, Any]:
"""Serializable form written to ``invalidated_memory_units.causal_links``.
The key names double as the column list of the ``jsonb_to_recordset``
read in ``snapshot_causal_links`` keep them in sync.
"""
return {
"from_unit_id": self.from_unit_id,
"to_unit_id": self.to_unit_id,
"link_type": self.link_type,
"weight": self.weight,
}
@classmethod
def from_json_dict(cls, raw: Any) -> "CausalLinkDescriptor | None":
"""Parse one stored descriptor, or None when it isn't a usable causal edge.
The archive column is plain JSON with no schema enforcement (a restore
from an older backup, or a hand-edited row, can put anything there), and
``memory_links`` has a ``link_type`` CHECK constraint so an unusable
entry is skipped rather than allowed to abort the whole revert.
"""
if not isinstance(raw, dict):
return None
from_unit_id = raw.get("from_unit_id")
to_unit_id = raw.get("to_unit_id")
link_type = raw.get("link_type")
if not from_unit_id or not to_unit_id or link_type not in CAUSAL_LINK_TYPES:
return None
return cls(
from_unit_id=str(from_unit_id),
to_unit_id=str(to_unit_id),
link_type=str(link_type),
weight=float(raw.get("weight") or DEFAULT_CAUSAL_LINK_WEIGHT),
)
@@ -109,41 +109,28 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
end.replace(hour=23, minute=59, second=59, microsecond=999999),
)
def safe_constraint(start: datetime | None, end: datetime | None) -> DateRange | NoTemporalConstraintSentinel:
if start is None or end is None:
return NO_TEMPORAL_CONSTRAINT
return constraint(start, end)
def subtract_months(months: int) -> datetime | 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)
def add_years(base_date: datetime, years: int) -> datetime | None:
def add_years(base_date: datetime, years: int) -> datetime:
year = base_date.year + years
if year < datetime.min.year or year > datetime.max.year:
return None
day = min(base_date.day, calendar.monthrange(year, base_date.month)[1])
return base_date.replace(year=year, day=day)
def add_days(base_date: datetime | None, days: int) -> datetime | None:
if base_date is None:
return None
try:
return base_date + timedelta(days=days)
except OverflowError:
return None
def has_chinese_temporal_context(match: re.Match[str]) -> bool:
if match.end() >= len(query):
return True
@@ -371,7 +358,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 +403,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)
@@ -453,11 +438,6 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
return NO_TEMPORAL_CONSTRAINT
return constraint(start, reference_date)
def safe_since_constraint(start: datetime | None) -> DateRange | NoTemporalConstraintSentinel:
if start is None:
return NO_TEMPORAL_CONSTRAINT
return since_constraint(start)
def since_from_period(
period: DateRange | None,
) -> DateRange | NoTemporalConstraintSentinel | None:
@@ -470,24 +450,24 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
return None
return since_constraint(day)
def relative_offset_datetime(amount: int, unit: str, direction: int) -> datetime | None:
def relative_offset_datetime(amount: int, unit: str, direction: int) -> datetime:
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)
def point_constraint_at_offset(amount: int, unit: str, direction: int) -> DateRange | NoTemporalConstraintSentinel:
def point_constraint_at_offset(amount: int, unit: str, direction: int) -> DateRange:
d = relative_offset_datetime(amount, unit, direction)
return safe_constraint(d, d)
return constraint(d, d)
def window_to_reference(amount: int, unit: str) -> DateRange | NoTemporalConstraintSentinel:
return safe_constraint(relative_offset_datetime(amount, unit, -1), reference_date)
def window_to_reference(amount: int, unit: str) -> DateRange:
return constraint(relative_offset_datetime(amount, unit, -1), reference_date)
def window_from_reference(amount: int, unit: str) -> DateRange | NoTemporalConstraintSentinel:
return safe_constraint(reference_date, relative_offset_datetime(amount, unit, 1))
def window_from_reference(amount: int, unit: str) -> DateRange:
return constraint(reference_date, relative_offset_datetime(amount, unit, 1))
# Chinese rule guide
#
@@ -616,8 +596,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))
@@ -803,8 +781,8 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
if relative_year_fixed_day_since_match:
year = relative_year_number(relative_year_fixed_day_since_match.group(1))
base = add_years(reference_date, year - reference_date.year)
d = add_days(base, fixed_day_offset(relative_year_fixed_day_since_match.group(2)))
return safe_since_constraint(d)
d = base + timedelta(days=fixed_day_offset(relative_year_fixed_day_since_match.group(2)))
return since_constraint(d)
fixed_day_since_match = chinese_search(
rf"(大大后天|大后天|后天|明天|明日|今天|今日|本日|当日|当天|昨天|昨日|大大前天|大前天|前天)"
@@ -821,7 +799,7 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
amount = parse_chinese_number(exact_relative_since_match.group(1))
unit = exact_relative_since_match.group(2)
if amount is not None:
return safe_since_constraint(relative_offset_datetime(amount, unit, -1))
return since_constraint(relative_offset_datetime(amount, unit, -1))
weekend_since_match = chinese_search(
rf"(?<![上下大小每个各隔])"
@@ -841,7 +819,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}"
@@ -921,8 +899,8 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
if relative_year_daypart_since_match:
year = relative_year_number(relative_year_daypart_since_match.group(1))
base = add_years(reference_date, year - reference_date.year)
d = add_days(base, daypart_day_offset(relative_year_daypart_since_match.group(2)))
return safe_since_constraint(d)
d = base + timedelta(days=daypart_day_offset(relative_year_daypart_since_match.group(2)))
return since_constraint(d)
daypart_since_match = chinese_search(
rf"(昨晚|昨夜|前晚|前夜|今晚|今早|今晨|明早|明晚|明夜){chinese_since_suffix_pattern}"
@@ -937,17 +915,17 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
if relative_year_daypart_match:
year = relative_year_number(relative_year_daypart_match.group(1))
base = add_years(reference_date, year - reference_date.year)
d = add_days(base, daypart_day_offset(relative_year_daypart_match.group(2)))
return safe_constraint(d, d)
d = base + timedelta(days=daypart_day_offset(relative_year_daypart_match.group(2)))
return constraint(d, d)
# Day-part abbreviations still resolve only to date granularity.
if chinese_search(r"昨晚|昨夜"):
d = add_days(reference_date, daypart_day_offset("昨晚"))
return safe_constraint(d, d)
d = reference_date + timedelta(days=daypart_day_offset("昨晚"))
return constraint(d, d)
if chinese_search(r"前晚|前夜"):
d = add_days(reference_date, daypart_day_offset("前晚"))
return safe_constraint(d, d)
d = reference_date + timedelta(days=daypart_day_offset("前晚"))
return constraint(d, d)
if chinese_search(r"今晚|今早|今晨"):
return constraint(reference_date, reference_date)
@@ -963,8 +941,8 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
if relative_year_fixed_day_match:
year = relative_year_number(relative_year_fixed_day_match.group(1))
base = add_years(reference_date, year - reference_date.year)
d = add_days(base, fixed_day_offset(relative_year_fixed_day_match.group(2)))
return safe_constraint(d, d)
d = base + timedelta(days=fixed_day_offset(relative_year_fixed_day_match.group(2)))
return constraint(d, d)
if chinese_search(r"昨天|昨日"):
d = reference_date - timedelta(days=1)
@@ -1037,7 +1015,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 +1023,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 +1039,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 +1048,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 +1056,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 +1065,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后])"
@@ -1107,7 +1085,7 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
end_amount = parse_chinese_number(amount_text[-1])
unit = adjacent_fuzzy_future_match.group(2)
if start_amount is not None and end_amount is not None:
return safe_constraint(
return constraint(
relative_offset_datetime(start_amount, unit, 1),
relative_offset_datetime(end_amount, unit, 1),
)
@@ -1115,7 +1093,7 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
few_future_match = chinese_search(rf"[几数]个?(天|日|周|星期|礼拜|月|年){chinese_relative_future_suffix_pattern}")
if few_future_match:
unit = few_future_match.group(1)
return safe_constraint(relative_offset_datetime(2, unit, 1), relative_offset_datetime(5, unit, 1))
return constraint(relative_offset_datetime(2, unit, 1), relative_offset_datetime(5, unit, 1))
exact_future_match = chinese_search(
rf"(?<![{_CHINESE_NUMERAL_PREFIX_CHARS}])([0-9]+|[{_CHINESE_NUMERAL_CHARS}]+)个?(天|日|周|星期|礼拜|月|年){chinese_relative_future_suffix_pattern}"
@@ -1135,7 +1113,7 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
second_amount = parse_chinese_number(adjacent_fuzzy_past_match.group(2))
unit = adjacent_fuzzy_past_match.group(3)
if first_amount is not None and second_amount is not None and second_amount == first_amount + 1:
return safe_constraint(
return constraint(
relative_offset_datetime(second_amount, unit, -1),
relative_offset_datetime(first_amount, unit, -1),
)
@@ -1166,7 +1144,7 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
return constraint(reference_date - timedelta(days=150), reference_date - timedelta(days=60))
if chinese_search(r"一两年前|[两二]三年前|三两年前"):
return safe_constraint(add_years(reference_date, -3), add_years(reference_date, -1))
return constraint(add_years(reference_date, -3), add_years(reference_date, -1))
rolling_this_adjacent_match = chinese_search(
r"这(一两|[两二]三|三两|三四|四五|五六|六七|七八|八九|九十)个?(天|日|周|星期|礼拜|月|年)"
@@ -1176,7 +1154,7 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
end_amount = 3 if amount_text in ("一两", "三两") else parse_chinese_number(amount_text[-1])
unit = rolling_this_adjacent_match.group(2)
if end_amount is not None:
return safe_constraint(relative_offset_datetime(end_amount, unit, -1), reference_date)
return constraint(relative_offset_datetime(end_amount, unit, -1), reference_date)
rolling_this_count_match = chinese_search(rf"这([0-9]+|[{_CHINESE_NUMERAL_CHARS}]+)个?(天|日|周|星期|礼拜|月|年)")
if rolling_this_count_match:
@@ -1213,7 +1191,7 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
end_amount = 3 if amount_text in ("一两", "三两") else parse_chinese_number(amount_text[-1])
unit = rolling_past_adjacent_match.group(3)
if end_amount is not None:
return safe_constraint(relative_offset_datetime(end_amount, unit, -1), reference_date)
return constraint(relative_offset_datetime(end_amount, unit, -1), reference_date)
rolling_past_few_match = chinese_search(r"(过去|近|最近)几个?(天|日|周|星期|礼拜|月|年)")
if rolling_past_few_match:
@@ -1234,14 +1212,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 +1282,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 +1419,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 +1427,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 +1445,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 +1570,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 +1577,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 +1591,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 +1710,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 +1743,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})"):
File diff suppressed because it is too large Load Diff
@@ -17,20 +17,6 @@ _MISSION_PRIORITY_NOTE = (
"DECISION GUIDE, or OUTPUT FORMAT below, the MISSION takes priority."
)
# Default language rule — used only when HINDSIGHT_API_LLM_OUTPUT_LANGUAGE is
# unset. Without it the whole prompt is English and multilingual models drift:
# Chinese source facts intermittently produce English observations. Retain's
# fact extraction carries the equivalent rule (see _BASE_FACT_EXTRACTION_PROMPT),
# so this makes "preserve the source language" the pipeline-wide default. When an
# output language IS configured, this section is omitted and
# output_language_directive() takes over — the two must never both be present or
# they contradict each other.
_DEFAULT_LANGUAGE_RULE = """## LANGUAGE
Write every observation in the language of its own source facts never translate them. Per observation, not per batch: when one merges facts of several languages, the majority wins. Proper nouns, identifiers, and units stay verbatim.
When an existing observation is written in a different language from the new facts updating it, do NOT edit its wording in place that is what produces an English sentence with a Chinese detail bolted on. Discard the old phrasing and compose the merged observation from scratch in the new facts' language."""
_PROCESSING_RULES = """## PROCESSING RULES
1. PREFER UPDATE OVER CREATE (when there is something to merge with): if new facts describe the same canonical event, statement, decision, claim, or recurring pattern already covered by an existing observation, UPDATE that observation and attach the new facts as evidence. Do NOT create a near-duplicate sibling. One canonical observation with many source facts is always better than many siblings with one source fact each. Merge aggressively on: same named event, same diagnostic finding, same architectural decision, same recurring claim. **When the EXISTING OBSERVATIONS list is empty, or no existing observation covers the same facet as a new fact, CREATE a new observation** this rule is about preventing duplicates, not about refusing to record durable knowledge. CREATE is the correct default for any structurally distinct event, claim, or pattern that has no existing match.
@@ -51,36 +37,19 @@ _PROCESSING_RULES = """## PROCESSING RULES
9. KEEP DISTINCT TOPICS DISTINCT: do not merge observations about different people, entities, or unrelated topics. Merging is for the same canonical fact recurring not for related-but-distinct claims."""
# Field-by-field definitions of the input shape used by the cached system
# prefix. The call site runs .format(), so these strings must contain no braces.
_FACT_FIELDS = """One per line, formatted as `[uuid] fact text (temporal fields)`:
- `[uuid]`: the fact's identifier — copy it verbatim into `source_fact_ids`
- `occurred_start` / `occurred_end`: when the described event happened. This can be long before the fact was stated a fact recorded today may describe a 2019 event.
- `mentioned_at`: when the source material that states this fact was written. This is the fact's recency: how up to date the statement is, NOT when it was added to memory. A fact taken from an old document keeps its old `mentioned_at` even if it was only just processed."""
_OBSERVATION_FIELDS = """- `id`: unique identifier — copy this exactly when issuing an UPDATE or DELETE
- `text`: the observation content
- `proof_count`: how many source facts this observation has already merged
- `occurred_start` / `occurred_end`: the span of the events behind the observation earliest start and latest end across its source facts
- `mentioned_at`: the latest of the `mentioned_at` values of its source facts the most recent point at which this observation was stated
- `source_memories`: the supporting facts behind this observation. May be partial or absent for large observations the count above remains the true total. Each entry carries the same `text` and temporal fields as a new fact, plus:
- `context`: optional surrounding context for that fact"""
# Stable description of the input shape. For the cached split path this lives in
# the system prefix (build_consolidation_system_prompt) so it is not re-sent on
# every batch; the per-batch user message then carries only the actual data.
_INPUT_FORMAT_NOTE = f"""## INPUT FORMAT
_INPUT_FORMAT_NOTE = """## INPUT FORMAT
Each request provides new facts and existing observations. Every temporal field is optional and is omitted when unknown.
### New facts
{_FACT_FIELDS}
### Existing observations
A JSON array pooled from recalls across the new facts. Each entry has:
{_OBSERVATION_FIELDS}"""
Each request provides new facts and existing observations:
- New facts: one per line, each prefixed with its `[uuid]`, followed by the fact text and optional temporal fields.
- Existing observations: a JSON array pooled from recalls across the new facts. Each entry has:
- `id`: unique identifier copy this exactly when issuing an UPDATE or DELETE
- `text`: the observation content
- `proof_count`: number of supporting memories
- `occurred_start` / `occurred_end`: temporal range of source facts
- `source_memories`: array of supporting facts with their text and dates"""
# Per-batch data section for the cached split path — the stable format
# explanation above is omitted here (it lives in the cached prefix); only the
@@ -95,6 +64,24 @@ _SPLIT_INPUT_SECTION = """## INPUT
{observations_text}"""
# Data section — format placeholders {facts_text} and {observations_text} are substituted at call time
_INPUT_SECTION = """## INPUT
### New facts
{facts_text}
### Existing observations
JSON array, pooled from recalls across all new facts above. Each entry has:
- `id`: unique identifier copy this exactly when issuing an UPDATE or DELETE
- `text`: the observation content
- `proof_count`: number of supporting memories
- `occurred_start` / `occurred_end`: temporal range of source facts
- `source_memories`: array of supporting facts with their text and dates
{observations_text}"""
_DECISION_GUIDE = """## DECISION GUIDE
- **Same canonical event, decision, claim, or facet as an existing observation UPDATE** (use `observation_id` + new `source_fact_ids`).
@@ -155,6 +142,39 @@ Expected output (UPDATE for the state change; CREATE for the unrelated work-hour
- Return `{{"creates": [], "updates": [], "deletes": []}}` if nothing durable is found."""
def build_batch_consolidation_prompt(
observations_mission: str | None = None,
observation_capacity_note: str | None = None,
llm_output_language: str | None = None,
) -> str:
"""
Build the consolidation prompt for batch mode (multiple facts per LLM call).
The mission defines *what* to track (customisable per bank) and takes
priority over the built-in processing rules when the two conflict.
Processing rules, decision guide, and output format are always present.
When ``llm_output_language`` is set, observations are emitted in that
language.
"""
mission = escape_for_prompt(observations_mission or _DEFAULT_MISSION)
capacity_section = ""
if observation_capacity_note:
capacity_section = f"\n\n## CAPACITY CONSTRAINT\n\n{escape_for_prompt(observation_capacity_note)}"
return (
"You are a memory consolidation system. Synthesize new facts into "
"observations, merging with existing observations when appropriate.\n\n"
f"## MISSION\n\n{mission}\n\n"
f"{_MISSION_PRIORITY_NOTE}"
f"{capacity_section}\n\n"
f"{_PROCESSING_RULES}\n\n"
f"{_INPUT_SECTION}\n\n"
f"{_DECISION_GUIDE}\n\n"
f"{_OUTPUT_SECTION}" + output_language_directive(llm_output_language)
)
def build_consolidation_system_prompt(
llm_output_language: str | None = None,
) -> str:
@@ -169,17 +189,11 @@ def build_consolidation_system_prompt(
bank and a single CachedContent serves them all. Returns final text
(brace-escaped examples already unescaped) for verbatim use as system message
and cached prefix.
``llm_output_language`` picks between two mutually exclusive language rules:
unset keeps each observation in the language of its own source facts (the
default), set forces every observation into that one configured language.
"""
language_section = "" if llm_output_language else f"{_DEFAULT_LANGUAGE_RULE}\n\n"
template = (
"You are a memory consolidation system. Synthesize new facts into "
"observations, merging with existing observations when appropriate.\n\n"
f"{_MISSION_PRIORITY_NOTE}\n\n"
f"{language_section}"
f"{_PROCESSING_RULES}\n\n"
f"{_INPUT_FORMAT_NOTE}\n\n"
f"{_DECISION_GUIDE}\n\n"
@@ -8,10 +8,10 @@ Configuration via environment variables - see hindsight_api.config for all env v
import asyncio
import logging
import os
import warnings
from abc import ABC, abstractmethod
from concurrent.futures import ThreadPoolExecutor
from typing import Any
import httpx
@@ -19,8 +19,8 @@ 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_CPU_MEM_ARENA,
DEFAULT_RERANKER_FLASHRANK_MODEL,
DEFAULT_RERANKER_GOOGLE_MODEL,
DEFAULT_RERANKER_LITELLM_MAX_TOKENS_PER_DOC,
@@ -34,19 +34,58 @@ from ..config import (
DEFAULT_RERANKER_TEI_MAX_CONCURRENT,
DEFAULT_RERANKER_ZEROENTROPY_MODEL,
DEFAULT_ZEROENTROPY_BASE_URL,
RerankerMemberConfig,
ENV_RERANKER_ALIBABA_API_KEY,
ENV_RERANKER_COHERE_API_KEY,
ENV_RERANKER_FLASHRANK_CACHE_DIR,
ENV_RERANKER_FLASHRANK_CPU_MEM_ARENA,
ENV_RERANKER_FLASHRANK_MODEL,
ENV_RERANKER_GOOGLE_PROJECT_ID,
ENV_RERANKER_PROVIDER,
ENV_RERANKER_SILICONFLOW_API_KEY,
ENV_RERANKER_TEI_URL,
ENV_RERANKER_ZEROENTROPY_API_KEY,
)
from .bank_attribution import reranker_bank_attribution_headers
from .local_device import (
release_local_inference_memory,
resolve_model_device_type,
select_local_device,
)
from .tei_retry import tei_retry_delay
logger = logging.getLogger(__name__)
def _resolve_malloc_trim():
"""Return a callable that asks glibc to release freed heap pages to the OS.
Local CPU rerankers (FlashRank/ONNX, SentenceTransformers/torch) allocate
large transient numpy/tensor buffers per call. On Linux glibc, those pages
are freed at the Python level but kept by the allocator as a high-water
mark RSS grows monotonically across many recalls (see issue #1717).
Calling `malloc_trim(0)` after each batch returns those pages to the OS.
Resolved once at import; returns a no-op on non-glibc platforms (macOS,
musl, Windows) where the call is unavailable or unnecessary.
"""
import sys
if sys.platform != "linux":
return lambda: None
import ctypes
import ctypes.util
libc_path = ctypes.util.find_library("c")
if libc_path is None:
return lambda: None
try:
libc = ctypes.CDLL(libc_path)
trim = libc.malloc_trim
except (OSError, AttributeError):
# Not glibc (musl has no malloc_trim) or libc lookup failed.
return lambda: None
trim.argtypes = [ctypes.c_size_t]
trim.restype = ctypes.c_int
return lambda: trim(0)
_malloc_trim = _resolve_malloc_trim()
class CrossEncoderModel(ABC):
"""
Abstract base class for cross-encoder reranking.
@@ -60,15 +99,6 @@ class CrossEncoderModel(ABC):
"""Return a human-readable name for this provider (e.g., 'local', 'tei')."""
pass
@property
def blocking_init(self) -> bool:
"""Whether ``initialize()`` blocks the event loop (loads a model in-process).
Callers run those in a thread pool. Remote providers leave this False, and
so does :class:`MultiCrossEncoder` it offloads its own members.
"""
return False
@abstractmethod
async def initialize(self) -> None:
"""
@@ -120,7 +150,6 @@ class LocalSTCrossEncoder(CrossEncoderModel):
fp16: bool = False,
bucket_batching: bool = False,
batch_size: int = DEFAULT_RERANKER_LOCAL_BATCH_SIZE,
allow_mps: bool = False,
):
"""
Initialize local SentenceTransformers cross-encoder.
@@ -142,9 +171,6 @@ class LocalSTCrossEncoder(CrossEncoderModel):
Default: False (opt-in via env var).
batch_size: Batch size for predict() calls. Optimal values vary by
hardware and model (MPS: 32, CUDA: 128+). Default: 32.
allow_mps: Opt in to the Apple Silicon MPS GPU. Disabled by default
because MPS leaks memory under variable-length workloads
(see engine/local_device.py). Default: False
"""
self.model_name = model_name or DEFAULT_RERANKER_LOCAL_MODEL
self.force_cpu = force_cpu
@@ -152,19 +178,13 @@ class LocalSTCrossEncoder(CrossEncoderModel):
self.fp16 = fp16
self.bucket_batching = bucket_batching
self.batch_size = batch_size
self.allow_mps = allow_mps
self._model = None
self._device_type: str = "cpu"
LocalSTCrossEncoder._max_concurrent = max_concurrent
@property
def provider_name(self) -> str:
return "local"
@property
def blocking_init(self) -> bool:
return True
async def initialize(self) -> None:
"""Load the cross-encoder model and initialize the executor."""
if self._model is not None:
@@ -180,13 +200,33 @@ class LocalSTCrossEncoder(CrossEncoderModel):
logger.info(f"Reranker: initializing local provider with model {self.model_name}")
# Determine device based on hardware availability. We always set
# low_cpu_mem_usage=False to prevent lazy loading (meta tensors) which can
# cause issues when accelerate is installed but no GPU is available.
# Determine device based on hardware availability.
# We always set low_cpu_mem_usage=False to prevent lazy loading (meta tensors)
# which can cause issues when accelerate is installed but no GPU is available.
# Note: We do NOT use device_map because CrossEncoder internally calls .to(device)
# after loading, which conflicts with accelerate's device_map handling.
# MPS is opt-in (allow_mps) — see engine/local_device.py for why.
device = select_local_device(self.force_cpu, self.allow_mps)
import torch
# Force CPU mode if configured (used in daemon mode to avoid MPS/XPC issues on macOS)
if self.force_cpu:
device = "cpu"
logger.info("Reranker: forcing CPU mode (HINDSIGHT_API_RERANKER_LOCAL_FORCE_CPU=1)")
else:
# Check for GPU (CUDA), Apple Silicon (MPS), or Intel XPU
# Wrap in try-except to gracefully handle any device detection issues
# (e.g., in CI environments or when PyTorch is built without GPU support)
device = "cpu" # Default to CPU
try:
has_gpu = torch.cuda.is_available() or (
hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
)
# Intel Arc XPU support — torch.xpu is available when the XPU build is loaded
if not has_gpu and hasattr(torch, "xpu"):
has_gpu = torch.xpu.is_available()
if has_gpu:
device = None # Let sentence-transformers auto-detect GPU/MPS/XPU
except Exception as e:
logger.warning(f"Failed to detect GPU/MPS/XPU, falling back to CPU: {e}")
# Patch transformers 5.x compatibility for models using XLM-RoBERTa
# (e.g., jina-reranker-v2-base-multilingual). transformers 5.x removed
@@ -230,11 +270,9 @@ class LocalSTCrossEncoder(CrossEncoderModel):
# Restore original logging level
transformers_logger.setLevel(original_level)
self._device_type = resolve_model_device_type(self._model)
# FP16 inference: convert model weights to half precision.
# Empirically validated: 27-36% faster on MPS, quality-identical (20/20 overlap).
if self.fp16 and self._device_type != "cpu":
if self.fp16 and device != "cpu":
self._model.model.half()
logger.info("Reranker: FP16 inference enabled")
@@ -277,7 +315,7 @@ class LocalSTCrossEncoder(CrossEncoderModel):
scores = self._model.predict(pairs, batch_size=self.batch_size, show_progress_bar=False)
return scores.tolist() if hasattr(scores, "tolist") else list(scores)
finally:
release_local_inference_memory(self._device_type)
_malloc_trim()
async def predict(self, pairs: list[tuple[str, str]]) -> list[float]:
"""
@@ -393,20 +431,14 @@ class RemoteTEICrossEncoder(CrossEncoderModel):
await asyncio.sleep(delay)
delay *= 2 # Exponential backoff
except httpx.HTTPStatusError as e:
# TEI uses 429 as normal overload backpressure. Retry it with
# the same bounded budget as transient server errors.
if (e.response.status_code == 429 or e.response.status_code >= 500) and attempt < self.max_retries:
# Retry on 5xx server errors
if e.response.status_code >= 500 and attempt < self.max_retries:
last_error = e
sleep_delay = tei_retry_delay(
e.response,
delay,
request_timeout=self.timeout,
)
logger.warning(
f"TEI transient error (attempt {attempt + 1}/{self.max_retries + 1}): {e}. "
f"Retrying in {sleep_delay:.2f}s..."
f"TEI server error (attempt {attempt + 1}/{self.max_retries + 1}): {e}. "
f"Retrying in {delay}s..."
)
await asyncio.sleep(sleep_delay)
await asyncio.sleep(delay)
delay *= 2
else:
raise
@@ -452,7 +484,6 @@ class RemoteTEICrossEncoder(CrossEncoderModel):
semaphore,
"POST",
f"{self.base_url}/rerank",
headers=reranker_bank_attribution_headers(),
json={
"query": query,
"texts": texts,
@@ -593,11 +624,7 @@ class _CohereCompatibleRerankClient:
if self.include_top_n:
body["top_n"] = len(texts)
response = await self._async_client.post(
self.rerank_url,
headers=reranker_bank_attribution_headers(),
json=body,
)
response = await self._async_client.post(self.rerank_url, json=body)
response.raise_for_status()
result = response.json()
@@ -873,7 +900,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,18 +913,12 @@ 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
@property
@@ -969,26 +989,12 @@ 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.
"""Synchronous predict - processes each query group."""
from flashrank import RerankRequest
FlashRank scores every passage of a request in one ONNX forward pass, and
that pass allocates attention tensors sized ``batch * heads * seq^2``. At
the default reranker candidate cap that is gigabytes per call, which OOM-
killed containers on large banks (issue #3355): the burst scales with the
candidate pool the retrieval arms produce, not with how much work the
caller asked for. FlashRank also pads a request to its longest passage, so
one long candidate inflates the sequence length for every other one.
Splitting into ``batch_size`` chunks bounds the peak the same way the
local and TEI providers already do. Scores are identical either way
passages are scored independently, so batching changes only the
allocation profile.
"""
if not pairs:
return []
from flashrank import RerankRequest
try:
# Group pairs by query
query_groups: dict[str, list[tuple[int, str]]] = {}
@@ -1000,29 +1006,24 @@ 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:
release_local_inference_memory(self._device_type)
_malloc_trim()
async def predict(self, pairs: list[tuple[str, str]]) -> list[float]:
"""
@@ -1150,7 +1151,6 @@ class LiteLLMCrossEncoder(CrossEncoderModel):
# LiteLLM /rerank follows Cohere API format
response = await self._async_client.post(
f"{self.api_base}/rerank",
headers=reranker_bank_attribution_headers(),
json={
"model": self.model,
"query": query,
@@ -1269,11 +1269,10 @@ class LiteLLMSDKCrossEncoder(CrossEncoderModel):
indices = [idx for idx, _ in indexed_texts]
# Build kwargs for rerank call
rerank_kwargs: dict[str, Any] = {
rerank_kwargs = {
"model": self.model,
"query": query,
"documents": texts,
"headers": reranker_bank_attribution_headers(),
}
if self.api_key:
rerank_kwargs["api_key"] = self.api_key
@@ -1282,9 +1281,21 @@ class LiteLLMSDKCrossEncoder(CrossEncoderModel):
response = await self._litellm.arerank(**rerank_kwargs)
for result in response.results:
original_idx = result["index"]
all_scores[indices[original_idx]] = result["relevance_score"]
# Map scores back to original positions
# Response format: RerankResponse with results list
# Each result is a TypedDict with "index" and "relevance_score"
if hasattr(response, "results") and response.results:
for result in response.results:
# Results are TypedDicts, use dict-style access
original_idx = result["index"]
score = result.get("relevance_score", result.get("score", 0.0))
all_scores[indices[original_idx]] = score
elif isinstance(response, list):
# Direct list of scores (unlikely but defensive)
for i, score in enumerate(response):
all_scores[indices[i]] = score
else:
logger.warning(f"Unexpected response format from LiteLLM rerank: {type(response)}")
return all_scores
@@ -1603,247 +1614,131 @@ class AlibabaCloudCrossEncoder(CrossEncoderModel):
return await self._client.predict(pairs)
class MultiCrossEncoder(CrossEncoderModel):
"""Failover across an ordered chain of cross-encoders.
Member 0 is the primary (the unindexed ``HINDSIGHT_API_RERANKER_*`` config);
members 1..N are the indexed fallbacks. Each ``predict`` tries members in order
and returns the first usable set of scores, so an unreachable reranker costs
ranking quality (whatever the next member gives) instead of the whole recall.
Put ``rrf`` last to degrade to the fusion order rather than failing.
Each member keeps its own retry budget, so we only advance after a member has
exhausted its retries and raised. A member that fails to initialize is not
fatal that is the point of the chain it is retried lazily on the next
request that reaches it.
def create_cross_encoder_from_env() -> CrossEncoderModel:
"""
Create a CrossEncoderModel instance based on configuration.
def __init__(self, members: list[CrossEncoderModel]) -> None:
if len(members) < 2:
raise ValueError("MultiCrossEncoder requires at least two members")
self._members = members
self._ready = [False] * len(members)
self._locks = [asyncio.Lock() for _ in members]
self._active = 0
@property
def provider_name(self) -> str:
"""The provider of the member that last served a request (primary before any).
Callers use this to detect a passthrough reranker, so it has to track the
member actually serving rather than name the chain: a chain that has
degraded to its ``rrf`` member is passthrough. Concurrent requests share it,
so a request that fails over can briefly mislabel a neighbour this only
tunes downstream scoring, never correctness.
"""
return self._members[self._active].provider_name
async def _initialize_member(self, index: int) -> None:
"""Initialize one member, off the event loop when it loads a model in-process."""
member = self._members[index]
if member.blocking_init:
loop = asyncio.get_running_loop()
await loop.run_in_executor(None, lambda: asyncio.run(member.initialize()))
else:
await member.initialize()
self._ready[index] = True
async def _ensure_member_ready(self, index: int) -> None:
async with self._locks[index]:
if not self._ready[index]:
await self._initialize_member(index)
async def initialize(self) -> None:
"""Initialize every member, tolerating members that are down.
Members initialize concurrently so one unreachable member cannot eat the
startup budget the others need. Failures are logged and retried on use.
"""
results = await asyncio.gather(
*(self._ensure_member_ready(i) for i in range(len(self._members))),
return_exceptions=True,
)
for index, result in enumerate(results):
if isinstance(result, BaseException):
logger.warning(
"Reranker member %d (%s) failed to initialize: %s; it will be retried on use",
index,
self._members[index].provider_name,
result,
)
if not any(self._ready):
logger.error("Reranker: no member of the failover chain initialized; recall will retry them per request")
async def predict(self, pairs: list[tuple[str, str]]) -> list[float]:
"""Score ``pairs`` with the first member that answers usably."""
last_exc: BaseException | None = None
for index, member in enumerate(self._members):
try:
if not self._ready[index]:
await self._ensure_member_ready(index)
scores = await member.predict(pairs)
if len(scores) != len(pairs):
raise RuntimeError(f"returned {len(scores)} scores for {len(pairs)} pairs")
except Exception as e: # noqa: BLE001 - re-raised below if no member answers
last_exc = e
remaining = len(self._members) - index - 1
logger.warning(
"Reranker member %d (%s) failed: %s%s",
index,
member.provider_name,
e,
f"; trying next member ({remaining} left)" if remaining else "; no members left",
)
continue
if index != self._active:
logger.info(
"Reranker: now serving from member %d (%s)",
index,
member.provider_name,
)
self._active = index
return scores
# All members failed; surface the last error (loop ran at least once).
assert last_exc is not None
raise last_exc
def create_cross_encoder(member: RerankerMemberConfig) -> CrossEncoderModel:
"""
Create a CrossEncoderModel for one member of the reranker chain.
``member`` is the primary (index 0, the unindexed ``HINDSIGHT_API_RERANKER_*``
config) or an indexed fallback. Missing-setting errors name the member's own
env var, so a chain misconfiguration points at the exact indexed variable.
Args:
member: Resolved settings for this member
Reads configuration via get_config() to ensure consistency across the codebase.
Returns:
Configured CrossEncoderModel instance
"""
provider = member.provider.lower()
from ..config import get_config
config = get_config()
provider = config.reranker_provider.lower()
if provider == "tei":
url = member.tei_url
url = config.reranker_tei_url
if not url:
raise ValueError(f"{member.env_name('TEI_URL')} is required when {member.env_name('PROVIDER')} is 'tei'")
raise ValueError(f"{ENV_RERANKER_TEI_URL} is required when {ENV_RERANKER_PROVIDER} is 'tei'")
return RemoteTEICrossEncoder(
base_url=url,
timeout=member.tei_http_timeout,
batch_size=member.tei_batch_size,
max_concurrent=member.tei_max_concurrent,
timeout=config.reranker_tei_http_timeout,
batch_size=config.reranker_tei_batch_size,
max_concurrent=config.reranker_tei_max_concurrent,
)
elif provider == "local":
return LocalSTCrossEncoder(
model_name=member.local_model,
max_concurrent=member.local_max_concurrent,
force_cpu=member.local_force_cpu,
trust_remote_code=member.local_trust_remote_code,
fp16=member.local_fp16,
bucket_batching=member.local_bucket_batching,
batch_size=member.local_batch_size,
allow_mps=member.local_allow_mps,
model_name=config.reranker_local_model,
max_concurrent=config.reranker_local_max_concurrent,
force_cpu=config.reranker_local_force_cpu,
trust_remote_code=config.reranker_local_trust_remote_code,
fp16=config.reranker_local_fp16,
bucket_batching=config.reranker_local_bucket_batching,
batch_size=config.reranker_local_batch_size,
)
elif provider == "cohere":
api_key = member.cohere_api_key
api_key = config.reranker_cohere_api_key
if not api_key:
raise ValueError(
f"{member.env_name('COHERE_API_KEY')} is required when {member.env_name('PROVIDER')} is 'cohere'"
)
raise ValueError(f"{ENV_RERANKER_COHERE_API_KEY} is required when {ENV_RERANKER_PROVIDER} is 'cohere'")
return CohereCrossEncoder(
api_key=api_key,
model=member.cohere_model,
base_url=member.cohere_base_url,
timeout=member.cohere_timeout,
model=config.reranker_cohere_model,
base_url=config.reranker_cohere_base_url,
timeout=config.reranker_cohere_timeout,
)
elif provider == "openrouter":
api_key = member.openrouter_api_key
api_key = config.reranker_openrouter_api_key
if not api_key:
shared = ", HINDSIGHT_API_OPENROUTER_API_KEY, or HINDSIGHT_API_LLM_API_KEY" if member.index == 0 else ""
raise ValueError(
f"{member.env_name('OPENROUTER_API_KEY')}{shared} is required "
f"when {member.env_name('PROVIDER')} is 'openrouter'"
"HINDSIGHT_API_RERANKER_OPENROUTER_API_KEY, HINDSIGHT_API_OPENROUTER_API_KEY, "
f"or HINDSIGHT_API_LLM_API_KEY is required when {ENV_RERANKER_PROVIDER} is 'openrouter'"
)
return CohereCrossEncoder(
api_key=api_key,
model=member.openrouter_model,
base_url=member.openrouter_base_url,
timeout=member.openrouter_timeout,
model=config.reranker_openrouter_model,
base_url=config.reranker_openrouter_base_url,
timeout=config.reranker_openrouter_timeout,
)
elif provider == "flashrank":
return FlashRankCrossEncoder(
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,
)
model = os.environ.get(ENV_RERANKER_FLASHRANK_MODEL, DEFAULT_RERANKER_FLASHRANK_MODEL)
cache_dir = os.environ.get(ENV_RERANKER_FLASHRANK_CACHE_DIR, DEFAULT_RERANKER_FLASHRANK_CACHE_DIR)
cpu_mem_arena = os.environ.get(
ENV_RERANKER_FLASHRANK_CPU_MEM_ARENA, str(DEFAULT_RERANKER_FLASHRANK_CPU_MEM_ARENA)
).lower() in ("true", "1", "yes")
return FlashRankCrossEncoder(model_name=model, cache_dir=cache_dir, cpu_mem_arena=cpu_mem_arena)
elif provider == "litellm":
return LiteLLMCrossEncoder(
api_base=member.litellm_api_base,
api_key=member.litellm_api_key,
model=member.litellm_model,
max_tokens_per_doc=member.litellm_max_tokens_per_doc,
timeout=member.litellm_timeout,
api_base=config.reranker_litellm_api_base,
api_key=config.reranker_litellm_api_key,
model=config.reranker_litellm_model,
max_tokens_per_doc=config.reranker_litellm_max_tokens_per_doc,
timeout=config.reranker_litellm_timeout,
)
elif provider == "litellm-sdk":
return LiteLLMSDKCrossEncoder(
api_key=member.litellm_sdk_api_key or None,
model=member.litellm_sdk_model,
api_base=member.litellm_sdk_api_base,
max_tokens_per_doc=member.litellm_max_tokens_per_doc,
timeout=member.litellm_sdk_timeout,
api_key=config.reranker_litellm_sdk_api_key or None,
model=config.reranker_litellm_sdk_model,
api_base=config.reranker_litellm_sdk_api_base,
max_tokens_per_doc=config.reranker_litellm_max_tokens_per_doc,
timeout=config.reranker_litellm_sdk_timeout,
)
elif provider == "zeroentropy":
api_key = member.zeroentropy_api_key
api_key = config.reranker_zeroentropy_api_key
if not api_key:
raise ValueError(
f"{member.env_name('ZEROENTROPY_API_KEY')} is required "
f"when {member.env_name('PROVIDER')} is 'zeroentropy'"
f"{ENV_RERANKER_ZEROENTROPY_API_KEY} is required when {ENV_RERANKER_PROVIDER} is 'zeroentropy'"
)
return ZeroEntropyCrossEncoder(
api_key=api_key,
model=member.zeroentropy_model,
base_url=member.zeroentropy_base_url,
timeout=member.zeroentropy_timeout,
model=config.reranker_zeroentropy_model,
base_url=config.reranker_zeroentropy_base_url,
timeout=config.reranker_zeroentropy_timeout,
)
elif provider == "siliconflow":
api_key = member.siliconflow_api_key
api_key = config.reranker_siliconflow_api_key
if not api_key:
raise ValueError(
f"{member.env_name('SILICONFLOW_API_KEY')} is required "
f"when {member.env_name('PROVIDER')} is 'siliconflow'"
f"{ENV_RERANKER_SILICONFLOW_API_KEY} is required when {ENV_RERANKER_PROVIDER} is 'siliconflow'"
)
return SiliconFlowCrossEncoder(
api_key=api_key,
model=member.siliconflow_model,
base_url=member.siliconflow_base_url,
timeout=member.siliconflow_timeout,
model=config.reranker_siliconflow_model,
base_url=config.reranker_siliconflow_base_url,
timeout=config.reranker_siliconflow_timeout,
)
elif provider == "google":
project_id = member.google_project_id
project_id = config.reranker_google_project_id
if not project_id:
shared = " (or HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID)" if member.index == 0 else ""
raise ValueError(
f"{member.env_name('GOOGLE_PROJECT_ID')}{shared} "
f"is required when {member.env_name('PROVIDER')} is 'google'"
f"{ENV_RERANKER_GOOGLE_PROJECT_ID} (or HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID) "
f"is required when {ENV_RERANKER_PROVIDER} is 'google'"
)
return GoogleCrossEncoder(
project_id=project_id,
model=member.google_model,
service_account_key=member.google_service_account_key,
timeout=member.google_timeout,
model=config.reranker_google_model,
service_account_key=config.reranker_google_service_account_key,
timeout=config.reranker_google_timeout,
)
elif provider == "alibaba":
api_key = member.alibaba_api_key
api_key = config.reranker_alibaba_api_key
if not api_key:
raise ValueError(
f"{member.env_name('ALIBABA_API_KEY')} is required when {member.env_name('PROVIDER')} is 'alibaba'"
)
raise ValueError(f"{ENV_RERANKER_ALIBABA_API_KEY} is required when {ENV_RERANKER_PROVIDER} is 'alibaba'")
return AlibabaCloudCrossEncoder(
api_key=api_key,
model=member.alibaba_model,
timeout=member.alibaba_timeout,
model=config.reranker_alibaba_model,
timeout=config.reranker_alibaba_timeout,
)
elif provider == "rrf":
return RRFPassthroughCrossEncoder()
@@ -1853,23 +1748,3 @@ def create_cross_encoder(member: RerankerMemberConfig) -> CrossEncoderModel:
raise ValueError(
f"Unknown reranker provider: {provider}. Supported: 'local', 'tei', 'cohere', 'zeroentropy', 'siliconflow', 'alibaba', 'google', 'flashrank', 'litellm', 'litellm-sdk', 'rrf', 'jina-mlx'"
)
def create_cross_encoder_from_env() -> CrossEncoderModel:
"""
Create the configured reranker, based on configuration.
Reads configuration via get_config() to ensure consistency across the codebase.
With no ``HINDSIGHT_API_RERANKER_<n>_*`` members configured (the default) this
is the single configured reranker; otherwise the chain is wrapped in a
:class:`MultiCrossEncoder` that fails over across members in order.
Returns:
Configured CrossEncoderModel instance
"""
from ..config import get_config
chain = get_config().reranker_chain()
if len(chain) == 1:
return create_cross_encoder(chain[0])
return MultiCrossEncoder([create_cross_encoder(member) for member in chain])
@@ -67,28 +67,16 @@ def create_database_backend(backend_type: str) -> DatabaseBackend:
return _get_backend_class(backend_type)()
_OPS_CACHE: dict[str, DataAccessOps] = {}
def create_data_access_ops(backend_type: str) -> DataAccessOps:
"""Factory: the DataAccessOps for a backend name.
Returns a per-dialect SINGLETON: ``DataAccessOps`` is stateless (it only builds and runs SQL),
so one shared instance per dialect is correct and it means the database backend and the
memories store hold the *same* ops object, so a test that patches a method on it (e.g.
``enqueue_graph_maintenance``) observes every caller regardless of which layer issued it.
"""Factory: create a DataAccessOps by backend name.
Args:
backend_type: One of "postgresql" or "oracle".
Returns:
The shared DataAccessOps instance for that backend.
A DataAccessOps instance.
Raises:
ValueError: If backend_type is not recognized.
"""
ops = _OPS_CACHE.get(backend_type)
if ops is None:
ops = _get_ops_class(backend_type)()
_OPS_CACHE[backend_type] = ops
return ops
return _get_ops_class(backend_type)()
@@ -112,23 +112,6 @@ class DatabaseConnection(ABC):
"""
...
async def execute_rows_affected(self, query: str, *args: Any, timeout: float | None = None) -> int:
"""Execute a DML statement and return the number of rows it affected.
Normalizes the dialect-specific execute result into a plain int so callers
never hand-parse an ``"UPDATE <n>"`` / ``"DELETE <n>"`` command tag in
business logic (mirrors ``parse_json`` above, which normalizes the other
dialect-divergent result shape). asyncpg returns the tag directly; the
Oracle connection reshapes ``cursor.rowcount`` into the same trailing-count
form, so parsing the last token is dialect-safe. Returns 0 when the status
has no trailing count (e.g. a non-DML statement).
"""
status = await self.execute(query, *args, timeout=timeout)
if not isinstance(status, str):
return 0
parts = status.split()
return int(parts[-1]) if parts and parts[-1].isdigit() else 0
@abstractmethod
async def executemany(self, query: str, args: list[tuple[Any, ...]], *, timeout: float | None = None) -> None:
"""Execute a query for each set of arguments.
@@ -324,17 +307,6 @@ class DatabaseBackend(ABC):
"""Close the connection pool and release all resources."""
...
@property
@abstractmethod
def is_ready(self) -> bool:
"""Whether the pool exists and can serve connections.
False before :meth:`initialize` and after :meth:`shutdown`. Best-effort
callers (tracing, auditing) check this to skip work during those windows
instead of acquiring and interpreting the resulting error.
"""
...
@abstractmethod
@asynccontextmanager
async def acquire(self) -> AsyncIterator[DatabaseConnection]:
@@ -18,128 +18,12 @@ and mirrors Django's ``DatabaseOperations`` architecture.
from abc import ABC, abstractmethod
from dataclasses import dataclass
from datetime import datetime
from typing import Any
from .base import DatabaseConnection
from .result import ResultRow
def 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
useless: ``claim_graph_maintenance_batch`` locks queue rows ``FOR UPDATE``
*without* ``SKIP LOCKED`` (it is written assuming a single runner per bank),
so the runs convoy on each other's row locks while each holds a worker slot.
Same guarantee ``consolidation`` already gets from its ``bank_id != ALL(busy)``
exclusion, and the same caveat: a row wedged in 'processing' holds its bank
until something releases it (``hindsight-admin recover``, or a restart with a
stable ``HINDSIGHT_API_WORKER_ID`` so ``recover_own_tasks`` matches it). That
is a general gap in claim recovery, not specific to graph_maintenance.
Two differences from the consolidation form, both forced by the shape of this
problem:
* It is a **predicate**, not a separate claim phase. Pulling graph_maintenance
into its own phase after the generic shared-pool query would drop it below
every other operation type: it has no reserved-slot floor
(``WORKER_SLOT_TYPE_DEFAULTS`` gives consolidation 2 and graph_maintenance
0), and the poller's fairness pass calls ``claim_tasks`` with
``shared_limit=1``, so a single pending retain would starve it indefinitely.
As a predicate it keeps competing by ``created_at``.
* It also suppresses every same-bank row but the oldest **within one batch**.
Excluding busy banks alone does not: with several pending rows and nothing
yet processing, one batch claims them all the convoy, unchanged. Several
pending rows per bank are reachable through the recovery paths
(``_reclaim_own_processing_tasks`` resets *all* of a worker's processing
rows in one statement, from ``recover_own_tasks`` at startup and
``release_own_tasks`` at shutdown, plus ``_schedule_retry`` /
``_defer_operation`` / ``hindsight-admin recover``).
The candidate row is always 'pending' and the 'pending' branch is
strictly-older, so the subquery can never match the candidate itself. The
fragment carries no SQL comments on purpose it is rewritten for Oracle by
regex (``db/oracle.py``).
Args:
table: Fully-qualified async_operations table.
alias: Alias of the outer candidate row in the calling query.
"""
return f"""
({alias}.operation_type <> 'graph_maintenance' OR NOT EXISTS (
SELECT 1 FROM {table} gm_peer
WHERE gm_peer.bank_id = {alias}.bank_id
AND gm_peer.operation_type = 'graph_maintenance'
AND (
gm_peer.status = 'processing'
OR (gm_peer.status = 'pending'
AND gm_peer.task_payload IS NOT NULL
AND (gm_peer.next_retry_at IS NULL OR gm_peer.next_retry_at <= NOW())
AND (gm_peer.created_at < {alias}.created_at
OR (gm_peer.created_at = {alias}.created_at
AND gm_peer.operation_id < {alias}.operation_id)))
)
))
"""
@dataclass
class TagListingParts:
"""Backend-specific SQL fragments for the tag listing query."""
@@ -150,57 +34,6 @@ class TagListingParts:
bank_prefix: str
@dataclass(frozen=True)
class UpdatedWindow:
"""Recall's ``created_after``/``created_before`` bounds, as SQL for graph expansion.
Recall applies the window to ``updated_at`` a consolidation touch makes a
fact current again so link expansion has to bound the same column its seed
query does. Filtering only the seeds is not enough: a single in-window seed
would otherwise drag its whole neighbourhood (shared entities, semantic kNN
links, causal links) into the results no matter how old those neighbours are.
``first_param_index`` is where the bounds land in the owning query's param
list, so each call site keeps the placeholder numbering next to the params it
binds. Rendering is per-alias because the same window is applied to several
correlation names within one query.
"""
after: datetime | None
before: datetime | None
first_param_index: int
def clause(self, alias: str) -> str:
"""``AND <alias>.updated_at > $n ...`` — empty when the window is unbounded."""
parts: list[str] = []
index = self.first_param_index
if self.after is not None:
parts.append(f" AND {alias}.updated_at > ${index}")
index += 1
if self.before is not None:
parts.append(f" AND {alias}.updated_at < ${index}")
return "".join(parts)
@property
def params(self) -> list[datetime]:
"""The bound values, in placeholder order. Append to the owning param list."""
return [bound for bound in (self.after, self.before) if bound is not None]
@dataclass(frozen=True)
class LinkExpansionRows:
"""The three link-expansion signals, kept apart until they are scored.
They cannot be concatenated at the SQL layer: each carries a different score
scale (shared-entity count, kNN weight, causal weight) and the caller applies
a different transformation to each before summing them.
"""
entity: list[ResultRow]
semantic: list[ResultRow]
causal: list[ResultRow]
class DataAccessOps(ABC):
"""Backend-specific multi-statement data access operations.
@@ -316,14 +149,9 @@ class DataAccessOps(ABC):
bank_id: str,
entity_names: list[str],
entity_dates: list,
entity_kinds: list[str],
) -> dict[str, str]:
"""Bulk insert entities with ON CONFLICT DO NOTHING, returning id-by-lowercase-name.
``entity_kinds`` ("regular"/"label", parallel to ``entity_names``) is
stored on the row so label entities stay out of the partial trigram
index (#3208).
PG uses INSERT ... SELECT FROM unnest() with RETURNING.
Non-PG inserts row-by-row then SELECTs.
"""
@@ -344,26 +172,6 @@ class DataAccessOps(ABC):
"""
...
@abstractmethod
async def bulk_reassert_entities(
self,
conn: DatabaseConnection,
table: str,
bank_id: str,
entity_ids: list[str],
canonical_names: list[str],
entity_kinds: list[str],
) -> None:
"""Lock resolved parents and re-create any pruned since Phase-1 resolution.
Closes the retain Phase-1/prune race (#2662): existing rows are locked
(PG ``FOR KEY SHARE`` / Oracle ``FOR UPDATE``) so a concurrent
``prune_orphan_entities`` blocks until the caller's transaction commits,
while rows already deleted are re-inserted idempotently. ``entity_ids``
must be sorted by the caller for a stable lock order.
"""
...
@abstractmethod
async def bulk_insert_unit_entities(
self,
@@ -422,16 +230,12 @@ class DataAccessOps(ABC):
mu_table: str,
ue_table: str,
per_entity_limit: int,
window: UpdatedWindow,
) -> str:
"""Build entity expansion CTE for link expansion retrieval.
PG uses DISTINCT ON with CROSS JOIN LATERAL and GROUP BY.
Non-PG splits into entity_scores subquery then JOINs for full columns
(can't GROUP BY CLOB).
``window`` narrows candidates *before* the per-entity cap, so out-of-window
neighbours don't consume an entity's bounded fan-out.
"""
...
@@ -440,7 +244,6 @@ class DataAccessOps(ABC):
self,
ml_table: str,
mu_table: str,
window: UpdatedWindow,
) -> str:
"""Build semantic + causal expansion CTEs.
@@ -459,8 +262,7 @@ class DataAccessOps(ABC):
seed_ids: list,
budget: int,
per_entity_limit: int,
window: UpdatedWindow,
) -> LinkExpansionRows:
) -> tuple[list[ResultRow], list[ResultRow], list[ResultRow]]:
"""Observation-specific graph expansion.
PG uses native array ops (source_memory_ids column) for performance.
@@ -646,43 +448,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,
@@ -690,10 +455,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.
@@ -706,10 +470,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
@@ -719,23 +484,6 @@ class DataAccessOps(ABC):
# -- Task claiming operations ------------------------------------------
@abstractmethod
async def prune_terminal_operations(
self,
conn: DatabaseConnection,
table: str,
cutoff: datetime,
*,
batch_size: int,
) -> int:
"""Delete one deterministic batch of terminal operations older than ``cutoff``.
Implementations must lock candidates without waiting on rows another
worker is pruning, never select pending/processing rows, and return the
number deleted. The caller provides a transaction around this method.
"""
...
@abstractmethod
async def claim_tasks(
self,
@@ -753,12 +501,6 @@ class DataAccessOps(ABC):
Oracle implementation uses two-step claims (query busy banks first, then
claim excluding them) to avoid ORA-02014.
Implementations must apply :func:`graph_maintenance_bank_serialization_sql`
to every query that can return a ``graph_maintenance`` row, so at most one
such row per bank is ever in flight, and :func:`document_serialization_sql`
to every query that can return a ``retain`` row, so at most one retain per
document is ever in flight.
Args:
consolidation_bank_priority: Per-bank priority for consolidation scheduling.
Maps bank name patterns to integer priorities (higher = claimed first).
@@ -767,48 +509,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.
"""
...
@@ -10,18 +10,9 @@ import uuid as uuid_mod
from datetime import UTC, datetime
from .base import DatabaseConnection
from .ops import (
DataAccessOps,
LinkExpansionRows,
TagListingParts,
UpdatedWindow,
document_serialization_sql,
graph_maintenance_bank_serialization_sql,
)
from .ops import DataAccessOps, TagListingParts
from .result import DictResultRow as ResultRow
ORACLE_IN_LIST_LIMIT = 1000
class OracleOps(DataAccessOps):
"""Oracle-specific data access operations."""
@@ -181,24 +172,22 @@ class OracleOps(DataAccessOps):
bank_id: str,
entity_names: list[str],
entity_dates: list,
entity_kinds: list[str],
) -> dict[str, str]:
# Row-by-row insert with duplicate suppression.
# Can't use RETURNING with ON CONFLICT DO NOTHING reliably,
# so INSERT (ignoring dups) then SELECT all IDs at the end.
id_by_name: dict[str, str] = {}
for name, event_date, kind in zip(entity_names, entity_dates, entity_kinds):
for name, event_date in zip(entity_names, entity_dates):
ts = event_date if event_date else datetime.now(UTC)
await conn.execute(
f"""
INSERT INTO {table} (bank_id, canonical_name, first_seen, last_seen, mention_count, entity_kind)
VALUES ($1, $2, $3, $3, 0, $4)
INSERT INTO {table} (bank_id, canonical_name, first_seen, last_seen, mention_count)
VALUES ($1, $2, $3, $3, 0)
ON CONFLICT (bank_id, LOWER(canonical_name)) DO NOTHING
""",
bank_id,
name,
ts,
kind,
)
# Now SELECT all the entities we just inserted (or that already existed)
for name in entity_names:
@@ -227,7 +216,7 @@ class OracleOps(DataAccessOps):
for orig_name in missing_names:
row = await conn.fetchrow(
f"""
SELECT id, canonical_name, LOWER(canonical_name) AS name_lower
SELECT id, LOWER(canonical_name) AS name_lower
FROM {table}
WHERE bank_id = $1 AND LOWER(canonical_name) = LOWER($2)
""",
@@ -235,41 +224,10 @@ class OracleOps(DataAccessOps):
orig_name,
)
if row:
# Wrap in a dict-like to include input_name for downstream compat
results.append(row)
return results
async def bulk_reassert_entities(
self,
conn: DatabaseConnection,
table: str,
bank_id: str,
entity_ids: list[str],
canonical_names: list[str],
entity_kinds: list[str],
) -> None:
# Oracle has no FOR KEY SHARE; FOR UPDATE is the row-lock equivalent that
# blocks a concurrent prune DELETE until this transaction commits. Lock
# each surviving parent in the caller's stable id order (pruned ids are
# simply absent here), then re-insert any that vanished. The translation
# layer rewrites ON CONFLICT DO NOTHING to strip-and-catch ORA-00001, so
# a name recreated under a new id is suppressed rather than raising.
for entity_id in entity_ids:
await conn.fetchrow(
f"SELECT id FROM {table} WHERE id = $1 FOR UPDATE",
entity_id,
)
await conn.executemany(
f"""
INSERT INTO {table} (id, bank_id, canonical_name, entity_kind)
VALUES ($1, $2, $3, $4)
ON CONFLICT DO NOTHING
""",
[
(entity_id, bank_id, canonical_name, kind)
for entity_id, canonical_name, kind in zip(entity_ids, canonical_names, entity_kinds)
],
)
async def bulk_insert_unit_entities(
self,
conn: DatabaseConnection,
@@ -295,29 +253,22 @@ class OracleOps(DataAccessOps):
) -> None:
if not unit_ids:
return
# Locking upsert (#3034), the Oracle analogue of the PG
# ``ON CONFLICT DO UPDATE``. The old IGNORE_ROW_ON_DUPKEY_INDEX insert
# skipped duplicates WITHOUT locking the existing row, so a mutation
# re-enqueueing an already-queued unit could not block a worker from
# concurrently claiming (deleting) that row and processing the unit's
# pre-mutation state — the re-enqueue signal was silently lost. MERGE
# WHEN MATCHED takes an exclusive row lock on the existing queue row
# (the SET is a deliberate no-op that preserves enqueued_at); WHEN NOT
# MATCHED inserts a fresh row. That serialises the mutation against the
# worker's claim for the same (bank_id, unit_id).
# Oracle doesn't support ON CONFLICT; rely on the PK and the
# IGNORE_ROW_ON_DUPKEY_INDEX hint to skip duplicates server-side.
# The hint name must match the PK constraint exactly.
#
# Sort to enforce a global (bank_id, unit_id) lock-acquisition order,
# matching claim_graph_maintenance_batch's delete order, so overlapping
# mutation/worker sets acquire the shared row locks ascending and cannot
# cycle.
# Sort to enforce a global lock-acquisition order on the
# (bank_id, unit_id) PK. Without this, two concurrent
# transactions inserting overlapping unit_id sets in different
# orders can deadlock on the unique-check row locks. Sorting
# gives every concurrent caller the same lock order, so
# conflicting inserts queue cleanly instead of cycling.
sorted_unit_ids = sorted(unit_ids)
await conn.executemany(
f"""
MERGE INTO {table} q
USING (SELECT $1 AS bank_id, $2 AS unit_id FROM dual) s
ON (q.bank_id = s.bank_id AND q.unit_id = s.unit_id)
WHEN MATCHED THEN UPDATE SET q.enqueued_at = q.enqueued_at
WHEN NOT MATCHED THEN INSERT (bank_id, unit_id) VALUES (s.bank_id, s.unit_id)
INSERT /*+ IGNORE_ROW_ON_DUPKEY_INDEX({table}, pk_graph_maintenance_queue) */
INTO {table} (bank_id, unit_id)
VALUES ($1, $2)
""",
[(bank_id, uid) for uid in sorted_unit_ids],
)
@@ -342,15 +293,7 @@ class OracleOps(DataAccessOps):
bank_id,
limit,
)
# Ordered locking (#3034): the per-row DELETE takes the queue rows'
# exclusive locks in executemany array order. Sort the claimed keys by
# unit_id so those locks are acquired in the same (bank_id, unit_id)
# order the enqueue MERGE uses — overlapping mutation/worker sets then
# lock the shared rows ascending and cannot cycle. (The batch is still
# *chosen* oldest-first by enqueued_at above; only the lock/delete order
# is normalised.) The Pass 1 retry wrap in run_graph_maintenance_job is
# the ORA-00060 backstop for any residual interleaving.
claimed = sorted(str(row["unit_id"]) for row in rows)
claimed = [str(row["unit_id"]) for row in rows]
if claimed:
await conn.executemany(
f"DELETE FROM {table} WHERE bank_id = $1 AND unit_id = $2",
@@ -358,80 +301,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 +315,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 +326,20 @@ 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
@@ -569,7 +430,6 @@ class OracleOps(DataAccessOps):
mu_table: str,
ue_table: str,
per_entity_limit: int,
window: UpdatedWindow,
) -> str:
# Oracle: can't GROUP BY CLOB columns (text, context).
# Restructure: count entities per unit_id in a subquery, then join to get full columns.
@@ -587,16 +447,6 @@ class OracleOps(DataAccessOps):
FROM {ue_table} ue_target
WHERE ue_target.entity_id = se.entity_id
AND ue_target.unit_id != ALL($1::uuid[])
-- Filter before applying the cap: candidates from other fact
-- types, or outside the recall window, must not consume this
-- entity's bounded fan-out.
AND EXISTS (
SELECT 1
FROM {mu_table} mu_target
WHERE mu_target.id = ue_target.unit_id
AND mu_target.fact_type = $2
{window.clause("mu_target")}
)
ORDER BY ue_target.unit_id DESC
FETCH FIRST {per_entity_limit} ROWS ONLY
) t
@@ -609,6 +459,7 @@ class OracleOps(DataAccessOps):
es.score, 'entity' AS source
FROM entity_scores es
JOIN {mu_table} mu ON mu.id = es.unit_id
WHERE mu.fact_type = $2
ORDER BY es.score DESC
FETCH FIRST $3 ROWS ONLY
)"""
@@ -617,7 +468,6 @@ class OracleOps(DataAccessOps):
self,
ml_table: str,
mu_table: str,
window: UpdatedWindow,
) -> str:
# Non-PG: can't GROUP BY CLOB columns, no DISTINCT ON.
# Restructure semantic: compute max weight per id, then join for full columns.
@@ -632,7 +482,6 @@ class OracleOps(DataAccessOps):
AND ml.link_type = 'semantic'
AND mu.fact_type = $2
AND mu.id != ALL($1::uuid[])
{window.clause("mu")}
UNION ALL
SELECT mu.id, ml.weight
FROM {ml_table} ml
@@ -641,7 +490,6 @@ class OracleOps(DataAccessOps):
AND ml.link_type = 'semantic'
AND mu.fact_type = $2
AND mu.id != ALL($1::uuid[])
{window.clause("mu")}
) sem_raw
GROUP BY id
),
@@ -668,7 +516,6 @@ class OracleOps(DataAccessOps):
WHERE ml.from_unit_id = ANY($1::uuid[])
AND ml.link_type IN ('causes', 'caused_by', 'enables', 'prevents')
AND mu.fact_type = $2
{window.clause("mu")}
),
causal_expanded AS (
SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at,
@@ -687,8 +534,7 @@ class OracleOps(DataAccessOps):
seed_ids: list,
budget: int,
per_entity_limit: int,
window: UpdatedWindow,
) -> LinkExpansionRows:
) -> tuple[list[ResultRow], list[ResultRow], list[ResultRow]]:
import logging
logger = logging.getLogger(__name__)
@@ -742,13 +588,11 @@ class OracleOps(DataAccessOps):
WHERE os3.observation_id = mu.id
AND os3.source_id IN (SELECT source_id FROM connected_sources)
)
{window.clause("mu")}
ORDER BY score DESC
FETCH FIRST $2 ROWS ONLY
""",
seed_ids,
budget,
*window.params,
)
logger.debug(f"[LinkExpansion] observation graph (Oracle): found {len(entity_rows)} connected observations")
@@ -764,14 +608,12 @@ class OracleOps(DataAccessOps):
WHERE ml.from_unit_id = ANY($1::uuid[])
AND ml.link_type = 'semantic' AND mu.fact_type = 'observation'
AND mu.id != ALL($1::uuid[])
{window.clause("mu")}
UNION ALL
SELECT mu.id, ml.weight
FROM {ml_table} ml JOIN {mu_table} mu ON mu.id = ml.from_unit_id
WHERE ml.to_unit_id = ANY($1::uuid[])
AND ml.link_type = 'semantic' AND mu.fact_type = 'observation'
AND mu.id != ALL($1::uuid[])
{window.clause("mu")}
) sem_raw
GROUP BY id
),
@@ -797,7 +639,6 @@ class OracleOps(DataAccessOps):
WHERE ml.from_unit_id = ANY($1::uuid[])
AND ml.link_type IN ('causes', 'caused_by', 'enables', 'prevents')
AND mu.fact_type = 'observation'
{window.clause("mu")}
),
causal_expanded AS (
SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at,
@@ -812,12 +653,11 @@ class OracleOps(DataAccessOps):
""",
seed_ids,
budget,
*window.params,
)
semantic_rows = [r for r in sem_causal_rows if r["source"] == "semantic"]
causal_rows = [r for r in sem_causal_rows if r["source"] == "causal"]
return LinkExpansionRows(entity=list(entity_rows), semantic=semantic_rows, causal=causal_rows)
return list(entity_rows), semantic_rows, causal_rows
def build_tag_listing_parts(self, mu_table: str) -> TagListingParts:
return TagListingParts(
@@ -984,157 +824,6 @@ class OracleOps(DataAccessOps):
# -- Task claiming operations ------------------------------------------
async def prune_terminal_operations(
self,
conn: DatabaseConnection,
table: str,
cutoff: datetime,
*,
batch_size: int,
) -> int:
# Oracle rejects a row-limited SELECT ... FOR UPDATE (ORA-02014). Pick
# the deterministic bounded IDs first, then lock only that candidate
# set and re-check eligibility before deleting in the same transaction.
# Clamp to Oracle's 1000-expression IN-list limit because the adapter
# expands the candidate UUID list into individual bind variables.
# Cancelled children cannot complete parent aggregation, so retain the
# parent guard only for completed/failed children. Before removing a
# cancelled child, preserve its signal by cancelling a pending parent
# in this transaction and refreshing the parent's retention window.
# Validate metadata before HEXTORAW: CASE makes malformed UUIDs yield
# NULL while keeping the indexed RAW parent.operation_id key unwrapped.
effective_batch_size = min(batch_size, ORACLE_IN_LIST_LIMIT)
candidates = await conn.fetch(
f"""
SELECT candidate_operation.operation_id
FROM {table} candidate_operation
WHERE candidate_operation.status IN ('completed', 'failed', 'cancelled')
AND candidate_operation.updated_at < $1
AND (
candidate_operation.status = 'cancelled'
OR NOT EXISTS (
SELECT 1
FROM {table} parent
WHERE parent.operation_id = CASE
WHEN REGEXP_LIKE(
JSON_VALUE(
candidate_operation.result_metadata,
'$.parent_operation_id' RETURNING VARCHAR2(36) NULL ON ERROR
),
'^[0-9A-Fa-f]{{8}}-[0-9A-Fa-f]{{4}}-[0-9A-Fa-f]{{4}}-[0-9A-Fa-f]{{4}}-[0-9A-Fa-f]{{12}}$'
)
THEN HEXTORAW(REPLACE(
JSON_VALUE(
candidate_operation.result_metadata,
'$.parent_operation_id' RETURNING VARCHAR2(36) NULL ON ERROR
),
'-',
''
))
ELSE NULL
END
AND parent.bank_id = candidate_operation.bank_id
)
)
ORDER BY candidate_operation.updated_at, candidate_operation.operation_id
LIMIT $2
""",
cutoff,
effective_batch_size,
)
if not candidates:
return 0
candidate_ids = [row["operation_id"] for row in candidates]
locked = await conn.fetch(
f"""
SELECT candidate_operation.operation_id
FROM {table} candidate_operation
WHERE candidate_operation.operation_id = ANY($1)
AND candidate_operation.status IN ('completed', 'failed', 'cancelled')
AND candidate_operation.updated_at < $2
AND (
candidate_operation.status = 'cancelled'
OR NOT EXISTS (
SELECT 1
FROM {table} parent
WHERE parent.operation_id = CASE
WHEN REGEXP_LIKE(
JSON_VALUE(
candidate_operation.result_metadata,
'$.parent_operation_id' RETURNING VARCHAR2(36) NULL ON ERROR
),
'^[0-9A-Fa-f]{{8}}-[0-9A-Fa-f]{{4}}-[0-9A-Fa-f]{{4}}-[0-9A-Fa-f]{{4}}-[0-9A-Fa-f]{{12}}$'
)
THEN HEXTORAW(REPLACE(
JSON_VALUE(
candidate_operation.result_metadata,
'$.parent_operation_id' RETURNING VARCHAR2(36) NULL ON ERROR
),
'-',
''
))
ELSE NULL
END
AND parent.bank_id = candidate_operation.bank_id
)
)
ORDER BY candidate_operation.updated_at, candidate_operation.operation_id
FOR UPDATE OF candidate_operation.operation_id SKIP LOCKED
""",
candidate_ids,
cutoff,
)
if not locked:
return 0
operation_ids = [row["operation_id"] for row in locked]
await conn.execute(
f"""
UPDATE {table} parent
SET status = 'cancelled',
updated_at = now(),
completed_at = COALESCE(parent.completed_at, now()),
error_message = COALESCE(
parent.error_message,
'Cancelled because a child operation was cancelled'
)
WHERE parent.status = 'pending'
AND EXISTS (
SELECT 1
FROM {table} candidate_operation
WHERE candidate_operation.operation_id = ANY($1)
AND candidate_operation.status = 'cancelled'
AND candidate_operation.updated_at < $2
AND candidate_operation.bank_id = parent.bank_id
AND parent.operation_id = CASE
WHEN REGEXP_LIKE(
JSON_VALUE(
candidate_operation.result_metadata,
'$.parent_operation_id' RETURNING VARCHAR2(36) NULL ON ERROR
),
'^[0-9A-Fa-f]{{8}}-[0-9A-Fa-f]{{4}}-[0-9A-Fa-f]{{4}}-[0-9A-Fa-f]{{4}}-[0-9A-Fa-f]{{12}}$'
)
THEN HEXTORAW(REPLACE(
JSON_VALUE(
candidate_operation.result_metadata,
'$.parent_operation_id' RETURNING VARCHAR2(36) NULL ON ERROR
),
'-',
''
))
ELSE NULL
END
)
""",
operation_ids,
cutoff,
)
await conn.execute(
f"DELETE FROM {table} WHERE operation_id = ANY($1)",
operation_ids,
)
return len(operation_ids)
async def _claim_consolidation_tasks(
self,
conn,
@@ -1221,7 +910,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
@@ -1240,7 +929,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
@@ -1258,7 +947,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
@@ -1275,7 +964,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
@@ -1316,7 +1005,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
@@ -1363,7 +1052,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
@@ -1415,15 +1104,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
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
SELECT operation_id, operation_type, task_payload, retry_count
FROM {table}
WHERE status = 'pending'
AND task_payload IS NOT NULL
AND operation_type = $1
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
ORDER BY created_at
LIMIT $2
FOR UPDATE SKIP LOCKED
""",
@@ -1438,22 +1125,18 @@ class OracleOps(DataAccessOps):
# --- Phase 2: claim from shared pool ---
remaining_shared = shared_limit
if remaining_shared > 0:
# 2a. Non-consolidation tasks. graph_maintenance stays in this
# created_at-ordered query — see graph_maintenance_bank_serialization_sql
# for why it is a predicate rather than a phase of its own.
# 2a. Non-consolidation tasks
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
FROM {table} o
WHERE o.status = 'pending'
AND o.task_payload IS NOT NULL
AND o.operation_type != 'consolidation'
AND (o.next_retry_at IS NULL OR o.next_retry_at <= NOW())
AND o.operation_id != ALL($1::uuid[])
AND {graph_maintenance_bank_serialization_sql(table, "o")}
AND {document_serialization_sql(table, "o")}
ORDER BY o.created_at
SELECT operation_id, operation_type, task_payload, retry_count
FROM {table}
WHERE status = 'pending'
AND task_payload IS NOT NULL
AND operation_type != 'consolidation'
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
AND operation_id != ALL($1::uuid[])
ORDER BY created_at
LIMIT $2
FOR UPDATE SKIP LOCKED
""",
@@ -1463,15 +1146,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
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
SELECT operation_id, operation_type, task_payload, retry_count
FROM {table}
WHERE status = 'pending'
AND task_payload IS NOT NULL
AND operation_type != 'consolidation'
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
ORDER BY created_at
LIMIT $1
FOR UPDATE SKIP LOCKED
""",
@@ -1511,19 +1192,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}
@@ -1534,34 +1202,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
File diff suppressed because it is too large Load Diff
@@ -23,8 +23,6 @@ from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from typing import Any, NamedTuple
from .pool_instrumentation import PoolStats, acquire_conn
class _OracleJSONEncoder(json.JSONEncoder):
"""JSON encoder that handles datetime and UUID objects."""
@@ -80,9 +78,7 @@ _LIKE_ANY_RE = re.compile(r"(\w+)\s+LIKE\s+ANY\s*\(\s*:(\d+)\s*\)", re.IGNORECAS
_NOT_LIKE_ALL_RE = re.compile(r"(\w+)\s+NOT\s+LIKE\s+ALL\s*\(\s*:(\d+)\s*\)", re.IGNORECASE)
_JSON_ARROW_TEXT_RE = re.compile(r'("?\w+"?)\s*->>\s*\'(\w+)\'') # handles both col and "col"
# Reserved-word columns ("trigger") are already quoted by the time this runs, so the
# column group must accept the quoted form too — same shape as the arrow regex above.
_JSON_HAS_KEY_RE = re.compile(r"(\"?\w+\"?)\s*\?\s*'(\w+)'")
_JSON_HAS_KEY_RE = re.compile(r"(\w+)\s*\?\s*'(\w+)'")
_JSONB_CONTAINS_RE = re.compile(r"(\w+)\s*@>\s*:(\d+)")
# ---------------------------------------------------------------------------
@@ -150,7 +146,6 @@ _JSON_COL_NAMES = {
"config",
"observation_scopes",
"source_memory_ids",
"causal_links",
"trigger",
"http_config",
"event_types",
@@ -449,6 +444,9 @@ def _rewrite_pg_to_oracle(query: str) -> RewriteResult:
if has_for_update:
# FOR UPDATE path: use ROWNUM instead of FETCH FIRST.
# Extract and remove LIMIT clause, inject ROWNUM into WHERE.
def _limit_to_rownum(m):
return "" # Remove the LIMIT clause; we'll add ROWNUM below
limit_val = None
limit_match = re.search(r"\bLIMIT\s+(\d+|:\w+)\b", query, re.IGNORECASE)
if limit_match:
@@ -689,6 +687,7 @@ class OracleConnection(DatabaseConnection):
"max_tokens",
"priority",
"proof_count",
"access_count",
"importance_score",
"decay_factor",
"chunk_index",
@@ -1243,11 +1242,6 @@ class OracleBackend(DatabaseBackend):
def __init__(self) -> None:
self._pool: Any = None
self._oracledb: Any = None
# Oracle pooled sessions retain CURRENT_SCHEMA across checkouts. Cache
# SESSION_USER so default-schema acquisitions can explicitly reset a
# connection that was previously used for a tenant schema.
self._default_schema: str | None = None
self._acquire_warn_threshold_s: float = 1.0
async def initialize(
self,
@@ -1263,10 +1257,6 @@ class OracleBackend(DatabaseBackend):
oracledb = _import_oracledb()
self._oracledb = oracledb
from ...config import get_config
self._acquire_warn_threshold_s = get_config().db_acquire_warn_threshold_ms / 1000.0
# Parse URL-format DSN (oracle://user:pass@host:port/service)
from urllib.parse import urlparse
@@ -1287,17 +1277,11 @@ class OracleBackend(DatabaseBackend):
logger.info(f"Oracle pool created (min={min_size}, max={max_size})")
async def shutdown(self) -> None:
# Drop the reference before awaiting close() so is_ready flips False for
# the whole teardown, not just after it completes (see PostgreSQLBackend).
pool, self._pool = self._pool, None
if pool is not None:
await pool.close(force=True)
if self._pool is not None:
await self._pool.close(force=True)
self._pool = None
logger.info("Oracle pool closed")
@property
def is_ready(self) -> bool:
return self._pool is not None
async def _set_session_schema(self, conn: Any) -> None:
"""Set the session schema on an Oracle connection.
@@ -1310,41 +1294,15 @@ class OracleBackend(DatabaseBackend):
from ..memory_engine import get_current_schema
schema = get_current_schema()
cursor = conn.cursor()
try:
if self._default_schema is None:
await cursor.execute("SELECT SYS_CONTEXT('USERENV', 'SESSION_USER') FROM DUAL")
row = await cursor.fetchone()
if not row or not row[0]:
raise RuntimeError("Oracle did not return SESSION_USER while resetting CURRENT_SCHEMA")
self._default_schema = str(row[0])
target_schema = self._default_schema if not schema or schema == "public" else schema
safe_schema = target_schema.replace('"', '""')
await cursor.execute(f'ALTER SESSION SET CURRENT_SCHEMA = "{safe_schema}"')
finally:
# oracledb's AsyncCursor.close() is synchronous (not a coroutine);
# awaiting it raises "object NoneType can't be used in 'await'
# expression" and aborts every acquire().
cursor.close()
def _pool_stats(self) -> PoolStats | None:
"""Snapshot for slow-acquire logs, from oracledb pool attributes."""
pool = self._pool
if pool is None:
return None
try:
busy = pool.busy
return PoolStats(in_use=busy, max=pool.max, idle=pool.opened - busy)
except Exception:
return None
if schema and schema != "public":
cursor = conn.cursor()
await cursor.execute(f'ALTER SESSION SET CURRENT_SCHEMA = "{schema}"')
await cursor.close()
@asynccontextmanager
async def acquire(self) -> AsyncIterator[OracleConnection]:
pool = self._ensure_pool()
conn = await acquire_conn(
pool.acquire(), pool_stats=self._pool_stats, warn_threshold_s=self._acquire_warn_threshold_s
)
conn = await pool.acquire()
try:
await self._set_session_schema(conn)
yield OracleConnection(conn)
@@ -1360,9 +1318,7 @@ class OracleBackend(DatabaseBackend):
@asynccontextmanager
async def transaction(self) -> AsyncIterator[OracleConnection]:
pool = self._ensure_pool()
conn = await acquire_conn(
pool.acquire(), pool_stats=self._pool_stats, warn_threshold_s=self._acquire_warn_threshold_s
)
conn = await pool.acquire()
try:
await self._set_session_schema(conn)
yield OracleConnection(conn)
@@ -1,137 +0,0 @@
"""Instrumentation for database connection-pool acquisition.
asyncpg exposes pool *size* and *idle* counts, but not how many callers are
currently **queued waiting** for a connection and that queue depth is the
signal that actually distinguishes a saturated pool from a healthy one. When the
pool is exhausted, ``/health`` (which itself acquires a connection to run
``SELECT 1``) blocks in ``pool.acquire()`` until a connection frees or the acquire
times out, so a liveness probe can fail **with the event loop completely idle**.
This module tracks the process-wide count of in-flight acquisitions that have not
yet obtained a connection, and times each acquire so a slow one logs with full
pool stats. It is the DB-side counterpart to ``loop_watchdog`` (which covers loop
stalls); together, a stuck ``/health`` can be attributed to either a blocked loop
or pool exhaustion from the logs alone.
The counter is a plain int mutated only from the event-loop thread (asyncpg
acquisitions are awaited on the loop), so no lock is needed.
"""
from __future__ import annotations
import logging
import time
from collections.abc import AsyncIterator, Callable
from contextlib import asynccontextmanager
from dataclasses import dataclass
from typing import Any
logger = logging.getLogger("hindsight.db.pool")
_waiting = 0 # callers currently blocked in pool.acquire(), process-wide
@dataclass(frozen=True, slots=True)
class PoolStats:
"""Point-in-time connection-pool utilization snapshot."""
in_use: int
max: int
idle: int
def waiting_count() -> int:
"""Number of callers currently blocked waiting to acquire a pooled connection."""
return _waiting
@asynccontextmanager
async def instrument_acquire(
acquire_cm: Any,
*,
pool_stats: Callable[[], PoolStats | None] | None = None,
warn_threshold_s: float,
) -> AsyncIterator[Any]:
"""Wrap a pool's ``acquire()`` context manager with wait tracking + slow-acquire logging.
Args:
acquire_cm: an async context manager yielding a connection (e.g. the object
returned by ``asyncpg.Pool.acquire()``).
pool_stats: optional zero-arg callable returning a ``PoolStats`` snapshot for
the slow-acquire log line.
warn_threshold_s: log a warning when the acquire itself takes at least this long.
Yields:
The acquired connection.
"""
global _waiting
_waiting += 1
start = time.monotonic()
acquired = False
try:
async with acquire_cm as conn:
acquired = True
_waiting -= 1
_record_acquire_wait(time.monotonic() - start, pool_stats, warn_threshold_s)
yield conn
finally:
# If __aenter__ raised (acquire timeout / cancellation), we never
# decremented above — do it here so the waiter count can't leak.
if not acquired:
_waiting -= 1
async def acquire_conn(
acquire_awaitable: Any,
*,
pool_stats: Callable[[], PoolStats | None] | None = None,
warn_threshold_s: float,
) -> Any:
"""Await a pool acquire that returns a connection, with wait tracking + slow log.
For pools whose acquire is ``conn = await pool.acquire()`` (oracledb) rather than
an async context manager (asyncpg use ``instrument_acquire`` for those). The
caller is responsible for releasing the returned connection.
"""
global _waiting
_waiting += 1
start = time.monotonic()
try:
conn = await acquire_awaitable
finally:
_waiting -= 1
_record_acquire_wait(time.monotonic() - start, pool_stats, warn_threshold_s)
return conn
def _record_acquire_wait(
wait_s: float,
pool_stats: Callable[[], PoolStats | None] | None,
warn_threshold_s: float,
) -> None:
try:
from ...metrics import get_metrics_collector
get_metrics_collector().record_db_acquire_wait(wait_s)
except Exception:
pass
if wait_s < warn_threshold_s:
return
stats: PoolStats | None = None
if pool_stats is not None:
try:
stats = pool_stats()
except Exception:
stats = None
logger.warning(
"slow DB pool acquire: waited %.3fs for a connection "
"(in_use=%s max=%s idle=%s waiting=%s). The pool is likely saturated; "
"/health can stall on connection acquisition while the event loop is free.",
wait_s,
stats.in_use if stats else None,
stats.max if stats else None,
stats.idle if stats else None,
_waiting,
)
@@ -8,55 +8,17 @@ 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
from .base import DatabaseBackend, DatabaseConnection
from .pool_instrumentation import PoolStats, instrument_acquire
logger = logging.getLogger(__name__)
async def apply_session_settings(conn: asyncpg.Connection, settings: list[tuple[str, str]]) -> None:
"""Apply session-scoped GUCs to ``conn`` in a single round trip.
Unless ``HINDSIGHT_API_DB_SESSION_SETUP_ON_ACQUIRE=false``, the pool passes
its init callback as ``setup=`` too, so this runs on *every* acquire, not
just on connection creation. Issued as N separate ``SET`` statements that
was N round trips and behind a transaction-mode pooler N server-side
transactions per acquire, which the worker's per-schema acquires
multiplied into a sustained commit-rate burn (#3499). One
``SELECT set_config(...)`` collapses them into one statement.
Some of the settings are extension-provided (``hnsw.ef_search``,
``pg_trgm.similarity_threshold``) and may not exist on the cluster; a single
statement fails as a whole, so on error fall back to applying them one by
one, skipping only the ones the server rejects.
"""
if not settings:
return
args: list[str] = [value for pair in settings for value in pair]
projection = ", ".join(f"set_config(${2 * i + 1}, ${2 * i + 2}, false)" for i in range(len(settings)))
try:
await conn.execute(f"SELECT {projection}", *args)
return
except asyncpg.exceptions.PostgresError:
# Narrow to PostgresError so genuine bugs in the pool/conn layer surface
# instead of being silently retried statement-by-statement.
logger.debug("Batched session setup failed — applying settings individually")
for name, value in settings:
try:
await conn.execute("SELECT set_config($1, $2, false)", name, value)
except asyncpg.exceptions.PostgresError:
logger.debug("Could not set %s — the server may not know this setting", name)
class PostgresConnection(DatabaseConnection):
"""DatabaseConnection wrapper around an asyncpg.Connection."""
@@ -101,54 +63,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."""
@@ -162,8 +76,6 @@ class PostgreSQLBackend(DatabaseBackend):
def __init__(self) -> None:
self._pool: asyncpg.Pool | None = None
self._acquire_warn_threshold_s: float = 1.0
self._acquire_timeout_s: float | None = None
async def initialize(
self,
@@ -176,41 +88,6 @@ class PostgreSQLBackend(DatabaseBackend):
statement_cache_size: int = 0,
init_callback: Any | None = None,
) -> None:
from ...config import get_config
config = get_config()
self._acquire_warn_threshold_s = config.db_acquire_warn_threshold_ms / 1000.0
# Kept for acquire() below: asyncpg's ``timeout`` create_pool kwarg is a
# *connect* kwarg (how long establishing a new connection may take), and
# ``Pool.acquire()`` defaults to waiting for a free connection forever.
# Passing it here alone made HINDSIGHT_API_DB_ACQUIRE_TIMEOUT a no-op for
# the wait it names: a pool-exhaustion stall never surfaced as an error,
# it just hung (#3002). 0 restores the unbounded behaviour.
self._acquire_timeout_s = acquire_timeout if acquire_timeout > 0 else None
# 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,
@@ -218,55 +95,29 @@ class PostgreSQLBackend(DatabaseBackend):
command_timeout=command_timeout,
statement_cache_size=statement_cache_size,
timeout=acquire_timeout,
init=pool_init,
setup=pool_setup,
init=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:
# Drop the reference *before* awaiting close(): closing is not
# instantaneous, and anything acquiring during that window would
# otherwise get an asyncpg "pool is closing" error rather than seeing
# is_ready False.
pool, self._pool = self._pool, None
if pool is not None:
await pool.close()
if self._pool is not None:
await self._pool.close()
self._pool = None
logger.info("PostgreSQL pool closed")
@property
def is_ready(self) -> bool:
return self._pool is not None
def _pool_stats(self) -> PoolStats | None:
"""Snapshot for slow-acquire logs. in_use = live connections minus idle ones."""
pool = self._pool
if pool is None:
return None
idle = pool.get_idle_size()
return PoolStats(in_use=pool.get_size() - idle, max=pool.get_max_size(), idle=idle)
@asynccontextmanager
async def acquire(self) -> AsyncIterator[PostgresConnection]:
pool = self._ensure_pool()
async with instrument_acquire(
pool.acquire(timeout=self._acquire_timeout_s),
pool_stats=self._pool_stats,
warn_threshold_s=self._acquire_warn_threshold_s,
) as conn:
async with pool.acquire() as conn:
yield PostgresConnection(conn)
@asynccontextmanager
async def transaction(self) -> AsyncIterator[PostgresConnection]:
pool = self._ensure_pool()
async with instrument_acquire(
pool.acquire(timeout=self._acquire_timeout_s),
pool_stats=self._pool_stats,
warn_threshold_s=self._acquire_warn_threshold_s,
) as conn:
async with pool.acquire() as conn:
async with conn.transaction():
yield PostgresConnection(conn)
@@ -164,6 +164,35 @@ class BudgetedOperation:
"""
return BudgetedPool(pool, self)
async def acquire_many(
self,
pool: Any,
count: int,
) -> AsyncIterator[list[Any]]:
"""
Acquire multiple connections within the budget.
Note: This acquires connections sequentially to respect the budget.
For parallel acquisition, use multiple acquire() calls with asyncio.gather().
This method is intended for use with raw asyncpg pools only, not DatabaseBackend.
Args:
pool: asyncpg connection pool (raw pool only)
count: Number of connections to acquire
Yields:
List of database connections
"""
connections = []
try:
for _ in range(count):
conn = await pool.acquire()
connections.append(conn)
yield connections
finally:
for conn in connections:
await pool.release(conn)
# Global default manager instance
_default_manager: ConnectionBudgetManager | None = None
@@ -4,7 +4,6 @@ Database utility functions for connection management with retry logic.
import asyncio
import logging
import random
import time
from collections.abc import AsyncIterator
from contextlib import AsyncExitStack, asynccontextmanager
@@ -17,20 +16,6 @@ DEFAULT_MAX_RETRIES = 3
DEFAULT_BASE_DELAY = 0.5 # seconds
DEFAULT_MAX_DELAY = 5.0 # seconds
def _backoff_delay(attempt: int, base_delay: float, max_delay: float) -> float:
"""Exponential backoff with equal jitter.
Deterministic backoff makes concurrent retriers wake in lock-step and
re-collide on the very same rows, re-triggering the deadlock they just
backed off from. "Equal jitter" half the window fixed, half random
keeps a floor (so we don't hot-spin) while decorrelating the wake-ups, so
two contenders that deadlocked together are very unlikely to retry in sync.
"""
ceil = min(base_delay * (2**attempt), max_delay)
return ceil / 2 + random.uniform(0, ceil / 2)
# Retryable exception types (checked by class name to avoid hard imports)
_RETRYABLE_EXCEPTION_NAMES = frozenset(
{
@@ -93,7 +78,7 @@ async def retry_with_backoff(
raise
last_exception = e
if attempt < max_retries:
delay = _backoff_delay(attempt, base_delay, max_delay)
delay = min(base_delay * (2**attempt), max_delay)
if type(e).__name__ == "DeadlockDetectedError" or _is_oracle_deadlock(e):
logger.warning(
"Deadlock detected during parallel document processing — "
@@ -151,7 +136,7 @@ async def acquire_with_retry(backend_or_pool: Any, max_retries: int = DEFAULT_MA
if not _is_retryable(e):
raise
if attempt < max_retries:
delay = _backoff_delay(attempt, DEFAULT_BASE_DELAY, DEFAULT_MAX_DELAY)
delay = min(DEFAULT_BASE_DELAY * (2**attempt), DEFAULT_MAX_DELAY)
logger.warning(
f"Database acquire failed (attempt {attempt + 1}/{max_retries + 1}): {e}. "
f"Retrying in {delay:.1f}s..."
@@ -48,12 +48,6 @@ from ..config import (
ENV_LLM_API_KEY,
)
from .bank_attribution import apply_bank_attribution
from .local_device import (
release_local_inference_memory,
resolve_model_device_type,
select_local_device,
)
from .tei_retry import tei_retry_delay
logger = logging.getLogger(__name__)
@@ -125,32 +119,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)
@@ -162,13 +136,7 @@ class LocalSTEmbeddings(Embeddings):
The embedding dimension is auto-detected from the model.
"""
def __init__(
self,
model_name: str | None = None,
force_cpu: bool = False,
trust_remote_code: bool = False,
allow_mps: bool = False,
):
def __init__(self, model_name: str | None = None, force_cpu: bool = False, trust_remote_code: bool = False):
"""
Initialize local SentenceTransformers embeddings.
@@ -180,17 +148,12 @@ class LocalSTEmbeddings(Embeddings):
trust_remote_code: Allow loading models with custom code (security risk).
Required for some models with custom architectures.
Default: False (disabled for security)
allow_mps: Opt in to the Apple Silicon MPS GPU. Disabled by default
because MPS leaks memory under variable-length workloads
(see engine/local_device.py). Default: False
"""
self.model_name = model_name or DEFAULT_EMBEDDINGS_LOCAL_MODEL
self.force_cpu = force_cpu
self.trust_remote_code = trust_remote_code
self.allow_mps = allow_mps
self._model = None
self._dimension: int | None = None
self._device_type: str = "cpu"
@property
def provider_name(self) -> str:
@@ -217,11 +180,31 @@ class LocalSTEmbeddings(Embeddings):
logger.info(f"Embeddings: initializing local provider with model {self.model_name}")
# Determine device based on hardware availability. We always set
# low_cpu_mem_usage=False to prevent lazy loading (meta tensors) which can
# cause issues when accelerate is installed but no GPU is available.
# MPS is opt-in (allow_mps) — see engine/local_device.py for why.
device = select_local_device(self.force_cpu, self.allow_mps)
# Determine device based on hardware availability.
# We always set low_cpu_mem_usage=False to prevent lazy loading (meta tensors)
# which can cause issues when accelerate is installed but no GPU is available.
import torch
# Force CPU mode if configured (used in daemon mode to avoid MPS/XPC issues on macOS)
if self.force_cpu:
device = "cpu"
logger.info("Embeddings: forcing CPU mode")
else:
# Check for GPU (CUDA), Apple Silicon (MPS), or Intel XPU
# Wrap in try-except to gracefully handle any device detection issues
# (e.g., in CI environments or when PyTorch is built without GPU support)
device = "cpu" # Default to CPU
try:
has_gpu = torch.cuda.is_available() or (
hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
)
# Intel Arc XPU support — torch.xpu is available when the XPU build is loaded
if not has_gpu and hasattr(torch, "xpu"):
has_gpu = torch.xpu.is_available()
if has_gpu:
device = None # Let sentence-transformers auto-detect GPU/MPS/XPU
except Exception as e:
logger.warning(f"Failed to detect GPU/MPS/XPU, falling back to CPU: {e}")
# Suppress verbose transformers warnings during model loading
# This suppresses the "UNEXPECTED" warnings from BertModel which are harmless
@@ -248,8 +231,7 @@ class LocalSTEmbeddings(Embeddings):
transformers_logger.setLevel(original_level)
self._dimension = self._model.get_sentence_embedding_dimension()
self._device_type = resolve_model_device_type(self._model)
logger.info(f"Embeddings: local provider initialized (dim: {self._dimension}, device: {self._device_type})")
logger.info(f"Embeddings: local provider initialized (dim: {self._dimension})")
def encode(self, texts: list[str]) -> list[list[float]]:
"""
@@ -261,49 +243,11 @@ class LocalSTEmbeddings(Embeddings):
Returns:
List of embedding vectors
"""
return self._encode_local(texts)
def encode_query(self, texts: list[str]) -> list[list[float]]:
return self._encode_local(texts, input_type="query")
def encode_documents(self, texts: list[str]) -> list[list[float]]:
return self._encode_local(texts, input_type="document")
def _encode_local(
self, texts: list[str], input_type: Literal["query", "document"] | None = None
) -> list[list[float]]:
if self._model is None:
raise RuntimeError("Embeddings not initialized. Call initialize() first.")
try:
# Delegate to SentenceTransformers' own asymmetric entry points rather than
# prefixing here: they apply whatever prompts the model ships with (and route
# the task for models exposing a Router module), so asymmetric models such as
# Qwen3-Embedding get their configured query prompt without Hindsight carrying
# per-model prefix config the way the ONNX provider has to. Models that declare
# no prompts are unaffected — SentenceTransformers defaults them to empty
# strings and skips prompt handling entirely, so this is byte-identical to
# encode() for e.g. the default BAAI/bge-small-en-v1.5.
# encode_query/encode_document exist only in sentence-transformers >= 5.0,
# which is why local-ml pins that floor.
if input_type == "query":
encode = self._model.encode_query
elif input_type == "document":
encode = self._model.encode_document
else:
encode = self._model.encode
embeddings = encode(texts, convert_to_numpy=True, show_progress_bar=False)
return [emb.tolist() for emb in embeddings]
finally:
# Only reclaim the GPU allocator pool here, and only when actually on a
# GPU (opt-in MPS/CUDA/XPU). encode() runs in tight retain loops, so a
# gc.collect()/malloc_trim on every call is too costly on the CPU default
# — and unnecessary: refcounting frees the small transient buffers
# immediately and the allocator reuses them for the next batch. (The
# reranker keeps its per-batch heap trim for the #1717 CPU case; it runs
# on the lighter recall path.) See engine/local_device.py.
if self._device_type != "cpu":
release_local_inference_memory(self._device_type)
embeddings = self._model.encode(texts, convert_to_numpy=True, show_progress_bar=False)
return [emb.tolist() for emb in embeddings]
class OnnxEmbeddings(Embeddings):
@@ -416,6 +360,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 +433,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 +443,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
@@ -529,7 +478,7 @@ class RemoteTEIEmbeddings(Embeddings):
response = self._client.post(url, **kwargs)
response.raise_for_status()
return response
except (httpx.ConnectError, httpx.ConnectTimeout, httpx.ReadTimeout, httpx.WriteTimeout) as e:
except (httpx.ConnectError, httpx.ReadTimeout, httpx.WriteTimeout) as e:
last_error = e
if attempt < self.max_retries:
logger.warning(
@@ -538,20 +487,13 @@ class RemoteTEIEmbeddings(Embeddings):
time.sleep(delay)
delay *= 2 # Exponential backoff
except httpx.HTTPStatusError as e:
# TEI uses 429 as normal overload backpressure. Retry it with
# the same bounded budget as transient server errors.
if (e.response.status_code == 429 or e.response.status_code >= 500) and attempt < self.max_retries:
# Retry on 5xx server errors
if e.response.status_code >= 500 and attempt < self.max_retries:
last_error = e
sleep_delay = tei_retry_delay(
e.response,
delay,
request_timeout=self.timeout,
)
logger.warning(
f"TEI transient error (attempt {attempt + 1}/{self.max_retries + 1}): {e}. "
f"Retrying in {sleep_delay:.2f}s..."
f"TEI server error (attempt {attempt + 1}/{self.max_retries + 1}): {e}. Retrying in {delay}s..."
)
time.sleep(sleep_delay)
time.sleep(delay)
delay *= 2
else:
raise
@@ -653,8 +595,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 +606,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 +613,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 +731,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 +742,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 +1086,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 +1097,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 +1202,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 +1216,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 +1224,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,30 +1575,16 @@ 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,
force_cpu=config.embeddings_local_force_cpu,
trust_remote_code=config.embeddings_local_trust_remote_code,
allow_mps=config.embeddings_local_allow_mps,
)
elif provider == "onnx":
return OnnxEmbeddings(
@@ -1710,8 +1616,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 +1623,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 +1637,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 +1651,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 +1683,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 +1691,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
@@ -6,16 +6,13 @@ to disambiguate entities across memory units.
"""
import asyncio
import heapq
import json
import logging
import re
from collections import defaultdict
from collections.abc import Iterator
from dataclasses import dataclass, field
from datetime import UTC, datetime
from difflib import SequenceMatcher
from typing import Any, Final, cast
from typing import Any, Final
from .db_utils import acquire_with_retry
from .memory_engine import fq_table
@@ -28,7 +25,6 @@ from .retain.entity_labels import (
from .retain.entity_labels import (
parse_entity_labels as _parse_entity_labels,
)
from .retain.types import ResolvedEntity
logger = logging.getLogger(__name__)
@@ -40,119 +36,6 @@ class _EntityToCreate:
idx: int
name: str
event_date: datetime | None
# Label entities (from entity_labels config) are never fuzzy-merged in-batch — their
# canonical names are user-defined (e.g. "use:use-001") and must stay distinct (GH-1558).
# Also stored on the row as entities.entity_kind so label rows stay out of the
# partial trigram index (#3208).
is_label: bool = False
@dataclass
class _SimilarNamePair:
"""A pair of in-batch new-entity names judged similar enough to be the same entity."""
name_a: str
name_b: str
# The in-batch dedup pass is O(N^2) over the batch's *new* names. It is sub-millisecond for a
# normal retain (a handful of new entities) but scales quadratically — measured on the retain hot
# path: ~0.8ms at 100 names, ~5ms at 250, ~22ms at 500, ~81ms at 1000. So skip it past this many
# unique new names and log rather than silently degrade; the cap sits well above any realistic
# single-retain new-entity count while bounding the tail.
_INTRABATCH_MAX_NAMES = 250
# A pg_trgm "word" is a maximal run of alphanumerics (Unicode letters/digits, underscore excluded);
# everything else (space, punctuation, emoji) is a separator. This is why decoration variants like
# "Wren <emoji>" collapse to the same trigram set.
_TRGM_WORD = re.compile(r"[^\W_]+", re.UNICODE)
def _trigram_set(text: str) -> set[str]:
"""Trigrams of ``text`` the way PostgreSQL pg_trgm generates them: lowercase, split into words,
pad each word with two leading + one trailing blank, and take every 3-char window."""
trigrams: set[str] = set()
for word in _TRGM_WORD.findall(text.lower()):
padded = f" {word} "
for i in range(len(padded) - 2):
trigrams.add(padded[i : i + 3])
return trigrams
def _trigram_similarity(a: str, b: str) -> float:
"""pg_trgm ``similarity(a, b)`` computed in-memory — the Jaccard index of the trigram sets.
Verified byte-for-byte against Postgres pg_trgm across emoji / accent / CJK / hyphen /
apostrophe cases (issue #3107), so the merge cutoff calibrated on pg_trgm transfers exactly.
Doing it in Python keeps the in-batch dedup off the retain transaction's DB connection and makes
it backend-agnostic (Postgres, Oracle, and the pg_trgm-absent "full" fallback all behave alike).
"""
ta, tb = _trigram_set(a), _trigram_set(b)
intersection = len(ta & tb)
union = len(ta) + len(tb) - intersection
return intersection / union if union else 0.0
def _find_intrabatch_similar_pairs(names: list[str], threshold: float) -> list[_SimilarNamePair]:
"""Every pair of ``names`` whose in-memory trigram similarity meets ``threshold``. O(N^2) over a
small, capped set of new names pure CPU, no DB round-trip."""
trigrams = [_trigram_set(n) for n in names]
pairs: list[_SimilarNamePair] = []
for i in range(len(names)):
ti = trigrams[i]
for j in range(i + 1, len(names)):
tj = trigrams[j]
intersection = len(ti & tj)
union = len(ti) + len(tj) - intersection
if union and intersection / union >= threshold:
pairs.append(_SimilarNamePair(name_a=names[i], name_b=names[j]))
return pairs
def _cluster_new_entity_names(
rep_by_lower: dict[str, str],
count_by_lower: dict[str, int],
pairs: list[_SimilarNamePair],
) -> dict[str, str]:
"""Union-find the similar-name pairs into clusters and pick one canonical name each.
Args:
rep_by_lower: lowercase name -> a representative original-case spelling of it.
count_by_lower: lowercase name -> how many mentions carry it (for canonical choice).
pairs: name pairs judged similar (order/case irrelevant; compared lowercased).
Returns:
lowercase name -> canonical original-case name for its cluster. Singletons map to
themselves, so the caller can look up every member uniformly.
"""
parent: dict[str, str] = {nl: nl for nl in rep_by_lower}
def find(x: str) -> str:
while parent[x] != x:
parent[x] = parent[parent[x]] # path halving
x = parent[x]
return x
for pair in pairs:
a, b = pair.name_a.lower(), pair.name_b.lower()
if a in parent and b in parent:
ra, rb = find(a), find(b)
if ra != rb:
parent[ra] = rb
clusters: dict[str, list[str]] = {}
for nl in rep_by_lower:
clusters.setdefault(find(nl), []).append(nl)
canonical_by_member: dict[str, str] = {}
for members in clusters.values():
# Canonical = most-mentioned, then shortest, then lexicographically smallest — a
# deterministic pick that prefers the plainest spelling in the cluster.
canonical_lower = min(members, key=lambda nl: (-count_by_lower[nl], len(rep_by_lower[nl]), rep_by_lower[nl]))
canonical_name = rep_by_lower[canonical_lower]
for nl in members:
canonical_by_member[nl] = canonical_name
return canonical_by_member
@dataclass
@@ -192,22 +75,6 @@ def _later_date(a: datetime | None, b: datetime | None) -> datetime | None:
return a if a > b else b
def _canonical_cooccurrence_pairs(entity_list: list[str]) -> Iterator[tuple[str, str]]:
"""Yield each distinct pair of ``entity_list`` as ``(a, b)`` with ``a < b``.
Canonical ordering matches the entity_cooccurrences PK and check constraint.
The pair is ordered into fresh locals rather than by swapping the loop
variables: ``entity_id_1`` is the outer iterate, so swapping it would leak
into the remaining inner iterations and build later pairs off the wrong
element.
"""
for i, entity_id_1 in enumerate(entity_list):
for entity_id_2 in entity_list[i + 1 :]:
if entity_id_1 == entity_id_2:
continue
yield (entity_id_1, entity_id_2) if entity_id_1 < entity_id_2 else (entity_id_2, entity_id_1)
@dataclass
class _CooccurrencePair:
"""A (entity_id_1, entity_id_2) pair observed in a retain batch (for post-txn flush)."""
@@ -225,32 +92,6 @@ class _CooccurrencePair:
_nlp = None
# Candidates scored between cooperative yields to the event loop. Scoring is
# synchronous CPU (one SequenceMatcher per candidate, ~50µs), so a batch with a
# large candidate set would otherwise hold the loop thread for minutes — health
# probes time out and the orchestrator kills the worker mid-op (GH-3211).
# 256 candidates ≈ 13ms of work between yields.
_SCORING_YIELD_EVERY: Final = 256
def _cheap_rank_key(entity_text_lower: str, candidate: tuple[Any, str, Any, datetime | None, int | None]) -> tuple:
"""Ordering key (not a multi-value return) approximating match quality cheaply.
Used only to truncate oversized candidate sets: the fuzzy strategies already
cap and pre-rank in SQL by real similarity, so this is the backstop for sets
built without a score (the "full" strategy's substring matching). Ranks an
exact match first, then a close name length, then a well-established entity
all O(1) per candidate, unlike the SequenceMatcher pass it protects.
"""
name_lower = candidate[1].lower()
return (
0 if name_lower == entity_text_lower else 1,
abs(len(name_lower) - len(entity_text_lower)),
-(candidate[4] or 0),
candidate[1],
)
class EntityResolver:
"""
Resolves entities to canonical IDs with disambiguation.
@@ -261,8 +102,6 @@ class EntityResolver:
pool: Any,
entity_lookup: str = "full",
entity_resolution_batch_size: int = 100,
intrabatch_merge_similarity: float = 0.5,
entity_resolution_max_candidates: int = 200,
):
"""
Initialize entity resolver.
@@ -274,22 +113,12 @@ class EntityResolver:
similar candidates per entity name (much faster for large banks).
entity_resolution_batch_size: Number of unique entity names to include
in each pg_trgm candidate lookup query.
intrabatch_merge_similarity: pg_trgm similarity at/above which two new
names created by the same retain are merged into one entity.
entity_resolution_max_candidates: Max candidates scored per entity
mention. Scoring is a synchronous SequenceMatcher call per
candidate, so an unbounded candidate set turns one resolution
batch into minutes of event-loop-blocking CPU (GH-3211).
"""
self.pool = pool
self.entity_lookup = entity_lookup
if entity_resolution_batch_size < 1:
raise ValueError("entity_resolution_batch_size must be >= 1")
self.entity_resolution_batch_size = entity_resolution_batch_size
self._intrabatch_merge_similarity = intrabatch_merge_similarity
if entity_resolution_max_candidates < 1:
raise ValueError("entity_resolution_max_candidates must be >= 1")
self.entity_resolution_max_candidates = entity_resolution_max_candidates
self._pg_trgm_checked = False
# Backend-specific operations — accessed via pool.ops (Django pattern).
self._ops = pool.ops if pool is not None else None
@@ -393,19 +222,6 @@ class EntityResolver:
"""Split values into fixed-size batches."""
return [values[i : i + size] for i in range(0, len(values), size)]
@staticmethod
def _label_texts(entity_texts: list[str], taxonomy_lookup: set[str] | None, labels_cfg) -> set[str]:
"""Subset of entity_texts that are label entities (resolved by exact match only).
Only gate on the config, not on the lookup set: text/map groups have no
fixed vocabulary, so a config with only those groups builds an EMPTY
lookup its labels are classified by key prefix inside
``is_label_entity``, and gating on the set would miss them entirely.
"""
if not labels_cfg:
return set()
return {t for t in entity_texts if _is_label_entity(t, labels_cfg, taxonomy_lookup or set())}
async def resolve_entities_batch(
self,
bank_id: str,
@@ -414,7 +230,7 @@ class EntityResolver:
unit_event_date,
conn=None,
entity_labels: list | None = None,
) -> list[ResolvedEntity]:
) -> list[str]:
"""
Resolve multiple entities in batch (MUCH faster than sequential).
@@ -429,8 +245,7 @@ class EntityResolver:
conn: Optional connection to use (if None, acquires from pool)
Returns:
Resolved entity identities (id + stored canonical name) in the same
order as input.
List of entity IDs in same order as input
"""
if not entities_data:
return []
@@ -456,7 +271,7 @@ class EntityResolver:
unit_event_date,
taxonomy_lookup: set[str] | None = None,
labels_cfg=None,
) -> list[ResolvedEntity]:
) -> list[str]:
if self.entity_lookup == "trigram":
# Route to backend-specific fuzzy strategy.
# Non-PG backends (Oracle) use UTL_MATCH instead of pg_trgm.
@@ -496,7 +311,7 @@ class EntityResolver:
unit_event_date,
taxonomy_lookup: set[str] | None = None,
labels_cfg=None,
) -> list[ResolvedEntity]:
) -> list[str]:
"""Original strategy: load all bank entities then match in Python."""
# Query ALL candidates for this bank
all_entities = await conn.fetch(
@@ -580,7 +395,7 @@ class EntityResolver:
unit_event_date,
taxonomy_lookup: set[str] | None = None,
labels_cfg=None,
) -> list[ResolvedEntity]:
) -> list[str]:
"""
Trigram strategy: fetch only similar candidates per entity name using pg_trgm.
@@ -590,78 +405,40 @@ class EntityResolver:
"""
entity_texts = list(set(e["text"] for e in entities_data))
# Label entities resolve by exact match only (their canonical names are
# user-defined and must not be fuzzy-merged). Probing them via the trigram
# index only returns similar-but-distinct label values that are always
# discarded, and that wasted work grows with the number of values a label
# accumulates. Resolve label texts with an exact lookup on the unique
# (bank_id, LOWER(canonical_name)) index and only fuzzy-match the rest.
label_set = self._label_texts(entity_texts, taxonomy_lookup, labels_cfg)
label_texts = [t for t in entity_texts if t in label_set]
fuzzy_texts = [t for t in entity_texts if t not in label_set]
rows = []
# Exact, index-only lookup for label texts.
for entity_text_batch in self._chunked(label_texts, self.entity_resolution_batch_size):
rows.extend(
await conn.fetch(
f"""
SELECT e.id, e.canonical_name, e.metadata, e.last_seen, e.mention_count,
q.query_text
FROM unnest($2::text[]) AS q(query_text)
JOIN {fq_table("entities")} e ON (
e.bank_id = $1
AND LOWER(e.canonical_name) = LOWER(q.query_text)
)
""",
bank_id,
entity_text_batch,
)
)
# Fetch candidates for the remaining texts in bounded batches.
# Fetch candidates for unique entity texts in bounded batches.
# Uses the GIN trigram index on LOWER(canonical_name) for case-insensitive
# similarity lookup. Previous version also had LIKE '%...' substring fallbacks,
# but those forced full sequential scans of the entities table and caused
# TimeoutErrors on banks with 10k+ entities. The pg_trgm similarity threshold
# that governs the `%` operator is applied once at pool-connection setup
# (HINDSIGHT_API_ENTITY_TRGM_SIMILARITY_THRESHOLD), so it is not toggled here.
# ``entity_kind != 'label'`` matches the predicate of the partial trigram
# index (label rows are exact-match-only, so they can never be a
# legitimate fuzzy result — without the filter they only inflate the
# candidate set and get discarded in the bitmap recheck, #3208). The
# clause must textually match the index predicate for the planner to
# choose the partial index, so it stays inside the LATERAL's WHERE
# alongside the `%` operator rather than moving out to the outer join.
#
# The LATERAL keeps only the best `max_candidates` per query text: on a bank
# with many near-identical names a single probe can otherwise return
# thousands of rows, and every one of them costs a SequenceMatcher call in
# _resolve_from_candidates (GH-3211). Ranking by pg_trgm similarity — which
# the index scan computes anyway — keeps the truncation at the noise end.
for entity_text_batch in self._chunked(fuzzy_texts, self.entity_resolution_batch_size):
rows.extend(
await conn.fetch(
f"""
SELECT c.id, c.canonical_name, c.metadata, c.last_seen, c.mention_count,
q.query_text
FROM unnest($2::text[]) AS q(query_text)
CROSS JOIN LATERAL (
SELECT e.id, e.canonical_name, e.metadata, e.last_seen, e.mention_count
FROM {fq_table("entities")} e
WHERE e.bank_id = $1
AND e.entity_kind != 'label'
AND LOWER(e.canonical_name) % LOWER(q.query_text)
ORDER BY similarity(LOWER(e.canonical_name), LOWER(q.query_text)) DESC, e.id
LIMIT $3
) c
""",
bank_id,
entity_text_batch,
self.entity_resolution_max_candidates,
# TimeoutErrors on banks with 10k+ entities. Lowering the similarity threshold
# to 0.15 (from default 0.3) catches most substring relationships while
# staying fully index-based.
await conn.execute("SET pg_trgm.similarity_threshold = 0.15")
try:
rows = []
for entity_text_batch in self._chunked(entity_texts, self.entity_resolution_batch_size):
rows.extend(
await conn.fetch(
f"""
SELECT DISTINCT ON (e.id)
e.id, e.canonical_name, e.metadata, e.last_seen, e.mention_count,
q.query_text
FROM unnest($2::text[]) AS q(query_text)
JOIN {fq_table("entities")} e ON (
e.bank_id = $1
AND LOWER(e.canonical_name) % LOWER(q.query_text)
)
""",
bank_id,
entity_text_batch,
)
)
)
finally:
# asyncpg returns connections to the pool with session state intact,
# so the lowered threshold would leak to future borrowers without RESET.
try:
await conn.execute("RESET pg_trgm.similarity_threshold")
except Exception:
logger.warning("Failed to reset pg_trgm similarity threshold after candidate lookup", exc_info=True)
# Group candidates by query_text
all_candidates: dict[str, list] = {t: [] for t in entity_texts}
@@ -722,7 +499,7 @@ class EntityResolver:
unit_event_date: datetime | None,
taxonomy_lookup: set[str] | None = None,
labels_cfg=None,
) -> list[ResolvedEntity]:
) -> list[str]:
"""
Oracle strategy: fetch similar candidates using UTL_MATCH.JARO_WINKLER_SIMILARITY.
@@ -734,14 +511,6 @@ class EntityResolver:
entity_texts = list(set(e["text"] for e in entities_data))
entities_table = fq_table("entities")
# Label entities resolve by exact match only, so the fuzzy Jaro-Winkler
# join only returns similar-but-distinct label values that are always
# discarded. Resolve label texts with an exact lookup on the unique
# (bank_id, LOWER(canonical_name)) index and only fuzzy-match the rest.
label_set = self._label_texts(entity_texts, taxonomy_lookup, labels_cfg)
label_texts = [t for t in entity_texts if t in label_set]
fuzzy_texts = [t for t in entity_texts if t not in label_set]
try:
# Batch entity texts into bounded sub-queries using JSON_TABLE to
# expand the list into rows. UTL_MATCH.JARO_WINKLER_SIMILARITY
@@ -749,7 +518,7 @@ class EntityResolver:
# Bounded batches mirror the PG trigram path so very wide retain
# batches don't time out a single JOIN on large banks.
rows = []
for entity_text_batch in self._chunked(label_texts, self.entity_resolution_batch_size):
for entity_text_batch in self._chunked(entity_texts, self.entity_resolution_batch_size):
rows.extend(
await conn.fetch(
f"""
@@ -758,46 +527,13 @@ class EntityResolver:
FROM JSON_TABLE($2, '$[*]' COLUMNS (query_text VARCHAR2(4000) PATH '$')) q
JOIN {entities_table} e ON (
e.bank_id = $1
AND LOWER(e.canonical_name) = LOWER(q.query_text)
AND UTL_MATCH.JARO_WINKLER_SIMILARITY(LOWER(e.canonical_name), LOWER(q.query_text)) > 70
)
""",
bank_id,
json.dumps(entity_text_batch),
)
)
# Only the best `max_candidates` per query text are returned: each
# candidate costs a synchronous SequenceMatcher call downstream, so an
# unbounded fuzzy match set blocks the event loop for minutes
# (GH-3211). Ranking by the same Jaro-Winkler score the join already
# computes keeps the truncation at the noise end.
for entity_text_batch in self._chunked(fuzzy_texts, self.entity_resolution_batch_size):
rows.extend(
await conn.fetch(
f"""
SELECT id, canonical_name, metadata, last_seen, mention_count, query_text
FROM (
SELECT e.id, e.canonical_name, e.metadata, e.last_seen, e.mention_count,
q.query_text,
ROW_NUMBER() OVER (
PARTITION BY q.query_text
ORDER BY UTL_MATCH.JARO_WINKLER_SIMILARITY(
LOWER(e.canonical_name), LOWER(q.query_text)
) DESC, e.id
) AS rn
FROM JSON_TABLE($2, '$[*]' COLUMNS (query_text VARCHAR2(4000) PATH '$')) q
JOIN {entities_table} e ON (
e.bank_id = $1
AND e.entity_kind != 'label'
AND UTL_MATCH.JARO_WINKLER_SIMILARITY(LOWER(e.canonical_name), LOWER(q.query_text)) > 70
)
)
WHERE rn <= $3
""",
bank_id,
json.dumps(entity_text_batch),
self.entity_resolution_max_candidates,
)
)
except Exception as e:
# UTL_MATCH may not be available (ORA-06550, ORA-00904, etc.)
# Catch broadly because Oracle error types vary depending on driver.
@@ -861,38 +597,6 @@ class EntityResolver:
labels_cfg,
)
def _intrabatch_canonical_map(self, entities_to_create: list[_EntityToCreate]) -> dict[str, str]:
"""Map each non-label new name (lowercased) to its cluster's canonical spelling.
Uses in-memory trigram similarity (``_trigram_similarity``, verified equal to Postgres
pg_trgm), so it is backend-agnostic no DB round-trip on the retain hot path, and it runs
identically on PostgreSQL, Oracle, and the pg_trgm-absent "full" fallback. Label entities
are excluded so distinct label values stay separate (GH-1558).
"""
rep_by_lower: dict[str, str] = {}
count_by_lower: dict[str, int] = {}
for e in entities_to_create:
if e.is_label:
continue
name_lower = e.name.lower()
rep_by_lower.setdefault(name_lower, e.name)
count_by_lower[name_lower] = count_by_lower.get(name_lower, 0) + 1
if len(rep_by_lower) < 2:
return {} # nothing to compare
if len(rep_by_lower) > _INTRABATCH_MAX_NAMES:
logger.warning(
"Skipping in-batch entity dedup: %d unique new names exceeds the %d cap "
"(O(N^2) trigram comparison); same-batch surface variants may not be merged.",
len(rep_by_lower),
_INTRABATCH_MAX_NAMES,
)
return {}
pairs = _find_intrabatch_similar_pairs(list(rep_by_lower.values()), self._intrabatch_merge_similarity)
if not pairs:
return {}
return _cluster_new_entity_names(rep_by_lower, count_by_lower, pairs)
async def _resolve_from_candidates(
self,
conn,
@@ -903,19 +607,13 @@ class EntityResolver:
cooccurrence_map: dict[str, set[str]],
taxonomy_lookup: set[str] | None = None,
labels_cfg=None,
) -> list[ResolvedEntity]:
) -> list[str]:
"""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
# inconsistency); it surfaces as a clear error at the reassert boundary
# rather than a silent NOT NULL violation deeper in Phase 2.
resolved: list[ResolvedEntity | None] = [None] * len(entities_data)
# Resolve each entity using pre-fetched candidates
entity_ids = [None] * len(entities_data)
entities_to_update: list[_EntityStat] = []
entities_to_create: list[_EntityToCreate] = []
# Candidates scored since the last yield, counted across mentions so a
# batch of many small candidate sets yields as often as one large set.
scored_since_yield = 0
for idx, entity_data in enumerate(entities_data):
entity_text = entity_data["text"]
@@ -925,85 +623,41 @@ class EntityResolver:
candidates = all_candidates.get(entity_text, [])
# Backstop truncation for candidate sets that were not capped at the
# source (the "full" strategy matches substrings in Python). The fuzzy
# strategies already return at most this many rows per query text, so
# this is normally a no-op.
if len(candidates) > self.entity_resolution_max_candidates:
logger.debug(
"Truncating %d candidates to %d for entity text %r",
len(candidates),
self.entity_resolution_max_candidates,
entity_text,
)
entity_text_lower_for_rank = entity_text.lower()
candidates = heapq.nsmallest(
self.entity_resolution_max_candidates,
candidates,
key=lambda c: _cheap_rank_key(entity_text_lower_for_rank, c),
)
# Label entities (from entity_labels config) use exact matching only.
# Their canonical names are user-defined (e.g., "use:use-001"),
# so fuzzy resolution must NOT merge distinct label values that
# happen to be textually similar (GH-1558). Don't gate on the
# lookup set — it is empty for text/map-only configs, whose labels
# classify by key prefix (see _label_texts).
is_label = bool(labels_cfg and _is_label_entity(entity_text, labels_cfg, taxonomy_lookup or set()))
# happen to be textually similar (GH-1558).
is_label = bool(
labels_cfg and taxonomy_lookup and _is_label_entity(entity_text, labels_cfg, taxonomy_lookup)
)
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)
)
entities_to_create.append(_EntityToCreate(idx=idx, name=entity_text, event_date=entity_event_date))
continue
if is_label:
# Exact case-insensitive match only for label entities
exact_match: ResolvedEntity | None = None
exact_match = None
entity_text_lower = entity_text.lower()
for candidate_id, canonical_name, metadata, last_seen, mention_count in candidates:
if canonical_name.lower() == entity_text_lower:
exact_match = ResolvedEntity(
entity_id=candidate_id, canonical_name=canonical_name, entity_kind="label"
)
exact_match = candidate_id
break
if exact_match:
resolved[idx] = exact_match
entities_to_update.append(
_EntityStat(entity_id=exact_match.entity_id, event_date=entity_event_date)
)
entity_ids[idx] = exact_match
entities_to_update.append(_EntityStat(entity_id=exact_match, event_date=entity_event_date))
else:
entities_to_create.append(
_EntityToCreate(idx=idx, name=entity_text, event_date=entity_event_date, is_label=True)
)
entities_to_create.append(_EntityToCreate(idx=idx, name=entity_text, event_date=entity_event_date))
continue
# Score candidates
best_candidate: ResolvedEntity | None = None
best_candidate = None
best_score = 0.0
nearby_entity_set = {e["text"].lower() for e in nearby_entities if e["text"] != entity_text}
for candidate_id, canonical_name, metadata, last_seen, mention_count in candidates:
# Hand the loop back periodically so /health (and every other task
# on this worker) still gets scheduled while a wide batch scores.
# Counted before the label skip below, so a candidate list that is
# entirely labels still yields — the skip runs _is_label_entity per
# row, which is cheap but not free.
scored_since_yield += 1
if scored_since_yield >= _SCORING_YIELD_EVERY:
scored_since_yield = 0
await asyncio.sleep(0)
# A label row can never be a fuzzy-match target (#1558): the
# trigram/UTL_MATCH probes exclude them in SQL via entity_kind,
# but the "full" fallback strategy loads every bank entity, so
# a textually-close label value could still outscore the 0.6
# threshold here (e.g. "topic empathy" vs "topic:empathy").
if labels_cfg and _is_label_entity(canonical_name, labels_cfg, taxonomy_lookup or set()):
continue
score = 0.0
# 1. Name similarity (0-0.5)
@@ -1031,17 +685,17 @@ class EntityResolver:
if score > best_score:
best_score = score
best_candidate = ResolvedEntity(entity_id=candidate_id, canonical_name=canonical_name)
best_candidate = candidate_id
# Apply unified threshold
threshold = 0.6
if best_score > threshold and best_candidate is not None:
resolved[idx] = best_candidate
entities_to_update.append(_EntityStat(entity_id=best_candidate.entity_id, event_date=entity_event_date))
if best_score > threshold:
entity_ids[idx] = best_candidate
entities_to_update.append(_EntityStat(entity_id=best_candidate, event_date=entity_event_date))
else:
entities_to_create.append(
_EntityToCreate(idx=idx, name=entity_data["text"], event_date=entity_event_date, is_label=is_label)
_EntityToCreate(idx=idx, name=entity_data["text"], event_date=entity_event_date)
)
# Existing entities: IDs already known from the candidate SELECT above.
@@ -1053,46 +707,24 @@ class EntityResolver:
# ON CONFLICT DO NOTHING returns nothing for rows that conflicted; we handle
# that rare case with a fallback SELECT.
if entities_to_create:
# Fuzzy-cluster the NON-label names about to be created so same-batch surface
# variants (case/emoji/suffix/typo of one name) collapse to a single entity. Without
# this, resolution only compares against already-persisted rows, so the first sighting
# of each variant in a batch always creates a distinct entity (issue #3107). Labels are
# excluded and keep exact grouping.
canonical_by_member = self._intrabatch_canonical_map(entities_to_create)
# Group by lowercase name — deduplicate within the batch.
@dataclass
class _NameGroup:
name: str
event_date: datetime | None
is_label: bool
indices: list[int] = field(default_factory=list)
groups: dict[str, _NameGroup] = {}
for e in entities_to_create:
# Non-label variants fold into their cluster's canonical name; everything else
# (labels, singletons) keys on itself, preserving the prior exact-match behavior.
canonical = canonical_by_member.get(e.name.lower(), e.name)
key = canonical.lower()
group = groups.get(key)
if group is None:
# Labels key on themselves and the dedup pass only clusters
# non-label names, so the first member's is_label holds for
# every member of the group.
group = _NameGroup(name=canonical, event_date=e.event_date, is_label=e.is_label)
groups[key] = group
elif e.event_date is not None and (group.event_date is None or e.event_date < group.event_date):
# Keep the earliest event_date across the cluster ("first seen").
group.event_date = e.event_date
group.indices.append(e.idx)
name_lower = e.name.lower()
if name_lower not in groups:
groups[name_lower] = _NameGroup(name=e.name, event_date=e.event_date)
groups[name_lower].indices.append(e.idx)
# Sort by lowercase name for deterministic ordering.
sorted_groups = sorted(groups.items())
entity_names = [g.name for _, g in sorted_groups]
entity_dates = [g.event_date for _, g in sorted_groups]
entity_kinds = ["label" if g.is_label else "regular" for _, g in sorted_groups]
# Stored canonical name per lowercase key, so a resurrected parent
# keeps the name it was created/matched with rather than a fallback.
canonical_by_name = {name_lower: g.name for name_lower, g in sorted_groups}
# INSERT ... ON CONFLICT DO NOTHING — no row lock on already-existing entities.
# mention_count starts at 0 here; flush_pending_stats() is the sole source of
@@ -1105,7 +737,6 @@ class EntityResolver:
bank_id,
entity_names,
entity_dates,
entity_kinds,
)
# Fallback SELECT for names that conflicted (another worker won the race).
@@ -1128,14 +759,11 @@ class EntityResolver:
)
for row in existing_rows:
id_by_name[row["name_lower"]] = row["id"]
canonical_by_name[row["name_lower"]] = row["canonical_name"]
# Also index by Python's lower() of the original input name so the
# assignment loop (which uses Python-lowercased keys) finds it even
# when Python and the database produce different lowercase strings.
if "input_name" in row:
input_name_lower = row["input_name"].lower()
id_by_name[input_name_lower] = row["id"]
canonical_by_name[input_name_lower] = row["canonical_name"]
id_by_name[row["input_name"].lower()] = row["id"]
# Assign entity IDs back and queue one stat per original mention so that
# flush_pending_stats() increments mention_count by the true mention count,
@@ -1143,74 +771,21 @@ class EntityResolver:
for name_lower, g in sorted_groups:
entity_id = id_by_name.get(name_lower)
if entity_id:
canonical_name = canonical_by_name.get(name_lower, g.name)
kind = "label" if g.is_label else "regular"
for original_idx in g.indices:
resolved[original_idx] = ResolvedEntity(
entity_id=entity_id, canonical_name=canonical_name, entity_kind=kind
)
pending.append(_EntityStat(entity_id=str(entity_id), event_date=g.event_date))
entity_ids[original_idx] = entity_id
pending.append(_EntityStat(entity_id=entity_id, event_date=g.event_date))
# Accumulate into the resolver's pending list; the orchestrator flushes
# these with await entity_resolver.flush_pending_stats() after the txn.
key = self._task_key()
self._pending_stats.setdefault(key, []).extend(pending)
missing = [i for i, entity in enumerate(resolved) if entity is None]
if missing:
raise RuntimeError(
f"Entity resolution produced no row for {len(missing)} mention(s) "
f"(indices {missing[:5]}); refusing to link units to a missing parent."
)
return cast(list[ResolvedEntity], resolved)
async def reassert_entities_batch(
self,
bank_id: str,
resolved_entities: list[ResolvedEntity],
conn,
) -> None:
"""Lock (and, if pruned, re-create) resolved parents before linking units.
Phase-1 resolution and the Phase-2 ``unit_entities`` insert run on
different transactions. In the gap, ``prune_orphan_entities`` can delete
a just-resolved parent it legitimately has no ``unit_entities`` row
yet and the Phase-2 FK insert then fails, dropping the whole batch as
non-retryable (silent memory loss, #2662).
Called on the Phase-2 connection immediately before
``link_units_to_entities_batch``, this locks the parents that still
exist (so the pruner blocks until we commit) and re-inserts any that
already vanished, in one round-trip. An entity referenced by a live unit
is by definition not an orphan, so resurrecting it is correct.
"""
# Deduplicate by id and lock in a stable order so concurrent reasserts
# acquire row locks consistently (same convention as bulk_insert_links).
seen: set[str] = set()
unique: list[ResolvedEntity] = []
for entity in sorted(resolved_entities, key=lambda e: e.entity_id):
if entity.entity_id in seen:
continue
seen.add(entity.entity_id)
unique.append(entity)
if not unique:
return
await self._ops.bulk_reassert_entities(
conn,
fq_table("entities"),
bank_id,
[entity.entity_id for entity in unique],
[entity.canonical_name for entity in unique],
[entity.entity_kind for entity in unique],
)
return entity_ids
async def link_units_to_entities_batch(
self,
unit_entity_pairs: list[tuple[str, str]] | list[tuple[str, str, datetime | None]],
conn=None,
bank_id: str | None = None,
):
"""
Link multiple memory units to entities in batch (MUCH faster than sequential).
@@ -1238,63 +813,22 @@ class EntityResolver:
if conn is None:
async with acquire_with_retry(self.pool) as conn:
return await self._link_units_to_entities_batch_impl(conn, normalized, bank_id)
return await self._link_units_to_entities_batch_impl(conn, normalized)
else:
return await self._link_units_to_entities_batch_impl(conn, normalized, bank_id)
return await self._link_units_to_entities_batch_impl(conn, normalized)
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 unitentity
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
):
async def _link_units_to_entities_batch_impl(self, conn, unit_entity_pairs: list[tuple[str, str, datetime | None]]):
# Sorted bulk insert to prevent deadlocks from inconsistent lock ordering
# across concurrent transactions on the unit_entities unique index.
sorted_pairs = sorted(unit_entity_pairs, key=lambda t: (t[0], t[1]))
unit_ids = [p[0] for p in sorted_pairs]
entity_ids = [p[1] for p in sorted_pairs]
# The unit→entity posting belongs to whoever stores the memory, so the
# memories store records it. Co-occurrence below is separate and unaffected:
# it references only `entities`, which stays in Postgres either way, and is
# read by the entity-graph endpoint and by resolution's disambiguation signal.
from .memories import get_memories
await get_memories().record_unit_entities(
conn=conn,
ops=self._ops,
fq_table=fq_table,
bank_id=bank_id,
unit_ids=unit_ids,
entity_ids=entity_ids,
txn=txn,
await self._ops.bulk_insert_unit_entities(
conn,
fq_table("unit_entities"),
unit_ids,
entity_ids,
)
# Build maps keyed by unit_id:
@@ -1319,12 +853,20 @@ class EntityResolver:
for unit_id, entity_ids in unit_to_entities.items():
entity_list = list(entity_ids)
event_date = unit_event_date.get(unit_id)
for key in _canonical_cooccurrence_pairs(entity_list):
prev = cooccurrence_pairs.get(key, _SENTINEL_MISSING)
if prev is _SENTINEL_MISSING:
cooccurrence_pairs[key] = event_date
else:
cooccurrence_pairs[key] = _later_date(prev, event_date)
for i, entity_id_1 in enumerate(entity_list):
for entity_id_2 in entity_list[i + 1 :]:
if entity_id_1 == entity_id_2:
continue
# Canonical ordering (entity_id_1 < entity_id_2) matches the
# entity_cooccurrences PK and check constraint.
if entity_id_1 > entity_id_2:
entity_id_1, entity_id_2 = entity_id_2, entity_id_1
key = (entity_id_1, entity_id_2)
prev = cooccurrence_pairs.get(key, _SENTINEL_MISSING)
if prev is _SENTINEL_MISSING:
cooccurrence_pairs[key] = event_date
else:
cooccurrence_pairs[key] = _later_date(prev, event_date)
# Accumulate co-occurrence pairs for post-transaction flush.
# The actual INSERT/UPDATE is deferred to flush_pending_stats() to avoid
@@ -1336,3 +878,58 @@ class EntityResolver:
_CooccurrencePair(entity_id_1=e1, entity_id_2=e2, event_date=ed)
for (e1, e2), ed in cooccurrence_pairs.items()
)
async def get_units_by_entity(self, entity_id: str, limit: int = 100) -> list[str]:
"""
Get all units that mention an entity.
Args:
entity_id: Entity ID
limit: Max results
Returns:
List of unit IDs
"""
async with acquire_with_retry(self.pool) as conn:
rows = await conn.fetch(
f"""
SELECT unit_id
FROM {fq_table("unit_entities")}
WHERE entity_id = $1
ORDER BY unit_id
LIMIT $2
""",
entity_id,
limit,
)
return [row["unit_id"] for row in rows]
async def get_entity_by_text(
self,
bank_id: str,
entity_text: str,
) -> str | None:
"""
Find an entity by text (for query resolution).
Args:
bank_id: bank ID
entity_text: Entity text to search for
Returns:
Entity ID if found, None otherwise
"""
async with acquire_with_retry(self.pool) as conn:
row = await conn.fetchrow(
f"""
SELECT id FROM {fq_table("entities")}
WHERE bank_id = $1
AND canonical_name ILIKE $2
ORDER BY mention_count DESC
LIMIT 1
""",
bank_id,
entity_text,
)
return row["id"] if row else None
@@ -1,62 +1,52 @@
"""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.
same probes retain uses (:func:`fetch_temporal_neighbors`,
:func:`compute_semantic_links_ann`) and insert the missing links.
``bulk_insert_links`` has ``ON CONFLICT DO NOTHING`` on the uniqueness
key, so we can re-probe freely and the DB de-dupes.
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 ``unit_entities`` 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_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.
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.
All three passes run on every invocation. The queue is the only source
of work for pass 1; passes 2 and 3 are bank-wide sweeps backed by indexes
on ``entities(bank_id)`` and ``unit_entities(entity_id)``, so they're
cheap when there's nothing to do.
The worker dedupes on bank: a second job for the same bank is dropped
while one is pending. Once processing starts, a new job becomes the
*next* pending slot so work enqueued during processing gets picked up
by the follow-up run.
That follow-up run is *deferred*, not parallel: ``claim_tasks`` will not claim a
graph_maintenance row for a bank that already has one in flight (#3230). Two
concurrent runs would do no extra work anyway each drains the same two
bank-scoped queues while convoying on each other's row locks and holding a
worker slot each.
"""
from __future__ import annotations
import logging
import time
import uuid as uuid_module
from dataclasses import dataclass
from typing import TYPE_CHECKING
from typing import TYPE_CHECKING, Any
from ..models import RequestContext
from .db.base import DatabaseConnection
# Re-exported for callers and tests that import the link caps from here; the caps
# themselves live with the link builders the relink pass mirrors — the temporal one
# with the retain-time builders, the semantic one with the store's relink pass — so
# there is a single definition of each and the two cannot drift.
from .memories.pg.graph import MAX_SEMANTIC_LINKS_PER_UNIT # noqa: F401
from .retain.link_utils import MAX_TEMPORAL_LINKS_PER_UNIT # noqa: F401
from .retain.link_utils import (
MAX_TEMPORAL_LINKS_PER_UNIT,
_bulk_insert_links,
_normalize_datetime,
compute_semantic_links_ann,
)
from .schema import fq_table
if TYPE_CHECKING:
@@ -64,14 +54,17 @@ 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
# Mirrors the ``top_k`` default in ``compute_semantic_links_ann`` at retain
# time. If you change one, change the other — otherwise victims would either
# never reach the cap (probe returns less than the cap) or stay perpetually
# under it (cap is higher than retain creates).
MAX_SEMANTIC_LINKS_PER_UNIT = 50
# Worker fetches this many rows per relink-loop iteration. Bounds
# per-iteration probe/insert latency so a 10k-row backlog doesn't hold a
# worker slot for minutes. Chosen so the typical iteration runs in well
# under 1s.
_DRAIN_BATCH_SIZE = 50
@dataclass
@@ -80,266 +73,285 @@ 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,
}
async def enqueue_relink_victims(
conn: DatabaseConnection,
bank_id: str,
affected_unit_ids: list[str],
include_affected_units: bool = False,
deleted_unit_ids: list[str],
ops: Any,
) -> int:
"""Enqueue surviving units whose outgoing temporal/semantic links pointed at
``affected_unit_ids`` for later link top-up.
``deleted_unit_ids`` for later link top-up.
Must run inside the same transaction that drops those links, *before* the
delete (or cascade) fires once the rows are gone, the join that finds the
victims returns nothing.
``include_affected_units`` covers the case where the affected units are NOT
being removed: an edit deletes every link incident to the edited unit but
leaves it live, so the unit needs its own outgoing adjacency rebuilt too.
Passing it for a unit that will be gone at commit is harmless but pointless
the drain skips queue rows with no live unit so callers should only set
it when the unit survives the transaction.
Delegated to the memories store: finding the victims is a `memory_links`
query, and a store whose links are inline has none, so it returns 0 and the
relink pass has nothing to do. The store resolves the dialect it needs from
``conn``.
Must run inside the same transaction that deletes the units, *before* the
cascade fires once the rows are gone, the join that finds the victims
returns nothing.
Args:
conn: Database connection inside the active transaction.
bank_id: Bank owning the affected units.
affected_unit_ids: Memory_unit IDs whose incident temporal/semantic
links are about to be (or are being) removed.
include_affected_units: Also enqueue ``affected_unit_ids`` themselves,
for callers that leave them live.
conn: Database connection inside the active delete transaction.
bank_id: Bank owning the deleted units.
deleted_unit_ids: Memory_unit IDs about to be (or being) deleted.
ops: ``DataAccessOps`` instance, supplies the dialect-specific
bulk-insert path.
Returns:
Number of distinct victim units enqueued (0 for a store with no links).
Number of distinct victim units enqueued (after dedup against rows
already in the queue).
"""
if not affected_unit_ids:
if not deleted_unit_ids:
return 0
from .memories import get_memories
deleted_uuids = [uuid_module.UUID(uid) if isinstance(uid, str) else uid for uid in deleted_unit_ids]
deleted_str_set = {str(uid) for uid in deleted_uuids}
return await get_memories().enqueue_relink_victims(
conn=conn,
fq_table=fq_table,
bank_id=bank_id,
affected_unit_ids=affected_unit_ids,
include_affected_units=include_affected_units,
# Find units (other than the ones being deleted) that have an outgoing
# temporal/semantic link pointing at a doomed unit. Entity links are
# intentionally excluded — they're scheduled for removal and would only
# add noise to the recompute job.
victim_rows = await conn.fetch(
f"""
SELECT DISTINCT from_unit_id
FROM {fq_table("memory_links")}
WHERE to_unit_id = ANY($1::uuid[])
AND bank_id = $2
AND link_type IN ('temporal', 'semantic')
""",
deleted_uuids,
bank_id,
)
victim_ids = [row["from_unit_id"] for row in victim_rows if str(row["from_unit_id"]) not in deleted_str_set]
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:
if not victim_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,
await ops.enqueue_graph_maintenance(
conn,
fq_table("graph_maintenance_queue"),
bank_id,
victim_ids,
)
logger.debug(
f"[GRAPH_MAINT] Enqueued {len(victim_ids)} relink victims in "
f"bank={bank_id} (deleted {len(deleted_unit_ids)} units)"
)
return len(victim_ids)
async def run_graph_maintenance_job(
memory_engine: "MemoryEngine",
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`.
"""
from ..config import get_config
from .memories import get_memories
del request_context # accepted for symmetry with other run_*_job helpers
backend = await memory_engine._get_backend()
store = get_memories()
config = get_config()
ops = backend.ops
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
# its own link table, so how it batches and re-probes is its business — including
# the #3034 serialisation (the claim takes queue rows FOR UPDATE in (bank_id,
# unit_id) order against a concurrent re-enqueue), which lives in the store's
# claim (`ops.claim_graph_maintenance_batch`). A store with no links returns an
# empty dict and this is a no-op.
relink = await store.relink_pass(
backend=backend, fq_table=fq_table, bank_id=bank_id, config=config, deadline=deadline
)
result.relink_units_processed = relink.units_processed
result.relink_links_added = relink.links_added
# Per-iteration loop: claim → top up → commit. We rely on submit-time
# dedup to keep at most one job per bank running, so no need for
# SKIP LOCKED.
iterations = 0
while True:
from .memory_engine import acquire_with_retry
# --- 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
async with acquire_with_retry(backend) as conn:
async with conn.transaction():
unit_ids = await ops.claim_graph_maintenance_batch(
conn,
fq_table("graph_maintenance_queue"),
bank_id,
_DRAIN_BATCH_SIZE,
)
if not unit_ids:
break
result.relink_links_added += await _relink_batch(conn, bank_id, unit_ids, ops, backend)
result.relink_units_processed += len(unit_ids)
iterations += 1
if iterations > 10000:
# Defensive guard against runaway loops — at 50 units/iter that's
# 500k targets, far beyond any realistic single-bank backlog.
logger.error(
f"[GRAPH_MAINT] bank={bank_id} hit iteration cap ({iterations}); aborting relink ({result.as_dict()})"
)
break
# --- Pass 2 & 3: entity / cooccurrence sweeps ---
# Bank-wide single-statement deletes. Cheap when there's nothing to do.
from .memory_engine import acquire_with_retry
async with acquire_with_retry(backend) as conn:
async with conn.transaction():
result.orphan_entities_pruned = await ops.prune_orphan_entities(
conn,
fq_table("entities"),
fq_table("unit_entities"),
bank_id,
)
# The orphan prune above cascades cooccurrences via FK. The
# explicit cooccurrence pass below catches the *stale-count*
# case: both entities still exist but no current unit witnesses
# them together.
result.stale_cooccurrences_pruned = await ops.prune_stale_cooccurrences(
conn,
fq_table("entity_cooccurrences"),
fq_table("unit_entities"),
fq_table("entities"),
bank_id,
)
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}"
)
return result.as_dict()
async def _relink_batch(
conn: DatabaseConnection,
bank_id: str,
victim_ids: list[str],
ops: Any,
backend: Any,
) -> int:
"""Top up temporal/semantic links for a batch of victim units. Returns rows inserted."""
# Load each victim's metadata. Victims whose units were deleted between
# enqueue and now silently drop out — exactly the no-op behaviour we want
# for stale queue rows.
victim_uuids = [uuid_module.UUID(vid) for vid in victim_ids]
victim_rows = await conn.fetch(
f"""
SELECT id::text AS id, event_date, fact_type, embedding::text AS embedding
FROM {fq_table("memory_units")}
WHERE id = ANY($1::uuid[])
AND bank_id = $2
AND fact_type IN ('experience', 'world')
""",
victim_uuids,
bank_id,
)
if not victim_rows:
return 0
alive_uuids = [uuid_module.UUID(row["id"]) for row in victim_rows]
# Count current outgoing temporal/semantic links per victim so we only
# probe for the ones genuinely below cap. Saves the bulk of the work when
# most victims still have plenty of links.
count_rows = await conn.fetch(
f"""
SELECT from_unit_id, link_type, COUNT(*) AS cnt
FROM {fq_table("memory_links")}
WHERE from_unit_id = ANY($1::uuid[])
AND bank_id = $2
AND link_type IN ('temporal', 'semantic')
GROUP BY from_unit_id, link_type
""",
alive_uuids,
bank_id,
)
counts: dict[tuple[str, str], int] = {}
for row in count_rows:
counts[(str(row["from_unit_id"]), row["link_type"])] = int(row["cnt"])
# --- Temporal top-up ---
temporal_needs = [r for r in victim_rows if counts.get((r["id"], "temporal"), 0) < MAX_TEMPORAL_LINKS_PER_UNIT]
new_links: list[tuple] = []
if temporal_needs:
lateral_unit_ids = [uuid_module.UUID(r["id"]) for r in temporal_needs if r["event_date"] is not None]
lateral_event_dates = [
_normalize_datetime(r["event_date"]) for r in temporal_needs if r["event_date"] is not None
]
lateral_fact_types = [r["fact_type"] for r in temporal_needs if r["event_date"] is not None]
if lateral_unit_ids:
rows = await ops.fetch_temporal_neighbors(
conn,
fq_table("memory_units"),
bank_id,
lateral_unit_ids,
lateral_event_dates,
lateral_fact_types,
MAX_TEMPORAL_LINKS_PER_UNIT,
)
for row in rows:
time_diff_h = float(row["time_diff_hours"])
# Mirror the 24h window enforced at retain time. The bidirectional
# index scan returns the K closest neighbours regardless of
# window, so we filter here.
if time_diff_h > 24:
continue
weight = max(0.3, 1.0 - (time_diff_h / 24))
new_links.append((row["from_id"], str(row["id"]), "temporal", weight, None))
# --- Semantic top-up ---
# ANN must run on its own connection: it opens a nested transaction with
# SET LOCAL hnsw.ef_search + CREATE TEMP TABLE ON COMMIT DROP, and nesting
# that inside our current write transaction would commit our writes early.
semantic_needs = [
r
for r in victim_rows
if counts.get((r["id"], "semantic"), 0) < MAX_SEMANTIC_LINKS_PER_UNIT and r["embedding"] is not None
]
if semantic_needs:
from .memory_engine import acquire_with_retry
seed_ids = [r["id"] for r in semantic_needs]
seed_embs = [r["embedding"] for r in semantic_needs]
seed_ftypes = [r["fact_type"] for r in semantic_needs]
async with acquire_with_retry(backend) as ann_conn:
try:
ann_links = await compute_semantic_links_ann(
ann_conn,
bank_id,
seed_ids,
seed_embs,
fact_types=seed_ftypes,
)
# Strip self-links (rare but possible because the ANN probe
# has no exclude list — see the comment in compute_semantic_links_ann).
ann_links = [lnk for lnk in ann_links if lnk[0] != lnk[1]]
new_links.extend(ann_links)
except Exception as e:
# ANN uses PG-specific HNSW syntax; on dialects/configs where
# it isn't available we still want the temporal top-up to land.
logger.warning(f"[GRAPH_MAINT] Semantic top-up failed for bank={bank_id}: {type(e).__name__}: {e}")
if not new_links:
return 0
await _bulk_insert_links(
conn,
new_links,
bank_id=bank_id,
skip_exists_check=False,
ops=ops,
)
return len(new_links)
@@ -6,7 +6,6 @@ authentication when a TenantExtension is configured.
"""
from abc import ABC, abstractmethod
from dataclasses import dataclass
from datetime import datetime
from typing import TYPE_CHECKING, Any
@@ -14,26 +13,9 @@ if TYPE_CHECKING:
from hindsight_api.engine.memory_engine import BankLlmHealthInfo, Budget
from hindsight_api.engine.response_models import RecallResult, ReflectResult
from hindsight_api.engine.search.tags import TagsMatch
from hindsight_api.extensions import BankWriteOperation
from hindsight_api.models import RequestContext
@dataclass(frozen=True)
class BankConfigState:
"""Resolved bank configuration and its bank-level overrides."""
config: dict[str, Any]
overrides: dict[str, Any]
@dataclass(frozen=True)
class BankTemplateImportWrite:
"""One bank-write decision reserved for a specific imported resource."""
operation: "BankWriteOperation"
target: str | None = None
class MemoryEngineInterface(ABC):
"""
Abstract interface for the Memory Engine.
@@ -198,37 +180,6 @@ class MemoryEngineInterface(ABC):
"""
...
@abstractmethod
async def get_bank_config(
self,
bank_id: str,
*,
request_context: "RequestContext",
) -> BankConfigState:
"""Return resolved configuration after authenticating and authorizing the read."""
...
@abstractmethod
async def update_bank_config(
self,
bank_id: str,
updates: dict[str, Any],
*,
request_context: "RequestContext",
) -> BankConfigState:
"""Create a bank if needed and persist validated configuration overrides."""
...
@abstractmethod
async def reset_bank_config(
self,
bank_id: str,
*,
request_context: "RequestContext",
) -> BankConfigState:
"""Remove all bank configuration overrides after authorization."""
...
@abstractmethod
async def update_bank_disposition(
self,
@@ -324,8 +275,6 @@ class MemoryEngineInterface(ABC):
*,
fact_type: str | None = None,
search_query: str | None = None,
entity_id: str | None = None,
created_before: datetime | None = None,
limit: int = 100,
offset: int = 0,
request_context: "RequestContext",
@@ -337,8 +286,6 @@ class MemoryEngineInterface(ABC):
bank_id: The memory bank ID.
fact_type: Filter by fact type.
search_query: Full-text search query.
entity_id: Filter to memory units linked to this entity ID.
created_before: Keep units with ``created_at`` before this instant.
limit: Maximum results.
offset: Pagination offset.
request_context: Request context for authentication.
@@ -531,15 +478,11 @@ class MemoryEngineInterface(ABC):
Get consolidation freshness for a bank.
Cheap alternative to get_bank_stats when callers only need
last_consolidated_at / last_memory_write_at / pending_consolidation /
failed_consolidation.
last_consolidated_at / pending_consolidation / failed_consolidation.
Returns:
Dict with last_consolidated_at and last_memory_write_at (ISO-8601
strings or None), pending_consolidation (int), and
failed_consolidation (int). last_memory_write_at is the newest write
across the bank's memories — a mental model refreshed at or after it
cannot be stale, whatever its scope.
Dict with last_consolidated_at (ISO-8601 string or None),
pending_consolidation (int), and failed_consolidation (int).
"""
...
@@ -622,30 +565,6 @@ class MemoryEngineInterface(ABC):
"""
...
@abstractmethod
async def delete_operation(
self,
bank_id: str,
operation_id: str,
*,
request_context: "RequestContext",
) -> dict[str, Any]:
"""
Delete a terminal async operation record.
Args:
bank_id: The memory bank ID.
operation_id: The operation ID to delete.
request_context: Request context for authentication.
Returns:
Dict with success status and message.
Raises:
ValueError: If operation not found.
"""
...
@abstractmethod
async def update_bank(
self,
@@ -653,8 +572,6 @@ class MemoryEngineInterface(ABC):
*,
name: str | None = None,
mission: str | None = None,
config_updates: dict[str, Any] | None = None,
create_if_missing: bool = True,
request_context: "RequestContext",
) -> dict[str, Any]:
"""
@@ -664,9 +581,6 @@ class MemoryEngineInterface(ABC):
bank_id: The memory bank ID.
name: New bank name (optional).
mission: New mission text (optional, replaces existing).
config_updates: Bank configuration overrides to apply with the profile update.
create_if_missing: Create a missing bank when True; otherwise raise
a 404 operation error.
request_context: Request context for authentication.
Returns:
@@ -5,57 +5,12 @@ 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
from datetime import datetime
from enum import StrEnum
from typing import Any, Callable, Self
from typing import Any
from .response_models import LLMToolCallResult
logger = logging.getLogger(__name__)
class LLMToolChoiceMode(StrEnum):
"""Canonical tool-selection modes shared by every LLM provider."""
AUTO = "auto"
NONE = "none"
REQUIRED = "required"
NAMED = "named"
@dataclass(frozen=True, slots=True)
class LLMToolChoice:
"""Typed internal tool selection serialized only at provider boundaries."""
mode: LLMToolChoiceMode
function_name: str | None = None
def __post_init__(self) -> None:
if self.mode is LLMToolChoiceMode.NAMED:
if self.function_name is None or not self.function_name or self.function_name != self.function_name.strip():
raise ValueError("Named tool choice requires a non-empty canonical function name")
elif self.function_name is not None:
raise ValueError(f"Tool choice mode {self.mode.value!r} cannot include a function name")
@classmethod
def named(cls, function_name: str) -> Self:
return cls(mode=LLMToolChoiceMode.NAMED, function_name=function_name)
@property
def selected_function_name(self) -> str:
if self.function_name is None:
raise ValueError("Tool choice does not select a named function")
return self.function_name
LLM_TOOL_CHOICE_AUTO = LLMToolChoice(mode=LLMToolChoiceMode.AUTO)
LLM_TOOL_CHOICE_NONE = LLMToolChoice(mode=LLMToolChoiceMode.NONE)
LLM_TOOL_CHOICE_REQUIRED = LLMToolChoice(mode=LLMToolChoiceMode.REQUIRED)
class LLMInterface(ABC):
"""
@@ -71,7 +26,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 +37,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:
@@ -139,7 +71,6 @@ class LLMInterface(ABC):
strict_schema: bool = False,
return_usage: bool = False,
cached_prefix: str | None = None,
attempt_context: Callable[[], AbstractAsyncContextManager[None]] | None = None,
) -> Any:
"""
Make an LLM API call with retry logic.
@@ -161,11 +92,6 @@ class LLMInterface(ABC):
cached_prefix: Opaque handle from ``get_or_create_cached_prefix`` for the
cacheable system prefix, or None. Providers without explicit prompt
caching ignore it (and the wrapper only forwards it when set).
attempt_context: Factory for an async context manager holding the shared
concurrency permits. Passed only when the provider declares
``supports_attempt_scoped_concurrency()``; the provider must enter it
around each individual upstream request so retry backoff never
occupies a permit.
Returns:
If return_usage=False: Parsed response if response_format is provided, otherwise text content.
@@ -188,10 +114,8 @@ class LLMInterface(ABC):
max_retries: int = 5,
initial_backoff: float = 1.0,
max_backoff: float = 30.0,
tool_choice: LLMToolChoice = LLM_TOOL_CHOICE_AUTO,
tool_choice: str | dict[str, Any] = "auto",
cached_prefix: str | None = None,
cached_prefix_message_count: int = 0,
attempt_context: Callable[[], AbstractAsyncContextManager[None]] | None = None,
) -> LLMToolCallResult:
"""
Make an LLM API call with tool/function calling support.
@@ -205,9 +129,7 @@ class LLMInterface(ABC):
max_retries: Maximum retry attempts.
initial_backoff: Initial backoff time in seconds.
max_backoff: Maximum backoff time in seconds.
tool_choice: Canonical tool-selection policy.
attempt_context: Factory for an async context manager holding the shared
concurrency permits see ``call``.
tool_choice: How to choose tools - "auto", "none", "required", or specific function.
Returns:
LLMToolCallResult with content and/or tool_calls.
@@ -223,10 +145,6 @@ class LLMInterface(ABC):
"""
return False
def supports_attempt_scoped_concurrency(self) -> bool:
"""Whether retries can acquire concurrency permits per upstream attempt."""
return False
# ── Prompt prefix caching (optional, per-provider) ─────────────────────────
def supports_prompt_caching(self) -> bool:
@@ -267,45 +185,6 @@ class LLMInterface(ABC):
"""
return None
# ── Step-by-step incremental prompt caching (optional) ─────────────────────
#
# For agentic loops (reflect) the dominant cost is the conversation prefix
# re-sent every turn, not the static system prefix. Providers that can cache
# a *growing* prefix implement these: the caller rolls one cache per step
# (each covering the previous step's full input), passes its handle plus the
# message count it covers to ``call_with_tools`` so only the new turns are
# sent fresh, and tears the caches down when the loop ends. Default no-ops so
# non-supporting providers transparently run uncached.
def supports_incremental_prompt_cache(self) -> bool:
"""Whether this provider can cache a growing multi-turn conversation prefix."""
return False
async def create_incremental_cache(
self,
*,
session_id: str,
messages: list[dict[str, Any]],
tools: list[dict[str, Any]] | None = None,
) -> str | None:
"""Cache ``system + tools + messages`` and return an opaque handle, or None.
The handle is passed back to ``call_with_tools(cached_prefix=...,
cached_prefix_message_count=len(messages))``. Caches are grouped under
``session_id`` for teardown via ``delete_cache_session``. Returns None
when caching is unavailable or the prefix is too small caller falls
back to an uncached call.
"""
return None
async def delete_cached_prefix(self, name: str) -> None:
"""Best-effort delete of a single cache handle (a superseded step)."""
return None
async def delete_cache_session(self, session_id: str) -> None:
"""Best-effort teardown of every cache created under ``session_id``."""
return None
async def submit_batch(
self,
requests: list[dict[str, Any]],
@@ -36,21 +36,6 @@ from .db_utils import acquire_with_retry
logger = logging.getLogger(__name__)
def _llm_requests_persistable() -> bool:
"""Whether the ``llm_requests`` table exists on the active backend.
``llm_requests`` is PostgreSQL-only: its migration is ``run_for_dialect(pg=...)``
with the Oracle slot intentionally absent, and MaintenanceLoop skips its
retention sweep on Oracle for the same reason. On Oracle the table does not
exist, so best-effort trace writes must be skipped rather than attempted
otherwise every LLM call fires an INSERT that fails with ORA-00903 and spams
the error log. Mirrors the ``_is_oracle()`` gate in MaintenanceLoop.start.
"""
from .schema import _is_oracle
return not _is_oracle()
# ── bank/operation attribution (carried across the async call chain) ──────────
@@ -391,32 +376,10 @@ class LLMTraceRecorder:
# INSERTs it patches — but it must not block on unrelated operations).
self._pending: dict[str | None, set[asyncio.Task]] = {}
def _writable(self) -> Any | None:
"""Return the pool to write through, or None if writing isn't possible.
Covers the two lifecycle windows in which best-effort trace writes must
be skipped rather than attempted: before the backend pool is created
(``initialize()`` verifies the LLM before the DB is up) and during/after
shutdown. Writes already in flight need no handling the pools close
gracefully, waiting for their connections to be released.
"""
pool = self._pool_getter()
if pool is None:
return None
# Backends declare readiness explicitly; a raw pool (some callers pass
# one directly) has no lifecycle flag and is assumed usable.
from .db.base import DatabaseBackend
if isinstance(pool, DatabaseBackend) and not pool.is_ready:
return None
return pool
def is_enabled(self, scope: str) -> bool:
"""Whether tracing is active for the given call scope."""
if not self._enabled:
return False
if not _llm_requests_persistable():
return False
if self._allowed_scopes is not None:
return scope in self._allowed_scopes
return True
@@ -510,7 +473,7 @@ class LLMTraceRecorder:
async def _safe_write(self, record: LLMRequestRecord) -> None:
"""Write a trace row. Errors are logged, never raised."""
pool = self._writable()
pool = self._pool_getter()
if pool is None:
logger.debug("LLM trace skipped: pool not available")
return
@@ -583,7 +546,7 @@ class LLMTraceRecorder:
ids are snapshotted synchronously here because the caller may reset the
context immediately after.
"""
if not self._enabled or not _llm_requests_persistable() or trace_ctx is None or not trace_ctx.trace_id:
if not self._enabled or trace_ctx is None or not trace_ctx.trace_id:
return
created_ids = list(dict.fromkeys([*(created or []), *trace_ctx.created_memory_ids]))
source_ids = list(dict.fromkeys([*(source or []), *trace_ctx.source_memory_ids]))
@@ -605,9 +568,8 @@ class LLMTraceRecorder:
# so the UPDATE patches rows that already exist rather than racing ahead
# of them (without blocking on unrelated operations' pending writes).
await self._flush_pending(trace_id)
pool = self._writable()
pool = self._pool_getter()
if pool is None:
logger.debug("LLM trace memory_id attach skipped: pool not available")
return
try:
schema = self._schema_getter()
@@ -9,11 +9,9 @@ import os
import re
import time
import uuid
from contextlib import AsyncExitStack, asynccontextmanager
from contextlib import AsyncExitStack
from typing import TYPE_CHECKING, Any
from json_repair import repair_json
# Vertex AI imports (conditional - for LLMProvider to pass credentials to GeminiLLM)
try:
from google.oauth2 import service_account
@@ -29,19 +27,13 @@ from ..config import (
ENV_REFLECT_LLM_MAX_CONCURRENT,
ENV_RETAIN_LLM_MAX_CONCURRENT,
)
from .cache_affinity import parse_cache_affinity
from .llm_interface import (
LLM_TOOL_CHOICE_AUTO,
LLMToolChoice,
LLMToolChoiceMode,
)
from .llm_interface import (
OutputTooLongError as OutputTooLongError,
)
if TYPE_CHECKING:
from .response_models import LLMToolCallResult
# Seed applied to every Groq request for deterministic behavior.
DEFAULT_LLM_SEED = 4242
logger = logging.getLogger(__name__)
# Disable httpx logging
@@ -115,34 +107,13 @@ def _semaphores_for_scope(scope: str) -> list[asyncio.Semaphore]:
return [per_op, _global_llm_semaphore]
@asynccontextmanager
async def _attempt_permits(scope: str):
"""Hold configured LLM concurrency permits for one upstream attempt."""
from ..worker.stage import get_stage, set_stage
async with AsyncExitStack() as stack:
for sem in _semaphores_for_scope(scope):
await stack.enter_async_context(sem)
try:
yield
except BaseException:
# A failed attempt exits here with its permits released while the
# provider classifies the error and sleeps out its backoff. Suffix
# the stage so `attempt=N` always means "permits held, request in
# flight" (#3002); the next attempt re-stamps after re-acquiring.
stage = get_stage()
if stage is not None and not stage.endswith(".backoff"):
set_stage(f"{stage}.backoff")
raise
def _request_params(
*,
max_completion_tokens: int | None = None,
temperature: float | None = None,
scope: str | None = None,
response_format: Any | None = None,
tool_choice: LLMToolChoice | None = None,
tool_choice: str | dict[str, Any] | None = None,
) -> dict[str, Any] | None:
"""Build the requested-params bag for tracing — only values the caller set.
@@ -157,8 +128,8 @@ def _request_params(
params["temperature"] = temperature
if response_format is not None:
params["response_schema"] = getattr(response_format, "__name__", None) or "structured"
if tool_choice is not None and tool_choice.mode is not LLMToolChoiceMode.AUTO:
params["tool_choice"] = tool_choice.function_name or tool_choice.mode.value
if tool_choice is not None and tool_choice != "auto":
params["tool_choice"] = tool_choice if isinstance(tool_choice, str) else "named"
return params or None
@@ -193,11 +164,16 @@ def sanitize_text(text: str | None) -> str | None:
sanitize_llm_output = sanitize_text
# ``OutputTooLongError`` is re-exported from ``llm_interface`` (the canonical
# definition the providers raise) so that ``fact_extraction`` and ``multi_llm``,
# which import it from here, catch/inspect the very same class. Do NOT redefine
# it locally: a shadow class silently breaks ``except OutputTooLongError`` on the
# real provider path (see issue #3172).
class OutputTooLongError(Exception):
"""
Bridge exception raised when LLM output exceeds token limits.
This wraps provider-specific errors (e.g., OpenAI's LengthFinishReasonError)
to allow callers to handle output length issues without depending on
provider-specific implementations.
"""
pass
def parse_llm_json(raw: str) -> Any:
@@ -208,14 +184,6 @@ def parse_llm_json(raw: str) -> Any:
1. Markdown code fences (```json ... ```) strip them before parsing.
2. Embedded control characters (\\x00-\\x1f, \\x7f) replace with space
and retry if the initial parse fails.
3. Structural malformation (trailing commas, unterminated strings, single
quotes, invalid ``\\escape`` sequences) repaired as a last resort via
``json_repair`` (#2547/#2544).
The repair pass is purely *structural*: it fixes JSON that ``json.loads``
cannot parse at all. It deliberately does NOT touch content semantics
degenerate-but-valid JSON (repetition loops or leaked scaffolding inside
string values) parses fine here and is out of scope for this helper.
Args:
raw: Raw text returned by the LLM.
@@ -224,8 +192,7 @@ def parse_llm_json(raw: str) -> Any:
Parsed Python object (dict, list, etc.).
Raises:
json.JSONDecodeError: If the text cannot be parsed even after cleanup
and structural repair (e.g. repair yields an empty result).
json.JSONDecodeError: If the text cannot be parsed even after cleanup.
"""
text = raw.strip()
@@ -242,19 +209,7 @@ def parse_llm_json(raw: str) -> Any:
# Some models (e.g. Gemini) embed raw control characters inside JSON
# string values. Replacing them with a space usually produces valid JSON.
cleaned = re.sub(r"[\x00-\x1f\x7f]", " ", text)
try:
return json.loads(cleaned)
except json.JSONDecodeError:
# Last resort: structural repair of malformed JSON. ``repair_json`` never
# raises — unrecoverable input yields an empty result ("" / {} / []). Keep
# failing loudly in that case rather than let an empty object masquerade
# as a successful parse: callers (retry ladders, the #1833 fail-loud path)
# rely on JSONDecodeError to retry or surface the failure.
repaired = repair_json(cleaned, return_objects=True)
if not repaired:
raise
return repaired
_PROVIDERS_WITHOUT_API_KEY = frozenset(
@@ -271,7 +226,6 @@ _PROVIDERS_WITHOUT_API_KEY = frozenset(
"litellmrouter",
"bedrock",
"nous",
"xai-oauth",
}
)
@@ -281,23 +235,12 @@ def requires_api_key(provider: str) -> bool:
return provider.lower() not in _PROVIDERS_WITHOUT_API_KEY
def _validate_ollama_num_ctx(value: Any) -> int | None:
"""Validate a native Ollama context-window override."""
if value is None:
return None
if isinstance(value, bool) or not isinstance(value, int):
raise ValueError(f"ollama_num_ctx must be a positive integer, got {value!r}")
if value < 1:
raise ValueError(f"ollama_num_ctx must be >= 1, got {value}")
return value
def create_llm_provider(
provider: str,
api_key: str,
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,
@@ -311,9 +254,6 @@ def create_llm_provider(
litellmrouter_config: dict[str, Any] | None = None,
gemini_service_tier: str | None = None,
timeout: float | None = None,
ollama_num_ctx: int | None = None,
cache_affinity: str | None = None,
structured_output_forced_tool: bool = False,
) -> Any: # Returns LLMInterface
"""
Factory function to create the appropriate LLM provider implementation.
@@ -323,15 +263,11 @@ 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".
gemini_service_tier: Gemini service tier (for Gemini provider) - None (default) or "flex" (50% cheaper).
ollama_num_ctx: Native Ollama context window override. None lets Ollama use the
model/server default.
extra_body: Extra request-body params merged into the provider's native
call. Threaded into OpenAI-compatible, Fireworks, Anthropic, Gemini/
VertexAI and LiteLLM providers (each merges them in its own parameter
@@ -339,20 +275,9 @@ def create_llm_provider(
for OpenAI/Anthropic vs ``max_output_tokens`` for Gemini).
default_headers: Custom headers passed to provider SDK clients (used by operators
routing through proxies / request-tracing middleware). Wired into the Anthropic
provider, the ``OpenAICompatibleLLM`` branch, ``fireworks``, ``nous`` and the
Responses API (SDK ``default_headers``), and into the LiteLLM-backed providers
``litellm``, ``litellmrouter`` and ``bedrock`` as the LiteLLM ``extra_headers``
completion kwarg; other providers may opt in as needed.
cache_affinity: Backend prompt-cache pinning mode, forwarded to the
``OpenAICompatibleLLM`` branch, ``fireworks`` and ``nous`` (all three share the
OpenAI-compatible wire format): "none" (default), "xai_conv_id",
"openai_prompt_cache_key", or "auto". Providers on other branches do their own
cache work or none at all. See ``engine/cache_affinity.py``.
structured_output_forced_tool: Ask the LiteLLM-backed providers (``litellm``,
``litellmrouter``, ``bedrock``) for structured output via a forced tool call
instead of ``response_format``. For backends that reject the response_format
route see ``HINDSIGHT_API_LLM_STRUCTURED_OUTPUT_FORCED_TOOL``. Other
providers ignore it.
provider (SDK ``default_headers``) and the LiteLLM-backed providers ``litellm``,
``litellmrouter`` and ``bedrock`` as the LiteLLM ``extra_headers`` completion
kwarg; other providers may opt in as needed.
vertexai_project_id: Vertex AI project ID (for VertexAI provider).
vertexai_region: Vertex AI region (for VertexAI provider).
vertexai_credentials: Vertex AI credentials object (for VertexAI provider).
@@ -366,8 +291,6 @@ def create_llm_provider(
Returns:
LLMInterface implementation for the specified provider.
"""
ollama_num_ctx = _validate_ollama_num_ctx(ollama_num_ctx)
from .providers import (
AnthropicLLM,
ClaudeCodeLLM,
@@ -380,7 +303,6 @@ def create_llm_provider(
MockLLM,
NoneLLM,
OpenAICompatibleLLM,
OpenAIResponsesLLM,
)
provider_lower = provider.lower()
@@ -398,7 +320,6 @@ def create_llm_provider(
base_url=base_url,
model=model,
reasoning_effort=reasoning_effort,
extra_body=extra_body,
)
elif provider_lower == "claude-code":
@@ -465,7 +386,6 @@ def create_llm_provider(
extra_body=extra_body,
default_headers=default_headers,
timeout=timeout,
structured_output_forced_tool=structured_output_forced_tool,
)
elif provider_lower == "litellmrouter":
@@ -486,7 +406,6 @@ def create_llm_provider(
extra_body=extra_body,
default_headers=default_headers,
timeout=timeout,
structured_output_forced_tool=structured_output_forced_tool,
)
elif provider_lower == "bedrock":
@@ -502,7 +421,6 @@ def create_llm_provider(
default_headers=default_headers,
bedrock_service_tier=bedrock_service_tier,
timeout=timeout,
structured_output_forced_tool=structured_output_forced_tool,
)
elif provider_lower == "llamacpp":
@@ -515,7 +433,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,
@@ -535,16 +452,12 @@ def create_llm_provider(
model=model,
reasoning_effort=reasoning_effort,
extra_body=extra_body,
default_headers=default_headers,
cache_affinity=cache_affinity,
)
elif provider_lower == "nous":
# Nous Portal is OpenAI-compatible on the wire; NousLLM adds rotating
# inference:invoke JWT auth read natively from ~/.hermes/auth.json
# (no static api_key, no hermes_cli dependency — same shape as Codex).
# default_headers/cache_affinity ride NousLLM's **kwargs passthrough to
# OpenAICompatibleLLM.__init__ unchanged (see NousLLM.__init__).
from hindsight_api.engine.providers.nous_llm import NousLLM
return NousLLM(
@@ -554,41 +467,6 @@ def create_llm_provider(
model=model,
reasoning_effort=reasoning_effort,
extra_body=extra_body,
default_headers=default_headers,
cache_affinity=cache_affinity,
timeout=timeout,
)
elif provider_lower == "xai-oauth":
# SuperGrok subscription lane: api.x.ai spoken plainly, but the
# credential is a device-code OAuth grant with proactive/reactive
# refresh over a shared on-disk store, and xAI's 403 shapes need their
# own classification — neither fits the OpenAI SDK client, hence its
# own provider.
from hindsight_api.engine.providers.xai_oauth_llm import XaiOAuthLLM
return XaiOAuthLLM(
provider=provider,
api_key=api_key,
base_url=base_url,
model=model,
reasoning_effort=reasoning_effort,
timeout=timeout,
)
elif provider_lower == "openai-responses":
# OpenAI Responses API (/v1/responses). Unlike chat/completions, it
# supports reasoning + function tools together, so reflect's tool loop
# can run with a real reasoning_effort. See OpenAIResponsesLLM.
return OpenAIResponsesLLM(
provider=provider,
api_key=api_key,
base_url=base_url,
model=model,
reasoning_effort=reasoning_effort,
openai_service_tier=openai_service_tier,
extra_body=extra_body,
default_headers=default_headers,
timeout=timeout,
)
@@ -616,9 +494,6 @@ def create_llm_provider(
groq_service_tier=groq_service_tier,
openai_service_tier=openai_service_tier,
extra_body=extra_body,
default_headers=default_headers,
cache_affinity=cache_affinity,
ollama_num_ctx=ollama_num_ctx,
timeout=timeout,
)
@@ -639,7 +514,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,
@@ -656,9 +531,6 @@ class LLMProvider:
max_retries: int | None = None,
initial_backoff: float | None = None,
max_backoff: float | None = None,
ollama_num_ctx: int | None = None,
cache_affinity: str | None = None,
structured_output_forced_tool: bool = False,
):
"""
Initialize LLM provider.
@@ -668,24 +540,16 @@ 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.
gemini_service_tier: Gemini service tier (None or "flex") - from config.
ollama_num_ctx: Native Ollama context window override. ``None`` lets Ollama
use the model/server default.
gemini_safety_settings: Safety settings for Gemini/VertexAI providers.
extra_body: Extra request-body params merged into the provider's native call
(OpenAI-compatible, Fireworks, Anthropic, Gemini/VertexAI, LiteLLM).
default_headers: Custom headers passed as ``default_headers`` to provider SDK clients.
Used by operators routing through proxies / request-tracing middleware.
cache_affinity: Backend prompt-cache pinning mode for the OpenAI-compatible and
Fireworks providers ("none", "xai_conv_id", "openai_prompt_cache_key",
"auto"). Validated here for every provider so a typo never fails silently;
providers on other factory branches ignore it. Used verbatim callers
resolve the per-operation/global fallback.
litellmrouter_config: Provider-specific config for ``provider="litellmrouter"``.
JSON object passed verbatim to ``litellm.Router(**config)`` see
https://docs.litellm.ai/docs/routing. Ignored unless ``provider == "litellmrouter"``.
@@ -706,9 +570,6 @@ class LLMProvider:
``max_retries``. ``None`` keeps each method's own fallback.
max_backoff: Default maximum retry backoff (seconds), same resolution as
``max_retries``. ``None`` keeps each method's own fallback.
structured_output_forced_tool: Structured output via a forced tool call
instead of ``response_format``, for the LiteLLM-backed providers - from
config (``HINDSIGHT_API_LLM_STRUCTURED_OUTPUT_FORCED_TOOL``).
This constructor uses every argument as passed and does not read global
``HindsightConfig``: resolving the server-level default for a ``None`` argument is the
@@ -737,10 +598,6 @@ class LLMProvider:
self.openai_service_tier = openai_service_tier
self.bedrock_service_tier = bedrock_service_tier
self.gemini_service_tier = gemini_service_tier
# Structured-output transport for the LiteLLM-backed providers. Used verbatim —
# the caller resolves the server-level default, like the fields above.
self.structured_output_forced_tool = structured_output_forced_tool
self.ollama_num_ctx = _validate_ollama_num_ctx(ollama_num_ctx)
# Gemini safety settings (instance default; can be overridden per-request via context var)
self.gemini_safety_settings = gemini_safety_settings
# Gemini prompt caching: when True, retain extraction (and any future
@@ -754,16 +611,10 @@ class LLMProvider:
# Used verbatim — callers resolve the global fallback (see _member_to_llm /
# the per-op builds in MemoryEngine, and LLMProvider.from_env).
self.default_headers = default_headers
# Backend prompt-cache pinning mode. Validated here rather than only at the
# provider so a typo fails for every provider, not just the ones that act on
# it — the setting has no visible effect in the response, so a silent
# fallback to "none" would be indistinguishable from it working.
self.cache_affinity = parse_cache_affinity(cache_affinity).value
# Validate provider
valid_providers = [
"openai",
"openai-responses",
"groq",
"ollama",
"ollama-cloud",
@@ -789,7 +640,6 @@ class LLMProvider:
"atlas",
"fireworks",
"nous",
"xai-oauth",
]
if self.provider not in valid_providers:
raise ValueError(f"Invalid LLM provider: {self.provider}. Must be one of: {', '.join(valid_providers)}")
@@ -892,10 +742,7 @@ class LLMProvider:
gemini_safety_settings=self.gemini_safety_settings,
prompt_cache_enabled=self.prompt_cache_enabled,
litellmrouter_config=router_config,
ollama_num_ctx=self.ollama_num_ctx,
timeout=self.timeout,
cache_affinity=self.cache_affinity,
structured_output_forced_tool=self.structured_output_forced_tool,
)
# Backward compatibility: Keep mock provider properties
@@ -956,7 +803,7 @@ class LLMProvider:
initial_backoff: float | None = None,
max_backoff: float | None = None,
skip_validation: bool = False,
strict_schema: bool | None = None,
strict_schema: bool = False,
return_usage: bool = False,
cached_prefix: str | None = None,
) -> Any:
@@ -977,10 +824,9 @@ class LLMProvider:
configured default (``llm_max_backoff``), else 60.0.
skip_validation: Return raw JSON without Pydantic validation.
strict_schema: Per-call override requesting grammar-enforced (json_schema strict)
structured output instead of the soft json_object path. None (the default)
inherits the server-level HINDSIGHT_API_LLM_STRICT_SCHEMA flag; an explicit
True or False wins over it, so a caller can force strict output on -- or off --
for its own scope. Providers without a strict mode ignore it.
structured output instead of the soft json_object path. The server-level
HINDSIGHT_API_LLM_STRICT_SCHEMA flag is OR-ed in here so it applies to every call;
providers without a strict mode ignore it.
return_usage: If True, return tuple (result, TokenUsage) instead of just result.
Returns:
@@ -998,13 +844,7 @@ class LLMProvider:
from ..worker.stage import set_stage
structured = "+structured" if response_format is not None else ""
# `.queued` until the concurrency permits are in hand — see the acquire
# below. Without it, a call waiting on a saturated semaphore is
# indistinguishable from one the provider is actively running, and the
# label points at the provider (#3002: an operator lost an hour to
# "llm.bedrock.*" for tasks that had never reached Bedrock).
base_stage = f"llm.{self.provider}.{scope}{structured}"
set_stage(f"{base_stage}.queued")
set_stage(f"llm.{self.provider}.{scope}{structured}")
# Resolve the retry policy: explicit per-call arg wins, else the provider's
# configured per-operation/global default, else this method's own fallback.
@@ -1021,18 +861,14 @@ class LLMProvider:
)
# Resolve strict-schema once, here, rather than in each provider: the
# per-call argument, falling back to the server-level
# HINDSIGHT_API_LLM_STRICT_SCHEMA flag when the caller expressed no
# preference. Providers with a json_schema response_format (OpenAI-compatible,
# per-call argument OR the server-level HINDSIGHT_API_LLM_STRICT_SCHEMA
# flag. Providers with a json_schema response_format (OpenAI-compatible,
# LiteLLM) then grammar-enforce structured output instead of the fragile
# soft json_object path; Gemini already enforces its native response_schema,
# and providers without a strict mode simply ignore the flag.
from ..config import get_config
# An explicit per-call value wins in BOTH directions -- `or` would have made a
# per-call False indistinguishable from "unset", silently ignoring any caller
# that opts out while the global flag is on.
strict_schema = strict_schema if strict_schema is not None else get_config().llm_strict_schema
strict_schema = strict_schema or get_config().llm_strict_schema
# LLM call observability flows through the OTel GenAI recorder
# (tracing.get_span_recorder().record_llm_call). Provider implementations
@@ -1061,18 +897,9 @@ class LLMProvider:
# hand so the error path below can attach it if parsing/validation fails.
usage_token = set_response_usage(None)
try:
# Providers that own retry loops acquire the shared permits for each
# upstream attempt so backoff never occupies request capacity.
attempt_gated = self._provider_impl.supports_attempt_scoped_concurrency()
async with AsyncExitStack() as stack:
if not attempt_gated:
for sem in _semaphores_for_scope(scope):
await stack.enter_async_context(sem)
# Permits in hand — only now leave `.queued`. Attempt-gated
# providers acquire permits per attempt instead, so they keep
# `.queued` until their first `attempt=N` stamp lands after
# the permit acquire inside attempt_context (#3002).
set_stage(base_stage)
for sem in _semaphores_for_scope(scope):
await stack.enter_async_context(sem)
# cached_prefix is only set for providers that returned a handle
# from get_or_create_cached_prefix() (e.g. Gemini); it's None for
@@ -1081,7 +908,6 @@ class LLMProvider:
cache_kwarg = {"cached_prefix": cached_prefix} if cached_prefix is not None else {}
try:
# Delegate to provider implementation
attempt_kwarg = {"attempt_context": lambda: _attempt_permits(scope)} if attempt_gated else {}
result = await self._provider_impl.call(
messages=messages,
response_format=response_format,
@@ -1095,7 +921,6 @@ class LLMProvider:
strict_schema=strict_schema,
return_usage=return_usage,
**cache_kwarg,
**attempt_kwarg,
)
except Exception as e:
# The provider call may have succeeded (and incurred token
@@ -1140,9 +965,8 @@ class LLMProvider:
max_retries: int | None = None,
initial_backoff: float | None = None,
max_backoff: float | None = None,
tool_choice: LLMToolChoice = LLM_TOOL_CHOICE_AUTO,
tool_choice: str | dict[str, Any] = "auto",
cached_prefix: str | None = None,
cached_prefix_message_count: int = 0,
) -> "LLMToolCallResult":
"""
Make an LLM API call with tool/function calling support.
@@ -1159,16 +983,14 @@ class LLMProvider:
configured default (``llm_initial_backoff``), else 1.0.
max_backoff: Maximum backoff time in seconds. ``None`` uses the provider's
configured default (``llm_max_backoff``), else 30.0.
tool_choice: Canonical tool-selection policy.
tool_choice: How to choose tools - "auto", "none", "required", or {"type": "function", "function": {"name": "..."}}
Returns:
LLMToolCallResult with content and/or tool_calls.
"""
from ..worker.stage import set_stage
# `.queued` until the permits are held — see the structured path above.
base_stage = f"llm.{self.provider}.{scope}+tools"
set_stage(f"{base_stage}.queued")
set_stage(f"llm.{self.provider}.{scope}+tools")
# Resolve the retry policy: explicit per-call arg wins, else the provider's
# configured per-operation/global default, else this method's own fallback.
@@ -1207,28 +1029,16 @@ class LLMProvider:
# hand so the error path below can attach it if parsing/validation fails.
usage_token = set_response_usage(None)
try:
attempt_gated = self._provider_impl.supports_attempt_scoped_concurrency()
async with AsyncExitStack() as stack:
if not attempt_gated:
for sem in _semaphores_for_scope(scope):
await stack.enter_async_context(sem)
# Permits in hand — only now leave `.queued`; attempt-gated
# providers stay `.queued` until their first post-acquire
# `attempt=N` stamp (see call() above, #3002).
set_stage(base_stage)
for sem in _semaphores_for_scope(scope):
await stack.enter_async_context(sem)
# cached_prefix is only set for providers that returned a handle
# from get_or_create_cached_prefix() / create_incremental_cache();
# forward it (plus how many leading messages it covers) only when
# present so non-caching providers keep their signature.
cache_kwarg = (
{"cached_prefix": cached_prefix, "cached_prefix_message_count": cached_prefix_message_count}
if cached_prefix is not None
else {}
)
# from get_or_create_cached_prefix(); forward it only when present
# so non-caching providers keep their signature (same as call()).
cache_kwarg = {"cached_prefix": cached_prefix} if cached_prefix is not None else {}
try:
# Delegate to provider implementation
attempt_kwarg = {"attempt_context": lambda: _attempt_permits(scope)} if attempt_gated else {}
result = await self._provider_impl.call_with_tools(
messages=messages,
tools=tools,
@@ -1240,7 +1050,6 @@ class LLMProvider:
max_backoff=max_backoff,
tool_choice=tool_choice,
**cache_kwarg,
**attempt_kwarg,
)
except Exception as e:
# The provider call may have succeeded (and incurred token
@@ -1435,17 +1244,15 @@ class LLMProvider:
# does so without building the full HindsightConfig, keeping from_env() a
# lightweight env-only loader (see test_llm_provider_from_env_keeps_lightweight_loader).
from ..config import (
DEFAULT_LLM_CACHE_AFFINITY,
DEFAULT_LLM_GROQ_SERVICE_TIER,
DEFAULT_LLM_OPENAI_SERVICE_TIER,
DEFAULT_LLM_PROMPT_CACHE_ENABLED,
DEFAULT_LLM_PROVIDER,
DEFAULT_LLM_STRUCTURED_OUTPUT_FORCED_TOOL,
DEFAULT_LLM_REASONING_EFFORT,
DEFAULT_LLM_TIMEOUT,
ENV_LLM_API_KEY,
ENV_LLM_BASE_URL,
ENV_LLM_BEDROCK_SERVICE_TIER,
ENV_LLM_CACHE_AFFINITY,
ENV_LLM_DEFAULT_HEADERS,
ENV_LLM_EXTRA_BODY,
ENV_LLM_GEMINI_SAFETY_SETTINGS,
@@ -1453,20 +1260,16 @@ class LLMProvider:
ENV_LLM_GROQ_SERVICE_TIER,
ENV_LLM_LITELLMROUTER_CONFIG,
ENV_LLM_MODEL,
ENV_LLM_OLLAMA_NUM_CTX,
ENV_LLM_OPENAI_SERVICE_TIER,
ENV_LLM_PROMPT_CACHE_ENABLED,
ENV_LLM_PROVIDER,
ENV_LLM_REASONING_EFFORT,
ENV_LLM_STRUCTURED_OUTPUT_FORCED_TOOL,
ENV_LLM_TIMEOUT,
ENV_LLM_VERTEXAI_PROJECT_ID,
ENV_LLM_VERTEXAI_REGION,
ENV_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY,
_get_default_model_for_provider,
_parse_boolean_env,
_parse_llm_router_config,
_parse_optional_positive_int,
parse_gemini_service_tier,
)
@@ -1484,9 +1287,6 @@ class LLMProvider:
model = os.getenv(ENV_LLM_MODEL) or _get_default_model_for_provider(provider)
extra_body = json.loads(os.getenv(ENV_LLM_EXTRA_BODY, "null"))
default_headers = json.loads(os.getenv(ENV_LLM_DEFAULT_HEADERS, "null"))
# Same default as HindsightConfig.from_env: this entry point must not
# resolve to a different mode than the engine's own config path.
cache_affinity = os.getenv(ENV_LLM_CACHE_AFFINITY, DEFAULT_LLM_CACHE_AFFINITY) or None
prompt_cache_enabled = os.getenv(
ENV_LLM_PROMPT_CACHE_ENABLED, str(DEFAULT_LLM_PROMPT_CACHE_ENABLED)
).lower() in (
@@ -1501,10 +1301,9 @@ 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,
groq_service_tier=os.getenv(ENV_LLM_GROQ_SERVICE_TIER, DEFAULT_LLM_GROQ_SERVICE_TIER),
openai_service_tier=os.getenv(ENV_LLM_OPENAI_SERVICE_TIER, DEFAULT_LLM_OPENAI_SERVICE_TIER),
bedrock_service_tier=os.getenv(ENV_LLM_BEDROCK_SERVICE_TIER) or None,
@@ -1515,16 +1314,11 @@ class LLMProvider:
),
gemini_safety_settings=json.loads(os.getenv(ENV_LLM_GEMINI_SAFETY_SETTINGS, "null")),
prompt_cache_enabled=prompt_cache_enabled,
ollama_num_ctx=_parse_optional_positive_int(ENV_LLM_OLLAMA_NUM_CTX, os.getenv(ENV_LLM_OLLAMA_NUM_CTX)),
litellmrouter_config=_parse_llm_router_config(ENV_LLM_LITELLMROUTER_CONFIG),
vertexai_project_id=os.getenv(ENV_LLM_VERTEXAI_PROJECT_ID) or None,
vertexai_region=os.getenv(ENV_LLM_VERTEXAI_REGION) or None,
vertexai_service_account_key=os.getenv(ENV_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY) or None,
timeout=float(os.getenv(ENV_LLM_TIMEOUT, str(DEFAULT_LLM_TIMEOUT))),
structured_output_forced_tool=_parse_boolean_env(
ENV_LLM_STRUCTURED_OUTPUT_FORCED_TOOL,
DEFAULT_LLM_STRUCTURED_OUTPUT_FORCED_TOOL,
),
)
@@ -1,172 +0,0 @@
"""Device selection and post-inference memory release for local (in-process)
SentenceTransformer / CrossEncoder models.
Two concerns live here, both about keeping a local API instance's memory flat:
**1. Device selection MPS is opt-in.**
On Apple Silicon the PyTorch **MPS** (Metal) backend caches a distinct compiled
kernel graph *and* allocator pool per unique input tensor shape, and never
releases them. Under the variable-length, high-volume recall/rerank/embed traffic
the engine generates (documents and candidate sets of every size), that per-shape
cache grows without bound. A local instance was observed idling at ~20 GB ~9.4 GB
of Metal graphics memory plus ~8 GB of native heap, essentially all of it stale
per-shape MPS cache. CPU inference has no per-shape cache: the same workload holds
flat at a few hundred MB, with negligible latency cost for the small default
models (and MPS actually *slows down* over time as it recompiles graphs for new
shapes). So MPS is excluded from auto-detection and must be opted into explicitly;
CUDA and Intel XPU still auto-select.
This is a confirmed, still-open PyTorch bug in the MPSGraph compilation cache
(keyed on tensor shape, no eviction path). We are tracking it upstream:
- https://github.com/pytorch/pytorch/issues/181213
([MPS] unbounded RSS growth with varying-shape inference our exact case)
- https://github.com/pytorch/pytorch/issues/164299 (graphCache identified as
the primary leak culprit)
- https://github.com/pytorch/pytorch/issues/182815 (proposes, but has not yet
shipped, a torch.mps.invalidate_graph_cache() API / PYTORCH_MPS_DISABLE_GRAPH_CACHE
env var that would let us keep MPS)
No released mitigation exists today: empty_cache(), synchronize(),
PYTORCH_MPS_HIGH_WATERMARK_RATIO, and autorelease pools were all confirmed
ineffective upstream. Revisit MPS-as-default once one of those knobs lands.
**2. Memory release after each batch.**
Local CPU inference allocates large transient numpy/tensor buffers per call. The
allocator keeps those freed pages as a high-water mark, so RSS grows monotonically
across many calls (issue #1717). We return them to the OS after each batch —
``malloc_trim`` on glibc/Linux, ``malloc_zone_pressure_relief`` on macOS (the
original #1717 fix covered only Linux). When the model ran on a GPU we also empty
that backend's allocator pool via ``torch.<backend>.empty_cache()``.
"""
from __future__ import annotations
import ctypes
import ctypes.util
import gc
import logging
import sys
logger = logging.getLogger(__name__)
def select_local_device(force_cpu: bool, allow_mps: bool) -> str | None:
"""Choose the device for a local SentenceTransformer / CrossEncoder.
Returns a value suitable to pass as the model's ``device`` argument:
- ``"cpu"`` forced CPU, or the only accelerator is MPS and it is not allowed.
- ``None`` let sentence-transformers auto-detect (picks CUDA / XPU,
handling multi-GPU correctly).
- ``"mps"`` Apple Silicon GPU, only when ``allow_mps`` is set.
MPS is never auto-selected because its per-shape cache leaks unbounded memory
under the engine's variable-length workload (see the module docstring). Set the
matching ``*_ALLOW_MPS`` config flag to opt back in.
"""
if force_cpu:
return "cpu"
try:
import torch
if torch.cuda.is_available():
return None # auto-detect CUDA
if hasattr(torch, "xpu") and torch.xpu.is_available():
return None # auto-detect Intel XPU
mps_available = hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
if mps_available:
if allow_mps:
return "mps"
logger.info(
"Local model: MPS (Apple Silicon GPU) is available but disabled by "
"default because its per-shape cache leaks memory under variable-length "
"workloads; running on CPU. Set the *_ALLOW_MPS flag to opt in."
)
return "cpu"
return "cpu"
except Exception as e: # pragma: no cover - defensive
logger.warning("Local device detection failed, falling back to CPU: %s", e)
return "cpu"
def resolve_model_device_type(model: object) -> str:
"""Best-effort device *type* ("cpu" / "cuda" / "mps" / "xpu") of a loaded model.
Used to decide which GPU allocator pool to empty after inference. Falls back to
``"cpu"`` (the safe no-op choice for release) if the device can't be read.
"""
device = getattr(model, "device", None)
if device is None:
inner = getattr(model, "model", None) # CrossEncoder wraps the HF model
device = getattr(inner, "device", None)
try:
return device.type if device is not None else "cpu"
except Exception: # pragma: no cover - defensive
return "cpu"
def _resolve_heap_trim():
"""Return a callable that asks the C allocator to release freed pages to the OS.
glibc (Linux) exposes ``malloc_trim``; macOS exposes
``malloc_zone_pressure_relief``. Resolved once at import; returns a no-op on
platforms where neither is available (musl, Windows).
"""
if sys.platform == "linux":
libc_path = ctypes.util.find_library("c")
if libc_path is None:
return lambda: None
try:
libc = ctypes.CDLL(libc_path)
trim = libc.malloc_trim
except (OSError, AttributeError):
# Not glibc (musl has no malloc_trim) or libc lookup failed.
return lambda: None
trim.argtypes = [ctypes.c_size_t]
trim.restype = ctypes.c_int
return lambda: trim(0)
if sys.platform == "darwin":
try:
libc = ctypes.CDLL("/usr/lib/libSystem.dylib")
default_zone = libc.malloc_default_zone
default_zone.restype = ctypes.c_void_p
relief = libc.malloc_zone_pressure_relief
relief.argtypes = [ctypes.c_void_p, ctypes.c_size_t]
relief.restype = ctypes.c_size_t
except (OSError, AttributeError):
return lambda: None
# pressure_relief(zone, goal=0) reclaims as much as possible.
return lambda: relief(default_zone(), 0)
return lambda: None
_heap_trim = _resolve_heap_trim()
def _empty_gpu_cache(device_type: str | None) -> None:
"""Empty the allocator pool of the GPU backend the model ran on, if any."""
if not device_type or device_type == "cpu":
return
try:
import torch
backend = getattr(torch, device_type, None) # torch.cuda / torch.mps / torch.xpu
if backend is not None and hasattr(backend, "empty_cache"):
backend.empty_cache()
except Exception: # pragma: no cover - defensive
pass
def release_local_inference_memory(device_type: str | None = None) -> None:
"""Release transient heap (and GPU allocator) memory after a local inference batch.
Frees Python objects, returns freed native pages to the OS, and empties the GPU
allocator pool when the model ran on a GPU. Safe to call on every platform and
device; the pieces that don't apply are cheap no-ops.
"""
gc.collect()
_heap_trim()
_empty_gpu_cache(device_type)
@@ -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
@@ -22,50 +21,24 @@ from one place, so we don't spawn a separate ``asyncio`` task per concern:
The loop wakes on a short fixed tick and runs each job when its own
``last_run + interval`` is due (run-at-start, then on interval), so adding jobs
with different cadences doesn't burst CPU. Cross-tenant discovery goes through
server-side PL/pgSQL routines (``schemas_with_expired_rows`` and
``banks_needing_consolidation``, in the configured schema see ``fq_routine``)
one round-trip each instead of a per-schema query storm, which matters at
thousands of tenants.
The loop runs in *every* API/worker process with no leader election, so a job that
enqueues work must make that enqueue idempotent or the fleet queues one wave per
process. 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.
server-side PL/pgSQL routines (``public.schemas_with_expired_rows`` and
``public.banks_needing_consolidation``) one round-trip each instead of a
per-schema query storm, which matters at thousands of tenants.
"""
from __future__ import annotations
import asyncio
import logging
import random
import time
from collections.abc import Coroutine
from datetime import datetime, timedelta, timezone
from datetime import datetime, timezone
from typing import TYPE_CHECKING, Any
from ..config import HindsightConfig, get_config
from ..models import RequestContext
from .db_utils import acquire_with_retry
from .schema import _is_oracle, fq_routine, fq_table, fq_table_explicit
from .schema import _is_oracle, fq_table
if TYPE_CHECKING:
from .memory_engine import MemoryEngine
@@ -74,40 +47,8 @@ logger = logging.getLogger(__name__)
# Short tick so jobs with different cadences share one loop without per-job tasks.
_TICK_SECONDS = 60
# Cross-store txn recovery (only when the memories store keeps its rows outside SQL): a backstop
# for a writer that crashed between its external writes and the decide. The happy path decides
# inline after commit, so this rarely finds work; five minutes bounds how long a crashed txn stalls
# its namespace's fold.
_TXN_RECOVERY_INTERVAL_SECONDS = 300
# A pending txn is left alone for this long from first sighting before the sweep aborts an
# unwitnessed one — the writer may still be mid-flight (PendingTxn carries no timestamp).
_TXN_RECOVERY_GRACE_SECONDS = 300
# ── 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
# Retention sweeps are not time-sensitive; hourly matches the previous per-sweep cadence.
_RETENTION_INTERVAL_SECONDS = 3600
class MaintenanceLoop:
@@ -119,9 +60,6 @@ class MaintenanceLoop:
self._stop = asyncio.Event()
# Monotonic timestamps of the last run per job, keyed by job name.
self._last_run: dict[str, float] = {}
# Cross-store txn recovery: first-sighting time per pending txn_id, so an unwitnessed
# txn gets a grace period before the sweep aborts it. Persists across ticks.
self._txn_first_seen: dict[str, float] = {}
# ── lifecycle ──────────────────────────────────────────────────────────
@@ -159,44 +97,14 @@ class MaintenanceLoop:
def _any_job_enabled() -> bool:
cfg = get_config()
reconcile_on = cfg.consolidation_reconcile_interval_seconds > 0
# Not gated on audit_log_enabled: that is per-bank overridable, so rows
# can exist even when the deployment default is off. Retention is driven
# purely by the (server-level) window.
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_enabled and 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
return (
reconcile_on
or audit_on
or llm_on
or mm_refresh_on
or op_cleanup_on
or MaintenanceLoop._cross_store_recovery_enabled()
)
@staticmethod
def _cross_store_recovery_enabled() -> bool:
"""True when the memories store keeps memories outside SQL and therefore has
cross-store write-group txns a crashed writer could leave undecided.
Deliberately reads the PROCESS-LEVEL class attribute, not the per-bank
``writes_memory_rows_in_sql_for(bank_id)`` this only decides whether the recovery LOOP
needs to run at all. A store that routes some banks outside SQL keeps the class attribute
False so the loop runs, then ``recover_pending_txns`` is bank-scoped inside it."""
try:
from .memories import get_memories
return not get_memories().writes_memory_rows_in_sql
except Exception:
return False
return reconcile_on or audit_on or llm_on or mm_refresh_on
# ── 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 +115,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 +126,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,15 +134,6 @@ 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)
):
await self._run_timed("operation cleanup", self._run_operation_cleanup(cfg))
if self._cross_store_recovery_enabled() and self._is_due("txn_recovery", _TXN_RECOVERY_INTERVAL_SECONDS):
await self._run_timed("cross-store txn recovery", self._run_txn_recovery())
async def _run_timed(self, name: str, coro: Coroutine[Any, Any, None]) -> None:
"""Run a maintenance job and emit one timing line for it.
@@ -274,205 +152,31 @@ class MaintenanceLoop:
async def _run_retention(self, cfg: HindsightConfig) -> None:
# Retention days are static server-level config, so one global cutoff
# applies to every tenant schema (the routine sweeps them all).
# Not gated on audit_log_enabled: it is per-bank overridable, so a bank
# may be writing audit rows while the deployment default is off. Gating
# the purge on the global flag would let those rows accumulate forever.
if cfg.audit_log_retention_days > 0:
if cfg.audit_log_enabled and cfg.audit_log_retention_days > 0:
await self._purge_expired("audit_log", "started_at", cfg.audit_log_retention_days)
if cfg.llm_trace_enabled and cfg.llm_trace_retention_days > 0:
await self._purge_expired("llm_requests", "started_at", cfg.llm_trace_retention_days)
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
"SELECT * FROM public.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
# ── terminal operation cleanup ─────────────────────────────────────────
async def _run_operation_cleanup(self, cfg: HindsightConfig) -> None:
"""Prune one bounded batch of expired terminal operations per tenant schema.
Previously this rode the worker's task-claiming loop, so it only fired
when that loop happened to iterate and was interleaved with claiming. It
is a periodic housekeeping sweep like the retention jobs above, so it
belongs on the same schedule.
Discovery is one cross-tenant round-trip (``schemas_with_expired_operations``)
rather than a connection + prune transaction per tenant; pending and
processing rows are never prunable, so a schema holding only in-flight
work is correctly reported as having nothing to do.
"""
engine = self._engine
backend = engine._backend
try:
async with acquire_with_retry(backend, max_retries=1) as conn:
rows = await conn.fetch(
f"SELECT * FROM {fq_routine('schemas_with_expired_operations')}($1)",
cfg.operation_retention_days,
)
if result and result != "DELETE 0":
logger.info(f"Retention sweep {schema}.{table}: {result}")
except Exception as e:
logger.warning(f"Operation cleanup discovery failed: {e}")
return
if not rows:
return
# Prune only schemas the deployment actually serves. The routine reports
# every schema owning an async_operations table, including ones tenant
# discovery doesn't claim.
try:
tenants = await engine._tenant_extension.list_tenants()
except Exception as e:
logger.warning(f"Operation cleanup tenant discovery failed: {e}")
return
known = {t.schema for t in tenants} | {get_config().database_schema}
from .memory_engine import _current_schema
cutoff = datetime.now(timezone.utc) - timedelta(days=cfg.operation_retention_days)
pruned = 0
for row in rows:
schema = row[0]
if schema not in known:
continue
# Oracle resolves unqualified names from a context-bound session
# schema; on PostgreSQL this is harmless and fq_table stays explicit.
token = _current_schema.set(schema)
try:
table = fq_table_explicit("async_operations", schema)
async with acquire_with_retry(backend, max_retries=1) as conn:
# 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
)
async with conn.transaction():
deleted = await backend.ops.prune_terminal_operations(
conn, table, cutoff, batch_size=cfg.operation_cleanup_batch_size
)
if deleted:
pruned += deleted
logger.info(f"Operation cleanup pruned {deleted} expired terminal operations from {schema}")
except Exception as e:
logger.warning(f"Operation cleanup failed for schema {schema}: {e}")
finally:
_current_schema.reset(token)
if pruned:
logger.info(f"Operation cleanup: pruned {pruned} operation(s) total")
# ── cross-store txn recovery ─────────────────────────────────────────────
async def _run_txn_recovery(self) -> None:
"""Resolve write-group txns a crashed writer left undecided, for a store that keeps its
rows outside SQL.
For each bank, the store lists its namespace's pending txns and decides each against the
Postgres witness table (present commit, absent past the grace abort never on
assumption), then reaps expired witness rows. A no-op for the SQL stores. Best-effort: a
failure here only delays a stalled fold until the next tick.
"""
from .memories import get_memories
store = get_memories()
if store.writes_memory_rows_in_sql:
return
backend = self._engine._backend
try:
async with acquire_with_retry(backend, max_retries=1) as conn:
bank_ids = [r[0] for r in await conn.fetch(f"SELECT bank_id FROM {fq_table('banks')}")]
if not bank_ids:
return
decided = await store.recover_pending_txns(
conn=conn,
fq_table=fq_table,
bank_ids=bank_ids,
first_seen=self._txn_first_seen,
now=time.monotonic(),
grace_seconds=_TXN_RECOVERY_GRACE_SECONDS,
)
except Exception as e:
logger.warning(f"Cross-store txn recovery failed: {e}")
return
if decided:
logger.info(f"Cross-store txn recovery: decided {decided} undecided txn(s)")
logger.warning(f"Retention sweep failed for {table}: {e}")
# ── consolidation reconcile ──────────────────────────────────────────────
@@ -481,9 +185,7 @@ class MaintenanceLoop:
engine = self._engine
try:
async with acquire_with_retry(engine._backend, max_retries=1) as conn:
rows = await conn.fetch(
f"SELECT schema_name, bank_id FROM {fq_routine('banks_needing_consolidation')}()"
)
rows = await conn.fetch("SELECT schema_name, bank_id FROM public.banks_needing_consolidation()")
except Exception as e:
logger.warning(f"Consolidation reconcile discovery failed: {e}")
return
@@ -542,7 +244,7 @@ class MaintenanceLoop:
Discovery (the set of cron-scheduled models, minus any with an in-flight
refresh) is one cross-tenant round-trip via
``mental_models_with_cron()``. Cron *due-ness* is evaluated here in
``public.mental_models_with_cron()``. Cron *due-ness* is evaluated here in
Python a scheduled fire has elapsed when the most recent cron boundary at
or before now is later than ``last_refreshed_at`` because cron arithmetic
isn't expressible in plain SQL. Each due model is refreshed only when it is
@@ -554,7 +256,7 @@ class MaintenanceLoop:
async with acquire_with_retry(engine._backend, max_retries=1) as conn:
rows = await conn.fetch(
"SELECT schema_name, bank_id, mental_model_id, refresh_cron, last_refreshed_at "
f"FROM {fq_routine('mental_models_with_cron')}()"
"FROM public.mental_models_with_cron()"
)
except Exception as e:
logger.warning(f"Scheduled mental model refresh discovery failed: {e}")
@@ -598,7 +300,6 @@ class MaintenanceLoop:
submitted = 0
skipped_unknown = 0
skipped_fresh = 0
skipped_in_flight = 0
for row in due:
schema = row["schema_name"]
bank_id = row["bank_id"]
@@ -618,8 +319,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,
@@ -630,28 +330,18 @@ class MaintenanceLoop:
if not is_stale:
skipped_fresh += 1
continue
# skip_if_in_flight makes the enqueue itself idempotent. The discovery
# routine already excludes models with a pending/processing refresh,
# but that exclusion is a *read*: this loop runs in every process, so
# every process saw the same "nothing in flight" snapshot and inserted
# its own operation — one queued wave per process (#3210). The insert
# now carries the check, so a second one is never created.
result = await engine.submit_async_refresh_mental_model(
bank_id=bank_id, mental_model_id=mm_id, request_context=context, skip_if_in_flight=True
await engine.submit_async_refresh_mental_model(
bank_id=bank_id, mental_model_id=mm_id, request_context=context
)
if result.get("deduplicated"):
skipped_in_flight += 1
else:
submitted += 1
submitted += 1
except Exception as e:
logger.warning(f"Scheduled mental model refresh failed for {mm_id} in {schema}: {e}")
finally:
_current_schema.reset(token)
if submitted or skipped_unknown or skipped_fresh or skipped_in_flight:
if submitted or skipped_unknown or skipped_fresh:
logger.info(
f"Scheduled mental model refresh: scheduled {submitted} model(s)"
+ (f", {skipped_fresh} up-to-date" if skipped_fresh else "")
+ (f", {skipped_in_flight} already in flight" if skipped_in_flight else "")
+ (f", skipped {skipped_unknown} in unrecognized schema(s)" if skipped_unknown else "")
)
@@ -1,88 +0,0 @@
"""The memories store: which one is installed, and how the engine reaches it.
Resolved through the ordinary extension loader ``HINDSIGHT_API_MEMORIES_EXTENSION``
names a ``module:Class``, and ``HINDSIGHT_API_MEMORIES_*`` becomes its config so
this behaves like every other extension point. Unset (the normal case) means
:class:`~hindsight_api.engine.memories.postgres.PostgresMemories`: rows in
`memory_units`, links in `memory_links` / `unit_entities`, retrieval as SQL.
"""
from __future__ import annotations
import logging
from .base import (
META_CHUNK_ID,
CausalEdgeRecord,
DeletePredicate,
FactRecord,
MemoriesExtension,
MemoryPatch,
RecallArms,
ScanPage,
StoredMemory,
build_fact_records,
build_text_signals,
source_key,
)
logger = logging.getLogger(__name__)
_memories: MemoriesExtension | None = None
def create_memories(context=None) -> MemoriesExtension:
"""Build the configured memories store, or the Postgres default."""
from ...extensions.loader import load_extension
loaded = load_extension("MEMORIES", MemoriesExtension, context=context)
if loaded is not None:
logger.info("[memories] store=%s (memory rows do not go to postgres)", loaded.name)
return loaded
from .postgres import PostgresMemories
return PostgresMemories({})
def get_memories() -> MemoriesExtension:
"""The process-wide memories store, built on first use.
Retrieval and the retain pipeline reach it through call chains that do not
carry the engine, so it is resolved here rather than threaded through every
signature.
"""
global _memories
if _memories is None:
_memories = create_memories()
return _memories
def set_memories(memories: MemoriesExtension | None) -> None:
"""Override the store (tests, and engine startup after initialize())."""
global _memories
_memories = memories
# The graph arm's retriever is chosen from the store and then cached, so it
# has to be re-resolved whenever the store changes.
from ..search.retrieval import set_default_graph_retriever
set_default_graph_retriever(None)
__all__ = [
"META_CHUNK_ID",
"CausalEdgeRecord",
"DeletePredicate",
"FactRecord",
"MemoriesExtension",
"MemoryPatch",
"RecallArms",
"ScanPage",
"StoredMemory",
"build_fact_records",
"build_text_signals",
"create_memories",
"get_memories",
"set_memories",
"source_key",
]
File diff suppressed because it is too large Load Diff
@@ -1,20 +0,0 @@
"""The Postgres memories implementation, split by what calls it.
:class:`~hindsight_api.engine.memories.postgres.PostgresMemories` is a thin class
over these modules; the queries live here, grouped by concern rather than piled
behind one object:
* :mod:`counts` the stats/admin aggregates (freshness, per-doc, timeseries, scopes)
* :mod:`curation` the memory/entity list and detail views
* :mod:`graph` the graph view, entity postings, and the maintenance passes
* :mod:`reads` addressed reads: get, scan, count, tags, consolidation state
* :mod:`writes` inserts, deletes, and observation invalidation
Every function here takes the live connection and Hindsight's ``fq_table``
resolver rather than reaching for globals, so each is callable from a
transaction the caller already owns.
"""
from __future__ import annotations
__all__ = ["counts", "curation", "graph", "reads", "writes"]
@@ -1,168 +0,0 @@
"""The count/aggregate surfaces: consolidation freshness, per-document counts,
ingestion over time, observation scopes.
Each is one ``GROUP BY`` (or filtered ``COUNT``) over `memory_units`. They back
the stats and admin views, not retrieval, so they are grouped here away from the
addressed reads. The SQL is lifted verbatim from the engine methods that used to
carry it; only the connection and ``fq_table`` resolver are now parameters.
"""
from __future__ import annotations
from collections.abc import Callable
from datetime import datetime
from typing import Any
async def consolidation_freshness(*, conn, fq_table: Callable[[str], str], bank_id: str) -> dict[str, Any]:
"""Last consolidation time, the pending / failed fact counts, and the write watermark, in one scan.
``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)
rides along for free. That watermark is what lets a caller decide a mental
model is up to date without running its own scoped scan: nothing in the bank
changed since the refresh, so nothing in the model's scope did either.
"""
row = await conn.fetchrow(
f"""
SELECT
MAX(consolidated_at) AS last_consolidated_at,
MAX(updated_at) AS last_memory_write_at,
COUNT(*) FILTER (
WHERE consolidated_at IS NULL
AND consolidation_failed_at IS NULL
AND fact_type IN ('experience', 'world')
) AS pending,
COUNT(*) FILTER (WHERE consolidation_failed_at IS NOT NULL AND fact_type IN ('experience', 'world')) AS failed
FROM {fq_table("memory_units")}
WHERE bank_id = $1
""",
bank_id,
)
if row is None:
return {"last_consolidated_at": None, "last_memory_write_at": None, "pending": 0, "failed": 0}
return {
"last_consolidated_at": row["last_consolidated_at"],
"last_memory_write_at": row["last_memory_write_at"],
"pending": row["pending"] or 0,
"failed": row["failed"] or 0,
}
async def link_counts(*, conn, fq_table: Callable[[str], str], bank_id: str) -> dict[str, int]:
"""``{link_type: count}`` of live links in a bank.
Non-entity links (temporal / semantic / caused_by) are a single ``GROUP BY`` over
``memory_links``. Entity links are no longer stored there they are derived on demand
from ``unit_entities``, replicating the historical writer cap of ``MAX_LINKS_PER_ENTITY``
bidirectional edges per shared entity so they are aggregated to one ``entity`` scalar.
"""
max_links_per_entity = 10
non_entity_link_rows = await conn.fetch(
f"""
SELECT link_type, COUNT(*) as count
FROM {fq_table("memory_links")}
WHERE bank_id = $1
GROUP BY link_type
""",
bank_id,
)
entity_total_row = await conn.fetchrow(
f"""
WITH per_entity AS (
SELECT ue.entity_id, COUNT(*) AS n
FROM {fq_table("unit_entities")} ue
JOIN {fq_table("memory_units")} mu ON mu.id = ue.unit_id
WHERE mu.bank_id = $1
GROUP BY ue.entity_id
)
SELECT COALESCE(SUM(LEAST(n - 1, $2)), 0)::bigint AS count
FROM per_entity
""",
bank_id,
max_links_per_entity,
)
entity_link_total = int(entity_total_row["count"] or 0) if entity_total_row else 0
counts: dict[str, int] = {row["link_type"]: row["count"] for row in non_entity_link_rows}
if entity_link_total > 0:
counts["entity"] = entity_link_total
return counts
async def document_memory_counts(
*, conn, fq_table: Callable[[str], str], bank_id: str, document_ids: list[str]
) -> dict[str, int]:
"""Live memory count per document id, for the ids given."""
if not document_ids:
return {}
rows = await conn.fetch(
f"""
SELECT document_id, COUNT(*) AS unit_count
FROM {fq_table("memory_units")}
WHERE bank_id = $1 AND document_id = ANY($2::text[])
GROUP BY document_id
""",
bank_id,
list(document_ids),
)
return {row["document_id"]: row["unit_count"] for row in rows}
async def memories_timeseries(
*, conn, fq_table: Callable[[str], str], bank_id: str, time_field: str, trunc: str, since: datetime
) -> list[dict[str, Any]]:
"""Memories bucketed by ``time_field`` (truncated to ``trunc``) and fact_type.
``time_field`` is whitelisted by the caller before it reaches here it is
interpolated into SQL. Event-time fields fall back to ``created_at`` per row so
rows without an event timestamp still appear.
"""
bucket_expr = time_field if time_field == "created_at" else f"COALESCE({time_field}, created_at)"
rows = await conn.fetch(
f"""
SELECT date_trunc('{trunc}', {bucket_expr} AT TIME ZONE 'UTC') AS bucket,
fact_type, COUNT(*) AS count
FROM {fq_table("memory_units")}
WHERE bank_id = $1 AND {bucket_expr} >= $2
GROUP BY bucket, fact_type
ORDER BY bucket
""",
bank_id,
since,
)
return [{"bucket": r["bucket"], "fact_type": r["fact_type"], "count": r["count"]} for r in rows]
async def observation_scope_counts(*, conn, fq_table: Callable[[str], str], bank_id: str) -> list[dict[str, Any]]:
"""Observations grouped by scope (their sorted tag set), most-populous first."""
rows = await conn.fetch(
f"""
SELECT scope, COUNT(*) AS count
FROM (
SELECT COALESCE(ARRAY(SELECT unnest(tags) ORDER BY 1), '{{}}'::text[]) AS scope
FROM {fq_table("memory_units")}
WHERE bank_id = $1 AND fact_type = 'observation'
) s
GROUP BY scope
ORDER BY count DESC, scope
""",
bank_id,
)
return [{"tags": list(r["scope"]), "count": r["count"]} for r in rows]
__all__ = [
"consolidation_freshness",
"document_memory_counts",
"link_counts",
"memories_timeseries",
"observation_scope_counts",
]
@@ -1,507 +0,0 @@
"""Curation reads: the memory list, the memory detail view, and the entity list.
These back the curation UI the table of memories a bank holds, the detail panel
for one of them, and the entity roster beside it. They are paged and filtered
rather than ranked: nothing here scores anything, and nothing walks the corpus.
Two things separate them from the addressed reads in :mod:`reads`. They render
*view* dicts (ISO strings, joined entity names, a ``state`` discriminator) rather
than :class:`~hindsight_api.engine.memories.base.StoredMemory`, because the HTTP
layer serialises what comes back verbatim. And they read the archive as well as
the live table: curation moves an invalidated fact to `invalidated_memory_units`,
so "show me the invalidated ones" is a different table, not a different predicate.
Authentication, operation validation and audit stay with the engine methods that
call these only the queries and their row rendering live here.
"""
from __future__ import annotations
import json
from datetime import datetime
from typing import Any
from ...search.tags import build_tags_where_clause
def _entity_rows_for_units_sql(*, ops, fq_table, unit_ids_placeholder: int) -> str:
"""SQL SELECT producing ``(unit_id, entity_id, canonical_name)`` rows for
the given unit IDs.
Direct rows come from ``unit_entities``. Observations rarely carry
direct rows there; their entity association lives transitively through
their source memories (``source_memory_ids`` on PG, the
``observation_sources`` junction on Oracle). When an observation has
no direct entity rows the SELECT inherits its source memories'
entities, so the result is the same set callers would get from
``get_memory_unit``.
``unit_ids_placeholder`` is the 1-based parameter index that holds the
``uuid[]`` of unit IDs. The placeholder is referenced twice both
sides of the UNION need it so callers should not reuse the slot.
"""
ue = fq_table("unit_entities")
ents = fq_table("entities")
mu = fq_table("memory_units")
p = unit_ids_placeholder
direct = (
f"SELECT ue.unit_id, e.id AS entity_id, e.canonical_name "
f"FROM {ue} ue "
f"JOIN {ents} e ON e.id = ue.entity_id "
f"WHERE ue.unit_id = ANY(${p}::uuid[])"
)
if ops.uses_observation_sources_table:
os_t = fq_table("observation_sources")
inherited = (
f"SELECT os.observation_id AS unit_id, e.id AS entity_id, e.canonical_name "
f"FROM {os_t} os "
f"JOIN {ue} src_ue ON src_ue.unit_id = os.source_id "
f"JOIN {ents} e ON e.id = src_ue.entity_id "
f"WHERE os.observation_id = ANY(${p}::uuid[]) "
f"AND NOT EXISTS (SELECT 1 FROM {ue} d WHERE d.unit_id = os.observation_id)"
)
else:
inherited = (
f"SELECT obs.id AS unit_id, e.id AS entity_id, e.canonical_name "
f"FROM {mu} obs "
f"CROSS JOIN LATERAL unnest(obs.source_memory_ids) AS src_id "
f"JOIN {ue} src_ue ON src_ue.unit_id = src_id "
f"JOIN {ents} e ON e.id = src_ue.entity_id "
f"WHERE obs.id = ANY(${p}::uuid[]) "
f"AND obs.fact_type = 'observation' "
f"AND obs.source_memory_ids IS NOT NULL "
f"AND NOT EXISTS (SELECT 1 FROM {ue} d WHERE d.unit_id = obs.id)"
)
return f"({direct}) UNION ({inherited})"
async def list_memory_units(
*,
conn,
ops,
fq_table,
bank_id: str,
fact_type: str | None = None,
search_query: str | None = None,
consolidation_state: str | None = None,
state: str | None = None,
document_id: str | None = None,
entity_id: str | None = None,
tags: list[str] | None = None,
tags_match: str = "any",
created_before: datetime | None = None,
limit: int = 100,
offset: int = 0,
) -> dict[str, Any]:
"""
List memory units for table view with optional full-text search.
Args:
conn: Open database connection (the caller owns the transaction).
ops: Dialect ops. Unused by this query; part of the interface signature.
fq_table: Table-name resolver.
bank_id: Filter by bank ID
fact_type: Filter by fact type (world, experience)
search_query: Full-text search query (searches text and context fields)
document_id: Optional filter to a single source document.
tags: Optional list of tag names to filter by. When omitted, no tag
filtering is applied (except tags_match='exact', which then selects
the untagged/global scope).
tags_match: How to combine tags (same modes as recall): 'any' (OR,
default) or 'all' (AND) both also include untagged units;
'any_strict'/'all_strict' exclude untagged units; 'exact' matches
units whose tag set equals the given tags exactly.
state: Optional curation-state filter ('valid' or 'invalidated').
Invalidated facts live in a separate archive table; 'invalidated'
reads that archive. Omitted/('valid') lists live facts.
consolidation_state: Optional filter on consolidation state. One of
'failed' (consolidation permanently failed and awaiting recovery),
'pending' (not yet consolidated, no failure), or
'done' (successfully consolidated). Only applies to source memory
types (world/experience).
limit: Maximum number of results to return
offset: Offset for pagination
Returns:
Dict with items (list of memory units) and total count
"""
if state is not None and state not in ("valid", "invalidated"):
raise ValueError(f"Invalid state '{state}': expected 'valid' or 'invalidated'.")
if entity_id is not None:
import uuid as _uuid
try:
_uuid.UUID(entity_id)
except ValueError:
raise ValueError(f"Invalid entity_id: '{entity_id}' is not a valid UUID") from None
# Invalidated facts live in a separate archive table; pick the source
# accordingly. Default (state is None) lists live facts.
is_archived = state == "invalidated"
source_table = fq_table("invalidated_memory_units") if is_archived else fq_table("memory_units")
# Build query conditions
query_conditions = []
query_params = []
param_count = 0
if bank_id:
param_count += 1
query_conditions.append(f"bank_id = ${param_count}")
query_params.append(bank_id)
if fact_type:
param_count += 1
query_conditions.append(f"fact_type = ${param_count}")
query_params.append(fact_type)
if document_id:
param_count += 1
query_conditions.append(f"document_id = ${param_count}")
query_params.append(document_id)
if entity_id:
# Reverse lookup via the stored entity links. Entity links reference live memory units, so
# this yields nothing against the invalidated archive (documented on the method).
param_count += 1
query_conditions.append(
f"id IN (SELECT unit_id FROM {fq_table('unit_entities')} WHERE entity_id = ${param_count}::uuid)"
)
query_params.append(entity_id)
if search_query:
# Full-text search on text and context fields using ILIKE
param_count += 1
query_conditions.append(f"(text ILIKE ${param_count} OR context ILIKE ${param_count})")
query_params.append(f"%{search_query}%")
if consolidation_state:
# Named apart from `state`, which the engine method used to shadow here;
# `is_archived` was already resolved above, so behaviour is unchanged.
wanted = consolidation_state.lower()
if wanted == "failed":
query_conditions.append("consolidation_failed_at IS NOT NULL AND fact_type IN ('experience', 'world')")
elif wanted == "pending":
query_conditions.append(
"consolidated_at IS NULL AND consolidation_failed_at IS NULL AND fact_type IN ('experience', 'world')"
)
elif wanted == "done":
query_conditions.append("consolidated_at IS NOT NULL AND fact_type IN ('experience', 'world')")
else:
raise ValueError(
f"Invalid consolidation_state '{consolidation_state}': expected 'failed', 'pending', or 'done'."
)
if tags:
tags_clause, tags_params, next_param = build_tags_where_clause(tags, param_count + 1, "", tags_match)
if tags_clause:
query_conditions.append(tags_clause.removeprefix("AND "))
query_params.extend(tags_params)
param_count = next_param - 1
elif tags_match == "exact":
# Exact match with no tags is the "global" scope: rows that carry no
# tags at all. (Other match modes treat empty tags as "no filter".)
query_conditions.append("(tags IS NULL OR tags = '{}')")
if created_before is not None:
param_count += 1
query_conditions.append(f"created_at < ${param_count}")
query_params.append(created_before)
where_clause = "WHERE " + " AND ".join(query_conditions) if query_conditions else ""
# Get total count
count_query = f"""
SELECT COUNT(*) as total
FROM {source_table}
{where_clause}
"""
count_result = await conn.fetchrow(count_query, *query_params)
total = count_result["total"]
# Get units with limit and offset
param_count += 1
limit_param = f"${param_count}"
query_params.append(limit)
param_count += 1
offset_param = f"${param_count}"
query_params.append(offset)
# The archive carries invalidation bookkeeping; the live table doesn't.
curation_cols = (
"invalidation_reason, invalidated_at"
if is_archived
else "NULL::text AS invalidation_reason, NULL::timestamptz AS invalidated_at"
)
units = await conn.fetch(
f"""
SELECT id, text, event_date, context, fact_type, document_id,
mentioned_at, occurred_start, occurred_end, chunk_id, proof_count,
tags, metadata, consolidated_at, consolidation_failed_at, edited_at, {curation_cols}
FROM {source_table}
{where_clause}
ORDER BY mentioned_at DESC NULLS LAST, created_at DESC
LIMIT {limit_param} OFFSET {offset_param}
""",
*query_params,
)
# Get entity information for these units
if units:
unit_ids = [row["id"] for row in units]
unit_entities = await conn.fetch(
f"""
SELECT ue.unit_id, e.canonical_name
FROM {fq_table("unit_entities")} ue
JOIN {fq_table("entities")} e ON ue.entity_id = e.id
WHERE ue.unit_id = ANY($1::uuid[])
ORDER BY ue.unit_id
""",
unit_ids,
)
else:
unit_entities = []
# Build entity mapping
entity_map: dict[Any, list[str]] = {}
for row in unit_entities:
unit_id = row["unit_id"]
entity_name = row["canonical_name"]
if unit_id not in entity_map:
entity_map[unit_id] = []
entity_map[unit_id].append(entity_name)
# Build result items
items = []
for row in units:
unit_id = row["id"]
entities = entity_map.get(unit_id, [])
items.append(
{
"id": str(unit_id),
"text": row["text"],
"context": row["context"] if row["context"] else "",
"date": row["event_date"].isoformat() if row["event_date"] else "",
"fact_type": row["fact_type"],
"document_id": row["document_id"],
"mentioned_at": row["mentioned_at"].isoformat() if row["mentioned_at"] else None,
"occurred_start": row["occurred_start"].isoformat() if row["occurred_start"] else None,
"occurred_end": row["occurred_end"].isoformat() if row["occurred_end"] else None,
"entities": ", ".join(entities) if entities else "",
"chunk_id": row["chunk_id"] if row["chunk_id"] else None,
"proof_count": row["proof_count"] if row["proof_count"] is not None else 1,
"tags": list(row["tags"]) if row["tags"] else [],
"metadata": conn.parse_json(row["metadata"]) if row["metadata"] is not None else {},
"consolidated_at": row["consolidated_at"].isoformat() if row["consolidated_at"] else None,
"consolidation_failed_at": (
row["consolidation_failed_at"].isoformat() if row["consolidation_failed_at"] else None
),
"state": "invalidated" if is_archived else "valid",
"invalidation_reason": row["invalidation_reason"],
"invalidated_at": row["invalidated_at"].isoformat() if row["invalidated_at"] else None,
"edited_at": row["edited_at"].isoformat() if row["edited_at"] else None,
}
)
return {"items": items, "total": total, "limit": limit, "offset": offset}
async def get_memory_unit(*, conn, ops, fq_table, bank_id: str, unit_id: str) -> dict[str, Any] | None:
"""
Get a single memory unit by ID.
Args:
conn: Open database connection (the caller owns the transaction).
ops: Dialect ops, for the observationsource entity inheritance shape.
fq_table: Table-name resolver.
bank_id: Bank ID
unit_id: Memory unit ID (the caller validates it is a UUID)
Returns:
Dict with memory unit data or None if not found
"""
# Get the memory unit (include source_memory_ids for mental models).
# Curation moves invalidated facts to invalidated_memory_units, so fall
# back to the archive (with its invalidation bookkeeping) on a miss.
select_cols = (
"id, text, context, event_date, occurred_start, occurred_end, "
"mentioned_at, fact_type, document_id, chunk_id, tags, metadata, source_memory_ids, "
"observation_scopes, edited_at"
)
row = await conn.fetchrow(
f"SELECT {select_cols}, NULL::text AS invalidation_reason, NULL::timestamptz AS invalidated_at "
f"FROM {fq_table('memory_units')} WHERE id = $1 AND bank_id = $2",
unit_id,
bank_id,
)
unit_state = "valid"
if not row:
row = await conn.fetchrow(
f"SELECT {select_cols}, invalidation_reason, invalidated_at "
f"FROM {fq_table('invalidated_memory_units')} WHERE id = $1 AND bank_id = $2",
unit_id,
bank_id,
)
unit_state = "invalidated"
if not row:
return None
# Get entity information. _entity_rows_for_units_sql handles the
# observation→source_memory_ids inheritance fallback in SQL, so a
# single query covers direct rows and inherited ones.
entities_rows = await conn.fetch(
_entity_rows_for_units_sql(ops=ops, fq_table=fq_table, unit_ids_placeholder=1),
[row["id"]],
)
entities = [r["canonical_name"] for r in entities_rows]
result: dict[str, Any] = {
"id": str(row["id"]),
"text": row["text"],
"context": row["context"] if row["context"] else "",
"date": row["event_date"].isoformat() if row["event_date"] else "",
"type": row["fact_type"],
"mentioned_at": row["mentioned_at"].isoformat() if row["mentioned_at"] else None,
"occurred_start": row["occurred_start"].isoformat() if row["occurred_start"] else None,
"occurred_end": row["occurred_end"].isoformat() if row["occurred_end"] else None,
"entities": entities,
"document_id": row["document_id"] if row["document_id"] else None,
"chunk_id": str(row["chunk_id"]) if row["chunk_id"] else None,
"tags": row["tags"] if row["tags"] else [],
"metadata": conn.parse_json(row["metadata"]) if row["metadata"] is not None else {},
"observation_scopes": (
conn.parse_json(row["observation_scopes"]) if row["observation_scopes"] is not None else None
),
"state": unit_state,
"invalidation_reason": row["invalidation_reason"],
"invalidated_at": row["invalidated_at"].isoformat() if row["invalidated_at"] else None,
"edited_at": row["edited_at"].isoformat() if row["edited_at"] else None,
}
# For observations, include source_memory_ids
# history is deprecated here - use GET /memories/{id}/history instead
if row["fact_type"] == "observation":
result["history"] = []
if row["fact_type"] == "observation" and row["source_memory_ids"]:
source_ids = row["source_memory_ids"]
result["source_memory_ids"] = [str(sid) for sid in source_ids]
# Fetch source memories
source_rows = await conn.fetch(
f"""
SELECT id, text, fact_type, context, occurred_start, mentioned_at
FROM {fq_table("memory_units")}
WHERE id = ANY($1::uuid[])
ORDER BY mentioned_at DESC NULLS LAST
""",
source_ids,
)
result["source_memories"] = [
{
"id": str(r["id"]),
"text": r["text"],
"type": r["fact_type"],
"context": r["context"],
"occurred_start": r["occurred_start"].isoformat() if r["occurred_start"] else None,
"mentioned_at": r["mentioned_at"].isoformat() if r["mentioned_at"] else None,
}
for r in source_rows
]
return result
async def list_entities(
*,
conn,
fq_table,
bank_id: str,
search: str | None = None,
limit: int = 100,
offset: int = 0,
) -> dict[str, Any]:
"""
List all entities for a bank with pagination.
Args:
conn: Open database connection (the caller owns the transaction).
fq_table: Table-name resolver.
bank_id: bank IDentifier
search: Optional case-insensitive substring match on canonical_name.
limit: Maximum number of entities to return
offset: Offset for pagination
Returns:
Dict with items, total, limit, offset
"""
conditions = ["bank_id = $1"]
params: list[Any] = [bank_id]
if search:
# Substring match, same ILIKE shape entity lookup uses elsewhere. Applied
# to the count too, so the UI pages over the filtered set.
params.append(f"%{search}%")
conditions.append(f"canonical_name ILIKE ${len(params)}")
where_clause = " AND ".join(conditions)
# Get total count
total_row = await conn.fetchrow(
f"""
SELECT COUNT(*) as total
FROM {fq_table("entities")}
WHERE {where_clause}
""",
*params,
)
total = total_row["total"] if total_row else 0
# Get paginated entities
rows = await conn.fetch(
f"""
SELECT id, canonical_name, mention_count, first_seen, last_seen, metadata
FROM {fq_table("entities")}
WHERE {where_clause}
ORDER BY mention_count DESC, last_seen DESC, id ASC
LIMIT ${len(params) + 1} OFFSET ${len(params) + 2}
""",
*params,
limit,
offset,
)
entities = []
for row in rows:
# Handle metadata - may be dict, JSON string, or None
metadata = row["metadata"]
if metadata is None:
metadata = {}
elif isinstance(metadata, str):
try:
metadata = json.loads(metadata)
except json.JSONDecodeError:
metadata = {}
entities.append(
{
"id": str(row["id"]),
"canonical_name": row["canonical_name"],
"mention_count": row["mention_count"],
"first_seen": row["first_seen"].isoformat() if row["first_seen"] else None,
"last_seen": row["last_seen"].isoformat() if row["last_seen"] else None,
"metadata": metadata,
}
)
return {
"items": entities,
"total": total,
"limit": limit,
"offset": offset,
}
__all__ = ["get_memory_unit", "list_entities", "list_memory_units"]

Some files were not shown because too many files have changed in this diff Show More