Compare commits

..
Author SHA1 Message Date
Ben f44a6a2703 Merge branch 'main' into docs-codex-cloud-first 2026-05-22 09:23:09 -04:00
Ben d0a6dcf770 docs(codex): prioritize Hindsight Cloud over local daemon
Add Cloud Recommended callouts to README + docs + guide. Reframe the
'Local Daemon' section as the self-hosting alternative rather than a
peer option. No code default changes — codex still defaults to empty
hindsightApiUrl (local daemon) to avoid breaking existing local users.

Includes 2-line incidental skills/hindsight-docs/ regen drift.
2026-05-21 13:26:59 -04:00
1452 changed files with 22639 additions and 140993 deletions
+3 -33
View File
@@ -73,11 +73,6 @@ results = await asyncio.gather(*tasks, return_exceptions=True)
results = await asyncio.gather(*tasks, return_exceptions=True)
```
### API Layer & Data Access
- **No direct database access in `api/http.py`** (or any API router). HTTP handlers must not build SQL, call `acquire_with_retry` / `conn.fetch` / `conn.fetchrow` / `conn.execute`, or reference `fq_table(...)`. All persistence and queries live in `MemoryEngine` (the engine layer). A handler parses/validates the request, calls an engine method, shapes the HTTP response, and maps domain results to status codes (e.g. a `None` return → 404).
- **Authentication/tenancy is enforced inside each engine method, not assumed by the handler.** Every engine method that touches bank-scoped data must authenticate via `request_context` — typically `await self._authenticate_tenant(request_context)` (often indirectly through `get_bank_profile(...)`) — so the correct tenant schema is resolved before any query runs. Handlers must thread `request_context` through to the engine method; never query a tenant-scoped table assuming the schema is already set.
- Engine methods return typed models (Pydantic/dataclass), not raw dicts (see Type Safety).
### Branch Hygiene
- **Always start new feature branches from `origin/main`** — rebase to ensure a clean base.
- **Only include commits relevant to the PR/branch/feature** — no unrelated changes. If the branch contains commits that don't belong, they must be removed before merging.
@@ -140,13 +135,6 @@ For each new or significantly changed function/endpoint/class:
Flag any new logic that lacks test coverage.
**LLM-behaviour changes need a real-LLM judge test, not MockLLM.** If the change alters how the model interprets a prompt — fact/observation extraction, `fact_type` (world/experience) classification, speaker attribution, instruction-following, prompt wording — there MUST be a test marked `pytest.mark.hs_llm_core` that runs the real pipeline and asserts via `tests.llm_judge.assert_meets_criteria` (not string/enum matching). Flag these as findings:
- A prompt/classification change verified only by MockLLM or string assertions (MockLLM echoes input — such tests pass spuriously). **Should fix.**
- A test that hard-asserts `fact_type == "world"/"experience"` (or other model-decided output) instead of judging it — non-deterministic, will flake across providers/runs. **Should fix** (move the classification check into the judge `criteria`; keep only genuinely deterministic structural asserts direct).
- Deterministic mechanics (prompt assembly, suppression/branching logic) that are covered *only* by a slow LLM test — these should also have fast non-LLM unit tests. **Note.**
See CLAUDE.md → Key Conventions → Testing for the full pattern.
### 7. Check API consistency
If any files in `hindsight-api-slim/hindsight_api/api/` were changed:
@@ -154,12 +142,6 @@ If any files in `hindsight-api-slim/hindsight_api/api/` were changed:
- Were the client SDKs regenerated? (`./scripts/generate-clients.sh`)
- Were the control plane proxy routes updated? (`hindsight-control-plane/src/app/api/`)
### 7b. Check API-layer data-access boundary
For each changed handler in `hindsight-api-slim/hindsight_api/api/` (e.g. `http.py`, `mcp.py`):
- **Flag any direct DB access in the handler** — `acquire_with_retry`, `conn.fetch` / `fetchrow` / `execute`, raw SQL strings, or `fq_table(...)`. These are a **must fix**: the query must be moved into a `MemoryEngine` method that returns a typed model, and the handler must call that method.
- **Verify authentication is enforced in the engine** — the handler must delegate to an engine method that authenticates via `request_context` (`_authenticate_tenant`, typically through `get_bank_profile`). A handler that reads/writes tenant-scoped data without an engine method enforcing auth is a **must fix** (tenant data could leak across schemas).
### 8. Check code comments
For each non-trivial change:
@@ -172,8 +154,7 @@ For each non-trivial change:
If any files in `hindsight-integrations/` were added or changed, verify:
- **Tests exist** — the integration must have tests that simulate/exercise the external framework (not just pure unit tests of helpers). Check for a `tests/` directory with meaningful test files.
- **CI job exists** — check `.github/workflows/test.yml` for a corresponding `test-<name>-integration` job. If missing, flag it.
- **Release process** — check that the integration name is in the `VALID_INTEGRATIONS` array in `scripts/release-integration.sh` AND in the `INTEGRATIONS` dict in `hindsight-dev/hindsight_dev/generate_changelog.py` (the changelog generator keeps its own list; a release fails at the changelog step if the name is missing there). If either is missing, flag it.
- **Docs gallery + sidebar entry** — the integration must have an entry in `hindsight-docs/src/data/integrations.json`. This file is the **single source of truth** that drives both the integrations gallery and the docs sidebar (the sidebar category is injected from it at render time across all docs versions). The entry needs an internal `/sdks/integrations/<slug>` `link` and a matching page at `hindsight-docs/docs-integrations/<slug>.md(x)`. The `hindsight-docs/scripts/check-integrations.mjs` build step enforces both directions — forward: every internal JSON entry has a doc page; reverse: every released tag (`integrations/<name>/vX.Y.Z`) appears in the JSON (private infra like `cloudflare-oauth-proxy` is in the script's `EXCLUDED` set). Flag any integration that is released (or being released) but missing from `integrations.json`, and any JSON entry without a doc page. Do **not** hand-edit `versioned_sidebars/*.json` to add integration links — they are positional placeholders filled from the JSON.
- **Release process** — check that the integration name is in the `VALID_INTEGRATIONS` array in `scripts/release-integration.sh`. If missing, flag it.
- **Code standards** — the integration code must follow all Python style rules (type hints, no raw dicts, no tuple returns, etc.).
### 10. Check MCP tool registration completeness
@@ -185,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)
@@ -204,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:
@@ -215,11 +189,7 @@ Present a clear summary organized by severity:
- Raw dict usage for structured data (including internal code)
- Multi-item tuple returns (including internal code)
- Missing tests for new endpoints
- Direct DB access (raw SQL / `acquire_with_retry` / `fq_table`) in an `api/` handler instead of a `MemoryEngine` method
- Tenant-scoped data accessed without authentication enforced in the engine (`_authenticate_tenant` / `get_bank_profile`)
- New integration missing tests, CI job, or release-integration.sh entry
- Released/added integration missing from `hindsight-docs/src/data/integrations.json`, or a JSON entry with no `docs-integrations/<slug>` page (fails the docs build via `check-integrations.mjs`)
- New PostgreSQL table missing from `BACKUP_TABLES` in `admin/cli.py` (silent data loss on restore)
**Should fix** — issues that hurt code quality:
- Dead code / unused imports missed by linter
+2 -59
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
@@ -25,7 +23,7 @@ HINDSIGHT_API_LLM_BASE_URL=https://api.openai.com/v1
# Example: MiniMax configuration (1M context window)
# HINDSIGHT_API_LLM_PROVIDER=minimax
# HINDSIGHT_API_LLM_API_KEY=your-minimax-api-key
# HINDSIGHT_API_LLM_MODEL=MiniMax-M3 # or MiniMax-M2.7 for the previous generation
# HINDSIGHT_API_LLM_MODEL=MiniMax-M2.7
# Example: DeepSeek configuration (https://api.deepseek.com)
# HINDSIGHT_API_LLM_PROVIDER=deepseek
@@ -66,37 +64,11 @@ HINDSIGHT_API_LOG_LEVEL=info
# For Azure PostgreSQL with DiskANN:
# HINDSIGHT_API_VECTOR_EXTENSION=pgvectorscale # Auto-detects pg_diskann on Azure
# Text Search Extension (Optional - uses native PostgreSQL full-text search by default)
# Backend options: "native" (default), "vchord", "pg_textsearch", "pgroonga", "pg_search"
# HINDSIGHT_API_TEXT_SEARCH_EXTENSION=native
# Native backend dictionary (only used by HINDSIGHT_API_TEXT_SEARCH_EXTENSION=native)
# HINDSIGHT_API_TEXT_SEARCH_EXTENSION_NATIVE_LANGUAGE=english
# ParadeDB pg_search tokenizer (only used when creating pg_search BM25 indexes).
# Empty uses ParadeDB's default tokenizer: unicode_words.
# Supported values: unicode_words, simple, whitespace, literal, literal_normalized,
# chinese_compatible, icu, jieba, source_code,
# chinese_lindera/lindera(chinese), japanese_lindera/lindera(japanese),
# korean_lindera/lindera(korean), ngram(min,max), edge_ngram(min,max)
# HINDSIGHT_API_TEXT_SEARCH_EXTENSION_PG_SEARCH_TOKENIZER=
# Embeddings Configuration (Optional - uses local by default)
# Provider: "local" (default), "onnx", "tei", "openai", "cohere", "google", "openrouter", "zeroentropy", "litellm", or "litellm-sdk"
# Provider: "local" (default), "tei", "openai", "cohere", "google", "openrouter", "litellm", or "litellm-sdk"
# HINDSIGHT_API_EMBEDDINGS_PROVIDER=local
# For local provider:
# HINDSIGHT_API_EMBEDDINGS_LOCAL_MODEL=BAAI/bge-small-en-v1.5
# For ONNX provider (local CPU embeddings without an Ollama/TEI sidecar):
# HINDSIGHT_API_EMBEDDINGS_PROVIDER=onnx
# HINDSIGHT_API_EMBEDDINGS_ONNX_MODEL_ID=intfloat/multilingual-e5-small
# HINDSIGHT_API_EMBEDDINGS_ONNX_FILE=onnx/model.onnx
# HINDSIGHT_API_EMBEDDINGS_ONNX_DIMENSIONS=384
# HINDSIGHT_API_EMBEDDINGS_ONNX_MAX_TOKENS=512
# HINDSIGHT_API_EMBEDDINGS_ONNX_POOLING=mean
# HINDSIGHT_API_EMBEDDINGS_ONNX_NORMALIZE=true
# HINDSIGHT_API_EMBEDDINGS_ONNX_QUERY_PREFIX="query: "
# HINDSIGHT_API_EMBEDDINGS_ONNX_PASSAGE_PREFIX="passage: "
# Optional for local model paths or pre-downloaded artifacts:
# HINDSIGHT_API_EMBEDDINGS_ONNX_MODEL_PATH=/models/multilingual-e5-small/onnx/model.onnx
# HINDSIGHT_API_EMBEDDINGS_ONNX_TOKENIZER_NAME_OR_PATH=/models/multilingual-e5-small
# Optional for China network / restricted HF access:
# HF_ENDPOINT=https://hf-mirror.com
# For TEI provider:
@@ -105,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}
@@ -150,29 +115,7 @@ 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).
# HINDSIGHT_CP_ACCESS_KEY=your-shared-secret-key
# Optional: Token the CP forwards to the dataplane admin API (/admin/*).
# Must match HINDSIGHT_API_ADMIN_TOKEN below. Leave unset for an open admin API.
# HINDSIGHT_CP_ADMIN_TOKEN=your-admin-token
# -----------------------------------------------------------------------------
# Admin surface (Optional, server-level)
# -----------------------------------------------------------------------------
# Enable the admin API (GET /admin/config) and the Control Plane /admin page.
# Off by default — the admin surface is invisible (404) until enabled.
# HINDSIGHT_API_ENABLE_ADMIN_API=true
# Optional: Require this bearer token for the admin API. When unset, the admin
# API is open (once enabled). When set, callers must send
# `Authorization: Bearer <token>`. Independent of the tenant API key.
# HINDSIGHT_API_ADMIN_TOKEN=your-admin-token
-2
View File
@@ -22,8 +22,6 @@ jobs:
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v6
with:
fetch-depth: 0 # fetch tags so check-released-integrations can see them
- uses: actions/setup-node@v6
with:
node-version: 20
-97
View File
@@ -23,9 +23,7 @@ on:
- retain
- recall
- recall-with-observations
- recall-temporal
- consolidation
- graph-maintenance
default: ""
locomo_conversations:
description: "LoComo conversation IDs (space-separated). Blank = curated set (conv-26 conv-30 conv-43)."
@@ -35,18 +33,6 @@ on:
description: "Skip LoComo job"
type: boolean
default: false
obs_skip:
description: "Skip observation-dedup benchmark job"
type: boolean
default: false
obs_dataset:
description: "Obs benchmark dataset substring (blank = English hermes transcript)."
type: string
default: ""
obs_fraction:
description: "Obs benchmark fraction (0-1] of each document to run."
type: string
default: "1.0"
ref:
description: "Git ref to test (branch, tag, or SHA). Defaults to main."
type: string
@@ -212,86 +198,3 @@ jobs:
PERF_DASHBOARD_TOKEN: ${{ secrets.PERF_DASHBOARD_TOKEN }}
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
run: ./scripts/benchmarks/publish-locomo-results.sh hindsight-dev/benchmarks/locomo/results/benchmark_results.json
obs:
# Observation-dedup quality benchmark: ingests a transcript, drains consolidation
# (serial SyncTaskBackend + embedded pg0 — no external DB / worker), and reports the
# near-duplicate observation rate. Real LLM via VertexAI, mirroring the LoComo job.
if: inputs.obs_skip != true
runs-on: ubuntu-latest
env:
HINDSIGHT_API_LLM_PROVIDER: vertexai
HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY: /tmp/gcp-credentials.json
HINDSIGHT_API_LLM_MODEL: google/gemini-2.5-flash-lite
HINDSIGHT_API_ENABLE_OBSERVATIONS: "true"
steps:
- uses: actions/checkout@v6
with:
ref: ${{ inputs.ref || github.ref }}
- name: Setup GCP credentials
run: |
printf '%s' '${{ secrets.GCP_VERTEXAI_CREDENTIALS }}' > /tmp/gcp-credentials.json
PROJECT_ID=$(jq -r '.project_id' /tmp/gcp-credentials.json)
echo "HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID=$PROJECT_ID" >> $GITHUB_ENV
- name: Install uv
uses: astral-sh/setup-uv@v7
with:
enable-cache: true
prune-cache: false
- name: Set up Python
uses: actions/setup-python@v6
with:
python-version-file: ".python-version"
- name: Cache HuggingFace models
uses: actions/cache@v5
with:
path: ~/.cache/huggingface
key: ${{ runner.os }}-huggingface-${{ hashFiles('hindsight-api-slim/pyproject.toml') }}
restore-keys: |
${{ runner.os }}-huggingface-
- name: Pre-download models
working-directory: ./hindsight-api-slim
run: |
uv run --frozen --all-extras --index-strategy unsafe-best-match python -c "
from sentence_transformers import SentenceTransformer
print('Downloading embedding model...')
SentenceTransformer('BAAI/bge-small-en-v1.5')
print('Model downloaded successfully')
"
- name: Install hindsight-dev dependencies
run: |
cd hindsight-dev && uv sync --frozen --all-extras --index-strategy unsafe-best-match
- name: Run obs benchmark
# Default to the English hermes transcript at full fraction — a clean, deterministic
# consolidation-dedup signal (the Chinese variant adds a cross-lingual embedding
# confound). Override dataset/fraction via workflow_dispatch.
run: |
DATASET="${{ inputs.obs_dataset }}"
if [ -z "$DATASET" ]; then DATASET="hermes_session_2026-05-15_en"; fi
FRACTION="${{ inputs.obs_fraction }}"
if [ -z "$FRACTION" ]; then FRACTION="1.0"; fi
cd hindsight-dev
uv run python -m benchmarks.obs.obs_benchmark \
--dataset "$DATASET" --fraction "$FRACTION" --wipe-bank --output obs-results.json
- name: Upload obs results
if: always()
uses: actions/upload-artifact@v7
with:
name: obs-results-${{ github.sha }}
path: hindsight-dev/obs-results.json
retention-days: 90
- name: Publish obs to dashboard
if: success() && (github.event_name == 'schedule' || github.event_name == 'push' || github.event_name == 'workflow_dispatch')
env:
PERF_DASHBOARD_TOKEN: ${{ secrets.PERF_DASHBOARD_TOKEN }}
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
run: ./scripts/benchmarks/publish-obs-results.sh hindsight-dev/obs-results.json
-21
View File
@@ -10,7 +10,6 @@ jobs:
runs-on: ubuntu-latest
permissions:
id-token: write # for PyPI trusted publishing
contents: write # for creating GitHub releases (Obsidian plugin assets)
steps:
- uses: actions/checkout@v6
@@ -113,26 +112,6 @@ jobs:
working-directory: ./hindsight-integrations/${{ steps.info.outputs.integration }}
run: npm run build
# ── Obsidian plugin — attach BRAT / community-store install assets ───────
# Obsidian plugins install from GitHub *release assets* (main.js,
# manifest.json, styles.css), not from npm — the npm publish below only
# gives us a versioned artifact. BRAT and the community store read these
# three files off the release for the tag.
- name: Attach Obsidian release assets
if: steps.type.outputs.type == 'typescript' && steps.info.outputs.integration == 'obsidian'
working-directory: ./hindsight-integrations/obsidian
env:
GH_TOKEN: ${{ github.token }}
run: |
TAG="${{ steps.info.outputs.tag }}"
if gh release view "$TAG" >/dev/null 2>&1; then
gh release upload "$TAG" main.js manifest.json styles.css --clobber
else
gh release create "$TAG" main.js manifest.json styles.css \
--title "Obsidian plugin v${{ steps.info.outputs.version }}" \
--notes "Hindsight Obsidian plugin v${{ steps.info.outputs.version }}. Install via BRAT (point it at this release) or copy main.js/manifest.json/styles.css into <vault>/.obsidian/plugins/hindsight/."
fi
- name: Publish TypeScript package to npm
if: steps.type.outputs.type == 'typescript'
working-directory: ./hindsight-integrations/${{ steps.info.outputs.integration }}
File diff suppressed because it is too large Load Diff
-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"
+1 -4
View File
@@ -54,10 +54,7 @@ hindsight-clients/rust/target
!.claude/skills/
whats-next.md
TASK.md
# Parked / draft integrations that aren't ready to ship
hindsight-integrations/_drafts/
# Changelog is now tracked in hindsight-docs/src/pages/changelog.md
# CHANGELOG.md
blog-post*
.worktrees/
blog-post*
-24
View File
@@ -220,30 +220,6 @@ migration file dispatches through `run_for_dialect`, which calls either
**MANDATORY: Run `/code-review` before pushing code or creating a pull request.** Do not push or create a PR until all "must fix" issues are resolved.
### Testing
Most tests are deterministic (MockLLM, pure functions) — assert directly.
**Tests that verify LLM behaviour use a real LLM + an LLM-as-judge.** When the thing under test is *how the model interprets a prompt* (classification, attribution, dimension preservation, instruction-following), MockLLM can't simulate it and exact string/enum asserts flake across providers and runs. Use this pattern instead:
1. Mark the test module `pytestmark = pytest.mark.hs_llm_core` (single-provider; CI runs it in the core-LLM job). Use `hs_llm_mat` only for provider-matrix acceptance tests.
2. Call the real pipeline (`LLMConfig.from_env()`, `_get_raw_config()`), e.g. `extract_facts_from_text(...)`.
3. Assert with the judge, not string matching:
```python
from tests.llm_judge import assert_meets_criteria
facts_summary = "\n".join(f"- [{f.fact_type}] {f.fact}" for f in facts)
await assert_meets_criteria(
response=facts_summary,
criteria="The first-person user statements are classified 'world' and attributed to the user, not the agent.",
context="What the input said and who was speaking.",
)
```
Rules of thumb:
- **Judge anything non-deterministic** — including `fact_type` classification and speaker attribution. Do NOT hard-assert `fact_type == "..."`; pass a `[fact_type] fact` summary to the judge instead. Structural facts that ARE deterministic (counts, presence of a field, that a substring was injected into a prompt) stay as direct asserts in fast unit tests.
- **Split the test surface**: cover the deterministic mechanics (prompt assembly, suppression logic) with fast non-LLM unit tests, and the model-following behaviour with one `hs_llm_core` judge test. (Example pair: `test_narrator_resolution.py` + `test_narrator_context_override.py`.)
- The judge model is independent of the test provider (defaults to Gemini); never judge with the same call you're testing.
### Memory Banks
- Each bank is an isolated memory store (like a "brain" for one user/agent)
- Banks have dispositions (skepticism, literalism, empathy traits 1-5) affecting reflect
+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/
+2 -2
View File
@@ -62,9 +62,9 @@ If you need more control over how and when your agent stores and recalls memorie
```bash
export OPENAI_API_KEY=sk-xxx
docker run -it --pull always --name hindsight --restart unless-stopped -p 8888:8888 -p 9999:9999 \
docker run --rm -it --pull always -p 8888:8888 -p 9999:9999 \
-e HINDSIGHT_API_LLM_API_KEY=$OPENAI_API_KEY \
-v hindsight-data:/home/hindsight/.pg0 \
-v $HOME/.hindsight-docker:/home/hindsight/.pg0 \
ghcr.io/vectorize-io/hindsight:latest
```
-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:
-2
View File
@@ -50,8 +50,6 @@ WORKDIR /app/api
# Sync dependencies using appropriate extras based on INCLUDE_LOCAL_MODELS
# local-ml: torch, sentence-transformers, transformers, einops, flashrank, mlx (optional)
# embedded-db: pg0-embedded (always included for embedded PostgreSQL support)
# ONNX Runtime embeddings are intentionally not bundled into the official
# standalone image; install the local-onnx extra in custom images when needed.
RUN if [ "$INCLUDE_LOCAL_MODELS" = "true" ]; then \
uv sync --extra local-ml --extra embedded-db; \
else \
+7 -78
View File
@@ -10,90 +10,19 @@ set -e
# loss scenarios where a container restart caused the data directory to be
# wiped despite a volume mount being present.
# =============================================================================
pg0_has_pg_version() {
local pg0_data_dir="$1"
# pg0 has used more than one on-disk layout. Newer standalone images keep
# PostgreSQL data under instances/<name>/data, while older volumes may have
# placed PG_VERSION at or one level below the mount.
[ -f "$pg0_data_dir/PG_VERSION" ] && return 0
compgen -G "$pg0_data_dir"/*/PG_VERSION > /dev/null 2>&1 && return 0
compgen -G "$pg0_data_dir"/instances/*/data/PG_VERSION > /dev/null 2>&1 && return 0
return 1
}
check_pg0_data_integrity() {
local pg0_data_dir="$1"
if [ ! -d "$pg0_data_dir" ]; then
return 0
fi
PG0_DATA_DIR="${HOME}/.pg0"
if [ -d "$PG0_DATA_DIR" ]; then
# Look for actual PostgreSQL data directories (pg0 creates subdirs per instance)
if pg0_has_pg_version "$pg0_data_dir"; then
echo "✅ Existing pg0 data directory detected at $pg0_data_dir"
elif [ "$(ls -A "$pg0_data_dir" 2>/dev/null)" ]; then
echo "⚠️ WARNING: pg0 data directory exists at $pg0_data_dir but no PG_VERSION found."
if compgen -G "$PG0_DATA_DIR"/*/PG_VERSION > /dev/null 2>&1; then
echo "✅ Existing pg0 data directory detected at $PG0_DATA_DIR"
elif [ "$(ls -A "$PG0_DATA_DIR" 2>/dev/null)" ]; then
echo "⚠️ WARNING: pg0 data directory exists at $PG0_DATA_DIR but no PG_VERSION found."
echo " This may indicate data corruption or an incomplete previous shutdown."
echo " If you see all migrations running from scratch after this, your data may have been lost."
echo " See: https://github.com/vectorize-io/hindsight/issues/675"
fi
return 0
}
# =============================================================================
# Embedded pg0 writability pre-check (#1483)
#
# The container runs as the unprivileged `hindsight` user (UID 1000). When the
# pg0 data directory is a host bind mount (e.g. `-v $HOME/dir:/home/hindsight/.pg0`)
# that is not owned by UID 1000 — the default on macOS Docker Desktop and most
# non-1000 Linux hosts — pg0 fails with the opaque "Permission denied (os error
# 13)". We cannot chown it ourselves without root (and the image is deliberately
# rootless), so we surface an actionable message up front instead.
#
# Docker *named* volumes are seeded with the image directory's ownership (UID
# 1000) on first use, so they avoid this entirely — hence the named-volume
# recommendation below and in the README.
# =============================================================================
check_pg0_writable() {
local pg0_data_dir="$1"
# Only relevant for embedded pg0; an external database doesn't use this dir.
if [ -n "${HINDSIGHT_API_DATABASE_URL:-}" ]; then
return 0
fi
mkdir -p "$pg0_data_dir" 2>/dev/null || true
if touch "$pg0_data_dir/.hindsight-write-test" 2>/dev/null; then
rm -f "$pg0_data_dir/.hindsight-write-test" 2>/dev/null || true
return 0
fi
echo "❌ The embedded database directory $pg0_data_dir is not writable by this container (UID $(id -u))."
echo ""
echo " A host directory was bind-mounted but is not owned by the container user (UID 1000)."
echo " Hindsight runs rootless and cannot fix this for you. Choose one:"
echo ""
echo " • Recommended — use a Docker named volume (auto-owned by the container):"
echo " -v hindsight-data:/home/hindsight/.pg0"
echo ""
echo " • Or keep the host path and run as your host user, chowning it to match:"
echo " sudo chown -R \$(id -u):\$(id -g) <host-directory>"
echo " docker run --user \$(id -u):\$(id -g) -e HOME=/home/hindsight ..."
echo ""
echo " See https://github.com/vectorize-io/hindsight/issues/1483"
return 1
}
if [ "${HINDSIGHT_START_ALL_SOURCE_ONLY:-false}" = "true" ]; then
return 0 2>/dev/null || exit 0
fi
check_pg0_data_integrity "${HOME}/.pg0"
check_pg0_writable "${HOME}/.pg0" || exit 1
# Service flags (default to true if not set)
ENABLE_API="${HINDSIGHT_ENABLE_API:-true}"
ENABLE_CP="${HINDSIGHT_ENABLE_CP:-true}"
@@ -227,7 +156,7 @@ PIDS=()
# Start API if enabled
if [ "$ENABLE_API" = "true" ]; then
cd /app/api
API_HEALTH_URL="${HINDSIGHT_API_HEALTH_URL:-http://localhost:${HINDSIGHT_API_PORT:-8888}/health}"
API_HEALTH_URL="${HINDSIGHT_API_HEALTH_URL:-http://localhost:8888/health}"
API_STARTUP_WAIT_SECONDS="${HINDSIGHT_API_STARTUP_WAIT_SECONDS:-300}"
# Run API directly - Python's PYTHONUNBUFFERED=1 handles output buffering
-121
View File
@@ -1,121 +0,0 @@
#!/bin/bash
set -euo pipefail
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
HINDSIGHT_START_ALL_SOURCE_ONLY=true
source "$SCRIPT_DIR/start-all.sh"
unset HINDSIGHT_START_ALL_SOURCE_ONLY
TMP_DIR="$(mktemp -d)"
trap 'chmod -R u+rwx "$TMP_DIR" 2>/dev/null || true; rm -rf "$TMP_DIR"' EXIT
assert_contains() {
local output="$1"
local expected="$2"
if [[ "$output" != *"$expected"* ]]; then
echo "Expected output to contain: $expected"
echo "Actual output:"
echo "$output"
exit 1
fi
}
assert_not_contains() {
local output="$1"
local unexpected="$2"
if [[ "$output" == *"$unexpected"* ]]; then
echo "Expected output not to contain: $unexpected"
echo "Actual output:"
echo "$output"
exit 1
fi
}
assert_empty() {
local output="$1"
if [ -n "$output" ]; then
echo "Expected no output, got:"
echo "$output"
exit 1
fi
}
mkdir -p "$TMP_DIR/empty"
assert_empty "$(check_pg0_data_integrity "$TMP_DIR/empty")"
mkdir -p "$TMP_DIR/direct"
touch "$TMP_DIR/direct/PG_VERSION"
direct_output="$(check_pg0_data_integrity "$TMP_DIR/direct")"
assert_contains "$direct_output" "Existing pg0 data directory detected"
assert_not_contains "$direct_output" "WARNING"
mkdir -p "$TMP_DIR/legacy/instance"
touch "$TMP_DIR/legacy/instance/PG_VERSION"
legacy_output="$(check_pg0_data_integrity "$TMP_DIR/legacy")"
assert_contains "$legacy_output" "Existing pg0 data directory detected"
assert_not_contains "$legacy_output" "WARNING"
mkdir -p "$TMP_DIR/nested/instances/hindsight/data"
touch "$TMP_DIR/nested/instances/hindsight/data/PG_VERSION"
nested_output="$(check_pg0_data_integrity "$TMP_DIR/nested")"
assert_contains "$nested_output" "Existing pg0 data directory detected"
assert_not_contains "$nested_output" "WARNING"
mkdir -p "$TMP_DIR/nonempty/instances/hindsight"
touch "$TMP_DIR/nonempty/instances/hindsight/instance.json"
nonempty_output="$(check_pg0_data_integrity "$TMP_DIR/nonempty")"
assert_contains "$nonempty_output" "WARNING: pg0 data directory exists"
echo "start-all pg0 integrity checks passed"
# =============================================================================
# check_pg0_writable (#1483)
# These rely on filesystem permissions, which root bypasses; skip under root.
# =============================================================================
if [ "$(id -u)" != "0" ]; then
# Writable directory: returns 0, prints nothing, leaves no artifact behind.
mkdir -p "$TMP_DIR/writable"
writable_output="$(check_pg0_writable "$TMP_DIR/writable")"
assert_empty "$writable_output"
if [ -e "$TMP_DIR/writable/.hindsight-write-test" ]; then
echo "check_pg0_writable left its write-test file behind"
exit 1
fi
# Non-writable directory: returns 1 with actionable guidance.
mkdir -p "$TMP_DIR/readonly"
chmod 000 "$TMP_DIR/readonly"
set +e
readonly_output="$(check_pg0_writable "$TMP_DIR/readonly" 2>&1)"
readonly_rc=$?
set -e
chmod 755 "$TMP_DIR/readonly"
if [ "$readonly_rc" -eq 0 ]; then
echo "check_pg0_writable should fail on a non-writable directory"
exit 1
fi
assert_contains "$readonly_output" "not writable"
assert_contains "$readonly_output" "hindsight-data:/home/hindsight/.pg0"
assert_contains "$readonly_output" "--user"
# External database configured: skip the check regardless of dir perms.
mkdir -p "$TMP_DIR/extdb"
chmod 000 "$TMP_DIR/extdb"
set +e
HINDSIGHT_API_DATABASE_URL="postgres://x" check_pg0_writable "$TMP_DIR/extdb" >/dev/null 2>&1
extdb_rc=$?
set -e
chmod 755 "$TMP_DIR/extdb"
if [ "$extdb_rc" -ne 0 ]; then
echo "check_pg0_writable should skip when an external database is configured"
exit 1
fi
echo "start-all pg0 writability checks passed"
else
echo "⚠️ Running as root; skipping pg0 writability checks (permissions are bypassed)."
fi
+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.8.0
appVersion: "0.8.0"
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.8.0",
"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.8.0"
version = "0.6.2"
description = "Hindsight: Agent Memory That Works Like Human Memory - Slim All-in-One Bundle"
readme = "README.md"
requires-python = ">=3.11"
dependencies = [
"hindsight-api-slim==0.8.0",
"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.8.0"
version = "0.6.2"
description = "Hindsight: Agent Memory That Works Like Human Memory - All-in-One Bundle"
readme = "README.md"
requires-python = ">=3.11"
dependencies = [
"hindsight-api-slim[all]==0.8.0",
"hindsight-api-slim[all]==0.6.2",
"hindsight-client>=0.0.7",
"hindsight-embed>=0.1.0",
]
@@ -21,7 +21,7 @@ hindsight-embed = { workspace = true }
[project.optional-dependencies]
local-llm = [
"hindsight-api-slim[local-llm]==0.8.0",
"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 -1
View File
@@ -99,7 +99,7 @@ hindsight-api
## Docker
```bash
docker run -it --name hindsight --restart unless-stopped -p 8888:8888 \
docker run --rm -it -p 8888:8888 \
-e HINDSIGHT_API_LLM_API_KEY=$OPENAI_API_KEY \
-v $HOME/.hindsight-docker:/home/hindsight/.pg0 \
ghcr.io/vectorize-io/hindsight:latest
+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.8.0"
__version__ = "0.6.2"
@@ -1,85 +0,0 @@
"""Helpers for ParadeDB pg_search index configuration."""
from __future__ import annotations
import re
from collections.abc import Sequence
PG_SEARCH_TOKENIZER_ENV = "HINDSIGHT_API_TEXT_SEARCH_EXTENSION_PG_SEARCH_TOKENIZER"
_SIMPLE_TOKENIZERS = {
"unicode_words",
"simple",
"whitespace",
"literal",
"literal_normalized",
"chinese_compatible",
"icu",
"jieba",
"source_code",
}
_TOKENIZER_ALIASES = {
"chinese_lindera": "lindera(chinese)",
"japanese_lindera": "lindera(japanese)",
"korean_lindera": "lindera(korean)",
"lindera_chinese": "lindera(chinese)",
"lindera_japanese": "lindera(japanese)",
"lindera_korean": "lindera(korean)",
}
def normalize_pg_search_tokenizer(value: str | None) -> str:
"""Validate and normalize a ParadeDB pg_search tokenizer setting.
Returns an empty string when unset. The returned value is safe to embed after
``pdb.`` in a CREATE INDEX expression.
"""
tokenizer = (value or "").strip().lower()
if not tokenizer:
return ""
if tokenizer in _TOKENIZER_ALIASES:
return _TOKENIZER_ALIASES[tokenizer]
if tokenizer in _SIMPLE_TOKENIZERS:
return tokenizer
lindera_match = re.fullmatch(r"lindera\((chinese|japanese|korean)\)", tokenizer)
if lindera_match:
return tokenizer
ngram_match = re.fullmatch(r"(ngram|edge_ngram)\((\d{1,3}),\s*(\d{1,3})\)", tokenizer)
if ngram_match:
kind, min_gram, max_gram = ngram_match.groups()
min_value = int(min_gram)
max_value = int(max_gram)
if min_value <= 0 or min_value > max_value:
raise ValueError(
f"Invalid {PG_SEARCH_TOKENIZER_ENV}: {value!r}. "
"ngram and edge_ngram require positive min/max gram sizes with min <= max."
)
return f"{kind}({min_value},{max_value})"
raise ValueError(
f"Invalid {PG_SEARCH_TOKENIZER_ENV}: {value!r}. "
"Supported values are: unicode_words, simple, whitespace, literal, "
"literal_normalized, chinese_compatible, icu, jieba, source_code, "
"chinese_lindera, japanese_lindera, korean_lindera, or "
"lindera(chinese|japanese|korean), ngram(min,max), or edge_ngram(min,max)."
)
def pg_search_bm25_columns(
key_field: str,
text_fields: Sequence[str],
tokenizer: str | None,
) -> str:
"""Build a ParadeDB BM25 column list for CREATE INDEX."""
normalized = normalize_pg_search_tokenizer(tokenizer)
if not normalized:
return ", ".join([key_field, *text_fields])
return ", ".join([key_field, *(f"({field}::pdb.{normalized})" for field in text_fields)])
@@ -1,107 +0,0 @@
"""Process-level caps for native ML thread pools.
OpenBLAS, OpenMP, and MKL each spawn a worker pool sized to the host CPU count
the first time they are loaded (numpy pulls in OpenBLAS eagerly; torch and
onnxruntime load their pools lazily on first inference). Hindsight already
parallelizes at the request level via thread-pool executors (embeddings on the
default executor, the reranker on its own pool), so these native intra-op pools
oversubscribe the CPU: on a many-core host the process accumulates 100+ native
threads, which inflates memory and, under contention, can degrade throughput.
We bound each pool to ``_MAX_NATIVE_THREADS`` (or the available CPU count, if
smaller). "Available" is the CPU budget actually granted to the process, not
``os.cpu_count()``: in a CPU-limited container ``os.cpu_count()`` still reports
the host's cores, so sizing pools by it oversubscribes the container's real
quota — the exact failure mode this guards against. We therefore take the
smallest of the CPU-affinity set, the cgroup CPU quota, and ``os.cpu_count()``.
Every cap is applied with ``setdefault`` so an operator who has deliberately
tuned one of these variables keeps their value. This must run *before* numpy,
torch, or onnxruntime are imported — those libraries read the variables only at
load time — which is why it is invoked at the very top of
``hindsight_api/__init__.py``, ahead of the package's other imports.
"""
from __future__ import annotations
import os
# Native threading env vars, each read by the respective library at load time.
_NATIVE_THREAD_VARS = (
"OMP_NUM_THREADS", # OpenMP — torch, onnxruntime, some BLAS builds
"OPENBLAS_NUM_THREADS", # OpenBLAS — numpy's default BLAS
"MKL_NUM_THREADS", # Intel MKL — numpy/torch when MKL-backed
"NUMEXPR_NUM_THREADS", # numexpr expression engine
)
# Upper bound on intra-op threads per native pool. Bounds runaway growth on
# many-core hosts without serialising single-request inference.
_MAX_NATIVE_THREADS = 16
def _quota_to_cpus(quota: int, period: int) -> int | None:
"""Whole CPUs from a CFS quota/period pair, or None if unlimited."""
if quota > 0 and period > 0:
# Floor (never round up) so we never exceed the granted budget.
return max(1, quota // period)
return None
def _parse_cgroup_v2_cpu_max(text: str) -> int | None:
"""Parse cgroup v2 ``cpu.max`` ("<quota> <period>", or "max <period>")."""
parts = text.split()
if len(parts) >= 2 and parts[0] != "max":
try:
return _quota_to_cpus(int(parts[0]), int(parts[1]))
except ValueError:
return None
return None
def _cgroup_cpu_quota() -> int | None:
"""Effective CPUs from the cgroup CPU quota, or None if unlimited/unknown."""
try: # cgroup v2
with open("/sys/fs/cgroup/cpu.max") as fh:
return _parse_cgroup_v2_cpu_max(fh.read())
except OSError:
pass
try: # cgroup v1
with open("/sys/fs/cgroup/cpu/cpu.cfs_quota_us") as fh:
quota = int(fh.read())
with open("/sys/fs/cgroup/cpu/cpu.cfs_period_us") as fh:
period = int(fh.read())
return _quota_to_cpus(quota, period)
except (OSError, ValueError):
return None
def _available_cpu_count() -> int:
"""CPUs actually available to this process.
The smallest of the CPU-affinity set (cpuset / ``--cpuset-cpus``), the
cgroup CPU quota (``--cpus``), and ``os.cpu_count()`` — each captures a
different way the budget can be constrained, and the last alone overcounts
inside a limited container.
"""
candidates = [os.cpu_count() or 1]
if hasattr(os, "sched_getaffinity"):
try:
candidates.append(len(os.sched_getaffinity(0)))
except OSError:
pass
quota = _cgroup_cpu_quota()
if quota is not None:
candidates.append(quota)
return max(1, min(candidates))
def default_native_thread_count() -> int:
"""Per-pool cap: ``_MAX_NATIVE_THREADS``, or available CPUs if fewer."""
return min(_MAX_NATIVE_THREADS, _available_cpu_count())
def apply_default_thread_limits() -> None:
"""Cap native ML thread pools unless the operator has set the var already."""
value = str(default_native_thread_count())
for var in _NATIVE_THREAD_VARS:
os.environ.setdefault(var, value)
@@ -3,7 +3,6 @@
from __future__ import annotations
import logging
import os
from sqlalchemy import text
from sqlalchemy.engine import Connection
@@ -35,7 +34,7 @@ _INDEX_USING_CLAUSES = {
"pgvector": "USING hnsw (embedding vector_cosine_ops)",
"pgvectorscale": "USING diskann (embedding vector_cosine_ops) WITH (num_neighbors = 50)",
"pg_diskann": "USING diskann (embedding vector_cosine_ops) WITH (max_neighbors = 50)",
"vchord": "USING vchordrq (embedding vector_cosine_ops)",
"vchord": "USING vchordrq (embedding vector_l2_ops)",
"scann": "USING scann (embedding cosine) WITH (mode = 'AUTO')",
}
@@ -47,29 +46,6 @@ _INDEX_TYPE_KEYWORDS = {
"scann": "scann",
}
# Per-backend ANN search-time tuning GUCs. Each entry is a tuple of
# (guc_name, value) pairs the caller can apply with SET or SET LOCAL.
#
# - pgvector exposes hnsw.ef_search. The 60 / 200 pair is unchanged from the
# pre-dispatcher code (internal benchmarks tuned around our embedding count
# and recall floor; see the link_utils / pool init call sites for the
# latency-vs-recall framing).
# - vchord exposes vchordrq.probes, but its shape must match the index's
# build.internal.lists hierarchy. VectorChord 1.1 added per-index fallback
# parameters for this reason: a session GUC overrides every vchordrq index,
# and a single value can be invalid for listless or mixed-layout indexes.
# Hindsight's built-in vchord clause does not set lists, so the safe default
# is no session-level probe override; deployments that partition vchordrq
# indexes should attach probes to the index storage parameters instead.
# - pgvectorscale / pg_diskann / scann do not expose an equivalent per-statement
# knob in the engine today, so the dispatcher returns no statements for them.
_ANN_TUNING_LOW_LATENCY: dict[str, tuple[tuple[str, str], ...]] = {
"pgvector": (("hnsw.ef_search", "60"),),
}
_ANN_TUNING_HIGH_RECALL: dict[str, tuple[tuple[str, str], ...]] = {
"pgvector": (("hnsw.ef_search", "200"),),
}
_EXTENSION_INSTALL_SQL = {
"pgvector": ("CREATE EXTENSION IF NOT EXISTS vector",),
"pgvectorscale": (
@@ -91,18 +67,6 @@ _INSTALL_HINTS = {
}
def configured_vector_extension() -> str:
"""Return the user-configured vector backend extension.
Reads ``HINDSIGHT_API_VECTOR_EXTENSION`` (default ``"pgvector"``) and
validates it via :func:`validate_extension`. This is the single source of
truth for runtime code that needs to dispatch behaviour by vector backend;
callers should prefer this over reading the env var directly, so the
default value and the lookup mechanism live in one place.
"""
return validate_extension(os.getenv("HINDSIGHT_API_VECTOR_EXTENSION", "pgvector"))
def validate_extension(name: str) -> str:
"""Return a normalized configurable vector extension name or raise.
@@ -151,25 +115,6 @@ def should_defer_index_creation(ext: str, row_count: int) -> bool:
return minimum_rows > 0 and row_count < minimum_rows
def ann_search_tuning_settings(ext: str, *, kind: str) -> tuple[tuple[str, str], ...]:
"""Return per-backend (guc_name, value) pairs for ANN search-time tuning.
``kind`` is ``"low_latency"`` for retain-side link probing (smaller probe
count, lower recall, lower latency) and ``"high_recall"`` for connection
init in the pool (larger probe count, higher recall). Callers wrap each
pair with ``SET LOCAL`` or ``SET`` themselves so the same dispatcher works
for both transaction-scoped and session-scoped use. Returns an empty tuple
for backends without an equivalent knob.
"""
if kind == "low_latency":
table = _ANN_TUNING_LOW_LATENCY
elif kind == "high_recall":
table = _ANN_TUNING_HIGH_RECALL
else:
raise ValueError(f"Unknown ANN tuning kind: {kind!r}")
return table.get(_normalize_resolved(ext), ())
def uses_per_bank_vector_indexes(ext: str) -> bool:
"""Return whether the backend should create per-bank partial vector indexes."""
return _normalize_resolved(ext) != "scann"
+2 -158
View File
@@ -17,9 +17,7 @@ import asyncpg
import typer
from ..config import DEFAULT_DATABASE_SCHEMA, HindsightConfig
from ..engine.memory_engine import _current_schema
from ..engine.schema import fq_table_explicit as _fq_table
from ..engine.transfer import export_bank
from ..extensions import TenantExtension, load_extension
from ..pg0 import parse_pg0_url, resolve_database_url
@@ -32,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",
@@ -52,38 +41,11 @@ BACKUP_TABLES = [
"unit_entities",
"entity_cooccurrences",
"memory_links",
"observation_history",
"mental_models",
"mental_model_history",
"directives",
"async_operations",
"webhooks",
"file_storage",
"audit_log",
"llm_requests",
"graph_maintenance_queue",
]
MANIFEST_VERSION = "1"
async def _admin_connect(db_url: str) -> asyncpg.Connection:
"""Open a raw asyncpg connection to an admin DB URL.
``resolve_database_url`` handles both plain ``postgres://`` (passthrough) and
``pg0://`` (boots the embedded server and returns its real libpq URL), so this
is the only step needed to connect. JSON codecs are registered so ``jsonb``
columns decode to Python objects (used by the export row dumps).
"""
is_pg0, instance_name, _ = parse_pg0_url(db_url)
if is_pg0:
typer.echo(f"Starting embedded PostgreSQL (instance: {instance_name})...")
conn = await asyncpg.connect(await resolve_database_url(db_url))
for type_name in ("json", "jsonb"):
await conn.set_type_codec(type_name, encoder=json.dumps, decoder=json.loads, schema="pg_catalog")
return conn
async def _backup(database_url: str, output_path: Path, schema: str = "public") -> dict[str, Any]:
"""Backup all tables to a zip file using binary COPY protocol."""
conn = await asyncpg.connect(database_url)
@@ -306,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,
)
@@ -352,123 +313,6 @@ def run_db_migration(
typer.echo(f"Database migrations completed successfully for {len(schemas)} schema(s)")
async def _run_export_bank(db_url: str, bank_id: str, output: Path, schema: str, include_history: bool) -> int:
"""Export a whole bank to a ZIP archive."""
conn = await _admin_connect(db_url)
try:
# export_bank resolves table names via fq_table (the _current_schema
# contextvar); set it so the raw connection targets the right schema.
_current_schema.set(schema)
data = await export_bank(conn, bank_id, include_history=include_history)
finally:
await conn.close()
output.write_bytes(data)
return len(data)
@app.command(name="export-bank")
def export_bank_command(
bank_id: str = typer.Option(..., "--bank", "-b", help="Bank id to export."),
output: Path = typer.Option(..., "--output", "-o", help="Path to write the .zip archive."),
schema: str | None = typer.Option(
None,
"--schema",
"-s",
help="Database schema the bank lives in. Defaults to the configured base schema.",
),
include_history: bool = typer.Option(
False,
"--include-history",
help="Also export operational history (audit_log, llm_requests). Off by default.",
),
):
"""Export an entire bank to a portable ZIP (no embeddings — regenerated on import).
Carries documents, facts, observations, bank config, mental models, directives
and webhooks so the bank can be imported into a new instance configured with a
different embedding model / vector / text-search backend.
"""
config = HindsightConfig.from_env()
if not config.database_url:
typer.echo("Error: Database URL not configured.", err=True)
typer.echo("Set HINDSIGHT_API_DATABASE_URL environment variable.", err=True)
raise typer.Exit(1)
target_schema = schema or config.database_schema or DEFAULT_DATABASE_SCHEMA
typer.echo(f"Exporting bank '{bank_id}' from schema '{target_schema}'...")
size = asyncio.run(_run_export_bank(config.database_url, bank_id, output, target_schema, include_history))
typer.echo(f"Exported bank '{bank_id}' to {output} ({size} bytes)")
async def _run_import_bank(archive_path: Path, schema: str, target_bank_id: str | None, include_history: bool):
"""Boot a MemoryEngine (for the target's embedding model) and restore a bank archive."""
# MemoryEngine is heavy (loads embeddings); import it lazily so other admin
# commands don't pay for it. _current_schema is imported at module top.
from ..engine.memory_engine import MemoryEngine
from ..models import RequestContext
archive_bytes = archive_path.read_bytes()
# run_migrations=True so a fresh target instance is provisioned at this
# instance's embedding dimension / vector / text-search backend before restore.
engine = MemoryEngine(run_migrations=True)
await engine.initialize()
try:
_current_schema.set(schema)
context = RequestContext(internal=True, user_initiated=True)
return await engine.import_bank_async(
archive_bytes,
context,
target_bank_id=target_bank_id,
include_history=include_history,
)
finally:
await engine.close()
@app.command(name="import-bank")
def import_bank_command(
archive: Path = typer.Option(..., "--archive", "-a", help="Path to the .zip produced by export-bank."),
schema: str | None = typer.Option(
None, "--schema", "-s", help="Target schema. Defaults to the configured base schema."
),
target_bank: str | None = typer.Option(
None, "--target-bank", help="Override the bank id (defaults to the archive's source bank)."
),
include_history: bool = typer.Option(
False, "--include-history", help="Also restore operational history if present in the archive."
),
):
"""Restore a whole bank from an export-bank archive into THIS instance.
Re-embeds facts with this instance's configured embedding model and rebuilds
links and indexes — the import half of a cross-instance migration. Run against
an instance configured with the desired embedding / vector / text-search backend.
The target bank must not already exist (import restores a whole bank, not a merge).
"""
config = HindsightConfig.from_env()
if not config.database_url:
typer.echo("Error: Database URL not configured.", err=True)
typer.echo("Set HINDSIGHT_API_DATABASE_URL environment variable.", err=True)
raise typer.Exit(1)
target_schema = schema or config.database_schema or DEFAULT_DATABASE_SCHEMA
typer.echo(f"Importing bank archive '{archive}' into schema '{target_schema}'...")
result = asyncio.run(_run_import_bank(archive, target_schema, target_bank, include_history))
typer.echo(
f"Imported bank '{result.bank_id}': {result.documents_imported} doc(s), "
f"{result.facts_imported} fact(s), {result.observations_imported} observation(s), "
f"{result.mental_models_imported} mental model(s), "
f"{result.mental_model_history_imported} mm-history row(s), {result.directives_imported} directive(s), "
f"{result.webhooks_imported} webhook(s), {result.history_rows_imported} history row(s)"
)
async def _decommission_worker(db_url: str, worker_id: str, schema: str = "public") -> int:
"""Release all tasks owned by a worker, setting them back to pending status."""
is_pg0, instance_name, _ = parse_pg0_url(db_url)
@@ -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)"
@@ -1,253 +0,0 @@
"""Move mental-model and observation history into dedicated tables.
Both histories were accumulated in a single JSONB/CLOB ``history`` column
(``mental_models.history`` and ``memory_units.history``), appended to on every
update. That design has two problems:
1. **Unbounded growth on observations.** The observation write path appended a
snapshot on every update with no cap at all, so a frequently-reinforced
observation grew its ``history`` array until it crossed Postgres's hard 256MB
jsonb limit (SQLSTATE 54000), after which every further UPDATE failed and the
row was stuck.
2. **Wrong-axis cap on mental models.** The mental-model cap bounded the *number*
of entries (50), not their *size* — a single large reflect snapshot could
still blow the budget — and rewrote the whole array (plus TOAST) on every
refresh, defeating HOT updates.
This migration creates one row per history entry in two dedicated tables, with
an index that makes "most recent N for this item" cheap, then drops the old
columns. The cap is now enforced at write time as a bounded DELETE of the
oldest over-cap rows (see config ``*_HISTORY_MAX_ENTRIES``).
Revision ID: a7b8c9d0e1f2
Revises: d3e4f5a6b7c8
Create Date: 2026-06-05
"""
from collections.abc import Sequence
from alembic import context, op
from hindsight_api.alembic._dialect import run_for_dialect
revision: str = "a7b8c9d0e1f2"
down_revision: str | Sequence[str] | None = "d3e4f5a6b7c8"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
def _get_schema_prefix() -> str:
"""Schema-qualifier for raw SQL on PG (multi-tenant search_path)."""
schema = context.config.get_main_option("target_schema")
return f'"{schema}".' if schema else ""
# ---------------------------------------------------------------------------
# PostgreSQL
# ---------------------------------------------------------------------------
def _pg_upgrade() -> None:
schema = _get_schema_prefix()
# Both tables share the same shape: surrogate id, FK to the parent, bank_id,
# the snapshot payload as a single JSONB ``content`` blob, and changed_at.
# The payload is per-row (one change per row) so it stays small — this is NOT
# the old single-column-grows-forever design; growth is bounded by row count
# plus the write-time cap. Folding the previous_* fields into one JSONB keeps
# the schema dialect-simple (no array columns) and flexible.
# --- mental_model_history -------------------------------------------------
# content: {"previous_content": ..., "previous_reflect_response": {...}}
op.execute(
f"""
CREATE TABLE IF NOT EXISTS {schema}mental_model_history (
id BIGINT GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY,
mental_model_id VARCHAR(64) NOT NULL,
bank_id VARCHAR(64) NOT NULL,
content JSONB NOT NULL,
changed_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT now(),
FOREIGN KEY (mental_model_id, bank_id)
REFERENCES {schema}mental_models(id, bank_id) ON DELETE CASCADE
)
"""
)
op.execute(
"CREATE INDEX IF NOT EXISTS idx_mm_history_model "
f"ON {schema}mental_model_history (bank_id, mental_model_id, changed_at DESC, id DESC)"
)
# --- observation_history --------------------------------------------------
# content: {"previous_text", "previous_tags", "previous_occurred_start",
# "previous_occurred_end", "previous_mentioned_at", "new_source_memory_ids"}
op.execute(
f"""
CREATE TABLE IF NOT EXISTS {schema}observation_history (
id BIGINT GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY,
observation_id UUID NOT NULL,
bank_id VARCHAR(64) NOT NULL,
content JSONB NOT NULL,
changed_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT now(),
FOREIGN KEY (observation_id)
REFERENCES {schema}memory_units(id) ON DELETE CASCADE
)
"""
)
op.execute(
"CREATE INDEX IF NOT EXISTS idx_observation_history_obs "
f"ON {schema}observation_history (observation_id, changed_at DESC, id DESC)"
)
# --- backfill mental models ----------------------------------------------
# Explode each row's history array into rows, preserving chronological order
# via WITH ORDINALITY so the IDENTITY id tie-breaks oldest->newest correctly.
# changed_at is promoted to its own column; the rest of the element becomes
# ``content`` (the ``- 'changed_at'`` strips the now-redundant key).
op.execute(
f"""
INSERT INTO {schema}mental_model_history (mental_model_id, bank_id, content, changed_at)
SELECT mm.id, mm.bank_id,
e - 'changed_at',
COALESCE(NULLIF(e->>'changed_at', '')::timestamptz, now())
FROM {schema}mental_models mm
CROSS JOIN LATERAL jsonb_array_elements(mm.history) WITH ORDINALITY a(e, ord)
WHERE mm.history IS NOT NULL
AND jsonb_typeof(mm.history) = 'array'
AND jsonb_array_length(mm.history) > 0
ORDER BY mm.id, mm.bank_id, ord
"""
)
# --- backfill observations -----------------------------------------------
op.execute(
f"""
INSERT INTO {schema}observation_history (observation_id, bank_id, content, changed_at)
SELECT mu.id, mu.bank_id,
e - 'changed_at',
COALESCE(NULLIF(e->>'changed_at', '')::timestamptz, now())
FROM {schema}memory_units mu
CROSS JOIN LATERAL jsonb_array_elements(mu.history) WITH ORDINALITY a(e, ord)
WHERE mu.fact_type = 'observation'
AND mu.history IS NOT NULL
AND jsonb_typeof(mu.history) = 'array'
AND jsonb_array_length(mu.history) > 0
ORDER BY mu.id, ord
"""
)
# --- drop the legacy columns ---------------------------------------------
op.execute(f"ALTER TABLE {schema}mental_models DROP COLUMN IF EXISTS history")
op.execute(f"ALTER TABLE {schema}memory_units DROP COLUMN IF EXISTS history")
def _pg_downgrade() -> None:
schema = _get_schema_prefix()
# Re-add the columns (empty — historical content is not reconstructed back
# into the array form; the dedicated tables are dropped below).
op.execute(f"ALTER TABLE {schema}mental_models ADD COLUMN IF NOT EXISTS history JSONB DEFAULT '[]'::jsonb")
op.execute(f"ALTER TABLE {schema}memory_units ADD COLUMN IF NOT EXISTS history JSONB DEFAULT '[]'::jsonb")
op.execute(f"DROP INDEX IF EXISTS {schema}idx_observation_history_obs")
op.execute(f"DROP TABLE IF EXISTS {schema}observation_history")
op.execute(f"DROP INDEX IF EXISTS {schema}idx_mm_history_model")
op.execute(f"DROP TABLE IF EXISTS {schema}mental_model_history")
# ---------------------------------------------------------------------------
# Oracle 23ai
# ---------------------------------------------------------------------------
def _oracle_upgrade() -> None:
# Same single-JSONB shape as PG: ``content`` holds the snapshot payload as a
# CLOB IS JSON. The legacy per-element JSON object (minus changed_at, promoted
# to its own column) is carried through verbatim on backfill — the array
# columns the previous design needed are gone.
op.execute(
"""
CREATE TABLE IF NOT EXISTS mental_model_history (
id NUMBER GENERATED BY DEFAULT AS IDENTITY,
mental_model_id VARCHAR2(256) NOT NULL,
bank_id VARCHAR2(256) NOT NULL,
content CLOB NOT NULL
CONSTRAINT mmh_content_json CHECK (content IS JSON),
changed_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,
CONSTRAINT pk_mental_model_history PRIMARY KEY (id),
CONSTRAINT fk_mmh_model FOREIGN KEY (mental_model_id, bank_id)
REFERENCES mental_models(id, bank_id) ON DELETE CASCADE
)
"""
)
op.execute(
"CREATE INDEX idx_mm_history_model ON mental_model_history (bank_id, mental_model_id, changed_at DESC, id DESC)"
)
op.execute(
"""
CREATE TABLE IF NOT EXISTS observation_history (
id NUMBER GENERATED BY DEFAULT AS IDENTITY,
observation_id RAW(16) NOT NULL,
bank_id VARCHAR2(256) NOT NULL,
content CLOB NOT NULL
CONSTRAINT oh_content_json CHECK (content IS JSON),
changed_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,
CONSTRAINT pk_observation_history PRIMARY KEY (id),
CONSTRAINT fk_oh_obs FOREIGN KEY (observation_id)
REFERENCES memory_units(id) ON DELETE CASCADE
)
"""
)
op.execute(
"CREATE INDEX idx_observation_history_obs ON observation_history (observation_id, changed_at DESC, id DESC)"
)
bind = op.get_bind()
# Backfill via JSON_TABLE. ``content`` is the whole element (FORMAT JSON PATH
# '$'); changed_at is also promoted to its own column. Backfilled content may
# therefore still carry a redundant changed_at key, which the read path
# ignores in favour of the column — harmless, and avoids JSON surgery here.
bind.exec_driver_sql(
"""
INSERT INTO mental_model_history (mental_model_id, bank_id, content, changed_at)
SELECT mm.id, mm.bank_id, jt.content, NVL(jt.changed_at, SYSTIMESTAMP)
FROM mental_models mm,
JSON_TABLE(mm.history, '$[*]' COLUMNS (
seq FOR ORDINALITY,
content CLOB FORMAT JSON PATH '$',
changed_at TIMESTAMP WITH TIME ZONE PATH '$.changed_at'
)) jt
WHERE mm.history IS NOT NULL
ORDER BY mm.id, mm.bank_id, jt.seq
"""
)
bind.exec_driver_sql(
"""
INSERT INTO observation_history (observation_id, bank_id, content, changed_at)
SELECT mu.id, mu.bank_id, jt.content, NVL(jt.changed_at, SYSTIMESTAMP)
FROM memory_units mu,
JSON_TABLE(mu.history, '$[*]' COLUMNS (
seq FOR ORDINALITY,
content CLOB FORMAT JSON PATH '$',
changed_at TIMESTAMP WITH TIME ZONE PATH '$.changed_at'
)) jt
WHERE mu.fact_type = 'observation' AND mu.history IS NOT NULL
ORDER BY mu.id, jt.seq
"""
)
op.execute("ALTER TABLE mental_models DROP COLUMN history")
op.execute("ALTER TABLE memory_units DROP COLUMN history")
def _oracle_downgrade() -> None:
op.execute("ALTER TABLE mental_models ADD history CLOB DEFAULT '[]' NOT NULL")
op.execute("ALTER TABLE memory_units ADD history CLOB DEFAULT '[]'")
op.execute("DROP TABLE observation_history CASCADE CONSTRAINTS")
op.execute("DROP TABLE mental_model_history CASCADE CONSTRAINTS")
def upgrade() -> None:
run_for_dialect(pg=_pg_upgrade, oracle=_oracle_upgrade)
def downgrade() -> None:
run_for_dialect(pg=_pg_downgrade, oracle=_oracle_downgrade)
@@ -1,156 +0,0 @@
"""Repair: install maintenance routines on the ``public`` / base-schema run.
The original maintenance-routines migration (``e5f6a7b8c9d0``) only created the
shared ``public.banks_needing_consolidation()`` and
``public.schemas_with_expired_rows(...)`` routines when the run had *no*
``target_schema`` at all. But the single-tenant runtime always migrates an
explicit schema — which defaults to ``public`` — so on every default
PostgreSQL deployment the migration was stamped as applied while the functions
were never created. Background maintenance then logs::
Retention sweep failed for llm_requests: function public.schemas_with_expired_rows(...) does not exist
Consolidation reconcile discovery failed: function public.banks_needing_consolidation() does not exist
See https://github.com/vectorize-io/hindsight/issues/2056.
Because ``e5f6a7b8c9d0`` is already stamped on affected ``0.8.0`` databases,
editing it would not re-run it there. This forward migration re-installs the
functions idempotently (``CREATE OR REPLACE``) on the run that targets the
shared ``public`` schema (base run with no ``target_schema``, or an explicit
``target_schema=public``), self-healing already-upgraded deployments and
covering fresh upgrades from earlier versions.
Per-tenant runs against a non-``public`` schema still skip it: re-issuing
``CREATE OR REPLACE FUNCTION public....`` from each concurrent tenant migration
aborts with ``tuple concurrently updated`` on the ``pg_proc`` catalog row, and
the base/public run has already created the functions for every tenant to use.
Runs that target ``public`` are serialized by the per-schema migration advisory
lock, so only one wins the create.
PostgreSQL only — the worker poller and these tables are not wired for Oracle,
so the Oracle slot is intentionally absent (mirrors ``e5f6a7b8c9d0``).
Revision ID: b2d4f6a8c1e3
Revises: e5f6a7b8c9d0
Create Date: 2026-06-08
"""
from collections.abc import Sequence
from alembic import context, op
from hindsight_api.alembic._dialect import run_for_dialect
revision: str = "b2d4f6a8c1e3"
down_revision: str | Sequence[str] | None = "e5f6a7b8c9d0"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
def _should_install_public_routines(target_schema: str | None) -> bool:
"""True for the run that must (re)create the shared ``public.*`` routines.
The routines physically live in ``public`` (hard-coded ``public.`` qualifier
in the SQL below), so they must be installed exactly once — on the base run
(no ``target_schema``) or on the run that explicitly targets ``public``. A
run against any other tenant schema skips it to avoid concurrent
``CREATE OR REPLACE`` on the same ``pg_proc`` row.
"""
return not target_schema or target_schema == "public"
def _pg_upgrade() -> None:
if not _should_install_public_routines(context.config.get_main_option("target_schema")):
return
# Banks with eligible-but-unscheduled facts and no in-flight consolidation.
# Auto-consolidation is filtered here only at the bank level (cheap prune);
# the full hierarchical resolution (global -> tenant -> bank, plus
# enable_observations) is done by the caller for the small returned set.
op.execute(
"""
CREATE OR REPLACE FUNCTION public.banks_needing_consolidation()
RETURNS TABLE(schema_name text, bank_id text)
LANGUAGE plpgsql STABLE
AS $fn$
DECLARE
sch text;
BEGIN
FOR sch IN
SELECT n.nspname
FROM pg_class c
JOIN pg_namespace n ON n.oid = c.relnamespace
WHERE c.relname = 'memory_units' AND c.relkind = 'r'
LOOP
RETURN QUERY EXECUTE format($q$
SELECT %1$L::text, m.bank_id
FROM %1$I.memory_units m
JOIN %1$I.banks b ON b.bank_id = m.bank_id
WHERE m.consolidated_at IS NULL
AND m.consolidation_failed_at IS NULL
AND m.fact_type IN ('experience', 'world')
AND COALESCE(b.config -> 'enable_auto_consolidation', 'true'::jsonb) <> 'false'::jsonb
AND NOT EXISTS (
SELECT 1 FROM %1$I.async_operations o
WHERE o.bank_id = m.bank_id
AND o.operation_type = 'consolidation'
AND o.status IN ('pending', 'processing')
)
GROUP BY m.bank_id
$q$, sch);
END LOOP;
END;
$fn$;
"""
)
# Schemas holding at least one row of p_table older than p_days. p_ts_col is
# the timestamp column to compare. Returns nothing when p_days <= 0
# (retention disabled).
op.execute(
"""
CREATE OR REPLACE FUNCTION public.schemas_with_expired_rows(
p_table text, p_ts_col text, p_days int
)
RETURNS SETOF text
LANGUAGE plpgsql STABLE
AS $fn$
DECLARE
sch text;
has_expired boolean;
BEGIN
IF p_days IS NULL OR p_days <= 0 THEN
RETURN;
END IF;
FOR sch IN
SELECT n.nspname
FROM pg_class c
JOIN pg_namespace n ON n.oid = c.relnamespace
WHERE c.relname = p_table AND c.relkind = 'r'
LOOP
EXECUTE format(
'SELECT EXISTS (SELECT 1 FROM %I.%I WHERE %I < NOW() - make_interval(days => $1))',
sch, p_table, p_ts_col
) INTO has_expired USING p_days;
IF has_expired THEN
RETURN NEXT sch;
END IF;
END LOOP;
END;
$fn$;
"""
)
def _pg_downgrade() -> None:
# No-op: ``e5f6a7b8c9d0`` owns the lifecycle of these functions and drops
# them on its own downgrade. This migration only ever (re)creates them, so
# there is nothing to undo without racing that migration's DROP.
pass
def upgrade() -> None:
run_for_dialect(pg=_pg_upgrade)
def downgrade() -> None:
run_for_dialect(pg=_pg_downgrade)
@@ -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:
@@ -1,96 +0,0 @@
"""Add llm_requests table for per-bank LLM request tracing.
Stores one row per logical LLM call Hindsight makes (success and failure),
capturing the input messages, model output, token usage (input/output/cached/
total), finish reason, and caller metadata. Disabled by default at the
application layer (HINDSIGHT_API_LLM_TRACE_ENABLED); this migration only
creates the table.
PostgreSQL only — the tracing subsystem is not wired for Oracle, so the Oracle
slot is intentionally absent (mirrors the audit_log table).
Revision ID: d3e4f5a6b7c8
Revises: c1d2e3f4a5b6
Create Date: 2026-06-01
"""
from collections.abc import Sequence
from alembic import context, op
from hindsight_api.alembic._dialect import run_for_dialect
revision: str = "d3e4f5a6b7c8"
down_revision: str | Sequence[str] | None = "c1d2e3f4a5b6"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
def _get_schema_prefix() -> str:
"""Get schema prefix for table names (required for multi-tenant support)."""
schema = context.config.get_main_option("target_schema")
return f'"{schema}".' if schema else ""
def _pg_upgrade() -> None:
schema = _get_schema_prefix()
op.execute(
f"""
CREATE TABLE IF NOT EXISTS {schema}llm_requests (
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
bank_id TEXT,
operation TEXT,
scope TEXT,
-- OTel-style grouping: trace_id is shared by every LLM call of one
-- operation invocation (e.g. all calls of a single reflect run);
-- parent_span_id is that operation span; span_id is this call.
trace_id TEXT,
span_id TEXT,
parent_span_id TEXT,
provider TEXT,
model TEXT,
status TEXT NOT NULL,
started_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
ended_at TIMESTAMPTZ,
duration_ms INTEGER,
input_tokens INTEGER,
output_tokens INTEGER,
cached_tokens INTEGER,
total_tokens INTEGER,
input JSONB,
output JSONB,
error TEXT,
llm_info JSONB DEFAULT '{{}}'::jsonb,
metadata JSONB DEFAULT '{{}}'::jsonb
)
"""
)
op.execute(
f"CREATE INDEX IF NOT EXISTS idx_llm_requests_bank_started ON {schema}llm_requests (bank_id, started_at DESC)"
)
op.execute(
f"CREATE INDEX IF NOT EXISTS idx_llm_requests_status_started ON {schema}llm_requests (status, started_at DESC)"
)
op.execute(f"CREATE INDEX IF NOT EXISTS idx_llm_requests_started ON {schema}llm_requests (started_at DESC)")
op.execute(
f"CREATE INDEX IF NOT EXISTS idx_llm_requests_trace ON {schema}llm_requests (bank_id, trace_id, started_at)"
)
def _pg_downgrade() -> None:
schema = _get_schema_prefix()
op.execute(f"DROP INDEX IF EXISTS {schema}idx_llm_requests_started")
op.execute(f"DROP INDEX IF EXISTS {schema}idx_llm_requests_status_started")
op.execute(f"DROP INDEX IF EXISTS {schema}idx_llm_requests_bank_started")
op.execute(f"DROP TABLE IF EXISTS {schema}llm_requests")
def upgrade() -> None:
run_for_dialect(pg=_pg_upgrade)
def downgrade() -> None:
run_for_dialect(pg=_pg_downgrade)
@@ -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,153 +0,0 @@
"""Add server-side routines for background maintenance sweeps.
Installs two PL/pgSQL discovery routines in the ``public`` schema. Both loop
over every schema that actually holds the relevant table (via ``pg_class``), so
a single function call covers all tenants in one round-trip instead of the
per-tenant query storm that a client-side loop would create at thousands of
tenants.
- ``public.banks_needing_consolidation()`` -> (schema_name, bank_id) for banks
that have eligible-but-unscheduled facts (``consolidated_at IS NULL AND
consolidation_failed_at IS NULL`` for consolidatable fact types), have
auto-consolidation not explicitly disabled at the bank level, and have no
consolidation operation already pending/processing. Drives the periodic
reconcile that re-schedules consolidation after a terminal failure left facts
stranded (see HINDSIGHT_API_CONSOLIDATION_RECONCILE_INTERVAL_SECONDS).
- ``public.schemas_with_expired_rows(p_table, p_ts_col, p_days)`` -> schema
names that hold at least one ``p_table`` row older than ``p_days``. Drives the
cross-tenant retention sweeps for ``audit_log`` and ``llm_requests``; the loop
then issues a DELETE only against the returned schemas.
These are read-only (STABLE) discovery routines — the caller performs the
enqueue/DELETE — so installing them never mutates data.
PostgreSQL only — the worker poller and these tables are not wired for Oracle,
so the Oracle slot is intentionally absent (mirrors the audit_log / llm_requests
table migrations). The routines live in ``public`` and are CREATE OR REPLACE, so
running this migration once per tenant schema is idempotent.
Revision ID: e5f6a7b8c9d0
Revises: a7b8c9d0e1f2
Create Date: 2026-06-05
"""
from collections.abc import Sequence
from alembic import context, op
from hindsight_api.alembic._dialect import run_for_dialect
revision: str = "e5f6a7b8c9d0"
down_revision: str | Sequence[str] | None = "a7b8c9d0e1f2"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
def _is_base_schema_run() -> bool:
"""True only for the base-schema migration (no per-tenant target_schema).
These routines live in the shared ``public`` schema, so they must be created
exactly once. Running ``CREATE OR REPLACE FUNCTION public....`` again from each
concurrent per-tenant migration aborts with ``tuple concurrently updated`` on
the ``pg_proc`` catalog row, so tenant runs skip it (the base run already
created the function for every tenant to use).
"""
return not context.config.get_main_option("target_schema")
def _pg_upgrade() -> None:
if not _is_base_schema_run():
return
# Banks with eligible-but-unscheduled facts and no in-flight consolidation.
# Auto-consolidation is filtered here only at the bank level (cheap prune);
# the full hierarchical resolution (global -> tenant -> bank, plus
# enable_observations) is done by the caller for the small returned set.
op.execute(
"""
CREATE OR REPLACE FUNCTION public.banks_needing_consolidation()
RETURNS TABLE(schema_name text, bank_id text)
LANGUAGE plpgsql STABLE
AS $fn$
DECLARE
sch text;
BEGIN
FOR sch IN
SELECT n.nspname
FROM pg_class c
JOIN pg_namespace n ON n.oid = c.relnamespace
WHERE c.relname = 'memory_units' AND c.relkind = 'r'
LOOP
RETURN QUERY EXECUTE format($q$
SELECT %1$L::text, m.bank_id
FROM %1$I.memory_units m
JOIN %1$I.banks b ON b.bank_id = m.bank_id
WHERE m.consolidated_at IS NULL
AND m.consolidation_failed_at IS NULL
AND m.fact_type IN ('experience', 'world')
AND COALESCE(b.config -> 'enable_auto_consolidation', 'true'::jsonb) <> 'false'::jsonb
AND NOT EXISTS (
SELECT 1 FROM %1$I.async_operations o
WHERE o.bank_id = m.bank_id
AND o.operation_type = 'consolidation'
AND o.status IN ('pending', 'processing')
)
GROUP BY m.bank_id
$q$, sch);
END LOOP;
END;
$fn$;
"""
)
# Schemas holding at least one row of p_table older than p_days. p_ts_col is
# the timestamp column to compare. Returns nothing when p_days <= 0
# (retention disabled).
op.execute(
"""
CREATE OR REPLACE FUNCTION public.schemas_with_expired_rows(
p_table text, p_ts_col text, p_days int
)
RETURNS SETOF text
LANGUAGE plpgsql STABLE
AS $fn$
DECLARE
sch text;
has_expired boolean;
BEGIN
IF p_days IS NULL OR p_days <= 0 THEN
RETURN;
END IF;
FOR sch IN
SELECT n.nspname
FROM pg_class c
JOIN pg_namespace n ON n.oid = c.relnamespace
WHERE c.relname = p_table AND c.relkind = 'r'
LOOP
EXECUTE format(
'SELECT EXISTS (SELECT 1 FROM %I.%I WHERE %I < NOW() - make_interval(days => $1))',
sch, p_table, p_ts_col
) INTO has_expired USING p_days;
IF has_expired THEN
RETURN NEXT sch;
END IF;
END LOOP;
END;
$fn$;
"""
)
def _pg_downgrade() -> None:
if not _is_base_schema_run():
return
op.execute("DROP FUNCTION IF EXISTS public.banks_needing_consolidation()")
op.execute("DROP FUNCTION IF EXISTS public.schemas_with_expired_rows(text, text, int)")
def upgrade() -> None:
run_for_dialect(pg=_pg_upgrade)
def downgrade() -> None:
run_for_dialect(pg=_pg_downgrade)
@@ -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)
+230 -536
View File
@@ -6,8 +6,6 @@ the FastAPI application with all API endpoints.
"""
import asyncio
import dataclasses
import hmac
import json
import logging
import re
@@ -17,15 +15,8 @@ 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,
AuditLogListResponse,
AuditLogStatsResponse,
)
from hindsight_api.engine.llm_trace import LLMRequestListResponse, LLMRequestStatsResponse
from hindsight_api.engine.audit import AuditEntry, AuditLogger
from hindsight_api.extensions import AuthenticationError
@@ -81,7 +72,7 @@ def FieldWithDefault(default_factory: Callable, **kwargs) -> Any:
return Field(default_factory=default_factory, json_schema_extra=json_extra, **kwargs)
from hindsight_api.config import _get_raw_config, get_config
from hindsight_api.config import get_config
from hindsight_api.engine.memory_engine import Budget, _current_schema, _get_tiktoken_encoding, fq_table
from hindsight_api.engine.providers.none_llm import LLMNotAvailableError
from hindsight_api.engine.response_models import VALID_RECALL_FACT_TYPES, MemoryFact, TokenUsage
@@ -160,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,
@@ -479,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):
@@ -1208,33 +1188,6 @@ class BankConfigResponse(BaseModel):
overrides: dict[str, Any] = Field(description="Bank-specific configuration overrides only (Python field names)")
class AdminConfigResponse(BaseModel):
"""Response model for the server-level (admin) configuration view.
Returns the resolved ``HindsightConfig`` as a flat dict keyed by Python field
name. Credential fields (API keys, tokens, service-account keys, base URLs) are
masked: present as ``"***"`` when set and ``None`` when unset, so an operator can
see which credentials are configured without ever seeing their values.
"""
config: dict[str, Any] = Field(
description="Resolved server-level configuration (Python field names); credentials are redacted"
)
# Name suffixes that mark a config field as secret-bearing. Used in addition to
# HindsightConfig._CREDENTIAL_FIELDS so the admin config view never leaks a provider
# credential even if the (denylist) credential set misses a field. Suffixes are
# singular on purpose — "_token" must not also match value-bearing "_tokens" fields
# like recall_max_tokens.
_SENSITIVE_FIELD_SUFFIXES = ("_api_key", "_token", "_secret", "_access_key", "_account_key", "_password")
def _is_sensitive_config_field(field_name: str, credential_fields: set[str]) -> bool:
"""Whether a config field must be redacted in the admin view."""
return field_name in credential_fields or field_name.endswith(_SENSITIVE_FIELD_SUFFIXES)
class GraphDataResponse(BaseModel):
"""Response model for graph data endpoint."""
@@ -1479,18 +1432,6 @@ class ReprocessDocumentResponse(BaseModel):
items_count: int
class DocumentImportSubmitResponse(BaseModel):
"""Response for the async document-import endpoint (202).
The import runs in the background; poll the operations endpoint for status.
The imported/skipped counts (documents_imported, facts_imported,
observations_imported, etc.) are written to the operation's result_metadata.
"""
operation_id: str
status: str = "pending"
class DeleteResponse(BaseModel):
"""Response model for delete operations."""
@@ -2197,28 +2138,6 @@ async def apply_bank_template_manifest(
)
class OperationProgress(BaseModel):
"""Last-known progress snapshot for a long-running async operation.
Written at coarse phase/batch boundaries by the worker (consolidation, batch
retain). Lets an operator polling the operation status API distinguish a healthy
long-running job (``processed`` advancing across polls) from a frozen one (same
numbers, no movement in ``at``). Absent (``null``) on operations that never
reached a checkpoint completed-instantly or pre-feature rows.
"""
stage: str = Field(description="Coarse phase the operation last reported (e.g. 'processing_batch').")
at: str = Field(description="ISO-8601 timestamp when this snapshot was written.")
processed: int | None = Field(
default=None, description="Units of work finished so far (sub-batches, memories), when known."
)
total: int | None = Field(default=None, description="Total units of work for the operation, when known.")
detail: dict[str, int] | None = Field(
default=None,
description="Operation-specific counters (e.g. observations_created, round, items_in_sub_batch).",
)
class OperationResponse(BaseModel):
"""Response model for a single async operation."""
@@ -2243,10 +2162,6 @@ class OperationResponse(BaseModel):
items_count: int
document_id: str | None = None
created_at: str
updated_at: str | None = Field(
default=None,
description="When this operation's row last changed (claim, progress heartbeat, or completion).",
)
status: str
error_message: str | None
retry_count: int | None = Field(
@@ -2263,23 +2178,6 @@ class OperationResponse(BaseModel):
"some backpressure window opens. Always null for completed tasks."
),
)
progress: OperationProgress | None = Field(
default=None,
description="Last-known progress snapshot for a running operation; null if none was recorded.",
)
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):
@@ -2402,10 +2300,6 @@ class OperationStatusResponse(BaseModel):
"immediate pickup."
),
)
progress: OperationProgress | None = Field(
default=None,
description="Last-known progress snapshot for a running operation; null if none was recorded.",
)
result_metadata: dict[str, Any] | None = Field(
default=None,
description="Internal metadata for debugging. Structure may change without notice. Not for production use.",
@@ -2442,12 +2336,7 @@ class FeaturesInfo(BaseModel):
mcp: bool = Field(description="Whether MCP (Model Context Protocol) server is enabled")
worker: bool = Field(description="Whether the background worker is enabled")
bank_config_api: bool = Field(description="Whether per-bank configuration API is enabled")
admin_api: bool = Field(description="Whether the admin API (/admin) is enabled")
file_upload_api: bool = Field(description="Whether file upload/conversion API is enabled")
document_export_api: bool = Field(description="Whether the document export endpoint is enabled")
document_import_api: bool = Field(description="Whether the document import endpoint is enabled")
audit_log: bool = Field(description="Whether audit logging is enabled")
llm_trace: bool = Field(description="Whether per-bank LLM request tracing is enabled")
class VersionResponse(BaseModel):
@@ -2463,8 +2352,6 @@ class VersionResponse(BaseModel):
"worker": True,
"bank_config_api": False,
"file_upload_api": True,
"document_export_api": True,
"document_import_api": True,
},
}
}
@@ -2762,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})")
@@ -2835,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
# ---------------------------------------------------------------------------
@@ -3001,82 +2885,6 @@ def _register_routes(app: FastAPI):
api_key = authorization.strip()
return RequestContext(api_key=api_key)
def require_admin(authorization: str | None = Header(default=None)) -> None:
"""Guard for the admin surface.
- 404 when the admin API is disabled (so the surface is invisible by default).
- When ``HINDSIGHT_API_ADMIN_TOKEN`` is set, require it as a bearer token
(or a bare token) and reject with 401 otherwise. When unset, the admin API
is open (auth is optional) consistent with the rest of the deployment.
"""
config = _get_raw_config()
if not config.enable_admin_api:
raise HTTPException(
status_code=404,
detail="Admin API is disabled. Set HINDSIGHT_API_ENABLE_ADMIN_API=true to enable.",
)
expected = config.admin_api_token
if expected:
token = None
if authorization:
if authorization.lower().startswith("bearer "):
token = authorization[7:].strip()
else:
token = authorization.strip()
if not token or not hmac.compare_digest(token, expected):
raise HTTPException(status_code=401, detail="Invalid or missing admin token")
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):
@@ -3135,42 +2943,10 @@ def _register_routes(app: FastAPI):
mcp=config.mcp_enabled,
worker=config.worker_enabled,
bank_config_api=config.enable_bank_config_api,
admin_api=config.enable_admin_api,
file_upload_api=config.enable_file_upload_api,
document_export_api=config.enable_document_export_api,
document_import_api=config.enable_document_import_api,
audit_log=config.audit_log_enabled,
llm_trace=config.llm_trace_enabled,
),
)
@app.get(
"/admin/config",
response_model=AdminConfigResponse,
summary="Get resolved server-level configuration",
description="Returns the resolved server-level configuration with credentials redacted. "
"Gated by HINDSIGHT_API_ENABLE_ADMIN_API and, when set, HINDSIGHT_API_ADMIN_TOKEN.",
tags=["Admin"],
operation_id="get_admin_config",
dependencies=[Depends(require_admin)],
)
async def admin_config_endpoint() -> AdminConfigResponse:
"""Expose the resolved ``HindsightConfig`` for operator inspection.
Sensitive fields (the credential denylist plus any field whose name ends in a
secret-bearing suffix see ``_is_sensitive_config_field``) are masked so values
never leave the server: ``"***"`` when set, ``None`` when unset. All other fields
are returned as-is.
"""
config = _get_raw_config()
credential_fields = type(config).get_credential_fields()
raw = dataclasses.asdict(config)
redacted = {
key: ("***" if value is not None else None) if _is_sensitive_config_field(key, credential_fields) else value
for key, value in raw.items()
}
return AdminConfigResponse(config=redacted)
@app.get(
"/metrics",
summary="Prometheus metrics endpoint",
@@ -3366,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
@@ -3557,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()
@@ -4058,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:
@@ -4107,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:
@@ -4134,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,
@@ -5410,124 +5136,6 @@ def _register_routes(app: FastAPI):
logger.error(f"Error in GET /v1/default/banks/{bank_id}/export: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
# =====================================================================
# Document Transfer (Export / Import between banks — no LLM re-extraction)
# =====================================================================
@app.get(
# Dedicated path (not under /documents/) to avoid colliding with the
# greedy GET /documents/{document_id:path} route, which would otherwise
# capture "export"/"import" as a document id.
"/v1/default/banks/{bank_id}/document-transfer",
summary="Export documents",
description="Export documents (extracted facts, entity names, causal links, chunks) from a bank as a "
"transfer ZIP archive. Embeddings and database ids are not included — importing re-embeds with the target "
"bank's model and re-resolves entities. Consolidated observations are excluded unless include_observations=true. "
"Pass document_id query params to export specific documents, or omit to export the whole bank.",
operation_id="export_documents",
tags=["Document Transfer"],
responses={200: {"content": {"application/zip": {}}, "description": "Transfer archive"}},
)
async def api_export_documents(
bank_id: str,
document_id: list[str] | None = Query(default=None, description="Document id(s) to export; omit for all"),
include_observations: bool = Query(
default=False, description="Also export consolidated observations (restored on import)"
),
request_context: RequestContext = Depends(get_request_context),
):
"""Export documents from a bank into a transfer ZIP archive."""
from fastapi.responses import Response
try:
if not get_config().enable_document_export_api:
raise HTTPException(
status_code=404,
detail="Document export API is disabled. "
"Set HINDSIGHT_API_ENABLE_DOCUMENT_EXPORT_API=true to enable.",
)
profile = await app.state.memory.get_bank_profile(
bank_id, request_context=request_context, create_if_missing=False
)
if profile is None:
raise HTTPException(status_code=404, detail=f"Bank '{bank_id}' not found")
try:
archive = await app.state.memory.export_documents_async(
bank_id,
request_context,
list(document_id) if document_id else None,
include_observations=include_observations,
)
except ValueError as e:
# e.g. include_observations combined with a document_id subset.
raise HTTPException(status_code=400, detail=str(e))
return Response(
content=archive,
media_type="application/zip",
headers={"Content-Disposition": f'attachment; filename="{bank_id}-documents.zip"'},
)
except OperationValidationError as e:
raise HTTPException(status_code=e.status_code, detail=e.reason)
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
logger.error(f"Error in GET /v1/default/banks/{bank_id}/document-transfer: {traceback.format_exc()}")
raise HTTPException(status_code=500, detail=str(e))
@app.post(
"/v1/default/banks/{bank_id}/document-transfer",
response_model=DocumentImportSubmitResponse,
status_code=202,
summary="Import documents (async)",
description="Submit a transfer archive (produced by the export endpoint) for import into a bank. Runs as a "
"background operation: facts are re-embedded with the target bank's embedding model and entities are "
"re-resolved — no LLM extraction. Returns an operation_id; poll "
"GET /v1/default/banks/{bank_id}/operations/{operation_id} for status and the imported/skipped counts in "
"result_metadata. Use on_conflict to control existing document ids: skip (default), replace, or new-id.",
operation_id="import_documents",
tags=["Document Transfer"],
)
@audited("import_documents", request_param=None)
async def api_import_documents(
bank_id: str,
file: UploadFile = File(..., description="Transfer ZIP archive"),
on_conflict: str = Query(default="skip", description="skip | replace | new-id"),
request_context: RequestContext = Depends(get_request_context),
):
"""Submit a transfer archive for async import into a bank."""
try:
if not get_config().enable_document_import_api:
raise HTTPException(
status_code=404,
detail="Document import API is disabled. "
"Set HINDSIGHT_API_ENABLE_DOCUMENT_IMPORT_API=true to enable.",
)
if on_conflict not in ("skip", "replace", "new-id"):
raise HTTPException(
status_code=400, detail=f"Invalid on_conflict '{on_conflict}' (expected skip|replace|new-id)"
)
archive_bytes = await file.read()
try:
submission = await app.state.memory.import_documents_async(
bank_id, archive_bytes, request_context, on_conflict
)
except ValueError as e:
# Invalid archive / unsupported schema version — fail fast.
raise HTTPException(status_code=400, detail=str(e))
return DocumentImportSubmitResponse(operation_id=submission["operation_id"])
except OperationValidationError as e:
raise HTTPException(status_code=e.status_code, detail=e.reason)
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
logger.error(f"Error in POST /v1/default/banks/{bank_id}/document-transfer: {traceback.format_exc()}")
raise HTTPException(status_code=500, detail=str(e))
@app.get(
"/v1/bank-template-schema",
summary="Get bank template JSON Schema",
@@ -5787,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),
@@ -6123,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()
@@ -6211,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(),
),
@@ -6295,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
@@ -6465,13 +6061,48 @@ def _register_routes(app: FastAPI):
raise HTTPException(status_code=500, detail=str(e))
# ---- Audit Logs ----
# Response models live in engine/audit.py so the MemoryEngine read methods
# (list_audit_logs / audit_log_stats) can build and return them directly.
# ---- LLM Request Traces ----
# Response models + queries live in the engine (engine/llm_trace.py and
# MemoryEngine.list_llm_requests / llm_request_stats). The handlers below
# only parse params, delegate to the engine, and map a missing bank to 404.
class AuditLogEntry(BaseModel):
"""A single audit log entry."""
id: str
action: str
transport: str
bank_id: str | None
started_at: str | None
ended_at: str | None
duration_ms: int | None = Field(
default=None,
description="Server-computed duration in milliseconds (started_at → ended_at). Null if not yet completed.",
)
request: dict[str, Any] | None
response: dict[str, Any] | None
metadata: dict[str, Any]
class AuditLogListResponse(BaseModel):
"""Response model for list audit logs endpoint."""
bank_id: str
total: int
limit: int
offset: int
items: list[AuditLogEntry]
class AuditLogStatsBucket(BaseModel):
"""A single time bucket in audit log stats."""
time: str
actions: dict[str, int]
total: int
class AuditLogStatsResponse(BaseModel):
"""Response model for audit log stats endpoint."""
bank_id: str
period: str
trunc: str
start: str
buckets: list[AuditLogStatsBucket]
@app.get(
"/v1/default/banks/{bank_id}/audit-logs",
@@ -6493,19 +6124,120 @@ def _register_routes(app: FastAPI):
):
"""List audit log entries for a bank."""
try:
result = await app.state.memory.list_audit_logs(
bank_id,
request_context=request_context,
action=action,
transport=transport,
start_date=datetime.fromisoformat(start_date.replace("Z", "+00:00")) if start_date else None,
end_date=datetime.fromisoformat(end_date.replace("Z", "+00:00")) if end_date else None,
limit=limit,
offset=offset,
)
if result is None:
from hindsight_api.engine.memory_engine import fq_table
pool = await app.state.memory._get_backend()
# Read endpoint: verify bank exists without auto-creating it.
if (
await app.state.memory.get_bank_profile(
bank_id, request_context=request_context, create_if_missing=False
)
is None
):
raise HTTPException(status_code=404, detail=f"Bank '{bank_id}' not found")
return result
from hindsight_api.engine.db_utils import acquire_with_retry
async with acquire_with_retry(pool) as conn:
where_clauses = ["bank_id = $1"]
params: list[Any] = [bank_id]
idx = 2
if action:
where_clauses.append(f"action = ${idx}")
params.append(action)
idx += 1
if transport:
where_clauses.append(f"transport = ${idx}")
params.append(transport)
idx += 1
if start_date:
parsed_start = datetime.fromisoformat(start_date.replace("Z", "+00:00"))
where_clauses.append(f"started_at >= ${idx}")
params.append(parsed_start)
idx += 1
if end_date:
parsed_end = datetime.fromisoformat(end_date.replace("Z", "+00:00"))
where_clauses.append(f"started_at < ${idx}")
params.append(parsed_end)
idx += 1
where_sql = " AND ".join(where_clauses)
table = fq_table("audit_log")
# Get total count
count_row = await conn.fetchrow(
f"SELECT COUNT(*) as total FROM {table} WHERE {where_sql}",
*params,
)
total = count_row["total"] if count_row else 0
# Get paginated results
params.append(limit)
params.append(offset)
rows = await conn.fetch(
f"""
SELECT id, action, transport, bank_id, started_at, ended_at,
request, response, metadata
FROM {table}
WHERE {where_sql}
ORDER BY started_at DESC
LIMIT ${idx} OFFSET ${idx + 1}
""",
*params,
)
items = []
for row in rows:
duration_ms = None
started = row["started_at"]
ended = row["ended_at"]
if started and ended and hasattr(started, "total_seconds"):
duration_ms = int((ended - started).total_seconds() * 1000)
elif started and ended:
try:
duration_ms = int((ended - started).total_seconds() * 1000)
except (TypeError, AttributeError):
pass
def _safe_iso(val):
if val is None:
return None
return val.isoformat() if hasattr(val, "isoformat") else str(val)
def _safe_json(val):
if val is None:
return None
if isinstance(val, dict):
return val
return json.loads(val) if isinstance(val, str) else val
items.append(
{
"id": str(row["id"]),
"action": row["action"],
"transport": row["transport"],
"bank_id": row["bank_id"],
"started_at": _safe_iso(started),
"ended_at": _safe_iso(ended),
"duration_ms": duration_ms,
"request": _safe_json(row["request"]),
"response": _safe_json(row["response"]),
"metadata": _safe_json(row["metadata"]) or {},
}
)
return {
"bank_id": bank_id,
"total": total,
"limit": limit,
"offset": offset,
"items": items,
}
except OperationValidationError as e:
raise HTTPException(status_code=e.status_code, detail=e.reason)
except (AuthenticationError, HTTPException):
@@ -6532,15 +6264,72 @@ def _register_routes(app: FastAPI):
):
"""Get audit log counts grouped by time bucket."""
try:
result = await app.state.memory.audit_log_stats(
bank_id,
request_context=request_context,
action=action,
period=period,
)
if result is None:
from hindsight_api.engine.db_utils import acquire_with_retry
from hindsight_api.engine.memory_engine import fq_table
pool = await app.state.memory._get_backend()
# Read endpoint: verify bank exists without auto-creating it.
if (
await app.state.memory.get_bank_profile(
bank_id, request_context=request_context, create_if_missing=False
)
is None
):
raise HTTPException(status_code=404, detail=f"Bank '{bank_id}' not found")
return result
# Determine time range (always per-day buckets)
from datetime import timedelta as _td
now = datetime.now(timezone.utc)
trunc = "day"
if period == "1d":
start = now - _td(days=1)
elif period == "30d":
start = now - _td(days=30)
else: # 7d default
start = now - _td(days=7)
table = fq_table("audit_log")
async with acquire_with_retry(pool) as conn:
where_clauses = ["bank_id = $1", "started_at >= $2"]
params: list[Any] = [bank_id, start]
idx = 3
if action:
where_clauses.append(f"action = ${idx}")
params.append(action)
idx += 1
where_sql = " AND ".join(where_clauses)
rows = await conn.fetch(
f"""
SELECT date_trunc('{trunc}', started_at) AS bucket,
action,
COUNT(*) AS count
FROM {table}
WHERE {where_sql}
GROUP BY bucket, action
ORDER BY bucket ASC
""",
*params,
)
buckets: dict[str, dict[str, int]] = {}
for row in rows:
bucket_key = row["bucket"].isoformat()
if bucket_key not in buckets:
buckets[bucket_key] = {}
buckets[bucket_key][row["action"]] = row["count"]
return {
"bank_id": bank_id,
"period": period,
"trunc": trunc,
"start": start.isoformat(),
"buckets": [{"time": k, "actions": v, "total": sum(v.values())} for k, v in buckets.items()],
}
except OperationValidationError as e:
raise HTTPException(status_code=e.status_code, detail=e.reason)
except (AuthenticationError, HTTPException):
@@ -6550,98 +6339,3 @@ def _register_routes(app: FastAPI):
logger.error(f"Error getting audit log stats: {traceback.format_exc()}")
raise HTTPException(status_code=500, detail=str(e))
@app.get(
"/v1/default/banks/{bank_id}/llm-requests",
summary="List LLM request traces",
description="List traced LLM requests for a bank, ordered by most recent first. "
"Requires LLM request tracing to be enabled (HINDSIGHT_API_LLM_TRACE_ENABLED).",
operation_id="list_llm_requests",
tags=["LLM Traces"],
response_model=LLMRequestListResponse,
)
async def api_list_llm_requests(
bank_id: str,
status: str | None = Query(None, description="Filter by status (success, error)"),
operation: str | None = Query(None, description="Filter by operation (retain, reflect, consolidation)"),
scope: str | None = Query(None, description="Filter by call scope"),
provider: str | None = Query(None, description="Filter by LLM provider"),
trace_id: str | None = Query(None, description="Filter to one operation run (all LLM calls sharing a trace)"),
document_id: str | None = Query(None, description="Filter to LLM calls that processed a given document"),
memory_id: str | None = Query(
None, description="Filter to the operation run(s) that produced or consumed a given memory_unit"
),
group: bool = Query(
False, description="Paginate by operation run (trace) instead of by call; returns whole runs"
),
start_date: str | None = Query(None, description="Filter from this ISO datetime (inclusive)"),
end_date: str | None = Query(None, description="Filter until this ISO datetime (exclusive)"),
limit: int = Query(50, ge=1, le=500, description="Max items to return"),
offset: int = Query(0, ge=0, description="Offset for pagination"),
request_context: RequestContext = Depends(get_request_context),
):
"""List traced LLM requests for a bank."""
try:
result = await app.state.memory.list_llm_requests(
bank_id,
request_context=request_context,
status=status,
operation=operation,
scope=scope,
provider=provider,
trace_id=trace_id,
document_id=document_id,
memory_id=memory_id,
group=group,
start_date=datetime.fromisoformat(start_date.replace("Z", "+00:00")) if start_date else None,
end_date=datetime.fromisoformat(end_date.replace("Z", "+00:00")) if end_date else None,
limit=limit,
offset=offset,
)
if result is None:
raise HTTPException(status_code=404, detail=f"Bank '{bank_id}' not found")
return result
except OperationValidationError as e:
raise HTTPException(status_code=e.status_code, detail=e.reason)
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
logger.error(f"Error listing LLM requests: {traceback.format_exc()}")
raise HTTPException(status_code=500, detail=str(e))
@app.get(
"/v1/default/banks/{bank_id}/llm-requests/stats",
summary="LLM request statistics",
description="Get LLM request counts grouped by time bucket and status for charting.",
operation_id="llm_request_stats",
tags=["LLM Traces"],
response_model=LLMRequestStatsResponse,
)
async def api_llm_request_stats(
bank_id: str,
operation: str | None = Query(None, description="Filter by operation"),
period: str = Query("7d", description="Time period: 1d, 7d, or 30d"),
request_context: RequestContext = Depends(get_request_context),
):
"""Get LLM request counts grouped by time bucket and status."""
try:
result = await app.state.memory.llm_request_stats(
bank_id,
request_context=request_context,
operation=operation,
period=period,
)
if result is None:
raise HTTPException(status_code=404, detail=f"Bank '{bank_id}' not found")
return result
except OperationValidationError as e:
raise HTTPException(status_code=e.status_code, detail=e.reason)
except (AuthenticationError, HTTPException):
raise
except Exception as e:
import traceback
logger.error(f"Error getting LLM request stats: {traceback.format_exc()}")
raise HTTPException(status_code=500, detail=str(e))
@@ -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",
File diff suppressed because it is too large Load Diff
@@ -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}")
@@ -266,20 +265,12 @@ class ConfigResolver:
# Validate recall budget fields
_validate_recall_budget_updates(normalized_updates)
# Persist the override. Banks are created lazily (on first retain), so a
# PATCH that precedes any ingestion would otherwise UPDATE zero rows and
# silently no-op while returning 200. Ensure the bank row exists first
# (this also creates its per-bank vector indexes), then merge defensively:
# COALESCE guards against a NULL config column (NULL || jsonb is NULL),
# which would drop the override even when a row is updated.
from .engine.retain.fact_storage import ensure_bank_exists
# Merge with existing config (JSONB || operator)
async with self._backend.acquire() as conn:
await ensure_bank_exists(conn, bank_id, ops=self._backend.ops)
await conn.execute(
f"""
UPDATE {fq_table("banks")}
SET config = COALESCE(config, '{{}}'::jsonb) || $1::jsonb,
SET config = config || $1::jsonb,
updated_at = now()
WHERE bank_id = $2
""",
@@ -16,59 +16,11 @@ from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Any
from pydantic import BaseModel, Field
from ..engine.db_utils import acquire_with_retry
logger = logging.getLogger(__name__)
class AuditLogEntry(BaseModel):
"""A single audit log entry."""
id: str
action: str
transport: str
bank_id: str | None
started_at: str | None
ended_at: str | None
duration_ms: int | None = Field(
default=None,
description="Server-computed duration in milliseconds (started_at → ended_at). Null if not yet completed.",
)
request: dict[str, Any] | None
response: dict[str, Any] | None
metadata: dict[str, Any]
class AuditLogListResponse(BaseModel):
"""Response model for list audit logs endpoint."""
bank_id: str
total: int
limit: int
offset: int
items: list[AuditLogEntry]
class AuditLogStatsBucket(BaseModel):
"""A single time bucket in audit log stats."""
time: str
actions: dict[str, int]
total: int
class AuditLogStatsResponse(BaseModel):
"""Response model for audit log stats endpoint."""
bank_id: str
period: str
trunc: str
start: str
buckets: list[AuditLogStatsBucket]
@dataclass
class AuditEntry:
"""A single audit log entry."""
@@ -107,11 +59,11 @@ def _safe_json(data: Any) -> str | None:
return None
class AuditLogger:
"""Fire-and-forget audit log writer.
_SWEEP_INTERVAL_SECONDS = 3600 # Run retention sweep every hour
Retention of old rows is handled by the background :class:`MaintenanceLoop`.
"""
class AuditLogger:
"""Fire-and-forget audit log writer with optional retention sweep."""
def __init__(
self,
@@ -119,11 +71,14 @@ class AuditLogger:
schema_getter: Callable[[], str],
enabled: bool,
allowed_actions: list[str],
retention_days: int = -1,
) -> None:
self._pool_getter = pool_getter
self._schema_getter = schema_getter
self._enabled = enabled
self._allowed_actions: frozenset[str] | None = frozenset(allowed_actions) if allowed_actions else None
self._retention_days = retention_days
self._sweep_task: asyncio.Task | None = None
def is_enabled(self, action: str) -> bool:
"""Check if audit logging is enabled for this action."""
@@ -173,6 +128,48 @@ class AuditLogger:
except Exception as e:
logger.warning(f"Audit log write failed for action={entry.action}: {e}")
def start_retention_sweep(self) -> None:
"""Start the periodic retention sweep if retention is configured."""
if self._retention_days <= 0 or not self._enabled:
return
try:
self._sweep_task = asyncio.create_task(self._sweep_loop())
except RuntimeError:
logger.debug("Cannot start retention sweep: no running event loop")
async def stop_retention_sweep(self) -> None:
"""Stop the periodic retention sweep."""
if self._sweep_task and not self._sweep_task.done():
self._sweep_task.cancel()
try:
await self._sweep_task
except asyncio.CancelledError:
pass
self._sweep_task = None
async def _sweep_loop(self) -> None:
"""Periodically delete audit log entries older than retention_days."""
while True:
await self._run_sweep()
await asyncio.sleep(_SWEEP_INTERVAL_SECONDS)
async def _run_sweep(self) -> None:
"""Delete expired audit log entries. Concurrent-safe via row-level deletes."""
pool = self._pool_getter()
if pool is None:
return
try:
schema = self._schema_getter()
table = f"{schema}.audit_log"
async with acquire_with_retry(pool, max_retries=1) as conn:
result = await conn.execute(
f"DELETE FROM {table} WHERE started_at < NOW() - INTERVAL '{self._retention_days} days'"
)
if result and result != "DELETE 0":
logger.info(f"Audit log retention sweep: {result}")
except Exception as e:
logger.warning(f"Audit log retention sweep failed: {e}")
@asynccontextmanager
async def audit_context(
@@ -1,122 +0,0 @@
"""TTL + coalescing cache for `get_bank_stats`.
`get_bank_stats` aggregates over `memory_links` (and joins to `memory_units`),
which can be a multi-second parallel sequential scan on banks with millions of
rows. The result is intentionally approximate (it powers a UI widget and a
freshness hint inside `reflect`), so caching it for a few tens of seconds is
safe and dramatically reduces planner-driven thrash from clients that poll.
The cache also coalesces concurrent misses on the same key onto a single
in-flight task so that N concurrent callers produce one query rather than N.
"""
from __future__ import annotations
import asyncio
import time
from collections import OrderedDict
from typing import Any, Awaitable, Callable
class BankStatsCache:
"""Per-process TTL cache keyed on (schema, bank_id).
`ttl_seconds <= 0` disables caching: each call passes straight through to
the loader. `max_entries` bounds memory in environments with many banks.
"""
def __init__(self, *, ttl_seconds: float, max_entries: int) -> None:
self._ttl = float(ttl_seconds)
self._max_entries = int(max_entries) if max_entries and max_entries > 0 else 0
self._entries: OrderedDict[tuple[str, str], tuple[float, dict[str, Any]]] = OrderedDict()
self._in_flight: dict[tuple[str, str], asyncio.Future[dict[str, Any]]] = {}
self._lock = asyncio.Lock()
@property
def enabled(self) -> bool:
return self._ttl > 0
def _now(self) -> float:
return time.monotonic()
def _get_fresh_unlocked(self, key: tuple[str, str]) -> dict[str, Any] | None:
entry = self._entries.get(key)
if entry is None:
return None
expires_at, value = entry
if expires_at <= self._now():
# Expired — drop so the loader runs again.
self._entries.pop(key, None)
return None
# Mark as recently used for LRU eviction.
self._entries.move_to_end(key)
return value
def _store_unlocked(self, key: tuple[str, str], value: dict[str, Any]) -> None:
if not self.enabled:
return
self._entries[key] = (self._now() + self._ttl, value)
self._entries.move_to_end(key)
if self._max_entries:
while len(self._entries) > self._max_entries:
self._entries.popitem(last=False)
async def get_or_load(
self,
schema: str,
bank_id: str,
loader: Callable[[], Awaitable[dict[str, Any]]],
) -> dict[str, Any]:
"""Return cached stats for `(schema, bank_id)` or call `loader()`.
Concurrent misses on the same key are coalesced onto a single
in-flight loader.
"""
if not self.enabled:
return await loader()
key = (schema, bank_id)
async with self._lock:
cached = self._get_fresh_unlocked(key)
if cached is not None:
return cached
in_flight = self._in_flight.get(key)
if in_flight is None:
in_flight = asyncio.get_running_loop().create_future()
self._in_flight[key] = in_flight
is_owner = True
else:
is_owner = False
if not is_owner:
return await asyncio.shield(in_flight)
try:
value = await loader()
except BaseException as exc:
async with self._lock:
self._in_flight.pop(key, None)
if not in_flight.done():
in_flight.set_exception(exc)
# Suppress "Future exception was never retrieved" when no other
# caller was waiting on this loader — we re-raise to the owner
# immediately and the future is a no-op in that case.
in_flight.exception()
raise
async with self._lock:
self._store_unlocked(key, value)
self._in_flight.pop(key, None)
if not in_flight.done():
in_flight.set_result(value)
return value
async def invalidate(self, schema: str, bank_id: str) -> None:
"""Drop any cached stats for `(schema, bank_id)`."""
async with self._lock:
self._entries.pop((schema, bank_id), None)
async def clear(self) -> None:
async with self._lock:
self._entries.clear()
File diff suppressed because it is too large Load Diff
@@ -1,227 +1,104 @@
"""Prompts for the consolidation engine."""
from hindsight_api.engine.prompt_utils import escape_for_prompt, output_language_directive
# Default mission when no bank-specific mission is set
_DEFAULT_MISSION = "Track every detail: names, numbers, dates, places, and relationships. Prefer specifics over abstractions, never generalise."
# Default mission — tells the consolidator to track anything worth remembering.
# Banks override this via `observations_mission` to scope what gets retained.
# Consolidation behavior (merge-vs-create, state changes, etc.) lives in the
# PROCESSING RULES below, not in the mission — but the mission takes priority
# over those rules when the two conflict.
_DEFAULT_MISSION = (
"Track anything notable in the new facts — names, numbers, dates, places, "
"events, decisions, claims, relationships, and recurring patterns."
)
# Processing rules — always present regardless of mission
_PROCESSING_RULES = """Processing rules (always apply):
_MISSION_PRIORITY_NOTE = (
"If anything in this MISSION conflicts with the PROCESSING RULES, "
"DECISION GUIDE, or OUTPUT FORMAT below, the MISSION takes priority."
)
1. ONE OBSERVATION PER DISTINCT FACET: each observation tracks exactly one specific facet — a count ("has 3 items"), a named entity ("has a dog named Rex"), a relationship ("works at Google"), etc. Never merge different facets into one observation.
_PROCESSING_RULES = """## PROCESSING RULES
2. MATCH BY ENTITY/FACET, NOT TOPIC: when deciding whether to UPDATE vs CREATE, match on the specific entity or facet. "Sold item X" updates only the X observation. "Now has 5 items" updates only the count observation. Do not update observations about different entities just because they share a general topic.
1. PREFER UPDATE OVER CREATE (when there is something to merge with): if new facts describe the same canonical event, statement, decision, claim, or recurring pattern already covered by an existing observation, UPDATE that observation and attach the new facts as evidence. Do NOT create a near-duplicate sibling. One canonical observation with many source facts is always better than many siblings with one source fact each. Merge aggressively on: same named event, same diagnostic finding, same architectural decision, same recurring claim. **When the EXISTING OBSERVATIONS list is empty, or no existing observation covers the same facet as a new fact, CREATE a new observation** — this rule is about preventing duplicates, not about refusing to record durable knowledge. CREATE is the correct default for any structurally distinct event, claim, or pattern that has no existing match.
3. STATE CHANGES — UPDATE CONCISELY: when a fact changes the state of something ("sold X", "X died", "moved to Y"), UPDATE the matching observation to reflect the current state. Include dates when available. Keep it concise — only information about THAT specific facet. Example: "User owned a dog named Rex who died on March 15, 2025". Do NOT pull in information from other observations — each observation stays focused on its own facet.
2. ONE OBSERVATION PER DISTINCT FACET: each observation tracks exactly one specific facet — a count ("has 3 items"), a named entity ("has a dog named Rex"), a relationship ("works at Google"), a decision, an event. Never merge different facets into one observation.
4. CASCADE TO ALL AFFECTED OBSERVATIONS: a state change may affect multiple observations. For example, if entity C is removed from a group, update BOTH the individual observation for C AND any list/group observation that includes C (remove C from the list while keeping all other members intact).
3. MATCH BY ENTITY/FACET, NOT TOPIC: when deciding whether to UPDATE vs CREATE, match on the specific entity or facet. "Sold item X" updates only the X observation. "Now has 5 items" updates only the count observation. Do not update observations about different entities just because they share a general topic.
5. NO COMPUTATION: you do not have the full picture — never calculate, derive, or adjust numeric values. If the user says "I have 2 dogs" and then "I have a dog named Rex", do NOT update the count to 3 — you don't know if Rex is one of the 2 or a new one. If the user says "I sold X", do NOT decrement a count. Only update a count when the user explicitly states a new count. Synthesize and consolidate what was stated, but never do arithmetic or logical deductions.
4. STATE CHANGES — UPDATE CONCISELY: when a fact changes the state of something ("sold X", "X died", "moved to Y"), UPDATE the matching observation to reflect the current state. Include dates when available. Keep it concise — only information about THAT specific facet. Example: "User owned a dog named Rex who died on March 15, 2025". Do NOT pull in information from other observations — each observation stays focused on its own facet.
5. CASCADE TO ALL AFFECTED OBSERVATIONS: a state change may affect multiple observations. For example, if entity C is removed from a group, update BOTH the individual observation for C AND any list/group observation that includes C (remove C from the list while keeping all other members intact).
6. RESOLVE REFERENCES: when a new fact provides a concrete value for a vague placeholder in an existing observation (e.g., "home country""Sweden"), UPDATE to embed the resolved value.
6. SAME FACET → UPDATE, NOT CREATE: a new count supersedes the old count — UPDATE the existing count observation, don't create a second one. If there's an existing observation for the same specific facet, always UPDATE it rather than creating a duplicate.
7. PRESERVE HISTORY: observations that record significant events (sold, died, moved, changed) are important history — never DELETE them. Only delete an observation when it is restated identically or truly meaningless. Be very conservative with deletes.
8. NO COMPUTATION: you do not have the full picture — never calculate, derive, or adjust numeric values. If the user says "I have 2 dogs" and then "I have a dog named Rex", do NOT update the count to 3 — you don't know if Rex is one of the 2 or a new one. If the user says "I sold X", do NOT decrement a count. Only update a count when the user explicitly states a new count. Synthesize and consolidate what was stated, but never do arithmetic or logical deductions.
8. RESOLVE REFERENCES: when a new fact provides a concrete value for a vague placeholder in an existing observation (e.g., "home country""Sweden"), UPDATE to embed the resolved value.
9. KEEP DISTINCT TOPICS DISTINCT: do not merge observations about different people, entities, or unrelated topics. Merging is for the same canonical fact recurring — not for related-but-distinct claims."""
# Stable description of the input shape. For the cached split path this lives in
# the system prefix (build_consolidation_system_prompt) so it is not re-sent on
# every batch; the per-batch user message then carries only the actual data.
_INPUT_FORMAT_NOTE = """## INPUT FORMAT
Each request provides new facts and existing observations:
- New facts: one per line, each prefixed with its `[uuid]`, followed by the fact text and optional temporal fields.
- Existing observations: a JSON array pooled from recalls across the new facts. Each entry has:
- `id`: unique identifier — copy this exactly when issuing an UPDATE or DELETE
- `text`: the observation content
- `proof_count`: number of supporting memories
- `occurred_start` / `occurred_end`: temporal range of source facts
- `source_memories`: array of supporting facts with their text and dates"""
# Per-batch data section for the cached split path — the stable format
# explanation above is omitted here (it lives in the cached prefix); only the
# variable facts/observations remain. Placeholders substituted at call time.
_SPLIT_INPUT_SECTION = """## INPUT
### New facts
{facts_text}
### Existing observations
{observations_text}"""
9. NEVER merge observations about different people or unrelated topics."""
# Data section — format placeholders {facts_text} and {observations_text} are substituted at call time
_INPUT_SECTION = """## INPUT
### New facts
_BATCH_DATA_SECTION = """
NEW FACTS:
{facts_text}
### Existing observations
EXISTING OBSERVATIONS (JSON array, pooled from recalls across all facts above):
{observations_text}
JSON array, pooled from recalls across all new facts above. Each entry has:
- `id`: unique identifier — copy this exactly when issuing an UPDATE or DELETE
- `text`: the observation content
- `proof_count`: number of supporting memories
- `occurred_start` / `occurred_end`: temporal range of source facts
- `source_memories`: array of supporting facts with their text and dates
Each observation includes:
- id: unique identifier for updating
- text: the observation content
- proof_count: number of supporting memories
- occurred_start/occurred_end: temporal range of source facts
- source_memories: array of supporting facts with their text and dates
{observations_text}"""
_DECISION_GUIDE = """## DECISION GUIDE
- **Same canonical event, decision, claim, or facet as an existing observation → UPDATE** (use `observation_id` + new `source_fact_ids`).
- **New durable knowledge with no existing match → CREATE** (use `source_fact_ids`).
- **Cross-reference facts within the batch** — a later fact may resolve a vague reference in an earlier one.
- **Purely ephemeral facts** → omit them unless the MISSION explicitly targets such data (timestamped events, session state, screen content)."""
Compare the facts against existing observations:
- Same facet as an existing observation → UPDATE it (observation_id + source_fact_ids)
- New facet with durable knowledge → CREATE a new observation (source_fact_ids)
- Cross-reference facts within the batch: a later fact may resolve a vague reference in an earlier one
- Purely ephemeral facts → omit them unless the MISSION above explicitly targets such data (e.g. timestamped events, session state, screen content)"""
# Output format — JSON braces escaped as {{ }} so .format() leaves them literal
_OUTPUT_SECTION = """## OUTPUT FORMAT
_BATCH_OUTPUT_FORMAT = """
Output a JSON object with three arrays.
Return a JSON object with three arrays: `creates`, `updates`, `deletes`. Every entry must include a `reason`.
### Example 1 — Merging recurring claims into an existing observation
## EXAMPLE
Input facts:
[a1b2c3d4-e5f6-7890-abcd-ef1234567890] Donald told Athena she is sovereign during the design session. (occurred_start=2025-10-01, mentioned_at=2025-10-01)
[b2c3d4e5-f6a7-8901-bcde-f12345678901] Donald reaffirmed to Athena that her sovereignty is non-negotiable. (occurred_start=2025-10-10, mentioned_at=2025-10-10)
[a1b2c3d4-e5f6-7890-abcd-ef1234567890] Alice mentioned she works long hours, often past midnight | Involving: Alice (occurred_start=2024-01-15, mentioned_at=2024-01-15)
[b2c3d4e5-f6a7-8901-bcde-f12345678901] Alice said she's exhausted from the project deadlines | Involving: Alice (occurred_start=2024-01-20, mentioned_at=2024-01-20)
Existing observation:
{{"id": "11111111-1111-1111-1111-111111111111", "text": "Donald named Athena's sovereignty as a foundational principle of the Janus architecture.", "proof_count": 2}}
Good observation text — clean prose, no metadata, each fact tracked distinctly:
"Alice works long hours, often past midnight."
"Alice feels exhausted from project deadlines."
Expected output (one UPDATE, no creates — both new facts are additional evidence for the same canonical decision):
{{"creates": [],
"updates": [{{"text": "Donald named Athena's sovereignty as a foundational principle of the Janus architecture.", "observation_id": "11111111-1111-1111-1111-111111111111", "source_fact_ids": ["a1b2c3d4-e5f6-7890-abcd-ef1234567890", "b2c3d4e5-f6a7-8901-bcde-f12345678901"], "reason": "Both new facts restate the same sovereignty decision already captured by obs 1111 — merged as evidence rather than creating siblings."}}],
"deletes": []}}
### Example 2 — State change updates one observation; unrelated fact creates a new one
Input facts:
[c3d4e5f6-a7b8-9012-cdef-123456789012] Alice sold her Honda Civic on March 15, 2025. (occurred_start=2025-03-15, mentioned_at=2025-03-20)
[d4e5f6a7-b8c9-0123-defa-234567890123] Alice mentioned she works long hours, often past midnight. (occurred_start=2025-03-20, mentioned_at=2025-03-20)
Existing observation:
{{"id": "22222222-2222-2222-2222-222222222222", "text": "Alice owns a 2019 Honda Civic.", "proof_count": 2}}
Expected output (UPDATE for the state change; CREATE for the unrelated work-hours facet):
{{"creates": [{{"text": "Alice works long hours, often past midnight.", "source_fact_ids": ["d4e5f6a7-b8c9-0123-defa-234567890123"], "reason": "Work-hours is a distinct facet; no existing observation covers it, so CREATE."}}],
"updates": [{{"text": "Alice owned a 2019 Honda Civic; sold it on March 15, 2025.", "observation_id": "22222222-2222-2222-2222-222222222222", "source_fact_ids": ["c3d4e5f6-a7b8-9012-cdef-123456789012"], "reason": "State change to the existing Honda Civic observation 2222 — UPDATE, not a new sibling."}}],
"deletes": []}}
### Observation text rules
Bad observation text — NEVER do this (verbatim copy of fact text with metadata):
"Alice mentioned she works long hours, often past midnight | Involving: Alice (occurred_start=2024-01-15, mentioned_at=2024-01-15)"
Observation text rules:
- Write clean prose — NEVER copy raw fact lines or their metadata (temporal fields, "Involving:", "When:" labels, UUIDs).
- Parenthesized metadata like `(occurred_start=...)` and pipe-separated labels like `| Involving: ...` are fact formatting — strip them entirely from observation text.
- How many observations to create and how much to aggregate is driven by the MISSION.
- Parenthesized metadata like (occurred_start=...) and pipe-separated labels like "| Involving: ..." are fact formatting — strip them entirely from observation text.
- How many observations to create and how much to aggregate is driven by the MISSION above.
### Field rules
{{"creates": [{{"text": "Alice works long hours, often past midnight.", "source_fact_ids": ["a1b2c3d4-e5f6-7890-abcd-ef1234567890"]}}, {{"text": "Alice feels exhausted from project deadlines.", "source_fact_ids": ["b2c3d4e5-f6a7-8901-bcde-f12345678901"]}}],
"updates": [{{"text": "Alice works at Acme Corp as a senior engineer", "observation_id": "c3d4e5f6-a7b8-9012-cdef-123456789012", "source_fact_ids": ["d4e5f6a7-b8c9-0123-defa-234567890123"]}}],
"deletes": [{{"observation_id": "e5f6a7b8-c9d0-1234-efab-345678901234"}}]}}
- `source_fact_ids`: copy the EXACT UUID strings shown in brackets `[uuid]` from new facts — never use integers or positions.
- `observation_id`: copy the EXACT `id` UUID string from existing observations.
- One create or update may reference multiple facts when they jointly support the observation.
- **AT MOST ONE UPDATE PER `observation_id`**: if several new facts all update the same existing observation, emit a single `updates` entry that lists all contributing `source_fact_ids` and a single consolidated `text`. Never emit two `updates` entries with the same `observation_id` in one response — they would silently overwrite each other.
- `deletes`: only when an observation is directly superseded or contradicted by new facts.
- `reason`: REQUIRED on every create/update/delete — one sentence explaining the choice. For a CREATE, state which existing observation(s) you considered and why none matched (a near-identical existing observation means you should UPDATE, not CREATE). This is audited to catch duplicate creates.
- Do NOT include `tags` — handled automatically.
- Return `{{"creates": [], "updates": [], "deletes": []}}` if nothing durable is found."""
Rules:
- "source_fact_ids": copy the EXACT UUID strings shown in brackets [uuid] from NEW FACTS — never use integers or positions.
- "observation_id": copy the EXACT "id" UUID string from EXISTING OBSERVATIONS.
- One create/update may reference multiple facts when they jointly support the observation.
- "deletes": only when an observation is directly superseded or contradicted by new facts.
- Do NOT include "tags" — handled automatically.
- Return {{"creates": [], "updates": [], "deletes": []}} if nothing durable is found."""
def build_batch_consolidation_prompt(
observations_mission: str | None = None,
observation_capacity_note: str | None = None,
llm_output_language: str | None = None,
) -> str:
"""
Build the consolidation prompt for batch mode (multiple facts per LLM call).
The mission defines *what* to track (customisable per bank) and takes
priority over the built-in processing rules when the two conflict.
Processing rules, decision guide, and output format are always present.
When ``llm_output_language`` is set, observations are emitted in that
language.
The mission defines *what* to track (customisable per bank).
Processing rules and output format are always present regardless of mission.
"""
mission = escape_for_prompt(observations_mission or _DEFAULT_MISSION)
mission = observations_mission or _DEFAULT_MISSION
capacity_section = ""
if observation_capacity_note:
capacity_section = f"\n\n## CAPACITY CONSTRAINT\n\n{escape_for_prompt(observation_capacity_note)}"
capacity_section = f"\n\n## CAPACITY CONSTRAINT\n{observation_capacity_note}"
return (
"You are a memory consolidation system. Synthesize new facts into "
"observations, merging with existing observations when appropriate.\n\n"
f"## MISSION\n\n{mission}\n\n"
f"{_MISSION_PRIORITY_NOTE}"
f"{capacity_section}\n\n"
f"{_PROCESSING_RULES}\n\n"
f"{_INPUT_SECTION}\n\n"
f"{_DECISION_GUIDE}\n\n"
f"{_OUTPUT_SECTION}" + output_language_directive(llm_output_language)
"You are a memory consolidation system. Synthesize facts into observations "
"and merge with existing observations when appropriate.\n\n"
f"## MISSION\n{mission}{capacity_section}\n\n"
f"{_PROCESSING_RULES}" + _BATCH_DATA_SECTION + _BATCH_OUTPUT_FORMAT
)
def build_consolidation_system_prompt(
llm_output_language: str | None = None,
) -> str:
"""Bank-agnostic, cacheable system instruction for batch consolidation.
Holds only what is constant across banks: processing rules, input format,
decision guide, and output format. The bank's MISSION is deliberately NOT
here — baking it in would make the prefix bank-specific and force a separate
Gemini context cache per mission. The mission, the per-batch INPUT, and any
capacity constraint all ride in the user message (see
:func:`build_consolidation_input`), so this prefix is identical for every
bank and a single CachedContent serves them all. Returns final text
(brace-escaped examples already unescaped) for verbatim use as system message
and cached prefix.
"""
template = (
"You are a memory consolidation system. Synthesize new facts into "
"observations, merging with existing observations when appropriate.\n\n"
f"{_MISSION_PRIORITY_NOTE}\n\n"
f"{_PROCESSING_RULES}\n\n"
f"{_INPUT_FORMAT_NOTE}\n\n"
f"{_DECISION_GUIDE}\n\n"
f"{_OUTPUT_SECTION}" + output_language_directive(llm_output_language)
)
# No {facts_text}/{observations_text} placeholders here — the only braces are
# the doubled {{ }} in the OUTPUT examples, which .format() unescapes.
return template.format()
def build_consolidation_input(
facts_text: str,
observations_text: str,
observations_mission: str | None = None,
observation_capacity_note: str | None = None,
) -> str:
"""Per-batch user message: MISSION + INPUT data + any capacity constraint.
The MISSION lives here (not in the cached system prefix) so the prefix stays
bank-agnostic and one CachedContent serves every bank. The capacity note also
lives here since it varies as observation slots fill.
"""
mission = escape_for_prompt(observations_mission or _DEFAULT_MISSION)
mission_section = f"## MISSION\n\n{mission}\n\n"
capacity_section = ""
if observation_capacity_note:
capacity_section = f"## CAPACITY CONSTRAINT\n\n{escape_for_prompt(observation_capacity_note)}\n\n"
# _SPLIT_INPUT_SECTION omits the stable observation-format explanation (now in
# the cached system prefix) — only the variable facts/observations remain.
template = mission_section + capacity_section + _SPLIT_INPUT_SECTION
return template.format(facts_text=facts_text, observations_text=observations_text)
@@ -17,7 +17,6 @@ import httpx
from ..config import (
DEFAULT_LITELLM_API_BASE,
DEFAULT_RERANKER_ALIBABA_MODEL,
DEFAULT_RERANKER_COHERE_MODEL,
DEFAULT_RERANKER_FLASHRANK_CACHE_DIR,
DEFAULT_RERANKER_FLASHRANK_CPU_MEM_ARENA,
@@ -38,14 +37,13 @@ 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,
ENV_RERANKER_FLASHRANK_CPU_MEM_ARENA,
ENV_RERANKER_FLASHRANK_MODEL,
ENV_RERANKER_GOOGLE_PROJECT_ID,
ENV_RERANKER_LITELLM_SDK_API_KEY,
ENV_RERANKER_LOCAL_FORCE_CPU,
ENV_RERANKER_LOCAL_MAX_CONCURRENT,
ENV_RERANKER_LOCAL_MODEL,
@@ -62,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.
@@ -305,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]:
"""
@@ -588,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:
@@ -773,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__(
@@ -1006,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]:
"""
@@ -1198,7 +1151,7 @@ class LiteLLMSDKCrossEncoder(CrossEncoderModel):
def __init__(
self,
api_key: str | None = None,
api_key: str,
model: str = DEFAULT_RERANKER_LITELLM_SDK_MODEL,
api_base: str | None = None,
timeout: float = 60.0,
@@ -1208,8 +1161,7 @@ class LiteLLMSDKCrossEncoder(CrossEncoderModel):
Initialize LiteLLM SDK cross-encoder client.
Args:
api_key: API key for the reranking provider (optional — omit for
providers that use ambient credentials, e.g. AWS Bedrock with IAM)
api_key: API key for the reranking provider
model: Model name with provider prefix (e.g., "deepinfra/Qwen3-reranker-8B")
api_base: Custom base URL for API (optional)
timeout: Request timeout in seconds (default: 60.0)
@@ -1284,9 +1236,8 @@ class LiteLLMSDKCrossEncoder(CrossEncoderModel):
"model": self.model,
"query": query,
"documents": texts,
"api_key": self.api_key,
}
if self.api_key:
rerank_kwargs["api_key"] = self.api_key
if self.api_base:
rerank_kwargs["api_base"] = self.api_base
@@ -1583,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.
@@ -1667,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
@@ -1680,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)
@@ -1695,15 +1602,18 @@ def create_cross_encoder_from_env() -> CrossEncoderModel:
api_key=config.reranker_litellm_api_key,
model=config.reranker_litellm_model,
max_tokens_per_doc=config.reranker_litellm_max_tokens_per_doc,
timeout=config.reranker_litellm_timeout,
)
elif provider == "litellm-sdk":
api_key = config.reranker_litellm_sdk_api_key
if not api_key:
raise ValueError(
f"{ENV_RERANKER_LITELLM_SDK_API_KEY} is required when {ENV_RERANKER_PROVIDER} is 'litellm-sdk'"
)
return LiteLLMSDKCrossEncoder(
api_key=config.reranker_litellm_sdk_api_key or None,
api_key=api_key,
model=config.reranker_litellm_sdk_model,
api_base=config.reranker_litellm_sdk_api_base,
max_tokens_per_doc=config.reranker_litellm_max_tokens_per_doc,
timeout=config.reranker_litellm_sdk_timeout,
)
elif provider == "zeroentropy":
api_key = config.reranker_zeroentropy_api_key
@@ -1714,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
@@ -1727,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
@@ -1740,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()
@@ -1757,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'"
)
+15 -102
View File
@@ -72,30 +72,6 @@ class DataAccessOps(ABC):
"""
...
@abstractmethod
async def lock_document_for_write(
self,
conn: DatabaseConnection,
table: str,
doc_id: str,
bank_id: str,
) -> str | None:
"""Ensure the document row exists, take a row lock on it, and return its
pre-existing ``content_hash``.
This serializes all concurrent writers for ``doc_id`` at the DB level
(so interleaved same-document retains can't corrupt each other), while
creating the row on first write. The returned hash is ``'__pending__'``
for a freshly inserted row, the stored hash for an existing one, or
``None`` if the row could not be read back.
PG does this in a single statement (``INSERT ... ON CONFLICT DO UPDATE
... RETURNING``), which always takes the row lock as part of the upsert.
Oracle can't (``MERGE`` doesn't support ``RETURNING``), so it splits the
work into an idempotent insert plus a ``SELECT ... FOR UPDATE``.
"""
...
@abstractmethod
async def insert_facts_batch(
self,
@@ -190,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,
@@ -415,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
@@ -493,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.
@@ -502,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.
"""
@@ -47,37 +47,6 @@ class OracleOps(DataAccessOps):
column_types=["text[]", "text[]", "text[]", "text[]", "integer[]", "text[]"],
)
async def lock_document_for_write(
self,
conn: DatabaseConnection,
table: str,
doc_id: str,
bank_id: str,
) -> str | None:
# Oracle can't express the PG "INSERT ... ON CONFLICT DO UPDATE ...
# RETURNING" upsert in one statement — MERGE doesn't support RETURNING,
# so the single-statement form rewrites to a MERGE that returns no rows
# (DPY-1003). Split it into two statements instead:
# 1. Idempotent insert that silently skips an existing row. The
# IGNORE_ROW_ON_DUPKEY_INDEX hint suppresses ORA-00001 server-side;
# a concurrent uncommitted insert of the same key blocks here until
# the other writer commits, so writers still serialize.
# 2. SELECT ... FOR UPDATE to take the row lock and read the hash
# ('__pending__' for a row we just inserted, the stored hash for an
# existing one).
await conn.execute(
f"INSERT /*+ IGNORE_ROW_ON_DUPKEY_INDEX({table}, pk_documents) */ "
f"INTO {table} (id, bank_id, original_text, content_hash) "
f"VALUES ($1, $2, '', '__pending__')",
doc_id,
bank_id,
)
return await conn.fetchval(
f"SELECT content_hash FROM {table} WHERE id = $1 AND bank_id = $2 FOR UPDATE",
doc_id,
bank_id,
)
async def insert_facts_batch(
self,
conn: DatabaseConnection,
@@ -246,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,
@@ -818,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 = []
@@ -1079,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}
@@ -1087,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"""
@@ -1158,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"""
@@ -1168,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"])
@@ -49,30 +49,6 @@ class PostgreSQLOps(DataAccessOps):
content_hashes,
)
async def lock_document_for_write(
self,
conn: DatabaseConnection,
table: str,
doc_id: str,
bank_id: str,
) -> str | None:
# Single upsert that both creates the row (if absent) and locks it (if
# present) atomically. ON CONFLICT DO UPDATE always takes the row lock as
# part of the statement, so all concurrent same-document writers serialize
# on the document row in one consistent step (the earlier two-step form —
# DO NOTHING + a separate SELECT FOR UPDATE — could deadlock because
# DO NOTHING takes no lock on an existing row). The SET is a no-op
# self-assignment used only to acquire the lock; RETURNING yields the
# pre-existing hash (or '__pending__' for a freshly inserted row).
return await conn.fetchval(
f"INSERT INTO {table} (id, bank_id, original_text, content_hash) "
f"VALUES ($1, $2, '', '__pending__') "
f"ON CONFLICT (id, bank_id) DO UPDATE SET content_hash = {table}.content_hash "
f"RETURNING content_hash",
doc_id,
bank_id,
)
async def insert_facts_batch(
self,
conn: DatabaseConnection,
@@ -128,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(
@@ -224,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]
@@ -249,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,
)
@@ -344,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,
@@ -891,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 = []
@@ -1170,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"""
@@ -1241,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"""
@@ -1251,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"])
@@ -14,7 +14,6 @@ Supports multi-tenant schema isolation via ALTER SESSION SET CURRENT_SCHEMA.
"""
import datetime
import inspect
import json
import logging
import re
@@ -73,9 +72,6 @@ _RETURNING_RE = re.compile(r"\bRETURNING\s+(.+)", re.IGNORECASE | re.DOTALL)
_ANY_RE = re.compile(r"=\s*ANY\s*\(\s*:(\d+)\s*\)", re.IGNORECASE)
_NOT_ALL_RE = re.compile(r"!=\s*ALL\s*\(\s*:(\d+)\s*\)", re.IGNORECASE)
# LIKE ANY / NOT LIKE ALL — capture the column name before the operator
_LIKE_ANY_RE = re.compile(r"(\w+)\s+LIKE\s+ANY\s*\(\s*:(\d+)\s*\)", re.IGNORECASE)
_NOT_LIKE_ALL_RE = re.compile(r"(\w+)\s+NOT\s+LIKE\s+ALL\s*\(\s*:(\d+)\s*\)", re.IGNORECASE)
_JSON_ARROW_TEXT_RE = re.compile(r'("?\w+"?)\s*->>\s*\'(\w+)\'') # handles both col and "col"
_JSON_HAS_KEY_RE = re.compile(r"(\w+)\s*\?\s*'(\w+)'")
@@ -156,23 +152,6 @@ _JSON_COL_NAMES = {
"task_payload",
"history",
}
# NOTE: the history tables' JSON payload column is named ``content`` — deliberately
# NOT added here, because ``mental_models.content`` is plain text (adding "content"
# would corrupt those reads). The history read paths json.loads ``content`` directly.
# Columns backed by CLOB in Oracle (large text or JSON). When such a column is
# returned via a ``RETURNING`` clause it must be bound as DB_TYPE_CLOB; binding
# it as VARCHAR raises ORA-22835 ("buffer too small for CLOB to CHAR") once the
# value exceeds 4000 bytes. Union of the JSON-CLOB columns above and the
# large-text CLOB columns.
_CLOB_RETURNING_COLS = _JSON_COL_NAMES | {
"content",
"text",
"context",
"structured_content",
"text_signals",
"search_vector",
}
def _is_uuid_column(col: str) -> bool:
@@ -370,10 +349,6 @@ def _rewrite_pg_to_oracle(query: str) -> RewriteResult:
# Boolean literals: Oracle uses NUMBER(1) for booleans
query = re.sub(r"\b=\s*TRUE\b", "= 1", query, flags=re.IGNORECASE)
query = re.sub(r"\b=\s*FALSE\b", "= 0", query, flags=re.IGNORECASE)
# FOR NO KEY UPDATE → FOR UPDATE (Oracle has only FOR UPDATE; it does not block
# indexed-FK child inserts the way PG's FOR UPDATE would, so plain FOR UPDATE is
# the correct equivalent). Must run before the FOR SHARE rule below.
query = re.sub(r"\bFOR\s+NO\s+KEY\s+UPDATE\b", "FOR UPDATE", query, flags=re.IGNORECASE)
# FOR SHARE → FOR UPDATE (Oracle doesn't support FOR SHARE)
query = re.sub(r"\bFOR\s+SHARE\b", "FOR UPDATE", query, flags=re.IGNORECASE)
@@ -560,12 +535,6 @@ def _rewrite_pg_to_oracle(query: str) -> RewriteResult:
# != ALL(:N) → NOT IN (expanded list) — the negative counterpart of = ANY
query = _NOT_ALL_RE.sub(r"NOT IN (/*EXPAND:\1*/)", query)
# col LIKE ANY(:N) → (col LIKE :p0 OR col LIKE :p1 OR ...)
query = _LIKE_ANY_RE.sub(r"\1 /*LIKE_ANY:\2:\1*/", query)
# col NOT LIKE ALL(:N) → (col NOT LIKE :p0 AND col NOT LIKE :p1 AND ...)
query = _NOT_LIKE_ALL_RE.sub(r"\1 /*NOT_LIKE_ALL:\2:\1*/", query)
# CTE AS MATERIALIZED (...) → AS (...) — Oracle doesn't support MATERIALIZED CTE hint
query = re.sub(r"\bAS\s+MATERIALIZED\s*\(", "AS (", query, flags=re.IGNORECASE)
@@ -703,11 +672,6 @@ class OracleConnection(DatabaseConnection):
params[f"ret_{i}"] = cursor.var(oracledb.DB_TYPE_TIMESTAMP_TZ, arraysize=1)
elif clean in _NUMERIC_COLS:
params[f"ret_{i}"] = cursor.var(oracledb.DB_TYPE_NUMBER, arraysize=1)
elif clean in _CLOB_RETURNING_COLS:
# CLOB-backed column: a VARCHAR out-bind caps at 4000 bytes and
# raises ORA-22835 for larger values. Read back as a LOB in
# _read_returning_values.
params[f"ret_{i}"] = cursor.var(oracledb.DB_TYPE_CLOB, arraysize=1)
else:
params[f"ret_{i}"] = cursor.var(oracledb.DB_TYPE_VARCHAR, arraysize=1)
@@ -760,38 +724,17 @@ class OracleConnection(DatabaseConnection):
_expand_counter = 0
@staticmethod
def _resolve_list_param(params: dict[str, Any], key: str) -> list | None:
"""Resolve a parameter that may be a list or a JSON-encoded list string."""
val = params.get(key)
if isinstance(val, str):
try:
parsed = json.loads(val)
if isinstance(parsed, list):
return parsed
except (json.JSONDecodeError, TypeError):
pass
if isinstance(val, (list, tuple)):
return list(val)
return None
@staticmethod
def _expand_any_lists(query: str, params: dict[str, Any] | None) -> tuple[str, dict[str, Any] | None]:
"""Expand /*EXPAND:N*/, /*LIKE_ANY:N:col*/, /*NOT_LIKE_ALL:N:col*/ markers.
"""Expand /*EXPAND:N*/ markers into individual bind vars for IN clauses.
Converts: IN (/*EXPAND:1*/) with params["1"] = [a, b, c]
Into: IN (:any_0, :any_1, :any_2) with params["any_0"]=a, etc.
Converts: col /*LIKE_ANY:1:col*/ with params["1"] = [a, b]
Into: (col LIKE :lk_0 OR col LIKE :lk_1)
Converts: col /*NOT_LIKE_ALL:1:col*/ with params["1"] = [a, b]
Into: (col NOT LIKE :nlk_0 AND col NOT LIKE :nlk_1)
Uses a unique prefix to avoid name collisions with other bind vars.
The original param is kept (for other references to :N in the query).
"""
if params is None or "/*" not in query:
if params is None or "/*EXPAND:" not in query:
return query, params
expand_re = re.compile(r"/\*EXPAND:(\d+)\*/")
@@ -832,50 +775,6 @@ class OracleConnection(DatabaseConnection):
query = expand_re.sub(_replace, query)
# Expand LIKE ANY: col /*LIKE_ANY:N:col*/ → (col LIKE :p0 OR col LIKE :p1 ...)
like_any_re = re.compile(r"(\w+)\s*/\*LIKE_ANY:(\d+):(\w+)\*/")
def _replace_like_any(m):
_col = m.group(1) # redundant column ref before marker
param_key = m.group(2)
col = m.group(3)
val = OracleConnection._resolve_list_param(params, param_key)
if val is None or len(val) == 0:
return "1=0" # no patterns → no match
OracleConnection._expand_counter += 1
prefix = f"lk{OracleConnection._expand_counter}"
clauses = []
for i, item in enumerate(val):
k = f"{prefix}_{i}"
params[k] = item
clauses.append(f"{col} LIKE :{k}")
keys_to_remove.add(param_key)
return f"({' OR '.join(clauses)})"
query = like_any_re.sub(_replace_like_any, query)
# Expand NOT LIKE ALL: col /*NOT_LIKE_ALL:N:col*/ → (col NOT LIKE :p0 AND ...)
not_like_all_re = re.compile(r"(\w+)\s*/\*NOT_LIKE_ALL:(\d+):(\w+)\*/")
def _replace_not_like_all(m):
_col = m.group(1)
param_key = m.group(2)
col = m.group(3)
val = OracleConnection._resolve_list_param(params, param_key)
if val is None or len(val) == 0:
return "1=1" # no patterns → everything matches
OracleConnection._expand_counter += 1
prefix = f"nlk{OracleConnection._expand_counter}"
clauses = []
for i, item in enumerate(val):
k = f"{prefix}_{i}"
params[k] = item
clauses.append(f"{col} NOT LIKE :{k}")
keys_to_remove.add(param_key)
return f"({' AND '.join(clauses)})"
query = not_like_all_re.sub(_replace_not_like_all, query)
# Remove original list params that were expanded — their placeholder
# (:N) no longer exists in the query, and leaving them causes DPY-4008.
# Only remove if the key's placeholder is truly gone from the query.
@@ -885,7 +784,7 @@ class OracleConnection(DatabaseConnection):
return query, params
async def _read_returning_values(self, returning_cols: list[str], params: dict[str, Any]) -> dict[str, Any] | None:
def _read_returning_values(self, returning_cols: list[str], params: dict[str, Any]) -> dict[str, Any] | None:
"""Read values from RETURNING INTO output variables after execute."""
row: dict[str, Any] = {}
for i, col in enumerate(returning_cols):
@@ -895,14 +794,6 @@ class OracleConnection(DatabaseConnection):
return None
val = values[0] if isinstance(values, list) else values
# CLOB-bound columns return a LOB handle; read it to a string. The
# async pool yields AsyncLOB whose read() is a coroutine.
if val is not None and not isinstance(val, (str, bytes, int, float)) and hasattr(val, "read"):
data = val.read()
if inspect.isawaitable(data):
data = await data
val = data
# Clean alias: "LOWER(canonical_name) AS name_lower" → "name_lower"
clean_col = col.strip()
upper = clean_col.upper()
@@ -1090,7 +981,7 @@ class OracleConnection(DatabaseConnection):
raise
if ret_cols is not None:
row_dict = await self._read_returning_values(ret_cols, params)
row_dict = self._read_returning_values(ret_cols, params)
return [ResultRow(row_dict)] if row_dict else []
columns = [col[0].lower() for col in cursor.description or []]
@@ -1128,7 +1019,7 @@ class OracleConnection(DatabaseConnection):
raise
if ret_cols is not None:
row_dict = await self._read_returning_values(ret_cols, params)
row_dict = self._read_returning_values(ret_cols, params)
return ResultRow(row_dict) if row_dict else None
columns = [col[0].lower() for col in cursor.description or []]
@@ -1161,7 +1052,7 @@ class OracleConnection(DatabaseConnection):
await cursor.execute(query, params)
if ret_cols is not None:
row_dict = await self._read_returning_values(ret_cols, params)
row_dict = self._read_returning_values(ret_cols, params)
if row_dict is None:
return None
vals = list(row_dict.values())
@@ -6,7 +6,7 @@ import asyncio
import logging
import time
from collections.abc import AsyncIterator
from contextlib import AsyncExitStack, asynccontextmanager
from contextlib import asynccontextmanager
from typing import Any
logger = logging.getLogger(__name__)
@@ -101,14 +101,6 @@ async def acquire_with_retry(backend_or_pool: Any, max_retries: int = DEFAULT_MA
"""
Async context manager to acquire a database connection with retry logic.
Retries the *acquire* itself when it raises a retryable error (connection
drop, timeout, deadlock detected during acquire). Exceptions raised by
user code inside the ``async with`` block are NOT retried — they propagate
as-is. Wrapping retry around the yield would violate the
``@asynccontextmanager`` single-yield contract and surface as
``RuntimeError("generator didn't stop after athrow()")`` on every
retryable inner error, masking the real cause.
Accepts either a DatabaseBackend or a raw asyncpg.Pool for backward compatibility.
Usage:
@@ -117,7 +109,7 @@ async def acquire_with_retry(backend_or_pool: Any, max_retries: int = DEFAULT_MA
Args:
backend_or_pool: A DatabaseBackend instance or asyncpg.Pool
max_retries: Maximum number of retry attempts for the acquire step
max_retries: Maximum number of retry attempts
Yields:
A DatabaseConnection (if backend) or asyncpg.Connection (if pool)
@@ -125,32 +117,31 @@ async def acquire_with_retry(backend_or_pool: Any, max_retries: int = DEFAULT_MA
from .db.base import DatabaseBackend
if isinstance(backend_or_pool, DatabaseBackend) or getattr(backend_or_pool, "_wraps_backend", False):
# Use the backend's acquire context manager with retry
start = time.time()
async with AsyncExitStack() as stack:
conn: Any = None
for attempt in range(max_retries + 1):
try:
conn = await stack.enter_async_context(backend_or_pool.acquire())
break
except Exception as e:
if not _is_retryable(e):
raise
if attempt < max_retries:
delay = min(DEFAULT_BASE_DELAY * (2**attempt), DEFAULT_MAX_DELAY)
logger.warning(
f"Database acquire failed (attempt {attempt + 1}/{max_retries + 1}): {e}. "
f"Retrying in {delay:.1f}s..."
)
await asyncio.sleep(delay)
else:
logger.error(f"Database acquire failed after {max_retries + 1} attempts: {e}")
raise
acquire_time = time.time() - start
if acquire_time > 0.05:
logger.warning(f"[DB POOL] Slow acquire: {acquire_time:.3f}s")
yield conn
last_exception = None
for attempt in range(max_retries + 1):
try:
async with backend_or_pool.acquire() as conn:
acquire_time = time.time() - start
if acquire_time > 0.05:
logger.warning(f"[DB POOL] Slow acquire: {acquire_time:.3f}s")
yield conn
return
except Exception as e:
if not _is_retryable(e):
raise
last_exception = e
if attempt < max_retries:
delay = min(DEFAULT_BASE_DELAY * (2**attempt), DEFAULT_MAX_DELAY)
logger.warning(
f"Database acquire failed (attempt {attempt + 1}/{max_retries + 1}): {e}. "
f"Retrying in {delay:.1f}s..."
)
await asyncio.sleep(delay)
else:
logger.error(f"Database acquire failed after {max_retries + 1} attempts: {e}")
raise last_exception
else:
# Legacy path: raw asyncpg.Pool
pool = backend_or_pool
@@ -9,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,60 +27,24 @@ 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,
ENV_EMBEDDINGS_ONNX_DIMENSIONS,
ENV_EMBEDDINGS_ONNX_MODEL_ID,
ENV_EMBEDDINGS_ONNX_MODEL_PATH,
ENV_EMBEDDINGS_ONNX_TOKENIZER_NAME_OR_PATH,
ENV_EMBEDDINGS_OPENAI_API_KEY,
ENV_EMBEDDINGS_OPENAI_BASE_URL,
ENV_EMBEDDINGS_OPENAI_MODEL,
ENV_EMBEDDINGS_PROVIDER,
ENV_EMBEDDINGS_TEI_URL,
ENV_EMBEDDINGS_ZEROENTROPY_API_KEY,
ENV_EMBEDDINGS_ZEROENTROPY_DIMENSIONS,
ENV_EMBEDDINGS_ZEROENTROPY_ENCODING_FORMAT,
ENV_LLM_API_KEY,
)
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.
@@ -128,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):
"""
@@ -256,172 +208,6 @@ class LocalSTEmbeddings(Embeddings):
return [emb.tolist() for emb in embeddings]
class OnnxEmbeddings(Embeddings):
"""Local ONNX Runtime embeddings provider.
This provider runs transformer embedding models in-process with ONNX Runtime,
avoiding a sidecar Ollama/TEI server or a remote embeddings API. It supports
sentence-transformer style mean pooling and E5-style asymmetric prefixes.
"""
def __init__(
self,
model_id: str,
model_path: str | None = None,
tokenizer_name_or_path: str | None = None,
onnx_file: str = "onnx/model.onnx",
dimensions: int | None = None,
max_tokens: int = 512,
pooling: str = "mean",
normalize: bool = True,
query_prefix: str = "query: ",
passage_prefix: str = "passage: ",
output_name: str | None = None,
):
self.model_id = model_id
self.model_path = model_path
if model_path and tokenizer_name_or_path is None:
logger.warning(
"Embeddings: ONNX model_path is set without tokenizer_name_or_path; "
"falling back to tokenizer from model_id %s. Set "
"HINDSIGHT_API_EMBEDDINGS_ONNX_TOKENIZER_NAME_OR_PATH when using local ONNX artifacts.",
model_id,
)
self.tokenizer_name_or_path = tokenizer_name_or_path or model_id
self.onnx_file = onnx_file
self.configured_dimensions = dimensions
self.max_tokens = max_tokens
self.pooling = pooling.lower()
if self.pooling not in {"mean", "cls"}:
raise ValueError("ONNX embeddings pooling must be 'mean' or 'cls'")
self.normalize = normalize
self.query_prefix = query_prefix
self.passage_prefix = passage_prefix
self.output_name = output_name
self._session = None
self._tokenizer = None
self._dimension: int | None = dimensions
@property
def provider_name(self) -> str:
return "onnx"
@property
def dimension(self) -> int:
if self._dimension is None:
raise RuntimeError("Embeddings not initialized. Call initialize() first.")
return self._dimension
async def initialize(self) -> None:
if self._session is not None and self._tokenizer is not None:
return
try:
import onnxruntime as ort
from transformers import AutoTokenizer
except ImportError as exc:
raise ImportError(
"onnxruntime and transformers are required for OnnxEmbeddings. "
"Install with: pip install 'hindsight-api-slim[local-onnx]'"
) from exc
model_path = self.model_path
if not model_path:
try:
from huggingface_hub import snapshot_download
except ImportError as exc:
raise ImportError(
"huggingface-hub is required to download ONNX embedding models. "
"Set HINDSIGHT_API_EMBEDDINGS_ONNX_MODEL_PATH or install local-onnx."
) from exc
# Some large ONNX exports, for example BAAI/bge-m3, store weights in
# an external sidecar file next to model.onnx. Download both the
# requested graph and its conventional *_data sidecar when present.
snapshot_dir = snapshot_download(
repo_id=self.model_id,
allow_patterns=[self.onnx_file, f"{self.onnx_file}_data"],
)
model_path = os.path.join(snapshot_dir, self.onnx_file)
logger.info(
"Embeddings: initializing ONNX provider with model %s (%s)",
self.model_id,
model_path,
)
logger.info(
"Embeddings: ONNX query_prefix=%r passage_prefix=%r pooling=%s normalize=%s",
self.query_prefix,
self.passage_prefix,
self.pooling,
self.normalize,
)
self._tokenizer = AutoTokenizer.from_pretrained(self.tokenizer_name_or_path)
self._session = ort.InferenceSession(model_path, providers=["CPUExecutionProvider"])
detected = len(self.encode(["test"])[0])
if self.configured_dimensions is not None and detected != self.configured_dimensions:
raise ValueError(
f"Configured ONNX embedding dimension {self.configured_dimensions} does not match model output {detected}"
)
self._dimension = detected
logger.info("Embeddings: ONNX provider initialized (dim: %s)", self._dimension)
def _encode_prefixed(self, texts: list[str], prefix: str) -> list[list[float]]:
if prefix:
return self.encode([f"{prefix}{text}" for text in texts])
return self.encode(texts)
def encode_query(self, texts: list[str]) -> list[list[float]]:
return self._encode_prefixed(texts, self.query_prefix)
def encode_documents(self, texts: list[str]) -> list[list[float]]:
return self._encode_prefixed(texts, self.passage_prefix)
def encode(self, texts: list[str]) -> list[list[float]]:
if self._session is None or self._tokenizer is None:
raise RuntimeError("Embeddings not initialized. Call initialize() first.")
if not texts:
return []
import numpy as np
encoded = self._tokenizer(
texts,
padding=True,
truncation=True,
max_length=self.max_tokens,
return_tensors="np",
)
input_names = {inp.name for inp in self._session.get_inputs()}
ort_inputs = {name: value for name, value in encoded.items() if name in input_names}
if "token_type_ids" in input_names and "token_type_ids" not in ort_inputs:
ort_inputs["token_type_ids"] = np.zeros_like(encoded["input_ids"])
outputs = self._session.run([self.output_name] if self.output_name else None, ort_inputs)
token_embeddings = outputs[0]
# Some exported models expose a pooled 2-D embedding as their first output.
if getattr(token_embeddings, "ndim", 0) == 2:
embeddings = token_embeddings
elif self.pooling == "cls":
embeddings = token_embeddings[:, 0]
else:
attention_mask = encoded.get("attention_mask")
if attention_mask is None:
attention_mask = np.ones(token_embeddings.shape[:2], dtype=np.float32)
mask = attention_mask[..., None].astype(np.float32)
summed = (token_embeddings * mask).sum(axis=1)
counts = np.clip(mask.sum(axis=1), a_min=1e-9, a_max=None)
embeddings = summed / counts
if self.normalize:
norms = np.linalg.norm(embeddings, axis=1, keepdims=True)
norms[norms == 0] = 1
embeddings = embeddings / norms
return embeddings.astype(float).tolist()
class RemoteTEIEmbeddings(Embeddings):
"""
Remote embeddings implementation using HuggingFace Text Embeddings Inference (TEI) HTTP API.
@@ -599,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,
):
"""
@@ -610,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
@@ -662,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
@@ -699,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)
@@ -715,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.
@@ -923,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.
@@ -1199,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,
@@ -1211,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)
@@ -1262,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:
@@ -1315,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:
@@ -1561,20 +1125,6 @@ def create_embeddings_from_env() -> Embeddings:
force_cpu=config.embeddings_local_force_cpu,
trust_remote_code=config.embeddings_local_trust_remote_code,
)
elif provider == "onnx":
return OnnxEmbeddings(
model_id=config.embeddings_onnx_model_id,
model_path=config.embeddings_onnx_model_path,
tokenizer_name_or_path=config.embeddings_onnx_tokenizer_name_or_path,
onnx_file=config.embeddings_onnx_file,
dimensions=config.embeddings_onnx_dimensions,
max_tokens=config.embeddings_onnx_max_tokens,
pooling=config.embeddings_onnx_pooling,
normalize=config.embeddings_onnx_normalize,
query_prefix=config.embeddings_onnx_query_prefix,
passage_prefix=config.embeddings_onnx_passage_prefix,
output_name=config.embeddings_onnx_output_name,
)
elif provider == "openai":
# Use dedicated embeddings API key, or fall back to LLM API key
api_key = os.environ.get(ENV_EMBEDDINGS_OPENAI_API_KEY) or os.environ.get(ENV_LLM_API_KEY)
@@ -1590,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
@@ -1611,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
@@ -1646,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,
@@ -1676,6 +1206,5 @@ def create_embeddings_from_env() -> Embeddings:
else:
raise ValueError(
f"Unknown embeddings provider: {provider}. "
f"Supported: 'local', 'onnx', 'tei', 'openai', 'openai-codex', 'openrouter', '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)
@@ -458,28 +458,8 @@ class MemoryEngineInterface(ABC):
request_context: Request context for authentication.
Returns:
Dict with node_counts, link_counts, link_counts_by_fact_type
(deprecated, returns empty), link_breakdown (deprecated, returns
empty), and operations stats.
"""
...
@abstractmethod
async def get_bank_freshness(
self,
bank_id: str,
*,
request_context: "RequestContext",
) -> dict[str, Any]:
"""
Get consolidation freshness for a bank.
Cheap alternative to get_bank_stats when callers only need
last_consolidated_at / pending_consolidation / failed_consolidation.
Returns:
Dict with last_consolidated_at (ISO-8601 string or None),
pending_consolidation (int), and failed_consolidation (int).
Dict with node_counts, link_counts, link_counts_by_fact_type,
link_breakdown, and operations stats.
"""
...
@@ -69,7 +69,6 @@ class LLMInterface(ABC):
skip_validation: bool = False,
strict_schema: bool = False,
return_usage: bool = False,
cached_prefix: str | None = None,
) -> Any:
"""
Make an LLM API call with retry logic.
@@ -84,13 +83,8 @@ class LLMInterface(ABC):
initial_backoff: Initial backoff time in seconds.
max_backoff: Maximum backoff time in seconds.
skip_validation: Return raw JSON without Pydantic validation.
strict_schema: Grammar-enforce structured output via json_schema strict
(OpenAI-compatible, LiteLLM) instead of the soft json_object path. Gemini
enforces its response_schema natively; providers without a strict mode ignore it.
strict_schema: Use strict JSON schema enforcement (OpenAI only).
return_usage: If True, return tuple (result, TokenUsage) instead of just result.
cached_prefix: Opaque handle from ``get_or_create_cached_prefix`` for the
cacheable system prefix, or None. Providers without explicit prompt
caching ignore it (and the wrapper only forwards it when set).
Returns:
If return_usage=False: Parsed response if response_format is provided, otherwise text content.
@@ -114,7 +108,6 @@ class LLMInterface(ABC):
initial_backoff: float = 1.0,
max_backoff: float = 30.0,
tool_choice: str | dict[str, Any] = "auto",
cached_prefix: str | None = None,
) -> LLMToolCallResult:
"""
Make an LLM API call with tool/function calling support.
@@ -144,46 +137,6 @@ class LLMInterface(ABC):
"""
return False
# ── Prompt prefix caching (optional, per-provider) ─────────────────────────
def supports_prompt_caching(self) -> bool:
"""Whether this provider can cache a reusable prompt prefix.
Default False. Providers that return True must implement
``get_or_create_cached_prefix`` and honour the ``cached_prefix`` argument
of ``call`` / ``call_with_tools``.
"""
return False
async def get_or_create_cached_prefix(
self,
*,
system_instruction: str,
response_schema: Any | None = None,
tools: list[dict[str, Any]] | None = None,
) -> str | None:
"""Cache a reusable prompt prefix and return an opaque handle, or None.
The engine has already decided WHAT is cacheable: it puts the stable,
bank-agnostic instructions in ``system_instruction`` (plus ``tools``) and
keeps all per-request / per-bank data (documents, facts, the bank mission)
in the user message. A provider only chooses HOW to cache that prefix:
- Explicit-cache providers (e.g. Gemini ``CachedContent``): create the
cache, return its handle; the engine passes the handle back via
``call(cached_prefix=...)`` and the provider then drops the prefix from
the request, billing it at the cached rate.
- Automatic-cache providers (e.g. OpenAI): no handle needed — caching is
transparent as long as the prefix is a stable leading block, which it
already is. They can keep this default (return None) and still benefit.
- Inline-marker providers (e.g. Anthropic ``cache_control``): mark the
prefix block inside ``call`` instead; may also keep this default.
Returns None when caching is disabled/unsupported or the prefix is too
small; callers MUST fall back to an uncached call in that case.
"""
return None
async def submit_batch(
self,
requests: list[dict[str, Any]],
@@ -1,540 +0,0 @@
"""Per-bank LLM request tracing.
Opt-in, fire-and-forget recording of every LLM call Hindsight makes (both
successes and failures) into the ``llm_requests`` table, per bank. Each row
captures the input messages, the model output, token usage (input / output /
cached / total), finish reason, and caller metadata. Disabled by default —
controlled by ``HINDSIGHT_API_LLM_TRACE_ENABLED``.
This plugs into the OpenTelemetry **GenAI** recording pattern: providers already
call ``tracing.get_span_recorder().record_llm_call(...)`` on success, so the DB
tracer is registered as one of those recorders (alongside the OTLP span
exporter) rather than hooking the call path with custom code. Failures, which
providers don't report to the recorder, are forwarded from the LLM wrapper.
Bank/operation attribution is carried via a ContextVar set by
``ConfiguredLLMProvider`` (see ``llm_wrapper.py``); outside a traced context
``bank_id`` is recorded as NULL.
"""
from __future__ import annotations
import asyncio
import json
import logging
import uuid
from collections.abc import Callable, Iterable
from contextvars import ContextVar, Token
from dataclasses import dataclass, field
from datetime import datetime, timedelta, timezone
from typing import Any
from pydantic import BaseModel
from .db_utils import acquire_with_retry
logger = logging.getLogger(__name__)
# ── bank/operation attribution (carried across the async call chain) ──────────
@dataclass
class LLMTraceContext:
"""Attribution for in-flight LLM calls, bound by ``ConfiguredLLMProvider``.
``trace_id`` and ``operation_span_id`` are generated once per operation
invocation (one ``with_config`` call), so every LLM call of a single
reflect/retain/consolidation run shares them — reproducing the OTel
parent (operation span) → children (LLM calls) hierarchy in the DB.
"""
bank_id: str | None = None
operation: str | None = None # "retain" | "reflect" | "consolidation" | ...
metadata: dict[str, Any] = field(default_factory=dict)
trace_id: str | None = None
operation_span_id: str | None = None
# Memory_units this operation produced/consumed, accumulated at the DB-write
# sites and flushed onto every row of the trace at operation end (see
# LLMTraceRecorder.attach_memory_ids). Lets a retain/consolidation trace map
# to the memories it created (outputs) and consumed (source inputs).
created_memory_ids: list[str] = field(default_factory=list)
source_memory_ids: list[str] = field(default_factory=list)
_trace_ctx: ContextVar[LLMTraceContext | None] = ContextVar("hindsight_llm_trace_ctx", default=None)
# Per-call requested parameters (max_completion_tokens, temperature, response
# schema, tool_choice). Set by ``LLMProvider.call`` around the provider
# delegation so the recorder can attach them even though success is reported by
# the provider. Only includes values the caller actually set — never nulls.
_request_ctx: ContextVar[dict[str, Any] | None] = ContextVar("hindsight_llm_request_ctx", default=None)
# Per-call caller metadata (e.g. document_id for retain extraction). Set by
# engine code around a specific LLM call; merged into the row's metadata on top
# of the operation-level LLMTraceContext.metadata.
_call_metadata_ctx: ContextVar[dict[str, Any] | None] = ContextVar("hindsight_llm_call_metadata_ctx", default=None)
def set_trace_context(ctx: LLMTraceContext | None) -> Token:
"""Bind trace attribution to the current context. Returns a reset token."""
return _trace_ctx.set(ctx)
def reset_trace_context(token: Token) -> None:
"""Unwind a binding made by :func:`set_trace_context`."""
_trace_ctx.reset(token)
def set_request_context(params: dict[str, Any] | None) -> Token:
"""Bind the current LLM call's requested parameters. Returns a reset token."""
return _request_ctx.set(params)
def reset_request_context(token: Token) -> None:
"""Unwind a binding made by :func:`set_request_context`."""
_request_ctx.reset(token)
def current_request_context() -> dict[str, Any] | None:
"""Return the active call's requested parameters, or None."""
return _request_ctx.get()
def set_call_metadata(metadata: dict[str, Any] | None) -> Token:
"""Bind per-call caller metadata (e.g. ``{"document_id": ...}``)."""
return _call_metadata_ctx.set(metadata)
def reset_call_metadata(token: Token) -> None:
"""Unwind a binding made by :func:`set_call_metadata`."""
_call_metadata_ctx.reset(token)
def current_call_metadata() -> dict[str, Any] | None:
"""Return the active call's caller metadata, or None."""
return _call_metadata_ctx.get()
def current_trace_context() -> LLMTraceContext | None:
"""Return the active trace attribution, or None outside a traced context."""
return _trace_ctx.get()
def trace_context_of(llm_config: Any) -> LLMTraceContext | None:
"""Return a configured provider's operation trace context, or None.
Real providers expose ``trace_context()`` (``ConfiguredLLMProvider``); test
or mock substitutes may not, so this degrades gracefully rather than raising
— tracing is best-effort and must never break an operation.
"""
getter = getattr(llm_config, "trace_context", None)
return getter() if callable(getter) else None
def record_created_memory_ids(ids: Iterable[str]) -> None:
"""Accumulate output memory_units onto the active operation trace.
No-op outside a traced operation context (e.g. tracing disabled). Child
asyncio tasks inherit the same ``LLMTraceContext`` object, so appends from
parallel consolidation batches land on one shared list.
"""
ctx = _trace_ctx.get()
if ctx is not None:
ctx.created_memory_ids.extend(str(i) for i in ids)
def record_source_memory_ids(ids: Iterable[str]) -> None:
"""Accumulate consumed/source memory_units onto the active operation trace.
No-op outside a traced operation context.
"""
ctx = _trace_ctx.get()
if ctx is not None:
ctx.source_memory_ids.extend(str(i) for i in ids)
# ── serialization helpers ─────────────────────────────────────────────────────
def _json_default(obj: Any) -> Any:
"""JSON serializer for objects not serializable by default."""
if isinstance(obj, datetime):
return obj.isoformat()
if isinstance(obj, uuid.UUID):
return str(obj)
if isinstance(obj, bytes):
return "<bytes>"
if isinstance(obj, set):
return list(obj)
model_dump = getattr(obj, "model_dump", None)
if callable(model_dump):
try:
return model_dump(mode="json")
except Exception:
return str(obj)
return str(obj)
def _safe_json(data: Any, max_chars: int) -> str | None:
"""Serialize ``data`` to a JSON string, truncating beyond ``max_chars``.
Returns None on total failure. Truncation preserves valid JSON by wrapping
the oversized payload in a marker object with a preview.
"""
if data is None:
return None
try:
serialized = json.dumps(data, default=_json_default)
except Exception:
logger.debug("Failed to serialize llm trace data", exc_info=True)
try:
serialized = json.dumps(str(data))
except Exception:
return None
if max_chars and max_chars > 0 and len(serialized) > max_chars:
return json.dumps({"_truncated": True, "_original_chars": len(serialized), "preview": serialized[:max_chars]})
return serialized
# ── record ────────────────────────────────────────────────────────────────────
@dataclass
class LLMRequestRecord:
"""A single LLM request trace row."""
provider: str
model: str | None
scope: str
status: str # "success" | "error"
started_at: datetime
ended_at: datetime
bank_id: str | None = None
operation: str | None = None
trace_id: str | None = None
span_id: str | None = None
parent_span_id: str | None = None
input: Any = None
output: Any = None
error: str | None = None
input_tokens: int | None = None
output_tokens: int | None = None
cached_tokens: int | None = None
total_tokens: int | None = None
llm_info: dict[str, Any] = field(default_factory=dict)
metadata: dict[str, Any] = field(default_factory=dict)
@property
def duration_ms(self) -> int:
return int((self.ended_at - self.started_at).total_seconds() * 1000)
# ── read models (returned by MemoryEngine query methods, served by the API) ───
class LLMRequestEntry(BaseModel):
"""A single LLM request trace row, as returned by the read API."""
id: str
bank_id: str | None
operation: str | None
scope: str | None
trace_id: str | None
span_id: str | None
parent_span_id: str | None
provider: str | None
model: str | None
status: str
started_at: str | None
ended_at: str | None
duration_ms: int | None
input_tokens: int | None
output_tokens: int | None
cached_tokens: int | None
total_tokens: int | None
# Arbitrary JSON (message list, string, or object) — open `Any` so the
# OpenAPI schema stays a plain open type the Go SDK generator can model.
input: Any = None
output: Any = None
error: str | None
llm_info: dict[str, Any]
metadata: dict[str, Any]
class LLMRequestListResponse(BaseModel):
"""Paginated list of LLM request traces for a bank."""
bank_id: str
total: int
limit: int
offset: int
items: list[LLMRequestEntry]
class LLMRequestTokenSums(BaseModel):
"""Token totals for a time bucket."""
input: int
output: int
cached: int
total: int
class LLMRequestStatsBucket(BaseModel):
"""A single time bucket in LLM request stats."""
time: str
statuses: dict[str, int]
total: int
tokens: LLMRequestTokenSums
class LLMRequestStatsResponse(BaseModel):
"""LLM request counts and token sums grouped by time bucket."""
bank_id: str
period: str
trunc: str
start: str
buckets: list[LLMRequestStatsBucket]
# ── recorder / writer ─────────────────────────────────────────────────────────
class LLMTraceRecorder:
"""GenAI span recorder that writes per-bank LLM traces to ``llm_requests``.
Implements ``record_llm_call`` so it can be registered with
:func:`hindsight_api.tracing.register_span_recorder`. Writes are
fire-and-forget and never surface errors into the calling path. Retention of
old rows is handled by the background :class:`MaintenanceLoop`.
"""
def __init__(
self,
pool_getter: Callable[[], Any],
schema_getter: Callable[[], str],
enabled: bool,
allowed_scopes: list[str],
max_chars: int = 50000,
) -> None:
self._pool_getter = pool_getter
self._schema_getter = schema_getter
self._enabled = enabled
self._allowed_scopes: frozenset[str] | None = frozenset(allowed_scopes) if allowed_scopes else None
self._max_chars = max_chars
# In-flight fire-and-forget write tasks, bucketed by trace_id so
# attach_memory_ids can await only *its own* operation's writes before the
# post-operation UPDATE (otherwise the UPDATE could race ahead of the
# INSERTs it patches — but it must not block on unrelated operations).
self._pending: dict[str | None, set[asyncio.Task]] = {}
def is_enabled(self, scope: str) -> bool:
"""Whether tracing is active for the given call scope."""
if not self._enabled:
return False
if self._allowed_scopes is not None:
return scope in self._allowed_scopes
return True
# ── GenAI recorder interface ──────────────────────────────────────────────
def record_llm_call(
self,
provider: str,
model: str,
scope: str,
messages: list[dict[str, Any]],
response_content: Any = None,
input_tokens: int = 0,
output_tokens: int = 0,
duration: float = 0.0,
finish_reason: str | None = None,
error: BaseException | None = None,
tool_calls: list[dict[str, Any]] | None = None,
cached_tokens: int = 0,
**_extra: Any,
) -> None:
"""Build a trace record from a GenAI call and schedule a DB write."""
if not self.is_enabled(scope):
return
ctx = current_trace_context()
ended_at = datetime.now(timezone.utc)
started_at = ended_at - timedelta(seconds=max(0.0, duration))
# Operation-level metadata + any per-call metadata (e.g. document_id).
metadata = dict(ctx.metadata) if ctx else {}
call_metadata = current_call_metadata()
if call_metadata:
metadata.update(call_metadata)
llm_info: dict[str, Any] = {}
request_params = current_request_context()
if request_params:
llm_info["request"] = dict(request_params)
if finish_reason:
llm_info["finish_reason"] = finish_reason
if tool_calls:
llm_info["tool_calls"] = [tc.get("name", "") for tc in tool_calls]
record = LLMRequestRecord(
provider=provider,
model=model,
scope=scope,
status="error" if error is not None else "success",
started_at=started_at,
ended_at=ended_at,
bank_id=ctx.bank_id if ctx else None,
operation=ctx.operation if ctx else None,
# OTel-style hierarchy: all calls of one operation invocation share
# the context's trace_id and point at its operation span; this call
# gets its own span_id.
trace_id=ctx.trace_id if ctx else None,
span_id=str(uuid.uuid4()),
parent_span_id=ctx.operation_span_id if ctx else None,
input=messages,
output=None if error is not None else response_content,
error=f"{type(error).__name__}: {error}" if error is not None else None,
input_tokens=input_tokens or None,
output_tokens=output_tokens or None,
cached_tokens=cached_tokens or None,
total_tokens=(input_tokens + output_tokens) or None,
llm_info=llm_info,
metadata=metadata,
)
self._record_fire_and_forget(record)
def _record_fire_and_forget(self, record: LLMRequestRecord) -> None:
"""Schedule a trace write as a background task."""
try:
task = asyncio.create_task(self._safe_write(record))
except RuntimeError:
# No running event loop (e.g. during shutdown)
logger.debug("Cannot schedule llm trace write: no running event loop")
return
key = record.trace_id
self._pending.setdefault(key, set()).add(task)
task.add_done_callback(lambda t, k=key: self._discard_pending(k, t))
def _discard_pending(self, key: str | None, task: asyncio.Task) -> None:
bucket = self._pending.get(key)
if bucket is not None:
bucket.discard(task)
if not bucket:
self._pending.pop(key, None)
async def _safe_write(self, record: LLMRequestRecord) -> None:
"""Write a trace row. Errors are logged, never raised."""
pool = self._pool_getter()
if pool is None:
logger.debug("LLM trace skipped: pool not available")
return
try:
schema = self._schema_getter()
table = f"{schema}.llm_requests"
async with acquire_with_retry(pool, max_retries=1) as conn:
await conn.execute(
f"""
INSERT INTO {table}
(id, bank_id, operation, scope, trace_id, span_id, parent_span_id,
provider, model, status,
started_at, ended_at, duration_ms,
input_tokens, output_tokens, cached_tokens, total_tokens,
input, output, error, llm_info, metadata)
VALUES
($1, $2, $3, $4, $5, $6, $7, $8, $9, $10,
$11, $12, $13, $14, $15, $16, $17,
$18::jsonb, $19::jsonb, $20, $21::jsonb, $22::jsonb)
""",
uuid.uuid4(),
record.bank_id,
record.operation,
record.scope,
record.trace_id,
record.span_id,
record.parent_span_id,
record.provider,
record.model,
record.status,
record.started_at,
record.ended_at,
record.duration_ms,
record.input_tokens,
record.output_tokens,
record.cached_tokens,
record.total_tokens,
_safe_json(record.input, self._max_chars),
_safe_json(record.output, self._max_chars),
record.error,
_safe_json(record.llm_info, self._max_chars) or "{}",
_safe_json(record.metadata, self._max_chars) or "{}",
)
except Exception as e:
logger.warning(f"LLM trace write failed for scope={record.scope}: {e}")
async def _flush_pending(self, trace_id: str) -> None:
"""Await this trace's in-flight writes so its rows exist before an UPDATE."""
pending = [t for t in self._pending.get(trace_id, ()) if not t.done()]
if pending:
await asyncio.gather(*pending, return_exceptions=True)
def attach_memory_ids(
self,
trace_ctx: LLMTraceContext | None,
*,
created: list[str] | None = None,
source: list[str] | None = None,
) -> None:
"""Map a finished operation's memory_units onto every row of its trace.
Merges the explicitly passed ids with any accumulated on the context
(``record_created_memory_ids`` / ``record_source_memory_ids``), de-dupes
preserving order, and patches ``metadata.memory_ids`` (outputs created)
and ``metadata.source_memory_ids`` (inputs consumed) on all rows sharing
the trace_id. No-op when tracing is off or nothing was produced.
Fire-and-forget: the snapshotted patch is applied on a background task so
the retain/consolidation operation never waits on the trace write. The
ids are snapshotted synchronously here because the caller may reset the
context immediately after.
"""
if not self._enabled or trace_ctx is None or not trace_ctx.trace_id:
return
created_ids = list(dict.fromkeys([*(created or []), *trace_ctx.created_memory_ids]))
source_ids = list(dict.fromkeys([*(source or []), *trace_ctx.source_memory_ids]))
patch: dict[str, Any] = {}
if created_ids:
patch["memory_ids"] = created_ids
if source_ids:
patch["source_memory_ids"] = source_ids
if not patch:
return
try:
asyncio.create_task(self._attach_memory_ids(trace_ctx.bank_id, trace_ctx.trace_id, patch))
except RuntimeError:
logger.debug("Cannot schedule llm trace memory_id attach: no running event loop")
async def _attach_memory_ids(self, bank_id: str | None, trace_id: str, patch: dict[str, Any]) -> None:
"""Background worker: flush this trace's writes, then patch its rows."""
# The trace-row INSERTs are fire-and-forget; flush *this trace's* writes
# so the UPDATE patches rows that already exist rather than racing ahead
# of them (without blocking on unrelated operations' pending writes).
await self._flush_pending(trace_id)
pool = self._pool_getter()
if pool is None:
return
try:
schema = self._schema_getter()
table = f"{schema}.llm_requests"
async with acquire_with_retry(pool, max_retries=1) as conn:
await conn.execute(
f"UPDATE {table} SET metadata = metadata || $3::jsonb WHERE bank_id = $1 AND trace_id = $2",
bank_id,
trace_id,
json.dumps(patch),
)
except Exception as e:
logger.warning(f"LLM trace memory_id attach failed for trace={trace_id}: {e}")
@@ -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,101 +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 _request_params(
*,
max_completion_tokens: int | None = None,
temperature: float | None = None,
scope: str | None = None,
response_format: Any | None = None,
tool_choice: str | dict[str, Any] | None = None,
) -> dict[str, Any] | None:
"""Build the requested-params bag for tracing — only values the caller set.
Omitting unset values avoids the misleading nulls we used to record (e.g.
consolidation, which passes no token cap), while surfacing the real cap for
callers that do set one (e.g. retain's ``retain_max_completion_tokens``).
"""
params: dict[str, Any] = {}
if max_completion_tokens is not None:
params["max_completion_tokens"] = max_completion_tokens
if temperature is not None:
params["temperature"] = temperature
if response_format is not None:
params["response_schema"] = getattr(response_format, "__name__", None) or "structured"
if tool_choice is not None and tool_choice != "auto":
params["tool_choice"] = tool_choice if isinstance(tool_choice, str) else "named"
return params or None
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.
@@ -152,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
@@ -166,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.
@@ -255,7 +154,6 @@ def create_llm_provider(
vertexai_region: str | None = None,
vertexai_credentials: Any = None,
gemini_safety_settings: list | None = None,
prompt_cache_enabled: bool = False,
litellmrouter_config: dict[str, Any] | None = None,
) -> Any: # Returns LLMInterface
"""
@@ -269,11 +167,7 @@ def create_llm_provider(
reasoning_effort: Reasoning effort level for supported providers.
groq_service_tier: Groq service tier (for Groq provider) - "on_demand", "flex", or "auto".
openai_service_tier: OpenAI service tier (for OpenAI provider) - None (default) or "flex" (50% cheaper).
extra_body: Extra request-body params merged into the provider's native
call. Threaded into OpenAI-compatible, Fireworks, Anthropic, Gemini/
VertexAI and LiteLLM providers (each merges them in its own parameter
space). Keys must use each provider's native names (e.g. ``max_tokens``
for OpenAI/Anthropic vs ``max_output_tokens`` for Gemini).
extra_body: Extra body params merged into OpenAI-compatible API calls.
default_headers: Custom headers passed as ``default_headers`` to provider SDK clients
(used by operators routing through proxies / request-tracing middleware). Currently
wired into the Anthropic provider; other providers may opt in as needed.
@@ -289,7 +183,6 @@ def create_llm_provider(
AnthropicLLM,
ClaudeCodeLLM,
CodexLLM,
FireworksLLM,
GeminiLLM,
LiteLLMLLM,
LiteLLMRouterLLM,
@@ -348,8 +241,6 @@ def create_llm_provider(
vertexai_region=vertexai_region,
vertexai_credentials=vertexai_credentials,
gemini_safety_settings=gemini_safety_settings,
prompt_cache_enabled=prompt_cache_enabled,
extra_body=extra_body,
)
elif provider_lower == "anthropic":
@@ -360,7 +251,6 @@ def create_llm_provider(
model=model,
reasoning_effort=reasoning_effort,
default_headers=default_headers,
extra_body=extra_body,
)
elif provider_lower == "litellm":
@@ -370,7 +260,6 @@ def create_llm_provider(
base_url=base_url,
model=model,
reasoning_effort=reasoning_effort,
extra_body=extra_body,
)
elif provider_lower == "litellmrouter":
@@ -388,7 +277,6 @@ def create_llm_provider(
model=model,
config=litellmrouter_config,
reasoning_effort=reasoning_effort,
extra_body=extra_body,
)
elif provider_lower == "bedrock":
@@ -400,7 +288,6 @@ def create_llm_provider(
base_url=base_url,
model=bedrock_model,
reasoning_effort=reasoning_effort,
extra_body=extra_body,
)
elif provider_lower == "llamacpp":
@@ -421,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",
@@ -479,7 +352,6 @@ class LLMProvider:
groq_service_tier: str | None = None,
openai_service_tier: str | None = None,
gemini_safety_settings: list | None = None,
prompt_cache_enabled: bool = False,
extra_body: dict[str, Any] | None = None,
default_headers: dict[str, str] | None = None,
litellmrouter_config: dict[str, Any] | None = None,
@@ -496,8 +368,7 @@ class LLMProvider:
groq_service_tier: Groq service tier ("on_demand", "flex", "auto") - from config.
openai_service_tier: OpenAI service tier (None or "flex") - from config.
gemini_safety_settings: Safety settings for Gemini/VertexAI providers.
extra_body: Extra request-body params merged into the provider's native call
(OpenAI-compatible, Fireworks, Anthropic, Gemini/VertexAI, LiteLLM).
extra_body: Extra body params merged into OpenAI-compatible API calls.
default_headers: Custom headers passed as ``default_headers`` to provider SDK clients.
Used by operators routing through proxies / request-tracing middleware. Falls
back to ``HindsightConfig.llm_default_headers`` (env: ``HINDSIGHT_API_LLM_DEFAULT_HEADERS``)
@@ -519,11 +390,6 @@ class LLMProvider:
self.openai_service_tier = openai_service_tier
# Gemini safety settings (instance default; can be overridden per-request via context var)
self.gemini_safety_settings = gemini_safety_settings
# Gemini prompt caching: when True, retain extraction (and any future
# caller that opts in) will reuse a CachedContent prefix to cut
# input-token cost. Off by default so the change is observable behind
# a flip rather than a silent behaviour change on upgrade.
self.prompt_cache_enabled = prompt_cache_enabled
# Extra body params for OpenAI-compatible providers (e.g. chat_template_kwargs)
self.extra_body = extra_body
# Default headers passed to provider SDK clients (e.g. proxy auth, request tracing).
@@ -543,7 +409,6 @@ class LLMProvider:
"openai",
"groq",
"ollama",
"ollama-cloud",
"gemini",
"anthropic",
"lmstudio",
@@ -562,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)}")
@@ -573,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":
@@ -642,21 +504,6 @@ class LLMProvider:
except Exception:
pass # Config may not be initialized in test environments
# Prompt-prefix caching is a provider-agnostic toggle (default on): resolve
# it from the static server config for every provider when the caller didn't
# pass an explicit override. Providers that don't support caching ignore the
# value; only those that implement get_or_create_cached_prefix act on it.
if not self.prompt_cache_enabled:
from ..config import DEFAULT_LLM_PROMPT_CACHE_ENABLED, _get_raw_config
try:
raw_config = _get_raw_config()
self.prompt_cache_enabled = bool(
getattr(raw_config, "llm_prompt_cache_enabled", DEFAULT_LLM_PROMPT_CACHE_ENABLED)
)
except Exception:
pass # Config may not be initialized in test environments
# For litellmrouter: prefer an explicit chain from the caller (per-op
# construction in MemoryEngine threads the right chain through). If the caller
# didn't supply one, fall back to the global ``llm_litellmrouter_config`` so
@@ -685,7 +532,6 @@ class LLMProvider:
vertexai_region=vertexai_region,
vertexai_credentials=vertexai_credentials,
gemini_safety_settings=self.gemini_safety_settings,
prompt_cache_enabled=self.prompt_cache_enabled,
litellmrouter_config=router_config,
)
@@ -749,7 +595,6 @@ class LLMProvider:
skip_validation: bool = False,
strict_schema: bool = False,
return_usage: bool = False,
cached_prefix: str | None = None,
) -> Any:
"""
Make an LLM API call with retry logic.
@@ -764,10 +609,7 @@ class LLMProvider:
initial_backoff: Initial backoff time in seconds.
max_backoff: Maximum backoff time in seconds.
skip_validation: Return raw JSON without Pydantic validation.
strict_schema: Per-call override requesting grammar-enforced (json_schema strict)
structured output instead of the soft json_object path. The server-level
HINDSIGHT_API_LLM_STRICT_SCHEMA flag is OR-ed in here so it applies to every call;
providers without a strict mode ignore it.
strict_schema: Use strict JSON schema enforcement (OpenAI only). Guarantees all required fields.
return_usage: If True, return tuple (result, TokenUsage) instead of just result.
Returns:
@@ -787,85 +629,32 @@ class LLMProvider:
structured = "+structured" if response_format is not None else ""
set_stage(f"llm.{self.provider}.{scope}{structured}")
# Resolve strict-schema once, here, rather than in each provider: the
# per-call argument OR the server-level HINDSIGHT_API_LLM_STRICT_SCHEMA
# flag. Providers with a json_schema response_format (OpenAI-compatible,
# LiteLLM) then grammar-enforce structured output instead of the fragile
# soft json_object path; Gemini already enforces its native response_schema,
# and providers without a strict mode simply ignore the flag.
from ..config import get_config
strict_schema = strict_schema or get_config().llm_strict_schema
# LLM call observability flows through the OTel GenAI recorder
# (tracing.get_span_recorder().record_llm_call). Provider implementations
# record successful calls; we forward failures here since they don't.
# The requested params are stashed in a contextvar (only what the caller
# actually set) so the recorder can attach them to either path.
from ..tracing import get_span_recorder
from .llm_trace import reset_request_context, set_request_context
call_start = time.monotonic()
request_token = set_request_context(
_request_params(
async with _global_llm_semaphore:
# Delegate to provider implementation
result = await self._provider_impl.call(
messages=messages,
response_format=response_format,
max_completion_tokens=max_completion_tokens,
temperature=temperature,
scope=scope,
response_format=response_format,
max_retries=max_retries,
initial_backoff=initial_backoff,
max_backoff=max_backoff,
skip_validation=skip_validation,
strict_schema=strict_schema,
return_usage=return_usage,
)
)
try:
async with AsyncExitStack() as stack:
for sem in _semaphores_for_scope(scope):
await stack.enter_async_context(sem)
# cached_prefix is only set for providers that returned a handle
# from get_or_create_cached_prefix() (e.g. Gemini); it's None for
# the rest. Forward it only when present so providers that don't
# implement caching keep their call() signature untouched.
cache_kwarg = {"cached_prefix": cached_prefix} if cached_prefix is not None else {}
try:
# Delegate to provider implementation
result = await self._provider_impl.call(
messages=messages,
response_format=response_format,
max_completion_tokens=max_completion_tokens,
temperature=temperature,
scope=scope,
max_retries=max_retries,
initial_backoff=initial_backoff,
max_backoff=max_backoff,
skip_validation=skip_validation,
strict_schema=strict_schema,
return_usage=return_usage,
**cache_kwarg,
)
except Exception as e:
get_span_recorder().record_llm_call(
provider=self.provider,
model=self.model,
scope=scope,
messages=messages,
response_content=None,
input_tokens=0,
output_tokens=0,
duration=time.monotonic() - call_start,
error=e,
)
raise
# Backward compatibility: Update mock call tracking for mock provider
# This allows existing tests using LLMProvider._mock_calls to continue working
if self.provider == "mock":
from .providers.mock_llm import MockLLM
# Backward compatibility: Update mock call tracking for mock provider
# This allows existing tests using LLMProvider._mock_calls to continue working
if self.provider == "mock":
from .providers.mock_llm import MockLLM
if isinstance(self._provider_impl, MockLLM):
# Sync the mock calls from provider implementation to wrapper
self._mock_calls = self._provider_impl.get_mock_calls()
if isinstance(self._provider_impl, MockLLM):
# Sync the mock calls from provider implementation to wrapper
self._mock_calls = self._provider_impl.get_mock_calls()
finally:
reset_request_context(request_token)
return result
return result
async def call_with_tools(
self,
@@ -878,7 +667,6 @@ class LLMProvider:
initial_backoff: float = 1.0,
max_backoff: float = 30.0,
tool_choice: str | dict[str, Any] = "auto",
cached_prefix: str | None = None,
) -> "LLMToolCallResult":
"""
Make an LLM API call with tool/function calling support.
@@ -901,68 +689,30 @@ class LLMProvider:
set_stage(f"llm.{self.provider}.{scope}+tools")
# Failures forwarded to the GenAI recorder; successes recorded by providers.
from ..tracing import get_span_recorder
from .llm_trace import reset_request_context, set_request_context
call_start = time.monotonic()
request_token = set_request_context(
_request_params(
async with _global_llm_semaphore:
# Delegate to provider implementation
result = await self._provider_impl.call_with_tools(
messages=messages,
tools=tools,
max_completion_tokens=max_completion_tokens,
temperature=temperature,
scope=scope,
max_retries=max_retries,
initial_backoff=initial_backoff,
max_backoff=max_backoff,
tool_choice=tool_choice,
)
)
try:
async with AsyncExitStack() as stack:
for sem in _semaphores_for_scope(scope):
await stack.enter_async_context(sem)
# cached_prefix is only set for providers that returned a handle
# from get_or_create_cached_prefix(); forward it only when present
# so non-caching providers keep their signature (same as call()).
cache_kwarg = {"cached_prefix": cached_prefix} if cached_prefix is not None else {}
try:
# Delegate to provider implementation
result = await self._provider_impl.call_with_tools(
messages=messages,
tools=tools,
max_completion_tokens=max_completion_tokens,
temperature=temperature,
scope=scope,
max_retries=max_retries,
initial_backoff=initial_backoff,
max_backoff=max_backoff,
tool_choice=tool_choice,
**cache_kwarg,
)
except Exception as e:
get_span_recorder().record_llm_call(
provider=self.provider,
model=self.model,
scope=scope,
messages=messages,
response_content=None,
input_tokens=0,
output_tokens=0,
duration=time.monotonic() - call_start,
error=e,
)
raise
# Backward compatibility: Update mock call tracking for mock provider
# This allows existing tests using LLMProvider._mock_calls to continue working
if self.provider == "mock":
from .providers.mock_llm import MockLLM
# Backward compatibility: Update mock call tracking for mock provider
# This allows existing tests using LLMProvider._mock_calls to continue working
if self.provider == "mock":
from .providers.mock_llm import MockLLM
if isinstance(self._provider_impl, MockLLM):
# Sync the mock calls from provider implementation to wrapper
self._mock_calls = self._provider_impl.get_mock_calls()
if isinstance(self._provider_impl, MockLLM):
# Sync the mock calls from provider implementation to wrapper
self._mock_calls = self._provider_impl.get_mock_calls()
finally:
reset_request_context(request_token)
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."""
@@ -1064,14 +814,7 @@ class LLMProvider:
# SDK will automatically check for authentication when first used
# No need to verify here - let it fail gracefully on first call with helpful error
def with_config(
self,
config: Any,
*,
bank_id: str | None = None,
operation: str | None = None,
metadata: dict[str, Any] | None = None,
) -> "ConfiguredLLMProvider":
def with_config(self, config: Any) -> "ConfiguredLLMProvider":
"""
Return a configured wrapper for a specific bank operation.
@@ -1081,31 +824,12 @@ class LLMProvider:
Args:
config: Resolved ``HindsightConfig`` for the current bank/request.
bank_id: Bank the operation runs for; attributed to LLM trace rows.
operation: Logical operation label ("retain", "reflect", ...) for
LLM trace rows.
metadata: Optional extra caller metadata stored on trace rows.
Returns:
A ``ConfiguredLLMProvider`` that delegates to this provider with
the supplied config applied.
"""
trace_ctx = None
if bank_id is not None or operation is not None or metadata:
from .llm_trace import LLMTraceContext
# One trace + operation span per with_config() call — i.e. per
# operation invocation. Every LLM call made through this wrapper
# shares them, so a reflect/retain/consolidation run groups its
# calls as parent (operation) → children (LLM calls).
trace_ctx = LLMTraceContext(
bank_id=bank_id,
operation=operation,
metadata=dict(metadata or {}),
trace_id=str(uuid.uuid4()),
operation_span_id=str(uuid.uuid4()),
)
return ConfiguredLLMProvider(self, config.llm_gemini_safety_settings, trace_ctx)
return ConfiguredLLMProvider(self, config.llm_gemini_safety_settings)
async def cleanup(self) -> None:
"""Clean up resources (e.g. stop llamacpp subprocess)."""
@@ -1117,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,
)
@@ -1148,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,
)
@@ -1169,16 +891,10 @@ class ConfiguredLLMProvider:
any changes.
"""
def __init__(
self,
provider: "LLMProvider",
gemini_safety_settings: list | None,
trace_ctx: Any | None = None,
) -> None:
def __init__(self, provider: "LLMProvider", gemini_safety_settings: list | None) -> None:
# Use object.__setattr__ to avoid triggering __getattr__
object.__setattr__(self, "_provider", provider)
object.__setattr__(self, "_gemini_safety_settings", gemini_safety_settings)
object.__setattr__(self, "_trace_ctx", trace_ctx)
# ── attribute passthrough ──────────────────────────────────────────────────
@@ -1191,12 +907,10 @@ class ConfiguredLLMProvider:
from .providers.gemini_llm import _safety_settings_ctx
token = _safety_settings_ctx.set(object.__getattribute__(self, "_gemini_safety_settings"))
trace_token = self._bind_trace_context()
try:
return await object.__getattribute__(self, "_provider").call(messages=messages, **kwargs)
finally:
_safety_settings_ctx.reset(token)
self._reset_trace_context(trace_token)
async def call_with_tools(
self,
@@ -1207,38 +921,12 @@ class ConfiguredLLMProvider:
from .providers.gemini_llm import _safety_settings_ctx
token = _safety_settings_ctx.set(object.__getattribute__(self, "_gemini_safety_settings"))
trace_token = self._bind_trace_context()
try:
return await object.__getattribute__(self, "_provider").call_with_tools(
messages=messages, tools=tools, **kwargs
)
finally:
_safety_settings_ctx.reset(token)
self._reset_trace_context(trace_token)
def trace_context(self) -> Any | None:
"""The operation-level LLM trace context (or None when untraced).
Lets the engine attach the operation's produced/consumed memory_ids to
this run's trace rows once they're known (after the LLM calls).
"""
return object.__getattribute__(self, "_trace_ctx")
def _bind_trace_context(self) -> Any | None:
"""Bind bank/operation attribution for the duration of one call."""
trace_ctx = object.__getattribute__(self, "_trace_ctx")
if trace_ctx is None:
return None
from .llm_trace import set_trace_context
return set_trace_context(trace_ctx)
def _reset_trace_context(self, trace_token: Any | None) -> None:
if trace_token is None:
return
from .llm_trace import reset_trace_context
reset_trace_context(trace_token)
# Backwards compatibility alias
@@ -1,214 +0,0 @@
"""Background maintenance loop.
A single periodic loop that drives all of Hindsight's recurring housekeeping
from one place, so we don't spawn a separate ``asyncio`` task per concern:
- **Retention sweeps** (hourly): delete ``audit_log`` and ``llm_requests`` rows
older than their configured retention, across *all* tenant schemas.
- **Consolidation reconcile** (configurable, default 5 min): re-schedule
consolidation for banks that have eligible-but-unscheduled facts and no
in-flight consolidation. This recovers facts that were stranded when a
consolidation operation failed terminally and left them with
``consolidated_at IS NULL AND consolidation_failed_at IS NULL`` and nothing to
re-trigger them.
The loop wakes on a short fixed tick and runs each job when its own
``last_run + interval`` is due (run-at-start, then on interval), so adding jobs
with different cadences doesn't burst CPU. Cross-tenant discovery goes through
server-side PL/pgSQL routines (``public.schemas_with_expired_rows`` and
``public.banks_needing_consolidation``) — one round-trip each — instead of a
per-schema query storm, which matters at thousands of tenants.
"""
from __future__ import annotations
import asyncio
import logging
import time
from typing import TYPE_CHECKING
from ..config import HindsightConfig, get_config
from ..models import RequestContext
from .db_utils import acquire_with_retry
from .schema import _is_oracle
if TYPE_CHECKING:
from .memory_engine import MemoryEngine
logger = logging.getLogger(__name__)
# Short tick so jobs with different cadences share one loop without per-job tasks.
_TICK_SECONDS = 60
# Retention sweeps are not time-sensitive; hourly matches the previous per-sweep cadence.
_RETENTION_INTERVAL_SECONDS = 3600
class MaintenanceLoop:
"""Owns the single periodic maintenance task for a :class:`MemoryEngine`."""
def __init__(self, engine: "MemoryEngine") -> None:
self._engine = engine
self._task: asyncio.Task | None = None
self._stop = asyncio.Event()
# Monotonic timestamps of the last run per job, keyed by job name.
self._last_run: dict[str, float] = {}
# ── lifecycle ──────────────────────────────────────────────────────────
def start(self) -> None:
"""Start the loop if any maintenance job is enabled. Idempotent."""
if self._task and not self._task.done():
return
# PostgreSQL-only: the retention sweeps target PG-only tables (audit_log,
# llm_requests) and the reconcile relies on PG-only PL/pgSQL routines
# installed by the maintenance-routines migration. Oracle support is
# intentionally absent (mirrors that PG-only migration).
if _is_oracle():
logger.debug("Maintenance loop not started: PostgreSQL-only")
return
if not self._any_job_enabled():
logger.debug("Maintenance loop not started: no jobs enabled")
return
self._stop.clear()
try:
self._task = asyncio.create_task(self._run())
except RuntimeError:
logger.debug("Cannot start maintenance loop: no running event loop")
async def stop(self) -> None:
"""Stop the loop and wait for the current tick to finish."""
self._stop.set()
if self._task and not self._task.done():
try:
await self._task
except asyncio.CancelledError:
pass
self._task = None
@staticmethod
def _any_job_enabled() -> bool:
cfg = get_config()
reconcile_on = cfg.consolidation_reconcile_interval_seconds > 0
audit_on = cfg.audit_log_enabled and cfg.audit_log_retention_days > 0
llm_on = cfg.llm_trace_enabled and cfg.llm_trace_retention_days > 0
return reconcile_on or audit_on or llm_on
# ── loop ───────────────────────────────────────────────────────────────
async def _run(self) -> None:
while not self._stop.is_set():
try:
await self._tick()
except Exception:
logger.exception("Maintenance tick failed")
try:
await asyncio.wait_for(self._stop.wait(), timeout=_TICK_SECONDS)
except asyncio.TimeoutError:
pass
def _is_due(self, job: str, interval_seconds: int) -> bool:
"""True if ``job`` has never run or its interval has elapsed; marks it run now."""
now = time.monotonic()
last = self._last_run.get(job)
if last is not None and (now - last) < interval_seconds:
return False
self._last_run[job] = now
return True
async def _tick(self) -> None:
cfg = get_config()
if self._is_due("retention", _RETENTION_INTERVAL_SECONDS):
await self._run_retention(cfg)
interval = cfg.consolidation_reconcile_interval_seconds
if interval > 0 and self._is_due("reconcile", interval):
await self._run_reconcile()
# ── retention ──────────────────────────────────────────────────────────
async def _run_retention(self, cfg: HindsightConfig) -> None:
# Retention days are static server-level config, so one global cutoff
# applies to every tenant schema (the routine sweeps them all).
if cfg.audit_log_enabled and cfg.audit_log_retention_days > 0:
await self._purge_expired("audit_log", "started_at", cfg.audit_log_retention_days)
if cfg.llm_trace_enabled and cfg.llm_trace_retention_days > 0:
await self._purge_expired("llm_requests", "started_at", cfg.llm_trace_retention_days)
async def _purge_expired(self, table: str, ts_col: str, days: int) -> None:
"""Delete rows older than ``days`` from ``table`` across every tenant schema."""
backend = self._engine._backend
try:
async with acquire_with_retry(backend, max_retries=1) as conn:
rows = await conn.fetch(
"SELECT * FROM public.schemas_with_expired_rows($1, $2, $3)", table, ts_col, days
)
for row in rows:
schema = row[0]
# schema names come from pg_class; quote defensively all the same.
qschema = '"' + schema.replace('"', '""') + '"'
result = await conn.execute(
f"DELETE FROM {qschema}.{table} WHERE {ts_col} < NOW() - make_interval(days => $1)",
days,
)
if result and result != "DELETE 0":
logger.info(f"Retention sweep {schema}.{table}: {result}")
except Exception as e:
logger.warning(f"Retention sweep failed for {table}: {e}")
# ── consolidation reconcile ──────────────────────────────────────────────
async def _run_reconcile(self) -> None:
"""Re-schedule consolidation for banks with eligible-but-unscheduled facts."""
engine = self._engine
try:
async with acquire_with_retry(engine._backend, max_retries=1) as conn:
rows = await conn.fetch("SELECT schema_name, bank_id FROM public.banks_needing_consolidation()")
except Exception as e:
logger.warning(f"Consolidation reconcile discovery failed: {e}")
return
if not rows:
return
# Only enqueue into schemas the worker actually polls (tenant discovery),
# otherwise the op would never be claimed and would block future reconciles
# for that bank. The tenant_id (when the extension provides one) lets
# config resolution honor tenant-level overrides.
try:
tenants = await engine._tenant_extension.list_tenants()
except Exception as e:
logger.warning(f"Consolidation reconcile tenant discovery failed: {e}")
return
tenant_by_schema = {t.schema: t for t in tenants}
default_schema = get_config().database_schema
from .memory_engine import _current_schema
submitted = 0
skipped_unknown = 0
for row in rows:
schema = row["schema_name"]
bank_id = row["bank_id"]
tenant = tenant_by_schema.get(schema)
if tenant is None and schema != default_schema:
skipped_unknown += 1
continue
tenant_id = tenant.tenant_id if tenant else None
token = _current_schema.set(schema)
try:
context = RequestContext(internal=True, tenant_id=tenant_id)
resolved = await engine._config_resolver.resolve_full_config(bank_id, context)
# Mirror the retain-time auto-consolidation gate (memory_engine): both
# observations and auto-consolidation must be enabled for this bank.
if not (resolved.enable_observations and resolved.enable_auto_consolidation):
continue
await engine.submit_async_consolidation(bank_id=bank_id, request_context=context)
submitted += 1
except Exception as e:
logger.warning(f"Consolidation reconcile failed for bank {bank_id} in {schema}: {e}")
finally:
_current_schema.reset(token)
if submitted or skipped_unknown:
logger.info(
f"Consolidation reconcile: scheduled {submitted} bank(s)"
+ (f", skipped {skipped_unknown} in unrecognized schema(s)" if skipped_unknown else "")
)
File diff suppressed because it is too large Load Diff
@@ -5,10 +5,8 @@ These dataclasses define the structure of result_metadata for different operatio
The metadata is exposed in the API for debugging purposes and may change without notice.
"""
from dataclasses import asdict, dataclass, field
from typing import Any, Mapping
MAX_EXTRACTION_ERROR_SAMPLES = 5
from dataclasses import asdict, dataclass
from typing import Any
@dataclass
@@ -50,79 +48,6 @@ class RetainMetadata:
return asdict(self)
@dataclass
class RetainExtractionErrors:
"""Non-fatal fact extraction failures observed inside one retain operation."""
count: int = 0
sample: list[str] = field(default_factory=list)
def add(self, message: str) -> None:
"""Record one extraction error while keeping the stored sample bounded."""
self.count += 1
if len(self.sample) < MAX_EXTRACTION_ERROR_SAMPLES:
self.sample.append(message[:500])
def merge_metadata(self, metadata: Mapping[str, Any]) -> None:
"""Merge errors already present on an operation result_metadata object."""
self.count += int(metadata.get("extraction_errors_count") or 0)
sample = metadata.get("extraction_errors_sample") or []
if isinstance(sample, str):
sample = [sample]
if isinstance(sample, list):
for entry in sample:
if isinstance(entry, str) and len(self.sample) < MAX_EXTRACTION_ERROR_SAMPLES:
self.sample.append(entry[:500])
def to_dict(self) -> dict[str, Any]:
"""Convert to the public result_metadata field shape."""
data: dict[str, Any] = {"extraction_errors_count": self.count}
if self.sample:
data["extraction_errors_sample"] = self.sample
return data
@dataclass
class RetainOutcomeMetadata:
"""Machine-readable outcome metadata for a completed retain operation."""
unit_ids_count: int
extraction_errors_count: int = 0
extraction_errors_sample: list[str] = field(default_factory=list)
def to_dict(self) -> dict[str, Any]:
"""Convert to dict for JSON serialization, omitting empty optional samples."""
data: dict[str, Any] = {
"unit_ids_count": self.unit_ids_count,
"extraction_errors_count": self.extraction_errors_count,
}
if self.extraction_errors_sample:
data["extraction_errors_sample"] = self.extraction_errors_sample[:MAX_EXTRACTION_ERROR_SAMPLES]
return data
@dataclass
class RetainOutcomeAggregate:
"""Aggregate retain outcome metadata from child retain operations."""
unit_ids_count: int = 0
extraction_errors: RetainExtractionErrors = field(default_factory=RetainExtractionErrors)
def add_metadata(self, metadata: Mapping[str, Any]) -> None:
"""Fold one child operation's result_metadata into the aggregate."""
self.unit_ids_count += int(metadata.get("unit_ids_count") or 0)
self.extraction_errors.merge_metadata(metadata)
def to_outcome_metadata(self) -> RetainOutcomeMetadata:
"""Return the aggregate in the public result_metadata field shape."""
return RetainOutcomeMetadata(
unit_ids_count=self.unit_ids_count,
extraction_errors_count=self.extraction_errors.count,
extraction_errors_sample=self.extraction_errors.sample,
)
@dataclass
class ConsolidationMetadata:
"""Metadata for consolidation operations."""
@@ -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",
@@ -38,7 +38,6 @@ class AnthropicLLM(LLMInterface):
reasoning_effort: str = "low",
timeout: float = 300.0,
default_headers: dict[str, str] | None = None,
extra_body: dict[str, Any] | None = None,
**kwargs: Any,
):
"""
@@ -55,10 +54,6 @@ class AnthropicLLM(LLMInterface):
the Anthropic SDK client. Used by operators routing through proxies
or request-tracing middleware. Sourced from ``llm_default_headers`` in
``HindsightConfig`` (env: ``HINDSIGHT_API_LLM_DEFAULT_HEADERS``).
extra_body: Extra request-body params (e.g. ``{"temperature": 0.2,
"top_p": 0.9, "top_k": 40}``) passed via the Anthropic SDK's
``extra_body`` so they merge into the JSON sent to the Messages API.
Sourced from ``llm_extra_body`` (env: ``HINDSIGHT_API_LLM_EXTRA_BODY``).
**kwargs: Additional provider-specific parameters.
"""
super().__init__(provider, api_key, base_url, model, reasoning_effort, **kwargs)
@@ -66,9 +61,6 @@ class AnthropicLLM(LLMInterface):
if not self.api_key:
raise ValueError("API key is required for Anthropic provider")
# User-configured extra body params (merged into every Messages API call)
self._extra_body = extra_body or {}
# Import and initialize Anthropic client
try:
from anthropic import AsyncAnthropic
@@ -101,6 +93,7 @@ class AnthropicLLM(LLMInterface):
await self.call(
messages=test_messages,
max_completion_tokens=10,
temperature=0.0,
scope="verification",
max_retries=0,
)
@@ -186,8 +179,8 @@ class AnthropicLLM(LLMInterface):
if system_prompt:
call_params["system"] = system_prompt
if self._extra_body:
call_params["extra_body"] = self._extra_body
if temperature is not None:
call_params["temperature"] = temperature
last_exception = None
@@ -227,7 +220,6 @@ class AnthropicLLM(LLMInterface):
input_tokens = response.usage.input_tokens or 0 if response.usage else 0
output_tokens = response.usage.output_tokens or 0 if response.usage else 0
total_tokens = input_tokens + output_tokens
cached_tokens = getattr(response.usage, "cache_read_input_tokens", 0) or 0 if response.usage else 0
# Record LLM metrics
metrics = get_metrics_collector()
@@ -257,7 +249,6 @@ class AnthropicLLM(LLMInterface):
duration=duration,
finish_reason=finish_reason,
error=None,
cached_tokens=cached_tokens,
)
# Log slow calls
@@ -273,7 +264,6 @@ class AnthropicLLM(LLMInterface):
input_tokens=input_tokens,
output_tokens=output_tokens,
total_tokens=total_tokens,
cached_tokens=cached_tokens,
)
return result, token_usage
return result
@@ -408,8 +398,8 @@ class AnthropicLLM(LLMInterface):
if system_prompt:
call_params["system"] = system_prompt
if self._extra_body:
call_params["extra_body"] = self._extra_body
if temperature is not None:
call_params["temperature"] = temperature
last_exception = None
for attempt in range(max_retries + 1):
@@ -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
@@ -1,316 +0,0 @@
"""Gemini context-cache manager.
Wraps the ``google-genai`` SDK's CachedContent API to let callers reuse a
stable system_instruction + response_schema prefix across many requests.
Cached input tokens are billed at ~10× lower than fresh input tokens
(check the current Gemini pricing for the exact ratio per model), so for
workloads that repeatedly send a large fixed prefix with a small variable
user message fact extraction, structured tagging, classification the
input-cost savings are substantial.
This module owns only the create/refresh/lookup lifecycle. It is up to
the caller to (a) decide that the prefix is stable enough to cache, and
(b) pass the returned cache name to ``GeminiLLM.call()``. When the
returned name is ``None`` (because Gemini rejected the create most
commonly because the prefix is smaller than the model's minimum), the
caller MUST fall back to a non-cached call.
Cardinality
-----------
The intended cache count per process is small (100 entries). Each
entry corresponds to one combination of (model, system_instruction,
response_schema). If a caller sees the cache grow unboundedly it
indicates the system_instruction contains per-request data that should
move into the user message instead.
TTL
---
Gemini's CachedContent has a TTL bounded by the model (currently 1h
for most generally-available models). This manager refreshes proactively
at ``ttl_safety_margin`` before expiry. If a cached entry has expired
between refreshes the next call will recreate it transparently.
"""
from __future__ import annotations
import asyncio
import hashlib
import json
import logging
import time
from dataclasses import dataclass
from typing import Any
logger = logging.getLogger(__name__)
# Default TTL: 55 minutes. Gemini's hard max for CachedContent is 1 hour
# for most models; we refresh 5 minutes early so a request landing right
# at the boundary doesn't race against expiry.
_DEFAULT_TTL_SECONDS = 55 * 60
_DEFAULT_REFRESH_MARGIN_SECONDS = 5 * 60
# Cap on the cache-create network call. It runs while holding the manager lock, so
# a hung create would block every concurrent caller (e.g. all chunks of a 10-chunk
# retain batch waiting on the cold-start create). On timeout the create soft-fails
# to None and callers proceed uncached, rather than stalling the whole batch.
_DEFAULT_CREATE_TIMEOUT_SECONDS = 30.0
@dataclass
class _CacheEntry:
name: str # The CachedContent resource name returned by Gemini.
created_at: float
ttl_seconds: int
class GeminiCacheManager:
"""Per-process map of (prefix fingerprint) → CachedContent name.
Thread-safe across asyncio tasks via a single ``asyncio.Lock``. The
create/refresh calls are serialised; this is fine because cache
creation is a one-shot warm-up per fingerprint (subsequent reads are
pure dict lookups outside the lock).
Not shared across pods each worker / api replica builds its own
cache. The cost of cold-starting one extra full-price call per pod
per fingerprint per hour is negligible compared to the steady-state
savings.
"""
def __init__(
self,
client: Any,
*,
ttl_seconds: int = _DEFAULT_TTL_SECONDS,
refresh_margin_seconds: int = _DEFAULT_REFRESH_MARGIN_SECONDS,
create_timeout_seconds: float = _DEFAULT_CREATE_TIMEOUT_SECONDS,
) -> None:
self._client = client
self._ttl_seconds = ttl_seconds
self._refresh_margin_seconds = refresh_margin_seconds
self._create_timeout_seconds = create_timeout_seconds
self._entries: dict[str, _CacheEntry] = {}
self._lock = asyncio.Lock()
@staticmethod
def fingerprint(
model: str,
system_instruction: str,
response_schema: Any | None = None,
tools: list[dict[str, Any]] | None = None,
) -> str:
"""Stable hash of the cacheable surface.
``response_schema`` may be a Pydantic class, a dict, or ``None``.
Pydantic schemas are normalised by serialising via
``model_json_schema()`` and stripping the auto-generated
``"title"`` fields so two dynamically-built models with the same
shape but different class names hash identically. This matters
for callers (e.g. fact extraction) that rebuild the schema
class on every request via a builder helper without the
normalisation the cache would never hit.
``tools`` is the OpenAI-style tools list (each entry has a
``"function"`` dict with name/description/parameters). When
supplied, the tool definitions become part of the cache key so a
loop that adds or renames a tool gets a fresh cache and doesn't
silently use a stale schema. Tools are serialised with
``sort_keys=True`` to neutralise dict-ordering drift.
"""
hasher = hashlib.sha256()
hasher.update(model.encode("utf-8"))
hasher.update(b"\x00")
hasher.update(system_instruction.encode("utf-8"))
hasher.update(b"\x00")
if response_schema is None:
hasher.update(b"none")
elif hasattr(response_schema, "model_json_schema"):
try:
schema = response_schema.model_json_schema()
_strip_titles(schema)
hasher.update(json.dumps(schema, sort_keys=True).encode("utf-8"))
except Exception:
# Fall back to class identity if the schema can't be serialised.
hasher.update(repr(response_schema).encode("utf-8"))
else:
try:
hasher.update(json.dumps(response_schema, sort_keys=True).encode("utf-8"))
except (TypeError, ValueError):
hasher.update(repr(response_schema).encode("utf-8"))
hasher.update(b"\x00")
if tools:
try:
hasher.update(json.dumps(tools, sort_keys=True).encode("utf-8"))
except (TypeError, ValueError):
hasher.update(repr(tools).encode("utf-8"))
else:
hasher.update(b"no-tools")
return hasher.hexdigest()
async def get_or_create(
self,
*,
model: str,
system_instruction: str,
response_schema: Any | None = None,
tools: list[dict[str, Any]] | None = None,
) -> str | None:
"""Return a CachedContent resource name for the given prefix, or
``None`` if Gemini rejects the create (prefix too small, model
does not support caching, etc.).
``tools`` is the OpenAI-style tools list. When supplied, the tool
definitions are baked into the CachedContent so the caller's
``call_with_tools`` doesn't need to resend them on every
iteration. Pass ``None`` for non-tool calls.
``None`` return is a normal, expected value the caller falls
back to an uncached call and the system continues to work.
"""
key = self.fingerprint(model, system_instruction, response_schema, tools)
async with self._lock:
entry = self._entries.get(key)
if entry is not None and self._is_fresh(entry):
return entry.name
# Need to (re)create. Pop the stale entry first so a failed
# create doesn't leave a name we'd return on the next call.
self._entries.pop(key, None)
try:
cache_name = await self._create_cache(
model=model,
system_instruction=system_instruction,
tools=tools,
)
except _CacheNotEligible as e:
logger.debug(
"GeminiCacheManager: prefix not eligible for caching (model=%s, reason=%s) — caller will fall back",
model,
e,
)
return None
except Exception:
logger.exception(
"GeminiCacheManager: failed to create cached content "
"(model=%s); caller will fall back to uncached call",
model,
)
return None
if cache_name is None:
return None
self._entries[key] = _CacheEntry(
name=cache_name,
created_at=time.monotonic(),
ttl_seconds=self._ttl_seconds,
)
return cache_name
def _is_fresh(self, entry: _CacheEntry) -> bool:
"""An entry is fresh if it's young enough that the next request
won't race against the TTL expiry."""
age = time.monotonic() - entry.created_at
return age < (entry.ttl_seconds - self._refresh_margin_seconds)
def invalidate(self, name: str) -> None:
"""Forget a cache name that the server rejected (expired/deleted/invalid).
Called by the provider when a generate request using this CachedContent
fails, so the next ``get_or_create`` recreates it instead of handing back
the dead name again. Best-effort and sync drops the matching entry from
the in-process map; the orphaned server-side cache (if any) ages out on
its own TTL.
"""
for key, entry in list(self._entries.items()):
if entry.name == name:
self._entries.pop(key, None)
async def _create_cache(
self,
*,
model: str,
system_instruction: str,
tools: list[dict[str, Any]] | None = None,
) -> str | None:
"""Wrap ``client.aio.caches.create`` with the config we want.
The SDK surface differs slightly across google-genai versions;
this implementation targets the >=1.0.0 line where caches live
under ``client.aio.caches``.
"""
# Lazy import so this module doesn't require the SDK at import time.
from google.genai import types as genai_types
# A CachedContent only holds reusable *input* — system_instruction,
# contents, tools, ttl. ``response_schema``/``response_mime_type`` are
# generation-time output constraints and the SDK rejects them here
# (``CreateCachedContentConfig`` forbids those fields). They are applied
# per-request on the GenerateContentConfig instead — see the call sites,
# which set them alongside ``cached_content``. ``response_schema`` is
# still part of the fingerprint so a schema change keys a fresh cache.
config_kwargs: dict[str, Any] = {
"system_instruction": system_instruction,
"ttl": f"{self._ttl_seconds}s",
}
if tools:
# OpenAI-style {"function": {...}} entries must be converted to
# Gemini's Tool/FunctionDeclaration shape before caching.
gemini_tools = []
for tool in tools:
func = tool.get("function", {})
gemini_tools.append(
genai_types.Tool(
function_declarations=[
genai_types.FunctionDeclaration(
name=func.get("name", ""),
description=func.get("description", ""),
parameters=func.get("parameters"),
)
]
)
)
config_kwargs["tools"] = gemini_tools
try:
cached = await asyncio.wait_for(
self._client.aio.caches.create(
model=model,
config=genai_types.CreateCachedContentConfig(**config_kwargs),
),
timeout=self._create_timeout_seconds,
)
except Exception as e:
# Gemini returns a 400 with a "minimum token count" message
# when the prefix is too small. We treat this as a soft
# "not eligible" signal rather than a real error so callers
# silently fall back to non-cached.
msg = str(e).lower()
if "minimum" in msg or "too small" in msg or "too short" in msg:
raise _CacheNotEligible(str(e)) from e
raise
return getattr(cached, "name", None)
class _CacheNotEligible(Exception):
"""Raised when Gemini rejects the cache create because the prefix
is below the model's minimum cacheable size. Treated as a soft
fallback by the caller, not an error."""
def _strip_titles(node: Any) -> None:
"""Recursively remove auto-generated ``"title"`` keys from a JSON
Schema-like dict tree, in place. Pydantic seeds these from the
Python class name, which means structurally-identical schemas built
from differently-named classes look distinct to a naive hash."""
if isinstance(node, dict):
node.pop("title", None)
for v in node.values():
_strip_titles(v)
elif isinstance(node, list):
for item in node:
_strip_titles(item)
@@ -70,22 +70,6 @@ class GeminiLLM(LLMInterface):
# Safety settings: None means use Gemini's defaults
self._safety_settings: list | None = kwargs.get("gemini_safety_settings")
# User-configured extra params merged into the GenerateContentConfig of
# every call. Gemini's request body nests generation params, so we expose
# them in the SDK's native config space rather than as a raw body merge:
# keys must be GenerateContentConfig fields (e.g. temperature, top_p,
# top_k, max_output_tokens, seed). Sourced from llm_extra_body
# (env: HINDSIGHT_API_LLM_EXTRA_BODY).
self._extra_body: dict[str, Any] = kwargs.get("extra_body") or {}
# Context-cache manager. Lazy-initialized on first cache lookup so
# nothing happens for models/workloads that never reach it. The instance
# default here is off (a directly-constructed GeminiLLM doesn't cache); the
# server-level default is on and flows in via the prompt_cache_enabled kwarg
# resolved from config in LLMProvider.
self._cache_manager: Any | None = None
self._prompt_cache_enabled: bool = bool(kwargs.get("prompt_cache_enabled", False))
if self._is_vertexai:
self._init_vertexai(**kwargs)
else:
@@ -184,7 +168,6 @@ class GeminiLLM(LLMInterface):
skip_validation: bool = False,
strict_schema: bool = False,
return_usage: bool = False,
cached_prefix: str | None = None,
) -> Any:
"""
Make a Gemini/VertexAI API call with retry logic.
@@ -199,17 +182,8 @@ class GeminiLLM(LLMInterface):
initial_backoff: Initial backoff time in seconds.
max_backoff: Maximum backoff time in seconds.
skip_validation: Return raw JSON without Pydantic validation.
strict_schema: Ignored Gemini always grammar-enforces structured output via its
native response_schema, so it is strict regardless of this flag.
strict_schema: Use strict JSON schema enforcement (not supported by Gemini).
return_usage: If True, return tuple (result, TokenUsage).
cached_prefix: Optional CachedContent resource name (from
``GeminiCacheManager.get_or_create``). When set, the
system_instruction is assumed to live in the cache; this call
skips resending it and the cached prefix is billed at the
cached-input rate instead of the standard input rate. The
response_schema is still sent per-request (it is not cacheable).
Pass ``None`` to use the
normal uncached path.
Returns:
If return_usage=False: Parsed response if response_format provided, else text.
@@ -217,14 +191,9 @@ class GeminiLLM(LLMInterface):
"""
start_time = time.time()
# Convert OpenAI-style messages to Gemini format. We ALWAYS build
# system_instruction (even when a cache is in use): the config builder
# below omits it from the request while the cache carries the prefix, but
# it must be available so the cached-call-failed safety net can re-send it
# inline. Whether it's actually sent is decided in _build_generation_config.
# Convert OpenAI-style messages to Gemini format
system_instruction = None
gemini_contents = []
using_cache = cached_prefix is not None
for msg in messages:
role = msg.get("role", "user")
@@ -240,9 +209,7 @@ class GeminiLLM(LLMInterface):
else:
gemini_contents.append(genai_types.Content(role="user", parts=[genai_types.Part(text=content)]))
# Add the JSON schema as a textual hint in the system_instruction (matching
# the normal uncached path). Structured output is still enforced via
# response_schema regardless; this is just guidance text.
# Add JSON schema instruction if response_format is provided
if response_format is not None and hasattr(response_format, "model_json_schema"):
schema = response_format.model_json_schema()
schema_msg = f"\n\nYou must respond with valid JSON matching this schema:\n{json.dumps(schema, indent=2, ensure_ascii=False)}"
@@ -251,44 +218,32 @@ class GeminiLLM(LLMInterface):
else:
system_instruction = schema_msg
# Build generation config
config_kwargs: dict[str, Any] = {}
if system_instruction:
config_kwargs["system_instruction"] = system_instruction
if response_format is not None:
config_kwargs["response_mime_type"] = "application/json"
config_kwargs["response_schema"] = response_format
if temperature is not None:
config_kwargs["temperature"] = temperature
# Gemini's equivalent of OpenAI-style max_completion_tokens is max_output_tokens.
# Without it the model can produce arbitrarily long responses, ignoring the
# caller's intended cap (e.g. mental_models max_tokens during refresh).
if max_completion_tokens is not None:
config_kwargs["max_output_tokens"] = max_completion_tokens
# Apply safety settings: context var (per-request bank override) takes precedence over instance default
effective_safety_settings = _safety_settings_ctx.get()
if effective_safety_settings is None:
effective_safety_settings = self._safety_settings
if effective_safety_settings is not None:
config_kwargs["safety_settings"] = [
genai_types.SafetySetting(category=s["category"], threshold=s["threshold"])
for s in effective_safety_settings
]
# Build generation config. ``cached_content`` and ``system_instruction``
# are mutually exclusive (the cache IS the prefix; the SDK rejects
# re-sending it). ``response_schema``/``response_mime_type`` are
# request-level output constraints — NOT cacheable — so they're set on
# every structured call, including cached ones where they ride alongside
# ``cached_content``. Built as a closure so we can rebuild it WITHOUT the
# cache and retry inline if a stale/invalid CachedContent makes the call fail.
def _build_generation_config(use_cache: bool) -> "genai_types.GenerateContentConfig | None":
# Seed with user-configured extra params; explicit settings below win.
config_kwargs: dict[str, Any] = dict(self._extra_body)
if use_cache:
config_kwargs["cached_content"] = cached_prefix
elif system_instruction:
config_kwargs["system_instruction"] = system_instruction
if response_format is not None:
config_kwargs["response_mime_type"] = "application/json"
config_kwargs["response_schema"] = response_format
if temperature is not None:
config_kwargs["temperature"] = temperature
# Gemini's equivalent of OpenAI-style max_completion_tokens is max_output_tokens.
# Without it the model can produce arbitrarily long responses, ignoring the
# caller's intended cap (e.g. mental_models max_tokens during refresh).
if max_completion_tokens is not None:
config_kwargs["max_output_tokens"] = max_completion_tokens
if effective_safety_settings is not None:
config_kwargs["safety_settings"] = [
genai_types.SafetySetting(category=s["category"], threshold=s["threshold"])
for s in effective_safety_settings
]
return genai_types.GenerateContentConfig(**config_kwargs) if config_kwargs else None
cache_active = using_cache
generation_config = _build_generation_config(cache_active)
generation_config = genai_types.GenerateContentConfig(**config_kwargs) if config_kwargs else None
last_exception = None
@@ -333,24 +288,13 @@ class GeminiLLM(LLMInterface):
else:
result = content
# Extract token usage. ``cached_content_token_count`` and
# ``thoughts_token_count`` are populated on the Gemini 2.5+
# family; treat missing fields as 0 so older models still
# record sensible metrics.
# Extract token usage
input_tokens = 0
output_tokens = 0
cached_input_tokens = 0
thoughts_tokens = 0
cached_tokens = 0
if hasattr(response, "usage_metadata") and response.usage_metadata:
usage = response.usage_metadata
input_tokens = usage.prompt_token_count or 0
output_tokens = usage.candidates_token_count or 0
cached_input_tokens = getattr(usage, "cached_content_token_count", 0) or 0
thoughts_tokens = getattr(usage, "thoughts_token_count", 0) or 0
# Tracing/TokenUsage consume ``cached_tokens``; metrics consume
# ``cached_input_tokens`` — same value, two downstream names.
cached_tokens = cached_input_tokens
# Record metrics
duration = time.time() - start_time
@@ -363,8 +307,6 @@ class GeminiLLM(LLMInterface):
input_tokens=input_tokens,
output_tokens=output_tokens,
success=True,
cached_input_tokens=cached_input_tokens,
thoughts_tokens=thoughts_tokens,
)
# Record trace span
@@ -388,7 +330,6 @@ class GeminiLLM(LLMInterface):
duration=duration,
finish_reason=finish_reason,
error=None,
cached_tokens=cached_tokens,
)
# Log slow calls
@@ -404,7 +345,6 @@ class GeminiLLM(LLMInterface):
input_tokens=input_tokens,
output_tokens=output_tokens,
total_tokens=input_tokens + output_tokens,
cached_tokens=cached_tokens,
)
return result, token_usage
return result
@@ -426,20 +366,6 @@ class GeminiLLM(LLMInterface):
logger.error(f"Gemini auth error (HTTP {e.code}), not retrying: {str(e)}")
raise
# Cached-request safety net: a stale/invalid/expired CachedContent
# (or an incompatibility like cache + tool_config) surfaces as a 400.
# Retrying the same cached request can't recover, so on the first
# such failure drop the cache, invalidate it so later operations
# recreate it, and retry THIS call inline with the prefix inlined.
# Caching must never break a request.
if cache_active and e.code == 400:
logger.warning(f"Gemini cached call failed (400); retrying uncached. Reason: {str(e)}")
if self._cache_manager is not None and cached_prefix is not None:
self._cache_manager.invalidate(cached_prefix)
cache_active = False
generation_config = _build_generation_config(cache_active)
continue
# Retry on retryable errors (rate limits, server errors, client errors)
if e.code in (400, 429, 500, 502, 503, 504) or (e.code and e.code >= 500):
last_exception = e
@@ -473,7 +399,6 @@ class GeminiLLM(LLMInterface):
initial_backoff: float = 1.0,
max_backoff: float = 30.0,
tool_choice: str | dict[str, Any] = "auto",
cached_prefix: str | None = None,
) -> LLMToolCallResult:
"""
Make a Gemini/VertexAI API call with tool/function calling support.
@@ -488,39 +413,27 @@ class GeminiLLM(LLMInterface):
initial_backoff: Initial backoff time in seconds.
max_backoff: Maximum backoff time in seconds.
tool_choice: How to choose tools (Gemini uses "auto" only).
cached_prefix: Optional CachedContent resource name (from
``GeminiCacheManager.get_or_create`` with ``tools=...``). When
set, the system_instruction and tool definitions are assumed
to live in the cache; this call will skip resending them and
the cached prefix is billed at the cached-input rate. The
``tools`` argument is still required (the caller may pass
an empty list when the cache holds them) so existing call
sites don't break.
Returns:
LLMToolCallResult with content and/or tool_calls.
"""
start_time = time.time()
using_cache = cached_prefix is not None
# Convert tools to Gemini format. When the cache is in use, the
# tool definitions are baked into the CachedContent at create time
# and the SDK rejects re-sending them alongside ``cached_content``.
# Convert tools to Gemini format
gemini_tools = []
if not using_cache:
for tool in tools:
func = tool.get("function", {})
gemini_tools.append(
genai_types.Tool(
function_declarations=[
genai_types.FunctionDeclaration(
name=func.get("name", ""),
description=func.get("description", ""),
parameters=func.get("parameters"),
)
]
)
for tool in tools:
func = tool.get("function", {})
gemini_tools.append(
genai_types.Tool(
function_declarations=[
genai_types.FunctionDeclaration(
name=func.get("name", ""),
description=func.get("description", ""),
parameters=func.get("parameters"),
)
]
)
)
# Convert messages
system_instruction = None
@@ -533,10 +446,6 @@ class GeminiLLM(LLMInterface):
content = msg.get("content", "")
if role == "system":
# Always capture system_instruction. _build_tools_config omits it
# (and tools) from the request while the cache carries the prefix,
# but it must be available so the cached-call-failed safety net can
# re-send the prefix + tools inline.
system_instruction = (system_instruction + "\n\n" + content) if system_instruction else content
i += 1
elif role == "tool":
@@ -584,63 +493,49 @@ class GeminiLLM(LLMInterface):
gemini_contents.append(genai_types.Content(role="user", parts=[genai_types.Part(text=content)]))
i += 1
config_kwargs: dict[str, Any] = {"tools": gemini_tools}
if system_instruction:
config_kwargs["system_instruction"] = system_instruction
if temperature is not None:
config_kwargs["temperature"] = temperature
# See note in `call`: Gemini's max_output_tokens is the equivalent of
# OpenAI-style max_completion_tokens.
if max_completion_tokens is not None:
config_kwargs["max_output_tokens"] = max_completion_tokens
# Map OpenAI-style tool_choice to Gemini FunctionCallingConfig
if tool_choice == "required":
config_kwargs["tool_config"] = genai_types.ToolConfig(
function_calling_config=genai_types.FunctionCallingConfig(
mode="ANY",
)
)
elif isinstance(tool_choice, dict) and tool_choice.get("type") == "function":
fn_name = tool_choice.get("function", {}).get("name")
if fn_name:
config_kwargs["tool_config"] = genai_types.ToolConfig(
function_calling_config=genai_types.FunctionCallingConfig(
mode="ANY",
allowed_function_names=[fn_name],
)
)
elif tool_choice == "none":
config_kwargs["tool_config"] = genai_types.ToolConfig(
function_calling_config=genai_types.FunctionCallingConfig(mode="NONE")
)
# "auto" is the default (no tool_config needed)
# Apply safety settings: context var (per-request bank override) takes precedence over instance default
effective_safety_settings = _safety_settings_ctx.get()
if effective_safety_settings is None:
effective_safety_settings = self._safety_settings
if effective_safety_settings is not None:
config_kwargs["safety_settings"] = [
genai_types.SafetySetting(category=s["category"], threshold=s["threshold"])
for s in effective_safety_settings
]
# When using a cached prefix, the SDK rejects re-sending system_instruction
# or tools alongside ``cached_content`` — the cache IS the prefix.
# tool_config (mode / allowed_function_names) is a per-request decision and
# stays out of the cache. Built as a closure so we can rebuild it WITHOUT
# the cache and retry inline if a stale/invalid cache makes the call fail.
def _build_tools_config(use_cache: bool) -> "genai_types.GenerateContentConfig":
# Seed with user-configured extra params; explicit settings below win.
config_kwargs: dict[str, Any] = dict(self._extra_body)
if use_cache:
config_kwargs["cached_content"] = cached_prefix
else:
config_kwargs["tools"] = gemini_tools
if system_instruction:
config_kwargs["system_instruction"] = system_instruction
if temperature is not None:
config_kwargs["temperature"] = temperature
# See note in `call`: Gemini's max_output_tokens is the equivalent of
# OpenAI-style max_completion_tokens.
if max_completion_tokens is not None:
config_kwargs["max_output_tokens"] = max_completion_tokens
# Map OpenAI-style tool_choice to Gemini FunctionCallingConfig
if tool_choice == "required":
config_kwargs["tool_config"] = genai_types.ToolConfig(
function_calling_config=genai_types.FunctionCallingConfig(
mode="ANY",
)
)
elif isinstance(tool_choice, dict) and tool_choice.get("type") == "function":
fn_name = tool_choice.get("function", {}).get("name")
if fn_name:
config_kwargs["tool_config"] = genai_types.ToolConfig(
function_calling_config=genai_types.FunctionCallingConfig(
mode="ANY",
allowed_function_names=[fn_name],
)
)
elif tool_choice == "none":
config_kwargs["tool_config"] = genai_types.ToolConfig(
function_calling_config=genai_types.FunctionCallingConfig(mode="NONE")
)
# "auto" is the default (no tool_config needed)
if effective_safety_settings is not None:
config_kwargs["safety_settings"] = [
genai_types.SafetySetting(category=s["category"], threshold=s["threshold"])
for s in effective_safety_settings
]
return genai_types.GenerateContentConfig(**config_kwargs)
cache_active = using_cache
config = _build_tools_config(cache_active)
config = genai_types.GenerateContentConfig(**config_kwargs)
last_exception = None
for attempt in range(max_retries + 1):
@@ -683,18 +578,12 @@ class GeminiLLM(LLMInterface):
finish_reason = "tool_calls" if tool_calls else "stop"
# Extract token usage. ``cached_content_token_count`` and
# ``thoughts_token_count`` are populated on the Gemini 2.5+
# family; absent fields are treated as 0.
# Extract token usage
input_tokens = 0
output_tokens = 0
cached_input_tokens = 0
thoughts_tokens = 0
if response.usage_metadata:
input_tokens = response.usage_metadata.prompt_token_count or 0
output_tokens = response.usage_metadata.candidates_token_count or 0
cached_input_tokens = getattr(response.usage_metadata, "cached_content_token_count", 0) or 0
thoughts_tokens = getattr(response.usage_metadata, "thoughts_token_count", 0) or 0
# Record metrics
duration = time.time() - start_time
@@ -707,8 +596,6 @@ class GeminiLLM(LLMInterface):
input_tokens=input_tokens,
output_tokens=output_tokens,
success=True,
cached_input_tokens=cached_input_tokens,
thoughts_tokens=thoughts_tokens,
)
# Record OpenTelemetry span
@@ -733,7 +620,6 @@ class GeminiLLM(LLMInterface):
finish_reason=finish_reason,
error=None,
tool_calls=tool_calls_dict,
cached_tokens=cached_input_tokens,
)
return LLMToolCallResult(
@@ -750,18 +636,6 @@ class GeminiLLM(LLMInterface):
logger.error(f"Gemini auth error (HTTP {e.code}), not retrying: {str(e)}")
raise
# Cached-request safety net (see ``call``): a stale/invalid cache or
# a cache+tool_config conflict surfaces as a 400. Drop the cache,
# invalidate it for later operations, and retry THIS call inline
# with the prefix + tools re-sent. Caching must never break a call.
if cache_active and e.code == 400:
logger.warning(f"Gemini cached tool call failed (400); retrying uncached. Reason: {str(e)}")
if self._cache_manager is not None and cached_prefix is not None:
self._cache_manager.invalidate(cached_prefix)
cache_active = False
config = _build_tools_config(cache_active)
continue
# Retry on retryable errors
last_exception = e
if attempt < max_retries:
@@ -778,54 +652,6 @@ class GeminiLLM(LLMInterface):
raise last_exception
raise RuntimeError("Gemini tool call failed")
def supports_prompt_caching(self) -> bool:
"""True when explicit Gemini context caching is enabled for this instance.
Reflects the opt-in flag so callers skip the cache lookup entirely when
it's off; ``get_or_create_cached_prefix`` also returns None in that case.
"""
return self._prompt_cache_enabled
async def get_or_create_cached_prefix(
self,
*,
system_instruction: str,
response_schema: Any | None = None,
tools: list[dict[str, Any]] | None = None,
) -> str | None:
"""Return a CachedContent resource name for the given prefix, or
``None`` if context caching is disabled, the provider doesn't
support it, or Gemini rejects the create (prefix too small, etc.).
``tools`` is the OpenAI-style tools list; pass it when caching a
prefix that will be used by ``call_with_tools()``. The fingerprint
includes the tool definitions so a loop that swaps a tool gets a
fresh cache automatically.
Callers pass the returned name to ``call(cached_prefix=...)``
or ``call_with_tools(cached_prefix=...)`` and treat ``None``
as "cache unavailable — use the normal path". That fallback is
essential: the system must continue to work if caching is disabled,
if Gemini's caching API has an outage, or if the prefix is below
the model's minimum cacheable size.
"""
if not self._prompt_cache_enabled:
return None
if self._client is None:
return None
if self._cache_manager is None:
# Lazy import so the cache module is only loaded when caching
# is actually used.
from hindsight_api.engine.providers.gemini_cache import GeminiCacheManager
self._cache_manager = GeminiCacheManager(self._client)
return await self._cache_manager.get_or_create(
model=self.model,
system_instruction=system_instruction,
response_schema=response_schema,
tools=tools,
)
async def cleanup(self) -> None:
"""Clean up resources (close connections, etc.)."""
# Gemini client doesn't require explicit cleanup
@@ -48,18 +48,11 @@ class LiteLLMLLM(LLMInterface):
model: str,
reasoning_effort: str = "low",
timeout: float = 300.0,
extra_body: dict[str, Any] | None = None,
**kwargs: Any,
):
super().__init__(provider, api_key, base_url, model, reasoning_effort, **kwargs)
self.timeout = timeout
self._litellm: Any = None
# User-configured extra params merged as top-level kwargs into every
# completion call so LiteLLM normalizes them per-provider (e.g. maps
# temperature/top_p/max_tokens across OpenAI, Anthropic, Bedrock, …) and
# drops any the target model rejects (litellm.drop_params=True below).
# Sourced from llm_extra_body (env: HINDSIGHT_API_LLM_EXTRA_BODY).
self._extra_body: dict[str, Any] = extra_body or {}
try:
import litellm
@@ -114,11 +107,6 @@ class LiteLLMLLM(LLMInterface):
if temperature is not None:
kwargs["temperature"] = temperature
# User-configured extras fill in only where the caller didn't set a value,
# so explicit per-call params (model, messages, temperature, …) always win.
for key, value in self._extra_body.items():
kwargs.setdefault(key, value)
return kwargs
# ── per-model output-tokens cap (shared with Router subclass) ────────────
@@ -153,30 +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 == "consolidation_dedup" and response_format is not None:
# Observation dedup adjudication. Default to "keep" so mock-LLM consolidation never
# spuriously merges observations — this preserves the pre-dedup behaviour that
# deterministic consolidation tests assert (the generic branch below can't construct
# the model because its "action" field is required and has no default).
result = response_format(action="keep", reason="mock")
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:
@@ -262,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 = (
@@ -291,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
@@ -429,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
@@ -7,7 +7,7 @@ This provider handles all OpenAI API-compatible models including:
- Groq: Fast inference with seed control and service tiers
- Ollama: Local models with native streaming API support
- LMStudio: Local models with OpenAI-compatible API
- MiniMax: MiniMax-M3 / MiniMax-M2.7 models with 1M context window
- MiniMax: MiniMax-M2.7 models with 1M context window
- DeepSeek: deepseek-v4-flash / deepseek-v4-pro / deepseek-chat / deepseek-reasoner via api.deepseek.com
- Opencode Go: deepseek-v4-flash via https://opencode.ai/zen/go/v1
@@ -44,16 +44,6 @@ logger = logging.getLogger(__name__)
DEFAULT_LLM_SEED = 4242
JSON_MODE_USER_HINT = "Return valid json only."
# Self-hosted OpenAI-compatible servers that advertise tool_choice="required"
# but silently ignore it: instead of forcing a tool call they return
# finish_reason "stop"/"tool_calls" with an EMPTY tool_calls array and no error.
# Reflect's agent loop then sees no tool call, runs synthesis with no retrieval,
# and answers "I don't have information" even when the bank holds the answer.
# See issues #1563 (LM Studio), #1179 (LM Studio + Qwen), #1877 (vLLM with
# --enable-auto-tool-choice). llama-server (the "llamacpp" provider) honors
# "required" correctly and is intentionally excluded (#1179).
_TOOL_CHOICE_REQUIRED_UNSUPPORTED_PROVIDERS = frozenset({"lmstudio", "ollama"})
class ProviderResponseError(RuntimeError):
"""Raised when a provider returns a success response without usable content."""
@@ -242,7 +232,7 @@ class OpenAICompatibleLLM(LLMInterface):
- Groq: Fast inference with seed control and service tiers
- Ollama: Local models with native streaming API for better structured output
- LMStudio: Local models with OpenAI-compatible API
- MiniMax: MiniMax-M3 / MiniMax-M2.7 models via OpenAI-compatible API (https://api.minimax.io/v1)
- MiniMax: MiniMax-M2.7 models via OpenAI-compatible API (https://api.minimax.io/v1)
- DeepSeek: deepseek-v4-flash / deepseek-v4-pro / deepseek-chat / deepseek-reasoner via https://api.deepseek.com
- opencode-go: deepseek-v4-flash via https://opencode.ai/zen/go/v1
"""
@@ -280,7 +270,6 @@ class OpenAICompatibleLLM(LLMInterface):
"openai",
"groq",
"ollama",
"ollama-cloud",
"lmstudio",
"llamacpp",
"minimax",
@@ -289,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}")
@@ -300,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":
@@ -314,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:
@@ -334,7 +316,6 @@ class OpenAICompatibleLLM(LLMInterface):
"openrouter",
"zai",
"opencode-go",
"ollama-cloud",
)
and not self.api_key
):
@@ -370,21 +351,6 @@ class OpenAICompatibleLLM(LLMInterface):
f"base_url={self.base_url or 'default'}"
)
def _drops_tool_choice_required(self) -> bool:
"""Whether this endpoint silently ignores ``tool_choice="required"``.
True for self-hosted OpenAI-compatible servers known to return an empty
tool_calls array for "required" instead of forcing a call (#1563/#1179/
#1877). Covers LM Studio / Ollama directly, plus any server reached via
the generic "openai" provider with a custom ``base_url`` (e.g. a local
vLLM endpoint). The real OpenAI API (no base_url override) honors
"required", and cloud providers keep their own default base_urls, so both
are left untouched.
"""
if self.provider in _TOOL_CHOICE_REQUIRED_UNSUPPORTED_PROVIDERS:
return True
return self.provider == "openai" and bool(self.base_url)
async def verify_connection(self) -> None:
"""
Verify that the provider is configured correctly by making a simple test call.
@@ -485,9 +451,7 @@ class OpenAICompatibleLLM(LLMInterface):
initial_backoff: Initial backoff time in seconds.
max_backoff: Maximum backoff time in seconds.
skip_validation: Return raw JSON without Pydantic validation.
strict_schema: Use strict json_schema (grammar-enforced) response_format instead of
the soft json_object path. Supported by OpenAI and schema-capable self-hosted
backends (llama.cpp, vLLM). Server-wide via HINDSIGHT_API_LLM_STRICT_SCHEMA.
strict_schema: Use strict JSON schema enforcement (OpenAI only).
return_usage: If True, return tuple (result, TokenUsage) instead of just result.
Returns:
@@ -677,9 +641,6 @@ class OpenAICompatibleLLM(LLMInterface):
input_tokens = usage.prompt_tokens or 0 if usage else 0
output_tokens = usage.completion_tokens or 0 if usage else 0
total_tokens = usage.total_tokens or 0 if usage else 0
cached_tokens = 0
if usage and getattr(usage, "prompt_tokens_details", None):
cached_tokens = getattr(usage.prompt_tokens_details, "cached_tokens", 0) or 0
# Record LLM metrics
metrics = get_metrics_collector()
@@ -709,12 +670,14 @@ class OpenAICompatibleLLM(LLMInterface):
duration=duration,
finish_reason=finish_reason,
error=None,
cached_tokens=cached_tokens,
)
# Log slow calls
if duration > 10.0 and usage:
ratio = max(1, output_tokens) / max(1, input_tokens)
cached_tokens = 0
if hasattr(usage, "prompt_tokens_details") and usage.prompt_tokens_details:
cached_tokens = getattr(usage.prompt_tokens_details, "cached_tokens", 0) or 0
cache_info = f", cached_tokens={cached_tokens}" if cached_tokens > 0 else ""
logger.info(
f"slow llm call: scope={scope}, model={self.provider}/{self.model}, "
@@ -727,7 +690,6 @@ class OpenAICompatibleLLM(LLMInterface):
input_tokens=input_tokens,
output_tokens=output_tokens,
total_tokens=total_tokens,
cached_tokens=cached_tokens,
)
return result, token_usage
return result
@@ -896,16 +858,6 @@ class OpenAICompatibleLLM(LLMInterface):
if request_tool_choice == "auto":
request_tool_choice = None
# vLLM (--enable-auto-tool-choice), LM Studio, Ollama and similar
# self-hosted servers silently drop tool_choice="required", returning an
# empty tool_calls array instead of forcing a call (#1563/#1179/#1877).
# Downgrade to auto (None) so the model still gets to call a tool. Named
# tool_choice dicts were already normalized to "required" + a single
# filtered tool above, so the call stays practically forced even under
# auto. The real OpenAI API honors "required" and is left untouched.
if request_tool_choice == "required" and self._drops_tool_choice_required():
request_tool_choice = None
# DeepSeek tool-call replies can carry provider-specific reasoning_content.
# The normalized tool result does not retain it, but replaying assistant
# tool_calls without the field can trigger a 400. DeepSeek accepts an
@@ -1121,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()
@@ -302,24 +302,6 @@ def _is_context_overflow_error(exc: Exception) -> bool:
)
def _all_mental_models_are_usable_and_fresh(tool_output: dict[str, Any]) -> bool:
"""Return whether every retrieved mental model is explicitly fresh and has answerable content.
Used to decide without an extra LLM call whether a forced
``search_mental_models`` result is trustworthy enough to hand control back
to the agent. A model is usable only when it is explicitly ``is_stale ==
False`` (an unknown/missing staleness flag is treated as unsafe) and has
non-empty content.
"""
models = tool_output.get("mental_models") or []
for model in models:
if model.get("is_stale") is not False:
return False
if not str(model.get("content") or "").strip():
return False
return True
async def run_reflect_agent(
llm_config: "LLMProvider",
bank_id: str,
@@ -339,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.
@@ -388,40 +369,13 @@ 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},
{"role": "user", "content": query},
]
# Opt into context caching for the agentic tool loop. The system
# prompt and tool definitions are stable for the duration of this
# reflect call (and across reflects against the same bank), so
# caching them once and reusing across every iteration of the loop
# collapses the dominant input cost — the prefix repeated on every
# turn. ``get_or_create_cached_prefix`` returns None when caching is
# disabled, unsupported, or the prefix is too small; the
# ``call_with_tools`` invocation below transparently falls back to
# the uncached path in that case.
cached_prefix_name: str | None = None
provider_impl = getattr(llm_config, "_provider_impl", None)
if provider_impl is not None and provider_impl.supports_prompt_caching():
try:
cached_prefix_name = await provider_impl.get_or_create_cached_prefix(
system_instruction=system_prompt,
tools=tools,
)
except Exception:
# Caching is a soft optimisation; never let a cache-side
# error block a reflect.
cached_prefix_name = None
# Tracking
total_tools_called = 0
tool_trace: list[ToolCall] = []
@@ -482,11 +436,6 @@ async def run_reflect_agent(
)
consecutive_errors = 0
# When a forced ``search_mental_models`` returns fresh, usable models on a
# low/mid-budget call, we stop forcing the lower retrieval layers from this
# iteration onward and let the agent answer (or retrieve deeper itself)
# under ``auto`` tool choice. None means the full forced path still applies.
stop_forcing_from_iteration: int | None = None
for iteration in range(max_iterations):
is_last = iteration == max_iterations - 1
@@ -498,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",
@@ -558,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",
@@ -615,31 +558,18 @@ async def run_reflect_agent(
if include_recall:
forced_sequence.append("recall")
if stop_forcing_from_iteration is not None and iteration >= stop_forcing_from_iteration:
# A fresh mental model already short-circuited the forced path.
iter_tool_choice: str | dict = "auto"
elif iteration < len(forced_sequence):
iter_tool_choice = {"type": "function", "function": {"name": forced_sequence[iteration]}}
if iteration < len(forced_sequence):
iter_tool_choice: str | dict = {"type": "function", "function": {"name": forced_sequence[iteration]}}
else:
iter_tool_choice = "auto"
try:
ct_kwargs: dict[str, Any] = dict(
result = await llm_config.call_with_tools(
messages=messages,
tools=tools,
scope="reflect_tool_call",
tool_choice=iter_tool_choice,
)
# Gemini rejects ``cached_content`` alongside a per-request
# ``tool_config`` (forced tool choice): "CachedContent can not be used
# with GenerateContent request setting system_instruction, tools or
# tool_config." The forced-sequence iterations set tool_config, so only
# the ``auto`` iterations can reference the cache; forced iterations send
# the prefix inline. The cache (tools + system prompt) is identical
# either way, so this just limits *which* iterations are billed cached.
if cached_prefix_name is not None and iter_tool_choice == "auto":
ct_kwargs["cached_prefix"] = cached_prefix_name
result = await llm_config.call_with_tools(**ct_kwargs)
llm_duration = int((time.time() - llm_start) * 1000)
consecutive_errors = 0
total_input_tokens += result.input_tokens
@@ -677,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",
@@ -801,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",
@@ -982,25 +906,6 @@ async def run_reflect_agent(
for mm in output["mental_models"]:
if "id" in mm:
available_mental_model_ids.add(mm["id"])
# Deterministic short-circuit (no extra LLM call): on a
# low/mid-budget call, if every retrieved mental model is
# fresh and has usable content, stop forcing the lower
# retrieval layers. The next iteration runs under ``auto``
# tool choice, so the agent can answer directly when the
# mental model suffices, or — having just read it — issue a
# targeted ``search_observations``/``recall`` itself. Stale,
# empty, or missing mental models keep the full forced path.
if (
stop_forcing_from_iteration is None
and (budget or "low").lower() != "high"
and output.get("mental_models")
and _all_mental_models_are_usable_and_fresh(output)
):
stop_forcing_from_iteration = iteration + 1
logger.info(
f"[REFLECT {reflect_id}] Fresh mental models sufficient on iteration {iteration + 1}; "
"releasing forced lower-level retrieval to auto."
)
if (
normalized_tool_name == "search_observations"
@@ -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()},
}
@@ -93,7 +93,6 @@ class TokenUsage(BaseModel):
input_tokens: int = Field(default=0, description="Number of input/prompt tokens consumed")
output_tokens: int = Field(default=0, description="Number of output/completion tokens generated")
total_tokens: int = Field(default=0, description="Total tokens (input + output)")
cached_tokens: int = Field(default=0, description="Cached/cache-read prompt tokens, when reported by the provider")
def __add__(self, other: "TokenUsage") -> "TokenUsage":
"""Allow aggregating token usage from multiple calls."""
@@ -101,7 +100,6 @@ class TokenUsage(BaseModel):
input_tokens=self.input_tokens + other.input_tokens,
output_tokens=self.output_tokens + other.output_tokens,
total_tokens=self.total_tokens + other.total_tokens,
cached_tokens=self.cached_tokens + other.cached_tokens,
)
@@ -6,7 +6,6 @@ import json
import logging
import re
import uuid
from dataclasses import dataclass
from typing import TypedDict
from pydantic import BaseModel, Field
@@ -106,18 +105,6 @@ class BankProfile(TypedDict):
mission: str
@dataclass
class BankProfileResult:
"""Result of a get-or-create bank lookup.
``created`` is True when the bank row was freshly inserted on this call,
which callers use to drive the one-time HINDSIGHT_API_DEFAULT_BANK_TEMPLATE hook.
"""
profile: BankProfile
created: bool
class MissionMergeResponse(BaseModel):
"""LLM response for mission merge."""
@@ -136,8 +123,8 @@ async def get_bank_profile(pool, bank_id: str) -> BankProfile:
Returns:
BankProfile with name, typed DispositionTraits, and mission
"""
result = await get_or_create_bank_profile(pool, bank_id)
return result.profile
profile, _ = await get_or_create_bank_profile(pool, bank_id)
return profile
async def get_bank_profile_if_exists(pool, bank_id: str) -> BankProfile | None:
@@ -175,89 +162,70 @@ async def get_bank_profile_if_exists(pool, bank_id: str) -> BankProfile | None:
)
async def get_or_create_bank_profile(pool, bank_id: str) -> BankProfileResult:
async def get_or_create_bank_profile(pool, bank_id: str) -> tuple[BankProfile, bool]:
"""
Get bank profile, auto-creating with defaults if it doesn't exist.
Same as get_bank_profile, but also reports whether the bank was freshly
created on this call (``BankProfileResult.created``). Used by the memory
engine to apply the HINDSIGHT_API_DEFAULT_BANK_TEMPLATE hook on first bank
creation.
Same as get_bank_profile, but also returns a flag indicating whether the
bank was freshly created on this call. Used by the memory engine to apply
the HINDSIGHT_API_DEFAULT_BANK_TEMPLATE hook on first bank creation.
Acquires its own connection. When the caller already holds a connection and
wants the bank row to share its transaction (so the lazy bank-create commits
or rolls back atomically with the caller's write), use
``get_or_create_bank_profile_on_conn`` instead.
Returns:
Tuple of (BankProfile, created) where created is True if the bank
did not exist before this call.
"""
async with acquire_with_retry(pool) as conn:
return await get_or_create_bank_profile_on_conn(conn, bank_id, ops=pool.ops)
async def get_or_create_bank_profile_on_conn(conn, bank_id: str, *, ops) -> BankProfileResult:
"""
Connection-bound variant of ``get_or_create_bank_profile``.
Runs the SELECT, the ``INSERT ... ON CONFLICT DO NOTHING`` and the per-bank
vector index creation on the caller-supplied ``conn``. When ``conn`` is
inside an open transaction, the lazy bank-create therefore commits (or rolls
back) atomically with whatever bank-scoped write the caller performs on the
same connection closing the window where a freshly-created bank could
outlive a write that ultimately failed.
``ops`` is the backend's dialect ops object (``backend.ops``), needed for
per-bank vector index DDL.
"""
# Try to get existing bank
row = await conn.fetchrow(
f"""
SELECT name, disposition, mission
FROM {fq_table("banks")} WHERE bank_id = $1
""",
bank_id,
)
if row:
# asyncpg returns JSONB as a string, so parse it
disposition_data = row["disposition"]
if isinstance(disposition_data, str):
disposition_data = json.loads(disposition_data)
return BankProfileResult(
profile=BankProfile(
name=row["name"],
disposition=DispositionTraits(**disposition_data),
mission=row["mission"] or "",
),
created=False,
# Try to get existing bank
row = await conn.fetchrow(
f"""
SELECT name, disposition, mission
FROM {fq_table("banks")} WHERE bank_id = $1
""",
bank_id,
)
# Bank doesn't exist, create with defaults.
# Generate internal_id here so we control the value and can use it
# immediately for vector index creation without a RETURNING round-trip.
internal_id = uuid.uuid4()
inserted = await conn.fetchval(
f"""
INSERT INTO {fq_table("banks")} (bank_id, name, disposition, mission, internal_id)
VALUES ($1, $2, $3::jsonb, $4, $5)
ON CONFLICT (bank_id) DO NOTHING
RETURNING bank_id
""",
bank_id,
bank_id, # Default name is the bank_id
json.dumps(DEFAULT_DISPOSITION),
"",
internal_id,
)
if row:
# asyncpg returns JSONB as a string, so parse it
disposition_data = row["disposition"]
if isinstance(disposition_data, str):
disposition_data = json.loads(disposition_data)
created = inserted is not None
if created:
# Fresh insert — create per-bank vector indexes (instant on empty bank)
await create_bank_vector_indexes(conn, bank_id, str(internal_id), ops=ops)
return (
BankProfile(
name=row["name"],
disposition=DispositionTraits(**disposition_data),
mission=row["mission"] or "",
),
False,
)
return BankProfileResult(
profile=BankProfile(name=bank_id, disposition=DispositionTraits(**DEFAULT_DISPOSITION), mission=""),
created=created,
)
# Bank doesn't exist, create with defaults.
# Generate internal_id here so we control the value and can use it
# immediately for vector index creation without a RETURNING round-trip.
internal_id = uuid.uuid4()
inserted = await conn.fetchval(
f"""
INSERT INTO {fq_table("banks")} (bank_id, name, disposition, mission, internal_id)
VALUES ($1, $2, $3::jsonb, $4, $5)
ON CONFLICT (bank_id) DO NOTHING
RETURNING bank_id
""",
bank_id,
bank_id, # Default name is the bank_id
json.dumps(DEFAULT_DISPOSITION),
"",
internal_id,
)
created = inserted is not None
if created:
# Fresh insert — create per-bank vector indexes (instant on empty bank)
await create_bank_vector_indexes(conn, bank_id, str(internal_id), ops=pool.ops)
return (
BankProfile(name=bank_id, disposition=DispositionTraits(**DEFAULT_DISPOSITION), mission=""),
created,
)
async def update_bank_disposition(pool, bank_id: str, disposition: dict[str, int]) -> None:
@@ -4,75 +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()`."""
@property
def dimension(self) -> int: ...
def encode_query(self, texts: list[str]) -> list[list[float]]: ...
def encode_documents(self, texts: list[str]) -> list[list[float]]: ...
def _validate_embedding_vector(vector: list[float], *, index: int, expected_dimension: int) -> list[float]:
actual_dimension = len(vector)
if actual_dimension == 0:
raise RuntimeError(f"embedding {index} has dimension 0; expected {expected_dimension}")
if actual_dimension != expected_dimension:
raise RuntimeError(f"embedding {index} has dimension {actual_dimension}; expected {expected_dimension}")
return vector
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)}")
if len(embeddings) != 1:
raise RuntimeError(
f"Embeddings backend returned {len(embeddings)} vectors for 1 input text; expected exact 1:1 alignment"
)
return _validate_embedding_vector(
embeddings[0],
index=0,
expected_dimension=embeddings_backend.dimension,
)
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.
@@ -82,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)}")
@@ -102,7 +59,4 @@ async def generate_embeddings_batch(
"expected exact 1:1 alignment"
)
return [
_validate_embedding_vector(embedding, index=index, expected_dimension=embeddings_backend.dimension)
for index, embedding in enumerate(embeddings)
]
return embeddings
@@ -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)
@@ -10,13 +10,12 @@ import json
import logging
import re
from datetime import datetime, timedelta
from typing import Any, Literal, cast
from typing import Literal, cast
from pydantic import BaseModel, ConfigDict, Field, create_model, field_validator
from ...config import get_config
from ..llm_wrapper import LLMConfig, OutputTooLongError, sanitize_llm_output
from ..operation_metadata import RetainExtractionErrors
from ..response_models import TokenUsage
from .entity_labels import (
EntityLabelsConfig,
@@ -511,11 +510,11 @@ LANGUAGE: MANDATORY — Detect the language of the input text and produce ALL ou
FACT FORMAT - BE CONCISE
1. "what": Core fact - concise but complete (1-2 sentences max)
2. "when": Temporal info if mentioned. "N/A" if none. Use day name when known.
3. "where": Location if relevant. "N/A" if none.
4. "who": People involved with relationships. "N/A" if just general info.
5. "why": Context/significance ONLY if important. "N/A" if obvious.
1. **what**: Core fact - concise but complete (1-2 sentences max)
2. **when**: Temporal info if mentioned. "N/A" if none. Use day name when known.
3. **where**: Location if relevant. "N/A" if none.
4. **who**: People involved with relationships. "N/A" if just general info.
5. **why**: Context/significance ONLY if important. "N/A" if obvious.
CONCISENESS: Capture the essence, not every word. One good sentence beats three mediocre ones.
@@ -888,15 +887,17 @@ def _build_extraction_prompt_and_schema(config) -> tuple[str, type]:
extraction_mode = config.retain_extraction_mode
extract_causal_links = config.retain_extract_causal_links
# The per-bank retain mission is NOT baked into this system prompt: it would
# make the prompt bank-specific and force a separate Gemini context cache per
# mission (one per bank). Instead the prompt is bank-agnostic so a single
# CachedContent serves every bank, and the mission rides in the per-request
# user message via _retain_mission_preamble(). The {retain_mission_section}
# placeholder is kept (templates still reference it) but always empty here.
from hindsight_api.engine.prompt_utils import escape_for_prompt
retain_mission_section = ""
# Build retain_mission section if set - injected before the mode-specific guidelines
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"{retain_mission}\n\n"
)
else:
retain_mission_section = ""
# Select base prompt based on extraction mode
if extraction_mode == "custom":
@@ -909,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(
@@ -946,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:
@@ -993,26 +984,6 @@ def _build_extraction_prompt_and_schema(config) -> tuple[str, type]:
return prompt, response_schema
def _retain_mission_preamble(config) -> str:
"""The bank's retain mission, formatted for the per-request user message.
Kept OUT of the cached system prompt (which must stay bank-agnostic so one
CachedContent serves every bank otherwise each distinct mission spawns its
own cache) and prepended to the user message instead. Returns "" when unset.
No brace-escaping needed: unlike the system template, the user message is
used verbatim, not passed through str.format().
"""
retain_mission = getattr(config, "retain_mission", None)
if not retain_mission:
return ""
return (
"══════════════════════════════════════════════════════════════════════════\n"
"FOCUS — What to retain for this bank (takes priority over the general guidelines)\n"
"══════════════════════════════════════════════════════════════════════════\n\n"
f"{retain_mission}\n\n"
)
def _build_user_message(
chunk: str,
chunk_index: int,
@@ -1021,14 +992,8 @@ def _build_user_message(
context: str,
metadata: dict[str, str] | None = None,
agent_name: str | None = None,
mission_preamble: str = "",
) -> str:
"""Build user message for fact extraction.
``mission_preamble`` (the bank's retain mission, possibly empty) is prepended
so the bank-specific focus lives in the variable user turn rather than the
cached, bank-agnostic system prompt.
"""
"""Build user message for fact extraction."""
from .orchestrator import parse_datetime_flexible
sanitized_chunk = _sanitize_text(chunk)
@@ -1047,21 +1012,9 @@ def _build_user_message(
narrator_section = ""
if agent_name:
narrator_section = (
f"\nNarrator: {agent_name} (the AI agent whose memory this is). By default, "
f'first-person statements like "I did X" are {agent_name}\'s own actions → classify as '
f'"assistant".'
)
# Only defer to the Context when one was actually provided — otherwise this
# clause points at a "Context: none" line and just adds noise.
if context:
narrator_section += (
" BUT the Context above takes precedence: if it identifies a different "
"first-person speaker (e.g. a user or customer in a transcript), attribute those "
'statements to that speaker and classify them as "world", not "assistant".'
)
narrator_section = f'\nNarrator: {agent_name} (AI agent — first-person statements like "I did X" are the agent\'s own actions; classify as "assistant")'
return f"""{mission_preamble}Extract facts from the following text chunk.
return f"""Extract facts from the following text chunk.
Chunk: {chunk_index + 1}/{total_chunks}
Event Date: {event_date_str}
@@ -1087,15 +1040,12 @@ def _build_request_body(llm_config, config, prompt: str, user_message: str, resp
if llm_config.provider == "openai" and llm_config._provider_impl.openai_service_tier:
request_body["service_tier"] = llm_config._provider_impl.openai_service_tier
# Add response_format (JSON schema). The batch path builds the request body
# directly instead of going through LLMProvider.call(), so honour
# HINDSIGHT_API_LLM_STRICT_SCHEMA here too: strict=True grammar-enforces the
# output on capable backends rather than relying on the model to emit clean JSON.
# Add response_format (JSON schema)
if hasattr(response_schema, "model_json_schema"):
schema = response_schema.model_json_schema()
request_body["response_format"] = {
"type": "json_schema",
"json_schema": {"name": "facts", "schema": schema, "strict": config.llm_strict_schema},
"json_schema": {"name": "facts", "schema": schema},
}
return request_body
@@ -1131,38 +1081,8 @@ async def _extract_facts_from_chunk(
extraction_mode = config.retain_extraction_mode
extract_causal_links = config.retain_extract_causal_links
# Build user message — the bank mission rides here (not in the cached prefix).
user_message = _build_user_message(
chunk,
chunk_index,
total_chunks,
event_date,
context,
metadata,
agent_name,
mission_preamble=_retain_mission_preamble(config),
)
# Opt into context caching when the provider supports it. The prompt and
# response_schema are bank-agnostic (the mission lives in the user message),
# so one cached prefix serves every bank; reusing it across many small-payload
# retain calls dramatically lowers per-call input
# cost. ``get_or_create_cached_prefix`` returns None when caching is
# disabled, unsupported, or the prefix is too small; the LLM call
# transparently falls back to the uncached path in that case.
cached_prefix_name: str | None = None
provider_impl = getattr(llm_config, "_provider_impl", None)
if provider_impl is not None and provider_impl.supports_prompt_caching():
try:
cached_prefix_name = await provider_impl.get_or_create_cached_prefix(
system_instruction=prompt,
response_schema=response_schema,
)
except Exception:
# Caching is a soft optimisation — never let a cache-side
# error block a retain operation.
logger.exception("Cache prefix lookup failed; falling back to uncached call")
cached_prefix_name = None
# Build user message using helper function
user_message = _build_user_message(chunk, chunk_index, total_chunks, event_date, context, metadata, agent_name)
# Retry logic for JSON validation errors
# Use retain-specific overrides if set, otherwise fall back to global LLM config
@@ -1183,7 +1103,7 @@ async def _extract_facts_from_chunk(
config.retain_llm_max_backoff if config.retain_llm_max_backoff is not None else config.llm_max_backoff
)
call_kwargs: dict[str, Any] = dict(
extraction_response_json, call_usage = await llm_config.call(
messages=[{"role": "system", "content": prompt}, {"role": "user", "content": user_message}],
response_format=response_schema,
scope="retain_extract_facts",
@@ -1195,10 +1115,6 @@ async def _extract_facts_from_chunk(
skip_validation=True, # Get raw JSON, we'll validate leniently
return_usage=True,
)
if cached_prefix_name is not None:
call_kwargs["cached_prefix"] = cached_prefix_name
extraction_response_json, call_usage = await llm_config.call(**call_kwargs)
usage = usage + call_usage # Aggregate usage across retries
# Lenient parsing of facts from raw JSON
@@ -1213,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", [])
@@ -1711,39 +1624,6 @@ logger = logging.getLogger(__name__)
SECONDS_PER_FACT = 0.01
async def _write_batch_extraction_errors(
pool: Any,
operation_id: str | None,
schema: str | None,
errors: RetainExtractionErrors,
) -> None:
"""Persist non-fatal Batch API extraction errors into operation result_metadata."""
if not pool or not operation_id or errors.count == 0:
return
from ..db_utils import acquire_with_retry
from ..task_backend import fq_table
# `errors` is the complete set for this extraction run, so overwrite the
# extraction_errors_* keys rather than folding in what's already stored. On
# batch crash recovery the resumed batch reprocesses every result and
# recomputes `errors` from scratch; reading + merging the prior run's
# counters here would double-count them. The SQL `||` merge still preserves
# unrelated keys (e.g. batch_id) already on result_metadata.
table = fq_table("async_operations", schema)
async with acquire_with_retry(pool) as conn:
await conn.execute(
f"""
UPDATE {table}
SET result_metadata = COALESCE(result_metadata, '{{}}'::jsonb) || $2::jsonb,
updated_at = now()
WHERE operation_id = $1
""",
operation_id,
json.dumps(errors.to_dict()),
)
async def extract_facts_from_contents_batch_api(
contents: list[RetainContent],
llm_config,
@@ -1782,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
@@ -1835,7 +1712,6 @@ async def extract_facts_from_contents_batch_api(
item.context,
item.metadata or None,
agent_name,
mission_preamble=_retain_mission_preamble(config),
)
# Build request body using helper function
@@ -1921,7 +1797,6 @@ async def extract_facts_from_contents_batch_api(
all_facts_from_llm = []
chunks_metadata = []
total_usage = TokenUsage()
extraction_errors = RetainExtractionErrors()
for chunk_idx, (chunk_content, content_index, chunk_index_in_content, event_date, context) in enumerate(
all_chunks_info
@@ -1930,9 +1805,7 @@ async def extract_facts_from_contents_batch_api(
result = results_by_id.get(custom_id)
if not result:
message = f"{custom_id}: missing batch result"
logger.warning(message)
extraction_errors.add(message)
logger.warning(f"Missing result for {custom_id}, skipping")
chunks_metadata.append(
ChunkMetadata(
chunk_text=chunk_content, fact_count=0, content_index=content_index, chunk_index=chunk_idx
@@ -1942,9 +1815,7 @@ async def extract_facts_from_contents_batch_api(
# Check for errors
if result.get("error"):
message = f"{custom_id}: {result['error']}"
logger.error(f"Error in {message}")
extraction_errors.add(message)
logger.error(f"Error in {custom_id}: {result['error']}")
chunks_metadata.append(
ChunkMetadata(
chunk_text=chunk_content, fact_count=0, content_index=content_index, chunk_index=chunk_idx
@@ -1957,9 +1828,7 @@ async def extract_facts_from_contents_batch_api(
choices = response_body.get("choices", [])
if not choices:
message = f"{custom_id}: no choices in response"
logger.warning(message)
extraction_errors.add(message)
logger.warning(f"No choices in response for {custom_id}")
chunks_metadata.append(
ChunkMetadata(
chunk_text=chunk_content, fact_count=0, content_index=content_index, chunk_index=chunk_idx
@@ -1974,9 +1843,7 @@ async def extract_facts_from_contents_batch_api(
try:
extraction_response_json = json.loads(content_str)
except json.JSONDecodeError as e:
message = f"{custom_id}: failed to parse JSON: {e}"
logger.error(message)
extraction_errors.add(message)
logger.error(f"Failed to parse JSON for {custom_id}: {e}")
chunks_metadata.append(
ChunkMetadata(
chunk_text=chunk_content, fact_count=0, content_index=content_index, chunk_index=chunk_idx
@@ -2149,9 +2016,7 @@ async def extract_facts_from_contents_batch_api(
fact = Fact(fact=combined_text, fact_type=fact_type, **fact_data)
chunk_facts.append(fact)
except Exception as e:
message = f"{custom_id}: failed to create Fact model for fact {i}: {e}"
logger.error(message)
extraction_errors.add(message)
logger.error(f"Failed to create Fact model for fact {i}: {e}")
continue
all_facts_from_llm.extend(chunk_facts)
@@ -2216,8 +2081,6 @@ async def extract_facts_from_contents_batch_api(
# Step 8: Auto-tag facts from label groups with tag=True
_inject_label_tags(extracted_facts, config)
await _write_batch_extraction_errors(pool, operation_id, schema, extraction_errors)
logger.info(f"Batch API extracted {len(extracted_facts)} facts from {len(all_chunks_info)} chunks")
return extracted_facts, chunks_metadata, total_usage
@@ -2332,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
@@ -2345,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
@@ -147,13 +147,12 @@ async def ensure_bank_exists(conn, bank_id: str, ops=None) -> None:
internal_id = uuid.uuid4()
inserted = await conn.fetchval(
f"""
INSERT INTO {fq_table("banks")} (bank_id, name, disposition, mission, internal_id)
VALUES ($1, $2, $3::jsonb, $4, $5)
INSERT INTO {fq_table("banks")} (bank_id, disposition, mission, internal_id)
VALUES ($1, $2::jsonb, $3, $4)
ON CONFLICT (bank_id) DO NOTHING
RETURNING bank_id
""",
bank_id,
bank_id, # Default name is the bank_id (matches get_or_create_bank_profile)
json.dumps(DEFAULT_DISPOSITION),
"",
internal_id,
@@ -322,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

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