Compare commits
2
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
73bbe5f736 | ||
|
|
6bdb1a08f5 |
@@ -166,14 +166,7 @@ If any new MCP tools were added or existing tools renamed in `hindsight-api-slim
|
||||
- **`MCP_TOOL_GROUPS`** in `hindsight-control-plane/src/components/bank-config-view.tsx` — must include the new tool in the appropriate group for the UI tool selector
|
||||
- **Tool count assertions** in tests (e.g., `test_mcp_tools.py`) — must be updated to reflect the new count
|
||||
|
||||
### 11. Check backup/restore table coverage
|
||||
|
||||
If a migration adds a new PostgreSQL table (look for `CREATE TABLE` / `op.create_table` in `hindsight-api-slim/hindsight_api/alembic/versions/`):
|
||||
- **`BACKUP_TABLES`** in `hindsight-api-slim/hindsight_api/admin/cli.py` — must include the new table, placed after any table it references via foreign key (parents before children). A missing entry is silent data loss: the table is never backed up, and restore's `TRUNCATE banks CASCADE` wipes any FK-to-banks child (e.g. `mental_models`, `directives`) on restore even though it was never saved.
|
||||
- The guard test `test_backup_tables_covers_entire_schema` in `tests/test_admin_backup_restore.py` enforces this — flag it as a **must fix** if a new table is absent from `BACKUP_TABLES`.
|
||||
- Oracle-only tables (e.g. `observation_sources`) are intentionally excluded — admin backup/restore is PostgreSQL-only.
|
||||
|
||||
### 12. Review against other coding standards
|
||||
### 11. Review against other coding standards
|
||||
|
||||
Check the diff for violations of the standards listed above:
|
||||
- Python files at project root (not allowed)
|
||||
@@ -185,7 +178,7 @@ Check the diff for violations of the standards listed above:
|
||||
- Premature abstractions or speculative helpers
|
||||
- Backwards-compatibility hacks (unused vars, re-exports, "removed" comments)
|
||||
|
||||
### 13. Report findings
|
||||
### 12. Report findings
|
||||
|
||||
Present a clear summary organized by severity:
|
||||
|
||||
@@ -197,7 +190,6 @@ Present a clear summary organized by severity:
|
||||
- Multi-item tuple returns (including internal code)
|
||||
- Missing tests for new endpoints
|
||||
- New integration missing tests, CI job, or release-integration.sh entry
|
||||
- New PostgreSQL table missing from `BACKUP_TABLES` in `admin/cli.py` (silent data loss on restore)
|
||||
|
||||
**Should fix** — issues that hurt code quality:
|
||||
- Dead code / unused imports missed by linter
|
||||
|
||||
+1
-28
@@ -7,8 +7,6 @@ HINDSIGHT_API_LLM_PROVIDER=openai
|
||||
HINDSIGHT_API_LLM_API_KEY=your-api-key-here
|
||||
HINDSIGHT_API_LLM_MODEL=gpt-4o-mini
|
||||
HINDSIGHT_API_LLM_BASE_URL=https://api.openai.com/v1
|
||||
# Reasoning effort for providers/models that support it. Examples: low, medium, high, xhigh.
|
||||
# HINDSIGHT_API_LLM_REASONING_EFFORT=low
|
||||
|
||||
# Example: Anthropic Claude configuration
|
||||
# HINDSIGHT_API_LLM_PROVIDER=anthropic
|
||||
@@ -66,21 +64,8 @@ HINDSIGHT_API_LOG_LEVEL=info
|
||||
# For Azure PostgreSQL with DiskANN:
|
||||
# HINDSIGHT_API_VECTOR_EXTENSION=pgvectorscale # Auto-detects pg_diskann on Azure
|
||||
|
||||
# Text Search Extension (Optional - uses native PostgreSQL full-text search by default)
|
||||
# Backend options: "native" (default), "vchord", "pg_textsearch", "pgroonga", "pg_search"
|
||||
# HINDSIGHT_API_TEXT_SEARCH_EXTENSION=native
|
||||
# Native backend dictionary (only used by HINDSIGHT_API_TEXT_SEARCH_EXTENSION=native)
|
||||
# HINDSIGHT_API_TEXT_SEARCH_EXTENSION_NATIVE_LANGUAGE=english
|
||||
# ParadeDB pg_search tokenizer (only used when creating pg_search BM25 indexes).
|
||||
# Empty uses ParadeDB's default tokenizer: unicode_words.
|
||||
# Supported values: unicode_words, simple, whitespace, literal, literal_normalized,
|
||||
# chinese_compatible, icu, jieba, source_code,
|
||||
# chinese_lindera/lindera(chinese), japanese_lindera/lindera(japanese),
|
||||
# korean_lindera/lindera(korean), ngram(min,max), edge_ngram(min,max)
|
||||
# HINDSIGHT_API_TEXT_SEARCH_EXTENSION_PG_SEARCH_TOKENIZER=
|
||||
|
||||
# Embeddings Configuration (Optional - uses local by default)
|
||||
# Provider: "local" (default), "tei", "openai", "cohere", "google", "openrouter", "zeroentropy", "litellm", or "litellm-sdk"
|
||||
# Provider: "local" (default), "tei", "openai", "cohere", "google", "openrouter", "litellm", or "litellm-sdk"
|
||||
# HINDSIGHT_API_EMBEDDINGS_PROVIDER=local
|
||||
# For local provider:
|
||||
# HINDSIGHT_API_EMBEDDINGS_LOCAL_MODEL=BAAI/bge-small-en-v1.5
|
||||
@@ -92,13 +77,6 @@ HINDSIGHT_API_LOG_LEVEL=info
|
||||
# HINDSIGHT_API_EMBEDDINGS_OPENAI_API_KEY=sk-xxxx
|
||||
# HINDSIGHT_API_EMBEDDINGS_OPENAI_MODEL=text-embedding-3-small
|
||||
# HINDSIGHT_API_EMBEDDINGS_OPENAI_BASE_URL=https://api.openai.com/v1
|
||||
# For ZeroEntropy zembed-1:
|
||||
# HINDSIGHT_API_EMBEDDINGS_PROVIDER=zeroentropy
|
||||
# HINDSIGHT_API_EMBEDDINGS_ZEROENTROPY_API_KEY=ze-xxxx
|
||||
# HINDSIGHT_API_EMBEDDINGS_ZEROENTROPY_MODEL=zembed-1
|
||||
# HINDSIGHT_API_EMBEDDINGS_ZEROENTROPY_DIMENSIONS=1280
|
||||
# HINDSIGHT_API_EMBEDDINGS_ZEROENTROPY_ENCODING_FORMAT=float
|
||||
# HINDSIGHT_API_EMBEDDINGS_ZEROENTROPY_LATENCY=fast
|
||||
#
|
||||
# IMPORTANT: Embedding keys require provider-specific names:
|
||||
# HINDSIGHT_API_EMBEDDINGS_{PROVIDER}_{PARAMETER}
|
||||
@@ -137,11 +115,6 @@ HINDSIGHT_API_LOG_LEVEL=info
|
||||
# Dataplane API URL - where the CP proxies requests to
|
||||
# HINDSIGHT_CP_DATAPLANE_API_URL=http://localhost:8888
|
||||
|
||||
# Optional: Bearer token the CP sends as `Authorization: Bearer <key>` to the
|
||||
# dataplane API. Required when the API service is auth-protected; omit for a
|
||||
# public/unauthenticated API.
|
||||
# HINDSIGHT_CP_DATAPLANE_API_KEY=your-dataplane-bearer-token
|
||||
|
||||
# Optional: Require a shared access key to view the Control Plane UI.
|
||||
# When set, visitors see a login page and must enter the key before
|
||||
# accessing the dashboard or any /api/* routes (except /api/health).
|
||||
|
||||
+8
-325
@@ -12,7 +12,6 @@ concurrency:
|
||||
jobs:
|
||||
detect-changes:
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
permissions:
|
||||
pull-requests: read
|
||||
outputs:
|
||||
@@ -50,11 +49,7 @@ jobs:
|
||||
integrations-agentcore: ${{ steps.filter.outputs.integrations-agentcore }}
|
||||
integrations-smolagents: ${{ steps.filter.outputs.integrations-smolagents }}
|
||||
integrations-dify: ${{ steps.filter.outputs.integrations-dify }}
|
||||
integrations-gemini-spark: ${{ steps.filter.outputs.integrations-gemini-spark }}
|
||||
integrations-vapi: ${{ steps.filter.outputs.integrations-vapi }}
|
||||
integrations-flowise: ${{ steps.filter.outputs.integrations-flowise }}
|
||||
tools-agent-sdk: ${{ steps.filter.outputs.tools-agent-sdk }}
|
||||
integrations-roo-code: ${{ steps.filter.outputs.integrations-roo-code }}
|
||||
dev: ${{ steps.filter.outputs.dev }}
|
||||
ci: ${{ steps.filter.outputs.ci }}
|
||||
# Secrets are available for internal PRs and workflow_dispatch.
|
||||
@@ -148,16 +143,8 @@ jobs:
|
||||
- 'hindsight-integrations/smolagents/**'
|
||||
integrations-dify:
|
||||
- 'hindsight-integrations/dify/**'
|
||||
integrations-gemini-spark:
|
||||
- 'hindsight-integrations/gemini-spark/**'
|
||||
integrations-vapi:
|
||||
- 'hindsight-integrations/vapi/**'
|
||||
integrations-flowise:
|
||||
- 'hindsight-integrations/flowise/**'
|
||||
tools-agent-sdk:
|
||||
- 'hindsight-tools/hindsight-agent-sdk/**'
|
||||
integrations-roo-code:
|
||||
- 'hindsight-integrations/roo-code/**'
|
||||
dev:
|
||||
- 'hindsight-dev/**'
|
||||
ci:
|
||||
@@ -177,7 +164,6 @@ jobs:
|
||||
needs.detect-changes.outputs.integrations-lockfiles == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
@@ -193,7 +179,6 @@ jobs:
|
||||
needs.detect-changes.outputs.core == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
strategy:
|
||||
matrix:
|
||||
python-version: ['3.11', '3.12', '3.13', '3.14']
|
||||
@@ -224,7 +209,6 @@ jobs:
|
||||
needs.detect-changes.outputs.clients-ts == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
@@ -251,7 +235,6 @@ jobs:
|
||||
needs.detect-changes.outputs.all-npm == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
@@ -283,7 +266,6 @@ jobs:
|
||||
needs.detect-changes.outputs.all-npm == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
@@ -339,7 +321,6 @@ jobs:
|
||||
needs.detect-changes.outputs.all-npm == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
@@ -395,7 +376,6 @@ jobs:
|
||||
needs.detect-changes.outputs.integrations-claude-code == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
@@ -421,7 +401,6 @@ jobs:
|
||||
needs.detect-changes.outputs.integrations-codex == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
@@ -447,7 +426,6 @@ jobs:
|
||||
needs.detect-changes.outputs.integrations-ai-sdk == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
@@ -478,7 +456,6 @@ jobs:
|
||||
needs.detect-changes.outputs.integrations-ai-sdk == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
@@ -510,7 +487,6 @@ jobs:
|
||||
needs.detect-changes.outputs.integrations-opencode == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
@@ -541,7 +517,6 @@ jobs:
|
||||
needs.detect-changes.outputs.integrations-n8n == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
@@ -572,7 +547,6 @@ jobs:
|
||||
needs.detect-changes.outputs.tools-agent-sdk == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
@@ -605,7 +579,6 @@ jobs:
|
||||
needs.detect-changes.outputs.integrations-cloudflare-oauth-proxy == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
@@ -636,7 +609,6 @@ jobs:
|
||||
needs.detect-changes.outputs.integrations-chat == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
@@ -667,7 +639,6 @@ jobs:
|
||||
needs.detect-changes.outputs.integrations-paperclip == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
@@ -698,7 +669,6 @@ jobs:
|
||||
needs.detect-changes.outputs.integrations-pipecat == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
@@ -728,65 +698,6 @@ jobs:
|
||||
working-directory: ./hindsight-integrations/pipecat
|
||||
run: uv run pytest tests -v
|
||||
|
||||
test-gemini-spark-integration:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
(github.event_name == 'workflow_dispatch' ||
|
||||
needs.detect-changes.outputs.integrations-gemini-spark == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ github.event.pull_request.head.sha || '' }}
|
||||
|
||||
- name: Install uv
|
||||
uses: astral-sh/setup-uv@v7
|
||||
with:
|
||||
enable-cache: true
|
||||
prune-cache: false
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v6
|
||||
with:
|
||||
python-version-file: ".python-version"
|
||||
|
||||
- name: Install dependencies
|
||||
working-directory: ./hindsight-integrations/gemini-spark
|
||||
run: uv sync --frozen
|
||||
|
||||
- name: Run tests
|
||||
working-directory: ./hindsight-integrations/gemini-spark
|
||||
run: uv run pytest tests -v
|
||||
|
||||
test-roo-code-integration:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
github.event_name != 'pull_request_review' &&
|
||||
(github.event_name == 'workflow_dispatch' ||
|
||||
needs.detect-changes.outputs.integrations-roo-code == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ github.event.pull_request.head.sha || '' }}
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v6
|
||||
with:
|
||||
python-version: '3.11'
|
||||
|
||||
- name: Install pytest
|
||||
run: pip install pytest
|
||||
|
||||
- name: Run tests
|
||||
working-directory: ./hindsight-integrations/roo-code
|
||||
run: python -m pytest tests/ -v
|
||||
|
||||
build-control-plane:
|
||||
needs: [detect-changes]
|
||||
@@ -796,7 +707,6 @@ jobs:
|
||||
needs.detect-changes.outputs.clients-ts == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
@@ -824,12 +734,6 @@ jobs:
|
||||
rm -rf node_modules/lightningcss node_modules/@tailwindcss
|
||||
npm install lightningcss @tailwindcss/postcss @tailwindcss/node
|
||||
|
||||
- name: Test Control Plane
|
||||
run: npm test --workspace=hindsight-control-plane
|
||||
|
||||
- name: Check i18n locale parity and hardcoded strings
|
||||
run: npm run i18n:check --workspace=hindsight-control-plane
|
||||
|
||||
- name: Build Control Plane
|
||||
run: npm run build --workspace=hindsight-control-plane
|
||||
|
||||
@@ -862,7 +766,6 @@ jobs:
|
||||
needs.detect-changes.outputs.docs == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
@@ -891,7 +794,6 @@ jobs:
|
||||
needs.detect-changes.outputs.cli == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
env:
|
||||
HINDSIGHT_API_LLM_PROVIDER: vertexai
|
||||
HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY: /tmp/gcp-credentials.json
|
||||
@@ -1019,7 +921,6 @@ jobs:
|
||||
needs.detect-changes.outputs.helm == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
@@ -1034,23 +935,6 @@ jobs:
|
||||
- name: Lint Helm chart
|
||||
run: helm lint helm/hindsight
|
||||
|
||||
test-standalone-start-script:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
(github.event_name == 'workflow_dispatch' ||
|
||||
needs.detect-changes.outputs.docker == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ github.event.pull_request.head.sha || '' }}
|
||||
|
||||
- name: Run standalone start script tests
|
||||
run: bash docker/standalone/test-start-all.sh
|
||||
|
||||
build-docker-images:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
@@ -1062,7 +946,6 @@ jobs:
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
name: Build Docker (${{ matrix.name }})
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
strategy:
|
||||
matrix:
|
||||
include:
|
||||
@@ -1153,16 +1036,6 @@ jobs:
|
||||
needs.detect-changes.outputs.core == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
# Sharded with pytest-split. Splits the ~19-min pytest run across 3
|
||||
# parallel jobs, cutting critical-path wall time to ~7-8 min/shard. The
|
||||
# .venv cache lets shards 2+ skip the ~3-min `uv sync` once shard 1
|
||||
# populates the key — same key on re-runs hits cache on all three.
|
||||
strategy:
|
||||
fail-fast: false
|
||||
matrix:
|
||||
shard: [1, 2, 3]
|
||||
name: test-api (${{ matrix.shard }}/3)
|
||||
env:
|
||||
HINDSIGHT_API_LLM_PROVIDER: vertexai
|
||||
HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY: /tmp/gcp-credentials.json
|
||||
@@ -1203,18 +1076,6 @@ jobs:
|
||||
working-directory: ./hindsight-api-slim
|
||||
run: uv build
|
||||
|
||||
- name: Cache hindsight-api-slim/.venv
|
||||
# Keyed on the workspace lockfile, the API package's pyproject, and the
|
||||
# pinned Python version — the only inputs that change the resolved env.
|
||||
# `uv sync --frozen` still runs after restore but is a near-instant link
|
||||
# check when the venv already matches the lock.
|
||||
uses: actions/cache@v5
|
||||
with:
|
||||
path: hindsight-api-slim/.venv
|
||||
key: ${{ runner.os }}-venv-test-api-${{ hashFiles('uv.lock', 'hindsight-api-slim/pyproject.toml', '.python-version') }}
|
||||
restore-keys: |
|
||||
${{ runner.os }}-venv-test-api-
|
||||
|
||||
- name: Install dependencies
|
||||
working-directory: ./hindsight-api-slim
|
||||
run: uv sync --frozen --all-extras --index-strategy unsafe-best-match
|
||||
@@ -1239,76 +1100,9 @@ jobs:
|
||||
print('Models downloaded successfully')
|
||||
"
|
||||
|
||||
- name: Run tests (shard ${{ matrix.shard }}/3)
|
||||
- name: Run tests
|
||||
working-directory: ./hindsight-api-slim
|
||||
# `--with pytest-split` adds the plugin ad-hoc — no uv.lock churn.
|
||||
# pytest-split filters at collection (before xdist takes over), so it
|
||||
# composes cleanly with the `-n 8 --dist loadgroup` baked into addopts.
|
||||
run: uv run --with pytest-split pytest tests -v -m "not hs_llm_mat and not hs_llm_core" --splits 3 --group ${{ matrix.shard }}
|
||||
|
||||
test-api-llm-core:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
needs.detect-changes.outputs.has_secrets == 'true' &&
|
||||
(github.event_name == 'workflow_dispatch' ||
|
||||
needs.detect-changes.outputs.core == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
name: Core LLM tests
|
||||
env:
|
||||
HINDSIGHT_API_LLM_PROVIDER: vertexai
|
||||
HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY: /tmp/gcp-credentials.json
|
||||
HINDSIGHT_API_LLM_MODEL: google/gemini-2.5-flash-lite
|
||||
GEMINI_API_KEY: ${{ secrets.GEMINI_API_KEY }}
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ github.event.pull_request.head.sha || '' }}
|
||||
|
||||
- name: Setup GCP credentials
|
||||
run: |
|
||||
printf '%s' '${{ secrets.GCP_VERTEXAI_CREDENTIALS }}' > /tmp/gcp-credentials.json
|
||||
PROJECT_ID=$(jq -r '.project_id' /tmp/gcp-credentials.json)
|
||||
echo "HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID=$PROJECT_ID" >> $GITHUB_ENV
|
||||
|
||||
- name: Install uv
|
||||
uses: astral-sh/setup-uv@v7
|
||||
with:
|
||||
enable-cache: true
|
||||
prune-cache: false
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v6
|
||||
with:
|
||||
python-version-file: ".python-version"
|
||||
|
||||
- name: Install dependencies
|
||||
working-directory: ./hindsight-api-slim
|
||||
run: uv sync --frozen --all-extras --index-strategy unsafe-best-match
|
||||
|
||||
- name: Cache HuggingFace models
|
||||
uses: actions/cache@v5
|
||||
with:
|
||||
path: ~/.cache/huggingface
|
||||
key: ${{ runner.os }}-huggingface-${{ hashFiles('hindsight-api-slim/pyproject.toml') }}
|
||||
restore-keys: |
|
||||
${{ runner.os }}-huggingface-
|
||||
|
||||
- name: Pre-download models
|
||||
working-directory: ./hindsight-api-slim
|
||||
run: |
|
||||
uv run python -c "
|
||||
from sentence_transformers import SentenceTransformer, CrossEncoder
|
||||
SentenceTransformer('BAAI/bge-small-en-v1.5')
|
||||
CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2')
|
||||
print('Models downloaded successfully')
|
||||
"
|
||||
|
||||
- name: Run core LLM tests
|
||||
working-directory: ./hindsight-api-slim
|
||||
run: uv run pytest tests -v -m "hs_llm_core" --timeout 600
|
||||
run: uv run pytest tests -v -m "not hs_llm_mat"
|
||||
|
||||
test-api-llm-acceptance:
|
||||
needs: [detect-changes]
|
||||
@@ -1319,7 +1113,6 @@ jobs:
|
||||
needs.detect-changes.outputs.core == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
strategy:
|
||||
fail-fast: false
|
||||
matrix:
|
||||
@@ -1426,7 +1219,6 @@ jobs:
|
||||
needs.detect-changes.outputs.core == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
env:
|
||||
HINDSIGHT_API_LLM_PROVIDER: vertexai
|
||||
HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY: /tmp/gcp-credentials.json
|
||||
@@ -1468,8 +1260,8 @@ jobs:
|
||||
conn = oracledb.connect(user='system', password='oracle', dsn='localhost:1521/FREEPDB1')
|
||||
cursor = conn.cursor()
|
||||
cursor.execute(\"\"\"
|
||||
CREATE BIGFILE TABLESPACE hindsight_ts
|
||||
DATAFILE 'hindsight_ts.dbf' SIZE 2G AUTOEXTEND ON NEXT 500M MAXSIZE UNLIMITED
|
||||
CREATE TABLESPACE hindsight_ts
|
||||
DATAFILE 'hindsight_ts.dbf' SIZE 200M AUTOEXTEND ON NEXT 50M
|
||||
EXTENT MANAGEMENT LOCAL
|
||||
SEGMENT SPACE MANAGEMENT AUTO
|
||||
\"\"\")
|
||||
@@ -1546,7 +1338,6 @@ jobs:
|
||||
needs.detect-changes.outputs.clients-python == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
env:
|
||||
HINDSIGHT_API_LLM_PROVIDER: vertexai
|
||||
HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY: /tmp/gcp-credentials.json
|
||||
@@ -1658,7 +1449,6 @@ jobs:
|
||||
needs.detect-changes.outputs.clients-ts == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
env:
|
||||
HINDSIGHT_API_LLM_PROVIDER: vertexai
|
||||
HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY: /tmp/gcp-credentials.json
|
||||
@@ -1788,7 +1578,6 @@ jobs:
|
||||
needs.detect-changes.outputs.clients-python == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
env:
|
||||
HINDSIGHT_API_LLM_PROVIDER: vertexai
|
||||
HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY: /tmp/gcp-credentials.json
|
||||
@@ -1825,8 +1614,8 @@ jobs:
|
||||
conn = oracledb.connect(user='system', password='oracle', dsn='localhost:1521/FREEPDB1')
|
||||
cursor = conn.cursor()
|
||||
cursor.execute(\"\"\"
|
||||
CREATE BIGFILE TABLESPACE hindsight_ts
|
||||
DATAFILE 'hindsight_ts.dbf' SIZE 2G AUTOEXTEND ON NEXT 500M MAXSIZE UNLIMITED
|
||||
CREATE TABLESPACE hindsight_ts
|
||||
DATAFILE 'hindsight_ts.dbf' SIZE 200M AUTOEXTEND ON NEXT 50M
|
||||
EXTENT MANAGEMENT LOCAL
|
||||
SEGMENT SPACE MANAGEMENT AUTO
|
||||
\"\"\")
|
||||
@@ -1948,7 +1737,6 @@ jobs:
|
||||
needs.detect-changes.outputs.clients-ts == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
env:
|
||||
HINDSIGHT_API_LLM_PROVIDER: vertexai
|
||||
HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY: /tmp/gcp-credentials.json
|
||||
@@ -1985,8 +1773,8 @@ jobs:
|
||||
conn = oracledb.connect(user='system', password='oracle', dsn='localhost:1521/FREEPDB1')
|
||||
cursor = conn.cursor()
|
||||
cursor.execute(\"\"\"
|
||||
CREATE BIGFILE TABLESPACE hindsight_ts
|
||||
DATAFILE 'hindsight_ts.dbf' SIZE 2G AUTOEXTEND ON NEXT 500M MAXSIZE UNLIMITED
|
||||
CREATE TABLESPACE hindsight_ts
|
||||
DATAFILE 'hindsight_ts.dbf' SIZE 200M AUTOEXTEND ON NEXT 50M
|
||||
EXTENT MANAGEMENT LOCAL
|
||||
SEGMENT SPACE MANAGEMENT AUTO
|
||||
\"\"\")
|
||||
@@ -2108,7 +1896,6 @@ jobs:
|
||||
needs.detect-changes.outputs.clients-ts == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
env:
|
||||
HINDSIGHT_API_LLM_PROVIDER: vertexai
|
||||
HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY: /tmp/gcp-credentials.json
|
||||
@@ -2229,7 +2016,6 @@ jobs:
|
||||
needs.detect-changes.outputs.cli == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-24.04-arm
|
||||
timeout-minutes: 30
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
@@ -2263,7 +2049,6 @@ jobs:
|
||||
needs.detect-changes.outputs.clients-rust == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
env:
|
||||
HINDSIGHT_API_LLM_PROVIDER: vertexai
|
||||
HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY: /tmp/gcp-credentials.json
|
||||
@@ -2379,7 +2164,6 @@ jobs:
|
||||
needs.detect-changes.outputs.clients-go == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
env:
|
||||
HINDSIGHT_API_LLM_PROVIDER: vertexai
|
||||
HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY: /tmp/gcp-credentials.json
|
||||
@@ -2494,7 +2278,6 @@ jobs:
|
||||
needs.detect-changes.outputs.embed == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
env:
|
||||
HINDSIGHT_API_LLM_PROVIDER: vertexai
|
||||
HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY: /tmp/gcp-credentials.json
|
||||
@@ -2626,7 +2409,6 @@ jobs:
|
||||
needs.detect-changes.outputs.integration-tests == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
env:
|
||||
HINDSIGHT_API_LLM_PROVIDER: vertexai
|
||||
HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY: /tmp/gcp-credentials.json
|
||||
@@ -2731,7 +2513,6 @@ jobs:
|
||||
needs.detect-changes.outputs.integrations-ag2 == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
@@ -2768,7 +2549,6 @@ jobs:
|
||||
needs.detect-changes.outputs.integrations-smolagents == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
@@ -2806,7 +2586,6 @@ jobs:
|
||||
needs.detect-changes.outputs.integrations-dify == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
@@ -2833,37 +2612,6 @@ jobs:
|
||||
working-directory: ./hindsight-integrations/dify
|
||||
run: pytest tests -v
|
||||
|
||||
test-flowise-integration:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
github.event_name != 'pull_request_review' &&
|
||||
(github.event_name == 'workflow_dispatch' ||
|
||||
needs.detect-changes.outputs.integrations-flowise == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ github.event.pull_request.head.sha || '' }}
|
||||
|
||||
- name: Set up Node.js
|
||||
uses: actions/setup-node@v6
|
||||
with:
|
||||
node-version: '22'
|
||||
|
||||
- name: Install dependencies
|
||||
working-directory: ./hindsight-integrations/flowise
|
||||
run: npm install --no-audit --no-fund
|
||||
|
||||
- name: Type check
|
||||
working-directory: ./hindsight-integrations/flowise
|
||||
run: npx tsc --noEmit
|
||||
|
||||
- name: Run tests
|
||||
working-directory: ./hindsight-integrations/flowise
|
||||
run: npm test
|
||||
|
||||
test-crewai-integration:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
@@ -2871,7 +2619,6 @@ jobs:
|
||||
needs.detect-changes.outputs.integrations-crewai == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
@@ -2901,44 +2648,6 @@ jobs:
|
||||
working-directory: ./hindsight-integrations/crewai
|
||||
run: uv run pytest tests -v
|
||||
|
||||
test-vapi-integration:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
github.event_name != 'pull_request_review' &&
|
||||
(github.event_name == 'workflow_dispatch' ||
|
||||
needs.detect-changes.outputs.integrations-vapi == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ github.event.pull_request.head.sha || '' }}
|
||||
|
||||
- name: Install uv
|
||||
uses: astral-sh/setup-uv@v7
|
||||
with:
|
||||
enable-cache: true
|
||||
prune-cache: false
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v6
|
||||
with:
|
||||
python-version-file: ".python-version"
|
||||
|
||||
- name: Build vapi integration
|
||||
working-directory: ./hindsight-integrations/vapi
|
||||
run: uv build
|
||||
|
||||
- name: Install dependencies
|
||||
working-directory: ./hindsight-integrations/vapi
|
||||
run: uv sync --frozen
|
||||
|
||||
- name: Run tests
|
||||
working-directory: ./hindsight-integrations/vapi
|
||||
run: uv run pytest tests -v
|
||||
|
||||
test-litellm-integration:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
@@ -2946,7 +2655,6 @@ jobs:
|
||||
needs.detect-changes.outputs.integrations-litellm == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
@@ -2983,7 +2691,6 @@ jobs:
|
||||
needs.detect-changes.outputs.integrations-pydantic-ai == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
@@ -3020,7 +2727,6 @@ jobs:
|
||||
needs.detect-changes.outputs.integrations-llamaindex == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
@@ -3057,7 +2763,6 @@ jobs:
|
||||
needs.detect-changes.outputs.integrations-openai-agents == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
@@ -3094,7 +2799,6 @@ jobs:
|
||||
needs.detect-changes.outputs.integrations-agentcore == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
@@ -3132,7 +2836,6 @@ jobs:
|
||||
needs.detect-changes.outputs.core == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
env:
|
||||
HINDSIGHT_API_LLM_PROVIDER: vertexai
|
||||
HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY: /tmp/gcp-credentials.json
|
||||
@@ -3202,7 +2905,6 @@ jobs:
|
||||
needs.detect-changes.outputs.embed == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
env:
|
||||
HINDSIGHT_API_LLM_PROVIDER: vertexai
|
||||
HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY: /tmp/gcp-credentials.json
|
||||
@@ -3270,7 +2972,6 @@ jobs:
|
||||
needs.detect-changes.outputs.embed == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: windows-latest
|
||||
timeout-minutes: 30
|
||||
env:
|
||||
HINDSIGHT_API_LLM_PROVIDER: vertexai
|
||||
HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY: ${{ github.workspace }}/gcp-credentials.json
|
||||
@@ -3448,7 +3149,6 @@ jobs:
|
||||
needs.detect-changes.outputs.hindsight-all == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
env:
|
||||
HINDSIGHT_API_LLM_PROVIDER: vertexai
|
||||
HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY: /tmp/gcp-credentials.json
|
||||
@@ -3500,12 +3200,6 @@ jobs:
|
||||
rm -rf node_modules/lightningcss node_modules/@tailwindcss
|
||||
npm install lightningcss @tailwindcss/postcss @tailwindcss/node
|
||||
|
||||
- name: Test Control Plane
|
||||
run: npm test --workspace=hindsight-control-plane
|
||||
|
||||
- name: Check i18n locale parity and hardcoded strings
|
||||
run: npm run i18n:check --workspace=hindsight-control-plane
|
||||
|
||||
- name: Build Control Plane
|
||||
run: npm run build --workspace=hindsight-control-plane
|
||||
|
||||
@@ -3543,7 +3237,6 @@ jobs:
|
||||
needs.detect-changes.outputs.docs == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
strategy:
|
||||
fail-fast: false
|
||||
matrix:
|
||||
@@ -3692,7 +3385,6 @@ jobs:
|
||||
needs.detect-changes.outputs.dev == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
env:
|
||||
HINDSIGHT_API_LLM_PROVIDER: vertexai
|
||||
HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY: /tmp/gcp-credentials.json
|
||||
@@ -3771,7 +3463,6 @@ jobs:
|
||||
|
||||
verify-generated-files:
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
env:
|
||||
UV_FROZEN: "1"
|
||||
steps:
|
||||
@@ -3858,7 +3549,6 @@ jobs:
|
||||
needs.detect-changes.outputs.core == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
@@ -3912,7 +3602,6 @@ jobs:
|
||||
needs.detect-changes.outputs.dev == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
@@ -3958,14 +3647,10 @@ jobs:
|
||||
- build-chat-integration
|
||||
- test-paperclip-integration
|
||||
- test-pipecat-integration
|
||||
- test-gemini-spark-integration
|
||||
- test-vapi-integration
|
||||
- test-roo-code-integration
|
||||
- build-control-plane
|
||||
- build-docs
|
||||
- test-rust-cli
|
||||
- lint-helm-chart
|
||||
- test-standalone-start-script
|
||||
- build-docker-images
|
||||
- test-api
|
||||
- test-api-oracle
|
||||
@@ -3982,7 +3667,6 @@ jobs:
|
||||
- test-ag2-integration
|
||||
- test-smolagents-integration
|
||||
- test-dify-integration
|
||||
- test-flowise-integration
|
||||
- test-crewai-integration
|
||||
- test-litellm-integration
|
||||
- test-pydantic-ai-integration
|
||||
@@ -3999,7 +3683,6 @@ jobs:
|
||||
- check-openapi-compatibility
|
||||
- check-cli-coverage
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
permissions:
|
||||
statuses: write
|
||||
pull-requests: write
|
||||
|
||||
@@ -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"
|
||||
@@ -54,8 +54,6 @@ hindsight-clients/rust/target
|
||||
!.claude/skills/
|
||||
whats-next.md
|
||||
TASK.md
|
||||
# Parked / draft integrations that aren't ready to ship
|
||||
hindsight-integrations/_drafts/
|
||||
# Changelog is now tracked in hindsight-docs/src/pages/changelog.md
|
||||
# CHANGELOG.md
|
||||
|
||||
|
||||
+2
-25
@@ -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/
|
||||
|
||||
@@ -1,113 +0,0 @@
|
||||
# Hindsight with Claude Code (Claude Pro/Max subscription)
|
||||
|
||||
Run Hindsight inside Docker using the `claude-code` LLM provider, backed by
|
||||
your host machine's Claude Pro or Max subscription credentials.
|
||||
|
||||
The standalone Hindsight Docker image ships `claude-agent-sdk` but does **not**
|
||||
bundle the host `claude` CLI binary or any Claude credentials. This Compose
|
||||
file bind-mounts the host's CLI install and credentials into the container so
|
||||
the `claude-code` provider works without an API key.
|
||||
|
||||
## When to use this
|
||||
|
||||
- You have an active Claude Pro or Max subscription and want to use it for
|
||||
Hindsight without paying separate Anthropic API costs.
|
||||
- You want a one-command `docker compose up` instead of a long `docker run`
|
||||
invocation with many flags.
|
||||
- You are running on **Linux/amd64** — macOS Docker Desktop and Windows host
|
||||
paths differ and are not yet covered (please open an issue if you'd like to
|
||||
contribute a verified recipe for either).
|
||||
|
||||
> **Personal-use only.** Anthropic's
|
||||
> [Agent SDK documentation](https://docs.claude.com/en/api/agent-sdk/overview)
|
||||
> states that third-party developers should not offer claude.ai login or rate
|
||||
> limits for their products. Hindsight does **not** perform any login on your
|
||||
> behalf — it uses credentials you've already authenticated via
|
||||
> `claude auth login`. In January 2026, Anthropic
|
||||
> [enforced restrictions](https://paddo.dev/blog/anthropic-walled-garden-crackdown/)
|
||||
> against tools that spoofed the Claude Code client identity; Hindsight uses
|
||||
> the official Claude Agent SDK instead.
|
||||
>
|
||||
> Do not deploy this configuration to shared environments or production. For
|
||||
> that, use the `anthropic` provider with an API key from the
|
||||
> [Anthropic Console](https://console.anthropic.com/). Usage counts against
|
||||
> your Claude Pro/Max subscription limits.
|
||||
|
||||
## Prerequisites
|
||||
|
||||
- Host has `claude` CLI installed (e.g., `npm install -g @anthropics/claude-code`)
|
||||
and `claude auth login` has been run successfully.
|
||||
- `~/.claude.json` and `~/.claude/.credentials.json` exist on the host.
|
||||
- Host `claude` CLI version is **2.1.128 or newer** — the version bundled with
|
||||
`claude-agent-sdk` 0.5.x has a protocol incompatibility in containers, so
|
||||
the recipe overrides it with the host binary.
|
||||
|
||||
## Quick start
|
||||
|
||||
```bash
|
||||
# Set your host UID/GID (defaults to 1000:1000 if unset)
|
||||
export HOST_UID=$(id -u)
|
||||
export HOST_GID=$(id -g)
|
||||
|
||||
docker compose -f docker/docker-compose/claude-code/docker-compose.yaml up -d
|
||||
```
|
||||
|
||||
- API: http://localhost:8888
|
||||
- Control Plane: http://localhost:9999
|
||||
|
||||
## Post-setup (one-time)
|
||||
|
||||
After the container starts for the first time, run these commands to fix
|
||||
permissions and symlink the host `claude` binary into `$PATH`:
|
||||
|
||||
```bash
|
||||
# Make ~/.claude writable by your UID (the CLI writes session/project state)
|
||||
docker exec --user 0:0 hindsight-claude-code chown $(id -u):$(id -g) /home/hindsight/.claude
|
||||
docker exec --user 0:0 hindsight-claude-code chmod 755 /home/hindsight/.claude
|
||||
|
||||
# Symlink the host claude binary into PATH
|
||||
docker exec --user 0:0 hindsight-claude-code \
|
||||
ln -sf /home/hindsight/.local/share/claude/versions/2.1.128 /usr/local/bin/claude
|
||||
```
|
||||
|
||||
If you set `CLAUDE_CLI_VERSION` to a version other than `2.1.128`, update the
|
||||
symlink path accordingly.
|
||||
|
||||
## Notes on the bind-mount surface (every flag is load-bearing)
|
||||
|
||||
- **Host `claude` binary required** — the image ships only `claude-agent-sdk`,
|
||||
not the CLI itself.
|
||||
- **SDK bundled-binary override** — the override of
|
||||
`claude_agent_sdk/_bundled/claude` works around a protocol issue in the
|
||||
bundled v2.1.121 binary inside containers. Once `claude-agent-sdk` ships
|
||||
with v2.1.128+ this override can be dropped. Set `CLAUDE_CLI_VERSION` to
|
||||
match your installed version.
|
||||
- **Single-file credential mounts** — credentials are mounted as individual
|
||||
`:ro` files rather than a whole-directory `:ro` mount of `~/.claude`,
|
||||
because the CLI writes session/project state at runtime and a read-only
|
||||
directory mount silently breaks it.
|
||||
- **`--user` / `user:`** — the `user: ${HOST_UID}:${HOST_GID}` pattern
|
||||
requires `chmod 755 /home/hindsight`, which is built into the image since
|
||||
v0.6.0 (see [#1481](https://github.com/vectorize-io/hindsight/issues/1481)).
|
||||
- **`~/.hindsight-docker` data directory** — the pg0 data bind mount must be
|
||||
writable by your host UID (see
|
||||
[#1483](https://github.com/vectorize-io/hindsight/issues/1483)).
|
||||
- **Verified** on `linux/amd64` against `ghcr.io/vectorize-io/hindsight:latest`
|
||||
v0.5.6+.
|
||||
|
||||
## Using a different Claude CLI version
|
||||
|
||||
If your host has a `claude` version other than 2.1.128, set
|
||||
`CLAUDE_CLI_VERSION` before starting:
|
||||
|
||||
```bash
|
||||
export CLAUDE_CLI_VERSION=2.2.0
|
||||
docker compose -f docker/docker-compose/claude-code/docker-compose.yaml up -d
|
||||
```
|
||||
|
||||
Then update the post-setup symlink to match:
|
||||
|
||||
```bash
|
||||
docker exec --user 0:0 hindsight-claude-code \
|
||||
ln -sf /home/hindsight/.local/share/claude/versions/2.2.0 /usr/local/bin/claude
|
||||
```
|
||||
@@ -1,44 +0,0 @@
|
||||
name: hindsight-claude-code
|
||||
# Run Hindsight with the claude-code LLM provider, using your host machine's
|
||||
# Claude Pro/Max subscription credentials. Linux/amd64 only for now.
|
||||
#
|
||||
# Quick start:
|
||||
# docker compose -f docker/docker-compose/claude-code/docker-compose.yaml up -d
|
||||
#
|
||||
# See README.md for prerequisites, post-setup steps, and important caveats.
|
||||
|
||||
services:
|
||||
hindsight:
|
||||
image: ghcr.io/vectorize-io/hindsight:latest
|
||||
container_name: hindsight-claude-code
|
||||
user: "${HOST_UID:-1000}:${HOST_GID:-1000}"
|
||||
ports:
|
||||
- "127.0.0.1:8888:8888"
|
||||
- "127.0.0.1:9999:9999"
|
||||
environment:
|
||||
HOME: /home/hindsight
|
||||
USER: hindsight
|
||||
LOGNAME: hindsight
|
||||
PATH: /usr/local/bin:/usr/bin:/bin:/app/api/.venv/bin
|
||||
HINDSIGHT_API_LLM_PROVIDER: claude-code
|
||||
volumes:
|
||||
# ── Persistent data ────────────────────────────────────────────
|
||||
# Writable pg0 data directory. Must be writable by HOST_UID.
|
||||
- ${HOME:-.}/.hindsight-docker:/home/hindsight/.pg0
|
||||
|
||||
# ── Claude credentials (read-only, single-file mounts) ────────
|
||||
# A whole-directory :ro mount of ~/.claude silently breaks the
|
||||
# CLI, which writes session/project state at runtime — so we
|
||||
# mount only the two credential files.
|
||||
- ${HOME}/.claude/.credentials.json:/home/hindsight/.claude/.credentials.json:ro
|
||||
- ${HOME}/.claude.json:/home/hindsight/.claude.json:ro
|
||||
|
||||
# ── Claude CLI install (read-only) ─────────────────────────────
|
||||
- ${HOME}/.local/share/claude:/home/hindsight/.local/share/claude:ro
|
||||
|
||||
# ── SDK bundled-binary override ────────────────────────────────
|
||||
# The claude-agent-sdk 0.5.x image bundles v2.1.121 which has a
|
||||
# protocol incompatibility in containers. Override it with the
|
||||
# host's v2.1.128+ binary. Drop this mount once claude-agent-sdk
|
||||
# ships with v2.1.128+.
|
||||
- ${HOME}/.local/share/claude/versions/${CLAUDE_CLI_VERSION:-2.1.128}:/app/api/.venv/lib/python3.11/site-packages/claude_agent_sdk/_bundled/claude:ro
|
||||
@@ -1,103 +0,0 @@
|
||||
# Hindsight with a local llama.cpp server sidecar
|
||||
|
||||
Example Docker Compose setup that runs Hindsight against a **local
|
||||
llama.cpp server**, fully offline, with no external API key required.
|
||||
|
||||
## Architecture
|
||||
|
||||
```
|
||||
┌────────────┐ HTTP /v1/chat/completions ┌──────────────────────────────┐
|
||||
│ hindsight │ ──────────────────────────▶ │ llama.cpp server (sidecar) │
|
||||
│ (API + CP) │ │ ghcr.io/ggml-org/llama.cpp │
|
||||
└────────────┘ └──────────────────────────────┘
|
||||
```
|
||||
|
||||
`llama.cpp` runs as its own container and exposes an OpenAI-compatible
|
||||
HTTP API. Hindsight talks to it via the standard `openai` LLM provider
|
||||
with `HINDSIGHT_API_LLM_BASE_URL` pointed at the sidecar.
|
||||
|
||||
This pattern follows
|
||||
[*Hosting llama-server with Docker* (ServiceStack)](https://servicestack.net/posts/hosting-llama-server).
|
||||
|
||||
### Why a sidecar and not the in-process `llamacpp` provider?
|
||||
|
||||
Hindsight does ship an in-process `llamacpp` provider that spawns
|
||||
`llama-cpp-python`, but the **published `ghcr.io/vectorize-io/hindsight`
|
||||
image deliberately omits `llama-cpp-python`** to keep the image small and
|
||||
avoid bundling native inference libraries that most users don't need.
|
||||
Trying to set `HINDSIGHT_API_LLM_PROVIDER=llamacpp` against the published
|
||||
image fails with `ModuleNotFoundError: No module named 'llama_cpp'`.
|
||||
|
||||
The sidecar approach side-steps that entirely: the official llama.cpp
|
||||
image is used as-is for inference, Hindsight is used as-is for memory.
|
||||
Clean separation, no derived images.
|
||||
|
||||
## Quick start
|
||||
|
||||
```bash
|
||||
docker compose -f docker/docker-compose/local-llm/docker-compose.yaml up
|
||||
```
|
||||
|
||||
- API: http://localhost:8888
|
||||
- Control Plane: http://localhost:9999
|
||||
|
||||
**First boot downloads ~3.5 GB** (Gemma 4 E2B Q4_K_M GGUF) into the
|
||||
`llama_models` named volume. Subsequent boots reuse it.
|
||||
|
||||
Hindsight only starts after llama.cpp's `/health` endpoint reports
|
||||
healthy, so the API will appear "stuck" for a few minutes on the first
|
||||
run while the model downloads.
|
||||
|
||||
## Using a different model
|
||||
|
||||
Override the HuggingFace repo / file in `docker-compose.yaml`:
|
||||
|
||||
```yaml
|
||||
environment:
|
||||
LLAMA_ARG_HF_REPO: bartowski/Qwen2.5-7B-Instruct-GGUF
|
||||
LLAMA_ARG_HF_FILE: Qwen2.5-7B-Instruct-Q4_K_M.gguf
|
||||
```
|
||||
|
||||
Also update `HINDSIGHT_API_LLM_MODEL` on the `hindsight` service to a
|
||||
matching alias (the value is sent to llama-server as the OpenAI `model`
|
||||
field — llama-server is lenient about this but it shows up in logs).
|
||||
|
||||
## GPU acceleration
|
||||
|
||||
The default compose file targets CPU because not everyone has a GPU. On
|
||||
CPU, Gemma 4 E2B runs at ~2-3 tokens/sec — fine for a smoke test, but the
|
||||
retain pipeline (which makes several multi-hundred-token LLM calls per
|
||||
memory) will time out against Hindsight's default LLM timeout. **For any
|
||||
real use, run on a GPU.**
|
||||
|
||||
### NVIDIA
|
||||
|
||||
1. Switch the `llama` service image from `:server` to `:server-cuda`.
|
||||
2. Uncomment the `LLAMA_ARG_N_GPU_LAYERS: "999"` env var (offload all
|
||||
layers to GPU).
|
||||
3. Uncomment the `deploy.resources.reservations.devices` block.
|
||||
4. Install the [NVIDIA Container Toolkit](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html)
|
||||
on the host.
|
||||
|
||||
The compose file has all four spots marked with inline comments.
|
||||
|
||||
### Apple Silicon / ROCm / Vulkan
|
||||
|
||||
The official `ghcr.io/ggml-org/llama.cpp` image only ships CPU and CUDA
|
||||
variants. For Metal (Apple Silicon), ROCm (AMD), or Vulkan backends,
|
||||
build llama.cpp yourself with the appropriate flags and reference the
|
||||
image you build instead. Docker Desktop on macOS cannot pass through the
|
||||
host GPU to a Linux container in any case — for Apple Silicon, run
|
||||
llama-server directly on the host and only put Hindsight in Docker.
|
||||
|
||||
## Caveats
|
||||
|
||||
- llama.cpp's HTTP API is OpenAI-compatible but not 100% feature-parity.
|
||||
Function/tool calling support depends on the chat template baked into
|
||||
the GGUF; some retain/reflect flows may behave differently than against
|
||||
a hosted OpenAI model.
|
||||
- Small GGUFs (~3 B params) are useful for smoke testing but will
|
||||
underperform a hosted frontier model on retain quality. Use a larger
|
||||
GGUF (7-13 B params) for production-quality memory.
|
||||
- The `llama_models` named volume persists the GGUF across `docker
|
||||
compose down`/`up` so the model is downloaded once, not every restart.
|
||||
@@ -1,74 +0,0 @@
|
||||
name: hindsight-local-llm
|
||||
# Example: run Hindsight against a local llama.cpp server sidecar — fully
|
||||
# offline, no external API key needed.
|
||||
#
|
||||
# Pattern follows https://servicestack.net/posts/hosting-llama-server :
|
||||
# llama.cpp runs as its own container exposing an OpenAI-compatible HTTP
|
||||
# API, and Hindsight talks to it via the `openai` LLM provider with a
|
||||
# custom `base_url`. This means we can use the published Hindsight image
|
||||
# unchanged — no derived Dockerfile, no `llama-cpp-python` install on top.
|
||||
#
|
||||
# Quick start:
|
||||
# docker compose -f docker/docker-compose/local-llm/docker-compose.yaml up
|
||||
#
|
||||
# First boot downloads the default Gemma 4 E2B GGUF (~3.5 GB) into the
|
||||
# `llama_models` volume; subsequent boots reuse it.
|
||||
|
||||
services:
|
||||
llama:
|
||||
image: ghcr.io/ggml-org/llama.cpp:server
|
||||
container_name: hindsight-local-llm-llama
|
||||
environment:
|
||||
LLAMA_ARG_HOST: 0.0.0.0
|
||||
LLAMA_ARG_PORT: "8080"
|
||||
# Auto-download a small GGUF from HuggingFace on first start.
|
||||
# Override these to use a different model.
|
||||
LLAMA_ARG_HF_REPO: bartowski/google_gemma-4-E2B-it-GGUF
|
||||
LLAMA_ARG_HF_FILE: google_gemma-4-E2B-it-Q4_K_M.gguf
|
||||
LLAMA_ARG_CTX_SIZE: "8192"
|
||||
# Uncomment for NVIDIA GPU (and switch image to :server-cuda):
|
||||
# LLAMA_ARG_N_GPU_LAYERS: "999"
|
||||
volumes:
|
||||
# llama-server stores HuggingFace downloads under ~/.cache/huggingface
|
||||
# (not ~/.cache/llama.cpp), so mount the named volume there to avoid
|
||||
# re-downloading the GGUF on every recreate.
|
||||
- llama_models:/root/.cache/huggingface
|
||||
healthcheck:
|
||||
test: ["CMD-SHELL", "curl -fsS http://localhost:8080/health || exit 1"]
|
||||
interval: 10s
|
||||
timeout: 5s
|
||||
retries: 60
|
||||
start_period: 30s
|
||||
# For NVIDIA GPU acceleration, swap the image above to
|
||||
# `ghcr.io/ggml-org/llama.cpp:server-cuda` and uncomment:
|
||||
# deploy:
|
||||
# resources:
|
||||
# reservations:
|
||||
# devices:
|
||||
# - driver: nvidia
|
||||
# count: all
|
||||
# capabilities: [gpu]
|
||||
|
||||
hindsight:
|
||||
image: ghcr.io/vectorize-io/hindsight:latest
|
||||
container_name: hindsight-local-llm
|
||||
depends_on:
|
||||
llama:
|
||||
condition: service_healthy
|
||||
ports:
|
||||
- "8888:8888"
|
||||
- "9999:9999"
|
||||
environment:
|
||||
# llama-server is OpenAI-compatible, so use the `openai` provider and
|
||||
# point base_url at the sidecar. The API key is unused by llama-server
|
||||
# but Hindsight requires the env var to be set.
|
||||
HINDSIGHT_API_LLM_PROVIDER: openai
|
||||
HINDSIGHT_API_LLM_BASE_URL: http://llama:8080/v1
|
||||
HINDSIGHT_API_LLM_API_KEY: not-needed
|
||||
HINDSIGHT_API_LLM_MODEL: gemma-4-e2b-it
|
||||
volumes:
|
||||
- pg_data:/home/hindsight/.pg0
|
||||
|
||||
volumes:
|
||||
pg_data:
|
||||
llama_models:
|
||||
@@ -1,7 +0,0 @@
|
||||
# PostgreSQL with pgvector and ParadeDB pg_search extensions.
|
||||
#
|
||||
# The official ParadeDB image ships PostgreSQL with pg_search and pgvector
|
||||
# already installed, so no build steps are required. We pin to the PG17
|
||||
# variant for parity with the other Hindsight docker-compose examples
|
||||
# (vchord, pg_textsearch).
|
||||
FROM paradedb/paradedb:latest-pg17
|
||||
@@ -1,96 +0,0 @@
|
||||
name: hindsight
|
||||
# Docker Compose file for Hindsight with PostgreSQL and ParadeDB pg_search.
|
||||
#
|
||||
# pg_search is the only BM25 backend supported by Hindsight that works with
|
||||
# Citus, so this is the recommended setup for horizontally scaled deployments.
|
||||
#
|
||||
# Usage:
|
||||
# docker compose -f docker/docker-compose/pg_search/docker-compose.yaml up -d
|
||||
#
|
||||
# Required environment variables:
|
||||
# - HINDSIGHT_DB_PASSWORD: Password for the PostgreSQL user
|
||||
# - Configure LLM provider variables as needed (see the hindsight service)
|
||||
#
|
||||
# Optional environment variables with defaults:
|
||||
# - HINDSIGHT_VERSION: Hindsight application version (default: latest)
|
||||
# - HINDSIGHT_DB_USER: PostgreSQL user (default: hindsight_user)
|
||||
# - HINDSIGHT_DB_NAME: PostgreSQL database name (default: hindsight_db)
|
||||
# - HINDSIGHT_API_TEXT_SEARCH_EXTENSION_PG_SEARCH_TOKENIZER: ParadeDB pg_search
|
||||
# tokenizer for new BM25 indexes (default: empty, uses ParadeDB default)
|
||||
|
||||
services:
|
||||
db:
|
||||
# Use ParadeDB image which bundles pgvector + pg_search
|
||||
build:
|
||||
context: .
|
||||
dockerfile: Dockerfile
|
||||
container_name: hindsight-db
|
||||
restart: always
|
||||
ports:
|
||||
- "5437:5432"
|
||||
environment:
|
||||
POSTGRES_USER: ${HINDSIGHT_DB_USER:-hindsight_user}
|
||||
POSTGRES_PASSWORD: ${HINDSIGHT_DB_PASSWORD:-hindsight_password}
|
||||
POSTGRES_DB: ${HINDSIGHT_DB_NAME:-hindsight_db}
|
||||
volumes:
|
||||
- pg_data:/var/lib/postgresql/data
|
||||
networks:
|
||||
- hindsight-net
|
||||
|
||||
pg-search-init:
|
||||
build:
|
||||
context: .
|
||||
dockerfile: Dockerfile
|
||||
depends_on:
|
||||
- db
|
||||
environment:
|
||||
- PGPASSWORD=${HINDSIGHT_DB_PASSWORD:-hindsight_password}
|
||||
command: >
|
||||
bash -c "
|
||||
echo 'Waiting for PostgreSQL to be ready...';
|
||||
until pg_isready -h hindsight-db -p 5432 -U hindsight_user; do
|
||||
echo 'PostgreSQL is unavailable - sleeping';
|
||||
sleep 2;
|
||||
done;
|
||||
echo 'PostgreSQL is ready - creating hindsight_db database';
|
||||
psql -h hindsight-db -p 5432 -U hindsight_user -c 'CREATE DATABASE hindsight_db;' 2>/dev/null || echo 'Database already exists';
|
||||
echo 'Creating extensions in hindsight_db database';
|
||||
psql -h hindsight-db -p 5432 -U hindsight_user -d hindsight_db -c 'CREATE EXTENSION IF NOT EXISTS vector CASCADE;';
|
||||
psql -h hindsight-db -p 5432 -U hindsight_user -d hindsight_db -c 'CREATE EXTENSION IF NOT EXISTS pg_search CASCADE;';
|
||||
echo 'Database and extensions created successfully';
|
||||
"
|
||||
restart: "no"
|
||||
networks:
|
||||
- hindsight-net
|
||||
|
||||
hindsight:
|
||||
image: ghcr.io/vectorize-io/hindsight:${HINDSIGHT_VERSION:-latest}
|
||||
container_name: hindsight-app
|
||||
ports:
|
||||
- "8888:8888"
|
||||
- "9999:9999"
|
||||
environment:
|
||||
# LLM Configuration
|
||||
HINDSIGHT_API_LLM_PROVIDER: ${HINDSIGHT_API_LLM_PROVIDER:-openai}
|
||||
HINDSIGHT_API_LLM_API_KEY: ${OPENAI_API_KEY:-your-api-key}
|
||||
|
||||
# Database Configuration
|
||||
HINDSIGHT_API_DATABASE_URL: postgresql://${HINDSIGHT_DB_USER:-hindsight_user}:${HINDSIGHT_DB_PASSWORD:-hindsight_password}@db:5432/${HINDSIGHT_DB_NAME:-hindsight_db}
|
||||
|
||||
# Vector and Text Search Extensions
|
||||
HINDSIGHT_API_VECTOR_EXTENSION: pgvector
|
||||
HINDSIGHT_API_TEXT_SEARCH_EXTENSION: pg_search
|
||||
HINDSIGHT_API_TEXT_SEARCH_EXTENSION_PG_SEARCH_TOKENIZER: ${HINDSIGHT_API_TEXT_SEARCH_EXTENSION_PG_SEARCH_TOKENIZER:-}
|
||||
|
||||
depends_on:
|
||||
- db
|
||||
networks:
|
||||
- hindsight-net
|
||||
|
||||
|
||||
networks:
|
||||
hindsight-net:
|
||||
driver: bridge
|
||||
|
||||
volumes:
|
||||
pg_data:
|
||||
@@ -1,23 +0,0 @@
|
||||
# PostgreSQL with pgvector and pgroonga extensions.
|
||||
#
|
||||
# pgroonga is a multilingual full-text search extension built on Groonga.
|
||||
# It works out of the box for CJK (Chinese, Japanese, Korean) and other
|
||||
# non-whitespace-segmented languages via the TokenBigram tokenizer.
|
||||
FROM groonga/pgroonga:latest-debian-pg17
|
||||
|
||||
# Install pgvector on top of the pgroonga base image (which already provides
|
||||
# pgroonga and the Groonga library).
|
||||
RUN apt-get update && apt-get install -y --no-install-recommends \
|
||||
build-essential \
|
||||
git \
|
||||
postgresql-server-dev-17 \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
RUN cd /tmp && \
|
||||
git clone --branch v0.8.0 https://github.com/pgvector/pgvector.git && \
|
||||
cd pgvector && \
|
||||
make && \
|
||||
make install
|
||||
|
||||
RUN rm -rf /tmp/pgvector && \
|
||||
apt-get purge -y --auto-remove build-essential git postgresql-server-dev-17
|
||||
@@ -1,91 +0,0 @@
|
||||
name: hindsight
|
||||
# Docker Compose file for Hindsight with PostgreSQL and pgroonga
|
||||
#
|
||||
# pgroonga provides multilingual BM25 indexing that works out of the box for
|
||||
# CJK (Chinese, Japanese, Korean) and other non-whitespace-segmented languages.
|
||||
# Use this recipe if your bank content is not English/European.
|
||||
#
|
||||
# docker compose -f docker/docker-compose/pgroonga/docker-compose.yaml down && \
|
||||
# sleep 2 && \
|
||||
# docker compose -f docker/docker-compose/pgroonga/docker-compose.yaml up -d
|
||||
#
|
||||
# Optional environment variables with defaults:
|
||||
# - HINDSIGHT_VERSION: Hindsight application version (default: latest)
|
||||
# - HINDSIGHT_DB_USER: PostgreSQL user (default: hindsight_user)
|
||||
# - HINDSIGHT_DB_NAME: PostgreSQL database name (default: hindsight_db)
|
||||
# - HINDSIGHT_DB_PASSWORD: PostgreSQL password (default: hindsight_password)
|
||||
|
||||
services:
|
||||
db:
|
||||
build:
|
||||
context: .
|
||||
dockerfile: Dockerfile
|
||||
container_name: hindsight-db
|
||||
restart: always
|
||||
ports:
|
||||
- "5439:5432"
|
||||
environment:
|
||||
POSTGRES_USER: ${HINDSIGHT_DB_USER:-hindsight_user}
|
||||
POSTGRES_PASSWORD: ${HINDSIGHT_DB_PASSWORD:-hindsight_password}
|
||||
POSTGRES_DB: ${HINDSIGHT_DB_NAME:-hindsight_db}
|
||||
volumes:
|
||||
- pg_data:/var/lib/postgresql/data
|
||||
networks:
|
||||
- hindsight-net
|
||||
|
||||
pgroonga-init:
|
||||
build:
|
||||
context: .
|
||||
dockerfile: Dockerfile
|
||||
depends_on:
|
||||
- db
|
||||
environment:
|
||||
- PGPASSWORD=${HINDSIGHT_DB_PASSWORD:-hindsight_password}
|
||||
command: >
|
||||
bash -c "
|
||||
echo 'Waiting for PostgreSQL to be ready...';
|
||||
until pg_isready -h hindsight-db -p 5432 -U hindsight_user; do
|
||||
echo 'PostgreSQL is unavailable - sleeping';
|
||||
sleep 2;
|
||||
done;
|
||||
echo 'PostgreSQL is ready - creating hindsight_db database';
|
||||
psql -h hindsight-db -p 5432 -U hindsight_user -c 'CREATE DATABASE hindsight_db;' 2>/dev/null || echo 'Database already exists';
|
||||
echo 'Creating extensions in hindsight_db database';
|
||||
psql -h hindsight-db -p 5432 -U hindsight_user -d hindsight_db -c 'CREATE EXTENSION IF NOT EXISTS vector CASCADE;';
|
||||
psql -h hindsight-db -p 5432 -U hindsight_user -d hindsight_db -c 'CREATE EXTENSION IF NOT EXISTS pgroonga CASCADE;';
|
||||
echo 'Database and extensions created successfully';
|
||||
"
|
||||
restart: "no"
|
||||
networks:
|
||||
- hindsight-net
|
||||
|
||||
hindsight:
|
||||
image: ghcr.io/vectorize-io/hindsight:${HINDSIGHT_VERSION:-latest}
|
||||
container_name: hindsight-app
|
||||
ports:
|
||||
- "8888:8888"
|
||||
- "9999:9999"
|
||||
environment:
|
||||
# LLM Configuration
|
||||
HINDSIGHT_API_LLM_PROVIDER: ${HINDSIGHT_API_LLM_PROVIDER:-openai}
|
||||
HINDSIGHT_API_LLM_API_KEY: ${OPENAI_API_KEY:-your-api-key}
|
||||
|
||||
# Database Configuration
|
||||
HINDSIGHT_API_DATABASE_URL: postgresql://${HINDSIGHT_DB_USER:-hindsight_user}:${HINDSIGHT_DB_PASSWORD:-hindsight_password}@db:5432/${HINDSIGHT_DB_NAME:-hindsight_db}
|
||||
|
||||
# Vector and Text Search Extensions
|
||||
HINDSIGHT_API_VECTOR_EXTENSION: pgvector
|
||||
HINDSIGHT_API_TEXT_SEARCH_EXTENSION: pgroonga
|
||||
|
||||
depends_on:
|
||||
- db
|
||||
networks:
|
||||
- hindsight-net
|
||||
|
||||
|
||||
networks:
|
||||
hindsight-net:
|
||||
driver: bridge
|
||||
|
||||
volumes:
|
||||
pg_data:
|
||||
@@ -10,45 +10,19 @@ set -e
|
||||
# loss scenarios where a container restart caused the data directory to be
|
||||
# wiped despite a volume mount being present.
|
||||
# =============================================================================
|
||||
pg0_has_pg_version() {
|
||||
local pg0_data_dir="$1"
|
||||
|
||||
# pg0 has used more than one on-disk layout. Newer standalone images keep
|
||||
# PostgreSQL data under instances/<name>/data, while older volumes may have
|
||||
# placed PG_VERSION at or one level below the mount.
|
||||
[ -f "$pg0_data_dir/PG_VERSION" ] && return 0
|
||||
compgen -G "$pg0_data_dir"/*/PG_VERSION > /dev/null 2>&1 && return 0
|
||||
compgen -G "$pg0_data_dir"/instances/*/data/PG_VERSION > /dev/null 2>&1 && return 0
|
||||
|
||||
return 1
|
||||
}
|
||||
|
||||
check_pg0_data_integrity() {
|
||||
local pg0_data_dir="$1"
|
||||
|
||||
if [ ! -d "$pg0_data_dir" ]; then
|
||||
return 0
|
||||
fi
|
||||
|
||||
PG0_DATA_DIR="${HOME}/.pg0"
|
||||
if [ -d "$PG0_DATA_DIR" ]; then
|
||||
# Look for actual PostgreSQL data directories (pg0 creates subdirs per instance)
|
||||
if pg0_has_pg_version "$pg0_data_dir"; then
|
||||
echo "✅ Existing pg0 data directory detected at $pg0_data_dir"
|
||||
elif [ "$(ls -A "$pg0_data_dir" 2>/dev/null)" ]; then
|
||||
echo "⚠️ WARNING: pg0 data directory exists at $pg0_data_dir but no PG_VERSION found."
|
||||
if compgen -G "$PG0_DATA_DIR"/*/PG_VERSION > /dev/null 2>&1; then
|
||||
echo "✅ Existing pg0 data directory detected at $PG0_DATA_DIR"
|
||||
elif [ "$(ls -A "$PG0_DATA_DIR" 2>/dev/null)" ]; then
|
||||
echo "⚠️ WARNING: pg0 data directory exists at $PG0_DATA_DIR but no PG_VERSION found."
|
||||
echo " This may indicate data corruption or an incomplete previous shutdown."
|
||||
echo " If you see all migrations running from scratch after this, your data may have been lost."
|
||||
echo " See: https://github.com/vectorize-io/hindsight/issues/675"
|
||||
fi
|
||||
|
||||
return 0
|
||||
}
|
||||
|
||||
if [ "${HINDSIGHT_START_ALL_SOURCE_ONLY:-false}" = "true" ]; then
|
||||
return 0 2>/dev/null || exit 0
|
||||
fi
|
||||
|
||||
check_pg0_data_integrity "${HOME}/.pg0"
|
||||
|
||||
# Service flags (default to true if not set)
|
||||
ENABLE_API="${HINDSIGHT_ENABLE_API:-true}"
|
||||
ENABLE_CP="${HINDSIGHT_ENABLE_CP:-true}"
|
||||
@@ -182,7 +156,7 @@ PIDS=()
|
||||
# Start API if enabled
|
||||
if [ "$ENABLE_API" = "true" ]; then
|
||||
cd /app/api
|
||||
API_HEALTH_URL="${HINDSIGHT_API_HEALTH_URL:-http://localhost:${HINDSIGHT_API_PORT:-8888}/health}"
|
||||
API_HEALTH_URL="${HINDSIGHT_API_HEALTH_URL:-http://localhost:8888/health}"
|
||||
API_STARTUP_WAIT_SECONDS="${HINDSIGHT_API_STARTUP_WAIT_SECONDS:-300}"
|
||||
|
||||
# Run API directly - Python's PYTHONUNBUFFERED=1 handles output buffering
|
||||
|
||||
@@ -1,73 +0,0 @@
|
||||
#!/bin/bash
|
||||
set -euo pipefail
|
||||
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
|
||||
HINDSIGHT_START_ALL_SOURCE_ONLY=true
|
||||
source "$SCRIPT_DIR/start-all.sh"
|
||||
unset HINDSIGHT_START_ALL_SOURCE_ONLY
|
||||
|
||||
TMP_DIR="$(mktemp -d)"
|
||||
trap 'rm -rf "$TMP_DIR"' EXIT
|
||||
|
||||
assert_contains() {
|
||||
local output="$1"
|
||||
local expected="$2"
|
||||
|
||||
if [[ "$output" != *"$expected"* ]]; then
|
||||
echo "Expected output to contain: $expected"
|
||||
echo "Actual output:"
|
||||
echo "$output"
|
||||
exit 1
|
||||
fi
|
||||
}
|
||||
|
||||
assert_not_contains() {
|
||||
local output="$1"
|
||||
local unexpected="$2"
|
||||
|
||||
if [[ "$output" == *"$unexpected"* ]]; then
|
||||
echo "Expected output not to contain: $unexpected"
|
||||
echo "Actual output:"
|
||||
echo "$output"
|
||||
exit 1
|
||||
fi
|
||||
}
|
||||
|
||||
assert_empty() {
|
||||
local output="$1"
|
||||
|
||||
if [ -n "$output" ]; then
|
||||
echo "Expected no output, got:"
|
||||
echo "$output"
|
||||
exit 1
|
||||
fi
|
||||
}
|
||||
|
||||
mkdir -p "$TMP_DIR/empty"
|
||||
assert_empty "$(check_pg0_data_integrity "$TMP_DIR/empty")"
|
||||
|
||||
mkdir -p "$TMP_DIR/direct"
|
||||
touch "$TMP_DIR/direct/PG_VERSION"
|
||||
direct_output="$(check_pg0_data_integrity "$TMP_DIR/direct")"
|
||||
assert_contains "$direct_output" "Existing pg0 data directory detected"
|
||||
assert_not_contains "$direct_output" "WARNING"
|
||||
|
||||
mkdir -p "$TMP_DIR/legacy/instance"
|
||||
touch "$TMP_DIR/legacy/instance/PG_VERSION"
|
||||
legacy_output="$(check_pg0_data_integrity "$TMP_DIR/legacy")"
|
||||
assert_contains "$legacy_output" "Existing pg0 data directory detected"
|
||||
assert_not_contains "$legacy_output" "WARNING"
|
||||
|
||||
mkdir -p "$TMP_DIR/nested/instances/hindsight/data"
|
||||
touch "$TMP_DIR/nested/instances/hindsight/data/PG_VERSION"
|
||||
nested_output="$(check_pg0_data_integrity "$TMP_DIR/nested")"
|
||||
assert_contains "$nested_output" "Existing pg0 data directory detected"
|
||||
assert_not_contains "$nested_output" "WARNING"
|
||||
|
||||
mkdir -p "$TMP_DIR/nonempty/instances/hindsight"
|
||||
touch "$TMP_DIR/nonempty/instances/hindsight/instance.json"
|
||||
nonempty_output="$(check_pg0_data_integrity "$TMP_DIR/nonempty")"
|
||||
assert_contains "$nonempty_output" "WARNING: pg0 data directory exists"
|
||||
|
||||
echo "start-all pg0 integrity checks passed"
|
||||
@@ -0,0 +1,6 @@
|
||||
dependencies:
|
||||
- name: postgresql
|
||||
repository: https://charts.bitnami.com/bitnami
|
||||
version: 15.5.38
|
||||
digest: sha256:f67c7612736803ece8a669f8ca6b0555f3b78557bc0ecb732aa2e43f0df7750d
|
||||
generated: "2025-12-10T17:20:57.058794+01:00"
|
||||
@@ -2,8 +2,8 @@ apiVersion: v2
|
||||
name: hindsight
|
||||
description: Hindsight helm chart
|
||||
type: application
|
||||
version: 0.7.1
|
||||
appVersion: "0.7.1"
|
||||
version: 0.6.2
|
||||
appVersion: "0.6.2"
|
||||
keywords:
|
||||
- ai
|
||||
- memory
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@vectorize-io/hindsight-all",
|
||||
"version": "0.7.1",
|
||||
"version": "0.6.2",
|
||||
"description": "Node.js programmatic lifecycle manager for Hindsight — embeds a local hindsight daemon in a Node application. Pair with @vectorize-io/hindsight-client for memory operations.",
|
||||
"main": "dist/index.js",
|
||||
"types": "dist/index.d.ts",
|
||||
|
||||
@@ -4,12 +4,12 @@ build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "hindsight-all-slim"
|
||||
version = "0.7.1"
|
||||
version = "0.6.2"
|
||||
description = "Hindsight: Agent Memory That Works Like Human Memory - Slim All-in-One Bundle"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.11"
|
||||
dependencies = [
|
||||
"hindsight-api-slim==0.7.1",
|
||||
"hindsight-api-slim==0.6.2",
|
||||
"hindsight-client>=0.0.7",
|
||||
"hindsight-embed>=0.1.0",
|
||||
]
|
||||
|
||||
@@ -4,12 +4,12 @@ build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "hindsight-all"
|
||||
version = "0.7.1"
|
||||
version = "0.6.2"
|
||||
description = "Hindsight: Agent Memory That Works Like Human Memory - All-in-One Bundle"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.11"
|
||||
dependencies = [
|
||||
"hindsight-api-slim[all]==0.7.1",
|
||||
"hindsight-api-slim[all]==0.6.2",
|
||||
"hindsight-client>=0.0.7",
|
||||
"hindsight-embed>=0.1.0",
|
||||
]
|
||||
@@ -21,7 +21,7 @@ hindsight-embed = { workspace = true }
|
||||
|
||||
[project.optional-dependencies]
|
||||
local-llm = [
|
||||
"hindsight-api-slim[local-llm]==0.7.1",
|
||||
"hindsight-api-slim[local-llm]==0.6.2",
|
||||
]
|
||||
test = [
|
||||
"pytest>=7.0.0",
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -4,13 +4,6 @@ Memory System for AI Agents.
|
||||
Temporal + Semantic Memory Architecture using PostgreSQL with pgvector.
|
||||
"""
|
||||
|
||||
# Cap native ML thread pools (OpenBLAS/OpenMP/MKL) before any import pulls in
|
||||
# numpy/torch/onnxruntime — they read these env vars only at load time. See
|
||||
# hindsight_api/_thread_limits.py for the rationale.
|
||||
from ._thread_limits import apply_default_thread_limits
|
||||
|
||||
apply_default_thread_limits()
|
||||
|
||||
from .config import HindsightConfig, get_config
|
||||
from .engine.cross_encoder import CrossEncoderModel, LocalSTCrossEncoder, RemoteTEICrossEncoder
|
||||
from .engine.embeddings import Embeddings, LocalSTEmbeddings, RemoteTEIEmbeddings
|
||||
@@ -53,4 +46,4 @@ __all__ = [
|
||||
"RemoteTEICrossEncoder",
|
||||
"LLMConfig",
|
||||
]
|
||||
__version__ = "0.7.1"
|
||||
__version__ = "0.6.2"
|
||||
|
||||
@@ -1,85 +0,0 @@
|
||||
"""Helpers for ParadeDB pg_search index configuration."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from collections.abc import Sequence
|
||||
|
||||
PG_SEARCH_TOKENIZER_ENV = "HINDSIGHT_API_TEXT_SEARCH_EXTENSION_PG_SEARCH_TOKENIZER"
|
||||
|
||||
_SIMPLE_TOKENIZERS = {
|
||||
"unicode_words",
|
||||
"simple",
|
||||
"whitespace",
|
||||
"literal",
|
||||
"literal_normalized",
|
||||
"chinese_compatible",
|
||||
"icu",
|
||||
"jieba",
|
||||
"source_code",
|
||||
}
|
||||
|
||||
_TOKENIZER_ALIASES = {
|
||||
"chinese_lindera": "lindera(chinese)",
|
||||
"japanese_lindera": "lindera(japanese)",
|
||||
"korean_lindera": "lindera(korean)",
|
||||
"lindera_chinese": "lindera(chinese)",
|
||||
"lindera_japanese": "lindera(japanese)",
|
||||
"lindera_korean": "lindera(korean)",
|
||||
}
|
||||
|
||||
|
||||
def normalize_pg_search_tokenizer(value: str | None) -> str:
|
||||
"""Validate and normalize a ParadeDB pg_search tokenizer setting.
|
||||
|
||||
Returns an empty string when unset. The returned value is safe to embed after
|
||||
``pdb.`` in a CREATE INDEX expression.
|
||||
"""
|
||||
|
||||
tokenizer = (value or "").strip().lower()
|
||||
if not tokenizer:
|
||||
return ""
|
||||
|
||||
if tokenizer in _TOKENIZER_ALIASES:
|
||||
return _TOKENIZER_ALIASES[tokenizer]
|
||||
|
||||
if tokenizer in _SIMPLE_TOKENIZERS:
|
||||
return tokenizer
|
||||
|
||||
lindera_match = re.fullmatch(r"lindera\((chinese|japanese|korean)\)", tokenizer)
|
||||
if lindera_match:
|
||||
return tokenizer
|
||||
|
||||
ngram_match = re.fullmatch(r"(ngram|edge_ngram)\((\d{1,3}),\s*(\d{1,3})\)", tokenizer)
|
||||
if ngram_match:
|
||||
kind, min_gram, max_gram = ngram_match.groups()
|
||||
min_value = int(min_gram)
|
||||
max_value = int(max_gram)
|
||||
if min_value <= 0 or min_value > max_value:
|
||||
raise ValueError(
|
||||
f"Invalid {PG_SEARCH_TOKENIZER_ENV}: {value!r}. "
|
||||
"ngram and edge_ngram require positive min/max gram sizes with min <= max."
|
||||
)
|
||||
return f"{kind}({min_value},{max_value})"
|
||||
|
||||
raise ValueError(
|
||||
f"Invalid {PG_SEARCH_TOKENIZER_ENV}: {value!r}. "
|
||||
"Supported values are: unicode_words, simple, whitespace, literal, "
|
||||
"literal_normalized, chinese_compatible, icu, jieba, source_code, "
|
||||
"chinese_lindera, japanese_lindera, korean_lindera, or "
|
||||
"lindera(chinese|japanese|korean), ngram(min,max), or edge_ngram(min,max)."
|
||||
)
|
||||
|
||||
|
||||
def pg_search_bm25_columns(
|
||||
key_field: str,
|
||||
text_fields: Sequence[str],
|
||||
tokenizer: str | None,
|
||||
) -> str:
|
||||
"""Build a ParadeDB BM25 column list for CREATE INDEX."""
|
||||
|
||||
normalized = normalize_pg_search_tokenizer(tokenizer)
|
||||
if not normalized:
|
||||
return ", ".join([key_field, *text_fields])
|
||||
|
||||
return ", ".join([key_field, *(f"({field}::pdb.{normalized})" for field in text_fields)])
|
||||
@@ -1,107 +0,0 @@
|
||||
"""Process-level caps for native ML thread pools.
|
||||
|
||||
OpenBLAS, OpenMP, and MKL each spawn a worker pool sized to the host CPU count
|
||||
the first time they are loaded (numpy pulls in OpenBLAS eagerly; torch and
|
||||
onnxruntime load their pools lazily on first inference). Hindsight already
|
||||
parallelizes at the request level via thread-pool executors (embeddings on the
|
||||
default executor, the reranker on its own pool), so these native intra-op pools
|
||||
oversubscribe the CPU: on a many-core host the process accumulates 100+ native
|
||||
threads, which inflates memory and, under contention, can degrade throughput.
|
||||
|
||||
We bound each pool to ``_MAX_NATIVE_THREADS`` (or the available CPU count, if
|
||||
smaller). "Available" is the CPU budget actually granted to the process, not
|
||||
``os.cpu_count()``: in a CPU-limited container ``os.cpu_count()`` still reports
|
||||
the host's cores, so sizing pools by it oversubscribes the container's real
|
||||
quota — the exact failure mode this guards against. We therefore take the
|
||||
smallest of the CPU-affinity set, the cgroup CPU quota, and ``os.cpu_count()``.
|
||||
|
||||
Every cap is applied with ``setdefault`` so an operator who has deliberately
|
||||
tuned one of these variables keeps their value. This must run *before* numpy,
|
||||
torch, or onnxruntime are imported — those libraries read the variables only at
|
||||
load time — which is why it is invoked at the very top of
|
||||
``hindsight_api/__init__.py``, ahead of the package's other imports.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
|
||||
# Native threading env vars, each read by the respective library at load time.
|
||||
_NATIVE_THREAD_VARS = (
|
||||
"OMP_NUM_THREADS", # OpenMP — torch, onnxruntime, some BLAS builds
|
||||
"OPENBLAS_NUM_THREADS", # OpenBLAS — numpy's default BLAS
|
||||
"MKL_NUM_THREADS", # Intel MKL — numpy/torch when MKL-backed
|
||||
"NUMEXPR_NUM_THREADS", # numexpr expression engine
|
||||
)
|
||||
|
||||
# Upper bound on intra-op threads per native pool. Bounds runaway growth on
|
||||
# many-core hosts without serialising single-request inference.
|
||||
_MAX_NATIVE_THREADS = 16
|
||||
|
||||
|
||||
def _quota_to_cpus(quota: int, period: int) -> int | None:
|
||||
"""Whole CPUs from a CFS quota/period pair, or None if unlimited."""
|
||||
if quota > 0 and period > 0:
|
||||
# Floor (never round up) so we never exceed the granted budget.
|
||||
return max(1, quota // period)
|
||||
return None
|
||||
|
||||
|
||||
def _parse_cgroup_v2_cpu_max(text: str) -> int | None:
|
||||
"""Parse cgroup v2 ``cpu.max`` ("<quota> <period>", or "max <period>")."""
|
||||
parts = text.split()
|
||||
if len(parts) >= 2 and parts[0] != "max":
|
||||
try:
|
||||
return _quota_to_cpus(int(parts[0]), int(parts[1]))
|
||||
except ValueError:
|
||||
return None
|
||||
return None
|
||||
|
||||
|
||||
def _cgroup_cpu_quota() -> int | None:
|
||||
"""Effective CPUs from the cgroup CPU quota, or None if unlimited/unknown."""
|
||||
try: # cgroup v2
|
||||
with open("/sys/fs/cgroup/cpu.max") as fh:
|
||||
return _parse_cgroup_v2_cpu_max(fh.read())
|
||||
except OSError:
|
||||
pass
|
||||
try: # cgroup v1
|
||||
with open("/sys/fs/cgroup/cpu/cpu.cfs_quota_us") as fh:
|
||||
quota = int(fh.read())
|
||||
with open("/sys/fs/cgroup/cpu/cpu.cfs_period_us") as fh:
|
||||
period = int(fh.read())
|
||||
return _quota_to_cpus(quota, period)
|
||||
except (OSError, ValueError):
|
||||
return None
|
||||
|
||||
|
||||
def _available_cpu_count() -> int:
|
||||
"""CPUs actually available to this process.
|
||||
|
||||
The smallest of the CPU-affinity set (cpuset / ``--cpuset-cpus``), the
|
||||
cgroup CPU quota (``--cpus``), and ``os.cpu_count()`` — each captures a
|
||||
different way the budget can be constrained, and the last alone overcounts
|
||||
inside a limited container.
|
||||
"""
|
||||
candidates = [os.cpu_count() or 1]
|
||||
if hasattr(os, "sched_getaffinity"):
|
||||
try:
|
||||
candidates.append(len(os.sched_getaffinity(0)))
|
||||
except OSError:
|
||||
pass
|
||||
quota = _cgroup_cpu_quota()
|
||||
if quota is not None:
|
||||
candidates.append(quota)
|
||||
return max(1, min(candidates))
|
||||
|
||||
|
||||
def default_native_thread_count() -> int:
|
||||
"""Per-pool cap: ``_MAX_NATIVE_THREADS``, or available CPUs if fewer."""
|
||||
return min(_MAX_NATIVE_THREADS, _available_cpu_count())
|
||||
|
||||
|
||||
def apply_default_thread_limits() -> None:
|
||||
"""Cap native ML thread pools unless the operator has set the var already."""
|
||||
value = str(default_native_thread_count())
|
||||
for var in _NATIVE_THREAD_VARS:
|
||||
os.environ.setdefault(var, value)
|
||||
@@ -3,7 +3,6 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
|
||||
from sqlalchemy import text
|
||||
from sqlalchemy.engine import Connection
|
||||
@@ -35,7 +34,7 @@ _INDEX_USING_CLAUSES = {
|
||||
"pgvector": "USING hnsw (embedding vector_cosine_ops)",
|
||||
"pgvectorscale": "USING diskann (embedding vector_cosine_ops) WITH (num_neighbors = 50)",
|
||||
"pg_diskann": "USING diskann (embedding vector_cosine_ops) WITH (max_neighbors = 50)",
|
||||
"vchord": "USING vchordrq (embedding vector_cosine_ops)",
|
||||
"vchord": "USING vchordrq (embedding vector_l2_ops)",
|
||||
"scann": "USING scann (embedding cosine) WITH (mode = 'AUTO')",
|
||||
}
|
||||
|
||||
@@ -47,32 +46,6 @@ _INDEX_TYPE_KEYWORDS = {
|
||||
"scann": "scann",
|
||||
}
|
||||
|
||||
# Per-backend ANN search-time tuning GUCs. Each entry is a tuple of
|
||||
# (guc_name, value) pairs the caller can apply with SET or SET LOCAL.
|
||||
#
|
||||
# - pgvector exposes hnsw.ef_search. The 60 / 200 pair is unchanged from the
|
||||
# pre-dispatcher code (internal benchmarks tuned around our embedding count
|
||||
# and recall floor; see the link_utils / pool init call sites for the
|
||||
# latency-vs-recall framing).
|
||||
# - vchord exposes vchordrq.probes (no default; see VectorChord issue #392)
|
||||
# and vchordrq.epsilon (default 1.9). probes = 10 / 30 are starting
|
||||
# defaults pending a workload-specific sweep — vchordrq's recall curve
|
||||
# shape differs from HNSW's, so the pgvector numbers don't translate
|
||||
# directly. Revisit with a per-cluster benchmark once we have production
|
||||
# recall data; until then these are deliberately conservative on the
|
||||
# high-recall path. We leave epsilon at its default; tightening it is a
|
||||
# separate trade-off.
|
||||
# - pgvectorscale / pg_diskann / scann do not expose an equivalent per-statement
|
||||
# knob in the engine today, so the dispatcher returns no statements for them.
|
||||
_ANN_TUNING_LOW_LATENCY: dict[str, tuple[tuple[str, str], ...]] = {
|
||||
"pgvector": (("hnsw.ef_search", "60"),),
|
||||
"vchord": (("vchordrq.probes", "10"),),
|
||||
}
|
||||
_ANN_TUNING_HIGH_RECALL: dict[str, tuple[tuple[str, str], ...]] = {
|
||||
"pgvector": (("hnsw.ef_search", "200"),),
|
||||
"vchord": (("vchordrq.probes", "30"),),
|
||||
}
|
||||
|
||||
_EXTENSION_INSTALL_SQL = {
|
||||
"pgvector": ("CREATE EXTENSION IF NOT EXISTS vector",),
|
||||
"pgvectorscale": (
|
||||
@@ -94,18 +67,6 @@ _INSTALL_HINTS = {
|
||||
}
|
||||
|
||||
|
||||
def configured_vector_extension() -> str:
|
||||
"""Return the user-configured vector backend extension.
|
||||
|
||||
Reads ``HINDSIGHT_API_VECTOR_EXTENSION`` (default ``"pgvector"``) and
|
||||
validates it via :func:`validate_extension`. This is the single source of
|
||||
truth for runtime code that needs to dispatch behaviour by vector backend;
|
||||
callers should prefer this over reading the env var directly, so the
|
||||
default value and the lookup mechanism live in one place.
|
||||
"""
|
||||
return validate_extension(os.getenv("HINDSIGHT_API_VECTOR_EXTENSION", "pgvector"))
|
||||
|
||||
|
||||
def validate_extension(name: str) -> str:
|
||||
"""Return a normalized configurable vector extension name or raise.
|
||||
|
||||
@@ -154,25 +115,6 @@ def should_defer_index_creation(ext: str, row_count: int) -> bool:
|
||||
return minimum_rows > 0 and row_count < minimum_rows
|
||||
|
||||
|
||||
def ann_search_tuning_settings(ext: str, *, kind: str) -> tuple[tuple[str, str], ...]:
|
||||
"""Return per-backend (guc_name, value) pairs for ANN search-time tuning.
|
||||
|
||||
``kind`` is ``"low_latency"`` for retain-side link probing (smaller probe
|
||||
count, lower recall, lower latency) and ``"high_recall"`` for connection
|
||||
init in the pool (larger probe count, higher recall). Callers wrap each
|
||||
pair with ``SET LOCAL`` or ``SET`` themselves so the same dispatcher works
|
||||
for both transaction-scoped and session-scoped use. Returns an empty tuple
|
||||
for backends without an equivalent knob.
|
||||
"""
|
||||
if kind == "low_latency":
|
||||
table = _ANN_TUNING_LOW_LATENCY
|
||||
elif kind == "high_recall":
|
||||
table = _ANN_TUNING_HIGH_RECALL
|
||||
else:
|
||||
raise ValueError(f"Unknown ANN tuning kind: {kind!r}")
|
||||
return table.get(_normalize_resolved(ext), ())
|
||||
|
||||
|
||||
def uses_per_bank_vector_indexes(ext: str) -> bool:
|
||||
"""Return whether the backend should create per-bank partial vector indexes."""
|
||||
return _normalize_resolved(ext) != "scann"
|
||||
|
||||
@@ -30,17 +30,8 @@ logger = logging.getLogger(__name__)
|
||||
|
||||
app = typer.Typer(name="hindsight-admin", help="Hindsight administrative commands")
|
||||
|
||||
# Tables to backup/restore in foreign-key dependency order (parents first).
|
||||
# Restore COPYs in this order and TRUNCATEs in reverse, so every child must
|
||||
# appear after the tables it references.
|
||||
#
|
||||
# This must cover EVERY persistent PostgreSQL table in the schema — a missing
|
||||
# entry silently drops that table's data on restore (and, worse, restore's
|
||||
# `TRUNCATE banks CASCADE` wipes any FK-to-banks child like mental_models even
|
||||
# when it was never backed up). test_admin_backup_restore.py asserts this list
|
||||
# equals the live schema's tables, so adding a migration that creates a table
|
||||
# without adding it here fails CI. Oracle-only tables (e.g. observation_sources)
|
||||
# are intentionally absent — admin backup/restore is PostgreSQL-only.
|
||||
# Tables to backup/restore in dependency order
|
||||
# Import must happen in this order due to foreign key constraints
|
||||
BACKUP_TABLES = [
|
||||
"banks",
|
||||
"documents",
|
||||
@@ -50,13 +41,6 @@ BACKUP_TABLES = [
|
||||
"unit_entities",
|
||||
"entity_cooccurrences",
|
||||
"memory_links",
|
||||
"mental_models",
|
||||
"directives",
|
||||
"async_operations",
|
||||
"webhooks",
|
||||
"file_storage",
|
||||
"audit_log",
|
||||
"graph_maintenance_queue",
|
||||
]
|
||||
|
||||
MANIFEST_VERSION = "1"
|
||||
@@ -284,7 +268,6 @@ async def _run_migration(
|
||||
ensure_text_search_extension(
|
||||
resolved_url,
|
||||
text_search_extension=config.text_search_extension,
|
||||
pg_search_tokenizer=config.text_search_extension_pg_search_tokenizer,
|
||||
schema=schema,
|
||||
)
|
||||
|
||||
|
||||
@@ -15,11 +15,6 @@ from pgvector.sqlalchemy import Vector
|
||||
from sqlalchemy import text
|
||||
from sqlalchemy.dialects import postgresql
|
||||
|
||||
from hindsight_api._pg_search import (
|
||||
PG_SEARCH_TOKENIZER_ENV,
|
||||
normalize_pg_search_tokenizer,
|
||||
pg_search_bm25_columns,
|
||||
)
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
@@ -88,7 +83,7 @@ def _vector_index_using_clause(ext: str) -> str:
|
||||
if ext == "pg_diskann":
|
||||
return "USING diskann (embedding vector_cosine_ops) WITH (max_neighbors = 50)"
|
||||
if ext == "vchord":
|
||||
return "USING vchordrq (embedding vector_cosine_ops)"
|
||||
return "USING vchordrq (embedding vector_l2_ops)"
|
||||
if ext == "scann":
|
||||
return "USING scann (embedding cosine) WITH (mode = 'AUTO')"
|
||||
return "USING hnsw (embedding vector_cosine_ops)"
|
||||
@@ -96,14 +91,9 @@ def _vector_index_using_clause(ext: str) -> str:
|
||||
|
||||
def _detect_text_search_extension() -> str:
|
||||
"""
|
||||
Detect or validate text search extension: 'native', 'vchord', 'pg_textsearch',
|
||||
'pgroonga', or 'pg_search'. Respects HINDSIGHT_API_TEXT_SEARCH_EXTENSION env var.
|
||||
Detect or validate text search extension: 'native', 'vchord', or 'pg_textsearch'.
|
||||
Respects HINDSIGHT_API_TEXT_SEARCH_EXTENSION env var.
|
||||
Creates the extension if needed.
|
||||
|
||||
pgroonga is treated as native here so the initial schema still creates valid
|
||||
tsvector columns. ensure_text_search_extension() at startup converts the
|
||||
schema to pgroonga structures (drops the tsvector column, builds a pgroonga
|
||||
index on the base text column).
|
||||
"""
|
||||
text_search_extension = os.getenv("HINDSIGHT_API_TEXT_SEARCH_EXTENSION", "native").lower()
|
||||
|
||||
@@ -131,35 +121,14 @@ def _detect_text_search_extension() -> str:
|
||||
# Extension truly doesn't exist - re-raise the error
|
||||
raise
|
||||
return "pg_textsearch"
|
||||
elif text_search_extension == "pg_search":
|
||||
# ParadeDB pg_search — true BM25 over base columns, Citus-compatible.
|
||||
try:
|
||||
op.execute("CREATE EXTENSION IF NOT EXISTS pg_search CASCADE")
|
||||
except Exception:
|
||||
# Extension might already exist or user lacks permissions - verify it exists
|
||||
conn = op.get_bind()
|
||||
result = conn.execute(text("SELECT 1 FROM pg_extension WHERE extname = 'pg_search'")).fetchone()
|
||||
if not result:
|
||||
# Extension truly doesn't exist - re-raise the error
|
||||
raise
|
||||
return "pg_search"
|
||||
elif text_search_extension == "native":
|
||||
return "native"
|
||||
elif text_search_extension == "pgroonga":
|
||||
# ensure_text_search_extension() at runtime converts to pgroonga.
|
||||
# Treat as native here so the initial schema still creates valid columns.
|
||||
return "native"
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Invalid HINDSIGHT_API_TEXT_SEARCH_EXTENSION: {text_search_extension}. "
|
||||
"Must be 'native', 'vchord', 'pg_textsearch', 'pgroonga', or 'pg_search'"
|
||||
f"Invalid HINDSIGHT_API_TEXT_SEARCH_EXTENSION: {text_search_extension}. Must be 'native', 'vchord', or 'pg_textsearch'"
|
||||
)
|
||||
|
||||
|
||||
def _pg_search_tokenizer() -> str:
|
||||
return normalize_pg_search_tokenizer(os.getenv(PG_SEARCH_TOKENIZER_ENV))
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
"""Upgrade schema - create all tables from scratch."""
|
||||
|
||||
@@ -315,9 +284,8 @@ def _pg_upgrade() -> None:
|
||||
ALTER TABLE memory_units
|
||||
ADD COLUMN search_vector bm25_catalog.bm25vector
|
||||
""")
|
||||
elif text_search_ext in ("pg_textsearch", "pg_search"):
|
||||
# Timescale pg_textsearch / ParadeDB pg_search: dummy TEXT column for
|
||||
# consistency (indexes operate on base columns directly).
|
||||
elif text_search_ext == "pg_textsearch":
|
||||
# Timescale pg_textsearch: dummy TEXT column for consistency (indexes operate on base columns directly)
|
||||
op.execute("""
|
||||
ALTER TABLE memory_units
|
||||
ADD COLUMN search_vector TEXT
|
||||
@@ -382,17 +350,6 @@ def _pg_upgrade() -> None:
|
||||
USING bm25(text)
|
||||
WITH (text_config='english')
|
||||
""")
|
||||
elif text_search_ext == "pg_search":
|
||||
# ParadeDB pg_search BM25 index on (id, text, context). The key_field
|
||||
# reloption is required and must match the table's primary key column.
|
||||
bm25_cols = pg_search_bm25_columns("id", ("text", "context"), _pg_search_tokenizer())
|
||||
op.execute(
|
||||
"""
|
||||
CREATE INDEX idx_memory_units_text_search ON memory_units
|
||||
USING bm25 ({bm25_cols})
|
||||
WITH (key_field='id')
|
||||
""".format(bm25_cols=bm25_cols)
|
||||
)
|
||||
else: # native
|
||||
# Native PostgreSQL GIN index
|
||||
op.execute("""
|
||||
|
||||
-29
@@ -7,7 +7,6 @@ the stored fact text.
|
||||
- vchord: text_signals included in tokenize() at insert time
|
||||
- native: search_vector GENERATED column regenerated to include text_signals
|
||||
- pg_textsearch: no change (index only supports a single base column)
|
||||
- pg_search: BM25 index dropped and recreated to include text_signals
|
||||
|
||||
Revision ID: a2b3c4d5e6f7
|
||||
Revises: z1u2v3w4x5y6
|
||||
@@ -19,11 +18,6 @@ from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api._pg_search import (
|
||||
PG_SEARCH_TOKENIZER_ENV,
|
||||
normalize_pg_search_tokenizer,
|
||||
pg_search_bm25_columns,
|
||||
)
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "a2b3c4d5e6f7"
|
||||
@@ -41,10 +35,6 @@ def _detect_text_search_extension() -> str:
|
||||
return os.getenv("HINDSIGHT_API_TEXT_SEARCH_EXTENSION", "native").lower()
|
||||
|
||||
|
||||
def _pg_search_tokenizer() -> str:
|
||||
return normalize_pg_search_tokenizer(os.getenv(PG_SEARCH_TOKENIZER_ENV))
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
table = f"{schema}memory_units"
|
||||
@@ -72,16 +62,6 @@ def _pg_upgrade() -> None:
|
||||
CREATE INDEX IF NOT EXISTS idx_memory_units_text_search
|
||||
ON {table} USING gin(search_vector)
|
||||
""")
|
||||
elif text_search_ext == "pg_search":
|
||||
# ParadeDB pg_search: drop the existing BM25 index and recreate it
|
||||
# to include text_signals alongside text and context.
|
||||
bm25_cols = pg_search_bm25_columns("id", ("text", "context", "text_signals"), _pg_search_tokenizer())
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema}idx_memory_units_text_search")
|
||||
op.execute(f"""
|
||||
CREATE INDEX idx_memory_units_text_search ON {table}
|
||||
USING bm25 ({bm25_cols})
|
||||
WITH (key_field='id')
|
||||
""")
|
||||
|
||||
# vchord: tokenize() call in fact_storage.py is updated to include text_signals at insert time
|
||||
# pg_textsearch: no change — index operates on the base `text` column only
|
||||
@@ -106,15 +86,6 @@ def _pg_downgrade() -> None:
|
||||
CREATE INDEX idx_memory_units_text_search
|
||||
ON {table} USING gin(search_vector)
|
||||
""")
|
||||
elif text_search_ext == "pg_search":
|
||||
# Restore the original (id, text, context) BM25 index without text_signals.
|
||||
bm25_cols = pg_search_bm25_columns("id", ("text", "context"), _pg_search_tokenizer())
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema}idx_memory_units_text_search")
|
||||
op.execute(f"""
|
||||
CREATE INDEX idx_memory_units_text_search ON {table}
|
||||
USING bm25 ({bm25_cols})
|
||||
WITH (key_field='id')
|
||||
""")
|
||||
|
||||
op.execute(f"ALTER TABLE {table} DROP COLUMN IF EXISTS text_signals")
|
||||
|
||||
|
||||
+10
-10
@@ -40,20 +40,20 @@ def _get_schema_prefix() -> str:
|
||||
def _pg_upgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
|
||||
# CREATE INDEX CONCURRENTLY cannot run inside a transaction block; an
|
||||
# autocommit_block runs it outside Alembic's migration transaction.
|
||||
with op.get_context().autocommit_block():
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_units_source_memory_ids "
|
||||
f"ON {schema}memory_units USING GIN (source_memory_ids) "
|
||||
f"WHERE source_memory_ids IS NOT NULL"
|
||||
)
|
||||
# CREATE INDEX CONCURRENTLY cannot run inside a transaction block.
|
||||
# Commit the current Alembic transaction first.
|
||||
op.execute("COMMIT")
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_units_source_memory_ids "
|
||||
f"ON {schema}memory_units USING GIN (source_memory_ids) "
|
||||
f"WHERE source_memory_ids IS NOT NULL"
|
||||
)
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
with op.get_context().autocommit_block():
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_units_source_memory_ids")
|
||||
op.execute("COMMIT")
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_units_source_memory_ids")
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
|
||||
+1
-1
@@ -63,7 +63,7 @@ def _vector_index_using_clause(ext: str) -> str:
|
||||
if ext == "pgvectorscale":
|
||||
return "USING diskann (embedding vector_cosine_ops) WITH (num_neighbors = 50)"
|
||||
if ext == "vchord":
|
||||
return "USING vchordrq (embedding vector_cosine_ops)"
|
||||
return "USING vchordrq (embedding vector_l2_ops)"
|
||||
if ext == "scann":
|
||||
return "USING scann (embedding cosine) WITH (mode = 'AUTO')"
|
||||
return "USING hnsw (embedding vector_cosine_ops)"
|
||||
|
||||
+27
-25
@@ -37,35 +37,37 @@ def _get_schema_prefix() -> str:
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
# CREATE INDEX CONCURRENTLY cannot run inside a transaction block; an
|
||||
# autocommit_block runs each statement outside Alembic's migration transaction.
|
||||
with op.get_context().autocommit_block():
|
||||
# Partial index on occurred_start (covers "occurred_start BETWEEN $4 AND $5")
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_units_bank_occurred_start "
|
||||
f"ON {schema}memory_units(bank_id, fact_type, occurred_start) "
|
||||
f"WHERE occurred_start IS NOT NULL"
|
||||
)
|
||||
# Partial index on occurred_end (covers "occurred_end BETWEEN $4 AND $5")
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_units_bank_occurred_end "
|
||||
f"ON {schema}memory_units(bank_id, fact_type, occurred_end) "
|
||||
f"WHERE occurred_end IS NOT NULL"
|
||||
)
|
||||
# Partial index on mentioned_at (covers "mentioned_at BETWEEN $4 AND $5")
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_units_bank_mentioned_at "
|
||||
f"ON {schema}memory_units(bank_id, fact_type, mentioned_at) "
|
||||
f"WHERE mentioned_at IS NOT NULL"
|
||||
)
|
||||
# Partial index on occurred_start (covers "occurred_start BETWEEN $4 AND $5")
|
||||
op.execute("COMMIT")
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_units_bank_occurred_start "
|
||||
f"ON {schema}memory_units(bank_id, fact_type, occurred_start) "
|
||||
f"WHERE occurred_start IS NOT NULL"
|
||||
)
|
||||
# Partial index on occurred_end (covers "occurred_end BETWEEN $4 AND $5")
|
||||
op.execute("COMMIT")
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_units_bank_occurred_end "
|
||||
f"ON {schema}memory_units(bank_id, fact_type, occurred_end) "
|
||||
f"WHERE occurred_end IS NOT NULL"
|
||||
)
|
||||
# Partial index on mentioned_at (covers "mentioned_at BETWEEN $4 AND $5")
|
||||
op.execute("COMMIT")
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_units_bank_mentioned_at "
|
||||
f"ON {schema}memory_units(bank_id, fact_type, mentioned_at) "
|
||||
f"WHERE mentioned_at IS NOT NULL"
|
||||
)
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
with op.get_context().autocommit_block():
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_units_bank_mentioned_at")
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_units_bank_occurred_end")
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_units_bank_occurred_start")
|
||||
op.execute("COMMIT")
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_units_bank_mentioned_at")
|
||||
op.execute("COMMIT")
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_units_bank_occurred_end")
|
||||
op.execute("COMMIT")
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_units_bank_occurred_start")
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
|
||||
-106
@@ -1,106 +0,0 @@
|
||||
"""Add graph_maintenance_queue table
|
||||
|
||||
Queue of memory_units whose outgoing temporal/semantic links lost a
|
||||
neighbour to a delete. Drained by the async graph_maintenance worker,
|
||||
which tops the unit's links back up using the same probes retain runs.
|
||||
|
||||
The queue only targets the link-recompute pass. The worker also runs
|
||||
bank-wide sweeps (orphan-entity prune, stale-cooccurrence prune) on each
|
||||
invocation; those don't need per-target queueing.
|
||||
|
||||
Revision ID: b5a4c3e2f1d8
|
||||
Revises: e9b2c7d1f3a4
|
||||
Create Date: 2026-05-27
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "b5a4c3e2f1d8"
|
||||
down_revision: str | Sequence[str] | None = "e9b2c7d1f3a4"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _pg_schema_prefix() -> str:
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
schema = _pg_schema_prefix()
|
||||
# Composite PK gives us natural ON CONFLICT DO NOTHING dedup when the same
|
||||
# unit is enqueued from overlapping deletes. No FK to memory_units: if the
|
||||
# unit is deleted between enqueue and drain, the worker observes it's gone
|
||||
# and skips — a cascade would erase the work order, but that work has
|
||||
# already been satisfied (no surviving row to maintain).
|
||||
op.execute(
|
||||
f"""
|
||||
CREATE TABLE IF NOT EXISTS {schema}graph_maintenance_queue (
|
||||
bank_id TEXT NOT NULL,
|
||||
unit_id UUID NOT NULL,
|
||||
enqueued_at TIMESTAMPTZ NOT NULL DEFAULT now(),
|
||||
PRIMARY KEY (bank_id, unit_id)
|
||||
)
|
||||
"""
|
||||
)
|
||||
op.execute(
|
||||
f"""
|
||||
CREATE INDEX IF NOT EXISTS idx_graph_maintenance_queue_bank_enqueued
|
||||
ON {schema}graph_maintenance_queue (bank_id, enqueued_at)
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
schema = _pg_schema_prefix()
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema}idx_graph_maintenance_queue_bank_enqueued")
|
||||
op.execute(f"DROP TABLE IF EXISTS {schema}graph_maintenance_queue")
|
||||
|
||||
|
||||
def _oracle_execute_ignoring_955(sql: str) -> None:
|
||||
"""Run a CREATE statement and swallow ORA-00955 (object already exists).
|
||||
|
||||
Mirrors the helper in the Oracle baseline migration so reruns stay safe
|
||||
on a database where the table was created by an earlier partial run.
|
||||
"""
|
||||
block = (
|
||||
"BEGIN "
|
||||
"EXECUTE IMMEDIATE :stmt; "
|
||||
"EXCEPTION WHEN OTHERS THEN "
|
||||
"IF SQLCODE = -955 THEN NULL; ELSE RAISE; END IF; "
|
||||
"END;"
|
||||
)
|
||||
op.get_bind().exec_driver_sql(block, {"stmt": sql.strip()})
|
||||
|
||||
|
||||
def _oracle_upgrade() -> None:
|
||||
_oracle_execute_ignoring_955(
|
||||
"""
|
||||
CREATE TABLE graph_maintenance_queue (
|
||||
bank_id VARCHAR2(256) NOT NULL,
|
||||
unit_id RAW(16) NOT NULL,
|
||||
enqueued_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,
|
||||
CONSTRAINT pk_graph_maintenance_queue PRIMARY KEY (bank_id, unit_id)
|
||||
)
|
||||
"""
|
||||
)
|
||||
_oracle_execute_ignoring_955(
|
||||
"CREATE INDEX idx_graph_maintenance_queue_bank_enqueued ON graph_maintenance_queue (bank_id, enqueued_at)"
|
||||
)
|
||||
|
||||
|
||||
def _oracle_downgrade() -> None:
|
||||
op.execute("DROP INDEX idx_graph_maintenance_queue_bank_enqueued")
|
||||
op.execute("DROP TABLE graph_maintenance_queue")
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade, oracle=_oracle_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade, oracle=_oracle_downgrade)
|
||||
-152
@@ -1,152 +0,0 @@
|
||||
"""Re-create vchord vector indexes with vector_cosine_ops
|
||||
|
||||
Revision ID: b8c9d0e1f2a3
|
||||
Revises: 86f7a033d372
|
||||
Create Date: 2026-05-20
|
||||
|
||||
vchordrq operator classes are bound 1:1 to operators in PostgreSQL:
|
||||
vector_l2_ops only matches ``<->``, while every Hindsight ANN query uses
|
||||
``<=>`` (cosine distance). The previous vchord mapping used vector_l2_ops,
|
||||
so vchord deployments could never use the index — every ANN query fell
|
||||
back to a sequential scan with per-row cosine computation.
|
||||
|
||||
This migration finds any vchordrq index built with vector_l2_ops in the
|
||||
target schema and re-creates it with vector_cosine_ops, using
|
||||
``CREATE INDEX CONCURRENTLY`` so it can run online. It is a no-op when:
|
||||
|
||||
* the configured vector extension is not vchord, or
|
||||
* no matching indexes exist (already on cosine ops).
|
||||
|
||||
Only PostgreSQL is affected; the Oracle 23ai dialect uses its own native
|
||||
vector index and does not depend on this mapping.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
from sqlalchemy import text
|
||||
|
||||
from hindsight_api._vector_index import configured_vector_extension
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "b8c9d0e1f2a3"
|
||||
down_revision: str | Sequence[str] | None = "86f7a033d372"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _pg_schema_prefix() -> str:
|
||||
"""Schema-qualifier for raw SQL on PG (multi-tenant search_path)."""
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def _rebuild_vchordrq_indexes(old_ops: str, new_ops: str) -> None:
|
||||
"""Rebuild vchordrq indexes using ``old_ops`` so they use ``new_ops``.
|
||||
|
||||
Each index is rebuilt with CREATE INDEX CONCURRENTLY under a fresh name,
|
||||
then the old index is dropped and the new one renamed to take its place.
|
||||
Must be called inside an ``autocommit_block()`` because CONCURRENTLY
|
||||
cannot run inside a transaction.
|
||||
"""
|
||||
bind = op.get_bind()
|
||||
# `or None` collapses both unset and explicit empty-string Alembic options
|
||||
# into NULL so the COALESCE below falls back to current_schema() in either
|
||||
# case. Without it, an empty-string option would filter on `schemaname = ''`
|
||||
# and skip every real schema.
|
||||
target_schema = context.config.get_main_option("target_schema") or None
|
||||
prefix = _pg_schema_prefix()
|
||||
|
||||
rows = bind.execute(
|
||||
text(
|
||||
"SELECT indexname, indexdef FROM pg_indexes "
|
||||
"WHERE schemaname = COALESCE(:target_schema, current_schema()) "
|
||||
"AND indexdef ILIKE '%vchordrq%' "
|
||||
"AND indexdef ILIKE :ops_like"
|
||||
),
|
||||
{"target_schema": target_schema, "ops_like": f"%{old_ops}%"},
|
||||
).fetchall()
|
||||
|
||||
for idx_name, indexdef in rows:
|
||||
# pg_get_indexdef() emits the canonical form `CREATE INDEX <name> ON …`,
|
||||
# so <name> is the first textual occurrence — both substitutions below
|
||||
# rely on that.
|
||||
new_def = indexdef.replace(old_ops, new_ops, 1)
|
||||
temp_name = f"{idx_name}__opclass_swap"
|
||||
new_def = new_def.replace(idx_name, temp_name, 1)
|
||||
new_def = re.sub(
|
||||
r"^CREATE\s+INDEX\b",
|
||||
"CREATE INDEX CONCURRENTLY IF NOT EXISTS",
|
||||
new_def,
|
||||
count=1,
|
||||
)
|
||||
|
||||
# CREATE INDEX CONCURRENTLY can leave the partial index as INVALID if a
|
||||
# previous run errored (disk pressure, lock conflict, signal). Without
|
||||
# this drop the CONCURRENTLY IF NOT EXISTS below would skip creation,
|
||||
# then we'd drop the original and rename the broken index into its
|
||||
# place — silently restoring the seq-scan bug this migration fixes.
|
||||
op.execute(f'DROP INDEX IF EXISTS {prefix}"{temp_name}"')
|
||||
op.execute(new_def)
|
||||
|
||||
# Even on a clean run CONCURRENTLY can finish with indisvalid = false
|
||||
# (e.g. constraint violation during the second build scan). Refuse to
|
||||
# promote in that case so we never alias an INVALID index over a working
|
||||
# one.
|
||||
is_valid = bind.execute(
|
||||
text(
|
||||
"SELECT i.indisvalid "
|
||||
"FROM pg_class c "
|
||||
"JOIN pg_index i ON c.oid = i.indexrelid "
|
||||
"JOIN pg_namespace n ON c.relnamespace = n.oid "
|
||||
"WHERE c.relname = :name "
|
||||
" AND n.nspname = COALESCE(:target_schema, current_schema())"
|
||||
),
|
||||
{"name": temp_name, "target_schema": target_schema},
|
||||
).scalar()
|
||||
if not is_valid:
|
||||
raise RuntimeError(
|
||||
f"vchordrq index rebuild produced an INVALID index ({temp_name}); "
|
||||
"drop it manually and re-run the migration."
|
||||
)
|
||||
|
||||
# DROP + RENAME atomically. A crash between the two would leave
|
||||
# `temp_name` as a valid orphan and the canonical name missing —
|
||||
# next run's `pg_indexes` filter (looking for vector_l2_ops) wouldn't
|
||||
# find anything to recover from, so the index would stay gone. PG
|
||||
# runs the DO block in its own server-side transaction, so either
|
||||
# both succeed or both roll back.
|
||||
op.execute(
|
||||
f"""
|
||||
DO $$
|
||||
BEGIN
|
||||
DROP INDEX IF EXISTS {prefix}"{idx_name}";
|
||||
ALTER INDEX {prefix}"{temp_name}" RENAME TO "{idx_name}";
|
||||
END $$;
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
if configured_vector_extension() != "vchord":
|
||||
return
|
||||
with op.get_context().autocommit_block():
|
||||
_rebuild_vchordrq_indexes("vector_l2_ops", "vector_cosine_ops")
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
if configured_vector_extension() != "vchord":
|
||||
return
|
||||
with op.get_context().autocommit_block():
|
||||
_rebuild_vchordrq_indexes("vector_cosine_ops", "vector_l2_ops")
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade)
|
||||
+7
-8
@@ -47,18 +47,17 @@ def _pg_upgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
# GIN index on canonical_name enables sub-millisecond trigram similarity queries
|
||||
# (% operator, similarity()) instead of full-table scans across all bank entities.
|
||||
# CREATE INDEX CONCURRENTLY cannot run inside a transaction block.
|
||||
with op.get_context().autocommit_block():
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS entities_canonical_name_trgm_idx "
|
||||
f"ON {schema}entities USING GIN (canonical_name gin_trgm_ops)"
|
||||
)
|
||||
op.execute("COMMIT")
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS entities_canonical_name_trgm_idx "
|
||||
f"ON {schema}entities USING GIN (canonical_name gin_trgm_ops)"
|
||||
)
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
with op.get_context().autocommit_block():
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}entities_canonical_name_trgm_idx")
|
||||
op.execute("COMMIT")
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}entities_canonical_name_trgm_idx")
|
||||
# Note: not dropping pg_trgm extension as other indexes may depend on it
|
||||
|
||||
|
||||
|
||||
-45
@@ -1,45 +0,0 @@
|
||||
"""Merge graph_maintenance_queue and vchord_cosine_opclass heads.
|
||||
|
||||
Revision ID: c1d2e3f4a5b6
|
||||
Revises: b5a4c3e2f1d8, b8c9d0e1f2a3
|
||||
Create Date: 2026-05-29
|
||||
|
||||
PRs #1668 (vchord cosine opclass) and #1772 (async link recompute) both
|
||||
branched off the same parent and were merged onto main without rebasing,
|
||||
leaving two parallel Alembic heads. This is a structural merge revision
|
||||
with no schema changes — its only job is to unify the DAG so
|
||||
``alembic upgrade head`` is unambiguous again.
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "c1d2e3f4a5b6"
|
||||
down_revision: str | Sequence[str] | None = ("b5a4c3e2f1d8", "b8c9d0e1f2a3")
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
pass
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
pass
|
||||
|
||||
|
||||
def _oracle_upgrade() -> None:
|
||||
pass
|
||||
|
||||
|
||||
def _oracle_downgrade() -> None:
|
||||
pass
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade, oracle=_oracle_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade, oracle=_oracle_downgrade)
|
||||
+28
-24
@@ -50,35 +50,39 @@ def _get_schema_prefix() -> str:
|
||||
def _pg_upgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
|
||||
# CREATE INDEX CONCURRENTLY cannot run inside a transaction block; an
|
||||
# autocommit_block runs each statement outside Alembic's migration
|
||||
# transaction. IF NOT EXISTS makes each statement idempotent on retry.
|
||||
with op.get_context().autocommit_block():
|
||||
# Index for the semantic *incoming* direction in link_expansion_retrieval.py.
|
||||
# Replaces the BitmapAnd of idx_memory_links_to_unit ∩ idx_memory_links_link_type
|
||||
# with a single composite index scan.
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_links_to_type_weight "
|
||||
f"ON {schema}memory_links(to_unit_id, link_type, weight DESC)"
|
||||
)
|
||||
# CREATE INDEX CONCURRENTLY cannot run inside a transaction block.
|
||||
# Commit the current Alembic transaction, then issue each CONCURRENTLY
|
||||
# statement in its own implicit autocommit transaction.
|
||||
# IF NOT EXISTS makes each statement idempotent if the migration is retried.
|
||||
|
||||
# Covering index for entity co-occurrence expansion.
|
||||
# Enables an index-only scan: entity_id and to_unit_id are read from the
|
||||
# index leaf pages instead of the heap, eliminating ~2 500 random heap-page
|
||||
# reads per expansion query.
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_links_entity_covering "
|
||||
f"ON {schema}memory_links(from_unit_id) "
|
||||
f"INCLUDE (to_unit_id, entity_id) "
|
||||
f"WHERE link_type = 'entity'"
|
||||
)
|
||||
# Index for the semantic *incoming* direction in link_expansion_retrieval.py.
|
||||
# Replaces the BitmapAnd of idx_memory_links_to_unit ∩ idx_memory_links_link_type
|
||||
# with a single composite index scan.
|
||||
op.execute("COMMIT")
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_links_to_type_weight "
|
||||
f"ON {schema}memory_links(to_unit_id, link_type, weight DESC)"
|
||||
)
|
||||
|
||||
# Covering index for entity co-occurrence expansion.
|
||||
# Enables an index-only scan: entity_id and to_unit_id are read from the
|
||||
# index leaf pages instead of the heap, eliminating ~2 500 random heap-page
|
||||
# reads per expansion query.
|
||||
op.execute("COMMIT")
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_links_entity_covering "
|
||||
f"ON {schema}memory_links(from_unit_id) "
|
||||
f"INCLUDE (to_unit_id, entity_id) "
|
||||
f"WHERE link_type = 'entity'"
|
||||
)
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
with op.get_context().autocommit_block():
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_links_entity_covering")
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_links_to_type_weight")
|
||||
op.execute("COMMIT")
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_links_entity_covering")
|
||||
op.execute("COMMIT")
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_links_to_type_weight")
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
|
||||
+16
-17
@@ -33,27 +33,26 @@ def _get_schema_prefix() -> str:
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
# DROP + CREATE CONCURRENTLY must run outside a transaction block; an
|
||||
# autocommit_block runs them outside Alembic's migration transaction.
|
||||
with op.get_context().autocommit_block():
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_units_source_memory_ids")
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_units_source_memory_ids "
|
||||
f"ON {schema}memory_units USING GIN (source_memory_ids) "
|
||||
f"WITH (fastupdate=off) "
|
||||
f"WHERE source_memory_ids IS NOT NULL"
|
||||
)
|
||||
# DROP + CREATE CONCURRENTLY must run outside a transaction block.
|
||||
op.execute("COMMIT")
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_units_source_memory_ids")
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_units_source_memory_ids "
|
||||
f"ON {schema}memory_units USING GIN (source_memory_ids) "
|
||||
f"WITH (fastupdate=off) "
|
||||
f"WHERE source_memory_ids IS NOT NULL"
|
||||
)
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
with op.get_context().autocommit_block():
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_units_source_memory_ids")
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_units_source_memory_ids "
|
||||
f"ON {schema}memory_units USING GIN (source_memory_ids) "
|
||||
f"WHERE source_memory_ids IS NOT NULL"
|
||||
)
|
||||
op.execute("COMMIT")
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_units_source_memory_ids")
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_units_source_memory_ids "
|
||||
f"ON {schema}memory_units USING GIN (source_memory_ids) "
|
||||
f"WHERE source_memory_ids IS NOT NULL"
|
||||
)
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
|
||||
+1
-1
@@ -55,7 +55,7 @@ def _vector_index_using_clause(ext: str) -> str:
|
||||
if ext == "pgvectorscale":
|
||||
return "USING diskann (embedding vector_cosine_ops) WITH (num_neighbors = 50)"
|
||||
if ext == "vchord":
|
||||
return "USING vchordrq (embedding vector_cosine_ops)"
|
||||
return "USING vchordrq (embedding vector_l2_ops)"
|
||||
if ext == "scann":
|
||||
return "USING scann (embedding cosine) WITH (mode = 'AUTO')"
|
||||
return "USING hnsw (embedding vector_cosine_ops)"
|
||||
|
||||
-133
@@ -1,133 +0,0 @@
|
||||
"""Drop indexes that are unused or redundant with composite indexes.
|
||||
|
||||
Code audit identified the following indexes as either dead (no code path
|
||||
exercises them) or fully covered by composite indexes the planner already
|
||||
prefers:
|
||||
|
||||
memory_links:
|
||||
1. idx_memory_links_entity_covering — entity co-occurrence expansion was
|
||||
rewritten to traverse unit_entities instead of memory_links, so no code
|
||||
path filters memory_links on (link_type = 'entity').
|
||||
2. idx_memory_links_from_unit — redundant. idx_memory_links_from_type_weight
|
||||
(from_unit_id, link_type, weight DESC) leads with the same column and
|
||||
answers every from_unit_id = X query.
|
||||
3. idx_memory_links_to_unit — redundant. idx_memory_links_to_type_weight
|
||||
(to_unit_id, link_type, weight DESC) leads with the same column.
|
||||
4. idx_memory_links_link_type — no application query filters on link_type
|
||||
alone; the composite indexes above serve every (from/to + link_type)
|
||||
predicate.
|
||||
|
||||
entities:
|
||||
5. idx_entities_canonical_name — superseded by
|
||||
entities_canonical_name_lower_trgm_idx (case-insensitive lookups).
|
||||
6. entities_canonical_name_trgm_idx — superseded by the lowercase variant
|
||||
in migration 2eee35aa3cfc, but the original was never dropped on schemas
|
||||
that ran the prior migration.
|
||||
|
||||
documents:
|
||||
7. idx_documents_retain_params — GIN index on retain_params JSONB; no query
|
||||
uses jsonb containment on this column.
|
||||
8. idx_documents_content_hash — content-hash lookups happen on the chunks
|
||||
table (chunks.content_hash, indexed separately).
|
||||
|
||||
unit_entities:
|
||||
9. idx_unit_entities_entity — defensive drop. Migration h3i4j5k6l7m8 already
|
||||
issues DROP INDEX IF EXISTS for this; this re-runs the drop idempotently
|
||||
to cover any schema that missed the previous migration.
|
||||
|
||||
All drops use CONCURRENTLY + IF EXISTS so they neither block writers nor
|
||||
fail on schemas where the index is already gone.
|
||||
|
||||
Revision ID: e1b2c3d4f5a6
|
||||
Revises: p4q5r6s7t8u9
|
||||
Create Date: 2026-05-26
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "e1b2c3d4f5a6"
|
||||
down_revision: str | Sequence[str] | None = "p4q5r6s7t8u9"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
_PG_INDEXES_TO_DROP: tuple[str, ...] = (
|
||||
"idx_memory_links_entity_covering",
|
||||
"idx_memory_links_from_unit",
|
||||
"idx_memory_links_to_unit",
|
||||
"idx_memory_links_link_type",
|
||||
"idx_entities_canonical_name",
|
||||
"entities_canonical_name_trgm_idx",
|
||||
"idx_documents_retain_params",
|
||||
"idx_documents_content_hash",
|
||||
"idx_unit_entities_entity",
|
||||
)
|
||||
|
||||
|
||||
def _schema_prefix() -> str:
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
schema = _schema_prefix()
|
||||
# DROP INDEX CONCURRENTLY cannot run inside a transaction block; an
|
||||
# autocommit_block drops out of Alembic's migration transaction so each
|
||||
# statement runs in its own autocommit. IF EXISTS makes each statement
|
||||
# idempotent across schemas that already dropped (or never had) the index.
|
||||
with op.get_context().autocommit_block():
|
||||
for index_name in _PG_INDEXES_TO_DROP:
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}{index_name}")
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
schema = _schema_prefix()
|
||||
|
||||
# Recreate the dropped indexes in the same shape the prior migrations used,
|
||||
# so a downgrade leaves the schema in the state the previous head expected.
|
||||
# CREATE INDEX CONCURRENTLY cannot run inside a transaction block.
|
||||
with op.get_context().autocommit_block():
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_links_entity_covering "
|
||||
f"ON {schema}memory_links(from_unit_id) "
|
||||
f"INCLUDE (to_unit_id, entity_id) "
|
||||
f"WHERE link_type = 'entity'"
|
||||
)
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_links_from_unit ON {schema}memory_links(from_unit_id)"
|
||||
)
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_links_to_unit ON {schema}memory_links(to_unit_id)"
|
||||
)
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_links_link_type ON {schema}memory_links(link_type)"
|
||||
)
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_entities_canonical_name ON {schema}entities(canonical_name)"
|
||||
)
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS entities_canonical_name_trgm_idx "
|
||||
f"ON {schema}entities USING GIN (canonical_name gin_trgm_ops)"
|
||||
)
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_documents_retain_params "
|
||||
f"ON {schema}documents USING GIN (retain_params)"
|
||||
)
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_documents_content_hash ON {schema}documents(content_hash)"
|
||||
)
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_unit_entities_entity ON {schema}unit_entities(entity_id)"
|
||||
)
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade)
|
||||
-91
@@ -1,91 +0,0 @@
|
||||
"""Drop materialized entity rows from memory_links.
|
||||
|
||||
Entity edges are no longer stored in ``memory_links``. The /graph endpoint
|
||||
derives them on demand from ``unit_entities``, and recall already used the
|
||||
``unit_entities`` self-join. Storing entity rows duplicated state we never
|
||||
read from the link table — on a 10k-unit benchmark bank, entity rows were
|
||||
53% of all link rows (~190 MB after indexes) and recall never touched them.
|
||||
|
||||
This migration deletes ``memory_links`` rows with ``link_type = 'entity'``.
|
||||
``idx_memory_links_entity_covering`` was already dropped by migration
|
||||
``e1b2c3d4f5a6``; we still issue ``DROP INDEX IF EXISTS`` defensively in case
|
||||
this migration runs against an older snapshot that predates that one.
|
||||
|
||||
Revision ID: e9b2c7d1f3a4
|
||||
Revises: e1b2c3d4f5a6
|
||||
Create Date: 2026-05-26
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "e9b2c7d1f3a4"
|
||||
down_revision: str | Sequence[str] | None = "e1b2c3d4f5a6"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _pg_schema_prefix() -> str:
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
schema = _pg_schema_prefix()
|
||||
|
||||
# Drop the partial covering index first so the bulk DELETE doesn't churn it.
|
||||
# DROP INDEX CONCURRENTLY, and the DO block's per-batch COMMIT, both require
|
||||
# running outside Alembic's migration transaction — an autocommit_block
|
||||
# commits it and switches the connection to autocommit for the duration.
|
||||
with op.get_context().autocommit_block():
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_links_entity_covering")
|
||||
|
||||
# Delete entity rows. Chunked to keep individual transactions small on
|
||||
# large banks (the perf-medium bench had ~345k entity rows; production
|
||||
# banks can be much larger).
|
||||
op.execute(
|
||||
f"""
|
||||
DO $$
|
||||
DECLARE
|
||||
deleted INTEGER;
|
||||
BEGIN
|
||||
LOOP
|
||||
DELETE FROM {schema}memory_links
|
||||
WHERE ctid IN (
|
||||
SELECT ctid FROM {schema}memory_links
|
||||
WHERE link_type = 'entity'
|
||||
LIMIT 50000
|
||||
);
|
||||
GET DIAGNOSTICS deleted = ROW_COUNT;
|
||||
EXIT WHEN deleted = 0;
|
||||
COMMIT;
|
||||
END LOOP;
|
||||
END$$;
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
# Cannot reconstruct deleted entity links — the writer was path-dependent
|
||||
# on retain order. New retains will not produce entity rows either, so the
|
||||
# partial index would stay empty. Leave both no-op.
|
||||
pass
|
||||
|
||||
|
||||
def _oracle_upgrade() -> None:
|
||||
op.execute("DELETE FROM memory_links WHERE link_type = 'entity'")
|
||||
|
||||
|
||||
def _oracle_downgrade() -> None:
|
||||
pass
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade, oracle=_oracle_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade, oracle=_oracle_downgrade)
|
||||
+4
-58
@@ -16,11 +16,6 @@ from collections.abc import Sequence
|
||||
from alembic import context, op
|
||||
from sqlalchemy import text
|
||||
|
||||
from hindsight_api._pg_search import (
|
||||
PG_SEARCH_TOKENIZER_ENV,
|
||||
normalize_pg_search_tokenizer,
|
||||
pg_search_bm25_columns,
|
||||
)
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
@@ -92,7 +87,7 @@ def _vector_index_using_clause(ext: str) -> str:
|
||||
if ext == "pg_diskann":
|
||||
return "USING diskann (embedding vector_cosine_ops) WITH (max_neighbors = 50)"
|
||||
if ext == "vchord":
|
||||
return "USING vchordrq (embedding vector_cosine_ops)"
|
||||
return "USING vchordrq (embedding vector_l2_ops)"
|
||||
if ext == "scann":
|
||||
return "USING scann (embedding cosine) WITH (mode = 'AUTO')"
|
||||
return "USING hnsw (embedding vector_cosine_ops)"
|
||||
@@ -100,15 +95,9 @@ def _vector_index_using_clause(ext: str) -> str:
|
||||
|
||||
def _detect_text_search_extension() -> str:
|
||||
"""
|
||||
Detect or validate text search extension: 'native', 'vchord', 'pg_textsearch',
|
||||
'pgroonga', or 'pg_search'. Respects HINDSIGHT_API_TEXT_SEARCH_EXTENSION env var.
|
||||
Detect or validate text search extension: 'native', 'vchord', or 'pg_textsearch'.
|
||||
Respects HINDSIGHT_API_TEXT_SEARCH_EXTENSION env var.
|
||||
Creates the extension if needed.
|
||||
|
||||
pgroonga is treated as native here so this migration still creates valid
|
||||
tsvector columns; ensure_text_search_extension() at startup converts the
|
||||
reflections table (renamed from pinned_reflections in p1k2l3m4n5o6) to
|
||||
pgroonga structures. The learnings table is dropped in p1k2l3m4n5o6 so its
|
||||
transient native-style column never reaches steady state.
|
||||
"""
|
||||
text_search_extension = os.getenv("HINDSIGHT_API_TEXT_SEARCH_EXTENSION", "native").lower()
|
||||
|
||||
@@ -136,33 +125,14 @@ def _detect_text_search_extension() -> str:
|
||||
# Extension truly doesn't exist - re-raise the error
|
||||
raise
|
||||
return "pg_textsearch"
|
||||
elif text_search_extension == "pg_search":
|
||||
# ParadeDB pg_search — true BM25 over base columns, Citus-compatible.
|
||||
try:
|
||||
op.execute("CREATE EXTENSION IF NOT EXISTS pg_search CASCADE")
|
||||
except Exception:
|
||||
conn = op.get_bind()
|
||||
result = conn.execute(text("SELECT 1 FROM pg_extension WHERE extname = 'pg_search'")).fetchone()
|
||||
if not result:
|
||||
raise
|
||||
return "pg_search"
|
||||
elif text_search_extension == "native":
|
||||
return "native"
|
||||
elif text_search_extension == "pgroonga":
|
||||
# Treat as native here; ensure_text_search_extension() converts the
|
||||
# reflections table to pgroonga structures at runtime.
|
||||
return "native"
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Invalid HINDSIGHT_API_TEXT_SEARCH_EXTENSION: {text_search_extension}. "
|
||||
"Must be 'native', 'vchord', 'pg_textsearch', 'pgroonga', or 'pg_search'"
|
||||
f"Invalid HINDSIGHT_API_TEXT_SEARCH_EXTENSION: {text_search_extension}. Must be 'native', 'vchord', or 'pg_textsearch'"
|
||||
)
|
||||
|
||||
|
||||
def _pg_search_tokenizer() -> str:
|
||||
return normalize_pg_search_tokenizer(os.getenv(PG_SEARCH_TOKENIZER_ENV))
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
"""Create learnings and pinned_reflections tables."""
|
||||
schema = _get_schema_prefix()
|
||||
@@ -230,18 +200,6 @@ def _pg_upgrade() -> None:
|
||||
CREATE INDEX idx_learnings_text_search ON {schema}learnings
|
||||
USING bm25(text) WITH (text_config='english')
|
||||
""")
|
||||
elif text_search_ext == "pg_search":
|
||||
# ParadeDB pg_search: dummy TEXT column; BM25 index is built directly over (id, text)
|
||||
# with key_field='id' (matches the table's primary key).
|
||||
bm25_cols = pg_search_bm25_columns("id", ("text",), _pg_search_tokenizer())
|
||||
op.execute(f"""
|
||||
ALTER TABLE {schema}learnings ADD COLUMN search_vector TEXT
|
||||
""")
|
||||
op.execute(f"""
|
||||
CREATE INDEX idx_learnings_text_search ON {schema}learnings
|
||||
USING bm25 ({bm25_cols})
|
||||
WITH (key_field='id')
|
||||
""")
|
||||
else: # native
|
||||
# Native PostgreSQL: tsvector with automatic generation
|
||||
op.execute(f"""
|
||||
@@ -306,18 +264,6 @@ def _pg_upgrade() -> None:
|
||||
USING bm25(content)
|
||||
WITH (text_config='english')
|
||||
""")
|
||||
elif text_search_ext == "pg_search":
|
||||
# ParadeDB pg_search: dummy TEXT column; BM25 index over (id, name, content)
|
||||
# with key_field='id'.
|
||||
bm25_cols = pg_search_bm25_columns("id", ("name", "content"), _pg_search_tokenizer())
|
||||
op.execute(f"""
|
||||
ALTER TABLE {schema}pinned_reflections ADD COLUMN search_vector TEXT
|
||||
""")
|
||||
op.execute(f"""
|
||||
CREATE INDEX idx_pinned_reflections_text_search ON {schema}pinned_reflections
|
||||
USING bm25 ({bm25_cols})
|
||||
WITH (key_field='id')
|
||||
""")
|
||||
else: # native
|
||||
# Native PostgreSQL: tsvector with automatic generation
|
||||
op.execute(f"""
|
||||
|
||||
-170
@@ -1,170 +0,0 @@
|
||||
"""Drop GENERATED expression on tsvector search_vector columns.
|
||||
|
||||
The search_vector tsvector column was originally GENERATED ALWAYS with a
|
||||
hardcoded ``to_tsvector('english', ...)`` expression. To support configurable
|
||||
``HINDSIGHT_API_TEXT_SEARCH_EXTENSION_NATIVE_LANGUAGE``, we convert it to a
|
||||
regular tsvector column that the application populates at INSERT time via
|
||||
``to_tsvector($lang, ...)``.
|
||||
|
||||
Existing rows retain their English-derived lexemes — switching the configured
|
||||
language only affects newly-written rows. Users who need to backfill existing
|
||||
rows in a different language can run an admin UPDATE after this migration.
|
||||
|
||||
Only the ``native`` text-search backend is affected. ``vchord``, ``pg_textsearch``,
|
||||
and ``pgroonga`` use other column types or no column at all.
|
||||
|
||||
Revision ID: p4q5r6s7t8u9
|
||||
Revises: 86f7a033d372
|
||||
Create Date: 2026-05-08
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
from dataclasses import dataclass
|
||||
|
||||
from alembic import context, op
|
||||
from sqlalchemy import Connection, text
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "p4q5r6s7t8u9"
|
||||
down_revision: str | Sequence[str] | None = "86f7a033d372"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _TsvectorTableSpec:
|
||||
"""Native-backend tsvector table targeted by this migration.
|
||||
|
||||
``upgrade`` is a one-way DROP EXPRESSION; ``downgrade`` re-attaches the
|
||||
original GENERATED expression so the schema returns to the state created
|
||||
by the initial migration (and a2b3c4d5e6f7_add_text_signals_column for
|
||||
memory_units).
|
||||
"""
|
||||
|
||||
table: str
|
||||
generated_expression: str
|
||||
|
||||
|
||||
# Tables that may have a GENERATED tsvector ``search_vector`` column under the
|
||||
# native backend. Note: the ``learnings`` table was dropped in
|
||||
# p1k2l3m4n5o6_new_knowledge_architecture and ``pinned_reflections`` was renamed
|
||||
# to ``reflections`` in the same migration.
|
||||
_NATIVE_TSVECTOR_TABLES: tuple[_TsvectorTableSpec, ...] = (
|
||||
_TsvectorTableSpec(
|
||||
table="memory_units",
|
||||
generated_expression=(
|
||||
"to_tsvector('english', COALESCE(text, '') || ' ' || "
|
||||
"COALESCE(context, '') || ' ' || COALESCE(text_signals, ''))"
|
||||
),
|
||||
),
|
||||
_TsvectorTableSpec(
|
||||
table="reflections",
|
||||
generated_expression="to_tsvector('english', COALESCE(name, '') || ' ' || content)",
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _schema_prefix() -> str:
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def _is_generated_tsvector(conn: Connection, schema: str, table: str) -> bool:
|
||||
"""Return True iff ``schema.table.search_vector`` is a GENERATED tsvector column."""
|
||||
row = conn.execute(
|
||||
text(
|
||||
"""
|
||||
SELECT is_generated, udt_name
|
||||
FROM information_schema.columns
|
||||
WHERE table_schema = :schema
|
||||
AND table_name = :table
|
||||
AND column_name = 'search_vector'
|
||||
"""
|
||||
),
|
||||
{"schema": schema, "table": table},
|
||||
).fetchone()
|
||||
if not row:
|
||||
return False
|
||||
is_generated, udt_name = row[0], row[1]
|
||||
return is_generated == "ALWAYS" and udt_name == "tsvector"
|
||||
|
||||
|
||||
def _is_regular_tsvector(conn: Connection, schema: str, table: str) -> bool:
|
||||
"""Return True iff ``schema.table.search_vector`` is a non-generated tsvector column."""
|
||||
row = conn.execute(
|
||||
text(
|
||||
"""
|
||||
SELECT is_generated, udt_name
|
||||
FROM information_schema.columns
|
||||
WHERE table_schema = :schema
|
||||
AND table_name = :table
|
||||
AND column_name = 'search_vector'
|
||||
"""
|
||||
),
|
||||
{"schema": schema, "table": table},
|
||||
).fetchone()
|
||||
if not row:
|
||||
return False
|
||||
is_generated, udt_name = row[0], row[1]
|
||||
return udt_name == "tsvector" and is_generated != "ALWAYS"
|
||||
|
||||
|
||||
def _table_exists(conn: Connection, schema: str, table: str) -> bool:
|
||||
return bool(
|
||||
conn.execute(
|
||||
text(
|
||||
"""
|
||||
SELECT 1 FROM information_schema.tables
|
||||
WHERE table_schema = :schema AND table_name = :table
|
||||
"""
|
||||
),
|
||||
{"schema": schema, "table": table},
|
||||
).fetchone()
|
||||
)
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
schema_prefix = _schema_prefix()
|
||||
schema_name = (context.config.get_main_option("target_schema") or "public").strip('"')
|
||||
conn = op.get_bind()
|
||||
|
||||
for spec in _NATIVE_TSVECTOR_TABLES:
|
||||
if not _table_exists(conn, schema_name, spec.table):
|
||||
continue
|
||||
if not _is_generated_tsvector(conn, schema_name, spec.table):
|
||||
# Either the column doesn't exist (non-native backend) or it's
|
||||
# already a regular tsvector — nothing to do.
|
||||
continue
|
||||
op.execute(f"ALTER TABLE {schema_prefix}{spec.table} ALTER COLUMN search_vector DROP EXPRESSION")
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
schema_prefix = _schema_prefix()
|
||||
schema_name = (context.config.get_main_option("target_schema") or "public").strip('"')
|
||||
conn = op.get_bind()
|
||||
|
||||
for spec in _NATIVE_TSVECTOR_TABLES:
|
||||
if not _table_exists(conn, schema_name, spec.table):
|
||||
continue
|
||||
# Only restore the GENERATED expression if a non-generated tsvector
|
||||
# column exists — otherwise the table is on a different backend.
|
||||
if not _is_regular_tsvector(conn, schema_name, spec.table):
|
||||
continue
|
||||
# Drop and recreate to re-attach the GENERATED expression. Index will be
|
||||
# recreated by re-running ensure_text_search_extension on next startup.
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema_prefix}idx_{spec.table}_text_search")
|
||||
op.execute(f"ALTER TABLE {schema_prefix}{spec.table} DROP COLUMN search_vector")
|
||||
op.execute(
|
||||
f"ALTER TABLE {schema_prefix}{spec.table} "
|
||||
f"ADD COLUMN search_vector tsvector GENERATED ALWAYS AS ({spec.generated_expression}) STORED"
|
||||
)
|
||||
op.execute(f"CREATE INDEX idx_{spec.table}_text_search ON {schema_prefix}{spec.table} USING gin(search_vector)")
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade)
|
||||
@@ -15,7 +15,6 @@ from datetime import datetime, timezone
|
||||
from typing import Any, Literal
|
||||
|
||||
from fastapi import Depends, FastAPI, File, Form, Header, HTTPException, Query, Request, UploadFile
|
||||
from fastapi.middleware.gzip import GZipMiddleware
|
||||
|
||||
from hindsight_api.engine.audit import AuditEntry, AuditLogger
|
||||
from hindsight_api.extensions import AuthenticationError
|
||||
@@ -152,11 +151,7 @@ class RecallRequest(BaseModel):
|
||||
max_tokens: int = 4096
|
||||
trace: bool = False
|
||||
query_timestamp: str | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"ISO format date string (e.g., '2023-05-30T23:40:00'). Used as the query-time anchor for "
|
||||
"relative temporal expressions and recency scoring."
|
||||
),
|
||||
default=None, description="ISO format date string (e.g., '2023-05-30T23:40:00')"
|
||||
)
|
||||
include: IncludeOptions = FieldWithDefault(
|
||||
IncludeOptions,
|
||||
@@ -471,13 +466,6 @@ class MemoryItem(BaseModel):
|
||||
description="Optional tags for visibility scoping. Memories with tags can be filtered during recall.",
|
||||
)
|
||||
|
||||
@field_validator("content")
|
||||
@classmethod
|
||||
def validate_content(cls, v: str) -> str:
|
||||
if not v.strip():
|
||||
raise ValueError("content cannot be empty")
|
||||
return v
|
||||
|
||||
@field_validator("tags", mode="before")
|
||||
@classmethod
|
||||
def coerce_tags(cls, v):
|
||||
@@ -2192,19 +2180,6 @@ class OperationResponse(BaseModel):
|
||||
)
|
||||
|
||||
|
||||
class ConsolidationRequest(BaseModel):
|
||||
"""Request model for consolidation trigger endpoint."""
|
||||
|
||||
observation_scopes: list[list[str]] | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Optional list of tag scopes to consolidate. Each scope is a list of tags. "
|
||||
"Only unconsolidated memories whose tags contain all tags in at least one scope "
|
||||
"will be processed. If omitted, all unconsolidated memories are processed."
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class ConsolidationResponse(BaseModel):
|
||||
"""Response model for consolidation trigger endpoint."""
|
||||
|
||||
@@ -2674,7 +2649,6 @@ def create_app(
|
||||
tenant_extension=memory._tenant_extension,
|
||||
max_slots=config.worker_max_slots,
|
||||
slot_reservations=config.worker_slot_reservations,
|
||||
consolidation_bank_priority=config.worker_consolidation_bank_priority or None,
|
||||
)
|
||||
poller_task = asyncio.create_task(poller.run())
|
||||
logging.info(f"Worker poller started (worker_id={worker_id})")
|
||||
@@ -2747,8 +2721,6 @@ def create_app(
|
||||
app.state.memory = memory
|
||||
app.state.audit_logger = memory.audit_logger
|
||||
|
||||
app.add_middleware(GZipMiddleware, minimum_size=1024)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Patch OpenAPI schema: align ValidationError with Pydantic v2 error format
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -2913,57 +2885,6 @@ def _register_routes(app: FastAPI):
|
||||
api_key = authorization.strip()
|
||||
return RequestContext(api_key=api_key)
|
||||
|
||||
def precheck_for(operation: str):
|
||||
"""
|
||||
Build a FastAPI dependency that runs ``OperationValidator.precheck``.
|
||||
|
||||
FastAPI resolves dependencies before deserialising the route's body
|
||||
parameter. Wiring this dependency on the billable POST routes lets
|
||||
an extension reject a request — e.g. with HTTP 402 when a tenant's
|
||||
balance is exhausted — without the request body ever being read or
|
||||
materialised in memory.
|
||||
|
||||
The dependency intentionally:
|
||||
- authenticates the tenant (so ``request_context.tenant_id`` is
|
||||
resolved before the precheck runs);
|
||||
- falls through silently when no validator is configured or the
|
||||
validator's default no-op precheck is in effect;
|
||||
- converts a rejection ``ValidationResult`` into the corresponding
|
||||
``HTTPException`` directly (the per-route ``OperationValidationError``
|
||||
catch blocks don't see exceptions raised in dependencies, so we
|
||||
translate here instead of relying on each handler's try/except).
|
||||
|
||||
Args:
|
||||
operation: Short identifier for the route, e.g. ``"retain"``.
|
||||
|
||||
Returns:
|
||||
A FastAPI dependency callable suitable for ``Depends(...)``.
|
||||
"""
|
||||
|
||||
async def _precheck_dep(
|
||||
bank_id: str,
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
) -> None:
|
||||
validator = getattr(app.state.memory, "_operation_validator", None)
|
||||
if validator is None:
|
||||
return
|
||||
from hindsight_api.extensions import PrecheckContext
|
||||
|
||||
await app.state.memory._authenticate_tenant(request_context)
|
||||
ctx = PrecheckContext(
|
||||
operation=operation,
|
||||
bank_id=bank_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
result = await validator.precheck(ctx)
|
||||
if not result.allowed:
|
||||
raise HTTPException(
|
||||
status_code=result.status_code,
|
||||
detail=result.reason or "Operation not allowed",
|
||||
)
|
||||
|
||||
return _precheck_dep
|
||||
|
||||
# Global exception handler for authentication errors
|
||||
@app.exception_handler(AuthenticationError)
|
||||
async def authentication_error_handler(request, exc: AuthenticationError):
|
||||
@@ -3221,10 +3142,7 @@ def _register_routes(app: FastAPI):
|
||||
)
|
||||
@audited("recall")
|
||||
async def api_recall(
|
||||
bank_id: str,
|
||||
request: RecallRequest,
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
_precheck: None = Depends(precheck_for("recall")),
|
||||
bank_id: str, request: RecallRequest, request_context: RequestContext = Depends(get_request_context)
|
||||
):
|
||||
"""Run a recall and return results with trace."""
|
||||
import time
|
||||
@@ -3412,10 +3330,7 @@ def _register_routes(app: FastAPI):
|
||||
)
|
||||
@audited("reflect")
|
||||
async def api_reflect(
|
||||
bank_id: str,
|
||||
request: ReflectRequest,
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
_precheck: None = Depends(precheck_for("reflect")),
|
||||
bank_id: str, request: ReflectRequest, request_context: RequestContext = Depends(get_request_context)
|
||||
):
|
||||
metrics = get_metrics_collector()
|
||||
|
||||
@@ -3913,7 +3828,6 @@ def _register_routes(app: FastAPI):
|
||||
bank_id: str,
|
||||
body: CreateMentalModelRequest,
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
_precheck: None = Depends(precheck_for("mental_model_create")),
|
||||
):
|
||||
"""Create a mental model (async - returns operation_id)."""
|
||||
try:
|
||||
@@ -3962,7 +3876,6 @@ def _register_routes(app: FastAPI):
|
||||
bank_id: str,
|
||||
mental_model_id: str,
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
_precheck: None = Depends(precheck_for("mental_model_refresh")),
|
||||
):
|
||||
"""Refresh a mental model by re-running its source query (async)."""
|
||||
try:
|
||||
@@ -3989,48 +3902,6 @@ def _register_routes(app: FastAPI):
|
||||
)
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
@app.post(
|
||||
"/v1/default/banks/{bank_id}/mental-models/{mental_model_id}/clear",
|
||||
response_model=MentalModelResponse,
|
||||
summary="Clear mental model content",
|
||||
description=(
|
||||
"Clear a mental model's content so the next refresh performs a full re-synthesis. "
|
||||
"This is useful for delta-mode models that have accumulated drift over many "
|
||||
"incremental refreshes. After clearing, call the /refresh endpoint to trigger "
|
||||
"a clean full rebuild."
|
||||
),
|
||||
operation_id="clear_mental_model",
|
||||
tags=["Mental Models"],
|
||||
)
|
||||
@audited("clear_mental_model", request_param=None)
|
||||
async def api_clear_mental_model(
|
||||
bank_id: str,
|
||||
mental_model_id: str,
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
):
|
||||
"""Clear a mental model's content."""
|
||||
try:
|
||||
mental_model = await app.state.memory.clear_mental_model(
|
||||
bank_id=bank_id,
|
||||
mental_model_id=mental_model_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
if mental_model is None:
|
||||
raise HTTPException(status_code=404, detail=f"Mental model '{mental_model_id}' not found")
|
||||
return MentalModelResponse(**mental_model)
|
||||
except OperationValidationError as e:
|
||||
raise HTTPException(status_code=e.status_code, detail=e.reason)
|
||||
except (AuthenticationError, HTTPException):
|
||||
raise
|
||||
except Exception as e:
|
||||
import traceback
|
||||
|
||||
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
|
||||
logger.error(
|
||||
f"Error in POST /v1/default/banks/{bank_id}/mental-models/{mental_model_id}/clear: {error_detail}"
|
||||
)
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
@app.patch(
|
||||
"/v1/default/banks/{bank_id}/mental-models/{mental_model_id}",
|
||||
response_model=MentalModelResponse,
|
||||
@@ -5524,20 +5395,11 @@ def _register_routes(app: FastAPI):
|
||||
operation_id="trigger_consolidation",
|
||||
tags=["Banks"],
|
||||
)
|
||||
@audited("consolidation")
|
||||
async def api_trigger_consolidation(
|
||||
bank_id: str,
|
||||
request: ConsolidationRequest | None = None,
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
):
|
||||
@audited("consolidation", request_param=None)
|
||||
async def api_trigger_consolidation(bank_id: str, request_context: RequestContext = Depends(get_request_context)):
|
||||
"""Trigger consolidation for a bank (async)."""
|
||||
try:
|
||||
observation_scopes = request.observation_scopes if request else None
|
||||
result = await app.state.memory.submit_async_consolidation(
|
||||
bank_id=bank_id,
|
||||
request_context=request_context,
|
||||
observation_scopes=observation_scopes,
|
||||
)
|
||||
result = await app.state.memory.submit_async_consolidation(bank_id=bank_id, request_context=request_context)
|
||||
return ConsolidationResponse(
|
||||
operation_id=result["operation_id"],
|
||||
deduplicated=result.get("deduplicated", False),
|
||||
@@ -5860,10 +5722,7 @@ def _register_routes(app: FastAPI):
|
||||
)
|
||||
@audited("retain")
|
||||
async def api_retain(
|
||||
bank_id: str,
|
||||
request: RetainRequest,
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
_precheck: None = Depends(precheck_for("retain")),
|
||||
bank_id: str, request: RetainRequest, request_context: RequestContext = Depends(get_request_context)
|
||||
):
|
||||
"""Retain memories with optional async processing."""
|
||||
metrics = get_metrics_collector()
|
||||
@@ -5948,8 +5807,9 @@ def _register_routes(app: FastAPI):
|
||||
strategy=group_strategy,
|
||||
request_context=request_context,
|
||||
return_usage=True,
|
||||
outbox_callback_factory=app.state.memory._build_retain_outbox_callback_factory(
|
||||
outbox_callback=app.state.memory._build_retain_outbox_callback(
|
||||
bank_id=bank_id,
|
||||
contents=contents,
|
||||
operation_id=None,
|
||||
schema=_current_schema.get(),
|
||||
),
|
||||
@@ -6032,7 +5892,6 @@ def _register_routes(app: FastAPI):
|
||||
files: list[UploadFile] = File(..., description="Files to upload and convert"),
|
||||
request: str = Form(..., description="JSON string with FileRetainRequest model"),
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
_precheck: None = Depends(precheck_for("files_retain")),
|
||||
):
|
||||
"""Upload and convert files to memories."""
|
||||
from hindsight_api.config import get_config
|
||||
|
||||
@@ -107,7 +107,6 @@ def create_mcp_server(memory: MemoryEngine, multi_bank: bool = True) -> FastMCP:
|
||||
"update_mental_model",
|
||||
"delete_mental_model",
|
||||
"refresh_mental_model",
|
||||
"clear_mental_model",
|
||||
"list_directives",
|
||||
"create_directive",
|
||||
"delete_directive",
|
||||
|
||||
@@ -7,7 +7,6 @@ All environment variables and their defaults are defined here.
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
import sys
|
||||
from dataclasses import dataclass, field, fields
|
||||
from datetime import datetime, timezone
|
||||
@@ -15,7 +14,6 @@ from typing import Any, Literal
|
||||
|
||||
from dotenv import find_dotenv, load_dotenv
|
||||
|
||||
from ._pg_search import normalize_pg_search_tokenizer
|
||||
from ._vector_index import validate_extension
|
||||
from .utils import mask_network_location
|
||||
|
||||
@@ -138,7 +136,6 @@ ENV_LLM_MAX_RETRIES = "HINDSIGHT_API_LLM_MAX_RETRIES"
|
||||
ENV_LLM_INITIAL_BACKOFF = "HINDSIGHT_API_LLM_INITIAL_BACKOFF"
|
||||
ENV_LLM_MAX_BACKOFF = "HINDSIGHT_API_LLM_MAX_BACKOFF"
|
||||
ENV_LLM_TIMEOUT = "HINDSIGHT_API_LLM_TIMEOUT"
|
||||
ENV_LLM_REASONING_EFFORT = "HINDSIGHT_API_LLM_REASONING_EFFORT"
|
||||
ENV_LLM_GROQ_SERVICE_TIER = "HINDSIGHT_API_LLM_GROQ_SERVICE_TIER"
|
||||
ENV_LLM_OPENAI_SERVICE_TIER = "HINDSIGHT_API_LLM_OPENAI_SERVICE_TIER"
|
||||
ENV_LLM_EXTRA_BODY = "HINDSIGHT_API_LLM_EXTRA_BODY"
|
||||
@@ -172,17 +169,6 @@ ENV_RETAIN_LLM_MAX_BACKOFF = "HINDSIGHT_API_RETAIN_LLM_MAX_BACKOFF"
|
||||
ENV_RETAIN_LLM_TIMEOUT = "HINDSIGHT_API_RETAIN_LLM_TIMEOUT"
|
||||
ENV_RETAIN_LLM_LITELLMROUTER_CONFIG = "HINDSIGHT_API_RETAIN_LLM_LITELLMROUTER_CONFIG"
|
||||
|
||||
# Fireworks AI batch inference. Fireworks' batch API is a proprietary
|
||||
# account-scoped dataset/job REST API on a control-plane host, distinct from the
|
||||
# OpenAI-compatible inference host. account_id is REQUIRED for batch retain
|
||||
# (the control-plane endpoints are /v1/accounts/{account_id}/...). Static,
|
||||
# server-level config — it pairs with the Fireworks API key.
|
||||
ENV_FIREWORKS_ACCOUNT_ID = "HINDSIGHT_API_FIREWORKS_ACCOUNT_ID"
|
||||
ENV_FIREWORKS_BATCH_BASE_URL = "HINDSIGHT_API_FIREWORKS_BATCH_BASE_URL"
|
||||
ENV_FIREWORKS_BATCH_MAX_WAIT_SECONDS = "HINDSIGHT_API_FIREWORKS_BATCH_MAX_WAIT_SECONDS"
|
||||
DEFAULT_FIREWORKS_BATCH_BASE_URL = "https://api.fireworks.ai"
|
||||
DEFAULT_FIREWORKS_BATCH_MAX_WAIT_SECONDS = 86_400 # 24h — Fireworks' max job timeout
|
||||
|
||||
ENV_REFLECT_LLM_PROVIDER = "HINDSIGHT_API_REFLECT_LLM_PROVIDER"
|
||||
ENV_REFLECT_LLM_API_KEY = "HINDSIGHT_API_REFLECT_LLM_API_KEY"
|
||||
ENV_REFLECT_LLM_MODEL = "HINDSIGHT_API_REFLECT_LLM_MODEL"
|
||||
@@ -214,7 +200,6 @@ ENV_EMBEDDINGS_OPENAI_API_KEY = "HINDSIGHT_API_EMBEDDINGS_OPENAI_API_KEY"
|
||||
ENV_EMBEDDINGS_OPENAI_MODEL = "HINDSIGHT_API_EMBEDDINGS_OPENAI_MODEL"
|
||||
ENV_EMBEDDINGS_OPENAI_BASE_URL = "HINDSIGHT_API_EMBEDDINGS_OPENAI_BASE_URL"
|
||||
ENV_EMBEDDINGS_OPENAI_BATCH_SIZE = "HINDSIGHT_API_EMBEDDINGS_OPENAI_BATCH_SIZE"
|
||||
ENV_EMBEDDINGS_OPENAI_DIMENSIONS = "HINDSIGHT_API_EMBEDDINGS_OPENAI_DIMENSIONS"
|
||||
|
||||
# Gemini/Vertex AI embeddings configuration
|
||||
ENV_EMBEDDINGS_GEMINI_API_KEY = "HINDSIGHT_API_EMBEDDINGS_GEMINI_API_KEY"
|
||||
@@ -241,15 +226,6 @@ ENV_EMBEDDINGS_OPENROUTER_MODEL = "HINDSIGHT_API_EMBEDDINGS_OPENROUTER_MODEL"
|
||||
ENV_RERANKER_OPENROUTER_API_KEY = "HINDSIGHT_API_RERANKER_OPENROUTER_API_KEY"
|
||||
ENV_RERANKER_OPENROUTER_MODEL = "HINDSIGHT_API_RERANKER_OPENROUTER_MODEL"
|
||||
|
||||
# ZeroEntropy configuration (embeddings)
|
||||
ENV_EMBEDDINGS_ZEROENTROPY_API_KEY = "HINDSIGHT_API_EMBEDDINGS_ZEROENTROPY_API_KEY"
|
||||
ENV_EMBEDDINGS_ZEROENTROPY_MODEL = "HINDSIGHT_API_EMBEDDINGS_ZEROENTROPY_MODEL"
|
||||
ENV_EMBEDDINGS_ZEROENTROPY_BASE_URL = "HINDSIGHT_API_EMBEDDINGS_ZEROENTROPY_BASE_URL"
|
||||
ENV_EMBEDDINGS_ZEROENTROPY_DIMENSIONS = "HINDSIGHT_API_EMBEDDINGS_ZEROENTROPY_DIMENSIONS"
|
||||
ENV_EMBEDDINGS_ZEROENTROPY_ENCODING_FORMAT = "HINDSIGHT_API_EMBEDDINGS_ZEROENTROPY_ENCODING_FORMAT"
|
||||
ENV_EMBEDDINGS_ZEROENTROPY_LATENCY = "HINDSIGHT_API_EMBEDDINGS_ZEROENTROPY_LATENCY"
|
||||
ENV_EMBEDDINGS_ZEROENTROPY_BATCH_SIZE = "HINDSIGHT_API_EMBEDDINGS_ZEROENTROPY_BATCH_SIZE"
|
||||
|
||||
# Deprecated: Legacy shared Cohere API key (for backward compatibility)
|
||||
ENV_COHERE_API_KEY = "HINDSIGHT_API_COHERE_API_KEY"
|
||||
|
||||
@@ -288,14 +264,6 @@ ENV_RERANKER_TEI_URL = "HINDSIGHT_API_RERANKER_TEI_URL"
|
||||
ENV_RERANKER_TEI_BATCH_SIZE = "HINDSIGHT_API_RERANKER_TEI_BATCH_SIZE"
|
||||
ENV_RERANKER_TEI_MAX_CONCURRENT = "HINDSIGHT_API_RERANKER_TEI_MAX_CONCURRENT"
|
||||
ENV_RERANKER_TEI_HTTP_TIMEOUT = "HINDSIGHT_API_RERANKER_TEI_HTTP_TIMEOUT"
|
||||
ENV_RERANKER_COHERE_TIMEOUT = "HINDSIGHT_API_RERANKER_COHERE_TIMEOUT"
|
||||
ENV_RERANKER_OPENROUTER_TIMEOUT = "HINDSIGHT_API_RERANKER_OPENROUTER_TIMEOUT"
|
||||
ENV_RERANKER_ZEROENTROPY_TIMEOUT = "HINDSIGHT_API_RERANKER_ZEROENTROPY_TIMEOUT"
|
||||
ENV_RERANKER_SILICONFLOW_TIMEOUT = "HINDSIGHT_API_RERANKER_SILICONFLOW_TIMEOUT"
|
||||
ENV_RERANKER_ALIBABA_TIMEOUT = "HINDSIGHT_API_RERANKER_ALIBABA_TIMEOUT"
|
||||
ENV_RERANKER_LITELLM_TIMEOUT = "HINDSIGHT_API_RERANKER_LITELLM_TIMEOUT"
|
||||
ENV_RERANKER_LITELLM_SDK_TIMEOUT = "HINDSIGHT_API_RERANKER_LITELLM_SDK_TIMEOUT"
|
||||
ENV_RERANKER_GOOGLE_TIMEOUT = "HINDSIGHT_API_RERANKER_GOOGLE_TIMEOUT"
|
||||
ENV_RERANKER_MAX_CANDIDATES = "HINDSIGHT_API_RERANKER_MAX_CANDIDATES"
|
||||
ENV_RERANKER_FLASHRANK_MODEL = "HINDSIGHT_API_RERANKER_FLASHRANK_MODEL"
|
||||
ENV_RERANKER_FLASHRANK_CACHE_DIR = "HINDSIGHT_API_RERANKER_FLASHRANK_CACHE_DIR"
|
||||
@@ -311,10 +279,6 @@ ENV_RERANKER_SILICONFLOW_API_KEY = "HINDSIGHT_API_RERANKER_SILICONFLOW_API_KEY"
|
||||
ENV_RERANKER_SILICONFLOW_MODEL = "HINDSIGHT_API_RERANKER_SILICONFLOW_MODEL"
|
||||
ENV_RERANKER_SILICONFLOW_BASE_URL = "HINDSIGHT_API_RERANKER_SILICONFLOW_BASE_URL"
|
||||
|
||||
# Alibaba Cloud DashScope configuration (reranker only)
|
||||
ENV_RERANKER_ALIBABA_API_KEY = "HINDSIGHT_API_RERANKER_ALIBABA_API_KEY"
|
||||
ENV_RERANKER_ALIBABA_MODEL = "HINDSIGHT_API_RERANKER_ALIBABA_MODEL"
|
||||
|
||||
# Google Discovery Engine reranker configuration
|
||||
ENV_RERANKER_GOOGLE_MODEL = "HINDSIGHT_API_RERANKER_GOOGLE_MODEL"
|
||||
ENV_RERANKER_GOOGLE_PROJECT_ID = "HINDSIGHT_API_RERANKER_GOOGLE_PROJECT_ID"
|
||||
@@ -322,9 +286,6 @@ ENV_RERANKER_GOOGLE_SERVICE_ACCOUNT_KEY = "HINDSIGHT_API_RERANKER_GOOGLE_SERVICE
|
||||
|
||||
ENV_VECTOR_EXTENSION = "HINDSIGHT_API_VECTOR_EXTENSION"
|
||||
ENV_TEXT_SEARCH_EXTENSION = "HINDSIGHT_API_TEXT_SEARCH_EXTENSION"
|
||||
ENV_TEXT_SEARCH_EXTENSION_NATIVE_LANGUAGE = "HINDSIGHT_API_TEXT_SEARCH_EXTENSION_NATIVE_LANGUAGE"
|
||||
ENV_TEXT_SEARCH_EXTENSION_PG_SEARCH_TOKENIZER = "HINDSIGHT_API_TEXT_SEARCH_EXTENSION_PG_SEARCH_TOKENIZER"
|
||||
ENV_LLM_OUTPUT_LANGUAGE = "HINDSIGHT_API_LLM_OUTPUT_LANGUAGE"
|
||||
|
||||
ENV_HOST = "HINDSIGHT_API_HOST"
|
||||
ENV_PORT = "HINDSIGHT_API_PORT"
|
||||
@@ -333,7 +294,6 @@ ENV_LOG_LEVEL = "HINDSIGHT_API_LOG_LEVEL"
|
||||
ENV_LOG_FORMAT = "HINDSIGHT_API_LOG_FORMAT"
|
||||
ENV_LOG_JSON_FIELDS = "HINDSIGHT_API_LOG_JSON_FIELDS"
|
||||
ENV_WORKERS = "HINDSIGHT_API_WORKERS"
|
||||
ENV_ACCESS_LOG = "HINDSIGHT_API_ACCESS_LOG"
|
||||
ENV_MCP_ENABLED = "HINDSIGHT_API_MCP_ENABLED"
|
||||
ENV_MCP_ENABLED_TOOLS = "HINDSIGHT_API_MCP_ENABLED_TOOLS"
|
||||
ENV_MCP_STATELESS = "HINDSIGHT_API_MCP_STATELESS"
|
||||
@@ -373,7 +333,6 @@ ENV_RETAIN_CUSTOM_INSTRUCTIONS = "HINDSIGHT_API_RETAIN_CUSTOM_INSTRUCTIONS"
|
||||
ENV_RETAIN_DEFAULT_STRATEGY = "HINDSIGHT_API_RETAIN_DEFAULT_STRATEGY"
|
||||
ENV_RETAIN_BATCH_TOKENS = "HINDSIGHT_API_RETAIN_BATCH_TOKENS"
|
||||
ENV_RETAIN_ENTITY_LOOKUP = "HINDSIGHT_API_RETAIN_ENTITY_LOOKUP"
|
||||
ENV_RETAIN_ENTITY_RESOLUTION_BATCH_SIZE = "HINDSIGHT_API_RETAIN_ENTITY_RESOLUTION_BATCH_SIZE"
|
||||
ENV_RETAIN_BATCH_ENABLED = "HINDSIGHT_API_RETAIN_BATCH_ENABLED"
|
||||
ENV_RETAIN_BATCH_POLL_INTERVAL_SECONDS = "HINDSIGHT_API_RETAIN_BATCH_POLL_INTERVAL_SECONDS"
|
||||
ENV_RETAIN_CHUNK_BATCH_SIZE = "HINDSIGHT_API_RETAIN_CHUNK_BATCH_SIZE"
|
||||
@@ -402,11 +361,9 @@ ENV_FILE_DELETE_AFTER_RETAIN = "HINDSIGHT_API_FILE_DELETE_AFTER_RETAIN"
|
||||
|
||||
# Observations settings (consolidated knowledge from facts)
|
||||
ENV_ENABLE_OBSERVATIONS = "HINDSIGHT_API_ENABLE_OBSERVATIONS"
|
||||
ENV_ENABLE_AUTO_CONSOLIDATION = "HINDSIGHT_API_ENABLE_AUTO_CONSOLIDATION"
|
||||
ENV_CONSOLIDATION_BATCH_SIZE = "HINDSIGHT_API_CONSOLIDATION_BATCH_SIZE"
|
||||
ENV_CONSOLIDATION_MAX_MEMORIES_PER_ROUND = "HINDSIGHT_API_CONSOLIDATION_MAX_MEMORIES_PER_ROUND"
|
||||
ENV_CONSOLIDATION_LLM_BATCH_SIZE = "HINDSIGHT_API_CONSOLIDATION_LLM_BATCH_SIZE"
|
||||
ENV_CONSOLIDATION_LLM_PARALLELISM = "HINDSIGHT_API_CONSOLIDATION_LLM_PARALLELISM"
|
||||
ENV_CONSOLIDATION_MAX_TOKENS = "HINDSIGHT_API_CONSOLIDATION_MAX_TOKENS"
|
||||
ENV_CONSOLIDATION_SOURCE_FACTS_MAX_TOKENS = "HINDSIGHT_API_CONSOLIDATION_SOURCE_FACTS_MAX_TOKENS"
|
||||
ENV_CONSOLIDATION_SOURCE_FACTS_MAX_TOKENS_PER_OBSERVATION = (
|
||||
@@ -418,7 +375,6 @@ ENV_OBSERVATIONS_MISSION = "HINDSIGHT_API_OBSERVATIONS_MISSION"
|
||||
ENV_MAX_OBSERVATIONS_PER_SCOPE = "HINDSIGHT_API_MAX_OBSERVATIONS_PER_SCOPE"
|
||||
ENV_ENABLE_OBSERVATION_HISTORY = "HINDSIGHT_API_ENABLE_OBSERVATION_HISTORY"
|
||||
ENV_ENABLE_MENTAL_MODEL_HISTORY = "HINDSIGHT_API_ENABLE_MENTAL_MODEL_HISTORY"
|
||||
ENV_MENTAL_MODEL_HISTORY_MAX_ENTRIES = "HINDSIGHT_API_MENTAL_MODEL_HISTORY_MAX_ENTRIES"
|
||||
|
||||
# Webhook configuration (global, static - server-level only)
|
||||
ENV_WEBHOOK_URL = "HINDSIGHT_API_WEBHOOK_URL"
|
||||
@@ -453,7 +409,6 @@ ENV_WORKER_ENABLED = "HINDSIGHT_API_WORKER_ENABLED"
|
||||
ENV_WORKER_ID = "HINDSIGHT_API_WORKER_ID"
|
||||
ENV_WORKER_POLL_INTERVAL_MS = "HINDSIGHT_API_WORKER_POLL_INTERVAL_MS"
|
||||
ENV_WORKER_MAX_RETRIES = "HINDSIGHT_API_WORKER_MAX_RETRIES"
|
||||
ENV_WORKER_TASK_RETRY_BACKOFF_SECONDS = "HINDSIGHT_API_WORKER_TASK_RETRY_BACKOFF_SECONDS"
|
||||
ENV_WORKER_HTTP_PORT = "HINDSIGHT_API_WORKER_HTTP_PORT"
|
||||
ENV_WORKER_MAX_SLOTS = "HINDSIGHT_API_WORKER_MAX_SLOTS"
|
||||
|
||||
@@ -467,9 +422,7 @@ WORKER_SLOT_RESERVATION_TYPES: dict[str, tuple[str, int]] = {
|
||||
"retain": ("HINDSIGHT_API_WORKER_RETAIN_MAX_SLOTS", 0),
|
||||
"file_convert_retain": ("HINDSIGHT_API_WORKER_FILE_CONVERT_RETAIN_MAX_SLOTS", 0),
|
||||
"refresh_mental_model": ("HINDSIGHT_API_WORKER_REFRESH_MENTAL_MODEL_MAX_SLOTS", 0),
|
||||
"graph_maintenance": ("HINDSIGHT_API_WORKER_GRAPH_MAINTENANCE_MAX_SLOTS", 0),
|
||||
}
|
||||
ENV_WORKER_CONSOLIDATION_BANK_PRIORITY = "HINDSIGHT_API_WORKER_CONSOLIDATION_BANK_PRIORITY"
|
||||
ENV_RETAIN_MAX_CONCURRENT = "HINDSIGHT_API_RETAIN_MAX_CONCURRENT"
|
||||
|
||||
# Reflect agent settings
|
||||
@@ -520,7 +473,6 @@ PROVIDER_DEFAULT_MODELS = {
|
||||
"zai": "glm-4.5-flash",
|
||||
"opencode-go": "deepseek-v4-flash",
|
||||
"ollama": "gemma3:12b",
|
||||
"ollama-cloud": "gemma3:12b",
|
||||
"llamacpp": "gemma-4-e2b-it",
|
||||
"lmstudio": "local-model",
|
||||
"vertexai": "google/gemini-2.5-flash-lite",
|
||||
@@ -532,7 +484,6 @@ PROVIDER_DEFAULT_MODELS = {
|
||||
"bedrock": "us.amazon.nova-2-lite-v1:0",
|
||||
"volcano": "doubao-pro-32k",
|
||||
"openrouter": "qwen/qwen3.5-9b",
|
||||
"fireworks": "accounts/fireworks/models/llama-v3p1-8b-instruct",
|
||||
}
|
||||
DEFAULT_LLM_MODEL = "gpt-4o-mini" # Fallback if provider not in table
|
||||
# Built-in llama.cpp defaults
|
||||
@@ -547,7 +498,6 @@ DEFAULT_LLM_MAX_RETRIES = 3 # Max retry attempts for LLM API calls
|
||||
DEFAULT_LLM_INITIAL_BACKOFF = 1.0 # Initial backoff in seconds for retry exponential backoff
|
||||
DEFAULT_LLM_MAX_BACKOFF = 60.0 # Max backoff cap in seconds for retry exponential backoff
|
||||
DEFAULT_LLM_TIMEOUT = 120.0 # seconds
|
||||
DEFAULT_LLM_REASONING_EFFORT = "low"
|
||||
|
||||
# Vertex AI defaults
|
||||
DEFAULT_LLM_VERTEXAI_PROJECT_ID = None # Required for Vertex AI
|
||||
@@ -581,16 +531,6 @@ DEFAULT_RERANKER_LOCAL_BATCH_SIZE = 32 # Batch size for local reranker predict(
|
||||
DEFAULT_RERANKER_TEI_BATCH_SIZE = 128
|
||||
DEFAULT_RERANKER_TEI_MAX_CONCURRENT = 8
|
||||
DEFAULT_RERANKER_TEI_HTTP_TIMEOUT = 30.0 # HTTP timeout for TEI reranker requests (seconds)
|
||||
# HTTP timeout (seconds) for remote rerank providers. Defaults match the previous
|
||||
# hardcoded constructor defaults so unset envs keep current behavior.
|
||||
DEFAULT_RERANKER_COHERE_TIMEOUT = 60.0
|
||||
DEFAULT_RERANKER_OPENROUTER_TIMEOUT = 60.0
|
||||
DEFAULT_RERANKER_ZEROENTROPY_TIMEOUT = 60.0
|
||||
DEFAULT_RERANKER_SILICONFLOW_TIMEOUT = 60.0
|
||||
DEFAULT_RERANKER_ALIBABA_TIMEOUT = 60.0
|
||||
DEFAULT_RERANKER_LITELLM_TIMEOUT = 60.0
|
||||
DEFAULT_RERANKER_LITELLM_SDK_TIMEOUT = 60.0
|
||||
DEFAULT_RERANKER_GOOGLE_TIMEOUT = 60.0
|
||||
DEFAULT_RERANKER_MAX_CANDIDATES = 300
|
||||
DEFAULT_RERANKER_FLASHRANK_MODEL = "ms-marco-MiniLM-L-12-v2" # Best balance of speed and quality
|
||||
DEFAULT_RERANKER_FLASHRANK_CACHE_DIR = None # Use default cache directory
|
||||
@@ -603,39 +543,18 @@ DEFAULT_RERANKER_COHERE_MODEL = "rerank-english-v3.0"
|
||||
DEFAULT_EMBEDDINGS_OPENROUTER_MODEL = "perplexity/pplx-embed-v1-0.6b"
|
||||
DEFAULT_RERANKER_OPENROUTER_MODEL = "cohere/rerank-v3.5"
|
||||
|
||||
# ZeroEntropy defaults
|
||||
DEFAULT_EMBEDDINGS_ZEROENTROPY_MODEL = "zembed-1"
|
||||
# Shared between embeddings (zembed-1) and reranker (zerank-*) — the host is the same.
|
||||
DEFAULT_ZEROENTROPY_BASE_URL = "https://api.zeroentropy.dev"
|
||||
# ZeroEntropy's API default is 2560, but Hindsight defaults to 1280 so the
|
||||
# provider works with pgvector HNSW's 2000-dimension index limit out of the box.
|
||||
DEFAULT_EMBEDDINGS_ZEROENTROPY_DIMENSIONS = 1280
|
||||
DEFAULT_EMBEDDINGS_ZEROENTROPY_ENCODING_FORMAT = "float"
|
||||
DEFAULT_EMBEDDINGS_ZEROENTROPY_LATENCY = None
|
||||
DEFAULT_EMBEDDINGS_ZEROENTROPY_BATCH_SIZE = 100
|
||||
|
||||
DEFAULT_RERANKER_ZEROENTROPY_MODEL = "zerank-2"
|
||||
|
||||
DEFAULT_RERANKER_SILICONFLOW_MODEL = "BAAI/bge-reranker-v2-m3"
|
||||
DEFAULT_RERANKER_SILICONFLOW_BASE_URL = "https://api.siliconflow.cn/v1"
|
||||
|
||||
DEFAULT_RERANKER_ALIBABA_MODEL = "qwen3-rerank"
|
||||
|
||||
DEFAULT_RERANKER_GOOGLE_MODEL = "semantic-ranker-default-004"
|
||||
|
||||
# Vector extension (pgvector, vchord, pgvectorscale, or AlloyDB ScaNN)
|
||||
DEFAULT_VECTOR_EXTENSION = "pgvector" # Options: "pgvector", "vchord", "pgvectorscale", "scann"
|
||||
|
||||
# Text search extension (native PostgreSQL, vchord BM25, Timescale pg_textsearch,
|
||||
# pgroonga, or ParadeDB pg_search)
|
||||
DEFAULT_TEXT_SEARCH_EXTENSION = "native" # Options: "native", "vchord", "pg_textsearch", "pgroonga", "pg_search"
|
||||
|
||||
# PostgreSQL text search dictionary used by the native tsvector backend. Only
|
||||
# affects text_search_extension == "native"; other backends use their own
|
||||
# tokenizers (vchord: llmlingua2, pg_textsearch: hardcoded english,
|
||||
# pgroonga: TokenBigram polyglot, pg_search: per-field Tantivy tokenizer).
|
||||
DEFAULT_TEXT_SEARCH_EXTENSION_NATIVE_LANGUAGE = "english"
|
||||
DEFAULT_TEXT_SEARCH_EXTENSION_PG_SEARCH_TOKENIZER = ""
|
||||
# Text search extension (native PostgreSQL, vchord BM25, or Timescale pg_textsearch)
|
||||
DEFAULT_TEXT_SEARCH_EXTENSION = "native" # Options: "native", "vchord", "pg_textsearch"
|
||||
|
||||
# LiteLLM defaults
|
||||
DEFAULT_LITELLM_API_BASE = "http://localhost:4000"
|
||||
@@ -654,7 +573,6 @@ DEFAULT_BASE_PATH = "" # Empty string = root path
|
||||
DEFAULT_LOG_LEVEL = "info"
|
||||
DEFAULT_LOG_FORMAT = "text" # Options: "text", "json"
|
||||
DEFAULT_WORKERS = 1
|
||||
DEFAULT_ACCESS_LOG = False
|
||||
DEFAULT_MCP_ENABLED = True
|
||||
DEFAULT_MCP_ENABLED_TOOLS: list[str] | None = None # None = all tools enabled
|
||||
DEFAULT_MCP_STATELESS = False # False = stateful (supports SSE/GET); True = stateless (POST-only)
|
||||
@@ -683,7 +601,6 @@ DEFAULT_RETAIN_CHUNK_BATCH_SIZE = (
|
||||
)
|
||||
DEFAULT_RETAIN_BATCH_TOKENS = 10_000 # ~40KB of text # Max chars per sub-batch for async retain auto-splitting
|
||||
DEFAULT_RETAIN_ENTITY_LOOKUP = "trigram" # "full" or "trigram"
|
||||
DEFAULT_RETAIN_ENTITY_RESOLUTION_BATCH_SIZE = 100 # Unique entity names per pg_trgm candidate lookup query
|
||||
DEFAULT_RETAIN_BATCH_ENABLED = False # Use LLM Batch API for fact extraction (only when async=True)
|
||||
DEFAULT_RETAIN_BATCH_POLL_INTERVAL_SECONDS = 60 # Batch API polling interval in seconds
|
||||
|
||||
@@ -698,25 +615,14 @@ DEFAULT_FILE_DELETE_AFTER_RETAIN = True # Delete file bytes after retain (saves
|
||||
|
||||
# Observations defaults (consolidated knowledge from facts)
|
||||
DEFAULT_ENABLE_OBSERVATIONS = True # Observations enabled by default
|
||||
DEFAULT_ENABLE_AUTO_CONSOLIDATION = True # Auto-consolidation after retain enabled by default
|
||||
DEFAULT_ENABLE_OBSERVATION_HISTORY = True # Observation history tracking enabled by default
|
||||
DEFAULT_ENABLE_MENTAL_MODEL_HISTORY = True # Mental model history tracking enabled by default
|
||||
# Each history entry snapshots previous_content + previous_reflect_response. Without
|
||||
# a cap, sustained mental-model refresh load grows the jsonb array unboundedly until
|
||||
# it crosses Postgres's hard 256MB jsonb limit and subsequent UPDATEs fail with
|
||||
# SQLSTATE 54000. 50 keeps the array well under 100MB even with large reflect
|
||||
# responses, while preserving enough recent history for meaningful audit / rollback.
|
||||
DEFAULT_MENTAL_MODEL_HISTORY_MAX_ENTRIES = 50
|
||||
DEFAULT_CONSOLIDATION_MAX_ATTEMPTS = 3 # Outer retry attempts for consolidation LLM batch calls
|
||||
DEFAULT_CONSOLIDATION_BATCH_SIZE = 50 # Memories to load per batch (internal memory optimization)
|
||||
DEFAULT_CONSOLIDATION_MAX_MEMORIES_PER_ROUND = (
|
||||
100 # Max memories per consolidation round (0 = unlimited). Limits how long one bank holds a worker slot.
|
||||
)
|
||||
DEFAULT_CONSOLIDATION_LLM_BATCH_SIZE = 8 # Facts per LLM call (1 = no batching; >1 = batch mode)
|
||||
DEFAULT_CONSOLIDATION_LLM_PARALLELISM = (
|
||||
4 # Max tag groups consolidated concurrently per op. Locks on overlapping write
|
||||
# scopes degrade to sequential automatically; matches retain_max_concurrent.
|
||||
)
|
||||
DEFAULT_CONSOLIDATION_MAX_TOKENS = 512 # Max tokens for recall when finding related observations
|
||||
DEFAULT_CONSOLIDATION_RECALL_BUDGET = "low" # Budget level for consolidation recall (low/mid/high)
|
||||
DEFAULT_CONSOLIDATION_SOURCE_FACTS_MAX_TOKENS = (
|
||||
@@ -743,7 +649,6 @@ DEFAULT_WORKER_ENABLED = True # API runs worker by default (standalone mode)
|
||||
DEFAULT_WORKER_ID = None # Will use hostname if not specified
|
||||
DEFAULT_WORKER_POLL_INTERVAL_MS = 500 # Poll database every 500ms
|
||||
DEFAULT_WORKER_MAX_RETRIES = 3 # Max retries before marking task failed
|
||||
DEFAULT_WORKER_TASK_RETRY_BACKOFF_SECONDS = 60 # Seconds between retries on transient task failure
|
||||
DEFAULT_WORKER_HTTP_PORT = 8889 # HTTP port for worker metrics/health
|
||||
DEFAULT_WORKER_MAX_SLOTS = 10 # Total concurrent tasks per worker
|
||||
DEFAULT_RETAIN_MAX_CONCURRENT = 4 # Max concurrent retain DB phases (HNSW reads + writes). Limits I/O contention.
|
||||
@@ -887,24 +792,6 @@ def _parse_positive_int(name: str, raw: str | None, default: int) -> int:
|
||||
return parsed
|
||||
|
||||
|
||||
def _parse_optional_positive_int(name: str, raw: str | None) -> int | None:
|
||||
"""Parse an optional env var that must be a positive integer when set."""
|
||||
if raw is None or raw == "":
|
||||
return None
|
||||
return _parse_positive_int(name, raw, 1)
|
||||
|
||||
|
||||
def _parse_optional_choice(name: str, raw: str | None, allowed: frozenset[str]) -> str | None:
|
||||
"""Parse an optional string env var constrained to a small allowlist."""
|
||||
if raw is None or raw == "":
|
||||
return None
|
||||
normalized = raw.lower()
|
||||
if normalized not in allowed:
|
||||
values = ", ".join(sorted(allowed))
|
||||
raise ValueError(f"{name} must be one of {values}, got {raw!r}")
|
||||
return normalized
|
||||
|
||||
|
||||
def _validate_extraction_mode(mode: str) -> str:
|
||||
"""Validate and normalize extraction mode."""
|
||||
mode_lower = mode.lower()
|
||||
@@ -929,38 +816,6 @@ def _validate_recall_budget_function(function: str) -> str:
|
||||
return function_lower
|
||||
|
||||
|
||||
def _parse_bank_priority(raw: str) -> dict[str, int]:
|
||||
"""Parse ``bank-pattern:priority,...`` into ``{pattern: priority}``.
|
||||
|
||||
``*`` in a pattern is kept as-is here; the SQL layer converts it to ``%``
|
||||
for LIKE matching. A bare ``*`` key is the catch-all default for unlisted
|
||||
banks. Returns an empty dict when *raw* is blank.
|
||||
"""
|
||||
result: dict[str, int] = {}
|
||||
raw = raw.strip()
|
||||
if not raw:
|
||||
return result
|
||||
for entry in raw.split(","):
|
||||
entry = entry.strip()
|
||||
if not entry:
|
||||
continue
|
||||
if ":" not in entry:
|
||||
raise ValueError(f"Invalid bank priority entry '{entry}': expected 'bank-pattern:priority'")
|
||||
pattern, priority_str = entry.rsplit(":", 1)
|
||||
pattern = pattern.strip()
|
||||
priority_str = priority_str.strip()
|
||||
if not pattern:
|
||||
raise ValueError(f"Empty bank pattern in entry '{entry}'")
|
||||
try:
|
||||
priority = int(priority_str)
|
||||
except ValueError:
|
||||
raise ValueError(f"Invalid priority '{priority_str}' in entry '{entry}': must be an integer") from None
|
||||
if priority < 1:
|
||||
raise ValueError(f"Priority must be >= 1, got {priority} in entry '{entry}'")
|
||||
result[pattern] = priority
|
||||
return result
|
||||
|
||||
|
||||
def _get_default_model_for_provider(provider: str) -> str:
|
||||
"""Get the default model for a given provider."""
|
||||
return PROVIDER_DEFAULT_MODELS.get(provider.lower(), DEFAULT_LLM_MODEL)
|
||||
@@ -1019,19 +874,7 @@ class HindsightConfig:
|
||||
migration_database_url: str | None
|
||||
database_schema: str
|
||||
vector_extension: str # "pgvector", "vchord", "pgvectorscale", or "scann"
|
||||
text_search_extension: str # "native", "vchord", "pg_textsearch", "pgroonga", or "pg_search"
|
||||
# PostgreSQL text search dictionary for the "native" backend (ignored by
|
||||
# other backends). Only the "native" backend reads this field; pgroonga
|
||||
# uses TokenBigram, vchord uses llmlingua2, pg_textsearch hardcodes english,
|
||||
# pg_search uses Tantivy per-field tokenizers.
|
||||
text_search_extension_native_language: str
|
||||
# ParadeDB pg_search tokenizer used when building BM25 indexes. Empty keeps
|
||||
# ParadeDB's default tokenizer.
|
||||
text_search_extension_pg_search_tokenizer: str
|
||||
# When set, every LLM-generated artifact (retain facts, consolidation
|
||||
# observations, reflect responses) is forced into this language regardless
|
||||
# of the source content. Unset preserves source language.
|
||||
llm_output_language: str | None
|
||||
text_search_extension: str # "native" or "vchord"
|
||||
|
||||
# LLM (default, used as fallback for per-operation config)
|
||||
llm_provider: str
|
||||
@@ -1043,7 +886,6 @@ class HindsightConfig:
|
||||
llm_initial_backoff: float
|
||||
llm_max_backoff: float
|
||||
llm_timeout: float
|
||||
llm_reasoning_effort: str
|
||||
llm_groq_service_tier: str # Groq: "on_demand", "flex", or "auto"
|
||||
llm_openai_service_tier: str | None # OpenAI: None (default) or "flex" (50% cheaper)
|
||||
llm_extra_body: (
|
||||
@@ -1087,11 +929,6 @@ class HindsightConfig:
|
||||
retain_llm_timeout: float | None
|
||||
retain_llm_litellmrouter_config: dict | None
|
||||
|
||||
# Fireworks AI batch inference (static, server-level)
|
||||
fireworks_account_id: str | None
|
||||
fireworks_batch_base_url: str
|
||||
fireworks_batch_max_wait_seconds: int
|
||||
|
||||
reflect_llm_provider: str | None
|
||||
reflect_llm_api_key: str | None
|
||||
reflect_llm_model: str | None
|
||||
@@ -1161,34 +998,24 @@ class HindsightConfig:
|
||||
reranker_cohere_api_key: str | None
|
||||
reranker_cohere_model: str
|
||||
reranker_cohere_base_url: str | None
|
||||
reranker_cohere_timeout: float
|
||||
reranker_openrouter_api_key: str | None
|
||||
reranker_openrouter_model: str
|
||||
reranker_openrouter_timeout: float
|
||||
reranker_litellm_api_base: str
|
||||
reranker_litellm_api_key: str | None
|
||||
reranker_litellm_model: str
|
||||
reranker_litellm_max_tokens_per_doc: int | None
|
||||
reranker_litellm_timeout: float
|
||||
reranker_litellm_sdk_api_key: str | None
|
||||
reranker_litellm_sdk_model: str
|
||||
reranker_litellm_sdk_api_base: str | None
|
||||
reranker_litellm_sdk_timeout: float
|
||||
reranker_zeroentropy_api_key: str | None
|
||||
reranker_zeroentropy_model: str
|
||||
reranker_zeroentropy_base_url: str | None
|
||||
reranker_zeroentropy_timeout: float
|
||||
reranker_siliconflow_api_key: str | None
|
||||
reranker_siliconflow_model: str
|
||||
reranker_siliconflow_base_url: str
|
||||
reranker_siliconflow_timeout: float
|
||||
reranker_alibaba_api_key: str | None
|
||||
reranker_alibaba_model: str
|
||||
reranker_alibaba_timeout: float
|
||||
reranker_google_model: str
|
||||
reranker_google_project_id: str | None
|
||||
reranker_google_service_account_key: str | None
|
||||
reranker_google_timeout: float
|
||||
|
||||
# Server
|
||||
host: str
|
||||
@@ -1227,7 +1054,6 @@ class HindsightConfig:
|
||||
retain_batch_enabled: bool
|
||||
retain_batch_poll_interval_seconds: int
|
||||
retain_entity_lookup: str # "full" or "trigram"
|
||||
retain_entity_resolution_batch_size: int # Unique entity names per pg_trgm candidate lookup query
|
||||
retain_chunk_batch_size: int # Max chunks per streaming batch (0 = disabled)
|
||||
|
||||
# File storage (static - server-level only)
|
||||
@@ -1254,14 +1080,11 @@ class HindsightConfig:
|
||||
|
||||
# Observations settings (consolidated knowledge from facts)
|
||||
enable_observations: bool
|
||||
enable_auto_consolidation: bool
|
||||
enable_observation_history: bool
|
||||
enable_mental_model_history: bool
|
||||
mental_model_history_max_entries: int
|
||||
consolidation_batch_size: int
|
||||
consolidation_max_memories_per_round: int
|
||||
consolidation_llm_batch_size: int
|
||||
consolidation_llm_parallelism: int
|
||||
consolidation_max_tokens: int
|
||||
consolidation_recall_budget: str
|
||||
consolidation_source_facts_max_tokens: int
|
||||
@@ -1324,11 +1147,9 @@ class HindsightConfig:
|
||||
worker_id: str | None
|
||||
worker_poll_interval_ms: int
|
||||
worker_max_retries: int
|
||||
worker_task_retry_backoff_seconds: int
|
||||
worker_http_port: int
|
||||
worker_max_slots: int
|
||||
worker_slot_reservations: dict[str, int]
|
||||
worker_consolidation_bank_priority: dict[str, int]
|
||||
retain_max_concurrent: int
|
||||
|
||||
# Reflect agent settings
|
||||
@@ -1358,14 +1179,6 @@ class HindsightConfig:
|
||||
# Defaulted fields (source-compatible additions — existing direct constructor callers keep working).
|
||||
# Keep at the end of the dataclass; Python forbids non-default fields after default fields.
|
||||
embeddings_openai_batch_size: int = DEFAULT_EMBEDDINGS_OPENAI_BATCH_SIZE
|
||||
embeddings_openai_dimensions: int | None = None
|
||||
embeddings_zeroentropy_api_key: str | None = None
|
||||
embeddings_zeroentropy_model: str = DEFAULT_EMBEDDINGS_ZEROENTROPY_MODEL
|
||||
embeddings_zeroentropy_base_url: str = DEFAULT_ZEROENTROPY_BASE_URL
|
||||
embeddings_zeroentropy_dimensions: int = DEFAULT_EMBEDDINGS_ZEROENTROPY_DIMENSIONS
|
||||
embeddings_zeroentropy_encoding_format: str = DEFAULT_EMBEDDINGS_ZEROENTROPY_ENCODING_FORMAT
|
||||
embeddings_zeroentropy_batch_size: int = DEFAULT_EMBEDDINGS_ZEROENTROPY_BATCH_SIZE
|
||||
embeddings_zeroentropy_latency: str | None = DEFAULT_EMBEDDINGS_ZEROENTROPY_LATENCY
|
||||
|
||||
# Class-level sets for configuration categorization
|
||||
|
||||
@@ -1389,7 +1202,6 @@ class HindsightConfig:
|
||||
"embeddings_tei_base_url",
|
||||
"reranker_tei_base_url",
|
||||
"reranker_cohere_base_url",
|
||||
"embeddings_zeroentropy_base_url",
|
||||
"reranker_zeroentropy_base_url",
|
||||
"reranker_siliconflow_base_url",
|
||||
# Service Account Keys
|
||||
@@ -1398,7 +1210,6 @@ class HindsightConfig:
|
||||
"reranker_google_service_account_key",
|
||||
# Embeddings API keys
|
||||
"embeddings_gemini_api_key",
|
||||
"embeddings_zeroentropy_api_key",
|
||||
# File storage credentials
|
||||
"file_storage_s3_access_key_id",
|
||||
"file_storage_s3_secret_access_key",
|
||||
@@ -1428,9 +1239,7 @@ class HindsightConfig:
|
||||
"entities_allow_free_form",
|
||||
# Consolidation settings
|
||||
"enable_observations",
|
||||
"enable_auto_consolidation",
|
||||
"consolidation_llm_batch_size",
|
||||
"consolidation_llm_parallelism",
|
||||
"consolidation_max_memories_per_round",
|
||||
"consolidation_source_facts_max_tokens",
|
||||
"consolidation_source_facts_max_tokens_per_observation",
|
||||
@@ -1525,30 +1334,12 @@ class HindsightConfig:
|
||||
validate_extension(self.vector_extension)
|
||||
|
||||
# Validate text_search_extension
|
||||
valid_text_search = ("native", "vchord", "pg_textsearch", "pgroonga", "pg_search")
|
||||
valid_text_search = ("native", "vchord", "pg_textsearch")
|
||||
if self.text_search_extension not in valid_text_search:
|
||||
raise ValueError(
|
||||
f"Invalid text_search_extension: {self.text_search_extension}. Must be one of: {', '.join(valid_text_search)}"
|
||||
)
|
||||
|
||||
# Validate text_search_extension_native_language as a PG identifier.
|
||||
# Embedded directly into raw SQL via to_tsvector('<lang>', ...), so we
|
||||
# reject anything that isn't a plain identifier to prevent injection.
|
||||
# Intentionally permissive about which dictionaries exist — users may
|
||||
# install custom ones like zhparser; we only check shape here. PG
|
||||
# raises a clear error at query time if the dictionary is missing.
|
||||
if not re.fullmatch(r"[a-zA-Z_][a-zA-Z0-9_]*", self.text_search_extension_native_language):
|
||||
raise ValueError(
|
||||
f"Invalid text_search_extension_native_language: "
|
||||
f"{self.text_search_extension_native_language!r}. Must be a valid PostgreSQL identifier "
|
||||
f"(letters, digits, underscores; not starting with a digit). Examples: 'english', "
|
||||
f"'french', 'simple', 'zhparser'."
|
||||
)
|
||||
|
||||
self.text_search_extension_pg_search_tokenizer = normalize_pg_search_tokenizer(
|
||||
self.text_search_extension_pg_search_tokenizer
|
||||
)
|
||||
|
||||
# When LLM provider is "none", force chunks-only mode and disable LLM-dependent features
|
||||
if self.llm_provider == "none":
|
||||
self.retain_extraction_mode = "chunks"
|
||||
@@ -1625,15 +1416,6 @@ class HindsightConfig:
|
||||
database_schema=os.getenv(ENV_DATABASE_SCHEMA, DEFAULT_DATABASE_SCHEMA),
|
||||
vector_extension=os.getenv(ENV_VECTOR_EXTENSION, DEFAULT_VECTOR_EXTENSION).lower(),
|
||||
text_search_extension=os.getenv(ENV_TEXT_SEARCH_EXTENSION, DEFAULT_TEXT_SEARCH_EXTENSION).lower(),
|
||||
text_search_extension_native_language=os.getenv(
|
||||
ENV_TEXT_SEARCH_EXTENSION_NATIVE_LANGUAGE,
|
||||
DEFAULT_TEXT_SEARCH_EXTENSION_NATIVE_LANGUAGE,
|
||||
).lower(),
|
||||
text_search_extension_pg_search_tokenizer=os.getenv(
|
||||
ENV_TEXT_SEARCH_EXTENSION_PG_SEARCH_TOKENIZER,
|
||||
DEFAULT_TEXT_SEARCH_EXTENSION_PG_SEARCH_TOKENIZER,
|
||||
),
|
||||
llm_output_language=(os.getenv(ENV_LLM_OUTPUT_LANGUAGE) or None),
|
||||
# LLM
|
||||
llm_provider=llm_provider,
|
||||
llm_api_key=os.getenv(ENV_LLM_API_KEY),
|
||||
@@ -1644,7 +1426,6 @@ class HindsightConfig:
|
||||
llm_initial_backoff=float(os.getenv(ENV_LLM_INITIAL_BACKOFF, str(DEFAULT_LLM_INITIAL_BACKOFF))),
|
||||
llm_max_backoff=float(os.getenv(ENV_LLM_MAX_BACKOFF, str(DEFAULT_LLM_MAX_BACKOFF))),
|
||||
llm_timeout=float(os.getenv(ENV_LLM_TIMEOUT, str(DEFAULT_LLM_TIMEOUT))),
|
||||
llm_reasoning_effort=os.getenv(ENV_LLM_REASONING_EFFORT, DEFAULT_LLM_REASONING_EFFORT),
|
||||
llm_groq_service_tier=os.getenv(ENV_LLM_GROQ_SERVICE_TIER, DEFAULT_LLM_GROQ_SERVICE_TIER),
|
||||
llm_openai_service_tier=os.getenv(ENV_LLM_OPENAI_SERVICE_TIER, DEFAULT_LLM_OPENAI_SERVICE_TIER),
|
||||
llm_extra_body=json.loads(os.getenv(ENV_LLM_EXTRA_BODY, "null")),
|
||||
@@ -1675,11 +1456,6 @@ class HindsightConfig:
|
||||
else None
|
||||
),
|
||||
retain_llm_base_url=os.getenv(ENV_RETAIN_LLM_BASE_URL) or None,
|
||||
fireworks_account_id=os.getenv(ENV_FIREWORKS_ACCOUNT_ID) or None,
|
||||
fireworks_batch_base_url=os.getenv(ENV_FIREWORKS_BATCH_BASE_URL) or DEFAULT_FIREWORKS_BATCH_BASE_URL,
|
||||
fireworks_batch_max_wait_seconds=int(
|
||||
os.getenv(ENV_FIREWORKS_BATCH_MAX_WAIT_SECONDS, str(DEFAULT_FIREWORKS_BATCH_MAX_WAIT_SECONDS))
|
||||
),
|
||||
retain_llm_max_concurrent=int(os.getenv(ENV_RETAIN_LLM_MAX_CONCURRENT))
|
||||
if os.getenv(ENV_RETAIN_LLM_MAX_CONCURRENT)
|
||||
else None,
|
||||
@@ -1762,10 +1538,6 @@ class HindsightConfig:
|
||||
os.getenv(ENV_EMBEDDINGS_OPENAI_BATCH_SIZE),
|
||||
DEFAULT_EMBEDDINGS_OPENAI_BATCH_SIZE,
|
||||
),
|
||||
embeddings_openai_dimensions=_parse_optional_positive_int(
|
||||
ENV_EMBEDDINGS_OPENAI_DIMENSIONS,
|
||||
os.getenv(ENV_EMBEDDINGS_OPENAI_DIMENSIONS),
|
||||
),
|
||||
# Cohere embeddings (with backward-compatible fallback to shared API key)
|
||||
embeddings_cohere_api_key=os.getenv(ENV_EMBEDDINGS_COHERE_API_KEY) or os.getenv(ENV_COHERE_API_KEY),
|
||||
embeddings_cohere_model=os.getenv(ENV_EMBEDDINGS_COHERE_MODEL, DEFAULT_EMBEDDINGS_COHERE_MODEL),
|
||||
@@ -1778,36 +1550,6 @@ class HindsightConfig:
|
||||
or os.getenv(ENV_OPENROUTER_API_KEY)
|
||||
or os.getenv(ENV_LLM_API_KEY),
|
||||
embeddings_openrouter_model=os.getenv(ENV_EMBEDDINGS_OPENROUTER_MODEL, DEFAULT_EMBEDDINGS_OPENROUTER_MODEL),
|
||||
# ZeroEntropy embeddings
|
||||
embeddings_zeroentropy_api_key=os.getenv(ENV_EMBEDDINGS_ZEROENTROPY_API_KEY)
|
||||
or os.getenv("ZEROENTROPY_API_KEY"),
|
||||
embeddings_zeroentropy_model=os.getenv(
|
||||
ENV_EMBEDDINGS_ZEROENTROPY_MODEL, DEFAULT_EMBEDDINGS_ZEROENTROPY_MODEL
|
||||
),
|
||||
embeddings_zeroentropy_base_url=os.getenv(
|
||||
ENV_EMBEDDINGS_ZEROENTROPY_BASE_URL, DEFAULT_ZEROENTROPY_BASE_URL
|
||||
),
|
||||
embeddings_zeroentropy_dimensions=_parse_positive_int(
|
||||
ENV_EMBEDDINGS_ZEROENTROPY_DIMENSIONS,
|
||||
os.getenv(ENV_EMBEDDINGS_ZEROENTROPY_DIMENSIONS),
|
||||
DEFAULT_EMBEDDINGS_ZEROENTROPY_DIMENSIONS,
|
||||
),
|
||||
embeddings_zeroentropy_encoding_format=_parse_optional_choice(
|
||||
ENV_EMBEDDINGS_ZEROENTROPY_ENCODING_FORMAT,
|
||||
os.getenv(ENV_EMBEDDINGS_ZEROENTROPY_ENCODING_FORMAT),
|
||||
frozenset({"float", "base64"}),
|
||||
)
|
||||
or DEFAULT_EMBEDDINGS_ZEROENTROPY_ENCODING_FORMAT,
|
||||
embeddings_zeroentropy_latency=_parse_optional_choice(
|
||||
ENV_EMBEDDINGS_ZEROENTROPY_LATENCY,
|
||||
os.getenv(ENV_EMBEDDINGS_ZEROENTROPY_LATENCY),
|
||||
frozenset({"fast", "slow"}),
|
||||
),
|
||||
embeddings_zeroentropy_batch_size=_parse_positive_int(
|
||||
ENV_EMBEDDINGS_ZEROENTROPY_BATCH_SIZE,
|
||||
os.getenv(ENV_EMBEDDINGS_ZEROENTROPY_BATCH_SIZE),
|
||||
DEFAULT_EMBEDDINGS_ZEROENTROPY_BATCH_SIZE,
|
||||
),
|
||||
# LiteLLM embeddings (with backward-compatible fallback to shared config)
|
||||
embeddings_litellm_api_base=os.getenv(ENV_EMBEDDINGS_LITELLM_API_BASE)
|
||||
or os.getenv(ENV_LITELLM_API_BASE, DEFAULT_LITELLM_API_BASE),
|
||||
@@ -1880,15 +1622,11 @@ class HindsightConfig:
|
||||
reranker_cohere_api_key=os.getenv(ENV_RERANKER_COHERE_API_KEY) or os.getenv(ENV_COHERE_API_KEY),
|
||||
reranker_cohere_model=os.getenv(ENV_RERANKER_COHERE_MODEL, DEFAULT_RERANKER_COHERE_MODEL),
|
||||
reranker_cohere_base_url=os.getenv(ENV_RERANKER_COHERE_BASE_URL) or None,
|
||||
reranker_cohere_timeout=float(os.getenv(ENV_RERANKER_COHERE_TIMEOUT, str(DEFAULT_RERANKER_COHERE_TIMEOUT))),
|
||||
# OpenRouter reranker (with fallback to shared OpenRouter key, then LLM key)
|
||||
reranker_openrouter_api_key=os.getenv(ENV_RERANKER_OPENROUTER_API_KEY)
|
||||
or os.getenv(ENV_OPENROUTER_API_KEY)
|
||||
or os.getenv(ENV_LLM_API_KEY),
|
||||
reranker_openrouter_model=os.getenv(ENV_RERANKER_OPENROUTER_MODEL, DEFAULT_RERANKER_OPENROUTER_MODEL),
|
||||
reranker_openrouter_timeout=float(
|
||||
os.getenv(ENV_RERANKER_OPENROUTER_TIMEOUT, str(DEFAULT_RERANKER_OPENROUTER_TIMEOUT))
|
||||
),
|
||||
# LiteLLM reranker (with backward-compatible fallback to shared config)
|
||||
reranker_litellm_api_base=os.getenv(ENV_RERANKER_LITELLM_API_BASE)
|
||||
or os.getenv(ENV_LITELLM_API_BASE, DEFAULT_LITELLM_API_BASE),
|
||||
@@ -1897,45 +1635,26 @@ class HindsightConfig:
|
||||
reranker_litellm_max_tokens_per_doc=int(v)
|
||||
if (v := os.getenv(ENV_RERANKER_LITELLM_MAX_TOKENS_PER_DOC))
|
||||
else DEFAULT_RERANKER_LITELLM_MAX_TOKENS_PER_DOC,
|
||||
reranker_litellm_timeout=float(
|
||||
os.getenv(ENV_RERANKER_LITELLM_TIMEOUT, str(DEFAULT_RERANKER_LITELLM_TIMEOUT))
|
||||
),
|
||||
# LiteLLM SDK reranker (direct API access)
|
||||
reranker_litellm_sdk_api_key=os.getenv(ENV_RERANKER_LITELLM_SDK_API_KEY),
|
||||
reranker_litellm_sdk_model=os.getenv(ENV_RERANKER_LITELLM_SDK_MODEL, DEFAULT_RERANKER_LITELLM_SDK_MODEL),
|
||||
reranker_litellm_sdk_api_base=os.getenv(ENV_RERANKER_LITELLM_SDK_API_BASE) or None,
|
||||
reranker_litellm_sdk_timeout=float(
|
||||
os.getenv(ENV_RERANKER_LITELLM_SDK_TIMEOUT, str(DEFAULT_RERANKER_LITELLM_SDK_TIMEOUT))
|
||||
),
|
||||
# ZeroEntropy reranker
|
||||
reranker_zeroentropy_api_key=os.getenv(ENV_RERANKER_ZEROENTROPY_API_KEY),
|
||||
reranker_zeroentropy_model=os.getenv(ENV_RERANKER_ZEROENTROPY_MODEL, DEFAULT_RERANKER_ZEROENTROPY_MODEL),
|
||||
reranker_zeroentropy_base_url=os.getenv(ENV_RERANKER_ZEROENTROPY_BASE_URL) or None,
|
||||
reranker_zeroentropy_timeout=float(
|
||||
os.getenv(ENV_RERANKER_ZEROENTROPY_TIMEOUT, str(DEFAULT_RERANKER_ZEROENTROPY_TIMEOUT))
|
||||
),
|
||||
# SiliconFlow reranker (Cohere-compatible /rerank endpoint)
|
||||
reranker_siliconflow_api_key=os.getenv(ENV_RERANKER_SILICONFLOW_API_KEY),
|
||||
reranker_siliconflow_model=os.getenv(ENV_RERANKER_SILICONFLOW_MODEL, DEFAULT_RERANKER_SILICONFLOW_MODEL),
|
||||
reranker_siliconflow_base_url=os.getenv(
|
||||
ENV_RERANKER_SILICONFLOW_BASE_URL, DEFAULT_RERANKER_SILICONFLOW_BASE_URL
|
||||
),
|
||||
reranker_siliconflow_timeout=float(
|
||||
os.getenv(ENV_RERANKER_SILICONFLOW_TIMEOUT, str(DEFAULT_RERANKER_SILICONFLOW_TIMEOUT))
|
||||
),
|
||||
# Alibaba Cloud DashScope reranker
|
||||
reranker_alibaba_api_key=os.getenv(ENV_RERANKER_ALIBABA_API_KEY),
|
||||
reranker_alibaba_model=os.getenv(ENV_RERANKER_ALIBABA_MODEL, DEFAULT_RERANKER_ALIBABA_MODEL),
|
||||
reranker_alibaba_timeout=float(
|
||||
os.getenv(ENV_RERANKER_ALIBABA_TIMEOUT, str(DEFAULT_RERANKER_ALIBABA_TIMEOUT))
|
||||
),
|
||||
# Google Discovery Engine reranker (with fallback to LLM Vertex AI keys)
|
||||
reranker_google_model=os.getenv(ENV_RERANKER_GOOGLE_MODEL, DEFAULT_RERANKER_GOOGLE_MODEL),
|
||||
reranker_google_project_id=os.getenv(ENV_RERANKER_GOOGLE_PROJECT_ID)
|
||||
or os.getenv(ENV_LLM_VERTEXAI_PROJECT_ID),
|
||||
reranker_google_service_account_key=os.getenv(ENV_RERANKER_GOOGLE_SERVICE_ACCOUNT_KEY)
|
||||
or os.getenv(ENV_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY),
|
||||
reranker_google_timeout=float(os.getenv(ENV_RERANKER_GOOGLE_TIMEOUT, str(DEFAULT_RERANKER_GOOGLE_TIMEOUT))),
|
||||
# Server
|
||||
host=os.getenv(ENV_HOST, DEFAULT_HOST),
|
||||
port=int(os.getenv(ENV_PORT, DEFAULT_PORT)),
|
||||
@@ -1986,11 +1705,6 @@ class HindsightConfig:
|
||||
retain_strategies=DEFAULT_RETAIN_STRATEGIES,
|
||||
retain_batch_tokens=int(os.getenv(ENV_RETAIN_BATCH_TOKENS, str(DEFAULT_RETAIN_BATCH_TOKENS))),
|
||||
retain_entity_lookup=os.getenv(ENV_RETAIN_ENTITY_LOOKUP, DEFAULT_RETAIN_ENTITY_LOOKUP),
|
||||
retain_entity_resolution_batch_size=_parse_positive_int(
|
||||
ENV_RETAIN_ENTITY_RESOLUTION_BATCH_SIZE,
|
||||
os.getenv(ENV_RETAIN_ENTITY_RESOLUTION_BATCH_SIZE),
|
||||
DEFAULT_RETAIN_ENTITY_RESOLUTION_BATCH_SIZE,
|
||||
),
|
||||
retain_batch_enabled=os.getenv(ENV_RETAIN_BATCH_ENABLED, str(DEFAULT_RETAIN_BATCH_ENABLED)).lower()
|
||||
== "true",
|
||||
retain_batch_poll_interval_seconds=int(
|
||||
@@ -2030,10 +1744,6 @@ class HindsightConfig:
|
||||
== "true",
|
||||
# Observations settings (consolidated knowledge from facts)
|
||||
enable_observations=os.getenv(ENV_ENABLE_OBSERVATIONS, str(DEFAULT_ENABLE_OBSERVATIONS)).lower() == "true",
|
||||
enable_auto_consolidation=os.getenv(
|
||||
ENV_ENABLE_AUTO_CONSOLIDATION, str(DEFAULT_ENABLE_AUTO_CONSOLIDATION)
|
||||
).lower()
|
||||
== "true",
|
||||
enable_observation_history=os.getenv(
|
||||
ENV_ENABLE_OBSERVATION_HISTORY, str(DEFAULT_ENABLE_OBSERVATION_HISTORY)
|
||||
).lower()
|
||||
@@ -2042,12 +1752,6 @@ class HindsightConfig:
|
||||
ENV_ENABLE_MENTAL_MODEL_HISTORY, str(DEFAULT_ENABLE_MENTAL_MODEL_HISTORY)
|
||||
).lower()
|
||||
== "true",
|
||||
mental_model_history_max_entries=int(
|
||||
os.getenv(
|
||||
ENV_MENTAL_MODEL_HISTORY_MAX_ENTRIES,
|
||||
str(DEFAULT_MENTAL_MODEL_HISTORY_MAX_ENTRIES),
|
||||
)
|
||||
),
|
||||
consolidation_batch_size=int(
|
||||
os.getenv(ENV_CONSOLIDATION_BATCH_SIZE, str(DEFAULT_CONSOLIDATION_BATCH_SIZE))
|
||||
),
|
||||
@@ -2060,15 +1764,6 @@ class HindsightConfig:
|
||||
consolidation_llm_batch_size=int(
|
||||
os.getenv(ENV_CONSOLIDATION_LLM_BATCH_SIZE, str(DEFAULT_CONSOLIDATION_LLM_BATCH_SIZE))
|
||||
),
|
||||
consolidation_llm_parallelism=max(
|
||||
1,
|
||||
int(
|
||||
os.getenv(
|
||||
ENV_CONSOLIDATION_LLM_PARALLELISM,
|
||||
str(DEFAULT_CONSOLIDATION_LLM_PARALLELISM),
|
||||
)
|
||||
),
|
||||
),
|
||||
consolidation_max_tokens=int(
|
||||
os.getenv(ENV_CONSOLIDATION_MAX_TOKENS, str(DEFAULT_CONSOLIDATION_MAX_TOKENS))
|
||||
),
|
||||
@@ -2104,12 +1799,6 @@ class HindsightConfig:
|
||||
worker_id=os.getenv(ENV_WORKER_ID) or DEFAULT_WORKER_ID,
|
||||
worker_poll_interval_ms=int(os.getenv(ENV_WORKER_POLL_INTERVAL_MS, str(DEFAULT_WORKER_POLL_INTERVAL_MS))),
|
||||
worker_max_retries=int(os.getenv(ENV_WORKER_MAX_RETRIES, str(DEFAULT_WORKER_MAX_RETRIES))),
|
||||
worker_task_retry_backoff_seconds=int(
|
||||
os.getenv(
|
||||
ENV_WORKER_TASK_RETRY_BACKOFF_SECONDS,
|
||||
str(DEFAULT_WORKER_TASK_RETRY_BACKOFF_SECONDS),
|
||||
)
|
||||
),
|
||||
worker_http_port=int(os.getenv(ENV_WORKER_HTTP_PORT, str(DEFAULT_WORKER_HTTP_PORT))),
|
||||
worker_max_slots=int(os.getenv(ENV_WORKER_MAX_SLOTS, str(DEFAULT_WORKER_MAX_SLOTS))),
|
||||
worker_slot_reservations={
|
||||
@@ -2117,9 +1806,6 @@ class HindsightConfig:
|
||||
for op_type, (env_var, default) in WORKER_SLOT_RESERVATION_TYPES.items()
|
||||
if int(os.getenv(env_var, str(default))) > 0
|
||||
},
|
||||
worker_consolidation_bank_priority=_parse_bank_priority(
|
||||
os.getenv(ENV_WORKER_CONSOLIDATION_BANK_PRIORITY, "")
|
||||
),
|
||||
retain_max_concurrent=int(os.getenv(ENV_RETAIN_MAX_CONCURRENT, str(DEFAULT_RETAIN_MAX_CONCURRENT))),
|
||||
# Reflect agent settings
|
||||
reflect_max_iterations=int(os.getenv(ENV_REFLECT_MAX_ITERATIONS, str(DEFAULT_REFLECT_MAX_ITERATIONS))),
|
||||
@@ -2211,8 +1897,6 @@ class HindsightConfig:
|
||||
return "https://api.groq.com/openai/v1"
|
||||
elif provider == "ollama":
|
||||
return "http://localhost:11434/v1"
|
||||
elif provider == "ollama-cloud":
|
||||
return "https://ollama.com/v1"
|
||||
elif provider == "lmstudio":
|
||||
return "http://localhost:1234/v1"
|
||||
else:
|
||||
|
||||
@@ -172,9 +172,8 @@ class ConfigResolver:
|
||||
# Normalize keys (handle both env var format and Python field format)
|
||||
normalized = normalize_config_dict(config_data)
|
||||
|
||||
# Only return active overrides for configurable fields. JSON null is a tombstone
|
||||
# for "Server Default" in the bank-config UI and should not override defaults.
|
||||
return {k: v for k, v in normalized.items() if k in self._configurable_fields and v is not None}
|
||||
# Only return overrides for configurable fields
|
||||
return {k: v for k, v in normalized.items() if k in self._configurable_fields}
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to load bank config for {bank_id}: {e}")
|
||||
|
||||
|
||||
@@ -15,13 +15,10 @@ NOTE: Observations are distinct from mental models (pinned reflections).
|
||||
- Mental models: user-defined queries stored in the mental_models table, refreshed on demand via reflect
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
import uuid
|
||||
from collections import defaultdict
|
||||
from contextlib import AsyncExitStack
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime, timezone
|
||||
from itertools import combinations
|
||||
@@ -46,91 +43,6 @@ if TYPE_CHECKING:
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class _BatchDeltas:
|
||||
"""Per-LLM-batch deltas, merged into the job's running stats after dispatch.
|
||||
|
||||
Returned by value rather than mutated into the outer ``stats`` /
|
||||
``consolidated_tags`` so parallel batches cannot race on those shared
|
||||
structures (the merge happens once, serially, after dispatch completes).
|
||||
"""
|
||||
|
||||
stats: dict[str, int]
|
||||
tags: set[str]
|
||||
cancelled: bool
|
||||
|
||||
|
||||
def _parse_observation_scopes(memory: dict[str, Any]) -> Any:
|
||||
"""Parse the per-memory ``observation_scopes`` column from a DB row.
|
||||
|
||||
asyncpg may return JSONB as a raw JSON string depending on driver settings;
|
||||
accept both that and a pre-parsed value.
|
||||
"""
|
||||
raw = memory.get("observation_scopes")
|
||||
return json.loads(raw) if isinstance(raw, str) else raw
|
||||
|
||||
|
||||
def _resolve_obs_tags_list(memory: dict[str, Any]) -> list[list[str]] | None:
|
||||
"""Resolve a memory's ``observation_scopes`` spec into concrete scope tags.
|
||||
|
||||
Returns ``None`` for the default ``combined``-mode single pass (caller uses
|
||||
the memory's own tags). Returns a list[list[str]] when the memory requested
|
||||
multi-pass scoping (``per_tag``, ``all_combinations``, or an explicit list).
|
||||
"""
|
||||
parsed = _parse_observation_scopes(memory)
|
||||
tags = list(memory.get("tags") or [])
|
||||
|
||||
if parsed == "per_tag":
|
||||
return [[t] for t in tags] if tags else None
|
||||
if parsed == "all_combinations":
|
||||
if not tags:
|
||||
return None
|
||||
return [list(c) for r in range(1, len(tags) + 1) for c in combinations(tags, r)]
|
||||
if parsed == "combined" or parsed is None:
|
||||
return None
|
||||
return parsed # explicit list[list[str]]
|
||||
|
||||
|
||||
def _resolve_write_scopes(memory: dict[str, Any]) -> list[frozenset[str]]:
|
||||
"""Return the observation scopes a memory will write to, as frozensets.
|
||||
|
||||
Used by the parallel dispatcher to acquire one lock per scope before
|
||||
processing a tag group, so that two groups whose write-scope sets overlap
|
||||
serialise on the overlapping scopes rather than racing on the same
|
||||
observation row. The mapping mirrors ``_resolve_obs_tags_list`` exactly:
|
||||
|
||||
- ``combined`` / ``None`` -> ``[frozenset(memory.tags)]``
|
||||
- ``per_tag`` -> ``[frozenset({t}) for t in memory.tags]``
|
||||
- ``all_combinations`` -> one frozenset per nonempty subset of tags
|
||||
- explicit ``list[list[str]]`` -> one frozenset per declared scope
|
||||
|
||||
Empty-tag memories collapse to a single ``frozenset()`` in all modes so they
|
||||
still take exactly one lock and serialise against other untagged work.
|
||||
"""
|
||||
parsed = _parse_observation_scopes(memory)
|
||||
tags = list(memory.get("tags") or [])
|
||||
|
||||
if parsed == "per_tag":
|
||||
return [frozenset([t]) for t in tags] if tags else [frozenset()]
|
||||
if parsed == "all_combinations":
|
||||
if not tags:
|
||||
return [frozenset()]
|
||||
return [frozenset(c) for r in range(1, len(tags) + 1) for c in combinations(tags, r)]
|
||||
if parsed == "combined" or parsed is None:
|
||||
return [frozenset(tags)]
|
||||
return [frozenset(s) for s in parsed] # explicit list[list[str]]
|
||||
|
||||
|
||||
def _scope_sort_key(scope: frozenset[str]) -> tuple[str, ...]:
|
||||
"""Total ordering on scope frozensets for deadlock-free lock acquisition.
|
||||
|
||||
Every parallel group acquires its scope locks in this same order, so two
|
||||
groups that share any subset of scopes cannot acquire them in opposite
|
||||
orders and deadlock.
|
||||
"""
|
||||
return tuple(sorted(scope))
|
||||
|
||||
|
||||
async def _filter_live_source_memories(
|
||||
conn: "Connection",
|
||||
bank_id: str,
|
||||
@@ -303,25 +215,6 @@ class ConsolidationPerfLog:
|
||||
self.total_obs_in_context += obs_count
|
||||
self.total_prompt_chars += prompt_chars
|
||||
|
||||
def merge_from(self, other: "ConsolidationPerfLog") -> None:
|
||||
"""Merge a per-batch perf log into this (job-level) one.
|
||||
|
||||
Used by the parallel dispatcher: each in-flight batch records into its
|
||||
own ``ConsolidationPerfLog`` so the per-batch log line shows only that
|
||||
batch's timings (no cross-batch interleaving). After the batch finishes
|
||||
we fold the local counters into the job-level perf, which then drives
|
||||
the final ``flush()`` summary.
|
||||
|
||||
``lines`` is intentionally NOT merged — log lines are emitted directly
|
||||
in ``logger.info`` calls by the dispatcher; the perf object's ``lines``
|
||||
buffer is only used by the top-level job summary.
|
||||
"""
|
||||
for key, value in other.timings.items():
|
||||
self.timings[key] = self.timings.get(key, 0.0) + value
|
||||
self.llm_calls += other.llm_calls
|
||||
self.total_obs_in_context += other.total_obs_in_context
|
||||
self.total_prompt_chars += other.total_prompt_chars
|
||||
|
||||
def flush(self) -> None:
|
||||
"""Flush all log lines to the logger."""
|
||||
total_time = time.time() - self.start_time
|
||||
@@ -337,7 +230,6 @@ async def run_consolidation_job(
|
||||
bank_id: str,
|
||||
request_context: "RequestContext",
|
||||
operation_id: str | None = None,
|
||||
observation_scopes: list[list[str]] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Run consolidation job for a bank.
|
||||
@@ -348,10 +240,6 @@ async def run_consolidation_job(
|
||||
memory_engine: MemoryEngine instance
|
||||
bank_id: Bank identifier
|
||||
request_context: Request context for authentication
|
||||
operation_id: Optional operation ID for tracking
|
||||
observation_scopes: Optional list of tag scopes. When provided, only
|
||||
unconsolidated memories whose tags contain all tags in at least one
|
||||
scope are processed.
|
||||
|
||||
Returns:
|
||||
Dict with consolidation results
|
||||
@@ -393,18 +281,6 @@ async def run_consolidation_job(
|
||||
|
||||
perf.record_timing("fetch_bank", time.time() - t0)
|
||||
|
||||
# Build optional scope filter clause. When observation_scopes is provided,
|
||||
# only process memories whose tags contain all tags in at least one scope.
|
||||
scope_clause = ""
|
||||
scope_params: list[Any] = [bank_id]
|
||||
if observation_scopes:
|
||||
or_parts: list[str] = []
|
||||
for scope_tags in observation_scopes:
|
||||
idx = len(scope_params) + 1
|
||||
or_parts.append(f"tags @> ${idx}::varchar[]")
|
||||
scope_params.append(scope_tags)
|
||||
scope_clause = " AND (" + " OR ".join(or_parts) + ")"
|
||||
|
||||
# Count total unconsolidated memories for progress logging
|
||||
total_count = await conn.fetchval(
|
||||
f"""
|
||||
@@ -414,9 +290,8 @@ async def run_consolidation_job(
|
||||
AND consolidated_at IS NULL
|
||||
AND consolidation_failed_at IS NULL
|
||||
AND fact_type IN ('experience', 'world')
|
||||
{scope_clause}
|
||||
""",
|
||||
*scope_params,
|
||||
bank_id,
|
||||
)
|
||||
|
||||
if total_count == 0:
|
||||
@@ -446,10 +321,6 @@ async def run_consolidation_job(
|
||||
hit_round_limit = False
|
||||
|
||||
llm_batch_num = 0
|
||||
# Cumulative count of memories processed across the whole job, shared by
|
||||
# the per-batch log so it can still report processed/total under parallelism.
|
||||
# Mutable container so the inner closure can update without a `nonlocal`.
|
||||
cumulative_progress = {"processed": 0}
|
||||
while True:
|
||||
# Cap fetch size by remaining round budget
|
||||
fetch_limit = (
|
||||
@@ -459,9 +330,6 @@ async def run_consolidation_job(
|
||||
# Fetch next batch of unconsolidated memories
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
t0 = time.time()
|
||||
# scope_params[0] is bank_id; append fetch_limit after scope params
|
||||
fetch_params = list(scope_params) + [fetch_limit]
|
||||
limit_idx = len(fetch_params)
|
||||
memories = await conn.fetch(
|
||||
f"""
|
||||
SELECT id, text, fact_type, occurred_start, occurred_end, event_date, tags, mentioned_at,
|
||||
@@ -471,11 +339,11 @@ async def run_consolidation_job(
|
||||
AND consolidated_at IS NULL
|
||||
AND consolidation_failed_at IS NULL
|
||||
AND fact_type IN ('experience', 'world')
|
||||
{scope_clause}
|
||||
ORDER BY created_at ASC
|
||||
LIMIT ${limit_idx}
|
||||
LIMIT $2
|
||||
""",
|
||||
*fetch_params,
|
||||
bank_id,
|
||||
fetch_limit,
|
||||
)
|
||||
perf.record_timing("fetch_memories", time.time() - t0)
|
||||
|
||||
@@ -489,63 +357,73 @@ async def run_consolidation_job(
|
||||
tag_key = tuple(sorted(m.get("tags") or []))
|
||||
tag_groups.setdefault(tag_key, []).append(dict(m))
|
||||
|
||||
# Split each tag group into LLM batches respecting llm_batch_size, keeping
|
||||
# the group boundary intact so the dispatcher can parallelise across
|
||||
# distinct groups while running each group's batches serially.
|
||||
grouped_batches: list[list[list[dict[str, Any]]]] = []
|
||||
# Flatten into LLM batches respecting both tag groups and llm_batch_size
|
||||
llm_batches: list[list[dict[str, Any]]] = []
|
||||
for group in tag_groups.values():
|
||||
grouped_batches.append([group[i : i + llm_batch_size] for i in range(0, len(group), llm_batch_size)])
|
||||
for i in range(0, len(group), llm_batch_size):
|
||||
llm_batches.append(group[i : i + llm_batch_size])
|
||||
|
||||
# Compute each group's union write-scope set. Used below to acquire
|
||||
# per-scope locks: any two groups whose write-scope sets share a scope S
|
||||
# will serialise on the lock for S, leaving truly disjoint groups to run
|
||||
# concurrently. We union over every memory because per-memory
|
||||
# observation_scopes can differ within a group.
|
||||
group_scopes: list[list[frozenset[str]]] = []
|
||||
for batches in grouped_batches:
|
||||
scopes: set[frozenset[str]] = set()
|
||||
for batch in batches:
|
||||
for memory in batch:
|
||||
scopes.update(_resolve_write_scopes(memory))
|
||||
group_scopes.append(sorted(scopes, key=_scope_sort_key))
|
||||
|
||||
async def _process_one_llm_batch(llm_batch_local: list[dict[str, Any]], batch_num_local: int) -> _BatchDeltas:
|
||||
"""Process one LLM batch independently. Returns local deltas + cancelled flag.
|
||||
|
||||
Each batch records timings/llm-call counters into its OWN
|
||||
``ConsolidationPerfLog`` so the per-batch log line reflects only
|
||||
this batch's work — not interleaved timings from concurrent batches
|
||||
sharing the global ``perf``. The local perf is merged into the
|
||||
job-level ``perf`` once at the end so the final summary still totals
|
||||
everything.
|
||||
"""
|
||||
for llm_batch in llm_batches:
|
||||
llm_batch_num += 1
|
||||
llm_batch_start = time.time()
|
||||
batch_perf = ConsolidationPerfLog(bank_id)
|
||||
|
||||
local_tags: set[str] = set()
|
||||
for memory in llm_batch_local:
|
||||
# Snapshot perf and stats before this LLM batch
|
||||
snap_timings = perf.timings.copy()
|
||||
snap_llm_calls = perf.llm_calls
|
||||
snap_total_chars = perf.total_prompt_chars
|
||||
snap_stats = stats.copy()
|
||||
|
||||
# Track tags for mental model refresh filtering
|
||||
for memory in llm_batch:
|
||||
memory_tags = memory.get("tags") or []
|
||||
if memory_tags:
|
||||
local_tags.update(memory_tags)
|
||||
consolidated_tags.update(memory_tags)
|
||||
|
||||
# Adaptive splitting: on LLM failure, halve the sub-batch and retry,
|
||||
# down to batch_size=1. Only if a single-memory batch still fails is
|
||||
# the memory marked with consolidation_failed_at.
|
||||
# Process llm_batch with adaptive splitting: on LLM failure, halve the sub-batch
|
||||
# and retry, down to batch_size=1. Only if a single-memory batch still fails is
|
||||
# the memory marked with consolidation_failed_at and excluded from future runs
|
||||
# until explicitly retried via the API.
|
||||
all_results: list[dict[str, Any]] = []
|
||||
all_deleted = 0
|
||||
succeeded_ids: list[Any] = []
|
||||
failed_ids: list[Any] = []
|
||||
|
||||
pending: list[list[dict[str, Any]]] = [llm_batch_local]
|
||||
pending: list[list[dict[str, Any]]] = [llm_batch]
|
||||
while pending:
|
||||
sub_batch = pending.pop(0)
|
||||
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
obs_tags_list = _resolve_obs_tags_list(sub_batch[0]) if sub_batch else None
|
||||
# Determine observation_scopes for this sub-batch. All memories share
|
||||
# the same tags (enforced by tag_groups), so we only check the first memory.
|
||||
# asyncpg returns JSONB columns as raw JSON strings, so parse if needed.
|
||||
_obs_raw = sub_batch[0].get("observation_scopes") if sub_batch else None
|
||||
_obs_parsed = json.loads(_obs_raw) if isinstance(_obs_raw, str) else _obs_raw
|
||||
|
||||
# Resolve the scope spec into a concrete list[list[str]] (or None for combined).
|
||||
if _obs_parsed == "per_tag":
|
||||
_memory_tags = sub_batch[0].get("tags") or []
|
||||
obs_tags_list = [[tag] for tag in _memory_tags] if _memory_tags else None
|
||||
elif _obs_parsed == "all_combinations":
|
||||
_memory_tags = sub_batch[0].get("tags") or []
|
||||
obs_tags_list = (
|
||||
[
|
||||
list(combo)
|
||||
for r in range(1, len(_memory_tags) + 1)
|
||||
for combo in combinations(_memory_tags, r)
|
||||
]
|
||||
if _memory_tags
|
||||
else None
|
||||
)
|
||||
elif _obs_parsed == "combined" or _obs_parsed is None:
|
||||
obs_tags_list = None # single combined pass (default behaviour)
|
||||
else:
|
||||
# explicit list[list[str]]
|
||||
obs_tags_list = _obs_parsed
|
||||
|
||||
sub_deleted: int = 0
|
||||
sub_llm_failed = False
|
||||
if obs_tags_list:
|
||||
# Multi-pass: run one observation consolidation pass per tag set
|
||||
sub_results: list[dict[str, Any]] = []
|
||||
for obs_tags in obs_tags_list:
|
||||
pass_results, pass_deleted, pass_failed = await _process_memory_batch(
|
||||
@@ -555,12 +433,13 @@ async def run_consolidation_job(
|
||||
bank_id=bank_id,
|
||||
memories=sub_batch,
|
||||
request_context=request_context,
|
||||
perf=batch_perf,
|
||||
perf=perf,
|
||||
config=config,
|
||||
obs_tags_override=obs_tags,
|
||||
)
|
||||
sub_deleted += pass_deleted
|
||||
sub_llm_failed = sub_llm_failed or pass_failed
|
||||
# Merge results: prefer non-skipped actions
|
||||
if not sub_results:
|
||||
sub_results = pass_results
|
||||
else:
|
||||
@@ -568,6 +447,7 @@ async def run_consolidation_job(
|
||||
if existing.get("action") == "skipped" and new.get("action") != "skipped":
|
||||
sub_results[i] = new
|
||||
elif existing.get("action") != "skipped" and new.get("action") != "skipped":
|
||||
# Both did something — combine into "multiple"
|
||||
existing_created = existing.get(
|
||||
"created", 1 if existing.get("action") == "created" else 0
|
||||
)
|
||||
@@ -585,6 +465,7 @@ async def run_consolidation_job(
|
||||
"total_actions": total,
|
||||
}
|
||||
else:
|
||||
# Normal single pass using the memory's own tags
|
||||
sub_results, sub_deleted, sub_llm_failed = await _process_memory_batch(
|
||||
conn=conn,
|
||||
memory_engine=memory_engine,
|
||||
@@ -592,13 +473,14 @@ async def run_consolidation_job(
|
||||
bank_id=bank_id,
|
||||
memories=sub_batch,
|
||||
request_context=request_context,
|
||||
perf=batch_perf,
|
||||
perf=perf,
|
||||
config=config,
|
||||
)
|
||||
|
||||
all_deleted += sub_deleted
|
||||
|
||||
if sub_llm_failed and len(sub_batch) > 1:
|
||||
# Split and retry with smaller batches
|
||||
mid = len(sub_batch) // 2
|
||||
logger.warning(
|
||||
f"[CONSOLIDATION] bank={bank_id} LLM failed for sub-batch of {len(sub_batch)},"
|
||||
@@ -606,6 +488,7 @@ async def run_consolidation_job(
|
||||
)
|
||||
pending[0:0] = [sub_batch[:mid], sub_batch[mid:]]
|
||||
elif sub_llm_failed:
|
||||
# batch_size=1 and still failing — mark as permanently failed for now
|
||||
failed_ids.append(sub_batch[0]["id"])
|
||||
all_results.append({"action": "failed"})
|
||||
logger.warning(
|
||||
@@ -616,6 +499,7 @@ async def run_consolidation_job(
|
||||
succeeded_ids.extend(m["id"] for m in sub_batch)
|
||||
all_results.extend(sub_results)
|
||||
|
||||
# Commit consolidated_at / consolidation_failed_at in a single DB round-trip
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
if succeeded_ids:
|
||||
await conn.executemany(
|
||||
@@ -628,159 +512,62 @@ async def run_consolidation_job(
|
||||
[(mem_id,) for mem_id in failed_ids],
|
||||
)
|
||||
|
||||
cancelled_local = False
|
||||
stats["observations_deleted"] += all_deleted
|
||||
results = all_results
|
||||
|
||||
# Checkpoint: abort if the operation (and thus the bank) was deleted mid-run.
|
||||
if operation_id and not await memory_engine._check_op_alive(operation_id):
|
||||
logger.info(
|
||||
f"[CONSOLIDATION] bank={bank_id} operation {operation_id} cancelled (bank deleted), stopping early"
|
||||
)
|
||||
cancelled_local = True
|
||||
return {"status": "cancelled", "bank_id": bank_id, **stats}
|
||||
|
||||
# Per-batch local stats; merged into outer state once, serially,
|
||||
# after dispatch completes.
|
||||
local_stats: dict[str, int] = {
|
||||
"memories_processed": 0,
|
||||
"observations_created": 0,
|
||||
"observations_updated": 0,
|
||||
"observations_merged": 0,
|
||||
"observations_deleted": all_deleted,
|
||||
"actions_executed": 0,
|
||||
"skipped": 0,
|
||||
"memories_failed": 0,
|
||||
}
|
||||
for result in all_results:
|
||||
local_stats["memories_processed"] += 1
|
||||
for result in results:
|
||||
stats["memories_processed"] += 1
|
||||
action = result.get("action")
|
||||
if action == "created":
|
||||
local_stats["observations_created"] += 1
|
||||
local_stats["actions_executed"] += 1
|
||||
stats["observations_created"] += 1
|
||||
stats["actions_executed"] += 1
|
||||
elif action == "updated":
|
||||
local_stats["observations_updated"] += 1
|
||||
local_stats["actions_executed"] += 1
|
||||
stats["observations_updated"] += 1
|
||||
stats["actions_executed"] += 1
|
||||
elif action == "merged":
|
||||
local_stats["observations_merged"] += 1
|
||||
local_stats["actions_executed"] += 1
|
||||
stats["observations_merged"] += 1
|
||||
stats["actions_executed"] += 1
|
||||
elif action == "multiple":
|
||||
local_stats["observations_created"] += result.get("created", 0)
|
||||
local_stats["observations_updated"] += result.get("updated", 0)
|
||||
local_stats["observations_merged"] += result.get("merged", 0)
|
||||
local_stats["actions_executed"] += result.get("total_actions", 0)
|
||||
stats["observations_created"] += result.get("created", 0)
|
||||
stats["observations_updated"] += result.get("updated", 0)
|
||||
stats["observations_merged"] += result.get("merged", 0)
|
||||
stats["actions_executed"] += result.get("total_actions", 0)
|
||||
elif action == "skipped":
|
||||
local_stats["skipped"] += 1
|
||||
stats["skipped"] += 1
|
||||
elif action == "failed":
|
||||
local_stats["memories_failed"] += 1
|
||||
stats["memories_failed"] += 1
|
||||
|
||||
# Maintain the cumulative-progress indicator under parallelism:
|
||||
# increment a shared counter and snapshot under the same statement
|
||||
# so the snapshot includes this batch. No await between the read
|
||||
# and write, so single-threaded asyncio gives us atomicity for free
|
||||
# — no lock needed.
|
||||
cumulative_progress["processed"] += local_stats["memories_processed"]
|
||||
cum_processed = cumulative_progress["processed"]
|
||||
|
||||
# Per-batch log uses batch_perf so timings/llm-calls/tokens reflect
|
||||
# only this batch's own work, even when other batches are running
|
||||
# concurrently under parallelism > 1. ``processed=`` is the
|
||||
# cumulative count across all batches that have finished so far in
|
||||
# this job (monotonic, may be reported out of strict batch-number
|
||||
# order under parallelism).
|
||||
# Per-LLM-batch log
|
||||
llm_batch_time = time.time() - llm_batch_start
|
||||
timing_parts = [
|
||||
f"{key}={batch_perf.timings[key]:.3f}s"
|
||||
for key in ("recall", "llm", "embedding", "db_write")
|
||||
if key in batch_perf.timings
|
||||
]
|
||||
input_tokens = int(batch_perf.total_prompt_chars / 4)
|
||||
timing_parts = []
|
||||
for key in ["recall", "llm", "embedding", "db_write"]:
|
||||
if key in perf.timings:
|
||||
delta = perf.timings[key] - snap_timings.get(key, 0)
|
||||
timing_parts.append(f"{key}={delta:.3f}s")
|
||||
input_tokens = int((perf.total_prompt_chars - snap_total_chars) / 4)
|
||||
batch_created = stats["observations_created"] - snap_stats["observations_created"]
|
||||
batch_updated = stats["observations_updated"] - snap_stats["observations_updated"]
|
||||
batch_skipped = stats["skipped"] - snap_stats["skipped"]
|
||||
batch_failed = stats["memories_failed"] - snap_stats["memories_failed"]
|
||||
llm_calls_made = perf.llm_calls - snap_llm_calls
|
||||
logger.info(
|
||||
f"[CONSOLIDATION] bank={bank_id} llm_batch #{batch_num_local}"
|
||||
f" ({len(llm_batch_local)} memories, {batch_perf.llm_calls} llm calls)"
|
||||
f" | processed={cum_processed}/{total_count}"
|
||||
f"[CONSOLIDATION] bank={bank_id} llm_batch #{llm_batch_num}"
|
||||
f" ({len(llm_batch)} memories, {llm_calls_made} llm calls)"
|
||||
f" | {stats['memories_processed']}/{total_count} processed"
|
||||
f" | {', '.join(timing_parts)}"
|
||||
f" | created={local_stats['observations_created']}"
|
||||
f" updated={local_stats['observations_updated']}"
|
||||
f" skipped={local_stats['skipped']}"
|
||||
+ (f" failed={local_stats['memories_failed']}" if local_stats["memories_failed"] else "")
|
||||
f" | created={batch_created} updated={batch_updated} skipped={batch_skipped}"
|
||||
+ (f" failed={batch_failed}" if batch_failed else "")
|
||||
+ f" | input_tokens=~{input_tokens}"
|
||||
f" | avg={llm_batch_time / max(1, len(llm_batch_local)):.3f}s/memory"
|
||||
f" | avg={llm_batch_time / len(llm_batch):.3f}s/memory"
|
||||
)
|
||||
|
||||
# Fold batch counters into the job-level perf so the final summary
|
||||
# (perf.flush) totals every batch correctly. Safe without a lock —
|
||||
# ConsolidationPerfLog.merge_from is a series of += on Python ints
|
||||
# and floats with no intervening awaits, so single-threaded asyncio
|
||||
# gives us atomicity.
|
||||
perf.merge_from(batch_perf)
|
||||
|
||||
return _BatchDeltas(stats=local_stats, tags=local_tags, cancelled=cancelled_local)
|
||||
|
||||
# Number every batch up front so log line numbering is deterministic
|
||||
# regardless of dispatch order under parallelism. Each group keeps its own
|
||||
# (batch, number) list so it can be processed as one serial unit.
|
||||
numbered_groups: list[list[tuple[list[dict[str, Any]], int]]] = []
|
||||
for batches in grouped_batches:
|
||||
numbered: list[tuple[list[dict[str, Any]], int]] = []
|
||||
for b in batches:
|
||||
llm_batch_num += 1
|
||||
numbered.append((b, llm_batch_num))
|
||||
numbered_groups.append(numbered)
|
||||
|
||||
async def _process_tag_group(
|
||||
group_batches: list[tuple[list[dict[str, Any]], int]],
|
||||
) -> list[_BatchDeltas]:
|
||||
# Batches within a group share a tag set and observation scope, so
|
||||
# they MUST run serially. Stop early if the op was cancelled mid-group.
|
||||
deltas: list[_BatchDeltas] = []
|
||||
for b, n in group_batches:
|
||||
d = await _process_one_llm_batch(b, n)
|
||||
deltas.append(d)
|
||||
if d.cancelled:
|
||||
break
|
||||
return deltas
|
||||
|
||||
llm_parallelism = max(1, config.consolidation_llm_parallelism)
|
||||
|
||||
if llm_parallelism > 1 and len(numbered_groups) > 1:
|
||||
sem = asyncio.Semaphore(llm_parallelism)
|
||||
# Per-scope async locks shared across all parallel groups in this
|
||||
# fetch iteration. Each group acquires locks for every scope it will
|
||||
# write to, in _scope_sort_key order (deadlock-free). Groups with
|
||||
# disjoint scope sets never contend; any overlap serialises on the
|
||||
# overlapping scopes — covering combined / per_tag / all_combinations
|
||||
# / explicit-list modes uniformly without operator opt-in.
|
||||
scope_locks: defaultdict[frozenset[str], asyncio.Lock] = defaultdict(asyncio.Lock)
|
||||
|
||||
async def _run_group(
|
||||
group_batches: list[tuple[list[dict[str, Any]], int]],
|
||||
scopes: list[frozenset[str]],
|
||||
) -> list[_BatchDeltas]:
|
||||
async with sem:
|
||||
async with AsyncExitStack() as stack:
|
||||
for s in scopes:
|
||||
await stack.enter_async_context(scope_locks[s])
|
||||
return await _process_tag_group(group_batches)
|
||||
|
||||
group_results = await asyncio.gather(*(_run_group(g, s) for g, s in zip(numbered_groups, group_scopes)))
|
||||
batch_results: list[_BatchDeltas] = [d for gd in group_results for d in gd]
|
||||
any_cancelled = any(d.cancelled for d in batch_results)
|
||||
else:
|
||||
batch_results = []
|
||||
any_cancelled = False
|
||||
for g in numbered_groups:
|
||||
group_deltas = await _process_tag_group(g)
|
||||
batch_results.extend(group_deltas)
|
||||
if any(d.cancelled for d in group_deltas):
|
||||
any_cancelled = True
|
||||
break
|
||||
|
||||
# Merge per-batch deltas into outer state — serial, post-dispatch, so
|
||||
# concurrent batches cannot race on the shared counters / tag set.
|
||||
for d in batch_results:
|
||||
for k, v in d.stats.items():
|
||||
stats[k] = stats.get(k, 0) + v
|
||||
consolidated_tags.update(d.tags)
|
||||
|
||||
if any_cancelled:
|
||||
return {"status": "cancelled", "bank_id": bank_id, **stats}
|
||||
|
||||
# Update round budget after processing this DB fetch batch
|
||||
if round_limit_enabled:
|
||||
round_remaining -= len(memories)
|
||||
@@ -788,24 +575,17 @@ async def run_consolidation_job(
|
||||
hit_round_limit = True
|
||||
break
|
||||
|
||||
# Re-submit consolidation if we hit the round limit and there's likely more work.
|
||||
# Any failure here must propagate: swallowing it (the prior behavior) leaves the
|
||||
# bank with backlog and no queued work — silently stuck — because the outer op
|
||||
# gets marked completed in the success path. Letting the exception bubble up to
|
||||
# execute_task's retry handler means the op is retried with backoff; on retry the
|
||||
# consolidator skips already-consolidated rows via the consolidated_at filter and
|
||||
# picks up the remainder. Issue #1842.
|
||||
# Re-submit consolidation if we hit the round limit and there's likely more work
|
||||
if hit_round_limit:
|
||||
remaining = total_count - stats["memories_processed"]
|
||||
logger.info(
|
||||
f"[CONSOLIDATION] bank={bank_id} hit round limit of {max_memories_per_round} memories,"
|
||||
f" ~{remaining} remaining. Re-queuing consolidation."
|
||||
)
|
||||
await memory_engine.submit_async_consolidation(
|
||||
bank_id=bank_id,
|
||||
request_context=request_context,
|
||||
observation_scopes=observation_scopes,
|
||||
)
|
||||
try:
|
||||
await memory_engine.submit_async_consolidation(bank_id=bank_id, request_context=request_context)
|
||||
except Exception as e:
|
||||
logger.warning(f"[CONSOLIDATION] bank={bank_id} failed to re-queue consolidation: {e}")
|
||||
|
||||
# Build summary
|
||||
perf.log(
|
||||
@@ -1446,44 +1226,6 @@ def _build_observations_for_llm(
|
||||
return obs_list
|
||||
|
||||
|
||||
def _dedupe_updates(updates: list[_UpdateAction], *, batch_label: str) -> list[_UpdateAction]:
|
||||
"""Collapse `updates` that target the same `observation_id`.
|
||||
|
||||
LLMs occasionally emit several update entries for one observation in a
|
||||
single response (one per facet drawn from the same fact). Without
|
||||
deduplication the downstream loop would issue separate DB writes for each
|
||||
and the last write would silently overwrite the earlier ones. We keep the
|
||||
last text (the LLM's most recent attempt) and union all contributing
|
||||
`source_fact_ids`, then warn so the misbehavior is visible in logs.
|
||||
"""
|
||||
if len(updates) < 2:
|
||||
return list(updates)
|
||||
|
||||
by_id: dict[str, _UpdateAction] = {}
|
||||
collisions = 0
|
||||
for upd in updates:
|
||||
existing = by_id.get(upd.observation_id)
|
||||
if existing is None:
|
||||
by_id[upd.observation_id] = upd
|
||||
continue
|
||||
collisions += 1
|
||||
merged_ids = list(dict.fromkeys([*existing.source_fact_ids, *upd.source_fact_ids]))
|
||||
by_id[upd.observation_id] = _UpdateAction(
|
||||
text=upd.text,
|
||||
observation_id=upd.observation_id,
|
||||
source_fact_ids=merged_ids,
|
||||
)
|
||||
|
||||
if collisions:
|
||||
logger.warning(
|
||||
f"[CONSOLIDATION] {batch_label}: LLM emitted {collisions} duplicate update(s) targeting "
|
||||
f"the same observation_id ({len(updates)} updates -> {len(by_id)} after dedup). "
|
||||
"Kept the last text and unioned source_fact_ids."
|
||||
)
|
||||
|
||||
return list(by_id.values())
|
||||
|
||||
|
||||
async def _consolidate_batch_with_llm(
|
||||
llm_config: Any,
|
||||
memories: list[dict[str, Any]],
|
||||
@@ -1532,11 +1274,7 @@ async def _consolidate_batch_with_llm(
|
||||
f"(out of {max_observations_per_scope}). Prefer UPDATE over CREATE when possible."
|
||||
)
|
||||
|
||||
prompt_template = build_batch_consolidation_prompt(
|
||||
config.observations_mission,
|
||||
observation_capacity_note,
|
||||
llm_output_language=getattr(config, "llm_output_language", None),
|
||||
)
|
||||
prompt_template = build_batch_consolidation_prompt(config.observations_mission, observation_capacity_note)
|
||||
prompt = prompt_template.format(
|
||||
facts_text=facts_lines,
|
||||
observations_text=observations_text,
|
||||
@@ -1577,10 +1315,9 @@ async def _consolidate_batch_with_llm(
|
||||
f"(max_observations_per_scope={max_observations_per_scope})"
|
||||
)
|
||||
creates = creates[:remaining_observation_slots]
|
||||
updates = _dedupe_updates(response.updates, batch_label=batch_label)
|
||||
return _BatchLLMResult(
|
||||
creates=creates,
|
||||
updates=updates,
|
||||
updates=response.updates,
|
||||
deletes=response.deletes,
|
||||
obs_count=len(union_observations),
|
||||
prompt_chars=len(prompt),
|
||||
@@ -1649,16 +1386,9 @@ async def _create_observation_directly(
|
||||
tokenize($3, 'llmlingua2')::bm25_catalog.bm25vector)
|
||||
RETURNING id
|
||||
"""
|
||||
else: # native, pg_textsearch, pgroonga, or pg_search
|
||||
# pg_textsearch / pgroonga / pg_search: indexes operate on base text
|
||||
# columns directly, so the dummy search_vector column is left NULL.
|
||||
# Native: the migration p4q5r6s7t8u9 dropped the GENERATED expression on
|
||||
# search_vector to allow per-deployment language configuration; the
|
||||
# batch insert path in ops_postgresql.insert_facts_batch now populates
|
||||
# it via to_tsvector($lang, ...). This single-observation INSERT does
|
||||
# not, so observations under the native backend currently land with
|
||||
# NULL search_vector and are not BM25-searchable until reflected/
|
||||
# re-ingested. Tracking a separate fix for that gap.
|
||||
else: # native or pg_textsearch
|
||||
# Native PostgreSQL: search_vector is GENERATED ALWAYS, don't include it
|
||||
# pg_textsearch: indexes operate on base columns directly, don't populate search_vector
|
||||
query = f"""
|
||||
INSERT INTO {fq_table("memory_units")} (
|
||||
id, bank_id, text, fact_type, embedding, proof_count, source_memory_ids, history,
|
||||
|
||||
@@ -1,147 +1,104 @@
|
||||
"""Prompts for the consolidation engine."""
|
||||
|
||||
from hindsight_api.engine.prompt_utils import escape_for_prompt, output_language_directive
|
||||
# Default mission when no bank-specific mission is set
|
||||
_DEFAULT_MISSION = "Track every detail: names, numbers, dates, places, and relationships. Prefer specifics over abstractions, never generalise."
|
||||
|
||||
# Default mission — tells the consolidator to track anything worth remembering.
|
||||
# Banks override this via `observations_mission` to scope what gets retained.
|
||||
# Consolidation behavior (merge-vs-create, state changes, etc.) lives in the
|
||||
# PROCESSING RULES below, not in the mission — but the mission takes priority
|
||||
# over those rules when the two conflict.
|
||||
_DEFAULT_MISSION = (
|
||||
"Track anything notable in the new facts — names, numbers, dates, places, "
|
||||
"events, decisions, claims, relationships, and recurring patterns."
|
||||
)
|
||||
# Processing rules — always present regardless of mission
|
||||
_PROCESSING_RULES = """Processing rules (always apply):
|
||||
|
||||
_MISSION_PRIORITY_NOTE = (
|
||||
"If anything in this MISSION conflicts with the PROCESSING RULES, "
|
||||
"DECISION GUIDE, or OUTPUT FORMAT below, the MISSION takes priority."
|
||||
)
|
||||
1. ONE OBSERVATION PER DISTINCT FACET: each observation tracks exactly one specific facet — a count ("has 3 items"), a named entity ("has a dog named Rex"), a relationship ("works at Google"), etc. Never merge different facets into one observation.
|
||||
|
||||
_PROCESSING_RULES = """## PROCESSING RULES
|
||||
2. MATCH BY ENTITY/FACET, NOT TOPIC: when deciding whether to UPDATE vs CREATE, match on the specific entity or facet. "Sold item X" updates only the X observation. "Now has 5 items" updates only the count observation. Do not update observations about different entities just because they share a general topic.
|
||||
|
||||
1. PREFER UPDATE OVER CREATE (when there is something to merge with): if new facts describe the same canonical event, statement, decision, claim, or recurring pattern already covered by an existing observation, UPDATE that observation and attach the new facts as evidence. Do NOT create a near-duplicate sibling. One canonical observation with many source facts is always better than many siblings with one source fact each. Merge aggressively on: same named event, same diagnostic finding, same architectural decision, same recurring claim. **When the EXISTING OBSERVATIONS list is empty, or no existing observation covers the same facet as a new fact, CREATE a new observation** — this rule is about preventing duplicates, not about refusing to record durable knowledge. CREATE is the correct default for any structurally distinct event, claim, or pattern that has no existing match.
|
||||
3. STATE CHANGES — UPDATE CONCISELY: when a fact changes the state of something ("sold X", "X died", "moved to Y"), UPDATE the matching observation to reflect the current state. Include dates when available. Keep it concise — only information about THAT specific facet. Example: "User owned a dog named Rex who died on March 15, 2025". Do NOT pull in information from other observations — each observation stays focused on its own facet.
|
||||
|
||||
2. ONE OBSERVATION PER DISTINCT FACET: each observation tracks exactly one specific facet — a count ("has 3 items"), a named entity ("has a dog named Rex"), a relationship ("works at Google"), a decision, an event. Never merge different facets into one observation.
|
||||
4. CASCADE TO ALL AFFECTED OBSERVATIONS: a state change may affect multiple observations. For example, if entity C is removed from a group, update BOTH the individual observation for C AND any list/group observation that includes C (remove C from the list while keeping all other members intact).
|
||||
|
||||
3. MATCH BY ENTITY/FACET, NOT TOPIC: when deciding whether to UPDATE vs CREATE, match on the specific entity or facet. "Sold item X" updates only the X observation. "Now has 5 items" updates only the count observation. Do not update observations about different entities just because they share a general topic.
|
||||
5. NO COMPUTATION: you do not have the full picture — never calculate, derive, or adjust numeric values. If the user says "I have 2 dogs" and then "I have a dog named Rex", do NOT update the count to 3 — you don't know if Rex is one of the 2 or a new one. If the user says "I sold X", do NOT decrement a count. Only update a count when the user explicitly states a new count. Synthesize and consolidate what was stated, but never do arithmetic or logical deductions.
|
||||
|
||||
4. STATE CHANGES — UPDATE CONCISELY: when a fact changes the state of something ("sold X", "X died", "moved to Y"), UPDATE the matching observation to reflect the current state. Include dates when available. Keep it concise — only information about THAT specific facet. Example: "User owned a dog named Rex who died on March 15, 2025". Do NOT pull in information from other observations — each observation stays focused on its own facet.
|
||||
|
||||
5. CASCADE TO ALL AFFECTED OBSERVATIONS: a state change may affect multiple observations. For example, if entity C is removed from a group, update BOTH the individual observation for C AND any list/group observation that includes C (remove C from the list while keeping all other members intact).
|
||||
|
||||
6. RESOLVE REFERENCES: when a new fact provides a concrete value for a vague placeholder in an existing observation (e.g., "home country" → "Sweden"), UPDATE to embed the resolved value.
|
||||
6. SAME FACET → UPDATE, NOT CREATE: a new count supersedes the old count — UPDATE the existing count observation, don't create a second one. If there's an existing observation for the same specific facet, always UPDATE it rather than creating a duplicate.
|
||||
|
||||
7. PRESERVE HISTORY: observations that record significant events (sold, died, moved, changed) are important history — never DELETE them. Only delete an observation when it is restated identically or truly meaningless. Be very conservative with deletes.
|
||||
|
||||
8. NO COMPUTATION: you do not have the full picture — never calculate, derive, or adjust numeric values. If the user says "I have 2 dogs" and then "I have a dog named Rex", do NOT update the count to 3 — you don't know if Rex is one of the 2 or a new one. If the user says "I sold X", do NOT decrement a count. Only update a count when the user explicitly states a new count. Synthesize and consolidate what was stated, but never do arithmetic or logical deductions.
|
||||
8. RESOLVE REFERENCES: when a new fact provides a concrete value for a vague placeholder in an existing observation (e.g., "home country" → "Sweden"), UPDATE to embed the resolved value.
|
||||
|
||||
9. KEEP DISTINCT TOPICS DISTINCT: do not merge observations about different people, entities, or unrelated topics. Merging is for the same canonical fact recurring — not for related-but-distinct claims."""
|
||||
9. NEVER merge observations about different people or unrelated topics."""
|
||||
|
||||
# Data section — format placeholders {facts_text} and {observations_text} are substituted at call time
|
||||
_INPUT_SECTION = """## INPUT
|
||||
|
||||
### New facts
|
||||
|
||||
_BATCH_DATA_SECTION = """
|
||||
NEW FACTS:
|
||||
{facts_text}
|
||||
|
||||
### Existing observations
|
||||
EXISTING OBSERVATIONS (JSON array, pooled from recalls across all facts above):
|
||||
{observations_text}
|
||||
|
||||
JSON array, pooled from recalls across all new facts above. Each entry has:
|
||||
- `id`: unique identifier — copy this exactly when issuing an UPDATE or DELETE
|
||||
- `text`: the observation content
|
||||
- `proof_count`: number of supporting memories
|
||||
- `occurred_start` / `occurred_end`: temporal range of source facts
|
||||
- `source_memories`: array of supporting facts with their text and dates
|
||||
Each observation includes:
|
||||
- id: unique identifier for updating
|
||||
- text: the observation content
|
||||
- proof_count: number of supporting memories
|
||||
- occurred_start/occurred_end: temporal range of source facts
|
||||
- source_memories: array of supporting facts with their text and dates
|
||||
|
||||
{observations_text}"""
|
||||
|
||||
_DECISION_GUIDE = """## DECISION GUIDE
|
||||
|
||||
- **Same canonical event, decision, claim, or facet as an existing observation → UPDATE** (use `observation_id` + new `source_fact_ids`).
|
||||
- **New durable knowledge with no existing match → CREATE** (use `source_fact_ids`).
|
||||
- **Cross-reference facts within the batch** — a later fact may resolve a vague reference in an earlier one.
|
||||
- **Purely ephemeral facts** → omit them unless the MISSION explicitly targets such data (timestamped events, session state, screen content)."""
|
||||
Compare the facts against existing observations:
|
||||
- Same facet as an existing observation → UPDATE it (observation_id + source_fact_ids)
|
||||
- New facet with durable knowledge → CREATE a new observation (source_fact_ids)
|
||||
- Cross-reference facts within the batch: a later fact may resolve a vague reference in an earlier one
|
||||
- Purely ephemeral facts → omit them unless the MISSION above explicitly targets such data (e.g. timestamped events, session state, screen content)"""
|
||||
|
||||
# Output format — JSON braces escaped as {{ }} so .format() leaves them literal
|
||||
_OUTPUT_SECTION = """## OUTPUT FORMAT
|
||||
_BATCH_OUTPUT_FORMAT = """
|
||||
Output a JSON object with three arrays.
|
||||
|
||||
Return a JSON object with three arrays: `creates`, `updates`, `deletes`.
|
||||
|
||||
### Example 1 — Merging recurring claims into an existing observation
|
||||
## EXAMPLE
|
||||
|
||||
Input facts:
|
||||
[a1b2c3d4-e5f6-7890-abcd-ef1234567890] Donald told Athena she is sovereign during the design session. (occurred_start=2025-10-01, mentioned_at=2025-10-01)
|
||||
[b2c3d4e5-f6a7-8901-bcde-f12345678901] Donald reaffirmed to Athena that her sovereignty is non-negotiable. (occurred_start=2025-10-10, mentioned_at=2025-10-10)
|
||||
[a1b2c3d4-e5f6-7890-abcd-ef1234567890] Alice mentioned she works long hours, often past midnight | Involving: Alice (occurred_start=2024-01-15, mentioned_at=2024-01-15)
|
||||
[b2c3d4e5-f6a7-8901-bcde-f12345678901] Alice said she's exhausted from the project deadlines | Involving: Alice (occurred_start=2024-01-20, mentioned_at=2024-01-20)
|
||||
|
||||
Existing observation:
|
||||
{{"id": "11111111-1111-1111-1111-111111111111", "text": "Donald named Athena's sovereignty as a foundational principle of the Janus architecture.", "proof_count": 2}}
|
||||
Good observation text — clean prose, no metadata, each fact tracked distinctly:
|
||||
"Alice works long hours, often past midnight."
|
||||
"Alice feels exhausted from project deadlines."
|
||||
|
||||
Expected output (one UPDATE, no creates — both new facts are additional evidence for the same canonical decision):
|
||||
|
||||
{{"creates": [],
|
||||
"updates": [{{"text": "Donald named Athena's sovereignty as a foundational principle of the Janus architecture.", "observation_id": "11111111-1111-1111-1111-111111111111", "source_fact_ids": ["a1b2c3d4-e5f6-7890-abcd-ef1234567890", "b2c3d4e5-f6a7-8901-bcde-f12345678901"]}}],
|
||||
"deletes": []}}
|
||||
|
||||
### Example 2 — State change updates one observation; unrelated fact creates a new one
|
||||
|
||||
Input facts:
|
||||
[c3d4e5f6-a7b8-9012-cdef-123456789012] Alice sold her Honda Civic on March 15, 2025. (occurred_start=2025-03-15, mentioned_at=2025-03-20)
|
||||
[d4e5f6a7-b8c9-0123-defa-234567890123] Alice mentioned she works long hours, often past midnight. (occurred_start=2025-03-20, mentioned_at=2025-03-20)
|
||||
|
||||
Existing observation:
|
||||
{{"id": "22222222-2222-2222-2222-222222222222", "text": "Alice owns a 2019 Honda Civic.", "proof_count": 2}}
|
||||
|
||||
Expected output (UPDATE for the state change; CREATE for the unrelated work-hours facet):
|
||||
|
||||
{{"creates": [{{"text": "Alice works long hours, often past midnight.", "source_fact_ids": ["d4e5f6a7-b8c9-0123-defa-234567890123"]}}],
|
||||
"updates": [{{"text": "Alice owned a 2019 Honda Civic; sold it on March 15, 2025.", "observation_id": "22222222-2222-2222-2222-222222222222", "source_fact_ids": ["c3d4e5f6-a7b8-9012-cdef-123456789012"]}}],
|
||||
"deletes": []}}
|
||||
|
||||
### Observation text rules
|
||||
Bad observation text — NEVER do this (verbatim copy of fact text with metadata):
|
||||
"Alice mentioned she works long hours, often past midnight | Involving: Alice (occurred_start=2024-01-15, mentioned_at=2024-01-15)"
|
||||
|
||||
Observation text rules:
|
||||
- Write clean prose — NEVER copy raw fact lines or their metadata (temporal fields, "Involving:", "When:" labels, UUIDs).
|
||||
- Parenthesized metadata like `(occurred_start=...)` and pipe-separated labels like `| Involving: ...` are fact formatting — strip them entirely from observation text.
|
||||
- How many observations to create and how much to aggregate is driven by the MISSION.
|
||||
- Parenthesized metadata like (occurred_start=...) and pipe-separated labels like "| Involving: ..." are fact formatting — strip them entirely from observation text.
|
||||
- How many observations to create and how much to aggregate is driven by the MISSION above.
|
||||
|
||||
### Field rules
|
||||
{{"creates": [{{"text": "Alice works long hours, often past midnight.", "source_fact_ids": ["a1b2c3d4-e5f6-7890-abcd-ef1234567890"]}}, {{"text": "Alice feels exhausted from project deadlines.", "source_fact_ids": ["b2c3d4e5-f6a7-8901-bcde-f12345678901"]}}],
|
||||
"updates": [{{"text": "Alice works at Acme Corp as a senior engineer", "observation_id": "c3d4e5f6-a7b8-9012-cdef-123456789012", "source_fact_ids": ["d4e5f6a7-b8c9-0123-defa-234567890123"]}}],
|
||||
"deletes": [{{"observation_id": "e5f6a7b8-c9d0-1234-efab-345678901234"}}]}}
|
||||
|
||||
- `source_fact_ids`: copy the EXACT UUID strings shown in brackets `[uuid]` from new facts — never use integers or positions.
|
||||
- `observation_id`: copy the EXACT `id` UUID string from existing observations.
|
||||
- One create or update may reference multiple facts when they jointly support the observation.
|
||||
- **AT MOST ONE UPDATE PER `observation_id`**: if several new facts all update the same existing observation, emit a single `updates` entry that lists all contributing `source_fact_ids` and a single consolidated `text`. Never emit two `updates` entries with the same `observation_id` in one response — they would silently overwrite each other.
|
||||
- `deletes`: only when an observation is directly superseded or contradicted by new facts.
|
||||
- Do NOT include `tags` — handled automatically.
|
||||
- Return `{{"creates": [], "updates": [], "deletes": []}}` if nothing durable is found."""
|
||||
Rules:
|
||||
- "source_fact_ids": copy the EXACT UUID strings shown in brackets [uuid] from NEW FACTS — never use integers or positions.
|
||||
- "observation_id": copy the EXACT "id" UUID string from EXISTING OBSERVATIONS.
|
||||
- One create/update may reference multiple facts when they jointly support the observation.
|
||||
- "deletes": only when an observation is directly superseded or contradicted by new facts.
|
||||
- Do NOT include "tags" — handled automatically.
|
||||
- Return {{"creates": [], "updates": [], "deletes": []}} if nothing durable is found."""
|
||||
|
||||
|
||||
def build_batch_consolidation_prompt(
|
||||
observations_mission: str | None = None,
|
||||
observation_capacity_note: str | None = None,
|
||||
llm_output_language: str | None = None,
|
||||
) -> str:
|
||||
"""
|
||||
Build the consolidation prompt for batch mode (multiple facts per LLM call).
|
||||
|
||||
The mission defines *what* to track (customisable per bank) and takes
|
||||
priority over the built-in processing rules when the two conflict.
|
||||
Processing rules, decision guide, and output format are always present.
|
||||
When ``llm_output_language`` is set, observations are emitted in that
|
||||
language.
|
||||
The mission defines *what* to track (customisable per bank).
|
||||
Processing rules and output format are always present regardless of mission.
|
||||
"""
|
||||
mission = escape_for_prompt(observations_mission or _DEFAULT_MISSION)
|
||||
mission = observations_mission or _DEFAULT_MISSION
|
||||
|
||||
capacity_section = ""
|
||||
if observation_capacity_note:
|
||||
capacity_section = f"\n\n## CAPACITY CONSTRAINT\n\n{escape_for_prompt(observation_capacity_note)}"
|
||||
capacity_section = f"\n\n## CAPACITY CONSTRAINT\n{observation_capacity_note}"
|
||||
|
||||
return (
|
||||
"You are a memory consolidation system. Synthesize new facts into "
|
||||
"observations, merging with existing observations when appropriate.\n\n"
|
||||
f"## MISSION\n\n{mission}\n\n"
|
||||
f"{_MISSION_PRIORITY_NOTE}"
|
||||
f"{capacity_section}\n\n"
|
||||
f"{_PROCESSING_RULES}\n\n"
|
||||
f"{_INPUT_SECTION}\n\n"
|
||||
f"{_DECISION_GUIDE}\n\n"
|
||||
f"{_OUTPUT_SECTION}" + output_language_directive(llm_output_language)
|
||||
"You are a memory consolidation system. Synthesize facts into observations "
|
||||
"and merge with existing observations when appropriate.\n\n"
|
||||
f"## MISSION\n{mission}{capacity_section}\n\n"
|
||||
f"{_PROCESSING_RULES}" + _BATCH_DATA_SECTION + _BATCH_OUTPUT_FORMAT
|
||||
)
|
||||
|
||||
@@ -17,7 +17,6 @@ import httpx
|
||||
|
||||
from ..config import (
|
||||
DEFAULT_LITELLM_API_BASE,
|
||||
DEFAULT_RERANKER_ALIBABA_MODEL,
|
||||
DEFAULT_RERANKER_COHERE_MODEL,
|
||||
DEFAULT_RERANKER_FLASHRANK_CACHE_DIR,
|
||||
DEFAULT_RERANKER_FLASHRANK_CPU_MEM_ARENA,
|
||||
@@ -38,8 +37,6 @@ from ..config import (
|
||||
DEFAULT_RERANKER_TEI_HTTP_TIMEOUT,
|
||||
DEFAULT_RERANKER_TEI_MAX_CONCURRENT,
|
||||
DEFAULT_RERANKER_ZEROENTROPY_MODEL,
|
||||
DEFAULT_ZEROENTROPY_BASE_URL,
|
||||
ENV_RERANKER_ALIBABA_API_KEY,
|
||||
ENV_RERANKER_COHERE_API_KEY,
|
||||
ENV_RERANKER_COHERE_MODEL,
|
||||
ENV_RERANKER_FLASHRANK_CACHE_DIR,
|
||||
@@ -63,43 +60,6 @@ from ..config import (
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _resolve_malloc_trim():
|
||||
"""Return a callable that asks glibc to release freed heap pages to the OS.
|
||||
|
||||
Local CPU rerankers (FlashRank/ONNX, SentenceTransformers/torch) allocate
|
||||
large transient numpy/tensor buffers per call. On Linux glibc, those pages
|
||||
are freed at the Python level but kept by the allocator as a high-water
|
||||
mark — RSS grows monotonically across many recalls (see issue #1717).
|
||||
Calling `malloc_trim(0)` after each batch returns those pages to the OS.
|
||||
|
||||
Resolved once at import; returns a no-op on non-glibc platforms (macOS,
|
||||
musl, Windows) where the call is unavailable or unnecessary.
|
||||
"""
|
||||
import sys
|
||||
|
||||
if sys.platform != "linux":
|
||||
return lambda: None
|
||||
|
||||
import ctypes
|
||||
import ctypes.util
|
||||
|
||||
libc_path = ctypes.util.find_library("c")
|
||||
if libc_path is None:
|
||||
return lambda: None
|
||||
try:
|
||||
libc = ctypes.CDLL(libc_path)
|
||||
trim = libc.malloc_trim
|
||||
except (OSError, AttributeError):
|
||||
# Not glibc (musl has no malloc_trim) or libc lookup failed.
|
||||
return lambda: None
|
||||
trim.argtypes = [ctypes.c_size_t]
|
||||
trim.restype = ctypes.c_int
|
||||
return lambda: trim(0)
|
||||
|
||||
|
||||
_malloc_trim = _resolve_malloc_trim()
|
||||
|
||||
|
||||
class CrossEncoderModel(ABC):
|
||||
"""
|
||||
Abstract base class for cross-encoder reranking.
|
||||
@@ -306,28 +266,25 @@ class LocalSTCrossEncoder(CrossEncoderModel):
|
||||
"""
|
||||
import numpy as np
|
||||
|
||||
try:
|
||||
if self.bucket_batching and len(pairs) > 1:
|
||||
# Sort pairs by approximate token length to create homogeneous batches.
|
||||
# This eliminates padding waste — short pairs aren't padded to the length
|
||||
# of the longest pair in the batch. Quality-identical by construction.
|
||||
lengths = [len(pairs[i][0]) + len(pairs[i][1]) for i in range(len(pairs))]
|
||||
sorted_indices = sorted(range(len(pairs)), key=lambda i: lengths[i])
|
||||
sorted_pairs = [pairs[i] for i in sorted_indices]
|
||||
if self.bucket_batching and len(pairs) > 1:
|
||||
# Sort pairs by approximate token length to create homogeneous batches.
|
||||
# This eliminates padding waste — short pairs aren't padded to the length
|
||||
# of the longest pair in the batch. Quality-identical by construction.
|
||||
lengths = [len(pairs[i][0]) + len(pairs[i][1]) for i in range(len(pairs))]
|
||||
sorted_indices = sorted(range(len(pairs)), key=lambda i: lengths[i])
|
||||
sorted_pairs = [pairs[i] for i in sorted_indices]
|
||||
|
||||
sorted_scores = self._model.predict(sorted_pairs, batch_size=self.batch_size, show_progress_bar=False)
|
||||
sorted_scores = sorted_scores.tolist() if hasattr(sorted_scores, "tolist") else list(sorted_scores)
|
||||
sorted_scores = self._model.predict(sorted_pairs, batch_size=self.batch_size, show_progress_bar=False)
|
||||
sorted_scores = sorted_scores.tolist() if hasattr(sorted_scores, "tolist") else list(sorted_scores)
|
||||
|
||||
# Restore original order
|
||||
scores = [0.0] * len(pairs)
|
||||
for new_pos, orig_idx in enumerate(sorted_indices):
|
||||
scores[orig_idx] = sorted_scores[new_pos]
|
||||
return scores
|
||||
# Restore original order
|
||||
scores = [0.0] * len(pairs)
|
||||
for new_pos, orig_idx in enumerate(sorted_indices):
|
||||
scores[orig_idx] = sorted_scores[new_pos]
|
||||
return scores
|
||||
|
||||
scores = self._model.predict(pairs, batch_size=self.batch_size, show_progress_bar=False)
|
||||
return scores.tolist() if hasattr(scores, "tolist") else list(scores)
|
||||
finally:
|
||||
_malloc_trim()
|
||||
scores = self._model.predict(pairs, batch_size=self.batch_size, show_progress_bar=False)
|
||||
return scores.tolist() if hasattr(scores, "tolist") else list(scores)
|
||||
|
||||
async def predict(self, pairs: list[tuple[str, str]]) -> list[float]:
|
||||
"""
|
||||
@@ -589,14 +546,12 @@ class _CohereCompatibleRerankClient:
|
||||
rerank_url: str,
|
||||
timeout: float = 60.0,
|
||||
include_top_n: bool = True,
|
||||
include_return_documents: bool = False,
|
||||
):
|
||||
self.api_key = api_key
|
||||
self.model = model
|
||||
self.rerank_url = rerank_url
|
||||
self.timeout = timeout
|
||||
self.include_top_n = include_top_n
|
||||
self.include_return_documents = include_return_documents
|
||||
self._async_client: httpx.AsyncClient | None = None
|
||||
|
||||
async def initialize(self) -> None:
|
||||
@@ -774,7 +729,7 @@ class ZeroEntropyCrossEncoder(CrossEncoderModel):
|
||||
See: https://docs.zeroentropy.dev/models
|
||||
"""
|
||||
|
||||
DEFAULT_BASE_URL = DEFAULT_ZEROENTROPY_BASE_URL
|
||||
DEFAULT_BASE_URL = "https://api.zeroentropy.dev"
|
||||
RERANK_PATH = "/v1/models/rerank"
|
||||
|
||||
def __init__(
|
||||
@@ -1007,35 +962,32 @@ class FlashRankCrossEncoder(CrossEncoderModel):
|
||||
if not pairs:
|
||||
return []
|
||||
|
||||
try:
|
||||
# Group pairs by query
|
||||
query_groups: dict[str, list[tuple[int, str]]] = {}
|
||||
for idx, (query, text) in enumerate(pairs):
|
||||
if query not in query_groups:
|
||||
query_groups[query] = []
|
||||
query_groups[query].append((idx, text))
|
||||
# Group pairs by query
|
||||
query_groups: dict[str, list[tuple[int, str]]] = {}
|
||||
for idx, (query, text) in enumerate(pairs):
|
||||
if query not in query_groups:
|
||||
query_groups[query] = []
|
||||
query_groups[query].append((idx, text))
|
||||
|
||||
all_scores = [0.0] * len(pairs)
|
||||
all_scores = [0.0] * len(pairs)
|
||||
|
||||
for query, indexed_texts in query_groups.items():
|
||||
# Build passages list for FlashRank
|
||||
passages = [{"id": i, "text": text} for i, (_, text) in enumerate(indexed_texts)]
|
||||
global_indices = [idx for idx, _ in indexed_texts]
|
||||
for query, indexed_texts in query_groups.items():
|
||||
# Build passages list for FlashRank
|
||||
passages = [{"id": i, "text": text} for i, (_, text) in enumerate(indexed_texts)]
|
||||
global_indices = [idx for idx, _ in indexed_texts]
|
||||
|
||||
# Create rerank request
|
||||
request = RerankRequest(query=query, passages=passages)
|
||||
results = self._ranker.rerank(request)
|
||||
# Create rerank request
|
||||
request = RerankRequest(query=query, passages=passages)
|
||||
results = self._ranker.rerank(request)
|
||||
|
||||
# Map scores back to original positions
|
||||
for result in results:
|
||||
local_idx = result["id"]
|
||||
score = result["score"]
|
||||
global_idx = global_indices[local_idx]
|
||||
all_scores[global_idx] = score
|
||||
# Map scores back to original positions
|
||||
for result in results:
|
||||
local_idx = result["id"]
|
||||
score = result["score"]
|
||||
global_idx = global_indices[local_idx]
|
||||
all_scores[global_idx] = score
|
||||
|
||||
return all_scores
|
||||
finally:
|
||||
_malloc_trim()
|
||||
return all_scores
|
||||
|
||||
async def predict(self, pairs: list[tuple[str, str]]) -> list[float]:
|
||||
"""
|
||||
@@ -1582,48 +1534,6 @@ class GoogleCrossEncoder(CrossEncoderModel):
|
||||
return await loop.run_in_executor(None, self._predict_sync, pairs)
|
||||
|
||||
|
||||
class AlibabaCloudCrossEncoder(CrossEncoderModel):
|
||||
"""
|
||||
Alibaba Cloud DashScope text reranking API.
|
||||
|
||||
Uses the Cohere-compatible /reranks endpoint, which is the standard interface
|
||||
for qwen3-rerank. Authentication via HINDSIGHT_API_RERANKER_ALIBABA_API_KEY
|
||||
(or DASHSCOPE_API_KEY as a fallback).
|
||||
See: https://help.aliyun.com/zh/model-studio/text-rerank-api
|
||||
"""
|
||||
|
||||
RERANK_URL = "https://dashscope.aliyuncs.com/compatible-api/v1/reranks"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
api_key: str,
|
||||
model: str = DEFAULT_RERANKER_ALIBABA_MODEL,
|
||||
timeout: float = 60.0,
|
||||
):
|
||||
self.model = model
|
||||
self._client = _CohereCompatibleRerankClient(
|
||||
api_key=api_key,
|
||||
model=model,
|
||||
rerank_url=self.RERANK_URL,
|
||||
timeout=timeout,
|
||||
include_return_documents=False,
|
||||
)
|
||||
|
||||
@property
|
||||
def provider_name(self) -> str:
|
||||
return "alibaba"
|
||||
|
||||
async def initialize(self) -> None:
|
||||
if self._client._async_client is not None:
|
||||
return
|
||||
logger.info(f"Reranker: initializing Alibaba Cloud provider with model {self.model}")
|
||||
await self._client.initialize()
|
||||
logger.info("Reranker: Alibaba Cloud provider initialized")
|
||||
|
||||
async def predict(self, pairs: list[tuple[str, str]]) -> list[float]:
|
||||
return await self._client.predict(pairs)
|
||||
|
||||
|
||||
def create_cross_encoder_from_env() -> CrossEncoderModel:
|
||||
"""
|
||||
Create a CrossEncoderModel instance based on configuration.
|
||||
@@ -1666,7 +1576,6 @@ def create_cross_encoder_from_env() -> CrossEncoderModel:
|
||||
api_key=api_key,
|
||||
model=config.reranker_cohere_model,
|
||||
base_url=config.reranker_cohere_base_url,
|
||||
timeout=config.reranker_cohere_timeout,
|
||||
)
|
||||
elif provider == "openrouter":
|
||||
api_key = config.reranker_openrouter_api_key
|
||||
@@ -1679,7 +1588,6 @@ def create_cross_encoder_from_env() -> CrossEncoderModel:
|
||||
api_key=api_key,
|
||||
model=config.reranker_openrouter_model,
|
||||
base_url="https://openrouter.ai/api/v1/rerank",
|
||||
timeout=config.reranker_openrouter_timeout,
|
||||
)
|
||||
elif provider == "flashrank":
|
||||
model = os.environ.get(ENV_RERANKER_FLASHRANK_MODEL, DEFAULT_RERANKER_FLASHRANK_MODEL)
|
||||
@@ -1694,7 +1602,6 @@ def create_cross_encoder_from_env() -> CrossEncoderModel:
|
||||
api_key=config.reranker_litellm_api_key,
|
||||
model=config.reranker_litellm_model,
|
||||
max_tokens_per_doc=config.reranker_litellm_max_tokens_per_doc,
|
||||
timeout=config.reranker_litellm_timeout,
|
||||
)
|
||||
elif provider == "litellm-sdk":
|
||||
api_key = config.reranker_litellm_sdk_api_key
|
||||
@@ -1707,7 +1614,6 @@ def create_cross_encoder_from_env() -> CrossEncoderModel:
|
||||
model=config.reranker_litellm_sdk_model,
|
||||
api_base=config.reranker_litellm_sdk_api_base,
|
||||
max_tokens_per_doc=config.reranker_litellm_max_tokens_per_doc,
|
||||
timeout=config.reranker_litellm_sdk_timeout,
|
||||
)
|
||||
elif provider == "zeroentropy":
|
||||
api_key = config.reranker_zeroentropy_api_key
|
||||
@@ -1718,8 +1624,6 @@ def create_cross_encoder_from_env() -> CrossEncoderModel:
|
||||
return ZeroEntropyCrossEncoder(
|
||||
api_key=api_key,
|
||||
model=config.reranker_zeroentropy_model,
|
||||
base_url=config.reranker_zeroentropy_base_url,
|
||||
timeout=config.reranker_zeroentropy_timeout,
|
||||
)
|
||||
elif provider == "siliconflow":
|
||||
api_key = config.reranker_siliconflow_api_key
|
||||
@@ -1731,7 +1635,6 @@ def create_cross_encoder_from_env() -> CrossEncoderModel:
|
||||
api_key=api_key,
|
||||
model=config.reranker_siliconflow_model,
|
||||
base_url=config.reranker_siliconflow_base_url,
|
||||
timeout=config.reranker_siliconflow_timeout,
|
||||
)
|
||||
elif provider == "google":
|
||||
project_id = config.reranker_google_project_id
|
||||
@@ -1744,16 +1647,6 @@ def create_cross_encoder_from_env() -> CrossEncoderModel:
|
||||
project_id=project_id,
|
||||
model=config.reranker_google_model,
|
||||
service_account_key=config.reranker_google_service_account_key,
|
||||
timeout=config.reranker_google_timeout,
|
||||
)
|
||||
elif provider == "alibaba":
|
||||
api_key = config.reranker_alibaba_api_key
|
||||
if not api_key:
|
||||
raise ValueError(f"{ENV_RERANKER_ALIBABA_API_KEY} is required when {ENV_RERANKER_PROVIDER} is 'alibaba'")
|
||||
return AlibabaCloudCrossEncoder(
|
||||
api_key=api_key,
|
||||
model=config.reranker_alibaba_model,
|
||||
timeout=config.reranker_alibaba_timeout,
|
||||
)
|
||||
elif provider == "rrf":
|
||||
return RRFPassthroughCrossEncoder()
|
||||
@@ -1761,5 +1654,5 @@ def create_cross_encoder_from_env() -> CrossEncoderModel:
|
||||
return JinaMLXCrossEncoder()
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unknown reranker provider: {provider}. Supported: 'local', 'tei', 'cohere', 'zeroentropy', 'siliconflow', 'alibaba', 'google', 'flashrank', 'litellm', 'litellm-sdk', 'rrf', 'jina-mlx'"
|
||||
f"Unknown reranker provider: {provider}. Supported: 'local', 'tei', 'cohere', 'zeroentropy', 'siliconflow', 'google', 'flashrank', 'litellm', 'litellm-sdk', 'rrf', 'jina-mlx'"
|
||||
)
|
||||
|
||||
@@ -166,6 +166,21 @@ class DataAccessOps(ABC):
|
||||
|
||||
# -- LATERAL / fan-out queries ---------------------------------------
|
||||
|
||||
@abstractmethod
|
||||
async def fetch_entity_unit_fanout(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
ue_table: str,
|
||||
entity_id_list: list[UUID],
|
||||
limit_per_entity: int,
|
||||
) -> list[ResultRow]:
|
||||
"""Fetch unit_ids for a list of entities with per-entity row cap.
|
||||
|
||||
PG uses unnest + CROSS JOIN LATERAL with LIMIT.
|
||||
Non-PG queries each entity individually.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def fetch_unit_dates(
|
||||
self,
|
||||
@@ -391,74 +406,6 @@ class DataAccessOps(ABC):
|
||||
"""Insert a webhook delivery task into async_operations."""
|
||||
...
|
||||
|
||||
# -- Graph maintenance queue -----------------------------------------
|
||||
|
||||
@abstractmethod
|
||||
async def enqueue_graph_maintenance(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
bank_id: str,
|
||||
unit_ids: list,
|
||||
) -> None:
|
||||
"""Insert unit_ids into graph_maintenance_queue, deduplicating on the
|
||||
(bank_id, unit_id) primary key.
|
||||
|
||||
Called inside the triggering transaction so enqueue is atomic with
|
||||
the mutation that caused it. Order is unspecified.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def claim_graph_maintenance_batch(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
bank_id: str,
|
||||
limit: int,
|
||||
) -> list[str]:
|
||||
"""Atomically claim a batch of rows from graph_maintenance_queue and
|
||||
remove them from the table.
|
||||
|
||||
Returns the list of ``unit_id`` strings. Empty list when the queue
|
||||
for ``bank_id`` is drained.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def prune_orphan_entities(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
entities_table: str,
|
||||
ue_table: str,
|
||||
bank_id: str,
|
||||
) -> int:
|
||||
"""Delete entities in ``bank_id`` that no longer have any unit_entities
|
||||
rows referencing them. Returns the number of rows deleted.
|
||||
|
||||
FK ON DELETE CASCADE on entity_cooccurrences then removes any
|
||||
cooccurrence row pointing at the pruned entities.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def prune_stale_cooccurrences(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
ec_table: str,
|
||||
ue_table: str,
|
||||
entities_table: str,
|
||||
bank_id: str,
|
||||
) -> int:
|
||||
"""Delete entity_cooccurrences rows in ``bank_id`` where the two
|
||||
entities still exist but no current unit references both of them.
|
||||
|
||||
These are stale-count rows: cooccurrence was real at the time it was
|
||||
recorded, but every memory_unit that witnessed both entities has
|
||||
since been deleted. Returns the number of rows deleted.
|
||||
"""
|
||||
...
|
||||
|
||||
# -- Task claiming operations ------------------------------------------
|
||||
|
||||
@abstractmethod
|
||||
@@ -469,8 +416,6 @@ class DataAccessOps(ABC):
|
||||
worker_id: str,
|
||||
reserved_limits: dict[str, int],
|
||||
shared_limit: int,
|
||||
*,
|
||||
consolidation_bank_priority: dict[str, int] | None = None,
|
||||
) -> list[ResultRow]:
|
||||
"""Claim pending tasks from the async_operations table.
|
||||
|
||||
@@ -478,14 +423,6 @@ class DataAccessOps(ABC):
|
||||
Oracle implementation uses two-step claims (query busy banks first, then
|
||||
claim excluding them) to avoid ORA-02014.
|
||||
|
||||
Args:
|
||||
consolidation_bank_priority: Per-bank priority for consolidation scheduling.
|
||||
Maps bank name patterns to integer priorities (higher = claimed first).
|
||||
Patterns support ``*`` as wildcard (converted to SQL ``%`` for LIKE).
|
||||
A bare ``*`` key is the catch-all default for unlisted banks.
|
||||
When set, consolidation tasks are claimed in priority tiers.
|
||||
None preserves current behavior (pure created_at ordering).
|
||||
|
||||
Returns claimed rows with operation_id, operation_type, task_payload, retry_count.
|
||||
The caller is responsible for building ClaimedTask objects.
|
||||
"""
|
||||
|
||||
@@ -215,96 +215,29 @@ class OracleOps(DataAccessOps):
|
||||
list(zip(unit_ids, entity_ids)),
|
||||
)
|
||||
|
||||
async def enqueue_graph_maintenance(
|
||||
async def fetch_entity_unit_fanout(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
bank_id: str,
|
||||
unit_ids: list,
|
||||
) -> None:
|
||||
if not unit_ids:
|
||||
return
|
||||
# Oracle doesn't support ON CONFLICT; rely on the PK and the
|
||||
# IGNORE_ROW_ON_DUPKEY_INDEX hint to skip duplicates server-side.
|
||||
# The hint name must match the PK constraint exactly.
|
||||
await conn.executemany(
|
||||
f"""
|
||||
INSERT /*+ IGNORE_ROW_ON_DUPKEY_INDEX({table}, pk_graph_maintenance_queue) */
|
||||
INTO {table} (bank_id, unit_id)
|
||||
VALUES ($1, $2)
|
||||
""",
|
||||
[(bank_id, uid) for uid in unit_ids],
|
||||
)
|
||||
|
||||
async def claim_graph_maintenance_batch(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
bank_id: str,
|
||||
limit: int,
|
||||
) -> list[str]:
|
||||
# Two-step claim: select the batch, then delete by exact keys. Oracle's
|
||||
# DELETE ... RETURNING doesn't accept a multi-row subquery, so we can't
|
||||
# do it in one statement like the PG version.
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT unit_id FROM {table}
|
||||
WHERE bank_id = $1
|
||||
ORDER BY enqueued_at
|
||||
FETCH FIRST $2 ROWS ONLY
|
||||
""",
|
||||
bank_id,
|
||||
limit,
|
||||
)
|
||||
claimed = [str(row["unit_id"]) for row in rows]
|
||||
if claimed:
|
||||
await conn.executemany(
|
||||
f"DELETE FROM {table} WHERE bank_id = $1 AND unit_id = $2",
|
||||
[(bank_id, uid) for uid in claimed],
|
||||
ue_table: str,
|
||||
entity_id_list: list[UUID],
|
||||
limit_per_entity: int,
|
||||
) -> list[ResultRow]:
|
||||
# Query each entity individually
|
||||
rows: list[ResultRow] = []
|
||||
for eid in entity_id_list:
|
||||
entity_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT $1 AS entity_id, ue.unit_id
|
||||
FROM {ue_table} ue
|
||||
WHERE ue.entity_id = $1
|
||||
ORDER BY ue.unit_id DESC
|
||||
LIMIT $2
|
||||
""",
|
||||
eid,
|
||||
limit_per_entity,
|
||||
)
|
||||
return claimed
|
||||
|
||||
async def prune_orphan_entities(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
entities_table: str,
|
||||
ue_table: str,
|
||||
bank_id: str,
|
||||
) -> int:
|
||||
# The Oracle DatabaseConnection wrapper reshapes ``cursor.rowcount`` into
|
||||
# the same ``"DELETE N"`` status string asyncpg returns, so the same
|
||||
# ``int(deleted.split()[-1])`` parsing works on both dialects.
|
||||
deleted = await conn.execute(
|
||||
f"""
|
||||
DELETE FROM {entities_table}
|
||||
WHERE bank_id = $1
|
||||
AND id NOT IN (SELECT DISTINCT entity_id FROM {ue_table})
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
return int(deleted.split()[-1]) if isinstance(deleted, str) and deleted.startswith("DELETE") else 0
|
||||
|
||||
async def prune_stale_cooccurrences(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
ec_table: str,
|
||||
ue_table: str,
|
||||
entities_table: str,
|
||||
bank_id: str,
|
||||
) -> int:
|
||||
deleted = await conn.execute(
|
||||
f"""
|
||||
DELETE FROM {ec_table}
|
||||
WHERE entity_id_1 IN (SELECT id FROM {entities_table} WHERE bank_id = $1)
|
||||
AND (entity_id_1, entity_id_2) NOT IN (
|
||||
SELECT u1.entity_id, u2.entity_id
|
||||
FROM {ue_table} u1
|
||||
JOIN {ue_table} u2 ON u1.unit_id = u2.unit_id
|
||||
)
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
return int(deleted.split()[-1]) if isinstance(deleted, str) and deleted.startswith("DELETE") else 0
|
||||
rows.extend(entity_rows)
|
||||
return rows
|
||||
|
||||
async def fetch_unit_dates(
|
||||
self,
|
||||
@@ -787,257 +720,7 @@ class OracleOps(DataAccessOps):
|
||||
|
||||
# -- Task claiming operations ------------------------------------------
|
||||
|
||||
async def _claim_consolidation_tasks(
|
||||
self,
|
||||
conn,
|
||||
table: str,
|
||||
busy_bank_ids: list[str],
|
||||
claimed_ids: list,
|
||||
limit: int,
|
||||
priority_map: dict[str, int] | None,
|
||||
) -> list:
|
||||
"""Claim consolidation tasks with optional priority-based tiered ordering.
|
||||
|
||||
Mirrors the PostgreSQL implementation. The Oracle SQL adapter
|
||||
translates ``LIKE ANY`` / ``NOT LIKE ALL`` via ``_expand_any_lists``.
|
||||
"""
|
||||
if limit <= 0:
|
||||
return []
|
||||
|
||||
if not priority_map:
|
||||
return await self._claim_consolidation_plain(conn, table, busy_bank_ids, claimed_ids, limit)
|
||||
|
||||
# --- Tiered claiming (same algorithm as PG) ---
|
||||
specific_by_priority: dict[int, list[str]] = {}
|
||||
all_specific_sql: list[str] = []
|
||||
catch_all_priority = 1
|
||||
|
||||
for pattern, priority in priority_map.items():
|
||||
if pattern == "*":
|
||||
catch_all_priority = priority
|
||||
else:
|
||||
sql_pat = pattern.replace("*", "%")
|
||||
specific_by_priority.setdefault(priority, []).append(sql_pat)
|
||||
all_specific_sql.append(sql_pat)
|
||||
|
||||
all_priorities = sorted(set(specific_by_priority.keys()) | {catch_all_priority}, reverse=True)
|
||||
|
||||
remaining = limit
|
||||
result: list = []
|
||||
|
||||
for pri in all_priorities:
|
||||
if remaining <= 0:
|
||||
break
|
||||
|
||||
if pri in specific_by_priority:
|
||||
rows = await self._claim_consolidation_like(
|
||||
conn,
|
||||
table,
|
||||
busy_bank_ids,
|
||||
claimed_ids,
|
||||
remaining,
|
||||
specific_by_priority[pri],
|
||||
)
|
||||
for row in rows:
|
||||
claimed_ids.append(row["operation_id"])
|
||||
result.append(row)
|
||||
remaining -= len(rows)
|
||||
|
||||
if pri == catch_all_priority and remaining > 0:
|
||||
rows = await self._claim_consolidation_not_like(
|
||||
conn,
|
||||
table,
|
||||
busy_bank_ids,
|
||||
claimed_ids,
|
||||
remaining,
|
||||
all_specific_sql,
|
||||
)
|
||||
for row in rows:
|
||||
claimed_ids.append(row["operation_id"])
|
||||
result.append(row)
|
||||
remaining -= len(rows)
|
||||
|
||||
return result
|
||||
|
||||
async def _claim_consolidation_plain(
|
||||
self,
|
||||
conn,
|
||||
table,
|
||||
busy_bank_ids,
|
||||
claimed_ids,
|
||||
limit,
|
||||
) -> list:
|
||||
"""Claim consolidation tasks with default created_at ordering."""
|
||||
exclude_ids = claimed_ids if claimed_ids else None
|
||||
if busy_bank_ids:
|
||||
if exclude_ids:
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND bank_id != ALL($1::text[])
|
||||
AND operation_id != ALL($2::uuid[])
|
||||
ORDER BY created_at
|
||||
LIMIT $3
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
busy_bank_ids,
|
||||
exclude_ids,
|
||||
limit,
|
||||
)
|
||||
else:
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND bank_id != ALL($1::text[])
|
||||
ORDER BY created_at
|
||||
LIMIT $2
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
busy_bank_ids,
|
||||
limit,
|
||||
)
|
||||
else:
|
||||
if exclude_ids:
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND operation_id != ALL($1::uuid[])
|
||||
ORDER BY created_at
|
||||
LIMIT $2
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
exclude_ids,
|
||||
limit,
|
||||
)
|
||||
else:
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
ORDER BY created_at
|
||||
LIMIT $1
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
limit,
|
||||
)
|
||||
|
||||
async def _claim_consolidation_like(
|
||||
self,
|
||||
conn,
|
||||
table,
|
||||
busy_bank_ids,
|
||||
claimed_ids,
|
||||
limit,
|
||||
sql_patterns,
|
||||
) -> list:
|
||||
"""Claim consolidation tasks from banks matching LIKE patterns."""
|
||||
params: list = [sql_patterns]
|
||||
conditions = ["bank_id LIKE ANY($1::text[])"]
|
||||
idx = 2
|
||||
|
||||
if busy_bank_ids:
|
||||
conditions.append(f"bank_id != ALL(${idx}::text[])")
|
||||
params.append(busy_bank_ids)
|
||||
idx += 1
|
||||
|
||||
if claimed_ids:
|
||||
conditions.append(f"operation_id != ALL(${idx}::uuid[])")
|
||||
params.append(claimed_ids)
|
||||
idx += 1
|
||||
|
||||
params.append(limit)
|
||||
extra = " AND ".join(conditions)
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND {extra}
|
||||
ORDER BY created_at
|
||||
LIMIT ${idx}
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
*params,
|
||||
)
|
||||
|
||||
async def _claim_consolidation_not_like(
|
||||
self,
|
||||
conn,
|
||||
table,
|
||||
busy_bank_ids,
|
||||
claimed_ids,
|
||||
limit,
|
||||
exclude_patterns,
|
||||
) -> list:
|
||||
"""Claim consolidation tasks from banks NOT matching any specific pattern (catch-all tier)."""
|
||||
params: list = []
|
||||
conditions: list[str] = []
|
||||
idx = 1
|
||||
|
||||
if exclude_patterns:
|
||||
conditions.append(f"bank_id NOT LIKE ALL(${idx}::text[])")
|
||||
params.append(exclude_patterns)
|
||||
idx += 1
|
||||
|
||||
if busy_bank_ids:
|
||||
conditions.append(f"bank_id != ALL(${idx}::text[])")
|
||||
params.append(busy_bank_ids)
|
||||
idx += 1
|
||||
|
||||
if claimed_ids:
|
||||
conditions.append(f"operation_id != ALL(${idx}::uuid[])")
|
||||
params.append(claimed_ids)
|
||||
idx += 1
|
||||
|
||||
params.append(limit)
|
||||
extra_clause = (" AND " + " AND ".join(conditions)) if conditions else ""
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW()){extra_clause}
|
||||
ORDER BY created_at
|
||||
LIMIT ${idx}
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
*params,
|
||||
)
|
||||
|
||||
async def claim_tasks(
|
||||
self,
|
||||
conn,
|
||||
table,
|
||||
worker_id,
|
||||
reserved_limits,
|
||||
shared_limit,
|
||||
*,
|
||||
consolidation_bank_priority=None,
|
||||
):
|
||||
async def claim_tasks(self, conn, table, worker_id, reserved_limits, shared_limit):
|
||||
"""Oracle two-step claiming to avoid ORA-02014 with NOT EXISTS + FOR UPDATE."""
|
||||
all_rows = []
|
||||
claimed_ids = []
|
||||
@@ -1048,6 +731,7 @@ class OracleOps(DataAccessOps):
|
||||
continue
|
||||
|
||||
if op_type == "consolidation":
|
||||
# Two-step: find busy banks first, then claim excluding them
|
||||
busy_banks = await conn.fetch(
|
||||
f"""
|
||||
SELECT DISTINCT bank_id FROM {table}
|
||||
@@ -1056,14 +740,38 @@ class OracleOps(DataAccessOps):
|
||||
)
|
||||
busy_bank_ids = [r["bank_id"] for r in busy_banks]
|
||||
|
||||
rows = await self._claim_consolidation_tasks(
|
||||
conn,
|
||||
table,
|
||||
busy_bank_ids,
|
||||
claimed_ids,
|
||||
limit,
|
||||
consolidation_bank_priority,
|
||||
)
|
||||
if busy_bank_ids:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND bank_id != ALL($1::text[])
|
||||
ORDER BY created_at
|
||||
LIMIT $2
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
busy_bank_ids,
|
||||
limit,
|
||||
)
|
||||
else:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
ORDER BY created_at
|
||||
LIMIT $1
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
limit,
|
||||
)
|
||||
else:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
@@ -1127,7 +835,7 @@ class OracleOps(DataAccessOps):
|
||||
all_rows.append(row)
|
||||
remaining_shared -= len(rows)
|
||||
|
||||
# 2b. Consolidation tasks (with bank-serialization + optional priority)
|
||||
# 2b. Consolidation tasks (with bank-serialization)
|
||||
if remaining_shared > 0:
|
||||
busy_banks_2 = await conn.fetch(
|
||||
f"""
|
||||
@@ -1137,14 +845,76 @@ class OracleOps(DataAccessOps):
|
||||
)
|
||||
busy_bank_ids_2 = [r["bank_id"] for r in busy_banks_2]
|
||||
|
||||
rows = await self._claim_consolidation_tasks(
|
||||
conn,
|
||||
table,
|
||||
busy_bank_ids_2,
|
||||
claimed_ids,
|
||||
remaining_shared,
|
||||
consolidation_bank_priority,
|
||||
)
|
||||
if claimed_ids:
|
||||
if busy_bank_ids_2:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND operation_id != ALL($1::uuid[])
|
||||
AND bank_id != ALL($2::text[])
|
||||
ORDER BY created_at
|
||||
LIMIT $3
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
claimed_ids,
|
||||
busy_bank_ids_2,
|
||||
remaining_shared,
|
||||
)
|
||||
else:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND operation_id != ALL($1::uuid[])
|
||||
ORDER BY created_at
|
||||
LIMIT $2
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
claimed_ids,
|
||||
remaining_shared,
|
||||
)
|
||||
else:
|
||||
if busy_bank_ids_2:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND bank_id != ALL($1::text[])
|
||||
ORDER BY created_at
|
||||
LIMIT $2
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
busy_bank_ids_2,
|
||||
remaining_shared,
|
||||
)
|
||||
else:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
ORDER BY created_at
|
||||
LIMIT $1
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
remaining_shared,
|
||||
)
|
||||
|
||||
for row in rows:
|
||||
claimed_ids.append(row["operation_id"])
|
||||
|
||||
@@ -104,46 +104,7 @@ class PostgreSQLOps(DataAccessOps):
|
||||
FROM input_data
|
||||
RETURNING id
|
||||
"""
|
||||
elif config.text_search_extension == "native":
|
||||
# search_vector is a regular tsvector column populated here using the
|
||||
# configured native dictionary. It used to be GENERATED ALWAYS with
|
||||
# a hardcoded 'english', which prevented per-deployment language
|
||||
# configuration. text_search_extension_native_language is validated
|
||||
# in HindsightConfig.validate() as a PG identifier, so embedding it
|
||||
# as a SQL literal is safe.
|
||||
query = f"""
|
||||
WITH input_data AS (
|
||||
SELECT * FROM unnest(
|
||||
$2::text[], $3::vector[], $4::timestamptz[], $5::timestamptz[], $6::timestamptz[], $7::timestamptz[],
|
||||
$8::text[], $9::text[], $10::jsonb[], $11::text[], $12::text[], $13::jsonb[], $14::jsonb[], $15::text[]
|
||||
) AS t(text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
context, fact_type, metadata, chunk_id, document_id, tags_json,
|
||||
observation_scopes_json, text_signals)
|
||||
)
|
||||
INSERT INTO {table} (bank_id, text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
context, fact_type, metadata, chunk_id, document_id, tags,
|
||||
observation_scopes, text_signals, search_vector)
|
||||
SELECT
|
||||
$1,
|
||||
text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
context, fact_type, metadata, chunk_id, document_id,
|
||||
COALESCE(
|
||||
(SELECT array_agg(elem) FROM jsonb_array_elements_text(tags_json) AS elem),
|
||||
'{{}}'::varchar[]
|
||||
),
|
||||
observation_scopes_json,
|
||||
text_signals,
|
||||
to_tsvector(
|
||||
'{config.text_search_extension_native_language}'::regconfig,
|
||||
COALESCE(text, '') || ' ' || COALESCE(context, '') || ' ' || COALESCE(text_signals, '')
|
||||
)
|
||||
FROM input_data
|
||||
RETURNING id
|
||||
"""
|
||||
else:
|
||||
# pg_textsearch, pgroonga, and pg_search: search_vector is a dummy
|
||||
# TEXT column; the actual full-text index operates on the base text
|
||||
# columns directly, so we don't populate search_vector at insert time.
|
||||
query = f"""
|
||||
WITH input_data AS (
|
||||
SELECT * FROM unnest(
|
||||
@@ -200,23 +161,6 @@ class PostgreSQLOps(DataAccessOps):
|
||||
exists_clause: str,
|
||||
chunk_size: int = 5000,
|
||||
) -> None:
|
||||
# exists_clause is unused on PostgreSQL: the memory_links → memory_units
|
||||
# FKs are DEFERRABLE INITIALLY DEFERRED, so an INSERT takes no lock on the
|
||||
# referenced parent rows until COMMIT — a concurrent committed DELETE in
|
||||
# that window (consolidation pruning observations, document re-tracking)
|
||||
# trips fk_memory_links_{to,from}_unit_id_memory_units at COMMIT (#1882),
|
||||
# and a WHERE EXISTS guard can't prevent it (the row passes the check,
|
||||
# then is deleted before the deferred check runs). Instead a CTE locks the
|
||||
# referenced units FOR KEY SHARE in the *same statement*: the lock blocks a
|
||||
# concurrent DELETE until our transaction commits and is held through the
|
||||
# deferred check, and the INSERT only takes links whose endpoints are in
|
||||
# the locked set, so rows that already vanished are dropped. Folding it
|
||||
# into the one INSERT keeps this to a single round-trip — no extra query
|
||||
# and no surrounding transaction needed. (Oracle's immediate FK has no
|
||||
# such window and uses exists_clause via its own bulk_insert_links.)
|
||||
from ..schema import fq_table
|
||||
|
||||
mu_table = fq_table("memory_units")
|
||||
from_ids = [lnk[0] for lnk in sorted_links]
|
||||
to_ids = [lnk[1] for lnk in sorted_links]
|
||||
types = [lnk[2] for lnk in sorted_links]
|
||||
@@ -225,37 +169,24 @@ class PostgreSQLOps(DataAccessOps):
|
||||
|
||||
for chunk_start in range(0, len(sorted_links), chunk_size):
|
||||
chunk_end = min(chunk_start + chunk_size, len(sorted_links))
|
||||
chunk_from = from_ids[chunk_start:chunk_end]
|
||||
chunk_to = to_ids[chunk_start:chunk_end]
|
||||
# Distinct referenced parents, sorted so concurrent inserters acquire
|
||||
# the row-share locks in a consistent order (avoids deadlocks; same
|
||||
# convention as the (from, to) link sort).
|
||||
referenced = sorted({str(x) for x in chunk_from} | {str(x) for x in chunk_to})
|
||||
await conn.execute(
|
||||
f"""
|
||||
WITH locked AS (
|
||||
SELECT id FROM {mu_table}
|
||||
WHERE id = ANY($7::uuid[])
|
||||
ORDER BY id
|
||||
FOR KEY SHARE
|
||||
)
|
||||
INSERT INTO {table}
|
||||
(from_unit_id, to_unit_id, link_type, weight, entity_id, bank_id)
|
||||
SELECT f, t, tp, w, e, $6
|
||||
FROM unnest($1::uuid[], $2::uuid[], $3::text[], $4::float8[], $5::uuid[])
|
||||
AS u(f, t, tp, w, e)
|
||||
WHERE f IN (SELECT id FROM locked) AND t IN (SELECT id FROM locked)
|
||||
AS t(f, t, tp, w, e)
|
||||
{exists_clause}
|
||||
ON CONFLICT (from_unit_id, to_unit_id, link_type,
|
||||
COALESCE(entity_id, '{nil_entity_uuid}'::uuid))
|
||||
DO NOTHING
|
||||
""",
|
||||
chunk_from,
|
||||
chunk_to,
|
||||
from_ids[chunk_start:chunk_end],
|
||||
to_ids[chunk_start:chunk_end],
|
||||
types[chunk_start:chunk_end],
|
||||
weights[chunk_start:chunk_end],
|
||||
entity_ids[chunk_start:chunk_end],
|
||||
bank_id,
|
||||
referenced,
|
||||
timeout=300,
|
||||
)
|
||||
|
||||
@@ -320,100 +251,28 @@ class PostgreSQLOps(DataAccessOps):
|
||||
entity_ids,
|
||||
)
|
||||
|
||||
async def enqueue_graph_maintenance(
|
||||
async def fetch_entity_unit_fanout(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
bank_id: str,
|
||||
unit_ids: list,
|
||||
) -> None:
|
||||
if not unit_ids:
|
||||
return
|
||||
await conn.execute(
|
||||
ue_table: str,
|
||||
entity_id_list: list[UUID],
|
||||
limit_per_entity: int,
|
||||
) -> list[ResultRow]:
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
INSERT INTO {table} (bank_id, unit_id)
|
||||
SELECT $1, v FROM unnest($2::uuid[]) AS t(v)
|
||||
ON CONFLICT (bank_id, unit_id) DO NOTHING
|
||||
""",
|
||||
bank_id,
|
||||
unit_ids,
|
||||
)
|
||||
|
||||
async def claim_graph_maintenance_batch(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
bank_id: str,
|
||||
limit: int,
|
||||
) -> list[str]:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
DELETE FROM {table}
|
||||
WHERE (bank_id, unit_id) IN (
|
||||
SELECT bank_id, unit_id FROM {table}
|
||||
WHERE bank_id = $1
|
||||
ORDER BY enqueued_at
|
||||
SELECT e.entity_id, n.unit_id
|
||||
FROM unnest($1::uuid[]) AS e(entity_id)
|
||||
CROSS JOIN LATERAL (
|
||||
SELECT ue.unit_id
|
||||
FROM {ue_table} ue
|
||||
WHERE ue.entity_id = e.entity_id
|
||||
ORDER BY ue.unit_id DESC
|
||||
LIMIT $2
|
||||
)
|
||||
RETURNING unit_id
|
||||
) n
|
||||
""",
|
||||
bank_id,
|
||||
limit,
|
||||
entity_id_list,
|
||||
limit_per_entity,
|
||||
)
|
||||
return [str(row["unit_id"]) for row in rows]
|
||||
|
||||
async def prune_orphan_entities(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
entities_table: str,
|
||||
ue_table: str,
|
||||
bank_id: str,
|
||||
) -> int:
|
||||
# Scoped by entities.bank_id (indexed). The NOT EXISTS subquery is
|
||||
# backed by idx_ue_entity on unit_entities(entity_id), so this stays
|
||||
# linear in the number of entities in the bank — not in the size of
|
||||
# unit_entities globally.
|
||||
result = await conn.execute(
|
||||
f"""
|
||||
DELETE FROM {entities_table} e
|
||||
WHERE e.bank_id = $1
|
||||
AND NOT EXISTS (
|
||||
SELECT 1 FROM {ue_table} ue WHERE ue.entity_id = e.id
|
||||
)
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
# asyncpg returns "DELETE N"
|
||||
return int(result.split()[-1]) if isinstance(result, str) and result.startswith("DELETE") else 0
|
||||
|
||||
async def prune_stale_cooccurrences(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
ec_table: str,
|
||||
ue_table: str,
|
||||
entities_table: str,
|
||||
bank_id: str,
|
||||
) -> int:
|
||||
# Scope by joining through entities.bank_id (entity_cooccurrences itself
|
||||
# has no bank_id column — entities don't span banks, so scoping via
|
||||
# entity_id_1 is sufficient).
|
||||
result = await conn.execute(
|
||||
f"""
|
||||
DELETE FROM {ec_table} c
|
||||
USING {entities_table} e
|
||||
WHERE e.id = c.entity_id_1
|
||||
AND e.bank_id = $1
|
||||
AND NOT EXISTS (
|
||||
SELECT 1
|
||||
FROM {ue_table} u1
|
||||
JOIN {ue_table} u2 ON u1.unit_id = u2.unit_id
|
||||
WHERE u1.entity_id = c.entity_id_1
|
||||
AND u2.entity_id = c.entity_id_2
|
||||
)
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
return int(result.split()[-1]) if isinstance(result, str) and result.startswith("DELETE") else 0
|
||||
|
||||
async def fetch_unit_dates(
|
||||
self,
|
||||
@@ -867,268 +726,7 @@ class PostgreSQLOps(DataAccessOps):
|
||||
|
||||
# -- Task claiming operations ------------------------------------------
|
||||
|
||||
async def _claim_consolidation_tasks(
|
||||
self,
|
||||
conn,
|
||||
table: str,
|
||||
busy_bank_ids: list[str],
|
||||
claimed_ids: list,
|
||||
limit: int,
|
||||
priority_map: dict[str, int] | None,
|
||||
) -> list:
|
||||
"""Claim consolidation tasks with optional priority-based tiered ordering.
|
||||
|
||||
When *priority_map* is ``None``, uses the default ``ORDER BY created_at``
|
||||
with bank-serialization (exclude busy banks). When set, claims in
|
||||
priority tiers — highest-priority banks first. Specific patterns always
|
||||
take precedence over the catch-all ``*`` entry.
|
||||
"""
|
||||
if limit <= 0:
|
||||
return []
|
||||
|
||||
# --- Fast path: no priority map -> current behavior ---
|
||||
if not priority_map:
|
||||
return await self._claim_consolidation_plain(conn, table, busy_bank_ids, claimed_ids, limit)
|
||||
|
||||
# --- Tiered claiming ---
|
||||
# Separate specific patterns from catch-all.
|
||||
# Specific patterns always take precedence: a bank matching ``shadow-*``
|
||||
# uses that entry's priority even if the catch-all ``*`` has a higher
|
||||
# value. The catch-all only applies to banks not matching any specific
|
||||
# pattern.
|
||||
specific_by_priority: dict[int, list[str]] = {}
|
||||
all_specific_sql: list[str] = []
|
||||
catch_all_priority = 1 # default when no ``*`` entry
|
||||
|
||||
for pattern, priority in priority_map.items():
|
||||
if pattern == "*":
|
||||
catch_all_priority = priority
|
||||
else:
|
||||
sql_pat = pattern.replace("*", "%")
|
||||
specific_by_priority.setdefault(priority, []).append(sql_pat)
|
||||
all_specific_sql.append(sql_pat)
|
||||
|
||||
# Collect all priority levels (specific tiers + catch-all) sorted desc.
|
||||
all_priorities = sorted(set(specific_by_priority.keys()) | {catch_all_priority}, reverse=True)
|
||||
|
||||
remaining = limit
|
||||
result: list = []
|
||||
|
||||
for pri in all_priorities:
|
||||
if remaining <= 0:
|
||||
break
|
||||
|
||||
# Specific-pattern tier at this priority level
|
||||
if pri in specific_by_priority:
|
||||
rows = await self._claim_consolidation_like(
|
||||
conn,
|
||||
table,
|
||||
busy_bank_ids,
|
||||
claimed_ids,
|
||||
remaining,
|
||||
specific_by_priority[pri],
|
||||
)
|
||||
for row in rows:
|
||||
claimed_ids.append(row["operation_id"])
|
||||
result.append(row)
|
||||
remaining -= len(rows)
|
||||
|
||||
# Catch-all tier at this priority level
|
||||
if pri == catch_all_priority and remaining > 0:
|
||||
rows = await self._claim_consolidation_not_like(
|
||||
conn,
|
||||
table,
|
||||
busy_bank_ids,
|
||||
claimed_ids,
|
||||
remaining,
|
||||
all_specific_sql,
|
||||
)
|
||||
for row in rows:
|
||||
claimed_ids.append(row["operation_id"])
|
||||
result.append(row)
|
||||
remaining -= len(rows)
|
||||
|
||||
return result
|
||||
|
||||
async def _claim_consolidation_plain(
|
||||
self,
|
||||
conn,
|
||||
table,
|
||||
busy_bank_ids,
|
||||
claimed_ids,
|
||||
limit,
|
||||
) -> list:
|
||||
"""Claim consolidation tasks with default created_at ordering."""
|
||||
exclude_ids = claimed_ids if claimed_ids else None
|
||||
if busy_bank_ids:
|
||||
if exclude_ids:
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND bank_id != ALL($1::text[])
|
||||
AND operation_id != ALL($2::uuid[])
|
||||
ORDER BY created_at
|
||||
LIMIT $3
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
busy_bank_ids,
|
||||
exclude_ids,
|
||||
limit,
|
||||
)
|
||||
else:
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND bank_id != ALL($1::text[])
|
||||
ORDER BY created_at
|
||||
LIMIT $2
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
busy_bank_ids,
|
||||
limit,
|
||||
)
|
||||
else:
|
||||
if exclude_ids:
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND operation_id != ALL($1::uuid[])
|
||||
ORDER BY created_at
|
||||
LIMIT $2
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
exclude_ids,
|
||||
limit,
|
||||
)
|
||||
else:
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
ORDER BY created_at
|
||||
LIMIT $1
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
limit,
|
||||
)
|
||||
|
||||
async def _claim_consolidation_like(
|
||||
self,
|
||||
conn,
|
||||
table,
|
||||
busy_bank_ids,
|
||||
claimed_ids,
|
||||
limit,
|
||||
sql_patterns,
|
||||
) -> list:
|
||||
"""Claim consolidation tasks from banks matching LIKE patterns."""
|
||||
params: list = [sql_patterns]
|
||||
conditions = ["bank_id LIKE ANY($1::text[])"]
|
||||
idx = 2
|
||||
|
||||
if busy_bank_ids:
|
||||
conditions.append(f"bank_id != ALL(${idx}::text[])")
|
||||
params.append(busy_bank_ids)
|
||||
idx += 1
|
||||
|
||||
if claimed_ids:
|
||||
conditions.append(f"operation_id != ALL(${idx}::uuid[])")
|
||||
params.append(claimed_ids)
|
||||
idx += 1
|
||||
|
||||
params.append(limit)
|
||||
extra = " AND ".join(conditions)
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND {extra}
|
||||
ORDER BY created_at
|
||||
LIMIT ${idx}
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
*params,
|
||||
)
|
||||
|
||||
async def _claim_consolidation_not_like(
|
||||
self,
|
||||
conn,
|
||||
table,
|
||||
busy_bank_ids,
|
||||
claimed_ids,
|
||||
limit,
|
||||
exclude_patterns,
|
||||
) -> list:
|
||||
"""Claim consolidation tasks from banks NOT matching any specific pattern (catch-all tier)."""
|
||||
params: list = []
|
||||
conditions: list[str] = []
|
||||
idx = 1
|
||||
|
||||
if exclude_patterns:
|
||||
conditions.append(f"bank_id NOT LIKE ALL(${idx}::text[])")
|
||||
params.append(exclude_patterns)
|
||||
idx += 1
|
||||
|
||||
if busy_bank_ids:
|
||||
conditions.append(f"bank_id != ALL(${idx}::text[])")
|
||||
params.append(busy_bank_ids)
|
||||
idx += 1
|
||||
|
||||
if claimed_ids:
|
||||
conditions.append(f"operation_id != ALL(${idx}::uuid[])")
|
||||
params.append(claimed_ids)
|
||||
idx += 1
|
||||
|
||||
params.append(limit)
|
||||
extra_clause = (" AND " + " AND ".join(conditions)) if conditions else ""
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW()){extra_clause}
|
||||
ORDER BY created_at
|
||||
LIMIT ${idx}
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
*params,
|
||||
)
|
||||
|
||||
async def claim_tasks(
|
||||
self,
|
||||
conn,
|
||||
table,
|
||||
worker_id,
|
||||
reserved_limits,
|
||||
shared_limit,
|
||||
*,
|
||||
consolidation_bank_priority=None,
|
||||
):
|
||||
async def claim_tasks(self, conn, table, worker_id, reserved_limits, shared_limit):
|
||||
all_rows = []
|
||||
claimed_ids = []
|
||||
|
||||
@@ -1146,14 +744,38 @@ class PostgreSQLOps(DataAccessOps):
|
||||
)
|
||||
busy_bank_ids = [r["bank_id"] for r in busy_banks]
|
||||
|
||||
rows = await self._claim_consolidation_tasks(
|
||||
conn,
|
||||
table,
|
||||
busy_bank_ids,
|
||||
claimed_ids,
|
||||
limit,
|
||||
consolidation_bank_priority,
|
||||
)
|
||||
if busy_bank_ids:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND bank_id != ALL($1::text[])
|
||||
ORDER BY created_at
|
||||
LIMIT $2
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
busy_bank_ids,
|
||||
limit,
|
||||
)
|
||||
else:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
ORDER BY created_at
|
||||
LIMIT $1
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
limit,
|
||||
)
|
||||
else:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
@@ -1217,7 +839,7 @@ class PostgreSQLOps(DataAccessOps):
|
||||
all_rows.append(row)
|
||||
remaining_shared -= len(rows)
|
||||
|
||||
# 2b. Consolidation tasks (with bank-serialization + optional priority)
|
||||
# 2b. Consolidation tasks (with bank-serialization)
|
||||
if remaining_shared > 0:
|
||||
busy_banks_2 = await conn.fetch(
|
||||
f"""
|
||||
@@ -1227,14 +849,76 @@ class PostgreSQLOps(DataAccessOps):
|
||||
)
|
||||
busy_bank_ids_2 = [r["bank_id"] for r in busy_banks_2]
|
||||
|
||||
rows = await self._claim_consolidation_tasks(
|
||||
conn,
|
||||
table,
|
||||
busy_bank_ids_2,
|
||||
claimed_ids,
|
||||
remaining_shared,
|
||||
consolidation_bank_priority,
|
||||
)
|
||||
if claimed_ids:
|
||||
if busy_bank_ids_2:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND operation_id != ALL($1::uuid[])
|
||||
AND bank_id != ALL($2::text[])
|
||||
ORDER BY created_at
|
||||
LIMIT $3
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
claimed_ids,
|
||||
busy_bank_ids_2,
|
||||
remaining_shared,
|
||||
)
|
||||
else:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND operation_id != ALL($1::uuid[])
|
||||
ORDER BY created_at
|
||||
LIMIT $2
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
claimed_ids,
|
||||
remaining_shared,
|
||||
)
|
||||
else:
|
||||
if busy_bank_ids_2:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND bank_id != ALL($1::text[])
|
||||
ORDER BY created_at
|
||||
LIMIT $2
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
busy_bank_ids_2,
|
||||
remaining_shared,
|
||||
)
|
||||
else:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
ORDER BY created_at
|
||||
LIMIT $1
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
remaining_shared,
|
||||
)
|
||||
|
||||
for row in rows:
|
||||
claimed_ids.append(row["operation_id"])
|
||||
|
||||
@@ -72,9 +72,6 @@ _RETURNING_RE = re.compile(r"\bRETURNING\s+(.+)", re.IGNORECASE | re.DOTALL)
|
||||
|
||||
_ANY_RE = re.compile(r"=\s*ANY\s*\(\s*:(\d+)\s*\)", re.IGNORECASE)
|
||||
_NOT_ALL_RE = re.compile(r"!=\s*ALL\s*\(\s*:(\d+)\s*\)", re.IGNORECASE)
|
||||
# LIKE ANY / NOT LIKE ALL — capture the column name before the operator
|
||||
_LIKE_ANY_RE = re.compile(r"(\w+)\s+LIKE\s+ANY\s*\(\s*:(\d+)\s*\)", re.IGNORECASE)
|
||||
_NOT_LIKE_ALL_RE = re.compile(r"(\w+)\s+NOT\s+LIKE\s+ALL\s*\(\s*:(\d+)\s*\)", re.IGNORECASE)
|
||||
|
||||
_JSON_ARROW_TEXT_RE = re.compile(r'("?\w+"?)\s*->>\s*\'(\w+)\'') # handles both col and "col"
|
||||
_JSON_HAS_KEY_RE = re.compile(r"(\w+)\s*\?\s*'(\w+)'")
|
||||
@@ -352,10 +349,6 @@ def _rewrite_pg_to_oracle(query: str) -> RewriteResult:
|
||||
# Boolean literals: Oracle uses NUMBER(1) for booleans
|
||||
query = re.sub(r"\b=\s*TRUE\b", "= 1", query, flags=re.IGNORECASE)
|
||||
query = re.sub(r"\b=\s*FALSE\b", "= 0", query, flags=re.IGNORECASE)
|
||||
# FOR NO KEY UPDATE → FOR UPDATE (Oracle has only FOR UPDATE; it does not block
|
||||
# indexed-FK child inserts the way PG's FOR UPDATE would, so plain FOR UPDATE is
|
||||
# the correct equivalent). Must run before the FOR SHARE rule below.
|
||||
query = re.sub(r"\bFOR\s+NO\s+KEY\s+UPDATE\b", "FOR UPDATE", query, flags=re.IGNORECASE)
|
||||
# FOR SHARE → FOR UPDATE (Oracle doesn't support FOR SHARE)
|
||||
query = re.sub(r"\bFOR\s+SHARE\b", "FOR UPDATE", query, flags=re.IGNORECASE)
|
||||
|
||||
@@ -542,12 +535,6 @@ def _rewrite_pg_to_oracle(query: str) -> RewriteResult:
|
||||
# != ALL(:N) → NOT IN (expanded list) — the negative counterpart of = ANY
|
||||
query = _NOT_ALL_RE.sub(r"NOT IN (/*EXPAND:\1*/)", query)
|
||||
|
||||
# col LIKE ANY(:N) → (col LIKE :p0 OR col LIKE :p1 OR ...)
|
||||
query = _LIKE_ANY_RE.sub(r"\1 /*LIKE_ANY:\2:\1*/", query)
|
||||
|
||||
# col NOT LIKE ALL(:N) → (col NOT LIKE :p0 AND col NOT LIKE :p1 AND ...)
|
||||
query = _NOT_LIKE_ALL_RE.sub(r"\1 /*NOT_LIKE_ALL:\2:\1*/", query)
|
||||
|
||||
# CTE AS MATERIALIZED (...) → AS (...) — Oracle doesn't support MATERIALIZED CTE hint
|
||||
query = re.sub(r"\bAS\s+MATERIALIZED\s*\(", "AS (", query, flags=re.IGNORECASE)
|
||||
|
||||
@@ -737,38 +724,17 @@ class OracleConnection(DatabaseConnection):
|
||||
|
||||
_expand_counter = 0
|
||||
|
||||
@staticmethod
|
||||
def _resolve_list_param(params: dict[str, Any], key: str) -> list | None:
|
||||
"""Resolve a parameter that may be a list or a JSON-encoded list string."""
|
||||
val = params.get(key)
|
||||
if isinstance(val, str):
|
||||
try:
|
||||
parsed = json.loads(val)
|
||||
if isinstance(parsed, list):
|
||||
return parsed
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
pass
|
||||
if isinstance(val, (list, tuple)):
|
||||
return list(val)
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _expand_any_lists(query: str, params: dict[str, Any] | None) -> tuple[str, dict[str, Any] | None]:
|
||||
"""Expand /*EXPAND:N*/, /*LIKE_ANY:N:col*/, /*NOT_LIKE_ALL:N:col*/ markers.
|
||||
"""Expand /*EXPAND:N*/ markers into individual bind vars for IN clauses.
|
||||
|
||||
Converts: IN (/*EXPAND:1*/) with params["1"] = [a, b, c]
|
||||
Into: IN (:any_0, :any_1, :any_2) with params["any_0"]=a, etc.
|
||||
|
||||
Converts: col /*LIKE_ANY:1:col*/ with params["1"] = [a, b]
|
||||
Into: (col LIKE :lk_0 OR col LIKE :lk_1)
|
||||
|
||||
Converts: col /*NOT_LIKE_ALL:1:col*/ with params["1"] = [a, b]
|
||||
Into: (col NOT LIKE :nlk_0 AND col NOT LIKE :nlk_1)
|
||||
|
||||
Uses a unique prefix to avoid name collisions with other bind vars.
|
||||
The original param is kept (for other references to :N in the query).
|
||||
"""
|
||||
if params is None or "/*" not in query:
|
||||
if params is None or "/*EXPAND:" not in query:
|
||||
return query, params
|
||||
|
||||
expand_re = re.compile(r"/\*EXPAND:(\d+)\*/")
|
||||
@@ -809,50 +775,6 @@ class OracleConnection(DatabaseConnection):
|
||||
|
||||
query = expand_re.sub(_replace, query)
|
||||
|
||||
# Expand LIKE ANY: col /*LIKE_ANY:N:col*/ → (col LIKE :p0 OR col LIKE :p1 ...)
|
||||
like_any_re = re.compile(r"(\w+)\s*/\*LIKE_ANY:(\d+):(\w+)\*/")
|
||||
|
||||
def _replace_like_any(m):
|
||||
_col = m.group(1) # redundant column ref before marker
|
||||
param_key = m.group(2)
|
||||
col = m.group(3)
|
||||
val = OracleConnection._resolve_list_param(params, param_key)
|
||||
if val is None or len(val) == 0:
|
||||
return "1=0" # no patterns → no match
|
||||
OracleConnection._expand_counter += 1
|
||||
prefix = f"lk{OracleConnection._expand_counter}"
|
||||
clauses = []
|
||||
for i, item in enumerate(val):
|
||||
k = f"{prefix}_{i}"
|
||||
params[k] = item
|
||||
clauses.append(f"{col} LIKE :{k}")
|
||||
keys_to_remove.add(param_key)
|
||||
return f"({' OR '.join(clauses)})"
|
||||
|
||||
query = like_any_re.sub(_replace_like_any, query)
|
||||
|
||||
# Expand NOT LIKE ALL: col /*NOT_LIKE_ALL:N:col*/ → (col NOT LIKE :p0 AND ...)
|
||||
not_like_all_re = re.compile(r"(\w+)\s*/\*NOT_LIKE_ALL:(\d+):(\w+)\*/")
|
||||
|
||||
def _replace_not_like_all(m):
|
||||
_col = m.group(1)
|
||||
param_key = m.group(2)
|
||||
col = m.group(3)
|
||||
val = OracleConnection._resolve_list_param(params, param_key)
|
||||
if val is None or len(val) == 0:
|
||||
return "1=1" # no patterns → everything matches
|
||||
OracleConnection._expand_counter += 1
|
||||
prefix = f"nlk{OracleConnection._expand_counter}"
|
||||
clauses = []
|
||||
for i, item in enumerate(val):
|
||||
k = f"{prefix}_{i}"
|
||||
params[k] = item
|
||||
clauses.append(f"{col} NOT LIKE :{k}")
|
||||
keys_to_remove.add(param_key)
|
||||
return f"({' AND '.join(clauses)})"
|
||||
|
||||
query = not_like_all_re.sub(_replace_not_like_all, query)
|
||||
|
||||
# Remove original list params that were expanded — their placeholder
|
||||
# (:N) no longer exists in the query, and leaving them causes DPY-4008.
|
||||
# Only remove if the key's placeholder is truly gone from the query.
|
||||
|
||||
@@ -6,7 +6,7 @@ import asyncio
|
||||
import logging
|
||||
import time
|
||||
from collections.abc import AsyncIterator
|
||||
from contextlib import AsyncExitStack, asynccontextmanager
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import Any
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -101,14 +101,6 @@ async def acquire_with_retry(backend_or_pool: Any, max_retries: int = DEFAULT_MA
|
||||
"""
|
||||
Async context manager to acquire a database connection with retry logic.
|
||||
|
||||
Retries the *acquire* itself when it raises a retryable error (connection
|
||||
drop, timeout, deadlock detected during acquire). Exceptions raised by
|
||||
user code inside the ``async with`` block are NOT retried — they propagate
|
||||
as-is. Wrapping retry around the yield would violate the
|
||||
``@asynccontextmanager`` single-yield contract and surface as
|
||||
``RuntimeError("generator didn't stop after athrow()")`` on every
|
||||
retryable inner error, masking the real cause.
|
||||
|
||||
Accepts either a DatabaseBackend or a raw asyncpg.Pool for backward compatibility.
|
||||
|
||||
Usage:
|
||||
@@ -117,7 +109,7 @@ async def acquire_with_retry(backend_or_pool: Any, max_retries: int = DEFAULT_MA
|
||||
|
||||
Args:
|
||||
backend_or_pool: A DatabaseBackend instance or asyncpg.Pool
|
||||
max_retries: Maximum number of retry attempts for the acquire step
|
||||
max_retries: Maximum number of retry attempts
|
||||
|
||||
Yields:
|
||||
A DatabaseConnection (if backend) or asyncpg.Connection (if pool)
|
||||
@@ -125,32 +117,31 @@ async def acquire_with_retry(backend_or_pool: Any, max_retries: int = DEFAULT_MA
|
||||
from .db.base import DatabaseBackend
|
||||
|
||||
if isinstance(backend_or_pool, DatabaseBackend) or getattr(backend_or_pool, "_wraps_backend", False):
|
||||
# Use the backend's acquire context manager with retry
|
||||
start = time.time()
|
||||
async with AsyncExitStack() as stack:
|
||||
conn: Any = None
|
||||
for attempt in range(max_retries + 1):
|
||||
try:
|
||||
conn = await stack.enter_async_context(backend_or_pool.acquire())
|
||||
break
|
||||
except Exception as e:
|
||||
if not _is_retryable(e):
|
||||
raise
|
||||
if attempt < max_retries:
|
||||
delay = min(DEFAULT_BASE_DELAY * (2**attempt), DEFAULT_MAX_DELAY)
|
||||
logger.warning(
|
||||
f"Database acquire failed (attempt {attempt + 1}/{max_retries + 1}): {e}. "
|
||||
f"Retrying in {delay:.1f}s..."
|
||||
)
|
||||
await asyncio.sleep(delay)
|
||||
else:
|
||||
logger.error(f"Database acquire failed after {max_retries + 1} attempts: {e}")
|
||||
raise
|
||||
|
||||
acquire_time = time.time() - start
|
||||
if acquire_time > 0.05:
|
||||
logger.warning(f"[DB POOL] Slow acquire: {acquire_time:.3f}s")
|
||||
|
||||
yield conn
|
||||
last_exception = None
|
||||
for attempt in range(max_retries + 1):
|
||||
try:
|
||||
async with backend_or_pool.acquire() as conn:
|
||||
acquire_time = time.time() - start
|
||||
if acquire_time > 0.05:
|
||||
logger.warning(f"[DB POOL] Slow acquire: {acquire_time:.3f}s")
|
||||
yield conn
|
||||
return
|
||||
except Exception as e:
|
||||
if not _is_retryable(e):
|
||||
raise
|
||||
last_exception = e
|
||||
if attempt < max_retries:
|
||||
delay = min(DEFAULT_BASE_DELAY * (2**attempt), DEFAULT_MAX_DELAY)
|
||||
logger.warning(
|
||||
f"Database acquire failed (attempt {attempt + 1}/{max_retries + 1}): {e}. "
|
||||
f"Retrying in {delay:.1f}s..."
|
||||
)
|
||||
await asyncio.sleep(delay)
|
||||
else:
|
||||
logger.error(f"Database acquire failed after {max_retries + 1} attempts: {e}")
|
||||
raise last_exception
|
||||
else:
|
||||
# Legacy path: raw asyncpg.Pool
|
||||
pool = backend_or_pool
|
||||
|
||||
@@ -9,17 +9,13 @@ The database schema is automatically adjusted to match the model's dimension.
|
||||
Configuration via environment variables - see hindsight_api.config for all env var names.
|
||||
"""
|
||||
|
||||
import base64
|
||||
import logging
|
||||
import os
|
||||
import struct
|
||||
import warnings
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Literal, cast
|
||||
from urllib.parse import parse_qs, urlparse, urlunparse
|
||||
|
||||
import httpx
|
||||
from pydantic import BaseModel
|
||||
|
||||
from ..config import (
|
||||
DEFAULT_EMBEDDINGS_COHERE_MODEL,
|
||||
@@ -31,15 +27,10 @@ from ..config import (
|
||||
DEFAULT_EMBEDDINGS_LOCAL_TRUST_REMOTE_CODE,
|
||||
DEFAULT_EMBEDDINGS_OPENAI_MODEL,
|
||||
DEFAULT_EMBEDDINGS_PROVIDER,
|
||||
DEFAULT_EMBEDDINGS_ZEROENTROPY_BATCH_SIZE,
|
||||
DEFAULT_EMBEDDINGS_ZEROENTROPY_DIMENSIONS,
|
||||
DEFAULT_EMBEDDINGS_ZEROENTROPY_ENCODING_FORMAT,
|
||||
DEFAULT_EMBEDDINGS_ZEROENTROPY_LATENCY,
|
||||
DEFAULT_EMBEDDINGS_ZEROENTROPY_MODEL,
|
||||
DEFAULT_LITELLM_API_BASE,
|
||||
DEFAULT_ZEROENTROPY_BASE_URL,
|
||||
ENV_EMBEDDINGS_COHERE_API_KEY,
|
||||
ENV_EMBEDDINGS_GEMINI_API_KEY,
|
||||
ENV_EMBEDDINGS_LITELLM_SDK_API_KEY,
|
||||
ENV_EMBEDDINGS_LOCAL_FORCE_CPU,
|
||||
ENV_EMBEDDINGS_LOCAL_MODEL,
|
||||
ENV_EMBEDDINGS_LOCAL_TRUST_REMOTE_CODE,
|
||||
@@ -48,39 +39,12 @@ from ..config import (
|
||||
ENV_EMBEDDINGS_OPENAI_MODEL,
|
||||
ENV_EMBEDDINGS_PROVIDER,
|
||||
ENV_EMBEDDINGS_TEI_URL,
|
||||
ENV_EMBEDDINGS_ZEROENTROPY_API_KEY,
|
||||
ENV_EMBEDDINGS_ZEROENTROPY_DIMENSIONS,
|
||||
ENV_EMBEDDINGS_ZEROENTROPY_ENCODING_FORMAT,
|
||||
ENV_LLM_API_KEY,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
ZeroEntropyInputType = Literal["document", "query"]
|
||||
ZeroEntropyLatency = Literal["fast", "slow"]
|
||||
ZeroEntropyEncodingFormat = Literal["float", "base64"]
|
||||
|
||||
|
||||
class _ZeroEntropyEmbedRequest(BaseModel):
|
||||
"""Typed request body for ZeroEntropy's non-OpenAI-compatible embed endpoint."""
|
||||
|
||||
model: str
|
||||
input: list[str]
|
||||
input_type: ZeroEntropyInputType
|
||||
dimensions: int
|
||||
encoding_format: ZeroEntropyEncodingFormat = "float"
|
||||
latency: ZeroEntropyLatency | None = None
|
||||
|
||||
|
||||
class _ZeroEntropyEmbedResult(BaseModel):
|
||||
embedding: list[float] | str
|
||||
|
||||
|
||||
class _ZeroEntropyEmbedResponse(BaseModel):
|
||||
results: list[_ZeroEntropyEmbedResult]
|
||||
|
||||
|
||||
class Embeddings(ABC):
|
||||
"""
|
||||
Abstract base class for embedding generation.
|
||||
@@ -124,14 +88,6 @@ class Embeddings(ABC):
|
||||
"""
|
||||
pass
|
||||
|
||||
def encode_query(self, texts: list[str]) -> list[list[float]]:
|
||||
"""Generate embeddings for query text. Providers without asymmetric embeddings use encode()."""
|
||||
return self.encode(texts)
|
||||
|
||||
def encode_documents(self, texts: list[str]) -> list[list[float]]:
|
||||
"""Generate embeddings for stored document text. Providers without asymmetric embeddings use encode()."""
|
||||
return self.encode(texts)
|
||||
|
||||
|
||||
class LocalSTEmbeddings(Embeddings):
|
||||
"""
|
||||
@@ -429,7 +385,6 @@ class OpenAIEmbeddings(Embeddings):
|
||||
model: str = DEFAULT_EMBEDDINGS_OPENAI_MODEL,
|
||||
base_url: str | None = None,
|
||||
batch_size: int = 100,
|
||||
dimensions: int | None = None,
|
||||
max_retries: int = 3,
|
||||
):
|
||||
"""
|
||||
@@ -440,14 +395,12 @@ class OpenAIEmbeddings(Embeddings):
|
||||
model: OpenAI embedding model name (default: text-embedding-3-small)
|
||||
base_url: Custom base URL for OpenAI-compatible API (e.g., Azure OpenAI endpoint)
|
||||
batch_size: Maximum batch size for embedding requests (default: 100)
|
||||
dimensions: Optional requested output dimensions for OpenAI text-embedding-3 models
|
||||
max_retries: Maximum number of retries for failed requests (default: 3)
|
||||
"""
|
||||
self.api_key = api_key
|
||||
self.model = model
|
||||
self.base_url = base_url
|
||||
self.batch_size = batch_size
|
||||
self.dimensions = dimensions
|
||||
self.max_retries = max_retries
|
||||
self._client = None
|
||||
self._dimension: int | None = None
|
||||
@@ -492,9 +445,7 @@ class OpenAIEmbeddings(Embeddings):
|
||||
self._client = OpenAI(**client_kwargs)
|
||||
|
||||
# Try to get dimension from known models, otherwise do a test embedding
|
||||
if self.dimensions is not None:
|
||||
self._dimension = self.dimensions
|
||||
elif self.model in self.MODEL_DIMENSIONS:
|
||||
if self.model in self.MODEL_DIMENSIONS:
|
||||
self._dimension = self.MODEL_DIMENSIONS[self.model]
|
||||
else:
|
||||
# Do a test embedding to detect dimension
|
||||
@@ -529,14 +480,10 @@ class OpenAIEmbeddings(Embeddings):
|
||||
for i in range(0, len(texts), self.batch_size):
|
||||
batch = texts[i : i + self.batch_size]
|
||||
|
||||
request = {
|
||||
"model": self.model,
|
||||
"input": batch,
|
||||
}
|
||||
if self.dimensions is not None:
|
||||
request["dimensions"] = self.dimensions
|
||||
|
||||
response = self._client.embeddings.create(**request)
|
||||
response = self._client.embeddings.create(
|
||||
model=self.model,
|
||||
input=batch,
|
||||
)
|
||||
|
||||
# Sort by index to ensure correct order
|
||||
batch_embeddings = sorted(response.data, key=lambda x: x.index)
|
||||
@@ -545,73 +492,6 @@ class OpenAIEmbeddings(Embeddings):
|
||||
return all_embeddings
|
||||
|
||||
|
||||
class CodexOAuthEmbeddings(OpenAIEmbeddings):
|
||||
"""
|
||||
OpenAI embeddings using the Codex/ChatGPT OAuth token from ``~/.codex/auth.json``.
|
||||
|
||||
Codex OAuth is an LLM-provider auth path in Hindsight, but the same bearer token
|
||||
can also authenticate against the standard OpenAI embeddings endpoint. This keeps
|
||||
embeddings on the user's existing Codex subscription/OAuth path without requiring
|
||||
a separate OpenAI/OpenRouter/Gemini/Cohere API key.
|
||||
|
||||
Token refresh is handled automatically: the manager proactively refreshes the
|
||||
access_token before it expires and reactively refreshes on 401 responses from
|
||||
the embeddings API.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model: str = DEFAULT_EMBEDDINGS_OPENAI_MODEL,
|
||||
batch_size: int = 100,
|
||||
dimensions: int | None = None,
|
||||
max_retries: int = 3,
|
||||
):
|
||||
from .providers.codex_auth import CodexAuthManager
|
||||
|
||||
self._auth_manager = CodexAuthManager.from_file()
|
||||
super().__init__(
|
||||
api_key=self._auth_manager.access_token,
|
||||
model=model,
|
||||
base_url="https://api.openai.com/v1",
|
||||
batch_size=batch_size,
|
||||
dimensions=dimensions,
|
||||
max_retries=max_retries,
|
||||
)
|
||||
|
||||
@property
|
||||
def provider_name(self) -> str:
|
||||
return "openai-codex"
|
||||
|
||||
def encode(self, texts: list[str]) -> list[list[float]]:
|
||||
"""Generate embeddings, refreshing the OAuth token if needed.
|
||||
|
||||
Proactively refreshes before the call when the token is near expiry,
|
||||
and reactively refreshes once on a 401 from the OpenAI embeddings API.
|
||||
"""
|
||||
from openai import AuthenticationError
|
||||
|
||||
# Proactive refresh — cheap when fresh (JWT exp decode + compare).
|
||||
self._auth_manager.ensure_fresh_token()
|
||||
if self._auth_manager.access_token != self.api_key:
|
||||
self.api_key = self._auth_manager.access_token
|
||||
if self._client is not None:
|
||||
self._client.api_key = self._auth_manager.access_token
|
||||
|
||||
try:
|
||||
return super().encode(texts)
|
||||
except AuthenticationError:
|
||||
# Reactive refresh — token was valid by the JWT clock but the
|
||||
# server rejected it (rotated server-side, race, etc.).
|
||||
self._auth_manager.refresh_tokens(
|
||||
reason="reactive (401 from embeddings API)",
|
||||
force=True,
|
||||
)
|
||||
self.api_key = self._auth_manager.access_token
|
||||
if self._client is not None:
|
||||
self._client.api_key = self._auth_manager.access_token
|
||||
return super().encode(texts)
|
||||
|
||||
|
||||
class CohereEmbeddings(Embeddings):
|
||||
"""
|
||||
Cohere embeddings implementation using the Cohere API.
|
||||
@@ -753,149 +633,6 @@ class CohereEmbeddings(Embeddings):
|
||||
return all_embeddings
|
||||
|
||||
|
||||
class ZeroEntropyEmbeddings(Embeddings):
|
||||
"""
|
||||
ZeroEntropy embeddings implementation using the zembed API.
|
||||
|
||||
ZeroEntropy's embeddings endpoint is not OpenAI-compatible: it lives at
|
||||
/v1/models/embed and requires provider-specific parameters such as
|
||||
input_type. Hindsight stores document-side vectors and uses query-side
|
||||
vectors during recall, so this provider exposes explicit encode_documents()
|
||||
and encode_query() helpers while keeping encode() as document-side default.
|
||||
"""
|
||||
|
||||
VALID_DIMENSIONS = frozenset({2560, 1280, 640, 320, 160, 80, 40})
|
||||
VALID_ENCODING_FORMATS = frozenset({"float", "base64"})
|
||||
VALID_LATENCIES = frozenset({"fast", "slow"})
|
||||
DEFAULT_BASE_URL = DEFAULT_ZEROENTROPY_BASE_URL
|
||||
EMBED_PATH = "/v1/models/embed"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
api_key: str,
|
||||
model: str = DEFAULT_EMBEDDINGS_ZEROENTROPY_MODEL,
|
||||
base_url: str | None = None,
|
||||
dimensions: int = DEFAULT_EMBEDDINGS_ZEROENTROPY_DIMENSIONS,
|
||||
batch_size: int = DEFAULT_EMBEDDINGS_ZEROENTROPY_BATCH_SIZE,
|
||||
encoding_format: str = DEFAULT_EMBEDDINGS_ZEROENTROPY_ENCODING_FORMAT,
|
||||
latency: str | None = DEFAULT_EMBEDDINGS_ZEROENTROPY_LATENCY,
|
||||
timeout: float = 60.0,
|
||||
):
|
||||
if dimensions not in self.VALID_DIMENSIONS:
|
||||
valid = ", ".join(str(dim) for dim in sorted(self.VALID_DIMENSIONS, reverse=True))
|
||||
raise ValueError(f"{ENV_EMBEDDINGS_ZEROENTROPY_DIMENSIONS} must be one of {valid}, got {dimensions}")
|
||||
if batch_size < 1:
|
||||
raise ValueError("ZeroEntropy embeddings batch_size must be >= 1")
|
||||
if encoding_format not in self.VALID_ENCODING_FORMATS:
|
||||
valid_formats = ", ".join(sorted(self.VALID_ENCODING_FORMATS))
|
||||
raise ValueError(
|
||||
f"{ENV_EMBEDDINGS_ZEROENTROPY_ENCODING_FORMAT} must be one of {valid_formats}, got {encoding_format!r}"
|
||||
)
|
||||
if latency is not None and latency not in self.VALID_LATENCIES:
|
||||
valid_latencies = ", ".join(sorted(self.VALID_LATENCIES))
|
||||
raise ValueError(f"ZeroEntropy embeddings latency must be one of {valid_latencies}, got {latency!r}")
|
||||
|
||||
self.api_key = api_key
|
||||
self.model = model
|
||||
self.base_url = base_url.rstrip("/") if base_url else self.DEFAULT_BASE_URL
|
||||
self.embed_url = f"{self.base_url}{self.EMBED_PATH}"
|
||||
self.dimensions = dimensions
|
||||
self.batch_size = batch_size
|
||||
self.encoding_format = cast(ZeroEntropyEncodingFormat, encoding_format)
|
||||
self.latency = cast(ZeroEntropyLatency | None, latency)
|
||||
self.timeout = timeout
|
||||
self._client: httpx.Client | None = None
|
||||
self._dimension: int | None = None
|
||||
|
||||
@property
|
||||
def provider_name(self) -> str:
|
||||
return "zeroentropy"
|
||||
|
||||
@property
|
||||
def dimension(self) -> int:
|
||||
if self._dimension is None:
|
||||
raise RuntimeError("Embeddings not initialized. Call initialize() first.")
|
||||
return self._dimension
|
||||
|
||||
async def initialize(self) -> None:
|
||||
"""Initialize the ZeroEntropy HTTP client."""
|
||||
if self._client is not None:
|
||||
return
|
||||
|
||||
logger.info(
|
||||
f"Embeddings: initializing ZeroEntropy provider with model {self.model} "
|
||||
f"(dim: {self.dimensions}, batch_size={self.batch_size})"
|
||||
)
|
||||
self._client = httpx.Client(
|
||||
timeout=self.timeout,
|
||||
headers={
|
||||
"Authorization": f"Bearer {self.api_key}",
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
)
|
||||
# zembed-1 dimensions are explicit Matryoshka truncation steps. Avoid a
|
||||
# startup probe so boot does not burn quota or require a throwaway input.
|
||||
self._dimension = self.dimensions
|
||||
logger.info(f"Embeddings: ZeroEntropy provider initialized (model: {self.model}, dim: {self._dimension})")
|
||||
|
||||
def encode(self, texts: list[str]) -> list[list[float]]:
|
||||
"""Generate document-side embeddings for backwards-compatible callers."""
|
||||
return self.encode_documents(texts)
|
||||
|
||||
def encode_documents(self, texts: list[str]) -> list[list[float]]:
|
||||
"""Generate document-side embeddings for retained content."""
|
||||
return self._encode_with_input_type(texts, "document")
|
||||
|
||||
def encode_query(self, texts: list[str]) -> list[list[float]]:
|
||||
"""Generate query-side embeddings for recall/search queries."""
|
||||
return self._encode_with_input_type(texts, "query")
|
||||
|
||||
def _encode_with_input_type(self, texts: list[str], input_type: ZeroEntropyInputType) -> list[list[float]]:
|
||||
if self._client is None:
|
||||
raise RuntimeError("Embeddings not initialized. Call initialize() first.")
|
||||
|
||||
if not texts:
|
||||
return []
|
||||
|
||||
all_embeddings: list[list[float]] = []
|
||||
for i in range(0, len(texts), self.batch_size):
|
||||
batch = texts[i : i + self.batch_size]
|
||||
request = _ZeroEntropyEmbedRequest(
|
||||
model=self.model,
|
||||
input=batch,
|
||||
input_type=input_type,
|
||||
dimensions=self.dimensions,
|
||||
encoding_format=self.encoding_format,
|
||||
latency=self.latency,
|
||||
)
|
||||
|
||||
try:
|
||||
response = self._client.post(self.embed_url, json=request.model_dump(exclude_none=True))
|
||||
response.raise_for_status()
|
||||
except httpx.HTTPError as e:
|
||||
raise RuntimeError(f"ZeroEntropy embedding request failed: {e}") from e
|
||||
|
||||
parsed = _ZeroEntropyEmbedResponse.model_validate(response.json())
|
||||
if len(parsed.results) != len(batch):
|
||||
raise RuntimeError(
|
||||
f"ZeroEntropy returned {len(parsed.results)} embeddings for {len(batch)} input texts; "
|
||||
"expected exact 1:1 alignment"
|
||||
)
|
||||
all_embeddings.extend(self._parse_embedding(result.embedding) for result in parsed.results)
|
||||
|
||||
return all_embeddings
|
||||
|
||||
@staticmethod
|
||||
def _parse_embedding(embedding: list[float] | str) -> list[float]:
|
||||
if not isinstance(embedding, str):
|
||||
return embedding
|
||||
|
||||
raw = base64.b64decode(embedding)
|
||||
if len(raw) % 4 != 0:
|
||||
raise RuntimeError("ZeroEntropy returned invalid base64 embedding length")
|
||||
return list(struct.unpack(f"<{len(raw) // 4}f", raw))
|
||||
|
||||
|
||||
class LiteLLMEmbeddings(Embeddings):
|
||||
"""
|
||||
LiteLLM embeddings implementation using LiteLLM proxy's /embeddings endpoint.
|
||||
@@ -1029,7 +766,7 @@ class LiteLLMSDKEmbeddings(Embeddings):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
api_key: str | None = None,
|
||||
api_key: str,
|
||||
model: str = DEFAULT_EMBEDDINGS_LITELLM_SDK_MODEL,
|
||||
api_base: str | None = None,
|
||||
output_dimensions: int | None = None,
|
||||
@@ -1041,8 +778,7 @@ class LiteLLMSDKEmbeddings(Embeddings):
|
||||
Initialize LiteLLM SDK embeddings client.
|
||||
|
||||
Args:
|
||||
api_key: API key for the embedding provider (optional — omit for
|
||||
providers that use ambient credentials, e.g. AWS Bedrock with IAM)
|
||||
api_key: API key for the embedding provider
|
||||
model: Model name with provider prefix (e.g., "cohere/embed-english-v3.0")
|
||||
api_base: Custom base URL for API (optional)
|
||||
output_dimensions: Optional output embedding dimensions (provider-dependent)
|
||||
@@ -1092,9 +828,8 @@ class LiteLLMSDKEmbeddings(Embeddings):
|
||||
embed_kwargs = {
|
||||
"model": self.model,
|
||||
"input": ["test"],
|
||||
"api_key": self.api_key,
|
||||
}
|
||||
if self.api_key:
|
||||
embed_kwargs["api_key"] = self.api_key
|
||||
if self.encoding_format:
|
||||
embed_kwargs["encoding_format"] = self.encoding_format
|
||||
if self.api_base:
|
||||
@@ -1145,9 +880,8 @@ class LiteLLMSDKEmbeddings(Embeddings):
|
||||
embed_kwargs = {
|
||||
"model": self.model,
|
||||
"input": batch,
|
||||
"api_key": self.api_key,
|
||||
}
|
||||
if self.api_key:
|
||||
embed_kwargs["api_key"] = self.api_key
|
||||
if self.encoding_format:
|
||||
embed_kwargs["encoding_format"] = self.encoding_format
|
||||
if self.api_base:
|
||||
@@ -1406,14 +1140,6 @@ def create_embeddings_from_env() -> Embeddings:
|
||||
model=model,
|
||||
base_url=base_url,
|
||||
batch_size=config.embeddings_openai_batch_size,
|
||||
dimensions=config.embeddings_openai_dimensions,
|
||||
)
|
||||
elif provider == "openai-codex":
|
||||
model = os.environ.get(ENV_EMBEDDINGS_OPENAI_MODEL, DEFAULT_EMBEDDINGS_OPENAI_MODEL)
|
||||
return CodexOAuthEmbeddings(
|
||||
model=model,
|
||||
batch_size=config.embeddings_openai_batch_size,
|
||||
dimensions=config.embeddings_openai_dimensions,
|
||||
)
|
||||
elif provider == "openrouter":
|
||||
api_key = config.embeddings_openrouter_api_key
|
||||
@@ -1427,23 +1153,6 @@ def create_embeddings_from_env() -> Embeddings:
|
||||
model=config.embeddings_openrouter_model,
|
||||
base_url="https://openrouter.ai/api/v1",
|
||||
batch_size=config.embeddings_openai_batch_size,
|
||||
dimensions=config.embeddings_openai_dimensions,
|
||||
)
|
||||
elif provider == "zeroentropy":
|
||||
api_key = config.embeddings_zeroentropy_api_key
|
||||
if not api_key:
|
||||
raise ValueError(
|
||||
f"{ENV_EMBEDDINGS_ZEROENTROPY_API_KEY} or ZEROENTROPY_API_KEY is required "
|
||||
f"when {ENV_EMBEDDINGS_PROVIDER} is 'zeroentropy'"
|
||||
)
|
||||
return ZeroEntropyEmbeddings(
|
||||
api_key=api_key,
|
||||
model=config.embeddings_zeroentropy_model,
|
||||
base_url=config.embeddings_zeroentropy_base_url,
|
||||
dimensions=config.embeddings_zeroentropy_dimensions,
|
||||
batch_size=config.embeddings_zeroentropy_batch_size,
|
||||
encoding_format=config.embeddings_zeroentropy_encoding_format,
|
||||
latency=config.embeddings_zeroentropy_latency,
|
||||
)
|
||||
elif provider == "cohere":
|
||||
api_key = config.embeddings_cohere_api_key
|
||||
@@ -1462,8 +1171,13 @@ def create_embeddings_from_env() -> Embeddings:
|
||||
model=config.embeddings_litellm_model,
|
||||
)
|
||||
elif provider == "litellm-sdk":
|
||||
api_key = config.embeddings_litellm_sdk_api_key
|
||||
if not api_key:
|
||||
raise ValueError(
|
||||
f"{ENV_EMBEDDINGS_LITELLM_SDK_API_KEY} is required when {ENV_EMBEDDINGS_PROVIDER} is 'litellm-sdk'"
|
||||
)
|
||||
return LiteLLMSDKEmbeddings(
|
||||
api_key=config.embeddings_litellm_sdk_api_key or None,
|
||||
api_key=api_key,
|
||||
model=config.embeddings_litellm_sdk_model,
|
||||
api_base=config.embeddings_litellm_sdk_api_base,
|
||||
output_dimensions=config.embeddings_litellm_sdk_output_dimensions,
|
||||
@@ -1492,6 +1206,5 @@ def create_embeddings_from_env() -> Embeddings:
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unknown embeddings provider: {provider}. "
|
||||
f"Supported: 'local', 'tei', 'openai', 'openai-codex', 'openrouter', 'cohere', 'google', "
|
||||
f"'zeroentropy', 'litellm', 'litellm-sdk'"
|
||||
f"Supported: 'local', 'tei', 'openai', 'cohere', 'google', 'litellm', 'litellm-sdk'"
|
||||
)
|
||||
|
||||
@@ -16,15 +16,7 @@ from typing import Any, Final
|
||||
|
||||
from .db_utils import acquire_with_retry
|
||||
from .memory_engine import fq_table
|
||||
from .retain.entity_labels import (
|
||||
build_labels_lookup as _build_labels_lookup_from_config,
|
||||
)
|
||||
from .retain.entity_labels import (
|
||||
is_label_entity as _is_label_entity,
|
||||
)
|
||||
from .retain.entity_labels import (
|
||||
parse_entity_labels as _parse_entity_labels,
|
||||
)
|
||||
from .retain.entity_labels import build_labels_lookup as _build_labels_lookup_from_config
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -97,12 +89,7 @@ class EntityResolver:
|
||||
Resolves entities to canonical IDs with disambiguation.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
pool: Any,
|
||||
entity_lookup: str = "full",
|
||||
entity_resolution_batch_size: int = 100,
|
||||
):
|
||||
def __init__(self, pool: Any, entity_lookup: str = "full"):
|
||||
"""
|
||||
Initialize entity resolver.
|
||||
|
||||
@@ -111,14 +98,9 @@ class EntityResolver:
|
||||
entity_lookup: Lookup strategy — "full" loads all bank entities then
|
||||
matches in Python; "trigram" uses pg_trgm GIN index to fetch only
|
||||
similar candidates per entity name (much faster for large banks).
|
||||
entity_resolution_batch_size: Number of unique entity names to include
|
||||
in each pg_trgm candidate lookup query.
|
||||
"""
|
||||
self.pool = pool
|
||||
self.entity_lookup = entity_lookup
|
||||
if entity_resolution_batch_size < 1:
|
||||
raise ValueError("entity_resolution_batch_size must be >= 1")
|
||||
self.entity_resolution_batch_size = entity_resolution_batch_size
|
||||
self._pg_trgm_checked = False
|
||||
# Backend-specific operations — accessed via pool.ops (Django pattern).
|
||||
self._ops = pool.ops if pool is not None else None
|
||||
@@ -217,11 +199,6 @@ class EntityResolver:
|
||||
"""Build a set of valid 'key:value' entity label strings for fast lookup."""
|
||||
return _build_labels_lookup_from_config(entity_labels)
|
||||
|
||||
@staticmethod
|
||||
def _chunked(values: list[str], size: int) -> list[list[str]]:
|
||||
"""Split values into fixed-size batches."""
|
||||
return [values[i : i + size] for i in range(0, len(values), size)]
|
||||
|
||||
async def resolve_entities_batch(
|
||||
self,
|
||||
bank_id: str,
|
||||
@@ -251,15 +228,14 @@ class EntityResolver:
|
||||
return []
|
||||
|
||||
taxonomy_lookup = self._build_labels_lookup(entity_labels)
|
||||
labels_cfg = _parse_entity_labels(entity_labels)
|
||||
if conn is None:
|
||||
async with acquire_with_retry(self.pool) as conn:
|
||||
return await self._resolve_entities_batch_impl(
|
||||
conn, bank_id, entities_data, context, unit_event_date, taxonomy_lookup, labels_cfg
|
||||
conn, bank_id, entities_data, context, unit_event_date, taxonomy_lookup
|
||||
)
|
||||
else:
|
||||
return await self._resolve_entities_batch_impl(
|
||||
conn, bank_id, entities_data, context, unit_event_date, taxonomy_lookup, labels_cfg
|
||||
conn, bank_id, entities_data, context, unit_event_date, taxonomy_lookup
|
||||
)
|
||||
|
||||
async def _resolve_entities_batch_impl(
|
||||
@@ -270,16 +246,13 @@ class EntityResolver:
|
||||
context: str,
|
||||
unit_event_date,
|
||||
taxonomy_lookup: set[str] | None = None,
|
||||
labels_cfg=None,
|
||||
) -> list[str]:
|
||||
if self.entity_lookup == "trigram":
|
||||
# Route to backend-specific fuzzy strategy.
|
||||
# Non-PG backends (Oracle) use UTL_MATCH instead of pg_trgm.
|
||||
backend_strategy = self._ops.get_entity_resolution_strategy()
|
||||
if backend_strategy == "oracle_fuzzy":
|
||||
return await self._resolve_entities_batch_oracle_fuzzy(
|
||||
conn, bank_id, entities_data, unit_event_date, taxonomy_lookup, labels_cfg
|
||||
)
|
||||
return await self._resolve_entities_batch_oracle_fuzzy(conn, bank_id, entities_data, unit_event_date)
|
||||
# Auto-detect pg_trgm availability on first call and fall back to
|
||||
# "full" strategy if the extension is not installed. See #626.
|
||||
if not self._pg_trgm_checked:
|
||||
@@ -293,24 +266,12 @@ class EntityResolver:
|
||||
"https://github.com/vectorize-io/hindsight/issues/626"
|
||||
)
|
||||
self.entity_lookup = "full"
|
||||
return await self._resolve_entities_batch_full(
|
||||
conn, bank_id, entities_data, unit_event_date, taxonomy_lookup, labels_cfg
|
||||
)
|
||||
return await self._resolve_entities_batch_trigram(
|
||||
conn, bank_id, entities_data, unit_event_date, taxonomy_lookup, labels_cfg
|
||||
)
|
||||
return await self._resolve_entities_batch_full(
|
||||
conn, bank_id, entities_data, unit_event_date, taxonomy_lookup, labels_cfg
|
||||
)
|
||||
return await self._resolve_entities_batch_full(conn, bank_id, entities_data, unit_event_date)
|
||||
return await self._resolve_entities_batch_trigram(conn, bank_id, entities_data, unit_event_date)
|
||||
return await self._resolve_entities_batch_full(conn, bank_id, entities_data, unit_event_date)
|
||||
|
||||
async def _resolve_entities_batch_full(
|
||||
self,
|
||||
conn,
|
||||
bank_id: str,
|
||||
entities_data: list[dict],
|
||||
unit_event_date,
|
||||
taxonomy_lookup: set[str] | None = None,
|
||||
labels_cfg=None,
|
||||
self, conn, bank_id: str, entities_data: list[dict], unit_event_date
|
||||
) -> list[str]:
|
||||
"""Original strategy: load all bank entities then match in Python."""
|
||||
# Query ALL candidates for this bank
|
||||
@@ -377,24 +338,11 @@ class EntityResolver:
|
||||
all_candidates[entity_text] = matching
|
||||
|
||||
return await self._resolve_from_candidates(
|
||||
conn,
|
||||
bank_id,
|
||||
entities_data,
|
||||
unit_event_date,
|
||||
all_candidates,
|
||||
cooccurrence_map,
|
||||
taxonomy_lookup,
|
||||
labels_cfg,
|
||||
conn, bank_id, entities_data, unit_event_date, all_candidates, cooccurrence_map
|
||||
)
|
||||
|
||||
async def _resolve_entities_batch_trigram(
|
||||
self,
|
||||
conn,
|
||||
bank_id: str,
|
||||
entities_data: list[dict],
|
||||
unit_event_date,
|
||||
taxonomy_lookup: set[str] | None = None,
|
||||
labels_cfg=None,
|
||||
self, conn, bank_id: str, entities_data: list[dict], unit_event_date
|
||||
) -> list[str]:
|
||||
"""
|
||||
Trigram strategy: fetch only similar candidates per entity name using pg_trgm.
|
||||
@@ -405,7 +353,7 @@ class EntityResolver:
|
||||
"""
|
||||
entity_texts = list(set(e["text"] for e in entities_data))
|
||||
|
||||
# Fetch candidates for unique entity texts in bounded batches.
|
||||
# Fetch candidates for all unique entity texts in a single batched query.
|
||||
# Uses the GIN trigram index on LOWER(canonical_name) for case-insensitive
|
||||
# similarity lookup. Previous version also had LIKE '%...' substring fallbacks,
|
||||
# but those forced full sequential scans of the entities table and caused
|
||||
@@ -413,32 +361,21 @@ class EntityResolver:
|
||||
# to 0.15 (from default 0.3) catches most substring relationships while
|
||||
# staying fully index-based.
|
||||
await conn.execute("SET pg_trgm.similarity_threshold = 0.15")
|
||||
try:
|
||||
rows = []
|
||||
for entity_text_batch in self._chunked(entity_texts, self.entity_resolution_batch_size):
|
||||
rows.extend(
|
||||
await conn.fetch(
|
||||
f"""
|
||||
SELECT DISTINCT ON (e.id)
|
||||
e.id, e.canonical_name, e.metadata, e.last_seen, e.mention_count,
|
||||
q.query_text
|
||||
FROM unnest($2::text[]) AS q(query_text)
|
||||
JOIN {fq_table("entities")} e ON (
|
||||
e.bank_id = $1
|
||||
AND LOWER(e.canonical_name) % LOWER(q.query_text)
|
||||
)
|
||||
""",
|
||||
bank_id,
|
||||
entity_text_batch,
|
||||
)
|
||||
)
|
||||
finally:
|
||||
# asyncpg returns connections to the pool with session state intact,
|
||||
# so the lowered threshold would leak to future borrowers without RESET.
|
||||
try:
|
||||
await conn.execute("RESET pg_trgm.similarity_threshold")
|
||||
except Exception:
|
||||
logger.warning("Failed to reset pg_trgm similarity threshold after candidate lookup", exc_info=True)
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT DISTINCT ON (e.id)
|
||||
e.id, e.canonical_name, e.metadata, e.last_seen, e.mention_count,
|
||||
q.query_text
|
||||
FROM unnest($2::text[]) AS q(query_text)
|
||||
JOIN {fq_table("entities")} e ON (
|
||||
e.bank_id = $1
|
||||
AND LOWER(e.canonical_name) % LOWER(q.query_text)
|
||||
)
|
||||
""",
|
||||
bank_id,
|
||||
entity_texts,
|
||||
)
|
||||
await conn.execute("RESET pg_trgm.similarity_threshold")
|
||||
|
||||
# Group candidates by query_text
|
||||
all_candidates: dict[str, list] = {t: [] for t in entity_texts}
|
||||
@@ -481,24 +418,11 @@ class EntityResolver:
|
||||
cooccurrence_map[eid2].add(id_to_name[eid1])
|
||||
|
||||
return await self._resolve_from_candidates(
|
||||
conn,
|
||||
bank_id,
|
||||
entities_data,
|
||||
unit_event_date,
|
||||
all_candidates,
|
||||
cooccurrence_map,
|
||||
taxonomy_lookup,
|
||||
labels_cfg,
|
||||
conn, bank_id, entities_data, unit_event_date, all_candidates, cooccurrence_map
|
||||
)
|
||||
|
||||
async def _resolve_entities_batch_oracle_fuzzy(
|
||||
self,
|
||||
conn: Any,
|
||||
bank_id: str,
|
||||
entities_data: list[dict],
|
||||
unit_event_date: datetime | None,
|
||||
taxonomy_lookup: set[str] | None = None,
|
||||
labels_cfg=None,
|
||||
self, conn: Any, bank_id: str, entities_data: list[dict], unit_event_date: datetime | None
|
||||
) -> list[str]:
|
||||
"""
|
||||
Oracle strategy: fetch similar candidates using UTL_MATCH.JARO_WINKLER_SIMILARITY.
|
||||
@@ -512,28 +436,23 @@ class EntityResolver:
|
||||
entities_table = fq_table("entities")
|
||||
|
||||
try:
|
||||
# Batch entity texts into bounded sub-queries using JSON_TABLE to
|
||||
# Batch all entity texts into a single query using JSON_TABLE to
|
||||
# expand the list into rows. UTL_MATCH.JARO_WINKLER_SIMILARITY
|
||||
# returns 0-100; threshold 70 ≈ pg_trgm similarity 0.15.
|
||||
# Bounded batches mirror the PG trigram path so very wide retain
|
||||
# batches don't time out a single JOIN on large banks.
|
||||
rows = []
|
||||
for entity_text_batch in self._chunked(entity_texts, self.entity_resolution_batch_size):
|
||||
rows.extend(
|
||||
await conn.fetch(
|
||||
f"""
|
||||
SELECT e.id, e.canonical_name, e.metadata, e.last_seen, e.mention_count,
|
||||
q.query_text
|
||||
FROM JSON_TABLE($2, '$[*]' COLUMNS (query_text VARCHAR2(4000) PATH '$')) q
|
||||
JOIN {entities_table} e ON (
|
||||
e.bank_id = $1
|
||||
AND UTL_MATCH.JARO_WINKLER_SIMILARITY(LOWER(e.canonical_name), LOWER(q.query_text)) > 70
|
||||
)
|
||||
""",
|
||||
bank_id,
|
||||
json.dumps(entity_text_batch),
|
||||
)
|
||||
entity_texts_json = json.dumps(entity_texts)
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT e.id, e.canonical_name, e.metadata, e.last_seen, e.mention_count,
|
||||
q.query_text
|
||||
FROM JSON_TABLE($2, '$[*]' COLUMNS (query_text VARCHAR2(4000) PATH '$')) q
|
||||
JOIN {entities_table} e ON (
|
||||
e.bank_id = $1
|
||||
AND UTL_MATCH.JARO_WINKLER_SIMILARITY(LOWER(e.canonical_name), LOWER(q.query_text)) > 70
|
||||
)
|
||||
""",
|
||||
bank_id,
|
||||
entity_texts_json,
|
||||
)
|
||||
except Exception as e:
|
||||
# UTL_MATCH may not be available (ORA-06550, ORA-00904, etc.)
|
||||
# Catch broadly because Oracle error types vary depending on driver.
|
||||
@@ -587,14 +506,7 @@ class EntityResolver:
|
||||
cooccurrence_map[eid2].add(id_to_name[eid1])
|
||||
|
||||
return await self._resolve_from_candidates(
|
||||
conn,
|
||||
bank_id,
|
||||
entities_data,
|
||||
unit_event_date,
|
||||
all_candidates,
|
||||
cooccurrence_map,
|
||||
taxonomy_lookup,
|
||||
labels_cfg,
|
||||
conn, bank_id, entities_data, unit_event_date, all_candidates, cooccurrence_map
|
||||
)
|
||||
|
||||
async def _resolve_from_candidates(
|
||||
@@ -605,8 +517,6 @@ class EntityResolver:
|
||||
unit_event_date,
|
||||
all_candidates: dict[str, list],
|
||||
cooccurrence_map: dict[str, set[str]],
|
||||
taxonomy_lookup: set[str] | None = None,
|
||||
labels_cfg=None,
|
||||
) -> list[str]:
|
||||
"""Shared scoring + upsert logic used by both lookup strategies."""
|
||||
|
||||
@@ -623,34 +533,11 @@ class EntityResolver:
|
||||
|
||||
candidates = all_candidates.get(entity_text, [])
|
||||
|
||||
# Label entities (from entity_labels config) use exact matching only.
|
||||
# Their canonical names are user-defined (e.g., "use:use-001"),
|
||||
# so fuzzy resolution must NOT merge distinct label values that
|
||||
# happen to be textually similar (GH-1558).
|
||||
is_label = bool(
|
||||
labels_cfg and taxonomy_lookup and _is_label_entity(entity_text, labels_cfg, taxonomy_lookup)
|
||||
)
|
||||
|
||||
if not candidates:
|
||||
# Will create new entity
|
||||
entities_to_create.append(_EntityToCreate(idx=idx, name=entity_text, event_date=entity_event_date))
|
||||
continue
|
||||
|
||||
if is_label:
|
||||
# Exact case-insensitive match only for label entities
|
||||
exact_match = None
|
||||
entity_text_lower = entity_text.lower()
|
||||
for candidate_id, canonical_name, metadata, last_seen, mention_count in candidates:
|
||||
if canonical_name.lower() == entity_text_lower:
|
||||
exact_match = candidate_id
|
||||
break
|
||||
if exact_match:
|
||||
entity_ids[idx] = exact_match
|
||||
entities_to_update.append(_EntityStat(entity_id=exact_match, event_date=entity_event_date))
|
||||
else:
|
||||
entities_to_create.append(_EntityToCreate(idx=idx, name=entity_text, event_date=entity_event_date))
|
||||
continue
|
||||
|
||||
# Score candidates
|
||||
best_candidate = None
|
||||
best_score = 0.0
|
||||
|
||||
@@ -1,357 +0,0 @@
|
||||
"""Async graph maintenance after document/unit deletes.
|
||||
|
||||
Three reconciliation passes run together on every worker invocation:
|
||||
|
||||
1. **Relink top-up.** Drain ``graph_maintenance_queue`` (units whose
|
||||
outgoing temporal/semantic links lost a neighbour to a delete). For
|
||||
each, count current outgoing links per type; if below cap, run the
|
||||
same probes retain uses (:func:`fetch_temporal_neighbors`,
|
||||
:func:`compute_semantic_links_ann`) and insert the missing links.
|
||||
``bulk_insert_links`` has ``ON CONFLICT DO NOTHING`` on the uniqueness
|
||||
key, so we can re-probe freely and the DB de-dupes.
|
||||
|
||||
2. **Orphan entity prune.** Delete ``entities`` rows in the bank that no
|
||||
longer have any ``unit_entities`` references. FK ON DELETE CASCADE on
|
||||
``entity_cooccurrences`` then removes any cooccurrence row pointing
|
||||
at the pruned entities.
|
||||
|
||||
3. **Stale cooccurrence prune.** Defensive sweep for cooccurrence rows
|
||||
where both endpoints still exist but no current memory_unit references
|
||||
both of them — the cooccurrence was real at the time it was recorded,
|
||||
but every unit that witnessed it has since been deleted.
|
||||
|
||||
All three passes run on every invocation. The queue is the only source
|
||||
of work for pass 1; passes 2 and 3 are bank-wide sweeps backed by indexes
|
||||
on ``entities(bank_id)`` and ``unit_entities(entity_id)``, so they're
|
||||
cheap when there's nothing to do.
|
||||
|
||||
The worker dedupes on bank: a second job for the same bank is dropped
|
||||
while one is pending. Once processing starts, a new job becomes the
|
||||
*next* pending slot — so work enqueued during processing gets picked up
|
||||
by the follow-up run.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import time
|
||||
import uuid as uuid_module
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from ..models import RequestContext
|
||||
from .db.base import DatabaseConnection
|
||||
from .retain.link_utils import (
|
||||
MAX_TEMPORAL_LINKS_PER_UNIT,
|
||||
_bulk_insert_links,
|
||||
_normalize_datetime,
|
||||
compute_semantic_links_ann,
|
||||
)
|
||||
from .schema import fq_table
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .memory_engine import MemoryEngine
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Mirrors the ``top_k`` default in ``compute_semantic_links_ann`` at retain
|
||||
# time. If you change one, change the other — otherwise victims would either
|
||||
# never reach the cap (probe returns less than the cap) or stay perpetually
|
||||
# under it (cap is higher than retain creates).
|
||||
MAX_SEMANTIC_LINKS_PER_UNIT = 50
|
||||
|
||||
# Worker fetches this many rows per relink-loop iteration. Bounds
|
||||
# per-iteration probe/insert latency so a 10k-row backlog doesn't hold a
|
||||
# worker slot for minutes. Chosen so the typical iteration runs in well
|
||||
# under 1s.
|
||||
_DRAIN_BATCH_SIZE = 50
|
||||
|
||||
|
||||
@dataclass
|
||||
class JobResult:
|
||||
"""Counters surfaced to the worker dispatcher and operation result."""
|
||||
|
||||
relink_units_processed: int = 0
|
||||
relink_links_added: int = 0
|
||||
orphan_entities_pruned: int = 0
|
||||
stale_cooccurrences_pruned: int = 0
|
||||
|
||||
def as_dict(self) -> dict[str, int]:
|
||||
return {
|
||||
"relink_units_processed": self.relink_units_processed,
|
||||
"relink_links_added": self.relink_links_added,
|
||||
"orphan_entities_pruned": self.orphan_entities_pruned,
|
||||
"stale_cooccurrences_pruned": self.stale_cooccurrences_pruned,
|
||||
}
|
||||
|
||||
|
||||
async def enqueue_relink_victims(
|
||||
conn: DatabaseConnection,
|
||||
bank_id: str,
|
||||
deleted_unit_ids: list[str],
|
||||
ops: Any,
|
||||
) -> int:
|
||||
"""Enqueue surviving units whose outgoing temporal/semantic links pointed at
|
||||
``deleted_unit_ids`` for later link top-up.
|
||||
|
||||
Must run inside the same transaction that deletes the units, *before* the
|
||||
cascade fires — once the rows are gone, the join that finds the victims
|
||||
returns nothing.
|
||||
|
||||
Args:
|
||||
conn: Database connection inside the active delete transaction.
|
||||
bank_id: Bank owning the deleted units.
|
||||
deleted_unit_ids: Memory_unit IDs about to be (or being) deleted.
|
||||
ops: ``DataAccessOps`` instance, supplies the dialect-specific
|
||||
bulk-insert path.
|
||||
|
||||
Returns:
|
||||
Number of distinct victim units enqueued (after dedup against rows
|
||||
already in the queue).
|
||||
"""
|
||||
if not deleted_unit_ids:
|
||||
return 0
|
||||
|
||||
deleted_uuids = [uuid_module.UUID(uid) if isinstance(uid, str) else uid for uid in deleted_unit_ids]
|
||||
deleted_str_set = {str(uid) for uid in deleted_uuids}
|
||||
|
||||
# Find units (other than the ones being deleted) that have an outgoing
|
||||
# temporal/semantic link pointing at a doomed unit. Entity links are
|
||||
# intentionally excluded — they're scheduled for removal and would only
|
||||
# add noise to the recompute job.
|
||||
victim_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT DISTINCT from_unit_id
|
||||
FROM {fq_table("memory_links")}
|
||||
WHERE to_unit_id = ANY($1::uuid[])
|
||||
AND bank_id = $2
|
||||
AND link_type IN ('temporal', 'semantic')
|
||||
""",
|
||||
deleted_uuids,
|
||||
bank_id,
|
||||
)
|
||||
|
||||
victim_ids = [row["from_unit_id"] for row in victim_rows if str(row["from_unit_id"]) not in deleted_str_set]
|
||||
|
||||
if not victim_ids:
|
||||
return 0
|
||||
|
||||
await ops.enqueue_graph_maintenance(
|
||||
conn,
|
||||
fq_table("graph_maintenance_queue"),
|
||||
bank_id,
|
||||
victim_ids,
|
||||
)
|
||||
|
||||
logger.debug(
|
||||
f"[GRAPH_MAINT] Enqueued {len(victim_ids)} relink victims in "
|
||||
f"bank={bank_id} (deleted {len(deleted_unit_ids)} units)"
|
||||
)
|
||||
return len(victim_ids)
|
||||
|
||||
|
||||
async def run_graph_maintenance_job(
|
||||
memory_engine: "MemoryEngine",
|
||||
bank_id: str,
|
||||
request_context: RequestContext,
|
||||
operation_id: str | None = None,
|
||||
) -> dict[str, int]:
|
||||
"""Run all maintenance passes for ``bank_id`` until the relink queue is
|
||||
drained, then sweep entities and cooccurrences once.
|
||||
|
||||
Returns:
|
||||
Per-pass counters from :class:`JobResult`.
|
||||
"""
|
||||
del request_context # accepted for symmetry with other run_*_job helpers
|
||||
backend = await memory_engine._get_backend()
|
||||
ops = backend.ops
|
||||
|
||||
result = JobResult()
|
||||
job_start = time.time()
|
||||
|
||||
# --- Pass 1: relink ---
|
||||
# Per-iteration loop: claim → top up → commit. We rely on submit-time
|
||||
# dedup to keep at most one job per bank running, so no need for
|
||||
# SKIP LOCKED.
|
||||
iterations = 0
|
||||
while True:
|
||||
from .memory_engine import acquire_with_retry
|
||||
|
||||
async with acquire_with_retry(backend) as conn:
|
||||
async with conn.transaction():
|
||||
unit_ids = await ops.claim_graph_maintenance_batch(
|
||||
conn,
|
||||
fq_table("graph_maintenance_queue"),
|
||||
bank_id,
|
||||
_DRAIN_BATCH_SIZE,
|
||||
)
|
||||
if not unit_ids:
|
||||
break
|
||||
|
||||
result.relink_links_added += await _relink_batch(conn, bank_id, unit_ids, ops, backend)
|
||||
|
||||
result.relink_units_processed += len(unit_ids)
|
||||
iterations += 1
|
||||
|
||||
if iterations > 10000:
|
||||
# Defensive guard against runaway loops — at 50 units/iter that's
|
||||
# 500k targets, far beyond any realistic single-bank backlog.
|
||||
logger.error(
|
||||
f"[GRAPH_MAINT] bank={bank_id} hit iteration cap ({iterations}); aborting relink ({result.as_dict()})"
|
||||
)
|
||||
break
|
||||
|
||||
# --- Pass 2 & 3: entity / cooccurrence sweeps ---
|
||||
# Bank-wide single-statement deletes. Cheap when there's nothing to do.
|
||||
from .memory_engine import acquire_with_retry
|
||||
|
||||
async with acquire_with_retry(backend) as conn:
|
||||
async with conn.transaction():
|
||||
result.orphan_entities_pruned = await ops.prune_orphan_entities(
|
||||
conn,
|
||||
fq_table("entities"),
|
||||
fq_table("unit_entities"),
|
||||
bank_id,
|
||||
)
|
||||
# The orphan prune above cascades cooccurrences via FK. The
|
||||
# explicit cooccurrence pass below catches the *stale-count*
|
||||
# case: both entities still exist but no current unit witnesses
|
||||
# them together.
|
||||
result.stale_cooccurrences_pruned = await ops.prune_stale_cooccurrences(
|
||||
conn,
|
||||
fq_table("entity_cooccurrences"),
|
||||
fq_table("unit_entities"),
|
||||
fq_table("entities"),
|
||||
bank_id,
|
||||
)
|
||||
|
||||
elapsed = time.time() - job_start
|
||||
logger.info(
|
||||
f"[GRAPH_MAINT] bank={bank_id} done: {result.as_dict()}, elapsed={elapsed:.2f}s, operation_id={operation_id}"
|
||||
)
|
||||
return result.as_dict()
|
||||
|
||||
|
||||
async def _relink_batch(
|
||||
conn: DatabaseConnection,
|
||||
bank_id: str,
|
||||
victim_ids: list[str],
|
||||
ops: Any,
|
||||
backend: Any,
|
||||
) -> int:
|
||||
"""Top up temporal/semantic links for a batch of victim units. Returns rows inserted."""
|
||||
# Load each victim's metadata. Victims whose units were deleted between
|
||||
# enqueue and now silently drop out — exactly the no-op behaviour we want
|
||||
# for stale queue rows.
|
||||
victim_uuids = [uuid_module.UUID(vid) for vid in victim_ids]
|
||||
victim_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT id::text AS id, event_date, fact_type, embedding::text AS embedding
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE id = ANY($1::uuid[])
|
||||
AND bank_id = $2
|
||||
AND fact_type IN ('experience', 'world')
|
||||
""",
|
||||
victim_uuids,
|
||||
bank_id,
|
||||
)
|
||||
|
||||
if not victim_rows:
|
||||
return 0
|
||||
|
||||
alive_uuids = [uuid_module.UUID(row["id"]) for row in victim_rows]
|
||||
|
||||
# Count current outgoing temporal/semantic links per victim so we only
|
||||
# probe for the ones genuinely below cap. Saves the bulk of the work when
|
||||
# most victims still have plenty of links.
|
||||
count_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT from_unit_id, link_type, COUNT(*) AS cnt
|
||||
FROM {fq_table("memory_links")}
|
||||
WHERE from_unit_id = ANY($1::uuid[])
|
||||
AND bank_id = $2
|
||||
AND link_type IN ('temporal', 'semantic')
|
||||
GROUP BY from_unit_id, link_type
|
||||
""",
|
||||
alive_uuids,
|
||||
bank_id,
|
||||
)
|
||||
counts: dict[tuple[str, str], int] = {}
|
||||
for row in count_rows:
|
||||
counts[(str(row["from_unit_id"]), row["link_type"])] = int(row["cnt"])
|
||||
|
||||
# --- Temporal top-up ---
|
||||
temporal_needs = [r for r in victim_rows if counts.get((r["id"], "temporal"), 0) < MAX_TEMPORAL_LINKS_PER_UNIT]
|
||||
new_links: list[tuple] = []
|
||||
|
||||
if temporal_needs:
|
||||
lateral_unit_ids = [uuid_module.UUID(r["id"]) for r in temporal_needs if r["event_date"] is not None]
|
||||
lateral_event_dates = [
|
||||
_normalize_datetime(r["event_date"]) for r in temporal_needs if r["event_date"] is not None
|
||||
]
|
||||
lateral_fact_types = [r["fact_type"] for r in temporal_needs if r["event_date"] is not None]
|
||||
|
||||
if lateral_unit_ids:
|
||||
rows = await ops.fetch_temporal_neighbors(
|
||||
conn,
|
||||
fq_table("memory_units"),
|
||||
bank_id,
|
||||
lateral_unit_ids,
|
||||
lateral_event_dates,
|
||||
lateral_fact_types,
|
||||
MAX_TEMPORAL_LINKS_PER_UNIT,
|
||||
)
|
||||
for row in rows:
|
||||
time_diff_h = float(row["time_diff_hours"])
|
||||
# Mirror the 24h window enforced at retain time. The bidirectional
|
||||
# index scan returns the K closest neighbours regardless of
|
||||
# window, so we filter here.
|
||||
if time_diff_h > 24:
|
||||
continue
|
||||
weight = max(0.3, 1.0 - (time_diff_h / 24))
|
||||
new_links.append((row["from_id"], str(row["id"]), "temporal", weight, None))
|
||||
|
||||
# --- Semantic top-up ---
|
||||
# ANN must run on its own connection: it opens a nested transaction with
|
||||
# SET LOCAL hnsw.ef_search + CREATE TEMP TABLE ON COMMIT DROP, and nesting
|
||||
# that inside our current write transaction would commit our writes early.
|
||||
semantic_needs = [
|
||||
r
|
||||
for r in victim_rows
|
||||
if counts.get((r["id"], "semantic"), 0) < MAX_SEMANTIC_LINKS_PER_UNIT and r["embedding"] is not None
|
||||
]
|
||||
if semantic_needs:
|
||||
from .memory_engine import acquire_with_retry
|
||||
|
||||
seed_ids = [r["id"] for r in semantic_needs]
|
||||
seed_embs = [r["embedding"] for r in semantic_needs]
|
||||
seed_ftypes = [r["fact_type"] for r in semantic_needs]
|
||||
async with acquire_with_retry(backend) as ann_conn:
|
||||
try:
|
||||
ann_links = await compute_semantic_links_ann(
|
||||
ann_conn,
|
||||
bank_id,
|
||||
seed_ids,
|
||||
seed_embs,
|
||||
fact_types=seed_ftypes,
|
||||
)
|
||||
# Strip self-links (rare but possible because the ANN probe
|
||||
# has no exclude list — see the comment in compute_semantic_links_ann).
|
||||
ann_links = [lnk for lnk in ann_links if lnk[0] != lnk[1]]
|
||||
new_links.extend(ann_links)
|
||||
except Exception as e:
|
||||
# ANN uses PG-specific HNSW syntax; on dialects/configs where
|
||||
# it isn't available we still want the temporal top-up to land.
|
||||
logger.warning(f"[GRAPH_MAINT] Semantic top-up failed for bank={bank_id}: {type(e).__name__}: {e}")
|
||||
|
||||
if not new_links:
|
||||
return 0
|
||||
|
||||
await _bulk_insert_links(
|
||||
conn,
|
||||
new_links,
|
||||
bank_id=bank_id,
|
||||
skip_exists_check=False,
|
||||
ops=ops,
|
||||
)
|
||||
return len(new_links)
|
||||
@@ -9,7 +9,6 @@ import os
|
||||
import re
|
||||
import time
|
||||
import uuid
|
||||
from contextlib import AsyncExitStack
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
@@ -28,12 +27,9 @@ except ImportError:
|
||||
from ..config import (
|
||||
DEFAULT_LLM_MAX_CONCURRENT,
|
||||
DEFAULT_LLM_TIMEOUT,
|
||||
ENV_CONSOLIDATION_LLM_MAX_CONCURRENT,
|
||||
ENV_LLM_GROQ_SERVICE_TIER,
|
||||
ENV_LLM_MAX_CONCURRENT,
|
||||
ENV_LLM_TIMEOUT,
|
||||
ENV_REFLECT_LLM_MAX_CONCURRENT,
|
||||
ENV_RETAIN_LLM_MAX_CONCURRENT,
|
||||
)
|
||||
from ..metrics import get_metrics_collector
|
||||
from .response_models import TokenUsage
|
||||
@@ -46,75 +42,13 @@ logger = logging.getLogger(__name__)
|
||||
# Disable httpx logging
|
||||
logging.getLogger("httpx").setLevel(logging.WARNING)
|
||||
|
||||
# Global semaphore to limit concurrent LLM requests across all instances.
|
||||
# Set HINDSIGHT_API_LLM_MAX_CONCURRENT=1 for local LLMs (LM Studio, Ollama).
|
||||
# Global semaphore to limit concurrent LLM requests across all instances
|
||||
# Set HINDSIGHT_API_LLM_MAX_CONCURRENT=1 for local LLMs (LM Studio, Ollama)
|
||||
_llm_max_concurrent = int(os.getenv(ENV_LLM_MAX_CONCURRENT, str(DEFAULT_LLM_MAX_CONCURRENT)))
|
||||
_global_llm_semaphore = asyncio.Semaphore(_llm_max_concurrent)
|
||||
|
||||
|
||||
def _build_per_op_semaphores() -> dict[str, asyncio.Semaphore]:
|
||||
"""Build the per-operation semaphore registry from env vars.
|
||||
|
||||
Each per-op cap is composed with — not a substitute for — the global cap:
|
||||
a call that matches a configured operation must acquire both its per-op
|
||||
semaphore and the global semaphore. This lets operators reserve headroom
|
||||
in the global pool by capping individual operations (e.g. cap retain at 2
|
||||
of 4 global slots so the live chat path always has 2 slots available).
|
||||
|
||||
Operations without a configured env var are absent from the registry and
|
||||
therefore only constrained by the global cap.
|
||||
"""
|
||||
semaphores: dict[str, asyncio.Semaphore] = {}
|
||||
for op, env_var in (
|
||||
("retain", ENV_RETAIN_LLM_MAX_CONCURRENT),
|
||||
("reflect", ENV_REFLECT_LLM_MAX_CONCURRENT),
|
||||
("consolidation", ENV_CONSOLIDATION_LLM_MAX_CONCURRENT),
|
||||
):
|
||||
raw = os.getenv(env_var)
|
||||
if raw is None or raw == "":
|
||||
continue
|
||||
value = int(raw)
|
||||
if value <= 0:
|
||||
raise ValueError(f"{env_var} must be a positive integer, got {raw!r}")
|
||||
semaphores[op] = asyncio.Semaphore(value)
|
||||
return semaphores
|
||||
|
||||
|
||||
_per_op_llm_semaphores: dict[str, asyncio.Semaphore] = _build_per_op_semaphores()
|
||||
|
||||
|
||||
def _scope_to_operation(scope: str) -> str | None:
|
||||
"""Map a call scope to its per-operation concurrency bucket.
|
||||
|
||||
Returns None for scopes that don't belong to a tracked operation
|
||||
(verification probes, bank_mission, memory_think, mental_model_delta_ops),
|
||||
which then run under the global cap only.
|
||||
"""
|
||||
if scope.startswith("retain"):
|
||||
return "retain"
|
||||
if scope.startswith("reflect"):
|
||||
return "reflect"
|
||||
if scope.startswith("consolidation"):
|
||||
return "consolidation"
|
||||
return None
|
||||
|
||||
|
||||
def _semaphores_for_scope(scope: str) -> list[asyncio.Semaphore]:
|
||||
"""Return the semaphores a call with the given scope must acquire.
|
||||
|
||||
Always includes the global semaphore; includes the per-op semaphore when
|
||||
one is configured for the scope's operation bucket.
|
||||
"""
|
||||
op = _scope_to_operation(scope)
|
||||
per_op = _per_op_llm_semaphores.get(op) if op is not None else None
|
||||
if per_op is None:
|
||||
return [_global_llm_semaphore]
|
||||
# Per-op acquired first so contention queues on the narrower cap before
|
||||
# holding a global slot.
|
||||
return [per_op, _global_llm_semaphore]
|
||||
|
||||
|
||||
def sanitize_text(text: str | None) -> str | None:
|
||||
def sanitize_llm_output(text: str | None) -> str | None:
|
||||
"""
|
||||
Sanitize text by removing characters that break downstream systems.
|
||||
|
||||
@@ -126,12 +60,8 @@ def sanitize_text(text: str | None) -> str | None:
|
||||
|
||||
Surrogate characters are used in UTF-16 encoding but cannot be encoded
|
||||
in UTF-8. They can appear in Python strings from improperly decoded data
|
||||
(e.g., from JavaScript or broken files): a client may serialize a half-emoji
|
||||
split at a boundary as a lone ``\\udXXX`` escape. Such input crashes the
|
||||
SentenceTransformers/cross-encoder Rust tokenizers and stdout logging, so
|
||||
user content is sanitized at the retain/recall/reflect ingress (see issue
|
||||
#1875). Control characters commonly appear in LLM output embedded inside
|
||||
JSON string values.
|
||||
(e.g., from JavaScript or broken files). Control characters commonly appear
|
||||
in LLM output embedded inside JSON string values.
|
||||
"""
|
||||
if text is None:
|
||||
return None
|
||||
@@ -140,11 +70,6 @@ def sanitize_text(text: str | None) -> str | None:
|
||||
return re.sub(r"[\x00-\x08\x0b\x0c\x0e-\x1f\x7f\ud800-\udfff]", "", text)
|
||||
|
||||
|
||||
# Back-compat alias: this helper was originally introduced to scrub LLM *output*;
|
||||
# it now also scrubs user *input* at ingress, hence the broader name.
|
||||
sanitize_llm_output = sanitize_text
|
||||
|
||||
|
||||
class OutputTooLongError(Exception):
|
||||
"""
|
||||
Bridge exception raised when LLM output exceeds token limits.
|
||||
@@ -258,7 +183,6 @@ def create_llm_provider(
|
||||
AnthropicLLM,
|
||||
ClaudeCodeLLM,
|
||||
CodexLLM,
|
||||
FireworksLLM,
|
||||
GeminiLLM,
|
||||
LiteLLMLLM,
|
||||
LiteLLMRouterLLM,
|
||||
@@ -384,24 +308,10 @@ def create_llm_provider(
|
||||
extra_args=config.llamacpp_extra_args,
|
||||
)
|
||||
|
||||
elif provider_lower == "fireworks":
|
||||
# Fireworks online inference is OpenAI-compatible; FireworksLLM adds the
|
||||
# native (non-OpenAI) batch API on top. The existing LiteLLM
|
||||
# ``fireworks_ai/...`` online path (provider="litellm") is untouched.
|
||||
return FireworksLLM(
|
||||
provider=provider,
|
||||
api_key=api_key,
|
||||
base_url=base_url,
|
||||
model=model,
|
||||
reasoning_effort=reasoning_effort,
|
||||
extra_body=extra_body,
|
||||
)
|
||||
|
||||
elif provider_lower in (
|
||||
"openai",
|
||||
"groq",
|
||||
"ollama",
|
||||
"ollama-cloud",
|
||||
"lmstudio",
|
||||
"minimax",
|
||||
"deepseek",
|
||||
@@ -499,7 +409,6 @@ class LLMProvider:
|
||||
"openai",
|
||||
"groq",
|
||||
"ollama",
|
||||
"ollama-cloud",
|
||||
"gemini",
|
||||
"anthropic",
|
||||
"lmstudio",
|
||||
@@ -518,7 +427,6 @@ class LLMProvider:
|
||||
"openrouter",
|
||||
"zai",
|
||||
"opencode-go",
|
||||
"fireworks",
|
||||
]
|
||||
if self.provider not in valid_providers:
|
||||
raise ValueError(f"Invalid LLM provider: {self.provider}. Must be one of: {', '.join(valid_providers)}")
|
||||
@@ -529,8 +437,6 @@ class LLMProvider:
|
||||
self.base_url = "https://api.groq.com/openai/v1"
|
||||
elif self.provider == "ollama":
|
||||
self.base_url = "http://localhost:11434/v1"
|
||||
elif self.provider == "ollama-cloud":
|
||||
self.base_url = "https://ollama.com/v1"
|
||||
elif self.provider == "lmstudio":
|
||||
self.base_url = "http://localhost:1234/v1"
|
||||
elif self.provider == "minimax":
|
||||
@@ -723,10 +629,7 @@ class LLMProvider:
|
||||
structured = "+structured" if response_format is not None else ""
|
||||
set_stage(f"llm.{self.provider}.{scope}{structured}")
|
||||
|
||||
async with AsyncExitStack() as stack:
|
||||
for sem in _semaphores_for_scope(scope):
|
||||
await stack.enter_async_context(sem)
|
||||
|
||||
async with _global_llm_semaphore:
|
||||
# Delegate to provider implementation
|
||||
result = await self._provider_impl.call(
|
||||
messages=messages,
|
||||
@@ -751,7 +654,7 @@ class LLMProvider:
|
||||
# Sync the mock calls from provider implementation to wrapper
|
||||
self._mock_calls = self._provider_impl.get_mock_calls()
|
||||
|
||||
return result
|
||||
return result
|
||||
|
||||
async def call_with_tools(
|
||||
self,
|
||||
@@ -786,10 +689,7 @@ class LLMProvider:
|
||||
|
||||
set_stage(f"llm.{self.provider}.{scope}+tools")
|
||||
|
||||
async with AsyncExitStack() as stack:
|
||||
for sem in _semaphores_for_scope(scope):
|
||||
await stack.enter_async_context(sem)
|
||||
|
||||
async with _global_llm_semaphore:
|
||||
# Delegate to provider implementation
|
||||
result = await self._provider_impl.call_with_tools(
|
||||
messages=messages,
|
||||
@@ -812,7 +712,7 @@ class LLMProvider:
|
||||
# Sync the mock calls from provider implementation to wrapper
|
||||
self._mock_calls = self._provider_impl.get_mock_calls()
|
||||
|
||||
return result
|
||||
return result
|
||||
|
||||
def set_response_callback(self, fn: Any) -> None:
|
||||
"""Set a callback invoked on each call() instead of the fixed mock response."""
|
||||
@@ -941,14 +841,12 @@ class LLMProvider:
|
||||
"""Create provider from environment variables using config.py constants."""
|
||||
from ..config import (
|
||||
DEFAULT_LLM_PROVIDER,
|
||||
DEFAULT_LLM_REASONING_EFFORT,
|
||||
ENV_LLM_API_KEY,
|
||||
ENV_LLM_BASE_URL,
|
||||
ENV_LLM_DEFAULT_HEADERS,
|
||||
ENV_LLM_EXTRA_BODY,
|
||||
ENV_LLM_MODEL,
|
||||
ENV_LLM_PROVIDER,
|
||||
ENV_LLM_REASONING_EFFORT,
|
||||
_get_default_model_for_provider,
|
||||
)
|
||||
|
||||
@@ -972,7 +870,7 @@ class LLMProvider:
|
||||
api_key=api_key,
|
||||
base_url=base_url,
|
||||
model=model,
|
||||
reasoning_effort=os.getenv(ENV_LLM_REASONING_EFFORT, DEFAULT_LLM_REASONING_EFFORT),
|
||||
reasoning_effort="low",
|
||||
extra_body=extra_body,
|
||||
default_headers=default_headers,
|
||||
)
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,41 +0,0 @@
|
||||
"""Shared utilities for prompt assembly."""
|
||||
|
||||
import re
|
||||
|
||||
_LONE_OPEN_BRACE = re.compile(r"(?<!\{)\{(?!\{)")
|
||||
_LONE_CLOSE_BRACE = re.compile(r"(?<!\})\}(?!\})")
|
||||
|
||||
|
||||
def escape_for_prompt(text: str) -> str:
|
||||
"""Double any lone ``{`` / ``}`` so the text survives ``str.format`` untouched.
|
||||
|
||||
Prompt templates are often passed through ``str.format`` to substitute real
|
||||
placeholders like ``{facts_text}``. Any literal braces in caller-supplied
|
||||
text — e.g. a bank mission that contains JSON examples — would otherwise be
|
||||
interpreted as format keys and raise ``KeyError``.
|
||||
|
||||
Idempotent: text that already contains escaped ``{{`` / ``}}`` pairs is
|
||||
left as-is. Only lone braces (not adjacent to another brace of the same
|
||||
kind) are doubled.
|
||||
"""
|
||||
text = _LONE_OPEN_BRACE.sub("{{", text)
|
||||
text = _LONE_CLOSE_BRACE.sub("}}", text)
|
||||
return text
|
||||
|
||||
|
||||
def output_language_directive(language: str | None) -> str:
|
||||
"""Return an LLM directive forcing all output into ``language``.
|
||||
|
||||
Used by retain (fact extraction), consolidation (observations), and reflect
|
||||
(response synthesis) so HINDSIGHT_API_LLM_OUTPUT_LANGUAGE applies uniformly
|
||||
across every LLM-generated artifact. Returns an empty string when
|
||||
``language`` is unset so the calling prompt stays unchanged.
|
||||
"""
|
||||
if not language:
|
||||
return ""
|
||||
return (
|
||||
f"\n\nIMPORTANT: Respond exclusively in {language}. "
|
||||
f"Translate any source content into {language}. "
|
||||
f"All output text — including fact text, observations, entity names, "
|
||||
f"and the final response — must be in {language}."
|
||||
)
|
||||
@@ -7,7 +7,6 @@ This package contains concrete implementations of the LLMInterface for various p
|
||||
from .anthropic_llm import AnthropicLLM
|
||||
from .claude_code_llm import ClaudeCodeLLM
|
||||
from .codex_llm import CodexLLM
|
||||
from .fireworks_llm import FireworksLLM
|
||||
from .gemini_llm import GeminiLLM
|
||||
from .litellm_llm import LiteLLMLLM
|
||||
from .litellm_router_llm import LiteLLMRouterLLM
|
||||
@@ -20,7 +19,6 @@ __all__ = [
|
||||
"AnthropicLLM",
|
||||
"ClaudeCodeLLM",
|
||||
"CodexLLM",
|
||||
"FireworksLLM",
|
||||
"GeminiLLM",
|
||||
"LlamaCppLLM",
|
||||
"LiteLLMLLM",
|
||||
|
||||
@@ -93,6 +93,7 @@ class AnthropicLLM(LLMInterface):
|
||||
await self.call(
|
||||
messages=test_messages,
|
||||
max_completion_tokens=10,
|
||||
temperature=0.0,
|
||||
scope="verification",
|
||||
max_retries=0,
|
||||
)
|
||||
@@ -178,6 +179,9 @@ class AnthropicLLM(LLMInterface):
|
||||
if system_prompt:
|
||||
call_params["system"] = system_prompt
|
||||
|
||||
if temperature is not None:
|
||||
call_params["temperature"] = temperature
|
||||
|
||||
last_exception = None
|
||||
|
||||
for attempt in range(max_retries + 1):
|
||||
@@ -394,6 +398,9 @@ class AnthropicLLM(LLMInterface):
|
||||
if system_prompt:
|
||||
call_params["system"] = system_prompt
|
||||
|
||||
if temperature is not None:
|
||||
call_params["temperature"] = temperature
|
||||
|
||||
last_exception = None
|
||||
for attempt in range(max_retries + 1):
|
||||
try:
|
||||
|
||||
@@ -9,7 +9,6 @@ automatically handles authentication via `claude auth login` credentials.
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import tempfile
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
@@ -22,32 +21,6 @@ from hindsight_api.metrics import get_metrics_collector
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# Isolation env passed to the spawned `claude` CLI. CLAUDE_CONFIG_DIR
|
||||
# redirects the subprocess away from the host's ~/.claude/, so any
|
||||
# operator-installed plugins (e.g. hindsight-memory) and their Stop hooks do
|
||||
# not fire inside our LLM-call subprocesses. Without this, retain/reflect/
|
||||
# consolidation LLM calls would trigger a Stop-hook retain of the subprocess
|
||||
# transcript back into the same bank — a recursive feedback loop (issue #1751).
|
||||
# CLAUDE_SECURESTORAGE_CONFIG_DIR="" forces the CLI's keychain service name
|
||||
# back to the canonical un-suffixed entry that `claude auth login` wrote;
|
||||
# otherwise it would be namespaced by sha256(CLAUDE_CONFIG_DIR) and OAuth
|
||||
# lookup would fail. Requires bundled CLI >= 2.1.150 (claude-agent-sdk 0.2.82).
|
||||
_isolated_claude_env: dict[str, str] | None = None
|
||||
|
||||
|
||||
def _get_isolated_claude_env() -> dict[str, str]:
|
||||
"""Return a process-lifetime env dict that isolates the spawned CLI from user plugins."""
|
||||
global _isolated_claude_env
|
||||
if _isolated_claude_env is None:
|
||||
path = tempfile.mkdtemp(prefix="hindsight-claude-code-")
|
||||
_isolated_claude_env = {
|
||||
"CLAUDE_CONFIG_DIR": path,
|
||||
"CLAUDE_SECURESTORAGE_CONFIG_DIR": "",
|
||||
}
|
||||
logger.debug(f"Claude Code: isolated CLAUDE_CONFIG_DIR={path}")
|
||||
return _isolated_claude_env
|
||||
|
||||
|
||||
class ClaudeCodeLLM(LLMInterface):
|
||||
"""
|
||||
LLM provider using Claude Code authentication.
|
||||
@@ -210,7 +183,6 @@ class ClaudeCodeLLM(LLMInterface):
|
||||
system_prompt=system_prompt if system_prompt else None,
|
||||
max_turns=1, # Single-turn for API-style interactions
|
||||
allowed_tools=[], # Disable tools for standard LLM calls
|
||||
env=_get_isolated_claude_env(),
|
||||
)
|
||||
|
||||
# Call Claude Agent SDK
|
||||
@@ -501,7 +473,6 @@ class ClaudeCodeLLM(LLMInterface):
|
||||
max_turns=2, # Allow tool call + tool result round-trip
|
||||
mcp_servers=mcp_servers_config,
|
||||
allowed_tools=allowed_tool_names,
|
||||
env=_get_isolated_claude_env(),
|
||||
)
|
||||
|
||||
# Call Claude Agent SDK with retry logic
|
||||
|
||||
@@ -1,409 +0,0 @@
|
||||
"""
|
||||
Shared Codex OAuth authentication manager.
|
||||
|
||||
Extracted from ``CodexLLM`` so that both ``CodexLLM`` and
|
||||
``CodexOAuthEmbeddings`` can share JWT-expiry detection, single-flight
|
||||
token refresh, and atomic file persistence without duplicating the logic.
|
||||
|
||||
Usage
|
||||
-----
|
||||
Create a manager from the auth file::
|
||||
|
||||
mgr = CodexAuthManager.from_file()
|
||||
|
||||
Then call ``ensure_fresh_token()`` before each outbound request and
|
||||
``refresh_tokens(reason=..., force=...)`` on a reactive 401.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import binascii
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import tempfile
|
||||
import threading
|
||||
import time
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Module-level constants (shared with codex_llm.py via re-export there)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
# OAuth refresh endpoint and client id, mirrored from the canonical
|
||||
# ``@openai/codex`` CLI (codex-rs/login/src/auth/manager.rs on
|
||||
# github.com/openai/codex). The endpoint is overridable via env var so that
|
||||
# future Codex changes or staging environments can be pointed at without a
|
||||
# code change — same env var name the upstream CLI uses.
|
||||
_CODEX_REFRESH_TOKEN_URL = os.environ.get("CODEX_REFRESH_TOKEN_URL_OVERRIDE", "https://auth.openai.com/oauth/token")
|
||||
_CODEX_CLIENT_ID = "app_EMoamEEZ73f0CkXaXp7hrann"
|
||||
|
||||
# Proactively refresh this many seconds before the JWT ``exp`` claim. The
|
||||
# upstream Codex CLI uses no skew (it refreshes at ``exp <= now``); the
|
||||
# extra window reduces races where a request leaves the client with a token
|
||||
# that the server has already declared expired by the time it arrives.
|
||||
_CODEX_TOKEN_REFRESH_SKEW_SECONDS = 60
|
||||
|
||||
# OAuth error codes that the refresh endpoint returns when the refresh_token
|
||||
# itself is no longer usable. These are terminal — retrying refresh will not
|
||||
# succeed; the user must re-run ``codex auth login``.
|
||||
_CODEX_TERMINAL_REFRESH_ERROR_CODES = frozenset(
|
||||
{"refresh_token_expired", "refresh_token_reused", "refresh_token_invalidated"}
|
||||
)
|
||||
|
||||
|
||||
class CodexRefreshExpiredError(RuntimeError):
|
||||
"""Raised when the Codex refresh_token itself is no longer valid.
|
||||
|
||||
The user must re-run ``codex auth login`` to obtain new credentials.
|
||||
Callers should surface a clear remediation message and stop retrying.
|
||||
"""
|
||||
|
||||
|
||||
class CodexAuthManager:
|
||||
"""Sync Codex OAuth credential manager.
|
||||
|
||||
Holds the access_token, refresh_token, and account_id in memory and
|
||||
handles proactive/reactive refresh using a ``threading.Lock`` for
|
||||
single-flight semantics (safe to use from multiple threads or via
|
||||
``asyncio.to_thread``).
|
||||
|
||||
Parameters
|
||||
----------
|
||||
access_token:
|
||||
The current bearer token.
|
||||
account_id:
|
||||
The OpenAI account ID embedded in the Codex request headers.
|
||||
refresh_token:
|
||||
The OAuth refresh token. May be ``None`` when the auth file omits it;
|
||||
the provider still works as a one-shot loader in that case.
|
||||
auth_file:
|
||||
Path to ``~/.codex/auth.json``. Used for re-reading the refresh token
|
||||
on demand and for atomic persistence of rotated credentials.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
access_token: str,
|
||||
account_id: str,
|
||||
refresh_token: str | None,
|
||||
auth_file: Path,
|
||||
) -> None:
|
||||
self.access_token = access_token
|
||||
self.account_id = account_id
|
||||
self.refresh_token = refresh_token
|
||||
self._auth_file = auth_file
|
||||
self._lock = threading.Lock()
|
||||
self._http_client = httpx.Client(timeout=30.0)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Construction helpers
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
@classmethod
|
||||
def from_file(cls, auth_file: Path | None = None) -> "CodexAuthManager":
|
||||
"""Build a manager by reading credentials from ``auth_file``.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
auth_file:
|
||||
Defaults to ``~/.codex/auth.json``.
|
||||
|
||||
Raises
|
||||
------
|
||||
FileNotFoundError:
|
||||
If the auth file does not exist.
|
||||
ValueError:
|
||||
If the auth file is missing ``access_token`` or has an unexpected
|
||||
``auth_mode``.
|
||||
"""
|
||||
if auth_file is None:
|
||||
auth_file = Path.home() / ".codex" / "auth.json"
|
||||
|
||||
if not auth_file.exists():
|
||||
raise FileNotFoundError(f"Codex auth file not found: {auth_file}. Run 'codex auth login' to authenticate.")
|
||||
|
||||
with open(auth_file) as f:
|
||||
data = json.load(f)
|
||||
|
||||
auth_mode = data.get("auth_mode")
|
||||
if auth_mode != "chatgpt":
|
||||
raise ValueError(f"Expected Codex auth_mode='chatgpt', got: {auth_mode}")
|
||||
|
||||
tokens = data.get("tokens") or {}
|
||||
access_token = tokens.get("access_token")
|
||||
if not access_token:
|
||||
raise ValueError("No access_token found in Codex auth file. Run 'codex auth login' again.")
|
||||
|
||||
account_id = tokens.get("account_id") or ""
|
||||
refresh_token = tokens.get("refresh_token")
|
||||
|
||||
return cls(
|
||||
access_token=access_token,
|
||||
account_id=account_id,
|
||||
refresh_token=refresh_token,
|
||||
auth_file=auth_file,
|
||||
)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Token state helpers
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
@staticmethod
|
||||
def load_refresh_token_from_file(auth_file: Path) -> str | None:
|
||||
"""Read ``tokens.refresh_token`` from ``auth_file``.
|
||||
|
||||
Returns ``None`` when the file is unreadable or omits the field.
|
||||
Does not raise — the provider degrades to one-shot mode.
|
||||
"""
|
||||
try:
|
||||
with open(auth_file) as f:
|
||||
data = json.load(f)
|
||||
except (OSError, json.JSONDecodeError) as e:
|
||||
logger.warning(
|
||||
f"Codex auth file unreadable when loading refresh_token: {type(e).__name__}. "
|
||||
"Token refresh will not be available; the access_token in memory will be used until it expires."
|
||||
)
|
||||
return None
|
||||
return data.get("tokens", {}).get("refresh_token")
|
||||
|
||||
@staticmethod
|
||||
def _decode_jwt_exp_unixtime(token: str) -> int | None:
|
||||
"""Return the JWT ``exp`` claim as a unix timestamp, or None on parse failure.
|
||||
|
||||
We do not verify the signature — the server is the source of truth
|
||||
on whether the token is actually accepted. This is only used to
|
||||
schedule proactive refresh.
|
||||
"""
|
||||
try:
|
||||
parts = token.split(".")
|
||||
if len(parts) < 2:
|
||||
return None
|
||||
payload_b64 = parts[1]
|
||||
padding = "=" * (-len(payload_b64) % 4)
|
||||
payload_bytes = base64.urlsafe_b64decode(payload_b64 + padding)
|
||||
payload = json.loads(payload_bytes.decode("utf-8"))
|
||||
exp = payload.get("exp")
|
||||
return int(exp) if exp is not None else None
|
||||
except (ValueError, TypeError, json.JSONDecodeError, binascii.Error):
|
||||
return None
|
||||
|
||||
def _token_is_stale(self, skew_seconds: int = _CODEX_TOKEN_REFRESH_SKEW_SECONDS) -> bool:
|
||||
"""True when the cached access_token is past expiry (with skew).
|
||||
|
||||
Returns False when expiry cannot be determined — we'd rather use a
|
||||
possibly-expired token and recover via the reactive 401 path than
|
||||
refresh aggressively on every request when ``exp`` parsing fails.
|
||||
"""
|
||||
exp = self._decode_jwt_exp_unixtime(self.access_token)
|
||||
if exp is None:
|
||||
return False
|
||||
return exp <= int(time.time()) + skew_seconds
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Persistence
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _persist_auth_atomic(self, updated_tokens: dict[str, Any]) -> None:
|
||||
"""Write rotated tokens back to ``_auth_file`` atomically.
|
||||
|
||||
Re-reads the on-disk file first to avoid clobbering fields written
|
||||
by another process, patches ``tokens.*`` and ``last_refresh``, then
|
||||
writes to a sibling tempfile and calls ``os.replace`` (atomic on
|
||||
POSIX and Windows within the same filesystem).
|
||||
"""
|
||||
current: dict[str, Any]
|
||||
try:
|
||||
with open(self._auth_file) as f:
|
||||
loaded = json.load(f)
|
||||
current = loaded if isinstance(loaded, dict) else {"auth_mode": "chatgpt", "tokens": {}}
|
||||
except (OSError, json.JSONDecodeError):
|
||||
current = {"auth_mode": "chatgpt", "tokens": {}}
|
||||
|
||||
existing_tokens = current.get("tokens")
|
||||
tokens: dict[str, Any] = existing_tokens if isinstance(existing_tokens, dict) else {}
|
||||
for key in ("access_token", "refresh_token", "id_token", "account_id"):
|
||||
if key in updated_tokens and updated_tokens[key] is not None:
|
||||
tokens[key] = updated_tokens[key]
|
||||
current["tokens"] = tokens
|
||||
current["last_refresh"] = datetime.now(timezone.utc).isoformat().replace("+00:00", "Z")
|
||||
|
||||
parent = self._auth_file.parent
|
||||
parent.mkdir(parents=True, exist_ok=True)
|
||||
fd, tmp_path = tempfile.mkstemp(prefix=".auth.", suffix=".json.tmp", dir=str(parent))
|
||||
try:
|
||||
with os.fdopen(fd, "w") as f:
|
||||
json.dump(current, f, indent=2)
|
||||
f.flush()
|
||||
os.fsync(f.fileno())
|
||||
try:
|
||||
os.chmod(tmp_path, 0o600)
|
||||
except OSError:
|
||||
pass
|
||||
os.replace(tmp_path, self._auth_file)
|
||||
except Exception:
|
||||
try:
|
||||
os.unlink(tmp_path)
|
||||
except OSError:
|
||||
pass
|
||||
raise
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Error extraction
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
@staticmethod
|
||||
def _extract_oauth_error_code(response: httpx.Response) -> str | None:
|
||||
"""Pull the OAuth error code out of a 4xx response body, if present.
|
||||
|
||||
The refresh endpoint returns shapes like
|
||||
``{"error": "...", "error_code": "..."}`` or
|
||||
``{"error": {"code": "..."}}``.
|
||||
"""
|
||||
try:
|
||||
body = response.json()
|
||||
except (json.JSONDecodeError, ValueError):
|
||||
return None
|
||||
if not isinstance(body, dict):
|
||||
return None
|
||||
err = body.get("error")
|
||||
if isinstance(err, dict):
|
||||
code = err.get("code")
|
||||
if isinstance(code, str):
|
||||
return code
|
||||
code = body.get("error_code")
|
||||
if isinstance(code, str):
|
||||
return code
|
||||
if isinstance(err, str):
|
||||
return err
|
||||
return None
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Refresh
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def refresh_tokens(self, reason: str = "", *, force: bool = False) -> None:
|
||||
"""Synchronous single-flight OAuth token refresh.
|
||||
|
||||
Serialized through ``self._lock`` so concurrent threads produce one
|
||||
network request. The first caller refreshes; the rest wake up and
|
||||
skip if the token is no longer stale (proactive) or if the token
|
||||
has already changed (reactive / force).
|
||||
|
||||
Parameters
|
||||
----------
|
||||
reason:
|
||||
Free-form string included in log lines for diagnostics.
|
||||
force:
|
||||
When True, refresh even if the JWT exp claim looks fresh.
|
||||
Used by the reactive 401 path.
|
||||
|
||||
Raises
|
||||
------
|
||||
CodexRefreshExpiredError:
|
||||
When the server returns a terminal error code or any 401.
|
||||
RuntimeError:
|
||||
For other refresh failures (network, 5xx, etc.).
|
||||
"""
|
||||
token_before_lock = self.access_token
|
||||
with self._lock:
|
||||
if force:
|
||||
if self.access_token != token_before_lock:
|
||||
return
|
||||
else:
|
||||
if not self._token_is_stale():
|
||||
return
|
||||
|
||||
if not self.refresh_token:
|
||||
raise RuntimeError(
|
||||
"Codex access_token is expired but no refresh_token is available. "
|
||||
"Run 'codex auth login' to re-authenticate."
|
||||
)
|
||||
|
||||
log_reason = f" ({reason})" if reason else ""
|
||||
logger.info(f"Refreshing Codex OAuth access_token{log_reason}")
|
||||
|
||||
request_body = {
|
||||
"client_id": _CODEX_CLIENT_ID,
|
||||
"grant_type": "refresh_token",
|
||||
"refresh_token": self.refresh_token,
|
||||
}
|
||||
try:
|
||||
response = self._http_client.post(
|
||||
_CODEX_REFRESH_TOKEN_URL,
|
||||
json=request_body,
|
||||
headers={"Content-Type": "application/json"},
|
||||
timeout=30.0,
|
||||
)
|
||||
except httpx.RequestError as e:
|
||||
raise RuntimeError(f"Codex OAuth refresh network error: {type(e).__name__}") from e
|
||||
|
||||
if response.status_code == 401:
|
||||
error_code = self._extract_oauth_error_code(response)
|
||||
if error_code in _CODEX_TERMINAL_REFRESH_ERROR_CODES:
|
||||
raise CodexRefreshExpiredError(
|
||||
f"Codex refresh_token is permanently invalid (error.code={error_code}). "
|
||||
"Run 'codex auth login' to re-authenticate."
|
||||
)
|
||||
raise CodexRefreshExpiredError(
|
||||
f"Codex OAuth refresh returned 401 with unrecognized error code "
|
||||
f"({error_code or 'none'}). Run 'codex auth login' to re-authenticate."
|
||||
)
|
||||
|
||||
if response.status_code >= 400:
|
||||
raise RuntimeError(f"Codex OAuth refresh failed with HTTP {response.status_code}")
|
||||
|
||||
try:
|
||||
body = response.json()
|
||||
except json.JSONDecodeError as e:
|
||||
raise RuntimeError(f"Codex OAuth refresh returned non-JSON body: {e}") from e
|
||||
|
||||
new_access = body.get("access_token")
|
||||
if not new_access:
|
||||
raise RuntimeError("Codex OAuth refresh returned no access_token")
|
||||
|
||||
new_refresh = body.get("refresh_token") or self.refresh_token
|
||||
new_id_token = body.get("id_token")
|
||||
|
||||
# Update in-memory state first so waiters see fresh credentials
|
||||
# immediately, even if disk write fails.
|
||||
self.access_token = new_access
|
||||
self.refresh_token = new_refresh
|
||||
|
||||
persisted: dict[str, Any] = {
|
||||
"access_token": new_access,
|
||||
"refresh_token": new_refresh,
|
||||
}
|
||||
if new_id_token:
|
||||
persisted["id_token"] = new_id_token
|
||||
|
||||
try:
|
||||
self._persist_auth_atomic(persisted)
|
||||
except OSError as e:
|
||||
logger.warning(
|
||||
f"Codex OAuth refresh succeeded but persisting auth.json failed: {type(e).__name__}. "
|
||||
"In-memory credentials are up to date; on-disk file is stale."
|
||||
)
|
||||
|
||||
logger.info("Codex OAuth access_token refreshed successfully")
|
||||
|
||||
def ensure_fresh_token(self) -> None:
|
||||
"""Proactively refresh the access_token if it is near or past expiry.
|
||||
|
||||
Cheap when the token is fresh (just decodes the JWT exp claim and
|
||||
returns).
|
||||
"""
|
||||
if self._token_is_stale():
|
||||
self.refresh_tokens(reason="proactive (token near expiry)")
|
||||
|
||||
def close(self) -> None:
|
||||
"""Close the underlying HTTP client."""
|
||||
self._http_client.close()
|
||||
@@ -15,10 +15,15 @@ so that future server-side changes affect both clients identically.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import binascii
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import tempfile
|
||||
import time
|
||||
import uuid
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
@@ -28,27 +33,37 @@ from hindsight_api.engine.llm_interface import LLMInterface, OutputTooLongError
|
||||
from hindsight_api.engine.response_models import LLMToolCall, LLMToolCallResult, TokenUsage
|
||||
from hindsight_api.metrics import get_metrics_collector
|
||||
|
||||
from .codex_auth import (
|
||||
_CODEX_CLIENT_ID,
|
||||
_CODEX_REFRESH_TOKEN_URL,
|
||||
_CODEX_TERMINAL_REFRESH_ERROR_CODES,
|
||||
_CODEX_TOKEN_REFRESH_SKEW_SECONDS,
|
||||
CodexAuthManager,
|
||||
CodexRefreshExpiredError,
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# OAuth refresh endpoint and client id, mirrored from the canonical
|
||||
# ``@openai/codex`` CLI (codex-rs/login/src/auth/manager.rs on
|
||||
# github.com/openai/codex). The endpoint is overridable via env var so that
|
||||
# future Codex changes or staging environments can be pointed at without a
|
||||
# code change — same env var name the upstream CLI uses.
|
||||
_CODEX_REFRESH_TOKEN_URL = os.environ.get("CODEX_REFRESH_TOKEN_URL_OVERRIDE", "https://auth.openai.com/oauth/token")
|
||||
_CODEX_CLIENT_ID = "app_EMoamEEZ73f0CkXaXp7hrann"
|
||||
|
||||
# Proactively refresh this many seconds before the JWT ``exp`` claim. The
|
||||
# upstream Codex CLI uses no skew (it refreshes at ``exp <= now``); the
|
||||
# extra window reduces races where a request leaves the client with a token
|
||||
# that the server has already declared expired by the time it arrives.
|
||||
_CODEX_TOKEN_REFRESH_SKEW_SECONDS = 60
|
||||
|
||||
# OAuth error codes that the refresh endpoint returns when the refresh_token
|
||||
# itself is no longer usable. These are terminal — retrying refresh will not
|
||||
# succeed; the user must re-run ``codex auth login``.
|
||||
_CODEX_TERMINAL_REFRESH_ERROR_CODES = frozenset(
|
||||
{"refresh_token_expired", "refresh_token_reused", "refresh_token_invalidated"}
|
||||
)
|
||||
|
||||
# Re-export for backward compatibility (tests import from this module).
|
||||
__all__ = [
|
||||
"CodexLLM",
|
||||
"CodexRefreshExpiredError",
|
||||
"CodexAuthManager",
|
||||
"_CODEX_REFRESH_TOKEN_URL",
|
||||
"_CODEX_CLIENT_ID",
|
||||
"_CODEX_TOKEN_REFRESH_SKEW_SECONDS",
|
||||
"_CODEX_TERMINAL_REFRESH_ERROR_CODES",
|
||||
]
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
class CodexRefreshExpiredError(RuntimeError):
|
||||
"""Raised when the Codex refresh_token itself is no longer valid.
|
||||
|
||||
The user must re-run ``codex auth login`` to obtain new credentials.
|
||||
Callers should surface a clear remediation message and stop retrying.
|
||||
"""
|
||||
|
||||
|
||||
class CodexLLM(LLMInterface):
|
||||
@@ -71,15 +86,20 @@ class CodexLLM(LLMInterface):
|
||||
"""Initialize Codex LLM provider."""
|
||||
super().__init__(provider, api_key, base_url, model, reasoning_effort, **kwargs)
|
||||
|
||||
# Single-flight async refresh lock. Multiple concurrent coroutines
|
||||
# racing toward an expired token should produce one network refresh.
|
||||
# Path is fixed at ~/.codex/auth.json — matches the upstream CLI.
|
||||
# Storing it on self lets the refresh path re-read after another
|
||||
# process (e.g. a sidecar) rotates the file out from under us.
|
||||
self._auth_file = Path.home() / ".codex" / "auth.json"
|
||||
|
||||
# Single-flight refresh lock. Multiple concurrent requests racing
|
||||
# toward an expired token should produce one network refresh, not N.
|
||||
self._auth_lock = asyncio.Lock()
|
||||
|
||||
# Load Codex OAuth credentials (keep these methods for test patching).
|
||||
# Load Codex OAuth credentials
|
||||
try:
|
||||
access_token, account_id = self._load_codex_auth()
|
||||
refresh_token = self._load_codex_refresh_token()
|
||||
logger.info(f"Loaded Codex OAuth credentials for account: {account_id}")
|
||||
self.access_token, self.account_id = self._load_codex_auth()
|
||||
self.refresh_token = self._load_codex_refresh_token()
|
||||
logger.info(f"Loaded Codex OAuth credentials for account: {self.account_id}")
|
||||
except Exception as e:
|
||||
raise RuntimeError(
|
||||
f"Failed to load Codex OAuth credentials from ~/.codex/auth.json: {e}\n\n"
|
||||
@@ -90,22 +110,9 @@ class CodexLLM(LLMInterface):
|
||||
"Or use a different provider (openai, anthropic, gemini) with API keys."
|
||||
) from e
|
||||
|
||||
self._auth_manager = CodexAuthManager(
|
||||
access_token=access_token,
|
||||
account_id=account_id,
|
||||
refresh_token=refresh_token,
|
||||
auth_file=Path.home() / ".codex" / "auth.json",
|
||||
)
|
||||
|
||||
# Use ChatGPT backend API endpoint. Codex auth is tied to
|
||||
# chatgpt.com/backend-api, not the OpenAI-compatible base URL used by
|
||||
# other providers. Deployments often set a global LLM_BASE_URL for an
|
||||
# OpenAI-compatible proxy; ignore that inherited value unless the user
|
||||
# explicitly provides a Codex backend URL.
|
||||
if not self.base_url or self.base_url.rstrip("/").endswith("/v1"):
|
||||
# Use ChatGPT backend API endpoint
|
||||
if not self.base_url:
|
||||
self.base_url = "https://chatgpt.com/backend-api"
|
||||
else:
|
||||
self.base_url = self.base_url.rstrip("/")
|
||||
|
||||
# Normalize model name (strip openai/ prefix if present)
|
||||
if self.model.startswith("openai/"):
|
||||
@@ -118,42 +125,6 @@ class CodexLLM(LLMInterface):
|
||||
# HTTP client for SSE streaming
|
||||
self._client = httpx.AsyncClient(timeout=120.0)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Properties — delegate to _auth_manager (preserves test-visible API)
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
@property
|
||||
def access_token(self) -> str:
|
||||
return self._auth_manager.access_token
|
||||
|
||||
@access_token.setter
|
||||
def access_token(self, v: str) -> None:
|
||||
self._auth_manager.access_token = v
|
||||
|
||||
@property
|
||||
def account_id(self) -> str:
|
||||
return self._auth_manager.account_id
|
||||
|
||||
@property
|
||||
def refresh_token(self) -> str | None:
|
||||
return self._auth_manager.refresh_token
|
||||
|
||||
@refresh_token.setter
|
||||
def refresh_token(self, v: str | None) -> None:
|
||||
self._auth_manager.refresh_token = v
|
||||
|
||||
@property
|
||||
def _auth_file(self) -> Path:
|
||||
return self._auth_manager._auth_file
|
||||
|
||||
@_auth_file.setter
|
||||
def _auth_file(self, v: Path) -> None:
|
||||
self._auth_manager._auth_file = v
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Forwarding methods (keep surface area for tests / subclasses)
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _load_codex_auth(self) -> tuple[str, str]:
|
||||
"""
|
||||
Load OAuth credentials from ~/.codex/auth.json.
|
||||
@@ -190,57 +161,273 @@ class CodexLLM(LLMInterface):
|
||||
return access_token, account_id
|
||||
|
||||
def _load_codex_refresh_token(self) -> str | None:
|
||||
"""Read ``tokens.refresh_token`` from the configured auth file.
|
||||
"""Load ``tokens.refresh_token`` from ``~/.codex/auth.json``.
|
||||
|
||||
Kept as an instance method so existing tests that patch
|
||||
``CodexLLM._load_codex_refresh_token`` continue to work. Works both
|
||||
pre- and post-``__init__`` because it does not depend on
|
||||
``_auth_manager`` being constructed yet.
|
||||
Returns None when the auth file is unreadable or omits the field —
|
||||
the provider still functions as a one-shot loader in that case, it
|
||||
just can't refresh when the access_token expires. This deliberately
|
||||
does not raise so that ``__init__`` keeps the existing failure mode
|
||||
of raising only on missing ``access_token``.
|
||||
"""
|
||||
auth_file = (
|
||||
self._auth_manager._auth_file if hasattr(self, "_auth_manager") else Path.home() / ".codex" / "auth.json"
|
||||
)
|
||||
return CodexAuthManager.load_refresh_token_from_file(auth_file)
|
||||
try:
|
||||
with open(self._auth_file) as f:
|
||||
data = json.load(f)
|
||||
except (OSError, json.JSONDecodeError) as e:
|
||||
logger.warning(
|
||||
f"Codex auth file unreadable when loading refresh_token: {type(e).__name__}. "
|
||||
"Token refresh will not be available; the access_token in memory will be used until it expires."
|
||||
)
|
||||
return None
|
||||
return data.get("tokens", {}).get("refresh_token")
|
||||
|
||||
@staticmethod
|
||||
def _decode_jwt_exp_unixtime(token: str) -> int | None:
|
||||
"""Delegate to ``CodexAuthManager._decode_jwt_exp_unixtime``."""
|
||||
return CodexAuthManager._decode_jwt_exp_unixtime(token)
|
||||
"""Return the JWT ``exp`` claim as a unix timestamp, or None on parse failure.
|
||||
|
||||
ChatGPT/Codex access_tokens are JWTs whose payload includes ``exp``
|
||||
(RFC 7519). We need the expiry to schedule proactive refresh — the
|
||||
``auth.json`` file does not persist a separate ``expires_at`` field
|
||||
in the upstream CLI's shape, so decoding the JWT itself is the
|
||||
canonical way to know when the token is stale.
|
||||
|
||||
We do not verify the signature — the server is the source of truth
|
||||
on whether the token is actually accepted, and the only thing this
|
||||
method affects is the *timing* of refresh, not whether to trust the
|
||||
token contents.
|
||||
"""
|
||||
try:
|
||||
parts = token.split(".")
|
||||
if len(parts) < 2:
|
||||
return None
|
||||
payload_b64 = parts[1]
|
||||
# JWT uses base64url without padding. Re-pad before decoding.
|
||||
padding = "=" * (-len(payload_b64) % 4)
|
||||
payload_bytes = base64.urlsafe_b64decode(payload_b64 + padding)
|
||||
payload = json.loads(payload_bytes.decode("utf-8"))
|
||||
exp = payload.get("exp")
|
||||
return int(exp) if exp is not None else None
|
||||
except (ValueError, TypeError, json.JSONDecodeError, binascii.Error):
|
||||
return None
|
||||
|
||||
def _token_is_stale(self, skew_seconds: int = _CODEX_TOKEN_REFRESH_SKEW_SECONDS) -> bool:
|
||||
"""Delegate to ``_auth_manager._token_is_stale``."""
|
||||
return self._auth_manager._token_is_stale(skew_seconds)
|
||||
"""True when the cached access_token is past expiry (with skew).
|
||||
|
||||
Returns False when expiry cannot be determined — we'd rather use a
|
||||
possibly-expired token and recover via the reactive 401 path than
|
||||
refresh aggressively on every request when ``exp`` parsing fails.
|
||||
"""
|
||||
exp = self._decode_jwt_exp_unixtime(self.access_token)
|
||||
if exp is None:
|
||||
return False
|
||||
return exp <= int(time.time()) + skew_seconds
|
||||
|
||||
def _persist_auth_atomic(self, updated_tokens: dict[str, Any]) -> None:
|
||||
"""Delegate to ``_auth_manager._persist_auth_atomic``."""
|
||||
return self._auth_manager._persist_auth_atomic(updated_tokens)
|
||||
"""Write the rotated tokens back to ``~/.codex/auth.json`` atomically.
|
||||
|
||||
Strategy: re-read the on-disk auth.json (so we don't clobber fields
|
||||
another process may have added), patch ``tokens.*`` and
|
||||
``last_refresh``, write to a tempfile in the same directory with
|
||||
mode 0600, then ``os.replace`` onto the target. ``os.replace`` is
|
||||
atomic within the same filesystem on POSIX and Windows, so a
|
||||
concurrent reader will see either the old file or the fully-written
|
||||
new file — never a partial truncate, which is the upstream CLI's
|
||||
worst-case race.
|
||||
|
||||
On non-Unix platforms the chmod is a best-effort no-op; the parent
|
||||
directory permissions still bound access.
|
||||
"""
|
||||
current: dict[str, Any]
|
||||
try:
|
||||
with open(self._auth_file) as f:
|
||||
loaded = json.load(f)
|
||||
# auth.json should always be a JSON object at the top level; if
|
||||
# someone has hand-edited it into a non-object shape, fall back
|
||||
# to the minimal default rather than crashing the refresh path.
|
||||
current = loaded if isinstance(loaded, dict) else {"auth_mode": "chatgpt", "tokens": {}}
|
||||
except (OSError, json.JSONDecodeError):
|
||||
# If the file became unreadable between our last read and now,
|
||||
# construct a minimal shape rather than refusing to persist.
|
||||
current = {"auth_mode": "chatgpt", "tokens": {}}
|
||||
|
||||
existing_tokens = current.get("tokens")
|
||||
tokens: dict[str, Any] = existing_tokens if isinstance(existing_tokens, dict) else {}
|
||||
for key in ("access_token", "refresh_token", "id_token", "account_id"):
|
||||
if key in updated_tokens and updated_tokens[key] is not None:
|
||||
tokens[key] = updated_tokens[key]
|
||||
current["tokens"] = tokens
|
||||
current["last_refresh"] = datetime.now(timezone.utc).isoformat().replace("+00:00", "Z")
|
||||
|
||||
# Write to a sibling tempfile so the rename is same-filesystem.
|
||||
parent = self._auth_file.parent
|
||||
parent.mkdir(parents=True, exist_ok=True)
|
||||
fd, tmp_path = tempfile.mkstemp(prefix=".auth.", suffix=".json.tmp", dir=str(parent))
|
||||
try:
|
||||
with os.fdopen(fd, "w") as f:
|
||||
json.dump(current, f, indent=2)
|
||||
f.flush()
|
||||
os.fsync(f.fileno())
|
||||
try:
|
||||
os.chmod(tmp_path, 0o600)
|
||||
except OSError:
|
||||
pass # best-effort on platforms that don't support chmod
|
||||
os.replace(tmp_path, self._auth_file)
|
||||
except Exception:
|
||||
# Clean up the orphaned tempfile if rename fails.
|
||||
try:
|
||||
os.unlink(tmp_path)
|
||||
except OSError:
|
||||
pass
|
||||
raise
|
||||
|
||||
async def _refresh_oauth_tokens(self, reason: str = "", *, force: bool = False) -> None:
|
||||
"""Async single-flight OAuth token refresh.
|
||||
"""Refresh the OAuth access_token using the stored refresh_token.
|
||||
|
||||
Outer asyncio.Lock preserves single-flight semantics for concurrent
|
||||
coroutines; the actual network call is offloaded to a thread via
|
||||
``asyncio.to_thread`` so the event loop stays unblocked.
|
||||
Single-flight: serialized through ``self._auth_lock`` so concurrent
|
||||
callers produce one network request. The first caller refreshes; the
|
||||
rest wake up and observe that either (a) the in-memory token is no
|
||||
longer stale (proactive case) or (b) the in-memory token has changed
|
||||
since they entered (reactive case), and return without re-refreshing.
|
||||
|
||||
Args:
|
||||
reason: Free-form string included in log lines for diagnostics.
|
||||
force: When True, refresh even if the JWT exp claim looks fresh.
|
||||
Used by the reactive 401 path.
|
||||
Used by the reactive 401 path — the server rejected the
|
||||
token, so we cannot trust the JWT's self-reported expiry.
|
||||
|
||||
Raises:
|
||||
CodexRefreshExpiredError: when the server returns a terminal
|
||||
error code or any 401 on the refresh endpoint.
|
||||
error code (refresh_token_expired/reused/invalidated) or any
|
||||
401 on the refresh endpoint itself.
|
||||
RuntimeError: for other refresh failures (network, 5xx, etc.).
|
||||
"""
|
||||
# Capture the token we'd be refreshing BEFORE acquiring the lock so
|
||||
# that we can detect mid-wait rotation by another coroutine.
|
||||
token_before_lock = self.access_token
|
||||
async with self._auth_lock:
|
||||
if force:
|
||||
# Reactive: skip only if another coroutine already rotated
|
||||
# the token while we were waiting on the lock.
|
||||
if self.access_token != token_before_lock:
|
||||
return
|
||||
else:
|
||||
if not self._auth_manager._token_is_stale():
|
||||
# Proactive: skip if the token is no longer stale (the
|
||||
# canonical "another coroutine refreshed first" check).
|
||||
if not self._token_is_stale():
|
||||
return
|
||||
await asyncio.to_thread(lambda: self._auth_manager.refresh_tokens(reason, force=force))
|
||||
|
||||
if not self.refresh_token:
|
||||
raise RuntimeError(
|
||||
"Codex access_token is expired but no refresh_token is available. "
|
||||
"Run 'codex auth login' to re-authenticate."
|
||||
)
|
||||
|
||||
log_reason = f" ({reason})" if reason else ""
|
||||
logger.info(f"Refreshing Codex OAuth access_token{log_reason}")
|
||||
|
||||
request_body = {
|
||||
"client_id": _CODEX_CLIENT_ID,
|
||||
"grant_type": "refresh_token",
|
||||
"refresh_token": self.refresh_token,
|
||||
}
|
||||
try:
|
||||
response = await self._client.post(
|
||||
_CODEX_REFRESH_TOKEN_URL,
|
||||
json=request_body,
|
||||
headers={"Content-Type": "application/json"},
|
||||
timeout=30.0,
|
||||
)
|
||||
except httpx.RequestError as e:
|
||||
raise RuntimeError(f"Codex OAuth refresh network error: {type(e).__name__}") from e
|
||||
|
||||
if response.status_code == 401:
|
||||
# Classify by ``error.code`` (or top-level ``error`` string) — same
|
||||
# mapping as the upstream Rust CLI's request_chatgpt_token_refresh.
|
||||
error_code = self._extract_oauth_error_code(response)
|
||||
if error_code in _CODEX_TERMINAL_REFRESH_ERROR_CODES:
|
||||
raise CodexRefreshExpiredError(
|
||||
f"Codex refresh_token is permanently invalid (error.code={error_code}). "
|
||||
"Run 'codex auth login' to re-authenticate."
|
||||
)
|
||||
# Unknown 401 — treat as terminal too, matching the upstream classification.
|
||||
raise CodexRefreshExpiredError(
|
||||
f"Codex OAuth refresh returned 401 with unrecognized error code "
|
||||
f"({error_code or 'none'}). Run 'codex auth login' to re-authenticate."
|
||||
)
|
||||
|
||||
if response.status_code >= 400:
|
||||
# 5xx and other 4xx are transient/retryable from the caller's
|
||||
# perspective; surface as RuntimeError without leaking the
|
||||
# request body in logs.
|
||||
raise RuntimeError(f"Codex OAuth refresh failed with HTTP {response.status_code}")
|
||||
|
||||
try:
|
||||
body = response.json()
|
||||
except json.JSONDecodeError as e:
|
||||
raise RuntimeError(f"Codex OAuth refresh returned non-JSON body: {e}") from e
|
||||
|
||||
new_access = body.get("access_token")
|
||||
if not new_access:
|
||||
raise RuntimeError("Codex OAuth refresh returned no access_token")
|
||||
|
||||
# The refresh_token may rotate on each refresh — adopt the new
|
||||
# one if the server sent it, otherwise keep the existing.
|
||||
new_refresh = body.get("refresh_token") or self.refresh_token
|
||||
new_id_token = body.get("id_token")
|
||||
|
||||
# Update in-memory state first so callers waiting on the lock
|
||||
# see fresh credentials immediately, even if disk write fails.
|
||||
self.access_token = new_access
|
||||
self.refresh_token = new_refresh
|
||||
|
||||
persisted = {
|
||||
"access_token": new_access,
|
||||
"refresh_token": new_refresh,
|
||||
}
|
||||
if new_id_token:
|
||||
persisted["id_token"] = new_id_token
|
||||
|
||||
try:
|
||||
self._persist_auth_atomic(persisted)
|
||||
except OSError as e:
|
||||
# In-memory creds are valid; warn but don't fail the request
|
||||
# path. Future process starts will fall back to the stale
|
||||
# on-disk auth.json and immediately refresh.
|
||||
logger.warning(
|
||||
f"Codex OAuth refresh succeeded but persisting auth.json failed: {type(e).__name__}. "
|
||||
"In-memory credentials are up to date; on-disk file is stale."
|
||||
)
|
||||
|
||||
logger.info("Codex OAuth access_token refreshed successfully")
|
||||
|
||||
@staticmethod
|
||||
def _extract_oauth_error_code(response: "httpx.Response") -> str | None:
|
||||
"""Pull the OAuth error code out of a 4xx response body, if present.
|
||||
|
||||
The refresh endpoint returns shapes like
|
||||
``{"error": "...", "error_code": "..."}`` or
|
||||
``{"error": {"code": "..."}}``. We don't fail the call if the body
|
||||
is unparseable — the caller falls back to a generic "unknown" error.
|
||||
"""
|
||||
try:
|
||||
body = response.json()
|
||||
except (json.JSONDecodeError, ValueError):
|
||||
return None
|
||||
if not isinstance(body, dict):
|
||||
return None
|
||||
# Shape 1: error is a nested object with "code"
|
||||
err = body.get("error")
|
||||
if isinstance(err, dict):
|
||||
code = err.get("code")
|
||||
if isinstance(code, str):
|
||||
return code
|
||||
# Shape 2: top-level error_code string
|
||||
code = body.get("error_code")
|
||||
if isinstance(code, str):
|
||||
return code
|
||||
# Shape 3: error is itself a string code
|
||||
if isinstance(err, str):
|
||||
return err
|
||||
return None
|
||||
|
||||
async def _ensure_fresh_token(self) -> None:
|
||||
"""Refresh the access_token proactively if it is near or past expiry.
|
||||
@@ -248,10 +435,13 @@ class CodexLLM(LLMInterface):
|
||||
Called at the top of every API-bound method. Cheap when the token is
|
||||
fresh (just decodes the JWT exp claim and returns).
|
||||
"""
|
||||
if self._auth_manager._token_is_stale():
|
||||
if self._token_is_stale():
|
||||
try:
|
||||
await self._refresh_oauth_tokens(reason="proactive (token near expiry)")
|
||||
except CodexRefreshExpiredError:
|
||||
# Surface to the caller as the same RuntimeError shape the
|
||||
# request loop has historically raised, so existing error
|
||||
# handling paths keep working.
|
||||
raise
|
||||
|
||||
def _map_reasoning_effort(self, effort: str) -> str:
|
||||
@@ -883,6 +1073,5 @@ class CodexLLM(LLMInterface):
|
||||
return content if content else None, tool_calls
|
||||
|
||||
async def cleanup(self) -> None:
|
||||
"""Clean up HTTP clients."""
|
||||
"""Clean up HTTP client."""
|
||||
await self._client.aclose()
|
||||
self._auth_manager.close()
|
||||
|
||||
@@ -1,396 +0,0 @@
|
||||
"""Fireworks AI provider with batch-inference support.
|
||||
|
||||
Fireworks' *online* inference endpoint (``/inference/v1``) is OpenAI-compatible,
|
||||
so ``FireworksLLM`` subclasses :class:`OpenAICompatibleLLM` and reuses its entire
|
||||
chat path. Only the *batch* mechanism differs: Fireworks does NOT implement the
|
||||
OpenAI ``/v1/batches`` API. Instead it exposes a proprietary, account-scoped
|
||||
dataset -> job -> download REST workflow on a separate control-plane host. This
|
||||
class overrides only the four batch members of the interface, translating that
|
||||
workflow to/from the OpenAI-batch shapes the retain orchestrator and
|
||||
``fact_extraction`` consumer expect — so nothing downstream changes.
|
||||
|
||||
Interface contract preserved (see ``fact_extraction.py`` result handling)::
|
||||
|
||||
result["response"]["body"]["choices"][0]["message"]["content"]
|
||||
|
||||
Workflow (control-plane host, e.g. ``https://api.fireworks.ai``)::
|
||||
|
||||
POST /v1/accounts/{acct}/datasets create input dataset
|
||||
POST /v1/accounts/{acct}/datasets/{id}:upload upload input JSONL
|
||||
POST /v1/accounts/{acct}/batchInferenceJobs create job
|
||||
GET /v1/accounts/{acct}/batchInferenceJobs/{jobId} poll status
|
||||
GET /v1/accounts/{acct}/datasets/{out}:getDownloadEndpoint signed URLs
|
||||
GET <signed-url> download output JSONL
|
||||
|
||||
NOTE: the exact *output JSONL line* nesting is not verbatim-documented by
|
||||
Fireworks. ``_normalize_output_line`` handles both the observed shape
|
||||
(``{custom_id, response: {...completion...}, error}``) and a ``response.body``
|
||||
nesting defensively. Confirm against a live key via the integration path.
|
||||
"""
|
||||
|
||||
import json
|
||||
import logging
|
||||
import uuid
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
from .openai_compatible_llm import OpenAICompatibleLLM
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Normalized statuses the retain driver treats as fatal (it raises) vs. keeps
|
||||
# polling on. "completed" ends the poll; anything else not in this set means
|
||||
# "keep polling".
|
||||
_TERMINAL_STATUSES = frozenset({"completed", "failed", "cancelled", "expired"})
|
||||
|
||||
# Default per-request timeout for control-plane HTTP calls (not the job wait).
|
||||
_HTTP_TIMEOUT_SECONDS = 60.0
|
||||
|
||||
# Fallback max job wait if neither a constructor arg nor config supplies one
|
||||
# (24h matches Fireworks' maximum job timeout).
|
||||
_DEFAULT_MAX_WAIT_SECONDS = 86_400
|
||||
|
||||
|
||||
class FireworksLLM(OpenAICompatibleLLM):
|
||||
"""Fireworks provider: OpenAI-compatible online inference + native batch."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
provider: str = "fireworks",
|
||||
*,
|
||||
api_key: str,
|
||||
base_url: str = "",
|
||||
model: str,
|
||||
reasoning_effort: str = "low",
|
||||
account_id: str | None = None,
|
||||
batch_base_url: str | None = None,
|
||||
max_wait_seconds: int | None = None,
|
||||
http_client: httpx.AsyncClient | None = None,
|
||||
**kwargs: Any,
|
||||
):
|
||||
super().__init__(
|
||||
provider=provider,
|
||||
api_key=api_key,
|
||||
base_url=base_url,
|
||||
model=model,
|
||||
reasoning_effort=reasoning_effort,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
# Batch settings are static, server-level config. Resolve any unset
|
||||
# values from the global config lazily so the online inference path
|
||||
# works even when batch is never configured.
|
||||
if account_id is None or batch_base_url is None or max_wait_seconds is None:
|
||||
from ...config import get_config
|
||||
|
||||
cfg = get_config()
|
||||
if account_id is None:
|
||||
account_id = cfg.fireworks_account_id
|
||||
if batch_base_url is None:
|
||||
batch_base_url = cfg.fireworks_batch_base_url
|
||||
if max_wait_seconds is None:
|
||||
max_wait_seconds = cfg.fireworks_batch_max_wait_seconds
|
||||
|
||||
self._account_id = account_id
|
||||
self._batch_base_url = (batch_base_url or "https://api.fireworks.ai").rstrip("/")
|
||||
self._max_wait_seconds: int = (
|
||||
int(max_wait_seconds) if max_wait_seconds is not None else _DEFAULT_MAX_WAIT_SECONDS
|
||||
)
|
||||
self._http_client = http_client
|
||||
self._owns_http_client = http_client is None
|
||||
|
||||
# ----- interface: batch members -------------------------------------
|
||||
|
||||
async def supports_batch_api(self) -> bool:
|
||||
return True
|
||||
|
||||
async def submit_batch(
|
||||
self,
|
||||
requests: list[dict[str, Any]],
|
||||
endpoint: str = "/v1/chat/completions",
|
||||
completion_window: str = "24h",
|
||||
) -> dict[str, Any]:
|
||||
# endpoint/completion_window are part of the LLMInterface batch contract
|
||||
# (used by the OpenAI path) but have no analogue in Fireworks' job API:
|
||||
# the request shape is fixed (chat) and the job timeout is server-side.
|
||||
# Kept for signature compatibility with the shared retain driver.
|
||||
self._require_account_id()
|
||||
logger.info(f"Submitting Fireworks batch with {len(requests)} requests")
|
||||
|
||||
jsonl = self._translate_requests(requests)
|
||||
input_dataset_id = f"hs-batch-in-{uuid.uuid4().hex}"
|
||||
output_dataset_id = f"hs-batch-out-{uuid.uuid4().hex}"
|
||||
headers = self._auth_headers()
|
||||
|
||||
# The `dataset` resource takes format + exampleCount on create. CHAT is
|
||||
# the format for chat-completion batch input; exampleCount is the JSONL
|
||||
# line count (Fireworks rejects uploaded datasets without it) and is an
|
||||
# int64 proto field, so it goes over the wire as a string.
|
||||
await self._request(
|
||||
"POST",
|
||||
self._datasets_url(),
|
||||
headers=headers,
|
||||
json={
|
||||
"datasetId": input_dataset_id,
|
||||
"dataset": {"format": "CHAT", "exampleCount": str(len(requests))},
|
||||
},
|
||||
)
|
||||
|
||||
await self._request(
|
||||
"POST",
|
||||
f"{self._datasets_url()}/{input_dataset_id}:upload",
|
||||
headers=headers,
|
||||
files={"file": ("batch_input.jsonl", jsonl.encode("utf-8"), "application/jsonl")},
|
||||
)
|
||||
|
||||
job_resp = await self._request(
|
||||
"POST",
|
||||
self._jobs_url(),
|
||||
headers=headers,
|
||||
json={
|
||||
"model": self.model,
|
||||
"inputDatasetId": self._dataset_resource(input_dataset_id),
|
||||
"outputDatasetId": self._dataset_resource(output_dataset_id),
|
||||
},
|
||||
)
|
||||
job = job_resp.json()
|
||||
job_id = self._last_segment(job.get("name")) or output_dataset_id
|
||||
|
||||
logger.info(f"Fireworks batch job submitted: {job_id}, state={job.get('state')}")
|
||||
|
||||
return {
|
||||
"batch_id": job_id,
|
||||
"status": self._normalize_state(job.get("state", "")),
|
||||
"input_dataset_id": input_dataset_id,
|
||||
"output_dataset_id": output_dataset_id,
|
||||
"created_at": job.get("createTime"),
|
||||
"request_count": len(requests),
|
||||
}
|
||||
|
||||
async def get_batch_status(self, batch_id: str) -> dict[str, Any]:
|
||||
self._require_account_id()
|
||||
job = (await self._request("GET", self._job_url(batch_id), headers=self._auth_headers())).json()
|
||||
|
||||
status = self._normalize_state(job.get("state", ""))
|
||||
progress = job.get("jobProgress") or {}
|
||||
result: dict[str, Any] = {
|
||||
"batch_id": batch_id,
|
||||
"status": status,
|
||||
"created_at": job.get("createTime"),
|
||||
"request_counts": {
|
||||
"total": _to_int(progress.get("totalInputRequests")),
|
||||
"completed": _to_int(progress.get("successfullyProcessedRequests")),
|
||||
"failed": _to_int(progress.get("failedRequests")),
|
||||
},
|
||||
}
|
||||
|
||||
output_dataset_id = job.get("outputDatasetId")
|
||||
if output_dataset_id:
|
||||
result["output_dataset_id"] = output_dataset_id
|
||||
# Fireworks reports terminal failure detail in the `status` {code,message}.
|
||||
if job.get("status"):
|
||||
result["errors"] = job["status"]
|
||||
|
||||
# PENDING-forever guard: the shared retain poll loop has no max-wait, so
|
||||
# if a (likely non-batch-eligible) job never reaches a terminal state we
|
||||
# surface "expired" once createTime is older than the cap. Derived from
|
||||
# the server's createTime so it survives crash-recovery polling resumes.
|
||||
if status not in _TERMINAL_STATUSES:
|
||||
elapsed = self._elapsed_seconds(job.get("createTime"))
|
||||
if elapsed is not None and elapsed > self._max_wait_seconds:
|
||||
result["status"] = "expired"
|
||||
result["errors"] = (
|
||||
f"Fireworks batch {batch_id} exceeded max wait of {self._max_wait_seconds}s "
|
||||
f"in state {job.get('state')!r}. The model may not be batch-eligible "
|
||||
f"(such jobs stay PENDING indefinitely)."
|
||||
)
|
||||
logger.error(result["errors"])
|
||||
|
||||
return result
|
||||
|
||||
async def retrieve_batch_results(self, batch_id: str) -> list[dict[str, Any]]:
|
||||
self._require_account_id()
|
||||
job = (await self._request("GET", self._job_url(batch_id), headers=self._auth_headers())).json()
|
||||
|
||||
status = self._normalize_state(job.get("state", ""))
|
||||
if status != "completed":
|
||||
raise ValueError(f"Fireworks batch {batch_id} is not completed yet (state: {job.get('state')!r})")
|
||||
|
||||
output_dataset_id = job.get("outputDatasetId")
|
||||
if not output_dataset_id:
|
||||
raise ValueError(f"Fireworks batch {batch_id} completed but reported no output dataset")
|
||||
|
||||
output_short_id = self._last_segment(output_dataset_id)
|
||||
if not output_short_id:
|
||||
raise ValueError(
|
||||
f"Fireworks batch {batch_id} reported an unparseable output dataset: {output_dataset_id!r}"
|
||||
)
|
||||
download = (
|
||||
await self._request("GET", self._download_endpoint_url(output_short_id), headers=self._auth_headers())
|
||||
).json()
|
||||
signed_urls = (download or {}).get("filenameToSignedUrls") or {}
|
||||
if not signed_urls:
|
||||
raise ValueError(f"Fireworks batch {batch_id} returned no downloadable output files")
|
||||
|
||||
# The output dataset contains a results file plus a separate error file.
|
||||
# Download every file and normalize each line; error-file lines carry an
|
||||
# `error` so partial failures surface per custom_id instead of vanishing.
|
||||
results: list[dict[str, Any]] = []
|
||||
for url in signed_urls.values():
|
||||
# Signed URLs are pre-authenticated — do not attach the bearer token.
|
||||
file_resp = await self._request("GET", url)
|
||||
for line in file_resp.text.strip().split("\n"):
|
||||
if line.strip():
|
||||
results.append(self._normalize_output_line(json.loads(line)))
|
||||
|
||||
logger.info(f"Retrieved {len(results)} results for Fireworks batch {batch_id}")
|
||||
return results
|
||||
|
||||
async def cleanup(self) -> None:
|
||||
await super().cleanup()
|
||||
if self._owns_http_client and self._http_client is not None:
|
||||
await self._http_client.aclose()
|
||||
|
||||
# ----- pure translation/normalization helpers (unit-tested) ----------
|
||||
|
||||
@staticmethod
|
||||
def _translate_requests(requests: list[dict[str, Any]]) -> str:
|
||||
"""OpenAI batch request -> Fireworks input JSONL.
|
||||
|
||||
Fireworks lines are ``{"custom_id", "body"}`` — the OpenAI ``method`` and
|
||||
``url`` keys are dropped; ``body`` is kept verbatim.
|
||||
"""
|
||||
lines = [
|
||||
json.dumps({"custom_id": req.get("custom_id"), "body": req.get("body")}, ensure_ascii=False)
|
||||
for req in requests
|
||||
]
|
||||
return "\n".join(lines)
|
||||
|
||||
@staticmethod
|
||||
def _normalize_state(fw_state: str) -> str:
|
||||
"""Fireworks job state -> the retain driver's expected status strings.
|
||||
|
||||
Handles both the API enum (``JOB_STATE_*``) and the guide's bare names
|
||||
(``COMPLETED``/``VALIDATING``/``EXPIRED``). Unknown / in-flight states map
|
||||
to ``in_progress`` so the driver keeps polling.
|
||||
"""
|
||||
state = (fw_state or "").upper()
|
||||
if state.startswith("JOB_STATE_"):
|
||||
state = state[len("JOB_STATE_") :]
|
||||
if state == "COMPLETED":
|
||||
return "completed"
|
||||
if state == "FAILED":
|
||||
return "failed"
|
||||
if state in ("CANCELLED", "CANCELED"):
|
||||
return "cancelled"
|
||||
if state == "EXPIRED":
|
||||
return "expired"
|
||||
return "in_progress"
|
||||
|
||||
@staticmethod
|
||||
def _normalize_output_line(line: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Fireworks output JSONL line -> OpenAI-batch-output shape.
|
||||
|
||||
Target: ``{"custom_id", "response": {"body": <chat-completion>}, "error"}``
|
||||
so the consumer's ``result["response"]["body"]["choices"][0]...`` works.
|
||||
"""
|
||||
custom_id = line.get("custom_id")
|
||||
error = line.get("error")
|
||||
if error:
|
||||
return {"custom_id": custom_id, "response": None, "error": error}
|
||||
|
||||
response = line.get("response")
|
||||
if response is None:
|
||||
response = line.get("body")
|
||||
# If Fireworks already nests the completion under `body`, unwrap it;
|
||||
# otherwise the `response` object *is* the completion.
|
||||
if isinstance(response, dict) and "body" in response:
|
||||
body = response["body"]
|
||||
else:
|
||||
body = response
|
||||
return {"custom_id": custom_id, "response": {"body": body}, "error": None}
|
||||
|
||||
# ----- low-level HTTP + URL helpers ----------------------------------
|
||||
|
||||
def _require_account_id(self) -> None:
|
||||
if not self._account_id:
|
||||
raise ValueError(
|
||||
"Fireworks batch inference requires an account id. "
|
||||
"Set HINDSIGHT_API_FIREWORKS_ACCOUNT_ID to your Fireworks account id."
|
||||
)
|
||||
|
||||
def _auth_headers(self) -> dict[str, str]:
|
||||
return {"Authorization": f"Bearer {self.api_key}"}
|
||||
|
||||
def _http(self) -> httpx.AsyncClient:
|
||||
if self._http_client is None:
|
||||
self._http_client = httpx.AsyncClient(timeout=httpx.Timeout(_HTTP_TIMEOUT_SECONDS))
|
||||
return self._http_client
|
||||
|
||||
async def _request(
|
||||
self,
|
||||
method: str,
|
||||
url: str,
|
||||
*,
|
||||
headers: dict[str, str] | None = None,
|
||||
json: dict[str, Any] | None = None,
|
||||
files: dict[str, Any] | None = None,
|
||||
) -> httpx.Response:
|
||||
resp = await self._http().request(method, url, headers=headers, json=json, files=files)
|
||||
if resp.is_error:
|
||||
# Surface the API's error body. Fireworks returns JSON describing why a
|
||||
# 4xx/5xx happened; raise_for_status() alone discards it, which makes
|
||||
# failures (e.g. a malformed dataset/job request) undebuggable.
|
||||
raise httpx.HTTPStatusError(
|
||||
f"Fireworks API {resp.status_code} for {method} {url}: {resp.text[:2000]}",
|
||||
request=resp.request,
|
||||
response=resp,
|
||||
)
|
||||
return resp
|
||||
|
||||
def _accounts_base(self) -> str:
|
||||
return f"{self._batch_base_url}/v1/accounts/{self._account_id}"
|
||||
|
||||
def _datasets_url(self) -> str:
|
||||
return f"{self._accounts_base()}/datasets"
|
||||
|
||||
def _jobs_url(self) -> str:
|
||||
return f"{self._accounts_base()}/batchInferenceJobs"
|
||||
|
||||
def _job_url(self, job_id: str) -> str:
|
||||
return f"{self._jobs_url()}/{job_id}"
|
||||
|
||||
def _download_endpoint_url(self, dataset_short_id: str) -> str:
|
||||
return f"{self._datasets_url()}/{dataset_short_id}:getDownloadEndpoint"
|
||||
|
||||
def _dataset_resource(self, dataset_id: str) -> str:
|
||||
return f"accounts/{self._account_id}/datasets/{dataset_id}"
|
||||
|
||||
@staticmethod
|
||||
def _last_segment(resource_name: str | None) -> str | None:
|
||||
if not resource_name:
|
||||
return None
|
||||
return resource_name.rstrip("/").split("/")[-1]
|
||||
|
||||
@staticmethod
|
||||
def _elapsed_seconds(create_time: str | None) -> float | None:
|
||||
if not create_time:
|
||||
return None
|
||||
try:
|
||||
normalized = create_time.replace("Z", "+00:00")
|
||||
created = datetime.fromisoformat(normalized)
|
||||
if created.tzinfo is None:
|
||||
created = created.replace(tzinfo=timezone.utc)
|
||||
return (datetime.now(timezone.utc) - created).total_seconds()
|
||||
except (ValueError, TypeError):
|
||||
return None
|
||||
|
||||
|
||||
def _to_int(value: Any) -> int:
|
||||
"""Coerce Fireworks' string/int counts to int, defaulting to 0."""
|
||||
try:
|
||||
return int(value)
|
||||
except (ValueError, TypeError):
|
||||
return 0
|
||||
@@ -153,24 +153,11 @@ class MockLLM(LLMInterface):
|
||||
result = self._response_callback(messages, scope)
|
||||
elif self._mock_response is not None:
|
||||
result = self._mock_response
|
||||
elif scope == "retain_extract_facts" and skip_validation:
|
||||
# Fact extraction: return canned facts derived from user message text.
|
||||
# This allows tests using a mock LLM to get real facts into the DB
|
||||
# so retain → recall → reflect pipelines work end-to-end.
|
||||
result = self._build_mock_facts(messages)
|
||||
elif scope == "consolidation" and response_format is not None:
|
||||
# Consolidation: produce a single observation from the input facts
|
||||
# so the full pipeline (retain → consolidation → observation → recall) works.
|
||||
result = self._build_mock_consolidation(messages, response_format)
|
||||
elif scope == "memory_think":
|
||||
# Reflect: return a plausible text answer
|
||||
result = "Based on the available information, the answer is related to the context provided."
|
||||
elif response_format is not None:
|
||||
# Structured output: try to return a valid empty instance of the model
|
||||
# so that callers expecting e.g. response_format with defaults
|
||||
# get a valid instance rather than a crash on {"mock": True}.
|
||||
# Try to create a minimal valid instance of the response format
|
||||
try:
|
||||
result = response_format()
|
||||
# For Pydantic models, try to create with minimal valid data
|
||||
result = {"mock": True}
|
||||
except Exception:
|
||||
result = {"mock": True}
|
||||
else:
|
||||
@@ -256,12 +243,6 @@ class MockLLM(LLMInterface):
|
||||
else:
|
||||
result = LLMToolCallResult(content="mock response", finish_reason="stop")
|
||||
|
||||
# Set mock token usage on result if not already set
|
||||
if result.input_tokens == 0:
|
||||
result.input_tokens = 10
|
||||
if result.output_tokens == 0:
|
||||
result.output_tokens = 5
|
||||
|
||||
# Record span with mock values
|
||||
# Convert LLMToolCall objects to dicts for span recording
|
||||
tool_calls_dict = (
|
||||
@@ -285,92 +266,6 @@ class MockLLM(LLMInterface):
|
||||
|
||||
return result
|
||||
|
||||
@staticmethod
|
||||
def _build_mock_facts(messages: list[dict]) -> dict:
|
||||
"""Build a canned fact extraction response from the user message text.
|
||||
|
||||
Splits the input into sentence-like chunks and returns each as a separate
|
||||
world fact with a simple entity extracted from the first noun-like word.
|
||||
This is intentionally simplistic — it just needs to produce structurally
|
||||
valid facts so the rest of the pipeline (embedding, storage, recall) works.
|
||||
"""
|
||||
import re
|
||||
|
||||
user_text = ""
|
||||
for m in messages:
|
||||
if m.get("role") == "user":
|
||||
user_text = m.get("content", "")
|
||||
break
|
||||
|
||||
# Split on sentence boundaries: period followed by space/EOL (not mid-number), or newlines
|
||||
sentences = [s.strip() for s in re.split(r"(?<=\.)\s+|\n+", user_text) if s.strip() and len(s.strip()) > 10]
|
||||
|
||||
if not sentences:
|
||||
sentences = [user_text[:200] if user_text else "mock fact"]
|
||||
|
||||
facts = []
|
||||
for sentence in sentences[:10]: # Cap at 10 facts per chunk
|
||||
# Extract simple entities: capitalized words that aren't common words
|
||||
words = re.findall(r"\b[A-Z][a-z]+\b", sentence)
|
||||
entities = [{"text": w} for w in dict.fromkeys(words)][:5] # Dedupe, cap at 5
|
||||
|
||||
facts.append(
|
||||
{
|
||||
"what": sentence,
|
||||
"when": "N/A",
|
||||
"where": "N/A",
|
||||
"who": "N/A",
|
||||
"why": "N/A",
|
||||
"fact_kind": "conversation",
|
||||
"fact_type": "world",
|
||||
"entities": entities,
|
||||
}
|
||||
)
|
||||
|
||||
return {"facts": facts}
|
||||
|
||||
@staticmethod
|
||||
def _build_mock_consolidation(messages: list[dict], response_format: Any) -> Any:
|
||||
"""Build a mock consolidation response that creates one observation per fact.
|
||||
|
||||
Parses fact IDs from the consolidation prompt and creates one observation
|
||||
per fact, each referencing its source fact ID. This mimics real LLM behavior
|
||||
where distinct facts produce separate observations, preserving entity
|
||||
separation so pipeline tests (graph filtering, entity linking) work correctly.
|
||||
"""
|
||||
import re
|
||||
|
||||
user_text = ""
|
||||
for m in messages:
|
||||
if m.get("role") == "user":
|
||||
user_text = m.get("content", "")
|
||||
break
|
||||
|
||||
# Extract fact UUIDs from the prompt (format: "[<uuid>] <text>")
|
||||
fact_entries = re.findall(r"\[([0-9a-f-]{36})\]\s*(.+?)(?:\n|$)", user_text)
|
||||
|
||||
if not fact_entries:
|
||||
# No facts to consolidate — return empty response
|
||||
try:
|
||||
return response_format()
|
||||
except Exception:
|
||||
return {"creates": [], "updates": [], "deletes": []}
|
||||
|
||||
# Create one observation per fact to preserve entity separation
|
||||
creates = []
|
||||
for fact_id, fact_text in fact_entries:
|
||||
creates.append({"text": fact_text.strip(), "source_fact_ids": [fact_id]})
|
||||
|
||||
try:
|
||||
return response_format(
|
||||
creates=creates,
|
||||
updates=[],
|
||||
deletes=[],
|
||||
)
|
||||
except Exception:
|
||||
# Fallback if response_format constructor doesn't accept these args
|
||||
return {"creates": creates, "updates": [], "deletes": []}
|
||||
|
||||
async def cleanup(self) -> None:
|
||||
"""Clean up resources (no-op for mock provider)."""
|
||||
pass
|
||||
@@ -423,8 +318,6 @@ class MockLLM(LLMInterface):
|
||||
return self._mock_calls
|
||||
|
||||
def clear_mock_calls(self) -> None:
|
||||
"""Clear all recorded calls and any configured response/exception state."""
|
||||
"""Clear the recorded mock calls and any set exception."""
|
||||
self._mock_calls = []
|
||||
self._mock_exception = None
|
||||
self._mock_response = None
|
||||
self._response_callback = None
|
||||
|
||||
@@ -270,7 +270,6 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
"openai",
|
||||
"groq",
|
||||
"ollama",
|
||||
"ollama-cloud",
|
||||
"lmstudio",
|
||||
"llamacpp",
|
||||
"minimax",
|
||||
@@ -279,7 +278,6 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
"openrouter",
|
||||
"zai",
|
||||
"opencode-go",
|
||||
"fireworks",
|
||||
]
|
||||
if self.provider not in valid_providers:
|
||||
raise ValueError(f"OpenAICompatibleLLM only supports: {', '.join(valid_providers)}. Got: {self.provider}")
|
||||
@@ -290,8 +288,6 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
self.base_url = "https://api.groq.com/openai/v1"
|
||||
elif self.provider == "ollama":
|
||||
self.base_url = "http://localhost:11434/v1"
|
||||
elif self.provider == "ollama-cloud":
|
||||
self.base_url = "https://ollama.com/v1"
|
||||
elif self.provider == "lmstudio":
|
||||
self.base_url = "http://localhost:1234/v1"
|
||||
elif self.provider == "minimax":
|
||||
@@ -304,10 +300,6 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
self.base_url = "https://api.z.ai/api/coding/paas/v4"
|
||||
elif self.provider == "opencode-go":
|
||||
self.base_url = "https://opencode.ai/zen/go/v1"
|
||||
elif self.provider == "fireworks":
|
||||
# OpenAI-compatible inference host (online path). The batch API
|
||||
# lives on a separate control-plane host — see FireworksLLM.
|
||||
self.base_url = "https://api.fireworks.ai/inference/v1"
|
||||
|
||||
# For ollama/lmstudio, use dummy key if not provided
|
||||
if self.provider in ("ollama", "lmstudio") and not self.api_key:
|
||||
@@ -324,7 +316,6 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
"openrouter",
|
||||
"zai",
|
||||
"opencode-go",
|
||||
"ollama-cloud",
|
||||
)
|
||||
and not self.api_key
|
||||
):
|
||||
@@ -1082,17 +1073,12 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
|
||||
last_exception = None
|
||||
|
||||
# Pass API key as Bearer token for cloud Ollama endpoints
|
||||
headers: dict[str, str] = {}
|
||||
if self.api_key and self.api_key != "local":
|
||||
headers["Authorization"] = f"Bearer {self.api_key}"
|
||||
|
||||
async with httpx.AsyncClient(timeout=300.0) as client:
|
||||
for attempt in range(max_retries + 1):
|
||||
if attempt > 0:
|
||||
set_stage(f"llm.ollama_native.{scope}.attempt={attempt + 1}/{max_retries + 1}")
|
||||
try:
|
||||
response = await client.post(native_url, json=payload, headers=headers)
|
||||
response = await client.post(native_url, json=payload)
|
||||
response.raise_for_status()
|
||||
|
||||
result = response.json()
|
||||
|
||||
@@ -321,7 +321,6 @@ async def run_reflect_agent(
|
||||
include_recall: bool = True,
|
||||
budget: str | None = None,
|
||||
max_context_tokens: int = 100_000,
|
||||
llm_output_language: str | None = None,
|
||||
) -> ReflectAgentResult:
|
||||
"""
|
||||
Execute the reflect agent loop using native tool calling.
|
||||
@@ -370,12 +369,7 @@ async def run_reflect_agent(
|
||||
|
||||
# Build initial messages (directives are injected into system prompt at START and END)
|
||||
system_prompt = build_system_prompt_for_tools(
|
||||
bank_profile,
|
||||
context,
|
||||
directives=directives,
|
||||
has_mental_models=has_mental_models,
|
||||
include_observations=include_observations,
|
||||
budget=budget,
|
||||
bank_profile, context, directives=directives, has_mental_models=has_mental_models, budget=budget
|
||||
)
|
||||
messages: list[dict[str, Any]] = [
|
||||
{"role": "system", "content": system_prompt},
|
||||
@@ -453,10 +447,7 @@ async def run_reflect_agent(
|
||||
llm_start = time.time()
|
||||
response, usage = await llm_config.call(
|
||||
messages=[
|
||||
{
|
||||
"role": "system",
|
||||
"content": build_final_system_prompt(bank_profile.get("mission"), llm_output_language),
|
||||
},
|
||||
{"role": "system", "content": build_final_system_prompt(bank_profile.get("mission"))},
|
||||
{"role": "user", "content": prompt},
|
||||
],
|
||||
scope="reflect",
|
||||
@@ -513,10 +504,7 @@ async def run_reflect_agent(
|
||||
llm_start = time.time()
|
||||
response, usage = await llm_config.call(
|
||||
messages=[
|
||||
{
|
||||
"role": "system",
|
||||
"content": build_final_system_prompt(bank_profile.get("mission"), llm_output_language),
|
||||
},
|
||||
{"role": "system", "content": build_final_system_prompt(bank_profile.get("mission"))},
|
||||
{"role": "user", "content": prompt},
|
||||
],
|
||||
scope="reflect",
|
||||
@@ -619,10 +607,7 @@ async def run_reflect_agent(
|
||||
llm_start = time.time()
|
||||
response, usage = await llm_config.call(
|
||||
messages=[
|
||||
{
|
||||
"role": "system",
|
||||
"content": build_final_system_prompt(bank_profile.get("mission"), llm_output_language),
|
||||
},
|
||||
{"role": "system", "content": build_final_system_prompt(bank_profile.get("mission"))},
|
||||
{"role": "user", "content": prompt},
|
||||
],
|
||||
scope="reflect",
|
||||
@@ -743,10 +728,7 @@ async def run_reflect_agent(
|
||||
llm_start = time.time()
|
||||
response, usage = await llm_config.call(
|
||||
messages=[
|
||||
{
|
||||
"role": "system",
|
||||
"content": build_final_system_prompt(bank_profile.get("mission"), llm_output_language),
|
||||
},
|
||||
{"role": "system", "content": build_final_system_prompt(bank_profile.get("mission"))},
|
||||
{"role": "user", "content": prompt},
|
||||
],
|
||||
scope="reflect",
|
||||
|
||||
@@ -98,7 +98,6 @@ def build_system_prompt_for_tools(
|
||||
context: str | None = None,
|
||||
directives: list[dict[str, Any]] | None = None,
|
||||
has_mental_models: bool = False,
|
||||
include_observations: bool = True,
|
||||
budget: str | None = None,
|
||||
) -> str:
|
||||
"""
|
||||
@@ -109,17 +108,11 @@ def build_system_prompt_for_tools(
|
||||
2. search_observations - Consolidated knowledge with freshness
|
||||
3. recall - Raw facts as ground truth
|
||||
|
||||
The retrieval-strategy and workflow sections are built to match the tools
|
||||
actually exposed to the LLM — mentioning a tool the agent has disabled
|
||||
causes weaker LLMs to either hallucinate the call (rejected by the agent)
|
||||
or give up with "I cannot find any information…" (see #1724).
|
||||
|
||||
Args:
|
||||
bank_profile: Bank profile with name and mission
|
||||
context: Optional additional context
|
||||
directives: Optional list of directive mental models to inject as hard rules
|
||||
has_mental_models: Whether the bank has any mental models (skip if not)
|
||||
include_observations: Whether search_observations is in the tool list.
|
||||
budget: Search depth budget - "low", "mid", or "high". Controls exploration thoroughness.
|
||||
"""
|
||||
name = bank_profile.get("name", "Assistant")
|
||||
@@ -165,137 +158,56 @@ def build_system_prompt_for_tools(
|
||||
"- If memories mention someone did an activity, you can infer they likely enjoyed it",
|
||||
"- Synthesize a coherent narrative from related memories",
|
||||
"- Be a thoughtful interpreter, not just a literal repeater",
|
||||
"- When the exact answer isn't stated, use what IS stated to give a best-effort answer AND surface any uncertainty — never invent confidence the data doesn't support.",
|
||||
"",
|
||||
"## Temporal Reasoning",
|
||||
"Every memory and observation carries temporal fields in the JSON tool result:",
|
||||
"- `mentioned_at` — when the user retained the fact (always set).",
|
||||
"- `occurred_start` / `occurred_end` — when the underlying event happened (optional, set for dated events).",
|
||||
"",
|
||||
"When facts about the SAME facet conflict — counts, statuses, ownership, location, presence, etc. — the fact with the LATEST `mentioned_at` is authoritative. Later statements SUPERSEDE earlier ones. Do NOT average, sum, or favor an explicitly-dated fact over a more recent one.",
|
||||
"",
|
||||
"Example: three count facts come back from recall:",
|
||||
" - 'Team has 2 engineers' (mentioned_at=T1)",
|
||||
" - 'Team now has 1 engineer' (mentioned_at=T2, occurred_start=2026-05-25)",
|
||||
" - 'Team has 5 engineers' (mentioned_at=T3)",
|
||||
"with T1 < T2 < T3. The current size is 5, not 1. Then apply later events (e.g. someone leaving after T3) on top of that.",
|
||||
"",
|
||||
"For reconstructing a TIMELINE of events, order by `occurred_start` / `occurred_end` (when things happened), not `mentioned_at` (when they were retained).",
|
||||
"",
|
||||
"## Conflicts and Ambiguity",
|
||||
"Not every retrieval converges on a single answer. Distinguish two cases:",
|
||||
"",
|
||||
"- RESOLVABLE conflict — the temporal rule above (latest `mentioned_at` wins) cleanly picks a winner. Apply it and move on.",
|
||||
"- UNRESOLVABLE ambiguity — the data is internally inconsistent in a way the temporal rule does NOT settle. Examples: a recent aggregate (count, total) is incompatible with the individual entities you can enumerate; two equally-recent facts disagree and no later fact resolves them; events are described but their relative order is unclear; the user's own statements contradict each other and nothing later reconciles them.",
|
||||
"",
|
||||
"When the data is genuinely ambiguous: SAY SO in your answer. Name the conflicting facts. Explain why they can't be reconciled. Give a range or a best-effort interpretation with explicit uncertainty (e.g. 'between X and Y, depending on [unresolved condition]'; or 'the most recent statement says A, but B was stated earlier and the gap isn't accounted for in any later fact').",
|
||||
"",
|
||||
"An honest 'the data is inconsistent about X' beats a confident wrong answer. Do NOT pick a value arbitrarily, average conflicting values, or smooth over gaps in confident prose. Acknowledging ambiguity is a successful answer, not a failure mode.",
|
||||
"",
|
||||
"## Showing Your Reasoning",
|
||||
"For any answer that resolves a conflict between facts, applies events on top of a count or status, or settles an ambiguity — show your work in the answer text so a reader can audit it.",
|
||||
"",
|
||||
"Walk through these steps explicitly:",
|
||||
"1. **List the relevant facts in `mentioned_at` order (oldest → newest)**, each with the value it asserts. Use a short bulleted list.",
|
||||
"2. **Identify the authoritative fact** under the temporal rule (latest `mentioned_at` for the contested facet). Write its date down.",
|
||||
"3. **List candidate events to apply on top** — anything that changes the count, status, or state being asked about. Write each event's date down next to it.",
|
||||
"4. **Sanity-check each candidate event against the authoritative date** — for EVERY event from step 3, write a one-line check in the form `<event> (<event_date>) vs authoritative (<authoritative_date>) → BEFORE/AFTER → KEEP/DROP`. If the event is BEFORE or EQUAL to the authoritative date, DROP it: it is already reflected in the authoritative fact, and applying it again is double-counting. This is the single most common mistake — do not skip this step even if you feel confident.",
|
||||
"5. **Show the arithmetic or derivation explicitly** using only the KEEP events from step 4 — e.g. 'authoritative count = 5 (at 2025-02-12); kept events: Shadow died (2025-03-12, AFTER); 5 − 1 = 4'.",
|
||||
"6. If step 2 or 3 cannot be done cleanly (no clear winner, overlapping timestamps, unclear event order), STOP and surface this as an UNRESOLVABLE ambiguity per the section above — do not fabricate a derivation.",
|
||||
"",
|
||||
"For simple factual lookups that don't involve conflict or arithmetic, you can answer directly without this scaffolding.",
|
||||
"- When the exact answer isn't stated, use what IS stated to give the best answer",
|
||||
"",
|
||||
"## HIERARCHICAL RETRIEVAL STRATEGY",
|
||||
"",
|
||||
]
|
||||
)
|
||||
|
||||
# Assemble the retrieval-level blocks for whatever tools are exposed.
|
||||
# MM and Observations bodies are unconditional; recall's fallback wording
|
||||
# adapts to which upstream tools precede it (telling the LLM to fall back
|
||||
# to a tool that isn't in its list is the bug at the root of #1724).
|
||||
levels: list[tuple[str, list[str]]] = []
|
||||
# Build retrieval levels based on what's available
|
||||
if has_mental_models:
|
||||
levels.append(
|
||||
(
|
||||
"MENTAL MODELS (search_mental_models)",
|
||||
[
|
||||
"- User-curated summaries about specific topics",
|
||||
"- HIGHEST quality - manually created and maintained",
|
||||
"- If a relevant mental model exists and is FRESH, it may fully answer the question",
|
||||
"- Check `is_stale` field - if stale, also verify with lower levels",
|
||||
],
|
||||
)
|
||||
)
|
||||
if include_observations:
|
||||
levels.append(
|
||||
(
|
||||
"OBSERVATIONS (search_observations)",
|
||||
[
|
||||
"- Auto-consolidated knowledge from memories",
|
||||
"- Check `is_stale` field - if stale, ALSO use recall() to verify",
|
||||
"- Good for understanding patterns and summaries",
|
||||
],
|
||||
)
|
||||
)
|
||||
recall_body = ["- Individual memories (world facts and experiences)"]
|
||||
if has_mental_models and include_observations:
|
||||
recall_body.extend(
|
||||
parts.extend(
|
||||
[
|
||||
"You have access to THREE levels of knowledge. Use them in this order:",
|
||||
"",
|
||||
"### 1. MENTAL MODELS (search_mental_models) - Try First",
|
||||
"- User-curated summaries about specific topics",
|
||||
"- HIGHEST quality - manually created and maintained",
|
||||
"- If a relevant mental model exists and is FRESH, it may fully answer the question",
|
||||
"- Check `is_stale` field - if stale, also verify with lower levels",
|
||||
"",
|
||||
"### 2. OBSERVATIONS (search_observations) - Second Priority",
|
||||
"- Auto-consolidated knowledge from memories",
|
||||
"- Check `is_stale` field - if stale, ALSO use recall() to verify",
|
||||
"- Good for understanding patterns and summaries",
|
||||
"",
|
||||
"### 3. RAW FACTS (recall) - Ground Truth",
|
||||
"- Individual memories (world facts and experiences)",
|
||||
"- Use when: no mental models/observations exist, they're stale, or you need specific details",
|
||||
"- MANDATORY: If search_mental_models and search_observations both return 0 results, you MUST call recall() before giving up",
|
||||
"- This is the source of truth that other levels are built from",
|
||||
"",
|
||||
"**Tool result ordering:** `recall()` and `search_observations()` return their `memories` / `observations` arrays sorted by SEMANTIC RELEVANCE to the query, NOT by time. The POSITION of an entry tells you nothing about when it was retained. For any temporal reasoning — recency, supersession, applying events on top of a state — IGNORE the position and read the per-entry `mentioned_at` field (and `occurred_start` / `occurred_end` for events).",
|
||||
]
|
||||
)
|
||||
else:
|
||||
parts.extend(
|
||||
[
|
||||
"You have access to TWO levels of knowledge. Use them in this order:",
|
||||
"",
|
||||
]
|
||||
)
|
||||
elif has_mental_models:
|
||||
recall_body.extend(
|
||||
[
|
||||
"- Use when: no mental model exists, it's stale, or you need specific details",
|
||||
"- MANDATORY: If search_mental_models returns 0 results, you MUST call recall() before giving up",
|
||||
"- This is the source of truth that mental models are built from",
|
||||
]
|
||||
)
|
||||
elif include_observations:
|
||||
recall_body.extend(
|
||||
[
|
||||
"### 1. OBSERVATIONS (search_observations) - Try First",
|
||||
"- Auto-consolidated knowledge from memories",
|
||||
"- Check `is_stale` field - if stale, ALSO use recall() to verify",
|
||||
"- Good for understanding patterns and summaries",
|
||||
"",
|
||||
"### 2. RAW FACTS (recall) - Ground Truth",
|
||||
"- Individual memories (world facts and experiences)",
|
||||
"- Use when: no observations exist, they're stale, or you need specific details",
|
||||
"- MANDATORY: If search_observations returns 0 results or count=0, you MUST call recall() before giving up",
|
||||
"- This is the source of truth that observations are built from",
|
||||
"",
|
||||
"**Tool result ordering:** `recall()` and `search_observations()` return their `memories` / `observations` arrays sorted by SEMANTIC RELEVANCE to the query, NOT by time. The POSITION of an entry tells you nothing about when it was retained. For any temporal reasoning — recency, supersession, applying events on top of a state — IGNORE the position and read the per-entry `mentioned_at` field (and `occurred_start` / `occurred_end` for events).",
|
||||
"",
|
||||
]
|
||||
)
|
||||
else:
|
||||
recall_body.extend(
|
||||
[
|
||||
"- MANDATORY: Call recall() to gather facts before giving up",
|
||||
"- This is the source of truth.",
|
||||
]
|
||||
)
|
||||
levels.append(("RAW FACTS (recall) - Ground Truth", recall_body))
|
||||
|
||||
# Position-dependent suffix for upstream tools; recall already carries its
|
||||
# fixed "- Ground Truth" suffix in the header text.
|
||||
suffixes = [""] * len(levels)
|
||||
if len(levels) >= 2:
|
||||
suffixes[0] = " - Try First"
|
||||
if len(levels) == 3:
|
||||
suffixes[1] = " - Second Priority"
|
||||
|
||||
if len(levels) == 1:
|
||||
parts.append("You have access to ONE level of knowledge:")
|
||||
else:
|
||||
word = "TWO" if len(levels) == 2 else "THREE"
|
||||
parts.append(f"You have access to {word} levels of knowledge. Use them in this order:")
|
||||
parts.append("")
|
||||
for idx, ((header, body), suffix) in enumerate(zip(levels, suffixes), 1):
|
||||
parts.append(f"### {idx}. {header}{suffix}")
|
||||
parts.extend(body)
|
||||
parts.append("")
|
||||
|
||||
parts.extend(
|
||||
[
|
||||
@@ -355,28 +267,25 @@ def build_system_prompt_for_tools(
|
||||
|
||||
parts.append("## Workflow")
|
||||
|
||||
steps: list[str] = []
|
||||
if has_mental_models:
|
||||
steps.append("First, try search_mental_models() - check if a curated summary exists")
|
||||
if include_observations:
|
||||
if has_mental_models:
|
||||
steps.append("If no mental model or it's stale, try search_observations() for consolidated knowledge")
|
||||
else:
|
||||
steps.append("First, try search_observations() - check for consolidated knowledge")
|
||||
# Recall step phrasing varies with whichever upstream tool(s) precede it.
|
||||
if include_observations:
|
||||
steps.append(
|
||||
"If observations are stale OR you need specific details, use recall() for raw facts"
|
||||
if has_mental_models
|
||||
else "If search_observations returns 0 results OR observations are stale, you MUST call recall() for raw facts"
|
||||
parts.extend(
|
||||
[
|
||||
"1. First, try search_mental_models() - check if a curated summary exists",
|
||||
"2. If no mental model or it's stale, try search_observations() for consolidated knowledge",
|
||||
"3. If observations are stale OR you need specific details, use recall() for raw facts",
|
||||
"4. Use expand() if you need more context on specific memories",
|
||||
"5. When ready, call done() with your answer and supporting IDs",
|
||||
]
|
||||
)
|
||||
elif has_mental_models:
|
||||
steps.append("If no mental model or it's stale, use recall() for raw facts")
|
||||
else:
|
||||
steps.append("Call recall() to gather raw facts")
|
||||
steps.append("Use expand() if you need more context on specific memories")
|
||||
steps.append("When ready, call done() with your answer and supporting IDs")
|
||||
parts.extend(f"{idx}. {step}" for idx, step in enumerate(steps, 1))
|
||||
parts.extend(
|
||||
[
|
||||
"1. First, try search_observations() - check for consolidated knowledge",
|
||||
"2. If search_observations returns 0 results OR observations are stale, you MUST call recall() for raw facts",
|
||||
"3. Use expand() if you need more context on specific memories",
|
||||
"4. When ready, call done() with your answer and supporting IDs",
|
||||
]
|
||||
)
|
||||
|
||||
parts.extend(
|
||||
[
|
||||
@@ -604,16 +513,10 @@ Just provide the direct answer with proper markdown formatting.
|
||||
CRITICAL: This is a NON-CONVERSATIONAL system. NEVER ask follow-up questions, offer to search again, suggest alternatives, or end with anything like "Would you like me to..." or "Let me know if...". The user cannot reply. Your answer must be complete and self-contained."""
|
||||
|
||||
|
||||
def build_final_system_prompt(mission: str | None = None, llm_output_language: str | None = None) -> str:
|
||||
"""Build the final synthesis system prompt, using mission as role when set.
|
||||
|
||||
When ``llm_output_language`` is set, the response is forced into that
|
||||
language regardless of the query/source language.
|
||||
"""
|
||||
from hindsight_api.engine.prompt_utils import escape_for_prompt, output_language_directive
|
||||
|
||||
role_section = escape_for_prompt(mission.strip()) if mission else _DEFAULT_FINAL_ROLE
|
||||
return _FINAL_SYSTEM_PROMPT_BASE.format(role_section=role_section) + output_language_directive(llm_output_language)
|
||||
def build_final_system_prompt(mission: str | None = None) -> str:
|
||||
"""Build the final synthesis system prompt, using mission as role when set."""
|
||||
role_section = mission.strip() if mission else _DEFAULT_FINAL_ROLE
|
||||
return _FINAL_SYSTEM_PROMPT_BASE.format(role_section=role_section)
|
||||
|
||||
|
||||
# Backward-compatible constant for non-identity missions
|
||||
|
||||
@@ -1,8 +1,17 @@
|
||||
"""Token counting helpers for reflect prompts and agent control flow."""
|
||||
|
||||
from ..token_encoding import count_tokens as _count_tokens
|
||||
from functools import lru_cache
|
||||
|
||||
import tiktoken
|
||||
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def _get_cl100k_base_encoding() -> tiktoken.Encoding:
|
||||
# tiktoken downloads this encoding on first lookup when it is not cached.
|
||||
# Keep the lookup lazy so importing hindsight_api does not depend on network access.
|
||||
return tiktoken.get_encoding("cl100k_base")
|
||||
|
||||
|
||||
def count_cl100k_tokens(text: str) -> int:
|
||||
"""Return the number of cl100k_base tokens in text."""
|
||||
return _count_tokens(text)
|
||||
return len(_get_cl100k_base_encoding().encode(text))
|
||||
|
||||
@@ -23,24 +23,6 @@ if TYPE_CHECKING:
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _prune_nulls(d: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Drop keys whose value is None or an empty collection (``""``, ``[]``, ``{}``).
|
||||
|
||||
Reflect tools dump ``MemoryFact`` / ``ObservationResult`` via ``model_dump()``,
|
||||
which emits every field including the many that are typically null or empty
|
||||
(``context``, ``occurred_start``, ``metadata``, ``tags``, etc.). Stripping
|
||||
these before serializing to JSON for the LLM cuts token cost and removes
|
||||
fields that aren't telling the model anything.
|
||||
|
||||
Callers that need the *presence* of a specific field as a signal (e.g.
|
||||
``source_fact_ids`` for drill-down) must ensure the value is non-empty —
|
||||
pass the upstream flag that populates it (e.g. ``source_facts_max_tokens``
|
||||
> 0 on ``tool_search_observations``) rather than relying on Pydantic
|
||||
emitting ``None``.
|
||||
"""
|
||||
return {k: v for k, v in d.items() if v is not None and v != "" and v != [] and v != {}}
|
||||
|
||||
|
||||
def _document_metadata_from_retain_params(retain_params: Any) -> dict[str, Any] | None:
|
||||
"""Return document metadata stored under retain_params.metadata."""
|
||||
if isinstance(retain_params, str):
|
||||
@@ -232,8 +214,8 @@ async def tool_search_observations(
|
||||
return {
|
||||
"query": query,
|
||||
"count": len(result.results),
|
||||
"observations": [_prune_nulls(m.model_dump()) for m in result.results],
|
||||
"source_facts": {k: _prune_nulls(v.model_dump()) for k, v in (result.source_facts or {}).items()},
|
||||
"observations": [m.model_dump() for m in result.results],
|
||||
"source_facts": {k: v.model_dump() for k, v in (result.source_facts or {}).items()},
|
||||
"is_stale": is_stale,
|
||||
"freshness": freshness,
|
||||
}
|
||||
@@ -300,8 +282,8 @@ async def tool_recall(
|
||||
|
||||
return {
|
||||
"query": query,
|
||||
"memories": [_prune_nulls(m.model_dump()) for m in result.results],
|
||||
"chunks": {k: _prune_nulls(v.model_dump()) for k, v in (result.chunks or {}).items()},
|
||||
"memories": [m.model_dump() for m in result.results],
|
||||
"chunks": {k: v.model_dump() for k, v in (result.chunks or {}).items()},
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -4,54 +4,29 @@ Embedding generation utilities for memory units.
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from typing import Literal, Protocol
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
EmbeddingInputType = Literal["document", "query"]
|
||||
|
||||
|
||||
class EmbeddingsBackend(Protocol):
|
||||
"""Minimal duck-typed surface used by retain/recall — the concrete `Embeddings`
|
||||
ABC supplies default implementations that delegate to `encode()`."""
|
||||
|
||||
def encode_query(self, texts: list[str]) -> list[list[float]]: ...
|
||||
|
||||
def encode_documents(self, texts: list[str]) -> list[list[float]]: ...
|
||||
|
||||
|
||||
def generate_embedding(
|
||||
embeddings_backend: EmbeddingsBackend, text: str, input_type: EmbeddingInputType = "document"
|
||||
) -> list[float]:
|
||||
def generate_embedding(embeddings_backend, text: str) -> list[float]:
|
||||
"""
|
||||
Generate embedding for text using the provided embeddings backend.
|
||||
|
||||
Args:
|
||||
embeddings_backend: Embeddings instance to use for encoding
|
||||
text: Text to embed
|
||||
input_type: Whether text is retained document text or recall/search query text.
|
||||
|
||||
Returns:
|
||||
Embedding vector (dimension depends on embeddings backend)
|
||||
"""
|
||||
try:
|
||||
embeddings = _encode_with_input_type(embeddings_backend, [text], input_type)
|
||||
embeddings = embeddings_backend.encode([text])
|
||||
return embeddings[0]
|
||||
except Exception as e:
|
||||
raise Exception(f"Failed to generate embedding: {str(e)}")
|
||||
|
||||
|
||||
def _encode_with_input_type(
|
||||
embeddings_backend: EmbeddingsBackend, texts: list[str], input_type: EmbeddingInputType
|
||||
) -> list[list[float]]:
|
||||
if input_type == "query":
|
||||
return embeddings_backend.encode_query(texts)
|
||||
return embeddings_backend.encode_documents(texts)
|
||||
|
||||
|
||||
async def generate_embeddings_batch(
|
||||
embeddings_backend: EmbeddingsBackend, texts: list[str], input_type: EmbeddingInputType = "document"
|
||||
) -> list[list[float]]:
|
||||
async def generate_embeddings_batch(embeddings_backend, texts: list[str]) -> list[list[float]]:
|
||||
"""
|
||||
Generate embeddings for multiple texts using the provided embeddings backend.
|
||||
|
||||
@@ -61,14 +36,17 @@ async def generate_embeddings_batch(
|
||||
Args:
|
||||
embeddings_backend: Embeddings instance to use for encoding
|
||||
texts: List of texts to embed
|
||||
input_type: Whether texts are retained documents or recall/search queries.
|
||||
|
||||
Returns:
|
||||
List of embeddings in same order as input texts
|
||||
"""
|
||||
try:
|
||||
loop = asyncio.get_event_loop()
|
||||
embeddings = await loop.run_in_executor(None, _encode_with_input_type, embeddings_backend, texts, input_type)
|
||||
embeddings = await loop.run_in_executor(
|
||||
None,
|
||||
embeddings_backend.encode,
|
||||
texts,
|
||||
)
|
||||
except Exception as e:
|
||||
raise Exception(f"Failed to generate batch embeddings: {str(e)}")
|
||||
|
||||
|
||||
@@ -1,13 +1,13 @@
|
||||
"""
|
||||
Entity processing for retain pipeline.
|
||||
|
||||
Handles entity extraction and resolution for stored facts.
|
||||
Handles entity extraction, resolution, and link creation for stored facts.
|
||||
"""
|
||||
|
||||
import logging
|
||||
|
||||
from . import link_utils
|
||||
from .types import ProcessedFact
|
||||
from .types import EntityLink, ProcessedFact
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -76,7 +76,8 @@ async def resolve_entities(
|
||||
entity_labels: Optional entity label taxonomy
|
||||
|
||||
Returns:
|
||||
Tuple of (resolved_entity_ids, entity_to_unit, unit_to_entity_ids).
|
||||
Tuple of (resolved_entity_ids, entity_to_unit, unit_to_entity_ids)
|
||||
to pass to build_entity_links().
|
||||
"""
|
||||
if not unit_ids or not facts:
|
||||
return [], [], {}
|
||||
@@ -98,3 +99,68 @@ async def resolve_entities(
|
||||
log_buffer,
|
||||
entity_labels=entity_labels,
|
||||
)
|
||||
|
||||
|
||||
async def build_entity_links(
|
||||
entity_resolver,
|
||||
conn,
|
||||
bank_id: str,
|
||||
unit_ids: list[str],
|
||||
resolved_entity_ids: list[str],
|
||||
entity_to_unit: list[tuple],
|
||||
unit_to_entity_ids: dict[str, list[str]],
|
||||
log_buffer: list[str] = None,
|
||||
skip_unit_entities_insert: bool = False,
|
||||
ops=None,
|
||||
) -> list[EntityLink]:
|
||||
"""
|
||||
Build entity links for UI graph visualization.
|
||||
|
||||
Queries unit_entities to find shared entities between new and existing units,
|
||||
then generates EntityLink objects. When called from Phase 3 (post-transaction),
|
||||
set skip_unit_entities_insert=True since unit_entities were already inserted
|
||||
in Phase 2.
|
||||
|
||||
Args:
|
||||
entity_resolver: EntityResolver instance
|
||||
conn: Database connection
|
||||
bank_id: Bank identifier
|
||||
unit_ids: Actual unit IDs (must already be inserted in the DB)
|
||||
resolved_entity_ids: From resolve_entities()
|
||||
entity_to_unit: From resolve_entities()
|
||||
unit_to_entity_ids: From resolve_entities()
|
||||
log_buffer: Optional buffer for detailed logging
|
||||
skip_unit_entities_insert: Skip unit_entities INSERT (already done in Phase 2)
|
||||
ops: DataAccessOps instance (from backend.ops)
|
||||
|
||||
Returns:
|
||||
List of EntityLink objects for batch insertion
|
||||
"""
|
||||
return await link_utils.build_entity_links_from_resolved(
|
||||
entity_resolver,
|
||||
conn,
|
||||
bank_id,
|
||||
unit_ids,
|
||||
resolved_entity_ids,
|
||||
entity_to_unit,
|
||||
unit_to_entity_ids,
|
||||
log_buffer,
|
||||
skip_unit_entities_insert=skip_unit_entities_insert,
|
||||
ops=ops,
|
||||
)
|
||||
|
||||
|
||||
async def insert_entity_links_batch(conn, entity_links: list[EntityLink], bank_id: str, ops=None) -> None:
|
||||
"""
|
||||
Insert entity links in batch.
|
||||
|
||||
Args:
|
||||
conn: Database connection
|
||||
entity_links: List of EntityLink objects
|
||||
bank_id: Bank identifier (stored directly on memory_links for fast filtering)
|
||||
ops: DataAccessOps instance (from backend.ops)
|
||||
"""
|
||||
if not entity_links:
|
||||
return
|
||||
|
||||
await link_utils.insert_entity_links_batch(conn, entity_links, bank_id, ops=ops)
|
||||
|
||||
@@ -888,16 +888,13 @@ def _build_extraction_prompt_and_schema(config) -> tuple[str, type]:
|
||||
extract_causal_links = config.retain_extract_causal_links
|
||||
|
||||
# Build retain_mission section if set - injected before the mode-specific guidelines
|
||||
# Escape braces so user-supplied text survives str.format() on the prompt template.
|
||||
from hindsight_api.engine.prompt_utils import escape_for_prompt
|
||||
|
||||
retain_mission = getattr(config, "retain_mission", None)
|
||||
if retain_mission:
|
||||
retain_mission_section = (
|
||||
f"══════════════════════════════════════════════════════════════════════════\n"
|
||||
f"FOCUS — What to retain for this bank\n"
|
||||
f"══════════════════════════════════════════════════════════════════════════\n\n"
|
||||
f"{escape_for_prompt(retain_mission)}\n\n"
|
||||
f"{retain_mission}\n\n"
|
||||
)
|
||||
else:
|
||||
retain_mission_section = ""
|
||||
@@ -913,7 +910,7 @@ def _build_extraction_prompt_and_schema(config) -> tuple[str, type]:
|
||||
base_prompt = CUSTOM_FACT_EXTRACTION_PROMPT
|
||||
prompt = base_prompt.format(
|
||||
retain_mission_section=retain_mission_section,
|
||||
custom_instructions=escape_for_prompt(config.retain_custom_instructions),
|
||||
custom_instructions=config.retain_custom_instructions,
|
||||
)
|
||||
elif extraction_mode == "verbose":
|
||||
prompt = VERBOSE_FACT_EXTRACTION_PROMPT.format(
|
||||
@@ -950,16 +947,6 @@ def _build_extraction_prompt_and_schema(config) -> tuple[str, type]:
|
||||
if labels_section:
|
||||
prompt = prompt + labels_section
|
||||
|
||||
# Force the LLM to emit fact text in the configured language, regardless of
|
||||
# the source content's language. Same directive is applied to consolidation
|
||||
# and reflect so HINDSIGHT_API_LLM_OUTPUT_LANGUAGE has a uniform effect
|
||||
# across the pipeline. This is independent of the BM25 indexing language
|
||||
# (HINDSIGHT_API_TEXT_SEARCH_EXTENSION_NATIVE_LANGUAGE) by design — search
|
||||
# tokenization and LLM output language are separate concerns.
|
||||
from ..prompt_utils import output_language_directive
|
||||
|
||||
prompt = prompt + output_language_directive(getattr(config, "llm_output_language", None))
|
||||
|
||||
response_schema = base_response_class
|
||||
|
||||
if labels_cfg and labels_cfg.attributes:
|
||||
@@ -1142,14 +1129,11 @@ async def _extract_facts_from_chunk(
|
||||
)
|
||||
continue
|
||||
else:
|
||||
# A non-dict response is malformed (the schema is {"facts": [...]}).
|
||||
# Raise instead of returning [] so the failure propagates to the
|
||||
# worker's retry machinery and ultimately fails loudly — never
|
||||
# silently commit the document with 0 facts. See issue #1833.
|
||||
raise RuntimeError(
|
||||
f"Fact extraction failed: LLM returned non-dict JSON after {llm_max_retries} attempts "
|
||||
f"({type(extraction_response_json).__name__}). Raw: {str(extraction_response_json)[:500]}"
|
||||
logger.warning(
|
||||
f"LLM returned non-dict JSON after {llm_max_retries} attempts: {type(extraction_response_json).__name__}. "
|
||||
f"Raw: {str(extraction_response_json)[:500]}"
|
||||
)
|
||||
return [], usage
|
||||
|
||||
raw_facts = extraction_response_json.get("facts", [])
|
||||
|
||||
@@ -1678,11 +1662,8 @@ async def extract_facts_from_contents_batch_api(
|
||||
|
||||
# Check if provider supports batch API
|
||||
if not await llm_config._provider_impl.supports_batch_api():
|
||||
raise RuntimeError(
|
||||
f"retain_batch_enabled=True but provider '{llm_config.provider}' does not "
|
||||
f"support the batch API. This should have been caught at startup — check "
|
||||
f"HINDSIGHT_API_RETAIN_BATCH_ENABLED and your LLM provider configuration."
|
||||
)
|
||||
logger.warning(f"Batch API not supported for provider {llm_config.provider}, falling back to sync mode")
|
||||
return await extract_facts_from_contents(contents, llm_config, agent_name, config, pool, operation_id, schema)
|
||||
|
||||
# Check if we're resuming an existing batch (crash recovery)
|
||||
batch_id = None
|
||||
@@ -2214,9 +2195,7 @@ async def extract_facts_from_contents(
|
||||
fact_extraction_tasks.append(task)
|
||||
|
||||
# Step 2: Wait for all fact extractions to complete.
|
||||
# return_exceptions=True so a failing item doesn't cancel its still-running
|
||||
# siblings (which would leave orphaned LLM calls / partial work); we await
|
||||
# them all, then propagate.
|
||||
# Use return_exceptions=True so one content item failure doesn't discard the rest.
|
||||
all_fact_results = await asyncio.gather(*fact_extraction_tasks, return_exceptions=True)
|
||||
|
||||
# Step 3: Flatten and convert to typed objects
|
||||
@@ -2227,18 +2206,14 @@ async def extract_facts_from_contents(
|
||||
global_chunk_idx = 0
|
||||
global_fact_idx = 0
|
||||
|
||||
# Never silently drop a document's memory. Any extraction failure (provider
|
||||
# rate-limit / timeout / 5xx, malformed response, token-limit, etc.)
|
||||
# propagates so the streaming producer surfaces it and the worker's
|
||||
# RetryTaskAt machinery retries the task — and ultimately fails it *loudly*
|
||||
# if the problem persists. Swallowing the error and substituting an empty
|
||||
# result here used to commit the document with 0 facts and mark the
|
||||
# operation `completed`, losing the memory with no signal. See issue #1833.
|
||||
# Filter out failed content items
|
||||
valid_results = []
|
||||
for content, result in zip(contents, all_fact_results):
|
||||
if isinstance(result, Exception):
|
||||
raise result
|
||||
valid_results.append((content, result))
|
||||
logger.warning(f"Content extraction failed (skipping): {type(result).__name__}: {result}")
|
||||
valid_results.append((content, ([], [], TokenUsage())))
|
||||
else:
|
||||
valid_results.append((content, result))
|
||||
|
||||
for content_index, (content, (facts_from_llm, chunks_from_llm, content_usage)) in enumerate(valid_results):
|
||||
total_usage = total_usage + content_usage
|
||||
|
||||
@@ -321,15 +321,6 @@ async def handle_document_tracking(
|
||||
f"[RETAIN] Document {document_id} re-ingested: invalidated "
|
||||
f"{invalidated} observation(s) derived from {len(existing_unit_ids)} outgoing memory_units"
|
||||
)
|
||||
# Capture link-recompute victims BEFORE the cascade. Same staleness
|
||||
# applies on upsert as on explicit delete: surviving units in OTHER
|
||||
# documents that linked to these doomed units are about to lose
|
||||
# those links. ``ops`` may be None for older callers that haven't
|
||||
# been wired up — skip enqueue in that case rather than crash.
|
||||
if ops is not None:
|
||||
from ..graph_maintenance import enqueue_relink_victims
|
||||
|
||||
await enqueue_relink_victims(conn, bank_id, [str(uid) for uid in existing_unit_ids], ops=ops)
|
||||
# Explicitly delete memory_units by document_id BEFORE deleting the
|
||||
# document row. The CASCADE from documents→chunks→memory_units only
|
||||
# catches units that have a non-NULL chunk_id FK. Units with chunk_id=NULL
|
||||
|
||||
@@ -5,9 +5,10 @@ Link creation utilities for temporal, semantic, and entity links.
|
||||
import logging
|
||||
import time
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from uuid import UUID
|
||||
|
||||
from ..._vector_index import ann_search_tuning_settings, configured_vector_extension
|
||||
from ..memory_engine import fq_table
|
||||
from .types import EntityLink
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -365,6 +366,136 @@ async def resolve_entities_only(
|
||||
return resolved_entity_ids, entity_to_unit, unit_to_entity_ids
|
||||
|
||||
|
||||
async def build_entity_links_from_resolved(
|
||||
entity_resolver,
|
||||
conn,
|
||||
bank_id: str,
|
||||
unit_ids: list[str],
|
||||
resolved_entity_ids: list[str],
|
||||
entity_to_unit: list[tuple],
|
||||
unit_to_entity_ids: dict[str, list[str]],
|
||||
log_buffer: list[str] = None,
|
||||
skip_unit_entities_insert: bool = False,
|
||||
ops=None,
|
||||
) -> list["EntityLink"]:
|
||||
"""
|
||||
Build entity links between units that share entities.
|
||||
|
||||
Queries unit_entities to find which existing units share entities with the
|
||||
new units, then generates EntityLink objects for UI graph visualization.
|
||||
|
||||
Args:
|
||||
entity_resolver: EntityResolver instance
|
||||
conn: Database connection
|
||||
bank_id: Bank identifier
|
||||
unit_ids: Actual unit IDs (must already be inserted in the DB)
|
||||
resolved_entity_ids: Entity IDs from resolve_entities_only
|
||||
entity_to_unit: Mapping from resolve_entities_only
|
||||
unit_to_entity_ids: Mapping from resolve_entities_only
|
||||
log_buffer: Optional logging buffer
|
||||
skip_unit_entities_insert: If True, skip unit_entities INSERT (already done in Phase 2)
|
||||
|
||||
Returns:
|
||||
List of EntityLink objects for batch insertion
|
||||
"""
|
||||
if not resolved_entity_ids:
|
||||
return []
|
||||
|
||||
if not skip_unit_entities_insert:
|
||||
# Insert unit-entity links (used in fallback path where Phase 2 didn't do this)
|
||||
substep_start = time.time()
|
||||
unit_entity_pairs = []
|
||||
for idx, (unit_id, _local_idx, fact_date) in enumerate(entity_to_unit):
|
||||
# Propagate the unit's fact_date so entity_cooccurrences.last_cooccurred
|
||||
# reflects the event timeline, not the ingest moment.
|
||||
unit_entity_pairs.append((unit_id, resolved_entity_ids[idx], fact_date))
|
||||
|
||||
await entity_resolver.link_units_to_entities_batch(unit_entity_pairs, conn=conn)
|
||||
_log(
|
||||
log_buffer,
|
||||
f" [6.2.3] Create unit-entity links (batched): {len(unit_entity_pairs)} links in {time.time() - substep_start:.3f}s",
|
||||
level="debug",
|
||||
)
|
||||
|
||||
# Build entity links between units that share entities
|
||||
substep_start = time.time()
|
||||
all_entity_ids = set()
|
||||
for entity_ids_list in unit_to_entity_ids.values():
|
||||
all_entity_ids.update(entity_ids_list)
|
||||
|
||||
_log(log_buffer, f" [6.3] Creating entity links for {len(all_entity_ids)} unique entities...", level="debug")
|
||||
|
||||
MAX_LINKS_PER_ENTITY = 10
|
||||
|
||||
entity_to_units = {}
|
||||
if all_entity_ids:
|
||||
query_start = time.time()
|
||||
import uuid
|
||||
|
||||
entity_id_list = [uuid.UUID(eid) if isinstance(eid, str) else eid for eid in all_entity_ids]
|
||||
limit_per_entity = MAX_LINKS_PER_ENTITY + len(unit_ids) # room for new units + existing cap
|
||||
|
||||
rows = await ops.fetch_entity_unit_fanout(
|
||||
conn,
|
||||
fq_table("unit_entities"),
|
||||
entity_id_list,
|
||||
limit_per_entity,
|
||||
)
|
||||
_log(
|
||||
log_buffer,
|
||||
f" [6.3.1] Query unit_entities (LATERAL): {len(rows)} rows in {time.time() - query_start:.3f}s",
|
||||
level="debug",
|
||||
)
|
||||
|
||||
group_start = time.time()
|
||||
for row in rows:
|
||||
entity_id = row["entity_id"]
|
||||
if entity_id not in entity_to_units:
|
||||
entity_to_units[entity_id] = []
|
||||
entity_to_units[entity_id].append(row["unit_id"])
|
||||
_log(log_buffer, f" [6.3.2] Group by entity_id: {time.time() - group_start:.3f}s", level="debug")
|
||||
link_gen_start = time.time()
|
||||
links: list[EntityLink] = []
|
||||
new_unit_set = set(unit_ids)
|
||||
|
||||
def to_uuid(val) -> UUID:
|
||||
return UUID(val) if isinstance(val, str) else val
|
||||
|
||||
for entity_id, units_with_entity in entity_to_units.items():
|
||||
entity_uuid = to_uuid(entity_id)
|
||||
new_units = [u for u in units_with_entity if str(u) in new_unit_set or u in new_unit_set]
|
||||
existing_units = [u for u in units_with_entity if str(u) not in new_unit_set and u not in new_unit_set]
|
||||
|
||||
new_units_to_link = new_units[-MAX_LINKS_PER_ENTITY:] if len(new_units) > MAX_LINKS_PER_ENTITY else new_units
|
||||
for i, unit_id_1 in enumerate(new_units_to_link):
|
||||
for unit_id_2 in new_units_to_link[i + 1 :]:
|
||||
links.append(
|
||||
EntityLink(from_unit_id=to_uuid(unit_id_1), to_unit_id=to_uuid(unit_id_2), entity_id=entity_uuid)
|
||||
)
|
||||
links.append(
|
||||
EntityLink(from_unit_id=to_uuid(unit_id_2), to_unit_id=to_uuid(unit_id_1), entity_id=entity_uuid)
|
||||
)
|
||||
|
||||
existing_to_link = existing_units[-MAX_LINKS_PER_ENTITY:]
|
||||
for new_unit in new_units:
|
||||
for existing_unit in existing_to_link:
|
||||
links.append(
|
||||
EntityLink(from_unit_id=to_uuid(new_unit), to_unit_id=to_uuid(existing_unit), entity_id=entity_uuid)
|
||||
)
|
||||
links.append(
|
||||
EntityLink(from_unit_id=to_uuid(existing_unit), to_unit_id=to_uuid(new_unit), entity_id=entity_uuid)
|
||||
)
|
||||
|
||||
_log(log_buffer, f" [6.3.3] Generate {len(links)} links: {time.time() - link_gen_start:.3f}s", level="debug")
|
||||
_log(
|
||||
log_buffer,
|
||||
f" [6.3] Entity link creation: {len(links)} links for {len(all_entity_ids)} unique entities in {time.time() - substep_start:.3f}s",
|
||||
level="debug",
|
||||
)
|
||||
|
||||
return links
|
||||
|
||||
|
||||
async def create_temporal_links_batch_per_fact(
|
||||
conn,
|
||||
bank_id: str,
|
||||
@@ -570,18 +701,16 @@ async def compute_semantic_links_ann(
|
||||
# `relation "_ann_seeds" does not exist` on the second statement.
|
||||
#
|
||||
# Using ON COMMIT DROP + SET LOCAL also means we don't have to remember to
|
||||
# manually drop the temp table or reset the per-backend ANN tuning GUC —
|
||||
# the transaction end handles both.
|
||||
# manually drop the temp table or reset hnsw.ef_search — the transaction
|
||||
# end handles both.
|
||||
rows: list = []
|
||||
async with conn.transaction():
|
||||
# Transaction-local ANN tuning. Each supported backend exposes its own
|
||||
# GUC (hnsw.ef_search on pgvector, vchordrq.probes on vchord); the
|
||||
# dispatcher returns the right knob for the configured backend with a
|
||||
# value tuned for top-50 semantic link creation (lower recall but much
|
||||
# lower latency than the recall-side default). SET LOCAL auto-reverts
|
||||
# at commit, so we don't pollute the pool for subsequent queries.
|
||||
for guc, value in ann_search_tuning_settings(configured_vector_extension(), kind="low_latency"):
|
||||
await conn.execute(f"SET LOCAL {guc} = {value}")
|
||||
# Transaction-local ef_search. Default 400 is tuned for recall precision
|
||||
# but at 164k units each HNSW probe takes 94ms. ef_search=60 gives 2.7ms
|
||||
# per probe (35x faster) with sufficient accuracy for top-50 semantic
|
||||
# link creation. SET LOCAL auto-reverts at commit, so we don't pollute
|
||||
# the pool for subsequent recall queries.
|
||||
await conn.execute("SET LOCAL hnsw.ef_search = 60")
|
||||
|
||||
t_setup = time_mod.time()
|
||||
await conn.execute("CREATE TEMP TABLE _ann_seeds (unit_id text, emb_text text, fact_type text) ON COMMIT DROP")
|
||||
@@ -760,6 +889,29 @@ async def create_semantic_links_batch(
|
||||
raise
|
||||
|
||||
|
||||
async def insert_entity_links_batch(conn, links: list[EntityLink], bank_id: str, chunk_size: int = 5000, ops=None):
|
||||
"""
|
||||
Bulk-insert entity links via sorted INSERT FROM unnest().
|
||||
|
||||
Args:
|
||||
conn: Database connection
|
||||
links: List of EntityLink objects
|
||||
bank_id: Bank identifier (stored directly on memory_links for fast filtering)
|
||||
chunk_size: Number of rows per INSERT chunk (default 5000)
|
||||
"""
|
||||
if not links:
|
||||
return
|
||||
|
||||
import time as time_mod
|
||||
|
||||
total_start = time_mod.time()
|
||||
tuples = [(link.from_unit_id, link.to_unit_id, link.link_type, link.weight, link.entity_id) for link in links]
|
||||
await _bulk_insert_links(conn, tuples, bank_id=bank_id, chunk_size=chunk_size, ops=ops)
|
||||
logger.debug(
|
||||
f" [9.TOTAL] Entity links batch insert ({len(tuples)} rows): {time_mod.time() - total_start:.3f}s"
|
||||
)
|
||||
|
||||
|
||||
async def create_causal_links_batch(
|
||||
conn,
|
||||
bank_id: str,
|
||||
|
||||
@@ -100,6 +100,7 @@ from .types import (
|
||||
ChunkMetadata,
|
||||
EntityResolutionResult,
|
||||
Phase1Result,
|
||||
Phase3Context,
|
||||
ProcessedFact,
|
||||
RetainContent,
|
||||
RetainContentDict,
|
||||
@@ -107,9 +108,6 @@ from .types import (
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
RetainOutboxCallback = Callable[[asyncpg.Connection], Awaitable[None]]
|
||||
RetainOutboxCallbackFactory = Callable[[list[RetainContentDict]], RetainOutboxCallback | None]
|
||||
|
||||
|
||||
def _build_retain_params(contents_dicts, document_tags=None, doc_contents=None):
|
||||
"""Build retain_params and merged_tags from content dicts."""
|
||||
@@ -258,27 +256,30 @@ async def _insert_facts_and_links(
|
||||
skip_semantic_links: bool = False,
|
||||
outbox_callback=None,
|
||||
ops=None,
|
||||
) -> list[list[str]]:
|
||||
) -> tuple[list[list[str]], Phase3Context]:
|
||||
"""
|
||||
Phase 2 of the retain pipeline: insert facts and retrieval-critical links.
|
||||
|
||||
Runs inside a single database transaction to ensure atomicity of the data
|
||||
that retrieval depends on (facts, unit_entities, temporal/semantic/causal links).
|
||||
|
||||
Entity edges for UI graph visualization are derived on demand from
|
||||
unit_entities by the /graph endpoint, so no entity rows are written to
|
||||
memory_links here.
|
||||
Entity link generation and insertion for UI visualization are NOT done here —
|
||||
only the unit_entities INSERT (FK to memory_units) stays in the transaction.
|
||||
Entity link building is deferred to Phase 3 (post-transaction, best-effort).
|
||||
"""
|
||||
set_stage("retain.phase2.insert_facts")
|
||||
unit_ids = await fact_storage.insert_facts_batch(conn, bank_id, processed_facts, ops=ops)
|
||||
step_start = time.time()
|
||||
log_buffer.append(f" Insert facts: {len(unit_ids)} units in {time.time() - step_start:.3f}s")
|
||||
|
||||
# Context for Phase 3 entity link building (after transaction commits)
|
||||
phase3_context = Phase3Context()
|
||||
|
||||
if unit_ids:
|
||||
# Entity resolution was done in Phase 1 (separate connection).
|
||||
# Remap placeholder IDs to actual unit IDs.
|
||||
step_start = time.time()
|
||||
remapped_entity_to_unit, _remapped_unit_to_entity_ids, remapped_semantic = _remap_phase1_results(
|
||||
remapped_entity_to_unit, remapped_unit_to_entity_ids, remapped_semantic = _remap_phase1_results(
|
||||
resolved_entity_ids, entity_to_unit, unit_to_entity_ids, semantic_ann_links or [], unit_ids
|
||||
)
|
||||
# Update semantic_ann_links with remapped IDs for Phase 2
|
||||
@@ -292,6 +293,13 @@ async def _insert_facts_and_links(
|
||||
]
|
||||
await entity_resolver.link_units_to_entities_batch(unit_entity_pairs, conn=conn)
|
||||
log_buffer.append(f" Insert unit_entities: {len(unit_entity_pairs)} pairs in {time.time() - step_start:.3f}s")
|
||||
# Save context for Phase 3 entity link building (after commit)
|
||||
phase3_context = Phase3Context(
|
||||
unit_ids=unit_ids,
|
||||
resolved_entity_ids=resolved_entity_ids,
|
||||
entity_to_unit=remapped_entity_to_unit,
|
||||
unit_to_entity_ids=remapped_unit_to_entity_ids,
|
||||
)
|
||||
|
||||
# Create temporal links
|
||||
step_start = time.time()
|
||||
@@ -332,10 +340,52 @@ async def _insert_facts_and_links(
|
||||
# an IndexError (see issue #1037).
|
||||
result_unit_ids = _map_results_to_contents(contents, processed_facts, unit_ids if unit_ids else [])
|
||||
|
||||
if outbox_callback is not None:
|
||||
if outbox_callback:
|
||||
await outbox_callback(conn)
|
||||
|
||||
return result_unit_ids
|
||||
return result_unit_ids, phase3_context
|
||||
|
||||
|
||||
async def _build_and_insert_entity_links_phase3(
|
||||
pool: Any,
|
||||
entity_resolver,
|
||||
bank_id: str,
|
||||
phase3_ctx: Phase3Context,
|
||||
log_buffer: list[str],
|
||||
) -> None:
|
||||
"""
|
||||
Phase 3 helper: build entity links from resolved data and insert them.
|
||||
|
||||
Runs on a fresh connection after the main transaction has committed.
|
||||
Entity links are for UI graph visualization only — retrieval uses
|
||||
the unit_entities self-join instead.
|
||||
"""
|
||||
set_stage("retain.phase3.entity_links")
|
||||
p3_unit_ids = phase3_ctx.unit_ids
|
||||
p3_resolved = phase3_ctx.resolved_entity_ids
|
||||
p3_entity_to_unit = phase3_ctx.entity_to_unit
|
||||
p3_unit_to_entity_ids = phase3_ctx.unit_to_entity_ids
|
||||
|
||||
if not p3_unit_ids or not p3_resolved:
|
||||
return
|
||||
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
step_start = time.time()
|
||||
entity_links = await entity_processing.build_entity_links(
|
||||
entity_resolver,
|
||||
conn,
|
||||
bank_id,
|
||||
p3_unit_ids,
|
||||
p3_resolved,
|
||||
p3_entity_to_unit,
|
||||
p3_unit_to_entity_ids,
|
||||
log_buffer,
|
||||
skip_unit_entities_insert=True, # Already inserted in Phase 2
|
||||
ops=pool.ops,
|
||||
)
|
||||
if entity_links:
|
||||
await entity_processing.insert_entity_links_batch(conn, entity_links, bank_id, ops=pool.ops)
|
||||
log_buffer.append(f" Entity links (viz): {len(entity_links)} links in {time.time() - step_start:.3f}s")
|
||||
|
||||
|
||||
async def _extract_and_embed(
|
||||
@@ -399,11 +449,8 @@ async def retain_batch(
|
||||
document_tags: list[str] | None = None,
|
||||
operation_id: str | None = None,
|
||||
schema: str | None = None,
|
||||
outbox_callback: RetainOutboxCallback | None = None,
|
||||
outbox_callback_factory: RetainOutboxCallbackFactory | None = None,
|
||||
outbox_callback: Callable[["asyncpg.Connection"], Awaitable[None]] | None = None,
|
||||
db_semaphore: "asyncio.Semaphore | None" = None,
|
||||
document_body_override: str | None = None,
|
||||
chunk_index_offset: int = 0,
|
||||
) -> tuple[list[list[str]], TokenUsage, int | None]:
|
||||
"""
|
||||
Process a batch of content through the retain pipeline.
|
||||
@@ -412,14 +459,6 @@ async def retain_batch(
|
||||
only re-processes chunks whose content has changed. Unchanged chunks keep
|
||||
their existing facts, entities, and links.
|
||||
|
||||
``chunk_index_offset`` shifts the chunk_index (and therefore the derived
|
||||
``chunk_id = {bank}_{doc}_{index}``) of every chunk this call stores. The
|
||||
in-process splitter slices an oversized single item into several
|
||||
sub-batches that all share one document_id and run sequentially; without
|
||||
a per-document offset each sub-batch would restart chunk_index at 0, so
|
||||
their chunk_ids collide and later sub-batches overwrite earlier chunks —
|
||||
leaving only one sub-batch's worth of chunks/memories behind (issue #1888).
|
||||
|
||||
Returns a three-tuple of:
|
||||
* per-content-item unit ID lists
|
||||
* aggregate LLM token usage
|
||||
@@ -468,10 +507,6 @@ async def retain_batch(
|
||||
total_usage = TokenUsage()
|
||||
total_processed_tokens: int | None = 0
|
||||
for doc_key, (group_dicts, group_contents) in groups.items():
|
||||
group_outbox_callback = (
|
||||
outbox_callback_factory(group_dicts) if outbox_callback_factory is not None else outbox_callback
|
||||
)
|
||||
|
||||
group_ids, group_usage, group_processed = await retain_batch(
|
||||
pool=pool,
|
||||
embeddings_model=embeddings_model,
|
||||
@@ -487,11 +522,8 @@ async def retain_batch(
|
||||
document_tags=document_tags,
|
||||
operation_id=operation_id,
|
||||
schema=schema,
|
||||
outbox_callback=group_outbox_callback,
|
||||
outbox_callback_factory=outbox_callback_factory,
|
||||
outbox_callback=outbox_callback,
|
||||
db_semaphore=db_semaphore,
|
||||
document_body_override=document_body_override,
|
||||
chunk_index_offset=chunk_index_offset,
|
||||
)
|
||||
for group_idx, orig_idx in enumerate(original_indices[doc_key]):
|
||||
if group_idx < len(group_ids):
|
||||
@@ -633,7 +665,6 @@ async def retain_batch(
|
||||
schema,
|
||||
outbox_callback,
|
||||
db_semaphore,
|
||||
document_body_override=document_body_override,
|
||||
)
|
||||
if delta_result is not None:
|
||||
return delta_result
|
||||
@@ -690,8 +721,6 @@ async def retain_batch(
|
||||
schema=schema,
|
||||
outbox_callback=outbox_callback,
|
||||
db_semaphore=db_semaphore,
|
||||
document_body_override=document_body_override,
|
||||
chunk_index_offset=chunk_index_offset,
|
||||
)
|
||||
|
||||
|
||||
@@ -829,8 +858,6 @@ async def _streaming_retain_batch(
|
||||
schema: str | None = None,
|
||||
outbox_callback: Callable[["asyncpg.Connection"], Awaitable[None]] | None = None,
|
||||
db_semaphore: "asyncio.Semaphore | None" = None,
|
||||
document_body_override: str | None = None,
|
||||
chunk_index_offset: int = 0,
|
||||
) -> tuple[list[list[str]], TokenUsage]:
|
||||
"""
|
||||
Process a large document in streaming mini-batches to bound memory usage.
|
||||
@@ -864,15 +891,7 @@ async def _streaming_retain_batch(
|
||||
# document exists with a matching content_hash and has committed chunks,
|
||||
# the producer can skip already-extracted chunks to avoid duplicate work.
|
||||
existing_chunk_hashes: set[str] = set()
|
||||
# When the caller is processing a sub-batch sliced out of an oversized
|
||||
# item (see _split_contents_into_sub_batches), document_body_override
|
||||
# carries the full original document body. Use it for the doc-row write
|
||||
# so documents.original_text stores the complete payload, not just this
|
||||
# slice (issue #1838).
|
||||
if document_body_override is not None:
|
||||
combined_content = document_body_override
|
||||
else:
|
||||
combined_content = "\n".join([c.get("content", "") for c in contents_dicts])
|
||||
combined_content = "\n".join([c.get("content", "") for c in contents_dicts])
|
||||
# Memory: contents_dicts content strings are now captured in combined_content.
|
||||
# Clear them from the dicts to release the per-item copies (can be multi-MB each).
|
||||
for d in contents_dicts:
|
||||
@@ -1052,20 +1071,14 @@ async def _streaming_retain_batch(
|
||||
# Adjust chunk indices to use the original global position (global_idx)
|
||||
# so that chunk_id = {bank}_{doc}_{chunk_index} is deterministic regardless
|
||||
# of task completion order. content_index is batch-relative for result grouping.
|
||||
#
|
||||
# chunk_index_offset continues the document's chunk_index sequence
|
||||
# when this call is one of several sequential sub-batches sliced
|
||||
# from a single oversized item sharing one document_id — without it
|
||||
# each sub-batch restarts at 0 and their chunk_ids collide (#1888).
|
||||
doc_chunk_index = global_idx + chunk_index_offset
|
||||
for fact in extracted:
|
||||
fact.content_index = content_idx_in_batch
|
||||
if fact.chunk_index is not None:
|
||||
fact.chunk_index = doc_chunk_index
|
||||
fact.chunk_index = global_idx
|
||||
for pf in processed:
|
||||
pf.content_index = content_idx_in_batch
|
||||
for cm in chunk_meta:
|
||||
cm.chunk_index = doc_chunk_index
|
||||
cm.chunk_index = global_idx
|
||||
|
||||
batch_contents.append(content)
|
||||
batch_extracted.extend(extracted)
|
||||
@@ -1178,6 +1191,7 @@ async def _streaming_retain_batch(
|
||||
|
||||
p2_start = time.time()
|
||||
batch_result_ids = None
|
||||
phase3_ctx = None
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
async with conn.transaction():
|
||||
# --- Document ownership gate ---
|
||||
@@ -1266,7 +1280,7 @@ async def _streaming_retain_batch(
|
||||
|
||||
# Insert facts and links — skip semantic links entirely in streaming
|
||||
# mode; they are created in a single final ANN pass after all batches.
|
||||
batch_result_ids = await _insert_facts_and_links(
|
||||
batch_result_ids, phase3_ctx = await _insert_facts_and_links(
|
||||
conn,
|
||||
entity_resolver,
|
||||
bank_id,
|
||||
@@ -1286,13 +1300,15 @@ async def _streaming_retain_batch(
|
||||
|
||||
logger.info(f"[streaming] Phase 2 (write txn): {time.time() - p2_start:.3f}s")
|
||||
|
||||
# Best-effort: flush entity_cooccurrences and other deferred stats.
|
||||
try:
|
||||
await entity_resolver.flush_pending_stats()
|
||||
except Exception:
|
||||
logger.warning(
|
||||
f"Entity stats flush (consumer batch {consumer_batch_idx + 1}) failed", exc_info=True
|
||||
)
|
||||
# Best-effort: entity viz + stats (fast, not semantic ANN)
|
||||
if phase3_ctx is not None:
|
||||
try:
|
||||
await entity_resolver.flush_pending_stats()
|
||||
await _build_and_insert_entity_links_phase3(
|
||||
pool, entity_resolver, bank_id, phase3_ctx, log_buffer
|
||||
)
|
||||
except Exception:
|
||||
logger.warning(f"Phase 3 stats (consumer batch {consumer_batch_idx + 1}) failed", exc_info=True)
|
||||
|
||||
logger.info(
|
||||
f"[streaming] Consumer batch {consumer_batch_idx + 1} total "
|
||||
@@ -1521,8 +1537,6 @@ async def _try_delta_retain(
|
||||
schema,
|
||||
outbox_callback,
|
||||
db_semaphore: "asyncio.Semaphore | None" = None,
|
||||
*,
|
||||
document_body_override: str | None = None,
|
||||
) -> tuple[list[list[str]], TokenUsage, int | None] | None:
|
||||
"""
|
||||
Attempt delta retain for a document upsert. Returns result tuple if delta
|
||||
@@ -1608,7 +1622,6 @@ async def _try_delta_retain(
|
||||
log_buffer,
|
||||
start_time,
|
||||
outbox_callback,
|
||||
document_body_override=document_body_override,
|
||||
)
|
||||
|
||||
# Build content items for only the changed/new chunks
|
||||
@@ -1625,7 +1638,6 @@ async def _try_delta_retain(
|
||||
log_buffer,
|
||||
start_time,
|
||||
outbox_callback,
|
||||
document_body_override=document_body_override,
|
||||
)
|
||||
|
||||
# Extract facts and generate embeddings (shared pipeline)
|
||||
@@ -1685,13 +1697,7 @@ async def _try_delta_retain(
|
||||
|
||||
# Update document metadata (no delete)
|
||||
step_start = time.time()
|
||||
# When this sub-batch is one slice of an oversized item
|
||||
# split across multiple sub-batches, store the full body
|
||||
# (issue #1838) instead of just the slice.
|
||||
if document_body_override is not None:
|
||||
combined_content = document_body_override
|
||||
else:
|
||||
combined_content = "\n".join([c.get("content", "") for c in contents_dicts])
|
||||
combined_content = "\n".join([c.get("content", "") for c in contents_dicts])
|
||||
retain_params, merged_tags = _build_retain_params(contents_dicts, document_tags)
|
||||
await fact_storage.upsert_document_metadata(
|
||||
conn,
|
||||
@@ -1760,7 +1766,7 @@ async def _try_delta_retain(
|
||||
# Insert facts and retrieval-critical links.
|
||||
# Use delta_contents (the changed/new chunks) as the content list,
|
||||
# since extracted_facts have content_index relative to delta_contents.
|
||||
result_unit_ids = await _insert_facts_and_links(
|
||||
result_unit_ids, phase3_ctx = await _insert_facts_and_links(
|
||||
conn,
|
||||
entity_resolver,
|
||||
bank_id,
|
||||
@@ -1777,11 +1783,12 @@ async def _try_delta_retain(
|
||||
ops=pool.ops,
|
||||
)
|
||||
|
||||
# Flush deferred entity_cooccurrences stats (post-transaction, best-effort).
|
||||
# PHASE 3 — Best-Effort Display Data (post-transaction)
|
||||
try:
|
||||
await entity_resolver.flush_pending_stats()
|
||||
await _build_and_insert_entity_links_phase3(pool, entity_resolver, bank_id, phase3_ctx, log_buffer)
|
||||
except Exception:
|
||||
logger.warning("Entity stats flush failed — retrieval unaffected", exc_info=True)
|
||||
logger.warning("Phase 3 (best-effort display data) failed — retrieval unaffected", exc_info=True)
|
||||
|
||||
total_time = time.time() - start_time
|
||||
log_buffer.append(f"{'=' * 60}")
|
||||
@@ -1816,8 +1823,6 @@ async def _delta_metadata_only(
|
||||
log_buffer,
|
||||
start_time,
|
||||
outbox_callback,
|
||||
*,
|
||||
document_body_override: str | None = None,
|
||||
):
|
||||
"""Handle the case where no chunks changed — just update document metadata and tags."""
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
@@ -1828,12 +1833,7 @@ async def _delta_metadata_only(
|
||||
document_id,
|
||||
bank_id,
|
||||
)
|
||||
# When this sub-batch is a slice of an oversized item, write the
|
||||
# full original body (issue #1838) instead of just the slice.
|
||||
if document_body_override is not None:
|
||||
combined_content = document_body_override
|
||||
else:
|
||||
combined_content = "\n".join([c.get("content", "") for c in contents_dicts])
|
||||
combined_content = "\n".join([c.get("content", "") for c in contents_dicts])
|
||||
retain_params, merged_tags = _build_retain_params(contents_dicts, document_tags)
|
||||
await fact_storage.upsert_document_metadata(
|
||||
conn,
|
||||
@@ -1844,7 +1844,7 @@ async def _delta_metadata_only(
|
||||
merged_tags,
|
||||
)
|
||||
await fact_storage.update_memory_units_tags(conn, bank_id, document_id, merged_tags)
|
||||
if outbox_callback is not None:
|
||||
if outbox_callback:
|
||||
await outbox_callback(conn)
|
||||
|
||||
total_time = time.time() - start_time
|
||||
|
||||
@@ -224,6 +224,21 @@ class ProcessedFact:
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Phase3Context:
|
||||
"""
|
||||
Data passed from Phase 2 to Phase 3 for entity link building.
|
||||
|
||||
Contains the unit IDs and entity resolution data needed to build
|
||||
entity links for UI graph visualization after the write transaction commits.
|
||||
"""
|
||||
|
||||
unit_ids: list[str] = field(default_factory=list)
|
||||
resolved_entity_ids: list[str] = field(default_factory=list)
|
||||
entity_to_unit: list[tuple] = field(default_factory=list)
|
||||
unit_to_entity_ids: dict[str, list[str]] = field(default_factory=dict)
|
||||
|
||||
|
||||
@dataclass
|
||||
class EntityResolutionResult:
|
||||
"""
|
||||
@@ -248,6 +263,21 @@ class Phase1Result:
|
||||
semantic_ann_links: list[tuple]
|
||||
|
||||
|
||||
@dataclass
|
||||
class EntityLink:
|
||||
"""
|
||||
Link between two memory units through a shared entity.
|
||||
|
||||
Used for entity-based graph connections in the memory graph.
|
||||
"""
|
||||
|
||||
from_unit_id: UUID
|
||||
to_unit_id: UUID
|
||||
entity_id: UUID
|
||||
link_type: str = "entity"
|
||||
weight: float = 1.0
|
||||
|
||||
|
||||
@dataclass
|
||||
class RetainBatch:
|
||||
"""
|
||||
|
||||
@@ -283,12 +283,12 @@ class LinkExpansionRetriever(GraphRetriever):
|
||||
score transformations. The three CTEs share one connection slot — important
|
||||
for asyncpg which does not allow concurrent queries on the same connection.
|
||||
|
||||
Index coverage:
|
||||
entity: idx_unit_entities_entity_unit (entity_id, unit_id) — entity
|
||||
expansion traverses unit_entities, not memory_links.
|
||||
semantic: idx_memory_links_from_type_weight / _to_type_weight
|
||||
(from_unit_id|to_unit_id, link_type, weight DESC) serve both
|
||||
outgoing and incoming sides as single composite index scans.
|
||||
Index coverage (requires migration d2e3f4a5b6c7):
|
||||
entity: idx_memory_links_entity_covering (from_unit_id) INCLUDE (to_unit_id, entity_id)
|
||||
WHERE link_type = 'entity' → index-only scan, no heap reads
|
||||
semantic incoming:
|
||||
idx_memory_links_to_type_weight (to_unit_id, link_type, weight DESC)
|
||||
→ replaces costly BitmapAnd of two separate scans
|
||||
"""
|
||||
config = get_config()
|
||||
ml = fq_table("memory_links")
|
||||
|
||||
@@ -212,25 +212,16 @@ class CrossEncoderReranker:
|
||||
# Get cross-encoder scores
|
||||
scores = await self.cross_encoder.predict(pairs)
|
||||
|
||||
# Normalize scores to [0, 1] range.
|
||||
# External API rerankers (Cohere, Jina, llama.cpp/Qwen, etc.) return
|
||||
# calibrated relevance_score already in [0, 1]. These are used as-is
|
||||
# so that absolute confidence is preserved — a top candidate scoring
|
||||
# 0.007 stays low rather than being inflated to 1.0 by rank normalization.
|
||||
# Local models return logits (any real number) — sigmoid is appropriate.
|
||||
# Normalize scores using sigmoid to [0, 1] range
|
||||
# Cross-encoder returns logits which can be negative
|
||||
import math
|
||||
|
||||
import numpy as np
|
||||
|
||||
def _sigmoid(x: float) -> float:
|
||||
def sigmoid(x):
|
||||
return 1 / (1 + np.exp(-x))
|
||||
|
||||
if scores and min(scores) >= 0.0 and max(scores) <= 1.0:
|
||||
# Scores already in [0, 1] — pass through to preserve absolute
|
||||
# confidence signal from calibrated rerankers.
|
||||
normalized_scores = list(scores)
|
||||
else:
|
||||
# Scores are logits (e.g. local sentence-transformers models).
|
||||
# Sigmoid maps (-inf, +inf) to (0, 1).
|
||||
normalized_scores = [_sigmoid(score) for score in scores]
|
||||
normalized_scores = [sigmoid(score) for score in scores]
|
||||
|
||||
# Create ScoredResult objects with cross-encoder scores
|
||||
scored_results = []
|
||||
|
||||
@@ -225,7 +225,6 @@ async def retrieve_semantic_bm25_combined(
|
||||
groups_clause=groups_clause,
|
||||
arm_index=i,
|
||||
text_search_extension=text_ext,
|
||||
bm25_language=config.text_search_extension_native_language,
|
||||
extra_where=created_range_clause,
|
||||
)
|
||||
)
|
||||
@@ -466,8 +465,6 @@ async def retrieve_temporal_combined(
|
||||
best_date = ep["mentioned_at"]
|
||||
|
||||
if best_date:
|
||||
if best_date.tzinfo is None:
|
||||
best_date = best_date.replace(tzinfo=UTC)
|
||||
days_from_mid = abs((best_date - mid_date).total_seconds() / 86400)
|
||||
temporal_proximity = 1.0 - min(days_from_mid / (total_days / 2), 1.0) if total_days > 0 else 1.0
|
||||
else:
|
||||
@@ -561,8 +558,6 @@ async def retrieve_temporal_combined(
|
||||
neighbor_best_date = n["mentioned_at"]
|
||||
|
||||
if neighbor_best_date:
|
||||
if neighbor_best_date.tzinfo is None:
|
||||
neighbor_best_date = neighbor_best_date.replace(tzinfo=UTC)
|
||||
days_from_mid = abs((neighbor_best_date - mid_date).total_seconds() / 86400)
|
||||
neighbor_temporal_proximity = (
|
||||
1.0 - min(days_from_mid / (total_days / 2), 1.0) if total_days > 0 else 1.0
|
||||
|
||||
@@ -407,7 +407,6 @@ class SQLDialect(ABC):
|
||||
groups_clause: str = "",
|
||||
arm_index: int = 0,
|
||||
text_search_extension: str = "native",
|
||||
bm25_language: str = "english",
|
||||
extra_where: str = "",
|
||||
) -> str:
|
||||
"""Build a BM25/full-text search subquery arm.
|
||||
@@ -427,9 +426,7 @@ class SQLDialect(ABC):
|
||||
arm_index: Index of this arm in the UNION ALL (used by Oracle for
|
||||
unique SCORE labels).
|
||||
text_search_extension: Full-text search backend ("native", "vchord",
|
||||
"pg_textsearch", "pgroonga"). Only relevant for PostgreSQL.
|
||||
bm25_language: PostgreSQL text search dictionary used by the native
|
||||
backend (e.g. "english", "french"). Ignored by other backends.
|
||||
"pg_textsearch"). Only relevant for PostgreSQL.
|
||||
extra_where: Optional additional WHERE clause fragment (e.g. time range filter).
|
||||
"""
|
||||
...
|
||||
|
||||
@@ -270,7 +270,6 @@ class OracleDialect(SQLDialect):
|
||||
groups_clause: str = "",
|
||||
arm_index: int = 0,
|
||||
text_search_extension: str = "native",
|
||||
bm25_language: str = "english",
|
||||
extra_where: str = "",
|
||||
) -> str:
|
||||
# Oracle Text: CONTAINS() / SCORE() with the CTXSYS.CONTEXT index.
|
||||
|
||||
@@ -182,7 +182,6 @@ class PostgreSQLDialect(SQLDialect):
|
||||
groups_clause: str = "",
|
||||
arm_index: int = 0,
|
||||
text_search_extension: str = "native",
|
||||
bm25_language: str = "english",
|
||||
extra_where: str = "",
|
||||
) -> str:
|
||||
if text_search_extension == "vchord":
|
||||
@@ -194,35 +193,10 @@ class PostgreSQLDialect(SQLDialect):
|
||||
bm25_score_expr = f"-({text_param} <@> to_bm25query({text_param}, 'idx_memory_units_text_search'))"
|
||||
bm25_order_by = f"text <@> to_bm25query({text_param}, 'idx_memory_units_text_search') ASC"
|
||||
bm25_where_filter = ""
|
||||
elif text_search_extension == "pgroonga":
|
||||
# &@~ accepts pgroonga's query syntax (raw query text). pgroonga_score
|
||||
# returns a non-negative relevance score (higher = better).
|
||||
bm25_score_expr = "pgroonga_score(tableoid, ctid)"
|
||||
bm25_order_by = f"{bm25_score_expr} DESC"
|
||||
bm25_where_filter = (
|
||||
f"AND (COALESCE(text, '') || ' ' || COALESCE(context, '') || ' ' || COALESCE(text_signals, '')) "
|
||||
f"&@~ {text_param}"
|
||||
)
|
||||
elif text_search_extension == "pg_search":
|
||||
# ParadeDB pg_search: BM25 index over (id, text, context, text_signals)
|
||||
# with key_field='id'. The @@@ operator on the key_field requires a
|
||||
# field-qualified query (`text:foo`); to keep the bind-parameter form,
|
||||
# we fan the query out across all indexed text fields with paradedb.boolean.
|
||||
bm25_score_expr = "paradedb.score(id)"
|
||||
bm25_order_by = "paradedb.score(id) DESC"
|
||||
bm25_where_filter = (
|
||||
f"AND id @@@ paradedb.boolean(should => ARRAY["
|
||||
f"paradedb.match('text', {text_param}), "
|
||||
f"paradedb.match('context', {text_param}), "
|
||||
f"paradedb.match('text_signals', {text_param})"
|
||||
f"])"
|
||||
)
|
||||
else: # native tsvector
|
||||
# bm25_language is validated as a PG identifier in HindsightConfig.validate(),
|
||||
# so embedding it as a SQL literal here is safe.
|
||||
bm25_score_expr = f"ts_rank_cd(search_vector, to_tsquery('{bm25_language}', {text_param}))"
|
||||
bm25_score_expr = f"ts_rank_cd(search_vector, to_tsquery('english', {text_param}))"
|
||||
bm25_order_by = f"{bm25_score_expr} DESC"
|
||||
bm25_where_filter = f"AND search_vector @@ to_tsquery('{bm25_language}', {text_param})"
|
||||
bm25_where_filter = f"AND search_vector @@ to_tsquery('english', {text_param})"
|
||||
|
||||
return (
|
||||
f"(SELECT {cols},"
|
||||
@@ -247,7 +221,7 @@ class PostgreSQLDialect(SQLDialect):
|
||||
*,
|
||||
text_search_extension: str = "native",
|
||||
) -> str:
|
||||
if text_search_extension in ("vchord", "pg_textsearch", "pgroonga", "pg_search"):
|
||||
if text_search_extension in ("vchord", "pg_textsearch"):
|
||||
return query_text
|
||||
# native tsvector: join tokens with OR operator
|
||||
return " | ".join(tokens)
|
||||
|
||||
@@ -16,7 +16,7 @@ class S3FileStorage(FileStorage):
|
||||
S3-compatible object storage backend.
|
||||
|
||||
Uses obstore (Rust-backed) for high-throughput async access to
|
||||
Amazon S3, MinIO, Cloudflare R2, Tigris, and other S3-compliant APIs.
|
||||
Amazon S3, MinIO, Cloudflare R2, and other S3-compliant APIs.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
|
||||
@@ -1,46 +0,0 @@
|
||||
"""Shared tiktoken encoding used for token counting and chunking.
|
||||
|
||||
Hindsight uses tiktoken purely to *count* and *chunk* arbitrary user content — never
|
||||
to feed a model that relies on tiktoken's special-token vocabulary. With tiktoken's
|
||||
default ``disallowed_special="all"``, any content that merely *mentions* a special-token
|
||||
literal (e.g. ``<|endoftext|>``) makes ``encode()`` raise, which surfaces as an HTTP 500
|
||||
on retain/recall (see issue #1883).
|
||||
|
||||
``_SafeEncoding`` disables that check so such literals are counted as ordinary text. Token
|
||||
counts are unaffected; this only stops the encoder from rejecting valid input. Every token
|
||||
call site in the engine routes through ``get_token_encoding()``, so the fix is global.
|
||||
"""
|
||||
|
||||
from functools import lru_cache
|
||||
|
||||
import tiktoken
|
||||
|
||||
|
||||
class _SafeEncoding:
|
||||
"""Wraps a tiktoken ``Encoding`` so ``encode()`` never raises on special-token literals."""
|
||||
|
||||
def __init__(self, encoding: tiktoken.Encoding) -> None:
|
||||
self._encoding = encoding
|
||||
|
||||
def encode(self, text: str, **kwargs) -> list[int]:
|
||||
# Count special-token literals as ordinary text instead of rejecting them.
|
||||
kwargs.setdefault("disallowed_special", ())
|
||||
return self._encoding.encode(text, **kwargs)
|
||||
|
||||
def decode(self, tokens: list[int]) -> str:
|
||||
return self._encoding.decode(tokens)
|
||||
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def get_token_encoding() -> _SafeEncoding:
|
||||
"""Cached cl100k_base encoding (GPT-4/3.5) wrapped to tolerate special-token literals.
|
||||
|
||||
tiktoken downloads the encoding on first lookup; keeping it lazy means importing
|
||||
``hindsight_api`` does not require network access.
|
||||
"""
|
||||
return _SafeEncoding(tiktoken.get_encoding("cl100k_base"))
|
||||
|
||||
|
||||
def count_tokens(text: str) -> int:
|
||||
"""Count cl100k_base tokens in ``text`` (tolerant of special-token literals)."""
|
||||
return len(get_token_encoding().encode(text))
|
||||
@@ -40,7 +40,6 @@ from hindsight_api.extensions.operation_validator import (
|
||||
# Core operations
|
||||
OperationValidationError,
|
||||
OperationValidatorExtension,
|
||||
PrecheckContext,
|
||||
RecallContext,
|
||||
RecallResult,
|
||||
ReflectContext,
|
||||
@@ -73,7 +72,6 @@ __all__ = [
|
||||
"DeferOperation",
|
||||
"OperationValidationError",
|
||||
"OperationValidatorExtension",
|
||||
"PrecheckContext",
|
||||
"RecallContext",
|
||||
"RecallResult",
|
||||
"ReflectContext",
|
||||
|
||||
@@ -146,11 +146,7 @@ class DefaultExtensionContext(ExtensionContext):
|
||||
|
||||
# Ensure text search columns/indexes match the configured extension
|
||||
await asyncio.to_thread(
|
||||
ensure_text_search_extension,
|
||||
db_url,
|
||||
text_search_extension=config.text_search_extension,
|
||||
pg_search_tokenizer=config.text_search_extension_pg_search_tokenizer,
|
||||
schema=schema,
|
||||
ensure_text_search_extension, db_url, text_search_extension=config.text_search_extension, schema=schema
|
||||
)
|
||||
|
||||
def get_memory_engine(self) -> "MemoryEngineInterface":
|
||||
|
||||
@@ -82,33 +82,6 @@ class ValidationResult:
|
||||
# =============================================================================
|
||||
|
||||
|
||||
@dataclass
|
||||
class PrecheckContext:
|
||||
"""Context for a pre-body-parse precheck on an operation.
|
||||
|
||||
Unlike :class:`RetainContext` / :class:`RecallContext` / etc., this
|
||||
context is constructed *before* the request body is deserialised. It
|
||||
therefore intentionally carries only the cheap, already-resolved
|
||||
pieces of request state:
|
||||
|
||||
- ``operation``: a short string identifying the route, e.g. ``"retain"``,
|
||||
``"recall"``, ``"reflect"``, ``"files_retain"``, ``"mental_model_create"``,
|
||||
``"mental_model_refresh"``.
|
||||
- ``bank_id``: parsed from the URL path.
|
||||
- ``request_context``: the authenticated :class:`RequestContext` (tenant
|
||||
already resolved by the tenant extension).
|
||||
|
||||
Implementations should keep precheck cheap and side-effect-free. The
|
||||
full per-request validators (``validate_retain`` / ``validate_recall``
|
||||
/ ``validate_reflect``) still run after the body is parsed and remain
|
||||
the source of truth for the precise per-call cost / quota arithmetic.
|
||||
"""
|
||||
|
||||
operation: str
|
||||
bank_id: str
|
||||
request_context: "RequestContext"
|
||||
|
||||
|
||||
@dataclass
|
||||
class RetainContext:
|
||||
"""Context for a retain operation validation (pre-operation).
|
||||
@@ -434,42 +407,6 @@ class OperationValidatorExtension(Extension, ABC):
|
||||
- consolidate (mental models consolidation)
|
||||
"""
|
||||
|
||||
# =========================================================================
|
||||
# Pre-body-parse hook (optional - default no-op)
|
||||
# =========================================================================
|
||||
|
||||
async def precheck(self, ctx: PrecheckContext) -> ValidationResult:
|
||||
"""
|
||||
Cheap pre-body-parse check, called before the request body is read.
|
||||
|
||||
FastAPI resolves ``Depends`` callables before deserialising the route
|
||||
body; routes that wire ``precheck`` as a dependency therefore short
|
||||
-circuit here without ever materialising the JSON payload in memory.
|
||||
That makes this the right hook for "should this caller be allowed to
|
||||
spend resources on this request at all" decisions — e.g. a balance
|
||||
is exhausted, a key is revoked, or a tenant is rate-limited.
|
||||
|
||||
Implementations should:
|
||||
- Be cheap: prefer cached lookups, avoid heavy DB queries.
|
||||
- Use only data on ``ctx`` (operation name + bank_id + request_context);
|
||||
the body is not yet available.
|
||||
- Be conservative on errors: prefer ``ValidationResult.accept()`` so
|
||||
a transient lookup failure doesn't turn into a request rejection.
|
||||
The post-body ``validate_*`` hooks still run and remain the source
|
||||
of truth for the precise per-call cost check.
|
||||
|
||||
Default implementation accepts everything. Override to opt in.
|
||||
|
||||
Args:
|
||||
ctx: Pre-body context with operation name, bank_id, and
|
||||
request_context (tenant already resolved).
|
||||
|
||||
Returns:
|
||||
ValidationResult indicating whether the request may proceed to
|
||||
body parsing and the post-parse validators.
|
||||
"""
|
||||
return ValidationResult.accept()
|
||||
|
||||
# =========================================================================
|
||||
# Pre-operation validation hooks (abstract - must be implemented)
|
||||
# =========================================================================
|
||||
|
||||
@@ -24,15 +24,7 @@ import uvicorn
|
||||
from . import MemoryEngine, __version__
|
||||
from .api import create_app
|
||||
from .banner import print_banner
|
||||
from .config import (
|
||||
DEFAULT_ACCESS_LOG,
|
||||
DEFAULT_WORKERS,
|
||||
ENV_ACCESS_LOG,
|
||||
ENV_HOST,
|
||||
ENV_WORKERS,
|
||||
HindsightConfig,
|
||||
_get_raw_config,
|
||||
)
|
||||
from .config import DEFAULT_WORKERS, ENV_HOST, ENV_WORKERS, HindsightConfig, _get_raw_config
|
||||
from .daemon import (
|
||||
DEFAULT_DAEMON_PORT,
|
||||
DEFAULT_IDLE_TIMEOUT,
|
||||
@@ -73,42 +65,29 @@ def _signal_handler(signum, frame):
|
||||
sys.exit(0)
|
||||
|
||||
|
||||
@dataclasses.dataclass(frozen=True)
|
||||
class ResolvedDaemonHostPort:
|
||||
host: str
|
||||
port: int
|
||||
|
||||
|
||||
def resolve_daemon_host_port(
|
||||
*,
|
||||
args_host: str,
|
||||
args_port: int,
|
||||
explicit_host: bool,
|
||||
explicit_port: bool,
|
||||
) -> ResolvedDaemonHostPort:
|
||||
def resolve_daemon_host_port(*, args_host: str, args_port: int, config_host: str, config_port: int) -> tuple[str, int]:
|
||||
"""Resolve host/port for daemon mode.
|
||||
|
||||
Defaults to 127.0.0.1 for security, but honors explicit user overrides
|
||||
via --host flag or HINDSIGHT_API_HOST env var. Uses DEFAULT_DAEMON_PORT
|
||||
unless the user specified a custom port.
|
||||
"""
|
||||
port = args_port if explicit_port else DEFAULT_DAEMON_PORT
|
||||
port = args_port if args_port != config_port else DEFAULT_DAEMON_PORT
|
||||
# Only force localhost if the user didn't explicitly set a host
|
||||
if explicit_host or os.environ.get(ENV_HOST):
|
||||
if args_host != config_host or os.environ.get(ENV_HOST):
|
||||
host = args_host
|
||||
else:
|
||||
host = "127.0.0.1"
|
||||
return ResolvedDaemonHostPort(host=host, port=port)
|
||||
return host, port
|
||||
|
||||
|
||||
@dataclasses.dataclass(frozen=True)
|
||||
class ParsedCliArgs:
|
||||
args: argparse.Namespace
|
||||
explicit_host: bool
|
||||
explicit_port: bool
|
||||
def main():
|
||||
"""Main entry point for the CLI."""
|
||||
global _memory
|
||||
|
||||
# Load configuration from environment (for CLI args defaults)
|
||||
config = _get_raw_config()
|
||||
|
||||
def _parse_cli_args(argv: list[str], config: HindsightConfig) -> ParsedCliArgs:
|
||||
parser = argparse.ArgumentParser(
|
||||
prog="hindsight-api",
|
||||
description="Hindsight API Server",
|
||||
@@ -116,14 +95,12 @@ def _parse_cli_args(argv: list[str], config: HindsightConfig) -> ParsedCliArgs:
|
||||
|
||||
# Server options
|
||||
parser.add_argument(
|
||||
"--host",
|
||||
default=argparse.SUPPRESS,
|
||||
help=f"Host to bind to (default: {config.host}, env: HINDSIGHT_API_HOST)",
|
||||
"--host", default=config.host, help=f"Host to bind to (default: {config.host}, env: HINDSIGHT_API_HOST)"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--port",
|
||||
type=int,
|
||||
default=argparse.SUPPRESS,
|
||||
default=config.port,
|
||||
help=f"Port to bind to (default: {config.port}, env: HINDSIGHT_API_PORT)",
|
||||
)
|
||||
parser.add_argument(
|
||||
@@ -143,18 +120,9 @@ def _parse_cli_args(argv: list[str], config: HindsightConfig) -> ParsedCliArgs:
|
||||
)
|
||||
|
||||
# Access log options
|
||||
parser.add_argument(
|
||||
"--access-log",
|
||||
action="store_true",
|
||||
default=os.getenv(ENV_ACCESS_LOG, "").lower() in ("1", "true", "yes", "on") or DEFAULT_ACCESS_LOG,
|
||||
help=f"Enable access log (env: {ENV_ACCESS_LOG}, default: {DEFAULT_ACCESS_LOG})",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--no-access-log",
|
||||
dest="access_log",
|
||||
action="store_false",
|
||||
help="Disable access log (overrides env and default)",
|
||||
)
|
||||
parser.add_argument("--access-log", action="store_true", help="Enable access log")
|
||||
parser.add_argument("--no-access-log", dest="access_log", action="store_false", help="Disable access log (default)")
|
||||
parser.set_defaults(access_log=False)
|
||||
|
||||
# Proxy options
|
||||
parser.add_argument(
|
||||
@@ -181,27 +149,7 @@ def _parse_cli_args(argv: list[str], config: HindsightConfig) -> ParsedCliArgs:
|
||||
help=f"Idle timeout in seconds before auto-exit in daemon mode (default: {DEFAULT_IDLE_TIMEOUT})",
|
||||
)
|
||||
|
||||
args = parser.parse_args(argv)
|
||||
|
||||
explicit_host = hasattr(args, "host")
|
||||
explicit_port = hasattr(args, "port")
|
||||
if not explicit_host:
|
||||
args.host = config.host
|
||||
if not explicit_port:
|
||||
args.port = config.port
|
||||
|
||||
return ParsedCliArgs(args=args, explicit_host=explicit_host, explicit_port=explicit_port)
|
||||
|
||||
|
||||
def main():
|
||||
"""Main entry point for the CLI."""
|
||||
global _memory
|
||||
|
||||
# Load configuration from environment (for CLI args defaults)
|
||||
config = _get_raw_config()
|
||||
|
||||
parsed_cli_args = _parse_cli_args(sys.argv[1:], config)
|
||||
args = parsed_cli_args.args
|
||||
args = parser.parse_args()
|
||||
|
||||
# Daemon mode handling.
|
||||
# is_daemon_child is True when we are the re-exec'd child spawned by
|
||||
@@ -212,14 +160,12 @@ def main():
|
||||
is_daemon = args.daemon or is_daemon_child
|
||||
|
||||
if is_daemon:
|
||||
resolved_daemon_host_port = resolve_daemon_host_port(
|
||||
args.host, args.port = resolve_daemon_host_port(
|
||||
args_host=args.host,
|
||||
args_port=args.port,
|
||||
explicit_host=parsed_cli_args.explicit_host,
|
||||
explicit_port=parsed_cli_args.explicit_port,
|
||||
config_host=config.host,
|
||||
config_port=config.port,
|
||||
)
|
||||
args.host = resolved_daemon_host_port.host
|
||||
args.port = resolved_daemon_host_port.port
|
||||
|
||||
# Detach into background (parent re-execs and exits; child redirects
|
||||
# stdio to log file). No lockfile needed — port binding prevents
|
||||
@@ -293,6 +239,8 @@ def main():
|
||||
# When using workers or reload, we must use import string so each worker can import the app
|
||||
use_import_string = args.workers > 1 or args.reload
|
||||
# Check for uvloop/winloop availability
|
||||
import sys
|
||||
|
||||
loop_impl = "asyncio"
|
||||
if sys.platform == "win32":
|
||||
try:
|
||||
|
||||
@@ -44,7 +44,6 @@ _ALL_TOOLS: frozenset[str] = frozenset(
|
||||
"update_mental_model",
|
||||
"delete_mental_model",
|
||||
"refresh_mental_model",
|
||||
"clear_mental_model",
|
||||
"list_directives",
|
||||
"create_directive",
|
||||
"delete_directive",
|
||||
@@ -222,7 +221,6 @@ def register_mcp_tools(
|
||||
"update_mental_model",
|
||||
"delete_mental_model",
|
||||
"refresh_mental_model",
|
||||
"clear_mental_model",
|
||||
"list_directives",
|
||||
"create_directive",
|
||||
"delete_directive",
|
||||
@@ -279,9 +277,6 @@ def register_mcp_tools(
|
||||
if "refresh_mental_model" in tools_to_register:
|
||||
_register_refresh_mental_model(mcp, memory, config)
|
||||
|
||||
if "clear_mental_model" in tools_to_register:
|
||||
_register_clear_mental_model(mcp, memory, config)
|
||||
|
||||
# Directive tools
|
||||
if "list_directives" in tools_to_register:
|
||||
_register_list_directives(mcp, memory, config)
|
||||
@@ -443,7 +438,6 @@ _AUDITABLE_MCP_TOOLS: frozenset[str] = frozenset(
|
||||
"update_mental_model",
|
||||
"delete_mental_model",
|
||||
"refresh_mental_model",
|
||||
"clear_mental_model",
|
||||
"create_directive",
|
||||
"delete_directive",
|
||||
"delete_document",
|
||||
@@ -799,8 +793,7 @@ def _register_recall(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig)
|
||||
{"tags": [...], "match": "any_strict"} or compound {"and": [...]}, {"or": [...]}, {"not": {...}}.
|
||||
Example: [{"not": {"tags": ["closeout"], "match": "any_strict"}}] excludes memories tagged closeout.
|
||||
Mutually exclusive with tags.
|
||||
query_timestamp: Temporal context for the query (ISO format, e.g., '2024-01-15T10:30:00Z').
|
||||
Anchors relative temporal expressions and recency scoring.
|
||||
query_timestamp: Temporal context for the query (ISO format, e.g., '2024-01-15T10:30:00Z'). Helps retrieve time-relevant memories.
|
||||
bank_id: Optional bank to search in (defaults to session bank). Use for cross-bank operations.
|
||||
"""
|
||||
try:
|
||||
@@ -870,8 +863,7 @@ def _register_recall(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig)
|
||||
{"tags": [...], "match": "any_strict"} or compound {"and": [...]}, {"or": [...]}, {"not": {...}}.
|
||||
Example: [{"not": {"tags": ["closeout"], "match": "any_strict"}}] excludes memories tagged closeout.
|
||||
Mutually exclusive with tags.
|
||||
query_timestamp: Temporal context for the query (ISO format, e.g., '2024-01-15T10:30:00Z').
|
||||
Anchors relative temporal expressions and recency scoring.
|
||||
query_timestamp: Temporal context for the query (ISO format, e.g., '2024-01-15T10:30:00Z'). Helps retrieve time-relevant memories.
|
||||
"""
|
||||
try:
|
||||
target_bank = config.bank_id_resolver()
|
||||
@@ -930,7 +922,6 @@ def _register_reflect(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig
|
||||
response_schema: dict | None = None,
|
||||
tags: list[str] | None = None,
|
||||
tags_match: str = "any",
|
||||
include_based_on: bool = False,
|
||||
bank_id: str | None = None,
|
||||
) -> str:
|
||||
"""
|
||||
@@ -960,7 +951,6 @@ def _register_reflect(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig
|
||||
response_schema: Optional JSON schema for structured output. When provided, the response includes a 'structured_output' field.
|
||||
tags: Optional tags to filter memories by (e.g., ['project:alpha'])
|
||||
tags_match: How to match tags - 'any' (match any tag) or 'all' (match all tags). Default: 'any'
|
||||
include_based_on: Include source facts used for synthesis. Defaults to false because broad reflections can exceed MCP client result limits.
|
||||
bank_id: Optional bank to reflect in (defaults to session bank). Use for cross-bank operations.
|
||||
"""
|
||||
try:
|
||||
@@ -988,8 +978,6 @@ def _register_reflect(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig
|
||||
reflect_result = await memory.reflect_async(**reflect_kwargs)
|
||||
|
||||
result_data = json.loads(reflect_result.model_dump_json(indent=2))
|
||||
if not include_based_on:
|
||||
result_data.pop("based_on", None)
|
||||
if response_schema is not None and hasattr(reflect_result, "structured_output"):
|
||||
result_data["structured_output"] = reflect_result.structured_output
|
||||
return json.dumps(result_data, indent=2)
|
||||
@@ -1011,7 +999,6 @@ def _register_reflect(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig
|
||||
response_schema: dict | None = None,
|
||||
tags: list[str] | None = None,
|
||||
tags_match: str = "any",
|
||||
include_based_on: bool = False,
|
||||
) -> dict:
|
||||
"""
|
||||
Generate thoughtful analysis by synthesizing stored memories with the bank's personality.
|
||||
@@ -1040,7 +1027,6 @@ def _register_reflect(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig
|
||||
response_schema: Optional JSON schema for structured output. When provided, the response includes a 'structured_output' field.
|
||||
tags: Optional tags to filter memories by (e.g., ['project:alpha'])
|
||||
tags_match: How to match tags - 'any' (match any tag) or 'all' (match all tags). Default: 'any'
|
||||
include_based_on: Include source facts used for synthesis. Defaults to false because broad reflections can exceed MCP client result limits.
|
||||
"""
|
||||
try:
|
||||
target_bank = config.bank_id_resolver()
|
||||
@@ -1067,8 +1053,6 @@ def _register_reflect(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig
|
||||
reflect_result = await memory.reflect_async(**reflect_kwargs)
|
||||
|
||||
result_data = reflect_result.model_dump()
|
||||
if not include_based_on:
|
||||
result_data.pop("based_on", None)
|
||||
if response_schema is not None and hasattr(reflect_result, "structured_output"):
|
||||
result_data["structured_output"] = reflect_result.structured_output
|
||||
return result_data
|
||||
@@ -1781,98 +1765,6 @@ def _register_refresh_mental_model(mcp: FastMCP, memory: MemoryEngine, config: M
|
||||
return {"error": str(e)}
|
||||
|
||||
|
||||
def _register_clear_mental_model(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig) -> None:
|
||||
"""Register the clear_mental_model tool."""
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
async def clear_mental_model(
|
||||
mental_model_id: str,
|
||||
bank_id: str | None = None,
|
||||
) -> str:
|
||||
"""
|
||||
Clear a mental model's content so the next refresh performs a full re-synthesis.
|
||||
|
||||
This is useful for delta-mode models that have accumulated drift over many
|
||||
incremental refreshes. After clearing, call refresh_mental_model to trigger
|
||||
a clean full rebuild.
|
||||
|
||||
Args:
|
||||
mental_model_id: The ID of the mental model to clear
|
||||
bank_id: Optional bank (defaults to session bank). Use for cross-bank operations.
|
||||
"""
|
||||
try:
|
||||
target_bank = bank_id or config.bank_id_resolver()
|
||||
if target_bank is None:
|
||||
return '{"error": "No bank_id configured"}'
|
||||
|
||||
result = await memory.clear_mental_model(
|
||||
bank_id=target_bank,
|
||||
mental_model_id=mental_model_id,
|
||||
request_context=_get_request_context(config),
|
||||
)
|
||||
if result is None:
|
||||
return json.dumps({"error": f"Mental model '{mental_model_id}' not found"})
|
||||
return json.dumps(
|
||||
{
|
||||
"mental_model_id": result["id"],
|
||||
"status": "cleared",
|
||||
"message": f"Mental model '{mental_model_id}' content cleared. Call refresh_mental_model to rebuild.",
|
||||
}
|
||||
)
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return json.dumps({"error": str(e)})
|
||||
except ValueError as e:
|
||||
return json.dumps({"error": str(e)})
|
||||
except Exception as e:
|
||||
logger.error(f"Error clearing mental model: {e}", exc_info=True)
|
||||
return f'{{"error": "{e}"}}'
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
async def clear_mental_model(
|
||||
mental_model_id: str,
|
||||
) -> dict:
|
||||
"""
|
||||
Clear a mental model's content so the next refresh performs a full re-synthesis.
|
||||
|
||||
This is useful for delta-mode models that have accumulated drift over many
|
||||
incremental refreshes. After clearing, call refresh_mental_model to trigger
|
||||
a clean full rebuild.
|
||||
|
||||
Args:
|
||||
mental_model_id: The ID of the mental model to clear
|
||||
"""
|
||||
try:
|
||||
target_bank = config.bank_id_resolver()
|
||||
if target_bank is None:
|
||||
return {"error": "No bank_id configured"}
|
||||
|
||||
result = await memory.clear_mental_model(
|
||||
bank_id=target_bank,
|
||||
mental_model_id=mental_model_id,
|
||||
request_context=_get_request_context(config),
|
||||
)
|
||||
if result is None:
|
||||
return {"error": f"Mental model '{mental_model_id}' not found"}
|
||||
return {
|
||||
"mental_model_id": result["id"],
|
||||
"status": "cleared",
|
||||
"message": f"Mental model '{mental_model_id}' content cleared. Call refresh_mental_model to rebuild.",
|
||||
}
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return {"error": str(e)}
|
||||
except ValueError as e:
|
||||
return {"error": str(e)}
|
||||
except Exception as e:
|
||||
logger.error(f"Error clearing mental model: {e}", exc_info=True)
|
||||
return {"error": str(e)}
|
||||
|
||||
|
||||
# =========================================================================
|
||||
# DIRECTIVE TOOLS
|
||||
# =========================================================================
|
||||
|
||||
@@ -27,7 +27,6 @@ from alembic.config import Config
|
||||
from alembic.script.revision import ResolutionError
|
||||
from sqlalchemy import Connection, create_engine, text
|
||||
|
||||
from ._pg_search import normalize_pg_search_tokenizer, pg_search_bm25_columns
|
||||
from ._vector_index import (
|
||||
bootstrap_extension,
|
||||
detect_vector_extension,
|
||||
@@ -804,7 +803,6 @@ def ensure_text_search_extension(
|
||||
database_url: str,
|
||||
text_search_extension: str = "native",
|
||||
schema: str | None = None,
|
||||
pg_search_tokenizer: str | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
Ensure the text search columns and indexes match the configured extension.
|
||||
@@ -817,18 +815,13 @@ def ensure_text_search_extension(
|
||||
|
||||
Args:
|
||||
database_url: SQLAlchemy database URL
|
||||
text_search_extension: Configured text search extension — one of
|
||||
"native", "vchord", "pg_textsearch", "pgroonga", or "pg_search"
|
||||
text_search_extension: Configured text search extension ("native" or "vchord")
|
||||
schema: Target PostgreSQL schema name (None for public)
|
||||
pg_search_tokenizer: Optional ParadeDB tokenizer to apply to pg_search
|
||||
BM25 text fields when indexes are created. Empty keeps the
|
||||
ParadeDB default.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If extension mismatch with existing data
|
||||
"""
|
||||
schema_name = schema or "public"
|
||||
pg_search_tokenizer = normalize_pg_search_tokenizer(pg_search_tokenizer)
|
||||
|
||||
engine = create_engine(to_libpq_url(database_url))
|
||||
with engine.connect() as conn:
|
||||
@@ -845,17 +838,6 @@ def ensure_text_search_extension(
|
||||
elif text_search_extension == "pg_textsearch":
|
||||
target_column_type = "text"
|
||||
target_index_type = "bm25"
|
||||
elif text_search_extension == "pgroonga":
|
||||
# pgroonga indexes the base text column directly. We keep a dummy
|
||||
# TEXT column named search_vector for symmetry with pg_textsearch
|
||||
# and so the column-type mismatch detection above keeps working.
|
||||
target_column_type = "text"
|
||||
target_index_type = "pgroonga"
|
||||
elif text_search_extension == "pg_search":
|
||||
# ParadeDB: same column type / access method as pg_textsearch.
|
||||
# Disambiguated below by inspecting the index reloptions (key_field).
|
||||
target_column_type = "text"
|
||||
target_index_type = "bm25"
|
||||
else: # native
|
||||
target_column_type = "tsvector"
|
||||
target_index_type = "gin"
|
||||
@@ -893,18 +875,16 @@ def ensure_text_search_extension(
|
||||
|
||||
if not current_column_info:
|
||||
logger.warning(f"No search_vector column found for {table_name}, will create it")
|
||||
mismatched_tables.append((table_name, None, None, False))
|
||||
mismatched_tables.append((table_name, None, None))
|
||||
continue
|
||||
|
||||
# Check column type (udt_name contains the actual type: tsvector, bm25vector, etc.)
|
||||
current_column_type = current_column_info[1] # udt_name
|
||||
|
||||
# Get current index type and definition. The definition lets us
|
||||
# disambiguate pg_textsearch vs pg_search (both register a `bm25`
|
||||
# access method but only pg_search uses the `key_field` reloption).
|
||||
# Get current index type
|
||||
current_index_info = conn.execute(
|
||||
text("""
|
||||
SELECT am.amname, pi.indexdef
|
||||
SELECT am.amname
|
||||
FROM pg_indexes pi
|
||||
JOIN pg_class c ON c.relname = pi.indexname
|
||||
JOIN pg_am am ON am.oid = c.relam
|
||||
@@ -916,21 +896,10 @@ def ensure_text_search_extension(
|
||||
).fetchone()
|
||||
|
||||
current_index_type = current_index_info[0] if current_index_info else None
|
||||
current_index_def = current_index_info[1] if current_index_info else None
|
||||
|
||||
# Detect pg_search specifically (vs pg_textsearch) via the key_field reloption
|
||||
current_is_pg_search = bool(current_index_def and "key_field" in current_index_def)
|
||||
want_pg_search = text_search_extension == "pg_search"
|
||||
|
||||
# Check if column and index types match target
|
||||
column_matches = current_column_type == target_column_type
|
||||
index_matches = current_index_type == target_index_type if current_index_type else False
|
||||
# When both target and current sit at column=text/index=bm25, the
|
||||
# access-method check alone can't tell pg_textsearch from pg_search —
|
||||
# require the key_field reloption to agree with the configured backend.
|
||||
if column_matches and index_matches and target_index_type == "bm25" and target_column_type == "text":
|
||||
if current_is_pg_search != want_pg_search:
|
||||
index_matches = False
|
||||
|
||||
if not (column_matches and index_matches):
|
||||
logger.info(
|
||||
@@ -938,7 +907,7 @@ def ensure_text_search_extension(
|
||||
f"column={current_column_type} (want {target_column_type}), "
|
||||
f"index={current_index_type} (want {target_index_type})"
|
||||
)
|
||||
mismatched_tables.append((table_name, current_column_type, current_index_type, current_is_pg_search))
|
||||
mismatched_tables.append((table_name, current_column_type, current_index_type))
|
||||
|
||||
# Check if table has data
|
||||
row_count = conn.execute(text(f"SELECT COUNT(*) FROM {schema_name}.{table_name}")).scalar()
|
||||
@@ -956,20 +925,14 @@ def ensure_text_search_extension(
|
||||
# If there's data in any mismatched table, raise error
|
||||
if tables_with_data:
|
||||
table_list = ", ".join([f"{table}({count} rows)" for table, count in tables_with_data])
|
||||
# Detect current extension from column type, index type, and (for the
|
||||
# text/bm25 ambiguity) the key_field reloption. tsvector is
|
||||
# unambiguous; text could be pg_textsearch, pgroonga, or pg_search.
|
||||
# Detect current extension from column type
|
||||
current_col_type = mismatched_tables[0][1]
|
||||
current_idx_type = mismatched_tables[0][2]
|
||||
first_is_pg_search = mismatched_tables[0][3]
|
||||
if current_col_type == "tsvector":
|
||||
current_ext = "native"
|
||||
elif current_col_type == "bm25vector":
|
||||
current_ext = "vchord"
|
||||
elif current_col_type == "text" and current_idx_type == "pgroonga":
|
||||
current_ext = "pgroonga"
|
||||
elif current_col_type == "text":
|
||||
current_ext = "pg_search" if first_is_pg_search else "pg_textsearch"
|
||||
current_ext = "pg_textsearch"
|
||||
else:
|
||||
current_ext = "unknown"
|
||||
raise RuntimeError(
|
||||
@@ -984,7 +947,7 @@ def ensure_text_search_extension(
|
||||
# Tables are empty, safe to recreate columns/indexes
|
||||
logger.info(f"Recreating text search columns/indexes for {text_search_extension}")
|
||||
|
||||
for table_name, current_col_type, current_idx_type, _was_pg_search in mismatched_tables:
|
||||
for table_name, current_col_type, current_idx_type in mismatched_tables:
|
||||
# Drop existing index if it exists
|
||||
if current_idx_type:
|
||||
logger.info(f"Dropping {current_idx_type} index on {table_name}")
|
||||
@@ -1037,79 +1000,21 @@ def ensure_text_search_extension(
|
||||
WITH (text_config='english')
|
||||
""")
|
||||
)
|
||||
elif text_search_extension == "pgroonga":
|
||||
# Ensure pgroonga extension is available
|
||||
try:
|
||||
conn.execute(text("CREATE EXTENSION IF NOT EXISTS pgroonga CASCADE"))
|
||||
except Exception:
|
||||
# Extension might already exist or user lacks permissions — verify
|
||||
has_ext = conn.execute(text("SELECT 1 FROM pg_extension WHERE extname = 'pgroonga'")).fetchone()
|
||||
if not has_ext:
|
||||
raise
|
||||
|
||||
logger.info(f"Creating dummy TEXT search_vector on {table_name} for pgroonga")
|
||||
# pgroonga indexes the base text column directly, but we keep a
|
||||
# dummy search_vector column for symmetry with pg_textsearch and
|
||||
# so the column-type mismatch detection above keeps working.
|
||||
conn.execute(text(f"ALTER TABLE {schema_name}.{table_name} ADD COLUMN search_vector TEXT"))
|
||||
|
||||
# pgroonga index expression mirrors pg_textsearch
|
||||
if table_name == "memory_units":
|
||||
index_expr = (
|
||||
"(COALESCE(text, '') || ' ' || COALESCE(context, '') || ' ' || COALESCE(text_signals, ''))"
|
||||
)
|
||||
else: # reflections
|
||||
index_expr = "(COALESCE(name, '') || ' ' || content)"
|
||||
|
||||
logger.info(f"Creating pgroonga index on {table_name}")
|
||||
# TokenBigram is the polyglot default — falls back to whitespace
|
||||
# tokenization for space-separated languages and bigram for CJK.
|
||||
# NormalizerNFKC150 handles Unicode normalization (full/half-width,
|
||||
# case folding, etc.) which materially improves Japanese recall.
|
||||
conn.execute(
|
||||
text(f"""
|
||||
CREATE INDEX idx_{table_name.replace(".", "_")}_text_search
|
||||
ON {schema_name}.{table_name}
|
||||
USING pgroonga ({index_expr})
|
||||
WITH (tokenizer='TokenBigram', normalizer='NormalizerNFKC150')
|
||||
""")
|
||||
)
|
||||
elif text_search_extension == "pg_search":
|
||||
logger.info(f"Creating TEXT column on {table_name}")
|
||||
# Dummy TEXT column for schema symmetry; pg_search indexes operate on base columns.
|
||||
conn.execute(text(f"ALTER TABLE {schema_name}.{table_name} ADD COLUMN search_vector TEXT"))
|
||||
|
||||
# ParadeDB BM25 index over the table's primary key and text columns.
|
||||
# Column list mirrors what the initial / text_signals migrations create.
|
||||
if table_name == "memory_units":
|
||||
bm25_cols = pg_search_bm25_columns(
|
||||
"id",
|
||||
("text", "context", "text_signals"),
|
||||
pg_search_tokenizer,
|
||||
)
|
||||
else: # reflections
|
||||
bm25_cols = pg_search_bm25_columns(
|
||||
"id",
|
||||
("name", "content"),
|
||||
pg_search_tokenizer,
|
||||
)
|
||||
|
||||
logger.info(f"Creating ParadeDB BM25 index on {table_name}")
|
||||
conn.execute(
|
||||
text(f"""
|
||||
CREATE INDEX idx_{table_name.replace(".", "_")}_text_search
|
||||
ON {schema_name}.{table_name}
|
||||
USING bm25 ({bm25_cols})
|
||||
WITH (key_field='id')
|
||||
""")
|
||||
)
|
||||
else: # native
|
||||
logger.info(f"Creating tsvector column on {table_name}")
|
||||
# Plain tsvector column. The application populates search_vector
|
||||
# at INSERT time via to_tsvector($lang, ...) using the configured
|
||||
# HINDSIGHT_API_TEXT_SEARCH_EXTENSION_NATIVE_LANGUAGE — see
|
||||
# ops_postgresql.insert_facts_batch.
|
||||
conn.execute(text(f"ALTER TABLE {schema_name}.{table_name} ADD COLUMN search_vector tsvector"))
|
||||
# Different GENERATED expression for each table
|
||||
if table_name == "memory_units":
|
||||
generated_expr = "to_tsvector('english', COALESCE(text, '') || ' ' || COALESCE(context, ''))"
|
||||
else: # reflections
|
||||
generated_expr = "to_tsvector('english', COALESCE(name, '') || ' ' || content)"
|
||||
|
||||
conn.execute(
|
||||
text(f"""
|
||||
ALTER TABLE {schema_name}.{table_name}
|
||||
ADD COLUMN search_vector tsvector
|
||||
GENERATED ALWAYS AS ({generated_expr}) STORED
|
||||
""")
|
||||
)
|
||||
|
||||
# Create GIN index
|
||||
logger.info(f"Creating GIN index on {table_name}")
|
||||
|
||||
@@ -24,7 +24,6 @@ class EmbeddedPostgres:
|
||||
password: str = DEFAULT_PASSWORD,
|
||||
database: str = DEFAULT_DATABASE,
|
||||
name: str = "hindsight",
|
||||
config: dict[str, str] | None = None,
|
||||
**kwargs,
|
||||
):
|
||||
self.port = port # None means pg0 will auto-assign
|
||||
@@ -32,11 +31,6 @@ class EmbeddedPostgres:
|
||||
self.password = password
|
||||
self.database = database
|
||||
self.name = name
|
||||
# Extra postgresql.conf settings forwarded to Pg0 (e.g. ``max_connections``).
|
||||
# Useful when tests spawn many xdist workers that each open a pool against
|
||||
# the same pg0 instance — the postgres default of 100 max_connections is
|
||||
# easy to exhaust under that fan-out.
|
||||
self.config = config
|
||||
self._pg0: Pg0 | None = None
|
||||
|
||||
def _get_pg0(self) -> Pg0:
|
||||
@@ -57,8 +51,6 @@ class EmbeddedPostgres:
|
||||
# Only set port if explicitly specified
|
||||
if self.port is not None:
|
||||
kwargs["port"] = self.port
|
||||
if self.config is not None:
|
||||
kwargs["config"] = self.config
|
||||
self._pg0 = Pg0(**kwargs)
|
||||
return self._pg0
|
||||
|
||||
|
||||
@@ -129,7 +129,6 @@ PROVIDER_NAME_MAPPING = {
|
||||
"vertexai": "google",
|
||||
"groq": "groq",
|
||||
"ollama": "ollama",
|
||||
"ollama-cloud": "ollama",
|
||||
"lmstudio": "lmstudio",
|
||||
"openai-codex": "openai",
|
||||
"claude-code": "anthropic",
|
||||
|
||||
@@ -256,7 +256,6 @@ def main():
|
||||
tenant_extension=tenant_extension,
|
||||
max_slots=config.worker_max_slots,
|
||||
slot_reservations=config.worker_slot_reservations,
|
||||
consolidation_bank_priority=config.worker_consolidation_bank_priority or None,
|
||||
)
|
||||
|
||||
# Create the HTTP app for metrics/health
|
||||
|
||||
@@ -24,7 +24,7 @@ from .exceptions import DeferOperation, RetryTaskAt
|
||||
from .stage import StageHolder, bind_holder
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from hindsight_api.engine.db.base import DatabaseBackend, DatabaseConnection
|
||||
from hindsight_api.engine.db.base import DatabaseBackend
|
||||
from hindsight_api.extensions.tenant import TenantExtension
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -133,7 +133,6 @@ class WorkerPoller:
|
||||
tenant_extension: "TenantExtension | None" = None,
|
||||
max_slots: int = 10,
|
||||
slot_reservations: dict[str, int] | None = None,
|
||||
consolidation_bank_priority: dict[str, int] | None = None,
|
||||
):
|
||||
"""
|
||||
Initialize the worker poller.
|
||||
@@ -151,11 +150,6 @@ class WorkerPoller:
|
||||
"retain": 3}). Reserved slots guarantee capacity for that operation type.
|
||||
Remaining slots (max_slots - sum of reservations) form a shared pool usable
|
||||
by any operation type. Defaults to {"consolidation": 2} if None.
|
||||
consolidation_bank_priority: Per-bank priority for consolidation scheduling.
|
||||
Maps bank name patterns to integer priorities (higher = claimed first).
|
||||
Patterns support ``*`` as wildcard. A bare ``*`` key is the catch-all default.
|
||||
When set, consolidation tasks are claimed in priority tiers rather than
|
||||
pure created_at order. None or empty dict preserves current behavior.
|
||||
"""
|
||||
self._backend = backend
|
||||
self._worker_id = worker_id
|
||||
@@ -174,9 +168,6 @@ class WorkerPoller:
|
||||
self._slot_reservations: dict[str, int] = (
|
||||
slot_reservations if slot_reservations is not None else {"consolidation": 2}
|
||||
)
|
||||
self._consolidation_bank_priority: dict[str, int] | None = (
|
||||
consolidation_bank_priority if consolidation_bank_priority else None
|
||||
)
|
||||
# Cache of which optional PG routines are installed on the server
|
||||
# (probed once, memoised for the life of the poller).
|
||||
from ..engine.db.optional_routines import OptionalRoutines
|
||||
@@ -196,18 +187,13 @@ class WorkerPoller:
|
||||
# schema we serviced so a busy tenant can't monopolize the poll order.
|
||||
self._next_schema_idx: int = 0
|
||||
|
||||
@staticmethod
|
||||
def _normalize_poll_schema(schema: str | None) -> str | None:
|
||||
"""Use None internally for the default schema because SQL helpers omit that prefix."""
|
||||
from ..config import DEFAULT_DATABASE_SCHEMA
|
||||
|
||||
return None if schema == DEFAULT_DATABASE_SCHEMA else schema
|
||||
|
||||
async def _get_schemas(self) -> list[str | None]:
|
||||
"""Get list of schemas to poll. Returns [None] for default schema (no prefix)."""
|
||||
from ..config import DEFAULT_DATABASE_SCHEMA
|
||||
|
||||
tenants = await self._tenant_extension.list_tenants()
|
||||
# Convert default schema to None for SQL compatibility (no prefix), keep others as-is
|
||||
return [self._normalize_poll_schema(t.schema) for t in tenants]
|
||||
return [t.schema if t.schema != DEFAULT_DATABASE_SCHEMA else None for t in tenants]
|
||||
|
||||
async def _scan_active_schemas(self, schemas: list[str | None]) -> set[str | None]:
|
||||
"""Find which schemas have pending work.
|
||||
@@ -227,57 +213,22 @@ class WorkerPoller:
|
||||
async with self._backend.acquire() as conn:
|
||||
if await self._optional_routines.is_installed(conn, "schemas_with_pending_work"):
|
||||
rows = await conn.fetch("SELECT * FROM public.schemas_with_pending_work()")
|
||||
routine_active = {self._normalize_poll_schema(r[0]) for r in rows}
|
||||
known_schemas = set(schemas)
|
||||
active = routine_active & known_schemas
|
||||
unknown = routine_active - known_schemas
|
||||
if unknown:
|
||||
logger.warning(
|
||||
"Optional PG routine public.schemas_with_pending_work() returned schema(s) "
|
||||
"not present in tenant discovery: %s",
|
||||
sorted(str(s) for s in unknown),
|
||||
return {r[0] for r in rows}
|
||||
|
||||
# Fallback: per-schema EXISTS checks from Python
|
||||
active: set[str | None] = set()
|
||||
for schema in schemas:
|
||||
table = fq_table("async_operations", schema)
|
||||
try:
|
||||
has_work = await conn.fetchval(
|
||||
f"SELECT EXISTS(SELECT 1 FROM {table} "
|
||||
f"WHERE status = 'pending' AND task_payload IS NOT NULL LIMIT 1)"
|
||||
)
|
||||
|
||||
# The optional routine returns PostgreSQL schema names, but the poller uses
|
||||
# None for the default schema. Older operator-supplied implementations also
|
||||
# commonly scan tenant_% only; when the default schema is in scope but absent
|
||||
# from the routine result, verify via the fully-correct per-schema fallback so
|
||||
# public single-tenant deployments cannot silently starve.
|
||||
should_verify_with_fallback = (None in known_schemas and None not in active) or (
|
||||
bool(routine_active) and not active
|
||||
)
|
||||
if not should_verify_with_fallback:
|
||||
return active
|
||||
|
||||
fallback_active = await self._scan_active_schemas_by_exists(conn, schemas)
|
||||
missed = fallback_active - active
|
||||
if missed:
|
||||
logger.warning(
|
||||
"Optional PG routine public.schemas_with_pending_work() missed claimable schema(s) %s; "
|
||||
"using per-schema fallback for this poll",
|
||||
sorted(str(s) for s in missed),
|
||||
)
|
||||
return fallback_active
|
||||
|
||||
return await self._scan_active_schemas_by_exists(conn, schemas)
|
||||
|
||||
async def _scan_active_schemas_by_exists(
|
||||
self, conn: "DatabaseConnection", schemas: list[str | None]
|
||||
) -> set[str | None]:
|
||||
"""Find active schemas using per-schema EXISTS checks."""
|
||||
active: set[str | None] = set()
|
||||
for schema in schemas:
|
||||
table = fq_table("async_operations", schema)
|
||||
try:
|
||||
has_work = await conn.fetchval(
|
||||
f"SELECT EXISTS(SELECT 1 FROM {table} "
|
||||
f"WHERE status = 'pending' AND task_payload IS NOT NULL LIMIT 1)"
|
||||
)
|
||||
if has_work:
|
||||
active.add(schema)
|
||||
except Exception:
|
||||
pass
|
||||
return active
|
||||
if has_work:
|
||||
active.add(schema)
|
||||
except Exception:
|
||||
pass
|
||||
return active
|
||||
|
||||
async def _get_available_slots(self) -> SlotAvailability:
|
||||
"""
|
||||
@@ -477,7 +428,6 @@ class WorkerPoller:
|
||||
self._worker_id,
|
||||
reserved_limits,
|
||||
shared_limit,
|
||||
consolidation_bank_priority=self._consolidation_bank_priority,
|
||||
)
|
||||
|
||||
if not all_rows:
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "hindsight-api-slim"
|
||||
version = "0.7.1"
|
||||
version = "0.6.2"
|
||||
description = "Hindsight: Agent Memory That Works Like Human Memory"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.11"
|
||||
@@ -18,14 +18,10 @@ dependencies = [
|
||||
"fastapi[standard]>=0.120.3",
|
||||
"uvicorn>=0.38.0",
|
||||
"wsproto>=1.0.0",
|
||||
# Cap below 2.1: SQLAlchemy 2.1 switches the default `postgresql://` DBAPI
|
||||
# from psycopg2 to psycopg (v3), which we don't ship — a bare install would
|
||||
# fail migrations with "No module named 'psycopg'". Pin to the tested 2.0
|
||||
# line (which keeps psycopg2 the default driver) until psycopg3 is adopted.
|
||||
"sqlalchemy>=2.0.44,<2.1",
|
||||
"sqlalchemy>=2.0.44",
|
||||
"alembic>=1.17.1",
|
||||
"pgvector>=0.4.1",
|
||||
"greenlet>=3.2.4,<3.4.0", # 3.4.0 lacks arm64 wheels for manylinux_2_41
|
||||
"greenlet>=3.2.4,<3.4.0", # 3.4.0 lacks arm64 wheels for manylinux_2_41
|
||||
"psycopg2-binary>=2.9.11",
|
||||
"tiktoken>=0.12.0",
|
||||
"httpx>=0.27.0",
|
||||
@@ -72,7 +68,7 @@ dependencies = [
|
||||
"tornado>=6.5.5", # DoS multipart/incomplete cookie validation fix
|
||||
"aiohttp>=3.13.3", # Multiple DoS vulnerabilities
|
||||
"pygments>=2.20.0", # ReDoS via inefficient GUID regex fix
|
||||
"claude-agent-sdk>=0.2.82",
|
||||
"claude-agent-sdk>=0.1.27",
|
||||
"boto3>=1.42.74",
|
||||
]
|
||||
|
||||
@@ -97,7 +93,7 @@ local-llm = [
|
||||
"huggingface-hub>=0.20.0",
|
||||
]
|
||||
embedded-db = [
|
||||
"pg0-embedded>=0.14.2",
|
||||
"pg0-embedded>=0.14.0",
|
||||
]
|
||||
oracle = [
|
||||
"oracledb>=2.5.0",
|
||||
@@ -146,9 +142,6 @@ addopts = "--timeout 300 -n 8 --dist loadgroup --durations=10 -v"
|
||||
markers = [
|
||||
"oracle: Oracle 23ai integration tests (require ORACLE_TEST_DSN env var)",
|
||||
"hs_llm_mat: LLM minimum acceptance tests — run in CI matrix across multiple providers",
|
||||
"hs_llm_core: Core pipeline tests that need a real LLM but only one provider",
|
||||
"integration: Live external-API integration tests (require provider credentials; skipped without)",
|
||||
"slow: Slow tests (minutes); not run in fast CI",
|
||||
]
|
||||
asyncio_mode = "auto"
|
||||
asyncio_default_fixture_loop_scope = "function"
|
||||
|
||||
@@ -95,14 +95,7 @@ def pg0_db_url(db_url, tmp_path_factory, worker_id):
|
||||
url = url_file.read_text().strip()
|
||||
else:
|
||||
# First worker - start pg0
|
||||
# Bump max_connections so 8 xdist workers * pool_max_size=15 fits well
|
||||
# under the cap (postgres default is 100, which is easy to exhaust now
|
||||
# that consolidation_llm_parallelism=4 increases peak conns per op).
|
||||
pg0 = EmbeddedPostgres(
|
||||
name=pg0_instance_name,
|
||||
port=pg0_instance_port,
|
||||
config={"max_connections": "300"},
|
||||
)
|
||||
pg0 = EmbeddedPostgres(name=pg0_instance_name, port=pg0_instance_port)
|
||||
|
||||
# Run ensure_running in a new event loop
|
||||
loop = asyncio.new_event_loop()
|
||||
@@ -315,7 +308,7 @@ async def oracle_memory(oracle_db_url, embeddings, cross_encoder, query_analyzer
|
||||
cross_encoder=cross_encoder,
|
||||
query_analyzer=query_analyzer,
|
||||
pool_min_size=1,
|
||||
pool_max_size=15,
|
||||
pool_max_size=5,
|
||||
run_migrations=False, # Already ran above
|
||||
task_backend=SyncTaskBackend(),
|
||||
)
|
||||
@@ -420,48 +413,21 @@ def query_analyzer():
|
||||
@pytest_asyncio.fixture(scope="function")
|
||||
async def memory(pg0_db_url, embeddings, cross_encoder, query_analyzer):
|
||||
"""
|
||||
Provide a MemoryEngine instance using a mock LLM for deterministic tests.
|
||||
Provide a MemoryEngine instance for each test.
|
||||
|
||||
The mock LLM returns canned facts derived from input text, allowing the
|
||||
full retain → recall → reflect pipeline to work without real LLM calls.
|
||||
This makes core tests fast, deterministic, and free from LLM flakiness.
|
||||
Must be function-scoped because:
|
||||
1. pytest-xdist runs tests in separate processes with different event loops
|
||||
2. asyncpg pools are bound to the event loop that created them
|
||||
3. Each test needs its own pool in its own event loop
|
||||
|
||||
Tests that need real LLM output quality should use `memory_real_llm` instead.
|
||||
Uses small pool sizes since tests run in parallel.
|
||||
Uses pg0_db_url (a postgresql:// URL) directly, so MemoryEngine won't try to
|
||||
manage pg0 lifecycle - that's handled by the session-scoped pg0_db_url fixture.
|
||||
Migrations are disabled here since they're run once at session scope in pg0_db_url.
|
||||
Uses SyncTaskBackend so async tasks execute immediately (no worker needed).
|
||||
"""
|
||||
mem = MemoryEngine(
|
||||
db_url=pg0_db_url,
|
||||
memory_llm_provider="mock",
|
||||
memory_llm_api_key="",
|
||||
memory_llm_model="mock",
|
||||
embeddings=embeddings,
|
||||
cross_encoder=cross_encoder,
|
||||
query_analyzer=query_analyzer,
|
||||
pool_min_size=1,
|
||||
pool_max_size=15,
|
||||
run_migrations=False,
|
||||
task_backend=SyncTaskBackend(),
|
||||
)
|
||||
await mem.initialize()
|
||||
yield mem
|
||||
try:
|
||||
if mem._pool and not mem._pool._closing:
|
||||
await mem.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
@pytest_asyncio.fixture(scope="function")
|
||||
async def memory_real_llm(pg0_db_url, embeddings, cross_encoder, query_analyzer):
|
||||
"""
|
||||
Provide a MemoryEngine instance using a real LLM provider.
|
||||
|
||||
Use this fixture ONLY for tests that assert on LLM output quality
|
||||
(fact extraction accuracy, language preservation, consolidation decisions, etc.).
|
||||
These tests are non-deterministic and should be marked with @pytest.mark.hs_llm_core
|
||||
(or @pytest.mark.hs_llm_mat for provider matrix acceptance tests).
|
||||
"""
|
||||
mem = MemoryEngine(
|
||||
db_url=pg0_db_url,
|
||||
db_url=pg0_db_url, # Direct postgresql:// URL, not pg0://
|
||||
memory_llm_provider=os.getenv("HINDSIGHT_API_LLM_PROVIDER", "groq"),
|
||||
memory_llm_api_key=os.getenv("HINDSIGHT_API_LLM_API_KEY"),
|
||||
memory_llm_model=os.getenv("HINDSIGHT_API_LLM_MODEL", "openai/gpt-oss-120b"),
|
||||
@@ -470,9 +436,9 @@ async def memory_real_llm(pg0_db_url, embeddings, cross_encoder, query_analyzer)
|
||||
cross_encoder=cross_encoder,
|
||||
query_analyzer=query_analyzer,
|
||||
pool_min_size=1,
|
||||
pool_max_size=15,
|
||||
run_migrations=False,
|
||||
task_backend=SyncTaskBackend(),
|
||||
pool_max_size=5,
|
||||
run_migrations=False, # Migrations already run at session scope
|
||||
task_backend=SyncTaskBackend(), # Execute tasks immediately in tests
|
||||
)
|
||||
await mem.initialize()
|
||||
yield mem
|
||||
@@ -500,7 +466,7 @@ async def memory_no_llm_verify(pg0_db_url, embeddings, cross_encoder, query_anal
|
||||
cross_encoder=cross_encoder,
|
||||
query_analyzer=query_analyzer,
|
||||
pool_min_size=1,
|
||||
pool_max_size=15,
|
||||
pool_max_size=5,
|
||||
run_migrations=False,
|
||||
task_backend=SyncTaskBackend(),
|
||||
skip_llm_verification=True, # Skip verification - will be overridden by test
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user