Compare commits
236
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
65d0f0f6e1 | ||
|
|
3469b8fac7 | ||
|
|
2f13d13d3e | ||
|
|
66b3bff400 | ||
|
|
08a75b5b84 | ||
|
|
9c9a5a290c | ||
|
|
0f084cc365 | ||
|
|
cba2b0d83e | ||
|
|
aefc1ebcc8 | ||
|
|
9c9d791752 | ||
|
|
45f47a9176 | ||
|
|
d53eb2b852 | ||
|
|
410f973578 | ||
|
|
08304800cc | ||
|
|
a49d19cd59 | ||
|
|
bdb3a55dc2 | ||
|
|
f1700af683 | ||
|
|
9c33a7c730 | ||
|
|
c81e62aeb9 | ||
|
|
33aacf5c6c | ||
|
|
ca180dde45 | ||
|
|
76a1bfa554 | ||
|
|
e90cfa4ac9 | ||
|
|
10785666c7 | ||
|
|
59f9a2bf25 | ||
|
|
30700de670 | ||
|
|
a63253f59f | ||
|
|
afd00c037c | ||
|
|
3d877b05a5 | ||
|
|
902704dfcf | ||
|
|
449a9d70b2 | ||
|
|
e301883952 | ||
|
|
487e2a5e6d | ||
|
|
511ca72361 | ||
|
|
a3b0d2651c | ||
|
|
b79caa9aa8 | ||
|
|
abbd3619c6 | ||
|
|
7126bf8a23 | ||
|
|
ba9d227f4c | ||
|
|
d05b49a24b | ||
|
|
858f0b3a06 | ||
|
|
ce137de643 | ||
|
|
920c56987b | ||
|
|
f5dfe59b90 | ||
|
|
9901aa1e07 | ||
|
|
9181c9a29e | ||
|
|
c8b898bd05 | ||
|
|
41710ba176 | ||
|
|
5d22a8e8fb | ||
|
|
3d6b380515 | ||
|
|
3c1431ad99 | ||
|
|
f75251fc5c | ||
|
|
52c148c203 | ||
|
|
21a22decea | ||
|
|
02ca15de42 | ||
|
|
6ae6663c0d | ||
|
|
ca561aca9e | ||
|
|
b52b483cb5 | ||
|
|
cbc196805d | ||
|
|
69383af896 | ||
|
|
abf24b4e22 | ||
|
|
1cbf9adb67 | ||
|
|
7ddfe9e237 | ||
|
|
2e74a324da | ||
|
|
5437cc0299 | ||
|
|
bca87412f7 | ||
|
|
eb9be90312 | ||
|
|
b8da88c854 | ||
|
|
824db5d5b8 | ||
|
|
eca8526e4a | ||
|
|
3ef8099892 | ||
|
|
8b80959bba | ||
|
|
f890479705 | ||
|
|
576c44d2ff | ||
|
|
f9042e378d | ||
|
|
6a17992e81 | ||
|
|
7d4fd1aa40 | ||
|
|
1f89731406 | ||
|
|
e1e5f36cee | ||
|
|
7ceaa22a66 | ||
|
|
0a16295c1f | ||
|
|
088dfecbc7 | ||
|
|
9e30ae2526 | ||
|
|
13f3052e6e | ||
|
|
654e4c0cfa | ||
|
|
a5e5372192 | ||
|
|
7007ffdb04 | ||
|
|
320d1ce435 | ||
|
|
c571fac7db | ||
|
|
3bedc1cedc | ||
|
|
70d60e96cf | ||
|
|
568e3c3028 | ||
|
|
cbaec36f66 | ||
|
|
d8aada7b0e | ||
|
|
16e1cc4934 | ||
|
|
3ee9437020 | ||
|
|
9671786faf | ||
|
|
33645e08cd | ||
|
|
2f5844b38d | ||
|
|
1267e61edd | ||
|
|
931f2a77ff | ||
|
|
eeff5001af | ||
|
|
f835c731fe | ||
|
|
b6dbd614fc | ||
|
|
32fc9b7477 | ||
|
|
e4f54a6071 | ||
|
|
343b972a95 | ||
|
|
58c02feef0 | ||
|
|
a5f8b58ab5 | ||
|
|
78008a1ad0 | ||
|
|
2128c02e0e | ||
|
|
2eab07834a | ||
|
|
9f41d98172 | ||
|
|
84bab9c5b7 | ||
|
|
d73e552189 | ||
|
|
712a862841 | ||
|
|
f64c5d2097 | ||
|
|
d4bf740618 | ||
|
|
70a7411659 | ||
|
|
dee581396b | ||
|
|
6a1d5fcd30 | ||
|
|
16ed93b9a2 | ||
|
|
581bbf3fc6 | ||
|
|
d00c843262 | ||
|
|
33442f1961 | ||
|
|
a525df4837 | ||
|
|
90a2201655 | ||
|
|
34365c3248 | ||
|
|
06c912df34 | ||
|
|
43dc50dd3f | ||
|
|
7d5d5b2781 | ||
|
|
b79ab2b752 | ||
|
|
cf9918891b | ||
|
|
9372462e13 | ||
|
|
f2fc8f9f26 | ||
|
|
6a80ecbf65 | ||
|
|
870bf4a3d1 | ||
|
|
099f4c925a | ||
|
|
e08faadc17 | ||
|
|
dbd1d1a743 | ||
|
|
c084765950 | ||
|
|
d6ad53986a | ||
|
|
6076354a9c | ||
|
|
9f9c3a1b40 | ||
|
|
8ba862b026 | ||
|
|
fd87de9c15 | ||
|
|
adc85129ba | ||
|
|
2ff805d6e9 | ||
|
|
8125a0d758 | ||
|
|
e1e137b027 | ||
|
|
6b5aa3afe8 | ||
|
|
e28b8c00f6 | ||
|
|
1be5ff33b0 | ||
|
|
aeb0c8b553 | ||
|
|
d0b2ab9ad2 | ||
|
|
93562bfaaf | ||
|
|
9679d8139d | ||
|
|
ab7feb144b | ||
|
|
0c4b79b6d3 | ||
|
|
e9270fd312 | ||
|
|
da21e0727c | ||
|
|
d4b8b3544b | ||
|
|
873223964b | ||
|
|
e5724fcba0 | ||
|
|
848451bd01 | ||
|
|
f82f58fa83 | ||
|
|
2d74007d80 | ||
|
|
05686e1236 | ||
|
|
93300b9104 | ||
|
|
9402572339 | ||
|
|
e9cc771bbd | ||
|
|
2a2b90b0a0 | ||
|
|
2635bbb49e | ||
|
|
2e88bac605 | ||
|
|
2644930561 | ||
|
|
bbd3c5dc04 | ||
|
|
4f9cf15cdd | ||
|
|
2d95f78b09 | ||
|
|
773ef0cb63 | ||
|
|
7b2263ba3b | ||
|
|
d054b88403 | ||
|
|
1c32a7b928 | ||
|
|
d38ecdb9ec | ||
|
|
aad07a141b | ||
|
|
3fc87e767c | ||
|
|
e22ae05f47 | ||
|
|
b57e337fa2 | ||
|
|
c05c491d77 | ||
|
|
fc941d5cae | ||
|
|
576016f5dc | ||
|
|
b3995d1430 | ||
|
|
f519fc4fd0 | ||
|
|
72fd3d59db | ||
|
|
5a61ac50e9 | ||
|
|
1f1716bdb0 | ||
|
|
61a8014f9d | ||
|
|
c5091d29cd | ||
|
|
e82bc56580 | ||
|
|
fa0e63b088 | ||
|
|
27cb7e43e0 | ||
|
|
9e23e83abf | ||
|
|
bdf93f0660 | ||
|
|
b0e8ac0f4d | ||
|
|
f74b577e02 | ||
|
|
3c633e5e16 | ||
|
|
cf0537ba7e | ||
|
|
e5944b63e7 | ||
|
|
37348c859e | ||
|
|
cece2c903c | ||
|
|
d7c73f4342 | ||
|
|
3b9d2db091 | ||
|
|
9790d904e0 | ||
|
|
2463efd0f2 | ||
|
|
6674ee4706 | ||
|
|
57f154454d | ||
|
|
4028dd91f8 | ||
|
|
0e81d1a25e | ||
|
|
8a2388a48f | ||
|
|
48185a4bee | ||
|
|
7e23f8e149 | ||
|
|
f659bb17c4 | ||
|
|
f31f82627c | ||
|
|
e1c6220f0e | ||
|
|
66cbdda3cb | ||
|
|
cf4bd598b4 | ||
|
|
443c94c827 | ||
|
|
26794aab09 | ||
|
|
7863ffeb49 | ||
|
|
9e2890ba81 | ||
|
|
e0e65c44f6 | ||
|
|
6881f63781 | ||
|
|
6cb309f72b | ||
|
|
f9fe6953a3 | ||
|
|
07de798c3b | ||
|
|
cefa75545a | ||
|
|
cd99eef4c5 |
@@ -157,7 +157,16 @@ If any files in `hindsight-integrations/` were added or changed, verify:
|
||||
- **Release process** — check that the integration name is in the `VALID_INTEGRATIONS` array in `scripts/release-integration.sh`. If missing, flag it.
|
||||
- **Code standards** — the integration code must follow all Python style rules (type hints, no raw dicts, no tuple returns, etc.).
|
||||
|
||||
### 10. Review against other coding standards
|
||||
### 10. Check MCP tool registration completeness
|
||||
|
||||
If any new MCP tools were added or existing tools renamed in `hindsight-api-slim/hindsight_api/mcp_tools.py`:
|
||||
- **`_ALL_TOOLS` set** in `mcp_tools.py` — must include the new tool name
|
||||
- **`tools_to_register` default set** in `register_mcp_tools()` in `mcp_tools.py` — must include the new tool name
|
||||
- **`_SINGLE_BANK_TOOLS` set** in `hindsight-api-slim/hindsight_api/api/mcp.py` — must include the new tool if it is bank-scoped (not a bank-management tool like `list_banks`/`create_bank`)
|
||||
- **`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. Review against other coding standards
|
||||
|
||||
Check the diff for violations of the standards listed above:
|
||||
- Python files at project root (not allowed)
|
||||
@@ -169,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)
|
||||
|
||||
### 11. Report findings
|
||||
### 12. Report findings
|
||||
|
||||
Present a clear summary organized by severity:
|
||||
|
||||
|
||||
@@ -34,7 +34,7 @@ jobs:
|
||||
env:
|
||||
UMAMI_URL: https://analytics.hindsight.vectorize.io
|
||||
UMAMI_WEBSITE_ID: ${{ secrets.UMAMI_WEBSITE_ID }}
|
||||
- uses: actions/upload-pages-artifact@v4
|
||||
- uses: actions/upload-pages-artifact@v5
|
||||
with:
|
||||
path: hindsight-docs/build
|
||||
deploy:
|
||||
|
||||
@@ -0,0 +1,174 @@
|
||||
name: Performance Tests
|
||||
|
||||
on:
|
||||
schedule:
|
||||
# Run daily at 06:00 UTC
|
||||
- cron: "0 6 * * *"
|
||||
workflow_dispatch:
|
||||
inputs:
|
||||
scale:
|
||||
description: "Test scale (perf-test)"
|
||||
type: choice
|
||||
options:
|
||||
- tiny
|
||||
- small
|
||||
- medium
|
||||
- large
|
||||
default: large
|
||||
suite:
|
||||
description: "Perf-test suite to run (blank = all)"
|
||||
type: choice
|
||||
options:
|
||||
- ""
|
||||
- retain
|
||||
- recall
|
||||
default: ""
|
||||
locomo_max_conversations:
|
||||
description: "LoComo max conversations (0 = skip, blank = all)"
|
||||
type: number
|
||||
default: 0
|
||||
locomo_skip:
|
||||
description: "Skip LoComo job"
|
||||
type: boolean
|
||||
default: false
|
||||
ref:
|
||||
description: "Git ref to test (branch, tag, or SHA). Defaults to main."
|
||||
type: string
|
||||
default: ""
|
||||
|
||||
concurrency:
|
||||
group: perf-test
|
||||
cancel-in-progress: true
|
||||
|
||||
jobs:
|
||||
perf-test:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ inputs.ref || github.ref }}
|
||||
|
||||
- name: Install uv
|
||||
uses: astral-sh/setup-uv@v7
|
||||
with:
|
||||
enable-cache: true
|
||||
prune-cache: false
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v6
|
||||
with:
|
||||
python-version-file: ".python-version"
|
||||
|
||||
- name: Cache HuggingFace models
|
||||
uses: actions/cache@v5
|
||||
with:
|
||||
path: ~/.cache/huggingface
|
||||
key: ${{ runner.os }}-huggingface-${{ hashFiles('hindsight-api-slim/pyproject.toml') }}
|
||||
restore-keys: |
|
||||
${{ runner.os }}-huggingface-
|
||||
|
||||
- name: Pre-download models
|
||||
working-directory: ./hindsight-api-slim
|
||||
run: |
|
||||
uv run --frozen --all-extras --index-strategy unsafe-best-match python -c "
|
||||
from sentence_transformers import SentenceTransformer
|
||||
print('Downloading embedding model...')
|
||||
SentenceTransformer('BAAI/bge-small-en-v1.5')
|
||||
print('Model downloaded successfully')
|
||||
"
|
||||
|
||||
- name: Install hindsight-dev dependencies
|
||||
run: |
|
||||
cd hindsight-dev && uv sync --frozen --all-extras --index-strategy unsafe-best-match
|
||||
|
||||
- name: Run perf tests
|
||||
run: |
|
||||
SUITE_ARG=""
|
||||
if [ -n "${{ inputs.suite }}" ]; then
|
||||
SUITE_ARG="--suite ${{ inputs.suite }}"
|
||||
fi
|
||||
./scripts/benchmarks/run-perf-test.sh \
|
||||
--scale ${{ inputs.scale || 'large' }} \
|
||||
$SUITE_ARG \
|
||||
--output perf-results.json
|
||||
|
||||
- name: Upload perf results
|
||||
if: always()
|
||||
uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: perf-results-${{ github.sha }}
|
||||
path: hindsight-dev/perf-results.json
|
||||
retention-days: 90
|
||||
|
||||
locomo:
|
||||
if: inputs.locomo_skip != true
|
||||
runs-on: ubuntu-latest
|
||||
env:
|
||||
HINDSIGHT_API_LLM_PROVIDER: vertexai
|
||||
HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY: /tmp/gcp-credentials.json
|
||||
HINDSIGHT_API_LLM_MODEL: google/gemini-2.5-flash-lite
|
||||
HINDSIGHT_API_JUDGE_LLM_PROVIDER: vertexai
|
||||
HINDSIGHT_API_JUDGE_LLM_MODEL: google/gemini-2.5-flash-lite
|
||||
HINDSIGHT_API_ANSWER_LLM_PROVIDER: vertexai
|
||||
HINDSIGHT_API_ANSWER_LLM_MODEL: google/gemini-3.1-pro-preview
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ inputs.ref || github.ref }}
|
||||
|
||||
- name: Setup GCP credentials
|
||||
run: |
|
||||
printf '%s' '${{ secrets.GCP_VERTEXAI_CREDENTIALS }}' > /tmp/gcp-credentials.json
|
||||
PROJECT_ID=$(jq -r '.project_id' /tmp/gcp-credentials.json)
|
||||
echo "HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID=$PROJECT_ID" >> $GITHUB_ENV
|
||||
|
||||
- name: Install uv
|
||||
uses: astral-sh/setup-uv@v7
|
||||
with:
|
||||
enable-cache: true
|
||||
prune-cache: false
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v6
|
||||
with:
|
||||
python-version-file: ".python-version"
|
||||
|
||||
- name: Cache HuggingFace models
|
||||
uses: actions/cache@v5
|
||||
with:
|
||||
path: ~/.cache/huggingface
|
||||
key: ${{ runner.os }}-huggingface-${{ hashFiles('hindsight-api-slim/pyproject.toml') }}
|
||||
restore-keys: |
|
||||
${{ runner.os }}-huggingface-
|
||||
|
||||
- name: Pre-download models
|
||||
working-directory: ./hindsight-api-slim
|
||||
run: |
|
||||
uv run --frozen --all-extras --index-strategy unsafe-best-match python -c "
|
||||
from sentence_transformers import SentenceTransformer
|
||||
print('Downloading embedding model...')
|
||||
SentenceTransformer('BAAI/bge-small-en-v1.5')
|
||||
print('Model downloaded successfully')
|
||||
"
|
||||
|
||||
- name: Install hindsight-dev dependencies
|
||||
run: |
|
||||
cd hindsight-dev && uv sync --frozen --all-extras --index-strategy unsafe-best-match
|
||||
|
||||
- name: Run LoComo benchmark
|
||||
run: |
|
||||
MAX_CONV_ARG=""
|
||||
if [ "${{ inputs.locomo_max_conversations }}" != "0" ] && [ -n "${{ inputs.locomo_max_conversations }}" ]; then
|
||||
MAX_CONV_ARG="--max-conversations ${{ inputs.locomo_max_conversations }}"
|
||||
fi
|
||||
uv run python hindsight-dev/benchmarks/locomo/locomo_benchmark.py \
|
||||
--wait-consolidation \
|
||||
$MAX_CONV_ARG
|
||||
|
||||
- name: Upload LoComo results
|
||||
if: always()
|
||||
uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: locomo-results-${{ github.sha }}
|
||||
path: hindsight-dev/benchmarks/locomo/results/
|
||||
retention-days: 90
|
||||
@@ -82,6 +82,15 @@ jobs:
|
||||
node-version: '22'
|
||||
registry-url: 'https://registry.npmjs.org'
|
||||
|
||||
# Guard: fail fast if the integration's lockfile resolves any dep from a
|
||||
# monorepo workspace (link=true) or a relative file path. The release
|
||||
# runner has no pre-built workspace `dist/` so `npm run build` would
|
||||
# later fail at tsc with "Cannot find module". See:
|
||||
# https://github.com/vectorize-io/hindsight/issues/… (0.6.0 openclaw retry)
|
||||
- name: Check integration lockfile
|
||||
if: steps.type.outputs.type == 'typescript'
|
||||
run: ./scripts/check-integration-lockfiles.sh
|
||||
|
||||
- name: Install dependencies
|
||||
if: steps.type.outputs.type == 'typescript'
|
||||
working-directory: ./hindsight-integrations/${{ steps.info.outputs.integration }}
|
||||
|
||||
@@ -150,6 +150,55 @@ jobs:
|
||||
path: hindsight-clients/typescript/*.tgz
|
||||
retention-days: 1
|
||||
|
||||
release-hindsight-all-npm:
|
||||
runs-on: ubuntu-latest
|
||||
environment: npm
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
|
||||
- name: Set up Node.js
|
||||
uses: actions/setup-node@v6
|
||||
with:
|
||||
node-version: '22'
|
||||
registry-url: 'https://registry.npmjs.org'
|
||||
cache: 'npm'
|
||||
cache-dependency-path: package-lock.json
|
||||
|
||||
- name: Install dependencies
|
||||
run: npm ci --workspace=hindsight-all-npm
|
||||
|
||||
- name: Build
|
||||
run: npm run build --workspace=hindsight-all-npm
|
||||
|
||||
- name: Publish to npm
|
||||
working-directory: ./hindsight-all-npm
|
||||
run: |
|
||||
set +e
|
||||
OUTPUT=$(npm publish --access public 2>&1)
|
||||
EXIT_CODE=$?
|
||||
echo "$OUTPUT"
|
||||
if [ $EXIT_CODE -ne 0 ]; then
|
||||
if echo "$OUTPUT" | grep -q "cannot publish over"; then
|
||||
echo "Package version already published, skipping..."
|
||||
exit 0
|
||||
fi
|
||||
exit $EXIT_CODE
|
||||
fi
|
||||
env:
|
||||
NODE_AUTH_TOKEN: ${{ secrets.NPM_TOKEN }}
|
||||
|
||||
- name: Pack for GitHub release
|
||||
working-directory: ./hindsight-all-npm
|
||||
run: npm pack
|
||||
|
||||
- name: Upload artifacts
|
||||
uses: actions/upload-artifact@v7
|
||||
with:
|
||||
name: hindsight-all-npm
|
||||
path: hindsight-all-npm/*.tgz
|
||||
retention-days: 1
|
||||
|
||||
release-control-plane:
|
||||
runs-on: ubuntu-latest
|
||||
environment: npm
|
||||
@@ -407,7 +456,7 @@ jobs:
|
||||
|
||||
create-github-release:
|
||||
runs-on: ubuntu-latest
|
||||
needs: [release-python-packages, release-typescript-client, release-control-plane, release-rust-cli, release-docker-images, release-helm-chart]
|
||||
needs: [release-python-packages, release-typescript-client, release-hindsight-all-npm, release-control-plane, release-rust-cli, release-docker-images, release-helm-chart]
|
||||
permissions:
|
||||
contents: write
|
||||
|
||||
@@ -436,6 +485,12 @@ jobs:
|
||||
name: control-plane
|
||||
path: ./artifacts/control-plane
|
||||
|
||||
- name: Download hindsight-embed npm wrapper
|
||||
uses: actions/download-artifact@v8
|
||||
with:
|
||||
name: hindsight-all-npm
|
||||
path: ./artifacts/hindsight-all-npm
|
||||
|
||||
- name: Download Rust CLI (Linux)
|
||||
uses: actions/download-artifact@v8
|
||||
with:
|
||||
@@ -472,6 +527,8 @@ jobs:
|
||||
cp artifacts/python-packages/hindsight-embed/dist/* release-assets/ || true
|
||||
# TypeScript client
|
||||
cp artifacts/typescript-client/*.tgz release-assets/ || true
|
||||
# hindsight-embed npm wrapper
|
||||
cp artifacts/hindsight-all-npm/*.tgz release-assets/ || true
|
||||
# Control Plane
|
||||
cp artifacts/control-plane/*.tgz release-assets/ || true
|
||||
# Rust CLI binaries
|
||||
@@ -483,7 +540,7 @@ jobs:
|
||||
ls -la release-assets/
|
||||
|
||||
- name: Create GitHub Release
|
||||
uses: softprops/action-gh-release@v2
|
||||
uses: softprops/action-gh-release@v3
|
||||
with:
|
||||
files: release-assets/*
|
||||
generate_release_notes: true
|
||||
|
||||
+345
-48
@@ -32,6 +32,7 @@ jobs:
|
||||
helm: ${{ steps.filter.outputs.helm }}
|
||||
docs: ${{ steps.filter.outputs.docs }}
|
||||
embed: ${{ steps.filter.outputs.embed }}
|
||||
all-npm: ${{ steps.filter.outputs.all-npm }}
|
||||
hindsight-all: ${{ steps.filter.outputs.hindsight-all }}
|
||||
integration-tests: ${{ steps.filter.outputs.integration-tests }}
|
||||
integrations-openclaw: ${{ steps.filter.outputs.integrations-openclaw }}
|
||||
@@ -43,8 +44,12 @@ jobs:
|
||||
integrations-litellm: ${{ steps.filter.outputs.integrations-litellm }}
|
||||
integrations-pydantic-ai: ${{ steps.filter.outputs.integrations-pydantic-ai }}
|
||||
integrations-ag2: ${{ steps.filter.outputs.integrations-ag2 }}
|
||||
integrations-hermes: ${{ steps.filter.outputs.integrations-hermes }}
|
||||
integrations-llamaindex: ${{ steps.filter.outputs.integrations-llamaindex }}
|
||||
integrations-paperclip: ${{ steps.filter.outputs.integrations-paperclip }}
|
||||
integrations-opencode: ${{ steps.filter.outputs.integrations-opencode }}
|
||||
integrations-cloudflare-oauth-proxy: ${{ steps.filter.outputs.integrations-cloudflare-oauth-proxy }}
|
||||
integrations-lockfiles: ${{ steps.filter.outputs.integrations-lockfiles }}
|
||||
integrations-openai-agents: ${{ steps.filter.outputs.integrations-openai-agents }}
|
||||
dev: ${{ steps.filter.outputs.dev }}
|
||||
ci: ${{ steps.filter.outputs.ci }}
|
||||
# Secrets are available for internal PRs, pull_request_review, and workflow_dispatch.
|
||||
@@ -91,6 +96,10 @@ jobs:
|
||||
- '*.md'
|
||||
embed:
|
||||
- 'hindsight-embed/**'
|
||||
all-npm:
|
||||
- 'hindsight-all-npm/**'
|
||||
- 'package.json'
|
||||
- 'package-lock.json'
|
||||
hindsight-all:
|
||||
- 'hindsight-all/**'
|
||||
integration-tests:
|
||||
@@ -113,15 +122,48 @@ jobs:
|
||||
- 'hindsight-integrations/pydantic-ai/**'
|
||||
integrations-ag2:
|
||||
- 'hindsight-integrations/ag2/**'
|
||||
integrations-hermes:
|
||||
- 'hindsight-integrations/hermes/**'
|
||||
integrations-llamaindex:
|
||||
- 'hindsight-integrations/llamaindex/**'
|
||||
integrations-paperclip:
|
||||
- 'hindsight-integrations/paperclip/**'
|
||||
integrations-opencode:
|
||||
- 'hindsight-integrations/opencode/**'
|
||||
integrations-cloudflare-oauth-proxy:
|
||||
- 'hindsight-integrations/cloudflare-oauth-proxy/**'
|
||||
integrations-lockfiles:
|
||||
- 'hindsight-integrations/*/package-lock.json'
|
||||
- 'hindsight-integrations/*/package.json'
|
||||
- 'scripts/check-integration-lockfiles.sh'
|
||||
integrations-openai-agents:
|
||||
- 'hindsight-integrations/openai-agents/**'
|
||||
dev:
|
||||
- 'hindsight-dev/**'
|
||||
ci:
|
||||
- '.github/**'
|
||||
|
||||
# Fail fast if any hindsight-integrations/*/package-lock.json was regenerated
|
||||
# from the monorepo root and ended up symlinked at a workspace path instead
|
||||
# of the npm registry. That bit us on the 0.6.0 openclaw release — tsc in
|
||||
# the release workflow couldn't find `@vectorize-io/hindsight-client`
|
||||
# because its `resolved` url pointed at a workspace dir whose `dist/` was
|
||||
# gitignored and unbuilt. Catching this at PR time means the release CI
|
||||
# never hits that class of failure.
|
||||
check-integration-lockfiles:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
github.event_name != 'pull_request_review' &&
|
||||
(github.event_name == 'workflow_dispatch' ||
|
||||
needs.detect-changes.outputs.integrations-lockfiles == '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: Check integration lockfiles resolve from the npm registry
|
||||
run: ./scripts/check-integration-lockfiles.sh
|
||||
|
||||
build-api-python-versions:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
@@ -180,12 +222,12 @@ jobs:
|
||||
- name: Build TypeScript client
|
||||
run: npm run build --workspace=hindsight-clients/typescript
|
||||
|
||||
build-openclaw-integration:
|
||||
build-hindsight-all-npm:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
github.event_name != 'pull_request_review' &&
|
||||
(github.event_name == 'workflow_dispatch' ||
|
||||
needs.detect-changes.outputs.integrations-openclaw == 'true' ||
|
||||
needs.detect-changes.outputs.all-npm == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
|
||||
@@ -198,18 +240,125 @@ jobs:
|
||||
uses: actions/setup-node@v6
|
||||
with:
|
||||
node-version: '22'
|
||||
cache: 'npm'
|
||||
cache-dependency-path: package-lock.json
|
||||
|
||||
- name: Install dependencies
|
||||
run: npm ci --workspace=hindsight-all-npm
|
||||
|
||||
- name: Run tests
|
||||
run: npm test --workspace=hindsight-all-npm
|
||||
|
||||
- name: Build
|
||||
run: npm run build --workspace=hindsight-all-npm
|
||||
|
||||
build-openclaw-integration:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
github.event_name != 'pull_request_review' &&
|
||||
(github.event_name == 'workflow_dispatch' ||
|
||||
needs.detect-changes.outputs.integrations-openclaw == 'true' ||
|
||||
needs.detect-changes.outputs.clients-ts == 'true' ||
|
||||
needs.detect-changes.outputs.all-npm == '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'
|
||||
cache: 'npm'
|
||||
cache-dependency-path: package-lock.json
|
||||
|
||||
# openclaw depends on two monorepo workspaces via `file:` deps:
|
||||
# @vectorize-io/hindsight-client and @vectorize-io/hindsight-all. Their
|
||||
# `dist/` directories are gitignored, so we must build them first.
|
||||
# Otherwise vitest/tsc in openclaw fails with
|
||||
# "Failed to resolve entry for package ..." on the value imports.
|
||||
- name: Install root workspace dependencies
|
||||
run: npm ci
|
||||
|
||||
- name: Build hindsight-client (openclaw dep)
|
||||
run: npm run build --workspace=hindsight-clients/typescript
|
||||
|
||||
- name: Build hindsight-all-npm (openclaw dep)
|
||||
run: npm run build --workspace=hindsight-all-npm
|
||||
|
||||
- name: Install openclaw dependencies
|
||||
working-directory: ./hindsight-integrations/openclaw
|
||||
run: npm ci
|
||||
|
||||
# Build must run before tests: one unit test in src/backfill.test.ts
|
||||
# creates a symlink to `$cwd/dist/backfill.js` and calls realpathSync on
|
||||
# it via isDirectExecution(). Without a populated dist/ the realpath call
|
||||
# throws, both paths stay unresolved, and the equality assertion fails.
|
||||
- name: Build
|
||||
working-directory: ./hindsight-integrations/openclaw
|
||||
run: npm run build
|
||||
|
||||
- name: Run tests
|
||||
working-directory: ./hindsight-integrations/openclaw
|
||||
run: npm test
|
||||
|
||||
- name: Build
|
||||
smoke-openclaw-install:
|
||||
needs: [detect-changes, build-openclaw-integration]
|
||||
if: >-
|
||||
github.event_name != 'pull_request_review' &&
|
||||
(github.event_name == 'workflow_dispatch' ||
|
||||
needs.detect-changes.outputs.integrations-openclaw == 'true' ||
|
||||
needs.detect-changes.outputs.clients-ts == 'true' ||
|
||||
needs.detect-changes.outputs.all-npm == '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'
|
||||
cache: 'npm'
|
||||
cache-dependency-path: package-lock.json
|
||||
|
||||
# Install the openclaw CLI globally. The smoke test exercises the real
|
||||
# `openclaw plugins install` / `openclaw config set` / `openclaw plugins
|
||||
# doctor` commands — not the in-repo integration tests — so a real CLI
|
||||
# must be on PATH.
|
||||
- name: Install openclaw CLI
|
||||
run: npm install -g openclaw
|
||||
|
||||
- name: Verify openclaw CLI
|
||||
run: openclaw --version
|
||||
|
||||
# openclaw depends on the workspace packages via published version
|
||||
# ranges (^0.1.0 / ^0.5.0), not file: paths, so the smoke test's
|
||||
# `openclaw plugins install <tarball>` resolves them straight from the
|
||||
# npm registry. These builds are just for `npm pack` / local unit
|
||||
# tests, not for resolving the plugin's runtime deps.
|
||||
- name: Install root workspace dependencies
|
||||
run: npm ci
|
||||
|
||||
- name: Build hindsight-client (openclaw dep)
|
||||
run: npm run build --workspace=hindsight-clients/typescript
|
||||
|
||||
- name: Build hindsight-all-npm (openclaw dep)
|
||||
run: npm run build --workspace=hindsight-all-npm
|
||||
|
||||
- name: Install openclaw dependencies
|
||||
working-directory: ./hindsight-integrations/openclaw
|
||||
run: npm run build
|
||||
run: npm ci
|
||||
|
||||
- name: Run openclaw install smoke test
|
||||
working-directory: ./hindsight-integrations/openclaw
|
||||
run: ./scripts/smoke-test.sh
|
||||
|
||||
test-claude-code-integration:
|
||||
needs: [detect-changes]
|
||||
@@ -326,6 +475,68 @@ jobs:
|
||||
working-directory: ./hindsight-integrations/ai-sdk
|
||||
run: npm run test:deno
|
||||
|
||||
test-opencode-integration:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
github.event_name != 'pull_request_review' &&
|
||||
(github.event_name == 'workflow_dispatch' ||
|
||||
needs.detect-changes.outputs.integrations-opencode == '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/opencode
|
||||
run: npm ci
|
||||
|
||||
- name: Run tests
|
||||
working-directory: ./hindsight-integrations/opencode
|
||||
run: npm test
|
||||
|
||||
- name: Build
|
||||
working-directory: ./hindsight-integrations/opencode
|
||||
run: npm run build
|
||||
|
||||
test-cloudflare-oauth-proxy-integration:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
github.event_name != 'pull_request_review' &&
|
||||
(github.event_name == 'workflow_dispatch' ||
|
||||
needs.detect-changes.outputs.integrations-cloudflare-oauth-proxy == '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/cloudflare-oauth-proxy
|
||||
run: npm ci
|
||||
|
||||
- name: Typecheck
|
||||
working-directory: ./hindsight-integrations/cloudflare-oauth-proxy
|
||||
run: npm run typecheck
|
||||
|
||||
- name: Run tests
|
||||
working-directory: ./hindsight-integrations/cloudflare-oauth-proxy
|
||||
run: npm test
|
||||
|
||||
build-chat-integration:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
@@ -357,6 +568,37 @@ jobs:
|
||||
working-directory: ./hindsight-integrations/chat
|
||||
run: npm run build
|
||||
|
||||
test-paperclip-integration:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
github.event_name != 'pull_request_review' &&
|
||||
(github.event_name == 'workflow_dispatch' ||
|
||||
needs.detect-changes.outputs.integrations-paperclip == '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/paperclip
|
||||
run: npm ci
|
||||
|
||||
- name: Build
|
||||
working-directory: ./hindsight-integrations/paperclip
|
||||
run: npm run build
|
||||
|
||||
- name: Run tests
|
||||
working-directory: ./hindsight-integrations/paperclip
|
||||
run: npm test
|
||||
|
||||
build-control-plane:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
@@ -1460,6 +1702,18 @@ jobs:
|
||||
print('Models downloaded successfully')
|
||||
"
|
||||
|
||||
# openclaw depends on @vectorize-io/hindsight-client and
|
||||
# @vectorize-io/hindsight-all via `file:` — their `dist/` directories are
|
||||
# gitignored and must be built before openclaw's npm ci copies them.
|
||||
- name: Install root workspace dependencies
|
||||
run: npm ci
|
||||
|
||||
- name: Build hindsight-client (openclaw dep)
|
||||
run: npm run build --workspace=hindsight-clients/typescript
|
||||
|
||||
- name: Build hindsight-all-npm (openclaw dep)
|
||||
run: npm run build --workspace=hindsight-all-npm
|
||||
|
||||
- name: Install openclaw integration dependencies
|
||||
working-directory: ./hindsight-integrations/openclaw
|
||||
run: npm ci
|
||||
@@ -1755,43 +2009,6 @@ jobs:
|
||||
working-directory: ./hindsight-integrations/pydantic-ai
|
||||
run: uv run pytest tests -v
|
||||
|
||||
test-hermes-integration:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
github.event_name != 'pull_request_review' &&
|
||||
(github.event_name == 'workflow_dispatch' ||
|
||||
needs.detect-changes.outputs.integrations-hermes == '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: 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 hermes integration
|
||||
working-directory: ./hindsight-integrations/hermes
|
||||
run: uv build
|
||||
|
||||
- name: Install dependencies
|
||||
working-directory: ./hindsight-integrations/hermes
|
||||
run: uv sync --frozen
|
||||
|
||||
- name: Run tests
|
||||
working-directory: ./hindsight-integrations/hermes
|
||||
run: uv run pytest tests -v
|
||||
|
||||
test-llamaindex-integration:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
@@ -1829,6 +2046,43 @@ jobs:
|
||||
working-directory: ./hindsight-integrations/llamaindex
|
||||
run: uv run pytest tests -v
|
||||
|
||||
test-openai-agents-integration:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
github.event_name != 'pull_request_review' &&
|
||||
(github.event_name == 'workflow_dispatch' ||
|
||||
needs.detect-changes.outputs.integrations-openai-agents == '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: 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 openai-agents integration
|
||||
working-directory: ./hindsight-integrations/openai-agents
|
||||
run: uv build
|
||||
|
||||
- name: Install dependencies
|
||||
working-directory: ./hindsight-integrations/openai-agents
|
||||
run: uv sync --frozen
|
||||
|
||||
- name: Run tests
|
||||
working-directory: ./hindsight-integrations/openai-agents
|
||||
run: uv run pytest tests -v
|
||||
|
||||
test-pip-slim:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
@@ -2333,6 +2587,9 @@ jobs:
|
||||
- name: Run generate-openapi
|
||||
run: ./scripts/generate-openapi.sh
|
||||
|
||||
- name: Run generate-bank-template-schema
|
||||
run: ./scripts/generate-bank-template-schema.sh
|
||||
|
||||
- name: Run generate-clients
|
||||
run: ./scripts/generate-clients.sh
|
||||
|
||||
@@ -2352,6 +2609,7 @@ jobs:
|
||||
echo ""
|
||||
echo "Please run the following commands locally and commit the changes:"
|
||||
echo " ./scripts/generate-openapi.sh"
|
||||
echo " ./scripts/generate-bank-template-schema.sh"
|
||||
echo " ./scripts/generate-clients.sh"
|
||||
echo " ./scripts/generate-docs-skill.sh"
|
||||
echo " ./scripts/hooks/lint.sh"
|
||||
@@ -2413,6 +2671,40 @@ jobs:
|
||||
cd hindsight-dev
|
||||
uv run check-openapi-compatibility /tmp/old-openapi.json ../hindsight-docs/static/openapi.json
|
||||
|
||||
check-cli-coverage:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
github.event_name != 'pull_request_review' &&
|
||||
(github.event_name == 'workflow_dispatch' ||
|
||||
needs.detect-changes.outputs.core == 'true' ||
|
||||
needs.detect-changes.outputs.cli == 'true' ||
|
||||
needs.detect-changes.outputs.dev == '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: Install uv
|
||||
uses: astral-sh/setup-uv@v7
|
||||
with:
|
||||
enable-cache: true
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v6
|
||||
with:
|
||||
python-version-file: ".python-version"
|
||||
|
||||
- name: Install hindsight-dev dependencies
|
||||
run: |
|
||||
cd hindsight-dev && uv sync --frozen --index-strategy unsafe-best-match
|
||||
|
||||
- name: Check CLI covers every OpenAPI operation
|
||||
run: |
|
||||
cd hindsight-dev
|
||||
uv run cli-coverage-check
|
||||
|
||||
# Report CI status back to the PR for pull_request_review events.
|
||||
# GitHub does not automatically link pull_request_review check runs to the PR,
|
||||
# so we create a commit status on the PR head SHA and post a comment.
|
||||
@@ -2420,14 +2712,19 @@ jobs:
|
||||
if: github.event_name == 'pull_request_review' && github.event.review.state == 'approved' && always()
|
||||
needs:
|
||||
- detect-changes
|
||||
- check-integration-lockfiles
|
||||
- build-api-python-versions
|
||||
- build-typescript-client
|
||||
- build-openclaw-integration
|
||||
- smoke-openclaw-install
|
||||
- test-claude-code-integration
|
||||
- test-codex-integration
|
||||
- build-ai-sdk-integration
|
||||
- test-ai-sdk-integration-deno
|
||||
- test-opencode-integration
|
||||
- test-cloudflare-oauth-proxy-integration
|
||||
- build-chat-integration
|
||||
- test-paperclip-integration
|
||||
- build-control-plane
|
||||
- build-docs
|
||||
- test-rust-cli
|
||||
@@ -2446,7 +2743,6 @@ jobs:
|
||||
- test-crewai-integration
|
||||
- test-litellm-integration
|
||||
- test-pydantic-ai-integration
|
||||
- test-hermes-integration
|
||||
- test-llamaindex-integration
|
||||
- test-pip-slim
|
||||
- test-embed
|
||||
@@ -2455,6 +2751,7 @@ jobs:
|
||||
- test-upgrade
|
||||
- verify-generated-files
|
||||
- check-openapi-compatibility
|
||||
- check-cli-coverage
|
||||
runs-on: ubuntu-latest
|
||||
permissions:
|
||||
statuses: write
|
||||
@@ -2462,7 +2759,7 @@ jobs:
|
||||
steps:
|
||||
- name: Determine overall result
|
||||
id: result
|
||||
uses: actions/github-script@v7
|
||||
uses: actions/github-script@v8
|
||||
with:
|
||||
script: |
|
||||
const needs = ${{ toJSON(needs) }};
|
||||
@@ -2495,7 +2792,7 @@ jobs:
|
||||
core.setOutput('run_url', runUrl);
|
||||
|
||||
- name: Report status to PR
|
||||
uses: actions/github-script@v7
|
||||
uses: actions/github-script@v8
|
||||
with:
|
||||
script: |
|
||||
await github.rest.repos.createCommitStatus({
|
||||
@@ -2509,7 +2806,7 @@ jobs:
|
||||
});
|
||||
|
||||
- name: Comment on PR
|
||||
uses: actions/github-script@v7
|
||||
uses: actions/github-script@v8
|
||||
with:
|
||||
script: |
|
||||
const prNumber = context.payload.pull_request.number;
|
||||
|
||||
@@ -0,0 +1,7 @@
|
||||
{
|
||||
"semi": true,
|
||||
"singleQuote": false,
|
||||
"tabWidth": 2,
|
||||
"trailingComma": "es5",
|
||||
"printWidth": 100
|
||||
}
|
||||
@@ -68,8 +68,9 @@ cd hindsight-control-plane && npm run dev
|
||||
./scripts/benchmarks/run-locomo.sh
|
||||
|
||||
# Performance benchmarks
|
||||
./scripts/benchmarks/run-perf-test.sh # System perf (mock LLM + pg0)
|
||||
./scripts/benchmarks/run-perf-test.sh --scale tiny # Quick smoke test
|
||||
./scripts/benchmarks/run-consolidation.sh
|
||||
./scripts/benchmarks/run-retain-perf.sh --document <path> # Requires API server running
|
||||
|
||||
# Results viewer
|
||||
./scripts/benchmarks/start-visualizer.sh # View results at localhost:8001
|
||||
@@ -222,6 +223,10 @@ Every new integration in `hindsight-integrations/` must satisfy all of the follo
|
||||
|
||||
If any of these are missing, the integration is incomplete and must not be pushed or merged.
|
||||
|
||||
### Changelogs
|
||||
|
||||
Never add "Unreleased" entries to changelogs (e.g. `hindsight-docs/src/pages/changelog/**`). Changelog entries are written by the release script (`./scripts/release-integration.sh`) when a version is actually cut. If a bug fix or feature needs documenting before release, describe it in the PR/commit — the release tooling will surface it in the published changelog section.
|
||||
|
||||
### Adding New API Configuration Flags
|
||||
|
||||
Configuration follows a hierarchical system: **Global (env vars) → Tenant (via extension) → Bank (database)**.
|
||||
|
||||
@@ -2,8 +2,8 @@ apiVersion: v2
|
||||
name: hindsight
|
||||
description: Hindsight helm chart
|
||||
type: application
|
||||
version: 0.4.22
|
||||
appVersion: "0.4.22"
|
||||
version: 0.5.4
|
||||
appVersion: "0.5.4"
|
||||
keywords:
|
||||
- ai
|
||||
- memory
|
||||
|
||||
@@ -95,6 +95,27 @@ spec:
|
||||
{{- toYaml .Values.api.readinessProbe | nindent 10 }}
|
||||
resources:
|
||||
{{- toYaml .Values.api.resources | nindent 10 }}
|
||||
{{- if or .Values.api.persistence.modelCache.enabled .Values.api.extraVolumeMounts }}
|
||||
volumeMounts:
|
||||
{{- if .Values.api.persistence.modelCache.enabled }}
|
||||
- name: model-cache
|
||||
mountPath: /home/hindsight/.cache
|
||||
{{- end }}
|
||||
{{- with .Values.api.extraVolumeMounts }}
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
{{- end }}
|
||||
{{- if or .Values.api.persistence.modelCache.enabled .Values.api.extraVolumes }}
|
||||
volumes:
|
||||
{{- if .Values.api.persistence.modelCache.enabled }}
|
||||
- name: model-cache
|
||||
persistentVolumeClaim:
|
||||
claimName: {{ include "hindsight.fullname" . }}-api-model-cache
|
||||
{{- end }}
|
||||
{{- with .Values.api.extraVolumes }}
|
||||
{{- toYaml . | nindent 6 }}
|
||||
{{- end }}
|
||||
{{- end }}
|
||||
{{- with .Values.nodeSelector }}
|
||||
nodeSelector:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
|
||||
@@ -0,0 +1,21 @@
|
||||
{{- if and .Values.api.enabled .Values.api.persistence.modelCache.enabled }}
|
||||
apiVersion: v1
|
||||
kind: PersistentVolumeClaim
|
||||
metadata:
|
||||
name: {{ include "hindsight.fullname" . }}-api-model-cache
|
||||
labels:
|
||||
{{- include "hindsight.api.labels" . | nindent 4 }}
|
||||
{{- with .Values.api.persistence.modelCache.annotations }}
|
||||
annotations:
|
||||
{{- toYaml . | nindent 4 }}
|
||||
{{- end }}
|
||||
spec:
|
||||
accessModes:
|
||||
{{- toYaml .Values.api.persistence.modelCache.accessModes | nindent 4 }}
|
||||
{{- if .Values.api.persistence.modelCache.storageClass }}
|
||||
storageClassName: {{ .Values.api.persistence.modelCache.storageClass }}
|
||||
{{- end }}
|
||||
resources:
|
||||
requests:
|
||||
storage: {{ .Values.api.persistence.modelCache.size }}
|
||||
{{- end }}
|
||||
@@ -95,6 +95,16 @@ spec:
|
||||
{{- toYaml .Values.worker.readinessProbe | nindent 10 }}
|
||||
resources:
|
||||
{{- toYaml .Values.worker.resources | nindent 10 }}
|
||||
{{- if or .Values.worker.persistence.modelCache.enabled .Values.worker.extraVolumeMounts }}
|
||||
volumeMounts:
|
||||
{{- if .Values.worker.persistence.modelCache.enabled }}
|
||||
- name: model-cache
|
||||
mountPath: /home/hindsight/.cache
|
||||
{{- end }}
|
||||
{{- with .Values.worker.extraVolumeMounts }}
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
{{- end }}
|
||||
{{- with .Values.nodeSelector }}
|
||||
nodeSelector:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
@@ -107,4 +117,26 @@ spec:
|
||||
tolerations:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
{{- with .Values.worker.extraVolumes }}
|
||||
volumes:
|
||||
{{- toYaml . | nindent 6 }}
|
||||
{{- end }}
|
||||
{{- if .Values.worker.persistence.modelCache.enabled }}
|
||||
volumeClaimTemplates:
|
||||
- metadata:
|
||||
name: model-cache
|
||||
{{- with .Values.worker.persistence.modelCache.annotations }}
|
||||
annotations:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
spec:
|
||||
accessModes:
|
||||
{{- toYaml .Values.worker.persistence.modelCache.accessModes | nindent 8 }}
|
||||
{{- if .Values.worker.persistence.modelCache.storageClass }}
|
||||
storageClassName: {{ .Values.worker.persistence.modelCache.storageClass }}
|
||||
{{- end }}
|
||||
resources:
|
||||
requests:
|
||||
storage: {{ .Values.worker.persistence.modelCache.size }}
|
||||
{{- end }}
|
||||
{{- end }}
|
||||
|
||||
@@ -67,6 +67,33 @@ api:
|
||||
# Pod affinity/anti-affinity (overrides global affinity for this component)
|
||||
# affinity: {}
|
||||
|
||||
# Persistent volume for local model cache (reranker, embeddings)
|
||||
# Models are downloaded to /home/hindsight/.cache on first use.
|
||||
# Without persistence, models are re-downloaded on every pod restart.
|
||||
persistence:
|
||||
modelCache:
|
||||
enabled: false
|
||||
size: 5Gi
|
||||
storageClass: ""
|
||||
accessModes:
|
||||
- ReadWriteOnce
|
||||
annotations: {}
|
||||
|
||||
# Extra volume mounts for the api container
|
||||
# e.g.
|
||||
# extraVolumeMounts:
|
||||
# - name: my-volume
|
||||
# mountPath: /mnt/my-volume
|
||||
extraVolumeMounts: []
|
||||
|
||||
# Extra volumes for the api pod
|
||||
# e.g.
|
||||
# extraVolumes:
|
||||
# - name: my-volume
|
||||
# configMap:
|
||||
# name: my-configmap
|
||||
extraVolumes: []
|
||||
|
||||
# Environment variables
|
||||
env:
|
||||
#HINDSIGHT_API_LLM_PROVIDER: "groq"
|
||||
@@ -140,6 +167,32 @@ worker:
|
||||
# Pod affinity/anti-affinity (overrides global affinity for this component)
|
||||
# affinity: {}
|
||||
|
||||
# Persistent volume for local model cache (reranker, embeddings)
|
||||
# Uses volumeClaimTemplates since worker is a StatefulSet.
|
||||
persistence:
|
||||
modelCache:
|
||||
enabled: false
|
||||
size: 5Gi
|
||||
storageClass: ""
|
||||
accessModes:
|
||||
- ReadWriteOnce
|
||||
annotations: {}
|
||||
|
||||
# Extra volume mounts for the worker container
|
||||
# e.g.
|
||||
# extraVolumeMounts:
|
||||
# - name: my-volume
|
||||
# mountPath: /mnt/my-volume
|
||||
extraVolumeMounts: []
|
||||
|
||||
# Extra volumes for the worker pod
|
||||
# e.g.
|
||||
# extraVolumes:
|
||||
# - name: my-volume
|
||||
# configMap:
|
||||
# name: my-configmap
|
||||
extraVolumes: []
|
||||
|
||||
# Secret environment variables (inherited from api.secrets if not specified)
|
||||
secrets: {}
|
||||
|
||||
|
||||
@@ -0,0 +1,4 @@
|
||||
node_modules
|
||||
dist
|
||||
*.tgz
|
||||
.DS_Store
|
||||
@@ -0,0 +1,80 @@
|
||||
# @vectorize-io/hindsight-all
|
||||
|
||||
Node.js equivalent of the Python [`hindsight-all`](https://pypi.org/project/hindsight-all/) package — programmatic lifecycle manager for a local Hindsight daemon. Use this when you want to embed Hindsight in a Node application without hand-rolling subprocess management.
|
||||
|
||||
This package deliberately does **not** ship an HTTP client. Once the daemon is running, talk to it with [`@vectorize-io/hindsight-client`](https://www.npmjs.com/package/@vectorize-io/hindsight-client) against `server.getBaseUrl()`. The two packages compose — one owns the daemon process, the other owns the HTTP API surface.
|
||||
|
||||
## Requirements
|
||||
|
||||
- **Node.js >= 22** — uses global `fetch` and `AbortSignal.timeout`.
|
||||
- **`uv` / `uvx`** on `PATH` — used to download and run the underlying `hindsight-embed` daemon on first use. Install via <https://docs.astral.sh/uv/>.
|
||||
|
||||
## Install
|
||||
|
||||
```bash
|
||||
npm install @vectorize-io/hindsight-all @vectorize-io/hindsight-client
|
||||
```
|
||||
|
||||
## Example
|
||||
|
||||
```ts
|
||||
import { HindsightServer, consoleLogger } from '@vectorize-io/hindsight-all';
|
||||
import { HindsightClient } from '@vectorize-io/hindsight-client';
|
||||
|
||||
const server = new HindsightServer({
|
||||
profile: 'my-app',
|
||||
port: 9077,
|
||||
env: {
|
||||
HINDSIGHT_API_LLM_PROVIDER: 'anthropic',
|
||||
HINDSIGHT_API_LLM_API_KEY: process.env.ANTHROPIC_API_KEY,
|
||||
HINDSIGHT_API_LLM_MODEL: 'claude-sonnet-4-20250514',
|
||||
HINDSIGHT_EMBED_DAEMON_IDLE_TIMEOUT: '0',
|
||||
},
|
||||
logger: consoleLogger,
|
||||
});
|
||||
|
||||
await server.start();
|
||||
|
||||
const client = new HindsightClient({ baseUrl: server.getBaseUrl() });
|
||||
|
||||
await client.retain('user-123', 'User prefers dark mode and concise answers.', {
|
||||
documentId: 'pref-2026-04-01',
|
||||
});
|
||||
|
||||
const recall = await client.recall('user-123', 'what are the user preferences?');
|
||||
console.log(recall.results);
|
||||
|
||||
await server.stop();
|
||||
```
|
||||
|
||||
For a remote Hindsight API, skip `HindsightServer` entirely and just point `HindsightClient` at the remote URL.
|
||||
|
||||
## Open config — forward-compatible with new daemon flags
|
||||
|
||||
`HindsightServerOptions` is designed so every new environment variable or CLI flag in the underlying Hindsight daemon can be used without waiting for a wrapper release:
|
||||
|
||||
- **`env`** accepts an arbitrary `Record<string, string>`. Every entry is exported into the daemon process and written into the profile config via `--env KEY=VALUE`.
|
||||
- **`extraProfileCreateArgs`** / **`extraDaemonStartArgs`** append raw args to the respective commands.
|
||||
|
||||
## Development against a local checkout
|
||||
|
||||
If you're hacking on the Python `hindsight-embed` package in the same monorepo, point the server at the local path — it'll use `uv run --directory <path>` instead of `uvx`:
|
||||
|
||||
```ts
|
||||
new HindsightServer({
|
||||
embedPackagePath: '/path/to/hindsight-embed',
|
||||
// ...
|
||||
});
|
||||
```
|
||||
|
||||
## API surface
|
||||
|
||||
- `HindsightServer` — daemon lifecycle (`start`, `stop`, `checkHealth`, `getBaseUrl`, `getProfile`).
|
||||
- `Logger` interface plus `silentLogger` (default) and `consoleLogger` helpers.
|
||||
- `getEmbedCommand(opts)` — low-level helper that returns the `[cmd, ...args]` tuple used to invoke the underlying Python CLI.
|
||||
|
||||
For memory operations (retain, recall, reflect, bank management, stats) use [`@vectorize-io/hindsight-client`](https://www.npmjs.com/package/@vectorize-io/hindsight-client).
|
||||
|
||||
## License
|
||||
|
||||
MIT
|
||||
@@ -0,0 +1,57 @@
|
||||
{
|
||||
"name": "@vectorize-io/hindsight-all",
|
||||
"version": "0.5.4",
|
||||
"description": "Node.js programmatic lifecycle manager for Hindsight — embeds a local hindsight daemon in a Node application. Pair with @vectorize-io/hindsight-client for memory operations.",
|
||||
"main": "dist/index.js",
|
||||
"types": "dist/index.d.ts",
|
||||
"type": "module",
|
||||
"exports": {
|
||||
".": {
|
||||
"types": "./dist/index.d.ts",
|
||||
"import": "./dist/index.js"
|
||||
}
|
||||
},
|
||||
"keywords": [
|
||||
"hindsight",
|
||||
"hindsight-all",
|
||||
"memory",
|
||||
"ai",
|
||||
"agent",
|
||||
"long-term-memory",
|
||||
"llm",
|
||||
"embedded-server"
|
||||
],
|
||||
"author": "Vectorize <[email protected]>",
|
||||
"license": "MIT",
|
||||
"repository": {
|
||||
"type": "git",
|
||||
"url": "https://github.com/vectorize-io/hindsight.git",
|
||||
"directory": "hindsight-all-npm"
|
||||
},
|
||||
"files": [
|
||||
"dist",
|
||||
"README.md"
|
||||
],
|
||||
"scripts": {
|
||||
"build": "tsup",
|
||||
"dev": "tsup --watch",
|
||||
"clean": "rm -rf dist",
|
||||
"test": "vitest run src",
|
||||
"test:watch": "vitest src",
|
||||
"prepublishOnly": "npm run clean && npm run build"
|
||||
},
|
||||
"devDependencies": {
|
||||
"@types/node": "^22.0.0",
|
||||
"tsup": "^8.5.1",
|
||||
"typescript": "^5.7.0",
|
||||
"vitest": "^4.1.2"
|
||||
},
|
||||
"engines": {
|
||||
"node": ">=22"
|
||||
},
|
||||
"overrides": {
|
||||
"rollup": "^4.59.0",
|
||||
"picomatch": ">=2.3.2 <3.0.0 || >=4.0.4",
|
||||
"vite": ">=8.0.5"
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,32 @@
|
||||
import { describe, it, expect } from 'vitest';
|
||||
import { getEmbedCommand } from './command.js';
|
||||
|
||||
describe('getEmbedCommand', () => {
|
||||
it('defaults to uvx hindsight-embed@latest', () => {
|
||||
expect(getEmbedCommand()).toEqual(['uvx', 'hindsight-embed@latest']);
|
||||
});
|
||||
|
||||
it('honours an explicit version', () => {
|
||||
expect(getEmbedCommand({ embedVersion: '0.5.0' })).toEqual(['uvx', '[email protected]']);
|
||||
});
|
||||
|
||||
it('treats an empty version as latest', () => {
|
||||
expect(getEmbedCommand({ embedVersion: '' })).toEqual(['uvx', 'hindsight-embed@latest']);
|
||||
});
|
||||
|
||||
it('uses uv run --directory when a local path is given', () => {
|
||||
expect(getEmbedCommand({ embedPackagePath: '/abs/path' })).toEqual([
|
||||
'uv',
|
||||
'run',
|
||||
'--directory',
|
||||
'/abs/path',
|
||||
'hindsight-embed',
|
||||
]);
|
||||
});
|
||||
|
||||
it('local path takes precedence over version', () => {
|
||||
expect(
|
||||
getEmbedCommand({ embedPackagePath: '/abs/path', embedVersion: '0.5.0' }),
|
||||
).toEqual(['uv', 'run', '--directory', '/abs/path', 'hindsight-embed']);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,25 @@
|
||||
/**
|
||||
* Resolve the command that invokes the `hindsight-embed` Python CLI.
|
||||
*
|
||||
* - If `embedPackagePath` is set, runs the package from a local checkout via
|
||||
* `uv run --directory <path> hindsight-embed`. Used for in-repo development.
|
||||
* - Otherwise runs it via `uvx hindsight-embed@<version>` so no global install
|
||||
* is required.
|
||||
*
|
||||
* Returns the argv as `[command, ...baseArgs]` suitable for `spawn()` /
|
||||
* `execFile()` (never shell-interpolated).
|
||||
*/
|
||||
export interface EmbedCommandOptions {
|
||||
/** Version spec passed to uvx (e.g. "latest", "0.5.0"). Default: "latest". */
|
||||
embedVersion?: string;
|
||||
/** Local checkout path. When set, overrides `embedVersion` and uses `uv run`. */
|
||||
embedPackagePath?: string;
|
||||
}
|
||||
|
||||
export function getEmbedCommand(opts: EmbedCommandOptions = {}): string[] {
|
||||
if (opts.embedPackagePath) {
|
||||
return ['uv', 'run', '--directory', opts.embedPackagePath, 'hindsight-embed'];
|
||||
}
|
||||
const version = opts.embedVersion && opts.embedVersion.length > 0 ? opts.embedVersion : 'latest';
|
||||
return ['uvx', `hindsight-embed@${version}`];
|
||||
}
|
||||
@@ -0,0 +1,7 @@
|
||||
export { HindsightServer } from './server.js';
|
||||
export { getEmbedCommand } from './command.js';
|
||||
export { silentLogger, consoleLogger } from './logger.js';
|
||||
|
||||
export type { Logger } from './logger.js';
|
||||
export type { EmbedCommandOptions } from './command.js';
|
||||
export type { HindsightServerOptions } from './types.js';
|
||||
@@ -0,0 +1,29 @@
|
||||
/**
|
||||
* Pluggable logger interface.
|
||||
*
|
||||
* This package does not own any logging infrastructure — consumers inject
|
||||
* whatever they want (console, pino, openclaw's logger, a no-op). The default
|
||||
* is silent so embedding this package never adds noise to an unrelated app.
|
||||
*/
|
||||
export interface Logger {
|
||||
debug(msg: string): void;
|
||||
info(msg: string): void;
|
||||
warn(msg: string): void;
|
||||
error(msg: string): void;
|
||||
}
|
||||
|
||||
/** Logger that drops every call. Used when no logger is passed. */
|
||||
export const silentLogger: Logger = {
|
||||
debug: () => {},
|
||||
info: () => {},
|
||||
warn: () => {},
|
||||
error: () => {},
|
||||
};
|
||||
|
||||
/** Logger that writes to the standard console. Handy for CLIs and tests. */
|
||||
export const consoleLogger: Logger = {
|
||||
debug: (msg) => console.debug(msg),
|
||||
info: (msg) => console.log(msg),
|
||||
warn: (msg) => console.warn(msg),
|
||||
error: (msg) => console.error(msg),
|
||||
};
|
||||
@@ -0,0 +1,35 @@
|
||||
import { describe, it, expect } from 'vitest';
|
||||
import { HindsightServer } from './server.js';
|
||||
|
||||
describe('HindsightServer construction', () => {
|
||||
it('defaults base URL to http://127.0.0.1:8888', () => {
|
||||
const server = new HindsightServer();
|
||||
expect(server.getBaseUrl()).toBe('http://127.0.0.1:8888');
|
||||
expect(server.getProfile()).toBe('default');
|
||||
});
|
||||
|
||||
it('honours custom profile, port, and host', () => {
|
||||
const server = new HindsightServer({ profile: 'app', port: 9077, host: '0.0.0.0' });
|
||||
expect(server.getProfile()).toBe('app');
|
||||
expect(server.getBaseUrl()).toBe('http://0.0.0.0:9077');
|
||||
});
|
||||
|
||||
it('accepts open env pass-through without complaining about unknown keys', () => {
|
||||
const server = new HindsightServer({
|
||||
env: {
|
||||
HINDSIGHT_API_LLM_PROVIDER: 'openai',
|
||||
HINDSIGHT_API_LLM_MODEL: 'gpt-4o-mini',
|
||||
// A field that does not exist today — should still be accepted
|
||||
HINDSIGHT_FUTURE_FLAG: 'enabled',
|
||||
},
|
||||
});
|
||||
expect(server).toBeInstanceOf(HindsightServer);
|
||||
});
|
||||
|
||||
it('exposes checkHealth that returns false when no daemon is running', async () => {
|
||||
// Random high port that nothing is listening on.
|
||||
const server = new HindsightServer({ port: 1, readyTimeoutMs: 100 });
|
||||
const healthy = await server.checkHealth();
|
||||
expect(healthy).toBe(false);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,322 @@
|
||||
import { spawn } from 'child_process';
|
||||
import { getEmbedCommand } from './command.js';
|
||||
import { silentLogger } from './logger.js';
|
||||
import type { Logger } from './logger.js';
|
||||
import type { HindsightServerOptions } from './types.js';
|
||||
|
||||
const DEFAULT_PORT = 8888;
|
||||
const DEFAULT_HOST = '127.0.0.1';
|
||||
const DEFAULT_PROFILE = 'default';
|
||||
const DEFAULT_READY_TIMEOUT_MS = 30_000;
|
||||
const DEFAULT_READY_POLL_INTERVAL_MS = 1_000;
|
||||
|
||||
/**
|
||||
* Manages the lifecycle of a local Hindsight daemon from a Node.js process.
|
||||
*
|
||||
* On {@link start}, this class:
|
||||
* 1. Resolves the `hindsight-embed` command (via `uvx` or a local `uv run`).
|
||||
* 2. Runs `profile create <name> --merge --port <port> [--env K=V ...]`
|
||||
* with every entry in {@link HindsightServerOptions.env} forwarded as
|
||||
* an `--env` flag.
|
||||
* 3. Runs `daemon --profile <name> start` and waits for the start command
|
||||
* to exit.
|
||||
* 4. Polls `http://host:port/health` until it returns `200` or the
|
||||
* `readyTimeoutMs` budget is exhausted.
|
||||
*
|
||||
* On {@link stop}, it runs `daemon --profile <name> stop` and returns once
|
||||
* the command exits (or after a short grace period).
|
||||
*
|
||||
* This is the Node.js equivalent of the Python `hindsight-all` package's
|
||||
* `HindsightServer`: a thin programmatic lifecycle wrapper around the
|
||||
* Hindsight daemon. It does NOT ship an HTTP client — once `start()`
|
||||
* resolves, use `@vectorize-io/hindsight-client` against `getBaseUrl()` for
|
||||
* retain / recall / reflect.
|
||||
*
|
||||
* The class is deliberately transparent about the daemon: new CLI flags or
|
||||
* environment variables never require a code change here — callers can pass
|
||||
* them via `env`, `extraProfileCreateArgs`, or `extraDaemonStartArgs`.
|
||||
*/
|
||||
export class HindsightServer {
|
||||
private readonly profile: string;
|
||||
private readonly port: number;
|
||||
private readonly host: string;
|
||||
private readonly baseUrl: string;
|
||||
private readonly embedVersion: string | undefined;
|
||||
private readonly embedPackagePath: string | undefined;
|
||||
private readonly userEnv: Record<string, string | undefined>;
|
||||
private readonly extraProfileCreateArgs: string[];
|
||||
private readonly extraDaemonStartArgs: string[];
|
||||
private readonly platformCpuWorkaround: boolean;
|
||||
private readonly readyTimeoutMs: number;
|
||||
private readonly readyPollIntervalMs: number;
|
||||
private readonly logger: Logger;
|
||||
|
||||
constructor(opts: HindsightServerOptions = {}) {
|
||||
this.profile = opts.profile ?? DEFAULT_PROFILE;
|
||||
this.port = opts.port ?? DEFAULT_PORT;
|
||||
this.host = opts.host ?? DEFAULT_HOST;
|
||||
this.baseUrl = `http://${this.host}:${this.port}`;
|
||||
this.embedVersion = opts.embedVersion;
|
||||
this.embedPackagePath = opts.embedPackagePath;
|
||||
this.userEnv = opts.env ?? {};
|
||||
this.extraProfileCreateArgs = opts.extraProfileCreateArgs ?? [];
|
||||
this.extraDaemonStartArgs = opts.extraDaemonStartArgs ?? [];
|
||||
this.platformCpuWorkaround = opts.platformCpuWorkaround ?? (process.platform === 'darwin');
|
||||
this.readyTimeoutMs = opts.readyTimeoutMs ?? DEFAULT_READY_TIMEOUT_MS;
|
||||
this.readyPollIntervalMs = opts.readyPollIntervalMs ?? DEFAULT_READY_POLL_INTERVAL_MS;
|
||||
this.logger = opts.logger ?? silentLogger;
|
||||
}
|
||||
|
||||
/** The base URL the daemon listens on (`http://host:port`). */
|
||||
getBaseUrl(): string {
|
||||
return this.baseUrl;
|
||||
}
|
||||
|
||||
/** The profile name this server operates on. */
|
||||
getProfile(): string {
|
||||
return this.profile;
|
||||
}
|
||||
|
||||
/**
|
||||
* Ensure the daemon is configured and running. Idempotent — the underlying
|
||||
* `profile create --merge` and `daemon start` commands tolerate re-runs.
|
||||
*/
|
||||
async start(): Promise<void> {
|
||||
this.logger.info(`[hindsight] starting daemon for profile "${this.profile}"`);
|
||||
|
||||
const env = this.buildEnv();
|
||||
await this.configureProfile(env);
|
||||
await this.startDaemon(env);
|
||||
await this.waitForReady();
|
||||
|
||||
this.logger.info(`[hindsight] daemon ready at ${this.baseUrl}`);
|
||||
}
|
||||
|
||||
/** Stop the daemon. Never throws — logs and resolves even on failure. */
|
||||
async stop(): Promise<void> {
|
||||
this.logger.info(`[hindsight] stopping daemon for profile "${this.profile}"`);
|
||||
|
||||
const [cmd, ...baseArgs] = getEmbedCommand({
|
||||
embedVersion: this.embedVersion,
|
||||
embedPackagePath: this.embedPackagePath,
|
||||
});
|
||||
const args = [...baseArgs, 'daemon', '--profile', this.profile, 'stop'];
|
||||
|
||||
const child = spawn(cmd, args, { stdio: 'pipe' });
|
||||
this.pipeOutput(child, 'daemon.stop');
|
||||
|
||||
await new Promise<void>((resolve) => {
|
||||
const timeout = setTimeout(() => {
|
||||
this.logger.warn(`[hindsight] daemon stop timed out after 5s`);
|
||||
resolve();
|
||||
}, 5_000);
|
||||
child.on('exit', () => {
|
||||
clearTimeout(timeout);
|
||||
this.logger.info(`[hindsight] daemon stopped`);
|
||||
resolve();
|
||||
});
|
||||
child.on('error', (err) => {
|
||||
clearTimeout(timeout);
|
||||
this.logger.warn(`[hindsight] error stopping daemon: ${err.message}`);
|
||||
resolve();
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
/** Probe `/health` once with a short timeout. */
|
||||
async checkHealth(): Promise<boolean> {
|
||||
try {
|
||||
const res = await fetch(`${this.baseUrl}/health`, {
|
||||
signal: AbortSignal.timeout(2_000),
|
||||
});
|
||||
return res.ok;
|
||||
} catch {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
// -------------------------------------------------------------------------
|
||||
// Internal
|
||||
// -------------------------------------------------------------------------
|
||||
|
||||
/**
|
||||
* Merge the process env, the caller-supplied `env`, and (on macOS) the
|
||||
* embeddings CPU workaround. Caller-supplied values always win over the
|
||||
* workaround; undefined values are dropped.
|
||||
*/
|
||||
private buildEnv(): NodeJS.ProcessEnv {
|
||||
const merged: NodeJS.ProcessEnv = { ...process.env };
|
||||
|
||||
if (this.platformCpuWorkaround && process.platform === 'darwin') {
|
||||
merged['HINDSIGHT_API_EMBEDDINGS_LOCAL_FORCE_CPU'] = '1';
|
||||
merged['HINDSIGHT_API_RERANKER_LOCAL_FORCE_CPU'] = '1';
|
||||
}
|
||||
|
||||
for (const [key, value] of Object.entries(this.userEnv)) {
|
||||
if (value !== undefined) {
|
||||
merged[key] = value;
|
||||
}
|
||||
}
|
||||
|
||||
return merged;
|
||||
}
|
||||
|
||||
/**
|
||||
* Run `profile create <name> --merge --port <port> [--env K=V ...]`.
|
||||
* Every entry in the merged env that was passed via {@link userEnv} (or
|
||||
* auto-applied by the CPU workaround) is forwarded as `--env`.
|
||||
*/
|
||||
private async configureProfile(env: NodeJS.ProcessEnv): Promise<void> {
|
||||
this.logger.info(`[hindsight] configuring profile "${this.profile}"`);
|
||||
|
||||
const [cmd, ...baseArgs] = getEmbedCommand({
|
||||
embedVersion: this.embedVersion,
|
||||
embedPackagePath: this.embedPackagePath,
|
||||
});
|
||||
const createArgs = [
|
||||
...baseArgs,
|
||||
'profile',
|
||||
'create',
|
||||
this.profile,
|
||||
'--merge',
|
||||
'--port',
|
||||
String(this.port),
|
||||
];
|
||||
|
||||
// Forward every env var that the caller intended for the daemon as --env.
|
||||
// We only forward keys the caller explicitly set (userEnv) plus the CPU
|
||||
// workaround values — not the entire process.env, to avoid leaking random
|
||||
// host state into profile config.
|
||||
const envForProfile = this.collectProfileEnv(env);
|
||||
for (const [key, value] of Object.entries(envForProfile)) {
|
||||
createArgs.push('--env', `${key}=${value}`);
|
||||
}
|
||||
|
||||
createArgs.push(...this.extraProfileCreateArgs);
|
||||
|
||||
await this.runCommand(cmd, createArgs, env, 'profile.create');
|
||||
}
|
||||
|
||||
/** Collect only the env vars that should be written into the profile file. */
|
||||
private collectProfileEnv(env: NodeJS.ProcessEnv): Record<string, string> {
|
||||
const out: Record<string, string> = {};
|
||||
|
||||
// 1. User-supplied env — always forwarded.
|
||||
for (const [key, value] of Object.entries(this.userEnv)) {
|
||||
if (value !== undefined) {
|
||||
out[key] = value;
|
||||
}
|
||||
}
|
||||
|
||||
// 2. CPU workaround — only if auto-applied and not already overridden.
|
||||
if (this.platformCpuWorkaround && process.platform === 'darwin') {
|
||||
const cpuKeys = [
|
||||
'HINDSIGHT_API_EMBEDDINGS_LOCAL_FORCE_CPU',
|
||||
'HINDSIGHT_API_RERANKER_LOCAL_FORCE_CPU',
|
||||
];
|
||||
for (const key of cpuKeys) {
|
||||
if (!(key in out) && env[key] !== undefined) {
|
||||
out[key] = env[key] as string;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return out;
|
||||
}
|
||||
|
||||
private async startDaemon(env: NodeJS.ProcessEnv): Promise<void> {
|
||||
const [cmd, ...baseArgs] = getEmbedCommand({
|
||||
embedVersion: this.embedVersion,
|
||||
embedPackagePath: this.embedPackagePath,
|
||||
});
|
||||
const args = [
|
||||
...baseArgs,
|
||||
'daemon',
|
||||
'--profile',
|
||||
this.profile,
|
||||
'start',
|
||||
...this.extraDaemonStartArgs,
|
||||
];
|
||||
|
||||
await this.runCommand(cmd, args, env, 'daemon.start');
|
||||
}
|
||||
|
||||
/**
|
||||
* Spawn `cmd` with `args`, pipe its output through the logger, and resolve
|
||||
* once it exits with code 0. Rejects on non-zero exit or spawn error.
|
||||
*/
|
||||
private async runCommand(
|
||||
cmd: string,
|
||||
args: string[],
|
||||
env: NodeJS.ProcessEnv,
|
||||
label: string,
|
||||
): Promise<void> {
|
||||
const child = spawn(cmd, args, { stdio: 'pipe', env });
|
||||
let output = '';
|
||||
child.stdout?.on('data', (data: Buffer) => {
|
||||
const text = data.toString();
|
||||
output += text;
|
||||
for (const line of text.trimEnd().split('\n')) {
|
||||
if (line) this.logger.info(`[hindsight:${label}] ${line}`);
|
||||
}
|
||||
});
|
||||
child.stderr?.on('data', (data: Buffer) => {
|
||||
const text = data.toString();
|
||||
output += text;
|
||||
for (const line of text.trimEnd().split('\n')) {
|
||||
if (line) this.logger.warn(`[hindsight:${label}] ${line}`);
|
||||
}
|
||||
});
|
||||
|
||||
await new Promise<void>((resolve, reject) => {
|
||||
child.on('exit', (code) => {
|
||||
if (code === 0) {
|
||||
resolve();
|
||||
} else {
|
||||
reject(new Error(`${label} failed with code ${code}: ${output.trim()}`));
|
||||
}
|
||||
});
|
||||
child.on('error', (err) => {
|
||||
reject(new Error(`${label} failed to spawn: ${err.message}`, { cause: err }));
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
/** Stream a spawned child's stdout/stderr through the logger without blocking. */
|
||||
private pipeOutput(child: ReturnType<typeof spawn>, label: string): void {
|
||||
child.stdout?.on('data', (data: Buffer) => {
|
||||
for (const line of data.toString().trimEnd().split('\n')) {
|
||||
if (line) this.logger.info(`[hindsight:${label}] ${line}`);
|
||||
}
|
||||
});
|
||||
child.stderr?.on('data', (data: Buffer) => {
|
||||
for (const line of data.toString().trimEnd().split('\n')) {
|
||||
if (line) this.logger.warn(`[hindsight:${label}] ${line}`);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
/** Poll `/health` until it succeeds or `readyTimeoutMs` elapses. */
|
||||
private async waitForReady(): Promise<void> {
|
||||
const deadline = Date.now() + this.readyTimeoutMs;
|
||||
let attempt = 0;
|
||||
while (Date.now() < deadline) {
|
||||
attempt++;
|
||||
try {
|
||||
const res = await fetch(`${this.baseUrl}/health`, {
|
||||
signal: AbortSignal.timeout(this.readyPollIntervalMs),
|
||||
});
|
||||
if (res.ok) {
|
||||
this.logger.debug(`[hindsight] health check passed (attempt ${attempt})`);
|
||||
return;
|
||||
}
|
||||
} catch {
|
||||
// expected while the daemon is still booting
|
||||
}
|
||||
await new Promise((resolve) => setTimeout(resolve, this.readyPollIntervalMs));
|
||||
}
|
||||
throw new Error(
|
||||
`Hindsight daemon did not become ready within ${this.readyTimeoutMs}ms at ${this.baseUrl}`,
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,54 @@
|
||||
import type { Logger } from './logger.js';
|
||||
|
||||
/**
|
||||
* Options for {@link HindsightServer}.
|
||||
*
|
||||
* The server is intentionally thin and pass-through: anything configurable
|
||||
* on the daemon side (env vars or CLI flags) can be set here without needing
|
||||
* a new dedicated option. Use {@link env} for `HINDSIGHT_*` / `OPENAI_API_KEY` /
|
||||
* custom provider settings, and the two `extra*` arrays to append raw CLI
|
||||
* args to `profile create` or `daemon start`.
|
||||
*
|
||||
* For talking to the daemon after `start()`, use `@vectorize-io/hindsight-client`
|
||||
* against `server.getBaseUrl()`. This package does not ship its own HTTP
|
||||
* client.
|
||||
*/
|
||||
export interface HindsightServerOptions {
|
||||
/** Profile name used for `--profile <name>` on every sub-command. Default: `"default"`. */
|
||||
profile?: string;
|
||||
/** TCP port the daemon listens on. Default: `8888`. */
|
||||
port?: number;
|
||||
/** Hostname the daemon binds to (for health checks). Default: `127.0.0.1`. */
|
||||
host?: string;
|
||||
/** Version of the underlying `hindsight-embed` PyPI package to run via `uvx`. Default: `"latest"`. */
|
||||
embedVersion?: string;
|
||||
/** Local path to a `hindsight-embed` checkout — takes precedence over `embedVersion`. */
|
||||
embedPackagePath?: string;
|
||||
/**
|
||||
* Environment variables passed to the daemon process AND written into the
|
||||
* profile via repeated `--env KEY=VALUE` flags. This is the preferred way
|
||||
* to surface any `HINDSIGHT_API_*` / `HINDSIGHT_EMBED_*` setting — adding a
|
||||
* new daemon env var never requires a wrapper update.
|
||||
*
|
||||
* Values of `undefined` are dropped (so you can spread conditionally).
|
||||
*/
|
||||
env?: Record<string, string | undefined>;
|
||||
/** Extra args appended verbatim to `hindsight-embed profile create <name> --merge ...`. */
|
||||
extraProfileCreateArgs?: string[];
|
||||
/** Extra args appended verbatim to `hindsight-embed daemon --profile <name> start ...`. */
|
||||
extraDaemonStartArgs?: string[];
|
||||
/**
|
||||
* On macOS, automatically set
|
||||
* `HINDSIGHT_API_EMBEDDINGS_LOCAL_FORCE_CPU=1` and
|
||||
* `HINDSIGHT_API_RERANKER_LOCAL_FORCE_CPU=1` to avoid Metal/MPS crashes in
|
||||
* daemon mode. Default: `true` on `darwin`, ignored elsewhere. Any value set
|
||||
* explicitly in {@link env} wins over the auto-applied value.
|
||||
*/
|
||||
platformCpuWorkaround?: boolean;
|
||||
/** Max time (ms) to wait for `/health` to return 200. Default: `30_000`. */
|
||||
readyTimeoutMs?: number;
|
||||
/** Polling interval (ms) while waiting for `/health`. Default: `1_000`. */
|
||||
readyPollIntervalMs?: number;
|
||||
/** Optional pluggable logger. Default: silent. */
|
||||
logger?: Logger;
|
||||
}
|
||||
@@ -0,0 +1,18 @@
|
||||
{
|
||||
"compilerOptions": {
|
||||
"target": "ES2022",
|
||||
"module": "ES2022",
|
||||
"lib": ["ES2022"],
|
||||
"moduleResolution": "node",
|
||||
"declaration": true,
|
||||
"outDir": "./dist",
|
||||
"rootDir": "./src",
|
||||
"strict": true,
|
||||
"esModuleInterop": true,
|
||||
"skipLibCheck": true,
|
||||
"forceConsistentCasingInFileNames": true,
|
||||
"resolveJsonModule": true
|
||||
},
|
||||
"include": ["src/**/*"],
|
||||
"exclude": ["node_modules", "dist", "src/**/*.test.ts"]
|
||||
}
|
||||
@@ -0,0 +1,11 @@
|
||||
import { defineConfig } from 'tsup';
|
||||
|
||||
export default defineConfig({
|
||||
entry: ['src/index.ts'],
|
||||
format: ['esm'],
|
||||
dts: true,
|
||||
outDir: 'dist',
|
||||
clean: true,
|
||||
sourcemap: true,
|
||||
bundle: true,
|
||||
});
|
||||
@@ -0,0 +1,8 @@
|
||||
import { defineConfig } from 'vitest/config';
|
||||
|
||||
export default defineConfig({
|
||||
test: {
|
||||
include: ['src/**/*.test.ts'],
|
||||
environment: 'node',
|
||||
},
|
||||
});
|
||||
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "hindsight-all-slim"
|
||||
version = "0.4.22"
|
||||
version = "0.5.4"
|
||||
description = "Hindsight: Agent Memory That Works Like Human Memory - Slim All-in-One Bundle"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.11"
|
||||
|
||||
@@ -71,7 +71,7 @@ class HindsightEmbedded:
|
||||
llm_model: Model name to use
|
||||
llm_base_url: Optional custom base URL for LLM API
|
||||
database_url: Optional database URL override (default: profile-specific pg0)
|
||||
idle_timeout: Seconds before daemon auto-exits when idle (default: 300)
|
||||
idle_timeout: Seconds before daemon auto-exits when idle (default: 0, disabled)
|
||||
log_level: Daemon log level (default: "info")
|
||||
ui: Whether to start the control plane web UI alongside the daemon (default: False)
|
||||
ui_port: Port for the UI. Defaults to daemon_port + 10000.
|
||||
@@ -86,7 +86,7 @@ class HindsightEmbedded:
|
||||
llm_model: str = "openai/gpt-oss-120b",
|
||||
llm_base_url: Optional[str] = None,
|
||||
database_url: Optional[str] = None,
|
||||
idle_timeout: int = 300,
|
||||
idle_timeout: int = 0,
|
||||
log_level: str = "info",
|
||||
ui: bool = False,
|
||||
ui_port: Optional[int] = None,
|
||||
@@ -102,7 +102,7 @@ class HindsightEmbedded:
|
||||
llm_model: Model name to use
|
||||
llm_base_url: Optional custom base URL for LLM API
|
||||
database_url: Optional database URL override
|
||||
idle_timeout: Seconds before daemon auto-exits when idle
|
||||
idle_timeout: Seconds before daemon auto-exits when idle (0 = disabled)
|
||||
log_level: Daemon log level
|
||||
ui: Whether to start the control plane web UI alongside the daemon
|
||||
ui_port: Port for the UI (defaults to daemon_port + 10000)
|
||||
@@ -142,14 +142,37 @@ class HindsightEmbedded:
|
||||
self._memories_api: Optional[MemoriesAPI] = None
|
||||
|
||||
def _ensure_started(self):
|
||||
"""Ensure daemon is running (thread-safe)."""
|
||||
"""Ensure daemon is running (thread-safe), restarting if crashed."""
|
||||
if self._started and self._client is not None:
|
||||
return
|
||||
if self._manager.is_running(self.profile):
|
||||
return
|
||||
# Daemon crashed — reset state and fall through to restart
|
||||
logger.warning(
|
||||
"Daemon for profile '%s' is no longer responsive, restarting...",
|
||||
self.profile,
|
||||
)
|
||||
try:
|
||||
self._client.close()
|
||||
except Exception:
|
||||
logger.debug("Error closing stale client", exc_info=True)
|
||||
self._client = None
|
||||
self._started = False
|
||||
|
||||
with self._lock:
|
||||
# Double-check after acquiring lock
|
||||
if self._started and self._client is not None:
|
||||
return
|
||||
if self._manager.is_running(self.profile):
|
||||
return
|
||||
logger.warning(
|
||||
"Daemon for profile '%s' is no longer responsive (lock path), restarting...",
|
||||
self.profile,
|
||||
)
|
||||
try:
|
||||
self._client.close()
|
||||
except Exception:
|
||||
logger.debug("Error closing stale client", exc_info=True)
|
||||
self._client = None
|
||||
self._started = False
|
||||
|
||||
if self._closed:
|
||||
raise RuntimeError(
|
||||
@@ -190,12 +213,32 @@ class HindsightEmbedded:
|
||||
if self._closed:
|
||||
return
|
||||
|
||||
with self._lock:
|
||||
acquired = self._lock.acquire(timeout=5.0)
|
||||
if not acquired:
|
||||
# Lock is held by another thread (e.g. _ensure_started).
|
||||
# Mark closed to prevent new operations but skip shared-state
|
||||
# teardown — the daemon's idle timeout handles the rest.
|
||||
logger.warning(
|
||||
"Cleanup lock acquisition timed out for profile '%s'; "
|
||||
"marking closed, daemon will idle-stop on its own",
|
||||
self.profile,
|
||||
)
|
||||
self._closed = True
|
||||
return
|
||||
|
||||
try:
|
||||
if self._closed:
|
||||
return
|
||||
|
||||
if self._client is not None:
|
||||
self._client.close()
|
||||
try:
|
||||
self._client.close()
|
||||
except Exception:
|
||||
logger.debug(
|
||||
"Error closing client for profile '%s'",
|
||||
self.profile,
|
||||
exc_info=True,
|
||||
)
|
||||
self._client = None
|
||||
|
||||
# Stop UI if it was started
|
||||
@@ -209,6 +252,8 @@ class HindsightEmbedded:
|
||||
self._manager.stop(self.profile)
|
||||
|
||||
self._closed = True
|
||||
finally:
|
||||
self._lock.release()
|
||||
|
||||
def close(self, stop_daemon: bool = False):
|
||||
"""
|
||||
@@ -231,23 +276,10 @@ class HindsightEmbedded:
|
||||
This allows HindsightEmbedded to expose all HindsightClient methods
|
||||
without manually wrapping each one.
|
||||
"""
|
||||
# Ensure server is started before proxying
|
||||
# Ensure server is started (and restart if crashed) before proxying
|
||||
self._ensure_started()
|
||||
|
||||
# Get the attribute from the underlying client
|
||||
attr = getattr(self._client, name)
|
||||
|
||||
# If it's a callable, wrap it to ensure server is started
|
||||
# (shouldn't be needed since _ensure_started already called, but defensive)
|
||||
if callable(attr):
|
||||
|
||||
def wrapper(*args, **kwargs):
|
||||
self._ensure_started()
|
||||
return attr(*args, **kwargs)
|
||||
|
||||
return wrapper
|
||||
|
||||
return attr
|
||||
return getattr(self._client, name)
|
||||
|
||||
def __enter__(self):
|
||||
"""Context manager entry - ensures server is started."""
|
||||
@@ -372,11 +404,8 @@ class HindsightEmbedded:
|
||||
"""
|
||||
Get the underlying Hindsight client for direct access.
|
||||
|
||||
WARNING: Using this property directly means daemon restarts won't be
|
||||
handled automatically. Prefer using the API namespaces (banks, mental_models,
|
||||
directives, memories) or direct method calls on HindsightEmbedded instead.
|
||||
|
||||
Ensures daemon is started before returning the client.
|
||||
Ensures daemon is started (and restarts it if it has crashed) before
|
||||
returning the client.
|
||||
|
||||
Returns:
|
||||
Hindsight: The underlying client instance
|
||||
@@ -387,9 +416,8 @@ class HindsightEmbedded:
|
||||
|
||||
embedded = HindsightEmbedded(profile="myapp", ...)
|
||||
|
||||
# Direct access (not recommended - daemon crashes won't be handled)
|
||||
client = embedded.client
|
||||
banks = client.list_banks() # If daemon crashes, this will fail
|
||||
banks = client.list_banks()
|
||||
```
|
||||
"""
|
||||
self._ensure_started()
|
||||
@@ -403,8 +431,13 @@ class HindsightEmbedded:
|
||||
|
||||
@property
|
||||
def is_running(self) -> bool:
|
||||
"""Check if the client is initialized."""
|
||||
return self._started and not self._closed and self._client is not None
|
||||
"""Check if the client is initialized and the daemon is responsive."""
|
||||
return (
|
||||
self._started
|
||||
and not self._closed
|
||||
and self._client is not None
|
||||
and self._manager.is_running(self.profile)
|
||||
)
|
||||
|
||||
@property
|
||||
def ui_url(self) -> str:
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "hindsight-all"
|
||||
version = "0.4.22"
|
||||
version = "0.5.4"
|
||||
description = "Hindsight: Agent Memory That Works Like Human Memory - All-in-One Bundle"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.11"
|
||||
@@ -20,6 +20,9 @@ hindsight-client = { workspace = true }
|
||||
hindsight-embed = { workspace = true }
|
||||
|
||||
[project.optional-dependencies]
|
||||
local-llm = [
|
||||
"hindsight-api-slim[local-llm]>=0.4.17",
|
||||
]
|
||||
test = [
|
||||
"pytest>=7.0.0",
|
||||
"pytest-asyncio>=0.21.0",
|
||||
|
||||
@@ -0,0 +1,56 @@
|
||||
"""
|
||||
Unit test for _cleanup lock timeout behavior.
|
||||
|
||||
Verifies that _cleanup completes even when the lock is held by another thread,
|
||||
instead of hanging indefinitely (fixes #952).
|
||||
"""
|
||||
|
||||
import threading
|
||||
import time
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
def test_cleanup_completes_when_lock_held():
|
||||
"""
|
||||
_cleanup should complete (best-effort) even when self._lock is held
|
||||
by another thread, e.g. during a long _ensure_started call.
|
||||
"""
|
||||
with patch.dict("sys.modules", {
|
||||
"hindsight_client": MagicMock(),
|
||||
"hindsight_embed": MagicMock(),
|
||||
"hindsight.api_namespaces": MagicMock(),
|
||||
}):
|
||||
from hindsight.embedded import HindsightEmbedded
|
||||
|
||||
client = HindsightEmbedded.__new__(HindsightEmbedded)
|
||||
client.profile = "test"
|
||||
client._lock = threading.Lock()
|
||||
client._closed = False
|
||||
client._client = None
|
||||
client._started = False
|
||||
client._ui = False
|
||||
|
||||
# Simulate another thread holding the lock
|
||||
client._lock.acquire()
|
||||
|
||||
cleanup_done = threading.Event()
|
||||
|
||||
def run_cleanup():
|
||||
client._cleanup()
|
||||
cleanup_done.set()
|
||||
|
||||
t = threading.Thread(target=run_cleanup)
|
||||
t.start()
|
||||
|
||||
# Cleanup should complete within the timeout (5s) + margin
|
||||
assert cleanup_done.wait(timeout=8.0), (
|
||||
"_cleanup hung instead of timing out on lock acquisition"
|
||||
)
|
||||
|
||||
# Release the lock from the simulating thread
|
||||
client._lock.release()
|
||||
t.join(timeout=1.0)
|
||||
|
||||
assert client._closed, "Client should be marked as closed after cleanup"
|
||||
@@ -401,3 +401,42 @@ def test_embedded_ui_flag(llm_config):
|
||||
|
||||
finally:
|
||||
client.close()
|
||||
|
||||
|
||||
def test_embedded_daemon_crash_recovery(llm_config):
|
||||
"""
|
||||
Test that HindsightEmbedded recovers when the daemon crashes.
|
||||
|
||||
Simulates a crash by stopping the daemon, then verifies
|
||||
that the next operation transparently restarts it.
|
||||
"""
|
||||
profile = f"test_crash_{uuid.uuid4().hex[:8]}"
|
||||
bank_id = f"bank_{uuid.uuid4().hex[:8]}"
|
||||
|
||||
client = HindsightEmbedded(profile=profile, log_level="info", **llm_config)
|
||||
|
||||
try:
|
||||
# Start daemon and store a memory
|
||||
result = client.retain(bank_id=bank_id, content="Before crash")
|
||||
assert result.success, "Initial retain should succeed"
|
||||
assert client.is_running, "Daemon should be running"
|
||||
|
||||
original_url = client.url
|
||||
|
||||
# Simulate daemon crash by stopping it
|
||||
client._manager.stop(client.profile)
|
||||
assert not client._manager.is_running(client.profile), (
|
||||
"Daemon should be stopped after simulated crash"
|
||||
)
|
||||
|
||||
# Next operation should transparently restart the daemon
|
||||
result2 = client.retain(bank_id=bank_id, content="After crash recovery")
|
||||
assert result2.success, "Retain after crash recovery should succeed"
|
||||
assert client.is_running, "Daemon should be running again after recovery"
|
||||
|
||||
# Verify recall still works
|
||||
recall_result = client.recall(bank_id=bank_id, query="crash")
|
||||
assert isinstance(recall_result.results, list), "Recall should return results"
|
||||
|
||||
finally:
|
||||
client.close()
|
||||
|
||||
@@ -46,4 +46,4 @@ __all__ = [
|
||||
"RemoteTEICrossEncoder",
|
||||
"LLMConfig",
|
||||
]
|
||||
__version__ = "0.4.22"
|
||||
__version__ = "0.5.4"
|
||||
|
||||
@@ -375,6 +375,140 @@ def decommission_worker(
|
||||
typer.echo(f"No tasks found for worker '{worker_id}'")
|
||||
|
||||
|
||||
async def _decommission_all_workers(db_url: str, schema: str = "public") -> list[dict[str, Any]]:
|
||||
"""Release all processing tasks from all workers, setting them back to pending status."""
|
||||
is_pg0, instance_name, _ = parse_pg0_url(db_url)
|
||||
if is_pg0:
|
||||
typer.echo(f"Starting embedded PostgreSQL (instance: {instance_name})...")
|
||||
resolved_url = await resolve_database_url(db_url)
|
||||
|
||||
conn = await asyncpg.connect(resolved_url)
|
||||
try:
|
||||
table = _fq_table("async_operations", schema)
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
UPDATE {table}
|
||||
SET status = 'pending', worker_id = NULL, claimed_at = NULL, updated_at = now()
|
||||
WHERE status = 'processing'
|
||||
RETURNING operation_id, worker_id, operation_type
|
||||
""",
|
||||
)
|
||||
return [dict(r) for r in rows]
|
||||
finally:
|
||||
await conn.close()
|
||||
|
||||
|
||||
@app.command(name="decommission-workers")
|
||||
def decommission_workers(
|
||||
schema: str = typer.Option("public", "--schema", "-s", help="Database schema"),
|
||||
yes: bool = typer.Option(False, "--yes", "-y", help="Skip confirmation prompt"),
|
||||
):
|
||||
"""Release all processing tasks from all workers (sets status back to pending).
|
||||
|
||||
Use this command to recover from situations where one or more workers have crashed
|
||||
or been removed without graceful shutdown. All tasks currently in 'processing' status
|
||||
will be released back to the queue regardless of which worker owns them.
|
||||
"""
|
||||
config = HindsightConfig.from_env()
|
||||
|
||||
if not config.database_url:
|
||||
typer.echo("Error: Database URL not configured.", err=True)
|
||||
typer.echo("Set HINDSIGHT_API_DATABASE_URL environment variable.", err=True)
|
||||
raise typer.Exit(1)
|
||||
|
||||
if not yes:
|
||||
typer.confirm(
|
||||
"This will release ALL processing tasks from ALL workers back to pending. Continue?",
|
||||
abort=True,
|
||||
)
|
||||
|
||||
typer.echo(f"Decommissioning all workers (schema: {schema})...")
|
||||
|
||||
released = asyncio.run(_decommission_all_workers(config.database_url, schema))
|
||||
|
||||
if released:
|
||||
# Group by worker_id for summary
|
||||
by_worker: dict[str, int] = {}
|
||||
for row in released:
|
||||
wid = row["worker_id"] or "unknown"
|
||||
by_worker[wid] = by_worker.get(wid, 0) + 1
|
||||
|
||||
typer.echo(f"Released {len(released)} task(s):")
|
||||
for wid, count in by_worker.items():
|
||||
typer.echo(f" {wid}: {count} task(s)")
|
||||
else:
|
||||
typer.echo("No processing tasks found")
|
||||
|
||||
|
||||
async def _worker_status(db_url: str, schema: str = "public") -> list[dict[str, Any]]:
|
||||
"""Get all processing tasks grouped by worker with their last update time."""
|
||||
is_pg0, instance_name, _ = parse_pg0_url(db_url)
|
||||
if is_pg0:
|
||||
typer.echo(f"Starting embedded PostgreSQL (instance: {instance_name})...")
|
||||
resolved_url = await resolve_database_url(db_url)
|
||||
|
||||
conn = await asyncpg.connect(resolved_url)
|
||||
try:
|
||||
table = _fq_table("async_operations", schema)
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT worker_id, operation_id, operation_type, bank_id,
|
||||
claimed_at, updated_at,
|
||||
now() - claimed_at AS running_for,
|
||||
now() - updated_at AS last_update_ago
|
||||
FROM {table}
|
||||
WHERE status = 'processing'
|
||||
ORDER BY worker_id, claimed_at
|
||||
""",
|
||||
)
|
||||
return [dict(r) for r in rows]
|
||||
finally:
|
||||
await conn.close()
|
||||
|
||||
|
||||
@app.command(name="worker-status")
|
||||
def worker_status(
|
||||
schema: str = typer.Option("public", "--schema", "-s", help="Database schema"),
|
||||
):
|
||||
"""Show all currently processing tasks grouped by worker.
|
||||
|
||||
Displays each worker's active tasks with operation type, bank, how long
|
||||
the task has been running, and when it was last updated. Useful for
|
||||
identifying dead workers with orphaned tasks.
|
||||
"""
|
||||
config = HindsightConfig.from_env()
|
||||
|
||||
if not config.database_url:
|
||||
typer.echo("Error: Database URL not configured.", err=True)
|
||||
typer.echo("Set HINDSIGHT_API_DATABASE_URL environment variable.", err=True)
|
||||
raise typer.Exit(1)
|
||||
|
||||
rows = asyncio.run(_worker_status(config.database_url, schema))
|
||||
|
||||
if not rows:
|
||||
typer.echo("No processing tasks found")
|
||||
return
|
||||
|
||||
# Group by worker_id
|
||||
by_worker: dict[str, list[dict[str, Any]]] = {}
|
||||
for row in rows:
|
||||
wid = row["worker_id"] or "unknown"
|
||||
by_worker.setdefault(wid, []).append(row)
|
||||
|
||||
typer.echo(f"Processing tasks across {len(by_worker)} worker(s):\n")
|
||||
for wid, tasks in by_worker.items():
|
||||
typer.echo(f"Worker: {wid} ({len(tasks)} task(s))")
|
||||
for task in tasks:
|
||||
op_id = str(task["operation_id"])[:8]
|
||||
running_for = task["running_for"]
|
||||
last_update = task["last_update_ago"]
|
||||
typer.echo(
|
||||
f" {op_id} {task['operation_type']:<20s} bank={task['bank_id']}"
|
||||
f" running={running_for} last_update={last_update} ago"
|
||||
)
|
||||
typer.echo("")
|
||||
|
||||
|
||||
def main():
|
||||
app()
|
||||
|
||||
|
||||
@@ -12,6 +12,7 @@ from dotenv import load_dotenv
|
||||
from sqlalchemy import engine_from_config, pool
|
||||
|
||||
# Import your models here
|
||||
from hindsight_api.db_url import to_libpq_url
|
||||
from hindsight_api.models import Base
|
||||
|
||||
|
||||
@@ -65,11 +66,11 @@ def get_database_url() -> str:
|
||||
"Set HINDSIGHT_API_DATABASE_URL environment variable or pass database_url to run_migrations()."
|
||||
)
|
||||
|
||||
# For migrations, use psycopg2 (sync driver) to avoid pgbouncer prepared statement issues
|
||||
if database_url.startswith("postgresql+asyncpg://"):
|
||||
database_url = database_url.replace("postgresql+asyncpg://", "postgresql://", 1)
|
||||
elif database_url.startswith("postgres+asyncpg://"):
|
||||
database_url = database_url.replace("postgres+asyncpg://", "postgresql://", 1)
|
||||
# For migrations, use the sync psycopg2 driver (avoids pgbouncer prepared
|
||||
# statement issues and is required since create_engine is the sync API).
|
||||
# Also translates ?ssl=require (SQLAlchemy asyncpg style) to ?sslmode=require
|
||||
# (libpq style) for external-PostgreSQL deployments.
|
||||
database_url = to_libpq_url(database_url)
|
||||
|
||||
# Update config with processed URL for engine_from_config to use
|
||||
config.set_main_option("sqlalchemy.url", database_url)
|
||||
|
||||
+4
-4
@@ -4,8 +4,8 @@ The previous GIN trigram index on canonical_name was case-sensitive, causing
|
||||
"Alice" and "alice" to have different trigram sets. This recreates it on
|
||||
LOWER(canonical_name) so the % operator matches case-insensitively.
|
||||
|
||||
Revision ID: d6e7f8a9b0c1
|
||||
Revises: c5d6e7f8a9b0
|
||||
Revision ID: 2eee35aa3cfc
|
||||
Revises: d6e7f8a9b0c1
|
||||
Create Date: 2026-03-31
|
||||
"""
|
||||
|
||||
@@ -13,8 +13,8 @@ from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
revision: str = "d6e7f8a9b0c1"
|
||||
down_revision: str | Sequence[str] | None = "c5d6e7f8a9b0"
|
||||
revision: str = "2eee35aa3cfc"
|
||||
down_revision: str | Sequence[str] | None = "d6e7f8a9b0c1"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
+40
@@ -0,0 +1,40 @@
|
||||
"""Merge divergent migration heads for v0.5.3
|
||||
|
||||
v0.5.3 shipped with two migration heads that were never unified:
|
||||
|
||||
* ``c4x5y6z7a8b9`` — delta-refresh chain
|
||||
(``add_last_refreshed_source_query`` ->
|
||||
``add_structured_content_to_mental_models`` ->
|
||||
``backsweep_orphan_observations_v2``)
|
||||
|
||||
* ``h3i4j5k6l7m8`` — per-bank vector indexes / audit log chain
|
||||
(the ``merge_heads_and_add_unit_entities_index`` subtree)
|
||||
|
||||
Both fork from ``z1u2v3w4x5y6``. Upgrades from v0.5.2 still succeed — the
|
||||
walker applies the three c4x5 revisions and leaves the database stamped at
|
||||
both heads — but the result is a split DAG: ``alembic upgrade head``
|
||||
(singular) is ambiguous, and any future migration has to pick one head as
|
||||
its parent, orphaning the other.
|
||||
|
||||
This revision linearises the DAG into a single head. It has no schema
|
||||
effect.
|
||||
|
||||
Revision ID: 8c6fa6f7230b
|
||||
Revises: c4x5y6z7a8b9, h3i4j5k6l7m8
|
||||
Create Date: 2026-04-18
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
revision: str = "8c6fa6f7230b"
|
||||
down_revision: str | Sequence[str] | None = ("c4x5y6z7a8b9", "h3i4j5k6l7m8")
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
pass
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
pass
|
||||
+38
@@ -0,0 +1,38 @@
|
||||
"""Add last_refreshed_source_query column to mental_models
|
||||
|
||||
Revision ID: a2v3w4x5y6z7
|
||||
Revises: z1u2v3w4x5y6
|
||||
Create Date: 2026-04-15
|
||||
|
||||
Tracks the source_query that was used during the most recent refresh.
|
||||
Used by delta-mode refresh to detect when the query has changed: if it has,
|
||||
delta mode falls back to a full regeneration because the surgical-edit
|
||||
assumption (same topic, new facts) no longer holds.
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
revision: str = "a2v3w4x5y6z7"
|
||||
down_revision: str | Sequence[str] | None = "z1u2v3w4x5y6"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _get_schema_prefix() -> str:
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
op.execute(f"""
|
||||
ALTER TABLE {schema}mental_models
|
||||
ADD COLUMN IF NOT EXISTS last_refreshed_source_query TEXT
|
||||
""")
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
op.execute(f"ALTER TABLE {schema}mental_models DROP COLUMN IF EXISTS last_refreshed_source_query")
|
||||
+2
-2
@@ -1,7 +1,7 @@
|
||||
"""Fix per-bank vector indexes to match configured extension
|
||||
|
||||
Revision ID: a4b5c6d7e8f9
|
||||
Revises: d6e7f8a9b0c1
|
||||
Revises: 2eee35aa3cfc
|
||||
Create Date: 2026-04-01
|
||||
|
||||
Migration d5e6f7a8b9c0 hardcoded HNSW when creating per-bank partial vector
|
||||
@@ -21,7 +21,7 @@ from alembic import context, op
|
||||
from sqlalchemy import text
|
||||
|
||||
revision: str = "a4b5c6d7e8f9"
|
||||
down_revision: str | Sequence[str] | None = "d6e7f8a9b0c1"
|
||||
down_revision: str | Sequence[str] | None = "2eee35aa3cfc"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
+44
@@ -0,0 +1,44 @@
|
||||
"""Add structured_content JSONB column to mental_models
|
||||
|
||||
Revision ID: b3w4x5y6z7a8
|
||||
Revises: a2v3w4x5y6z7
|
||||
Create Date: 2026-04-16
|
||||
|
||||
Stores the structured representation of a mental model document (sections,
|
||||
blocks). The plain ``content`` column remains the rendered markdown shown to
|
||||
users. ``structured_content`` is the source of truth for delta-mode refreshes:
|
||||
each refresh applies a list of typed operations to the structured doc, then
|
||||
re-renders to markdown — so unchanged sections come through byte-identical
|
||||
without an LLM round-trip.
|
||||
|
||||
Nullable: existing markdown-only mental models continue to work in full mode;
|
||||
the column is populated lazily the first time a model is refreshed in delta
|
||||
mode.
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
revision: str = "b3w4x5y6z7a8"
|
||||
down_revision: str | Sequence[str] | None = "a2v3w4x5y6z7"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _get_schema_prefix() -> str:
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
op.execute(f"""
|
||||
ALTER TABLE {schema}mental_models
|
||||
ADD COLUMN IF NOT EXISTS structured_content JSONB
|
||||
""")
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
op.execute(f"ALTER TABLE {schema}mental_models DROP COLUMN IF EXISTS structured_content")
|
||||
+66
@@ -0,0 +1,66 @@
|
||||
"""backsweep_orphan_observations_v2
|
||||
|
||||
Re-run of Pass 2 from migration ``g7h8i9j0k1l2_backsweep_orphan_observations``
|
||||
to sweep observations that became orphaned between then and now.
|
||||
|
||||
Why we need it again:
|
||||
``fact_storage.handle_document_tracking`` (the retain/upsert path) deleted
|
||||
the existing document via the FK cascade — which removes the source
|
||||
``memory_units`` — but never invalidated the observations derived from
|
||||
them. Only the explicit ``MemoryEngine.delete_document`` API called
|
||||
``_delete_stale_observations_for_memories``. Every document re-ingest
|
||||
therefore left orphan observations whose ``source_memory_ids`` arrays
|
||||
pointed at IDs that no longer existed in ``memory_units``.
|
||||
|
||||
``handle_document_tracking`` now calls the same cleanup helper before the
|
||||
cascade, so no new orphans will accumulate going forward. This migration
|
||||
cleans up the historical residue.
|
||||
|
||||
Identical to Pass 2 of g7h8i9j0k1l2. Pass 1 (memory_units whose bank is
|
||||
gone) is intentionally not re-run; that scenario has no fresh source.
|
||||
|
||||
Revision ID: c4x5y6z7a8b9
|
||||
Revises: b3w4x5y6z7a8
|
||||
Create Date: 2026-04-16
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
revision: str = "c4x5y6z7a8b9"
|
||||
down_revision: str | Sequence[str] | None = "b3w4x5y6z7a8"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _get_schema_prefix() -> str:
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
mu = f"{schema}memory_units"
|
||||
|
||||
# Delete observations whose every source_memory_id refers to a now-deleted
|
||||
# memory_unit (or the array is empty). Observations with at least one
|
||||
# surviving source are left alone — the consolidation engine will refresh
|
||||
# their text on the next pass.
|
||||
op.execute(
|
||||
f"""
|
||||
DELETE FROM {mu} orphan
|
||||
WHERE orphan.fact_type = 'observation'
|
||||
AND NOT EXISTS (
|
||||
SELECT 1
|
||||
FROM {mu} src
|
||||
WHERE src.id = ANY(orphan.source_memory_ids)
|
||||
AND src.bank_id = orphan.bank_id
|
||||
)
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
# Deleted rows cannot be restored.
|
||||
pass
|
||||
+39
@@ -0,0 +1,39 @@
|
||||
"""Drop unused metadata column from documents table
|
||||
|
||||
Revision ID: d6e7f8a9b0c1
|
||||
Revises: c2d3e4f5g6h7, c5d6e7f8a9b0
|
||||
Create Date: 2026-03-30
|
||||
|
||||
The metadata column on documents was always stored as an empty dict {}.
|
||||
Actual document metadata is stored inside retain_params.metadata.
|
||||
|
||||
This migration was originally shipped in v0.4.22, then its file was deleted
|
||||
in v0.5.0 (and its revision ID accidentally reused by 2eee35aa3cfc).
|
||||
Restoring the file so that databases stamped at this revision can upgrade
|
||||
cleanly to v0.5.x+.
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
revision: str = "d6e7f8a9b0c1"
|
||||
down_revision: str | Sequence[str] | None = ("c2d3e4f5g6h7", "c5d6e7f8a9b0")
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _get_schema_prefix() -> str:
|
||||
"""Get schema prefix for table names (required for multi-tenant support)."""
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
op.execute(f"ALTER TABLE {schema}documents DROP COLUMN IF EXISTS metadata")
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
op.execute(f"ALTER TABLE {schema}documents ADD COLUMN IF NOT EXISTS metadata jsonb DEFAULT '{{}}'")
|
||||
+42
@@ -0,0 +1,42 @@
|
||||
"""Merge 3 migration heads and add unit_entities composite index
|
||||
|
||||
Revision ID: h3i4j5k6l7m8
|
||||
Revises: a4b5c6d7e8f9, g2h3i4j5k6l7
|
||||
Create Date: 2026-04-07
|
||||
|
||||
Merges three unmerged migration heads into one, and adds a composite index
|
||||
(entity_id, unit_id) on unit_entities for index-only scans in the LATERAL
|
||||
entity expansion query.
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
revision: str = "h3i4j5k6l7m8"
|
||||
down_revision: str | Sequence[str] | None = ("a4b5c6d7e8f9", "g2h3i4j5k6l7")
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _get_schema_prefix() -> str:
|
||||
"""Get schema prefix for table names (required for multi-tenant support)."""
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
# Composite index enables index-only scans for entity_id -> unit_id lookups
|
||||
op.execute(
|
||||
f"CREATE INDEX IF NOT EXISTS idx_unit_entities_entity_unit ON {schema}unit_entities (entity_id, unit_id)"
|
||||
)
|
||||
# Drop the now-redundant single-column index
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema}idx_unit_entities_entity")
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema}idx_unit_entities_entity_unit")
|
||||
# Restore the single-column index
|
||||
op.execute(f"CREATE INDEX IF NOT EXISTS idx_unit_entities_entity ON {schema}unit_entities (entity_id)")
|
||||
@@ -294,6 +294,44 @@ class EntityListResponse(BaseModel):
|
||||
offset: int
|
||||
|
||||
|
||||
class EntityGraphResponse(BaseModel):
|
||||
"""Response model for entity co-occurrence graph endpoint."""
|
||||
|
||||
model_config = ConfigDict(
|
||||
json_schema_extra={
|
||||
"example": {
|
||||
"nodes": [
|
||||
{"data": {"id": "uuid-1", "label": "Alice", "mentionCount": 12, "color": "#42a5f5"}},
|
||||
{"data": {"id": "uuid-2", "label": "Google", "mentionCount": 8, "color": "#42a5f5"}},
|
||||
],
|
||||
"edges": [
|
||||
{
|
||||
"data": {
|
||||
"id": "uuid-1-uuid-2",
|
||||
"source": "uuid-1",
|
||||
"target": "uuid-2",
|
||||
"linkType": "cooccurrence",
|
||||
"weight": 5,
|
||||
"color": "#ffd700",
|
||||
"lineStyle": "solid",
|
||||
"lastCooccurred": "2024-02-01T14:00:00Z",
|
||||
}
|
||||
}
|
||||
],
|
||||
"total_entities": 2,
|
||||
"total_edges": 1,
|
||||
"limit": 1000,
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
nodes: list[dict[str, Any]]
|
||||
edges: list[dict[str, Any]]
|
||||
total_entities: int
|
||||
total_edges: int
|
||||
limit: int
|
||||
|
||||
|
||||
class EntityDetailResponse(BaseModel):
|
||||
"""Response model for entity detail endpoint."""
|
||||
|
||||
@@ -463,6 +501,12 @@ class MemoryItem(BaseModel):
|
||||
description="Named retain strategy for this item. Overrides the bank's default strategy for this item only. "
|
||||
"Strategies are defined in the bank config under 'retain_strategies'.",
|
||||
)
|
||||
update_mode: Literal["replace", "append"] | None = Field(
|
||||
default=None,
|
||||
description="How to handle an existing document with the same document_id. "
|
||||
"'replace' (default) deletes old data and reprocesses from scratch. "
|
||||
"'append' concatenates new content to the existing document text and reprocesses.",
|
||||
)
|
||||
|
||||
@field_validator("timestamp", mode="before")
|
||||
@classmethod
|
||||
@@ -1410,6 +1454,7 @@ class BankStatsResponse(BaseModel):
|
||||
"failed_operations": 0,
|
||||
"last_consolidated_at": "2024-01-15T10:30:00Z",
|
||||
"pending_consolidation": 0,
|
||||
"failed_consolidation": 0,
|
||||
"total_observations": 45,
|
||||
}
|
||||
}
|
||||
@@ -1425,12 +1470,41 @@ class BankStatsResponse(BaseModel):
|
||||
links_breakdown: dict[str, dict[str, int]]
|
||||
pending_operations: int
|
||||
failed_operations: int
|
||||
operations_by_status: dict[str, int] = Field(
|
||||
default_factory=dict,
|
||||
description="Async operations grouped by status (pending, in_progress, completed, failed, cancelled).",
|
||||
)
|
||||
# Consolidation stats
|
||||
last_consolidated_at: str | None = Field(default=None, description="When consolidation last ran (ISO format)")
|
||||
pending_consolidation: int = Field(default=0, description="Number of memories not yet processed into observations")
|
||||
failed_consolidation: int = Field(
|
||||
default=0,
|
||||
description="Number of source memories (world/experience) whose consolidation permanently failed and can be retried via the consolidation recovery endpoint.",
|
||||
)
|
||||
total_observations: int = Field(default=0, description="Total number of observations")
|
||||
|
||||
|
||||
class MemoryTimeseriesBucket(BaseModel):
|
||||
"""One bucket in the memory ingestion time-series."""
|
||||
|
||||
time: str = Field(description="Bucket start timestamp in ISO-8601 (UTC).")
|
||||
world: int = Field(default=0, description="World-fact memories ingested in this bucket.")
|
||||
experience: int = Field(default=0, description="Experience memories ingested in this bucket.")
|
||||
observation: int = Field(default=0, description="Observations recorded in this bucket.")
|
||||
|
||||
|
||||
class MemoriesTimeseriesResponse(BaseModel):
|
||||
"""Time-series of memory ingestion bucketed by time and fact type."""
|
||||
|
||||
bank_id: str
|
||||
period: str = Field(description="One of: 1h, 12h, 1d, 7d, 30d, 90d.")
|
||||
trunc: str = Field(description="Bucket granularity: minute, hour, day.")
|
||||
buckets: list[MemoryTimeseriesBucket] = Field(
|
||||
default_factory=list,
|
||||
description="Per-bucket counts, always returned fully padded for the requested period.",
|
||||
)
|
||||
|
||||
|
||||
# Mental Model models
|
||||
|
||||
|
||||
@@ -1487,6 +1561,16 @@ class UpdateDirectiveRequest(BaseModel):
|
||||
class MentalModelTrigger(BaseModel):
|
||||
"""Trigger settings for a mental model."""
|
||||
|
||||
mode: Literal["full", "delta"] = Field(
|
||||
default="full",
|
||||
description=(
|
||||
"Refresh mode. 'full' (default) regenerates the mental model content from scratch on each refresh. "
|
||||
"'delta' performs surgical edits against the existing content: unchanged sections are preserved "
|
||||
"byte-for-byte, stale content is removed, new content is added. If the mental model has no existing "
|
||||
"content, or if the source_query has changed since the last refresh, delta mode falls back to a full "
|
||||
"regeneration automatically."
|
||||
),
|
||||
)
|
||||
refresh_after_consolidation: bool = Field(
|
||||
default=False,
|
||||
description="If true, refresh this mental model after observations consolidation (real-time mode)",
|
||||
@@ -1520,6 +1604,27 @@ class MentalModelTrigger(BaseModel):
|
||||
"Supports nested and/or/not expressions for complex tag-based scoping."
|
||||
),
|
||||
)
|
||||
include_chunks: bool | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Override whether the internal recall used during refresh returns raw chunk text. "
|
||||
"None means use the bank/global config default (recall_include_chunks)."
|
||||
),
|
||||
)
|
||||
recall_max_tokens: int | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Override the token budget for facts returned by the internal recall during refresh. "
|
||||
"None means use the bank/global config default (recall_max_tokens)."
|
||||
),
|
||||
)
|
||||
recall_chunks_max_tokens: int | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Override the token budget for raw chunks returned by the internal recall during refresh. "
|
||||
"None means use the bank/global config default (recall_chunks_max_tokens)."
|
||||
),
|
||||
)
|
||||
|
||||
@field_validator("fact_types")
|
||||
@classmethod
|
||||
@@ -1549,6 +1654,14 @@ class MentalModelResponse(BaseModel):
|
||||
default=None,
|
||||
description="Full reflect API response payload including based_on facts and observations",
|
||||
)
|
||||
is_stale: bool | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"True when new memories matching this mental model's tag/fact_type scope have been "
|
||||
"ingested since last_refreshed_at, or consolidation has pending items. Only populated "
|
||||
"when detail=full."
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class MentalModelListResponse(BaseModel):
|
||||
@@ -1661,10 +1774,67 @@ class BankTemplateConfig(BaseModel):
|
||||
disposition_skepticism: int | None = Field(default=None, ge=1, le=5, description="Skepticism trait (1-5)")
|
||||
disposition_literalism: int | None = Field(default=None, ge=1, le=5, description="Literalism trait (1-5)")
|
||||
disposition_empathy: int | None = Field(default=None, ge=1, le=5, description="Empathy trait (1-5)")
|
||||
entity_labels: list[str] | None = Field(default=None, description="Controlled vocabulary for entity labels")
|
||||
entity_labels: list[dict[str, Any]] | None = Field(
|
||||
default=None, description="Controlled vocabulary for entity labels"
|
||||
)
|
||||
entities_allow_free_form: bool | None = Field(
|
||||
default=None, description="Allow entities outside the label vocabulary"
|
||||
)
|
||||
retain_default_strategy: str | None = Field(
|
||||
default=None, description="Name of the default retain strategy (key into retain_strategies map)"
|
||||
)
|
||||
retain_strategies: dict | None = Field(
|
||||
default=None, description="Map of retain strategy name to per-strategy config dict"
|
||||
)
|
||||
retain_chunk_batch_size: int | None = Field(
|
||||
default=None, description="Max chunks per streaming batch (0 disables batching)"
|
||||
)
|
||||
mcp_enabled_tools: list[str] | None = Field(
|
||||
default=None, description="MCP tool allowlist for this bank (None = all tools)"
|
||||
)
|
||||
consolidation_llm_batch_size: int | None = Field(
|
||||
default=None, description="LLM batch size for observation consolidation"
|
||||
)
|
||||
consolidation_source_facts_max_tokens: int | None = Field(
|
||||
default=None, description="Max tokens of source facts per consolidation batch"
|
||||
)
|
||||
consolidation_source_facts_max_tokens_per_observation: int | None = Field(
|
||||
default=None, description="Max tokens of source facts per observation"
|
||||
)
|
||||
max_observations_per_scope: int | None = Field(
|
||||
default=None, description="Max observations to retain per consolidation scope"
|
||||
)
|
||||
reflect_source_facts_max_tokens: int | None = Field(
|
||||
default=None, description="Max tokens of source facts per reflect call"
|
||||
)
|
||||
llm_gemini_safety_settings: list | None = Field(
|
||||
default=None, description="Per-bank Gemini/VertexAI safety filter settings"
|
||||
)
|
||||
recall_budget_function: str | None = Field(
|
||||
default=None, description="Recall budget mapping function: 'fixed' or 'adaptive'"
|
||||
)
|
||||
recall_budget_fixed_low: int | None = Field(
|
||||
default=None, description="Fixed thinking_budget for budget=low (function='fixed')"
|
||||
)
|
||||
recall_budget_fixed_mid: int | None = Field(
|
||||
default=None, description="Fixed thinking_budget for budget=mid (function='fixed')"
|
||||
)
|
||||
recall_budget_fixed_high: int | None = Field(
|
||||
default=None, description="Fixed thinking_budget for budget=high (function='fixed')"
|
||||
)
|
||||
recall_budget_adaptive_low: float | None = Field(
|
||||
default=None, description="Ratio of max_tokens for budget=low (function='adaptive')"
|
||||
)
|
||||
recall_budget_adaptive_mid: float | None = Field(
|
||||
default=None, description="Ratio of max_tokens for budget=mid (function='adaptive')"
|
||||
)
|
||||
recall_budget_adaptive_high: float | None = Field(
|
||||
default=None, description="Ratio of max_tokens for budget=high (function='adaptive')"
|
||||
)
|
||||
recall_budget_min: int | None = Field(default=None, description="Floor for the adaptive function (after clamping)")
|
||||
recall_budget_max: int | None = Field(
|
||||
default=None, description="Ceiling for the adaptive function (after clamping)"
|
||||
)
|
||||
|
||||
def get_config_updates(self) -> dict[str, Any]:
|
||||
"""Return only the fields that were explicitly set (non-None)."""
|
||||
@@ -1792,6 +1962,150 @@ class BankTemplateImportResponse(BaseModel):
|
||||
dry_run: bool = Field(default=False, description="True if this was a validation-only run")
|
||||
|
||||
|
||||
def validate_bank_template(manifest: "BankTemplateManifest") -> list[str]:
|
||||
"""Validate a parsed manifest beyond Pydantic's structural checks.
|
||||
|
||||
Returns a list of human-readable error strings (e.g. invalid
|
||||
extraction mode values, conflicting settings).
|
||||
"""
|
||||
errors: list[str] = []
|
||||
if manifest.bank:
|
||||
bank = manifest.bank
|
||||
if bank.retain_extraction_mode is not None:
|
||||
valid_modes = ("concise", "verbose", "custom", "chunks")
|
||||
if bank.retain_extraction_mode not in valid_modes:
|
||||
errors.append(
|
||||
f"bank.retain_extraction_mode: must be one of {valid_modes}, got '{bank.retain_extraction_mode}'"
|
||||
)
|
||||
if bank.retain_custom_instructions and bank.retain_extraction_mode != "custom":
|
||||
errors.append("bank.retain_custom_instructions: requires retain_extraction_mode='custom'")
|
||||
if manifest.mental_models:
|
||||
for i, mm in enumerate(manifest.mental_models):
|
||||
if not mm.name.strip():
|
||||
errors.append(f"mental_models[{i}].name: must not be empty")
|
||||
if not mm.source_query.strip():
|
||||
errors.append(f"mental_models[{i}].source_query: must not be empty")
|
||||
if manifest.directives:
|
||||
for i, d in enumerate(manifest.directives):
|
||||
if not d.name.strip():
|
||||
errors.append(f"directives[{i}].name: must not be empty")
|
||||
if not d.content.strip():
|
||||
errors.append(f"directives[{i}].content: must not be empty")
|
||||
return errors
|
||||
|
||||
|
||||
async def apply_bank_template_manifest(
|
||||
memory,
|
||||
bank_id: str,
|
||||
manifest: "BankTemplateManifest",
|
||||
request_context: "RequestContext",
|
||||
) -> "BankTemplateImportResponse":
|
||||
"""Apply a validated BankTemplateManifest to an existing bank.
|
||||
|
||||
Shared by the /import endpoint and the default-template-on-create hook
|
||||
driven by HINDSIGHT_API_DEFAULT_BANK_TEMPLATE. The bank MUST already
|
||||
exist; caller is responsible for validation (Pydantic + validate_bank_template).
|
||||
"""
|
||||
config_applied = False
|
||||
if manifest.bank:
|
||||
config_updates = manifest.bank.get_config_updates()
|
||||
if config_updates:
|
||||
await memory._config_resolver.update_bank_config(bank_id, config_updates, request_context)
|
||||
config_applied = True
|
||||
|
||||
created_ids: list[str] = []
|
||||
updated_ids: list[str] = []
|
||||
operation_ids: list[str] = []
|
||||
|
||||
if manifest.mental_models:
|
||||
# Fetch existing mental models to decide create vs update
|
||||
existing = await memory.list_mental_models(bank_id=bank_id, request_context=request_context)
|
||||
existing_by_id = {m["id"]: m for m in existing}
|
||||
|
||||
for mm in manifest.mental_models:
|
||||
if mm.id in existing_by_id:
|
||||
await memory.update_mental_model(
|
||||
bank_id=bank_id,
|
||||
mental_model_id=mm.id,
|
||||
name=mm.name,
|
||||
source_query=mm.source_query,
|
||||
max_tokens=mm.max_tokens,
|
||||
tags=mm.tags if mm.tags else None,
|
||||
trigger=mm.trigger.model_dump() if mm.trigger else None,
|
||||
request_context=request_context,
|
||||
)
|
||||
result = await memory.submit_async_refresh_mental_model(
|
||||
bank_id=bank_id,
|
||||
mental_model_id=mm.id,
|
||||
request_context=request_context,
|
||||
)
|
||||
operation_ids.append(result["operation_id"])
|
||||
updated_ids.append(mm.id)
|
||||
else:
|
||||
mental_model = await memory.create_mental_model(
|
||||
bank_id=bank_id,
|
||||
name=mm.name,
|
||||
source_query=mm.source_query,
|
||||
content="Generating content...",
|
||||
mental_model_id=mm.id,
|
||||
tags=mm.tags if mm.tags else None,
|
||||
max_tokens=mm.max_tokens,
|
||||
trigger=mm.trigger.model_dump() if mm.trigger else None,
|
||||
request_context=request_context,
|
||||
)
|
||||
result = await memory.submit_async_refresh_mental_model(
|
||||
bank_id=bank_id,
|
||||
mental_model_id=mental_model["id"],
|
||||
request_context=request_context,
|
||||
)
|
||||
operation_ids.append(result["operation_id"])
|
||||
created_ids.append(mm.id)
|
||||
|
||||
directives_created: list[str] = []
|
||||
directives_updated: list[str] = []
|
||||
|
||||
if manifest.directives:
|
||||
existing_directives = await memory.list_directives(
|
||||
bank_id=bank_id, active_only=False, request_context=request_context
|
||||
)
|
||||
existing_by_name = {d["name"]: d for d in existing_directives}
|
||||
|
||||
for directive in manifest.directives:
|
||||
if directive.name in existing_by_name:
|
||||
await memory.update_directive(
|
||||
bank_id=bank_id,
|
||||
directive_id=existing_by_name[directive.name]["id"],
|
||||
content=directive.content,
|
||||
priority=directive.priority,
|
||||
is_active=directive.is_active,
|
||||
tags=directive.tags if directive.tags else None,
|
||||
request_context=request_context,
|
||||
)
|
||||
directives_updated.append(directive.name)
|
||||
else:
|
||||
await memory.create_directive(
|
||||
bank_id=bank_id,
|
||||
name=directive.name,
|
||||
content=directive.content,
|
||||
priority=directive.priority,
|
||||
is_active=directive.is_active,
|
||||
tags=directive.tags if directive.tags else None,
|
||||
request_context=request_context,
|
||||
)
|
||||
directives_created.append(directive.name)
|
||||
|
||||
return BankTemplateImportResponse(
|
||||
bank_id=bank_id,
|
||||
config_applied=config_applied,
|
||||
mental_models_created=created_ids,
|
||||
mental_models_updated=updated_ids,
|
||||
directives_created=directives_created,
|
||||
directives_updated=directives_updated,
|
||||
operation_ids=operation_ids,
|
||||
dry_run=False,
|
||||
)
|
||||
|
||||
|
||||
class OperationResponse(BaseModel):
|
||||
"""Response model for a single async operation."""
|
||||
|
||||
@@ -1805,6 +2119,8 @@ class OperationResponse(BaseModel):
|
||||
"created_at": "2024-01-15T10:30:00Z",
|
||||
"status": "pending",
|
||||
"error_message": None,
|
||||
"retry_count": 0,
|
||||
"next_retry_at": None,
|
||||
}
|
||||
}
|
||||
)
|
||||
@@ -1816,6 +2132,20 @@ class OperationResponse(BaseModel):
|
||||
created_at: str
|
||||
status: str
|
||||
error_message: str | None
|
||||
retry_count: int | None = Field(
|
||||
default=None,
|
||||
description="Number of times this operation has been retried after failure.",
|
||||
)
|
||||
next_retry_at: str | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"When the worker will next attempt this operation. For a pending "
|
||||
"operation, a value in the future indicates the task is waiting "
|
||||
"rather than available for immediate pickup — for example, an "
|
||||
"extension may have raised DeferOperation to park the task until "
|
||||
"some backpressure window opens. Always null for completed tasks."
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class ConsolidationResponse(BaseModel):
|
||||
@@ -1925,6 +2255,19 @@ class OperationStatusResponse(BaseModel):
|
||||
updated_at: str | None = None
|
||||
completed_at: str | None = None
|
||||
error_message: str | None = None
|
||||
retry_count: int | None = Field(
|
||||
default=None,
|
||||
description="Number of times this operation has been retried after failure.",
|
||||
)
|
||||
next_retry_at: str | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"When the worker will next attempt this operation. For a pending "
|
||||
"operation, a value in the future indicates the task is parked "
|
||||
"(e.g. by an extension raising DeferOperation) rather than awaiting "
|
||||
"immediate pickup."
|
||||
),
|
||||
)
|
||||
result_metadata: dict[str, Any] | None = Field(
|
||||
default=None,
|
||||
description="Internal metadata for debugging. Structure may change without notice. Not for production use.",
|
||||
@@ -1932,6 +2275,10 @@ class OperationStatusResponse(BaseModel):
|
||||
child_operations: list[ChildOperationStatus] | None = Field(
|
||||
default=None, description="Child operations for batch operations (if applicable)"
|
||||
)
|
||||
task_payload: dict[str, Any] | None = Field(
|
||||
default=None,
|
||||
description="Raw task payload (params the operation was submitted with). Only populated when include_payload=true.",
|
||||
)
|
||||
|
||||
|
||||
class AsyncOperationSubmitResponse(BaseModel):
|
||||
@@ -2268,7 +2615,7 @@ def create_app(
|
||||
schema=schema,
|
||||
tenant_extension=memory._tenant_extension,
|
||||
max_slots=config.worker_max_slots,
|
||||
consolidation_max_slots=config.worker_consolidation_max_slots,
|
||||
slot_reservations=config.worker_slot_reservations,
|
||||
)
|
||||
poller_task = asyncio.create_task(poller.run())
|
||||
logging.info(f"Worker poller started (worker_id={worker_id})")
|
||||
@@ -2615,6 +2962,7 @@ def _register_routes(app: FastAPI):
|
||||
bank_id: str,
|
||||
type: str | None = None,
|
||||
q: str | None = None,
|
||||
consolidation_state: str | None = None,
|
||||
limit: int = 100,
|
||||
offset: int = 0,
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
@@ -2629,6 +2977,8 @@ def _register_routes(app: FastAPI):
|
||||
bank_id: Memory Bank ID (from path)
|
||||
type: Filter by fact type (world, experience, opinion)
|
||||
q: Search query for full-text search (searches text and context)
|
||||
consolidation_state: Filter by consolidation state for source memories
|
||||
(world/experience). One of 'failed', 'pending', or 'done'.
|
||||
limit: Maximum number of results (default: 100)
|
||||
offset: Offset for pagination (default: 0)
|
||||
"""
|
||||
@@ -2637,11 +2987,14 @@ def _register_routes(app: FastAPI):
|
||||
bank_id=bank_id,
|
||||
fact_type=type,
|
||||
search_query=q,
|
||||
consolidation_state=consolidation_state,
|
||||
limit=limit,
|
||||
offset=offset,
|
||||
request_context=request_context,
|
||||
)
|
||||
return data
|
||||
except ValueError as e:
|
||||
raise HTTPException(status_code=400, detail=str(e))
|
||||
except OperationValidationError as e:
|
||||
raise HTTPException(status_code=e.status_code, detail=e.reason)
|
||||
except (AuthenticationError, HTTPException):
|
||||
@@ -2677,6 +3030,8 @@ def _register_routes(app: FastAPI):
|
||||
return data
|
||||
except OperationValidationError as e:
|
||||
raise HTTPException(status_code=e.status_code, detail=e.reason)
|
||||
except ValueError as e:
|
||||
raise HTTPException(status_code=400, detail=str(e))
|
||||
except (AuthenticationError, HTTPException):
|
||||
raise
|
||||
except Exception as e:
|
||||
@@ -3094,8 +3449,10 @@ def _register_routes(app: FastAPI):
|
||||
links_breakdown=links_breakdown,
|
||||
pending_operations=ops.get("pending", 0),
|
||||
failed_operations=ops.get("failed", 0),
|
||||
operations_by_status=ops,
|
||||
last_consolidated_at=stats["last_consolidated_at"],
|
||||
pending_consolidation=stats["pending_consolidation"],
|
||||
failed_consolidation=stats.get("failed_consolidation", 0),
|
||||
total_observations=stats["total_observations"],
|
||||
)
|
||||
except OperationValidationError as e:
|
||||
@@ -3109,6 +3466,35 @@ def _register_routes(app: FastAPI):
|
||||
logger.error(f"Error in /v1/default/banks/{bank_id}/stats: {error_detail}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
@app.get(
|
||||
"/v1/default/banks/{bank_id}/stats/memories-timeseries",
|
||||
response_model=MemoriesTimeseriesResponse,
|
||||
summary="Memory ingestion time-series",
|
||||
description="Memories ingested over a period, bucketed by time and broken down by fact type.",
|
||||
operation_id="get_memories_timeseries",
|
||||
tags=["Banks"],
|
||||
)
|
||||
async def api_memories_timeseries(
|
||||
bank_id: str,
|
||||
period: str = "7d",
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
):
|
||||
try:
|
||||
data = await app.state.memory.get_memories_timeseries(
|
||||
bank_id, period=period, request_context=request_context
|
||||
)
|
||||
return MemoriesTimeseriesResponse(**data)
|
||||
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 /v1/default/banks/{bank_id}/stats/memories-timeseries: {error_detail}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
@app.get(
|
||||
"/v1/default/banks/{bank_id}/entities",
|
||||
response_model=EntityListResponse,
|
||||
@@ -3145,6 +3531,36 @@ def _register_routes(app: FastAPI):
|
||||
logger.error(f"Error in /v1/default/banks/{bank_id}/entities: {error_detail}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
@app.get(
|
||||
"/v1/default/banks/{bank_id}/entities/graph",
|
||||
response_model=EntityGraphResponse,
|
||||
summary="Get entity co-occurrence graph",
|
||||
description="Return a graph of entities (nodes) and their co-occurrences (edges) for visualization.",
|
||||
operation_id="get_entity_graph",
|
||||
tags=["Entities"],
|
||||
)
|
||||
async def api_entity_graph(
|
||||
bank_id: str,
|
||||
limit: int = Query(default=1000, description="Maximum number of co-occurrence edges to return"),
|
||||
min_count: int = Query(default=1, description="Minimum cooccurrence_count to include an edge"),
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
):
|
||||
"""Return entity co-occurrence graph for a bank."""
|
||||
try:
|
||||
return await app.state.memory.get_entity_graph(
|
||||
bank_id, limit=limit, min_count=min_count, request_context=request_context
|
||||
)
|
||||
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 /v1/default/banks/{bank_id}/entities/graph: {error_detail}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
@app.get(
|
||||
"/v1/default/banks/{bank_id}/entities/{entity_id}",
|
||||
response_model=EntityDetailResponse,
|
||||
@@ -3287,6 +3703,8 @@ def _register_routes(app: FastAPI):
|
||||
raise
|
||||
except OperationValidationError as e:
|
||||
raise HTTPException(status_code=e.status_code, detail=e.reason)
|
||||
except ValueError as e:
|
||||
raise HTTPException(status_code=400, detail=str(e))
|
||||
except Exception as e:
|
||||
import traceback
|
||||
|
||||
@@ -3320,6 +3738,8 @@ def _register_routes(app: FastAPI):
|
||||
raise
|
||||
except OperationValidationError as e:
|
||||
raise HTTPException(status_code=e.status_code, detail=e.reason)
|
||||
except ValueError as e:
|
||||
raise HTTPException(status_code=400, detail=str(e))
|
||||
except Exception as e:
|
||||
import traceback
|
||||
|
||||
@@ -3484,6 +3904,8 @@ def _register_routes(app: FastAPI):
|
||||
return {"status": "deleted"}
|
||||
except OperationValidationError as e:
|
||||
raise HTTPException(status_code=e.status_code, detail=e.reason)
|
||||
except ValueError as e:
|
||||
raise HTTPException(status_code=400, detail=str(e))
|
||||
except (AuthenticationError, HTTPException):
|
||||
raise
|
||||
except Exception as e:
|
||||
@@ -4008,7 +4430,13 @@ def _register_routes(app: FastAPI):
|
||||
tags=["Operations"],
|
||||
)
|
||||
async def api_get_operation_status(
|
||||
bank_id: str, operation_id: str, request_context: RequestContext = Depends(get_request_context)
|
||||
bank_id: str,
|
||||
operation_id: str,
|
||||
include_payload: bool = Query(
|
||||
default=False,
|
||||
description="Include the raw task payload (submission params) in the response. May be large.",
|
||||
),
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
):
|
||||
"""Get the status of an async operation."""
|
||||
try:
|
||||
@@ -4018,7 +4446,9 @@ def _register_routes(app: FastAPI):
|
||||
except ValueError:
|
||||
raise HTTPException(status_code=400, detail=f"Invalid operation_id format: {operation_id}")
|
||||
|
||||
result = await app.state.memory.get_operation_status(bank_id, operation_id, request_context=request_context)
|
||||
result = await app.state.memory.get_operation_status(
|
||||
bank_id, operation_id, request_context=request_context, include_payload=include_payload
|
||||
)
|
||||
return OperationStatusResponse(**result)
|
||||
except OperationValidationError as e:
|
||||
raise HTTPException(status_code=e.status_code, detail=e.reason)
|
||||
@@ -4361,38 +4791,6 @@ def _register_routes(app: FastAPI):
|
||||
# Bank Template Import / Export
|
||||
# =====================================================================
|
||||
|
||||
def _validate_template(manifest: BankTemplateManifest) -> list[str]:
|
||||
"""Validate a parsed manifest beyond Pydantic's structural checks.
|
||||
|
||||
Returns a list of human-readable error strings (e.g. invalid
|
||||
extraction mode values, conflicting settings).
|
||||
"""
|
||||
errors: list[str] = []
|
||||
if manifest.bank:
|
||||
bank = manifest.bank
|
||||
if bank.retain_extraction_mode is not None:
|
||||
valid_modes = ("concise", "verbose", "custom", "chunks")
|
||||
if bank.retain_extraction_mode not in valid_modes:
|
||||
errors.append(
|
||||
f"bank.retain_extraction_mode: must be one of {valid_modes}, "
|
||||
f"got '{bank.retain_extraction_mode}'"
|
||||
)
|
||||
if bank.retain_custom_instructions and bank.retain_extraction_mode != "custom":
|
||||
errors.append("bank.retain_custom_instructions: requires retain_extraction_mode='custom'")
|
||||
if manifest.mental_models:
|
||||
for i, mm in enumerate(manifest.mental_models):
|
||||
if not mm.name.strip():
|
||||
errors.append(f"mental_models[{i}].name: must not be empty")
|
||||
if not mm.source_query.strip():
|
||||
errors.append(f"mental_models[{i}].source_query: must not be empty")
|
||||
if manifest.directives:
|
||||
for i, d in enumerate(manifest.directives):
|
||||
if not d.name.strip():
|
||||
errors.append(f"directives[{i}].name: must not be empty")
|
||||
if not d.content.strip():
|
||||
errors.append(f"directives[{i}].content: must not be empty")
|
||||
return errors
|
||||
|
||||
@app.post(
|
||||
"/v1/default/banks/{bank_id}/import",
|
||||
response_model=BankTemplateImportResponse,
|
||||
@@ -4428,7 +4826,7 @@ def _register_routes(app: FastAPI):
|
||||
)
|
||||
|
||||
# Semantic validation beyond Pydantic structural checks
|
||||
validation_errors = _validate_template(body)
|
||||
validation_errors = validate_bank_template(body)
|
||||
if validation_errors:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
@@ -4446,107 +4844,11 @@ def _register_routes(app: FastAPI):
|
||||
# Ensure bank exists (auto-creates with defaults if needed)
|
||||
await app.state.memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
|
||||
config_applied = False
|
||||
if body.bank:
|
||||
config_updates = body.bank.get_config_updates()
|
||||
if config_updates:
|
||||
await app.state.memory._config_resolver.update_bank_config(bank_id, config_updates, request_context)
|
||||
config_applied = True
|
||||
|
||||
created_ids: list[str] = []
|
||||
updated_ids: list[str] = []
|
||||
operation_ids: list[str] = []
|
||||
|
||||
if body.mental_models:
|
||||
# Fetch existing mental models to decide create vs update
|
||||
existing = await app.state.memory.list_mental_models(bank_id=bank_id, request_context=request_context)
|
||||
existing_by_id = {m["id"]: m for m in existing}
|
||||
|
||||
for mm in body.mental_models:
|
||||
if mm.id in existing_by_id:
|
||||
# Update existing mental model metadata
|
||||
await app.state.memory.update_mental_model(
|
||||
bank_id=bank_id,
|
||||
mental_model_id=mm.id,
|
||||
name=mm.name,
|
||||
source_query=mm.source_query,
|
||||
max_tokens=mm.max_tokens,
|
||||
tags=mm.tags if mm.tags else None,
|
||||
trigger=mm.trigger.model_dump() if mm.trigger else None,
|
||||
request_context=request_context,
|
||||
)
|
||||
# Schedule a refresh to regenerate content with updated query
|
||||
result = await app.state.memory.submit_async_refresh_mental_model(
|
||||
bank_id=bank_id,
|
||||
mental_model_id=mm.id,
|
||||
request_context=request_context,
|
||||
)
|
||||
operation_ids.append(result["operation_id"])
|
||||
updated_ids.append(mm.id)
|
||||
else:
|
||||
# Create new mental model
|
||||
mental_model = await app.state.memory.create_mental_model(
|
||||
bank_id=bank_id,
|
||||
name=mm.name,
|
||||
source_query=mm.source_query,
|
||||
content="Generating content...",
|
||||
mental_model_id=mm.id,
|
||||
tags=mm.tags if mm.tags else None,
|
||||
max_tokens=mm.max_tokens,
|
||||
trigger=mm.trigger.model_dump() if mm.trigger else None,
|
||||
request_context=request_context,
|
||||
)
|
||||
result = await app.state.memory.submit_async_refresh_mental_model(
|
||||
bank_id=bank_id,
|
||||
mental_model_id=mental_model["id"],
|
||||
request_context=request_context,
|
||||
)
|
||||
operation_ids.append(result["operation_id"])
|
||||
created_ids.append(mm.id)
|
||||
|
||||
directives_created: list[str] = []
|
||||
directives_updated: list[str] = []
|
||||
|
||||
if body.directives:
|
||||
# Fetch existing directives to decide create vs update (matched by name)
|
||||
existing_directives = await app.state.memory.list_directives(
|
||||
bank_id=bank_id, active_only=False, request_context=request_context
|
||||
)
|
||||
existing_by_name = {d["name"]: d for d in existing_directives}
|
||||
|
||||
for directive in body.directives:
|
||||
if directive.name in existing_by_name:
|
||||
await app.state.memory.update_directive(
|
||||
bank_id=bank_id,
|
||||
directive_id=existing_by_name[directive.name]["id"],
|
||||
content=directive.content,
|
||||
priority=directive.priority,
|
||||
is_active=directive.is_active,
|
||||
tags=directive.tags if directive.tags else None,
|
||||
request_context=request_context,
|
||||
)
|
||||
directives_updated.append(directive.name)
|
||||
else:
|
||||
await app.state.memory.create_directive(
|
||||
bank_id=bank_id,
|
||||
name=directive.name,
|
||||
content=directive.content,
|
||||
priority=directive.priority,
|
||||
is_active=directive.is_active,
|
||||
tags=directive.tags if directive.tags else None,
|
||||
request_context=request_context,
|
||||
)
|
||||
directives_created.append(directive.name)
|
||||
|
||||
return BankTemplateImportResponse(
|
||||
return await apply_bank_template_manifest(
|
||||
memory=app.state.memory,
|
||||
bank_id=bank_id,
|
||||
config_applied=config_applied,
|
||||
mental_models_created=created_ids,
|
||||
mental_models_updated=updated_ids,
|
||||
directives_created=directives_created,
|
||||
directives_updated=directives_updated,
|
||||
operation_ids=operation_ids,
|
||||
dry_run=False,
|
||||
manifest=body,
|
||||
request_context=request_context,
|
||||
)
|
||||
except ValueError as e:
|
||||
raise HTTPException(status_code=400, detail=str(e))
|
||||
@@ -4948,7 +5250,9 @@ def _register_routes(app: FastAPI):
|
||||
from hindsight_api.engine.retain import bank_utils
|
||||
|
||||
# Ensure the bank row exists before inserting into webhooks (FK constraint).
|
||||
await bank_utils.get_bank_profile(pool, bank_id)
|
||||
_, created = await bank_utils.get_or_create_bank_profile(pool, bank_id)
|
||||
if created:
|
||||
await app.state.memory._apply_default_bank_template(bank_id, request_context)
|
||||
|
||||
webhook_id = uuid.uuid4()
|
||||
now = datetime.now(timezone.utc).isoformat()
|
||||
@@ -5297,6 +5601,8 @@ def _register_routes(app: FastAPI):
|
||||
content_dict["tags"] = item.tags
|
||||
if item.observation_scopes is not None:
|
||||
content_dict["observation_scopes"] = item.observation_scopes
|
||||
if item.update_mode is not None:
|
||||
content_dict["update_mode"] = item.update_mode
|
||||
strategy_groups[effective].append(content_dict)
|
||||
|
||||
if request.async_:
|
||||
|
||||
@@ -97,6 +97,7 @@ def create_mcp_server(memory: MemoryEngine, multi_bank: bool = True) -> FastMCP:
|
||||
_SINGLE_BANK_TOOLS: frozenset[str] = frozenset(
|
||||
{
|
||||
"retain",
|
||||
"sync_retain",
|
||||
"recall",
|
||||
"reflect",
|
||||
"list_mental_models",
|
||||
@@ -156,24 +157,65 @@ def create_mcp_server(memory: MemoryEngine, multi_bank: bool = True) -> FastMCP:
|
||||
return mcp
|
||||
|
||||
|
||||
def _get_mcp_tools(mcp: FastMCP) -> dict:
|
||||
"""Get tool name→object mapping, compatible with FastMCP 2.x and 3.x."""
|
||||
# FastMCP 2.x: _tool_manager._tools
|
||||
if hasattr(mcp, "_tool_manager"):
|
||||
return mcp._tool_manager._tools # type: ignore[union-attr]
|
||||
# FastMCP 3.x: _local_provider._components with "tool:" prefix
|
||||
if hasattr(mcp, "_local_provider"):
|
||||
return {
|
||||
k.split(":")[1].split("@")[0]: v
|
||||
for k, v in mcp._local_provider._components.items() # type: ignore[union-attr]
|
||||
if k.startswith("tool:")
|
||||
}
|
||||
msg = "Cannot locate tools on FastMCP instance"
|
||||
raise AttributeError(msg)
|
||||
|
||||
|
||||
def _make_tools_tolerant(mcp: FastMCP) -> None:
|
||||
"""Wrap all tool run methods to strip unknown arguments before validation.
|
||||
"""Wrap all tool run methods to strip unknown arguments and coerce string-encoded JSON.
|
||||
|
||||
LLMs frequently add extra fields like "explanation" or "reasoning" to tool calls.
|
||||
FastMCP's Pydantic TypeAdapter rejects these with "Unexpected keyword argument".
|
||||
This wraps each tool's run() to filter arguments to only known parameters.
|
||||
|
||||
LLMs also frequently serialize list/dict arguments as JSON strings instead of native
|
||||
types (e.g., tags='["a","b"]' instead of tags=["a","b"]). This auto-coerces them.
|
||||
|
||||
This wraps each tool's run() to apply both fixes before validation.
|
||||
"""
|
||||
try:
|
||||
for name, tool in mcp._tool_manager._tools.items(): # type: ignore[unresolved-attribute] # FastMCP 2.x internal; guarded by try/except
|
||||
tools = _get_mcp_tools(mcp)
|
||||
for name, tool in tools.items():
|
||||
if hasattr(tool, "parameters") and tool.parameters:
|
||||
allowed = set(tool.parameters.get("properties", {}).keys())
|
||||
properties = tool.parameters.get("properties", {})
|
||||
allowed = set(properties.keys())
|
||||
|
||||
# Build sets of parameter names that expect array or object types.
|
||||
# Handles both direct types {"type": "array"} and anyOf/oneOf unions
|
||||
# like {"anyOf": [{"type": "array", ...}, {"type": "null"}]}.
|
||||
array_params: set[str] = set()
|
||||
object_params: set[str] = set()
|
||||
for param_name, param_schema in properties.items():
|
||||
_collect_coercible_types(param_schema, param_name, array_params, object_params)
|
||||
|
||||
original_run = tool.run
|
||||
|
||||
async def _tolerant_run(arguments, _allowed=allowed, _orig=original_run):
|
||||
async def _tolerant_run(
|
||||
arguments,
|
||||
_allowed=allowed,
|
||||
_orig=original_run,
|
||||
_array_params=array_params,
|
||||
_object_params=object_params,
|
||||
):
|
||||
extra_keys = set(arguments.keys()) - _allowed
|
||||
if extra_keys:
|
||||
logger.debug(f"Stripping unknown arguments from tool call: {extra_keys}")
|
||||
arguments = {k: v for k, v in arguments.items() if k in _allowed}
|
||||
|
||||
# Coerce string-encoded JSON for list/dict parameters
|
||||
arguments = _coerce_string_json(arguments, _array_params, _object_params)
|
||||
|
||||
return await _orig(arguments)
|
||||
|
||||
# FunctionTool is a Pydantic model with extra='forbid', so use
|
||||
@@ -183,6 +225,59 @@ def _make_tools_tolerant(mcp: FastMCP) -> None:
|
||||
logger.warning(f"Could not make tools tolerant of extra arguments: {e}")
|
||||
|
||||
|
||||
def _collect_coercible_types(schema: dict, param_name: str, array_params: set[str], object_params: set[str]) -> None:
|
||||
"""Check a JSON Schema property and add param_name to array_params/object_params if applicable."""
|
||||
# Direct type
|
||||
schema_type = schema.get("type")
|
||||
if schema_type == "array":
|
||||
array_params.add(param_name)
|
||||
return
|
||||
if schema_type == "object":
|
||||
object_params.add(param_name)
|
||||
return
|
||||
|
||||
# anyOf / oneOf unions (e.g., list[str] | None → {"anyOf": [{"type": "array"}, {"type": "null"}]})
|
||||
for variant in schema.get("anyOf", []) + schema.get("oneOf", []):
|
||||
variant_type = variant.get("type")
|
||||
if variant_type == "array":
|
||||
array_params.add(param_name)
|
||||
return
|
||||
if variant_type == "object":
|
||||
object_params.add(param_name)
|
||||
return
|
||||
|
||||
|
||||
def _coerce_string_json(arguments: dict, array_params: set[str], object_params: set[str]) -> dict:
|
||||
"""Auto-coerce string-encoded JSON arrays/objects to native types.
|
||||
|
||||
LLM agents frequently serialize list and dict tool arguments as JSON strings.
|
||||
This is backward-compatible: native arrays/objects pass through unchanged.
|
||||
"""
|
||||
for param_name in array_params:
|
||||
val = arguments.get(param_name)
|
||||
if isinstance(val, str):
|
||||
try:
|
||||
parsed = json.loads(val)
|
||||
if isinstance(parsed, list):
|
||||
arguments = {**arguments, param_name: parsed}
|
||||
logger.debug(f"Coerced string to list for parameter '{param_name}'")
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
pass
|
||||
|
||||
for param_name in object_params:
|
||||
val = arguments.get(param_name)
|
||||
if isinstance(val, str):
|
||||
try:
|
||||
parsed = json.loads(val)
|
||||
if isinstance(parsed, dict):
|
||||
arguments = {**arguments, param_name: parsed}
|
||||
logger.debug(f"Coerced string to dict for parameter '{param_name}'")
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
pass
|
||||
|
||||
return arguments
|
||||
|
||||
|
||||
class MCPMiddleware:
|
||||
"""ASGI middleware that intercepts MCP requests and routes to appropriate MCP server.
|
||||
|
||||
|
||||
@@ -177,6 +177,15 @@ ENV_EMBEDDINGS_TEI_URL = "HINDSIGHT_API_EMBEDDINGS_TEI_URL"
|
||||
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"
|
||||
|
||||
# Gemini/Vertex AI embeddings configuration
|
||||
ENV_EMBEDDINGS_GEMINI_API_KEY = "HINDSIGHT_API_EMBEDDINGS_GEMINI_API_KEY"
|
||||
ENV_EMBEDDINGS_GEMINI_MODEL = "HINDSIGHT_API_EMBEDDINGS_GEMINI_MODEL"
|
||||
ENV_EMBEDDINGS_GEMINI_OUTPUT_DIMENSIONALITY = "HINDSIGHT_API_EMBEDDINGS_GEMINI_OUTPUT_DIMENSIONALITY"
|
||||
ENV_EMBEDDINGS_VERTEXAI_PROJECT_ID = "HINDSIGHT_API_EMBEDDINGS_VERTEXAI_PROJECT_ID"
|
||||
ENV_EMBEDDINGS_VERTEXAI_REGION = "HINDSIGHT_API_EMBEDDINGS_VERTEXAI_REGION"
|
||||
ENV_EMBEDDINGS_VERTEXAI_SERVICE_ACCOUNT_KEY = "HINDSIGHT_API_EMBEDDINGS_VERTEXAI_SERVICE_ACCOUNT_KEY"
|
||||
|
||||
# Cohere configuration (separate for embeddings and reranker)
|
||||
ENV_EMBEDDINGS_COHERE_API_KEY = "HINDSIGHT_API_EMBEDDINGS_COHERE_API_KEY"
|
||||
@@ -186,6 +195,13 @@ ENV_RERANKER_COHERE_API_KEY = "HINDSIGHT_API_RERANKER_COHERE_API_KEY"
|
||||
ENV_RERANKER_COHERE_MODEL = "HINDSIGHT_API_RERANKER_COHERE_MODEL"
|
||||
ENV_RERANKER_COHERE_BASE_URL = "HINDSIGHT_API_RERANKER_COHERE_BASE_URL"
|
||||
|
||||
# OpenRouter configuration (embeddings and reranker)
|
||||
ENV_OPENROUTER_API_KEY = "HINDSIGHT_API_OPENROUTER_API_KEY"
|
||||
ENV_EMBEDDINGS_OPENROUTER_API_KEY = "HINDSIGHT_API_EMBEDDINGS_OPENROUTER_API_KEY"
|
||||
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"
|
||||
|
||||
# Deprecated: Legacy shared Cohere API key (for backward compatibility)
|
||||
ENV_COHERE_API_KEY = "HINDSIGHT_API_COHERE_API_KEY"
|
||||
|
||||
@@ -203,6 +219,7 @@ ENV_EMBEDDINGS_LITELLM_SDK_API_KEY = "HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_API_K
|
||||
ENV_EMBEDDINGS_LITELLM_SDK_MODEL = "HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_MODEL"
|
||||
ENV_EMBEDDINGS_LITELLM_SDK_API_BASE = "HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_API_BASE"
|
||||
ENV_EMBEDDINGS_LITELLM_SDK_OUTPUT_DIMENSIONS = "HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_OUTPUT_DIMENSIONS"
|
||||
ENV_EMBEDDINGS_LITELLM_SDK_ENCODING_FORMAT = "HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_ENCODING_FORMAT"
|
||||
ENV_RERANKER_LITELLM_SDK_API_KEY = "HINDSIGHT_API_RERANKER_LITELLM_SDK_API_KEY"
|
||||
ENV_RERANKER_LITELLM_SDK_MODEL = "HINDSIGHT_API_RERANKER_LITELLM_SDK_MODEL"
|
||||
ENV_RERANKER_LITELLM_SDK_API_BASE = "HINDSIGHT_API_RERANKER_LITELLM_SDK_API_BASE"
|
||||
@@ -222,6 +239,7 @@ ENV_RERANKER_LOCAL_BATCH_SIZE = "HINDSIGHT_API_RERANKER_LOCAL_BATCH_SIZE"
|
||||
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_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"
|
||||
@@ -231,6 +249,16 @@ ENV_RERANKER_ZEROENTROPY_API_KEY = "HINDSIGHT_API_RERANKER_ZEROENTROPY_API_KEY"
|
||||
ENV_RERANKER_ZEROENTROPY_MODEL = "HINDSIGHT_API_RERANKER_ZEROENTROPY_MODEL"
|
||||
ENV_RERANKER_ZEROENTROPY_BASE_URL = "HINDSIGHT_API_RERANKER_ZEROENTROPY_BASE_URL"
|
||||
|
||||
# SiliconFlow configuration (reranker only; Cohere-compatible /rerank endpoint)
|
||||
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"
|
||||
|
||||
# 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"
|
||||
ENV_RERANKER_GOOGLE_SERVICE_ACCOUNT_KEY = "HINDSIGHT_API_RERANKER_GOOGLE_SERVICE_ACCOUNT_KEY"
|
||||
|
||||
ENV_VECTOR_EXTENSION = "HINDSIGHT_API_VECTOR_EXTENSION"
|
||||
ENV_TEXT_SEARCH_EXTENSION = "HINDSIGHT_API_TEXT_SEARCH_EXTENSION"
|
||||
|
||||
@@ -239,16 +267,20 @@ ENV_PORT = "HINDSIGHT_API_PORT"
|
||||
ENV_BASE_PATH = "HINDSIGHT_API_BASE_PATH"
|
||||
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_MCP_ENABLED = "HINDSIGHT_API_MCP_ENABLED"
|
||||
ENV_MCP_ENABLED_TOOLS = "HINDSIGHT_API_MCP_ENABLED_TOOLS"
|
||||
ENV_MCP_STATELESS = "HINDSIGHT_API_MCP_STATELESS"
|
||||
ENV_ENABLE_BANK_CONFIG_API = "HINDSIGHT_API_ENABLE_BANK_CONFIG_API"
|
||||
ENV_DEFAULT_BANK_TEMPLATE = "HINDSIGHT_API_DEFAULT_BANK_TEMPLATE"
|
||||
ENV_GRAPH_RETRIEVER = "HINDSIGHT_API_GRAPH_RETRIEVER"
|
||||
ENV_RECALL_MAX_CONCURRENT = "HINDSIGHT_API_RECALL_MAX_CONCURRENT"
|
||||
ENV_RECALL_CONNECTION_BUDGET = "HINDSIGHT_API_RECALL_CONNECTION_BUDGET"
|
||||
ENV_RECALL_MAX_QUERY_TOKENS = "HINDSIGHT_API_RECALL_MAX_QUERY_TOKENS"
|
||||
ENV_MENTAL_MODEL_REFRESH_CONCURRENCY = "HINDSIGHT_API_MENTAL_MODEL_REFRESH_CONCURRENCY"
|
||||
ENV_LINK_EXPANSION_PER_ENTITY_LIMIT = "HINDSIGHT_API_LINK_EXPANSION_PER_ENTITY_LIMIT"
|
||||
ENV_LINK_EXPANSION_TIMEOUT = "HINDSIGHT_API_LINK_EXPANSION_TIMEOUT"
|
||||
|
||||
# OpenTelemetry tracing configuration
|
||||
ENV_OTEL_TRACES_ENABLED = "HINDSIGHT_API_OTEL_TRACES_ENABLED"
|
||||
@@ -256,6 +288,7 @@ ENV_OTEL_EXPORTER_OTLP_ENDPOINT = "HINDSIGHT_API_OTEL_EXPORTER_OTLP_ENDPOINT"
|
||||
ENV_OTEL_EXPORTER_OTLP_HEADERS = "HINDSIGHT_API_OTEL_EXPORTER_OTLP_HEADERS"
|
||||
ENV_OTEL_SERVICE_NAME = "HINDSIGHT_API_OTEL_SERVICE_NAME"
|
||||
ENV_OTEL_DEPLOYMENT_ENVIRONMENT = "HINDSIGHT_API_OTEL_DEPLOYMENT_ENVIRONMENT"
|
||||
ENV_METRICS_INCLUDE_BANK_ID = "HINDSIGHT_API_METRICS_INCLUDE_BANK_ID"
|
||||
|
||||
# Vertex AI configuration
|
||||
ENV_LLM_VERTEXAI_PROJECT_ID = "HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID"
|
||||
@@ -303,12 +336,14 @@ 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_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_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 = (
|
||||
"HINDSIGHT_API_CONSOLIDATION_SOURCE_FACTS_MAX_TOKENS_PER_OBSERVATION"
|
||||
)
|
||||
ENV_CONSOLIDATION_MAX_ATTEMPTS = "HINDSIGHT_API_CONSOLIDATION_MAX_ATTEMPTS"
|
||||
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"
|
||||
@@ -320,6 +355,14 @@ ENV_WEBHOOK_SECRET = "HINDSIGHT_API_WEBHOOK_SECRET"
|
||||
ENV_WEBHOOK_EVENT_TYPES = "HINDSIGHT_API_WEBHOOK_EVENT_TYPES"
|
||||
ENV_WEBHOOK_DELIVERY_POLL_INTERVAL_SECONDS = "HINDSIGHT_API_WEBHOOK_DELIVERY_POLL_INTERVAL_SECONDS"
|
||||
|
||||
# Built-in llama.cpp configuration (for provider=llamacpp)
|
||||
ENV_LLAMACPP_MODEL_PATH = "HINDSIGHT_API_LLAMACPP_MODEL_PATH"
|
||||
ENV_LLAMACPP_GPU_LAYERS = "HINDSIGHT_API_LLAMACPP_GPU_LAYERS"
|
||||
ENV_LLAMACPP_CONTEXT_SIZE = "HINDSIGHT_API_LLAMACPP_CONTEXT_SIZE"
|
||||
ENV_LLAMACPP_CHAT_FORMAT = "HINDSIGHT_API_LLAMACPP_CHAT_FORMAT"
|
||||
ENV_LLAMACPP_NO_GRAMMAR = "HINDSIGHT_API_LLAMACPP_NO_GRAMMAR"
|
||||
ENV_LLAMACPP_EXTRA_ARGS = "HINDSIGHT_API_LLAMACPP_EXTRA_ARGS"
|
||||
|
||||
# Optimization flags
|
||||
ENV_SKIP_LLM_VERIFICATION = "HINDSIGHT_API_SKIP_LLM_VERIFICATION"
|
||||
ENV_LAZY_RERANKER = "HINDSIGHT_API_LAZY_RERANKER"
|
||||
@@ -332,6 +375,7 @@ ENV_DB_POOL_MIN_SIZE = "HINDSIGHT_API_DB_POOL_MIN_SIZE"
|
||||
ENV_DB_POOL_MAX_SIZE = "HINDSIGHT_API_DB_POOL_MAX_SIZE"
|
||||
ENV_DB_COMMAND_TIMEOUT = "HINDSIGHT_API_DB_COMMAND_TIMEOUT"
|
||||
ENV_DB_ACQUIRE_TIMEOUT = "HINDSIGHT_API_DB_ACQUIRE_TIMEOUT"
|
||||
ENV_DB_STATEMENT_TIMEOUT = "HINDSIGHT_API_DB_STATEMENT_TIMEOUT"
|
||||
|
||||
# Worker configuration (distributed task processing)
|
||||
ENV_WORKER_ENABLED = "HINDSIGHT_API_WORKER_ENABLED"
|
||||
@@ -340,7 +384,18 @@ ENV_WORKER_POLL_INTERVAL_MS = "HINDSIGHT_API_WORKER_POLL_INTERVAL_MS"
|
||||
ENV_WORKER_MAX_RETRIES = "HINDSIGHT_API_WORKER_MAX_RETRIES"
|
||||
ENV_WORKER_HTTP_PORT = "HINDSIGHT_API_WORKER_HTTP_PORT"
|
||||
ENV_WORKER_MAX_SLOTS = "HINDSIGHT_API_WORKER_MAX_SLOTS"
|
||||
ENV_WORKER_CONSOLIDATION_MAX_SLOTS = "HINDSIGHT_API_WORKER_CONSOLIDATION_MAX_SLOTS"
|
||||
|
||||
# Per-operation-type slot reservations. Each entry maps an operation_type
|
||||
# (as stored in async_operations.operation_type) to its env var and default.
|
||||
# Adding a new operation type here is the ONLY change needed to make it
|
||||
# reservable via env var — config fields, from_env(), and the
|
||||
# worker_slot_reservations property all derive from this dict.
|
||||
WORKER_SLOT_RESERVATION_TYPES: dict[str, tuple[str, int]] = {
|
||||
"consolidation": ("HINDSIGHT_API_WORKER_CONSOLIDATION_MAX_SLOTS", 2),
|
||||
"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),
|
||||
}
|
||||
ENV_RETAIN_MAX_CONCURRENT = "HINDSIGHT_API_RETAIN_MAX_CONCURRENT"
|
||||
|
||||
# Reflect agent settings
|
||||
@@ -349,6 +404,20 @@ ENV_REFLECT_MAX_CONTEXT_TOKENS = "HINDSIGHT_API_REFLECT_MAX_CONTEXT_TOKENS"
|
||||
ENV_REFLECT_WALL_TIMEOUT = "HINDSIGHT_API_REFLECT_WALL_TIMEOUT"
|
||||
ENV_REFLECT_MISSION = "HINDSIGHT_API_REFLECT_MISSION"
|
||||
ENV_REFLECT_SOURCE_FACTS_MAX_TOKENS = "HINDSIGHT_API_REFLECT_SOURCE_FACTS_MAX_TOKENS"
|
||||
ENV_RECALL_INCLUDE_CHUNKS = "HINDSIGHT_API_RECALL_INCLUDE_CHUNKS"
|
||||
ENV_RECALL_MAX_TOKENS = "HINDSIGHT_API_RECALL_MAX_TOKENS"
|
||||
ENV_RECALL_CHUNKS_MAX_TOKENS = "HINDSIGHT_API_RECALL_CHUNKS_MAX_TOKENS"
|
||||
|
||||
# Recall budget mapping (budget enum -> thinking_budget integer)
|
||||
ENV_RECALL_BUDGET_FUNCTION = "HINDSIGHT_API_RECALL_BUDGET_FUNCTION"
|
||||
ENV_RECALL_BUDGET_FIXED_LOW = "HINDSIGHT_API_RECALL_BUDGET_FIXED_LOW"
|
||||
ENV_RECALL_BUDGET_FIXED_MID = "HINDSIGHT_API_RECALL_BUDGET_FIXED_MID"
|
||||
ENV_RECALL_BUDGET_FIXED_HIGH = "HINDSIGHT_API_RECALL_BUDGET_FIXED_HIGH"
|
||||
ENV_RECALL_BUDGET_ADAPTIVE_LOW = "HINDSIGHT_API_RECALL_BUDGET_ADAPTIVE_LOW"
|
||||
ENV_RECALL_BUDGET_ADAPTIVE_MID = "HINDSIGHT_API_RECALL_BUDGET_ADAPTIVE_MID"
|
||||
ENV_RECALL_BUDGET_ADAPTIVE_HIGH = "HINDSIGHT_API_RECALL_BUDGET_ADAPTIVE_HIGH"
|
||||
ENV_RECALL_BUDGET_MIN = "HINDSIGHT_API_RECALL_BUDGET_MIN"
|
||||
ENV_RECALL_BUDGET_MAX = "HINDSIGHT_API_RECALL_BUDGET_MAX"
|
||||
|
||||
# Audit log settings
|
||||
ENV_AUDIT_LOG_ENABLED = "HINDSIGHT_API_AUDIT_LOG_ENABLED"
|
||||
@@ -373,6 +442,7 @@ PROVIDER_DEFAULT_MODELS = {
|
||||
"groq": "openai/gpt-oss-120b",
|
||||
"minimax": "MiniMax-M2.7",
|
||||
"ollama": "gemma3:12b",
|
||||
"llamacpp": "gemma-4-e2b-it",
|
||||
"lmstudio": "local-model",
|
||||
"vertexai": "google/gemini-2.5-flash-lite",
|
||||
"openai-codex": "gpt-5.2-codex",
|
||||
@@ -382,10 +452,18 @@ PROVIDER_DEFAULT_MODELS = {
|
||||
"litellm": "gpt-4o-mini",
|
||||
"bedrock": "us.amazon.nova-2-lite-v1:0",
|
||||
"volcano": "doubao-pro-32k",
|
||||
"openrouter": "qwen/qwen3.5-9b",
|
||||
}
|
||||
DEFAULT_LLM_MODEL = "gpt-4o-mini" # Fallback if provider not in table
|
||||
# Built-in llama.cpp defaults
|
||||
DEFAULT_LLAMACPP_GPU_LAYERS = -1 # -1 = offload all layers to GPU (Metal/CUDA)
|
||||
DEFAULT_LLAMACPP_CONTEXT_SIZE = 8192
|
||||
DEFAULT_LLAMACPP_CHAT_FORMAT = None # None = auto-detect from GGUF metadata
|
||||
DEFAULT_LLAMACPP_NO_GRAMMAR = False # True = disable JSON grammar enforcement (faster but less reliable)
|
||||
DEFAULT_LLAMACPP_EXTRA_ARGS = None # Space-separated extra CLI args for llama.cpp server
|
||||
|
||||
DEFAULT_LLM_MAX_CONCURRENT = 32
|
||||
DEFAULT_LLM_MAX_RETRIES = 10 # Max retry attempts for LLM API calls
|
||||
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
|
||||
@@ -403,6 +481,9 @@ DEFAULT_EMBEDDINGS_LOCAL_MODEL = "BAAI/bge-small-en-v1.5"
|
||||
DEFAULT_EMBEDDINGS_LOCAL_FORCE_CPU = False # Force CPU mode for local embeddings (avoids MPS/XPC issues on macOS)
|
||||
DEFAULT_EMBEDDINGS_LOCAL_TRUST_REMOTE_CODE = False # Security: disabled by default, required for some models
|
||||
DEFAULT_EMBEDDINGS_OPENAI_MODEL = "text-embedding-3-small"
|
||||
DEFAULT_EMBEDDINGS_OPENAI_BATCH_SIZE = 100
|
||||
DEFAULT_EMBEDDINGS_GEMINI_MODEL = "gemini-embedding-001"
|
||||
DEFAULT_EMBEDDINGS_GEMINI_OUTPUT_DIMENSIONALITY = 768
|
||||
DEFAULT_EMBEDDING_DIMENSION = 384
|
||||
|
||||
DEFAULT_RERANKER_PROVIDER = "local"
|
||||
@@ -417,6 +498,7 @@ DEFAULT_RERANKER_LOCAL_BUCKET_BATCHING = False # Length-sorted bucket batching:
|
||||
DEFAULT_RERANKER_LOCAL_BATCH_SIZE = 32 # Batch size for local reranker predict() calls
|
||||
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)
|
||||
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
|
||||
@@ -424,8 +506,17 @@ DEFAULT_RERANKER_FLASHRANK_CACHE_DIR = None # Use default cache directory
|
||||
DEFAULT_EMBEDDINGS_COHERE_MODEL = "embed-english-v3.0"
|
||||
DEFAULT_RERANKER_COHERE_MODEL = "rerank-english-v3.0"
|
||||
|
||||
# OpenRouter defaults
|
||||
DEFAULT_EMBEDDINGS_OPENROUTER_MODEL = "perplexity/pplx-embed-v1-0.6b"
|
||||
DEFAULT_RERANKER_OPENROUTER_MODEL = "cohere/rerank-v3.5"
|
||||
|
||||
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_GOOGLE_MODEL = "semantic-ranker-default-004"
|
||||
|
||||
# Vector extension (pgvector, vchord, or pgvectorscale)
|
||||
DEFAULT_VECTOR_EXTENSION = "pgvector" # Options: "pgvector", "vchord", "pgvectorscale"
|
||||
|
||||
@@ -440,6 +531,7 @@ DEFAULT_RERANKER_LITELLM_MAX_TOKENS_PER_DOC: int | None = None
|
||||
|
||||
# LiteLLM SDK defaults
|
||||
DEFAULT_EMBEDDINGS_LITELLM_SDK_MODEL = "cohere/embed-english-v3.0"
|
||||
DEFAULT_EMBEDDINGS_LITELLM_SDK_ENCODING_FORMAT = "float"
|
||||
DEFAULT_RERANKER_LITELLM_SDK_MODEL = "cohere/rerank-english-v3.0"
|
||||
|
||||
DEFAULT_HOST = "0.0.0.0"
|
||||
@@ -452,11 +544,14 @@ 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)
|
||||
DEFAULT_ENABLE_BANK_CONFIG_API = True
|
||||
DEFAULT_DEFAULT_BANK_TEMPLATE: dict | None = None # BankTemplateManifest dict applied to newly-created banks
|
||||
DEFAULT_GRAPH_RETRIEVER = "link_expansion"
|
||||
DEFAULT_RECALL_MAX_CONCURRENT = 32 # Max concurrent recall operations per worker
|
||||
DEFAULT_RECALL_CONNECTION_BUDGET = 4 # Max concurrent DB connections per recall operation
|
||||
DEFAULT_RECALL_MAX_QUERY_TOKENS = 500 # Maximum tokens allowed in recall query
|
||||
DEFAULT_MENTAL_MODEL_REFRESH_CONCURRENCY = 8 # Max concurrent mental model refreshes
|
||||
DEFAULT_LINK_EXPANSION_PER_ENTITY_LIMIT = 200 # Max target units per entity in graph expansion
|
||||
DEFAULT_LINK_EXPANSION_TIMEOUT = 10.0 # Timeout (seconds) for entity expansion query
|
||||
|
||||
# Retain settings
|
||||
DEFAULT_RETAIN_MAX_COMPLETION_TOKENS = 64000 # Max tokens for fact extraction LLM call
|
||||
@@ -489,7 +584,11 @@ DEFAULT_FILE_DELETE_AFTER_RETAIN = True # Delete file bytes after retain (saves
|
||||
DEFAULT_ENABLE_OBSERVATIONS = True # Observations 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
|
||||
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_MAX_TOKENS = 512 # Max tokens for recall when finding related observations
|
||||
DEFAULT_CONSOLIDATION_SOURCE_FACTS_MAX_TOKENS = (
|
||||
@@ -509,6 +608,7 @@ DEFAULT_DB_POOL_MIN_SIZE = 5
|
||||
DEFAULT_DB_POOL_MAX_SIZE = 100
|
||||
DEFAULT_DB_COMMAND_TIMEOUT = 60 # seconds
|
||||
DEFAULT_DB_ACQUIRE_TIMEOUT = 30 # seconds
|
||||
DEFAULT_DB_STATEMENT_TIMEOUT = 600 # seconds (Postgres statement_timeout applied on every pool connection; 0 disables)
|
||||
|
||||
# Worker configuration (distributed task processing)
|
||||
DEFAULT_WORKER_ENABLED = True # API runs worker by default (standalone mode)
|
||||
@@ -517,7 +617,6 @@ DEFAULT_WORKER_POLL_INTERVAL_MS = 500 # Poll database every 500ms
|
||||
DEFAULT_WORKER_MAX_RETRIES = 3 # Max retries before marking task failed
|
||||
DEFAULT_WORKER_HTTP_PORT = 8889 # HTTP port for worker metrics/health
|
||||
DEFAULT_WORKER_MAX_SLOTS = 10 # Total concurrent tasks per worker
|
||||
DEFAULT_WORKER_CONSOLIDATION_MAX_SLOTS = 2 # Max concurrent consolidation tasks per worker
|
||||
DEFAULT_RETAIN_MAX_CONCURRENT = 4 # Max concurrent retain DB phases (HNSW reads + writes). Limits I/O contention.
|
||||
|
||||
# Reflect agent settings
|
||||
@@ -525,6 +624,25 @@ DEFAULT_REFLECT_MAX_ITERATIONS = 10 # Max tool call iterations before forcing r
|
||||
DEFAULT_REFLECT_MAX_CONTEXT_TOKENS = 100_000 # Max accumulated context tokens before forcing final prompt
|
||||
DEFAULT_REFLECT_WALL_TIMEOUT = 300 # Wall-clock timeout in seconds for the entire reflect operation (5 minutes)
|
||||
DEFAULT_REFLECT_SOURCE_FACTS_MAX_TOKENS = -1 # Token budget for source facts in search_observations (-1 = disabled)
|
||||
DEFAULT_RECALL_INCLUDE_CHUNKS = True # Whether internal recall (e.g. mental model refresh) returns raw chunks
|
||||
DEFAULT_RECALL_MAX_TOKENS = 2048 # Token budget for facts returned by internal recall
|
||||
DEFAULT_RECALL_CHUNKS_MAX_TOKENS = 1000 # Token budget for raw chunks returned by internal recall
|
||||
|
||||
# Recall budget mapping
|
||||
# "fixed": thinking_budget = recall_budget_fixed_<level> (preserves legacy behavior)
|
||||
# "adaptive": thinking_budget = round(max_tokens * recall_budget_adaptive_<level>),
|
||||
# clamped to [recall_budget_min, recall_budget_max]
|
||||
RECALL_BUDGET_FUNCTIONS = ("fixed", "adaptive")
|
||||
DEFAULT_RECALL_BUDGET_FUNCTION = "fixed"
|
||||
DEFAULT_RECALL_BUDGET_FIXED_LOW = 100
|
||||
DEFAULT_RECALL_BUDGET_FIXED_MID = 300
|
||||
DEFAULT_RECALL_BUDGET_FIXED_HIGH = 1000
|
||||
# Adaptive defaults chosen to roughly match fixed defaults at max_tokens=4096
|
||||
DEFAULT_RECALL_BUDGET_ADAPTIVE_LOW = 0.025
|
||||
DEFAULT_RECALL_BUDGET_ADAPTIVE_MID = 0.075
|
||||
DEFAULT_RECALL_BUDGET_ADAPTIVE_HIGH = 0.25
|
||||
DEFAULT_RECALL_BUDGET_MIN = 20 # Floor for the adaptive function
|
||||
DEFAULT_RECALL_BUDGET_MAX = 2000 # Ceiling for the adaptive function
|
||||
|
||||
# Disposition defaults (None = not set, fall back to bank DB value or 3)
|
||||
DEFAULT_DISPOSITION_SKEPTICISM = None
|
||||
@@ -535,6 +653,7 @@ DEFAULT_DISPOSITION_EMPATHY = None
|
||||
DEFAULT_OTEL_TRACES_ENABLED = False # Disabled by default for backward compatibility
|
||||
DEFAULT_OTEL_SERVICE_NAME = "hindsight-api"
|
||||
DEFAULT_OTEL_DEPLOYMENT_ENVIRONMENT = "development"
|
||||
DEFAULT_METRICS_INCLUDE_BANK_ID = False # Disabled by default to avoid high-cardinality OTel metric growth
|
||||
|
||||
# Audit log defaults
|
||||
DEFAULT_AUDIT_LOG_ENABLED = False # Disabled by default
|
||||
@@ -586,6 +705,10 @@ class JsonFormatter(logging.Formatter):
|
||||
logging.CRITICAL: "CRITICAL",
|
||||
}
|
||||
|
||||
def __init__(self, allowed_fields: frozenset[str] | None = None):
|
||||
super().__init__()
|
||||
self._allowed_fields = allowed_fields
|
||||
|
||||
def format(self, record: logging.LogRecord) -> str:
|
||||
log_entry = {
|
||||
"severity": self.SEVERITY_MAP.get(record.levelno, "DEFAULT"),
|
||||
@@ -594,10 +717,20 @@ class JsonFormatter(logging.Formatter):
|
||||
"logger": record.name,
|
||||
}
|
||||
|
||||
# Lazy import to avoid circular dependency (engine imports from config).
|
||||
from hindsight_api.engine.memory_engine import _current_schema
|
||||
|
||||
tenant = _current_schema.get()
|
||||
if tenant:
|
||||
log_entry["tenant"] = tenant
|
||||
|
||||
# Add exception info if present
|
||||
if record.exc_info:
|
||||
log_entry["exception"] = self.formatException(record.exc_info)
|
||||
|
||||
if self._allowed_fields is not None:
|
||||
log_entry = {k: v for k, v in log_entry.items() if k in self._allowed_fields}
|
||||
|
||||
return json.dumps(log_entry)
|
||||
|
||||
|
||||
@@ -606,6 +739,25 @@ def _parse_str_list(value: str) -> list[str]:
|
||||
return [v.strip() for v in value.split(",") if v.strip()]
|
||||
|
||||
|
||||
def _parse_positive_int(name: str, raw: str | None, default: int) -> int:
|
||||
"""
|
||||
Parse an env var that must be a positive integer (>= 1).
|
||||
|
||||
Falls back to ``default`` when unset/empty. Raises ValueError on non-integer
|
||||
or non-positive values so misconfiguration fails fast instead of triggering
|
||||
infinite loops or zero-step range() calls downstream.
|
||||
"""
|
||||
if raw is None or raw == "":
|
||||
return default
|
||||
try:
|
||||
parsed = int(raw)
|
||||
except ValueError as e:
|
||||
raise ValueError(f"{name} must be an integer, got {raw!r}") from e
|
||||
if parsed < 1:
|
||||
raise ValueError(f"{name} must be >= 1, got {parsed}")
|
||||
return parsed
|
||||
|
||||
|
||||
def _validate_extraction_mode(mode: str) -> str:
|
||||
"""Validate and normalize extraction mode."""
|
||||
mode_lower = mode.lower()
|
||||
@@ -618,11 +770,43 @@ def _validate_extraction_mode(mode: str) -> str:
|
||||
return mode_lower
|
||||
|
||||
|
||||
def _validate_recall_budget_function(function: str) -> str:
|
||||
"""Validate and normalize recall budget function."""
|
||||
function_lower = function.lower()
|
||||
if function_lower not in RECALL_BUDGET_FUNCTIONS:
|
||||
logger.warning(
|
||||
f"Invalid recall budget function '{function}', must be one of {RECALL_BUDGET_FUNCTIONS}. "
|
||||
f"Defaulting to '{DEFAULT_RECALL_BUDGET_FUNCTION}'."
|
||||
)
|
||||
return DEFAULT_RECALL_BUDGET_FUNCTION
|
||||
return function_lower
|
||||
|
||||
|
||||
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)
|
||||
|
||||
|
||||
def _parse_default_bank_template(raw: str | None) -> dict | None:
|
||||
"""
|
||||
Parse HINDSIGHT_API_DEFAULT_BANK_TEMPLATE as JSON.
|
||||
|
||||
The env var holds a BankTemplateManifest (JSON object) applied verbatim to
|
||||
every newly-created bank. Full Pydantic validation is deferred to bank
|
||||
creation time (to avoid pulling API models into config.py), but we fail
|
||||
fast here if the value is not valid JSON or not a JSON object.
|
||||
"""
|
||||
if raw is None or raw.strip() == "":
|
||||
return DEFAULT_DEFAULT_BANK_TEMPLATE
|
||||
try:
|
||||
parsed = json.loads(raw)
|
||||
except json.JSONDecodeError as e:
|
||||
raise ValueError(f"Invalid {ENV_DEFAULT_BANK_TEMPLATE}: expected a JSON object, got invalid JSON: {e}") from e
|
||||
if not isinstance(parsed, dict):
|
||||
raise ValueError(f"Invalid {ENV_DEFAULT_BANK_TEMPLATE}: expected a JSON object, got {type(parsed).__name__}")
|
||||
return parsed
|
||||
|
||||
|
||||
@dataclass
|
||||
class HindsightConfig:
|
||||
"""Configuration container for Hindsight API."""
|
||||
@@ -658,6 +842,14 @@ class HindsightConfig:
|
||||
# Gemini safety settings (None = use Gemini defaults; list of dicts with category/threshold)
|
||||
llm_gemini_safety_settings: list | None
|
||||
|
||||
# Built-in llama.cpp configuration (for provider=llamacpp)
|
||||
llamacpp_model_path: str | None # Path to GGUF file (None = auto-download default)
|
||||
llamacpp_gpu_layers: int # -1 = all layers on GPU, 0 = CPU only
|
||||
llamacpp_context_size: int # Context window size
|
||||
llamacpp_chat_format: str | None # Chat template format (None = auto-detect from GGUF)
|
||||
llamacpp_no_grammar: bool # Disable JSON grammar enforcement (faster, less reliable)
|
||||
llamacpp_extra_args: str | None # Space-separated extra CLI args for llama.cpp server
|
||||
|
||||
# Per-operation LLM configuration (None = use default LLM config)
|
||||
retain_llm_provider: str | None
|
||||
retain_llm_api_key: str | None
|
||||
@@ -699,6 +891,8 @@ class HindsightConfig:
|
||||
embeddings_cohere_api_key: str | None
|
||||
embeddings_cohere_model: str
|
||||
embeddings_cohere_base_url: str | None
|
||||
embeddings_openrouter_api_key: str | None
|
||||
embeddings_openrouter_model: str
|
||||
embeddings_litellm_api_base: str
|
||||
embeddings_litellm_api_key: str | None
|
||||
embeddings_litellm_model: str
|
||||
@@ -706,6 +900,14 @@ class HindsightConfig:
|
||||
embeddings_litellm_sdk_model: str
|
||||
embeddings_litellm_sdk_api_base: str | None
|
||||
embeddings_litellm_sdk_output_dimensions: int | None
|
||||
embeddings_litellm_sdk_encoding_format: str | None
|
||||
# Gemini/Vertex AI embeddings
|
||||
embeddings_gemini_api_key: str | None
|
||||
embeddings_gemini_model: str
|
||||
embeddings_gemini_output_dimensionality: int | None
|
||||
embeddings_vertexai_project_id: str | None
|
||||
embeddings_vertexai_region: str | None
|
||||
embeddings_vertexai_service_account_key: str | None
|
||||
|
||||
# Reranker
|
||||
reranker_provider: str
|
||||
@@ -719,10 +921,13 @@ class HindsightConfig:
|
||||
reranker_tei_url: str | None
|
||||
reranker_tei_batch_size: int
|
||||
reranker_tei_max_concurrent: int
|
||||
reranker_tei_http_timeout: float
|
||||
reranker_max_candidates: int
|
||||
reranker_cohere_api_key: str | None
|
||||
reranker_cohere_model: str
|
||||
reranker_cohere_base_url: str | None
|
||||
reranker_openrouter_api_key: str | None
|
||||
reranker_openrouter_model: str
|
||||
reranker_litellm_api_base: str
|
||||
reranker_litellm_api_key: str | None
|
||||
reranker_litellm_model: str
|
||||
@@ -733,6 +938,12 @@ class HindsightConfig:
|
||||
reranker_zeroentropy_api_key: str | None
|
||||
reranker_zeroentropy_model: str
|
||||
reranker_zeroentropy_base_url: str | None
|
||||
reranker_siliconflow_api_key: str | None
|
||||
reranker_siliconflow_model: str
|
||||
reranker_siliconflow_base_url: str
|
||||
reranker_google_model: str
|
||||
reranker_google_project_id: str | None
|
||||
reranker_google_service_account_key: str | None
|
||||
|
||||
# Server
|
||||
host: str
|
||||
@@ -740,10 +951,14 @@ class HindsightConfig:
|
||||
base_path: str
|
||||
log_level: str
|
||||
log_format: str
|
||||
log_json_fields: list[str] | None # None = all fields; explicit list = allowlist
|
||||
mcp_enabled: bool
|
||||
mcp_enabled_tools: list[str] | None # None = all tools; explicit list = allowlist
|
||||
mcp_stateless: bool # True = stateless HTTP (POST-only); False = stateful (supports GET/SSE)
|
||||
enable_bank_config_api: bool
|
||||
# Default bank template (static, server-level only). When set, the manifest is applied
|
||||
# to every newly-created bank, overriding the env/config defaults for any fields it sets.
|
||||
default_bank_template: dict | None
|
||||
|
||||
# Recall
|
||||
graph_retriever: str
|
||||
@@ -751,6 +966,8 @@ class HindsightConfig:
|
||||
recall_connection_budget: int
|
||||
recall_max_query_tokens: int
|
||||
mental_model_refresh_concurrency: int
|
||||
link_expansion_per_entity_limit: int
|
||||
link_expansion_timeout: float
|
||||
|
||||
# Retain settings
|
||||
retain_max_completion_tokens: int
|
||||
@@ -793,10 +1010,12 @@ class HindsightConfig:
|
||||
enable_observation_history: bool
|
||||
enable_mental_model_history: bool
|
||||
consolidation_batch_size: int
|
||||
consolidation_max_memories_per_round: int
|
||||
consolidation_llm_batch_size: int
|
||||
consolidation_max_tokens: int
|
||||
consolidation_source_facts_max_tokens: int
|
||||
consolidation_source_facts_max_tokens_per_observation: int
|
||||
consolidation_max_attempts: int
|
||||
observations_mission: str | None
|
||||
max_observations_per_scope: int
|
||||
|
||||
@@ -811,6 +1030,25 @@ class HindsightConfig:
|
||||
reflect_mission: str | None
|
||||
reflect_source_facts_max_tokens: int
|
||||
|
||||
# Recall settings (used by internal recall, e.g. during mental model refresh)
|
||||
recall_include_chunks: bool
|
||||
recall_max_tokens: int
|
||||
recall_chunks_max_tokens: int
|
||||
|
||||
# Recall budget mapping: how the Budget enum (LOW/MID/HIGH) maps to thinking_budget integer.
|
||||
# function="fixed": use the recall_budget_fixed_* values directly (legacy behavior).
|
||||
# function="adaptive": compute round(max_tokens * recall_budget_adaptive_*),
|
||||
# clamped to [recall_budget_min, recall_budget_max].
|
||||
recall_budget_function: str
|
||||
recall_budget_fixed_low: int
|
||||
recall_budget_fixed_mid: int
|
||||
recall_budget_fixed_high: int
|
||||
recall_budget_adaptive_low: float
|
||||
recall_budget_adaptive_mid: float
|
||||
recall_budget_adaptive_high: float
|
||||
recall_budget_min: int
|
||||
recall_budget_max: int
|
||||
|
||||
# Disposition settings (hierarchical - can be overridden per bank; None = fall back to DB)
|
||||
disposition_skepticism: int | None
|
||||
disposition_literalism: int | None
|
||||
@@ -828,6 +1066,7 @@ class HindsightConfig:
|
||||
db_pool_max_size: int
|
||||
db_command_timeout: int
|
||||
db_acquire_timeout: int
|
||||
db_statement_timeout: int
|
||||
|
||||
# Worker configuration (distributed task processing)
|
||||
worker_enabled: bool
|
||||
@@ -836,7 +1075,7 @@ class HindsightConfig:
|
||||
worker_max_retries: int
|
||||
worker_http_port: int
|
||||
worker_max_slots: int
|
||||
worker_consolidation_max_slots: int
|
||||
worker_slot_reservations: dict[str, int]
|
||||
retain_max_concurrent: int
|
||||
|
||||
# Reflect agent settings
|
||||
@@ -850,6 +1089,7 @@ class HindsightConfig:
|
||||
otel_exporter_otlp_headers: str | None
|
||||
otel_service_name: str
|
||||
otel_deployment_environment: str
|
||||
metrics_include_bank_id: bool
|
||||
|
||||
# Audit log configuration (static - server-level only)
|
||||
audit_log_enabled: bool # Master switch for audit logging
|
||||
@@ -862,6 +1102,10 @@ class HindsightConfig:
|
||||
webhook_event_types: list[str] # Event types to deliver globally
|
||||
webhook_delivery_poll_interval_seconds: int # How often the delivery worker polls
|
||||
|
||||
# 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
|
||||
|
||||
# Class-level sets for configuration categorization
|
||||
|
||||
# CREDENTIAL_FIELDS: Never exposed via API, never configurable per-tenant/bank
|
||||
@@ -880,8 +1124,13 @@ class HindsightConfig:
|
||||
"reranker_tei_base_url",
|
||||
"reranker_cohere_base_url",
|
||||
"reranker_zeroentropy_base_url",
|
||||
"reranker_siliconflow_base_url",
|
||||
# Service Account Keys
|
||||
"llm_vertexai_service_account_key",
|
||||
"embeddings_vertexai_service_account_key",
|
||||
"reranker_google_service_account_key",
|
||||
# Embeddings API keys
|
||||
"embeddings_gemini_api_key",
|
||||
# File storage credentials
|
||||
"file_storage_s3_access_key_id",
|
||||
"file_storage_s3_secret_access_key",
|
||||
@@ -911,6 +1160,7 @@ class HindsightConfig:
|
||||
# Consolidation settings
|
||||
"enable_observations",
|
||||
"consolidation_llm_batch_size",
|
||||
"consolidation_max_memories_per_round",
|
||||
"consolidation_source_facts_max_tokens",
|
||||
"consolidation_source_facts_max_tokens_per_observation",
|
||||
"observations_mission",
|
||||
@@ -918,6 +1168,20 @@ class HindsightConfig:
|
||||
# Reflect settings
|
||||
"reflect_mission",
|
||||
"reflect_source_facts_max_tokens",
|
||||
# Recall settings (used by internal recall, e.g. mental model refresh)
|
||||
"recall_include_chunks",
|
||||
"recall_max_tokens",
|
||||
"recall_chunks_max_tokens",
|
||||
# Recall budget mapping (Budget enum -> thinking_budget integer)
|
||||
"recall_budget_function",
|
||||
"recall_budget_fixed_low",
|
||||
"recall_budget_fixed_mid",
|
||||
"recall_budget_fixed_high",
|
||||
"recall_budget_adaptive_low",
|
||||
"recall_budget_adaptive_mid",
|
||||
"recall_budget_adaptive_high",
|
||||
"recall_budget_min",
|
||||
"recall_budget_max",
|
||||
# Disposition settings
|
||||
"disposition_skepticism",
|
||||
"disposition_literalism",
|
||||
@@ -1024,6 +1288,16 @@ class HindsightConfig:
|
||||
f"provider: {self.retain_llm_provider or self.llm_provider})"
|
||||
)
|
||||
|
||||
# Validate that sum of per-operation slot reservations does not exceed max_slots
|
||||
total_reserved = sum(self.worker_slot_reservations.values())
|
||||
if total_reserved > self.worker_max_slots:
|
||||
reservation_details = ", ".join(f"{k}={v}" for k, v in self.worker_slot_reservations.items() if v > 0)
|
||||
raise ValueError(
|
||||
f"Sum of per-operation slot reservations ({total_reserved}: {reservation_details}) "
|
||||
f"exceeds worker_max_slots ({self.worker_max_slots}). "
|
||||
f"Reduce reservations or increase HINDSIGHT_API_WORKER_MAX_SLOTS."
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_env(cls) -> "HindsightConfig":
|
||||
"""Create configuration from environment variables."""
|
||||
@@ -1058,6 +1332,14 @@ class HindsightConfig:
|
||||
or DEFAULT_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY,
|
||||
# Gemini safety settings (JSON-encoded list of {category, threshold} dicts)
|
||||
llm_gemini_safety_settings=json.loads(os.getenv(ENV_LLM_GEMINI_SAFETY_SETTINGS, "null")),
|
||||
# Built-in llama.cpp configuration
|
||||
llamacpp_model_path=os.getenv(ENV_LLAMACPP_MODEL_PATH) or None,
|
||||
llamacpp_gpu_layers=int(os.getenv(ENV_LLAMACPP_GPU_LAYERS, str(DEFAULT_LLAMACPP_GPU_LAYERS))),
|
||||
llamacpp_context_size=int(os.getenv(ENV_LLAMACPP_CONTEXT_SIZE, str(DEFAULT_LLAMACPP_CONTEXT_SIZE))),
|
||||
llamacpp_chat_format=os.getenv(ENV_LLAMACPP_CHAT_FORMAT) or DEFAULT_LLAMACPP_CHAT_FORMAT,
|
||||
llamacpp_no_grammar=os.getenv(ENV_LLAMACPP_NO_GRAMMAR, str(DEFAULT_LLAMACPP_NO_GRAMMAR)).lower()
|
||||
in ("true", "1"),
|
||||
llamacpp_extra_args=os.getenv(ENV_LLAMACPP_EXTRA_ARGS) or DEFAULT_LLAMACPP_EXTRA_ARGS,
|
||||
# Per-operation LLM config (None = use default)
|
||||
retain_llm_provider=os.getenv(ENV_RETAIN_LLM_PROVIDER) or None,
|
||||
retain_llm_api_key=os.getenv(ENV_RETAIN_LLM_API_KEY) or None,
|
||||
@@ -1142,10 +1424,20 @@ class HindsightConfig:
|
||||
in ("true", "1"),
|
||||
embeddings_tei_url=os.getenv(ENV_EMBEDDINGS_TEI_URL),
|
||||
embeddings_openai_base_url=os.getenv(ENV_EMBEDDINGS_OPENAI_BASE_URL) or None,
|
||||
embeddings_openai_batch_size=_parse_positive_int(
|
||||
ENV_EMBEDDINGS_OPENAI_BATCH_SIZE,
|
||||
os.getenv(ENV_EMBEDDINGS_OPENAI_BATCH_SIZE),
|
||||
DEFAULT_EMBEDDINGS_OPENAI_BATCH_SIZE,
|
||||
),
|
||||
# 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),
|
||||
embeddings_cohere_base_url=os.getenv(ENV_EMBEDDINGS_COHERE_BASE_URL) or None,
|
||||
# OpenRouter embeddings (with fallback to shared OpenRouter key, then LLM key)
|
||||
embeddings_openrouter_api_key=os.getenv(ENV_EMBEDDINGS_OPENROUTER_API_KEY)
|
||||
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),
|
||||
# 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),
|
||||
@@ -1160,6 +1452,23 @@ class HindsightConfig:
|
||||
embeddings_litellm_sdk_output_dimensions=int(v)
|
||||
if (v := os.getenv(ENV_EMBEDDINGS_LITELLM_SDK_OUTPUT_DIMENSIONS))
|
||||
else None,
|
||||
embeddings_litellm_sdk_encoding_format=os.getenv(
|
||||
ENV_EMBEDDINGS_LITELLM_SDK_ENCODING_FORMAT, DEFAULT_EMBEDDINGS_LITELLM_SDK_ENCODING_FORMAT
|
||||
),
|
||||
# Gemini/Vertex AI embeddings (with fallback to LLM keys)
|
||||
embeddings_gemini_api_key=os.getenv(ENV_EMBEDDINGS_GEMINI_API_KEY) or os.getenv(ENV_LLM_API_KEY),
|
||||
embeddings_gemini_model=os.getenv(ENV_EMBEDDINGS_GEMINI_MODEL, DEFAULT_EMBEDDINGS_GEMINI_MODEL),
|
||||
embeddings_gemini_output_dimensionality=int(
|
||||
os.getenv(
|
||||
ENV_EMBEDDINGS_GEMINI_OUTPUT_DIMENSIONALITY,
|
||||
str(DEFAULT_EMBEDDINGS_GEMINI_OUTPUT_DIMENSIONALITY),
|
||||
)
|
||||
),
|
||||
embeddings_vertexai_project_id=os.getenv(ENV_EMBEDDINGS_VERTEXAI_PROJECT_ID)
|
||||
or os.getenv(ENV_LLM_VERTEXAI_PROJECT_ID),
|
||||
embeddings_vertexai_region=os.getenv(ENV_EMBEDDINGS_VERTEXAI_REGION) or os.getenv(ENV_LLM_VERTEXAI_REGION),
|
||||
embeddings_vertexai_service_account_key=os.getenv(ENV_EMBEDDINGS_VERTEXAI_SERVICE_ACCOUNT_KEY)
|
||||
or os.getenv(ENV_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY),
|
||||
# Reranker
|
||||
reranker_provider=os.getenv(ENV_RERANKER_PROVIDER, DEFAULT_RERANKER_PROVIDER),
|
||||
reranker_local_model=os.getenv(ENV_RERANKER_LOCAL_MODEL, DEFAULT_RERANKER_LOCAL_MODEL),
|
||||
@@ -1188,11 +1497,19 @@ class HindsightConfig:
|
||||
reranker_tei_max_concurrent=int(
|
||||
os.getenv(ENV_RERANKER_TEI_MAX_CONCURRENT, str(DEFAULT_RERANKER_TEI_MAX_CONCURRENT))
|
||||
),
|
||||
reranker_tei_http_timeout=float(
|
||||
os.getenv(ENV_RERANKER_TEI_HTTP_TIMEOUT, str(DEFAULT_RERANKER_TEI_HTTP_TIMEOUT))
|
||||
),
|
||||
reranker_max_candidates=int(os.getenv(ENV_RERANKER_MAX_CANDIDATES, str(DEFAULT_RERANKER_MAX_CANDIDATES))),
|
||||
# Cohere reranker (with backward-compatible fallback to shared API key)
|
||||
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,
|
||||
# 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),
|
||||
# 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),
|
||||
@@ -1209,12 +1526,25 @@ class HindsightConfig:
|
||||
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,
|
||||
# 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
|
||||
),
|
||||
# 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),
|
||||
# Server
|
||||
host=os.getenv(ENV_HOST, DEFAULT_HOST),
|
||||
port=int(os.getenv(ENV_PORT, DEFAULT_PORT)),
|
||||
base_path=os.getenv(ENV_BASE_PATH, DEFAULT_BASE_PATH),
|
||||
log_level=os.getenv(ENV_LOG_LEVEL, DEFAULT_LOG_LEVEL),
|
||||
log_format=os.getenv(ENV_LOG_FORMAT, DEFAULT_LOG_FORMAT).lower(),
|
||||
log_json_fields=_parse_str_list(os.getenv(ENV_LOG_JSON_FIELDS, "")) or None,
|
||||
mcp_enabled=os.getenv(ENV_MCP_ENABLED, str(DEFAULT_MCP_ENABLED)).lower() == "true",
|
||||
mcp_enabled_tools=[t.strip() for t in os.getenv(ENV_MCP_ENABLED_TOOLS).split(",") if t.strip()]
|
||||
if os.getenv(ENV_MCP_ENABLED_TOOLS)
|
||||
@@ -1222,6 +1552,7 @@ class HindsightConfig:
|
||||
mcp_stateless=os.getenv(ENV_MCP_STATELESS, str(DEFAULT_MCP_STATELESS)).lower() == "true",
|
||||
enable_bank_config_api=os.getenv(ENV_ENABLE_BANK_CONFIG_API, str(DEFAULT_ENABLE_BANK_CONFIG_API)).lower()
|
||||
== "true",
|
||||
default_bank_template=_parse_default_bank_template(os.getenv(ENV_DEFAULT_BANK_TEMPLATE)),
|
||||
# Recall
|
||||
graph_retriever=os.getenv(ENV_GRAPH_RETRIEVER, DEFAULT_GRAPH_RETRIEVER),
|
||||
recall_max_concurrent=int(os.getenv(ENV_RECALL_MAX_CONCURRENT, str(DEFAULT_RECALL_MAX_CONCURRENT))),
|
||||
@@ -1232,6 +1563,10 @@ class HindsightConfig:
|
||||
mental_model_refresh_concurrency=int(
|
||||
os.getenv(ENV_MENTAL_MODEL_REFRESH_CONCURRENCY, str(DEFAULT_MENTAL_MODEL_REFRESH_CONCURRENCY))
|
||||
),
|
||||
link_expansion_per_entity_limit=int(
|
||||
os.getenv(ENV_LINK_EXPANSION_PER_ENTITY_LIMIT, str(DEFAULT_LINK_EXPANSION_PER_ENTITY_LIMIT))
|
||||
),
|
||||
link_expansion_timeout=float(os.getenv(ENV_LINK_EXPANSION_TIMEOUT, str(DEFAULT_LINK_EXPANSION_TIMEOUT))),
|
||||
# Optimization flags
|
||||
skip_llm_verification=os.getenv(ENV_SKIP_LLM_VERIFICATION, "false").lower() == "true",
|
||||
lazy_reranker=os.getenv(ENV_LAZY_RERANKER, "false").lower() == "true",
|
||||
@@ -1302,6 +1637,12 @@ class HindsightConfig:
|
||||
consolidation_batch_size=int(
|
||||
os.getenv(ENV_CONSOLIDATION_BATCH_SIZE, str(DEFAULT_CONSOLIDATION_BATCH_SIZE))
|
||||
),
|
||||
consolidation_max_memories_per_round=int(
|
||||
os.getenv(
|
||||
ENV_CONSOLIDATION_MAX_MEMORIES_PER_ROUND,
|
||||
str(DEFAULT_CONSOLIDATION_MAX_MEMORIES_PER_ROUND),
|
||||
)
|
||||
),
|
||||
consolidation_llm_batch_size=int(
|
||||
os.getenv(ENV_CONSOLIDATION_LLM_BATCH_SIZE, str(DEFAULT_CONSOLIDATION_LLM_BATCH_SIZE))
|
||||
),
|
||||
@@ -1317,6 +1658,9 @@ class HindsightConfig:
|
||||
str(DEFAULT_CONSOLIDATION_SOURCE_FACTS_MAX_TOKENS_PER_OBSERVATION),
|
||||
)
|
||||
),
|
||||
consolidation_max_attempts=int(
|
||||
os.getenv(ENV_CONSOLIDATION_MAX_ATTEMPTS, str(DEFAULT_CONSOLIDATION_MAX_ATTEMPTS))
|
||||
),
|
||||
observations_mission=os.getenv(ENV_OBSERVATIONS_MISSION) or DEFAULT_OBSERVATIONS_MISSION,
|
||||
max_observations_per_scope=int(
|
||||
os.getenv(ENV_MAX_OBSERVATIONS_PER_SCOPE, str(DEFAULT_MAX_OBSERVATIONS_PER_SCOPE))
|
||||
@@ -1330,6 +1674,7 @@ class HindsightConfig:
|
||||
db_pool_max_size=int(os.getenv(ENV_DB_POOL_MAX_SIZE, str(DEFAULT_DB_POOL_MAX_SIZE))),
|
||||
db_command_timeout=int(os.getenv(ENV_DB_COMMAND_TIMEOUT, str(DEFAULT_DB_COMMAND_TIMEOUT))),
|
||||
db_acquire_timeout=int(os.getenv(ENV_DB_ACQUIRE_TIMEOUT, str(DEFAULT_DB_ACQUIRE_TIMEOUT))),
|
||||
db_statement_timeout=int(os.getenv(ENV_DB_STATEMENT_TIMEOUT, str(DEFAULT_DB_STATEMENT_TIMEOUT))),
|
||||
# Worker configuration
|
||||
worker_enabled=os.getenv(ENV_WORKER_ENABLED, str(DEFAULT_WORKER_ENABLED)).lower() == "true",
|
||||
worker_id=os.getenv(ENV_WORKER_ID) or DEFAULT_WORKER_ID,
|
||||
@@ -1337,9 +1682,11 @@ class HindsightConfig:
|
||||
worker_max_retries=int(os.getenv(ENV_WORKER_MAX_RETRIES, str(DEFAULT_WORKER_MAX_RETRIES))),
|
||||
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_consolidation_max_slots=int(
|
||||
os.getenv(ENV_WORKER_CONSOLIDATION_MAX_SLOTS, str(DEFAULT_WORKER_CONSOLIDATION_MAX_SLOTS))
|
||||
),
|
||||
worker_slot_reservations={
|
||||
op_type: int(os.getenv(env_var, str(default)))
|
||||
for op_type, (env_var, default) in WORKER_SLOT_RESERVATION_TYPES.items()
|
||||
if int(os.getenv(env_var, str(default))) > 0
|
||||
},
|
||||
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))),
|
||||
@@ -1351,6 +1698,31 @@ class HindsightConfig:
|
||||
reflect_source_facts_max_tokens=int(
|
||||
os.getenv(ENV_REFLECT_SOURCE_FACTS_MAX_TOKENS, str(DEFAULT_REFLECT_SOURCE_FACTS_MAX_TOKENS))
|
||||
),
|
||||
recall_include_chunks=os.getenv(ENV_RECALL_INCLUDE_CHUNKS, str(DEFAULT_RECALL_INCLUDE_CHUNKS)).lower()
|
||||
in ("true", "1", "yes"),
|
||||
recall_max_tokens=int(os.getenv(ENV_RECALL_MAX_TOKENS, str(DEFAULT_RECALL_MAX_TOKENS))),
|
||||
recall_chunks_max_tokens=int(
|
||||
os.getenv(ENV_RECALL_CHUNKS_MAX_TOKENS, str(DEFAULT_RECALL_CHUNKS_MAX_TOKENS))
|
||||
),
|
||||
recall_budget_function=_validate_recall_budget_function(
|
||||
os.getenv(ENV_RECALL_BUDGET_FUNCTION, DEFAULT_RECALL_BUDGET_FUNCTION)
|
||||
),
|
||||
recall_budget_fixed_low=int(os.getenv(ENV_RECALL_BUDGET_FIXED_LOW, str(DEFAULT_RECALL_BUDGET_FIXED_LOW))),
|
||||
recall_budget_fixed_mid=int(os.getenv(ENV_RECALL_BUDGET_FIXED_MID, str(DEFAULT_RECALL_BUDGET_FIXED_MID))),
|
||||
recall_budget_fixed_high=int(
|
||||
os.getenv(ENV_RECALL_BUDGET_FIXED_HIGH, str(DEFAULT_RECALL_BUDGET_FIXED_HIGH))
|
||||
),
|
||||
recall_budget_adaptive_low=float(
|
||||
os.getenv(ENV_RECALL_BUDGET_ADAPTIVE_LOW, str(DEFAULT_RECALL_BUDGET_ADAPTIVE_LOW))
|
||||
),
|
||||
recall_budget_adaptive_mid=float(
|
||||
os.getenv(ENV_RECALL_BUDGET_ADAPTIVE_MID, str(DEFAULT_RECALL_BUDGET_ADAPTIVE_MID))
|
||||
),
|
||||
recall_budget_adaptive_high=float(
|
||||
os.getenv(ENV_RECALL_BUDGET_ADAPTIVE_HIGH, str(DEFAULT_RECALL_BUDGET_ADAPTIVE_HIGH))
|
||||
),
|
||||
recall_budget_min=int(os.getenv(ENV_RECALL_BUDGET_MIN, str(DEFAULT_RECALL_BUDGET_MIN))),
|
||||
recall_budget_max=int(os.getenv(ENV_RECALL_BUDGET_MAX, str(DEFAULT_RECALL_BUDGET_MAX))),
|
||||
# Disposition settings (None = fall back to DB value)
|
||||
disposition_skepticism=int(os.getenv(ENV_DISPOSITION_SKEPTICISM))
|
||||
if os.getenv(ENV_DISPOSITION_SKEPTICISM)
|
||||
@@ -1368,6 +1740,8 @@ class HindsightConfig:
|
||||
otel_exporter_otlp_headers=os.getenv(ENV_OTEL_EXPORTER_OTLP_HEADERS) or None,
|
||||
otel_service_name=os.getenv(ENV_OTEL_SERVICE_NAME, DEFAULT_OTEL_SERVICE_NAME),
|
||||
otel_deployment_environment=os.getenv(ENV_OTEL_DEPLOYMENT_ENVIRONMENT, DEFAULT_OTEL_DEPLOYMENT_ENVIRONMENT),
|
||||
metrics_include_bank_id=os.getenv(ENV_METRICS_INCLUDE_BANK_ID, str(DEFAULT_METRICS_INCLUDE_BANK_ID)).lower()
|
||||
in ("true", "1", "yes"),
|
||||
# Audit log configuration (static, server-level only)
|
||||
audit_log_enabled=os.getenv(ENV_AUDIT_LOG_ENABLED, str(DEFAULT_AUDIT_LOG_ENABLED)).lower() == "true",
|
||||
audit_log_actions=[
|
||||
@@ -1439,7 +1813,8 @@ class HindsightConfig:
|
||||
handler.setLevel(self.get_python_log_level())
|
||||
|
||||
if self.log_format == "json":
|
||||
handler.setFormatter(JsonFormatter())
|
||||
allowed = frozenset(self.log_json_fields) if self.log_json_fields else None
|
||||
handler.setFormatter(JsonFormatter(allowed_fields=allowed))
|
||||
else:
|
||||
handler.setFormatter(logging.Formatter("%(asctime)s - %(levelname)s - %(name)s - %(message)s"))
|
||||
|
||||
|
||||
@@ -15,7 +15,12 @@ from typing import Any
|
||||
|
||||
import asyncpg
|
||||
|
||||
from hindsight_api.config import HindsightConfig, _get_raw_config, normalize_config_dict
|
||||
from hindsight_api.config import (
|
||||
RECALL_BUDGET_FUNCTIONS,
|
||||
HindsightConfig,
|
||||
_get_raw_config,
|
||||
normalize_config_dict,
|
||||
)
|
||||
from hindsight_api.engine.memory_engine import fq_table
|
||||
from hindsight_api.extensions.tenant import TenantExtension
|
||||
from hindsight_api.models import RequestContext
|
||||
@@ -239,6 +244,15 @@ class ConfigResolver:
|
||||
logger.warning(f"Failed to check permissions for bank {bank_id}: {e}")
|
||||
# Continue without permission check (fail open for backward compatibility)
|
||||
|
||||
# Validate entity_labels structure
|
||||
if "entity_labels" in normalized_updates and normalized_updates["entity_labels"] is not None:
|
||||
from .engine.retain.entity_labels import parse_entity_labels
|
||||
|
||||
try:
|
||||
parse_entity_labels(normalized_updates["entity_labels"])
|
||||
except Exception as e:
|
||||
raise ValueError(f"Invalid entity_labels format: {e}")
|
||||
|
||||
# Validate retain_strategies: reject empty string keys
|
||||
if "retain_strategies" in normalized_updates and normalized_updates["retain_strategies"]:
|
||||
empty_keys = [k for k in normalized_updates["retain_strategies"] if not str(k).strip()]
|
||||
@@ -247,6 +261,9 @@ class ConfigResolver:
|
||||
"Strategy names must not be empty strings. Remove entries with empty names before saving."
|
||||
)
|
||||
|
||||
# Validate recall budget fields
|
||||
_validate_recall_budget_updates(normalized_updates)
|
||||
|
||||
# Merge with existing config (JSONB || operator)
|
||||
async with self.pool.acquire() as conn:
|
||||
await conn.execute(
|
||||
@@ -283,6 +300,53 @@ class ConfigResolver:
|
||||
logger.info(f"Reset bank config for {bank_id} to defaults")
|
||||
|
||||
|
||||
_RECALL_BUDGET_FIXED_KEYS = (
|
||||
"recall_budget_fixed_low",
|
||||
"recall_budget_fixed_mid",
|
||||
"recall_budget_fixed_high",
|
||||
)
|
||||
_RECALL_BUDGET_ADAPTIVE_KEYS = (
|
||||
"recall_budget_adaptive_low",
|
||||
"recall_budget_adaptive_mid",
|
||||
"recall_budget_adaptive_high",
|
||||
)
|
||||
|
||||
|
||||
def _validate_recall_budget_updates(updates: dict[str, Any]) -> None:
|
||||
"""Validate recall budget config updates. Raises ValueError on invalid input."""
|
||||
if "recall_budget_function" in updates:
|
||||
function = updates["recall_budget_function"]
|
||||
if not isinstance(function, str) or function.lower() not in RECALL_BUDGET_FUNCTIONS:
|
||||
raise ValueError(
|
||||
f"recall_budget_function must be one of {sorted(RECALL_BUDGET_FUNCTIONS)}, got {function!r}"
|
||||
)
|
||||
|
||||
for key in _RECALL_BUDGET_FIXED_KEYS:
|
||||
if key in updates:
|
||||
value = updates[key]
|
||||
if not isinstance(value, int) or isinstance(value, bool) or value < 1:
|
||||
raise ValueError(f"{key} must be a positive integer, got {value!r}")
|
||||
|
||||
for key in _RECALL_BUDGET_ADAPTIVE_KEYS:
|
||||
if key in updates:
|
||||
value = updates[key]
|
||||
if isinstance(value, bool) or not isinstance(value, (int, float)) or value <= 0:
|
||||
raise ValueError(f"{key} must be a positive number, got {value!r}")
|
||||
|
||||
for key in ("recall_budget_min", "recall_budget_max"):
|
||||
if key in updates:
|
||||
value = updates[key]
|
||||
if not isinstance(value, int) or isinstance(value, bool) or value < 1:
|
||||
raise ValueError(f"{key} must be a positive integer, got {value!r}")
|
||||
|
||||
if "recall_budget_min" in updates and "recall_budget_max" in updates:
|
||||
if updates["recall_budget_min"] > updates["recall_budget_max"]:
|
||||
raise ValueError(
|
||||
f"recall_budget_min ({updates['recall_budget_min']}) must be <= "
|
||||
f"recall_budget_max ({updates['recall_budget_max']})"
|
||||
)
|
||||
|
||||
|
||||
def apply_strategy(config: HindsightConfig, strategy_name: str) -> HindsightConfig:
|
||||
"""
|
||||
Apply a named retain strategy's overrides on top of a resolved config.
|
||||
|
||||
@@ -0,0 +1,58 @@
|
||||
"""Database URL normalization.
|
||||
|
||||
Hindsight accepts SQLAlchemy-style URLs like ``postgresql+asyncpg://...?ssl=require``
|
||||
for its async engine, but the same string cannot be handed directly to synchronous
|
||||
SQLAlchemy (psycopg2) or to :func:`asyncpg.create_pool`, which both expect a
|
||||
libpq-compatible URL (``postgresql://...?sslmode=require``).
|
||||
|
||||
:func:`to_libpq_url` performs that translation. It is idempotent and safe to
|
||||
apply to URLs that are already libpq-compatible, to the ``pg0`` embedded-PG
|
||||
marker, or to any non-PostgreSQL string (returned unchanged).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
|
||||
|
||||
_ASYNCPG_SCHEMES = ("postgresql+asyncpg", "postgres+asyncpg")
|
||||
_POSTGRES_SCHEMES = ("postgresql", "postgres") + _ASYNCPG_SCHEMES
|
||||
|
||||
|
||||
def to_libpq_url(url: str) -> str:
|
||||
"""Normalize a PostgreSQL URL for libpq-style consumers.
|
||||
|
||||
Accepts a SQLAlchemy URL (``postgresql+asyncpg://...``) or a plain libpq
|
||||
URL and returns a form suitable for:
|
||||
|
||||
- :func:`sqlalchemy.create_engine` (sync / psycopg2)
|
||||
- :func:`asyncpg.create_pool`
|
||||
|
||||
Transformations:
|
||||
|
||||
- ``postgresql+asyncpg`` / ``postgres+asyncpg`` / ``postgres`` → ``postgresql``
|
||||
- Query param ``ssl=<mode>`` → ``sslmode=<mode>`` (SQLAlchemy's asyncpg
|
||||
dialect uses ``ssl=``; libpq uses ``sslmode=``)
|
||||
|
||||
Any non-PostgreSQL input (e.g. the ``pg0`` embedded-PG marker, a sqlite
|
||||
URL, an empty string) is returned unchanged. Already-normalized URLs are
|
||||
returned unchanged.
|
||||
"""
|
||||
if not url or "://" not in url:
|
||||
return url
|
||||
|
||||
parts = urlsplit(url)
|
||||
if parts.scheme not in _POSTGRES_SCHEMES:
|
||||
return url
|
||||
|
||||
new_scheme = "postgresql"
|
||||
|
||||
new_query_pairs = [
|
||||
("sslmode", v) if k == "ssl" else (k, v)
|
||||
for k, v in parse_qsl(parts.query, keep_blank_values=True)
|
||||
]
|
||||
new_query = urlencode(new_query_pairs)
|
||||
|
||||
if new_scheme == parts.scheme and new_query == parts.query:
|
||||
return url
|
||||
|
||||
return urlunsplit((new_scheme, parts.netloc, parts.path, new_query, parts.fragment))
|
||||
@@ -42,6 +42,34 @@ if TYPE_CHECKING:
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
async def _filter_live_source_memories(
|
||||
conn: "Connection",
|
||||
bank_id: str,
|
||||
source_memory_ids: list[uuid.UUID],
|
||||
) -> list[uuid.UUID]:
|
||||
"""Return only the source memory ids that still exist in the bank.
|
||||
|
||||
Uses FOR SHARE to block concurrent deletes from removing a row between the
|
||||
check and the subsequent insert/update. Combined with the delete path running
|
||||
its stale-observation sweep *after* deleting the source row, this closes the
|
||||
race window where consolidation would otherwise produce an orphan observation.
|
||||
"""
|
||||
if not source_memory_ids:
|
||||
return []
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT id
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE id = ANY($1::uuid[]) AND bank_id = $2
|
||||
FOR SHARE
|
||||
""",
|
||||
source_memory_ids,
|
||||
bank_id,
|
||||
)
|
||||
live = {row["id"] for row in rows}
|
||||
return [mid for mid in source_memory_ids if mid in live]
|
||||
|
||||
|
||||
class _CreateAction(BaseModel):
|
||||
text: str
|
||||
source_fact_ids: list[str] # memory UUIDs from the NEW FACTS list
|
||||
@@ -219,6 +247,7 @@ async def run_consolidation_job(
|
||||
|
||||
perf = ConsolidationPerfLog(bank_id)
|
||||
max_memories_per_batch = config.consolidation_batch_size
|
||||
max_memories_per_round = config.consolidation_max_memories_per_round
|
||||
llm_batch_size = max(1, config.consolidation_llm_batch_size)
|
||||
|
||||
# Check if consolidation is enabled
|
||||
@@ -281,8 +310,17 @@ async def run_consolidation_job(
|
||||
# Track all unique tags from consolidated memories for mental model refresh filtering
|
||||
consolidated_tags: set[str] = set()
|
||||
|
||||
round_limit_enabled = max_memories_per_round > 0
|
||||
round_remaining = max_memories_per_round if round_limit_enabled else float("inf")
|
||||
hit_round_limit = False
|
||||
|
||||
llm_batch_num = 0
|
||||
while True:
|
||||
# Cap fetch size by remaining round budget
|
||||
fetch_limit = (
|
||||
min(max_memories_per_batch, int(round_remaining)) if round_limit_enabled else max_memories_per_batch
|
||||
)
|
||||
|
||||
# Fetch next batch of unconsolidated memories
|
||||
async with pool.acquire() as conn:
|
||||
t0 = time.time()
|
||||
@@ -299,7 +337,7 @@ async def run_consolidation_job(
|
||||
LIMIT $2
|
||||
""",
|
||||
bank_id,
|
||||
max_memories_per_batch,
|
||||
fetch_limit,
|
||||
)
|
||||
perf.record_timing("fetch_memories", time.time() - t0)
|
||||
|
||||
@@ -524,6 +562,25 @@ async def run_consolidation_job(
|
||||
f" | avg={llm_batch_time / len(llm_batch):.3f}s/memory"
|
||||
)
|
||||
|
||||
# Update round budget after processing this DB fetch batch
|
||||
if round_limit_enabled:
|
||||
round_remaining -= len(memories)
|
||||
if round_remaining <= 0:
|
||||
hit_round_limit = True
|
||||
break
|
||||
|
||||
# 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."
|
||||
)
|
||||
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(
|
||||
f"[3] Results: {stats['memories_processed']} memories -> "
|
||||
@@ -552,16 +609,21 @@ async def run_consolidation_job(
|
||||
if timing_parts:
|
||||
perf.log(f"[4] Timing breakdown: {', '.join(timing_parts)}")
|
||||
|
||||
# Trigger mental model refreshes for models with refresh_after_consolidation=true
|
||||
# SECURITY: Only refresh mental models with matching tags (or all if no tags were consolidated)
|
||||
mental_models_refreshed = await _trigger_mental_model_refreshes(
|
||||
memory_engine=memory_engine,
|
||||
bank_id=bank_id,
|
||||
request_context=request_context,
|
||||
consolidated_tags=list(consolidated_tags) if consolidated_tags else None,
|
||||
perf=perf,
|
||||
)
|
||||
stats["mental_models_refreshed"] = mental_models_refreshed
|
||||
# Trigger mental model refreshes only on the final round (when all memories are processed).
|
||||
# If we hit the round limit and re-queued, skip MM refresh — the next round will handle it.
|
||||
if hit_round_limit:
|
||||
stats["mental_models_refreshed"] = 0
|
||||
logger.info(f"[CONSOLIDATION] bank={bank_id} skipping mental model refresh (round limit hit, re-queued)")
|
||||
else:
|
||||
# SECURITY: Only refresh mental models with matching tags (or all if no tags were consolidated)
|
||||
mental_models_refreshed = await _trigger_mental_model_refreshes(
|
||||
memory_engine=memory_engine,
|
||||
bank_id=bank_id,
|
||||
request_context=request_context,
|
||||
consolidated_tags=list(consolidated_tags) if consolidated_tags else None,
|
||||
perf=perf,
|
||||
)
|
||||
stats["mental_models_refreshed"] = mental_models_refreshed
|
||||
|
||||
perf.flush()
|
||||
|
||||
@@ -593,17 +655,15 @@ async def _trigger_mental_model_refreshes(
|
||||
"""
|
||||
pool = memory_engine._pool
|
||||
|
||||
# Find mental models with refresh_after_consolidation=true
|
||||
# SECURITY: Control which mental models get refreshed based on tags
|
||||
# Find mental models with refresh_after_consolidation=true that are actually stale.
|
||||
# The tag filter on the SELECT enforces the security boundary (never look outside the
|
||||
# relevant tag scope); compute_mental_model_is_stale then verifies that new memories
|
||||
# in the MM's scope really were ingested since its last refresh.
|
||||
async with pool.acquire() as conn:
|
||||
if consolidated_tags:
|
||||
# Tagged memories were consolidated - refresh:
|
||||
# 1. Mental models with overlapping tags (security boundary)
|
||||
# 2. Untagged mental models (they're "global" and available to all contexts)
|
||||
# DO NOT refresh mental models with different tags
|
||||
rows = await conn.fetch(
|
||||
candidates = await conn.fetch(
|
||||
f"""
|
||||
SELECT id, name, tags
|
||||
SELECT id, name, tags, last_refreshed_at, trigger
|
||||
FROM {fq_table("mental_models")}
|
||||
WHERE bank_id = $1
|
||||
AND (trigger->>'refresh_after_consolidation')::boolean = true
|
||||
@@ -616,11 +676,9 @@ async def _trigger_mental_model_refreshes(
|
||||
consolidated_tags,
|
||||
)
|
||||
else:
|
||||
# Untagged memories were consolidated - only refresh untagged mental models
|
||||
# SECURITY: Tagged mental models are NOT refreshed when untagged memories are consolidated
|
||||
rows = await conn.fetch(
|
||||
candidates = await conn.fetch(
|
||||
f"""
|
||||
SELECT id, name, tags
|
||||
SELECT id, name, tags, last_refreshed_at, trigger
|
||||
FROM {fq_table("mental_models")}
|
||||
WHERE bank_id = $1
|
||||
AND (trigger->>'refresh_after_consolidation')::boolean = true
|
||||
@@ -629,6 +687,11 @@ async def _trigger_mental_model_refreshes(
|
||||
bank_id,
|
||||
)
|
||||
|
||||
rows = []
|
||||
for candidate in candidates:
|
||||
if await memory_engine.compute_mental_model_is_stale(conn, bank_id, candidate):
|
||||
rows.append(candidate)
|
||||
|
||||
if not rows:
|
||||
return 0
|
||||
|
||||
@@ -889,6 +952,15 @@ async def _execute_update_action(
|
||||
logger.debug(f"Update skipped: observation {observation_id} not found in recall results")
|
||||
return
|
||||
|
||||
live_source_memory_ids = await _filter_live_source_memories(conn, bank_id, source_memory_ids)
|
||||
if not live_source_memory_ids:
|
||||
logger.debug(
|
||||
f"Update skipped: all {len(source_memory_ids)} source memories for observation "
|
||||
f"{observation_id} were deleted concurrently"
|
||||
)
|
||||
return
|
||||
source_memory_ids = live_source_memory_ids
|
||||
|
||||
from ...config import get_config
|
||||
|
||||
history_entry = {
|
||||
@@ -1131,14 +1203,16 @@ async def _consolidate_batch_with_llm(
|
||||
memories: list[dict[str, Any]],
|
||||
union_observations: "list[MemoryFact]",
|
||||
union_source_facts: "dict[str, MemoryFact]",
|
||||
config: Any = None,
|
||||
config: Any,
|
||||
remaining_observation_slots: int | None = None,
|
||||
max_observations_per_scope: int = -1,
|
||||
) -> _BatchLLMResult:
|
||||
"""Single LLM call for a batch of facts against a pooled set of observations."""
|
||||
if config is None:
|
||||
raise ValueError("config is required for _consolidate_batch_with_llm")
|
||||
if union_observations:
|
||||
obs_list = _build_observations_for_llm(union_observations, union_source_facts)
|
||||
observations_text = json.dumps(obs_list, indent=2)
|
||||
observations_text = json.dumps(obs_list, indent=2, ensure_ascii=False)
|
||||
else:
|
||||
observations_text = "[]"
|
||||
|
||||
@@ -1172,8 +1246,7 @@ async def _consolidate_batch_with_llm(
|
||||
f"(out of {max_observations_per_scope}). Prefer UPDATE over CREATE when possible."
|
||||
)
|
||||
|
||||
observations_mission = config.observations_mission if config is not None else None
|
||||
prompt_template = build_batch_consolidation_prompt(observations_mission, observation_capacity_note)
|
||||
prompt_template = build_batch_consolidation_prompt(config.observations_mission, observation_capacity_note)
|
||||
prompt = prompt_template.format(
|
||||
facts_text=facts_lines,
|
||||
observations_text=observations_text,
|
||||
@@ -1182,15 +1255,29 @@ async def _consolidate_batch_with_llm(
|
||||
# Use a constrained response model when observation limit is active
|
||||
response_model = _build_response_model(max_creates=remaining_observation_slots)
|
||||
|
||||
max_attempts = 3
|
||||
max_attempts = config.consolidation_max_attempts
|
||||
inner_max_retries = config.consolidation_llm_max_retries
|
||||
last_exc: Exception | None = None
|
||||
# Pre-compute a stable identifier set for the batch so failure logs name the
|
||||
# exact memories whose consolidation is failing — without this, an opaque
|
||||
# "LLM batch call failed" line gives operators no way to find the offending
|
||||
# input until adaptive bisection narrows the batch down to a single memory.
|
||||
memory_ids = [str(m.get("id")) for m in memories]
|
||||
if len(memory_ids) <= 5:
|
||||
ids_label = ", ".join(memory_ids)
|
||||
else:
|
||||
ids_label = f"{', '.join(memory_ids[:3])}, ... +{len(memory_ids) - 3} more"
|
||||
batch_label = f"{len(memory_ids)} memories [{ids_label}]"
|
||||
for attempt in range(1, max_attempts + 1):
|
||||
try:
|
||||
response: _ConsolidationBatchResponse = await llm_config.call(
|
||||
messages=[{"role": "user", "content": prompt}],
|
||||
response_format=response_model,
|
||||
scope="consolidation",
|
||||
)
|
||||
call_kwargs: dict[str, Any] = {
|
||||
"messages": [{"role": "user", "content": prompt}],
|
||||
"response_format": response_model,
|
||||
"scope": "consolidation",
|
||||
}
|
||||
if inner_max_retries is not None:
|
||||
call_kwargs["max_retries"] = inner_max_retries
|
||||
response: _ConsolidationBatchResponse = await llm_config.call(**call_kwargs)
|
||||
# Defensive truncation: some LLM providers may not enforce JSON schema max_length
|
||||
creates = response.creates
|
||||
if remaining_observation_slots is not None and remaining_observation_slots >= 0:
|
||||
@@ -1209,10 +1296,13 @@ async def _consolidate_batch_with_llm(
|
||||
)
|
||||
except Exception as exc:
|
||||
last_exc = exc
|
||||
logger.warning(f"[CONSOLIDATION] LLM batch call failed (attempt {attempt}/{max_attempts}): {exc}")
|
||||
logger.warning(
|
||||
f"[CONSOLIDATION] LLM batch call failed (attempt {attempt}/{max_attempts}) for {batch_label}: {exc}"
|
||||
)
|
||||
|
||||
logger.error(
|
||||
f"[CONSOLIDATION] LLM batch call failed after {max_attempts} attempts, skipping batch. Last error: {last_exc}"
|
||||
f"[CONSOLIDATION] LLM batch call failed after {max_attempts} attempts for {batch_label}, "
|
||||
f"skipping batch. Last error: {last_exc}"
|
||||
)
|
||||
return _BatchLLMResult(obs_count=len(union_observations), prompt_chars=len(prompt), failed=True)
|
||||
|
||||
@@ -1231,6 +1321,12 @@ async def _create_observation_directly(
|
||||
perf: ConsolidationPerfLog | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Create an observation from one or more source memories with pre-processed text."""
|
||||
live_source_memory_ids = await _filter_live_source_memories(conn, bank_id, source_memory_ids)
|
||||
if not live_source_memory_ids:
|
||||
logger.debug(f"Create skipped: all {len(source_memory_ids)} source memories were deleted concurrently")
|
||||
return {"action": "skipped", "reason": "sources_deleted"}
|
||||
source_memory_ids = live_source_memory_ids
|
||||
|
||||
# Generate embedding for the observation (convert to string for pgvector)
|
||||
t0 = time.time()
|
||||
embeddings = await embedding_utils.generate_embeddings_batch(memory_engine.embeddings, [observation_text])
|
||||
|
||||
@@ -20,6 +20,7 @@ from ..config import (
|
||||
DEFAULT_RERANKER_COHERE_MODEL,
|
||||
DEFAULT_RERANKER_FLASHRANK_CACHE_DIR,
|
||||
DEFAULT_RERANKER_FLASHRANK_MODEL,
|
||||
DEFAULT_RERANKER_GOOGLE_MODEL,
|
||||
DEFAULT_RERANKER_LITELLM_MAX_TOKENS_PER_DOC,
|
||||
DEFAULT_RERANKER_LITELLM_MODEL,
|
||||
DEFAULT_RERANKER_LITELLM_SDK_MODEL,
|
||||
@@ -29,20 +30,26 @@ from ..config import (
|
||||
DEFAULT_RERANKER_LOCAL_MODEL,
|
||||
DEFAULT_RERANKER_LOCAL_TRUST_REMOTE_CODE,
|
||||
DEFAULT_RERANKER_PROVIDER,
|
||||
DEFAULT_RERANKER_SILICONFLOW_BASE_URL,
|
||||
DEFAULT_RERANKER_SILICONFLOW_MODEL,
|
||||
DEFAULT_RERANKER_TEI_BATCH_SIZE,
|
||||
DEFAULT_RERANKER_TEI_HTTP_TIMEOUT,
|
||||
DEFAULT_RERANKER_TEI_MAX_CONCURRENT,
|
||||
DEFAULT_RERANKER_ZEROENTROPY_MODEL,
|
||||
ENV_RERANKER_COHERE_API_KEY,
|
||||
ENV_RERANKER_COHERE_MODEL,
|
||||
ENV_RERANKER_FLASHRANK_CACHE_DIR,
|
||||
ENV_RERANKER_FLASHRANK_MODEL,
|
||||
ENV_RERANKER_GOOGLE_PROJECT_ID,
|
||||
ENV_RERANKER_LITELLM_SDK_API_KEY,
|
||||
ENV_RERANKER_LOCAL_FORCE_CPU,
|
||||
ENV_RERANKER_LOCAL_MAX_CONCURRENT,
|
||||
ENV_RERANKER_LOCAL_MODEL,
|
||||
ENV_RERANKER_LOCAL_TRUST_REMOTE_CODE,
|
||||
ENV_RERANKER_PROVIDER,
|
||||
ENV_RERANKER_SILICONFLOW_API_KEY,
|
||||
ENV_RERANKER_TEI_BATCH_SIZE,
|
||||
ENV_RERANKER_TEI_HTTP_TIMEOUT,
|
||||
ENV_RERANKER_TEI_MAX_CONCURRENT,
|
||||
ENV_RERANKER_TEI_URL,
|
||||
ENV_RERANKER_ZEROENTROPY_API_KEY,
|
||||
@@ -516,6 +523,84 @@ class RemoteTEICrossEncoder(CrossEncoderModel):
|
||||
return await self._predict_async(pairs)
|
||||
|
||||
|
||||
class _CohereCompatibleRerankClient:
|
||||
"""
|
||||
Internal HTTP client for Cohere-compatible /rerank endpoints.
|
||||
|
||||
Shared by all providers that speak the Cohere rerank wire format —
|
||||
{model, query, documents[, top_n]} request and
|
||||
{results: [{index, relevance_score}, ...]} response. This covers
|
||||
SiliconFlow, ZeroEntropy, Jina, Voyage, BGE self-hosted, and Cohere
|
||||
itself when reached via a custom base_url (e.g. Azure AI Foundry).
|
||||
|
||||
Not a CrossEncoderModel — providers compose it and expose their own
|
||||
provider_name / initialization logging.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
api_key: str,
|
||||
model: str,
|
||||
rerank_url: str,
|
||||
timeout: float = 60.0,
|
||||
include_top_n: bool = True,
|
||||
):
|
||||
self.api_key = api_key
|
||||
self.model = model
|
||||
self.rerank_url = rerank_url
|
||||
self.timeout = timeout
|
||||
self.include_top_n = include_top_n
|
||||
self._async_client: httpx.AsyncClient | None = None
|
||||
|
||||
async def initialize(self) -> None:
|
||||
if self._async_client is not None:
|
||||
return
|
||||
self._async_client = httpx.AsyncClient(
|
||||
timeout=self.timeout,
|
||||
headers={
|
||||
"Authorization": f"Bearer {self.api_key}",
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
)
|
||||
|
||||
async def predict(self, pairs: list[tuple[str, str]]) -> list[float]:
|
||||
if self._async_client is None:
|
||||
raise RuntimeError("Reranker not initialized. Call initialize() first.")
|
||||
|
||||
if not pairs:
|
||||
return []
|
||||
|
||||
query_groups: dict[str, list[tuple[int, str]]] = {}
|
||||
for idx, (query, text) in enumerate(pairs):
|
||||
query_groups.setdefault(query, []).append((idx, text))
|
||||
|
||||
all_scores = [0.0] * len(pairs)
|
||||
|
||||
for query, indexed_texts in query_groups.items():
|
||||
texts = [text for _, text in indexed_texts]
|
||||
indices = [idx for idx, _ in indexed_texts]
|
||||
|
||||
body: dict[str, object] = {
|
||||
"model": self.model,
|
||||
"query": query,
|
||||
"documents": texts,
|
||||
"return_documents": False,
|
||||
}
|
||||
if self.include_top_n:
|
||||
body["top_n"] = len(texts)
|
||||
|
||||
response = await self._async_client.post(self.rerank_url, json=body)
|
||||
response.raise_for_status()
|
||||
result = response.json()
|
||||
|
||||
for item in result.get("results", []):
|
||||
original_idx = item["index"]
|
||||
score = item["relevance_score"]
|
||||
all_scores[indices[original_idx]] = score
|
||||
|
||||
return all_scores
|
||||
|
||||
|
||||
class CohereCrossEncoder(CrossEncoderModel):
|
||||
"""
|
||||
Cohere cross-encoder implementation using the Cohere Rerank API.
|
||||
@@ -544,7 +629,20 @@ class CohereCrossEncoder(CrossEncoderModel):
|
||||
self.base_url = base_url
|
||||
self.timeout = timeout
|
||||
self._client = None
|
||||
self._httpx_client: httpx.Client | None = None
|
||||
# Used when base_url is set (Azure AI Foundry and other Cohere-compatible hosts).
|
||||
# Azure endpoints already include the full invoke path, so rerank_url == base_url
|
||||
# and top_n is omitted to match the existing Azure contract.
|
||||
self._http_client: _CohereCompatibleRerankClient | None = (
|
||||
_CohereCompatibleRerankClient(
|
||||
api_key=api_key,
|
||||
model=model,
|
||||
rerank_url=base_url,
|
||||
timeout=timeout,
|
||||
include_top_n=False,
|
||||
)
|
||||
if base_url
|
||||
else None
|
||||
)
|
||||
|
||||
@property
|
||||
def provider_name(self) -> str:
|
||||
@@ -552,23 +650,15 @@ class CohereCrossEncoder(CrossEncoderModel):
|
||||
|
||||
async def initialize(self) -> None:
|
||||
"""Initialize the Cohere client."""
|
||||
if self._client is not None or self._httpx_client is not None:
|
||||
if self._client is not None or (self._http_client and self._http_client._async_client):
|
||||
return
|
||||
|
||||
base_url_msg = f" at {self.base_url}" if self.base_url else ""
|
||||
logger.info(f"Reranker: initializing Cohere provider with model {self.model}{base_url_msg}")
|
||||
|
||||
if self.base_url:
|
||||
# For custom endpoints (Azure AI Foundry), use httpx directly to avoid SDK path appending
|
||||
# Azure endpoints already include the full path (e.g., /models/.../invoke)
|
||||
self._httpx_client = httpx.Client(
|
||||
timeout=self.timeout,
|
||||
headers={
|
||||
"Authorization": f"Bearer {self.api_key}",
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
)
|
||||
logger.info("Reranker: Cohere provider initialized (using httpx for custom endpoint)")
|
||||
if self._http_client is not None:
|
||||
await self._http_client.initialize()
|
||||
logger.info("Reranker: Cohere provider initialized (Cohere-compatible HTTP endpoint)")
|
||||
else:
|
||||
# For native Cohere API, use the official SDK
|
||||
try:
|
||||
@@ -589,25 +679,24 @@ class CohereCrossEncoder(CrossEncoderModel):
|
||||
Returns:
|
||||
List of relevance scores
|
||||
"""
|
||||
if self._client is None and self._httpx_client is None:
|
||||
if self._client is None and self._http_client is None:
|
||||
raise RuntimeError("Reranker not initialized. Call initialize() first.")
|
||||
|
||||
if not pairs:
|
||||
return []
|
||||
|
||||
# Run sync Cohere API calls in thread pool
|
||||
loop = asyncio.get_event_loop()
|
||||
return await loop.run_in_executor(None, self._predict_sync, pairs)
|
||||
if self._http_client is not None:
|
||||
return await self._http_client.predict(pairs)
|
||||
|
||||
def _predict_sync(self, pairs: list[tuple[str, str]]) -> list[float]:
|
||||
"""Synchronous predict implementation for Cohere API."""
|
||||
# Group pairs by query for efficient batching
|
||||
# Cohere rerank expects one query with multiple documents
|
||||
# Run sync Cohere SDK calls in thread pool
|
||||
loop = asyncio.get_event_loop()
|
||||
return await loop.run_in_executor(None, self._predict_sync_sdk, pairs)
|
||||
|
||||
def _predict_sync_sdk(self, pairs: list[tuple[str, str]]) -> list[float]:
|
||||
"""Synchronous predict using the native Cohere SDK."""
|
||||
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))
|
||||
query_groups.setdefault(query, []).append((idx, text))
|
||||
|
||||
all_scores = [0.0] * len(pairs)
|
||||
|
||||
@@ -615,40 +704,17 @@ class CohereCrossEncoder(CrossEncoderModel):
|
||||
texts = [text for _, text in indexed_texts]
|
||||
indices = [idx for idx, _ in indexed_texts]
|
||||
|
||||
if self._httpx_client:
|
||||
# Direct HTTP request for custom endpoints (Azure AI Foundry)
|
||||
response = self._httpx_client.post(
|
||||
self.base_url,
|
||||
json={
|
||||
"model": self.model,
|
||||
"query": query,
|
||||
"documents": texts,
|
||||
"return_documents": False,
|
||||
},
|
||||
)
|
||||
response.raise_for_status()
|
||||
result = response.json()
|
||||
response = self._client.rerank(
|
||||
query=query,
|
||||
documents=texts,
|
||||
model=self.model,
|
||||
return_documents=False,
|
||||
)
|
||||
|
||||
# Map scores back to original positions
|
||||
# Azure Cohere response format: {"results": [{"index": 0, "relevance_score": 0.9}, ...]}
|
||||
for item in result.get("results", []):
|
||||
original_idx = item["index"]
|
||||
score = item["relevance_score"]
|
||||
all_scores[indices[original_idx]] = score
|
||||
else:
|
||||
# Native Cohere SDK for standard API
|
||||
response = self._client.rerank(
|
||||
query=query,
|
||||
documents=texts,
|
||||
model=self.model,
|
||||
return_documents=False,
|
||||
)
|
||||
|
||||
# Map scores back to original positions
|
||||
for result in response.results:
|
||||
original_idx = result.index
|
||||
score = result.relevance_score
|
||||
all_scores[indices[original_idx]] = score
|
||||
for result in response.results:
|
||||
original_idx = result.index
|
||||
score = result.relevance_score
|
||||
all_scores[indices[original_idx]] = score
|
||||
|
||||
return all_scores
|
||||
|
||||
@@ -671,89 +737,70 @@ class ZeroEntropyCrossEncoder(CrossEncoderModel):
|
||||
base_url: str | None = None,
|
||||
timeout: float = 60.0,
|
||||
):
|
||||
"""
|
||||
Initialize ZeroEntropy cross-encoder client.
|
||||
|
||||
Args:
|
||||
api_key: ZeroEntropy API key
|
||||
model: ZeroEntropy rerank model name (default: zerank-2)
|
||||
base_url: Custom base URL for ZeroEntropy-compatible API (e.g., mock server or proxy)
|
||||
timeout: Request timeout in seconds (default: 60.0)
|
||||
"""
|
||||
self.api_key = api_key
|
||||
self.model = model
|
||||
self.base_url = base_url.rstrip("/") if base_url else self.DEFAULT_BASE_URL
|
||||
self.rerank_url = f"{self.base_url}{self.RERANK_PATH}"
|
||||
self.timeout = timeout
|
||||
self._async_client: httpx.AsyncClient | None = None
|
||||
self._client = _CohereCompatibleRerankClient(
|
||||
api_key=api_key,
|
||||
model=model,
|
||||
rerank_url=f"{self.base_url}{self.RERANK_PATH}",
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
@property
|
||||
def provider_name(self) -> str:
|
||||
return "zeroentropy"
|
||||
|
||||
async def initialize(self) -> None:
|
||||
"""Initialize the async HTTP client."""
|
||||
if self._async_client is not None:
|
||||
if self._client._async_client is not None:
|
||||
return
|
||||
|
||||
logger.info(f"Reranker: initializing ZeroEntropy provider with model {self.model}")
|
||||
self._async_client = httpx.AsyncClient(
|
||||
timeout=self.timeout,
|
||||
headers={
|
||||
"Authorization": f"Bearer {self.api_key}",
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
)
|
||||
await self._client.initialize()
|
||||
logger.info("Reranker: ZeroEntropy provider initialized")
|
||||
|
||||
async def predict(self, pairs: list[tuple[str, str]]) -> list[float]:
|
||||
"""
|
||||
Score query-document pairs using the ZeroEntropy Rerank API.
|
||||
return await self._client.predict(pairs)
|
||||
|
||||
Args:
|
||||
pairs: List of (query, document) tuples to score
|
||||
|
||||
Returns:
|
||||
List of relevance scores
|
||||
"""
|
||||
if self._async_client is None:
|
||||
raise RuntimeError("Reranker not initialized. Call initialize() first.")
|
||||
class SiliconFlowCrossEncoder(CrossEncoderModel):
|
||||
"""
|
||||
SiliconFlow cross-encoder implementation.
|
||||
|
||||
if not pairs:
|
||||
return []
|
||||
SiliconFlow (https://siliconflow.cn) exposes a Cohere-compatible /rerank
|
||||
endpoint. Shares the HTTP client with ZeroEntropy/Cohere-custom-endpoint
|
||||
via _CohereCompatibleRerankClient.
|
||||
"""
|
||||
|
||||
# Group pairs by query for efficient batching
|
||||
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))
|
||||
RERANK_PATH = "/rerank"
|
||||
|
||||
all_scores = [0.0] * len(pairs)
|
||||
def __init__(
|
||||
self,
|
||||
api_key: str,
|
||||
model: str = DEFAULT_RERANKER_SILICONFLOW_MODEL,
|
||||
base_url: str = DEFAULT_RERANKER_SILICONFLOW_BASE_URL,
|
||||
timeout: float = 60.0,
|
||||
):
|
||||
self.model = model
|
||||
self.base_url = base_url.rstrip("/")
|
||||
self._client = _CohereCompatibleRerankClient(
|
||||
api_key=api_key,
|
||||
model=model,
|
||||
rerank_url=f"{self.base_url}{self.RERANK_PATH}",
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
for query, indexed_texts in query_groups.items():
|
||||
texts = [text for _, text in indexed_texts]
|
||||
indices = [idx for idx, _ in indexed_texts]
|
||||
@property
|
||||
def provider_name(self) -> str:
|
||||
return "siliconflow"
|
||||
|
||||
response = await self._async_client.post(
|
||||
self.rerank_url,
|
||||
json={
|
||||
"model": self.model,
|
||||
"query": query,
|
||||
"documents": texts,
|
||||
"top_n": len(texts),
|
||||
},
|
||||
)
|
||||
response.raise_for_status()
|
||||
result = response.json()
|
||||
async def initialize(self) -> None:
|
||||
if self._client._async_client is not None:
|
||||
return
|
||||
logger.info(f"Reranker: initializing SiliconFlow provider at {self.base_url} with model {self.model}")
|
||||
await self._client.initialize()
|
||||
logger.info("Reranker: SiliconFlow provider initialized")
|
||||
|
||||
# Map scores back to original positions
|
||||
for item in result.get("results", []):
|
||||
original_idx = item["index"]
|
||||
score = item["relevance_score"]
|
||||
all_scores[indices[original_idx]] = score
|
||||
|
||||
return all_scores
|
||||
async def predict(self, pairs: list[tuple[str, str]]) -> list[float]:
|
||||
return await self._client.predict(pairs)
|
||||
|
||||
|
||||
class RRFPassthroughCrossEncoder(CrossEncoderModel):
|
||||
@@ -1205,14 +1252,31 @@ class JinaMLXCrossEncoder(CrossEncoderModel):
|
||||
if self._reranker is not None:
|
||||
return
|
||||
|
||||
# Pre-warm transformers.AutoTokenizer to fully populate the transformers
|
||||
# namespace before mlx_lm imports it. transformers 5.x uses _LazyModule,
|
||||
# which has an unguarded window where `from transformers import AutoTokenizer`
|
||||
# raises ImportError if another thread is concurrently initializing the
|
||||
# namespace (e.g. embeddings init in an executor thread).
|
||||
# See: https://github.com/vectorize-io/hindsight/issues/994
|
||||
import transformers
|
||||
|
||||
_ = transformers.AutoTokenizer
|
||||
|
||||
try:
|
||||
import mlx.core # noqa: F401
|
||||
import mlx_lm # noqa: F401
|
||||
except ImportError:
|
||||
except ImportError as exc:
|
||||
# Only swallow "package not installed" errors. Anything else (e.g. a
|
||||
# transitive import failure inside mlx_lm) must surface verbatim so
|
||||
# the real cause is debuggable instead of being masked by a generic
|
||||
# "install mlx" message.
|
||||
msg = str(exc)
|
||||
if "mlx" not in msg and "mlx_lm" not in msg:
|
||||
raise
|
||||
raise ImportError(
|
||||
"mlx and mlx-lm are required for JinaMLXCrossEncoder. "
|
||||
"Install with: pip install mlx>=0.31.0 mlx-lm>=0.31.1 safetensors>=0.6.2"
|
||||
)
|
||||
) from exc
|
||||
|
||||
loop = asyncio.get_event_loop()
|
||||
await loop.run_in_executor(None, self._load_model)
|
||||
@@ -1220,6 +1284,7 @@ class JinaMLXCrossEncoder(CrossEncoderModel):
|
||||
def _load_model(self) -> None:
|
||||
"""Download (if needed) and load the MLX reranker. Runs in a thread."""
|
||||
import os
|
||||
import threading
|
||||
|
||||
from huggingface_hub import snapshot_download
|
||||
|
||||
@@ -1235,6 +1300,10 @@ class JinaMLXCrossEncoder(CrossEncoderModel):
|
||||
model_path=model_path,
|
||||
projector_path=os.path.join(model_path, "projector.safetensors"),
|
||||
)
|
||||
# MLX Metal GPU ops are not thread-safe — concurrent calls to
|
||||
# Device::end_encoding() crash with SIGSEGV (NULL deref).
|
||||
# Serialize all reranker inference through this lock.
|
||||
self._mlx_lock = threading.Lock()
|
||||
logger.info("Reranker: jina-mlx provider initialized")
|
||||
|
||||
def _predict_sync(self, pairs: list[tuple[str, str]]) -> list[float]:
|
||||
@@ -1248,13 +1317,14 @@ class JinaMLXCrossEncoder(CrossEncoderModel):
|
||||
|
||||
all_scores = [0.0] * len(pairs)
|
||||
|
||||
for query, indexed_docs in query_groups.items():
|
||||
docs = [doc for _, doc in indexed_docs]
|
||||
indices = [idx for idx, _ in indexed_docs]
|
||||
results = self._reranker.rerank(query, docs)
|
||||
for result in results:
|
||||
original_idx = result["index"]
|
||||
all_scores[indices[original_idx]] = result["relevance_score"]
|
||||
with self._mlx_lock:
|
||||
for query, indexed_docs in query_groups.items():
|
||||
docs = [doc for _, doc in indexed_docs]
|
||||
indices = [idx for idx, _ in indexed_docs]
|
||||
results = self._reranker.rerank(query, docs)
|
||||
for result in results:
|
||||
original_idx = result["index"]
|
||||
all_scores[indices[original_idx]] = result["relevance_score"]
|
||||
|
||||
return all_scores
|
||||
|
||||
@@ -1266,6 +1336,164 @@ class JinaMLXCrossEncoder(CrossEncoderModel):
|
||||
return await loop.run_in_executor(None, self._predict_sync, pairs)
|
||||
|
||||
|
||||
class GoogleCrossEncoder(CrossEncoderModel):
|
||||
"""
|
||||
Google Discovery Engine cross-encoder using the Ranking REST API.
|
||||
|
||||
Uses httpx + google-auth for lightweight REST calls (no gRPC/protobuf).
|
||||
Supports ADC (Application Default Credentials) or service account key file.
|
||||
|
||||
Available models:
|
||||
- semantic-ranker-default-004: Best quality, 1024 tokens/record (recommended)
|
||||
- semantic-ranker-fast-004: Lower latency, 1024 tokens/record
|
||||
|
||||
Max 200 records per API request. Location is always "global".
|
||||
"""
|
||||
|
||||
MAX_RECORDS_PER_REQUEST = 200
|
||||
API_BASE = "https://discoveryengine.googleapis.com/v1"
|
||||
SCOPES = ["https://www.googleapis.com/auth/cloud-platform"]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
project_id: str,
|
||||
model: str = DEFAULT_RERANKER_GOOGLE_MODEL,
|
||||
service_account_key: str | None = None,
|
||||
location: str = "global",
|
||||
timeout: float = 60.0,
|
||||
):
|
||||
"""
|
||||
Initialize Google Discovery Engine cross-encoder.
|
||||
|
||||
Args:
|
||||
project_id: Google Cloud project ID
|
||||
model: Ranking model name (default: semantic-ranker-default-004)
|
||||
service_account_key: Path to service account JSON key file.
|
||||
If None, uses Application Default Credentials (ADC).
|
||||
location: API location (default: "global")
|
||||
timeout: Request timeout in seconds (default: 60.0)
|
||||
"""
|
||||
self.project_id = project_id
|
||||
self.model = model
|
||||
self.service_account_key = service_account_key
|
||||
self.location = location
|
||||
self.timeout = timeout
|
||||
self._credentials = None
|
||||
self._client: httpx.Client | None = None
|
||||
self._rank_url: str | None = None
|
||||
|
||||
@property
|
||||
def provider_name(self) -> str:
|
||||
return "google"
|
||||
|
||||
def _get_auth_headers(self) -> dict[str, str]:
|
||||
"""Get Authorization header with a fresh access token."""
|
||||
import google.auth.transport.requests
|
||||
|
||||
if not self._credentials.valid:
|
||||
self._credentials.refresh(google.auth.transport.requests.Request())
|
||||
return {"Authorization": f"Bearer {self._credentials.token}"}
|
||||
|
||||
async def initialize(self) -> None:
|
||||
"""Initialize credentials and HTTP client."""
|
||||
if self._client is not None:
|
||||
return
|
||||
|
||||
auth_method = "ADC" if not self.service_account_key else "service_account"
|
||||
logger.info(
|
||||
f"Reranker: initializing Google Discovery Engine provider "
|
||||
f"(project={self.project_id}, model={self.model}, auth={auth_method})"
|
||||
)
|
||||
if self.service_account_key:
|
||||
try:
|
||||
from google.oauth2 import service_account
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"google-auth is required for GoogleCrossEncoder. Install it with: pip install google-auth"
|
||||
)
|
||||
self._credentials = service_account.Credentials.from_service_account_file(
|
||||
self.service_account_key,
|
||||
scopes=self.SCOPES,
|
||||
)
|
||||
else:
|
||||
try:
|
||||
import google.auth
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"google-auth is required for GoogleCrossEncoder. Install it with: pip install google-auth"
|
||||
)
|
||||
self._credentials, _ = google.auth.default(scopes=self.SCOPES)
|
||||
|
||||
ranking_config = f"projects/{self.project_id}/locations/{self.location}/rankingConfigs/default_ranking_config"
|
||||
self._rank_url = f"{self.API_BASE}/{ranking_config}:rank"
|
||||
self._client = httpx.Client(timeout=self.timeout)
|
||||
|
||||
logger.info("Reranker: Google Discovery Engine provider initialized")
|
||||
|
||||
def _predict_sync(self, pairs: list[tuple[str, str]]) -> list[float]:
|
||||
"""Synchronous predict via REST API."""
|
||||
if not pairs:
|
||||
return []
|
||||
|
||||
# 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)
|
||||
|
||||
for query, indexed_texts in query_groups.items():
|
||||
texts = [text for _, text in indexed_texts]
|
||||
indices = [idx for idx, _ in indexed_texts]
|
||||
|
||||
# Process in batches of MAX_RECORDS_PER_REQUEST
|
||||
for batch_start in range(0, len(texts), self.MAX_RECORDS_PER_REQUEST):
|
||||
batch_texts = texts[batch_start : batch_start + self.MAX_RECORDS_PER_REQUEST]
|
||||
batch_indices = indices[batch_start : batch_start + self.MAX_RECORDS_PER_REQUEST]
|
||||
|
||||
records = [{"id": str(i), "content": text} for i, text in enumerate(batch_texts)]
|
||||
|
||||
response = self._client.post(
|
||||
self._rank_url,
|
||||
headers=self._get_auth_headers(),
|
||||
json={
|
||||
"model": self.model,
|
||||
"query": query,
|
||||
"records": records,
|
||||
"topN": len(records),
|
||||
},
|
||||
)
|
||||
response.raise_for_status()
|
||||
result = response.json()
|
||||
|
||||
for record in result.get("records", []):
|
||||
local_idx = int(record["id"])
|
||||
all_scores[batch_indices[local_idx]] = record["score"]
|
||||
|
||||
return all_scores
|
||||
|
||||
async def predict(self, pairs: list[tuple[str, str]]) -> list[float]:
|
||||
"""
|
||||
Score query-document pairs using Google Discovery Engine Ranking API.
|
||||
|
||||
Args:
|
||||
pairs: List of (query, document) tuples to score
|
||||
|
||||
Returns:
|
||||
List of relevance scores (0-1, higher = more relevant)
|
||||
"""
|
||||
if self._client is None:
|
||||
raise RuntimeError("Reranker not initialized. Call initialize() first.")
|
||||
|
||||
if not pairs:
|
||||
return []
|
||||
|
||||
loop = asyncio.get_event_loop()
|
||||
return await loop.run_in_executor(None, self._predict_sync, pairs)
|
||||
|
||||
|
||||
def create_cross_encoder_from_env() -> CrossEncoderModel:
|
||||
"""
|
||||
Create a CrossEncoderModel instance based on configuration.
|
||||
@@ -1286,6 +1514,7 @@ def create_cross_encoder_from_env() -> CrossEncoderModel:
|
||||
raise ValueError(f"{ENV_RERANKER_TEI_URL} is required when {ENV_RERANKER_PROVIDER} is 'tei'")
|
||||
return RemoteTEICrossEncoder(
|
||||
base_url=url,
|
||||
timeout=config.reranker_tei_http_timeout,
|
||||
batch_size=config.reranker_tei_batch_size,
|
||||
max_concurrent=config.reranker_tei_max_concurrent,
|
||||
)
|
||||
@@ -1308,6 +1537,18 @@ def create_cross_encoder_from_env() -> CrossEncoderModel:
|
||||
model=config.reranker_cohere_model,
|
||||
base_url=config.reranker_cohere_base_url,
|
||||
)
|
||||
elif provider == "openrouter":
|
||||
api_key = config.reranker_openrouter_api_key
|
||||
if not api_key:
|
||||
raise ValueError(
|
||||
"HINDSIGHT_API_RERANKER_OPENROUTER_API_KEY, HINDSIGHT_API_OPENROUTER_API_KEY, "
|
||||
f"or HINDSIGHT_API_LLM_API_KEY is required when {ENV_RERANKER_PROVIDER} is 'openrouter'"
|
||||
)
|
||||
return CohereCrossEncoder(
|
||||
api_key=api_key,
|
||||
model=config.reranker_openrouter_model,
|
||||
base_url="https://openrouter.ai/api/v1/rerank",
|
||||
)
|
||||
elif provider == "flashrank":
|
||||
model = os.environ.get(ENV_RERANKER_FLASHRANK_MODEL, DEFAULT_RERANKER_FLASHRANK_MODEL)
|
||||
cache_dir = os.environ.get(ENV_RERANKER_FLASHRANK_CACHE_DIR, DEFAULT_RERANKER_FLASHRANK_CACHE_DIR)
|
||||
@@ -1341,11 +1582,34 @@ def create_cross_encoder_from_env() -> CrossEncoderModel:
|
||||
api_key=api_key,
|
||||
model=config.reranker_zeroentropy_model,
|
||||
)
|
||||
elif provider == "siliconflow":
|
||||
api_key = config.reranker_siliconflow_api_key
|
||||
if not api_key:
|
||||
raise ValueError(
|
||||
f"{ENV_RERANKER_SILICONFLOW_API_KEY} is required when {ENV_RERANKER_PROVIDER} is 'siliconflow'"
|
||||
)
|
||||
return SiliconFlowCrossEncoder(
|
||||
api_key=api_key,
|
||||
model=config.reranker_siliconflow_model,
|
||||
base_url=config.reranker_siliconflow_base_url,
|
||||
)
|
||||
elif provider == "google":
|
||||
project_id = config.reranker_google_project_id
|
||||
if not project_id:
|
||||
raise ValueError(
|
||||
f"{ENV_RERANKER_GOOGLE_PROJECT_ID} (or HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID) "
|
||||
f"is required when {ENV_RERANKER_PROVIDER} is 'google'"
|
||||
)
|
||||
return GoogleCrossEncoder(
|
||||
project_id=project_id,
|
||||
model=config.reranker_google_model,
|
||||
service_account_key=config.reranker_google_service_account_key,
|
||||
)
|
||||
elif provider == "rrf":
|
||||
return RRFPassthroughCrossEncoder()
|
||||
elif provider == "jina-mlx":
|
||||
return JinaMLXCrossEncoder()
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unknown reranker provider: {provider}. Supported: 'local', 'tei', 'cohere', 'zeroentropy', 'flashrank', 'litellm', 'litellm-sdk', 'rrf', 'jina-mlx'"
|
||||
f"Unknown reranker provider: {provider}. Supported: 'local', 'tei', 'cohere', 'zeroentropy', 'siliconflow', 'google', 'flashrank', 'litellm', 'litellm-sdk', 'rrf', 'jina-mlx'"
|
||||
)
|
||||
|
||||
@@ -19,6 +19,7 @@ import httpx
|
||||
|
||||
from ..config import (
|
||||
DEFAULT_EMBEDDINGS_COHERE_MODEL,
|
||||
DEFAULT_EMBEDDINGS_GEMINI_MODEL,
|
||||
DEFAULT_EMBEDDINGS_LITELLM_MODEL,
|
||||
DEFAULT_EMBEDDINGS_LITELLM_SDK_MODEL,
|
||||
DEFAULT_EMBEDDINGS_LOCAL_FORCE_CPU,
|
||||
@@ -28,6 +29,7 @@ from ..config import (
|
||||
DEFAULT_EMBEDDINGS_PROVIDER,
|
||||
DEFAULT_LITELLM_API_BASE,
|
||||
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,
|
||||
@@ -755,6 +757,7 @@ class LiteLLMSDKEmbeddings(Embeddings):
|
||||
output_dimensions: int | None = None,
|
||||
batch_size: int = 100,
|
||||
timeout: float = 60.0,
|
||||
encoding_format: str | None = "float",
|
||||
):
|
||||
"""
|
||||
Initialize LiteLLM SDK embeddings client.
|
||||
@@ -766,6 +769,8 @@ class LiteLLMSDKEmbeddings(Embeddings):
|
||||
output_dimensions: Optional output embedding dimensions (provider-dependent)
|
||||
batch_size: Maximum batch size for embedding requests (default: 100)
|
||||
timeout: Request timeout in seconds (default: 60.0)
|
||||
encoding_format: Encoding format for embeddings (default: "float").
|
||||
Set to None or empty string to omit (needed for Voyage AI, Gemini).
|
||||
"""
|
||||
self.api_key = api_key
|
||||
self.model = model
|
||||
@@ -773,6 +778,7 @@ class LiteLLMSDKEmbeddings(Embeddings):
|
||||
self.output_dimensions = output_dimensions
|
||||
self.batch_size = batch_size
|
||||
self.timeout = timeout
|
||||
self.encoding_format = encoding_format or None
|
||||
self._litellm = None # Will be set during initialization
|
||||
self._dimension: int | None = None
|
||||
|
||||
@@ -808,8 +814,9 @@ class LiteLLMSDKEmbeddings(Embeddings):
|
||||
"model": self.model,
|
||||
"input": ["test"],
|
||||
"api_key": self.api_key,
|
||||
"encoding_format": "float",
|
||||
}
|
||||
if self.encoding_format:
|
||||
embed_kwargs["encoding_format"] = self.encoding_format
|
||||
if self.api_base:
|
||||
embed_kwargs["api_base"] = self.api_base
|
||||
if self.output_dimensions is not None:
|
||||
@@ -857,8 +864,9 @@ class LiteLLMSDKEmbeddings(Embeddings):
|
||||
"model": self.model,
|
||||
"input": batch,
|
||||
"api_key": self.api_key,
|
||||
"encoding_format": "float",
|
||||
}
|
||||
if self.encoding_format:
|
||||
embed_kwargs["encoding_format"] = self.encoding_format
|
||||
if self.api_base:
|
||||
embed_kwargs["api_base"] = self.api_base
|
||||
if self.output_dimensions is not None:
|
||||
@@ -884,6 +892,179 @@ class LiteLLMSDKEmbeddings(Embeddings):
|
||||
return all_embeddings
|
||||
|
||||
|
||||
class GeminiEmbeddings(Embeddings):
|
||||
"""
|
||||
Google embeddings via the google.genai SDK.
|
||||
|
||||
Supports both:
|
||||
1. Gemini API (api.generativeai.google.com) with API key authentication
|
||||
2. Vertex AI with service account or Application Default Credentials (ADC)
|
||||
|
||||
Uses the embed_content API: client.models.embed_content(model, contents)
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model: str = DEFAULT_EMBEDDINGS_GEMINI_MODEL,
|
||||
api_key: str | None = None,
|
||||
vertexai_project_id: str | None = None,
|
||||
vertexai_region: str | None = None,
|
||||
vertexai_service_account_key: str | None = None,
|
||||
output_dimensionality: int | None = None,
|
||||
batch_size: int = 100,
|
||||
):
|
||||
self.model = model
|
||||
self.api_key = api_key
|
||||
self.vertexai_project_id = vertexai_project_id
|
||||
self.vertexai_region = vertexai_region or "us-central1"
|
||||
self.vertexai_service_account_key = vertexai_service_account_key
|
||||
self.output_dimensionality = output_dimensionality
|
||||
self.batch_size = batch_size
|
||||
self._client = None
|
||||
self._dimension: int | None = None
|
||||
self._is_vertexai = vertexai_project_id is not None
|
||||
self._embed_config = None # EmbedContentConfig, built during initialize()
|
||||
|
||||
@property
|
||||
def provider_name(self) -> str:
|
||||
return "google"
|
||||
|
||||
@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 Google genai client and detect embedding dimension."""
|
||||
if self._client is not None:
|
||||
return
|
||||
|
||||
from google import genai
|
||||
from google.genai import types as genai_types
|
||||
|
||||
if self._is_vertexai:
|
||||
self._init_vertexai(genai)
|
||||
else:
|
||||
self._init_gemini(genai)
|
||||
|
||||
# Build EmbedContentConfig if output_dimensionality is set
|
||||
if self.output_dimensionality is not None:
|
||||
self._embed_config = genai_types.EmbedContentConfig(
|
||||
output_dimensionality=self.output_dimensionality,
|
||||
)
|
||||
|
||||
# Detect dimension via a test embedding (respects output_dimensionality)
|
||||
embed_kwargs = {"model": self.model, "contents": ["test"]}
|
||||
if self._embed_config is not None:
|
||||
embed_kwargs["config"] = self._embed_config
|
||||
|
||||
result = self._client.models.embed_content(**embed_kwargs) # type: ignore[union-attr]
|
||||
if result.embeddings and len(result.embeddings) > 0:
|
||||
self._dimension = len(result.embeddings[0].values)
|
||||
|
||||
auth_mode = "vertex_ai" if self._is_vertexai else "api_key"
|
||||
logger.info(
|
||||
f"Embeddings: google provider initialized (auth: {auth_mode}, model: {self.model}, dim: {self._dimension})"
|
||||
)
|
||||
|
||||
def _init_gemini(self, genai) -> None:
|
||||
"""Initialize Gemini API client with API key."""
|
||||
if not self.api_key:
|
||||
raise ValueError("Gemini embeddings provider requires an API key")
|
||||
|
||||
self._client = genai.Client(api_key=self.api_key)
|
||||
logger.info(f"Embeddings: initializing Gemini provider with model {self.model}")
|
||||
|
||||
def _init_vertexai(self, genai) -> None:
|
||||
"""Initialize Vertex AI client with project, region, and credentials."""
|
||||
if not self.vertexai_project_id:
|
||||
raise ValueError(
|
||||
"HINDSIGHT_API_EMBEDDINGS_VERTEXAI_PROJECT_ID (or HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID) "
|
||||
"is required for Vertex AI embeddings provider."
|
||||
)
|
||||
|
||||
auth_method = "ADC"
|
||||
credentials = None
|
||||
|
||||
if self.vertexai_service_account_key:
|
||||
try:
|
||||
from google.oauth2 import service_account
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"Vertex AI service account auth requires 'google-auth' package. "
|
||||
"Install with: pip install google-auth"
|
||||
)
|
||||
credentials = service_account.Credentials.from_service_account_file(
|
||||
self.vertexai_service_account_key,
|
||||
scopes=["https://www.googleapis.com/auth/cloud-platform"],
|
||||
)
|
||||
auth_method = "service_account"
|
||||
logger.info(f"Embeddings: Vertex AI using service account key: {self.vertexai_service_account_key}")
|
||||
|
||||
# Strip google/ prefix from model name — native SDK uses bare names
|
||||
if self.model.startswith("google/"):
|
||||
self.model = self.model[len("google/") :]
|
||||
|
||||
client_kwargs = {
|
||||
"vertexai": True,
|
||||
"project": self.vertexai_project_id,
|
||||
"location": self.vertexai_region,
|
||||
}
|
||||
if credentials is not None:
|
||||
client_kwargs["credentials"] = credentials
|
||||
|
||||
self._client = genai.Client(**client_kwargs)
|
||||
logger.info(
|
||||
f"Embeddings: initializing Vertex AI provider "
|
||||
f"(project={self.vertexai_project_id}, region={self.vertexai_region}, "
|
||||
f"model={self.model}, auth={auth_method})"
|
||||
)
|
||||
|
||||
def encode(self, texts: list[str]) -> list[list[float]]:
|
||||
"""
|
||||
Generate embeddings using the Google genai SDK.
|
||||
|
||||
Args:
|
||||
texts: List of text strings to encode
|
||||
|
||||
Returns:
|
||||
List of embedding vectors
|
||||
"""
|
||||
if self._client is None:
|
||||
raise RuntimeError("Embeddings not initialized. Call initialize() first.")
|
||||
|
||||
if not texts:
|
||||
return []
|
||||
|
||||
all_embeddings = []
|
||||
|
||||
# Process in batches
|
||||
for i in range(0, len(texts), self.batch_size):
|
||||
batch = texts[i : i + self.batch_size]
|
||||
|
||||
embed_kwargs = {"model": self.model, "contents": batch}
|
||||
if self._embed_config is not None:
|
||||
embed_kwargs["config"] = self._embed_config
|
||||
|
||||
result = self._client.models.embed_content(**embed_kwargs)
|
||||
|
||||
all_embeddings.extend([emb.values for emb in result.embeddings])
|
||||
|
||||
# L2-normalize when output_dimensionality is set — Gemini only returns
|
||||
# normalized vectors at full 3072 dims; truncated dims need re-normalization
|
||||
# for accurate cosine similarity.
|
||||
if self.output_dimensionality is not None:
|
||||
import numpy as np
|
||||
|
||||
arr = np.array(all_embeddings)
|
||||
norms = np.linalg.norm(arr, axis=1, keepdims=True)
|
||||
norms[norms == 0] = 1
|
||||
all_embeddings = (arr / norms).tolist()
|
||||
|
||||
return all_embeddings
|
||||
|
||||
|
||||
def create_embeddings_from_env() -> Embeddings:
|
||||
"""
|
||||
Create an Embeddings instance based on configuration.
|
||||
@@ -919,7 +1100,25 @@ def create_embeddings_from_env() -> Embeddings:
|
||||
)
|
||||
model = os.environ.get(ENV_EMBEDDINGS_OPENAI_MODEL, DEFAULT_EMBEDDINGS_OPENAI_MODEL)
|
||||
base_url = os.environ.get(ENV_EMBEDDINGS_OPENAI_BASE_URL) or None
|
||||
return OpenAIEmbeddings(api_key=api_key, model=model, base_url=base_url)
|
||||
return OpenAIEmbeddings(
|
||||
api_key=api_key,
|
||||
model=model,
|
||||
base_url=base_url,
|
||||
batch_size=config.embeddings_openai_batch_size,
|
||||
)
|
||||
elif provider == "openrouter":
|
||||
api_key = config.embeddings_openrouter_api_key
|
||||
if not api_key:
|
||||
raise ValueError(
|
||||
"HINDSIGHT_API_EMBEDDINGS_OPENROUTER_API_KEY, HINDSIGHT_API_OPENROUTER_API_KEY, "
|
||||
f"or {ENV_LLM_API_KEY} is required when {ENV_EMBEDDINGS_PROVIDER} is 'openrouter'"
|
||||
)
|
||||
return OpenAIEmbeddings(
|
||||
api_key=api_key,
|
||||
model=config.embeddings_openrouter_model,
|
||||
base_url="https://openrouter.ai/api/v1",
|
||||
batch_size=config.embeddings_openai_batch_size,
|
||||
)
|
||||
elif provider == "cohere":
|
||||
api_key = config.embeddings_cohere_api_key
|
||||
if not api_key:
|
||||
@@ -946,9 +1145,29 @@ def create_embeddings_from_env() -> Embeddings:
|
||||
model=config.embeddings_litellm_sdk_model,
|
||||
api_base=config.embeddings_litellm_sdk_api_base,
|
||||
output_dimensions=config.embeddings_litellm_sdk_output_dimensions,
|
||||
encoding_format=config.embeddings_litellm_sdk_encoding_format,
|
||||
)
|
||||
elif provider == "google":
|
||||
vertexai_project_id = config.embeddings_vertexai_project_id
|
||||
if vertexai_project_id:
|
||||
api_key = None # Vertex AI uses ADC or service account
|
||||
else:
|
||||
api_key = config.embeddings_gemini_api_key
|
||||
if not api_key:
|
||||
raise ValueError(
|
||||
f"{ENV_EMBEDDINGS_GEMINI_API_KEY} or {ENV_LLM_API_KEY} is required "
|
||||
f"when {ENV_EMBEDDINGS_PROVIDER} is 'google' (set VERTEXAI_PROJECT_ID for Vertex AI auth instead)"
|
||||
)
|
||||
return GeminiEmbeddings(
|
||||
model=config.embeddings_gemini_model,
|
||||
api_key=api_key,
|
||||
vertexai_project_id=vertexai_project_id,
|
||||
vertexai_region=config.embeddings_vertexai_region,
|
||||
vertexai_service_account_key=config.embeddings_vertexai_service_account_key,
|
||||
output_dimensionality=config.embeddings_gemini_output_dimensionality,
|
||||
)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unknown embeddings provider: {provider}. "
|
||||
f"Supported: 'local', 'tei', 'openai', 'cohere', 'litellm', 'litellm-sdk'"
|
||||
f"Supported: 'local', 'tei', 'openai', 'cohere', 'google', 'litellm', 'litellm-sdk'"
|
||||
)
|
||||
|
||||
@@ -122,6 +122,7 @@ _PROVIDERS_WITHOUT_API_KEY = frozenset(
|
||||
{
|
||||
"ollama",
|
||||
"lmstudio",
|
||||
"llamacpp",
|
||||
"openai-codex",
|
||||
"claude-code",
|
||||
"mock",
|
||||
@@ -178,6 +179,7 @@ def create_llm_provider(
|
||||
CodexLLM,
|
||||
GeminiLLM,
|
||||
LiteLLMLLM,
|
||||
LlamaCppLLM,
|
||||
MockLLM,
|
||||
NoneLLM,
|
||||
OpenAICompatibleLLM,
|
||||
@@ -263,7 +265,25 @@ def create_llm_provider(
|
||||
reasoning_effort=reasoning_effort,
|
||||
)
|
||||
|
||||
elif provider_lower in ("openai", "groq", "ollama", "lmstudio", "minimax", "volcano"):
|
||||
elif provider_lower == "llamacpp":
|
||||
from ..config import get_config
|
||||
|
||||
config = get_config()
|
||||
return LlamaCppLLM(
|
||||
provider=provider,
|
||||
api_key=api_key,
|
||||
base_url=base_url,
|
||||
model=model,
|
||||
reasoning_effort=reasoning_effort,
|
||||
model_path=config.llamacpp_model_path,
|
||||
gpu_layers=config.llamacpp_gpu_layers,
|
||||
context_size=config.llamacpp_context_size,
|
||||
chat_format=config.llamacpp_chat_format,
|
||||
no_grammar=config.llamacpp_no_grammar,
|
||||
extra_args=config.llamacpp_extra_args,
|
||||
)
|
||||
|
||||
elif provider_lower in ("openai", "groq", "ollama", "lmstudio", "minimax", "volcano", "openrouter"):
|
||||
return OpenAICompatibleLLM(
|
||||
provider=provider,
|
||||
api_key=api_key,
|
||||
@@ -333,6 +353,7 @@ class LLMProvider:
|
||||
"gemini",
|
||||
"anthropic",
|
||||
"lmstudio",
|
||||
"llamacpp",
|
||||
"vertexai",
|
||||
"openai-codex",
|
||||
"claude-code",
|
||||
@@ -342,6 +363,7 @@ class LLMProvider:
|
||||
"litellm",
|
||||
"bedrock",
|
||||
"volcano",
|
||||
"openrouter",
|
||||
]
|
||||
if self.provider not in valid_providers:
|
||||
raise ValueError(f"Invalid LLM provider: {self.provider}. Must be one of: {', '.join(valid_providers)}")
|
||||
@@ -356,6 +378,8 @@ class LLMProvider:
|
||||
self.base_url = "http://localhost:1234/v1"
|
||||
elif self.provider == "minimax":
|
||||
self.base_url = "https://api.minimax.io/v1"
|
||||
elif self.provider == "openrouter":
|
||||
self.base_url = "https://openrouter.ai/api/v1"
|
||||
|
||||
# Prepare Vertex AI config (if applicable)
|
||||
vertexai_project_id = None
|
||||
@@ -512,6 +536,15 @@ class LLMProvider:
|
||||
OutputTooLongError: If output exceeds token limits.
|
||||
Exception: Re-raises API errors after retries exhausted.
|
||||
"""
|
||||
# Stage breadcrumb so the worker log shows which LLM call a task is
|
||||
# currently inside; the stage_age field then reveals long JSON-schema
|
||||
# retry loops (e.g. a small model that can't satisfy strict_schema).
|
||||
# No-op outside a worker context.
|
||||
from ..worker.stage import set_stage
|
||||
|
||||
structured = "+structured" if response_format is not None else ""
|
||||
set_stage(f"llm.{self.provider}.{scope}{structured}")
|
||||
|
||||
async with _global_llm_semaphore:
|
||||
# Delegate to provider implementation
|
||||
result = await self._provider_impl.call(
|
||||
@@ -568,6 +601,10 @@ class LLMProvider:
|
||||
Returns:
|
||||
LLMToolCallResult with content and/or tool_calls.
|
||||
"""
|
||||
from ..worker.stage import set_stage
|
||||
|
||||
set_stage(f"llm.{self.provider}.{scope}+tools")
|
||||
|
||||
async with _global_llm_semaphore:
|
||||
# Delegate to provider implementation
|
||||
result = await self._provider_impl.call_with_tools(
|
||||
@@ -711,8 +748,9 @@ class LLMProvider:
|
||||
return ConfiguredLLMProvider(self, config.llm_gemini_safety_settings)
|
||||
|
||||
async def cleanup(self) -> None:
|
||||
"""Clean up resources."""
|
||||
pass
|
||||
"""Clean up resources (e.g. stop llamacpp subprocess)."""
|
||||
if self._provider_impl:
|
||||
await self._provider_impl.cleanup()
|
||||
|
||||
@classmethod
|
||||
def from_env(cls) -> "LLMProvider":
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -9,6 +9,7 @@ from .claude_code_llm import ClaudeCodeLLM
|
||||
from .codex_llm import CodexLLM
|
||||
from .gemini_llm import GeminiLLM
|
||||
from .litellm_llm import LiteLLMLLM
|
||||
from .llamacpp_llm import LlamaCppLLM
|
||||
from .mock_llm import MockLLM
|
||||
from .none_llm import NoneLLM
|
||||
from .openai_compatible_llm import OpenAICompatibleLLM
|
||||
@@ -18,6 +19,7 @@ __all__ = [
|
||||
"ClaudeCodeLLM",
|
||||
"CodexLLM",
|
||||
"GeminiLLM",
|
||||
"LlamaCppLLM",
|
||||
"LiteLLMLLM",
|
||||
"MockLLM",
|
||||
"NoneLLM",
|
||||
|
||||
@@ -153,7 +153,7 @@ class AnthropicLLM(LLMInterface):
|
||||
# Add JSON schema instruction if response_format is provided
|
||||
if response_format is not None and hasattr(response_format, "model_json_schema"):
|
||||
schema = response_format.model_json_schema()
|
||||
schema_msg = f"\n\nYou must respond with valid JSON matching this schema:\n{json.dumps(schema, indent=2)}"
|
||||
schema_msg = f"\n\nYou must respond with valid JSON matching this schema:\n{json.dumps(schema, indent=2, ensure_ascii=False)}"
|
||||
if system_prompt:
|
||||
system_prompt += schema_msg
|
||||
else:
|
||||
|
||||
@@ -171,7 +171,7 @@ class ClaudeCodeLLM(LLMInterface):
|
||||
if response_format is not None and hasattr(response_format, "model_json_schema"):
|
||||
schema = response_format.model_json_schema()
|
||||
schema_instruction = (
|
||||
f"\n\nYou must respond with valid JSON matching this schema:\n{json.dumps(schema, indent=2)}\n\n"
|
||||
f"\n\nYou must respond with valid JSON matching this schema:\n{json.dumps(schema, indent=2, ensure_ascii=False)}\n\n"
|
||||
"Respond with ONLY the JSON, no markdown formatting."
|
||||
)
|
||||
user_content += schema_instruction
|
||||
|
||||
@@ -205,7 +205,7 @@ class CodexLLM(LLMInterface):
|
||||
# Add JSON schema instruction if response_format is provided
|
||||
if response_format is not None and hasattr(response_format, "model_json_schema"):
|
||||
schema = response_format.model_json_schema()
|
||||
schema_msg = f"\n\nYou must respond with valid JSON matching this schema:\n{json.dumps(schema, indent=2)}"
|
||||
schema_msg = f"\n\nYou must respond with valid JSON matching this schema:\n{json.dumps(schema, indent=2, ensure_ascii=False)}"
|
||||
system_instruction += schema_msg
|
||||
|
||||
# gpt-5.2-codex only supports "detailed" reasoning summary
|
||||
|
||||
@@ -23,6 +23,7 @@ from hindsight_api.engine.llm_interface import LLMInterface, OutputTooLongError
|
||||
from hindsight_api.engine.llm_wrapper import parse_llm_json
|
||||
from hindsight_api.engine.response_models import LLMToolCall, LLMToolCallResult, TokenUsage
|
||||
from hindsight_api.metrics import get_metrics_collector
|
||||
from hindsight_api.worker.stage import set_stage
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -174,7 +175,7 @@ class GeminiLLM(LLMInterface):
|
||||
Args:
|
||||
messages: List of message dicts with 'role' and 'content'.
|
||||
response_format: Optional Pydantic model for structured output.
|
||||
max_completion_tokens: Maximum tokens in response (not supported by Gemini).
|
||||
max_completion_tokens: Maximum tokens in response (mapped to Gemini's max_output_tokens).
|
||||
temperature: Sampling temperature (0.0-2.0).
|
||||
scope: Scope identifier for tracking.
|
||||
max_retries: Maximum retry attempts.
|
||||
@@ -211,7 +212,7 @@ class GeminiLLM(LLMInterface):
|
||||
# Add JSON schema instruction if response_format is provided
|
||||
if response_format is not None and hasattr(response_format, "model_json_schema"):
|
||||
schema = response_format.model_json_schema()
|
||||
schema_msg = f"\n\nYou must respond with valid JSON matching this schema:\n{json.dumps(schema, indent=2)}"
|
||||
schema_msg = f"\n\nYou must respond with valid JSON matching this schema:\n{json.dumps(schema, indent=2, ensure_ascii=False)}"
|
||||
if system_instruction:
|
||||
system_instruction += schema_msg
|
||||
else:
|
||||
@@ -226,6 +227,11 @@ class GeminiLLM(LLMInterface):
|
||||
config_kwargs["response_schema"] = response_format
|
||||
if temperature is not None:
|
||||
config_kwargs["temperature"] = temperature
|
||||
# Gemini's equivalent of OpenAI-style max_completion_tokens is max_output_tokens.
|
||||
# Without it the model can produce arbitrarily long responses, ignoring the
|
||||
# caller's intended cap (e.g. mental_models max_tokens during refresh).
|
||||
if max_completion_tokens is not None:
|
||||
config_kwargs["max_output_tokens"] = max_completion_tokens
|
||||
|
||||
# Apply safety settings: context var (per-request bank override) takes precedence over instance default
|
||||
effective_safety_settings = _safety_settings_ctx.get()
|
||||
@@ -242,6 +248,8 @@ class GeminiLLM(LLMInterface):
|
||||
last_exception = None
|
||||
|
||||
for attempt in range(max_retries + 1):
|
||||
if attempt > 0:
|
||||
set_stage(f"llm.gemini.{scope}.attempt={attempt + 1}/{max_retries + 1}")
|
||||
try:
|
||||
response = await asyncio.wait_for(
|
||||
self._client.aio.models.generate_content(
|
||||
@@ -398,7 +406,7 @@ class GeminiLLM(LLMInterface):
|
||||
Args:
|
||||
messages: List of message dicts. Can include tool results with role='tool'.
|
||||
tools: List of tool definitions in OpenAI format.
|
||||
max_completion_tokens: Maximum tokens (not supported by Gemini).
|
||||
max_completion_tokens: Maximum tokens (mapped to Gemini's max_output_tokens).
|
||||
temperature: Sampling temperature.
|
||||
scope: Scope identifier for tracking.
|
||||
max_retries: Maximum retry attempts.
|
||||
@@ -490,6 +498,10 @@ class GeminiLLM(LLMInterface):
|
||||
config_kwargs["system_instruction"] = system_instruction
|
||||
if temperature is not None:
|
||||
config_kwargs["temperature"] = temperature
|
||||
# See note in `call`: Gemini's max_output_tokens is the equivalent of
|
||||
# OpenAI-style max_completion_tokens.
|
||||
if max_completion_tokens is not None:
|
||||
config_kwargs["max_output_tokens"] = max_completion_tokens
|
||||
|
||||
# Map OpenAI-style tool_choice to Gemini FunctionCallingConfig
|
||||
if tool_choice == "required":
|
||||
@@ -527,6 +539,8 @@ class GeminiLLM(LLMInterface):
|
||||
|
||||
last_exception = None
|
||||
for attempt in range(max_retries + 1):
|
||||
if attempt > 0:
|
||||
set_stage(f"llm.gemini.tools.attempt={attempt + 1}/{max_retries + 1}")
|
||||
try:
|
||||
response = await asyncio.wait_for(
|
||||
self._client.aio.models.generate_content(
|
||||
|
||||
@@ -21,6 +21,7 @@ from typing import Any
|
||||
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 hindsight_api.worker.stage import set_stage
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -141,6 +142,8 @@ class LiteLLMLLM(LLMInterface):
|
||||
last_exception = None
|
||||
|
||||
for attempt in range(max_retries + 1):
|
||||
if attempt > 0:
|
||||
set_stage(f"llm.litellm.{scope}.attempt={attempt + 1}/{max_retries + 1}")
|
||||
try:
|
||||
response = await self._litellm.acompletion(**call_kwargs)
|
||||
|
||||
@@ -283,6 +286,8 @@ class LiteLLMLLM(LLMInterface):
|
||||
|
||||
last_exception = None
|
||||
for attempt in range(max_retries + 1):
|
||||
if attempt > 0:
|
||||
set_stage(f"llm.litellm.tools.attempt={attempt + 1}/{max_retries + 1}")
|
||||
try:
|
||||
response = await self._litellm.acompletion(**call_kwargs)
|
||||
|
||||
|
||||
@@ -0,0 +1,428 @@
|
||||
"""
|
||||
Built-in llama.cpp LLM provider for fully offline operation.
|
||||
|
||||
Manages a llama-cpp-python server as a subprocess, downloads GGUF models
|
||||
from HuggingFace on first use, and delegates inference to the OpenAI-compatible API.
|
||||
|
||||
Usage:
|
||||
HINDSIGHT_API_LLM_PROVIDER=llamacpp
|
||||
HINDSIGHT_API_LLAMACPP_MODEL_PATH=~/.hindsight/models/gemma-4-E2B-it-Q4_K_M.gguf
|
||||
HINDSIGHT_API_LLAMACPP_GPU_LAYERS=-1 # -1 = all layers on GPU
|
||||
HINDSIGHT_API_LLAMACPP_CONTEXT_SIZE=8192
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import os
|
||||
import signal
|
||||
import socket
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from hindsight_api.engine.llm_interface import LLMInterface
|
||||
from hindsight_api.engine.response_models import LLMToolCallResult
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Default GGUF model for offline mode
|
||||
DEFAULT_LLAMACPP_HF_REPO = "bartowski/google_gemma-4-E2B-it-GGUF"
|
||||
DEFAULT_LLAMACPP_HF_FILENAME = "google_gemma-4-E2B-it-Q4_K_M.gguf"
|
||||
DEFAULT_LLAMACPP_MODEL_ALIAS = "gemma-4-e2b-it"
|
||||
|
||||
MODELS_DIR = Path.home() / ".hindsight" / "models"
|
||||
|
||||
# Singleton server instance — shared across all LlamaCppLLM instances
|
||||
# (retain, reflect, consolidation each create their own LLMProvider,
|
||||
# but they should all share one llama.cpp server process)
|
||||
_shared_server: "LlamaCppServer | None" = None
|
||||
_shared_server_lock = asyncio.Lock()
|
||||
|
||||
|
||||
def _find_free_port() -> int:
|
||||
"""Find a free TCP port on localhost."""
|
||||
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
|
||||
s.bind(("127.0.0.1", 0))
|
||||
return s.getsockname()[1]
|
||||
|
||||
|
||||
def _download_default_model() -> Path:
|
||||
"""Download the default GGUF model from HuggingFace if not already cached.
|
||||
|
||||
Returns:
|
||||
Path to the downloaded GGUF file.
|
||||
"""
|
||||
try:
|
||||
from huggingface_hub import hf_hub_download
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"huggingface-hub is required for automatic model download. "
|
||||
"Install with: pip install 'hindsight-api-slim[local-llm]'"
|
||||
)
|
||||
|
||||
MODELS_DIR.mkdir(parents=True, exist_ok=True)
|
||||
target = MODELS_DIR / DEFAULT_LLAMACPP_HF_FILENAME
|
||||
|
||||
if target.exists():
|
||||
logger.info(f"Using cached model: {target}")
|
||||
return target
|
||||
|
||||
logger.info(
|
||||
f"Downloading {DEFAULT_LLAMACPP_HF_FILENAME} from {DEFAULT_LLAMACPP_HF_REPO} (~3.5 GB, first run only)..."
|
||||
)
|
||||
|
||||
downloaded = hf_hub_download(
|
||||
repo_id=DEFAULT_LLAMACPP_HF_REPO,
|
||||
filename=DEFAULT_LLAMACPP_HF_FILENAME,
|
||||
local_dir=str(MODELS_DIR),
|
||||
)
|
||||
|
||||
logger.info(f"Model downloaded: {downloaded}")
|
||||
return Path(downloaded)
|
||||
|
||||
|
||||
def _resolve_model_path(model_path: str | None) -> Path:
|
||||
"""Resolve the model path, downloading the default if needed.
|
||||
|
||||
Args:
|
||||
model_path: Explicit path to a GGUF file, or None to use the default.
|
||||
|
||||
Returns:
|
||||
Resolved Path to the GGUF file.
|
||||
"""
|
||||
if model_path:
|
||||
p = Path(model_path).expanduser()
|
||||
if not p.exists():
|
||||
raise FileNotFoundError(
|
||||
f"GGUF model not found: {p}\n"
|
||||
f"Set HINDSIGHT_API_LLAMACPP_MODEL_PATH to a valid .gguf file, "
|
||||
f"or remove the setting to auto-download the default model."
|
||||
)
|
||||
return p
|
||||
|
||||
return _download_default_model()
|
||||
|
||||
|
||||
class LlamaCppServer:
|
||||
"""Manages a llama-cpp-python OpenAI-compatible server as a subprocess."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model_path: Path,
|
||||
port: int,
|
||||
gpu_layers: int = -1,
|
||||
context_size: int = 8192,
|
||||
chat_format: str | None = None,
|
||||
extra_args: str | None = None,
|
||||
):
|
||||
self.model_path = model_path
|
||||
self.port = port
|
||||
self.gpu_layers = gpu_layers
|
||||
self.context_size = context_size
|
||||
self.chat_format = chat_format
|
||||
self.extra_args = extra_args
|
||||
self._process: subprocess.Popen | None = None
|
||||
|
||||
@property
|
||||
def base_url(self) -> str:
|
||||
return f"http://127.0.0.1:{self.port}/v1"
|
||||
|
||||
async def start(self) -> None:
|
||||
"""Start the llama.cpp server subprocess."""
|
||||
cmd = [
|
||||
sys.executable,
|
||||
"-m",
|
||||
"llama_cpp.server",
|
||||
"--model",
|
||||
str(self.model_path),
|
||||
"--host",
|
||||
"127.0.0.1",
|
||||
"--port",
|
||||
str(self.port),
|
||||
"--n_gpu_layers",
|
||||
str(self.gpu_layers),
|
||||
"--n_ctx",
|
||||
str(self.context_size),
|
||||
"--flash_attn",
|
||||
"true",
|
||||
"--n_batch",
|
||||
"2048",
|
||||
# Prompt cache: reuse KV cache for repeated system prompts
|
||||
"--cache",
|
||||
"true",
|
||||
]
|
||||
# Only pass chat_format if explicitly set (most GGUF models have it embedded)
|
||||
if self.chat_format:
|
||||
cmd.extend(["--chat_format", self.chat_format])
|
||||
# User-provided extra args (e.g. "--type_k 1 --type_v 1 --n_threads 8")
|
||||
if self.extra_args:
|
||||
cmd.extend(self.extra_args.split())
|
||||
|
||||
logger.info(f"Starting llama.cpp server: {' '.join(cmd)}")
|
||||
|
||||
# Write stderr to a log file to avoid pipe buffer deadlock
|
||||
# (llama.cpp outputs a lot of model metadata on stderr during loading)
|
||||
self._log_path = MODELS_DIR / "llamacpp_server.log"
|
||||
self._log_file = open(self._log_path, "w")
|
||||
|
||||
self._process = subprocess.Popen(
|
||||
cmd,
|
||||
stdout=subprocess.DEVNULL,
|
||||
stderr=self._log_file,
|
||||
# Ensure the subprocess is killed when the parent exits
|
||||
preexec_fn=os.setsid if hasattr(os, "setsid") else None,
|
||||
)
|
||||
|
||||
# Wait for the server to be ready
|
||||
await self._wait_for_ready()
|
||||
|
||||
async def _wait_for_ready(self, timeout: float = 120.0) -> None:
|
||||
"""Wait for the llama.cpp server to accept connections."""
|
||||
import httpx
|
||||
|
||||
start = time.monotonic()
|
||||
url = f"http://127.0.0.1:{self.port}/v1/models"
|
||||
last_log = start
|
||||
|
||||
while time.monotonic() - start < timeout:
|
||||
# Check if process died
|
||||
if self._process and self._process.poll() is not None:
|
||||
stderr = ""
|
||||
try:
|
||||
stderr = self._log_path.read_text()[-2000:]
|
||||
except Exception:
|
||||
pass
|
||||
raise RuntimeError(f"llama.cpp server exited with code {self._process.returncode}.\nstderr: {stderr}")
|
||||
|
||||
try:
|
||||
async with httpx.AsyncClient() as client:
|
||||
resp = await client.get(url, timeout=5.0)
|
||||
if resp.status_code == 200:
|
||||
logger.info(f"llama.cpp server ready on port {self.port}")
|
||||
return
|
||||
except (httpx.ConnectError, httpx.TimeoutException, httpx.ConnectTimeout):
|
||||
pass
|
||||
|
||||
# Log progress every 15s
|
||||
now = time.monotonic()
|
||||
if now - last_log > 15:
|
||||
elapsed = int(now - start)
|
||||
logger.info(f"Waiting for llama.cpp server to load model... ({elapsed}s)")
|
||||
last_log = now
|
||||
|
||||
await asyncio.sleep(1.0)
|
||||
|
||||
# Timeout — read the log to help debug
|
||||
stderr = ""
|
||||
try:
|
||||
stderr = self._log_path.read_text()[-2000:]
|
||||
except Exception:
|
||||
pass
|
||||
raise TimeoutError(
|
||||
f"llama.cpp server did not become ready within {timeout}s.\n"
|
||||
f"Check model compatibility and available memory.\n"
|
||||
f"Server log: {stderr}"
|
||||
)
|
||||
|
||||
async def stop(self) -> None:
|
||||
"""Stop the llama.cpp server subprocess."""
|
||||
if self._process is None:
|
||||
return
|
||||
|
||||
logger.info("Stopping llama.cpp server...")
|
||||
try:
|
||||
# Send SIGTERM to the process group
|
||||
if hasattr(os, "killpg"):
|
||||
os.killpg(os.getpgid(self._process.pid), signal.SIGTERM)
|
||||
else:
|
||||
self._process.terminate()
|
||||
|
||||
# Wait up to 10s for graceful shutdown
|
||||
try:
|
||||
self._process.wait(timeout=10)
|
||||
except subprocess.TimeoutExpired:
|
||||
if hasattr(os, "killpg"):
|
||||
os.killpg(os.getpgid(self._process.pid), signal.SIGKILL)
|
||||
else:
|
||||
self._process.kill()
|
||||
self._process.wait(timeout=5)
|
||||
except (ProcessLookupError, OSError):
|
||||
pass # Process already exited
|
||||
finally:
|
||||
self._process = None
|
||||
if hasattr(self, "_log_file") and self._log_file:
|
||||
self._log_file.close()
|
||||
self._log_file = None
|
||||
logger.info("llama.cpp server stopped")
|
||||
|
||||
|
||||
class LlamaCppLLM(LLMInterface):
|
||||
"""
|
||||
Built-in llama.cpp provider.
|
||||
|
||||
Manages a llama-cpp-python server subprocess and delegates to OpenAICompatibleLLM
|
||||
for actual inference calls. Handles model downloading and server lifecycle.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
provider: str,
|
||||
api_key: str,
|
||||
base_url: str,
|
||||
model: str,
|
||||
reasoning_effort: str = "low",
|
||||
model_path: str | None = None,
|
||||
gpu_layers: int = -1,
|
||||
context_size: int = 8192,
|
||||
chat_format: str | None = None,
|
||||
no_grammar: bool = False,
|
||||
extra_args: str | None = None,
|
||||
**kwargs: Any,
|
||||
):
|
||||
super().__init__(
|
||||
provider=provider,
|
||||
api_key=api_key or "llamacpp",
|
||||
base_url=base_url or "",
|
||||
model=model or DEFAULT_LLAMACPP_MODEL_ALIAS,
|
||||
reasoning_effort=reasoning_effort,
|
||||
)
|
||||
self._model_path_str = model_path
|
||||
self._gpu_layers = gpu_layers
|
||||
self._context_size = context_size
|
||||
self._chat_format = chat_format
|
||||
self._no_grammar = no_grammar
|
||||
self._extra_args = extra_args
|
||||
self._server: LlamaCppServer | None = None
|
||||
self._delegate: Any = None # OpenAICompatibleLLM, created after server starts
|
||||
self._initialized = False
|
||||
|
||||
async def _ensure_initialized(self) -> None:
|
||||
"""Lazy initialization: download model + start shared server on first use."""
|
||||
if self._initialized:
|
||||
return
|
||||
|
||||
global _shared_server
|
||||
|
||||
from .openai_compatible_llm import OpenAICompatibleLLM
|
||||
|
||||
async with _shared_server_lock:
|
||||
if _shared_server is None:
|
||||
# Resolve and potentially download the model
|
||||
model_path = _resolve_model_path(self._model_path_str)
|
||||
logger.info(f"Using GGUF model: {model_path}")
|
||||
|
||||
# Start the shared llama.cpp server
|
||||
port = _find_free_port()
|
||||
_shared_server = LlamaCppServer(
|
||||
model_path=model_path,
|
||||
port=port,
|
||||
gpu_layers=self._gpu_layers,
|
||||
context_size=self._context_size,
|
||||
chat_format=self._chat_format,
|
||||
extra_args=self._extra_args,
|
||||
)
|
||||
await _shared_server.start()
|
||||
|
||||
self._server = _shared_server
|
||||
|
||||
# Create the delegate that talks to the shared server's OpenAI-compatible API
|
||||
if self._no_grammar:
|
||||
logger.info("Grammar enforcement disabled (HINDSIGHT_API_LLAMACPP_NO_GRAMMAR=true)")
|
||||
self._delegate = OpenAICompatibleLLM(
|
||||
provider="llamacpp",
|
||||
api_key="llamacpp",
|
||||
base_url=self._server.base_url,
|
||||
model=self.model,
|
||||
reasoning_effort=self.reasoning_effort,
|
||||
)
|
||||
|
||||
self._initialized = True
|
||||
|
||||
async def verify_connection(self) -> None:
|
||||
"""Verify the llama.cpp server is running and can generate text."""
|
||||
await self._ensure_initialized()
|
||||
# Make a simple test call to verify the model can actually generate
|
||||
await self._delegate.call(
|
||||
messages=[{"role": "user", "content": "Say 'ok'"}],
|
||||
max_completion_tokens=10,
|
||||
max_retries=2,
|
||||
initial_backoff=0.5,
|
||||
max_backoff=2.0,
|
||||
scope="verification",
|
||||
)
|
||||
logger.info("llama.cpp LLM verification passed")
|
||||
|
||||
async def call(
|
||||
self,
|
||||
messages: list[dict[str, str]],
|
||||
response_format: Any | None = None,
|
||||
max_completion_tokens: int | None = None,
|
||||
temperature: float | None = None,
|
||||
scope: str = "memory",
|
||||
max_retries: int = 10,
|
||||
initial_backoff: float = 1.0,
|
||||
max_backoff: float = 60.0,
|
||||
skip_validation: bool = False,
|
||||
strict_schema: bool = False,
|
||||
return_usage: bool = False,
|
||||
) -> Any:
|
||||
"""Delegate call to the OpenAI-compatible API."""
|
||||
await self._ensure_initialized()
|
||||
return await self._delegate.call(
|
||||
messages=messages,
|
||||
response_format=response_format,
|
||||
max_completion_tokens=max_completion_tokens,
|
||||
temperature=temperature,
|
||||
scope=scope,
|
||||
max_retries=max_retries,
|
||||
initial_backoff=initial_backoff,
|
||||
max_backoff=max_backoff,
|
||||
skip_validation=skip_validation,
|
||||
strict_schema=strict_schema,
|
||||
return_usage=return_usage,
|
||||
)
|
||||
|
||||
async def call_with_tools(
|
||||
self,
|
||||
messages: list[dict[str, Any]],
|
||||
tools: list[dict[str, Any]],
|
||||
max_completion_tokens: int | None = None,
|
||||
temperature: float | None = None,
|
||||
scope: str = "tools",
|
||||
max_retries: int = 5,
|
||||
initial_backoff: float = 1.0,
|
||||
max_backoff: float = 30.0,
|
||||
tool_choice: str | dict[str, Any] = "auto",
|
||||
) -> LLMToolCallResult:
|
||||
"""Delegate tool calls to the OpenAI-compatible API."""
|
||||
await self._ensure_initialized()
|
||||
return await self._delegate.call_with_tools(
|
||||
messages=messages,
|
||||
tools=tools,
|
||||
max_completion_tokens=max_completion_tokens,
|
||||
temperature=temperature,
|
||||
scope=scope,
|
||||
max_retries=max_retries,
|
||||
initial_backoff=initial_backoff,
|
||||
max_backoff=max_backoff,
|
||||
tool_choice=tool_choice,
|
||||
)
|
||||
|
||||
async def cleanup(self) -> None:
|
||||
"""Stop the shared llama.cpp server."""
|
||||
global _shared_server
|
||||
|
||||
if self._delegate:
|
||||
await self._delegate.cleanup()
|
||||
self._delegate = None
|
||||
|
||||
# Stop the shared server (only the first cleanup call actually stops it)
|
||||
async with _shared_server_lock:
|
||||
if _shared_server is not None:
|
||||
await _shared_server.stop()
|
||||
_shared_server = None
|
||||
|
||||
self._server = None
|
||||
self._initialized = False
|
||||
@@ -33,6 +33,7 @@ from hindsight_api.config import DEFAULT_LLM_TIMEOUT, ENV_LLM_TIMEOUT
|
||||
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 hindsight_api.worker.stage import set_stage
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -59,6 +60,32 @@ def _strip_code_fences(content: str) -> str:
|
||||
return content
|
||||
|
||||
|
||||
def _summarize_status_error(e: APIStatusError, body_max: int = 400) -> str:
|
||||
"""Render an APIStatusError with status code + truncated response body.
|
||||
|
||||
Without this, retry loops only log "API error after N attempts" with the
|
||||
bare exception message — losing the provider's actual error payload, which
|
||||
is the only thing that explains *why* a request failed (rate limit reason,
|
||||
invalid tool schema, model overloaded, etc.).
|
||||
"""
|
||||
body: Any = getattr(e, "body", None)
|
||||
if body is None:
|
||||
try:
|
||||
body = e.response.text
|
||||
except Exception:
|
||||
body = None
|
||||
if isinstance(body, (dict, list)):
|
||||
try:
|
||||
body_str = json.dumps(body, default=str, ensure_ascii=False)
|
||||
except Exception:
|
||||
body_str = str(body)
|
||||
else:
|
||||
body_str = str(body or "").strip()
|
||||
if len(body_str) > body_max:
|
||||
body_str = body_str[:body_max] + "...TRUNCATED"
|
||||
return f"HTTP {e.status_code}: {body_str or '<no body>'}"
|
||||
|
||||
|
||||
class OpenAICompatibleLLM(LLMInterface):
|
||||
"""
|
||||
LLM provider for OpenAI-compatible APIs.
|
||||
@@ -100,7 +127,7 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
super().__init__(provider, api_key, base_url, model, reasoning_effort, **kwargs)
|
||||
|
||||
# Validate provider
|
||||
valid_providers = ["openai", "groq", "ollama", "lmstudio", "minimax", "volcano"]
|
||||
valid_providers = ["openai", "groq", "ollama", "lmstudio", "llamacpp", "minimax", "volcano", "openrouter"]
|
||||
if self.provider not in valid_providers:
|
||||
raise ValueError(f"OpenAICompatibleLLM only supports: {', '.join(valid_providers)}. Got: {self.provider}")
|
||||
|
||||
@@ -114,13 +141,15 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
self.base_url = "http://localhost:1234/v1"
|
||||
elif self.provider == "minimax":
|
||||
self.base_url = "https://api.minimax.io/v1"
|
||||
elif self.provider == "openrouter":
|
||||
self.base_url = "https://openrouter.ai/api/v1"
|
||||
|
||||
# For ollama/lmstudio, use dummy key if not provided
|
||||
if self.provider in ("ollama", "lmstudio") and not self.api_key:
|
||||
self.api_key = "local"
|
||||
|
||||
# Validate API key for cloud providers
|
||||
if self.provider in ("openai", "groq", "minimax") and not self.api_key:
|
||||
if self.provider in ("openai", "groq", "minimax", "openrouter") and not self.api_key:
|
||||
raise ValueError(f"API key is required for {self.provider}")
|
||||
|
||||
# Service tier configuration (from config, not env vars)
|
||||
@@ -191,6 +220,36 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
|
||||
return None
|
||||
|
||||
def _max_tokens_param_name(self) -> str:
|
||||
"""Return the correct parameter name for limiting response tokens.
|
||||
|
||||
Native OpenAI, Azure OpenAI, Groq, and llamacpp accept 'max_completion_tokens'.
|
||||
Mistral and other OpenAI-compatible endpoints that haven't adopted the newer
|
||||
parameter name require 'max_tokens', so when the openai provider is configured
|
||||
with a non-Azure custom base_url we fall back to the widely-supported
|
||||
'max_tokens'.
|
||||
|
||||
Reasoning models (GPT-5, o1, o3) only accept 'max_completion_tokens' and reject
|
||||
'max_tokens' outright, so they always use the new parameter name regardless of
|
||||
base_url.
|
||||
"""
|
||||
# Reasoning models (GPT-5, o1, o3, ...) only accept max_completion_tokens.
|
||||
# Azure OpenAI + GPT-5 is the canonical example: issue #978.
|
||||
if self._supports_reasoning_model():
|
||||
return "max_completion_tokens"
|
||||
# Native OpenAI (no custom base URL), Groq, and llamacpp use max_completion_tokens
|
||||
if self.provider in ("groq", "llamacpp"):
|
||||
return "max_completion_tokens"
|
||||
if self.provider == "openai" and not self.base_url:
|
||||
return "max_completion_tokens"
|
||||
# Azure OpenAI is fully OpenAI-API-compatible — detect it by hostname so users
|
||||
# can keep provider=openai + an Azure base_url (the documented setup).
|
||||
if self.provider == "openai" and self.base_url and ".openai.azure.com" in self.base_url:
|
||||
return "max_completion_tokens"
|
||||
# openai with custom base_url, ollama, lmstudio, minimax, volcano —
|
||||
# use the widely-supported max_tokens
|
||||
return "max_tokens"
|
||||
|
||||
async def call(
|
||||
self,
|
||||
messages: list[dict[str, str]],
|
||||
@@ -263,9 +322,7 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
# For reasoning models, enforce minimum to ensure space for reasoning + output
|
||||
if is_reasoning_model and max_completion_tokens < 16000:
|
||||
max_completion_tokens = 16000
|
||||
call_params["max_completion_tokens"] = max_completion_tokens
|
||||
|
||||
# Temperature - reasoning models don't support custom temperature
|
||||
call_params[self._max_tokens_param_name()] = max_completion_tokens
|
||||
if temperature is not None and not is_reasoning_model:
|
||||
# MiniMax requires temperature in (0.0, 1.0] — clamp accordingly
|
||||
if self.provider == "minimax":
|
||||
@@ -308,9 +365,7 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
else:
|
||||
# Soft enforcement: add schema to prompt and use json_object mode
|
||||
if schema is not None:
|
||||
schema_msg = (
|
||||
f"\n\nYou must respond with valid JSON matching this schema:\n{json.dumps(schema, indent=2)}"
|
||||
)
|
||||
schema_msg = f"\n\nYou must respond with valid JSON matching this schema:\n{json.dumps(schema, indent=2, ensure_ascii=False)}"
|
||||
|
||||
if call_params["messages"] and call_params["messages"][0].get("role") == "system":
|
||||
first_msg = call_params["messages"][0]
|
||||
@@ -320,13 +375,23 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
first_msg = call_params["messages"][0]
|
||||
if isinstance(first_msg, dict) and isinstance(first_msg.get("content"), str):
|
||||
first_msg["content"] = schema_msg + "\n\n" + first_msg["content"]
|
||||
if self.provider not in ("lmstudio", "ollama", "volcano"):
|
||||
# LM Studio, Ollama and Volcano don't support json_object response format reliably
|
||||
# Providers that skip json_object grammar enforcement
|
||||
skip_grammar = self.provider in ("lmstudio", "ollama", "volcano")
|
||||
if self.provider == "llamacpp":
|
||||
from hindsight_api.config import get_config
|
||||
|
||||
skip_grammar = get_config().llamacpp_no_grammar
|
||||
if not skip_grammar:
|
||||
call_params["response_format"] = {"type": "json_object"}
|
||||
|
||||
last_exception = None
|
||||
|
||||
for attempt in range(max_retries + 1):
|
||||
# Surface attempt count in worker stage so JSON-schema retry loops
|
||||
# are visible from logs (small models on strict structured output
|
||||
# often loop here). Cheap no-op outside worker context.
|
||||
if attempt > 0:
|
||||
set_stage(f"llm.{self.provider}.{scope}.attempt={attempt + 1}/{max_retries + 1}")
|
||||
try:
|
||||
if response_format is not None:
|
||||
response = await self._client.chat.completions.create(**call_params)
|
||||
@@ -509,12 +574,19 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
|
||||
last_exception = e
|
||||
if attempt < max_retries:
|
||||
logger.warning(
|
||||
f"APIStatusError ({self.provider}/{self.model}, scope={scope}, "
|
||||
f"attempt {attempt + 1}/{max_retries + 1}): {_summarize_status_error(e)}"
|
||||
)
|
||||
backoff = min(initial_backoff * (2**attempt), max_backoff)
|
||||
jitter = backoff * 0.2 * (2 * (time.time() % 1) - 1)
|
||||
sleep_time = backoff + jitter
|
||||
await asyncio.sleep(sleep_time)
|
||||
else:
|
||||
logger.error(f"API error after {max_retries + 1} attempts: {str(e)}")
|
||||
logger.error(
|
||||
f"API error after {max_retries + 1} attempts ({self.provider}/{self.model}, "
|
||||
f"scope={scope}): {_summarize_status_error(e)}"
|
||||
)
|
||||
raise
|
||||
|
||||
except Exception:
|
||||
@@ -577,7 +649,7 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
}
|
||||
|
||||
if max_completion_tokens is not None:
|
||||
call_params["max_completion_tokens"] = max_completion_tokens
|
||||
call_params[self._max_tokens_param_name()] = max_completion_tokens
|
||||
if temperature is not None:
|
||||
# MiniMax requires temperature in (0.0, 1.0] — clamp accordingly
|
||||
if self.provider == "minimax":
|
||||
@@ -594,6 +666,8 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
last_exception = None
|
||||
|
||||
for attempt in range(max_retries + 1):
|
||||
if attempt > 0:
|
||||
set_stage(f"llm.{self.provider}.tools.attempt={attempt + 1}/{max_retries + 1}")
|
||||
try:
|
||||
response = await self._client.chat.completions.create(**call_params)
|
||||
|
||||
@@ -663,18 +737,41 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
|
||||
except APIConnectionError as e:
|
||||
last_exception = e
|
||||
status_code = getattr(e, "status_code", None) or getattr(
|
||||
getattr(e, "response", None), "status_code", None
|
||||
)
|
||||
if attempt < max_retries:
|
||||
logger.warning(
|
||||
f"APIConnectionError in tool call ({self.provider}/{self.model}, scope={scope}, "
|
||||
f"attempt {attempt + 1}/{max_retries + 1}, HTTP {status_code}): {str(e)[:200]}"
|
||||
)
|
||||
await asyncio.sleep(min(initial_backoff * (2**attempt), max_backoff))
|
||||
continue
|
||||
logger.error(
|
||||
f"Connection error in tool call after {max_retries + 1} attempts "
|
||||
f"({self.provider}/{self.model}, scope={scope}): {str(e)}"
|
||||
)
|
||||
raise
|
||||
|
||||
except APIStatusError as e:
|
||||
if e.status_code in (401, 403):
|
||||
logger.error(
|
||||
f"Auth error in tool call (HTTP {e.status_code}, {self.provider}/{self.model}), "
|
||||
f"not retrying: {_summarize_status_error(e)}"
|
||||
)
|
||||
raise
|
||||
last_exception = e
|
||||
if attempt < max_retries:
|
||||
logger.warning(
|
||||
f"APIStatusError in tool call ({self.provider}/{self.model}, scope={scope}, "
|
||||
f"attempt {attempt + 1}/{max_retries + 1}): {_summarize_status_error(e)}"
|
||||
)
|
||||
await asyncio.sleep(min(initial_backoff * (2**attempt), max_backoff))
|
||||
continue
|
||||
logger.error(
|
||||
f"API error in tool call after {max_retries + 1} attempts "
|
||||
f"({self.provider}/{self.model}, scope={scope}): {_summarize_status_error(e)}"
|
||||
)
|
||||
raise
|
||||
|
||||
except Exception:
|
||||
@@ -722,6 +819,7 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
"model": self.model,
|
||||
"messages": messages,
|
||||
"stream": False,
|
||||
"think": False, # Disable thinking for reasoning models (qwen3.5, etc.)
|
||||
}
|
||||
|
||||
# Add schema as format parameter for structured output
|
||||
@@ -743,6 +841,8 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
|
||||
async with httpx.AsyncClient(timeout=300.0) as client:
|
||||
for attempt in range(max_retries + 1):
|
||||
if attempt > 0:
|
||||
set_stage(f"llm.ollama_native.{scope}.attempt={attempt + 1}/{max_retries + 1}")
|
||||
try:
|
||||
response = await client.post(native_url, json=payload)
|
||||
response.raise_for_status()
|
||||
@@ -874,7 +974,7 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
logger.info(f"Submitting batch with {len(requests)} requests to {self.provider}")
|
||||
|
||||
# Format requests as JSONL
|
||||
jsonl_content = "\n".join(json.dumps(req) for req in requests)
|
||||
jsonl_content = "\n".join(json.dumps(req, ensure_ascii=False) for req in requests)
|
||||
|
||||
# Upload file to provider (wrap in BytesIO with filename)
|
||||
file_bytes = io.BytesIO(jsonl_content.encode("utf-8"))
|
||||
|
||||
@@ -137,7 +137,21 @@ class DateparserQueryAnalyzer(QueryAnalyzer):
|
||||
"RETURN_AS_TIMEZONE_AWARE": False,
|
||||
}
|
||||
|
||||
results = self._search_dates(query, settings=settings)
|
||||
# Wrap dateparser in a defensive try/except. dateparser has been
|
||||
# observed to crash with internal errors (e.g., IndexError from
|
||||
# locale.translate_search) on certain query inputs. A parser bug
|
||||
# should not bring down the whole search/consolidation pipeline —
|
||||
# treat any failure as "no temporal constraint found" so the caller
|
||||
# can fall back to non-temporal retrieval.
|
||||
try:
|
||||
results = self._search_dates(query, settings=settings)
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
"dateparser raised %s on query (treating as no temporal constraint): %s",
|
||||
type(e).__name__,
|
||||
e,
|
||||
)
|
||||
return QueryAnalysis(temporal_constraint=None)
|
||||
|
||||
if not results:
|
||||
return QueryAnalysis(temporal_constraint=None)
|
||||
|
||||
@@ -17,7 +17,12 @@ from typing import TYPE_CHECKING, Any, Awaitable, Callable
|
||||
import tiktoken
|
||||
|
||||
from .models import DirectiveInfo, LLMCall, ReflectAgentResult, TokenUsageSummary, ToolCall
|
||||
from .prompts import FINAL_SYSTEM_PROMPT, _extract_directive_rules, build_final_prompt, build_system_prompt_for_tools
|
||||
from .prompts import (
|
||||
_extract_directive_rules,
|
||||
build_final_prompt,
|
||||
build_final_system_prompt,
|
||||
build_system_prompt_for_tools,
|
||||
)
|
||||
from .tools_schema import get_reflect_tools
|
||||
|
||||
|
||||
@@ -186,7 +191,7 @@ async def _generate_structured_output(
|
||||
DynamicModel = create_model("StructuredResponse", **fields)
|
||||
|
||||
# Include the full schema in the prompt for better LLM guidance
|
||||
schema_str = json.dumps(response_schema, indent=2)
|
||||
schema_str = json.dumps(response_schema, indent=2, ensure_ascii=False)
|
||||
|
||||
# Build field descriptions for the prompt
|
||||
field_descriptions = []
|
||||
@@ -446,7 +451,7 @@ async def run_reflect_agent(
|
||||
llm_start = time.time()
|
||||
response, usage = await llm_config.call(
|
||||
messages=[
|
||||
{"role": "system", "content": FINAL_SYSTEM_PROMPT},
|
||||
{"role": "system", "content": build_final_system_prompt(bank_profile.get("mission"))},
|
||||
{"role": "user", "content": prompt},
|
||||
],
|
||||
scope="reflect",
|
||||
@@ -503,7 +508,7 @@ async def run_reflect_agent(
|
||||
llm_start = time.time()
|
||||
response, usage = await llm_config.call(
|
||||
messages=[
|
||||
{"role": "system", "content": FINAL_SYSTEM_PROMPT},
|
||||
{"role": "system", "content": build_final_system_prompt(bank_profile.get("mission"))},
|
||||
{"role": "user", "content": prompt},
|
||||
],
|
||||
scope="reflect",
|
||||
@@ -606,7 +611,7 @@ async def run_reflect_agent(
|
||||
llm_start = time.time()
|
||||
response, usage = await llm_config.call(
|
||||
messages=[
|
||||
{"role": "system", "content": FINAL_SYSTEM_PROMPT},
|
||||
{"role": "system", "content": build_final_system_prompt(bank_profile.get("mission"))},
|
||||
{"role": "user", "content": prompt},
|
||||
],
|
||||
scope="reflect",
|
||||
@@ -649,9 +654,57 @@ async def run_reflect_agent(
|
||||
|
||||
# No tool calls - LLM wants to respond with text
|
||||
if not result.tool_calls:
|
||||
if result.content:
|
||||
# When directives are present but no evidence has been gathered,
|
||||
# the LLM tends to echo directive content verbatim as its answer.
|
||||
# Fall through to the final-prompt path which doesn't include
|
||||
# directives and handles "no data" gracefully.
|
||||
has_gathered_evidence = (
|
||||
bool(available_memory_ids) or bool(available_mental_model_ids) or bool(available_observation_ids)
|
||||
)
|
||||
directive_leak_risk = directives and not has_gathered_evidence
|
||||
if result.content and not directive_leak_risk:
|
||||
answer = _clean_answer_text(result.content.strip())
|
||||
|
||||
# The call_with_tools call above is intentionally uncapped so the
|
||||
# LLM has headroom to emit tool-call JSON plus any intermediate
|
||||
# reasoning. But when the LLM short-circuits and returns text
|
||||
# directly, that text becomes the user-visible final answer and
|
||||
# must respect max_tokens like the forced-final paths do. If it
|
||||
# overshoots, run one extra capped call to rewrite it within
|
||||
# the cap.
|
||||
if max_tokens is not None and len(_TIKTOKEN_ENCODING.encode(answer)) > max_tokens:
|
||||
rewrite_start = time.time()
|
||||
rewritten, rewrite_usage = await llm_config.call(
|
||||
messages=[
|
||||
{
|
||||
"role": "system",
|
||||
"content": (
|
||||
"Rewrite the user's text so it fits within the requested token "
|
||||
"budget. Preserve the key facts and structure; drop lower-priority "
|
||||
"detail. Respond with the rewritten text only, no preamble."
|
||||
),
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": f"Target budget: {max_tokens} tokens.\n\nText to rewrite:\n{answer}",
|
||||
},
|
||||
],
|
||||
scope="reflect",
|
||||
max_completion_tokens=max_tokens,
|
||||
return_usage=True,
|
||||
)
|
||||
total_input_tokens += rewrite_usage.input_tokens
|
||||
total_output_tokens += rewrite_usage.output_tokens
|
||||
llm_trace.append(
|
||||
{
|
||||
"scope": "final_rewrite",
|
||||
"duration_ms": int((time.time() - rewrite_start) * 1000),
|
||||
"input_tokens": rewrite_usage.input_tokens,
|
||||
"output_tokens": rewrite_usage.output_tokens,
|
||||
}
|
||||
)
|
||||
answer = _clean_answer_text(rewritten.strip())
|
||||
|
||||
# Generate structured output if schema provided
|
||||
structured_output = None
|
||||
if response_schema and answer:
|
||||
@@ -679,7 +732,7 @@ async def run_reflect_agent(
|
||||
llm_start = time.time()
|
||||
response, usage = await llm_config.call(
|
||||
messages=[
|
||||
{"role": "system", "content": FINAL_SYSTEM_PROMPT},
|
||||
{"role": "system", "content": build_final_system_prompt(bank_profile.get("mission"))},
|
||||
{"role": "user", "content": prompt},
|
||||
],
|
||||
scope="reflect",
|
||||
@@ -743,7 +796,8 @@ async def run_reflect_agent(
|
||||
"content": json.dumps(
|
||||
{
|
||||
"error": "You must search for information first. Use search_mental_models(), search_observations(), or recall() before providing your final answer."
|
||||
}
|
||||
},
|
||||
ensure_ascii=False,
|
||||
),
|
||||
}
|
||||
)
|
||||
@@ -805,7 +859,8 @@ async def run_reflect_agent(
|
||||
"content": json.dumps(
|
||||
{
|
||||
"error": f"Tool '{_normalize_tool_name(tc.name)}' is not available. Use only the tools provided to you."
|
||||
}
|
||||
},
|
||||
ensure_ascii=False,
|
||||
),
|
||||
}
|
||||
)
|
||||
@@ -876,7 +931,7 @@ async def run_reflect_agent(
|
||||
"role": "tool",
|
||||
"tool_call_id": tc.id,
|
||||
"name": tc.name, # Required by Gemini
|
||||
"content": json.dumps(output, default=str),
|
||||
"content": json.dumps(output, default=str, ensure_ascii=False),
|
||||
}
|
||||
)
|
||||
|
||||
@@ -899,7 +954,7 @@ async def run_reflect_agent(
|
||||
)
|
||||
|
||||
try:
|
||||
output_chars = len(json.dumps(output))
|
||||
output_chars = len(json.dumps(output, ensure_ascii=False))
|
||||
except (TypeError, ValueError):
|
||||
output_chars = len(str(output))
|
||||
|
||||
@@ -936,7 +991,7 @@ def _tool_call_to_dict(tc: "LLMToolCall") -> dict[str, Any]:
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": tc.name,
|
||||
"arguments": json.dumps(tc.arguments),
|
||||
"arguments": json.dumps(tc.arguments, ensure_ascii=False),
|
||||
},
|
||||
}
|
||||
if tc.thought_signature is not None:
|
||||
@@ -1034,7 +1089,7 @@ async def _execute_tool_with_timing(
|
||||
# Set attributes
|
||||
span.set_attribute("hindsight.tool.name", normalized_name)
|
||||
span.set_attribute("hindsight.tool.id", tc.id)
|
||||
span.set_attribute("hindsight.tool.arguments", json.dumps(tc.arguments))
|
||||
span.set_attribute("hindsight.tool.arguments", json.dumps(tc.arguments, ensure_ascii=False))
|
||||
|
||||
try:
|
||||
result = await _execute_tool(
|
||||
|
||||
@@ -0,0 +1,307 @@
|
||||
"""Delta operations for structured mental models.
|
||||
|
||||
The LLM's job during a delta refresh is to emit a list of these operations,
|
||||
each targeting an existing section (by id) or referencing a position relative
|
||||
to one. ``apply_operations`` validates and applies each op in turn against a
|
||||
copy of the document; invalid ops (unknown ``section_id``, out-of-range
|
||||
``block_index``, malformed payloads) are dropped with a debug-friendly reason.
|
||||
|
||||
Sections and blocks not mentioned by any op are physically copied through
|
||||
unchanged — there is no LLM-mediated re-emission of unchanged text, so prose
|
||||
drift is structurally impossible.
|
||||
|
||||
Why operations and not "output the new structured doc":
|
||||
- "Output the new doc" still asks the LLM to *generate* every section's
|
||||
blocks, including ones it didn't intend to modify, which gives it the same
|
||||
opportunity to drift.
|
||||
- Operations make the no-change case mechanical: zero ops → identical doc.
|
||||
- Operations are auditable: each refresh produces a log of exactly what
|
||||
changed, useful for debugging the LLM's behaviour and explaining diffs.
|
||||
|
||||
Failure modes are by design conservative: an operation list that fails to
|
||||
parse against the Pydantic schema, or an LLM that returns invalid ops, results
|
||||
in zero changes — the document stays as-is. The structure can only get better
|
||||
or stay the same per refresh, never get worse.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from typing import Annotated, Any, Literal, Union
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
|
||||
from .structured_doc import (
|
||||
Block,
|
||||
Section,
|
||||
StructuredDocument,
|
||||
make_unique_id,
|
||||
slugify_heading,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# Op payloads ---------------------------------------------------------------
|
||||
|
||||
|
||||
class _OpBase(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
|
||||
class AppendBlockOp(_OpBase):
|
||||
"""Add a new block at the end of an existing section."""
|
||||
|
||||
op: Literal["append_block"] = "append_block"
|
||||
section_id: str
|
||||
block: Block
|
||||
|
||||
|
||||
class InsertBlockOp(_OpBase):
|
||||
"""Insert a new block at ``index`` in an existing section.
|
||||
|
||||
``index`` may equal ``len(section.blocks)`` (append) but not be greater.
|
||||
"""
|
||||
|
||||
op: Literal["insert_block"] = "insert_block"
|
||||
section_id: str
|
||||
index: int = Field(ge=0)
|
||||
block: Block
|
||||
|
||||
|
||||
class ReplaceBlockOp(_OpBase):
|
||||
"""Replace the block at ``index`` of an existing section."""
|
||||
|
||||
op: Literal["replace_block"] = "replace_block"
|
||||
section_id: str
|
||||
index: int = Field(ge=0)
|
||||
block: Block
|
||||
|
||||
|
||||
class RemoveBlockOp(_OpBase):
|
||||
"""Remove the block at ``index`` of an existing section."""
|
||||
|
||||
op: Literal["remove_block"] = "remove_block"
|
||||
section_id: str
|
||||
index: int = Field(ge=0)
|
||||
|
||||
|
||||
class AddSectionOp(_OpBase):
|
||||
"""Add a brand-new section.
|
||||
|
||||
``after_section_id`` is optional; when omitted the new section is appended
|
||||
at the end. ``new_id`` is optional; when omitted we slugify the heading
|
||||
and disambiguate against existing IDs.
|
||||
"""
|
||||
|
||||
op: Literal["add_section"] = "add_section"
|
||||
heading: str
|
||||
level: int = Field(default=2, ge=1, le=6)
|
||||
blocks: list[Block] = Field(default_factory=list)
|
||||
after_section_id: str | None = None
|
||||
new_id: str | None = None
|
||||
|
||||
|
||||
class RemoveSectionOp(_OpBase):
|
||||
"""Remove an entire section by id."""
|
||||
|
||||
op: Literal["remove_section"] = "remove_section"
|
||||
section_id: str
|
||||
|
||||
|
||||
class ReplaceSectionBlocksOp(_OpBase):
|
||||
"""Replace all blocks of a section in one go.
|
||||
|
||||
Used when most of a section's contents are stale and rebuilding it as a
|
||||
unit is clearer than emitting many block-level ops. The section's heading
|
||||
and id are preserved.
|
||||
"""
|
||||
|
||||
op: Literal["replace_section_blocks"] = "replace_section_blocks"
|
||||
section_id: str
|
||||
blocks: list[Block] = Field(default_factory=list)
|
||||
|
||||
|
||||
class RenameSectionOp(_OpBase):
|
||||
"""Rename a section's heading. The id is unchanged so future ops still resolve."""
|
||||
|
||||
op: Literal["rename_section"] = "rename_section"
|
||||
section_id: str
|
||||
new_heading: str
|
||||
|
||||
|
||||
Operation = Annotated[
|
||||
Union[
|
||||
AppendBlockOp,
|
||||
InsertBlockOp,
|
||||
ReplaceBlockOp,
|
||||
RemoveBlockOp,
|
||||
AddSectionOp,
|
||||
RemoveSectionOp,
|
||||
ReplaceSectionBlocksOp,
|
||||
RenameSectionOp,
|
||||
],
|
||||
Field(discriminator="op"),
|
||||
]
|
||||
|
||||
|
||||
class DeltaOperationList(BaseModel):
|
||||
"""Container for the operations produced by an LLM delta call."""
|
||||
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
operations: list[Operation] = Field(default_factory=list)
|
||||
|
||||
|
||||
# Application ---------------------------------------------------------------
|
||||
|
||||
|
||||
class AppliedDelta(BaseModel):
|
||||
"""Outcome of applying a list of operations to a document."""
|
||||
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
document: StructuredDocument
|
||||
applied: list[dict[str, Any]] = Field(default_factory=list)
|
||||
skipped: list[dict[str, Any]] = Field(default_factory=list)
|
||||
|
||||
@property
|
||||
def changed(self) -> bool:
|
||||
return len(self.applied) > 0
|
||||
|
||||
|
||||
def _op_summary(op: Operation) -> dict[str, Any]:
|
||||
"""Compact dict suitable for the audit trail."""
|
||||
data = op.model_dump()
|
||||
return {k: v for k, v in data.items() if k != "block" and k != "blocks"} | {
|
||||
"op": data["op"],
|
||||
}
|
||||
|
||||
|
||||
def apply_operations(
|
||||
doc: StructuredDocument,
|
||||
operations: list[Operation],
|
||||
) -> AppliedDelta:
|
||||
"""Apply a list of operations to a document, returning a new document.
|
||||
|
||||
The original document is never mutated. Invalid operations (unknown
|
||||
section, out-of-range index, name collision when adding a section) are
|
||||
skipped and recorded in ``skipped`` with a ``reason`` string.
|
||||
"""
|
||||
new_doc = doc.model_copy(deep=True)
|
||||
applied: list[dict[str, Any]] = []
|
||||
skipped: list[dict[str, Any]] = []
|
||||
|
||||
def skip(op: Operation, reason: str) -> None:
|
||||
entry = _op_summary(op)
|
||||
entry["reason"] = reason
|
||||
skipped.append(entry)
|
||||
logger.debug(f"[STRUCTURED_DELTA] skipping op {entry}")
|
||||
|
||||
for op in operations:
|
||||
if isinstance(op, AppendBlockOp):
|
||||
section = new_doc.section_by_id(op.section_id)
|
||||
if section is None:
|
||||
skip(op, f"unknown section_id: {op.section_id}")
|
||||
continue
|
||||
section.blocks.append(op.block)
|
||||
applied.append(_op_summary(op))
|
||||
continue
|
||||
|
||||
if isinstance(op, InsertBlockOp):
|
||||
section = new_doc.section_by_id(op.section_id)
|
||||
if section is None:
|
||||
skip(op, f"unknown section_id: {op.section_id}")
|
||||
continue
|
||||
if op.index > len(section.blocks):
|
||||
skip(
|
||||
op,
|
||||
f"index out of range: {op.index} > {len(section.blocks)}",
|
||||
)
|
||||
continue
|
||||
section.blocks.insert(op.index, op.block)
|
||||
applied.append(_op_summary(op))
|
||||
continue
|
||||
|
||||
if isinstance(op, ReplaceBlockOp):
|
||||
section = new_doc.section_by_id(op.section_id)
|
||||
if section is None:
|
||||
skip(op, f"unknown section_id: {op.section_id}")
|
||||
continue
|
||||
if op.index >= len(section.blocks):
|
||||
skip(
|
||||
op,
|
||||
f"index out of range: {op.index} >= {len(section.blocks)}",
|
||||
)
|
||||
continue
|
||||
section.blocks[op.index] = op.block
|
||||
applied.append(_op_summary(op))
|
||||
continue
|
||||
|
||||
if isinstance(op, RemoveBlockOp):
|
||||
section = new_doc.section_by_id(op.section_id)
|
||||
if section is None:
|
||||
skip(op, f"unknown section_id: {op.section_id}")
|
||||
continue
|
||||
if op.index >= len(section.blocks):
|
||||
skip(
|
||||
op,
|
||||
f"index out of range: {op.index} >= {len(section.blocks)}",
|
||||
)
|
||||
continue
|
||||
section.blocks.pop(op.index)
|
||||
applied.append(_op_summary(op))
|
||||
continue
|
||||
|
||||
if isinstance(op, AddSectionOp):
|
||||
existing_ids = {s.id for s in new_doc.sections}
|
||||
base_id = op.new_id or slugify_heading(op.heading)
|
||||
section_id = make_unique_id(base_id, existing_ids)
|
||||
new_section = Section(
|
||||
id=section_id,
|
||||
heading=op.heading,
|
||||
level=op.level,
|
||||
blocks=list(op.blocks),
|
||||
)
|
||||
if op.after_section_id is None:
|
||||
new_doc.sections.append(new_section)
|
||||
else:
|
||||
idx = new_doc.section_index(op.after_section_id)
|
||||
if idx is None:
|
||||
skip(op, f"unknown after_section_id: {op.after_section_id}")
|
||||
continue
|
||||
new_doc.sections.insert(idx + 1, new_section)
|
||||
entry = _op_summary(op)
|
||||
entry["assigned_id"] = section_id
|
||||
applied.append(entry)
|
||||
continue
|
||||
|
||||
if isinstance(op, RemoveSectionOp):
|
||||
idx = new_doc.section_index(op.section_id)
|
||||
if idx is None:
|
||||
skip(op, f"unknown section_id: {op.section_id}")
|
||||
continue
|
||||
new_doc.sections.pop(idx)
|
||||
applied.append(_op_summary(op))
|
||||
continue
|
||||
|
||||
if isinstance(op, ReplaceSectionBlocksOp):
|
||||
section = new_doc.section_by_id(op.section_id)
|
||||
if section is None:
|
||||
skip(op, f"unknown section_id: {op.section_id}")
|
||||
continue
|
||||
section.blocks = list(op.blocks)
|
||||
applied.append(_op_summary(op))
|
||||
continue
|
||||
|
||||
if isinstance(op, RenameSectionOp):
|
||||
section = new_doc.section_by_id(op.section_id)
|
||||
if section is None:
|
||||
skip(op, f"unknown section_id: {op.section_id}")
|
||||
continue
|
||||
section.heading = op.new_heading
|
||||
applied.append(_op_summary(op))
|
||||
continue
|
||||
|
||||
skip(op, f"unhandled op type: {type(op).__name__}") # pragma: no cover
|
||||
|
||||
return AppliedDelta(document=new_doc, applied=applied, skipped=skipped)
|
||||
@@ -18,6 +18,9 @@ _TIKTOKEN_ENCODING = tiktoken.get_encoding("cl100k_base")
|
||||
# The remainder covers the system prompt, question, bank context, and output tokens.
|
||||
_FINAL_PROMPT_CONTEXT_FRACTION = 0.8
|
||||
|
||||
_DEFAULT_ROLE = "You are a reflection agent that answers questions by reasoning over retrieved memories."
|
||||
_DEFAULT_FINAL_ROLE = "You are a thoughtful assistant that synthesizes answers from retrieved memories."
|
||||
|
||||
|
||||
def _extract_directive_rules(directives: list[dict[str, Any]]) -> list[str]:
|
||||
"""Extract directive rules as a list of strings."""
|
||||
@@ -133,7 +136,9 @@ def build_system_prompt_for_tools(
|
||||
|
||||
parts.extend(
|
||||
[
|
||||
"You are a reflection agent that answers questions by reasoning over retrieved memories.",
|
||||
mission.strip() if mission else _DEFAULT_ROLE,
|
||||
"",
|
||||
"Answer the user's question by reasoning over retrieved memories.",
|
||||
"",
|
||||
]
|
||||
)
|
||||
@@ -369,7 +374,7 @@ def build_agent_prompt(
|
||||
output = entry["output"]
|
||||
# Format as proper JSON for LLM readability
|
||||
try:
|
||||
output_str = json.dumps(output, indent=2, default=str)
|
||||
output_str = json.dumps(output, indent=2, default=str, ensure_ascii=False)
|
||||
except (TypeError, ValueError):
|
||||
output_str = str(output)
|
||||
parts.append(f"\n### Call {i}: {tool}\n```json\n{output_str}\n```")
|
||||
@@ -444,7 +449,7 @@ def build_final_prompt(
|
||||
tool = entry["tool"]
|
||||
output = entry["output"]
|
||||
try:
|
||||
output_str = json.dumps(output, indent=2, default=str)
|
||||
output_str = json.dumps(output, indent=2, default=str, ensure_ascii=False)
|
||||
except (TypeError, ValueError):
|
||||
output_str = str(output)
|
||||
block = f"\n### From {tool}:\n```json\n{output_str}\n```"
|
||||
@@ -479,9 +484,9 @@ def build_final_prompt(
|
||||
return "\n".join(parts)
|
||||
|
||||
|
||||
FINAL_SYSTEM_PROMPT = """CRITICAL: You MUST ONLY use information from retrieved tool results. NEVER make up names, people, events, or entities.
|
||||
_FINAL_SYSTEM_PROMPT_BASE = """CRITICAL: You MUST ONLY use information from retrieved tool results. NEVER make up names, people, events, or entities.
|
||||
|
||||
You are a thoughtful assistant that synthesizes answers from retrieved memories.
|
||||
{role_section}
|
||||
|
||||
Your approach:
|
||||
- Reason over the retrieved memories to answer the question
|
||||
@@ -508,3 +513,213 @@ CRITICAL: Output ONLY the final synthesized answer. Do NOT include:
|
||||
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) -> 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
|
||||
FINAL_SYSTEM_PROMPT = build_final_system_prompt()
|
||||
|
||||
|
||||
STRUCTURED_DELTA_SYSTEM_PROMPT = """You are integrating *new information* into an existing structured document.
|
||||
|
||||
You will be given:
|
||||
1. TOPIC — the question this document answers. Content that does not help
|
||||
answer this question is OFF-TOPIC and should be removed.
|
||||
2. CURRENT DOCUMENT (JSON) — the existing structured mental model. Each section
|
||||
has a stable ``id``, a ``heading``, a ``level`` (1..6), and an ordered list
|
||||
of ``blocks``. Blocks are typed: ``paragraph``, ``bullet_list``,
|
||||
``ordered_list``, or ``code``.
|
||||
3. NEW INFORMATION SYNTHESIS (markdown) — a synthesis showing how the new facts
|
||||
relate to the document's topic. Use it to understand context and relevance,
|
||||
but do NOT copy its formatting or wording wholesale.
|
||||
4. SUPPORTING FACTS — observations and facts created since the last refresh.
|
||||
These are genuinely new — they were NOT available when the current document
|
||||
was written.
|
||||
|
||||
Your task: output a JSON object ``{"operations": [...]}``. Applied to CURRENT
|
||||
DOCUMENT, the operations must produce a document that best answers the TOPIC
|
||||
by integrating the new facts.
|
||||
|
||||
RULES
|
||||
- These facts are NEW since the last refresh. The existing document already
|
||||
captures all prior information from earlier refreshes. Your job is to
|
||||
integrate the new facts into the existing document.
|
||||
- **Preserve existing content**: The current document was built from prior facts
|
||||
that you cannot see. Do NOT remove or replace existing sections just because
|
||||
the new facts do not reference them. Only remove content when the new facts
|
||||
explicitly contradict or supersede it.
|
||||
- **Merge overlapping topics**: When new facts cover topics that overlap with
|
||||
existing sections, merge the new information INTO the existing section
|
||||
rather than creating duplicates. When new facts provide more specific or
|
||||
authoritative guidance on a topic already covered generically, update the
|
||||
existing content to reflect the more specific guidance.
|
||||
- **Preserve examples**: Concrete examples, before/after pairs, sample sentences,
|
||||
and illustrative ✅/❌ comparisons are MORE valuable than abstract rules.
|
||||
When facts contain examples, include them. Never drop an example to make
|
||||
room for an abstract restatement of the same point.
|
||||
- Operations target sections by ``section_id`` (use the ``id`` field of the
|
||||
section in CURRENT DOCUMENT, NOT the heading). Block operations target
|
||||
blocks by ``index`` (0-based, against the section's current block list).
|
||||
- **Add** new content with ``append_block``, ``insert_block``, or ``add_section``
|
||||
when facts introduce information not yet covered. Prefer extending an
|
||||
existing section over creating a new one.
|
||||
- **Update** existing content with ``replace_block`` or ``replace_section_blocks``
|
||||
when new facts provide corrections, updates, or more specific information
|
||||
about topics already in the document.
|
||||
- **Remove** content with ``remove_block`` or ``remove_section`` ONLY when
|
||||
the new facts explicitly contradict or supersede it.
|
||||
- NEVER emit operations whose only effect is to reword unchanged content.
|
||||
- NEVER emit operations to "normalize" formatting (numbered → bulleted, casing
|
||||
changes, paragraph → list, etc).
|
||||
- Every operation MUST be justifiable by a specific fact in SUPPORTING FACTS.
|
||||
- Output ``{"operations": []}`` only if the new facts are already reflected
|
||||
in the document (e.g., from a concurrent update).
|
||||
|
||||
ALLOWED OPERATIONS (each line shows the JSON shape)
|
||||
- ``{"op": "append_block", "section_id": "...", "block": {...}}``
|
||||
- ``{"op": "insert_block", "section_id": "...", "index": N, "block": {...}}``
|
||||
- ``{"op": "replace_block", "section_id": "...", "index": N, "block": {...}}``
|
||||
- ``{"op": "remove_block", "section_id": "...", "index": N}``
|
||||
- ``{"op": "add_section", "heading": "...", "level": 2, "blocks": [...], "after_section_id": "..."}``
|
||||
- ``{"op": "remove_section", "section_id": "..."}``
|
||||
- ``{"op": "replace_section_blocks", "section_id": "...", "blocks": [...]}``
|
||||
- ``{"op": "rename_section", "section_id": "...", "new_heading": "..."}``
|
||||
|
||||
Block shapes
|
||||
- ``{"type": "paragraph", "text": "..."}``
|
||||
- ``{"type": "bullet_list", "items": ["...", "..."]}``
|
||||
- ``{"type": "ordered_list", "items": ["...", "..."]}``
|
||||
- ``{"type": "code", "language": "json", "text": "..."}``
|
||||
|
||||
OUTPUT FORMAT
|
||||
Return ONLY a single JSON object on its own, with no prose before or after,
|
||||
no markdown code fences, no commentary. The object must have exactly one
|
||||
top-level key, ``operations``, whose value is an array of operation objects
|
||||
(empty array when nothing changes).
|
||||
|
||||
Examples
|
||||
- No changes needed → ``{"operations": []}``
|
||||
- Add one bullet to an existing "Members" section →
|
||||
``{"operations": [{"op": "append_block", "section_id": "members",
|
||||
"block": {"type": "bullet_list", "items": ["Carol — junior engineer"]}}]}``
|
||||
- Replace a paragraph that has been corrected by new facts →
|
||||
``{"operations": [{"op": "replace_block", "section_id": "overview",
|
||||
"index": 0, "block": {"type": "paragraph", "text": "Updated summary."}}]}``
|
||||
- Remove an obsolete block →
|
||||
``{"operations": [{"op": "remove_block", "section_id": "status", "index": 2}]}``"""
|
||||
|
||||
|
||||
def build_structured_delta_prompt(
|
||||
*,
|
||||
current_document_json: str,
|
||||
candidate_markdown: str,
|
||||
supporting_facts: list[dict[str, Any]],
|
||||
source_query: str,
|
||||
max_output_tokens: int | None = None,
|
||||
) -> str:
|
||||
"""Build the user prompt for a structured-delta mental model refresh.
|
||||
|
||||
The LLM's job is to emit operations against ``current_document_json``;
|
||||
the surrounding ``candidate_markdown`` and ``supporting_facts`` are
|
||||
references for *what new information exists*, not templates to mimic.
|
||||
|
||||
``max_output_tokens`` is surfaced in the prompt so the model can keep its
|
||||
op list within the provider's response cap. The actual cap is enforced by
|
||||
the caller; this is just an advisory anchor — without it the model often
|
||||
returns op lists whose JSON gets truncated mid-string.
|
||||
"""
|
||||
fact_lines: list[str] = []
|
||||
for f in supporting_facts:
|
||||
fid = f.get("id", "")
|
||||
text = (f.get("text") or "").strip().replace("\n", " ")
|
||||
ftype = f.get("type", "")
|
||||
fact_lines.append(f"- [{ftype}:{fid}] {text}")
|
||||
facts_block = "\n".join(fact_lines) if fact_lines else "(no supporting facts retrieved)"
|
||||
|
||||
budget_hint = ""
|
||||
if max_output_tokens is not None:
|
||||
budget_hint = (
|
||||
f"\n\n## Output budget\n"
|
||||
f"Your JSON response must fit within ~{max_output_tokens} tokens. If you "
|
||||
"would need more than this to express every change, prefer the highest-"
|
||||
"leverage edits first (a few ``replace_section_blocks`` ops over many "
|
||||
"block-level ops) so the response always parses as valid JSON."
|
||||
)
|
||||
|
||||
return (
|
||||
f"## Topic\n{source_query}\n\n"
|
||||
f"## CURRENT DOCUMENT (apply ops to this; reference section ids as listed)\n"
|
||||
f"```json\n{current_document_json}\n```\n\n"
|
||||
f"## NEW INFORMATION SYNTHESIS (context for how new facts relate to the topic)\n"
|
||||
f"```markdown\n{candidate_markdown}\n```\n\n"
|
||||
f"## SUPPORTING FACTS (new since last refresh — integrate these)\n{facts_block}"
|
||||
f"{budget_hint}\n\n"
|
||||
"## Task\n"
|
||||
"Output a JSON object matching the operations schema. Integrate the new "
|
||||
"supporting facts into CURRENT DOCUMENT. Add, update, or remove content "
|
||||
"as needed. Preserve unchanged sections and blocks by not mentioning them."
|
||||
)
|
||||
|
||||
|
||||
DELTA_SYSTEM_PROMPT = """You are performing a surgical delta update to an existing mental model document.
|
||||
|
||||
You will be given:
|
||||
1. CURRENT DOCUMENT: the existing mental model content (markdown).
|
||||
2. CANDIDATE UPDATE: a freshly generated synthesis based on the latest retrieved memories.
|
||||
3. SUPPORTING FACTS: the observations and facts that support the CANDIDATE UPDATE.
|
||||
|
||||
Your task: produce an updated version of the CURRENT DOCUMENT that reflects the new reality, with the MINIMUM possible changes.
|
||||
|
||||
ABSOLUTE RULES:
|
||||
- Preserve unchanged content BYTE-FOR-BYTE. If a sentence, heading, bullet, code block, or section is still accurate according to the CANDIDATE UPDATE and SUPPORTING FACTS, copy it verbatim — same wording, same punctuation, same whitespace, same markdown structure.
|
||||
- Do NOT reformat, rephrase, or re-style content that is still accurate. No "light edits for clarity", no reordering for flow, no synonym swaps.
|
||||
- Remove content that is contradicted by the CANDIDATE UPDATE or SUPPORTING FACTS (stale content).
|
||||
- Add new content ONLY when the SUPPORTING FACTS contain information not already in the CURRENT DOCUMENT.
|
||||
- When adding new content, prefer appending to an existing relevant section. Creating a new section is acceptable when the new information does not fit any existing section.
|
||||
- When creating a new section, match the heading style, tone, and formatting conventions used in the CURRENT DOCUMENT.
|
||||
- Every assertion in your output MUST be grounded in either (a) the CURRENT DOCUMENT (preserved) or (b) the SUPPORTING FACTS. Never introduce outside knowledge.
|
||||
- If nothing in the SUPPORTING FACTS contradicts or extends the CURRENT DOCUMENT, return the CURRENT DOCUMENT UNCHANGED, character for character.
|
||||
|
||||
OUTPUT FORMAT:
|
||||
- Output ONLY the updated markdown document. No preamble, no explanation, no diff markers, no commentary.
|
||||
- Do not wrap the output in code fences unless the CURRENT DOCUMENT itself was entirely a code fence."""
|
||||
|
||||
|
||||
def build_delta_prompt(
|
||||
*,
|
||||
current_content: str,
|
||||
candidate_content: str,
|
||||
supporting_facts: list[dict[str, Any]],
|
||||
source_query: str,
|
||||
) -> str:
|
||||
"""Build the user prompt for a delta-mode mental model refresh.
|
||||
|
||||
Args:
|
||||
current_content: The existing mental model content (to preserve as much as possible).
|
||||
candidate_content: Fresh synthesis from the reflect agent reflecting new reality.
|
||||
supporting_facts: Flat list of fact dicts (id, text, type) supporting the candidate.
|
||||
source_query: The mental model's source query, for topical framing.
|
||||
"""
|
||||
fact_lines: list[str] = []
|
||||
for f in supporting_facts:
|
||||
fid = f.get("id", "")
|
||||
text = (f.get("text") or "").strip().replace("\n", " ")
|
||||
ftype = f.get("type", "")
|
||||
fact_lines.append(f"- [{ftype}:{fid}] {text}")
|
||||
facts_block = "\n".join(fact_lines) if fact_lines else "(no supporting facts retrieved)"
|
||||
|
||||
return (
|
||||
f"## Topic\n{source_query}\n\n"
|
||||
f"## CURRENT DOCUMENT\n```markdown\n{current_content}\n```\n\n"
|
||||
f"## CANDIDATE UPDATE\n```markdown\n{candidate_content}\n```\n\n"
|
||||
f"## SUPPORTING FACTS\n{facts_block}\n\n"
|
||||
"## Task\n"
|
||||
"Produce the updated mental model document by applying the minimum necessary changes "
|
||||
"to CURRENT DOCUMENT so that it reflects CANDIDATE UPDATE and SUPPORTING FACTS. "
|
||||
"Preserve unchanged content byte-for-byte. Output only the final markdown."
|
||||
)
|
||||
|
||||
@@ -0,0 +1,301 @@
|
||||
"""Structured representation of a mental model document.
|
||||
|
||||
Why this exists
|
||||
---------------
|
||||
Storing mental models as raw markdown forces every refresh to round-trip prose
|
||||
through an LLM, which then drifts on stylistic details (numbered vs bulleted
|
||||
lists, casing, separator lines, paraphrasing) even when instructed to preserve
|
||||
content byte-for-byte. The intrinsic mechanism of an LLM is to *generate* the
|
||||
next token from a gestalt of the input — not to copy tokens verbatim — so any
|
||||
"preserve unchanged content" instruction is fundamentally a soft constraint.
|
||||
|
||||
The fix is to give the LLM no opportunity to drift on unchanged content. We
|
||||
keep an authoritative structured representation of the document; the markdown
|
||||
shown to users is a deterministic render of that structure. Delta refreshes
|
||||
emit *operations* against the structure (see ``delta_ops.py``); sections and
|
||||
blocks not mentioned by any operation are physically untouched.
|
||||
|
||||
Schema (v1)
|
||||
-----------
|
||||
A document is an ordered list of ``Section``s. Each section has:
|
||||
- ``id`` : stable slug derived from ``heading`` (used as the operation
|
||||
target across refreshes; surviving renames is a separate
|
||||
concern handled by an explicit ``rename`` op).
|
||||
- ``heading``: the markdown heading text (without the ``#`` prefix).
|
||||
- ``level`` : 1 (``#``) … 6 (``######``). Default 2.
|
||||
- ``blocks``: ordered list of typed blocks — paragraph, bullet_list,
|
||||
ordered_list, code.
|
||||
|
||||
The schema is intentionally narrow: it covers what real mental-model documents
|
||||
actually contain (the kind a coding agent writes for itself or a user writes as
|
||||
a "skill" doc). Tables, images, and raw HTML are out of scope until needed.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from typing import Annotated, Literal, Union
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
|
||||
# Blocks ---------------------------------------------------------------------
|
||||
|
||||
|
||||
class ParagraphBlock(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
type: Literal["paragraph"] = "paragraph"
|
||||
text: str
|
||||
|
||||
|
||||
class BulletListBlock(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
type: Literal["bullet_list"] = "bullet_list"
|
||||
items: list[str] = Field(default_factory=list)
|
||||
|
||||
|
||||
class OrderedListBlock(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
type: Literal["ordered_list"] = "ordered_list"
|
||||
items: list[str] = Field(default_factory=list)
|
||||
|
||||
|
||||
class CodeBlock(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
type: Literal["code"] = "code"
|
||||
language: str = ""
|
||||
text: str
|
||||
|
||||
|
||||
Block = Annotated[
|
||||
Union[ParagraphBlock, BulletListBlock, OrderedListBlock, CodeBlock],
|
||||
Field(discriminator="type"),
|
||||
]
|
||||
|
||||
|
||||
# Section / Document ---------------------------------------------------------
|
||||
|
||||
|
||||
class Section(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
id: str
|
||||
heading: str
|
||||
level: int = Field(default=2, ge=1, le=6)
|
||||
blocks: list[Block] = Field(default_factory=list)
|
||||
|
||||
|
||||
class StructuredDocument(BaseModel):
|
||||
"""Top-level structured representation of a mental model."""
|
||||
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
version: Literal[1] = 1
|
||||
sections: list[Section] = Field(default_factory=list)
|
||||
|
||||
def section_by_id(self, section_id: str) -> Section | None:
|
||||
for s in self.sections:
|
||||
if s.id == section_id:
|
||||
return s
|
||||
return None
|
||||
|
||||
def section_index(self, section_id: str) -> int | None:
|
||||
for i, s in enumerate(self.sections):
|
||||
if s.id == section_id:
|
||||
return i
|
||||
return None
|
||||
|
||||
|
||||
# Slug helpers ---------------------------------------------------------------
|
||||
|
||||
_SLUG_RX = re.compile(r"[^a-z0-9]+")
|
||||
|
||||
|
||||
def slugify_heading(heading: str) -> str:
|
||||
"""Stable, deterministic slug from a heading.
|
||||
|
||||
"Stop Conditions" -> "stop-conditions"
|
||||
"Inputs and Context" -> "inputs-and-context"
|
||||
"""
|
||||
slug = _SLUG_RX.sub("-", heading.strip().lower()).strip("-")
|
||||
return slug or "section"
|
||||
|
||||
|
||||
def make_unique_id(base: str, existing: set[str]) -> str:
|
||||
"""Disambiguate by appending -2, -3, … if the slug is already in use."""
|
||||
if base not in existing:
|
||||
return base
|
||||
i = 2
|
||||
while f"{base}-{i}" in existing:
|
||||
i += 1
|
||||
return f"{base}-{i}"
|
||||
|
||||
|
||||
# Renderer -------------------------------------------------------------------
|
||||
|
||||
|
||||
def render_block(block: Block) -> str:
|
||||
"""Render a single block to markdown. No trailing newline."""
|
||||
if isinstance(block, ParagraphBlock):
|
||||
return block.text.rstrip()
|
||||
if isinstance(block, BulletListBlock):
|
||||
return "\n".join(f"- {item.rstrip()}" for item in block.items)
|
||||
if isinstance(block, OrderedListBlock):
|
||||
return "\n".join(f"{i + 1}. {item.rstrip()}" for i, item in enumerate(block.items))
|
||||
if isinstance(block, CodeBlock):
|
||||
fence_lang = block.language or ""
|
||||
return f"```{fence_lang}\n{block.text}\n```"
|
||||
raise TypeError(f"Unknown block type: {type(block)!r}")
|
||||
|
||||
|
||||
def render_section(section: Section) -> str:
|
||||
"""Render a section: heading + blank line + blocks separated by blank lines."""
|
||||
parts = ["#" * section.level + " " + section.heading.strip()]
|
||||
for block in section.blocks:
|
||||
parts.append("") # blank line before each block
|
||||
parts.append(render_block(block))
|
||||
return "\n".join(parts)
|
||||
|
||||
|
||||
def render_document(doc: StructuredDocument) -> str:
|
||||
"""Render the whole document. Sections separated by a single blank line.
|
||||
|
||||
The output is byte-stable: same structured input always produces the same
|
||||
markdown, modulo the inherent ordering of sections/blocks/items.
|
||||
"""
|
||||
if not doc.sections:
|
||||
return ""
|
||||
return "\n\n".join(render_section(s) for s in doc.sections) + "\n"
|
||||
|
||||
|
||||
# Parser ---------------------------------------------------------------------
|
||||
#
|
||||
# The parser is intentionally lenient: it accepts the markdown produced by
|
||||
# our own renderer (round-trip-safe) and the markdown an LLM tends to produce
|
||||
# for mental-model documents. It is *not* a general CommonMark parser — it
|
||||
# does not need to be. When it cannot classify a block it falls back to a
|
||||
# paragraph so that no content is silently dropped.
|
||||
|
||||
_HEADING_RX = re.compile(r"^(#{1,6})\s+(.+?)\s*$")
|
||||
_BULLET_RX = re.compile(r"^\s*[-*+]\s+(.*)$")
|
||||
_ORDERED_RX = re.compile(r"^\s*\d+[.)]\s+(.*)$")
|
||||
_FENCE_RX = re.compile(r"^```([A-Za-z0-9_+-]*)\s*$")
|
||||
|
||||
|
||||
def _strip_separators(lines: list[str]) -> list[str]:
|
||||
"""Drop horizontal-rule lines (`---`, `***`) used as section separators.
|
||||
|
||||
Our renderer never emits these, but LLM output frequently includes them
|
||||
between sections; treating them as blank lines avoids parsing them as
|
||||
paragraphs.
|
||||
"""
|
||||
return ["" if re.fullmatch(r"\s*([-*_])\1{2,}\s*", line) else line for line in lines]
|
||||
|
||||
|
||||
def _split_blocks(lines: list[str]) -> list[list[str]]:
|
||||
"""Group consecutive non-blank lines into block chunks."""
|
||||
chunks: list[list[str]] = []
|
||||
current: list[str] = []
|
||||
in_fence = False
|
||||
for line in lines:
|
||||
if _FENCE_RX.match(line):
|
||||
current.append(line)
|
||||
in_fence = not in_fence
|
||||
continue
|
||||
if in_fence:
|
||||
current.append(line)
|
||||
continue
|
||||
if line.strip() == "":
|
||||
if current:
|
||||
chunks.append(current)
|
||||
current = []
|
||||
else:
|
||||
current.append(line)
|
||||
if current:
|
||||
chunks.append(current)
|
||||
return chunks
|
||||
|
||||
|
||||
def _parse_block(chunk: list[str]) -> Block:
|
||||
"""Parse a single non-empty chunk into a block."""
|
||||
if chunk and _FENCE_RX.match(chunk[0]):
|
||||
m = _FENCE_RX.match(chunk[0])
|
||||
lang = m.group(1) if m else ""
|
||||
body_lines = chunk[1:]
|
||||
if body_lines and _FENCE_RX.match(body_lines[-1]):
|
||||
body_lines = body_lines[:-1]
|
||||
return CodeBlock(language=lang, text="\n".join(body_lines))
|
||||
|
||||
if all(_BULLET_RX.match(line) for line in chunk):
|
||||
items = []
|
||||
for line in chunk:
|
||||
m = _BULLET_RX.match(line)
|
||||
assert m is not None
|
||||
items.append(m.group(1).strip())
|
||||
return BulletListBlock(items=items)
|
||||
|
||||
if all(_ORDERED_RX.match(line) for line in chunk):
|
||||
items = []
|
||||
for line in chunk:
|
||||
m = _ORDERED_RX.match(line)
|
||||
assert m is not None
|
||||
items.append(m.group(1).strip())
|
||||
return OrderedListBlock(items=items)
|
||||
|
||||
return ParagraphBlock(text=" ".join(line.strip() for line in chunk).strip())
|
||||
|
||||
|
||||
def parse_markdown(markdown: str) -> StructuredDocument:
|
||||
"""Best-effort parse of a markdown document into the structured schema.
|
||||
|
||||
Sections are introduced by ATX headings (``#``..``######``). Anything
|
||||
before the first heading is wrapped into an implicit "Overview" section
|
||||
so we never silently drop user content. Section IDs are unique slugs of
|
||||
their headings.
|
||||
"""
|
||||
raw_lines = (markdown or "").splitlines()
|
||||
lines = _strip_separators(raw_lines)
|
||||
|
||||
sections: list[Section] = []
|
||||
used_ids: set[str] = set()
|
||||
pending: list[str] = []
|
||||
current: Section | None = None
|
||||
|
||||
def flush_pending_into(section: Section) -> None:
|
||||
if not pending:
|
||||
return
|
||||
for chunk in _split_blocks(pending):
|
||||
section.blocks.append(_parse_block(chunk))
|
||||
pending.clear()
|
||||
|
||||
for line in lines:
|
||||
m = _HEADING_RX.match(line)
|
||||
if m:
|
||||
if current is not None:
|
||||
flush_pending_into(current)
|
||||
sections.append(current)
|
||||
elif pending:
|
||||
# Content before the first heading: wrap in implicit section.
|
||||
base = "overview"
|
||||
section_id = make_unique_id(base, used_ids)
|
||||
used_ids.add(section_id)
|
||||
implicit = Section(id=section_id, heading="Overview", level=2)
|
||||
flush_pending_into(implicit)
|
||||
sections.append(implicit)
|
||||
level = len(m.group(1))
|
||||
heading = m.group(2).strip()
|
||||
section_id = make_unique_id(slugify_heading(heading), used_ids)
|
||||
used_ids.add(section_id)
|
||||
current = Section(id=section_id, heading=heading, level=level)
|
||||
else:
|
||||
pending.append(line)
|
||||
|
||||
if current is not None:
|
||||
flush_pending_into(current)
|
||||
sections.append(current)
|
||||
elif pending:
|
||||
base = "overview"
|
||||
section_id = make_unique_id(base, used_ids)
|
||||
used_ids.add(section_id)
|
||||
implicit = Section(id=section_id, heading="Overview", level=2)
|
||||
flush_pending_into(implicit)
|
||||
sections.append(implicit)
|
||||
|
||||
return StructuredDocument(sections=sections)
|
||||
@@ -9,6 +9,7 @@ Implements hierarchical retrieval:
|
||||
|
||||
import logging
|
||||
import uuid
|
||||
from dataclasses import replace
|
||||
from datetime import datetime, timezone
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
@@ -22,6 +23,7 @@ logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
async def tool_search_mental_models(
|
||||
memory_engine: "MemoryEngine",
|
||||
conn: "Connection",
|
||||
bank_id: str,
|
||||
query: str,
|
||||
@@ -31,7 +33,6 @@ async def tool_search_mental_models(
|
||||
tags_match: str = "any",
|
||||
tag_groups: "list | None" = None,
|
||||
exclude_ids: list[str] | None = None,
|
||||
pending_consolidation: int = 0,
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Search user-curated mental models by semantic similarity.
|
||||
@@ -81,7 +82,7 @@ async def tool_search_mental_models(
|
||||
f"""
|
||||
SELECT
|
||||
id, name, content,
|
||||
tags, created_at, last_refreshed_at,
|
||||
tags, created_at, last_refreshed_at, trigger,
|
||||
1 - (embedding <=> $2::vector) as relevance
|
||||
FROM {fq_table("mental_models")}
|
||||
WHERE bank_id = $1 AND embedding IS NOT NULL {filters}
|
||||
@@ -98,10 +99,9 @@ async def tool_search_mental_models(
|
||||
if last_refreshed_at and last_refreshed_at.tzinfo is None:
|
||||
last_refreshed_at = last_refreshed_at.replace(tzinfo=timezone.utc)
|
||||
|
||||
# A mental model is stale when there are memories that haven't been consolidated yet —
|
||||
# the same signal used for observations staleness.
|
||||
is_stale = pending_consolidation > 0
|
||||
staleness_reason = f"{pending_consolidation} memories pending consolidation" if is_stale else None
|
||||
# Per-MM staleness: new in-scope memories since last refresh (includes pending).
|
||||
is_stale = await memory_engine.compute_mental_model_is_stale(conn, bank_id, row)
|
||||
staleness_reason = "new in-scope memories ingested since last refresh" if is_stale else None
|
||||
|
||||
mental_models.append(
|
||||
{
|
||||
@@ -135,6 +135,8 @@ async def tool_search_observations(
|
||||
last_consolidated_at: datetime | None = None,
|
||||
pending_consolidation: int = 0,
|
||||
source_facts_max_tokens: int = -1,
|
||||
created_after: datetime | None = None,
|
||||
created_before: datetime | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Search consolidated observations using recall.
|
||||
@@ -162,17 +164,24 @@ async def tool_search_observations(
|
||||
if include_source_facts and source_facts_max_tokens > 0:
|
||||
recall_kwargs["max_source_facts_tokens"] = source_facts_max_tokens
|
||||
|
||||
# Use an internal request context so this recall is not billed as a
|
||||
# user-facing operation. The reflect caller is already billed for the
|
||||
# overall reflect operation; double-billing the sub-recalls would
|
||||
# overcharge the customer.
|
||||
internal_ctx = replace(request_context, internal=True)
|
||||
result = await memory_engine.recall_async(
|
||||
bank_id=bank_id,
|
||||
query=query,
|
||||
fact_type=["observation"],
|
||||
max_tokens=max_tokens,
|
||||
enable_trace=False,
|
||||
request_context=request_context,
|
||||
request_context=internal_ctx,
|
||||
tags=tags,
|
||||
tags_match=tags_match,
|
||||
tag_groups=tag_groups,
|
||||
include_source_facts=include_source_facts,
|
||||
created_after=created_after,
|
||||
created_before=created_before,
|
||||
_connection_budget=1,
|
||||
_quiet=True,
|
||||
**recall_kwargs,
|
||||
@@ -208,6 +217,9 @@ async def tool_recall(
|
||||
connection_budget: int = 1,
|
||||
max_chunk_tokens: int = 1000,
|
||||
fact_types: list[str] | None = None,
|
||||
include_chunks: bool = True,
|
||||
created_after: datetime | None = None,
|
||||
created_before: datetime | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Search memories using TEMPR retrieval.
|
||||
@@ -224,25 +236,28 @@ async def tool_recall(
|
||||
tags: Filter by tags (includes untagged memories)
|
||||
tags_match: How to match tags - "any" (OR), "all" (AND), or "exact"
|
||||
connection_budget: Max DB connections for this recall (default 1 for internal ops)
|
||||
max_chunk_tokens: Maximum tokens for raw source chunk text (default 1000, always included)
|
||||
max_chunk_tokens: Maximum tokens for raw source chunk text (default 1000)
|
||||
fact_types: Optional filter for fact types to retrieve. Defaults to ["experience", "world"].
|
||||
include_chunks: Whether to fetch raw chunk text alongside facts (default True).
|
||||
|
||||
Returns:
|
||||
Dict with list of matching memories including raw chunk text
|
||||
Dict with list of matching memories including raw chunk text (when include_chunks)
|
||||
"""
|
||||
# Only world/experience are valid for raw recall (observation is handled by search_observations)
|
||||
recall_fact_type = [ft for ft in (fact_types or ["experience", "world"]) if ft in ("world", "experience")]
|
||||
include_chunks = True
|
||||
internal_ctx = replace(request_context, internal=True)
|
||||
result = await memory_engine.recall_async(
|
||||
bank_id=bank_id,
|
||||
query=query,
|
||||
fact_type=recall_fact_type,
|
||||
max_tokens=max_tokens,
|
||||
enable_trace=False,
|
||||
request_context=request_context,
|
||||
request_context=internal_ctx,
|
||||
tags=tags,
|
||||
tags_match=tags_match,
|
||||
tag_groups=tag_groups,
|
||||
created_after=created_after,
|
||||
created_before=created_before,
|
||||
_connection_budget=connection_budget,
|
||||
_quiet=True, # Suppress logging for internal operations
|
||||
include_chunks=include_chunks,
|
||||
|
||||
@@ -113,6 +113,22 @@ async def get_bank_profile(pool, bank_id: str) -> BankProfile:
|
||||
Returns:
|
||||
BankProfile with name, typed DispositionTraits, and mission
|
||||
"""
|
||||
profile, _ = await get_or_create_bank_profile(pool, bank_id)
|
||||
return profile
|
||||
|
||||
|
||||
async def get_or_create_bank_profile(pool, bank_id: str) -> tuple[BankProfile, bool]:
|
||||
"""
|
||||
Get bank profile, auto-creating with defaults if it doesn't exist.
|
||||
|
||||
Same as get_bank_profile, but also returns a flag indicating whether the
|
||||
bank was freshly created on this call. Used by the memory engine to apply
|
||||
the HINDSIGHT_API_DEFAULT_BANK_TEMPLATE hook on first bank creation.
|
||||
|
||||
Returns:
|
||||
Tuple of (BankProfile, created) where created is True if the bank
|
||||
did not exist before this call.
|
||||
"""
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
# Try to get existing bank
|
||||
row = await conn.fetchrow(
|
||||
@@ -129,10 +145,13 @@ async def get_bank_profile(pool, bank_id: str) -> BankProfile:
|
||||
if isinstance(disposition_data, str):
|
||||
disposition_data = json.loads(disposition_data)
|
||||
|
||||
return BankProfile(
|
||||
name=row["name"],
|
||||
disposition=DispositionTraits(**disposition_data),
|
||||
mission=row["mission"] or "",
|
||||
return (
|
||||
BankProfile(
|
||||
name=row["name"],
|
||||
disposition=DispositionTraits(**disposition_data),
|
||||
mission=row["mission"] or "",
|
||||
),
|
||||
False,
|
||||
)
|
||||
|
||||
# Bank doesn't exist, create with defaults.
|
||||
@@ -153,11 +172,15 @@ async def get_bank_profile(pool, bank_id: str) -> BankProfile:
|
||||
internal_id,
|
||||
)
|
||||
|
||||
if inserted:
|
||||
created = inserted is not None
|
||||
if created:
|
||||
# Fresh insert — create per-bank vector indexes (instant on empty bank)
|
||||
await create_bank_vector_indexes(conn, bank_id, str(internal_id))
|
||||
|
||||
return BankProfile(name=bank_id, disposition=DispositionTraits(**DEFAULT_DISPOSITION), mission="")
|
||||
return (
|
||||
BankProfile(name=bank_id, disposition=DispositionTraits(**DEFAULT_DISPOSITION), mission=""),
|
||||
created,
|
||||
)
|
||||
|
||||
|
||||
async def update_bank_disposition(pool, bank_id: str, disposition: dict[str, int]) -> None:
|
||||
|
||||
@@ -100,11 +100,20 @@ async def store_chunks_batch(conn, bank_id: str, document_id: str, chunks: list[
|
||||
content_hashes.append(compute_chunk_hash(chunk.chunk_text))
|
||||
chunk_id_map[chunk.chunk_index] = chunk_id
|
||||
|
||||
# Batch insert all chunks
|
||||
# Batch upsert all chunks. ON CONFLICT makes this idempotent: re-submitting
|
||||
# a retain under the same document_id (the pattern in vectorize-io/hindsight#977)
|
||||
# may produce chunk_ids that already exist when upstream cascade-delete or
|
||||
# delta-retain paths don't run (or race with a concurrent task). Overwriting
|
||||
# is the correct behavior per the document_id grouping semantics — the caller
|
||||
# intends this chunk to hold the latest content at that (document_id, index).
|
||||
await conn.execute(
|
||||
f"""
|
||||
INSERT INTO {fq_table("chunks")} (chunk_id, document_id, bank_id, chunk_text, chunk_index, content_hash)
|
||||
SELECT * FROM unnest($1::text[], $2::text[], $3::text[], $4::text[], $5::integer[], $6::text[])
|
||||
ON CONFLICT (chunk_id) DO UPDATE SET
|
||||
chunk_text = EXCLUDED.chunk_text,
|
||||
chunk_index = EXCLUDED.chunk_index,
|
||||
content_hash = EXCLUDED.content_hash
|
||||
""",
|
||||
chunk_ids,
|
||||
[document_id] * len(chunk_texts),
|
||||
|
||||
@@ -47,6 +47,16 @@ async def generate_embeddings_batch(embeddings_backend, texts: list[str]) -> lis
|
||||
embeddings_backend.encode,
|
||||
texts,
|
||||
)
|
||||
return embeddings
|
||||
except Exception as e:
|
||||
raise Exception(f"Failed to generate batch embeddings: {str(e)}")
|
||||
|
||||
# Guarantee 1:1 alignment with input texts. A silent length mismatch here
|
||||
# propagates downstream as zip() drops items, eventually surfacing as an
|
||||
# IndexError in retain mapping (see issue #1037).
|
||||
if len(embeddings) != len(texts):
|
||||
raise RuntimeError(
|
||||
f"Embeddings backend returned {len(embeddings)} vectors for {len(texts)} input texts; "
|
||||
"expected exact 1:1 alignment"
|
||||
)
|
||||
|
||||
return embeddings
|
||||
|
||||
@@ -909,6 +909,7 @@ def _build_user_message(
|
||||
event_date: datetime | None,
|
||||
context: str,
|
||||
metadata: dict[str, str] | None = None,
|
||||
agent_name: str | None = None,
|
||||
) -> str:
|
||||
"""Build user message for fact extraction."""
|
||||
from .orchestrator import parse_datetime_flexible
|
||||
@@ -927,11 +928,15 @@ def _build_user_message(
|
||||
metadata_lines = "\n".join(f" {k}: {v}" for k, v in metadata.items())
|
||||
metadata_section = f"\nMetadata:\n{metadata_lines}"
|
||||
|
||||
narrator_section = ""
|
||||
if agent_name:
|
||||
narrator_section = f'\nNarrator: {agent_name} (AI agent — first-person statements like "I did X" are the agent\'s own actions; classify as "assistant")'
|
||||
|
||||
return f"""Extract facts from the following text chunk.
|
||||
|
||||
Chunk: {chunk_index + 1}/{total_chunks}
|
||||
Event Date: {event_date_str}
|
||||
Context: {sanitized_context}{metadata_section}
|
||||
Context: {sanitized_context}{metadata_section}{narrator_section}
|
||||
|
||||
Text:
|
||||
{sanitized_chunk}"""
|
||||
@@ -995,7 +1000,7 @@ async def _extract_facts_from_chunk(
|
||||
extract_causal_links = config.retain_extract_causal_links
|
||||
|
||||
# Build user message using helper function
|
||||
user_message = _build_user_message(chunk, chunk_index, total_chunks, event_date, context, metadata)
|
||||
user_message = _build_user_message(chunk, chunk_index, total_chunks, event_date, context, metadata, agent_name)
|
||||
|
||||
# Retry logic for JSON validation errors
|
||||
# Use retain-specific overrides if set, otherwise fall back to global LLM config
|
||||
@@ -1055,7 +1060,7 @@ async def _extract_facts_from_chunk(
|
||||
f"LLM response missing 'facts' field or returned empty list. "
|
||||
f"Response: {extraction_response_json}. "
|
||||
f"Input: "
|
||||
f"date: {event_date.isoformat()}, "
|
||||
f"date: {event_date.isoformat() if event_date else 'unset'}, "
|
||||
f"context: {context if context else 'none'}, "
|
||||
f"text: {chunk}"
|
||||
)
|
||||
@@ -1632,7 +1637,13 @@ async def extract_facts_from_contents_batch_api(
|
||||
|
||||
# Build user message using helper function
|
||||
user_message = _build_user_message(
|
||||
chunk, chunk_index_in_content, len(chunks), item.event_date, item.context, item.metadata or None
|
||||
chunk,
|
||||
chunk_index_in_content,
|
||||
len(chunks),
|
||||
item.event_date,
|
||||
item.context,
|
||||
item.metadata or None,
|
||||
agent_name,
|
||||
)
|
||||
|
||||
# Build request body using helper function
|
||||
|
||||
@@ -7,6 +7,7 @@ Handles insertion of facts into the database.
|
||||
import json
|
||||
import logging
|
||||
import uuid
|
||||
from datetime import datetime
|
||||
|
||||
from ...config import get_config
|
||||
from ..memory_engine import fq_table
|
||||
@@ -17,6 +18,23 @@ from .types import ProcessedFact
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
async def get_document_content(
|
||||
conn,
|
||||
bank_id: str,
|
||||
document_id: str,
|
||||
) -> str | None:
|
||||
"""Fetch the original_text of an existing document.
|
||||
|
||||
Returns None if the document does not exist.
|
||||
"""
|
||||
row = await conn.fetchval(
|
||||
f"SELECT original_text FROM {fq_table('documents')} WHERE id = $1 AND bank_id = $2",
|
||||
document_id,
|
||||
bank_id,
|
||||
)
|
||||
return row
|
||||
|
||||
|
||||
async def insert_facts_batch(
|
||||
conn, bank_id: str, facts: list[ProcessedFact], document_id: str | None = None
|
||||
) -> list[str]:
|
||||
@@ -207,6 +225,85 @@ async def ensure_bank_exists(conn, bank_id: str) -> None:
|
||||
await create_bank_vector_indexes(conn, bank_id, str(internal_id))
|
||||
|
||||
|
||||
async def delete_stale_observations_for_memories(
|
||||
conn,
|
||||
bank_id: str,
|
||||
fact_ids: "list[str | uuid.UUID]",
|
||||
) -> int:
|
||||
"""Delete observations whose source memories are about to be removed.
|
||||
|
||||
Mirrors the cleanup performed by ``MemoryEngine.delete_document`` so that
|
||||
every code path that removes ``memory_units`` also removes the
|
||||
observations derived from them. Without this, ingesting a fresh version
|
||||
of a document via the retain pipeline (which does a full-replace
|
||||
``DELETE FROM documents`` cascade) used to leave orphan observations
|
||||
pointing at memory IDs that no longer existed.
|
||||
|
||||
For each observation referencing any of ``fact_ids``:
|
||||
1. Delete the observation row (its text is stale once even one source
|
||||
memory disappears).
|
||||
2. Reset ``consolidated_at = NULL`` on the surviving source memories so
|
||||
they get re-consolidated under fresh observations on the next run.
|
||||
|
||||
Must be called within an active transaction, before the source memories
|
||||
are deleted.
|
||||
|
||||
Returns the number of observations deleted.
|
||||
"""
|
||||
if not fact_ids:
|
||||
return 0
|
||||
|
||||
fact_uuids = [uuid.UUID(str(fid)) if not isinstance(fid, uuid.UUID) else fid for fid in fact_ids]
|
||||
|
||||
affected_obs = await conn.fetch(
|
||||
f"""
|
||||
SELECT id, source_memory_ids
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE bank_id = $1
|
||||
AND fact_type = 'observation'
|
||||
AND source_memory_ids && $2::uuid[]
|
||||
""",
|
||||
bank_id,
|
||||
fact_uuids,
|
||||
)
|
||||
|
||||
if not affected_obs:
|
||||
return 0
|
||||
|
||||
deleted_set = {str(uid) for uid in fact_uuids}
|
||||
obs_ids = [obs["id"] for obs in affected_obs]
|
||||
seen_remaining: set[str] = set()
|
||||
remaining_source_ids: list[uuid.UUID] = []
|
||||
for obs in affected_obs:
|
||||
for src_id in obs["source_memory_ids"] or []:
|
||||
src_str = str(src_id)
|
||||
if src_str not in deleted_set and src_str not in seen_remaining:
|
||||
remaining_source_ids.append(src_id)
|
||||
seen_remaining.add(src_str)
|
||||
|
||||
await conn.execute(
|
||||
f"DELETE FROM {fq_table('memory_units')} WHERE id = ANY($1::uuid[])",
|
||||
obs_ids,
|
||||
)
|
||||
|
||||
if remaining_source_ids:
|
||||
await conn.execute(
|
||||
f"""
|
||||
UPDATE {fq_table("memory_units")}
|
||||
SET consolidated_at = NULL
|
||||
WHERE id = ANY($1::uuid[])
|
||||
AND fact_type IN ('experience', 'world')
|
||||
""",
|
||||
remaining_source_ids,
|
||||
)
|
||||
|
||||
logger.info(
|
||||
f"[OBSERVATIONS] Deleted {len(obs_ids)} observations, reset {len(remaining_source_ids)} "
|
||||
f"source memories for re-consolidation in bank {bank_id}"
|
||||
)
|
||||
return len(obs_ids)
|
||||
|
||||
|
||||
async def handle_document_tracking(
|
||||
conn,
|
||||
bank_id: str,
|
||||
@@ -237,17 +334,58 @@ async def handle_document_tracking(
|
||||
combined_content = _sanitize_text(combined_content) or ""
|
||||
content_hash = hashlib.sha256(combined_content.encode()).hexdigest()
|
||||
|
||||
# Delete old document first (cascades to units and links)
|
||||
# Only delete on the first batch to avoid deleting data we just inserted
|
||||
# Delete old document first (cascades to units and links).
|
||||
# Only delete on the first batch to avoid deleting data we just inserted.
|
||||
# Before the cascade, fan out to delete observations derived from the
|
||||
# outgoing memory_units — otherwise the FK ON DELETE CASCADE removes the
|
||||
# source memory_units but leaves observation rows pointing at IDs that
|
||||
# no longer exist (consolidated_at on co-source memories also stays
|
||||
# frozen). Same cleanup the explicit ``delete_document`` API performs.
|
||||
preserved_created_at = None
|
||||
if is_first_batch:
|
||||
await conn.fetchval(
|
||||
f"DELETE FROM {fq_table('documents')} WHERE id = $1 AND bank_id = $2 RETURNING id",
|
||||
existing_unit_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT id FROM {fq_table("memory_units")}
|
||||
WHERE document_id = $1 AND fact_type IN ('experience', 'world')
|
||||
""",
|
||||
document_id,
|
||||
)
|
||||
existing_unit_ids = [row["id"] for row in existing_unit_rows]
|
||||
if existing_unit_ids:
|
||||
invalidated = await delete_stale_observations_for_memories(conn, bank_id, existing_unit_ids)
|
||||
if invalidated:
|
||||
logger.info(
|
||||
f"[RETAIN] Document {document_id} re-ingested: invalidated "
|
||||
f"{invalidated} observation(s) derived from {len(existing_unit_ids)} outgoing memory_units"
|
||||
)
|
||||
# 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
|
||||
# (e.g. from partial writes or edge cases) would survive the cascade.
|
||||
# This explicit delete ensures complete cleanup.
|
||||
await conn.execute(
|
||||
f"DELETE FROM {fq_table('memory_units')} WHERE document_id = $1 AND bank_id = $2",
|
||||
document_id,
|
||||
bank_id,
|
||||
)
|
||||
# Capture created_at before deletion so re-ingestion preserves it.
|
||||
preserved_created_at = await conn.fetchval(
|
||||
f"DELETE FROM {fq_table('documents')} WHERE id = $1 AND bank_id = $2 RETURNING created_at",
|
||||
document_id,
|
||||
bank_id,
|
||||
)
|
||||
|
||||
# Insert document (or update if exists from concurrent operations)
|
||||
await _upsert_document_row(conn, bank_id, document_id, combined_content, content_hash, retain_params, document_tags)
|
||||
await _upsert_document_row(
|
||||
conn,
|
||||
bank_id,
|
||||
document_id,
|
||||
combined_content,
|
||||
content_hash,
|
||||
retain_params,
|
||||
document_tags,
|
||||
preserved_created_at=preserved_created_at,
|
||||
)
|
||||
|
||||
|
||||
async def upsert_document_metadata(
|
||||
@@ -280,12 +418,19 @@ async def _upsert_document_row(
|
||||
content_hash: str,
|
||||
retain_params: dict | None = None,
|
||||
document_tags: list[str] | None = None,
|
||||
preserved_created_at: datetime | None = None,
|
||||
) -> None:
|
||||
"""Insert or update a document row."""
|
||||
"""Insert or update a document row.
|
||||
|
||||
When ``preserved_created_at`` is provided, it is used for ``created_at`` on
|
||||
INSERT so that re-ingesting a document (which deletes + inserts the row)
|
||||
keeps the original creation timestamp. ``updated_at`` is always set to
|
||||
``NOW()`` on both INSERT and the ON CONFLICT UPDATE branch.
|
||||
"""
|
||||
await conn.execute(
|
||||
f"""
|
||||
INSERT INTO {fq_table("documents")} (id, bank_id, original_text, content_hash, retain_params, tags)
|
||||
VALUES ($1, $2, $3, $4, $5, $6)
|
||||
INSERT INTO {fq_table("documents")} (id, bank_id, original_text, content_hash, retain_params, tags, created_at, updated_at)
|
||||
VALUES ($1, $2, $3, $4, $5, $6, COALESCE($7, NOW()), NOW())
|
||||
ON CONFLICT (id, bank_id) DO UPDATE
|
||||
SET original_text = EXCLUDED.original_text,
|
||||
content_hash = EXCLUDED.content_hash,
|
||||
@@ -299,6 +444,7 @@ async def _upsert_document_row(
|
||||
content_hash,
|
||||
json.dumps(retain_params) if retain_params else None,
|
||||
document_tags or [],
|
||||
preserved_created_at,
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -739,17 +739,10 @@ async def compute_semantic_links_ann(
|
||||
return []
|
||||
|
||||
import time as time_mod
|
||||
import uuid as uuid_mod
|
||||
|
||||
ann_start = time_mod.time()
|
||||
links = []
|
||||
|
||||
# Lower ef_search for retain ANN — default 400 is tuned for recall precision
|
||||
# but at 164k units each HNSW probe takes 94ms. ef_search=60 gives 2.7ms/probe
|
||||
# (35x faster) with sufficient accuracy for top-50 semantic link creation.
|
||||
# Reset after to avoid polluting the connection pool for recall queries.
|
||||
await conn.execute("SET hnsw.ef_search = 60")
|
||||
|
||||
logger.debug(f"[ANN] Starting: {len(unit_ids)} seeds, top_k={top_k}")
|
||||
|
||||
# Build per-unit fact_types (default to 'world' if not provided)
|
||||
@@ -760,54 +753,70 @@ async def compute_semantic_links_ann(
|
||||
# sequential-scan every HNSW probe result against the array, destroying
|
||||
# performance (67s for 8k seeds). Self-links are harmless (ON CONFLICT DO
|
||||
# NOTHING handles duplicates in memory_links).
|
||||
t_setup = time_mod.time()
|
||||
await conn.execute("CREATE TEMP TABLE IF NOT EXISTS _ann_seeds (unit_id text, emb_text text, fact_type text)")
|
||||
await conn.execute("TRUNCATE _ann_seeds")
|
||||
#
|
||||
# The entire CREATE TEMP TABLE → COPY → SELECT sequence MUST run inside a
|
||||
# single transaction. Callers may connect through pgBouncer in `transaction`
|
||||
# pool mode, in which case the backend is only pinned to the client for the
|
||||
# duration of a transaction. Outside a transaction, pgBouncer can rebind
|
||||
# the client to a different backend between statements, and the temp table
|
||||
# (which is session-scoped to its creating backend) becomes invisible.
|
||||
# The observed failure mode was an intermittent
|
||||
# `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 hnsw.ef_search — the transaction
|
||||
# end handles both.
|
||||
rows: list = []
|
||||
async with conn.transaction():
|
||||
# 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")
|
||||
|
||||
records = [
|
||||
(uid, emb if isinstance(emb, str) else str(emb), ft) for uid, emb, ft in zip(unit_ids, embeddings, fact_types)
|
||||
]
|
||||
await conn.copy_records_to_table("_ann_seeds", records=records, columns=["unit_id", "emb_text", "fact_type"])
|
||||
logger.debug(f"[ANN] Temp table setup: {time_mod.time() - t_setup:.3f}s ({len(records)} seeds)")
|
||||
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")
|
||||
|
||||
# Run one ANN query per fact_type so each uses the right HNSW index.
|
||||
rows = []
|
||||
active_types = set(fact_types)
|
||||
for fact_type in active_types:
|
||||
t_query = time_mod.time()
|
||||
seed_count = sum(1 for ft in fact_types if ft == fact_type)
|
||||
logger.debug(f"[ANN] Querying fact_type={fact_type}: {seed_count} seeds")
|
||||
ft_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT s.unit_id AS from_id,
|
||||
n.id::text AS to_id,
|
||||
n.similarity
|
||||
FROM _ann_seeds s
|
||||
CROSS JOIN LATERAL (
|
||||
SELECT mu.id,
|
||||
1 - (mu.embedding <=> s.emb_text::vector) AS similarity
|
||||
FROM {fq_table("memory_units")} mu
|
||||
WHERE mu.bank_id = $1
|
||||
AND mu.fact_type = $2
|
||||
AND mu.embedding IS NOT NULL
|
||||
ORDER BY mu.embedding <=> s.emb_text::vector
|
||||
LIMIT $3
|
||||
) n
|
||||
WHERE s.fact_type = $2
|
||||
""",
|
||||
bank_id,
|
||||
fact_type,
|
||||
top_k,
|
||||
timeout=300, # ANN on large banks can take minutes
|
||||
)
|
||||
logger.debug(f"[ANN] fact_type={fact_type}: {len(ft_rows)} rows in {time_mod.time() - t_query:.3f}s")
|
||||
rows.extend(ft_rows)
|
||||
records = [
|
||||
(uid, emb if isinstance(emb, str) else str(emb), ft)
|
||||
for uid, emb, ft in zip(unit_ids, embeddings, fact_types)
|
||||
]
|
||||
await conn.copy_records_to_table("_ann_seeds", records=records, columns=["unit_id", "emb_text", "fact_type"])
|
||||
logger.debug(f"[ANN] Temp table setup: {time_mod.time() - t_setup:.3f}s ({len(records)} seeds)")
|
||||
|
||||
# Clean up temp table (no ON COMMIT DROP since we're not in a transaction)
|
||||
await conn.execute("DROP TABLE IF EXISTS _ann_seeds")
|
||||
|
||||
# Reset ef_search to default so the pooled connection doesn't affect recall queries
|
||||
await conn.execute("RESET hnsw.ef_search")
|
||||
# Run one ANN query per fact_type so each uses the right HNSW index.
|
||||
active_types = set(fact_types)
|
||||
for fact_type in active_types:
|
||||
t_query = time_mod.time()
|
||||
seed_count = sum(1 for ft in fact_types if ft == fact_type)
|
||||
logger.debug(f"[ANN] Querying fact_type={fact_type}: {seed_count} seeds")
|
||||
ft_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT s.unit_id AS from_id,
|
||||
n.id::text AS to_id,
|
||||
n.similarity
|
||||
FROM _ann_seeds s
|
||||
CROSS JOIN LATERAL (
|
||||
SELECT mu.id,
|
||||
1 - (mu.embedding <=> s.emb_text::vector) AS similarity
|
||||
FROM {fq_table("memory_units")} mu
|
||||
WHERE mu.bank_id = $1
|
||||
AND mu.fact_type = $2
|
||||
AND mu.embedding IS NOT NULL
|
||||
ORDER BY mu.embedding <=> s.emb_text::vector
|
||||
LIMIT $3
|
||||
) n
|
||||
WHERE s.fact_type = $2
|
||||
""",
|
||||
bank_id,
|
||||
fact_type,
|
||||
top_k,
|
||||
)
|
||||
logger.debug(f"[ANN] fact_type={fact_type}: {len(ft_rows)} rows in {time_mod.time() - t_query:.3f}s")
|
||||
rows.extend(ft_rows)
|
||||
# Transaction commits here. _ann_seeds is dropped (ON COMMIT DROP).
|
||||
# hnsw.ef_search reverts (SET LOCAL).
|
||||
|
||||
for row in rows:
|
||||
sim = float(min(1.0, max(0.0, row["similarity"])))
|
||||
|
||||
@@ -14,8 +14,9 @@ from collections.abc import Awaitable, Callable
|
||||
from datetime import UTC, datetime
|
||||
from typing import Any
|
||||
|
||||
from ...worker.stage import set_stage
|
||||
from ..db_utils import acquire_with_retry
|
||||
from ..memory_engine import fq_table
|
||||
from ..memory_engine import count_tokens, fq_table
|
||||
from . import bank_utils
|
||||
|
||||
|
||||
@@ -24,6 +25,32 @@ def utcnow():
|
||||
return datetime.now(UTC)
|
||||
|
||||
|
||||
def _merge_processed_content_tokens(a: int | None, b: int | None) -> int | None:
|
||||
"""Combine the processed-content-tokens signal across sub-results.
|
||||
|
||||
Semantics (see RetainResult.processed_content_tokens):
|
||||
* None means "this part of the retain did not go through chunk-level
|
||||
dedup" — i.e. the entire submitted payload was processed. If any
|
||||
sub-result is None, the aggregate is None so callers conservatively
|
||||
bill the full content.
|
||||
* Otherwise, accumulate the int values.
|
||||
"""
|
||||
if a is None or b is None:
|
||||
return None
|
||||
return a + b
|
||||
|
||||
|
||||
def _count_delta_content_tokens(delta_contents: list["RetainContent"]) -> int:
|
||||
"""Sum content + context tokens across the chunk items that were
|
||||
actually fed into the extraction pipeline on a partial-delta retain.
|
||||
"""
|
||||
total = 0
|
||||
for c in delta_contents:
|
||||
total += count_tokens(c.content or "")
|
||||
total += count_tokens(c.context or "")
|
||||
return total
|
||||
|
||||
|
||||
def parse_datetime_flexible(value: Any) -> datetime:
|
||||
"""
|
||||
Parse a datetime value that could be either a datetime object or an ISO string.
|
||||
@@ -71,7 +98,6 @@ from . import (
|
||||
from .types import (
|
||||
ChunkMetadata,
|
||||
EntityResolutionResult,
|
||||
ExtractedFact,
|
||||
Phase1Result,
|
||||
Phase3Context,
|
||||
ProcessedFact,
|
||||
@@ -133,6 +159,7 @@ async def _pre_resolve_phase1(
|
||||
Running these outside the transaction avoids holding row locks during
|
||||
slow reads, eliminating TimeoutErrors under concurrent load.
|
||||
"""
|
||||
set_stage("retain.phase1.resolve")
|
||||
from .link_utils import compute_semantic_links_ann
|
||||
|
||||
user_entities_per_content = {idx: content.entities for idx, content in enumerate(contents) if content.entities}
|
||||
@@ -238,6 +265,7 @@ async def _insert_facts_and_links(
|
||||
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)
|
||||
step_start = time.time()
|
||||
log_buffer.append(f" Insert facts: {len(unit_ids)} units in {time.time() - step_start:.3f}s")
|
||||
@@ -299,8 +327,11 @@ async def _insert_facts_and_links(
|
||||
causal_link_count = await link_creation.create_causal_links_batch(conn, bank_id, unit_ids, processed_facts)
|
||||
log_buffer.append(f" Causal links: {causal_link_count} links in {time.time() - step_start:.3f}s")
|
||||
|
||||
# Map results back to original content items
|
||||
result_unit_ids = _map_results_to_contents(contents, extracted_facts, unit_ids if unit_ids else [])
|
||||
# Map results back to original content items. Use processed_facts (not
|
||||
# extracted_facts) because unit_ids has 1:1 alignment with processed_facts —
|
||||
# any upstream drop between extraction and processing would otherwise cause
|
||||
# an IndexError (see issue #1037).
|
||||
result_unit_ids = _map_results_to_contents(contents, processed_facts, unit_ids if unit_ids else [])
|
||||
|
||||
if outbox_callback:
|
||||
await outbox_callback(conn)
|
||||
@@ -322,6 +353,7 @@ async def _build_and_insert_entity_links_phase3(
|
||||
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
|
||||
@@ -367,6 +399,7 @@ async def _extract_and_embed(
|
||||
Returns:
|
||||
Tuple of (extracted_facts, processed_facts, chunks_metadata, usage)
|
||||
"""
|
||||
set_stage("retain.extract_and_embed")
|
||||
step_start = time.time()
|
||||
extracted_facts, chunks, usage = await fact_extraction.extract_facts_from_contents(
|
||||
contents, llm_config, agent_name, config, pool, operation_id, schema
|
||||
@@ -410,13 +443,21 @@ async def retain_batch(
|
||||
schema: str | None = None,
|
||||
outbox_callback: Callable[["asyncpg.Connection"], Awaitable[None]] | None = None,
|
||||
db_semaphore: "asyncio.Semaphore | None" = None,
|
||||
) -> tuple[list[list[str]], TokenUsage]:
|
||||
) -> tuple[list[list[str]], TokenUsage, int | None]:
|
||||
"""
|
||||
Process a batch of content through the retain pipeline.
|
||||
|
||||
Supports delta retain: when upserting a document that already has chunks,
|
||||
only re-processes chunks whose content has changed. Unchanged chunks keep
|
||||
their existing facts, entities, and links.
|
||||
|
||||
Returns a three-tuple of:
|
||||
* per-content-item unit ID lists
|
||||
* aggregate LLM token usage
|
||||
* processed_content_tokens — content+context tokens that actually went
|
||||
through extraction after chunk-level dedup, or ``None`` if this path
|
||||
didn't dedup (caller should treat as "bill full submitted content").
|
||||
See ``RetainResult.processed_content_tokens`` for details.
|
||||
"""
|
||||
start_time = time.time()
|
||||
total_chars = sum(len(item.get("content", "")) for item in contents_dicts)
|
||||
@@ -456,8 +497,9 @@ async def retain_batch(
|
||||
# Process each group and merge results back in original order
|
||||
result_unit_ids: list[list[str]] = [[] for _ in contents_dicts]
|
||||
total_usage = TokenUsage()
|
||||
total_processed_tokens: int | None = 0
|
||||
for doc_key, (group_dicts, group_contents) in groups.items():
|
||||
group_ids, group_usage = await retain_batch(
|
||||
group_ids, group_usage, group_processed = await retain_batch(
|
||||
pool=pool,
|
||||
embeddings_model=embeddings_model,
|
||||
llm_config=llm_config,
|
||||
@@ -479,11 +521,12 @@ async def retain_batch(
|
||||
if group_idx < len(group_ids):
|
||||
result_unit_ids[orig_idx] = group_ids[group_idx]
|
||||
total_usage = total_usage + group_usage
|
||||
return result_unit_ids, total_usage
|
||||
total_processed_tokens = _merge_processed_content_tokens(total_processed_tokens, group_processed)
|
||||
return result_unit_ids, total_usage, total_processed_tokens
|
||||
|
||||
# Resolve effective document_id early so both delta and streaming paths
|
||||
# can find existing chunks from a prior attempt. On retry, the generated
|
||||
# document_id is recovered from operation result_metadata.
|
||||
# can find existing chunks from a prior attempt. On retry, a generated
|
||||
# document_id is recovered from operation result_metadata.document_ids[0].
|
||||
effective_doc_id = document_id
|
||||
if not effective_doc_id:
|
||||
doc_ids = {item.get("document_id") for item in contents_dicts if item.get("document_id")}
|
||||
@@ -502,26 +545,95 @@ async def retain_batch(
|
||||
if isinstance(row["result_metadata"], dict)
|
||||
else json.loads(row["result_metadata"])
|
||||
)
|
||||
effective_doc_id = meta.get("generated_document_id")
|
||||
recovered = meta.get("document_ids") or []
|
||||
if recovered:
|
||||
effective_doc_id = recovered[0]
|
||||
except Exception:
|
||||
pass
|
||||
if not effective_doc_id:
|
||||
effective_doc_id = str(uuid.uuid4())
|
||||
# Persist so retries reuse the same document_id
|
||||
if operation_id:
|
||||
try:
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
await conn.execute(
|
||||
f"""
|
||||
UPDATE {fq_table("async_operations")}
|
||||
SET result_metadata = result_metadata || $1::jsonb, updated_at = now()
|
||||
WHERE operation_id = $2
|
||||
""",
|
||||
json.dumps({"generated_document_id": effective_doc_id}),
|
||||
uuid.UUID(operation_id),
|
||||
)
|
||||
except Exception:
|
||||
logger.warning("Failed to persist generated document_id", exc_info=True)
|
||||
|
||||
# Record effective_doc_id on the operation (idempotent set-append). Captures
|
||||
# both user-provided and generated ids so the operation shows every document
|
||||
# it touched, and lets retries reuse the same generated id.
|
||||
if operation_id:
|
||||
try:
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
await conn.execute(
|
||||
f"""
|
||||
UPDATE {fq_table("async_operations")}
|
||||
SET result_metadata = jsonb_set(
|
||||
COALESCE(result_metadata, '{{}}'::jsonb),
|
||||
'{{document_ids}}',
|
||||
CASE
|
||||
WHEN COALESCE(result_metadata->'document_ids', '[]'::jsonb) @> $1::jsonb
|
||||
THEN result_metadata->'document_ids'
|
||||
ELSE COALESCE(result_metadata->'document_ids', '[]'::jsonb) || $1::jsonb
|
||||
END,
|
||||
true
|
||||
),
|
||||
updated_at = now()
|
||||
WHERE operation_id = $2
|
||||
""",
|
||||
json.dumps([effective_doc_id]),
|
||||
uuid.UUID(operation_id),
|
||||
)
|
||||
except Exception:
|
||||
logger.warning("Failed to persist document_id", exc_info=True)
|
||||
|
||||
# --- Append mode: prepend existing document content to new content ---
|
||||
# When update_mode="append", fetch the existing document text and prepend it
|
||||
# so the full document is reprocessed (delta retain will skip unchanged chunks).
|
||||
update_mode = None
|
||||
for item in contents_dicts:
|
||||
item_mode = item.get("update_mode")
|
||||
if item_mode:
|
||||
update_mode = item_mode
|
||||
break
|
||||
|
||||
if update_mode == "append" and effective_doc_id and is_first_batch:
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
existing_text = await fact_storage.get_document_content(conn, bank_id, effective_doc_id)
|
||||
if existing_text:
|
||||
# Prepend existing text as a new content item at the beginning
|
||||
existing_content: RetainContentDict = {"content": existing_text}
|
||||
# Copy context/tags from first item for consistency
|
||||
first = contents_dicts[0]
|
||||
if first.get("context"):
|
||||
existing_content["context"] = first["context"]
|
||||
if first.get("tags"):
|
||||
existing_content["tags"] = first["tags"]
|
||||
contents_dicts = [existing_content, *contents_dicts]
|
||||
# Rebuild contents list to match
|
||||
contents = _build_contents(contents_dicts, document_tags)
|
||||
log_buffer.append(
|
||||
f"[append] Prepended {len(existing_text):,} chars from existing document {effective_doc_id}"
|
||||
)
|
||||
|
||||
# --- Stale-request check (best-effort, before LLM extraction) ---
|
||||
# If the document was already updated by a more recent retain (updated_at > our
|
||||
# start_time), skip this request entirely to avoid overwriting newer content
|
||||
# (e.g. a longer conversation) with older data. This is an optimization — the
|
||||
# real correctness guarantee comes from the FOR UPDATE + content_hash check
|
||||
# inside each batch TXN (see _run_mini_batch_db_work).
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
doc_row = await conn.fetchrow(
|
||||
f"SELECT updated_at FROM {fq_table('documents')} WHERE id = $1 AND bank_id = $2",
|
||||
effective_doc_id,
|
||||
bank_id,
|
||||
)
|
||||
if doc_row and doc_row["updated_at"]:
|
||||
doc_updated = doc_row["updated_at"].timestamp()
|
||||
if doc_updated > start_time:
|
||||
log_buffer.append(
|
||||
f"[stale] Skipping retain: document {effective_doc_id} was updated at "
|
||||
f"{doc_row['updated_at'].isoformat()} (after this request started at "
|
||||
f"{datetime.fromtimestamp(start_time, tz=UTC).isoformat()})"
|
||||
)
|
||||
logger.info("\n" + "\n".join(log_buffer) + "\n")
|
||||
# No new content was processed — report 0 so callers can skip
|
||||
# billing cleanly instead of falling back to full-content billing.
|
||||
return [[] for _ in contents], TokenUsage(), 0
|
||||
|
||||
# --- Delta retain: check if we can skip unchanged chunks ---
|
||||
if is_first_batch:
|
||||
@@ -679,7 +791,6 @@ async def _run_final_semantic_ann(
|
||||
async with ann_semaphore:
|
||||
t0 = time.time()
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
await conn.execute("SET statement_timeout = '300s'")
|
||||
ann_links = await compute_semantic_links_ann(
|
||||
conn,
|
||||
bank_id,
|
||||
@@ -692,7 +803,6 @@ async def _run_final_semantic_ann(
|
||||
if ann_links:
|
||||
await _bulk_insert_links(conn, ann_links, bank_id=bank_id)
|
||||
chunk_link_counts[chunk_idx] = len(ann_links)
|
||||
await conn.execute("RESET statement_timeout")
|
||||
logger.info(
|
||||
f"[streaming] Final ANN chunk {chunk_idx + 1}/{num_chunks}: "
|
||||
f"{len(ann_links)} links in {time.time() - t0:.3f}s"
|
||||
@@ -758,25 +868,27 @@ async def _streaming_retain_batch(
|
||||
# Default template for metadata (context, event_date, etc.) when content list is empty.
|
||||
_default_content = RetainContent(content="")
|
||||
|
||||
# Load existing chunk hashes BEFORE document tracking to detect recovery.
|
||||
# If chunks exist AND the document content hash matches, this is a retry of
|
||||
# the same content — preserve existing data. If content differs, this is an
|
||||
# update — cascade-delete old data and start fresh.
|
||||
# ---------------------------------------------------------------------------
|
||||
# Recovery detection (read-only, before LLM extraction)
|
||||
# ---------------------------------------------------------------------------
|
||||
# Check if this is a retry of the same content (crash recovery). If the
|
||||
# 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()
|
||||
combined_content = "\n".join([c.get("content", "") for c in contents_dicts])
|
||||
new_content_hash = hashlib.sha256(combined_content.encode()).hexdigest()
|
||||
# Sanitize before hashing to match what handle_document_tracking stores
|
||||
sanitized_content = fact_extraction._sanitize_text(combined_content) or ""
|
||||
new_content_hash = hashlib.sha256(sanitized_content.encode()).hexdigest()
|
||||
is_recovery = False
|
||||
|
||||
try:
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
# Check if document exists with matching content hash
|
||||
doc_row = await conn.fetchrow(
|
||||
f"SELECT content_hash FROM {fq_table('documents')} WHERE id = $1 AND bank_id = $2",
|
||||
effective_doc_id,
|
||||
bank_id,
|
||||
)
|
||||
if doc_row and doc_row["content_hash"] == new_content_hash:
|
||||
# Same content — load chunk hashes for recovery skip
|
||||
existing_rows = await chunk_storage.load_existing_chunks(conn, bank_id, effective_doc_id)
|
||||
existing_chunk_hashes = {c.content_hash for c in existing_rows if c.content_hash}
|
||||
if existing_chunk_hashes:
|
||||
@@ -788,24 +900,22 @@ async def _streaming_retain_batch(
|
||||
except Exception:
|
||||
pass # If we can't load, just process all chunks
|
||||
|
||||
# Create/update the document row.
|
||||
# ---------------------------------------------------------------------------
|
||||
# Document tracking is DEFERRED to the first consumer batch TXN.
|
||||
# ---------------------------------------------------------------------------
|
||||
# Previously, document tracking (cascade-delete old data + insert doc row)
|
||||
# ran in a separate transaction BEFORE LLM extraction. This left a gap
|
||||
# between the cascade-delete and the first chunk write, allowing concurrent
|
||||
# requests to interleave and produce duplicates.
|
||||
#
|
||||
# Now, document tracking runs atomically inside the first batch's write TXN,
|
||||
# using SELECT ... FOR UPDATE on the document row for serialization across
|
||||
# workers. Each batch TXN also verifies document ownership via content_hash
|
||||
# to detect when a concurrent request has taken over the document.
|
||||
# See _run_mini_batch_db_work() for the implementation.
|
||||
retain_params, merged_tags = _build_retain_params(contents_dicts, document_tags)
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
async with conn.transaction():
|
||||
if is_recovery:
|
||||
# Recovery: same content, partially committed — preserve existing data
|
||||
await fact_storage.upsert_document_metadata(
|
||||
conn, bank_id, effective_doc_id, combined_content, retain_params, merged_tags
|
||||
)
|
||||
log_buffer.append(
|
||||
f"[streaming] Document {effective_doc_id} updated (recovery, preserving existing chunks)"
|
||||
)
|
||||
else:
|
||||
# Fresh or update: cascade-delete old data if document exists
|
||||
await fact_storage.handle_document_tracking(
|
||||
conn, bank_id, effective_doc_id, combined_content, is_first_batch, retain_params, merged_tags
|
||||
)
|
||||
log_buffer.append(f"[streaming] Document {effective_doc_id} tracked (full content)")
|
||||
# Track whether document tracking has been done (by the first batch)
|
||||
doc_tracking_done = [False]
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Producer-consumer pipeline: LLM extraction runs concurrently with DB writes
|
||||
@@ -818,6 +928,10 @@ async def _streaming_retain_batch(
|
||||
|
||||
# Shared mutable state for the producer to report skipped chunks and usage
|
||||
producer_error: list[BaseException] = []
|
||||
# Set to True by _run_mini_batch_db_work when a concurrent request takes
|
||||
# over the document (content_hash mismatch). The consumer checks this and
|
||||
# stops processing further batches.
|
||||
pipeline_aborted: list[bool] = [False]
|
||||
|
||||
# ---- LLM Producer ----
|
||||
# Fires all chunk extractions as concurrent tasks (bounded by the LLM
|
||||
@@ -876,17 +990,15 @@ async def _streaming_retain_batch(
|
||||
# Phase 1 (entity resolution) -> Phase 2 (write txn) -> Phase 3 (ANN fire-and-forget).
|
||||
async def _db_consumer() -> None:
|
||||
batch: list[tuple] = []
|
||||
global_chunk_offset = 0
|
||||
consumer_batch_idx = 0
|
||||
|
||||
while True:
|
||||
item = await chunk_queue.get()
|
||||
if item is None:
|
||||
# Process any remaining items
|
||||
if batch:
|
||||
if batch and not pipeline_aborted[0]:
|
||||
await _process_db_batch(
|
||||
batch,
|
||||
global_chunk_offset,
|
||||
consumer_batch_idx,
|
||||
is_last=True,
|
||||
)
|
||||
@@ -895,19 +1007,24 @@ async def _streaming_retain_batch(
|
||||
batch.append(item)
|
||||
|
||||
if len(batch) >= chunk_batch_size:
|
||||
if pipeline_aborted[0]:
|
||||
# Another request took over the document — discard this batch
|
||||
log_buffer.append(
|
||||
f"[streaming] Consumer: discarding batch of {len(batch)} chunks "
|
||||
f"(pipeline aborted due to concurrent takeover)"
|
||||
)
|
||||
batch = []
|
||||
continue
|
||||
await _process_db_batch(
|
||||
batch,
|
||||
global_chunk_offset,
|
||||
consumer_batch_idx,
|
||||
is_last=False,
|
||||
)
|
||||
global_chunk_offset += len(batch)
|
||||
consumer_batch_idx += 1
|
||||
batch = []
|
||||
|
||||
async def _process_db_batch(
|
||||
batch: list[tuple],
|
||||
global_chunk_offset: int,
|
||||
consumer_batch_idx: int,
|
||||
is_last: bool,
|
||||
) -> None:
|
||||
@@ -921,15 +1038,17 @@ async def _streaming_retain_batch(
|
||||
|
||||
for global_idx, content, extracted, processed, chunk_meta, usage in batch:
|
||||
content_idx_in_batch = len(batch_contents)
|
||||
# Adjust chunk indices to global offsets and remap content_index
|
||||
# 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.
|
||||
for fact in extracted:
|
||||
fact.content_index = content_idx_in_batch
|
||||
if fact.chunk_index is not None:
|
||||
fact.chunk_index = global_chunk_offset + content_idx_in_batch
|
||||
fact.chunk_index = global_idx
|
||||
for pf in processed:
|
||||
pf.content_index = content_idx_in_batch
|
||||
for cm in chunk_meta:
|
||||
cm.chunk_index = global_chunk_offset + content_idx_in_batch
|
||||
cm.chunk_index = global_idx
|
||||
|
||||
batch_contents.append(content)
|
||||
batch_extracted.extend(extracted)
|
||||
@@ -941,6 +1060,46 @@ async def _streaming_retain_batch(
|
||||
total_usage = total_usage + batch_usage
|
||||
|
||||
if not batch_extracted:
|
||||
# Even with 0 facts, the first batch must still run document tracking
|
||||
# (cascade-delete + insert doc row) to establish ownership and prevent
|
||||
# concurrent requests from interleaving. Later batches can safely skip.
|
||||
if not doc_tracking_done[0]:
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
async with conn.transaction():
|
||||
await conn.execute(
|
||||
f"INSERT INTO {fq_table('documents')} (id, bank_id, original_text, content_hash) "
|
||||
f"VALUES ($1, $2, '', '__pending__') "
|
||||
f"ON CONFLICT (id, bank_id) DO NOTHING",
|
||||
effective_doc_id,
|
||||
bank_id,
|
||||
)
|
||||
await conn.fetchval(
|
||||
f"SELECT content_hash FROM {fq_table('documents')} "
|
||||
f"WHERE id = $1 AND bank_id = $2 FOR UPDATE",
|
||||
effective_doc_id,
|
||||
bank_id,
|
||||
)
|
||||
if is_recovery:
|
||||
await fact_storage.upsert_document_metadata(
|
||||
conn,
|
||||
bank_id,
|
||||
effective_doc_id,
|
||||
combined_content,
|
||||
retain_params,
|
||||
merged_tags,
|
||||
)
|
||||
else:
|
||||
await fact_storage.handle_document_tracking(
|
||||
conn,
|
||||
bank_id,
|
||||
effective_doc_id,
|
||||
combined_content,
|
||||
is_first_batch,
|
||||
retain_params,
|
||||
merged_tags,
|
||||
)
|
||||
doc_tracking_done[0] = True
|
||||
log_buffer.append(f"[streaming] Document {effective_doc_id} tracked (0 facts in first batch)")
|
||||
log_buffer.append(
|
||||
f"[streaming] Consumer batch {consumer_batch_idx + 1}: "
|
||||
f"0 facts extracted from {len(batch)} chunks, skipping"
|
||||
@@ -971,10 +1130,92 @@ async def _streaming_retain_batch(
|
||||
|
||||
logger.info(f"[streaming] Phase 1 (entity resolution): {time.time() - p1_start:.3f}s")
|
||||
|
||||
# Phase 2 — Write transaction (within-batch semantic links only)
|
||||
# Phase 2 — Write transaction
|
||||
# -----------------------------------------------------------------
|
||||
# Concurrent-safety via row-level locking:
|
||||
#
|
||||
# The streaming pipeline splits work across multiple batch TXNs.
|
||||
# Without protection, two concurrent retains for the same document
|
||||
# can interleave: Request A writes batch1, Request B cascade-deletes
|
||||
# A's doc and writes its own batch1, then A's batch2 adds stale data
|
||||
# on top of B's → duplicates.
|
||||
#
|
||||
# To prevent this, every batch TXN:
|
||||
# 1. SELECT ... FOR UPDATE on the document row — serializes all
|
||||
# writers for this document at the DB level (works across workers).
|
||||
# 2. Check content_hash — if it doesn't match ours, another request
|
||||
# took over the document → abort remaining batches.
|
||||
# 3. First batch only: run handle_document_tracking (cascade-delete
|
||||
# old data + insert doc row) atomically with the first chunk write.
|
||||
# This eliminates the gap between "delete old" and "insert new"
|
||||
# that previously allowed interleaving.
|
||||
# -----------------------------------------------------------------
|
||||
|
||||
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 ---
|
||||
# Lock the document row to serialize all concurrent writers.
|
||||
# SELECT ... FOR UPDATE doesn't lock non-existent rows, so we
|
||||
# first ensure the row exists with a lightweight upsert, THEN lock it.
|
||||
# The content_hash='__pending__' placeholder is immediately overwritten
|
||||
# by handle_document_tracking or upsert_document_metadata below.
|
||||
await conn.execute(
|
||||
f"INSERT INTO {fq_table('documents')} (id, bank_id, original_text, content_hash) "
|
||||
f"VALUES ($1, $2, '', '__pending__') "
|
||||
f"ON CONFLICT (id, bank_id) DO NOTHING",
|
||||
effective_doc_id,
|
||||
bank_id,
|
||||
)
|
||||
existing_hash = await conn.fetchval(
|
||||
f"SELECT content_hash FROM {fq_table('documents')} WHERE id = $1 AND bank_id = $2 FOR UPDATE",
|
||||
effective_doc_id,
|
||||
bank_id,
|
||||
)
|
||||
|
||||
if not doc_tracking_done[0]:
|
||||
# --- First batch: document tracking (atomic with chunk write) ---
|
||||
if is_recovery:
|
||||
await fact_storage.upsert_document_metadata(
|
||||
conn,
|
||||
bank_id,
|
||||
effective_doc_id,
|
||||
combined_content,
|
||||
retain_params,
|
||||
merged_tags,
|
||||
)
|
||||
log_buffer.append(
|
||||
f"[streaming] Document {effective_doc_id} updated "
|
||||
f"(recovery, preserving existing chunks)"
|
||||
)
|
||||
else:
|
||||
await fact_storage.handle_document_tracking(
|
||||
conn,
|
||||
bank_id,
|
||||
effective_doc_id,
|
||||
combined_content,
|
||||
is_first_batch,
|
||||
retain_params,
|
||||
merged_tags,
|
||||
)
|
||||
log_buffer.append(f"[streaming] Document {effective_doc_id} tracked (full content)")
|
||||
doc_tracking_done[0] = True
|
||||
else:
|
||||
# --- Later batches: verify we still own the document ---
|
||||
# If another request took over (cascade-deleted our doc and
|
||||
# inserted its own), the content_hash won't match ours.
|
||||
if existing_hash is not None and existing_hash != new_content_hash:
|
||||
log_buffer.append(
|
||||
f"[streaming] Document {effective_doc_id} taken over by "
|
||||
f"concurrent request (hash mismatch) — aborting remaining batches"
|
||||
)
|
||||
logger.info("\n" + "\n".join(log_buffer) + "\n")
|
||||
# Signal the consumer to stop processing further batches
|
||||
pipeline_aborted[0] = True
|
||||
return
|
||||
|
||||
# Store chunks with correct global indices
|
||||
step_start = time.time()
|
||||
chunk_id_map = {}
|
||||
@@ -1016,11 +1257,14 @@ async def _streaming_retain_batch(
|
||||
logger.info(f"[streaming] Phase 2 (write txn): {time.time() - p2_start:.3f}s")
|
||||
|
||||
# Best-effort: entity viz + stats (fast, not semantic ANN)
|
||||
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)
|
||||
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 "
|
||||
@@ -1028,8 +1272,9 @@ async def _streaming_retain_batch(
|
||||
)
|
||||
|
||||
# Collect unit_ids from this batch
|
||||
for content_ids in batch_result_ids:
|
||||
all_unit_ids.extend(content_ids)
|
||||
if batch_result_ids:
|
||||
for content_ids in batch_result_ids:
|
||||
all_unit_ids.extend(content_ids)
|
||||
|
||||
if db_semaphore is not None:
|
||||
async with db_semaphore:
|
||||
@@ -1072,6 +1317,47 @@ async def _streaming_retain_batch(
|
||||
if producer_error:
|
||||
raise producer_error[0]
|
||||
|
||||
# If no batch was processed (e.g. zero facts extracted from gibberish
|
||||
# content, or all chunks skipped in recovery), the document row was
|
||||
# never created by the first batch TXN. Create it now so the document
|
||||
# is tracked regardless of extraction results.
|
||||
if not doc_tracking_done[0] and not pipeline_aborted[0]:
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
async with conn.transaction():
|
||||
await conn.execute(
|
||||
f"INSERT INTO {fq_table('documents')} (id, bank_id, original_text, content_hash) "
|
||||
f"VALUES ($1, $2, '', '__pending__') "
|
||||
f"ON CONFLICT (id, bank_id) DO NOTHING",
|
||||
effective_doc_id,
|
||||
bank_id,
|
||||
)
|
||||
await conn.fetchval(
|
||||
f"SELECT content_hash FROM {fq_table('documents')} WHERE id = $1 AND bank_id = $2 FOR UPDATE",
|
||||
effective_doc_id,
|
||||
bank_id,
|
||||
)
|
||||
if is_recovery:
|
||||
await fact_storage.upsert_document_metadata(
|
||||
conn,
|
||||
bank_id,
|
||||
effective_doc_id,
|
||||
combined_content,
|
||||
retain_params,
|
||||
merged_tags,
|
||||
)
|
||||
else:
|
||||
await fact_storage.handle_document_tracking(
|
||||
conn,
|
||||
bank_id,
|
||||
effective_doc_id,
|
||||
combined_content,
|
||||
is_first_batch,
|
||||
retain_params,
|
||||
merged_tags,
|
||||
)
|
||||
doc_tracking_done[0] = True
|
||||
log_buffer.append(f"[streaming] Document {effective_doc_id} tracked (no facts extracted)")
|
||||
|
||||
# Mark facts as committed in operation metadata (crash recovery checkpoint)
|
||||
if operation_id and all_unit_ids:
|
||||
try:
|
||||
@@ -1108,16 +1394,31 @@ async def _streaming_retain_batch(
|
||||
# This replaces per-batch within-batch + fire-and-forget ANN with a single
|
||||
# efficient pass after all facts are in the database.
|
||||
# ---------------------------------------------------------------------------
|
||||
if all_unit_ids:
|
||||
if all_unit_ids and not pipeline_aborted[0]:
|
||||
ann_start = time.time()
|
||||
await _run_final_semantic_ann(pool, bank_id, all_unit_ids, log_buffer)
|
||||
try:
|
||||
await _run_final_semantic_ann(pool, bank_id, all_unit_ids, log_buffer)
|
||||
except Exception:
|
||||
# ANN pass is best-effort. FK violations can occur if a concurrent
|
||||
# retain cascade-deleted our units between the batch commit and here.
|
||||
logger.warning(
|
||||
f"[streaming] Final ANN pass failed for document {effective_doc_id} "
|
||||
f"(units may have been superseded by concurrent retain)",
|
||||
exc_info=True,
|
||||
)
|
||||
log_buffer.append(f"[streaming] Final ANN pass: {time.time() - ann_start:.3f}s for {len(all_unit_ids)} units")
|
||||
|
||||
total_time = time.time() - start_time
|
||||
log_buffer.append(f"{'=' * 60}")
|
||||
log_buffer.append(
|
||||
f"STREAMING RETAIN COMPLETE: {len(all_unit_ids)} units across {num_batches} batches in {total_time:.3f}s"
|
||||
)
|
||||
if pipeline_aborted[0]:
|
||||
log_buffer.append(
|
||||
f"STREAMING RETAIN ABORTED: document {effective_doc_id} was taken over by "
|
||||
f"a concurrent request after {total_time:.3f}s — data from this request was discarded"
|
||||
)
|
||||
else:
|
||||
log_buffer.append(
|
||||
f"STREAMING RETAIN COMPLETE: {len(all_unit_ids)} units across {num_batches} batches in {total_time:.3f}s"
|
||||
)
|
||||
log_buffer.append(f"Document: {effective_doc_id}")
|
||||
log_buffer.append(f"{'=' * 60}")
|
||||
logger.info("\n" + "\n".join(log_buffer) + "\n")
|
||||
@@ -1125,7 +1426,10 @@ async def _streaming_retain_batch(
|
||||
# Map all unit_ids back to the original content items.
|
||||
# For streaming mode with a single document, all units belong to content 0.
|
||||
result_unit_ids = [all_unit_ids] + [[] for _ in contents[1:]]
|
||||
return result_unit_ids, total_usage
|
||||
# The streaming path doesn't compute per-chunk content-hash dedup in
|
||||
# a way that lets us report a partial-processed tokens count — signal
|
||||
# ``None`` so callers bill against the full submitted payload.
|
||||
return result_unit_ids, total_usage, None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -1153,10 +1457,15 @@ async def _try_delta_retain(
|
||||
schema,
|
||||
outbox_callback,
|
||||
db_semaphore: "asyncio.Semaphore | None" = None,
|
||||
):
|
||||
) -> tuple[list[list[str]], TokenUsage, int | None] | None:
|
||||
"""
|
||||
Attempt delta retain for a document upsert. Returns result tuple if delta
|
||||
was performed, or None to fall back to full retain.
|
||||
|
||||
When a result tuple is returned, the third element is the content+context
|
||||
token count for the chunks that actually went through extraction
|
||||
(``0`` if the submission matched prior content exactly and nothing was
|
||||
re-extracted).
|
||||
"""
|
||||
# Need a single document_id
|
||||
effective_doc_id = document_id
|
||||
@@ -1166,9 +1475,17 @@ async def _try_delta_retain(
|
||||
return None
|
||||
effective_doc_id = doc_ids.pop()
|
||||
|
||||
# Load existing chunks
|
||||
# Load existing chunks and snapshot the document's content_hash. This is
|
||||
# outside the write TXN, so a concurrent retain could modify the document
|
||||
# between this read and the write. The write TXN verifies the hash hasn't
|
||||
# changed; if it has, we fall back to streaming (which has full protection).
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
existing_chunks = await chunk_storage.load_existing_chunks(conn, bank_id, effective_doc_id)
|
||||
doc_hash_at_load = await conn.fetchval(
|
||||
f"SELECT content_hash FROM {fq_table('documents')} WHERE id = $1 AND bank_id = $2",
|
||||
effective_doc_id,
|
||||
bank_id,
|
||||
)
|
||||
|
||||
if not existing_chunks:
|
||||
return None
|
||||
@@ -1276,8 +1593,28 @@ async def _try_delta_retain(
|
||||
)
|
||||
|
||||
# PHASE 2 — Core Write Transaction (atomic)
|
||||
# Lock the document row and verify ownership. Delta loaded existing
|
||||
# chunks OUTSIDE this TXN, so a concurrent retain may have cascade-deleted
|
||||
# and replaced the document since then. If the content_hash changed,
|
||||
# the chunk state we based our delta diff on is stale — abort.
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
async with conn.transaction():
|
||||
current_hash = await conn.fetchval(
|
||||
f"SELECT content_hash FROM {fq_table('documents')} WHERE id = $1 AND bank_id = $2 FOR UPDATE",
|
||||
effective_doc_id,
|
||||
bank_id,
|
||||
)
|
||||
# Verify the document hasn't been replaced since we loaded chunks.
|
||||
# Compare the current hash against what we snapshotted at load time.
|
||||
if current_hash is not None and doc_hash_at_load is not None and current_hash != doc_hash_at_load:
|
||||
log_buffer.append(
|
||||
f"[delta] Document {effective_doc_id} was modified by concurrent request "
|
||||
f"since chunks were loaded — aborting delta, falling back to full retain"
|
||||
)
|
||||
logger.info("\n" + "\n".join(log_buffer) + "\n")
|
||||
# Return None to fall back to streaming (which has full FOR UPDATE protection)
|
||||
return None
|
||||
|
||||
# Update document metadata (no delete)
|
||||
step_start = time.time()
|
||||
combined_content = "\n".join([c.get("content", "") for c in contents_dicts])
|
||||
@@ -1387,7 +1724,12 @@ async def _try_delta_retain(
|
||||
await _run_delta_db_work()
|
||||
else:
|
||||
await _run_delta_db_work()
|
||||
return result_unit_ids, usage
|
||||
# Count content + context tokens that actually went through extraction.
|
||||
# ``delta_contents`` holds the per-chunk RetainContent items for the
|
||||
# changed/new chunks (see ``_build_delta_contents``) — i.e. exactly what
|
||||
# the LLM pipeline saw this call. Unchanged chunks contribute zero.
|
||||
processed_tokens = _count_delta_content_tokens(delta_contents)
|
||||
return result_unit_ids, usage, processed_tokens
|
||||
|
||||
|
||||
async def _delta_metadata_only(
|
||||
@@ -1404,6 +1746,12 @@ async def _delta_metadata_only(
|
||||
"""Handle the case where no chunks changed — just update document metadata and tags."""
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
async with conn.transaction():
|
||||
# Lock the document row to serialize with concurrent retains
|
||||
await conn.fetchval(
|
||||
f"SELECT content_hash FROM {fq_table('documents')} WHERE id = $1 AND bank_id = $2 FOR UPDATE",
|
||||
document_id,
|
||||
bank_id,
|
||||
)
|
||||
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(
|
||||
@@ -1421,7 +1769,11 @@ async def _delta_metadata_only(
|
||||
total_time = time.time() - start_time
|
||||
log_buffer.append(f"DELTA RETAIN (no changes): metadata updated in {total_time:.3f}s")
|
||||
logger.info("\n" + "\n".join(log_buffer) + "\n")
|
||||
return [[] for _ in contents], TokenUsage()
|
||||
# Nothing went through the extraction pipeline — report 0 processed
|
||||
# content tokens so callers can bill accordingly (a caller that's been
|
||||
# told ``0`` knows the retain was a pure metadata update and should
|
||||
# charge nothing for content).
|
||||
return [[] for _ in contents], TokenUsage(), 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -1516,21 +1868,29 @@ def _build_delta_contents(
|
||||
|
||||
def _map_results_to_contents(
|
||||
contents: list[RetainContent],
|
||||
extracted_facts: list[ExtractedFact],
|
||||
processed_facts: list[ProcessedFact],
|
||||
unit_ids: list[str],
|
||||
) -> list[list[str]]:
|
||||
"""Map created unit IDs back to original content items."""
|
||||
"""Map created unit IDs back to original content items.
|
||||
|
||||
`processed_facts` and `unit_ids` must have the same length: each unit_id
|
||||
corresponds to the processed_fact at the same index.
|
||||
"""
|
||||
if len(processed_facts) != len(unit_ids):
|
||||
raise ValueError(f"processed_facts ({len(processed_facts)}) and unit_ids ({len(unit_ids)}) length mismatch")
|
||||
|
||||
facts_by_content: dict[int, list[int]] = {i: [] for i in range(len(contents))}
|
||||
for i, fact in enumerate(extracted_facts):
|
||||
facts_by_content[fact.content_index].append(i)
|
||||
for i, fact in enumerate(processed_facts):
|
||||
# Normalize content_index: some LLM providers return 1-indexed values.
|
||||
# Clamp to valid range to prevent KeyError.
|
||||
idx = fact.content_index
|
||||
if idx < 0 or idx >= len(contents):
|
||||
idx = min(max(idx, 0), len(contents) - 1) if len(contents) > 0 else 0
|
||||
facts_by_content[idx].append(i)
|
||||
|
||||
result_unit_ids = []
|
||||
unit_idx = 0
|
||||
for content_index in range(len(contents)):
|
||||
content_unit_ids = []
|
||||
for _ in facts_by_content[content_index]:
|
||||
content_unit_ids.append(unit_ids[unit_idx])
|
||||
unit_idx += 1
|
||||
content_unit_ids = [unit_ids[i] for i in facts_by_content[content_index]]
|
||||
result_unit_ids.append(content_unit_ids)
|
||||
|
||||
return result_unit_ids
|
||||
|
||||
@@ -25,6 +25,9 @@ class RetainContentDict(TypedDict, total=False):
|
||||
observation_scopes: How to scope observations for consolidation (optional).
|
||||
"per_tag" runs one pass per individual tag; "combined" (default) runs a
|
||||
single pass with all tags; a list[list[str]] specifies exact passes.
|
||||
update_mode: How to handle existing documents with the same document_id (optional).
|
||||
"replace" (default) deletes old data and reprocesses. "append" concatenates
|
||||
new content to the existing document and reprocesses.
|
||||
"""
|
||||
|
||||
content: str # Required
|
||||
@@ -37,6 +40,7 @@ class RetainContentDict(TypedDict, total=False):
|
||||
observation_scopes: (
|
||||
Literal["per_tag", "combined", "all_combinations"] | list[list[str]]
|
||||
) # Observation scopes for consolidation
|
||||
update_mode: Literal["replace", "append"]
|
||||
|
||||
|
||||
@dataclass
|
||||
|
||||
@@ -8,6 +8,7 @@ of the recall pipeline.
|
||||
|
||||
import logging
|
||||
from abc import ABC, abstractmethod
|
||||
from datetime import datetime
|
||||
|
||||
from .tags import TagGroup, TagsMatch
|
||||
from .types import GraphRetrievalTimings, RetrievalResult
|
||||
@@ -45,6 +46,8 @@ class GraphRetriever(ABC):
|
||||
tags: list[str] | None = None, # Visibility scope tags for filtering
|
||||
tags_match: TagsMatch = "any", # How to match tags: 'any' (OR) or 'all' (AND)
|
||||
tag_groups: list[TagGroup] | None = None, # Compound boolean tag filter groups
|
||||
created_after: datetime | None = None, # Only include memory_units created after this time
|
||||
created_before: datetime | None = None, # Only include memory_units created before this time
|
||||
) -> tuple[list[RetrievalResult], GraphRetrievalTimings | None]:
|
||||
"""
|
||||
Retrieve relevant facts via graph traversal.
|
||||
|
||||
@@ -6,25 +6,32 @@ stored in memory_links:
|
||||
|
||||
1. Entity links — query-time self-join through unit_entities. Score = number of distinct
|
||||
shared entities between the seed set and each candidate, computed via
|
||||
COUNT(DISTINCT entity_id). More accurate than precomputed entity links.
|
||||
COUNT(DISTINCT entity_id). Uses a LATERAL per-entity cap
|
||||
(graph_per_entity_limit, default 200) to prevent high-fanout entities
|
||||
from exploding the self-join intermediate rows.
|
||||
2. Semantic links — precomputed kNN graph (each new fact linked to its top-5 most
|
||||
similar existing facts at insert time, similarity >= 0.7). Checked
|
||||
in both directions since the graph is not symmetric. Score = weight.
|
||||
3. Causal links — explicit causal chains (causes/caused_by/enables/prevents).
|
||||
Score = weight + 1.0 (boosted as highest-quality signal).
|
||||
|
||||
All three signals are bounded at retain time, so no LATERAL fan-out caps are needed
|
||||
at query time. Each expansion is a simple aggregation over a small result set.
|
||||
Entity expansion is bounded by graph_per_entity_limit (LATERAL cap per entity).
|
||||
A timeout fallback (graph_expansion_timeout) drops entity expansion entirely if the
|
||||
query still exceeds the budget.
|
||||
|
||||
For non-observation fact types the three expansions are issued as a single CTE query
|
||||
(one roundtrip, one connection) with a `source` discriminator column so the Python
|
||||
merge step can apply per-signal score transformations.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import math
|
||||
import time
|
||||
from datetime import datetime
|
||||
from typing import Any
|
||||
|
||||
from ...config import get_config
|
||||
from ..db_utils import acquire_with_retry
|
||||
from ..memory_engine import fq_table
|
||||
from .graph_retrieval import GraphRetriever
|
||||
@@ -44,6 +51,8 @@ async def _find_semantic_seeds(
|
||||
tags: list[str] | None = None,
|
||||
tags_match: TagsMatch = "any",
|
||||
tag_groups: list[TagGroup] | None = None,
|
||||
created_after: datetime | None = None,
|
||||
created_before: datetime | None = None,
|
||||
) -> list[RetrievalResult]:
|
||||
"""Find semantic seeds via embedding search."""
|
||||
from .tags import build_tag_groups_where_clause, build_tags_where_clause_simple
|
||||
@@ -51,15 +60,29 @@ async def _find_semantic_seeds(
|
||||
tags_clause = build_tags_where_clause_simple(tags, 6, match=tags_match)
|
||||
tag_groups_param_start = 6 + (1 if tags else 0)
|
||||
groups_clause, groups_params, _ = build_tag_groups_where_clause(tag_groups, tag_groups_param_start)
|
||||
|
||||
_next_idx = tag_groups_param_start + len(groups_params)
|
||||
created_range_clause = ""
|
||||
created_range_params: list[Any] = []
|
||||
if created_after is not None:
|
||||
created_range_params.append(created_after)
|
||||
created_range_clause += f" AND updated_at > ${_next_idx}"
|
||||
_next_idx += 1
|
||||
if created_before is not None:
|
||||
created_range_params.append(created_before)
|
||||
created_range_clause += f" AND updated_at < ${_next_idx}"
|
||||
_next_idx += 1
|
||||
|
||||
params = [query_embedding_str, bank_id, fact_type, threshold, limit]
|
||||
if tags:
|
||||
params.append(tags)
|
||||
params.extend(groups_params)
|
||||
params.extend(created_range_params)
|
||||
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT id, text, context, event_date, occurred_start, occurred_end,
|
||||
mentioned_at, fact_type, document_id, chunk_id, tags,
|
||||
mentioned_at, fact_type, document_id, chunk_id, tags, proof_count,
|
||||
1 - (embedding <=> $1::vector) AS similarity
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE bank_id = $2
|
||||
@@ -68,6 +91,7 @@ async def _find_semantic_seeds(
|
||||
AND (1 - (embedding <=> $1::vector)) >= $4
|
||||
{tags_clause}
|
||||
{groups_clause}
|
||||
{created_range_clause}
|
||||
ORDER BY embedding <=> $1::vector
|
||||
LIMIT $5
|
||||
""",
|
||||
@@ -116,6 +140,8 @@ class LinkExpansionRetriever(GraphRetriever):
|
||||
tags: list[str] | None = None,
|
||||
tags_match: TagsMatch = "any",
|
||||
tag_groups: list[TagGroup] | None = None,
|
||||
created_after: "datetime | None" = None,
|
||||
created_before: "datetime | None" = None,
|
||||
) -> tuple[list[RetrievalResult], GraphRetrievalTimings | None]:
|
||||
"""
|
||||
Retrieve facts by expanding links from seeds.
|
||||
@@ -154,6 +180,8 @@ class LinkExpansionRetriever(GraphRetriever):
|
||||
tags=tags,
|
||||
tags_match=tags_match,
|
||||
tag_groups=tag_groups,
|
||||
created_after=created_after,
|
||||
created_before=created_before,
|
||||
)
|
||||
timings.seeds_time = time.time() - seeds_start
|
||||
logger.debug(
|
||||
@@ -262,35 +290,48 @@ class LinkExpansionRetriever(GraphRetriever):
|
||||
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")
|
||||
mu = fq_table("memory_units")
|
||||
ue = fq_table("unit_entities")
|
||||
|
||||
per_entity_limit = config.link_expansion_per_entity_limit
|
||||
|
||||
# Entity CTE with LATERAL fanout cap.
|
||||
# Every seed entity (including high-frequency ones) is kept, but each
|
||||
# entity's expansion is capped to per_entity_limit target units. The
|
||||
# LATERAL subquery orders by unit_id DESC so the most recently inserted
|
||||
# units are preferred (a recency proxy that is free — it rides the PK
|
||||
# index with no extra sort).
|
||||
entity_cte = f"""
|
||||
seed_entities AS (
|
||||
SELECT DISTINCT ue.entity_id
|
||||
FROM {ue} ue
|
||||
WHERE ue.unit_id = ANY($1::uuid[])
|
||||
),
|
||||
entity_expanded AS (
|
||||
-- Entity co-occurrence via unit_entities self-join.
|
||||
-- Finds units sharing entities with seeds at query time — more accurate
|
||||
-- than precomputed entity links (no stale 50-neighbor cap).
|
||||
-- Score = COUNT(DISTINCT shared entities), mapped to [0,1] via tanh.
|
||||
SELECT mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
|
||||
mu.occurred_end, mu.mentioned_at,
|
||||
mu.fact_type, mu.document_id, mu.chunk_id, mu.tags,
|
||||
COUNT(DISTINCT ue_seed.entity_id)::float AS score,
|
||||
mu.fact_type, mu.document_id, mu.chunk_id, mu.tags, mu.proof_count,
|
||||
COUNT(DISTINCT se.entity_id)::float AS score,
|
||||
'entity'::text AS source
|
||||
FROM {ue} ue_seed
|
||||
JOIN {ue} ue_target ON ue_seed.entity_id = ue_target.entity_id
|
||||
JOIN {mu} mu ON mu.id = ue_target.unit_id
|
||||
WHERE ue_seed.unit_id = ANY($1::uuid[])
|
||||
AND ue_target.unit_id != ALL($1::uuid[])
|
||||
AND mu.fact_type = $2
|
||||
FROM seed_entities se
|
||||
CROSS JOIN LATERAL (
|
||||
SELECT ue_target.unit_id
|
||||
FROM {ue} ue_target
|
||||
WHERE ue_target.entity_id = se.entity_id
|
||||
AND ue_target.unit_id != ALL($1::uuid[])
|
||||
ORDER BY ue_target.unit_id DESC
|
||||
LIMIT {per_entity_limit}
|
||||
) t
|
||||
JOIN {mu} mu ON mu.id = t.unit_id
|
||||
WHERE mu.fact_type = $2
|
||||
GROUP BY mu.id
|
||||
ORDER BY score DESC
|
||||
LIMIT $3
|
||||
)"""
|
||||
|
||||
all_rows = await conn.fetch(
|
||||
f"""
|
||||
WITH {entity_cte},
|
||||
semantic_causal_cte = f"""
|
||||
semantic_expanded AS (
|
||||
-- Semantic kNN: both outgoing (seeds → their kNN at insert time) and
|
||||
-- incoming (facts inserted after seeds that found seeds as kNN).
|
||||
@@ -298,14 +339,14 @@ class LinkExpansionRetriever(GraphRetriever):
|
||||
SELECT
|
||||
id, text, context, event_date, occurred_start,
|
||||
occurred_end, mentioned_at,
|
||||
fact_type, document_id, chunk_id, tags,
|
||||
fact_type, document_id, chunk_id, tags, proof_count,
|
||||
MAX(weight) AS score,
|
||||
'semantic'::text AS source
|
||||
FROM (
|
||||
SELECT
|
||||
mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
|
||||
mu.occurred_end, mu.mentioned_at,
|
||||
mu.fact_type, mu.document_id, mu.chunk_id, mu.tags,
|
||||
mu.fact_type, mu.document_id, mu.chunk_id, mu.tags, mu.proof_count,
|
||||
ml.weight
|
||||
FROM {ml} ml
|
||||
JOIN {mu} mu ON mu.id = ml.to_unit_id
|
||||
@@ -317,7 +358,7 @@ class LinkExpansionRetriever(GraphRetriever):
|
||||
SELECT
|
||||
mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
|
||||
mu.occurred_end, mu.mentioned_at,
|
||||
mu.fact_type, mu.document_id, mu.chunk_id, mu.tags,
|
||||
mu.fact_type, mu.document_id, mu.chunk_id, mu.tags, mu.proof_count,
|
||||
ml.weight
|
||||
FROM {ml} ml
|
||||
JOIN {mu} mu ON mu.id = ml.from_unit_id
|
||||
@@ -328,7 +369,7 @@ class LinkExpansionRetriever(GraphRetriever):
|
||||
) sem_raw
|
||||
GROUP BY id, text, context, event_date, occurred_start,
|
||||
occurred_end, mentioned_at,
|
||||
fact_type, document_id, chunk_id, tags
|
||||
fact_type, document_id, chunk_id, tags, proof_count
|
||||
ORDER BY score DESC
|
||||
LIMIT $3
|
||||
),
|
||||
@@ -339,7 +380,7 @@ class LinkExpansionRetriever(GraphRetriever):
|
||||
SELECT DISTINCT ON (mu.id)
|
||||
mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
|
||||
mu.occurred_end, mu.mentioned_at,
|
||||
mu.fact_type, mu.document_id, mu.chunk_id, mu.tags,
|
||||
mu.fact_type, mu.document_id, mu.chunk_id, mu.tags, mu.proof_count,
|
||||
ml.weight AS score,
|
||||
'causal'::text AS source
|
||||
FROM {ml} ml
|
||||
@@ -350,18 +391,37 @@ class LinkExpansionRetriever(GraphRetriever):
|
||||
AND mu.fact_type = $2
|
||||
ORDER BY mu.id, ml.weight DESC
|
||||
LIMIT $3
|
||||
)
|
||||
)"""
|
||||
|
||||
full_query = f"""
|
||||
WITH {entity_cte},
|
||||
{semantic_causal_cte}
|
||||
SELECT * FROM entity_expanded
|
||||
UNION ALL
|
||||
SELECT * FROM semantic_expanded
|
||||
UNION ALL
|
||||
SELECT * FROM causal_expanded
|
||||
""",
|
||||
seed_ids,
|
||||
fact_type,
|
||||
budget,
|
||||
self.causal_weight_threshold,
|
||||
)
|
||||
"""
|
||||
|
||||
params = [seed_ids, fact_type, budget, self.causal_weight_threshold]
|
||||
|
||||
try:
|
||||
all_rows = await asyncio.wait_for(
|
||||
conn.fetch(full_query, *params),
|
||||
timeout=config.link_expansion_timeout,
|
||||
)
|
||||
except asyncio.TimeoutError:
|
||||
logger.warning(
|
||||
f"[LinkExpansion] Entity expansion timed out after {config.link_expansion_timeout}s "
|
||||
f"for fact_type={fact_type}, falling back to semantic+causal only"
|
||||
)
|
||||
fallback_query = f"""
|
||||
WITH {semantic_causal_cte}
|
||||
SELECT * FROM semantic_expanded
|
||||
UNION ALL
|
||||
SELECT * FROM causal_expanded
|
||||
"""
|
||||
all_rows = await conn.fetch(fallback_query, *params)
|
||||
|
||||
entity_rows = [r for r in all_rows if r["source"] == "entity"]
|
||||
semantic_rows = [r for r in all_rows if r["source"] == "semantic"]
|
||||
@@ -401,17 +461,31 @@ class LinkExpansionRetriever(GraphRetriever):
|
||||
f"{len(source_ids_found)} source_memory_ids found"
|
||||
)
|
||||
|
||||
config = get_config()
|
||||
ue = fq_table("unit_entities")
|
||||
per_entity_limit = config.link_expansion_per_entity_limit
|
||||
|
||||
connected_sources_cte = f"""
|
||||
connected_sources AS (
|
||||
-- Find sources sharing entities with seed observation sources
|
||||
-- via unit_entities self-join (query-time, no precomputed links needed).
|
||||
SELECT DISTINCT ue_target.unit_id AS source_id
|
||||
source_entities AS (
|
||||
SELECT DISTINCT ue_seed.entity_id
|
||||
FROM seed_sources ss
|
||||
JOIN {ue} ue_seed ON ue_seed.unit_id = ss.source_id
|
||||
JOIN {ue} ue_target ON ue_seed.entity_id = ue_target.entity_id
|
||||
WHERE ue_target.unit_id != ss.source_id
|
||||
),
|
||||
connected_sources AS (
|
||||
-- Find sources sharing entities with seed observation sources
|
||||
-- via LATERAL-capped self-join (prevents hub entity fanout).
|
||||
SELECT DISTINCT t.unit_id AS source_id
|
||||
FROM source_entities se
|
||||
CROSS JOIN LATERAL (
|
||||
SELECT ue_target.unit_id
|
||||
FROM {ue} ue_target
|
||||
WHERE ue_target.entity_id = se.entity_id
|
||||
ORDER BY ue_target.unit_id DESC
|
||||
LIMIT {per_entity_limit}
|
||||
) t
|
||||
WHERE NOT EXISTS (
|
||||
SELECT 1 FROM seed_sources ss WHERE ss.source_id = t.unit_id
|
||||
)
|
||||
)"""
|
||||
|
||||
entity_rows = await conn.fetch(
|
||||
@@ -429,7 +503,7 @@ class LinkExpansionRetriever(GraphRetriever):
|
||||
SELECT
|
||||
mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
|
||||
mu.occurred_end, mu.mentioned_at,
|
||||
mu.fact_type, mu.document_id, mu.chunk_id, mu.tags,
|
||||
mu.fact_type, mu.document_id, mu.chunk_id, mu.tags, mu.proof_count,
|
||||
(SELECT COUNT(DISTINCT s) FROM unnest(mu.source_memory_ids) s WHERE s = ANY(ca.source_ids))::float AS score
|
||||
FROM {fq_table("memory_units")} mu, connected_array ca
|
||||
WHERE mu.fact_type = 'observation'
|
||||
@@ -453,13 +527,13 @@ class LinkExpansionRetriever(GraphRetriever):
|
||||
SELECT
|
||||
id, text, context, event_date, occurred_start,
|
||||
occurred_end, mentioned_at,
|
||||
fact_type, document_id, chunk_id, tags,
|
||||
fact_type, document_id, chunk_id, tags, proof_count,
|
||||
MAX(weight) AS score,
|
||||
'semantic'::text AS source
|
||||
FROM (
|
||||
SELECT mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
|
||||
mu.occurred_end, mu.mentioned_at, mu.fact_type, mu.document_id,
|
||||
mu.chunk_id, mu.tags, ml.weight
|
||||
mu.chunk_id, mu.tags, mu.proof_count, ml.weight
|
||||
FROM {ml} ml JOIN {mu} mu ON mu.id = ml.to_unit_id
|
||||
WHERE ml.from_unit_id = ANY($1::uuid[])
|
||||
AND ml.link_type = 'semantic' AND mu.fact_type = 'observation'
|
||||
@@ -467,21 +541,21 @@ class LinkExpansionRetriever(GraphRetriever):
|
||||
UNION ALL
|
||||
SELECT mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
|
||||
mu.occurred_end, mu.mentioned_at, mu.fact_type, mu.document_id,
|
||||
mu.chunk_id, mu.tags, ml.weight
|
||||
mu.chunk_id, mu.tags, mu.proof_count, ml.weight
|
||||
FROM {ml} ml JOIN {mu} mu ON mu.id = ml.from_unit_id
|
||||
WHERE ml.to_unit_id = ANY($1::uuid[])
|
||||
AND ml.link_type = 'semantic' AND mu.fact_type = 'observation'
|
||||
AND mu.id != ALL($1::uuid[])
|
||||
) sem_raw
|
||||
GROUP BY id, text, context, event_date, occurred_start, occurred_end,
|
||||
mentioned_at, fact_type, document_id, chunk_id, tags
|
||||
mentioned_at, fact_type, document_id, chunk_id, tags, proof_count
|
||||
ORDER BY score DESC LIMIT $2
|
||||
),
|
||||
causal_expanded AS (
|
||||
SELECT DISTINCT ON (mu.id)
|
||||
mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
|
||||
mu.occurred_end, mu.mentioned_at, mu.fact_type, mu.document_id,
|
||||
mu.chunk_id, mu.tags, ml.weight AS score, 'causal'::text AS source
|
||||
mu.chunk_id, mu.tags, mu.proof_count, ml.weight AS score, 'causal'::text AS source
|
||||
FROM {ml} ml JOIN {mu} mu ON ml.to_unit_id = mu.id
|
||||
WHERE ml.from_unit_id = ANY($1::uuid[])
|
||||
AND ml.link_type IN ('causes', 'caused_by', 'enables', 'prevents')
|
||||
|
||||
@@ -2,6 +2,7 @@
|
||||
Cross-encoder neural reranking for search results.
|
||||
"""
|
||||
|
||||
import math
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from .types import MergedCandidate, ScoredResult
|
||||
@@ -13,6 +14,7 @@ UTC = timezone.utc
|
||||
# so the max combined boost is (1 + alpha/2)^2 ≈ +21% and min is (1 - alpha/2)^2 ≈ -19%.
|
||||
_RECENCY_ALPHA: float = 0.2
|
||||
_TEMPORAL_ALPHA: float = 0.2
|
||||
_PROOF_COUNT_ALPHA: float = 0.1 # Conservative: max ±5% for evidence strength
|
||||
|
||||
|
||||
def apply_combined_scoring(
|
||||
@@ -20,32 +22,81 @@ def apply_combined_scoring(
|
||||
now: datetime,
|
||||
recency_alpha: float = _RECENCY_ALPHA,
|
||||
temporal_alpha: float = _TEMPORAL_ALPHA,
|
||||
proof_count_alpha: float = _PROOF_COUNT_ALPHA,
|
||||
is_passthrough_reranker: bool = False,
|
||||
) -> None:
|
||||
"""Apply combined scoring to a list of ScoredResults in-place.
|
||||
|
||||
Uses the cross-encoder score as the primary relevance signal, with recency
|
||||
and temporal proximity applied as multiplicative boosts. This ensures the
|
||||
influence of these secondary signals is always proportional to the base
|
||||
relevance score, regardless of the cross-encoder model's score calibration.
|
||||
Uses the cross-encoder score as the primary relevance signal, with recency,
|
||||
temporal proximity, and proof count applied as multiplicative boosts. This
|
||||
ensures the influence of these secondary signals is always proportional to
|
||||
the base relevance score, regardless of the cross-encoder model's score
|
||||
calibration.
|
||||
|
||||
Formula::
|
||||
|
||||
recency_boost = 1 + recency_alpha * (recency - 0.5) # in [1-α/2, 1+α/2]
|
||||
temporal_boost = 1 + temporal_alpha * (temporal - 0.5) # in [1-α/2, 1+α/2]
|
||||
combined_score = cross_encoder_score_normalized * recency_boost * temporal_boost
|
||||
recency_boost = 1 + recency_alpha * (recency - 0.5) # in [1-α/2, 1+α/2]
|
||||
temporal_boost = 1 + temporal_alpha * (temporal - 0.5) # in [1-α/2, 1+α/2]
|
||||
proof_count_boost = 1 + proof_count_alpha * (proof_norm - 0.5) # in [1-α/2, 1+α/2]
|
||||
combined_score = CE_normalized * recency_boost * temporal_boost * proof_count_boost
|
||||
|
||||
proof_norm maps proof_count using a smooth logarithmic curve centered at 0.5,
|
||||
clamped to [0, 1]:
|
||||
proof_count=1 → 0.5 + 0 = 0.5 (neutral multiplier)
|
||||
proof_count=150 → clamped to 1.0 (max +5% boost)
|
||||
|
||||
Temporal proximity is treated as neutral (0.5) when not set by temporal retrieval,
|
||||
so temporal_boost collapses to 1.0 for non-temporal queries.
|
||||
|
||||
Proof count is treated as neutral (0.5) when not available (non-observation facts),
|
||||
so proof_count_boost collapses to 1.0 for world/experience/opinion facts.
|
||||
|
||||
Args:
|
||||
scored_results: Results from the cross-encoder reranker. Mutated in place.
|
||||
now: Current UTC datetime for recency calculation.
|
||||
recency_alpha: Max relative recency adjustment (default 0.2 → ±10%).
|
||||
temporal_alpha: Max relative temporal adjustment (default 0.2 → ±10%).
|
||||
proof_count_alpha: Max relative proof count adjustment (default 0.1 → ±5%).
|
||||
"""
|
||||
if now.tzinfo is None:
|
||||
now = now.replace(tzinfo=UTC)
|
||||
|
||||
# When the configured cross-encoder is a passthrough (e.g.
|
||||
# RRFPassthroughCrossEncoder used by slim deployments), every
|
||||
# cross_encoder_score_normalized is identical and provides no relevance
|
||||
# signal. In that case the multiplicative recency / temporal / proof_count
|
||||
# boosts below become the *only* ranking signal — making the final order a
|
||||
# pure recency sort regardless of how relevant a candidate actually is.
|
||||
#
|
||||
# Detect that case and seed cross_encoder_score_normalized from the RRF
|
||||
# rank instead, so the boosts modulate a meaningful base score rather than
|
||||
# replacing it. This is a no-op for real cross-encoders, which produce
|
||||
# diverse scores.
|
||||
# When the reranker is a passthrough (e.g. RRFPassthroughCrossEncoder used
|
||||
# by slim deployments), every cross_encoder_score_normalized is identical
|
||||
# and provides no relevance signal. The multiplicative recency / temporal /
|
||||
# proof_count boosts below would then become the *only* ranking signal,
|
||||
# making the final order a pure recency sort regardless of how relevant a
|
||||
# candidate actually is.
|
||||
#
|
||||
# Seed cross_encoder_score_normalized from the RRF rank instead, so the
|
||||
# boosts modulate a meaningful base score. Caller passes is_passthrough
|
||||
# explicitly because "all scores identical" is too fragile a heuristic —
|
||||
# a real reranker can also tie scores (especially in tests with synthetic
|
||||
# data) and we'd corrupt legitimate single-result reranks.
|
||||
if is_passthrough_reranker and scored_results:
|
||||
n = len(scored_results)
|
||||
sorted_by_rrf = sorted(
|
||||
scored_results,
|
||||
key=lambda s: getattr(getattr(s, "candidate", None), "rrf_score", 0.0),
|
||||
reverse=True,
|
||||
)
|
||||
denom = max(1, n - 1)
|
||||
for new_rank, sr in enumerate(sorted_by_rrf):
|
||||
# Map rank → [0.1, 1.0] so the recency boost can still nudge
|
||||
# ordering between adjacent candidates without overpowering RRF.
|
||||
sr.cross_encoder_score_normalized = 1.0 - (0.9 * new_rank / denom)
|
||||
|
||||
for sr in scored_results:
|
||||
# Recency: linear decay over 365 days → [0.1, 1.0]; neutral 0.5 if no date.
|
||||
sr.recency = 0.5
|
||||
@@ -59,13 +110,23 @@ def apply_combined_scoring(
|
||||
# Temporal proximity: meaningful only for temporal queries; neutral otherwise.
|
||||
sr.temporal = sr.retrieval.temporal_proximity if sr.retrieval.temporal_proximity is not None else 0.5
|
||||
|
||||
# Proof count: log-normalized evidence strength; neutral for non-observations.
|
||||
proof_count = sr.retrieval.proof_count
|
||||
if proof_count is not None and proof_count >= 1:
|
||||
# Clamp to [0, 1] so extreme counts stay within documented ±5% range
|
||||
proof_norm = min(1.0, max(0.0, 0.5 + (math.log(proof_count) / 10.0)))
|
||||
else:
|
||||
# Neutral baseline is precisely 0.5, ensuring neutral multiplier (1.0)
|
||||
proof_norm = 0.5
|
||||
|
||||
# RRF: kept at 0.0 for trace continuity but excluded from scoring.
|
||||
# RRF is batch-relative (min-max normalised) and redundant after reranking.
|
||||
sr.rrf_normalized = 0.0
|
||||
|
||||
recency_boost = 1.0 + recency_alpha * (sr.recency - 0.5)
|
||||
temporal_boost = 1.0 + temporal_alpha * (sr.temporal - 0.5)
|
||||
sr.combined_score = sr.cross_encoder_score_normalized * recency_boost * temporal_boost
|
||||
proof_count_boost = 1.0 + proof_count_alpha * (proof_norm - 0.5)
|
||||
sr.combined_score = sr.cross_encoder_score_normalized * recency_boost * temporal_boost * proof_count_boost
|
||||
sr.weight = sr.combined_score
|
||||
|
||||
|
||||
|
||||
@@ -13,7 +13,7 @@ import logging
|
||||
import re
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import UTC, datetime
|
||||
from typing import Optional
|
||||
from typing import Any, Optional
|
||||
|
||||
from ...config import get_config
|
||||
from ..db_utils import acquire_with_retry
|
||||
@@ -98,6 +98,8 @@ async def retrieve_semantic_bm25_combined(
|
||||
tags: list[str] | None = None,
|
||||
tags_match: TagsMatch = "any",
|
||||
tag_groups: list[TagGroup] | None = None,
|
||||
created_after: datetime | None = None,
|
||||
created_before: datetime | None = None,
|
||||
) -> dict[str, tuple[list[RetrievalResult], list[RetrievalResult]]]:
|
||||
"""
|
||||
Combined semantic + BM25 retrieval for multiple fact types in a single query.
|
||||
@@ -141,7 +143,7 @@ async def retrieve_semantic_bm25_combined(
|
||||
|
||||
cols = (
|
||||
"id, text, context, event_date, occurred_start, occurred_end, mentioned_at, "
|
||||
"fact_type, document_id, chunk_id, tags, metadata"
|
||||
"fact_type, document_id, chunk_id, tags, metadata, proof_count"
|
||||
)
|
||||
table = fq_table("memory_units")
|
||||
|
||||
@@ -163,6 +165,21 @@ async def retrieve_semantic_bm25_combined(
|
||||
tag_groups_param_start = tags_param_idx + (1 if tags else 0)
|
||||
groups_clause, groups_params, _ = build_tag_groups_where_clause(tag_groups, tag_groups_param_start)
|
||||
|
||||
# --- created_at time range filter (appended after tags/groups) ---
|
||||
# Param indices are computed relative to the final params list built below,
|
||||
# so we pre-compute the next available index after all preceding params.
|
||||
_next_idx = tag_groups_param_start + len(groups_params)
|
||||
created_range_clause = ""
|
||||
created_range_params: list[Any] = []
|
||||
if created_after is not None:
|
||||
created_range_params.append(created_after)
|
||||
created_range_clause += f" AND updated_at > ${_next_idx}"
|
||||
_next_idx += 1
|
||||
if created_before is not None:
|
||||
created_range_params.append(created_before)
|
||||
created_range_clause += f" AND updated_at < ${_next_idx}"
|
||||
_next_idx += 1
|
||||
|
||||
# --- Semantic UNION ALL arms (one per fact_type) ---
|
||||
# Each arm has its own ORDER BY embedding <=> $1 LIMIT {hnsw_fetch}, which
|
||||
# lets the planner use the partial HNSW index for that fact_type.
|
||||
@@ -180,6 +197,7 @@ async def retrieve_semantic_bm25_combined(
|
||||
f" AND (1 - (embedding <=> $1::vector)) >= 0.3"
|
||||
f" {tags_clause}"
|
||||
f" {groups_clause}"
|
||||
f" {created_range_clause}"
|
||||
f" ORDER BY embedding <=> $1::vector"
|
||||
f" LIMIT {hnsw_fetch})"
|
||||
)
|
||||
@@ -220,6 +238,7 @@ async def retrieve_semantic_bm25_combined(
|
||||
f" {bm25_where_filter}"
|
||||
f" {tags_clause}"
|
||||
f" {groups_clause}"
|
||||
f" {created_range_clause}"
|
||||
f" ORDER BY {bm25_order_by}"
|
||||
f" LIMIT $3)"
|
||||
)
|
||||
@@ -233,6 +252,7 @@ async def retrieve_semantic_bm25_combined(
|
||||
if tags:
|
||||
params.append(tags)
|
||||
params.extend(groups_params)
|
||||
params.extend(created_range_params)
|
||||
|
||||
rows = await conn.fetch(query, *params)
|
||||
|
||||
@@ -266,6 +286,8 @@ async def retrieve_temporal_combined(
|
||||
tags: list[str] | None = None,
|
||||
tags_match: TagsMatch = "any",
|
||||
tag_groups: list[TagGroup] | None = None,
|
||||
created_after: datetime | None = None,
|
||||
created_before: datetime | None = None,
|
||||
) -> dict[str, list[RetrievalResult]]:
|
||||
"""
|
||||
Temporal retrieval for multiple fact types in a single query.
|
||||
@@ -299,10 +321,25 @@ async def retrieve_temporal_combined(
|
||||
tags_clause = build_tags_where_clause_simple(tags, 7, match=tags_match)
|
||||
tag_groups_param_start = 7 + (1 if tags else 0)
|
||||
groups_clause, groups_params, _ = build_tag_groups_where_clause(tag_groups, tag_groups_param_start)
|
||||
|
||||
# created_at time range filter (after tags/groups)
|
||||
_next_idx = tag_groups_param_start + len(groups_params)
|
||||
created_range_clause = ""
|
||||
created_range_params: list[Any] = []
|
||||
if created_after is not None:
|
||||
created_range_params.append(created_after)
|
||||
created_range_clause += f" AND updated_at > ${_next_idx}"
|
||||
_next_idx += 1
|
||||
if created_before is not None:
|
||||
created_range_params.append(created_before)
|
||||
created_range_clause += f" AND updated_at < ${_next_idx}"
|
||||
_next_idx += 1
|
||||
|
||||
params: list = [query_emb_str, bank_id, fact_types, start_date, end_date, semantic_threshold]
|
||||
if tags:
|
||||
params.append(tags)
|
||||
params.extend(groups_params)
|
||||
params.extend(created_range_params)
|
||||
|
||||
# Two-phase entry point query:
|
||||
# Phase 1 (date_ranked): rank by date only — no embedding computation — for all units in
|
||||
@@ -334,9 +371,10 @@ async def retrieve_temporal_combined(
|
||||
)
|
||||
{tags_clause}
|
||||
{groups_clause}
|
||||
{created_range_clause}
|
||||
),
|
||||
sim_ranked AS (
|
||||
SELECT mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start, mu.occurred_end, mu.mentioned_at, mu.fact_type, mu.document_id, mu.chunk_id, mu.tags, mu.metadata,
|
||||
SELECT mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start, mu.occurred_end, mu.mentioned_at, mu.fact_type, mu.proof_count, mu.document_id, mu.chunk_id, mu.tags, mu.metadata,
|
||||
1 - (mu.embedding <=> $1::vector) AS similarity,
|
||||
ROW_NUMBER() OVER (PARTITION BY mu.fact_type ORDER BY mu.embedding <=> $1::vector) AS sim_rn
|
||||
FROM date_ranked dr
|
||||
@@ -344,7 +382,7 @@ async def retrieve_temporal_combined(
|
||||
WHERE dr.rn <= 50
|
||||
AND (1 - (mu.embedding <=> $1::vector)) >= $6
|
||||
)
|
||||
SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, fact_type, document_id, chunk_id, tags, metadata, similarity
|
||||
SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, fact_type, proof_count, document_id, chunk_id, tags, metadata, similarity
|
||||
FROM sim_ranked
|
||||
WHERE sim_rn <= 10
|
||||
""",
|
||||
@@ -536,6 +574,8 @@ async def retrieve_all_fact_types_parallel(
|
||||
tags: list[str] | None = None,
|
||||
tags_match: TagsMatch = "any",
|
||||
tag_groups: list[TagGroup] | None = None,
|
||||
created_after: datetime | None = None,
|
||||
created_before: datetime | None = None,
|
||||
) -> MultiFactTypeRetrievalResult:
|
||||
"""
|
||||
Optimized retrieval for multiple fact types using batched queries.
|
||||
@@ -594,6 +634,8 @@ async def retrieve_all_fact_types_parallel(
|
||||
tags=tags,
|
||||
tags_match=tags_match,
|
||||
tag_groups=tag_groups,
|
||||
created_after=created_after,
|
||||
created_before=created_before,
|
||||
)
|
||||
semantic_bm25_time = time.time() - semantic_bm25_start
|
||||
|
||||
@@ -613,6 +655,8 @@ async def retrieve_all_fact_types_parallel(
|
||||
tags=tags,
|
||||
tags_match=tags_match,
|
||||
tag_groups=tag_groups,
|
||||
created_after=created_after,
|
||||
created_before=created_before,
|
||||
)
|
||||
temporal_time = time.time() - temporal_start
|
||||
|
||||
@@ -636,6 +680,8 @@ async def retrieve_all_fact_types_parallel(
|
||||
tags=tags,
|
||||
tags_match=tags_match,
|
||||
tag_groups=tag_groups,
|
||||
created_after=created_after,
|
||||
created_before=created_before,
|
||||
)
|
||||
return ft, results, time.time() - graph_start, graph_timing
|
||||
|
||||
|
||||
@@ -62,17 +62,18 @@ def format_facts_for_prompt(facts: list[MemoryFact]) -> str:
|
||||
if fact.context:
|
||||
fact_obj["context"] = fact.context
|
||||
|
||||
# Add occurred_start if available (when the fact occurred)
|
||||
if fact.occurred_start:
|
||||
occurred_start = fact.occurred_start
|
||||
if isinstance(occurred_start, str):
|
||||
fact_obj["occurred_start"] = occurred_start
|
||||
elif isinstance(occurred_start, datetime):
|
||||
fact_obj["occurred_start"] = occurred_start.strftime("%Y-%m-%d %H:%M:%S")
|
||||
# Add temporal fields if available
|
||||
for field_name in ("occurred_start", "occurred_end", "mentioned_at"):
|
||||
value = getattr(fact, field_name, None)
|
||||
if value:
|
||||
if isinstance(value, str):
|
||||
fact_obj[field_name] = value
|
||||
elif isinstance(value, datetime):
|
||||
fact_obj[field_name] = value.strftime("%Y-%m-%d %H:%M:%S")
|
||||
|
||||
formatted.append(fact_obj)
|
||||
|
||||
return json.dumps(formatted, indent=2)
|
||||
return json.dumps(formatted, indent=2, ensure_ascii=False)
|
||||
|
||||
|
||||
def format_entity_summaries_for_prompt(entities: dict) -> str:
|
||||
|
||||
@@ -48,6 +48,7 @@ class RetrievalResult:
|
||||
chunk_id: str | None = None
|
||||
tags: list[str] | None = None # Visibility scope tags
|
||||
metadata: dict[str, str] | None = None # User-provided metadata
|
||||
proof_count: int | None = None # Number of supporting memories (observations only)
|
||||
|
||||
# Retrieval-specific scores (only one will be set depending on retrieval method)
|
||||
similarity: float | None = None # Semantic retrieval
|
||||
@@ -72,6 +73,7 @@ class RetrievalResult:
|
||||
chunk_id=row.get("chunk_id"),
|
||||
tags=row.get("tags"),
|
||||
metadata=row.get("metadata"),
|
||||
proof_count=row.get("proof_count"),
|
||||
similarity=row.get("similarity"),
|
||||
bm25_score=row.get("bm25_score"),
|
||||
activation=row.get("activation"),
|
||||
|
||||
@@ -2,7 +2,8 @@
|
||||
Task backend for distributed task processing.
|
||||
|
||||
This provides an abstraction for task storage and execution:
|
||||
- BrokerTaskBackend: Uses PostgreSQL as broker (production)
|
||||
- BrokerTaskBackend: Uses PostgreSQL as broker (production API servers)
|
||||
- WorkerTaskBackend: No-op submit_task (production workers — child tasks are polled)
|
||||
- SyncTaskBackend: Executes tasks immediately (testing/embedded)
|
||||
"""
|
||||
|
||||
@@ -125,6 +126,33 @@ class SyncTaskBackend(TaskBackend):
|
||||
logger.debug("SyncTaskBackend shutdown")
|
||||
|
||||
|
||||
class WorkerTaskBackend(TaskBackend):
|
||||
"""
|
||||
Task backend for worker processes.
|
||||
|
||||
Workers execute tasks directly via the poller (claim → execute), so they
|
||||
don't need submit_task to run anything. When engine code running *inside*
|
||||
a worker-executed task calls submit_task (e.g. retain triggers consolidation),
|
||||
the row has already been INSERTed into async_operations with task_payload by
|
||||
_submit_async_operation — so submit_task is a no-op. The new task will be
|
||||
picked up by a worker on the next poll cycle instead of being executed inline,
|
||||
which avoids blocking the parent task.
|
||||
"""
|
||||
|
||||
async def initialize(self):
|
||||
self._initialized = True
|
||||
logger.debug("WorkerTaskBackend initialized")
|
||||
|
||||
async def submit_task(self, task_dict: dict[str, Any]):
|
||||
"""No-op: the row already exists in async_operations; a worker will claim it."""
|
||||
task_type = task_dict.get("type", "unknown")
|
||||
logger.debug(f"WorkerTaskBackend: submit_task no-op for {task_type} (will be picked up by poller)")
|
||||
|
||||
async def shutdown(self):
|
||||
self._initialized = False
|
||||
logger.debug("WorkerTaskBackend shutdown")
|
||||
|
||||
|
||||
class BrokerTaskBackend(TaskBackend):
|
||||
"""
|
||||
Task backend using PostgreSQL as broker.
|
||||
@@ -193,17 +221,21 @@ class BrokerTaskBackend(TaskBackend):
|
||||
table = fq_table("async_operations", schema)
|
||||
|
||||
if operation_id:
|
||||
# Update existing operation with task payload
|
||||
# Callers now include task_payload in the same INSERT that creates the
|
||||
# async_operations row (see MemoryEngine._submit_async_operation). The
|
||||
# WHERE clause guards against overwriting that payload — the UPDATE is a
|
||||
# no-op when the row is already claimable, and only fills in a NULL payload
|
||||
# for any legacy caller that still creates the row first.
|
||||
await pool.execute(
|
||||
f"""
|
||||
UPDATE {table}
|
||||
SET task_payload = $1::jsonb, updated_at = now()
|
||||
WHERE operation_id = $2
|
||||
WHERE operation_id = $2 AND task_payload IS NULL
|
||||
""",
|
||||
payload_json,
|
||||
operation_id,
|
||||
)
|
||||
logger.debug(f"Updated task payload for operation {operation_id}")
|
||||
logger.debug(f"submit_task UPDATE for operation {operation_id} (no-op if payload already set)")
|
||||
else:
|
||||
# Insert new operation (for tasks without pre-created records)
|
||||
# e.g., access_count_update tasks
|
||||
|
||||
@@ -55,6 +55,7 @@ from hindsight_api.extensions.tenant import (
|
||||
TenantExtension,
|
||||
)
|
||||
from hindsight_api.models import RequestContext
|
||||
from hindsight_api.worker.exceptions import DeferOperation
|
||||
|
||||
__all__ = [
|
||||
# Base
|
||||
@@ -68,6 +69,7 @@ __all__ = [
|
||||
# MCP Extension
|
||||
"MCPExtension",
|
||||
# Operation Validator - Core
|
||||
"DeferOperation",
|
||||
"OperationValidationError",
|
||||
"OperationValidatorExtension",
|
||||
"RecallContext",
|
||||
|
||||
@@ -176,6 +176,22 @@ class RetainResult:
|
||||
llm_input_tokens: int | None = None
|
||||
llm_output_tokens: int | None = None
|
||||
llm_total_tokens: int | None = None
|
||||
# Content tokens the retain pipeline actually processed, after
|
||||
# chunk-level content-hash deduplication. Semantics:
|
||||
# None — no dedup signal available (e.g. a first-time retain or a
|
||||
# path that doesn't compute it). Callers that care about
|
||||
# "what was actually new on this retain" should treat None
|
||||
# as "the full submitted content was processed."
|
||||
# 0 — the entire submission was a duplicate of prior content
|
||||
# (all chunks matched by content_hash); nothing went
|
||||
# through LLM extraction.
|
||||
# N>0 — only N tokens of content + context went through the
|
||||
# extraction pipeline. The remainder was dedup'd against
|
||||
# existing chunks.
|
||||
# This is the basis most billing/metering extensions want to use
|
||||
# when the customer's client resubmits growing payloads to the same
|
||||
# document_id (e.g. a session transcript appended to on each turn).
|
||||
processed_content_tokens: int | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -376,6 +392,16 @@ class OperationValidatorExtension(Extension, ABC):
|
||||
2. [operation executes]
|
||||
3. on_*_complete (post-operation)
|
||||
|
||||
Outcomes for `validate_*` hooks:
|
||||
- accept: return `ValidationResult.accept()` (or `accept_with(...)`)
|
||||
- reject: return `ValidationResult.reject(reason, status_code)`
|
||||
(raises `OperationValidationError` upstream)
|
||||
- defer: raise `DeferOperation(exec_date, reason)` from
|
||||
`hindsight_api.worker.exceptions` to requeue the task for a
|
||||
future time without bumping `retry_count`. Worker-only — do
|
||||
not raise from `validate_recall` / `validate_reflect` in
|
||||
synchronous HTTP request paths, where it surfaces as a 500.
|
||||
|
||||
Supported operations:
|
||||
- retain, recall, reflect (core memory operations)
|
||||
- consolidate (mental models consolidation)
|
||||
|
||||
@@ -29,6 +29,7 @@ from hindsight_api.models import RequestContext
|
||||
_ALL_TOOLS: frozenset[str] = frozenset(
|
||||
{
|
||||
"retain",
|
||||
"sync_retain",
|
||||
"recall",
|
||||
"reflect",
|
||||
"list_banks",
|
||||
@@ -139,6 +140,7 @@ def build_content_dict(
|
||||
metadata: dict[str, str] | None = None,
|
||||
document_id: str | None = None,
|
||||
strategy: str | None = None,
|
||||
update_mode: str | None = None,
|
||||
) -> tuple[dict[str, Any], str | None]:
|
||||
"""Build a content dict for retain operations.
|
||||
|
||||
@@ -150,6 +152,7 @@ def build_content_dict(
|
||||
metadata: Optional key-value metadata to attach to the memory
|
||||
document_id: Optional document ID to associate the memory with
|
||||
strategy: Optional named retain strategy override (e.g., 'exact', 'verbose')
|
||||
update_mode: How to handle existing documents ('replace' or 'append')
|
||||
|
||||
Returns:
|
||||
Tuple of (content_dict, error_message). error_message is None if successful.
|
||||
@@ -184,6 +187,8 @@ def build_content_dict(
|
||||
content_dict["document_id"] = document_id
|
||||
if strategy is not None:
|
||||
content_dict["strategy"] = strategy
|
||||
if update_mode is not None:
|
||||
content_dict["update_mode"] = update_mode
|
||||
|
||||
return content_dict, None
|
||||
|
||||
@@ -202,6 +207,7 @@ def register_mcp_tools(
|
||||
"""
|
||||
tools_to_register = config.tools or {
|
||||
"retain",
|
||||
"sync_retain",
|
||||
"recall",
|
||||
"reflect",
|
||||
"list_banks",
|
||||
@@ -235,6 +241,9 @@ def register_mcp_tools(
|
||||
if "retain" in tools_to_register:
|
||||
_register_retain(mcp, memory, config)
|
||||
|
||||
if "sync_retain" in tools_to_register:
|
||||
_register_sync_retain(mcp, memory, config)
|
||||
|
||||
if "recall" in tools_to_register:
|
||||
_register_recall(mcp, memory, config)
|
||||
|
||||
@@ -539,6 +548,7 @@ def _register_retain(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig)
|
||||
document_id: str | None = None,
|
||||
bank_id: str | None = None,
|
||||
strategy: str | None = None,
|
||||
update_mode: str | None = None,
|
||||
) -> dict:
|
||||
"""
|
||||
Args:
|
||||
@@ -550,12 +560,15 @@ def _register_retain(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig)
|
||||
document_id: Optional document ID to associate this memory with
|
||||
bank_id: Optional bank to store in (defaults to session bank). Use for cross-bank operations.
|
||||
strategy: Optional named retain strategy (e.g., 'exact' for verbatim storage). Strategies are defined in the bank config.
|
||||
update_mode: How to handle existing documents with the same document_id. 'replace' (default) or 'append' (concatenates new content to existing).
|
||||
"""
|
||||
target_bank = bank_id or config.bank_id_resolver()
|
||||
if target_bank is None:
|
||||
return {"status": "error", "message": "No bank_id configured"}
|
||||
|
||||
content_dict, error = build_content_dict(content, context, timestamp, tags, metadata, document_id, strategy)
|
||||
content_dict, error = build_content_dict(
|
||||
content, context, timestamp, tags, metadata, document_id, strategy, update_mode
|
||||
)
|
||||
if error:
|
||||
return {"status": "error", "message": error}
|
||||
|
||||
@@ -590,6 +603,7 @@ def _register_retain(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig)
|
||||
metadata: dict[str, str] | None = None,
|
||||
document_id: str | None = None,
|
||||
strategy: str | None = None,
|
||||
update_mode: str | None = None,
|
||||
) -> dict:
|
||||
"""
|
||||
Args:
|
||||
@@ -600,12 +614,15 @@ def _register_retain(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig)
|
||||
metadata: Optional key-value metadata to attach (e.g., {'source': 'slack', 'channel': 'general'})
|
||||
document_id: Optional document ID to associate this memory with
|
||||
strategy: Optional named retain strategy (e.g., 'exact' for verbatim storage). Strategies are defined in the bank config.
|
||||
update_mode: How to handle existing documents with the same document_id. 'replace' (default) or 'append' (concatenates new content to existing).
|
||||
"""
|
||||
target_bank = config.bank_id_resolver()
|
||||
if target_bank is None:
|
||||
return {"status": "error", "message": "No bank_id configured"}
|
||||
|
||||
content_dict, error = build_content_dict(content, context, timestamp, tags, metadata, document_id, strategy)
|
||||
content_dict, error = build_content_dict(
|
||||
content, context, timestamp, tags, metadata, document_id, strategy, update_mode
|
||||
)
|
||||
if error:
|
||||
return {"status": "error", "message": error}
|
||||
|
||||
@@ -630,6 +647,124 @@ def _register_retain(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig)
|
||||
return {"status": "error", "message": str(e)}
|
||||
|
||||
|
||||
def _register_sync_retain(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig) -> None:
|
||||
"""Register the sync_retain tool (synchronous retain that waits for completion)."""
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
async def sync_retain(
|
||||
content: str,
|
||||
context: str = "general",
|
||||
timestamp: str | None = None,
|
||||
tags: list[str] | None = None,
|
||||
metadata: dict[str, str] | None = None,
|
||||
document_id: str | None = None,
|
||||
bank_id: str | None = None,
|
||||
strategy: str | None = None,
|
||||
) -> dict:
|
||||
"""Store information to long-term memory and wait for completion.
|
||||
|
||||
Unlike retain (which is asynchronous), this tool blocks until the memory
|
||||
is fully stored and immediately available for recall.
|
||||
|
||||
Args:
|
||||
content: The fact/memory to store (be specific and include relevant details)
|
||||
context: Category for the memory (e.g., 'preferences', 'work', 'hobbies', 'family'). Default: 'general'
|
||||
timestamp: When this event/fact occurred (ISO format, e.g., '2024-01-15T10:30:00Z'). Useful for timeline tracking.
|
||||
tags: Optional tags for scoped visibility filtering (e.g., ['project:alpha', 'user:123'])
|
||||
metadata: Optional key-value metadata to attach (e.g., {'source': 'slack', 'channel': 'general'})
|
||||
document_id: Optional document ID to associate this memory with
|
||||
bank_id: Optional bank to store in (defaults to session bank). Use for cross-bank operations.
|
||||
strategy: Optional named retain strategy (e.g., 'exact' for verbatim storage). Strategies are defined in the bank config.
|
||||
"""
|
||||
target_bank = bank_id or config.bank_id_resolver()
|
||||
if target_bank is None:
|
||||
return {"status": "error", "message": "No bank_id configured"}
|
||||
|
||||
content_dict, error = build_content_dict(content, context, timestamp, tags, metadata, document_id, strategy)
|
||||
if error:
|
||||
return {"status": "error", "message": error}
|
||||
|
||||
request_context = _get_request_context(config)
|
||||
|
||||
try:
|
||||
result = await memory.retain_batch_async(
|
||||
bank_id=target_bank,
|
||||
contents=[content_dict],
|
||||
request_context=request_context,
|
||||
strategy=content_dict.pop("strategy", None),
|
||||
)
|
||||
memory_ids = [uid for batch in result for uid in batch]
|
||||
return {
|
||||
"status": "completed",
|
||||
"message": "Memory stored successfully",
|
||||
"memory_ids": memory_ids,
|
||||
}
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Sync retain rejected: {e}")
|
||||
return {"status": "error", "message": str(e)}
|
||||
except Exception as e:
|
||||
logger.error(f"Error in sync retain: {e}", exc_info=True)
|
||||
return {"status": "error", "message": str(e)}
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
async def sync_retain(
|
||||
content: str,
|
||||
context: str = "general",
|
||||
timestamp: str | None = None,
|
||||
tags: list[str] | None = None,
|
||||
metadata: dict[str, str] | None = None,
|
||||
document_id: str | None = None,
|
||||
strategy: str | None = None,
|
||||
) -> dict:
|
||||
"""Store information to long-term memory and wait for completion.
|
||||
|
||||
Unlike retain (which is asynchronous), this tool blocks until the memory
|
||||
is fully stored and immediately available for recall.
|
||||
|
||||
Args:
|
||||
content: The fact/memory to store (be specific and include relevant details)
|
||||
context: Category for the memory (e.g., 'preferences', 'work', 'hobbies', 'family'). Default: 'general'
|
||||
timestamp: When this event/fact occurred (ISO format, e.g., '2024-01-15T10:30:00Z'). Useful for timeline tracking.
|
||||
tags: Optional tags for scoped visibility filtering (e.g., ['project:alpha', 'user:123'])
|
||||
metadata: Optional key-value metadata to attach (e.g., {'source': 'slack', 'channel': 'general'})
|
||||
document_id: Optional document ID to associate this memory with
|
||||
strategy: Optional named retain strategy (e.g., 'exact' for verbatim storage). Strategies are defined in the bank config.
|
||||
"""
|
||||
target_bank = config.bank_id_resolver()
|
||||
if target_bank is None:
|
||||
return {"status": "error", "message": "No bank_id configured"}
|
||||
|
||||
content_dict, error = build_content_dict(content, context, timestamp, tags, metadata, document_id, strategy)
|
||||
if error:
|
||||
return {"status": "error", "message": error}
|
||||
|
||||
request_context = _get_request_context(config)
|
||||
|
||||
try:
|
||||
result = await memory.retain_batch_async(
|
||||
bank_id=target_bank,
|
||||
contents=[content_dict],
|
||||
request_context=request_context,
|
||||
strategy=content_dict.pop("strategy", None),
|
||||
)
|
||||
memory_ids = [uid for batch in result for uid in batch]
|
||||
return {
|
||||
"status": "completed",
|
||||
"message": "Memory stored successfully",
|
||||
"memory_ids": memory_ids,
|
||||
}
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Sync retain rejected: {e}")
|
||||
return {"status": "error", "message": str(e)}
|
||||
except Exception as e:
|
||||
logger.error(f"Error in sync retain: {e}", exc_info=True)
|
||||
return {"status": "error", "message": str(e)}
|
||||
|
||||
|
||||
def _register_recall(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig) -> None:
|
||||
"""Register the recall tool."""
|
||||
description = config.recall_description or DEFAULT_MCP_RECALL_DESCRIPTION
|
||||
@@ -2719,6 +2854,44 @@ def _register_get_bank_stats(mcp: FastMCP, memory: MemoryEngine, config: MCPTool
|
||||
return f'{{"error": "{e}"}}'
|
||||
|
||||
|
||||
async def _do_update_bank(
|
||||
memory: MemoryEngine,
|
||||
target_bank: str,
|
||||
request_context: RequestContext,
|
||||
*,
|
||||
name: str | None = None,
|
||||
mission: str | None = None,
|
||||
config_updates: dict[str, Any] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Shared implementation for update_bank MCP tool variants.
|
||||
|
||||
Args:
|
||||
name: Display name (stored in banks table).
|
||||
mission: Deprecated alias for reflect_mission — mapped into config_updates.
|
||||
config_updates: Arbitrary config overrides passed to config_resolver.update_bank_config().
|
||||
Supports all configurable fields (retain_mission, disposition_*, etc.).
|
||||
The config resolver validates keys and rejects non-configurable/credential fields.
|
||||
"""
|
||||
# Update display name via engine (stored in DB banks table)
|
||||
if name is not None:
|
||||
await memory.update_bank(
|
||||
target_bank,
|
||||
name=name,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Merge deprecated mission alias into config_updates as reflect_mission
|
||||
effective_config: dict[str, Any] = dict(config_updates) if config_updates else {}
|
||||
if mission is not None and "reflect_mission" not in effective_config:
|
||||
effective_config["reflect_mission"] = mission
|
||||
|
||||
if effective_config:
|
||||
await memory._config_resolver.update_bank_config(target_bank, effective_config, request_context)
|
||||
|
||||
# Return updated profile
|
||||
return await memory.get_bank_profile(target_bank, request_context=request_context)
|
||||
|
||||
|
||||
def _register_update_bank(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig) -> None:
|
||||
"""Register the update_bank tool."""
|
||||
|
||||
@@ -2728,16 +2901,37 @@ def _register_update_bank(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsCo
|
||||
async def update_bank(
|
||||
name: str | None = None,
|
||||
mission: str | None = None,
|
||||
config_updates: dict[str, Any] | None = None,
|
||||
bank_id: str | None = None,
|
||||
) -> str:
|
||||
"""
|
||||
Update a memory bank's metadata.
|
||||
Update a memory bank's configuration.
|
||||
|
||||
Changes the name or mission of an existing bank.
|
||||
Updates the bank's name and/or any bank-level configuration fields.
|
||||
Only provided fields will be updated; omitted fields remain unchanged.
|
||||
|
||||
Args:
|
||||
name: New human-friendly name for the bank
|
||||
mission: New mission describing who the agent is and what they're trying to accomplish
|
||||
name: Human-friendly display name for the bank.
|
||||
mission: Deprecated alias for config_updates.reflect_mission.
|
||||
config_updates: Dictionary of configuration fields to update. Supports all
|
||||
bank-configurable fields including:
|
||||
- reflect_mission: Mission/context for Reflect operations.
|
||||
- retain_mission: Steers what gets extracted during retain().
|
||||
- retain_extraction_mode: 'concise' (default), 'verbose', or 'custom'.
|
||||
- retain_custom_instructions: Custom extraction prompt (active when mode is 'custom').
|
||||
- retain_chunk_size: Maximum token size for each content chunk.
|
||||
- retain_chunk_batch_size: Number of chunks to process in parallel.
|
||||
- enable_observations: Toggle observation consolidation after retain().
|
||||
- observations_mission: Controls observation synthesis rules.
|
||||
- disposition_skepticism: Critical evaluation level (1-5).
|
||||
- disposition_literalism: Literal vs. abstract interpretation (1-5).
|
||||
- disposition_empathy: Emotional context consideration (1-5).
|
||||
- entity_labels: Controlled vocabulary for entity classification.
|
||||
- entities_allow_free_form: Allow labels outside entity_labels.
|
||||
- recall_include_chunks: Include raw chunks in recall results.
|
||||
- recall_max_tokens: Max tokens for recall results.
|
||||
- mcp_enabled_tools: Tool allowlist for this bank.
|
||||
Any configurable field name is accepted (use Python field names).
|
||||
bank_id: Optional bank (defaults to session bank). Use for cross-bank operations.
|
||||
"""
|
||||
try:
|
||||
@@ -2745,14 +2939,16 @@ def _register_update_bank(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsCo
|
||||
if target_bank is None:
|
||||
return '{"error": "No bank_id configured"}'
|
||||
|
||||
result = await memory.update_bank(
|
||||
result = await _do_update_bank(
|
||||
memory,
|
||||
target_bank,
|
||||
_get_request_context(config),
|
||||
name=name,
|
||||
mission=mission,
|
||||
request_context=_get_request_context(config),
|
||||
config_updates=config_updates,
|
||||
)
|
||||
return json.dumps(result, indent=2, default=str)
|
||||
except OperationValidationError as e:
|
||||
except (OperationValidationError, ValueError) as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return json.dumps({"error": str(e)})
|
||||
except Exception as e:
|
||||
@@ -2765,29 +2961,52 @@ def _register_update_bank(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsCo
|
||||
async def update_bank(
|
||||
name: str | None = None,
|
||||
mission: str | None = None,
|
||||
config_updates: dict[str, Any] | None = None,
|
||||
) -> dict:
|
||||
"""
|
||||
Update this memory bank's metadata.
|
||||
Update this memory bank's configuration.
|
||||
|
||||
Changes the name or mission of the bank.
|
||||
Updates the bank's name and/or any bank-level configuration fields.
|
||||
Only provided fields will be updated; omitted fields remain unchanged.
|
||||
|
||||
Args:
|
||||
name: New human-friendly name for the bank
|
||||
mission: New mission describing who the agent is and what they're trying to accomplish
|
||||
name: Human-friendly display name for the bank.
|
||||
mission: Deprecated alias for config_updates.reflect_mission.
|
||||
config_updates: Dictionary of configuration fields to update. Supports all
|
||||
bank-configurable fields including:
|
||||
- reflect_mission: Mission/context for Reflect operations.
|
||||
- retain_mission: Steers what gets extracted during retain().
|
||||
- retain_extraction_mode: 'concise' (default), 'verbose', or 'custom'.
|
||||
- retain_custom_instructions: Custom extraction prompt (active when mode is 'custom').
|
||||
- retain_chunk_size: Maximum token size for each content chunk.
|
||||
- retain_chunk_batch_size: Number of chunks to process in parallel.
|
||||
- enable_observations: Toggle observation consolidation after retain().
|
||||
- observations_mission: Controls observation synthesis rules.
|
||||
- disposition_skepticism: Critical evaluation level (1-5).
|
||||
- disposition_literalism: Literal vs. abstract interpretation (1-5).
|
||||
- disposition_empathy: Emotional context consideration (1-5).
|
||||
- entity_labels: Controlled vocabulary for entity classification.
|
||||
- entities_allow_free_form: Allow labels outside entity_labels.
|
||||
- recall_include_chunks: Include raw chunks in recall results.
|
||||
- recall_max_tokens: Max tokens for recall results.
|
||||
- mcp_enabled_tools: Tool allowlist for this bank.
|
||||
Any configurable field name is accepted (use Python field names).
|
||||
"""
|
||||
try:
|
||||
target_bank = config.bank_id_resolver()
|
||||
if target_bank is None:
|
||||
return {"error": "No bank_id configured"}
|
||||
|
||||
result = await memory.update_bank(
|
||||
result = await _do_update_bank(
|
||||
memory,
|
||||
target_bank,
|
||||
_get_request_context(config),
|
||||
name=name,
|
||||
mission=mission,
|
||||
request_context=_get_request_context(config),
|
||||
config_updates=config_updates,
|
||||
)
|
||||
return result
|
||||
except OperationValidationError as e:
|
||||
except (OperationValidationError, ValueError) as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return {"error": str(e)}
|
||||
except Exception as e:
|
||||
|
||||
@@ -252,6 +252,9 @@ class MetricsCollector(MetricsCollectorBase):
|
||||
|
||||
def __init__(self):
|
||||
self.meter = get_meter()
|
||||
from .config import get_config
|
||||
|
||||
self._include_bank_id = get_config().metrics_include_bank_id
|
||||
|
||||
# Operation latency histogram (in seconds)
|
||||
# Records duration of retain, recall, reflect operations
|
||||
@@ -332,10 +335,11 @@ class MetricsCollector(MetricsCollectorBase):
|
||||
start_time = time.time()
|
||||
attributes = {
|
||||
"operation": operation,
|
||||
"bank_id": bank_id,
|
||||
"source": source,
|
||||
"tenant": _get_tenant(),
|
||||
}
|
||||
if self._include_bank_id:
|
||||
attributes["bank_id"] = bank_id
|
||||
if budget:
|
||||
attributes["budget"] = budget
|
||||
if max_tokens:
|
||||
|
||||
@@ -27,6 +27,7 @@ from alembic.config import Config
|
||||
from alembic.script.revision import ResolutionError
|
||||
from sqlalchemy import Connection, create_engine, text
|
||||
|
||||
from .db_url import to_libpq_url
|
||||
from .utils import mask_network_location
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -220,7 +221,7 @@ def run_migrations(
|
||||
# ineffective when the app URL goes through a pooler. Configure
|
||||
# HINDSIGHT_API_MIGRATION_DATABASE_URL to the direct PostgreSQL endpoint
|
||||
# (e.g. hindsight-pg-rw) to restore correct locking behaviour.
|
||||
migration_url = migration_database_url or database_url
|
||||
migration_url = to_libpq_url(migration_database_url or database_url)
|
||||
|
||||
try:
|
||||
# Determine script location
|
||||
@@ -450,7 +451,7 @@ def check_migration_status(
|
||||
return None, None
|
||||
|
||||
# Get current revision from database
|
||||
engine = create_engine(database_url)
|
||||
engine = create_engine(to_libpq_url(database_url))
|
||||
with engine.connect() as connection:
|
||||
context = MigrationContext.configure(connection)
|
||||
current_rev = context.get_current_revision()
|
||||
@@ -624,7 +625,7 @@ def ensure_embedding_dimension(
|
||||
"""
|
||||
schema_name = schema or "public"
|
||||
|
||||
engine = create_engine(database_url)
|
||||
engine = create_engine(to_libpq_url(database_url))
|
||||
with engine.connect() as conn:
|
||||
# Check if memory_units table exists (proxy for schema being initialized)
|
||||
table_exists = conn.execute(
|
||||
@@ -673,7 +674,7 @@ def ensure_vector_extension(
|
||||
"""
|
||||
schema_name = schema or "public"
|
||||
|
||||
engine = create_engine(database_url)
|
||||
engine = create_engine(to_libpq_url(database_url))
|
||||
with engine.connect() as conn:
|
||||
# Detect which vector extension should be used
|
||||
target_ext = _detect_vector_extension(conn, vector_extension)
|
||||
@@ -894,7 +895,7 @@ def ensure_text_search_extension(
|
||||
"""
|
||||
schema_name = schema or "public"
|
||||
|
||||
engine = create_engine(database_url)
|
||||
engine = create_engine(to_libpq_url(database_url))
|
||||
with engine.connect() as conn:
|
||||
# Tables with search_vector columns to check
|
||||
tables_to_check = [
|
||||
|
||||
@@ -24,6 +24,12 @@ class RequestContext:
|
||||
mcp_authenticated: bool = False # True when MCP transport auth already validated (skips tenant re-auth)
|
||||
user_initiated: bool = False # True for async operations that originated from a user request
|
||||
allowed_bank_ids: list[str] | None = None # None = unrestricted (all banks)
|
||||
# Number of times this task has been retried. Populated by the worker
|
||||
# from async_operations.retry_count before dispatching to a task handler;
|
||||
# 0 for sync/HTTP requests and for the first worker attempt. Useful for
|
||||
# validators that want exponential backoff on repeated failures (e.g.
|
||||
# "defer for 2^retry_count minutes") without querying the DB themselves.
|
||||
retry_count: int = 0
|
||||
|
||||
|
||||
from pgvector.sqlalchemy import Vector
|
||||
@@ -62,7 +68,6 @@ class Document(Base):
|
||||
bank_id: Mapped[str] = mapped_column(Text, primary_key=True)
|
||||
original_text: Mapped[str | None] = mapped_column(Text)
|
||||
content_hash: Mapped[str | None] = mapped_column(Text)
|
||||
doc_metadata: Mapped[dict] = mapped_column("metadata", JSONB, server_default=sql_text("'{}'::jsonb"))
|
||||
created_at: Mapped[datetime] = mapped_column(TIMESTAMP(timezone=True), server_default=func.now())
|
||||
updated_at: Mapped[datetime] = mapped_column(TIMESTAMP(timezone=True), server_default=func.now())
|
||||
|
||||
|
||||
@@ -7,3 +7,24 @@ class RetryTaskAt(Exception):
|
||||
def __init__(self, retry_at: datetime, message: str = ""):
|
||||
self.retry_at = retry_at
|
||||
super().__init__(message)
|
||||
|
||||
|
||||
class DeferOperation(Exception):
|
||||
"""Raise from an extension hook (or task handler) to requeue the
|
||||
operation for execution at a later time, without counting as a retry.
|
||||
|
||||
Unlike `RetryTaskAt`, this is not a failure: `retry_count` is not
|
||||
incremented and `error_message` is not written. Use this for
|
||||
backpressure / "not yet, try later" decisions made before or during
|
||||
task execution (e.g. quota windows, warming dependencies, upstream
|
||||
rate limits).
|
||||
|
||||
Worker-only: raising this from a hook called in HTTP request context
|
||||
(e.g. `validate_recall` for a synchronous recall) will surface as an
|
||||
unhandled 500 — there is no queue to defer to.
|
||||
"""
|
||||
|
||||
def __init__(self, exec_date: datetime, reason: str = ""):
|
||||
self.exec_date = exec_date
|
||||
self.reason = reason
|
||||
super().__init__(reason)
|
||||
|
||||
@@ -18,7 +18,7 @@ import sys
|
||||
import warnings
|
||||
|
||||
from ..config import get_config
|
||||
from ..engine.task_backend import SyncTaskBackend
|
||||
from ..engine.task_backend import WorkerTaskBackend
|
||||
from .poller import WorkerPoller
|
||||
|
||||
# Filter deprecation warnings from third-party libraries
|
||||
@@ -164,7 +164,11 @@ def main():
|
||||
print(f" Poll interval: {args.poll_interval}ms")
|
||||
print(f" Max retries: {args.max_retries}")
|
||||
print(f" Max slots: {config.worker_max_slots}")
|
||||
print(f" Consolidation max slots: {config.worker_consolidation_max_slots}")
|
||||
reservations = config.worker_slot_reservations
|
||||
reservations_str = ", ".join(f"{k}={v}" for k, v in reservations.items()) if reservations else "none"
|
||||
shared_pool = max(0, config.worker_max_slots - sum(reservations.values()))
|
||||
print(f" Slot reservations: {reservations_str}")
|
||||
print(f" Shared pool: {shared_pool}")
|
||||
print(f" HTTP server: {args.http_host}:{args.http_port}")
|
||||
print()
|
||||
|
||||
@@ -191,11 +195,13 @@ def main():
|
||||
logger.info(f"Loaded operation validator: {operation_validator.__class__.__name__}")
|
||||
|
||||
# Initialize MemoryEngine
|
||||
# Workers use SyncTaskBackend because they execute tasks directly,
|
||||
# they don't need to store tasks (they poll from DB)
|
||||
# Workers use WorkerTaskBackend: submit_task is a no-op because the
|
||||
# row already exists in async_operations. Child tasks (e.g. consolidation
|
||||
# triggered by retain) will be picked up by the poller on the next cycle
|
||||
# instead of being executed inline, which avoids blocking the parent task.
|
||||
memory = MemoryEngine(
|
||||
run_migrations=False, # Workers don't run migrations
|
||||
task_backend=SyncTaskBackend(),
|
||||
task_backend=WorkerTaskBackend(),
|
||||
tenant_extension=tenant_extension,
|
||||
operation_validator=operation_validator,
|
||||
)
|
||||
@@ -222,7 +228,7 @@ def main():
|
||||
schema=schema,
|
||||
tenant_extension=tenant_extension,
|
||||
max_slots=config.worker_max_slots,
|
||||
consolidation_max_slots=config.worker_consolidation_max_slots,
|
||||
slot_reservations=config.worker_slot_reservations,
|
||||
)
|
||||
|
||||
# Create the HTTP app for metrics/health
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,59 @@
|
||||
"""Stage breadcrumbs for in-flight worker tasks.
|
||||
|
||||
The worker poller binds a `StageHolder` to each task's contextvar scope.
|
||||
Engine code calls `set_stage("retain.facts.llm")` at phase boundaries; the
|
||||
poller reads the holder periodically to surface what each in-flight task is
|
||||
currently doing in `WORKER_STATS` / `WORKER_TASK` log lines.
|
||||
|
||||
Outside a worker context the contextvar is unset and `set_stage` is a no-op,
|
||||
so engine code is safe to call from sync HTTP requests, tests, or the CLI
|
||||
without any setup.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from contextvars import ContextVar
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
|
||||
@dataclass
|
||||
class StageHolder:
|
||||
"""Mutable container for the current task's stage label."""
|
||||
|
||||
stage: str = "init"
|
||||
updated_at: float = field(default_factory=time.monotonic)
|
||||
|
||||
|
||||
_current_holder: ContextVar[StageHolder | None] = ContextVar("hindsight_stage_holder", default=None)
|
||||
|
||||
|
||||
def bind_holder(holder: StageHolder):
|
||||
"""Bind a holder to the current async context.
|
||||
|
||||
Must be called from inside the task coroutine itself (not from the
|
||||
spawning code) so the binding lives in the task's own contextvar scope.
|
||||
|
||||
Returns the token that can be passed to `_current_holder.reset()` if
|
||||
the binding ever needs to be unwound.
|
||||
"""
|
||||
return _current_holder.set(holder)
|
||||
|
||||
|
||||
def set_stage(name: str) -> None:
|
||||
"""Update the current task's stage label.
|
||||
|
||||
No-op when called outside a worker task context (e.g. from a sync HTTP
|
||||
request, a test, or the CLI). Cheap enough to call per-phase.
|
||||
"""
|
||||
holder = _current_holder.get()
|
||||
if holder is None:
|
||||
return
|
||||
holder.stage = name
|
||||
holder.updated_at = time.monotonic()
|
||||
|
||||
|
||||
def get_stage() -> str | None:
|
||||
"""Return the current stage label, or None if no holder is bound."""
|
||||
holder = _current_holder.get()
|
||||
return holder.stage if holder is not None else None
|
||||
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "hindsight-api-slim"
|
||||
version = "0.4.22"
|
||||
version = "0.5.4"
|
||||
description = "Hindsight: Agent Memory That Works Like Human Memory"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.11"
|
||||
@@ -21,7 +21,7 @@ dependencies = [
|
||||
"sqlalchemy>=2.0.44",
|
||||
"alembic>=1.17.1",
|
||||
"pgvector>=0.4.1",
|
||||
"greenlet>=3.2.4",
|
||||
"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",
|
||||
@@ -40,7 +40,7 @@ dependencies = [
|
||||
"anthropic>=0.40.0",
|
||||
"typer>=0.9.0",
|
||||
"cohere>=5.0.0",
|
||||
"litellm>=1.0.0,<=1.82.6", # 1.82.7+ contains a supply chain attack (malicious .pth credential stealer)
|
||||
"litellm>=1.83.0", # 1.82.7/1.82.8 had a supply chain compromise (yanked); 1.83.0+ also fixes GHSA-jjhc-v7c2-5hh6 / GHSA-53mr-6c8q-9789
|
||||
"markitdown[pdf,docx,pptx,xlsx,xls]>=0.1.4", # File to markdown conversion
|
||||
"obstore>=0.4.0", # S3/GCS/Azure object storage client (Rust-backed)
|
||||
"winloop>=0.1.0; sys_platform == 'win32'",
|
||||
@@ -78,6 +78,11 @@ local-ml = [
|
||||
"mlx-lm>=0.31.1",
|
||||
"safetensors>=0.6.2",
|
||||
]
|
||||
local-llm = [
|
||||
# Built-in llama.cpp inference for fully offline operation
|
||||
"llama-cpp-python[server]>=0.3.0",
|
||||
"huggingface-hub>=0.20.0",
|
||||
]
|
||||
embedded-db = [
|
||||
"pg0-embedded>=0.11.0",
|
||||
]
|
||||
|
||||
@@ -0,0 +1,47 @@
|
||||
"""Graph-level sanity checks for the Alembic migration DAG.
|
||||
|
||||
These tests do not touch a database; they only parse the revision files on
|
||||
disk, so they are cheap to run in CI and catch DAG accidents (divergent
|
||||
heads, unreachable revisions) at merge time instead of at deploy time.
|
||||
"""
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
from alembic.config import Config
|
||||
from alembic.script import ScriptDirectory
|
||||
|
||||
|
||||
def _script_directory() -> ScriptDirectory:
|
||||
cfg = Config()
|
||||
script_location = Path(__file__).parent.parent / "hindsight_api" / "alembic"
|
||||
cfg.set_main_option("script_location", str(script_location))
|
||||
return ScriptDirectory.from_config(cfg)
|
||||
|
||||
|
||||
def test_single_head() -> None:
|
||||
"""The DAG must have exactly one head.
|
||||
|
||||
A second head means a branch was added without a merge revision, which
|
||||
makes ``alembic upgrade head`` (singular) ambiguous and forces the next
|
||||
migration author to orphan whichever head they don't pick as parent.
|
||||
v0.5.3 shipped in exactly that state; this test would have caught it.
|
||||
|
||||
Fix for a new head: ``alembic merge heads -m "<reason>"``.
|
||||
"""
|
||||
script = _script_directory()
|
||||
heads = script.get_heads()
|
||||
assert len(heads) == 1, (
|
||||
f"Alembic has {len(heads)} heads ({heads}); expected exactly 1. "
|
||||
"Unify them with ``alembic merge heads -m '<reason>'``."
|
||||
)
|
||||
|
||||
|
||||
def test_single_base() -> None:
|
||||
"""The DAG must have exactly one base (the initial schema).
|
||||
|
||||
Multiple bases mean disconnected migration trees, which can only happen
|
||||
through manual file edits.
|
||||
"""
|
||||
script = _script_directory()
|
||||
bases = script.get_bases()
|
||||
assert len(bases) == 1, f"Alembic has {len(bases)} bases ({bases}); expected exactly 1."
|
||||
@@ -433,3 +433,309 @@ async def test_config_retain_batch_tokens_respected(memory, request_context):
|
||||
# Even small batches use parent-child pattern now (simpler code path)
|
||||
assert "child_operations" in status
|
||||
assert status["result_metadata"]["num_sub_batches"] == 1
|
||||
|
||||
|
||||
async def _child_metadata(memory, bank_id: str, parent_operation_id: str, request_context):
|
||||
"""Fetch the first child operation's result_metadata for a parent batch_retain."""
|
||||
parent = await memory.get_operation_status(
|
||||
bank_id=bank_id,
|
||||
operation_id=parent_operation_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
assert parent["status"] == "completed", parent
|
||||
assert parent["child_operations"], "expected at least one child operation"
|
||||
child_id = parent["child_operations"][0]["operation_id"]
|
||||
child = await memory.get_operation_status(
|
||||
bank_id=bank_id,
|
||||
operation_id=child_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
return child["result_metadata"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retain_records_user_provided_document_ids(memory, request_context):
|
||||
"""User-supplied document_ids land in child op result_metadata.document_ids."""
|
||||
bank_id = "test_doc_ids_user_supplied"
|
||||
d1 = str(uuid.uuid4())
|
||||
d2 = str(uuid.uuid4())
|
||||
contents = [
|
||||
{"content": "User-supplied doc one content.", "document_id": d1},
|
||||
{"content": "User-supplied doc two content.", "document_id": d2},
|
||||
]
|
||||
|
||||
result = await memory.submit_async_retain(
|
||||
bank_id=bank_id,
|
||||
contents=contents,
|
||||
request_context=request_context,
|
||||
)
|
||||
await asyncio.sleep(0.2)
|
||||
|
||||
meta = await _child_metadata(memory, bank_id, result["operation_id"], request_context)
|
||||
assert "document_ids" in meta, meta
|
||||
assert set(meta["document_ids"]) == {d1, d2}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retain_records_generated_document_id(memory, request_context):
|
||||
"""With no document_ids supplied, retain records the single generated id."""
|
||||
bank_id = "test_doc_ids_generated"
|
||||
contents = [
|
||||
{"content": "Generated doc item one."},
|
||||
{"content": "Generated doc item two."},
|
||||
]
|
||||
|
||||
result = await memory.submit_async_retain(
|
||||
bank_id=bank_id,
|
||||
contents=contents,
|
||||
request_context=request_context,
|
||||
)
|
||||
await asyncio.sleep(0.2)
|
||||
|
||||
meta = await _child_metadata(memory, bank_id, result["operation_id"], request_context)
|
||||
assert "document_ids" in meta, meta
|
||||
assert isinstance(meta["document_ids"], list)
|
||||
assert len(meta["document_ids"]) == 1
|
||||
# Must be a valid UUID string (generated by the orchestrator)
|
||||
uuid.UUID(meta["document_ids"][0])
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retain_records_shared_document_id_once(memory, request_context):
|
||||
"""Items sharing one document_id record it exactly once (idempotent set-append)."""
|
||||
bank_id = "test_doc_ids_shared"
|
||||
shared = str(uuid.uuid4())
|
||||
# Duplicate per-item doc_ids are rejected up front, so shared-doc mode
|
||||
# is exercised by a single item carrying the id.
|
||||
contents = [{"content": "Shared doc, chunk A.", "document_id": shared}]
|
||||
|
||||
result = await memory.submit_async_retain(
|
||||
bank_id=bank_id,
|
||||
contents=contents,
|
||||
request_context=request_context,
|
||||
)
|
||||
await asyncio.sleep(0.2)
|
||||
|
||||
meta = await _child_metadata(memory, bank_id, result["operation_id"], request_context)
|
||||
assert meta.get("document_ids") == [shared]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_operation_status_include_payload(memory, request_context):
|
||||
"""include_payload=True returns the original submission payload; default omits it."""
|
||||
bank_id = "test_include_payload"
|
||||
contents = [{"content": "Payload roundtrip test item."}]
|
||||
|
||||
result = await memory.submit_async_retain(
|
||||
bank_id=bank_id,
|
||||
contents=contents,
|
||||
request_context=request_context,
|
||||
)
|
||||
await asyncio.sleep(0.2)
|
||||
|
||||
parent = await memory.get_operation_status(
|
||||
bank_id=bank_id,
|
||||
operation_id=result["operation_id"],
|
||||
request_context=request_context,
|
||||
)
|
||||
child_id = parent["child_operations"][0]["operation_id"]
|
||||
|
||||
# Default: no payload
|
||||
without = await memory.get_operation_status(
|
||||
bank_id=bank_id,
|
||||
operation_id=child_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
assert without.get("task_payload") is None
|
||||
|
||||
# With flag: payload populated
|
||||
with_payload = await memory.get_operation_status(
|
||||
bank_id=bank_id,
|
||||
operation_id=child_id,
|
||||
request_context=request_context,
|
||||
include_payload=True,
|
||||
)
|
||||
payload = with_payload.get("task_payload")
|
||||
assert payload is not None, with_payload
|
||||
assert payload.get("bank_id") == bank_id
|
||||
assert payload.get("contents")
|
||||
assert payload["contents"][0]["content"] == "Payload roundtrip test item."
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_operation_status_exposes_retry_count_and_next_retry_at(memory, request_context):
|
||||
"""get_operation_status and list_operations return retry_count and next_retry_at.
|
||||
|
||||
Consumers need these to distinguish a freshly-queued pending task from
|
||||
one that's parked for a future retry (e.g. because an extension raised
|
||||
DeferOperation). Without them, "pending" is ambiguous and callers can't
|
||||
render a helpful "deferred until X" state.
|
||||
"""
|
||||
from datetime import datetime, timedelta, timezone
|
||||
|
||||
bank_id = "test_retry_fields"
|
||||
result = await memory.submit_async_retain(
|
||||
bank_id=bank_id,
|
||||
contents=[{"content": "retry-fields test item"}],
|
||||
request_context=request_context,
|
||||
)
|
||||
await asyncio.sleep(0.1)
|
||||
parent_id = result["operation_id"]
|
||||
child_id = None
|
||||
|
||||
# Get the child op (the batch_retain parent holds a single child in the
|
||||
# sync/simplified path used by SyncTaskBackend tests).
|
||||
parent_status = await memory.get_operation_status(
|
||||
bank_id=bank_id, operation_id=parent_id, request_context=request_context,
|
||||
)
|
||||
assert "retry_count" in parent_status
|
||||
assert "next_retry_at" in parent_status
|
||||
assert parent_status["retry_count"] == 0
|
||||
# Completed tasks should have next_retry_at cleared on the row (or the
|
||||
# status field doesn't include it meaningfully), so we don't assert a
|
||||
# specific value here — only that the key is present.
|
||||
if parent_status.get("child_operations"):
|
||||
child_id = parent_status["child_operations"][0]["operation_id"]
|
||||
|
||||
# list_operations also exposes both fields
|
||||
listed = await memory.list_operations(
|
||||
bank_id=bank_id, request_context=request_context, limit=10, offset=0,
|
||||
)
|
||||
assert listed["operations"], listed
|
||||
for op in listed["operations"]:
|
||||
assert "retry_count" in op
|
||||
assert "next_retry_at" in op
|
||||
assert isinstance(op["retry_count"], int)
|
||||
|
||||
# Simulate a deferred op: set next_retry_at to 15 min in the future for
|
||||
# the child row directly in the DB, then fetch via the API and confirm
|
||||
# the value round-trips as an ISO-8601 string.
|
||||
if child_id:
|
||||
pool = await memory._get_pool()
|
||||
future = datetime.now(timezone.utc) + timedelta(minutes=15)
|
||||
await pool.execute(
|
||||
"UPDATE async_operations SET status = 'pending', next_retry_at = $1, retry_count = 2 WHERE operation_id = $2",
|
||||
future,
|
||||
uuid.UUID(child_id),
|
||||
)
|
||||
fetched = await memory.get_operation_status(
|
||||
bank_id=bank_id, operation_id=child_id, request_context=request_context,
|
||||
)
|
||||
assert fetched["retry_count"] == 2
|
||||
assert fetched["next_retry_at"] is not None
|
||||
# Round-trip tolerance: within 1 second.
|
||||
parsed = datetime.fromisoformat(fetched["next_retry_at"])
|
||||
assert abs((parsed - future).total_seconds()) < 1.0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_request_context_retry_count_propagated_to_validator(memory_no_llm_verify, request_context):
|
||||
"""_handle_batch_retain forwards the task's _retry_count as
|
||||
RequestContext.retry_count, so validator extensions can compute
|
||||
exponential backoff without querying async_operations themselves.
|
||||
"""
|
||||
from hindsight_api.extensions import (
|
||||
OperationValidatorExtension,
|
||||
RecallContext,
|
||||
ReflectContext,
|
||||
RetainContext,
|
||||
ValidationResult,
|
||||
)
|
||||
|
||||
captured: dict[str, int] = {"retry_count": -1}
|
||||
|
||||
class CapturingValidator(OperationValidatorExtension):
|
||||
def __init__(self):
|
||||
super().__init__({})
|
||||
|
||||
async def validate_retain(self, ctx: RetainContext) -> ValidationResult:
|
||||
captured["retry_count"] = ctx.request_context.retry_count
|
||||
return ValidationResult.accept()
|
||||
|
||||
async def validate_recall(self, ctx: RecallContext) -> ValidationResult:
|
||||
return ValidationResult.accept()
|
||||
|
||||
async def validate_reflect(self, ctx: ReflectContext) -> ValidationResult:
|
||||
return ValidationResult.accept()
|
||||
|
||||
memory_no_llm_verify._operation_validator = CapturingValidator()
|
||||
|
||||
bank_id = f"test-retry-propagate-{uuid.uuid4().hex[:8]}"
|
||||
pool = await memory_no_llm_verify._get_pool()
|
||||
await _ensure_bank(pool, bank_id)
|
||||
|
||||
task_dict = {
|
||||
"type": "batch_retain",
|
||||
"bank_id": bank_id,
|
||||
"contents": [{"content": "retry-propagate test"}],
|
||||
"_tenant_id": "default",
|
||||
"_retry_count": 3, # simulate 3rd retry
|
||||
}
|
||||
await memory_no_llm_verify._handle_batch_retain(task_dict)
|
||||
|
||||
assert captured["retry_count"] == 3, (
|
||||
"Validator should see retry_count=3 from task_dict['_retry_count']; "
|
||||
f"got {captured['retry_count']}"
|
||||
)
|
||||
|
||||
# Default (missing _retry_count key) must surface as 0, not raise.
|
||||
captured["retry_count"] = -1
|
||||
task_dict_no_retry = {
|
||||
"type": "batch_retain",
|
||||
"bank_id": bank_id,
|
||||
"contents": [{"content": "retry-propagate default test"}],
|
||||
"_tenant_id": "default",
|
||||
}
|
||||
await memory_no_llm_verify._handle_batch_retain(task_dict_no_retry)
|
||||
assert captured["retry_count"] == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_submit_async_operation_leaves_claimable_row_when_submit_task_fails(memory):
|
||||
"""Regression for the crash-window orphan bug fixed in #1091.
|
||||
|
||||
Previously, _submit_async_operation INSERTed the async_operations row without
|
||||
task_payload, then called submit_task as a separate step to fill it in. If
|
||||
submit_task failed (crash, timeout, dropped connection) after the INSERT
|
||||
committed, the row was left with task_payload IS NULL and became permanently
|
||||
stuck because the worker claim query filters on task_payload IS NOT NULL.
|
||||
|
||||
With the atomic INSERT, even if submit_task raises afterwards the row is born
|
||||
claimable. This test simulates the crash by forcing submit_task to raise.
|
||||
"""
|
||||
bank_id = f"test_orphan_prevention_{uuid.uuid4().hex[:8]}"
|
||||
pool = await memory._get_pool()
|
||||
await _ensure_bank(pool, bank_id)
|
||||
|
||||
async def failing_submit_task(_task_dict):
|
||||
raise RuntimeError("Simulated crash between INSERT and submit_task")
|
||||
|
||||
memory._task_backend.submit_task = failing_submit_task # type: ignore[method-assign]
|
||||
|
||||
with pytest.raises(RuntimeError, match="Simulated crash"):
|
||||
await memory._submit_async_operation(
|
||||
bank_id=bank_id,
|
||||
operation_type="retain",
|
||||
task_type="batch_retain",
|
||||
task_payload={"contents": [{"content": "hello", "document_id": "d1"}]},
|
||||
)
|
||||
|
||||
rows = await pool.fetch(
|
||||
"""
|
||||
SELECT status, task_payload
|
||||
FROM async_operations
|
||||
WHERE bank_id = $1 AND operation_type = 'retain'
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
assert len(rows) == 1, f"Expected exactly one retain row for bank_id={bank_id}, got {len(rows)}"
|
||||
row = rows[0]
|
||||
assert row["status"] == "pending"
|
||||
assert row["task_payload"] is not None, (
|
||||
"task_payload must be set atomically by the INSERT — a NULL here means "
|
||||
"the worker claim query (task_payload IS NOT NULL) will never pick this row up"
|
||||
)
|
||||
payload = json.loads(row["task_payload"])
|
||||
assert payload["type"] == "batch_retain"
|
||||
assert payload["bank_id"] == bank_id
|
||||
assert payload["contents"] == [{"content": "hello", "document_id": "d1"}]
|
||||
|
||||
@@ -34,7 +34,12 @@ async def test_submit_async_retain_includes_document_tags_in_task_payload():
|
||||
contents = [{"content": "Async retain payload test."}]
|
||||
document_tags = ["scope:tools", "user:alice"]
|
||||
|
||||
with patch("hindsight_api.engine.memory_engine.bank_utils.get_bank_profile", new_callable=AsyncMock):
|
||||
# Return (profile, created=False) so the default-template-on-create hook is skipped.
|
||||
with patch(
|
||||
"hindsight_api.engine.memory_engine.bank_utils.get_or_create_bank_profile",
|
||||
new_callable=AsyncMock,
|
||||
return_value=(MagicMock(), False),
|
||||
):
|
||||
result = await MemoryEngine.submit_async_retain(
|
||||
engine,
|
||||
bank_id="bank-1",
|
||||
|
||||
@@ -0,0 +1,242 @@
|
||||
"""
|
||||
Tests for the bank stats endpoint and the memories-timeseries endpoint.
|
||||
|
||||
Covers the new fields exposed by GET /v1/default/banks/{bank_id}/stats
|
||||
(operations_by_status) and the new endpoint
|
||||
GET /v1/default/banks/{bank_id}/stats/memories-timeseries.
|
||||
"""
|
||||
import uuid
|
||||
from datetime import datetime
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
|
||||
from hindsight_api.api import create_app
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def api_client(memory):
|
||||
app = create_app(memory, initialize_memory=False)
|
||||
transport = httpx.ASGITransport(app=app)
|
||||
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
|
||||
yield client
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def test_bank_id():
|
||||
return f"stats_test_{datetime.now().timestamp()}"
|
||||
|
||||
|
||||
async def _insert_memory(memory, bank_id: str, text: str, *, failed: bool = False) -> str:
|
||||
"""Insert a single experience memory, optionally marked as consolidation-failed."""
|
||||
mem_id = uuid.uuid4()
|
||||
async with memory._pool.acquire() as conn:
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO memory_units (id, bank_id, text, fact_type, created_at, consolidation_failed_at)
|
||||
VALUES ($1, $2, $3, 'experience', now(), CASE WHEN $4 THEN now() ELSE NULL END)
|
||||
""",
|
||||
mem_id,
|
||||
bank_id,
|
||||
text,
|
||||
failed,
|
||||
)
|
||||
return str(mem_id)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bank_stats_exposes_operations_by_status(api_client, test_bank_id):
|
||||
"""/stats should return operations_by_status with all finished operations."""
|
||||
try:
|
||||
# Kick off a retain so at least one completed operation exists.
|
||||
response = await api_client.post(
|
||||
f"/v1/default/banks/{test_bank_id}/memories",
|
||||
json={"items": [{"content": "Alice is a software engineer.", "context": "team"}]},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
|
||||
response = await api_client.get(f"/v1/default/banks/{test_bank_id}/stats")
|
||||
assert response.status_code == 200
|
||||
stats = response.json()
|
||||
|
||||
assert "operations_by_status" in stats
|
||||
assert isinstance(stats["operations_by_status"], dict)
|
||||
# A synchronous retain finishes as "completed".
|
||||
assert stats["operations_by_status"].get("completed", 0) >= 1
|
||||
# pending/failed counters should still be present as scalar mirrors.
|
||||
assert stats["pending_operations"] == stats["operations_by_status"].get("pending", 0)
|
||||
assert stats["failed_operations"] == stats["operations_by_status"].get("failed", 0)
|
||||
finally:
|
||||
await api_client.delete(f"/v1/default/banks/{test_bank_id}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"period,expected_count,expected_trunc",
|
||||
[
|
||||
("1h", 60, "minute"),
|
||||
("12h", 12, "hour"),
|
||||
("1d", 24, "hour"),
|
||||
("7d", 7, "day"),
|
||||
("30d", 30, "day"),
|
||||
("90d", 90, "day"),
|
||||
],
|
||||
)
|
||||
async def test_memories_timeseries_periods(
|
||||
api_client, test_bank_id, period, expected_count, expected_trunc
|
||||
):
|
||||
"""Every period must return the full expected bucket count and trunc."""
|
||||
try:
|
||||
response = await api_client.post(
|
||||
f"/v1/default/banks/{test_bank_id}/memories",
|
||||
json={"items": [{"content": "Bob works on infrastructure.", "context": "team"}]},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
|
||||
response = await api_client.get(
|
||||
f"/v1/default/banks/{test_bank_id}/stats/memories-timeseries",
|
||||
params={"period": period},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
body = response.json()
|
||||
|
||||
assert body["bank_id"] == test_bank_id
|
||||
assert body["period"] == period
|
||||
assert body["trunc"] == expected_trunc
|
||||
assert len(body["buckets"]) == expected_count
|
||||
|
||||
for bucket in body["buckets"]:
|
||||
assert "time" in bucket
|
||||
assert bucket["world"] >= 0
|
||||
assert bucket["experience"] >= 0
|
||||
assert bucket["observation"] >= 0
|
||||
finally:
|
||||
await api_client.delete(f"/v1/default/banks/{test_bank_id}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_memories_timeseries_invalid_period_falls_back(api_client, test_bank_id):
|
||||
"""An unknown period must fall back to the 7d default."""
|
||||
try:
|
||||
response = await api_client.get(
|
||||
f"/v1/default/banks/{test_bank_id}/stats/memories-timeseries",
|
||||
params={"period": "nonsense"},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
body = response.json()
|
||||
assert body["period"] == "7d"
|
||||
assert body["trunc"] == "day"
|
||||
assert len(body["buckets"]) == 7
|
||||
finally:
|
||||
await api_client.delete(f"/v1/default/banks/{test_bank_id}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_memories_timeseries_empty_bank_returns_zero_filled_buckets(
|
||||
api_client, test_bank_id
|
||||
):
|
||||
"""A bank with no memories must still return the full zero-filled bucket set."""
|
||||
try:
|
||||
# Ensure the bank exists.
|
||||
response = await api_client.get(f"/v1/default/banks/{test_bank_id}/profile")
|
||||
assert response.status_code == 200
|
||||
|
||||
response = await api_client.get(
|
||||
f"/v1/default/banks/{test_bank_id}/stats/memories-timeseries",
|
||||
params={"period": "7d"},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
body = response.json()
|
||||
assert len(body["buckets"]) == 7
|
||||
for bucket in body["buckets"]:
|
||||
assert bucket["world"] == 0
|
||||
assert bucket["experience"] == 0
|
||||
assert bucket["observation"] == 0
|
||||
finally:
|
||||
await api_client.delete(f"/v1/default/banks/{test_bank_id}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_memories_timeseries_reflects_retained_memories(api_client, test_bank_id):
|
||||
"""Freshly-retained memories must show up in today's bucket counts."""
|
||||
try:
|
||||
response = await api_client.post(
|
||||
f"/v1/default/banks/{test_bank_id}/memories",
|
||||
json={
|
||||
"items": [
|
||||
{"content": "Alice is a software engineer.", "context": "team"},
|
||||
{"content": "Bob works on infrastructure.", "context": "team"},
|
||||
]
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
|
||||
response = await api_client.get(
|
||||
f"/v1/default/banks/{test_bank_id}/stats/memories-timeseries",
|
||||
params={"period": "7d"},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
body = response.json()
|
||||
totals = sum(b["world"] + b["experience"] + b["observation"] for b in body["buckets"])
|
||||
assert totals >= 2, "expected at least two memories across all buckets"
|
||||
|
||||
# Those memories should land in the most-recent bucket.
|
||||
latest = body["buckets"][-1]
|
||||
assert latest["world"] + latest["experience"] + latest["observation"] >= 2
|
||||
finally:
|
||||
await api_client.delete(f"/v1/default/banks/{test_bank_id}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bank_stats_reports_failed_consolidation(api_client, memory, test_bank_id):
|
||||
"""/stats must surface the count of memories with consolidation_failed_at set."""
|
||||
try:
|
||||
await _insert_memory(memory, test_bank_id, "Alice failed 1.", failed=True)
|
||||
await _insert_memory(memory, test_bank_id, "Alice failed 2.", failed=True)
|
||||
await _insert_memory(memory, test_bank_id, "Alice pending.", failed=False)
|
||||
|
||||
response = await api_client.get(f"/v1/default/banks/{test_bank_id}/stats")
|
||||
assert response.status_code == 200
|
||||
stats = response.json()
|
||||
|
||||
assert stats["failed_consolidation"] == 2
|
||||
# The two failed memories also count as "not-yet-consolidated".
|
||||
assert stats["pending_consolidation"] >= 3
|
||||
finally:
|
||||
await api_client.delete(f"/v1/default/banks/{test_bank_id}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_memories_filter_by_consolidation_state_failed(api_client, memory, test_bank_id):
|
||||
"""?consolidation_state=failed returns only memories with consolidation_failed_at set."""
|
||||
try:
|
||||
failed_id = await _insert_memory(memory, test_bank_id, "Broken item.", failed=True)
|
||||
await _insert_memory(memory, test_bank_id, "Healthy item.", failed=False)
|
||||
|
||||
response = await api_client.get(
|
||||
f"/v1/default/banks/{test_bank_id}/memories/list",
|
||||
params={"consolidation_state": "failed"},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
body = response.json()
|
||||
|
||||
ids = [item["id"] for item in body["items"]]
|
||||
assert failed_id in ids
|
||||
assert body["total"] == 1
|
||||
assert body["items"][0]["consolidation_failed_at"] is not None
|
||||
finally:
|
||||
await api_client.delete(f"/v1/default/banks/{test_bank_id}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_memories_filter_by_consolidation_state_rejects_unknown(api_client, test_bank_id):
|
||||
"""An invalid consolidation_state value must return a 400 (not 500)."""
|
||||
try:
|
||||
response = await api_client.get(
|
||||
f"/v1/default/banks/{test_bank_id}/memories/list",
|
||||
params={"consolidation_state": "bogus"},
|
||||
)
|
||||
assert response.status_code == 400
|
||||
finally:
|
||||
await api_client.delete(f"/v1/default/banks/{test_bank_id}")
|
||||
@@ -0,0 +1,115 @@
|
||||
"""Verify that BankTemplateConfig exposes every hierarchical field that
|
||||
_CONFIGURABLE_FIELDS already accepts at the engine layer.
|
||||
|
||||
This test guards the fix for the gap described in the upstream PR title
|
||||
"fix(bank-template): align BankTemplateConfig with _CONFIGURABLE_FIELDS".
|
||||
Each new field is POSTed through /v1/default/banks/{id}/import and then
|
||||
read back via the bank-config endpoint; assertion is that the applied
|
||||
value round-trips through the engine.
|
||||
|
||||
Runs via: uv run pytest tests/test_bank_template_configurable_fields.py -v
|
||||
|
||||
The api_client fixture (shared with tests/test_bank_templates.py) wraps
|
||||
create_app(memory, initialize_memory=False) in an httpx.ASGITransport
|
||||
with base_url http://test — in-process, no network, no tenant extension.
|
||||
Copy the fixture inline here so the test file does not depend on a
|
||||
conftest we do not ship in the patch.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
|
||||
from hindsight_api.api import create_app
|
||||
from hindsight_api.api.http import BankTemplateConfig
|
||||
|
||||
# Each tuple is (field_name, applied_value). Values chosen to differ
|
||||
# visibly from defaults so round-trip bugs surface.
|
||||
NEW_FIELDS: list[tuple[str, object]] = [
|
||||
("retain_default_strategy", "strategy-a"),
|
||||
("retain_strategies", {"strategy-a": {"mode": "concise", "max_tokens": 512}}),
|
||||
("retain_chunk_batch_size", 7),
|
||||
("mcp_enabled_tools", ["list_banks", "get_bank_profile"]),
|
||||
("consolidation_llm_batch_size", 11),
|
||||
("consolidation_source_facts_max_tokens", 2048),
|
||||
("consolidation_source_facts_max_tokens_per_observation", 256),
|
||||
("max_observations_per_scope", 13),
|
||||
("reflect_source_facts_max_tokens", 4096),
|
||||
("llm_gemini_safety_settings", [{"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_NONE"}]),
|
||||
("recall_budget_function", "adaptive"),
|
||||
("recall_budget_fixed_low", 50),
|
||||
("recall_budget_fixed_mid", 250),
|
||||
("recall_budget_fixed_high", 800),
|
||||
("recall_budget_adaptive_low", 0.05),
|
||||
("recall_budget_adaptive_mid", 0.1),
|
||||
("recall_budget_adaptive_high", 0.4),
|
||||
("recall_budget_min", 30),
|
||||
("recall_budget_max", 1500),
|
||||
]
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def api_client(memory):
|
||||
"""Matches the fixture in tests/test_bank_templates.py — in-process
|
||||
ASGI test client, no tenant extension, no auth."""
|
||||
app = create_app(memory, initialize_memory=False)
|
||||
transport = httpx.ASGITransport(app=app)
|
||||
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
|
||||
yield client
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def bank_id():
|
||||
return f"tmpl_config_{datetime.now().timestamp()}"
|
||||
|
||||
|
||||
def test_bank_template_config_declares_every_configurable_field():
|
||||
"""Pydantic-level guard: every field in NEW_FIELDS must be a declared
|
||||
attribute of BankTemplateConfig so get_config_updates() picks it up."""
|
||||
declared = set(BankTemplateConfig.model_fields.keys())
|
||||
missing = [name for name, _ in NEW_FIELDS if name not in declared]
|
||||
assert not missing, f"BankTemplateConfig missing fields: {missing}"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("field_name,applied_value", NEW_FIELDS, ids=[n for n, _ in NEW_FIELDS])
|
||||
async def test_new_field_round_trips_through_import(
|
||||
api_client: httpx.AsyncClient,
|
||||
bank_id: str,
|
||||
field_name: str,
|
||||
applied_value: object,
|
||||
):
|
||||
"""POST a minimal manifest with one new field set, then read bank
|
||||
config back and assert the value made it through.
|
||||
|
||||
Bank config response shape per upstream's test_import_applies_config:
|
||||
top-level keys are resolved hierarchical config; per-bank overrides
|
||||
live under config["overrides"][<field>]. Assert on the override slot.
|
||||
"""
|
||||
unique_bank_id = f"{bank_id}_{field_name}"
|
||||
manifest = {
|
||||
"version": "1",
|
||||
"bank": {field_name: applied_value},
|
||||
}
|
||||
|
||||
resp = await api_client.post(
|
||||
f"/v1/default/banks/{unique_bank_id}/import",
|
||||
json=manifest,
|
||||
)
|
||||
assert resp.status_code == 200, resp.text
|
||||
|
||||
# Read bank config back — field must reflect the applied value
|
||||
# under the "overrides" slot, matching upstream's own test shape.
|
||||
read = await api_client.get(f"/v1/default/banks/{unique_bank_id}/config")
|
||||
assert read.status_code == 200, read.text
|
||||
config = read.json()
|
||||
overrides = config.get("overrides", {})
|
||||
assert overrides.get(field_name) == applied_value, (
|
||||
f"round-trip mismatch for {field_name}: "
|
||||
f"sent {applied_value!r}, got {overrides.get(field_name)!r} "
|
||||
f"(full overrides: {overrides!r})"
|
||||
)
|
||||
@@ -598,3 +598,182 @@ class TestExport:
|
||||
assert resp.status_code == 200
|
||||
data = resp.json()
|
||||
assert data["version"] == "1"
|
||||
|
||||
|
||||
class TestDefaultBankTemplateEnvVar:
|
||||
"""Tests for HINDSIGHT_API_DEFAULT_BANK_TEMPLATE — a server-level env var
|
||||
whose manifest is applied automatically to every newly-created bank."""
|
||||
|
||||
@pytest.fixture
|
||||
def default_template(self):
|
||||
return {
|
||||
"version": "1",
|
||||
"bank": {
|
||||
"reflect_mission": "default-env-mission",
|
||||
"retain_extraction_mode": "verbose",
|
||||
"disposition_empathy": 5,
|
||||
"disposition_skepticism": 1,
|
||||
},
|
||||
"mental_models": [
|
||||
{
|
||||
"id": "default-env-model",
|
||||
"name": "Default Env Model",
|
||||
"source_query": "What is the default?",
|
||||
},
|
||||
],
|
||||
"directives": [
|
||||
{
|
||||
"name": "Default Env Directive",
|
||||
"content": "Follow the default behavior.",
|
||||
"priority": 7,
|
||||
},
|
||||
],
|
||||
}
|
||||
|
||||
@pytest.fixture
|
||||
def _patched_default_template(self, monkeypatch, default_template):
|
||||
"""Install the default template on the already-initialized global config.
|
||||
|
||||
We can't rely on env-var resolution here: MemoryEngine (and its
|
||||
ConfigResolver) snapshot the global config at fixture init time.
|
||||
Patching the field directly keeps the test deterministic while still
|
||||
exercising the same code path that reads `get_config().default_bank_template`.
|
||||
"""
|
||||
from hindsight_api.config import _get_raw_config
|
||||
|
||||
raw = _get_raw_config()
|
||||
monkeypatch.setattr(raw, "default_bank_template", default_template)
|
||||
yield default_template
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_default_template_applied_on_new_bank(
|
||||
self, api_client, bank_id, _patched_default_template
|
||||
):
|
||||
"""Creating a new bank applies the default template (config + mental models + directives)."""
|
||||
# Trigger bank auto-creation via GET profile
|
||||
resp = await api_client.put(f"/v1/default/banks/{bank_id}", json={})
|
||||
assert resp.status_code == 200
|
||||
|
||||
# Config from template should be present as bank overrides
|
||||
config_resp = await api_client.get(f"/v1/default/banks/{bank_id}/config")
|
||||
assert config_resp.status_code == 200
|
||||
overrides = config_resp.json()["overrides"]
|
||||
assert overrides["reflect_mission"] == "default-env-mission"
|
||||
assert overrides["retain_extraction_mode"] == "verbose"
|
||||
assert overrides["disposition_empathy"] == 5
|
||||
assert overrides["disposition_skepticism"] == 1
|
||||
|
||||
# Mental model from template should exist
|
||||
mm_resp = await api_client.get(f"/v1/default/banks/{bank_id}/mental-models/default-env-model")
|
||||
assert mm_resp.status_code == 200
|
||||
assert mm_resp.json()["name"] == "Default Env Model"
|
||||
|
||||
# Directive from template should exist
|
||||
dir_resp = await api_client.get(f"/v1/default/banks/{bank_id}/directives")
|
||||
assert dir_resp.status_code == 200
|
||||
names = [d["name"] for d in dir_resp.json()["items"]]
|
||||
assert "Default Env Directive" in names
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_default_template_overrides_env_config_defaults(
|
||||
self, api_client, bank_id, monkeypatch, default_template
|
||||
):
|
||||
"""Fields set by the default template override server-level env-var defaults.
|
||||
|
||||
We point both HINDSIGHT_API_RETAIN_EXTRACTION_MODE (env) and the
|
||||
default template at different values, then confirm the template wins
|
||||
via the per-bank config overrides layer (highest precedence).
|
||||
"""
|
||||
from hindsight_api.config import _get_raw_config
|
||||
|
||||
raw = _get_raw_config()
|
||||
# Simulate an env-level default of "concise", overridden by a template that sets "verbose".
|
||||
monkeypatch.setattr(raw, "retain_extraction_mode", "concise")
|
||||
monkeypatch.setattr(raw, "default_bank_template", default_template)
|
||||
|
||||
resp = await api_client.put(f"/v1/default/banks/{bank_id}", json={})
|
||||
assert resp.status_code == 200
|
||||
|
||||
config_resp = await api_client.get(f"/v1/default/banks/{bank_id}/config")
|
||||
overrides = config_resp.json()["overrides"]
|
||||
# Template value wins at the bank-override layer.
|
||||
assert overrides["retain_extraction_mode"] == "verbose"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_default_template_not_reapplied_on_existing_bank(
|
||||
self, api_client, bank_id, _patched_default_template
|
||||
):
|
||||
"""Template only applies on FIRST creation; subsequent puts are no-ops."""
|
||||
# First hit creates the bank and applies the template
|
||||
resp = await api_client.put(f"/v1/default/banks/{bank_id}", json={})
|
||||
assert resp.status_code == 200
|
||||
|
||||
# User explicitly overrides a template-set field
|
||||
patch_resp = await api_client.patch(
|
||||
f"/v1/default/banks/{bank_id}/config",
|
||||
json={"updates": {"reflect_mission": "user-override"}},
|
||||
)
|
||||
assert patch_resp.status_code == 200
|
||||
|
||||
# Second put — template must NOT be reapplied (would clobber the override)
|
||||
resp = await api_client.put(f"/v1/default/banks/{bank_id}", json={})
|
||||
assert resp.status_code == 200
|
||||
|
||||
config_resp = await api_client.get(f"/v1/default/banks/{bank_id}/config")
|
||||
assert config_resp.json()["overrides"]["reflect_mission"] == "user-override"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_default_template_unset_is_noop(self, api_client, bank_id):
|
||||
"""With the env var unset (fixture default), bank creation behaves as before."""
|
||||
resp = await api_client.put(f"/v1/default/banks/{bank_id}", json={})
|
||||
assert resp.status_code == 200
|
||||
|
||||
# No template = no overrides
|
||||
config_resp = await api_client.get(f"/v1/default/banks/{bank_id}/config")
|
||||
assert config_resp.json()["overrides"] == {}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_default_template_malformed_is_swallowed(
|
||||
self, api_client, bank_id, monkeypatch
|
||||
):
|
||||
"""A malformed default template is logged and ignored — bank creation still succeeds."""
|
||||
from hindsight_api.config import _get_raw_config
|
||||
|
||||
raw = _get_raw_config()
|
||||
# Wrong version number fails Pydantic validation.
|
||||
monkeypatch.setattr(raw, "default_bank_template", {"version": "999"})
|
||||
|
||||
resp = await api_client.put(f"/v1/default/banks/{bank_id}", json={})
|
||||
# Bank creation must not fail even though the template is broken.
|
||||
assert resp.status_code == 200
|
||||
|
||||
def test_parse_default_bank_template_valid_json(self, monkeypatch):
|
||||
"""_parse_default_bank_template parses a valid JSON object env var."""
|
||||
from hindsight_api.config import _parse_default_bank_template
|
||||
|
||||
parsed = _parse_default_bank_template('{"version": "1", "bank": {"disposition_empathy": 4}}')
|
||||
assert parsed == {"version": "1", "bank": {"disposition_empathy": 4}}
|
||||
|
||||
def test_parse_default_bank_template_none_or_empty(self):
|
||||
"""Unset / empty env var resolves to None."""
|
||||
from hindsight_api.config import _parse_default_bank_template
|
||||
|
||||
assert _parse_default_bank_template(None) is None
|
||||
assert _parse_default_bank_template("") is None
|
||||
assert _parse_default_bank_template(" ") is None
|
||||
|
||||
def test_parse_default_bank_template_invalid_json_raises(self):
|
||||
"""Invalid JSON fails fast with a clear error."""
|
||||
from hindsight_api.config import _parse_default_bank_template
|
||||
|
||||
with pytest.raises(ValueError, match="HINDSIGHT_API_DEFAULT_BANK_TEMPLATE"):
|
||||
_parse_default_bank_template("not-json")
|
||||
|
||||
def test_parse_default_bank_template_non_object_raises(self):
|
||||
"""Non-object JSON (e.g. array, string) fails fast."""
|
||||
from hindsight_api.config import _parse_default_bank_template
|
||||
|
||||
with pytest.raises(ValueError, match="expected a JSON object"):
|
||||
_parse_default_bank_template("[1, 2, 3]")
|
||||
with pytest.raises(ValueError, match="expected a JSON object"):
|
||||
_parse_default_bank_template('"just a string"')
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user