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
1627 changed files with 25667 additions and 148489 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"
}
]
}
-25
View File
@@ -78,11 +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).
### 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.
@@ -159,18 +154,6 @@ 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`):
@@ -221,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:
+1 -140
View File
@@ -2,7 +2,7 @@
# 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
@@ -21,25 +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
# 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
@@ -57,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
@@ -84,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
@@ -114,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
@@ -139,13 +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_ENTITY_TRGM_SIMILARITY_THRESHOLD=0.15 # Postgres pg_trgm.similarity_threshold applied on every pool connection, used by entity resolution's % trigram match. Must be in (0, 1]. Lower catches more substring-ish matches at higher CPU cost on large entity sets; higher is stricter and cheaper.
# HINDSIGHT_API_ENTITY_INTRABATCH_MERGE_SIMILARITY=0.5 # Trigram similarity (pg_trgm-equivalent, computed in-memory) at/above which two new names created by the SAME retain are merged into one entity (in-batch dedup of surface-form variants). Must be in (0, 1]. A merge cutoff, stricter than the recall threshold above; raise toward 1.0 to merge only near-identical forms.
# HINDSIGHT_API_RETAIN_ENTITY_RESOLUTION_MAX_CANDIDATES=200 # Max candidates scored per entity mention during retain. The fuzzy lookup keeps only this many best matches per name (ranked by trigram/Jaro-Winkler similarity) before scoring them one by one. On banks holding thousands of near-identical names an uncapped set turns one retain into minutes of CPU that stall the worker's health checks. Raise only if entities that should merge are being duplicated.
# HINDSIGHT_API_MIGRATION_CONCURRENCY=1 # Tenant schemas to migrate concurrently (PG only, each in its own process; per-schema work stays sequential). Each worker has ~1-2s startup cost + uses ~3 DB connections, so it only pays off with many schemas (tens+) or slow migrations; keep concurrency*3 <= spare max_connections. Default: 1 (sequential).
# HINDSIGHT_API_OPERATION_RETENTION_DAYS=30 # Prune terminal operation rows, payloads, and metadata after this many days; 0 (the default) keeps them forever.
# HINDSIGHT_API_OPERATION_CLEANUP_BATCH_SIZE=1000 # Maximum expired terminal rows deleted per tenant schema in each cleanup cycle; must be positive.
# Vector Extension (Optional - uses pgvector by default)
# Options: "pgvector" (default), "vchord", "pgvectorscale" (DiskANN)
@@ -165,11 +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=
# Optional cap on the number of terms in the native PostgreSQL BM25 tsquery.
# Long queries OR-join every normalized token, which can match too much of a
# large bank. 0 (default) keeps the historical uncapped behavior; a positive
# value bounds only the native backend (other BM25 backends get the raw query).
# HINDSIGHT_API_BM25_MAX_QUERY_TERMS=0
# File Parser (Optional - uses markitdown by default)
# HINDSIGHT_API_FILE_PARSER=markitdown
@@ -182,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
@@ -210,11 +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
# For TEI provider:
# HINDSIGHT_API_EMBEDDINGS_TEI_URL=http://localhost:8080
# For OpenAI-compatible embeddings:
@@ -236,58 +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
# 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)
@@ -307,33 +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
# -----------------------------------------------------------------------------
# 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: 40 KiB

File diff suppressed because it is too large Load Diff
-32
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
-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
+17 -208
View File
@@ -36,14 +36,11 @@ jobs:
integrations-composio: ${{ steps.filter.outputs.integrations-composio }}
integrations-chat: ${{ steps.filter.outputs.integrations-chat }}
integrations-claude-code: ${{ steps.filter.outputs.integrations-claude-code }}
integrations-coding-agents: ${{ steps.filter.outputs.integrations-coding-agents }}
integrations-cline: ${{ steps.filter.outputs.integrations-cline }}
integrations-codex: ${{ steps.filter.outputs.integrations-codex }}
integrations-github-copilot: ${{ steps.filter.outputs.integrations-github-copilot }}
integrations-continue: ${{ steps.filter.outputs.integrations-continue }}
integrations-cursor-cli: ${{ steps.filter.outputs.integrations-cursor-cli }}
integrations-zcode: ${{ steps.filter.outputs.integrations-zcode }}
integrations-copilot-cli: ${{ steps.filter.outputs.integrations-copilot-cli }}
integrations-crewai: ${{ steps.filter.outputs.integrations-crewai }}
integrations-litellm: ${{ steps.filter.outputs.integrations-litellm }}
integrations-pydantic-ai: ${{ steps.filter.outputs.integrations-pydantic-ai }}
@@ -145,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:
@@ -157,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:
@@ -187,8 +180,6 @@ jobs:
- 'hindsight-integrations/cursor/**'
integrations-zed:
- 'hindsight-integrations/zed/**'
integrations-zcode:
- 'hindsight-integrations/zcode/**'
integrations-n8n:
- 'hindsight-integrations/n8n/**'
integrations-zapier:
@@ -292,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: >-
@@ -475,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: >-
@@ -577,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]
@@ -754,43 +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
build-ai-sdk-integration:
needs: [detect-changes]
if: >-
@@ -1318,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
@@ -1925,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
@@ -2306,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
@@ -2487,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
@@ -3633,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: >-
@@ -4960,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
@@ -5099,7 +4909,6 @@ jobs:
- test-github-copilot-integration
- test-codex-integration
- test-cursor-cli-integration
- test-zcode-integration
- build-ai-sdk-integration
- test-ai-sdk-integration-deno
- test-opencode-integration
-8
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
@@ -48,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

@@ -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'
+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
@@ -1,341 +0,0 @@
# v2 Knowledge Pages — Implementation Plan
> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking.
**Goal:** Make v2 knowledge pages a reliable, cleanly-tiered "wiki" surface: passive `entity_labels` tier-tagging + tag-scoped seeded pages, a `hindsight_*` MCP surface with one active `capture_initiative` verb that creates per-initiative pages linked by tag, and SessionStart/UserPromptSubmit page-roster injection.
**Architecture:** Shared TS core (`hindsight-integrations/hindsight-coding-agents`). Extraction stays blind to "pages"; classification is intrinsic (`knowledge:<tier>` tags), pages are tag-scoped saved views. Per-initiative navigation via a `relatedPageId:<id>` tag the synthesizer renders into `[[page:<id>]]`, with the Initiatives folder/roster as the guaranteed fallback.
**Tech Stack:** TypeScript, vitest, tsup bundling. Hindsight REST API (`/knowledge-base/*`, `/mental-models`, `/memories`, bank `/config`).
**Spec:** `docs/superpowers/specs/2026-07-25-v2-knowledge-pages-design.md`
**Working dir for all commands:** `hindsight-integrations/hindsight-coding-agents`
**Test command:** `npx vitest run <file>` (fast suite; excludes `*.live.test.ts`). Full check: `npx vitest run && npx tsc --noEmit`.
**Conventions to follow (existing patterns):**
- `HindsightClient` HTTP via `this.req("METHOD", this.bankUrl(path), body?)`; JSON via `await r.json()`.
- MCP tools are SDK-free `ToolSpec { name, description, inputSchema (ZodRawShape), handler }`; wrap handler bodies in `guarded(...)`; `ok(value)` / `err(e)` result helpers.
- Fail-open everywhere in hooks; pure logic separated from stdin/stdout plumbing.
- Do NOT add a Claude co-author trailer to any commit.
---
## Task 1: Config field `pageRefreshEveryTurns`
**Files:**
- Modify: `src/core/config.ts`
- Test: `src/core/config.test.ts`
- [ ] **Step 1: Write failing test** — assert the default resolves to 10 and an override wins.
```ts
it("pageRefreshEveryTurns defaults to 10 and is overridable", () => {
expect(loadConfig({ harness: "claude-code", projectDir: process.cwd() }).pageRefreshEveryTurns).toBe(10);
});
```
(Add an override case mirroring the existing override tests in this file.)
- [ ] **Step 2: Run** `npx vitest run src/core/config.test.ts` → FAIL (property missing).
- [ ] **Step 3: Implement** — add `pageRefreshEveryTurns: number` to the `Config` type and default `10` in the same place `recallMaxTokens`/`reflectTimeoutMs` are defined/merged. Follow the exact merge/layering pattern already used for numeric fields.
- [ ] **Step 4: Run** the test → PASS.
- [ ] **Step 5: Commit** `git add src/core/config.ts src/core/config.test.ts && git commit -m "feat(core): add pageRefreshEveryTurns config (default 10)"`
---
## Task 2: `knowledge-injection.ts` — roster/preamble formatting (pure, new)
**Files:**
- Create: `src/core/knowledge-injection.ts`
- Test: `src/core/knowledge-injection.test.ts`
Pure, SDK-free, no network. Parses the `listPages()` payload and formats the two injections.
- [ ] **Step 1: Write failing tests**
```ts
import { describe, expect, it } from "vitest";
import { parsePageList, buildKnowledgePreamble, buildRosterRefresh } from "./knowledge-injection";
describe("parsePageList", () => {
it("extracts {id,title} from the mental-model list shape, tolerating junk", () => {
const raw = { items: [{ id: "p1", name: "Component map" }, { id: "p2", name: "Core concepts" }, { nope: 1 }] };
expect(parsePageList(raw)).toEqual([{ id: "p1", title: "Component map" }, { id: "p2", title: "Core concepts" }]);
});
it("returns [] for null/garbage", () => {
expect(parsePageList(null)).toEqual([]);
expect(parsePageList(42 as unknown)).toEqual([]);
});
});
describe("buildKnowledgePreamble", () => {
it("includes guidance, a roster of pages, and a refresh note", () => {
const out = buildKnowledgePreamble([{ id: "p1", title: "Component map" }]);
expect(out).toContain("<hindsight_knowledge>");
expect(out).toContain("Component map");
expect(out).toContain("p1");
expect(out).toMatch(/hindsight_read_knowledge_page/);
});
it("has an empty-state line when there are no pages", () => {
const out = buildKnowledgePreamble([]);
expect(out).toMatch(/no knowledge pages yet|still learning/i);
});
});
describe("buildRosterRefresh", () => {
it("is a compact 'current pages' block listing ids+titles", () => {
const out = buildRosterRefresh([{ id: "p1", title: "Component map" }]);
expect(out).toContain("Component map");
expect(out).toContain("p1");
});
it("returns undefined when there are no pages (nothing to refresh)", () => {
expect(buildRosterRefresh([])).toBeUndefined();
});
});
```
- [ ] **Step 2: Run** `npx vitest run src/core/knowledge-injection.test.ts` → FAIL.
- [ ] **Step 3: Implement**
```ts
export interface PageRef { id: string; title: string; }
/** Defensive parse of HindsightClient.listPages() (GET /mental-models?detail=metadata → {items:[{id,name}]}). */
export function parsePageList(raw: unknown): PageRef[] {
const items = (raw as { items?: unknown })?.items;
if (!Array.isArray(items)) return [];
const out: PageRef[] = [];
for (const it of items) {
const id = (it as { id?: unknown })?.id;
const name = (it as { name?: unknown })?.name;
if (typeof id === "string" && typeof name === "string") out.push({ id, title: name });
}
return out;
}
function roster(pages: PageRef[]): string {
return pages.map((p) => `- ${p.title} (${p.id})`).join("\n");
}
/** SessionStart: teach when/why to use pages + list what exists. Empty-state aware. */
export function buildKnowledgePreamble(pages: PageRef[]): string {
const body = pages.length
? `Knowledge pages available in this repository:\n${roster(pages)}`
: "No knowledge pages yet — Hindsight is still learning this repo; they'll appear as it processes.";
return (
"<hindsight_knowledge>\n" +
"This repository has a Hindsight knowledge base: curated, continuously-updated pages summarizing its " +
"durable engineering knowledge (architecture, components, conventions, key decisions, and in-flight initiatives).\n" +
"Before substantial work, consult the relevant pages instead of re-deriving understanding from the code: read " +
"Conventions before writing new code, the Component map before changing a subsystem, and an initiative's page " +
"before continuing that feature.\n" +
`${body}\n` +
"Read one with hindsight_read_knowledge_page(page_id). Follow any [[page:<id>]] links you see. The list is " +
"re-injected for you periodically as it changes.\n" +
"</hindsight_knowledge>"
);
}
/** Periodic UserPromptSubmit refresh — compact, or undefined when there's nothing to show. */
export function buildRosterRefresh(pages: PageRef[]): string | undefined {
if (!pages.length) return undefined;
return (
"<hindsight_knowledge_refresh>\n" +
`Current Hindsight knowledge pages (may have changed):\n${roster(pages)}\n` +
"Read any with hindsight_read_knowledge_page(page_id).\n" +
"</hindsight_knowledge_refresh>"
);
}
```
- [ ] **Step 4: Run** the test → PASS.
- [ ] **Step 5: Commit** `git add src/core/knowledge-injection.ts src/core/knowledge-injection.test.ts && git commit -m "feat(core): knowledge-injection roster/preamble formatting"`
---
## Task 3: `entity_labels` tier vocabulary + configureBank wiring
**Files:**
- Modify: `src/core/missions.ts` (add `KNOWLEDGE_LABELS`)
- Modify: `src/core/hindsight.ts` (`configureBank` PATCH sets `entity_labels`)
- Test: `src/core/hindsight.*.test.ts` (add/extend a config test with a mock client)
- [ ] **Step 1: Write failing test** — assert `configureBank` PATCHes `/config` with `entity_labels` containing the `knowledge` group and its five values, and `entities_allow_free_form: true`. Use the existing fetch/req mock pattern from `hindsight.*.test.ts`; capture the PATCH body to `/config` and assert on it.
- [ ] **Step 2: Run** → FAIL.
- [ ] **Step 3: Implement**
- In `missions.ts`, export `KNOWLEDGE_LABELS` — the exact object from the spec §4 (`key:"knowledge"`, `type:"multi-values"`, `optional:true`, `tag:true`, the verbose group `description`, and the five value `{value,description}` entries: feature-work, decision, convention, component, concept).
- In `hindsight.ts::configureBank`, extend the existing `PATCH .../config` `updates` object to include `entity_labels: [KNOWLEDGE_LABELS]` and `entities_allow_free_form: true`. Import `KNOWLEDGE_LABELS`.
- Update the `[bank] configured …` log to mention `entity_labels`.
- [ ] **Step 4: Run** → PASS.
- [ ] **Step 5: Commit** `git add src/core/missions.ts src/core/hindsight.ts src/core/hindsight.*.test.ts && git commit -m "feat(core): passive knowledge entity_labels tier vocabulary + configureBank wiring"`
---
## Task 4: Tag-scoped seeded pages + Initiatives folder + link source_query
**Files:**
- Modify: `src/core/missions.ts` (`PAGES` gain `tags`; Initiatives `source_query` link instruction)
- Modify: `src/core/hindsight.ts` (`ensureFolder`, `createPages` sets page `tags` + parents Initiatives under the folder)
- Test: `src/core/hindsight.pages.test.ts`
- [ ] **Step 1: Write failing tests** (mock client `req`):
- Each seeded page POST to `/knowledge-base/pages` includes `tags: ["knowledge:<tier>"]` mapped per the spec §5 table.
- The Initiatives page is created with `parent_id` equal to the id returned by an Initiatives folder POST to `/knowledge-base/folders`.
- `ensureFolder("Initiatives")` returns an existing root folder's id when the tree already contains it (GET `/knowledge-base/tree`) and does NOT POST a duplicate.
- [ ] **Step 2: Run** → FAIL.
- [ ] **Step 3: Implement**
- `missions.ts`: add `tags: string[]` to each `PAGES` entry (feature-work/decision/convention/component/concept mapping). Append to the Initiatives `source_query`: *"When a source memory carries a tag of the form `relatedPageId:<id>`, include a Markdown link `[[page:<id>]]` to that page in the summary, so each initiative links to its detailed page."*
- `hindsight.ts`: add
```ts
/** Find a root folder by name (case-insensitive) or create it; returns its id. */
async ensureFolder(name: string): Promise<string | undefined> {
try {
const tree = (await (await this.req("GET", this.bankUrl("/knowledge-base/tree"))).json()) as
{ roots?: { id?: string; kind?: string; name?: string }[] };
const hit = (tree.roots || []).find((n) => n.kind === "folder" && (n.name || "").toLowerCase() === name.toLowerCase());
if (hit?.id) return hit.id;
} catch { /* fall through to create */ }
try {
const r = await this.req("POST", this.bankUrl("/knowledge-base/folders"), { name });
return ((await r.json()) as { id?: string }).id;
} catch { return undefined; }
}
```
- In `createPages()`: before the loop, `const initiativesFolderId = await this.ensureFolder("Initiatives");`. For each page, build body `{ name, source_query, tags: p.tags, parent_id: <initiativesFolderId if this is the Initiatives page else undefined>, trigger: { fact_types:[...], refresh_after_consolidation:true } }`. (Page-level `tags` drives synthesis scoping via `RefreshTagFiltering`; `tags_match` defaults to `all_strict` when tags present.)
- [ ] **Step 4: Run** → PASS.
- [ ] **Step 5: Commit** `git add src/core/missions.ts src/core/hindsight.ts src/core/hindsight.pages.test.ts && git commit -m "feat(core): tag-scope seeded pages, Initiatives folder, relatedPageId link source_query"`
---
## Task 5: Client helpers — per-initiative page + marker retain
**Files:**
- Modify: `src/core/hindsight.ts` (`captureInitiative`)
- Test: `src/core/hindsight.pages.test.ts`
- [ ] **Step 1: Write failing tests** (mock `req`):
- `captureInitiative({title:"Retry backoff for the uploader", summary:"…"})` → derives slug `retry-backoff-for-the-uploader`, POSTs a page id `initiative-<slug>` to `/knowledge-base/pages` with `parent_id` = the Initiatives folder and `tags: ["knowledge:feature-work"]`, AND POSTs a marker to `/memories` (via `retain`) tagged `["knowledge:feature-work","relatedPageId:initiative-<slug>"]`, strategy `session` or `document` (pick `document`), `async:true`. Returns `{ page_id: "initiative-<slug>" }`.
- Slug is deterministic and identical between the page id and the `relatedPageId:` tag value.
- Enhancement path: `captureInitiative({title, summary, relatesToPageId:"initiative-x"})` POSTs NO new page; marker tagged `relatedPageId:initiative-x`; returns `{ page_id: "initiative-x" }`.
- [ ] **Step 2: Run** → FAIL.
- [ ] **Step 3: Implement**
```ts
private slugify(s: string): string {
return s.toLowerCase().replace(/[^a-z0-9]+/g, "-").replace(/^-+|-+$/g, "").slice(0, 60) || "initiative";
}
/** Active-path capture: register a major feature as a per-initiative page + a tagged marker memory. */
async captureInitiative(args: { title: string; summary: string; relatesToPageId?: string }): Promise<{ page_id: string }> {
const pageId = args.relatesToPageId ?? `initiative-${this.slugify(args.title)}`;
if (!args.relatesToPageId) {
const folderId = await this.ensureFolder("Initiatives");
await this.req("POST", this.bankUrl("/knowledge-base/pages"), {
name: args.title,
source_query: `Summarize the "${args.title}" initiative: what is being built or changed and why, and its current state — drawn from the project's memory.`,
parent_id: folderId,
tags: ["knowledge:feature-work", `relatedPageId:${pageId}`],
trigger: { fact_types: ["world", "experience", "observation"], refresh_after_consolidation: true },
});
}
const verb = args.relatesToPageId ? "Enhancement to an existing initiative" : "New initiative";
const content = `${verb}: ${args.title}. ${args.summary}`;
await this.retain(content, "initiative marker", pageId /* not a stable doc id requirement; see note */,
["knowledge:feature-work", `relatedPageId:${pageId}`], "document", { async: true });
return { page_id: pageId };
}
```
- NOTE: use a UNIQUE document id per marker (e.g. `initiative-marker-<slug>-<n>`), NOT `pageId`, so repeated enhancement captures accrue instead of replacing. Since `Date.now()` is fine here (runtime, not a workflow script), suffix with a timestamp: `initiative-marker-${this.slugify(args.title)}-${Date.now()}`. Keep the `relatedPageId` tag equal to `pageId`.
- Confirm `retain(content, context, documentId, tags, strategy, opts)` signature matches current `HindsightClient.retain`.
- [ ] **Step 4: Run** → PASS.
- [ ] **Step 5: Commit** `git add src/core/hindsight.ts src/core/hindsight.pages.test.ts && git commit -m "feat(core): captureInitiative — per-initiative page + relatedPageId marker"`
---
## Task 6: MCP surface — `hindsight_*` grounding + `capture_initiative`; drop page CRUD
**Files:**
- Modify: `src/core/knowledge-tools.ts`
- Modify: `src/mcp-server.ts` (only if it references removed tool names)
- Test: `src/core/knowledge-tools.test.ts`, `src/mcp-server.test.ts` (tool-count assertions)
- [ ] **Step 1: Write failing tests**
- `buildKnowledgeTools(client, bankId)` returns exactly these tool names: `hindsight_get_current_bank`, `hindsight_list_knowledge_pages`, `hindsight_read_knowledge_page`, `hindsight_search_memory`, `hindsight_capture_initiative`, `hindsight_ingest_document`. (Assert the set; update any count assertion.)
- `hindsight_capture_initiative` handler calls `client.captureInitiative` with `{title, summary, relatesToPageId?}` and returns the page id (mock client).
- No `create_page` / `update_page` / `delete_page` tools are present.
- Each tool still fails closed via `guarded` (a thrown client error → `isError:true`, no throw).
- [ ] **Step 2: Run** → FAIL.
- [ ] **Step 3: Implement**
- Rebuild the `buildKnowledgeTools` list: rename read/recall/ingest/bank tools to the `hindsight_*` names; drop `create_page`/`update_page`/`delete_page`; add `hindsight_capture_initiative` with `inputSchema { title: z.string(), summary: z.string(), relates_to_page_id: z.string().optional() }` calling `client.captureInitiative({ title, summary, relatesToPageId: relates_to_page_id })`.
- Use the **verbatim agent-facing `description` strings** from the spec §6 / the brainstorm (grounding tools + the explicit WHEN/WHEN-NOT `capture_initiative` description).
- Update `mcp-server.ts` only if it enumerates tool names; otherwise it consumes `buildKnowledgeTools` generically and needs no change.
- [ ] **Step 4: Run** `npx vitest run src/core/knowledge-tools.test.ts src/mcp-server.test.ts` → PASS.
- [ ] **Step 5: Commit** `git add src/core/knowledge-tools.ts src/mcp-server.ts src/core/knowledge-tools.test.ts src/mcp-server.test.ts && git commit -m "feat(mcp): hindsight_* grounding tools + capture_initiative; remove raw page CRUD from agent"`
---
## Task 7: SessionStart — preamble + roster
**Files:**
- Modify: `src/core/session-start.ts`
- Test: `src/core/session-start.test.ts`
- [ ] **Step 1: Write failing tests**
- `buildSessionStartContext` now fetches pages via the client and injects `buildKnowledgePreamble(...)` instead of the static `KNOWLEDGE_MISSION`. Extend the `SeedContextClient` interface with `listPages(): Promise<unknown>`; the mock returns `{items:[{id:"p1",name:"Component map"}]}` and the output contains "Component map".
- listPages failure is fail-open: the preamble still renders (empty-state) and the seed logic is unaffected.
- [ ] **Step 2: Run** → FAIL.
- [ ] **Step 3: Implement**
- Add `listPages` to `SeedContextClient`.
- Replace the `parts.push(KNOWLEDGE_MISSION)` line with: fetch `const pages = parsePageList(await client.listPages().catch(() => null));` then `parts.push(buildKnowledgePreamble(pages));`. Import from `./knowledge-injection`.
- Remove the now-unused `KNOWLEDGE_MISSION` export if nothing else references it (grep first; keep if referenced).
- [ ] **Step 4: Run** → PASS.
- [ ] **Step 5: Commit** `git add src/core/session-start.ts src/core/session-start.test.ts && git commit -m "feat(core): SessionStart injects page roster + guidance preamble"`
---
## Task 8: UserPromptSubmit — hook-counted periodic roster refresh
**Files:**
- Modify: `src/core/hook.ts`
- Test: `src/core/hook.test.ts`
- [ ] **Step 1: Write failing tests**
- The session cache round-trips `{answer, turns}`; each `buildHookOutput` call increments `turns`.
- Add `listPages` to the `HookClient` interface. On a turn where `turns % cfg.pageRefreshEveryTurns === 0`, the output includes `buildRosterRefresh(...)` content (assert "Component map" appears); on other turns it does not.
- Refresh is fail-open (a `listPages` rejection doesn't break recall/injection).
- First-turn behavior (reflect) unchanged.
- [ ] **Step 2: Run** → FAIL.
- [ ] **Step 3: Implement**
- Extend the cache read/write to `{ answer?: string; turns?: number }`. Compute `const turns = (cached.turns ?? 0) + 1;` and persist it alongside `answer`.
- Add `listPages(): Promise<unknown>` to `HookClient`.
- After computing `memBlock`, if `cfg.pageRefreshEveryTurns > 0 && turns % cfg.pageRefreshEveryTurns === 0`, `try { const refresh = buildRosterRefresh(parsePageList(await client.listPages())); if (refresh) blocks.push(refresh); } catch { /* fail-open */ }`. Kick the `listPages` call off concurrently with recall to avoid added latency.
- Import from `./knowledge-injection`.
- [ ] **Step 4: Run** `npx vitest run src/core/hook.test.ts` → PASS.
- [ ] **Step 5: Commit** `git add src/core/hook.ts src/core/hook.test.ts && git commit -m "feat(core): UserPromptSubmit hook-counted periodic page-roster refresh"`
---
## Task 9: Full check + LLM behavior (live) verification
**Files:**
- Modify: `src/system.live.test.ts` (add coverage; runs only under `HINDSIGHT_LIVE_E2E=1`)
- [ ] **Step 1: Full fast suite + types** — `npx vitest run && npx tsc --noEmit` → all green.
- [ ] **Step 2: Add a live assertion** (guarded by the existing live env flag) that after seeding a small repo + one `captureInitiative`, the Initiatives page content contains a `[[page:initiative-…]]` link (verifies the `relatedPageId` → link rendering end-to-end). Keep it in the live suite; do not run in the fast job.
- [ ] **Step 3: Manual/live run** (optional, operator): `HINDSIGHT_API_URL=http://localhost:8888 npm run test:live`.
- [ ] **Step 4: Commit** `git add src/system.live.test.ts && git commit -m "test(live): initiative page renders relatedPageId link end-to-end"`
---
## Final review
- [ ] Dispatch a final code-reviewer over the whole change set against the spec (`docs/superpowers/specs/2026-07-25-v2-knowledge-pages-design.md`).
- [ ] Rebuild + dev-install the `claude-code-v2` bundle so the running plugin picks up the new hooks/MCP (`bash scripts/dev-install.sh`); do not push/PR without explicit consent.
- [ ] Note deferred follow-ups: session drill-down tag, `capture_decision`, `gotcha` tier, older-bank reseed requirement.
@@ -1,135 +0,0 @@
# v2 Knowledge Pages — Design Spec
**Status:** approved in brainstorm (2026-07-25), pending implementation plan
**Scope:** `hindsight-integrations/hindsight-coding-agents` (shared TS core) + `claude-code-v2` wrapper
**Motivation:** make knowledge pages a real, trustworthy "wiki" surface for the vectorize-crm demo (the `knowledge-pages-as-trust-surface` principle) — the agent reliably knows what pages exist, pages are cleanly tiered instead of blended, and major initiatives become first-class, linkable pages.
---
## 1. Problem
Three gaps in the current v2 branch:
1. **Page discovery is a blind fetch.** SessionStart injects a static `KNOWLEDGE_MISSION` telling the agent to call `agent_knowledge_list_pages`, but hands it **no roster** — the agent never learns a page exists unless it independently decides to call the tool. Per-turn recall injects facts, not pages.
2. **Pages are blended.** Neither the seeded `PAGES` nor the agent's `create_page` tool scope synthesis by tag, so every page synthesizes from the whole bank filtered only by `fact_type`. Git-log, session, and survey memories all bleed into every page.
3. **No first-class initiative tracking / linking.** Hindsight has no native page-to-page links. A "major feature" leaves no durable, navigable page a future session can pick up.
## 2. Principles applied
- Automatic/visible value; zero out-of-band CLI; memory beats code search; knowledge pages as a trust surface; minimal post-setup burden.
- Modular units, small files, follow existing patterns (per-hook specs, fail-open, unit-testable pure cores).
- The **memory extractor never knows what a "page" is.** Classification is by the fact's *intrinsic* nature; pages are application-side saved views. No abstraction leak into extraction.
## 3. Architecture overview
Two complementary curation paths + a discovery layer:
- **Passive (automatic):** `entity_labels` schema-forces the extractor to tag qualifying facts `knowledge:<tier>`. Seeded **tier pages** each filter on one tier tag. No agent effort.
- **Active (high-signal):** one intent-named MCP verb, `hindsight_capture_initiative`, lets the agent register a major feature as a **per-initiative page** with a tag-based link back from the aggregate Initiatives page.
- **Discovery:** SessionStart injects guidance + the page roster; the UserPromptSubmit hook re-injects a fresh roster on a fixed cadence (hook-counted, not model-counted).
## 4. `entity_labels` — passive tier tagging
One hierarchical bank config group, set by `configureBank` at seed time:
```jsonc
{
"key": "knowledge",
"type": "multi-values", // 0, 1, or several — empty is normal
"optional": true,
"tag": true, // emits knowledge:<value> onto the fact's tags
"description": "Routing labels for this project's Hindsight KNOWLEDGE PAGES — curated, human-readable summaries of the repo's DURABLE engineering knowledge (architecture, key decisions, conventions, ongoing initiatives), each page rebuilt automatically from the facts labeled for it. Mark a fact only when it is durable, reusable knowledge a developer would still want surfaced in future sessions. IMPORTANT: leave this EMPTY for routine, transient, or operational facts — a passing test, a one-off command, a status update, a debugging dead-end. MOST facts should get no label here. Assign more than one value only when the fact genuinely fits several.",
"values": [
{ "value": "feature-work", "description": "A new feature, initiative, or enhancement being planned or built — the capability being added and the intent behind it. Not routine bug-fixes or chores." },
{ "value": "decision", "description": "A technical decision that will constrain future work, with its rationale — why this approach was chosen over alternatives, or a rule deliberately adopted." },
{ "value": "convention", "description": "An established way this project does things — naming, structure, testing, error handling, or another recurring pattern a contributor is expected to follow." },
{ "value": "component", "description": "What a specific module, file, service, or subsystem is responsible for, or how components depend on and connect to one another." },
{ "value": "concept", "description": "A domain concept, key abstraction, or piece of project vocabulary a new contributor must understand to work effectively." }
]
}
```
Notes:
- `tag: true``_inject_label_tags` copies each `knowledge:<value>` onto the fact's `tags` (no extra query infra).
- Selectivity (multi-values + "mostly empty" instruction) prevents force-fitting routine facts into a tier.
## 5. Seeded tier pages (tag-scoped)
Created via `/knowledge-base/pages` (supports `tags`, `trigger`, `parent_id`) — **not** `/mental-models`. Each `PAGES` entry gains a `trigger.tags` pin:
| Page | `trigger.tags` |
| --- | --- |
| Initiatives and enhancements | `["knowledge:feature-work"]` |
| Key decisions and rationale | `["knowledge:decision"]` |
| Conventions and patterns | `["knowledge:convention"]` |
| Component map | `["knowledge:component"]` |
| Core concepts | `["knowledge:concept"]` |
`tags_match` strict enough to exclude untagged facts (`all_strict`/`any_strict`). Tag matching is exact set-ops (no wildcards) — this is *why* the vocabulary is fixed, not per-feature.
## 6. MCP surface
Raw page CRUD (`create_page`/`update_page`/`delete_page`) is **removed** from the agent. The agent sees grounding tools + one capture verb. Naming convention: `hindsight_*`.
**Grounding**
- `hindsight_list_knowledge_pages` `{}` — roster: id, title, one-line coverage. (agent-facing description as drafted in brainstorm)
- `hindsight_read_knowledge_page` `{ page_id }` — full page content; follow `[[page:<id>]]` links by re-calling.
- `hindsight_search_memory` `{ query, max_tokens? }` — raw fact recall for specifics pages don't cover.
- `hindsight_get_current_bank` `{}` — minor introspection (kept).
**Capture**
- `hindsight_capture_initiative` `{ title, summary, relates_to_page_id? }` — the one active verb. Explicit WHEN / WHEN-NOT description (as drafted). Returns the initiative page id.
- `hindsight_ingest_document` `{ title, content }` — existing `agent_knowledge_ingest`, reframed.
(Full agent-facing descriptions are captured verbatim in the brainstorm thread and will be reproduced in the implementation plan.)
## 7. `hindsight_capture_initiative` mechanism
- Derive one slug `S` from `title`. Page id = `initiative-<S>`. **The slug in the tag and the page id are the same token, derived once** (cannot drift).
- **New initiative** (`relates_to_page_id` omitted):
1. Create page `initiative-<S>` (title from `title`, `source_query` about that initiative) under an **"Initiatives" folder** (tag-scoped).
2. Retain a marker memory (text = title + summary) tagged `["knowledge:feature-work", "relatedPageId:initiative-<S>"]`. **No session tag** (decided — the MCP server has no Claude session id; faking one wouldn't link to the Stop write-back's `conversation:<sessionId>` doc anyway).
- **Enhancement** (`relates_to_page_id` given): marker only, `relatedPageId = relates_to_page_id`; no new page. Re-invoking for the same initiative accrues markers → the page re-synthesizes with progress.
### Link survival (why `relatedPageId` as a tag, not in prose)
A tag is set directly via the retain `tags` param — it **bypasses LLM extraction entirely**, so it's guaranteed present verbatim (no REF-ID-style preservation needed at extraction). Verified: the reflect/synthesis path SELECTs `tags` and serializes facts via `_prune_nulls(model_dump())`, which keeps non-empty tags → **the synthesis LLM sees the tag.** The **Initiatives page `source_query`** instructs: *"when a memory carries a `relatedPageId:<id>` tag, emit a `[[page:<id>]]` link to it."* The link id is generated from the tag value at synthesis time, so it always matches the created page id.
- Only **Stage 2 (synthesis)** is probabilistic now (bounded token budget may omit some entries when there are many).
- **Guaranteed fallback:** the per-initiative page always exists (created via API, independent of any LLM stage) and appears in the **Initiatives folder / injected roster**, so navigation works even if a synthesized inline link drops.
## 8. Page-access injection
- **SessionStart** (`session-start.ts`): replace static `KNOWLEDGE_MISSION` with a preamble = (a) guidance on *when/why* to consult pages, (b) the roster fetched via `client.listPages()` (`- <title> (<id>)`, empty-state aware), (c) a note that the list refreshes periodically. Cold repo → empty roster line; roster comes alive mid-session as seeding/survey complete.
- **UserPromptSubmit** (`hook.ts`): extend the per-session cache (`{answer}``{answer, turns}`); the **hook** counts user turns and, roughly every `pageRefreshEveryTurns` (default 10, approximate), calls `listPages()` and injects a compact roster refresh. Runs concurrently with recall; **fail-open** (a refresh error never blocks the turn).
- **Shared formatting** (new `core/knowledge-injection.ts`, SDK-free/unit-testable): `parsePageList(raw) -> {id,title}[]`, `buildKnowledgePreamble(pages)`, `buildRosterRefresh(pages)`.
- **Config:** `pageRefreshEveryTurns` (default 10).
## 9. Non-goals / deferred
- Session drill-down tag on captured markers (dropped — see §7).
- `hindsight_capture_decision` and other capture verbs (passive path covers those tiers; revisit if the aggregate pages aren't sharp enough).
- A `gotcha`/`pitfall` tier (five tiers for now).
- Native page-to-page links / backlinks (Hindsight has none; we approximate via folder tree + `relatedPageId`-driven `[[page:<id>]]`).
## 10. Risks / migration
- **Older banks** need re-seeding to pick up the new `entity_labels`, the `session` retain strategy, and the tag-scoped page triggers (`configureBank` sets them). User is starting fresh with v2 banks, so acceptable; live retain fails open otherwise.
- **Stage-2 synthesis omission** for large initiative counts — mitigated by the folder/roster fallback.
- **Instruction adherence** for the `source_query` link-rendering and the label selectivity — both are LLM-following behaviors; cover with an `hs_llm_core` judge test, and the deterministic mechanics (tag injection, roster formatting, slug/id equality, hook turn-counting) with fast unit tests.
## 11. Testing
- **Deterministic unit tests:** `knowledge-injection` formatting + empty-state; hook turn-counter + cadence; `capture_initiative` slug→id→tag equality and request shape (mock client); tag-scoped page request bodies; entity_labels config emitted by `configureBank`.
- **LLM judge test (`hs_llm_core`):** label selectivity (routine facts get no `knowledge:*`), and `relatedPageId``[[page:<id>]]` rendering in a synthesized Initiatives page.
## 12. File map (anticipated)
- `src/core/knowledge-injection.ts` (new) — roster/preamble formatting.
- `src/core/session-start.ts` — preamble + roster.
- `src/core/hook.ts` — cache `{answer,turns}` + periodic roster refresh.
- `src/core/config.ts``pageRefreshEveryTurns`.
- `src/core/missions.ts``entity_labels` group; tag-scoped `PAGES`; Initiatives `source_query` link instruction.
- `src/core/hindsight.ts``configureBank` sets `entity_labels`; `createPages` pins `trigger.tags` + Initiatives folder; new `createInitiativePage`/marker retain helpers.
- `src/core/knowledge-tools.ts` — new `hindsight_*` grounding + `capture_initiative` tools; remove raw page CRUD from agent surface.
- Tests alongside each.
@@ -1,139 +0,0 @@
# Reflect + Pages Runtime — Design Spec
**Status:** decided (2026-07-27), reconciles the earlier reflect-based runtime with the recall-based v2 into one opinionated path
**Scope:** `hindsight-integrations/hindsight-coding-agents` (shared TS core) + `claude-code-v2` wrapper
**Motivation:** the 33-task coding benchmark showed the v2 recall-per-prompt runtime *underperforms no memory* (35.0 mean corrections vs 32.0 baseline), while the earlier reflect-injection runtime beats baseline by 22% (25.0). This spec restores reflect as the only deep-memory path and replaces raw per-turn recall with lightweight injection from knowledge pages — "fast like recall, organized like reflect" — keeping v2's page/curation machinery where it earned its place and deleting it where it didn't.
---
## 1. Problem
Two prior iterations, each half right:
1. **Reflect runtime (v1):** one agentic REFLECT over the bank at session start, cached and re-injected every turn. Benchmark-proven (25.0 mean corrections) — but nothing surfaced mid-session; a task that drifted away from the first message got stale context.
2. **Recall runtime (v2):** per-prompt recall injection for turn-by-turn visibility, plus knowledge pages as a trust surface. But raw recall injects unsynthesized fact fragments — noise that *hurt*: 35.0 mean corrections, worse than running with no memory at all.
| Runtime | Mean corrections (33-task benchmark) | vs no-memory (32.0) |
| --- | --- | --- |
| Reflect-injection (v1) | **25.0** | **22%** |
| Recall-per-prompt (v2) | 35.0 | +9% (regression) |
| No memory | 32.0 | baseline |
The reconciliation: keep reflect's synthesis quality as the deep path, keep v2's per-turn visibility principle, but source the per-turn material from the already-synthesized knowledge pages instead of raw recall.
## 2. Decisions
Explicit, decided — not options:
1. **Reflect restored** as the only deep-memory path (session-start, agentic synthesis, cached + re-injected every turn).
2. **Recall removed from the runtime** entirely. No per-prompt `recall` call.
3. **No `memoryMode` flag.** One opinionated path; config is for environment, naming, and harness wiring only — never behavior selection.
4. **Sections, not pages, are the per-turn injection unit** — locally matched, budget-trimmed, provenance-labeled.
5. **JSON turn transcripts** replace the markdown tool-call transcript in the Stop-hook write-back, with compact action entries.
6. **No tags / no `entity_labels`.** The server re-synthesizes pages after consolidation; "living pages" needs no client-side tagging machinery.
## 3. Runtime path — session start
Three steps, in order, all inside existing hooks (no out-of-band CLI):
### 3a. Cold-repo bootstrap (kept from v2)
On a bank with no prior memories: automatic shallow gitlog seed + codebase survey, exactly as v2 does it. The user never runs a setup command; the first session self-seeds. (Deep ingestion of that history is §7 — the seed here stays instant.)
### 3b. REFLECT once, on the first task message
The benchmark-proven core:
- On the first user prompt of the session, run one **REFLECT** — agentic synthesis over the whole bank, prompted to return the *root-cause decision with exact values* (concrete file paths, config values, version numbers — not summaries of summaries).
- Cache the result per session; **re-inject it every turn**. It is the session's durable deep context.
- One LLM-backed call per session, on the message that actually states the task — not on session-open, where there is nothing to reflect about.
### 3c. Page index build
Fetch all knowledge pages once (existing `listPages` + page reads), split each page at headings into **sections**, and build a **local section index** in the hook process. This index is what every subsequent turn matches against (§4) — no further server calls on the hot path.
## 4. Runtime path — every turn
Per-turn visibility, satisfied at ~zero latency and ~zero cost. Injection sources from **knowledge pages, not raw recall** — the material is already synthesized and organized; the turn hook only *selects* from it.
Mechanism (local, deterministic — no server call, no LLM call):
| Aspect | Design |
| --- | --- |
| Unit | Page **sections** (pages split at headings at index-build time) |
| Matching | Lexical: prompt scored against each section by weighted term overlap; **heading hits weighted higher** than body hits |
| Selection | Top 23 sections |
| Budget | Trimmed to a **~700-token total** |
| Provenance | Each snippet labeled `From <page> <section>` + a tool pointer to read the full page |
| Floor | A minimum-score threshold below which **nothing is injected** — silence over noise |
| Refresh | Section index rebuilt on the existing 10-turn roster cadence (`pageRefreshEveryTurns`) |
The score floor is load-bearing: the benchmark showed that injecting weak matches is worse than injecting nothing (v2's regression). An empty injection is a correct outcome, not a failure mode.
## 5. Write-back
The Stop-hook session retain is **kept** — same trigger, same fail-open behavior. What changes is the transcript format handed to extraction:
- **JSON turns**, not markdown: an array of `{ "role": "user" | "assistant", "text": ... }` entries for the conversational content.
- Tool calls collapse to **compact one-line action entries**: `{ "role": "action", "text": "Edit boltons/strutils.py" }` — tool name + primary target only, **no arguments, no outputs**.
Rationale: extraction keeps the concrete artifacts (which files were touched, what actions occurred) without the transcript noise of full tool payloads — the markdown tool-call dumps were volume without signal.
## 6. Knowledge pages
Simplified from the v2 spec:
- **Dropped: tags and `entity_labels`** (v2 spec §45). The server already re-synthesizes pages after consolidation, so pages stay "living" with no client-side routing machinery. The extractor-never-knows-about-pages principle now holds trivially — there is nothing to route.
- **Creation paths:**
1. **Seeded taxonomy** at bank creation (the fixed page set, as today, minus tag triggers).
2. **Agent-driven `capture_initiative`** at plan approval — the one active capture verb survives from v2.
3. **Organic splitting** of pages that outgrow their scope is a **server/curator concern**, not a client feature.
## 7. Ingestion — progressive background deepening
*Status: design accepted, implementation phased separately.*
Replaces the manual backfill CLI as the user-facing path (the CLI was out-of-band burden; nobody runs it). The principle: converge to full-depth history through normal usage, with zero user action.
1. **Instant shallow seed** — the gitlog seed from §3a; the session is useful immediately.
2. **Background deepening** — a background worker deep-ingests **per-commit-with-diffs, incrementally**, never blocking a turn.
3. **Working-set prioritization** — commits are ingested in order of relevance to what the agent is actually doing: files the agent reads/edits get their commit histories ingested **first**. Depth arrives where it pays off.
4. **Checkpointing** — progress persists across sessions; each session resumes deepening where the last left off, converging to full depth over normal usage.
The **backfill CLI survives as an internal tool** (benchmark setup, CI bank preparation) — it is no longer a documented user path.
## 8. Gap analysis — v2 principles under this design
| v2 principle | How this design satisfies it |
| --- | --- |
| See-it-working (automatic, visible value) | Reflect answer visible from turn 1; page-section snippets appear with explicit `From <page> <section>` provenance, so the user sees memory working — and the score floor keeps it from visibly misfiring. |
| No out-of-band CLI | Cold-repo auto-seed kept (§3a); backfill CLI demoted to internal-only, replaced by background deepening (§7). Nothing requires a terminal command. |
| Reuse-over-reinvent | Reflect, `listPages`, Stop-hook retain, `capture_initiative`, and the 10-turn refresh cadence are all existing machinery recombined; the only new code is the local section index and matcher — deliberately dumb (lexical, no LLM). |
| Preserve-intent | Reflect is prompted for root-cause decisions with exact values; JSON transcripts keep concrete action artifacts; per-commit-with-diffs deepening captures *why* the code changed, not just that it did. |
| Near-zero-burden | No config flags to choose, no CLI to run, no tags to maintain; one LLM call per session start, everything else local. |
## 9. Verification gates
Ship gates, in order:
1. **Reflect-restored benchmark:** the restored runtime must recover **~25 mean corrections at n=2 on identical banks** to the original reflect run. This proves the restoration is faithful before anything is layered on.
2. **Reflect+pages benchmark:** with per-turn section injection enabled, the score **must not regress** vs reflect-alone. Section injection earns its place by not hurting; any regression points at the floor/budget tuning.
3. **Live system suite:** existing hook/integration suite updated for the new path — reflect caching + per-turn re-injection, section index build/refresh, score-floor silence, JSON transcript shape, action-entry compaction. Deterministic pieces (matcher scoring, budget trim, provenance formatting, transcript serialization) as fast unit tests.
## 10. Non-goals / deferred
- Any per-turn LLM or server call for injection (explicitly excluded — the local matcher is the whole point).
- Semantic/embedding-based section matching (revisit only if lexical matching demonstrably misses; start dumb).
- Client-side page splitting or curation (server/curator concern, §6).
- Progressive-deepening implementation details (worker scheduling, checkpoint format) — phased separately per §7.
## 11. File map (anticipated)
- `src/core/reflect.ts` (restored) — session reflect call + per-session cache.
- `src/core/section-index.ts` (new) — page → sections split, lexical scorer, budget trim, provenance formatting; pure/unit-testable.
- `src/core/hook.ts` — drop recall; inject cached reflect + matched sections; index refresh on roster cadence.
- `src/core/session-start.ts` — cold-repo seed (unchanged) + reflect trigger wiring + initial index build.
- `src/core/transcript.ts` (new or reworked) — JSON turn serialization + action-entry compaction for the Stop hook.
- `src/core/missions.ts` / `src/core/hindsight.ts` — remove `entity_labels` and tag-scoped page triggers; keep seeded taxonomy + `capture_initiative`.
- `src/core/config.ts` — remove any behavior flags; keep env/naming/harness + `pageRefreshEveryTurns`.
- Tests alongside each.
+2 -2
View File
@@ -2,8 +2,8 @@ apiVersion: v2
name: hindsight
description: Hindsight helm chart
type: application
version: 0.8.6
appVersion: "0.8.6"
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 }}
+1 -1
View File
@@ -1,6 +1,6 @@
{
"name": "@vectorize-io/hindsight-all",
"version": "0.8.6",
"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",
+3 -4
View File
@@ -1,16 +1,15 @@
[build-system]
requires = ["setuptools>=77"]
requires = ["setuptools>=61"]
build-backend = "setuptools.build_meta"
[project]
name = "hindsight-all-slim"
version = "0.8.6"
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.8.6",
"hindsight-api-slim==0.8.4",
"hindsight-client>=0.0.7",
"hindsight-embed>=0.1.0",
]
+4 -5
View File
@@ -1,16 +1,15 @@
[build-system]
requires = ["hatchling>=1.27"]
requires = ["hatchling"]
build-backend = "hatchling.build"
[project]
name = "hindsight-all"
version = "0.8.6"
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.8.6",
"hindsight-api-slim[all]==0.8.4",
"hindsight-client>=0.0.7",
"hindsight-embed>=0.1.0",
]
@@ -22,7 +21,7 @@ hindsight-embed = { workspace = true }
[project.optional-dependencies]
local-llm = [
"hindsight-api-slim[local-llm]==0.8.6",
"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.8.6"
__version__ = "0.8.4"
+31 -416
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
@@ -70,173 +66,7 @@ BACKUP_TABLES = [
"graph_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:
@@ -247,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))
@@ -257,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] = {}
@@ -285,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)
@@ -313,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")
@@ -325,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:
@@ -347,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...")
@@ -395,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()
@@ -434,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}")
@@ -467,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")
@@ -481,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()
@@ -516,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(
@@ -600,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)
@@ -735,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()
@@ -854,8 +473,7 @@ def import_bank_command(
async def _decommission_worker(db_url: str, worker_id: str, schema: str = "public") -> int:
"""Release all tasks owned by a worker, setting them back to pending status."""
_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)
@@ -914,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)
@@ -980,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)
@@ -1049,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)
@@ -1,268 +0,0 @@
"""Drop the deprecated entity schema from ``memory_links``.
Entity edges are no longer materialized in ``memory_links``. Retain stores
memory-to-entity associations in ``unit_entities``, and both read paths derive
entity edges on demand from that table — the /graph endpoint builds them from
shared ``unit_entities`` rows, and recall expands via the ``unit_entities``
self-join. Migration ``e9b2c7d1f3a4`` deleted the materialized entity rows and
current writers only ever pass ``entity_id = NULL``, so the entity-specific
schema on ``memory_links`` is now dead weight:
- the ``entity_id`` column and its FK to ``entities``
- ``link_type = 'entity'`` in the table CHECK constraint
- the entity index (``idx_memory_links_entity`` on PG, ``idx_ml_entity`` on Oracle)
- the ``entity_id`` term in the function-based ``idx_memory_links_unique``
Once entity edges are gone, every remaining row (temporal, semantic, and the
causal types) has a single meaningful identity — ``(from_unit_id, to_unit_id,
link_type)`` — so the unique index collapses to those three columns and
preserves the existing effective uniqueness semantics.
Production ``memory_links`` tables and the entity index can be very large, so
this migration is written to avoid long exclusive locks: the residual delete is
chunked with per-batch commits, every index is swapped with ``CONCURRENTLY``,
and the new CHECK is added ``NOT VALID`` then validated separately so writers are
never blocked on a full-table scan.
Downgrade restores the former schema *shape* (column, FK, index, CHECK, and the
expression unique index) but cannot reconstruct the historical entity rows that
``e9b2c7d1f3a4`` already deleted — new retains do not produce them either, so the
restored entity index would simply stay empty.
Revision ID: c1e7a9d3f5b2
Revises: e4a7c1b9d2f6
Create Date: 2026-08-04
"""
from collections.abc import Sequence
from alembic import context, op
from hindsight_api.alembic._dialect import run_for_dialect
revision: str = "c1e7a9d3f5b2"
down_revision: str | Sequence[str] | None = "e4a7c1b9d2f6"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
_NIL_ENTITY_UUID = "00000000-0000-0000-0000-000000000000"
_LINK_TYPES_WITHOUT_ENTITY = "'temporal', 'semantic', 'causes', 'caused_by', 'enables', 'prevents'"
_LINK_TYPES_WITH_ENTITY = "'temporal', 'semantic', 'entity', 'causes', 'caused_by', 'enables', 'prevents'"
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()
# CONCURRENTLY index swaps and the DO block's per-batch COMMIT both require
# running outside Alembic's migration transaction — an autocommit_block
# commits it and switches the connection to autocommit for the duration.
with op.get_context().autocommit_block():
# 1. Defensively delete any residual entity rows. e9b2c7d1f3a4 already
# did this, but a bank that predates its deployment can still carry
# them, and they must be gone before the three-column unique index
# (which no longer distinguishes entity rows) and the entity-free
# CHECK are created. Chunked with per-batch commits so a very large
# table drains in bounded transactions instead of one long-locking
# delete.
op.execute(
f"""
DO $$
DECLARE
deleted INTEGER;
BEGIN
LOOP
DELETE FROM {schema}memory_links
WHERE ctid IN (
SELECT ctid FROM {schema}memory_links
WHERE link_type = 'entity'
LIMIT 50000
);
GET DIAGNOSTICS deleted = ROW_COUNT;
EXIT WHEN deleted = 0;
COMMIT;
END LOOP;
END$$;
"""
)
# 2. Drop the entity indexes. idx_memory_links_entity_covering was
# already removed by e1b2c3d4f5a6; dropped again defensively.
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_links_entity")
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_links_entity_covering")
# 3. Replace the expression unique index with the three-column form.
# Build the new one first (under a temporary name) so uniqueness is
# never unprotected, then drop the old expression index and rename.
# Non-entity rows already collapse to (from, to, link_type) under the
# old COALESCE(entity_id, nil) expression, so this cannot find a
# duplicate once the entity rows are gone.
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_links_unique_new")
op.execute(
f"CREATE UNIQUE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_links_unique_new "
f"ON {schema}memory_links (from_unit_id, to_unit_id, link_type)"
)
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_links_unique")
op.execute(f"ALTER INDEX IF EXISTS {schema}idx_memory_links_unique_new RENAME TO idx_memory_links_unique")
# 4. Drop the FK and the now-unreferenced column. Both are metadata-only
# on PostgreSQL (DROP COLUMN marks the attribute dropped, no rewrite).
op.execute(f"ALTER TABLE {schema}memory_links DROP CONSTRAINT IF EXISTS fk_memory_links_entity_id_entities")
op.execute(f"ALTER TABLE {schema}memory_links DROP COLUMN IF EXISTS entity_id")
# 5. Recreate the link_type CHECK without 'entity'. Added NOT VALID then
# validated separately: the ADD takes a brief lock without scanning,
# and VALIDATE takes only SHARE UPDATE EXCLUSIVE, so concurrent reads
# and writes are never blocked on the scan.
op.execute(f"ALTER TABLE {schema}memory_links DROP CONSTRAINT IF EXISTS memory_links_link_type_check")
op.execute(
f"ALTER TABLE {schema}memory_links ADD CONSTRAINT memory_links_link_type_check "
f"CHECK (link_type IN ({_LINK_TYPES_WITHOUT_ENTITY})) NOT VALID"
)
op.execute(f"ALTER TABLE {schema}memory_links VALIDATE CONSTRAINT memory_links_link_type_check")
def _pg_downgrade() -> None:
schema = _pg_schema_prefix()
with op.get_context().autocommit_block():
# Restore the column and FK (nullable, as in the initial schema).
op.execute(f"ALTER TABLE {schema}memory_links ADD COLUMN IF NOT EXISTS entity_id uuid")
op.execute(f"ALTER TABLE {schema}memory_links DROP CONSTRAINT IF EXISTS fk_memory_links_entity_id_entities")
op.execute(
f"ALTER TABLE {schema}memory_links ADD CONSTRAINT fk_memory_links_entity_id_entities "
f"FOREIGN KEY (entity_id) REFERENCES {schema}entities (id) ON DELETE CASCADE"
)
# Restore the entity partial index.
op.execute(
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_links_entity "
f"ON {schema}memory_links (entity_id) WHERE entity_id IS NOT NULL"
)
# Restore the expression unique index (entity_id back in the key).
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_links_unique_old")
op.execute(
f"CREATE UNIQUE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_links_unique_old "
f"ON {schema}memory_links (from_unit_id, to_unit_id, link_type, "
f"COALESCE(entity_id, '{_NIL_ENTITY_UUID}'::uuid))"
)
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_links_unique")
op.execute(f"ALTER INDEX IF EXISTS {schema}idx_memory_links_unique_old RENAME TO idx_memory_links_unique")
# Restore 'entity' as a permitted link_type.
op.execute(f"ALTER TABLE {schema}memory_links DROP CONSTRAINT IF EXISTS memory_links_link_type_check")
op.execute(
f"ALTER TABLE {schema}memory_links ADD CONSTRAINT memory_links_link_type_check "
f"CHECK (link_type IN ({_LINK_TYPES_WITH_ENTITY})) NOT VALID"
)
op.execute(f"ALTER TABLE {schema}memory_links VALIDATE CONSTRAINT memory_links_link_type_check")
def _oracle_exec_ignoring(sql: str, *ignore_codes: int) -> None:
"""Run a DDL statement, swallowing the given ORA-NNNNN codes for idempotency.
Oracle has no ``IF EXISTS``/``IF NOT EXISTS`` for most objects; the standard
pattern (see e4a7c1b9d2f6) is EXECUTE IMMEDIATE inside a PL/SQL block that
re-raises anything but the expected "already absent"/"already present" code.
"""
conditions = " AND ".join(f"SQLCODE != {code}" for code in ignore_codes)
escaped = sql.replace("'", "''")
op.execute(
f"""
BEGIN
EXECUTE IMMEDIATE '{escaped}';
EXCEPTION WHEN OTHERS THEN
IF {conditions} THEN RAISE; END IF;
END;
"""
)
def _oracle_upgrade() -> None:
# 1. Defensively drain residual entity rows in bounded, committed chunks so a
# large table doesn't delete under one long-held lock / huge undo segment.
op.execute(
"""
BEGIN
LOOP
DELETE FROM memory_links WHERE link_type = 'entity' AND ROWNUM <= 50000;
EXIT WHEN SQL%ROWCOUNT = 0;
COMMIT;
END LOOP;
END;
"""
)
# 2. Build the three-column unique index under a temporary name before
# dropping the old expression index, then rename. Oracle DDL implicitly
# commits, so ordering it this way keeps duplicate protection continuous
# instead of leaving a gap between drop and create. ONLINE avoids blocking
# concurrent DML during the (potentially large) index builds/drops.
# ORA-00955: name already in use; ORA-01418: index does not exist.
_oracle_exec_ignoring("DROP INDEX idx_memory_links_unique_new", -1418)
_oracle_exec_ignoring(
"CREATE UNIQUE INDEX idx_memory_links_unique_new ON memory_links (from_unit_id, to_unit_id, link_type) ONLINE",
-955,
)
_oracle_exec_ignoring("DROP INDEX idx_memory_links_unique", -1418)
_oracle_exec_ignoring("ALTER INDEX idx_memory_links_unique_new RENAME TO idx_memory_links_unique", -1418)
# 3. Drop the entity index. ORA-01418: index does not exist.
_oracle_exec_ignoring("DROP INDEX idx_ml_entity ONLINE", -1418)
# 4. Drop the entity FK, then the column. ORA-02443: constraint does not
# exist; ORA-00904: column does not exist.
_oracle_exec_ignoring("ALTER TABLE memory_links DROP CONSTRAINT fk_ml_entity", -2443)
_oracle_exec_ignoring("ALTER TABLE memory_links DROP COLUMN entity_id", -904)
# 5. Recreate the link_type CHECK without 'entity'. ORA-02443: constraint
# does not exist; ORA-02264: name already used by an existing constraint.
_oracle_exec_ignoring("ALTER TABLE memory_links DROP CONSTRAINT chk_ml_link_type", -2443)
_oracle_exec_ignoring(
f"ALTER TABLE memory_links ADD CONSTRAINT chk_ml_link_type CHECK (link_type IN ({_LINK_TYPES_WITHOUT_ENTITY}))",
-2264,
)
def _oracle_downgrade() -> None:
# Restore the column, FK, entity index, expression unique index, and the
# 'entity'-permitting CHECK. ORA-01430: column already exists; ORA-00955:
# name already in use; ORA-01418: index does not exist; ORA-02443/-2264:
# constraint absent / name in use.
_oracle_exec_ignoring("ALTER TABLE memory_links ADD (entity_id RAW(16))", -1430)
_oracle_exec_ignoring(
"ALTER TABLE memory_links ADD CONSTRAINT fk_ml_entity "
"FOREIGN KEY (entity_id) REFERENCES entities (id) ON DELETE CASCADE",
-2264,
)
_oracle_exec_ignoring("CREATE INDEX idx_ml_entity ON memory_links (entity_id)", -955)
_oracle_exec_ignoring("DROP INDEX idx_memory_links_unique_old", -1418)
_oracle_exec_ignoring(
"CREATE UNIQUE INDEX idx_memory_links_unique_old ON memory_links ("
"from_unit_id, to_unit_id, link_type, "
"NVL(entity_id, HEXTORAW('00000000000000000000000000000000'))) ONLINE",
-955,
)
_oracle_exec_ignoring("DROP INDEX idx_memory_links_unique", -1418)
_oracle_exec_ignoring("ALTER INDEX idx_memory_links_unique_old RENAME TO idx_memory_links_unique", -1418)
_oracle_exec_ignoring("ALTER TABLE memory_links DROP CONSTRAINT chk_ml_link_type", -2443)
_oracle_exec_ignoring(
f"ALTER TABLE memory_links ADD CONSTRAINT chk_ml_link_type CHECK (link_type IN ({_LINK_TYPES_WITH_ENTITY}))",
-2264,
)
def upgrade() -> None:
run_for_dialect(pg=_pg_upgrade, oracle=_oracle_upgrade)
def downgrade() -> None:
run_for_dialect(pg=_pg_downgrade, oracle=_oracle_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,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,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,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,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: c1e7a9d3f5b2
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 = "c1e7a9d3f5b2"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
def _pg_schema_prefix() -> str:
"""Schema-qualifier for raw SQL on PG (multi-tenant search_path)."""
schema = context.config.get_main_option("target_schema")
return f'"{schema}".' if schema else ""
def _configured_vector_extension() -> str:
ext = os.getenv("HINDSIGHT_API_VECTOR_EXTENSION", "pgvector").lower()
if ext not in {"pgvector", "pgvectorscale", "vchord", "scann"}:
raise ValueError(
f"Invalid HINDSIGHT_API_VECTOR_EXTENSION: {ext}. Must be 'pgvector', 'vchord', 'pgvectorscale', or 'scann'"
)
return ext
def _pg_upgrade() -> None:
# ScaNN uses a global vector index by design — nothing stale to repair.
if _configured_vector_extension() == "scann":
return
schema = _pg_schema_prefix()
# DROP INDEX needs ACCESS EXCLUSIVE on memory_units. While it waits for
# in-flight transactions, every new query on the table queues behind it,
# so on a write-busy schema an unbounded wait can pile up traffic. Fail
# fast instead: the migration errors, the schema stays below head, and
# the next migration pass retries — preferable to freezing the table.
# SET LOCAL scopes the timeout to this migration's transaction.
op.execute("SET LOCAL lock_timeout = '10s'")
op.execute(f"DROP INDEX IF EXISTS {schema}idx_memory_units_embedding")
def _pg_downgrade() -> None:
# Intentional no-op: recreating a potentially multi-GB vector index that no
# query uses is not a safe downgrade action. Downgrading past d5e6f7a8b9c0
# restores the global index for deployments that genuinely need it.
pass
def upgrade() -> None:
# PG-only repair: the stale index is a PostgreSQL artifact of the old
# reconcile; Oracle deployments never had a reconcile that created it.
run_for_dialect(pg=_pg_upgrade)
def downgrade() -> None:
run_for_dialect(pg=_pg_downgrade)
File diff suppressed because it is too large Load Diff
@@ -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)
@@ -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
@@ -32,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):
@@ -136,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,
@@ -340,17 +331,11 @@ class ConfigResolver:
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
@@ -359,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)
@@ -400,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"
@@ -410,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:
@@ -449,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
@@ -469,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,
@@ -499,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:
+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,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,11 +109,6 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
end.replace(hour=23, minute=59, second=59, microsecond=999999),
)
def safe_constraint(start: datetime | None, end: datetime | None) -> DateRange | NoTemporalConstraintSentinel:
if start is None or end is None:
return NO_TEMPORAL_CONSTRAINT
return constraint(start, end)
def subtract_months(months: int) -> datetime:
month_index = reference_date.month - months - 1
year = reference_date.year + month_index // 12
@@ -131,21 +126,11 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
day = min(base_date.day, calendar.monthrange(year, month)[1])
return base_date.replace(year=year, month=month, day=day)
def add_years(base_date: datetime, years: int) -> datetime | 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
@@ -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,7 +450,7 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
return None
return since_constraint(day)
def relative_offset_datetime(amount: int, unit: str, direction: int) -> datetime | None:
def relative_offset_datetime(amount: int, unit: str, direction: int) -> datetime:
if unit in ("", ""):
return reference_date + timedelta(days=direction * amount)
if unit in ("", "星期", "礼拜"):
@@ -479,15 +459,15 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
return add_months(reference_date, direction * amount)
return add_years(reference_date, direction * amount)
def point_constraint_at_offset(amount: int, unit: str, direction: int) -> DateRange | 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
#
@@ -801,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"(大大后天|大后天|后天|明天|明日|今天|今日|本日|当日|当天|昨天|昨日|大大前天|大前天|前天)"
@@ -819,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"(?<![上下大小每个各隔])"
@@ -919,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}"
@@ -935,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)
@@ -961,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)
@@ -1105,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),
)
@@ -1113,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}"
@@ -1133,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),
)
@@ -1164,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"这(一两|[两二]三|三两|三四|四五|五六|六七|七八|八九|九十)个?(天|日|周|星期|礼拜|月|年)"
@@ -1174,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:
@@ -1211,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:
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
@@ -20,6 +20,7 @@ from ..config import (
DEFAULT_RERANKER_ALIBABA_MODEL,
DEFAULT_RERANKER_COHERE_MODEL,
DEFAULT_RERANKER_FLASHRANK_CACHE_DIR,
DEFAULT_RERANKER_FLASHRANK_CPU_MEM_ARENA,
DEFAULT_RERANKER_FLASHRANK_MODEL,
DEFAULT_RERANKER_GOOGLE_MODEL,
DEFAULT_RERANKER_LITELLM_MAX_TOKENS_PER_DOC,
@@ -33,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.
@@ -59,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:
"""
@@ -119,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.
@@ -141,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
@@ -151,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:
@@ -179,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
@@ -229,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")
@@ -276,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]:
"""
@@ -392,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
@@ -451,7 +484,6 @@ class RemoteTEICrossEncoder(CrossEncoderModel):
semaphore,
"POST",
f"{self.base_url}/rerank",
headers=reranker_bank_attribution_headers(),
json={
"query": query,
"texts": texts,
@@ -592,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()
@@ -891,7 +919,6 @@ class FlashRankCrossEncoder(CrossEncoderModel):
self.max_length = max_length
self.cpu_mem_arena = cpu_mem_arena
self._ranker = None
self._device_type: str = "cpu" # FlashRank runs on CPU via ONNX Runtime
FlashRankCrossEncoder._max_concurrent = max_concurrent
@property
@@ -963,11 +990,11 @@ class FlashRankCrossEncoder(CrossEncoderModel):
def _predict_sync(self, pairs: list[tuple[str, str]]) -> list[float]:
"""Synchronous predict - processes each query group."""
from flashrank import RerankRequest
if not pairs:
return []
from flashrank import RerankRequest
try:
# Group pairs by query
query_groups: dict[str, list[tuple[int, str]]] = {}
@@ -996,7 +1023,7 @@ class FlashRankCrossEncoder(CrossEncoderModel):
return all_scores
finally:
release_local_inference_memory(self._device_type)
_malloc_trim()
async def predict(self, pairs: list[tuple[str, str]]) -> list[float]:
"""
@@ -1124,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,
@@ -1243,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
@@ -1256,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
@@ -1577,246 +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,
)
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()
@@ -1826,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,7 +18,6 @@ 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
@@ -35,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.
@@ -182,6 +130,7 @@ class DataAccessOps(ABC):
table: str,
sorted_links: list[tuple],
bank_id: str,
nil_entity_uuid: str,
exists_clause: str,
chunk_size: int = 5000,
) -> None:
@@ -200,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.
"""
@@ -228,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,
@@ -306,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.
"""
...
@@ -324,7 +244,6 @@ class DataAccessOps(ABC):
self,
ml_table: str,
mu_table: str,
window: UpdatedWindow,
) -> str:
"""Build semantic + causal expansion CTEs.
@@ -343,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.
@@ -566,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,
@@ -10,11 +10,9 @@ import uuid as uuid_mod
from datetime import UTC, datetime
from .base import DatabaseConnection
from .ops import DataAccessOps, LinkExpansionRows, TagListingParts, UpdatedWindow
from .ops import DataAccessOps, TagListingParts
from .result import DictResultRow as ResultRow
ORACLE_IN_LIST_LIMIT = 1000
class OracleOps(DataAccessOps):
"""Oracle-specific data access operations."""
@@ -142,6 +140,7 @@ class OracleOps(DataAccessOps):
table: str,
sorted_links: list[tuple],
bank_id: str,
nil_entity_uuid: str,
exists_clause: str,
chunk_size: int = 5000,
) -> None:
@@ -152,16 +151,18 @@ class OracleOps(DataAccessOps):
to_ids = [lnk[1] for lnk in sorted_links]
types = [lnk[2] for lnk in sorted_links]
weights = [lnk[3] for lnk in sorted_links]
entity_ids = [lnk[4] for lnk in sorted_links]
await conn.executemany(
f"""
INSERT INTO {table}
(from_unit_id, to_unit_id, link_type, weight, bank_id)
VALUES ($1, $2, $3, $4, $5)
ON CONFLICT (from_unit_id, to_unit_id, link_type)
(from_unit_id, to_unit_id, link_type, weight, entity_id, bank_id)
VALUES ($1, $2, $3, $4, $5, $6)
ON CONFLICT (from_unit_id, to_unit_id, link_type,
COALESCE(entity_id, '{nil_entity_uuid}'::uuid))
DO NOTHING
""",
[(from_ids[i], to_ids[i], types[i], weights[i], bank_id) for i in range(len(sorted_links))],
[(from_ids[i], to_ids[i], types[i], weights[i], entity_ids[i], bank_id) for i in range(len(sorted_links))],
)
async def bulk_insert_entities(
@@ -171,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:
@@ -217,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)
""",
@@ -225,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,
@@ -285,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],
)
@@ -332,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",
@@ -376,12 +329,6 @@ class OracleOps(DataAccessOps):
entities_table: str,
bank_id: str,
) -> int:
# NB: the Postgres path additionally selects victims FOR UPDATE in sorted
# (entity_id_1, entity_id_2) order to prevent the #2529 deadlock against
# retain's sorted cooccurrence upsert. Oracle's DELETE can't carry that
# ordered-lock CTE the same way, so here we rely on the Pass 2/3 retry
# wrap in run_graph_maintenance_job (retry_with_backoff is ORA-00060
# deadlock-aware) to recover instead. Deliberate dialect asymmetry.
deleted = await conn.execute(
f"""
DELETE FROM {ec_table}
@@ -483,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.
@@ -501,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
@@ -523,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
)"""
@@ -531,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.
@@ -546,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
@@ -555,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
),
@@ -582,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,
@@ -601,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__)
@@ -656,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")
@@ -678,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
),
@@ -711,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,
@@ -726,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(
@@ -898,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,
@@ -4,40 +4,11 @@ Uses unnest(), LATERAL, DISTINCT ON, and native array operations for
efficient batch operations.
"""
from datetime import datetime
from .base import DatabaseConnection
from .ops import DataAccessOps, LinkExpansionRows, TagListingParts, UpdatedWindow
from .ops import DataAccessOps, TagListingParts
from .result import ResultRow
def pg_search_vector_expr(
config,
*,
text_col: str = "text",
context_col: str = "context",
signals_col: str = "text_signals",
) -> str | None:
"""SQL expression that builds ``search_vector`` for the configured PG text-search backend.
Single source of truth shared by the batch insert (over the ``input_data``
CTE columns) and the curation revert recompute (over a ``memory_units`` row),
so the two can never drift. Returns ``None`` for backends that leave
``search_vector`` unpopulated — pgroonga / pg_textsearch / pg_search index the
base text columns directly and keep only a dummy column, so there is nothing
to build.
``text_search_extension_native_language`` is validated as a PG identifier in
``HindsightConfig.validate()``, so embedding it as a SQL literal is safe.
"""
combined = f"COALESCE({text_col}, '') || ' ' || COALESCE({context_col}, '') || ' ' || COALESCE({signals_col}, '')"
if config.text_search_extension == "vchord":
return f"tokenize({combined}, 'llmlingua2')::bm25_catalog.bm25vector"
if config.text_search_extension == "native":
return f"to_tsvector('{config.text_search_extension_native_language}'::regconfig, {combined})"
return None
class PostgreSQLOps(DataAccessOps):
"""PostgreSQL-specific data access operations using unnest and LATERAL."""
@@ -122,39 +93,101 @@ class PostgreSQLOps(DataAccessOps):
config = get_config()
table = self._get_mu_table()
# search_vector is populated inline for backends that store a real vector
# (native tsvector, vchord bm25vector). pgroonga / pg_textsearch / pg_search
# index the base text columns directly and keep only a dummy column, so the
# expression is None and the column is left out of the insert entirely.
# Same expression is reused by curation revert (see pg_search_vector_expr).
sv_expr = pg_search_vector_expr(config)
sv_insert_col = ", search_vector" if sv_expr else ""
sv_select_val = f",\n {sv_expr}" if sv_expr else ""
query = f"""
WITH input_data AS (
SELECT * FROM unnest(
$2::text[], $3::vector[], $4::timestamptz[], $5::timestamptz[], $6::timestamptz[], $7::timestamptz[],
$8::text[], $9::text[], $10::jsonb[], $11::text[], $12::text[], $13::jsonb[], $14::jsonb[], $15::text[]
) AS t(text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
context, fact_type, metadata, chunk_id, document_id, tags_json,
observation_scopes_json, text_signals)
)
INSERT INTO {table} (bank_id, text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
context, fact_type, metadata, chunk_id, document_id, tags,
observation_scopes, text_signals{sv_insert_col})
SELECT
$1,
text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
context, fact_type, metadata, chunk_id, document_id,
COALESCE(
(SELECT array_agg(elem) FROM jsonb_array_elements_text(tags_json) AS elem),
'{{}}'::varchar[]
),
observation_scopes_json,
text_signals{sv_select_val}
FROM input_data
RETURNING id
"""
if config.text_search_extension == "vchord":
query = f"""
WITH input_data AS (
SELECT * FROM unnest(
$2::text[], $3::vector[], $4::timestamptz[], $5::timestamptz[], $6::timestamptz[], $7::timestamptz[],
$8::text[], $9::text[], $10::jsonb[], $11::text[], $12::text[], $13::jsonb[], $14::jsonb[], $15::text[]
) AS t(text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
context, fact_type, metadata, chunk_id, document_id, tags_json,
observation_scopes_json, text_signals)
)
INSERT INTO {table} (bank_id, text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
context, fact_type, metadata, chunk_id, document_id, tags,
observation_scopes, text_signals, search_vector)
SELECT
$1,
text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
context, fact_type, metadata, chunk_id, document_id,
COALESCE(
(SELECT array_agg(elem) FROM jsonb_array_elements_text(tags_json) AS elem),
'{{}}'::varchar[]
),
observation_scopes_json,
text_signals,
tokenize(
COALESCE(text, '') || ' ' || COALESCE(context, '') || ' ' || COALESCE(text_signals, ''),
'llmlingua2'
)::bm25_catalog.bm25vector
FROM input_data
RETURNING id
"""
elif config.text_search_extension == "native":
# search_vector is a regular tsvector column populated here using the
# configured native dictionary. It used to be GENERATED ALWAYS with
# a hardcoded 'english', which prevented per-deployment language
# configuration. text_search_extension_native_language is validated
# in HindsightConfig.validate() as a PG identifier, so embedding it
# as a SQL literal is safe.
query = f"""
WITH input_data AS (
SELECT * FROM unnest(
$2::text[], $3::vector[], $4::timestamptz[], $5::timestamptz[], $6::timestamptz[], $7::timestamptz[],
$8::text[], $9::text[], $10::jsonb[], $11::text[], $12::text[], $13::jsonb[], $14::jsonb[], $15::text[]
) AS t(text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
context, fact_type, metadata, chunk_id, document_id, tags_json,
observation_scopes_json, text_signals)
)
INSERT INTO {table} (bank_id, text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
context, fact_type, metadata, chunk_id, document_id, tags,
observation_scopes, text_signals, search_vector)
SELECT
$1,
text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
context, fact_type, metadata, chunk_id, document_id,
COALESCE(
(SELECT array_agg(elem) FROM jsonb_array_elements_text(tags_json) AS elem),
'{{}}'::varchar[]
),
observation_scopes_json,
text_signals,
to_tsvector(
'{config.text_search_extension_native_language}'::regconfig,
COALESCE(text, '') || ' ' || COALESCE(context, '') || ' ' || COALESCE(text_signals, '')
)
FROM input_data
RETURNING id
"""
else:
# pg_textsearch, pgroonga, and pg_search: search_vector is a dummy
# TEXT column; the actual full-text index operates on the base text
# columns directly, so we don't populate search_vector at insert time.
query = f"""
WITH input_data AS (
SELECT * FROM unnest(
$2::text[], $3::vector[], $4::timestamptz[], $5::timestamptz[], $6::timestamptz[], $7::timestamptz[],
$8::text[], $9::text[], $10::jsonb[], $11::text[], $12::text[], $13::jsonb[], $14::jsonb[], $15::text[]
) AS t(text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
context, fact_type, metadata, chunk_id, document_id, tags_json,
observation_scopes_json, text_signals)
)
INSERT INTO {table} (bank_id, text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
context, fact_type, metadata, chunk_id, document_id, tags,
observation_scopes, text_signals)
SELECT
$1,
text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
context, fact_type, metadata, chunk_id, document_id,
COALESCE(
(SELECT array_agg(elem) FROM jsonb_array_elements_text(tags_json) AS elem),
'{{}}'::varchar[]
),
observation_scopes_json,
text_signals
FROM input_data
RETURNING id
"""
results = await conn.fetch(
query,
@@ -182,6 +215,7 @@ class PostgreSQLOps(DataAccessOps):
table: str,
sorted_links: list[tuple],
bank_id: str,
nil_entity_uuid: str,
exists_clause: str,
chunk_size: int = 5000,
) -> None:
@@ -206,6 +240,7 @@ class PostgreSQLOps(DataAccessOps):
to_ids = [lnk[1] for lnk in sorted_links]
types = [lnk[2] for lnk in sorted_links]
weights = [lnk[3] for lnk in sorted_links]
entity_ids = [lnk[4] for lnk in sorted_links]
for chunk_start in range(0, len(sorted_links), chunk_size):
chunk_end = min(chunk_start + chunk_size, len(sorted_links))
@@ -219,23 +254,25 @@ class PostgreSQLOps(DataAccessOps):
f"""
WITH locked AS (
SELECT id FROM {mu_table}
WHERE id = ANY($6::uuid[])
WHERE id = ANY($7::uuid[])
ORDER BY id
FOR KEY SHARE
)
INSERT INTO {table}
(from_unit_id, to_unit_id, link_type, weight, bank_id)
SELECT f, t, tp, w, $5
FROM unnest($1::uuid[], $2::uuid[], $3::text[], $4::float8[])
AS u(f, t, tp, w)
(from_unit_id, to_unit_id, link_type, weight, entity_id, bank_id)
SELECT f, t, tp, w, e, $6
FROM unnest($1::uuid[], $2::uuid[], $3::text[], $4::float8[], $5::uuid[])
AS u(f, t, tp, w, e)
WHERE f IN (SELECT id FROM locked) AND t IN (SELECT id FROM locked)
ON CONFLICT (from_unit_id, to_unit_id, link_type)
ON CONFLICT (from_unit_id, to_unit_id, link_type,
COALESCE(entity_id, '{nil_entity_uuid}'::uuid))
DO NOTHING
""",
chunk_from,
chunk_to,
types[chunk_start:chunk_end],
weights[chunk_start:chunk_end],
entity_ids[chunk_start:chunk_end],
bank_id,
referenced,
timeout=300,
@@ -248,22 +285,12 @@ class PostgreSQLOps(DataAccessOps):
bank_id: str,
entity_names: list[str],
entity_dates: list,
entity_kinds: list[str],
) -> dict[str, str]:
# ORDER BY LOWER(name) so every concurrent batch inserts in the same order
# as the conflict target (bank_id, LOWER(canonical_name)). ON CONFLICT DO
# NOTHING takes a ShareLock on the inserting transaction of any speculative
# row it collides with, so two batches with overlapping names inserting in
# different orders deadlock. The caller already sorts by Python's
# ``str.lower()``, which agrees with the index for ASCII but not for every
# locale (see the Turkish-İ note in entity_resolver) — ordering in SQL makes
# the database's own collation the single arbiter for all writers.
inserted_rows = await conn.fetch(
f"""
INSERT INTO {table} (bank_id, canonical_name, first_seen, last_seen, mention_count, entity_kind)
SELECT $1, name, COALESCE(event_date, now()), COALESCE(event_date, now()), 0, kind
FROM unnest($2::text[], $3::timestamptz[], $4::text[]) AS t(name, event_date, kind)
ORDER BY LOWER(name)
INSERT INTO {table} (bank_id, canonical_name, first_seen, last_seen, mention_count)
SELECT $1, name, COALESCE(event_date, now()), COALESCE(event_date, now()), 0
FROM unnest($2::text[], $3::timestamptz[]) AS t(name, event_date)
ON CONFLICT (bank_id, LOWER(canonical_name))
DO NOTHING
RETURNING id, LOWER(canonical_name) AS name_lower
@@ -271,7 +298,6 @@ class PostgreSQLOps(DataAccessOps):
bank_id,
entity_names,
entity_dates,
entity_kinds,
)
return {row["name_lower"]: row["id"] for row in inserted_rows}
@@ -284,7 +310,7 @@ class PostgreSQLOps(DataAccessOps):
) -> list[ResultRow]:
return await conn.fetch(
f"""
SELECT e.id, e.canonical_name, LOWER(e.canonical_name) AS name_lower, inputs.input_name
SELECT e.id, LOWER(e.canonical_name) AS name_lower, inputs.input_name
FROM {table} e
JOIN (
SELECT LOWER(n) AS input_name_lower, n AS input_name
@@ -296,44 +322,6 @@ class PostgreSQLOps(DataAccessOps):
missing_names,
)
async def bulk_reassert_entities(
self,
conn: DatabaseConnection,
table: str,
bank_id: str,
entity_ids: list[str],
canonical_names: list[str],
entity_kinds: list[str],
) -> None:
# One statement, one round-trip (same shape as bulk_insert_links):
# * the CTE takes FOR KEY SHARE on every parent that still exists,
# held to COMMIT, so a concurrent prune_orphan_entities DELETE blocks
# until the caller's unit_entities insert has committed;
# * the INSERT re-creates only the parents that were already pruned
# (NOT IN locked), carrying the canonical_name resolved in Phase 1.
# ON CONFLICT DO NOTHING (no target) keeps the rare case where another
# worker recreated the name under a new id from raising — that row stays
# absent and its unit link is the sole casualty, never the whole batch.
await conn.execute(
f"""
WITH locked AS (
SELECT id FROM {table}
WHERE id = ANY($2::uuid[])
ORDER BY id
FOR KEY SHARE
)
INSERT INTO {table} (id, bank_id, canonical_name, entity_kind)
SELECT t.entity_id, $1, t.canonical_name, t.entity_kind
FROM unnest($2::uuid[], $3::text[], $4::text[]) AS t(entity_id, canonical_name, entity_kind)
WHERE t.entity_id NOT IN (SELECT id FROM locked)
ON CONFLICT DO NOTHING
""",
bank_id,
entity_ids,
canonical_names,
entity_kinds,
)
async def bulk_insert_unit_entities(
self,
conn: DatabaseConnection,
@@ -369,24 +357,11 @@ class PostgreSQLOps(DataAccessOps):
# concurrent caller the same lock order, so conflicting inserts
# queue cleanly instead of cycling.
sorted_unit_ids = sorted(unit_ids)
# DO UPDATE (not DO NOTHING) on a duplicate enqueue — #3034. The SET is a
# deliberate no-op that preserves enqueued_at; its only purpose is to take
# the existing row's lock. DO NOTHING does NOT lock the conflicting row, so
# a mutation that re-enqueues an already-queued unit could not block a
# worker from concurrently claiming (deleting) that row and processing the
# unit's pre-mutation state; the re-enqueue signal was then silently lost
# and the unit's derived links stayed stale with an empty queue. Locking
# the row serialises the mutation against the worker's claim for that
# (bank_id, unit_id): the worker either waits for the committed post-mutation
# state, or (if it claimed first) this INSERT lands a fresh row after the
# worker's delete commits. Row locks are acquired in sorted unit_id order,
# matching claim_graph_maintenance_batch, so the two never cycle.
await conn.execute(
f"""
INSERT INTO {table} (bank_id, unit_id)
SELECT $1, v FROM unnest($2::uuid[]) AS t(v)
ON CONFLICT (bank_id, unit_id)
DO UPDATE SET enqueued_at = {table}.enqueued_at
ON CONFLICT (bank_id, unit_id) DO NOTHING
""",
bank_id,
sorted_unit_ids,
@@ -399,35 +374,16 @@ class PostgreSQLOps(DataAccessOps):
bank_id: str,
limit: int,
) -> list[str]:
# Ordered locking (#3034). Choose the oldest batch by enqueued_at, but
# acquire the row locks in (bank_id, unit_id) order — the same order the
# enqueue upsert takes them — so a foreground mutation re-enqueueing an
# overlapping unit set can never cycle against a worker draining it. The
# `chosen` CTE is MATERIALIZED so the enqueued_at pick is fenced from the
# locking clause; `FOR UPDATE OF q ... ORDER BY q.unit_id` then puts
# LockRows above the Sort, so locks are taken ascending by unit_id (same
# idiom as prune_stale_cooccurrences' #2529 ordered lock). A concurrent
# enqueue holding one of these rows blocks this claim until it commits, at
# which point the worker deletes and processes the committed state.
rows = await conn.fetch(
f"""
WITH chosen AS MATERIALIZED (
DELETE FROM {table}
WHERE (bank_id, unit_id) IN (
SELECT bank_id, unit_id FROM {table}
WHERE bank_id = $1
ORDER BY enqueued_at
LIMIT $2
),
locked AS (
SELECT q.bank_id, q.unit_id
FROM {table} q
JOIN chosen c ON c.bank_id = q.bank_id AND c.unit_id = q.unit_id
ORDER BY q.unit_id
FOR UPDATE OF q
)
DELETE FROM {table} q
USING locked l
WHERE q.bank_id = l.bank_id AND q.unit_id = l.unit_id
RETURNING q.unit_id
RETURNING unit_id
""",
bank_id,
limit,
@@ -469,48 +425,19 @@ class PostgreSQLOps(DataAccessOps):
# Scope by joining through entities.bank_id (entity_cooccurrences itself
# has no bank_id column — entities don't span banks, so scoping via
# entity_id_1 is sufficient).
#
# Ordered locking (deadlock avoidance, #2529): retain's concurrent
# cooccurrence upsert (entity_resolver._flush_pending) locks rows in
# sorted (entity_id_1, entity_id_2) order — sorted specifically to give
# every writer one consistent lock-acquisition order. A plain
# `DELETE ... USING` scans/locks in whatever order the join plan picks,
# so it could lock the same rows in the opposite order and cycle. We
# instead select the victims in that same sorted order `FOR UPDATE`
# first — the locking clause materialises the CTE and places LockRows
# above the Sort, so locks are acquired ascending, matching the upsert —
# then delete the already-locked rows. Same order on both sides ⇒ no
# cycle (the deadlock is prevented, not merely retried). The Pass 2/3
# retry wrap in run_graph_maintenance_job stays as a backstop for the
# residual paths (FK cascade from prune_orphan_entities, Oracle).
#
# The staleness predicate is an INTERSECT of the two entities' unit sets
# rather than the equivalent `unit_entities u1 JOIN u2 ON u1.unit_id =
# u2.unit_id` self-join (#2473): both INTERSECT branches resolve as Index
# Only Scans on idx_unit_entities_entity_unit (entity_id, unit_id), so the
# per-pair cost is bounded by the two entities' degrees. The self-join let
# the planner pick an anti-join that rescanned a high-degree hub entity's
# membership set for every pair — 28-30min on a bank with a ~100K-membership
# hub, even when zero rows were stale. Don't "simplify" it back.
result = await conn.execute(
f"""
WITH victims AS (
SELECT c.entity_id_1, c.entity_id_2
FROM {ec_table} c
JOIN {entities_table} e ON e.id = c.entity_id_1
WHERE e.bank_id = $1
AND NOT EXISTS (
SELECT unit_id FROM {ue_table} WHERE entity_id = c.entity_id_1
INTERSECT
SELECT unit_id FROM {ue_table} WHERE entity_id = c.entity_id_2
)
ORDER BY c.entity_id_1, c.entity_id_2
FOR UPDATE OF c
)
DELETE FROM {ec_table} c
USING victims v
WHERE c.entity_id_1 = v.entity_id_1
AND c.entity_id_2 = v.entity_id_2
USING {entities_table} e
WHERE e.id = c.entity_id_1
AND e.bank_id = $1
AND NOT EXISTS (
SELECT 1
FROM {ue_table} u1
JOIN {ue_table} u2 ON u1.unit_id = u2.unit_id
WHERE u1.entity_id = c.entity_id_1
AND u2.entity_id = c.entity_id_2
)
""",
bank_id,
)
@@ -522,21 +449,11 @@ class PostgreSQLOps(DataAccessOps):
mu_table: str,
unit_ids: list[str],
) -> list[ResultRow]:
# Cast only canonical UUID text inputs, never the indexed column. The old
# ``id::text`` predicate silently ignored malformed, uppercase, braced,
# and unhyphenated inputs; filtering before the cast preserves that
# behavior while allowing the primary-key index to serve the lookup.
return await conn.fetch(
f"""
SELECT id, event_date, fact_type
FROM {mu_table}
WHERE id = ANY(
ARRAY(
SELECT input.unit_id::uuid
FROM unnest($1::text[]) AS input(unit_id)
WHERE input.unit_id ~ '^[0-9a-f]{{8}}-[0-9a-f]{{4}}-[0-9a-f]{{4}}-[0-9a-f]{{4}}-[0-9a-f]{{12}}$'
)
)
WHERE id::text = ANY($1)
""",
unit_ids,
)
@@ -605,7 +522,6 @@ class PostgreSQLOps(DataAccessOps):
mu_table: str,
ue_table: str,
per_entity_limit: int,
window: UpdatedWindow,
) -> str:
return f"""
seed_entities AS (
@@ -625,20 +541,11 @@ class PostgreSQLOps(DataAccessOps):
FROM {ue_table} ue_target
WHERE ue_target.entity_id = se.entity_id
AND ue_target.unit_id != ALL($1::uuid[])
-- Filter before applying the cap: candidates from other fact
-- types, or outside the recall window, must not consume this
-- entity's bounded fan-out.
AND EXISTS (
SELECT 1
FROM {mu_table} mu_target
WHERE mu_target.id = ue_target.unit_id
AND mu_target.fact_type = $2
{window.clause("mu_target")}
)
ORDER BY ue_target.unit_id DESC
LIMIT {per_entity_limit}
) t
JOIN {mu_table} mu ON mu.id = t.unit_id
WHERE mu.fact_type = $2
GROUP BY mu.id
ORDER BY score DESC
LIMIT $3
@@ -648,7 +555,6 @@ class PostgreSQLOps(DataAccessOps):
self,
ml_table: str,
mu_table: str,
window: UpdatedWindow,
) -> str:
# Exact v0.5.6 query shape: GROUP BY + MAX(weight) for semantic,
# DISTINCT ON for causal.
@@ -672,7 +578,6 @@ class PostgreSQLOps(DataAccessOps):
AND ml.link_type = 'semantic'
AND mu.fact_type = $2
AND mu.id != ALL($1::uuid[])
{window.clause("mu")}
UNION ALL
SELECT
mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
@@ -685,7 +590,6 @@ class PostgreSQLOps(DataAccessOps):
AND ml.link_type = 'semantic'
AND mu.fact_type = $2
AND mu.id != ALL($1::uuid[])
{window.clause("mu")}
) sem_raw
GROUP BY id, text, context, event_date, occurred_start,
occurred_end, mentioned_at,
@@ -705,7 +609,6 @@ class PostgreSQLOps(DataAccessOps):
WHERE ml.from_unit_id = ANY($1::uuid[])
AND ml.link_type IN ('causes', 'caused_by', 'enables', 'prevents')
AND mu.fact_type = $2
{window.clause("mu")}
ORDER BY mu.id, ml.weight DESC
LIMIT $3
)"""
@@ -719,13 +622,8 @@ class PostgreSQLOps(DataAccessOps):
seed_ids: list,
budget: int,
per_entity_limit: int,
window: UpdatedWindow,
) -> LinkExpansionRows:
) -> tuple[list[ResultRow], list[ResultRow], list[ResultRow]]:
# v0.5.6 array ops: unnest, &&, COUNT(DISTINCT) on source_memory_ids.
#
# The window bounds the observations that come *back*, not the source facts
# traversed to reach them: an observation is in the window when it was itself
# written or refreshed there, regardless of how old the facts underneath it are.
entity_rows = await conn.fetch(
f"""
@@ -767,13 +665,11 @@ class PostgreSQLOps(DataAccessOps):
AND mu.id != ALL($1::uuid[])
AND ca.source_ids IS NOT NULL
AND mu.source_memory_ids && ca.source_ids
{window.clause("mu")}
ORDER BY score DESC
LIMIT $2
""",
seed_ids,
budget,
*window.params,
)
# Exact v0.5.6 query shape: GROUP BY + MAX(weight) for semantic,
@@ -795,7 +691,6 @@ class PostgreSQLOps(DataAccessOps):
WHERE ml.from_unit_id = ANY($1::uuid[])
AND ml.link_type = 'semantic' AND mu.fact_type = 'observation'
AND mu.id != ALL($1::uuid[])
{window.clause("mu")}
UNION ALL
SELECT mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
mu.occurred_end, mu.mentioned_at, mu.fact_type, mu.document_id,
@@ -804,7 +699,6 @@ class PostgreSQLOps(DataAccessOps):
WHERE ml.to_unit_id = ANY($1::uuid[])
AND ml.link_type = 'semantic' AND mu.fact_type = 'observation'
AND mu.id != ALL($1::uuid[])
{window.clause("mu")}
) sem_raw
GROUP BY id, text, context, event_date, occurred_start, occurred_end,
mentioned_at, fact_type, document_id, chunk_id, tags, proof_count
@@ -819,7 +713,6 @@ class PostgreSQLOps(DataAccessOps):
WHERE ml.from_unit_id = ANY($1::uuid[])
AND ml.link_type IN ('causes', 'caused_by', 'enables', 'prevents')
AND mu.fact_type = 'observation'
{window.clause("mu")}
ORDER BY mu.id, ml.weight DESC LIMIT $2
)
SELECT * FROM semantic_expanded
@@ -828,12 +721,11 @@ class PostgreSQLOps(DataAccessOps):
""",
seed_ids,
budget,
*window.params,
)
semantic_rows = [r for r in sem_causal_rows if r["source"] == "semantic"]
causal_rows = [r for r in sem_causal_rows if r["source"] == "causal"]
return 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(
@@ -869,16 +761,10 @@ class PostgreSQLOps(DataAccessOps):
internal_id: str,
fact_types: dict[str, str],
) -> None:
# CONCURRENTLY so the drop takes ShareUpdateExclusive, not ACCESS
# EXCLUSIVE, on the shared memory_units table. A plain DROP INDEX blocks
# (and deadlocks with) every other bank's concurrent reads/writes on the
# table; CONCURRENTLY does not conflict with DML. The caller
# (delete_bank) runs this on an autocommit connection after its delete
# transaction has committed — CONCURRENTLY cannot run inside a tx.
for ft, suffix in fact_types.items():
uid = str(internal_id).replace("-", "")[:16]
idx = f"idx_mu_emb_{suffix}_{uid}"
await conn.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}.{idx}")
await conn.execute(f"DROP INDEX IF EXISTS {schema}.{idx}")
def get_entity_resolution_strategy(self) -> str:
return "trigram"
@@ -1008,93 +894,6 @@ class PostgreSQLOps(DataAccessOps):
# -- Task claiming operations ------------------------------------------
async def prune_terminal_operations(
self,
conn: DatabaseConnection,
table: str,
cutoff: datetime,
*,
batch_size: int,
) -> int:
# Lock only the bounded candidate set. SKIP LOCKED lets multiple
# workers prune disjoint batches without waiting or double-deleting.
# Cancelled children cannot complete parent aggregation, so retain the
# parent guard only for completed/failed children. Before removing a
# cancelled child, preserve its signal by cancelling a pending parent
# in this transaction and refreshing the parent's retention window.
candidates = await conn.fetch(
f"""
SELECT candidate_operation.operation_id
FROM {table} candidate_operation
WHERE candidate_operation.status IN ('completed', 'failed', 'cancelled')
AND candidate_operation.updated_at < $1
AND (
candidate_operation.status = 'cancelled'
OR NOT EXISTS (
SELECT 1
FROM {table} parent
WHERE parent.operation_id = CASE
WHEN candidate_operation.result_metadata->>'parent_operation_id'
~* '^[0-9a-f]{{8}}-[0-9a-f]{{4}}-[0-9a-f]{{4}}-[0-9a-f]{{4}}-[0-9a-f]{{12}}$'
THEN (candidate_operation.result_metadata->>'parent_operation_id')::uuid
ELSE NULL
END
AND parent.bank_id = candidate_operation.bank_id
)
)
ORDER BY candidate_operation.updated_at, candidate_operation.operation_id
LIMIT $2
FOR UPDATE OF candidate_operation SKIP LOCKED
""",
cutoff,
batch_size,
)
if not candidates:
return 0
candidate_ids = [row["operation_id"] for row in candidates]
await conn.execute(
f"""
UPDATE {table} parent
SET status = 'cancelled',
updated_at = now(),
completed_at = COALESCE(parent.completed_at, now()),
error_message = COALESCE(
parent.error_message,
'Cancelled because a child operation was cancelled'
)
WHERE parent.status = 'pending'
AND EXISTS (
SELECT 1
FROM {table} candidate_operation
WHERE candidate_operation.operation_id = ANY($1)
AND candidate_operation.status = 'cancelled'
AND candidate_operation.updated_at < $2
AND candidate_operation.bank_id = parent.bank_id
AND parent.operation_id = CASE
WHEN candidate_operation.result_metadata->>'parent_operation_id'
~* '^[0-9a-f]{{8}}-[0-9a-f]{{4}}-[0-9a-f]{{4}}-[0-9a-f]{{4}}-[0-9a-f]{{12}}$'
THEN (candidate_operation.result_metadata->>'parent_operation_id')::uuid
ELSE NULL
END
)
""",
candidate_ids,
cutoff,
)
rows = await conn.fetch(
f"""
DELETE FROM {table}
WHERE operation_id = ANY($1)
AND status IN ('completed', 'failed', 'cancelled')
AND updated_at < $2
RETURNING operation_id
""",
candidate_ids,
cutoff,
)
return len(rows)
async def _claim_consolidation_tasks(
self,
conn,
@@ -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,
)
@@ -15,7 +15,6 @@ from typing import Any
import asyncpg # noqa: F401
from .base import DatabaseBackend, DatabaseConnection
from .pool_instrumentation import PoolStats, instrument_acquire
logger = logging.getLogger(__name__)
@@ -77,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,
@@ -91,16 +88,6 @@ class PostgreSQLBackend(DatabaseBackend):
statement_cache_size: int = 0,
init_callback: Any | None = None,
) -> None:
from ...config import get_config
self._acquire_warn_threshold_s = get_config().db_acquire_warn_threshold_ms / 1000.0
# Kept for acquire() below: asyncpg's ``timeout`` create_pool kwarg is a
# *connect* kwarg (how long establishing a new connection may take), and
# ``Pool.acquire()`` defaults to waiting for a free connection forever.
# Passing it here alone made HINDSIGHT_API_DB_ACQUIRE_TIMEOUT a no-op for
# the wait it names: a pool-exhaustion stall never surfaced as an error,
# it just hung (#3002). 0 restores the unbounded behaviour.
self._acquire_timeout_s = acquire_timeout if acquire_timeout > 0 else None
self._pool = await asyncpg.create_pool(
dsn,
min_size=min_size,
@@ -108,12 +95,7 @@ class PostgreSQLBackend(DatabaseBackend):
command_timeout=command_timeout,
statement_cache_size=statement_cache_size,
timeout=acquire_timeout,
# init runs once per new connection; setup runs on every acquire,
# after asyncpg's release-time RESET ALL. Passing init_callback as
# both keeps the per-connection session GUCs (hnsw.ef_search, etc.)
# applied after a connection is reused, not just on first creation.
init=init_callback,
setup=init_callback,
)
logger.info(
f"PostgreSQL pool created (min={min_size}, max={max_size}, "
@@ -121,45 +103,21 @@ class PostgreSQLBackend(DatabaseBackend):
)
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__)
@@ -142,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.
@@ -160,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:
@@ -197,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
@@ -228,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]]:
"""
@@ -241,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):
@@ -514,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(
@@ -523,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
@@ -1628,7 +1585,6 @@ def create_embeddings_from_env() -> Embeddings:
model_name=config.embeddings_local_model,
force_cpu=config.embeddings_local_force_cpu,
trust_remote_code=config.embeddings_local_trust_remote_code,
allow_mps=config.embeddings_local_allow_mps,
)
elif provider == "onnx":
return OnnxEmbeddings(
@@ -6,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,32 +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 _link_units_to_entities_batch_impl(
self, conn, unit_entity_pairs: list[tuple[str, str, datetime | None]], bank_id: str | None = 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,
await self._ops.bulk_insert_unit_entities(
conn,
fq_table("unit_entities"),
unit_ids,
entity_ids,
)
# Build maps keyed by unit_id:
@@ -1288,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
@@ -1305,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
@@ -5,25 +5,25 @@ 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. **Orphan entity prune.** Delete ``entities`` rows in the bank that no
longer have any live memory references. FK ON DELETE CASCADE on
longer have any ``unit_entities`` references. FK ON DELETE CASCADE on
``entity_cooccurrences`` then removes any cooccurrence row pointing
at the pruned entities.
3. **Stale cooccurrence prune.** Defensive sweep for cooccurrence rows
where both endpoints still exist but no current memory references
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 (drain the queue, wrap the sweep in a
deadlock-retry) and asks the store to do the part that touches storage. A store
whose links travel inside its memories has no `memory_links` to dangle and no
join table to sweep, so its relink and cooccurrence passes are no-ops and the
job simply prunes the orphan `entities` rows, which stay in Postgres regardless.
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
@@ -35,18 +35,18 @@ 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:
@@ -54,19 +54,17 @@ if TYPE_CHECKING:
logger = logging.getLogger(__name__)
# Retry budget for the idempotent Pass 2/3 entity/cooccurrence sweep. Higher
# than db_utils' default (3) because the sweep has no client waiting on it and
# is safe to rerun, so we'd rather spend a longer jittered-backoff tail than
# drop a maintenance pass and leak stale graph rows (see run_graph_maintenance_job).
_SWEEP_MAX_RETRIES = 8
# 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
@dataclass
class _SweepCounts:
"""Prune counts returned by the Pass 2/3 sweep (avoids a bare tuple return)."""
orphan_entities_pruned: int
stale_cooccurrences_pruned: int
# 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
@@ -90,52 +88,67 @@ class JobResult:
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]
if not victim_ids:
return 0
await ops.enqueue_graph_maintenance(
conn,
fq_table("graph_maintenance_queue"),
bank_id,
victim_ids,
)
logger.debug(
f"[GRAPH_MAINT] Enqueued {len(victim_ids)} relink victims in "
f"bank={bank_id} (deleted {len(deleted_unit_ids)} units)"
)
return len(victim_ids)
async def run_graph_maintenance_job(
memory_engine: "MemoryEngine",
@@ -150,69 +163,195 @@ async def run_graph_maintenance_job(
Per-pass counters from :class:`JobResult`.
"""
del request_context # accepted for symmetry with other run_*_job helpers
from ..config import get_config
from .memories import get_memories
backend = await memory_engine._get_backend()
store = get_memories()
config = get_config()
ops = backend.ops
result = JobResult()
job_start = time.time()
# --- 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)
result.relink_units_processed = relink.get("relink_units_processed", 0)
result.relink_links_added = relink.get("relink_links_added", 0)
# Per-iteration loop: claim → top up → commit. We rely on submit-time
# dedup to keep at most one job per bank running, so no need for
# SKIP LOCKED.
iterations = 0
while True:
from .memory_engine import acquire_with_retry
async with acquire_with_retry(backend) as conn:
async with conn.transaction():
unit_ids = await ops.claim_graph_maintenance_batch(
conn,
fq_table("graph_maintenance_queue"),
bank_id,
_DRAIN_BATCH_SIZE,
)
if not unit_ids:
break
result.relink_links_added += await _relink_batch(conn, bank_id, unit_ids, ops, backend)
result.relink_units_processed += len(unit_ids)
iterations += 1
if iterations > 10000:
# Defensive guard against runaway loops — at 50 units/iter that's
# 500k targets, far beyond any realistic single-bank backlog.
logger.error(
f"[GRAPH_MAINT] bank={bank_id} hit iteration cap ({iterations}); aborting relink ({result.as_dict()})"
)
break
# --- Pass 2 & 3: entity / cooccurrence sweeps ---
# Bank-wide single-statement deletes. Cheap when there's nothing to do.
#
# Unlike Pass 1's queue claim, these DELETEs aren't protected by any
# consistent lock-ordering guarantee: the stale-cooccurrence prune scans
# entity_cooccurrences via a join/NOT EXISTS plan, while retain's concurrent
# cooccurrence upserts (entity_resolver._flush_pending) lock the same rows in
# sorted (entity_id_1, entity_id_2) order. When a sweep and a concurrent
# upsert touch overlapping rows in opposite orders, Postgres detects a
# genuine circular wait and aborts one side with DeadlockDetectedError. Both
# prunes are idempotent bank-wide sweeps — rerunning only deletes what's
# still stale — so retrying the whole transaction on deadlock is safe.
#
# The prunes themselves are the store's: the orphan-`entities` sweep applies
# to every store (that registry stays in Postgres), while the cooccurrence
# sweep is a no-op for a store that never wrote `unit_entities`.
from .db_utils import retry_with_backoff
from .memory_engine import acquire_with_retry
async def _run_sweep() -> _SweepCounts:
async with acquire_with_retry(backend) as conn:
async with conn.transaction():
orphan_pruned = await store.prune_orphan_entities(conn=conn, fq_table=fq_table, bank_id=bank_id)
# The orphan prune above cascades cooccurrences via FK. The
# explicit cooccurrence pass below catches the *stale-count*
# case: both entities still exist but no current unit witnesses
# them together.
stale_pruned = await store.prune_stale_cooccurrences(conn=conn, fq_table=fq_table, bank_id=bank_id)
return _SweepCounts(orphan_entities_pruned=orphan_pruned, stale_cooccurrences_pruned=stale_pruned)
# A larger retry budget than the default (3): this is idempotent background
# maintenance with no client waiting on it, so a longer retry tail costs
# nothing, whereas a dropped sweep silently leaks orphan entities / stale
# cooccurrences until the next run. With jittered backoff a single sweep
# contending against continuous retain upserts effectively never exhausts
# this budget (each retry independently clears with high probability).
sweep = await retry_with_backoff(_run_sweep, max_retries=_SWEEP_MAX_RETRIES)
result.orphan_entities_pruned = sweep.orphan_entities_pruned
result.stale_cooccurrences_pruned = sweep.stale_cooccurrences_pruned
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
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:
@@ -6,54 +6,12 @@ enabling support for multiple LLM backends (OpenAI, Anthropic, Gemini, Codex, et
"""
from abc import ABC, abstractmethod
from contextlib import AbstractAsyncContextManager
from dataclasses import dataclass
from datetime import datetime
from enum import StrEnum
from typing import Any, Callable, Self
from typing import Any
from .response_models import LLMToolCallResult
class LLMToolChoiceMode(StrEnum):
"""Canonical tool-selection modes shared by every LLM provider."""
AUTO = "auto"
NONE = "none"
REQUIRED = "required"
NAMED = "named"
@dataclass(frozen=True, slots=True)
class LLMToolChoice:
"""Typed internal tool selection serialized only at provider boundaries."""
mode: LLMToolChoiceMode
function_name: str | None = None
def __post_init__(self) -> None:
if self.mode is LLMToolChoiceMode.NAMED:
if self.function_name is None or not self.function_name or self.function_name != self.function_name.strip():
raise ValueError("Named tool choice requires a non-empty canonical function name")
elif self.function_name is not None:
raise ValueError(f"Tool choice mode {self.mode.value!r} cannot include a function name")
@classmethod
def named(cls, function_name: str) -> Self:
return cls(mode=LLMToolChoiceMode.NAMED, function_name=function_name)
@property
def selected_function_name(self) -> str:
if self.function_name is None:
raise ValueError("Tool choice does not select a named function")
return self.function_name
LLM_TOOL_CHOICE_AUTO = LLMToolChoice(mode=LLMToolChoiceMode.AUTO)
LLM_TOOL_CHOICE_NONE = LLMToolChoice(mode=LLMToolChoiceMode.NONE)
LLM_TOOL_CHOICE_REQUIRED = LLMToolChoice(mode=LLMToolChoiceMode.REQUIRED)
class LLMInterface(ABC):
"""
Abstract interface for LLM providers.
@@ -113,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.
@@ -135,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.
@@ -162,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.
@@ -179,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.
@@ -197,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:
@@ -241,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,18 +27,13 @@ from ..config import (
ENV_REFLECT_LLM_MAX_CONCURRENT,
ENV_RETAIN_LLM_MAX_CONCURRENT,
)
from .llm_interface import (
LLM_TOOL_CHOICE_AUTO,
LLMToolChoice,
LLMToolChoiceMode,
)
from .llm_interface import (
OutputTooLongError as OutputTooLongError,
)
if TYPE_CHECKING:
from .response_models import LLMToolCallResult
# Seed applied to every Groq request for deterministic behavior.
DEFAULT_LLM_SEED = 4242
logger = logging.getLogger(__name__)
# Disable httpx logging
@@ -114,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.
@@ -156,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
@@ -192,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:
@@ -207,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.
@@ -223,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()
@@ -241,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(
@@ -279,17 +235,6 @@ 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,
@@ -309,7 +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,
) -> Any: # Returns LLMInterface
"""
Factory function to create the appropriate LLM provider implementation.
@@ -324,8 +268,6 @@ def create_llm_provider(
openai_service_tier: OpenAI service tier (for OpenAI provider) - None (default) or "flex" (50% cheaper).
bedrock_service_tier: Bedrock service tier (for Bedrock provider) - None (default), "flex", "priority", or "reserved".
gemini_service_tier: Gemini service tier (for Gemini provider) - None (default) or "flex" (50% cheaper).
ollama_num_ctx: Native Ollama context window override. None lets Ollama use the
model/server default.
extra_body: Extra request-body params merged into the provider's native
call. Threaded into OpenAI-compatible, Fireworks, Anthropic, Gemini/
VertexAI and LiteLLM providers (each merges them in its own parameter
@@ -349,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,
@@ -363,7 +303,6 @@ def create_llm_provider(
MockLLM,
NoneLLM,
OpenAICompatibleLLM,
OpenAIResponsesLLM,
)
provider_lower = provider.lower()
@@ -531,22 +470,6 @@ def create_llm_provider(
timeout=timeout,
)
elif provider_lower == "openai-responses":
# OpenAI Responses API (/v1/responses). Unlike chat/completions, it
# supports reasoning + function tools together, so reflect's tool loop
# can run with a real reasoning_effort. See OpenAIResponsesLLM.
return OpenAIResponsesLLM(
provider=provider,
api_key=api_key,
base_url=base_url,
model=model,
reasoning_effort=reasoning_effort,
openai_service_tier=openai_service_tier,
extra_body=extra_body,
default_headers=default_headers,
timeout=timeout,
)
elif provider_lower in (
"openai",
"groq",
@@ -571,7 +494,6 @@ def create_llm_provider(
groq_service_tier=groq_service_tier,
openai_service_tier=openai_service_tier,
extra_body=extra_body,
ollama_num_ctx=ollama_num_ctx,
timeout=timeout,
)
@@ -609,7 +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,
):
"""
Initialize LLM provider.
@@ -624,8 +545,6 @@ class LLMProvider:
openai_service_tier: OpenAI service tier (None or "flex") - from config.
bedrock_service_tier: Bedrock service tier (None, "flex", "priority", "reserved") - from config.
gemini_service_tier: Gemini service tier (None or "flex") - from config.
ollama_num_ctx: Native Ollama context window override. ``None`` lets Ollama
use the model/server default.
gemini_safety_settings: Safety settings for Gemini/VertexAI providers.
extra_body: Extra request-body params merged into the provider's native call
(OpenAI-compatible, Fireworks, Anthropic, Gemini/VertexAI, LiteLLM).
@@ -679,7 +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
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
@@ -697,7 +615,6 @@ class LLMProvider:
# Validate provider
valid_providers = [
"openai",
"openai-responses",
"groq",
"ollama",
"ollama-cloud",
@@ -825,7 +742,6 @@ 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,
)
@@ -887,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:
@@ -908,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:
@@ -929,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.
@@ -952,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
@@ -992,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
@@ -1012,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,
@@ -1026,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
@@ -1071,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.
@@ -1090,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.
@@ -1138,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,
@@ -1171,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
@@ -1382,7 +1260,6 @@ 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,
@@ -1393,7 +1270,6 @@ class LLMProvider:
ENV_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY,
_get_default_model_for_provider,
_parse_llm_router_config,
_parse_optional_positive_int,
parse_gemini_service_tier,
)
@@ -1438,7 +1314,6 @@ 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,
@@ -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)
@@ -21,16 +21,9 @@ 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. Retention and operation cleanup are deletes; the consolidation reconcile
and the scheduled mental model refresh both dedupe against in-flight operations
inside the inserting transaction (see ``_submit_async_operation``).
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
@@ -39,13 +32,13 @@ import asyncio
import logging
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
@@ -56,18 +49,6 @@ logger = logging.getLogger(__name__)
_TICK_SECONDS = 60
# Retention sweeps are not time-sensitive; hourly matches the previous per-sweep cadence.
_RETENTION_INTERVAL_SECONDS = 3600
# Operation cleanup deletes one bounded batch per schema per run, so its cadence
# sets the drain rate for a backlog. Kept at one-per-tick (the value it used while
# it rode the worker's poll loop) so throughput is unchanged by the move.
_OPERATION_CLEANUP_INTERVAL_SECONDS = 60
# Cross-store txn recovery (only when the memories store keeps its rows outside SQL): a backstop
# for a writer that crashed between its external writes and the decide. The happy path decides
# inline after commit, so this rarely finds work; five minutes bounds how long a crashed txn stalls
# its namespace's fold.
_TXN_RECOVERY_INTERVAL_SECONDS = 300
# A pending txn is left alone for this long from first sighting before the sweep aborts an
# unwitnessed one — the writer may still be mid-flight (PendingTxn carries no timestamp).
_TXN_RECOVERY_GRACE_SECONDS = 300
class MaintenanceLoop:
@@ -79,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 ──────────────────────────────────────────────────────────
@@ -119,32 +97,10 @@ 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.
audit_on = cfg.audit_log_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_retention_days > 0
return (
reconcile_on
or audit_on
or llm_on
or mm_refresh_on
or op_cleanup_on
or MaintenanceLoop._cross_store_recovery_enabled()
)
@staticmethod
def _cross_store_recovery_enabled() -> bool:
"""True when the memories store keeps memories outside SQL and therefore has
cross-store write-group txns a crashed writer could leave undecided."""
try:
from .memories import get_memories
return not get_memories().writes_memory_rows_in_sql
except Exception:
return False
return reconcile_on or audit_on or llm_on or mm_refresh_on
# ── loop ───────────────────────────────────────────────────────────────
@@ -178,10 +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())
if cfg.operation_retention_days > 0 and self._is_due("operation_cleanup", _OPERATION_CLEANUP_INTERVAL_SECONDS):
await self._run_timed("operation cleanup", self._run_operation_cleanup(cfg))
if self._cross_store_recovery_enabled() and self._is_due("txn_recovery", _TXN_RECOVERY_INTERVAL_SECONDS):
await self._run_timed("cross-store txn recovery", self._run_txn_recovery())
async def _run_timed(self, name: str, coro: Coroutine[Any, Any, None]) -> None:
"""Run a maintenance job and emit one timing line for it.
@@ -200,10 +152,7 @@ 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)
@@ -214,7 +163,7 @@ class MaintenanceLoop:
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
)
for row in rows:
schema = row[0]
@@ -229,109 +178,6 @@ class MaintenanceLoop:
except Exception as e:
logger.warning(f"Retention sweep failed for {table}: {e}")
# ── terminal operation cleanup ─────────────────────────────────────────
async def _run_operation_cleanup(self, cfg: HindsightConfig) -> None:
"""Prune one bounded batch of expired terminal operations per tenant schema.
Previously this rode the worker's task-claiming loop, so it only fired
when that loop happened to iterate and was interleaved with claiming. It
is a periodic housekeeping sweep like the retention jobs above, so it
belongs on the same schedule.
Discovery is one cross-tenant round-trip (``schemas_with_expired_operations``)
rather than a connection + prune transaction per tenant; pending and
processing rows are never prunable, so a schema holding only in-flight
work is correctly reported as having nothing to do.
"""
engine = self._engine
backend = engine._backend
try:
async with acquire_with_retry(backend, max_retries=1) as conn:
rows = await conn.fetch(
f"SELECT * FROM {fq_routine('schemas_with_expired_operations')}($1)",
cfg.operation_retention_days,
)
except Exception as e:
logger.warning(f"Operation cleanup discovery failed: {e}")
return
if not rows:
return
# Prune only schemas the deployment actually serves. The routine reports
# every schema owning an async_operations table, including ones tenant
# discovery doesn't claim.
try:
tenants = await engine._tenant_extension.list_tenants()
except Exception as e:
logger.warning(f"Operation cleanup tenant discovery failed: {e}")
return
known = {t.schema for t in tenants} | {get_config().database_schema}
from .memory_engine import _current_schema
cutoff = datetime.now(timezone.utc) - timedelta(days=cfg.operation_retention_days)
pruned = 0
for row in rows:
schema = row[0]
if schema not in known:
continue
# Oracle resolves unqualified names from a context-bound session
# schema; on PostgreSQL this is harmless and fq_table stays explicit.
token = _current_schema.set(schema)
try:
table = fq_table_explicit("async_operations", schema)
async with acquire_with_retry(backend, max_retries=1) as conn:
async with conn.transaction():
deleted = await backend.ops.prune_terminal_operations(
conn, table, cutoff, batch_size=cfg.operation_cleanup_batch_size
)
if deleted:
pruned += deleted
logger.info(f"Operation cleanup pruned {deleted} expired terminal operations from {schema}")
except Exception as e:
logger.warning(f"Operation cleanup failed for schema {schema}: {e}")
finally:
_current_schema.reset(token)
if pruned:
logger.info(f"Operation cleanup: pruned {pruned} operation(s) total")
# ── cross-store txn recovery ─────────────────────────────────────────────
async def _run_txn_recovery(self) -> None:
"""Resolve write-group txns a crashed writer left undecided, for a store that keeps its
rows outside SQL.
For each bank, the store lists its namespace's pending txns and decides each against the
Postgres witness table (present commit, absent past the grace abort never on
assumption), then reaps expired witness rows. A no-op for the SQL stores. Best-effort: a
failure here only delays a stalled fold until the next tick.
"""
from .memories import get_memories
store = get_memories()
if store.writes_memory_rows_in_sql:
return
backend = self._engine._backend
try:
async with acquire_with_retry(backend, max_retries=1) as conn:
bank_ids = [r[0] for r in await conn.fetch(f"SELECT bank_id FROM {fq_table('banks')}")]
if not bank_ids:
return
decided = await store.recover_pending_txns(
conn=conn,
fq_table=fq_table,
bank_ids=bank_ids,
first_seen=self._txn_first_seen,
now=time.monotonic(),
grace_seconds=_TXN_RECOVERY_GRACE_SECONDS,
)
except Exception as e:
logger.warning(f"Cross-store txn recovery failed: {e}")
return
if decided:
logger.info(f"Cross-store txn recovery: decided {decided} undecided txn(s)")
# ── consolidation reconcile ──────────────────────────────────────────────
async def _run_reconcile(self) -> None:
@@ -339,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
@@ -400,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
@@ -412,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}")
@@ -456,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"]
@@ -487,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,86 +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,
ScanPage,
StoredMemory,
build_fact_records,
build_text_signals,
source_key,
)
logger = logging.getLogger(__name__)
_memories: MemoriesExtension | None = None
def create_memories(context=None) -> MemoriesExtension:
"""Build the configured memories store, or the Postgres default."""
from ...extensions.loader import load_extension
loaded = load_extension("MEMORIES", MemoriesExtension, context=context)
if loaded is not None:
logger.info("[memories] store=%s (memory rows do not go to postgres)", loaded.name)
return loaded
from .postgres import PostgresMemories
return PostgresMemories({})
def get_memories() -> MemoriesExtension:
"""The process-wide memories store, built on first use.
Retrieval and the retain pipeline reach it through call chains that do not
carry the engine, so it is resolved here rather than threaded through every
signature.
"""
global _memories
if _memories is None:
_memories = create_memories()
return _memories
def set_memories(memories: MemoriesExtension | None) -> None:
"""Override the store (tests, and engine startup after initialize())."""
global _memories
_memories = memories
# The graph arm's retriever is chosen from the store and then cached, so it
# has to be re-resolved whenever the store changes.
from ..search.retrieval import set_default_graph_retriever
set_default_graph_retriever(None)
__all__ = [
"META_CHUNK_ID",
"CausalEdgeRecord",
"DeletePredicate",
"FactRecord",
"MemoriesExtension",
"MemoryPatch",
"ScanPage",
"StoredMemory",
"build_fact_records",
"build_text_signals",
"create_memories",
"get_memories",
"set_memories",
"source_key",
]
File diff suppressed because it is too large Load Diff
@@ -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,157 +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.
All four come from a single pass so keeping ``failed`` part of the
published contract costs nothing over reflect()'s ``pending`` read, and
``last_memory_write_at`` (the newest ``updated_at`` anywhere in the bank)
rides along for free. That watermark is what lets a caller decide a mental
model is up to date without running its own scoped scan: nothing in the bank
changed since the refresh, so nothing in the model's scope did either.
"""
row = await conn.fetchrow(
f"""
SELECT
MAX(consolidated_at) AS last_consolidated_at,
MAX(updated_at) AS last_memory_write_at,
COUNT(*) FILTER (WHERE consolidated_at IS NULL AND fact_type IN ('experience', 'world')) AS pending,
COUNT(*) FILTER (WHERE consolidation_failed_at IS NOT NULL AND fact_type IN ('experience', 'world')) AS failed
FROM {fq_table("memory_units")}
WHERE bank_id = $1
""",
bank_id,
)
if row is None:
return {"last_consolidated_at": None, "last_memory_write_at": None, "pending": 0, "failed": 0}
return {
"last_consolidated_at": row["last_consolidated_at"],
"last_memory_write_at": row["last_memory_write_at"],
"pending": row["pending"] or 0,
"failed": row["failed"] or 0,
}
async def link_counts(*, conn, fq_table: Callable[[str], str], bank_id: str) -> dict[str, int]:
"""``{link_type: count}`` of live links in a bank.
Non-entity links (temporal / semantic / caused_by) are a single ``GROUP BY`` over
``memory_links``. Entity links are no longer stored there they are derived on demand
from ``unit_entities``, replicating the historical writer cap of ``MAX_LINKS_PER_ENTITY``
bidirectional edges per shared entity so they are aggregated to one ``entity`` scalar.
"""
max_links_per_entity = 10
non_entity_link_rows = await conn.fetch(
f"""
SELECT link_type, COUNT(*) as count
FROM {fq_table("memory_links")}
WHERE bank_id = $1
GROUP BY link_type
""",
bank_id,
)
entity_total_row = await conn.fetchrow(
f"""
WITH per_entity AS (
SELECT ue.entity_id, COUNT(*) AS n
FROM {fq_table("unit_entities")} ue
JOIN {fq_table("memory_units")} mu ON mu.id = ue.unit_id
WHERE mu.bank_id = $1
GROUP BY ue.entity_id
)
SELECT COALESCE(SUM(LEAST(n - 1, $2)), 0)::bigint AS count
FROM per_entity
""",
bank_id,
max_links_per_entity,
)
entity_link_total = int(entity_total_row["count"] or 0) if entity_total_row else 0
counts: dict[str, int] = {row["link_type"]: row["count"] for row in non_entity_link_rows}
if entity_link_total > 0:
counts["entity"] = entity_link_total
return counts
async def document_memory_counts(
*, conn, fq_table: Callable[[str], str], bank_id: str, document_ids: list[str]
) -> dict[str, int]:
"""Live memory count per document id, for the ids given."""
if not document_ids:
return {}
rows = await conn.fetch(
f"""
SELECT document_id, COUNT(*) AS unit_count
FROM {fq_table("memory_units")}
WHERE bank_id = $1 AND document_id = ANY($2::text[])
GROUP BY document_id
""",
bank_id,
list(document_ids),
)
return {row["document_id"]: row["unit_count"] for row in rows}
async def memories_timeseries(
*, conn, fq_table: Callable[[str], str], bank_id: str, time_field: str, trunc: str, since: datetime
) -> list[dict[str, Any]]:
"""Memories bucketed by ``time_field`` (truncated to ``trunc``) and fact_type.
``time_field`` is whitelisted by the caller before it reaches here it is
interpolated into SQL. Event-time fields fall back to ``created_at`` per row so
rows without an event timestamp still appear.
"""
bucket_expr = time_field if time_field == "created_at" else f"COALESCE({time_field}, created_at)"
rows = await conn.fetch(
f"""
SELECT date_trunc('{trunc}', {bucket_expr} AT TIME ZONE 'UTC') AS bucket,
fact_type, COUNT(*) AS count
FROM {fq_table("memory_units")}
WHERE bank_id = $1 AND {bucket_expr} >= $2
GROUP BY bucket, fact_type
ORDER BY bucket
""",
bank_id,
since,
)
return [{"bucket": r["bucket"], "fact_type": r["fact_type"], "count": r["count"]} for r in rows]
async def observation_scope_counts(*, conn, fq_table: Callable[[str], str], bank_id: str) -> list[dict[str, Any]]:
"""Observations grouped by scope (their sorted tag set), most-populous first."""
rows = await conn.fetch(
f"""
SELECT scope, COUNT(*) AS count
FROM (
SELECT COALESCE(ARRAY(SELECT unnest(tags) ORDER BY 1), '{{}}'::text[]) AS scope
FROM {fq_table("memory_units")}
WHERE bank_id = $1 AND fact_type = 'observation'
) s
GROUP BY scope
ORDER BY count DESC, scope
""",
bank_id,
)
return [{"tags": list(r["scope"]), "count": r["count"]} for r in rows]
__all__ = [
"consolidation_freshness",
"document_memory_counts",
"link_counts",
"memories_timeseries",
"observation_scope_counts",
]
@@ -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"]
@@ -1,779 +0,0 @@
"""Graph-shaped reads and the link-maintenance passes, in SQL.
Everything here is a query over the *joins* around `memory_units` rather than
over the memories themselves: `unit_entities` (which entities a memory mentions)
and `memory_links` (memory-to-memory temporal/semantic/causal edges).
Two groups of callers:
* **The graph view.** :func:`graph_units`, :func:`graph_entity_rows` and
:func:`graph_direct_links` return raw rows; the engine still owns the
filtering, the observation inheritance, the derived entity edges, the
colouring and the response assembly. These functions answer only "which
memories", "which entity postings" and "which stored edges".
* **The graph-maintenance job.** :func:`enqueue_relink_victims` runs inside the
delete transaction; :func:`relink_pass`, :func:`prune_orphan_entities` and
:func:`prune_stale_cooccurrences` are the three reconciliation passes the job
drives. The job keeps the orchestration (pass ordering, the deadlock retry
around the sweeps, the timing log); each function here does the pass's work.
:func:`entity_memory_counts` and :func:`entities_for_units` are the two entity
postings reads that are not part of the graph view but read the same join table.
A store whose links travel inside the memory has nothing to relink and no join
table to sweep, which is why these are methods on the interface at all: it
answers them with zeroes rather than with SQL.
"""
from __future__ import annotations
import logging
import uuid as uuid_module
from collections.abc import Callable
from typing import Any
from ....config import get_config
from ...db.base import DatabaseConnection
from ...retain.link_utils import (
MAX_TEMPORAL_LINKS_PER_UNIT,
_bulk_insert_links,
_normalize_datetime,
compute_semantic_links_ann,
)
logger = logging.getLogger(__name__)
# Mirrors the ``top_k`` default in ``compute_semantic_links_ann`` at retain
# time. If you change one, change the other — otherwise victims would either
# never reach the cap (probe returns less than the cap) or stay perpetually
# under it (cap is higher than retain creates).
MAX_SEMANTIC_LINKS_PER_UNIT = 50
# Worker fetches this many rows per relink-loop iteration. Bounds
# per-iteration probe/insert latency so a 10k-row backlog doesn't hold a
# worker slot for minutes. Chosen so the typical iteration runs in well
# under 1s.
_DRAIN_BATCH_SIZE = 50
# Defensive guard against runaway relink loops — at _DRAIN_BATCH_SIZE units per
# iteration that's 500k targets, far beyond any realistic single-bank backlog.
_RELINK_ITERATION_CAP = 10000
# Cap at 10k edges — the UI can't usefully render more, and uncapped queries
# on highly-connected graphs (e.g. 1000 nodes with 500k+ edges) are too slow.
_GRAPH_MAX_EDGES = 10000
# Columns the graph view renders: nodes take id/text/date/context/entities,
# the table rows take the rest, and `source_memory_ids` is what lets the caller
# inherit an observation's links and entities from the facts behind it.
_GRAPH_UNIT_COLUMNS = (
"id, text, event_date, context, occurred_start, occurred_end, mentioned_at, "
"document_id, chunk_id, fact_type, tags, created_at, proof_count, source_memory_ids"
)
def _ops_for(conn: DatabaseConnection) -> Any:
"""The ``DataAccessOps`` matching the connection's SQL dialect.
This is the SQL memories store, and SQL means Postgres *or* Oracle the two
speak different dialects (Oracle inherits entity links through the
``observation_sources`` junction, Postgres through ``source_memory_ids``
arrays), so the ops must follow the connection rather than assume Postgres.
The ops go by ``conn.backend_type`` the connection objects carry the dialect
but not the backend's ``ops`` handle, so resolve through the per-dialect cache
of ``create_data_access_ops`` (a dict lookup after the first call, and the same
instance the backend holds). The default covers callers that hand in a bare
asyncpg connection with no dialect to report.
"""
from ...db import create_data_access_ops
return create_data_access_ops(getattr(conn, "backend_type", "postgresql"))
def _as_uuids(unit_ids: list) -> list:
"""Coerce a mixed list of uuid strings / UUIDs to UUIDs for a ``uuid[]`` bind."""
return [uuid_module.UUID(uid) if isinstance(uid, str) else uid for uid in unit_ids]
# ---------------------------------------------------------------- graph view
def _observations_via_source_match(
fq_table: Callable[[str], str],
ops: Any,
source_column: str,
source_placeholder: int,
bank_placeholder: int | None,
) -> str:
"""A predicate matching observations whose *sources* satisfy ``<col> = $n``.
Observations carry no `document_id` / `chunk_id` of their own; the link to a
source row lives in `source_memory_ids` (native array) or the
`observation_sources` junction, depending on the dialect.
"""
if ops.uses_observation_sources_table:
bank_clause = f" AND src.bank_id = ${bank_placeholder}" if bank_placeholder else ""
return (
f"id IN (SELECT os.observation_id "
f"FROM {fq_table('observation_sources')} os "
f"JOIN {fq_table('memory_units')} src ON src.id = os.source_id "
f"WHERE src.{source_column} = ${source_placeholder}{bank_clause})"
)
bank_clause = f" AND bank_id = ${bank_placeholder}" if bank_placeholder else ""
return (
f"source_memory_ids && (SELECT array_agg(id) "
f"FROM {fq_table('memory_units')} "
f"WHERE {source_column} = ${source_placeholder}{bank_clause})"
)
async def graph_units(
*,
conn: DatabaseConnection,
fq_table: Callable[[str], str],
bank_id: str | None = None,
fact_type: str | None = None,
search_query: str | None = None,
document_id: str | None = None,
chunk_id: str | None = None,
tags: list[str] | None = None,
tags_match: str = "all_strict",
limit: int = 1000,
) -> dict[str, Any]:
"""Memory nodes for the graph view, plus the total matching count.
Returns ``{"units": [...], "total": int}``: ``units`` is the page (newest
first, capped at ``limit``); ``total`` is how many match the filters, which
the UI shows alongside the page. ``document_id`` / ``chunk_id`` also match an
observation whose *sources* carry them, since observations have neither of
their own.
"""
from ...search.tags import build_tags_where_clause_simple
ops = _ops_for(conn)
conditions: list[str] = []
params: list[Any] = []
bank_placeholder: int | None = None
if bank_id:
params.append(bank_id)
bank_placeholder = len(params)
conditions.append(f"bank_id = ${bank_placeholder}")
if fact_type:
params.append(fact_type)
conditions.append(f"fact_type = ${len(params)}")
if document_id:
params.append(document_id)
obs = _observations_via_source_match(fq_table, ops, "document_id", len(params), bank_placeholder)
conditions.append(f"(document_id = ${len(params)} OR (fact_type = 'observation' AND {obs}))")
if chunk_id:
params.append(chunk_id)
obs = _observations_via_source_match(fq_table, ops, "chunk_id", len(params), bank_placeholder)
conditions.append(f"(chunk_id = ${len(params)} OR (fact_type = 'observation' AND {obs}))")
if search_query:
params.append(f"%{search_query}%")
conditions.append(f"(text ILIKE ${len(params)} OR context ILIKE ${len(params)})")
if tags:
tag_clause = build_tags_where_clause_simple(tags, len(params) + 1, match=tags_match)
if tag_clause:
conditions.append(tag_clause.removeprefix("AND "))
params.append(tags)
elif tags_match == "exact":
# Exact match with no tags is the "global" scope: rows carrying no tags at
# all. (Other modes treat empty tags as "no filter".)
conditions.append("(tags IS NULL OR tags = '{}')")
where_clause = "WHERE " + " AND ".join(conditions) if conditions else ""
total_row = await conn.fetchrow(
f"SELECT COUNT(*) AS total FROM {fq_table('memory_units')} {where_clause}",
*params,
)
total = total_row["total"] if total_row else 0
params.append(limit)
rows = await conn.fetch(
f"""
SELECT {_GRAPH_UNIT_COLUMNS}
FROM {fq_table("memory_units")}
{where_clause}
ORDER BY mentioned_at DESC NULLS LAST, event_date DESC
LIMIT ${len(params)}
""",
*params,
)
return {"units": [dict(row) for row in rows], "total": total}
async def graph_entity_rows(
*,
conn: DatabaseConnection,
fq_table: Callable[[str], str],
bank_id: str,
unit_ids: list[str],
) -> list[dict[str, Any]]:
"""``(unit_id, entity_id, canonical_name)`` rows for the graph view's entity edges.
Direct `unit_entities` postings only. An observation's entities are inherited
from its source memories by the caller, which is why the ids it passes here
are the visible units *plus* their source memories.
Scoped by unit id rather than by bank: the ids already came from a
bank-scoped :func:`graph_units`, and `unit_entities` carries no bank column.
"""
if not unit_ids:
return []
rows = await conn.fetch(
f"""
SELECT ue.unit_id, e.id AS entity_id, e.canonical_name
FROM {fq_table("unit_entities")} ue
JOIN {fq_table("entities")} e ON ue.entity_id = e.id
WHERE ue.unit_id = ANY($1::uuid[])
ORDER BY ue.unit_id
""",
_as_uuids(unit_ids),
)
return [dict(row) for row in rows]
async def graph_direct_links(
*,
conn: DatabaseConnection,
fq_table: Callable[[str], str],
bank_id: str,
unit_ids: list[str],
) -> list[dict[str, Any]]:
"""Memory-to-memory edges with *both* endpoints in ``unit_ids``.
Entity edges are derived by the caller from `unit_entities` so we don't
materialize them in `memory_links` anymore (dropped in migration
e9b2c7d1f3a4) no link_type filter is needed. ``entity_name`` is selected as
NULL so the row shape matches the derived edges the caller mixes these with.
Pass the visible units *and* the source memories they inherit from: the
caller copies a source memory's links onto the observations built on it.
"""
if not unit_ids:
return []
rows = await conn.fetch(
f"""
SELECT ml.from_unit_id,
ml.to_unit_id,
ml.link_type,
ml.weight,
NULL::text AS entity_name
FROM {fq_table("memory_links")} ml
WHERE ml.from_unit_id = ANY($1::uuid[])
AND ml.to_unit_id = ANY($1::uuid[])
ORDER BY ml.weight DESC NULLS LAST
LIMIT $2
""",
_as_uuids(unit_ids),
_GRAPH_MAX_EDGES,
)
return [dict(row) for row in rows]
# ------------------------------------------------------------ entity postings
async def entity_memory_counts(
*,
conn: DatabaseConnection,
fq_table: Callable[[str], str],
bank_id: str,
entity_ids: list[str] | None = None,
) -> dict[str, int]:
"""Live memory count per entity id, for the entities in ``bank_id``.
The GROUP BY is what makes this an orphan test: an entity with no surviving
`unit_entities` row produces no group, so it is simply absent from the
result rather than present with a zero.
Scoped through ``memory_units.bank_id`` `unit_entities` has no bank column,
and joining is what keeps the count to *live* memories (deleted units take
their postings with them via ON DELETE CASCADE).
"""
params: list[Any] = [bank_id]
entity_filter = ""
if entity_ids is not None:
if not entity_ids:
return {}
params.append(_as_uuids(entity_ids))
entity_filter = f"AND ue.entity_id = ANY(${len(params)}::uuid[])"
rows = await conn.fetch(
f"""
SELECT ue.entity_id, COUNT(*) AS n
FROM {fq_table("unit_entities")} ue
JOIN {fq_table("memory_units")} mu ON mu.id = ue.unit_id
WHERE mu.bank_id = $1
{entity_filter}
GROUP BY ue.entity_id
""",
*params,
)
return {str(row["entity_id"]): int(row["n"]) for row in rows}
def _entity_rows_for_units_sql(
fq_table: Callable[[str], str],
ops: Any,
unit_ids_placeholder: int,
) -> str:
"""SQL SELECT producing ``(unit_id, entity_id, canonical_name)`` rows for
the given unit IDs.
Direct rows come from ``unit_entities``. Observations rarely carry
direct rows there; their entity association lives transitively through
their source memories (``source_memory_ids`` on PG, the
``observation_sources`` junction on Oracle). When an observation has
no direct entity rows the SELECT inherits its source memories'
entities, so the result is the same set callers would get from
``get_memory_unit``.
``unit_ids_placeholder`` is the 1-based parameter index that holds the
``uuid[]`` of unit IDs. The placeholder is referenced twice both
sides of the UNION need it so callers should not reuse the slot.
"""
ue = fq_table("unit_entities")
ents = fq_table("entities")
mu = fq_table("memory_units")
p = unit_ids_placeholder
direct = (
f"SELECT ue.unit_id, e.id AS entity_id, e.canonical_name "
f"FROM {ue} ue "
f"JOIN {ents} e ON e.id = ue.entity_id "
f"WHERE ue.unit_id = ANY(${p}::uuid[])"
)
if ops.uses_observation_sources_table:
os_t = fq_table("observation_sources")
inherited = (
f"SELECT os.observation_id AS unit_id, e.id AS entity_id, e.canonical_name "
f"FROM {os_t} os "
f"JOIN {ue} src_ue ON src_ue.unit_id = os.source_id "
f"JOIN {ents} e ON e.id = src_ue.entity_id "
f"WHERE os.observation_id = ANY(${p}::uuid[]) "
f"AND NOT EXISTS (SELECT 1 FROM {ue} d WHERE d.unit_id = os.observation_id)"
)
else:
inherited = (
f"SELECT obs.id AS unit_id, e.id AS entity_id, e.canonical_name "
f"FROM {mu} obs "
f"CROSS JOIN LATERAL unnest(obs.source_memory_ids) AS src_id "
f"JOIN {ue} src_ue ON src_ue.unit_id = src_id "
f"JOIN {ents} e ON e.id = src_ue.entity_id "
f"WHERE obs.id = ANY(${p}::uuid[]) "
f"AND obs.fact_type = 'observation' "
f"AND obs.source_memory_ids IS NOT NULL "
f"AND NOT EXISTS (SELECT 1 FROM {ue} d WHERE d.unit_id = obs.id)"
)
return f"({direct}) UNION ({inherited})"
async def entities_for_units(
*,
conn: DatabaseConnection,
fq_table: Callable[[str], str],
bank_id: str,
unit_ids: list[str],
) -> dict[str, list[str]]:
"""The entity ids each unit carries, keyed by unit id.
Observations inherit their source memories' entities when they carry no
direct postings of their own see :func:`_entity_rows_for_units_sql`. Units
with no entities are absent rather than mapped to an empty list.
"""
if not unit_ids:
return {}
rows = await conn.fetch(
_entity_rows_for_units_sql(fq_table, _ops_for(conn), unit_ids_placeholder=1),
_as_uuids(unit_ids),
)
# UNION already de-duplicates whole rows, but a unit can reach the same
# entity through more than one source memory, so dedupe per unit while
# preserving the order the rows arrived in.
by_unit: dict[str, list[str]] = {}
for row in rows:
unit_key = str(row["unit_id"])
entity_id = str(row["entity_id"])
bucket = by_unit.setdefault(unit_key, [])
if entity_id not in bucket:
bucket.append(entity_id)
return by_unit
async def entity_map_for_units(
*,
conn: DatabaseConnection,
fq_table: Callable[[str], str],
bank_id: str,
unit_ids: list[str],
) -> dict[str, list[dict[str, str]]]:
"""``{unit_id: [{entity_id, canonical_name}]}`` — the recall/curation shape.
The named twin of :func:`entities_for_units`: recall renders the entity name
on each fact, so it needs the label, not just the id. Observation-via-source
inheritance and the per-unit dedupe are identical.
"""
if not unit_ids:
return {}
rows = await conn.fetch(
_entity_rows_for_units_sql(fq_table, _ops_for(conn), unit_ids_placeholder=1),
_as_uuids(unit_ids),
)
by_unit: dict[str, list[dict[str, str]]] = {}
for row in rows:
unit_key = str(row["unit_id"])
entity_id = str(row["entity_id"])
bucket = by_unit.setdefault(unit_key, [])
if not any(existing["entity_id"] == entity_id for existing in bucket):
bucket.append({"entity_id": entity_id, "canonical_name": row["canonical_name"]})
return by_unit
# --------------------------------------------------------------- maintenance
async def enqueue_relink_victims(
*,
conn: DatabaseConnection,
fq_table: Callable[[str], str],
bank_id: str,
affected_unit_ids: list,
include_affected_units: bool = False,
) -> int:
"""Enqueue surviving units whose outgoing temporal/semantic links pointed at
``affected_unit_ids`` for later link top-up.
Must run inside the same transaction that drops those links, *before* the
delete (or cascade) fires once the rows are gone, the join that finds the
victims returns nothing.
Args:
conn: Database connection inside the active transaction.
fq_table: Schema-qualifying table-name resolver.
bank_id: Bank owning the affected units.
affected_unit_ids: Memory_unit IDs whose incident temporal/semantic links
are about to be (or are being) removed.
include_affected_units: Also enqueue ``affected_unit_ids`` themselves for
an edit that deletes a unit's links but leaves the unit live, so its own
outgoing adjacency is rebuilt too. One combined insert keeps the queue's
sorted lock ordering intact.
Returns:
Number of distinct victim units enqueued (after dedup against rows
already in the queue).
"""
if not affected_unit_ids:
return 0
ops = _ops_for(conn)
affected_uuids = _as_uuids(affected_unit_ids)
affected_str_set = {str(uid) for uid in affected_uuids}
# Find units (other than the affected ones) that have an outgoing
# temporal/semantic link pointing at an affected unit. Only those two link
# types are relinked by graph maintenance; entity edges are not stored in
# memory_links (they're derived from unit_entities), so nothing else applies.
victim_rows = await conn.fetch(
f"""
SELECT DISTINCT from_unit_id
FROM {fq_table("memory_links")}
WHERE to_unit_id = ANY($1::uuid[])
AND bank_id = $2
AND link_type IN ('temporal', 'semantic')
""",
affected_uuids,
bank_id,
)
victim_ids = {row["from_unit_id"] for row in victim_rows if str(row["from_unit_id"]) not in affected_str_set}
if include_affected_units:
victim_ids.update(affected_uuids)
if not victim_ids:
return 0
await ops.enqueue_graph_maintenance(
conn,
fq_table("graph_maintenance_queue"),
bank_id,
list(victim_ids),
)
logger.debug(
f"[GRAPH_MAINT] Enqueued {len(victim_ids)} relink victims in "
f"bank={bank_id} ({len(affected_unit_ids)} units affected)"
)
return len(victim_ids)
async def relink_pass(
*,
backend: Any,
fq_table: Callable[[str], str],
bank_id: str,
config: Any,
) -> dict:
"""Drain ``graph_maintenance_queue`` for ``bank_id``, topping up lost links.
Per-iteration loop: claim top up commit. We rely on submit-time
dedup to keep at most one job per bank running, so no need for
SKIP LOCKED.
Takes ``backend`` rather than a connection because the loop spans several
transactions one per claimed batch, plus a separate connection for the ANN
probe so it has to acquire its own.
``config`` is the caller's resolved configuration. The Postgres pass takes
its caps from retain's link_utils (so relink and retain agree on what "full"
means) and never reads it; it is accepted so a store that *does* tune its
relinking gets it.
Returns:
``{"relink_units_processed": int, "relink_links_added": int}``.
"""
del config # accepted for symmetry with stores that tune their own relinking
ops = backend.ops
units_processed = 0
links_added = 0
iterations = 0
while True:
from ...memory_engine import acquire_with_retry
async with acquire_with_retry(backend) as conn:
async with conn.transaction():
unit_ids = await ops.claim_graph_maintenance_batch(
conn,
fq_table("graph_maintenance_queue"),
bank_id,
_DRAIN_BATCH_SIZE,
)
if not unit_ids:
break
links_added += await _relink_batch(conn, fq_table, bank_id, unit_ids, ops, backend)
units_processed += len(unit_ids)
iterations += 1
if iterations > _RELINK_ITERATION_CAP:
# Defensive guard against runaway loops — at 50 units/iter that's
# 500k targets, far beyond any realistic single-bank backlog.
logger.error(
f"[GRAPH_MAINT] bank={bank_id} hit iteration cap ({iterations}); aborting relink "
f"(units_processed={units_processed}, links_added={links_added})"
)
break
return {"relink_units_processed": units_processed, "relink_links_added": links_added}
async def _relink_batch(
conn: DatabaseConnection,
fq_table: Callable[[str], str],
bank_id: str,
victim_ids: list[str],
ops: Any,
backend: Any,
) -> int:
"""Top up temporal/semantic links for a batch of victim units. Returns rows inserted."""
# Load each victim's metadata. Victims whose units were deleted between
# enqueue and now silently drop out — exactly the no-op behaviour we want
# for stale queue rows.
victim_uuids = [uuid_module.UUID(vid) for vid in victim_ids]
victim_rows = await conn.fetch(
f"""
SELECT id::text AS id, event_date, fact_type, embedding::text AS embedding
FROM {fq_table("memory_units")}
WHERE id = ANY($1::uuid[])
AND bank_id = $2
AND fact_type IN ('experience', 'world')
""",
victim_uuids,
bank_id,
)
if not victim_rows:
return 0
alive_uuids = [uuid_module.UUID(row["id"]) for row in victim_rows]
# Count current outgoing temporal/semantic links per victim so we only
# probe for the ones genuinely below cap. Saves the bulk of the work when
# most victims still have plenty of links.
count_rows = await conn.fetch(
f"""
SELECT from_unit_id, link_type, COUNT(*) AS cnt
FROM {fq_table("memory_links")}
WHERE from_unit_id = ANY($1::uuid[])
AND bank_id = $2
AND link_type IN ('temporal', 'semantic')
GROUP BY from_unit_id, link_type
""",
alive_uuids,
bank_id,
)
counts: dict[tuple[str, str], int] = {}
for row in count_rows:
counts[(str(row["from_unit_id"]), row["link_type"])] = int(row["cnt"])
# --- Temporal top-up ---
temporal_needs = [r for r in victim_rows if counts.get((r["id"], "temporal"), 0) < MAX_TEMPORAL_LINKS_PER_UNIT]
new_links: list[tuple] = []
if temporal_needs:
lateral_unit_ids = [uuid_module.UUID(r["id"]) for r in temporal_needs if r["event_date"] is not None]
lateral_event_dates = [
_normalize_datetime(r["event_date"]) for r in temporal_needs if r["event_date"] is not None
]
lateral_fact_types = [r["fact_type"] for r in temporal_needs if r["event_date"] is not None]
if lateral_unit_ids:
rows = await ops.fetch_temporal_neighbors(
conn,
fq_table("memory_units"),
bank_id,
lateral_unit_ids,
lateral_event_dates,
lateral_fact_types,
MAX_TEMPORAL_LINKS_PER_UNIT,
)
for row in rows:
time_diff_h = float(row["time_diff_hours"])
# Mirror the 24h window enforced at retain time. The bidirectional
# index scan returns the K closest neighbours regardless of
# window, so we filter here.
if time_diff_h > 24:
continue
weight = max(0.3, 1.0 - (time_diff_h / 24))
new_links.append((row["from_id"], str(row["id"]), "temporal", weight))
# --- Semantic top-up ---
# ANN must run on its own connection: it opens a nested transaction with
# SET LOCAL hnsw.ef_search + CREATE TEMP TABLE ON COMMIT DROP, and nesting
# that inside our current write transaction would commit our writes early.
semantic_needs = [
r
for r in victim_rows
if counts.get((r["id"], "semantic"), 0) < MAX_SEMANTIC_LINKS_PER_UNIT and r["embedding"] is not None
]
if semantic_needs:
from ...memory_engine import acquire_with_retry
seed_ids = [r["id"] for r in semantic_needs]
seed_embs = [r["embedding"] for r in semantic_needs]
seed_ftypes = [r["fact_type"] for r in semantic_needs]
async with acquire_with_retry(backend) as ann_conn:
try:
ann_links = await compute_semantic_links_ann(
ann_conn,
bank_id,
seed_ids,
seed_embs,
fact_types=seed_ftypes,
threshold=get_config().semantic_link_min_similarity,
)
# Strip self-links (rare but possible because the ANN probe
# has no exclude list — see the comment in compute_semantic_links_ann).
ann_links = [lnk for lnk in ann_links if lnk[0] != lnk[1]]
new_links.extend(ann_links)
except Exception as e:
# ANN uses PG-specific HNSW syntax; on dialects/configs where
# it isn't available we still want the temporal top-up to land.
logger.warning(f"[GRAPH_MAINT] Semantic top-up failed for bank={bank_id}: {type(e).__name__}: {e}")
if not new_links:
return 0
await _bulk_insert_links(
conn,
new_links,
bank_id=bank_id,
skip_exists_check=False,
ops=ops,
)
return len(new_links)
async def prune_orphan_entities(
*,
conn: DatabaseConnection,
fq_table: Callable[[str], str],
bank_id: str,
) -> int:
"""Delete ``entities`` rows in the bank with no remaining ``unit_entities``
references. Returns the number pruned.
FK ON DELETE CASCADE on ``entity_cooccurrences`` then removes any
cooccurrence row pointing at the pruned entities which is why this runs
before :func:`prune_stale_cooccurrences` rather than after.
A bank-wide single-statement delete, cheap when there's nothing to do. It is
idempotent (rerunning only deletes what is still orphaned), so the caller is
free to retry the whole transaction on deadlock.
"""
ops = _ops_for(conn)
return await ops.prune_orphan_entities(
conn,
fq_table("entities"),
fq_table("unit_entities"),
bank_id,
)
async def prune_stale_cooccurrences(
*,
conn: DatabaseConnection,
fq_table: Callable[[str], str],
bank_id: str,
) -> int:
"""Delete cooccurrence rows no current memory witnesses. Returns the count.
Defensive sweep for rows where both endpoints still exist but no current
memory_unit references both of them the cooccurrence was real at the time
it was recorded, but every unit that witnessed it has since been deleted.
:func:`prune_orphan_entities` cascades the *missing-entity* case via FK; this
pass catches the *stale-count* case it cannot see.
Like the orphan prune, a bank-wide idempotent sweep backed by indexes, so
it's cheap when there's nothing to do and safe for the caller to retry.
"""
ops = _ops_for(conn)
return await ops.prune_stale_cooccurrences(
conn,
fq_table("entity_cooccurrences"),
fq_table("unit_entities"),
fq_table("entities"),
bank_id,
)
__all__ = [
"MAX_SEMANTIC_LINKS_PER_UNIT",
"enqueue_relink_victims",
"entities_for_units",
"entity_map_for_units",
"entity_memory_counts",
"graph_direct_links",
"graph_entity_rows",
"graph_units",
"prune_orphan_entities",
"prune_stale_cooccurrences",
"relink_pass",
]
@@ -1,539 +0,0 @@
"""Addressed reads over `memory_units`: get, scan, count, tags, consolidation state.
Not retrieval nothing here ranks. These are the queries behind the curation
detail view, export, the bank-stats panel and the consolidation queue, lifted out
of the call sites that used to issue them inline (``memory_engine``,
``transfer/export``, ``consolidation/consolidator``) so
:class:`~hindsight_api.engine.memories.postgres.PostgresMemories` can delegate
rather than embed SQL.
Every function takes the live connection and Hindsight's ``fq_table`` resolver, so
each one runs inside whatever transaction the caller already holds; none of them
acquires a connection of its own.
**Cursor semantics.** ``scan_memories``'s ``page_token`` is opaque to callers, and
for Postgres it is simply a *numeric offset rendered as a decimal string* against
the scan's fixed ``ORDER BY created_at, id``. An empty token means "start at the
beginning", and an empty token comes back once the walk is exhausted (i.e. the
final short page). An offset cursor is a position rather than a snapshot exactly
the guarantee :class:`~hindsight_api.engine.memories.base.ScanPage` documents:
rows written or deleted mid-walk can shift later pages, so a scan is
eventually-complete browsing rather than a consistent iterator. ``skip`` is applied
*on top of* the decoded cursor, so a caller that pages with both should pass
``skip`` only on the first call the returned token already accounts for it.
"""
from __future__ import annotations
import json
import uuid
from collections.abc import Callable
from datetime import datetime
from typing import Any
from ...search.tags import (
build_tag_groups_where_clause,
build_tags_where_clause,
build_tags_where_clause_simple,
)
from ..base import ScanPage, StoredMemory
# The `memory_units` projection every read here shares. Superset of the by-id
# SELECT the recall source-facts path used (text/fact_type/context/timestamps/
# document_id/chunk_id/tags/metadata), plus the observation bookkeeping columns
# `StoredMemory` carries: source_memory_ids and consolidated_at.
_MEMORY_COLUMNS = """
id, text, fact_type, context, document_id, chunk_id, tags, metadata,
proof_count, event_date, occurred_start, occurred_end, mentioned_at,
created_at, source_memory_ids, consolidated_at, observation_scopes
"""
# The scan's order. Fixed (created_at, id) like the export loader's, because an
# offset cursor is only meaningful against a total order.
_SCAN_ORDER = "ORDER BY created_at, id"
def _as_json(value: Any) -> Any:
"""Coerce an asyncpg JSONB column (str or already-decoded) to a Python object.
Connections differ in whether a JSONB codec is registered, so the column
arrives either as text or as the decoded object.
"""
if value is None:
return None
if isinstance(value, str):
try:
return json.loads(value)
except json.JSONDecodeError:
# A valid scalar such as `"combined"` arrives already decoded on
# connections that do register a decoder.
return value
return value
def _as_uuids(unit_ids: list[Any]) -> list[uuid.UUID]:
"""Unit ids as UUIDs, dropping anything unparseable.
A malformed id is treated the same way a deleted one is simply absent from
the result rather than failing the whole read.
"""
out: list[uuid.UUID] = []
for unit_id in unit_ids or []:
if isinstance(unit_id, uuid.UUID):
out.append(unit_id)
continue
try:
out.append(uuid.UUID(str(unit_id)))
except (ValueError, AttributeError, TypeError):
continue
return out
def _column(row: Any, name: str, default: Any = None) -> Any:
"""One column of an asyncpg Record, tolerating a narrower projection."""
try:
return row[name]
except (KeyError, IndexError):
return default
def _stored_from_row(row: Any) -> StoredMemory:
"""Map a `memory_units` row onto :class:`StoredMemory`.
Shared by every read in this module so the row dataclass mapping exists
once. ``entity_ids`` stays empty: the unitentity posting lives in
`unit_entities` and is served by ``entities_for_units``, not by a join here.
"""
source_ids = _column(row, "source_memory_ids") or []
return StoredMemory(
unit_id=str(row["id"]),
text=row["text"],
fact_type=row["fact_type"],
context=_column(row, "context"),
document_id=_column(row, "document_id"),
chunk_id=str(_column(row, "chunk_id")) if _column(row, "chunk_id") else None,
tags=list(_column(row, "tags") or []),
metadata=_as_json(_column(row, "metadata")),
proof_count=_column(row, "proof_count") or 1,
event_date=_column(row, "event_date"),
occurred_start=_column(row, "occurred_start"),
occurred_end=_column(row, "occurred_end"),
mentioned_at=_column(row, "mentioned_at"),
created_at=_column(row, "created_at"),
source_memory_ids=[str(sid) for sid in source_ids],
consolidated_at=_column(row, "consolidated_at"),
# Consolidation routes a candidate by its scopes, so this has to survive
# the trip through the store rather than being re-queried per memory.
observation_scopes=_as_json(_column(row, "observation_scopes")),
)
def _decode_page_token(page_token: str) -> int:
"""Decode the offset cursor. Empty, malformed or negative all mean "start"."""
if not page_token:
return 0
try:
offset = int(page_token)
except (TypeError, ValueError):
return 0
return offset if offset > 0 else 0
async def get_memories(
*, conn, fq_table: Callable[[str], str], bank_id: str, unit_ids: list[str]
) -> list[StoredMemory]:
"""Fetch memories by id. Missing or deleted ids are simply absent."""
ids = _as_uuids(unit_ids)
if not ids:
return []
rows = await conn.fetch(
f"""
SELECT {_MEMORY_COLUMNS}
FROM {fq_table("memory_units")}
WHERE bank_id = $1 AND id = ANY($2::uuid[])
""",
bank_id,
ids,
)
return [_stored_from_row(row) for row in rows]
async def _semantic_edges(
*, conn, fq_table: Callable[[str], str], bank_id: str, unit_ids: list[uuid.UUID]
) -> dict[str, list[tuple[str, float]]]:
"""Derived kNN edges for ``unit_ids``, keyed by unit id.
Walked in both directions, like the graph arm's semantic expansion: a
`memory_links` row is written once, so a unit's neighbourhood is the union of
the edges leaving it and those arriving at it.
"""
if not unit_ids:
return {}
rows = await conn.fetch(
f"""
SELECT from_unit_id AS unit_id, to_unit_id AS target_id, weight
FROM {fq_table("memory_links")}
WHERE bank_id = $1 AND link_type = 'semantic' AND from_unit_id = ANY($2::uuid[])
UNION ALL
SELECT to_unit_id AS unit_id, from_unit_id AS target_id, weight
FROM {fq_table("memory_links")}
WHERE bank_id = $1 AND link_type = 'semantic' AND to_unit_id = ANY($2::uuid[])
""",
bank_id,
unit_ids,
)
edges: dict[str, list[tuple[str, float]]] = {}
for row in rows:
edges.setdefault(str(row["unit_id"]), []).append((str(row["target_id"]), float(row["weight"] or 0.0)))
return edges
async def scan_memories(
*,
conn,
fq_table: Callable[[str], str],
bank_id: str,
fact_types: list[str] | None = None,
limit: int = 100,
page_token: str = "",
tags: list[str] | None = None,
tags_match: str = "any",
tag_groups: list | None = None,
document_id: str | None = None,
metadata_equals: dict[str, str] | None = None,
skip: int = 0,
include_edges: bool = False,
) -> ScanPage:
"""Page through stored memories. A full walk — for browsing and export only.
See the module docstring for the ``page_token`` (offset) cursor semantics.
"""
if limit is None or limit <= 0:
return ScanPage()
where: list[str] = ["bank_id = $1"]
params: list[Any] = [bank_id]
if fact_types:
params.append(list(fact_types))
where.append(f"fact_type = ANY(${len(params)})")
if document_id is not None:
# A real column here, which is why it is not folded into
# `metadata_equals`: only a store without the column keeps it in the bag.
params.append(document_id)
where.append(f"document_id = ${len(params)}")
if metadata_equals:
# str→str equality across every key, which is exactly JSONB containment.
params.append(json.dumps(metadata_equals))
where.append(f"metadata @> ${len(params)}::jsonb")
# The tags clause owns its own `AND` prefix and, per the helper's contract,
# only consumes a bind param when `tags` is non-empty (match="exact" with no
# tags is the untagged/global scope and needs none).
tags_clause = build_tags_where_clause_simple(tags, len(params) + 1, match=tags_match)
if tags:
params.append(list(tags))
# Compound tag groups (AND/OR/NOT trees), AND-ed on. Also owns its `AND` prefix and appends
# one bind param per leaf; empty/absent groups yield no clause and no params.
groups_clause, group_params, _ = build_tag_groups_where_clause(tag_groups, param_offset=len(params) + 1)
params.extend(group_params)
offset = _decode_page_token(page_token) + max(int(skip or 0), 0)
params.append(limit)
limit_idx = len(params)
params.append(offset)
offset_idx = len(params)
rows = await conn.fetch(
f"""
SELECT {_MEMORY_COLUMNS}
FROM {fq_table("memory_units")}
WHERE {" AND ".join(where)} {tags_clause} {groups_clause}
{_SCAN_ORDER}
LIMIT ${limit_idx} OFFSET ${offset_idx}
""",
*params,
)
memories = [_stored_from_row(row) for row in rows]
if include_edges and memories:
edges = await _semantic_edges(
conn=conn, fq_table=fq_table, bank_id=bank_id, unit_ids=_as_uuids([m.unit_id for m in memories])
)
for memory in memories:
memory.semantic_edges = edges.get(memory.unit_id, [])
# A short page means the walk is exhausted, so the cursor goes empty.
next_token = str(offset + len(rows)) if len(rows) == limit else ""
return ScanPage(memories=memories, next_page_token=next_token)
async def count_memories(*, conn, fq_table: Callable[[str], str], bank_id: str) -> dict[str, int]:
"""Live memory count per fact_type. The bank-stats node counts."""
rows = await conn.fetch(
f"""
SELECT fact_type, COUNT(*) as count
FROM {fq_table("memory_units")}
WHERE bank_id = $1
GROUP BY fact_type
""",
bank_id,
)
return {row["fact_type"]: int(row["count"]) for row in rows}
async def list_tags(
*,
conn,
fq_table: Callable[[str], str],
bank_id: str,
pattern: str | None = None,
limit: int = 100,
offset: int = 0,
) -> dict[str, Any]:
"""One page of a bank's tag histogram: ``{"items": [{tag, count}], "total", "limit", "offset"}``.
``memory_units`` lives in SQL for this store, so the wildcard filter, the
``count DESC, tag ASC`` ordering and the paging all run in SQL the whole
histogram never crosses the wire. The dialect fragments come from
``build_tag_listing_parts`` (``unnest`` on Postgres, ``JSON_TABLE`` on Oracle):
this module backs both dialects, so it must not inline either one's SQL.
"""
from ...db import create_data_access_ops
ops = create_data_access_ops(getattr(conn, "backend_type", "postgresql"))
tag_parts = ops.build_tag_listing_parts(fq_table("memory_units"))
tag_source = tag_parts.tag_source
non_empty_check = tag_parts.non_empty_check
tag_col = tag_parts.tag_col
bank_prefix = tag_parts.bank_prefix
params: list[Any] = [bank_id]
pattern_clause = ""
if pattern:
# '*' is the wildcard, matched case-insensitively — same anchored ILIKE semantics as before.
params.append(pattern.replace("*", "%"))
pattern_clause = f"AND {tag_col} ILIKE $2"
total_row = await conn.fetchrow(
f"""
SELECT COUNT(DISTINCT {tag_col}) as total
FROM {tag_source}
WHERE {bank_prefix}bank_id = $1 {non_empty_check}
{pattern_clause}
""",
*params,
)
total = int(total_row["total"]) if total_row else 0
limit_param = len(params) + 1
offset_param = len(params) + 2
params.extend([limit, offset])
rows = await conn.fetch(
f"""
SELECT {tag_col} as tag, COUNT(*) as count
FROM {tag_source}
WHERE {bank_prefix}bank_id = $1 {non_empty_check}
{pattern_clause}
GROUP BY {tag_col}
ORDER BY count DESC, {tag_col} ASC
LIMIT ${limit_param} OFFSET ${offset_param}
""",
*params,
)
return {
"items": [{"tag": row["tag"], "count": int(row["count"])} for row in rows],
"total": total,
"limit": limit,
"offset": offset,
}
async def find_unconsolidated(
*,
conn,
fq_table: Callable[[str], str],
bank_id: str,
fact_types: list[str],
limit: int,
scope_tags: list[str] | None = None,
) -> list[StoredMemory]:
"""Memories not yet folded into an observation, oldest first.
The consolidator's candidate query: never consolidated, never *failed* to
consolidate (a memory the LLM could not handle must not be retried forever),
ordered by ``created_at`` so the queue drains in arrival order. ``scope_tags``
is the same ``tags @> scope`` containment the job's scope filter uses — the
job ORs several scopes together; one scope is passed here.
"""
where = [
"bank_id = $1",
"consolidated_at IS NULL",
"consolidation_failed_at IS NULL",
]
params: list[Any] = [bank_id]
if fact_types:
params.append(list(fact_types))
where.append(f"fact_type = ANY(${len(params)})")
if scope_tags:
params.append(list(scope_tags))
where.append(f"tags @> ${len(params)}::varchar[]")
params.append(limit)
rows = await conn.fetch(
f"""
SELECT {_MEMORY_COLUMNS}
FROM {fq_table("memory_units")}
WHERE {" AND ".join(where)}
ORDER BY created_at ASC
LIMIT ${len(params)}
""",
*params,
)
return [_stored_from_row(row) for row in rows]
async def count_unconsolidated(
*,
conn,
fq_table: Callable[[str], str],
bank_id: str,
fact_types: list[str],
scopes: list[list[str] | None],
limit: int,
) -> int:
"""Bounded ``COUNT(*)`` of unconsolidated candidates matching any scope — the cheap counterpart
to :func:`find_unconsolidated` that never ships a row.
Same predicates as ``find_unconsolidated`` (never consolidated, never failed, matching
fact_type), with the scopes OR'd as ``tags @> scope`` containment. ``id`` is the PK so each row
counts once; the inner ``LIMIT`` floors the count at ``limit`` exactly as walking that many rows
would, so a huge backlog stays a single index count instead of a 17-column fetch.
"""
where = ["bank_id = $1", "consolidated_at IS NULL", "consolidation_failed_at IS NULL"]
params: list[Any] = [bank_id]
if fact_types:
params.append(list(fact_types))
where.append(f"fact_type = ANY(${len(params)})")
# An unscoped entry (None) matches every row, collapsing the OR to no tag filter at all.
scope_clauses: list[str] = []
unscoped = any(scope is None for scope in scopes)
if not unscoped:
for scope in scopes:
params.append(list(scope or []))
scope_clauses.append(f"tags @> ${len(params)}::varchar[]")
if scope_clauses:
where.append("(" + " OR ".join(scope_clauses) + ")")
params.append(limit)
row = await conn.fetchrow(
f"""
SELECT COUNT(*) AS c FROM (
SELECT 1 FROM {fq_table("memory_units")}
WHERE {" AND ".join(where)}
LIMIT ${len(params)}
) sub
""",
*params,
)
return int(row["c"]) if row else 0
async def mark_consolidated(
*,
conn,
fq_table: Callable[[str], str],
bank_id: str,
unit_ids: list[str],
when: datetime | None,
failed: bool = False,
) -> None:
"""Stamp (or clear, with ``when=None``) the consolidated marker on sources.
``failed`` writes ``consolidation_failed_at`` instead of ``consolidated_at``,
which is what keeps a memory the LLM could not consolidate out of the queue.
``when=None`` clears the column rather than stamping it that is how a source
is requeued once the observation built on it is deleted. The clear keeps the
``fact_type IN ('experience', 'world')`` guard the requeue sites carry:
observations are never themselves consolidated, so nothing about them should
be reset by a requeue.
``updated_at`` is deliberately left alone, matching the consolidator's own
statements: consolidation bookkeeping is not an edit to the memory, and
bumping it would make every consolidation pass look like a write to the
staleness check below.
"""
ids = _as_uuids(unit_ids)
if not ids:
return
column = "consolidation_failed_at" if failed else "consolidated_at"
guard = "" if when is not None else " AND fact_type IN ('experience', 'world')"
await conn.execute(
f"""
UPDATE {fq_table("memory_units")}
SET {column} = $1
WHERE bank_id = $2 AND id = ANY($3::uuid[]){guard}
""",
when,
bank_id,
ids,
)
async def any_memory_updated_since(
*,
conn,
fq_table: Callable[[str], str],
bank_id: str,
since: datetime,
fact_types: list[str] | None = None,
tags: list[str] | None = None,
tags_match: str = "any",
tag_groups: list | None = None,
) -> bool:
"""Whether any memory in ``bank_id``'s scope was written after ``since``.
Backs the mental-model staleness check, so it is a bounded existence test
``LIMIT 1``, never a COUNT: the answer is "is there one", and the planner can
stop at the first hit. The scope is the mental model's: its flat tags (or the
compound ``tag_groups``) plus an optional ``fact_types`` restriction. This is
where the staleness query's WHERE lives, so the same scope that gates a
refresh decides whether one is due.
"""
params: list[Any] = [bank_id, since]
where = ["bank_id = $1", "updated_at > $2"]
tag_clause, tag_params, next_param = build_tags_where_clause(tags, param_offset=len(params) + 1, match=tags_match)
if tag_clause:
where.append(tag_clause.removeprefix("AND "))
params.extend(tag_params)
group_clause, group_params, _ = build_tag_groups_where_clause(tag_groups, param_offset=next_param)
if group_clause:
where.append(group_clause.removeprefix("AND "))
params.extend(group_params)
# Untagged, no tag_groups → no tag constraint, matching any memory in the bank.
if fact_types:
params.append(list(fact_types))
where.append(f"fact_type = ANY(${len(params)}::text[])")
row = await conn.fetchval(
f"SELECT 1 FROM {fq_table('memory_units')} WHERE {' AND '.join(where)} LIMIT 1",
*params,
)
return row is not None
__all__ = [
"any_memory_updated_since",
"count_memories",
"find_unconsolidated",
"get_memories",
"list_tags",
"mark_consolidated",
"scan_memories",
]
@@ -1,584 +0,0 @@
"""Writes against `memory_units`: the fact insert, the deletes, and observation invalidation.
Everything here mutates the memories slice and nothing else. The document row,
the chunks, the entity registry and the link tables stay with their own callers
what lands in this module is only the statements that touch `memory_units` (and,
on backends that keep one, the `observation_sources` junction that hangs off it).
Each function takes the live connection and Hindsight's ``fq_table`` resolver, so
it runs inside whatever transaction the caller already holds; ``ops`` is the
dialect ops object, which is what lets the same code serve the PG (native array)
and Oracle (junction table) shapes of the observationsource relation.
"""
from __future__ import annotations
import json
import logging
import uuid
from collections.abc import Callable
from typing import TYPE_CHECKING, Any
from ....config import get_config
from ..base import StoredMemory
if TYPE_CHECKING: # pragma: no cover - typing only
from ...retain.types import ProcessedFact
logger = logging.getLogger(__name__)
async def insert_facts(
*,
conn,
ops,
bank_id: str,
facts: list[ProcessedFact],
document_id: str | None = None,
) -> list[str]:
"""Insert facts into the database in batch.
Args:
conn: Database connection
bank_id: Bank identifier
facts: List of ProcessedFact objects to insert
document_id: Optional document ID to associate with facts
Returns:
List of unit IDs (UUIDs as strings) for the inserted facts, in the same
order as ``facts``.
"""
if not facts:
return []
# Imported here: `retain` reaches back into the engine for `fq_table`, so a
# module-level import would close the cycle once the engine imports this store.
from ...retain.fact_extraction import _sanitize_text
# Prepare data for batch insert
fact_texts = []
embeddings = []
event_dates = []
occurred_starts = []
occurred_ends = []
mentioned_ats = []
contexts = []
fact_types = []
metadata_jsons = []
chunk_ids = []
document_ids = []
tags_list = []
observation_scopes_list = []
text_signals_list = []
for fact in facts:
fact_texts.append(_sanitize_text(fact.fact_text))
# Convert embedding to string for asyncpg vector type
embeddings.append(str(fact.embedding))
# event_date: Use occurred_start if available, otherwise use mentioned_at
# This maintains backward compatibility while handling None occurred_start
event_dates.append(fact.occurred_start if fact.occurred_start is not None else fact.mentioned_at)
occurred_starts.append(fact.occurred_start)
occurred_ends.append(fact.occurred_end)
mentioned_ats.append(fact.mentioned_at)
contexts.append(_sanitize_text(fact.context))
fact_types.append(fact.fact_type)
metadata_jsons.append(json.dumps(fact.metadata))
chunk_ids.append(fact.chunk_id)
# Use per-fact document_id if available, otherwise fallback to batch-level document_id
document_ids.append(fact.document_id if fact.document_id else document_id)
# Convert tags to JSON string for proper batch insertion (PostgreSQL unnest doesn't handle 2D arrays well)
tags_list.append(json.dumps(fact.tags if fact.tags else []))
# observation_scopes: stored as JSONB (string or 2D array), None if not provided
observation_scopes_list.append(
json.dumps(fact.observation_scopes) if fact.observation_scopes is not None else None
)
# Build text_signals: entity names + date tokens for enriched BM25 indexing
signal_parts = []
if fact.entities:
signal_parts.extend(e.name for e in fact.entities)
if fact.occurred_start:
try:
signal_parts.append(fact.occurred_start.strftime("%B %d %Y").lstrip("0").replace(" 0", " "))
except (ValueError, AttributeError):
pass
if fact.occurred_end and fact.occurred_end != fact.occurred_start:
try:
signal_parts.append(fact.occurred_end.strftime("%B %d %Y").lstrip("0").replace(" 0", " "))
except (ValueError, AttributeError):
pass
text_signals_list.append(" ".join(signal_parts) if signal_parts else None)
# Batch insert all facts — delegates to DataAccessOps which handles
# unnest (PG) vs row-by-row (Oracle) transparently.
config = get_config()
return await ops.insert_facts_batch(
conn,
bank_id,
fact_texts,
embeddings,
event_dates,
occurred_starts,
occurred_ends,
mentioned_ats,
contexts,
fact_types,
metadata_jsons,
chunk_ids,
document_ids,
tags_list,
observation_scopes_list,
text_signals_list,
text_search_extension=config.text_search_extension,
)
async def delete_document(*, conn, fq_table: Callable[[str], str], bank_id: str, document_id: str) -> None:
"""Delete every memory unit belonging to ``document_id``.
Explicitly delete memory_units by document_id BEFORE deleting the
document row. The CASCADE from documentschunksmemory_units only
catches units that have a non-NULL chunk_id FK. Units with chunk_id=NULL
(e.g. from partial writes or edge cases) would survive the cascade.
This explicit delete ensures complete cleanup.
Called when a document is replaced, so it races the replacement's writes: it
must remove only what was written *before* this call, never the facts
arriving moments later which the ``document_id``/``bank_id`` predicate
gives for free inside the caller's transaction.
"""
await conn.execute(
f"DELETE FROM {fq_table('memory_units')} WHERE document_id = $1 AND bank_id = $2",
document_id,
bank_id,
)
async def delete_observations(*, conn, fq_table: Callable[[str], str], bank_id: str) -> None:
"""Delete all observations in a bank, leaving the facts behind them.
Only the observation rows: requeuing the surviving sources (clearing
``consolidated_at``) and resetting the bank's consolidation timestamp belong
to the caller, which owns the bank row.
"""
await conn.execute(
f"DELETE FROM {fq_table('memory_units')} WHERE bank_id = $1 AND fact_type = 'observation'",
bank_id,
)
async def observations_for_sources(
*,
conn,
ops,
fq_table: Callable[[str], str],
bank_id: str,
unit_ids: list[str | uuid.UUID],
) -> list[StoredMemory]:
"""Observations consolidated from any of ``unit_ids``.
Only ``unit_id`` and ``source_memory_ids`` are populated the caller uses
them to delete the observations and to work out which sources survive, and
the rest of the row is about to be deleted anyway.
"""
if not unit_ids:
return []
fact_uuids = [uuid.UUID(str(fid)) if not isinstance(fid, uuid.UUID) else fid for fid in unit_ids]
if ops is not None and not ops.uses_observation_sources_table:
# PG: use native array overlap operator
rows = await conn.fetch(
f"""
SELECT id, source_memory_ids
FROM {fq_table("memory_units")}
WHERE bank_id = $1
AND fact_type = 'observation'
AND source_memory_ids && $2::uuid[]
""",
bank_id,
fact_uuids,
)
else:
# Oracle / default: use observation_sources junction table
rows = await conn.fetch(
f"""
SELECT mu.id, mu.source_memory_ids
FROM {fq_table("memory_units")} mu
WHERE mu.bank_id = $1
AND mu.fact_type = 'observation'
AND EXISTS (
SELECT 1 FROM {fq_table("observation_sources")} os
WHERE os.observation_id = mu.id
AND os.source_id = ANY($2::uuid[])
)
""",
bank_id,
fact_uuids,
)
return [
StoredMemory(
unit_id=str(row["id"]),
text="",
fact_type="observation",
source_memory_ids=[str(src_id) for src_id in (row["source_memory_ids"] or [])],
)
for row in rows
]
async def delete_stale_observations(
*,
conn,
ops,
fq_table: Callable[[str], str],
bank_id: str,
fact_ids: list[str | uuid.UUID],
) -> int:
"""Delete observations whose source memories are about to be removed.
Mirrors the cleanup performed by ``MemoryEngine.delete_document`` so that
every code path that removes ``memory_units`` also removes the
observations derived from them. Without this, ingesting a fresh version
of a document via the retain pipeline (which does a full-replace
``DELETE FROM documents`` cascade) used to leave orphan observations
pointing at memory IDs that no longer existed.
For each observation referencing any of ``fact_ids``:
1. Delete the observation row (its text is stale once even one source
memory disappears).
2. Reset ``consolidated_at = NULL`` on the surviving source memories so
they get re-consolidated under fresh observations on the next run.
Must be called within an active transaction, before the source memories
are deleted.
Returns the number of observations deleted.
"""
if not fact_ids:
return 0
fact_uuids = [uuid.UUID(str(fid)) if not isinstance(fid, uuid.UUID) else fid for fid in fact_ids]
affected_obs = await observations_for_sources(
conn=conn, ops=ops, fq_table=fq_table, bank_id=bank_id, unit_ids=fact_uuids
)
if not affected_obs:
return 0
deleted_set = {str(uid) for uid in fact_uuids}
obs_ids = [uuid.UUID(obs.unit_id) for obs in affected_obs]
seen_remaining: set[str] = set()
remaining_source_ids: list[uuid.UUID] = []
for obs in affected_obs:
for src_str in obs.source_memory_ids:
if src_str not in deleted_set and src_str not in seen_remaining:
remaining_source_ids.append(uuid.UUID(src_str))
seen_remaining.add(src_str)
await conn.execute(
f"DELETE FROM {fq_table('memory_units')} WHERE id = ANY($1::uuid[])",
obs_ids,
)
# Their history is keyed by observation_id and no longer cascades from memory_units (that FK
# was dropped so history can be recorded for observations kept outside SQL), so drop the
# deleted observations' snapshots explicitly rather than leaving them to accumulate.
await conn.execute(
f"DELETE FROM {fq_table('observation_history')} WHERE bank_id = $1 AND observation_id = ANY($2::uuid[])",
bank_id,
obs_ids,
)
if remaining_source_ids:
await conn.execute(
f"""
UPDATE {fq_table("memory_units")}
SET consolidated_at = NULL
WHERE id = ANY($1::uuid[])
AND fact_type IN ('experience', 'world')
""",
remaining_source_ids,
)
logger.info(
f"[OBSERVATIONS] Deleted {len(obs_ids)} observations, reset {len(remaining_source_ids)} "
f"source memories for re-consolidation in bank {bank_id}"
)
return len(obs_ids)
# --------------------------------------------------------------------- curation archive
#
# Invalidation moves a rejected memory between two tables rather than flagging it,
# so recall / consolidation / graph never carry a "valid?" predicate: live facts
# live in `memory_units`, invalidated ones in `invalidated_memory_units`. The
# archive is cold storage — no index, so it drops the `embedding` and
# `search_vector` columns, which are recomputed on the way back.
# The two recall-surface columns the archive omits. Both follow server config
# (embedding dimension, search backend), so keeping them out of the INSERT…SELECT
# round-trip makes a model or text-backend switch structurally unable to trip a
# type/dimension mismatch (#2209, #2503); each is recomputed on revert.
_ARCHIVE_OMITTED = ('"embedding"', '"search_vector"')
async def _memory_unit_columns(conn, fq_table: Callable[[str], str]) -> str:
"""The quoted, ordinal column list of `memory_units`.
Read from the catalog rather than hardcoded so a schema migration cannot make
the archive round-trip drift from the live table (the archive is created via
``LIKE memory_units``, so the lists line up).
"""
rows = await conn.fetch(
f"SELECT a.attname FROM pg_attribute a "
f"WHERE a.attrelid = '{fq_table('memory_units')}'::regclass "
f"AND a.attnum > 0 AND NOT a.attisdropped ORDER BY a.attnum"
)
return ", ".join(f'"{r["attname"]}"' for r in rows)
async def _archive_columns(conn, fq_table: Callable[[str], str]) -> str:
"""`_memory_unit_columns` minus the two the archive does not carry."""
collist = await _memory_unit_columns(conn, fq_table)
return ", ".join(c for c in (s.strip() for s in collist.split(",")) if c not in _ARCHIVE_OMITTED)
_ARCHIVE_SELECT = (
"id, text, fact_type, context, occurred_start, occurred_end, mentioned_at, "
"document_id, chunk_id, tags, metadata, proof_count, event_date, created_at, "
"consolidated_at, entity_ids"
)
def _archived_stored(row: Any) -> StoredMemory:
"""Map an `invalidated_memory_units` row onto :class:`StoredMemory`."""
return StoredMemory(
unit_id=str(row["id"]),
text=row["text"],
fact_type=row["fact_type"],
context=row["context"],
document_id=row["document_id"],
chunk_id=str(row["chunk_id"]) if row["chunk_id"] else None,
tags=list(row["tags"] or []),
metadata=row["metadata"] if isinstance(row["metadata"], dict) else None,
proof_count=row["proof_count"] or 1,
event_date=row["event_date"],
occurred_start=row["occurred_start"],
occurred_end=row["occurred_end"],
mentioned_at=row["mentioned_at"],
created_at=row["created_at"],
consolidated_at=row["consolidated_at"],
entity_ids=[str(e) for e in (row["entity_ids"] or [])],
)
async def get_archived_memory(*, conn, fq_table, bank_id: str, unit_id: str) -> StoredMemory | None:
row = await conn.fetchrow(
f"SELECT {_ARCHIVE_SELECT} FROM {fq_table('invalidated_memory_units')} WHERE id = $1 AND bank_id = $2",
str(unit_id),
bank_id,
)
return _archived_stored(row) if row else None
async def invalidate_memory(*, conn, fq_table, bank_id: str, unit_id: str, reason: str | None) -> bool:
mu = fq_table("memory_units")
arch = fq_table("invalidated_memory_units")
ue = fq_table("unit_entities")
arch_cols = await _archive_columns(conn, fq_table)
# Snapshot the entity ids before the delete cascade takes `unit_entities`, so
# revert can restore the postings the move is about to drop.
entity_ids = [
r["entity_id"] for r in await conn.fetch(f"SELECT entity_id FROM {ue} WHERE unit_id = $1", str(unit_id))
]
# Causal edges are retain-time extraction output the FK cascade would destroy for good —
# unlike temporal/semantic links they can't be recomputed, so snapshot their descriptors onto
# the archive row and revert rematerializes them (#2864).
from ...retain.link_utils import snapshot_causal_links
causal_links = await snapshot_causal_links(conn, bank_id, str(unit_id))
inserted = await conn.fetchval(
f"INSERT INTO {arch} ({arch_cols}, invalidation_reason, invalidated_at, entity_ids, causal_links) "
f"SELECT {arch_cols}, $2, now(), $3::uuid[], $5::jsonb FROM {mu} WHERE id = $1 AND bank_id = $4 "
f"RETURNING id",
str(unit_id),
reason,
entity_ids,
bank_id,
json.dumps([descriptor.as_json_dict() for descriptor in causal_links]),
)
if inserted is None:
return False
# The cascade prunes `unit_entities` and `memory_links` with the row.
await conn.execute(f"DELETE FROM {mu} WHERE id = $1 AND bank_id = $2", str(unit_id), bank_id)
return True
async def set_invalidation_reason(*, conn, fq_table, bank_id: str, unit_id: str, reason: str | None) -> None:
await conn.execute(
f"UPDATE {fq_table('invalidated_memory_units')} SET invalidation_reason = $3 WHERE id = $1 AND bank_id = $2",
str(unit_id),
bank_id,
reason,
)
async def restore_memory(*, conn, fq_table, bank_id: str, unit_id: str) -> StoredMemory | None:
mu = fq_table("memory_units")
arch = fq_table("invalidated_memory_units")
ue = fq_table("unit_entities")
ent = fq_table("entities")
arch_cols = await _archive_columns(conn, fq_table)
arch_row = await conn.fetchrow(
f"SELECT {_ARCHIVE_SELECT} FROM {arch} WHERE id = $1 AND bank_id = $2", str(unit_id), bank_id
)
if arch_row is None:
return None
# Move the row back. The archive omits embedding/search_vector, so both default
# to NULL here; search_vector is rebuilt now, the embedding by the caller.
await conn.execute(
f"INSERT INTO {mu} ({arch_cols}) SELECT {arch_cols} FROM {arch} WHERE id = $1 AND bank_id = $2",
str(unit_id),
bank_id,
)
# Rebuild search_vector with the *current* backend, so a backend change while
# the fact sat archived cannot leave a stale/wrong-type vector (#2503). None
# means the backend indexes base columns directly and leaves it empty.
from ...db.ops_postgresql import pg_search_vector_expr
sv_expr = pg_search_vector_expr(get_config())
if sv_expr is not None:
await conn.execute(
f"UPDATE {mu} SET search_vector = {sv_expr} WHERE id = $1 AND bank_id = $2", str(unit_id), bank_id
)
# Re-consolidate from scratch; links are rebuilt by graph maintenance.
await conn.execute(
f"UPDATE {mu} SET consolidated_at = NULL, consolidation_failed_at = NULL, updated_at = now() "
f"WHERE id = $1 AND bank_id = $2",
str(unit_id),
bank_id,
)
# Restore the entity postings for entities that still exist — some may have
# been swept as orphans while the memory was archived.
if arch_row["entity_ids"]:
await conn.execute(
f"INSERT INTO {ue} (unit_id, entity_id) "
f"SELECT $1, eid FROM unnest($2::uuid[]) AS eid "
f"WHERE EXISTS (SELECT 1 FROM {ent} e WHERE e.id = eid AND e.bank_id = $3) "
f"ON CONFLICT DO NOTHING",
str(unit_id),
arch_row["entity_ids"],
bank_id,
)
# Rematerialize the causal edges parked at invalidation (#2864). Edges whose peer is still
# archived or permanently deleted are skipped — the peer keeps its own copy and recreates the
# edge when it reverts, so the restore is order-independent and idempotent.
from ...retain.link_utils import rematerialize_causal_links
from .graph import _ops_for
causal_json = await conn.fetchval(
f"SELECT causal_links FROM {arch} WHERE id = $1 AND bank_id = $2", str(unit_id), bank_id
)
if causal_json:
await rematerialize_causal_links(conn, bank_id, conn.parse_json(causal_json) or [], ops=_ops_for(conn))
# Invalidation cascaded away this unit's derived outgoing links; queue it so graph maintenance
# rebuilds them (the drain only touches queued units — it never scans for missing adjacency).
await _ops_for(conn).enqueue_graph_maintenance(
conn, fq_table("graph_maintenance_queue"), bank_id, [uuid.UUID(str(unit_id))]
)
await conn.execute(f"DELETE FROM {arch} WHERE id = $1 AND bank_id = $2", str(unit_id), bank_id)
return _archived_stored(arch_row)
async def set_memory_embedding(*, conn, fq_table, bank_id: str, unit_id: str, embedding) -> None:
await conn.execute(
f"UPDATE {fq_table('memory_units')} SET embedding = $3::vector WHERE id = $1 AND bank_id = $2",
str(unit_id),
bank_id,
embedding,
)
async def clear_unit_entities(*, conn, fq_table, bank_id: str, unit_id: str) -> None:
await conn.execute(f"DELETE FROM {fq_table('unit_entities')} WHERE unit_id = $1", str(unit_id))
async def apply_edit(
*,
conn,
fq_table,
bank_id: str,
unit_id: str,
text: str,
context: str | None,
fact_type: str,
occurred_start,
occurred_end,
event_date,
mentioned_at,
entity_ids: list[str] | None,
) -> None:
# `entity_ids` and `mentioned_at` are unused here: the entity postings are
# re-linked into `unit_entities` by the caller, and an edit does not move the
# mention time. Both are on the signature for a store that carries entities on
# the memory and rebuilds it wholesale.
from ...causal_links import CAUSAL_LINK_TYPES
from ...db.ops_postgresql import pg_search_vector_expr
mu = fq_table("memory_units")
ml = fq_table("memory_links")
# The caller enqueues the relink victims (and the edited unit itself, via
# ``include_affected_units``) before invoking this — one combined queue insert keeps the
# graph-maintenance queue's lock ordering intact.
# Keep the stored text-search vector in sync with the edited text/context.
# Reference the bind parameters, not the columns: PostgreSQL evaluates the
# UPDATE's RHS before the sibling SET assignments land, so a column reference
# would see the pre-edit values.
sv_expr = pg_search_vector_expr(get_config(), text_col="$3", context_col="$4")
sv_clause = f", search_vector = {sv_expr}" if sv_expr else ""
await conn.execute(
f"""
UPDATE {mu}
SET text = $3, context = $4, fact_type = $5, occurred_start = $6, occurred_end = $7,
event_date = $8, consolidated_at = NULL, consolidation_failed_at = NULL,
edited_at = now(), updated_at = now(){sv_clause}
WHERE id = $1 AND bank_id = $2
""",
str(unit_id),
bank_id,
text,
context,
fact_type,
occurred_start,
occurred_end,
event_date,
)
# Drop only the DERIVED links — graph maintenance recomputes temporal/semantic. Causal edges
# are retain-time extraction output that nothing recreates, so an edit preserves them (#2864).
await conn.execute(
f"DELETE FROM {ml} WHERE (from_unit_id = $1 OR to_unit_id = $1) AND NOT (link_type = ANY($2::text[]))",
str(unit_id),
list(CAUSAL_LINK_TYPES),
)
__all__ = [
"apply_edit",
"clear_unit_entities",
"delete_document",
"delete_observations",
"delete_stale_observations",
"get_archived_memory",
"insert_facts",
"invalidate_memory",
"observations_for_sources",
"restore_memory",
"set_invalidation_reason",
"set_memory_embedding",
]
@@ -1,509 +0,0 @@
"""The default memories store: Postgres holds the memories and the links.
This is the behaviour Hindsight has always had, stated as an implementation of
:class:`~hindsight_api.engine.memories.base.MemoriesExtension` rather than as the
absence of one. Rows go in `memory_units`, the joins around it are `memory_links`
and `unit_entities`, and every read is SQL writing a row *is* indexing it, so
:meth:`index_facts` has nothing left to do.
The class is deliberately thin. Each method delegates to a plain function in
:mod:`hindsight_api.engine.memories.pg`, split by what calls it curation,
graph, reads, writes so a change to one area is a change to one file, and the
SQL is grouped by concern rather than piled behind a class. The two retrieval
arms delegate further out still, to the query functions that already own them in
:mod:`hindsight_api.engine.search.retrieval`.
Keeping this as an explicit store (rather than an ``if store is None`` branch at
each call site) means the default path is the one the whole test suite exercises,
and a second implementation cannot change it by accident.
"""
from __future__ import annotations
from datetime import datetime
from typing import Any
from .base import DeletePredicate, MemoriesExtension, MemoryPatch, ScanPage, StoredMemory
from .pg import counts, curation, graph, reads, writes
class PostgresMemories(MemoriesExtension):
"""Memories in `memory_units`, links in `memory_links` / `unit_entities`."""
name = "postgres"
# ------------------------------------------------------------------ writes
async def insert_facts(
self,
*,
conn,
ops,
bank_id: str,
facts: list,
document_id: str | None = None,
defer_index: bool = False,
txn=None,
) -> list[str]:
# `txn` is ignored: Postgres memories live in the caller's own transaction, so the
# write is already atomic with it — there is no separate store to hold invisible.
# `defer_index` is meaningless here: the INSERT that returns the ids is
# also what indexes the facts, so there is nothing to defer.
return await writes.insert_facts(conn=conn, ops=ops, bank_id=bank_id, facts=facts, document_id=document_id)
async def delete_facts(self, bank_id: str, unit_ids: list[str], *, txn=None) -> None:
"""No-op: the caller's `memory_units` DELETE (or its FK cascade) removed them."""
async def delete_where(self, bank_id: str, predicate: DeletePredicate, txn=None) -> int:
"""No-op: predicate deletes are issued as SQL by the caller that owns the transaction."""
return 0
async def delete_document(self, *, conn, fq_table, bank_id: str, document_id: str, txn=None) -> None:
# `txn` ignored: Postgres memories are covered by the caller's own transaction.
await writes.delete_document(conn=conn, fq_table=fq_table, bank_id=bank_id, document_id=document_id)
async def drop_bank_storage(self, bank_id: str) -> None:
"""No-op: deleting the bank cascades to its memories."""
async def delete_observations(self, *, conn, fq_table, bank_id: str, txn=None) -> None:
await writes.delete_observations(conn=conn, fq_table=fq_table, bank_id=bank_id)
async def update_memories(self, bank_id: str, patches: list[MemoryPatch], txn=None) -> None:
"""No-op: the caller's UPDATE already wrote the row it holds open."""
# ------------------------------------------------------------------ recall arms
async def search(
self,
*,
conn,
bank_id: str,
fact_types: list[str],
query_embedding: str,
query_text: str,
limit: int,
tags: list[str] | None = None,
tags_match: str = "any",
tag_groups: list | None = None,
created_after: datetime | None = None,
created_before: datetime | None = None,
min_semantic: float | None = None,
min_keyword: float | None = None,
graph_seed_min_similarity: float | None = None,
) -> "dict[str, SemanticBm25Result]":
# Imported here: retrieval imports this package, so a module-level import
# would close the cycle.
from ..search.retrieval import retrieve_semantic_bm25_combined_sql
return await retrieve_semantic_bm25_combined_sql(
conn,
query_embedding,
query_text,
bank_id,
fact_types,
limit,
tags=tags,
tags_match=tags_match,
tag_groups=tag_groups,
created_after=created_after,
created_before=created_before,
min_semantic=min_semantic,
min_keyword=min_keyword,
graph_seed_min_similarity=graph_seed_min_similarity,
)
async def temporal_search(
self,
*,
conn,
bank_id: str,
fact_types: list[str],
query_embedding: str,
start_date: datetime,
end_date: datetime,
limit: int,
semantic_threshold: float = 0.1,
tags: list[str] | None = None,
tags_match: str = "any",
tag_groups: list | None = None,
created_after: datetime | None = None,
created_before: datetime | None = None,
) -> dict[str, list]:
from ..search.retrieval import retrieve_temporal_combined_sql
return await retrieve_temporal_combined_sql(
conn,
query_embedding,
bank_id,
fact_types,
start_date,
end_date,
limit,
semantic_threshold=semantic_threshold,
tags=tags,
tags_match=tags_match,
tag_groups=tag_groups,
created_after=created_after,
created_before=created_before,
)
# ------------------------------------------------------------------ addressed reads
async def get_memories(self, *, conn, fq_table, bank_id: str, unit_ids: list[str]) -> list[StoredMemory]:
return await reads.get_memories(conn=conn, fq_table=fq_table, bank_id=bank_id, unit_ids=unit_ids)
async def scan_memories(
self,
*,
conn,
fq_table,
bank_id: str,
fact_types: list[str] | None = None,
limit: int = 100,
page_token: str = "",
tags: list[str] | None = None,
tags_match: str = "any",
tag_groups: list | None = None,
document_id: str | None = None,
metadata_equals: dict[str, str] | None = None,
skip: int = 0,
include_edges: bool = False,
) -> ScanPage:
return await reads.scan_memories(
conn=conn,
fq_table=fq_table,
bank_id=bank_id,
fact_types=fact_types,
limit=limit,
page_token=page_token,
tags=tags,
tags_match=tags_match,
tag_groups=tag_groups,
document_id=document_id,
metadata_equals=metadata_equals,
skip=skip,
include_edges=include_edges,
)
async def count_memories(self, *, conn, fq_table, bank_id: str) -> dict[str, int]:
return await reads.count_memories(conn=conn, fq_table=fq_table, bank_id=bank_id)
async def list_tags(
self,
*,
conn,
fq_table,
bank_id: str,
pattern: str | None = None,
limit: int = 100,
offset: int = 0,
) -> dict[str, Any]:
return await reads.list_tags(
conn=conn, fq_table=fq_table, bank_id=bank_id, pattern=pattern, limit=limit, offset=offset
)
async def find_unconsolidated(
self,
*,
conn,
fq_table,
bank_id: str,
fact_types: list[str],
limit: int,
scope_tags: list[str] | None = None,
) -> list[StoredMemory]:
return await reads.find_unconsolidated(
conn=conn,
fq_table=fq_table,
bank_id=bank_id,
fact_types=fact_types,
limit=limit,
scope_tags=scope_tags,
)
async def count_unconsolidated(
self,
*,
conn,
fq_table,
bank_id: str,
fact_types: list[str],
scopes: list[list[str] | None],
limit: int,
) -> int:
return await reads.count_unconsolidated(
conn=conn, fq_table=fq_table, bank_id=bank_id, fact_types=fact_types, scopes=scopes, limit=limit
)
async def mark_consolidated(
self,
*,
conn,
fq_table,
bank_id: str,
unit_ids: list[str],
when: datetime | None,
failed: bool = False,
txn=None,
) -> None:
await reads.mark_consolidated(
conn=conn, fq_table=fq_table, bank_id=bank_id, unit_ids=unit_ids, when=when, failed=failed
)
async def any_memory_updated_since(
self,
*,
conn,
fq_table,
bank_id: str,
since: datetime,
fact_types: list[str] | None = None,
tags: list[str] | None = None,
tags_match: str = "any",
tag_groups: list | None = None,
) -> bool:
return await reads.any_memory_updated_since(
conn=conn,
fq_table=fq_table,
bank_id=bank_id,
since=since,
fact_types=fact_types,
tags=tags,
tags_match=tags_match,
tag_groups=tag_groups,
)
# -- count surfaces --
async def consolidation_freshness(self, *, conn, fq_table, bank_id: str) -> dict[str, Any]:
return await counts.consolidation_freshness(conn=conn, fq_table=fq_table, bank_id=bank_id)
async def document_memory_counts(self, *, conn, fq_table, bank_id: str, document_ids: list[str]) -> dict[str, int]:
return await counts.document_memory_counts(
conn=conn, fq_table=fq_table, bank_id=bank_id, document_ids=document_ids
)
async def link_counts(self, *, conn, fq_table, bank_id: str) -> dict[str, int]:
return await counts.link_counts(conn=conn, fq_table=fq_table, bank_id=bank_id)
async def memories_timeseries(
self, *, conn, fq_table, bank_id: str, time_field: str, trunc: str, since: datetime
) -> list[dict[str, Any]]:
return await counts.memories_timeseries(
conn=conn, fq_table=fq_table, bank_id=bank_id, time_field=time_field, trunc=trunc, since=since
)
async def observation_scope_counts(self, *, conn, fq_table, bank_id: str) -> list[dict[str, Any]]:
return await counts.observation_scope_counts(conn=conn, fq_table=fq_table, bank_id=bank_id)
# ------------------------------------------------------------------ observations
async def upsert_observation(self, *, conn, bank_id: str, record, txn=None) -> None:
"""No-op: the observation was written as a `memory_units` row by the caller."""
async def observations_for_sources(
self, *, conn, ops, fq_table, bank_id: str, unit_ids: list[str]
) -> list[StoredMemory]:
return await writes.observations_for_sources(
conn=conn, ops=ops, fq_table=fq_table, bank_id=bank_id, unit_ids=unit_ids
)
async def delete_stale_observations(self, *, conn, ops, fq_table, bank_id: str, fact_ids: list) -> int:
return await writes.delete_stale_observations(
conn=conn, ops=ops, fq_table=fq_table, bank_id=bank_id, fact_ids=fact_ids
)
# ------------------------------------------------------------------ curation reads
async def list_memory_units(
self,
*,
conn,
ops,
fq_table,
bank_id: str,
fact_type: str | None = None,
search_query: str | None = None,
consolidation_state: str | None = None,
state: str | None = None,
document_id: str | None = None,
entity_id: str | None = None,
tags: list[str] | None = None,
tags_match: str = "any",
created_before: datetime | None = None,
limit: int = 100,
offset: int = 0,
) -> dict[str, Any]:
return await curation.list_memory_units(
conn=conn,
ops=ops,
fq_table=fq_table,
bank_id=bank_id,
fact_type=fact_type,
search_query=search_query,
consolidation_state=consolidation_state,
state=state,
document_id=document_id,
entity_id=entity_id,
tags=tags,
tags_match=tags_match,
created_before=created_before,
limit=limit,
offset=offset,
)
async def get_memory_unit(self, *, conn, ops, fq_table, bank_id: str, unit_id: str) -> dict[str, Any] | None:
return await curation.get_memory_unit(conn=conn, ops=ops, fq_table=fq_table, bank_id=bank_id, unit_id=unit_id)
# -- curation archive --
async def get_archived_memory(self, *, conn, fq_table, bank_id: str, unit_id: str) -> StoredMemory | None:
return await writes.get_archived_memory(conn=conn, fq_table=fq_table, bank_id=bank_id, unit_id=unit_id)
async def invalidate_memory(
self, *, conn, fq_table, bank_id: str, unit_id: str, reason: str | None, txn=None
) -> bool:
return await writes.invalidate_memory(
conn=conn, fq_table=fq_table, bank_id=bank_id, unit_id=unit_id, reason=reason
)
async def set_invalidation_reason(self, *, conn, fq_table, bank_id: str, unit_id: str, reason: str | None) -> None:
await writes.set_invalidation_reason(
conn=conn, fq_table=fq_table, bank_id=bank_id, unit_id=unit_id, reason=reason
)
async def restore_memory(self, *, conn, fq_table, bank_id: str, unit_id: str, txn=None) -> StoredMemory | None:
return await writes.restore_memory(conn=conn, fq_table=fq_table, bank_id=bank_id, unit_id=unit_id)
async def set_memory_embedding(self, *, conn, fq_table, bank_id: str, unit_id: str, embedding, txn=None) -> None:
await writes.set_memory_embedding(
conn=conn, fq_table=fq_table, bank_id=bank_id, unit_id=unit_id, embedding=embedding
)
async def clear_unit_entities(self, *, conn, fq_table, bank_id: str, unit_id: str) -> None:
await writes.clear_unit_entities(conn=conn, fq_table=fq_table, bank_id=bank_id, unit_id=unit_id)
async def apply_edit(
self,
*,
conn,
fq_table,
bank_id: str,
unit_id: str,
text: str,
context: str | None,
fact_type: str,
occurred_start,
occurred_end,
event_date,
mentioned_at,
entity_ids: list[str] | None,
txn=None,
) -> None:
await writes.apply_edit(
conn=conn,
fq_table=fq_table,
bank_id=bank_id,
unit_id=unit_id,
text=text,
context=context,
fact_type=fact_type,
occurred_start=occurred_start,
occurred_end=occurred_end,
event_date=event_date,
mentioned_at=mentioned_at,
entity_ids=entity_ids,
)
async def list_entities(
self,
*,
conn,
fq_table,
bank_id: str,
search: str | None = None,
limit: int = 100,
offset: int = 0,
) -> dict[str, Any]:
return await curation.list_entities(
conn=conn, fq_table=fq_table, bank_id=bank_id, search=search, limit=limit, offset=offset
)
# ------------------------------------------------------------------ graph
async def graph_units(
self,
*,
conn,
fq_table,
bank_id: str,
fact_type: str | None = None,
search_query: str | None = None,
document_id: str | None = None,
chunk_id: str | None = None,
tags: list[str] | None = None,
tags_match: str = "all_strict",
limit: int = 1000,
) -> dict[str, Any]:
return await graph.graph_units(
conn=conn,
fq_table=fq_table,
bank_id=bank_id,
fact_type=fact_type,
search_query=search_query,
document_id=document_id,
chunk_id=chunk_id,
tags=tags,
tags_match=tags_match,
limit=limit,
)
async def graph_entity_rows(self, *, conn, fq_table, bank_id: str, unit_ids: list[str]) -> list[dict[str, Any]]:
return await graph.graph_entity_rows(conn=conn, fq_table=fq_table, bank_id=bank_id, unit_ids=unit_ids)
async def graph_direct_links(self, *, conn, fq_table, bank_id: str, unit_ids: list[str]) -> list[dict[str, Any]]:
return await graph.graph_direct_links(conn=conn, fq_table=fq_table, bank_id=bank_id, unit_ids=unit_ids)
async def entity_memory_counts(
self, *, conn, fq_table, bank_id: str, entity_ids: list[str] | None = None
) -> dict[str, int]:
return await graph.entity_memory_counts(conn=conn, fq_table=fq_table, bank_id=bank_id, entity_ids=entity_ids)
async def entities_for_units(self, *, conn, fq_table, bank_id: str, unit_ids: list[str]) -> dict[str, list[str]]:
return await graph.entities_for_units(conn=conn, fq_table=fq_table, bank_id=bank_id, unit_ids=unit_ids)
async def entity_map_for_units(
self, *, conn, fq_table, bank_id: str, unit_ids: list[str]
) -> dict[str, list[dict[str, str]]]:
return await graph.entity_map_for_units(conn=conn, fq_table=fq_table, bank_id=bank_id, unit_ids=unit_ids)
# ------------------------------------------------------------------ maintenance
async def record_unit_entities(
self, *, conn, ops, fq_table, bank_id: str | None = None, unit_ids: list[Any], entity_ids: list[Any]
) -> None:
# The join is keyed by global unit id, so bank_id is not needed here.
await ops.bulk_insert_unit_entities(conn, fq_table("unit_entities"), unit_ids, entity_ids)
async def enqueue_relink_victims(
self, *, conn, fq_table, bank_id: str, affected_unit_ids: list, include_affected_units: bool = False
) -> int:
return await graph.enqueue_relink_victims(
conn=conn,
fq_table=fq_table,
bank_id=bank_id,
affected_unit_ids=affected_unit_ids,
include_affected_units=include_affected_units,
)
async def relink_pass(self, *, backend, fq_table, bank_id: str, config) -> dict:
return await graph.relink_pass(backend=backend, fq_table=fq_table, bank_id=bank_id, config=config)
async def prune_orphan_entities(self, *, conn, fq_table, bank_id: str) -> int:
return await graph.prune_orphan_entities(conn=conn, fq_table=fq_table, bank_id=bank_id)
async def prune_stale_cooccurrences(self, *, conn, fq_table, bank_id: str) -> int:
return await graph.prune_stale_cooccurrences(conn=conn, fq_table=fq_table, bank_id=bank_id)
__all__ = ["PostgresMemories"]
File diff suppressed because it is too large Load Diff
@@ -1,242 +0,0 @@
"""Models describing what a mental model refresh did.
A refresh resolves a scope, picks full-vs-delta, runs reflect over a bounded
snapshot, and (in delta mode) applies structured operations to the existing
document. Every one of those steps can quietly produce a document that isn't
what the user expected, and until now the reasoning behind each only ever
reached a log line.
These models carry that reasoning out to callers, so both the dry run (preview,
nothing persisted) and ``trigger.keep_trace`` (recorded on every real refresh,
including the cron- and consolidation-driven ones no human is watching) can
report it.
Kept out of ``response_models`` on purpose: these reference the tag-group types
from ``search.tags``, and ``response_models`` is imported early enough in the
engine's import graph that pulling the search package in from there is a cycle.
"""
from datetime import datetime
from typing import Any, Literal
from pydantic import BaseModel, ConfigDict, Field
from .response_models import LLMCallTrace, TokenUsage
from .search.tags import TagGroup, TagsMatch
RefreshMode = Literal["full", "delta"]
ModeFallbackReason = Literal[
"no_baseline_content",
"source_query_changed",
"structured_doc_unreadable",
"delta_ops_failed",
"delta_ops_all_skipped",
]
RefreshOutcome = Literal[
"content_written",
"content_preserved_no_new_facts",
"refresh_failed_empty_candidate",
"refresh_failed_delta_not_applied",
]
class MentalModelRefreshScope(BaseModel):
"""The memory scope a refresh actually resolved to.
A model's stored ``tags`` are not what filters memories — ``tags_match``
defaults to ``all_strict`` when tags are present, and ``tag_groups``
override flat tags entirely. This reports the resolved result.
"""
tags: list[str] | None = Field(default=None, description="Flat tags used to filter memories (null when unused).")
tags_match: TagsMatch = Field(description="Resolved tag match mode.")
tag_groups: list[TagGroup] | None = Field(
default=None, description="Compound tag expressions used instead of flat tags, when set."
)
fact_types: list[str] | None = Field(default=None, description="Fact types retrieved (null means all).")
exclude_mental_models: bool = Field(description="Whether other mental models were excluded from the reflect loop.")
exclude_mental_model_ids: list[str] = Field(
default_factory=list, description="Mental models excluded by ID (always includes the model being refreshed)."
)
class MentalModelRefreshWindow(BaseModel):
"""The time window a refresh read memories from."""
created_after: datetime | None = Field(
default=None,
description=(
"Lower bound on memory creation time. Set only in delta mode, where it is the model's "
"last_refreshed_at — so a delta refresh only sees memories newer than the last one."
),
)
created_before: datetime = Field(
description=(
"Database-time snapshot bounding the refresh. Memories committed after this are not read, "
"so they stay newer than the persisted watermark and are caught by the next refresh."
)
)
watermark: datetime | None = Field(
default=None,
description=(
"The last_refreshed_at a real refresh would persist: the newest in-scope memory visible at "
"the snapshot, not now(). Null means no in-scope memory was visible."
),
)
class MentalModelFactCounts(BaseModel):
"""Facts the refresh saw, keyed by fact type.
``retrieved`` and ``used`` diverging is the single most common cause of a
disappointing refresh: recall found plenty, but the reflect agent declared
none of it relevant to the topic, so none of it reached the document.
"""
retrieved: dict[str, int] = Field(
default_factory=dict, description="Facts the reflect agent's tool calls returned, by fact type."
)
used: dict[str, int] = Field(
default_factory=dict, description="Facts the agent declared it actually based the answer on, by fact type."
)
class MentalModelDeltaOperations(BaseModel):
"""Structured operations a delta refresh emitted against the existing document."""
applied: list[dict[str, Any]] = Field(
default_factory=list, description="Operations applied to the document, in order."
)
skipped: list[dict[str, Any]] = Field(
default_factory=list, description="Operations dropped as invalid, each with a reason."
)
class MentalModelTraceToolCall(BaseModel):
"""One reflect tool call made during a refresh.
``output`` is carried only by the dry run, which persists nothing. The trace
stored on the model row keeps ``result_count`` instead: it is re-read on every
fetch, so embedding full recall payloads there would bloat the row without
bound. Raw prompts and responses are available separately via LLM request
tracing.
"""
tool: str = Field(description="Tool name: recall, search_observations, get_mental_model, expand, …")
reason: str | None = Field(default=None, description="The agent's stated reason for the call.")
input: dict[str, Any] = Field(default_factory=dict, description="Tool input parameters.")
output: dict[str, Any] | None = Field(
default=None,
description=(
"What the tool returned. Present on a dry run, which stores nothing; omitted from the "
"trace persisted by a real refresh to keep that row bounded."
),
)
updated_at: datetime | None = Field(
default=None,
description=(
"The refresh window's lower bound as given to this call — the delta watermark. Named "
"for what it actually filters: the predicate is on the memory's updated_at, so a "
"memory merely touched since the last refresh qualifies. Null means the tool applies "
"no time bound at all, so its results are not limited to the window (mental-model "
"lookup and chunk expansion behave this way)."
),
)
result_count: int | None = Field(default=None, description="Number of items the tool returned, when countable.")
duration_ms: int = Field(description="Execution time in milliseconds.")
iteration: int = Field(default=0, description="Agent loop iteration (1-based) this call belongs to.")
class MentalModelRefreshTrace(BaseModel):
"""Execution trace of a mental model refresh, recorded when trigger.keep_trace is on.
Deliberately shaped like reflect's trace — the calls the agent made, plus the
refresh-specific decision and nothing more. This is persisted on the mental
model row and re-read on every fetch, so anything derivable from elsewhere is
left out: the evidence lives in ``reflect_response.based_on``, and the
resolved scope and snapshot window are reported by the dry run.
"""
recorded_at: datetime | None = Field(default=None, description="When this trace was recorded.")
effective_mode: RefreshMode = Field(description="Whether the refresh ran as full or delta.")
mode_fallback_reason: ModeFallbackReason | None = Field(
default=None, description="Why delta was requested but not applied, if that happened."
)
outcome: RefreshOutcome = Field(description="What the refresh did with the document.")
tool_calls: list[MentalModelTraceToolCall] = Field(
default_factory=list, description="Reflect tool calls made during the refresh."
)
llm_calls: list[LLMCallTrace] = Field(default_factory=list, description="LLM calls made during the refresh.")
delta_operations: MentalModelDeltaOperations | None = Field(
default=None, description="Structured operations emitted, in delta mode."
)
usage: TokenUsage | None = Field(default=None, description="Token usage across the refresh's LLM calls.")
duration_ms: int = Field(default=0, description="Wall-clock duration of the refresh.")
warnings: list[str] = Field(
default_factory=list, description="Conditions worth a human's attention, in plain language."
)
class MentalModelDryRunRefreshResult(BaseModel):
"""Preview of what a mental model refresh would do, having changed nothing.
Runs the real pipeline same scope resolution, same reflect call, same
delta operations then reports the result instead of persisting it. The
model's content, structured content, watermark, and last_refreshed_at are
all left untouched, so a delta dry run is repeatable: it reads the same
window the next real refresh would.
"""
model_config = ConfigDict(
json_schema_extra={
"example": {
"mental_model_id": "coding-style",
"name": "Coding Style",
"requested_mode": "delta",
"effective_mode": "full",
"mode_fallback_reason": "source_query_changed",
"outcome": "content_written",
"would_persist": True,
"facts": {"retrieved": {"observation": 12}, "used": {"observation": 4}},
"warnings": [],
}
}
)
mental_model_id: str = Field(description="The mental model previewed.")
name: str = Field(description="Display name of the mental model.")
requested_mode: RefreshMode = Field(description="The mode asked for (from the model's trigger, or overridden).")
effective_mode: RefreshMode = Field(description="The mode the refresh actually ran in.")
mode_fallback_reason: ModeFallbackReason | None = Field(
default=None, description="Why delta was requested but not applied, if that happened."
)
outcome: RefreshOutcome = Field(description="What a real refresh would do with the document.")
would_persist: bool = Field(description="Whether a real refresh would write new content.")
scope: MentalModelRefreshScope = Field(description="The resolved memory scope.")
window: MentalModelRefreshWindow = Field(description="The snapshot window read from.")
facts: MentalModelFactCounts = Field(description="Facts retrieved versus actually used.")
based_on: dict[str, list[dict[str, Any]]] = Field(
default_factory=dict,
description=(
"The evidence this run would ground the document on, keyed by fact type — the same "
"shape a refresh persists under reflect_response.based_on. Returned so a preview can "
"show its sources without having to write them anywhere."
),
)
current_content: str = Field(description="The model's content as it stands now.")
candidate_content: str = Field(description="Raw reflect synthesis, before any delta operations.")
preview_content: str = Field(
description="The content a real refresh would store: the delta-edited document, or the candidate in full mode."
)
diff: str = Field(description="Unified diff from current_content to preview_content. Empty when identical.")
delta_operations: MentalModelDeltaOperations | None = Field(
default=None, description="Structured operations emitted, in delta mode."
)
trace: MentalModelRefreshTrace = Field(description="Execution trace of the run, always included for a dry run.")
usage: TokenUsage = Field(default_factory=TokenUsage, description="Token usage across the run's LLM calls.")
duration_ms: int = Field(default=0, description="Wall-clock duration of the run.")
warnings: list[str] = Field(
default_factory=list, description="Conditions worth a human's attention, in plain language."
)
@@ -19,18 +19,10 @@ class BatchRetainParentMetadata:
total_tokens: int
num_sub_batches: int
is_parent: bool = True
# Set only when the whole batch targets a single document, so the operations
# list surfaces which document an in-flight retain is (re)writing. The
# documents UI cross-checks this to badge rows as "updating". Multi-document
# batches leave it None and are matched per single-document child instead.
document_id: str | None = None
def to_dict(self) -> dict[str, Any]:
"""Convert to dict for JSON serialization, omitting document_id when unset."""
data = asdict(self)
if data.get("document_id") is None:
data.pop("document_id", None)
return data
"""Convert to dict for JSON serialization."""
return asdict(self)
@dataclass
@@ -41,15 +33,10 @@ class BatchRetainChildMetadata:
parent_operation_id: str
sub_batch_index: int
total_sub_batches: int
# Set only when this child processes a single document (see the parent's note).
document_id: str | None = None
def to_dict(self) -> dict[str, Any]:
"""Convert to dict for JSON serialization, omitting document_id when unset."""
data = asdict(self)
if data.get("document_id") is None:
data.pop("document_id", None)
return data
"""Convert to dict for JSON serialization."""
return asdict(self)
@dataclass
@@ -155,27 +142,3 @@ class RefreshMentalModelMetadata:
def to_dict(self) -> dict[str, Any]:
"""Convert to dict for JSON serialization."""
return asdict(self)
@dataclass
class RefreshMentalModelOutcomeMetadata:
"""Machine-readable outcome metadata for a completed refresh_mental_model operation.
Refresh parity with RetainOutcomeMetadata (#2605): lets a monitoring layer
distinguish "refreshed with real content" from "refreshed empty" by reading
result_metadata alone, without a follow-up content fetch.
"""
content_len: int
populated_content: bool
based_on_counts: dict[str, int] = field(default_factory=dict)
# Delta operations the model emitted, as applied vs rejected. A refresh whose
# ops are routinely rejected still completes successfully with a plausible
# document, so the count is the only signal that some of this run's new facts
# never reached it. Both are 0 for a full-mode refresh, which emits no ops.
delta_ops_applied: int = 0
delta_ops_skipped: int = 0
def to_dict(self) -> dict[str, Any]:
"""Convert to dict for JSON serialization."""
return asdict(self)
@@ -72,7 +72,6 @@ class MarkitdownParser(FileParser):
ocr_base_url: str | None = None,
ocr_model: str | None = None,
ocr_prompt: str | None = None,
ocr_default_headers: dict[str, str] | None = None,
):
"""Initialize markitdown parser."""
# Lazy import to avoid requiring markitdown for all users
@@ -90,7 +89,6 @@ class MarkitdownParser(FileParser):
base_url=ocr_base_url,
model=ocr_model,
prompt=ocr_prompt,
default_headers=ocr_default_headers,
)
self._markitdown = MarkItDown(
llm_client=ocr_options.llm_client,
@@ -107,7 +105,6 @@ class MarkitdownParser(FileParser):
base_url: str | None,
model: str | None,
prompt: str | None,
default_headers: dict[str, str] | None,
) -> MarkitdownOcrOptions:
"""Build MarkItDown options for OpenAI-compatible image OCR."""
if not model or not model.strip():
@@ -132,15 +129,8 @@ class MarkitdownParser(FileParser):
except ImportError as e:
raise RuntimeError("openai package is required when Markitdown OCR is enabled.") from e
client_kwargs: dict[str, object] = {
"api_key": api_key,
"base_url": base_url.strip(),
}
if default_headers:
client_kwargs["default_headers"] = default_headers
return MarkitdownOcrOptions(
llm_client=OpenAI(**client_kwargs),
llm_client=OpenAI(api_key=api_key, base_url=base_url.strip()),
llm_model=model.strip(),
llm_prompt=prompt or DEFAULT_FILE_PARSER_MARKITDOWN_OCR_PROMPT,
)
@@ -196,11 +186,7 @@ class MarkitdownParser(FileParser):
if Path(filename).suffix.lower() not in _TEXT_EXTENSIONS:
return None
try:
# file_data may arrive as a non-``bytes`` buffer (e.g. a memoryview or
# a native/Rust-backed buffer object) that has no ``.decode``; coerce
# through the buffer protocol before the UTF-8 probe. The ``tmp.write``
# in the caller already relies only on the same buffer protocol.
bytes(file_data).decode("utf-8")
file_data.decode("utf-8")
except UnicodeDecodeError:
return None
from markitdown import StreamInfo
@@ -15,7 +15,6 @@ from .llamacpp_llm import LlamaCppLLM
from .mock_llm import MockLLM
from .none_llm import NoneLLM
from .openai_compatible_llm import OpenAICompatibleLLM
from .openai_responses_llm import OpenAIResponsesLLM
__all__ = [
"AnthropicLLM",
@@ -29,5 +28,4 @@ __all__ = [
"MockLLM",
"NoneLLM",
"OpenAICompatibleLLM",
"OpenAIResponsesLLM",
]
@@ -12,15 +12,12 @@ import asyncio
import json
import logging
import time
from contextlib import AbstractAsyncContextManager, nullcontext
from typing import Any, Callable
from typing import Any
from hindsight_api.engine.llm_interface import LLM_TOOL_CHOICE_AUTO, LLMInterface, LLMToolChoice
from hindsight_api.engine.llm_interface import LLMInterface
from hindsight_api.engine.llm_trace import LLMResponseUsage, stash_response_usage
from hindsight_api.engine.providers.llm_debug import dump_request_on_4xx
from hindsight_api.engine.response_models import LLMToolCall, LLMToolCallResult, TokenUsage
from hindsight_api.metrics import get_metrics_collector
from hindsight_api.worker.stage import set_stage
logger = logging.getLogger(__name__)
@@ -37,43 +34,6 @@ def _usage_from_anthropic_response(response: Any) -> LLMResponseUsage:
)
_EPHEMERAL_CACHE = {"type": "ephemeral"}
def _cached_system_blocks(system_prompt: str) -> list[dict[str, Any]]:
"""Render the system prompt as a block list with a cache_control marker.
Anthropic prompt caching is a prefix match: marking the (single) system
block caches tools + system together. The system prompt is stable per
scope fact extraction reuses it across every chunk, reflect and
consolidation keep their stable instructions there so repeat calls read
it at ~10% of the base input price. Markers below the model's minimum
cacheable prefix are silently ignored (no write premium), so marking is
safe unconditionally. This is the "inline-marker provider" strategy that
``LLMInterface.get_or_create_cached_prefix`` documents for Anthropic.
"""
return [{"type": "text", "text": system_prompt, "cache_control": _EPHEMERAL_CACHE}]
def _mark_last_message_for_caching(messages: list[dict[str, Any]]) -> None:
"""Add a cache_control marker to the final content block, in place.
Used on the multi-turn (tool-calling) path: the reflect agent loop resends
the entire growing conversation each iteration, so this request's
end-marker becomes the next iteration's cache read point. Together with
the system marker this uses 2 of the 4 allowed breakpoints.
"""
if not messages:
return
last = messages[-1]
content = last.get("content")
if isinstance(content, str):
if content.strip(): # the API rejects empty text blocks
last["content"] = [{"type": "text", "text": content, "cache_control": _EPHEMERAL_CACHE}]
elif isinstance(content, list) and content and isinstance(content[-1], dict):
content[-1]["cache_control"] = _EPHEMERAL_CACHE
class AnthropicLLM(LLMInterface):
"""
LLM provider using Anthropic's Claude models.
@@ -175,7 +135,6 @@ class AnthropicLLM(LLMInterface):
skip_validation: bool = False,
strict_schema: bool = False,
return_usage: bool = False,
attempt_context: Callable[[], AbstractAsyncContextManager[None]] | None = None,
) -> Any:
"""
Make an LLM API call with retry logic.
@@ -247,9 +206,7 @@ class AnthropicLLM(LLMInterface):
}
if system_prompt:
# One-shot calls share only the system prompt with each other, so
# that is the sole cache breakpoint on this path.
call_params["system"] = _cached_system_blocks(system_prompt)
call_params["system"] = system_prompt
if use_forced_tool:
# Single tool whose input_schema IS the response schema; force the model to
@@ -266,9 +223,7 @@ class AnthropicLLM(LLMInterface):
for attempt in range(max_retries + 1):
try:
async with attempt_context() if attempt_context is not None else nullcontext():
set_stage(f"llm.{self.provider}.{scope}.attempt={attempt + 1}/{max_retries + 1}")
response = await self._client.messages.create(**call_params)
response = await self._client.messages.create(**call_params)
# Stash usage before parse/validate, which may raise locally
# even though the provider charged for these tokens (#2387).
stash_response_usage(_usage_from_anthropic_response(response))
@@ -391,9 +346,6 @@ class AnthropicLLM(LLMInterface):
logger.error(f"Anthropic auth error (HTTP {e.status_code}), not retrying: {str(e)}")
raise
# Diagnostic dump (opt-in) of the exact request behind any 4xx.
dump_request_on_4xx(scope=scope, provider=self.provider, model=self.model, err=e, request=call_params)
last_exception = e
if attempt < max_retries:
# Check if it's a rate limit or server error
@@ -428,8 +380,7 @@ class AnthropicLLM(LLMInterface):
max_retries: int = 5,
initial_backoff: float = 1.0,
max_backoff: float = 30.0,
tool_choice: LLMToolChoice = LLM_TOOL_CHOICE_AUTO,
attempt_context: Callable[[], AbstractAsyncContextManager[None]] | None = None,
tool_choice: str | dict[str, Any] = "auto",
) -> LLMToolCallResult:
"""
Make an LLM API call with tool/function calling support.
@@ -499,11 +450,6 @@ class AnthropicLLM(LLMInterface):
else:
anthropic_messages.append({"role": role, "content": content})
# Multi-turn tool loop: cache the stable prefix (tools + system) via
# the system marker, and the growing conversation via an end-marker
# that the next iteration reads back.
_mark_last_message_for_caching(anthropic_messages)
call_params: dict[str, Any] = {
"model": self.model,
"messages": anthropic_messages,
@@ -511,7 +457,7 @@ class AnthropicLLM(LLMInterface):
"max_tokens": max_completion_tokens or 4096,
}
if system_prompt:
call_params["system"] = _cached_system_blocks(system_prompt)
call_params["system"] = system_prompt
if self._extra_body:
call_params["extra_body"] = self._extra_body
@@ -519,9 +465,7 @@ class AnthropicLLM(LLMInterface):
last_exception = None
for attempt in range(max_retries + 1):
try:
async with attempt_context() if attempt_context is not None else nullcontext():
set_stage(f"llm.{self.provider}.{scope}.attempt={attempt + 1}/{max_retries + 1}")
response = await self._client.messages.create(**call_params)
response = await self._client.messages.create(**call_params)
stash_response_usage(_usage_from_anthropic_response(response))
# Extract content and tool calls
@@ -589,8 +533,6 @@ class AnthropicLLM(LLMInterface):
except (APIConnectionError, APIStatusError) as e:
if isinstance(e, APIStatusError) and e.status_code in (401, 403):
raise
# Diagnostic dump (opt-in) of the exact request behind any 4xx.
dump_request_on_4xx(scope=scope, provider=self.provider, model=self.model, err=e, request=call_params)
last_exception = e
if attempt < max_retries:
await asyncio.sleep(min(initial_backoff * (2**attempt), max_backoff))
@@ -601,221 +543,7 @@ class AnthropicLLM(LLMInterface):
raise last_exception
raise RuntimeError("Anthropic tool call failed")
# ── Message Batches API (50% token discount) ─────────────────────────────
_BATCH_TOOL_NAME = "structured_response"
async def supports_batch_api(self) -> bool:
"""Anthropic supports batch operations via the Message Batches API."""
return True
@staticmethod
def _map_batch_status(processing_status: str) -> str:
"""Map Anthropic ``processing_status`` onto the OpenAI vocabulary.
The engine's poll loop breaks on "completed" and hard-fails on
"failed"/"expired"/"cancelled"; anything else keeps polling. Anthropic
batches only end as "ended" (per-request failures surface in the
results, mirroring OpenAI's "completed"-with-errors semantics), so
"ended" maps to "completed" and the non-terminal states pass through.
"""
return "completed" if processing_status == "ended" else processing_status
def _translate_batch_body(self, body: dict[str, Any]) -> dict[str, Any]:
"""Translate one OpenAI-shaped request body into Messages API params.
Mirrors the conversion rules of ``call()``: system messages fold into
the ``system`` param; ``max_completion_tokens`` becomes ``max_tokens``
(default 4096); ``temperature`` is dropped (the sync path never sends
it either current Claude models reject non-default sampling params);
an OpenAI ``response_format`` json_schema becomes a single forced
tool_use tool when strict (native constrained decoding, issue #1002),
else the schema is injected into the system prompt.
The system prompt carries the same cache_control marker as the sync
one-shot path (its sole breakpoint): every request in a retain batch
shares the fact-extraction system prompt, so the first item's cache
write serves the remaining items as best-effort reads and the
cache-read discount stacks with the 50% batch discount.
"""
system_prompt: str | None = None
messages: list[dict[str, Any]] = []
for msg in body.get("messages", []):
role = msg.get("role", "user")
content = msg.get("content", "")
if role == "system":
system_prompt = (system_prompt + "\n\n" + content) if system_prompt else content
else:
messages.append({"role": role, "content": content})
params: dict[str, Any] = {
"model": body.get("model") or self.model,
"messages": messages,
"max_tokens": body.get("max_completion_tokens") or 4096,
}
json_schema = (body.get("response_format") or {}).get("json_schema") or {}
schema = json_schema.get("schema")
if schema is not None:
if json_schema.get("strict"):
params["tools"] = [
{
"name": self._BATCH_TOOL_NAME,
"description": "Return the structured response.",
"input_schema": schema,
}
]
params["tool_choice"] = {"type": "tool", "name": self._BATCH_TOOL_NAME}
else:
schema_msg = "\n\nYou must respond with valid JSON matching this schema:\n" + json.dumps(
schema, indent=2, ensure_ascii=False
)
system_prompt = (system_prompt + schema_msg) if system_prompt else schema_msg
if system_prompt:
params["system"] = _cached_system_blocks(system_prompt)
# Batch params ARE the raw Messages body, so operator-configured extra
# body params merge directly (the sync path routes them through the
# SDK's extra_body, which does the same merge server-side).
if self._extra_body:
params.update(self._extra_body)
return params
def _translate_batch_message(self, message: Any) -> dict[str, Any]:
"""Render an Anthropic Message as the OpenAI response body the engine parses.
The engine reads ``choices[0].message.content`` (json.loads'ing it when
a schema was requested) and sums ``usage`` under the OpenAI key names.
Forced-tool responses carry their JSON in the tool_use block's input,
so that is re-serialized as the content string.
"""
content = ""
tool_input = None
for block in message.content:
if block.type == "tool_use" and block.name == self._BATCH_TOOL_NAME:
tool_input = block.input or {}
elif block.type == "text":
content += block.text
if tool_input is not None:
content = json.dumps(tool_input, ensure_ascii=False)
usage = getattr(message, "usage", None)
input_tokens = (usage.input_tokens or 0) if usage else 0
output_tokens = (usage.output_tokens or 0) if usage else 0
return {
"choices": [
{
"message": {"role": "assistant", "content": content},
"finish_reason": getattr(message, "stop_reason", None),
}
],
"usage": {
"prompt_tokens": input_tokens,
"completion_tokens": output_tokens,
"total_tokens": input_tokens + output_tokens,
},
}
async def submit_batch(
self,
requests: list[dict[str, Any]],
endpoint: str = "/v1/chat/completions",
completion_window: str = "24h",
) -> dict[str, Any]:
"""Submit a batch of requests to the Message Batches API.
Accepts the engine's OpenAI-JSONL-shaped entries. ``endpoint`` and
``completion_window`` belong to that shared shape and have no Anthropic
equivalent (batches always resolve within 24 hours); both are ignored.
"""
batch_requests = [
{
"custom_id": req["custom_id"],
"params": self._translate_batch_body(req.get("body") or {}),
}
for req in requests
]
logger.info(f"Submitting Anthropic message batch with {len(batch_requests)} requests")
batch = await self._client.messages.batches.create(requests=batch_requests)
logger.info(f"Anthropic batch submitted: {batch.id}, status={batch.processing_status}")
return {
"batch_id": batch.id,
"status": self._map_batch_status(batch.processing_status),
"created_at": batch.created_at,
"request_count": len(batch_requests),
}
async def get_batch_status(self, batch_id: str) -> dict[str, Any]:
"""Get batch status in the shape the engine's poll loop expects."""
batch = await self._client.messages.batches.retrieve(batch_id)
counts = batch.request_counts
processing = getattr(counts, "processing", 0) or 0
succeeded = getattr(counts, "succeeded", 0) or 0
errored = getattr(counts, "errored", 0) or 0
canceled = getattr(counts, "canceled", 0) or 0
expired = getattr(counts, "expired", 0) or 0
resolved = succeeded + errored + canceled + expired
result: dict[str, Any] = {
"batch_id": batch.id,
"status": self._map_batch_status(batch.processing_status),
"created_at": batch.created_at,
"request_counts": {
"total": processing + resolved,
"completed": resolved,
"failed": errored,
},
}
ended_at = getattr(batch, "ended_at", None)
if ended_at:
result["completed_at"] = ended_at
return result
async def retrieve_batch_results(self, batch_id: str) -> list[dict[str, Any]]:
"""Retrieve completed batch results, translated to the OpenAI shape.
Succeeded entries become ``{"custom_id", "response": {"body": ...}}``;
errored/canceled/expired entries become ``{"custom_id", "error": ...}``
so the engine's per-result error handling applies unchanged.
"""
batch = await self._client.messages.batches.retrieve(batch_id)
if batch.processing_status != "ended":
raise ValueError(f"Batch {batch_id} is not completed yet (status: {batch.processing_status})")
decoder = await self._client.messages.batches.results(batch_id)
results: list[dict[str, Any]] = []
async for entry in decoder:
outcome = entry.result
if outcome.type == "succeeded":
results.append(
{
"custom_id": entry.custom_id,
"response": {"body": self._translate_batch_message(outcome.message)},
}
)
else:
error = getattr(outcome, "error", None)
if error is not None:
detail = f"{getattr(error, 'type', 'error')}: {getattr(error, 'message', error)}"
else:
detail = f"batch request {outcome.type}"
results.append({"custom_id": entry.custom_id, "error": detail})
logger.info(f"Retrieved {len(results)} results for Anthropic batch {batch_id}")
return results
async def cleanup(self) -> None:
"""Clean up resources (close Anthropic client connections)."""
if hasattr(self, "_client") and self._client:
await self._client.close()
def supports_attempt_scoped_concurrency(self) -> bool:
return True
@@ -11,16 +11,14 @@ import json
import logging
import tempfile
import time
from contextlib import AbstractAsyncContextManager, nullcontext
from typing import Any, Callable
from typing import Any
from pydantic import ValidationError
from hindsight_api.engine.llm_interface import LLM_TOOL_CHOICE_AUTO, LLMInterface, LLMToolChoice, LLMToolChoiceMode
from hindsight_api.engine.llm_interface import LLMInterface
from hindsight_api.engine.llm_trace import LLMResponseUsage, stash_response_usage
from hindsight_api.engine.response_models import LLMToolCall, LLMToolCallResult, TokenUsage
from hindsight_api.metrics import get_metrics_collector
from hindsight_api.worker.stage import set_stage
logger = logging.getLogger(__name__)
@@ -51,20 +49,6 @@ def _get_isolated_claude_env() -> dict[str, str]:
return _isolated_claude_env
def _result_error_detail(message: Any) -> str:
"""Build an actionable error string from an ``is_error`` ResultMessage.
The CLI can report a failure with ``is_error=True`` while ``subtype``
still reads ``"success"``, putting the real detail in ``result`` (e.g.
quota exhaustion: ``You've hit your weekly limit · resets ...`` with
``api_error_status: 429``). The SDK's own fallback exception surfaces
only the subtype, producing the misleading "Claude Code returned an
error result: success" (issue #2702) — so prefer ``result``.
"""
detail = (message.result or "").strip() or message.subtype or "unknown error"
return f"Claude Code reported an error: {detail}"
class ClaudeCodeLLM(LLMInterface):
"""
LLM provider using Claude Code authentication.
@@ -164,7 +148,6 @@ class ClaudeCodeLLM(LLMInterface):
skip_validation: bool = False,
strict_schema: bool = False,
return_usage: bool = False,
attempt_context: Callable[[], AbstractAsyncContextManager[None]] | None = None,
) -> Any:
"""
Make an LLM API call with retry logic.
@@ -193,7 +176,6 @@ class ClaudeCodeLLM(LLMInterface):
from claude_agent_sdk import ( # type: ignore[unresolved-import]
AssistantMessage,
ClaudeAgentOptions,
ResultMessage,
TextBlock,
query,
)
@@ -227,27 +209,10 @@ class ClaudeCodeLLM(LLMInterface):
user_content += schema_instruction
# Configure SDK options
#
# tools=[] is required here for the same reason call_with_tools() below
# already sets it: with `tools` left at its default (None -> full
# "claude_code" built-in preset), allowed_tools=[] alone does not stop
# the CLI from loading the full built-in toolset and deferring into
# ToolSearch before answering, which burns the single max_turns=1
# budget on a tool-deferral step instead of a text response. Without
# this, single-turn calls intermittently fail with "Reached maximum
# number of turns (1)" even though the prompt itself needs no tools.
options = ClaudeAgentOptions(
system_prompt=system_prompt if system_prompt else None,
max_turns=1, # Single-turn for API-style interactions
tools=[], # Disable built-in tools so nothing forces a ToolSearch deferral
allowed_tools=[], # Disable tools for standard LLM calls
# Pin the configured model (issue #2881). Without this the spawned CLI
# runs its own default model — an Opus-class model on Pro/Max OAuth —
# regardless of HINDSIGHT_API_*_LLM_MODEL, while metrics/logs still print
# self.model, so the mismatch is invisible. The isolated CLAUDE_CONFIG_DIR
# (fresh temp dir) means a host settings.json can't reach the CLI either,
# so passing it through here is the only channel.
model=self.model or None,
env=_get_isolated_claude_env(),
)
@@ -258,18 +223,11 @@ class ClaudeCodeLLM(LLMInterface):
# Collect streaming response
full_text = ""
async with attempt_context() if attempt_context is not None else nullcontext():
set_stage(f"llm.claude_code.{scope}.attempt={attempt + 1}/{max_retries + 1}")
async for message in query(prompt=user_content, options=options):
if isinstance(message, AssistantMessage):
for block in message.content:
if isinstance(block, TextBlock):
full_text += block.text
elif isinstance(message, ResultMessage) and message.is_error:
# Surface the CLI's actual error text (e.g. quota
# exhaustion) instead of the SDK's subtype-based
# fallback exception (issue #2702).
raise RuntimeError(_result_error_detail(message))
async for message in query(prompt=user_content, options=options):
if isinstance(message, AssistantMessage):
for block in message.content:
if isinstance(block, TextBlock):
full_text += block.text
# The Claude Agent SDK doesn't report exact counts; stash the same
# char/4 estimate the success path traces so a later parse/validate
@@ -329,7 +287,7 @@ class ClaudeCodeLLM(LLMInterface):
# Record trace span
try:
from hindsight_api.tracing import _serialize_for_span, get_span_recorder
from hindsight_api.tracing import get_span_recorder
span_recorder = get_span_recorder()
span_recorder.record_llm_call(
@@ -337,17 +295,15 @@ class ClaudeCodeLLM(LLMInterface):
model=self.model,
scope=scope,
messages=messages,
response_content=_serialize_for_span(result),
response_content=result if isinstance(result, str) else result.model_dump_json(),
input_tokens=estimated_input,
output_tokens=estimated_output,
duration=duration,
finish_reason=None,
error=None,
)
except Exception as span_error:
# Tracing must remain best-effort, but expose instrumentation
# bugs that would otherwise silently erase spans (#3025).
logger.debug("Claude Code span recording failed: %s", span_error, exc_info=True)
except Exception:
pass # logging failure must never affect the operation
# Log slow calls
if duration > 10.0:
@@ -406,8 +362,7 @@ class ClaudeCodeLLM(LLMInterface):
max_retries: int = 5,
initial_backoff: float = 1.0,
max_backoff: float = 30.0,
tool_choice: LLMToolChoice = LLM_TOOL_CHOICE_AUTO,
attempt_context: Callable[[], AbstractAsyncContextManager[None]] | None = None,
tool_choice: str | dict[str, Any] = "auto",
) -> LLMToolCallResult:
"""
Make an LLM API call with tool/function calling support using Claude Agent SDK.
@@ -425,7 +380,7 @@ class ClaudeCodeLLM(LLMInterface):
max_retries: Maximum retry attempts.
initial_backoff: Initial backoff time in seconds.
max_backoff: Maximum backoff time in seconds.
tool_choice: Canonical tool-selection policy.
tool_choice: How to choose tools - "auto", "none", "required", or specific function dict.
- "auto": Model decides whether to call tools (default)
- "required": Model must call at least one tool
- "none": Model must not call any tools
@@ -438,7 +393,6 @@ class ClaudeCodeLLM(LLMInterface):
AssistantMessage,
ClaudeAgentOptions,
ClaudeSDKClient,
ResultMessage,
SdkMcpTool,
TextBlock,
ToolUseBlock,
@@ -519,27 +473,30 @@ class ClaudeCodeLLM(LLMInterface):
mcp_servers_config = {"hindsight_tools": mcp_server} if sdk_tools else {}
# Process tool_choice
if tool_choice.mode is LLMToolChoiceMode.NAMED:
if isinstance(tool_choice, dict) and tool_choice.get("type") == "function":
# Force a specific tool: filter allowed_tools to only that tool and add instruction
forced_name = tool_choice.selected_function_name
forced_tool_mcp_name = f"mcp__hindsight_tools__{forced_name}"
if forced_tool_mcp_name in allowed_tool_names:
allowed_tool_names = [forced_tool_mcp_name]
force_instruction = (
f"\n\nIMPORTANT: You MUST call the '{forced_name}' tool. Do not respond with text only."
)
system_prompt += force_instruction
logger.debug(f"Claude Code: Forcing tool call to '{forced_name}'")
else:
logger.warning(f"Claude Code: Forced tool '{forced_name}' not found in available tools")
elif tool_choice.mode is LLMToolChoiceMode.REQUIRED:
forced_name = tool_choice.get("function", {}).get("name")
if forced_name:
# Filter to only the forced tool (with MCP prefix)
forced_tool_mcp_name = f"mcp__hindsight_tools__{forced_name}"
if forced_tool_mcp_name in allowed_tool_names:
allowed_tool_names = [forced_tool_mcp_name]
# Add strong instruction to system prompt
force_instruction = (
f"\n\nIMPORTANT: You MUST call the '{forced_name}' tool. Do not respond with text only."
)
system_prompt += force_instruction
logger.debug(f"Claude Code: Forcing tool call to '{forced_name}'")
else:
logger.warning(f"Claude Code: Forced tool '{forced_name}' not found in available tools")
elif tool_choice == "required":
# Must call at least one tool
tool_instruction = (
"\n\nIMPORTANT: You MUST call at least one of the available tools. Do not respond with text only."
)
system_prompt += tool_instruction
logger.debug("Claude Code: Tool call required")
elif tool_choice.mode is LLMToolChoiceMode.NONE:
elif tool_choice == "none":
# No tools should be called - disable all tools
allowed_tool_names = []
mcp_servers_config = {}
@@ -547,33 +504,16 @@ class ClaudeCodeLLM(LLMInterface):
# else: tool_choice == "auto" or unspecified - use default behavior (no changes needed)
# Configure SDK options with MCP server
#
# tools=[] disables built-in CLI tools (Read, Write, Bash, ToolSearch, etc.)
# Without this, Claude Code CLI defers MCP tools when too many built-in tools
# are loaded, forcing Claude to use ToolSearch first — which wastes the turn
# are loaded, forcing Claude to use ToolSearch first — which wastes the max_turns
# budget and prevents direct MCP tool calls.
#
# max_turns=1 is critical (issue #2966). call_with_tools() is one *round* of
# an agentic loop the caller drives: the model proposes tool calls, we return
# them, and the orchestrator (reflect/agent.py) executes the REAL tools and
# feeds the results back on the next call. The SDK, however, runs its own
# in-process loop: it invokes our SDK MCP handlers — which are deliberate
# placeholders returning "[Tool <name> called successfully]" (no real data) —
# and lets the model react. With max_turns >= 2 the model calls recall, sees
# the empty placeholder, re-queries with reworded searches, exhausts the turn
# budget, and the run ends in error_max_turns with its tool calls discarded —
# exactly the "0 tool calls / no information" failure in #2966. Capping at a
# single turn stops the SDK from acting on the placeholder results: the model
# emits its first tool call (or a text answer) and we return that to the caller
# unchanged, matching how every other provider's call_with_tools() behaves.
options = ClaudeAgentOptions(
system_prompt=system_prompt if system_prompt else None,
tools=[], # Disable built-in tools so MCP tools load eagerly
max_turns=1, # One round: propose tool calls (or answer); caller drives the loop
max_turns=2, # Allow tool call + tool result round-trip
mcp_servers=mcp_servers_config,
allowed_tools=allowed_tool_names,
# Pin the configured model (issue #2881) — see the call() options block.
model=self.model or None,
env=_get_isolated_claude_env(),
)
@@ -584,49 +524,32 @@ class ClaudeCodeLLM(LLMInterface):
full_text = ""
tool_calls: list[LLMToolCall] = []
async with attempt_context() if attempt_context is not None else nullcontext():
set_stage(f"llm.claude_code.tools.attempt={attempt + 1}/{max_retries + 1}")
# Use ClaudeSDKClient for tool calling support
# Note: query() does NOT support custom tools, only ClaudeSDKClient does
async with ClaudeSDKClient(options=options) as client:
# Send the query
await client.query(user_content)
# Use ClaudeSDKClient for tool calling support
# Note: query() does NOT support custom tools, only ClaudeSDKClient does
async with ClaudeSDKClient(options=options) as client:
# Send the query
await client.query(user_content)
# Receive response
async for message in client.receive_response():
if isinstance(message, AssistantMessage):
for block in message.content:
if isinstance(block, TextBlock):
full_text += block.text
elif isinstance(block, ToolUseBlock):
# SDK returns tool names with MCP prefix (mcp__hindsight_tools__{name})
# Strip the prefix to return original tool name expected by caller
tool_name = block.name
if tool_name.startswith("mcp__hindsight_tools__"):
tool_name = tool_name.replace("mcp__hindsight_tools__", "", 1)
# Receive response
async for message in client.receive_response():
if isinstance(message, AssistantMessage):
for block in message.content:
if isinstance(block, TextBlock):
full_text += block.text
elif isinstance(block, ToolUseBlock):
# SDK returns tool names with MCP prefix (mcp__hindsight_tools__{name})
# Strip the prefix to return original tool name expected by caller
tool_name = block.name
if tool_name.startswith("mcp__hindsight_tools__"):
tool_name = tool_name.replace("mcp__hindsight_tools__", "", 1)
tool_calls.append(
LLMToolCall(
id=block.id,
name=tool_name,
arguments=block.input,
)
tool_calls.append(
LLMToolCall(
id=block.id,
name=tool_name,
arguments=block.input,
)
if tool_calls:
# This round proposed tool call(s). Stop consuming the
# stream so the SDK does not run another turn against our
# placeholder handlers (issue #2966) — the caller executes
# the real tools and calls us again with the results.
break
elif isinstance(message, ResultMessage) and message.is_error:
# With max_turns=1 the CLI reports error_max_turns whenever
# the model spent its single turn issuing a tool call (there
# was no follow-up turn to emit final text). That is expected
# here and not a failure: we already captured the tool call
# above and break before reaching this branch. Only a genuine
# error with nothing to return should surface (issue #2702).
if not tool_calls:
raise RuntimeError(_result_error_detail(message))
)
# Record metrics
duration = time.time() - start_time
@@ -688,6 +611,3 @@ class ClaudeCodeLLM(LLMInterface):
async def cleanup(self) -> None:
"""Clean up resources (no HTTP client to close for Claude Agent SDK)."""
pass
def supports_attempt_scoped_concurrency(self) -> bool:
return True
@@ -19,7 +19,6 @@ from __future__ import annotations
import base64
import binascii
import contextlib
import json
import logging
import os
@@ -32,11 +31,6 @@ from typing import Any
import httpx
try:
import fcntl
except ImportError: # pragma: no cover - Windows
fcntl = None # type: ignore[assignment]
logger = logging.getLogger(__name__)
@@ -64,9 +58,6 @@ _CODEX_TOKEN_REFRESH_SKEW_SECONDS = 60
_CODEX_TERMINAL_REFRESH_ERROR_CODES = frozenset(
{"refresh_token_expired", "refresh_token_reused", "refresh_token_invalidated"}
)
_CODEX_AUTH_LOCK_TIMEOUT_SECONDS = 20.0
_CODEX_AUTH_LOCKS_GUARD = threading.Lock()
_CODEX_AUTH_LOCKS: dict[Path, threading.Lock] = {}
def default_codex_auth_file() -> Path:
@@ -85,44 +76,6 @@ def default_codex_auth_file() -> Path:
return Path.home() / ".codex" / "auth.json"
def _path_scoped_lock(auth_file: Path) -> threading.Lock:
key = auth_file.expanduser().resolve(strict=False)
with _CODEX_AUTH_LOCKS_GUARD:
lock = _CODEX_AUTH_LOCKS.get(key)
if lock is None:
lock = threading.Lock()
_CODEX_AUTH_LOCKS[key] = lock
return lock
@contextlib.contextmanager
def _codex_auth_lock(auth_file: Path, timeout_seconds: float = _CODEX_AUTH_LOCK_TIMEOUT_SECONDS):
"""Cross-process advisory lock for one Codex auth store."""
with _path_scoped_lock(auth_file):
if fcntl is None: # pragma: no cover - Windows
logger.debug("fcntl unavailable; Codex refresh proceeds without a cross-process lock.")
yield
return
lock_path = auth_file.with_suffix(".lock")
lock_path.parent.mkdir(parents=True, exist_ok=True)
with open(lock_path, "a+") as lock_file:
deadline = time.monotonic() + max(1.0, timeout_seconds)
while True:
try:
fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
break
except (BlockingIOError, OSError):
if time.monotonic() >= deadline:
raise TimeoutError("Timed out waiting for the Codex auth store lock") from None
time.sleep(0.05)
try:
yield
finally:
with contextlib.suppress(OSError):
fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN)
class CodexRefreshExpiredError(RuntimeError):
"""Raised when the Codex refresh_token itself is no longer valid.
@@ -239,34 +192,6 @@ class CodexAuthManager:
return None
return data.get("tokens", {}).get("refresh_token")
@staticmethod
def _load_tokens_from_file(auth_file: Path) -> dict[str, Any] | None:
try:
with open(auth_file) as f:
data = json.load(f)
except (OSError, json.JSONDecodeError):
return None
tokens = data.get("tokens")
return tokens if isinstance(tokens, dict) else None
def _adopt_tokens(self, tokens: dict[str, Any]) -> bool:
"""Adopt a newer on-disk Codex token set if present."""
access_token = tokens.get("access_token")
refresh_token = tokens.get("refresh_token")
account_id = tokens.get("account_id")
changed = False
if isinstance(access_token, str) and access_token and access_token != self.access_token:
self.access_token = access_token
changed = True
if isinstance(refresh_token, str) and refresh_token and refresh_token != self.refresh_token:
self.refresh_token = refresh_token
changed = True
if isinstance(account_id, str) and account_id and account_id != self.account_id:
self.account_id = account_id
changed = True
return changed
@staticmethod
def _decode_jwt_exp_unixtime(token: str) -> int | None:
"""Return the JWT ``exp`` claim as a unix timestamp, or None on parse failure.
@@ -300,11 +225,6 @@ class CodexAuthManager:
return False
return exp <= int(time.time()) + skew_seconds
def _token_is_fresh_with_known_expiry(self, skew_seconds: int = _CODEX_TOKEN_REFRESH_SKEW_SECONDS) -> bool:
"""True only when the cached token has a known expiry outside the skew window."""
exp = self._decode_jwt_exp_unixtime(self.access_token)
return exp is not None and exp > int(time.time()) + skew_seconds
# ------------------------------------------------------------------
# Persistence
# ------------------------------------------------------------------
@@ -419,93 +339,78 @@ class CodexAuthManager:
if not self._token_is_stale():
return
with _codex_auth_lock(self._auth_file):
disk_tokens = self._load_tokens_from_file(self._auth_file)
if disk_tokens and self._adopt_tokens(disk_tokens):
if force or self._token_is_fresh_with_known_expiry():
return
if not self.refresh_token:
raise RuntimeError(
"Codex access_token is expired but no refresh_token is available. "
"Run 'codex auth login' to re-authenticate."
)
if not self.refresh_token:
raise RuntimeError(
"Codex access_token is expired but no refresh_token is available. "
log_reason = f" ({reason})" if reason else ""
logger.info(f"Refreshing Codex OAuth access_token{log_reason}")
request_body = {
"client_id": _CODEX_CLIENT_ID,
"grant_type": "refresh_token",
"refresh_token": self.refresh_token,
}
try:
response = self._http_client.post(
_CODEX_REFRESH_TOKEN_URL,
json=request_body,
headers={"Content-Type": "application/json"},
timeout=30.0,
)
except httpx.RequestError as e:
raise RuntimeError(f"Codex OAuth refresh network error: {type(e).__name__}") from e
if response.status_code == 401:
error_code = self._extract_oauth_error_code(response)
if error_code in _CODEX_TERMINAL_REFRESH_ERROR_CODES:
raise CodexRefreshExpiredError(
f"Codex refresh_token is permanently invalid (error.code={error_code}). "
"Run 'codex auth login' to re-authenticate."
)
raise CodexRefreshExpiredError(
f"Codex OAuth refresh returned 401 with unrecognized error code "
f"({error_code or 'none'}). Run 'codex auth login' to re-authenticate."
)
log_reason = f" ({reason})" if reason else ""
logger.info(f"Refreshing Codex OAuth access_token{log_reason}")
if response.status_code >= 400:
raise RuntimeError(f"Codex OAuth refresh failed with HTTP {response.status_code}")
request_access_token = self.access_token
request_refresh_token = self.refresh_token
request_body = {
"client_id": _CODEX_CLIENT_ID,
"grant_type": "refresh_token",
"refresh_token": request_refresh_token,
}
try:
response = self._http_client.post(
_CODEX_REFRESH_TOKEN_URL,
json=request_body,
headers={"Content-Type": "application/json"},
timeout=30.0,
)
except httpx.RequestError as e:
raise RuntimeError(f"Codex OAuth refresh network error: {type(e).__name__}") from e
try:
body = response.json()
except json.JSONDecodeError as e:
raise RuntimeError(f"Codex OAuth refresh returned non-JSON body: {e}") from e
if response.status_code == 401:
error_code = self._extract_oauth_error_code(response)
disk_tokens = self._load_tokens_from_file(self._auth_file)
if disk_tokens and (
disk_tokens.get("access_token") != request_access_token
or disk_tokens.get("refresh_token") != request_refresh_token
):
self._adopt_tokens(disk_tokens)
return
if error_code in _CODEX_TERMINAL_REFRESH_ERROR_CODES:
raise CodexRefreshExpiredError(
f"Codex refresh_token is permanently invalid (error.code={error_code}). "
"Run 'codex auth login' to re-authenticate."
)
raise CodexRefreshExpiredError(
f"Codex OAuth refresh returned 401 with unrecognized error code "
f"({error_code or 'none'}). Run 'codex auth login' to re-authenticate."
)
new_access = body.get("access_token")
if not new_access:
raise RuntimeError("Codex OAuth refresh returned no access_token")
if response.status_code >= 400:
raise RuntimeError(f"Codex OAuth refresh failed with HTTP {response.status_code}")
new_refresh = body.get("refresh_token") or self.refresh_token
new_id_token = body.get("id_token")
try:
body = response.json()
except json.JSONDecodeError as e:
raise RuntimeError(f"Codex OAuth refresh returned non-JSON body: {e}") from e
# Update in-memory state first so waiters see fresh credentials
# immediately, even if disk write fails.
self.access_token = new_access
self.refresh_token = new_refresh
new_access = body.get("access_token")
if not new_access:
raise RuntimeError("Codex OAuth refresh returned no access_token")
persisted: dict[str, Any] = {
"access_token": new_access,
"refresh_token": new_refresh,
}
if new_id_token:
persisted["id_token"] = new_id_token
new_refresh = body.get("refresh_token") or self.refresh_token
new_id_token = body.get("id_token")
try:
self._persist_auth_atomic(persisted)
except OSError as e:
logger.warning(
f"Codex OAuth refresh succeeded but persisting auth.json failed: {type(e).__name__}. "
"In-memory credentials are up to date; on-disk file is stale."
)
# Update in-memory state first so waiters see fresh credentials
# immediately, even if disk write fails.
self.access_token = new_access
self.refresh_token = new_refresh
persisted: dict[str, Any] = {
"access_token": new_access,
"refresh_token": new_refresh,
}
if new_id_token:
persisted["id_token"] = new_id_token
try:
self._persist_auth_atomic(persisted)
except OSError as e:
logger.warning(
f"Codex OAuth refresh succeeded but persisting auth.json failed: {type(e).__name__}. "
"In-memory credentials are up to date; on-disk file is stale."
)
logger.info("Codex OAuth access_token refreshed successfully")
logger.info("Codex OAuth access_token refreshed successfully")
def ensure_fresh_token(self) -> None:
"""Proactively refresh the access_token if it is near or past expiry.
@@ -20,19 +20,15 @@ import json
import logging
import time
import uuid
from contextlib import AbstractAsyncContextManager, nullcontext
from pathlib import Path
from typing import Any, Callable
from typing import Any
import httpx
from hindsight_api.engine.llm_interface import LLM_TOOL_CHOICE_AUTO, LLMInterface, LLMToolChoice, LLMToolChoiceMode
from hindsight_api.engine.llm_interface import LLMInterface
from hindsight_api.engine.llm_trace import LLMResponseUsage, stash_response_usage
from hindsight_api.engine.providers.llm_debug import dump_request_on_4xx
from hindsight_api.engine.response_models import LLMToolCall, LLMToolCallResult, TokenUsage
from hindsight_api.engine.structured_output import strict_json_schema
from hindsight_api.metrics import get_metrics_collector
from hindsight_api.worker.stage import set_stage
from .codex_auth import (
_CODEX_CLIENT_ID,
@@ -57,59 +53,6 @@ __all__ = [
logger = logging.getLogger(__name__)
# Newer Codex models are gated on the first-party client identity; the previous
# browser-shaped User-Agent returned "Model not found" for Luna (#2643).
# Use a neutral version because Hindsight must not claim a specific Codex release.
_CODEX_ORIGINATOR = "codex_cli_rs"
_CODEX_USER_AGENT = "codex_cli_rs/0.0.0 (Hindsight)"
# Name of the single forced function tool used to carry structured output when
# strict_schema is on. The Codex backend speaks the OpenAI Responses API, so a
# forced function call gives us constrained decoding straight into the response
# schema — no prompt-injected schema, no raw json.loads on free-form model text,
# no invalid-\escape retry storm (issue #2504, same class as #1002 / #2339).
_STRUCTURED_TOOL_NAME = "structured_response"
# Valid JSON string escape characters (the char that may follow a backslash).
_VALID_JSON_ESCAPE_CHARS = set('"\\/bfnrtu')
def _repair_invalid_json_escapes(text: str) -> str:
"""Best-effort repair of invalid ``\\escape`` sequences in a JSON string.
Escape-heavy content (code, serial/CLI commands, Windows paths, regexes)
makes weaker models emit backslashes that aren't valid JSON escapes (e.g.
``\\d``, ``\\s``, ``C:\\Users``), so ``json.loads`` fails deterministically
and every retry re-fails the same way (issue #2504). This doubles any
backslash that isn't part of a valid escape so the payload parses. It is a
lenient fallback only the strict_schema forced-tool path is the real fix.
"""
result: list[str] = []
i = 0
n = len(text)
while i < n:
ch = text[i]
if ch == "\\" and i + 1 < n:
nxt = text[i + 1]
if nxt in _VALID_JSON_ESCAPE_CHARS:
# Preserve the valid escape (both chars) verbatim.
result.append(ch)
result.append(nxt)
i += 2
continue
# Invalid escape: escape the lone backslash so JSON parses.
result.append("\\\\")
i += 1
continue
if ch == "\\" and i + 1 == n:
# Trailing lone backslash — escape it.
result.append("\\\\")
i += 1
continue
result.append(ch)
i += 1
return "".join(result)
class CodexLLM(LLMInterface):
"""
@@ -174,8 +117,8 @@ class CodexLLM(LLMInterface):
if self.model.startswith("openai/"):
self.model = self.model[len("openai/") :]
# Reasoning summary controls presentation separately from the backend's
# reasoning effort, which is sent unchanged in each request payload.
# Map reasoning effort to Codex reasoning summary format
# Codex supports: "auto", "concise", "detailed"
self.reasoning_summary = self._map_reasoning_effort(reasoning_effort)
# HTTP client for SSE streaming
@@ -197,18 +140,6 @@ class CodexLLM(LLMInterface):
def account_id(self) -> str:
return self._auth_manager.account_id
def _build_request_headers(self) -> httpx.Headers:
return httpx.Headers(
{
"Authorization": f"Bearer {self.access_token}",
"Content-Type": "application/json",
"OpenAI-Account-ID": self.account_id,
"User-Agent": _CODEX_USER_AGENT,
"Origin": "https://chatgpt.com",
"originator": _CODEX_ORIGINATOR,
}
)
@property
def refresh_token(self) -> str | None:
return self._auth_manager.refresh_token
@@ -345,6 +276,32 @@ class CodexLLM(LLMInterface):
}
return mapping.get(effort.lower(), "auto")
def _normalize_tool_choice(self, tool_choice: str | dict[str, Any]) -> str | dict[str, Any]:
"""Normalize forced function tool choice for the Codex Responses API.
Older agent paths may still pass OpenAI chat-completions style named
tool choice payloads such as:
{"type": "function", "function": {"name": "recall"}}
Codex Responses expects the named function at the top level instead:
{"type": "function", "name": "recall"}
"""
if not isinstance(tool_choice, dict):
return tool_choice
if str(tool_choice.get("type") or "").strip() != "function":
return tool_choice
function_payload = tool_choice.get("function")
if isinstance(function_payload, dict):
function_name = str(function_payload.get("name") or "").strip()
if function_name:
return {"type": "function", "name": function_name}
function_name = str(tool_choice.get("name") or "").strip()
if function_name:
return {"type": "function", "name": function_name}
return tool_choice
async def verify_connection(self) -> None:
"""Verify Codex connection by making a simple test call."""
try:
@@ -378,20 +335,8 @@ class CodexLLM(LLMInterface):
skip_validation: bool = False,
strict_schema: bool = False,
return_usage: bool = False,
attempt_context: Callable[[], AbstractAsyncContextManager[None]] | None = None,
) -> Any:
"""Make API call to Codex backend with SSE streaming.
Args:
strict_schema: Route structured output through a single forced
function tool (constrained decoding) instead of prompt-injecting
the schema and parsing free-form text. The Codex backend speaks
the OpenAI Responses API, so the forced function call emits the
response schema directly as tool arguments eliminating the
invalid-``\\escape`` retry storm (issue #2504). When False, falls
back to schema-in-prompt + JSON parse, now hardened with a lenient
invalid-escape repair before giving up.
"""
"""Make API call to Codex backend with SSE streaming."""
start_time = time.time()
# Proactively refresh the OAuth access_token if it's near expiry.
@@ -416,22 +361,11 @@ class CodexLLM(LLMInterface):
else:
user_messages.append(msg)
# Structured output: prefer a single forced function tool (constrained
# decoding) over text-injecting the schema and parsing the reply. The
# forced tool guarantees schema-shaped JSON in the tool arguments,
# eliminating the invalid-\escape retry storm (issue #2504). When
# strict_schema is off we keep the schema-in-prompt + json.loads
# fallback (now hardened with a lenient escape repair) for callers that
# can't force tools.
schema = None
use_forced_tool = False
# Add JSON schema instruction if response_format is provided
if response_format is not None and hasattr(response_format, "model_json_schema"):
schema = strict_json_schema(response_format) if strict_schema else response_format.model_json_schema()
if strict_schema:
use_forced_tool = True
else:
schema_msg = f"\n\nYou must respond with valid JSON matching this schema:\n{json.dumps(schema, indent=2, ensure_ascii=False)}"
system_instruction += schema_msg
schema = response_format.model_json_schema()
schema_msg = f"\n\nYou must respond with valid JSON matching this schema:\n{json.dumps(schema, indent=2, ensure_ascii=False)}"
system_instruction += schema_msg
# gpt-5.2-codex only supports "detailed" reasoning summary
reasoning_summary = "detailed" if "5.2" in self.model else self.reasoning_summary
@@ -451,28 +385,20 @@ class CodexLLM(LLMInterface):
"tools": [],
"tool_choice": "auto",
"parallel_tool_calls": True,
"reasoning": {"effort": self.reasoning_effort, "summary": reasoning_summary},
"reasoning": {"summary": reasoning_summary},
"store": False, # Codex uses stateless mode
"stream": True, # SSE streaming
"include": ["reasoning.encrypted_content"],
"prompt_cache_key": str(uuid.uuid4()),
}
if use_forced_tool and schema is not None:
# Single function tool whose parameters ARE the response schema;
# force it via tool_choice so the backend does constrained decoding.
payload["tools"] = [
{
"type": "function",
"name": _STRUCTURED_TOOL_NAME,
"description": "Return the structured response.",
"parameters": schema,
}
]
payload["tool_choice"] = {"type": "function", "name": _STRUCTURED_TOOL_NAME}
payload["parallel_tool_calls"] = False
headers = self._build_request_headers()
headers = {
"Authorization": f"Bearer {self.access_token}",
"Content-Type": "application/json",
"OpenAI-Account-ID": self.account_id,
"User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7)",
"Origin": "https://chatgpt.com",
}
url = f"{self.base_url}/codex/responses"
@@ -483,20 +409,11 @@ class CodexLLM(LLMInterface):
attempt = 0
while True:
try:
async with attempt_context() if attempt_context is not None else nullcontext():
set_stage(f"llm.codex.{scope}.attempt={attempt + 1}/{max_retries + 1}")
response = await self._client.post(url, json=payload, headers=headers, timeout=120.0)
response.raise_for_status()
response = await self._client.post(url, json=payload, headers=headers, timeout=120.0)
response.raise_for_status()
# Forced-tool path: read structured output from the function-call
# arguments (already a JSON string in a dedicated channel) rather
# than from free-form assistant text.
if use_forced_tool:
text_content, tool_calls = await self._parse_sse_tool_stream(response)
content = text_content or ""
else:
tool_calls = []
content = await self._parse_sse_stream(response)
# Parse SSE stream
content = await self._parse_sse_stream(response)
# Codex SSE carries no usage block; stash the same char/4 estimate
# the success path traces so a later parse/validate failure records
@@ -509,28 +426,7 @@ class CodexLLM(LLMInterface):
)
# Handle structured output
if use_forced_tool:
tool_input = None
for tc in tool_calls:
if tc.name == _STRUCTURED_TOOL_NAME:
tool_input = tc.arguments if isinstance(tc.arguments, dict) else None
break
if tool_input is None:
# Model ignored the forced tool (rare — e.g. a gateway that
# drops tool_choice). Retry so we don't hard-fail.
logger.warning(
f"Codex forced structured tool missing from response "
f"(attempt {attempt + 1}/{max_retries + 1})"
)
if attempt < max_retries:
backoff = min(initial_backoff * (2**attempt), max_backoff)
await asyncio.sleep(backoff)
attempt += 1
continue
raise RuntimeError("Codex did not return the forced structured_response tool call")
content = json.dumps(tool_input)
result = tool_input if skip_validation else response_format.model_validate(tool_input)
elif response_format is not None:
if response_format is not None:
# Models may wrap JSON in markdown
clean_content = content
if "```json" in content:
@@ -541,20 +437,13 @@ class CodexLLM(LLMInterface):
try:
json_data = json.loads(clean_content)
except json.JSONDecodeError as e:
# Escape-heavy content deterministically re-fails every
# retry (issue #2504). Try a lenient invalid-escape repair
# before burning a retry / re-raising.
try:
json_data = json.loads(_repair_invalid_json_escapes(clean_content))
logger.info("Codex JSON parsed after repairing invalid escape sequences")
except json.JSONDecodeError:
logger.warning(f"Codex JSON parse error (attempt {attempt + 1}/{max_retries + 1}): {e}")
if attempt < max_retries:
backoff = min(initial_backoff * (2**attempt), max_backoff)
await asyncio.sleep(backoff)
attempt += 1
continue
raise
logger.warning(f"Codex JSON parse error (attempt {attempt + 1}/{max_retries + 1}): {e}")
if attempt < max_retries:
backoff = min(initial_backoff * (2**attempt), max_backoff)
await asyncio.sleep(backoff)
attempt += 1
continue
raise
if skip_validation:
result = json_data
@@ -578,7 +467,7 @@ class CodexLLM(LLMInterface):
# Record trace span
try:
from hindsight_api.tracing import _serialize_for_span, get_span_recorder
from hindsight_api.tracing import get_span_recorder
# Estimate tokens for tracing
estimated_input = sum(len(m.get("content", "")) for m in messages) // 4
@@ -589,17 +478,15 @@ class CodexLLM(LLMInterface):
model=self.model,
scope=scope,
messages=messages,
response_content=_serialize_for_span(result),
response_content=result if isinstance(result, str) else result.model_dump_json(),
input_tokens=estimated_input,
output_tokens=estimated_output,
duration=duration,
finish_reason=None,
error=None,
)
except Exception as span_error:
# Tracing must remain best-effort, but expose instrumentation
# bugs that would otherwise silently erase spans (#3025).
logger.debug("Codex span recording failed: %s", span_error, exc_info=True)
except Exception:
pass # logging failure must never affect the operation
if return_usage:
# Codex doesn't provide token counts, estimate based on content
@@ -655,9 +542,6 @@ class CodexLLM(LLMInterface):
"Run 'codex auth login' to re-authenticate."
) from e
# Diagnostic dump (opt-in) of the exact request behind any 4xx.
dump_request_on_4xx(scope=scope, provider=self.provider, model=self.model, err=e, request=payload)
# Log the actual error message from the API
error_detail = e.response.text[:500] if hasattr(e.response, "text") else str(e)
@@ -753,8 +637,7 @@ class CodexLLM(LLMInterface):
max_retries: int = 5,
initial_backoff: float = 1.0,
max_backoff: float = 30.0,
tool_choice: LLMToolChoice = LLM_TOOL_CHOICE_AUTO,
attempt_context: Callable[[], AbstractAsyncContextManager[None]] | None = None,
tool_choice: str | dict[str, Any] = "auto",
) -> LLMToolCallResult:
"""
Make API call with tool calling support.
@@ -771,7 +654,7 @@ class CodexLLM(LLMInterface):
max_retries: Maximum retry attempts.
initial_backoff: Initial backoff time in seconds.
max_backoff: Maximum backoff time in seconds.
tool_choice: Canonical tool-selection policy.
tool_choice: How to choose tools - "auto", "none", "required", or a specific function.
Returns:
LLMToolCallResult with content and/or tool_calls.
@@ -833,20 +716,22 @@ class CodexLLM(LLMInterface):
"instructions": system_instruction,
"input": user_messages,
"tools": codex_tools,
"tool_choice": (
{"type": "function", "name": tool_choice.selected_function_name}
if tool_choice.mode is LLMToolChoiceMode.NAMED
else tool_choice.mode.value
),
"tool_choice": self._normalize_tool_choice(tool_choice),
"parallel_tool_calls": True,
"reasoning": {"effort": self.reasoning_effort, "summary": reasoning_summary},
"reasoning": {"summary": reasoning_summary},
"store": False,
"stream": True,
"include": ["reasoning.encrypted_content"],
"prompt_cache_key": str(uuid.uuid4()),
}
headers = self._build_request_headers()
headers = {
"Authorization": f"Bearer {self.access_token}",
"Content-Type": "application/json",
"OpenAI-Account-ID": self.account_id,
"User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7)",
"Origin": "https://chatgpt.com",
}
url = f"{self.base_url}/codex/responses"
@@ -859,28 +744,10 @@ class CodexLLM(LLMInterface):
# surfaces immediately to keep behavior identical for callers.
attempted_refresh_after_auth_error = False
async def _request_attempt(attempt: int) -> tuple[str | None, list[LLMToolCall]]:
async with attempt_context() if attempt_context is not None else nullcontext():
set_stage(f"llm.codex.tools.attempt={attempt}/2")
response = await self._client.post(url, json=payload, headers=headers, timeout=120.0)
if response.status_code != 200:
# 401/403 on the first attempt may still be recovered by the
# reactive token refresh below — don't log those as errors yet.
detail = f"Codex API error {response.status_code}: {response.text[:500]}"
if response.status_code in (401, 403) and not attempted_refresh_after_auth_error:
logger.warning(f"{detail} (will attempt token refresh)")
else:
logger.error(detail)
response.raise_for_status()
return await self._parse_sse_tool_stream(response)
try:
try:
content, tool_calls = await _request_attempt(1)
except httpx.HTTPStatusError as auth_error:
response = auth_error.response
if response.status_code not in (401, 403) or attempted_refresh_after_auth_error:
raise
response = await self._client.post(url, json=payload, headers=headers, timeout=120.0)
if response.status_code in (401, 403) and not attempted_refresh_after_auth_error:
attempted_refresh_after_auth_error = True
try:
await self._refresh_oauth_tokens(
@@ -889,7 +756,7 @@ class CodexLLM(LLMInterface):
)
headers["Authorization"] = f"Bearer {self.access_token}"
logger.info("Codex auth refreshed after auth error; retrying tool-call request once")
content, tool_calls = await _request_attempt(2)
response = await self._client.post(url, json=payload, headers=headers, timeout=120.0)
except CodexRefreshExpiredError as refresh_err:
logger.error(
"Codex refresh_token is permanently invalid; cannot recover from auth error in tool-call path"
@@ -902,7 +769,16 @@ class CodexLLM(LLMInterface):
logger.error(
f"Codex token refresh attempt failed in tool-call path: {type(refresh_err).__name__}: {refresh_err}"
)
raise auth_error
# Fall through to the normal error path below.
# Log response details on error
if response.status_code != 200:
logger.error(f"Codex API error {response.status_code}: {response.text[:500]}")
response.raise_for_status()
# Parse SSE for tool calls and content
content, tool_calls = await self._parse_sse_tool_stream(response)
duration = time.time() - start_time
metrics = get_metrics_collector()
@@ -952,8 +828,6 @@ class CodexLLM(LLMInterface):
)
except Exception as e:
# Diagnostic dump (opt-in) of the exact request behind any 4xx.
dump_request_on_4xx(scope=scope, provider=self.provider, model=self.model, err=e, request=payload)
logger.error(f"Codex tool call error: {e}")
raise
@@ -998,13 +872,8 @@ class CodexLLM(LLMInterface):
try:
arguments = json.loads(arguments_str)
except json.JSONDecodeError:
# Escape-heavy content can emit invalid \escape
# sequences (issue #2504); repair before giving up.
try:
arguments = json.loads(_repair_invalid_json_escapes(arguments_str))
except json.JSONDecodeError:
logger.warning(f"Failed to parse tool arguments: {arguments_str}")
arguments = {}
logger.warning(f"Failed to parse tool arguments: {arguments_str}")
arguments = {}
tool_calls.append(
LLMToolCall(
@@ -1023,6 +892,3 @@ class CodexLLM(LLMInterface):
"""Clean up HTTP clients."""
await self._client.aclose()
self._auth_manager.close()
def supports_attempt_scoped_concurrency(self) -> bool:
return True
@@ -56,14 +56,6 @@ _DEFAULT_REFRESH_MARGIN_SECONDS = 5 * 60
# to None and callers proceed uncached, rather than stalling the whole batch.
_DEFAULT_CREATE_TIMEOUT_SECONDS = 30.0
# TTL for the per-step reflect caches created by ``create_incremental``. These
# live only for the duration of one reflect (seconds), so the TTL is just a
# storage backstop in case the explicit ``delete_session`` at reflect end is
# missed (crash / event-loop teardown). Short so orphaned caches age out fast —
# storage is billed per token-hour, so a 5-minute cap keeps the cost of a leaked
# cache negligible.
_DEFAULT_INCREMENTAL_TTL_SECONDS = 5 * 60
@dataclass
class _CacheEntry:
@@ -100,10 +92,6 @@ class GeminiCacheManager:
self._create_timeout_seconds = create_timeout_seconds
self._entries: dict[str, _CacheEntry] = {}
self._lock = asyncio.Lock()
# session_id -> CachedContent names created via ``create_incremental``.
# A reflect creates a fresh rolling cache per step under one session id;
# ``delete_session`` tears them all down when the reflect finishes.
self._sessions: dict[str, list[str]] = {}
@staticmethod
def fingerprint(
@@ -241,99 +229,18 @@ class GeminiCacheManager:
if entry.name == name:
self._entries.pop(key, None)
async def create_incremental(
self,
*,
session_id: str,
model: str,
system_instruction: str,
contents: list[Any],
tools: list[dict[str, Any]] | None = None,
) -> str | None:
"""Create a fresh CachedContent holding ``system + tools + contents`` and
track it under ``session_id`` for later teardown.
Unlike ``get_or_create``, this does NOT deduplicate by fingerprint: each
step of a reflect grows the conversation prefix, so every call is a
distinct, single-use cache. The reflect loop creates one per step (each
covering the previous step's full input) and reuses it for exactly the
next model turn, then supersedes it. All caches for the session are
deleted by ``delete_session`` when the reflect ends; the short TTL is
only a backstop.
Returns the cache resource name, or ``None`` when caching is disabled,
the prefix is below the model minimum, or the create otherwise fails
callers MUST fall back to an uncached call in that case.
"""
try:
name = await self._create_cache(
model=model,
system_instruction=system_instruction,
tools=tools,
contents=contents,
ttl_seconds=_DEFAULT_INCREMENTAL_TTL_SECONDS,
)
except _CacheNotEligible as e:
logger.debug(
"GeminiCacheManager: incremental prefix not eligible (model=%s, reason=%s) — caller falls back",
model,
e,
)
return None
except Exception:
logger.exception(
"GeminiCacheManager: failed to create incremental cache (model=%s); caller falls back",
model,
)
return None
if name is not None:
self._sessions.setdefault(session_id, []).append(name)
return name
async def delete(self, name: str) -> None:
"""Best-effort server-side delete of a single CachedContent.
Swallows all errors: a failed delete just means the cache ages out on
its TTL. Also drops any matching in-process entry.
"""
self.invalidate(name)
try:
await self._client.aio.caches.delete(name=name)
except Exception:
logger.debug("GeminiCacheManager: delete of cache %s failed (will age out on TTL)", name, exc_info=True)
async def delete_session(self, session_id: str) -> None:
"""Delete every CachedContent created for ``session_id`` (reflect teardown).
Deletes concurrently and best-effort a reflect must never fail because
a cache couldn't be torn down; the short TTL is the backstop.
"""
names = self._sessions.pop(session_id, [])
if not names:
return
await asyncio.gather(*(self.delete(n) for n in names), return_exceptions=True)
async def _create_cache(
self,
*,
model: str,
system_instruction: str,
tools: list[dict[str, Any]] | None = None,
contents: list[Any] | None = None,
ttl_seconds: int | None = None,
) -> str | None:
"""Wrap ``client.aio.caches.create`` with the config we want.
The SDK surface differs slightly across google-genai versions;
this implementation targets the >=1.0.0 line where caches live
under ``client.aio.caches``.
``contents`` (already-converted ``genai_types.Content`` turns) is
appended after the system_instruction/tools so the cache can hold a
growing multi-turn conversation prefix, not just the static prefix
this is what the step-by-step reflect cache relies on. ``ttl_seconds``
overrides the manager default (used to give per-step reflect caches a
short backstop TTL).
"""
# Lazy import so this module doesn't require the SDK at import time.
from google.genai import types as genai_types
@@ -347,10 +254,8 @@ class GeminiCacheManager:
# still part of the fingerprint so a schema change keys a fresh cache.
config_kwargs: dict[str, Any] = {
"system_instruction": system_instruction,
"ttl": f"{ttl_seconds if ttl_seconds is not None else self._ttl_seconds}s",
"ttl": f"{self._ttl_seconds}s",
}
if contents:
config_kwargs["contents"] = contents
if tools:
# OpenAI-style {"function": {...}} entries must be converted to
# Gemini's Tool/FunctionDeclaration shape before caching.
@@ -12,19 +12,16 @@ import io
import json
import logging
import time
from contextlib import AbstractAsyncContextManager, nullcontext
from contextvars import ContextVar
from dataclasses import dataclass
from typing import Any, Callable
from typing import Any
from google import genai
from google.genai import errors as genai_errors
from google.genai import types as genai_types
from hindsight_api.engine.llm_interface import LLM_TOOL_CHOICE_AUTO, LLMInterface, LLMToolChoice, LLMToolChoiceMode
from hindsight_api.engine.llm_interface import LLMInterface
from hindsight_api.engine.llm_trace import LLMResponseUsage, stash_response_usage
from hindsight_api.engine.llm_wrapper import parse_llm_json
from hindsight_api.engine.providers.llm_debug import dump_request_on_4xx
from hindsight_api.engine.response_models import LLMToolCall, LLMToolCallResult, TokenUsage
from hindsight_api.metrics import get_metrics_collector
from hindsight_api.worker.stage import set_stage
@@ -66,99 +63,6 @@ def _usage_from_gemini_response(response: Any) -> LLMResponseUsage:
)
@dataclass(frozen=True)
class _GeminiConversation:
"""A message list converted to Gemini's request shape."""
system_instruction: str | None
contents: list["genai_types.Content"]
def _convert_messages_to_gemini(msg_list: list[dict[str, Any]]) -> _GeminiConversation:
"""Convert OpenAI-style messages to a Gemini (system_instruction, contents) pair.
Shared by ``call_with_tools`` (request body) and the incremental cache
builder so a cached prefix and the live request serialise turns identically
any drift would fingerprint differently and defeat the cache. Consecutive
``role="tool"`` messages are grouped into a single ``user`` Content with
multiple FunctionResponse parts, matching Gemini's multi-turn requirement.
"""
system_instruction: str | None = None
gemini_contents: list[genai_types.Content] = []
pending_tool_names_by_call_id: dict[str, str] = {}
i = 0
while i < len(msg_list):
msg = msg_list[i]
role = msg.get("role", "user")
content = msg.get("content", "")
if role != "tool" and pending_tool_names_by_call_id:
missing_ids = ", ".join(sorted(pending_tool_names_by_call_id))
raise ValueError(f"Gemini assistant tool calls require results before the next message: {missing_ids}")
if role == "system":
system_instruction = (system_instruction + "\n\n" + content) if system_instruction else content
i += 1
elif role == "tool":
parts = []
while i < len(msg_list) and msg_list[i].get("role") == "tool":
tool_msg = msg_list[i]
tool_content = tool_msg.get("content", "")
tool_call_id = tool_msg["tool_call_id"]
tool_name = pending_tool_names_by_call_id.pop(tool_call_id, None)
if tool_name is None:
raise ValueError(f"Gemini tool result references unknown tool_call_id {tool_call_id!r}")
parts.append(
genai_types.Part(
function_response=genai_types.FunctionResponse(
name=tool_name,
response={"result": tool_content},
)
)
)
i += 1
if pending_tool_names_by_call_id:
missing_ids = ", ".join(sorted(pending_tool_names_by_call_id))
raise ValueError(f"Gemini assistant tool calls are missing results: {missing_ids}")
gemini_contents.append(genai_types.Content(role="user", parts=parts))
elif role == "assistant":
tool_calls_in_msg = msg.get("tool_calls", [])
if tool_calls_in_msg:
parts = []
if content:
parts.append(genai_types.Part(text=content))
for tc in tool_calls_in_msg:
tool_call_id = tc["id"]
fn = tc["function"]
fn_name = fn["name"]
if tool_call_id in pending_tool_names_by_call_id:
raise ValueError(
f"Gemini assistant tool call id {tool_call_id!r} must be unique within its turn"
)
pending_tool_names_by_call_id[tool_call_id] = fn_name
fn_args_str = fn.get("arguments", "{}")
fn_args = parse_llm_json(fn_args_str)
thought_signature = tc.get("thought_signature")
fc_kwargs: dict[str, Any] = {"name": fn_name, "args": fn_args}
part_kwargs: dict[str, Any] = {"function_call": genai_types.FunctionCall(**fc_kwargs)}
if thought_signature:
part_kwargs["thought_signature"] = base64.b64decode(thought_signature)
parts.append(genai_types.Part(**part_kwargs))
gemini_contents.append(genai_types.Content(role="model", parts=parts))
else:
gemini_contents.append(genai_types.Content(role="model", parts=[genai_types.Part(text=content)]))
i += 1
else:
gemini_contents.append(genai_types.Content(role="user", parts=[genai_types.Part(text=content)]))
i += 1
if pending_tool_names_by_call_id:
missing_ids = ", ".join(sorted(pending_tool_names_by_call_id))
raise ValueError(f"Gemini assistant tool calls are missing results: {missing_ids}")
return _GeminiConversation(system_instruction=system_instruction, contents=gemini_contents)
class GeminiLLM(LLMInterface):
"""
LLM provider for Google Gemini and Vertex AI.
@@ -312,7 +216,6 @@ class GeminiLLM(LLMInterface):
strict_schema: bool = False,
return_usage: bool = False,
cached_prefix: str | None = None,
attempt_context: Callable[[], AbstractAsyncContextManager[None]] | None = None,
) -> Any:
"""
Make a Gemini/VertexAI API call with retry logic.
@@ -426,17 +329,17 @@ class GeminiLLM(LLMInterface):
last_exception = None
for attempt in range(max_retries + 1):
if attempt > 0:
set_stage(f"llm.gemini.{scope}.attempt={attempt + 1}/{max_retries + 1}")
try:
async with attempt_context() if attempt_context is not None else nullcontext():
set_stage(f"llm.gemini.{scope}.attempt={attempt + 1}/{max_retries + 1}")
response = await asyncio.wait_for(
self._client.aio.models.generate_content(
model=self.model,
contents=gemini_contents,
config=generation_config,
),
timeout=90.0, # Safety net for network hangs; valid slow responses are <90s
)
response = await asyncio.wait_for(
self._client.aio.models.generate_content(
model=self.model,
contents=gemini_contents,
config=generation_config,
),
timeout=90.0, # Safety net for network hangs; valid slow responses are <90s
)
# Stash usage before parse/validate, which may raise locally
# even though the provider charged for these tokens (#2387).
stash_response_usage(_usage_from_gemini_response(response))
@@ -576,17 +479,6 @@ class GeminiLLM(LLMInterface):
logger.error(f"Gemini auth error (HTTP {e.code}), not retrying: {str(e)}")
raise
# Diagnostic dump (opt-in) of the exact request behind any 4xx, captured
# before the cache-drop retry below rebuilds the config so we see what failed.
dump_request_on_4xx(
scope=scope,
provider=self.provider,
model=self.model,
err=e,
request=generation_config,
messages=gemini_contents,
)
# Cached-request safety net: a stale/invalid/expired CachedContent
# (or an incompatibility like cache + tool_config) surfaces as a 400.
# Retrying the same cached request can't recover, so on the first
@@ -633,10 +525,8 @@ class GeminiLLM(LLMInterface):
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 a Gemini/VertexAI API call with tool/function calling support.
@@ -650,22 +540,15 @@ class GeminiLLM(LLMInterface):
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.
tool_choice: How to choose tools (Gemini uses "auto" only).
cached_prefix: Optional CachedContent resource name (from
``GeminiCacheManager.get_or_create`` or ``create_incremental``).
When set, the system_instruction and tool definitions are assumed
``GeminiCacheManager.get_or_create`` with ``tools=...``). When
set, the system_instruction and tool definitions are assumed
to live in the cache; this call will skip resending them and
the cached prefix is billed at the cached-input rate. The
``tools`` argument is still required (the caller may pass
an empty list when the cache holds them) so existing call
sites don't break.
cached_prefix_message_count: Number of leading ``messages`` already
baked into ``cached_prefix`` (the step-by-step reflect cache holds
a growing conversation prefix, not just system+tools). Only the
messages AFTER this index are sent as request contents the rest
come from the cache and bill at the cached rate. 0 means the cache
holds only the static prefix (system+tools), so the full
conversation is still sent (legacy behaviour).
Returns:
LLMToolCallResult with content and/or tool_calls.
@@ -673,45 +556,86 @@ class GeminiLLM(LLMInterface):
start_time = time.time()
using_cache = cached_prefix is not None
# Convert tools to Gemini format. While the cache is in use the tool
# definitions live in the CachedContent and the SDK rejects re-sending
# them alongside ``cached_content`` (see ``_build_tools_config``), but we
# still build them unconditionally so the cached-call-failed fallback —
# which drops the cache and re-sends prefix + tools inline — has real
# tools to send rather than an empty list.
# Convert tools to Gemini format. When the cache is in use, the
# tool definitions are baked into the CachedContent at create time
# and the SDK rejects re-sending them alongside ``cached_content``.
gemini_tools = []
for tool in tools:
func = tool.get("function", {})
gemini_tools.append(
genai_types.Tool(
function_declarations=[
genai_types.FunctionDeclaration(
name=func.get("name", ""),
description=func.get("description", ""),
parameters=func.get("parameters"),
)
]
if not using_cache:
for tool in tools:
func = tool.get("function", {})
gemini_tools.append(
genai_types.Tool(
function_declarations=[
genai_types.FunctionDeclaration(
name=func.get("name", ""),
description=func.get("description", ""),
parameters=func.get("parameters"),
)
]
)
)
)
# Convert messages. ``system_instruction`` and the FULL contents are always
# computed: _build_tools_config omits system/tools from the request while
# the cache carries the prefix, but the cached-call-failed safety net must
# be able to re-send the whole prefix + tools inline.
converted = _convert_messages_to_gemini(list(messages))
system_instruction = converted.system_instruction
full_contents = converted.contents
# Convert messages
system_instruction = None
gemini_contents = []
msg_list = list(messages)
i = 0
while i < len(msg_list):
msg = msg_list[i]
role = msg.get("role", "user")
content = msg.get("content", "")
# Step-by-step reflect cache: when the cache already holds the first
# ``cached_prefix_message_count`` messages, send ONLY the newer turns as
# request contents — the cached prefix supplies the rest at the cached
# rate. The split is always at a whole-turn boundary (the reflect loop
# advances the cache one completed turn at a time), so slicing the raw
# messages before conversion never splits a grouped tool turn.
if using_cache and cached_prefix_message_count > 0:
delta_contents = _convert_messages_to_gemini(list(messages)[cached_prefix_message_count:]).contents
else:
delta_contents = full_contents
if role == "system":
# Always capture system_instruction. _build_tools_config omits it
# (and tools) from the request while the cache carries the prefix,
# but it must be available so the cached-call-failed safety net can
# re-send the prefix + tools inline.
system_instruction = (system_instruction + "\n\n" + content) if system_instruction else content
i += 1
elif role == "tool":
# Gemini requires ALL tool responses for a given model turn to be grouped
# into a single Content with multiple FunctionResponse parts.
# Consecutive role="tool" messages correspond to one model turn's tool calls.
parts = []
while i < len(msg_list) and msg_list[i].get("role") == "tool":
tool_msg = msg_list[i]
tool_content = tool_msg.get("content", "")
parts.append(
genai_types.Part(
function_response=genai_types.FunctionResponse(
name=tool_msg.get("name", ""),
response={"result": tool_content},
)
)
)
i += 1
gemini_contents.append(genai_types.Content(role="user", parts=parts))
elif role == "assistant":
tool_calls_in_msg = msg.get("tool_calls", [])
if tool_calls_in_msg:
# Convert OpenAI-style tool_calls to Gemini function_call parts
# This is required for proper multi-turn conversation history
parts = []
if content:
parts.append(genai_types.Part(text=content))
for tc in tool_calls_in_msg:
fn = tc.get("function", {})
fn_name = fn.get("name", "")
fn_args_str = fn.get("arguments", "{}")
fn_args = parse_llm_json(fn_args_str)
thought_signature = tc.get("thought_signature")
fc_kwargs: dict[str, Any] = {"name": fn_name, "args": fn_args}
part_kwargs: dict[str, Any] = {"function_call": genai_types.FunctionCall(**fc_kwargs)}
if thought_signature:
part_kwargs["thought_signature"] = base64.b64decode(thought_signature)
parts.append(genai_types.Part(**part_kwargs))
gemini_contents.append(genai_types.Content(role="model", parts=parts))
else:
gemini_contents.append(genai_types.Content(role="model", parts=[genai_types.Part(text=content)]))
i += 1
else:
gemini_contents.append(genai_types.Content(role="user", parts=[genai_types.Part(text=content)]))
i += 1
# Apply safety settings: context var (per-request bank override) takes precedence over instance default
effective_safety_settings = _safety_settings_ctx.get()
@@ -741,20 +665,22 @@ class GeminiLLM(LLMInterface):
config_kwargs["max_output_tokens"] = max_completion_tokens
# Map OpenAI-style tool_choice to Gemini FunctionCallingConfig
if tool_choice.mode is LLMToolChoiceMode.REQUIRED:
if tool_choice == "required":
config_kwargs["tool_config"] = genai_types.ToolConfig(
function_calling_config=genai_types.FunctionCallingConfig(
mode="ANY",
)
)
elif tool_choice.mode is LLMToolChoiceMode.NAMED:
config_kwargs["tool_config"] = genai_types.ToolConfig(
function_calling_config=genai_types.FunctionCallingConfig(
mode="ANY",
allowed_function_names=[tool_choice.selected_function_name],
elif isinstance(tool_choice, dict) and tool_choice.get("type") == "function":
fn_name = tool_choice.get("function", {}).get("name")
if fn_name:
config_kwargs["tool_config"] = genai_types.ToolConfig(
function_calling_config=genai_types.FunctionCallingConfig(
mode="ANY",
allowed_function_names=[fn_name],
)
)
)
elif tool_choice.mode is LLMToolChoiceMode.NONE:
elif tool_choice == "none":
config_kwargs["tool_config"] = genai_types.ToolConfig(
function_calling_config=genai_types.FunctionCallingConfig(mode="NONE")
)
@@ -772,21 +698,17 @@ class GeminiLLM(LLMInterface):
last_exception = None
for attempt in range(max_retries + 1):
if attempt > 0:
set_stage(f"llm.gemini.tools.attempt={attempt + 1}/{max_retries + 1}")
try:
# With the cache active, send only the un-cached tail (delta);
# on the uncached fallback path send the full conversation so the
# re-inlined system+tools prefix has its whole context.
active_contents = delta_contents if cache_active else full_contents
async with attempt_context() if attempt_context is not None else nullcontext():
set_stage(f"llm.gemini.tools.attempt={attempt + 1}/{max_retries + 1}")
response = await asyncio.wait_for(
self._client.aio.models.generate_content(
model=self.model,
contents=active_contents,
config=config,
),
timeout=90.0, # Safety net for network hangs; valid slow responses are <90s
)
response = await asyncio.wait_for(
self._client.aio.models.generate_content(
model=self.model,
contents=gemini_contents,
config=config,
),
timeout=90.0, # Safety net for network hangs; valid slow responses are <90s
)
stash_response_usage(_usage_from_gemini_response(response))
# Extract content and tool calls
@@ -885,17 +807,6 @@ class GeminiLLM(LLMInterface):
logger.error(f"Gemini auth error (HTTP {e.code}), not retrying: {str(e)}")
raise
# Diagnostic dump (opt-in) of the exact request behind any 4xx, captured
# before the cache-drop retry below rebuilds the config so we see what failed.
dump_request_on_4xx(
scope=scope,
provider=self.provider,
model=self.model,
err=e,
request=config,
messages=active_contents,
)
# Cached-request safety net (see ``call``): a stale/invalid cache or
# a cache+tool_config conflict surfaces as a 400. Drop the cache,
# invalidate it for later operations, and retry THIS call inline
@@ -972,56 +883,6 @@ class GeminiLLM(LLMInterface):
tools=tools,
)
# ── Step-by-step incremental prompt caching (reflect tool loop) ──────────
def supports_incremental_prompt_cache(self) -> bool:
"""True when explicit caching is on — the reflect loop can then roll a
per-step CachedContent that grows with the conversation."""
return self._prompt_cache_enabled
def _ensure_cache_manager(self) -> Any:
if self._cache_manager is None:
from hindsight_api.engine.providers.gemini_cache import GeminiCacheManager
self._cache_manager = GeminiCacheManager(self._client)
return self._cache_manager
async def create_incremental_cache(
self,
*,
session_id: str,
messages: list[dict[str, Any]],
tools: list[dict[str, Any]] | None = None,
) -> str | None:
"""Cache ``system + tools + messages`` as a conversation prefix and return
its resource name (or ``None`` caller falls back to an uncached call).
The reflect loop calls this once per step with the growing message list so
each step's cache entirely contains the previous step's input; the next
model turn then references it and re-sends only its own delta. Caches are
tracked under ``session_id`` and torn down by ``delete_cache_session``.
"""
if not self._prompt_cache_enabled or self._client is None:
return None
converted = _convert_messages_to_gemini(list(messages))
return await self._ensure_cache_manager().create_incremental(
session_id=session_id,
model=self.model,
system_instruction=converted.system_instruction or "",
contents=converted.contents,
tools=tools,
)
async def delete_cached_prefix(self, name: str) -> None:
"""Best-effort delete of a single CachedContent (superseded reflect step)."""
if self._cache_manager is not None:
await self._cache_manager.delete(name)
async def delete_cache_session(self, session_id: str) -> None:
"""Tear down every CachedContent created for a reflect session."""
if self._cache_manager is not None:
await self._cache_manager.delete_session(session_id)
# ── Batch API (Gemini API only — not Vertex AI) ─────────────────────────
#
# Google's Gemini Batch API gives a flat 50% discount on input + output
@@ -1169,7 +1030,7 @@ class GeminiLLM(LLMInterface):
Mirrors the synchronous ``call`` path: system messages become
``systemInstruction``; a ``response_format`` json_schema forces JSON
output (``responseMimeType``), appends the schema as a textual hint, and
grammar-enforces via ``responseJsonSchema`` whenever a schema is present.
grammar-enforces via ``responseJsonSchema`` when ``strict`` is set.
"""
system_texts: list[str] = []
contents: list[dict[str, Any]] = []
@@ -1198,13 +1059,8 @@ class GeminiLLM(LLMInterface):
system_texts.append(
"You must respond with valid JSON matching this schema:\n" + json.dumps(schema, ensure_ascii=False)
)
# #2699: Gemini always grammar-enforces structured output via its native
# response_schema (``strict`` is an OpenAI concept, meaningless here). Set
# the native schema whenever one is present so the batch path mirrors the
# interactive path; otherwise batch requests at default config
# (HINDSIGHT_API_LLM_STRICT_SCHEMA=False) get only a textual hint and
# intermittently emit malformed JSON, losing every fact in the chunk.
generation_config["responseJsonSchema"] = schema
if json_schema.get("strict"):
generation_config["responseJsonSchema"] = schema
request: dict[str, Any] = {"contents": contents}
if system_texts:
@@ -1307,6 +1163,3 @@ class GeminiLLM(LLMInterface):
"""Clean up resources (close connections, etc.)."""
# Gemini client doesn't require explicit cleanup
pass
def supports_attempt_scoped_concurrency(self) -> bool:
return True
@@ -17,24 +17,14 @@ import json
import logging
import os
import time
from contextlib import AbstractAsyncContextManager, nullcontext
from typing import Any, Callable
from typing import Any
from litellm.exceptions import Timeout as LiteLLMTimeout
from hindsight_api.config import DEFAULT_LLM_TIMEOUT, ENV_LLM_TIMEOUT
from hindsight_api.engine.llm_interface import (
LLM_TOOL_CHOICE_AUTO,
LLMInterface,
LLMToolChoice,
LLMToolChoiceMode,
OutputTooLongError,
)
from hindsight_api.engine.llm_interface import LLMInterface, OutputTooLongError
from hindsight_api.engine.llm_trace import LLMResponseUsage, stash_response_usage
from hindsight_api.engine.llm_wrapper import parse_llm_json
from hindsight_api.engine.providers.llm_debug import dump_request_on_4xx
from hindsight_api.engine.response_models import LLMToolCall, LLMToolCallResult, TokenUsage
from hindsight_api.engine.structured_output import strict_json_schema
from hindsight_api.metrics import get_metrics_collector
from hindsight_api.worker.stage import set_stage
@@ -236,7 +226,6 @@ class LiteLLMLLM(LLMInterface):
skip_validation: bool = False,
strict_schema: bool = False,
return_usage: bool = False,
attempt_context: Callable[[], AbstractAsyncContextManager[None]] | None = None,
) -> Any:
start_time = time.time()
@@ -244,7 +233,7 @@ class LiteLLMLLM(LLMInterface):
# Add JSON schema response format if provided
if response_format is not None and hasattr(response_format, "model_json_schema"):
schema = strict_json_schema(response_format) if strict_schema else response_format.model_json_schema()
schema = response_format.model_json_schema()
call_kwargs["response_format"] = {
"type": "json_schema",
"json_schema": {
@@ -257,13 +246,13 @@ class LiteLLMLLM(LLMInterface):
last_exception = None
for attempt in range(max_retries + 1):
if attempt > 0:
set_stage(f"llm.{self._stage_label}.{scope}.attempt={attempt + 1}/{max_retries + 1}")
try:
async with attempt_context() if attempt_context is not None else nullcontext():
set_stage(f"llm.{self._stage_label}.{scope}.attempt={attempt + 1}/{max_retries + 1}")
response = await asyncio.wait_for(
self._acompletion(**call_kwargs),
timeout=self.timeout,
)
response = await asyncio.wait_for(
self._acompletion(**call_kwargs),
timeout=self.timeout,
)
# Stash usage before the length check and parse/validate below,
# which may raise locally even though the provider charged for
# these tokens (#2387).
@@ -288,17 +277,7 @@ class LiteLLMLLM(LLMInterface):
try:
json_data = json.loads(clean_content)
except json.JSONDecodeError:
try:
json_data = json.loads(content)
except json.JSONDecodeError:
if attempt < max_retries:
# Prefer a clean re-roll first — a fresh generation
# usually beats repairing a malformed one.
raise
# Retry budget spent: structural repair as a last
# resort (#2547/#2544). Raises again if unrecoverable,
# which the outer handler surfaces loudly.
json_data = parse_llm_json(content)
json_data = json.loads(content)
if skip_validation:
result = json_data
@@ -399,9 +378,6 @@ class LiteLLMLLM(LLMInterface):
logger.error(f"LiteLLM auth error, not retrying: {e}")
raise
# Diagnostic dump (opt-in) of the exact request behind any 4xx.
dump_request_on_4xx(scope=scope, provider=self.provider, model=self.model, err=e, request=call_kwargs)
last_exception = e
if attempt < max_retries:
# Retry on rate limits, connection errors, server errors
@@ -432,31 +408,23 @@ class LiteLLMLLM(LLMInterface):
max_retries: int = 5,
initial_backoff: float = 1.0,
max_backoff: float = 30.0,
tool_choice: LLMToolChoice = LLM_TOOL_CHOICE_AUTO,
attempt_context: Callable[[], AbstractAsyncContextManager[None]] | None = None,
tool_choice: str | dict[str, Any] = "auto",
) -> LLMToolCallResult:
start_time = time.time()
call_kwargs = self._build_common_kwargs(messages, max_completion_tokens, temperature)
call_kwargs["tools"] = tools
call_kwargs["tool_choice"] = (
{
"type": "function",
"function": {"name": tool_choice.selected_function_name},
}
if tool_choice.mode is LLMToolChoiceMode.NAMED
else tool_choice.mode.value
)
call_kwargs["tool_choice"] = tool_choice
last_exception = None
for attempt in range(max_retries + 1):
if attempt > 0:
set_stage(f"llm.{self._stage_label}.tools.attempt={attempt + 1}/{max_retries + 1}")
try:
async with attempt_context() if attempt_context is not None else nullcontext():
set_stage(f"llm.{self._stage_label}.tools.attempt={attempt + 1}/{max_retries + 1}")
response = await asyncio.wait_for(
self._acompletion(**call_kwargs),
timeout=self.timeout,
)
response = await asyncio.wait_for(
self._acompletion(**call_kwargs),
timeout=self.timeout,
)
# Stash usage before the tool-call argument parse below, which
# can raise json.JSONDecodeError locally even though the provider
# already billed for these tokens; without this the error trace
@@ -556,9 +524,6 @@ class LiteLLMLLM(LLMInterface):
if "401" in error_str or "403" in error_str or "unauthorized" in error_str:
raise
# Diagnostic dump (opt-in) of the exact request behind any 4xx.
dump_request_on_4xx(scope=scope, provider=self.provider, model=self.model, err=e, request=call_kwargs)
last_exception = e
if attempt < max_retries:
is_retryable = any(
@@ -579,6 +544,3 @@ class LiteLLMLLM(LLMInterface):
async def cleanup(self) -> None:
"""Clean up resources."""
pass
def supports_attempt_scoped_concurrency(self) -> bool:
return True
@@ -19,11 +19,10 @@ import socket
import subprocess
import sys
import time
from contextlib import AbstractAsyncContextManager
from pathlib import Path
from typing import Any, Callable
from typing import Any
from hindsight_api.engine.llm_interface import LLM_TOOL_CHOICE_AUTO, LLMInterface, LLMToolChoice
from hindsight_api.engine.llm_interface import LLMInterface
from hindsight_api.engine.response_models import LLMToolCallResult
logger = logging.getLogger(__name__)
@@ -368,7 +367,6 @@ class LlamaCppLLM(LLMInterface):
skip_validation: bool = False,
strict_schema: bool = False,
return_usage: bool = False,
attempt_context: Callable[[], AbstractAsyncContextManager[None]] | None = None,
) -> Any:
"""Delegate call to the OpenAI-compatible API."""
await self._ensure_initialized()
@@ -384,7 +382,6 @@ class LlamaCppLLM(LLMInterface):
skip_validation=skip_validation,
strict_schema=strict_schema,
return_usage=return_usage,
attempt_context=attempt_context,
)
async def call_with_tools(
@@ -397,8 +394,7 @@ class LlamaCppLLM(LLMInterface):
max_retries: int = 5,
initial_backoff: float = 1.0,
max_backoff: float = 30.0,
tool_choice: LLMToolChoice = LLM_TOOL_CHOICE_AUTO,
attempt_context: Callable[[], AbstractAsyncContextManager[None]] | None = None,
tool_choice: str | dict[str, Any] = "auto",
) -> LLMToolCallResult:
"""Delegate tool calls to the OpenAI-compatible API."""
await self._ensure_initialized()
@@ -412,12 +408,8 @@ class LlamaCppLLM(LLMInterface):
initial_backoff=initial_backoff,
max_backoff=max_backoff,
tool_choice=tool_choice,
attempt_context=attempt_context,
)
def supports_attempt_scoped_concurrency(self) -> bool:
return True
async def cleanup(self) -> None:
"""Stop the shared llama.cpp server."""
global _shared_server
@@ -1,168 +0,0 @@
"""Opt-in diagnostic: dump the exact request behind an LLM 4xx rejection.
Some ``400 INVALID_ARGUMENT`` / ``400 Bad Request`` rejections of structured-output
calls are not reproducible by reconstructing the request after the fact the failing
factor lives in the request as it was actually assembled at runtime. Reconstructed
replays of the same inputs return ``200``, so the only reliable way to see what the
model rejected is to capture the real request at the moment it fails.
This helper is provider-agnostic. Every provider's error handler calls
``dump_request_on_4xx`` with whatever it assembled a Pydantic config
(google-genai ``GenerateContentConfig``), a kwargs dict (OpenAI / Anthropic /
LiteLLM ``**call_params``), etc. plus the raised error. The helper self-gates:
it is a no-op unless the ``llm_debug_dump_4xx`` config flag
(``HINDSIGHT_API_LLM_DEBUG_DUMP_4XX``) is enabled AND the error carries a 4xx
status, so callers can drop one unconditional call into each ``except`` block.
Safety / scope:
- Off by default the config flag is unset in normal operation.
- The serialized config omits message bodies (the ``messages``/``contents``/``input``
keys are stripped); message previews are length-capped, so an enabled dump can't
flood logs or spill large bodies.
- Never raises diagnostics must not break the request path (falls back to ``repr``).
"""
from __future__ import annotations
import json
import logging
from dataclasses import dataclass
from typing import Any
logger = logging.getLogger(__name__)
# Top-level request keys whose values are message bodies. Stripped from the config
# view so the dump never spills large user content — previews are logged separately.
_CONTENT_KEYS = ("messages", "contents", "input")
_PREVIEW_CHARS = 1500
_CONFIG_REPR_CAP = 8000
_ERR_CAP = 200
def _enabled() -> bool:
from hindsight_api.config import get_config
return bool(get_config().llm_debug_dump_4xx)
def status_code_of(err: Any) -> int | None:
"""Best-effort HTTP status of a provider error, across SDK error shapes.
OpenAI/Anthropic expose ``status_code``; google-genai uses ``code``; some wrap the
status on a ``response``. Returns None when no integer status is discoverable.
"""
for attr in ("status_code", "code", "http_status"):
value = getattr(err, attr, None)
if isinstance(value, int):
return value
response = getattr(err, "response", None)
if response is not None:
value = getattr(response, "status_code", None)
if isinstance(value, int):
return value
return None
def _serialize_config(request: Any) -> str:
"""Render the request config to a string without message bodies, never raising."""
try:
if request is None:
return "null"
# Pydantic models (google-genai GenerateContentConfig, SDK params objects).
dump = getattr(request, "model_dump_json", None)
if callable(dump):
return dump(exclude_none=True)
if isinstance(request, dict):
view = {k: v for k, v in request.items() if k not in _CONTENT_KEYS}
return json.dumps(view, ensure_ascii=False, default=str)
return repr(request)[:_CONFIG_REPR_CAP]
except Exception:
return repr(request)[:_CONFIG_REPR_CAP]
@dataclass
class _MessagePreview:
"""A message rendered for the dump: role + extracted text (not yet length-capped)."""
role: str
text: str
def _message_preview(msg: Any) -> _MessagePreview:
"""Extract role + text from a message across dict and provider-object shapes."""
# OpenAI / Anthropic dict: {"role": ..., "content": str | list[block]}
if isinstance(msg, dict):
role = str(msg.get("role", "?"))
content = msg.get("content")
if isinstance(content, str):
return _MessagePreview(role, content)
if isinstance(content, list):
text = ""
for block in content:
if isinstance(block, dict):
text += block.get("text") or ""
else:
text += getattr(block, "text", "") or ""
return _MessagePreview(role, text)
return _MessagePreview(role, "" if content is None else str(content))
# google-genai Content: role + parts[].text
role = str(getattr(msg, "role", "?"))
text = ""
for part in getattr(msg, "parts", None) or []:
text += getattr(part, "text", None) or ""
if not text:
text = getattr(msg, "content", "") or ""
return _MessagePreview(role, text)
def _resolve_messages(request: Any, messages: Any) -> Any:
"""Where per-message previews come from: explicit ``messages``, else inside ``request``."""
if messages is not None:
return messages
if isinstance(request, dict):
for key in _CONTENT_KEYS:
if key in request:
return request[key]
return []
def dump_request_on_4xx(
*,
scope: str,
provider: str,
model: str,
err: Any,
request: Any = None,
messages: Any = None,
) -> None:
"""Log the exact request behind an LLM 4xx when the diagnostic is enabled.
No-op unless ``HINDSIGHT_API_LLM_DEBUG_DUMP_4XX`` is truthy and ``err`` carries a
4xx status. ``request`` is whatever the provider assembled (a Pydantic config, a
kwargs dict, ...); ``messages`` overrides where the per-message previews come from
(defaults to the message list found inside ``request``).
"""
if not _enabled():
return
code = status_code_of(err)
if code is None or not (400 <= code < 500):
return
try:
cfg_repr = _serialize_config(request)
summary = []
for msg in _resolve_messages(request, messages) or []:
m = _message_preview(msg)
summary.append({"role": m.role, "chars": len(m.text), "preview": m.text[:_PREVIEW_CHARS]})
logger.error(
"[LLM_4XX_DUMP] provider=%s model=%s scope=%s code=%s err=%s config=%s contents=%s",
provider,
model,
scope,
code,
str(err)[:_ERR_CAP],
cfg_repr,
json.dumps(summary, ensure_ascii=False),
)
except Exception as dump_exc: # never let diagnostics break the request path
logger.warning("[LLM_4XX_DUMP] failed to serialize rejected request: %s", dump_exc)
@@ -7,10 +7,9 @@ without making actual API calls to external LLM services.
import logging
from collections.abc import Callable
from contextlib import AbstractAsyncContextManager
from typing import Any
from ..llm_interface import LLM_TOOL_CHOICE_AUTO, LLMInterface, LLMToolChoice, LLMToolChoiceMode
from ..llm_interface import LLMInterface
from ..response_models import LLMToolCall, LLMToolCallResult, TokenUsage
logger = logging.getLogger(__name__)
@@ -92,7 +91,6 @@ class MockLLM(LLMInterface):
skip_validation: bool = False,
strict_schema: bool = False,
return_usage: bool = False,
attempt_context: Callable[[], AbstractAsyncContextManager[None]] | None = None,
) -> Any:
"""
Make a mock LLM API call.
@@ -202,8 +200,7 @@ class MockLLM(LLMInterface):
max_retries: int = 5,
initial_backoff: float = 1.0,
max_backoff: float = 30.0,
tool_choice: LLMToolChoice = LLM_TOOL_CHOICE_AUTO,
attempt_context: Callable[[], AbstractAsyncContextManager[None]] | None = None,
tool_choice: str | dict[str, Any] = "auto",
) -> LLMToolCallResult:
"""
Make a mock LLM API call with tool/function calling support.
@@ -269,7 +266,7 @@ class MockLLM(LLMInterface):
else:
result = LLMToolCallResult(content="mock response", finish_reason="stop")
else:
result = self._compliant_tool_call(tools, tool_choice, messages)
result = LLMToolCallResult(content="mock response", finish_reason="stop")
# Set mock token usage on result if not already set
if result.input_tokens == 0:
@@ -300,61 +297,6 @@ class MockLLM(LLMInterface):
return result
@staticmethod
def _compliant_tool_call(
tools: list[dict[str, Any]],
tool_choice: LLMToolChoice,
messages: list[dict[str, Any]],
) -> LLMToolCallResult:
"""Default tool response: simulate a compliant tool-calling model.
Real providers drive the reflect loop entirely through tool calls -- they
honor a forced tool choice, then finish via ``done`` -- and the reflect
agent now rejects a turn that yields no tool call at all (a transport that
can't tool-call raises ReflectToolCallError). So the mock must behave like a
working provider here rather than returning bare "mock response" prose,
which used to be salvaged as the answer. Only this default path is affected;
tests that script turns via ``_response_callback`` / ``_mock_response`` are not.
"""
tool_names = {t.get("function", {}).get("name") for t in tools}
def _mock_query() -> str:
for message in reversed(messages):
content = message.get("content")
if message.get("role") == "user" and isinstance(content, str) and content.strip():
return content[:200]
return "mock query"
# Honor a forced retrieval tool so the loop actually runs recall/search and
# gathers evidence (populates based_on for tests that assert on it).
if tool_choice.mode is LLMToolChoiceMode.NAMED and tool_choice.function_name in {
"search_mental_models",
"search_observations",
"recall",
}:
return LLMToolCallResult(
tool_calls=[
LLMToolCall(
id="mock_forced",
name=tool_choice.function_name,
arguments={"reason": "mock", "query": _mock_query()},
)
],
finish_reason="tool_calls",
)
# Auto turn: finish via the done tool, mirroring a model that has gathered
# enough. The reflect evidence guardrail handles the empty-bank case (no
# evidence -> forced text synthesis on the final iteration).
if "done" in tool_names:
return LLMToolCallResult(
tool_calls=[LLMToolCall(id="mock_done", name="done", arguments={"answer": "mock response"})],
finish_reason="tool_calls",
)
# No done tool offered (non-reflect tool call): fall back to plain text.
return LLMToolCallResult(content="mock response", finish_reason="stop")
@staticmethod
def _build_mock_facts(messages: list[dict]) -> dict:
"""Build a canned fact extraction response from the user message text.
@@ -8,10 +8,9 @@ it raises a clear error instead of a confusing connection failure.
"""
import logging
from contextlib import AbstractAsyncContextManager
from typing import Any, Callable
from typing import Any
from ..llm_interface import LLM_TOOL_CHOICE_AUTO, LLMInterface, LLMToolChoice
from ..llm_interface import LLMInterface
from ..response_models import LLMToolCallResult
logger = logging.getLogger(__name__)
@@ -49,7 +48,6 @@ class NoneLLM(LLMInterface):
skip_validation: bool = False,
strict_schema: bool = False,
return_usage: bool = False,
attempt_context: Callable[[], AbstractAsyncContextManager[None]] | None = None,
) -> Any:
"""Raise LLMNotAvailableError — no LLM is configured."""
raise LLMNotAvailableError(
@@ -67,8 +65,7 @@ class NoneLLM(LLMInterface):
max_retries: int = 5,
initial_backoff: float = 1.0,
max_backoff: float = 30.0,
tool_choice: LLMToolChoice = LLM_TOOL_CHOICE_AUTO,
attempt_context: Callable[[], AbstractAsyncContextManager[None]] | None = None,
tool_choice: str | dict[str, Any] = "auto",
) -> LLMToolCallResult:
"""Raise LLMNotAvailableError — no LLM is configured."""
raise LLMNotAvailableError(
@@ -26,10 +26,9 @@ import logging
import os
import re
import time
from contextlib import AbstractAsyncContextManager, nullcontext
from datetime import UTC, datetime, timedelta
from email.utils import parsedate_to_datetime
from typing import Any, Callable
from typing import Any
from urllib.parse import parse_qs, urlparse, urlunparse
import httpx
@@ -37,18 +36,9 @@ from openai import APIConnectionError, APIStatusError, AsyncOpenAI, LengthFinish
from hindsight_api.config import DEFAULT_LLM_TIMEOUT, ENV_LLM_TIMEOUT
from hindsight_api.engine.bank_attribution import apply_bank_attribution
from hindsight_api.engine.llm_interface import (
LLM_TOOL_CHOICE_AUTO,
LLMInterface,
LLMToolChoice,
LLMToolChoiceMode,
OutputTooLongError,
ProviderRateLimitResetError,
)
from hindsight_api.engine.llm_interface import LLMInterface, OutputTooLongError, ProviderRateLimitResetError
from hindsight_api.engine.llm_trace import LLMResponseUsage, stash_response_usage
from hindsight_api.engine.providers.llm_debug import dump_request_on_4xx
from hindsight_api.engine.response_models import LLMToolCall, LLMToolCallResult, TokenUsage
from hindsight_api.engine.structured_output import strict_json_schema
from hindsight_api.metrics import get_metrics_collector
from hindsight_api.worker.stage import set_stage
@@ -57,57 +47,17 @@ logger = logging.getLogger(__name__)
# Seed applied to every Groq request for deterministic behavior
DEFAULT_LLM_SEED = 4242
JSON_MODE_USER_HINT = "Return valid json only."
DEFAULT_VERIFICATION_MAX_COMPLETION_TOKENS = 512
def _validate_ollama_num_ctx(value: Any) -> int | None:
"""Validate a native Ollama context-window override."""
if value is None:
return None
if isinstance(value, bool) or not isinstance(value, int):
raise ValueError(f"ollama_num_ctx must be a positive integer, got {value!r}")
if value < 1:
raise ValueError(f"ollama_num_ctx must be >= 1, got {value}")
return value
# Provider implementations that advertise tool_choice="required"
# Self-hosted OpenAI-compatible servers that advertise tool_choice="required"
# but silently ignore it: instead of forcing a tool call they return
# finish_reason "stop"/"tool_calls" with an EMPTY tool_calls array and no error.
# Reflect's agent loop then sees no tool call, runs synthesis with no retrieval,
# and answers "I don't have information" even when the bank holds the answer.
# See issues #1563 (LM Studio), #1179 (LM Studio + Qwen), #1877 (vLLM with
# --enable-auto-tool-choice). The generic OpenAI provider is intentionally not
# inferred from its URL: custom OpenAI-compatible endpoints can implement the
# required-tool contract, and silently downgrading them changes request semantics.
# llama-server (the "llamacpp" provider) honors "required" correctly and is
# intentionally excluded (#1179).
# --enable-auto-tool-choice). llama-server (the "llamacpp" provider) honors
# "required" correctly and is intentionally excluded (#1179).
_TOOL_CHOICE_REQUIRED_UNSUPPORTED_PROVIDERS = frozenset({"lmstudio", "ollama"})
# Local providers whose OpenAI-compatible surface always lives under a `/v1`
# path (LM Studio: http://localhost:1234/v1, Ollama: http://localhost:11434/v1).
# For these we know the exact endpoint shape, so a bare host base URL can be
# normalized safely. Cloud/proxy endpoints are left untouched — their path is
# provider-specific and must be supplied verbatim.
_V1_PATH_LOCAL_PROVIDERS = frozenset({"lmstudio", "ollama"})
def _ensure_v1_base_url(base_url: str) -> str:
"""Append the OpenAI-compatible ``/v1`` prefix to a bare local base URL.
LM Studio's server UI advertises its address as ``http://localhost:1234``,
so users commonly set ``HINDSIGHT_API_LLM_BASE_URL`` to that bare host. The
OpenAI SDK then POSTs to ``<host>/chat/completions`` and LM Studio rejects it
with ``Unexpected endpoint or method`` its OpenAI-compatible routes live
under ``/v1``. Only a base URL with no meaningful path (bare host or a lone
trailing slash) is rewritten; anything with an explicit path (e.g. a reverse
proxy mount or an already-correct ``/v1``) is returned unchanged. See #2922.
"""
parsed = urlparse(base_url)
if parsed.path.strip("/"):
return base_url
return urlunparse(parsed._replace(path="/v1"))
class ProviderResponseError(RuntimeError):
"""Raised when a provider returns a success response without usable content."""
@@ -117,68 +67,23 @@ class ProviderResponseError(RuntimeError):
self.retryable = retryable
def _is_json(text: str) -> bool:
"""True if ``text`` parses as a JSON value."""
try:
json.loads(text)
except (json.JSONDecodeError, ValueError):
return False
return True
def _outer_json_span(content: str) -> str | None:
"""Return the outermost ``{...}`` / ``[...]`` span if it parses as JSON, else None.
Fallback for responses where fences are partial/absent or the model wrapped
the JSON in surrounding prose. Only returned when it is valid JSON so callers
never receive a worse candidate than the raw content.
"""
starts = [i for i in (content.find("{"), content.find("[")) if i >= 0]
ends = [i for i in (content.rfind("}"), content.rfind("]")) if i >= 0]
if not starts or not ends:
return None
start, end = min(starts), max(ends)
if end <= start:
return None
candidate = content[start : end + 1].strip()
return candidate if _is_json(candidate) else None
def _strip_code_fences(content: str) -> str:
"""Strip markdown code fences from LLM response if present.
Many LLM providers (MiniMax, some Ollama models, Claude via proxies)
wrap JSON responses in ```json ... ``` fences even when json_object
response format is requested. Fences are detected by line (a closing
``` must sit alone on its line) so triple-backticks *inside* JSON string
values do not truncate the payload. When the stripped candidate is not
valid JSON (partial fence, prose-wrapped output, truncated response), fall
back to the outermost parseable JSON span. Returns the original content
unchanged if no better candidate is found.
response format is requested. This strips the fences while preserving
the JSON content inside. Returns the original content unchanged if
no fences are detected.
"""
candidate = content
if "```" in content:
lines = content.split("\n")
# Find first line that starts a code fence (``` optionally followed by language)
fence_start = next((i for i, line in enumerate(lines) if line.startswith("```")), None)
if fence_start is not None:
# Find matching closing fence (``` alone or with trailing whitespace)
fence_end = next(
(j for j in range(fence_start + 1, len(lines)) if lines[j].strip() == "```"),
None,
)
if fence_end is not None:
candidate = "\n".join(lines[fence_start + 1 : fence_end]).strip()
if _is_json(candidate):
return candidate
# Fence stripping did not yield valid JSON — try to recover the outer JSON span.
span = _outer_json_span(content)
if span is not None:
return span
return candidate
if "```" not in content:
return content
try:
if "```json" in content:
return content.split("```json")[1].split("```")[0].strip()
return content.split("```")[1].split("```")[0].strip()
except (IndexError, ValueError):
return content
# Reasoning/thinking tags emitted by extended-thinking models. Some providers
@@ -530,8 +435,6 @@ class OpenAICompatibleLLM(LLMInterface):
timeout: float | None = None,
groq_service_tier: str | None = None,
extra_body: dict[str, Any] | None = None,
*,
ollama_num_ctx: int | None = None,
**kwargs: Any,
):
"""
@@ -542,15 +445,10 @@ class OpenAICompatibleLLM(LLMInterface):
api_key: API key (optional for ollama/lmstudio).
base_url: Base URL for the API (uses defaults for groq/ollama/lmstudio if empty).
model: Model name.
reasoning_effort: Reasoning effort level for supported models
("none", "low", "medium", "high"). "none" is required when calling
function tools on some reasoning models, which reject every other
value including omitting the parameter entirely.
reasoning_effort: Reasoning effort level for supported models ("low", "medium", "high").
timeout: Request timeout in seconds (uses env var or 120s default).
groq_service_tier: Groq service tier ("on_demand", "flex", "auto").
extra_body: Extra body params merged into every API call.
ollama_num_ctx: Native Ollama context window override. None lets Ollama use
the model/server default.
**kwargs: Additional provider-specific parameters.
"""
super().__init__(provider, api_key, base_url, model, reasoning_effort, **kwargs)
@@ -605,11 +503,6 @@ class OpenAICompatibleLLM(LLMInterface):
# lives on a separate control-plane host — see FireworksLLM.
self.base_url = "https://api.fireworks.ai/inference/v1"
# Normalize bare local base URLs (e.g. a user pasting the address shown
# in the LM Studio UI) so the OpenAI SDK targets the `/v1` routes. See #2922.
if self.provider in _V1_PATH_LOCAL_PROVIDERS and self.base_url:
self.base_url = _ensure_v1_base_url(self.base_url)
# For ollama/lmstudio, use dummy key if not provided
if self.provider in ("ollama", "lmstudio") and not self.api_key:
self.api_key = "local"
@@ -636,7 +529,6 @@ class OpenAICompatibleLLM(LLMInterface):
# Service tier configuration (from config, not env vars)
self.groq_service_tier = groq_service_tier
self.openai_service_tier = kwargs.get("openai_service_tier")
self.ollama_num_ctx = _validate_ollama_num_ctx(ollama_num_ctx)
# User-configured extra body params (merged into every API call)
self._config_extra_body = extra_body or {}
@@ -667,17 +559,17 @@ class OpenAICompatibleLLM(LLMInterface):
def _drops_tool_choice_required(self) -> bool:
"""Whether this endpoint silently ignores ``tool_choice="required"``.
Only explicitly identified provider implementations are classified as
unsupported. A custom base URL does not identify endpoint capabilities:
an OpenAI-compatible endpoint may correctly enforce required tool calls,
and replacing ``required`` with ``auto`` would violate the caller's named
tool choice after the tools list has been narrowed.
True for self-hosted OpenAI-compatible servers known to return an empty
tool_calls array for "required" instead of forcing a call (#1563/#1179/
#1877). Covers LM Studio / Ollama directly, plus any server reached via
the generic "openai" provider with a custom ``base_url`` (e.g. a local
vLLM endpoint). The real OpenAI API (no base_url override) honors
"required", and cloud providers keep their own default base_urls, so both
are left untouched.
"""
return self.provider in _TOOL_CHOICE_REQUIRED_UNSUPPORTED_PROVIDERS
def _verification_max_completion_tokens(self) -> int:
"""Return the startup verification budget for OpenAI-compatible gateways."""
return DEFAULT_VERIFICATION_MAX_COMPLETION_TOKENS
if self.provider in _TOOL_CHOICE_REQUIRED_UNSUPPORTED_PROVIDERS:
return True
return self.provider == "openai" and bool(self.base_url)
async def verify_connection(self) -> None:
"""
@@ -690,7 +582,7 @@ class OpenAICompatibleLLM(LLMInterface):
logger.info(f"Verifying connection: {self.provider}/{self.model}")
await self.call(
messages=[{"role": "user", "content": "Say 'ok'"}],
max_completion_tokens=self._verification_max_completion_tokens(),
max_completion_tokens=100,
max_retries=2,
initial_backoff=0.5,
max_backoff=2.0,
@@ -752,11 +644,6 @@ class OpenAICompatibleLLM(LLMInterface):
# use the widely-supported max_tokens
return "max_tokens"
def _apply_provider_extra_body_defaults(self, extra_body: dict[str, Any]) -> None:
"""Apply provider-specific extra_body defaults while preserving user overrides."""
if self.provider == "minimax":
extra_body.setdefault("thinking", {"type": "disabled"})
async def call(
self,
messages: list[dict[str, str]],
@@ -770,7 +657,6 @@ class OpenAICompatibleLLM(LLMInterface):
skip_validation: bool = False,
strict_schema: bool = False,
return_usage: bool = False,
attempt_context: Callable[[], AbstractAsyncContextManager[None]] | None = None,
) -> Any:
"""
Make an LLM API call with retry logic.
@@ -811,7 +697,6 @@ class OpenAICompatibleLLM(LLMInterface):
skip_validation=skip_validation,
scope=scope,
return_usage=return_usage,
attempt_context=attempt_context,
)
start_time = time.time()
@@ -846,7 +731,6 @@ class OpenAICompatibleLLM(LLMInterface):
# Provider-specific parameters
extra_body: dict[str, Any] = {**self._config_extra_body}
self._apply_provider_extra_body_defaults(extra_body)
if self.provider == "groq":
call_params["seed"] = DEFAULT_LLM_SEED
# Add service_tier if configured
@@ -862,7 +746,7 @@ class OpenAICompatibleLLM(LLMInterface):
if response_format is not None:
schema = None
if hasattr(response_format, "model_json_schema"):
schema = strict_json_schema(response_format) if strict_schema else response_format.model_json_schema()
schema = response_format.model_json_schema()
if strict_schema and schema is not None:
# Use OpenAI's strict JSON schema enforcement
@@ -905,11 +789,11 @@ class OpenAICompatibleLLM(LLMInterface):
# Surface attempt count in worker stage so JSON-schema retry loops
# are visible from logs (small models on strict structured output
# often loop here). Cheap no-op outside worker context.
if attempt > 0:
set_stage(f"llm.{self.provider}.{scope}.attempt={attempt + 1}/{max_retries + 1}")
try:
if response_format is not None:
async with attempt_context() if attempt_context is not None else nullcontext():
set_stage(f"llm.{self.provider}.{scope}.attempt={attempt + 1}/{max_retries + 1}")
response = await self._client.chat.completions.create(**call_params)
response = await self._client.chat.completions.create(**call_params)
# Stash usage before parse/validate, which may raise locally
# even though the provider charged for these tokens (#2387).
stash_response_usage(_usage_from_openai_response(response))
@@ -965,9 +849,7 @@ class OpenAICompatibleLLM(LLMInterface):
else:
result = response_format.model_validate(json_data)
else:
async with attempt_context() if attempt_context is not None else nullcontext():
set_stage(f"llm.{self.provider}.{scope}.attempt={attempt + 1}/{max_retries + 1}")
response = await self._client.chat.completions.create(**call_params)
response = await self._client.chat.completions.create(**call_params)
stash_response_usage(_usage_from_openai_response(response))
result, first_choice = _content_or_error(
response,
@@ -1004,9 +886,7 @@ class OpenAICompatibleLLM(LLMInterface):
output_tokens = max(0, output_tokens - thoughts_tokens)
total_tokens = max(0, total_tokens - thoughts_tokens)
# Record LLM metrics. ``output_tokens`` is visible-only by now, so
# ``thoughts_tokens`` has to be recorded alongside it or the reasoning
# half of the billed output reaches no counter at all.
# Record LLM metrics
metrics = get_metrics_collector()
metrics.record_llm_call(
provider=self.provider,
@@ -1016,8 +896,6 @@ class OpenAICompatibleLLM(LLMInterface):
input_tokens=input_tokens,
output_tokens=output_tokens,
success=True,
cached_input_tokens=cached_tokens,
thoughts_tokens=thoughts_tokens,
)
# Record trace span
@@ -1086,9 +964,6 @@ class OpenAICompatibleLLM(LLMInterface):
logger.error(f"Auth error (HTTP {e.status_code}), not retrying: {str(e)}")
raise
# Diagnostic dump (opt-in) of the exact request behind any 4xx.
dump_request_on_4xx(scope=scope, provider=self.provider, model=self.model, err=e, request=call_params)
_raise_provider_quota_defer(
e, provider=self.provider, model=self.model, scope=scope, max_backoff=max_backoff
)
@@ -1179,8 +1054,7 @@ class OpenAICompatibleLLM(LLMInterface):
max_retries: int = 5,
initial_backoff: float = 1.0,
max_backoff: float = 30.0,
tool_choice: LLMToolChoice = LLM_TOOL_CHOICE_AUTO,
attempt_context: Callable[[], AbstractAsyncContextManager[None]] | None = None,
tool_choice: str | dict[str, Any] = "auto",
) -> LLMToolCallResult:
"""
Make an LLM API call with tool/function calling support.
@@ -1194,43 +1068,51 @@ class OpenAICompatibleLLM(LLMInterface):
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.
tool_choice: How to choose tools - "auto", "none", "required", or specific function.
Returns:
LLMToolCallResult with content and/or tool_calls.
"""
start_time = time.time()
request_tool_choice: str | None
if tool_choice.mode is LLMToolChoiceMode.NAMED:
forced_name = tool_choice.selected_function_name
filtered = [tool for tool in tools if tool.get("function", {}).get("name") == forced_name]
if len(filtered) != 1:
raise ValueError(
f"Named tool_choice must reference exactly one declared tool; "
f"found {len(filtered)} definitions for {forced_name!r}"
)
tools = filtered
request_tool_choice = LLMToolChoiceMode.REQUIRED.value
elif tool_choice.mode is LLMToolChoiceMode.AUTO:
request_tool_choice = None
else:
request_tool_choice = tool_choice.mode.value
request_tool_choice: str | dict[str, Any] | None = tool_choice
# Normalize named tool_choice dicts to "required" + filter tools.
# Some providers (e.g. LM Studio, Ollama) reject the OpenAI named format
# {"type": "function", "function": {"name": "..."}}. The semantics are
# identical to tool_choice="required" with the tools list restricted to
# just the requested tool, so we apply that transformation where supported.
if isinstance(request_tool_choice, dict) and request_tool_choice.get("type") == "function":
forced_name = request_tool_choice.get("function", {}).get("name")
if forced_name:
filtered = [t for t in tools if t.get("function", {}).get("name") == forced_name]
if filtered:
tools = filtered
request_tool_choice = "required"
# DeepSeek accepts tool calls but rejects explicit required/named
# tool_choice values. The tools list has already been narrowed for
# forced calls, so omitting tool_choice preserves the practical behavior.
if "deepseek" in self.model.lower() and tool_choice.mode is not LLMToolChoiceMode.AUTO:
if "deepseek" in self.model.lower() and request_tool_choice != "auto":
request_tool_choice = None
# LM Studio and Ollama silently drop tool_choice="required", returning an
# empty tool_calls array instead of forcing a call (#1563/#1179).
# "auto" is the OpenAI API default — omitting tool_choice is semantically
# identical. Some providers (e.g. DeepSeek's reasoner pathway, which
# deepseek-v4-flash falls into when thinking mode is enabled) reject the
# parameter outright, returning HTTP 400 even for value "auto". Sending it
# only when the caller asks for a non-default behaviour avoids those 400s
# without changing semantics for compliant providers.
if request_tool_choice == "auto":
request_tool_choice = None
# vLLM (--enable-auto-tool-choice), LM Studio, Ollama and similar
# self-hosted servers silently drop tool_choice="required", returning an
# empty tool_calls array instead of forcing a call (#1563/#1179/#1877).
# Downgrade to auto (None) so the model still gets to call a tool. Named
# tool_choice dicts were already normalized to "required" + a single
# filtered tool above, so the call stays practically forced even under
# auto. Generic OpenAI-compatible endpoints retain the canonical
# ``required`` contract regardless of whether they use a custom base URL.
if request_tool_choice == LLMToolChoiceMode.REQUIRED.value and self._drops_tool_choice_required():
# auto. The real OpenAI API honors "required" and is left untouched.
if request_tool_choice == "required" and self._drops_tool_choice_required():
request_tool_choice = None
# DeepSeek tool-call replies can carry provider-specific reasoning_content.
@@ -1265,16 +1147,8 @@ class OpenAICompatibleLLM(LLMInterface):
temperature = max(0.01, min(temperature, 1.0))
call_params["temperature"] = temperature
# Set reasoning_effort for reasoning models, matching call(). Omitting it
# here is not a neutral default: OpenAI rejects function tools on a
# reasoning model unless reasoning_effort is present and set to "none",
# so leaving it out fails exactly like sending an unsupported value.
if self._supports_reasoning_model():
call_params["reasoning_effort"] = self.reasoning_effort
# Provider-specific parameters
extra_body: dict[str, Any] = {**self._config_extra_body}
self._apply_provider_extra_body_defaults(extra_body)
if self.provider == "groq":
call_params["seed"] = DEFAULT_LLM_SEED
if extra_body:
@@ -1285,10 +1159,10 @@ class OpenAICompatibleLLM(LLMInterface):
last_exception = None
for attempt in range(max_retries + 1):
if attempt > 0:
set_stage(f"llm.{self.provider}.tools.attempt={attempt + 1}/{max_retries + 1}")
try:
async with attempt_context() if attempt_context is not None else nullcontext():
set_stage(f"llm.{self.provider}.tools.attempt={attempt + 1}/{max_retries + 1}")
response = await self._client.chat.completions.create(**call_params)
response = await self._client.chat.completions.create(**call_params)
message = response.choices[0].message
finish_reason = response.choices[0].finish_reason
@@ -1322,8 +1196,6 @@ class OpenAICompatibleLLM(LLMInterface):
if thoughts_tokens:
output_tokens = max(0, output_tokens - thoughts_tokens)
# See ``call()``: record the reasoning and cached counts too, so no
# billed token is dropped from the metrics counters.
metrics = get_metrics_collector()
metrics.record_llm_call(
provider=self.provider,
@@ -1333,8 +1205,6 @@ class OpenAICompatibleLLM(LLMInterface):
input_tokens=input_tokens,
output_tokens=output_tokens,
success=True,
cached_input_tokens=cached_tokens,
thoughts_tokens=thoughts_tokens,
)
# Record OpenTelemetry span
@@ -1396,10 +1266,6 @@ class OpenAICompatibleLLM(LLMInterface):
f"not retrying: {_summarize_status_error(e)}"
)
raise
# Diagnostic dump (opt-in) of the exact request behind any 4xx.
dump_request_on_4xx(scope=scope, provider=self.provider, model=self.model, err=e, request=call_params)
_raise_provider_quota_defer(
e, provider=self.provider, model=self.model, scope=scope, max_backoff=max_backoff
)
@@ -1436,7 +1302,6 @@ class OpenAICompatibleLLM(LLMInterface):
skip_validation: bool,
scope: str = "memory",
return_usage: bool = False,
attempt_context: Callable[[], AbstractAsyncContextManager[None]] | None = None,
) -> Any:
"""
Call Ollama using native API with JSON schema enforcement.
@@ -1472,10 +1337,9 @@ class OpenAICompatibleLLM(LLMInterface):
# Add optional parameters with optimized defaults for Ollama
options: dict[str, Any] = {
"num_ctx": 16384, # 16k context window for larger prompts
"num_batch": 512, # Optimal batch size for prompt processing
}
if self.ollama_num_ctx is not None:
options["num_ctx"] = self.ollama_num_ctx
if max_completion_tokens:
options["num_predict"] = max_completion_tokens
if temperature is not None:
@@ -1491,10 +1355,10 @@ class OpenAICompatibleLLM(LLMInterface):
async with httpx.AsyncClient(timeout=300.0) as client:
for attempt in range(max_retries + 1):
if attempt > 0:
set_stage(f"llm.ollama_native.{scope}.attempt={attempt + 1}/{max_retries + 1}")
try:
async with attempt_context() if attempt_context is not None else nullcontext():
set_stage(f"llm.ollama_native.{scope}.attempt={attempt + 1}/{max_retries + 1}")
response = await client.post(native_url, json=payload, headers=headers)
response = await client.post(native_url, json=payload, headers=headers)
response.raise_for_status()
result = response.json()
@@ -1731,6 +1595,3 @@ class OpenAICompatibleLLM(LLMInterface):
"""Clean up resources (close OpenAI client connections)."""
if hasattr(self, "_client") and self._client:
await self._client.close()
def supports_attempt_scoped_concurrency(self) -> bool:
return True

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