Compare commits

..
781 changed files with 9340 additions and 68782 deletions
+2 -10
View File
@@ -166,14 +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.
### 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)
@@ -185,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:
@@ -197,7 +190,6 @@ Present a clear summary organized by severity:
- Multi-item tuple returns (including internal code)
- Missing tests for new endpoints
- New integration missing tests, CI job, or release-integration.sh entry
- 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 -28
View File
@@ -7,8 +7,6 @@ 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
# Example: Anthropic Claude configuration
# HINDSIGHT_API_LLM_PROVIDER=anthropic
@@ -66,21 +64,8 @@ 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=
# Embeddings Configuration (Optional - uses local by default)
# Provider: "local" (default), "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
@@ -92,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}
@@ -137,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).
+8 -325
View File
@@ -12,7 +12,6 @@ concurrency:
jobs:
detect-changes:
runs-on: ubuntu-latest
timeout-minutes: 30
permissions:
pull-requests: read
outputs:
@@ -50,11 +49,7 @@ jobs:
integrations-agentcore: ${{ steps.filter.outputs.integrations-agentcore }}
integrations-smolagents: ${{ steps.filter.outputs.integrations-smolagents }}
integrations-dify: ${{ steps.filter.outputs.integrations-dify }}
integrations-gemini-spark: ${{ steps.filter.outputs.integrations-gemini-spark }}
integrations-vapi: ${{ steps.filter.outputs.integrations-vapi }}
integrations-flowise: ${{ steps.filter.outputs.integrations-flowise }}
tools-agent-sdk: ${{ steps.filter.outputs.tools-agent-sdk }}
integrations-roo-code: ${{ steps.filter.outputs.integrations-roo-code }}
dev: ${{ steps.filter.outputs.dev }}
ci: ${{ steps.filter.outputs.ci }}
# Secrets are available for internal PRs and workflow_dispatch.
@@ -148,16 +143,8 @@ jobs:
- 'hindsight-integrations/smolagents/**'
integrations-dify:
- 'hindsight-integrations/dify/**'
integrations-gemini-spark:
- 'hindsight-integrations/gemini-spark/**'
integrations-vapi:
- 'hindsight-integrations/vapi/**'
integrations-flowise:
- 'hindsight-integrations/flowise/**'
tools-agent-sdk:
- 'hindsight-tools/hindsight-agent-sdk/**'
integrations-roo-code:
- 'hindsight-integrations/roo-code/**'
dev:
- 'hindsight-dev/**'
ci:
@@ -177,7 +164,6 @@ jobs:
needs.detect-changes.outputs.integrations-lockfiles == 'true' ||
needs.detect-changes.outputs.ci == 'true')
runs-on: ubuntu-latest
timeout-minutes: 30
steps:
- uses: actions/checkout@v6
with:
@@ -193,7 +179,6 @@ jobs:
needs.detect-changes.outputs.core == 'true' ||
needs.detect-changes.outputs.ci == 'true')
runs-on: ubuntu-latest
timeout-minutes: 30
strategy:
matrix:
python-version: ['3.11', '3.12', '3.13', '3.14']
@@ -224,7 +209,6 @@ jobs:
needs.detect-changes.outputs.clients-ts == 'true' ||
needs.detect-changes.outputs.ci == 'true')
runs-on: ubuntu-latest
timeout-minutes: 30
steps:
- uses: actions/checkout@v6
@@ -251,7 +235,6 @@ jobs:
needs.detect-changes.outputs.all-npm == 'true' ||
needs.detect-changes.outputs.ci == 'true')
runs-on: ubuntu-latest
timeout-minutes: 30
steps:
- uses: actions/checkout@v6
@@ -283,7 +266,6 @@ jobs:
needs.detect-changes.outputs.all-npm == 'true' ||
needs.detect-changes.outputs.ci == 'true')
runs-on: ubuntu-latest
timeout-minutes: 30
steps:
- uses: actions/checkout@v6
@@ -339,7 +321,6 @@ jobs:
needs.detect-changes.outputs.all-npm == 'true' ||
needs.detect-changes.outputs.ci == 'true')
runs-on: ubuntu-latest
timeout-minutes: 30
steps:
- uses: actions/checkout@v6
@@ -395,7 +376,6 @@ jobs:
needs.detect-changes.outputs.integrations-claude-code == 'true' ||
needs.detect-changes.outputs.ci == 'true')
runs-on: ubuntu-latest
timeout-minutes: 30
steps:
- uses: actions/checkout@v6
@@ -421,7 +401,6 @@ jobs:
needs.detect-changes.outputs.integrations-codex == 'true' ||
needs.detect-changes.outputs.ci == 'true')
runs-on: ubuntu-latest
timeout-minutes: 30
steps:
- uses: actions/checkout@v6
@@ -447,7 +426,6 @@ jobs:
needs.detect-changes.outputs.integrations-ai-sdk == 'true' ||
needs.detect-changes.outputs.ci == 'true')
runs-on: ubuntu-latest
timeout-minutes: 30
steps:
- uses: actions/checkout@v6
@@ -478,7 +456,6 @@ jobs:
needs.detect-changes.outputs.integrations-ai-sdk == 'true' ||
needs.detect-changes.outputs.ci == 'true')
runs-on: ubuntu-latest
timeout-minutes: 30
steps:
- uses: actions/checkout@v6
@@ -510,7 +487,6 @@ jobs:
needs.detect-changes.outputs.integrations-opencode == 'true' ||
needs.detect-changes.outputs.ci == 'true')
runs-on: ubuntu-latest
timeout-minutes: 30
steps:
- uses: actions/checkout@v6
@@ -541,7 +517,6 @@ jobs:
needs.detect-changes.outputs.integrations-n8n == 'true' ||
needs.detect-changes.outputs.ci == 'true')
runs-on: ubuntu-latest
timeout-minutes: 30
steps:
- uses: actions/checkout@v6
@@ -572,7 +547,6 @@ jobs:
needs.detect-changes.outputs.tools-agent-sdk == 'true' ||
needs.detect-changes.outputs.ci == 'true')
runs-on: ubuntu-latest
timeout-minutes: 30
steps:
- uses: actions/checkout@v6
@@ -605,7 +579,6 @@ jobs:
needs.detect-changes.outputs.integrations-cloudflare-oauth-proxy == 'true' ||
needs.detect-changes.outputs.ci == 'true')
runs-on: ubuntu-latest
timeout-minutes: 30
steps:
- uses: actions/checkout@v6
@@ -636,7 +609,6 @@ jobs:
needs.detect-changes.outputs.integrations-chat == 'true' ||
needs.detect-changes.outputs.ci == 'true')
runs-on: ubuntu-latest
timeout-minutes: 30
steps:
- uses: actions/checkout@v6
@@ -667,7 +639,6 @@ jobs:
needs.detect-changes.outputs.integrations-paperclip == 'true' ||
needs.detect-changes.outputs.ci == 'true')
runs-on: ubuntu-latest
timeout-minutes: 30
steps:
- uses: actions/checkout@v6
@@ -698,7 +669,6 @@ jobs:
needs.detect-changes.outputs.integrations-pipecat == 'true' ||
needs.detect-changes.outputs.ci == 'true')
runs-on: ubuntu-latest
timeout-minutes: 30
steps:
- uses: actions/checkout@v6
@@ -728,65 +698,6 @@ jobs:
working-directory: ./hindsight-integrations/pipecat
run: uv run pytest tests -v
test-gemini-spark-integration:
needs: [detect-changes]
if: >-
(github.event_name == 'workflow_dispatch' ||
needs.detect-changes.outputs.integrations-gemini-spark == 'true' ||
needs.detect-changes.outputs.ci == 'true')
runs-on: ubuntu-latest
timeout-minutes: 30
steps:
- uses: actions/checkout@v6
with:
ref: ${{ github.event.pull_request.head.sha || '' }}
- name: Install uv
uses: astral-sh/setup-uv@v7
with:
enable-cache: true
prune-cache: false
- name: Set up Python
uses: actions/setup-python@v6
with:
python-version-file: ".python-version"
- name: Install dependencies
working-directory: ./hindsight-integrations/gemini-spark
run: uv sync --frozen
- name: Run tests
working-directory: ./hindsight-integrations/gemini-spark
run: uv run pytest tests -v
test-roo-code-integration:
needs: [detect-changes]
if: >-
github.event_name != 'pull_request_review' &&
(github.event_name == 'workflow_dispatch' ||
needs.detect-changes.outputs.integrations-roo-code == 'true' ||
needs.detect-changes.outputs.ci == 'true')
runs-on: ubuntu-latest
timeout-minutes: 30
steps:
- uses: actions/checkout@v6
with:
ref: ${{ github.event.pull_request.head.sha || '' }}
- name: Set up Python
uses: actions/setup-python@v6
with:
python-version: '3.11'
- name: Install pytest
run: pip install pytest
- name: Run tests
working-directory: ./hindsight-integrations/roo-code
run: python -m pytest tests/ -v
build-control-plane:
needs: [detect-changes]
@@ -796,7 +707,6 @@ jobs:
needs.detect-changes.outputs.clients-ts == 'true' ||
needs.detect-changes.outputs.ci == 'true')
runs-on: ubuntu-latest
timeout-minutes: 30
steps:
- uses: actions/checkout@v6
@@ -824,12 +734,6 @@ jobs:
rm -rf node_modules/lightningcss node_modules/@tailwindcss
npm install lightningcss @tailwindcss/postcss @tailwindcss/node
- name: Test Control Plane
run: npm test --workspace=hindsight-control-plane
- name: Check i18n locale parity and hardcoded strings
run: npm run i18n:check --workspace=hindsight-control-plane
- name: Build Control Plane
run: npm run build --workspace=hindsight-control-plane
@@ -862,7 +766,6 @@ jobs:
needs.detect-changes.outputs.docs == 'true' ||
needs.detect-changes.outputs.ci == 'true')
runs-on: ubuntu-latest
timeout-minutes: 30
steps:
- uses: actions/checkout@v6
@@ -891,7 +794,6 @@ jobs:
needs.detect-changes.outputs.cli == 'true' ||
needs.detect-changes.outputs.ci == 'true')
runs-on: ubuntu-latest
timeout-minutes: 30
env:
HINDSIGHT_API_LLM_PROVIDER: vertexai
HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY: /tmp/gcp-credentials.json
@@ -1019,7 +921,6 @@ jobs:
needs.detect-changes.outputs.helm == 'true' ||
needs.detect-changes.outputs.ci == 'true')
runs-on: ubuntu-latest
timeout-minutes: 30
steps:
- uses: actions/checkout@v6
@@ -1034,23 +935,6 @@ jobs:
- name: Lint Helm chart
run: helm lint helm/hindsight
test-standalone-start-script:
needs: [detect-changes]
if: >-
(github.event_name == 'workflow_dispatch' ||
needs.detect-changes.outputs.docker == 'true' ||
needs.detect-changes.outputs.ci == 'true')
runs-on: ubuntu-latest
timeout-minutes: 30
steps:
- uses: actions/checkout@v6
with:
ref: ${{ github.event.pull_request.head.sha || '' }}
- name: Run standalone start script tests
run: bash docker/standalone/test-start-all.sh
build-docker-images:
needs: [detect-changes]
if: >-
@@ -1062,7 +946,6 @@ jobs:
needs.detect-changes.outputs.ci == 'true')
name: Build Docker (${{ matrix.name }})
runs-on: ubuntu-latest
timeout-minutes: 30
strategy:
matrix:
include:
@@ -1153,16 +1036,6 @@ jobs:
needs.detect-changes.outputs.core == 'true' ||
needs.detect-changes.outputs.ci == 'true')
runs-on: ubuntu-latest
timeout-minutes: 30
# Sharded with pytest-split. Splits the ~19-min pytest run across 3
# parallel jobs, cutting critical-path wall time to ~7-8 min/shard. The
# .venv cache lets shards 2+ skip the ~3-min `uv sync` once shard 1
# populates the key — same key on re-runs hits cache on all three.
strategy:
fail-fast: false
matrix:
shard: [1, 2, 3]
name: test-api (${{ matrix.shard }}/3)
env:
HINDSIGHT_API_LLM_PROVIDER: vertexai
HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY: /tmp/gcp-credentials.json
@@ -1203,18 +1076,6 @@ jobs:
working-directory: ./hindsight-api-slim
run: uv build
- name: Cache hindsight-api-slim/.venv
# Keyed on the workspace lockfile, the API package's pyproject, and the
# pinned Python version — the only inputs that change the resolved env.
# `uv sync --frozen` still runs after restore but is a near-instant link
# check when the venv already matches the lock.
uses: actions/cache@v5
with:
path: hindsight-api-slim/.venv
key: ${{ runner.os }}-venv-test-api-${{ hashFiles('uv.lock', 'hindsight-api-slim/pyproject.toml', '.python-version') }}
restore-keys: |
${{ runner.os }}-venv-test-api-
- name: Install dependencies
working-directory: ./hindsight-api-slim
run: uv sync --frozen --all-extras --index-strategy unsafe-best-match
@@ -1239,76 +1100,9 @@ jobs:
print('Models downloaded successfully')
"
- name: Run tests (shard ${{ matrix.shard }}/3)
- name: Run tests
working-directory: ./hindsight-api-slim
# `--with pytest-split` adds the plugin ad-hoc — no uv.lock churn.
# pytest-split filters at collection (before xdist takes over), so it
# composes cleanly with the `-n 8 --dist loadgroup` baked into addopts.
run: uv run --with pytest-split pytest tests -v -m "not hs_llm_mat and not hs_llm_core" --splits 3 --group ${{ matrix.shard }}
test-api-llm-core:
needs: [detect-changes]
if: >-
needs.detect-changes.outputs.has_secrets == 'true' &&
(github.event_name == 'workflow_dispatch' ||
needs.detect-changes.outputs.core == 'true' ||
needs.detect-changes.outputs.ci == 'true')
runs-on: ubuntu-latest
timeout-minutes: 30
name: Core LLM tests
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
GEMINI_API_KEY: ${{ secrets.GEMINI_API_KEY }}
steps:
- uses: actions/checkout@v6
with:
ref: ${{ github.event.pull_request.head.sha || '' }}
- 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: Install dependencies
working-directory: ./hindsight-api-slim
run: uv sync --frozen --all-extras --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-
- name: Pre-download models
working-directory: ./hindsight-api-slim
run: |
uv run 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 successfully')
"
- name: Run core LLM tests
working-directory: ./hindsight-api-slim
run: uv run pytest tests -v -m "hs_llm_core" --timeout 600
run: uv run pytest tests -v -m "not hs_llm_mat"
test-api-llm-acceptance:
needs: [detect-changes]
@@ -1319,7 +1113,6 @@ jobs:
needs.detect-changes.outputs.core == 'true' ||
needs.detect-changes.outputs.ci == 'true')
runs-on: ubuntu-latest
timeout-minutes: 30
strategy:
fail-fast: false
matrix:
@@ -1426,7 +1219,6 @@ jobs:
needs.detect-changes.outputs.core == 'true' ||
needs.detect-changes.outputs.ci == 'true')
runs-on: ubuntu-latest
timeout-minutes: 30
env:
HINDSIGHT_API_LLM_PROVIDER: vertexai
HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY: /tmp/gcp-credentials.json
@@ -1468,8 +1260,8 @@ jobs:
conn = oracledb.connect(user='system', password='oracle', dsn='localhost:1521/FREEPDB1')
cursor = conn.cursor()
cursor.execute(\"\"\"
CREATE BIGFILE TABLESPACE hindsight_ts
DATAFILE 'hindsight_ts.dbf' SIZE 2G AUTOEXTEND ON NEXT 500M MAXSIZE UNLIMITED
CREATE TABLESPACE hindsight_ts
DATAFILE 'hindsight_ts.dbf' SIZE 200M AUTOEXTEND ON NEXT 50M
EXTENT MANAGEMENT LOCAL
SEGMENT SPACE MANAGEMENT AUTO
\"\"\")
@@ -1546,7 +1338,6 @@ jobs:
needs.detect-changes.outputs.clients-python == 'true' ||
needs.detect-changes.outputs.ci == 'true')
runs-on: ubuntu-latest
timeout-minutes: 30
env:
HINDSIGHT_API_LLM_PROVIDER: vertexai
HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY: /tmp/gcp-credentials.json
@@ -1658,7 +1449,6 @@ jobs:
needs.detect-changes.outputs.clients-ts == 'true' ||
needs.detect-changes.outputs.ci == 'true')
runs-on: ubuntu-latest
timeout-minutes: 30
env:
HINDSIGHT_API_LLM_PROVIDER: vertexai
HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY: /tmp/gcp-credentials.json
@@ -1788,7 +1578,6 @@ jobs:
needs.detect-changes.outputs.clients-python == 'true' ||
needs.detect-changes.outputs.ci == 'true')
runs-on: ubuntu-latest
timeout-minutes: 30
env:
HINDSIGHT_API_LLM_PROVIDER: vertexai
HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY: /tmp/gcp-credentials.json
@@ -1825,8 +1614,8 @@ jobs:
conn = oracledb.connect(user='system', password='oracle', dsn='localhost:1521/FREEPDB1')
cursor = conn.cursor()
cursor.execute(\"\"\"
CREATE BIGFILE TABLESPACE hindsight_ts
DATAFILE 'hindsight_ts.dbf' SIZE 2G AUTOEXTEND ON NEXT 500M MAXSIZE UNLIMITED
CREATE TABLESPACE hindsight_ts
DATAFILE 'hindsight_ts.dbf' SIZE 200M AUTOEXTEND ON NEXT 50M
EXTENT MANAGEMENT LOCAL
SEGMENT SPACE MANAGEMENT AUTO
\"\"\")
@@ -1948,7 +1737,6 @@ jobs:
needs.detect-changes.outputs.clients-ts == 'true' ||
needs.detect-changes.outputs.ci == 'true')
runs-on: ubuntu-latest
timeout-minutes: 30
env:
HINDSIGHT_API_LLM_PROVIDER: vertexai
HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY: /tmp/gcp-credentials.json
@@ -1985,8 +1773,8 @@ jobs:
conn = oracledb.connect(user='system', password='oracle', dsn='localhost:1521/FREEPDB1')
cursor = conn.cursor()
cursor.execute(\"\"\"
CREATE BIGFILE TABLESPACE hindsight_ts
DATAFILE 'hindsight_ts.dbf' SIZE 2G AUTOEXTEND ON NEXT 500M MAXSIZE UNLIMITED
CREATE TABLESPACE hindsight_ts
DATAFILE 'hindsight_ts.dbf' SIZE 200M AUTOEXTEND ON NEXT 50M
EXTENT MANAGEMENT LOCAL
SEGMENT SPACE MANAGEMENT AUTO
\"\"\")
@@ -2108,7 +1896,6 @@ jobs:
needs.detect-changes.outputs.clients-ts == 'true' ||
needs.detect-changes.outputs.ci == 'true')
runs-on: ubuntu-latest
timeout-minutes: 30
env:
HINDSIGHT_API_LLM_PROVIDER: vertexai
HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY: /tmp/gcp-credentials.json
@@ -2229,7 +2016,6 @@ jobs:
needs.detect-changes.outputs.cli == 'true' ||
needs.detect-changes.outputs.ci == 'true')
runs-on: ubuntu-24.04-arm
timeout-minutes: 30
steps:
- uses: actions/checkout@v6
@@ -2263,7 +2049,6 @@ jobs:
needs.detect-changes.outputs.clients-rust == 'true' ||
needs.detect-changes.outputs.ci == 'true')
runs-on: ubuntu-latest
timeout-minutes: 30
env:
HINDSIGHT_API_LLM_PROVIDER: vertexai
HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY: /tmp/gcp-credentials.json
@@ -2379,7 +2164,6 @@ jobs:
needs.detect-changes.outputs.clients-go == 'true' ||
needs.detect-changes.outputs.ci == 'true')
runs-on: ubuntu-latest
timeout-minutes: 30
env:
HINDSIGHT_API_LLM_PROVIDER: vertexai
HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY: /tmp/gcp-credentials.json
@@ -2494,7 +2278,6 @@ jobs:
needs.detect-changes.outputs.embed == 'true' ||
needs.detect-changes.outputs.ci == 'true')
runs-on: ubuntu-latest
timeout-minutes: 30
env:
HINDSIGHT_API_LLM_PROVIDER: vertexai
HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY: /tmp/gcp-credentials.json
@@ -2626,7 +2409,6 @@ jobs:
needs.detect-changes.outputs.integration-tests == 'true' ||
needs.detect-changes.outputs.ci == 'true')
runs-on: ubuntu-latest
timeout-minutes: 30
env:
HINDSIGHT_API_LLM_PROVIDER: vertexai
HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY: /tmp/gcp-credentials.json
@@ -2731,7 +2513,6 @@ jobs:
needs.detect-changes.outputs.integrations-ag2 == 'true' ||
needs.detect-changes.outputs.ci == 'true')
runs-on: ubuntu-latest
timeout-minutes: 30
steps:
- uses: actions/checkout@v6
@@ -2768,7 +2549,6 @@ jobs:
needs.detect-changes.outputs.integrations-smolagents == 'true' ||
needs.detect-changes.outputs.ci == 'true')
runs-on: ubuntu-latest
timeout-minutes: 30
steps:
- uses: actions/checkout@v6
@@ -2806,7 +2586,6 @@ jobs:
needs.detect-changes.outputs.integrations-dify == 'true' ||
needs.detect-changes.outputs.ci == 'true')
runs-on: ubuntu-latest
timeout-minutes: 30
steps:
- uses: actions/checkout@v6
@@ -2833,37 +2612,6 @@ jobs:
working-directory: ./hindsight-integrations/dify
run: pytest tests -v
test-flowise-integration:
needs: [detect-changes]
if: >-
github.event_name != 'pull_request_review' &&
(github.event_name == 'workflow_dispatch' ||
needs.detect-changes.outputs.integrations-flowise == 'true' ||
needs.detect-changes.outputs.ci == 'true')
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v6
with:
ref: ${{ github.event.pull_request.head.sha || '' }}
- name: Set up Node.js
uses: actions/setup-node@v6
with:
node-version: '22'
- name: Install dependencies
working-directory: ./hindsight-integrations/flowise
run: npm install --no-audit --no-fund
- name: Type check
working-directory: ./hindsight-integrations/flowise
run: npx tsc --noEmit
- name: Run tests
working-directory: ./hindsight-integrations/flowise
run: npm test
test-crewai-integration:
needs: [detect-changes]
if: >-
@@ -2871,7 +2619,6 @@ jobs:
needs.detect-changes.outputs.integrations-crewai == 'true' ||
needs.detect-changes.outputs.ci == 'true')
runs-on: ubuntu-latest
timeout-minutes: 30
steps:
- uses: actions/checkout@v6
@@ -2901,44 +2648,6 @@ jobs:
working-directory: ./hindsight-integrations/crewai
run: uv run pytest tests -v
test-vapi-integration:
needs: [detect-changes]
if: >-
github.event_name != 'pull_request_review' &&
(github.event_name == 'workflow_dispatch' ||
needs.detect-changes.outputs.integrations-vapi == 'true' ||
needs.detect-changes.outputs.ci == 'true')
runs-on: ubuntu-latest
timeout-minutes: 30
steps:
- uses: actions/checkout@v6
with:
ref: ${{ github.event.pull_request.head.sha || '' }}
- name: Install uv
uses: astral-sh/setup-uv@v7
with:
enable-cache: true
prune-cache: false
- name: Set up Python
uses: actions/setup-python@v6
with:
python-version-file: ".python-version"
- name: Build vapi integration
working-directory: ./hindsight-integrations/vapi
run: uv build
- name: Install dependencies
working-directory: ./hindsight-integrations/vapi
run: uv sync --frozen
- name: Run tests
working-directory: ./hindsight-integrations/vapi
run: uv run pytest tests -v
test-litellm-integration:
needs: [detect-changes]
if: >-
@@ -2946,7 +2655,6 @@ jobs:
needs.detect-changes.outputs.integrations-litellm == 'true' ||
needs.detect-changes.outputs.ci == 'true')
runs-on: ubuntu-latest
timeout-minutes: 30
steps:
- uses: actions/checkout@v6
@@ -2983,7 +2691,6 @@ jobs:
needs.detect-changes.outputs.integrations-pydantic-ai == 'true' ||
needs.detect-changes.outputs.ci == 'true')
runs-on: ubuntu-latest
timeout-minutes: 30
steps:
- uses: actions/checkout@v6
@@ -3020,7 +2727,6 @@ jobs:
needs.detect-changes.outputs.integrations-llamaindex == 'true' ||
needs.detect-changes.outputs.ci == 'true')
runs-on: ubuntu-latest
timeout-minutes: 30
steps:
- uses: actions/checkout@v6
@@ -3057,7 +2763,6 @@ jobs:
needs.detect-changes.outputs.integrations-openai-agents == 'true' ||
needs.detect-changes.outputs.ci == 'true')
runs-on: ubuntu-latest
timeout-minutes: 30
steps:
- uses: actions/checkout@v6
@@ -3094,7 +2799,6 @@ jobs:
needs.detect-changes.outputs.integrations-agentcore == 'true' ||
needs.detect-changes.outputs.ci == 'true')
runs-on: ubuntu-latest
timeout-minutes: 30
steps:
- uses: actions/checkout@v6
@@ -3132,7 +2836,6 @@ jobs:
needs.detect-changes.outputs.core == 'true' ||
needs.detect-changes.outputs.ci == 'true')
runs-on: ubuntu-latest
timeout-minutes: 30
env:
HINDSIGHT_API_LLM_PROVIDER: vertexai
HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY: /tmp/gcp-credentials.json
@@ -3202,7 +2905,6 @@ jobs:
needs.detect-changes.outputs.embed == 'true' ||
needs.detect-changes.outputs.ci == 'true')
runs-on: ubuntu-latest
timeout-minutes: 30
env:
HINDSIGHT_API_LLM_PROVIDER: vertexai
HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY: /tmp/gcp-credentials.json
@@ -3270,7 +2972,6 @@ jobs:
needs.detect-changes.outputs.embed == 'true' ||
needs.detect-changes.outputs.ci == 'true')
runs-on: windows-latest
timeout-minutes: 30
env:
HINDSIGHT_API_LLM_PROVIDER: vertexai
HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY: ${{ github.workspace }}/gcp-credentials.json
@@ -3448,7 +3149,6 @@ jobs:
needs.detect-changes.outputs.hindsight-all == 'true' ||
needs.detect-changes.outputs.ci == 'true')
runs-on: ubuntu-latest
timeout-minutes: 30
env:
HINDSIGHT_API_LLM_PROVIDER: vertexai
HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY: /tmp/gcp-credentials.json
@@ -3500,12 +3200,6 @@ jobs:
rm -rf node_modules/lightningcss node_modules/@tailwindcss
npm install lightningcss @tailwindcss/postcss @tailwindcss/node
- name: Test Control Plane
run: npm test --workspace=hindsight-control-plane
- name: Check i18n locale parity and hardcoded strings
run: npm run i18n:check --workspace=hindsight-control-plane
- name: Build Control Plane
run: npm run build --workspace=hindsight-control-plane
@@ -3543,7 +3237,6 @@ jobs:
needs.detect-changes.outputs.docs == 'true' ||
needs.detect-changes.outputs.ci == 'true')
runs-on: ubuntu-latest
timeout-minutes: 30
strategy:
fail-fast: false
matrix:
@@ -3692,7 +3385,6 @@ jobs:
needs.detect-changes.outputs.dev == 'true' ||
needs.detect-changes.outputs.ci == 'true')
runs-on: ubuntu-latest
timeout-minutes: 30
env:
HINDSIGHT_API_LLM_PROVIDER: vertexai
HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY: /tmp/gcp-credentials.json
@@ -3771,7 +3463,6 @@ jobs:
verify-generated-files:
runs-on: ubuntu-latest
timeout-minutes: 30
env:
UV_FROZEN: "1"
steps:
@@ -3858,7 +3549,6 @@ jobs:
needs.detect-changes.outputs.core == 'true' ||
needs.detect-changes.outputs.ci == 'true')
runs-on: ubuntu-latest
timeout-minutes: 30
steps:
- uses: actions/checkout@v6
with:
@@ -3912,7 +3602,6 @@ jobs:
needs.detect-changes.outputs.dev == 'true' ||
needs.detect-changes.outputs.ci == 'true')
runs-on: ubuntu-latest
timeout-minutes: 30
steps:
- uses: actions/checkout@v6
with:
@@ -3958,14 +3647,10 @@ jobs:
- build-chat-integration
- test-paperclip-integration
- test-pipecat-integration
- test-gemini-spark-integration
- test-vapi-integration
- test-roo-code-integration
- build-control-plane
- build-docs
- test-rust-cli
- lint-helm-chart
- test-standalone-start-script
- build-docker-images
- test-api
- test-api-oracle
@@ -3982,7 +3667,6 @@ jobs:
- test-ag2-integration
- test-smolagents-integration
- test-dify-integration
- test-flowise-integration
- test-crewai-integration
- test-litellm-integration
- test-pydantic-ai-integration
@@ -3999,7 +3683,6 @@ jobs:
- check-openapi-compatibility
- check-cli-coverage
runs-on: ubuntu-latest
timeout-minutes: 30
permissions:
statuses: write
pull-requests: write
-115
View File
@@ -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"
-2
View File
@@ -54,8 +54,6 @@ 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
+2 -25
View File
@@ -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/
-113
View File
@@ -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
-103
View File
@@ -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:
-23
View File
@@ -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:
+7 -33
View File
@@ -10,45 +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
}
if [ "${HINDSIGHT_START_ALL_SOURCE_ONLY:-false}" = "true" ]; then
return 0 2>/dev/null || exit 0
fi
check_pg0_data_integrity "${HOME}/.pg0"
# Service flags (default to true if not set)
ENABLE_API="${HINDSIGHT_ENABLE_API:-true}"
ENABLE_CP="${HINDSIGHT_ENABLE_CP:-true}"
@@ -182,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
-73
View File
@@ -1,73 +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 '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"
+6
View File
@@ -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 -2
View File
@@ -2,8 +2,8 @@ apiVersion: v2
name: hindsight
description: Hindsight helm chart
type: application
version: 0.7.1
appVersion: "0.7.1"
version: 0.6.2
appVersion: "0.6.2"
keywords:
- ai
- memory
+1 -1
View File
@@ -1,6 +1,6 @@
{
"name": "@vectorize-io/hindsight-all",
"version": "0.7.1",
"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",
+2 -2
View File
@@ -4,12 +4,12 @@ build-backend = "setuptools.build_meta"
[project]
name = "hindsight-all-slim"
version = "0.7.1"
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.7.1",
"hindsight-api-slim==0.6.2",
"hindsight-client>=0.0.7",
"hindsight-embed>=0.1.0",
]
+3 -3
View File
@@ -4,12 +4,12 @@ build-backend = "hatchling.build"
[project]
name = "hindsight-all"
version = "0.7.1"
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.7.1",
"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.7.1",
"hindsight-api-slim[local-llm]==0.6.2",
]
test = [
"pytest>=7.0.0",
+1 -1
View File
@@ -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:
+1 -8
View File
@@ -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.7.1"
__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,32 +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 (no default; see VectorChord issue #392)
# and vchordrq.epsilon (default 1.9). probes = 10 / 30 are starting
# defaults pending a workload-specific sweep — vchordrq's recall curve
# shape differs from HNSW's, so the pgvector numbers don't translate
# directly. Revisit with a per-cluster benchmark once we have production
# recall data; until then these are deliberately conservative on the
# high-recall path. We leave epsilon at its default; tightening it is a
# separate trade-off.
# - 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"),),
"vchord": (("vchordrq.probes", "10"),),
}
_ANN_TUNING_HIGH_RECALL: dict[str, tuple[tuple[str, str], ...]] = {
"pgvector": (("hnsw.ef_search", "200"),),
"vchord": (("vchordrq.probes", "30"),),
}
_EXTENSION_INSTALL_SQL = {
"pgvector": ("CREATE EXTENSION IF NOT EXISTS vector",),
"pgvectorscale": (
@@ -94,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.
@@ -154,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"
+2 -19
View File
@@ -30,17 +30,8 @@ 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",
@@ -50,13 +41,6 @@ BACKUP_TABLES = [
"unit_entities",
"entity_cooccurrences",
"memory_links",
"mental_models",
"directives",
"async_operations",
"webhooks",
"file_storage",
"audit_log",
"graph_maintenance_queue",
]
MANIFEST_VERSION = "1"
@@ -284,7 +268,6 @@ async def _run_migration(
ensure_text_search_extension(
resolved_url,
text_search_extension=config.text_search_extension,
pg_search_tokenizer=config.text_search_extension_pg_search_tokenizer,
schema=schema,
)
@@ -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("""
@@ -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")
@@ -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:
@@ -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)"
@@ -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:
@@ -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)
@@ -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)
@@ -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
@@ -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)
@@ -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:
@@ -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:
@@ -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)"
@@ -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)
@@ -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)
@@ -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"""
@@ -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)
+9 -150
View File
@@ -15,7 +15,6 @@ from datetime import datetime, timezone
from typing import Any, Literal
from fastapi import Depends, FastAPI, File, Form, Header, HTTPException, Query, Request, UploadFile
from fastapi.middleware.gzip import GZipMiddleware
from hindsight_api.engine.audit import AuditEntry, AuditLogger
from hindsight_api.extensions import AuthenticationError
@@ -152,11 +151,7 @@ class RecallRequest(BaseModel):
max_tokens: int = 4096
trace: bool = False
query_timestamp: str | None = Field(
default=None,
description=(
"ISO format date string (e.g., '2023-05-30T23:40:00'). Used as the query-time anchor for "
"relative temporal expressions and recency scoring."
),
default=None, description="ISO format date string (e.g., '2023-05-30T23:40:00')"
)
include: IncludeOptions = FieldWithDefault(
IncludeOptions,
@@ -471,13 +466,6 @@ class MemoryItem(BaseModel):
description="Optional tags for visibility scoping. Memories with tags can be filtered during recall.",
)
@field_validator("content")
@classmethod
def validate_content(cls, v: str) -> str:
if not v.strip():
raise ValueError("content cannot be empty")
return v
@field_validator("tags", mode="before")
@classmethod
def coerce_tags(cls, v):
@@ -2192,19 +2180,6 @@ class OperationResponse(BaseModel):
)
class ConsolidationRequest(BaseModel):
"""Request model for consolidation trigger endpoint."""
observation_scopes: list[list[str]] | None = Field(
default=None,
description=(
"Optional list of tag scopes to consolidate. Each scope is a list of tags. "
"Only unconsolidated memories whose tags contain all tags in at least one scope "
"will be processed. If omitted, all unconsolidated memories are processed."
),
)
class ConsolidationResponse(BaseModel):
"""Response model for consolidation trigger endpoint."""
@@ -2674,7 +2649,6 @@ def create_app(
tenant_extension=memory._tenant_extension,
max_slots=config.worker_max_slots,
slot_reservations=config.worker_slot_reservations,
consolidation_bank_priority=config.worker_consolidation_bank_priority or None,
)
poller_task = asyncio.create_task(poller.run())
logging.info(f"Worker poller started (worker_id={worker_id})")
@@ -2747,8 +2721,6 @@ def create_app(
app.state.memory = memory
app.state.audit_logger = memory.audit_logger
app.add_middleware(GZipMiddleware, minimum_size=1024)
# ---------------------------------------------------------------------------
# Patch OpenAPI schema: align ValidationError with Pydantic v2 error format
# ---------------------------------------------------------------------------
@@ -2913,57 +2885,6 @@ def _register_routes(app: FastAPI):
api_key = authorization.strip()
return RequestContext(api_key=api_key)
def precheck_for(operation: str):
"""
Build a FastAPI dependency that runs ``OperationValidator.precheck``.
FastAPI resolves dependencies before deserialising the route's body
parameter. Wiring this dependency on the billable POST routes lets
an extension reject a request e.g. with HTTP 402 when a tenant's
balance is exhausted without the request body ever being read or
materialised in memory.
The dependency intentionally:
- authenticates the tenant (so ``request_context.tenant_id`` is
resolved before the precheck runs);
- falls through silently when no validator is configured or the
validator's default no-op precheck is in effect;
- converts a rejection ``ValidationResult`` into the corresponding
``HTTPException`` directly (the per-route ``OperationValidationError``
catch blocks don't see exceptions raised in dependencies, so we
translate here instead of relying on each handler's try/except).
Args:
operation: Short identifier for the route, e.g. ``"retain"``.
Returns:
A FastAPI dependency callable suitable for ``Depends(...)``.
"""
async def _precheck_dep(
bank_id: str,
request_context: RequestContext = Depends(get_request_context),
) -> None:
validator = getattr(app.state.memory, "_operation_validator", None)
if validator is None:
return
from hindsight_api.extensions import PrecheckContext
await app.state.memory._authenticate_tenant(request_context)
ctx = PrecheckContext(
operation=operation,
bank_id=bank_id,
request_context=request_context,
)
result = await validator.precheck(ctx)
if not result.allowed:
raise HTTPException(
status_code=result.status_code,
detail=result.reason or "Operation not allowed",
)
return _precheck_dep
# Global exception handler for authentication errors
@app.exception_handler(AuthenticationError)
async def authentication_error_handler(request, exc: AuthenticationError):
@@ -3221,10 +3142,7 @@ def _register_routes(app: FastAPI):
)
@audited("recall")
async def api_recall(
bank_id: str,
request: RecallRequest,
request_context: RequestContext = Depends(get_request_context),
_precheck: None = Depends(precheck_for("recall")),
bank_id: str, request: RecallRequest, request_context: RequestContext = Depends(get_request_context)
):
"""Run a recall and return results with trace."""
import time
@@ -3412,10 +3330,7 @@ def _register_routes(app: FastAPI):
)
@audited("reflect")
async def api_reflect(
bank_id: str,
request: ReflectRequest,
request_context: RequestContext = Depends(get_request_context),
_precheck: None = Depends(precheck_for("reflect")),
bank_id: str, request: ReflectRequest, request_context: RequestContext = Depends(get_request_context)
):
metrics = get_metrics_collector()
@@ -3913,7 +3828,6 @@ def _register_routes(app: FastAPI):
bank_id: str,
body: CreateMentalModelRequest,
request_context: RequestContext = Depends(get_request_context),
_precheck: None = Depends(precheck_for("mental_model_create")),
):
"""Create a mental model (async - returns operation_id)."""
try:
@@ -3962,7 +3876,6 @@ def _register_routes(app: FastAPI):
bank_id: str,
mental_model_id: str,
request_context: RequestContext = Depends(get_request_context),
_precheck: None = Depends(precheck_for("mental_model_refresh")),
):
"""Refresh a mental model by re-running its source query (async)."""
try:
@@ -3989,48 +3902,6 @@ def _register_routes(app: FastAPI):
)
raise HTTPException(status_code=500, detail=str(e))
@app.post(
"/v1/default/banks/{bank_id}/mental-models/{mental_model_id}/clear",
response_model=MentalModelResponse,
summary="Clear mental model content",
description=(
"Clear a mental model's content so the next refresh performs a full re-synthesis. "
"This is useful for delta-mode models that have accumulated drift over many "
"incremental refreshes. After clearing, call the /refresh endpoint to trigger "
"a clean full rebuild."
),
operation_id="clear_mental_model",
tags=["Mental Models"],
)
@audited("clear_mental_model", request_param=None)
async def api_clear_mental_model(
bank_id: str,
mental_model_id: str,
request_context: RequestContext = Depends(get_request_context),
):
"""Clear a mental model's content."""
try:
mental_model = await app.state.memory.clear_mental_model(
bank_id=bank_id,
mental_model_id=mental_model_id,
request_context=request_context,
)
if mental_model is None:
raise HTTPException(status_code=404, detail=f"Mental model '{mental_model_id}' not found")
return MentalModelResponse(**mental_model)
except OperationValidationError as e:
raise HTTPException(status_code=e.status_code, detail=e.reason)
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
logger.error(
f"Error in POST /v1/default/banks/{bank_id}/mental-models/{mental_model_id}/clear: {error_detail}"
)
raise HTTPException(status_code=500, detail=str(e))
@app.patch(
"/v1/default/banks/{bank_id}/mental-models/{mental_model_id}",
response_model=MentalModelResponse,
@@ -5524,20 +5395,11 @@ def _register_routes(app: FastAPI):
operation_id="trigger_consolidation",
tags=["Banks"],
)
@audited("consolidation")
async def api_trigger_consolidation(
bank_id: str,
request: ConsolidationRequest | None = None,
request_context: RequestContext = Depends(get_request_context),
):
@audited("consolidation", request_param=None)
async def api_trigger_consolidation(bank_id: str, request_context: RequestContext = Depends(get_request_context)):
"""Trigger consolidation for a bank (async)."""
try:
observation_scopes = request.observation_scopes if request else None
result = await app.state.memory.submit_async_consolidation(
bank_id=bank_id,
request_context=request_context,
observation_scopes=observation_scopes,
)
result = await app.state.memory.submit_async_consolidation(bank_id=bank_id, request_context=request_context)
return ConsolidationResponse(
operation_id=result["operation_id"],
deduplicated=result.get("deduplicated", False),
@@ -5860,10 +5722,7 @@ def _register_routes(app: FastAPI):
)
@audited("retain")
async def api_retain(
bank_id: str,
request: RetainRequest,
request_context: RequestContext = Depends(get_request_context),
_precheck: None = Depends(precheck_for("retain")),
bank_id: str, request: RetainRequest, request_context: RequestContext = Depends(get_request_context)
):
"""Retain memories with optional async processing."""
metrics = get_metrics_collector()
@@ -5948,8 +5807,9 @@ def _register_routes(app: FastAPI):
strategy=group_strategy,
request_context=request_context,
return_usage=True,
outbox_callback_factory=app.state.memory._build_retain_outbox_callback_factory(
outbox_callback=app.state.memory._build_retain_outbox_callback(
bank_id=bank_id,
contents=contents,
operation_id=None,
schema=_current_schema.get(),
),
@@ -6032,7 +5892,6 @@ def _register_routes(app: FastAPI):
files: list[UploadFile] = File(..., description="Files to upload and convert"),
request: str = Form(..., description="JSON string with FileRetainRequest model"),
request_context: RequestContext = Depends(get_request_context),
_precheck: None = Depends(precheck_for("files_retain")),
):
"""Upload and convert files to memories."""
from hindsight_api.config import get_config
@@ -107,7 +107,6 @@ 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",
+4 -320
View File
@@ -7,7 +7,6 @@ All environment variables and their defaults are defined here.
import json
import logging
import os
import re
import sys
from dataclasses import dataclass, field, fields
from datetime import datetime, timezone
@@ -15,7 +14,6 @@ from typing import Any, Literal
from dotenv import find_dotenv, load_dotenv
from ._pg_search import normalize_pg_search_tokenizer
from ._vector_index import validate_extension
from .utils import mask_network_location
@@ -138,7 +136,6 @@ ENV_LLM_MAX_RETRIES = "HINDSIGHT_API_LLM_MAX_RETRIES"
ENV_LLM_INITIAL_BACKOFF = "HINDSIGHT_API_LLM_INITIAL_BACKOFF"
ENV_LLM_MAX_BACKOFF = "HINDSIGHT_API_LLM_MAX_BACKOFF"
ENV_LLM_TIMEOUT = "HINDSIGHT_API_LLM_TIMEOUT"
ENV_LLM_REASONING_EFFORT = "HINDSIGHT_API_LLM_REASONING_EFFORT"
ENV_LLM_GROQ_SERVICE_TIER = "HINDSIGHT_API_LLM_GROQ_SERVICE_TIER"
ENV_LLM_OPENAI_SERVICE_TIER = "HINDSIGHT_API_LLM_OPENAI_SERVICE_TIER"
ENV_LLM_EXTRA_BODY = "HINDSIGHT_API_LLM_EXTRA_BODY"
@@ -172,17 +169,6 @@ ENV_RETAIN_LLM_MAX_BACKOFF = "HINDSIGHT_API_RETAIN_LLM_MAX_BACKOFF"
ENV_RETAIN_LLM_TIMEOUT = "HINDSIGHT_API_RETAIN_LLM_TIMEOUT"
ENV_RETAIN_LLM_LITELLMROUTER_CONFIG = "HINDSIGHT_API_RETAIN_LLM_LITELLMROUTER_CONFIG"
# Fireworks AI batch inference. Fireworks' batch API is a proprietary
# account-scoped dataset/job REST API on a control-plane host, distinct from the
# OpenAI-compatible inference host. account_id is REQUIRED for batch retain
# (the control-plane endpoints are /v1/accounts/{account_id}/...). Static,
# server-level config — it pairs with the Fireworks API key.
ENV_FIREWORKS_ACCOUNT_ID = "HINDSIGHT_API_FIREWORKS_ACCOUNT_ID"
ENV_FIREWORKS_BATCH_BASE_URL = "HINDSIGHT_API_FIREWORKS_BATCH_BASE_URL"
ENV_FIREWORKS_BATCH_MAX_WAIT_SECONDS = "HINDSIGHT_API_FIREWORKS_BATCH_MAX_WAIT_SECONDS"
DEFAULT_FIREWORKS_BATCH_BASE_URL = "https://api.fireworks.ai"
DEFAULT_FIREWORKS_BATCH_MAX_WAIT_SECONDS = 86_400 # 24h — Fireworks' max job timeout
ENV_REFLECT_LLM_PROVIDER = "HINDSIGHT_API_REFLECT_LLM_PROVIDER"
ENV_REFLECT_LLM_API_KEY = "HINDSIGHT_API_REFLECT_LLM_API_KEY"
ENV_REFLECT_LLM_MODEL = "HINDSIGHT_API_REFLECT_LLM_MODEL"
@@ -214,7 +200,6 @@ ENV_EMBEDDINGS_OPENAI_API_KEY = "HINDSIGHT_API_EMBEDDINGS_OPENAI_API_KEY"
ENV_EMBEDDINGS_OPENAI_MODEL = "HINDSIGHT_API_EMBEDDINGS_OPENAI_MODEL"
ENV_EMBEDDINGS_OPENAI_BASE_URL = "HINDSIGHT_API_EMBEDDINGS_OPENAI_BASE_URL"
ENV_EMBEDDINGS_OPENAI_BATCH_SIZE = "HINDSIGHT_API_EMBEDDINGS_OPENAI_BATCH_SIZE"
ENV_EMBEDDINGS_OPENAI_DIMENSIONS = "HINDSIGHT_API_EMBEDDINGS_OPENAI_DIMENSIONS"
# Gemini/Vertex AI embeddings configuration
ENV_EMBEDDINGS_GEMINI_API_KEY = "HINDSIGHT_API_EMBEDDINGS_GEMINI_API_KEY"
@@ -241,15 +226,6 @@ ENV_EMBEDDINGS_OPENROUTER_MODEL = "HINDSIGHT_API_EMBEDDINGS_OPENROUTER_MODEL"
ENV_RERANKER_OPENROUTER_API_KEY = "HINDSIGHT_API_RERANKER_OPENROUTER_API_KEY"
ENV_RERANKER_OPENROUTER_MODEL = "HINDSIGHT_API_RERANKER_OPENROUTER_MODEL"
# ZeroEntropy configuration (embeddings)
ENV_EMBEDDINGS_ZEROENTROPY_API_KEY = "HINDSIGHT_API_EMBEDDINGS_ZEROENTROPY_API_KEY"
ENV_EMBEDDINGS_ZEROENTROPY_MODEL = "HINDSIGHT_API_EMBEDDINGS_ZEROENTROPY_MODEL"
ENV_EMBEDDINGS_ZEROENTROPY_BASE_URL = "HINDSIGHT_API_EMBEDDINGS_ZEROENTROPY_BASE_URL"
ENV_EMBEDDINGS_ZEROENTROPY_DIMENSIONS = "HINDSIGHT_API_EMBEDDINGS_ZEROENTROPY_DIMENSIONS"
ENV_EMBEDDINGS_ZEROENTROPY_ENCODING_FORMAT = "HINDSIGHT_API_EMBEDDINGS_ZEROENTROPY_ENCODING_FORMAT"
ENV_EMBEDDINGS_ZEROENTROPY_LATENCY = "HINDSIGHT_API_EMBEDDINGS_ZEROENTROPY_LATENCY"
ENV_EMBEDDINGS_ZEROENTROPY_BATCH_SIZE = "HINDSIGHT_API_EMBEDDINGS_ZEROENTROPY_BATCH_SIZE"
# Deprecated: Legacy shared Cohere API key (for backward compatibility)
ENV_COHERE_API_KEY = "HINDSIGHT_API_COHERE_API_KEY"
@@ -288,14 +264,6 @@ ENV_RERANKER_TEI_URL = "HINDSIGHT_API_RERANKER_TEI_URL"
ENV_RERANKER_TEI_BATCH_SIZE = "HINDSIGHT_API_RERANKER_TEI_BATCH_SIZE"
ENV_RERANKER_TEI_MAX_CONCURRENT = "HINDSIGHT_API_RERANKER_TEI_MAX_CONCURRENT"
ENV_RERANKER_TEI_HTTP_TIMEOUT = "HINDSIGHT_API_RERANKER_TEI_HTTP_TIMEOUT"
ENV_RERANKER_COHERE_TIMEOUT = "HINDSIGHT_API_RERANKER_COHERE_TIMEOUT"
ENV_RERANKER_OPENROUTER_TIMEOUT = "HINDSIGHT_API_RERANKER_OPENROUTER_TIMEOUT"
ENV_RERANKER_ZEROENTROPY_TIMEOUT = "HINDSIGHT_API_RERANKER_ZEROENTROPY_TIMEOUT"
ENV_RERANKER_SILICONFLOW_TIMEOUT = "HINDSIGHT_API_RERANKER_SILICONFLOW_TIMEOUT"
ENV_RERANKER_ALIBABA_TIMEOUT = "HINDSIGHT_API_RERANKER_ALIBABA_TIMEOUT"
ENV_RERANKER_LITELLM_TIMEOUT = "HINDSIGHT_API_RERANKER_LITELLM_TIMEOUT"
ENV_RERANKER_LITELLM_SDK_TIMEOUT = "HINDSIGHT_API_RERANKER_LITELLM_SDK_TIMEOUT"
ENV_RERANKER_GOOGLE_TIMEOUT = "HINDSIGHT_API_RERANKER_GOOGLE_TIMEOUT"
ENV_RERANKER_MAX_CANDIDATES = "HINDSIGHT_API_RERANKER_MAX_CANDIDATES"
ENV_RERANKER_FLASHRANK_MODEL = "HINDSIGHT_API_RERANKER_FLASHRANK_MODEL"
ENV_RERANKER_FLASHRANK_CACHE_DIR = "HINDSIGHT_API_RERANKER_FLASHRANK_CACHE_DIR"
@@ -311,10 +279,6 @@ ENV_RERANKER_SILICONFLOW_API_KEY = "HINDSIGHT_API_RERANKER_SILICONFLOW_API_KEY"
ENV_RERANKER_SILICONFLOW_MODEL = "HINDSIGHT_API_RERANKER_SILICONFLOW_MODEL"
ENV_RERANKER_SILICONFLOW_BASE_URL = "HINDSIGHT_API_RERANKER_SILICONFLOW_BASE_URL"
# Alibaba Cloud DashScope configuration (reranker only)
ENV_RERANKER_ALIBABA_API_KEY = "HINDSIGHT_API_RERANKER_ALIBABA_API_KEY"
ENV_RERANKER_ALIBABA_MODEL = "HINDSIGHT_API_RERANKER_ALIBABA_MODEL"
# Google Discovery Engine reranker configuration
ENV_RERANKER_GOOGLE_MODEL = "HINDSIGHT_API_RERANKER_GOOGLE_MODEL"
ENV_RERANKER_GOOGLE_PROJECT_ID = "HINDSIGHT_API_RERANKER_GOOGLE_PROJECT_ID"
@@ -322,9 +286,6 @@ ENV_RERANKER_GOOGLE_SERVICE_ACCOUNT_KEY = "HINDSIGHT_API_RERANKER_GOOGLE_SERVICE
ENV_VECTOR_EXTENSION = "HINDSIGHT_API_VECTOR_EXTENSION"
ENV_TEXT_SEARCH_EXTENSION = "HINDSIGHT_API_TEXT_SEARCH_EXTENSION"
ENV_TEXT_SEARCH_EXTENSION_NATIVE_LANGUAGE = "HINDSIGHT_API_TEXT_SEARCH_EXTENSION_NATIVE_LANGUAGE"
ENV_TEXT_SEARCH_EXTENSION_PG_SEARCH_TOKENIZER = "HINDSIGHT_API_TEXT_SEARCH_EXTENSION_PG_SEARCH_TOKENIZER"
ENV_LLM_OUTPUT_LANGUAGE = "HINDSIGHT_API_LLM_OUTPUT_LANGUAGE"
ENV_HOST = "HINDSIGHT_API_HOST"
ENV_PORT = "HINDSIGHT_API_PORT"
@@ -333,7 +294,6 @@ ENV_LOG_LEVEL = "HINDSIGHT_API_LOG_LEVEL"
ENV_LOG_FORMAT = "HINDSIGHT_API_LOG_FORMAT"
ENV_LOG_JSON_FIELDS = "HINDSIGHT_API_LOG_JSON_FIELDS"
ENV_WORKERS = "HINDSIGHT_API_WORKERS"
ENV_ACCESS_LOG = "HINDSIGHT_API_ACCESS_LOG"
ENV_MCP_ENABLED = "HINDSIGHT_API_MCP_ENABLED"
ENV_MCP_ENABLED_TOOLS = "HINDSIGHT_API_MCP_ENABLED_TOOLS"
ENV_MCP_STATELESS = "HINDSIGHT_API_MCP_STATELESS"
@@ -373,7 +333,6 @@ ENV_RETAIN_CUSTOM_INSTRUCTIONS = "HINDSIGHT_API_RETAIN_CUSTOM_INSTRUCTIONS"
ENV_RETAIN_DEFAULT_STRATEGY = "HINDSIGHT_API_RETAIN_DEFAULT_STRATEGY"
ENV_RETAIN_BATCH_TOKENS = "HINDSIGHT_API_RETAIN_BATCH_TOKENS"
ENV_RETAIN_ENTITY_LOOKUP = "HINDSIGHT_API_RETAIN_ENTITY_LOOKUP"
ENV_RETAIN_ENTITY_RESOLUTION_BATCH_SIZE = "HINDSIGHT_API_RETAIN_ENTITY_RESOLUTION_BATCH_SIZE"
ENV_RETAIN_BATCH_ENABLED = "HINDSIGHT_API_RETAIN_BATCH_ENABLED"
ENV_RETAIN_BATCH_POLL_INTERVAL_SECONDS = "HINDSIGHT_API_RETAIN_BATCH_POLL_INTERVAL_SECONDS"
ENV_RETAIN_CHUNK_BATCH_SIZE = "HINDSIGHT_API_RETAIN_CHUNK_BATCH_SIZE"
@@ -402,11 +361,9 @@ ENV_FILE_DELETE_AFTER_RETAIN = "HINDSIGHT_API_FILE_DELETE_AFTER_RETAIN"
# Observations settings (consolidated knowledge from facts)
ENV_ENABLE_OBSERVATIONS = "HINDSIGHT_API_ENABLE_OBSERVATIONS"
ENV_ENABLE_AUTO_CONSOLIDATION = "HINDSIGHT_API_ENABLE_AUTO_CONSOLIDATION"
ENV_CONSOLIDATION_BATCH_SIZE = "HINDSIGHT_API_CONSOLIDATION_BATCH_SIZE"
ENV_CONSOLIDATION_MAX_MEMORIES_PER_ROUND = "HINDSIGHT_API_CONSOLIDATION_MAX_MEMORIES_PER_ROUND"
ENV_CONSOLIDATION_LLM_BATCH_SIZE = "HINDSIGHT_API_CONSOLIDATION_LLM_BATCH_SIZE"
ENV_CONSOLIDATION_LLM_PARALLELISM = "HINDSIGHT_API_CONSOLIDATION_LLM_PARALLELISM"
ENV_CONSOLIDATION_MAX_TOKENS = "HINDSIGHT_API_CONSOLIDATION_MAX_TOKENS"
ENV_CONSOLIDATION_SOURCE_FACTS_MAX_TOKENS = "HINDSIGHT_API_CONSOLIDATION_SOURCE_FACTS_MAX_TOKENS"
ENV_CONSOLIDATION_SOURCE_FACTS_MAX_TOKENS_PER_OBSERVATION = (
@@ -418,7 +375,6 @@ ENV_OBSERVATIONS_MISSION = "HINDSIGHT_API_OBSERVATIONS_MISSION"
ENV_MAX_OBSERVATIONS_PER_SCOPE = "HINDSIGHT_API_MAX_OBSERVATIONS_PER_SCOPE"
ENV_ENABLE_OBSERVATION_HISTORY = "HINDSIGHT_API_ENABLE_OBSERVATION_HISTORY"
ENV_ENABLE_MENTAL_MODEL_HISTORY = "HINDSIGHT_API_ENABLE_MENTAL_MODEL_HISTORY"
ENV_MENTAL_MODEL_HISTORY_MAX_ENTRIES = "HINDSIGHT_API_MENTAL_MODEL_HISTORY_MAX_ENTRIES"
# Webhook configuration (global, static - server-level only)
ENV_WEBHOOK_URL = "HINDSIGHT_API_WEBHOOK_URL"
@@ -453,7 +409,6 @@ ENV_WORKER_ENABLED = "HINDSIGHT_API_WORKER_ENABLED"
ENV_WORKER_ID = "HINDSIGHT_API_WORKER_ID"
ENV_WORKER_POLL_INTERVAL_MS = "HINDSIGHT_API_WORKER_POLL_INTERVAL_MS"
ENV_WORKER_MAX_RETRIES = "HINDSIGHT_API_WORKER_MAX_RETRIES"
ENV_WORKER_TASK_RETRY_BACKOFF_SECONDS = "HINDSIGHT_API_WORKER_TASK_RETRY_BACKOFF_SECONDS"
ENV_WORKER_HTTP_PORT = "HINDSIGHT_API_WORKER_HTTP_PORT"
ENV_WORKER_MAX_SLOTS = "HINDSIGHT_API_WORKER_MAX_SLOTS"
@@ -467,9 +422,7 @@ WORKER_SLOT_RESERVATION_TYPES: dict[str, tuple[str, int]] = {
"retain": ("HINDSIGHT_API_WORKER_RETAIN_MAX_SLOTS", 0),
"file_convert_retain": ("HINDSIGHT_API_WORKER_FILE_CONVERT_RETAIN_MAX_SLOTS", 0),
"refresh_mental_model": ("HINDSIGHT_API_WORKER_REFRESH_MENTAL_MODEL_MAX_SLOTS", 0),
"graph_maintenance": ("HINDSIGHT_API_WORKER_GRAPH_MAINTENANCE_MAX_SLOTS", 0),
}
ENV_WORKER_CONSOLIDATION_BANK_PRIORITY = "HINDSIGHT_API_WORKER_CONSOLIDATION_BANK_PRIORITY"
ENV_RETAIN_MAX_CONCURRENT = "HINDSIGHT_API_RETAIN_MAX_CONCURRENT"
# Reflect agent settings
@@ -520,7 +473,6 @@ PROVIDER_DEFAULT_MODELS = {
"zai": "glm-4.5-flash",
"opencode-go": "deepseek-v4-flash",
"ollama": "gemma3:12b",
"ollama-cloud": "gemma3:12b",
"llamacpp": "gemma-4-e2b-it",
"lmstudio": "local-model",
"vertexai": "google/gemini-2.5-flash-lite",
@@ -532,7 +484,6 @@ PROVIDER_DEFAULT_MODELS = {
"bedrock": "us.amazon.nova-2-lite-v1:0",
"volcano": "doubao-pro-32k",
"openrouter": "qwen/qwen3.5-9b",
"fireworks": "accounts/fireworks/models/llama-v3p1-8b-instruct",
}
DEFAULT_LLM_MODEL = "gpt-4o-mini" # Fallback if provider not in table
# Built-in llama.cpp defaults
@@ -547,7 +498,6 @@ DEFAULT_LLM_MAX_RETRIES = 3 # Max retry attempts for LLM API calls
DEFAULT_LLM_INITIAL_BACKOFF = 1.0 # Initial backoff in seconds for retry exponential backoff
DEFAULT_LLM_MAX_BACKOFF = 60.0 # Max backoff cap in seconds for retry exponential backoff
DEFAULT_LLM_TIMEOUT = 120.0 # seconds
DEFAULT_LLM_REASONING_EFFORT = "low"
# Vertex AI defaults
DEFAULT_LLM_VERTEXAI_PROJECT_ID = None # Required for Vertex AI
@@ -581,16 +531,6 @@ DEFAULT_RERANKER_LOCAL_BATCH_SIZE = 32 # Batch size for local reranker predict(
DEFAULT_RERANKER_TEI_BATCH_SIZE = 128
DEFAULT_RERANKER_TEI_MAX_CONCURRENT = 8
DEFAULT_RERANKER_TEI_HTTP_TIMEOUT = 30.0 # HTTP timeout for TEI reranker requests (seconds)
# HTTP timeout (seconds) for remote rerank providers. Defaults match the previous
# hardcoded constructor defaults so unset envs keep current behavior.
DEFAULT_RERANKER_COHERE_TIMEOUT = 60.0
DEFAULT_RERANKER_OPENROUTER_TIMEOUT = 60.0
DEFAULT_RERANKER_ZEROENTROPY_TIMEOUT = 60.0
DEFAULT_RERANKER_SILICONFLOW_TIMEOUT = 60.0
DEFAULT_RERANKER_ALIBABA_TIMEOUT = 60.0
DEFAULT_RERANKER_LITELLM_TIMEOUT = 60.0
DEFAULT_RERANKER_LITELLM_SDK_TIMEOUT = 60.0
DEFAULT_RERANKER_GOOGLE_TIMEOUT = 60.0
DEFAULT_RERANKER_MAX_CANDIDATES = 300
DEFAULT_RERANKER_FLASHRANK_MODEL = "ms-marco-MiniLM-L-12-v2" # Best balance of speed and quality
DEFAULT_RERANKER_FLASHRANK_CACHE_DIR = None # Use default cache directory
@@ -603,39 +543,18 @@ DEFAULT_RERANKER_COHERE_MODEL = "rerank-english-v3.0"
DEFAULT_EMBEDDINGS_OPENROUTER_MODEL = "perplexity/pplx-embed-v1-0.6b"
DEFAULT_RERANKER_OPENROUTER_MODEL = "cohere/rerank-v3.5"
# ZeroEntropy defaults
DEFAULT_EMBEDDINGS_ZEROENTROPY_MODEL = "zembed-1"
# Shared between embeddings (zembed-1) and reranker (zerank-*) — the host is the same.
DEFAULT_ZEROENTROPY_BASE_URL = "https://api.zeroentropy.dev"
# ZeroEntropy's API default is 2560, but Hindsight defaults to 1280 so the
# provider works with pgvector HNSW's 2000-dimension index limit out of the box.
DEFAULT_EMBEDDINGS_ZEROENTROPY_DIMENSIONS = 1280
DEFAULT_EMBEDDINGS_ZEROENTROPY_ENCODING_FORMAT = "float"
DEFAULT_EMBEDDINGS_ZEROENTROPY_LATENCY = None
DEFAULT_EMBEDDINGS_ZEROENTROPY_BATCH_SIZE = 100
DEFAULT_RERANKER_ZEROENTROPY_MODEL = "zerank-2"
DEFAULT_RERANKER_SILICONFLOW_MODEL = "BAAI/bge-reranker-v2-m3"
DEFAULT_RERANKER_SILICONFLOW_BASE_URL = "https://api.siliconflow.cn/v1"
DEFAULT_RERANKER_ALIBABA_MODEL = "qwen3-rerank"
DEFAULT_RERANKER_GOOGLE_MODEL = "semantic-ranker-default-004"
# Vector extension (pgvector, vchord, pgvectorscale, or AlloyDB ScaNN)
DEFAULT_VECTOR_EXTENSION = "pgvector" # Options: "pgvector", "vchord", "pgvectorscale", "scann"
# Text search extension (native PostgreSQL, vchord BM25, Timescale pg_textsearch,
# pgroonga, or ParadeDB pg_search)
DEFAULT_TEXT_SEARCH_EXTENSION = "native" # Options: "native", "vchord", "pg_textsearch", "pgroonga", "pg_search"
# PostgreSQL text search dictionary used by the native tsvector backend. Only
# affects text_search_extension == "native"; other backends use their own
# tokenizers (vchord: llmlingua2, pg_textsearch: hardcoded english,
# pgroonga: TokenBigram polyglot, pg_search: per-field Tantivy tokenizer).
DEFAULT_TEXT_SEARCH_EXTENSION_NATIVE_LANGUAGE = "english"
DEFAULT_TEXT_SEARCH_EXTENSION_PG_SEARCH_TOKENIZER = ""
# Text search extension (native PostgreSQL, vchord BM25, or Timescale pg_textsearch)
DEFAULT_TEXT_SEARCH_EXTENSION = "native" # Options: "native", "vchord", "pg_textsearch"
# LiteLLM defaults
DEFAULT_LITELLM_API_BASE = "http://localhost:4000"
@@ -654,7 +573,6 @@ DEFAULT_BASE_PATH = "" # Empty string = root path
DEFAULT_LOG_LEVEL = "info"
DEFAULT_LOG_FORMAT = "text" # Options: "text", "json"
DEFAULT_WORKERS = 1
DEFAULT_ACCESS_LOG = False
DEFAULT_MCP_ENABLED = True
DEFAULT_MCP_ENABLED_TOOLS: list[str] | None = None # None = all tools enabled
DEFAULT_MCP_STATELESS = False # False = stateful (supports SSE/GET); True = stateless (POST-only)
@@ -683,7 +601,6 @@ DEFAULT_RETAIN_CHUNK_BATCH_SIZE = (
)
DEFAULT_RETAIN_BATCH_TOKENS = 10_000 # ~40KB of text # Max chars per sub-batch for async retain auto-splitting
DEFAULT_RETAIN_ENTITY_LOOKUP = "trigram" # "full" or "trigram"
DEFAULT_RETAIN_ENTITY_RESOLUTION_BATCH_SIZE = 100 # Unique entity names per pg_trgm candidate lookup query
DEFAULT_RETAIN_BATCH_ENABLED = False # Use LLM Batch API for fact extraction (only when async=True)
DEFAULT_RETAIN_BATCH_POLL_INTERVAL_SECONDS = 60 # Batch API polling interval in seconds
@@ -698,25 +615,14 @@ DEFAULT_FILE_DELETE_AFTER_RETAIN = True # Delete file bytes after retain (saves
# Observations defaults (consolidated knowledge from facts)
DEFAULT_ENABLE_OBSERVATIONS = True # Observations enabled by default
DEFAULT_ENABLE_AUTO_CONSOLIDATION = True # Auto-consolidation after retain enabled by default
DEFAULT_ENABLE_OBSERVATION_HISTORY = True # Observation history tracking enabled by default
DEFAULT_ENABLE_MENTAL_MODEL_HISTORY = True # Mental model history tracking enabled by default
# Each history entry snapshots previous_content + previous_reflect_response. Without
# a cap, sustained mental-model refresh load grows the jsonb array unboundedly until
# it crosses Postgres's hard 256MB jsonb limit and subsequent UPDATEs fail with
# SQLSTATE 54000. 50 keeps the array well under 100MB even with large reflect
# responses, while preserving enough recent history for meaningful audit / rollback.
DEFAULT_MENTAL_MODEL_HISTORY_MAX_ENTRIES = 50
DEFAULT_CONSOLIDATION_MAX_ATTEMPTS = 3 # Outer retry attempts for consolidation LLM batch calls
DEFAULT_CONSOLIDATION_BATCH_SIZE = 50 # Memories to load per batch (internal memory optimization)
DEFAULT_CONSOLIDATION_MAX_MEMORIES_PER_ROUND = (
100 # Max memories per consolidation round (0 = unlimited). Limits how long one bank holds a worker slot.
)
DEFAULT_CONSOLIDATION_LLM_BATCH_SIZE = 8 # Facts per LLM call (1 = no batching; >1 = batch mode)
DEFAULT_CONSOLIDATION_LLM_PARALLELISM = (
4 # Max tag groups consolidated concurrently per op. Locks on overlapping write
# scopes degrade to sequential automatically; matches retain_max_concurrent.
)
DEFAULT_CONSOLIDATION_MAX_TOKENS = 512 # Max tokens for recall when finding related observations
DEFAULT_CONSOLIDATION_RECALL_BUDGET = "low" # Budget level for consolidation recall (low/mid/high)
DEFAULT_CONSOLIDATION_SOURCE_FACTS_MAX_TOKENS = (
@@ -743,7 +649,6 @@ DEFAULT_WORKER_ENABLED = True # API runs worker by default (standalone mode)
DEFAULT_WORKER_ID = None # Will use hostname if not specified
DEFAULT_WORKER_POLL_INTERVAL_MS = 500 # Poll database every 500ms
DEFAULT_WORKER_MAX_RETRIES = 3 # Max retries before marking task failed
DEFAULT_WORKER_TASK_RETRY_BACKOFF_SECONDS = 60 # Seconds between retries on transient task failure
DEFAULT_WORKER_HTTP_PORT = 8889 # HTTP port for worker metrics/health
DEFAULT_WORKER_MAX_SLOTS = 10 # Total concurrent tasks per worker
DEFAULT_RETAIN_MAX_CONCURRENT = 4 # Max concurrent retain DB phases (HNSW reads + writes). Limits I/O contention.
@@ -887,24 +792,6 @@ def _parse_positive_int(name: str, raw: str | None, default: int) -> int:
return parsed
def _parse_optional_positive_int(name: str, raw: str | None) -> int | None:
"""Parse an optional env var that must be a positive integer when set."""
if raw is None or raw == "":
return None
return _parse_positive_int(name, raw, 1)
def _parse_optional_choice(name: str, raw: str | None, allowed: frozenset[str]) -> str | None:
"""Parse an optional string env var constrained to a small allowlist."""
if raw is None or raw == "":
return None
normalized = raw.lower()
if normalized not in allowed:
values = ", ".join(sorted(allowed))
raise ValueError(f"{name} must be one of {values}, got {raw!r}")
return normalized
def _validate_extraction_mode(mode: str) -> str:
"""Validate and normalize extraction mode."""
mode_lower = mode.lower()
@@ -929,38 +816,6 @@ def _validate_recall_budget_function(function: str) -> str:
return function_lower
def _parse_bank_priority(raw: str) -> dict[str, int]:
"""Parse ``bank-pattern:priority,...`` into ``{pattern: priority}``.
``*`` in a pattern is kept as-is here; the SQL layer converts it to ``%``
for LIKE matching. A bare ``*`` key is the catch-all default for unlisted
banks. Returns an empty dict when *raw* is blank.
"""
result: dict[str, int] = {}
raw = raw.strip()
if not raw:
return result
for entry in raw.split(","):
entry = entry.strip()
if not entry:
continue
if ":" not in entry:
raise ValueError(f"Invalid bank priority entry '{entry}': expected 'bank-pattern:priority'")
pattern, priority_str = entry.rsplit(":", 1)
pattern = pattern.strip()
priority_str = priority_str.strip()
if not pattern:
raise ValueError(f"Empty bank pattern in entry '{entry}'")
try:
priority = int(priority_str)
except ValueError:
raise ValueError(f"Invalid priority '{priority_str}' in entry '{entry}': must be an integer") from None
if priority < 1:
raise ValueError(f"Priority must be >= 1, got {priority} in entry '{entry}'")
result[pattern] = priority
return result
def _get_default_model_for_provider(provider: str) -> str:
"""Get the default model for a given provider."""
return PROVIDER_DEFAULT_MODELS.get(provider.lower(), DEFAULT_LLM_MODEL)
@@ -1019,19 +874,7 @@ class HindsightConfig:
migration_database_url: str | None
database_schema: str
vector_extension: str # "pgvector", "vchord", "pgvectorscale", or "scann"
text_search_extension: str # "native", "vchord", "pg_textsearch", "pgroonga", or "pg_search"
# PostgreSQL text search dictionary for the "native" backend (ignored by
# other backends). Only the "native" backend reads this field; pgroonga
# uses TokenBigram, vchord uses llmlingua2, pg_textsearch hardcodes english,
# pg_search uses Tantivy per-field tokenizers.
text_search_extension_native_language: str
# ParadeDB pg_search tokenizer used when building BM25 indexes. Empty keeps
# ParadeDB's default tokenizer.
text_search_extension_pg_search_tokenizer: str
# When set, every LLM-generated artifact (retain facts, consolidation
# observations, reflect responses) is forced into this language regardless
# of the source content. Unset preserves source language.
llm_output_language: str | None
text_search_extension: str # "native" or "vchord"
# LLM (default, used as fallback for per-operation config)
llm_provider: str
@@ -1043,7 +886,6 @@ class HindsightConfig:
llm_initial_backoff: float
llm_max_backoff: float
llm_timeout: float
llm_reasoning_effort: str
llm_groq_service_tier: str # Groq: "on_demand", "flex", or "auto"
llm_openai_service_tier: str | None # OpenAI: None (default) or "flex" (50% cheaper)
llm_extra_body: (
@@ -1087,11 +929,6 @@ class HindsightConfig:
retain_llm_timeout: float | None
retain_llm_litellmrouter_config: dict | None
# Fireworks AI batch inference (static, server-level)
fireworks_account_id: str | None
fireworks_batch_base_url: str
fireworks_batch_max_wait_seconds: int
reflect_llm_provider: str | None
reflect_llm_api_key: str | None
reflect_llm_model: str | None
@@ -1161,34 +998,24 @@ class HindsightConfig:
reranker_cohere_api_key: str | None
reranker_cohere_model: str
reranker_cohere_base_url: str | None
reranker_cohere_timeout: float
reranker_openrouter_api_key: str | None
reranker_openrouter_model: str
reranker_openrouter_timeout: float
reranker_litellm_api_base: str
reranker_litellm_api_key: str | None
reranker_litellm_model: str
reranker_litellm_max_tokens_per_doc: int | None
reranker_litellm_timeout: float
reranker_litellm_sdk_api_key: str | None
reranker_litellm_sdk_model: str
reranker_litellm_sdk_api_base: str | None
reranker_litellm_sdk_timeout: float
reranker_zeroentropy_api_key: str | None
reranker_zeroentropy_model: str
reranker_zeroentropy_base_url: str | None
reranker_zeroentropy_timeout: float
reranker_siliconflow_api_key: str | None
reranker_siliconflow_model: str
reranker_siliconflow_base_url: str
reranker_siliconflow_timeout: float
reranker_alibaba_api_key: str | None
reranker_alibaba_model: str
reranker_alibaba_timeout: float
reranker_google_model: str
reranker_google_project_id: str | None
reranker_google_service_account_key: str | None
reranker_google_timeout: float
# Server
host: str
@@ -1227,7 +1054,6 @@ class HindsightConfig:
retain_batch_enabled: bool
retain_batch_poll_interval_seconds: int
retain_entity_lookup: str # "full" or "trigram"
retain_entity_resolution_batch_size: int # Unique entity names per pg_trgm candidate lookup query
retain_chunk_batch_size: int # Max chunks per streaming batch (0 = disabled)
# File storage (static - server-level only)
@@ -1254,14 +1080,11 @@ class HindsightConfig:
# Observations settings (consolidated knowledge from facts)
enable_observations: bool
enable_auto_consolidation: bool
enable_observation_history: bool
enable_mental_model_history: bool
mental_model_history_max_entries: int
consolidation_batch_size: int
consolidation_max_memories_per_round: int
consolidation_llm_batch_size: int
consolidation_llm_parallelism: int
consolidation_max_tokens: int
consolidation_recall_budget: str
consolidation_source_facts_max_tokens: int
@@ -1324,11 +1147,9 @@ class HindsightConfig:
worker_id: str | None
worker_poll_interval_ms: int
worker_max_retries: int
worker_task_retry_backoff_seconds: int
worker_http_port: int
worker_max_slots: int
worker_slot_reservations: dict[str, int]
worker_consolidation_bank_priority: dict[str, int]
retain_max_concurrent: int
# Reflect agent settings
@@ -1358,14 +1179,6 @@ class HindsightConfig:
# Defaulted fields (source-compatible additions — existing direct constructor callers keep working).
# Keep at the end of the dataclass; Python forbids non-default fields after default fields.
embeddings_openai_batch_size: int = DEFAULT_EMBEDDINGS_OPENAI_BATCH_SIZE
embeddings_openai_dimensions: int | None = None
embeddings_zeroentropy_api_key: str | None = None
embeddings_zeroentropy_model: str = DEFAULT_EMBEDDINGS_ZEROENTROPY_MODEL
embeddings_zeroentropy_base_url: str = DEFAULT_ZEROENTROPY_BASE_URL
embeddings_zeroentropy_dimensions: int = DEFAULT_EMBEDDINGS_ZEROENTROPY_DIMENSIONS
embeddings_zeroentropy_encoding_format: str = DEFAULT_EMBEDDINGS_ZEROENTROPY_ENCODING_FORMAT
embeddings_zeroentropy_batch_size: int = DEFAULT_EMBEDDINGS_ZEROENTROPY_BATCH_SIZE
embeddings_zeroentropy_latency: str | None = DEFAULT_EMBEDDINGS_ZEROENTROPY_LATENCY
# Class-level sets for configuration categorization
@@ -1389,7 +1202,6 @@ class HindsightConfig:
"embeddings_tei_base_url",
"reranker_tei_base_url",
"reranker_cohere_base_url",
"embeddings_zeroentropy_base_url",
"reranker_zeroentropy_base_url",
"reranker_siliconflow_base_url",
# Service Account Keys
@@ -1398,7 +1210,6 @@ class HindsightConfig:
"reranker_google_service_account_key",
# Embeddings API keys
"embeddings_gemini_api_key",
"embeddings_zeroentropy_api_key",
# File storage credentials
"file_storage_s3_access_key_id",
"file_storage_s3_secret_access_key",
@@ -1428,9 +1239,7 @@ class HindsightConfig:
"entities_allow_free_form",
# Consolidation settings
"enable_observations",
"enable_auto_consolidation",
"consolidation_llm_batch_size",
"consolidation_llm_parallelism",
"consolidation_max_memories_per_round",
"consolidation_source_facts_max_tokens",
"consolidation_source_facts_max_tokens_per_observation",
@@ -1525,30 +1334,12 @@ class HindsightConfig:
validate_extension(self.vector_extension)
# Validate text_search_extension
valid_text_search = ("native", "vchord", "pg_textsearch", "pgroonga", "pg_search")
valid_text_search = ("native", "vchord", "pg_textsearch")
if self.text_search_extension not in valid_text_search:
raise ValueError(
f"Invalid text_search_extension: {self.text_search_extension}. Must be one of: {', '.join(valid_text_search)}"
)
# Validate text_search_extension_native_language as a PG identifier.
# Embedded directly into raw SQL via to_tsvector('<lang>', ...), so we
# reject anything that isn't a plain identifier to prevent injection.
# Intentionally permissive about which dictionaries exist — users may
# install custom ones like zhparser; we only check shape here. PG
# raises a clear error at query time if the dictionary is missing.
if not re.fullmatch(r"[a-zA-Z_][a-zA-Z0-9_]*", self.text_search_extension_native_language):
raise ValueError(
f"Invalid text_search_extension_native_language: "
f"{self.text_search_extension_native_language!r}. Must be a valid PostgreSQL identifier "
f"(letters, digits, underscores; not starting with a digit). Examples: 'english', "
f"'french', 'simple', 'zhparser'."
)
self.text_search_extension_pg_search_tokenizer = normalize_pg_search_tokenizer(
self.text_search_extension_pg_search_tokenizer
)
# When LLM provider is "none", force chunks-only mode and disable LLM-dependent features
if self.llm_provider == "none":
self.retain_extraction_mode = "chunks"
@@ -1625,15 +1416,6 @@ class HindsightConfig:
database_schema=os.getenv(ENV_DATABASE_SCHEMA, DEFAULT_DATABASE_SCHEMA),
vector_extension=os.getenv(ENV_VECTOR_EXTENSION, DEFAULT_VECTOR_EXTENSION).lower(),
text_search_extension=os.getenv(ENV_TEXT_SEARCH_EXTENSION, DEFAULT_TEXT_SEARCH_EXTENSION).lower(),
text_search_extension_native_language=os.getenv(
ENV_TEXT_SEARCH_EXTENSION_NATIVE_LANGUAGE,
DEFAULT_TEXT_SEARCH_EXTENSION_NATIVE_LANGUAGE,
).lower(),
text_search_extension_pg_search_tokenizer=os.getenv(
ENV_TEXT_SEARCH_EXTENSION_PG_SEARCH_TOKENIZER,
DEFAULT_TEXT_SEARCH_EXTENSION_PG_SEARCH_TOKENIZER,
),
llm_output_language=(os.getenv(ENV_LLM_OUTPUT_LANGUAGE) or None),
# LLM
llm_provider=llm_provider,
llm_api_key=os.getenv(ENV_LLM_API_KEY),
@@ -1644,7 +1426,6 @@ class HindsightConfig:
llm_initial_backoff=float(os.getenv(ENV_LLM_INITIAL_BACKOFF, str(DEFAULT_LLM_INITIAL_BACKOFF))),
llm_max_backoff=float(os.getenv(ENV_LLM_MAX_BACKOFF, str(DEFAULT_LLM_MAX_BACKOFF))),
llm_timeout=float(os.getenv(ENV_LLM_TIMEOUT, str(DEFAULT_LLM_TIMEOUT))),
llm_reasoning_effort=os.getenv(ENV_LLM_REASONING_EFFORT, DEFAULT_LLM_REASONING_EFFORT),
llm_groq_service_tier=os.getenv(ENV_LLM_GROQ_SERVICE_TIER, DEFAULT_LLM_GROQ_SERVICE_TIER),
llm_openai_service_tier=os.getenv(ENV_LLM_OPENAI_SERVICE_TIER, DEFAULT_LLM_OPENAI_SERVICE_TIER),
llm_extra_body=json.loads(os.getenv(ENV_LLM_EXTRA_BODY, "null")),
@@ -1675,11 +1456,6 @@ class HindsightConfig:
else None
),
retain_llm_base_url=os.getenv(ENV_RETAIN_LLM_BASE_URL) or None,
fireworks_account_id=os.getenv(ENV_FIREWORKS_ACCOUNT_ID) or None,
fireworks_batch_base_url=os.getenv(ENV_FIREWORKS_BATCH_BASE_URL) or DEFAULT_FIREWORKS_BATCH_BASE_URL,
fireworks_batch_max_wait_seconds=int(
os.getenv(ENV_FIREWORKS_BATCH_MAX_WAIT_SECONDS, str(DEFAULT_FIREWORKS_BATCH_MAX_WAIT_SECONDS))
),
retain_llm_max_concurrent=int(os.getenv(ENV_RETAIN_LLM_MAX_CONCURRENT))
if os.getenv(ENV_RETAIN_LLM_MAX_CONCURRENT)
else None,
@@ -1762,10 +1538,6 @@ class HindsightConfig:
os.getenv(ENV_EMBEDDINGS_OPENAI_BATCH_SIZE),
DEFAULT_EMBEDDINGS_OPENAI_BATCH_SIZE,
),
embeddings_openai_dimensions=_parse_optional_positive_int(
ENV_EMBEDDINGS_OPENAI_DIMENSIONS,
os.getenv(ENV_EMBEDDINGS_OPENAI_DIMENSIONS),
),
# Cohere embeddings (with backward-compatible fallback to shared API key)
embeddings_cohere_api_key=os.getenv(ENV_EMBEDDINGS_COHERE_API_KEY) or os.getenv(ENV_COHERE_API_KEY),
embeddings_cohere_model=os.getenv(ENV_EMBEDDINGS_COHERE_MODEL, DEFAULT_EMBEDDINGS_COHERE_MODEL),
@@ -1778,36 +1550,6 @@ class HindsightConfig:
or os.getenv(ENV_OPENROUTER_API_KEY)
or os.getenv(ENV_LLM_API_KEY),
embeddings_openrouter_model=os.getenv(ENV_EMBEDDINGS_OPENROUTER_MODEL, DEFAULT_EMBEDDINGS_OPENROUTER_MODEL),
# ZeroEntropy embeddings
embeddings_zeroentropy_api_key=os.getenv(ENV_EMBEDDINGS_ZEROENTROPY_API_KEY)
or os.getenv("ZEROENTROPY_API_KEY"),
embeddings_zeroentropy_model=os.getenv(
ENV_EMBEDDINGS_ZEROENTROPY_MODEL, DEFAULT_EMBEDDINGS_ZEROENTROPY_MODEL
),
embeddings_zeroentropy_base_url=os.getenv(
ENV_EMBEDDINGS_ZEROENTROPY_BASE_URL, DEFAULT_ZEROENTROPY_BASE_URL
),
embeddings_zeroentropy_dimensions=_parse_positive_int(
ENV_EMBEDDINGS_ZEROENTROPY_DIMENSIONS,
os.getenv(ENV_EMBEDDINGS_ZEROENTROPY_DIMENSIONS),
DEFAULT_EMBEDDINGS_ZEROENTROPY_DIMENSIONS,
),
embeddings_zeroentropy_encoding_format=_parse_optional_choice(
ENV_EMBEDDINGS_ZEROENTROPY_ENCODING_FORMAT,
os.getenv(ENV_EMBEDDINGS_ZEROENTROPY_ENCODING_FORMAT),
frozenset({"float", "base64"}),
)
or DEFAULT_EMBEDDINGS_ZEROENTROPY_ENCODING_FORMAT,
embeddings_zeroentropy_latency=_parse_optional_choice(
ENV_EMBEDDINGS_ZEROENTROPY_LATENCY,
os.getenv(ENV_EMBEDDINGS_ZEROENTROPY_LATENCY),
frozenset({"fast", "slow"}),
),
embeddings_zeroentropy_batch_size=_parse_positive_int(
ENV_EMBEDDINGS_ZEROENTROPY_BATCH_SIZE,
os.getenv(ENV_EMBEDDINGS_ZEROENTROPY_BATCH_SIZE),
DEFAULT_EMBEDDINGS_ZEROENTROPY_BATCH_SIZE,
),
# LiteLLM embeddings (with backward-compatible fallback to shared config)
embeddings_litellm_api_base=os.getenv(ENV_EMBEDDINGS_LITELLM_API_BASE)
or os.getenv(ENV_LITELLM_API_BASE, DEFAULT_LITELLM_API_BASE),
@@ -1880,15 +1622,11 @@ class HindsightConfig:
reranker_cohere_api_key=os.getenv(ENV_RERANKER_COHERE_API_KEY) or os.getenv(ENV_COHERE_API_KEY),
reranker_cohere_model=os.getenv(ENV_RERANKER_COHERE_MODEL, DEFAULT_RERANKER_COHERE_MODEL),
reranker_cohere_base_url=os.getenv(ENV_RERANKER_COHERE_BASE_URL) or None,
reranker_cohere_timeout=float(os.getenv(ENV_RERANKER_COHERE_TIMEOUT, str(DEFAULT_RERANKER_COHERE_TIMEOUT))),
# OpenRouter reranker (with fallback to shared OpenRouter key, then LLM key)
reranker_openrouter_api_key=os.getenv(ENV_RERANKER_OPENROUTER_API_KEY)
or os.getenv(ENV_OPENROUTER_API_KEY)
or os.getenv(ENV_LLM_API_KEY),
reranker_openrouter_model=os.getenv(ENV_RERANKER_OPENROUTER_MODEL, DEFAULT_RERANKER_OPENROUTER_MODEL),
reranker_openrouter_timeout=float(
os.getenv(ENV_RERANKER_OPENROUTER_TIMEOUT, str(DEFAULT_RERANKER_OPENROUTER_TIMEOUT))
),
# LiteLLM reranker (with backward-compatible fallback to shared config)
reranker_litellm_api_base=os.getenv(ENV_RERANKER_LITELLM_API_BASE)
or os.getenv(ENV_LITELLM_API_BASE, DEFAULT_LITELLM_API_BASE),
@@ -1897,45 +1635,26 @@ class HindsightConfig:
reranker_litellm_max_tokens_per_doc=int(v)
if (v := os.getenv(ENV_RERANKER_LITELLM_MAX_TOKENS_PER_DOC))
else DEFAULT_RERANKER_LITELLM_MAX_TOKENS_PER_DOC,
reranker_litellm_timeout=float(
os.getenv(ENV_RERANKER_LITELLM_TIMEOUT, str(DEFAULT_RERANKER_LITELLM_TIMEOUT))
),
# LiteLLM SDK reranker (direct API access)
reranker_litellm_sdk_api_key=os.getenv(ENV_RERANKER_LITELLM_SDK_API_KEY),
reranker_litellm_sdk_model=os.getenv(ENV_RERANKER_LITELLM_SDK_MODEL, DEFAULT_RERANKER_LITELLM_SDK_MODEL),
reranker_litellm_sdk_api_base=os.getenv(ENV_RERANKER_LITELLM_SDK_API_BASE) or None,
reranker_litellm_sdk_timeout=float(
os.getenv(ENV_RERANKER_LITELLM_SDK_TIMEOUT, str(DEFAULT_RERANKER_LITELLM_SDK_TIMEOUT))
),
# ZeroEntropy reranker
reranker_zeroentropy_api_key=os.getenv(ENV_RERANKER_ZEROENTROPY_API_KEY),
reranker_zeroentropy_model=os.getenv(ENV_RERANKER_ZEROENTROPY_MODEL, DEFAULT_RERANKER_ZEROENTROPY_MODEL),
reranker_zeroentropy_base_url=os.getenv(ENV_RERANKER_ZEROENTROPY_BASE_URL) or None,
reranker_zeroentropy_timeout=float(
os.getenv(ENV_RERANKER_ZEROENTROPY_TIMEOUT, str(DEFAULT_RERANKER_ZEROENTROPY_TIMEOUT))
),
# SiliconFlow reranker (Cohere-compatible /rerank endpoint)
reranker_siliconflow_api_key=os.getenv(ENV_RERANKER_SILICONFLOW_API_KEY),
reranker_siliconflow_model=os.getenv(ENV_RERANKER_SILICONFLOW_MODEL, DEFAULT_RERANKER_SILICONFLOW_MODEL),
reranker_siliconflow_base_url=os.getenv(
ENV_RERANKER_SILICONFLOW_BASE_URL, DEFAULT_RERANKER_SILICONFLOW_BASE_URL
),
reranker_siliconflow_timeout=float(
os.getenv(ENV_RERANKER_SILICONFLOW_TIMEOUT, str(DEFAULT_RERANKER_SILICONFLOW_TIMEOUT))
),
# Alibaba Cloud DashScope reranker
reranker_alibaba_api_key=os.getenv(ENV_RERANKER_ALIBABA_API_KEY),
reranker_alibaba_model=os.getenv(ENV_RERANKER_ALIBABA_MODEL, DEFAULT_RERANKER_ALIBABA_MODEL),
reranker_alibaba_timeout=float(
os.getenv(ENV_RERANKER_ALIBABA_TIMEOUT, str(DEFAULT_RERANKER_ALIBABA_TIMEOUT))
),
# Google Discovery Engine reranker (with fallback to LLM Vertex AI keys)
reranker_google_model=os.getenv(ENV_RERANKER_GOOGLE_MODEL, DEFAULT_RERANKER_GOOGLE_MODEL),
reranker_google_project_id=os.getenv(ENV_RERANKER_GOOGLE_PROJECT_ID)
or os.getenv(ENV_LLM_VERTEXAI_PROJECT_ID),
reranker_google_service_account_key=os.getenv(ENV_RERANKER_GOOGLE_SERVICE_ACCOUNT_KEY)
or os.getenv(ENV_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY),
reranker_google_timeout=float(os.getenv(ENV_RERANKER_GOOGLE_TIMEOUT, str(DEFAULT_RERANKER_GOOGLE_TIMEOUT))),
# Server
host=os.getenv(ENV_HOST, DEFAULT_HOST),
port=int(os.getenv(ENV_PORT, DEFAULT_PORT)),
@@ -1986,11 +1705,6 @@ class HindsightConfig:
retain_strategies=DEFAULT_RETAIN_STRATEGIES,
retain_batch_tokens=int(os.getenv(ENV_RETAIN_BATCH_TOKENS, str(DEFAULT_RETAIN_BATCH_TOKENS))),
retain_entity_lookup=os.getenv(ENV_RETAIN_ENTITY_LOOKUP, DEFAULT_RETAIN_ENTITY_LOOKUP),
retain_entity_resolution_batch_size=_parse_positive_int(
ENV_RETAIN_ENTITY_RESOLUTION_BATCH_SIZE,
os.getenv(ENV_RETAIN_ENTITY_RESOLUTION_BATCH_SIZE),
DEFAULT_RETAIN_ENTITY_RESOLUTION_BATCH_SIZE,
),
retain_batch_enabled=os.getenv(ENV_RETAIN_BATCH_ENABLED, str(DEFAULT_RETAIN_BATCH_ENABLED)).lower()
== "true",
retain_batch_poll_interval_seconds=int(
@@ -2030,10 +1744,6 @@ class HindsightConfig:
== "true",
# Observations settings (consolidated knowledge from facts)
enable_observations=os.getenv(ENV_ENABLE_OBSERVATIONS, str(DEFAULT_ENABLE_OBSERVATIONS)).lower() == "true",
enable_auto_consolidation=os.getenv(
ENV_ENABLE_AUTO_CONSOLIDATION, str(DEFAULT_ENABLE_AUTO_CONSOLIDATION)
).lower()
== "true",
enable_observation_history=os.getenv(
ENV_ENABLE_OBSERVATION_HISTORY, str(DEFAULT_ENABLE_OBSERVATION_HISTORY)
).lower()
@@ -2042,12 +1752,6 @@ class HindsightConfig:
ENV_ENABLE_MENTAL_MODEL_HISTORY, str(DEFAULT_ENABLE_MENTAL_MODEL_HISTORY)
).lower()
== "true",
mental_model_history_max_entries=int(
os.getenv(
ENV_MENTAL_MODEL_HISTORY_MAX_ENTRIES,
str(DEFAULT_MENTAL_MODEL_HISTORY_MAX_ENTRIES),
)
),
consolidation_batch_size=int(
os.getenv(ENV_CONSOLIDATION_BATCH_SIZE, str(DEFAULT_CONSOLIDATION_BATCH_SIZE))
),
@@ -2060,15 +1764,6 @@ class HindsightConfig:
consolidation_llm_batch_size=int(
os.getenv(ENV_CONSOLIDATION_LLM_BATCH_SIZE, str(DEFAULT_CONSOLIDATION_LLM_BATCH_SIZE))
),
consolidation_llm_parallelism=max(
1,
int(
os.getenv(
ENV_CONSOLIDATION_LLM_PARALLELISM,
str(DEFAULT_CONSOLIDATION_LLM_PARALLELISM),
)
),
),
consolidation_max_tokens=int(
os.getenv(ENV_CONSOLIDATION_MAX_TOKENS, str(DEFAULT_CONSOLIDATION_MAX_TOKENS))
),
@@ -2104,12 +1799,6 @@ class HindsightConfig:
worker_id=os.getenv(ENV_WORKER_ID) or DEFAULT_WORKER_ID,
worker_poll_interval_ms=int(os.getenv(ENV_WORKER_POLL_INTERVAL_MS, str(DEFAULT_WORKER_POLL_INTERVAL_MS))),
worker_max_retries=int(os.getenv(ENV_WORKER_MAX_RETRIES, str(DEFAULT_WORKER_MAX_RETRIES))),
worker_task_retry_backoff_seconds=int(
os.getenv(
ENV_WORKER_TASK_RETRY_BACKOFF_SECONDS,
str(DEFAULT_WORKER_TASK_RETRY_BACKOFF_SECONDS),
)
),
worker_http_port=int(os.getenv(ENV_WORKER_HTTP_PORT, str(DEFAULT_WORKER_HTTP_PORT))),
worker_max_slots=int(os.getenv(ENV_WORKER_MAX_SLOTS, str(DEFAULT_WORKER_MAX_SLOTS))),
worker_slot_reservations={
@@ -2117,9 +1806,6 @@ class HindsightConfig:
for op_type, (env_var, default) in WORKER_SLOT_RESERVATION_TYPES.items()
if int(os.getenv(env_var, str(default))) > 0
},
worker_consolidation_bank_priority=_parse_bank_priority(
os.getenv(ENV_WORKER_CONSOLIDATION_BANK_PRIORITY, "")
),
retain_max_concurrent=int(os.getenv(ENV_RETAIN_MAX_CONCURRENT, str(DEFAULT_RETAIN_MAX_CONCURRENT))),
# Reflect agent settings
reflect_max_iterations=int(os.getenv(ENV_REFLECT_MAX_ITERATIONS, str(DEFAULT_REFLECT_MAX_ITERATIONS))),
@@ -2211,8 +1897,6 @@ class HindsightConfig:
return "https://api.groq.com/openai/v1"
elif provider == "ollama":
return "http://localhost:11434/v1"
elif provider == "ollama-cloud":
return "https://ollama.com/v1"
elif provider == "lmstudio":
return "http://localhost:1234/v1"
else:
@@ -172,9 +172,8 @@ 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}")
@@ -15,13 +15,10 @@ NOTE: Observations are distinct from mental models (pinned reflections).
- Mental models: user-defined queries stored in the mental_models table, refreshed on demand via reflect
"""
import asyncio
import json
import logging
import time
import uuid
from collections import defaultdict
from contextlib import AsyncExitStack
from dataclasses import dataclass, field
from datetime import datetime, timezone
from itertools import combinations
@@ -46,91 +43,6 @@ if TYPE_CHECKING:
logger = logging.getLogger(__name__)
@dataclass
class _BatchDeltas:
"""Per-LLM-batch deltas, merged into the job's running stats after dispatch.
Returned by value rather than mutated into the outer ``stats`` /
``consolidated_tags`` so parallel batches cannot race on those shared
structures (the merge happens once, serially, after dispatch completes).
"""
stats: dict[str, int]
tags: set[str]
cancelled: bool
def _parse_observation_scopes(memory: dict[str, Any]) -> Any:
"""Parse the per-memory ``observation_scopes`` column from a DB row.
asyncpg may return JSONB as a raw JSON string depending on driver settings;
accept both that and a pre-parsed value.
"""
raw = memory.get("observation_scopes")
return json.loads(raw) if isinstance(raw, str) else raw
def _resolve_obs_tags_list(memory: dict[str, Any]) -> list[list[str]] | None:
"""Resolve a memory's ``observation_scopes`` spec into concrete scope tags.
Returns ``None`` for the default ``combined``-mode single pass (caller uses
the memory's own tags). Returns a list[list[str]] when the memory requested
multi-pass scoping (``per_tag``, ``all_combinations``, or an explicit list).
"""
parsed = _parse_observation_scopes(memory)
tags = list(memory.get("tags") or [])
if parsed == "per_tag":
return [[t] for t in tags] if tags else None
if parsed == "all_combinations":
if not tags:
return None
return [list(c) for r in range(1, len(tags) + 1) for c in combinations(tags, r)]
if parsed == "combined" or parsed is None:
return None
return parsed # explicit list[list[str]]
def _resolve_write_scopes(memory: dict[str, Any]) -> list[frozenset[str]]:
"""Return the observation scopes a memory will write to, as frozensets.
Used by the parallel dispatcher to acquire one lock per scope before
processing a tag group, so that two groups whose write-scope sets overlap
serialise on the overlapping scopes rather than racing on the same
observation row. The mapping mirrors ``_resolve_obs_tags_list`` exactly:
- ``combined`` / ``None`` -> ``[frozenset(memory.tags)]``
- ``per_tag`` -> ``[frozenset({t}) for t in memory.tags]``
- ``all_combinations`` -> one frozenset per nonempty subset of tags
- explicit ``list[list[str]]`` -> one frozenset per declared scope
Empty-tag memories collapse to a single ``frozenset()`` in all modes so they
still take exactly one lock and serialise against other untagged work.
"""
parsed = _parse_observation_scopes(memory)
tags = list(memory.get("tags") or [])
if parsed == "per_tag":
return [frozenset([t]) for t in tags] if tags else [frozenset()]
if parsed == "all_combinations":
if not tags:
return [frozenset()]
return [frozenset(c) for r in range(1, len(tags) + 1) for c in combinations(tags, r)]
if parsed == "combined" or parsed is None:
return [frozenset(tags)]
return [frozenset(s) for s in parsed] # explicit list[list[str]]
def _scope_sort_key(scope: frozenset[str]) -> tuple[str, ...]:
"""Total ordering on scope frozensets for deadlock-free lock acquisition.
Every parallel group acquires its scope locks in this same order, so two
groups that share any subset of scopes cannot acquire them in opposite
orders and deadlock.
"""
return tuple(sorted(scope))
async def _filter_live_source_memories(
conn: "Connection",
bank_id: str,
@@ -303,25 +215,6 @@ class ConsolidationPerfLog:
self.total_obs_in_context += obs_count
self.total_prompt_chars += prompt_chars
def merge_from(self, other: "ConsolidationPerfLog") -> None:
"""Merge a per-batch perf log into this (job-level) one.
Used by the parallel dispatcher: each in-flight batch records into its
own ``ConsolidationPerfLog`` so the per-batch log line shows only that
batch's timings (no cross-batch interleaving). After the batch finishes
we fold the local counters into the job-level perf, which then drives
the final ``flush()`` summary.
``lines`` is intentionally NOT merged — log lines are emitted directly
in ``logger.info`` calls by the dispatcher; the perf object's ``lines``
buffer is only used by the top-level job summary.
"""
for key, value in other.timings.items():
self.timings[key] = self.timings.get(key, 0.0) + value
self.llm_calls += other.llm_calls
self.total_obs_in_context += other.total_obs_in_context
self.total_prompt_chars += other.total_prompt_chars
def flush(self) -> None:
"""Flush all log lines to the logger."""
total_time = time.time() - self.start_time
@@ -337,7 +230,6 @@ async def run_consolidation_job(
bank_id: str,
request_context: "RequestContext",
operation_id: str | None = None,
observation_scopes: list[list[str]] | None = None,
) -> dict[str, Any]:
"""
Run consolidation job for a bank.
@@ -348,10 +240,6 @@ async def run_consolidation_job(
memory_engine: MemoryEngine instance
bank_id: Bank identifier
request_context: Request context for authentication
operation_id: Optional operation ID for tracking
observation_scopes: Optional list of tag scopes. When provided, only
unconsolidated memories whose tags contain all tags in at least one
scope are processed.
Returns:
Dict with consolidation results
@@ -393,18 +281,6 @@ async def run_consolidation_job(
perf.record_timing("fetch_bank", time.time() - t0)
# Build optional scope filter clause. When observation_scopes is provided,
# only process memories whose tags contain all tags in at least one scope.
scope_clause = ""
scope_params: list[Any] = [bank_id]
if observation_scopes:
or_parts: list[str] = []
for scope_tags in observation_scopes:
idx = len(scope_params) + 1
or_parts.append(f"tags @> ${idx}::varchar[]")
scope_params.append(scope_tags)
scope_clause = " AND (" + " OR ".join(or_parts) + ")"
# Count total unconsolidated memories for progress logging
total_count = await conn.fetchval(
f"""
@@ -414,9 +290,8 @@ async def run_consolidation_job(
AND consolidated_at IS NULL
AND consolidation_failed_at IS NULL
AND fact_type IN ('experience', 'world')
{scope_clause}
""",
*scope_params,
bank_id,
)
if total_count == 0:
@@ -446,10 +321,6 @@ async def run_consolidation_job(
hit_round_limit = False
llm_batch_num = 0
# Cumulative count of memories processed across the whole job, shared by
# the per-batch log so it can still report processed/total under parallelism.
# Mutable container so the inner closure can update without a `nonlocal`.
cumulative_progress = {"processed": 0}
while True:
# Cap fetch size by remaining round budget
fetch_limit = (
@@ -459,9 +330,6 @@ async def run_consolidation_job(
# Fetch next batch of unconsolidated memories
async with acquire_with_retry(pool) as conn:
t0 = time.time()
# scope_params[0] is bank_id; append fetch_limit after scope params
fetch_params = list(scope_params) + [fetch_limit]
limit_idx = len(fetch_params)
memories = await conn.fetch(
f"""
SELECT id, text, fact_type, occurred_start, occurred_end, event_date, tags, mentioned_at,
@@ -471,11 +339,11 @@ async def run_consolidation_job(
AND consolidated_at IS NULL
AND consolidation_failed_at IS NULL
AND fact_type IN ('experience', 'world')
{scope_clause}
ORDER BY created_at ASC
LIMIT ${limit_idx}
LIMIT $2
""",
*fetch_params,
bank_id,
fetch_limit,
)
perf.record_timing("fetch_memories", time.time() - t0)
@@ -489,63 +357,73 @@ async def run_consolidation_job(
tag_key = tuple(sorted(m.get("tags") or []))
tag_groups.setdefault(tag_key, []).append(dict(m))
# Split each tag group into LLM batches respecting llm_batch_size, keeping
# the group boundary intact so the dispatcher can parallelise across
# distinct groups while running each group's batches serially.
grouped_batches: list[list[list[dict[str, Any]]]] = []
# Flatten into LLM batches respecting both tag groups and llm_batch_size
llm_batches: list[list[dict[str, Any]]] = []
for group in tag_groups.values():
grouped_batches.append([group[i : i + llm_batch_size] for i in range(0, len(group), llm_batch_size)])
for i in range(0, len(group), llm_batch_size):
llm_batches.append(group[i : i + llm_batch_size])
# Compute each group's union write-scope set. Used below to acquire
# per-scope locks: any two groups whose write-scope sets share a scope S
# will serialise on the lock for S, leaving truly disjoint groups to run
# concurrently. We union over every memory because per-memory
# observation_scopes can differ within a group.
group_scopes: list[list[frozenset[str]]] = []
for batches in grouped_batches:
scopes: set[frozenset[str]] = set()
for batch in batches:
for memory in batch:
scopes.update(_resolve_write_scopes(memory))
group_scopes.append(sorted(scopes, key=_scope_sort_key))
async def _process_one_llm_batch(llm_batch_local: list[dict[str, Any]], batch_num_local: int) -> _BatchDeltas:
"""Process one LLM batch independently. Returns local deltas + cancelled flag.
Each batch records timings/llm-call counters into its OWN
``ConsolidationPerfLog`` so the per-batch log line reflects only
this batch's work — not interleaved timings from concurrent batches
sharing the global ``perf``. The local perf is merged into the
job-level ``perf`` once at the end so the final summary still totals
everything.
"""
for llm_batch in llm_batches:
llm_batch_num += 1
llm_batch_start = time.time()
batch_perf = ConsolidationPerfLog(bank_id)
local_tags: set[str] = set()
for memory in llm_batch_local:
# Snapshot perf and stats before this LLM batch
snap_timings = perf.timings.copy()
snap_llm_calls = perf.llm_calls
snap_total_chars = perf.total_prompt_chars
snap_stats = stats.copy()
# Track tags for mental model refresh filtering
for memory in llm_batch:
memory_tags = memory.get("tags") or []
if memory_tags:
local_tags.update(memory_tags)
consolidated_tags.update(memory_tags)
# Adaptive splitting: on LLM failure, halve the sub-batch and retry,
# down to batch_size=1. Only if a single-memory batch still fails is
# the memory marked with consolidation_failed_at.
# Process llm_batch with adaptive splitting: on LLM failure, halve the sub-batch
# and retry, down to batch_size=1. Only if a single-memory batch still fails is
# the memory marked with consolidation_failed_at and excluded from future runs
# until explicitly retried via the API.
all_results: list[dict[str, Any]] = []
all_deleted = 0
succeeded_ids: list[Any] = []
failed_ids: list[Any] = []
pending: list[list[dict[str, Any]]] = [llm_batch_local]
pending: list[list[dict[str, Any]]] = [llm_batch]
while pending:
sub_batch = pending.pop(0)
async with acquire_with_retry(pool) as conn:
obs_tags_list = _resolve_obs_tags_list(sub_batch[0]) if sub_batch else None
# Determine observation_scopes for this sub-batch. All memories share
# the same tags (enforced by tag_groups), so we only check the first memory.
# asyncpg returns JSONB columns as raw JSON strings, so parse if needed.
_obs_raw = sub_batch[0].get("observation_scopes") if sub_batch else None
_obs_parsed = json.loads(_obs_raw) if isinstance(_obs_raw, str) else _obs_raw
# Resolve the scope spec into a concrete list[list[str]] (or None for combined).
if _obs_parsed == "per_tag":
_memory_tags = sub_batch[0].get("tags") or []
obs_tags_list = [[tag] for tag in _memory_tags] if _memory_tags else None
elif _obs_parsed == "all_combinations":
_memory_tags = sub_batch[0].get("tags") or []
obs_tags_list = (
[
list(combo)
for r in range(1, len(_memory_tags) + 1)
for combo in combinations(_memory_tags, r)
]
if _memory_tags
else None
)
elif _obs_parsed == "combined" or _obs_parsed is None:
obs_tags_list = None # single combined pass (default behaviour)
else:
# explicit list[list[str]]
obs_tags_list = _obs_parsed
sub_deleted: int = 0
sub_llm_failed = False
if obs_tags_list:
# Multi-pass: run one observation consolidation pass per tag set
sub_results: list[dict[str, Any]] = []
for obs_tags in obs_tags_list:
pass_results, pass_deleted, pass_failed = await _process_memory_batch(
@@ -555,12 +433,13 @@ async def run_consolidation_job(
bank_id=bank_id,
memories=sub_batch,
request_context=request_context,
perf=batch_perf,
perf=perf,
config=config,
obs_tags_override=obs_tags,
)
sub_deleted += pass_deleted
sub_llm_failed = sub_llm_failed or pass_failed
# Merge results: prefer non-skipped actions
if not sub_results:
sub_results = pass_results
else:
@@ -568,6 +447,7 @@ async def run_consolidation_job(
if existing.get("action") == "skipped" and new.get("action") != "skipped":
sub_results[i] = new
elif existing.get("action") != "skipped" and new.get("action") != "skipped":
# Both did something — combine into "multiple"
existing_created = existing.get(
"created", 1 if existing.get("action") == "created" else 0
)
@@ -585,6 +465,7 @@ async def run_consolidation_job(
"total_actions": total,
}
else:
# Normal single pass using the memory's own tags
sub_results, sub_deleted, sub_llm_failed = await _process_memory_batch(
conn=conn,
memory_engine=memory_engine,
@@ -592,13 +473,14 @@ async def run_consolidation_job(
bank_id=bank_id,
memories=sub_batch,
request_context=request_context,
perf=batch_perf,
perf=perf,
config=config,
)
all_deleted += sub_deleted
if sub_llm_failed and len(sub_batch) > 1:
# Split and retry with smaller batches
mid = len(sub_batch) // 2
logger.warning(
f"[CONSOLIDATION] bank={bank_id} LLM failed for sub-batch of {len(sub_batch)},"
@@ -606,6 +488,7 @@ async def run_consolidation_job(
)
pending[0:0] = [sub_batch[:mid], sub_batch[mid:]]
elif sub_llm_failed:
# batch_size=1 and still failing — mark as permanently failed for now
failed_ids.append(sub_batch[0]["id"])
all_results.append({"action": "failed"})
logger.warning(
@@ -616,6 +499,7 @@ async def run_consolidation_job(
succeeded_ids.extend(m["id"] for m in sub_batch)
all_results.extend(sub_results)
# Commit consolidated_at / consolidation_failed_at in a single DB round-trip
async with acquire_with_retry(pool) as conn:
if succeeded_ids:
await conn.executemany(
@@ -628,159 +512,62 @@ async def run_consolidation_job(
[(mem_id,) for mem_id in failed_ids],
)
cancelled_local = False
stats["observations_deleted"] += all_deleted
results = all_results
# Checkpoint: abort if the operation (and thus the bank) was deleted mid-run.
if operation_id and not await memory_engine._check_op_alive(operation_id):
logger.info(
f"[CONSOLIDATION] bank={bank_id} operation {operation_id} cancelled (bank deleted), stopping early"
)
cancelled_local = True
return {"status": "cancelled", "bank_id": bank_id, **stats}
# Per-batch local stats; merged into outer state once, serially,
# after dispatch completes.
local_stats: dict[str, int] = {
"memories_processed": 0,
"observations_created": 0,
"observations_updated": 0,
"observations_merged": 0,
"observations_deleted": all_deleted,
"actions_executed": 0,
"skipped": 0,
"memories_failed": 0,
}
for result in all_results:
local_stats["memories_processed"] += 1
for result in results:
stats["memories_processed"] += 1
action = result.get("action")
if action == "created":
local_stats["observations_created"] += 1
local_stats["actions_executed"] += 1
stats["observations_created"] += 1
stats["actions_executed"] += 1
elif action == "updated":
local_stats["observations_updated"] += 1
local_stats["actions_executed"] += 1
stats["observations_updated"] += 1
stats["actions_executed"] += 1
elif action == "merged":
local_stats["observations_merged"] += 1
local_stats["actions_executed"] += 1
stats["observations_merged"] += 1
stats["actions_executed"] += 1
elif action == "multiple":
local_stats["observations_created"] += result.get("created", 0)
local_stats["observations_updated"] += result.get("updated", 0)
local_stats["observations_merged"] += result.get("merged", 0)
local_stats["actions_executed"] += result.get("total_actions", 0)
stats["observations_created"] += result.get("created", 0)
stats["observations_updated"] += result.get("updated", 0)
stats["observations_merged"] += result.get("merged", 0)
stats["actions_executed"] += result.get("total_actions", 0)
elif action == "skipped":
local_stats["skipped"] += 1
stats["skipped"] += 1
elif action == "failed":
local_stats["memories_failed"] += 1
stats["memories_failed"] += 1
# Maintain the cumulative-progress indicator under parallelism:
# increment a shared counter and snapshot under the same statement
# so the snapshot includes this batch. No await between the read
# and write, so single-threaded asyncio gives us atomicity for free
# — no lock needed.
cumulative_progress["processed"] += local_stats["memories_processed"]
cum_processed = cumulative_progress["processed"]
# Per-batch log uses batch_perf so timings/llm-calls/tokens reflect
# only this batch's own work, even when other batches are running
# concurrently under parallelism > 1. ``processed=`` is the
# cumulative count across all batches that have finished so far in
# this job (monotonic, may be reported out of strict batch-number
# order under parallelism).
# Per-LLM-batch log
llm_batch_time = time.time() - llm_batch_start
timing_parts = [
f"{key}={batch_perf.timings[key]:.3f}s"
for key in ("recall", "llm", "embedding", "db_write")
if key in batch_perf.timings
]
input_tokens = int(batch_perf.total_prompt_chars / 4)
timing_parts = []
for key in ["recall", "llm", "embedding", "db_write"]:
if key in perf.timings:
delta = perf.timings[key] - snap_timings.get(key, 0)
timing_parts.append(f"{key}={delta:.3f}s")
input_tokens = int((perf.total_prompt_chars - snap_total_chars) / 4)
batch_created = stats["observations_created"] - snap_stats["observations_created"]
batch_updated = stats["observations_updated"] - snap_stats["observations_updated"]
batch_skipped = stats["skipped"] - snap_stats["skipped"]
batch_failed = stats["memories_failed"] - snap_stats["memories_failed"]
llm_calls_made = perf.llm_calls - snap_llm_calls
logger.info(
f"[CONSOLIDATION] bank={bank_id} llm_batch #{batch_num_local}"
f" ({len(llm_batch_local)} memories, {batch_perf.llm_calls} llm calls)"
f" | processed={cum_processed}/{total_count}"
f"[CONSOLIDATION] bank={bank_id} llm_batch #{llm_batch_num}"
f" ({len(llm_batch)} memories, {llm_calls_made} llm calls)"
f" | {stats['memories_processed']}/{total_count} processed"
f" | {', '.join(timing_parts)}"
f" | created={local_stats['observations_created']}"
f" updated={local_stats['observations_updated']}"
f" skipped={local_stats['skipped']}"
+ (f" failed={local_stats['memories_failed']}" if local_stats["memories_failed"] else "")
f" | created={batch_created} updated={batch_updated} skipped={batch_skipped}"
+ (f" failed={batch_failed}" if batch_failed else "")
+ f" | input_tokens=~{input_tokens}"
f" | avg={llm_batch_time / max(1, len(llm_batch_local)):.3f}s/memory"
f" | avg={llm_batch_time / len(llm_batch):.3f}s/memory"
)
# Fold batch counters into the job-level perf so the final summary
# (perf.flush) totals every batch correctly. Safe without a lock —
# ConsolidationPerfLog.merge_from is a series of += on Python ints
# and floats with no intervening awaits, so single-threaded asyncio
# gives us atomicity.
perf.merge_from(batch_perf)
return _BatchDeltas(stats=local_stats, tags=local_tags, cancelled=cancelled_local)
# Number every batch up front so log line numbering is deterministic
# regardless of dispatch order under parallelism. Each group keeps its own
# (batch, number) list so it can be processed as one serial unit.
numbered_groups: list[list[tuple[list[dict[str, Any]], int]]] = []
for batches in grouped_batches:
numbered: list[tuple[list[dict[str, Any]], int]] = []
for b in batches:
llm_batch_num += 1
numbered.append((b, llm_batch_num))
numbered_groups.append(numbered)
async def _process_tag_group(
group_batches: list[tuple[list[dict[str, Any]], int]],
) -> list[_BatchDeltas]:
# Batches within a group share a tag set and observation scope, so
# they MUST run serially. Stop early if the op was cancelled mid-group.
deltas: list[_BatchDeltas] = []
for b, n in group_batches:
d = await _process_one_llm_batch(b, n)
deltas.append(d)
if d.cancelled:
break
return deltas
llm_parallelism = max(1, config.consolidation_llm_parallelism)
if llm_parallelism > 1 and len(numbered_groups) > 1:
sem = asyncio.Semaphore(llm_parallelism)
# Per-scope async locks shared across all parallel groups in this
# fetch iteration. Each group acquires locks for every scope it will
# write to, in _scope_sort_key order (deadlock-free). Groups with
# disjoint scope sets never contend; any overlap serialises on the
# overlapping scopes — covering combined / per_tag / all_combinations
# / explicit-list modes uniformly without operator opt-in.
scope_locks: defaultdict[frozenset[str], asyncio.Lock] = defaultdict(asyncio.Lock)
async def _run_group(
group_batches: list[tuple[list[dict[str, Any]], int]],
scopes: list[frozenset[str]],
) -> list[_BatchDeltas]:
async with sem:
async with AsyncExitStack() as stack:
for s in scopes:
await stack.enter_async_context(scope_locks[s])
return await _process_tag_group(group_batches)
group_results = await asyncio.gather(*(_run_group(g, s) for g, s in zip(numbered_groups, group_scopes)))
batch_results: list[_BatchDeltas] = [d for gd in group_results for d in gd]
any_cancelled = any(d.cancelled for d in batch_results)
else:
batch_results = []
any_cancelled = False
for g in numbered_groups:
group_deltas = await _process_tag_group(g)
batch_results.extend(group_deltas)
if any(d.cancelled for d in group_deltas):
any_cancelled = True
break
# Merge per-batch deltas into outer state — serial, post-dispatch, so
# concurrent batches cannot race on the shared counters / tag set.
for d in batch_results:
for k, v in d.stats.items():
stats[k] = stats.get(k, 0) + v
consolidated_tags.update(d.tags)
if any_cancelled:
return {"status": "cancelled", "bank_id": bank_id, **stats}
# Update round budget after processing this DB fetch batch
if round_limit_enabled:
round_remaining -= len(memories)
@@ -788,24 +575,17 @@ async def run_consolidation_job(
hit_round_limit = True
break
# Re-submit consolidation if we hit the round limit and there's likely more work.
# Any failure here must propagate: swallowing it (the prior behavior) leaves the
# bank with backlog and no queued work — silently stuck — because the outer op
# gets marked completed in the success path. Letting the exception bubble up to
# execute_task's retry handler means the op is retried with backoff; on retry the
# consolidator skips already-consolidated rows via the consolidated_at filter and
# picks up the remainder. Issue #1842.
# Re-submit consolidation if we hit the round limit and there's likely more work
if hit_round_limit:
remaining = total_count - stats["memories_processed"]
logger.info(
f"[CONSOLIDATION] bank={bank_id} hit round limit of {max_memories_per_round} memories,"
f" ~{remaining} remaining. Re-queuing consolidation."
)
await memory_engine.submit_async_consolidation(
bank_id=bank_id,
request_context=request_context,
observation_scopes=observation_scopes,
)
try:
await memory_engine.submit_async_consolidation(bank_id=bank_id, request_context=request_context)
except Exception as e:
logger.warning(f"[CONSOLIDATION] bank={bank_id} failed to re-queue consolidation: {e}")
# Build summary
perf.log(
@@ -1446,44 +1226,6 @@ def _build_observations_for_llm(
return obs_list
def _dedupe_updates(updates: list[_UpdateAction], *, batch_label: str) -> list[_UpdateAction]:
"""Collapse `updates` that target the same `observation_id`.
LLMs occasionally emit several update entries for one observation in a
single response (one per facet drawn from the same fact). Without
deduplication the downstream loop would issue separate DB writes for each
and the last write would silently overwrite the earlier ones. We keep the
last text (the LLM's most recent attempt) and union all contributing
`source_fact_ids`, then warn so the misbehavior is visible in logs.
"""
if len(updates) < 2:
return list(updates)
by_id: dict[str, _UpdateAction] = {}
collisions = 0
for upd in updates:
existing = by_id.get(upd.observation_id)
if existing is None:
by_id[upd.observation_id] = upd
continue
collisions += 1
merged_ids = list(dict.fromkeys([*existing.source_fact_ids, *upd.source_fact_ids]))
by_id[upd.observation_id] = _UpdateAction(
text=upd.text,
observation_id=upd.observation_id,
source_fact_ids=merged_ids,
)
if collisions:
logger.warning(
f"[CONSOLIDATION] {batch_label}: LLM emitted {collisions} duplicate update(s) targeting "
f"the same observation_id ({len(updates)} updates -> {len(by_id)} after dedup). "
"Kept the last text and unioned source_fact_ids."
)
return list(by_id.values())
async def _consolidate_batch_with_llm(
llm_config: Any,
memories: list[dict[str, Any]],
@@ -1532,11 +1274,7 @@ async def _consolidate_batch_with_llm(
f"(out of {max_observations_per_scope}). Prefer UPDATE over CREATE when possible."
)
prompt_template = build_batch_consolidation_prompt(
config.observations_mission,
observation_capacity_note,
llm_output_language=getattr(config, "llm_output_language", None),
)
prompt_template = build_batch_consolidation_prompt(config.observations_mission, observation_capacity_note)
prompt = prompt_template.format(
facts_text=facts_lines,
observations_text=observations_text,
@@ -1577,10 +1315,9 @@ async def _consolidate_batch_with_llm(
f"(max_observations_per_scope={max_observations_per_scope})"
)
creates = creates[:remaining_observation_slots]
updates = _dedupe_updates(response.updates, batch_label=batch_label)
return _BatchLLMResult(
creates=creates,
updates=updates,
updates=response.updates,
deletes=response.deletes,
obs_count=len(union_observations),
prompt_chars=len(prompt),
@@ -1649,16 +1386,9 @@ async def _create_observation_directly(
tokenize($3, 'llmlingua2')::bm25_catalog.bm25vector)
RETURNING id
"""
else: # native, pg_textsearch, pgroonga, or pg_search
# pg_textsearch / pgroonga / pg_search: indexes operate on base text
# columns directly, so the dummy search_vector column is left NULL.
# Native: the migration p4q5r6s7t8u9 dropped the GENERATED expression on
# search_vector to allow per-deployment language configuration; the
# batch insert path in ops_postgresql.insert_facts_batch now populates
# it via to_tsvector($lang, ...). This single-observation INSERT does
# not, so observations under the native backend currently land with
# NULL search_vector and are not BM25-searchable until reflected/
# re-ingested. Tracking a separate fix for that gap.
else: # native or pg_textsearch
# Native PostgreSQL: search_vector is GENERATED ALWAYS, don't include it
# pg_textsearch: indexes operate on base columns directly, don't populate search_vector
query = f"""
INSERT INTO {fq_table("memory_units")} (
id, bank_id, text, fact_type, embedding, proof_count, source_memory_ids, history,
@@ -1,147 +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."""
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`.
### 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"]}}],
"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"]}}],
"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"]}}],
"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.
- 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
)
@@ -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,
@@ -38,8 +37,6 @@ from ..config import (
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,
@@ -63,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.
@@ -306,28 +266,25 @@ class LocalSTCrossEncoder(CrossEncoderModel):
"""
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]:
"""
@@ -589,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:
@@ -774,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__(
@@ -1007,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]:
"""
@@ -1582,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.
@@ -1666,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
@@ -1679,7 +1588,6 @@ def create_cross_encoder_from_env() -> CrossEncoderModel:
api_key=api_key,
model=config.reranker_openrouter_model,
base_url="https://openrouter.ai/api/v1/rerank",
timeout=config.reranker_openrouter_timeout,
)
elif provider == "flashrank":
model = os.environ.get(ENV_RERANKER_FLASHRANK_MODEL, DEFAULT_RERANKER_FLASHRANK_MODEL)
@@ -1694,7 +1602,6 @@ 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
@@ -1707,7 +1614,6 @@ def create_cross_encoder_from_env() -> CrossEncoderModel:
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
@@ -1718,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
@@ -1731,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
@@ -1744,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()
@@ -1761,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'"
)
@@ -166,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,
@@ -391,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
@@ -469,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.
@@ -478,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.
"""
@@ -215,96 +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.
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 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,
@@ -787,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 = []
@@ -1048,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}
@@ -1056,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"""
@@ -1127,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"""
@@ -1137,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"])
@@ -104,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(
@@ -200,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]
@@ -225,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,
)
@@ -320,100 +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
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,
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,
@@ -867,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 = []
@@ -1146,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"""
@@ -1217,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"""
@@ -1227,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"])
@@ -72,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+)'")
@@ -352,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)
@@ -542,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)
@@ -737,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+)\*/")
@@ -809,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.
@@ -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,17 +9,13 @@ 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,
@@ -31,15 +27,10 @@ from ..config import (
DEFAULT_EMBEDDINGS_LOCAL_TRUST_REMOTE_CODE,
DEFAULT_EMBEDDINGS_OPENAI_MODEL,
DEFAULT_EMBEDDINGS_PROVIDER,
DEFAULT_EMBEDDINGS_ZEROENTROPY_BATCH_SIZE,
DEFAULT_EMBEDDINGS_ZEROENTROPY_DIMENSIONS,
DEFAULT_EMBEDDINGS_ZEROENTROPY_ENCODING_FORMAT,
DEFAULT_EMBEDDINGS_ZEROENTROPY_LATENCY,
DEFAULT_EMBEDDINGS_ZEROENTROPY_MODEL,
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,
@@ -48,39 +39,12 @@ from ..config import (
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,
)
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.
@@ -124,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):
"""
@@ -429,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,
):
"""
@@ -440,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
@@ -492,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
@@ -529,14 +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
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)
@@ -545,73 +492,6 @@ class OpenAIEmbeddings(Embeddings):
return all_embeddings
class CodexOAuthEmbeddings(OpenAIEmbeddings):
"""
OpenAI embeddings using the Codex/ChatGPT OAuth token from ``~/.codex/auth.json``.
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.
@@ -753,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.
@@ -1029,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,
@@ -1041,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)
@@ -1092,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:
@@ -1145,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:
@@ -1406,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
@@ -1427,23 +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 == "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
@@ -1462,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,
@@ -1492,6 +1206,5 @@ def create_embeddings_from_env() -> Embeddings:
else:
raise ValueError(
f"Unknown embeddings provider: {provider}. "
f"Supported: 'local', 'tei', 'openai', 'openai-codex', 'openrouter', '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
@@ -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)
@@ -9,7 +9,6 @@ import os
import re
import time
import uuid
from contextlib import AsyncExitStack
from pathlib import Path
from typing import Any
@@ -28,12 +27,9 @@ except ImportError:
from ..config import (
DEFAULT_LLM_MAX_CONCURRENT,
DEFAULT_LLM_TIMEOUT,
ENV_CONSOLIDATION_LLM_MAX_CONCURRENT,
ENV_LLM_GROQ_SERVICE_TIER,
ENV_LLM_MAX_CONCURRENT,
ENV_LLM_TIMEOUT,
ENV_REFLECT_LLM_MAX_CONCURRENT,
ENV_RETAIN_LLM_MAX_CONCURRENT,
)
from ..metrics import get_metrics_collector
from .response_models import TokenUsage
@@ -46,75 +42,13 @@ logger = logging.getLogger(__name__)
# Disable httpx logging
logging.getLogger("httpx").setLevel(logging.WARNING)
# Global semaphore to limit concurrent LLM requests across all instances.
# Set HINDSIGHT_API_LLM_MAX_CONCURRENT=1 for local LLMs (LM Studio, Ollama).
# Global semaphore to limit concurrent LLM requests across all instances
# Set HINDSIGHT_API_LLM_MAX_CONCURRENT=1 for local LLMs (LM Studio, Ollama)
_llm_max_concurrent = int(os.getenv(ENV_LLM_MAX_CONCURRENT, str(DEFAULT_LLM_MAX_CONCURRENT)))
_global_llm_semaphore = asyncio.Semaphore(_llm_max_concurrent)
def _build_per_op_semaphores() -> dict[str, asyncio.Semaphore]:
"""Build the per-operation semaphore registry from env vars.
Each per-op cap is composed with — not a substitute for — the global cap:
a call that matches a configured operation must acquire both its per-op
semaphore and the global semaphore. This lets operators reserve headroom
in the global pool by capping individual operations (e.g. cap retain at 2
of 4 global slots so the live chat path always has 2 slots available).
Operations without a configured env var are absent from the registry and
therefore only constrained by the global cap.
"""
semaphores: dict[str, asyncio.Semaphore] = {}
for op, env_var in (
("retain", ENV_RETAIN_LLM_MAX_CONCURRENT),
("reflect", ENV_REFLECT_LLM_MAX_CONCURRENT),
("consolidation", ENV_CONSOLIDATION_LLM_MAX_CONCURRENT),
):
raw = os.getenv(env_var)
if raw is None or raw == "":
continue
value = int(raw)
if value <= 0:
raise ValueError(f"{env_var} must be a positive integer, got {raw!r}")
semaphores[op] = asyncio.Semaphore(value)
return semaphores
_per_op_llm_semaphores: dict[str, asyncio.Semaphore] = _build_per_op_semaphores()
def _scope_to_operation(scope: str) -> str | None:
"""Map a call scope to its per-operation concurrency bucket.
Returns None for scopes that don't belong to a tracked operation
(verification probes, bank_mission, memory_think, mental_model_delta_ops),
which then run under the global cap only.
"""
if scope.startswith("retain"):
return "retain"
if scope.startswith("reflect"):
return "reflect"
if scope.startswith("consolidation"):
return "consolidation"
return None
def _semaphores_for_scope(scope: str) -> list[asyncio.Semaphore]:
"""Return the semaphores a call with the given scope must acquire.
Always includes the global semaphore; includes the per-op semaphore when
one is configured for the scope's operation bucket.
"""
op = _scope_to_operation(scope)
per_op = _per_op_llm_semaphores.get(op) if op is not None else None
if per_op is None:
return [_global_llm_semaphore]
# Per-op acquired first so contention queues on the narrower cap before
# holding a global slot.
return [per_op, _global_llm_semaphore]
def sanitize_text(text: str | None) -> str | None:
def sanitize_llm_output(text: str | None) -> str | None:
"""
Sanitize text by removing characters that break downstream systems.
@@ -126,12 +60,8 @@ def sanitize_text(text: str | None) -> str | None:
Surrogate characters are used in UTF-16 encoding but cannot be encoded
in UTF-8. They can appear in Python strings from improperly decoded data
(e.g., from JavaScript or broken files): a client may serialize a half-emoji
split at a boundary as a lone ``\\udXXX`` escape. Such input crashes the
SentenceTransformers/cross-encoder Rust tokenizers and stdout logging, so
user content is sanitized at the retain/recall/reflect ingress (see issue
#1875). Control characters commonly appear in LLM output embedded inside
JSON string values.
(e.g., from JavaScript or broken files). Control characters commonly appear
in LLM output embedded inside JSON string values.
"""
if text is None:
return None
@@ -140,11 +70,6 @@ def sanitize_text(text: str | None) -> str | None:
return re.sub(r"[\x00-\x08\x0b\x0c\x0e-\x1f\x7f\ud800-\udfff]", "", text)
# Back-compat alias: this helper was originally introduced to scrub LLM *output*;
# it now also scrubs user *input* at ingress, hence the broader name.
sanitize_llm_output = sanitize_text
class OutputTooLongError(Exception):
"""
Bridge exception raised when LLM output exceeds token limits.
@@ -258,7 +183,6 @@ def create_llm_provider(
AnthropicLLM,
ClaudeCodeLLM,
CodexLLM,
FireworksLLM,
GeminiLLM,
LiteLLMLLM,
LiteLLMRouterLLM,
@@ -384,24 +308,10 @@ def create_llm_provider(
extra_args=config.llamacpp_extra_args,
)
elif provider_lower == "fireworks":
# Fireworks online inference is OpenAI-compatible; FireworksLLM adds the
# native (non-OpenAI) batch API on top. The existing LiteLLM
# ``fireworks_ai/...`` online path (provider="litellm") is untouched.
return FireworksLLM(
provider=provider,
api_key=api_key,
base_url=base_url,
model=model,
reasoning_effort=reasoning_effort,
extra_body=extra_body,
)
elif provider_lower in (
"openai",
"groq",
"ollama",
"ollama-cloud",
"lmstudio",
"minimax",
"deepseek",
@@ -499,7 +409,6 @@ class LLMProvider:
"openai",
"groq",
"ollama",
"ollama-cloud",
"gemini",
"anthropic",
"lmstudio",
@@ -518,7 +427,6 @@ class LLMProvider:
"openrouter",
"zai",
"opencode-go",
"fireworks",
]
if self.provider not in valid_providers:
raise ValueError(f"Invalid LLM provider: {self.provider}. Must be one of: {', '.join(valid_providers)}")
@@ -529,8 +437,6 @@ class LLMProvider:
self.base_url = "https://api.groq.com/openai/v1"
elif self.provider == "ollama":
self.base_url = "http://localhost:11434/v1"
elif self.provider == "ollama-cloud":
self.base_url = "https://ollama.com/v1"
elif self.provider == "lmstudio":
self.base_url = "http://localhost:1234/v1"
elif self.provider == "minimax":
@@ -723,10 +629,7 @@ class LLMProvider:
structured = "+structured" if response_format is not None else ""
set_stage(f"llm.{self.provider}.{scope}{structured}")
async with AsyncExitStack() as stack:
for sem in _semaphores_for_scope(scope):
await stack.enter_async_context(sem)
async with _global_llm_semaphore:
# Delegate to provider implementation
result = await self._provider_impl.call(
messages=messages,
@@ -751,7 +654,7 @@ class LLMProvider:
# Sync the mock calls from provider implementation to wrapper
self._mock_calls = self._provider_impl.get_mock_calls()
return result
return result
async def call_with_tools(
self,
@@ -786,10 +689,7 @@ class LLMProvider:
set_stage(f"llm.{self.provider}.{scope}+tools")
async with AsyncExitStack() as stack:
for sem in _semaphores_for_scope(scope):
await stack.enter_async_context(sem)
async with _global_llm_semaphore:
# Delegate to provider implementation
result = await self._provider_impl.call_with_tools(
messages=messages,
@@ -812,7 +712,7 @@ class LLMProvider:
# Sync the mock calls from provider implementation to wrapper
self._mock_calls = self._provider_impl.get_mock_calls()
return result
return result
def set_response_callback(self, fn: Any) -> None:
"""Set a callback invoked on each call() instead of the fixed mock response."""
@@ -941,14 +841,12 @@ class LLMProvider:
"""Create provider from environment variables using config.py constants."""
from ..config import (
DEFAULT_LLM_PROVIDER,
DEFAULT_LLM_REASONING_EFFORT,
ENV_LLM_API_KEY,
ENV_LLM_BASE_URL,
ENV_LLM_DEFAULT_HEADERS,
ENV_LLM_EXTRA_BODY,
ENV_LLM_MODEL,
ENV_LLM_PROVIDER,
ENV_LLM_REASONING_EFFORT,
_get_default_model_for_provider,
)
@@ -972,7 +870,7 @@ class LLMProvider:
api_key=api_key,
base_url=base_url,
model=model,
reasoning_effort=os.getenv(ENV_LLM_REASONING_EFFORT, DEFAULT_LLM_REASONING_EFFORT),
reasoning_effort="low",
extra_body=extra_body,
default_headers=default_headers,
)
File diff suppressed because it is too large Load Diff
@@ -1,41 +0,0 @@
"""Shared utilities for prompt assembly."""
import re
_LONE_OPEN_BRACE = re.compile(r"(?<!\{)\{(?!\{)")
_LONE_CLOSE_BRACE = re.compile(r"(?<!\})\}(?!\})")
def escape_for_prompt(text: str) -> str:
"""Double any lone ``{`` / ``}`` so the text survives ``str.format`` untouched.
Prompt templates are often passed through ``str.format`` to substitute real
placeholders like ``{facts_text}``. Any literal braces in caller-supplied
text — e.g. a bank mission that contains JSON examples — would otherwise be
interpreted as format keys and raise ``KeyError``.
Idempotent: text that already contains escaped ``{{`` / ``}}`` pairs is
left as-is. Only lone braces (not adjacent to another brace of the same
kind) are doubled.
"""
text = _LONE_OPEN_BRACE.sub("{{", text)
text = _LONE_CLOSE_BRACE.sub("}}", text)
return text
def output_language_directive(language: str | None) -> str:
"""Return an LLM directive forcing all output into ``language``.
Used by retain (fact extraction), consolidation (observations), and reflect
(response synthesis) so HINDSIGHT_API_LLM_OUTPUT_LANGUAGE applies uniformly
across every LLM-generated artifact. Returns an empty string when
``language`` is unset so the calling prompt stays unchanged.
"""
if not language:
return ""
return (
f"\n\nIMPORTANT: Respond exclusively in {language}. "
f"Translate any source content into {language}. "
f"All output text — including fact text, observations, entity names, "
f"and the final response — must be in {language}."
)
@@ -7,7 +7,6 @@ This package contains concrete implementations of the LLMInterface for various p
from .anthropic_llm import AnthropicLLM
from .claude_code_llm import ClaudeCodeLLM
from .codex_llm import CodexLLM
from .fireworks_llm import FireworksLLM
from .gemini_llm import GeminiLLM
from .litellm_llm import LiteLLMLLM
from .litellm_router_llm import LiteLLMRouterLLM
@@ -20,7 +19,6 @@ __all__ = [
"AnthropicLLM",
"ClaudeCodeLLM",
"CodexLLM",
"FireworksLLM",
"GeminiLLM",
"LlamaCppLLM",
"LiteLLMLLM",
@@ -93,6 +93,7 @@ class AnthropicLLM(LLMInterface):
await self.call(
messages=test_messages,
max_completion_tokens=10,
temperature=0.0,
scope="verification",
max_retries=0,
)
@@ -178,6 +179,9 @@ class AnthropicLLM(LLMInterface):
if system_prompt:
call_params["system"] = system_prompt
if temperature is not None:
call_params["temperature"] = temperature
last_exception = None
for attempt in range(max_retries + 1):
@@ -394,6 +398,9 @@ class AnthropicLLM(LLMInterface):
if system_prompt:
call_params["system"] = system_prompt
if temperature is not None:
call_params["temperature"] = temperature
last_exception = None
for attempt in range(max_retries + 1):
try:
@@ -9,7 +9,6 @@ automatically handles authentication via `claude auth login` credentials.
import asyncio
import json
import logging
import tempfile
import time
from typing import Any
@@ -22,32 +21,6 @@ from hindsight_api.metrics import get_metrics_collector
logger = logging.getLogger(__name__)
# Isolation env passed to the spawned `claude` CLI. CLAUDE_CONFIG_DIR
# redirects the subprocess away from the host's ~/.claude/, so any
# operator-installed plugins (e.g. hindsight-memory) and their Stop hooks do
# not fire inside our LLM-call subprocesses. Without this, retain/reflect/
# consolidation LLM calls would trigger a Stop-hook retain of the subprocess
# transcript back into the same bank — a recursive feedback loop (issue #1751).
# CLAUDE_SECURESTORAGE_CONFIG_DIR="" forces the CLI's keychain service name
# back to the canonical un-suffixed entry that `claude auth login` wrote;
# otherwise it would be namespaced by sha256(CLAUDE_CONFIG_DIR) and OAuth
# lookup would fail. Requires bundled CLI >= 2.1.150 (claude-agent-sdk 0.2.82).
_isolated_claude_env: dict[str, str] | None = None
def _get_isolated_claude_env() -> dict[str, str]:
"""Return a process-lifetime env dict that isolates the spawned CLI from user plugins."""
global _isolated_claude_env
if _isolated_claude_env is None:
path = tempfile.mkdtemp(prefix="hindsight-claude-code-")
_isolated_claude_env = {
"CLAUDE_CONFIG_DIR": path,
"CLAUDE_SECURESTORAGE_CONFIG_DIR": "",
}
logger.debug(f"Claude Code: isolated CLAUDE_CONFIG_DIR={path}")
return _isolated_claude_env
class ClaudeCodeLLM(LLMInterface):
"""
LLM provider using Claude Code authentication.
@@ -210,7 +183,6 @@ class ClaudeCodeLLM(LLMInterface):
system_prompt=system_prompt if system_prompt else None,
max_turns=1, # Single-turn for API-style interactions
allowed_tools=[], # Disable tools for standard LLM calls
env=_get_isolated_claude_env(),
)
# Call Claude Agent SDK
@@ -501,7 +473,6 @@ class ClaudeCodeLLM(LLMInterface):
max_turns=2, # Allow tool call + tool result round-trip
mcp_servers=mcp_servers_config,
allowed_tools=allowed_tool_names,
env=_get_isolated_claude_env(),
)
# Call Claude Agent SDK with retry logic
@@ -1,409 +0,0 @@
"""
Shared Codex OAuth authentication manager.
Extracted from ``CodexLLM`` so that both ``CodexLLM`` and
``CodexOAuthEmbeddings`` can share JWT-expiry detection, single-flight
token refresh, and atomic file persistence without duplicating the logic.
Usage
-----
Create a manager from the auth file::
mgr = CodexAuthManager.from_file()
Then call ``ensure_fresh_token()`` before each outbound request and
``refresh_tokens(reason=..., force=...)`` on a reactive 401.
"""
from __future__ import annotations
import base64
import binascii
import json
import logging
import os
import tempfile
import threading
import time
from datetime import datetime, timezone
from pathlib import Path
from typing import Any
import httpx
logger = logging.getLogger(__name__)
# ---------------------------------------------------------------------------
# Module-level constants (shared with codex_llm.py via re-export there)
# ---------------------------------------------------------------------------
# OAuth refresh endpoint and client id, mirrored from the canonical
# ``@openai/codex`` CLI (codex-rs/login/src/auth/manager.rs on
# github.com/openai/codex). The endpoint is overridable via env var so that
# future Codex changes or staging environments can be pointed at without a
# code change — same env var name the upstream CLI uses.
_CODEX_REFRESH_TOKEN_URL = os.environ.get("CODEX_REFRESH_TOKEN_URL_OVERRIDE", "https://auth.openai.com/oauth/token")
_CODEX_CLIENT_ID = "app_EMoamEEZ73f0CkXaXp7hrann"
# Proactively refresh this many seconds before the JWT ``exp`` claim. The
# upstream Codex CLI uses no skew (it refreshes at ``exp <= now``); the
# extra window reduces races where a request leaves the client with a token
# that the server has already declared expired by the time it arrives.
_CODEX_TOKEN_REFRESH_SKEW_SECONDS = 60
# OAuth error codes that the refresh endpoint returns when the refresh_token
# itself is no longer usable. These are terminal — retrying refresh will not
# succeed; the user must re-run ``codex auth login``.
_CODEX_TERMINAL_REFRESH_ERROR_CODES = frozenset(
{"refresh_token_expired", "refresh_token_reused", "refresh_token_invalidated"}
)
class CodexRefreshExpiredError(RuntimeError):
"""Raised when the Codex refresh_token itself is no longer valid.
The user must re-run ``codex auth login`` to obtain new credentials.
Callers should surface a clear remediation message and stop retrying.
"""
class CodexAuthManager:
"""Sync Codex OAuth credential manager.
Holds the access_token, refresh_token, and account_id in memory and
handles proactive/reactive refresh using a ``threading.Lock`` for
single-flight semantics (safe to use from multiple threads or via
``asyncio.to_thread``).
Parameters
----------
access_token:
The current bearer token.
account_id:
The OpenAI account ID embedded in the Codex request headers.
refresh_token:
The OAuth refresh token. May be ``None`` when the auth file omits it;
the provider still works as a one-shot loader in that case.
auth_file:
Path to ``~/.codex/auth.json``. Used for re-reading the refresh token
on demand and for atomic persistence of rotated credentials.
"""
def __init__(
self,
access_token: str,
account_id: str,
refresh_token: str | None,
auth_file: Path,
) -> None:
self.access_token = access_token
self.account_id = account_id
self.refresh_token = refresh_token
self._auth_file = auth_file
self._lock = threading.Lock()
self._http_client = httpx.Client(timeout=30.0)
# ------------------------------------------------------------------
# Construction helpers
# ------------------------------------------------------------------
@classmethod
def from_file(cls, auth_file: Path | None = None) -> "CodexAuthManager":
"""Build a manager by reading credentials from ``auth_file``.
Parameters
----------
auth_file:
Defaults to ``~/.codex/auth.json``.
Raises
------
FileNotFoundError:
If the auth file does not exist.
ValueError:
If the auth file is missing ``access_token`` or has an unexpected
``auth_mode``.
"""
if auth_file is None:
auth_file = Path.home() / ".codex" / "auth.json"
if not auth_file.exists():
raise FileNotFoundError(f"Codex auth file not found: {auth_file}. Run 'codex auth login' to authenticate.")
with open(auth_file) as f:
data = json.load(f)
auth_mode = data.get("auth_mode")
if auth_mode != "chatgpt":
raise ValueError(f"Expected Codex auth_mode='chatgpt', got: {auth_mode}")
tokens = data.get("tokens") or {}
access_token = tokens.get("access_token")
if not access_token:
raise ValueError("No access_token found in Codex auth file. Run 'codex auth login' again.")
account_id = tokens.get("account_id") or ""
refresh_token = tokens.get("refresh_token")
return cls(
access_token=access_token,
account_id=account_id,
refresh_token=refresh_token,
auth_file=auth_file,
)
# ------------------------------------------------------------------
# Token state helpers
# ------------------------------------------------------------------
@staticmethod
def load_refresh_token_from_file(auth_file: Path) -> str | None:
"""Read ``tokens.refresh_token`` from ``auth_file``.
Returns ``None`` when the file is unreadable or omits the field.
Does not raise — the provider degrades to one-shot mode.
"""
try:
with open(auth_file) as f:
data = json.load(f)
except (OSError, json.JSONDecodeError) as e:
logger.warning(
f"Codex auth file unreadable when loading refresh_token: {type(e).__name__}. "
"Token refresh will not be available; the access_token in memory will be used until it expires."
)
return None
return data.get("tokens", {}).get("refresh_token")
@staticmethod
def _decode_jwt_exp_unixtime(token: str) -> int | None:
"""Return the JWT ``exp`` claim as a unix timestamp, or None on parse failure.
We do not verify the signature — the server is the source of truth
on whether the token is actually accepted. This is only used to
schedule proactive refresh.
"""
try:
parts = token.split(".")
if len(parts) < 2:
return None
payload_b64 = parts[1]
padding = "=" * (-len(payload_b64) % 4)
payload_bytes = base64.urlsafe_b64decode(payload_b64 + padding)
payload = json.loads(payload_bytes.decode("utf-8"))
exp = payload.get("exp")
return int(exp) if exp is not None else None
except (ValueError, TypeError, json.JSONDecodeError, binascii.Error):
return None
def _token_is_stale(self, skew_seconds: int = _CODEX_TOKEN_REFRESH_SKEW_SECONDS) -> bool:
"""True when the cached access_token is past expiry (with skew).
Returns False when expiry cannot be determined — we'd rather use a
possibly-expired token and recover via the reactive 401 path than
refresh aggressively on every request when ``exp`` parsing fails.
"""
exp = self._decode_jwt_exp_unixtime(self.access_token)
if exp is None:
return False
return exp <= int(time.time()) + skew_seconds
# ------------------------------------------------------------------
# Persistence
# ------------------------------------------------------------------
def _persist_auth_atomic(self, updated_tokens: dict[str, Any]) -> None:
"""Write rotated tokens back to ``_auth_file`` atomically.
Re-reads the on-disk file first to avoid clobbering fields written
by another process, patches ``tokens.*`` and ``last_refresh``, then
writes to a sibling tempfile and calls ``os.replace`` (atomic on
POSIX and Windows within the same filesystem).
"""
current: dict[str, Any]
try:
with open(self._auth_file) as f:
loaded = json.load(f)
current = loaded if isinstance(loaded, dict) else {"auth_mode": "chatgpt", "tokens": {}}
except (OSError, json.JSONDecodeError):
current = {"auth_mode": "chatgpt", "tokens": {}}
existing_tokens = current.get("tokens")
tokens: dict[str, Any] = existing_tokens if isinstance(existing_tokens, dict) else {}
for key in ("access_token", "refresh_token", "id_token", "account_id"):
if key in updated_tokens and updated_tokens[key] is not None:
tokens[key] = updated_tokens[key]
current["tokens"] = tokens
current["last_refresh"] = datetime.now(timezone.utc).isoformat().replace("+00:00", "Z")
parent = self._auth_file.parent
parent.mkdir(parents=True, exist_ok=True)
fd, tmp_path = tempfile.mkstemp(prefix=".auth.", suffix=".json.tmp", dir=str(parent))
try:
with os.fdopen(fd, "w") as f:
json.dump(current, f, indent=2)
f.flush()
os.fsync(f.fileno())
try:
os.chmod(tmp_path, 0o600)
except OSError:
pass
os.replace(tmp_path, self._auth_file)
except Exception:
try:
os.unlink(tmp_path)
except OSError:
pass
raise
# ------------------------------------------------------------------
# Error extraction
# ------------------------------------------------------------------
@staticmethod
def _extract_oauth_error_code(response: httpx.Response) -> str | None:
"""Pull the OAuth error code out of a 4xx response body, if present.
The refresh endpoint returns shapes like
``{"error": "...", "error_code": "..."}`` or
``{"error": {"code": "..."}}``.
"""
try:
body = response.json()
except (json.JSONDecodeError, ValueError):
return None
if not isinstance(body, dict):
return None
err = body.get("error")
if isinstance(err, dict):
code = err.get("code")
if isinstance(code, str):
return code
code = body.get("error_code")
if isinstance(code, str):
return code
if isinstance(err, str):
return err
return None
# ------------------------------------------------------------------
# Refresh
# ------------------------------------------------------------------
def refresh_tokens(self, reason: str = "", *, force: bool = False) -> None:
"""Synchronous single-flight OAuth token refresh.
Serialized through ``self._lock`` so concurrent threads produce one
network request. The first caller refreshes; the rest wake up and
skip if the token is no longer stale (proactive) or if the token
has already changed (reactive / force).
Parameters
----------
reason:
Free-form string included in log lines for diagnostics.
force:
When True, refresh even if the JWT exp claim looks fresh.
Used by the reactive 401 path.
Raises
------
CodexRefreshExpiredError:
When the server returns a terminal error code or any 401.
RuntimeError:
For other refresh failures (network, 5xx, etc.).
"""
token_before_lock = self.access_token
with self._lock:
if force:
if self.access_token != token_before_lock:
return
else:
if not self._token_is_stale():
return
if not self.refresh_token:
raise RuntimeError(
"Codex access_token is expired but no refresh_token is available. "
"Run 'codex auth login' to re-authenticate."
)
log_reason = f" ({reason})" if reason else ""
logger.info(f"Refreshing Codex OAuth access_token{log_reason}")
request_body = {
"client_id": _CODEX_CLIENT_ID,
"grant_type": "refresh_token",
"refresh_token": self.refresh_token,
}
try:
response = self._http_client.post(
_CODEX_REFRESH_TOKEN_URL,
json=request_body,
headers={"Content-Type": "application/json"},
timeout=30.0,
)
except httpx.RequestError as e:
raise RuntimeError(f"Codex OAuth refresh network error: {type(e).__name__}") from e
if response.status_code == 401:
error_code = self._extract_oauth_error_code(response)
if error_code in _CODEX_TERMINAL_REFRESH_ERROR_CODES:
raise CodexRefreshExpiredError(
f"Codex refresh_token is permanently invalid (error.code={error_code}). "
"Run 'codex auth login' to re-authenticate."
)
raise CodexRefreshExpiredError(
f"Codex OAuth refresh returned 401 with unrecognized error code "
f"({error_code or 'none'}). Run 'codex auth login' to re-authenticate."
)
if response.status_code >= 400:
raise RuntimeError(f"Codex OAuth refresh failed with HTTP {response.status_code}")
try:
body = response.json()
except json.JSONDecodeError as e:
raise RuntimeError(f"Codex OAuth refresh returned non-JSON body: {e}") from e
new_access = body.get("access_token")
if not new_access:
raise RuntimeError("Codex OAuth refresh returned no access_token")
new_refresh = body.get("refresh_token") or self.refresh_token
new_id_token = body.get("id_token")
# Update in-memory state first so waiters see fresh credentials
# immediately, even if disk write fails.
self.access_token = new_access
self.refresh_token = new_refresh
persisted: dict[str, Any] = {
"access_token": new_access,
"refresh_token": new_refresh,
}
if new_id_token:
persisted["id_token"] = new_id_token
try:
self._persist_auth_atomic(persisted)
except OSError as e:
logger.warning(
f"Codex OAuth refresh succeeded but persisting auth.json failed: {type(e).__name__}. "
"In-memory credentials are up to date; on-disk file is stale."
)
logger.info("Codex OAuth access_token refreshed successfully")
def ensure_fresh_token(self) -> None:
"""Proactively refresh the access_token if it is near or past expiry.
Cheap when the token is fresh (just decodes the JWT exp claim and
returns).
"""
if self._token_is_stale():
self.refresh_tokens(reason="proactive (token near expiry)")
def close(self) -> None:
"""Close the underlying HTTP client."""
self._http_client.close()
@@ -15,10 +15,15 @@ so that future server-side changes affect both clients identically.
"""
import asyncio
import base64
import binascii
import json
import logging
import os
import tempfile
import time
import uuid
from datetime import datetime, timezone
from pathlib import Path
from typing import Any
@@ -28,27 +33,37 @@ from hindsight_api.engine.llm_interface import LLMInterface, OutputTooLongError
from hindsight_api.engine.response_models import LLMToolCall, LLMToolCallResult, TokenUsage
from hindsight_api.metrics import get_metrics_collector
from .codex_auth import (
_CODEX_CLIENT_ID,
_CODEX_REFRESH_TOKEN_URL,
_CODEX_TERMINAL_REFRESH_ERROR_CODES,
_CODEX_TOKEN_REFRESH_SKEW_SECONDS,
CodexAuthManager,
CodexRefreshExpiredError,
logger = logging.getLogger(__name__)
# OAuth refresh endpoint and client id, mirrored from the canonical
# ``@openai/codex`` CLI (codex-rs/login/src/auth/manager.rs on
# github.com/openai/codex). The endpoint is overridable via env var so that
# future Codex changes or staging environments can be pointed at without a
# code change — same env var name the upstream CLI uses.
_CODEX_REFRESH_TOKEN_URL = os.environ.get("CODEX_REFRESH_TOKEN_URL_OVERRIDE", "https://auth.openai.com/oauth/token")
_CODEX_CLIENT_ID = "app_EMoamEEZ73f0CkXaXp7hrann"
# Proactively refresh this many seconds before the JWT ``exp`` claim. The
# upstream Codex CLI uses no skew (it refreshes at ``exp <= now``); the
# extra window reduces races where a request leaves the client with a token
# that the server has already declared expired by the time it arrives.
_CODEX_TOKEN_REFRESH_SKEW_SECONDS = 60
# OAuth error codes that the refresh endpoint returns when the refresh_token
# itself is no longer usable. These are terminal — retrying refresh will not
# succeed; the user must re-run ``codex auth login``.
_CODEX_TERMINAL_REFRESH_ERROR_CODES = frozenset(
{"refresh_token_expired", "refresh_token_reused", "refresh_token_invalidated"}
)
# Re-export for backward compatibility (tests import from this module).
__all__ = [
"CodexLLM",
"CodexRefreshExpiredError",
"CodexAuthManager",
"_CODEX_REFRESH_TOKEN_URL",
"_CODEX_CLIENT_ID",
"_CODEX_TOKEN_REFRESH_SKEW_SECONDS",
"_CODEX_TERMINAL_REFRESH_ERROR_CODES",
]
logger = logging.getLogger(__name__)
class CodexRefreshExpiredError(RuntimeError):
"""Raised when the Codex refresh_token itself is no longer valid.
The user must re-run ``codex auth login`` to obtain new credentials.
Callers should surface a clear remediation message and stop retrying.
"""
class CodexLLM(LLMInterface):
@@ -71,15 +86,20 @@ class CodexLLM(LLMInterface):
"""Initialize Codex LLM provider."""
super().__init__(provider, api_key, base_url, model, reasoning_effort, **kwargs)
# Single-flight async refresh lock. Multiple concurrent coroutines
# racing toward an expired token should produce one network refresh.
# Path is fixed at ~/.codex/auth.json — matches the upstream CLI.
# Storing it on self lets the refresh path re-read after another
# process (e.g. a sidecar) rotates the file out from under us.
self._auth_file = Path.home() / ".codex" / "auth.json"
# Single-flight refresh lock. Multiple concurrent requests racing
# toward an expired token should produce one network refresh, not N.
self._auth_lock = asyncio.Lock()
# Load Codex OAuth credentials (keep these methods for test patching).
# Load Codex OAuth credentials
try:
access_token, account_id = self._load_codex_auth()
refresh_token = self._load_codex_refresh_token()
logger.info(f"Loaded Codex OAuth credentials for account: {account_id}")
self.access_token, self.account_id = self._load_codex_auth()
self.refresh_token = self._load_codex_refresh_token()
logger.info(f"Loaded Codex OAuth credentials for account: {self.account_id}")
except Exception as e:
raise RuntimeError(
f"Failed to load Codex OAuth credentials from ~/.codex/auth.json: {e}\n\n"
@@ -90,22 +110,9 @@ class CodexLLM(LLMInterface):
"Or use a different provider (openai, anthropic, gemini) with API keys."
) from e
self._auth_manager = CodexAuthManager(
access_token=access_token,
account_id=account_id,
refresh_token=refresh_token,
auth_file=Path.home() / ".codex" / "auth.json",
)
# Use ChatGPT backend API endpoint. Codex auth is tied to
# chatgpt.com/backend-api, not the OpenAI-compatible base URL used by
# other providers. Deployments often set a global LLM_BASE_URL for an
# OpenAI-compatible proxy; ignore that inherited value unless the user
# explicitly provides a Codex backend URL.
if not self.base_url or self.base_url.rstrip("/").endswith("/v1"):
# Use ChatGPT backend API endpoint
if not self.base_url:
self.base_url = "https://chatgpt.com/backend-api"
else:
self.base_url = self.base_url.rstrip("/")
# Normalize model name (strip openai/ prefix if present)
if self.model.startswith("openai/"):
@@ -118,42 +125,6 @@ class CodexLLM(LLMInterface):
# HTTP client for SSE streaming
self._client = httpx.AsyncClient(timeout=120.0)
# ------------------------------------------------------------------
# Properties — delegate to _auth_manager (preserves test-visible API)
# ------------------------------------------------------------------
@property
def access_token(self) -> str:
return self._auth_manager.access_token
@access_token.setter
def access_token(self, v: str) -> None:
self._auth_manager.access_token = v
@property
def account_id(self) -> str:
return self._auth_manager.account_id
@property
def refresh_token(self) -> str | None:
return self._auth_manager.refresh_token
@refresh_token.setter
def refresh_token(self, v: str | None) -> None:
self._auth_manager.refresh_token = v
@property
def _auth_file(self) -> Path:
return self._auth_manager._auth_file
@_auth_file.setter
def _auth_file(self, v: Path) -> None:
self._auth_manager._auth_file = v
# ------------------------------------------------------------------
# Forwarding methods (keep surface area for tests / subclasses)
# ------------------------------------------------------------------
def _load_codex_auth(self) -> tuple[str, str]:
"""
Load OAuth credentials from ~/.codex/auth.json.
@@ -190,57 +161,273 @@ class CodexLLM(LLMInterface):
return access_token, account_id
def _load_codex_refresh_token(self) -> str | None:
"""Read ``tokens.refresh_token`` from the configured auth file.
"""Load ``tokens.refresh_token`` from ``~/.codex/auth.json``.
Kept as an instance method so existing tests that patch
``CodexLLM._load_codex_refresh_token`` continue to work. Works both
pre- and post-``__init__`` because it does not depend on
``_auth_manager`` being constructed yet.
Returns None when the auth file is unreadable or omits the field —
the provider still functions as a one-shot loader in that case, it
just can't refresh when the access_token expires. This deliberately
does not raise so that ``__init__`` keeps the existing failure mode
of raising only on missing ``access_token``.
"""
auth_file = (
self._auth_manager._auth_file if hasattr(self, "_auth_manager") else Path.home() / ".codex" / "auth.json"
)
return CodexAuthManager.load_refresh_token_from_file(auth_file)
try:
with open(self._auth_file) as f:
data = json.load(f)
except (OSError, json.JSONDecodeError) as e:
logger.warning(
f"Codex auth file unreadable when loading refresh_token: {type(e).__name__}. "
"Token refresh will not be available; the access_token in memory will be used until it expires."
)
return None
return data.get("tokens", {}).get("refresh_token")
@staticmethod
def _decode_jwt_exp_unixtime(token: str) -> int | None:
"""Delegate to ``CodexAuthManager._decode_jwt_exp_unixtime``."""
return CodexAuthManager._decode_jwt_exp_unixtime(token)
"""Return the JWT ``exp`` claim as a unix timestamp, or None on parse failure.
ChatGPT/Codex access_tokens are JWTs whose payload includes ``exp``
(RFC 7519). We need the expiry to schedule proactive refresh — the
``auth.json`` file does not persist a separate ``expires_at`` field
in the upstream CLI's shape, so decoding the JWT itself is the
canonical way to know when the token is stale.
We do not verify the signature — the server is the source of truth
on whether the token is actually accepted, and the only thing this
method affects is the *timing* of refresh, not whether to trust the
token contents.
"""
try:
parts = token.split(".")
if len(parts) < 2:
return None
payload_b64 = parts[1]
# JWT uses base64url without padding. Re-pad before decoding.
padding = "=" * (-len(payload_b64) % 4)
payload_bytes = base64.urlsafe_b64decode(payload_b64 + padding)
payload = json.loads(payload_bytes.decode("utf-8"))
exp = payload.get("exp")
return int(exp) if exp is not None else None
except (ValueError, TypeError, json.JSONDecodeError, binascii.Error):
return None
def _token_is_stale(self, skew_seconds: int = _CODEX_TOKEN_REFRESH_SKEW_SECONDS) -> bool:
"""Delegate to ``_auth_manager._token_is_stale``."""
return self._auth_manager._token_is_stale(skew_seconds)
"""True when the cached access_token is past expiry (with skew).
Returns False when expiry cannot be determined — we'd rather use a
possibly-expired token and recover via the reactive 401 path than
refresh aggressively on every request when ``exp`` parsing fails.
"""
exp = self._decode_jwt_exp_unixtime(self.access_token)
if exp is None:
return False
return exp <= int(time.time()) + skew_seconds
def _persist_auth_atomic(self, updated_tokens: dict[str, Any]) -> None:
"""Delegate to ``_auth_manager._persist_auth_atomic``."""
return self._auth_manager._persist_auth_atomic(updated_tokens)
"""Write the rotated tokens back to ``~/.codex/auth.json`` atomically.
Strategy: re-read the on-disk auth.json (so we don't clobber fields
another process may have added), patch ``tokens.*`` and
``last_refresh``, write to a tempfile in the same directory with
mode 0600, then ``os.replace`` onto the target. ``os.replace`` is
atomic within the same filesystem on POSIX and Windows, so a
concurrent reader will see either the old file or the fully-written
new file — never a partial truncate, which is the upstream CLI's
worst-case race.
On non-Unix platforms the chmod is a best-effort no-op; the parent
directory permissions still bound access.
"""
current: dict[str, Any]
try:
with open(self._auth_file) as f:
loaded = json.load(f)
# auth.json should always be a JSON object at the top level; if
# someone has hand-edited it into a non-object shape, fall back
# to the minimal default rather than crashing the refresh path.
current = loaded if isinstance(loaded, dict) else {"auth_mode": "chatgpt", "tokens": {}}
except (OSError, json.JSONDecodeError):
# If the file became unreadable between our last read and now,
# construct a minimal shape rather than refusing to persist.
current = {"auth_mode": "chatgpt", "tokens": {}}
existing_tokens = current.get("tokens")
tokens: dict[str, Any] = existing_tokens if isinstance(existing_tokens, dict) else {}
for key in ("access_token", "refresh_token", "id_token", "account_id"):
if key in updated_tokens and updated_tokens[key] is not None:
tokens[key] = updated_tokens[key]
current["tokens"] = tokens
current["last_refresh"] = datetime.now(timezone.utc).isoformat().replace("+00:00", "Z")
# Write to a sibling tempfile so the rename is same-filesystem.
parent = self._auth_file.parent
parent.mkdir(parents=True, exist_ok=True)
fd, tmp_path = tempfile.mkstemp(prefix=".auth.", suffix=".json.tmp", dir=str(parent))
try:
with os.fdopen(fd, "w") as f:
json.dump(current, f, indent=2)
f.flush()
os.fsync(f.fileno())
try:
os.chmod(tmp_path, 0o600)
except OSError:
pass # best-effort on platforms that don't support chmod
os.replace(tmp_path, self._auth_file)
except Exception:
# Clean up the orphaned tempfile if rename fails.
try:
os.unlink(tmp_path)
except OSError:
pass
raise
async def _refresh_oauth_tokens(self, reason: str = "", *, force: bool = False) -> None:
"""Async single-flight OAuth token refresh.
"""Refresh the OAuth access_token using the stored refresh_token.
Outer asyncio.Lock preserves single-flight semantics for concurrent
coroutines; the actual network call is offloaded to a thread via
``asyncio.to_thread`` so the event loop stays unblocked.
Single-flight: serialized through ``self._auth_lock`` so concurrent
callers produce one network request. The first caller refreshes; the
rest wake up and observe that either (a) the in-memory token is no
longer stale (proactive case) or (b) the in-memory token has changed
since they entered (reactive case), and return without re-refreshing.
Args:
reason: Free-form string included in log lines for diagnostics.
force: When True, refresh even if the JWT exp claim looks fresh.
Used by the reactive 401 path.
Used by the reactive 401 path — the server rejected the
token, so we cannot trust the JWT's self-reported expiry.
Raises:
CodexRefreshExpiredError: when the server returns a terminal
error code or any 401 on the refresh endpoint.
error code (refresh_token_expired/reused/invalidated) or any
401 on the refresh endpoint itself.
RuntimeError: for other refresh failures (network, 5xx, etc.).
"""
# Capture the token we'd be refreshing BEFORE acquiring the lock so
# that we can detect mid-wait rotation by another coroutine.
token_before_lock = self.access_token
async with self._auth_lock:
if force:
# Reactive: skip only if another coroutine already rotated
# the token while we were waiting on the lock.
if self.access_token != token_before_lock:
return
else:
if not self._auth_manager._token_is_stale():
# Proactive: skip if the token is no longer stale (the
# canonical "another coroutine refreshed first" check).
if not self._token_is_stale():
return
await asyncio.to_thread(lambda: self._auth_manager.refresh_tokens(reason, force=force))
if not self.refresh_token:
raise RuntimeError(
"Codex access_token is expired but no refresh_token is available. "
"Run 'codex auth login' to re-authenticate."
)
log_reason = f" ({reason})" if reason else ""
logger.info(f"Refreshing Codex OAuth access_token{log_reason}")
request_body = {
"client_id": _CODEX_CLIENT_ID,
"grant_type": "refresh_token",
"refresh_token": self.refresh_token,
}
try:
response = await self._client.post(
_CODEX_REFRESH_TOKEN_URL,
json=request_body,
headers={"Content-Type": "application/json"},
timeout=30.0,
)
except httpx.RequestError as e:
raise RuntimeError(f"Codex OAuth refresh network error: {type(e).__name__}") from e
if response.status_code == 401:
# Classify by ``error.code`` (or top-level ``error`` string) — same
# mapping as the upstream Rust CLI's request_chatgpt_token_refresh.
error_code = self._extract_oauth_error_code(response)
if error_code in _CODEX_TERMINAL_REFRESH_ERROR_CODES:
raise CodexRefreshExpiredError(
f"Codex refresh_token is permanently invalid (error.code={error_code}). "
"Run 'codex auth login' to re-authenticate."
)
# Unknown 401 — treat as terminal too, matching the upstream classification.
raise CodexRefreshExpiredError(
f"Codex OAuth refresh returned 401 with unrecognized error code "
f"({error_code or 'none'}). Run 'codex auth login' to re-authenticate."
)
if response.status_code >= 400:
# 5xx and other 4xx are transient/retryable from the caller's
# perspective; surface as RuntimeError without leaking the
# request body in logs.
raise RuntimeError(f"Codex OAuth refresh failed with HTTP {response.status_code}")
try:
body = response.json()
except json.JSONDecodeError as e:
raise RuntimeError(f"Codex OAuth refresh returned non-JSON body: {e}") from e
new_access = body.get("access_token")
if not new_access:
raise RuntimeError("Codex OAuth refresh returned no access_token")
# The refresh_token may rotate on each refresh — adopt the new
# one if the server sent it, otherwise keep the existing.
new_refresh = body.get("refresh_token") or self.refresh_token
new_id_token = body.get("id_token")
# Update in-memory state first so callers waiting on the lock
# see fresh credentials immediately, even if disk write fails.
self.access_token = new_access
self.refresh_token = new_refresh
persisted = {
"access_token": new_access,
"refresh_token": new_refresh,
}
if new_id_token:
persisted["id_token"] = new_id_token
try:
self._persist_auth_atomic(persisted)
except OSError as e:
# In-memory creds are valid; warn but don't fail the request
# path. Future process starts will fall back to the stale
# on-disk auth.json and immediately refresh.
logger.warning(
f"Codex OAuth refresh succeeded but persisting auth.json failed: {type(e).__name__}. "
"In-memory credentials are up to date; on-disk file is stale."
)
logger.info("Codex OAuth access_token refreshed successfully")
@staticmethod
def _extract_oauth_error_code(response: "httpx.Response") -> str | None:
"""Pull the OAuth error code out of a 4xx response body, if present.
The refresh endpoint returns shapes like
``{"error": "...", "error_code": "..."}`` or
``{"error": {"code": "..."}}``. We don't fail the call if the body
is unparseable — the caller falls back to a generic "unknown" error.
"""
try:
body = response.json()
except (json.JSONDecodeError, ValueError):
return None
if not isinstance(body, dict):
return None
# Shape 1: error is a nested object with "code"
err = body.get("error")
if isinstance(err, dict):
code = err.get("code")
if isinstance(code, str):
return code
# Shape 2: top-level error_code string
code = body.get("error_code")
if isinstance(code, str):
return code
# Shape 3: error is itself a string code
if isinstance(err, str):
return err
return None
async def _ensure_fresh_token(self) -> None:
"""Refresh the access_token proactively if it is near or past expiry.
@@ -248,10 +435,13 @@ class CodexLLM(LLMInterface):
Called at the top of every API-bound method. Cheap when the token is
fresh (just decodes the JWT exp claim and returns).
"""
if self._auth_manager._token_is_stale():
if self._token_is_stale():
try:
await self._refresh_oauth_tokens(reason="proactive (token near expiry)")
except CodexRefreshExpiredError:
# Surface to the caller as the same RuntimeError shape the
# request loop has historically raised, so existing error
# handling paths keep working.
raise
def _map_reasoning_effort(self, effort: str) -> str:
@@ -883,6 +1073,5 @@ class CodexLLM(LLMInterface):
return content if content else None, tool_calls
async def cleanup(self) -> None:
"""Clean up HTTP clients."""
"""Clean up HTTP client."""
await self._client.aclose()
self._auth_manager.close()
@@ -1,396 +0,0 @@
"""Fireworks AI provider with batch-inference support.
Fireworks' *online* inference endpoint (``/inference/v1``) is OpenAI-compatible,
so ``FireworksLLM`` subclasses :class:`OpenAICompatibleLLM` and reuses its entire
chat path. Only the *batch* mechanism differs: Fireworks does NOT implement the
OpenAI ``/v1/batches`` API. Instead it exposes a proprietary, account-scoped
dataset -> job -> download REST workflow on a separate control-plane host. This
class overrides only the four batch members of the interface, translating that
workflow to/from the OpenAI-batch shapes the retain orchestrator and
``fact_extraction`` consumer expect — so nothing downstream changes.
Interface contract preserved (see ``fact_extraction.py`` result handling)::
result["response"]["body"]["choices"][0]["message"]["content"]
Workflow (control-plane host, e.g. ``https://api.fireworks.ai``)::
POST /v1/accounts/{acct}/datasets create input dataset
POST /v1/accounts/{acct}/datasets/{id}:upload upload input JSONL
POST /v1/accounts/{acct}/batchInferenceJobs create job
GET /v1/accounts/{acct}/batchInferenceJobs/{jobId} poll status
GET /v1/accounts/{acct}/datasets/{out}:getDownloadEndpoint signed URLs
GET <signed-url> download output JSONL
NOTE: the exact *output JSONL line* nesting is not verbatim-documented by
Fireworks. ``_normalize_output_line`` handles both the observed shape
(``{custom_id, response: {...completion...}, error}``) and a ``response.body``
nesting defensively. Confirm against a live key via the integration path.
"""
import json
import logging
import uuid
from datetime import datetime, timezone
from typing import Any
import httpx
from .openai_compatible_llm import OpenAICompatibleLLM
logger = logging.getLogger(__name__)
# Normalized statuses the retain driver treats as fatal (it raises) vs. keeps
# polling on. "completed" ends the poll; anything else not in this set means
# "keep polling".
_TERMINAL_STATUSES = frozenset({"completed", "failed", "cancelled", "expired"})
# Default per-request timeout for control-plane HTTP calls (not the job wait).
_HTTP_TIMEOUT_SECONDS = 60.0
# Fallback max job wait if neither a constructor arg nor config supplies one
# (24h matches Fireworks' maximum job timeout).
_DEFAULT_MAX_WAIT_SECONDS = 86_400
class FireworksLLM(OpenAICompatibleLLM):
"""Fireworks provider: OpenAI-compatible online inference + native batch."""
def __init__(
self,
provider: str = "fireworks",
*,
api_key: str,
base_url: str = "",
model: str,
reasoning_effort: str = "low",
account_id: str | None = None,
batch_base_url: str | None = None,
max_wait_seconds: int | None = None,
http_client: httpx.AsyncClient | None = None,
**kwargs: Any,
):
super().__init__(
provider=provider,
api_key=api_key,
base_url=base_url,
model=model,
reasoning_effort=reasoning_effort,
**kwargs,
)
# Batch settings are static, server-level config. Resolve any unset
# values from the global config lazily so the online inference path
# works even when batch is never configured.
if account_id is None or batch_base_url is None or max_wait_seconds is None:
from ...config import get_config
cfg = get_config()
if account_id is None:
account_id = cfg.fireworks_account_id
if batch_base_url is None:
batch_base_url = cfg.fireworks_batch_base_url
if max_wait_seconds is None:
max_wait_seconds = cfg.fireworks_batch_max_wait_seconds
self._account_id = account_id
self._batch_base_url = (batch_base_url or "https://api.fireworks.ai").rstrip("/")
self._max_wait_seconds: int = (
int(max_wait_seconds) if max_wait_seconds is not None else _DEFAULT_MAX_WAIT_SECONDS
)
self._http_client = http_client
self._owns_http_client = http_client is None
# ----- interface: batch members -------------------------------------
async def supports_batch_api(self) -> bool:
return True
async def submit_batch(
self,
requests: list[dict[str, Any]],
endpoint: str = "/v1/chat/completions",
completion_window: str = "24h",
) -> dict[str, Any]:
# endpoint/completion_window are part of the LLMInterface batch contract
# (used by the OpenAI path) but have no analogue in Fireworks' job API:
# the request shape is fixed (chat) and the job timeout is server-side.
# Kept for signature compatibility with the shared retain driver.
self._require_account_id()
logger.info(f"Submitting Fireworks batch with {len(requests)} requests")
jsonl = self._translate_requests(requests)
input_dataset_id = f"hs-batch-in-{uuid.uuid4().hex}"
output_dataset_id = f"hs-batch-out-{uuid.uuid4().hex}"
headers = self._auth_headers()
# The `dataset` resource takes format + exampleCount on create. CHAT is
# the format for chat-completion batch input; exampleCount is the JSONL
# line count (Fireworks rejects uploaded datasets without it) and is an
# int64 proto field, so it goes over the wire as a string.
await self._request(
"POST",
self._datasets_url(),
headers=headers,
json={
"datasetId": input_dataset_id,
"dataset": {"format": "CHAT", "exampleCount": str(len(requests))},
},
)
await self._request(
"POST",
f"{self._datasets_url()}/{input_dataset_id}:upload",
headers=headers,
files={"file": ("batch_input.jsonl", jsonl.encode("utf-8"), "application/jsonl")},
)
job_resp = await self._request(
"POST",
self._jobs_url(),
headers=headers,
json={
"model": self.model,
"inputDatasetId": self._dataset_resource(input_dataset_id),
"outputDatasetId": self._dataset_resource(output_dataset_id),
},
)
job = job_resp.json()
job_id = self._last_segment(job.get("name")) or output_dataset_id
logger.info(f"Fireworks batch job submitted: {job_id}, state={job.get('state')}")
return {
"batch_id": job_id,
"status": self._normalize_state(job.get("state", "")),
"input_dataset_id": input_dataset_id,
"output_dataset_id": output_dataset_id,
"created_at": job.get("createTime"),
"request_count": len(requests),
}
async def get_batch_status(self, batch_id: str) -> dict[str, Any]:
self._require_account_id()
job = (await self._request("GET", self._job_url(batch_id), headers=self._auth_headers())).json()
status = self._normalize_state(job.get("state", ""))
progress = job.get("jobProgress") or {}
result: dict[str, Any] = {
"batch_id": batch_id,
"status": status,
"created_at": job.get("createTime"),
"request_counts": {
"total": _to_int(progress.get("totalInputRequests")),
"completed": _to_int(progress.get("successfullyProcessedRequests")),
"failed": _to_int(progress.get("failedRequests")),
},
}
output_dataset_id = job.get("outputDatasetId")
if output_dataset_id:
result["output_dataset_id"] = output_dataset_id
# Fireworks reports terminal failure detail in the `status` {code,message}.
if job.get("status"):
result["errors"] = job["status"]
# PENDING-forever guard: the shared retain poll loop has no max-wait, so
# if a (likely non-batch-eligible) job never reaches a terminal state we
# surface "expired" once createTime is older than the cap. Derived from
# the server's createTime so it survives crash-recovery polling resumes.
if status not in _TERMINAL_STATUSES:
elapsed = self._elapsed_seconds(job.get("createTime"))
if elapsed is not None and elapsed > self._max_wait_seconds:
result["status"] = "expired"
result["errors"] = (
f"Fireworks batch {batch_id} exceeded max wait of {self._max_wait_seconds}s "
f"in state {job.get('state')!r}. The model may not be batch-eligible "
f"(such jobs stay PENDING indefinitely)."
)
logger.error(result["errors"])
return result
async def retrieve_batch_results(self, batch_id: str) -> list[dict[str, Any]]:
self._require_account_id()
job = (await self._request("GET", self._job_url(batch_id), headers=self._auth_headers())).json()
status = self._normalize_state(job.get("state", ""))
if status != "completed":
raise ValueError(f"Fireworks batch {batch_id} is not completed yet (state: {job.get('state')!r})")
output_dataset_id = job.get("outputDatasetId")
if not output_dataset_id:
raise ValueError(f"Fireworks batch {batch_id} completed but reported no output dataset")
output_short_id = self._last_segment(output_dataset_id)
if not output_short_id:
raise ValueError(
f"Fireworks batch {batch_id} reported an unparseable output dataset: {output_dataset_id!r}"
)
download = (
await self._request("GET", self._download_endpoint_url(output_short_id), headers=self._auth_headers())
).json()
signed_urls = (download or {}).get("filenameToSignedUrls") or {}
if not signed_urls:
raise ValueError(f"Fireworks batch {batch_id} returned no downloadable output files")
# The output dataset contains a results file plus a separate error file.
# Download every file and normalize each line; error-file lines carry an
# `error` so partial failures surface per custom_id instead of vanishing.
results: list[dict[str, Any]] = []
for url in signed_urls.values():
# Signed URLs are pre-authenticated — do not attach the bearer token.
file_resp = await self._request("GET", url)
for line in file_resp.text.strip().split("\n"):
if line.strip():
results.append(self._normalize_output_line(json.loads(line)))
logger.info(f"Retrieved {len(results)} results for Fireworks batch {batch_id}")
return results
async def cleanup(self) -> None:
await super().cleanup()
if self._owns_http_client and self._http_client is not None:
await self._http_client.aclose()
# ----- pure translation/normalization helpers (unit-tested) ----------
@staticmethod
def _translate_requests(requests: list[dict[str, Any]]) -> str:
"""OpenAI batch request -> Fireworks input JSONL.
Fireworks lines are ``{"custom_id", "body"}`` — the OpenAI ``method`` and
``url`` keys are dropped; ``body`` is kept verbatim.
"""
lines = [
json.dumps({"custom_id": req.get("custom_id"), "body": req.get("body")}, ensure_ascii=False)
for req in requests
]
return "\n".join(lines)
@staticmethod
def _normalize_state(fw_state: str) -> str:
"""Fireworks job state -> the retain driver's expected status strings.
Handles both the API enum (``JOB_STATE_*``) and the guide's bare names
(``COMPLETED``/``VALIDATING``/``EXPIRED``). Unknown / in-flight states map
to ``in_progress`` so the driver keeps polling.
"""
state = (fw_state or "").upper()
if state.startswith("JOB_STATE_"):
state = state[len("JOB_STATE_") :]
if state == "COMPLETED":
return "completed"
if state == "FAILED":
return "failed"
if state in ("CANCELLED", "CANCELED"):
return "cancelled"
if state == "EXPIRED":
return "expired"
return "in_progress"
@staticmethod
def _normalize_output_line(line: dict[str, Any]) -> dict[str, Any]:
"""Fireworks output JSONL line -> OpenAI-batch-output shape.
Target: ``{"custom_id", "response": {"body": <chat-completion>}, "error"}``
so the consumer's ``result["response"]["body"]["choices"][0]...`` works.
"""
custom_id = line.get("custom_id")
error = line.get("error")
if error:
return {"custom_id": custom_id, "response": None, "error": error}
response = line.get("response")
if response is None:
response = line.get("body")
# If Fireworks already nests the completion under `body`, unwrap it;
# otherwise the `response` object *is* the completion.
if isinstance(response, dict) and "body" in response:
body = response["body"]
else:
body = response
return {"custom_id": custom_id, "response": {"body": body}, "error": None}
# ----- low-level HTTP + URL helpers ----------------------------------
def _require_account_id(self) -> None:
if not self._account_id:
raise ValueError(
"Fireworks batch inference requires an account id. "
"Set HINDSIGHT_API_FIREWORKS_ACCOUNT_ID to your Fireworks account id."
)
def _auth_headers(self) -> dict[str, str]:
return {"Authorization": f"Bearer {self.api_key}"}
def _http(self) -> httpx.AsyncClient:
if self._http_client is None:
self._http_client = httpx.AsyncClient(timeout=httpx.Timeout(_HTTP_TIMEOUT_SECONDS))
return self._http_client
async def _request(
self,
method: str,
url: str,
*,
headers: dict[str, str] | None = None,
json: dict[str, Any] | None = None,
files: dict[str, Any] | None = None,
) -> httpx.Response:
resp = await self._http().request(method, url, headers=headers, json=json, files=files)
if resp.is_error:
# Surface the API's error body. Fireworks returns JSON describing why a
# 4xx/5xx happened; raise_for_status() alone discards it, which makes
# failures (e.g. a malformed dataset/job request) undebuggable.
raise httpx.HTTPStatusError(
f"Fireworks API {resp.status_code} for {method} {url}: {resp.text[:2000]}",
request=resp.request,
response=resp,
)
return resp
def _accounts_base(self) -> str:
return f"{self._batch_base_url}/v1/accounts/{self._account_id}"
def _datasets_url(self) -> str:
return f"{self._accounts_base()}/datasets"
def _jobs_url(self) -> str:
return f"{self._accounts_base()}/batchInferenceJobs"
def _job_url(self, job_id: str) -> str:
return f"{self._jobs_url()}/{job_id}"
def _download_endpoint_url(self, dataset_short_id: str) -> str:
return f"{self._datasets_url()}/{dataset_short_id}:getDownloadEndpoint"
def _dataset_resource(self, dataset_id: str) -> str:
return f"accounts/{self._account_id}/datasets/{dataset_id}"
@staticmethod
def _last_segment(resource_name: str | None) -> str | None:
if not resource_name:
return None
return resource_name.rstrip("/").split("/")[-1]
@staticmethod
def _elapsed_seconds(create_time: str | None) -> float | None:
if not create_time:
return None
try:
normalized = create_time.replace("Z", "+00:00")
created = datetime.fromisoformat(normalized)
if created.tzinfo is None:
created = created.replace(tzinfo=timezone.utc)
return (datetime.now(timezone.utc) - created).total_seconds()
except (ValueError, TypeError):
return None
def _to_int(value: Any) -> int:
"""Coerce Fireworks' string/int counts to int, defaulting to 0."""
try:
return int(value)
except (ValueError, TypeError):
return 0
@@ -153,24 +153,11 @@ class MockLLM(LLMInterface):
result = self._response_callback(messages, scope)
elif self._mock_response is not None:
result = self._mock_response
elif scope == "retain_extract_facts" and skip_validation:
# Fact extraction: return canned facts derived from user message text.
# This allows tests using a mock LLM to get real facts into the DB
# so retain → recall → reflect pipelines work end-to-end.
result = self._build_mock_facts(messages)
elif scope == "consolidation" and response_format is not None:
# Consolidation: produce a single observation from the input facts
# so the full pipeline (retain → consolidation → observation → recall) works.
result = self._build_mock_consolidation(messages, response_format)
elif scope == "memory_think":
# Reflect: return a plausible text answer
result = "Based on the available information, the answer is related to the context provided."
elif response_format is not None:
# Structured output: try to return a valid empty instance of the model
# so that callers expecting e.g. response_format with defaults
# get a valid instance rather than a crash on {"mock": True}.
# Try to create a minimal valid instance of the response format
try:
result = response_format()
# For Pydantic models, try to create with minimal valid data
result = {"mock": True}
except Exception:
result = {"mock": True}
else:
@@ -256,12 +243,6 @@ class MockLLM(LLMInterface):
else:
result = LLMToolCallResult(content="mock response", finish_reason="stop")
# Set mock token usage on result if not already set
if result.input_tokens == 0:
result.input_tokens = 10
if result.output_tokens == 0:
result.output_tokens = 5
# Record span with mock values
# Convert LLMToolCall objects to dicts for span recording
tool_calls_dict = (
@@ -285,92 +266,6 @@ class MockLLM(LLMInterface):
return result
@staticmethod
def _build_mock_facts(messages: list[dict]) -> dict:
"""Build a canned fact extraction response from the user message text.
Splits the input into sentence-like chunks and returns each as a separate
world fact with a simple entity extracted from the first noun-like word.
This is intentionally simplistic — it just needs to produce structurally
valid facts so the rest of the pipeline (embedding, storage, recall) works.
"""
import re
user_text = ""
for m in messages:
if m.get("role") == "user":
user_text = m.get("content", "")
break
# Split on sentence boundaries: period followed by space/EOL (not mid-number), or newlines
sentences = [s.strip() for s in re.split(r"(?<=\.)\s+|\n+", user_text) if s.strip() and len(s.strip()) > 10]
if not sentences:
sentences = [user_text[:200] if user_text else "mock fact"]
facts = []
for sentence in sentences[:10]: # Cap at 10 facts per chunk
# Extract simple entities: capitalized words that aren't common words
words = re.findall(r"\b[A-Z][a-z]+\b", sentence)
entities = [{"text": w} for w in dict.fromkeys(words)][:5] # Dedupe, cap at 5
facts.append(
{
"what": sentence,
"when": "N/A",
"where": "N/A",
"who": "N/A",
"why": "N/A",
"fact_kind": "conversation",
"fact_type": "world",
"entities": entities,
}
)
return {"facts": facts}
@staticmethod
def _build_mock_consolidation(messages: list[dict], response_format: Any) -> Any:
"""Build a mock consolidation response that creates one observation per fact.
Parses fact IDs from the consolidation prompt and creates one observation
per fact, each referencing its source fact ID. This mimics real LLM behavior
where distinct facts produce separate observations, preserving entity
separation so pipeline tests (graph filtering, entity linking) work correctly.
"""
import re
user_text = ""
for m in messages:
if m.get("role") == "user":
user_text = m.get("content", "")
break
# Extract fact UUIDs from the prompt (format: "[<uuid>] <text>")
fact_entries = re.findall(r"\[([0-9a-f-]{36})\]\s*(.+?)(?:\n|$)", user_text)
if not fact_entries:
# No facts to consolidate — return empty response
try:
return response_format()
except Exception:
return {"creates": [], "updates": [], "deletes": []}
# Create one observation per fact to preserve entity separation
creates = []
for fact_id, fact_text in fact_entries:
creates.append({"text": fact_text.strip(), "source_fact_ids": [fact_id]})
try:
return response_format(
creates=creates,
updates=[],
deletes=[],
)
except Exception:
# Fallback if response_format constructor doesn't accept these args
return {"creates": creates, "updates": [], "deletes": []}
async def cleanup(self) -> None:
"""Clean up resources (no-op for mock provider)."""
pass
@@ -423,8 +318,6 @@ class MockLLM(LLMInterface):
return self._mock_calls
def clear_mock_calls(self) -> None:
"""Clear all recorded calls and any configured response/exception state."""
"""Clear the recorded mock calls and any set exception."""
self._mock_calls = []
self._mock_exception = None
self._mock_response = None
self._response_callback = None
@@ -270,7 +270,6 @@ class OpenAICompatibleLLM(LLMInterface):
"openai",
"groq",
"ollama",
"ollama-cloud",
"lmstudio",
"llamacpp",
"minimax",
@@ -279,7 +278,6 @@ class OpenAICompatibleLLM(LLMInterface):
"openrouter",
"zai",
"opencode-go",
"fireworks",
]
if self.provider not in valid_providers:
raise ValueError(f"OpenAICompatibleLLM only supports: {', '.join(valid_providers)}. Got: {self.provider}")
@@ -290,8 +288,6 @@ class OpenAICompatibleLLM(LLMInterface):
self.base_url = "https://api.groq.com/openai/v1"
elif self.provider == "ollama":
self.base_url = "http://localhost:11434/v1"
elif self.provider == "ollama-cloud":
self.base_url = "https://ollama.com/v1"
elif self.provider == "lmstudio":
self.base_url = "http://localhost:1234/v1"
elif self.provider == "minimax":
@@ -304,10 +300,6 @@ class OpenAICompatibleLLM(LLMInterface):
self.base_url = "https://api.z.ai/api/coding/paas/v4"
elif self.provider == "opencode-go":
self.base_url = "https://opencode.ai/zen/go/v1"
elif self.provider == "fireworks":
# OpenAI-compatible inference host (online path). The batch API
# lives on a separate control-plane host — see FireworksLLM.
self.base_url = "https://api.fireworks.ai/inference/v1"
# For ollama/lmstudio, use dummy key if not provided
if self.provider in ("ollama", "lmstudio") and not self.api_key:
@@ -324,7 +316,6 @@ class OpenAICompatibleLLM(LLMInterface):
"openrouter",
"zai",
"opencode-go",
"ollama-cloud",
)
and not self.api_key
):
@@ -1082,17 +1073,12 @@ class OpenAICompatibleLLM(LLMInterface):
last_exception = None
# Pass API key as Bearer token for cloud Ollama endpoints
headers: dict[str, str] = {}
if self.api_key and self.api_key != "local":
headers["Authorization"] = f"Bearer {self.api_key}"
async with httpx.AsyncClient(timeout=300.0) as client:
for attempt in range(max_retries + 1):
if attempt > 0:
set_stage(f"llm.ollama_native.{scope}.attempt={attempt + 1}/{max_retries + 1}")
try:
response = await client.post(native_url, json=payload, headers=headers)
response = await client.post(native_url, json=payload)
response.raise_for_status()
result = response.json()
@@ -321,7 +321,6 @@ async def run_reflect_agent(
include_recall: bool = True,
budget: str | None = None,
max_context_tokens: int = 100_000,
llm_output_language: str | None = None,
) -> ReflectAgentResult:
"""
Execute the reflect agent loop using native tool calling.
@@ -370,12 +369,7 @@ async def run_reflect_agent(
# Build initial messages (directives are injected into system prompt at START and END)
system_prompt = build_system_prompt_for_tools(
bank_profile,
context,
directives=directives,
has_mental_models=has_mental_models,
include_observations=include_observations,
budget=budget,
bank_profile, context, directives=directives, has_mental_models=has_mental_models, budget=budget
)
messages: list[dict[str, Any]] = [
{"role": "system", "content": system_prompt},
@@ -453,10 +447,7 @@ async def run_reflect_agent(
llm_start = time.time()
response, usage = await llm_config.call(
messages=[
{
"role": "system",
"content": build_final_system_prompt(bank_profile.get("mission"), llm_output_language),
},
{"role": "system", "content": build_final_system_prompt(bank_profile.get("mission"))},
{"role": "user", "content": prompt},
],
scope="reflect",
@@ -513,10 +504,7 @@ async def run_reflect_agent(
llm_start = time.time()
response, usage = await llm_config.call(
messages=[
{
"role": "system",
"content": build_final_system_prompt(bank_profile.get("mission"), llm_output_language),
},
{"role": "system", "content": build_final_system_prompt(bank_profile.get("mission"))},
{"role": "user", "content": prompt},
],
scope="reflect",
@@ -619,10 +607,7 @@ async def run_reflect_agent(
llm_start = time.time()
response, usage = await llm_config.call(
messages=[
{
"role": "system",
"content": build_final_system_prompt(bank_profile.get("mission"), llm_output_language),
},
{"role": "system", "content": build_final_system_prompt(bank_profile.get("mission"))},
{"role": "user", "content": prompt},
],
scope="reflect",
@@ -743,10 +728,7 @@ async def run_reflect_agent(
llm_start = time.time()
response, usage = await llm_config.call(
messages=[
{
"role": "system",
"content": build_final_system_prompt(bank_profile.get("mission"), llm_output_language),
},
{"role": "system", "content": build_final_system_prompt(bank_profile.get("mission"))},
{"role": "user", "content": prompt},
],
scope="reflect",
@@ -98,7 +98,6 @@ def build_system_prompt_for_tools(
context: str | None = None,
directives: list[dict[str, Any]] | None = None,
has_mental_models: bool = False,
include_observations: bool = True,
budget: str | None = None,
) -> str:
"""
@@ -109,17 +108,11 @@ def build_system_prompt_for_tools(
2. search_observations - Consolidated knowledge with freshness
3. recall - Raw facts as ground truth
The retrieval-strategy and workflow sections are built to match the tools
actually exposed to the LLM mentioning a tool the agent has disabled
causes weaker LLMs to either hallucinate the call (rejected by the agent)
or give up with "I cannot find any information…" (see #1724).
Args:
bank_profile: Bank profile with name and mission
context: Optional additional context
directives: Optional list of directive mental models to inject as hard rules
has_mental_models: Whether the bank has any mental models (skip if not)
include_observations: Whether search_observations is in the tool list.
budget: Search depth budget - "low", "mid", or "high". Controls exploration thoroughness.
"""
name = bank_profile.get("name", "Assistant")
@@ -165,137 +158,56 @@ def build_system_prompt_for_tools(
"- If memories mention someone did an activity, you can infer they likely enjoyed it",
"- Synthesize a coherent narrative from related memories",
"- Be a thoughtful interpreter, not just a literal repeater",
"- When the exact answer isn't stated, use what IS stated to give a best-effort answer AND surface any uncertainty — never invent confidence the data doesn't support.",
"",
"## Temporal Reasoning",
"Every memory and observation carries temporal fields in the JSON tool result:",
"- `mentioned_at` — when the user retained the fact (always set).",
"- `occurred_start` / `occurred_end` — when the underlying event happened (optional, set for dated events).",
"",
"When facts about the SAME facet conflict — counts, statuses, ownership, location, presence, etc. — the fact with the LATEST `mentioned_at` is authoritative. Later statements SUPERSEDE earlier ones. Do NOT average, sum, or favor an explicitly-dated fact over a more recent one.",
"",
"Example: three count facts come back from recall:",
" - 'Team has 2 engineers' (mentioned_at=T1)",
" - 'Team now has 1 engineer' (mentioned_at=T2, occurred_start=2026-05-25)",
" - 'Team has 5 engineers' (mentioned_at=T3)",
"with T1 < T2 < T3. The current size is 5, not 1. Then apply later events (e.g. someone leaving after T3) on top of that.",
"",
"For reconstructing a TIMELINE of events, order by `occurred_start` / `occurred_end` (when things happened), not `mentioned_at` (when they were retained).",
"",
"## Conflicts and Ambiguity",
"Not every retrieval converges on a single answer. Distinguish two cases:",
"",
"- RESOLVABLE conflict — the temporal rule above (latest `mentioned_at` wins) cleanly picks a winner. Apply it and move on.",
"- UNRESOLVABLE ambiguity — the data is internally inconsistent in a way the temporal rule does NOT settle. Examples: a recent aggregate (count, total) is incompatible with the individual entities you can enumerate; two equally-recent facts disagree and no later fact resolves them; events are described but their relative order is unclear; the user's own statements contradict each other and nothing later reconciles them.",
"",
"When the data is genuinely ambiguous: SAY SO in your answer. Name the conflicting facts. Explain why they can't be reconciled. Give a range or a best-effort interpretation with explicit uncertainty (e.g. 'between X and Y, depending on [unresolved condition]'; or 'the most recent statement says A, but B was stated earlier and the gap isn't accounted for in any later fact').",
"",
"An honest 'the data is inconsistent about X' beats a confident wrong answer. Do NOT pick a value arbitrarily, average conflicting values, or smooth over gaps in confident prose. Acknowledging ambiguity is a successful answer, not a failure mode.",
"",
"## Showing Your Reasoning",
"For any answer that resolves a conflict between facts, applies events on top of a count or status, or settles an ambiguity — show your work in the answer text so a reader can audit it.",
"",
"Walk through these steps explicitly:",
"1. **List the relevant facts in `mentioned_at` order (oldest → newest)**, each with the value it asserts. Use a short bulleted list.",
"2. **Identify the authoritative fact** under the temporal rule (latest `mentioned_at` for the contested facet). Write its date down.",
"3. **List candidate events to apply on top** — anything that changes the count, status, or state being asked about. Write each event's date down next to it.",
"4. **Sanity-check each candidate event against the authoritative date** — for EVERY event from step 3, write a one-line check in the form `<event> (<event_date>) vs authoritative (<authoritative_date>) → BEFORE/AFTER → KEEP/DROP`. If the event is BEFORE or EQUAL to the authoritative date, DROP it: it is already reflected in the authoritative fact, and applying it again is double-counting. This is the single most common mistake — do not skip this step even if you feel confident.",
"5. **Show the arithmetic or derivation explicitly** using only the KEEP events from step 4 — e.g. 'authoritative count = 5 (at 2025-02-12); kept events: Shadow died (2025-03-12, AFTER); 5 1 = 4'.",
"6. If step 2 or 3 cannot be done cleanly (no clear winner, overlapping timestamps, unclear event order), STOP and surface this as an UNRESOLVABLE ambiguity per the section above — do not fabricate a derivation.",
"",
"For simple factual lookups that don't involve conflict or arithmetic, you can answer directly without this scaffolding.",
"- When the exact answer isn't stated, use what IS stated to give the best answer",
"",
"## HIERARCHICAL RETRIEVAL STRATEGY",
"",
]
)
# Assemble the retrieval-level blocks for whatever tools are exposed.
# MM and Observations bodies are unconditional; recall's fallback wording
# adapts to which upstream tools precede it (telling the LLM to fall back
# to a tool that isn't in its list is the bug at the root of #1724).
levels: list[tuple[str, list[str]]] = []
# Build retrieval levels based on what's available
if has_mental_models:
levels.append(
(
"MENTAL MODELS (search_mental_models)",
[
"- User-curated summaries about specific topics",
"- HIGHEST quality - manually created and maintained",
"- If a relevant mental model exists and is FRESH, it may fully answer the question",
"- Check `is_stale` field - if stale, also verify with lower levels",
],
)
)
if include_observations:
levels.append(
(
"OBSERVATIONS (search_observations)",
[
"- Auto-consolidated knowledge from memories",
"- Check `is_stale` field - if stale, ALSO use recall() to verify",
"- Good for understanding patterns and summaries",
],
)
)
recall_body = ["- Individual memories (world facts and experiences)"]
if has_mental_models and include_observations:
recall_body.extend(
parts.extend(
[
"You have access to THREE levels of knowledge. Use them in this order:",
"",
"### 1. MENTAL MODELS (search_mental_models) - Try First",
"- User-curated summaries about specific topics",
"- HIGHEST quality - manually created and maintained",
"- If a relevant mental model exists and is FRESH, it may fully answer the question",
"- Check `is_stale` field - if stale, also verify with lower levels",
"",
"### 2. OBSERVATIONS (search_observations) - Second Priority",
"- Auto-consolidated knowledge from memories",
"- Check `is_stale` field - if stale, ALSO use recall() to verify",
"- Good for understanding patterns and summaries",
"",
"### 3. RAW FACTS (recall) - Ground Truth",
"- Individual memories (world facts and experiences)",
"- Use when: no mental models/observations exist, they're stale, or you need specific details",
"- MANDATORY: If search_mental_models and search_observations both return 0 results, you MUST call recall() before giving up",
"- This is the source of truth that other levels are built from",
"",
"**Tool result ordering:** `recall()` and `search_observations()` return their `memories` / `observations` arrays sorted by SEMANTIC RELEVANCE to the query, NOT by time. The POSITION of an entry tells you nothing about when it was retained. For any temporal reasoning — recency, supersession, applying events on top of a state — IGNORE the position and read the per-entry `mentioned_at` field (and `occurred_start` / `occurred_end` for events).",
]
)
else:
parts.extend(
[
"You have access to TWO levels of knowledge. Use them in this order:",
"",
]
)
elif has_mental_models:
recall_body.extend(
[
"- Use when: no mental model exists, it's stale, or you need specific details",
"- MANDATORY: If search_mental_models returns 0 results, you MUST call recall() before giving up",
"- This is the source of truth that mental models are built from",
]
)
elif include_observations:
recall_body.extend(
[
"### 1. OBSERVATIONS (search_observations) - Try First",
"- Auto-consolidated knowledge from memories",
"- Check `is_stale` field - if stale, ALSO use recall() to verify",
"- Good for understanding patterns and summaries",
"",
"### 2. RAW FACTS (recall) - Ground Truth",
"- Individual memories (world facts and experiences)",
"- Use when: no observations exist, they're stale, or you need specific details",
"- MANDATORY: If search_observations returns 0 results or count=0, you MUST call recall() before giving up",
"- This is the source of truth that observations are built from",
"",
"**Tool result ordering:** `recall()` and `search_observations()` return their `memories` / `observations` arrays sorted by SEMANTIC RELEVANCE to the query, NOT by time. The POSITION of an entry tells you nothing about when it was retained. For any temporal reasoning — recency, supersession, applying events on top of a state — IGNORE the position and read the per-entry `mentioned_at` field (and `occurred_start` / `occurred_end` for events).",
"",
]
)
else:
recall_body.extend(
[
"- MANDATORY: Call recall() to gather facts before giving up",
"- This is the source of truth.",
]
)
levels.append(("RAW FACTS (recall) - Ground Truth", recall_body))
# Position-dependent suffix for upstream tools; recall already carries its
# fixed "- Ground Truth" suffix in the header text.
suffixes = [""] * len(levels)
if len(levels) >= 2:
suffixes[0] = " - Try First"
if len(levels) == 3:
suffixes[1] = " - Second Priority"
if len(levels) == 1:
parts.append("You have access to ONE level of knowledge:")
else:
word = "TWO" if len(levels) == 2 else "THREE"
parts.append(f"You have access to {word} levels of knowledge. Use them in this order:")
parts.append("")
for idx, ((header, body), suffix) in enumerate(zip(levels, suffixes), 1):
parts.append(f"### {idx}. {header}{suffix}")
parts.extend(body)
parts.append("")
parts.extend(
[
@@ -355,28 +267,25 @@ def build_system_prompt_for_tools(
parts.append("## Workflow")
steps: list[str] = []
if has_mental_models:
steps.append("First, try search_mental_models() - check if a curated summary exists")
if include_observations:
if has_mental_models:
steps.append("If no mental model or it's stale, try search_observations() for consolidated knowledge")
else:
steps.append("First, try search_observations() - check for consolidated knowledge")
# Recall step phrasing varies with whichever upstream tool(s) precede it.
if include_observations:
steps.append(
"If observations are stale OR you need specific details, use recall() for raw facts"
if has_mental_models
else "If search_observations returns 0 results OR observations are stale, you MUST call recall() for raw facts"
parts.extend(
[
"1. First, try search_mental_models() - check if a curated summary exists",
"2. If no mental model or it's stale, try search_observations() for consolidated knowledge",
"3. If observations are stale OR you need specific details, use recall() for raw facts",
"4. Use expand() if you need more context on specific memories",
"5. When ready, call done() with your answer and supporting IDs",
]
)
elif has_mental_models:
steps.append("If no mental model or it's stale, use recall() for raw facts")
else:
steps.append("Call recall() to gather raw facts")
steps.append("Use expand() if you need more context on specific memories")
steps.append("When ready, call done() with your answer and supporting IDs")
parts.extend(f"{idx}. {step}" for idx, step in enumerate(steps, 1))
parts.extend(
[
"1. First, try search_observations() - check for consolidated knowledge",
"2. If search_observations returns 0 results OR observations are stale, you MUST call recall() for raw facts",
"3. Use expand() if you need more context on specific memories",
"4. When ready, call done() with your answer and supporting IDs",
]
)
parts.extend(
[
@@ -604,16 +513,10 @@ Just provide the direct answer with proper markdown formatting.
CRITICAL: This is a NON-CONVERSATIONAL system. NEVER ask follow-up questions, offer to search again, suggest alternatives, or end with anything like "Would you like me to..." or "Let me know if...". The user cannot reply. Your answer must be complete and self-contained."""
def build_final_system_prompt(mission: str | None = None, llm_output_language: str | None = None) -> str:
"""Build the final synthesis system prompt, using mission as role when set.
When ``llm_output_language`` is set, the response is forced into that
language regardless of the query/source language.
"""
from hindsight_api.engine.prompt_utils import escape_for_prompt, output_language_directive
role_section = escape_for_prompt(mission.strip()) if mission else _DEFAULT_FINAL_ROLE
return _FINAL_SYSTEM_PROMPT_BASE.format(role_section=role_section) + output_language_directive(llm_output_language)
def build_final_system_prompt(mission: str | None = None) -> str:
"""Build the final synthesis system prompt, using mission as role when set."""
role_section = mission.strip() if mission else _DEFAULT_FINAL_ROLE
return _FINAL_SYSTEM_PROMPT_BASE.format(role_section=role_section)
# Backward-compatible constant for non-identity missions
@@ -1,8 +1,17 @@
"""Token counting helpers for reflect prompts and agent control flow."""
from ..token_encoding import count_tokens as _count_tokens
from functools import lru_cache
import tiktoken
@lru_cache(maxsize=1)
def _get_cl100k_base_encoding() -> tiktoken.Encoding:
# tiktoken downloads this encoding on first lookup when it is not cached.
# Keep the lookup lazy so importing hindsight_api does not depend on network access.
return tiktoken.get_encoding("cl100k_base")
def count_cl100k_tokens(text: str) -> int:
"""Return the number of cl100k_base tokens in text."""
return _count_tokens(text)
return len(_get_cl100k_base_encoding().encode(text))
@@ -23,24 +23,6 @@ if TYPE_CHECKING:
logger = logging.getLogger(__name__)
def _prune_nulls(d: dict[str, Any]) -> dict[str, Any]:
"""Drop keys whose value is None or an empty collection (``""``, ``[]``, ``{}``).
Reflect tools dump ``MemoryFact`` / ``ObservationResult`` via ``model_dump()``,
which emits every field including the many that are typically null or empty
(``context``, ``occurred_start``, ``metadata``, ``tags``, etc.). Stripping
these before serializing to JSON for the LLM cuts token cost and removes
fields that aren't telling the model anything.
Callers that need the *presence* of a specific field as a signal (e.g.
``source_fact_ids`` for drill-down) must ensure the value is non-empty
pass the upstream flag that populates it (e.g. ``source_facts_max_tokens``
> 0 on ``tool_search_observations``) rather than relying on Pydantic
emitting ``None``.
"""
return {k: v for k, v in d.items() if v is not None and v != "" and v != [] and v != {}}
def _document_metadata_from_retain_params(retain_params: Any) -> dict[str, Any] | None:
"""Return document metadata stored under retain_params.metadata."""
if isinstance(retain_params, str):
@@ -232,8 +214,8 @@ async def tool_search_observations(
return {
"query": query,
"count": len(result.results),
"observations": [_prune_nulls(m.model_dump()) for m in result.results],
"source_facts": {k: _prune_nulls(v.model_dump()) for k, v in (result.source_facts or {}).items()},
"observations": [m.model_dump() for m in result.results],
"source_facts": {k: v.model_dump() for k, v in (result.source_facts or {}).items()},
"is_stale": is_stale,
"freshness": freshness,
}
@@ -300,8 +282,8 @@ async def tool_recall(
return {
"query": query,
"memories": [_prune_nulls(m.model_dump()) for m in result.results],
"chunks": {k: _prune_nulls(v.model_dump()) for k, v in (result.chunks or {}).items()},
"memories": [m.model_dump() for m in result.results],
"chunks": {k: v.model_dump() for k, v in (result.chunks or {}).items()},
}
@@ -4,54 +4,29 @@ Embedding generation utilities for memory units.
import asyncio
import logging
from typing import Literal, Protocol
logger = logging.getLogger(__name__)
EmbeddingInputType = Literal["document", "query"]
class EmbeddingsBackend(Protocol):
"""Minimal duck-typed surface used by retain/recall — the concrete `Embeddings`
ABC supplies default implementations that delegate to `encode()`."""
def encode_query(self, texts: list[str]) -> list[list[float]]: ...
def encode_documents(self, texts: list[str]) -> list[list[float]]: ...
def generate_embedding(
embeddings_backend: EmbeddingsBackend, text: str, input_type: EmbeddingInputType = "document"
) -> list[float]:
def generate_embedding(embeddings_backend, text: str) -> list[float]:
"""
Generate embedding for text using the provided embeddings backend.
Args:
embeddings_backend: Embeddings instance to use for encoding
text: Text to embed
input_type: Whether text is retained document text or recall/search query text.
Returns:
Embedding vector (dimension depends on embeddings backend)
"""
try:
embeddings = _encode_with_input_type(embeddings_backend, [text], input_type)
embeddings = embeddings_backend.encode([text])
return embeddings[0]
except Exception as e:
raise Exception(f"Failed to generate embedding: {str(e)}")
def _encode_with_input_type(
embeddings_backend: EmbeddingsBackend, texts: list[str], input_type: EmbeddingInputType
) -> list[list[float]]:
if input_type == "query":
return embeddings_backend.encode_query(texts)
return embeddings_backend.encode_documents(texts)
async def generate_embeddings_batch(
embeddings_backend: EmbeddingsBackend, texts: list[str], input_type: EmbeddingInputType = "document"
) -> list[list[float]]:
async def generate_embeddings_batch(embeddings_backend, texts: list[str]) -> list[list[float]]:
"""
Generate embeddings for multiple texts using the provided embeddings backend.
@@ -61,14 +36,17 @@ async def generate_embeddings_batch(
Args:
embeddings_backend: Embeddings instance to use for encoding
texts: List of texts to embed
input_type: Whether texts are retained documents or recall/search queries.
Returns:
List of embeddings in same order as input texts
"""
try:
loop = asyncio.get_event_loop()
embeddings = await loop.run_in_executor(None, _encode_with_input_type, embeddings_backend, texts, input_type)
embeddings = await loop.run_in_executor(
None,
embeddings_backend.encode,
texts,
)
except Exception as e:
raise Exception(f"Failed to generate batch embeddings: {str(e)}")
@@ -1,13 +1,13 @@
"""
Entity processing for retain pipeline.
Handles entity extraction and resolution for stored facts.
Handles entity extraction, resolution, and link creation for stored facts.
"""
import logging
from . import link_utils
from .types import ProcessedFact
from .types import EntityLink, ProcessedFact
logger = logging.getLogger(__name__)
@@ -76,7 +76,8 @@ async def resolve_entities(
entity_labels: Optional entity label taxonomy
Returns:
Tuple of (resolved_entity_ids, entity_to_unit, unit_to_entity_ids).
Tuple of (resolved_entity_ids, entity_to_unit, unit_to_entity_ids)
to pass to build_entity_links().
"""
if not unit_ids or not facts:
return [], [], {}
@@ -98,3 +99,68 @@ async def resolve_entities(
log_buffer,
entity_labels=entity_labels,
)
async def build_entity_links(
entity_resolver,
conn,
bank_id: str,
unit_ids: list[str],
resolved_entity_ids: list[str],
entity_to_unit: list[tuple],
unit_to_entity_ids: dict[str, list[str]],
log_buffer: list[str] = None,
skip_unit_entities_insert: bool = False,
ops=None,
) -> list[EntityLink]:
"""
Build entity links for UI graph visualization.
Queries unit_entities to find shared entities between new and existing units,
then generates EntityLink objects. When called from Phase 3 (post-transaction),
set skip_unit_entities_insert=True since unit_entities were already inserted
in Phase 2.
Args:
entity_resolver: EntityResolver instance
conn: Database connection
bank_id: Bank identifier
unit_ids: Actual unit IDs (must already be inserted in the DB)
resolved_entity_ids: From resolve_entities()
entity_to_unit: From resolve_entities()
unit_to_entity_ids: From resolve_entities()
log_buffer: Optional buffer for detailed logging
skip_unit_entities_insert: Skip unit_entities INSERT (already done in Phase 2)
ops: DataAccessOps instance (from backend.ops)
Returns:
List of EntityLink objects for batch insertion
"""
return await link_utils.build_entity_links_from_resolved(
entity_resolver,
conn,
bank_id,
unit_ids,
resolved_entity_ids,
entity_to_unit,
unit_to_entity_ids,
log_buffer,
skip_unit_entities_insert=skip_unit_entities_insert,
ops=ops,
)
async def insert_entity_links_batch(conn, entity_links: list[EntityLink], bank_id: str, ops=None) -> None:
"""
Insert entity links in batch.
Args:
conn: Database connection
entity_links: List of EntityLink objects
bank_id: Bank identifier (stored directly on memory_links for fast filtering)
ops: DataAccessOps instance (from backend.ops)
"""
if not entity_links:
return
await link_utils.insert_entity_links_batch(conn, entity_links, bank_id, ops=ops)
@@ -888,16 +888,13 @@ def _build_extraction_prompt_and_schema(config) -> tuple[str, type]:
extract_causal_links = config.retain_extract_causal_links
# Build retain_mission section if set - injected before the mode-specific guidelines
# Escape braces so user-supplied text survives str.format() on the prompt template.
from hindsight_api.engine.prompt_utils import escape_for_prompt
retain_mission = getattr(config, "retain_mission", None)
if retain_mission:
retain_mission_section = (
f"══════════════════════════════════════════════════════════════════════════\n"
f"FOCUS — What to retain for this bank\n"
f"══════════════════════════════════════════════════════════════════════════\n\n"
f"{escape_for_prompt(retain_mission)}\n\n"
f"{retain_mission}\n\n"
)
else:
retain_mission_section = ""
@@ -913,7 +910,7 @@ def _build_extraction_prompt_and_schema(config) -> tuple[str, type]:
base_prompt = CUSTOM_FACT_EXTRACTION_PROMPT
prompt = base_prompt.format(
retain_mission_section=retain_mission_section,
custom_instructions=escape_for_prompt(config.retain_custom_instructions),
custom_instructions=config.retain_custom_instructions,
)
elif extraction_mode == "verbose":
prompt = VERBOSE_FACT_EXTRACTION_PROMPT.format(
@@ -950,16 +947,6 @@ def _build_extraction_prompt_and_schema(config) -> tuple[str, type]:
if labels_section:
prompt = prompt + labels_section
# Force the LLM to emit fact text in the configured language, regardless of
# the source content's language. Same directive is applied to consolidation
# and reflect so HINDSIGHT_API_LLM_OUTPUT_LANGUAGE has a uniform effect
# across the pipeline. This is independent of the BM25 indexing language
# (HINDSIGHT_API_TEXT_SEARCH_EXTENSION_NATIVE_LANGUAGE) by design — search
# tokenization and LLM output language are separate concerns.
from ..prompt_utils import output_language_directive
prompt = prompt + output_language_directive(getattr(config, "llm_output_language", None))
response_schema = base_response_class
if labels_cfg and labels_cfg.attributes:
@@ -1142,14 +1129,11 @@ async def _extract_facts_from_chunk(
)
continue
else:
# A non-dict response is malformed (the schema is {"facts": [...]}).
# Raise instead of returning [] so the failure propagates to the
# worker's retry machinery and ultimately fails loudly — never
# silently commit the document with 0 facts. See issue #1833.
raise RuntimeError(
f"Fact extraction failed: LLM returned non-dict JSON after {llm_max_retries} attempts "
f"({type(extraction_response_json).__name__}). Raw: {str(extraction_response_json)[:500]}"
logger.warning(
f"LLM returned non-dict JSON after {llm_max_retries} attempts: {type(extraction_response_json).__name__}. "
f"Raw: {str(extraction_response_json)[:500]}"
)
return [], usage
raw_facts = extraction_response_json.get("facts", [])
@@ -1678,11 +1662,8 @@ async def extract_facts_from_contents_batch_api(
# Check if provider supports batch API
if not await llm_config._provider_impl.supports_batch_api():
raise RuntimeError(
f"retain_batch_enabled=True but provider '{llm_config.provider}' does not "
f"support the batch API. This should have been caught at startup — check "
f"HINDSIGHT_API_RETAIN_BATCH_ENABLED and your LLM provider configuration."
)
logger.warning(f"Batch API not supported for provider {llm_config.provider}, falling back to sync mode")
return await extract_facts_from_contents(contents, llm_config, agent_name, config, pool, operation_id, schema)
# Check if we're resuming an existing batch (crash recovery)
batch_id = None
@@ -2214,9 +2195,7 @@ async def extract_facts_from_contents(
fact_extraction_tasks.append(task)
# Step 2: Wait for all fact extractions to complete.
# return_exceptions=True so a failing item doesn't cancel its still-running
# siblings (which would leave orphaned LLM calls / partial work); we await
# them all, then propagate.
# Use return_exceptions=True so one content item failure doesn't discard the rest.
all_fact_results = await asyncio.gather(*fact_extraction_tasks, return_exceptions=True)
# Step 3: Flatten and convert to typed objects
@@ -2227,18 +2206,14 @@ async def extract_facts_from_contents(
global_chunk_idx = 0
global_fact_idx = 0
# Never silently drop a document's memory. Any extraction failure (provider
# rate-limit / timeout / 5xx, malformed response, token-limit, etc.)
# propagates so the streaming producer surfaces it and the worker's
# RetryTaskAt machinery retries the task — and ultimately fails it *loudly*
# if the problem persists. Swallowing the error and substituting an empty
# result here used to commit the document with 0 facts and mark the
# operation `completed`, losing the memory with no signal. See issue #1833.
# Filter out failed content items
valid_results = []
for content, result in zip(contents, all_fact_results):
if isinstance(result, Exception):
raise result
valid_results.append((content, result))
logger.warning(f"Content extraction failed (skipping): {type(result).__name__}: {result}")
valid_results.append((content, ([], [], TokenUsage())))
else:
valid_results.append((content, result))
for content_index, (content, (facts_from_llm, chunks_from_llm, content_usage)) in enumerate(valid_results):
total_usage = total_usage + content_usage
@@ -321,15 +321,6 @@ async def handle_document_tracking(
f"[RETAIN] Document {document_id} re-ingested: invalidated "
f"{invalidated} observation(s) derived from {len(existing_unit_ids)} outgoing memory_units"
)
# Capture link-recompute victims BEFORE the cascade. Same staleness
# applies on upsert as on explicit delete: surviving units in OTHER
# documents that linked to these doomed units are about to lose
# those links. ``ops`` may be None for older callers that haven't
# been wired up — skip enqueue in that case rather than crash.
if ops is not None:
from ..graph_maintenance import enqueue_relink_victims
await enqueue_relink_victims(conn, bank_id, [str(uid) for uid in existing_unit_ids], ops=ops)
# Explicitly delete memory_units by document_id BEFORE deleting the
# document row. The CASCADE from documents→chunks→memory_units only
# catches units that have a non-NULL chunk_id FK. Units with chunk_id=NULL
@@ -5,9 +5,10 @@ Link creation utilities for temporal, semantic, and entity links.
import logging
import time
from datetime import UTC, datetime, timedelta
from uuid import UUID
from ..._vector_index import ann_search_tuning_settings, configured_vector_extension
from ..memory_engine import fq_table
from .types import EntityLink
logger = logging.getLogger(__name__)
@@ -365,6 +366,136 @@ async def resolve_entities_only(
return resolved_entity_ids, entity_to_unit, unit_to_entity_ids
async def build_entity_links_from_resolved(
entity_resolver,
conn,
bank_id: str,
unit_ids: list[str],
resolved_entity_ids: list[str],
entity_to_unit: list[tuple],
unit_to_entity_ids: dict[str, list[str]],
log_buffer: list[str] = None,
skip_unit_entities_insert: bool = False,
ops=None,
) -> list["EntityLink"]:
"""
Build entity links between units that share entities.
Queries unit_entities to find which existing units share entities with the
new units, then generates EntityLink objects for UI graph visualization.
Args:
entity_resolver: EntityResolver instance
conn: Database connection
bank_id: Bank identifier
unit_ids: Actual unit IDs (must already be inserted in the DB)
resolved_entity_ids: Entity IDs from resolve_entities_only
entity_to_unit: Mapping from resolve_entities_only
unit_to_entity_ids: Mapping from resolve_entities_only
log_buffer: Optional logging buffer
skip_unit_entities_insert: If True, skip unit_entities INSERT (already done in Phase 2)
Returns:
List of EntityLink objects for batch insertion
"""
if not resolved_entity_ids:
return []
if not skip_unit_entities_insert:
# Insert unit-entity links (used in fallback path where Phase 2 didn't do this)
substep_start = time.time()
unit_entity_pairs = []
for idx, (unit_id, _local_idx, fact_date) in enumerate(entity_to_unit):
# Propagate the unit's fact_date so entity_cooccurrences.last_cooccurred
# reflects the event timeline, not the ingest moment.
unit_entity_pairs.append((unit_id, resolved_entity_ids[idx], fact_date))
await entity_resolver.link_units_to_entities_batch(unit_entity_pairs, conn=conn)
_log(
log_buffer,
f" [6.2.3] Create unit-entity links (batched): {len(unit_entity_pairs)} links in {time.time() - substep_start:.3f}s",
level="debug",
)
# Build entity links between units that share entities
substep_start = time.time()
all_entity_ids = set()
for entity_ids_list in unit_to_entity_ids.values():
all_entity_ids.update(entity_ids_list)
_log(log_buffer, f" [6.3] Creating entity links for {len(all_entity_ids)} unique entities...", level="debug")
MAX_LINKS_PER_ENTITY = 10
entity_to_units = {}
if all_entity_ids:
query_start = time.time()
import uuid
entity_id_list = [uuid.UUID(eid) if isinstance(eid, str) else eid for eid in all_entity_ids]
limit_per_entity = MAX_LINKS_PER_ENTITY + len(unit_ids) # room for new units + existing cap
rows = await ops.fetch_entity_unit_fanout(
conn,
fq_table("unit_entities"),
entity_id_list,
limit_per_entity,
)
_log(
log_buffer,
f" [6.3.1] Query unit_entities (LATERAL): {len(rows)} rows in {time.time() - query_start:.3f}s",
level="debug",
)
group_start = time.time()
for row in rows:
entity_id = row["entity_id"]
if entity_id not in entity_to_units:
entity_to_units[entity_id] = []
entity_to_units[entity_id].append(row["unit_id"])
_log(log_buffer, f" [6.3.2] Group by entity_id: {time.time() - group_start:.3f}s", level="debug")
link_gen_start = time.time()
links: list[EntityLink] = []
new_unit_set = set(unit_ids)
def to_uuid(val) -> UUID:
return UUID(val) if isinstance(val, str) else val
for entity_id, units_with_entity in entity_to_units.items():
entity_uuid = to_uuid(entity_id)
new_units = [u for u in units_with_entity if str(u) in new_unit_set or u in new_unit_set]
existing_units = [u for u in units_with_entity if str(u) not in new_unit_set and u not in new_unit_set]
new_units_to_link = new_units[-MAX_LINKS_PER_ENTITY:] if len(new_units) > MAX_LINKS_PER_ENTITY else new_units
for i, unit_id_1 in enumerate(new_units_to_link):
for unit_id_2 in new_units_to_link[i + 1 :]:
links.append(
EntityLink(from_unit_id=to_uuid(unit_id_1), to_unit_id=to_uuid(unit_id_2), entity_id=entity_uuid)
)
links.append(
EntityLink(from_unit_id=to_uuid(unit_id_2), to_unit_id=to_uuid(unit_id_1), entity_id=entity_uuid)
)
existing_to_link = existing_units[-MAX_LINKS_PER_ENTITY:]
for new_unit in new_units:
for existing_unit in existing_to_link:
links.append(
EntityLink(from_unit_id=to_uuid(new_unit), to_unit_id=to_uuid(existing_unit), entity_id=entity_uuid)
)
links.append(
EntityLink(from_unit_id=to_uuid(existing_unit), to_unit_id=to_uuid(new_unit), entity_id=entity_uuid)
)
_log(log_buffer, f" [6.3.3] Generate {len(links)} links: {time.time() - link_gen_start:.3f}s", level="debug")
_log(
log_buffer,
f" [6.3] Entity link creation: {len(links)} links for {len(all_entity_ids)} unique entities in {time.time() - substep_start:.3f}s",
level="debug",
)
return links
async def create_temporal_links_batch_per_fact(
conn,
bank_id: str,
@@ -570,18 +701,16 @@ async def compute_semantic_links_ann(
# `relation "_ann_seeds" does not exist` on the second statement.
#
# Using ON COMMIT DROP + SET LOCAL also means we don't have to remember to
# manually drop the temp table or reset the per-backend ANN tuning GUC —
# the transaction end handles both.
# manually drop the temp table or reset hnsw.ef_search — the transaction
# end handles both.
rows: list = []
async with conn.transaction():
# Transaction-local ANN tuning. Each supported backend exposes its own
# GUC (hnsw.ef_search on pgvector, vchordrq.probes on vchord); the
# dispatcher returns the right knob for the configured backend with a
# value tuned for top-50 semantic link creation (lower recall but much
# lower latency than the recall-side default). SET LOCAL auto-reverts
# at commit, so we don't pollute the pool for subsequent queries.
for guc, value in ann_search_tuning_settings(configured_vector_extension(), kind="low_latency"):
await conn.execute(f"SET LOCAL {guc} = {value}")
# Transaction-local ef_search. Default 400 is tuned for recall precision
# but at 164k units each HNSW probe takes 94ms. ef_search=60 gives 2.7ms
# per probe (35x faster) with sufficient accuracy for top-50 semantic
# link creation. SET LOCAL auto-reverts at commit, so we don't pollute
# the pool for subsequent recall queries.
await conn.execute("SET LOCAL hnsw.ef_search = 60")
t_setup = time_mod.time()
await conn.execute("CREATE TEMP TABLE _ann_seeds (unit_id text, emb_text text, fact_type text) ON COMMIT DROP")
@@ -760,6 +889,29 @@ async def create_semantic_links_batch(
raise
async def insert_entity_links_batch(conn, links: list[EntityLink], bank_id: str, chunk_size: int = 5000, ops=None):
"""
Bulk-insert entity links via sorted INSERT FROM unnest().
Args:
conn: Database connection
links: List of EntityLink objects
bank_id: Bank identifier (stored directly on memory_links for fast filtering)
chunk_size: Number of rows per INSERT chunk (default 5000)
"""
if not links:
return
import time as time_mod
total_start = time_mod.time()
tuples = [(link.from_unit_id, link.to_unit_id, link.link_type, link.weight, link.entity_id) for link in links]
await _bulk_insert_links(conn, tuples, bank_id=bank_id, chunk_size=chunk_size, ops=ops)
logger.debug(
f" [9.TOTAL] Entity links batch insert ({len(tuples)} rows): {time_mod.time() - total_start:.3f}s"
)
async def create_causal_links_batch(
conn,
bank_id: str,
@@ -100,6 +100,7 @@ from .types import (
ChunkMetadata,
EntityResolutionResult,
Phase1Result,
Phase3Context,
ProcessedFact,
RetainContent,
RetainContentDict,
@@ -107,9 +108,6 @@ from .types import (
logger = logging.getLogger(__name__)
RetainOutboxCallback = Callable[[asyncpg.Connection], Awaitable[None]]
RetainOutboxCallbackFactory = Callable[[list[RetainContentDict]], RetainOutboxCallback | None]
def _build_retain_params(contents_dicts, document_tags=None, doc_contents=None):
"""Build retain_params and merged_tags from content dicts."""
@@ -258,27 +256,30 @@ async def _insert_facts_and_links(
skip_semantic_links: bool = False,
outbox_callback=None,
ops=None,
) -> list[list[str]]:
) -> tuple[list[list[str]], Phase3Context]:
"""
Phase 2 of the retain pipeline: insert facts and retrieval-critical links.
Runs inside a single database transaction to ensure atomicity of the data
that retrieval depends on (facts, unit_entities, temporal/semantic/causal links).
Entity edges for UI graph visualization are derived on demand from
unit_entities by the /graph endpoint, so no entity rows are written to
memory_links here.
Entity link generation and insertion for UI visualization are NOT done here
only the unit_entities INSERT (FK to memory_units) stays in the transaction.
Entity link building is deferred to Phase 3 (post-transaction, best-effort).
"""
set_stage("retain.phase2.insert_facts")
unit_ids = await fact_storage.insert_facts_batch(conn, bank_id, processed_facts, ops=ops)
step_start = time.time()
log_buffer.append(f" Insert facts: {len(unit_ids)} units in {time.time() - step_start:.3f}s")
# Context for Phase 3 entity link building (after transaction commits)
phase3_context = Phase3Context()
if unit_ids:
# Entity resolution was done in Phase 1 (separate connection).
# Remap placeholder IDs to actual unit IDs.
step_start = time.time()
remapped_entity_to_unit, _remapped_unit_to_entity_ids, remapped_semantic = _remap_phase1_results(
remapped_entity_to_unit, remapped_unit_to_entity_ids, remapped_semantic = _remap_phase1_results(
resolved_entity_ids, entity_to_unit, unit_to_entity_ids, semantic_ann_links or [], unit_ids
)
# Update semantic_ann_links with remapped IDs for Phase 2
@@ -292,6 +293,13 @@ async def _insert_facts_and_links(
]
await entity_resolver.link_units_to_entities_batch(unit_entity_pairs, conn=conn)
log_buffer.append(f" Insert unit_entities: {len(unit_entity_pairs)} pairs in {time.time() - step_start:.3f}s")
# Save context for Phase 3 entity link building (after commit)
phase3_context = Phase3Context(
unit_ids=unit_ids,
resolved_entity_ids=resolved_entity_ids,
entity_to_unit=remapped_entity_to_unit,
unit_to_entity_ids=remapped_unit_to_entity_ids,
)
# Create temporal links
step_start = time.time()
@@ -332,10 +340,52 @@ async def _insert_facts_and_links(
# an IndexError (see issue #1037).
result_unit_ids = _map_results_to_contents(contents, processed_facts, unit_ids if unit_ids else [])
if outbox_callback is not None:
if outbox_callback:
await outbox_callback(conn)
return result_unit_ids
return result_unit_ids, phase3_context
async def _build_and_insert_entity_links_phase3(
pool: Any,
entity_resolver,
bank_id: str,
phase3_ctx: Phase3Context,
log_buffer: list[str],
) -> None:
"""
Phase 3 helper: build entity links from resolved data and insert them.
Runs on a fresh connection after the main transaction has committed.
Entity links are for UI graph visualization only retrieval uses
the unit_entities self-join instead.
"""
set_stage("retain.phase3.entity_links")
p3_unit_ids = phase3_ctx.unit_ids
p3_resolved = phase3_ctx.resolved_entity_ids
p3_entity_to_unit = phase3_ctx.entity_to_unit
p3_unit_to_entity_ids = phase3_ctx.unit_to_entity_ids
if not p3_unit_ids or not p3_resolved:
return
async with acquire_with_retry(pool) as conn:
step_start = time.time()
entity_links = await entity_processing.build_entity_links(
entity_resolver,
conn,
bank_id,
p3_unit_ids,
p3_resolved,
p3_entity_to_unit,
p3_unit_to_entity_ids,
log_buffer,
skip_unit_entities_insert=True, # Already inserted in Phase 2
ops=pool.ops,
)
if entity_links:
await entity_processing.insert_entity_links_batch(conn, entity_links, bank_id, ops=pool.ops)
log_buffer.append(f" Entity links (viz): {len(entity_links)} links in {time.time() - step_start:.3f}s")
async def _extract_and_embed(
@@ -399,11 +449,8 @@ async def retain_batch(
document_tags: list[str] | None = None,
operation_id: str | None = None,
schema: str | None = None,
outbox_callback: RetainOutboxCallback | None = None,
outbox_callback_factory: RetainOutboxCallbackFactory | None = None,
outbox_callback: Callable[["asyncpg.Connection"], Awaitable[None]] | None = None,
db_semaphore: "asyncio.Semaphore | None" = None,
document_body_override: str | None = None,
chunk_index_offset: int = 0,
) -> tuple[list[list[str]], TokenUsage, int | None]:
"""
Process a batch of content through the retain pipeline.
@@ -412,14 +459,6 @@ async def retain_batch(
only re-processes chunks whose content has changed. Unchanged chunks keep
their existing facts, entities, and links.
``chunk_index_offset`` shifts the chunk_index (and therefore the derived
``chunk_id = {bank}_{doc}_{index}``) of every chunk this call stores. The
in-process splitter slices an oversized single item into several
sub-batches that all share one document_id and run sequentially; without
a per-document offset each sub-batch would restart chunk_index at 0, so
their chunk_ids collide and later sub-batches overwrite earlier chunks
leaving only one sub-batch's worth of chunks/memories behind (issue #1888).
Returns a three-tuple of:
* per-content-item unit ID lists
* aggregate LLM token usage
@@ -468,10 +507,6 @@ async def retain_batch(
total_usage = TokenUsage()
total_processed_tokens: int | None = 0
for doc_key, (group_dicts, group_contents) in groups.items():
group_outbox_callback = (
outbox_callback_factory(group_dicts) if outbox_callback_factory is not None else outbox_callback
)
group_ids, group_usage, group_processed = await retain_batch(
pool=pool,
embeddings_model=embeddings_model,
@@ -487,11 +522,8 @@ async def retain_batch(
document_tags=document_tags,
operation_id=operation_id,
schema=schema,
outbox_callback=group_outbox_callback,
outbox_callback_factory=outbox_callback_factory,
outbox_callback=outbox_callback,
db_semaphore=db_semaphore,
document_body_override=document_body_override,
chunk_index_offset=chunk_index_offset,
)
for group_idx, orig_idx in enumerate(original_indices[doc_key]):
if group_idx < len(group_ids):
@@ -633,7 +665,6 @@ async def retain_batch(
schema,
outbox_callback,
db_semaphore,
document_body_override=document_body_override,
)
if delta_result is not None:
return delta_result
@@ -690,8 +721,6 @@ async def retain_batch(
schema=schema,
outbox_callback=outbox_callback,
db_semaphore=db_semaphore,
document_body_override=document_body_override,
chunk_index_offset=chunk_index_offset,
)
@@ -829,8 +858,6 @@ async def _streaming_retain_batch(
schema: str | None = None,
outbox_callback: Callable[["asyncpg.Connection"], Awaitable[None]] | None = None,
db_semaphore: "asyncio.Semaphore | None" = None,
document_body_override: str | None = None,
chunk_index_offset: int = 0,
) -> tuple[list[list[str]], TokenUsage]:
"""
Process a large document in streaming mini-batches to bound memory usage.
@@ -864,15 +891,7 @@ async def _streaming_retain_batch(
# document exists with a matching content_hash and has committed chunks,
# the producer can skip already-extracted chunks to avoid duplicate work.
existing_chunk_hashes: set[str] = set()
# When the caller is processing a sub-batch sliced out of an oversized
# item (see _split_contents_into_sub_batches), document_body_override
# carries the full original document body. Use it for the doc-row write
# so documents.original_text stores the complete payload, not just this
# slice (issue #1838).
if document_body_override is not None:
combined_content = document_body_override
else:
combined_content = "\n".join([c.get("content", "") for c in contents_dicts])
combined_content = "\n".join([c.get("content", "") for c in contents_dicts])
# Memory: contents_dicts content strings are now captured in combined_content.
# Clear them from the dicts to release the per-item copies (can be multi-MB each).
for d in contents_dicts:
@@ -1052,20 +1071,14 @@ async def _streaming_retain_batch(
# Adjust chunk indices to use the original global position (global_idx)
# so that chunk_id = {bank}_{doc}_{chunk_index} is deterministic regardless
# of task completion order. content_index is batch-relative for result grouping.
#
# chunk_index_offset continues the document's chunk_index sequence
# when this call is one of several sequential sub-batches sliced
# from a single oversized item sharing one document_id — without it
# each sub-batch restarts at 0 and their chunk_ids collide (#1888).
doc_chunk_index = global_idx + chunk_index_offset
for fact in extracted:
fact.content_index = content_idx_in_batch
if fact.chunk_index is not None:
fact.chunk_index = doc_chunk_index
fact.chunk_index = global_idx
for pf in processed:
pf.content_index = content_idx_in_batch
for cm in chunk_meta:
cm.chunk_index = doc_chunk_index
cm.chunk_index = global_idx
batch_contents.append(content)
batch_extracted.extend(extracted)
@@ -1178,6 +1191,7 @@ async def _streaming_retain_batch(
p2_start = time.time()
batch_result_ids = None
phase3_ctx = None
async with acquire_with_retry(pool) as conn:
async with conn.transaction():
# --- Document ownership gate ---
@@ -1266,7 +1280,7 @@ async def _streaming_retain_batch(
# Insert facts and links — skip semantic links entirely in streaming
# mode; they are created in a single final ANN pass after all batches.
batch_result_ids = await _insert_facts_and_links(
batch_result_ids, phase3_ctx = await _insert_facts_and_links(
conn,
entity_resolver,
bank_id,
@@ -1286,13 +1300,15 @@ async def _streaming_retain_batch(
logger.info(f"[streaming] Phase 2 (write txn): {time.time() - p2_start:.3f}s")
# Best-effort: flush entity_cooccurrences and other deferred stats.
try:
await entity_resolver.flush_pending_stats()
except Exception:
logger.warning(
f"Entity stats flush (consumer batch {consumer_batch_idx + 1}) failed", exc_info=True
)
# Best-effort: entity viz + stats (fast, not semantic ANN)
if phase3_ctx is not None:
try:
await entity_resolver.flush_pending_stats()
await _build_and_insert_entity_links_phase3(
pool, entity_resolver, bank_id, phase3_ctx, log_buffer
)
except Exception:
logger.warning(f"Phase 3 stats (consumer batch {consumer_batch_idx + 1}) failed", exc_info=True)
logger.info(
f"[streaming] Consumer batch {consumer_batch_idx + 1} total "
@@ -1521,8 +1537,6 @@ async def _try_delta_retain(
schema,
outbox_callback,
db_semaphore: "asyncio.Semaphore | None" = None,
*,
document_body_override: str | None = None,
) -> tuple[list[list[str]], TokenUsage, int | None] | None:
"""
Attempt delta retain for a document upsert. Returns result tuple if delta
@@ -1608,7 +1622,6 @@ async def _try_delta_retain(
log_buffer,
start_time,
outbox_callback,
document_body_override=document_body_override,
)
# Build content items for only the changed/new chunks
@@ -1625,7 +1638,6 @@ async def _try_delta_retain(
log_buffer,
start_time,
outbox_callback,
document_body_override=document_body_override,
)
# Extract facts and generate embeddings (shared pipeline)
@@ -1685,13 +1697,7 @@ async def _try_delta_retain(
# Update document metadata (no delete)
step_start = time.time()
# When this sub-batch is one slice of an oversized item
# split across multiple sub-batches, store the full body
# (issue #1838) instead of just the slice.
if document_body_override is not None:
combined_content = document_body_override
else:
combined_content = "\n".join([c.get("content", "") for c in contents_dicts])
combined_content = "\n".join([c.get("content", "") for c in contents_dicts])
retain_params, merged_tags = _build_retain_params(contents_dicts, document_tags)
await fact_storage.upsert_document_metadata(
conn,
@@ -1760,7 +1766,7 @@ async def _try_delta_retain(
# Insert facts and retrieval-critical links.
# Use delta_contents (the changed/new chunks) as the content list,
# since extracted_facts have content_index relative to delta_contents.
result_unit_ids = await _insert_facts_and_links(
result_unit_ids, phase3_ctx = await _insert_facts_and_links(
conn,
entity_resolver,
bank_id,
@@ -1777,11 +1783,12 @@ async def _try_delta_retain(
ops=pool.ops,
)
# Flush deferred entity_cooccurrences stats (post-transaction, best-effort).
# PHASE 3 — Best-Effort Display Data (post-transaction)
try:
await entity_resolver.flush_pending_stats()
await _build_and_insert_entity_links_phase3(pool, entity_resolver, bank_id, phase3_ctx, log_buffer)
except Exception:
logger.warning("Entity stats flush failed — retrieval unaffected", exc_info=True)
logger.warning("Phase 3 (best-effort display data) failed — retrieval unaffected", exc_info=True)
total_time = time.time() - start_time
log_buffer.append(f"{'=' * 60}")
@@ -1816,8 +1823,6 @@ async def _delta_metadata_only(
log_buffer,
start_time,
outbox_callback,
*,
document_body_override: str | None = None,
):
"""Handle the case where no chunks changed — just update document metadata and tags."""
async with acquire_with_retry(pool) as conn:
@@ -1828,12 +1833,7 @@ async def _delta_metadata_only(
document_id,
bank_id,
)
# When this sub-batch is a slice of an oversized item, write the
# full original body (issue #1838) instead of just the slice.
if document_body_override is not None:
combined_content = document_body_override
else:
combined_content = "\n".join([c.get("content", "") for c in contents_dicts])
combined_content = "\n".join([c.get("content", "") for c in contents_dicts])
retain_params, merged_tags = _build_retain_params(contents_dicts, document_tags)
await fact_storage.upsert_document_metadata(
conn,
@@ -1844,7 +1844,7 @@ async def _delta_metadata_only(
merged_tags,
)
await fact_storage.update_memory_units_tags(conn, bank_id, document_id, merged_tags)
if outbox_callback is not None:
if outbox_callback:
await outbox_callback(conn)
total_time = time.time() - start_time
@@ -224,6 +224,21 @@ class ProcessedFact:
)
@dataclass
class Phase3Context:
"""
Data passed from Phase 2 to Phase 3 for entity link building.
Contains the unit IDs and entity resolution data needed to build
entity links for UI graph visualization after the write transaction commits.
"""
unit_ids: list[str] = field(default_factory=list)
resolved_entity_ids: list[str] = field(default_factory=list)
entity_to_unit: list[tuple] = field(default_factory=list)
unit_to_entity_ids: dict[str, list[str]] = field(default_factory=dict)
@dataclass
class EntityResolutionResult:
"""
@@ -248,6 +263,21 @@ class Phase1Result:
semantic_ann_links: list[tuple]
@dataclass
class EntityLink:
"""
Link between two memory units through a shared entity.
Used for entity-based graph connections in the memory graph.
"""
from_unit_id: UUID
to_unit_id: UUID
entity_id: UUID
link_type: str = "entity"
weight: float = 1.0
@dataclass
class RetainBatch:
"""
@@ -283,12 +283,12 @@ class LinkExpansionRetriever(GraphRetriever):
score transformations. The three CTEs share one connection slot important
for asyncpg which does not allow concurrent queries on the same connection.
Index coverage:
entity: idx_unit_entities_entity_unit (entity_id, unit_id) entity
expansion traverses unit_entities, not memory_links.
semantic: idx_memory_links_from_type_weight / _to_type_weight
(from_unit_id|to_unit_id, link_type, weight DESC) serve both
outgoing and incoming sides as single composite index scans.
Index coverage (requires migration d2e3f4a5b6c7):
entity: idx_memory_links_entity_covering (from_unit_id) INCLUDE (to_unit_id, entity_id)
WHERE link_type = 'entity' index-only scan, no heap reads
semantic incoming:
idx_memory_links_to_type_weight (to_unit_id, link_type, weight DESC)
replaces costly BitmapAnd of two separate scans
"""
config = get_config()
ml = fq_table("memory_links")
@@ -212,25 +212,16 @@ class CrossEncoderReranker:
# Get cross-encoder scores
scores = await self.cross_encoder.predict(pairs)
# Normalize scores to [0, 1] range.
# External API rerankers (Cohere, Jina, llama.cpp/Qwen, etc.) return
# calibrated relevance_score already in [0, 1]. These are used as-is
# so that absolute confidence is preserved — a top candidate scoring
# 0.007 stays low rather than being inflated to 1.0 by rank normalization.
# Local models return logits (any real number) — sigmoid is appropriate.
# Normalize scores using sigmoid to [0, 1] range
# Cross-encoder returns logits which can be negative
import math
import numpy as np
def _sigmoid(x: float) -> float:
def sigmoid(x):
return 1 / (1 + np.exp(-x))
if scores and min(scores) >= 0.0 and max(scores) <= 1.0:
# Scores already in [0, 1] — pass through to preserve absolute
# confidence signal from calibrated rerankers.
normalized_scores = list(scores)
else:
# Scores are logits (e.g. local sentence-transformers models).
# Sigmoid maps (-inf, +inf) to (0, 1).
normalized_scores = [_sigmoid(score) for score in scores]
normalized_scores = [sigmoid(score) for score in scores]
# Create ScoredResult objects with cross-encoder scores
scored_results = []
@@ -225,7 +225,6 @@ async def retrieve_semantic_bm25_combined(
groups_clause=groups_clause,
arm_index=i,
text_search_extension=text_ext,
bm25_language=config.text_search_extension_native_language,
extra_where=created_range_clause,
)
)
@@ -466,8 +465,6 @@ async def retrieve_temporal_combined(
best_date = ep["mentioned_at"]
if best_date:
if best_date.tzinfo is None:
best_date = best_date.replace(tzinfo=UTC)
days_from_mid = abs((best_date - mid_date).total_seconds() / 86400)
temporal_proximity = 1.0 - min(days_from_mid / (total_days / 2), 1.0) if total_days > 0 else 1.0
else:
@@ -561,8 +558,6 @@ async def retrieve_temporal_combined(
neighbor_best_date = n["mentioned_at"]
if neighbor_best_date:
if neighbor_best_date.tzinfo is None:
neighbor_best_date = neighbor_best_date.replace(tzinfo=UTC)
days_from_mid = abs((neighbor_best_date - mid_date).total_seconds() / 86400)
neighbor_temporal_proximity = (
1.0 - min(days_from_mid / (total_days / 2), 1.0) if total_days > 0 else 1.0
@@ -407,7 +407,6 @@ class SQLDialect(ABC):
groups_clause: str = "",
arm_index: int = 0,
text_search_extension: str = "native",
bm25_language: str = "english",
extra_where: str = "",
) -> str:
"""Build a BM25/full-text search subquery arm.
@@ -427,9 +426,7 @@ class SQLDialect(ABC):
arm_index: Index of this arm in the UNION ALL (used by Oracle for
unique SCORE labels).
text_search_extension: Full-text search backend ("native", "vchord",
"pg_textsearch", "pgroonga"). Only relevant for PostgreSQL.
bm25_language: PostgreSQL text search dictionary used by the native
backend (e.g. "english", "french"). Ignored by other backends.
"pg_textsearch"). Only relevant for PostgreSQL.
extra_where: Optional additional WHERE clause fragment (e.g. time range filter).
"""
...
@@ -270,7 +270,6 @@ class OracleDialect(SQLDialect):
groups_clause: str = "",
arm_index: int = 0,
text_search_extension: str = "native",
bm25_language: str = "english",
extra_where: str = "",
) -> str:
# Oracle Text: CONTAINS() / SCORE() with the CTXSYS.CONTEXT index.
@@ -182,7 +182,6 @@ class PostgreSQLDialect(SQLDialect):
groups_clause: str = "",
arm_index: int = 0,
text_search_extension: str = "native",
bm25_language: str = "english",
extra_where: str = "",
) -> str:
if text_search_extension == "vchord":
@@ -194,35 +193,10 @@ class PostgreSQLDialect(SQLDialect):
bm25_score_expr = f"-({text_param} <@> to_bm25query({text_param}, 'idx_memory_units_text_search'))"
bm25_order_by = f"text <@> to_bm25query({text_param}, 'idx_memory_units_text_search') ASC"
bm25_where_filter = ""
elif text_search_extension == "pgroonga":
# &@~ accepts pgroonga's query syntax (raw query text). pgroonga_score
# returns a non-negative relevance score (higher = better).
bm25_score_expr = "pgroonga_score(tableoid, ctid)"
bm25_order_by = f"{bm25_score_expr} DESC"
bm25_where_filter = (
f"AND (COALESCE(text, '') || ' ' || COALESCE(context, '') || ' ' || COALESCE(text_signals, '')) "
f"&@~ {text_param}"
)
elif text_search_extension == "pg_search":
# ParadeDB pg_search: BM25 index over (id, text, context, text_signals)
# with key_field='id'. The @@@ operator on the key_field requires a
# field-qualified query (`text:foo`); to keep the bind-parameter form,
# we fan the query out across all indexed text fields with paradedb.boolean.
bm25_score_expr = "paradedb.score(id)"
bm25_order_by = "paradedb.score(id) DESC"
bm25_where_filter = (
f"AND id @@@ paradedb.boolean(should => ARRAY["
f"paradedb.match('text', {text_param}), "
f"paradedb.match('context', {text_param}), "
f"paradedb.match('text_signals', {text_param})"
f"])"
)
else: # native tsvector
# bm25_language is validated as a PG identifier in HindsightConfig.validate(),
# so embedding it as a SQL literal here is safe.
bm25_score_expr = f"ts_rank_cd(search_vector, to_tsquery('{bm25_language}', {text_param}))"
bm25_score_expr = f"ts_rank_cd(search_vector, to_tsquery('english', {text_param}))"
bm25_order_by = f"{bm25_score_expr} DESC"
bm25_where_filter = f"AND search_vector @@ to_tsquery('{bm25_language}', {text_param})"
bm25_where_filter = f"AND search_vector @@ to_tsquery('english', {text_param})"
return (
f"(SELECT {cols},"
@@ -247,7 +221,7 @@ class PostgreSQLDialect(SQLDialect):
*,
text_search_extension: str = "native",
) -> str:
if text_search_extension in ("vchord", "pg_textsearch", "pgroonga", "pg_search"):
if text_search_extension in ("vchord", "pg_textsearch"):
return query_text
# native tsvector: join tokens with OR operator
return " | ".join(tokens)
@@ -16,7 +16,7 @@ class S3FileStorage(FileStorage):
S3-compatible object storage backend.
Uses obstore (Rust-backed) for high-throughput async access to
Amazon S3, MinIO, Cloudflare R2, Tigris, and other S3-compliant APIs.
Amazon S3, MinIO, Cloudflare R2, and other S3-compliant APIs.
"""
def __init__(
@@ -1,46 +0,0 @@
"""Shared tiktoken encoding used for token counting and chunking.
Hindsight uses tiktoken purely to *count* and *chunk* arbitrary user content never
to feed a model that relies on tiktoken's special-token vocabulary. With tiktoken's
default ``disallowed_special="all"``, any content that merely *mentions* a special-token
literal (e.g. ``<|endoftext|>``) makes ``encode()`` raise, which surfaces as an HTTP 500
on retain/recall (see issue #1883).
``_SafeEncoding`` disables that check so such literals are counted as ordinary text. Token
counts are unaffected; this only stops the encoder from rejecting valid input. Every token
call site in the engine routes through ``get_token_encoding()``, so the fix is global.
"""
from functools import lru_cache
import tiktoken
class _SafeEncoding:
"""Wraps a tiktoken ``Encoding`` so ``encode()`` never raises on special-token literals."""
def __init__(self, encoding: tiktoken.Encoding) -> None:
self._encoding = encoding
def encode(self, text: str, **kwargs) -> list[int]:
# Count special-token literals as ordinary text instead of rejecting them.
kwargs.setdefault("disallowed_special", ())
return self._encoding.encode(text, **kwargs)
def decode(self, tokens: list[int]) -> str:
return self._encoding.decode(tokens)
@lru_cache(maxsize=1)
def get_token_encoding() -> _SafeEncoding:
"""Cached cl100k_base encoding (GPT-4/3.5) wrapped to tolerate special-token literals.
tiktoken downloads the encoding on first lookup; keeping it lazy means importing
``hindsight_api`` does not require network access.
"""
return _SafeEncoding(tiktoken.get_encoding("cl100k_base"))
def count_tokens(text: str) -> int:
"""Count cl100k_base tokens in ``text`` (tolerant of special-token literals)."""
return len(get_token_encoding().encode(text))
@@ -40,7 +40,6 @@ from hindsight_api.extensions.operation_validator import (
# Core operations
OperationValidationError,
OperationValidatorExtension,
PrecheckContext,
RecallContext,
RecallResult,
ReflectContext,
@@ -73,7 +72,6 @@ __all__ = [
"DeferOperation",
"OperationValidationError",
"OperationValidatorExtension",
"PrecheckContext",
"RecallContext",
"RecallResult",
"ReflectContext",
@@ -146,11 +146,7 @@ class DefaultExtensionContext(ExtensionContext):
# Ensure text search columns/indexes match the configured extension
await asyncio.to_thread(
ensure_text_search_extension,
db_url,
text_search_extension=config.text_search_extension,
pg_search_tokenizer=config.text_search_extension_pg_search_tokenizer,
schema=schema,
ensure_text_search_extension, db_url, text_search_extension=config.text_search_extension, schema=schema
)
def get_memory_engine(self) -> "MemoryEngineInterface":
@@ -82,33 +82,6 @@ class ValidationResult:
# =============================================================================
@dataclass
class PrecheckContext:
"""Context for a pre-body-parse precheck on an operation.
Unlike :class:`RetainContext` / :class:`RecallContext` / etc., this
context is constructed *before* the request body is deserialised. It
therefore intentionally carries only the cheap, already-resolved
pieces of request state:
- ``operation``: a short string identifying the route, e.g. ``"retain"``,
``"recall"``, ``"reflect"``, ``"files_retain"``, ``"mental_model_create"``,
``"mental_model_refresh"``.
- ``bank_id``: parsed from the URL path.
- ``request_context``: the authenticated :class:`RequestContext` (tenant
already resolved by the tenant extension).
Implementations should keep precheck cheap and side-effect-free. The
full per-request validators (``validate_retain`` / ``validate_recall``
/ ``validate_reflect``) still run after the body is parsed and remain
the source of truth for the precise per-call cost / quota arithmetic.
"""
operation: str
bank_id: str
request_context: "RequestContext"
@dataclass
class RetainContext:
"""Context for a retain operation validation (pre-operation).
@@ -434,42 +407,6 @@ class OperationValidatorExtension(Extension, ABC):
- consolidate (mental models consolidation)
"""
# =========================================================================
# Pre-body-parse hook (optional - default no-op)
# =========================================================================
async def precheck(self, ctx: PrecheckContext) -> ValidationResult:
"""
Cheap pre-body-parse check, called before the request body is read.
FastAPI resolves ``Depends`` callables before deserialising the route
body; routes that wire ``precheck`` as a dependency therefore short
-circuit here without ever materialising the JSON payload in memory.
That makes this the right hook for "should this caller be allowed to
spend resources on this request at all" decisions — e.g. a balance
is exhausted, a key is revoked, or a tenant is rate-limited.
Implementations should:
- Be cheap: prefer cached lookups, avoid heavy DB queries.
- Use only data on ``ctx`` (operation name + bank_id + request_context);
the body is not yet available.
- Be conservative on errors: prefer ``ValidationResult.accept()`` so
a transient lookup failure doesn't turn into a request rejection.
The post-body ``validate_*`` hooks still run and remain the source
of truth for the precise per-call cost check.
Default implementation accepts everything. Override to opt in.
Args:
ctx: Pre-body context with operation name, bank_id, and
request_context (tenant already resolved).
Returns:
ValidationResult indicating whether the request may proceed to
body parsing and the post-parse validators.
"""
return ValidationResult.accept()
# =========================================================================
# Pre-operation validation hooks (abstract - must be implemented)
# =========================================================================
+21 -73
View File
@@ -24,15 +24,7 @@ import uvicorn
from . import MemoryEngine, __version__
from .api import create_app
from .banner import print_banner
from .config import (
DEFAULT_ACCESS_LOG,
DEFAULT_WORKERS,
ENV_ACCESS_LOG,
ENV_HOST,
ENV_WORKERS,
HindsightConfig,
_get_raw_config,
)
from .config import DEFAULT_WORKERS, ENV_HOST, ENV_WORKERS, HindsightConfig, _get_raw_config
from .daemon import (
DEFAULT_DAEMON_PORT,
DEFAULT_IDLE_TIMEOUT,
@@ -73,42 +65,29 @@ def _signal_handler(signum, frame):
sys.exit(0)
@dataclasses.dataclass(frozen=True)
class ResolvedDaemonHostPort:
host: str
port: int
def resolve_daemon_host_port(
*,
args_host: str,
args_port: int,
explicit_host: bool,
explicit_port: bool,
) -> ResolvedDaemonHostPort:
def resolve_daemon_host_port(*, args_host: str, args_port: int, config_host: str, config_port: int) -> tuple[str, int]:
"""Resolve host/port for daemon mode.
Defaults to 127.0.0.1 for security, but honors explicit user overrides
via --host flag or HINDSIGHT_API_HOST env var. Uses DEFAULT_DAEMON_PORT
unless the user specified a custom port.
"""
port = args_port if explicit_port else DEFAULT_DAEMON_PORT
port = args_port if args_port != config_port else DEFAULT_DAEMON_PORT
# Only force localhost if the user didn't explicitly set a host
if explicit_host or os.environ.get(ENV_HOST):
if args_host != config_host or os.environ.get(ENV_HOST):
host = args_host
else:
host = "127.0.0.1"
return ResolvedDaemonHostPort(host=host, port=port)
return host, port
@dataclasses.dataclass(frozen=True)
class ParsedCliArgs:
args: argparse.Namespace
explicit_host: bool
explicit_port: bool
def main():
"""Main entry point for the CLI."""
global _memory
# Load configuration from environment (for CLI args defaults)
config = _get_raw_config()
def _parse_cli_args(argv: list[str], config: HindsightConfig) -> ParsedCliArgs:
parser = argparse.ArgumentParser(
prog="hindsight-api",
description="Hindsight API Server",
@@ -116,14 +95,12 @@ def _parse_cli_args(argv: list[str], config: HindsightConfig) -> ParsedCliArgs:
# Server options
parser.add_argument(
"--host",
default=argparse.SUPPRESS,
help=f"Host to bind to (default: {config.host}, env: HINDSIGHT_API_HOST)",
"--host", default=config.host, help=f"Host to bind to (default: {config.host}, env: HINDSIGHT_API_HOST)"
)
parser.add_argument(
"--port",
type=int,
default=argparse.SUPPRESS,
default=config.port,
help=f"Port to bind to (default: {config.port}, env: HINDSIGHT_API_PORT)",
)
parser.add_argument(
@@ -143,18 +120,9 @@ def _parse_cli_args(argv: list[str], config: HindsightConfig) -> ParsedCliArgs:
)
# Access log options
parser.add_argument(
"--access-log",
action="store_true",
default=os.getenv(ENV_ACCESS_LOG, "").lower() in ("1", "true", "yes", "on") or DEFAULT_ACCESS_LOG,
help=f"Enable access log (env: {ENV_ACCESS_LOG}, default: {DEFAULT_ACCESS_LOG})",
)
parser.add_argument(
"--no-access-log",
dest="access_log",
action="store_false",
help="Disable access log (overrides env and default)",
)
parser.add_argument("--access-log", action="store_true", help="Enable access log")
parser.add_argument("--no-access-log", dest="access_log", action="store_false", help="Disable access log (default)")
parser.set_defaults(access_log=False)
# Proxy options
parser.add_argument(
@@ -181,27 +149,7 @@ def _parse_cli_args(argv: list[str], config: HindsightConfig) -> ParsedCliArgs:
help=f"Idle timeout in seconds before auto-exit in daemon mode (default: {DEFAULT_IDLE_TIMEOUT})",
)
args = parser.parse_args(argv)
explicit_host = hasattr(args, "host")
explicit_port = hasattr(args, "port")
if not explicit_host:
args.host = config.host
if not explicit_port:
args.port = config.port
return ParsedCliArgs(args=args, explicit_host=explicit_host, explicit_port=explicit_port)
def main():
"""Main entry point for the CLI."""
global _memory
# Load configuration from environment (for CLI args defaults)
config = _get_raw_config()
parsed_cli_args = _parse_cli_args(sys.argv[1:], config)
args = parsed_cli_args.args
args = parser.parse_args()
# Daemon mode handling.
# is_daemon_child is True when we are the re-exec'd child spawned by
@@ -212,14 +160,12 @@ def main():
is_daemon = args.daemon or is_daemon_child
if is_daemon:
resolved_daemon_host_port = resolve_daemon_host_port(
args.host, args.port = resolve_daemon_host_port(
args_host=args.host,
args_port=args.port,
explicit_host=parsed_cli_args.explicit_host,
explicit_port=parsed_cli_args.explicit_port,
config_host=config.host,
config_port=config.port,
)
args.host = resolved_daemon_host_port.host
args.port = resolved_daemon_host_port.port
# Detach into background (parent re-execs and exits; child redirects
# stdio to log file). No lockfile needed — port binding prevents
@@ -293,6 +239,8 @@ def main():
# When using workers or reload, we must use import string so each worker can import the app
use_import_string = args.workers > 1 or args.reload
# Check for uvloop/winloop availability
import sys
loop_impl = "asyncio"
if sys.platform == "win32":
try:
+2 -110
View File
@@ -44,7 +44,6 @@ _ALL_TOOLS: frozenset[str] = frozenset(
"update_mental_model",
"delete_mental_model",
"refresh_mental_model",
"clear_mental_model",
"list_directives",
"create_directive",
"delete_directive",
@@ -222,7 +221,6 @@ def register_mcp_tools(
"update_mental_model",
"delete_mental_model",
"refresh_mental_model",
"clear_mental_model",
"list_directives",
"create_directive",
"delete_directive",
@@ -279,9 +277,6 @@ def register_mcp_tools(
if "refresh_mental_model" in tools_to_register:
_register_refresh_mental_model(mcp, memory, config)
if "clear_mental_model" in tools_to_register:
_register_clear_mental_model(mcp, memory, config)
# Directive tools
if "list_directives" in tools_to_register:
_register_list_directives(mcp, memory, config)
@@ -443,7 +438,6 @@ _AUDITABLE_MCP_TOOLS: frozenset[str] = frozenset(
"update_mental_model",
"delete_mental_model",
"refresh_mental_model",
"clear_mental_model",
"create_directive",
"delete_directive",
"delete_document",
@@ -799,8 +793,7 @@ def _register_recall(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig)
{"tags": [...], "match": "any_strict"} or compound {"and": [...]}, {"or": [...]}, {"not": {...}}.
Example: [{"not": {"tags": ["closeout"], "match": "any_strict"}}] excludes memories tagged closeout.
Mutually exclusive with tags.
query_timestamp: Temporal context for the query (ISO format, e.g., '2024-01-15T10:30:00Z').
Anchors relative temporal expressions and recency scoring.
query_timestamp: Temporal context for the query (ISO format, e.g., '2024-01-15T10:30:00Z'). Helps retrieve time-relevant memories.
bank_id: Optional bank to search in (defaults to session bank). Use for cross-bank operations.
"""
try:
@@ -870,8 +863,7 @@ def _register_recall(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig)
{"tags": [...], "match": "any_strict"} or compound {"and": [...]}, {"or": [...]}, {"not": {...}}.
Example: [{"not": {"tags": ["closeout"], "match": "any_strict"}}] excludes memories tagged closeout.
Mutually exclusive with tags.
query_timestamp: Temporal context for the query (ISO format, e.g., '2024-01-15T10:30:00Z').
Anchors relative temporal expressions and recency scoring.
query_timestamp: Temporal context for the query (ISO format, e.g., '2024-01-15T10:30:00Z'). Helps retrieve time-relevant memories.
"""
try:
target_bank = config.bank_id_resolver()
@@ -930,7 +922,6 @@ def _register_reflect(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig
response_schema: dict | None = None,
tags: list[str] | None = None,
tags_match: str = "any",
include_based_on: bool = False,
bank_id: str | None = None,
) -> str:
"""
@@ -960,7 +951,6 @@ def _register_reflect(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig
response_schema: Optional JSON schema for structured output. When provided, the response includes a 'structured_output' field.
tags: Optional tags to filter memories by (e.g., ['project:alpha'])
tags_match: How to match tags - 'any' (match any tag) or 'all' (match all tags). Default: 'any'
include_based_on: Include source facts used for synthesis. Defaults to false because broad reflections can exceed MCP client result limits.
bank_id: Optional bank to reflect in (defaults to session bank). Use for cross-bank operations.
"""
try:
@@ -988,8 +978,6 @@ def _register_reflect(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig
reflect_result = await memory.reflect_async(**reflect_kwargs)
result_data = json.loads(reflect_result.model_dump_json(indent=2))
if not include_based_on:
result_data.pop("based_on", None)
if response_schema is not None and hasattr(reflect_result, "structured_output"):
result_data["structured_output"] = reflect_result.structured_output
return json.dumps(result_data, indent=2)
@@ -1011,7 +999,6 @@ def _register_reflect(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig
response_schema: dict | None = None,
tags: list[str] | None = None,
tags_match: str = "any",
include_based_on: bool = False,
) -> dict:
"""
Generate thoughtful analysis by synthesizing stored memories with the bank's personality.
@@ -1040,7 +1027,6 @@ def _register_reflect(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig
response_schema: Optional JSON schema for structured output. When provided, the response includes a 'structured_output' field.
tags: Optional tags to filter memories by (e.g., ['project:alpha'])
tags_match: How to match tags - 'any' (match any tag) or 'all' (match all tags). Default: 'any'
include_based_on: Include source facts used for synthesis. Defaults to false because broad reflections can exceed MCP client result limits.
"""
try:
target_bank = config.bank_id_resolver()
@@ -1067,8 +1053,6 @@ def _register_reflect(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig
reflect_result = await memory.reflect_async(**reflect_kwargs)
result_data = reflect_result.model_dump()
if not include_based_on:
result_data.pop("based_on", None)
if response_schema is not None and hasattr(reflect_result, "structured_output"):
result_data["structured_output"] = reflect_result.structured_output
return result_data
@@ -1781,98 +1765,6 @@ def _register_refresh_mental_model(mcp: FastMCP, memory: MemoryEngine, config: M
return {"error": str(e)}
def _register_clear_mental_model(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig) -> None:
"""Register the clear_mental_model tool."""
if config.include_bank_id_param:
@mcp.tool()
async def clear_mental_model(
mental_model_id: str,
bank_id: str | None = None,
) -> str:
"""
Clear a mental model's content so the next refresh performs a full re-synthesis.
This is useful for delta-mode models that have accumulated drift over many
incremental refreshes. After clearing, call refresh_mental_model to trigger
a clean full rebuild.
Args:
mental_model_id: The ID of the mental model to clear
bank_id: Optional bank (defaults to session bank). Use for cross-bank operations.
"""
try:
target_bank = bank_id or config.bank_id_resolver()
if target_bank is None:
return '{"error": "No bank_id configured"}'
result = await memory.clear_mental_model(
bank_id=target_bank,
mental_model_id=mental_model_id,
request_context=_get_request_context(config),
)
if result is None:
return json.dumps({"error": f"Mental model '{mental_model_id}' not found"})
return json.dumps(
{
"mental_model_id": result["id"],
"status": "cleared",
"message": f"Mental model '{mental_model_id}' content cleared. Call refresh_mental_model to rebuild.",
}
)
except OperationValidationError as e:
logger.warning(f"Operation rejected: {e}")
return json.dumps({"error": str(e)})
except ValueError as e:
return json.dumps({"error": str(e)})
except Exception as e:
logger.error(f"Error clearing mental model: {e}", exc_info=True)
return f'{{"error": "{e}"}}'
else:
@mcp.tool()
async def clear_mental_model(
mental_model_id: str,
) -> dict:
"""
Clear a mental model's content so the next refresh performs a full re-synthesis.
This is useful for delta-mode models that have accumulated drift over many
incremental refreshes. After clearing, call refresh_mental_model to trigger
a clean full rebuild.
Args:
mental_model_id: The ID of the mental model to clear
"""
try:
target_bank = config.bank_id_resolver()
if target_bank is None:
return {"error": "No bank_id configured"}
result = await memory.clear_mental_model(
bank_id=target_bank,
mental_model_id=mental_model_id,
request_context=_get_request_context(config),
)
if result is None:
return {"error": f"Mental model '{mental_model_id}' not found"}
return {
"mental_model_id": result["id"],
"status": "cleared",
"message": f"Mental model '{mental_model_id}' content cleared. Call refresh_mental_model to rebuild.",
}
except OperationValidationError as e:
logger.warning(f"Operation rejected: {e}")
return {"error": str(e)}
except ValueError as e:
return {"error": str(e)}
except Exception as e:
logger.error(f"Error clearing mental model: {e}", exc_info=True)
return {"error": str(e)}
# =========================================================================
# DIRECTIVE TOOLS
# =========================================================================
+21 -116
View File
@@ -27,7 +27,6 @@ from alembic.config import Config
from alembic.script.revision import ResolutionError
from sqlalchemy import Connection, create_engine, text
from ._pg_search import normalize_pg_search_tokenizer, pg_search_bm25_columns
from ._vector_index import (
bootstrap_extension,
detect_vector_extension,
@@ -804,7 +803,6 @@ def ensure_text_search_extension(
database_url: str,
text_search_extension: str = "native",
schema: str | None = None,
pg_search_tokenizer: str | None = None,
) -> None:
"""
Ensure the text search columns and indexes match the configured extension.
@@ -817,18 +815,13 @@ def ensure_text_search_extension(
Args:
database_url: SQLAlchemy database URL
text_search_extension: Configured text search extension one of
"native", "vchord", "pg_textsearch", "pgroonga", or "pg_search"
text_search_extension: Configured text search extension ("native" or "vchord")
schema: Target PostgreSQL schema name (None for public)
pg_search_tokenizer: Optional ParadeDB tokenizer to apply to pg_search
BM25 text fields when indexes are created. Empty keeps the
ParadeDB default.
Raises:
RuntimeError: If extension mismatch with existing data
"""
schema_name = schema or "public"
pg_search_tokenizer = normalize_pg_search_tokenizer(pg_search_tokenizer)
engine = create_engine(to_libpq_url(database_url))
with engine.connect() as conn:
@@ -845,17 +838,6 @@ def ensure_text_search_extension(
elif text_search_extension == "pg_textsearch":
target_column_type = "text"
target_index_type = "bm25"
elif text_search_extension == "pgroonga":
# pgroonga indexes the base text column directly. We keep a dummy
# TEXT column named search_vector for symmetry with pg_textsearch
# and so the column-type mismatch detection above keeps working.
target_column_type = "text"
target_index_type = "pgroonga"
elif text_search_extension == "pg_search":
# ParadeDB: same column type / access method as pg_textsearch.
# Disambiguated below by inspecting the index reloptions (key_field).
target_column_type = "text"
target_index_type = "bm25"
else: # native
target_column_type = "tsvector"
target_index_type = "gin"
@@ -893,18 +875,16 @@ def ensure_text_search_extension(
if not current_column_info:
logger.warning(f"No search_vector column found for {table_name}, will create it")
mismatched_tables.append((table_name, None, None, False))
mismatched_tables.append((table_name, None, None))
continue
# Check column type (udt_name contains the actual type: tsvector, bm25vector, etc.)
current_column_type = current_column_info[1] # udt_name
# Get current index type and definition. The definition lets us
# disambiguate pg_textsearch vs pg_search (both register a `bm25`
# access method but only pg_search uses the `key_field` reloption).
# Get current index type
current_index_info = conn.execute(
text("""
SELECT am.amname, pi.indexdef
SELECT am.amname
FROM pg_indexes pi
JOIN pg_class c ON c.relname = pi.indexname
JOIN pg_am am ON am.oid = c.relam
@@ -916,21 +896,10 @@ def ensure_text_search_extension(
).fetchone()
current_index_type = current_index_info[0] if current_index_info else None
current_index_def = current_index_info[1] if current_index_info else None
# Detect pg_search specifically (vs pg_textsearch) via the key_field reloption
current_is_pg_search = bool(current_index_def and "key_field" in current_index_def)
want_pg_search = text_search_extension == "pg_search"
# Check if column and index types match target
column_matches = current_column_type == target_column_type
index_matches = current_index_type == target_index_type if current_index_type else False
# When both target and current sit at column=text/index=bm25, the
# access-method check alone can't tell pg_textsearch from pg_search —
# require the key_field reloption to agree with the configured backend.
if column_matches and index_matches and target_index_type == "bm25" and target_column_type == "text":
if current_is_pg_search != want_pg_search:
index_matches = False
if not (column_matches and index_matches):
logger.info(
@@ -938,7 +907,7 @@ def ensure_text_search_extension(
f"column={current_column_type} (want {target_column_type}), "
f"index={current_index_type} (want {target_index_type})"
)
mismatched_tables.append((table_name, current_column_type, current_index_type, current_is_pg_search))
mismatched_tables.append((table_name, current_column_type, current_index_type))
# Check if table has data
row_count = conn.execute(text(f"SELECT COUNT(*) FROM {schema_name}.{table_name}")).scalar()
@@ -956,20 +925,14 @@ def ensure_text_search_extension(
# If there's data in any mismatched table, raise error
if tables_with_data:
table_list = ", ".join([f"{table}({count} rows)" for table, count in tables_with_data])
# Detect current extension from column type, index type, and (for the
# text/bm25 ambiguity) the key_field reloption. tsvector is
# unambiguous; text could be pg_textsearch, pgroonga, or pg_search.
# Detect current extension from column type
current_col_type = mismatched_tables[0][1]
current_idx_type = mismatched_tables[0][2]
first_is_pg_search = mismatched_tables[0][3]
if current_col_type == "tsvector":
current_ext = "native"
elif current_col_type == "bm25vector":
current_ext = "vchord"
elif current_col_type == "text" and current_idx_type == "pgroonga":
current_ext = "pgroonga"
elif current_col_type == "text":
current_ext = "pg_search" if first_is_pg_search else "pg_textsearch"
current_ext = "pg_textsearch"
else:
current_ext = "unknown"
raise RuntimeError(
@@ -984,7 +947,7 @@ def ensure_text_search_extension(
# Tables are empty, safe to recreate columns/indexes
logger.info(f"Recreating text search columns/indexes for {text_search_extension}")
for table_name, current_col_type, current_idx_type, _was_pg_search in mismatched_tables:
for table_name, current_col_type, current_idx_type in mismatched_tables:
# Drop existing index if it exists
if current_idx_type:
logger.info(f"Dropping {current_idx_type} index on {table_name}")
@@ -1037,79 +1000,21 @@ def ensure_text_search_extension(
WITH (text_config='english')
""")
)
elif text_search_extension == "pgroonga":
# Ensure pgroonga extension is available
try:
conn.execute(text("CREATE EXTENSION IF NOT EXISTS pgroonga CASCADE"))
except Exception:
# Extension might already exist or user lacks permissions — verify
has_ext = conn.execute(text("SELECT 1 FROM pg_extension WHERE extname = 'pgroonga'")).fetchone()
if not has_ext:
raise
logger.info(f"Creating dummy TEXT search_vector on {table_name} for pgroonga")
# pgroonga indexes the base text column directly, but we keep a
# dummy search_vector column for symmetry with pg_textsearch and
# so the column-type mismatch detection above keeps working.
conn.execute(text(f"ALTER TABLE {schema_name}.{table_name} ADD COLUMN search_vector TEXT"))
# pgroonga index expression mirrors pg_textsearch
if table_name == "memory_units":
index_expr = (
"(COALESCE(text, '') || ' ' || COALESCE(context, '') || ' ' || COALESCE(text_signals, ''))"
)
else: # reflections
index_expr = "(COALESCE(name, '') || ' ' || content)"
logger.info(f"Creating pgroonga index on {table_name}")
# TokenBigram is the polyglot default — falls back to whitespace
# tokenization for space-separated languages and bigram for CJK.
# NormalizerNFKC150 handles Unicode normalization (full/half-width,
# case folding, etc.) which materially improves Japanese recall.
conn.execute(
text(f"""
CREATE INDEX idx_{table_name.replace(".", "_")}_text_search
ON {schema_name}.{table_name}
USING pgroonga ({index_expr})
WITH (tokenizer='TokenBigram', normalizer='NormalizerNFKC150')
""")
)
elif text_search_extension == "pg_search":
logger.info(f"Creating TEXT column on {table_name}")
# Dummy TEXT column for schema symmetry; pg_search indexes operate on base columns.
conn.execute(text(f"ALTER TABLE {schema_name}.{table_name} ADD COLUMN search_vector TEXT"))
# ParadeDB BM25 index over the table's primary key and text columns.
# Column list mirrors what the initial / text_signals migrations create.
if table_name == "memory_units":
bm25_cols = pg_search_bm25_columns(
"id",
("text", "context", "text_signals"),
pg_search_tokenizer,
)
else: # reflections
bm25_cols = pg_search_bm25_columns(
"id",
("name", "content"),
pg_search_tokenizer,
)
logger.info(f"Creating ParadeDB BM25 index on {table_name}")
conn.execute(
text(f"""
CREATE INDEX idx_{table_name.replace(".", "_")}_text_search
ON {schema_name}.{table_name}
USING bm25 ({bm25_cols})
WITH (key_field='id')
""")
)
else: # native
logger.info(f"Creating tsvector column on {table_name}")
# Plain tsvector column. The application populates search_vector
# at INSERT time via to_tsvector($lang, ...) using the configured
# HINDSIGHT_API_TEXT_SEARCH_EXTENSION_NATIVE_LANGUAGE — see
# ops_postgresql.insert_facts_batch.
conn.execute(text(f"ALTER TABLE {schema_name}.{table_name} ADD COLUMN search_vector tsvector"))
# Different GENERATED expression for each table
if table_name == "memory_units":
generated_expr = "to_tsvector('english', COALESCE(text, '') || ' ' || COALESCE(context, ''))"
else: # reflections
generated_expr = "to_tsvector('english', COALESCE(name, '') || ' ' || content)"
conn.execute(
text(f"""
ALTER TABLE {schema_name}.{table_name}
ADD COLUMN search_vector tsvector
GENERATED ALWAYS AS ({generated_expr}) STORED
""")
)
# Create GIN index
logger.info(f"Creating GIN index on {table_name}")
-8
View File
@@ -24,7 +24,6 @@ class EmbeddedPostgres:
password: str = DEFAULT_PASSWORD,
database: str = DEFAULT_DATABASE,
name: str = "hindsight",
config: dict[str, str] | None = None,
**kwargs,
):
self.port = port # None means pg0 will auto-assign
@@ -32,11 +31,6 @@ class EmbeddedPostgres:
self.password = password
self.database = database
self.name = name
# Extra postgresql.conf settings forwarded to Pg0 (e.g. ``max_connections``).
# Useful when tests spawn many xdist workers that each open a pool against
# the same pg0 instance — the postgres default of 100 max_connections is
# easy to exhaust under that fan-out.
self.config = config
self._pg0: Pg0 | None = None
def _get_pg0(self) -> Pg0:
@@ -57,8 +51,6 @@ class EmbeddedPostgres:
# Only set port if explicitly specified
if self.port is not None:
kwargs["port"] = self.port
if self.config is not None:
kwargs["config"] = self.config
self._pg0 = Pg0(**kwargs)
return self._pg0
@@ -129,7 +129,6 @@ PROVIDER_NAME_MAPPING = {
"vertexai": "google",
"groq": "groq",
"ollama": "ollama",
"ollama-cloud": "ollama",
"lmstudio": "lmstudio",
"openai-codex": "openai",
"claude-code": "anthropic",
@@ -256,7 +256,6 @@ def main():
tenant_extension=tenant_extension,
max_slots=config.worker_max_slots,
slot_reservations=config.worker_slot_reservations,
consolidation_bank_priority=config.worker_consolidation_bank_priority or None,
)
# Create the HTTP app for metrics/health
@@ -24,7 +24,7 @@ from .exceptions import DeferOperation, RetryTaskAt
from .stage import StageHolder, bind_holder
if TYPE_CHECKING:
from hindsight_api.engine.db.base import DatabaseBackend, DatabaseConnection
from hindsight_api.engine.db.base import DatabaseBackend
from hindsight_api.extensions.tenant import TenantExtension
logger = logging.getLogger(__name__)
@@ -133,7 +133,6 @@ class WorkerPoller:
tenant_extension: "TenantExtension | None" = None,
max_slots: int = 10,
slot_reservations: dict[str, int] | None = None,
consolidation_bank_priority: dict[str, int] | None = None,
):
"""
Initialize the worker poller.
@@ -151,11 +150,6 @@ class WorkerPoller:
"retain": 3}). Reserved slots guarantee capacity for that operation type.
Remaining slots (max_slots - sum of reservations) form a shared pool usable
by any operation type. Defaults to {"consolidation": 2} if None.
consolidation_bank_priority: Per-bank priority for consolidation scheduling.
Maps bank name patterns to integer priorities (higher = claimed first).
Patterns support ``*`` as wildcard. A bare ``*`` key is the catch-all default.
When set, consolidation tasks are claimed in priority tiers rather than
pure created_at order. None or empty dict preserves current behavior.
"""
self._backend = backend
self._worker_id = worker_id
@@ -174,9 +168,6 @@ class WorkerPoller:
self._slot_reservations: dict[str, int] = (
slot_reservations if slot_reservations is not None else {"consolidation": 2}
)
self._consolidation_bank_priority: dict[str, int] | None = (
consolidation_bank_priority if consolidation_bank_priority else None
)
# Cache of which optional PG routines are installed on the server
# (probed once, memoised for the life of the poller).
from ..engine.db.optional_routines import OptionalRoutines
@@ -196,18 +187,13 @@ class WorkerPoller:
# schema we serviced so a busy tenant can't monopolize the poll order.
self._next_schema_idx: int = 0
@staticmethod
def _normalize_poll_schema(schema: str | None) -> str | None:
"""Use None internally for the default schema because SQL helpers omit that prefix."""
from ..config import DEFAULT_DATABASE_SCHEMA
return None if schema == DEFAULT_DATABASE_SCHEMA else schema
async def _get_schemas(self) -> list[str | None]:
"""Get list of schemas to poll. Returns [None] for default schema (no prefix)."""
from ..config import DEFAULT_DATABASE_SCHEMA
tenants = await self._tenant_extension.list_tenants()
# Convert default schema to None for SQL compatibility (no prefix), keep others as-is
return [self._normalize_poll_schema(t.schema) for t in tenants]
return [t.schema if t.schema != DEFAULT_DATABASE_SCHEMA else None for t in tenants]
async def _scan_active_schemas(self, schemas: list[str | None]) -> set[str | None]:
"""Find which schemas have pending work.
@@ -227,57 +213,22 @@ class WorkerPoller:
async with self._backend.acquire() as conn:
if await self._optional_routines.is_installed(conn, "schemas_with_pending_work"):
rows = await conn.fetch("SELECT * FROM public.schemas_with_pending_work()")
routine_active = {self._normalize_poll_schema(r[0]) for r in rows}
known_schemas = set(schemas)
active = routine_active & known_schemas
unknown = routine_active - known_schemas
if unknown:
logger.warning(
"Optional PG routine public.schemas_with_pending_work() returned schema(s) "
"not present in tenant discovery: %s",
sorted(str(s) for s in unknown),
return {r[0] for r in rows}
# Fallback: per-schema EXISTS checks from Python
active: set[str | None] = set()
for schema in schemas:
table = fq_table("async_operations", schema)
try:
has_work = await conn.fetchval(
f"SELECT EXISTS(SELECT 1 FROM {table} "
f"WHERE status = 'pending' AND task_payload IS NOT NULL LIMIT 1)"
)
# The optional routine returns PostgreSQL schema names, but the poller uses
# None for the default schema. Older operator-supplied implementations also
# commonly scan tenant_% only; when the default schema is in scope but absent
# from the routine result, verify via the fully-correct per-schema fallback so
# public single-tenant deployments cannot silently starve.
should_verify_with_fallback = (None in known_schemas and None not in active) or (
bool(routine_active) and not active
)
if not should_verify_with_fallback:
return active
fallback_active = await self._scan_active_schemas_by_exists(conn, schemas)
missed = fallback_active - active
if missed:
logger.warning(
"Optional PG routine public.schemas_with_pending_work() missed claimable schema(s) %s; "
"using per-schema fallback for this poll",
sorted(str(s) for s in missed),
)
return fallback_active
return await self._scan_active_schemas_by_exists(conn, schemas)
async def _scan_active_schemas_by_exists(
self, conn: "DatabaseConnection", schemas: list[str | None]
) -> set[str | None]:
"""Find active schemas using per-schema EXISTS checks."""
active: set[str | None] = set()
for schema in schemas:
table = fq_table("async_operations", schema)
try:
has_work = await conn.fetchval(
f"SELECT EXISTS(SELECT 1 FROM {table} "
f"WHERE status = 'pending' AND task_payload IS NOT NULL LIMIT 1)"
)
if has_work:
active.add(schema)
except Exception:
pass
return active
if has_work:
active.add(schema)
except Exception:
pass
return active
async def _get_available_slots(self) -> SlotAvailability:
"""
@@ -477,7 +428,6 @@ class WorkerPoller:
self._worker_id,
reserved_limits,
shared_limit,
consolidation_bank_priority=self._consolidation_bank_priority,
)
if not all_rows:
+5 -12
View File
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
[project]
name = "hindsight-api-slim"
version = "0.7.1"
version = "0.6.2"
description = "Hindsight: Agent Memory That Works Like Human Memory"
readme = "README.md"
requires-python = ">=3.11"
@@ -18,14 +18,10 @@ dependencies = [
"fastapi[standard]>=0.120.3",
"uvicorn>=0.38.0",
"wsproto>=1.0.0",
# Cap below 2.1: SQLAlchemy 2.1 switches the default `postgresql://` DBAPI
# from psycopg2 to psycopg (v3), which we don't ship — a bare install would
# fail migrations with "No module named 'psycopg'". Pin to the tested 2.0
# line (which keeps psycopg2 the default driver) until psycopg3 is adopted.
"sqlalchemy>=2.0.44,<2.1",
"sqlalchemy>=2.0.44",
"alembic>=1.17.1",
"pgvector>=0.4.1",
"greenlet>=3.2.4,<3.4.0", # 3.4.0 lacks arm64 wheels for manylinux_2_41
"greenlet>=3.2.4,<3.4.0", # 3.4.0 lacks arm64 wheels for manylinux_2_41
"psycopg2-binary>=2.9.11",
"tiktoken>=0.12.0",
"httpx>=0.27.0",
@@ -72,7 +68,7 @@ dependencies = [
"tornado>=6.5.5", # DoS multipart/incomplete cookie validation fix
"aiohttp>=3.13.3", # Multiple DoS vulnerabilities
"pygments>=2.20.0", # ReDoS via inefficient GUID regex fix
"claude-agent-sdk>=0.2.82",
"claude-agent-sdk>=0.1.27",
"boto3>=1.42.74",
]
@@ -97,7 +93,7 @@ local-llm = [
"huggingface-hub>=0.20.0",
]
embedded-db = [
"pg0-embedded>=0.14.2",
"pg0-embedded>=0.14.0",
]
oracle = [
"oracledb>=2.5.0",
@@ -146,9 +142,6 @@ addopts = "--timeout 300 -n 8 --dist loadgroup --durations=10 -v"
markers = [
"oracle: Oracle 23ai integration tests (require ORACLE_TEST_DSN env var)",
"hs_llm_mat: LLM minimum acceptance tests — run in CI matrix across multiple providers",
"hs_llm_core: Core pipeline tests that need a real LLM but only one provider",
"integration: Live external-API integration tests (require provider credentials; skipped without)",
"slow: Slow tests (minutes); not run in fast CI",
]
asyncio_mode = "auto"
asyncio_default_fixture_loop_scope = "function"
+17 -51
View File
@@ -95,14 +95,7 @@ def pg0_db_url(db_url, tmp_path_factory, worker_id):
url = url_file.read_text().strip()
else:
# First worker - start pg0
# Bump max_connections so 8 xdist workers * pool_max_size=15 fits well
# under the cap (postgres default is 100, which is easy to exhaust now
# that consolidation_llm_parallelism=4 increases peak conns per op).
pg0 = EmbeddedPostgres(
name=pg0_instance_name,
port=pg0_instance_port,
config={"max_connections": "300"},
)
pg0 = EmbeddedPostgres(name=pg0_instance_name, port=pg0_instance_port)
# Run ensure_running in a new event loop
loop = asyncio.new_event_loop()
@@ -315,7 +308,7 @@ async def oracle_memory(oracle_db_url, embeddings, cross_encoder, query_analyzer
cross_encoder=cross_encoder,
query_analyzer=query_analyzer,
pool_min_size=1,
pool_max_size=15,
pool_max_size=5,
run_migrations=False, # Already ran above
task_backend=SyncTaskBackend(),
)
@@ -420,48 +413,21 @@ def query_analyzer():
@pytest_asyncio.fixture(scope="function")
async def memory(pg0_db_url, embeddings, cross_encoder, query_analyzer):
"""
Provide a MemoryEngine instance using a mock LLM for deterministic tests.
Provide a MemoryEngine instance for each test.
The mock LLM returns canned facts derived from input text, allowing the
full retain recall reflect pipeline to work without real LLM calls.
This makes core tests fast, deterministic, and free from LLM flakiness.
Must be function-scoped because:
1. pytest-xdist runs tests in separate processes with different event loops
2. asyncpg pools are bound to the event loop that created them
3. Each test needs its own pool in its own event loop
Tests that need real LLM output quality should use `memory_real_llm` instead.
Uses small pool sizes since tests run in parallel.
Uses pg0_db_url (a postgresql:// URL) directly, so MemoryEngine won't try to
manage pg0 lifecycle - that's handled by the session-scoped pg0_db_url fixture.
Migrations are disabled here since they're run once at session scope in pg0_db_url.
Uses SyncTaskBackend so async tasks execute immediately (no worker needed).
"""
mem = MemoryEngine(
db_url=pg0_db_url,
memory_llm_provider="mock",
memory_llm_api_key="",
memory_llm_model="mock",
embeddings=embeddings,
cross_encoder=cross_encoder,
query_analyzer=query_analyzer,
pool_min_size=1,
pool_max_size=15,
run_migrations=False,
task_backend=SyncTaskBackend(),
)
await mem.initialize()
yield mem
try:
if mem._pool and not mem._pool._closing:
await mem.close()
except Exception:
pass
@pytest_asyncio.fixture(scope="function")
async def memory_real_llm(pg0_db_url, embeddings, cross_encoder, query_analyzer):
"""
Provide a MemoryEngine instance using a real LLM provider.
Use this fixture ONLY for tests that assert on LLM output quality
(fact extraction accuracy, language preservation, consolidation decisions, etc.).
These tests are non-deterministic and should be marked with @pytest.mark.hs_llm_core
(or @pytest.mark.hs_llm_mat for provider matrix acceptance tests).
"""
mem = MemoryEngine(
db_url=pg0_db_url,
db_url=pg0_db_url, # Direct postgresql:// URL, not pg0://
memory_llm_provider=os.getenv("HINDSIGHT_API_LLM_PROVIDER", "groq"),
memory_llm_api_key=os.getenv("HINDSIGHT_API_LLM_API_KEY"),
memory_llm_model=os.getenv("HINDSIGHT_API_LLM_MODEL", "openai/gpt-oss-120b"),
@@ -470,9 +436,9 @@ async def memory_real_llm(pg0_db_url, embeddings, cross_encoder, query_analyzer)
cross_encoder=cross_encoder,
query_analyzer=query_analyzer,
pool_min_size=1,
pool_max_size=15,
run_migrations=False,
task_backend=SyncTaskBackend(),
pool_max_size=5,
run_migrations=False, # Migrations already run at session scope
task_backend=SyncTaskBackend(), # Execute tasks immediately in tests
)
await mem.initialize()
yield mem
@@ -500,7 +466,7 @@ async def memory_no_llm_verify(pg0_db_url, embeddings, cross_encoder, query_anal
cross_encoder=cross_encoder,
query_analyzer=query_analyzer,
pool_min_size=1,
pool_max_size=15,
pool_max_size=5,
run_migrations=False,
task_backend=SyncTaskBackend(),
skip_llm_verification=True, # Skip verification - will be overridden by test

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