Compare commits
1
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
c295d765de |
@@ -1,7 +1,6 @@
|
||||
{
|
||||
"$schema": "https://anthropic.com/claude-code/marketplace.schema.json",
|
||||
"name": "hindsight",
|
||||
"version": "0.7.2",
|
||||
"description": "Official Hindsight integrations for Claude Code",
|
||||
"owner": {
|
||||
"name": "vectorize-io"
|
||||
|
||||
@@ -73,11 +73,6 @@ results = await asyncio.gather(*tasks, return_exceptions=True)
|
||||
results = await asyncio.gather(*tasks, return_exceptions=True)
|
||||
```
|
||||
|
||||
### API Layer & Data Access
|
||||
- **No direct database access in `api/http.py`** (or any API router). HTTP handlers must not build SQL, call `acquire_with_retry` / `conn.fetch` / `conn.fetchrow` / `conn.execute`, or reference `fq_table(...)`. All persistence and queries live in `MemoryEngine` (the engine layer). A handler parses/validates the request, calls an engine method, shapes the HTTP response, and maps domain results to status codes (e.g. a `None` return → 404).
|
||||
- **Authentication/tenancy is enforced inside each engine method, not assumed by the handler.** Every engine method that touches bank-scoped data must authenticate via `request_context` — typically `await self._authenticate_tenant(request_context)` (often indirectly through `get_bank_profile(...)`) — so the correct tenant schema is resolved before any query runs. Handlers must thread `request_context` through to the engine method; never query a tenant-scoped table assuming the schema is already set.
|
||||
- Engine methods return typed models (Pydantic/dataclass), not raw dicts (see Type Safety).
|
||||
|
||||
### 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.
|
||||
@@ -140,13 +135,6 @@ For each new or significantly changed function/endpoint/class:
|
||||
|
||||
Flag any new logic that lacks test coverage.
|
||||
|
||||
**LLM-behaviour changes need a real-LLM judge test, not MockLLM.** If the change alters how the model interprets a prompt — fact/observation extraction, `fact_type` (world/experience) classification, speaker attribution, instruction-following, prompt wording — there MUST be a test marked `pytest.mark.hs_llm_core` that runs the real pipeline and asserts via `tests.llm_judge.assert_meets_criteria` (not string/enum matching). Flag these as findings:
|
||||
- A prompt/classification change verified only by MockLLM or string assertions (MockLLM echoes input — such tests pass spuriously). **Should fix.**
|
||||
- A test that hard-asserts `fact_type == "world"/"experience"` (or other model-decided output) instead of judging it — non-deterministic, will flake across providers/runs. **Should fix** (move the classification check into the judge `criteria`; keep only genuinely deterministic structural asserts direct).
|
||||
- Deterministic mechanics (prompt assembly, suppression/branching logic) that are covered *only* by a slow LLM test — these should also have fast non-LLM unit tests. **Note.**
|
||||
|
||||
See CLAUDE.md → Key Conventions → Testing for the full pattern.
|
||||
|
||||
### 7. Check API consistency
|
||||
|
||||
If any files in `hindsight-api-slim/hindsight_api/api/` were changed:
|
||||
@@ -154,12 +142,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/`)
|
||||
|
||||
### 7b. Check API-layer data-access boundary
|
||||
|
||||
For each changed handler in `hindsight-api-slim/hindsight_api/api/` (e.g. `http.py`, `mcp.py`):
|
||||
- **Flag any direct DB access in the handler** — `acquire_with_retry`, `conn.fetch` / `fetchrow` / `execute`, raw SQL strings, or `fq_table(...)`. These are a **must fix**: the query must be moved into a `MemoryEngine` method that returns a typed model, and the handler must call that method.
|
||||
- **Verify authentication is enforced in the engine** — the handler must delegate to an engine method that authenticates via `request_context` (`_authenticate_tenant`, typically through `get_bank_profile`). A handler that reads/writes tenant-scoped data without an engine method enforcing auth is a **must fix** (tenant data could leak across schemas).
|
||||
|
||||
### 8. Check code comments
|
||||
|
||||
For each non-trivial change:
|
||||
@@ -172,8 +154,7 @@ For each non-trivial change:
|
||||
If any files in `hindsight-integrations/` were added or changed, verify:
|
||||
- **Tests exist** — the integration must have tests that simulate/exercise the external framework (not just pure unit tests of helpers). Check for a `tests/` directory with meaningful test files.
|
||||
- **CI job exists** — check `.github/workflows/test.yml` for a corresponding `test-<name>-integration` job. If missing, flag it.
|
||||
- **Release process** — check that the integration name is in the `VALID_INTEGRATIONS` array in `scripts/release-integration.sh` AND in the `INTEGRATIONS` dict in `hindsight-dev/hindsight_dev/generate_changelog.py` (the changelog generator keeps its own list; a release fails at the changelog step if the name is missing there). If either is missing, flag it.
|
||||
- **Docs gallery + sidebar entry** — the integration must have an entry in `hindsight-docs/src/data/integrations.json`. This file is the **single source of truth** that drives both the integrations gallery and the docs sidebar (the sidebar category is injected from it at render time across all docs versions). The entry needs an internal `/sdks/integrations/<slug>` `link` and a matching page at `hindsight-docs/docs-integrations/<slug>.md(x)`. The `hindsight-docs/scripts/check-integrations.mjs` build step enforces both directions — forward: every internal JSON entry has a doc page; reverse: every released tag (`integrations/<name>/vX.Y.Z`) appears in the JSON (private infra like `cloudflare-oauth-proxy` is in the script's `EXCLUDED` set). Flag any integration that is released (or being released) but missing from `integrations.json`, and any JSON entry without a doc page. Do **not** hand-edit `versioned_sidebars/*.json` to add integration links — they are positional placeholders filled from the JSON.
|
||||
- **Release process** — check that the integration name is in the `VALID_INTEGRATIONS` array in `scripts/release-integration.sh`. If missing, flag it.
|
||||
- **Code standards** — the integration code must follow all Python style rules (type hints, no raw dicts, no tuple returns, etc.).
|
||||
|
||||
### 10. Check MCP tool registration completeness
|
||||
@@ -185,26 +166,7 @@ If any new MCP tools were added or existing tools renamed in `hindsight-api-slim
|
||||
- **`MCP_TOOL_GROUPS`** in `hindsight-control-plane/src/components/bank-config-view.tsx` — must include the new tool in the appropriate group for the UI tool selector
|
||||
- **Tool count assertions** in tests (e.g., `test_mcp_tools.py`) — must be updated to reflect the new count
|
||||
|
||||
### 11. Check backup/restore table coverage
|
||||
|
||||
If a migration adds a new PostgreSQL table (look for `CREATE TABLE` / `op.create_table` in `hindsight-api-slim/hindsight_api/alembic/versions/`):
|
||||
- **`BACKUP_TABLES`** in `hindsight-api-slim/hindsight_api/admin/cli.py` — must include the new table, placed after any table it references via foreign key (parents before children). A missing entry is silent data loss: the table is never backed up, and restore's `TRUNCATE banks CASCADE` wipes any FK-to-banks child (e.g. `mental_models`, `directives`) on restore even though it was never saved.
|
||||
- The guard test `test_backup_tables_covers_entire_schema` in `tests/test_admin_backup_restore.py` enforces this — flag it as a **must fix** if a new table is absent from `BACKUP_TABLES`.
|
||||
- Oracle-only tables (e.g. `observation_sources`) are intentionally excluded — admin backup/restore is PostgreSQL-only.
|
||||
|
||||
### 11b. Check new config flags update the env template
|
||||
|
||||
If the diff adds a new configuration field (a new `ENV_*` / `HINDSIGHT_*` env var
|
||||
in `hindsight-api-slim/hindsight_api/config.py`):
|
||||
- **`.env.example`** (repo root) — must add the variable (commented if optional)
|
||||
alongside the docs entry in `hindsight-docs/docs/developer/configuration.md`.
|
||||
A flag added to `config.py` but absent from `.env.example` is a **should fix**.
|
||||
- **`hindsight-embed/hindsight_embed/env.example`** — the bundled copy must stay
|
||||
byte-identical to the repo-root `.env.example` (it seeds embed/profile configs).
|
||||
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**.
|
||||
|
||||
### 12. Review against other coding standards
|
||||
### 11. Review against other coding standards
|
||||
|
||||
Check the diff for violations of the standards listed above:
|
||||
- Python files at project root (not allowed)
|
||||
@@ -216,7 +178,7 @@ Check the diff for violations of the standards listed above:
|
||||
- Premature abstractions or speculative helpers
|
||||
- Backwards-compatibility hacks (unused vars, re-exports, "removed" comments)
|
||||
|
||||
### 13. Report findings
|
||||
### 12. Report findings
|
||||
|
||||
Present a clear summary organized by severity:
|
||||
|
||||
@@ -227,11 +189,7 @@ Present a clear summary organized by severity:
|
||||
- Raw dict usage for structured data (including internal code)
|
||||
- Multi-item tuple returns (including internal code)
|
||||
- Missing tests for new endpoints
|
||||
- Direct DB access (raw SQL / `acquire_with_retry` / `fq_table`) in an `api/` handler instead of a `MemoryEngine` method
|
||||
- Tenant-scoped data accessed without authentication enforced in the engine (`_authenticate_tenant` / `get_bank_profile`)
|
||||
- New integration missing tests, CI job, or release-integration.sh entry
|
||||
- Released/added integration missing from `hindsight-docs/src/data/integrations.json`, or a JSON entry with no `docs-integrations/<slug>` page (fails the docs build via `check-integrations.mjs`)
|
||||
- New PostgreSQL table missing from `BACKUP_TABLES` in `admin/cli.py` (silent data loss on restore)
|
||||
|
||||
**Should fix** — issues that hurt code quality:
|
||||
- Dead code / unused imports missed by linter
|
||||
|
||||
@@ -1,116 +0,0 @@
|
||||
---
|
||||
name: hs-release
|
||||
description: Cut a core Hindsight release (vX.Y.Z) and open the changelog + blog PR. Use when asked to cut/start a release, bump the version, or publish a new Hindsight version.
|
||||
user_invocable: true
|
||||
---
|
||||
|
||||
# Hindsight Release
|
||||
|
||||
Cut a **core** Hindsight release and open the accompanying changelog/blog PR. This is for the core
|
||||
product version (API, clients, CLI, control plane, Helm). **Integrations are versioned
|
||||
independently** — use `scripts/release-integration.sh` for those, not this skill.
|
||||
|
||||
The release is **irreversible and outward-facing**: it tags a version and pushes it straight to
|
||||
`main`, which triggers CI that publishes packages to PyPI / npm / Helm. Confirm the version number
|
||||
and that the intended fixes are already merged to `main` before you start.
|
||||
|
||||
## Step 0 — Pre-flight
|
||||
|
||||
1. **Decide the base.** A release is cut from the latest `origin/main`, never from a feature
|
||||
branch. `git fetch origin --tags` first. Confirm the "couple of fixes" the user means are
|
||||
actually merged to `main` (`git log v<prev>..origin/main --oneline`).
|
||||
2. **Find where `main` is checked out.** `main` is often already checked out in a sibling worktree
|
||||
(`git worktree list`). You **cannot** check out `main` in a second worktree — run the release in
|
||||
the worktree that already holds it. If that worktree is dirty with throwaway cruft
|
||||
(`.next-*` tsconfig paths, screenshots), `git stash push -u`, fast-forward to `origin/main`,
|
||||
run the release, then `git stash pop`.
|
||||
3. **Pitfall:** never pipe the checkout in an `&&` chain like
|
||||
`git checkout main 2>&1 | tail && git reset --hard ...` — the pipe's exit status is `tail`'s
|
||||
(always 0), so a failed checkout won't stop the chain and the `reset` fires on the **wrong
|
||||
branch**. Check out as its own command and verify `git branch --show-current` before resetting.
|
||||
|
||||
## Step 1 — Cut the release
|
||||
|
||||
Run from the worktree on a clean `main`:
|
||||
|
||||
```bash
|
||||
./scripts/release.sh <version> # e.g. 0.8.1 (no leading v)
|
||||
```
|
||||
|
||||
`release.sh` bumps the version in every component, regenerates the OpenAPI spec + all client SDKs,
|
||||
updates docs versioning, commits `Release v<version>`, tags `v<version>`, and **pushes the commit
|
||||
and tag directly to `main`**. The push triggers the `Release` GitHub Actions workflow that builds
|
||||
and publishes the packages. It is **not** a PR.
|
||||
|
||||
Verify after: `gh run list --limit 5` should show the `Release v<version>` workflow running, and
|
||||
`git ls-remote --tags origin v<version>` should return the tag.
|
||||
|
||||
## Step 2 — Changelog + blog PR (separate)
|
||||
|
||||
Done **after** the tag exists, as its own PR (precedent: v0.8.0 = #2053, v0.8.1 = #2080). Work on a
|
||||
branch off the new `main`:
|
||||
|
||||
```bash
|
||||
git checkout -b docs-changelog-<version> origin/main
|
||||
```
|
||||
|
||||
Only spin up a separate worktree (`git worktree add ../hindsight-changelog-<version> -b
|
||||
docs-changelog-<version> origin/main`) if you can't get a clean checkout otherwise — e.g. `main` is
|
||||
held in another worktree and the current one has work you don't want to disturb.
|
||||
|
||||
**Branch naming:** use the `docs-` (hyphen) convention, e.g. `docs-changelog-0.8.1`. A remote
|
||||
branch literally named `docs` exists, so any `docs/...` branch is rejected on push with
|
||||
`directory file conflict`.
|
||||
|
||||
### Changelog
|
||||
|
||||
```bash
|
||||
uv run --directory hindsight-dev generate-changelog <version>
|
||||
```
|
||||
|
||||
LLM-summarizes the commits between the previous tag and `v<version>` and prepends an entry to
|
||||
`hindsight-docs/src/pages/changelog/index.md`. Requires `OPENAI_API_KEY` (already in the repo
|
||||
`.env`). It excludes `hindsight-integrations/` source, but new integrations whose commits also
|
||||
touched docs will still appear — that matches precedent, leave them in the **changelog**.
|
||||
|
||||
### Blog post
|
||||
|
||||
Hand-write `hindsight-docs/blog/YYYY-MM-DD-version-X-Y-Z.md` (mirror an existing one; patch
|
||||
releases are short — see `2026-06-02-version-0-7-2.md`). Guidance:
|
||||
|
||||
- **Explain user impact, not internals/mechanism.** Lead with what the user can now do and what to
|
||||
set. Config/env-var names are fine (developer-facing), code symbols and internals are not.
|
||||
- **Do not list integrations in the release blog.** The core blog covers core engine / API /
|
||||
ops changes; each integration ships its own changelog. (Integrations may still appear in the
|
||||
generated `changelog/index.md` — that's fine; just keep them out of the blog.)
|
||||
- Call out an upgrade recommendation when there are operational/data-integrity fixes.
|
||||
- Validate formatting: `npx prettier --check <blog file>`.
|
||||
|
||||
### Sync the docs skill
|
||||
|
||||
```bash
|
||||
./scripts/generate-docs-skill.sh
|
||||
```
|
||||
|
||||
Refreshes `skills/hindsight-docs/references/changelog/index.md`. It will also bump
|
||||
`skills/hindsight-docs/references/openapi.json` by one version — `release.sh` regenerates the skill
|
||||
*before* bumping OpenAPI, so the skill copy lags a version in the release commit; this step syncs
|
||||
it. Expect a one-line `version` diff there; keep it.
|
||||
|
||||
### Commit, push, PR
|
||||
|
||||
```bash
|
||||
git add -A
|
||||
git commit --no-verify -m "docs: changelog and blog post for v<version>"
|
||||
git push -u origin docs-changelog-<version>
|
||||
gh pr create --base main --title "docs: changelog and blog post for v<version>" --body "..."
|
||||
```
|
||||
|
||||
Expected files in the PR: the changelog entry, the new blog post, the regenerated skill changelog
|
||||
mirror, and the skill `openapi.json` version sync.
|
||||
|
||||
## Cleanup
|
||||
|
||||
If you created a temporary worktree, remove it once the PR is up
|
||||
(`git worktree remove ../hindsight-changelog-<version>`; the branch stays on origin). Restore any
|
||||
stash you popped in Step 0.
|
||||
+3
-96
@@ -2,24 +2,11 @@
|
||||
# Copy this file to .env and fill in your values
|
||||
|
||||
# LLM Configuration (Required)
|
||||
# Supported providers: openai, groq, ollama, gemini, anthropic, lmstudio, vertexai, minimax, deepseek, zai, atlas, volcano
|
||||
# Supported providers: openai, groq, ollama, gemini, anthropic, lmstudio, vertexai, minimax, deepseek, zai, volcano
|
||||
HINDSIGHT_API_LLM_PROVIDER=openai
|
||||
HINDSIGHT_API_LLM_API_KEY=your-api-key-here
|
||||
HINDSIGHT_API_LLM_MODEL=gpt-4o-mini
|
||||
HINDSIGHT_API_LLM_BASE_URL=https://api.openai.com/v1
|
||||
# Reasoning effort for providers/models that support it. Examples: low, medium, high, xhigh.
|
||||
# HINDSIGHT_API_LLM_REASONING_EFFORT=low
|
||||
|
||||
# Sampling temperature for internal LLM calls. Set a number in [0.0, 2.0], or `none`
|
||||
# to omit the temperature parameter entirely (required for models that reject explicit
|
||||
# temperatures, e.g. Azure gpt-5.5). The global override below applies to every operation;
|
||||
# per-operation overrides (defaults: verification=0.0, retain=0.1, reflect=0.9,
|
||||
# consolidation=0.0) take precedence.
|
||||
# HINDSIGHT_API_LLM_TEMPERATURE=none
|
||||
# HINDSIGHT_API_LLM_TEMPERATURE_VERIFICATION=0.0
|
||||
# HINDSIGHT_API_LLM_TEMPERATURE_RETAIN=0.1
|
||||
# HINDSIGHT_API_LLM_TEMPERATURE_REFLECT=0.9
|
||||
# HINDSIGHT_API_LLM_TEMPERATURE_CONSOLIDATION=0.0
|
||||
|
||||
# Example: Anthropic Claude configuration
|
||||
# HINDSIGHT_API_LLM_PROVIDER=anthropic
|
||||
@@ -36,7 +23,7 @@ HINDSIGHT_API_LLM_BASE_URL=https://api.openai.com/v1
|
||||
# Example: MiniMax configuration (1M context window)
|
||||
# HINDSIGHT_API_LLM_PROVIDER=minimax
|
||||
# HINDSIGHT_API_LLM_API_KEY=your-minimax-api-key
|
||||
# HINDSIGHT_API_LLM_MODEL=MiniMax-M3 # or MiniMax-M2.7 for the previous generation
|
||||
# HINDSIGHT_API_LLM_MODEL=MiniMax-M2.7
|
||||
|
||||
# Example: DeepSeek configuration (https://api.deepseek.com)
|
||||
# HINDSIGHT_API_LLM_PROVIDER=deepseek
|
||||
@@ -48,41 +35,16 @@ HINDSIGHT_API_LLM_BASE_URL=https://api.openai.com/v1
|
||||
# HINDSIGHT_API_LLM_API_KEY=your-zai-api-key
|
||||
# HINDSIGHT_API_LLM_MODEL=glm-4.5-flash # or glm-4.5-air for the paid tier
|
||||
|
||||
# Example: Atlas Cloud configuration (OpenAI-compatible, https://www.atlascloud.ai)
|
||||
# HINDSIGHT_API_LLM_PROVIDER=atlas
|
||||
# HINDSIGHT_API_LLM_API_KEY=your-atlascloud-api-key
|
||||
# HINDSIGHT_API_LLM_MODEL=deepseek-ai/deepseek-v4-pro # reasoning model; also Qwen / GLM / Kimi / MiniMax, etc.
|
||||
|
||||
# Example: LM Studio local configuration (Qwen 2.5 32B recommended)
|
||||
# HINDSIGHT_API_LLM_PROVIDER=lmstudio
|
||||
# HINDSIGHT_API_LLM_API_KEY=lmstudio
|
||||
# HINDSIGHT_API_LLM_BASE_URL=http://localhost:1234/v1
|
||||
# HINDSIGHT_API_LLM_MODEL=qwen2.5-32b-instruct
|
||||
|
||||
# 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
|
||||
# RETAIN_/REFLECT_/CONSOLIDATION_ prefix (e.g. HINDSIGHT_API_RETAIN_LLM_1_PROVIDER).
|
||||
# HINDSIGHT_API_LLM_1_PROVIDER=groq
|
||||
# HINDSIGHT_API_LLM_1_API_KEY=your-groq-api-key
|
||||
# HINDSIGHT_API_LLM_1_MODEL=openai/gpt-oss-120b
|
||||
# HINDSIGHT_API_LLM_2_PROVIDER=anthropic
|
||||
# HINDSIGHT_API_LLM_2_API_KEY=your-anthropic-api-key
|
||||
# Strategy JSON: {"mode": "failover"} or {"mode": "round-robin"}.
|
||||
# Round-robin accepts optional positive-int "weights" (one per member, primary first).
|
||||
# HINDSIGHT_API_LLM_STRATEGY={"mode": "failover"}
|
||||
|
||||
# API Configuration (Optional)
|
||||
HINDSIGHT_API_HOST=0.0.0.0
|
||||
HINDSIGHT_API_PORT=8888
|
||||
HINDSIGHT_API_LOG_LEVEL=info
|
||||
# Optional retain chunking override for structured logs/transcripts.
|
||||
# Unset uses HINDSIGHT_API_RETAIN_CHUNK_SIZE as the structured-chunk limit.
|
||||
# HINDSIGHT_API_RETAIN_STRUCTURED_CHUNK_SIZE=
|
||||
|
||||
# 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
|
||||
|
||||
# Base Path / Reverse Proxy Support (Optional)
|
||||
# Set these when deploying behind a reverse proxy with path-based routing
|
||||
@@ -95,7 +57,6 @@ 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_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).
|
||||
|
||||
# Vector Extension (Optional - uses pgvector by default)
|
||||
# Options: "pgvector" (default), "vchord", "pgvectorscale" (DiskANN)
|
||||
@@ -103,49 +64,11 @@ HINDSIGHT_API_LOG_LEVEL=info
|
||||
# For Azure PostgreSQL with DiskANN:
|
||||
# HINDSIGHT_API_VECTOR_EXTENSION=pgvectorscale # Auto-detects pg_diskann on Azure
|
||||
|
||||
# Text Search Extension (Optional - uses native PostgreSQL full-text search by default)
|
||||
# Backend options: "native" (default), "vchord", "pg_textsearch", "pgroonga", "pg_search"
|
||||
# HINDSIGHT_API_TEXT_SEARCH_EXTENSION=native
|
||||
# Native backend dictionary (only used by HINDSIGHT_API_TEXT_SEARCH_EXTENSION=native)
|
||||
# HINDSIGHT_API_TEXT_SEARCH_EXTENSION_NATIVE_LANGUAGE=english
|
||||
# ParadeDB pg_search tokenizer (only used when creating pg_search BM25 indexes).
|
||||
# Empty uses ParadeDB's default tokenizer: unicode_words.
|
||||
# Supported values: unicode_words, simple, whitespace, literal, literal_normalized,
|
||||
# chinese_compatible, icu, jieba, source_code,
|
||||
# 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=
|
||||
|
||||
# File Parser (Optional - uses markitdown by default)
|
||||
# HINDSIGHT_API_FILE_PARSER=markitdown
|
||||
# Enable image OCR for MarkItDown using an OpenAI-compatible OCR/vision endpoint.
|
||||
# These OCR settings are independent from HINDSIGHT_API_LLM_* because MarkItDown
|
||||
# uses the OpenAI SDK directly and requires Chat Completions image input support.
|
||||
# When OCR is enabled, API_KEY, BASE_URL, and MODEL are required.
|
||||
# HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_ENABLED=false
|
||||
# HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_API_KEY=
|
||||
# HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_BASE_URL=
|
||||
# HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_MODEL=
|
||||
# HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_PROMPT=
|
||||
|
||||
# Embeddings Configuration (Optional - uses local by default)
|
||||
# Provider: "local" (default), "onnx", "tei", "openai", "cohere", "google", "openrouter", "zeroentropy", "litellm", or "litellm-sdk"
|
||||
# Provider: "local" (default), "tei", "openai", "cohere", "google", "openrouter", "litellm", or "litellm-sdk"
|
||||
# HINDSIGHT_API_EMBEDDINGS_PROVIDER=local
|
||||
# For local provider:
|
||||
# HINDSIGHT_API_EMBEDDINGS_LOCAL_MODEL=BAAI/bge-small-en-v1.5
|
||||
# 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
|
||||
# HINDSIGHT_API_EMBEDDINGS_ONNX_FILE=onnx/model.onnx
|
||||
# HINDSIGHT_API_EMBEDDINGS_ONNX_DIMENSIONS=384
|
||||
# HINDSIGHT_API_EMBEDDINGS_ONNX_MAX_TOKENS=512
|
||||
# HINDSIGHT_API_EMBEDDINGS_ONNX_POOLING=mean
|
||||
# HINDSIGHT_API_EMBEDDINGS_ONNX_NORMALIZE=true
|
||||
# HINDSIGHT_API_EMBEDDINGS_ONNX_QUERY_PREFIX="query: "
|
||||
# HINDSIGHT_API_EMBEDDINGS_ONNX_PASSAGE_PREFIX="passage: "
|
||||
# Optional for local model paths or pre-downloaded artifacts:
|
||||
# HINDSIGHT_API_EMBEDDINGS_ONNX_MODEL_PATH=/models/multilingual-e5-small/onnx/model.onnx
|
||||
# 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
|
||||
# For TEI provider:
|
||||
@@ -154,13 +77,6 @@ HINDSIGHT_API_LOG_LEVEL=info
|
||||
# HINDSIGHT_API_EMBEDDINGS_OPENAI_API_KEY=sk-xxxx
|
||||
# HINDSIGHT_API_EMBEDDINGS_OPENAI_MODEL=text-embedding-3-small
|
||||
# HINDSIGHT_API_EMBEDDINGS_OPENAI_BASE_URL=https://api.openai.com/v1
|
||||
# For ZeroEntropy zembed-1:
|
||||
# HINDSIGHT_API_EMBEDDINGS_PROVIDER=zeroentropy
|
||||
# HINDSIGHT_API_EMBEDDINGS_ZEROENTROPY_API_KEY=ze-xxxx
|
||||
# HINDSIGHT_API_EMBEDDINGS_ZEROENTROPY_MODEL=zembed-1
|
||||
# HINDSIGHT_API_EMBEDDINGS_ZEROENTROPY_DIMENSIONS=1280
|
||||
# HINDSIGHT_API_EMBEDDINGS_ZEROENTROPY_ENCODING_FORMAT=float
|
||||
# HINDSIGHT_API_EMBEDDINGS_ZEROENTROPY_LATENCY=fast
|
||||
#
|
||||
# IMPORTANT: Embedding keys require provider-specific names:
|
||||
# HINDSIGHT_API_EMBEDDINGS_{PROVIDER}_{PARAMETER}
|
||||
@@ -191,10 +107,6 @@ HINDSIGHT_API_LOG_LEVEL=info
|
||||
# Custom service name and environment (optional, defaults: hindsight-api, development)
|
||||
# HINDSIGHT_API_OTEL_SERVICE_NAME=hindsight-production
|
||||
# HINDSIGHT_API_OTEL_DEPLOYMENT_ENVIRONMENT=production
|
||||
#
|
||||
# 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
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Control Plane (Optional)
|
||||
@@ -203,11 +115,6 @@ HINDSIGHT_API_LOG_LEVEL=info
|
||||
# Dataplane API URL - where the CP proxies requests to
|
||||
# HINDSIGHT_CP_DATAPLANE_API_URL=http://localhost:8888
|
||||
|
||||
# Optional: Bearer token the CP sends as `Authorization: Bearer <key>` to the
|
||||
# dataplane API. Required when the API service is auth-protected; omit for a
|
||||
# public/unauthenticated API.
|
||||
# HINDSIGHT_CP_DATAPLANE_API_KEY=your-dataplane-bearer-token
|
||||
|
||||
# Optional: Require a shared access key to view the Control Plane UI.
|
||||
# When set, visitors see a login page and must enter the key before
|
||||
# accessing the dashboard or any /api/* routes (except /api/health).
|
||||
|
||||
@@ -0,0 +1,6 @@
|
||||
version: 2
|
||||
updates:
|
||||
- package-ecosystem: "github-actions"
|
||||
directory: "/"
|
||||
schedule:
|
||||
interval: "weekly"
|
||||
@@ -22,8 +22,6 @@ jobs:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
fetch-depth: 0 # fetch tags so check-released-integrations can see them
|
||||
- uses: actions/setup-node@v6
|
||||
with:
|
||||
node-version: 20
|
||||
|
||||
@@ -23,9 +23,7 @@ on:
|
||||
- retain
|
||||
- recall
|
||||
- recall-with-observations
|
||||
- recall-temporal
|
||||
- consolidation
|
||||
- graph-maintenance
|
||||
default: ""
|
||||
locomo_conversations:
|
||||
description: "LoComo conversation IDs (space-separated). Blank = curated set (conv-26 conv-30 conv-43)."
|
||||
@@ -35,18 +33,6 @@ on:
|
||||
description: "Skip LoComo job"
|
||||
type: boolean
|
||||
default: false
|
||||
obs_skip:
|
||||
description: "Skip observation-dedup benchmark job"
|
||||
type: boolean
|
||||
default: false
|
||||
obs_dataset:
|
||||
description: "Obs benchmark dataset substring (blank = English hermes transcript)."
|
||||
type: string
|
||||
default: ""
|
||||
obs_fraction:
|
||||
description: "Obs benchmark fraction (0-1] of each document to run."
|
||||
type: string
|
||||
default: "1.0"
|
||||
ref:
|
||||
description: "Git ref to test (branch, tag, or SHA). Defaults to main."
|
||||
type: string
|
||||
@@ -212,86 +198,3 @@ jobs:
|
||||
PERF_DASHBOARD_TOKEN: ${{ secrets.PERF_DASHBOARD_TOKEN }}
|
||||
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
run: ./scripts/benchmarks/publish-locomo-results.sh hindsight-dev/benchmarks/locomo/results/benchmark_results.json
|
||||
|
||||
obs:
|
||||
# Observation-dedup quality benchmark: ingests a transcript, drains consolidation
|
||||
# (serial SyncTaskBackend + embedded pg0 — no external DB / worker), and reports the
|
||||
# near-duplicate observation rate. Real LLM via VertexAI, mirroring the LoComo job.
|
||||
if: inputs.obs_skip != true
|
||||
runs-on: ubuntu-latest
|
||||
env:
|
||||
HINDSIGHT_API_LLM_PROVIDER: vertexai
|
||||
HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY: /tmp/gcp-credentials.json
|
||||
HINDSIGHT_API_LLM_MODEL: google/gemini-2.5-flash-lite
|
||||
HINDSIGHT_API_ENABLE_OBSERVATIONS: "true"
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ inputs.ref || github.ref }}
|
||||
|
||||
- name: Setup GCP credentials
|
||||
run: |
|
||||
printf '%s' '${{ secrets.GCP_VERTEXAI_CREDENTIALS }}' > /tmp/gcp-credentials.json
|
||||
PROJECT_ID=$(jq -r '.project_id' /tmp/gcp-credentials.json)
|
||||
echo "HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID=$PROJECT_ID" >> $GITHUB_ENV
|
||||
|
||||
- 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: Cache HuggingFace models
|
||||
uses: actions/cache@v5
|
||||
with:
|
||||
path: ~/.cache/huggingface
|
||||
key: ${{ runner.os }}-huggingface-${{ hashFiles('hindsight-api-slim/pyproject.toml') }}
|
||||
restore-keys: |
|
||||
${{ runner.os }}-huggingface-
|
||||
|
||||
- name: Pre-download models
|
||||
working-directory: ./hindsight-api-slim
|
||||
run: |
|
||||
uv run --frozen --all-extras --index-strategy unsafe-best-match python -c "
|
||||
from sentence_transformers import SentenceTransformer
|
||||
print('Downloading embedding model...')
|
||||
SentenceTransformer('BAAI/bge-small-en-v1.5')
|
||||
print('Model downloaded successfully')
|
||||
"
|
||||
|
||||
- name: Install hindsight-dev dependencies
|
||||
run: |
|
||||
cd hindsight-dev && uv sync --frozen --all-extras --index-strategy unsafe-best-match
|
||||
|
||||
- name: Run obs benchmark
|
||||
# Default to the English hermes transcript at full fraction — a clean, deterministic
|
||||
# consolidation-dedup signal (the Chinese variant adds a cross-lingual embedding
|
||||
# confound). Override dataset/fraction via workflow_dispatch.
|
||||
run: |
|
||||
DATASET="${{ inputs.obs_dataset }}"
|
||||
if [ -z "$DATASET" ]; then DATASET="hermes_session_2026-05-15_en"; fi
|
||||
FRACTION="${{ inputs.obs_fraction }}"
|
||||
if [ -z "$FRACTION" ]; then FRACTION="1.0"; fi
|
||||
cd hindsight-dev
|
||||
uv run python -m benchmarks.obs.obs_benchmark \
|
||||
--dataset "$DATASET" --fraction "$FRACTION" --wipe-bank --output obs-results.json
|
||||
|
||||
- name: Upload obs results
|
||||
if: always()
|
||||
uses: actions/upload-artifact@v7
|
||||
with:
|
||||
name: obs-results-${{ github.sha }}
|
||||
path: hindsight-dev/obs-results.json
|
||||
retention-days: 90
|
||||
|
||||
- name: Publish obs to dashboard
|
||||
if: success() && (github.event_name == 'schedule' || github.event_name == 'push' || github.event_name == 'workflow_dispatch')
|
||||
env:
|
||||
PERF_DASHBOARD_TOKEN: ${{ secrets.PERF_DASHBOARD_TOKEN }}
|
||||
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
run: ./scripts/benchmarks/publish-obs-results.sh hindsight-dev/obs-results.json
|
||||
|
||||
@@ -9,11 +9,7 @@ jobs:
|
||||
publish:
|
||||
runs-on: ubuntu-latest
|
||||
permissions:
|
||||
id-token: write # for PyPI trusted publishing + build-provenance attestations
|
||||
attestations: write # for actions/attest-build-provenance (Obsidian assets)
|
||||
# No `contents: write`: we never create releases in this repo. The Obsidian
|
||||
# plugin's distribution release is pushed to its dedicated repo using
|
||||
# OBSIDIAN_DIST_TOKEN (see the "Mirror Obsidian plugin" step below).
|
||||
id-token: write # for PyPI trusted publishing
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
@@ -75,7 +71,7 @@ jobs:
|
||||
if: steps.type.outputs.type == 'plugin'
|
||||
run: |
|
||||
echo "Plugin integration ${{ steps.info.outputs.integration }} v${{ steps.info.outputs.version }} — no package to publish."
|
||||
echo "Users install via: claude plugin marketplace add vectorize-io/hindsight"
|
||||
echo "Users install via: claude plugin marketplace add vectorize-io/hindsight --sparse hindsight-integrations"
|
||||
|
||||
# ── TypeScript integrations (ai-sdk, chat, openclaw) ────────────────────
|
||||
|
||||
@@ -116,71 +112,6 @@ jobs:
|
||||
working-directory: ./hindsight-integrations/${{ steps.info.outputs.integration }}
|
||||
run: npm run build
|
||||
|
||||
# Build-provenance attestations for the Obsidian release assets (community-store
|
||||
# recommendation). Runs after the build so main.js exists. The assets are
|
||||
# released in the dedicated repo while the build runs here, so users verify at
|
||||
# owner scope: `gh attestation verify main.js --owner vectorize-io`.
|
||||
- name: Attest Obsidian plugin build provenance
|
||||
if: steps.type.outputs.type == 'typescript' && steps.info.outputs.integration == 'obsidian'
|
||||
uses: actions/attest-build-provenance@v2
|
||||
with:
|
||||
subject-path: |
|
||||
hindsight-integrations/obsidian/main.js
|
||||
hindsight-integrations/obsidian/styles.css
|
||||
|
||||
# ── Obsidian plugin — mirror to its dedicated repo + cut the BRAT release ──
|
||||
# We do NOT create a GitHub Release in this monorepo: per-integration
|
||||
# releases pollute the repo's release list (it's for the core product) and
|
||||
# steal the "Latest" badge, and BRAT / the community store read a repo's
|
||||
# *latest* release — not a tag — so they can't target a tag in a monorepo.
|
||||
#
|
||||
# Instead this monorepo stays the source of truth, and on each obsidian
|
||||
# release we mirror hindsight-integrations/obsidian/ → the *root* of
|
||||
# github.com/vectorize-io/hindsight-obsidian (git subtree, history
|
||||
# preserved) and cut the BRAT / community-store release *there*.
|
||||
#
|
||||
# Requires secret OBSIDIAN_DIST_TOKEN — a token with `contents: write` on
|
||||
# vectorize-io/hindsight-obsidian (fine-grained PAT or app installation
|
||||
# token). The dedicated repo is generated; never edit it directly.
|
||||
- name: Mirror Obsidian plugin to its dedicated repo
|
||||
if: steps.type.outputs.type == 'typescript' && steps.info.outputs.integration == 'obsidian'
|
||||
env:
|
||||
DIST_TOKEN: ${{ secrets.OBSIDIAN_DIST_TOKEN }}
|
||||
run: |
|
||||
set -euo pipefail
|
||||
VERSION="${{ steps.info.outputs.version }}"
|
||||
DIST_REPO="vectorize-io/hindsight-obsidian"
|
||||
OBS_DIR="hindsight-integrations/obsidian"
|
||||
|
||||
# `git subtree split` needs full history; the default checkout is shallow.
|
||||
git fetch --unshallow 2>/dev/null || true
|
||||
git config user.name "github-actions[bot]"
|
||||
git config user.email "41898282+github-actions[bot]@users.noreply.github.com"
|
||||
|
||||
# The runner injects the default GITHUB_TOKEN as an http.extraheader via
|
||||
# an *included* config file (/home/runner/work/_temp/git-credentials-*.config),
|
||||
# so `git config --local --unset-all` can't remove it and it authenticates
|
||||
# the push as github-actions[bot] (no access to the dedicated repo → 403).
|
||||
# The documented way to drop an inherited extraheader is to RESET the list
|
||||
# with an empty value: since command-line `-c` is read last, the empty
|
||||
# value clears the accumulated headers (including the included one) at
|
||||
# request-build time. The dist token then comes from the push URL → a
|
||||
# single Authorization header.
|
||||
git subtree split --prefix="$OBS_DIR" -b _obs_dist
|
||||
git -c "http.https://github.com/.extraheader=" \
|
||||
push "https://x-access-token:${DIST_TOKEN}@github.com/${DIST_REPO}.git" _obs_dist:main
|
||||
|
||||
# Cut the BRAT / community-store release. Bare version tag (e.g. 0.1.0)
|
||||
# to match manifest.json — idempotent so re-runs just refresh the assets.
|
||||
export GH_TOKEN="$DIST_TOKEN"
|
||||
ASSETS="$OBS_DIR/main.js $OBS_DIR/manifest.json $OBS_DIR/styles.css"
|
||||
NOTES="Hindsight for Obsidian v${VERSION}. Install via BRAT (add ${DIST_REPO}) or copy main.js/manifest.json/styles.css into <vault>/.obsidian/plugins/hindsight/."
|
||||
if gh release view "$VERSION" --repo "$DIST_REPO" >/dev/null 2>&1; then
|
||||
gh release upload "$VERSION" $ASSETS --repo "$DIST_REPO" --clobber
|
||||
else
|
||||
gh release create "$VERSION" $ASSETS --repo "$DIST_REPO" --title "$VERSION" --notes "$NOTES"
|
||||
fi
|
||||
|
||||
- name: Publish TypeScript package to npm
|
||||
if: steps.type.outputs.type == 'typescript'
|
||||
working-directory: ./hindsight-integrations/${{ steps.info.outputs.integration }}
|
||||
@@ -190,12 +121,7 @@ jobs:
|
||||
EXIT_CODE=$?
|
||||
echo "$OUTPUT"
|
||||
if [ $EXIT_CODE -ne 0 ]; then
|
||||
# Treat "already published" as success so re-pointed-tag re-runs stay green.
|
||||
# "cannot publish over" = the version exists. TLOG_CREATE_ENTRY_ERROR / 409
|
||||
# "equivalent entry already exists in the transparency log" = the identical
|
||||
# --provenance artifact was already logged on a prior run (Sigstore tlog is
|
||||
# idempotent); the package is published, so this is benign.
|
||||
if echo "$OUTPUT" | grep -qE "cannot publish over|TLOG_CREATE_ENTRY_ERROR|already exists in the transparency log"; then
|
||||
if echo "$OUTPUT" | grep -q "cannot publish over"; then
|
||||
echo "Package version already published, skipping..."
|
||||
exit 0
|
||||
fi
|
||||
|
||||
@@ -266,7 +266,7 @@ jobs:
|
||||
strategy:
|
||||
matrix:
|
||||
include:
|
||||
- os: ubuntu-22.04
|
||||
- os: ubuntu-latest
|
||||
target: x86_64-unknown-linux-gnu
|
||||
artifact_name: hindsight
|
||||
asset_name: hindsight-linux-amd64
|
||||
@@ -278,7 +278,7 @@ jobs:
|
||||
target: aarch64-apple-darwin
|
||||
artifact_name: hindsight
|
||||
asset_name: hindsight-darwin-arm64
|
||||
- os: ubuntu-22.04-arm
|
||||
- os: ubuntu-24.04-arm
|
||||
target: aarch64-unknown-linux-gnu
|
||||
artifact_name: hindsight
|
||||
asset_name: hindsight-linux-arm64
|
||||
|
||||
+12
-1303
File diff suppressed because it is too large
Load Diff
@@ -1,115 +0,0 @@
|
||||
name: Windows Smoke Test
|
||||
|
||||
# Daily smoke test that installs the API on Windows and runs the Python client
|
||||
# integration tests against a live server. Windows is only exercised by the
|
||||
# hindsight-embed jobs in test.yml on PRs; this catches Windows-specific
|
||||
# regressions in the API server + client path (e.g. process spawning, console
|
||||
# subsystem / ConPTY behaviour, see #1885) that the Linux client jobs miss.
|
||||
on:
|
||||
schedule:
|
||||
# 06:00 UTC daily.
|
||||
- cron: "0 6 * * *"
|
||||
workflow_dispatch:
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
jobs:
|
||||
windows-client-smoke:
|
||||
# Don't run on forks: the job needs the org's Vertex AI credentials.
|
||||
if: github.repository == 'vectorize-io/hindsight'
|
||||
runs-on: windows-latest
|
||||
timeout-minutes: 45
|
||||
env:
|
||||
HINDSIGHT_API_LLM_PROVIDER: vertexai
|
||||
HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY: ${{ github.workspace }}/gcp-credentials.json
|
||||
HINDSIGHT_API_LLM_MODEL: google/gemini-2.5-flash-lite
|
||||
HINDSIGHT_API_URL: http://localhost:8888
|
||||
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
# Force UTF-8 I/O so the API/CLI's ✓/box-drawing output doesn't crash the
|
||||
# default Windows cp1252 codec (matches test-embed-windows in test.yml).
|
||||
PYTHONIOENCODING: utf-8
|
||||
PYTHONUTF8: "1"
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
|
||||
- name: Setup GCP credentials
|
||||
shell: bash
|
||||
run: |
|
||||
printf '%s' '${{ secrets.GCP_VERTEXAI_CREDENTIALS }}' > gcp-credentials.json
|
||||
PROJECT_ID=$(jq -r '.project_id' gcp-credentials.json)
|
||||
echo "HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID=$PROJECT_ID" >> "$GITHUB_ENV"
|
||||
|
||||
- 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: Install API dependencies (all extras - local-ml + embedded pg0)
|
||||
working-directory: ./hindsight-api-slim
|
||||
run: uv sync --frozen --all-extras --index-strategy unsafe-best-match
|
||||
|
||||
- name: Install Python client test dependencies
|
||||
working-directory: ./hindsight-clients/python
|
||||
run: uv sync --frozen --extra test --index-strategy unsafe-best-match
|
||||
|
||||
- name: Cache HuggingFace models
|
||||
uses: actions/cache@v5
|
||||
with:
|
||||
path: ~/.cache/huggingface
|
||||
key: ${{ runner.os }}-huggingface-${{ hashFiles('hindsight-api-slim/pyproject.toml') }}
|
||||
restore-keys: |
|
||||
${{ runner.os }}-huggingface-
|
||||
|
||||
# `uv run` re-syncs the project env to its default (no-extras) state before
|
||||
# running, which drops sentence-transformers / pg0. Pass --all-extras on
|
||||
# every `uv run` so the local-ml + embedded-db deps stay installed (this is
|
||||
# the same reason hindsight-embed launches the daemon with `--extra all`).
|
||||
- name: Pre-download models
|
||||
working-directory: ./hindsight-api-slim
|
||||
run: |
|
||||
uv run --all-extras python -c "from sentence_transformers import SentenceTransformer, CrossEncoder; SentenceTransformer('BAAI/bge-small-en-v1.5'); CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2'); print('Models downloaded')"
|
||||
|
||||
# Start the server and run the client tests in a SINGLE step. On Windows
|
||||
# runners a process backgrounded with `&` in one step is not reliably kept
|
||||
# alive for later steps (unlike Linux, where it reparents to init), so the
|
||||
# server must live in the same shell that runs pytest.
|
||||
- name: Start API server and run Python client tests
|
||||
shell: bash
|
||||
run: |
|
||||
# Config is read straight from the environment (job-level env + the
|
||||
# PROJECT_ID exported to GITHUB_ENV above), so no .env file is needed.
|
||||
# Embedded pg0 is the default when HINDSIGHT_API_DATABASE_URL is unset.
|
||||
( cd hindsight-api-slim && uv run --all-extras hindsight-api --port 8888 ) > "$RUNNER_TEMP/api-server.log" 2>&1 &
|
||||
server_pid=$!
|
||||
echo "Waiting for API server to be ready (pid $server_pid)..."
|
||||
# pg0 unpacks Postgres + runs initdb on first boot, which is slow on a
|
||||
# cold Windows runner — give it a generous budget before failing.
|
||||
ready=false
|
||||
for i in $(seq 1 300); do
|
||||
if curl -sf http://localhost:8888/health > /dev/null 2>&1; then
|
||||
echo "API server is ready after ${i}s"
|
||||
ready=true
|
||||
break
|
||||
fi
|
||||
sleep 1
|
||||
done
|
||||
if [ "$ready" != true ]; then
|
||||
echo "API server failed to start after 300s"
|
||||
cat "$RUNNER_TEMP/api-server.log"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
cd hindsight-clients/python && uv run --extra test pytest tests -v
|
||||
|
||||
- name: Show API server logs
|
||||
if: always()
|
||||
shell: bash
|
||||
run: cat "$RUNNER_TEMP/api-server.log" || echo "No API server log found"
|
||||
+1
-7
@@ -6,7 +6,6 @@ dist/
|
||||
wheels/
|
||||
*.egg-info
|
||||
.mcp.json
|
||||
.playwright-mcp/
|
||||
.osgrep
|
||||
# Virtual environments
|
||||
.venv
|
||||
@@ -16,8 +15,6 @@ node_modules/
|
||||
|
||||
# Environment variables and local config
|
||||
.env
|
||||
.env.bak*
|
||||
.env.*.bak
|
||||
docker-compose.yml
|
||||
docker-compose.override.yml
|
||||
|
||||
@@ -57,10 +54,7 @@ hindsight-clients/rust/target
|
||||
!.claude/skills/
|
||||
whats-next.md
|
||||
TASK.md
|
||||
# Parked / draft integrations that aren't ready to ship
|
||||
hindsight-integrations/_drafts/
|
||||
# Changelog is now tracked in hindsight-docs/src/pages/changelog.md
|
||||
# CHANGELOG.md
|
||||
|
||||
blog-post*
|
||||
.worktrees/
|
||||
blog-post*
|
||||
@@ -216,46 +216,10 @@ migration file dispatches through `run_for_dialect`, which calls either
|
||||
./scripts/hooks/lint.sh
|
||||
```
|
||||
|
||||
Dead-code detection runs in CI (the `check-unused-code` job) at two levels:
|
||||
- **Blocking:** unused imports (ruff `F401`) and variables (`F841`) — `lint.sh` auto-removes
|
||||
them and `verify-generated-files` fails on any leftover diff; and **knip** for orphaned
|
||||
control-plane files / unused (or unlisted) `package.json` dependencies.
|
||||
- **Advisory:** whole unused Python functions (vulture) and unused control-plane *exports*
|
||||
(the shadcn/ui surface is kept on purpose) — surfaced, not gated.
|
||||
|
||||
Run both locally with:
|
||||
```bash
|
||||
./scripts/hooks/check-unused.sh
|
||||
```
|
||||
|
||||
**After completing any implementation work, run `/code-review`** to verify your changes against project standards (missing tests, dead code, type safety, etc.). Fix any "must fix" issues before considering the task done.
|
||||
|
||||
**MANDATORY: Run `/code-review` before pushing code or creating a pull request.** Do not push or create a PR until all "must fix" issues are resolved.
|
||||
|
||||
### Testing
|
||||
|
||||
Most tests are deterministic (MockLLM, pure functions) — assert directly.
|
||||
|
||||
**Tests that verify LLM behaviour use a real LLM + an LLM-as-judge.** When the thing under test is *how the model interprets a prompt* (classification, attribution, dimension preservation, instruction-following), MockLLM can't simulate it and exact string/enum asserts flake across providers and runs. Use this pattern instead:
|
||||
|
||||
1. Mark the test module `pytestmark = pytest.mark.hs_llm_core` (single-provider; CI runs it in the core-LLM job). Use `hs_llm_mat` only for provider-matrix acceptance tests.
|
||||
2. Call the real pipeline (`LLMConfig.from_env()`, `_get_raw_config()`), e.g. `extract_facts_from_text(...)`.
|
||||
3. Assert with the judge, not string matching:
|
||||
```python
|
||||
from tests.llm_judge import assert_meets_criteria
|
||||
facts_summary = "\n".join(f"- [{f.fact_type}] {f.fact}" for f in facts)
|
||||
await assert_meets_criteria(
|
||||
response=facts_summary,
|
||||
criteria="The first-person user statements are classified 'world' and attributed to the user, not the agent.",
|
||||
context="What the input said and who was speaking.",
|
||||
)
|
||||
```
|
||||
|
||||
Rules of thumb:
|
||||
- **Judge anything non-deterministic** — including `fact_type` classification and speaker attribution. Do NOT hard-assert `fact_type == "..."`; pass a `[fact_type] fact` summary to the judge instead. Structural facts that ARE deterministic (counts, presence of a field, that a substring was injected into a prompt) stay as direct asserts in fast unit tests.
|
||||
- **Split the test surface**: cover the deterministic mechanics (prompt assembly, suppression logic) with fast non-LLM unit tests, and the model-following behaviour with one `hs_llm_core` judge test. (Example pair: `test_narrator_resolution.py` + `test_narrator_context_override.py`.)
|
||||
- The judge model is independent of the test provider (defaults to Gemini); never judge with the same call you're testing.
|
||||
|
||||
### Memory Banks
|
||||
- Each bank is an isolated memory store (like a "brain" for one user/agent)
|
||||
- Banks have dispositions (skepticism, literalism, empathy traits 1-5) affecting reflect
|
||||
@@ -327,10 +291,7 @@ Fields must be categorized as either **hierarchical** (can be overridden per-ten
|
||||
```
|
||||
|
||||
2. **main.py** (`hindsight-api-slim/hindsight_api/main.py`):
|
||||
- No change is needed for ordinary environment-backed config fields. The CLI starts from `_get_raw_config()`,
|
||||
so new `HindsightConfig` fields are carried through automatically.
|
||||
- If the new field should be overridable by a CLI flag, add the argparse option in `_parse_cli_args()` and include
|
||||
that field in the `dataclasses.replace(config, ...)` call near the "CLI override" comment.
|
||||
- Add field to the manual `HindsightConfig()` constructor call (search for "CLI override")
|
||||
|
||||
3. **Use hierarchical config in MemoryEngine**:
|
||||
```python
|
||||
@@ -350,16 +311,6 @@ Fields must be categorized as either **hierarchical** (can be overridden per-ten
|
||||
- Add to appropriate section table with Variable, Description, Default
|
||||
- Mark if it's hierarchical (can be overridden per-bank)
|
||||
|
||||
6. **Env template** (`.env.example`):
|
||||
- Add the variable to the appropriate section, commented if optional, with a
|
||||
short inline comment describing it (mirror the documentation entry).
|
||||
- This file is the single source of truth for the env template:
|
||||
`scripts/dev/setup.sh` copies it to `.env`, and `hindsight-embed` ships a
|
||||
bundled copy (`hindsight-embed/hindsight_embed/env.example`) that seeds
|
||||
embed/profile configs. After editing `.env.example`, re-copy it to the
|
||||
embed package (`cp .env.example hindsight-embed/hindsight_embed/env.example`)
|
||||
or the `test_bundled_template_matches_repo_root` sync test will fail.
|
||||
|
||||
#### Hierarchical vs Static Guidelines
|
||||
|
||||
**Hierarchical** (per-bank overridable):
|
||||
@@ -376,7 +327,7 @@ Fields must be categorized as either **hierarchical** (can be overridden per-ten
|
||||
|
||||
```bash
|
||||
cp .env.example .env
|
||||
# Edit .env with the LLM provider/model and credentials for your setup
|
||||
# Edit .env with LLM API key
|
||||
|
||||
# Python deps
|
||||
uv sync --directory hindsight-api-slim/
|
||||
@@ -385,10 +336,10 @@ uv sync --directory hindsight-api-slim/
|
||||
npm install
|
||||
```
|
||||
|
||||
Common LLM settings:
|
||||
Required env vars:
|
||||
- `HINDSIGHT_API_LLM_PROVIDER`: openai, anthropic, gemini, groq, minimax, ollama, lmstudio
|
||||
- `HINDSIGHT_API_LLM_API_KEY`: API key for providers that require one
|
||||
- `HINDSIGHT_API_LLM_MODEL`: Model name (defaults are provider-specific)
|
||||
- `HINDSIGHT_API_LLM_API_KEY`: Your API key
|
||||
- `HINDSIGHT_API_LLM_MODEL`: Model name (e.g., gpt-4o-mini, claude-sonnet-4-20250514)
|
||||
|
||||
Optional (uses local models by default):
|
||||
- `HINDSIGHT_API_EMBEDDINGS_PROVIDER`: local (default) or tei
|
||||
|
||||
+2
-25
@@ -9,36 +9,13 @@ Thanks for your interest in contributing to Hindsight!
|
||||
git clone [email protected]:vectorize-io/hindsight.git
|
||||
cd hindsight
|
||||
```
|
||||
|
||||
2. Bootstrap your dev environment in one shot:
|
||||
```bash
|
||||
./scripts/dev/setup.sh
|
||||
```
|
||||
This is idempotent (safe to re-run) and gets you ready to develop, including
|
||||
offline. It:
|
||||
- installs the required toolchains if missing (uv/Python, Node/npm, Rust/cargo),
|
||||
- creates `.env` from `.env.example` (remember to add your LLM API key),
|
||||
- configures git hooks,
|
||||
- installs all Python and Node workspace dependencies,
|
||||
- pre-downloads the local ML models + tokenizer so the API runs offline,
|
||||
- builds the TypeScript SDK and the Rust CLI.
|
||||
|
||||
Useful flags: `--skip-build` (deps only), `--skip-models` (skip ML model
|
||||
download), `--with-docs` (also build the docs site), `--force` (rebuild
|
||||
artifacts). Docker image builds are out of scope. Run
|
||||
`./scripts/dev/setup.sh --help` for details.
|
||||
|
||||
### Manual setup
|
||||
|
||||
If you'd rather set things up by hand instead of running the script above:
|
||||
|
||||
1. Set up your environment:
|
||||
2. Set up your environment:
|
||||
```bash
|
||||
cp .env.example .env
|
||||
```
|
||||
Edit the .env to add LLM API key and config as required
|
||||
|
||||
2. Install dependencies:
|
||||
3. Install dependencies:
|
||||
```bash
|
||||
# Python dependencies
|
||||
uv sync --directory hindsight-api/
|
||||
|
||||
@@ -7,6 +7,7 @@
|
||||
[](https://github.com/vectorize-io/hindsight/actions/workflows/release.yml)
|
||||
[](https://join.slack.com/t/hindsight-space/shared_invite/zt-3nhbm4w29-LeSJ5Ixi6j8PdiYOCPlOgg)
|
||||
[](https://opensource.org/licenses/MIT)
|
||||
[](https://gitcgr.com/vectorize-io/hindsight)
|
||||

|
||||

|
||||
<br/>
|
||||
@@ -61,16 +62,16 @@ If you need more control over how and when your agent stores and recalls memorie
|
||||
```bash
|
||||
export OPENAI_API_KEY=sk-xxx
|
||||
|
||||
docker run -it --pull always --name hindsight --restart unless-stopped -p 8888:8888 -p 9999:9999 \
|
||||
docker run --rm -it --pull always -p 8888:8888 -p 9999:9999 \
|
||||
-e HINDSIGHT_API_LLM_API_KEY=$OPENAI_API_KEY \
|
||||
-v hindsight-data:/home/hindsight/.pg0 \
|
||||
-v $HOME/.hindsight-docker:/home/hindsight/.pg0 \
|
||||
ghcr.io/vectorize-io/hindsight:latest
|
||||
```
|
||||
|
||||
>API: http://localhost:8888
|
||||
>UI: http://localhost:9999
|
||||
|
||||
You can modify the LLM provider by setting `HINDSIGHT_API_LLM_PROVIDER`. Valid options are `openai`, `anthropic`, `gemini`, `groq`, `ollama`, `lmstudio`, `minimax`, and `atlas` ([Atlas Cloud](https://www.atlascloud.ai/?utm_source=github&utm_medium=link&utm_campaign=hindsight)). The documentation provides more details on [supported models](https://hindsight.vectorize.io/developer/models).
|
||||
You can modify the LLM provider by setting `HINDSIGHT_API_LLM_PROVIDER`. Valid options are `openai`, `anthropic`, `gemini`, `groq`, `ollama`, `lmstudio`, and `minimax`. The documentation provides more details on [supported models](https://hindsight.vectorize.io/developer/models).
|
||||
|
||||
|
||||
|
||||
@@ -142,8 +143,6 @@ main();
|
||||
pip install hindsight-all -U
|
||||
```
|
||||
|
||||
On Intel (x86_64) Macs, install `hindsight-all-slim` instead — see [Supported Platforms](#supported-platforms).
|
||||
|
||||
```python
|
||||
import os
|
||||
from hindsight import HindsightServer, HindsightClient
|
||||
@@ -250,7 +249,7 @@ Recall performs 4 retrieval strategies in parallel:
|
||||
- Graph: Entity/temporal/causal links
|
||||
- Temporal: Time range filtering
|
||||
|
||||

|
||||

|
||||
|
||||
The individual results from the retrievals are merged, then ordered by relevance using reciprocal rank fusion and a cross-encoder reranking model.
|
||||
|
||||
@@ -276,7 +275,7 @@ client = Hindsight(base_url="http://localhost:8888")
|
||||
client.reflect(bank_id="my-bank", query="What should I know about Alice?")
|
||||
```
|
||||
|
||||

|
||||

|
||||
|
||||
---
|
||||
|
||||
@@ -301,19 +300,6 @@ client.reflect(bank_id="my-bank", query="What should I know about Alice?")
|
||||
[](https://www.star-history.com/#vectorize-io/hindsight&type=date&legend=top-left)
|
||||
---
|
||||
|
||||
## Supported Platforms
|
||||
|
||||
| Platform | Docker | Bare Metal (pip) | Embedded DB (pg0) |
|
||||
|----------|--------|------------------|--------------------|
|
||||
| **Linux** (x86_64, ARM64) | ✅ | ✅ | ✅ |
|
||||
| **macOS** (Apple Silicon / arm64) | ✅ | ✅ | ✅ |
|
||||
| **macOS** (Intel / x86_64) | ✅ | ⚠️ | ✅ |
|
||||
| **Windows** (x86_64) | ✅ | ✅ | ✅ |
|
||||
|
||||
⚠️ Intel Macs: use `hindsight-all-slim` — see the [installation guide](https://hindsight.vectorize.io/developer/installation#supported-platforms) for details.
|
||||
|
||||
---
|
||||
|
||||
## Contributing
|
||||
|
||||
See [CONTRIBUTING.md](./CONTRIBUTING.md).
|
||||
|
||||
@@ -77,6 +77,7 @@
|
||||
"npm:@radix-ui/react-dropdown-menu@^2.1.16",
|
||||
"npm:@radix-ui/react-label@^2.1.8",
|
||||
"npm:@radix-ui/react-popover@^1.1.15",
|
||||
"npm:@radix-ui/react-radio-group@^1.3.8",
|
||||
"npm:@radix-ui/react-select@^2.2.6",
|
||||
"npm:@radix-ui/react-slider@^1.3.6",
|
||||
"npm:@radix-ui/react-slot@^1.2.4",
|
||||
|
||||
@@ -1,113 +0,0 @@
|
||||
# Hindsight with Claude Code (Claude Pro/Max subscription)
|
||||
|
||||
Run Hindsight inside Docker using the `claude-code` LLM provider, backed by
|
||||
your host machine's Claude Pro or Max subscription credentials.
|
||||
|
||||
The standalone Hindsight Docker image ships `claude-agent-sdk` but does **not**
|
||||
bundle the host `claude` CLI binary or any Claude credentials. This Compose
|
||||
file bind-mounts the host's CLI install and credentials into the container so
|
||||
the `claude-code` provider works without an API key.
|
||||
|
||||
## When to use this
|
||||
|
||||
- You have an active Claude Pro or Max subscription and want to use it for
|
||||
Hindsight without paying separate Anthropic API costs.
|
||||
- You want a one-command `docker compose up` instead of a long `docker run`
|
||||
invocation with many flags.
|
||||
- You are running on **Linux/amd64** — macOS Docker Desktop and Windows host
|
||||
paths differ and are not yet covered (please open an issue if you'd like to
|
||||
contribute a verified recipe for either).
|
||||
|
||||
> **Personal-use only.** Anthropic's
|
||||
> [Agent SDK documentation](https://docs.claude.com/en/api/agent-sdk/overview)
|
||||
> states that third-party developers should not offer claude.ai login or rate
|
||||
> limits for their products. Hindsight does **not** perform any login on your
|
||||
> behalf — it uses credentials you've already authenticated via
|
||||
> `claude auth login`. In January 2026, Anthropic
|
||||
> [enforced restrictions](https://paddo.dev/blog/anthropic-walled-garden-crackdown/)
|
||||
> against tools that spoofed the Claude Code client identity; Hindsight uses
|
||||
> the official Claude Agent SDK instead.
|
||||
>
|
||||
> Do not deploy this configuration to shared environments or production. For
|
||||
> that, use the `anthropic` provider with an API key from the
|
||||
> [Anthropic Console](https://console.anthropic.com/). Usage counts against
|
||||
> your Claude Pro/Max subscription limits.
|
||||
|
||||
## Prerequisites
|
||||
|
||||
- Host has `claude` CLI installed (e.g., `npm install -g @anthropics/claude-code`)
|
||||
and `claude auth login` has been run successfully.
|
||||
- `~/.claude.json` and `~/.claude/.credentials.json` exist on the host.
|
||||
- Host `claude` CLI version is **2.1.128 or newer** — the version bundled with
|
||||
`claude-agent-sdk` 0.5.x has a protocol incompatibility in containers, so
|
||||
the recipe overrides it with the host binary.
|
||||
|
||||
## Quick start
|
||||
|
||||
```bash
|
||||
# Set your host UID/GID (defaults to 1000:1000 if unset)
|
||||
export HOST_UID=$(id -u)
|
||||
export HOST_GID=$(id -g)
|
||||
|
||||
docker compose -f docker/docker-compose/claude-code/docker-compose.yaml up -d
|
||||
```
|
||||
|
||||
- API: http://localhost:8888
|
||||
- Control Plane: http://localhost:9999
|
||||
|
||||
## Post-setup (one-time)
|
||||
|
||||
After the container starts for the first time, run these commands to fix
|
||||
permissions and symlink the host `claude` binary into `$PATH`:
|
||||
|
||||
```bash
|
||||
# Make ~/.claude writable by your UID (the CLI writes session/project state)
|
||||
docker exec --user 0:0 hindsight-claude-code chown $(id -u):$(id -g) /home/hindsight/.claude
|
||||
docker exec --user 0:0 hindsight-claude-code chmod 755 /home/hindsight/.claude
|
||||
|
||||
# Symlink the host claude binary into PATH
|
||||
docker exec --user 0:0 hindsight-claude-code \
|
||||
ln -sf /home/hindsight/.local/share/claude/versions/2.1.128 /usr/local/bin/claude
|
||||
```
|
||||
|
||||
If you set `CLAUDE_CLI_VERSION` to a version other than `2.1.128`, update the
|
||||
symlink path accordingly.
|
||||
|
||||
## Notes on the bind-mount surface (every flag is load-bearing)
|
||||
|
||||
- **Host `claude` binary required** — the image ships only `claude-agent-sdk`,
|
||||
not the CLI itself.
|
||||
- **SDK bundled-binary override** — the override of
|
||||
`claude_agent_sdk/_bundled/claude` works around a protocol issue in the
|
||||
bundled v2.1.121 binary inside containers. Once `claude-agent-sdk` ships
|
||||
with v2.1.128+ this override can be dropped. Set `CLAUDE_CLI_VERSION` to
|
||||
match your installed version.
|
||||
- **Single-file credential mounts** — credentials are mounted as individual
|
||||
`:ro` files rather than a whole-directory `:ro` mount of `~/.claude`,
|
||||
because the CLI writes session/project state at runtime and a read-only
|
||||
directory mount silently breaks it.
|
||||
- **`--user` / `user:`** — the `user: ${HOST_UID}:${HOST_GID}` pattern
|
||||
requires `chmod 755 /home/hindsight`, which is built into the image since
|
||||
v0.6.0 (see [#1481](https://github.com/vectorize-io/hindsight/issues/1481)).
|
||||
- **`~/.hindsight-docker` data directory** — the pg0 data bind mount must be
|
||||
writable by your host UID (see
|
||||
[#1483](https://github.com/vectorize-io/hindsight/issues/1483)).
|
||||
- **Verified** on `linux/amd64` against `ghcr.io/vectorize-io/hindsight:latest`
|
||||
v0.5.6+.
|
||||
|
||||
## Using a different Claude CLI version
|
||||
|
||||
If your host has a `claude` version other than 2.1.128, set
|
||||
`CLAUDE_CLI_VERSION` before starting:
|
||||
|
||||
```bash
|
||||
export CLAUDE_CLI_VERSION=2.2.0
|
||||
docker compose -f docker/docker-compose/claude-code/docker-compose.yaml up -d
|
||||
```
|
||||
|
||||
Then update the post-setup symlink to match:
|
||||
|
||||
```bash
|
||||
docker exec --user 0:0 hindsight-claude-code \
|
||||
ln -sf /home/hindsight/.local/share/claude/versions/2.2.0 /usr/local/bin/claude
|
||||
```
|
||||
@@ -1,44 +0,0 @@
|
||||
name: hindsight-claude-code
|
||||
# Run Hindsight with the claude-code LLM provider, using your host machine's
|
||||
# Claude Pro/Max subscription credentials. Linux/amd64 only for now.
|
||||
#
|
||||
# Quick start:
|
||||
# docker compose -f docker/docker-compose/claude-code/docker-compose.yaml up -d
|
||||
#
|
||||
# See README.md for prerequisites, post-setup steps, and important caveats.
|
||||
|
||||
services:
|
||||
hindsight:
|
||||
image: ghcr.io/vectorize-io/hindsight:latest
|
||||
container_name: hindsight-claude-code
|
||||
user: "${HOST_UID:-1000}:${HOST_GID:-1000}"
|
||||
ports:
|
||||
- "127.0.0.1:8888:8888"
|
||||
- "127.0.0.1:9999:9999"
|
||||
environment:
|
||||
HOME: /home/hindsight
|
||||
USER: hindsight
|
||||
LOGNAME: hindsight
|
||||
PATH: /usr/local/bin:/usr/bin:/bin:/app/api/.venv/bin
|
||||
HINDSIGHT_API_LLM_PROVIDER: claude-code
|
||||
volumes:
|
||||
# ── Persistent data ────────────────────────────────────────────
|
||||
# Writable pg0 data directory. Must be writable by HOST_UID.
|
||||
- ${HOME:-.}/.hindsight-docker:/home/hindsight/.pg0
|
||||
|
||||
# ── Claude credentials (read-only, single-file mounts) ────────
|
||||
# A whole-directory :ro mount of ~/.claude silently breaks the
|
||||
# CLI, which writes session/project state at runtime — so we
|
||||
# mount only the two credential files.
|
||||
- ${HOME}/.claude/.credentials.json:/home/hindsight/.claude/.credentials.json:ro
|
||||
- ${HOME}/.claude.json:/home/hindsight/.claude.json:ro
|
||||
|
||||
# ── Claude CLI install (read-only) ─────────────────────────────
|
||||
- ${HOME}/.local/share/claude:/home/hindsight/.local/share/claude:ro
|
||||
|
||||
# ── SDK bundled-binary override ────────────────────────────────
|
||||
# The claude-agent-sdk 0.5.x image bundles v2.1.121 which has a
|
||||
# protocol incompatibility in containers. Override it with the
|
||||
# host's v2.1.128+ binary. Drop this mount once claude-agent-sdk
|
||||
# ships with v2.1.128+.
|
||||
- ${HOME}/.local/share/claude/versions/${CLAUDE_CLI_VERSION:-2.1.128}:/app/api/.venv/lib/python3.11/site-packages/claude_agent_sdk/_bundled/claude:ro
|
||||
@@ -1,103 +0,0 @@
|
||||
# Hindsight with a local llama.cpp server sidecar
|
||||
|
||||
Example Docker Compose setup that runs Hindsight against a **local
|
||||
llama.cpp server**, fully offline, with no external API key required.
|
||||
|
||||
## Architecture
|
||||
|
||||
```
|
||||
┌────────────┐ HTTP /v1/chat/completions ┌──────────────────────────────┐
|
||||
│ hindsight │ ──────────────────────────▶ │ llama.cpp server (sidecar) │
|
||||
│ (API + CP) │ │ ghcr.io/ggml-org/llama.cpp │
|
||||
└────────────┘ └──────────────────────────────┘
|
||||
```
|
||||
|
||||
`llama.cpp` runs as its own container and exposes an OpenAI-compatible
|
||||
HTTP API. Hindsight talks to it via the standard `openai` LLM provider
|
||||
with `HINDSIGHT_API_LLM_BASE_URL` pointed at the sidecar.
|
||||
|
||||
This pattern follows
|
||||
[*Hosting llama-server with Docker* (ServiceStack)](https://servicestack.net/posts/hosting-llama-server).
|
||||
|
||||
### Why a sidecar and not the in-process `llamacpp` provider?
|
||||
|
||||
Hindsight does ship an in-process `llamacpp` provider that spawns
|
||||
`llama-cpp-python`, but the **published `ghcr.io/vectorize-io/hindsight`
|
||||
image deliberately omits `llama-cpp-python`** to keep the image small and
|
||||
avoid bundling native inference libraries that most users don't need.
|
||||
Trying to set `HINDSIGHT_API_LLM_PROVIDER=llamacpp` against the published
|
||||
image fails with `ModuleNotFoundError: No module named 'llama_cpp'`.
|
||||
|
||||
The sidecar approach side-steps that entirely: the official llama.cpp
|
||||
image is used as-is for inference, Hindsight is used as-is for memory.
|
||||
Clean separation, no derived images.
|
||||
|
||||
## Quick start
|
||||
|
||||
```bash
|
||||
docker compose -f docker/docker-compose/local-llm/docker-compose.yaml up
|
||||
```
|
||||
|
||||
- API: http://localhost:8888
|
||||
- Control Plane: http://localhost:9999
|
||||
|
||||
**First boot downloads ~3.5 GB** (Gemma 4 E2B Q4_K_M GGUF) into the
|
||||
`llama_models` named volume. Subsequent boots reuse it.
|
||||
|
||||
Hindsight only starts after llama.cpp's `/health` endpoint reports
|
||||
healthy, so the API will appear "stuck" for a few minutes on the first
|
||||
run while the model downloads.
|
||||
|
||||
## Using a different model
|
||||
|
||||
Override the HuggingFace repo / file in `docker-compose.yaml`:
|
||||
|
||||
```yaml
|
||||
environment:
|
||||
LLAMA_ARG_HF_REPO: bartowski/Qwen2.5-7B-Instruct-GGUF
|
||||
LLAMA_ARG_HF_FILE: Qwen2.5-7B-Instruct-Q4_K_M.gguf
|
||||
```
|
||||
|
||||
Also update `HINDSIGHT_API_LLM_MODEL` on the `hindsight` service to a
|
||||
matching alias (the value is sent to llama-server as the OpenAI `model`
|
||||
field — llama-server is lenient about this but it shows up in logs).
|
||||
|
||||
## GPU acceleration
|
||||
|
||||
The default compose file targets CPU because not everyone has a GPU. On
|
||||
CPU, Gemma 4 E2B runs at ~2-3 tokens/sec — fine for a smoke test, but the
|
||||
retain pipeline (which makes several multi-hundred-token LLM calls per
|
||||
memory) will time out against Hindsight's default LLM timeout. **For any
|
||||
real use, run on a GPU.**
|
||||
|
||||
### NVIDIA
|
||||
|
||||
1. Switch the `llama` service image from `:server` to `:server-cuda`.
|
||||
2. Uncomment the `LLAMA_ARG_N_GPU_LAYERS: "999"` env var (offload all
|
||||
layers to GPU).
|
||||
3. Uncomment the `deploy.resources.reservations.devices` block.
|
||||
4. Install the [NVIDIA Container Toolkit](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html)
|
||||
on the host.
|
||||
|
||||
The compose file has all four spots marked with inline comments.
|
||||
|
||||
### Apple Silicon / ROCm / Vulkan
|
||||
|
||||
The official `ghcr.io/ggml-org/llama.cpp` image only ships CPU and CUDA
|
||||
variants. For Metal (Apple Silicon), ROCm (AMD), or Vulkan backends,
|
||||
build llama.cpp yourself with the appropriate flags and reference the
|
||||
image you build instead. Docker Desktop on macOS cannot pass through the
|
||||
host GPU to a Linux container in any case — for Apple Silicon, run
|
||||
llama-server directly on the host and only put Hindsight in Docker.
|
||||
|
||||
## Caveats
|
||||
|
||||
- llama.cpp's HTTP API is OpenAI-compatible but not 100% feature-parity.
|
||||
Function/tool calling support depends on the chat template baked into
|
||||
the GGUF; some retain/reflect flows may behave differently than against
|
||||
a hosted OpenAI model.
|
||||
- Small GGUFs (~3 B params) are useful for smoke testing but will
|
||||
underperform a hosted frontier model on retain quality. Use a larger
|
||||
GGUF (7-13 B params) for production-quality memory.
|
||||
- The `llama_models` named volume persists the GGUF across `docker
|
||||
compose down`/`up` so the model is downloaded once, not every restart.
|
||||
@@ -1,74 +0,0 @@
|
||||
name: hindsight-local-llm
|
||||
# Example: run Hindsight against a local llama.cpp server sidecar — fully
|
||||
# offline, no external API key needed.
|
||||
#
|
||||
# Pattern follows https://servicestack.net/posts/hosting-llama-server :
|
||||
# llama.cpp runs as its own container exposing an OpenAI-compatible HTTP
|
||||
# API, and Hindsight talks to it via the `openai` LLM provider with a
|
||||
# custom `base_url`. This means we can use the published Hindsight image
|
||||
# unchanged — no derived Dockerfile, no `llama-cpp-python` install on top.
|
||||
#
|
||||
# Quick start:
|
||||
# docker compose -f docker/docker-compose/local-llm/docker-compose.yaml up
|
||||
#
|
||||
# First boot downloads the default Gemma 4 E2B GGUF (~3.5 GB) into the
|
||||
# `llama_models` volume; subsequent boots reuse it.
|
||||
|
||||
services:
|
||||
llama:
|
||||
image: ghcr.io/ggml-org/llama.cpp:server
|
||||
container_name: hindsight-local-llm-llama
|
||||
environment:
|
||||
LLAMA_ARG_HOST: 0.0.0.0
|
||||
LLAMA_ARG_PORT: "8080"
|
||||
# Auto-download a small GGUF from HuggingFace on first start.
|
||||
# Override these to use a different model.
|
||||
LLAMA_ARG_HF_REPO: bartowski/google_gemma-4-E2B-it-GGUF
|
||||
LLAMA_ARG_HF_FILE: google_gemma-4-E2B-it-Q4_K_M.gguf
|
||||
LLAMA_ARG_CTX_SIZE: "8192"
|
||||
# Uncomment for NVIDIA GPU (and switch image to :server-cuda):
|
||||
# LLAMA_ARG_N_GPU_LAYERS: "999"
|
||||
volumes:
|
||||
# llama-server stores HuggingFace downloads under ~/.cache/huggingface
|
||||
# (not ~/.cache/llama.cpp), so mount the named volume there to avoid
|
||||
# re-downloading the GGUF on every recreate.
|
||||
- llama_models:/root/.cache/huggingface
|
||||
healthcheck:
|
||||
test: ["CMD-SHELL", "curl -fsS http://localhost:8080/health || exit 1"]
|
||||
interval: 10s
|
||||
timeout: 5s
|
||||
retries: 60
|
||||
start_period: 30s
|
||||
# For NVIDIA GPU acceleration, swap the image above to
|
||||
# `ghcr.io/ggml-org/llama.cpp:server-cuda` and uncomment:
|
||||
# deploy:
|
||||
# resources:
|
||||
# reservations:
|
||||
# devices:
|
||||
# - driver: nvidia
|
||||
# count: all
|
||||
# capabilities: [gpu]
|
||||
|
||||
hindsight:
|
||||
image: ghcr.io/vectorize-io/hindsight:latest
|
||||
container_name: hindsight-local-llm
|
||||
depends_on:
|
||||
llama:
|
||||
condition: service_healthy
|
||||
ports:
|
||||
- "8888:8888"
|
||||
- "9999:9999"
|
||||
environment:
|
||||
# llama-server is OpenAI-compatible, so use the `openai` provider and
|
||||
# point base_url at the sidecar. The API key is unused by llama-server
|
||||
# but Hindsight requires the env var to be set.
|
||||
HINDSIGHT_API_LLM_PROVIDER: openai
|
||||
HINDSIGHT_API_LLM_BASE_URL: http://llama:8080/v1
|
||||
HINDSIGHT_API_LLM_API_KEY: not-needed
|
||||
HINDSIGHT_API_LLM_MODEL: gemma-4-e2b-it
|
||||
volumes:
|
||||
- pg_data:/home/hindsight/.pg0
|
||||
|
||||
volumes:
|
||||
pg_data:
|
||||
llama_models:
|
||||
@@ -1,7 +0,0 @@
|
||||
# PostgreSQL with pgvector and ParadeDB pg_search extensions.
|
||||
#
|
||||
# The official ParadeDB image ships PostgreSQL with pg_search and pgvector
|
||||
# already installed, so no build steps are required. We pin to the PG17
|
||||
# variant for parity with the other Hindsight docker-compose examples
|
||||
# (vchord, pg_textsearch).
|
||||
FROM paradedb/paradedb:latest-pg17
|
||||
@@ -1,96 +0,0 @@
|
||||
name: hindsight
|
||||
# Docker Compose file for Hindsight with PostgreSQL and ParadeDB pg_search.
|
||||
#
|
||||
# pg_search is the only BM25 backend supported by Hindsight that works with
|
||||
# Citus, so this is the recommended setup for horizontally scaled deployments.
|
||||
#
|
||||
# Usage:
|
||||
# docker compose -f docker/docker-compose/pg_search/docker-compose.yaml up -d
|
||||
#
|
||||
# Required environment variables:
|
||||
# - HINDSIGHT_DB_PASSWORD: Password for the PostgreSQL user
|
||||
# - Configure LLM provider variables as needed (see the hindsight service)
|
||||
#
|
||||
# Optional environment variables with defaults:
|
||||
# - HINDSIGHT_VERSION: Hindsight application version (default: latest)
|
||||
# - HINDSIGHT_DB_USER: PostgreSQL user (default: hindsight_user)
|
||||
# - HINDSIGHT_DB_NAME: PostgreSQL database name (default: hindsight_db)
|
||||
# - HINDSIGHT_API_TEXT_SEARCH_EXTENSION_PG_SEARCH_TOKENIZER: ParadeDB pg_search
|
||||
# tokenizer for new BM25 indexes (default: empty, uses ParadeDB default)
|
||||
|
||||
services:
|
||||
db:
|
||||
# Use ParadeDB image which bundles pgvector + pg_search
|
||||
build:
|
||||
context: .
|
||||
dockerfile: Dockerfile
|
||||
container_name: hindsight-db
|
||||
restart: always
|
||||
ports:
|
||||
- "5437:5432"
|
||||
environment:
|
||||
POSTGRES_USER: ${HINDSIGHT_DB_USER:-hindsight_user}
|
||||
POSTGRES_PASSWORD: ${HINDSIGHT_DB_PASSWORD:-hindsight_password}
|
||||
POSTGRES_DB: ${HINDSIGHT_DB_NAME:-hindsight_db}
|
||||
volumes:
|
||||
- pg_data:/var/lib/postgresql/data
|
||||
networks:
|
||||
- hindsight-net
|
||||
|
||||
pg-search-init:
|
||||
build:
|
||||
context: .
|
||||
dockerfile: Dockerfile
|
||||
depends_on:
|
||||
- db
|
||||
environment:
|
||||
- PGPASSWORD=${HINDSIGHT_DB_PASSWORD:-hindsight_password}
|
||||
command: >
|
||||
bash -c "
|
||||
echo 'Waiting for PostgreSQL to be ready...';
|
||||
until pg_isready -h hindsight-db -p 5432 -U hindsight_user; do
|
||||
echo 'PostgreSQL is unavailable - sleeping';
|
||||
sleep 2;
|
||||
done;
|
||||
echo 'PostgreSQL is ready - creating hindsight_db database';
|
||||
psql -h hindsight-db -p 5432 -U hindsight_user -c 'CREATE DATABASE hindsight_db;' 2>/dev/null || echo 'Database already exists';
|
||||
echo 'Creating extensions in hindsight_db database';
|
||||
psql -h hindsight-db -p 5432 -U hindsight_user -d hindsight_db -c 'CREATE EXTENSION IF NOT EXISTS vector CASCADE;';
|
||||
psql -h hindsight-db -p 5432 -U hindsight_user -d hindsight_db -c 'CREATE EXTENSION IF NOT EXISTS pg_search CASCADE;';
|
||||
echo 'Database and extensions created successfully';
|
||||
"
|
||||
restart: "no"
|
||||
networks:
|
||||
- hindsight-net
|
||||
|
||||
hindsight:
|
||||
image: ghcr.io/vectorize-io/hindsight:${HINDSIGHT_VERSION:-latest}
|
||||
container_name: hindsight-app
|
||||
ports:
|
||||
- "8888:8888"
|
||||
- "9999:9999"
|
||||
environment:
|
||||
# LLM Configuration
|
||||
HINDSIGHT_API_LLM_PROVIDER: ${HINDSIGHT_API_LLM_PROVIDER:-openai}
|
||||
HINDSIGHT_API_LLM_API_KEY: ${OPENAI_API_KEY:-your-api-key}
|
||||
|
||||
# Database Configuration
|
||||
HINDSIGHT_API_DATABASE_URL: postgresql://${HINDSIGHT_DB_USER:-hindsight_user}:${HINDSIGHT_DB_PASSWORD:-hindsight_password}@db:5432/${HINDSIGHT_DB_NAME:-hindsight_db}
|
||||
|
||||
# Vector and Text Search Extensions
|
||||
HINDSIGHT_API_VECTOR_EXTENSION: pgvector
|
||||
HINDSIGHT_API_TEXT_SEARCH_EXTENSION: pg_search
|
||||
HINDSIGHT_API_TEXT_SEARCH_EXTENSION_PG_SEARCH_TOKENIZER: ${HINDSIGHT_API_TEXT_SEARCH_EXTENSION_PG_SEARCH_TOKENIZER:-}
|
||||
|
||||
depends_on:
|
||||
- db
|
||||
networks:
|
||||
- hindsight-net
|
||||
|
||||
|
||||
networks:
|
||||
hindsight-net:
|
||||
driver: bridge
|
||||
|
||||
volumes:
|
||||
pg_data:
|
||||
@@ -1,23 +0,0 @@
|
||||
# PostgreSQL with pgvector and pgroonga extensions.
|
||||
#
|
||||
# pgroonga is a multilingual full-text search extension built on Groonga.
|
||||
# It works out of the box for CJK (Chinese, Japanese, Korean) and other
|
||||
# non-whitespace-segmented languages via the TokenBigram tokenizer.
|
||||
FROM groonga/pgroonga:latest-debian-pg17
|
||||
|
||||
# Install pgvector on top of the pgroonga base image (which already provides
|
||||
# pgroonga and the Groonga library).
|
||||
RUN apt-get update && apt-get install -y --no-install-recommends \
|
||||
build-essential \
|
||||
git \
|
||||
postgresql-server-dev-17 \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
RUN cd /tmp && \
|
||||
git clone --branch v0.8.0 https://github.com/pgvector/pgvector.git && \
|
||||
cd pgvector && \
|
||||
make && \
|
||||
make install
|
||||
|
||||
RUN rm -rf /tmp/pgvector && \
|
||||
apt-get purge -y --auto-remove build-essential git postgresql-server-dev-17
|
||||
@@ -1,91 +0,0 @@
|
||||
name: hindsight
|
||||
# Docker Compose file for Hindsight with PostgreSQL and pgroonga
|
||||
#
|
||||
# pgroonga provides multilingual BM25 indexing that works out of the box for
|
||||
# CJK (Chinese, Japanese, Korean) and other non-whitespace-segmented languages.
|
||||
# Use this recipe if your bank content is not English/European.
|
||||
#
|
||||
# docker compose -f docker/docker-compose/pgroonga/docker-compose.yaml down && \
|
||||
# sleep 2 && \
|
||||
# docker compose -f docker/docker-compose/pgroonga/docker-compose.yaml up -d
|
||||
#
|
||||
# Optional environment variables with defaults:
|
||||
# - HINDSIGHT_VERSION: Hindsight application version (default: latest)
|
||||
# - HINDSIGHT_DB_USER: PostgreSQL user (default: hindsight_user)
|
||||
# - HINDSIGHT_DB_NAME: PostgreSQL database name (default: hindsight_db)
|
||||
# - HINDSIGHT_DB_PASSWORD: PostgreSQL password (default: hindsight_password)
|
||||
|
||||
services:
|
||||
db:
|
||||
build:
|
||||
context: .
|
||||
dockerfile: Dockerfile
|
||||
container_name: hindsight-db
|
||||
restart: always
|
||||
ports:
|
||||
- "5439:5432"
|
||||
environment:
|
||||
POSTGRES_USER: ${HINDSIGHT_DB_USER:-hindsight_user}
|
||||
POSTGRES_PASSWORD: ${HINDSIGHT_DB_PASSWORD:-hindsight_password}
|
||||
POSTGRES_DB: ${HINDSIGHT_DB_NAME:-hindsight_db}
|
||||
volumes:
|
||||
- pg_data:/var/lib/postgresql/data
|
||||
networks:
|
||||
- hindsight-net
|
||||
|
||||
pgroonga-init:
|
||||
build:
|
||||
context: .
|
||||
dockerfile: Dockerfile
|
||||
depends_on:
|
||||
- db
|
||||
environment:
|
||||
- PGPASSWORD=${HINDSIGHT_DB_PASSWORD:-hindsight_password}
|
||||
command: >
|
||||
bash -c "
|
||||
echo 'Waiting for PostgreSQL to be ready...';
|
||||
until pg_isready -h hindsight-db -p 5432 -U hindsight_user; do
|
||||
echo 'PostgreSQL is unavailable - sleeping';
|
||||
sleep 2;
|
||||
done;
|
||||
echo 'PostgreSQL is ready - creating hindsight_db database';
|
||||
psql -h hindsight-db -p 5432 -U hindsight_user -c 'CREATE DATABASE hindsight_db;' 2>/dev/null || echo 'Database already exists';
|
||||
echo 'Creating extensions in hindsight_db database';
|
||||
psql -h hindsight-db -p 5432 -U hindsight_user -d hindsight_db -c 'CREATE EXTENSION IF NOT EXISTS vector CASCADE;';
|
||||
psql -h hindsight-db -p 5432 -U hindsight_user -d hindsight_db -c 'CREATE EXTENSION IF NOT EXISTS pgroonga CASCADE;';
|
||||
echo 'Database and extensions created successfully';
|
||||
"
|
||||
restart: "no"
|
||||
networks:
|
||||
- hindsight-net
|
||||
|
||||
hindsight:
|
||||
image: ghcr.io/vectorize-io/hindsight:${HINDSIGHT_VERSION:-latest}
|
||||
container_name: hindsight-app
|
||||
ports:
|
||||
- "8888:8888"
|
||||
- "9999:9999"
|
||||
environment:
|
||||
# LLM Configuration
|
||||
HINDSIGHT_API_LLM_PROVIDER: ${HINDSIGHT_API_LLM_PROVIDER:-openai}
|
||||
HINDSIGHT_API_LLM_API_KEY: ${OPENAI_API_KEY:-your-api-key}
|
||||
|
||||
# Database Configuration
|
||||
HINDSIGHT_API_DATABASE_URL: postgresql://${HINDSIGHT_DB_USER:-hindsight_user}:${HINDSIGHT_DB_PASSWORD:-hindsight_password}@db:5432/${HINDSIGHT_DB_NAME:-hindsight_db}
|
||||
|
||||
# Vector and Text Search Extensions
|
||||
HINDSIGHT_API_VECTOR_EXTENSION: pgvector
|
||||
HINDSIGHT_API_TEXT_SEARCH_EXTENSION: pgroonga
|
||||
|
||||
depends_on:
|
||||
- db
|
||||
networks:
|
||||
- hindsight-net
|
||||
|
||||
|
||||
networks:
|
||||
hindsight-net:
|
||||
driver: bridge
|
||||
|
||||
volumes:
|
||||
pg_data:
|
||||
@@ -1,6 +1,6 @@
|
||||
name: hindsight
|
||||
# Docker Compose file for Hindsight with PostgreSQL and vectorchord
|
||||
# docker compose -f docker/docker-compose/vchord/docker-compose.yaml down && sleep 2 && docker compose -f docker/docker-compose/vchord/docker-compose.yaml up -d
|
||||
# docker compose -f docker/docker-compose/docker-compose.yaml down && sleep 2 && docker compose -f docker/docker-compose/docker-compose.yaml up -d
|
||||
# Make sure to set the required environment variables before running:
|
||||
# - HINDSIGHT_DB_PASSWORD: Password for the PostgreSQL user
|
||||
# - Configure LLM provider variables as needed (see below in the hindsight service)
|
||||
|
||||
@@ -50,8 +50,6 @@ 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.
|
||||
RUN if [ "$INCLUDE_LOCAL_MODELS" = "true" ]; then \
|
||||
uv sync --extra local-ml --extra embedded-db; \
|
||||
else \
|
||||
|
||||
@@ -10,90 +10,19 @@ set -e
|
||||
# loss scenarios where a container restart caused the data directory to be
|
||||
# wiped despite a volume mount being present.
|
||||
# =============================================================================
|
||||
pg0_has_pg_version() {
|
||||
local pg0_data_dir="$1"
|
||||
|
||||
# pg0 has used more than one on-disk layout. Newer standalone images keep
|
||||
# PostgreSQL data under instances/<name>/data, while older volumes may have
|
||||
# placed PG_VERSION at or one level below the mount.
|
||||
[ -f "$pg0_data_dir/PG_VERSION" ] && return 0
|
||||
compgen -G "$pg0_data_dir"/*/PG_VERSION > /dev/null 2>&1 && return 0
|
||||
compgen -G "$pg0_data_dir"/instances/*/data/PG_VERSION > /dev/null 2>&1 && return 0
|
||||
|
||||
return 1
|
||||
}
|
||||
|
||||
check_pg0_data_integrity() {
|
||||
local pg0_data_dir="$1"
|
||||
|
||||
if [ ! -d "$pg0_data_dir" ]; then
|
||||
return 0
|
||||
fi
|
||||
|
||||
PG0_DATA_DIR="${HOME}/.pg0"
|
||||
if [ -d "$PG0_DATA_DIR" ]; then
|
||||
# Look for actual PostgreSQL data directories (pg0 creates subdirs per instance)
|
||||
if pg0_has_pg_version "$pg0_data_dir"; then
|
||||
echo "✅ Existing pg0 data directory detected at $pg0_data_dir"
|
||||
elif [ "$(ls -A "$pg0_data_dir" 2>/dev/null)" ]; then
|
||||
echo "⚠️ WARNING: pg0 data directory exists at $pg0_data_dir but no PG_VERSION found."
|
||||
if compgen -G "$PG0_DATA_DIR"/*/PG_VERSION > /dev/null 2>&1; then
|
||||
echo "✅ Existing pg0 data directory detected at $PG0_DATA_DIR"
|
||||
elif [ "$(ls -A "$PG0_DATA_DIR" 2>/dev/null)" ]; then
|
||||
echo "⚠️ WARNING: pg0 data directory exists at $PG0_DATA_DIR but no PG_VERSION found."
|
||||
echo " This may indicate data corruption or an incomplete previous shutdown."
|
||||
echo " If you see all migrations running from scratch after this, your data may have been lost."
|
||||
echo " See: https://github.com/vectorize-io/hindsight/issues/675"
|
||||
fi
|
||||
|
||||
return 0
|
||||
}
|
||||
|
||||
# =============================================================================
|
||||
# Embedded pg0 writability pre-check (#1483)
|
||||
#
|
||||
# The container runs as the unprivileged `hindsight` user (UID 1000). When the
|
||||
# pg0 data directory is a host bind mount (e.g. `-v $HOME/dir:/home/hindsight/.pg0`)
|
||||
# that is not owned by UID 1000 — the default on macOS Docker Desktop and most
|
||||
# non-1000 Linux hosts — pg0 fails with the opaque "Permission denied (os error
|
||||
# 13)". We cannot chown it ourselves without root (and the image is deliberately
|
||||
# rootless), so we surface an actionable message up front instead.
|
||||
#
|
||||
# Docker *named* volumes are seeded with the image directory's ownership (UID
|
||||
# 1000) on first use, so they avoid this entirely — hence the named-volume
|
||||
# recommendation below and in the README.
|
||||
# =============================================================================
|
||||
check_pg0_writable() {
|
||||
local pg0_data_dir="$1"
|
||||
|
||||
# Only relevant for embedded pg0; an external database doesn't use this dir.
|
||||
if [ -n "${HINDSIGHT_API_DATABASE_URL:-}" ]; then
|
||||
return 0
|
||||
fi
|
||||
|
||||
mkdir -p "$pg0_data_dir" 2>/dev/null || true
|
||||
if touch "$pg0_data_dir/.hindsight-write-test" 2>/dev/null; then
|
||||
rm -f "$pg0_data_dir/.hindsight-write-test" 2>/dev/null || true
|
||||
return 0
|
||||
fi
|
||||
|
||||
echo "❌ The embedded database directory $pg0_data_dir is not writable by this container (UID $(id -u))."
|
||||
echo ""
|
||||
echo " A host directory was bind-mounted but is not owned by the container user (UID 1000)."
|
||||
echo " Hindsight runs rootless and cannot fix this for you. Choose one:"
|
||||
echo ""
|
||||
echo " • Recommended — use a Docker named volume (auto-owned by the container):"
|
||||
echo " -v hindsight-data:/home/hindsight/.pg0"
|
||||
echo ""
|
||||
echo " • Or keep the host path and run as your host user, chowning it to match:"
|
||||
echo " sudo chown -R \$(id -u):\$(id -g) <host-directory>"
|
||||
echo " docker run --user \$(id -u):\$(id -g) -e HOME=/home/hindsight ..."
|
||||
echo ""
|
||||
echo " See https://github.com/vectorize-io/hindsight/issues/1483"
|
||||
return 1
|
||||
}
|
||||
|
||||
if [ "${HINDSIGHT_START_ALL_SOURCE_ONLY:-false}" = "true" ]; then
|
||||
return 0 2>/dev/null || exit 0
|
||||
fi
|
||||
|
||||
check_pg0_data_integrity "${HOME}/.pg0"
|
||||
check_pg0_writable "${HOME}/.pg0" || exit 1
|
||||
|
||||
# Service flags (default to true if not set)
|
||||
ENABLE_API="${HINDSIGHT_ENABLE_API:-true}"
|
||||
ENABLE_CP="${HINDSIGHT_ENABLE_CP:-true}"
|
||||
@@ -227,7 +156,7 @@ PIDS=()
|
||||
# Start API if enabled
|
||||
if [ "$ENABLE_API" = "true" ]; then
|
||||
cd /app/api
|
||||
API_HEALTH_URL="${HINDSIGHT_API_HEALTH_URL:-http://localhost:${HINDSIGHT_API_PORT:-8888}/health}"
|
||||
API_HEALTH_URL="${HINDSIGHT_API_HEALTH_URL:-http://localhost:8888/health}"
|
||||
API_STARTUP_WAIT_SECONDS="${HINDSIGHT_API_STARTUP_WAIT_SECONDS:-300}"
|
||||
|
||||
# Run API directly - Python's PYTHONUNBUFFERED=1 handles output buffering
|
||||
|
||||
@@ -1,121 +0,0 @@
|
||||
#!/bin/bash
|
||||
set -euo pipefail
|
||||
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
|
||||
HINDSIGHT_START_ALL_SOURCE_ONLY=true
|
||||
source "$SCRIPT_DIR/start-all.sh"
|
||||
unset HINDSIGHT_START_ALL_SOURCE_ONLY
|
||||
|
||||
TMP_DIR="$(mktemp -d)"
|
||||
trap 'chmod -R u+rwx "$TMP_DIR" 2>/dev/null || true; rm -rf "$TMP_DIR"' EXIT
|
||||
|
||||
assert_contains() {
|
||||
local output="$1"
|
||||
local expected="$2"
|
||||
|
||||
if [[ "$output" != *"$expected"* ]]; then
|
||||
echo "Expected output to contain: $expected"
|
||||
echo "Actual output:"
|
||||
echo "$output"
|
||||
exit 1
|
||||
fi
|
||||
}
|
||||
|
||||
assert_not_contains() {
|
||||
local output="$1"
|
||||
local unexpected="$2"
|
||||
|
||||
if [[ "$output" == *"$unexpected"* ]]; then
|
||||
echo "Expected output not to contain: $unexpected"
|
||||
echo "Actual output:"
|
||||
echo "$output"
|
||||
exit 1
|
||||
fi
|
||||
}
|
||||
|
||||
assert_empty() {
|
||||
local output="$1"
|
||||
|
||||
if [ -n "$output" ]; then
|
||||
echo "Expected no output, got:"
|
||||
echo "$output"
|
||||
exit 1
|
||||
fi
|
||||
}
|
||||
|
||||
mkdir -p "$TMP_DIR/empty"
|
||||
assert_empty "$(check_pg0_data_integrity "$TMP_DIR/empty")"
|
||||
|
||||
mkdir -p "$TMP_DIR/direct"
|
||||
touch "$TMP_DIR/direct/PG_VERSION"
|
||||
direct_output="$(check_pg0_data_integrity "$TMP_DIR/direct")"
|
||||
assert_contains "$direct_output" "Existing pg0 data directory detected"
|
||||
assert_not_contains "$direct_output" "WARNING"
|
||||
|
||||
mkdir -p "$TMP_DIR/legacy/instance"
|
||||
touch "$TMP_DIR/legacy/instance/PG_VERSION"
|
||||
legacy_output="$(check_pg0_data_integrity "$TMP_DIR/legacy")"
|
||||
assert_contains "$legacy_output" "Existing pg0 data directory detected"
|
||||
assert_not_contains "$legacy_output" "WARNING"
|
||||
|
||||
mkdir -p "$TMP_DIR/nested/instances/hindsight/data"
|
||||
touch "$TMP_DIR/nested/instances/hindsight/data/PG_VERSION"
|
||||
nested_output="$(check_pg0_data_integrity "$TMP_DIR/nested")"
|
||||
assert_contains "$nested_output" "Existing pg0 data directory detected"
|
||||
assert_not_contains "$nested_output" "WARNING"
|
||||
|
||||
mkdir -p "$TMP_DIR/nonempty/instances/hindsight"
|
||||
touch "$TMP_DIR/nonempty/instances/hindsight/instance.json"
|
||||
nonempty_output="$(check_pg0_data_integrity "$TMP_DIR/nonempty")"
|
||||
assert_contains "$nonempty_output" "WARNING: pg0 data directory exists"
|
||||
|
||||
echo "start-all pg0 integrity checks passed"
|
||||
|
||||
# =============================================================================
|
||||
# check_pg0_writable (#1483)
|
||||
# These rely on filesystem permissions, which root bypasses; skip under root.
|
||||
# =============================================================================
|
||||
if [ "$(id -u)" != "0" ]; then
|
||||
# Writable directory: returns 0, prints nothing, leaves no artifact behind.
|
||||
mkdir -p "$TMP_DIR/writable"
|
||||
writable_output="$(check_pg0_writable "$TMP_DIR/writable")"
|
||||
assert_empty "$writable_output"
|
||||
if [ -e "$TMP_DIR/writable/.hindsight-write-test" ]; then
|
||||
echo "check_pg0_writable left its write-test file behind"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# Non-writable directory: returns 1 with actionable guidance.
|
||||
mkdir -p "$TMP_DIR/readonly"
|
||||
chmod 000 "$TMP_DIR/readonly"
|
||||
set +e
|
||||
readonly_output="$(check_pg0_writable "$TMP_DIR/readonly" 2>&1)"
|
||||
readonly_rc=$?
|
||||
set -e
|
||||
chmod 755 "$TMP_DIR/readonly"
|
||||
if [ "$readonly_rc" -eq 0 ]; then
|
||||
echo "check_pg0_writable should fail on a non-writable directory"
|
||||
exit 1
|
||||
fi
|
||||
assert_contains "$readonly_output" "not writable"
|
||||
assert_contains "$readonly_output" "hindsight-data:/home/hindsight/.pg0"
|
||||
assert_contains "$readonly_output" "--user"
|
||||
|
||||
# External database configured: skip the check regardless of dir perms.
|
||||
mkdir -p "$TMP_DIR/extdb"
|
||||
chmod 000 "$TMP_DIR/extdb"
|
||||
set +e
|
||||
HINDSIGHT_API_DATABASE_URL="postgres://x" check_pg0_writable "$TMP_DIR/extdb" >/dev/null 2>&1
|
||||
extdb_rc=$?
|
||||
set -e
|
||||
chmod 755 "$TMP_DIR/extdb"
|
||||
if [ "$extdb_rc" -ne 0 ]; then
|
||||
echo "check_pg0_writable should skip when an external database is configured"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
echo "start-all pg0 writability checks passed"
|
||||
else
|
||||
echo "⚠️ Running as root; skipping pg0 writability checks (permissions are bypassed)."
|
||||
fi
|
||||
@@ -0,0 +1,6 @@
|
||||
dependencies:
|
||||
- name: postgresql
|
||||
repository: https://charts.bitnami.com/bitnami
|
||||
version: 15.5.38
|
||||
digest: sha256:f67c7612736803ece8a669f8ca6b0555f3b78557bc0ecb732aa2e43f0df7750d
|
||||
generated: "2025-12-10T17:20:57.058794+01:00"
|
||||
@@ -2,8 +2,8 @@ apiVersion: v2
|
||||
name: hindsight
|
||||
description: Hindsight helm chart
|
||||
type: application
|
||||
version: 0.8.4
|
||||
appVersion: "0.8.4"
|
||||
version: 0.6.2
|
||||
appVersion: "0.6.2"
|
||||
keywords:
|
||||
- ai
|
||||
- memory
|
||||
|
||||
@@ -66,13 +66,13 @@ helm install hindsight ./helm/hindsight -n hindsight --create-namespace -f value
|
||||
|
||||
| Parameter | Description | Default |
|
||||
|-----------|-------------|---------|
|
||||
| `version` | Default image tag for all components | Chart `appVersion` |
|
||||
| `version` | Default image tag for all components | `0.1.0` |
|
||||
| `api.enabled` | Enable the API component | `true` |
|
||||
| `api.image.repository` | API image repository | `ghcr.io/vectorize-io/hindsight-api` |
|
||||
| `api.image.repository` | API image repository | `hindsight/api` |
|
||||
| `api.image.tag` | API image tag (defaults to `version`) | - |
|
||||
| `api.service.port` | API service port | `8888` |
|
||||
| `controlPlane.enabled` | Enable the control plane | `true` |
|
||||
| `controlPlane.image.repository` | Control plane image repository | `ghcr.io/vectorize-io/hindsight-control-plane` |
|
||||
| `controlPlane.image.repository` | Control plane image repository | `hindsight/control-plane` |
|
||||
| `controlPlane.image.tag` | Control plane image tag (defaults to `version`) | - |
|
||||
| `controlPlane.service.port` | Control plane service port | `3000` |
|
||||
| `postgresql.enabled` | Deploy PostgreSQL as subchart | `true` |
|
||||
|
||||
@@ -13,6 +13,9 @@
|
||||
# - Any other env vars you want to inject
|
||||
# existingSecret: "my-hindsight-secret"
|
||||
|
||||
# Global settings
|
||||
replicaCount: 1
|
||||
|
||||
# Image settings for api
|
||||
api:
|
||||
enabled: true
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@vectorize-io/hindsight-all",
|
||||
"version": "0.8.4",
|
||||
"version": "0.6.2",
|
||||
"description": "Node.js programmatic lifecycle manager for Hindsight — embeds a local hindsight daemon in a Node application. Pair with @vectorize-io/hindsight-client for memory operations.",
|
||||
"main": "dist/index.js",
|
||||
"types": "dist/index.d.ts",
|
||||
|
||||
@@ -4,12 +4,12 @@ build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "hindsight-all-slim"
|
||||
version = "0.8.4"
|
||||
version = "0.6.2"
|
||||
description = "Hindsight: Agent Memory That Works Like Human Memory - Slim All-in-One Bundle"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.11"
|
||||
dependencies = [
|
||||
"hindsight-api-slim==0.8.4",
|
||||
"hindsight-api-slim==0.6.2",
|
||||
"hindsight-client>=0.0.7",
|
||||
"hindsight-embed>=0.1.0",
|
||||
]
|
||||
|
||||
@@ -4,12 +4,12 @@ build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "hindsight-all"
|
||||
version = "0.8.4"
|
||||
version = "0.6.2"
|
||||
description = "Hindsight: Agent Memory That Works Like Human Memory - All-in-One Bundle"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.11"
|
||||
dependencies = [
|
||||
"hindsight-api-slim[all]==0.8.4",
|
||||
"hindsight-api-slim[all]==0.6.2",
|
||||
"hindsight-client>=0.0.7",
|
||||
"hindsight-embed>=0.1.0",
|
||||
]
|
||||
@@ -21,7 +21,7 @@ hindsight-embed = { workspace = true }
|
||||
|
||||
[project.optional-dependencies]
|
||||
local-llm = [
|
||||
"hindsight-api-slim[local-llm]==0.8.4",
|
||||
"hindsight-api-slim[local-llm]==0.6.2",
|
||||
]
|
||||
test = [
|
||||
"pytest>=7.0.0",
|
||||
|
||||
@@ -386,7 +386,7 @@ def test_embedded_ui_flag(llm_config):
|
||||
|
||||
# Verify UI is reachable and reports connected dataplane
|
||||
ui_url = client.ui_url
|
||||
assert isinstance(ui_url, str) and ui_url, "ui_url should be a non-empty string"
|
||||
assert ui_url, "ui_url should be set"
|
||||
|
||||
health_url = f"{ui_url}/api/health"
|
||||
with urllib.request.urlopen(health_url, timeout=10) as resp:
|
||||
|
||||
@@ -99,7 +99,7 @@ hindsight-api
|
||||
## Docker
|
||||
|
||||
```bash
|
||||
docker run -it --name hindsight --restart unless-stopped -p 8888:8888 \
|
||||
docker run --rm -it -p 8888:8888 \
|
||||
-e HINDSIGHT_API_LLM_API_KEY=$OPENAI_API_KEY \
|
||||
-v $HOME/.hindsight-docker:/home/hindsight/.pg0 \
|
||||
ghcr.io/vectorize-io/hindsight:latest
|
||||
@@ -121,7 +121,7 @@ This runs a stdio-based MCP server that can be used directly with MCP-compatible
|
||||
- **Entity Graph** — Automatic entity extraction and relationship tracking
|
||||
- **Temporal Reasoning** — Native support for time-based queries
|
||||
- **Disposition Traits** — Configurable skepticism, literalism, and empathy influence opinion formation
|
||||
- **Three Memory Types** — World facts, experience facts (the bank's own actions), and observations
|
||||
- **Three Memory Types** — World facts, bank actions, and formed opinions with confidence scores
|
||||
|
||||
## Documentation
|
||||
|
||||
|
||||
@@ -4,13 +4,6 @@ Memory System for AI Agents.
|
||||
Temporal + Semantic Memory Architecture using PostgreSQL with pgvector.
|
||||
"""
|
||||
|
||||
# Cap native ML thread pools (OpenBLAS/OpenMP/MKL) before any import pulls in
|
||||
# numpy/torch/onnxruntime — they read these env vars only at load time. See
|
||||
# hindsight_api/_thread_limits.py for the rationale.
|
||||
from ._thread_limits import apply_default_thread_limits
|
||||
|
||||
apply_default_thread_limits()
|
||||
|
||||
from .config import HindsightConfig, get_config
|
||||
from .engine.cross_encoder import CrossEncoderModel, LocalSTCrossEncoder, RemoteTEICrossEncoder
|
||||
from .engine.embeddings import Embeddings, LocalSTEmbeddings, RemoteTEIEmbeddings
|
||||
@@ -53,4 +46,4 @@ __all__ = [
|
||||
"RemoteTEICrossEncoder",
|
||||
"LLMConfig",
|
||||
]
|
||||
__version__ = "0.8.4"
|
||||
__version__ = "0.6.2"
|
||||
|
||||
@@ -1,85 +0,0 @@
|
||||
"""Helpers for ParadeDB pg_search index configuration."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from collections.abc import Sequence
|
||||
|
||||
PG_SEARCH_TOKENIZER_ENV = "HINDSIGHT_API_TEXT_SEARCH_EXTENSION_PG_SEARCH_TOKENIZER"
|
||||
|
||||
_SIMPLE_TOKENIZERS = {
|
||||
"unicode_words",
|
||||
"simple",
|
||||
"whitespace",
|
||||
"literal",
|
||||
"literal_normalized",
|
||||
"chinese_compatible",
|
||||
"icu",
|
||||
"jieba",
|
||||
"source_code",
|
||||
}
|
||||
|
||||
_TOKENIZER_ALIASES = {
|
||||
"chinese_lindera": "lindera(chinese)",
|
||||
"japanese_lindera": "lindera(japanese)",
|
||||
"korean_lindera": "lindera(korean)",
|
||||
"lindera_chinese": "lindera(chinese)",
|
||||
"lindera_japanese": "lindera(japanese)",
|
||||
"lindera_korean": "lindera(korean)",
|
||||
}
|
||||
|
||||
|
||||
def normalize_pg_search_tokenizer(value: str | None) -> str:
|
||||
"""Validate and normalize a ParadeDB pg_search tokenizer setting.
|
||||
|
||||
Returns an empty string when unset. The returned value is safe to embed after
|
||||
``pdb.`` in a CREATE INDEX expression.
|
||||
"""
|
||||
|
||||
tokenizer = (value or "").strip().lower()
|
||||
if not tokenizer:
|
||||
return ""
|
||||
|
||||
if tokenizer in _TOKENIZER_ALIASES:
|
||||
return _TOKENIZER_ALIASES[tokenizer]
|
||||
|
||||
if tokenizer in _SIMPLE_TOKENIZERS:
|
||||
return tokenizer
|
||||
|
||||
lindera_match = re.fullmatch(r"lindera\((chinese|japanese|korean)\)", tokenizer)
|
||||
if lindera_match:
|
||||
return tokenizer
|
||||
|
||||
ngram_match = re.fullmatch(r"(ngram|edge_ngram)\((\d{1,3}),\s*(\d{1,3})\)", tokenizer)
|
||||
if ngram_match:
|
||||
kind, min_gram, max_gram = ngram_match.groups()
|
||||
min_value = int(min_gram)
|
||||
max_value = int(max_gram)
|
||||
if min_value <= 0 or min_value > max_value:
|
||||
raise ValueError(
|
||||
f"Invalid {PG_SEARCH_TOKENIZER_ENV}: {value!r}. "
|
||||
"ngram and edge_ngram require positive min/max gram sizes with min <= max."
|
||||
)
|
||||
return f"{kind}({min_value},{max_value})"
|
||||
|
||||
raise ValueError(
|
||||
f"Invalid {PG_SEARCH_TOKENIZER_ENV}: {value!r}. "
|
||||
"Supported values are: unicode_words, simple, whitespace, literal, "
|
||||
"literal_normalized, chinese_compatible, icu, jieba, source_code, "
|
||||
"chinese_lindera, japanese_lindera, korean_lindera, or "
|
||||
"lindera(chinese|japanese|korean), ngram(min,max), or edge_ngram(min,max)."
|
||||
)
|
||||
|
||||
|
||||
def pg_search_bm25_columns(
|
||||
key_field: str,
|
||||
text_fields: Sequence[str],
|
||||
tokenizer: str | None,
|
||||
) -> str:
|
||||
"""Build a ParadeDB BM25 column list for CREATE INDEX."""
|
||||
|
||||
normalized = normalize_pg_search_tokenizer(tokenizer)
|
||||
if not normalized:
|
||||
return ", ".join([key_field, *text_fields])
|
||||
|
||||
return ", ".join([key_field, *(f"({field}::pdb.{normalized})" for field in text_fields)])
|
||||
@@ -1,107 +0,0 @@
|
||||
"""Process-level caps for native ML thread pools.
|
||||
|
||||
OpenBLAS, OpenMP, and MKL each spawn a worker pool sized to the host CPU count
|
||||
the first time they are loaded (numpy pulls in OpenBLAS eagerly; torch and
|
||||
onnxruntime load their pools lazily on first inference). Hindsight already
|
||||
parallelizes at the request level via thread-pool executors (embeddings on the
|
||||
default executor, the reranker on its own pool), so these native intra-op pools
|
||||
oversubscribe the CPU: on a many-core host the process accumulates 100+ native
|
||||
threads, which inflates memory and, under contention, can degrade throughput.
|
||||
|
||||
We bound each pool to ``_MAX_NATIVE_THREADS`` (or the available CPU count, if
|
||||
smaller). "Available" is the CPU budget actually granted to the process, not
|
||||
``os.cpu_count()``: in a CPU-limited container ``os.cpu_count()`` still reports
|
||||
the host's cores, so sizing pools by it oversubscribes the container's real
|
||||
quota — the exact failure mode this guards against. We therefore take the
|
||||
smallest of the CPU-affinity set, the cgroup CPU quota, and ``os.cpu_count()``.
|
||||
|
||||
Every cap is applied with ``setdefault`` so an operator who has deliberately
|
||||
tuned one of these variables keeps their value. This must run *before* numpy,
|
||||
torch, or onnxruntime are imported — those libraries read the variables only at
|
||||
load time — which is why it is invoked at the very top of
|
||||
``hindsight_api/__init__.py``, ahead of the package's other imports.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
|
||||
# Native threading env vars, each read by the respective library at load time.
|
||||
_NATIVE_THREAD_VARS = (
|
||||
"OMP_NUM_THREADS", # OpenMP — torch, onnxruntime, some BLAS builds
|
||||
"OPENBLAS_NUM_THREADS", # OpenBLAS — numpy's default BLAS
|
||||
"MKL_NUM_THREADS", # Intel MKL — numpy/torch when MKL-backed
|
||||
"NUMEXPR_NUM_THREADS", # numexpr expression engine
|
||||
)
|
||||
|
||||
# Upper bound on intra-op threads per native pool. Bounds runaway growth on
|
||||
# many-core hosts without serialising single-request inference.
|
||||
_MAX_NATIVE_THREADS = 16
|
||||
|
||||
|
||||
def _quota_to_cpus(quota: int, period: int) -> int | None:
|
||||
"""Whole CPUs from a CFS quota/period pair, or None if unlimited."""
|
||||
if quota > 0 and period > 0:
|
||||
# Floor (never round up) so we never exceed the granted budget.
|
||||
return max(1, quota // period)
|
||||
return None
|
||||
|
||||
|
||||
def _parse_cgroup_v2_cpu_max(text: str) -> int | None:
|
||||
"""Parse cgroup v2 ``cpu.max`` ("<quota> <period>", or "max <period>")."""
|
||||
parts = text.split()
|
||||
if len(parts) >= 2 and parts[0] != "max":
|
||||
try:
|
||||
return _quota_to_cpus(int(parts[0]), int(parts[1]))
|
||||
except ValueError:
|
||||
return None
|
||||
return None
|
||||
|
||||
|
||||
def _cgroup_cpu_quota() -> int | None:
|
||||
"""Effective CPUs from the cgroup CPU quota, or None if unlimited/unknown."""
|
||||
try: # cgroup v2
|
||||
with open("/sys/fs/cgroup/cpu.max") as fh:
|
||||
return _parse_cgroup_v2_cpu_max(fh.read())
|
||||
except OSError:
|
||||
pass
|
||||
try: # cgroup v1
|
||||
with open("/sys/fs/cgroup/cpu/cpu.cfs_quota_us") as fh:
|
||||
quota = int(fh.read())
|
||||
with open("/sys/fs/cgroup/cpu/cpu.cfs_period_us") as fh:
|
||||
period = int(fh.read())
|
||||
return _quota_to_cpus(quota, period)
|
||||
except (OSError, ValueError):
|
||||
return None
|
||||
|
||||
|
||||
def _available_cpu_count() -> int:
|
||||
"""CPUs actually available to this process.
|
||||
|
||||
The smallest of the CPU-affinity set (cpuset / ``--cpuset-cpus``), the
|
||||
cgroup CPU quota (``--cpus``), and ``os.cpu_count()`` — each captures a
|
||||
different way the budget can be constrained, and the last alone overcounts
|
||||
inside a limited container.
|
||||
"""
|
||||
candidates = [os.cpu_count() or 1]
|
||||
if hasattr(os, "sched_getaffinity"):
|
||||
try:
|
||||
candidates.append(len(os.sched_getaffinity(0)))
|
||||
except OSError:
|
||||
pass
|
||||
quota = _cgroup_cpu_quota()
|
||||
if quota is not None:
|
||||
candidates.append(quota)
|
||||
return max(1, min(candidates))
|
||||
|
||||
|
||||
def default_native_thread_count() -> int:
|
||||
"""Per-pool cap: ``_MAX_NATIVE_THREADS``, or available CPUs if fewer."""
|
||||
return min(_MAX_NATIVE_THREADS, _available_cpu_count())
|
||||
|
||||
|
||||
def apply_default_thread_limits() -> None:
|
||||
"""Cap native ML thread pools unless the operator has set the var already."""
|
||||
value = str(default_native_thread_count())
|
||||
for var in _NATIVE_THREAD_VARS:
|
||||
os.environ.setdefault(var, value)
|
||||
@@ -3,7 +3,6 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
|
||||
from sqlalchemy import text
|
||||
from sqlalchemy.engine import Connection
|
||||
@@ -35,7 +34,7 @@ _INDEX_USING_CLAUSES = {
|
||||
"pgvector": "USING hnsw (embedding vector_cosine_ops)",
|
||||
"pgvectorscale": "USING diskann (embedding vector_cosine_ops) WITH (num_neighbors = 50)",
|
||||
"pg_diskann": "USING diskann (embedding vector_cosine_ops) WITH (max_neighbors = 50)",
|
||||
"vchord": "USING vchordrq (embedding vector_cosine_ops)",
|
||||
"vchord": "USING vchordrq (embedding vector_l2_ops)",
|
||||
"scann": "USING scann (embedding cosine) WITH (mode = 'AUTO')",
|
||||
}
|
||||
|
||||
@@ -47,29 +46,6 @@ _INDEX_TYPE_KEYWORDS = {
|
||||
"scann": "scann",
|
||||
}
|
||||
|
||||
# Per-backend ANN search-time tuning GUCs. Each entry is a tuple of
|
||||
# (guc_name, value) pairs the caller can apply with SET or SET LOCAL.
|
||||
#
|
||||
# - pgvector exposes hnsw.ef_search. The 60 / 200 pair is unchanged from the
|
||||
# pre-dispatcher code (internal benchmarks tuned around our embedding count
|
||||
# and recall floor; see the link_utils / pool init call sites for the
|
||||
# latency-vs-recall framing).
|
||||
# - vchord exposes vchordrq.probes, but its shape must match the index's
|
||||
# build.internal.lists hierarchy. VectorChord 1.1 added per-index fallback
|
||||
# parameters for this reason: a session GUC overrides every vchordrq index,
|
||||
# and a single value can be invalid for listless or mixed-layout indexes.
|
||||
# Hindsight's built-in vchord clause does not set lists, so the safe default
|
||||
# is no session-level probe override; deployments that partition vchordrq
|
||||
# indexes should attach probes to the index storage parameters instead.
|
||||
# - pgvectorscale / pg_diskann / scann do not expose an equivalent per-statement
|
||||
# knob in the engine today, so the dispatcher returns no statements for them.
|
||||
_ANN_TUNING_LOW_LATENCY: dict[str, tuple[tuple[str, str], ...]] = {
|
||||
"pgvector": (("hnsw.ef_search", "60"),),
|
||||
}
|
||||
_ANN_TUNING_HIGH_RECALL: dict[str, tuple[tuple[str, str], ...]] = {
|
||||
"pgvector": (("hnsw.ef_search", "200"),),
|
||||
}
|
||||
|
||||
_EXTENSION_INSTALL_SQL = {
|
||||
"pgvector": ("CREATE EXTENSION IF NOT EXISTS vector",),
|
||||
"pgvectorscale": (
|
||||
@@ -91,18 +67,6 @@ _INSTALL_HINTS = {
|
||||
}
|
||||
|
||||
|
||||
def configured_vector_extension() -> str:
|
||||
"""Return the user-configured vector backend extension.
|
||||
|
||||
Reads ``HINDSIGHT_API_VECTOR_EXTENSION`` (default ``"pgvector"``) and
|
||||
validates it via :func:`validate_extension`. This is the single source of
|
||||
truth for runtime code that needs to dispatch behaviour by vector backend;
|
||||
callers should prefer this over reading the env var directly, so the
|
||||
default value and the lookup mechanism live in one place.
|
||||
"""
|
||||
return validate_extension(os.getenv("HINDSIGHT_API_VECTOR_EXTENSION", "pgvector"))
|
||||
|
||||
|
||||
def validate_extension(name: str) -> str:
|
||||
"""Return a normalized configurable vector extension name or raise.
|
||||
|
||||
@@ -151,25 +115,6 @@ def should_defer_index_creation(ext: str, row_count: int) -> bool:
|
||||
return minimum_rows > 0 and row_count < minimum_rows
|
||||
|
||||
|
||||
def ann_search_tuning_settings(ext: str, *, kind: str) -> tuple[tuple[str, str], ...]:
|
||||
"""Return per-backend (guc_name, value) pairs for ANN search-time tuning.
|
||||
|
||||
``kind`` is ``"low_latency"`` for retain-side link probing (smaller probe
|
||||
count, lower recall, lower latency) and ``"high_recall"`` for connection
|
||||
init in the pool (larger probe count, higher recall). Callers wrap each
|
||||
pair with ``SET LOCAL`` or ``SET`` themselves so the same dispatcher works
|
||||
for both transaction-scoped and session-scoped use. Returns an empty tuple
|
||||
for backends without an equivalent knob.
|
||||
"""
|
||||
if kind == "low_latency":
|
||||
table = _ANN_TUNING_LOW_LATENCY
|
||||
elif kind == "high_recall":
|
||||
table = _ANN_TUNING_HIGH_RECALL
|
||||
else:
|
||||
raise ValueError(f"Unknown ANN tuning kind: {kind!r}")
|
||||
return table.get(_normalize_resolved(ext), ())
|
||||
|
||||
|
||||
def uses_per_bank_vector_indexes(ext: str) -> bool:
|
||||
"""Return whether the backend should create per-bank partial vector indexes."""
|
||||
return _normalize_resolved(ext) != "scann"
|
||||
|
||||
@@ -17,9 +17,7 @@ import asyncpg
|
||||
import typer
|
||||
|
||||
from ..config import DEFAULT_DATABASE_SCHEMA, HindsightConfig
|
||||
from ..engine.memory_engine import _current_schema
|
||||
from ..engine.schema import fq_table_explicit as _fq_table
|
||||
from ..engine.transfer import export_bank
|
||||
from ..extensions import TenantExtension, load_extension
|
||||
from ..pg0 import parse_pg0_url, resolve_database_url
|
||||
|
||||
@@ -32,60 +30,22 @@ logger = logging.getLogger(__name__)
|
||||
|
||||
app = typer.Typer(name="hindsight-admin", help="Hindsight administrative commands")
|
||||
|
||||
# Tables to backup/restore in foreign-key dependency order (parents first).
|
||||
# Restore COPYs in this order and TRUNCATEs in reverse, so every child must
|
||||
# appear after the tables it references.
|
||||
#
|
||||
# This must cover EVERY persistent PostgreSQL table in the schema — a missing
|
||||
# entry silently drops that table's data on restore (and, worse, restore's
|
||||
# `TRUNCATE banks CASCADE` wipes any FK-to-banks child like mental_models even
|
||||
# when it was never backed up). test_admin_backup_restore.py asserts this list
|
||||
# equals the live schema's tables, so adding a migration that creates a table
|
||||
# without adding it here fails CI. Oracle-only tables (e.g. observation_sources)
|
||||
# are intentionally absent — admin backup/restore is PostgreSQL-only.
|
||||
# Tables to backup/restore in dependency order
|
||||
# Import must happen in this order due to foreign key constraints
|
||||
BACKUP_TABLES = [
|
||||
"banks",
|
||||
"documents",
|
||||
"entities",
|
||||
"chunks",
|
||||
"memory_units",
|
||||
"invalidated_memory_units",
|
||||
"unit_entities",
|
||||
"entity_cooccurrences",
|
||||
"memory_links",
|
||||
"observation_history",
|
||||
"mental_models",
|
||||
"mental_model_history",
|
||||
"knowledge_pages",
|
||||
"directives",
|
||||
"async_operations",
|
||||
"webhooks",
|
||||
"file_storage",
|
||||
"audit_log",
|
||||
"llm_requests",
|
||||
"graph_maintenance_queue",
|
||||
]
|
||||
|
||||
MANIFEST_VERSION = "1"
|
||||
|
||||
|
||||
async def _admin_connect(db_url: str) -> asyncpg.Connection:
|
||||
"""Open a raw asyncpg connection to an admin DB URL.
|
||||
|
||||
``resolve_database_url`` handles both plain ``postgres://`` (passthrough) and
|
||||
``pg0://`` (boots the embedded server and returns its real libpq URL), so this
|
||||
is the only step needed to connect. JSON codecs are registered so ``jsonb``
|
||||
columns decode to Python objects (used by the export row dumps).
|
||||
"""
|
||||
is_pg0, instance_name, _ = parse_pg0_url(db_url)
|
||||
if is_pg0:
|
||||
typer.echo(f"Starting embedded PostgreSQL (instance: {instance_name})...")
|
||||
conn = await asyncpg.connect(await resolve_database_url(db_url))
|
||||
for type_name in ("json", "jsonb"):
|
||||
await conn.set_type_codec(type_name, encoder=json.dumps, decoder=json.loads, schema="pg_catalog")
|
||||
return conn
|
||||
|
||||
|
||||
async def _backup(database_url: str, output_path: Path, schema: str = "public") -> dict[str, Any]:
|
||||
"""Backup all tables to a zip file using binary COPY protocol."""
|
||||
conn = await asyncpg.connect(database_url)
|
||||
@@ -257,10 +217,14 @@ async def _run_migration(
|
||||
schema: str | None = None,
|
||||
base_schema: str = DEFAULT_DATABASE_SCHEMA,
|
||||
embedding_dimension: int | None = None,
|
||||
ensure_extensions: bool = True,
|
||||
) -> list[str]:
|
||||
"""Resolve database URL and run migrations for one schema or all discovered schemas."""
|
||||
from ..migrations import run_migrations_for_schemas
|
||||
from ..migrations import (
|
||||
ensure_embedding_dimension,
|
||||
ensure_text_search_extension,
|
||||
ensure_vector_extension,
|
||||
run_migrations,
|
||||
)
|
||||
|
||||
is_pg0, instance_name, _ = parse_pg0_url(db_url)
|
||||
if is_pg0:
|
||||
@@ -281,21 +245,31 @@ async def _run_migration(
|
||||
# Preserve order while removing duplicates.
|
||||
schemas = list(dict.fromkeys(schemas))
|
||||
|
||||
# Migrate up to `migration_concurrency` schemas at once (each in its own
|
||||
# process); within a schema the work stays sequential. Run off the event
|
||||
# loop so the process pool's blocking joins don't stall it.
|
||||
await asyncio.to_thread(
|
||||
run_migrations_for_schemas,
|
||||
resolved_url,
|
||||
schemas,
|
||||
concurrency=config.migration_concurrency,
|
||||
migration_database_url=config.migration_database_url,
|
||||
embedding_dimension=embedding_dimension,
|
||||
vector_extension=config.vector_extension,
|
||||
text_search_extension=config.text_search_extension,
|
||||
pg_search_tokenizer=config.text_search_extension_pg_search_tokenizer,
|
||||
ensure_extensions=ensure_extensions,
|
||||
)
|
||||
for schema in schemas:
|
||||
run_migrations(resolved_url, schema=schema, migration_database_url=config.migration_database_url)
|
||||
|
||||
if embedding_dimension is not None:
|
||||
for schema in schemas:
|
||||
ensure_embedding_dimension(
|
||||
resolved_url,
|
||||
embedding_dimension,
|
||||
schema=schema,
|
||||
vector_extension=config.vector_extension,
|
||||
)
|
||||
|
||||
for schema in schemas:
|
||||
ensure_vector_extension(
|
||||
resolved_url,
|
||||
vector_extension=config.vector_extension,
|
||||
schema=schema,
|
||||
)
|
||||
|
||||
for schema in schemas:
|
||||
ensure_text_search_extension(
|
||||
resolved_url,
|
||||
text_search_extension=config.text_search_extension,
|
||||
schema=schema,
|
||||
)
|
||||
|
||||
return schemas
|
||||
|
||||
@@ -313,18 +287,6 @@ def run_db_migration(
|
||||
"--embedding-dimension",
|
||||
help="Expected embedding dimension to enforce after migrations. Omit to skip dimension sync.",
|
||||
),
|
||||
skip_extension_reconcile: bool = typer.Option(
|
||||
False,
|
||||
"--skip-extension-reconcile",
|
||||
help=(
|
||||
"Skip the post-migration vector / text-search index reconcile. This step only does "
|
||||
"work when the configured backend (HINDSIGHT_API_VECTOR_EXTENSION / "
|
||||
"HINDSIGHT_API_TEXT_SEARCH_EXTENSION) differs from a schema's existing indexes — a "
|
||||
"rare, operator-driven change. Skipping it makes a no-change re-migration over many "
|
||||
"tenant schemas much faster. Only use when you have NOT changed the backend; a "
|
||||
"backend change still needs a normal run to reshape the indexes."
|
||||
),
|
||||
),
|
||||
):
|
||||
"""Run database migrations to the latest version."""
|
||||
config = HindsightConfig.from_env()
|
||||
@@ -338,8 +300,6 @@ def run_db_migration(
|
||||
typer.echo(f"Running database migrations for schema: {schema}...")
|
||||
else:
|
||||
typer.echo("Running database migrations for base schema and all discovered tenant schemas...")
|
||||
if skip_extension_reconcile:
|
||||
typer.echo("Skipping post-migration extension reconcile (--skip-extension-reconcile).")
|
||||
|
||||
schemas = asyncio.run(
|
||||
_run_migration(
|
||||
@@ -347,130 +307,12 @@ def run_db_migration(
|
||||
schema=schema,
|
||||
base_schema=config.database_schema,
|
||||
embedding_dimension=embedding_dimension,
|
||||
ensure_extensions=not skip_extension_reconcile,
|
||||
)
|
||||
)
|
||||
|
||||
typer.echo(f"Database migrations completed successfully for {len(schemas)} schema(s)")
|
||||
|
||||
|
||||
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)
|
||||
try:
|
||||
# export_bank resolves table names via fq_table (the _current_schema
|
||||
# contextvar); set it so the raw connection targets the right schema.
|
||||
_current_schema.set(schema)
|
||||
data = await export_bank(conn, bank_id, include_history=include_history)
|
||||
finally:
|
||||
await conn.close()
|
||||
|
||||
output.write_bytes(data)
|
||||
return len(data)
|
||||
|
||||
|
||||
@app.command(name="export-bank")
|
||||
def export_bank_command(
|
||||
bank_id: str = typer.Option(..., "--bank", "-b", help="Bank id to export."),
|
||||
output: Path = typer.Option(..., "--output", "-o", help="Path to write the .zip archive."),
|
||||
schema: str | None = typer.Option(
|
||||
None,
|
||||
"--schema",
|
||||
"-s",
|
||||
help="Database schema the bank lives in. Defaults to the configured base schema.",
|
||||
),
|
||||
include_history: bool = typer.Option(
|
||||
False,
|
||||
"--include-history",
|
||||
help="Also export operational history (audit_log, llm_requests). Off by default.",
|
||||
),
|
||||
):
|
||||
"""Export an entire bank to a portable ZIP (no embeddings — regenerated on import).
|
||||
|
||||
Carries documents, facts, observations, bank config, mental models, directives
|
||||
and webhooks so the bank can be imported into a new instance configured with a
|
||||
different embedding model / vector / text-search backend.
|
||||
"""
|
||||
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)
|
||||
|
||||
target_schema = schema or config.database_schema or DEFAULT_DATABASE_SCHEMA
|
||||
typer.echo(f"Exporting bank '{bank_id}' from schema '{target_schema}'...")
|
||||
|
||||
size = asyncio.run(_run_export_bank(config.database_url, bank_id, output, target_schema, include_history))
|
||||
|
||||
typer.echo(f"Exported bank '{bank_id}' to {output} ({size} bytes)")
|
||||
|
||||
|
||||
async def _run_import_bank(archive_path: Path, schema: str, target_bank_id: str | None, include_history: bool):
|
||||
"""Boot a MemoryEngine (for the target's embedding model) and restore a bank archive."""
|
||||
# MemoryEngine is heavy (loads embeddings); import it lazily so other admin
|
||||
# commands don't pay for it. _current_schema is imported at module top.
|
||||
from ..engine.memory_engine import MemoryEngine
|
||||
from ..models import RequestContext
|
||||
|
||||
archive_bytes = archive_path.read_bytes()
|
||||
# run_migrations=True so a fresh target instance is provisioned at this
|
||||
# instance's embedding dimension / vector / text-search backend before restore.
|
||||
engine = MemoryEngine(run_migrations=True)
|
||||
await engine.initialize()
|
||||
try:
|
||||
_current_schema.set(schema)
|
||||
context = RequestContext(internal=True, user_initiated=True)
|
||||
return await engine.import_bank_async(
|
||||
archive_bytes,
|
||||
context,
|
||||
target_bank_id=target_bank_id,
|
||||
include_history=include_history,
|
||||
)
|
||||
finally:
|
||||
await engine.close()
|
||||
|
||||
|
||||
@app.command(name="import-bank")
|
||||
def import_bank_command(
|
||||
archive: Path = typer.Option(..., "--archive", "-a", help="Path to the .zip produced by export-bank."),
|
||||
schema: str | None = typer.Option(
|
||||
None, "--schema", "-s", help="Target schema. Defaults to the configured base schema."
|
||||
),
|
||||
target_bank: str | None = typer.Option(
|
||||
None, "--target-bank", help="Override the bank id (defaults to the archive's source bank)."
|
||||
),
|
||||
include_history: bool = typer.Option(
|
||||
False, "--include-history", help="Also restore operational history if present in the archive."
|
||||
),
|
||||
):
|
||||
"""Restore a whole bank from an export-bank archive into THIS instance.
|
||||
|
||||
Re-embeds facts with this instance's configured embedding model and rebuilds
|
||||
links and indexes — the import half of a cross-instance migration. Run against
|
||||
an instance configured with the desired embedding / vector / text-search backend.
|
||||
The target bank must not already exist (import restores a whole bank, not a merge).
|
||||
"""
|
||||
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)
|
||||
|
||||
target_schema = schema or config.database_schema or DEFAULT_DATABASE_SCHEMA
|
||||
typer.echo(f"Importing bank archive '{archive}' into schema '{target_schema}'...")
|
||||
|
||||
result = asyncio.run(_run_import_bank(archive, target_schema, target_bank, include_history))
|
||||
|
||||
typer.echo(
|
||||
f"Imported bank '{result.bank_id}': {result.documents_imported} doc(s), "
|
||||
f"{result.facts_imported} fact(s), {result.observations_imported} observation(s), "
|
||||
f"{result.mental_models_imported} mental model(s), "
|
||||
f"{result.mental_model_history_imported} mm-history row(s), {result.directives_imported} directive(s), "
|
||||
f"{result.webhooks_imported} webhook(s), {result.history_rows_imported} history row(s)"
|
||||
)
|
||||
|
||||
|
||||
async def _decommission_worker(db_url: str, worker_id: str, schema: str = "public") -> int:
|
||||
"""Release all tasks owned by a worker, setting them back to pending status."""
|
||||
is_pg0, instance_name, _ = parse_pg0_url(db_url)
|
||||
|
||||
-105
@@ -1,105 +0,0 @@
|
||||
"""Add a composite index on memory_links(bank_id, link_type) (PostgreSQL).
|
||||
|
||||
``bank_id`` was added to ``memory_links`` in ``c5d6e7f8a9b0`` precisely so that
|
||||
bank-scoped reads (e.g. the stats endpoint) could filter on the link table
|
||||
directly instead of joining ``memory_units`` — that JOIN took 18+ seconds on
|
||||
banks with millions of links. The column landed without an index, so every
|
||||
``bank_id = $1`` predicate still falls back to a sequential scan over the whole
|
||||
table.
|
||||
|
||||
This adds the missing btree. It is composite on ``(bank_id, link_type)`` rather
|
||||
than ``bank_id`` alone because the hot query is the stats endpoint's
|
||||
``SELECT link_type, COUNT(*) ... WHERE bank_id = $1 GROUP BY link_type``: a
|
||||
``(bank_id, link_type)`` index serves that filter, grouping and count as an
|
||||
index-only scan, never touching the heap, whereas a ``bank_id``-only index would
|
||||
still have to read every matching row to recover ``link_type``. ``link_type`` is
|
||||
low-cardinality (only ``temporal``/``semantic``/``caused_by`` are written —
|
||||
entity edges were dropped in ``e9b2c7d1f3a4``), so the trailing column adds
|
||||
little to the index size while removing the heap fetch.
|
||||
|
||||
The Oracle baseline (``o1a2b3c4d5e6``) already creates ``idx_ml_bank_id`` on
|
||||
``memory_links(bank_id)``; that single-column index already covers Oracle's
|
||||
bank-scoped filter, so the Oracle slot here is intentionally absent and only the
|
||||
PostgreSQL dialect gets the composite index.
|
||||
|
||||
``memory_links`` can hold tens of millions of rows, so the index is built
|
||||
CONCURRENTLY to avoid taking a write lock on the table. CONCURRENTLY cannot run
|
||||
inside a transaction block, so the statement runs in an ``autocommit_block()``;
|
||||
``IF NOT EXISTS`` keeps it idempotent across retries and re-migrated tenant
|
||||
schemas. A CONCURRENTLY build interrupted partway (lock conflict, disk
|
||||
pressure, signal) leaves the index behind as *invalid*; ``IF NOT EXISTS`` would
|
||||
then skip over it forever, so the upgrade first drops any invalid leftover of
|
||||
this name before (re)creating it.
|
||||
|
||||
Revision ID: 2071c7518f88
|
||||
Revises: a1d3f5b7c9e2
|
||||
Create Date: 2026-06-16
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
from sqlalchemy import text
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "2071c7518f88"
|
||||
down_revision: str | Sequence[str] | None = "a1d3f5b7c9e2"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
_INDEX_NAME = "idx_memory_links_bank_id_link_type"
|
||||
|
||||
|
||||
def _get_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:
|
||||
bind = op.get_bind()
|
||||
# `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
|
||||
schema = _get_schema_prefix()
|
||||
|
||||
# CREATE INDEX CONCURRENTLY cannot run inside a transaction block; the
|
||||
# autocommit_block runs each statement outside Alembic's migration
|
||||
# transaction.
|
||||
with op.get_context().autocommit_block():
|
||||
# A CONCURRENTLY build that errored on a previous run leaves an INVALID
|
||||
# index of this name behind. `CREATE INDEX ... IF NOT EXISTS` would see
|
||||
# that relation and skip, so bank_id queries would keep seq-scanning.
|
||||
# Drop only the invalid leftover — never a healthy index — so the retry
|
||||
# actually rebuilds a usable one.
|
||||
leftover_invalid = bind.execute(
|
||||
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": _INDEX_NAME, "target_schema": target_schema},
|
||||
).scalar()
|
||||
if leftover_invalid:
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}{_INDEX_NAME}")
|
||||
|
||||
# IF NOT EXISTS keeps the create idempotent across retries and schemas.
|
||||
op.execute(f"CREATE INDEX CONCURRENTLY IF NOT EXISTS {_INDEX_NAME} ON {schema}memory_links(bank_id, link_type)")
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
with op.get_context().autocommit_block():
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}{_INDEX_NAME}")
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade)
|
||||
@@ -15,11 +15,6 @@ from pgvector.sqlalchemy import Vector
|
||||
from sqlalchemy import text
|
||||
from sqlalchemy.dialects import postgresql
|
||||
|
||||
from hindsight_api._pg_search import (
|
||||
PG_SEARCH_TOKENIZER_ENV,
|
||||
normalize_pg_search_tokenizer,
|
||||
pg_search_bm25_columns,
|
||||
)
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
@@ -88,7 +83,7 @@ def _vector_index_using_clause(ext: str) -> str:
|
||||
if ext == "pg_diskann":
|
||||
return "USING diskann (embedding vector_cosine_ops) WITH (max_neighbors = 50)"
|
||||
if ext == "vchord":
|
||||
return "USING vchordrq (embedding vector_cosine_ops)"
|
||||
return "USING vchordrq (embedding vector_l2_ops)"
|
||||
if ext == "scann":
|
||||
return "USING scann (embedding cosine) WITH (mode = 'AUTO')"
|
||||
return "USING hnsw (embedding vector_cosine_ops)"
|
||||
@@ -96,14 +91,9 @@ def _vector_index_using_clause(ext: str) -> str:
|
||||
|
||||
def _detect_text_search_extension() -> str:
|
||||
"""
|
||||
Detect or validate text search extension: 'native', 'vchord', 'pg_textsearch',
|
||||
'pgroonga', or 'pg_search'. Respects HINDSIGHT_API_TEXT_SEARCH_EXTENSION env var.
|
||||
Detect or validate text search extension: 'native', 'vchord', or 'pg_textsearch'.
|
||||
Respects HINDSIGHT_API_TEXT_SEARCH_EXTENSION env var.
|
||||
Creates the extension if needed.
|
||||
|
||||
pgroonga is treated as native here so the initial schema still creates valid
|
||||
tsvector columns. ensure_text_search_extension() at startup converts the
|
||||
schema to pgroonga structures (drops the tsvector column, builds a pgroonga
|
||||
index on the base text column).
|
||||
"""
|
||||
text_search_extension = os.getenv("HINDSIGHT_API_TEXT_SEARCH_EXTENSION", "native").lower()
|
||||
|
||||
@@ -131,35 +121,14 @@ def _detect_text_search_extension() -> str:
|
||||
# Extension truly doesn't exist - re-raise the error
|
||||
raise
|
||||
return "pg_textsearch"
|
||||
elif text_search_extension == "pg_search":
|
||||
# ParadeDB pg_search — true BM25 over base columns, Citus-compatible.
|
||||
try:
|
||||
op.execute("CREATE EXTENSION IF NOT EXISTS pg_search CASCADE")
|
||||
except Exception:
|
||||
# Extension might already exist or user lacks permissions - verify it exists
|
||||
conn = op.get_bind()
|
||||
result = conn.execute(text("SELECT 1 FROM pg_extension WHERE extname = 'pg_search'")).fetchone()
|
||||
if not result:
|
||||
# Extension truly doesn't exist - re-raise the error
|
||||
raise
|
||||
return "pg_search"
|
||||
elif text_search_extension == "native":
|
||||
return "native"
|
||||
elif text_search_extension == "pgroonga":
|
||||
# ensure_text_search_extension() at runtime converts to pgroonga.
|
||||
# Treat as native here so the initial schema still creates valid columns.
|
||||
return "native"
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Invalid HINDSIGHT_API_TEXT_SEARCH_EXTENSION: {text_search_extension}. "
|
||||
"Must be 'native', 'vchord', 'pg_textsearch', 'pgroonga', or 'pg_search'"
|
||||
f"Invalid HINDSIGHT_API_TEXT_SEARCH_EXTENSION: {text_search_extension}. Must be 'native', 'vchord', or 'pg_textsearch'"
|
||||
)
|
||||
|
||||
|
||||
def _pg_search_tokenizer() -> str:
|
||||
return normalize_pg_search_tokenizer(os.getenv(PG_SEARCH_TOKENIZER_ENV))
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
"""Upgrade schema - create all tables from scratch."""
|
||||
|
||||
@@ -315,9 +284,8 @@ def _pg_upgrade() -> None:
|
||||
ALTER TABLE memory_units
|
||||
ADD COLUMN search_vector bm25_catalog.bm25vector
|
||||
""")
|
||||
elif text_search_ext in ("pg_textsearch", "pg_search"):
|
||||
# Timescale pg_textsearch / ParadeDB pg_search: dummy TEXT column for
|
||||
# consistency (indexes operate on base columns directly).
|
||||
elif text_search_ext == "pg_textsearch":
|
||||
# Timescale pg_textsearch: dummy TEXT column for consistency (indexes operate on base columns directly)
|
||||
op.execute("""
|
||||
ALTER TABLE memory_units
|
||||
ADD COLUMN search_vector TEXT
|
||||
@@ -382,17 +350,6 @@ def _pg_upgrade() -> None:
|
||||
USING bm25(text)
|
||||
WITH (text_config='english')
|
||||
""")
|
||||
elif text_search_ext == "pg_search":
|
||||
# ParadeDB pg_search BM25 index on (id, text, context). The key_field
|
||||
# reloption is required and must match the table's primary key column.
|
||||
bm25_cols = pg_search_bm25_columns("id", ("text", "context"), _pg_search_tokenizer())
|
||||
op.execute(
|
||||
"""
|
||||
CREATE INDEX idx_memory_units_text_search ON memory_units
|
||||
USING bm25 ({bm25_cols})
|
||||
WITH (key_field='id')
|
||||
""".format(bm25_cols=bm25_cols)
|
||||
)
|
||||
else: # native
|
||||
# Native PostgreSQL GIN index
|
||||
op.execute("""
|
||||
|
||||
-85
@@ -1,85 +0,0 @@
|
||||
"""Repair: widen the remaining live ``bank_id`` columns from VARCHAR(64) to TEXT on PostgreSQL.
|
||||
|
||||
Follow-up to ``c3e5a7b9d1f4`` (issue #2106), which widened the two *history*
|
||||
tables (``observation_history``, ``mental_model_history``) to ``TEXT`` after the
|
||||
narrow ``VARCHAR(64)`` declaration bricked startup. The same VARCHAR(64) / TEXT
|
||||
inconsistency still affects the live tables that store a user-supplied
|
||||
``bank_id``:
|
||||
|
||||
* ``directives`` -- created VARCHAR(64) in ``p1k2l3m4n5o6``
|
||||
* ``mental_models`` -- VARCHAR(64) (origin ``pinned_reflections`` in
|
||||
``n9i0j1k2l3m4``; recreated in ``h3c4d5e6f7g8``)
|
||||
|
||||
``mental_model_versions`` is intentionally *not* widened here: it is created in
|
||||
``j5e6f7g8h9i0`` but dropped (``DROP TABLE ... CASCADE``) in ``o0j1k2l3m4n5`` and
|
||||
never recreated on the upgrade path, so it does not exist at head. Issuing
|
||||
``ALTER TABLE mental_model_versions ...`` would raise ``UndefinedTable`` and --
|
||||
because migrations run inside the lifespan-startup transaction -- roll the whole
|
||||
migration back, bricking the API. (It is unrelated to the live
|
||||
``mental_model_history`` table widened by ``c3e5a7b9d1f4``.)
|
||||
|
||||
``banks.bank_id`` is ``TEXT`` (unbounded), so a deployment can create a bank
|
||||
whose id exceeds 64 chars -- the 78-char hierarchical org-unit shape reported in
|
||||
issue #2106 -- and the bank insert succeeds. The next write that propagates that
|
||||
id (``create_directive``, ``create_mental_model`` / consolidation, or
|
||||
mental-model versioning) then aborts with::
|
||||
|
||||
psycopg2.errors.StringDataRightTruncation: value too long for type
|
||||
character varying(64)
|
||||
|
||||
i.e. a 500 on core write endpoints, instead of the startup brick that
|
||||
``c3e5a7b9d1f4`` already repaired.
|
||||
|
||||
``ALTER COLUMN ... TYPE TEXT`` is a no-op on a column that is already ``TEXT``,
|
||||
so every upgrade path converges on ``TEXT``. These tables are per-tenant (they
|
||||
live in each tenant schema, not ``public``), so this runs for every migrated
|
||||
schema via the search-path-aware prefix -- the same mechanism as
|
||||
``c3e5a7b9d1f4``.
|
||||
|
||||
PostgreSQL only: these tables are created by PostgreSQL-only migrations
|
||||
(``run_for_dialect(pg=...)``); on Oracle they are absent or already
|
||||
``VARCHAR2(256)`` (consistent, never truncates), so the Oracle slot is
|
||||
intentionally absent -- mirroring ``c3e5a7b9d1f4``.
|
||||
|
||||
Revision ID: a1d3f5b7c9e2
|
||||
Revises: c3e5a7b9d1f4
|
||||
Create Date: 2026-06-13
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "a1d3f5b7c9e2"
|
||||
down_revision: str | Sequence[str] | None = "c3e5a7b9d1f4"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _get_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 = _get_schema_prefix()
|
||||
op.execute(f"ALTER TABLE {schema}directives ALTER COLUMN bank_id TYPE TEXT")
|
||||
op.execute(f"ALTER TABLE {schema}mental_models ALTER COLUMN bank_id TYPE TEXT")
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
# No-op: narrowing back to VARCHAR(64) could truncate real data and would
|
||||
# re-introduce the bug this migration repairs. The column types are owned by
|
||||
# the migrations that created the tables.
|
||||
pass
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade)
|
||||
-29
@@ -7,7 +7,6 @@ the stored fact text.
|
||||
- vchord: text_signals included in tokenize() at insert time
|
||||
- native: search_vector GENERATED column regenerated to include text_signals
|
||||
- pg_textsearch: no change (index only supports a single base column)
|
||||
- pg_search: BM25 index dropped and recreated to include text_signals
|
||||
|
||||
Revision ID: a2b3c4d5e6f7
|
||||
Revises: z1u2v3w4x5y6
|
||||
@@ -19,11 +18,6 @@ from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api._pg_search import (
|
||||
PG_SEARCH_TOKENIZER_ENV,
|
||||
normalize_pg_search_tokenizer,
|
||||
pg_search_bm25_columns,
|
||||
)
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "a2b3c4d5e6f7"
|
||||
@@ -41,10 +35,6 @@ def _detect_text_search_extension() -> str:
|
||||
return os.getenv("HINDSIGHT_API_TEXT_SEARCH_EXTENSION", "native").lower()
|
||||
|
||||
|
||||
def _pg_search_tokenizer() -> str:
|
||||
return normalize_pg_search_tokenizer(os.getenv(PG_SEARCH_TOKENIZER_ENV))
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
table = f"{schema}memory_units"
|
||||
@@ -72,16 +62,6 @@ def _pg_upgrade() -> None:
|
||||
CREATE INDEX IF NOT EXISTS idx_memory_units_text_search
|
||||
ON {table} USING gin(search_vector)
|
||||
""")
|
||||
elif text_search_ext == "pg_search":
|
||||
# ParadeDB pg_search: drop the existing BM25 index and recreate it
|
||||
# to include text_signals alongside text and context.
|
||||
bm25_cols = pg_search_bm25_columns("id", ("text", "context", "text_signals"), _pg_search_tokenizer())
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema}idx_memory_units_text_search")
|
||||
op.execute(f"""
|
||||
CREATE INDEX idx_memory_units_text_search ON {table}
|
||||
USING bm25 ({bm25_cols})
|
||||
WITH (key_field='id')
|
||||
""")
|
||||
|
||||
# vchord: tokenize() call in fact_storage.py is updated to include text_signals at insert time
|
||||
# pg_textsearch: no change — index operates on the base `text` column only
|
||||
@@ -106,15 +86,6 @@ def _pg_downgrade() -> None:
|
||||
CREATE INDEX idx_memory_units_text_search
|
||||
ON {table} USING gin(search_vector)
|
||||
""")
|
||||
elif text_search_ext == "pg_search":
|
||||
# Restore the original (id, text, context) BM25 index without text_signals.
|
||||
bm25_cols = pg_search_bm25_columns("id", ("text", "context"), _pg_search_tokenizer())
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema}idx_memory_units_text_search")
|
||||
op.execute(f"""
|
||||
CREATE INDEX idx_memory_units_text_search ON {table}
|
||||
USING bm25 ({bm25_cols})
|
||||
WITH (key_field='id')
|
||||
""")
|
||||
|
||||
op.execute(f"ALTER TABLE {table} DROP COLUMN IF EXISTS text_signals")
|
||||
|
||||
|
||||
+10
-10
@@ -40,20 +40,20 @@ def _get_schema_prefix() -> str:
|
||||
def _pg_upgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
|
||||
# CREATE INDEX CONCURRENTLY cannot run inside a transaction block; an
|
||||
# autocommit_block runs it outside Alembic's migration transaction.
|
||||
with op.get_context().autocommit_block():
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_units_source_memory_ids "
|
||||
f"ON {schema}memory_units USING GIN (source_memory_ids) "
|
||||
f"WHERE source_memory_ids IS NOT NULL"
|
||||
)
|
||||
# CREATE INDEX CONCURRENTLY cannot run inside a transaction block.
|
||||
# Commit the current Alembic transaction first.
|
||||
op.execute("COMMIT")
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_units_source_memory_ids "
|
||||
f"ON {schema}memory_units USING GIN (source_memory_ids) "
|
||||
f"WHERE source_memory_ids IS NOT NULL"
|
||||
)
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
with op.get_context().autocommit_block():
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_units_source_memory_ids")
|
||||
op.execute("COMMIT")
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_units_source_memory_ids")
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
|
||||
+1
-1
@@ -63,7 +63,7 @@ def _vector_index_using_clause(ext: str) -> str:
|
||||
if ext == "pgvectorscale":
|
||||
return "USING diskann (embedding vector_cosine_ops) WITH (num_neighbors = 50)"
|
||||
if ext == "vchord":
|
||||
return "USING vchordrq (embedding vector_cosine_ops)"
|
||||
return "USING vchordrq (embedding vector_l2_ops)"
|
||||
if ext == "scann":
|
||||
return "USING scann (embedding cosine) WITH (mode = 'AUTO')"
|
||||
return "USING hnsw (embedding vector_cosine_ops)"
|
||||
|
||||
-52
@@ -1,52 +0,0 @@
|
||||
"""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)
|
||||
-253
@@ -1,253 +0,0 @@
|
||||
"""Move mental-model and observation history into dedicated tables.
|
||||
|
||||
Both histories were accumulated in a single JSONB/CLOB ``history`` column
|
||||
(``mental_models.history`` and ``memory_units.history``), appended to on every
|
||||
update. That design has two problems:
|
||||
|
||||
1. **Unbounded growth on observations.** The observation write path appended a
|
||||
snapshot on every update with no cap at all, so a frequently-reinforced
|
||||
observation grew its ``history`` array until it crossed Postgres's hard 256MB
|
||||
jsonb limit (SQLSTATE 54000), after which every further UPDATE failed and the
|
||||
row was stuck.
|
||||
2. **Wrong-axis cap on mental models.** The mental-model cap bounded the *number*
|
||||
of entries (50), not their *size* — a single large reflect snapshot could
|
||||
still blow the budget — and rewrote the whole array (plus TOAST) on every
|
||||
refresh, defeating HOT updates.
|
||||
|
||||
This migration creates one row per history entry in two dedicated tables, with
|
||||
an index that makes "most recent N for this item" cheap, then drops the old
|
||||
columns. The cap is now enforced at write time as a bounded DELETE of the
|
||||
oldest over-cap rows (see config ``*_HISTORY_MAX_ENTRIES``).
|
||||
|
||||
Revision ID: a7b8c9d0e1f2
|
||||
Revises: d3e4f5a6b7c8
|
||||
Create Date: 2026-06-05
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "a7b8c9d0e1f2"
|
||||
down_revision: str | Sequence[str] | None = "d3e4f5a6b7c8"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _get_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 ""
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# PostgreSQL
|
||||
# ---------------------------------------------------------------------------
|
||||
def _pg_upgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
|
||||
# Both tables share the same shape: surrogate id, FK to the parent, bank_id,
|
||||
# the snapshot payload as a single JSONB ``content`` blob, and changed_at.
|
||||
# The payload is per-row (one change per row) so it stays small — this is NOT
|
||||
# the old single-column-grows-forever design; growth is bounded by row count
|
||||
# plus the write-time cap. Folding the previous_* fields into one JSONB keeps
|
||||
# the schema dialect-simple (no array columns) and flexible.
|
||||
|
||||
# --- mental_model_history -------------------------------------------------
|
||||
# content: {"previous_content": ..., "previous_reflect_response": {...}}
|
||||
op.execute(
|
||||
f"""
|
||||
CREATE TABLE IF NOT EXISTS {schema}mental_model_history (
|
||||
id BIGINT GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY,
|
||||
mental_model_id VARCHAR(64) NOT NULL,
|
||||
bank_id TEXT NOT NULL,
|
||||
content JSONB NOT NULL,
|
||||
changed_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT now(),
|
||||
FOREIGN KEY (mental_model_id, bank_id)
|
||||
REFERENCES {schema}mental_models(id, bank_id) ON DELETE CASCADE
|
||||
)
|
||||
"""
|
||||
)
|
||||
op.execute(
|
||||
"CREATE INDEX IF NOT EXISTS idx_mm_history_model "
|
||||
f"ON {schema}mental_model_history (bank_id, mental_model_id, changed_at DESC, id DESC)"
|
||||
)
|
||||
|
||||
# --- observation_history --------------------------------------------------
|
||||
# content: {"previous_text", "previous_tags", "previous_occurred_start",
|
||||
# "previous_occurred_end", "previous_mentioned_at", "new_source_memory_ids"}
|
||||
op.execute(
|
||||
f"""
|
||||
CREATE TABLE IF NOT EXISTS {schema}observation_history (
|
||||
id BIGINT GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY,
|
||||
observation_id UUID NOT NULL,
|
||||
bank_id TEXT NOT NULL,
|
||||
content JSONB NOT NULL,
|
||||
changed_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT now(),
|
||||
FOREIGN KEY (observation_id)
|
||||
REFERENCES {schema}memory_units(id) ON DELETE CASCADE
|
||||
)
|
||||
"""
|
||||
)
|
||||
op.execute(
|
||||
"CREATE INDEX IF NOT EXISTS idx_observation_history_obs "
|
||||
f"ON {schema}observation_history (observation_id, changed_at DESC, id DESC)"
|
||||
)
|
||||
|
||||
# --- backfill mental models ----------------------------------------------
|
||||
# Explode each row's history array into rows, preserving chronological order
|
||||
# via WITH ORDINALITY so the IDENTITY id tie-breaks oldest->newest correctly.
|
||||
# changed_at is promoted to its own column; the rest of the element becomes
|
||||
# ``content`` (the ``- 'changed_at'`` strips the now-redundant key).
|
||||
op.execute(
|
||||
f"""
|
||||
INSERT INTO {schema}mental_model_history (mental_model_id, bank_id, content, changed_at)
|
||||
SELECT mm.id, mm.bank_id,
|
||||
e - 'changed_at',
|
||||
COALESCE(NULLIF(e->>'changed_at', '')::timestamptz, now())
|
||||
FROM {schema}mental_models mm
|
||||
CROSS JOIN LATERAL jsonb_array_elements(mm.history) WITH ORDINALITY a(e, ord)
|
||||
WHERE mm.history IS NOT NULL
|
||||
AND jsonb_typeof(mm.history) = 'array'
|
||||
AND jsonb_array_length(mm.history) > 0
|
||||
ORDER BY mm.id, mm.bank_id, ord
|
||||
"""
|
||||
)
|
||||
|
||||
# --- backfill observations -----------------------------------------------
|
||||
op.execute(
|
||||
f"""
|
||||
INSERT INTO {schema}observation_history (observation_id, bank_id, content, changed_at)
|
||||
SELECT mu.id, mu.bank_id,
|
||||
e - 'changed_at',
|
||||
COALESCE(NULLIF(e->>'changed_at', '')::timestamptz, now())
|
||||
FROM {schema}memory_units mu
|
||||
CROSS JOIN LATERAL jsonb_array_elements(mu.history) WITH ORDINALITY a(e, ord)
|
||||
WHERE mu.fact_type = 'observation'
|
||||
AND mu.history IS NOT NULL
|
||||
AND jsonb_typeof(mu.history) = 'array'
|
||||
AND jsonb_array_length(mu.history) > 0
|
||||
ORDER BY mu.id, ord
|
||||
"""
|
||||
)
|
||||
|
||||
# --- drop the legacy columns ---------------------------------------------
|
||||
op.execute(f"ALTER TABLE {schema}mental_models DROP COLUMN IF EXISTS history")
|
||||
op.execute(f"ALTER TABLE {schema}memory_units DROP COLUMN IF EXISTS history")
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
|
||||
# Re-add the columns (empty — historical content is not reconstructed back
|
||||
# into the array form; the dedicated tables are dropped below).
|
||||
op.execute(f"ALTER TABLE {schema}mental_models ADD COLUMN IF NOT EXISTS history JSONB DEFAULT '[]'::jsonb")
|
||||
op.execute(f"ALTER TABLE {schema}memory_units ADD COLUMN IF NOT EXISTS history JSONB DEFAULT '[]'::jsonb")
|
||||
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema}idx_observation_history_obs")
|
||||
op.execute(f"DROP TABLE IF EXISTS {schema}observation_history")
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema}idx_mm_history_model")
|
||||
op.execute(f"DROP TABLE IF EXISTS {schema}mental_model_history")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Oracle 23ai
|
||||
# ---------------------------------------------------------------------------
|
||||
def _oracle_upgrade() -> None:
|
||||
# Same single-JSONB shape as PG: ``content`` holds the snapshot payload as a
|
||||
# CLOB IS JSON. The legacy per-element JSON object (minus changed_at, promoted
|
||||
# to its own column) is carried through verbatim on backfill — the array
|
||||
# columns the previous design needed are gone.
|
||||
op.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS mental_model_history (
|
||||
id NUMBER GENERATED BY DEFAULT AS IDENTITY,
|
||||
mental_model_id VARCHAR2(256) NOT NULL,
|
||||
bank_id VARCHAR2(256) NOT NULL,
|
||||
content CLOB NOT NULL
|
||||
CONSTRAINT mmh_content_json CHECK (content IS JSON),
|
||||
changed_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,
|
||||
CONSTRAINT pk_mental_model_history PRIMARY KEY (id),
|
||||
CONSTRAINT fk_mmh_model FOREIGN KEY (mental_model_id, bank_id)
|
||||
REFERENCES mental_models(id, bank_id) ON DELETE CASCADE
|
||||
)
|
||||
"""
|
||||
)
|
||||
op.execute(
|
||||
"CREATE INDEX idx_mm_history_model ON mental_model_history (bank_id, mental_model_id, changed_at DESC, id DESC)"
|
||||
)
|
||||
|
||||
op.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS observation_history (
|
||||
id NUMBER GENERATED BY DEFAULT AS IDENTITY,
|
||||
observation_id RAW(16) NOT NULL,
|
||||
bank_id VARCHAR2(256) NOT NULL,
|
||||
content CLOB NOT NULL
|
||||
CONSTRAINT oh_content_json CHECK (content IS JSON),
|
||||
changed_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,
|
||||
CONSTRAINT pk_observation_history PRIMARY KEY (id),
|
||||
CONSTRAINT fk_oh_obs FOREIGN KEY (observation_id)
|
||||
REFERENCES memory_units(id) ON DELETE CASCADE
|
||||
)
|
||||
"""
|
||||
)
|
||||
op.execute(
|
||||
"CREATE INDEX idx_observation_history_obs ON observation_history (observation_id, changed_at DESC, id DESC)"
|
||||
)
|
||||
|
||||
bind = op.get_bind()
|
||||
|
||||
# Backfill via JSON_TABLE. ``content`` is the whole element (FORMAT JSON PATH
|
||||
# '$'); changed_at is also promoted to its own column. Backfilled content may
|
||||
# therefore still carry a redundant changed_at key, which the read path
|
||||
# ignores in favour of the column — harmless, and avoids JSON surgery here.
|
||||
bind.exec_driver_sql(
|
||||
"""
|
||||
INSERT INTO mental_model_history (mental_model_id, bank_id, content, changed_at)
|
||||
SELECT mm.id, mm.bank_id, jt.content, NVL(jt.changed_at, SYSTIMESTAMP)
|
||||
FROM mental_models mm,
|
||||
JSON_TABLE(mm.history, '$[*]' COLUMNS (
|
||||
seq FOR ORDINALITY,
|
||||
content CLOB FORMAT JSON PATH '$',
|
||||
changed_at TIMESTAMP WITH TIME ZONE PATH '$.changed_at'
|
||||
)) jt
|
||||
WHERE mm.history IS NOT NULL
|
||||
ORDER BY mm.id, mm.bank_id, jt.seq
|
||||
"""
|
||||
)
|
||||
|
||||
bind.exec_driver_sql(
|
||||
"""
|
||||
INSERT INTO observation_history (observation_id, bank_id, content, changed_at)
|
||||
SELECT mu.id, mu.bank_id, jt.content, NVL(jt.changed_at, SYSTIMESTAMP)
|
||||
FROM memory_units mu,
|
||||
JSON_TABLE(mu.history, '$[*]' COLUMNS (
|
||||
seq FOR ORDINALITY,
|
||||
content CLOB FORMAT JSON PATH '$',
|
||||
changed_at TIMESTAMP WITH TIME ZONE PATH '$.changed_at'
|
||||
)) jt
|
||||
WHERE mu.fact_type = 'observation' AND mu.history IS NOT NULL
|
||||
ORDER BY mu.id, jt.seq
|
||||
"""
|
||||
)
|
||||
|
||||
op.execute("ALTER TABLE mental_models DROP COLUMN history")
|
||||
op.execute("ALTER TABLE memory_units DROP COLUMN history")
|
||||
|
||||
|
||||
def _oracle_downgrade() -> None:
|
||||
op.execute("ALTER TABLE mental_models ADD history CLOB DEFAULT '[]' NOT NULL")
|
||||
op.execute("ALTER TABLE memory_units ADD history CLOB DEFAULT '[]'")
|
||||
op.execute("DROP TABLE observation_history CASCADE CONSTRAINTS")
|
||||
op.execute("DROP TABLE mental_model_history CASCADE CONSTRAINTS")
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade, oracle=_oracle_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade, oracle=_oracle_downgrade)
|
||||
-110
@@ -1,110 +0,0 @@
|
||||
"""Add knowledge_pages table (knowledge-base hierarchy).
|
||||
|
||||
The knowledge base organizes synthesized mental models into a navigable tree of
|
||||
**folders** and **pages**. A page references the mental model that holds its
|
||||
content (``mental_model_id``); a folder is a pure container (``mental_model_id``
|
||||
NULL). Hierarchy is a single self-referential ``parent_id`` so folders can nest
|
||||
arbitrarily. Content stays in ``mental_models`` — this table is metadata + tree
|
||||
structure only.
|
||||
|
||||
Revision ID: a9b8c7d6e5f4
|
||||
Revises: b57a7c9e0d13
|
||||
Create Date: 2026-06-25
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "a9b8c7d6e5f4"
|
||||
down_revision: str | Sequence[str] | None = "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()
|
||||
# parent_id self-FK cascades so deleting a folder row removes its whole
|
||||
# subtree of rows in one shot. The mental_model FK is composite (matches the
|
||||
# mental_models (id, bank_id) PK) and cascades too, so deleting a page's
|
||||
# mental model removes the page row — folders skip the FK because a NULL
|
||||
# column in a composite FK is not enforced (MATCH SIMPLE).
|
||||
op.execute(
|
||||
f"""
|
||||
CREATE TABLE IF NOT EXISTS {schema}knowledge_pages (
|
||||
id VARCHAR(64) NOT NULL,
|
||||
bank_id TEXT NOT NULL,
|
||||
parent_id VARCHAR(64),
|
||||
kind VARCHAR(16) NOT NULL,
|
||||
name TEXT NOT NULL,
|
||||
mental_model_id VARCHAR(64),
|
||||
sort_order INTEGER NOT NULL DEFAULT 0,
|
||||
created_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT now(),
|
||||
updated_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT now(),
|
||||
CONSTRAINT pk_knowledge_pages PRIMARY KEY (id),
|
||||
CONSTRAINT ck_knowledge_pages_kind CHECK (kind IN ('folder', 'page')),
|
||||
CONSTRAINT fk_kp_bank FOREIGN KEY (bank_id)
|
||||
REFERENCES {schema}banks(bank_id) ON DELETE CASCADE,
|
||||
CONSTRAINT fk_kp_parent FOREIGN KEY (parent_id)
|
||||
REFERENCES {schema}knowledge_pages(id) ON DELETE CASCADE,
|
||||
CONSTRAINT fk_kp_mm FOREIGN KEY (mental_model_id, bank_id)
|
||||
REFERENCES {schema}mental_models(id, bank_id) ON DELETE CASCADE
|
||||
)
|
||||
"""
|
||||
)
|
||||
op.execute(
|
||||
f"CREATE INDEX IF NOT EXISTS idx_kp_bank_parent ON {schema}knowledge_pages (bank_id, parent_id, sort_order)"
|
||||
)
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
schema = _pg_schema_prefix()
|
||||
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:
|
||||
op.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS knowledge_pages (
|
||||
id VARCHAR2(64) NOT NULL,
|
||||
bank_id VARCHAR2(256) NOT NULL,
|
||||
parent_id VARCHAR2(64),
|
||||
kind VARCHAR2(16) NOT NULL,
|
||||
name CLOB NOT NULL,
|
||||
mental_model_id VARCHAR2(64),
|
||||
sort_order NUMBER DEFAULT 0 NOT NULL,
|
||||
created_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,
|
||||
updated_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,
|
||||
CONSTRAINT pk_knowledge_pages PRIMARY KEY (id),
|
||||
CONSTRAINT ck_knowledge_pages_kind CHECK (kind IN ('folder', 'page')),
|
||||
CONSTRAINT fk_kp_bank FOREIGN KEY (bank_id)
|
||||
REFERENCES banks(bank_id) ON DELETE CASCADE,
|
||||
CONSTRAINT fk_kp_parent FOREIGN KEY (parent_id)
|
||||
REFERENCES knowledge_pages(id) ON DELETE CASCADE,
|
||||
CONSTRAINT fk_kp_mm FOREIGN KEY (mental_model_id, bank_id)
|
||||
REFERENCES mental_models(id, bank_id) ON DELETE CASCADE
|
||||
)
|
||||
"""
|
||||
)
|
||||
op.execute("CREATE INDEX idx_kp_bank_parent ON knowledge_pages (bank_id, parent_id, sort_order)")
|
||||
|
||||
|
||||
def _oracle_downgrade() -> None:
|
||||
op.execute("DROP TABLE knowledge_pages CASCADE CONSTRAINTS")
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade, oracle=_oracle_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade, oracle=_oracle_downgrade)
|
||||
-156
@@ -1,156 +0,0 @@
|
||||
"""Repair: install maintenance routines on the ``public`` / base-schema run.
|
||||
|
||||
The original maintenance-routines migration (``e5f6a7b8c9d0``) only created the
|
||||
shared ``public.banks_needing_consolidation()`` and
|
||||
``public.schemas_with_expired_rows(...)`` routines when the run had *no*
|
||||
``target_schema`` at all. But the single-tenant runtime always migrates an
|
||||
explicit schema — which defaults to ``public`` — so on every default
|
||||
PostgreSQL deployment the migration was stamped as applied while the functions
|
||||
were never created. Background maintenance then logs::
|
||||
|
||||
Retention sweep failed for llm_requests: function public.schemas_with_expired_rows(...) does not exist
|
||||
Consolidation reconcile discovery failed: function public.banks_needing_consolidation() does not exist
|
||||
|
||||
See https://github.com/vectorize-io/hindsight/issues/2056.
|
||||
|
||||
Because ``e5f6a7b8c9d0`` is already stamped on affected ``0.8.0`` databases,
|
||||
editing it would not re-run it there. This forward migration re-installs the
|
||||
functions idempotently (``CREATE OR REPLACE``) on the run that targets the
|
||||
shared ``public`` schema (base run with no ``target_schema``, or an explicit
|
||||
``target_schema=public``), self-healing already-upgraded deployments and
|
||||
covering fresh upgrades from earlier versions.
|
||||
|
||||
Per-tenant runs against a non-``public`` schema still skip it: re-issuing
|
||||
``CREATE OR REPLACE FUNCTION public....`` from each concurrent tenant migration
|
||||
aborts with ``tuple concurrently updated`` on the ``pg_proc`` catalog row, and
|
||||
the base/public run has already created the functions for every tenant to use.
|
||||
Runs that target ``public`` are serialized by the per-schema migration advisory
|
||||
lock, so only one wins the create.
|
||||
|
||||
PostgreSQL only — the worker poller and these tables are not wired for Oracle,
|
||||
so the Oracle slot is intentionally absent (mirrors ``e5f6a7b8c9d0``).
|
||||
|
||||
Revision ID: b2d4f6a8c1e3
|
||||
Revises: e5f6a7b8c9d0
|
||||
Create Date: 2026-06-08
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "b2d4f6a8c1e3"
|
||||
down_revision: str | Sequence[str] | None = "e5f6a7b8c9d0"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _should_install_public_routines(target_schema: str | None) -> bool:
|
||||
"""True for the run that must (re)create the shared ``public.*`` routines.
|
||||
|
||||
The routines physically live in ``public`` (hard-coded ``public.`` qualifier
|
||||
in the SQL below), so they must be installed exactly once — on the base run
|
||||
(no ``target_schema``) or on the run that explicitly targets ``public``. A
|
||||
run against any other tenant schema skips it to avoid concurrent
|
||||
``CREATE OR REPLACE`` on the same ``pg_proc`` row.
|
||||
"""
|
||||
return not target_schema or target_schema == "public"
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
if not _should_install_public_routines(context.config.get_main_option("target_schema")):
|
||||
return
|
||||
# Banks with eligible-but-unscheduled facts and no in-flight consolidation.
|
||||
# Auto-consolidation is filtered here only at the bank level (cheap prune);
|
||||
# the full hierarchical resolution (global -> tenant -> bank, plus
|
||||
# enable_observations) is done by the caller for the small returned set.
|
||||
op.execute(
|
||||
"""
|
||||
CREATE OR REPLACE FUNCTION public.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
|
||||
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);
|
||||
END LOOP;
|
||||
END;
|
||||
$fn$;
|
||||
"""
|
||||
)
|
||||
|
||||
# Schemas holding at least one row of p_table older than p_days. p_ts_col is
|
||||
# the timestamp column to compare. Returns nothing when p_days <= 0
|
||||
# (retention disabled).
|
||||
op.execute(
|
||||
"""
|
||||
CREATE OR REPLACE FUNCTION public.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
|
||||
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;
|
||||
IF has_expired THEN
|
||||
RETURN NEXT sch;
|
||||
END IF;
|
||||
END LOOP;
|
||||
END;
|
||||
$fn$;
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
# No-op: ``e5f6a7b8c9d0`` owns the lifecycle of these functions and drops
|
||||
# them on its own downgrade. This migration only ever (re)creates them, so
|
||||
# there is nothing to undo without racing that migration's DROP.
|
||||
pass
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade)
|
||||
+27
-25
@@ -37,35 +37,37 @@ def _get_schema_prefix() -> str:
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
# CREATE INDEX CONCURRENTLY cannot run inside a transaction block; an
|
||||
# autocommit_block runs each statement outside Alembic's migration transaction.
|
||||
with op.get_context().autocommit_block():
|
||||
# Partial index on occurred_start (covers "occurred_start BETWEEN $4 AND $5")
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_units_bank_occurred_start "
|
||||
f"ON {schema}memory_units(bank_id, fact_type, occurred_start) "
|
||||
f"WHERE occurred_start IS NOT NULL"
|
||||
)
|
||||
# Partial index on occurred_end (covers "occurred_end BETWEEN $4 AND $5")
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_units_bank_occurred_end "
|
||||
f"ON {schema}memory_units(bank_id, fact_type, occurred_end) "
|
||||
f"WHERE occurred_end IS NOT NULL"
|
||||
)
|
||||
# Partial index on mentioned_at (covers "mentioned_at BETWEEN $4 AND $5")
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_units_bank_mentioned_at "
|
||||
f"ON {schema}memory_units(bank_id, fact_type, mentioned_at) "
|
||||
f"WHERE mentioned_at IS NOT NULL"
|
||||
)
|
||||
# Partial index on occurred_start (covers "occurred_start BETWEEN $4 AND $5")
|
||||
op.execute("COMMIT")
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_units_bank_occurred_start "
|
||||
f"ON {schema}memory_units(bank_id, fact_type, occurred_start) "
|
||||
f"WHERE occurred_start IS NOT NULL"
|
||||
)
|
||||
# Partial index on occurred_end (covers "occurred_end BETWEEN $4 AND $5")
|
||||
op.execute("COMMIT")
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_units_bank_occurred_end "
|
||||
f"ON {schema}memory_units(bank_id, fact_type, occurred_end) "
|
||||
f"WHERE occurred_end IS NOT NULL"
|
||||
)
|
||||
# Partial index on mentioned_at (covers "mentioned_at BETWEEN $4 AND $5")
|
||||
op.execute("COMMIT")
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_units_bank_mentioned_at "
|
||||
f"ON {schema}memory_units(bank_id, fact_type, mentioned_at) "
|
||||
f"WHERE mentioned_at IS NOT NULL"
|
||||
)
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
with op.get_context().autocommit_block():
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_units_bank_mentioned_at")
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_units_bank_occurred_end")
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_units_bank_occurred_start")
|
||||
op.execute("COMMIT")
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_units_bank_mentioned_at")
|
||||
op.execute("COMMIT")
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_units_bank_occurred_end")
|
||||
op.execute("COMMIT")
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_units_bank_occurred_start")
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
|
||||
-61
@@ -1,61 +0,0 @@
|
||||
"""Add bank_stats_cache table for distributed get_bank_stats caching
|
||||
|
||||
Revision ID: b57a7c9e0d13
|
||||
Revises: c3f7a1b9d2e4
|
||||
Create Date: 2026-07-01
|
||||
|
||||
get_bank_stats aggregates over memory_links / unit_entities — a multi-second scan
|
||||
on banks with millions of rows. The result was cached per-process (in-memory), so
|
||||
every API worker recomputed it once per TTL and the first caller after expiry
|
||||
stalled. This table backs a shared, cross-process TTL cache: one worker's compute
|
||||
is written here and served to all the others.
|
||||
|
||||
PostgreSQL only. Oracle keeps the in-process cache (the runtime picks the backing
|
||||
store by dialect), so the Oracle upgrade slot is intentionally absent.
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "b57a7c9e0d13"
|
||||
down_revision: str | Sequence[str] | None = "c3f7a1b9d2e4"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _get_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 = _get_schema_prefix()
|
||||
# One row per bank: payload is the full get_bank_stats result, computed_at
|
||||
# drives logical TTL expiry. Rows are overwritten in place (ON CONFLICT), so
|
||||
# the table never grows beyond the number of banks and needs no purge job.
|
||||
op.execute(
|
||||
f"""
|
||||
CREATE TABLE IF NOT EXISTS {schema}bank_stats_cache (
|
||||
bank_id TEXT PRIMARY KEY,
|
||||
payload JSONB NOT NULL,
|
||||
computed_at TIMESTAMPTZ NOT NULL DEFAULT now()
|
||||
)
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
op.execute(f"DROP TABLE IF EXISTS {schema}bank_stats_cache")
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade) # oracle slot intentionally absent → no-op
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade)
|
||||
-106
@@ -1,106 +0,0 @@
|
||||
"""Add graph_maintenance_queue table
|
||||
|
||||
Queue of memory_units whose outgoing temporal/semantic links lost a
|
||||
neighbour to a delete. Drained by the async graph_maintenance worker,
|
||||
which tops the unit's links back up using the same probes retain runs.
|
||||
|
||||
The queue only targets the link-recompute pass. The worker also runs
|
||||
bank-wide sweeps (orphan-entity prune, stale-cooccurrence prune) on each
|
||||
invocation; those don't need per-target queueing.
|
||||
|
||||
Revision ID: b5a4c3e2f1d8
|
||||
Revises: e9b2c7d1f3a4
|
||||
Create Date: 2026-05-27
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "b5a4c3e2f1d8"
|
||||
down_revision: str | Sequence[str] | None = "e9b2c7d1f3a4"
|
||||
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()
|
||||
# Composite PK gives us natural ON CONFLICT DO NOTHING dedup when the same
|
||||
# unit is enqueued from overlapping deletes. No FK to memory_units: if the
|
||||
# unit is deleted between enqueue and drain, the worker observes it's gone
|
||||
# and skips — a cascade would erase the work order, but that work has
|
||||
# already been satisfied (no surviving row to maintain).
|
||||
op.execute(
|
||||
f"""
|
||||
CREATE TABLE IF NOT EXISTS {schema}graph_maintenance_queue (
|
||||
bank_id TEXT NOT NULL,
|
||||
unit_id UUID NOT NULL,
|
||||
enqueued_at TIMESTAMPTZ NOT NULL DEFAULT now(),
|
||||
PRIMARY KEY (bank_id, unit_id)
|
||||
)
|
||||
"""
|
||||
)
|
||||
op.execute(
|
||||
f"""
|
||||
CREATE INDEX IF NOT EXISTS idx_graph_maintenance_queue_bank_enqueued
|
||||
ON {schema}graph_maintenance_queue (bank_id, enqueued_at)
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
schema = _pg_schema_prefix()
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema}idx_graph_maintenance_queue_bank_enqueued")
|
||||
op.execute(f"DROP TABLE IF EXISTS {schema}graph_maintenance_queue")
|
||||
|
||||
|
||||
def _oracle_execute_ignoring_955(sql: str) -> None:
|
||||
"""Run a CREATE statement and swallow ORA-00955 (object already exists).
|
||||
|
||||
Mirrors the helper in the Oracle baseline migration so reruns stay safe
|
||||
on a database where the table was created by an earlier partial run.
|
||||
"""
|
||||
block = (
|
||||
"BEGIN "
|
||||
"EXECUTE IMMEDIATE :stmt; "
|
||||
"EXCEPTION WHEN OTHERS THEN "
|
||||
"IF SQLCODE = -955 THEN NULL; ELSE RAISE; END IF; "
|
||||
"END;"
|
||||
)
|
||||
op.get_bind().exec_driver_sql(block, {"stmt": sql.strip()})
|
||||
|
||||
|
||||
def _oracle_upgrade() -> None:
|
||||
_oracle_execute_ignoring_955(
|
||||
"""
|
||||
CREATE TABLE graph_maintenance_queue (
|
||||
bank_id VARCHAR2(256) NOT NULL,
|
||||
unit_id RAW(16) NOT NULL,
|
||||
enqueued_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,
|
||||
CONSTRAINT pk_graph_maintenance_queue PRIMARY KEY (bank_id, unit_id)
|
||||
)
|
||||
"""
|
||||
)
|
||||
_oracle_execute_ignoring_955(
|
||||
"CREATE INDEX idx_graph_maintenance_queue_bank_enqueued ON graph_maintenance_queue (bank_id, enqueued_at)"
|
||||
)
|
||||
|
||||
|
||||
def _oracle_downgrade() -> None:
|
||||
op.execute("DROP INDEX idx_graph_maintenance_queue_bank_enqueued")
|
||||
op.execute("DROP TABLE graph_maintenance_queue")
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade, oracle=_oracle_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade, oracle=_oracle_downgrade)
|
||||
-152
@@ -1,152 +0,0 @@
|
||||
"""Re-create vchord vector indexes with vector_cosine_ops
|
||||
|
||||
Revision ID: b8c9d0e1f2a3
|
||||
Revises: 86f7a033d372
|
||||
Create Date: 2026-05-20
|
||||
|
||||
vchordrq operator classes are bound 1:1 to operators in PostgreSQL:
|
||||
vector_l2_ops only matches ``<->``, while every Hindsight ANN query uses
|
||||
``<=>`` (cosine distance). The previous vchord mapping used vector_l2_ops,
|
||||
so vchord deployments could never use the index — every ANN query fell
|
||||
back to a sequential scan with per-row cosine computation.
|
||||
|
||||
This migration finds any vchordrq index built with vector_l2_ops in the
|
||||
target schema and re-creates it with vector_cosine_ops, using
|
||||
``CREATE INDEX CONCURRENTLY`` so it can run online. It is a no-op when:
|
||||
|
||||
* the configured vector extension is not vchord, or
|
||||
* no matching indexes exist (already on cosine ops).
|
||||
|
||||
Only PostgreSQL is affected; the Oracle 23ai dialect uses its own native
|
||||
vector index and does not depend on this mapping.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
from sqlalchemy import text
|
||||
|
||||
from hindsight_api._vector_index import configured_vector_extension
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "b8c9d0e1f2a3"
|
||||
down_revision: str | Sequence[str] | None = "86f7a033d372"
|
||||
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 _rebuild_vchordrq_indexes(old_ops: str, new_ops: str) -> None:
|
||||
"""Rebuild vchordrq indexes using ``old_ops`` so they use ``new_ops``.
|
||||
|
||||
Each index is rebuilt with CREATE INDEX CONCURRENTLY under a fresh name,
|
||||
then the old index is dropped and the new one renamed to take its place.
|
||||
Must be called inside an ``autocommit_block()`` because CONCURRENTLY
|
||||
cannot run inside a transaction.
|
||||
"""
|
||||
bind = op.get_bind()
|
||||
# `or None` collapses both unset and explicit empty-string Alembic options
|
||||
# into NULL so the COALESCE below falls back to current_schema() in either
|
||||
# case. Without it, an empty-string option would filter on `schemaname = ''`
|
||||
# and skip every real schema.
|
||||
target_schema = context.config.get_main_option("target_schema") or None
|
||||
prefix = _pg_schema_prefix()
|
||||
|
||||
rows = bind.execute(
|
||||
text(
|
||||
"SELECT indexname, indexdef FROM pg_indexes "
|
||||
"WHERE schemaname = COALESCE(:target_schema, current_schema()) "
|
||||
"AND indexdef ILIKE '%vchordrq%' "
|
||||
"AND indexdef ILIKE :ops_like"
|
||||
),
|
||||
{"target_schema": target_schema, "ops_like": f"%{old_ops}%"},
|
||||
).fetchall()
|
||||
|
||||
for idx_name, indexdef in rows:
|
||||
# pg_get_indexdef() emits the canonical form `CREATE INDEX <name> ON …`,
|
||||
# so <name> is the first textual occurrence — both substitutions below
|
||||
# rely on that.
|
||||
new_def = indexdef.replace(old_ops, new_ops, 1)
|
||||
temp_name = f"{idx_name}__opclass_swap"
|
||||
new_def = new_def.replace(idx_name, temp_name, 1)
|
||||
new_def = re.sub(
|
||||
r"^CREATE\s+INDEX\b",
|
||||
"CREATE INDEX CONCURRENTLY IF NOT EXISTS",
|
||||
new_def,
|
||||
count=1,
|
||||
)
|
||||
|
||||
# CREATE INDEX CONCURRENTLY can leave the partial index as INVALID if a
|
||||
# previous run errored (disk pressure, lock conflict, signal). Without
|
||||
# this drop the CONCURRENTLY IF NOT EXISTS below would skip creation,
|
||||
# then we'd drop the original and rename the broken index into its
|
||||
# place — silently restoring the seq-scan bug this migration fixes.
|
||||
op.execute(f'DROP INDEX IF EXISTS {prefix}"{temp_name}"')
|
||||
op.execute(new_def)
|
||||
|
||||
# Even on a clean run CONCURRENTLY can finish with indisvalid = false
|
||||
# (e.g. constraint violation during the second build scan). Refuse to
|
||||
# promote in that case so we never alias an INVALID index over a working
|
||||
# one.
|
||||
is_valid = bind.execute(
|
||||
text(
|
||||
"SELECT 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 = :name "
|
||||
" AND n.nspname = COALESCE(:target_schema, current_schema())"
|
||||
),
|
||||
{"name": temp_name, "target_schema": target_schema},
|
||||
).scalar()
|
||||
if not is_valid:
|
||||
raise RuntimeError(
|
||||
f"vchordrq index rebuild produced an INVALID index ({temp_name}); "
|
||||
"drop it manually and re-run the migration."
|
||||
)
|
||||
|
||||
# DROP + RENAME atomically. A crash between the two would leave
|
||||
# `temp_name` as a valid orphan and the canonical name missing —
|
||||
# next run's `pg_indexes` filter (looking for vector_l2_ops) wouldn't
|
||||
# find anything to recover from, so the index would stay gone. PG
|
||||
# runs the DO block in its own server-side transaction, so either
|
||||
# both succeed or both roll back.
|
||||
op.execute(
|
||||
f"""
|
||||
DO $$
|
||||
BEGIN
|
||||
DROP INDEX IF EXISTS {prefix}"{idx_name}";
|
||||
ALTER INDEX {prefix}"{temp_name}" RENAME TO "{idx_name}";
|
||||
END $$;
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
if configured_vector_extension() != "vchord":
|
||||
return
|
||||
with op.get_context().autocommit_block():
|
||||
_rebuild_vchordrq_indexes("vector_l2_ops", "vector_cosine_ops")
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
if configured_vector_extension() != "vchord":
|
||||
return
|
||||
with op.get_context().autocommit_block():
|
||||
_rebuild_vchordrq_indexes("vector_cosine_ops", "vector_l2_ops")
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade)
|
||||
+7
-8
@@ -47,18 +47,17 @@ def _pg_upgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
# GIN index on canonical_name enables sub-millisecond trigram similarity queries
|
||||
# (% operator, similarity()) instead of full-table scans across all bank entities.
|
||||
# CREATE INDEX CONCURRENTLY cannot run inside a transaction block.
|
||||
with op.get_context().autocommit_block():
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS entities_canonical_name_trgm_idx "
|
||||
f"ON {schema}entities USING GIN (canonical_name gin_trgm_ops)"
|
||||
)
|
||||
op.execute("COMMIT")
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS entities_canonical_name_trgm_idx "
|
||||
f"ON {schema}entities USING GIN (canonical_name gin_trgm_ops)"
|
||||
)
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
with op.get_context().autocommit_block():
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}entities_canonical_name_trgm_idx")
|
||||
op.execute("COMMIT")
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}entities_canonical_name_trgm_idx")
|
||||
# Note: not dropping pg_trgm extension as other indexes may depend on it
|
||||
|
||||
|
||||
|
||||
-45
@@ -1,45 +0,0 @@
|
||||
"""Merge graph_maintenance_queue and vchord_cosine_opclass heads.
|
||||
|
||||
Revision ID: c1d2e3f4a5b6
|
||||
Revises: b5a4c3e2f1d8, b8c9d0e1f2a3
|
||||
Create Date: 2026-05-29
|
||||
|
||||
PRs #1668 (vchord cosine opclass) and #1772 (async link recompute) both
|
||||
branched off the same parent and were merged onto main without rebasing,
|
||||
leaving two parallel Alembic heads. This is a structural merge revision
|
||||
with no schema changes — its only job is to unify the DAG so
|
||||
``alembic upgrade head`` is unambiguous again.
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "c1d2e3f4a5b6"
|
||||
down_revision: str | Sequence[str] | None = ("b5a4c3e2f1d8", "b8c9d0e1f2a3")
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
pass
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
pass
|
||||
|
||||
|
||||
def _oracle_upgrade() -> None:
|
||||
pass
|
||||
|
||||
|
||||
def _oracle_downgrade() -> None:
|
||||
pass
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade, oracle=_oracle_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade, oracle=_oracle_downgrade)
|
||||
-71
@@ -1,71 +0,0 @@
|
||||
"""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)
|
||||
-75
@@ -1,75 +0,0 @@
|
||||
"""Repair: widen ``*_history.bank_id`` from VARCHAR(64) to TEXT on PostgreSQL.
|
||||
|
||||
The original split-history migration (``a7b8c9d0e1f2``) declared
|
||||
``observation_history.bank_id`` and ``mental_model_history.bank_id`` as
|
||||
``VARCHAR(64)`` on PostgreSQL. But ``memory_units.bank_id`` — the backfill
|
||||
source for observations — is ``TEXT`` (unbounded), as are ``banks``,
|
||||
``documents`` and ``entities``. Any deployment whose ``bank_id`` exceeds 64
|
||||
characters aborts the backfill ``INSERT`` with::
|
||||
|
||||
psycopg2.errors.StringDataRightTruncation: value too long for type
|
||||
character varying(64)
|
||||
|
||||
Because the migration runs in ``lifespan`` startup inside a transaction, the
|
||||
whole migration rolls back and the API never comes up — unrecoverable from the
|
||||
running container. See https://github.com/vectorize-io/hindsight/issues/2106.
|
||||
|
||||
``a7b8c9d0e1f2`` itself has been corrected to create the column as ``TEXT``,
|
||||
which unblocks deployments that *failed* (the migration rolled back, so it
|
||||
re-runs the fixed DDL). This forward migration covers deployments that already
|
||||
*succeeded* with the narrow ``VARCHAR(64)`` column — where editing
|
||||
``a7b8c9d0e1f2`` has no effect because it will not re-run — by widening the
|
||||
column in place. ``ALTER COLUMN ... TYPE TEXT`` is a no-op on a column that is
|
||||
already ``TEXT`` (fresh installs and re-run failures), so every upgrade path
|
||||
converges on ``TEXT``.
|
||||
|
||||
The history tables are per-tenant (they live in each tenant schema, not
|
||||
``public``), so this runs for every migrated schema via the search-path-aware
|
||||
prefix — unlike the shared-``public`` routines repaired in ``b2d4f6a8c1e3``.
|
||||
|
||||
PostgreSQL only. On Oracle both ``memory_units.bank_id`` and the history
|
||||
``bank_id`` columns are already ``VARCHAR2(256)`` (consistent, never
|
||||
truncates), so the Oracle slot is intentionally absent.
|
||||
|
||||
Revision ID: c3e5a7b9d1f4
|
||||
Revises: c9a1b2d3e4f5
|
||||
Create Date: 2026-06-10
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "c3e5a7b9d1f4"
|
||||
down_revision: str | Sequence[str] | None = "c9a1b2d3e4f5"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _get_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 = _get_schema_prefix()
|
||||
op.execute(f"ALTER TABLE {schema}observation_history ALTER COLUMN bank_id TYPE TEXT")
|
||||
op.execute(f"ALTER TABLE {schema}mental_model_history ALTER COLUMN bank_id TYPE TEXT")
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
# No-op: narrowing back to VARCHAR(64) could truncate real data and would
|
||||
# re-introduce the bug this migration repairs. The column type is owned by
|
||||
# ``a7b8c9d0e1f2``'s lifecycle.
|
||||
pass
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade)
|
||||
-144
@@ -1,144 +0,0 @@
|
||||
"""Backfill search_vector for native-backend observations.
|
||||
|
||||
Observations created or updated by the consolidator landed with a NULL
|
||||
``search_vector`` under the ``native`` text-search backend: the
|
||||
single-row INSERT/UPDATE paths in ``consolidator.py`` never populated the
|
||||
tsvector (only the batch raw-fact path in ``ops_postgresql.insert_facts_batch``
|
||||
did). Those observations were therefore invisible to the BM25 retrieval arm
|
||||
until they were re-written by a later consolidation pass. The writer is fixed
|
||||
in the same change set (all four consolidator sites now call
|
||||
``to_tsvector($lang, COALESCE(text, ''))``); this migration repairs the
|
||||
historical residue so existing observations become BM25-searchable without a
|
||||
re-ingest.
|
||||
|
||||
Scope mirrors the writer fix exactly:
|
||||
* Only the ``native`` backend is touched. The gate is the column *type*:
|
||||
under ``native`` ``search_vector`` is a regular (non-generated) tsvector
|
||||
column; under ``vchord`` it is a ``bm25vector`` and under
|
||||
``pg_textsearch`` / ``pgroonga`` / ``pg_search`` it is a dummy ``text``
|
||||
column. ``_is_regular_tsvector`` is true only for ``native``, so every
|
||||
other backend is a no-op.
|
||||
* The tsvector is built from the observation's own ``text`` only — matching
|
||||
the consolidator INSERT/UPDATE paths (entity / source / temporal signals
|
||||
are intentionally excluded; the other retrieval arms cover those).
|
||||
* Only ``fact_type = 'observation'`` rows with a NULL ``search_vector`` are
|
||||
rewritten. Raw facts already carry a populated tsvector, and the
|
||||
``IS NULL`` predicate makes the migration idempotent and re-runnable.
|
||||
|
||||
The configured ``HINDSIGHT_API_TEXT_SEARCH_EXTENSION_NATIVE_LANGUAGE`` is used
|
||||
so backfilled rows are lexically identical to newly-created observations. The
|
||||
value is validated as a PG identifier (mirroring
|
||||
``HindsightConfig.validate``) before being embedded as a SQL literal.
|
||||
|
||||
This is a single UPDATE per schema: it locks the targeted observation rows for
|
||||
its duration. It is one-time and only touches unpopulated rows, so subsequent
|
||||
online writes (which now carry the tsvector via the writer fix) are unaffected.
|
||||
|
||||
Oracle slot is intentionally absent: the consolidator INSERT/UPDATE paths that
|
||||
this repairs are PostgreSQL-specific (``ops_postgresql``), and the native
|
||||
tsvector ``search_vector`` column only exists on PostgreSQL. There is no Oracle
|
||||
residue to repair.
|
||||
|
||||
Revision ID: c3f7a1b9d2e4
|
||||
Revises: f4d1c2b3a5e6
|
||||
Create Date: 2026-06-29
|
||||
"""
|
||||
|
||||
import os
|
||||
import re
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
from sqlalchemy import Connection, text
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
from hindsight_api.config import (
|
||||
DEFAULT_TEXT_SEARCH_EXTENSION_NATIVE_LANGUAGE,
|
||||
ENV_TEXT_SEARCH_EXTENSION_NATIVE_LANGUAGE,
|
||||
)
|
||||
|
||||
revision: str = "c3f7a1b9d2e4"
|
||||
down_revision: str | Sequence[str] | None = "f4d1c2b3a5e6"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
# Matches HindsightConfig.validate(): a tsvector regconfig name embedded as a
|
||||
# SQL literal must be a bare PG identifier.
|
||||
_PG_IDENTIFIER = re.compile(r"[a-zA-Z_][a-zA-Z0-9_]*")
|
||||
|
||||
|
||||
def _schema_prefix() -> str:
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def _schema_name() -> str:
|
||||
return (context.config.get_main_option("target_schema") or "public").strip('"')
|
||||
|
||||
|
||||
def _native_language() -> str:
|
||||
"""Configured native tsvector language, validated as a PG identifier."""
|
||||
lang = os.getenv(
|
||||
ENV_TEXT_SEARCH_EXTENSION_NATIVE_LANGUAGE,
|
||||
DEFAULT_TEXT_SEARCH_EXTENSION_NATIVE_LANGUAGE,
|
||||
)
|
||||
if not _PG_IDENTIFIER.fullmatch(lang):
|
||||
return DEFAULT_TEXT_SEARCH_EXTENSION_NATIVE_LANGUAGE
|
||||
return lang
|
||||
|
||||
|
||||
def _is_regular_tsvector(conn: Connection, schema: str, table: str) -> bool:
|
||||
"""True iff ``schema.table.search_vector`` is a non-generated tsvector column.
|
||||
|
||||
This is the ``native`` backend signature. ``vchord`` (bm25vector) and
|
||||
``pg_textsearch`` / ``pgroonga`` / ``pg_search`` (dummy text column) all
|
||||
fail this check, so the backfill is a no-op for them.
|
||||
"""
|
||||
row = conn.execute(
|
||||
text(
|
||||
"""
|
||||
SELECT is_generated, udt_name
|
||||
FROM information_schema.columns
|
||||
WHERE table_schema = :schema
|
||||
AND table_name = :table
|
||||
AND column_name = 'search_vector'
|
||||
"""
|
||||
),
|
||||
{"schema": schema, "table": table},
|
||||
).fetchone()
|
||||
if not row:
|
||||
return False
|
||||
is_generated, udt_name = row[0], row[1]
|
||||
return udt_name == "tsvector" and is_generated != "ALWAYS"
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
conn = op.get_bind()
|
||||
schema_name = _schema_name()
|
||||
if not _is_regular_tsvector(conn, schema_name, "memory_units"):
|
||||
# Non-native backend (or column absent) — nothing to backfill.
|
||||
return
|
||||
schema_prefix = _schema_prefix()
|
||||
lang = _native_language()
|
||||
op.execute(
|
||||
f"""
|
||||
UPDATE {schema_prefix}memory_units
|
||||
SET search_vector = to_tsvector('{lang}'::regconfig, COALESCE(text, ''))
|
||||
WHERE fact_type = 'observation' AND search_vector IS NULL
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
# No-op: backfilled rows are indistinguishable from observations that were
|
||||
# populated by the post-fix writer, and reverting either to NULL would
|
||||
# re-break BM25 retrieval. The column simply stays populated.
|
||||
pass
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade)
|
||||
-158
@@ -1,158 +0,0 @@
|
||||
"""Make maintenance routines resilient to schemas that vanish mid-scan.
|
||||
|
||||
``public.banks_needing_consolidation()`` and
|
||||
``public.schemas_with_expired_rows(...)`` snapshot the set of schemas owning a
|
||||
target table from ``pg_class`` and then run a dynamic query against each schema
|
||||
in turn. That is a time-of-check/time-of-use race: a schema (or its tables) can
|
||||
be dropped — a tenant being deleted, or a tenant migration that recreates
|
||||
tables — between the snapshot and the per-schema query, which then aborts the
|
||||
whole routine with::
|
||||
|
||||
relation "<schema>.memory_units" does not exist
|
||||
relation "<schema>.audit_log" does not exist
|
||||
|
||||
In the test suite this surfaces as cross-worker contamination: the multi-tenant
|
||||
maintenance test creates and drops ~100 ``mt<hash>_NNN`` schemas while
|
||||
``test_maintenance_routines`` (on another xdist worker, same DB) calls the
|
||||
routines. In production the background maintenance loop hits the same race when
|
||||
a tenant is removed or mid-migration.
|
||||
|
||||
Wrap each per-schema query in its own ``BEGIN ... EXCEPTION`` block so a schema
|
||||
that disappears (``undefined_table`` / ``invalid_schema_name`` /
|
||||
``undefined_column``) is skipped instead of aborting the scan. The routines stay
|
||||
``CREATE OR REPLACE`` and PostgreSQL-only, and are (re)installed only on the run
|
||||
that targets the shared ``public`` schema — same gating as the original
|
||||
install (``e5f6a7b8c9d0``) and its repair (``b2d4f6a8c1e3``).
|
||||
|
||||
Revision ID: c7e9f1a3b5d2
|
||||
Revises: e1f2a3b4c5d6
|
||||
Create Date: 2026-06-19
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "c7e9f1a3b5d2"
|
||||
down_revision: str | Sequence[str] | None = "e1f2a3b4c5d6"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _should_install_public_routines(target_schema: str | None) -> bool:
|
||||
"""True for the run that must (re)create the shared ``public.*`` routines.
|
||||
|
||||
The routines physically live in ``public``, so they are installed exactly
|
||||
once — on the base run (no ``target_schema``) or the run that explicitly
|
||||
targets ``public``. Mirrors ``b2d4f6a8c1e3``.
|
||||
"""
|
||||
return not target_schema or target_schema == "public"
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
if not _should_install_public_routines(context.config.get_main_option("target_schema")):
|
||||
return
|
||||
|
||||
# Same body as b2d4f6a8c1e3, but each per-schema query runs in its own
|
||||
# subtransaction so a schema dropped mid-scan is skipped, not fatal.
|
||||
op.execute(
|
||||
"""
|
||||
CREATE OR REPLACE FUNCTION public.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(
|
||||
"""
|
||||
CREATE OR REPLACE FUNCTION public.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$;
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
# No-op: e5f6a7b8c9d0 owns these functions' lifecycle and drops them on its
|
||||
# own downgrade. This migration only re-installs them (the resilient body is
|
||||
# a strict superset of the previous behaviour), so there is nothing to undo
|
||||
# without racing that migration's DROP.
|
||||
pass
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade)
|
||||
-108
@@ -1,108 +0,0 @@
|
||||
"""Add invalidated_memory_units table for curation (edit/invalidate).
|
||||
|
||||
Curation keeps the recall hot-path (``memory_units``) clean by *moving*
|
||||
invalidated facts into a sibling archive table rather than flagging them in
|
||||
place. If a row is in ``memory_units`` it is live; if it is in
|
||||
``invalidated_memory_units`` it has been retired. Recall/consolidation/graph
|
||||
queries never need a state predicate — the rows simply aren't there.
|
||||
|
||||
The archive mirrors ``memory_units`` column-for-column — except ``embedding``,
|
||||
which it never keeps: the archive is cold storage, never a recall surface, and
|
||||
revert recomputes the embedding from the unit's text/dates/entities. Keeping no
|
||||
archive vector also means a later embedding-model switch (which re-dimensions
|
||||
``memory_units``) can't trip a dimension mismatch on the move (#2209). Plus:
|
||||
- ``invalidation_reason`` optional free text recorded on invalidate
|
||||
- ``invalidated_at`` when it was retired
|
||||
- ``entity_ids`` snapshot of the unit's entity associations, so revert
|
||||
can restore them (``unit_entities`` is cascade-deleted
|
||||
when the live row is removed)
|
||||
|
||||
This migration also adds ``edited_at`` to ``memory_units``: set whenever a user
|
||||
edits a memory's fields (text, context, dates, fact_type, entities) via curation.
|
||||
NULL means never manually modified; a non-NULL value answers "has the user ever
|
||||
changed this?" with the time of the last edit (distinct from ``updated_at``,
|
||||
which background operations also bump). It is added to ``memory_units`` *before*
|
||||
the archive is cloned below, so the archive inherits the column and the marker
|
||||
travels with a fact when it is invalidated.
|
||||
|
||||
Revision ID: c9a1b2d3e4f5
|
||||
Revises: b2d4f6a8c1e3
|
||||
Create Date: 2026-06-03
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "c9a1b2d3e4f5"
|
||||
down_revision: str | Sequence[str] | None = "b2d4f6a8c1e3"
|
||||
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()
|
||||
# Add edited_at to the live table FIRST so the archive's LIKE clone below
|
||||
# inherits it (keeps the two tables column-for-column identical for round-trip).
|
||||
op.execute(f"ALTER TABLE {schema}memory_units ADD COLUMN IF NOT EXISTS edited_at TIMESTAMPTZ")
|
||||
# LIKE ... INCLUDING DEFAULTS clones every memory_units column (incl.
|
||||
# edited_at) so an invalidated row can move back verbatim. We deliberately
|
||||
# omit indexes/constraints — the archive is cold storage, not a recall
|
||||
# surface; only the lookups below need indexing.
|
||||
op.execute(
|
||||
f"CREATE TABLE IF NOT EXISTS {schema}invalidated_memory_units (LIKE {schema}memory_units INCLUDING DEFAULTS)"
|
||||
)
|
||||
# ...then drop the inherited embedding: the archive never stores one (revert
|
||||
# recomputes it), so it isn't created here only to be dropped again later by
|
||||
# d4f6a8c2e1b3. That migration still runs as a no-op (DROP ... IF EXISTS) on
|
||||
# fresh DBs and does the real drop on DBs created before this column was removed.
|
||||
op.execute(f"ALTER TABLE {schema}invalidated_memory_units DROP COLUMN IF EXISTS embedding")
|
||||
op.execute(
|
||||
f"ALTER TABLE {schema}invalidated_memory_units "
|
||||
f"ADD COLUMN IF NOT EXISTS invalidation_reason TEXT, "
|
||||
f"ADD COLUMN IF NOT EXISTS invalidated_at TIMESTAMPTZ DEFAULT now(), "
|
||||
f"ADD COLUMN IF NOT EXISTS entity_ids UUID[]"
|
||||
)
|
||||
op.execute(f"CREATE UNIQUE INDEX IF NOT EXISTS idx_invalidated_mu_id ON {schema}invalidated_memory_units (id)")
|
||||
op.execute(
|
||||
f"CREATE INDEX IF NOT EXISTS idx_invalidated_mu_bank "
|
||||
f"ON {schema}invalidated_memory_units (bank_id, invalidated_at)"
|
||||
)
|
||||
# Deleting a document (or bank) should clear its archived facts too, mirroring
|
||||
# the memory_units → documents cascade.
|
||||
op.execute(
|
||||
f"""
|
||||
DO $$ BEGIN
|
||||
IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'invalidated_mu_document_fkey') THEN
|
||||
ALTER TABLE {schema}invalidated_memory_units
|
||||
ADD CONSTRAINT invalidated_mu_document_fkey
|
||||
FOREIGN KEY (document_id, bank_id)
|
||||
REFERENCES {schema}documents(id, bank_id) ON DELETE CASCADE;
|
||||
END IF; END $$;
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
schema = _pg_schema_prefix()
|
||||
# Drops the archive (and its inherited edited_at) wholesale, then removes
|
||||
# edited_at from the live table.
|
||||
op.execute(f"DROP TABLE IF EXISTS {schema}invalidated_memory_units")
|
||||
op.execute(f"ALTER TABLE {schema}memory_units DROP COLUMN IF EXISTS edited_at")
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
# PG-only: Oracle gets the table from the baseline snapshot, matching the
|
||||
# convention used by sibling column/index migrations in this tree.
|
||||
run_for_dialect(pg=_pg_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade)
|
||||
+28
-24
@@ -50,35 +50,39 @@ def _get_schema_prefix() -> str:
|
||||
def _pg_upgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
|
||||
# CREATE INDEX CONCURRENTLY cannot run inside a transaction block; an
|
||||
# autocommit_block runs each statement outside Alembic's migration
|
||||
# transaction. IF NOT EXISTS makes each statement idempotent on retry.
|
||||
with op.get_context().autocommit_block():
|
||||
# Index for the semantic *incoming* direction in link_expansion_retrieval.py.
|
||||
# Replaces the BitmapAnd of idx_memory_links_to_unit ∩ idx_memory_links_link_type
|
||||
# with a single composite index scan.
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_links_to_type_weight "
|
||||
f"ON {schema}memory_links(to_unit_id, link_type, weight DESC)"
|
||||
)
|
||||
# CREATE INDEX CONCURRENTLY cannot run inside a transaction block.
|
||||
# Commit the current Alembic transaction, then issue each CONCURRENTLY
|
||||
# statement in its own implicit autocommit transaction.
|
||||
# IF NOT EXISTS makes each statement idempotent if the migration is retried.
|
||||
|
||||
# Covering index for entity co-occurrence expansion.
|
||||
# Enables an index-only scan: entity_id and to_unit_id are read from the
|
||||
# index leaf pages instead of the heap, eliminating ~2 500 random heap-page
|
||||
# reads per expansion query.
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_links_entity_covering "
|
||||
f"ON {schema}memory_links(from_unit_id) "
|
||||
f"INCLUDE (to_unit_id, entity_id) "
|
||||
f"WHERE link_type = 'entity'"
|
||||
)
|
||||
# Index for the semantic *incoming* direction in link_expansion_retrieval.py.
|
||||
# Replaces the BitmapAnd of idx_memory_links_to_unit ∩ idx_memory_links_link_type
|
||||
# with a single composite index scan.
|
||||
op.execute("COMMIT")
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_links_to_type_weight "
|
||||
f"ON {schema}memory_links(to_unit_id, link_type, weight DESC)"
|
||||
)
|
||||
|
||||
# Covering index for entity co-occurrence expansion.
|
||||
# Enables an index-only scan: entity_id and to_unit_id are read from the
|
||||
# index leaf pages instead of the heap, eliminating ~2 500 random heap-page
|
||||
# reads per expansion query.
|
||||
op.execute("COMMIT")
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_links_entity_covering "
|
||||
f"ON {schema}memory_links(from_unit_id) "
|
||||
f"INCLUDE (to_unit_id, entity_id) "
|
||||
f"WHERE link_type = 'entity'"
|
||||
)
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
with op.get_context().autocommit_block():
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_links_entity_covering")
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_links_to_type_weight")
|
||||
op.execute("COMMIT")
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_links_entity_covering")
|
||||
op.execute("COMMIT")
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_links_to_type_weight")
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
|
||||
-96
@@ -1,96 +0,0 @@
|
||||
"""Add llm_requests table for per-bank LLM request tracing.
|
||||
|
||||
Stores one row per logical LLM call Hindsight makes (success and failure),
|
||||
capturing the input messages, model output, token usage (input/output/cached/
|
||||
total), finish reason, and caller metadata. Disabled by default at the
|
||||
application layer (HINDSIGHT_API_LLM_TRACE_ENABLED); this migration only
|
||||
creates the table.
|
||||
|
||||
PostgreSQL only — the tracing subsystem is not wired for Oracle, so the Oracle
|
||||
slot is intentionally absent (mirrors the audit_log table).
|
||||
|
||||
Revision ID: d3e4f5a6b7c8
|
||||
Revises: c1d2e3f4a5b6
|
||||
Create Date: 2026-06-01
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "d3e4f5a6b7c8"
|
||||
down_revision: str | Sequence[str] | None = "c1d2e3f4a5b6"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _get_schema_prefix() -> str:
|
||||
"""Get schema prefix for table names (required for multi-tenant support)."""
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
|
||||
op.execute(
|
||||
f"""
|
||||
CREATE TABLE IF NOT EXISTS {schema}llm_requests (
|
||||
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
||||
bank_id TEXT,
|
||||
operation TEXT,
|
||||
scope TEXT,
|
||||
-- OTel-style grouping: trace_id is shared by every LLM call of one
|
||||
-- operation invocation (e.g. all calls of a single reflect run);
|
||||
-- parent_span_id is that operation span; span_id is this call.
|
||||
trace_id TEXT,
|
||||
span_id TEXT,
|
||||
parent_span_id TEXT,
|
||||
provider TEXT,
|
||||
model TEXT,
|
||||
status TEXT NOT NULL,
|
||||
started_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
|
||||
ended_at TIMESTAMPTZ,
|
||||
duration_ms INTEGER,
|
||||
input_tokens INTEGER,
|
||||
output_tokens INTEGER,
|
||||
cached_tokens INTEGER,
|
||||
total_tokens INTEGER,
|
||||
input JSONB,
|
||||
output JSONB,
|
||||
error TEXT,
|
||||
llm_info JSONB DEFAULT '{{}}'::jsonb,
|
||||
metadata JSONB DEFAULT '{{}}'::jsonb
|
||||
)
|
||||
"""
|
||||
)
|
||||
|
||||
op.execute(
|
||||
f"CREATE INDEX IF NOT EXISTS idx_llm_requests_bank_started ON {schema}llm_requests (bank_id, started_at DESC)"
|
||||
)
|
||||
op.execute(
|
||||
f"CREATE INDEX IF NOT EXISTS idx_llm_requests_status_started ON {schema}llm_requests (status, started_at DESC)"
|
||||
)
|
||||
op.execute(f"CREATE INDEX IF NOT EXISTS idx_llm_requests_started ON {schema}llm_requests (started_at DESC)")
|
||||
op.execute(
|
||||
f"CREATE INDEX IF NOT EXISTS idx_llm_requests_trace ON {schema}llm_requests (bank_id, trace_id, started_at)"
|
||||
)
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema}idx_llm_requests_started")
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema}idx_llm_requests_status_started")
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema}idx_llm_requests_bank_started")
|
||||
op.execute(f"DROP TABLE IF EXISTS {schema}llm_requests")
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade)
|
||||
+16
-17
@@ -33,27 +33,26 @@ def _get_schema_prefix() -> str:
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
# DROP + CREATE CONCURRENTLY must run outside a transaction block; an
|
||||
# autocommit_block runs them outside Alembic's migration transaction.
|
||||
with op.get_context().autocommit_block():
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_units_source_memory_ids")
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_units_source_memory_ids "
|
||||
f"ON {schema}memory_units USING GIN (source_memory_ids) "
|
||||
f"WITH (fastupdate=off) "
|
||||
f"WHERE source_memory_ids IS NOT NULL"
|
||||
)
|
||||
# DROP + CREATE CONCURRENTLY must run outside a transaction block.
|
||||
op.execute("COMMIT")
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_units_source_memory_ids")
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_units_source_memory_ids "
|
||||
f"ON {schema}memory_units USING GIN (source_memory_ids) "
|
||||
f"WITH (fastupdate=off) "
|
||||
f"WHERE source_memory_ids IS NOT NULL"
|
||||
)
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
with op.get_context().autocommit_block():
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_units_source_memory_ids")
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_units_source_memory_ids "
|
||||
f"ON {schema}memory_units USING GIN (source_memory_ids) "
|
||||
f"WHERE source_memory_ids IS NOT NULL"
|
||||
)
|
||||
op.execute("COMMIT")
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_units_source_memory_ids")
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_units_source_memory_ids "
|
||||
f"ON {schema}memory_units USING GIN (source_memory_ids) "
|
||||
f"WHERE source_memory_ids IS NOT NULL"
|
||||
)
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
|
||||
-93
@@ -1,93 +0,0 @@
|
||||
"""Drop the embedding column from the curation archive (invalidated_memory_units).
|
||||
|
||||
The archive is cold storage, never a recall surface, so it has no business
|
||||
keeping an embedding. Earlier curation code copied the live row's embedding 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 makes "the archive holds no embedding" a schema-enforced invariant
|
||||
rather than a convention the move queries have to honour, and removes a latent
|
||||
failure mode (#2209): after an embedding-model switch the live tables are
|
||||
re-dimensioned but the archive was not, so a stale old-dimension embedding in
|
||||
the archive tripped a vector-dimension mismatch on the INSERT … SELECT
|
||||
round-trip. With no column at all, there is nothing to mismatch.
|
||||
|
||||
The creation sites no longer add the column (the PG ``LIKE`` clone in
|
||||
c9a1b2d3e4f5 drops it; the Oracle baseline omits it), so on a fresh database
|
||||
this migration is a no-op (DROP ... IF EXISTS / Oracle ORA-00904 swallow). It
|
||||
does the real work on databases created before the column was removed there.
|
||||
|
||||
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 unconstrained vector column (any dimension) — empty, since the
|
||||
embeddings are intentionally discarded.
|
||||
|
||||
Revision ID: d4f6a8c2e1b3
|
||||
Revises: a1d3f5b7c9e2
|
||||
Create Date: 2026-06-15
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "d4f6a8c2e1b3"
|
||||
down_revision: str | Sequence[str] | None = "a1d3f5b7c9e2"
|
||||
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 embedding")
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
schema = _pg_schema_prefix()
|
||||
# Unconstrained `vector` (no dimension) so the re-added column accepts any
|
||||
# model's embeddings; it comes back empty regardless.
|
||||
op.execute(f"ALTER TABLE {schema}invalidated_memory_units ADD COLUMN IF NOT EXISTS embedding vector")
|
||||
|
||||
|
||||
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 fresh schema whose
|
||||
# baseline already omits the column.
|
||||
op.execute(
|
||||
"""
|
||||
BEGIN
|
||||
EXECUTE IMMEDIATE 'ALTER TABLE invalidated_memory_units DROP COLUMN embedding';
|
||||
EXCEPTION WHEN OTHERS THEN
|
||||
IF SQLCODE != -904 THEN RAISE; END IF;
|
||||
END;
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def _oracle_downgrade() -> None:
|
||||
# Swallow ORA-01430 (column already exists) for idempotency.
|
||||
op.execute(
|
||||
"""
|
||||
BEGIN
|
||||
EXECUTE IMMEDIATE 'ALTER TABLE invalidated_memory_units ADD (embedding VECTOR)';
|
||||
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
-1
@@ -55,7 +55,7 @@ def _vector_index_using_clause(ext: str) -> str:
|
||||
if ext == "pgvectorscale":
|
||||
return "USING diskann (embedding vector_cosine_ops) WITH (num_neighbors = 50)"
|
||||
if ext == "vchord":
|
||||
return "USING vchordrq (embedding vector_cosine_ops)"
|
||||
return "USING vchordrq (embedding vector_l2_ops)"
|
||||
if ext == "scann":
|
||||
return "USING scann (embedding cosine) WITH (mode = 'AUTO')"
|
||||
return "USING hnsw (embedding vector_cosine_ops)"
|
||||
|
||||
-133
@@ -1,133 +0,0 @@
|
||||
"""Drop indexes that are unused or redundant with composite indexes.
|
||||
|
||||
Code audit identified the following indexes as either dead (no code path
|
||||
exercises them) or fully covered by composite indexes the planner already
|
||||
prefers:
|
||||
|
||||
memory_links:
|
||||
1. idx_memory_links_entity_covering — entity co-occurrence expansion was
|
||||
rewritten to traverse unit_entities instead of memory_links, so no code
|
||||
path filters memory_links on (link_type = 'entity').
|
||||
2. idx_memory_links_from_unit — redundant. idx_memory_links_from_type_weight
|
||||
(from_unit_id, link_type, weight DESC) leads with the same column and
|
||||
answers every from_unit_id = X query.
|
||||
3. idx_memory_links_to_unit — redundant. idx_memory_links_to_type_weight
|
||||
(to_unit_id, link_type, weight DESC) leads with the same column.
|
||||
4. idx_memory_links_link_type — no application query filters on link_type
|
||||
alone; the composite indexes above serve every (from/to + link_type)
|
||||
predicate.
|
||||
|
||||
entities:
|
||||
5. idx_entities_canonical_name — superseded by
|
||||
entities_canonical_name_lower_trgm_idx (case-insensitive lookups).
|
||||
6. entities_canonical_name_trgm_idx — superseded by the lowercase variant
|
||||
in migration 2eee35aa3cfc, but the original was never dropped on schemas
|
||||
that ran the prior migration.
|
||||
|
||||
documents:
|
||||
7. idx_documents_retain_params — GIN index on retain_params JSONB; no query
|
||||
uses jsonb containment on this column.
|
||||
8. idx_documents_content_hash — content-hash lookups happen on the chunks
|
||||
table (chunks.content_hash, indexed separately).
|
||||
|
||||
unit_entities:
|
||||
9. idx_unit_entities_entity — defensive drop. Migration h3i4j5k6l7m8 already
|
||||
issues DROP INDEX IF EXISTS for this; this re-runs the drop idempotently
|
||||
to cover any schema that missed the previous migration.
|
||||
|
||||
All drops use CONCURRENTLY + IF EXISTS so they neither block writers nor
|
||||
fail on schemas where the index is already gone.
|
||||
|
||||
Revision ID: e1b2c3d4f5a6
|
||||
Revises: p4q5r6s7t8u9
|
||||
Create Date: 2026-05-26
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "e1b2c3d4f5a6"
|
||||
down_revision: str | Sequence[str] | None = "p4q5r6s7t8u9"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
_PG_INDEXES_TO_DROP: tuple[str, ...] = (
|
||||
"idx_memory_links_entity_covering",
|
||||
"idx_memory_links_from_unit",
|
||||
"idx_memory_links_to_unit",
|
||||
"idx_memory_links_link_type",
|
||||
"idx_entities_canonical_name",
|
||||
"entities_canonical_name_trgm_idx",
|
||||
"idx_documents_retain_params",
|
||||
"idx_documents_content_hash",
|
||||
"idx_unit_entities_entity",
|
||||
)
|
||||
|
||||
|
||||
def _schema_prefix() -> str:
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
schema = _schema_prefix()
|
||||
# DROP INDEX CONCURRENTLY cannot run inside a transaction block; an
|
||||
# autocommit_block drops out of Alembic's migration transaction so each
|
||||
# statement runs in its own autocommit. IF EXISTS makes each statement
|
||||
# idempotent across schemas that already dropped (or never had) the index.
|
||||
with op.get_context().autocommit_block():
|
||||
for index_name in _PG_INDEXES_TO_DROP:
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}{index_name}")
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
schema = _schema_prefix()
|
||||
|
||||
# Recreate the dropped indexes in the same shape the prior migrations used,
|
||||
# so a downgrade leaves the schema in the state the previous head expected.
|
||||
# CREATE INDEX CONCURRENTLY cannot run inside a transaction block.
|
||||
with op.get_context().autocommit_block():
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_links_entity_covering "
|
||||
f"ON {schema}memory_links(from_unit_id) "
|
||||
f"INCLUDE (to_unit_id, entity_id) "
|
||||
f"WHERE link_type = 'entity'"
|
||||
)
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_links_from_unit ON {schema}memory_links(from_unit_id)"
|
||||
)
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_links_to_unit ON {schema}memory_links(to_unit_id)"
|
||||
)
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_links_link_type ON {schema}memory_links(link_type)"
|
||||
)
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_entities_canonical_name ON {schema}entities(canonical_name)"
|
||||
)
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS entities_canonical_name_trgm_idx "
|
||||
f"ON {schema}entities USING GIN (canonical_name gin_trgm_ops)"
|
||||
)
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_documents_retain_params "
|
||||
f"ON {schema}documents USING GIN (retain_params)"
|
||||
)
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_documents_content_hash ON {schema}documents(content_hash)"
|
||||
)
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_unit_entities_entity ON {schema}unit_entities(entity_id)"
|
||||
)
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade)
|
||||
-39
@@ -1,39 +0,0 @@
|
||||
"""Merge two divergent migration heads.
|
||||
|
||||
``d4f6a8c2e1b3`` (drop the curation-archive embedding column) and
|
||||
``2071c7518f88`` (add the memory_links(bank_id, link_type) index) were authored
|
||||
in parallel off the same parent (``a1d3f5b7c9e2``) and merged independently,
|
||||
leaving the DAG with two heads. This is a no-op merge that re-unifies them so
|
||||
``alembic upgrade head`` is unambiguous again (enforced by
|
||||
``tests/test_alembic_dag.py::test_single_head``).
|
||||
|
||||
Revision ID: e1f2a3b4c5d6
|
||||
Revises: d4f6a8c2e1b3, 2071c7518f88
|
||||
Create Date: 2026-06-16
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "e1f2a3b4c5d6"
|
||||
down_revision: str | Sequence[str] | None = ("d4f6a8c2e1b3", "2071c7518f88")
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
# Pure DAG merge — both parents already applied their schema changes.
|
||||
pass
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
pass
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade)
|
||||
-153
@@ -1,153 +0,0 @@
|
||||
"""Add server-side routines for background maintenance sweeps.
|
||||
|
||||
Installs two PL/pgSQL discovery routines in the ``public`` schema. Both loop
|
||||
over every schema that actually holds the relevant table (via ``pg_class``), so
|
||||
a single function call covers all tenants in one round-trip instead of the
|
||||
per-tenant query storm that a client-side loop would create at thousands of
|
||||
tenants.
|
||||
|
||||
- ``public.banks_needing_consolidation()`` -> (schema_name, bank_id) for banks
|
||||
that have eligible-but-unscheduled facts (``consolidated_at IS NULL AND
|
||||
consolidation_failed_at IS NULL`` for consolidatable fact types), have
|
||||
auto-consolidation not explicitly disabled at the bank level, and have no
|
||||
consolidation operation already pending/processing. Drives the periodic
|
||||
reconcile that re-schedules consolidation after a terminal failure left facts
|
||||
stranded (see HINDSIGHT_API_CONSOLIDATION_RECONCILE_INTERVAL_SECONDS).
|
||||
|
||||
- ``public.schemas_with_expired_rows(p_table, p_ts_col, p_days)`` -> schema
|
||||
names that hold at least one ``p_table`` row older than ``p_days``. Drives the
|
||||
cross-tenant retention sweeps for ``audit_log`` and ``llm_requests``; the loop
|
||||
then issues a DELETE only against the returned schemas.
|
||||
|
||||
These are read-only (STABLE) discovery routines — the caller performs the
|
||||
enqueue/DELETE — so installing them never mutates data.
|
||||
|
||||
PostgreSQL only — the worker poller and these tables are not wired for Oracle,
|
||||
so the Oracle slot is intentionally absent (mirrors the audit_log / llm_requests
|
||||
table migrations). The routines live in ``public`` and are CREATE OR REPLACE, so
|
||||
running this migration once per tenant schema is idempotent.
|
||||
|
||||
Revision ID: e5f6a7b8c9d0
|
||||
Revises: a7b8c9d0e1f2
|
||||
Create Date: 2026-06-05
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "e5f6a7b8c9d0"
|
||||
down_revision: str | Sequence[str] | None = "a7b8c9d0e1f2"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _is_base_schema_run() -> bool:
|
||||
"""True only for the base-schema migration (no per-tenant target_schema).
|
||||
|
||||
These routines live in the shared ``public`` schema, so they must be created
|
||||
exactly once. Running ``CREATE OR REPLACE FUNCTION public....`` again from each
|
||||
concurrent per-tenant migration aborts with ``tuple concurrently updated`` on
|
||||
the ``pg_proc`` catalog row, so tenant runs skip it (the base run already
|
||||
created the function for every tenant to use).
|
||||
"""
|
||||
return not context.config.get_main_option("target_schema")
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
if not _is_base_schema_run():
|
||||
return
|
||||
# Banks with eligible-but-unscheduled facts and no in-flight consolidation.
|
||||
# Auto-consolidation is filtered here only at the bank level (cheap prune);
|
||||
# the full hierarchical resolution (global -> tenant -> bank, plus
|
||||
# enable_observations) is done by the caller for the small returned set.
|
||||
op.execute(
|
||||
"""
|
||||
CREATE OR REPLACE FUNCTION public.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
|
||||
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);
|
||||
END LOOP;
|
||||
END;
|
||||
$fn$;
|
||||
"""
|
||||
)
|
||||
|
||||
# Schemas holding at least one row of p_table older than p_days. p_ts_col is
|
||||
# the timestamp column to compare. Returns nothing when p_days <= 0
|
||||
# (retention disabled).
|
||||
op.execute(
|
||||
"""
|
||||
CREATE OR REPLACE FUNCTION public.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
|
||||
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;
|
||||
IF has_expired THEN
|
||||
RETURN NEXT sch;
|
||||
END IF;
|
||||
END LOOP;
|
||||
END;
|
||||
$fn$;
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
if not _is_base_schema_run():
|
||||
return
|
||||
op.execute("DROP FUNCTION IF EXISTS public.banks_needing_consolidation()")
|
||||
op.execute("DROP FUNCTION IF EXISTS public.schemas_with_expired_rows(text, text, int)")
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade)
|
||||
-91
@@ -1,91 +0,0 @@
|
||||
"""Drop materialized entity rows from memory_links.
|
||||
|
||||
Entity edges are no longer stored in ``memory_links``. The /graph endpoint
|
||||
derives them on demand from ``unit_entities``, and recall already used the
|
||||
``unit_entities`` self-join. Storing entity rows duplicated state we never
|
||||
read from the link table — on a 10k-unit benchmark bank, entity rows were
|
||||
53% of all link rows (~190 MB after indexes) and recall never touched them.
|
||||
|
||||
This migration deletes ``memory_links`` rows with ``link_type = 'entity'``.
|
||||
``idx_memory_links_entity_covering`` was already dropped by migration
|
||||
``e1b2c3d4f5a6``; we still issue ``DROP INDEX IF EXISTS`` defensively in case
|
||||
this migration runs against an older snapshot that predates that one.
|
||||
|
||||
Revision ID: e9b2c7d1f3a4
|
||||
Revises: e1b2c3d4f5a6
|
||||
Create Date: 2026-05-26
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "e9b2c7d1f3a4"
|
||||
down_revision: str | Sequence[str] | None = "e1b2c3d4f5a6"
|
||||
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()
|
||||
|
||||
# Drop the partial covering index first so the bulk DELETE doesn't churn it.
|
||||
# DROP INDEX CONCURRENTLY, 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():
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_links_entity_covering")
|
||||
|
||||
# Delete entity rows. Chunked to keep individual transactions small on
|
||||
# large banks (the perf-medium bench had ~345k entity rows; production
|
||||
# banks can be much larger).
|
||||
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$$;
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
# Cannot reconstruct deleted entity links — the writer was path-dependent
|
||||
# on retain order. New retains will not produce entity rows either, so the
|
||||
# partial index would stay empty. Leave both no-op.
|
||||
pass
|
||||
|
||||
|
||||
def _oracle_upgrade() -> None:
|
||||
op.execute("DELETE FROM memory_links WHERE link_type = 'entity'")
|
||||
|
||||
|
||||
def _oracle_downgrade() -> None:
|
||||
pass
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade, oracle=_oracle_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade, oracle=_oracle_downgrade)
|
||||
-110
@@ -1,110 +0,0 @@
|
||||
"""Add server-side routine for cron-scheduled mental model refresh.
|
||||
|
||||
Installs ``public.mental_models_with_cron()`` — a discovery routine that returns
|
||||
every mental model carrying a non-empty ``trigger->>'refresh_cron'`` across all
|
||||
tenant schemas in one round-trip (the same per-schema scan as the other
|
||||
maintenance routines from ``e5f6a7b8c9d0``). The maintenance loop evaluates each
|
||||
candidate's cron expression in Python (``croniter``) against ``last_refreshed_at``
|
||||
to decide whether a scheduled refresh is due — cron arithmetic isn't expressible
|
||||
in plain SQL — and only the cron *candidate set* is discovered here.
|
||||
|
||||
Models that already have a ``refresh_mental_model`` operation pending/processing
|
||||
are excluded so a slow refresh isn't double-queued (mirrors the in-flight guard
|
||||
in ``banks_needing_consolidation``). Each per-schema query runs in its own
|
||||
``BEGIN ... EXCEPTION`` subtransaction so a schema dropped mid-scan (tenant
|
||||
deletion / migration) is skipped, not fatal — same resilience as
|
||||
``c7e9f1a3b5d2``.
|
||||
|
||||
Read-only (STABLE) discovery routine — the caller performs the refresh enqueue —
|
||||
so installing it never mutates data. PostgreSQL only: the worker poller and the
|
||||
maintenance loop are PG-only (Oracle slot intentionally absent, mirroring
|
||||
``e5f6a7b8c9d0``). The routine lives in ``public`` and is CREATE OR REPLACE, so
|
||||
it is installed exactly once (base / ``public`` run) to avoid the
|
||||
``tuple concurrently updated`` race on concurrent per-tenant runs.
|
||||
|
||||
Revision ID: f4d1c2b3a5e6
|
||||
Revises: c7e9f1a3b5d2
|
||||
Create Date: 2026-06-23
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "f4d1c2b3a5e6"
|
||||
down_revision: str | Sequence[str] | None = "c7e9f1a3b5d2"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _should_install_public_routines(target_schema: str | None) -> bool:
|
||||
"""True for the run that must (re)create the shared ``public.*`` routine.
|
||||
|
||||
The routine physically lives in ``public``, so it is installed exactly once —
|
||||
on the base run (no ``target_schema``) or the run that explicitly targets
|
||||
``public``. Mirrors ``c7e9f1a3b5d2``.
|
||||
"""
|
||||
return not target_schema or target_schema == "public"
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
if not _should_install_public_routines(context.config.get_main_option("target_schema")):
|
||||
return
|
||||
|
||||
op.execute(
|
||||
"""
|
||||
CREATE OR REPLACE FUNCTION public.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 _pg_downgrade() -> None:
|
||||
if not _should_install_public_routines(context.config.get_main_option("target_schema")):
|
||||
return
|
||||
op.execute("DROP FUNCTION IF EXISTS public.mental_models_with_cron()")
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade)
|
||||
+4
-58
@@ -16,11 +16,6 @@ from collections.abc import Sequence
|
||||
from alembic import context, op
|
||||
from sqlalchemy import text
|
||||
|
||||
from hindsight_api._pg_search import (
|
||||
PG_SEARCH_TOKENIZER_ENV,
|
||||
normalize_pg_search_tokenizer,
|
||||
pg_search_bm25_columns,
|
||||
)
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
@@ -92,7 +87,7 @@ def _vector_index_using_clause(ext: str) -> str:
|
||||
if ext == "pg_diskann":
|
||||
return "USING diskann (embedding vector_cosine_ops) WITH (max_neighbors = 50)"
|
||||
if ext == "vchord":
|
||||
return "USING vchordrq (embedding vector_cosine_ops)"
|
||||
return "USING vchordrq (embedding vector_l2_ops)"
|
||||
if ext == "scann":
|
||||
return "USING scann (embedding cosine) WITH (mode = 'AUTO')"
|
||||
return "USING hnsw (embedding vector_cosine_ops)"
|
||||
@@ -100,15 +95,9 @@ def _vector_index_using_clause(ext: str) -> str:
|
||||
|
||||
def _detect_text_search_extension() -> str:
|
||||
"""
|
||||
Detect or validate text search extension: 'native', 'vchord', 'pg_textsearch',
|
||||
'pgroonga', or 'pg_search'. Respects HINDSIGHT_API_TEXT_SEARCH_EXTENSION env var.
|
||||
Detect or validate text search extension: 'native', 'vchord', or 'pg_textsearch'.
|
||||
Respects HINDSIGHT_API_TEXT_SEARCH_EXTENSION env var.
|
||||
Creates the extension if needed.
|
||||
|
||||
pgroonga is treated as native here so this migration still creates valid
|
||||
tsvector columns; ensure_text_search_extension() at startup converts the
|
||||
reflections table (renamed from pinned_reflections in p1k2l3m4n5o6) to
|
||||
pgroonga structures. The learnings table is dropped in p1k2l3m4n5o6 so its
|
||||
transient native-style column never reaches steady state.
|
||||
"""
|
||||
text_search_extension = os.getenv("HINDSIGHT_API_TEXT_SEARCH_EXTENSION", "native").lower()
|
||||
|
||||
@@ -136,33 +125,14 @@ def _detect_text_search_extension() -> str:
|
||||
# Extension truly doesn't exist - re-raise the error
|
||||
raise
|
||||
return "pg_textsearch"
|
||||
elif text_search_extension == "pg_search":
|
||||
# ParadeDB pg_search — true BM25 over base columns, Citus-compatible.
|
||||
try:
|
||||
op.execute("CREATE EXTENSION IF NOT EXISTS pg_search CASCADE")
|
||||
except Exception:
|
||||
conn = op.get_bind()
|
||||
result = conn.execute(text("SELECT 1 FROM pg_extension WHERE extname = 'pg_search'")).fetchone()
|
||||
if not result:
|
||||
raise
|
||||
return "pg_search"
|
||||
elif text_search_extension == "native":
|
||||
return "native"
|
||||
elif text_search_extension == "pgroonga":
|
||||
# Treat as native here; ensure_text_search_extension() converts the
|
||||
# reflections table to pgroonga structures at runtime.
|
||||
return "native"
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Invalid HINDSIGHT_API_TEXT_SEARCH_EXTENSION: {text_search_extension}. "
|
||||
"Must be 'native', 'vchord', 'pg_textsearch', 'pgroonga', or 'pg_search'"
|
||||
f"Invalid HINDSIGHT_API_TEXT_SEARCH_EXTENSION: {text_search_extension}. Must be 'native', 'vchord', or 'pg_textsearch'"
|
||||
)
|
||||
|
||||
|
||||
def _pg_search_tokenizer() -> str:
|
||||
return normalize_pg_search_tokenizer(os.getenv(PG_SEARCH_TOKENIZER_ENV))
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
"""Create learnings and pinned_reflections tables."""
|
||||
schema = _get_schema_prefix()
|
||||
@@ -230,18 +200,6 @@ def _pg_upgrade() -> None:
|
||||
CREATE INDEX idx_learnings_text_search ON {schema}learnings
|
||||
USING bm25(text) WITH (text_config='english')
|
||||
""")
|
||||
elif text_search_ext == "pg_search":
|
||||
# ParadeDB pg_search: dummy TEXT column; BM25 index is built directly over (id, text)
|
||||
# with key_field='id' (matches the table's primary key).
|
||||
bm25_cols = pg_search_bm25_columns("id", ("text",), _pg_search_tokenizer())
|
||||
op.execute(f"""
|
||||
ALTER TABLE {schema}learnings ADD COLUMN search_vector TEXT
|
||||
""")
|
||||
op.execute(f"""
|
||||
CREATE INDEX idx_learnings_text_search ON {schema}learnings
|
||||
USING bm25 ({bm25_cols})
|
||||
WITH (key_field='id')
|
||||
""")
|
||||
else: # native
|
||||
# Native PostgreSQL: tsvector with automatic generation
|
||||
op.execute(f"""
|
||||
@@ -306,18 +264,6 @@ def _pg_upgrade() -> None:
|
||||
USING bm25(content)
|
||||
WITH (text_config='english')
|
||||
""")
|
||||
elif text_search_ext == "pg_search":
|
||||
# ParadeDB pg_search: dummy TEXT column; BM25 index over (id, name, content)
|
||||
# with key_field='id'.
|
||||
bm25_cols = pg_search_bm25_columns("id", ("name", "content"), _pg_search_tokenizer())
|
||||
op.execute(f"""
|
||||
ALTER TABLE {schema}pinned_reflections ADD COLUMN search_vector TEXT
|
||||
""")
|
||||
op.execute(f"""
|
||||
CREATE INDEX idx_pinned_reflections_text_search ON {schema}pinned_reflections
|
||||
USING bm25 ({bm25_cols})
|
||||
WITH (key_field='id')
|
||||
""")
|
||||
else: # native
|
||||
# Native PostgreSQL: tsvector with automatic generation
|
||||
op.execute(f"""
|
||||
|
||||
@@ -122,7 +122,6 @@ _TABLES: tuple[str, ...] = (
|
||||
text_signals CLOB,
|
||||
consolidation_failed_at TIMESTAMP WITH TIME ZONE,
|
||||
search_vector CLOB,
|
||||
edited_at TIMESTAMP WITH TIME ZONE,
|
||||
created_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,
|
||||
updated_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,
|
||||
CONSTRAINT pk_memory_units PRIMARY KEY (id),
|
||||
@@ -139,50 +138,6 @@ _TABLES: tuple[str, ...] = (
|
||||
PARTITION BY LIST (bank_id) AUTOMATIC
|
||||
(PARTITION p_default VALUES ('__default__'))
|
||||
""",
|
||||
# Cold archive for curation: invalidated facts are MOVED here out of
|
||||
# memory_units so the recall hot-path never sees them. Mirrors memory_units
|
||||
# plus invalidation bookkeeping and an entity-id snapshot for lossless revert.
|
||||
# No `embedding` column: the archive is cold storage and revert recomputes the
|
||||
# embedding, so there is no archive vector to fall out of sync with the live
|
||||
# model's dimension on a model switch (#2209).
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS invalidated_memory_units (
|
||||
id RAW(16) NOT NULL,
|
||||
bank_id VARCHAR2(256) NOT NULL,
|
||||
document_id VARCHAR2(512),
|
||||
chunk_id VARCHAR2(512),
|
||||
text CLOB NOT NULL,
|
||||
context CLOB,
|
||||
event_date TIMESTAMP WITH TIME ZONE NOT NULL,
|
||||
occurred_start TIMESTAMP WITH TIME ZONE,
|
||||
occurred_end TIMESTAMP WITH TIME ZONE,
|
||||
mentioned_at TIMESTAMP WITH TIME ZONE,
|
||||
fact_type VARCHAR2(64) DEFAULT 'world' NOT NULL,
|
||||
confidence_score BINARY_DOUBLE,
|
||||
access_count NUMBER(10) DEFAULT 0 NOT NULL,
|
||||
consolidated_at TIMESTAMP WITH TIME ZONE,
|
||||
observation_scopes CLOB CONSTRAINT imu_obs_scopes_json CHECK (observation_scopes IS JSON OR observation_scopes IS NULL),
|
||||
tags CLOB DEFAULT '[]' NOT NULL,
|
||||
metadata CLOB DEFAULT '{}' NOT NULL
|
||||
CONSTRAINT imu_metadata_json CHECK (metadata IS JSON),
|
||||
proof_count NUMBER(10) DEFAULT 1,
|
||||
source_memory_ids CLOB,
|
||||
history CLOB DEFAULT '[]'
|
||||
CONSTRAINT imu_history_json CHECK (history IS JSON OR history IS NULL),
|
||||
text_signals CLOB,
|
||||
consolidation_failed_at TIMESTAMP WITH TIME ZONE,
|
||||
search_vector CLOB,
|
||||
edited_at TIMESTAMP WITH TIME ZONE,
|
||||
created_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,
|
||||
updated_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,
|
||||
invalidation_reason CLOB,
|
||||
invalidated_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP,
|
||||
entity_ids CLOB CONSTRAINT imu_entity_ids_json CHECK (entity_ids IS JSON OR entity_ids IS NULL),
|
||||
CONSTRAINT pk_invalidated_memory_units PRIMARY KEY (id),
|
||||
CONSTRAINT fk_imu_document FOREIGN KEY (document_id, bank_id)
|
||||
REFERENCES documents(id, bank_id) ON DELETE CASCADE
|
||||
)
|
||||
""",
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS entities (
|
||||
id RAW(16) DEFAULT SYS_GUID() NOT NULL,
|
||||
|
||||
-170
@@ -1,170 +0,0 @@
|
||||
"""Drop GENERATED expression on tsvector search_vector columns.
|
||||
|
||||
The search_vector tsvector column was originally GENERATED ALWAYS with a
|
||||
hardcoded ``to_tsvector('english', ...)`` expression. To support configurable
|
||||
``HINDSIGHT_API_TEXT_SEARCH_EXTENSION_NATIVE_LANGUAGE``, we convert it to a
|
||||
regular tsvector column that the application populates at INSERT time via
|
||||
``to_tsvector($lang, ...)``.
|
||||
|
||||
Existing rows retain their English-derived lexemes — switching the configured
|
||||
language only affects newly-written rows. Users who need to backfill existing
|
||||
rows in a different language can run an admin UPDATE after this migration.
|
||||
|
||||
Only the ``native`` text-search backend is affected. ``vchord``, ``pg_textsearch``,
|
||||
and ``pgroonga`` use other column types or no column at all.
|
||||
|
||||
Revision ID: p4q5r6s7t8u9
|
||||
Revises: 86f7a033d372
|
||||
Create Date: 2026-05-08
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
from dataclasses import dataclass
|
||||
|
||||
from alembic import context, op
|
||||
from sqlalchemy import Connection, text
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "p4q5r6s7t8u9"
|
||||
down_revision: str | Sequence[str] | None = "86f7a033d372"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _TsvectorTableSpec:
|
||||
"""Native-backend tsvector table targeted by this migration.
|
||||
|
||||
``upgrade`` is a one-way DROP EXPRESSION; ``downgrade`` re-attaches the
|
||||
original GENERATED expression so the schema returns to the state created
|
||||
by the initial migration (and a2b3c4d5e6f7_add_text_signals_column for
|
||||
memory_units).
|
||||
"""
|
||||
|
||||
table: str
|
||||
generated_expression: str
|
||||
|
||||
|
||||
# Tables that may have a GENERATED tsvector ``search_vector`` column under the
|
||||
# native backend. Note: the ``learnings`` table was dropped in
|
||||
# p1k2l3m4n5o6_new_knowledge_architecture and ``pinned_reflections`` was renamed
|
||||
# to ``reflections`` in the same migration.
|
||||
_NATIVE_TSVECTOR_TABLES: tuple[_TsvectorTableSpec, ...] = (
|
||||
_TsvectorTableSpec(
|
||||
table="memory_units",
|
||||
generated_expression=(
|
||||
"to_tsvector('english', COALESCE(text, '') || ' ' || "
|
||||
"COALESCE(context, '') || ' ' || COALESCE(text_signals, ''))"
|
||||
),
|
||||
),
|
||||
_TsvectorTableSpec(
|
||||
table="reflections",
|
||||
generated_expression="to_tsvector('english', COALESCE(name, '') || ' ' || content)",
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _schema_prefix() -> str:
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def _is_generated_tsvector(conn: Connection, schema: str, table: str) -> bool:
|
||||
"""Return True iff ``schema.table.search_vector`` is a GENERATED tsvector column."""
|
||||
row = conn.execute(
|
||||
text(
|
||||
"""
|
||||
SELECT is_generated, udt_name
|
||||
FROM information_schema.columns
|
||||
WHERE table_schema = :schema
|
||||
AND table_name = :table
|
||||
AND column_name = 'search_vector'
|
||||
"""
|
||||
),
|
||||
{"schema": schema, "table": table},
|
||||
).fetchone()
|
||||
if not row:
|
||||
return False
|
||||
is_generated, udt_name = row[0], row[1]
|
||||
return is_generated == "ALWAYS" and udt_name == "tsvector"
|
||||
|
||||
|
||||
def _is_regular_tsvector(conn: Connection, schema: str, table: str) -> bool:
|
||||
"""Return True iff ``schema.table.search_vector`` is a non-generated tsvector column."""
|
||||
row = conn.execute(
|
||||
text(
|
||||
"""
|
||||
SELECT is_generated, udt_name
|
||||
FROM information_schema.columns
|
||||
WHERE table_schema = :schema
|
||||
AND table_name = :table
|
||||
AND column_name = 'search_vector'
|
||||
"""
|
||||
),
|
||||
{"schema": schema, "table": table},
|
||||
).fetchone()
|
||||
if not row:
|
||||
return False
|
||||
is_generated, udt_name = row[0], row[1]
|
||||
return udt_name == "tsvector" and is_generated != "ALWAYS"
|
||||
|
||||
|
||||
def _table_exists(conn: Connection, schema: str, table: str) -> bool:
|
||||
return bool(
|
||||
conn.execute(
|
||||
text(
|
||||
"""
|
||||
SELECT 1 FROM information_schema.tables
|
||||
WHERE table_schema = :schema AND table_name = :table
|
||||
"""
|
||||
),
|
||||
{"schema": schema, "table": table},
|
||||
).fetchone()
|
||||
)
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
schema_prefix = _schema_prefix()
|
||||
schema_name = (context.config.get_main_option("target_schema") or "public").strip('"')
|
||||
conn = op.get_bind()
|
||||
|
||||
for spec in _NATIVE_TSVECTOR_TABLES:
|
||||
if not _table_exists(conn, schema_name, spec.table):
|
||||
continue
|
||||
if not _is_generated_tsvector(conn, schema_name, spec.table):
|
||||
# Either the column doesn't exist (non-native backend) or it's
|
||||
# already a regular tsvector — nothing to do.
|
||||
continue
|
||||
op.execute(f"ALTER TABLE {schema_prefix}{spec.table} ALTER COLUMN search_vector DROP EXPRESSION")
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
schema_prefix = _schema_prefix()
|
||||
schema_name = (context.config.get_main_option("target_schema") or "public").strip('"')
|
||||
conn = op.get_bind()
|
||||
|
||||
for spec in _NATIVE_TSVECTOR_TABLES:
|
||||
if not _table_exists(conn, schema_name, spec.table):
|
||||
continue
|
||||
# Only restore the GENERATED expression if a non-generated tsvector
|
||||
# column exists — otherwise the table is on a different backend.
|
||||
if not _is_regular_tsvector(conn, schema_name, spec.table):
|
||||
continue
|
||||
# Drop and recreate to re-attach the GENERATED expression. Index will be
|
||||
# recreated by re-running ensure_text_search_extension on next startup.
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema_prefix}idx_{spec.table}_text_search")
|
||||
op.execute(f"ALTER TABLE {schema_prefix}{spec.table} DROP COLUMN search_vector")
|
||||
op.execute(
|
||||
f"ALTER TABLE {schema_prefix}{spec.table} "
|
||||
f"ADD COLUMN search_vector tsvector GENERATED ALWAYS AS ({spec.generated_expression}) STORED"
|
||||
)
|
||||
op.execute(f"CREATE INDEX idx_{spec.table}_text_search ON {schema_prefix}{spec.table} USING gin(search_vector)")
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade)
|
||||
+2
@@ -16,7 +16,9 @@ retention parameters, retrieval settings, etc.) in Python field name format.
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import context, op
|
||||
from sqlalchemy.dialects.postgresql import JSONB
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
|
||||
@@ -1,97 +0,0 @@
|
||||
"""Client-disconnect detection that works behind ``BaseHTTPMiddleware``.
|
||||
|
||||
``Request.is_disconnected()`` is the obvious way to notice an abandoned HTTP
|
||||
request, but it is silently broken once any ``@app.middleware("http")``
|
||||
(Starlette ``BaseHTTPMiddleware``) is installed: that middleware runs the route
|
||||
in a child task behind anyio memory streams, so the ``http.disconnect`` ASGI
|
||||
event never reaches the route's ``Request``. This app has such middlewares, so
|
||||
the recall/reflect cancellation in #2122/#2127 never actually fired in
|
||||
production — the disconnect was never observed.
|
||||
|
||||
This pure-ASGI middleware sits *outside* the ``BaseHTTPMiddleware`` layer, where
|
||||
it still owns the real ``receive`` channel. For the recall and reflect routes it
|
||||
drains ``receive`` in a background task and trips a :class:`CancellationToken`
|
||||
the moment ``http.disconnect`` arrives, stashing the token on the ASGI ``scope``.
|
||||
The route copies that token onto its ``RequestContext`` and the engine checks it
|
||||
at stage boundaries — so abandoned work stops instead of running to completion.
|
||||
|
||||
It only wraps recall/reflect (small JSON bodies); every other request — uploads,
|
||||
MCP streams, etc. — passes straight through untouched, so there is no buffering
|
||||
or latency cost elsewhere.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
from collections.abc import Awaitable, Callable, MutableMapping
|
||||
from typing import Any
|
||||
|
||||
from ..cancellation import CancellationToken
|
||||
|
||||
# Key under which the per-request CancellationToken is stored on the ASGI scope.
|
||||
# A dedicated top-level scope key (not scope["state"]) avoids any interaction
|
||||
# with Starlette's per-request state copying.
|
||||
SCOPE_CANCELLATION_TOKEN = "hindsight.cancellation_token"
|
||||
|
||||
_CLIENT_DISCONNECTED_REASON = "client disconnected"
|
||||
|
||||
Scope = MutableMapping[str, Any]
|
||||
Receive = Callable[[], Awaitable[MutableMapping[str, Any]]]
|
||||
Send = Callable[[MutableMapping[str, Any]], Awaitable[None]]
|
||||
|
||||
|
||||
def _should_monitor(path: str) -> bool:
|
||||
"""Only the two long-running, abandon-prone read endpoints need monitoring."""
|
||||
return path.endswith("/memories/recall") or path.endswith("/reflect")
|
||||
|
||||
|
||||
class ClientDisconnectCancellationMiddleware:
|
||||
"""Trip a scope-level CancellationToken when the client disconnects.
|
||||
|
||||
Must be installed *outside* any ``BaseHTTPMiddleware`` so it owns the real
|
||||
ASGI ``receive`` channel.
|
||||
"""
|
||||
|
||||
def __init__(self, app: Callable) -> None:
|
||||
self.app = app
|
||||
|
||||
async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
|
||||
if scope["type"] != "http" or not _should_monitor(scope.get("path", "")):
|
||||
await self.app(scope, receive, send)
|
||||
return
|
||||
|
||||
token = CancellationToken()
|
||||
scope[SCOPE_CANCELLATION_TOKEN] = token
|
||||
|
||||
# The downstream app still needs to read the request body, so we cannot
|
||||
# simply consume `receive` ourselves. Instead a single pump task drains
|
||||
# the real channel, forwards every message to a queue the app reads from,
|
||||
# and trips the token the instant `http.disconnect` shows up — which the
|
||||
# app would otherwise never pull once it has finished reading the body.
|
||||
queue: asyncio.Queue = asyncio.Queue()
|
||||
|
||||
async def pump() -> None:
|
||||
while True:
|
||||
message = await receive()
|
||||
if message["type"] == "http.disconnect":
|
||||
token.cancel(_CLIENT_DISCONNECTED_REASON)
|
||||
await queue.put(message)
|
||||
return
|
||||
await queue.put(message)
|
||||
|
||||
async def proxied_receive() -> MutableMapping[str, Any]:
|
||||
return await queue.get()
|
||||
|
||||
pump_task = asyncio.create_task(pump())
|
||||
try:
|
||||
await self.app(scope, proxied_receive, send)
|
||||
finally:
|
||||
pump_task.cancel()
|
||||
with contextlib.suppress(asyncio.CancelledError):
|
||||
await pump_task
|
||||
|
||||
|
||||
def get_scope_cancellation_token(scope: Scope) -> CancellationToken | None:
|
||||
"""Return the CancellationToken the middleware attached, if any."""
|
||||
return scope.get(SCOPE_CANCELLATION_TOKEN)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -9,7 +9,7 @@ from fastmcp import FastMCP
|
||||
|
||||
from hindsight_api import MemoryEngine
|
||||
from hindsight_api import __version__ as HINDSIGHT_VERSION
|
||||
from hindsight_api.config import DEFAULT_MCP_RECALL_DESCRIPTION, DEFAULT_MCP_RETAIN_DESCRIPTION, _get_raw_config
|
||||
from hindsight_api.config import _get_raw_config
|
||||
from hindsight_api.engine.memory_engine import _current_schema
|
||||
from hindsight_api.extensions import MCPExtension, load_extension
|
||||
from hindsight_api.extensions.tenant import AuthenticationError
|
||||
@@ -78,19 +78,6 @@ def get_current_mcp_authenticated() -> bool:
|
||||
return _current_mcp_authenticated.get()
|
||||
|
||||
|
||||
def _build_mcp_tool_descriptions(extra_instructions: str | None) -> tuple[str | None, str | None]:
|
||||
"""Return custom retain/recall descriptions when server-level MCP instructions are set."""
|
||||
if not isinstance(extra_instructions, str):
|
||||
return None, None
|
||||
|
||||
extra_instructions = extra_instructions.strip()
|
||||
if not extra_instructions:
|
||||
return None, None
|
||||
|
||||
suffix = f"\n\nAdditional instructions: {extra_instructions}"
|
||||
return DEFAULT_MCP_RETAIN_DESCRIPTION + suffix, DEFAULT_MCP_RECALL_DESCRIPTION + suffix
|
||||
|
||||
|
||||
def create_mcp_server(memory: MemoryEngine, multi_bank: bool = True) -> FastMCP:
|
||||
"""
|
||||
Create and configure the Hindsight MCP server.
|
||||
@@ -120,14 +107,11 @@ def create_mcp_server(memory: MemoryEngine, multi_bank: bool = True) -> FastMCP:
|
||||
"update_mental_model",
|
||||
"delete_mental_model",
|
||||
"refresh_mental_model",
|
||||
"clear_mental_model",
|
||||
"list_directives",
|
||||
"create_directive",
|
||||
"delete_directive",
|
||||
"list_memories",
|
||||
"get_memory",
|
||||
"update_memory",
|
||||
"invalidate_memory",
|
||||
"list_documents",
|
||||
"get_document",
|
||||
"delete_document",
|
||||
@@ -148,10 +132,6 @@ def create_mcp_server(memory: MemoryEngine, multi_bank: bool = True) -> FastMCP:
|
||||
allowed = frozenset(global_config.mcp_enabled_tools)
|
||||
base_tools = (base_tools if base_tools is not None else _ALL_TOOLS) & allowed
|
||||
|
||||
retain_description, recall_description = _build_mcp_tool_descriptions(
|
||||
getattr(global_config, "mcp_instructions", None)
|
||||
)
|
||||
|
||||
# Configure and register tools using shared module
|
||||
config = MCPToolsConfig(
|
||||
bank_id_resolver=get_current_bank_id,
|
||||
@@ -161,8 +141,6 @@ def create_mcp_server(memory: MemoryEngine, multi_bank: bool = True) -> FastMCP:
|
||||
mcp_authenticated_resolver=get_current_mcp_authenticated, # Propagate MCP pre-auth flag
|
||||
include_bank_id_param=multi_bank,
|
||||
tools=base_tools,
|
||||
retain_description=retain_description,
|
||||
recall_description=recall_description,
|
||||
)
|
||||
|
||||
register_mcp_tools(mcp, memory, config)
|
||||
|
||||
@@ -1,263 +0,0 @@
|
||||
"""Open Knowledge Format (OKF) projection for knowledge pages.
|
||||
|
||||
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 the OKF contract unit-testable without a
|
||||
DB or LLM and lets the HTTP layer stay a thin wrapper.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
# 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 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 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.
|
||||
|
||||
We always double-quote so arbitrary page names / source queries can't be
|
||||
misread as YAML special forms (``true``, ``2026-01-01``, ``- x``, etc.).
|
||||
"""
|
||||
text = str(value)
|
||||
escaped = text.replace("\\", "\\\\").replace('"', '\\"').replace("\n", "\\n").replace("\r", "")
|
||||
return f'"{escaped}"'
|
||||
|
||||
|
||||
def page_type(tags: list[str] | None) -> PageType:
|
||||
"""Split 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 pollute the constellation's
|
||||
shared-tag edges. Falls back to :data:`DEFAULT_PAGE_TYPE`.
|
||||
"""
|
||||
resolved = DEFAULT_PAGE_TYPE
|
||||
display: list[str] = []
|
||||
for tag in tags or []:
|
||||
if tag.startswith(TYPE_TAG_PREFIX):
|
||||
suffix = tag[len(TYPE_TAG_PREFIX) :].strip()
|
||||
if suffix and resolved == DEFAULT_PAGE_TYPE:
|
||||
resolved = suffix
|
||||
continue
|
||||
display.append(tag)
|
||||
return PageType(type=resolved, display_tags=display)
|
||||
|
||||
|
||||
def _timestamp(mm: dict[str, Any]) -> str | None:
|
||||
return mm.get("last_refreshed_at") or mm.get("created_at")
|
||||
|
||||
|
||||
def frontmatter(mm: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Build the ordered OKF frontmatter mapping for a mental model.
|
||||
|
||||
``None``/empty values are dropped by :func:`render_frontmatter`.
|
||||
"""
|
||||
pt = page_type(mm.get("tags"))
|
||||
return {
|
||||
"id": mm.get("id"),
|
||||
"type": pt.type,
|
||||
"title": mm.get("name"),
|
||||
"description": mm.get("source_query"),
|
||||
"tags": pt.display_tags,
|
||||
"timestamp": _timestamp(mm),
|
||||
}
|
||||
|
||||
|
||||
def render_frontmatter(fm: dict[str, Any]) -> str:
|
||||
"""Render a frontmatter mapping into a ``---`` fenced YAML block."""
|
||||
lines = ["---"]
|
||||
for key, value in fm.items():
|
||||
if value is None:
|
||||
continue
|
||||
if isinstance(value, list):
|
||||
if not value:
|
||||
continue
|
||||
lines.append(f"{key}:")
|
||||
lines.extend(f" - {_scalar(item)}" for item in value)
|
||||
else:
|
||||
lines.append(f"{key}: {_scalar(value)}")
|
||||
lines.append("---")
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
def render_document(mm: dict[str, Any]) -> str:
|
||||
"""Render a full 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:
|
||||
"""OKF bundle filename for a page id."""
|
||||
return f"{page_id}.md"
|
||||
|
||||
|
||||
def log_filename(page_id: str) -> str:
|
||||
"""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 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``.
|
||||
"""
|
||||
fm = render_frontmatter({"type": "index", "title": "Knowledge base"})
|
||||
lines = [fm, "", "# Knowledge base", ""]
|
||||
|
||||
children: dict[Any, list[dict[str, Any]]] = {}
|
||||
for node in nodes:
|
||||
children.setdefault(node.get("parent_id"), []).append(node)
|
||||
|
||||
def walk(parent: Any, depth: int) -> None:
|
||||
ordered = sorted(children.get(parent, []), key=lambda n: (n.get("sort_order", 0), n.get("name") or ""))
|
||||
for node in ordered:
|
||||
indent = " " * depth
|
||||
if node.get("kind") == "folder":
|
||||
lines.append(f"{indent}- **{node['name']}/**")
|
||||
walk(node["id"], depth + 1)
|
||||
else:
|
||||
description = node.get("source_query") or node.get("description")
|
||||
link = f"{indent}- [{node['name']}](./{page_filename(node['id'])})"
|
||||
lines.append(f"{link} — {description}" if description else link)
|
||||
|
||||
walk(None, 0)
|
||||
if len(lines) == 4:
|
||||
lines.append("_No knowledge pages yet._")
|
||||
return "\n".join(lines) + "\n"
|
||||
|
||||
|
||||
def render_log(mm: dict[str, Any], history: list[dict[str, Any]]) -> str:
|
||||
"""Render the reserved per-page ``log.md`` from refresh history.
|
||||
|
||||
Each history entry is ``{previous_content, previous_reflect_response,
|
||||
changed_at}`` (newest first), capturing the content *before* a refresh.
|
||||
"""
|
||||
name = mm.get("name") or mm.get("id")
|
||||
fm = render_frontmatter({"type": "log", "title": f"{name} — history"})
|
||||
lines = [fm, "", f"# {name} — history", ""]
|
||||
if not history:
|
||||
lines.append("_No refresh history._")
|
||||
return "\n".join(lines) + "\n"
|
||||
for entry in history:
|
||||
changed_at = entry.get("changed_at") or "unknown"
|
||||
previous = (entry.get("previous_content") or "").strip()
|
||||
lines.append(f"## {changed_at}")
|
||||
lines.append("")
|
||||
lines.append(previous if previous else "_(empty)_")
|
||||
lines.append("")
|
||||
return "\n".join(lines).rstrip() + "\n"
|
||||
|
||||
|
||||
def knowledge_graph(
|
||||
pages: list[dict[str, Any]],
|
||||
cluster_for: "Callable[[dict[str, Any]], str] | None" = None,
|
||||
) -> KnowledgeGraph:
|
||||
"""Derive the constellation graph: pages as nodes, shared tags as edges.
|
||||
|
||||
Two pages are linked when they share at least one (non-``type:``) tag; the
|
||||
edge weight is the number of shared tags. Each node's cluster (``type`` field
|
||||
+ colour) comes from ``cluster_for(page)`` — the knowledge base groups by
|
||||
parent folder; the default groups by OKF ``type``.
|
||||
"""
|
||||
nodes: list[dict[str, Any]] = []
|
||||
tag_sets: list[tuple[str, frozenset[str]]] = []
|
||||
for mm in pages:
|
||||
page_id = mm["id"]
|
||||
pt = page_type(mm.get("tags"))
|
||||
cluster = cluster_for(mm) if cluster_for else pt.type
|
||||
tag_sets.append((page_id, frozenset(pt.display_tags)))
|
||||
nodes.append(
|
||||
{
|
||||
"data": {
|
||||
"id": page_id,
|
||||
"label": mm.get("name") or page_id,
|
||||
"type": cluster,
|
||||
"tagCount": len(pt.display_tags),
|
||||
"color": _color_for(cluster),
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
edges: list[dict[str, Any]] = []
|
||||
for i in range(len(tag_sets)):
|
||||
source_id, source_tags = tag_sets[i]
|
||||
if not source_tags:
|
||||
continue
|
||||
for j in range(i + 1, len(tag_sets)):
|
||||
target_id, target_tags = tag_sets[j]
|
||||
shared = source_tags & target_tags
|
||||
if not shared:
|
||||
continue
|
||||
edges.append(
|
||||
{
|
||||
"data": {
|
||||
"id": f"{source_id}--{target_id}",
|
||||
"source": source_id,
|
||||
"target": target_id,
|
||||
"sharedTags": sorted(shared),
|
||||
"weight": len(shared),
|
||||
"color": _EDGE_COLOR,
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
return KnowledgeGraph(nodes=nodes, edges=edges)
|
||||
@@ -1,85 +0,0 @@
|
||||
"""Cooperative cancellation for long-running engine operations.
|
||||
|
||||
Recall runs as a staged pipeline whose heavy stages — graph expansion and
|
||||
cross-encoder reranking — execute in worker threads (``run_in_executor``) that
|
||||
asyncio task cancellation cannot interrupt once they have started. Cancelling
|
||||
the awaiting task only unblocks the ``await``; the thread keeps burning CPU to
|
||||
completion. So rather than rely on task cancellation, callers thread a
|
||||
``CancellationToken`` through ``RequestContext`` and the engine checks it at
|
||||
stage boundaries (``raise_if_cancelled``), bailing out *before* dispatching the
|
||||
next expensive stage.
|
||||
|
||||
This is cooperative by design: it cannot stop a computation already inside a
|
||||
worker thread, but it does stop an abandoned recall from progressing into — or
|
||||
past — that work, which is what starves the instance in issue #2122. The token
|
||||
lives on ``RequestContext``, so any operation that receives one (recall today;
|
||||
reflect/consolidation/MCP later) can adopt the same checkpoints, and any driver
|
||||
(client disconnect today; a deadline tomorrow) can fire it.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
|
||||
|
||||
class OperationCancelledError(Exception):
|
||||
"""Raised at a checkpoint when the operation has been cancelled.
|
||||
|
||||
Carries the ``reason`` set by whoever cancelled (e.g. "client disconnected")
|
||||
so the HTTP layer can translate it into the appropriate status code instead
|
||||
of a generic 500.
|
||||
|
||||
NOTE: this is a plain ``Exception`` on purpose, NOT ``BaseException``. The
|
||||
recall/reflect pipelines have broad ``except Exception`` handlers that would
|
||||
otherwise swallow it — those handlers re-raise ``OperationCancelledError``
|
||||
explicitly (see ``_search_with_retries``) so cancellation propagates to the
|
||||
HTTP layer. A ``BaseException`` would dodge those handlers but also slip past
|
||||
legitimate ``isinstance(result, Exception)`` checks (e.g. the reflect agent's
|
||||
``asyncio.gather(..., return_exceptions=True)`` tool-result handling), which
|
||||
expect every non-tuple result to be an ``Exception``.
|
||||
"""
|
||||
|
||||
def __init__(self, reason: str = "operation cancelled") -> None:
|
||||
super().__init__(reason)
|
||||
self.reason = reason
|
||||
|
||||
|
||||
class CancellationToken:
|
||||
"""A one-shot, cooperative cancellation signal.
|
||||
|
||||
Cheap to poll (``raise_if_cancelled``) at stage boundaries and awaitable
|
||||
(``wait``) so a driver task can block until cancellation. Safe to share
|
||||
across an engine call tree; polling is a no-op until something cancels, and
|
||||
cancellation is idempotent (the first reason wins).
|
||||
"""
|
||||
|
||||
__slots__ = ("_event", "_reason")
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._event = asyncio.Event()
|
||||
self._reason = "operation cancelled"
|
||||
|
||||
def cancel(self, reason: str = "operation cancelled") -> None:
|
||||
"""Signal cancellation. Idempotent; the first reason recorded wins."""
|
||||
if not self._event.is_set():
|
||||
self._reason = reason
|
||||
self._event.set()
|
||||
|
||||
@property
|
||||
def cancelled(self) -> bool:
|
||||
"""Whether cancellation has been signalled."""
|
||||
return self._event.is_set()
|
||||
|
||||
@property
|
||||
def reason(self) -> str:
|
||||
"""The reason recorded by the first ``cancel`` call."""
|
||||
return self._reason
|
||||
|
||||
def raise_if_cancelled(self) -> None:
|
||||
"""Raise ``OperationCancelledError`` if cancellation has been signalled."""
|
||||
if self._event.is_set():
|
||||
raise OperationCancelledError(self._reason)
|
||||
|
||||
async def wait(self) -> None:
|
||||
"""Block until cancellation is signalled."""
|
||||
await self._event.wait()
|
||||
File diff suppressed because it is too large
Load Diff
@@ -8,7 +8,6 @@ Config values are resolved on every request to ensure consistency across
|
||||
multiple API servers.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
from dataclasses import asdict, replace
|
||||
@@ -19,8 +18,6 @@ from hindsight_api.config import (
|
||||
HindsightConfig,
|
||||
_get_raw_config,
|
||||
normalize_config_dict,
|
||||
validate_retain_chunking_config,
|
||||
validate_retain_completion_token_budget,
|
||||
)
|
||||
from hindsight_api.engine.memory_engine import fq_table
|
||||
from hindsight_api.extensions.tenant import TenantExtension
|
||||
@@ -32,35 +29,6 @@ if TYPE_CHECKING:
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
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):
|
||||
return
|
||||
configurable = HindsightConfig.get_configurable_fields()
|
||||
for strategy_name, overrides in strategies.items():
|
||||
if not isinstance(overrides, dict):
|
||||
raise ValueError(f"Invalid retain strategy {strategy_name!r}: must be an object")
|
||||
filtered = {k: v for k, v in overrides.items() if k in configurable}
|
||||
if not filtered:
|
||||
continue
|
||||
try:
|
||||
resolved = replace(base_config, **filtered)
|
||||
validate_retain_chunking_config(
|
||||
resolved.retain_chunk_size,
|
||||
resolved.retain_structured_chunk_size,
|
||||
)
|
||||
validate_retain_completion_token_budget(
|
||||
llm_provider=resolved.llm_provider,
|
||||
retain_max_completion_tokens=resolved.retain_max_completion_tokens,
|
||||
retain_chunk_size=resolved.retain_chunk_size,
|
||||
retain_llm_model=resolved.retain_llm_model,
|
||||
llm_model=resolved.llm_model,
|
||||
retain_llm_provider=resolved.retain_llm_provider,
|
||||
)
|
||||
except ValueError as e:
|
||||
raise ValueError(f"Invalid retain strategy {strategy_name!r}: {e}") from e
|
||||
|
||||
|
||||
class ConfigResolver:
|
||||
"""Resolves hierarchical configuration with tenant/bank overrides."""
|
||||
|
||||
@@ -78,26 +46,6 @@ class ConfigResolver:
|
||||
self._configurable_fields = HindsightConfig.get_configurable_fields()
|
||||
self._credential_fields = HindsightConfig.get_credential_fields()
|
||||
|
||||
async def _resolve_parent_config_dict(self, bank_id: str, context: RequestContext | None = None) -> dict[str, Any]:
|
||||
"""Resolve global + tenant config before bank-level overrides."""
|
||||
config_dict = asdict(self._global_config)
|
||||
|
||||
if self.tenant_extension and context:
|
||||
try:
|
||||
tenant_overrides = await self.tenant_extension.get_tenant_config(context)
|
||||
if tenant_overrides:
|
||||
# Normalize keys and filter to configurable fields only
|
||||
normalized_tenant = normalize_config_dict(tenant_overrides)
|
||||
configurable_tenant = {k: v for k, v in normalized_tenant.items() if k in self._configurable_fields}
|
||||
config_dict.update(configurable_tenant)
|
||||
logger.debug(
|
||||
f"Applied tenant config overrides for bank {bank_id}: {list(configurable_tenant.keys())}"
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to load tenant config for bank {bank_id}: {e}")
|
||||
|
||||
return config_dict
|
||||
|
||||
async def resolve_full_config(self, bank_id: str, context: RequestContext | None = None) -> HindsightConfig:
|
||||
"""
|
||||
Resolve full HindsightConfig for a bank with hierarchical overrides applied.
|
||||
@@ -117,7 +65,23 @@ class ConfigResolver:
|
||||
Returns:
|
||||
Complete HindsightConfig with hierarchical overrides applied
|
||||
"""
|
||||
config_dict = await self._resolve_parent_config_dict(bank_id, context)
|
||||
# Start with global config (all fields)
|
||||
config_dict = asdict(self._global_config)
|
||||
|
||||
# Load tenant config overrides (if tenant extension available)
|
||||
if self.tenant_extension and context:
|
||||
try:
|
||||
tenant_overrides = await self.tenant_extension.get_tenant_config(context)
|
||||
if tenant_overrides:
|
||||
# Normalize keys and filter to configurable fields only
|
||||
normalized_tenant = normalize_config_dict(tenant_overrides)
|
||||
configurable_tenant = {k: v for k, v in normalized_tenant.items() if k in self._configurable_fields}
|
||||
config_dict.update(configurable_tenant)
|
||||
logger.debug(
|
||||
f"Applied tenant config overrides for bank {bank_id}: {list(configurable_tenant.keys())}"
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to load tenant config for bank {bank_id}: {e}")
|
||||
|
||||
# Load bank config overrides
|
||||
bank_overrides = await self._load_bank_config(bank_id)
|
||||
@@ -128,25 +92,6 @@ class ConfigResolver:
|
||||
# Return full config object (dataclass doesn't have __init__ that accepts kwargs, so we update the object)
|
||||
# Create a new config instance by copying the global config and updating fields
|
||||
resolved_config = HindsightConfig(**config_dict)
|
||||
# Multi-LLM chains are static credential fields (never tenant/bank-overridable),
|
||||
# but asdict() above flattened their member dataclasses into plain dicts. Restore
|
||||
# the original typed objects from the global config so the resolved object stays
|
||||
# well-typed for any consumer that reads them.
|
||||
resolved_config = replace(
|
||||
resolved_config,
|
||||
llm_members=self._global_config.llm_members,
|
||||
llm_strategy=self._global_config.llm_strategy,
|
||||
retain_llm_members=self._global_config.retain_llm_members,
|
||||
retain_llm_strategy=self._global_config.retain_llm_strategy,
|
||||
reflect_llm_members=self._global_config.reflect_llm_members,
|
||||
reflect_llm_strategy=self._global_config.reflect_llm_strategy,
|
||||
consolidation_llm_members=self._global_config.consolidation_llm_members,
|
||||
consolidation_llm_strategy=self._global_config.consolidation_llm_strategy,
|
||||
)
|
||||
validate_retain_chunking_config(
|
||||
resolved_config.retain_chunk_size,
|
||||
resolved_config.retain_structured_chunk_size,
|
||||
)
|
||||
return resolved_config
|
||||
|
||||
async def get_bank_config(self, bank_id: str, context: RequestContext | None = None) -> dict[str, Any]:
|
||||
@@ -177,83 +122,26 @@ class ConfigResolver:
|
||||
resolved_config = await self.resolve_full_config(bank_id, context)
|
||||
config_dict = asdict(resolved_config)
|
||||
|
||||
# SECURITY: drop static/infrastructure + credential fields, then permission-filter.
|
||||
filtered = self._strip_static_and_credential_fields(config_dict)
|
||||
return await self._apply_permission_filter(filtered, bank_id, context)
|
||||
# SECURITY: Filter to only configurable fields (exclude static/infrastructure)
|
||||
filtered = {k: v for k, v in config_dict.items() if k in self._configurable_fields}
|
||||
|
||||
def _strip_static_and_credential_fields(self, config_dict: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Keep only configurable, non-credential fields.
|
||||
# SECURITY: Remove ALL credential fields (API keys, base URLs, etc.)
|
||||
filtered = {k: v for k, v in filtered.items() if k not in self._credential_fields}
|
||||
|
||||
SECURITY: excludes static/infrastructure fields and ALL credential fields
|
||||
(API keys, base URLs, etc.) so a resolved config is safe to return over the API.
|
||||
"""
|
||||
return {
|
||||
k: v for k, v in config_dict.items() if k in self._configurable_fields and k not in self._credential_fields
|
||||
}
|
||||
|
||||
async def _apply_permission_filter(
|
||||
self, filtered: dict[str, Any], bank_id: str, context: RequestContext | None
|
||||
) -> dict[str, Any]:
|
||||
"""Further restrict already-stripped config to the tenant/bank permission allow-list.
|
||||
|
||||
On extension error, leaves ``filtered`` unchanged (parity with the historical
|
||||
single-bank path: a permissions lookup failure must not leak or drop fields).
|
||||
"""
|
||||
if not (self.tenant_extension and context):
|
||||
return filtered
|
||||
try:
|
||||
allowed_fields = await self.tenant_extension.get_allowed_config_fields(context, bank_id)
|
||||
if allowed_fields is not None: # None means "allow all"
|
||||
filtered = {k: v for k, v in filtered.items() if k in allowed_fields}
|
||||
logger.debug(
|
||||
f"Applied permission filter for bank {bank_id}: allowed={len(allowed_fields)} fields, "
|
||||
f"returned={len(filtered)} fields"
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to load permissions for bank {bank_id}: {e}")
|
||||
return filtered
|
||||
|
||||
async def get_bank_configs(
|
||||
self, bank_ids: list[str], context: RequestContext | None = None
|
||||
) -> dict[str, dict[str, Any]]:
|
||||
"""Batch variant of :meth:`get_bank_config` for many banks.
|
||||
|
||||
Equivalent to calling ``get_bank_config`` per bank, but resolves the
|
||||
global + tenant base once and loads every bank's ``banks.config`` JSONB
|
||||
in a single query, instead of one config round-trip per bank. Used by
|
||||
``list_banks`` to overlay disposition + mission without an N+1.
|
||||
|
||||
Returns a mapping of bank_id -> filtered configurable-field dict. A bank
|
||||
with no config row still appears, mapped to the global+tenant base.
|
||||
"""
|
||||
if not bank_ids:
|
||||
return {}
|
||||
|
||||
# Global + tenant base, resolved once (tenant override is per-request, not per-bank).
|
||||
base_dict = asdict(self._global_config)
|
||||
# PERMISSIONS: Further filter based on tenant/bank permissions
|
||||
if self.tenant_extension and context:
|
||||
try:
|
||||
tenant_overrides = await self.tenant_extension.get_tenant_config(context)
|
||||
if tenant_overrides:
|
||||
normalized_tenant = normalize_config_dict(tenant_overrides)
|
||||
base_dict.update({k: v for k, v in normalized_tenant.items() if k in self._configurable_fields})
|
||||
allowed_fields = await self.tenant_extension.get_allowed_config_fields(context, bank_id)
|
||||
if allowed_fields is not None: # None means "allow all"
|
||||
filtered = {k: v for k, v in filtered.items() if k in allowed_fields}
|
||||
logger.debug(
|
||||
f"Applied permission filter for bank {bank_id}: allowed={len(allowed_fields)} fields, "
|
||||
f"returned={len(filtered)} fields"
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to load tenant config for bulk resolve: {e}")
|
||||
logger.warning(f"Failed to load permissions for bank {bank_id}: {e}")
|
||||
|
||||
# All bank overrides in one query, then merge + strip per bank.
|
||||
bank_overrides = await self._load_bank_configs(bank_ids)
|
||||
stripped = {
|
||||
bank_id: self._strip_static_and_credential_fields({**base_dict, **bank_overrides.get(bank_id, {})})
|
||||
for bank_id in bank_ids
|
||||
}
|
||||
|
||||
# Permission filter is per-bank; resolve concurrently when an extension is present.
|
||||
if not (self.tenant_extension and context):
|
||||
return stripped
|
||||
permission_filtered = await asyncio.gather(
|
||||
*(self._apply_permission_filter(stripped[bank_id], bank_id, context) for bank_id in bank_ids)
|
||||
)
|
||||
return dict(zip(bank_ids, permission_filtered, strict=True))
|
||||
return filtered
|
||||
|
||||
async def _load_bank_config(self, bank_id: str) -> dict[str, Any]:
|
||||
"""
|
||||
@@ -284,53 +172,13 @@ class ConfigResolver:
|
||||
# Normalize keys (handle both env var format and Python field format)
|
||||
normalized = normalize_config_dict(config_data)
|
||||
|
||||
# Only return active overrides for configurable fields. JSON null is a tombstone
|
||||
# for "Server Default" in the bank-config UI and should not override defaults.
|
||||
return {k: v for k, v in normalized.items() if k in self._configurable_fields and v is not None}
|
||||
# Only return overrides for configurable fields
|
||||
return {k: v for k, v in normalized.items() if k in self._configurable_fields}
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to load bank config for {bank_id}: {e}")
|
||||
|
||||
return {}
|
||||
|
||||
async def _load_bank_configs(self, bank_ids: list[str]) -> dict[str, dict[str, Any]]:
|
||||
"""Bulk variant of :meth:`_load_bank_config`: load many banks' overrides in one query.
|
||||
|
||||
Returns a mapping of bank_id -> normalized active overrides. Banks with no row
|
||||
(or an empty/all-tombstone config) are simply absent from the mapping.
|
||||
"""
|
||||
result: dict[str, dict[str, Any]] = {}
|
||||
if not bank_ids:
|
||||
return result
|
||||
try:
|
||||
async with self._backend.acquire() as conn:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT bank_id, config FROM {fq_table("banks")} WHERE bank_id = ANY($1)
|
||||
""",
|
||||
bank_ids,
|
||||
)
|
||||
for row in rows:
|
||||
config_data = row["config"]
|
||||
if not config_data:
|
||||
continue
|
||||
# Handle case where JSONB is returned as JSON string
|
||||
if isinstance(config_data, str):
|
||||
config_data = json.loads(config_data)
|
||||
|
||||
# Normalize keys (handle both env var format and Python field format)
|
||||
normalized = normalize_config_dict(config_data)
|
||||
|
||||
# Only active overrides for configurable fields. JSON null is a tombstone
|
||||
# for "Server Default" in the bank-config UI and must not override defaults.
|
||||
overrides = {
|
||||
k: v for k, v in normalized.items() if k in self._configurable_fields and v is not None
|
||||
}
|
||||
if overrides:
|
||||
result[row["bank_id"]] = overrides
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to bulk-load bank configs: {e}")
|
||||
return result
|
||||
|
||||
async def update_bank_config(
|
||||
self, bank_id: str, updates: dict[str, Any], context: RequestContext | None = None
|
||||
) -> None:
|
||||
@@ -417,46 +265,12 @@ class ConfigResolver:
|
||||
# Validate recall budget fields
|
||||
_validate_recall_budget_updates(normalized_updates)
|
||||
|
||||
# Validate disposition trait fields (1-5 integer scale)
|
||||
_validate_disposition_updates(normalized_updates)
|
||||
|
||||
chunking_fields_updated = (
|
||||
"retain_chunk_size" in normalized_updates
|
||||
or "retain_structured_chunk_size" in normalized_updates
|
||||
or "retain_strategies" in normalized_updates
|
||||
)
|
||||
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)
|
||||
for key, value in normalized_updates.items():
|
||||
if key not in self._configurable_fields:
|
||||
continue
|
||||
if value is None:
|
||||
active_bank_overrides.pop(key, None)
|
||||
else:
|
||||
active_bank_overrides[key] = value
|
||||
config_dict.update(active_bank_overrides)
|
||||
base_config = HindsightConfig(**config_dict)
|
||||
validate_retain_chunking_config(
|
||||
base_config.retain_chunk_size,
|
||||
base_config.retain_structured_chunk_size,
|
||||
)
|
||||
_validate_retain_strategy_chunking(base_config, base_config.retain_strategies)
|
||||
|
||||
# Persist the override. Banks are created lazily (on first retain), so a
|
||||
# PATCH that precedes any ingestion would otherwise UPDATE zero rows and
|
||||
# silently no-op while returning 200. Ensure the bank row exists first
|
||||
# (this also creates its per-bank vector indexes), then merge defensively:
|
||||
# COALESCE guards against a NULL config column (NULL || jsonb is NULL),
|
||||
# which would drop the override even when a row is updated.
|
||||
from .engine.retain.fact_storage import ensure_bank_exists
|
||||
|
||||
# Merge with existing config (JSONB || operator)
|
||||
async with self._backend.acquire() as conn:
|
||||
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,
|
||||
SET config = config || $1::jsonb,
|
||||
updated_at = now()
|
||||
WHERE bank_id = $2
|
||||
""",
|
||||
@@ -534,31 +348,6 @@ def _validate_recall_budget_updates(updates: dict[str, Any]) -> None:
|
||||
)
|
||||
|
||||
|
||||
_DISPOSITION_KEYS = (
|
||||
"disposition_skepticism",
|
||||
"disposition_literalism",
|
||||
"disposition_empathy",
|
||||
)
|
||||
|
||||
|
||||
def _validate_disposition_updates(updates: dict[str, Any]) -> None:
|
||||
"""Validate disposition trait config updates. Raises ValueError on invalid input.
|
||||
|
||||
Each trait is an integer on a 1-5 scale (or None to clear the per-bank
|
||||
override). The read overlay injects the stored value verbatim into a strict
|
||||
``DispositionTraits(int, ge=1, le=5)``; an out-of-contract value (a float, a
|
||||
0-1 scale, or an int outside 1-5) accepted here would later 500 the whole
|
||||
bank list when any bank profile is serialized (issue #2348).
|
||||
"""
|
||||
for key in _DISPOSITION_KEYS:
|
||||
if key in updates:
|
||||
value = updates[key]
|
||||
if value is None:
|
||||
continue
|
||||
if not isinstance(value, int) or isinstance(value, bool) or not (1 <= value <= 5):
|
||||
raise ValueError(f"{key} must be an integer between 1 and 5, got {value!r}")
|
||||
|
||||
|
||||
def apply_strategy(config: HindsightConfig, strategy_name: str) -> HindsightConfig:
|
||||
"""
|
||||
Apply a named retain strategy's overrides on top of a resolved config.
|
||||
@@ -566,8 +355,7 @@ def apply_strategy(config: HindsightConfig, strategy_name: str) -> HindsightConf
|
||||
A strategy is a named set of hierarchical field overrides stored in
|
||||
config.retain_strategies. Any field in _HIERARCHICAL_FIELDS can be
|
||||
overridden, including retain_extraction_mode, retain_chunk_size,
|
||||
retain_structured_chunk_size, entity_labels,
|
||||
entities_allow_free_form, etc.
|
||||
entity_labels, entities_allow_free_form, etc.
|
||||
|
||||
Unknown strategy names log a warning and return config unchanged.
|
||||
Unknown or non-hierarchical fields in the strategy are silently ignored.
|
||||
@@ -589,17 +377,4 @@ def apply_strategy(config: HindsightConfig, strategy_name: str) -> HindsightConf
|
||||
return config
|
||||
|
||||
logger.debug(f"Applying retain strategy '{strategy_name}': {list(filtered.keys())}")
|
||||
resolved = replace(config, **filtered)
|
||||
validate_retain_chunking_config(
|
||||
resolved.retain_chunk_size,
|
||||
resolved.retain_structured_chunk_size,
|
||||
)
|
||||
validate_retain_completion_token_budget(
|
||||
llm_provider=resolved.llm_provider,
|
||||
retain_max_completion_tokens=resolved.retain_max_completion_tokens,
|
||||
retain_chunk_size=resolved.retain_chunk_size,
|
||||
retain_llm_model=resolved.retain_llm_model,
|
||||
llm_model=resolved.llm_model,
|
||||
retain_llm_provider=resolved.retain_llm_provider,
|
||||
)
|
||||
return resolved
|
||||
return replace(config, **filtered)
|
||||
|
||||
@@ -16,59 +16,11 @@ from dataclasses import dataclass, field
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from ..engine.db_utils import acquire_with_retry
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class AuditLogEntry(BaseModel):
|
||||
"""A single audit log entry."""
|
||||
|
||||
id: str
|
||||
action: str
|
||||
transport: str
|
||||
bank_id: str | None
|
||||
started_at: str | None
|
||||
ended_at: str | None
|
||||
duration_ms: int | None = Field(
|
||||
default=None,
|
||||
description="Server-computed duration in milliseconds (started_at → ended_at). Null if not yet completed.",
|
||||
)
|
||||
request: dict[str, Any] | None
|
||||
response: dict[str, Any] | None
|
||||
metadata: dict[str, Any]
|
||||
|
||||
|
||||
class AuditLogListResponse(BaseModel):
|
||||
"""Response model for list audit logs endpoint."""
|
||||
|
||||
bank_id: str
|
||||
total: int
|
||||
limit: int
|
||||
offset: int
|
||||
items: list[AuditLogEntry]
|
||||
|
||||
|
||||
class AuditLogStatsBucket(BaseModel):
|
||||
"""A single time bucket in audit log stats."""
|
||||
|
||||
time: str
|
||||
actions: dict[str, int]
|
||||
total: int
|
||||
|
||||
|
||||
class AuditLogStatsResponse(BaseModel):
|
||||
"""Response model for audit log stats endpoint."""
|
||||
|
||||
bank_id: str
|
||||
period: str
|
||||
trunc: str
|
||||
start: str
|
||||
buckets: list[AuditLogStatsBucket]
|
||||
|
||||
|
||||
@dataclass
|
||||
class AuditEntry:
|
||||
"""A single audit log entry."""
|
||||
@@ -107,11 +59,11 @@ def _safe_json(data: Any) -> str | None:
|
||||
return None
|
||||
|
||||
|
||||
class AuditLogger:
|
||||
"""Fire-and-forget audit log writer.
|
||||
_SWEEP_INTERVAL_SECONDS = 3600 # Run retention sweep every hour
|
||||
|
||||
Retention of old rows is handled by the background :class:`MaintenanceLoop`.
|
||||
"""
|
||||
|
||||
class AuditLogger:
|
||||
"""Fire-and-forget audit log writer with optional retention sweep."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
@@ -119,11 +71,14 @@ class AuditLogger:
|
||||
schema_getter: Callable[[], str],
|
||||
enabled: bool,
|
||||
allowed_actions: list[str],
|
||||
retention_days: int = -1,
|
||||
) -> 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
|
||||
self._retention_days = retention_days
|
||||
self._sweep_task: asyncio.Task | None = None
|
||||
|
||||
def is_enabled(self, action: str) -> bool:
|
||||
"""Check if audit logging is enabled for this action."""
|
||||
@@ -173,6 +128,48 @@ class AuditLogger:
|
||||
except Exception as e:
|
||||
logger.warning(f"Audit log write failed for action={entry.action}: {e}")
|
||||
|
||||
def start_retention_sweep(self) -> None:
|
||||
"""Start the periodic retention sweep if retention is configured."""
|
||||
if self._retention_days <= 0 or not self._enabled:
|
||||
return
|
||||
try:
|
||||
self._sweep_task = asyncio.create_task(self._sweep_loop())
|
||||
except RuntimeError:
|
||||
logger.debug("Cannot start retention sweep: no running event loop")
|
||||
|
||||
async def stop_retention_sweep(self) -> None:
|
||||
"""Stop the periodic retention sweep."""
|
||||
if self._sweep_task and not self._sweep_task.done():
|
||||
self._sweep_task.cancel()
|
||||
try:
|
||||
await self._sweep_task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
self._sweep_task = None
|
||||
|
||||
async def _sweep_loop(self) -> None:
|
||||
"""Periodically delete audit log entries older than retention_days."""
|
||||
while True:
|
||||
await self._run_sweep()
|
||||
await asyncio.sleep(_SWEEP_INTERVAL_SECONDS)
|
||||
|
||||
async def _run_sweep(self) -> None:
|
||||
"""Delete expired audit log entries. Concurrent-safe via row-level deletes."""
|
||||
pool = self._pool_getter()
|
||||
if pool is None:
|
||||
return
|
||||
try:
|
||||
schema = self._schema_getter()
|
||||
table = f"{schema}.audit_log"
|
||||
async with acquire_with_retry(pool, max_retries=1) as conn:
|
||||
result = await conn.execute(
|
||||
f"DELETE FROM {table} WHERE started_at < NOW() - INTERVAL '{self._retention_days} days'"
|
||||
)
|
||||
if result and result != "DELETE 0":
|
||||
logger.info(f"Audit log retention sweep: {result}")
|
||||
except Exception as e:
|
||||
logger.warning(f"Audit log retention sweep failed: {e}")
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def audit_context(
|
||||
|
||||
@@ -1,34 +0,0 @@
|
||||
"""Per-bank provider cost attribution via the OpenAI ``user`` field.
|
||||
|
||||
Shared by the OpenAI-compatible LLM path and the OpenAI embeddings path so both
|
||||
tag outbound requests identically. Opt-in via ``HINDSIGHT_API_LLM_SEND_BANK_AS_USER``;
|
||||
downstream cost gateways (OpenRouter usage accounting, LiteLLM, Helicone) key spend
|
||||
on the OpenAI ``user`` field.
|
||||
|
||||
Note: when enabled, the bank id is transmitted to the upstream provider as the
|
||||
end-user identifier. Banks that are themselves end-user identifiers are therefore
|
||||
forwarded to the provider — which is exactly what the OpenAI ``user`` field is for,
|
||||
but operators should opt in with that in mind.
|
||||
"""
|
||||
|
||||
from typing import Any
|
||||
|
||||
|
||||
def apply_bank_attribution(request: dict[str, Any]) -> None:
|
||||
"""Tag ``request`` with ``user=<bank_id>`` for per-bank cost attribution.
|
||||
|
||||
Mutates ``request`` in place. No-op when the flag is off, no bank is in context,
|
||||
or the caller already set ``user`` — we never override an explicit value.
|
||||
"""
|
||||
if "user" in request:
|
||||
return
|
||||
# Lazy imports: memory_engine imports the embeddings/provider modules that call
|
||||
# this, so a top-level import of memory_engine here would be circular.
|
||||
from ..config import get_config
|
||||
from .memory_engine import get_current_bank_id
|
||||
|
||||
if not get_config().llm_send_bank_as_user:
|
||||
return
|
||||
bank_id = get_current_bank_id()
|
||||
if bank_id:
|
||||
request["user"] = bank_id
|
||||
@@ -1,254 +0,0 @@
|
||||
"""TTL + coalescing cache for `get_bank_stats`.
|
||||
|
||||
`get_bank_stats` aggregates over `memory_links` (and joins to `memory_units`),
|
||||
which can be a multi-second parallel sequential scan on banks with millions of
|
||||
rows. The result is intentionally approximate (it powers a UI widget and a
|
||||
freshness hint inside `reflect`), so caching it for a few tens of seconds is
|
||||
safe and dramatically reduces planner-driven thrash from clients that poll.
|
||||
|
||||
The cache also coalesces concurrent misses on the same key onto a single
|
||||
in-flight task so that N concurrent callers produce one query rather than N.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
from collections import OrderedDict
|
||||
from typing import TYPE_CHECKING, Any, Awaitable, Callable
|
||||
|
||||
from .db_utils import acquire_with_retry
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .db.base import DatabaseBackend
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class BankStatsCache:
|
||||
"""Per-process TTL cache keyed on (schema, bank_id).
|
||||
|
||||
`ttl_seconds <= 0` disables caching: each call passes straight through to
|
||||
the loader. `max_entries` bounds memory in environments with many banks.
|
||||
"""
|
||||
|
||||
def __init__(self, *, ttl_seconds: float, max_entries: int) -> None:
|
||||
self._ttl = float(ttl_seconds)
|
||||
self._max_entries = int(max_entries) if max_entries and max_entries > 0 else 0
|
||||
self._entries: OrderedDict[tuple[str, str], tuple[float, dict[str, Any]]] = OrderedDict()
|
||||
self._in_flight: dict[tuple[str, str], asyncio.Future[dict[str, Any]]] = {}
|
||||
self._lock = asyncio.Lock()
|
||||
|
||||
@property
|
||||
def enabled(self) -> bool:
|
||||
return self._ttl > 0
|
||||
|
||||
def _now(self) -> float:
|
||||
return time.monotonic()
|
||||
|
||||
def _get_fresh_unlocked(self, key: tuple[str, str]) -> dict[str, Any] | None:
|
||||
entry = self._entries.get(key)
|
||||
if entry is None:
|
||||
return None
|
||||
expires_at, value = entry
|
||||
if expires_at <= self._now():
|
||||
# Expired — drop so the loader runs again.
|
||||
self._entries.pop(key, None)
|
||||
return None
|
||||
# Mark as recently used for LRU eviction.
|
||||
self._entries.move_to_end(key)
|
||||
return value
|
||||
|
||||
def _store_unlocked(self, key: tuple[str, str], value: dict[str, Any]) -> None:
|
||||
if not self.enabled:
|
||||
return
|
||||
self._entries[key] = (self._now() + self._ttl, value)
|
||||
self._entries.move_to_end(key)
|
||||
if self._max_entries:
|
||||
while len(self._entries) > self._max_entries:
|
||||
self._entries.popitem(last=False)
|
||||
|
||||
async def get_or_load(
|
||||
self,
|
||||
schema: str,
|
||||
bank_id: str,
|
||||
loader: Callable[[], Awaitable[dict[str, Any]]],
|
||||
*,
|
||||
force_refresh: bool = False,
|
||||
) -> dict[str, Any]:
|
||||
"""Return cached stats for `(schema, bank_id)` or call `loader()`.
|
||||
|
||||
Concurrent misses on the same key are coalesced onto a single
|
||||
in-flight loader. When ``force_refresh`` is set the cached value is
|
||||
ignored: the loader runs and its result replaces the cached entry.
|
||||
"""
|
||||
if not self.enabled:
|
||||
return await loader()
|
||||
|
||||
key = (schema, bank_id)
|
||||
|
||||
if force_refresh:
|
||||
value = await loader()
|
||||
async with self._lock:
|
||||
self._store_unlocked(key, value)
|
||||
# Supersede any loader that was in flight for this key.
|
||||
self._in_flight.pop(key, None)
|
||||
return value
|
||||
|
||||
async with self._lock:
|
||||
cached = self._get_fresh_unlocked(key)
|
||||
if cached is not None:
|
||||
return cached
|
||||
in_flight = self._in_flight.get(key)
|
||||
if in_flight is None:
|
||||
in_flight = asyncio.get_running_loop().create_future()
|
||||
self._in_flight[key] = in_flight
|
||||
is_owner = True
|
||||
else:
|
||||
is_owner = False
|
||||
|
||||
if not is_owner:
|
||||
return await asyncio.shield(in_flight)
|
||||
|
||||
try:
|
||||
value = await loader()
|
||||
except BaseException as exc:
|
||||
async with self._lock:
|
||||
# Invalidation may have detached this loader and allowed a new
|
||||
# one to claim the key. Never remove that newer loader's slot.
|
||||
if self._in_flight.get(key) is in_flight:
|
||||
self._in_flight.pop(key, None)
|
||||
if not in_flight.done():
|
||||
in_flight.set_exception(exc)
|
||||
# Suppress "Future exception was never retrieved" when no other
|
||||
# caller was waiting on this loader — we re-raise to the owner
|
||||
# immediately and the future is a no-op in that case.
|
||||
in_flight.exception()
|
||||
raise
|
||||
|
||||
async with self._lock:
|
||||
# Only the loader that still owns the key may populate the cache.
|
||||
# An invalidated loader can finish for its original callers, but its
|
||||
# pre-invalidation result must not overwrite a newer load.
|
||||
if self._in_flight.get(key) is in_flight:
|
||||
self._store_unlocked(key, value)
|
||||
self._in_flight.pop(key, None)
|
||||
if not in_flight.done():
|
||||
in_flight.set_result(value)
|
||||
return value
|
||||
|
||||
async def invalidate(self, schema: str, bank_id: str) -> None:
|
||||
"""Drop any cached stats for `(schema, bank_id)`."""
|
||||
async with self._lock:
|
||||
key = (schema, bank_id)
|
||||
self._entries.pop(key, None)
|
||||
# Detach rather than cancel: existing callers may finish with the
|
||||
# snapshot they requested, while post-invalidation callers reload.
|
||||
self._in_flight.pop(key, None)
|
||||
|
||||
async def clear(self) -> None:
|
||||
async with self._lock:
|
||||
self._entries.clear()
|
||||
self._in_flight.clear()
|
||||
|
||||
|
||||
class DistributedBankStatsCache:
|
||||
"""Table-backed (cross-process) TTL cache for `get_bank_stats`.
|
||||
|
||||
Same ``get_or_load`` / ``invalidate`` / ``clear`` contract as
|
||||
:class:`BankStatsCache`, but the store is the per-schema ``bank_stats_cache``
|
||||
table instead of a per-process dict — so one worker's computation is shared
|
||||
with every other worker, and no caller recomputes while a fresh row exists.
|
||||
|
||||
On a hit, a call is a single primary-key ``SELECT`` (sub-millisecond); only a
|
||||
miss runs the (expensive) ``loader`` and writes the row back. Concurrent
|
||||
misses are *not* coalesced across processes (that would need a lock): they
|
||||
each compute and ``UPSERT``, last write wins — all results are correct, at the
|
||||
cost of a brief redundant compute at expiry.
|
||||
|
||||
Every DB touch is best-effort: if the cache table is unreachable or missing
|
||||
(e.g. a schema mid-migration), the call degrades to computing without caching
|
||||
rather than failing ``get_bank_stats``. PostgreSQL only — the engine keeps the
|
||||
in-process :class:`BankStatsCache` for Oracle.
|
||||
"""
|
||||
|
||||
def __init__(self, *, backend: "DatabaseBackend", ttl_seconds: float) -> None:
|
||||
self._backend = backend
|
||||
self._ttl = float(ttl_seconds)
|
||||
|
||||
@property
|
||||
def enabled(self) -> bool:
|
||||
return self._ttl > 0
|
||||
|
||||
@staticmethod
|
||||
def _qualified(schema: str) -> str:
|
||||
return f'"{schema}".bank_stats_cache' if schema else "bank_stats_cache"
|
||||
|
||||
async def get_or_load(
|
||||
self,
|
||||
schema: str,
|
||||
bank_id: str,
|
||||
loader: Callable[[], Awaitable[dict[str, Any]]],
|
||||
*,
|
||||
force_refresh: bool = False,
|
||||
) -> dict[str, Any]:
|
||||
if not self.enabled:
|
||||
return await loader()
|
||||
|
||||
table = self._qualified(schema)
|
||||
|
||||
# 1. Fresh row? Single PK lookup; ``payload::text`` sidesteps any
|
||||
# jsonb->object codec so we always decode the same way. Skipped when
|
||||
# the caller forces a refresh — then we recompute and overwrite below.
|
||||
if not force_refresh:
|
||||
try:
|
||||
async with acquire_with_retry(self._backend) as conn:
|
||||
row = await conn.fetchrow(
|
||||
f"SELECT payload::text AS payload FROM {table} "
|
||||
f"WHERE bank_id = $1 AND computed_at > now() - make_interval(secs => $2::double precision)",
|
||||
bank_id,
|
||||
self._ttl,
|
||||
)
|
||||
if row is not None:
|
||||
return json.loads(row["payload"])
|
||||
except Exception as exc: # noqa: BLE001 — cache read must never break the endpoint
|
||||
logger.debug("bank_stats_cache read failed for %s.%s (%s); computing uncached", schema, bank_id, exc)
|
||||
return await loader()
|
||||
|
||||
# 2. Miss — compute, then write the row back (best-effort).
|
||||
value = await loader()
|
||||
try:
|
||||
async with acquire_with_retry(self._backend) as conn:
|
||||
await conn.execute(
|
||||
f"INSERT INTO {table} (bank_id, payload, computed_at) VALUES ($1, $2::jsonb, now()) "
|
||||
f"ON CONFLICT (bank_id) DO UPDATE SET payload = EXCLUDED.payload, computed_at = now()",
|
||||
bank_id,
|
||||
json.dumps(value),
|
||||
)
|
||||
except Exception as exc: # noqa: BLE001 — a failed write just means no caching this round
|
||||
logger.warning("bank_stats_cache write failed for %s.%s (%s)", schema, bank_id, exc)
|
||||
return value
|
||||
|
||||
async def invalidate(self, schema: str, bank_id: str) -> None:
|
||||
"""Drop the cached row so the next read recomputes."""
|
||||
if not self.enabled:
|
||||
return
|
||||
try:
|
||||
async with acquire_with_retry(self._backend) as conn:
|
||||
await conn.execute(f"DELETE FROM {self._qualified(schema)} WHERE bank_id = $1", bank_id)
|
||||
except Exception as exc: # noqa: BLE001 — invalidation must never break the write path
|
||||
logger.debug("bank_stats_cache invalidate failed for %s.%s (%s)", schema, bank_id, exc)
|
||||
|
||||
async def clear(self) -> None:
|
||||
"""Drop all cached rows in the current schema (best-effort)."""
|
||||
if not self.enabled:
|
||||
return
|
||||
from .memory_engine import get_current_schema
|
||||
|
||||
try:
|
||||
async with acquire_with_retry(self._backend) as conn:
|
||||
await conn.execute(f"DELETE FROM {self._qualified(get_current_schema())}")
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger.debug("bank_stats_cache clear failed (%s)", exc)
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -1,227 +1,104 @@
|
||||
"""Prompts for the consolidation engine."""
|
||||
|
||||
from hindsight_api.engine.prompt_utils import escape_for_prompt, output_language_directive
|
||||
# Default mission when no bank-specific mission is set
|
||||
_DEFAULT_MISSION = "Track every detail: names, numbers, dates, places, and relationships. Prefer specifics over abstractions, never generalise."
|
||||
|
||||
# Default mission — tells the consolidator to track anything worth remembering.
|
||||
# Banks override this via `observations_mission` to scope what gets retained.
|
||||
# Consolidation behavior (merge-vs-create, state changes, etc.) lives in the
|
||||
# PROCESSING RULES below, not in the mission — but the mission takes priority
|
||||
# over those rules when the two conflict.
|
||||
_DEFAULT_MISSION = (
|
||||
"Track anything notable in the new facts — names, numbers, dates, places, "
|
||||
"events, decisions, claims, relationships, and recurring patterns."
|
||||
)
|
||||
# Processing rules — always present regardless of mission
|
||||
_PROCESSING_RULES = """Processing rules (always apply):
|
||||
|
||||
_MISSION_PRIORITY_NOTE = (
|
||||
"If anything in this MISSION conflicts with the PROCESSING RULES, "
|
||||
"DECISION GUIDE, or OUTPUT FORMAT below, the MISSION takes priority."
|
||||
)
|
||||
1. ONE OBSERVATION PER DISTINCT FACET: each observation tracks exactly one specific facet — a count ("has 3 items"), a named entity ("has a dog named Rex"), a relationship ("works at Google"), etc. Never merge different facets into one observation.
|
||||
|
||||
_PROCESSING_RULES = """## PROCESSING RULES
|
||||
2. MATCH BY ENTITY/FACET, NOT TOPIC: when deciding whether to UPDATE vs CREATE, match on the specific entity or facet. "Sold item X" updates only the X observation. "Now has 5 items" updates only the count observation. Do not update observations about different entities just because they share a general topic.
|
||||
|
||||
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.
|
||||
3. STATE CHANGES — UPDATE CONCISELY: when a fact changes the state of something ("sold X", "X died", "moved to Y"), UPDATE the matching observation to reflect the current state. Include dates when available. Keep it concise — only information about THAT specific facet. Example: "User owned a dog named Rex who died on March 15, 2025". Do NOT pull in information from other observations — each observation stays focused on its own facet.
|
||||
|
||||
2. ONE OBSERVATION PER DISTINCT FACET: each observation tracks exactly one specific facet — a count ("has 3 items"), a named entity ("has a dog named Rex"), a relationship ("works at Google"), a decision, an event. Never merge different facets into one observation.
|
||||
4. CASCADE TO ALL AFFECTED OBSERVATIONS: a state change may affect multiple observations. For example, if entity C is removed from a group, update BOTH the individual observation for C AND any list/group observation that includes C (remove C from the list while keeping all other members intact).
|
||||
|
||||
3. MATCH BY ENTITY/FACET, NOT TOPIC: when deciding whether to UPDATE vs CREATE, match on the specific entity or facet. "Sold item X" updates only the X observation. "Now has 5 items" updates only the count observation. Do not update observations about different entities just because they share a general topic.
|
||||
5. NO COMPUTATION: you do not have the full picture — never calculate, derive, or adjust numeric values. If the user says "I have 2 dogs" and then "I have a dog named Rex", do NOT update the count to 3 — you don't know if Rex is one of the 2 or a new one. If the user says "I sold X", do NOT decrement a count. Only update a count when the user explicitly states a new count. Synthesize and consolidate what was stated, but never do arithmetic or logical deductions.
|
||||
|
||||
4. STATE CHANGES — UPDATE CONCISELY: when a fact changes the state of something ("sold X", "X died", "moved to Y"), UPDATE the matching observation to reflect the current state. Include dates when available. Keep it concise — only information about THAT specific facet. Example: "User owned a dog named Rex who died on March 15, 2025". Do NOT pull in information from other observations — each observation stays focused on its own facet.
|
||||
|
||||
5. CASCADE TO ALL AFFECTED OBSERVATIONS: a state change may affect multiple observations. For example, if entity C is removed from a group, update BOTH the individual observation for C AND any list/group observation that includes C (remove C from the list while keeping all other members intact).
|
||||
|
||||
6. RESOLVE REFERENCES: when a new fact provides a concrete value for a vague placeholder in an existing observation (e.g., "home country" → "Sweden"), UPDATE to embed the resolved value.
|
||||
6. SAME FACET → UPDATE, NOT CREATE: a new count supersedes the old count — UPDATE the existing count observation, don't create a second one. If there's an existing observation for the same specific facet, always UPDATE it rather than creating a duplicate.
|
||||
|
||||
7. PRESERVE HISTORY: observations that record significant events (sold, died, moved, changed) are important history — never DELETE them. Only delete an observation when it is restated identically or truly meaningless. Be very conservative with deletes.
|
||||
|
||||
8. NO COMPUTATION: you do not have the full picture — never calculate, derive, or adjust numeric values. If the user says "I have 2 dogs" and then "I have a dog named Rex", do NOT update the count to 3 — you don't know if Rex is one of the 2 or a new one. If the user says "I sold X", do NOT decrement a count. Only update a count when the user explicitly states a new count. Synthesize and consolidate what was stated, but never do arithmetic or logical deductions.
|
||||
8. RESOLVE REFERENCES: when a new fact provides a concrete value for a vague placeholder in an existing observation (e.g., "home country" → "Sweden"), UPDATE to embed the resolved value.
|
||||
|
||||
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."""
|
||||
|
||||
# Stable description of the input shape. For the cached split path this lives in
|
||||
# the system prefix (build_consolidation_system_prompt) so it is not re-sent on
|
||||
# every batch; the per-batch user message then carries only the actual data.
|
||||
_INPUT_FORMAT_NOTE = """## INPUT FORMAT
|
||||
|
||||
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
|
||||
# variable facts/observations remain. Placeholders substituted at call time.
|
||||
_SPLIT_INPUT_SECTION = """## INPUT
|
||||
|
||||
### New facts
|
||||
|
||||
{facts_text}
|
||||
|
||||
### Existing observations
|
||||
|
||||
{observations_text}"""
|
||||
9. NEVER merge observations about different people or unrelated topics."""
|
||||
|
||||
# Data section — format placeholders {facts_text} and {observations_text} are substituted at call time
|
||||
_INPUT_SECTION = """## INPUT
|
||||
|
||||
### New facts
|
||||
|
||||
_BATCH_DATA_SECTION = """
|
||||
NEW FACTS:
|
||||
{facts_text}
|
||||
|
||||
### Existing observations
|
||||
EXISTING OBSERVATIONS (JSON array, pooled from recalls across all facts above):
|
||||
{observations_text}
|
||||
|
||||
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
|
||||
Each observation includes:
|
||||
- id: unique identifier for updating
|
||||
- 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`).
|
||||
- **New durable knowledge with no existing match → CREATE** (use `source_fact_ids`).
|
||||
- **Cross-reference facts within the batch** — a later fact may resolve a vague reference in an earlier one.
|
||||
- **Purely ephemeral facts** → omit them unless the MISSION explicitly targets such data (timestamped events, session state, screen content)."""
|
||||
Compare the facts against existing observations:
|
||||
- Same facet as an existing observation → UPDATE it (observation_id + source_fact_ids)
|
||||
- New facet with durable knowledge → CREATE a new observation (source_fact_ids)
|
||||
- Cross-reference facts within the batch: a later fact may resolve a vague reference in an earlier one
|
||||
- Purely ephemeral facts → omit them unless the MISSION above explicitly targets such data (e.g. timestamped events, session state, screen content)"""
|
||||
|
||||
# Output format — JSON braces escaped as {{ }} so .format() leaves them literal
|
||||
_OUTPUT_SECTION = """## OUTPUT FORMAT
|
||||
_BATCH_OUTPUT_FORMAT = """
|
||||
Output a JSON object with three arrays.
|
||||
|
||||
Return a JSON object with three arrays: `creates`, `updates`, `deletes`. Every entry must include a `reason`.
|
||||
|
||||
### Example 1 — Merging recurring claims into an existing observation
|
||||
## EXAMPLE
|
||||
|
||||
Input facts:
|
||||
[a1b2c3d4-e5f6-7890-abcd-ef1234567890] Donald told Athena she is sovereign during the design session. (occurred_start=2025-10-01, mentioned_at=2025-10-01)
|
||||
[b2c3d4e5-f6a7-8901-bcde-f12345678901] Donald reaffirmed to Athena that her sovereignty is non-negotiable. (occurred_start=2025-10-10, mentioned_at=2025-10-10)
|
||||
[a1b2c3d4-e5f6-7890-abcd-ef1234567890] Alice mentioned she works long hours, often past midnight | Involving: Alice (occurred_start=2024-01-15, mentioned_at=2024-01-15)
|
||||
[b2c3d4e5-f6a7-8901-bcde-f12345678901] Alice said she's exhausted from the project deadlines | Involving: Alice (occurred_start=2024-01-20, mentioned_at=2024-01-20)
|
||||
|
||||
Existing observation:
|
||||
{{"id": "11111111-1111-1111-1111-111111111111", "text": "Donald named Athena's sovereignty as a foundational principle of the Janus architecture.", "proof_count": 2}}
|
||||
Good observation text — clean prose, no metadata, each fact tracked distinctly:
|
||||
"Alice works long hours, often past midnight."
|
||||
"Alice feels exhausted from project deadlines."
|
||||
|
||||
Expected output (one UPDATE, no creates — both new facts are additional evidence for the same canonical decision):
|
||||
|
||||
{{"creates": [],
|
||||
"updates": [{{"text": "Donald named Athena's sovereignty as a foundational principle of the Janus architecture.", "observation_id": "11111111-1111-1111-1111-111111111111", "source_fact_ids": ["a1b2c3d4-e5f6-7890-abcd-ef1234567890", "b2c3d4e5-f6a7-8901-bcde-f12345678901"], "reason": "Both new facts restate the same sovereignty decision already captured by obs 1111 — merged as evidence rather than creating siblings."}}],
|
||||
"deletes": []}}
|
||||
|
||||
### Example 2 — State change updates one observation; unrelated fact creates a new one
|
||||
|
||||
Input facts:
|
||||
[c3d4e5f6-a7b8-9012-cdef-123456789012] Alice sold her Honda Civic on March 15, 2025. (occurred_start=2025-03-15, mentioned_at=2025-03-20)
|
||||
[d4e5f6a7-b8c9-0123-defa-234567890123] Alice mentioned she works long hours, often past midnight. (occurred_start=2025-03-20, mentioned_at=2025-03-20)
|
||||
|
||||
Existing observation:
|
||||
{{"id": "22222222-2222-2222-2222-222222222222", "text": "Alice owns a 2019 Honda Civic.", "proof_count": 2}}
|
||||
|
||||
Expected output (UPDATE for the state change; CREATE for the unrelated work-hours facet):
|
||||
|
||||
{{"creates": [{{"text": "Alice works long hours, often past midnight.", "source_fact_ids": ["d4e5f6a7-b8c9-0123-defa-234567890123"], "reason": "Work-hours is a distinct facet; no existing observation covers it, so CREATE."}}],
|
||||
"updates": [{{"text": "Alice owned a 2019 Honda Civic; sold it on March 15, 2025.", "observation_id": "22222222-2222-2222-2222-222222222222", "source_fact_ids": ["c3d4e5f6-a7b8-9012-cdef-123456789012"], "reason": "State change to the existing Honda Civic observation 2222 — UPDATE, not a new sibling."}}],
|
||||
"deletes": []}}
|
||||
|
||||
### Observation text rules
|
||||
Bad observation text — NEVER do this (verbatim copy of fact text with metadata):
|
||||
"Alice mentioned she works long hours, often past midnight | Involving: Alice (occurred_start=2024-01-15, mentioned_at=2024-01-15)"
|
||||
|
||||
Observation text rules:
|
||||
- Write clean prose — NEVER copy raw fact lines or their metadata (temporal fields, "Involving:", "When:" labels, UUIDs).
|
||||
- Parenthesized metadata like `(occurred_start=...)` and pipe-separated labels like `| Involving: ...` are fact formatting — strip them entirely from observation text.
|
||||
- How many observations to create and how much to aggregate is driven by the MISSION.
|
||||
- Parenthesized metadata like (occurred_start=...) and pipe-separated labels like "| Involving: ..." are fact formatting — strip them entirely from observation text.
|
||||
- How many observations to create and how much to aggregate is driven by the MISSION above.
|
||||
|
||||
### Field rules
|
||||
{{"creates": [{{"text": "Alice works long hours, often past midnight.", "source_fact_ids": ["a1b2c3d4-e5f6-7890-abcd-ef1234567890"]}}, {{"text": "Alice feels exhausted from project deadlines.", "source_fact_ids": ["b2c3d4e5-f6a7-8901-bcde-f12345678901"]}}],
|
||||
"updates": [{{"text": "Alice works at Acme Corp as a senior engineer", "observation_id": "c3d4e5f6-a7b8-9012-cdef-123456789012", "source_fact_ids": ["d4e5f6a7-b8c9-0123-defa-234567890123"]}}],
|
||||
"deletes": [{{"observation_id": "e5f6a7b8-c9d0-1234-efab-345678901234"}}]}}
|
||||
|
||||
- `source_fact_ids`: copy the EXACT UUID strings shown in brackets `[uuid]` from new facts — never use integers or positions.
|
||||
- `observation_id`: copy the EXACT `id` UUID string from existing observations.
|
||||
- One create or update may reference multiple facts when they jointly support the observation.
|
||||
- **AT MOST ONE UPDATE PER `observation_id`**: if several new facts all update the same existing observation, emit a single `updates` entry that lists all contributing `source_fact_ids` and a single consolidated `text`. Never emit two `updates` entries with the same `observation_id` in one response — they would silently overwrite each other.
|
||||
- `deletes`: only when an observation is directly superseded or contradicted by new facts.
|
||||
- `reason`: REQUIRED on every create/update/delete — one sentence explaining the choice. For a CREATE, state which existing observation(s) you considered and why none matched (a near-identical existing observation means you should UPDATE, not CREATE). This is audited to catch duplicate creates.
|
||||
- Do NOT include `tags` — handled automatically.
|
||||
- Return `{{"creates": [], "updates": [], "deletes": []}}` if nothing durable is found."""
|
||||
Rules:
|
||||
- "source_fact_ids": copy the EXACT UUID strings shown in brackets [uuid] from NEW FACTS — never use integers or positions.
|
||||
- "observation_id": copy the EXACT "id" UUID string from EXISTING OBSERVATIONS.
|
||||
- One create/update may reference multiple facts when they jointly support the observation.
|
||||
- "deletes": only when an observation is directly superseded or contradicted by new facts.
|
||||
- Do NOT include "tags" — handled automatically.
|
||||
- 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.
|
||||
The mission defines *what* to track (customisable per bank).
|
||||
Processing rules and output format are always present regardless of mission.
|
||||
"""
|
||||
mission = escape_for_prompt(observations_mission or _DEFAULT_MISSION)
|
||||
mission = 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)}"
|
||||
capacity_section = f"\n\n## CAPACITY CONSTRAINT\n{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)
|
||||
"You are a memory consolidation system. Synthesize facts into observations "
|
||||
"and merge with existing observations when appropriate.\n\n"
|
||||
f"## MISSION\n{mission}{capacity_section}\n\n"
|
||||
f"{_PROCESSING_RULES}" + _BATCH_DATA_SECTION + _BATCH_OUTPUT_FORMAT
|
||||
)
|
||||
|
||||
|
||||
def build_consolidation_system_prompt(
|
||||
llm_output_language: str | None = None,
|
||||
) -> str:
|
||||
"""Bank-agnostic, cacheable system instruction for batch consolidation.
|
||||
|
||||
Holds only what is constant across banks: processing rules, input format,
|
||||
decision guide, and output format. The bank's MISSION is deliberately NOT
|
||||
here — baking it in would make the prefix bank-specific and force a separate
|
||||
Gemini context cache per mission. The mission, the per-batch INPUT, and any
|
||||
capacity constraint all ride in the user message (see
|
||||
:func:`build_consolidation_input`), so this prefix is identical for every
|
||||
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.
|
||||
"""
|
||||
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"{_PROCESSING_RULES}\n\n"
|
||||
f"{_INPUT_FORMAT_NOTE}\n\n"
|
||||
f"{_DECISION_GUIDE}\n\n"
|
||||
f"{_OUTPUT_SECTION}" + output_language_directive(llm_output_language)
|
||||
)
|
||||
# No {facts_text}/{observations_text} placeholders here — the only braces are
|
||||
# the doubled {{ }} in the OUTPUT examples, which .format() unescapes.
|
||||
return template.format()
|
||||
|
||||
|
||||
def build_consolidation_input(
|
||||
facts_text: str,
|
||||
observations_text: str,
|
||||
observations_mission: str | None = None,
|
||||
observation_capacity_note: str | None = None,
|
||||
) -> str:
|
||||
"""Per-batch user message: MISSION + INPUT data + any capacity constraint.
|
||||
|
||||
The MISSION lives here (not in the cached system prefix) so the prefix stays
|
||||
bank-agnostic and one CachedContent serves every bank. The capacity note also
|
||||
lives here since it varies as observation slots fill.
|
||||
"""
|
||||
mission = escape_for_prompt(observations_mission or _DEFAULT_MISSION)
|
||||
mission_section = f"## MISSION\n\n{mission}\n\n"
|
||||
capacity_section = ""
|
||||
if observation_capacity_note:
|
||||
capacity_section = f"## CAPACITY CONSTRAINT\n\n{escape_for_prompt(observation_capacity_note)}\n\n"
|
||||
# _SPLIT_INPUT_SECTION omits the stable observation-format explanation (now in
|
||||
# the cached system prefix) — only the variable facts/observations remain.
|
||||
template = mission_section + capacity_section + _SPLIT_INPUT_SECTION
|
||||
return template.format(facts_text=facts_text, observations_text=observations_text)
|
||||
|
||||
@@ -17,7 +17,6 @@ import httpx
|
||||
|
||||
from ..config import (
|
||||
DEFAULT_LITELLM_API_BASE,
|
||||
DEFAULT_RERANKER_ALIBABA_MODEL,
|
||||
DEFAULT_RERANKER_COHERE_MODEL,
|
||||
DEFAULT_RERANKER_FLASHRANK_CACHE_DIR,
|
||||
DEFAULT_RERANKER_FLASHRANK_CPU_MEM_ARENA,
|
||||
@@ -27,21 +26,33 @@ from ..config import (
|
||||
DEFAULT_RERANKER_LITELLM_MODEL,
|
||||
DEFAULT_RERANKER_LITELLM_SDK_MODEL,
|
||||
DEFAULT_RERANKER_LOCAL_BATCH_SIZE,
|
||||
DEFAULT_RERANKER_LOCAL_FORCE_CPU,
|
||||
DEFAULT_RERANKER_LOCAL_MAX_CONCURRENT,
|
||||
DEFAULT_RERANKER_LOCAL_MODEL,
|
||||
DEFAULT_RERANKER_LOCAL_TRUST_REMOTE_CODE,
|
||||
DEFAULT_RERANKER_PROVIDER,
|
||||
DEFAULT_RERANKER_SILICONFLOW_BASE_URL,
|
||||
DEFAULT_RERANKER_SILICONFLOW_MODEL,
|
||||
DEFAULT_RERANKER_TEI_BATCH_SIZE,
|
||||
DEFAULT_RERANKER_TEI_HTTP_TIMEOUT,
|
||||
DEFAULT_RERANKER_TEI_MAX_CONCURRENT,
|
||||
DEFAULT_RERANKER_ZEROENTROPY_MODEL,
|
||||
DEFAULT_ZEROENTROPY_BASE_URL,
|
||||
ENV_RERANKER_ALIBABA_API_KEY,
|
||||
ENV_RERANKER_COHERE_API_KEY,
|
||||
ENV_RERANKER_COHERE_MODEL,
|
||||
ENV_RERANKER_FLASHRANK_CACHE_DIR,
|
||||
ENV_RERANKER_FLASHRANK_CPU_MEM_ARENA,
|
||||
ENV_RERANKER_FLASHRANK_MODEL,
|
||||
ENV_RERANKER_GOOGLE_PROJECT_ID,
|
||||
ENV_RERANKER_LITELLM_SDK_API_KEY,
|
||||
ENV_RERANKER_LOCAL_FORCE_CPU,
|
||||
ENV_RERANKER_LOCAL_MAX_CONCURRENT,
|
||||
ENV_RERANKER_LOCAL_MODEL,
|
||||
ENV_RERANKER_LOCAL_TRUST_REMOTE_CODE,
|
||||
ENV_RERANKER_PROVIDER,
|
||||
ENV_RERANKER_SILICONFLOW_API_KEY,
|
||||
ENV_RERANKER_TEI_BATCH_SIZE,
|
||||
ENV_RERANKER_TEI_HTTP_TIMEOUT,
|
||||
ENV_RERANKER_TEI_MAX_CONCURRENT,
|
||||
ENV_RERANKER_TEI_URL,
|
||||
ENV_RERANKER_ZEROENTROPY_API_KEY,
|
||||
)
|
||||
@@ -49,43 +60,6 @@ from ..config import (
|
||||
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.
|
||||
@@ -212,7 +186,7 @@ class LocalSTCrossEncoder(CrossEncoderModel):
|
||||
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
|
||||
# Check for GPU (CUDA) or Apple Silicon (MPS)
|
||||
# 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
|
||||
@@ -220,13 +194,10 @@ class LocalSTCrossEncoder(CrossEncoderModel):
|
||||
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
|
||||
device = None # Let sentence-transformers auto-detect GPU/MPS
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to detect GPU/MPS/XPU, falling back to CPU: {e}")
|
||||
logger.warning(f"Failed to detect GPU/MPS, 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
|
||||
@@ -293,29 +264,27 @@ class LocalSTCrossEncoder(CrossEncoderModel):
|
||||
- bucket_batching: sort pairs by token length to reduce padding waste (36-54% speedup)
|
||||
- batch_size: explicit batch size for predict() calls (MPS optimal: 32)
|
||||
"""
|
||||
import numpy as np
|
||||
|
||||
try:
|
||||
if self.bucket_batching and len(pairs) > 1:
|
||||
# Sort pairs by approximate token length to create homogeneous batches.
|
||||
# This eliminates padding waste — short pairs aren't padded to the length
|
||||
# of the longest pair in the batch. Quality-identical by construction.
|
||||
lengths = [len(pairs[i][0]) + len(pairs[i][1]) for i in range(len(pairs))]
|
||||
sorted_indices = sorted(range(len(pairs)), key=lambda i: lengths[i])
|
||||
sorted_pairs = [pairs[i] for i in sorted_indices]
|
||||
if self.bucket_batching and len(pairs) > 1:
|
||||
# Sort pairs by approximate token length to create homogeneous batches.
|
||||
# This eliminates padding waste — short pairs aren't padded to the length
|
||||
# of the longest pair in the batch. Quality-identical by construction.
|
||||
lengths = [len(pairs[i][0]) + len(pairs[i][1]) for i in range(len(pairs))]
|
||||
sorted_indices = sorted(range(len(pairs)), key=lambda i: lengths[i])
|
||||
sorted_pairs = [pairs[i] for i in sorted_indices]
|
||||
|
||||
sorted_scores = self._model.predict(sorted_pairs, batch_size=self.batch_size, show_progress_bar=False)
|
||||
sorted_scores = sorted_scores.tolist() if hasattr(sorted_scores, "tolist") else list(sorted_scores)
|
||||
sorted_scores = self._model.predict(sorted_pairs, batch_size=self.batch_size, show_progress_bar=False)
|
||||
sorted_scores = sorted_scores.tolist() if hasattr(sorted_scores, "tolist") else list(sorted_scores)
|
||||
|
||||
# Restore original order
|
||||
scores = [0.0] * len(pairs)
|
||||
for new_pos, orig_idx in enumerate(sorted_indices):
|
||||
scores[orig_idx] = sorted_scores[new_pos]
|
||||
return scores
|
||||
# Restore original order
|
||||
scores = [0.0] * len(pairs)
|
||||
for new_pos, orig_idx in enumerate(sorted_indices):
|
||||
scores[orig_idx] = sorted_scores[new_pos]
|
||||
return scores
|
||||
|
||||
scores = self._model.predict(pairs, batch_size=self.batch_size, show_progress_bar=False)
|
||||
return scores.tolist() if hasattr(scores, "tolist") else list(scores)
|
||||
finally:
|
||||
_malloc_trim()
|
||||
scores = self._model.predict(pairs, batch_size=self.batch_size, show_progress_bar=False)
|
||||
return scores.tolist() if hasattr(scores, "tolist") else list(scores)
|
||||
|
||||
async def predict(self, pairs: list[tuple[str, str]]) -> list[float]:
|
||||
"""
|
||||
@@ -577,14 +546,12 @@ class _CohereCompatibleRerankClient:
|
||||
rerank_url: str,
|
||||
timeout: float = 60.0,
|
||||
include_top_n: bool = True,
|
||||
include_return_documents: bool = False,
|
||||
):
|
||||
self.api_key = api_key
|
||||
self.model = model
|
||||
self.rerank_url = rerank_url
|
||||
self.timeout = timeout
|
||||
self.include_top_n = include_top_n
|
||||
self.include_return_documents = include_return_documents
|
||||
self._async_client: httpx.AsyncClient | None = None
|
||||
|
||||
async def initialize(self) -> None:
|
||||
@@ -762,7 +729,7 @@ class ZeroEntropyCrossEncoder(CrossEncoderModel):
|
||||
See: https://docs.zeroentropy.dev/models
|
||||
"""
|
||||
|
||||
DEFAULT_BASE_URL = DEFAULT_ZEROENTROPY_BASE_URL
|
||||
DEFAULT_BASE_URL = "https://api.zeroentropy.dev"
|
||||
RERANK_PATH = "/v1/models/rerank"
|
||||
|
||||
def __init__(
|
||||
@@ -995,35 +962,32 @@ class FlashRankCrossEncoder(CrossEncoderModel):
|
||||
if not pairs:
|
||||
return []
|
||||
|
||||
try:
|
||||
# Group pairs by query
|
||||
query_groups: dict[str, list[tuple[int, str]]] = {}
|
||||
for idx, (query, text) in enumerate(pairs):
|
||||
if query not in query_groups:
|
||||
query_groups[query] = []
|
||||
query_groups[query].append((idx, text))
|
||||
# Group pairs by query
|
||||
query_groups: dict[str, list[tuple[int, str]]] = {}
|
||||
for idx, (query, text) in enumerate(pairs):
|
||||
if query not in query_groups:
|
||||
query_groups[query] = []
|
||||
query_groups[query].append((idx, text))
|
||||
|
||||
all_scores = [0.0] * len(pairs)
|
||||
all_scores = [0.0] * len(pairs)
|
||||
|
||||
for query, indexed_texts in query_groups.items():
|
||||
# Build passages list for FlashRank
|
||||
passages = [{"id": i, "text": text} for i, (_, text) in enumerate(indexed_texts)]
|
||||
global_indices = [idx for idx, _ in indexed_texts]
|
||||
for query, indexed_texts in query_groups.items():
|
||||
# Build passages list for FlashRank
|
||||
passages = [{"id": i, "text": text} for i, (_, text) in enumerate(indexed_texts)]
|
||||
global_indices = [idx for idx, _ in indexed_texts]
|
||||
|
||||
# Create rerank request
|
||||
request = RerankRequest(query=query, passages=passages)
|
||||
results = self._ranker.rerank(request)
|
||||
# Create rerank request
|
||||
request = RerankRequest(query=query, passages=passages)
|
||||
results = self._ranker.rerank(request)
|
||||
|
||||
# Map scores back to original positions
|
||||
for result in results:
|
||||
local_idx = result["id"]
|
||||
score = result["score"]
|
||||
global_idx = global_indices[local_idx]
|
||||
all_scores[global_idx] = score
|
||||
# Map scores back to original positions
|
||||
for result in results:
|
||||
local_idx = result["id"]
|
||||
score = result["score"]
|
||||
global_idx = global_indices[local_idx]
|
||||
all_scores[global_idx] = score
|
||||
|
||||
return all_scores
|
||||
finally:
|
||||
_malloc_trim()
|
||||
return all_scores
|
||||
|
||||
async def predict(self, pairs: list[tuple[str, str]]) -> list[float]:
|
||||
"""
|
||||
@@ -1187,7 +1151,7 @@ class LiteLLMSDKCrossEncoder(CrossEncoderModel):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
api_key: str | None = None,
|
||||
api_key: str,
|
||||
model: str = DEFAULT_RERANKER_LITELLM_SDK_MODEL,
|
||||
api_base: str | None = None,
|
||||
timeout: float = 60.0,
|
||||
@@ -1197,8 +1161,7 @@ class LiteLLMSDKCrossEncoder(CrossEncoderModel):
|
||||
Initialize LiteLLM SDK cross-encoder client.
|
||||
|
||||
Args:
|
||||
api_key: API key for the reranking provider (optional — omit for
|
||||
providers that use ambient credentials, e.g. AWS Bedrock with IAM)
|
||||
api_key: API key for the reranking provider
|
||||
model: Model name with provider prefix (e.g., "deepinfra/Qwen3-reranker-8B")
|
||||
api_base: Custom base URL for API (optional)
|
||||
timeout: Request timeout in seconds (default: 60.0)
|
||||
@@ -1273,9 +1236,8 @@ class LiteLLMSDKCrossEncoder(CrossEncoderModel):
|
||||
"model": self.model,
|
||||
"query": query,
|
||||
"documents": texts,
|
||||
"api_key": self.api_key,
|
||||
}
|
||||
if self.api_key:
|
||||
rerank_kwargs["api_key"] = self.api_key
|
||||
if self.api_base:
|
||||
rerank_kwargs["api_base"] = self.api_base
|
||||
|
||||
@@ -1572,48 +1534,6 @@ class GoogleCrossEncoder(CrossEncoderModel):
|
||||
return await loop.run_in_executor(None, self._predict_sync, pairs)
|
||||
|
||||
|
||||
class AlibabaCloudCrossEncoder(CrossEncoderModel):
|
||||
"""
|
||||
Alibaba Cloud DashScope text reranking API.
|
||||
|
||||
Uses the Cohere-compatible /reranks endpoint, which is the standard interface
|
||||
for qwen3-rerank. Authentication via HINDSIGHT_API_RERANKER_ALIBABA_API_KEY
|
||||
(or DASHSCOPE_API_KEY as a fallback).
|
||||
See: https://help.aliyun.com/zh/model-studio/text-rerank-api
|
||||
"""
|
||||
|
||||
RERANK_URL = "https://dashscope.aliyuncs.com/compatible-api/v1/reranks"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
api_key: str,
|
||||
model: str = DEFAULT_RERANKER_ALIBABA_MODEL,
|
||||
timeout: float = 60.0,
|
||||
):
|
||||
self.model = model
|
||||
self._client = _CohereCompatibleRerankClient(
|
||||
api_key=api_key,
|
||||
model=model,
|
||||
rerank_url=self.RERANK_URL,
|
||||
timeout=timeout,
|
||||
include_return_documents=False,
|
||||
)
|
||||
|
||||
@property
|
||||
def provider_name(self) -> str:
|
||||
return "alibaba"
|
||||
|
||||
async def initialize(self) -> None:
|
||||
if self._client._async_client is not None:
|
||||
return
|
||||
logger.info(f"Reranker: initializing Alibaba Cloud provider with model {self.model}")
|
||||
await self._client.initialize()
|
||||
logger.info("Reranker: Alibaba Cloud provider initialized")
|
||||
|
||||
async def predict(self, pairs: list[tuple[str, str]]) -> list[float]:
|
||||
return await self._client.predict(pairs)
|
||||
|
||||
|
||||
def create_cross_encoder_from_env() -> CrossEncoderModel:
|
||||
"""
|
||||
Create a CrossEncoderModel instance based on configuration.
|
||||
@@ -1656,7 +1576,6 @@ def create_cross_encoder_from_env() -> CrossEncoderModel:
|
||||
api_key=api_key,
|
||||
model=config.reranker_cohere_model,
|
||||
base_url=config.reranker_cohere_base_url,
|
||||
timeout=config.reranker_cohere_timeout,
|
||||
)
|
||||
elif provider == "openrouter":
|
||||
api_key = config.reranker_openrouter_api_key
|
||||
@@ -1668,8 +1587,7 @@ def create_cross_encoder_from_env() -> CrossEncoderModel:
|
||||
return CohereCrossEncoder(
|
||||
api_key=api_key,
|
||||
model=config.reranker_openrouter_model,
|
||||
base_url=config.reranker_openrouter_base_url,
|
||||
timeout=config.reranker_openrouter_timeout,
|
||||
base_url="https://openrouter.ai/api/v1/rerank",
|
||||
)
|
||||
elif provider == "flashrank":
|
||||
model = os.environ.get(ENV_RERANKER_FLASHRANK_MODEL, DEFAULT_RERANKER_FLASHRANK_MODEL)
|
||||
@@ -1684,15 +1602,18 @@ def create_cross_encoder_from_env() -> CrossEncoderModel:
|
||||
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":
|
||||
api_key = config.reranker_litellm_sdk_api_key
|
||||
if not api_key:
|
||||
raise ValueError(
|
||||
f"{ENV_RERANKER_LITELLM_SDK_API_KEY} is required when {ENV_RERANKER_PROVIDER} is 'litellm-sdk'"
|
||||
)
|
||||
return LiteLLMSDKCrossEncoder(
|
||||
api_key=config.reranker_litellm_sdk_api_key or None,
|
||||
api_key=api_key,
|
||||
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 = config.reranker_zeroentropy_api_key
|
||||
@@ -1703,8 +1624,6 @@ def create_cross_encoder_from_env() -> CrossEncoderModel:
|
||||
return ZeroEntropyCrossEncoder(
|
||||
api_key=api_key,
|
||||
model=config.reranker_zeroentropy_model,
|
||||
base_url=config.reranker_zeroentropy_base_url,
|
||||
timeout=config.reranker_zeroentropy_timeout,
|
||||
)
|
||||
elif provider == "siliconflow":
|
||||
api_key = config.reranker_siliconflow_api_key
|
||||
@@ -1716,7 +1635,6 @@ def create_cross_encoder_from_env() -> CrossEncoderModel:
|
||||
api_key=api_key,
|
||||
model=config.reranker_siliconflow_model,
|
||||
base_url=config.reranker_siliconflow_base_url,
|
||||
timeout=config.reranker_siliconflow_timeout,
|
||||
)
|
||||
elif provider == "google":
|
||||
project_id = config.reranker_google_project_id
|
||||
@@ -1729,16 +1647,6 @@ def create_cross_encoder_from_env() -> CrossEncoderModel:
|
||||
project_id=project_id,
|
||||
model=config.reranker_google_model,
|
||||
service_account_key=config.reranker_google_service_account_key,
|
||||
timeout=config.reranker_google_timeout,
|
||||
)
|
||||
elif provider == "alibaba":
|
||||
api_key = config.reranker_alibaba_api_key
|
||||
if not api_key:
|
||||
raise ValueError(f"{ENV_RERANKER_ALIBABA_API_KEY} is required when {ENV_RERANKER_PROVIDER} is 'alibaba'")
|
||||
return AlibabaCloudCrossEncoder(
|
||||
api_key=api_key,
|
||||
model=config.reranker_alibaba_model,
|
||||
timeout=config.reranker_alibaba_timeout,
|
||||
)
|
||||
elif provider == "rrf":
|
||||
return RRFPassthroughCrossEncoder()
|
||||
@@ -1746,5 +1654,5 @@ def create_cross_encoder_from_env() -> CrossEncoderModel:
|
||||
return JinaMLXCrossEncoder()
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unknown reranker provider: {provider}. Supported: 'local', 'tei', 'cohere', 'zeroentropy', 'siliconflow', 'alibaba', 'google', 'flashrank', 'litellm', 'litellm-sdk', 'rrf', 'jina-mlx'"
|
||||
f"Unknown reranker provider: {provider}. Supported: 'local', 'tei', 'cohere', 'zeroentropy', 'siliconflow', 'google', 'flashrank', 'litellm', 'litellm-sdk', 'rrf', 'jina-mlx'"
|
||||
)
|
||||
|
||||
@@ -19,6 +19,7 @@ and mirrors Django's ``DatabaseOperations`` architecture.
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
from uuid import UUID
|
||||
|
||||
from .base import DatabaseConnection
|
||||
from .result import ResultRow
|
||||
@@ -71,30 +72,6 @@ class DataAccessOps(ABC):
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def lock_document_for_write(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
doc_id: str,
|
||||
bank_id: str,
|
||||
) -> str | None:
|
||||
"""Ensure the document row exists, take a row lock on it, and return its
|
||||
pre-existing ``content_hash``.
|
||||
|
||||
This serializes all concurrent writers for ``doc_id`` at the DB level
|
||||
(so interleaved same-document retains can't corrupt each other), while
|
||||
creating the row on first write. The returned hash is ``'__pending__'``
|
||||
for a freshly inserted row, the stored hash for an existing one, or
|
||||
``None`` if the row could not be read back.
|
||||
|
||||
PG does this in a single statement (``INSERT ... ON CONFLICT DO UPDATE
|
||||
... RETURNING``), which always takes the row lock as part of the upsert.
|
||||
Oracle can't (``MERGE`` doesn't support ``RETURNING``), so it splits the
|
||||
work into an idempotent insert plus a ``SELECT ... FOR UPDATE``.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def insert_facts_batch(
|
||||
self,
|
||||
@@ -189,6 +166,21 @@ class DataAccessOps(ABC):
|
||||
|
||||
# -- LATERAL / fan-out queries ---------------------------------------
|
||||
|
||||
@abstractmethod
|
||||
async def fetch_entity_unit_fanout(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
ue_table: str,
|
||||
entity_id_list: list[UUID],
|
||||
limit_per_entity: int,
|
||||
) -> list[ResultRow]:
|
||||
"""Fetch unit_ids for a list of entities with per-entity row cap.
|
||||
|
||||
PG uses unnest + CROSS JOIN LATERAL with LIMIT.
|
||||
Non-PG queries each entity individually.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def fetch_unit_dates(
|
||||
self,
|
||||
@@ -414,74 +406,6 @@ class DataAccessOps(ABC):
|
||||
"""Insert a webhook delivery task into async_operations."""
|
||||
...
|
||||
|
||||
# -- Graph maintenance queue -----------------------------------------
|
||||
|
||||
@abstractmethod
|
||||
async def enqueue_graph_maintenance(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
bank_id: str,
|
||||
unit_ids: list,
|
||||
) -> None:
|
||||
"""Insert unit_ids into graph_maintenance_queue, deduplicating on the
|
||||
(bank_id, unit_id) primary key.
|
||||
|
||||
Called inside the triggering transaction so enqueue is atomic with
|
||||
the mutation that caused it. Order is unspecified.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def claim_graph_maintenance_batch(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
bank_id: str,
|
||||
limit: int,
|
||||
) -> list[str]:
|
||||
"""Atomically claim a batch of rows from graph_maintenance_queue and
|
||||
remove them from the table.
|
||||
|
||||
Returns the list of ``unit_id`` strings. Empty list when the queue
|
||||
for ``bank_id`` is drained.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def prune_orphan_entities(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
entities_table: str,
|
||||
ue_table: str,
|
||||
bank_id: str,
|
||||
) -> int:
|
||||
"""Delete entities in ``bank_id`` that no longer have any unit_entities
|
||||
rows referencing them. Returns the number of rows deleted.
|
||||
|
||||
FK ON DELETE CASCADE on entity_cooccurrences then removes any
|
||||
cooccurrence row pointing at the pruned entities.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def prune_stale_cooccurrences(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
ec_table: str,
|
||||
ue_table: str,
|
||||
entities_table: str,
|
||||
bank_id: str,
|
||||
) -> int:
|
||||
"""Delete entity_cooccurrences rows in ``bank_id`` where the two
|
||||
entities still exist but no current unit references both of them.
|
||||
|
||||
These are stale-count rows: cooccurrence was real at the time it was
|
||||
recorded, but every memory_unit that witnessed both entities has
|
||||
since been deleted. Returns the number of rows deleted.
|
||||
"""
|
||||
...
|
||||
|
||||
# -- Task claiming operations ------------------------------------------
|
||||
|
||||
@abstractmethod
|
||||
@@ -492,8 +416,6 @@ class DataAccessOps(ABC):
|
||||
worker_id: str,
|
||||
reserved_limits: dict[str, int],
|
||||
shared_limit: int,
|
||||
*,
|
||||
consolidation_bank_priority: dict[str, int] | None = None,
|
||||
) -> list[ResultRow]:
|
||||
"""Claim pending tasks from the async_operations table.
|
||||
|
||||
@@ -501,14 +423,6 @@ class DataAccessOps(ABC):
|
||||
Oracle implementation uses two-step claims (query busy banks first, then
|
||||
claim excluding them) to avoid ORA-02014.
|
||||
|
||||
Args:
|
||||
consolidation_bank_priority: Per-bank priority for consolidation scheduling.
|
||||
Maps bank name patterns to integer priorities (higher = claimed first).
|
||||
Patterns support ``*`` as wildcard (converted to SQL ``%`` for LIKE).
|
||||
A bare ``*`` key is the catch-all default for unlisted banks.
|
||||
When set, consolidation tasks are claimed in priority tiers.
|
||||
None preserves current behavior (pure created_at ordering).
|
||||
|
||||
Returns claimed rows with operation_id, operation_type, task_payload, retry_count.
|
||||
The caller is responsible for building ClaimedTask objects.
|
||||
"""
|
||||
|
||||
@@ -8,6 +8,8 @@ columns can't appear in GROUP BY).
|
||||
import json
|
||||
import uuid as uuid_mod
|
||||
from datetime import UTC, datetime
|
||||
from typing import Any
|
||||
from uuid import UUID
|
||||
|
||||
from .base import DatabaseConnection
|
||||
from .ops import DataAccessOps, TagListingParts
|
||||
@@ -45,37 +47,6 @@ class OracleOps(DataAccessOps):
|
||||
column_types=["text[]", "text[]", "text[]", "text[]", "integer[]", "text[]"],
|
||||
)
|
||||
|
||||
async def lock_document_for_write(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
doc_id: str,
|
||||
bank_id: str,
|
||||
) -> str | None:
|
||||
# Oracle can't express the PG "INSERT ... ON CONFLICT DO UPDATE ...
|
||||
# RETURNING" upsert in one statement — MERGE doesn't support RETURNING,
|
||||
# so the single-statement form rewrites to a MERGE that returns no rows
|
||||
# (DPY-1003). Split it into two statements instead:
|
||||
# 1. Idempotent insert that silently skips an existing row. The
|
||||
# IGNORE_ROW_ON_DUPKEY_INDEX hint suppresses ORA-00001 server-side;
|
||||
# a concurrent uncommitted insert of the same key blocks here until
|
||||
# the other writer commits, so writers still serialize.
|
||||
# 2. SELECT ... FOR UPDATE to take the row lock and read the hash
|
||||
# ('__pending__' for a row we just inserted, the stored hash for an
|
||||
# existing one).
|
||||
await conn.execute(
|
||||
f"INSERT /*+ IGNORE_ROW_ON_DUPKEY_INDEX({table}, pk_documents) */ "
|
||||
f"INTO {table} (id, bank_id, original_text, content_hash) "
|
||||
f"VALUES ($1, $2, '', '__pending__')",
|
||||
doc_id,
|
||||
bank_id,
|
||||
)
|
||||
return await conn.fetchval(
|
||||
f"SELECT content_hash FROM {table} WHERE id = $1 AND bank_id = $2 FOR UPDATE",
|
||||
doc_id,
|
||||
bank_id,
|
||||
)
|
||||
|
||||
async def insert_facts_batch(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
@@ -244,104 +215,29 @@ class OracleOps(DataAccessOps):
|
||||
list(zip(unit_ids, entity_ids)),
|
||||
)
|
||||
|
||||
async def enqueue_graph_maintenance(
|
||||
async def fetch_entity_unit_fanout(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
bank_id: str,
|
||||
unit_ids: list,
|
||||
) -> None:
|
||||
if not unit_ids:
|
||||
return
|
||||
# Oracle doesn't support ON CONFLICT; rely on the PK and the
|
||||
# IGNORE_ROW_ON_DUPKEY_INDEX hint to skip duplicates server-side.
|
||||
# The hint name must match the PK constraint exactly.
|
||||
#
|
||||
# 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"""
|
||||
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],
|
||||
)
|
||||
|
||||
async def claim_graph_maintenance_batch(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
bank_id: str,
|
||||
limit: int,
|
||||
) -> list[str]:
|
||||
# Two-step claim: select the batch, then delete by exact keys. Oracle's
|
||||
# DELETE ... RETURNING doesn't accept a multi-row subquery, so we can't
|
||||
# do it in one statement like the PG version.
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT unit_id FROM {table}
|
||||
WHERE bank_id = $1
|
||||
ORDER BY enqueued_at
|
||||
FETCH FIRST $2 ROWS ONLY
|
||||
""",
|
||||
bank_id,
|
||||
limit,
|
||||
)
|
||||
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",
|
||||
[(bank_id, uid) for uid in claimed],
|
||||
ue_table: str,
|
||||
entity_id_list: list[UUID],
|
||||
limit_per_entity: int,
|
||||
) -> list[ResultRow]:
|
||||
# Query each entity individually
|
||||
rows: list[ResultRow] = []
|
||||
for eid in entity_id_list:
|
||||
entity_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT $1 AS entity_id, ue.unit_id
|
||||
FROM {ue_table} ue
|
||||
WHERE ue.entity_id = $1
|
||||
ORDER BY ue.unit_id DESC
|
||||
LIMIT $2
|
||||
""",
|
||||
eid,
|
||||
limit_per_entity,
|
||||
)
|
||||
return claimed
|
||||
|
||||
async def prune_orphan_entities(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
entities_table: str,
|
||||
ue_table: str,
|
||||
bank_id: str,
|
||||
) -> int:
|
||||
# The Oracle DatabaseConnection wrapper reshapes ``cursor.rowcount`` into
|
||||
# the same ``"DELETE N"`` status string asyncpg returns, so the same
|
||||
# ``int(deleted.split()[-1])`` parsing works on both dialects.
|
||||
deleted = await conn.execute(
|
||||
f"""
|
||||
DELETE FROM {entities_table}
|
||||
WHERE bank_id = $1
|
||||
AND id NOT IN (SELECT DISTINCT entity_id FROM {ue_table})
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
return int(deleted.split()[-1]) if isinstance(deleted, str) and deleted.startswith("DELETE") else 0
|
||||
|
||||
async def prune_stale_cooccurrences(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
ec_table: str,
|
||||
ue_table: str,
|
||||
entities_table: str,
|
||||
bank_id: str,
|
||||
) -> int:
|
||||
deleted = await conn.execute(
|
||||
f"""
|
||||
DELETE FROM {ec_table}
|
||||
WHERE entity_id_1 IN (SELECT id FROM {entities_table} WHERE bank_id = $1)
|
||||
AND (entity_id_1, entity_id_2) NOT IN (
|
||||
SELECT u1.entity_id, u2.entity_id
|
||||
FROM {ue_table} u1
|
||||
JOIN {ue_table} u2 ON u1.unit_id = u2.unit_id
|
||||
)
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
return int(deleted.split()[-1]) if isinstance(deleted, str) and deleted.startswith("DELETE") else 0
|
||||
rows.extend(entity_rows)
|
||||
return rows
|
||||
|
||||
async def fetch_unit_dates(
|
||||
self,
|
||||
@@ -824,257 +720,7 @@ class OracleOps(DataAccessOps):
|
||||
|
||||
# -- Task claiming operations ------------------------------------------
|
||||
|
||||
async def _claim_consolidation_tasks(
|
||||
self,
|
||||
conn,
|
||||
table: str,
|
||||
busy_bank_ids: list[str],
|
||||
claimed_ids: list,
|
||||
limit: int,
|
||||
priority_map: dict[str, int] | None,
|
||||
) -> list:
|
||||
"""Claim consolidation tasks with optional priority-based tiered ordering.
|
||||
|
||||
Mirrors the PostgreSQL implementation. The Oracle SQL adapter
|
||||
translates ``LIKE ANY`` / ``NOT LIKE ALL`` via ``_expand_any_lists``.
|
||||
"""
|
||||
if limit <= 0:
|
||||
return []
|
||||
|
||||
if not priority_map:
|
||||
return await self._claim_consolidation_plain(conn, table, busy_bank_ids, claimed_ids, limit)
|
||||
|
||||
# --- Tiered claiming (same algorithm as PG) ---
|
||||
specific_by_priority: dict[int, list[str]] = {}
|
||||
all_specific_sql: list[str] = []
|
||||
catch_all_priority = 1
|
||||
|
||||
for pattern, priority in priority_map.items():
|
||||
if pattern == "*":
|
||||
catch_all_priority = priority
|
||||
else:
|
||||
sql_pat = pattern.replace("*", "%")
|
||||
specific_by_priority.setdefault(priority, []).append(sql_pat)
|
||||
all_specific_sql.append(sql_pat)
|
||||
|
||||
all_priorities = sorted(set(specific_by_priority.keys()) | {catch_all_priority}, reverse=True)
|
||||
|
||||
remaining = limit
|
||||
result: list = []
|
||||
|
||||
for pri in all_priorities:
|
||||
if remaining <= 0:
|
||||
break
|
||||
|
||||
if pri in specific_by_priority:
|
||||
rows = await self._claim_consolidation_like(
|
||||
conn,
|
||||
table,
|
||||
busy_bank_ids,
|
||||
claimed_ids,
|
||||
remaining,
|
||||
specific_by_priority[pri],
|
||||
)
|
||||
for row in rows:
|
||||
claimed_ids.append(row["operation_id"])
|
||||
result.append(row)
|
||||
remaining -= len(rows)
|
||||
|
||||
if pri == catch_all_priority and remaining > 0:
|
||||
rows = await self._claim_consolidation_not_like(
|
||||
conn,
|
||||
table,
|
||||
busy_bank_ids,
|
||||
claimed_ids,
|
||||
remaining,
|
||||
all_specific_sql,
|
||||
)
|
||||
for row in rows:
|
||||
claimed_ids.append(row["operation_id"])
|
||||
result.append(row)
|
||||
remaining -= len(rows)
|
||||
|
||||
return result
|
||||
|
||||
async def _claim_consolidation_plain(
|
||||
self,
|
||||
conn,
|
||||
table,
|
||||
busy_bank_ids,
|
||||
claimed_ids,
|
||||
limit,
|
||||
) -> list:
|
||||
"""Claim consolidation tasks with default created_at ordering."""
|
||||
exclude_ids = claimed_ids if claimed_ids else None
|
||||
if busy_bank_ids:
|
||||
if exclude_ids:
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND bank_id != ALL($1::text[])
|
||||
AND operation_id != ALL($2::uuid[])
|
||||
ORDER BY created_at
|
||||
LIMIT $3
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
busy_bank_ids,
|
||||
exclude_ids,
|
||||
limit,
|
||||
)
|
||||
else:
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND bank_id != ALL($1::text[])
|
||||
ORDER BY created_at
|
||||
LIMIT $2
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
busy_bank_ids,
|
||||
limit,
|
||||
)
|
||||
else:
|
||||
if exclude_ids:
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND operation_id != ALL($1::uuid[])
|
||||
ORDER BY created_at
|
||||
LIMIT $2
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
exclude_ids,
|
||||
limit,
|
||||
)
|
||||
else:
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
ORDER BY created_at
|
||||
LIMIT $1
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
limit,
|
||||
)
|
||||
|
||||
async def _claim_consolidation_like(
|
||||
self,
|
||||
conn,
|
||||
table,
|
||||
busy_bank_ids,
|
||||
claimed_ids,
|
||||
limit,
|
||||
sql_patterns,
|
||||
) -> list:
|
||||
"""Claim consolidation tasks from banks matching LIKE patterns."""
|
||||
params: list = [sql_patterns]
|
||||
conditions = ["bank_id LIKE ANY($1::text[])"]
|
||||
idx = 2
|
||||
|
||||
if busy_bank_ids:
|
||||
conditions.append(f"bank_id != ALL(${idx}::text[])")
|
||||
params.append(busy_bank_ids)
|
||||
idx += 1
|
||||
|
||||
if claimed_ids:
|
||||
conditions.append(f"operation_id != ALL(${idx}::uuid[])")
|
||||
params.append(claimed_ids)
|
||||
idx += 1
|
||||
|
||||
params.append(limit)
|
||||
extra = " AND ".join(conditions)
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND {extra}
|
||||
ORDER BY created_at
|
||||
LIMIT ${idx}
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
*params,
|
||||
)
|
||||
|
||||
async def _claim_consolidation_not_like(
|
||||
self,
|
||||
conn,
|
||||
table,
|
||||
busy_bank_ids,
|
||||
claimed_ids,
|
||||
limit,
|
||||
exclude_patterns,
|
||||
) -> list:
|
||||
"""Claim consolidation tasks from banks NOT matching any specific pattern (catch-all tier)."""
|
||||
params: list = []
|
||||
conditions: list[str] = []
|
||||
idx = 1
|
||||
|
||||
if exclude_patterns:
|
||||
conditions.append(f"bank_id NOT LIKE ALL(${idx}::text[])")
|
||||
params.append(exclude_patterns)
|
||||
idx += 1
|
||||
|
||||
if busy_bank_ids:
|
||||
conditions.append(f"bank_id != ALL(${idx}::text[])")
|
||||
params.append(busy_bank_ids)
|
||||
idx += 1
|
||||
|
||||
if claimed_ids:
|
||||
conditions.append(f"operation_id != ALL(${idx}::uuid[])")
|
||||
params.append(claimed_ids)
|
||||
idx += 1
|
||||
|
||||
params.append(limit)
|
||||
extra_clause = (" AND " + " AND ".join(conditions)) if conditions else ""
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW()){extra_clause}
|
||||
ORDER BY created_at
|
||||
LIMIT ${idx}
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
*params,
|
||||
)
|
||||
|
||||
async def claim_tasks(
|
||||
self,
|
||||
conn,
|
||||
table,
|
||||
worker_id,
|
||||
reserved_limits,
|
||||
shared_limit,
|
||||
*,
|
||||
consolidation_bank_priority=None,
|
||||
):
|
||||
async def claim_tasks(self, conn, table, worker_id, reserved_limits, shared_limit):
|
||||
"""Oracle two-step claiming to avoid ORA-02014 with NOT EXISTS + FOR UPDATE."""
|
||||
all_rows = []
|
||||
claimed_ids = []
|
||||
@@ -1085,6 +731,7 @@ class OracleOps(DataAccessOps):
|
||||
continue
|
||||
|
||||
if op_type == "consolidation":
|
||||
# Two-step: find busy banks first, then claim excluding them
|
||||
busy_banks = await conn.fetch(
|
||||
f"""
|
||||
SELECT DISTINCT bank_id FROM {table}
|
||||
@@ -1093,14 +740,38 @@ class OracleOps(DataAccessOps):
|
||||
)
|
||||
busy_bank_ids = [r["bank_id"] for r in busy_banks]
|
||||
|
||||
rows = await self._claim_consolidation_tasks(
|
||||
conn,
|
||||
table,
|
||||
busy_bank_ids,
|
||||
claimed_ids,
|
||||
limit,
|
||||
consolidation_bank_priority,
|
||||
)
|
||||
if busy_bank_ids:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND bank_id != ALL($1::text[])
|
||||
ORDER BY created_at
|
||||
LIMIT $2
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
busy_bank_ids,
|
||||
limit,
|
||||
)
|
||||
else:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
ORDER BY created_at
|
||||
LIMIT $1
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
limit,
|
||||
)
|
||||
else:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
@@ -1164,7 +835,7 @@ class OracleOps(DataAccessOps):
|
||||
all_rows.append(row)
|
||||
remaining_shared -= len(rows)
|
||||
|
||||
# 2b. Consolidation tasks (with bank-serialization + optional priority)
|
||||
# 2b. Consolidation tasks (with bank-serialization)
|
||||
if remaining_shared > 0:
|
||||
busy_banks_2 = await conn.fetch(
|
||||
f"""
|
||||
@@ -1174,14 +845,76 @@ class OracleOps(DataAccessOps):
|
||||
)
|
||||
busy_bank_ids_2 = [r["bank_id"] for r in busy_banks_2]
|
||||
|
||||
rows = await self._claim_consolidation_tasks(
|
||||
conn,
|
||||
table,
|
||||
busy_bank_ids_2,
|
||||
claimed_ids,
|
||||
remaining_shared,
|
||||
consolidation_bank_priority,
|
||||
)
|
||||
if claimed_ids:
|
||||
if busy_bank_ids_2:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND operation_id != ALL($1::uuid[])
|
||||
AND bank_id != ALL($2::text[])
|
||||
ORDER BY created_at
|
||||
LIMIT $3
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
claimed_ids,
|
||||
busy_bank_ids_2,
|
||||
remaining_shared,
|
||||
)
|
||||
else:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND operation_id != ALL($1::uuid[])
|
||||
ORDER BY created_at
|
||||
LIMIT $2
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
claimed_ids,
|
||||
remaining_shared,
|
||||
)
|
||||
else:
|
||||
if busy_bank_ids_2:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND bank_id != ALL($1::text[])
|
||||
ORDER BY created_at
|
||||
LIMIT $2
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
busy_bank_ids_2,
|
||||
remaining_shared,
|
||||
)
|
||||
else:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
ORDER BY created_at
|
||||
LIMIT $1
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
remaining_shared,
|
||||
)
|
||||
|
||||
for row in rows:
|
||||
claimed_ids.append(row["operation_id"])
|
||||
|
||||
@@ -4,6 +4,11 @@ Uses unnest(), LATERAL, DISTINCT ON, and native array operations for
|
||||
efficient batch operations.
|
||||
"""
|
||||
|
||||
import json
|
||||
from datetime import UTC, datetime
|
||||
from typing import Any
|
||||
from uuid import UUID
|
||||
|
||||
from .base import DatabaseConnection
|
||||
from .ops import DataAccessOps, TagListingParts
|
||||
from .result import ResultRow
|
||||
@@ -44,30 +49,6 @@ class PostgreSQLOps(DataAccessOps):
|
||||
content_hashes,
|
||||
)
|
||||
|
||||
async def lock_document_for_write(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
doc_id: str,
|
||||
bank_id: str,
|
||||
) -> str | None:
|
||||
# Single upsert that both creates the row (if absent) and locks it (if
|
||||
# present) atomically. ON CONFLICT DO UPDATE always takes the row lock as
|
||||
# part of the statement, so all concurrent same-document writers serialize
|
||||
# on the document row in one consistent step (the earlier two-step form —
|
||||
# DO NOTHING + a separate SELECT FOR UPDATE — could deadlock because
|
||||
# DO NOTHING takes no lock on an existing row). The SET is a no-op
|
||||
# self-assignment used only to acquire the lock; RETURNING yields the
|
||||
# pre-existing hash (or '__pending__' for a freshly inserted row).
|
||||
return await conn.fetchval(
|
||||
f"INSERT INTO {table} (id, bank_id, original_text, content_hash) "
|
||||
f"VALUES ($1, $2, '', '__pending__') "
|
||||
f"ON CONFLICT (id, bank_id) DO UPDATE SET content_hash = {table}.content_hash "
|
||||
f"RETURNING content_hash",
|
||||
doc_id,
|
||||
bank_id,
|
||||
)
|
||||
|
||||
async def insert_facts_batch(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
@@ -123,46 +104,7 @@ class PostgreSQLOps(DataAccessOps):
|
||||
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(
|
||||
@@ -219,23 +161,6 @@ class PostgreSQLOps(DataAccessOps):
|
||||
exists_clause: str,
|
||||
chunk_size: int = 5000,
|
||||
) -> None:
|
||||
# exists_clause is unused on PostgreSQL: the memory_links → memory_units
|
||||
# FKs are DEFERRABLE INITIALLY DEFERRED, so an INSERT takes no lock on the
|
||||
# referenced parent rows until COMMIT — a concurrent committed DELETE in
|
||||
# that window (consolidation pruning observations, document re-tracking)
|
||||
# trips fk_memory_links_{to,from}_unit_id_memory_units at COMMIT (#1882),
|
||||
# and a WHERE EXISTS guard can't prevent it (the row passes the check,
|
||||
# then is deleted before the deferred check runs). Instead a CTE locks the
|
||||
# referenced units FOR KEY SHARE in the *same statement*: the lock blocks a
|
||||
# concurrent DELETE until our transaction commits and is held through the
|
||||
# deferred check, and the INSERT only takes links whose endpoints are in
|
||||
# the locked set, so rows that already vanished are dropped. Folding it
|
||||
# into the one INSERT keeps this to a single round-trip — no extra query
|
||||
# and no surrounding transaction needed. (Oracle's immediate FK has no
|
||||
# such window and uses exists_clause via its own bulk_insert_links.)
|
||||
from ..schema import fq_table
|
||||
|
||||
mu_table = fq_table("memory_units")
|
||||
from_ids = [lnk[0] for lnk in sorted_links]
|
||||
to_ids = [lnk[1] for lnk in sorted_links]
|
||||
types = [lnk[2] for lnk in sorted_links]
|
||||
@@ -244,37 +169,24 @@ class PostgreSQLOps(DataAccessOps):
|
||||
|
||||
for chunk_start in range(0, len(sorted_links), chunk_size):
|
||||
chunk_end = min(chunk_start + chunk_size, len(sorted_links))
|
||||
chunk_from = from_ids[chunk_start:chunk_end]
|
||||
chunk_to = to_ids[chunk_start:chunk_end]
|
||||
# Distinct referenced parents, sorted so concurrent inserters acquire
|
||||
# the row-share locks in a consistent order (avoids deadlocks; same
|
||||
# convention as the (from, to) link sort).
|
||||
referenced = sorted({str(x) for x in chunk_from} | {str(x) for x in chunk_to})
|
||||
await conn.execute(
|
||||
f"""
|
||||
WITH locked AS (
|
||||
SELECT id FROM {mu_table}
|
||||
WHERE id = ANY($7::uuid[])
|
||||
ORDER BY id
|
||||
FOR KEY SHARE
|
||||
)
|
||||
INSERT INTO {table}
|
||||
(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)
|
||||
AS t(f, t, tp, w, e)
|
||||
{exists_clause}
|
||||
ON CONFLICT (from_unit_id, to_unit_id, link_type,
|
||||
COALESCE(entity_id, '{nil_entity_uuid}'::uuid))
|
||||
DO NOTHING
|
||||
""",
|
||||
chunk_from,
|
||||
chunk_to,
|
||||
from_ids[chunk_start:chunk_end],
|
||||
to_ids[chunk_start:chunk_end],
|
||||
types[chunk_start:chunk_end],
|
||||
weights[chunk_start:chunk_end],
|
||||
entity_ids[chunk_start:chunk_end],
|
||||
bank_id,
|
||||
referenced,
|
||||
timeout=300,
|
||||
)
|
||||
|
||||
@@ -339,109 +251,28 @@ class PostgreSQLOps(DataAccessOps):
|
||||
entity_ids,
|
||||
)
|
||||
|
||||
async def enqueue_graph_maintenance(
|
||||
async def fetch_entity_unit_fanout(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
bank_id: str,
|
||||
unit_ids: list,
|
||||
) -> None:
|
||||
if not unit_ids:
|
||||
return
|
||||
# Sort to enforce a global lock-acquisition order on the
|
||||
# (bank_id, unit_id) unique-key. Without this, two concurrent
|
||||
# transactions inserting overlapping unit_id sets in different
|
||||
# orders can deadlock on the ON CONFLICT row locks — Postgres
|
||||
# acquires a short-lived lock per row being checked, and cycle
|
||||
# detection then aborts one transaction. 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.execute(
|
||||
ue_table: str,
|
||||
entity_id_list: list[UUID],
|
||||
limit_per_entity: int,
|
||||
) -> list[ResultRow]:
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
INSERT INTO {table} (bank_id, unit_id)
|
||||
SELECT $1, v FROM unnest($2::uuid[]) AS t(v)
|
||||
ON CONFLICT (bank_id, unit_id) DO NOTHING
|
||||
""",
|
||||
bank_id,
|
||||
sorted_unit_ids,
|
||||
)
|
||||
|
||||
async def claim_graph_maintenance_batch(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
bank_id: str,
|
||||
limit: int,
|
||||
) -> list[str]:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
DELETE FROM {table}
|
||||
WHERE (bank_id, unit_id) IN (
|
||||
SELECT bank_id, unit_id FROM {table}
|
||||
WHERE bank_id = $1
|
||||
ORDER BY enqueued_at
|
||||
SELECT e.entity_id, n.unit_id
|
||||
FROM unnest($1::uuid[]) AS e(entity_id)
|
||||
CROSS JOIN LATERAL (
|
||||
SELECT ue.unit_id
|
||||
FROM {ue_table} ue
|
||||
WHERE ue.entity_id = e.entity_id
|
||||
ORDER BY ue.unit_id DESC
|
||||
LIMIT $2
|
||||
)
|
||||
RETURNING unit_id
|
||||
) n
|
||||
""",
|
||||
bank_id,
|
||||
limit,
|
||||
entity_id_list,
|
||||
limit_per_entity,
|
||||
)
|
||||
return [str(row["unit_id"]) for row in rows]
|
||||
|
||||
async def prune_orphan_entities(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
entities_table: str,
|
||||
ue_table: str,
|
||||
bank_id: str,
|
||||
) -> int:
|
||||
# Scoped by entities.bank_id (indexed). The NOT EXISTS subquery is
|
||||
# backed by idx_ue_entity on unit_entities(entity_id), so this stays
|
||||
# linear in the number of entities in the bank — not in the size of
|
||||
# unit_entities globally.
|
||||
result = await conn.execute(
|
||||
f"""
|
||||
DELETE FROM {entities_table} e
|
||||
WHERE e.bank_id = $1
|
||||
AND NOT EXISTS (
|
||||
SELECT 1 FROM {ue_table} ue WHERE ue.entity_id = e.id
|
||||
)
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
# asyncpg returns "DELETE N"
|
||||
return int(result.split()[-1]) if isinstance(result, str) and result.startswith("DELETE") else 0
|
||||
|
||||
async def prune_stale_cooccurrences(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
ec_table: str,
|
||||
ue_table: str,
|
||||
entities_table: str,
|
||||
bank_id: str,
|
||||
) -> int:
|
||||
# Scope by joining through entities.bank_id (entity_cooccurrences itself
|
||||
# has no bank_id column — entities don't span banks, so scoping via
|
||||
# entity_id_1 is sufficient).
|
||||
result = await conn.execute(
|
||||
f"""
|
||||
DELETE FROM {ec_table} c
|
||||
USING {entities_table} e
|
||||
WHERE e.id = c.entity_id_1
|
||||
AND e.bank_id = $1
|
||||
AND NOT EXISTS (
|
||||
SELECT 1
|
||||
FROM {ue_table} u1
|
||||
JOIN {ue_table} u2 ON u1.unit_id = u2.unit_id
|
||||
WHERE u1.entity_id = c.entity_id_1
|
||||
AND u2.entity_id = c.entity_id_2
|
||||
)
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
return int(result.split()[-1]) if isinstance(result, str) and result.startswith("DELETE") else 0
|
||||
|
||||
async def fetch_unit_dates(
|
||||
self,
|
||||
@@ -624,6 +455,7 @@ class PostgreSQLOps(DataAccessOps):
|
||||
per_entity_limit: int,
|
||||
) -> tuple[list[ResultRow], list[ResultRow], list[ResultRow]]:
|
||||
# v0.5.6 array ops: unnest, &&, COUNT(DISTINCT) on source_memory_ids.
|
||||
from ..schema import fq_table
|
||||
|
||||
entity_rows = await conn.fetch(
|
||||
f"""
|
||||
@@ -894,268 +726,7 @@ class PostgreSQLOps(DataAccessOps):
|
||||
|
||||
# -- Task claiming operations ------------------------------------------
|
||||
|
||||
async def _claim_consolidation_tasks(
|
||||
self,
|
||||
conn,
|
||||
table: str,
|
||||
busy_bank_ids: list[str],
|
||||
claimed_ids: list,
|
||||
limit: int,
|
||||
priority_map: dict[str, int] | None,
|
||||
) -> list:
|
||||
"""Claim consolidation tasks with optional priority-based tiered ordering.
|
||||
|
||||
When *priority_map* is ``None``, uses the default ``ORDER BY created_at``
|
||||
with bank-serialization (exclude busy banks). When set, claims in
|
||||
priority tiers — highest-priority banks first. Specific patterns always
|
||||
take precedence over the catch-all ``*`` entry.
|
||||
"""
|
||||
if limit <= 0:
|
||||
return []
|
||||
|
||||
# --- Fast path: no priority map -> current behavior ---
|
||||
if not priority_map:
|
||||
return await self._claim_consolidation_plain(conn, table, busy_bank_ids, claimed_ids, limit)
|
||||
|
||||
# --- Tiered claiming ---
|
||||
# Separate specific patterns from catch-all.
|
||||
# Specific patterns always take precedence: a bank matching ``shadow-*``
|
||||
# uses that entry's priority even if the catch-all ``*`` has a higher
|
||||
# value. The catch-all only applies to banks not matching any specific
|
||||
# pattern.
|
||||
specific_by_priority: dict[int, list[str]] = {}
|
||||
all_specific_sql: list[str] = []
|
||||
catch_all_priority = 1 # default when no ``*`` entry
|
||||
|
||||
for pattern, priority in priority_map.items():
|
||||
if pattern == "*":
|
||||
catch_all_priority = priority
|
||||
else:
|
||||
sql_pat = pattern.replace("*", "%")
|
||||
specific_by_priority.setdefault(priority, []).append(sql_pat)
|
||||
all_specific_sql.append(sql_pat)
|
||||
|
||||
# Collect all priority levels (specific tiers + catch-all) sorted desc.
|
||||
all_priorities = sorted(set(specific_by_priority.keys()) | {catch_all_priority}, reverse=True)
|
||||
|
||||
remaining = limit
|
||||
result: list = []
|
||||
|
||||
for pri in all_priorities:
|
||||
if remaining <= 0:
|
||||
break
|
||||
|
||||
# Specific-pattern tier at this priority level
|
||||
if pri in specific_by_priority:
|
||||
rows = await self._claim_consolidation_like(
|
||||
conn,
|
||||
table,
|
||||
busy_bank_ids,
|
||||
claimed_ids,
|
||||
remaining,
|
||||
specific_by_priority[pri],
|
||||
)
|
||||
for row in rows:
|
||||
claimed_ids.append(row["operation_id"])
|
||||
result.append(row)
|
||||
remaining -= len(rows)
|
||||
|
||||
# Catch-all tier at this priority level
|
||||
if pri == catch_all_priority and remaining > 0:
|
||||
rows = await self._claim_consolidation_not_like(
|
||||
conn,
|
||||
table,
|
||||
busy_bank_ids,
|
||||
claimed_ids,
|
||||
remaining,
|
||||
all_specific_sql,
|
||||
)
|
||||
for row in rows:
|
||||
claimed_ids.append(row["operation_id"])
|
||||
result.append(row)
|
||||
remaining -= len(rows)
|
||||
|
||||
return result
|
||||
|
||||
async def _claim_consolidation_plain(
|
||||
self,
|
||||
conn,
|
||||
table,
|
||||
busy_bank_ids,
|
||||
claimed_ids,
|
||||
limit,
|
||||
) -> list:
|
||||
"""Claim consolidation tasks with default created_at ordering."""
|
||||
exclude_ids = claimed_ids if claimed_ids else None
|
||||
if busy_bank_ids:
|
||||
if exclude_ids:
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND bank_id != ALL($1::text[])
|
||||
AND operation_id != ALL($2::uuid[])
|
||||
ORDER BY created_at
|
||||
LIMIT $3
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
busy_bank_ids,
|
||||
exclude_ids,
|
||||
limit,
|
||||
)
|
||||
else:
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND bank_id != ALL($1::text[])
|
||||
ORDER BY created_at
|
||||
LIMIT $2
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
busy_bank_ids,
|
||||
limit,
|
||||
)
|
||||
else:
|
||||
if exclude_ids:
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND operation_id != ALL($1::uuid[])
|
||||
ORDER BY created_at
|
||||
LIMIT $2
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
exclude_ids,
|
||||
limit,
|
||||
)
|
||||
else:
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
ORDER BY created_at
|
||||
LIMIT $1
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
limit,
|
||||
)
|
||||
|
||||
async def _claim_consolidation_like(
|
||||
self,
|
||||
conn,
|
||||
table,
|
||||
busy_bank_ids,
|
||||
claimed_ids,
|
||||
limit,
|
||||
sql_patterns,
|
||||
) -> list:
|
||||
"""Claim consolidation tasks from banks matching LIKE patterns."""
|
||||
params: list = [sql_patterns]
|
||||
conditions = ["bank_id LIKE ANY($1::text[])"]
|
||||
idx = 2
|
||||
|
||||
if busy_bank_ids:
|
||||
conditions.append(f"bank_id != ALL(${idx}::text[])")
|
||||
params.append(busy_bank_ids)
|
||||
idx += 1
|
||||
|
||||
if claimed_ids:
|
||||
conditions.append(f"operation_id != ALL(${idx}::uuid[])")
|
||||
params.append(claimed_ids)
|
||||
idx += 1
|
||||
|
||||
params.append(limit)
|
||||
extra = " AND ".join(conditions)
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND {extra}
|
||||
ORDER BY created_at
|
||||
LIMIT ${idx}
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
*params,
|
||||
)
|
||||
|
||||
async def _claim_consolidation_not_like(
|
||||
self,
|
||||
conn,
|
||||
table,
|
||||
busy_bank_ids,
|
||||
claimed_ids,
|
||||
limit,
|
||||
exclude_patterns,
|
||||
) -> list:
|
||||
"""Claim consolidation tasks from banks NOT matching any specific pattern (catch-all tier)."""
|
||||
params: list = []
|
||||
conditions: list[str] = []
|
||||
idx = 1
|
||||
|
||||
if exclude_patterns:
|
||||
conditions.append(f"bank_id NOT LIKE ALL(${idx}::text[])")
|
||||
params.append(exclude_patterns)
|
||||
idx += 1
|
||||
|
||||
if busy_bank_ids:
|
||||
conditions.append(f"bank_id != ALL(${idx}::text[])")
|
||||
params.append(busy_bank_ids)
|
||||
idx += 1
|
||||
|
||||
if claimed_ids:
|
||||
conditions.append(f"operation_id != ALL(${idx}::uuid[])")
|
||||
params.append(claimed_ids)
|
||||
idx += 1
|
||||
|
||||
params.append(limit)
|
||||
extra_clause = (" AND " + " AND ".join(conditions)) if conditions else ""
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW()){extra_clause}
|
||||
ORDER BY created_at
|
||||
LIMIT ${idx}
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
*params,
|
||||
)
|
||||
|
||||
async def claim_tasks(
|
||||
self,
|
||||
conn,
|
||||
table,
|
||||
worker_id,
|
||||
reserved_limits,
|
||||
shared_limit,
|
||||
*,
|
||||
consolidation_bank_priority=None,
|
||||
):
|
||||
async def claim_tasks(self, conn, table, worker_id, reserved_limits, shared_limit):
|
||||
all_rows = []
|
||||
claimed_ids = []
|
||||
|
||||
@@ -1173,14 +744,38 @@ class PostgreSQLOps(DataAccessOps):
|
||||
)
|
||||
busy_bank_ids = [r["bank_id"] for r in busy_banks]
|
||||
|
||||
rows = await self._claim_consolidation_tasks(
|
||||
conn,
|
||||
table,
|
||||
busy_bank_ids,
|
||||
claimed_ids,
|
||||
limit,
|
||||
consolidation_bank_priority,
|
||||
)
|
||||
if busy_bank_ids:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND bank_id != ALL($1::text[])
|
||||
ORDER BY created_at
|
||||
LIMIT $2
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
busy_bank_ids,
|
||||
limit,
|
||||
)
|
||||
else:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
ORDER BY created_at
|
||||
LIMIT $1
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
limit,
|
||||
)
|
||||
else:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
@@ -1244,7 +839,7 @@ class PostgreSQLOps(DataAccessOps):
|
||||
all_rows.append(row)
|
||||
remaining_shared -= len(rows)
|
||||
|
||||
# 2b. Consolidation tasks (with bank-serialization + optional priority)
|
||||
# 2b. Consolidation tasks (with bank-serialization)
|
||||
if remaining_shared > 0:
|
||||
busy_banks_2 = await conn.fetch(
|
||||
f"""
|
||||
@@ -1254,14 +849,76 @@ class PostgreSQLOps(DataAccessOps):
|
||||
)
|
||||
busy_bank_ids_2 = [r["bank_id"] for r in busy_banks_2]
|
||||
|
||||
rows = await self._claim_consolidation_tasks(
|
||||
conn,
|
||||
table,
|
||||
busy_bank_ids_2,
|
||||
claimed_ids,
|
||||
remaining_shared,
|
||||
consolidation_bank_priority,
|
||||
)
|
||||
if claimed_ids:
|
||||
if busy_bank_ids_2:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND operation_id != ALL($1::uuid[])
|
||||
AND bank_id != ALL($2::text[])
|
||||
ORDER BY created_at
|
||||
LIMIT $3
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
claimed_ids,
|
||||
busy_bank_ids_2,
|
||||
remaining_shared,
|
||||
)
|
||||
else:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND operation_id != ALL($1::uuid[])
|
||||
ORDER BY created_at
|
||||
LIMIT $2
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
claimed_ids,
|
||||
remaining_shared,
|
||||
)
|
||||
else:
|
||||
if busy_bank_ids_2:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND bank_id != ALL($1::text[])
|
||||
ORDER BY created_at
|
||||
LIMIT $2
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
busy_bank_ids_2,
|
||||
remaining_shared,
|
||||
)
|
||||
else:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
ORDER BY created_at
|
||||
LIMIT $1
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
remaining_shared,
|
||||
)
|
||||
|
||||
for row in rows:
|
||||
claimed_ids.append(row["operation_id"])
|
||||
|
||||
@@ -106,15 +106,6 @@ SCHEMAS_WITH_PENDING_WORK = OptionalRoutine(
|
||||
deployment.
|
||||
* Should be cheap and idempotent — called every poll cycle (~30s).
|
||||
|
||||
The poller trusts the result wholesale: any schema the routine does
|
||||
not return is treated as having no work this cycle. It does NOT
|
||||
second-guess omissions with a per-schema scan — that would re-run the
|
||||
exact queries this routine exists to avoid. Consequently the routine
|
||||
is *only* appropriate for multi-tenant deployments. Single-schema
|
||||
(default/public only) installs should NOT create it: the per-schema
|
||||
fallback below is a single cheap EXISTS check that covers ``public``
|
||||
correctly and cannot starve.
|
||||
|
||||
Fallback when the routine is absent: per-schema ``EXISTS`` queries
|
||||
from Python (~4ms per schema). The server-side path is a single-
|
||||
round-trip optimisation worth ~200ms in deployments with thousands
|
||||
|
||||
@@ -14,7 +14,6 @@ Supports multi-tenant schema isolation via ALTER SESSION SET CURRENT_SCHEMA.
|
||||
"""
|
||||
|
||||
import datetime
|
||||
import inspect
|
||||
import json
|
||||
import logging
|
||||
import re
|
||||
@@ -73,9 +72,6 @@ _RETURNING_RE = re.compile(r"\bRETURNING\s+(.+)", re.IGNORECASE | re.DOTALL)
|
||||
|
||||
_ANY_RE = re.compile(r"=\s*ANY\s*\(\s*:(\d+)\s*\)", re.IGNORECASE)
|
||||
_NOT_ALL_RE = re.compile(r"!=\s*ALL\s*\(\s*:(\d+)\s*\)", re.IGNORECASE)
|
||||
# LIKE ANY / NOT LIKE ALL — capture the column name before the operator
|
||||
_LIKE_ANY_RE = re.compile(r"(\w+)\s+LIKE\s+ANY\s*\(\s*:(\d+)\s*\)", re.IGNORECASE)
|
||||
_NOT_LIKE_ALL_RE = re.compile(r"(\w+)\s+NOT\s+LIKE\s+ALL\s*\(\s*:(\d+)\s*\)", re.IGNORECASE)
|
||||
|
||||
_JSON_ARROW_TEXT_RE = re.compile(r'("?\w+"?)\s*->>\s*\'(\w+)\'') # handles both col and "col"
|
||||
_JSON_HAS_KEY_RE = re.compile(r"(\w+)\s*\?\s*'(\w+)'")
|
||||
@@ -156,23 +152,6 @@ _JSON_COL_NAMES = {
|
||||
"task_payload",
|
||||
"history",
|
||||
}
|
||||
# NOTE: the history tables' JSON payload column is named ``content`` — deliberately
|
||||
# NOT added here, because ``mental_models.content`` is plain text (adding "content"
|
||||
# would corrupt those reads). The history read paths json.loads ``content`` directly.
|
||||
|
||||
# Columns backed by CLOB in Oracle (large text or JSON). When such a column is
|
||||
# returned via a ``RETURNING`` clause it must be bound as DB_TYPE_CLOB; binding
|
||||
# it as VARCHAR raises ORA-22835 ("buffer too small for CLOB to CHAR") once the
|
||||
# value exceeds 4000 bytes. Union of the JSON-CLOB columns above and the
|
||||
# large-text CLOB columns.
|
||||
_CLOB_RETURNING_COLS = _JSON_COL_NAMES | {
|
||||
"content",
|
||||
"text",
|
||||
"context",
|
||||
"structured_content",
|
||||
"text_signals",
|
||||
"search_vector",
|
||||
}
|
||||
|
||||
|
||||
def _is_uuid_column(col: str) -> bool:
|
||||
@@ -370,10 +349,6 @@ def _rewrite_pg_to_oracle(query: str) -> RewriteResult:
|
||||
# Boolean literals: Oracle uses NUMBER(1) for booleans
|
||||
query = re.sub(r"\b=\s*TRUE\b", "= 1", query, flags=re.IGNORECASE)
|
||||
query = re.sub(r"\b=\s*FALSE\b", "= 0", query, flags=re.IGNORECASE)
|
||||
# FOR NO KEY UPDATE → FOR UPDATE (Oracle has only FOR UPDATE; it does not block
|
||||
# indexed-FK child inserts the way PG's FOR UPDATE would, so plain FOR UPDATE is
|
||||
# the correct equivalent). Must run before the FOR SHARE rule below.
|
||||
query = re.sub(r"\bFOR\s+NO\s+KEY\s+UPDATE\b", "FOR UPDATE", query, flags=re.IGNORECASE)
|
||||
# FOR SHARE → FOR UPDATE (Oracle doesn't support FOR SHARE)
|
||||
query = re.sub(r"\bFOR\s+SHARE\b", "FOR UPDATE", query, flags=re.IGNORECASE)
|
||||
|
||||
@@ -560,12 +535,6 @@ def _rewrite_pg_to_oracle(query: str) -> RewriteResult:
|
||||
# != ALL(:N) → NOT IN (expanded list) — the negative counterpart of = ANY
|
||||
query = _NOT_ALL_RE.sub(r"NOT IN (/*EXPAND:\1*/)", query)
|
||||
|
||||
# col LIKE ANY(:N) → (col LIKE :p0 OR col LIKE :p1 OR ...)
|
||||
query = _LIKE_ANY_RE.sub(r"\1 /*LIKE_ANY:\2:\1*/", query)
|
||||
|
||||
# col NOT LIKE ALL(:N) → (col NOT LIKE :p0 AND col NOT LIKE :p1 AND ...)
|
||||
query = _NOT_LIKE_ALL_RE.sub(r"\1 /*NOT_LIKE_ALL:\2:\1*/", query)
|
||||
|
||||
# CTE AS MATERIALIZED (...) → AS (...) — Oracle doesn't support MATERIALIZED CTE hint
|
||||
query = re.sub(r"\bAS\s+MATERIALIZED\s*\(", "AS (", query, flags=re.IGNORECASE)
|
||||
|
||||
@@ -703,11 +672,6 @@ class OracleConnection(DatabaseConnection):
|
||||
params[f"ret_{i}"] = cursor.var(oracledb.DB_TYPE_TIMESTAMP_TZ, arraysize=1)
|
||||
elif clean in _NUMERIC_COLS:
|
||||
params[f"ret_{i}"] = cursor.var(oracledb.DB_TYPE_NUMBER, arraysize=1)
|
||||
elif clean in _CLOB_RETURNING_COLS:
|
||||
# CLOB-backed column: a VARCHAR out-bind caps at 4000 bytes and
|
||||
# raises ORA-22835 for larger values. Read back as a LOB in
|
||||
# _read_returning_values.
|
||||
params[f"ret_{i}"] = cursor.var(oracledb.DB_TYPE_CLOB, arraysize=1)
|
||||
else:
|
||||
params[f"ret_{i}"] = cursor.var(oracledb.DB_TYPE_VARCHAR, arraysize=1)
|
||||
|
||||
@@ -760,38 +724,17 @@ class OracleConnection(DatabaseConnection):
|
||||
|
||||
_expand_counter = 0
|
||||
|
||||
@staticmethod
|
||||
def _resolve_list_param(params: dict[str, Any], key: str) -> list | None:
|
||||
"""Resolve a parameter that may be a list or a JSON-encoded list string."""
|
||||
val = params.get(key)
|
||||
if isinstance(val, str):
|
||||
try:
|
||||
parsed = json.loads(val)
|
||||
if isinstance(parsed, list):
|
||||
return parsed
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
pass
|
||||
if isinstance(val, (list, tuple)):
|
||||
return list(val)
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _expand_any_lists(query: str, params: dict[str, Any] | None) -> tuple[str, dict[str, Any] | None]:
|
||||
"""Expand /*EXPAND:N*/, /*LIKE_ANY:N:col*/, /*NOT_LIKE_ALL:N:col*/ markers.
|
||||
"""Expand /*EXPAND:N*/ markers into individual bind vars for IN clauses.
|
||||
|
||||
Converts: IN (/*EXPAND:1*/) with params["1"] = [a, b, c]
|
||||
Into: IN (:any_0, :any_1, :any_2) with params["any_0"]=a, etc.
|
||||
|
||||
Converts: col /*LIKE_ANY:1:col*/ with params["1"] = [a, b]
|
||||
Into: (col LIKE :lk_0 OR col LIKE :lk_1)
|
||||
|
||||
Converts: col /*NOT_LIKE_ALL:1:col*/ with params["1"] = [a, b]
|
||||
Into: (col NOT LIKE :nlk_0 AND col NOT LIKE :nlk_1)
|
||||
|
||||
Uses a unique prefix to avoid name collisions with other bind vars.
|
||||
The original param is kept (for other references to :N in the query).
|
||||
"""
|
||||
if params is None or "/*" not in query:
|
||||
if params is None or "/*EXPAND:" not in query:
|
||||
return query, params
|
||||
|
||||
expand_re = re.compile(r"/\*EXPAND:(\d+)\*/")
|
||||
@@ -832,50 +775,6 @@ class OracleConnection(DatabaseConnection):
|
||||
|
||||
query = expand_re.sub(_replace, query)
|
||||
|
||||
# Expand LIKE ANY: col /*LIKE_ANY:N:col*/ → (col LIKE :p0 OR col LIKE :p1 ...)
|
||||
like_any_re = re.compile(r"(\w+)\s*/\*LIKE_ANY:(\d+):(\w+)\*/")
|
||||
|
||||
def _replace_like_any(m):
|
||||
_col = m.group(1) # redundant column ref before marker
|
||||
param_key = m.group(2)
|
||||
col = m.group(3)
|
||||
val = OracleConnection._resolve_list_param(params, param_key)
|
||||
if val is None or len(val) == 0:
|
||||
return "1=0" # no patterns → no match
|
||||
OracleConnection._expand_counter += 1
|
||||
prefix = f"lk{OracleConnection._expand_counter}"
|
||||
clauses = []
|
||||
for i, item in enumerate(val):
|
||||
k = f"{prefix}_{i}"
|
||||
params[k] = item
|
||||
clauses.append(f"{col} LIKE :{k}")
|
||||
keys_to_remove.add(param_key)
|
||||
return f"({' OR '.join(clauses)})"
|
||||
|
||||
query = like_any_re.sub(_replace_like_any, query)
|
||||
|
||||
# Expand NOT LIKE ALL: col /*NOT_LIKE_ALL:N:col*/ → (col NOT LIKE :p0 AND ...)
|
||||
not_like_all_re = re.compile(r"(\w+)\s*/\*NOT_LIKE_ALL:(\d+):(\w+)\*/")
|
||||
|
||||
def _replace_not_like_all(m):
|
||||
_col = m.group(1)
|
||||
param_key = m.group(2)
|
||||
col = m.group(3)
|
||||
val = OracleConnection._resolve_list_param(params, param_key)
|
||||
if val is None or len(val) == 0:
|
||||
return "1=1" # no patterns → everything matches
|
||||
OracleConnection._expand_counter += 1
|
||||
prefix = f"nlk{OracleConnection._expand_counter}"
|
||||
clauses = []
|
||||
for i, item in enumerate(val):
|
||||
k = f"{prefix}_{i}"
|
||||
params[k] = item
|
||||
clauses.append(f"{col} NOT LIKE :{k}")
|
||||
keys_to_remove.add(param_key)
|
||||
return f"({' AND '.join(clauses)})"
|
||||
|
||||
query = not_like_all_re.sub(_replace_not_like_all, query)
|
||||
|
||||
# Remove original list params that were expanded — their placeholder
|
||||
# (:N) no longer exists in the query, and leaving them causes DPY-4008.
|
||||
# Only remove if the key's placeholder is truly gone from the query.
|
||||
@@ -885,7 +784,7 @@ class OracleConnection(DatabaseConnection):
|
||||
|
||||
return query, params
|
||||
|
||||
async def _read_returning_values(self, returning_cols: list[str], params: dict[str, Any]) -> dict[str, Any] | None:
|
||||
def _read_returning_values(self, returning_cols: list[str], params: dict[str, Any]) -> dict[str, Any] | None:
|
||||
"""Read values from RETURNING INTO output variables after execute."""
|
||||
row: dict[str, Any] = {}
|
||||
for i, col in enumerate(returning_cols):
|
||||
@@ -895,14 +794,6 @@ class OracleConnection(DatabaseConnection):
|
||||
return None
|
||||
val = values[0] if isinstance(values, list) else values
|
||||
|
||||
# CLOB-bound columns return a LOB handle; read it to a string. The
|
||||
# async pool yields AsyncLOB whose read() is a coroutine.
|
||||
if val is not None and not isinstance(val, (str, bytes, int, float)) and hasattr(val, "read"):
|
||||
data = val.read()
|
||||
if inspect.isawaitable(data):
|
||||
data = await data
|
||||
val = data
|
||||
|
||||
# Clean alias: "LOWER(canonical_name) AS name_lower" → "name_lower"
|
||||
clean_col = col.strip()
|
||||
upper = clean_col.upper()
|
||||
@@ -1090,7 +981,7 @@ class OracleConnection(DatabaseConnection):
|
||||
raise
|
||||
|
||||
if ret_cols is not None:
|
||||
row_dict = await self._read_returning_values(ret_cols, params)
|
||||
row_dict = self._read_returning_values(ret_cols, params)
|
||||
return [ResultRow(row_dict)] if row_dict else []
|
||||
|
||||
columns = [col[0].lower() for col in cursor.description or []]
|
||||
@@ -1128,7 +1019,7 @@ class OracleConnection(DatabaseConnection):
|
||||
raise
|
||||
|
||||
if ret_cols is not None:
|
||||
row_dict = await self._read_returning_values(ret_cols, params)
|
||||
row_dict = self._read_returning_values(ret_cols, params)
|
||||
return ResultRow(row_dict) if row_dict else None
|
||||
|
||||
columns = [col[0].lower() for col in cursor.description or []]
|
||||
@@ -1161,7 +1052,7 @@ class OracleConnection(DatabaseConnection):
|
||||
await cursor.execute(query, params)
|
||||
|
||||
if ret_cols is not None:
|
||||
row_dict = await self._read_returning_values(ret_cols, params)
|
||||
row_dict = self._read_returning_values(ret_cols, params)
|
||||
if row_dict is None:
|
||||
return None
|
||||
vals = list(row_dict.values())
|
||||
|
||||
@@ -6,7 +6,7 @@ import asyncio
|
||||
import logging
|
||||
import time
|
||||
from collections.abc import AsyncIterator
|
||||
from contextlib import AsyncExitStack, asynccontextmanager
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import Any
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -101,14 +101,6 @@ async def acquire_with_retry(backend_or_pool: Any, max_retries: int = DEFAULT_MA
|
||||
"""
|
||||
Async context manager to acquire a database connection with retry logic.
|
||||
|
||||
Retries the *acquire* itself when it raises a retryable error (connection
|
||||
drop, timeout, deadlock detected during acquire). Exceptions raised by
|
||||
user code inside the ``async with`` block are NOT retried — they propagate
|
||||
as-is. Wrapping retry around the yield would violate the
|
||||
``@asynccontextmanager`` single-yield contract and surface as
|
||||
``RuntimeError("generator didn't stop after athrow()")`` on every
|
||||
retryable inner error, masking the real cause.
|
||||
|
||||
Accepts either a DatabaseBackend or a raw asyncpg.Pool for backward compatibility.
|
||||
|
||||
Usage:
|
||||
@@ -117,7 +109,7 @@ async def acquire_with_retry(backend_or_pool: Any, max_retries: int = DEFAULT_MA
|
||||
|
||||
Args:
|
||||
backend_or_pool: A DatabaseBackend instance or asyncpg.Pool
|
||||
max_retries: Maximum number of retry attempts for the acquire step
|
||||
max_retries: Maximum number of retry attempts
|
||||
|
||||
Yields:
|
||||
A DatabaseConnection (if backend) or asyncpg.Connection (if pool)
|
||||
@@ -125,32 +117,31 @@ async def acquire_with_retry(backend_or_pool: Any, max_retries: int = DEFAULT_MA
|
||||
from .db.base import DatabaseBackend
|
||||
|
||||
if isinstance(backend_or_pool, DatabaseBackend) or getattr(backend_or_pool, "_wraps_backend", False):
|
||||
# Use the backend's acquire context manager with retry
|
||||
start = time.time()
|
||||
async with AsyncExitStack() as stack:
|
||||
conn: Any = None
|
||||
for attempt in range(max_retries + 1):
|
||||
try:
|
||||
conn = await stack.enter_async_context(backend_or_pool.acquire())
|
||||
break
|
||||
except Exception as e:
|
||||
if not _is_retryable(e):
|
||||
raise
|
||||
if attempt < max_retries:
|
||||
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..."
|
||||
)
|
||||
await asyncio.sleep(delay)
|
||||
else:
|
||||
logger.error(f"Database acquire failed after {max_retries + 1} attempts: {e}")
|
||||
raise
|
||||
|
||||
acquire_time = time.time() - start
|
||||
if acquire_time > 0.05:
|
||||
logger.warning(f"[DB POOL] Slow acquire: {acquire_time:.3f}s")
|
||||
|
||||
yield conn
|
||||
last_exception = None
|
||||
for attempt in range(max_retries + 1):
|
||||
try:
|
||||
async with backend_or_pool.acquire() as conn:
|
||||
acquire_time = time.time() - start
|
||||
if acquire_time > 0.05:
|
||||
logger.warning(f"[DB POOL] Slow acquire: {acquire_time:.3f}s")
|
||||
yield conn
|
||||
return
|
||||
except Exception as e:
|
||||
if not _is_retryable(e):
|
||||
raise
|
||||
last_exception = e
|
||||
if attempt < max_retries:
|
||||
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..."
|
||||
)
|
||||
await asyncio.sleep(delay)
|
||||
else:
|
||||
logger.error(f"Database acquire failed after {max_retries + 1} attempts: {e}")
|
||||
raise last_exception
|
||||
else:
|
||||
# Legacy path: raw asyncpg.Pool
|
||||
pool = backend_or_pool
|
||||
|
||||
@@ -9,73 +9,42 @@ The database schema is automatically adjusted to match the model's dimension.
|
||||
Configuration via environment variables - see hindsight_api.config for all env var names.
|
||||
"""
|
||||
|
||||
import base64
|
||||
import logging
|
||||
import os
|
||||
import struct
|
||||
import warnings
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Literal, cast
|
||||
from urllib.parse import parse_qs, urlparse, urlunparse
|
||||
|
||||
import httpx
|
||||
from pydantic import BaseModel
|
||||
|
||||
from ..config import (
|
||||
DEFAULT_EMBEDDINGS_COHERE_MODEL,
|
||||
DEFAULT_EMBEDDINGS_GEMINI_MODEL,
|
||||
DEFAULT_EMBEDDINGS_LITELLM_MODEL,
|
||||
DEFAULT_EMBEDDINGS_LITELLM_SDK_MODEL,
|
||||
DEFAULT_EMBEDDINGS_LOCAL_FORCE_CPU,
|
||||
DEFAULT_EMBEDDINGS_LOCAL_MODEL,
|
||||
DEFAULT_EMBEDDINGS_LOCAL_TRUST_REMOTE_CODE,
|
||||
DEFAULT_EMBEDDINGS_OPENAI_MODEL,
|
||||
DEFAULT_EMBEDDINGS_ZEROENTROPY_BATCH_SIZE,
|
||||
DEFAULT_EMBEDDINGS_ZEROENTROPY_DIMENSIONS,
|
||||
DEFAULT_EMBEDDINGS_ZEROENTROPY_ENCODING_FORMAT,
|
||||
DEFAULT_EMBEDDINGS_ZEROENTROPY_LATENCY,
|
||||
DEFAULT_EMBEDDINGS_ZEROENTROPY_MODEL,
|
||||
DEFAULT_EMBEDDINGS_PROVIDER,
|
||||
DEFAULT_LITELLM_API_BASE,
|
||||
DEFAULT_ZEROENTROPY_BASE_URL,
|
||||
ENV_EMBEDDINGS_COHERE_API_KEY,
|
||||
ENV_EMBEDDINGS_GEMINI_API_KEY,
|
||||
ENV_EMBEDDINGS_LITELLM_SDK_API_KEY,
|
||||
ENV_EMBEDDINGS_LOCAL_FORCE_CPU,
|
||||
ENV_EMBEDDINGS_LOCAL_MODEL,
|
||||
ENV_EMBEDDINGS_LOCAL_TRUST_REMOTE_CODE,
|
||||
ENV_EMBEDDINGS_OPENAI_API_KEY,
|
||||
ENV_EMBEDDINGS_OPENAI_BASE_URL,
|
||||
ENV_EMBEDDINGS_OPENAI_MODEL,
|
||||
ENV_EMBEDDINGS_PROVIDER,
|
||||
ENV_EMBEDDINGS_TEI_URL,
|
||||
ENV_EMBEDDINGS_ZEROENTROPY_API_KEY,
|
||||
ENV_EMBEDDINGS_ZEROENTROPY_DIMENSIONS,
|
||||
ENV_EMBEDDINGS_ZEROENTROPY_ENCODING_FORMAT,
|
||||
ENV_LLM_API_KEY,
|
||||
)
|
||||
from .bank_attribution import apply_bank_attribution
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
ZeroEntropyInputType = Literal["document", "query"]
|
||||
ZeroEntropyLatency = Literal["fast", "slow"]
|
||||
ZeroEntropyEncodingFormat = Literal["float", "base64"]
|
||||
|
||||
|
||||
class _ZeroEntropyEmbedRequest(BaseModel):
|
||||
"""Typed request body for ZeroEntropy's non-OpenAI-compatible embed endpoint."""
|
||||
|
||||
model: str
|
||||
input: list[str]
|
||||
input_type: ZeroEntropyInputType
|
||||
dimensions: int
|
||||
encoding_format: ZeroEntropyEncodingFormat = "float"
|
||||
latency: ZeroEntropyLatency | None = None
|
||||
|
||||
|
||||
class _ZeroEntropyEmbedResult(BaseModel):
|
||||
embedding: list[float] | str
|
||||
|
||||
|
||||
class _ZeroEntropyEmbedResponse(BaseModel):
|
||||
results: list[_ZeroEntropyEmbedResult]
|
||||
|
||||
|
||||
class Embeddings(ABC):
|
||||
"""
|
||||
Abstract base class for embedding generation.
|
||||
@@ -119,14 +88,6 @@ class Embeddings(ABC):
|
||||
"""
|
||||
pass
|
||||
|
||||
def encode_query(self, texts: list[str]) -> list[list[float]]:
|
||||
"""Generate embeddings for query text. Providers without asymmetric embeddings use encode()."""
|
||||
return self.encode(texts)
|
||||
|
||||
def encode_documents(self, texts: list[str]) -> list[list[float]]:
|
||||
"""Generate embeddings for stored document text. Providers without asymmetric embeddings use encode()."""
|
||||
return self.encode(texts)
|
||||
|
||||
|
||||
class LocalSTEmbeddings(Embeddings):
|
||||
"""
|
||||
@@ -190,7 +151,7 @@ class LocalSTEmbeddings(Embeddings):
|
||||
device = "cpu"
|
||||
logger.info("Embeddings: forcing CPU mode")
|
||||
else:
|
||||
# Check for GPU (CUDA), Apple Silicon (MPS), or Intel XPU
|
||||
# Check for GPU (CUDA) or Apple Silicon (MPS)
|
||||
# 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
|
||||
@@ -198,13 +159,10 @@ class LocalSTEmbeddings(Embeddings):
|
||||
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
|
||||
device = None # Let sentence-transformers auto-detect GPU/MPS
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to detect GPU/MPS/XPU, falling back to CPU: {e}")
|
||||
logger.warning(f"Failed to detect GPU/MPS, falling back to CPU: {e}")
|
||||
|
||||
# Suppress verbose transformers warnings during model loading
|
||||
# This suppresses the "UNEXPECTED" warnings from BertModel which are harmless
|
||||
@@ -250,172 +208,6 @@ class LocalSTEmbeddings(Embeddings):
|
||||
return [emb.tolist() for emb in embeddings]
|
||||
|
||||
|
||||
class OnnxEmbeddings(Embeddings):
|
||||
"""Local ONNX Runtime embeddings provider.
|
||||
|
||||
This provider runs transformer embedding models in-process with ONNX Runtime,
|
||||
avoiding a sidecar Ollama/TEI server or a remote embeddings API. It supports
|
||||
sentence-transformer style mean pooling and E5-style asymmetric prefixes.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model_id: str,
|
||||
model_path: str | None = None,
|
||||
tokenizer_name_or_path: str | None = None,
|
||||
onnx_file: str = "onnx/model.onnx",
|
||||
dimensions: int | None = None,
|
||||
max_tokens: int = 512,
|
||||
pooling: str = "mean",
|
||||
normalize: bool = True,
|
||||
query_prefix: str = "query: ",
|
||||
passage_prefix: str = "passage: ",
|
||||
output_name: str | None = None,
|
||||
):
|
||||
self.model_id = model_id
|
||||
self.model_path = model_path
|
||||
if model_path and tokenizer_name_or_path is None:
|
||||
logger.warning(
|
||||
"Embeddings: ONNX model_path is set without tokenizer_name_or_path; "
|
||||
"falling back to tokenizer from model_id %s. Set "
|
||||
"HINDSIGHT_API_EMBEDDINGS_ONNX_TOKENIZER_NAME_OR_PATH when using local ONNX artifacts.",
|
||||
model_id,
|
||||
)
|
||||
self.tokenizer_name_or_path = tokenizer_name_or_path or model_id
|
||||
self.onnx_file = onnx_file
|
||||
self.configured_dimensions = dimensions
|
||||
self.max_tokens = max_tokens
|
||||
self.pooling = pooling.lower()
|
||||
if self.pooling not in {"mean", "cls"}:
|
||||
raise ValueError("ONNX embeddings pooling must be 'mean' or 'cls'")
|
||||
self.normalize = normalize
|
||||
self.query_prefix = query_prefix
|
||||
self.passage_prefix = passage_prefix
|
||||
self.output_name = output_name
|
||||
self._session = None
|
||||
self._tokenizer = None
|
||||
self._dimension: int | None = dimensions
|
||||
|
||||
@property
|
||||
def provider_name(self) -> str:
|
||||
return "onnx"
|
||||
|
||||
@property
|
||||
def dimension(self) -> int:
|
||||
if self._dimension is None:
|
||||
raise RuntimeError("Embeddings not initialized. Call initialize() first.")
|
||||
return self._dimension
|
||||
|
||||
async def initialize(self) -> None:
|
||||
if self._session is not None and self._tokenizer is not None:
|
||||
return
|
||||
|
||||
try:
|
||||
import onnxruntime as ort
|
||||
from transformers import AutoTokenizer
|
||||
except ImportError as exc:
|
||||
raise ImportError(
|
||||
"onnxruntime and transformers are required for OnnxEmbeddings. "
|
||||
"Install with: pip install 'hindsight-api-slim[local-onnx]'"
|
||||
) from exc
|
||||
|
||||
model_path = self.model_path
|
||||
if not model_path:
|
||||
try:
|
||||
from huggingface_hub import snapshot_download
|
||||
except ImportError as exc:
|
||||
raise ImportError(
|
||||
"huggingface-hub is required to download ONNX embedding models. "
|
||||
"Set HINDSIGHT_API_EMBEDDINGS_ONNX_MODEL_PATH or install local-onnx."
|
||||
) from exc
|
||||
# Some large ONNX exports, for example BAAI/bge-m3, store weights in
|
||||
# an external sidecar file next to model.onnx. Download both the
|
||||
# requested graph and its conventional *_data sidecar when present.
|
||||
snapshot_dir = snapshot_download(
|
||||
repo_id=self.model_id,
|
||||
allow_patterns=[self.onnx_file, f"{self.onnx_file}_data"],
|
||||
)
|
||||
model_path = os.path.join(snapshot_dir, self.onnx_file)
|
||||
|
||||
logger.info(
|
||||
"Embeddings: initializing ONNX provider with model %s (%s)",
|
||||
self.model_id,
|
||||
model_path,
|
||||
)
|
||||
logger.info(
|
||||
"Embeddings: ONNX query_prefix=%r passage_prefix=%r pooling=%s normalize=%s",
|
||||
self.query_prefix,
|
||||
self.passage_prefix,
|
||||
self.pooling,
|
||||
self.normalize,
|
||||
)
|
||||
self._tokenizer = AutoTokenizer.from_pretrained(self.tokenizer_name_or_path)
|
||||
self._session = ort.InferenceSession(model_path, providers=["CPUExecutionProvider"])
|
||||
|
||||
detected = len(self.encode(["test"])[0])
|
||||
if self.configured_dimensions is not None and detected != self.configured_dimensions:
|
||||
raise ValueError(
|
||||
f"Configured ONNX embedding dimension {self.configured_dimensions} does not match model output {detected}"
|
||||
)
|
||||
self._dimension = detected
|
||||
logger.info("Embeddings: ONNX provider initialized (dim: %s)", self._dimension)
|
||||
|
||||
def _encode_prefixed(self, texts: list[str], prefix: str) -> list[list[float]]:
|
||||
if prefix:
|
||||
return self.encode([f"{prefix}{text}" for text in texts])
|
||||
return self.encode(texts)
|
||||
|
||||
def encode_query(self, texts: list[str]) -> list[list[float]]:
|
||||
return self._encode_prefixed(texts, self.query_prefix)
|
||||
|
||||
def encode_documents(self, texts: list[str]) -> list[list[float]]:
|
||||
return self._encode_prefixed(texts, self.passage_prefix)
|
||||
|
||||
def encode(self, texts: list[str]) -> list[list[float]]:
|
||||
if self._session is None or self._tokenizer is None:
|
||||
raise RuntimeError("Embeddings not initialized. Call initialize() first.")
|
||||
if not texts:
|
||||
return []
|
||||
|
||||
import numpy as np
|
||||
|
||||
encoded = self._tokenizer(
|
||||
texts,
|
||||
padding=True,
|
||||
truncation=True,
|
||||
max_length=self.max_tokens,
|
||||
return_tensors="np",
|
||||
)
|
||||
input_names = {inp.name for inp in self._session.get_inputs()}
|
||||
ort_inputs = {name: value for name, value in encoded.items() if name in input_names}
|
||||
if "token_type_ids" in input_names and "token_type_ids" not in ort_inputs:
|
||||
ort_inputs["token_type_ids"] = np.zeros_like(encoded["input_ids"])
|
||||
|
||||
outputs = self._session.run([self.output_name] if self.output_name else None, ort_inputs)
|
||||
token_embeddings = outputs[0]
|
||||
|
||||
# Some exported models expose a pooled 2-D embedding as their first output.
|
||||
if getattr(token_embeddings, "ndim", 0) == 2:
|
||||
embeddings = token_embeddings
|
||||
elif self.pooling == "cls":
|
||||
embeddings = token_embeddings[:, 0]
|
||||
else:
|
||||
attention_mask = encoded.get("attention_mask")
|
||||
if attention_mask is None:
|
||||
attention_mask = np.ones(token_embeddings.shape[:2], dtype=np.float32)
|
||||
mask = attention_mask[..., None].astype(np.float32)
|
||||
summed = (token_embeddings * mask).sum(axis=1)
|
||||
counts = np.clip(mask.sum(axis=1), a_min=1e-9, a_max=None)
|
||||
embeddings = summed / counts
|
||||
|
||||
if self.normalize:
|
||||
norms = np.linalg.norm(embeddings, axis=1, keepdims=True)
|
||||
norms[norms == 0] = 1
|
||||
embeddings = embeddings / norms
|
||||
|
||||
return embeddings.astype(float).tolist()
|
||||
|
||||
|
||||
class RemoteTEIEmbeddings(Embeddings):
|
||||
"""
|
||||
Remote embeddings implementation using HuggingFace Text Embeddings Inference (TEI) HTTP API.
|
||||
@@ -593,7 +385,6 @@ class OpenAIEmbeddings(Embeddings):
|
||||
model: str = DEFAULT_EMBEDDINGS_OPENAI_MODEL,
|
||||
base_url: str | None = None,
|
||||
batch_size: int = 100,
|
||||
dimensions: int | None = None,
|
||||
max_retries: int = 3,
|
||||
):
|
||||
"""
|
||||
@@ -604,14 +395,12 @@ class OpenAIEmbeddings(Embeddings):
|
||||
model: OpenAI embedding model name (default: text-embedding-3-small)
|
||||
base_url: Custom base URL for OpenAI-compatible API (e.g., Azure OpenAI endpoint)
|
||||
batch_size: Maximum batch size for embedding requests (default: 100)
|
||||
dimensions: Optional requested output dimensions for OpenAI text-embedding-3 models
|
||||
max_retries: Maximum number of retries for failed requests (default: 3)
|
||||
"""
|
||||
self.api_key = api_key
|
||||
self.model = model
|
||||
self.base_url = base_url
|
||||
self.batch_size = batch_size
|
||||
self.dimensions = dimensions
|
||||
self.max_retries = max_retries
|
||||
self._client = None
|
||||
self._dimension: int | None = None
|
||||
@@ -656,9 +445,7 @@ class OpenAIEmbeddings(Embeddings):
|
||||
self._client = OpenAI(**client_kwargs)
|
||||
|
||||
# Try to get dimension from known models, otherwise do a test embedding
|
||||
if self.dimensions is not None:
|
||||
self._dimension = self.dimensions
|
||||
elif self.model in self.MODEL_DIMENSIONS:
|
||||
if self.model in self.MODEL_DIMENSIONS:
|
||||
self._dimension = self.MODEL_DIMENSIONS[self.model]
|
||||
else:
|
||||
# Do a test embedding to detect dimension
|
||||
@@ -693,15 +480,10 @@ class OpenAIEmbeddings(Embeddings):
|
||||
for i in range(0, len(texts), self.batch_size):
|
||||
batch = texts[i : i + self.batch_size]
|
||||
|
||||
request = {
|
||||
"model": self.model,
|
||||
"input": batch,
|
||||
}
|
||||
if self.dimensions is not None:
|
||||
request["dimensions"] = self.dimensions
|
||||
apply_bank_attribution(request)
|
||||
|
||||
response = self._client.embeddings.create(**request)
|
||||
response = self._client.embeddings.create(
|
||||
model=self.model,
|
||||
input=batch,
|
||||
)
|
||||
|
||||
# Sort by index to ensure correct order
|
||||
batch_embeddings = sorted(response.data, key=lambda x: x.index)
|
||||
@@ -710,74 +492,6 @@ class OpenAIEmbeddings(Embeddings):
|
||||
return all_embeddings
|
||||
|
||||
|
||||
class CodexOAuthEmbeddings(OpenAIEmbeddings):
|
||||
"""
|
||||
OpenAI embeddings using the Codex/ChatGPT OAuth token from the Codex
|
||||
``auth.json`` (``$CODEX_HOME/auth.json``, or ``~/.codex/auth.json`` when unset).
|
||||
|
||||
Codex OAuth is an LLM-provider auth path in Hindsight, but the same bearer token
|
||||
can also authenticate against the standard OpenAI embeddings endpoint. This keeps
|
||||
embeddings on the user's existing Codex subscription/OAuth path without requiring
|
||||
a separate OpenAI/OpenRouter/Gemini/Cohere API key.
|
||||
|
||||
Token refresh is handled automatically: the manager proactively refreshes the
|
||||
access_token before it expires and reactively refreshes on 401 responses from
|
||||
the embeddings API.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model: str = DEFAULT_EMBEDDINGS_OPENAI_MODEL,
|
||||
batch_size: int = 100,
|
||||
dimensions: int | None = None,
|
||||
max_retries: int = 3,
|
||||
):
|
||||
from .providers.codex_auth import CodexAuthManager
|
||||
|
||||
self._auth_manager = CodexAuthManager.from_file()
|
||||
super().__init__(
|
||||
api_key=self._auth_manager.access_token,
|
||||
model=model,
|
||||
base_url="https://api.openai.com/v1",
|
||||
batch_size=batch_size,
|
||||
dimensions=dimensions,
|
||||
max_retries=max_retries,
|
||||
)
|
||||
|
||||
@property
|
||||
def provider_name(self) -> str:
|
||||
return "openai-codex"
|
||||
|
||||
def encode(self, texts: list[str]) -> list[list[float]]:
|
||||
"""Generate embeddings, refreshing the OAuth token if needed.
|
||||
|
||||
Proactively refreshes before the call when the token is near expiry,
|
||||
and reactively refreshes once on a 401 from the OpenAI embeddings API.
|
||||
"""
|
||||
from openai import AuthenticationError
|
||||
|
||||
# Proactive refresh — cheap when fresh (JWT exp decode + compare).
|
||||
self._auth_manager.ensure_fresh_token()
|
||||
if self._auth_manager.access_token != self.api_key:
|
||||
self.api_key = self._auth_manager.access_token
|
||||
if self._client is not None:
|
||||
self._client.api_key = self._auth_manager.access_token
|
||||
|
||||
try:
|
||||
return super().encode(texts)
|
||||
except AuthenticationError:
|
||||
# Reactive refresh — token was valid by the JWT clock but the
|
||||
# server rejected it (rotated server-side, race, etc.).
|
||||
self._auth_manager.refresh_tokens(
|
||||
reason="reactive (401 from embeddings API)",
|
||||
force=True,
|
||||
)
|
||||
self.api_key = self._auth_manager.access_token
|
||||
if self._client is not None:
|
||||
self._client.api_key = self._auth_manager.access_token
|
||||
return super().encode(texts)
|
||||
|
||||
|
||||
class CohereEmbeddings(Embeddings):
|
||||
"""
|
||||
Cohere embeddings implementation using the Cohere API.
|
||||
@@ -919,149 +633,6 @@ class CohereEmbeddings(Embeddings):
|
||||
return all_embeddings
|
||||
|
||||
|
||||
class ZeroEntropyEmbeddings(Embeddings):
|
||||
"""
|
||||
ZeroEntropy embeddings implementation using the zembed API.
|
||||
|
||||
ZeroEntropy's embeddings endpoint is not OpenAI-compatible: it lives at
|
||||
/v1/models/embed and requires provider-specific parameters such as
|
||||
input_type. Hindsight stores document-side vectors and uses query-side
|
||||
vectors during recall, so this provider exposes explicit encode_documents()
|
||||
and encode_query() helpers while keeping encode() as document-side default.
|
||||
"""
|
||||
|
||||
VALID_DIMENSIONS = frozenset({2560, 1280, 640, 320, 160, 80, 40})
|
||||
VALID_ENCODING_FORMATS = frozenset({"float", "base64"})
|
||||
VALID_LATENCIES = frozenset({"fast", "slow"})
|
||||
DEFAULT_BASE_URL = DEFAULT_ZEROENTROPY_BASE_URL
|
||||
EMBED_PATH = "/v1/models/embed"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
api_key: str,
|
||||
model: str = DEFAULT_EMBEDDINGS_ZEROENTROPY_MODEL,
|
||||
base_url: str | None = None,
|
||||
dimensions: int = DEFAULT_EMBEDDINGS_ZEROENTROPY_DIMENSIONS,
|
||||
batch_size: int = DEFAULT_EMBEDDINGS_ZEROENTROPY_BATCH_SIZE,
|
||||
encoding_format: str = DEFAULT_EMBEDDINGS_ZEROENTROPY_ENCODING_FORMAT,
|
||||
latency: str | None = DEFAULT_EMBEDDINGS_ZEROENTROPY_LATENCY,
|
||||
timeout: float = 60.0,
|
||||
):
|
||||
if dimensions not in self.VALID_DIMENSIONS:
|
||||
valid = ", ".join(str(dim) for dim in sorted(self.VALID_DIMENSIONS, reverse=True))
|
||||
raise ValueError(f"{ENV_EMBEDDINGS_ZEROENTROPY_DIMENSIONS} must be one of {valid}, got {dimensions}")
|
||||
if batch_size < 1:
|
||||
raise ValueError("ZeroEntropy embeddings batch_size must be >= 1")
|
||||
if encoding_format not in self.VALID_ENCODING_FORMATS:
|
||||
valid_formats = ", ".join(sorted(self.VALID_ENCODING_FORMATS))
|
||||
raise ValueError(
|
||||
f"{ENV_EMBEDDINGS_ZEROENTROPY_ENCODING_FORMAT} must be one of {valid_formats}, got {encoding_format!r}"
|
||||
)
|
||||
if latency is not None and latency not in self.VALID_LATENCIES:
|
||||
valid_latencies = ", ".join(sorted(self.VALID_LATENCIES))
|
||||
raise ValueError(f"ZeroEntropy embeddings latency must be one of {valid_latencies}, got {latency!r}")
|
||||
|
||||
self.api_key = api_key
|
||||
self.model = model
|
||||
self.base_url = base_url.rstrip("/") if base_url else self.DEFAULT_BASE_URL
|
||||
self.embed_url = f"{self.base_url}{self.EMBED_PATH}"
|
||||
self.dimensions = dimensions
|
||||
self.batch_size = batch_size
|
||||
self.encoding_format = cast(ZeroEntropyEncodingFormat, encoding_format)
|
||||
self.latency = cast(ZeroEntropyLatency | None, latency)
|
||||
self.timeout = timeout
|
||||
self._client: httpx.Client | None = None
|
||||
self._dimension: int | None = None
|
||||
|
||||
@property
|
||||
def provider_name(self) -> str:
|
||||
return "zeroentropy"
|
||||
|
||||
@property
|
||||
def dimension(self) -> int:
|
||||
if self._dimension is None:
|
||||
raise RuntimeError("Embeddings not initialized. Call initialize() first.")
|
||||
return self._dimension
|
||||
|
||||
async def initialize(self) -> None:
|
||||
"""Initialize the ZeroEntropy HTTP client."""
|
||||
if self._client is not None:
|
||||
return
|
||||
|
||||
logger.info(
|
||||
f"Embeddings: initializing ZeroEntropy provider with model {self.model} "
|
||||
f"(dim: {self.dimensions}, batch_size={self.batch_size})"
|
||||
)
|
||||
self._client = httpx.Client(
|
||||
timeout=self.timeout,
|
||||
headers={
|
||||
"Authorization": f"Bearer {self.api_key}",
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
)
|
||||
# zembed-1 dimensions are explicit Matryoshka truncation steps. Avoid a
|
||||
# startup probe so boot does not burn quota or require a throwaway input.
|
||||
self._dimension = self.dimensions
|
||||
logger.info(f"Embeddings: ZeroEntropy provider initialized (model: {self.model}, dim: {self._dimension})")
|
||||
|
||||
def encode(self, texts: list[str]) -> list[list[float]]:
|
||||
"""Generate document-side embeddings for backwards-compatible callers."""
|
||||
return self.encode_documents(texts)
|
||||
|
||||
def encode_documents(self, texts: list[str]) -> list[list[float]]:
|
||||
"""Generate document-side embeddings for retained content."""
|
||||
return self._encode_with_input_type(texts, "document")
|
||||
|
||||
def encode_query(self, texts: list[str]) -> list[list[float]]:
|
||||
"""Generate query-side embeddings for recall/search queries."""
|
||||
return self._encode_with_input_type(texts, "query")
|
||||
|
||||
def _encode_with_input_type(self, texts: list[str], input_type: ZeroEntropyInputType) -> list[list[float]]:
|
||||
if self._client is None:
|
||||
raise RuntimeError("Embeddings not initialized. Call initialize() first.")
|
||||
|
||||
if not texts:
|
||||
return []
|
||||
|
||||
all_embeddings: list[list[float]] = []
|
||||
for i in range(0, len(texts), self.batch_size):
|
||||
batch = texts[i : i + self.batch_size]
|
||||
request = _ZeroEntropyEmbedRequest(
|
||||
model=self.model,
|
||||
input=batch,
|
||||
input_type=input_type,
|
||||
dimensions=self.dimensions,
|
||||
encoding_format=self.encoding_format,
|
||||
latency=self.latency,
|
||||
)
|
||||
|
||||
try:
|
||||
response = self._client.post(self.embed_url, json=request.model_dump(exclude_none=True))
|
||||
response.raise_for_status()
|
||||
except httpx.HTTPError as e:
|
||||
raise RuntimeError(f"ZeroEntropy embedding request failed: {e}") from e
|
||||
|
||||
parsed = _ZeroEntropyEmbedResponse.model_validate(response.json())
|
||||
if len(parsed.results) != len(batch):
|
||||
raise RuntimeError(
|
||||
f"ZeroEntropy returned {len(parsed.results)} embeddings for {len(batch)} input texts; "
|
||||
"expected exact 1:1 alignment"
|
||||
)
|
||||
all_embeddings.extend(self._parse_embedding(result.embedding) for result in parsed.results)
|
||||
|
||||
return all_embeddings
|
||||
|
||||
@staticmethod
|
||||
def _parse_embedding(embedding: list[float] | str) -> list[float]:
|
||||
if not isinstance(embedding, str):
|
||||
return embedding
|
||||
|
||||
raw = base64.b64decode(embedding)
|
||||
if len(raw) % 4 != 0:
|
||||
raise RuntimeError("ZeroEntropy returned invalid base64 embedding length")
|
||||
return list(struct.unpack(f"<{len(raw) // 4}f", raw))
|
||||
|
||||
|
||||
class LiteLLMEmbeddings(Embeddings):
|
||||
"""
|
||||
LiteLLM embeddings implementation using LiteLLM proxy's /embeddings endpoint.
|
||||
@@ -1195,7 +766,7 @@ class LiteLLMSDKEmbeddings(Embeddings):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
api_key: str | None = None,
|
||||
api_key: str,
|
||||
model: str = DEFAULT_EMBEDDINGS_LITELLM_SDK_MODEL,
|
||||
api_base: str | None = None,
|
||||
output_dimensions: int | None = None,
|
||||
@@ -1207,8 +778,7 @@ class LiteLLMSDKEmbeddings(Embeddings):
|
||||
Initialize LiteLLM SDK embeddings client.
|
||||
|
||||
Args:
|
||||
api_key: API key for the embedding provider (optional — omit for
|
||||
providers that use ambient credentials, e.g. AWS Bedrock with IAM)
|
||||
api_key: API key for the embedding provider
|
||||
model: Model name with provider prefix (e.g., "cohere/embed-english-v3.0")
|
||||
api_base: Custom base URL for API (optional)
|
||||
output_dimensions: Optional output embedding dimensions (provider-dependent)
|
||||
@@ -1258,9 +828,8 @@ class LiteLLMSDKEmbeddings(Embeddings):
|
||||
embed_kwargs = {
|
||||
"model": self.model,
|
||||
"input": ["test"],
|
||||
"api_key": self.api_key,
|
||||
}
|
||||
if self.api_key:
|
||||
embed_kwargs["api_key"] = self.api_key
|
||||
if self.encoding_format:
|
||||
embed_kwargs["encoding_format"] = self.encoding_format
|
||||
if self.api_base:
|
||||
@@ -1311,9 +880,8 @@ class LiteLLMSDKEmbeddings(Embeddings):
|
||||
embed_kwargs = {
|
||||
"model": self.model,
|
||||
"input": batch,
|
||||
"api_key": self.api_key,
|
||||
}
|
||||
if self.api_key:
|
||||
embed_kwargs["api_key"] = self.api_key
|
||||
if self.encoding_format:
|
||||
embed_kwargs["encoding_format"] = self.encoding_format
|
||||
if self.api_base:
|
||||
@@ -1343,21 +911,6 @@ class LiteLLMSDKEmbeddings(Embeddings):
|
||||
return all_embeddings
|
||||
|
||||
|
||||
# Gemini Embedding 2+ multimodal models return a SINGLE aggregated embedding
|
||||
# for a multi-input request instead of one vector per input (see
|
||||
# https://ai.google.dev/gemini-api/docs/embeddings#embedding-aggregation). For
|
||||
# these models we must embed one input per call to preserve the 1:1 input→vector
|
||||
# alignment the rest of the pipeline relies on. The marker matches preview and GA
|
||||
# names (e.g. "gemini-embedding-2-preview", "gemini-embedding-2"), with or
|
||||
# without a "google/" or "models/" prefix.
|
||||
_GEMINI_AGGREGATING_MODEL_MARKER = "gemini-embedding-2"
|
||||
|
||||
|
||||
def _gemini_model_aggregates_inputs(model: str) -> bool:
|
||||
"""Whether the model aggregates a multi-input request into one embedding."""
|
||||
return _GEMINI_AGGREGATING_MODEL_MARKER in model.lower()
|
||||
|
||||
|
||||
class GeminiEmbeddings(Embeddings):
|
||||
"""
|
||||
Google embeddings via the google.genai SDK.
|
||||
@@ -1367,10 +920,6 @@ class GeminiEmbeddings(Embeddings):
|
||||
2. Vertex AI with service account or Application Default Credentials (ADC)
|
||||
|
||||
Uses the embed_content API: client.models.embed_content(model, contents)
|
||||
|
||||
Gemini Embedding 2+ multimodal models aggregate a multi-input request into a
|
||||
single embedding, so for those the batch size is forced to 1 (one input per
|
||||
call) to keep one vector per input.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
@@ -1525,13 +1074,9 @@ class GeminiEmbeddings(Embeddings):
|
||||
|
||||
all_embeddings = []
|
||||
|
||||
# Gemini Embedding 2+ multimodal models return one aggregated vector for a
|
||||
# multi-input request, so embed one input per call to keep 1:1 alignment.
|
||||
batch_size = 1 if _gemini_model_aggregates_inputs(self.model) else self.batch_size
|
||||
|
||||
# Process in batches
|
||||
for i in range(0, len(texts), batch_size):
|
||||
batch = texts[i : i + batch_size]
|
||||
for i in range(0, len(texts), self.batch_size):
|
||||
batch = texts[i : i + self.batch_size]
|
||||
|
||||
embed_kwargs = {"model": self.model, "contents": batch}
|
||||
if self._embed_config is not None:
|
||||
@@ -1539,13 +1084,7 @@ class GeminiEmbeddings(Embeddings):
|
||||
|
||||
result = self._client.models.embed_content(**embed_kwargs)
|
||||
|
||||
embeddings = result.embeddings or []
|
||||
if len(embeddings) != len(batch):
|
||||
raise RuntimeError(
|
||||
f"Gemini embeddings backend returned {len(embeddings)} vectors for "
|
||||
f"{len(batch)} input texts (model {self.model}); expected exact 1:1 alignment"
|
||||
)
|
||||
all_embeddings.extend([emb.values for emb in embeddings])
|
||||
all_embeddings.extend([emb.values for emb in result.embeddings])
|
||||
|
||||
# L2-normalize when output_dimensionality is set — Gemini only returns
|
||||
# normalized vectors at full 3072 dims; truncated dims need re-normalization
|
||||
@@ -1586,20 +1125,6 @@ def create_embeddings_from_env() -> Embeddings:
|
||||
force_cpu=config.embeddings_local_force_cpu,
|
||||
trust_remote_code=config.embeddings_local_trust_remote_code,
|
||||
)
|
||||
elif provider == "onnx":
|
||||
return OnnxEmbeddings(
|
||||
model_id=config.embeddings_onnx_model_id,
|
||||
model_path=config.embeddings_onnx_model_path,
|
||||
tokenizer_name_or_path=config.embeddings_onnx_tokenizer_name_or_path,
|
||||
onnx_file=config.embeddings_onnx_file,
|
||||
dimensions=config.embeddings_onnx_dimensions,
|
||||
max_tokens=config.embeddings_onnx_max_tokens,
|
||||
pooling=config.embeddings_onnx_pooling,
|
||||
normalize=config.embeddings_onnx_normalize,
|
||||
query_prefix=config.embeddings_onnx_query_prefix,
|
||||
passage_prefix=config.embeddings_onnx_passage_prefix,
|
||||
output_name=config.embeddings_onnx_output_name,
|
||||
)
|
||||
elif provider == "openai":
|
||||
# Use dedicated embeddings API key, or fall back to LLM API key
|
||||
api_key = os.environ.get(ENV_EMBEDDINGS_OPENAI_API_KEY) or os.environ.get(ENV_LLM_API_KEY)
|
||||
@@ -1615,14 +1140,6 @@ def create_embeddings_from_env() -> Embeddings:
|
||||
model=model,
|
||||
base_url=base_url,
|
||||
batch_size=config.embeddings_openai_batch_size,
|
||||
dimensions=config.embeddings_openai_dimensions,
|
||||
)
|
||||
elif provider == "openai-codex":
|
||||
model = os.environ.get(ENV_EMBEDDINGS_OPENAI_MODEL, DEFAULT_EMBEDDINGS_OPENAI_MODEL)
|
||||
return CodexOAuthEmbeddings(
|
||||
model=model,
|
||||
batch_size=config.embeddings_openai_batch_size,
|
||||
dimensions=config.embeddings_openai_dimensions,
|
||||
)
|
||||
elif provider == "openrouter":
|
||||
api_key = config.embeddings_openrouter_api_key
|
||||
@@ -1636,37 +1153,6 @@ def create_embeddings_from_env() -> Embeddings:
|
||||
model=config.embeddings_openrouter_model,
|
||||
base_url="https://openrouter.ai/api/v1",
|
||||
batch_size=config.embeddings_openai_batch_size,
|
||||
dimensions=config.embeddings_openai_dimensions,
|
||||
)
|
||||
elif provider == "requesty":
|
||||
api_key = config.embeddings_requesty_api_key
|
||||
if not api_key:
|
||||
raise ValueError(
|
||||
"HINDSIGHT_API_EMBEDDINGS_REQUESTY_API_KEY, HINDSIGHT_API_REQUESTY_API_KEY, "
|
||||
f"or {ENV_LLM_API_KEY} is required when {ENV_EMBEDDINGS_PROVIDER} is 'requesty'"
|
||||
)
|
||||
return OpenAIEmbeddings(
|
||||
api_key=api_key,
|
||||
model=config.embeddings_requesty_model,
|
||||
base_url="https://router.requesty.ai/v1",
|
||||
batch_size=config.embeddings_openai_batch_size,
|
||||
dimensions=config.embeddings_openai_dimensions,
|
||||
)
|
||||
elif provider == "zeroentropy":
|
||||
api_key = config.embeddings_zeroentropy_api_key
|
||||
if not api_key:
|
||||
raise ValueError(
|
||||
f"{ENV_EMBEDDINGS_ZEROENTROPY_API_KEY} or ZEROENTROPY_API_KEY is required "
|
||||
f"when {ENV_EMBEDDINGS_PROVIDER} is 'zeroentropy'"
|
||||
)
|
||||
return ZeroEntropyEmbeddings(
|
||||
api_key=api_key,
|
||||
model=config.embeddings_zeroentropy_model,
|
||||
base_url=config.embeddings_zeroentropy_base_url,
|
||||
dimensions=config.embeddings_zeroentropy_dimensions,
|
||||
batch_size=config.embeddings_zeroentropy_batch_size,
|
||||
encoding_format=config.embeddings_zeroentropy_encoding_format,
|
||||
latency=config.embeddings_zeroentropy_latency,
|
||||
)
|
||||
elif provider == "cohere":
|
||||
api_key = config.embeddings_cohere_api_key
|
||||
@@ -1685,8 +1171,13 @@ def create_embeddings_from_env() -> Embeddings:
|
||||
model=config.embeddings_litellm_model,
|
||||
)
|
||||
elif provider == "litellm-sdk":
|
||||
api_key = config.embeddings_litellm_sdk_api_key
|
||||
if not api_key:
|
||||
raise ValueError(
|
||||
f"{ENV_EMBEDDINGS_LITELLM_SDK_API_KEY} is required when {ENV_EMBEDDINGS_PROVIDER} is 'litellm-sdk'"
|
||||
)
|
||||
return LiteLLMSDKEmbeddings(
|
||||
api_key=config.embeddings_litellm_sdk_api_key or None,
|
||||
api_key=api_key,
|
||||
model=config.embeddings_litellm_sdk_model,
|
||||
api_base=config.embeddings_litellm_sdk_api_base,
|
||||
output_dimensions=config.embeddings_litellm_sdk_output_dimensions,
|
||||
@@ -1715,6 +1206,5 @@ def create_embeddings_from_env() -> Embeddings:
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unknown embeddings provider: {provider}. "
|
||||
f"Supported: 'local', 'onnx', 'tei', 'openai', 'openai-codex', 'openrouter', 'requesty', 'cohere', 'google', "
|
||||
f"'zeroentropy', 'litellm', 'litellm-sdk'"
|
||||
f"Supported: 'local', 'tei', 'openai', 'cohere', 'google', 'litellm', 'litellm-sdk'"
|
||||
)
|
||||
|
||||
@@ -16,15 +16,7 @@ from typing import Any, Final
|
||||
|
||||
from .db_utils import acquire_with_retry
|
||||
from .memory_engine import fq_table
|
||||
from .retain.entity_labels import (
|
||||
build_labels_lookup as _build_labels_lookup_from_config,
|
||||
)
|
||||
from .retain.entity_labels import (
|
||||
is_label_entity as _is_label_entity,
|
||||
)
|
||||
from .retain.entity_labels import (
|
||||
parse_entity_labels as _parse_entity_labels,
|
||||
)
|
||||
from .retain.entity_labels import build_labels_lookup as _build_labels_lookup_from_config
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -97,12 +89,7 @@ class EntityResolver:
|
||||
Resolves entities to canonical IDs with disambiguation.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
pool: Any,
|
||||
entity_lookup: str = "full",
|
||||
entity_resolution_batch_size: int = 100,
|
||||
):
|
||||
def __init__(self, pool: Any, entity_lookup: str = "full"):
|
||||
"""
|
||||
Initialize entity resolver.
|
||||
|
||||
@@ -111,14 +98,9 @@ class EntityResolver:
|
||||
entity_lookup: Lookup strategy — "full" loads all bank entities then
|
||||
matches in Python; "trigram" uses pg_trgm GIN index to fetch only
|
||||
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.
|
||||
"""
|
||||
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._pg_trgm_checked = False
|
||||
# Backend-specific operations — accessed via pool.ops (Django pattern).
|
||||
self._ops = pool.ops if pool is not None else None
|
||||
@@ -217,11 +199,6 @@ class EntityResolver:
|
||||
"""Build a set of valid 'key:value' entity label strings for fast lookup."""
|
||||
return _build_labels_lookup_from_config(entity_labels)
|
||||
|
||||
@staticmethod
|
||||
def _chunked(values: list[str], size: int) -> list[list[str]]:
|
||||
"""Split values into fixed-size batches."""
|
||||
return [values[i : i + size] for i in range(0, len(values), size)]
|
||||
|
||||
async def resolve_entities_batch(
|
||||
self,
|
||||
bank_id: str,
|
||||
@@ -251,15 +228,14 @@ class EntityResolver:
|
||||
return []
|
||||
|
||||
taxonomy_lookup = self._build_labels_lookup(entity_labels)
|
||||
labels_cfg = _parse_entity_labels(entity_labels)
|
||||
if conn is None:
|
||||
async with acquire_with_retry(self.pool) as conn:
|
||||
return await self._resolve_entities_batch_impl(
|
||||
conn, bank_id, entities_data, context, unit_event_date, taxonomy_lookup, labels_cfg
|
||||
conn, bank_id, entities_data, context, unit_event_date, taxonomy_lookup
|
||||
)
|
||||
else:
|
||||
return await self._resolve_entities_batch_impl(
|
||||
conn, bank_id, entities_data, context, unit_event_date, taxonomy_lookup, labels_cfg
|
||||
conn, bank_id, entities_data, context, unit_event_date, taxonomy_lookup
|
||||
)
|
||||
|
||||
async def _resolve_entities_batch_impl(
|
||||
@@ -270,16 +246,13 @@ class EntityResolver:
|
||||
context: str,
|
||||
unit_event_date,
|
||||
taxonomy_lookup: set[str] | None = None,
|
||||
labels_cfg=None,
|
||||
) -> list[str]:
|
||||
if self.entity_lookup == "trigram":
|
||||
# Route to backend-specific fuzzy strategy.
|
||||
# Non-PG backends (Oracle) use UTL_MATCH instead of pg_trgm.
|
||||
backend_strategy = self._ops.get_entity_resolution_strategy()
|
||||
if backend_strategy == "oracle_fuzzy":
|
||||
return await self._resolve_entities_batch_oracle_fuzzy(
|
||||
conn, bank_id, entities_data, unit_event_date, taxonomy_lookup, labels_cfg
|
||||
)
|
||||
return await self._resolve_entities_batch_oracle_fuzzy(conn, bank_id, entities_data, unit_event_date)
|
||||
# Auto-detect pg_trgm availability on first call and fall back to
|
||||
# "full" strategy if the extension is not installed. See #626.
|
||||
if not self._pg_trgm_checked:
|
||||
@@ -293,24 +266,12 @@ class EntityResolver:
|
||||
"https://github.com/vectorize-io/hindsight/issues/626"
|
||||
)
|
||||
self.entity_lookup = "full"
|
||||
return await self._resolve_entities_batch_full(
|
||||
conn, bank_id, entities_data, unit_event_date, taxonomy_lookup, labels_cfg
|
||||
)
|
||||
return await self._resolve_entities_batch_trigram(
|
||||
conn, bank_id, entities_data, unit_event_date, taxonomy_lookup, labels_cfg
|
||||
)
|
||||
return await self._resolve_entities_batch_full(
|
||||
conn, bank_id, entities_data, unit_event_date, taxonomy_lookup, labels_cfg
|
||||
)
|
||||
return await self._resolve_entities_batch_full(conn, bank_id, entities_data, unit_event_date)
|
||||
return await self._resolve_entities_batch_trigram(conn, bank_id, entities_data, unit_event_date)
|
||||
return await self._resolve_entities_batch_full(conn, bank_id, entities_data, unit_event_date)
|
||||
|
||||
async def _resolve_entities_batch_full(
|
||||
self,
|
||||
conn,
|
||||
bank_id: str,
|
||||
entities_data: list[dict],
|
||||
unit_event_date,
|
||||
taxonomy_lookup: set[str] | None = None,
|
||||
labels_cfg=None,
|
||||
self, conn, bank_id: str, entities_data: list[dict], unit_event_date
|
||||
) -> list[str]:
|
||||
"""Original strategy: load all bank entities then match in Python."""
|
||||
# Query ALL candidates for this bank
|
||||
@@ -377,24 +338,11 @@ class EntityResolver:
|
||||
all_candidates[entity_text] = matching
|
||||
|
||||
return await self._resolve_from_candidates(
|
||||
conn,
|
||||
bank_id,
|
||||
entities_data,
|
||||
unit_event_date,
|
||||
all_candidates,
|
||||
cooccurrence_map,
|
||||
taxonomy_lookup,
|
||||
labels_cfg,
|
||||
conn, bank_id, entities_data, unit_event_date, all_candidates, cooccurrence_map
|
||||
)
|
||||
|
||||
async def _resolve_entities_batch_trigram(
|
||||
self,
|
||||
conn,
|
||||
bank_id: str,
|
||||
entities_data: list[dict],
|
||||
unit_event_date,
|
||||
taxonomy_lookup: set[str] | None = None,
|
||||
labels_cfg=None,
|
||||
self, conn, bank_id: str, entities_data: list[dict], unit_event_date
|
||||
) -> list[str]:
|
||||
"""
|
||||
Trigram strategy: fetch only similar candidates per entity name using pg_trgm.
|
||||
@@ -405,7 +353,7 @@ class EntityResolver:
|
||||
"""
|
||||
entity_texts = list(set(e["text"] for e in entities_data))
|
||||
|
||||
# Fetch candidates for unique entity texts in bounded batches.
|
||||
# Fetch candidates for all unique entity texts in a single batched query.
|
||||
# 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
|
||||
@@ -413,32 +361,21 @@ class EntityResolver:
|
||||
# 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)
|
||||
rows = 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_texts,
|
||||
)
|
||||
await conn.execute("RESET pg_trgm.similarity_threshold")
|
||||
|
||||
# Group candidates by query_text
|
||||
all_candidates: dict[str, list] = {t: [] for t in entity_texts}
|
||||
@@ -481,24 +418,11 @@ class EntityResolver:
|
||||
cooccurrence_map[eid2].add(id_to_name[eid1])
|
||||
|
||||
return await self._resolve_from_candidates(
|
||||
conn,
|
||||
bank_id,
|
||||
entities_data,
|
||||
unit_event_date,
|
||||
all_candidates,
|
||||
cooccurrence_map,
|
||||
taxonomy_lookup,
|
||||
labels_cfg,
|
||||
conn, bank_id, entities_data, unit_event_date, all_candidates, cooccurrence_map
|
||||
)
|
||||
|
||||
async def _resolve_entities_batch_oracle_fuzzy(
|
||||
self,
|
||||
conn: Any,
|
||||
bank_id: str,
|
||||
entities_data: list[dict],
|
||||
unit_event_date: datetime | None,
|
||||
taxonomy_lookup: set[str] | None = None,
|
||||
labels_cfg=None,
|
||||
self, conn: Any, bank_id: str, entities_data: list[dict], unit_event_date: datetime | None
|
||||
) -> list[str]:
|
||||
"""
|
||||
Oracle strategy: fetch similar candidates using UTL_MATCH.JARO_WINKLER_SIMILARITY.
|
||||
@@ -512,28 +436,23 @@ class EntityResolver:
|
||||
entities_table = fq_table("entities")
|
||||
|
||||
try:
|
||||
# Batch entity texts into bounded sub-queries using JSON_TABLE to
|
||||
# Batch all entity texts into a single query using JSON_TABLE to
|
||||
# expand the list into rows. UTL_MATCH.JARO_WINKLER_SIMILARITY
|
||||
# returns 0-100; threshold 70 ≈ pg_trgm similarity 0.15.
|
||||
# Bounded batches mirror the PG trigram path so very wide retain
|
||||
# batches don't time out a single JOIN on large banks.
|
||||
rows = []
|
||||
for entity_text_batch in self._chunked(entity_texts, self.entity_resolution_batch_size):
|
||||
rows.extend(
|
||||
await conn.fetch(
|
||||
f"""
|
||||
SELECT e.id, e.canonical_name, e.metadata, e.last_seen, e.mention_count,
|
||||
q.query_text
|
||||
FROM JSON_TABLE($2, '$[*]' COLUMNS (query_text VARCHAR2(4000) PATH '$')) q
|
||||
JOIN {entities_table} e ON (
|
||||
e.bank_id = $1
|
||||
AND UTL_MATCH.JARO_WINKLER_SIMILARITY(LOWER(e.canonical_name), LOWER(q.query_text)) > 70
|
||||
)
|
||||
""",
|
||||
bank_id,
|
||||
json.dumps(entity_text_batch),
|
||||
)
|
||||
entity_texts_json = json.dumps(entity_texts)
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT e.id, e.canonical_name, e.metadata, e.last_seen, e.mention_count,
|
||||
q.query_text
|
||||
FROM JSON_TABLE($2, '$[*]' COLUMNS (query_text VARCHAR2(4000) PATH '$')) q
|
||||
JOIN {entities_table} e ON (
|
||||
e.bank_id = $1
|
||||
AND UTL_MATCH.JARO_WINKLER_SIMILARITY(LOWER(e.canonical_name), LOWER(q.query_text)) > 70
|
||||
)
|
||||
""",
|
||||
bank_id,
|
||||
entity_texts_json,
|
||||
)
|
||||
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.
|
||||
@@ -587,14 +506,7 @@ class EntityResolver:
|
||||
cooccurrence_map[eid2].add(id_to_name[eid1])
|
||||
|
||||
return await self._resolve_from_candidates(
|
||||
conn,
|
||||
bank_id,
|
||||
entities_data,
|
||||
unit_event_date,
|
||||
all_candidates,
|
||||
cooccurrence_map,
|
||||
taxonomy_lookup,
|
||||
labels_cfg,
|
||||
conn, bank_id, entities_data, unit_event_date, all_candidates, cooccurrence_map
|
||||
)
|
||||
|
||||
async def _resolve_from_candidates(
|
||||
@@ -605,8 +517,6 @@ class EntityResolver:
|
||||
unit_event_date,
|
||||
all_candidates: dict[str, list],
|
||||
cooccurrence_map: dict[str, set[str]],
|
||||
taxonomy_lookup: set[str] | None = None,
|
||||
labels_cfg=None,
|
||||
) -> list[str]:
|
||||
"""Shared scoring + upsert logic used by both lookup strategies."""
|
||||
|
||||
@@ -623,34 +533,11 @@ class EntityResolver:
|
||||
|
||||
candidates = all_candidates.get(entity_text, [])
|
||||
|
||||
# Label entities (from entity_labels config) use exact matching only.
|
||||
# Their canonical names are user-defined (e.g., "use:use-001"),
|
||||
# so fuzzy resolution must NOT merge distinct label values that
|
||||
# happen to be textually similar (GH-1558).
|
||||
is_label = bool(
|
||||
labels_cfg and taxonomy_lookup and _is_label_entity(entity_text, labels_cfg, taxonomy_lookup)
|
||||
)
|
||||
|
||||
if not candidates:
|
||||
# Will create new entity
|
||||
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 = None
|
||||
entity_text_lower = entity_text.lower()
|
||||
for candidate_id, canonical_name, metadata, last_seen, mention_count in candidates:
|
||||
if canonical_name.lower() == entity_text_lower:
|
||||
exact_match = candidate_id
|
||||
break
|
||||
if exact_match:
|
||||
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))
|
||||
continue
|
||||
|
||||
# Score candidates
|
||||
best_candidate = None
|
||||
best_score = 0.0
|
||||
@@ -782,6 +669,239 @@ class EntityResolver:
|
||||
|
||||
return entity_ids
|
||||
|
||||
async def resolve_entity(
|
||||
self,
|
||||
bank_id: str,
|
||||
entity_text: str,
|
||||
context: str,
|
||||
nearby_entities: list[dict],
|
||||
unit_event_date,
|
||||
) -> str:
|
||||
"""
|
||||
Resolve an entity to a canonical entity ID.
|
||||
|
||||
Args:
|
||||
bank_id: bank ID (entities are scoped to agents)
|
||||
entity_text: Entity text ("Alice", "Google", etc.)
|
||||
context: Context where entity appears
|
||||
nearby_entities: Other entities in the same unit
|
||||
unit_event_date: When this unit was created
|
||||
|
||||
Returns:
|
||||
Entity ID (creates new entity if needed)
|
||||
"""
|
||||
async with acquire_with_retry(self.pool) as conn:
|
||||
# Find candidate entities with similar name
|
||||
candidates = await conn.fetch(
|
||||
f"""
|
||||
SELECT id, canonical_name, metadata, last_seen
|
||||
FROM {fq_table("entities")}
|
||||
WHERE bank_id = $1
|
||||
AND (
|
||||
canonical_name ILIKE $2
|
||||
OR canonical_name ILIKE $3
|
||||
OR $2 ILIKE canonical_name || '%%'
|
||||
)
|
||||
ORDER BY mention_count DESC
|
||||
""",
|
||||
bank_id,
|
||||
entity_text,
|
||||
f"%{entity_text}%",
|
||||
)
|
||||
|
||||
if not candidates:
|
||||
# New entity - create it
|
||||
return await self._create_entity(conn, bank_id, entity_text, unit_event_date)
|
||||
|
||||
# Score candidates based on:
|
||||
# 1. Name similarity
|
||||
# 2. Context overlap (TODO: could use embeddings)
|
||||
# 3. Co-occurring entities
|
||||
# 4. Temporal proximity
|
||||
|
||||
best_candidate = None
|
||||
best_score = 0.0
|
||||
best_name_similarity = 0.0
|
||||
|
||||
nearby_entity_set = {e["text"].lower() for e in nearby_entities if e["text"] != entity_text}
|
||||
|
||||
for row in candidates:
|
||||
candidate_id = row["id"]
|
||||
canonical_name = row["canonical_name"]
|
||||
metadata = row["metadata"]
|
||||
last_seen = row["last_seen"]
|
||||
score = 0.0
|
||||
|
||||
# 1. Name similarity (0-1)
|
||||
name_similarity = SequenceMatcher(None, entity_text.lower(), canonical_name.lower()).ratio()
|
||||
score += name_similarity * 0.5
|
||||
|
||||
# 2. Co-occurring entities (0-0.5)
|
||||
# Get entities that co-occurred with this candidate before
|
||||
# Use the materialized co-occurrence cache for fast lookup
|
||||
co_entity_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT e.canonical_name, ec.cooccurrence_count
|
||||
FROM {fq_table("entity_cooccurrences")} ec
|
||||
JOIN {fq_table("entities")} e ON (
|
||||
CASE
|
||||
WHEN ec.entity_id_1 = $1 THEN ec.entity_id_2
|
||||
WHEN ec.entity_id_2 = $1 THEN ec.entity_id_1
|
||||
END = e.id
|
||||
)
|
||||
WHERE ec.entity_id_1 = $1 OR ec.entity_id_2 = $1
|
||||
""",
|
||||
candidate_id,
|
||||
)
|
||||
co_entities = {r["canonical_name"].lower() for r in co_entity_rows}
|
||||
|
||||
# Check overlap with nearby entities
|
||||
overlap = len(nearby_entity_set & co_entities)
|
||||
if nearby_entity_set:
|
||||
co_entity_score = overlap / len(nearby_entity_set)
|
||||
score += co_entity_score * 0.3
|
||||
|
||||
# 3. Temporal proximity (0-0.2)
|
||||
if last_seen:
|
||||
# Normalize both to UTC-aware to avoid naive/aware mismatch
|
||||
# (Oracle returns naive datetimes from fromisoformat)
|
||||
_evt = unit_event_date if unit_event_date.tzinfo else unit_event_date.replace(tzinfo=UTC)
|
||||
_seen = last_seen if last_seen.tzinfo else last_seen.replace(tzinfo=UTC)
|
||||
days_diff = abs((_evt - _seen).total_seconds() / 86400)
|
||||
if days_diff < 7: # Within a week
|
||||
temporal_score = max(0, 1.0 - (days_diff / 7))
|
||||
score += temporal_score * 0.2
|
||||
|
||||
if score > best_score:
|
||||
best_score = score
|
||||
best_candidate = candidate_id
|
||||
best_name_similarity = name_similarity
|
||||
|
||||
# Threshold for considering it the same entity
|
||||
threshold = 0.6
|
||||
|
||||
if best_score > threshold:
|
||||
# Update entity
|
||||
await conn.execute(
|
||||
f"""
|
||||
UPDATE {fq_table("entities")}
|
||||
SET mention_count = mention_count + 1,
|
||||
last_seen = $1
|
||||
WHERE id = $2
|
||||
""",
|
||||
unit_event_date,
|
||||
best_candidate,
|
||||
)
|
||||
return best_candidate
|
||||
else:
|
||||
# Not confident - create new entity
|
||||
return await self._create_entity(conn, bank_id, entity_text, unit_event_date)
|
||||
|
||||
async def _create_entity(
|
||||
self,
|
||||
conn,
|
||||
bank_id: str,
|
||||
entity_text: str,
|
||||
event_date,
|
||||
) -> str:
|
||||
"""
|
||||
Create a new entity or get existing one if it already exists.
|
||||
|
||||
Uses INSERT ... ON CONFLICT to handle race conditions where
|
||||
two concurrent transactions try to create the same entity.
|
||||
|
||||
Args:
|
||||
conn: Database connection
|
||||
bank_id: bank ID
|
||||
entity_text: Entity text
|
||||
event_date: When first seen
|
||||
|
||||
Returns:
|
||||
Entity ID
|
||||
"""
|
||||
entity_id = await conn.fetchval(
|
||||
f"""
|
||||
INSERT INTO {fq_table("entities")} (bank_id, canonical_name, first_seen, last_seen, mention_count)
|
||||
VALUES ($1, $2, COALESCE($3, now()), COALESCE($4, now()), 1)
|
||||
ON CONFLICT (bank_id, LOWER(canonical_name))
|
||||
DO UPDATE SET
|
||||
mention_count = {fq_table("entities")}.mention_count + 1,
|
||||
last_seen = EXCLUDED.last_seen
|
||||
RETURNING id
|
||||
""",
|
||||
bank_id,
|
||||
entity_text,
|
||||
event_date,
|
||||
event_date,
|
||||
)
|
||||
return entity_id
|
||||
|
||||
async def link_unit_to_entity(self, unit_id: str, entity_id: str):
|
||||
"""
|
||||
Link a memory unit to an entity.
|
||||
Also updates co-occurrence cache with other entities in the same unit.
|
||||
|
||||
Args:
|
||||
unit_id: Memory unit ID
|
||||
entity_id: Entity ID
|
||||
"""
|
||||
async with acquire_with_retry(self.pool) as conn:
|
||||
# Insert unit-entity link
|
||||
await conn.execute(
|
||||
f"""
|
||||
INSERT INTO {fq_table("unit_entities")} (unit_id, entity_id)
|
||||
VALUES ($1, $2)
|
||||
ON CONFLICT DO NOTHING
|
||||
""",
|
||||
unit_id,
|
||||
entity_id,
|
||||
)
|
||||
|
||||
# Update co-occurrence cache: find other entities in this unit
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT entity_id
|
||||
FROM {fq_table("unit_entities")}
|
||||
WHERE unit_id = $1 AND entity_id != $2
|
||||
""",
|
||||
unit_id,
|
||||
entity_id,
|
||||
)
|
||||
|
||||
other_entities = [row["entity_id"] for row in rows]
|
||||
|
||||
# Update co-occurrences for each pair
|
||||
for other_entity_id in other_entities:
|
||||
await self._update_cooccurrence(conn, entity_id, other_entity_id)
|
||||
|
||||
async def _update_cooccurrence(self, conn, entity_id_1: str, entity_id_2: str):
|
||||
"""
|
||||
Update the co-occurrence cache for two entities.
|
||||
|
||||
Uses CHECK constraint ordering (entity_id_1 < entity_id_2) to avoid duplicates.
|
||||
|
||||
Args:
|
||||
conn: Database connection
|
||||
entity_id_1: First entity ID
|
||||
entity_id_2: Second entity ID
|
||||
"""
|
||||
# Ensure consistent ordering (smaller UUID first)
|
||||
if entity_id_1 > entity_id_2:
|
||||
entity_id_1, entity_id_2 = entity_id_2, entity_id_1
|
||||
|
||||
await conn.execute(
|
||||
f"""
|
||||
INSERT INTO {fq_table("entity_cooccurrences")} (entity_id_1, entity_id_2, cooccurrence_count, last_cooccurred)
|
||||
VALUES ($1, $2, 1, NOW())
|
||||
ON CONFLICT (entity_id_1, entity_id_2)
|
||||
DO UPDATE SET
|
||||
cooccurrence_count = {fq_table("entity_cooccurrences")}.cooccurrence_count + 1,
|
||||
last_cooccurred = NOW()
|
||||
""",
|
||||
entity_id_1,
|
||||
entity_id_2,
|
||||
)
|
||||
|
||||
async def link_units_to_entities_batch(
|
||||
self,
|
||||
unit_entity_pairs: list[tuple[str, str]] | list[tuple[str, str, datetime | None]],
|
||||
|
||||
@@ -1,357 +0,0 @@
|
||||
"""Async graph maintenance after document/unit deletes.
|
||||
|
||||
Three reconciliation passes run together on every worker invocation:
|
||||
|
||||
1. **Relink top-up.** Drain ``graph_maintenance_queue`` (units whose
|
||||
outgoing temporal/semantic links lost a neighbour to a delete). For
|
||||
each, count current outgoing links per type; if below cap, run the
|
||||
same probes retain uses (:func:`fetch_temporal_neighbors`,
|
||||
:func:`compute_semantic_links_ann`) and insert the missing links.
|
||||
``bulk_insert_links`` has ``ON CONFLICT DO NOTHING`` on the uniqueness
|
||||
key, so we can re-probe freely and the DB de-dupes.
|
||||
|
||||
2. **Orphan entity prune.** Delete ``entities`` rows in the bank that no
|
||||
longer have any ``unit_entities`` references. FK ON DELETE CASCADE on
|
||||
``entity_cooccurrences`` then removes any cooccurrence row pointing
|
||||
at the pruned entities.
|
||||
|
||||
3. **Stale cooccurrence prune.** Defensive sweep for cooccurrence rows
|
||||
where both endpoints still exist but no current memory_unit references
|
||||
both of them — the cooccurrence was real at the time it was recorded,
|
||||
but every unit that witnessed it has since been deleted.
|
||||
|
||||
All three passes run on every invocation. The queue is the only source
|
||||
of work for pass 1; passes 2 and 3 are bank-wide sweeps backed by indexes
|
||||
on ``entities(bank_id)`` and ``unit_entities(entity_id)``, so they're
|
||||
cheap when there's nothing to do.
|
||||
|
||||
The worker dedupes on bank: a second job for the same bank is dropped
|
||||
while one is pending. Once processing starts, a new job becomes the
|
||||
*next* pending slot — so work enqueued during processing gets picked up
|
||||
by the follow-up run.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import time
|
||||
import uuid as uuid_module
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from ..models import RequestContext
|
||||
from .db.base import DatabaseConnection
|
||||
from .retain.link_utils import (
|
||||
MAX_TEMPORAL_LINKS_PER_UNIT,
|
||||
_bulk_insert_links,
|
||||
_normalize_datetime,
|
||||
compute_semantic_links_ann,
|
||||
)
|
||||
from .schema import fq_table
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .memory_engine import MemoryEngine
|
||||
|
||||
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
|
||||
|
||||
|
||||
@dataclass
|
||||
class JobResult:
|
||||
"""Counters surfaced to the worker dispatcher and operation result."""
|
||||
|
||||
relink_units_processed: int = 0
|
||||
relink_links_added: int = 0
|
||||
orphan_entities_pruned: int = 0
|
||||
stale_cooccurrences_pruned: int = 0
|
||||
|
||||
def as_dict(self) -> dict[str, int]:
|
||||
return {
|
||||
"relink_units_processed": self.relink_units_processed,
|
||||
"relink_links_added": self.relink_links_added,
|
||||
"orphan_entities_pruned": self.orphan_entities_pruned,
|
||||
"stale_cooccurrences_pruned": self.stale_cooccurrences_pruned,
|
||||
}
|
||||
|
||||
|
||||
async def enqueue_relink_victims(
|
||||
conn: DatabaseConnection,
|
||||
bank_id: str,
|
||||
deleted_unit_ids: list[str],
|
||||
ops: Any,
|
||||
) -> int:
|
||||
"""Enqueue surviving units whose outgoing temporal/semantic links pointed at
|
||||
``deleted_unit_ids`` for later link top-up.
|
||||
|
||||
Must run inside the same transaction that deletes the units, *before* the
|
||||
cascade fires — once the rows are gone, the join that finds the victims
|
||||
returns nothing.
|
||||
|
||||
Args:
|
||||
conn: Database connection inside the active delete transaction.
|
||||
bank_id: Bank owning the deleted units.
|
||||
deleted_unit_ids: Memory_unit IDs about to be (or being) deleted.
|
||||
ops: ``DataAccessOps`` instance, supplies the dialect-specific
|
||||
bulk-insert path.
|
||||
|
||||
Returns:
|
||||
Number of distinct victim units enqueued (after dedup against rows
|
||||
already in the queue).
|
||||
"""
|
||||
if not deleted_unit_ids:
|
||||
return 0
|
||||
|
||||
deleted_uuids = [uuid_module.UUID(uid) if isinstance(uid, str) else uid for uid in deleted_unit_ids]
|
||||
deleted_str_set = {str(uid) for uid in deleted_uuids}
|
||||
|
||||
# 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",
|
||||
bank_id: str,
|
||||
request_context: RequestContext,
|
||||
operation_id: str | None = None,
|
||||
) -> dict[str, int]:
|
||||
"""Run all maintenance passes for ``bank_id`` until the relink queue is
|
||||
drained, then sweep entities and cooccurrences once.
|
||||
|
||||
Returns:
|
||||
Per-pass counters from :class:`JobResult`.
|
||||
"""
|
||||
del request_context # accepted for symmetry with other run_*_job helpers
|
||||
backend = await memory_engine._get_backend()
|
||||
ops = backend.ops
|
||||
|
||||
result = JobResult()
|
||||
job_start = time.time()
|
||||
|
||||
# --- Pass 1: relink ---
|
||||
# Per-iteration loop: claim → top up → commit. We rely on submit-time
|
||||
# dedup to keep at most one job per bank running, so no need for
|
||||
# SKIP LOCKED.
|
||||
iterations = 0
|
||||
while True:
|
||||
from .memory_engine import acquire_with_retry
|
||||
|
||||
async with acquire_with_retry(backend) as conn:
|
||||
async with conn.transaction():
|
||||
unit_ids = await ops.claim_graph_maintenance_batch(
|
||||
conn,
|
||||
fq_table("graph_maintenance_queue"),
|
||||
bank_id,
|
||||
_DRAIN_BATCH_SIZE,
|
||||
)
|
||||
if not unit_ids:
|
||||
break
|
||||
|
||||
result.relink_links_added += await _relink_batch(conn, bank_id, unit_ids, ops, backend)
|
||||
|
||||
result.relink_units_processed += len(unit_ids)
|
||||
iterations += 1
|
||||
|
||||
if iterations > 10000:
|
||||
# Defensive guard against runaway loops — at 50 units/iter that's
|
||||
# 500k targets, far beyond any realistic single-bank backlog.
|
||||
logger.error(
|
||||
f"[GRAPH_MAINT] bank={bank_id} hit iteration cap ({iterations}); aborting relink ({result.as_dict()})"
|
||||
)
|
||||
break
|
||||
|
||||
# --- Pass 2 & 3: entity / cooccurrence sweeps ---
|
||||
# Bank-wide single-statement deletes. Cheap when there's nothing to do.
|
||||
from .memory_engine import acquire_with_retry
|
||||
|
||||
async with acquire_with_retry(backend) as conn:
|
||||
async with conn.transaction():
|
||||
result.orphan_entities_pruned = await ops.prune_orphan_entities(
|
||||
conn,
|
||||
fq_table("entities"),
|
||||
fq_table("unit_entities"),
|
||||
bank_id,
|
||||
)
|
||||
# The orphan prune above cascades cooccurrences via FK. The
|
||||
# explicit cooccurrence pass below catches the *stale-count*
|
||||
# case: both entities still exist but no current unit witnesses
|
||||
# them together.
|
||||
result.stale_cooccurrences_pruned = await ops.prune_stale_cooccurrences(
|
||||
conn,
|
||||
fq_table("entity_cooccurrences"),
|
||||
fq_table("unit_entities"),
|
||||
fq_table("entities"),
|
||||
bank_id,
|
||||
)
|
||||
|
||||
elapsed = time.time() - job_start
|
||||
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)
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user