Compare commits
246
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
9e16dadaa4 | ||
|
|
300a45c2aa | ||
|
|
f6d522a650 | ||
|
|
3485ea2541 | ||
|
|
12e5d2a217 | ||
|
|
086bcac64d | ||
|
|
552f9d19c8 | ||
|
|
cd2da82819 | ||
|
|
402f0d3dab | ||
|
|
2bf399c2ee | ||
|
|
71ad98c7e0 | ||
|
|
0017ee0877 | ||
|
|
1b7357d99a | ||
|
|
c88e6e4cac | ||
|
|
858809ec06 | ||
|
|
ac7c1fcd8c | ||
|
|
dc1a741642 | ||
|
|
729639a557 | ||
|
|
b66b36c080 | ||
|
|
b92346b696 | ||
|
|
28a7caa63b | ||
|
|
46a4db3b24 | ||
|
|
fd7f218488 | ||
|
|
d46e4b1e61 | ||
|
|
eb48593e6e | ||
|
|
838f8e9be1 | ||
|
|
ba8f119448 | ||
|
|
3ad390d834 | ||
|
|
514697a81f | ||
|
|
e2f5bd0380 | ||
|
|
d18e83b083 | ||
|
|
b837e66ce6 | ||
|
|
daae8223c3 | ||
|
|
3455460e0b | ||
|
|
2bada2dbec | ||
|
|
324b4b0a59 | ||
|
|
9be2e0503a | ||
|
|
526c61a170 | ||
|
|
3ce26866d2 | ||
|
|
8314de5e06 | ||
|
|
1a37ad15d1 | ||
|
|
300a8c1e81 | ||
|
|
b50a86a87f | ||
|
|
1d85f5a0d6 | ||
|
|
2e5ed7f936 | ||
|
|
76bcd93156 | ||
|
|
b153541e27 | ||
|
|
c91696f53d | ||
|
|
50f559c9e4 | ||
|
|
4ea650f6c4 | ||
|
|
df6662fe04 | ||
|
|
75dd70fa8a | ||
|
|
92f3ee4671 | ||
|
|
e9b187330c | ||
|
|
28c9aa6151 | ||
|
|
98593f9a20 | ||
|
|
868d5e2ffd | ||
|
|
c308e473a8 | ||
|
|
8fbe85f0ca | ||
|
|
e97a5c9a6e | ||
|
|
a50567f864 | ||
|
|
461b00d4d9 | ||
|
|
4bc772d8f8 | ||
|
|
99a8978905 | ||
|
|
91106f30ef | ||
|
|
5b1c3486f3 | ||
|
|
685e4cf0ef | ||
|
|
73a0ad0cc3 | ||
|
|
f7cc9ad663 | ||
|
|
843dcec77b | ||
|
|
ae0e3cec8d | ||
|
|
a9967627ae | ||
|
|
f6d659c927 | ||
|
|
794b83d839 | ||
|
|
dcc2d69d6f | ||
|
|
4ba2fffe8d | ||
|
|
0cffa43cbd | ||
|
|
70677457d9 | ||
|
|
ee2d8f7540 | ||
|
|
b962de50d8 | ||
|
|
130bb2d616 | ||
|
|
b0c1afb19f | ||
|
|
9654a06e22 | ||
|
|
e1c6092785 | ||
|
|
6fb8ac97a0 | ||
|
|
ecd0b846ed | ||
|
|
0bbc058336 | ||
|
|
4ba54d8c8f | ||
|
|
da55dbb694 | ||
|
|
ab5d2b783b | ||
|
|
13b1d92297 | ||
|
|
a7514e1868 | ||
|
|
db7f492103 | ||
|
|
cd1ab497c5 | ||
|
|
6034e5383d | ||
|
|
cdc26daa2a | ||
|
|
b67b688635 | ||
|
|
ac5181f565 | ||
|
|
9f4b3b670f | ||
|
|
42ed681440 | ||
|
|
06147ef2e9 | ||
|
|
8eb6e69a75 | ||
|
|
8f6e0e5bec | ||
|
|
80982da577 | ||
|
|
90674aef17 | ||
|
|
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 |
+6
-1
@@ -2,7 +2,7 @@
|
||||
# Copy this file to .env and fill in your values
|
||||
|
||||
# LLM Configuration (Required)
|
||||
# Supported providers: openai, groq, ollama, gemini, anthropic, lmstudio, vertexai, minimax, volcano
|
||||
# Supported providers: openai, groq, ollama, gemini, anthropic, lmstudio, vertexai, minimax, deepseek, volcano
|
||||
HINDSIGHT_API_LLM_PROVIDER=openai
|
||||
HINDSIGHT_API_LLM_API_KEY=your-api-key-here
|
||||
HINDSIGHT_API_LLM_MODEL=gpt-4o-mini
|
||||
@@ -25,6 +25,11 @@ HINDSIGHT_API_LLM_BASE_URL=https://api.openai.com/v1
|
||||
# HINDSIGHT_API_LLM_API_KEY=your-minimax-api-key
|
||||
# HINDSIGHT_API_LLM_MODEL=MiniMax-M2.7
|
||||
|
||||
# Example: DeepSeek configuration (https://api.deepseek.com)
|
||||
# HINDSIGHT_API_LLM_PROVIDER=deepseek
|
||||
# HINDSIGHT_API_LLM_API_KEY=your-deepseek-api-key
|
||||
# HINDSIGHT_API_LLM_MODEL=deepseek-v4-flash # or deepseek-v4-pro / deepseek-chat / deepseek-reasoner
|
||||
|
||||
# Example: LM Studio local configuration (Qwen 2.5 32B recommended)
|
||||
# HINDSIGHT_API_LLM_PROVIDER=lmstudio
|
||||
# HINDSIGHT_API_LLM_API_KEY=lmstudio
|
||||
|
||||
@@ -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@v7
|
||||
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@v7
|
||||
with:
|
||||
name: locomo-results-${{ github.sha }}
|
||||
path: hindsight-dev/benchmarks/locomo/results/
|
||||
retention-days: 90
|
||||
@@ -497,6 +497,12 @@ jobs:
|
||||
name: rust-cli-hindsight-linux-amd64
|
||||
path: ./artifacts/rust-cli-linux
|
||||
|
||||
- name: Download Rust CLI (Linux ARM)
|
||||
uses: actions/download-artifact@v8
|
||||
with:
|
||||
name: rust-cli-hindsight-linux-arm64
|
||||
path: ./artifacts/rust-cli-linux-arm64
|
||||
|
||||
- name: Download Rust CLI (macOS Intel)
|
||||
uses: actions/download-artifact@v8
|
||||
with:
|
||||
@@ -533,6 +539,7 @@ jobs:
|
||||
cp artifacts/control-plane/*.tgz release-assets/ || true
|
||||
# Rust CLI binaries
|
||||
cp artifacts/rust-cli-linux/hindsight-linux-amd64 release-assets/ || true
|
||||
cp artifacts/rust-cli-linux-arm64/hindsight-linux-arm64 release-assets/ || true
|
||||
cp artifacts/rust-cli-darwin-amd64/hindsight-darwin-amd64 release-assets/ || true
|
||||
cp artifacts/rust-cli-darwin-arm64/hindsight-darwin-arm64 release-assets/ || true
|
||||
# Helm chart
|
||||
|
||||
+561
-3
@@ -49,6 +49,12 @@ jobs:
|
||||
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 }}
|
||||
integrations-pipecat: ${{ steps.filter.outputs.integrations-pipecat }}
|
||||
integrations-agentcore: ${{ steps.filter.outputs.integrations-agentcore }}
|
||||
integrations-smolagents: ${{ steps.filter.outputs.integrations-smolagents }}
|
||||
tools-agent-sdk: ${{ steps.filter.outputs.tools-agent-sdk }}
|
||||
tools-self-driving-agents: ${{ steps.filter.outputs.tools-self-driving-agents }}
|
||||
dev: ${{ steps.filter.outputs.dev }}
|
||||
ci: ${{ steps.filter.outputs.ci }}
|
||||
# Secrets are available for internal PRs, pull_request_review, and workflow_dispatch.
|
||||
@@ -133,6 +139,18 @@ jobs:
|
||||
- 'hindsight-integrations/*/package-lock.json'
|
||||
- 'hindsight-integrations/*/package.json'
|
||||
- 'scripts/check-integration-lockfiles.sh'
|
||||
integrations-openai-agents:
|
||||
- 'hindsight-integrations/openai-agents/**'
|
||||
integrations-pipecat:
|
||||
- 'hindsight-integrations/pipecat/**'
|
||||
integrations-agentcore:
|
||||
- 'hindsight-integrations/agentcore/**'
|
||||
integrations-smolagents:
|
||||
- 'hindsight-integrations/smolagents/**'
|
||||
tools-agent-sdk:
|
||||
- 'hindsight-tools/hindsight-agent-sdk/**'
|
||||
tools-self-driving-agents:
|
||||
- 'hindsight-tools/self-driving-agents/**'
|
||||
dev:
|
||||
- 'hindsight-dev/**'
|
||||
ci:
|
||||
@@ -286,6 +304,9 @@ jobs:
|
||||
- name: Build hindsight-all-npm (openclaw dep)
|
||||
run: npm run build --workspace=hindsight-all-npm
|
||||
|
||||
- name: Build hindsight-agent-sdk (openclaw dep)
|
||||
run: npm run build --workspace=hindsight-tools/hindsight-agent-sdk
|
||||
|
||||
- name: Install openclaw dependencies
|
||||
working-directory: ./hindsight-integrations/openclaw
|
||||
run: npm ci
|
||||
@@ -349,6 +370,9 @@ jobs:
|
||||
- name: Build hindsight-all-npm (openclaw dep)
|
||||
run: npm run build --workspace=hindsight-all-npm
|
||||
|
||||
- name: Build hindsight-agent-sdk (openclaw dep)
|
||||
run: npm run build --workspace=hindsight-tools/hindsight-agent-sdk
|
||||
|
||||
- name: Install openclaw dependencies
|
||||
working-directory: ./hindsight-integrations/openclaw
|
||||
run: npm ci
|
||||
@@ -503,6 +527,72 @@ jobs:
|
||||
working-directory: ./hindsight-integrations/opencode
|
||||
run: npm run build
|
||||
|
||||
test-hindsight-agent-sdk:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
github.event_name != 'pull_request_review' &&
|
||||
(github.event_name == 'workflow_dispatch' ||
|
||||
needs.detect-changes.outputs.tools-agent-sdk == '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
|
||||
|
||||
- name: Install root workspace dependencies
|
||||
run: npm ci
|
||||
|
||||
- name: Build hindsight-client (agent-sdk dep)
|
||||
run: npm run build --workspace=hindsight-clients/typescript
|
||||
|
||||
- name: Run tests
|
||||
run: npm test --workspace=hindsight-tools/hindsight-agent-sdk
|
||||
|
||||
- name: Build
|
||||
run: npm run build --workspace=hindsight-tools/hindsight-agent-sdk
|
||||
|
||||
test-self-driving-agents:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
github.event_name != 'pull_request_review' &&
|
||||
(github.event_name == 'workflow_dispatch' ||
|
||||
needs.detect-changes.outputs.tools-self-driving-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: Set up Node.js
|
||||
uses: actions/setup-node@v6
|
||||
with:
|
||||
node-version: '22'
|
||||
cache: 'npm'
|
||||
cache-dependency-path: package-lock.json
|
||||
|
||||
- name: Install root workspace dependencies
|
||||
run: npm ci
|
||||
|
||||
- name: Build hindsight-client (self-driving-agents dep)
|
||||
run: npm run build --workspace=hindsight-clients/typescript
|
||||
|
||||
- name: Run tests
|
||||
run: npm test --workspace=hindsight-tools/self-driving-agents
|
||||
|
||||
- name: Build
|
||||
run: npm run build --workspace=hindsight-tools/self-driving-agents
|
||||
|
||||
test-cloudflare-oauth-proxy-integration:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
@@ -596,6 +686,44 @@ jobs:
|
||||
working-directory: ./hindsight-integrations/paperclip
|
||||
run: npm test
|
||||
|
||||
test-pipecat-integration:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
github.event_name != 'pull_request_review' &&
|
||||
(github.event_name == 'workflow_dispatch' ||
|
||||
needs.detect-changes.outputs.integrations-pipecat == '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 pipecat integration
|
||||
working-directory: ./hindsight-integrations/pipecat
|
||||
run: uv build
|
||||
|
||||
- name: Install dependencies
|
||||
working-directory: ./hindsight-integrations/pipecat
|
||||
run: uv sync --frozen
|
||||
|
||||
- name: Run tests
|
||||
working-directory: ./hindsight-integrations/pipecat
|
||||
run: uv run pytest tests -v
|
||||
|
||||
|
||||
build-control-plane:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
@@ -1007,6 +1135,128 @@ jobs:
|
||||
working-directory: ./hindsight-api-slim
|
||||
run: uv run pytest tests -v
|
||||
|
||||
test-api-oracle:
|
||||
needs: [detect-changes]
|
||||
# Gated behind the "oracle-tests" PR label so it doesn't run by default.
|
||||
# Add the label to any PR that needs Oracle validation.
|
||||
if: >-
|
||||
needs.detect-changes.outputs.has_secrets == 'true' &&
|
||||
contains(github.event.pull_request.labels.*.name, 'oracle-tests') &&
|
||||
((github.event_name == 'pull_request_review' && github.event.review.state == 'approved') ||
|
||||
github.event_name == 'workflow_dispatch' ||
|
||||
needs.detect-changes.outputs.core == 'true' ||
|
||||
needs.detect-changes.outputs.ci == '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
|
||||
GEMINI_API_KEY: ${{ secrets.GEMINI_API_KEY }}
|
||||
OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
|
||||
COHERE_API_KEY: ${{ secrets.COHERE_API_KEY }}
|
||||
HINDSIGHT_API_EMBEDDINGS_OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
|
||||
HINDSIGHT_API_DATABASE_BACKEND: oracle
|
||||
ORACLE_TEST_DSN: oracle+oracledb://hindsight_test:hindsight_test@localhost:1521/FREEPDB1
|
||||
|
||||
services:
|
||||
oracle:
|
||||
image: container-registry.oracle.com/database/free:latest
|
||||
env:
|
||||
ORACLE_PWD: oracle
|
||||
ports:
|
||||
- 1521:1521
|
||||
options: >-
|
||||
--health-cmd "echo 'SELECT 1 FROM DUAL;' | sqlplus -s system/oracle@localhost:1521/FREEPDB1 || exit 1"
|
||||
--health-interval 30s
|
||||
--health-timeout 10s
|
||||
--health-retries 10
|
||||
--health-start-period 120s
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ github.event.pull_request.head.sha || '' }}
|
||||
|
||||
- name: Setup Oracle test user
|
||||
# The SYSTEM tablespace uses manual segment space management which
|
||||
# doesn't support VECTOR types. Create an ASSM tablespace and a
|
||||
# dedicated test user so VECTOR columns work correctly.
|
||||
run: |
|
||||
pip install oracledb
|
||||
python3 -c "
|
||||
import oracledb
|
||||
conn = oracledb.connect(user='system', password='oracle', dsn='localhost:1521/FREEPDB1')
|
||||
cursor = conn.cursor()
|
||||
cursor.execute(\"\"\"
|
||||
CREATE TABLESPACE hindsight_ts
|
||||
DATAFILE 'hindsight_ts.dbf' SIZE 200M AUTOEXTEND ON NEXT 50M
|
||||
EXTENT MANAGEMENT LOCAL
|
||||
SEGMENT SPACE MANAGEMENT AUTO
|
||||
\"\"\")
|
||||
cursor.execute(\"\"\"
|
||||
CREATE USER hindsight_test IDENTIFIED BY hindsight_test
|
||||
DEFAULT TABLESPACE hindsight_ts
|
||||
TEMPORARY TABLESPACE temp
|
||||
QUOTA UNLIMITED ON hindsight_ts
|
||||
\"\"\")
|
||||
cursor.execute('GRANT CONNECT, RESOURCE, CREATE TABLE, CREATE SEQUENCE, CREATE VIEW, CREATE PROCEDURE TO hindsight_test')
|
||||
cursor.execute('GRANT CTXAPP TO hindsight_test')
|
||||
conn.commit()
|
||||
conn.close()
|
||||
print('Oracle test user created successfully')
|
||||
"
|
||||
|
||||
- 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: Build API
|
||||
working-directory: ./hindsight-api-slim
|
||||
run: uv build
|
||||
|
||||
- name: Install dependencies
|
||||
working-directory: ./hindsight-api-slim
|
||||
run: uv sync --frozen --all-extras --index-strategy unsafe-best-match
|
||||
|
||||
- name: Cache HuggingFace models
|
||||
uses: actions/cache@v5
|
||||
with:
|
||||
path: ~/.cache/huggingface
|
||||
key: ${{ runner.os }}-huggingface-${{ hashFiles('hindsight-api-slim/pyproject.toml') }}
|
||||
restore-keys: |
|
||||
${{ runner.os }}-huggingface-
|
||||
|
||||
- name: Pre-download models
|
||||
working-directory: ./hindsight-api-slim
|
||||
run: |
|
||||
uv run python -c "
|
||||
from sentence_transformers import SentenceTransformer, CrossEncoder
|
||||
print('Downloading embedding model...')
|
||||
SentenceTransformer('BAAI/bge-small-en-v1.5')
|
||||
print('Downloading cross-encoder model...')
|
||||
CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2')
|
||||
print('Models downloaded successfully')
|
||||
"
|
||||
|
||||
- name: Run Oracle tests
|
||||
working-directory: ./hindsight-api-slim
|
||||
# -n0: run sequentially to avoid ORA-00060 deadlocks from concurrent
|
||||
# test transactions against the same Oracle Free container.
|
||||
run: uv run pytest tests -v -m oracle -n0
|
||||
|
||||
test-python-client:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
@@ -1711,6 +1961,9 @@ jobs:
|
||||
- name: Build hindsight-all-npm (openclaw dep)
|
||||
run: npm run build --workspace=hindsight-all-npm
|
||||
|
||||
- name: Build hindsight-agent-sdk (openclaw dep)
|
||||
run: npm run build --workspace=hindsight-tools/hindsight-agent-sdk
|
||||
|
||||
- name: Install openclaw integration dependencies
|
||||
working-directory: ./hindsight-integrations/openclaw
|
||||
run: npm ci
|
||||
@@ -1895,6 +2148,43 @@ jobs:
|
||||
working-directory: ./hindsight-integrations/ag2
|
||||
run: uv run pytest tests -v
|
||||
|
||||
test-smolagents-integration:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
github.event_name != 'pull_request_review' &&
|
||||
(github.event_name == 'workflow_dispatch' ||
|
||||
needs.detect-changes.outputs.integrations-smolagents == '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 smolagents integration
|
||||
working-directory: ./hindsight-integrations/smolagents
|
||||
run: uv build
|
||||
|
||||
- name: Install dependencies
|
||||
working-directory: ./hindsight-integrations/smolagents
|
||||
run: uv sync --frozen
|
||||
|
||||
- name: Run tests
|
||||
working-directory: ./hindsight-integrations/smolagents
|
||||
run: uv run pytest tests -v
|
||||
|
||||
test-crewai-integration:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
@@ -2043,6 +2333,80 @@ 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-agentcore-integration:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
github.event_name != 'pull_request_review' &&
|
||||
(github.event_name == 'workflow_dispatch' ||
|
||||
needs.detect-changes.outputs.integrations-agentcore == '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 agentcore integration
|
||||
working-directory: ./hindsight-integrations/agentcore
|
||||
run: uv build
|
||||
|
||||
- name: Install dependencies
|
||||
working-directory: ./hindsight-integrations/agentcore
|
||||
run: uv sync --frozen
|
||||
|
||||
- name: Run tests
|
||||
working-directory: ./hindsight-integrations/agentcore
|
||||
run: uv run pytest tests -v
|
||||
|
||||
test-pip-slim:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
@@ -2175,6 +2539,189 @@ jobs:
|
||||
working-directory: ./hindsight-embed
|
||||
run: ./test.sh
|
||||
|
||||
test-embed-windows:
|
||||
# Windows coverage for hindsight-embed. Runs the same unit tests + smoke
|
||||
# test as the Linux `test-embed` job, plus a `uv pip install --target`
|
||||
# sanity check that validates the sibling-binary resolution used by
|
||||
# users who install via `uv pip install hindsight-all` on Windows
|
||||
# (closes #1240).
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
needs.detect-changes.outputs.has_secrets == 'true' &&
|
||||
((github.event_name == 'pull_request_review' && github.event.review.state == 'approved') ||
|
||||
github.event_name == 'workflow_dispatch' ||
|
||||
needs.detect-changes.outputs.core == 'true' ||
|
||||
needs.detect-changes.outputs.embed == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: windows-latest
|
||||
env:
|
||||
HINDSIGHT_API_LLM_PROVIDER: vertexai
|
||||
HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY: ${{ github.workspace }}/gcp-credentials.json
|
||||
HINDSIGHT_API_LLM_MODEL: google/gemini-2.5-flash-lite
|
||||
# Force UTF-8 I/O so the CLI's ✓/box-drawing output doesn't crash the
|
||||
# default Windows cp1252 codec. Also applied at runtime via
|
||||
# sys.stdout.reconfigure in cli.py; this belt-and-suspenders covers
|
||||
# subprocesses the daemon spawns.
|
||||
PYTHONIOENCODING: utf-8
|
||||
PYTHONUTF8: "1"
|
||||
# pg0-embedded unpacks Postgres on first boot — noticeably slower on a
|
||||
# cold Windows runner than POSIX. Double the embed startup budget.
|
||||
HINDSIGHT_EMBED_DAEMON_STARTUP_TIMEOUT: "360"
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ github.event.pull_request.head.sha || '' }}
|
||||
|
||||
- name: Setup GCP credentials
|
||||
shell: bash
|
||||
run: |
|
||||
printf '%s' '${{ secrets.GCP_VERTEXAI_CREDENTIALS }}' > gcp-credentials.json
|
||||
PROJECT_ID=$(jq -r '.project_id' gcp-credentials.json)
|
||||
echo "HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID=$PROJECT_ID" >> "$GITHUB_ENV"
|
||||
|
||||
- name: Install uv
|
||||
uses: astral-sh/setup-uv@v7
|
||||
with:
|
||||
enable-cache: true
|
||||
prune-cache: false
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v6
|
||||
with:
|
||||
python-version-file: ".python-version"
|
||||
|
||||
- name: Install embed dependencies
|
||||
working-directory: ./hindsight-embed
|
||||
run: uv sync --frozen --index-strategy unsafe-best-match
|
||||
|
||||
- name: Install API dependencies (with local-ml and embedded-db for smoke test)
|
||||
working-directory: ./hindsight-api-slim
|
||||
run: uv sync --frozen --all-extras --index-strategy unsafe-best-match
|
||||
|
||||
- name: Cache HuggingFace models
|
||||
uses: actions/cache@v5
|
||||
with:
|
||||
path: ~/.cache/huggingface
|
||||
key: ${{ runner.os }}-huggingface-embed-${{ hashFiles('hindsight-embed/pyproject.toml') }}
|
||||
restore-keys: |
|
||||
${{ runner.os }}-huggingface-embed-
|
||||
${{ runner.os }}-huggingface-
|
||||
|
||||
- name: Run unit and integration tests
|
||||
working-directory: ./hindsight-embed
|
||||
run: uv run pytest tests/ -v
|
||||
|
||||
# Smoke test's retain/recall commands delegate to the Rust hindsight CLI.
|
||||
# On POSIX, hindsight-embed auto-installs the CLI via curl|bash; on
|
||||
# Windows that installer isn't available (and `bash` on windows-latest
|
||||
# routes to WSL which isn't provisioned). Build the CLI from source and
|
||||
# drop it into ~/.local/bin where find_cli_binary() looks first.
|
||||
- name: Install Rust
|
||||
uses: dtolnay/rust-toolchain@stable
|
||||
|
||||
- name: Cache cargo build
|
||||
uses: actions/cache@v5
|
||||
with:
|
||||
path: |
|
||||
~/.cargo/registry
|
||||
~/.cargo/git
|
||||
hindsight-cli/target
|
||||
key: ${{ runner.os }}-cargo-embed-smoke-${{ hashFiles('**/Cargo.lock') }}
|
||||
restore-keys: |
|
||||
${{ runner.os }}-cargo-embed-smoke-
|
||||
${{ runner.os }}-cargo-
|
||||
|
||||
- name: Build hindsight CLI
|
||||
working-directory: ./hindsight-cli
|
||||
run: cargo build --release
|
||||
|
||||
- name: Stage hindsight CLI where find_cli_binary expects it
|
||||
shell: bash
|
||||
run: |
|
||||
set -euo pipefail
|
||||
install_dir="$HOME/.local/bin"
|
||||
mkdir -p "$install_dir"
|
||||
cp hindsight-cli/target/release/hindsight.exe "$install_dir/hindsight.exe"
|
||||
"$install_dir/hindsight.exe" --version
|
||||
|
||||
- name: Run smoke test
|
||||
shell: bash
|
||||
working-directory: ./hindsight-embed
|
||||
run: ./test.sh
|
||||
|
||||
# Real-world install test for issue #1240: drop both packages into a
|
||||
# --target directory (the layout you get from `uv pip install hindsight-all`
|
||||
# or NixOS) and verify the sibling binary is discovered (not the uvx
|
||||
# fallback). Exercises a different code path than the smoke test, which
|
||||
# uses `uv run --project` via the monorepo branch of _find_api_command.
|
||||
#
|
||||
# IMPORTANT: install outside the repo checkout. `_find_api_command` first
|
||||
# probes `<pkg>/../../hindsight-api-slim` for dev mode; if the target dir
|
||||
# lives inside the monorepo, that branch matches and we never exercise
|
||||
# the sibling-binary path we actually want to test.
|
||||
- name: Install hindsight-embed and hindsight-api into --target directory
|
||||
shell: bash
|
||||
run: |
|
||||
set -euo pipefail
|
||||
target="$RUNNER_TEMP/install-test"
|
||||
rm -rf "$target"
|
||||
mkdir -p "$target"
|
||||
uv pip install --target "$target" ./hindsight-embed ./hindsight-api-slim
|
||||
|
||||
- name: Verify sibling hindsight-api.exe is present
|
||||
shell: bash
|
||||
run: |
|
||||
set -euo pipefail
|
||||
target="$RUNNER_TEMP/install-test"
|
||||
if [ -f "$target/Scripts/hindsight-api.exe" ]; then
|
||||
echo "Found $target/Scripts/hindsight-api.exe"
|
||||
elif [ -f "$target/bin/hindsight-api.exe" ]; then
|
||||
echo "Found $target/bin/hindsight-api.exe"
|
||||
else
|
||||
echo "::error::hindsight-api.exe not found in install target"
|
||||
ls "$target/"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
- name: Verify _find_api_command resolves the sibling binary (not uvx)
|
||||
shell: bash
|
||||
run: |
|
||||
set -euo pipefail
|
||||
target="$RUNNER_TEMP/install-test"
|
||||
PYTHONPATH="$target" python -c "
|
||||
from hindsight_embed.daemon_embed_manager import DaemonEmbedManager
|
||||
cmd = DaemonEmbedManager()._find_api_command()
|
||||
print('Resolved command:', cmd)
|
||||
assert len(cmd) == 1 and cmd[0].endswith('hindsight-api.exe'), (
|
||||
f'Expected sibling hindsight-api.exe, got {cmd!r}. '
|
||||
'Falling back to uvx on --target installs reintroduces issue #1240.'
|
||||
)
|
||||
"
|
||||
|
||||
- name: Smoke-check installed hindsight-embed binary runs
|
||||
shell: bash
|
||||
run: |
|
||||
set -euo pipefail
|
||||
target="$RUNNER_TEMP/install-test"
|
||||
export PYTHONPATH="$target"
|
||||
if [ -f "$target/Scripts/hindsight-embed.exe" ]; then
|
||||
"$target/Scripts/hindsight-embed.exe" --help
|
||||
else
|
||||
"$target/bin/hindsight-embed.exe" --help
|
||||
fi
|
||||
|
||||
- name: Collect daemon logs on failure
|
||||
if: failure()
|
||||
shell: bash
|
||||
run: |
|
||||
for f in ~/.hindsight/daemon.log ~/.hindsight/profiles/*.log ~/.hindsight/profiles/*.stderr.log; do
|
||||
if [ -f "$f" ]; then
|
||||
echo "=== $f ==="
|
||||
cat "$f"
|
||||
fi
|
||||
done || true
|
||||
|
||||
test-hindsight-all:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
@@ -2547,6 +3094,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
|
||||
|
||||
@@ -2566,6 +3116,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"
|
||||
@@ -2681,12 +3232,14 @@ jobs:
|
||||
- test-cloudflare-oauth-proxy-integration
|
||||
- build-chat-integration
|
||||
- test-paperclip-integration
|
||||
- test-pipecat-integration
|
||||
- build-control-plane
|
||||
- build-docs
|
||||
- test-rust-cli
|
||||
- lint-helm-chart
|
||||
- build-docker-images
|
||||
- test-api
|
||||
- test-api-oracle
|
||||
- test-python-client
|
||||
- test-typescript-client
|
||||
- test-typescript-client-deno
|
||||
@@ -2696,13 +3249,18 @@ jobs:
|
||||
- test-openclaw-integration
|
||||
- test-integration
|
||||
- test-ag2-integration
|
||||
- test-smolagents-integration
|
||||
- test-crewai-integration
|
||||
- test-litellm-integration
|
||||
- test-pydantic-ai-integration
|
||||
- test-llamaindex-integration
|
||||
- test-agentcore-integration
|
||||
- test-pip-slim
|
||||
- test-embed
|
||||
- test-embed-windows
|
||||
- test-hindsight-all
|
||||
- test-hindsight-agent-sdk
|
||||
- test-self-driving-agents
|
||||
- test-doc-examples
|
||||
- test-upgrade
|
||||
- verify-generated-files
|
||||
@@ -2715,7 +3273,7 @@ jobs:
|
||||
steps:
|
||||
- name: Determine overall result
|
||||
id: result
|
||||
uses: actions/github-script@v8
|
||||
uses: actions/github-script@v9
|
||||
with:
|
||||
script: |
|
||||
const needs = ${{ toJSON(needs) }};
|
||||
@@ -2748,7 +3306,7 @@ jobs:
|
||||
core.setOutput('run_url', runUrl);
|
||||
|
||||
- name: Report status to PR
|
||||
uses: actions/github-script@v8
|
||||
uses: actions/github-script@v9
|
||||
with:
|
||||
script: |
|
||||
await github.rest.repos.createCommitStatus({
|
||||
@@ -2762,7 +3320,7 @@ jobs:
|
||||
});
|
||||
|
||||
- name: Comment on PR
|
||||
uses: actions/github-script@v8
|
||||
uses: actions/github-script@v9
|
||||
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
|
||||
|
||||
@@ -2,8 +2,8 @@ apiVersion: v2
|
||||
name: hindsight
|
||||
description: Hindsight helm chart
|
||||
type: application
|
||||
version: 0.5.1
|
||||
appVersion: "0.5.1"
|
||||
version: 0.5.6
|
||||
appVersion: "0.5.6"
|
||||
keywords:
|
||||
- ai
|
||||
- memory
|
||||
|
||||
+10
-10
@@ -18,17 +18,17 @@ 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';
|
||||
import { HindsightServer, consoleLogger } from "@vectorize-io/hindsight-all";
|
||||
import { HindsightClient } from "@vectorize-io/hindsight-client";
|
||||
|
||||
const server = new HindsightServer({
|
||||
profile: 'my-app',
|
||||
profile: "my-app",
|
||||
port: 9077,
|
||||
env: {
|
||||
HINDSIGHT_API_LLM_PROVIDER: 'anthropic',
|
||||
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',
|
||||
HINDSIGHT_API_LLM_MODEL: "claude-sonnet-4-20250514",
|
||||
HINDSIGHT_EMBED_DAEMON_IDLE_TIMEOUT: "0",
|
||||
},
|
||||
logger: consoleLogger,
|
||||
});
|
||||
@@ -37,11 +37,11 @@ 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',
|
||||
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?');
|
||||
const recall = await client.recall("user-123", "what are the user preferences?");
|
||||
console.log(recall.results);
|
||||
|
||||
await server.stop();
|
||||
@@ -62,7 +62,7 @@ If you're hacking on the Python `hindsight-embed` package in the same monorepo,
|
||||
|
||||
```ts
|
||||
new HindsightServer({
|
||||
embedPackagePath: '/path/to/hindsight-embed',
|
||||
embedPackagePath: "/path/to/hindsight-embed",
|
||||
// ...
|
||||
});
|
||||
```
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@vectorize-io/hindsight-all",
|
||||
"version": "0.5.1",
|
||||
"version": "0.5.6",
|
||||
"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",
|
||||
|
||||
@@ -1,32 +1,36 @@
|
||||
import { describe, it, expect } from 'vitest';
|
||||
import { getEmbedCommand } from './command.js';
|
||||
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']);
|
||||
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("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("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("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']);
|
||||
it("local path takes precedence over version", () => {
|
||||
expect(getEmbedCommand({ embedPackagePath: "/abs/path", embedVersion: "0.5.0" })).toEqual([
|
||||
"uv",
|
||||
"run",
|
||||
"--directory",
|
||||
"/abs/path",
|
||||
"hindsight-embed",
|
||||
]);
|
||||
});
|
||||
});
|
||||
|
||||
@@ -18,8 +18,8 @@ export interface EmbedCommandOptions {
|
||||
|
||||
export function getEmbedCommand(opts: EmbedCommandOptions = {}): string[] {
|
||||
if (opts.embedPackagePath) {
|
||||
return ['uv', 'run', '--directory', opts.embedPackagePath, 'hindsight-embed'];
|
||||
return ["uv", "run", "--directory", opts.embedPackagePath, "hindsight-embed"];
|
||||
}
|
||||
const version = opts.embedVersion && opts.embedVersion.length > 0 ? opts.embedVersion : 'latest';
|
||||
return ['uvx', `hindsight-embed@${version}`];
|
||||
const version = opts.embedVersion && opts.embedVersion.length > 0 ? opts.embedVersion : "latest";
|
||||
return ["uvx", `hindsight-embed@${version}`];
|
||||
}
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
export { HindsightServer } from './server.js';
|
||||
export { getEmbedCommand } from './command.js';
|
||||
export { silentLogger, consoleLogger } from './logger.js';
|
||||
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';
|
||||
export type { Logger } from "./logger.js";
|
||||
export type { EmbedCommandOptions } from "./command.js";
|
||||
export type { HindsightServerOptions } from "./types.js";
|
||||
|
||||
@@ -1,32 +1,32 @@
|
||||
import { describe, it, expect } from 'vitest';
|
||||
import { HindsightServer } from './server.js';
|
||||
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', () => {
|
||||
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');
|
||||
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("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', () => {
|
||||
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',
|
||||
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',
|
||||
HINDSIGHT_FUTURE_FLAG: "enabled",
|
||||
},
|
||||
});
|
||||
expect(server).toBeInstanceOf(HindsightServer);
|
||||
});
|
||||
|
||||
it('exposes checkHealth that returns false when no daemon is running', async () => {
|
||||
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();
|
||||
|
||||
@@ -1,12 +1,12 @@
|
||||
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';
|
||||
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_HOST = "127.0.0.1";
|
||||
const DEFAULT_PROFILE = "default";
|
||||
const DEFAULT_READY_TIMEOUT_MS = 30_000;
|
||||
const DEFAULT_READY_POLL_INTERVAL_MS = 1_000;
|
||||
|
||||
@@ -61,7 +61,7 @@ export class HindsightServer {
|
||||
this.userEnv = opts.env ?? {};
|
||||
this.extraProfileCreateArgs = opts.extraProfileCreateArgs ?? [];
|
||||
this.extraDaemonStartArgs = opts.extraDaemonStartArgs ?? [];
|
||||
this.platformCpuWorkaround = opts.platformCpuWorkaround ?? (process.platform === 'darwin');
|
||||
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;
|
||||
@@ -100,22 +100,22 @@ export class HindsightServer {
|
||||
embedVersion: this.embedVersion,
|
||||
embedPackagePath: this.embedPackagePath,
|
||||
});
|
||||
const args = [...baseArgs, 'daemon', '--profile', this.profile, 'stop'];
|
||||
const args = [...baseArgs, "daemon", "--profile", this.profile, "stop"];
|
||||
|
||||
const child = spawn(cmd, args, { stdio: 'pipe' });
|
||||
this.pipeOutput(child, 'daemon.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', () => {
|
||||
child.on("exit", () => {
|
||||
clearTimeout(timeout);
|
||||
this.logger.info(`[hindsight] daemon stopped`);
|
||||
resolve();
|
||||
});
|
||||
child.on('error', (err) => {
|
||||
child.on("error", (err) => {
|
||||
clearTimeout(timeout);
|
||||
this.logger.warn(`[hindsight] error stopping daemon: ${err.message}`);
|
||||
resolve();
|
||||
@@ -147,9 +147,9 @@ export class HindsightServer {
|
||||
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';
|
||||
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)) {
|
||||
@@ -175,11 +175,11 @@ export class HindsightServer {
|
||||
});
|
||||
const createArgs = [
|
||||
...baseArgs,
|
||||
'profile',
|
||||
'create',
|
||||
"profile",
|
||||
"create",
|
||||
this.profile,
|
||||
'--merge',
|
||||
'--port',
|
||||
"--merge",
|
||||
"--port",
|
||||
String(this.port),
|
||||
];
|
||||
|
||||
@@ -189,12 +189,12 @@ export class HindsightServer {
|
||||
// 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("--env", `${key}=${value}`);
|
||||
}
|
||||
|
||||
createArgs.push(...this.extraProfileCreateArgs);
|
||||
|
||||
await this.runCommand(cmd, createArgs, env, 'profile.create');
|
||||
await this.runCommand(cmd, createArgs, env, "profile.create");
|
||||
}
|
||||
|
||||
/** Collect only the env vars that should be written into the profile file. */
|
||||
@@ -209,10 +209,10 @@ export class HindsightServer {
|
||||
}
|
||||
|
||||
// 2. CPU workaround — only if auto-applied and not already overridden.
|
||||
if (this.platformCpuWorkaround && process.platform === 'darwin') {
|
||||
if (this.platformCpuWorkaround && process.platform === "darwin") {
|
||||
const cpuKeys = [
|
||||
'HINDSIGHT_API_EMBEDDINGS_LOCAL_FORCE_CPU',
|
||||
'HINDSIGHT_API_RERANKER_LOCAL_FORCE_CPU',
|
||||
"HINDSIGHT_API_EMBEDDINGS_LOCAL_FORCE_CPU",
|
||||
"HINDSIGHT_API_RERANKER_LOCAL_FORCE_CPU",
|
||||
];
|
||||
for (const key of cpuKeys) {
|
||||
if (!(key in out) && env[key] !== undefined) {
|
||||
@@ -231,14 +231,14 @@ export class HindsightServer {
|
||||
});
|
||||
const args = [
|
||||
...baseArgs,
|
||||
'daemon',
|
||||
'--profile',
|
||||
"daemon",
|
||||
"--profile",
|
||||
this.profile,
|
||||
'start',
|
||||
"start",
|
||||
...this.extraDaemonStartArgs,
|
||||
];
|
||||
|
||||
await this.runCommand(cmd, args, env, 'daemon.start');
|
||||
await this.runCommand(cmd, args, env, "daemon.start");
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -249,34 +249,34 @@ export class HindsightServer {
|
||||
cmd: string,
|
||||
args: string[],
|
||||
env: NodeJS.ProcessEnv,
|
||||
label: string,
|
||||
label: string
|
||||
): Promise<void> {
|
||||
const child = spawn(cmd, args, { stdio: 'pipe', env });
|
||||
let output = '';
|
||||
child.stdout?.on('data', (data: Buffer) => {
|
||||
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')) {
|
||||
for (const line of text.trimEnd().split("\n")) {
|
||||
if (line) this.logger.info(`[hindsight:${label}] ${line}`);
|
||||
}
|
||||
});
|
||||
child.stderr?.on('data', (data: Buffer) => {
|
||||
child.stderr?.on("data", (data: Buffer) => {
|
||||
const text = data.toString();
|
||||
output += text;
|
||||
for (const line of text.trimEnd().split('\n')) {
|
||||
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) => {
|
||||
child.on("exit", (code) => {
|
||||
if (code === 0) {
|
||||
resolve();
|
||||
} else {
|
||||
reject(new Error(`${label} failed with code ${code}: ${output.trim()}`));
|
||||
}
|
||||
});
|
||||
child.on('error', (err) => {
|
||||
child.on("error", (err) => {
|
||||
reject(new Error(`${label} failed to spawn: ${err.message}`, { cause: err }));
|
||||
});
|
||||
});
|
||||
@@ -284,13 +284,13 @@ export class HindsightServer {
|
||||
|
||||
/** 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')) {
|
||||
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')) {
|
||||
child.stderr?.on("data", (data: Buffer) => {
|
||||
for (const line of data.toString().trimEnd().split("\n")) {
|
||||
if (line) this.logger.warn(`[hindsight:${label}] ${line}`);
|
||||
}
|
||||
});
|
||||
@@ -316,7 +316,7 @@ export class HindsightServer {
|
||||
await new Promise((resolve) => setTimeout(resolve, this.readyPollIntervalMs));
|
||||
}
|
||||
throw new Error(
|
||||
`Hindsight daemon did not become ready within ${this.readyTimeoutMs}ms at ${this.baseUrl}`,
|
||||
`Hindsight daemon did not become ready within ${this.readyTimeoutMs}ms at ${this.baseUrl}`
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
import type { Logger } from './logger.js';
|
||||
import type { Logger } from "./logger.js";
|
||||
|
||||
/**
|
||||
* Options for {@link HindsightServer}.
|
||||
|
||||
@@ -1,10 +1,10 @@
|
||||
import { defineConfig } from 'tsup';
|
||||
import { defineConfig } from "tsup";
|
||||
|
||||
export default defineConfig({
|
||||
entry: ['src/index.ts'],
|
||||
format: ['esm'],
|
||||
entry: ["src/index.ts"],
|
||||
format: ["esm"],
|
||||
dts: true,
|
||||
outDir: 'dist',
|
||||
outDir: "dist",
|
||||
clean: true,
|
||||
sourcemap: true,
|
||||
bundle: true,
|
||||
|
||||
@@ -1,8 +1,8 @@
|
||||
import { defineConfig } from 'vitest/config';
|
||||
import { defineConfig } from "vitest/config";
|
||||
|
||||
export default defineConfig({
|
||||
test: {
|
||||
include: ['src/**/*.test.ts'],
|
||||
environment: 'node',
|
||||
include: ["src/**/*.test.ts"],
|
||||
environment: "node",
|
||||
},
|
||||
});
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "hindsight-all-slim"
|
||||
version = "0.5.1"
|
||||
version = "0.5.6"
|
||||
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(
|
||||
@@ -253,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."""
|
||||
@@ -394,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
|
||||
@@ -409,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()
|
||||
@@ -425,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.5.1"
|
||||
version = "0.5.6"
|
||||
description = "Hindsight: Agent Memory That Works Like Human Memory - All-in-One Bundle"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.11"
|
||||
|
||||
@@ -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.5.1"
|
||||
__version__ = "0.5.6"
|
||||
|
||||
@@ -1,5 +1,7 @@
|
||||
"""
|
||||
Hindsight Admin CLI - backup and restore operations.
|
||||
"""PostgreSQL-only admin utilities (backup, restore, migration, worker management).
|
||||
|
||||
Not supported on Oracle backends. Uses asyncpg.connect() directly, binary COPY,
|
||||
TRUNCATE CASCADE, and REFRESH MATERIALIZED VIEW — all inherently PG-specific.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
@@ -15,15 +17,10 @@ import asyncpg
|
||||
import typer
|
||||
|
||||
from ..config import DEFAULT_DATABASE_SCHEMA, HindsightConfig
|
||||
from ..engine.schema import fq_table_explicit as _fq_table
|
||||
from ..extensions import TenantExtension, load_extension
|
||||
from ..pg0 import parse_pg0_url, resolve_database_url
|
||||
|
||||
|
||||
def _fq_table(table: str, schema: str) -> str:
|
||||
"""Get fully-qualified table name with schema prefix."""
|
||||
return f"{schema}.{table}"
|
||||
|
||||
|
||||
# Setup logging
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
@@ -375,6 +372,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
|
||||
+57
@@ -0,0 +1,57 @@
|
||||
"""Backfill mental_models.subtype for databases that ran h3c4d5e6f7g8 before the fix
|
||||
|
||||
Migration h3c4d5e6f7g8 used CREATE TABLE IF NOT EXISTS to create the
|
||||
mental_models table with a subtype column. But on databases where the table
|
||||
already existed (from the reflections -> mental_models rename chain), the
|
||||
CREATE was a no-op and subtype was never added. A fix was later added to
|
||||
h3c4d5e6f7g8 (Step 4b), but databases that had already run the migration
|
||||
never re-execute it. This migration adds the missing columns idempotently.
|
||||
|
||||
Revision ID: d5y6z7a8b9c0
|
||||
Revises: 8c6fa6f7230b
|
||||
Create Date: 2026-04-18
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
revision: str = "d5y6z7a8b9c0"
|
||||
down_revision: str | Sequence[str] | None = "8c6fa6f7230b"
|
||||
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()
|
||||
|
||||
# Add columns that h3c4d5e6f7g8 intended to create but missed when
|
||||
# the table already existed from the reflections rename chain.
|
||||
for col_ddl in [
|
||||
"subtype VARCHAR(32) NOT NULL DEFAULT 'structural'",
|
||||
"description TEXT NOT NULL DEFAULT ''",
|
||||
"entity_id UUID",
|
||||
"observations JSONB DEFAULT '{\"observations\": []}'::jsonb",
|
||||
"links VARCHAR[]",
|
||||
"last_updated TIMESTAMP WITH TIME ZONE",
|
||||
]:
|
||||
op.execute(f"ALTER TABLE {schema}mental_models ADD COLUMN IF NOT EXISTS {col_ddl}")
|
||||
|
||||
# Ensure the CHECK constraint exists
|
||||
op.execute(f"ALTER TABLE {schema}mental_models DROP CONSTRAINT IF EXISTS ck_mental_models_subtype")
|
||||
op.execute(f"""
|
||||
ALTER TABLE {schema}mental_models
|
||||
ADD CONSTRAINT ck_mental_models_subtype CHECK (subtype IN ('structural', 'emergent', 'pinned', 'learned'))
|
||||
""")
|
||||
|
||||
op.execute(f"CREATE INDEX IF NOT EXISTS idx_mental_models_subtype ON {schema}mental_models(bank_id, subtype)")
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
# No-op: these columns are part of the intended schema
|
||||
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 '{{}}'")
|
||||
@@ -85,6 +85,26 @@ def upgrade() -> None:
|
||||
)
|
||||
""")
|
||||
|
||||
# Step 4b: If the table already existed (from reflections rename chain),
|
||||
# it won't have the v4 columns. Add them idempotently so the migration
|
||||
# works regardless of whether CREATE TABLE above was a no-op.
|
||||
for col_ddl in [
|
||||
"subtype VARCHAR(32) NOT NULL DEFAULT 'directive'",
|
||||
"description TEXT NOT NULL DEFAULT ''",
|
||||
"entity_id UUID",
|
||||
"observations JSONB DEFAULT '{\"observations\": []}'::jsonb",
|
||||
"links VARCHAR[]",
|
||||
"last_updated TIMESTAMP WITH TIME ZONE",
|
||||
]:
|
||||
op.execute(f"ALTER TABLE {schema}mental_models ADD COLUMN IF NOT EXISTS {col_ddl}")
|
||||
|
||||
# Ensure the subtype CHECK constraint exists (may not if table was renamed)
|
||||
op.execute(f"ALTER TABLE {schema}mental_models DROP CONSTRAINT IF EXISTS ck_mental_models_subtype")
|
||||
op.execute(f"""
|
||||
ALTER TABLE {schema}mental_models
|
||||
ADD CONSTRAINT ck_mental_models_subtype CHECK (subtype IN ('structural', 'emergent', 'pinned', 'learned'))
|
||||
""")
|
||||
|
||||
# Step 5: Create indexes for efficient queries (if not exist)
|
||||
op.execute(f"CREATE INDEX IF NOT EXISTS idx_mental_models_bank_id ON {schema}mental_models(bank_id)")
|
||||
op.execute(f"CREATE INDEX IF NOT EXISTS idx_mental_models_subtype ON {schema}mental_models(bank_id, subtype)")
|
||||
|
||||
+2
-2
@@ -1,7 +1,7 @@
|
||||
"""Merge 3 migration heads and add unit_entities composite index
|
||||
|
||||
Revision ID: h3i4j5k6l7m8
|
||||
Revises: a4b5c6d7e8f9, c2d3e4f5g6h7, g2h3i4j5k6l7
|
||||
Revises: a4b5c6d7e8f9, g2h3i4j5k6l7
|
||||
Create Date: 2026-04-07
|
||||
|
||||
Merges three unmerged migration heads into one, and adds a composite index
|
||||
@@ -14,7 +14,7 @@ from collections.abc import Sequence
|
||||
from alembic import context, op
|
||||
|
||||
revision: str = "h3i4j5k6l7m8"
|
||||
down_revision: str | Sequence[str] | None = ("a4b5c6d7e8f9", "c2d3e4f5g6h7", "g2h3i4j5k6l7")
|
||||
down_revision: str | Sequence[str] | None = ("a4b5c6d7e8f9", "g2h3i4j5k6l7")
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
+39
@@ -0,0 +1,39 @@
|
||||
"""Add 'cancelled' to async_operations status check constraint
|
||||
|
||||
Revision ID: i4j5k6l7m8n9
|
||||
Revises: d5y6z7a8b9c0
|
||||
Create Date: 2026-04-23
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
revision: str = "i4j5k6l7m8n9"
|
||||
down_revision: str | Sequence[str] | None = "d5y6z7a8b9c0"
|
||||
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}async_operations DROP CONSTRAINT IF EXISTS async_operations_status_check")
|
||||
op.execute(
|
||||
f"ALTER TABLE {schema}async_operations ADD CONSTRAINT async_operations_status_check "
|
||||
f"CHECK (status IN ('pending', 'processing', 'completed', 'failed', 'cancelled'))"
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
op.execute(f"ALTER TABLE {schema}async_operations DROP CONSTRAINT IF EXISTS async_operations_status_check")
|
||||
op.execute(
|
||||
f"ALTER TABLE {schema}async_operations ADD CONSTRAINT async_operations_status_check "
|
||||
f"CHECK (status IN ('pending', 'processing', 'completed', 'failed'))"
|
||||
)
|
||||
+70
@@ -0,0 +1,70 @@
|
||||
"""Create observation_sources junction table
|
||||
|
||||
Replaces the source_memory_ids UUID[] column (PG) / CLOB (Oracle) with a
|
||||
proper junction table. This eliminates dialect-specific array operators
|
||||
(&&, unnest, JSON_TABLE) and enables standard SQL joins for all backends.
|
||||
|
||||
The old source_memory_ids column is retained for now (dual-write) and will
|
||||
be dropped in a future migration once all read paths are migrated.
|
||||
|
||||
Revision ID: k6l7m8n9o0p1
|
||||
Revises: i4j5k6l7m8n9
|
||||
Create Date: 2026-04-24
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
revision: str = "k6l7m8n9o0p1"
|
||||
down_revision: str | Sequence[str] | None = "i4j5k6l7m8n9"
|
||||
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()
|
||||
|
||||
# Create junction table.
|
||||
# observation_id has ON DELETE CASCADE so deleting an observation cleans up its rows.
|
||||
# source_id intentionally has NO FK — when a source memory is deleted, we need
|
||||
# observation_sources rows to still exist so delete_stale_observations_for_memories()
|
||||
# can find affected observations. Those observations are then deleted, which cascades
|
||||
# to observation_sources via the observation_id FK.
|
||||
op.execute(f"""
|
||||
CREATE TABLE IF NOT EXISTS {schema}observation_sources (
|
||||
observation_id UUID NOT NULL,
|
||||
source_id UUID NOT NULL,
|
||||
PRIMARY KEY (observation_id, source_id),
|
||||
FOREIGN KEY (observation_id) REFERENCES {schema}memory_units(id) ON DELETE CASCADE
|
||||
)
|
||||
""")
|
||||
|
||||
# Index on source_id for reverse lookups (find observations referencing a given source)
|
||||
op.execute(f"""
|
||||
CREATE INDEX IF NOT EXISTS idx_obs_sources_source_id
|
||||
ON {schema}observation_sources(source_id, observation_id)
|
||||
""")
|
||||
|
||||
# Backfill from existing source_memory_ids array column
|
||||
op.execute(f"""
|
||||
INSERT INTO {schema}observation_sources (observation_id, source_id)
|
||||
SELECT mu.id, unnest(mu.source_memory_ids)
|
||||
FROM {schema}memory_units mu
|
||||
WHERE mu.fact_type = 'observation'
|
||||
AND mu.source_memory_ids IS NOT NULL
|
||||
AND array_length(mu.source_memory_ids, 1) > 0
|
||||
ON CONFLICT DO NOTHING
|
||||
""")
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema}idx_obs_sources_source_id")
|
||||
op.execute(f"DROP TABLE IF EXISTS {schema}observation_sources")
|
||||
+4
-2
@@ -80,11 +80,13 @@ def upgrade() -> None:
|
||||
# 4. Drop the mental_model_versions table (no longer used)
|
||||
op.execute(f"DROP TABLE IF EXISTS {schema}mental_model_versions CASCADE")
|
||||
|
||||
# 5. Drop old constraints and add new one that only allows 'directive'
|
||||
# 5. Drop old constraints and add new one that allows current subtypes.
|
||||
# 'pinned' is still used by the code for user-created mental models;
|
||||
# 'directive' is used for system directives.
|
||||
op.execute(f"ALTER TABLE {schema}mental_models DROP CONSTRAINT IF EXISTS ck_mental_models_subtype")
|
||||
op.execute(f"""
|
||||
ALTER TABLE {schema}mental_models
|
||||
ADD CONSTRAINT ck_mental_models_subtype CHECK (subtype = 'directive')
|
||||
ADD CONSTRAINT ck_mental_models_subtype CHECK (subtype IN ('directive', 'pinned'))
|
||||
""")
|
||||
|
||||
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -8,6 +8,7 @@ from contextvars import ContextVar
|
||||
from fastmcp import FastMCP
|
||||
|
||||
from hindsight_api import MemoryEngine
|
||||
from hindsight_api import __version__ as HINDSIGHT_VERSION
|
||||
from hindsight_api.config import _get_raw_config
|
||||
from hindsight_api.engine.memory_engine import _current_schema
|
||||
from hindsight_api.extensions import MCPExtension, load_extension
|
||||
@@ -89,7 +90,7 @@ def create_mcp_server(memory: MemoryEngine, multi_bank: bool = True) -> FastMCP:
|
||||
Returns:
|
||||
Configured FastMCP server instance
|
||||
"""
|
||||
mcp = FastMCP("hindsight-mcp-server")
|
||||
mcp = FastMCP("hindsight-mcp-server", version=HINDSIGHT_VERSION)
|
||||
|
||||
global_config = _get_raw_config()
|
||||
|
||||
@@ -111,7 +112,6 @@ def create_mcp_server(memory: MemoryEngine, multi_bank: bool = True) -> FastMCP:
|
||||
"delete_directive",
|
||||
"list_memories",
|
||||
"get_memory",
|
||||
"delete_memory",
|
||||
"list_documents",
|
||||
"get_document",
|
||||
"delete_document",
|
||||
|
||||
@@ -10,10 +10,12 @@ import os
|
||||
import sys
|
||||
from dataclasses import dataclass, field, fields
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any
|
||||
from typing import Any, Literal
|
||||
|
||||
from dotenv import find_dotenv, load_dotenv
|
||||
|
||||
from .utils import mask_network_location
|
||||
|
||||
# Load .env file, searching current and parent directories (overrides existing env vars)
|
||||
load_dotenv(find_dotenv(usecwd=True), override=True)
|
||||
|
||||
@@ -117,6 +119,7 @@ def normalize_config_dict(config: dict[str, Any]) -> dict[str, Any]:
|
||||
|
||||
|
||||
# Environment variable names
|
||||
ENV_DATABASE_BACKEND = "HINDSIGHT_API_DATABASE_BACKEND"
|
||||
ENV_DATABASE_URL = "HINDSIGHT_API_DATABASE_URL"
|
||||
ENV_MIGRATION_DATABASE_URL = "HINDSIGHT_API_MIGRATION_DATABASE_URL"
|
||||
ENV_DATABASE_SCHEMA = "HINDSIGHT_API_DATABASE_SCHEMA"
|
||||
@@ -177,11 +180,13 @@ 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_GEMINI_FORCE_IPV4 = "HINDSIGHT_API_EMBEDDINGS_GEMINI_FORCE_IPV4"
|
||||
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"
|
||||
@@ -190,6 +195,7 @@ ENV_EMBEDDINGS_VERTEXAI_SERVICE_ACCOUNT_KEY = "HINDSIGHT_API_EMBEDDINGS_VERTEXAI
|
||||
ENV_EMBEDDINGS_COHERE_API_KEY = "HINDSIGHT_API_EMBEDDINGS_COHERE_API_KEY"
|
||||
ENV_EMBEDDINGS_COHERE_MODEL = "HINDSIGHT_API_EMBEDDINGS_COHERE_MODEL"
|
||||
ENV_EMBEDDINGS_COHERE_BASE_URL = "HINDSIGHT_API_EMBEDDINGS_COHERE_BASE_URL"
|
||||
ENV_EMBEDDINGS_COHERE_OUTPUT_DIMENSIONS = "HINDSIGHT_API_EMBEDDINGS_COHERE_OUTPUT_DIMENSIONS"
|
||||
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"
|
||||
@@ -238,9 +244,11 @@ 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"
|
||||
ENV_RERANKER_FLASHRANK_CPU_MEM_ARENA = "HINDSIGHT_API_RERANKER_FLASHRANK_CPU_MEM_ARENA"
|
||||
|
||||
# ZeroEntropy configuration (reranker only)
|
||||
ENV_RERANKER_ZEROENTROPY_API_KEY = "HINDSIGHT_API_RERANKER_ZEROENTROPY_API_KEY"
|
||||
@@ -265,6 +273,7 @@ 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"
|
||||
@@ -325,6 +334,7 @@ ENV_FILE_PARSER = "HINDSIGHT_API_FILE_PARSER"
|
||||
ENV_FILE_PARSER_ALLOWLIST = "HINDSIGHT_API_FILE_PARSER_ALLOWLIST"
|
||||
ENV_FILE_PARSER_IRIS_TOKEN = "HINDSIGHT_API_FILE_PARSER_IRIS_TOKEN"
|
||||
ENV_FILE_PARSER_IRIS_ORG_ID = "HINDSIGHT_API_FILE_PARSER_IRIS_ORG_ID"
|
||||
ENV_FILE_PARSER_LLAMA_PARSE_API_KEY = "HINDSIGHT_API_FILE_PARSER_LLAMA_PARSE_API_KEY"
|
||||
ENV_FILE_CONVERSION_MAX_BATCH_SIZE_MB = "HINDSIGHT_API_FILE_CONVERSION_MAX_BATCH_SIZE_MB"
|
||||
ENV_FILE_CONVERSION_MAX_BATCH_SIZE = "HINDSIGHT_API_FILE_CONVERSION_MAX_BATCH_SIZE"
|
||||
ENV_ENABLE_FILE_UPLOAD_API = "HINDSIGHT_API_ENABLE_FILE_UPLOAD_API"
|
||||
@@ -333,12 +343,15 @@ 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_RECALL_BUDGET = "HINDSIGHT_API_CONSOLIDATION_RECALL_BUDGET"
|
||||
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"
|
||||
@@ -370,6 +383,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"
|
||||
@@ -378,7 +392,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
|
||||
@@ -387,6 +412,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"
|
||||
@@ -399,6 +438,7 @@ ENV_DISPOSITION_LITERALISM = "HINDSIGHT_API_DISPOSITION_LITERALISM"
|
||||
ENV_DISPOSITION_EMPATHY = "HINDSIGHT_API_DISPOSITION_EMPATHY"
|
||||
|
||||
# Default values
|
||||
DEFAULT_DATABASE_BACKEND = "postgresql"
|
||||
DEFAULT_DATABASE_URL = "pg0"
|
||||
DEFAULT_DATABASE_SCHEMA = "public"
|
||||
DEFAULT_LLM_PROVIDER = "openai"
|
||||
@@ -410,6 +450,7 @@ PROVIDER_DEFAULT_MODELS = {
|
||||
"gemini": "gemini-2.5-flash",
|
||||
"groq": "openai/gpt-oss-120b",
|
||||
"minimax": "MiniMax-M2.7",
|
||||
"deepseek": "deepseek-v4-flash",
|
||||
"ollama": "gemma3:12b",
|
||||
"llamacpp": "gemma-4-e2b-it",
|
||||
"lmstudio": "local-model",
|
||||
@@ -432,7 +473,7 @@ DEFAULT_LLAMACPP_NO_GRAMMAR = False # True = disable JSON grammar enforcement (
|
||||
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
|
||||
@@ -450,8 +491,10 @@ 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_EMBEDDINGS_GEMINI_FORCE_IPV4 = False
|
||||
DEFAULT_EMBEDDING_DIMENSION = 384
|
||||
|
||||
DEFAULT_RERANKER_PROVIDER = "local"
|
||||
@@ -466,9 +509,11 @@ 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
|
||||
DEFAULT_RERANKER_FLASHRANK_CPU_MEM_ARENA = False # Disable ONNX CPU memory arena to bound RSS
|
||||
|
||||
DEFAULT_EMBEDDINGS_COHERE_MODEL = "embed-english-v3.0"
|
||||
DEFAULT_RERANKER_COHERE_MODEL = "rerank-english-v3.0"
|
||||
@@ -551,12 +596,17 @@ 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_RECALL_BUDGET = "low" # Budget level for consolidation recall (low/mid/high)
|
||||
DEFAULT_CONSOLIDATION_SOURCE_FACTS_MAX_TOKENS = (
|
||||
-1
|
||||
) # Total token budget for source facts in consolidation recall (-1 = unlimited)
|
||||
4096 # Total token budget for source facts in consolidation recall (-1 = unlimited)
|
||||
)
|
||||
DEFAULT_CONSOLIDATION_SOURCE_FACTS_MAX_TOKENS_PER_OBSERVATION = (
|
||||
256 # Max tokens of source facts per observation in consolidation prompt (-1 = unlimited)
|
||||
)
|
||||
@@ -571,6 +621,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)
|
||||
@@ -579,7 +630,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
|
||||
@@ -587,6 +637,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
|
||||
@@ -649,6 +718,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"),
|
||||
@@ -657,10 +730,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)
|
||||
|
||||
|
||||
@@ -669,6 +752,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()
|
||||
@@ -681,6 +783,18 @@ 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)
|
||||
@@ -711,6 +825,7 @@ class HindsightConfig:
|
||||
"""Configuration container for Hindsight API."""
|
||||
|
||||
# Database
|
||||
database_backend: Literal["postgresql", "oracle"]
|
||||
database_url: str
|
||||
migration_database_url: str | None
|
||||
database_schema: str
|
||||
@@ -790,6 +905,7 @@ class HindsightConfig:
|
||||
embeddings_cohere_api_key: str | None
|
||||
embeddings_cohere_model: str
|
||||
embeddings_cohere_base_url: str | None
|
||||
embeddings_cohere_output_dimensions: int | None
|
||||
embeddings_openrouter_api_key: str | None
|
||||
embeddings_openrouter_model: str
|
||||
embeddings_litellm_api_base: str
|
||||
@@ -804,6 +920,7 @@ class HindsightConfig:
|
||||
embeddings_gemini_api_key: str | None
|
||||
embeddings_gemini_model: str
|
||||
embeddings_gemini_output_dimensionality: int | None
|
||||
embeddings_gemini_force_ipv4: bool
|
||||
embeddings_vertexai_project_id: str | None
|
||||
embeddings_vertexai_region: str | None
|
||||
embeddings_vertexai_service_account_key: str | None
|
||||
@@ -820,6 +937,7 @@ 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
|
||||
@@ -849,6 +967,7 @@ 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)
|
||||
@@ -897,6 +1016,7 @@ class HindsightConfig:
|
||||
file_parser_allowlist: list[str] | None # Parsers clients may request (None = all registered)
|
||||
file_parser_iris_token: str | None # Vectorize API token for iris parser (VECTORIZE_TOKEN)
|
||||
file_parser_iris_org_id: str | None # Vectorize org ID for iris parser (VECTORIZE_ORG_ID)
|
||||
file_parser_llama_parse_api_key: str | None # LlamaCloud API key for llama_parse parser
|
||||
file_conversion_max_batch_size_mb: int # Max total batch size in MB (all files combined)
|
||||
file_conversion_max_batch_size: int # Max files per request
|
||||
enable_file_upload_api: bool
|
||||
@@ -907,10 +1027,13 @@ 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_recall_budget: str
|
||||
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
|
||||
|
||||
@@ -925,6 +1048,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
|
||||
@@ -942,6 +1084,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
|
||||
@@ -950,7 +1093,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
|
||||
@@ -977,6 +1120,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
|
||||
@@ -1009,6 +1156,7 @@ class HindsightConfig:
|
||||
"file_storage_azure_account_key",
|
||||
# File parser credentials
|
||||
"file_parser_iris_token",
|
||||
"file_parser_llama_parse_api_key",
|
||||
}
|
||||
|
||||
# CONFIGURABLE_FIELDS: Safe behavioral settings that can be customized per-tenant/bank
|
||||
@@ -1031,6 +1179,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",
|
||||
@@ -1038,6 +1187,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",
|
||||
@@ -1144,6 +1307,40 @@ class HindsightConfig:
|
||||
f"provider: {self.retain_llm_provider or self.llm_provider})"
|
||||
)
|
||||
|
||||
# Warn if local ML dependencies are missing when configured.
|
||||
# Don't hard-fail here — the actual ImportError fires at model init time
|
||||
# with a clear message. This early warning catches it before startup proceeds.
|
||||
if self.embeddings_provider == "local" or self.reranker_provider == "local":
|
||||
try:
|
||||
import importlib
|
||||
|
||||
importlib.import_module("sentence_transformers")
|
||||
except ImportError:
|
||||
missing = []
|
||||
if self.embeddings_provider == "local":
|
||||
missing.append("embeddings")
|
||||
if self.reranker_provider == "local":
|
||||
missing.append("reranker")
|
||||
logger.warning(
|
||||
"Local ML provider configured for %s, but 'sentence-transformers' "
|
||||
"is not installed. The API will fail at startup. Either:\n"
|
||||
" 1. Install local ML deps: pip install hindsight-api[local-ml]\n"
|
||||
" 2. Use a remote provider instead:\n"
|
||||
" HINDSIGHT_API_EMBEDDINGS_PROVIDER=openai (or gemini, tei)\n"
|
||||
" HINDSIGHT_API_RERANKER_PROVIDER=none (or tei)",
|
||||
" and ".join(missing),
|
||||
)
|
||||
|
||||
# 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."""
|
||||
@@ -1153,6 +1350,7 @@ class HindsightConfig:
|
||||
|
||||
config = cls(
|
||||
# Database
|
||||
database_backend=os.getenv(ENV_DATABASE_BACKEND, DEFAULT_DATABASE_BACKEND).lower(),
|
||||
database_url=os.getenv(ENV_DATABASE_URL, DEFAULT_DATABASE_URL),
|
||||
migration_database_url=os.getenv(ENV_MIGRATION_DATABASE_URL) or None,
|
||||
database_schema=os.getenv(ENV_DATABASE_SCHEMA, DEFAULT_DATABASE_SCHEMA),
|
||||
@@ -1270,10 +1468,18 @@ 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,
|
||||
embeddings_cohere_output_dimensions=int(v)
|
||||
if (v := os.getenv(ENV_EMBEDDINGS_COHERE_OUTPUT_DIMENSIONS))
|
||||
else 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)
|
||||
@@ -1305,6 +1511,11 @@ class HindsightConfig:
|
||||
str(DEFAULT_EMBEDDINGS_GEMINI_OUTPUT_DIMENSIONALITY),
|
||||
)
|
||||
),
|
||||
embeddings_gemini_force_ipv4=os.getenv(
|
||||
ENV_EMBEDDINGS_GEMINI_FORCE_IPV4,
|
||||
str(DEFAULT_EMBEDDINGS_GEMINI_FORCE_IPV4),
|
||||
).lower()
|
||||
in ("true", "1"),
|
||||
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),
|
||||
@@ -1338,6 +1549,9 @@ 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),
|
||||
@@ -1382,6 +1596,7 @@ class HindsightConfig:
|
||||
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)
|
||||
@@ -1449,6 +1664,7 @@ class HindsightConfig:
|
||||
else None,
|
||||
file_parser_iris_token=os.getenv(ENV_FILE_PARSER_IRIS_TOKEN) or None,
|
||||
file_parser_iris_org_id=os.getenv(ENV_FILE_PARSER_IRIS_ORG_ID) or None,
|
||||
file_parser_llama_parse_api_key=os.getenv(ENV_FILE_PARSER_LLAMA_PARSE_API_KEY) or None,
|
||||
file_conversion_max_batch_size_mb=int(
|
||||
os.getenv(ENV_FILE_CONVERSION_MAX_BATCH_SIZE_MB, str(DEFAULT_FILE_CONVERSION_MAX_BATCH_SIZE_MB))
|
||||
),
|
||||
@@ -1474,12 +1690,19 @@ 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))
|
||||
),
|
||||
consolidation_max_tokens=int(
|
||||
os.getenv(ENV_CONSOLIDATION_MAX_TOKENS, str(DEFAULT_CONSOLIDATION_MAX_TOKENS))
|
||||
),
|
||||
consolidation_recall_budget=os.getenv(ENV_CONSOLIDATION_RECALL_BUDGET, DEFAULT_CONSOLIDATION_RECALL_BUDGET),
|
||||
consolidation_source_facts_max_tokens=int(
|
||||
os.getenv(ENV_CONSOLIDATION_SOURCE_FACTS_MAX_TOKENS, str(DEFAULT_CONSOLIDATION_SOURCE_FACTS_MAX_TOKENS))
|
||||
),
|
||||
@@ -1489,6 +1712,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))
|
||||
@@ -1502,6 +1728,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,
|
||||
@@ -1509,9 +1736,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))),
|
||||
@@ -1523,6 +1752,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)
|
||||
@@ -1613,7 +1867,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"))
|
||||
|
||||
@@ -1624,9 +1879,9 @@ class HindsightConfig:
|
||||
|
||||
def log_config(self) -> None:
|
||||
"""Log the current configuration (without sensitive values)."""
|
||||
logger.info(f"Database: {self.database_url} (schema: {self.database_schema})")
|
||||
logger.info(f"Database: {mask_network_location(self.database_url)} (schema: {self.database_schema})")
|
||||
if self.migration_database_url:
|
||||
logger.info(f"Migration database: {self.migration_database_url}")
|
||||
logger.info(f"Migration database: {mask_network_location(self.migration_database_url)}")
|
||||
logger.info(f"LLM: provider={self.llm_provider}, model={self.llm_model}")
|
||||
if self.retain_llm_provider or self.retain_llm_model:
|
||||
retain_provider = self.retain_llm_provider or self.llm_provider
|
||||
|
||||
@@ -11,30 +11,36 @@ multiple API servers.
|
||||
import json
|
||||
import logging
|
||||
from dataclasses import asdict, replace
|
||||
from typing import Any
|
||||
from typing import TYPE_CHECKING, 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
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from hindsight_api.engine.db.base import DatabaseBackend
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ConfigResolver:
|
||||
"""Resolves hierarchical configuration with tenant/bank overrides."""
|
||||
|
||||
def __init__(self, pool: asyncpg.Pool, tenant_extension: TenantExtension | None = None):
|
||||
def __init__(self, backend: "DatabaseBackend", tenant_extension: TenantExtension | None = None):
|
||||
"""
|
||||
Initialize config resolver.
|
||||
|
||||
Args:
|
||||
pool: Database connection pool
|
||||
backend: Database backend for connection acquisition
|
||||
tenant_extension: Optional tenant extension for tenant-level config and permissions
|
||||
"""
|
||||
self.pool = pool
|
||||
self._backend = backend
|
||||
self.tenant_extension = tenant_extension
|
||||
self._global_config = _get_raw_config()
|
||||
self._configurable_fields = HindsightConfig.get_configurable_fields()
|
||||
@@ -148,7 +154,7 @@ class ConfigResolver:
|
||||
Dict of config overrides (only configurable fields, normalized keys)
|
||||
"""
|
||||
try:
|
||||
async with self.pool.acquire() as conn:
|
||||
async with self._backend.acquire() as conn:
|
||||
row = await conn.fetchrow(
|
||||
f"""
|
||||
SELECT config FROM {fq_table("banks")} WHERE bank_id = $1
|
||||
@@ -256,8 +262,11 @@ 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:
|
||||
async with self._backend.acquire() as conn:
|
||||
await conn.execute(
|
||||
f"""
|
||||
UPDATE {fq_table("banks")}
|
||||
@@ -278,7 +287,7 @@ class ConfigResolver:
|
||||
Args:
|
||||
bank_id: Bank identifier
|
||||
"""
|
||||
async with self.pool.acquire() as conn:
|
||||
async with self._backend.acquire() as conn:
|
||||
await conn.execute(
|
||||
f"""
|
||||
UPDATE {fq_table("banks")}
|
||||
@@ -292,6 +301,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.
|
||||
|
||||
@@ -63,7 +63,17 @@ def daemonize():
|
||||
Fork the current process into a background daemon.
|
||||
|
||||
Uses double-fork technique to properly detach from terminal.
|
||||
|
||||
On Windows there is no fork model: the spawning parent is expected to
|
||||
detach us via `CREATE_NEW_PROCESS_GROUP | DETACHED_PROCESS` and to
|
||||
redirect stdout/stderr to HINDSIGHT_API_DAEMON_LOG before exec. We
|
||||
still ensure the log directory exists so that any file handlers set
|
||||
up by the calling app have a valid target.
|
||||
"""
|
||||
if sys.platform == "win32":
|
||||
DAEMON_LOG_PATH.parent.mkdir(parents=True, exist_ok=True)
|
||||
return
|
||||
|
||||
# First fork - detach from parent
|
||||
try:
|
||||
pid = os.fork()
|
||||
|
||||
@@ -0,0 +1,57 @@
|
||||
"""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))
|
||||
@@ -16,8 +16,6 @@ from dataclasses import dataclass, field
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any
|
||||
|
||||
import asyncpg
|
||||
|
||||
from ..engine.db_utils import acquire_with_retry
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -69,7 +67,7 @@ class AuditLogger:
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
pool_getter: Callable[[], asyncpg.Pool | None],
|
||||
pool_getter: Callable[[], Any],
|
||||
schema_getter: Callable[[], str],
|
||||
enabled: bool,
|
||||
allowed_actions: list[str],
|
||||
|
||||
@@ -27,8 +27,9 @@ from typing import TYPE_CHECKING, Any
|
||||
from pydantic import BaseModel, field_validator
|
||||
|
||||
from ...config import get_config
|
||||
from ..db_utils import acquire_with_retry
|
||||
from ..llm_wrapper import sanitize_llm_output
|
||||
from ..memory_engine import fq_table
|
||||
from ..memory_engine import Budget, fq_table
|
||||
from ..retain import embedding_utils
|
||||
from .prompts import build_batch_consolidation_prompt
|
||||
|
||||
@@ -42,6 +43,39 @@ 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.
|
||||
|
||||
Oracle note: Oracle doesn't support FOR SHARE, so the SQL rewriter promotes
|
||||
it to FOR UPDATE. Oracle's MVCC consistent-read semantics make FOR SHARE
|
||||
unnecessary (the sweep runs AFTER deletion), but FOR UPDATE is more
|
||||
conservative and still correct.
|
||||
"""
|
||||
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 +253,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
|
||||
@@ -226,10 +261,10 @@ async def run_consolidation_job(
|
||||
logger.debug(f"Consolidation disabled for bank {bank_id}")
|
||||
return {"status": "disabled", "bank_id": bank_id}
|
||||
|
||||
pool = memory_engine._pool
|
||||
pool = memory_engine._backend
|
||||
|
||||
# Get bank profile
|
||||
async with pool.acquire() as conn:
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
t0 = time.time()
|
||||
bank_row = await conn.fetchrow(
|
||||
f"""
|
||||
@@ -281,10 +316,19 @@ 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:
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
t0 = time.time()
|
||||
memories = await conn.fetch(
|
||||
f"""
|
||||
@@ -299,7 +343,7 @@ async def run_consolidation_job(
|
||||
LIMIT $2
|
||||
""",
|
||||
bank_id,
|
||||
max_memories_per_batch,
|
||||
fetch_limit,
|
||||
)
|
||||
perf.record_timing("fetch_memories", time.time() - t0)
|
||||
|
||||
@@ -348,7 +392,7 @@ async def run_consolidation_job(
|
||||
while pending:
|
||||
sub_batch = pending.pop(0)
|
||||
|
||||
async with pool.acquire() as conn:
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
# Determine observation_scopes for this sub-batch. All memories share
|
||||
# the same tags (enforced by tag_groups), so we only check the first memory.
|
||||
# asyncpg returns JSONB columns as raw JSON strings, so parse if needed.
|
||||
@@ -456,7 +500,7 @@ async def run_consolidation_job(
|
||||
all_results.extend(sub_results)
|
||||
|
||||
# Commit consolidated_at / consolidation_failed_at in a single DB round-trip
|
||||
async with pool.acquire() as conn:
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
if succeeded_ids:
|
||||
await conn.executemany(
|
||||
f"UPDATE {fq_table('memory_units')} SET consolidated_at = NOW() WHERE id = $1",
|
||||
@@ -524,6 +568,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 +615,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()
|
||||
|
||||
@@ -591,19 +659,17 @@ async def _trigger_mental_model_refreshes(
|
||||
Returns:
|
||||
Number of mental models scheduled for refresh
|
||||
"""
|
||||
pool = memory_engine._pool
|
||||
pool = memory_engine._backend
|
||||
|
||||
# Find mental models with refresh_after_consolidation=true
|
||||
# SECURITY: Control which mental models get refreshed based on tags
|
||||
async with pool.acquire() as conn:
|
||||
# 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 acquire_with_retry(pool) 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 +682,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 +693,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 +958,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 = {
|
||||
@@ -946,6 +1024,23 @@ async def _execute_update_action(
|
||||
source_mentioned_at,
|
||||
merged_tags,
|
||||
)
|
||||
|
||||
# Dual-write: sync observation_sources junction table with updated source_ids.
|
||||
# DELETE + INSERT is simpler than diffing, and this runs inside a transaction.
|
||||
obs_uuid = uuid.UUID(observation_id)
|
||||
await conn.execute(
|
||||
f"DELETE FROM {fq_table('observation_sources')} WHERE observation_id = $1",
|
||||
obs_uuid,
|
||||
)
|
||||
if source_ids:
|
||||
await conn.executemany(
|
||||
f"""
|
||||
INSERT INTO {fq_table("observation_sources")} (observation_id, source_id)
|
||||
VALUES ($1, $2)
|
||||
""",
|
||||
[(obs_uuid, sid) for sid in source_ids],
|
||||
)
|
||||
|
||||
if perf:
|
||||
perf.record_timing("db_write", time.time() - t0)
|
||||
|
||||
@@ -1066,10 +1161,14 @@ async def _find_related_observations(
|
||||
else:
|
||||
recall_span = None
|
||||
|
||||
# Resolve budget: consolidation doesn't need deep recall, default to LOW to reduce memory fan-out
|
||||
recall_budget = Budget(config.consolidation_recall_budget)
|
||||
|
||||
try:
|
||||
recall_result = await memory_engine.recall_async(
|
||||
bank_id=bank_id,
|
||||
query=query,
|
||||
budget=recall_budget,
|
||||
max_tokens=config.consolidation_max_tokens, # Token budget for observations (configurable)
|
||||
fact_type=["observation"], # Only retrieve observations
|
||||
request_context=request_context,
|
||||
@@ -1131,14 +1230,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 +1273,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 +1282,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 +1323,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 +1348,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])
|
||||
@@ -1288,6 +1411,18 @@ async def _create_observation_directly(
|
||||
obs_mentioned_at,
|
||||
)
|
||||
|
||||
# Dual-write: populate observation_sources junction table alongside
|
||||
# the source_memory_ids column. The junction table enables portable SQL
|
||||
# joins, replacing PG-specific array operators and Oracle JSON_TABLE.
|
||||
if source_memory_ids:
|
||||
await conn.executemany(
|
||||
f"""
|
||||
INSERT INTO {fq_table("observation_sources")} (observation_id, source_id)
|
||||
VALUES ($1, $2)
|
||||
""",
|
||||
[(observation_id, sid) for sid in source_memory_ids],
|
||||
)
|
||||
|
||||
if perf:
|
||||
perf.record_timing("db_write", time.time() - t0)
|
||||
|
||||
|
||||
@@ -19,6 +19,7 @@ from ..config import (
|
||||
DEFAULT_LITELLM_API_BASE,
|
||||
DEFAULT_RERANKER_COHERE_MODEL,
|
||||
DEFAULT_RERANKER_FLASHRANK_CACHE_DIR,
|
||||
DEFAULT_RERANKER_FLASHRANK_CPU_MEM_ARENA,
|
||||
DEFAULT_RERANKER_FLASHRANK_MODEL,
|
||||
DEFAULT_RERANKER_GOOGLE_MODEL,
|
||||
DEFAULT_RERANKER_LITELLM_MAX_TOKENS_PER_DOC,
|
||||
@@ -33,11 +34,13 @@ from ..config import (
|
||||
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_CPU_MEM_ARENA,
|
||||
ENV_RERANKER_FLASHRANK_MODEL,
|
||||
ENV_RERANKER_GOOGLE_PROJECT_ID,
|
||||
ENV_RERANKER_LITELLM_SDK_API_KEY,
|
||||
@@ -48,6 +51,7 @@ from ..config import (
|
||||
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,
|
||||
@@ -862,6 +866,7 @@ class FlashRankCrossEncoder(CrossEncoderModel):
|
||||
cache_dir: str | None = None,
|
||||
max_length: int = 512,
|
||||
max_concurrent: int = 4,
|
||||
cpu_mem_arena: bool = False,
|
||||
):
|
||||
"""
|
||||
Initialize FlashRank cross-encoder.
|
||||
@@ -871,10 +876,15 @@ class FlashRankCrossEncoder(CrossEncoderModel):
|
||||
cache_dir: Directory to cache downloaded models. Default: system cache
|
||||
max_length: Maximum sequence length for reranking. Default: 512
|
||||
max_concurrent: Maximum concurrent reranking calls. Default: 4
|
||||
cpu_mem_arena: Enable ONNX Runtime CPU memory arena. Default: False.
|
||||
When True, ONNX pre-allocates a memory arena that never
|
||||
shrinks, causing RSS to grow monotonically. False trades
|
||||
slightly slower per-call allocation for bounded RSS.
|
||||
"""
|
||||
self.model_name = model_name or DEFAULT_RERANKER_FLASHRANK_MODEL
|
||||
self.cache_dir = cache_dir or DEFAULT_RERANKER_FLASHRANK_CACHE_DIR
|
||||
self.max_length = max_length
|
||||
self.cpu_mem_arena = cpu_mem_arena
|
||||
self._ranker = None
|
||||
FlashRankCrossEncoder._max_concurrent = max_concurrent
|
||||
|
||||
@@ -892,15 +902,47 @@ class FlashRankCrossEncoder(CrossEncoderModel):
|
||||
except ImportError:
|
||||
raise ImportError("flashrank is required for FlashRankCrossEncoder. Install it with: pip install flashrank")
|
||||
|
||||
logger.info(f"Reranker: initializing FlashRank provider with model {self.model_name}")
|
||||
logger.info(
|
||||
f"Reranker: initializing FlashRank provider with model {self.model_name}"
|
||||
f" (cpu_mem_arena={self.cpu_mem_arena})"
|
||||
)
|
||||
|
||||
# Configure ONNX session options before Ranker creates the session.
|
||||
# When cpu_mem_arena=False (default), ONNX won't pre-allocate an arena
|
||||
# that grows monotonically, keeping RSS bounded after rerank batches.
|
||||
if not self.cpu_mem_arena:
|
||||
import onnxruntime as ort
|
||||
|
||||
session_options = ort.SessionOptions()
|
||||
session_options.enable_cpu_mem_arena = False
|
||||
else:
|
||||
session_options = None
|
||||
|
||||
# Initialize ranker with optional cache directory
|
||||
ranker_kwargs = {"model_name": self.model_name, "max_length": self.max_length}
|
||||
ranker_kwargs: dict = {"model_name": self.model_name, "max_length": self.max_length}
|
||||
if self.cache_dir:
|
||||
ranker_kwargs["cache_dir"] = self.cache_dir
|
||||
|
||||
self._ranker = Ranker(**ranker_kwargs)
|
||||
|
||||
# Patch the ONNX session options if arena is disabled.
|
||||
# FlashRank's Ranker doesn't expose SessionOptions in its API,
|
||||
# so we replace the session after initialization.
|
||||
if session_options is not None and hasattr(self._ranker, "session"):
|
||||
import onnxruntime as ort
|
||||
|
||||
model_file = None
|
||||
model_dir = getattr(self._ranker, "model_dir", None)
|
||||
if model_dir:
|
||||
from pathlib import Path
|
||||
|
||||
for candidate in Path(model_dir).glob("*.onnx"):
|
||||
model_file = str(candidate)
|
||||
break
|
||||
if model_file:
|
||||
self._ranker.session = ort.InferenceSession(model_file, sess_options=session_options)
|
||||
logger.info("Reranker: replaced FlashRank ONNX session with cpu_mem_arena=False")
|
||||
|
||||
# Initialize shared executor
|
||||
if FlashRankCrossEncoder._executor is None:
|
||||
FlashRankCrossEncoder._executor = ThreadPoolExecutor(
|
||||
@@ -1282,6 +1324,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
|
||||
|
||||
@@ -1297,6 +1340,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]:
|
||||
@@ -1310,13 +1357,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
|
||||
|
||||
@@ -1506,6 +1554,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,
|
||||
)
|
||||
@@ -1543,7 +1592,10 @@ def create_cross_encoder_from_env() -> CrossEncoderModel:
|
||||
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)
|
||||
return FlashRankCrossEncoder(model_name=model, cache_dir=cache_dir)
|
||||
cpu_mem_arena = os.environ.get(
|
||||
ENV_RERANKER_FLASHRANK_CPU_MEM_ARENA, str(DEFAULT_RERANKER_FLASHRANK_CPU_MEM_ARENA)
|
||||
).lower() in ("true", "1", "yes")
|
||||
return FlashRankCrossEncoder(model_name=model, cache_dir=cache_dir, cpu_mem_arena=cpu_mem_arena)
|
||||
elif provider == "litellm":
|
||||
return LiteLLMCrossEncoder(
|
||||
api_base=config.reranker_litellm_api_base,
|
||||
|
||||
@@ -0,0 +1,82 @@
|
||||
"""Database backend abstraction layer.
|
||||
|
||||
Provides a uniform interface over different database drivers (asyncpg, oracledb, etc.)
|
||||
so that business logic is decoupled from any specific database platform.
|
||||
|
||||
Usage:
|
||||
from hindsight_api.engine.db import create_database_backend, DatabaseBackend
|
||||
|
||||
backend = create_database_backend("postgresql")
|
||||
await backend.initialize(dsn="postgresql://...")
|
||||
async with backend.acquire() as conn:
|
||||
rows = await conn.fetch("SELECT ...")
|
||||
"""
|
||||
|
||||
from .base import DatabaseBackend, DatabaseConnection
|
||||
from .ops import DataAccessOps
|
||||
from .result import ResultRow
|
||||
|
||||
__all__ = [
|
||||
"DataAccessOps",
|
||||
"DatabaseBackend",
|
||||
"DatabaseConnection",
|
||||
"ResultRow",
|
||||
"create_data_access_ops",
|
||||
"create_database_backend",
|
||||
]
|
||||
|
||||
|
||||
def _get_backend_class(backend_type: str) -> type[DatabaseBackend]:
|
||||
"""Resolve backend class by name using lazy imports."""
|
||||
if backend_type == "postgresql":
|
||||
from .postgresql import PostgreSQLBackend
|
||||
|
||||
return PostgreSQLBackend
|
||||
if backend_type == "oracle":
|
||||
from .oracle import OracleBackend
|
||||
|
||||
return OracleBackend
|
||||
raise ValueError(f"Unknown database backend: {backend_type!r}. Supported: 'postgresql', 'oracle'.")
|
||||
|
||||
|
||||
def _get_ops_class(backend_type: str) -> type[DataAccessOps]:
|
||||
"""Resolve ops class by name using lazy imports."""
|
||||
if backend_type == "postgresql":
|
||||
from .ops_postgresql import PostgreSQLOps
|
||||
|
||||
return PostgreSQLOps
|
||||
if backend_type == "oracle":
|
||||
from .ops_oracle import OracleOps
|
||||
|
||||
return OracleOps
|
||||
raise ValueError(f"Unknown data access ops: {backend_type!r}. Supported: 'postgresql', 'oracle'.")
|
||||
|
||||
|
||||
def create_database_backend(backend_type: str) -> DatabaseBackend:
|
||||
"""Factory: create a DatabaseBackend by name.
|
||||
|
||||
Args:
|
||||
backend_type: One of "postgresql" or "oracle".
|
||||
|
||||
Returns:
|
||||
An uninitialized DatabaseBackend instance.
|
||||
|
||||
Raises:
|
||||
ValueError: If backend_type is not recognized.
|
||||
"""
|
||||
return _get_backend_class(backend_type)()
|
||||
|
||||
|
||||
def create_data_access_ops(backend_type: str) -> DataAccessOps:
|
||||
"""Factory: create a DataAccessOps by backend name.
|
||||
|
||||
Args:
|
||||
backend_type: One of "postgresql" or "oracle".
|
||||
|
||||
Returns:
|
||||
A DataAccessOps instance.
|
||||
|
||||
Raises:
|
||||
ValueError: If backend_type is not recognized.
|
||||
"""
|
||||
return _get_ops_class(backend_type)()
|
||||
@@ -0,0 +1,342 @@
|
||||
"""Abstract base classes for database backend abstraction.
|
||||
|
||||
Defines the interfaces that all database backends (PostgreSQL, Oracle, etc.)
|
||||
must implement. Business logic depends only on these interfaces.
|
||||
"""
|
||||
|
||||
import json
|
||||
from abc import ABC, abstractmethod
|
||||
from collections.abc import AsyncIterator
|
||||
from contextlib import asynccontextmanager
|
||||
|
||||
# TYPE_CHECKING-only import to avoid circular import at runtime.
|
||||
# DataAccessOps lives in ops.py which imports nothing from base.py,
|
||||
# so the cycle is: base -> ops (type-only) and ops -> (nothing from base).
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from .result import ResultRow
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .ops import DataAccessOps
|
||||
|
||||
|
||||
class DatabaseConnection(ABC):
|
||||
"""Wraps a single connection from the pool.
|
||||
|
||||
Provides a uniform interface over asyncpg.Connection, oracledb cursor, etc.
|
||||
Methods mirror asyncpg's connection API for minimal migration friction.
|
||||
"""
|
||||
|
||||
@property
|
||||
def backend_type(self) -> str:
|
||||
"""Return ``"postgresql"`` or ``"oracle"``."""
|
||||
return "postgresql"
|
||||
|
||||
def parse_json(self, value: Any) -> Any:
|
||||
"""Parse a JSON column value into a Python object.
|
||||
|
||||
PG (asyncpg) returns JSON columns as strings that need json.loads().
|
||||
Oracle returns them as pre-parsed dicts/lists (via OracleConnection
|
||||
row conversion). This method normalizes both to Python objects.
|
||||
"""
|
||||
if value is None:
|
||||
return None
|
||||
if isinstance(value, str):
|
||||
try:
|
||||
return json.loads(value)
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
return value
|
||||
# Already a dict/list (Oracle pre-parses JSON columns)
|
||||
return value
|
||||
|
||||
async def bulk_insert_from_arrays(
|
||||
self,
|
||||
table: str,
|
||||
columns: list[str],
|
||||
arrays: list[list],
|
||||
*,
|
||||
column_types: list[str] | None = None,
|
||||
returning: str | None = None,
|
||||
) -> list[ResultRow] | str:
|
||||
"""Insert multiple rows from parallel arrays.
|
||||
|
||||
Default implementation uses ``INSERT ... SELECT * FROM unnest(...)``
|
||||
(PostgreSQL). Oracle overrides this with ``executemany``.
|
||||
|
||||
Args:
|
||||
table: Fully-qualified table name.
|
||||
columns: Column names matching the arrays.
|
||||
arrays: Parallel lists of values, one per column.
|
||||
column_types: PG type suffixes for unnest casting (e.g. ``["text[]", "uuid[]"]``).
|
||||
Ignored by backends that don't use unnest.
|
||||
returning: Optional column expression for a RETURNING clause.
|
||||
|
||||
Returns:
|
||||
If *returning* is set, a list of ResultRow; otherwise a status string.
|
||||
"""
|
||||
# Default: PostgreSQL unnest path
|
||||
col_list = ", ".join(columns)
|
||||
n_cols = len(columns)
|
||||
types = column_types or ["text[]"] * n_cols
|
||||
unnest_args = ", ".join(f"${i + 1}::{types[i]}" for i in range(n_cols))
|
||||
query = f"INSERT INTO {table} ({col_list}) SELECT * FROM unnest({unnest_args})"
|
||||
if returning:
|
||||
query += f" RETURNING {returning}"
|
||||
return await self.fetch(query, *arrays)
|
||||
result = await self.execute(query, *arrays)
|
||||
return result
|
||||
|
||||
@abstractmethod
|
||||
@asynccontextmanager
|
||||
async def transaction(self) -> AsyncIterator["DatabaseConnection"]:
|
||||
"""Start a transaction (or savepoint if already in a transaction).
|
||||
|
||||
Yields:
|
||||
Self — the same connection, now inside a transaction scope.
|
||||
On clean exit the transaction is committed; on exception it is rolled back.
|
||||
"""
|
||||
... # pragma: no cover
|
||||
yield # type: ignore[misc]
|
||||
|
||||
@abstractmethod
|
||||
async def execute(self, query: str, *args: Any, timeout: float | None = None) -> str:
|
||||
"""Execute a query and return a status string (e.g. 'INSERT 0 1').
|
||||
|
||||
Args:
|
||||
query: SQL query with dialect-appropriate placeholders.
|
||||
*args: Positional bind parameters.
|
||||
timeout: Optional statement timeout in seconds.
|
||||
|
||||
Returns:
|
||||
Command status string.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def executemany(self, query: str, args: list[tuple[Any, ...]], *, timeout: float | None = None) -> None:
|
||||
"""Execute a query for each set of arguments.
|
||||
|
||||
Args:
|
||||
query: SQL query with dialect-appropriate placeholders.
|
||||
args: List of argument tuples, one per execution.
|
||||
timeout: Optional statement timeout in seconds.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def fetch(self, query: str, *args: Any, timeout: float | None = None) -> list[ResultRow]:
|
||||
"""Execute a query and return all rows.
|
||||
|
||||
Args:
|
||||
query: SQL query with dialect-appropriate placeholders.
|
||||
*args: Positional bind parameters.
|
||||
timeout: Optional statement timeout in seconds.
|
||||
|
||||
Returns:
|
||||
List of ResultRow objects.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def fetchrow(self, query: str, *args: Any, timeout: float | None = None) -> ResultRow | None:
|
||||
"""Execute a query and return a single row (or None).
|
||||
|
||||
Args:
|
||||
query: SQL query with dialect-appropriate placeholders.
|
||||
*args: Positional bind parameters.
|
||||
timeout: Optional statement timeout in seconds.
|
||||
|
||||
Returns:
|
||||
A single ResultRow, or None if no rows match.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def fetchval(self, query: str, *args: Any, column: int = 0, timeout: float | None = None) -> Any:
|
||||
"""Execute a query and return a single value from the first row.
|
||||
|
||||
Args:
|
||||
query: SQL query with dialect-appropriate placeholders.
|
||||
*args: Positional bind parameters.
|
||||
column: Column index to return (default 0).
|
||||
timeout: Optional statement timeout in seconds.
|
||||
|
||||
Returns:
|
||||
The value from the specified column of the first row, or None.
|
||||
"""
|
||||
...
|
||||
|
||||
async def copy_records_to_table(
|
||||
self,
|
||||
table_name: str,
|
||||
*,
|
||||
records: list[tuple[Any, ...]],
|
||||
columns: list[str],
|
||||
timeout: float | None = None,
|
||||
) -> None:
|
||||
"""Bulk-load records into a table.
|
||||
|
||||
Default implementation uses executemany INSERT. Backends with native
|
||||
bulk-load support (e.g. asyncpg COPY) should override for performance.
|
||||
"""
|
||||
cols = ", ".join(columns)
|
||||
placeholders = ", ".join(f"${i + 1}" for i in range(len(columns)))
|
||||
query = f"INSERT INTO {table_name} ({cols}) VALUES ({placeholders})"
|
||||
await self.executemany(query, list(records))
|
||||
|
||||
|
||||
class DatabaseBackend(ABC):
|
||||
"""Database pool lifecycle and connection acquisition.
|
||||
|
||||
Manages the connection pool and provides context managers for
|
||||
acquiring connections and running transactions.
|
||||
|
||||
The ``ops`` property provides backend-specific data access operations
|
||||
(the Strategy pattern — like Django's ``connection.ops``). All business
|
||||
logic should use ``backend.ops`` instead of creating DataAccessOps
|
||||
instances directly.
|
||||
"""
|
||||
|
||||
_ops_instance: "DataAccessOps | None" = None
|
||||
|
||||
# -- Backend capabilities --------------------------------------------
|
||||
# Subclasses override these to advertise what the platform supports.
|
||||
# Callers use these instead of checking ``config.database_backend``.
|
||||
|
||||
@property
|
||||
def backend_type(self) -> str:
|
||||
"""Return ``"postgresql"`` or ``"oracle"``."""
|
||||
return "postgresql"
|
||||
|
||||
@property
|
||||
def ops(self) -> "DataAccessOps":
|
||||
"""Backend-specific data access operations (cached).
|
||||
|
||||
Follows the Django pattern: ``connection.ops`` provides the
|
||||
operations handler for the current backend. Created lazily on
|
||||
first access and cached for the lifetime of the backend.
|
||||
"""
|
||||
if self._ops_instance is None:
|
||||
from . import create_data_access_ops
|
||||
|
||||
self._ops_instance = create_data_access_ops(self.backend_type)
|
||||
return self._ops_instance
|
||||
|
||||
@property
|
||||
def supports_partial_indexes(self) -> bool:
|
||||
"""Can CREATE INDEX … WHERE <predicate>."""
|
||||
return True
|
||||
|
||||
@property
|
||||
def supports_bm25(self) -> bool:
|
||||
"""Has BM25 / tsvector full-text search."""
|
||||
return True
|
||||
|
||||
@property
|
||||
def supports_unnest(self) -> bool:
|
||||
"""Supports ``unnest()`` for expanding arrays into rows."""
|
||||
return True
|
||||
|
||||
@property
|
||||
def supports_pg_trgm(self) -> bool:
|
||||
"""Platform *might* have pg_trgm (must still be checked at runtime)."""
|
||||
return True
|
||||
|
||||
@property
|
||||
def supports_worker_poller(self) -> bool:
|
||||
"""Whether this backend supports the async WorkerPoller.
|
||||
|
||||
WorkerPoller is backend-agnostic (uses DatabaseBackend.acquire()).
|
||||
All current backends (PostgreSQL, Oracle) support it.
|
||||
"""
|
||||
return True
|
||||
|
||||
def normalize_schema(self, schema: str | None) -> str | None:
|
||||
"""Normalize a schema name for this backend.
|
||||
|
||||
Returns the schema as-is by default. Oracle overrides this to
|
||||
convert ``"public"`` (a PG-specific default) to ``None`` (use the
|
||||
connecting user's default schema).
|
||||
"""
|
||||
return schema
|
||||
|
||||
def run_migrations(self, dsn: str, *, schema: str | None = None) -> None:
|
||||
"""Run database migrations for this backend.
|
||||
|
||||
PG uses Alembic migrations. Oracle uses its own idempotent DDL runner.
|
||||
Subclasses must override this method.
|
||||
"""
|
||||
raise NotImplementedError(f"{type(self).__name__} must implement run_migrations()")
|
||||
|
||||
def create_task_backend(self, *, pool_getter: Any = None, schema_getter: Any = None) -> Any:
|
||||
"""Create the task backend for this database.
|
||||
|
||||
All backends use BrokerTaskBackend for async worker/poller execution.
|
||||
"""
|
||||
from ..task_backend import BrokerTaskBackend
|
||||
|
||||
return BrokerTaskBackend(pool_getter=pool_getter, schema_getter=schema_getter)
|
||||
|
||||
@abstractmethod
|
||||
async def initialize(
|
||||
self,
|
||||
dsn: str,
|
||||
*,
|
||||
min_size: int = 5,
|
||||
max_size: int = 20,
|
||||
command_timeout: float = 300,
|
||||
acquire_timeout: float = 30,
|
||||
statement_cache_size: int = 0,
|
||||
init_callback: Any | None = None,
|
||||
) -> None:
|
||||
"""Create the connection pool.
|
||||
|
||||
Args:
|
||||
dsn: Database connection string.
|
||||
min_size: Minimum number of connections in the pool.
|
||||
max_size: Maximum number of connections in the pool.
|
||||
command_timeout: Default command timeout in seconds.
|
||||
acquire_timeout: Timeout for acquiring a connection from the pool.
|
||||
statement_cache_size: Size of the prepared-statement cache (0 to disable).
|
||||
init_callback: Optional async callback invoked on each new connection.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def shutdown(self) -> None:
|
||||
"""Close the connection pool and release all resources."""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
@asynccontextmanager
|
||||
async def acquire(self) -> AsyncIterator[DatabaseConnection]:
|
||||
"""Acquire a connection from the pool.
|
||||
|
||||
Yields:
|
||||
A DatabaseConnection wrapper.
|
||||
"""
|
||||
... # pragma: no cover
|
||||
yield # type: ignore[misc]
|
||||
|
||||
@abstractmethod
|
||||
@asynccontextmanager
|
||||
async def transaction(self) -> AsyncIterator[DatabaseConnection]:
|
||||
"""Acquire a connection and start a transaction.
|
||||
|
||||
The transaction is committed on clean exit, rolled back on exception.
|
||||
|
||||
Yields:
|
||||
A DatabaseConnection wrapper inside a transaction.
|
||||
"""
|
||||
... # pragma: no cover
|
||||
yield # type: ignore[misc]
|
||||
|
||||
@abstractmethod
|
||||
def get_pool(self) -> Any:
|
||||
"""Return the underlying raw pool object.
|
||||
|
||||
Escape hatch for gradual migration — callers that still need direct
|
||||
pool access (e.g. asyncpg-specific features) can use this during
|
||||
the transition period.
|
||||
"""
|
||||
...
|
||||
@@ -0,0 +1,428 @@
|
||||
"""Abstract base class for backend-specific data access operations.
|
||||
|
||||
SQLDialect handles SQL *fragment* generation (param placeholders, JSON ops, vector
|
||||
distance, etc.) — stateless, no I/O.
|
||||
|
||||
DataAccessOps handles multi-statement *execution* patterns that differ between
|
||||
backends (unnest batch insert vs executemany, LATERAL fan-out vs per-row query,
|
||||
DISTINCT ON vs GROUP BY workarounds, etc.). Methods receive a DatabaseConnection
|
||||
and execute complete operations.
|
||||
|
||||
This eliminates scattered ``if backend_type == "postgresql"`` conditionals from
|
||||
business logic. Adding a new backend (e.g. Neon, Databricks) means implementing
|
||||
this ABC — consumer code never checks the backend directly.
|
||||
|
||||
Follows the Strategy pattern (Fowler's "Replace Conditional with Polymorphism")
|
||||
and mirrors Django's ``DatabaseOperations`` architecture.
|
||||
"""
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
from uuid import UUID
|
||||
|
||||
from .base import DatabaseConnection
|
||||
from .result import ResultRow
|
||||
|
||||
|
||||
@dataclass
|
||||
class TagListingParts:
|
||||
"""Backend-specific SQL fragments for the tag listing query."""
|
||||
|
||||
tag_source: str
|
||||
non_empty_check: str
|
||||
tag_col: str
|
||||
bank_prefix: str
|
||||
|
||||
|
||||
class DataAccessOps(ABC):
|
||||
"""Backend-specific multi-statement data access operations.
|
||||
|
||||
Each method encapsulates a complete DB operation that differs
|
||||
in execution strategy between backends.
|
||||
"""
|
||||
|
||||
# -- Bulk insert operations ------------------------------------------
|
||||
|
||||
@abstractmethod
|
||||
async def bulk_upsert_chunks(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
chunk_ids: list[str],
|
||||
document_ids: list[str],
|
||||
bank_ids: list[str],
|
||||
chunk_texts: list[str],
|
||||
chunk_indices: list[int],
|
||||
content_hashes: list[str],
|
||||
) -> None:
|
||||
"""Bulk upsert chunks with ON CONFLICT handling.
|
||||
|
||||
PG uses INSERT ... SELECT FROM unnest() with ON CONFLICT DO UPDATE.
|
||||
Non-PG uses bulk_insert_from_arrays (executemany).
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def insert_facts_batch(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
bank_id: str,
|
||||
fact_texts: list[str],
|
||||
embeddings: list[str],
|
||||
event_dates: list,
|
||||
occurred_starts: list,
|
||||
occurred_ends: list,
|
||||
mentioned_ats: list,
|
||||
contexts: list[str],
|
||||
fact_types: list[str],
|
||||
metadata_jsons: list[str],
|
||||
chunk_ids: list,
|
||||
document_ids: list,
|
||||
tags_list: list[str],
|
||||
observation_scopes_list: list,
|
||||
text_signals_list: list,
|
||||
text_search_extension: str = "native",
|
||||
) -> list[str]:
|
||||
"""Batch-insert facts, returning IDs.
|
||||
|
||||
PG uses INSERT ... SELECT FROM unnest() with RETURNING.
|
||||
Non-PG inserts row-by-row with individual RETURNING.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def bulk_insert_links(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
sorted_links: list[tuple],
|
||||
bank_id: str,
|
||||
nil_entity_uuid: str,
|
||||
exists_clause: str,
|
||||
chunk_size: int = 5000,
|
||||
) -> None:
|
||||
"""Bulk insert memory_links with conflict handling.
|
||||
|
||||
PG uses INSERT ... SELECT FROM unnest() with chunking.
|
||||
Non-PG uses executemany.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def bulk_insert_entities(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
bank_id: str,
|
||||
entity_names: list[str],
|
||||
entity_dates: list,
|
||||
) -> dict[str, str]:
|
||||
"""Bulk insert entities with ON CONFLICT DO NOTHING, returning id-by-lowercase-name.
|
||||
|
||||
PG uses INSERT ... SELECT FROM unnest() with RETURNING.
|
||||
Non-PG inserts row-by-row then SELECTs.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def fetch_missing_entity_ids(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
bank_id: str,
|
||||
missing_names: list[str],
|
||||
) -> list[ResultRow]:
|
||||
"""Fetch entity IDs for names that conflicted during insert.
|
||||
|
||||
PG uses unnest + JOIN.
|
||||
Non-PG queries each name individually.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def bulk_insert_unit_entities(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
unit_ids: list,
|
||||
entity_ids: list,
|
||||
) -> None:
|
||||
"""Bulk insert unit_entities links with ON CONFLICT DO NOTHING.
|
||||
|
||||
PG uses INSERT ... SELECT FROM unnest().
|
||||
Non-PG uses executemany.
|
||||
"""
|
||||
...
|
||||
|
||||
# -- LATERAL / fan-out queries ---------------------------------------
|
||||
|
||||
@abstractmethod
|
||||
async def fetch_entity_unit_fanout(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
ue_table: str,
|
||||
entity_id_list: list[UUID],
|
||||
limit_per_entity: int,
|
||||
) -> list[ResultRow]:
|
||||
"""Fetch unit_ids for a list of entities with per-entity row cap.
|
||||
|
||||
PG uses unnest + CROSS JOIN LATERAL with LIMIT.
|
||||
Non-PG queries each entity individually.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def fetch_unit_dates(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
mu_table: str,
|
||||
unit_ids: list[str],
|
||||
) -> list[ResultRow]:
|
||||
"""Fetch event_date/fact_type for a list of unit IDs.
|
||||
|
||||
PG uses ANY($1) array binding.
|
||||
Non-PG queries each unit individually.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def fetch_temporal_neighbors(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
mu_table: str,
|
||||
bank_id: str,
|
||||
lateral_unit_ids: list,
|
||||
lateral_event_dates: list,
|
||||
lateral_fact_types: list,
|
||||
half_limit: int,
|
||||
batch_size: int = 500,
|
||||
) -> list[ResultRow]:
|
||||
"""Fetch temporal neighbors using bidirectional index scan.
|
||||
|
||||
PG uses unnest + CROSS JOIN LATERAL for batched bidirectional scan.
|
||||
Non-PG queries each unit individually with backward/forward scans.
|
||||
"""
|
||||
...
|
||||
|
||||
# -- CTE builders for graph retrieval --------------------------------
|
||||
|
||||
@abstractmethod
|
||||
def build_entity_expansion_cte(
|
||||
self,
|
||||
mu_table: str,
|
||||
ue_table: str,
|
||||
per_entity_limit: int,
|
||||
) -> str:
|
||||
"""Build entity expansion CTE for link expansion retrieval.
|
||||
|
||||
PG uses DISTINCT ON with CROSS JOIN LATERAL and GROUP BY.
|
||||
Non-PG splits into entity_scores subquery then JOINs for full columns
|
||||
(can't GROUP BY CLOB).
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
def build_semantic_causal_cte(
|
||||
self,
|
||||
ml_table: str,
|
||||
mu_table: str,
|
||||
) -> str:
|
||||
"""Build semantic + causal expansion CTEs.
|
||||
|
||||
PG uses DISTINCT ON for deduplication.
|
||||
Non-PG computes MAX(weight) in subquery then JOINs for full columns.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def expand_observations(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
mu_table: str,
|
||||
ue_table: str,
|
||||
ml_table: str,
|
||||
seed_ids: list,
|
||||
budget: int,
|
||||
per_entity_limit: int,
|
||||
) -> tuple[list[ResultRow], list[ResultRow], list[ResultRow]]:
|
||||
"""Observation-specific graph expansion.
|
||||
|
||||
Both backends use the observation_sources junction table with standard
|
||||
SQL joins. Previously PG used native array ops and Oracle used JSON_TABLE.
|
||||
"""
|
||||
...
|
||||
|
||||
# -- Tag listing -----------------------------------------------------
|
||||
|
||||
@abstractmethod
|
||||
def build_tag_listing_parts(self, mu_table: str) -> TagListingParts:
|
||||
"""Build SQL fragments for the tag listing query.
|
||||
|
||||
PG uses unnest(tags) to expand the VARCHAR[] column.
|
||||
Non-PG uses CROSS APPLY JSON_TABLE on the CLOB column.
|
||||
"""
|
||||
...
|
||||
|
||||
# -- Bank index management -------------------------------------------
|
||||
|
||||
@abstractmethod
|
||||
async def create_bank_vector_indexes(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
bank_id: str,
|
||||
internal_id: str,
|
||||
index_clause: str,
|
||||
fact_types: dict[str, str],
|
||||
) -> None:
|
||||
"""Create per-bank partial vector indexes.
|
||||
|
||||
PG creates per-(bank, fact_type) partial indexes.
|
||||
Non-PG is a no-op (uses global index).
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def drop_bank_vector_indexes(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
schema: str,
|
||||
internal_id: str,
|
||||
fact_types: dict[str, str],
|
||||
) -> None:
|
||||
"""Drop per-bank partial vector indexes.
|
||||
|
||||
PG drops per-(bank, fact_type) indexes.
|
||||
Non-PG is a no-op.
|
||||
"""
|
||||
...
|
||||
|
||||
# -- Entity resolution strategy routing ------------------------------
|
||||
|
||||
@abstractmethod
|
||||
def get_entity_resolution_strategy(self) -> str:
|
||||
"""Return the fuzzy entity matching strategy name.
|
||||
|
||||
PG uses "trigram" (pg_trgm).
|
||||
Non-PG uses "oracle_fuzzy" (UTL_MATCH) or falls back to "full".
|
||||
"""
|
||||
...
|
||||
|
||||
# -- Webhook operations ------------------------------------------------
|
||||
|
||||
@abstractmethod
|
||||
async def create_webhook(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
webhook_id: Any,
|
||||
bank_id: str,
|
||||
url: str,
|
||||
secret: str | None,
|
||||
event_types: list[str],
|
||||
enabled: bool,
|
||||
http_config_json: str,
|
||||
) -> ResultRow | None:
|
||||
"""Insert a webhook row and return the created row."""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def list_webhooks_for_bank(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
bank_id: str,
|
||||
) -> list[ResultRow]:
|
||||
"""List all webhooks for a bank, ordered by created_at."""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def get_webhooks_for_dispatch(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
webhook_table: str,
|
||||
bank_id: str,
|
||||
) -> list[ResultRow]:
|
||||
"""Get enabled webhooks matching a bank (bank-specific + global NULL rows)."""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def update_webhook(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
webhook_id: Any,
|
||||
bank_id: str,
|
||||
set_clauses: list[str],
|
||||
params: list[Any],
|
||||
) -> ResultRow | None:
|
||||
"""Update a webhook and return the updated row, or None if not found."""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def delete_webhook(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
webhook_id: Any,
|
||||
bank_id: str,
|
||||
) -> bool:
|
||||
"""Delete a webhook. Returns True if a row was deleted."""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def list_webhook_deliveries(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
ops_table: str,
|
||||
webhook_id: str,
|
||||
bank_id: str,
|
||||
limit: int,
|
||||
cursor: str | None,
|
||||
) -> list[ResultRow]:
|
||||
"""List webhook delivery operations for a specific webhook, newest first."""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def insert_webhook_delivery_task(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
ops_table: str,
|
||||
operation_id: Any,
|
||||
bank_id: str,
|
||||
payload_json: str,
|
||||
timestamp: Any,
|
||||
) -> None:
|
||||
"""Insert a webhook delivery task into async_operations."""
|
||||
...
|
||||
|
||||
# -- Task claiming operations ------------------------------------------
|
||||
|
||||
@abstractmethod
|
||||
async def claim_tasks(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
worker_id: str,
|
||||
reserved_limits: dict[str, int],
|
||||
shared_limit: int,
|
||||
) -> list[ResultRow]:
|
||||
"""Claim pending tasks from the async_operations table.
|
||||
|
||||
PG implementation can use NOT EXISTS + FOR UPDATE SKIP LOCKED in one query.
|
||||
Oracle implementation uses two-step claims (query busy banks first, then
|
||||
claim excluding them) to avoid ORA-02014.
|
||||
|
||||
Returns claimed rows with operation_id, operation_type, task_payload, retry_count.
|
||||
The caller is responsible for building ClaimedTask objects.
|
||||
"""
|
||||
...
|
||||
|
||||
# -- Shared helpers (concrete) -----------------------------------------
|
||||
|
||||
def _get_mu_table(self) -> str:
|
||||
"""Get the fully-qualified memory_units table name."""
|
||||
from ..schema import fq_table
|
||||
|
||||
return fq_table("memory_units")
|
||||
@@ -0,0 +1,938 @@
|
||||
"""Oracle 23ai implementation of DataAccessOps.
|
||||
|
||||
Uses executemany, per-row queries, JSON_TABLE, and ROW_NUMBER() workarounds
|
||||
for Oracle-specific syntax requirements (no unnest, no DISTINCT ON, CLOB
|
||||
columns can't appear in GROUP BY).
|
||||
"""
|
||||
|
||||
import json
|
||||
import uuid as uuid_mod
|
||||
from datetime import UTC, datetime
|
||||
from typing import Any
|
||||
from uuid import UUID
|
||||
|
||||
from .base import DatabaseConnection
|
||||
from .ops import DataAccessOps, TagListingParts
|
||||
from .result import ResultRow
|
||||
|
||||
|
||||
class OracleOps(DataAccessOps):
|
||||
"""Oracle-specific data access operations."""
|
||||
|
||||
async def bulk_upsert_chunks(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
chunk_ids: list[str],
|
||||
document_ids: list[str],
|
||||
bank_ids: list[str],
|
||||
chunk_texts: list[str],
|
||||
chunk_indices: list[int],
|
||||
content_hashes: list[str],
|
||||
) -> None:
|
||||
# Oracle's thin-client executemany with array binds is already well-optimized —
|
||||
# it batches network round-trips into a single call, so INSERT ALL or other
|
||||
# patterns would not provide a meaningful improvement.
|
||||
await conn.bulk_insert_from_arrays(
|
||||
table,
|
||||
["chunk_id", "document_id", "bank_id", "chunk_text", "chunk_index", "content_hash"],
|
||||
[
|
||||
chunk_ids,
|
||||
document_ids,
|
||||
bank_ids,
|
||||
chunk_texts,
|
||||
chunk_indices,
|
||||
content_hashes,
|
||||
],
|
||||
column_types=["text[]", "text[]", "text[]", "text[]", "integer[]", "text[]"],
|
||||
)
|
||||
|
||||
async def insert_facts_batch(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
bank_id: str,
|
||||
fact_texts: list[str],
|
||||
embeddings: list[str],
|
||||
event_dates: list,
|
||||
occurred_starts: list,
|
||||
occurred_ends: list,
|
||||
mentioned_ats: list,
|
||||
contexts: list[str],
|
||||
fact_types: list[str],
|
||||
metadata_jsons: list[str],
|
||||
chunk_ids: list,
|
||||
document_ids: list,
|
||||
tags_list: list[str],
|
||||
observation_scopes_list: list,
|
||||
text_signals_list: list,
|
||||
text_search_extension: str = "native",
|
||||
) -> list[str]:
|
||||
table = self._get_mu_table()
|
||||
# Generate UUIDs client-side so we can use executemany (single network
|
||||
# round-trip) instead of N individual INSERT+RETURNING calls.
|
||||
unit_ids = [str(uuid_mod.uuid4()) for _ in range(len(fact_texts))]
|
||||
rows_data = []
|
||||
for i in range(len(fact_texts)):
|
||||
tags_value = json.loads(tags_list[i]) if tags_list[i] else []
|
||||
rows_data.append(
|
||||
(
|
||||
unit_ids[i],
|
||||
bank_id,
|
||||
fact_texts[i],
|
||||
embeddings[i],
|
||||
event_dates[i],
|
||||
occurred_starts[i],
|
||||
occurred_ends[i],
|
||||
mentioned_ats[i],
|
||||
contexts[i],
|
||||
fact_types[i],
|
||||
metadata_jsons[i],
|
||||
chunk_ids[i],
|
||||
document_ids[i],
|
||||
tags_value,
|
||||
observation_scopes_list[i],
|
||||
text_signals_list[i],
|
||||
)
|
||||
)
|
||||
await conn.executemany(
|
||||
f"""
|
||||
INSERT INTO {table} (id, bank_id, text, embedding, event_date, occurred_start,
|
||||
occurred_end, mentioned_at, context, fact_type, metadata, chunk_id, document_id,
|
||||
tags, observation_scopes, text_signals)
|
||||
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16)
|
||||
""",
|
||||
rows_data,
|
||||
)
|
||||
return unit_ids
|
||||
|
||||
async def bulk_insert_links(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
sorted_links: list[tuple],
|
||||
bank_id: str,
|
||||
nil_entity_uuid: str,
|
||||
exists_clause: str,
|
||||
chunk_size: int = 5000,
|
||||
) -> None:
|
||||
# The backend rewrites ON CONFLICT DO NOTHING for duplicate suppression.
|
||||
# WHERE EXISTS checks are intentionally skipped: executemany does not support
|
||||
# correlated subqueries in this form, and callers guarantee unit validity.
|
||||
from_ids = [lnk[0] for lnk in sorted_links]
|
||||
to_ids = [lnk[1] for lnk in sorted_links]
|
||||
types = [lnk[2] for lnk in sorted_links]
|
||||
weights = [lnk[3] for lnk in sorted_links]
|
||||
entity_ids = [lnk[4] for lnk in sorted_links]
|
||||
|
||||
await conn.executemany(
|
||||
f"""
|
||||
INSERT INTO {table}
|
||||
(from_unit_id, to_unit_id, link_type, weight, entity_id, bank_id)
|
||||
VALUES ($1, $2, $3, $4, $5, $6)
|
||||
ON CONFLICT (from_unit_id, to_unit_id, link_type,
|
||||
COALESCE(entity_id, '{nil_entity_uuid}'::uuid))
|
||||
DO NOTHING
|
||||
""",
|
||||
[(from_ids[i], to_ids[i], types[i], weights[i], entity_ids[i], bank_id) for i in range(len(sorted_links))],
|
||||
)
|
||||
|
||||
async def bulk_insert_entities(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
bank_id: str,
|
||||
entity_names: list[str],
|
||||
entity_dates: list,
|
||||
) -> dict[str, str]:
|
||||
# Row-by-row insert with duplicate suppression.
|
||||
# Can't use RETURNING with ON CONFLICT DO NOTHING reliably,
|
||||
# so INSERT (ignoring dups) then SELECT all IDs at the end.
|
||||
id_by_name: dict[str, str] = {}
|
||||
for name, event_date in zip(entity_names, entity_dates):
|
||||
ts = event_date if event_date else datetime.now(UTC)
|
||||
await conn.execute(
|
||||
f"""
|
||||
INSERT INTO {table} (bank_id, canonical_name, first_seen, last_seen, mention_count)
|
||||
VALUES ($1, $2, $3, $3, 0)
|
||||
ON CONFLICT (bank_id, LOWER(canonical_name)) DO NOTHING
|
||||
""",
|
||||
bank_id,
|
||||
name,
|
||||
ts,
|
||||
)
|
||||
# Now SELECT all the entities we just inserted (or that already existed)
|
||||
for name in entity_names:
|
||||
row = await conn.fetchrow(
|
||||
f"""
|
||||
SELECT id, LOWER(canonical_name) AS name_lower
|
||||
FROM {table}
|
||||
WHERE bank_id = $1 AND LOWER(canonical_name) = LOWER($2)
|
||||
""",
|
||||
bank_id,
|
||||
name,
|
||||
)
|
||||
if row:
|
||||
id_by_name[row["name_lower"]] = row["id"]
|
||||
return id_by_name
|
||||
|
||||
async def fetch_missing_entity_ids(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
bank_id: str,
|
||||
missing_names: list[str],
|
||||
) -> list[ResultRow]:
|
||||
# Query each missing entity individually
|
||||
results: list[ResultRow] = []
|
||||
for orig_name in missing_names:
|
||||
row = await conn.fetchrow(
|
||||
f"""
|
||||
SELECT id, LOWER(canonical_name) AS name_lower
|
||||
FROM {table}
|
||||
WHERE bank_id = $1 AND LOWER(canonical_name) = LOWER($2)
|
||||
""",
|
||||
bank_id,
|
||||
orig_name,
|
||||
)
|
||||
if row:
|
||||
# Wrap in a dict-like to include input_name for downstream compat
|
||||
results.append(row)
|
||||
return results
|
||||
|
||||
async def bulk_insert_unit_entities(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
unit_ids: list,
|
||||
entity_ids: list,
|
||||
) -> None:
|
||||
await conn.executemany(
|
||||
f"""
|
||||
INSERT INTO {table} (unit_id, entity_id)
|
||||
VALUES ($1, $2)
|
||||
ON CONFLICT DO NOTHING
|
||||
""",
|
||||
list(zip(unit_ids, entity_ids)),
|
||||
)
|
||||
|
||||
async def fetch_entity_unit_fanout(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
ue_table: str,
|
||||
entity_id_list: list[UUID],
|
||||
limit_per_entity: int,
|
||||
) -> list[ResultRow]:
|
||||
# Query each entity individually
|
||||
rows: list[ResultRow] = []
|
||||
for eid in entity_id_list:
|
||||
entity_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT $1 AS entity_id, ue.unit_id
|
||||
FROM {ue_table} ue
|
||||
WHERE ue.entity_id = $1
|
||||
ORDER BY ue.unit_id DESC
|
||||
LIMIT $2
|
||||
""",
|
||||
eid,
|
||||
limit_per_entity,
|
||||
)
|
||||
rows.extend(entity_rows)
|
||||
return rows
|
||||
|
||||
async def fetch_unit_dates(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
mu_table: str,
|
||||
unit_ids: list[str],
|
||||
) -> list[ResultRow]:
|
||||
# No ANY() array binding; query each unit individually
|
||||
rows: list[ResultRow] = []
|
||||
for uid in unit_ids:
|
||||
row = await conn.fetchrow(
|
||||
f"""
|
||||
SELECT id, event_date, fact_type
|
||||
FROM {mu_table}
|
||||
WHERE id = $1
|
||||
""",
|
||||
uid,
|
||||
)
|
||||
if row:
|
||||
rows.append(row)
|
||||
return rows
|
||||
|
||||
async def fetch_temporal_neighbors(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
mu_table: str,
|
||||
bank_id: str,
|
||||
lateral_unit_ids: list,
|
||||
lateral_event_dates: list,
|
||||
lateral_fact_types: list,
|
||||
half_limit: int,
|
||||
batch_size: int = 500,
|
||||
) -> list[ResultRow]:
|
||||
# Per-unit queries (no unnest/LATERAL in Oracle).
|
||||
# Fetch up to half_limit in each direction, then combine and keep the
|
||||
# half_limit closest overall via ROW_NUMBER — matching the PG behavior.
|
||||
rows: list[ResultRow] = []
|
||||
for uid, edate, ftype in zip(lateral_unit_ids, lateral_event_dates, lateral_fact_types):
|
||||
uid_str = str(uid) if not isinstance(uid, str) else uid
|
||||
unit_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT from_id, id, event_date, time_diff_hours FROM (
|
||||
SELECT combined.*, ROW_NUMBER() OVER (ORDER BY combined.time_diff_hours) AS rn
|
||||
FROM (
|
||||
SELECT * FROM (
|
||||
SELECT $1 AS from_id, mu.id, mu.event_date,
|
||||
ABS(EXTRACT(DAY FROM (mu.event_date - $2)) * 24
|
||||
+ EXTRACT(HOUR FROM (mu.event_date - $2))) AS time_diff_hours
|
||||
FROM {mu_table} mu
|
||||
WHERE mu.bank_id = $4
|
||||
AND mu.fact_type = $3
|
||||
AND mu.event_date <= $2
|
||||
AND mu.id != $6
|
||||
ORDER BY mu.event_date DESC
|
||||
FETCH FIRST $5 ROWS ONLY
|
||||
) bwd
|
||||
UNION ALL
|
||||
SELECT * FROM (
|
||||
SELECT $1 AS from_id, mu.id, mu.event_date,
|
||||
ABS(EXTRACT(DAY FROM (mu.event_date - $2)) * 24
|
||||
+ EXTRACT(HOUR FROM (mu.event_date - $2))) AS time_diff_hours
|
||||
FROM {mu_table} mu
|
||||
WHERE mu.bank_id = $4
|
||||
AND mu.fact_type = $3
|
||||
AND mu.event_date > $2
|
||||
AND mu.id != $6
|
||||
ORDER BY mu.event_date ASC
|
||||
FETCH FIRST $5 ROWS ONLY
|
||||
) fwd
|
||||
) combined
|
||||
) ranked
|
||||
WHERE rn <= $5
|
||||
""",
|
||||
uid_str,
|
||||
edate,
|
||||
ftype,
|
||||
bank_id,
|
||||
half_limit,
|
||||
uid,
|
||||
)
|
||||
rows.extend(unit_rows)
|
||||
return rows
|
||||
|
||||
def build_entity_expansion_cte(
|
||||
self,
|
||||
mu_table: str,
|
||||
ue_table: str,
|
||||
per_entity_limit: int,
|
||||
) -> str:
|
||||
# Oracle: can't GROUP BY CLOB columns (text, context).
|
||||
# Restructure: count entities per unit_id in a subquery, then join to get full columns.
|
||||
return f"""
|
||||
seed_entities AS (
|
||||
SELECT DISTINCT ue.entity_id
|
||||
FROM {ue_table} ue
|
||||
WHERE ue.unit_id = ANY($1::uuid[])
|
||||
),
|
||||
entity_scores AS (
|
||||
SELECT t.unit_id, COUNT(DISTINCT se.entity_id) AS score
|
||||
FROM seed_entities se
|
||||
CROSS JOIN LATERAL (
|
||||
SELECT ue_target.unit_id
|
||||
FROM {ue_table} 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
|
||||
FETCH FIRST {per_entity_limit} ROWS ONLY
|
||||
) t
|
||||
GROUP BY t.unit_id
|
||||
),
|
||||
entity_expanded 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.proof_count,
|
||||
es.score, 'entity' AS source
|
||||
FROM entity_scores es
|
||||
JOIN {mu_table} mu ON mu.id = es.unit_id
|
||||
WHERE mu.fact_type = $2
|
||||
ORDER BY es.score DESC
|
||||
FETCH FIRST $3 ROWS ONLY
|
||||
)"""
|
||||
|
||||
def build_semantic_causal_cte(
|
||||
self,
|
||||
ml_table: str,
|
||||
mu_table: str,
|
||||
) -> str:
|
||||
# Non-PG: can't GROUP BY CLOB columns, no DISTINCT ON.
|
||||
# Restructure semantic: compute max weight per id, then join for full columns.
|
||||
return f"""
|
||||
sem_scores AS (
|
||||
SELECT id, MAX(weight) AS score
|
||||
FROM (
|
||||
SELECT mu.id, ml.weight
|
||||
FROM {ml_table} ml
|
||||
JOIN {mu_table} 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 = $2
|
||||
AND mu.id != ALL($1::uuid[])
|
||||
UNION ALL
|
||||
SELECT mu.id, ml.weight
|
||||
FROM {ml_table} ml
|
||||
JOIN {mu_table} 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 = $2
|
||||
AND mu.id != ALL($1::uuid[])
|
||||
) sem_raw
|
||||
GROUP BY id
|
||||
),
|
||||
semantic_expanded 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.proof_count,
|
||||
ss.score, 'semantic' AS source
|
||||
FROM sem_scores ss
|
||||
JOIN {mu_table} mu ON mu.id = ss.id
|
||||
ORDER BY ss.score DESC
|
||||
FETCH FIRST $3 ROWS ONLY
|
||||
),
|
||||
causal_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.proof_count,
|
||||
ml.weight AS score,
|
||||
'causal' AS source,
|
||||
ROW_NUMBER() OVER (PARTITION BY mu.id ORDER BY ml.weight DESC) AS rn_
|
||||
FROM {ml_table} ml
|
||||
JOIN {mu_table} 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')
|
||||
AND mu.fact_type = $2
|
||||
),
|
||||
causal_expanded AS (
|
||||
SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
fact_type, document_id, chunk_id, tags, proof_count, score, source
|
||||
FROM causal_ranked WHERE rn_ = 1
|
||||
ORDER BY score DESC
|
||||
FETCH FIRST $3 ROWS ONLY
|
||||
)"""
|
||||
|
||||
async def expand_observations(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
mu_table: str,
|
||||
ue_table: str,
|
||||
ml_table: str,
|
||||
seed_ids: list,
|
||||
budget: int,
|
||||
per_entity_limit: int,
|
||||
) -> tuple[list[ResultRow], list[ResultRow], list[ResultRow]]:
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Entity expansion via observation_sources junction table.
|
||||
# Previously used JSON_TABLE to explode source_memory_ids CLOB. The junction
|
||||
# table approach uses standard SQL joins, identical to the PG backend.
|
||||
from ..schema import fq_table
|
||||
|
||||
obs_sources_table = fq_table("observation_sources")
|
||||
entity_rows = await conn.fetch(
|
||||
f"""
|
||||
WITH seed_sources AS (
|
||||
SELECT DISTINCT os.source_id
|
||||
FROM {obs_sources_table} os
|
||||
WHERE os.observation_id = ANY($1::uuid[])
|
||||
),
|
||||
source_entities AS (
|
||||
SELECT DISTINCT ue_seed.entity_id
|
||||
FROM seed_sources ss
|
||||
JOIN {ue_table} ue_seed ON ue_seed.unit_id = ss.source_id
|
||||
),
|
||||
connected_sources AS (
|
||||
SELECT DISTINCT t.unit_id AS source_id
|
||||
FROM source_entities se
|
||||
CROSS JOIN LATERAL (
|
||||
SELECT ue_target.unit_id
|
||||
FROM {ue_table} ue_target
|
||||
WHERE ue_target.entity_id = se.entity_id
|
||||
ORDER BY ue_target.unit_id DESC
|
||||
FETCH FIRST {per_entity_limit} ROWS ONLY
|
||||
) t
|
||||
WHERE NOT EXISTS (
|
||||
SELECT 1 FROM seed_sources ss WHERE ss.source_id = t.unit_id
|
||||
)
|
||||
)
|
||||
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.proof_count,
|
||||
(SELECT COUNT(*)
|
||||
FROM {obs_sources_table} os2
|
||||
WHERE os2.observation_id = mu.id
|
||||
AND os2.source_id IN (SELECT source_id FROM connected_sources)
|
||||
) AS score
|
||||
FROM {mu_table} mu
|
||||
WHERE mu.fact_type = 'observation'
|
||||
AND mu.id != ALL($1::uuid[])
|
||||
AND EXISTS (
|
||||
SELECT 1 FROM {obs_sources_table} os3
|
||||
WHERE os3.observation_id = mu.id
|
||||
AND os3.source_id IN (SELECT source_id FROM connected_sources)
|
||||
)
|
||||
ORDER BY score DESC
|
||||
FETCH FIRST $2 ROWS ONLY
|
||||
""",
|
||||
seed_ids,
|
||||
budget,
|
||||
)
|
||||
logger.debug(f"[LinkExpansion] observation graph (Oracle): found {len(entity_rows)} connected observations")
|
||||
|
||||
# Semantic + causal for observations (Oracle path)
|
||||
# Avoids GROUP BY CLOB and DISTINCT ON — mirrors _expand_world_facts Oracle strategy.
|
||||
sem_causal_rows = await conn.fetch(
|
||||
f"""
|
||||
WITH sem_scores AS (
|
||||
SELECT id, MAX(weight) AS score
|
||||
FROM (
|
||||
SELECT mu.id, ml.weight
|
||||
FROM {ml_table} ml JOIN {mu_table} 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'
|
||||
AND mu.id != ALL($1::uuid[])
|
||||
UNION ALL
|
||||
SELECT mu.id, ml.weight
|
||||
FROM {ml_table} ml JOIN {mu_table} 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
|
||||
),
|
||||
semantic_expanded 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.proof_count,
|
||||
ss.score, 'semantic' AS source
|
||||
FROM sem_scores ss
|
||||
JOIN {mu_table} mu ON mu.id = ss.id
|
||||
ORDER BY ss.score DESC
|
||||
FETCH FIRST $2 ROWS ONLY
|
||||
),
|
||||
causal_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.proof_count, ml.weight AS score,
|
||||
'causal' AS source,
|
||||
ROW_NUMBER() OVER (PARTITION BY mu.id ORDER BY ml.weight DESC) AS rn_
|
||||
FROM {ml_table} ml
|
||||
JOIN {mu_table} 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')
|
||||
AND mu.fact_type = 'observation'
|
||||
),
|
||||
causal_expanded AS (
|
||||
SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
fact_type, document_id, chunk_id, tags, proof_count, score, source
|
||||
FROM causal_ranked WHERE rn_ = 1
|
||||
ORDER BY score DESC
|
||||
FETCH FIRST $2 ROWS ONLY
|
||||
)
|
||||
SELECT * FROM semantic_expanded
|
||||
UNION ALL
|
||||
SELECT * FROM causal_expanded
|
||||
""",
|
||||
seed_ids,
|
||||
budget,
|
||||
)
|
||||
|
||||
semantic_rows = [r for r in sem_causal_rows if r["source"] == "semantic"]
|
||||
causal_rows = [r for r in sem_causal_rows if r["source"] == "causal"]
|
||||
return list(entity_rows), semantic_rows, causal_rows
|
||||
|
||||
def build_tag_listing_parts(self, mu_table: str) -> TagListingParts:
|
||||
return TagListingParts(
|
||||
tag_source=(
|
||||
f"{mu_table} mu CROSS APPLY JSON_TABLE(mu.tags, '$[*]' COLUMNS (tag VARCHAR2(256) PATH '$')) jt"
|
||||
),
|
||||
non_empty_check="AND mu.tags IS NOT NULL AND DBMS_LOB.GETLENGTH(mu.tags) > 2",
|
||||
tag_col="jt.tag",
|
||||
bank_prefix="mu.",
|
||||
)
|
||||
|
||||
async def create_bank_vector_indexes(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
bank_id: str,
|
||||
internal_id: str,
|
||||
index_clause: str,
|
||||
fact_types: dict[str, str],
|
||||
) -> None:
|
||||
# Oracle 23ai supports HNSW vector indexes but does NOT support partial
|
||||
# indexes (WHERE clause on CREATE INDEX for vector indexes). Uses a single
|
||||
# global HNSW index with ORGANIZATION NEIGHBOR PARTITIONS created during
|
||||
# migrations. memory_units is partitioned by LIST (bank_id) AUTOMATIC,
|
||||
# so Oracle creates partitions per bank on INSERT and the optimizer can
|
||||
# prune partitions on bank_id-scoped queries.
|
||||
return
|
||||
|
||||
async def drop_bank_vector_indexes(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
schema: str,
|
||||
internal_id: str,
|
||||
fact_types: dict[str, str],
|
||||
) -> None:
|
||||
# Oracle uses a single global vector index (no per-bank indexes to drop).
|
||||
return
|
||||
|
||||
def get_entity_resolution_strategy(self) -> str:
|
||||
return "oracle_fuzzy"
|
||||
|
||||
# -- Webhook operations ------------------------------------------------
|
||||
|
||||
async def create_webhook(
|
||||
self,
|
||||
conn,
|
||||
table,
|
||||
webhook_id,
|
||||
bank_id,
|
||||
url,
|
||||
secret,
|
||||
event_types,
|
||||
enabled,
|
||||
http_config_json,
|
||||
):
|
||||
return await conn.fetchrow(
|
||||
f"""
|
||||
INSERT INTO {table}
|
||||
(id, bank_id, url, secret, event_types, enabled, http_config, created_at, updated_at)
|
||||
VALUES ($1, $2, $3, $4, $5, $6, $7::jsonb, NOW(), NOW())
|
||||
RETURNING id, bank_id, url, secret, event_types, enabled,
|
||||
http_config::text, created_at::text, updated_at::text
|
||||
""",
|
||||
webhook_id,
|
||||
bank_id,
|
||||
url,
|
||||
secret,
|
||||
event_types,
|
||||
enabled,
|
||||
http_config_json,
|
||||
)
|
||||
|
||||
async def list_webhooks_for_bank(self, conn, table, bank_id):
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT id, bank_id, url, secret, event_types, enabled,
|
||||
http_config::text, created_at::text, updated_at::text
|
||||
FROM {table}
|
||||
WHERE bank_id = $1
|
||||
ORDER BY created_at
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
|
||||
async def get_webhooks_for_dispatch(self, conn, webhook_table, bank_id):
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT id, bank_id, url, secret, event_types, enabled, http_config::text
|
||||
FROM {webhook_table}
|
||||
WHERE (bank_id = $1 OR bank_id IS NULL) AND enabled = true
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
|
||||
async def update_webhook(self, conn, table, webhook_id, bank_id, set_clauses, params):
|
||||
set_clauses_with_ts = set_clauses + ["updated_at = NOW()"]
|
||||
return await conn.fetchrow(
|
||||
f"""
|
||||
UPDATE {table}
|
||||
SET {", ".join(set_clauses_with_ts)}
|
||||
WHERE id = $1 AND bank_id = $2
|
||||
RETURNING id, bank_id, url, secret, event_types, enabled,
|
||||
http_config::text, created_at::text, updated_at::text
|
||||
""",
|
||||
*params,
|
||||
)
|
||||
|
||||
async def delete_webhook(self, conn, table, webhook_id, bank_id):
|
||||
result = await conn.execute(
|
||||
f"DELETE FROM {table} WHERE id = $1 AND bank_id = $2",
|
||||
webhook_id,
|
||||
bank_id,
|
||||
)
|
||||
return int(result.split()[-1]) > 0 if result else False
|
||||
|
||||
async def list_webhook_deliveries(self, conn, ops_table, webhook_id, bank_id, limit, cursor):
|
||||
fetch_limit = limit + 1
|
||||
if cursor:
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, status, retry_count, next_retry_at::text,
|
||||
error_message, task_payload, result_metadata::text, created_at::text, updated_at::text
|
||||
FROM {ops_table}
|
||||
WHERE operation_type = 'webhook_delivery'
|
||||
AND bank_id = $1
|
||||
AND task_payload->>'webhook_id' = $2
|
||||
AND created_at < $3::timestamptz
|
||||
ORDER BY created_at DESC
|
||||
LIMIT $4
|
||||
""",
|
||||
bank_id,
|
||||
webhook_id,
|
||||
cursor,
|
||||
fetch_limit,
|
||||
)
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, status, retry_count, next_retry_at::text,
|
||||
error_message, task_payload, result_metadata::text, created_at::text, updated_at::text
|
||||
FROM {ops_table}
|
||||
WHERE operation_type = 'webhook_delivery'
|
||||
AND bank_id = $1
|
||||
AND task_payload->>'webhook_id' = $2
|
||||
ORDER BY created_at DESC
|
||||
LIMIT $3
|
||||
""",
|
||||
bank_id,
|
||||
webhook_id,
|
||||
fetch_limit,
|
||||
)
|
||||
|
||||
async def insert_webhook_delivery_task(self, conn, ops_table, operation_id, bank_id, payload_json, timestamp):
|
||||
await conn.execute(
|
||||
f"""
|
||||
INSERT INTO {ops_table}
|
||||
(operation_id, bank_id, operation_type, status, task_payload, result_metadata, created_at, updated_at)
|
||||
VALUES ($1, $2, 'webhook_delivery', 'pending', $3::jsonb, '{{}}'::jsonb, $4, $4)
|
||||
""",
|
||||
operation_id,
|
||||
bank_id,
|
||||
payload_json,
|
||||
timestamp,
|
||||
)
|
||||
|
||||
# -- Task claiming operations ------------------------------------------
|
||||
|
||||
async def claim_tasks(self, conn, table, worker_id, reserved_limits, shared_limit):
|
||||
"""Oracle two-step claiming to avoid ORA-02014 with NOT EXISTS + FOR UPDATE."""
|
||||
all_rows = []
|
||||
claimed_ids = []
|
||||
|
||||
# --- Phase 1: claim from reserved pools ---
|
||||
for op_type, limit in reserved_limits.items():
|
||||
if limit <= 0:
|
||||
continue
|
||||
|
||||
if op_type == "consolidation":
|
||||
# Two-step: find busy banks first, then claim excluding them
|
||||
busy_banks = await conn.fetch(
|
||||
f"""
|
||||
SELECT DISTINCT bank_id FROM {table}
|
||||
WHERE operation_type = 'consolidation' AND status = 'processing'
|
||||
""",
|
||||
)
|
||||
busy_bank_ids = [r["bank_id"] for r in busy_banks]
|
||||
|
||||
if busy_bank_ids:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND bank_id != ALL($1::text[])
|
||||
ORDER BY created_at
|
||||
LIMIT $2
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
busy_bank_ids,
|
||||
limit,
|
||||
)
|
||||
else:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
ORDER BY created_at
|
||||
LIMIT $1
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
limit,
|
||||
)
|
||||
else:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = $1
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
ORDER BY created_at
|
||||
LIMIT $2
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
op_type,
|
||||
limit,
|
||||
)
|
||||
|
||||
for row in rows:
|
||||
claimed_ids.append(row["operation_id"])
|
||||
all_rows.append(row)
|
||||
|
||||
# --- Phase 2: claim from shared pool ---
|
||||
remaining_shared = shared_limit
|
||||
if remaining_shared > 0:
|
||||
# 2a. Non-consolidation tasks
|
||||
if claimed_ids:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type != 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND operation_id != ALL($1::uuid[])
|
||||
ORDER BY created_at
|
||||
LIMIT $2
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
claimed_ids,
|
||||
remaining_shared,
|
||||
)
|
||||
else:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type != 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
ORDER BY created_at
|
||||
LIMIT $1
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
remaining_shared,
|
||||
)
|
||||
|
||||
for row in rows:
|
||||
claimed_ids.append(row["operation_id"])
|
||||
all_rows.append(row)
|
||||
remaining_shared -= len(rows)
|
||||
|
||||
# 2b. Consolidation tasks (with bank-serialization)
|
||||
if remaining_shared > 0:
|
||||
busy_banks_2 = await conn.fetch(
|
||||
f"""
|
||||
SELECT DISTINCT bank_id FROM {table}
|
||||
WHERE operation_type = 'consolidation' AND status = 'processing'
|
||||
""",
|
||||
)
|
||||
busy_bank_ids_2 = [r["bank_id"] for r in busy_banks_2]
|
||||
|
||||
if claimed_ids:
|
||||
if busy_bank_ids_2:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND operation_id != ALL($1::uuid[])
|
||||
AND bank_id != ALL($2::text[])
|
||||
ORDER BY created_at
|
||||
LIMIT $3
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
claimed_ids,
|
||||
busy_bank_ids_2,
|
||||
remaining_shared,
|
||||
)
|
||||
else:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND operation_id != ALL($1::uuid[])
|
||||
ORDER BY created_at
|
||||
LIMIT $2
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
claimed_ids,
|
||||
remaining_shared,
|
||||
)
|
||||
else:
|
||||
if busy_bank_ids_2:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND bank_id != ALL($1::text[])
|
||||
ORDER BY created_at
|
||||
LIMIT $2
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
busy_bank_ids_2,
|
||||
remaining_shared,
|
||||
)
|
||||
else:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
ORDER BY created_at
|
||||
LIMIT $1
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
remaining_shared,
|
||||
)
|
||||
|
||||
for row in rows:
|
||||
claimed_ids.append(row["operation_id"])
|
||||
all_rows.append(row)
|
||||
|
||||
if not all_rows:
|
||||
return []
|
||||
|
||||
# Mark all claimed rows as processing
|
||||
operation_ids = [row["operation_id"] for row in all_rows]
|
||||
await conn.execute(
|
||||
f"""
|
||||
UPDATE {table}
|
||||
SET status = 'processing', worker_id = $1, claimed_at = now(), updated_at = now()
|
||||
WHERE operation_id = ANY($2)
|
||||
""",
|
||||
worker_id,
|
||||
operation_ids,
|
||||
)
|
||||
|
||||
return all_rows
|
||||
@@ -0,0 +1,942 @@
|
||||
"""PostgreSQL implementation of DataAccessOps.
|
||||
|
||||
Uses unnest(), LATERAL, DISTINCT ON, and native array operations for
|
||||
efficient batch operations.
|
||||
"""
|
||||
|
||||
import json
|
||||
from datetime import UTC, datetime
|
||||
from typing import Any
|
||||
from uuid import UUID
|
||||
|
||||
from .base import DatabaseConnection
|
||||
from .ops import DataAccessOps, TagListingParts
|
||||
from .result import ResultRow
|
||||
|
||||
|
||||
class PostgreSQLOps(DataAccessOps):
|
||||
"""PostgreSQL-specific data access operations using unnest and LATERAL."""
|
||||
|
||||
async def bulk_upsert_chunks(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
chunk_ids: list[str],
|
||||
document_ids: list[str],
|
||||
bank_ids: list[str],
|
||||
chunk_texts: list[str],
|
||||
chunk_indices: list[int],
|
||||
content_hashes: list[str],
|
||||
) -> None:
|
||||
await conn.execute(
|
||||
f"""
|
||||
INSERT INTO {table} (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_ids,
|
||||
bank_ids,
|
||||
chunk_texts,
|
||||
chunk_indices,
|
||||
content_hashes,
|
||||
)
|
||||
|
||||
async def insert_facts_batch(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
bank_id: str,
|
||||
fact_texts: list[str],
|
||||
embeddings: list[str],
|
||||
event_dates: list,
|
||||
occurred_starts: list,
|
||||
occurred_ends: list,
|
||||
mentioned_ats: list,
|
||||
contexts: list[str],
|
||||
fact_types: list[str],
|
||||
metadata_jsons: list[str],
|
||||
chunk_ids: list,
|
||||
document_ids: list,
|
||||
tags_list: list[str],
|
||||
observation_scopes_list: list,
|
||||
text_signals_list: list,
|
||||
text_search_extension: str = "native",
|
||||
) -> list[str]:
|
||||
from ...config import get_config
|
||||
|
||||
config = get_config()
|
||||
table = self._get_mu_table()
|
||||
|
||||
if config.text_search_extension == "vchord":
|
||||
query = f"""
|
||||
WITH input_data AS (
|
||||
SELECT * FROM unnest(
|
||||
$2::text[], $3::vector[], $4::timestamptz[], $5::timestamptz[], $6::timestamptz[], $7::timestamptz[],
|
||||
$8::text[], $9::text[], $10::jsonb[], $11::text[], $12::text[], $13::jsonb[], $14::jsonb[], $15::text[]
|
||||
) AS t(text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
context, fact_type, metadata, chunk_id, document_id, tags_json,
|
||||
observation_scopes_json, text_signals)
|
||||
)
|
||||
INSERT INTO {table} (bank_id, text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
context, fact_type, metadata, chunk_id, document_id, tags,
|
||||
observation_scopes, text_signals, search_vector)
|
||||
SELECT
|
||||
$1,
|
||||
text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
context, fact_type, metadata, chunk_id, document_id,
|
||||
COALESCE(
|
||||
(SELECT array_agg(elem) FROM jsonb_array_elements_text(tags_json) AS elem),
|
||||
'{{}}'::varchar[]
|
||||
),
|
||||
observation_scopes_json,
|
||||
text_signals,
|
||||
tokenize(
|
||||
COALESCE(text, '') || ' ' || COALESCE(context, '') || ' ' || COALESCE(text_signals, ''),
|
||||
'llmlingua2'
|
||||
)::bm25_catalog.bm25vector
|
||||
FROM input_data
|
||||
RETURNING id
|
||||
"""
|
||||
else:
|
||||
query = f"""
|
||||
WITH input_data AS (
|
||||
SELECT * FROM unnest(
|
||||
$2::text[], $3::vector[], $4::timestamptz[], $5::timestamptz[], $6::timestamptz[], $7::timestamptz[],
|
||||
$8::text[], $9::text[], $10::jsonb[], $11::text[], $12::text[], $13::jsonb[], $14::jsonb[], $15::text[]
|
||||
) AS t(text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
context, fact_type, metadata, chunk_id, document_id, tags_json,
|
||||
observation_scopes_json, text_signals)
|
||||
)
|
||||
INSERT INTO {table} (bank_id, text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
context, fact_type, metadata, chunk_id, document_id, tags,
|
||||
observation_scopes, text_signals)
|
||||
SELECT
|
||||
$1,
|
||||
text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
context, fact_type, metadata, chunk_id, document_id,
|
||||
COALESCE(
|
||||
(SELECT array_agg(elem) FROM jsonb_array_elements_text(tags_json) AS elem),
|
||||
'{{}}'::varchar[]
|
||||
),
|
||||
observation_scopes_json,
|
||||
text_signals
|
||||
FROM input_data
|
||||
RETURNING id
|
||||
"""
|
||||
|
||||
results = await conn.fetch(
|
||||
query,
|
||||
bank_id,
|
||||
fact_texts,
|
||||
embeddings,
|
||||
event_dates,
|
||||
occurred_starts,
|
||||
occurred_ends,
|
||||
mentioned_ats,
|
||||
contexts,
|
||||
fact_types,
|
||||
metadata_jsons,
|
||||
chunk_ids,
|
||||
document_ids,
|
||||
tags_list,
|
||||
observation_scopes_list,
|
||||
text_signals_list,
|
||||
)
|
||||
return [str(row["id"]) for row in results]
|
||||
|
||||
async def bulk_insert_links(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
sorted_links: list[tuple],
|
||||
bank_id: str,
|
||||
nil_entity_uuid: str,
|
||||
exists_clause: str,
|
||||
chunk_size: int = 5000,
|
||||
) -> None:
|
||||
from_ids = [lnk[0] for lnk in sorted_links]
|
||||
to_ids = [lnk[1] for lnk in sorted_links]
|
||||
types = [lnk[2] for lnk in sorted_links]
|
||||
weights = [lnk[3] for lnk in sorted_links]
|
||||
entity_ids = [lnk[4] for lnk in sorted_links]
|
||||
|
||||
for chunk_start in range(0, len(sorted_links), chunk_size):
|
||||
chunk_end = min(chunk_start + chunk_size, len(sorted_links))
|
||||
await conn.execute(
|
||||
f"""
|
||||
INSERT INTO {table}
|
||||
(from_unit_id, to_unit_id, link_type, weight, entity_id, bank_id)
|
||||
SELECT f, t, tp, w, e, $6
|
||||
FROM unnest($1::uuid[], $2::uuid[], $3::text[], $4::float8[], $5::uuid[])
|
||||
AS t(f, t, tp, w, e)
|
||||
{exists_clause}
|
||||
ON CONFLICT (from_unit_id, to_unit_id, link_type,
|
||||
COALESCE(entity_id, '{nil_entity_uuid}'::uuid))
|
||||
DO NOTHING
|
||||
""",
|
||||
from_ids[chunk_start:chunk_end],
|
||||
to_ids[chunk_start:chunk_end],
|
||||
types[chunk_start:chunk_end],
|
||||
weights[chunk_start:chunk_end],
|
||||
entity_ids[chunk_start:chunk_end],
|
||||
bank_id,
|
||||
timeout=300,
|
||||
)
|
||||
|
||||
async def bulk_insert_entities(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
bank_id: str,
|
||||
entity_names: list[str],
|
||||
entity_dates: list,
|
||||
) -> dict[str, str]:
|
||||
inserted_rows = await conn.fetch(
|
||||
f"""
|
||||
INSERT INTO {table} (bank_id, canonical_name, first_seen, last_seen, mention_count)
|
||||
SELECT $1, name, COALESCE(event_date, now()), COALESCE(event_date, now()), 0
|
||||
FROM unnest($2::text[], $3::timestamptz[]) AS t(name, event_date)
|
||||
ON CONFLICT (bank_id, LOWER(canonical_name))
|
||||
DO NOTHING
|
||||
RETURNING id, LOWER(canonical_name) AS name_lower
|
||||
""",
|
||||
bank_id,
|
||||
entity_names,
|
||||
entity_dates,
|
||||
)
|
||||
return {row["name_lower"]: row["id"] for row in inserted_rows}
|
||||
|
||||
async def fetch_missing_entity_ids(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
bank_id: str,
|
||||
missing_names: list[str],
|
||||
) -> list[ResultRow]:
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT e.id, LOWER(e.canonical_name) AS name_lower, inputs.input_name
|
||||
FROM {table} e
|
||||
JOIN (
|
||||
SELECT LOWER(n) AS input_name_lower, n AS input_name
|
||||
FROM unnest($2::text[]) AS n
|
||||
) AS inputs ON LOWER(e.canonical_name) = inputs.input_name_lower
|
||||
WHERE e.bank_id = $1
|
||||
""",
|
||||
bank_id,
|
||||
missing_names,
|
||||
)
|
||||
|
||||
async def bulk_insert_unit_entities(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
unit_ids: list,
|
||||
entity_ids: list,
|
||||
) -> None:
|
||||
await conn.execute(
|
||||
f"""
|
||||
INSERT INTO {table} (unit_id, entity_id)
|
||||
SELECT u, e FROM unnest($1::uuid[], $2::uuid[]) AS t(u, e)
|
||||
ON CONFLICT DO NOTHING
|
||||
""",
|
||||
unit_ids,
|
||||
entity_ids,
|
||||
)
|
||||
|
||||
async def fetch_entity_unit_fanout(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
ue_table: str,
|
||||
entity_id_list: list[UUID],
|
||||
limit_per_entity: int,
|
||||
) -> list[ResultRow]:
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT e.entity_id, n.unit_id
|
||||
FROM unnest($1::uuid[]) AS e(entity_id)
|
||||
CROSS JOIN LATERAL (
|
||||
SELECT ue.unit_id
|
||||
FROM {ue_table} ue
|
||||
WHERE ue.entity_id = e.entity_id
|
||||
ORDER BY ue.unit_id DESC
|
||||
LIMIT $2
|
||||
) n
|
||||
""",
|
||||
entity_id_list,
|
||||
limit_per_entity,
|
||||
)
|
||||
|
||||
async def fetch_unit_dates(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
mu_table: str,
|
||||
unit_ids: list[str],
|
||||
) -> list[ResultRow]:
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT id, event_date, fact_type
|
||||
FROM {mu_table}
|
||||
WHERE id::text = ANY($1)
|
||||
""",
|
||||
unit_ids,
|
||||
)
|
||||
|
||||
async def fetch_temporal_neighbors(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
mu_table: str,
|
||||
bank_id: str,
|
||||
lateral_unit_ids: list,
|
||||
lateral_event_dates: list,
|
||||
lateral_fact_types: list,
|
||||
half_limit: int,
|
||||
batch_size: int = 500,
|
||||
) -> list[ResultRow]:
|
||||
rows: list[ResultRow] = []
|
||||
for start in range(0, len(lateral_unit_ids), batch_size):
|
||||
end = min(start + batch_size, len(lateral_unit_ids))
|
||||
# Exact v0.5.6 query shape: src.unit_id::text AS from_id,
|
||||
# combined.*, ABS(EXTRACT(...)), ROW_NUMBER PARTITION BY src.unit_id.
|
||||
batch_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT from_id, id, event_date, time_diff_hours FROM (
|
||||
SELECT src.unit_id::text AS from_id, combined.*,
|
||||
ROW_NUMBER() OVER (
|
||||
PARTITION BY src.unit_id
|
||||
ORDER BY combined.time_diff_hours
|
||||
) AS rn
|
||||
FROM unnest($1::uuid[], $2::timestamptz[], $3::text[])
|
||||
AS src(unit_id, event_date, fact_type)
|
||||
CROSS JOIN LATERAL (
|
||||
(SELECT mu.id, mu.event_date,
|
||||
ABS(EXTRACT(EPOCH FROM mu.event_date - src.event_date)) / 3600.0 AS time_diff_hours
|
||||
FROM {mu_table} mu
|
||||
WHERE mu.bank_id = $4
|
||||
AND mu.fact_type = src.fact_type
|
||||
AND mu.event_date <= src.event_date
|
||||
AND mu.id != src.unit_id
|
||||
ORDER BY mu.event_date DESC
|
||||
LIMIT $5)
|
||||
UNION ALL
|
||||
(SELECT mu.id, mu.event_date,
|
||||
ABS(EXTRACT(EPOCH FROM mu.event_date - src.event_date)) / 3600.0 AS time_diff_hours
|
||||
FROM {mu_table} mu
|
||||
WHERE mu.bank_id = $4
|
||||
AND mu.fact_type = src.fact_type
|
||||
AND mu.event_date > src.event_date
|
||||
AND mu.id != src.unit_id
|
||||
ORDER BY mu.event_date ASC
|
||||
LIMIT $5)
|
||||
) combined
|
||||
) ranked
|
||||
WHERE rn <= $5
|
||||
""",
|
||||
lateral_unit_ids[start:end],
|
||||
lateral_event_dates[start:end],
|
||||
lateral_fact_types[start:end],
|
||||
bank_id,
|
||||
half_limit,
|
||||
)
|
||||
rows.extend(batch_rows)
|
||||
return rows
|
||||
|
||||
def build_entity_expansion_cte(
|
||||
self,
|
||||
mu_table: str,
|
||||
ue_table: str,
|
||||
per_entity_limit: int,
|
||||
) -> str:
|
||||
return f"""
|
||||
seed_entities AS (
|
||||
SELECT DISTINCT ue.entity_id
|
||||
FROM {ue_table} ue
|
||||
WHERE ue.unit_id = ANY($1::uuid[])
|
||||
),
|
||||
entity_expanded 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.proof_count,
|
||||
COUNT(DISTINCT se.entity_id)::float AS score,
|
||||
'entity'::text AS source
|
||||
FROM seed_entities se
|
||||
CROSS JOIN LATERAL (
|
||||
SELECT ue_target.unit_id
|
||||
FROM {ue_table} 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_table} mu ON mu.id = t.unit_id
|
||||
WHERE mu.fact_type = $2
|
||||
GROUP BY mu.id
|
||||
ORDER BY score DESC
|
||||
LIMIT $3
|
||||
)"""
|
||||
|
||||
def build_semantic_causal_cte(
|
||||
self,
|
||||
ml_table: str,
|
||||
mu_table: str,
|
||||
) -> str:
|
||||
# Exact v0.5.6 query shape: GROUP BY + MAX(weight) for semantic,
|
||||
# DISTINCT ON for causal.
|
||||
return f"""
|
||||
semantic_expanded AS (
|
||||
SELECT
|
||||
id, text, context, event_date, occurred_start,
|
||||
occurred_end, mentioned_at,
|
||||
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.proof_count,
|
||||
ml.weight
|
||||
FROM {ml_table} ml
|
||||
JOIN {mu_table} 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 = $2
|
||||
AND mu.id != ALL($1::uuid[])
|
||||
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, mu.proof_count,
|
||||
ml.weight
|
||||
FROM {ml_table} ml
|
||||
JOIN {mu_table} 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 = $2
|
||||
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, proof_count
|
||||
ORDER BY score DESC
|
||||
LIMIT $3
|
||||
),
|
||||
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, mu.proof_count,
|
||||
ml.weight AS score,
|
||||
'causal'::text AS source
|
||||
FROM {ml_table} ml
|
||||
JOIN {mu_table} 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')
|
||||
AND mu.fact_type = $2
|
||||
ORDER BY mu.id, ml.weight DESC
|
||||
LIMIT $3
|
||||
)"""
|
||||
|
||||
async def expand_observations(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
mu_table: str,
|
||||
ue_table: str,
|
||||
ml_table: str,
|
||||
seed_ids: list,
|
||||
budget: int,
|
||||
per_entity_limit: int,
|
||||
) -> tuple[list[ResultRow], list[ResultRow], list[ResultRow]]:
|
||||
# Entity expansion via observation_sources junction table.
|
||||
# Previously used PG-specific unnest(source_memory_ids) and array
|
||||
# overlap (&&). The junction table approach is portable across backends.
|
||||
from ..schema import fq_table
|
||||
|
||||
obs_sources_table = fq_table("observation_sources")
|
||||
entity_rows = await conn.fetch(
|
||||
f"""
|
||||
WITH source_ids AS (
|
||||
SELECT DISTINCT os.source_id
|
||||
FROM {obs_sources_table} os
|
||||
WHERE os.observation_id = ANY($1::uuid[])
|
||||
),
|
||||
source_entities AS (
|
||||
SELECT DISTINCT ue_seed.entity_id
|
||||
FROM source_ids si
|
||||
JOIN {ue_table} ue_seed ON ue_seed.unit_id = si.source_id
|
||||
),
|
||||
connected_sources AS (
|
||||
SELECT DISTINCT t.unit_id AS source_id
|
||||
FROM source_entities se
|
||||
CROSS JOIN LATERAL (
|
||||
SELECT ue_target.unit_id
|
||||
FROM {ue_table} ue_target
|
||||
WHERE ue_target.entity_id = se.entity_id
|
||||
ORDER BY ue_target.unit_id DESC
|
||||
LIMIT {per_entity_limit}
|
||||
) t
|
||||
WHERE t.unit_id NOT IN (SELECT source_id FROM source_ids)
|
||||
)
|
||||
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.proof_count,
|
||||
(SELECT COUNT(*)
|
||||
FROM {obs_sources_table} os2
|
||||
WHERE os2.observation_id = mu.id
|
||||
AND os2.source_id IN (SELECT source_id FROM connected_sources)
|
||||
)::float AS score
|
||||
FROM {mu_table} mu
|
||||
WHERE mu.fact_type = 'observation'
|
||||
AND mu.id != ALL($1::uuid[])
|
||||
AND EXISTS (
|
||||
SELECT 1 FROM {obs_sources_table} os3
|
||||
WHERE os3.observation_id = mu.id
|
||||
AND os3.source_id IN (SELECT source_id FROM connected_sources)
|
||||
)
|
||||
ORDER BY score DESC
|
||||
LIMIT $2
|
||||
""",
|
||||
seed_ids,
|
||||
budget,
|
||||
)
|
||||
|
||||
# Exact v0.5.6 query shape: GROUP BY + MAX(weight) for semantic,
|
||||
# DISTINCT ON for causal, hardcoded to fact_type='observation'.
|
||||
sem_causal_rows = await conn.fetch(
|
||||
f"""
|
||||
WITH semantic_expanded AS (
|
||||
SELECT
|
||||
id, text, context, event_date, occurred_start,
|
||||
occurred_end, mentioned_at,
|
||||
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.proof_count, ml.weight
|
||||
FROM {ml_table} ml JOIN {mu_table} 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'
|
||||
AND mu.id != ALL($1::uuid[])
|
||||
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, mu.proof_count, ml.weight
|
||||
FROM {ml_table} ml JOIN {mu_table} 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, 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, mu.proof_count, ml.weight AS score, 'causal'::text AS source
|
||||
FROM {ml_table} ml JOIN {mu_table} 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')
|
||||
AND mu.fact_type = 'observation'
|
||||
ORDER BY mu.id, ml.weight DESC LIMIT $2
|
||||
)
|
||||
SELECT * FROM semantic_expanded
|
||||
UNION ALL
|
||||
SELECT * FROM causal_expanded
|
||||
""",
|
||||
seed_ids,
|
||||
budget,
|
||||
)
|
||||
|
||||
semantic_rows = [r for r in sem_causal_rows if r["source"] == "semantic"]
|
||||
causal_rows = [r for r in sem_causal_rows if r["source"] == "causal"]
|
||||
return list(entity_rows), semantic_rows, causal_rows
|
||||
|
||||
def build_tag_listing_parts(self, mu_table: str) -> TagListingParts:
|
||||
return TagListingParts(
|
||||
tag_source=f"{mu_table}, unnest(tags) AS tag",
|
||||
non_empty_check="AND tags IS NOT NULL AND tags != '{}'",
|
||||
tag_col="tag",
|
||||
bank_prefix="",
|
||||
)
|
||||
|
||||
async def create_bank_vector_indexes(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
bank_id: str,
|
||||
internal_id: str,
|
||||
index_clause: str,
|
||||
fact_types: dict[str, str],
|
||||
) -> None:
|
||||
escaped = bank_id.replace("'", "''")
|
||||
for ft, suffix in fact_types.items():
|
||||
uid = str(internal_id).replace("-", "")[:16]
|
||||
idx = f"idx_mu_emb_{suffix}_{uid}"
|
||||
await conn.execute(
|
||||
f"CREATE INDEX IF NOT EXISTS {idx} "
|
||||
f"ON {table} {index_clause} "
|
||||
f"WHERE fact_type = '{ft}' AND bank_id = '{escaped}'"
|
||||
)
|
||||
|
||||
async def drop_bank_vector_indexes(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
schema: str,
|
||||
internal_id: str,
|
||||
fact_types: dict[str, str],
|
||||
) -> None:
|
||||
for ft, suffix in fact_types.items():
|
||||
uid = str(internal_id).replace("-", "")[:16]
|
||||
idx = f"idx_mu_emb_{suffix}_{uid}"
|
||||
await conn.execute(f"DROP INDEX IF EXISTS {schema}.{idx}")
|
||||
|
||||
def get_entity_resolution_strategy(self) -> str:
|
||||
return "trigram"
|
||||
|
||||
# -- Webhook operations ------------------------------------------------
|
||||
|
||||
async def create_webhook(
|
||||
self,
|
||||
conn,
|
||||
table,
|
||||
webhook_id,
|
||||
bank_id,
|
||||
url,
|
||||
secret,
|
||||
event_types,
|
||||
enabled,
|
||||
http_config_json,
|
||||
):
|
||||
return await conn.fetchrow(
|
||||
f"""
|
||||
INSERT INTO {table}
|
||||
(id, bank_id, url, secret, event_types, enabled, http_config, created_at, updated_at)
|
||||
VALUES ($1, $2, $3, $4, $5, $6, $7::jsonb, NOW(), NOW())
|
||||
RETURNING id, bank_id, url, secret, event_types, enabled,
|
||||
http_config::text, created_at::text, updated_at::text
|
||||
""",
|
||||
webhook_id,
|
||||
bank_id,
|
||||
url,
|
||||
secret,
|
||||
event_types,
|
||||
enabled,
|
||||
http_config_json,
|
||||
)
|
||||
|
||||
async def list_webhooks_for_bank(self, conn, table, bank_id):
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT id, bank_id, url, secret, event_types, enabled,
|
||||
http_config::text, created_at::text, updated_at::text
|
||||
FROM {table}
|
||||
WHERE bank_id = $1
|
||||
ORDER BY created_at
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
|
||||
async def get_webhooks_for_dispatch(self, conn, webhook_table, bank_id):
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT id, bank_id, url, secret, event_types, enabled, http_config::text
|
||||
FROM {webhook_table}
|
||||
WHERE (bank_id = $1 OR bank_id IS NULL) AND enabled = true
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
|
||||
async def update_webhook(self, conn, table, webhook_id, bank_id, set_clauses, params):
|
||||
set_clauses_with_ts = set_clauses + ["updated_at = NOW()"]
|
||||
return await conn.fetchrow(
|
||||
f"""
|
||||
UPDATE {table}
|
||||
SET {", ".join(set_clauses_with_ts)}
|
||||
WHERE id = $1 AND bank_id = $2
|
||||
RETURNING id, bank_id, url, secret, event_types, enabled,
|
||||
http_config::text, created_at::text, updated_at::text
|
||||
""",
|
||||
*params,
|
||||
)
|
||||
|
||||
async def delete_webhook(self, conn, table, webhook_id, bank_id):
|
||||
result = await conn.execute(
|
||||
f"DELETE FROM {table} WHERE id = $1 AND bank_id = $2",
|
||||
webhook_id,
|
||||
bank_id,
|
||||
)
|
||||
return int(result.split()[-1]) > 0 if result else False
|
||||
|
||||
async def list_webhook_deliveries(self, conn, ops_table, webhook_id, bank_id, limit, cursor):
|
||||
fetch_limit = limit + 1
|
||||
if cursor:
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, status, retry_count, next_retry_at::text,
|
||||
error_message, task_payload, result_metadata::text, created_at::text, updated_at::text
|
||||
FROM {ops_table}
|
||||
WHERE operation_type = 'webhook_delivery'
|
||||
AND bank_id = $1
|
||||
AND task_payload->>'webhook_id' = $2
|
||||
AND created_at < $3::timestamptz
|
||||
ORDER BY created_at DESC
|
||||
LIMIT $4
|
||||
""",
|
||||
bank_id,
|
||||
webhook_id,
|
||||
cursor,
|
||||
fetch_limit,
|
||||
)
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, status, retry_count, next_retry_at::text,
|
||||
error_message, task_payload, result_metadata::text, created_at::text, updated_at::text
|
||||
FROM {ops_table}
|
||||
WHERE operation_type = 'webhook_delivery'
|
||||
AND bank_id = $1
|
||||
AND task_payload->>'webhook_id' = $2
|
||||
ORDER BY created_at DESC
|
||||
LIMIT $3
|
||||
""",
|
||||
bank_id,
|
||||
webhook_id,
|
||||
fetch_limit,
|
||||
)
|
||||
|
||||
async def insert_webhook_delivery_task(self, conn, ops_table, operation_id, bank_id, payload_json, timestamp):
|
||||
await conn.execute(
|
||||
f"""
|
||||
INSERT INTO {ops_table}
|
||||
(operation_id, bank_id, operation_type, status, task_payload, result_metadata, created_at, updated_at)
|
||||
VALUES ($1, $2, 'webhook_delivery', 'pending', $3::jsonb, '{{}}'::jsonb, $4, $4)
|
||||
""",
|
||||
operation_id,
|
||||
bank_id,
|
||||
payload_json,
|
||||
timestamp,
|
||||
)
|
||||
|
||||
# -- Task claiming operations ------------------------------------------
|
||||
|
||||
async def claim_tasks(self, conn, table, worker_id, reserved_limits, shared_limit):
|
||||
all_rows = []
|
||||
claimed_ids = []
|
||||
|
||||
# --- Phase 1: claim from reserved pools ---
|
||||
for op_type, limit in reserved_limits.items():
|
||||
if limit <= 0:
|
||||
continue
|
||||
|
||||
if op_type == "consolidation":
|
||||
busy_banks = await conn.fetch(
|
||||
f"""
|
||||
SELECT DISTINCT bank_id FROM {table}
|
||||
WHERE operation_type = 'consolidation' AND status = 'processing'
|
||||
""",
|
||||
)
|
||||
busy_bank_ids = [r["bank_id"] for r in busy_banks]
|
||||
|
||||
if busy_bank_ids:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND bank_id != ALL($1::text[])
|
||||
ORDER BY created_at
|
||||
LIMIT $2
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
busy_bank_ids,
|
||||
limit,
|
||||
)
|
||||
else:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
ORDER BY created_at
|
||||
LIMIT $1
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
limit,
|
||||
)
|
||||
else:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = $1
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
ORDER BY created_at
|
||||
LIMIT $2
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
op_type,
|
||||
limit,
|
||||
)
|
||||
|
||||
for row in rows:
|
||||
claimed_ids.append(row["operation_id"])
|
||||
all_rows.append(row)
|
||||
|
||||
# --- Phase 2: claim from shared pool ---
|
||||
remaining_shared = shared_limit
|
||||
if remaining_shared > 0:
|
||||
# 2a. Non-consolidation tasks
|
||||
if claimed_ids:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type != 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND operation_id != ALL($1::uuid[])
|
||||
ORDER BY created_at
|
||||
LIMIT $2
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
claimed_ids,
|
||||
remaining_shared,
|
||||
)
|
||||
else:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type != 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
ORDER BY created_at
|
||||
LIMIT $1
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
remaining_shared,
|
||||
)
|
||||
|
||||
for row in rows:
|
||||
claimed_ids.append(row["operation_id"])
|
||||
all_rows.append(row)
|
||||
remaining_shared -= len(rows)
|
||||
|
||||
# 2b. Consolidation tasks (with bank-serialization)
|
||||
if remaining_shared > 0:
|
||||
busy_banks_2 = await conn.fetch(
|
||||
f"""
|
||||
SELECT DISTINCT bank_id FROM {table}
|
||||
WHERE operation_type = 'consolidation' AND status = 'processing'
|
||||
""",
|
||||
)
|
||||
busy_bank_ids_2 = [r["bank_id"] for r in busy_banks_2]
|
||||
|
||||
if claimed_ids:
|
||||
if busy_bank_ids_2:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND operation_id != ALL($1::uuid[])
|
||||
AND bank_id != ALL($2::text[])
|
||||
ORDER BY created_at
|
||||
LIMIT $3
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
claimed_ids,
|
||||
busy_bank_ids_2,
|
||||
remaining_shared,
|
||||
)
|
||||
else:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND operation_id != ALL($1::uuid[])
|
||||
ORDER BY created_at
|
||||
LIMIT $2
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
claimed_ids,
|
||||
remaining_shared,
|
||||
)
|
||||
else:
|
||||
if busy_bank_ids_2:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND bank_id != ALL($1::text[])
|
||||
ORDER BY created_at
|
||||
LIMIT $2
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
busy_bank_ids_2,
|
||||
remaining_shared,
|
||||
)
|
||||
else:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
ORDER BY created_at
|
||||
LIMIT $1
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
remaining_shared,
|
||||
)
|
||||
|
||||
for row in rows:
|
||||
claimed_ids.append(row["operation_id"])
|
||||
all_rows.append(row)
|
||||
|
||||
if not all_rows:
|
||||
return []
|
||||
|
||||
# Mark all claimed rows as processing
|
||||
operation_ids = [row["operation_id"] for row in all_rows]
|
||||
await conn.execute(
|
||||
f"""
|
||||
UPDATE {table}
|
||||
SET status = 'processing', worker_id = $1, claimed_at = now(), updated_at = now()
|
||||
WHERE operation_id = ANY($2)
|
||||
""",
|
||||
worker_id,
|
||||
operation_ids,
|
||||
)
|
||||
|
||||
return all_rows
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,128 @@
|
||||
"""PostgreSQL backend implementation using asyncpg.
|
||||
|
||||
Wraps asyncpg's pool and connection objects behind the DatabaseBackend
|
||||
and DatabaseConnection interfaces.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from collections.abc import AsyncIterator
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import Any
|
||||
|
||||
import asyncpg # noqa: F401
|
||||
|
||||
from .base import DatabaseBackend, DatabaseConnection
|
||||
from .result import ResultRow
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class PostgresConnection(DatabaseConnection):
|
||||
"""DatabaseConnection wrapper around an asyncpg.Connection."""
|
||||
|
||||
__slots__ = ("_conn",)
|
||||
|
||||
def __init__(self, conn: asyncpg.Connection) -> None:
|
||||
self._conn = conn
|
||||
|
||||
@asynccontextmanager
|
||||
async def transaction(self) -> AsyncIterator["PostgresConnection"]:
|
||||
async with self._conn.transaction():
|
||||
yield self
|
||||
|
||||
async def execute(self, query: str, *args: Any, timeout: float | None = None) -> str:
|
||||
return await self._conn.execute(query, *args, timeout=timeout)
|
||||
|
||||
async def executemany(self, query: str, args: list[tuple[Any, ...]], *, timeout: float | None = None) -> None:
|
||||
await self._conn.executemany(query, args, timeout=timeout)
|
||||
|
||||
async def fetch(self, query: str, *args: Any, timeout: float | None = None) -> list[ResultRow]:
|
||||
rows = await self._conn.fetch(query, *args, timeout=timeout)
|
||||
return [ResultRow(row) for row in rows]
|
||||
|
||||
async def fetchrow(self, query: str, *args: Any, timeout: float | None = None) -> ResultRow | None:
|
||||
row = await self._conn.fetchrow(query, *args, timeout=timeout)
|
||||
if row is None:
|
||||
return None
|
||||
return ResultRow(row)
|
||||
|
||||
async def fetchval(self, query: str, *args: Any, column: int = 0, timeout: float | None = None) -> Any:
|
||||
return await self._conn.fetchval(query, *args, column=column, timeout=timeout)
|
||||
|
||||
async def copy_records_to_table(
|
||||
self,
|
||||
table_name: str,
|
||||
*,
|
||||
records: list[tuple[Any, ...]],
|
||||
columns: list[str],
|
||||
timeout: float | None = None,
|
||||
) -> None:
|
||||
"""Use asyncpg's native COPY for fast bulk loading."""
|
||||
await self._conn.copy_records_to_table(table_name, records=records, columns=columns, timeout=timeout)
|
||||
|
||||
|
||||
class PostgreSQLBackend(DatabaseBackend):
|
||||
"""DatabaseBackend implementation wrapping an asyncpg connection pool."""
|
||||
|
||||
def run_migrations(self, dsn: str, *, schema: str | None = None) -> None:
|
||||
"""Run Alembic migrations for PostgreSQL."""
|
||||
from ...config import get_config
|
||||
from ...migrations import run_migrations
|
||||
|
||||
config = get_config()
|
||||
run_migrations(dsn, schema=schema, migration_database_url=config.migration_database_url)
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._pool: asyncpg.Pool | None = None
|
||||
|
||||
async def initialize(
|
||||
self,
|
||||
dsn: str,
|
||||
*,
|
||||
min_size: int = 5,
|
||||
max_size: int = 20,
|
||||
command_timeout: float = 300,
|
||||
acquire_timeout: float = 30,
|
||||
statement_cache_size: int = 0,
|
||||
init_callback: Any | None = None,
|
||||
) -> None:
|
||||
self._pool = await asyncpg.create_pool(
|
||||
dsn,
|
||||
min_size=min_size,
|
||||
max_size=max_size,
|
||||
command_timeout=command_timeout,
|
||||
statement_cache_size=statement_cache_size,
|
||||
timeout=acquire_timeout,
|
||||
init=init_callback,
|
||||
)
|
||||
logger.info(
|
||||
f"PostgreSQL pool created (min={min_size}, max={max_size}, "
|
||||
f"cmd_timeout={command_timeout}s, acquire_timeout={acquire_timeout}s)"
|
||||
)
|
||||
|
||||
async def shutdown(self) -> None:
|
||||
if self._pool is not None:
|
||||
await self._pool.close()
|
||||
self._pool = None
|
||||
logger.info("PostgreSQL pool closed")
|
||||
|
||||
@asynccontextmanager
|
||||
async def acquire(self) -> AsyncIterator[PostgresConnection]:
|
||||
pool = self._ensure_pool()
|
||||
async with pool.acquire() as conn:
|
||||
yield PostgresConnection(conn)
|
||||
|
||||
@asynccontextmanager
|
||||
async def transaction(self) -> AsyncIterator[PostgresConnection]:
|
||||
pool = self._ensure_pool()
|
||||
async with pool.acquire() as conn:
|
||||
async with conn.transaction():
|
||||
yield PostgresConnection(conn)
|
||||
|
||||
def get_pool(self) -> asyncpg.Pool:
|
||||
return self._ensure_pool()
|
||||
|
||||
def _ensure_pool(self) -> asyncpg.Pool:
|
||||
if self._pool is None:
|
||||
raise RuntimeError("PostgreSQLBackend is not initialized. Call initialize() first.")
|
||||
return self._pool
|
||||
@@ -0,0 +1,105 @@
|
||||
"""Uniform row wrapper over heterogeneous database drivers.
|
||||
|
||||
ResultRow provides dict-like access to database rows regardless of whether
|
||||
the underlying driver returns asyncpg.Record, oracledb rows, or plain dicts.
|
||||
"""
|
||||
|
||||
from typing import Any
|
||||
|
||||
|
||||
class ResultRow:
|
||||
"""Dict-like wrapper over database rows.
|
||||
|
||||
Supports both key-based access (row["col"]) and attribute access (row.col).
|
||||
Wraps asyncpg.Record, oracledb named-tuple rows, or plain dicts.
|
||||
"""
|
||||
|
||||
__slots__ = ("_data",)
|
||||
|
||||
def __init__(self, data: Any) -> None:
|
||||
"""Wrap a row from any database driver.
|
||||
|
||||
Args:
|
||||
data: The raw row object (asyncpg.Record, dict, named tuple, etc.)
|
||||
"""
|
||||
object.__setattr__(self, "_data", data)
|
||||
|
||||
# -- dict-like access ------------------------------------------------
|
||||
|
||||
def __getitem__(self, key: str | int) -> Any:
|
||||
"""Get a value by column name or index."""
|
||||
data = object.__getattribute__(self, "_data")
|
||||
if isinstance(data, dict):
|
||||
return data[key]
|
||||
return data[key]
|
||||
|
||||
def __getattr__(self, key: str) -> Any:
|
||||
"""Get a value by attribute name (for convenience)."""
|
||||
data = object.__getattribute__(self, "_data")
|
||||
if isinstance(data, dict):
|
||||
try:
|
||||
return data[key]
|
||||
except KeyError:
|
||||
raise AttributeError(key) from None
|
||||
# asyncpg.Record and named tuples support key-based access
|
||||
try:
|
||||
return data[key]
|
||||
except (KeyError, TypeError):
|
||||
raise AttributeError(key) from None
|
||||
|
||||
def get(self, key: str, default: Any = None) -> Any:
|
||||
"""Get a value with a default (like dict.get)."""
|
||||
try:
|
||||
return self[key]
|
||||
except (KeyError, IndexError):
|
||||
return default
|
||||
|
||||
def keys(self) -> list[str]:
|
||||
"""Return column names."""
|
||||
data = object.__getattribute__(self, "_data")
|
||||
if isinstance(data, dict):
|
||||
return list(data.keys())
|
||||
# asyncpg.Record has .keys()
|
||||
if hasattr(data, "keys"):
|
||||
return list(data.keys())
|
||||
return []
|
||||
|
||||
def values(self) -> list[Any]:
|
||||
"""Return column values."""
|
||||
data = object.__getattribute__(self, "_data")
|
||||
if isinstance(data, dict):
|
||||
return list(data.values())
|
||||
if hasattr(data, "values"):
|
||||
return list(data.values())
|
||||
return []
|
||||
|
||||
def items(self) -> list[tuple[str, Any]]:
|
||||
"""Return (key, value) pairs."""
|
||||
data = object.__getattribute__(self, "_data")
|
||||
if isinstance(data, dict):
|
||||
return list(data.items())
|
||||
if hasattr(data, "items"):
|
||||
return list(data.items())
|
||||
return list(zip(self.keys(), self.values()))
|
||||
|
||||
# -- representation --------------------------------------------------
|
||||
|
||||
def __repr__(self) -> str:
|
||||
data = object.__getattribute__(self, "_data")
|
||||
return f"ResultRow({data!r})"
|
||||
|
||||
def __contains__(self, key: str) -> bool:
|
||||
data = object.__getattribute__(self, "_data")
|
||||
if isinstance(data, dict):
|
||||
return key in data
|
||||
if hasattr(data, "keys"):
|
||||
return key in data.keys()
|
||||
return False
|
||||
|
||||
def __len__(self) -> int:
|
||||
data = object.__getattribute__(self, "_data")
|
||||
return len(data)
|
||||
|
||||
def __bool__(self) -> bool:
|
||||
data = object.__getattribute__(self, "_data")
|
||||
return bool(data)
|
||||
@@ -11,10 +11,7 @@ import logging
|
||||
import uuid
|
||||
from contextlib import asynccontextmanager
|
||||
from dataclasses import dataclass, field
|
||||
from typing import TYPE_CHECKING, AsyncIterator
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import asyncpg
|
||||
from typing import Any, AsyncIterator
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -122,14 +119,14 @@ class BudgetedOperation:
|
||||
return self._manager._get_budget(self.operation_id)
|
||||
|
||||
@asynccontextmanager
|
||||
async def acquire(self, pool: "asyncpg.Pool") -> AsyncIterator["asyncpg.Connection"]:
|
||||
async def acquire(self, pool: Any) -> AsyncIterator[Any]:
|
||||
"""
|
||||
Acquire a connection within the operation's budget.
|
||||
|
||||
Blocks if the operation has reached its connection limit.
|
||||
|
||||
Args:
|
||||
pool: asyncpg connection pool
|
||||
pool: asyncpg connection pool or DatabaseBackend
|
||||
|
||||
Yields:
|
||||
Database connection
|
||||
@@ -137,14 +134,22 @@ class BudgetedOperation:
|
||||
budget = self.budget
|
||||
async with budget.semaphore:
|
||||
budget.active_count += 1
|
||||
conn = await pool.acquire()
|
||||
try:
|
||||
yield conn
|
||||
from .db.base import DatabaseBackend
|
||||
|
||||
if isinstance(pool, DatabaseBackend):
|
||||
async with pool.acquire() as conn:
|
||||
yield conn
|
||||
else:
|
||||
conn = await pool.acquire()
|
||||
try:
|
||||
yield conn
|
||||
finally:
|
||||
await pool.release(conn)
|
||||
finally:
|
||||
budget.active_count -= 1
|
||||
await pool.release(conn)
|
||||
|
||||
def wrap_pool(self, pool: "asyncpg.Pool") -> "BudgetedPool":
|
||||
def wrap_pool(self, pool: Any) -> "BudgetedPool":
|
||||
"""
|
||||
Wrap a pool with this operation's budget.
|
||||
|
||||
@@ -161,17 +166,18 @@ class BudgetedOperation:
|
||||
|
||||
async def acquire_many(
|
||||
self,
|
||||
pool: "asyncpg.Pool",
|
||||
pool: Any,
|
||||
count: int,
|
||||
) -> AsyncIterator[list["asyncpg.Connection"]]:
|
||||
) -> AsyncIterator[list[Any]]:
|
||||
"""
|
||||
Acquire multiple connections within the budget.
|
||||
|
||||
Note: This acquires connections sequentially to respect the budget.
|
||||
For parallel acquisition, use multiple acquire() calls with asyncio.gather().
|
||||
This method is intended for use with raw asyncpg pools only, not DatabaseBackend.
|
||||
|
||||
Args:
|
||||
pool: asyncpg connection pool
|
||||
pool: asyncpg connection pool (raw pool only)
|
||||
count: Number of connections to acquire
|
||||
|
||||
Yields:
|
||||
@@ -249,29 +255,42 @@ class BudgetedPool:
|
||||
await some_function(budgeted_pool, ...)
|
||||
"""
|
||||
|
||||
def __init__(self, pool: "asyncpg.Pool", operation: BudgetedOperation):
|
||||
_wraps_backend = True
|
||||
|
||||
def __init__(self, pool: Any, operation: BudgetedOperation):
|
||||
self._pool = pool
|
||||
self._operation = operation
|
||||
|
||||
async def acquire(self) -> "asyncpg.Connection":
|
||||
@asynccontextmanager
|
||||
async def acquire(self) -> AsyncIterator[Any]:
|
||||
"""
|
||||
Acquire a connection within the budget.
|
||||
Acquire a connection within the budget as an async context manager.
|
||||
|
||||
Note: Caller must release the connection when done.
|
||||
Prefer using as context manager via acquire_with_retry or op.acquire().
|
||||
The connection is automatically released when the context exits.
|
||||
"""
|
||||
budget = self._operation.budget
|
||||
await budget.semaphore.acquire()
|
||||
budget.active_count += 1
|
||||
try:
|
||||
return await self._pool.acquire()
|
||||
from .db.base import DatabaseBackend
|
||||
|
||||
if isinstance(self._pool, DatabaseBackend):
|
||||
async with self._pool.acquire() as conn:
|
||||
yield conn
|
||||
else:
|
||||
conn = await self._pool.acquire()
|
||||
try:
|
||||
yield conn
|
||||
finally:
|
||||
await self._pool.release(conn)
|
||||
except Exception:
|
||||
raise
|
||||
finally:
|
||||
budget.active_count -= 1
|
||||
budget.semaphore.release()
|
||||
raise
|
||||
|
||||
async def release(self, conn: "asyncpg.Connection") -> None:
|
||||
"""Release a connection back to the pool."""
|
||||
async def release(self, conn: Any) -> None:
|
||||
"""Release a connection back to the pool (legacy path only)."""
|
||||
budget = self._operation.budget
|
||||
try:
|
||||
await self._pool.release(conn)
|
||||
|
||||
@@ -4,9 +4,10 @@ Database utility functions for connection management with retry logic.
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import time
|
||||
from collections.abc import AsyncIterator
|
||||
from contextlib import asynccontextmanager
|
||||
|
||||
import asyncpg
|
||||
from typing import Any
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -15,24 +16,43 @@ DEFAULT_MAX_RETRIES = 3
|
||||
DEFAULT_BASE_DELAY = 0.5 # seconds
|
||||
DEFAULT_MAX_DELAY = 5.0 # seconds
|
||||
|
||||
# Exceptions that indicate transient connection issues worth retrying
|
||||
RETRYABLE_EXCEPTIONS = (
|
||||
asyncpg.exceptions.InterfaceError,
|
||||
asyncpg.exceptions.ConnectionDoesNotExistError,
|
||||
asyncpg.exceptions.TooManyConnectionsError,
|
||||
asyncpg.exceptions.DeadlockDetectedError,
|
||||
OSError,
|
||||
ConnectionError,
|
||||
asyncio.TimeoutError,
|
||||
# Retryable exception types (checked by class name to avoid hard imports)
|
||||
_RETRYABLE_EXCEPTION_NAMES = frozenset(
|
||||
{
|
||||
"InterfaceError",
|
||||
"ConnectionDoesNotExistError",
|
||||
"TooManyConnectionsError",
|
||||
"DeadlockDetectedError",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _is_oracle_deadlock(exc: BaseException) -> bool:
|
||||
"""Check if an exception is an Oracle ORA-00060 deadlock."""
|
||||
try:
|
||||
import oracledb # type: ignore[import-not-found]
|
||||
except ImportError:
|
||||
return False
|
||||
if isinstance(exc, oracledb.DatabaseError) and exc.args:
|
||||
err = exc.args[0]
|
||||
return getattr(err, "code", None) == 60 # ORA-00060
|
||||
return False
|
||||
|
||||
|
||||
def _is_retryable(exc: BaseException) -> bool:
|
||||
"""Check if an exception is retryable (transient connection issue)."""
|
||||
if isinstance(exc, (OSError, ConnectionError, asyncio.TimeoutError)):
|
||||
return True
|
||||
if type(exc).__name__ in _RETRYABLE_EXCEPTION_NAMES:
|
||||
return True
|
||||
return _is_oracle_deadlock(exc)
|
||||
|
||||
|
||||
async def retry_with_backoff(
|
||||
func,
|
||||
max_retries: int = DEFAULT_MAX_RETRIES,
|
||||
base_delay: float = DEFAULT_BASE_DELAY,
|
||||
max_delay: float = DEFAULT_MAX_DELAY,
|
||||
retryable_exceptions: tuple = RETRYABLE_EXCEPTIONS,
|
||||
):
|
||||
"""
|
||||
Execute an async function with exponential backoff retry.
|
||||
@@ -42,7 +62,6 @@ async def retry_with_backoff(
|
||||
max_retries: Maximum number of retry attempts
|
||||
base_delay: Initial delay between retries (seconds)
|
||||
max_delay: Maximum delay between retries (seconds)
|
||||
retryable_exceptions: Tuple of exception types to retry on
|
||||
|
||||
Returns:
|
||||
Result of the function
|
||||
@@ -54,13 +73,16 @@ async def retry_with_backoff(
|
||||
for attempt in range(max_retries + 1):
|
||||
try:
|
||||
return await func()
|
||||
except retryable_exceptions as e:
|
||||
except Exception as e:
|
||||
if not _is_retryable(e):
|
||||
raise
|
||||
last_exception = e
|
||||
if attempt < max_retries:
|
||||
delay = min(base_delay * (2**attempt), max_delay)
|
||||
if isinstance(e, asyncpg.exceptions.DeadlockDetectedError):
|
||||
if type(e).__name__ == "DeadlockDetectedError" or _is_oracle_deadlock(e):
|
||||
logger.warning(
|
||||
f"Deadlock detected during parallel document processing — this is expected and will resolve automatically "
|
||||
"Deadlock detected during parallel document processing — "
|
||||
"this is expected and will resolve automatically "
|
||||
f"(attempt {attempt + 1}/{max_retries + 1}, retrying in {delay:.1f}s)"
|
||||
)
|
||||
else:
|
||||
@@ -75,38 +97,68 @@ async def retry_with_backoff(
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def acquire_with_retry(pool: asyncpg.Pool, max_retries: int = DEFAULT_MAX_RETRIES):
|
||||
async def acquire_with_retry(backend_or_pool: Any, max_retries: int = DEFAULT_MAX_RETRIES) -> AsyncIterator[Any]:
|
||||
"""
|
||||
Async context manager to acquire a connection with retry logic.
|
||||
Async context manager to acquire a database connection with retry logic.
|
||||
|
||||
Accepts either a DatabaseBackend or a raw asyncpg.Pool for backward compatibility.
|
||||
|
||||
Usage:
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
async with acquire_with_retry(backend) as conn:
|
||||
await conn.execute(...)
|
||||
|
||||
Args:
|
||||
pool: The asyncpg connection pool
|
||||
backend_or_pool: A DatabaseBackend instance or asyncpg.Pool
|
||||
max_retries: Maximum number of retry attempts
|
||||
|
||||
Yields:
|
||||
An asyncpg connection
|
||||
A DatabaseConnection (if backend) or asyncpg.Connection (if pool)
|
||||
"""
|
||||
import time
|
||||
from .db.base import DatabaseBackend
|
||||
|
||||
start = time.time()
|
||||
if isinstance(backend_or_pool, DatabaseBackend) or getattr(backend_or_pool, "_wraps_backend", False):
|
||||
# Use the backend's acquire context manager with retry
|
||||
start = time.time()
|
||||
last_exception = None
|
||||
for attempt in range(max_retries + 1):
|
||||
try:
|
||||
async with backend_or_pool.acquire() as conn:
|
||||
acquire_time = time.time() - start
|
||||
if acquire_time > 0.05:
|
||||
logger.warning(f"[DB POOL] Slow acquire: {acquire_time:.3f}s")
|
||||
yield conn
|
||||
return
|
||||
except Exception as e:
|
||||
if not _is_retryable(e):
|
||||
raise
|
||||
last_exception = e
|
||||
if attempt < max_retries:
|
||||
delay = min(DEFAULT_BASE_DELAY * (2**attempt), DEFAULT_MAX_DELAY)
|
||||
logger.warning(
|
||||
f"Database acquire failed (attempt {attempt + 1}/{max_retries + 1}): {e}. "
|
||||
f"Retrying in {delay:.1f}s..."
|
||||
)
|
||||
await asyncio.sleep(delay)
|
||||
else:
|
||||
logger.error(f"Database acquire failed after {max_retries + 1} attempts: {e}")
|
||||
raise last_exception
|
||||
else:
|
||||
# Legacy path: raw asyncpg.Pool
|
||||
pool = backend_or_pool
|
||||
start = time.time()
|
||||
|
||||
async def acquire():
|
||||
return await pool.acquire()
|
||||
async def acquire():
|
||||
return await pool.acquire()
|
||||
|
||||
conn = await retry_with_backoff(acquire, max_retries=max_retries)
|
||||
acquire_time = time.time() - start
|
||||
conn = await retry_with_backoff(acquire, max_retries=max_retries)
|
||||
acquire_time = time.time() - start
|
||||
|
||||
# Log slow connection acquisitions (indicates pool contention)
|
||||
if acquire_time > 0.05: # 50ms threshold
|
||||
pool_size = pool.get_size()
|
||||
pool_free = pool.get_idle_size()
|
||||
logger.warning(f"[DB POOL] Slow acquire: {acquire_time:.3f}s | size={pool_size}, idle={pool_free}")
|
||||
if acquire_time > 0.05:
|
||||
pool_size = pool.get_size()
|
||||
pool_free = pool.get_idle_size()
|
||||
logger.warning(f"[DB POOL] Slow acquire: {acquire_time:.3f}s | size={pool_size}, idle={pool_free}")
|
||||
|
||||
try:
|
||||
yield conn
|
||||
finally:
|
||||
await pool.release(conn)
|
||||
try:
|
||||
yield conn
|
||||
finally:
|
||||
await pool.release(conn)
|
||||
|
||||
@@ -516,6 +516,7 @@ class CohereEmbeddings(Embeddings):
|
||||
api_key: str,
|
||||
model: str = DEFAULT_EMBEDDINGS_COHERE_MODEL,
|
||||
base_url: str | None = None,
|
||||
output_dimensions: int | None = None,
|
||||
batch_size: int = 96,
|
||||
timeout: float = 60.0,
|
||||
input_type: str = "search_document",
|
||||
@@ -527,6 +528,7 @@ class CohereEmbeddings(Embeddings):
|
||||
api_key: Cohere API key
|
||||
model: Cohere embedding model name (default: embed-english-v3.0)
|
||||
base_url: Custom base URL for Cohere-compatible API (e.g., Azure-hosted endpoint)
|
||||
output_dimensions: Optional output embedding dimensions (for Matryoshka-capable models)
|
||||
batch_size: Maximum batch size for embedding requests (default: 96, Cohere's limit)
|
||||
timeout: Request timeout in seconds (default: 60.0)
|
||||
input_type: Input type for embeddings (default: search_document).
|
||||
@@ -535,6 +537,7 @@ class CohereEmbeddings(Embeddings):
|
||||
self.api_key = api_key
|
||||
self.model = model
|
||||
self.base_url = base_url
|
||||
self.output_dimensions = output_dimensions
|
||||
self.batch_size = batch_size
|
||||
self.timeout = timeout
|
||||
self.input_type = input_type
|
||||
@@ -570,8 +573,10 @@ class CohereEmbeddings(Embeddings):
|
||||
client_kwargs["base_url"] = self.base_url
|
||||
self._client = cohere.Client(**client_kwargs)
|
||||
|
||||
# Try to get dimension from known models, otherwise do a test embedding
|
||||
if self.model in self.MODEL_DIMENSIONS:
|
||||
# If output_dimensions is explicitly set, use that as the dimension
|
||||
if self.output_dimensions is not None:
|
||||
self._dimension = self.output_dimensions
|
||||
elif self.model in self.MODEL_DIMENSIONS:
|
||||
self._dimension = self.MODEL_DIMENSIONS[self.model]
|
||||
else:
|
||||
# Do a test embedding to detect dimension
|
||||
@@ -607,13 +612,23 @@ class CohereEmbeddings(Embeddings):
|
||||
for i in range(0, len(texts), self.batch_size):
|
||||
batch = texts[i : i + self.batch_size]
|
||||
|
||||
response = self._client.embed(
|
||||
texts=batch,
|
||||
model=self.model,
|
||||
input_type=self.input_type,
|
||||
)
|
||||
|
||||
all_embeddings.extend(response.embeddings)
|
||||
if self.output_dimensions is not None:
|
||||
# Use v2 API which supports output_dimension
|
||||
response = self._client.v2.embed(
|
||||
texts=batch,
|
||||
model=self.model,
|
||||
input_type=self.input_type,
|
||||
output_dimension=self.output_dimensions,
|
||||
embedding_types=["float"],
|
||||
)
|
||||
all_embeddings.extend(response.embeddings.float_)
|
||||
else:
|
||||
response = self._client.embed(
|
||||
texts=batch,
|
||||
model=self.model,
|
||||
input_type=self.input_type,
|
||||
)
|
||||
all_embeddings.extend(response.embeddings)
|
||||
|
||||
return all_embeddings
|
||||
|
||||
@@ -821,6 +836,8 @@ class LiteLLMSDKEmbeddings(Embeddings):
|
||||
embed_kwargs["api_base"] = self.api_base
|
||||
if self.output_dimensions is not None:
|
||||
embed_kwargs["dimensions"] = self.output_dimensions
|
||||
if self.model.startswith("openai/"):
|
||||
embed_kwargs["allowed_openai_params"] = ["dimensions"]
|
||||
|
||||
# Use async embedding method (standard in litellm)
|
||||
response = await self._litellm.aembedding(**embed_kwargs)
|
||||
@@ -871,6 +888,8 @@ class LiteLLMSDKEmbeddings(Embeddings):
|
||||
embed_kwargs["api_base"] = self.api_base
|
||||
if self.output_dimensions is not None:
|
||||
embed_kwargs["dimensions"] = self.output_dimensions
|
||||
if self.model.startswith("openai/"):
|
||||
embed_kwargs["allowed_openai_params"] = ["dimensions"]
|
||||
|
||||
# Use sync embedding (litellm doesn't have async in thread-safe way)
|
||||
response = self._litellm.embedding(**embed_kwargs)
|
||||
@@ -912,6 +931,7 @@ class GeminiEmbeddings(Embeddings):
|
||||
vertexai_service_account_key: str | None = None,
|
||||
output_dimensionality: int | None = None,
|
||||
batch_size: int = 100,
|
||||
force_ipv4: bool = False,
|
||||
):
|
||||
self.model = model
|
||||
self.api_key = api_key
|
||||
@@ -920,7 +940,9 @@ class GeminiEmbeddings(Embeddings):
|
||||
self.vertexai_service_account_key = vertexai_service_account_key
|
||||
self.output_dimensionality = output_dimensionality
|
||||
self.batch_size = batch_size
|
||||
self.force_ipv4 = force_ipv4
|
||||
self._client = None
|
||||
self._httpx_client = None
|
||||
self._dimension: int | None = None
|
||||
self._is_vertexai = vertexai_project_id is not None
|
||||
self._embed_config = None # EmbedContentConfig, built during initialize()
|
||||
@@ -946,7 +968,7 @@ class GeminiEmbeddings(Embeddings):
|
||||
if self._is_vertexai:
|
||||
self._init_vertexai(genai)
|
||||
else:
|
||||
self._init_gemini(genai)
|
||||
self._init_gemini(genai, genai_types)
|
||||
|
||||
# Build EmbedContentConfig if output_dimensionality is set
|
||||
if self.output_dimensionality is not None:
|
||||
@@ -968,12 +990,25 @@ class GeminiEmbeddings(Embeddings):
|
||||
f"Embeddings: google provider initialized (auth: {auth_mode}, model: {self.model}, dim: {self._dimension})"
|
||||
)
|
||||
|
||||
def _init_gemini(self, genai) -> None:
|
||||
def _init_gemini(self, genai, genai_types) -> 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)
|
||||
client_kwargs = {"api_key": self.api_key}
|
||||
if self.force_ipv4:
|
||||
import httpx
|
||||
|
||||
self._httpx_client = httpx.Client(
|
||||
timeout=10,
|
||||
transport=httpx.HTTPTransport(local_address="0.0.0.0"),
|
||||
)
|
||||
client_kwargs["http_options"] = genai_types.HttpOptions(
|
||||
timeout=10000,
|
||||
httpxClient=self._httpx_client,
|
||||
)
|
||||
|
||||
self._client = genai.Client(**client_kwargs)
|
||||
logger.info(f"Embeddings: initializing Gemini provider with model {self.model}")
|
||||
|
||||
def _init_vertexai(self, genai) -> None:
|
||||
@@ -1100,7 +1135,12 @@ 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:
|
||||
@@ -1112,6 +1152,7 @@ def create_embeddings_from_env() -> Embeddings:
|
||||
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
|
||||
@@ -1121,6 +1162,7 @@ def create_embeddings_from_env() -> Embeddings:
|
||||
api_key=api_key,
|
||||
model=config.embeddings_cohere_model,
|
||||
base_url=config.embeddings_cohere_base_url,
|
||||
output_dimensions=config.embeddings_cohere_output_dimensions,
|
||||
)
|
||||
elif provider == "litellm":
|
||||
return LiteLLMEmbeddings(
|
||||
@@ -1159,6 +1201,7 @@ def create_embeddings_from_env() -> Embeddings:
|
||||
vertexai_region=config.embeddings_vertexai_region,
|
||||
vertexai_service_account_key=config.embeddings_vertexai_service_account_key,
|
||||
output_dimensionality=config.embeddings_gemini_output_dimensionality,
|
||||
force_ipv4=config.embeddings_gemini_force_ipv4,
|
||||
)
|
||||
else:
|
||||
raise ValueError(
|
||||
|
||||
@@ -6,13 +6,13 @@ to disambiguate entities across memory units.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
from collections import defaultdict
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import UTC, datetime
|
||||
from difflib import SequenceMatcher
|
||||
|
||||
import asyncpg
|
||||
from typing import Any
|
||||
|
||||
from .db_utils import acquire_with_retry
|
||||
from .memory_engine import fq_table
|
||||
@@ -63,7 +63,7 @@ class EntityResolver:
|
||||
Resolves entities to canonical IDs with disambiguation.
|
||||
"""
|
||||
|
||||
def __init__(self, pool: asyncpg.Pool, entity_lookup: str = "full"):
|
||||
def __init__(self, pool: Any, entity_lookup: str = "full"):
|
||||
"""
|
||||
Initialize entity resolver.
|
||||
|
||||
@@ -76,6 +76,8 @@ class EntityResolver:
|
||||
self.pool = pool
|
||||
self.entity_lookup = entity_lookup
|
||||
self._pg_trgm_checked = False
|
||||
# Backend-specific operations — accessed via pool.ops (Django pattern).
|
||||
self._ops = pool.ops if pool is not None else None
|
||||
# Keyed by asyncio task id so concurrent retain batches never mix their
|
||||
# pending updates. flush_pending_stats() pops only the calling task's items.
|
||||
self._pending_stats: dict[int, list[_EntityStat]] = {}
|
||||
@@ -216,6 +218,11 @@ class EntityResolver:
|
||||
taxonomy_lookup: set[str] | None = None,
|
||||
) -> list[str]:
|
||||
if self.entity_lookup == "trigram":
|
||||
# Route to backend-specific fuzzy strategy.
|
||||
# Non-PG backends (Oracle) use UTL_MATCH instead of pg_trgm.
|
||||
backend_strategy = self._ops.get_entity_resolution_strategy()
|
||||
if backend_strategy == "oracle_fuzzy":
|
||||
return await self._resolve_entities_batch_oracle_fuzzy(conn, bank_id, entities_data, unit_event_date)
|
||||
# Auto-detect pg_trgm availability on first call and fall back to
|
||||
# "full" strategy if the extension is not installed. See #626.
|
||||
if not self._pg_trgm_checked:
|
||||
@@ -384,6 +391,94 @@ class EntityResolver:
|
||||
conn, bank_id, entities_data, unit_event_date, all_candidates, cooccurrence_map
|
||||
)
|
||||
|
||||
async def _resolve_entities_batch_oracle_fuzzy(
|
||||
self, conn: Any, bank_id: str, entities_data: list[dict], unit_event_date: datetime | None
|
||||
) -> list[str]:
|
||||
"""
|
||||
Oracle strategy: fetch similar candidates using UTL_MATCH.JARO_WINKLER_SIMILARITY.
|
||||
|
||||
Replaces pg_trgm for Oracle backends. Uses JSON_TABLE to expand the
|
||||
entity text list into rows (Oracle equivalent of PG's unnest), then
|
||||
joins with a Jaro-Winkler threshold of 70/100 (≈ pg_trgm 0.15).
|
||||
Falls back to the "full" strategy if UTL_MATCH is unavailable.
|
||||
"""
|
||||
entity_texts = list(set(e["text"] for e in entities_data))
|
||||
entities_table = fq_table("entities")
|
||||
|
||||
try:
|
||||
# Batch all entity texts into a single query using JSON_TABLE to
|
||||
# expand the list into rows. UTL_MATCH.JARO_WINKLER_SIMILARITY
|
||||
# returns 0-100; threshold 70 ≈ pg_trgm similarity 0.15.
|
||||
entity_texts_json = json.dumps(entity_texts)
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT e.id, e.canonical_name, e.metadata, e.last_seen, e.mention_count,
|
||||
q.query_text
|
||||
FROM JSON_TABLE($2, '$[*]' COLUMNS (query_text VARCHAR2(4000) PATH '$')) q
|
||||
JOIN {entities_table} e ON (
|
||||
e.bank_id = $1
|
||||
AND UTL_MATCH.JARO_WINKLER_SIMILARITY(LOWER(e.canonical_name), LOWER(q.query_text)) > 70
|
||||
)
|
||||
""",
|
||||
bank_id,
|
||||
entity_texts_json,
|
||||
)
|
||||
except Exception as e:
|
||||
# UTL_MATCH may not be available (ORA-06550, ORA-00904, etc.)
|
||||
# Catch broadly because Oracle error types vary depending on driver.
|
||||
# Fall back to the "full" strategy which works on any backend.
|
||||
logger.warning(
|
||||
"UTL_MATCH.JARO_WINKLER_SIMILARITY not available on Oracle — "
|
||||
"falling back to 'full' entity lookup strategy. Error: %s",
|
||||
e,
|
||||
)
|
||||
self.entity_lookup = "full"
|
||||
return await self._resolve_entities_batch_full(conn, bank_id, entities_data, unit_event_date)
|
||||
|
||||
# Group candidates by query_text (same structure as trigram strategy)
|
||||
all_candidates: dict[str, list] = {t: [] for t in entity_texts}
|
||||
candidate_ids: set = set()
|
||||
for row in rows:
|
||||
query_text = row["query_text"]
|
||||
all_candidates[query_text].append(
|
||||
(row["id"], row["canonical_name"], row["metadata"], row["last_seen"], row["mention_count"])
|
||||
)
|
||||
candidate_ids.add(row["id"])
|
||||
|
||||
# Fetch co-occurrences only for the candidate entities (not all bank entities)
|
||||
cooccurrence_map: dict[str, set[str]] = {}
|
||||
if candidate_ids:
|
||||
candidate_id_list = list(candidate_ids)
|
||||
cooc_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT ec.entity_id_1, ec.entity_id_2
|
||||
FROM {fq_table("entity_cooccurrences")} ec
|
||||
WHERE ec.entity_id_1 = ANY($1::uuid[])
|
||||
OR ec.entity_id_2 = ANY($1::uuid[])
|
||||
""",
|
||||
candidate_id_list,
|
||||
)
|
||||
# Build name lookup for co-occurrence mapping
|
||||
id_to_name = {
|
||||
row["id"]: row["canonical_name"].lower()
|
||||
for cands in all_candidates.values()
|
||||
for row in [{"id": c[0], "canonical_name": c[1]} for c in cands]
|
||||
}
|
||||
for row in cooc_rows:
|
||||
eid1, eid2 = row["entity_id_1"], row["entity_id_2"]
|
||||
if eid1 not in cooccurrence_map:
|
||||
cooccurrence_map[eid1] = set()
|
||||
if eid2 not in cooccurrence_map:
|
||||
cooccurrence_map[eid2] = set()
|
||||
if eid2 in id_to_name:
|
||||
cooccurrence_map[eid1].add(id_to_name[eid2])
|
||||
if eid1 in id_to_name:
|
||||
cooccurrence_map[eid2].add(id_to_name[eid1])
|
||||
|
||||
return await self._resolve_from_candidates(
|
||||
conn, bank_id, entities_data, unit_event_date, all_candidates, cooccurrence_map
|
||||
)
|
||||
|
||||
async def _resolve_from_candidates(
|
||||
self,
|
||||
conn,
|
||||
@@ -491,24 +586,19 @@ class EntityResolver:
|
||||
# INSERT ... ON CONFLICT DO NOTHING — no row lock on already-existing entities.
|
||||
# mention_count starts at 0 here; flush_pending_stats() is the sole source of
|
||||
# truth for mention counting (one stat per original mention in the batch).
|
||||
inserted_rows = await conn.fetch(
|
||||
f"""
|
||||
INSERT INTO {fq_table("entities")} (bank_id, canonical_name, first_seen, last_seen, mention_count)
|
||||
SELECT $1, name, COALESCE(event_date, now()), COALESCE(event_date, now()), 0
|
||||
FROM unnest($2::text[], $3::timestamptz[]) AS t(name, event_date)
|
||||
ON CONFLICT (bank_id, LOWER(canonical_name))
|
||||
DO NOTHING
|
||||
RETURNING id, LOWER(canonical_name) AS name_lower
|
||||
""",
|
||||
entities_table = fq_table("entities")
|
||||
|
||||
id_by_name = await self._ops.bulk_insert_entities(
|
||||
conn,
|
||||
entities_table,
|
||||
bank_id,
|
||||
entity_names,
|
||||
entity_dates,
|
||||
)
|
||||
id_by_name: dict[str, str] = {row["name_lower"]: row["id"] for row in inserted_rows}
|
||||
|
||||
# Fallback SELECT for names that conflicted (another worker won the race).
|
||||
#
|
||||
# IMPORTANT: we must let PostgreSQL do the lowercasing on BOTH sides of the
|
||||
# IMPORTANT: we must let the database do the lowercasing on BOTH sides of the
|
||||
# comparison. Python's str.lower() and PostgreSQL's LOWER() differ for some
|
||||
# Unicode characters — most notably Turkish İ (U+0130):
|
||||
# Python: 'İstanbul'.lower() == 'i\u0307stanbul' (i + combining dot, 2 chars)
|
||||
@@ -516,24 +606,11 @@ class EntityResolver:
|
||||
# Passing a Python-lowercased name to "LOWER(canonical_name) = ANY($2::text[])"
|
||||
# would fail to match the stored entity, leaving entity_id as None and causing
|
||||
# a NOT NULL constraint violation on unit_entities.entity_id.
|
||||
#
|
||||
# Fix: pass the original (mixed-case) input names and use
|
||||
# "LOWER(canonical_name) = ANY(SELECT LOWER(n) FROM unnest($2) AS n)" so
|
||||
# PostgreSQL lowercases both sides identically. The query also returns the
|
||||
# original input_name so we can index id_by_name by Python's lower() of that
|
||||
# name, which is what the assignment loop below uses as its lookup key.
|
||||
missing_original = [g.name for name_lower, g in sorted_groups if name_lower not in id_by_name]
|
||||
if missing_original:
|
||||
existing_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT e.id, LOWER(e.canonical_name) AS name_lower, inputs.input_name
|
||||
FROM {fq_table("entities")} e
|
||||
JOIN (
|
||||
SELECT LOWER(n) AS input_name_lower, n AS input_name
|
||||
FROM unnest($2::text[]) AS n
|
||||
) AS inputs ON LOWER(e.canonical_name) = inputs.input_name_lower
|
||||
WHERE e.bank_id = $1
|
||||
""",
|
||||
existing_rows = await self._ops.fetch_missing_entity_ids(
|
||||
conn,
|
||||
entities_table,
|
||||
bank_id,
|
||||
missing_original,
|
||||
)
|
||||
@@ -541,8 +618,9 @@ class EntityResolver:
|
||||
id_by_name[row["name_lower"]] = row["id"]
|
||||
# Also index by Python's lower() of the original input name so the
|
||||
# assignment loop (which uses Python-lowercased keys) finds it even
|
||||
# when Python and PostgreSQL produce different lowercase strings.
|
||||
id_by_name[row["input_name"].lower()] = row["id"]
|
||||
# when Python and the database produce different lowercase strings.
|
||||
if "input_name" in row:
|
||||
id_by_name[row["input_name"].lower()] = row["id"]
|
||||
|
||||
# Assign entity IDs back and queue one stat per original mention so that
|
||||
# flush_pending_stats() increments mention_count by the true mention count,
|
||||
@@ -655,7 +733,11 @@ class EntityResolver:
|
||||
|
||||
# 3. Temporal proximity (0-0.2)
|
||||
if last_seen:
|
||||
days_diff = abs((unit_event_date - last_seen).total_seconds() / 86400)
|
||||
# Normalize both to UTC-aware to avoid naive/aware mismatch
|
||||
# (Oracle returns naive datetimes from fromisoformat)
|
||||
_evt = unit_event_date if unit_event_date.tzinfo else unit_event_date.replace(tzinfo=UTC)
|
||||
_seen = last_seen if last_seen.tzinfo else last_seen.replace(tzinfo=UTC)
|
||||
days_diff = abs((_evt - _seen).total_seconds() / 86400)
|
||||
if days_diff < 7: # Within a week
|
||||
temporal_score = max(0, 1.0 - (days_diff / 7))
|
||||
score += temporal_score * 0.2
|
||||
@@ -815,12 +897,10 @@ class EntityResolver:
|
||||
sorted_pairs = sorted(unit_entity_pairs)
|
||||
unit_ids = [p[0] for p in sorted_pairs]
|
||||
entity_ids = [p[1] for p in sorted_pairs]
|
||||
await conn.execute(
|
||||
f"""
|
||||
INSERT INTO {fq_table("unit_entities")} (unit_id, entity_id)
|
||||
SELECT u, e FROM unnest($1::uuid[], $2::uuid[]) AS t(u, e)
|
||||
ON CONFLICT DO NOTHING
|
||||
""",
|
||||
|
||||
await self._ops.bulk_insert_unit_entities(
|
||||
conn,
|
||||
fq_table("unit_entities"),
|
||||
unit_ids,
|
||||
entity_ids,
|
||||
)
|
||||
|
||||
@@ -161,16 +161,22 @@ class MemoryEngineInterface(ABC):
|
||||
bank_id: str,
|
||||
*,
|
||||
request_context: "RequestContext",
|
||||
) -> dict[str, Any]:
|
||||
create_if_missing: bool = True,
|
||||
) -> dict[str, Any] | None:
|
||||
"""
|
||||
Get bank profile including disposition and mission.
|
||||
|
||||
Args:
|
||||
bank_id: The memory bank ID.
|
||||
request_context: Request context for authentication.
|
||||
create_if_missing: If True (default), the bank is auto-created
|
||||
with defaults if it does not exist. Pass False to make this
|
||||
a strict read — returns None if the bank does not exist.
|
||||
|
||||
Returns:
|
||||
Bank profile dict with bank_id, name, disposition, and mission.
|
||||
Bank profile dict with bank_id, name, disposition, and mission,
|
||||
or None when create_if_missing=False and the bank does not
|
||||
exist.
|
||||
"""
|
||||
...
|
||||
|
||||
@@ -289,25 +295,6 @@ class MemoryEngineInterface(ABC):
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def delete_memory_unit(
|
||||
self,
|
||||
unit_id: str,
|
||||
*,
|
||||
request_context: "RequestContext",
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Delete a specific memory unit.
|
||||
|
||||
Args:
|
||||
unit_id: The memory unit ID.
|
||||
request_context: Request context for authentication.
|
||||
|
||||
Returns:
|
||||
Deletion result.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def get_graph_data(
|
||||
self,
|
||||
|
||||
@@ -283,7 +283,7 @@ def create_llm_provider(
|
||||
extra_args=config.llamacpp_extra_args,
|
||||
)
|
||||
|
||||
elif provider_lower in ("openai", "groq", "ollama", "lmstudio", "minimax", "volcano", "openrouter"):
|
||||
elif provider_lower in ("openai", "groq", "ollama", "lmstudio", "minimax", "deepseek", "volcano", "openrouter"):
|
||||
return OpenAICompatibleLLM(
|
||||
provider=provider,
|
||||
api_key=api_key,
|
||||
@@ -360,6 +360,7 @@ class LLMProvider:
|
||||
"mock",
|
||||
"none",
|
||||
"minimax",
|
||||
"deepseek",
|
||||
"litellm",
|
||||
"bedrock",
|
||||
"volcano",
|
||||
@@ -378,6 +379,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 == "deepseek":
|
||||
self.base_url = "https://api.deepseek.com"
|
||||
elif self.provider == "openrouter":
|
||||
self.base_url = "https://openrouter.ai/api/v1"
|
||||
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -5,12 +5,14 @@ from dataclasses import dataclass
|
||||
|
||||
from .base import FileParser, UnsupportedFileTypeError
|
||||
from .iris import IrisParser
|
||||
from .llama_parse import LlamaParseParser
|
||||
from .markitdown import MarkitdownParser
|
||||
|
||||
__all__ = [
|
||||
"FileParser",
|
||||
"UnsupportedFileTypeError",
|
||||
"IrisParser",
|
||||
"LlamaParseParser",
|
||||
"MarkitdownParser",
|
||||
"FileParserRegistry",
|
||||
"ConvertResult",
|
||||
|
||||
@@ -0,0 +1,125 @@
|
||||
"""LlamaParse parser implementation using the LlamaIndex Cloud parsing API."""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import mimetypes
|
||||
import time
|
||||
|
||||
import httpx
|
||||
|
||||
from .base import FileParser, UnsupportedFileTypeError
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_LLAMA_PARSE_BASE_URL = "https://api.cloud.llamaindex.ai/api/parsing"
|
||||
_DEFAULT_POLL_INTERVAL = 2.0 # seconds
|
||||
_DEFAULT_TIMEOUT = 300.0 # seconds
|
||||
|
||||
# HTTP status codes that indicate the file type is not supported.
|
||||
# Other 4xx codes (401, 403, 429, etc.) are operational errors, not file-type issues.
|
||||
_UNSUPPORTED_FILE_STATUS_CODES = {400, 415, 422}
|
||||
|
||||
|
||||
class LlamaParseParser(FileParser):
|
||||
"""
|
||||
LlamaParse file parser using LlamaIndex's hosted parsing service.
|
||||
|
||||
Uploads files to the LlamaParse API, polls until the parse job completes,
|
||||
and returns the resulting markdown. The API determines which file types
|
||||
are supported — UnsupportedFileTypeError is raised if the file is rejected.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
api_key: str,
|
||||
poll_interval: float = _DEFAULT_POLL_INTERVAL,
|
||||
timeout: float = _DEFAULT_TIMEOUT,
|
||||
):
|
||||
"""
|
||||
Initialize llama_parse parser.
|
||||
|
||||
Args:
|
||||
api_key: LlamaCloud API key (typically starts with "llx-")
|
||||
poll_interval: Seconds between status poll requests (default: 2)
|
||||
timeout: Maximum seconds to wait for parsing (default: 300)
|
||||
"""
|
||||
self._api_key = api_key
|
||||
self._poll_interval = poll_interval
|
||||
self._timeout = timeout
|
||||
self._auth_headers = {"Authorization": f"Bearer {api_key}"}
|
||||
self._client = httpx.AsyncClient(timeout=httpx.Timeout(30.0, read=120.0))
|
||||
|
||||
async def convert(self, file_data: bytes, filename: str) -> str:
|
||||
"""
|
||||
Parse file to markdown using the LlamaParse API.
|
||||
|
||||
Raises:
|
||||
UnsupportedFileTypeError: If the LlamaParse API rejects the file type
|
||||
RuntimeError: If parsing fails for another reason
|
||||
"""
|
||||
content_type = mimetypes.guess_type(filename)[0] or "application/octet-stream"
|
||||
|
||||
# Step 1: Upload file and start parse job
|
||||
upload_resp = await self._client.post(
|
||||
f"{_LLAMA_PARSE_BASE_URL}/upload",
|
||||
headers=self._auth_headers,
|
||||
# Ensure file_data is plain bytes (storage backends may return obstore.Bytes)
|
||||
files={"file": (filename, bytes(file_data), content_type)},
|
||||
)
|
||||
_raise_for_status(upload_resp, filename, "upload")
|
||||
job_id: str = upload_resp.json()["id"]
|
||||
|
||||
# Step 2: Poll job status until SUCCESS or ERROR
|
||||
deadline = time.monotonic() + self._timeout
|
||||
while True:
|
||||
status_resp = await self._client.get(
|
||||
f"{_LLAMA_PARSE_BASE_URL}/job/{job_id}",
|
||||
headers=self._auth_headers,
|
||||
)
|
||||
_raise_for_status(status_resp, filename, "poll job status")
|
||||
status_data = status_resp.json()
|
||||
status = status_data.get("status")
|
||||
|
||||
if status == "SUCCESS":
|
||||
break
|
||||
if status in ("ERROR", "CANCELLED"):
|
||||
error = status_data.get("error_code") or status_data.get("error") or "unknown error"
|
||||
raise RuntimeError(f"LlamaParse job failed for '{filename}': {error}")
|
||||
|
||||
if time.monotonic() >= deadline:
|
||||
raise RuntimeError(f"LlamaParse job timed out after {self._timeout}s for '{filename}'")
|
||||
|
||||
await asyncio.sleep(self._poll_interval)
|
||||
|
||||
# Step 3: Fetch the markdown result
|
||||
result_resp = await self._client.get(
|
||||
f"{_LLAMA_PARSE_BASE_URL}/job/{job_id}/result/markdown",
|
||||
headers=self._auth_headers,
|
||||
)
|
||||
_raise_for_status(result_resp, filename, "fetch markdown result")
|
||||
markdown = result_resp.json().get("markdown")
|
||||
if not markdown:
|
||||
raise RuntimeError(f"No content extracted from '{filename}'")
|
||||
return markdown
|
||||
|
||||
def name(self) -> str:
|
||||
"""Get parser name."""
|
||||
return "llama_parse"
|
||||
|
||||
|
||||
def _raise_for_status(response: httpx.Response, filename: str, step: str) -> None:
|
||||
"""
|
||||
Raise an appropriate error for HTTP errors.
|
||||
|
||||
Raises UnsupportedFileTypeError for 400/415/422 (file rejected by the API).
|
||||
Raises RuntimeError for all other errors (auth, rate-limit, server errors).
|
||||
"""
|
||||
if not response.is_error:
|
||||
return
|
||||
body = response.text or "<empty>"
|
||||
msg = (
|
||||
f"LlamaParse API error during {step} for '{filename}': {response.status_code} {response.reason_phrase} — {body}"
|
||||
)
|
||||
if response.status_code in _UNSUPPORTED_FILE_STATUS_CODES:
|
||||
raise UnsupportedFileTypeError(msg)
|
||||
raise RuntimeError(msg)
|
||||
@@ -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
|
||||
|
||||
@@ -175,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.
|
||||
@@ -212,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:
|
||||
@@ -227,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()
|
||||
@@ -401,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.
|
||||
@@ -493,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":
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
"""
|
||||
OpenAI-compatible LLM provider supporting OpenAI, Groq, Ollama, LMStudio, and MiniMax.
|
||||
OpenAI-compatible LLM provider supporting OpenAI, Groq, Ollama, LMStudio, MiniMax, and DeepSeek.
|
||||
|
||||
This provider handles all OpenAI API-compatible models including:
|
||||
- OpenAI: GPT-4, GPT-4o, GPT-5, o1, o3 (reasoning models)
|
||||
@@ -7,6 +7,7 @@ This provider handles all OpenAI API-compatible models including:
|
||||
- Ollama: Local models with native streaming API support
|
||||
- LMStudio: Local models with OpenAI-compatible API
|
||||
- MiniMax: MiniMax-M2.7 models with 1M context window
|
||||
- DeepSeek: deepseek-v4-flash / deepseek-v4-pro / deepseek-chat / deepseek-reasoner via api.deepseek.com
|
||||
|
||||
Features:
|
||||
- Reasoning models with extended thinking (o1, o3, GPT-5 families)
|
||||
@@ -60,6 +61,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.
|
||||
@@ -70,6 +97,7 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
- Ollama: Local models with native streaming API for better structured output
|
||||
- LMStudio: Local models with OpenAI-compatible API
|
||||
- MiniMax: MiniMax-M2.7 models via OpenAI-compatible API (https://api.minimax.io/v1)
|
||||
- DeepSeek: deepseek-v4-flash / deepseek-v4-pro / deepseek-chat / deepseek-reasoner via https://api.deepseek.com
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
@@ -101,7 +129,17 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
super().__init__(provider, api_key, base_url, model, reasoning_effort, **kwargs)
|
||||
|
||||
# Validate provider
|
||||
valid_providers = ["openai", "groq", "ollama", "lmstudio", "llamacpp", "minimax", "volcano", "openrouter"]
|
||||
valid_providers = [
|
||||
"openai",
|
||||
"groq",
|
||||
"ollama",
|
||||
"lmstudio",
|
||||
"llamacpp",
|
||||
"minimax",
|
||||
"deepseek",
|
||||
"volcano",
|
||||
"openrouter",
|
||||
]
|
||||
if self.provider not in valid_providers:
|
||||
raise ValueError(f"OpenAICompatibleLLM only supports: {', '.join(valid_providers)}. Got: {self.provider}")
|
||||
|
||||
@@ -115,6 +153,8 @@ 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 == "deepseek":
|
||||
self.base_url = "https://api.deepseek.com"
|
||||
elif self.provider == "openrouter":
|
||||
self.base_url = "https://openrouter.ai/api/v1"
|
||||
|
||||
@@ -123,7 +163,7 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
self.api_key = "local"
|
||||
|
||||
# Validate API key for cloud providers
|
||||
if self.provider in ("openai", "groq", "minimax", "openrouter") and not self.api_key:
|
||||
if self.provider in ("openai", "groq", "minimax", "deepseek", "openrouter") and not self.api_key:
|
||||
raise ValueError(f"API key is required for {self.provider}")
|
||||
|
||||
# Service tier configuration (from config, not env vars)
|
||||
@@ -180,7 +220,12 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
def _supports_reasoning_model(self) -> bool:
|
||||
"""Check if the current model is a reasoning model (o1, o3, GPT-5, DeepSeek)."""
|
||||
model_lower = self.model.lower()
|
||||
return any(x in model_lower for x in ["gpt-5", "o1", "o3", "deepseek"])
|
||||
if "deepseek" in model_lower:
|
||||
# DeepSeek v4-flash is the non-thinking route. Treating every
|
||||
# DeepSeek model as a reasoning model injects reasoning_effort,
|
||||
# which conflicts with thinking-disabled flash calls.
|
||||
return any(x in model_lower for x in ["v4-pro", "reasoner", "r1", "thinking"])
|
||||
return any(x in model_lower for x in ["gpt-5", "o1", "o3"])
|
||||
|
||||
def _get_max_reasoning_tokens(self) -> int | None:
|
||||
"""Get max reasoning tokens for reasoning models."""
|
||||
@@ -339,9 +384,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]
|
||||
@@ -550,12 +593,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:
|
||||
@@ -596,26 +646,59 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
"""
|
||||
start_time = time.time()
|
||||
|
||||
request_tool_choice: str | dict[str, Any] | None = tool_choice
|
||||
|
||||
# Normalize named tool_choice dicts to "required" + filter tools.
|
||||
# Some providers (e.g. LM Studio, Ollama) reject the OpenAI named format
|
||||
# {"type": "function", "function": {"name": "..."}}. The semantics are
|
||||
# identical to tool_choice="required" with the tools list restricted to
|
||||
# just the requested tool, so we apply that transformation universally.
|
||||
if isinstance(tool_choice, dict) and tool_choice.get("type") == "function":
|
||||
forced_name = tool_choice.get("function", {}).get("name")
|
||||
# just the requested tool, so we apply that transformation where supported.
|
||||
if isinstance(request_tool_choice, dict) and request_tool_choice.get("type") == "function":
|
||||
forced_name = request_tool_choice.get("function", {}).get("name")
|
||||
if forced_name:
|
||||
filtered = [t for t in tools if t.get("function", {}).get("name") == forced_name]
|
||||
if filtered:
|
||||
tools = filtered
|
||||
tool_choice = "required"
|
||||
request_tool_choice = "required"
|
||||
|
||||
# DeepSeek accepts tool calls but rejects explicit required/named
|
||||
# tool_choice values. The tools list has already been narrowed for
|
||||
# forced calls, so omitting tool_choice preserves the practical behavior.
|
||||
if "deepseek" in self.model.lower() and request_tool_choice != "auto":
|
||||
request_tool_choice = None
|
||||
|
||||
# "auto" is the OpenAI API default — omitting tool_choice is semantically
|
||||
# identical. Some providers (e.g. DeepSeek's reasoner pathway, which
|
||||
# deepseek-v4-flash falls into when thinking mode is enabled) reject the
|
||||
# parameter outright, returning HTTP 400 even for value "auto". Sending it
|
||||
# only when the caller asks for a non-default behaviour avoids those 400s
|
||||
# without changing semantics for compliant providers.
|
||||
if request_tool_choice == "auto":
|
||||
request_tool_choice = None
|
||||
|
||||
# DeepSeek tool-call replies can carry provider-specific reasoning_content.
|
||||
# The normalized tool result does not retain it, but replaying assistant
|
||||
# tool_calls without the field can trigger a 400. DeepSeek accepts an
|
||||
# empty-string fallback, matching the provider's history-replay contract.
|
||||
if "deepseek" in self.model.lower():
|
||||
normalized_messages: list[dict[str, Any]] = []
|
||||
for msg in messages:
|
||||
if msg.get("role") == "assistant" and msg.get("tool_calls") and "reasoning_content" not in msg:
|
||||
normalized_msg = dict(msg)
|
||||
normalized_msg["reasoning_content"] = ""
|
||||
normalized_messages.append(normalized_msg)
|
||||
else:
|
||||
normalized_messages.append(msg)
|
||||
messages = normalized_messages
|
||||
|
||||
# Build call parameters
|
||||
call_params: dict[str, Any] = {
|
||||
"model": self.model,
|
||||
"messages": messages,
|
||||
"tools": tools,
|
||||
"tool_choice": tool_choice,
|
||||
}
|
||||
if request_tool_choice is not None:
|
||||
call_params["tool_choice"] = request_tool_choice
|
||||
|
||||
if max_completion_tokens is not None:
|
||||
call_params[self._max_tokens_param_name()] = max_completion_tokens
|
||||
@@ -706,18 +789,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:
|
||||
@@ -765,6 +871,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
|
||||
@@ -919,7 +1026,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"))
|
||||
|
||||
@@ -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)
|
||||
@@ -23,6 +23,7 @@ logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
async def tool_search_mental_models(
|
||||
memory_engine: "MemoryEngine",
|
||||
conn: "Connection",
|
||||
bank_id: str,
|
||||
query: str,
|
||||
@@ -32,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.
|
||||
@@ -82,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}
|
||||
@@ -99,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(
|
||||
{
|
||||
@@ -136,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.
|
||||
@@ -179,6 +180,8 @@ async def tool_search_observations(
|
||||
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,
|
||||
@@ -214,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.
|
||||
@@ -230,15 +236,15 @@ 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,
|
||||
@@ -250,6 +256,8 @@ async def tool_recall(
|
||||
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,
|
||||
|
||||
@@ -46,7 +46,7 @@ def _vector_index_clause() -> str:
|
||||
return "USING hnsw (embedding vector_cosine_ops)"
|
||||
|
||||
|
||||
async def create_bank_vector_indexes(conn, bank_id: str, internal_id: str) -> None:
|
||||
async def create_bank_vector_indexes(conn, bank_id: str, internal_id: str, ops=None) -> None:
|
||||
"""Create per-(bank, fact_type) partial vector indexes for a newly created bank.
|
||||
|
||||
Respects the HINDSIGHT_API_VECTOR_EXTENSION config to use the appropriate
|
||||
@@ -55,29 +55,35 @@ async def create_bank_vector_indexes(conn, bank_id: str, internal_id: str) -> No
|
||||
Called immediately after the bank row is first inserted. Safe on empty banks
|
||||
(index build is instant). Idempotent via CREATE INDEX IF NOT EXISTS.
|
||||
bank_id is escaped for SQL literal safety (apostrophes doubled).
|
||||
|
||||
On Oracle 23ai, this is a no-op — Oracle uses a single global vector index
|
||||
created during migrations. Partial indexes (WHERE clause) are not supported
|
||||
for Oracle vector indexes.
|
||||
"""
|
||||
table = fq_table("memory_units")
|
||||
escaped = bank_id.replace("'", "''")
|
||||
using_clause = _vector_index_clause()
|
||||
for ft in _BANK_INDEX_FACT_TYPES:
|
||||
idx = _bank_index_name(ft, internal_id)
|
||||
await conn.execute(
|
||||
f"CREATE INDEX IF NOT EXISTS {idx} "
|
||||
f"ON {table} {using_clause} "
|
||||
f"WHERE fact_type = '{ft}' AND bank_id = '{escaped}'"
|
||||
)
|
||||
await ops.create_bank_vector_indexes(
|
||||
conn,
|
||||
fq_table("memory_units"),
|
||||
bank_id,
|
||||
internal_id,
|
||||
_vector_index_clause(),
|
||||
_BANK_INDEX_FACT_TYPES,
|
||||
)
|
||||
|
||||
|
||||
async def drop_bank_vector_indexes(conn, internal_id: str) -> None:
|
||||
async def drop_bank_vector_indexes(conn, internal_id: str, ops=None) -> None:
|
||||
"""Drop per-(bank, fact_type) partial vector indexes for a bank being deleted.
|
||||
|
||||
Called before the bank row is deleted so internal_id is still known.
|
||||
Idempotent via DROP INDEX IF EXISTS.
|
||||
|
||||
On Oracle, this is a no-op (uses single global vector index).
|
||||
"""
|
||||
schema = get_current_schema()
|
||||
for ft in _BANK_INDEX_FACT_TYPES:
|
||||
idx = _bank_index_name(ft, internal_id)
|
||||
await conn.execute(f"DROP INDEX IF EXISTS {schema}.{idx}")
|
||||
await ops.drop_bank_vector_indexes(
|
||||
conn,
|
||||
get_current_schema(),
|
||||
internal_id,
|
||||
_BANK_INDEX_FACT_TYPES,
|
||||
)
|
||||
|
||||
|
||||
DEFAULT_DISPOSITION = {
|
||||
@@ -117,6 +123,41 @@ async def get_bank_profile(pool, bank_id: str) -> BankProfile:
|
||||
return profile
|
||||
|
||||
|
||||
async def get_bank_profile_if_exists(pool, bank_id: str) -> BankProfile | None:
|
||||
"""
|
||||
Get bank profile (name, disposition + mission) without auto-creating.
|
||||
|
||||
Returns None if the bank does not exist. This is the read-only variant
|
||||
of get_bank_profile, intended for read endpoints where a bank that
|
||||
doesn't exist should surface as 404 rather than be silently created.
|
||||
|
||||
Args:
|
||||
pool: Database connection pool
|
||||
bank_id: bank IDentifier
|
||||
|
||||
Returns:
|
||||
BankProfile if the bank exists, otherwise None.
|
||||
"""
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
row = await conn.fetchrow(
|
||||
f"""
|
||||
SELECT name, disposition, mission
|
||||
FROM {fq_table("banks")} WHERE bank_id = $1
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
if not row:
|
||||
return None
|
||||
disposition_data = row["disposition"]
|
||||
if isinstance(disposition_data, str):
|
||||
disposition_data = json.loads(disposition_data)
|
||||
return BankProfile(
|
||||
name=row["name"],
|
||||
disposition=DispositionTraits(**disposition_data),
|
||||
mission=row["mission"] or "",
|
||||
)
|
||||
|
||||
|
||||
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.
|
||||
@@ -175,7 +216,7 @@ async def get_or_create_bank_profile(pool, bank_id: str) -> tuple[BankProfile, b
|
||||
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))
|
||||
await create_bank_vector_indexes(conn, bank_id, str(internal_id), ops=pool.ops)
|
||||
|
||||
return (
|
||||
BankProfile(name=bank_id, disposition=DispositionTraits(**DEFAULT_DISPOSITION), mission=""),
|
||||
|
||||
@@ -69,7 +69,9 @@ async def delete_chunks_by_ids(conn, chunk_ids: list[str]) -> None:
|
||||
)
|
||||
|
||||
|
||||
async def store_chunks_batch(conn, bank_id: str, document_id: str, chunks: list[ChunkMetadata]) -> dict[int, str]:
|
||||
async def store_chunks_batch(
|
||||
conn, bank_id: str, document_id: str, chunks: list[ChunkMetadata], ops=None
|
||||
) -> dict[int, str]:
|
||||
"""
|
||||
Store document chunks in the database.
|
||||
|
||||
@@ -78,6 +80,7 @@ async def store_chunks_batch(conn, bank_id: str, document_id: str, chunks: list[
|
||||
bank_id: Bank identifier
|
||||
document_id: Document identifier
|
||||
chunks: List of ChunkMetadata objects
|
||||
ops: DataAccessOps instance (from backend.ops)
|
||||
|
||||
Returns:
|
||||
Dictionary mapping global chunk index to chunk_id
|
||||
@@ -101,20 +104,11 @@ async def store_chunks_batch(conn, bank_id: str, document_id: str, chunks: list[
|
||||
chunk_id_map[chunk.chunk_index] = chunk_id
|
||||
|
||||
# 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
|
||||
""",
|
||||
# a retain under the same document_id may produce chunk_ids that already exist.
|
||||
# Overwriting is the correct behavior per document_id grouping semantics.
|
||||
await ops.bulk_upsert_chunks(
|
||||
conn,
|
||||
fq_table("chunks"),
|
||||
chunk_ids,
|
||||
[document_id] * len(chunk_texts),
|
||||
[bank_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
|
||||
|
||||
@@ -111,6 +111,7 @@ async def build_entity_links(
|
||||
unit_to_entity_ids: dict[str, list[str]],
|
||||
log_buffer: list[str] = None,
|
||||
skip_unit_entities_insert: bool = False,
|
||||
ops=None,
|
||||
) -> list[EntityLink]:
|
||||
"""
|
||||
Build entity links for UI graph visualization.
|
||||
@@ -130,6 +131,7 @@ async def build_entity_links(
|
||||
unit_to_entity_ids: From resolve_entities()
|
||||
log_buffer: Optional buffer for detailed logging
|
||||
skip_unit_entities_insert: Skip unit_entities INSERT (already done in Phase 2)
|
||||
ops: DataAccessOps instance (from backend.ops)
|
||||
|
||||
Returns:
|
||||
List of EntityLink objects for batch insertion
|
||||
@@ -144,10 +146,11 @@ async def build_entity_links(
|
||||
unit_to_entity_ids,
|
||||
log_buffer,
|
||||
skip_unit_entities_insert=skip_unit_entities_insert,
|
||||
ops=ops,
|
||||
)
|
||||
|
||||
|
||||
async def insert_entity_links_batch(conn, entity_links: list[EntityLink], bank_id: str) -> None:
|
||||
async def insert_entity_links_batch(conn, entity_links: list[EntityLink], bank_id: str, ops=None) -> None:
|
||||
"""
|
||||
Insert entity links in batch.
|
||||
|
||||
@@ -155,8 +158,9 @@ async def insert_entity_links_batch(conn, entity_links: list[EntityLink], bank_i
|
||||
conn: Database connection
|
||||
entity_links: List of EntityLink objects
|
||||
bank_id: Bank identifier (stored directly on memory_links for fast filtering)
|
||||
ops: DataAccessOps instance (from backend.ops)
|
||||
"""
|
||||
if not entity_links:
|
||||
return
|
||||
|
||||
await link_utils.insert_entity_links_batch(conn, entity_links, bank_id)
|
||||
await link_utils.insert_entity_links_batch(conn, entity_links, bank_id, ops=ops)
|
||||
|
||||
@@ -110,12 +110,6 @@ class CausalRelation(BaseModel):
|
||||
relation_type: Literal["caused_by"] = Field(
|
||||
description="How this fact relates to the target: 'caused_by' = this fact was caused by the target"
|
||||
)
|
||||
strength: float = Field(
|
||||
description="Strength of relationship (0.0 to 1.0)",
|
||||
ge=0.0,
|
||||
le=1.0,
|
||||
default=1.0,
|
||||
)
|
||||
|
||||
|
||||
class FactCausalRelation(BaseModel):
|
||||
@@ -134,12 +128,6 @@ class FactCausalRelation(BaseModel):
|
||||
relation_type: Literal["caused_by"] = Field(
|
||||
description="How this fact relates to the target fact: 'caused_by' = this fact was caused by the target fact"
|
||||
)
|
||||
strength: float = Field(
|
||||
description="Strength of relationship (0.0 to 1.0). 1.0 = strong, 0.5 = moderate",
|
||||
ge=0.0,
|
||||
le=1.0,
|
||||
default=1.0,
|
||||
)
|
||||
|
||||
|
||||
class ExtractedFact(BaseModel):
|
||||
@@ -1216,7 +1204,6 @@ async def _extract_facts_from_chunk(
|
||||
# New schema uses target_index
|
||||
target_idx = rel.get("target_index")
|
||||
relation_type = rel.get("relation_type")
|
||||
strength = rel.get("strength", 1.0)
|
||||
|
||||
if target_idx is None or relation_type is None:
|
||||
continue
|
||||
@@ -1233,7 +1220,6 @@ async def _extract_facts_from_chunk(
|
||||
CausalRelation(
|
||||
target_fact_index=target_idx,
|
||||
relation_type=relation_type,
|
||||
strength=strength,
|
||||
)
|
||||
)
|
||||
except Exception as e:
|
||||
@@ -1602,13 +1588,15 @@ async def extract_facts_from_contents_batch_api(
|
||||
# Check if we're resuming an existing batch (crash recovery)
|
||||
batch_id = None
|
||||
if operation_id and pool:
|
||||
from ..db_utils import acquire_with_retry
|
||||
from ..task_backend import fq_table
|
||||
|
||||
table = fq_table("async_operations", schema)
|
||||
row = await pool.fetchrow(
|
||||
f"SELECT result_metadata FROM {table} WHERE operation_id = $1",
|
||||
operation_id,
|
||||
)
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
row = await conn.fetchrow(
|
||||
f"SELECT result_metadata FROM {table} WHERE operation_id = $1",
|
||||
operation_id,
|
||||
)
|
||||
|
||||
if row and row["result_metadata"]:
|
||||
metadata = row["result_metadata"]
|
||||
@@ -1675,18 +1663,20 @@ async def extract_facts_from_contents_batch_api(
|
||||
}
|
||||
|
||||
# Update operation result_metadata
|
||||
from ..db_utils import acquire_with_retry
|
||||
from ..task_backend import fq_table
|
||||
|
||||
table = fq_table("async_operations", schema)
|
||||
await pool.execute(
|
||||
f"""
|
||||
UPDATE {table}
|
||||
SET result_metadata = result_metadata || $1::jsonb, updated_at = now()
|
||||
WHERE operation_id = $2
|
||||
""",
|
||||
json.dumps(batch_state),
|
||||
operation_id,
|
||||
)
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
await conn.execute(
|
||||
f"""
|
||||
UPDATE {table}
|
||||
SET result_metadata = result_metadata || $1::jsonb, updated_at = now()
|
||||
WHERE operation_id = $2
|
||||
""",
|
||||
json.dumps(batch_state),
|
||||
operation_id,
|
||||
)
|
||||
logger.info(f"Stored batch state for operation {operation_id} (crash recovery enabled)")
|
||||
else:
|
||||
logger.info(f"Resuming polling for existing batch: {batch_id}")
|
||||
@@ -1909,7 +1899,6 @@ async def extract_facts_from_contents_batch_api(
|
||||
continue
|
||||
target_idx = rel.get("target_index")
|
||||
relation_type = rel.get("relation_type")
|
||||
strength = rel.get("strength", 1.0)
|
||||
|
||||
if target_idx is None or relation_type is None:
|
||||
continue
|
||||
@@ -1918,9 +1907,7 @@ async def extract_facts_from_contents_batch_api(
|
||||
|
||||
try:
|
||||
validated_relations.append(
|
||||
CausalRelation(
|
||||
target_fact_index=target_idx, relation_type=relation_type, strength=strength
|
||||
)
|
||||
CausalRelation(target_fact_index=target_idx, relation_type=relation_type)
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
@@ -2250,7 +2237,6 @@ def _convert_causal_relations(relations_from_llm, fact_start_idx: int) -> list[C
|
||||
causal_relation = CausalRelationType(
|
||||
relation_type=rel.relation_type,
|
||||
target_fact_index=fact_start_idx + rel.target_fact_index,
|
||||
strength=rel.strength,
|
||||
)
|
||||
causal_relations.append(causal_relation)
|
||||
return causal_relations
|
||||
|
||||
@@ -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
|
||||
@@ -35,7 +36,7 @@ async def get_document_content(
|
||||
|
||||
|
||||
async def insert_facts_batch(
|
||||
conn, bank_id: str, facts: list[ProcessedFact], document_id: str | None = None
|
||||
conn, bank_id: str, facts: list[ProcessedFact], document_id: str | None = None, ops=None
|
||||
) -> list[str]:
|
||||
"""
|
||||
Insert facts into the database in batch.
|
||||
@@ -106,77 +107,16 @@ async def insert_facts_batch(
|
||||
pass
|
||||
text_signals_list.append(" ".join(signal_parts) if signal_parts else None)
|
||||
|
||||
# Batch insert all facts
|
||||
# Note: tags are passed as JSON strings and converted back to varchar[] via jsonb_array_elements_text + array_agg
|
||||
# Query varies based on text search backend
|
||||
# Batch insert all facts — delegates to DataAccessOps which handles
|
||||
# unnest (PG) vs row-by-row (Oracle) transparently.
|
||||
config = get_config()
|
||||
if config.text_search_extension == "vchord":
|
||||
# VectorChord: manually tokenize and insert search_vector
|
||||
# text_signals (entity names etc.) are included in the tokenize input for enriched BM25
|
||||
query = f"""
|
||||
WITH input_data AS (
|
||||
SELECT * FROM unnest(
|
||||
$2::text[], $3::vector[], $4::timestamptz[], $5::timestamptz[], $6::timestamptz[], $7::timestamptz[],
|
||||
$8::text[], $9::text[], $10::jsonb[], $11::text[], $12::text[], $13::jsonb[], $14::jsonb[], $15::text[]
|
||||
) AS t(text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
context, fact_type, metadata, chunk_id, document_id, tags_json,
|
||||
observation_scopes_json, text_signals)
|
||||
)
|
||||
INSERT INTO {fq_table("memory_units")} (bank_id, text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
context, fact_type, metadata, chunk_id, document_id, tags,
|
||||
observation_scopes, text_signals, search_vector)
|
||||
SELECT
|
||||
$1,
|
||||
text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
context, fact_type, metadata, chunk_id, document_id,
|
||||
COALESCE(
|
||||
(SELECT array_agg(elem) FROM jsonb_array_elements_text(tags_json) AS elem),
|
||||
'{{}}'::varchar[]
|
||||
),
|
||||
observation_scopes_json,
|
||||
text_signals,
|
||||
tokenize(
|
||||
COALESCE(text, '') || ' ' || COALESCE(context, '') || ' ' || COALESCE(text_signals, ''),
|
||||
'llmlingua2'
|
||||
)::bm25_catalog.bm25vector
|
||||
FROM input_data
|
||||
RETURNING id
|
||||
"""
|
||||
else: # native or pg_textsearch
|
||||
# Native PostgreSQL: search_vector is GENERATED ALWAYS (expression includes text_signals), don't include it
|
||||
# pg_textsearch: indexes operate on base columns directly, don't populate search_vector
|
||||
query = f"""
|
||||
WITH input_data AS (
|
||||
SELECT * FROM unnest(
|
||||
$2::text[], $3::vector[], $4::timestamptz[], $5::timestamptz[], $6::timestamptz[], $7::timestamptz[],
|
||||
$8::text[], $9::text[], $10::jsonb[], $11::text[], $12::text[], $13::jsonb[], $14::jsonb[], $15::text[]
|
||||
) AS t(text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
context, fact_type, metadata, chunk_id, document_id, tags_json,
|
||||
observation_scopes_json, text_signals)
|
||||
)
|
||||
INSERT INTO {fq_table("memory_units")} (bank_id, text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
context, fact_type, metadata, chunk_id, document_id, tags,
|
||||
observation_scopes, text_signals)
|
||||
SELECT
|
||||
$1,
|
||||
text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
context, fact_type, metadata, chunk_id, document_id,
|
||||
COALESCE(
|
||||
(SELECT array_agg(elem) FROM jsonb_array_elements_text(tags_json) AS elem),
|
||||
'{{}}'::varchar[]
|
||||
),
|
||||
observation_scopes_json,
|
||||
text_signals
|
||||
FROM input_data
|
||||
RETURNING id
|
||||
"""
|
||||
|
||||
results = await conn.fetch(
|
||||
query,
|
||||
return await ops.insert_facts_batch(
|
||||
conn,
|
||||
bank_id,
|
||||
fact_texts,
|
||||
embeddings,
|
||||
event_dates, # event_date: occurred_start if available, else mentioned_at
|
||||
event_dates,
|
||||
occurred_starts,
|
||||
occurred_ends,
|
||||
mentioned_ats,
|
||||
@@ -188,13 +128,11 @@ async def insert_facts_batch(
|
||||
tags_list,
|
||||
observation_scopes_list,
|
||||
text_signals_list,
|
||||
text_search_extension=config.text_search_extension,
|
||||
)
|
||||
|
||||
unit_ids = [str(row["id"]) for row in results]
|
||||
return unit_ids
|
||||
|
||||
|
||||
async def ensure_bank_exists(conn, bank_id: str) -> None:
|
||||
async def ensure_bank_exists(conn, bank_id: str, ops=None) -> None:
|
||||
"""
|
||||
Ensure bank exists in the database.
|
||||
|
||||
@@ -221,7 +159,92 @@ async def ensure_bank_exists(conn, bank_id: str) -> None:
|
||||
)
|
||||
if inserted:
|
||||
# Fresh insert — create per-bank vector indexes
|
||||
await create_bank_vector_indexes(conn, bank_id, str(internal_id))
|
||||
await create_bank_vector_indexes(conn, bank_id, str(internal_id), ops=ops)
|
||||
|
||||
|
||||
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]
|
||||
|
||||
# Use observation_sources junction table instead of PG-specific array
|
||||
# overlap operator (&&). This is portable across all backends.
|
||||
affected_obs = await conn.fetch(
|
||||
f"""
|
||||
SELECT mu.id, mu.source_memory_ids
|
||||
FROM {fq_table("memory_units")} mu
|
||||
WHERE mu.bank_id = $1
|
||||
AND mu.fact_type = 'observation'
|
||||
AND EXISTS (
|
||||
SELECT 1 FROM {fq_table("observation_sources")} os
|
||||
WHERE os.observation_id = mu.id
|
||||
AND os.source_id = ANY($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(
|
||||
@@ -254,17 +277,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(
|
||||
@@ -297,12 +361,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,
|
||||
@@ -316,6 +387,7 @@ async def _upsert_document_row(
|
||||
content_hash,
|
||||
json.dumps(retain_params) if retain_params else None,
|
||||
document_tags or [],
|
||||
preserved_created_at,
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -12,7 +12,7 @@ from .types import ProcessedFact
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
async def create_temporal_links_batch(conn, bank_id: str, unit_ids: list[str]) -> int:
|
||||
async def create_temporal_links_batch(conn, bank_id: str, unit_ids: list[str], ops=None) -> int:
|
||||
"""
|
||||
Create temporal links between facts.
|
||||
|
||||
@@ -29,7 +29,7 @@ async def create_temporal_links_batch(conn, bank_id: str, unit_ids: list[str]) -
|
||||
if not unit_ids:
|
||||
return 0
|
||||
|
||||
return await link_utils.create_temporal_links_batch_per_fact(conn, bank_id, unit_ids, log_buffer=[])
|
||||
return await link_utils.create_temporal_links_batch_per_fact(conn, bank_id, unit_ids, log_buffer=[], ops=ops)
|
||||
|
||||
|
||||
async def create_semantic_links_batch(
|
||||
@@ -38,6 +38,7 @@ async def create_semantic_links_batch(
|
||||
unit_ids: list[str],
|
||||
embeddings: list[list[float]],
|
||||
pre_computed_ann_links: list[tuple] | None = None,
|
||||
ops=None,
|
||||
) -> int:
|
||||
"""
|
||||
Create semantic links between facts.
|
||||
@@ -63,11 +64,13 @@ async def create_semantic_links_batch(
|
||||
raise ValueError(f"Mismatch between unit_ids ({len(unit_ids)}) and embeddings ({len(embeddings)})")
|
||||
|
||||
return await link_utils.create_semantic_links_batch(
|
||||
conn, bank_id, unit_ids, embeddings, log_buffer=[], pre_computed_ann_links=pre_computed_ann_links
|
||||
conn, bank_id, unit_ids, embeddings, log_buffer=[], pre_computed_ann_links=pre_computed_ann_links, ops=ops
|
||||
)
|
||||
|
||||
|
||||
async def create_causal_links_batch(conn, bank_id: str, unit_ids: list[str], facts: list[ProcessedFact]) -> int:
|
||||
async def create_causal_links_batch(
|
||||
conn, bank_id: str, unit_ids: list[str], facts: list[ProcessedFact], ops=None
|
||||
) -> int:
|
||||
"""
|
||||
Create causal links between facts.
|
||||
|
||||
@@ -97,7 +100,6 @@ async def create_causal_links_batch(conn, bank_id: str, unit_ids: list[str], fac
|
||||
{
|
||||
"relation_type": rel.relation_type,
|
||||
"target_fact_index": rel.target_fact_index,
|
||||
"strength": rel.strength,
|
||||
}
|
||||
for rel in fact.causal_relations
|
||||
]
|
||||
@@ -105,6 +107,6 @@ async def create_causal_links_batch(conn, bank_id: str, unit_ids: list[str], fac
|
||||
else:
|
||||
causal_relations_per_fact.append([])
|
||||
|
||||
link_count = await link_utils.create_causal_links_batch(conn, bank_id, unit_ids, causal_relations_per_fact)
|
||||
link_count = await link_utils.create_causal_links_batch(conn, bank_id, unit_ids, causal_relations_per_fact, ops=ops)
|
||||
|
||||
return link_count
|
||||
|
||||
@@ -57,16 +57,13 @@ async def _bulk_insert_links(
|
||||
bank_id: str = "",
|
||||
chunk_size: int = 5000,
|
||||
skip_exists_check: bool = False,
|
||||
ops=None,
|
||||
) -> None:
|
||||
"""Bulk-insert links using sorted INSERT FROM unnest().
|
||||
|
||||
Sorting by (from_unit_id, to_unit_id) ensures all concurrent transactions
|
||||
acquire index locks in the same order, eliminating circular-wait deadlocks.
|
||||
|
||||
A single INSERT ... SELECT FROM unnest() is also faster than executemany
|
||||
(one round-trip vs N), and acquires all locks within one statement execution
|
||||
rather than interleaving with other transactions between rows.
|
||||
|
||||
Args:
|
||||
conn: Database connection (must be inside a transaction).
|
||||
links: List of (from_unit_id, to_unit_id, link_type, weight, entity_id) tuples.
|
||||
@@ -76,6 +73,7 @@ async def _bulk_insert_links(
|
||||
skip_exists_check: Skip WHERE EXISTS checks on memory_units. Use when
|
||||
all referenced unit IDs are guaranteed to exist (e.g., within
|
||||
the same transaction that inserted them).
|
||||
ops: DataAccessOps instance for backend-specific bulk operations.
|
||||
"""
|
||||
if not links:
|
||||
return
|
||||
@@ -84,12 +82,6 @@ async def _bulk_insert_links(
|
||||
# across concurrent transactions — prevents deadlocks.
|
||||
sorted_links = sorted(links, key=lambda lnk: (str(lnk[0]), str(lnk[1])))
|
||||
|
||||
from_ids = [lnk[0] for lnk in sorted_links]
|
||||
to_ids = [lnk[1] for lnk in sorted_links]
|
||||
types = [lnk[2] for lnk in sorted_links]
|
||||
weights = [lnk[3] for lnk in sorted_links]
|
||||
entity_ids = [lnk[4] for lnk in sorted_links]
|
||||
|
||||
exists_clause = ""
|
||||
if not skip_exists_check:
|
||||
exists_clause = (
|
||||
@@ -97,28 +89,15 @@ async def _bulk_insert_links(
|
||||
f" AND EXISTS (SELECT 1 FROM {fq_table('memory_units')} mu WHERE mu.id = t)"
|
||||
)
|
||||
|
||||
for chunk_start in range(0, len(sorted_links), chunk_size):
|
||||
chunk_end = min(chunk_start + chunk_size, len(sorted_links))
|
||||
await conn.execute(
|
||||
f"""
|
||||
INSERT INTO {fq_table("memory_links")}
|
||||
(from_unit_id, to_unit_id, link_type, weight, entity_id, bank_id)
|
||||
SELECT f, t, tp, w, e, $6
|
||||
FROM unnest($1::uuid[], $2::uuid[], $3::text[], $4::float8[], $5::uuid[])
|
||||
AS t(f, t, tp, w, e)
|
||||
{exists_clause}
|
||||
ON CONFLICT (from_unit_id, to_unit_id, link_type,
|
||||
COALESCE(entity_id, '{_NIL_ENTITY_UUID}'::uuid))
|
||||
DO NOTHING
|
||||
""",
|
||||
from_ids[chunk_start:chunk_end],
|
||||
to_ids[chunk_start:chunk_end],
|
||||
types[chunk_start:chunk_end],
|
||||
weights[chunk_start:chunk_end],
|
||||
entity_ids[chunk_start:chunk_end],
|
||||
bank_id,
|
||||
timeout=300,
|
||||
)
|
||||
await ops.bulk_insert_links(
|
||||
conn,
|
||||
fq_table("memory_links"),
|
||||
sorted_links,
|
||||
bank_id,
|
||||
_NIL_ENTITY_UUID,
|
||||
exists_clause,
|
||||
chunk_size,
|
||||
)
|
||||
|
||||
|
||||
def _normalize_datetime(dt):
|
||||
@@ -397,6 +376,7 @@ async def build_entity_links_from_resolved(
|
||||
unit_to_entity_ids: dict[str, list[str]],
|
||||
log_buffer: list[str] = None,
|
||||
skip_unit_entities_insert: bool = False,
|
||||
ops=None,
|
||||
) -> list["EntityLink"]:
|
||||
"""
|
||||
Build entity links between units that share entities.
|
||||
@@ -451,22 +431,13 @@ async def build_entity_links_from_resolved(
|
||||
import uuid
|
||||
|
||||
entity_id_list = [uuid.UUID(eid) if isinstance(eid, str) else eid for eid in all_entity_ids]
|
||||
# Use LATERAL with LIMIT to cap rows fetched per entity at the SQL level,
|
||||
# avoiding transfer of thousands of rows for high-cardinality entities.
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT e.entity_id, n.unit_id
|
||||
FROM unnest($1::uuid[]) AS e(entity_id)
|
||||
CROSS JOIN LATERAL (
|
||||
SELECT ue.unit_id
|
||||
FROM {fq_table("unit_entities")} ue
|
||||
WHERE ue.entity_id = e.entity_id
|
||||
ORDER BY ue.unit_id DESC
|
||||
LIMIT $2
|
||||
) n
|
||||
""",
|
||||
limit_per_entity = MAX_LINKS_PER_ENTITY + len(unit_ids) # room for new units + existing cap
|
||||
|
||||
rows = await ops.fetch_entity_unit_fanout(
|
||||
conn,
|
||||
fq_table("unit_entities"),
|
||||
entity_id_list,
|
||||
MAX_LINKS_PER_ENTITY + len(unit_ids), # room for new units + existing cap
|
||||
limit_per_entity,
|
||||
)
|
||||
_log(
|
||||
log_buffer,
|
||||
@@ -529,6 +500,7 @@ async def create_temporal_links_batch_per_fact(
|
||||
unit_ids: list[str],
|
||||
time_window_hours: int = 24,
|
||||
log_buffer: list[str] = None,
|
||||
ops=None,
|
||||
) -> int:
|
||||
"""
|
||||
Create temporal links for multiple units, each with their own event_date.
|
||||
@@ -554,14 +526,7 @@ async def create_temporal_links_batch_per_fact(
|
||||
|
||||
# Get the event_date for each new unit
|
||||
fetch_dates_start = time_mod.time()
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT id, event_date, fact_type
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE id::text = ANY($1)
|
||||
""",
|
||||
unit_ids,
|
||||
)
|
||||
rows = await ops.fetch_unit_dates(conn, fq_table("memory_units"), unit_ids)
|
||||
new_units = {str(row["id"]): (row["event_date"], row["fact_type"]) for row in rows}
|
||||
_log(
|
||||
log_buffer,
|
||||
@@ -590,52 +555,22 @@ async def create_temporal_links_batch_per_fact(
|
||||
TEMPORAL_LATERAL_BATCH = 500
|
||||
half_limit = MAX_TEMPORAL_LINKS_PER_UNIT # fetch K in each direction, take top K combined
|
||||
mu = fq_table("memory_units")
|
||||
rows = []
|
||||
for batch_start in range(0, len(new_unit_entries), TEMPORAL_LATERAL_BATCH):
|
||||
batch_end = batch_start + TEMPORAL_LATERAL_BATCH
|
||||
batch_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT from_id, id, event_date, time_diff_hours FROM (
|
||||
SELECT src.unit_id::text AS from_id, combined.*,
|
||||
ROW_NUMBER() OVER (
|
||||
PARTITION BY src.unit_id
|
||||
ORDER BY combined.time_diff_hours
|
||||
) AS rn
|
||||
FROM unnest($1::uuid[], $2::timestamptz[], $3::text[])
|
||||
AS src(unit_id, event_date, fact_type)
|
||||
CROSS JOIN LATERAL (
|
||||
-- Scan backward (older events) using index order
|
||||
(SELECT mu.id, mu.event_date,
|
||||
ABS(EXTRACT(EPOCH FROM mu.event_date - src.event_date)) / 3600.0 AS time_diff_hours
|
||||
FROM {mu} mu
|
||||
WHERE mu.bank_id = $4
|
||||
AND mu.fact_type = src.fact_type
|
||||
AND mu.event_date <= src.event_date
|
||||
AND mu.id != src.unit_id
|
||||
ORDER BY mu.event_date DESC
|
||||
LIMIT $5)
|
||||
UNION ALL
|
||||
-- Scan forward (newer events) using index order
|
||||
(SELECT mu.id, mu.event_date,
|
||||
ABS(EXTRACT(EPOCH FROM mu.event_date - src.event_date)) / 3600.0 AS time_diff_hours
|
||||
FROM {mu} mu
|
||||
WHERE mu.bank_id = $4
|
||||
AND mu.fact_type = src.fact_type
|
||||
AND mu.event_date > src.event_date
|
||||
AND mu.id != src.unit_id
|
||||
ORDER BY mu.event_date ASC
|
||||
LIMIT $5)
|
||||
) combined
|
||||
) ranked
|
||||
WHERE rn <= $5
|
||||
""",
|
||||
lateral_unit_ids[batch_start:batch_end],
|
||||
lateral_event_dates[batch_start:batch_end],
|
||||
lateral_fact_types[batch_start:batch_end],
|
||||
bank_id,
|
||||
half_limit,
|
||||
)
|
||||
rows.extend(batch_rows)
|
||||
|
||||
# Bidirectional index scan: instead of scanning all units in the 24h
|
||||
# window (O(N) — 164k rows at scale) and sorting by proximity, we scan
|
||||
# the nearest K units in each direction using the B-tree index on
|
||||
# (bank_id, fact_type, event_date). This reads only 2×K rows per probe
|
||||
# regardless of bank size — 120x faster at 164k units (0.6ms vs 74ms).
|
||||
rows = await ops.fetch_temporal_neighbors(
|
||||
conn,
|
||||
mu,
|
||||
bank_id,
|
||||
lateral_unit_ids,
|
||||
lateral_event_dates,
|
||||
lateral_fact_types,
|
||||
half_limit,
|
||||
batch_size=TEMPORAL_LATERAL_BATCH,
|
||||
)
|
||||
else:
|
||||
rows = []
|
||||
|
||||
@@ -686,7 +621,7 @@ async def create_temporal_links_batch_per_fact(
|
||||
|
||||
if links:
|
||||
insert_start = time_mod.time()
|
||||
await _bulk_insert_links(conn, links, bank_id=bank_id, skip_exists_check=True)
|
||||
await _bulk_insert_links(conn, links, bank_id=bank_id, skip_exists_check=True, ops=ops)
|
||||
_log(log_buffer, f" [7.4] Insert {len(links)} temporal links: {time_mod.time() - insert_start:.3f}s")
|
||||
|
||||
return len(links)
|
||||
@@ -812,7 +747,6 @@ async def compute_semantic_links_ann(
|
||||
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)
|
||||
@@ -889,6 +823,7 @@ async def create_semantic_links_batch(
|
||||
threshold: float = 0.7,
|
||||
log_buffer: list[str] = None,
|
||||
pre_computed_ann_links: list[tuple] | None = None,
|
||||
ops=None,
|
||||
) -> int:
|
||||
"""
|
||||
Phase 2: Create semantic links (within-batch + pre-computed ANN results).
|
||||
@@ -937,7 +872,7 @@ async def create_semantic_links_batch(
|
||||
|
||||
if all_links:
|
||||
insert_start = time_mod.time()
|
||||
await _bulk_insert_links(conn, all_links, bank_id=bank_id)
|
||||
await _bulk_insert_links(conn, all_links, bank_id=bank_id, ops=ops)
|
||||
_log(
|
||||
log_buffer, f" [8.3] Insert {len(all_links)} semantic links: {time_mod.time() - insert_start:.3f}s"
|
||||
)
|
||||
@@ -952,7 +887,7 @@ async def create_semantic_links_batch(
|
||||
raise
|
||||
|
||||
|
||||
async def insert_entity_links_batch(conn, links: list[EntityLink], bank_id: str, chunk_size: int = 5000):
|
||||
async def insert_entity_links_batch(conn, links: list[EntityLink], bank_id: str, chunk_size: int = 5000, ops=None):
|
||||
"""
|
||||
Bulk-insert entity links via sorted INSERT FROM unnest().
|
||||
|
||||
@@ -969,7 +904,7 @@ async def insert_entity_links_batch(conn, links: list[EntityLink], bank_id: str,
|
||||
|
||||
total_start = time_mod.time()
|
||||
tuples = [(link.from_unit_id, link.to_unit_id, link.link_type, link.weight, link.entity_id) for link in links]
|
||||
await _bulk_insert_links(conn, tuples, bank_id=bank_id, chunk_size=chunk_size)
|
||||
await _bulk_insert_links(conn, tuples, bank_id=bank_id, chunk_size=chunk_size, ops=ops)
|
||||
logger.debug(
|
||||
f" [9.TOTAL] Entity links batch insert ({len(tuples)} rows): {time_mod.time() - total_start:.3f}s"
|
||||
)
|
||||
@@ -980,6 +915,7 @@ async def create_causal_links_batch(
|
||||
bank_id: str,
|
||||
unit_ids: list[str],
|
||||
causal_relations_per_fact: list[list[dict]],
|
||||
ops=None,
|
||||
) -> int:
|
||||
"""
|
||||
Create causal links between facts based on LLM-extracted causal relationships.
|
||||
@@ -991,7 +927,6 @@ async def create_causal_links_batch(
|
||||
Each element is a list of dicts with:
|
||||
- target_fact_index: Index into unit_ids for the target fact
|
||||
- relation_type: "caused_by"
|
||||
- strength: Float in [0.0, 1.0] representing relationship strength
|
||||
|
||||
Returns:
|
||||
Number of causal links created
|
||||
@@ -1018,7 +953,6 @@ async def create_causal_links_batch(
|
||||
for relation in causal_relations:
|
||||
target_idx = relation["target_fact_index"]
|
||||
relation_type = relation["relation_type"]
|
||||
strength = relation.get("strength", 1.0)
|
||||
|
||||
# Validate relation_type - only "caused_by" is supported (DB constraint)
|
||||
valid_types = {"caused_by"}
|
||||
@@ -1041,14 +975,11 @@ async def create_causal_links_batch(
|
||||
if from_unit_id == to_unit_id:
|
||||
continue
|
||||
|
||||
# Add the causal link
|
||||
# link_type is the relation_type (e.g., "causes", "caused_by")
|
||||
# weight is the strength of the relationship
|
||||
links.append((from_unit_id, to_unit_id, relation_type, strength, None))
|
||||
links.append((from_unit_id, to_unit_id, relation_type, 1.0, None))
|
||||
|
||||
if links:
|
||||
insert_start = time_mod.time()
|
||||
await _bulk_insert_links(conn, links, bank_id=bank_id, skip_exists_check=True)
|
||||
await _bulk_insert_links(conn, links, bank_id=bank_id, skip_exists_check=True, ops=ops)
|
||||
logger.debug(f" [10.1] Insert {len(links)} causal links: {time_mod.time() - insert_start:.3f}s")
|
||||
|
||||
return len(links)
|
||||
|
||||
@@ -15,8 +15,9 @@ from datetime import UTC, datetime
|
||||
from typing import Any
|
||||
|
||||
from ...worker.stage import set_stage
|
||||
from ..db.base import DatabaseBackend
|
||||
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
|
||||
|
||||
|
||||
@@ -25,6 +26,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.
|
||||
@@ -72,7 +99,6 @@ from . import (
|
||||
from .types import (
|
||||
ChunkMetadata,
|
||||
EntityResolutionResult,
|
||||
ExtractedFact,
|
||||
Phase1Result,
|
||||
Phase3Context,
|
||||
ProcessedFact,
|
||||
@@ -115,7 +141,7 @@ def _build_retain_params(contents_dicts, document_tags=None, doc_contents=None):
|
||||
|
||||
|
||||
async def _pre_resolve_phase1(
|
||||
pool,
|
||||
pool: Any,
|
||||
entity_resolver,
|
||||
bank_id: str,
|
||||
contents: list[RetainContent],
|
||||
@@ -229,6 +255,7 @@ async def _insert_facts_and_links(
|
||||
semantic_ann_links: list[tuple],
|
||||
skip_semantic_links: bool = False,
|
||||
outbox_callback=None,
|
||||
ops=None,
|
||||
) -> tuple[list[list[str]], Phase3Context]:
|
||||
"""
|
||||
Phase 2 of the retain pipeline: insert facts and retrieval-critical links.
|
||||
@@ -241,7 +268,7 @@ async def _insert_facts_and_links(
|
||||
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)
|
||||
unit_ids = await fact_storage.insert_facts_batch(conn, bank_id, processed_facts, ops=ops)
|
||||
step_start = time.time()
|
||||
log_buffer.append(f" Insert facts: {len(unit_ids)} units in {time.time() - step_start:.3f}s")
|
||||
|
||||
@@ -274,7 +301,7 @@ async def _insert_facts_and_links(
|
||||
|
||||
# Create temporal links
|
||||
step_start = time.time()
|
||||
temporal_link_count = await link_creation.create_temporal_links_batch(conn, bank_id, unit_ids)
|
||||
temporal_link_count = await link_creation.create_temporal_links_batch(conn, bank_id, unit_ids, ops=ops)
|
||||
log_buffer.append(f" Temporal links: {temporal_link_count} links in {time.time() - step_start:.3f}s")
|
||||
|
||||
# Create semantic links (within-batch + pre-computed ANN from Phase 1)
|
||||
@@ -290,6 +317,7 @@ async def _insert_facts_and_links(
|
||||
unit_ids,
|
||||
embeddings_for_links,
|
||||
pre_computed_ann_links=semantic_ann_links,
|
||||
ops=ops,
|
||||
)
|
||||
log_buffer.append(f" Semantic links: {semantic_link_count} links in {time.time() - step_start:.3f}s")
|
||||
|
||||
@@ -299,11 +327,16 @@ async def _insert_facts_and_links(
|
||||
|
||||
# Create causal links
|
||||
step_start = time.time()
|
||||
causal_link_count = await link_creation.create_causal_links_batch(conn, bank_id, unit_ids, processed_facts)
|
||||
causal_link_count = await link_creation.create_causal_links_batch(
|
||||
conn, bank_id, unit_ids, processed_facts, ops=ops
|
||||
)
|
||||
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)
|
||||
@@ -312,7 +345,7 @@ async def _insert_facts_and_links(
|
||||
|
||||
|
||||
async def _build_and_insert_entity_links_phase3(
|
||||
pool,
|
||||
pool: Any,
|
||||
entity_resolver,
|
||||
bank_id: str,
|
||||
phase3_ctx: Phase3Context,
|
||||
@@ -346,9 +379,10 @@ async def _build_and_insert_entity_links_phase3(
|
||||
p3_unit_to_entity_ids,
|
||||
log_buffer,
|
||||
skip_unit_entities_insert=True, # Already inserted in Phase 2
|
||||
ops=pool.ops,
|
||||
)
|
||||
if entity_links:
|
||||
await entity_processing.insert_entity_links_batch(conn, entity_links, bank_id)
|
||||
await entity_processing.insert_entity_links_batch(conn, entity_links, bank_id, ops=pool.ops)
|
||||
log_buffer.append(f" Entity links (viz): {len(entity_links)} links in {time.time() - step_start:.3f}s")
|
||||
|
||||
|
||||
@@ -361,7 +395,7 @@ async def _extract_and_embed(
|
||||
format_date_fn,
|
||||
fact_type_override: str | None,
|
||||
log_buffer: list[str],
|
||||
pool=None,
|
||||
pool: Any = None,
|
||||
operation_id: str | None = None,
|
||||
schema: str | None = None,
|
||||
) -> tuple[list, list[ProcessedFact], list[ChunkMetadata], TokenUsage]:
|
||||
@@ -399,7 +433,7 @@ async def _extract_and_embed(
|
||||
|
||||
|
||||
async def retain_batch(
|
||||
pool,
|
||||
pool: Any,
|
||||
embeddings_model,
|
||||
llm_config,
|
||||
entity_resolver,
|
||||
@@ -415,13 +449,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)
|
||||
@@ -461,8 +503,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,
|
||||
@@ -484,11 +527,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")}
|
||||
@@ -507,26 +551,41 @@ 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
|
||||
@@ -557,6 +616,31 @@ async def retain_batch(
|
||||
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:
|
||||
delta_result = await _try_delta_retain(
|
||||
@@ -639,7 +723,7 @@ _ANN_PARALLELISM = 4 # Max concurrent ANN chunks to avoid pool saturation
|
||||
|
||||
|
||||
async def _run_final_semantic_ann(
|
||||
pool,
|
||||
pool: Any,
|
||||
bank_id: str,
|
||||
unit_ids: list[str],
|
||||
log_buffer: list[str],
|
||||
@@ -713,7 +797,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,
|
||||
@@ -724,9 +807,8 @@ async def _run_final_semantic_ann(
|
||||
log_buffer=log_buffer,
|
||||
)
|
||||
if ann_links:
|
||||
await _bulk_insert_links(conn, ann_links, bank_id=bank_id)
|
||||
await _bulk_insert_links(conn, ann_links, bank_id=bank_id, ops=pool.ops)
|
||||
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"
|
||||
@@ -743,7 +825,7 @@ async def _run_final_semantic_ann(
|
||||
|
||||
|
||||
async def _streaming_retain_batch(
|
||||
pool,
|
||||
pool: Any,
|
||||
embeddings_model,
|
||||
llm_config,
|
||||
entity_resolver,
|
||||
@@ -792,25 +874,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:
|
||||
@@ -822,24 +906,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
|
||||
@@ -852,6 +934,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
|
||||
@@ -910,17 +996,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,
|
||||
)
|
||||
@@ -929,19 +1013,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:
|
||||
@@ -955,15 +1044,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)
|
||||
@@ -975,6 +1066,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"
|
||||
@@ -1005,16 +1136,98 @@ 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 = {}
|
||||
if batch_chunk_meta:
|
||||
chunk_id_map = await chunk_storage.store_chunks_batch(
|
||||
conn, bank_id, effective_doc_id, batch_chunk_meta
|
||||
conn, bank_id, effective_doc_id, batch_chunk_meta, ops=pool.ops
|
||||
)
|
||||
log_buffer.append(
|
||||
f" Store chunks: {len(batch_chunk_meta)} chunks in {time.time() - step_start:.3f}s"
|
||||
@@ -1045,16 +1258,20 @@ async def _streaming_retain_batch(
|
||||
semantic_ann_links=[],
|
||||
skip_semantic_links=True,
|
||||
outbox_callback=outbox_callback if is_last else None,
|
||||
ops=pool.ops,
|
||||
)
|
||||
|
||||
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 "
|
||||
@@ -1062,8 +1279,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:
|
||||
@@ -1106,6 +1324,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:
|
||||
@@ -1142,16 +1401,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")
|
||||
@@ -1159,7 +1433,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
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -1168,7 +1445,7 @@ async def _streaming_retain_batch(
|
||||
|
||||
|
||||
async def _try_delta_retain(
|
||||
pool,
|
||||
pool: Any,
|
||||
embeddings_model,
|
||||
llm_config,
|
||||
entity_resolver,
|
||||
@@ -1187,10 +1464,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
|
||||
@@ -1200,9 +1482,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
|
||||
@@ -1310,8 +1600,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])
|
||||
@@ -1363,7 +1673,7 @@ async def _try_delta_retain(
|
||||
for cm in new_chunk_metadata
|
||||
]
|
||||
chunk_id_map = await chunk_storage.store_chunks_batch(
|
||||
conn, bank_id, effective_doc_id, remapped_chunks
|
||||
conn, bank_id, effective_doc_id, remapped_chunks, ops=pool.ops
|
||||
)
|
||||
for chunk_idx, chunk_id in chunk_id_map.items():
|
||||
chunk_id_map_by_doc[(effective_doc_id, chunk_idx)] = chunk_id
|
||||
@@ -1397,6 +1707,7 @@ async def _try_delta_retain(
|
||||
unit_to_entity_ids=phase1.entities.unit_to_entity_ids,
|
||||
semantic_ann_links=phase1.semantic_ann_links,
|
||||
outbox_callback=outbox_callback,
|
||||
ops=pool.ops,
|
||||
)
|
||||
|
||||
# PHASE 3 — Best-Effort Display Data (post-transaction)
|
||||
@@ -1421,11 +1732,16 @@ 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(
|
||||
pool,
|
||||
pool: Any,
|
||||
bank_id,
|
||||
contents_dicts,
|
||||
contents,
|
||||
@@ -1438,6 +1754,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(
|
||||
@@ -1455,7 +1777,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
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -1550,12 +1876,19 @@ 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):
|
||||
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
|
||||
@@ -1564,12 +1897,8 @@ def _map_results_to_contents(
|
||||
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
|
||||
|
||||
@@ -99,7 +99,6 @@ class CausalRelation:
|
||||
|
||||
relation_type: str # "caused_by"
|
||||
target_fact_index: int # Index of the target fact in the batch
|
||||
strength: float = 1.0 # Strength of the causal relationship
|
||||
|
||||
|
||||
@dataclass
|
||||
|
||||
@@ -0,0 +1,41 @@
|
||||
"""
|
||||
Centralized schema-qualified table name helpers.
|
||||
|
||||
Single source of truth for producing ``"schema".table_name`` references
|
||||
that respect both the active schema context and the database backend.
|
||||
"""
|
||||
|
||||
from ..config import get_config
|
||||
|
||||
|
||||
def _is_oracle() -> bool:
|
||||
"""Return True when the configured database backend is Oracle."""
|
||||
return get_config().database_backend == "oracle"
|
||||
|
||||
|
||||
def fq_table(table_name: str) -> str:
|
||||
"""Get fully-qualified table name using the current schema context.
|
||||
|
||||
On Oracle the schema is set at the session level (``ALTER SESSION SET
|
||||
CURRENT_SCHEMA``), so we return the bare table name. On PostgreSQL
|
||||
we prefix with the schema from :func:`memory_engine.get_current_schema`.
|
||||
"""
|
||||
if _is_oracle():
|
||||
return table_name
|
||||
from .memory_engine import get_current_schema
|
||||
|
||||
return f"{get_current_schema()}.{table_name}"
|
||||
|
||||
|
||||
def fq_table_explicit(table: str, schema: str | None = None) -> str:
|
||||
"""Get fully-qualified table name with an explicit schema override.
|
||||
|
||||
Used by modules that don't rely on the context-variable schema
|
||||
(e.g. task_backend, worker poller) and instead pass the schema
|
||||
explicitly.
|
||||
"""
|
||||
if _is_oracle():
|
||||
return table
|
||||
if schema:
|
||||
return f'"{schema}".{table}'
|
||||
return table
|
||||
@@ -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.
|
||||
|
||||
@@ -28,6 +28,8 @@ 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
|
||||
@@ -49,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
|
||||
@@ -56,10 +60,24 @@ 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"""
|
||||
@@ -73,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
|
||||
""",
|
||||
@@ -93,15 +112,8 @@ class LinkExpansionRetriever(GraphRetriever):
|
||||
The Python merge step applies per-signal score transformations.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
causal_weight_threshold: float = 0.3,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
causal_weight_threshold: Minimum weight for causal links to follow.
|
||||
"""
|
||||
self.causal_weight_threshold = causal_weight_threshold
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
@@ -121,6 +133,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.
|
||||
@@ -159,6 +173,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(
|
||||
@@ -177,10 +193,15 @@ class LinkExpansionRetriever(GraphRetriever):
|
||||
|
||||
query_start = time.time()
|
||||
|
||||
ops = pool.ops
|
||||
if fact_type == "observation":
|
||||
entity_rows, semantic_rows, causal_rows = await self._expand_observations(conn, seed_ids, budget)
|
||||
entity_rows, semantic_rows, causal_rows = await self._expand_observations(
|
||||
conn, seed_ids, budget, ops=ops
|
||||
)
|
||||
else:
|
||||
entity_rows, semantic_rows, causal_rows = await self._expand_combined(conn, seed_ids, fact_type, budget)
|
||||
entity_rows, semantic_rows, causal_rows = await self._expand_combined(
|
||||
conn, seed_ids, fact_type, budget, ops=ops
|
||||
)
|
||||
|
||||
timings.edge_load_time = time.time() - query_start
|
||||
timings.db_queries = 1
|
||||
@@ -252,6 +273,8 @@ class LinkExpansionRetriever(GraphRetriever):
|
||||
seed_ids: list,
|
||||
fact_type: str,
|
||||
budget: int,
|
||||
*,
|
||||
ops,
|
||||
) -> tuple[list, list, list]:
|
||||
"""
|
||||
Single-roundtrip CTE query combining entity, semantic, and causal expansions.
|
||||
@@ -274,101 +297,8 @@ class LinkExpansionRetriever(GraphRetriever):
|
||||
|
||||
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 (
|
||||
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.proof_count,
|
||||
COUNT(DISTINCT se.entity_id)::float AS score,
|
||||
'entity'::text AS source
|
||||
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
|
||||
)"""
|
||||
|
||||
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).
|
||||
-- Score = max similarity weight across both directions.
|
||||
SELECT
|
||||
id, text, context, event_date, occurred_start,
|
||||
occurred_end, mentioned_at,
|
||||
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.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 = $2
|
||||
AND mu.id != ALL($1::uuid[])
|
||||
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, 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 = $2
|
||||
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, proof_count
|
||||
ORDER BY score DESC
|
||||
LIMIT $3
|
||||
),
|
||||
causal_expanded AS (
|
||||
-- Causal chains: explicit causes/enables/prevents links from seeds.
|
||||
-- DISTINCT ON handles the case where a seed has multiple causal links
|
||||
-- to the same target; best weight wins.
|
||||
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.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')
|
||||
AND ml.weight >= $4
|
||||
AND mu.fact_type = $2
|
||||
ORDER BY mu.id, ml.weight DESC
|
||||
LIMIT $3
|
||||
)"""
|
||||
entity_cte = ops.build_entity_expansion_cte(mu, ue, per_entity_limit)
|
||||
semantic_causal_cte = ops.build_semantic_causal_cte(ml, mu)
|
||||
|
||||
full_query = f"""
|
||||
WITH {entity_cte},
|
||||
@@ -380,7 +310,7 @@ class LinkExpansionRetriever(GraphRetriever):
|
||||
SELECT * FROM causal_expanded
|
||||
"""
|
||||
|
||||
params = [seed_ids, fact_type, budget, self.causal_weight_threshold]
|
||||
params = [seed_ids, fact_type, budget]
|
||||
|
||||
try:
|
||||
all_rows = await asyncio.wait_for(
|
||||
@@ -397,6 +327,7 @@ class LinkExpansionRetriever(GraphRetriever):
|
||||
SELECT * FROM semantic_expanded
|
||||
UNION ALL
|
||||
SELECT * FROM causal_expanded
|
||||
LIMIT $3
|
||||
"""
|
||||
all_rows = await conn.fetch(fallback_query, *params)
|
||||
|
||||
@@ -410,6 +341,8 @@ class LinkExpansionRetriever(GraphRetriever):
|
||||
conn,
|
||||
seed_ids: list,
|
||||
budget: int,
|
||||
*,
|
||||
ops,
|
||||
) -> tuple[list, list, list]:
|
||||
"""
|
||||
Observation-specific expansion.
|
||||
@@ -440,114 +373,19 @@ class LinkExpansionRetriever(GraphRetriever):
|
||||
|
||||
config = get_config()
|
||||
ue = fq_table("unit_entities")
|
||||
per_entity_limit = config.link_expansion_per_entity_limit
|
||||
|
||||
connected_sources_cte = f"""
|
||||
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
|
||||
),
|
||||
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(
|
||||
f"""
|
||||
WITH seed_sources AS (
|
||||
SELECT DISTINCT unnest(source_memory_ids) AS source_id
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE id = ANY($1::uuid[])
|
||||
AND source_memory_ids IS NOT NULL
|
||||
),
|
||||
{connected_sources_cte},
|
||||
connected_array AS (
|
||||
SELECT array_agg(source_id) AS source_ids FROM connected_sources
|
||||
)
|
||||
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.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'
|
||||
AND mu.id != ALL($1::uuid[])
|
||||
AND ca.source_ids IS NOT NULL
|
||||
AND mu.source_memory_ids && ca.source_ids
|
||||
ORDER BY score DESC
|
||||
LIMIT $2
|
||||
""",
|
||||
seed_ids,
|
||||
budget,
|
||||
)
|
||||
logger.debug(f"[LinkExpansion] observation graph: found {len(entity_rows)} connected observations")
|
||||
|
||||
# Semantic + causal for observations in one query
|
||||
ml = fq_table("memory_links")
|
||||
mu = fq_table("memory_units")
|
||||
sem_causal_rows = await conn.fetch(
|
||||
f"""
|
||||
WITH semantic_expanded AS (
|
||||
SELECT
|
||||
id, text, context, event_date, occurred_start,
|
||||
occurred_end, mentioned_at,
|
||||
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.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'
|
||||
AND mu.id != ALL($1::uuid[])
|
||||
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, 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, 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, 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')
|
||||
AND ml.weight >= $3 AND mu.fact_type = 'observation'
|
||||
ORDER BY mu.id, ml.weight DESC LIMIT $2
|
||||
)
|
||||
SELECT * FROM semantic_expanded
|
||||
UNION ALL
|
||||
SELECT * FROM causal_expanded
|
||||
""",
|
||||
per_entity_limit = config.link_expansion_per_entity_limit
|
||||
|
||||
# Delegate to DataAccessOps. Both backends now use the observation_sources
|
||||
# junction table with standard SQL joins (previously PG used native array
|
||||
# ops and Oracle used JSON_TABLE).
|
||||
return await ops.expand_observations(
|
||||
conn,
|
||||
mu,
|
||||
ue,
|
||||
ml,
|
||||
seed_ids,
|
||||
budget,
|
||||
self.causal_weight_threshold,
|
||||
per_entity_limit,
|
||||
)
|
||||
|
||||
semantic_rows = [r for r in sem_causal_rows if r["source"] == "semantic"]
|
||||
causal_rows = [r for r in sem_causal_rows if r["source"] == "causal"]
|
||||
return entity_rows, semantic_rows, causal_rows
|
||||
|
||||
@@ -13,11 +13,12 @@ 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
|
||||
from ..memory_engine import fq_table
|
||||
from ..sql import create_sql_dialect
|
||||
from .graph_retrieval import GraphRetriever
|
||||
from .link_expansion_retrieval import LinkExpansionRetriever
|
||||
from .tags import TagGroup, TagsMatch, build_tag_groups_where_clause, build_tags_where_clause_simple
|
||||
@@ -98,6 +99,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.
|
||||
@@ -145,6 +148,14 @@ async def retrieve_semantic_bm25_combined(
|
||||
)
|
||||
table = fq_table("memory_units")
|
||||
|
||||
config = get_config()
|
||||
|
||||
# Use the SQL dialect to build backend-specific query arms, avoiding
|
||||
# inline if/else branches for each database.
|
||||
# Use getattr for backward compat: raw asyncpg connections (used in some
|
||||
# tests) lack backend_type; default to "postgresql".
|
||||
dialect = create_sql_dialect(getattr(conn, "backend_type", "postgresql"))
|
||||
|
||||
# --- Parameter layout ---
|
||||
# $1 = query_emb_str (semantic arms)
|
||||
# $2 = bank_id
|
||||
@@ -153,88 +164,129 @@ async def retrieve_semantic_bm25_combined(
|
||||
# $4 = bm25_text
|
||||
# $5 = tags (if present)
|
||||
# $6+ = tag_groups params (one per leaf)
|
||||
# When no tokens ($3 is skipped — not included in params to avoid type inference gap):
|
||||
# When no tokens:
|
||||
# $3 = tags (if present)
|
||||
# $4+ = tag_groups params (one per leaf)
|
||||
tags_param_idx = 5 if tokens else 3
|
||||
_include_bm25 = bool(tokens)
|
||||
tags_param_idx = 5 if _include_bm25 else 3
|
||||
tags_clause = build_tags_where_clause_simple(tags, tags_param_idx, match=tags_match)
|
||||
|
||||
# tag_groups params start immediately after the tags param slot
|
||||
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)
|
||||
|
||||
# --- 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.
|
||||
sem_arms = []
|
||||
for ft in fact_types:
|
||||
sem_arms.append(
|
||||
f"(SELECT {cols},"
|
||||
f" 1 - (embedding <=> $1::vector) AS similarity,"
|
||||
f" NULL::float AS bm25_score,"
|
||||
f" 'semantic' AS source"
|
||||
f" FROM {table}"
|
||||
f" WHERE bank_id = $2"
|
||||
f" AND fact_type = '{ft}'"
|
||||
f" AND embedding IS NOT NULL"
|
||||
f" AND (1 - (embedding <=> $1::vector)) >= 0.3"
|
||||
f" {tags_clause}"
|
||||
f" {groups_clause}"
|
||||
f" ORDER BY embedding <=> $1::vector"
|
||||
f" LIMIT {hnsw_fetch})"
|
||||
)
|
||||
# --- 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
|
||||
|
||||
arms = sem_arms
|
||||
# --- Semantic UNION ALL arms (one per fact_type) ---
|
||||
# Each arm has its own ORDER BY ... LIMIT, enabling the partial HNSW indexes
|
||||
# per fact_type instead of forcing a full sequential scan.
|
||||
arms = [
|
||||
dialect.build_semantic_arm(
|
||||
table=table,
|
||||
cols=cols,
|
||||
fact_type=ft,
|
||||
embedding_param="$1",
|
||||
bank_id_param="$2",
|
||||
fetch_limit=hnsw_fetch,
|
||||
tags_clause=tags_clause,
|
||||
groups_clause=groups_clause,
|
||||
extra_where=created_range_clause,
|
||||
)
|
||||
for ft in fact_types
|
||||
]
|
||||
|
||||
# --- BM25 UNION ALL arms (one per fact_type, only when tokens present) ---
|
||||
if tokens:
|
||||
config = get_config()
|
||||
if config.text_search_extension == "vchord":
|
||||
bm25_score_expr = (
|
||||
"search_vector <&> to_bm25query('idx_memory_units_text_search', tokenize($4, 'llmlingua2'))"
|
||||
)
|
||||
bm25_order_by = f"{bm25_score_expr} DESC"
|
||||
bm25_where_filter = ""
|
||||
bm25_text_param: str = query_text
|
||||
elif config.text_search_extension == "pg_textsearch":
|
||||
bm25_score_expr = "-(text <@> to_bm25query($4, 'idx_memory_units_text_search'))"
|
||||
bm25_order_by = "text <@> to_bm25query($4, 'idx_memory_units_text_search') ASC"
|
||||
bm25_where_filter = ""
|
||||
bm25_text_param = query_text
|
||||
else: # native
|
||||
query_tsquery = " | ".join(tokens)
|
||||
bm25_score_expr = "ts_rank_cd(search_vector, to_tsquery('english', $4))"
|
||||
bm25_order_by = f"{bm25_score_expr} DESC"
|
||||
bm25_where_filter = "AND search_vector @@ to_tsquery('english', $4)"
|
||||
bm25_text_param = query_tsquery
|
||||
|
||||
for ft in fact_types:
|
||||
if _include_bm25:
|
||||
text_ext = config.text_search_extension
|
||||
bm25_text_param: str = dialect.prepare_bm25_text(tokens, query_text, text_search_extension=text_ext)
|
||||
for i, ft in enumerate(fact_types):
|
||||
arms.append(
|
||||
f"(SELECT {cols},"
|
||||
f" NULL::float AS similarity,"
|
||||
f" {bm25_score_expr} AS bm25_score,"
|
||||
f" 'bm25' AS source"
|
||||
f" FROM {table}"
|
||||
f" WHERE bank_id = $2"
|
||||
f" AND fact_type = '{ft}'"
|
||||
f" {bm25_where_filter}"
|
||||
f" {tags_clause}"
|
||||
f" {groups_clause}"
|
||||
f" ORDER BY {bm25_order_by}"
|
||||
f" LIMIT $3)"
|
||||
dialect.build_bm25_arm(
|
||||
table=table,
|
||||
cols=cols,
|
||||
fact_type=ft,
|
||||
bank_id_param="$2",
|
||||
limit_param="$3",
|
||||
text_param="$4",
|
||||
tags_clause=tags_clause,
|
||||
groups_clause=groups_clause,
|
||||
arm_index=i,
|
||||
text_search_extension=text_ext,
|
||||
extra_where=created_range_clause,
|
||||
)
|
||||
)
|
||||
|
||||
query = "\nUNION ALL\n".join(arms)
|
||||
|
||||
params: list = [query_emb_str, bank_id]
|
||||
if tokens:
|
||||
if _include_bm25:
|
||||
params.append(limit) # $3: BM25 LIMIT (only referenced when tokens are present)
|
||||
params.append(bm25_text_param) # $4
|
||||
if tags:
|
||||
params.append(tags)
|
||||
params.extend(groups_params)
|
||||
params.extend(created_range_params)
|
||||
|
||||
rows = await conn.fetch(query, *params)
|
||||
try:
|
||||
rows = await conn.fetch(query, *params)
|
||||
except Exception as e:
|
||||
# Oracle Text CONTAINS can fail with DRG-10599 ("column is not indexed")
|
||||
# if the CTXSYS text index hasn't synced yet or is unavailable. Fall
|
||||
# back to semantic-only so the search still returns results.
|
||||
# We must rebuild the semantic arms with no-BM25 param indices because
|
||||
# Oracle requires every bind param to be referenced in the query (DPY-4008).
|
||||
err_str = str(e)
|
||||
if _include_bm25 and ("DRG-10599" in err_str or "ORA-30600" in err_str or "ORA-29902" in err_str):
|
||||
logger.warning("Oracle Text CONTAINS failed (%s), falling back to semantic-only search", err_str[:120])
|
||||
# Rebuild with no-BM25 param layout: $1=embedding, $2=bank_id, $3=tags, ...
|
||||
fb_tags_idx = 3
|
||||
fb_tags_clause = build_tags_where_clause_simple(tags, fb_tags_idx, match=tags_match)
|
||||
fb_groups_start = fb_tags_idx + (1 if tags else 0)
|
||||
fb_groups_clause, _, _ = build_tag_groups_where_clause(tag_groups, fb_groups_start)
|
||||
fb_next_idx = fb_groups_start + len(groups_params)
|
||||
fb_created_clause = ""
|
||||
if created_after is not None:
|
||||
fb_created_clause += f" AND updated_at > ${fb_next_idx}"
|
||||
fb_next_idx += 1
|
||||
if created_before is not None:
|
||||
fb_created_clause += f" AND updated_at < ${fb_next_idx}"
|
||||
fb_next_idx += 1
|
||||
fb_arms = [
|
||||
dialect.build_semantic_arm(
|
||||
table=table,
|
||||
cols=cols,
|
||||
fact_type=ft,
|
||||
embedding_param="$1",
|
||||
bank_id_param="$2",
|
||||
fetch_limit=hnsw_fetch,
|
||||
tags_clause=fb_tags_clause,
|
||||
groups_clause=fb_groups_clause,
|
||||
extra_where=fb_created_clause,
|
||||
)
|
||||
for ft in fact_types
|
||||
]
|
||||
fb_query = "\nUNION ALL\n".join(fb_arms)
|
||||
fb_params: list = [query_emb_str, bank_id]
|
||||
if tags:
|
||||
fb_params.append(tags)
|
||||
fb_params.extend(groups_params)
|
||||
fb_params.extend(created_range_params)
|
||||
rows = await conn.fetch(fb_query, *fb_params)
|
||||
else:
|
||||
raise
|
||||
|
||||
# Group results; trim semantic to limit (over-fetched for HNSW approximation).
|
||||
sem_counts: dict[str, int] = {ft: 0 for ft in fact_types}
|
||||
@@ -266,6 +318,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 +353,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,6 +403,7 @@ 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.proof_count, mu.document_id, mu.chunk_id, mu.tags, mu.metadata,
|
||||
@@ -536,6 +606,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 +666,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 +687,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 +712,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
|
||||
|
||||
|
||||
@@ -73,7 +73,7 @@ def format_facts_for_prompt(facts: list[MemoryFact]) -> str:
|
||||
|
||||
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:
|
||||
|
||||
@@ -0,0 +1,41 @@
|
||||
"""SQL dialect abstraction layer.
|
||||
|
||||
Isolates database-specific SQL syntax (parameter placeholders, JSON operators,
|
||||
vector distance functions, etc.) behind a common interface.
|
||||
|
||||
Usage:
|
||||
from hindsight_api.engine.sql import create_sql_dialect, SQLDialect
|
||||
|
||||
dialect = create_sql_dialect("postgresql")
|
||||
placeholder = dialect.param(1) # "$1" for PG, ":1" for Oracle
|
||||
"""
|
||||
|
||||
from .base import SQLDialect
|
||||
|
||||
__all__ = [
|
||||
"SQLDialect",
|
||||
"create_sql_dialect",
|
||||
]
|
||||
|
||||
|
||||
def create_sql_dialect(backend_type: str) -> SQLDialect:
|
||||
"""Factory: create a SQLDialect by backend name.
|
||||
|
||||
Args:
|
||||
backend_type: One of "postgresql" or "oracle".
|
||||
|
||||
Returns:
|
||||
A SQLDialect instance.
|
||||
|
||||
Raises:
|
||||
ValueError: If backend_type is not recognized.
|
||||
"""
|
||||
if backend_type == "postgresql":
|
||||
from .postgresql import PostgreSQLDialect
|
||||
|
||||
return PostgreSQLDialect()
|
||||
elif backend_type == "oracle":
|
||||
from .oracle import OracleDialect
|
||||
|
||||
return OracleDialect()
|
||||
raise ValueError(f"Unknown SQL dialect: {backend_type!r}. Supported dialects: 'postgresql', 'oracle'.")
|
||||
@@ -0,0 +1,455 @@
|
||||
"""Abstract base class for SQL dialect modules.
|
||||
|
||||
Each method encapsulates a SQL pattern that differs between database platforms.
|
||||
Business logic calls these methods instead of embedding raw SQL fragments.
|
||||
"""
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
|
||||
|
||||
class SQLDialect(ABC):
|
||||
"""SQL dialect interface for portable query construction.
|
||||
|
||||
Implementors provide database-specific SQL fragments for operations that
|
||||
are not standard across PostgreSQL and Oracle (parameter binding, JSON
|
||||
operators, vector distance, full-text search, etc.).
|
||||
"""
|
||||
|
||||
# -- Parameter binding -----------------------------------------------
|
||||
|
||||
@abstractmethod
|
||||
def param(self, n: int) -> str:
|
||||
"""Return the nth positional parameter placeholder.
|
||||
|
||||
Args:
|
||||
n: 1-based parameter index.
|
||||
|
||||
Returns:
|
||||
"$1" for PostgreSQL, ":1" for Oracle.
|
||||
"""
|
||||
...
|
||||
|
||||
# -- Type casting ----------------------------------------------------
|
||||
|
||||
@abstractmethod
|
||||
def cast(self, param: str, type_name: str) -> str:
|
||||
"""Cast a parameter or expression to the given type.
|
||||
|
||||
Args:
|
||||
param: The expression to cast (e.g. "$1" or a column name).
|
||||
type_name: Target type (e.g. "jsonb", "uuid[]", "vector").
|
||||
|
||||
Returns:
|
||||
Cast expression (e.g. "$1::jsonb" for PG, "CAST(:1 AS ...)" for Oracle).
|
||||
"""
|
||||
...
|
||||
|
||||
# -- Vector operations -----------------------------------------------
|
||||
|
||||
@abstractmethod
|
||||
def vector_distance(self, col: str, param: str) -> str:
|
||||
"""Cosine distance expression between a column and a parameter.
|
||||
|
||||
Args:
|
||||
col: Column name containing the vector.
|
||||
param: Parameter placeholder for the query vector.
|
||||
|
||||
Returns:
|
||||
Distance expression (lower = more similar).
|
||||
PG: "col <=> $1::vector"
|
||||
Oracle: "VECTOR_DISTANCE(col, :1, COSINE)"
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
def vector_similarity(self, col: str, param: str) -> str:
|
||||
"""Cosine similarity expression (1 - distance).
|
||||
|
||||
Args:
|
||||
col: Column name.
|
||||
param: Parameter placeholder.
|
||||
|
||||
Returns:
|
||||
Similarity expression (higher = more similar).
|
||||
"""
|
||||
...
|
||||
|
||||
# -- JSON operations -------------------------------------------------
|
||||
|
||||
@abstractmethod
|
||||
def json_extract_text(self, col: str, key: str) -> str:
|
||||
"""Extract a text value from a JSON/JSONB column.
|
||||
|
||||
Args:
|
||||
col: Column name.
|
||||
key: JSON key to extract.
|
||||
|
||||
Returns:
|
||||
PG: "col ->> 'key'"
|
||||
Oracle: "JSON_VALUE(col, '$.key')"
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
def json_contains(self, col: str, param: str) -> str:
|
||||
"""Test whether a JSON column contains the given JSON object.
|
||||
|
||||
Args:
|
||||
col: Column name.
|
||||
param: Parameter placeholder for the JSON object to test.
|
||||
|
||||
Returns:
|
||||
PG: "col @> $1::jsonb"
|
||||
Oracle: "JSON_EXISTS(col, ...)"
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
def json_merge(self, col: str, param: str) -> str:
|
||||
"""Merge (concatenate) a JSON object into a JSON column.
|
||||
|
||||
Args:
|
||||
col: Column name.
|
||||
param: Parameter placeholder for the JSON to merge.
|
||||
|
||||
Returns:
|
||||
PG: "col || $1::jsonb"
|
||||
Oracle: "JSON_MERGEPATCH(col, :1)"
|
||||
"""
|
||||
...
|
||||
|
||||
# -- Text search -----------------------------------------------------
|
||||
|
||||
@abstractmethod
|
||||
def text_search_score(self, col: str, query_param: str, *, index_name: str | None = None) -> str:
|
||||
"""Relevance score expression for full-text search.
|
||||
|
||||
Args:
|
||||
col: Column name (text or tsvector/bm25vector).
|
||||
query_param: Parameter placeholder for the search query.
|
||||
index_name: Optional index name (needed by some backends).
|
||||
|
||||
Returns:
|
||||
Score expression (higher = more relevant).
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
def text_search_order(self, col: str, query_param: str, *, index_name: str | None = None) -> str:
|
||||
"""ORDER BY expression for full-text search (ascending = best first).
|
||||
|
||||
Args:
|
||||
col: Column name.
|
||||
query_param: Parameter placeholder for the search query.
|
||||
index_name: Optional index name.
|
||||
|
||||
Returns:
|
||||
Expression suitable for ORDER BY ... ASC.
|
||||
"""
|
||||
...
|
||||
|
||||
# -- Fuzzy string matching -------------------------------------------
|
||||
|
||||
@abstractmethod
|
||||
def similarity(self, col: str, param: str) -> str:
|
||||
"""Fuzzy string similarity score between a column and a parameter.
|
||||
|
||||
Args:
|
||||
col: Column name.
|
||||
param: Parameter placeholder.
|
||||
|
||||
Returns:
|
||||
PG: "similarity(col, $1)"
|
||||
Oracle: "UTL_MATCH.EDIT_DISTANCE_SIMILARITY(col, :1) / 100.0"
|
||||
"""
|
||||
...
|
||||
|
||||
# -- Upsert ----------------------------------------------------------
|
||||
|
||||
@abstractmethod
|
||||
def upsert(
|
||||
self,
|
||||
table: str,
|
||||
columns: list[str],
|
||||
conflict_columns: list[str],
|
||||
update_columns: list[str],
|
||||
) -> str:
|
||||
"""Generate an upsert statement.
|
||||
|
||||
Args:
|
||||
table: Fully-qualified table name.
|
||||
columns: All columns in the INSERT.
|
||||
conflict_columns: Columns that form the unique constraint.
|
||||
update_columns: Columns to update on conflict.
|
||||
|
||||
Returns:
|
||||
Complete INSERT ... ON CONFLICT DO UPDATE (PG)
|
||||
or MERGE INTO ... (Oracle) statement.
|
||||
"""
|
||||
...
|
||||
|
||||
# -- Bulk operations -------------------------------------------------
|
||||
|
||||
@abstractmethod
|
||||
def bulk_unnest(self, param_types: list[tuple[str, str]]) -> str:
|
||||
"""Generate a bulk unnest/table-value expression.
|
||||
|
||||
Converts parallel arrays into rows.
|
||||
|
||||
Args:
|
||||
param_types: List of (param_placeholder, sql_type) pairs
|
||||
e.g. [("$1", "text[]"), ("$2", "uuid[]")]
|
||||
|
||||
Returns:
|
||||
PG: "unnest($1::text[], $2::uuid[])"
|
||||
Oracle: JSON_TABLE-based equivalent.
|
||||
"""
|
||||
...
|
||||
|
||||
# -- Pagination ------------------------------------------------------
|
||||
|
||||
@abstractmethod
|
||||
def limit_offset(self, limit_param: str, offset_param: str) -> str:
|
||||
"""Generate LIMIT/OFFSET clause.
|
||||
|
||||
Args:
|
||||
limit_param: Parameter placeholder for row limit.
|
||||
offset_param: Parameter placeholder for row offset.
|
||||
|
||||
Returns:
|
||||
PG: "LIMIT $1 OFFSET $2"
|
||||
Oracle: "OFFSET :2 ROWS FETCH FIRST :1 ROWS ONLY"
|
||||
"""
|
||||
...
|
||||
|
||||
# -- RETURNING clause ------------------------------------------------
|
||||
|
||||
@abstractmethod
|
||||
def returning(self, columns: list[str]) -> str:
|
||||
"""Generate a RETURNING clause.
|
||||
|
||||
Args:
|
||||
columns: Column names to return.
|
||||
|
||||
Returns:
|
||||
PG: "RETURNING col1, col2"
|
||||
Oracle: "RETURNING col1, col2 INTO :out1, :out2" (handled by backend).
|
||||
"""
|
||||
...
|
||||
|
||||
# -- Pattern matching ------------------------------------------------
|
||||
|
||||
@abstractmethod
|
||||
def ilike(self, col: str, param: str) -> str:
|
||||
"""Case-insensitive LIKE expression.
|
||||
|
||||
Args:
|
||||
col: Column name.
|
||||
param: Parameter placeholder for the pattern.
|
||||
|
||||
Returns:
|
||||
PG: "col ILIKE $1"
|
||||
Oracle: "UPPER(col) LIKE UPPER(:1)"
|
||||
"""
|
||||
...
|
||||
|
||||
# -- Array operations ------------------------------------------------
|
||||
|
||||
@abstractmethod
|
||||
def array_any(self, param: str) -> str:
|
||||
"""IN-array membership expression.
|
||||
|
||||
Args:
|
||||
param: Parameter placeholder for the array.
|
||||
|
||||
Returns:
|
||||
PG: "= ANY($1)"
|
||||
Oracle: "IN (SELECT ... FROM JSON_TABLE(...))"
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
def array_all(self, param: str) -> str:
|
||||
"""NOT-IN-array expression (not equal to all elements).
|
||||
|
||||
Args:
|
||||
param: Parameter placeholder for the array.
|
||||
|
||||
Returns:
|
||||
PG: "!= ALL($1)"
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
def array_contains(self, col: str, param: str) -> str:
|
||||
"""Test whether an array column contains all elements in the parameter.
|
||||
|
||||
Args:
|
||||
col: Array column name.
|
||||
param: Parameter placeholder for the array to test.
|
||||
|
||||
Returns:
|
||||
PG: "col @> $1::varchar[]"
|
||||
"""
|
||||
...
|
||||
|
||||
# -- Locking ---------------------------------------------------------
|
||||
|
||||
@abstractmethod
|
||||
def for_update_skip_locked(self) -> str:
|
||||
"""FOR UPDATE SKIP LOCKED clause (same on both PG and Oracle)."""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
def advisory_lock(self, id_param: str) -> str:
|
||||
"""Advisory lock expression.
|
||||
|
||||
Args:
|
||||
id_param: Parameter placeholder for the lock ID.
|
||||
|
||||
Returns:
|
||||
PG: "pg_try_advisory_lock($1)"
|
||||
Oracle: "SELECT ... FOR UPDATE NOWAIT" equivalent.
|
||||
"""
|
||||
...
|
||||
|
||||
# -- UUID generation -------------------------------------------------
|
||||
|
||||
@abstractmethod
|
||||
def generate_uuid(self) -> str:
|
||||
"""SQL expression to generate a random UUID.
|
||||
|
||||
Returns:
|
||||
PG: "gen_random_uuid()"
|
||||
Oracle: "SYS_GUID()"
|
||||
"""
|
||||
...
|
||||
|
||||
# -- Misc ------------------------------------------------------------
|
||||
|
||||
@abstractmethod
|
||||
def greatest(self, *args: str) -> str:
|
||||
"""GREATEST() function (same on both platforms)."""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
def current_timestamp(self) -> str:
|
||||
"""Current timestamp expression.
|
||||
|
||||
Returns:
|
||||
PG: "now()"
|
||||
Oracle: "SYSTIMESTAMP"
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
def array_agg(self, expr: str) -> str:
|
||||
"""Aggregate values into an array.
|
||||
|
||||
Args:
|
||||
expr: Expression to aggregate.
|
||||
|
||||
Returns:
|
||||
PG: "array_agg(expr)"
|
||||
Oracle: "CAST(COLLECT(expr) AS ...)" or JSON_ARRAYAGG.
|
||||
"""
|
||||
...
|
||||
|
||||
# -- Retrieval query arms ----------------------------------------------
|
||||
# These build complete subquery arms for the UNION ALL retrieval query.
|
||||
# Each database has significantly different syntax for vector search and
|
||||
# full-text search, so these belong in the dialect rather than inline
|
||||
# conditionals in retrieval.py.
|
||||
|
||||
@abstractmethod
|
||||
def build_semantic_arm(
|
||||
self,
|
||||
*,
|
||||
table: str,
|
||||
cols: str,
|
||||
fact_type: str,
|
||||
embedding_param: str,
|
||||
bank_id_param: str,
|
||||
fetch_limit: int,
|
||||
tags_clause: str = "",
|
||||
groups_clause: str = "",
|
||||
extra_where: str = "",
|
||||
) -> str:
|
||||
"""Build a semantic (vector similarity) search subquery arm.
|
||||
|
||||
Returns a complete subquery suitable for UNION ALL that selects
|
||||
matching rows ordered by cosine similarity.
|
||||
|
||||
Args:
|
||||
table: Fully-qualified table name.
|
||||
cols: Column list expression.
|
||||
fact_type: Fact type literal (inlined, not parameterized).
|
||||
embedding_param: Parameter placeholder for query embedding.
|
||||
bank_id_param: Parameter placeholder for bank_id.
|
||||
fetch_limit: Max rows to fetch (over-fetched for HNSW approximation).
|
||||
tags_clause: Optional WHERE clause fragment for tag filtering.
|
||||
groups_clause: Optional WHERE clause fragment for tag group filtering.
|
||||
extra_where: Optional additional WHERE clause fragment (e.g. time range filter).
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
def build_bm25_arm(
|
||||
self,
|
||||
*,
|
||||
table: str,
|
||||
cols: str,
|
||||
fact_type: str,
|
||||
bank_id_param: str,
|
||||
limit_param: str,
|
||||
text_param: str,
|
||||
tags_clause: str = "",
|
||||
groups_clause: str = "",
|
||||
arm_index: int = 0,
|
||||
text_search_extension: str = "native",
|
||||
extra_where: str = "",
|
||||
) -> str:
|
||||
"""Build a BM25/full-text search subquery arm.
|
||||
|
||||
Returns a complete subquery suitable for UNION ALL that selects
|
||||
matching rows ordered by text relevance score.
|
||||
|
||||
Args:
|
||||
table: Fully-qualified table name.
|
||||
cols: Column list expression.
|
||||
fact_type: Fact type literal (inlined, not parameterized).
|
||||
bank_id_param: Parameter placeholder for bank_id.
|
||||
limit_param: Parameter placeholder for result limit.
|
||||
text_param: Parameter placeholder for the search text.
|
||||
tags_clause: Optional WHERE clause fragment for tag filtering.
|
||||
groups_clause: Optional WHERE clause fragment for tag group filtering.
|
||||
arm_index: Index of this arm in the UNION ALL (used by Oracle for
|
||||
unique SCORE labels).
|
||||
text_search_extension: Full-text search backend ("native", "vchord",
|
||||
"pg_textsearch"). Only relevant for PostgreSQL.
|
||||
extra_where: Optional additional WHERE clause fragment (e.g. time range filter).
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
def prepare_bm25_text(
|
||||
self,
|
||||
tokens: list[str],
|
||||
query_text: str,
|
||||
*,
|
||||
text_search_extension: str = "native",
|
||||
) -> str:
|
||||
"""Prepare the text parameter value for BM25 search.
|
||||
|
||||
Transforms tokens/query text into the format expected by the backend's
|
||||
full-text search engine.
|
||||
|
||||
Args:
|
||||
tokens: Tokenized query words.
|
||||
query_text: Original query text.
|
||||
text_search_extension: Full-text search backend variant.
|
||||
|
||||
Returns:
|
||||
Prepared text string to bind as the BM25 text parameter.
|
||||
"""
|
||||
...
|
||||
@@ -0,0 +1,317 @@
|
||||
"""Oracle 23ai SQL dialect implementation.
|
||||
|
||||
Provides Oracle-specific SQL fragments for parameter binding, JSON operators,
|
||||
vector distance (VECTOR_DISTANCE), full-text search (Oracle Text), and
|
||||
other non-portable patterns.
|
||||
"""
|
||||
|
||||
from .base import SQLDialect
|
||||
|
||||
|
||||
class OracleDialect(SQLDialect):
|
||||
"""SQL dialect for Oracle 23ai (python-oracledb)."""
|
||||
|
||||
# Characters that need escaping in Oracle Text CONTAINS queries.
|
||||
_ORACLE_TEXT_SPECIAL = frozenset("&|!{}()[]~*?%-$>")
|
||||
|
||||
# Oracle Text reserved words that must be escaped with curly braces
|
||||
# when used as plain search terms. Full list from Oracle Text docs:
|
||||
# ABOUT, AND, BT, BTG, BTI, BTP, EQUIV, FUZZY, HASPATH, INPATH,
|
||||
# MINUS, NEAR, NOT, NT, NTG, NTI, NTP, OR, PT, RT, SQE, SYN,
|
||||
# TR, TRSYN, TT, WITHIN.
|
||||
_ORACLE_TEXT_RESERVED = frozenset(
|
||||
{
|
||||
"about",
|
||||
"and",
|
||||
"bt",
|
||||
"btg",
|
||||
"bti",
|
||||
"btp",
|
||||
"equiv",
|
||||
"fuzzy",
|
||||
"haspath",
|
||||
"inpath",
|
||||
"minus",
|
||||
"near",
|
||||
"not",
|
||||
"nt",
|
||||
"ntg",
|
||||
"nti",
|
||||
"ntp",
|
||||
"or",
|
||||
"pt",
|
||||
"rt",
|
||||
"sqe",
|
||||
"syn",
|
||||
"tr",
|
||||
"trsyn",
|
||||
"tt",
|
||||
"within",
|
||||
}
|
||||
)
|
||||
|
||||
# -- Parameter binding -----------------------------------------------
|
||||
|
||||
def param(self, n: int) -> str:
|
||||
return f":{n}"
|
||||
|
||||
# -- Type casting ----------------------------------------------------
|
||||
|
||||
def cast(self, param: str, type_name: str) -> str:
|
||||
# Oracle uses standard CAST syntax
|
||||
oracle_type = self._map_type(type_name)
|
||||
return f"CAST({param} AS {oracle_type})"
|
||||
|
||||
@staticmethod
|
||||
def _map_type(pg_type: str) -> str:
|
||||
"""Map PostgreSQL type names to Oracle equivalents."""
|
||||
mapping = {
|
||||
"jsonb": "CLOB", # Oracle stores JSON in CLOB
|
||||
"json": "CLOB",
|
||||
"text": "VARCHAR2(4000)",
|
||||
"text[]": "CLOB", # JSON array
|
||||
"uuid": "RAW(16)",
|
||||
"uuid[]": "CLOB", # JSON array
|
||||
"varchar[]": "CLOB", # JSON array
|
||||
"float8": "BINARY_DOUBLE",
|
||||
"float8[]": "CLOB",
|
||||
"timestamptz": "TIMESTAMP WITH TIME ZONE",
|
||||
"timestamptz[]": "CLOB",
|
||||
"vector": "VECTOR",
|
||||
"vector[]": "CLOB",
|
||||
"integer": "NUMBER",
|
||||
"bigint": "NUMBER",
|
||||
"boolean": "NUMBER(1)",
|
||||
}
|
||||
return mapping.get(pg_type, pg_type.upper())
|
||||
|
||||
# -- Vector operations -----------------------------------------------
|
||||
|
||||
def vector_distance(self, col: str, param: str) -> str:
|
||||
return f"VECTOR_DISTANCE({col}, {param}, COSINE)"
|
||||
|
||||
def vector_similarity(self, col: str, param: str) -> str:
|
||||
return f"(1 - VECTOR_DISTANCE({col}, {param}, COSINE))"
|
||||
|
||||
# -- JSON operations -------------------------------------------------
|
||||
|
||||
def json_extract_text(self, col: str, key: str) -> str:
|
||||
return f"JSON_VALUE({col}, '$.{key}')"
|
||||
|
||||
def json_contains(self, col: str, param: str) -> str:
|
||||
return f"JSON_EXISTS({col}, '$?(@ == {param})')"
|
||||
|
||||
def json_merge(self, col: str, param: str) -> str:
|
||||
return f"JSON_MERGEPATCH({col}, {param})"
|
||||
|
||||
# -- Text search -----------------------------------------------------
|
||||
|
||||
def text_search_score(self, col: str, query_param: str, *, index_name: str | None = None) -> str:
|
||||
# Oracle Text: CONTAINS with SCORE
|
||||
return "SCORE(1)"
|
||||
|
||||
def text_search_order(self, col: str, query_param: str, *, index_name: str | None = None) -> str:
|
||||
return "SCORE(1) DESC"
|
||||
|
||||
# -- Fuzzy string matching -------------------------------------------
|
||||
|
||||
def similarity(self, col: str, param: str) -> str:
|
||||
return f"UTL_MATCH.EDIT_DISTANCE_SIMILARITY({col}, {param}) / 100.0"
|
||||
|
||||
# -- Upsert ----------------------------------------------------------
|
||||
|
||||
def upsert(
|
||||
self,
|
||||
table: str,
|
||||
columns: list[str],
|
||||
conflict_columns: list[str],
|
||||
update_columns: list[str],
|
||||
) -> str:
|
||||
col_list = ", ".join(columns)
|
||||
src_cols = ", ".join(f":{i + 1} AS {c}" for i, c in enumerate(columns))
|
||||
on_clause = " AND ".join(f"t.{c} = s.{c}" for c in conflict_columns)
|
||||
|
||||
if not update_columns:
|
||||
return (
|
||||
f"MERGE INTO {table} t "
|
||||
f"USING (SELECT {src_cols} FROM DUAL) s "
|
||||
f"ON ({on_clause}) "
|
||||
f"WHEN NOT MATCHED THEN INSERT ({col_list}) "
|
||||
f"VALUES ({', '.join(f's.{c}' for c in columns)})"
|
||||
)
|
||||
|
||||
updates = ", ".join(f"t.{c} = s.{c}" for c in update_columns)
|
||||
return (
|
||||
f"MERGE INTO {table} t "
|
||||
f"USING (SELECT {src_cols} FROM DUAL) s "
|
||||
f"ON ({on_clause}) "
|
||||
f"WHEN MATCHED THEN UPDATE SET {updates} "
|
||||
f"WHEN NOT MATCHED THEN INSERT ({col_list}) "
|
||||
f"VALUES ({', '.join(f's.{c}' for c in columns)})"
|
||||
)
|
||||
|
||||
# -- Bulk operations -------------------------------------------------
|
||||
|
||||
def bulk_unnest(self, param_types: list[tuple[str, str]]) -> str:
|
||||
# Oracle: use JSON_TABLE to expand a JSON array into rows
|
||||
# Caller passes a JSON array as the parameter
|
||||
columns = []
|
||||
for i, (param, sql_type) in enumerate(param_types):
|
||||
oracle_type = self._map_type(sql_type.rstrip("[]"))
|
||||
columns.append(f"c{i} {oracle_type} PATH '$[{i}]'")
|
||||
cols_spec = ", ".join(columns)
|
||||
# Using first param as the JSON array source
|
||||
first_param = param_types[0][0]
|
||||
return f"JSON_TABLE({first_param}, '$[*]' COLUMNS ({cols_spec}))"
|
||||
|
||||
# -- Pagination ------------------------------------------------------
|
||||
|
||||
def limit_offset(self, limit_param: str, offset_param: str) -> str:
|
||||
return f"OFFSET {offset_param} ROWS FETCH FIRST {limit_param} ROWS ONLY"
|
||||
|
||||
# -- RETURNING clause ------------------------------------------------
|
||||
|
||||
def returning(self, columns: list[str]) -> str:
|
||||
# Oracle RETURNING requires INTO clause with output bind variables.
|
||||
# The backend layer handles the output variable binding.
|
||||
return f"RETURNING {', '.join(columns)} INTO {', '.join(f':out_{c}' for c in columns)}"
|
||||
|
||||
# -- Pattern matching ------------------------------------------------
|
||||
|
||||
def ilike(self, col: str, param: str) -> str:
|
||||
return f"UPPER({col}) LIKE UPPER({param})"
|
||||
|
||||
# -- Array operations ------------------------------------------------
|
||||
|
||||
def array_any(self, param: str) -> str:
|
||||
# Oracle: expand JSON array to rows for IN clause
|
||||
return f"IN (SELECT value FROM JSON_TABLE({param}, '$[*]' COLUMNS (value PATH '$')))"
|
||||
|
||||
def array_all(self, param: str) -> str:
|
||||
return f"NOT IN (SELECT value FROM JSON_TABLE({param}, '$[*]' COLUMNS (value PATH '$')))"
|
||||
|
||||
def array_contains(self, col: str, param: str) -> str:
|
||||
# Oracle: check all elements of param array exist in col JSON array
|
||||
return (
|
||||
f"(SELECT COUNT(*) FROM JSON_TABLE({param}, '$[*]' COLUMNS (v PATH '$')) "
|
||||
f"WHERE JSON_EXISTS({col}, '$[*]?(@ == v)')) = "
|
||||
f"(SELECT COUNT(*) FROM JSON_TABLE({param}, '$[*]' COLUMNS (v PATH '$')))"
|
||||
)
|
||||
|
||||
# -- Locking ---------------------------------------------------------
|
||||
|
||||
def for_update_skip_locked(self) -> str:
|
||||
return "FOR UPDATE SKIP LOCKED"
|
||||
|
||||
def advisory_lock(self, id_param: str) -> str:
|
||||
# Oracle doesn't have advisory locks. Use SELECT FOR UPDATE NOWAIT on a lock row.
|
||||
return "SELECT 1 FROM dual FOR UPDATE NOWAIT"
|
||||
|
||||
# -- UUID generation -------------------------------------------------
|
||||
|
||||
def generate_uuid(self) -> str:
|
||||
return "SYS_GUID()"
|
||||
|
||||
# -- Misc ------------------------------------------------------------
|
||||
|
||||
def greatest(self, *args: str) -> str:
|
||||
return f"GREATEST({', '.join(args)})"
|
||||
|
||||
def current_timestamp(self) -> str:
|
||||
return "SYSTIMESTAMP"
|
||||
|
||||
def array_agg(self, expr: str) -> str:
|
||||
return f"JSON_ARRAYAGG({expr})"
|
||||
|
||||
# -- Retrieval query arms ----------------------------------------------
|
||||
|
||||
def build_semantic_arm(
|
||||
self,
|
||||
*,
|
||||
table: str,
|
||||
cols: str,
|
||||
fact_type: str,
|
||||
embedding_param: str,
|
||||
bank_id_param: str,
|
||||
fetch_limit: int,
|
||||
tags_clause: str = "",
|
||||
groups_clause: str = "",
|
||||
extra_where: str = "",
|
||||
) -> str:
|
||||
# Oracle 23ai: VECTOR_DISTANCE for cosine, FETCH FIRST for limiting.
|
||||
# Wrapped in a derived table to work within UNION ALL.
|
||||
return (
|
||||
f"SELECT * FROM (SELECT {cols},"
|
||||
f" 1 - VECTOR_DISTANCE(embedding, {embedding_param}, COSINE) AS similarity,"
|
||||
f" NULL AS bm25_score,"
|
||||
f" 'semantic' AS source"
|
||||
f" FROM {table}"
|
||||
f" WHERE bank_id = {bank_id_param}"
|
||||
f" AND fact_type = '{fact_type}'"
|
||||
f" AND embedding IS NOT NULL"
|
||||
f" AND (1 - VECTOR_DISTANCE(embedding, {embedding_param}, COSINE)) >= 0.3"
|
||||
f" {tags_clause}"
|
||||
f" {groups_clause}"
|
||||
f" {extra_where}"
|
||||
f" ORDER BY VECTOR_DISTANCE(embedding, {embedding_param}, COSINE)"
|
||||
f" FETCH FIRST {fetch_limit} ROWS ONLY) t"
|
||||
)
|
||||
|
||||
def build_bm25_arm(
|
||||
self,
|
||||
*,
|
||||
table: str,
|
||||
cols: str,
|
||||
fact_type: str,
|
||||
bank_id_param: str,
|
||||
limit_param: str,
|
||||
text_param: str,
|
||||
tags_clause: str = "",
|
||||
groups_clause: str = "",
|
||||
arm_index: int = 0,
|
||||
text_search_extension: str = "native",
|
||||
extra_where: str = "",
|
||||
) -> str:
|
||||
# Oracle Text: CONTAINS() / SCORE() with the CTXSYS.CONTEXT index.
|
||||
# Each arm gets a unique SCORE label (10 + arm_index) to avoid
|
||||
# conflicts within the UNION ALL.
|
||||
label = 10 + arm_index
|
||||
return (
|
||||
f"SELECT * FROM (SELECT {cols},"
|
||||
f" NULL AS similarity,"
|
||||
f" SCORE({label}) AS bm25_score,"
|
||||
f" 'bm25' AS source"
|
||||
f" FROM {table}"
|
||||
f" WHERE bank_id = {bank_id_param}"
|
||||
f" AND fact_type = '{fact_type}'"
|
||||
f" AND CONTAINS(text, {text_param}, {label}) > 0"
|
||||
f" {tags_clause}"
|
||||
f" {groups_clause}"
|
||||
f" {extra_where}"
|
||||
f" ORDER BY SCORE({label}) DESC"
|
||||
f" FETCH FIRST {limit_param} ROWS ONLY) t{arm_index}"
|
||||
)
|
||||
|
||||
def prepare_bm25_text(
|
||||
self,
|
||||
tokens: list[str],
|
||||
query_text: str,
|
||||
*,
|
||||
text_search_extension: str = "native",
|
||||
) -> str:
|
||||
# Oracle Text: filter tokens with special chars, escape reserved words
|
||||
# with curly braces (e.g. "about" → "{about}"), and join with OR.
|
||||
safe: list[str] = []
|
||||
for t in tokens:
|
||||
if any(c in self._ORACLE_TEXT_SPECIAL for c in t):
|
||||
continue
|
||||
if t.lower() in self._ORACLE_TEXT_RESERVED:
|
||||
safe.append(f"{{{t}}}")
|
||||
else:
|
||||
safe.append(t)
|
||||
if safe:
|
||||
return " OR ".join(safe)
|
||||
# All tokens were filtered out — escape the original query text as a
|
||||
# single term so we still attempt a search rather than erroring out.
|
||||
fallback = query_text.strip() or tokens[0]
|
||||
return f"{{{fallback}}}"
|
||||
@@ -0,0 +1,228 @@
|
||||
"""PostgreSQL SQL dialect implementation.
|
||||
|
||||
Provides PostgreSQL-specific SQL fragments for parameter binding, JSON operators,
|
||||
vector distance (pgvector), full-text search (VectorChord BM25 / tsvector),
|
||||
and other non-portable patterns.
|
||||
"""
|
||||
|
||||
from .base import SQLDialect
|
||||
|
||||
|
||||
class PostgreSQLDialect(SQLDialect):
|
||||
"""SQL dialect for PostgreSQL (asyncpg)."""
|
||||
|
||||
# -- Parameter binding -----------------------------------------------
|
||||
|
||||
def param(self, n: int) -> str:
|
||||
return f"${n}"
|
||||
|
||||
# -- Type casting ----------------------------------------------------
|
||||
|
||||
def cast(self, param: str, type_name: str) -> str:
|
||||
return f"{param}::{type_name}"
|
||||
|
||||
# -- Vector operations -----------------------------------------------
|
||||
|
||||
def vector_distance(self, col: str, param: str) -> str:
|
||||
return f"{col} <=> {param}::vector"
|
||||
|
||||
def vector_similarity(self, col: str, param: str) -> str:
|
||||
return f"1 - ({col} <=> {param}::vector)"
|
||||
|
||||
# -- JSON operations -------------------------------------------------
|
||||
|
||||
def json_extract_text(self, col: str, key: str) -> str:
|
||||
return f"{col} ->> '{key}'"
|
||||
|
||||
def json_contains(self, col: str, param: str) -> str:
|
||||
return f"{col} @> {param}::jsonb"
|
||||
|
||||
def json_merge(self, col: str, param: str) -> str:
|
||||
return f"{col} || {param}::jsonb"
|
||||
|
||||
# -- Text search -----------------------------------------------------
|
||||
|
||||
def text_search_score(self, col: str, query_param: str, *, index_name: str | None = None) -> str:
|
||||
if index_name:
|
||||
# VectorChord BM25
|
||||
return f"-({col} <@> to_bm25query({query_param}, '{index_name}'))"
|
||||
# Fallback to tsvector
|
||||
return f"ts_rank_cd({col}, to_tsquery({query_param}))"
|
||||
|
||||
def text_search_order(self, col: str, query_param: str, *, index_name: str | None = None) -> str:
|
||||
if index_name:
|
||||
# VectorChord BM25 — lower distance = better, so ASC
|
||||
return f"{col} <@> to_bm25query({query_param}, '{index_name}') ASC"
|
||||
return f"ts_rank_cd({col}, to_tsquery({query_param})) DESC"
|
||||
|
||||
# -- Fuzzy string matching -------------------------------------------
|
||||
|
||||
def similarity(self, col: str, param: str) -> str:
|
||||
return f"similarity({col}, {param})"
|
||||
|
||||
# -- Upsert ----------------------------------------------------------
|
||||
|
||||
def upsert(
|
||||
self,
|
||||
table: str,
|
||||
columns: list[str],
|
||||
conflict_columns: list[str],
|
||||
update_columns: list[str],
|
||||
) -> str:
|
||||
col_list = ", ".join(columns)
|
||||
placeholders = ", ".join(f"${i + 1}" for i in range(len(columns)))
|
||||
conflict = ", ".join(conflict_columns)
|
||||
|
||||
if not update_columns:
|
||||
return f"INSERT INTO {table} ({col_list}) VALUES ({placeholders}) ON CONFLICT ({conflict}) DO NOTHING"
|
||||
|
||||
updates = ", ".join(f"{c} = EXCLUDED.{c}" for c in update_columns)
|
||||
return (
|
||||
f"INSERT INTO {table} ({col_list}) VALUES ({placeholders}) ON CONFLICT ({conflict}) DO UPDATE SET {updates}"
|
||||
)
|
||||
|
||||
# -- Bulk operations -------------------------------------------------
|
||||
|
||||
def bulk_unnest(self, param_types: list[tuple[str, str]]) -> str:
|
||||
args = ", ".join(f"{p}::{t}" for p, t in param_types)
|
||||
return f"unnest({args})"
|
||||
|
||||
# -- Pagination ------------------------------------------------------
|
||||
|
||||
def limit_offset(self, limit_param: str, offset_param: str) -> str:
|
||||
return f"LIMIT {limit_param} OFFSET {offset_param}"
|
||||
|
||||
# -- RETURNING clause ------------------------------------------------
|
||||
|
||||
def returning(self, columns: list[str]) -> str:
|
||||
return f"RETURNING {', '.join(columns)}"
|
||||
|
||||
# -- Pattern matching ------------------------------------------------
|
||||
|
||||
def ilike(self, col: str, param: str) -> str:
|
||||
return f"{col} ILIKE {param}"
|
||||
|
||||
# -- Array operations ------------------------------------------------
|
||||
|
||||
def array_any(self, param: str) -> str:
|
||||
return f"= ANY({param})"
|
||||
|
||||
def array_all(self, param: str) -> str:
|
||||
return f"!= ALL({param})"
|
||||
|
||||
def array_contains(self, col: str, param: str) -> str:
|
||||
return f"{col} @> {param}::varchar[]"
|
||||
|
||||
# -- Locking ---------------------------------------------------------
|
||||
|
||||
def for_update_skip_locked(self) -> str:
|
||||
return "FOR UPDATE SKIP LOCKED"
|
||||
|
||||
def advisory_lock(self, id_param: str) -> str:
|
||||
return f"pg_try_advisory_lock({id_param})"
|
||||
|
||||
# -- UUID generation -------------------------------------------------
|
||||
|
||||
def generate_uuid(self) -> str:
|
||||
return "gen_random_uuid()"
|
||||
|
||||
# -- Misc ------------------------------------------------------------
|
||||
|
||||
def greatest(self, *args: str) -> str:
|
||||
return f"GREATEST({', '.join(args)})"
|
||||
|
||||
def current_timestamp(self) -> str:
|
||||
return "now()"
|
||||
|
||||
def array_agg(self, expr: str) -> str:
|
||||
return f"array_agg({expr})"
|
||||
|
||||
# -- Retrieval query arms ----------------------------------------------
|
||||
|
||||
def build_semantic_arm(
|
||||
self,
|
||||
*,
|
||||
table: str,
|
||||
cols: str,
|
||||
fact_type: str,
|
||||
embedding_param: str,
|
||||
bank_id_param: str,
|
||||
fetch_limit: int,
|
||||
tags_clause: str = "",
|
||||
groups_clause: str = "",
|
||||
extra_where: str = "",
|
||||
) -> str:
|
||||
return (
|
||||
f"(SELECT {cols},"
|
||||
f" 1 - (embedding <=> {embedding_param}::vector) AS similarity,"
|
||||
f" NULL::float AS bm25_score,"
|
||||
f" 'semantic' AS source"
|
||||
f" FROM {table}"
|
||||
f" WHERE bank_id = {bank_id_param}"
|
||||
f" AND fact_type = '{fact_type}'"
|
||||
f" AND embedding IS NOT NULL"
|
||||
f" AND (1 - (embedding <=> {embedding_param}::vector)) >= 0.3"
|
||||
f" {tags_clause}"
|
||||
f" {groups_clause}"
|
||||
f" {extra_where}"
|
||||
f" ORDER BY embedding <=> {embedding_param}::vector"
|
||||
f" LIMIT {fetch_limit})"
|
||||
)
|
||||
|
||||
def build_bm25_arm(
|
||||
self,
|
||||
*,
|
||||
table: str,
|
||||
cols: str,
|
||||
fact_type: str,
|
||||
bank_id_param: str,
|
||||
limit_param: str,
|
||||
text_param: str,
|
||||
tags_clause: str = "",
|
||||
groups_clause: str = "",
|
||||
arm_index: int = 0,
|
||||
text_search_extension: str = "native",
|
||||
extra_where: str = "",
|
||||
) -> str:
|
||||
if text_search_extension == "vchord":
|
||||
bm25_score_expr = (
|
||||
f"search_vector <&> to_bm25query('idx_memory_units_text_search', tokenize({text_param}, 'llmlingua2'))"
|
||||
)
|
||||
bm25_order_by = f"{bm25_score_expr} DESC"
|
||||
bm25_where_filter = ""
|
||||
elif text_search_extension == "pg_textsearch":
|
||||
bm25_score_expr = f"-({text_param} <@> to_bm25query({text_param}, 'idx_memory_units_text_search'))"
|
||||
bm25_order_by = f"text <@> to_bm25query({text_param}, 'idx_memory_units_text_search') ASC"
|
||||
bm25_where_filter = ""
|
||||
else: # native tsvector
|
||||
bm25_score_expr = f"ts_rank_cd(search_vector, to_tsquery('english', {text_param}))"
|
||||
bm25_order_by = f"{bm25_score_expr} DESC"
|
||||
bm25_where_filter = f"AND search_vector @@ to_tsquery('english', {text_param})"
|
||||
|
||||
return (
|
||||
f"(SELECT {cols},"
|
||||
f" NULL::float AS similarity,"
|
||||
f" {bm25_score_expr} AS bm25_score,"
|
||||
f" 'bm25' AS source"
|
||||
f" FROM {table}"
|
||||
f" WHERE bank_id = {bank_id_param}"
|
||||
f" AND fact_type = '{fact_type}'"
|
||||
f" {bm25_where_filter}"
|
||||
f" {tags_clause}"
|
||||
f" {groups_clause}"
|
||||
f" {extra_where}"
|
||||
f" ORDER BY {bm25_order_by}"
|
||||
f" LIMIT {limit_param})"
|
||||
)
|
||||
|
||||
def prepare_bm25_text(
|
||||
self,
|
||||
tokens: list[str],
|
||||
query_text: str,
|
||||
*,
|
||||
text_search_extension: str = "native",
|
||||
) -> str:
|
||||
if text_search_extension in ("vchord", "pg_textsearch"):
|
||||
return query_text
|
||||
# native tsvector: join tokens with OR operator
|
||||
return " | ".join(tokens)
|
||||
@@ -2,23 +2,15 @@
|
||||
|
||||
import logging
|
||||
from collections.abc import Callable
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import asyncpg
|
||||
from typing import Any
|
||||
|
||||
from ..db_utils import acquire_with_retry
|
||||
from ..schema import fq_table_explicit as fq_table
|
||||
from .base import FileStorage
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def fq_table(table: str, schema: str | None = None) -> str:
|
||||
"""Get fully-qualified table name with optional schema prefix."""
|
||||
if schema:
|
||||
return f'"{schema}".{table}'
|
||||
return table
|
||||
|
||||
|
||||
class PostgreSQLFileStorage(FileStorage):
|
||||
"""
|
||||
PostgreSQL BYTEA-based file storage.
|
||||
@@ -42,7 +34,7 @@ class PostgreSQLFileStorage(FileStorage):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
pool_getter: Callable[[], "asyncpg.Pool"],
|
||||
pool_getter: Callable[[], Any],
|
||||
schema: str | None = None,
|
||||
schema_getter: Callable[[], str] | None = None,
|
||||
):
|
||||
@@ -74,7 +66,7 @@ class PostgreSQLFileStorage(FileStorage):
|
||||
"""Store file in PostgreSQL."""
|
||||
pool = self._pool_getter()
|
||||
|
||||
async with pool.acquire() as conn:
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
await conn.execute(
|
||||
f"""
|
||||
INSERT INTO {fq_table("file_storage", self._schema)}
|
||||
@@ -94,7 +86,7 @@ class PostgreSQLFileStorage(FileStorage):
|
||||
"""Retrieve file from PostgreSQL."""
|
||||
pool = self._pool_getter()
|
||||
|
||||
async with pool.acquire() as conn:
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
row = await conn.fetchrow(
|
||||
f"""
|
||||
SELECT data FROM {fq_table("file_storage", self._schema)}
|
||||
@@ -112,7 +104,7 @@ class PostgreSQLFileStorage(FileStorage):
|
||||
"""Delete file from PostgreSQL."""
|
||||
pool = self._pool_getter()
|
||||
|
||||
async with pool.acquire() as conn:
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
result = await conn.execute(
|
||||
f"""
|
||||
DELETE FROM {fq_table("file_storage", self._schema)}
|
||||
@@ -129,7 +121,7 @@ class PostgreSQLFileStorage(FileStorage):
|
||||
"""Check if file exists in PostgreSQL."""
|
||||
pool = self._pool_getter()
|
||||
|
||||
async with pool.acquire() as conn:
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
row = await conn.fetchrow(
|
||||
f"""
|
||||
SELECT 1 FROM {fq_table("file_storage", self._schema)}
|
||||
|
||||
@@ -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)
|
||||
"""
|
||||
|
||||
@@ -10,19 +11,16 @@ import json
|
||||
import logging
|
||||
from abc import ABC, abstractmethod
|
||||
from collections.abc import Awaitable, Callable
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import asyncpg
|
||||
from typing import Any
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def fq_table(table: str, schema: str | None = None) -> str:
|
||||
"""Get fully-qualified table name with optional schema prefix."""
|
||||
if schema:
|
||||
return f'"{schema}".{table}'
|
||||
return table
|
||||
from .schema import fq_table_explicit
|
||||
|
||||
return fq_table_explicit(table, schema)
|
||||
|
||||
|
||||
class TaskBackend(ABC):
|
||||
@@ -125,6 +123,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 async-operation row has already been persisted (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.
|
||||
@@ -138,7 +163,7 @@ class BrokerTaskBackend(TaskBackend):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
pool_getter: Callable[[], "asyncpg.Pool"],
|
||||
pool_getter: Callable[[], Any],
|
||||
schema: str | None = None,
|
||||
schema_getter: Callable[[], str | None] | None = None,
|
||||
):
|
||||
@@ -192,34 +217,42 @@ class BrokerTaskBackend(TaskBackend):
|
||||
schema = self._schema_getter() if self._schema_getter else self._schema
|
||||
table = fq_table("async_operations", schema)
|
||||
|
||||
from .db_utils import acquire_with_retry
|
||||
|
||||
if operation_id:
|
||||
# Update existing operation with task payload
|
||||
await pool.execute(
|
||||
f"""
|
||||
UPDATE {table}
|
||||
SET task_payload = $1::jsonb, updated_at = now()
|
||||
WHERE operation_id = $2
|
||||
""",
|
||||
payload_json,
|
||||
operation_id,
|
||||
)
|
||||
logger.debug(f"Updated task payload for operation {operation_id}")
|
||||
# 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.
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
await conn.execute(
|
||||
f"""
|
||||
UPDATE {table}
|
||||
SET task_payload = $1::jsonb, updated_at = now()
|
||||
WHERE operation_id = $2 AND task_payload IS NULL
|
||||
""",
|
||||
payload_json,
|
||||
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
|
||||
import uuid
|
||||
|
||||
new_id = uuid.uuid4()
|
||||
await pool.execute(
|
||||
f"""
|
||||
INSERT INTO {table} (operation_id, bank_id, operation_type, status, task_payload)
|
||||
VALUES ($1, $2, $3, 'pending', $4::jsonb)
|
||||
""",
|
||||
new_id,
|
||||
bank_id,
|
||||
task_type,
|
||||
payload_json,
|
||||
)
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
await conn.execute(
|
||||
f"""
|
||||
INSERT INTO {table} (operation_id, bank_id, operation_type, status, task_payload)
|
||||
VALUES ($1, $2, $3, 'pending', $4::jsonb)
|
||||
""",
|
||||
new_id,
|
||||
bank_id,
|
||||
task_type,
|
||||
payload_json,
|
||||
)
|
||||
logger.debug(f"Created new operation {new_id} for task type {task_type}")
|
||||
|
||||
async def shutdown(self):
|
||||
@@ -240,6 +273,8 @@ class BrokerTaskBackend(TaskBackend):
|
||||
"""
|
||||
import asyncio
|
||||
|
||||
from .db_utils import acquire_with_retry
|
||||
|
||||
pool = self._pool_getter()
|
||||
schema = self._schema_getter() if self._schema_getter else self._schema
|
||||
table = fq_table("async_operations", schema)
|
||||
@@ -247,12 +282,13 @@ class BrokerTaskBackend(TaskBackend):
|
||||
start_time = asyncio.get_event_loop().time()
|
||||
while asyncio.get_event_loop().time() - start_time < timeout:
|
||||
# Check if there are any pending tasks with payloads
|
||||
count = await pool.fetchval(
|
||||
f"""
|
||||
SELECT COUNT(*) FROM {table}
|
||||
WHERE status = 'pending' AND task_payload IS NOT NULL
|
||||
"""
|
||||
)
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
count = await conn.fetchval(
|
||||
f"""
|
||||
SELECT COUNT(*) FROM {table}
|
||||
WHERE status = 'pending' AND task_payload IS NOT NULL
|
||||
"""
|
||||
)
|
||||
|
||||
if count == 0:
|
||||
return
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -45,7 +45,6 @@ _ALL_TOOLS: frozenset[str] = frozenset(
|
||||
"delete_directive",
|
||||
"list_memories",
|
||||
"get_memory",
|
||||
"delete_memory",
|
||||
"list_documents",
|
||||
"get_document",
|
||||
"delete_document",
|
||||
@@ -223,7 +222,6 @@ def register_mcp_tools(
|
||||
"delete_directive",
|
||||
"list_memories",
|
||||
"get_memory",
|
||||
"delete_memory",
|
||||
"list_documents",
|
||||
"get_document",
|
||||
"delete_document",
|
||||
@@ -292,9 +290,6 @@ def register_mcp_tools(
|
||||
if "get_memory" in tools_to_register:
|
||||
_register_get_memory(mcp, memory, config)
|
||||
|
||||
if "delete_memory" in tools_to_register:
|
||||
_register_delete_memory(mcp, memory, config)
|
||||
|
||||
# Document tools
|
||||
if "list_documents" in tools_to_register:
|
||||
_register_list_documents(mcp, memory, config)
|
||||
@@ -441,7 +436,6 @@ _AUDITABLE_MCP_TOOLS: frozenset[str] = frozenset(
|
||||
"refresh_mental_model",
|
||||
"create_directive",
|
||||
"delete_directive",
|
||||
"delete_memory",
|
||||
"delete_document",
|
||||
"cancel_operation",
|
||||
}
|
||||
@@ -2163,74 +2157,6 @@ def _register_get_memory(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsCon
|
||||
return {"error": str(e)}
|
||||
|
||||
|
||||
def _register_delete_memory(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig) -> None:
|
||||
"""Register the delete_memory tool."""
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
async def delete_memory(
|
||||
memory_id: str,
|
||||
bank_id: str | None = None,
|
||||
) -> str:
|
||||
"""
|
||||
Delete a specific memory by ID.
|
||||
|
||||
Permanently removes a memory unit and its associated data.
|
||||
|
||||
Args:
|
||||
memory_id: The ID of the memory to delete
|
||||
bank_id: Optional bank (accepted for consistency, not used in deletion).
|
||||
"""
|
||||
try:
|
||||
target_bank = bank_id or config.bank_id_resolver()
|
||||
if target_bank is None:
|
||||
return '{"error": "No bank_id configured"}'
|
||||
|
||||
result = await memory.delete_memory_unit(
|
||||
unit_id=memory_id,
|
||||
request_context=_get_request_context(config),
|
||||
)
|
||||
return json.dumps({"status": "deleted", "memory_id": memory_id, **result}, default=str)
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return json.dumps({"error": str(e)})
|
||||
except Exception as e:
|
||||
logger.error(f"Error deleting memory: {e}", exc_info=True)
|
||||
return f'{{"error": "{e}"}}'
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
async def delete_memory(
|
||||
memory_id: str,
|
||||
) -> dict:
|
||||
"""
|
||||
Delete a specific memory by ID.
|
||||
|
||||
Permanently removes a memory unit and its associated data.
|
||||
|
||||
Args:
|
||||
memory_id: The ID of the memory to delete
|
||||
"""
|
||||
try:
|
||||
target_bank = config.bank_id_resolver()
|
||||
if target_bank is None:
|
||||
return {"error": "No bank_id configured"}
|
||||
|
||||
result = await memory.delete_memory_unit(
|
||||
unit_id=memory_id,
|
||||
request_context=_get_request_context(config),
|
||||
)
|
||||
return {"status": "deleted", "memory_id": memory_id, **result}
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return {"error": str(e)}
|
||||
except Exception as e:
|
||||
logger.error(f"Error deleting memory: {e}", exc_info=True)
|
||||
return {"error": str(e)}
|
||||
|
||||
|
||||
# =========================================================================
|
||||
# DOCUMENT TOOLS
|
||||
# =========================================================================
|
||||
@@ -2854,6 +2780,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."""
|
||||
|
||||
@@ -2863,16 +2827,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:
|
||||
@@ -2880,14 +2865,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:
|
||||
@@ -2900,29 +2887,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:
|
||||
|
||||
@@ -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 = [
|
||||
|
||||
@@ -0,0 +1,636 @@
|
||||
"""
|
||||
Oracle 23ai database migrations.
|
||||
|
||||
Uses idempotent DDL (CREATE TABLE IF NOT EXISTS) so migrations can safely
|
||||
run multiple times. Oracle 23ai natively supports IF NOT EXISTS for DDL.
|
||||
|
||||
Tables mirror the PostgreSQL schema defined in alembic/versions/ but use
|
||||
Oracle-native types:
|
||||
- UUID → RAW(16) with DEFAULT SYS_GUID()
|
||||
- TEXT/VARCHAR → VARCHAR2 / CLOB
|
||||
- JSONB → CLOB (with IS JSON CHECK)
|
||||
- BOOLEAN → NUMBER(1)
|
||||
- FLOAT → BINARY_DOUBLE
|
||||
- TIMESTAMP WITH TIME ZONE → TIMESTAMP WITH TIME ZONE
|
||||
- VARCHAR[] → CLOB (JSON array stored as string)
|
||||
- BYTEA → BLOB
|
||||
- vector(384) → VECTOR(384, FLOAT32) (Oracle 23ai native)
|
||||
"""
|
||||
|
||||
import logging
|
||||
from urllib.parse import urlparse
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# DDL statements — executed in dependency order
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_DDL_TABLES = [
|
||||
# -----------------------------------------------------------------------
|
||||
# 1. BANKS
|
||||
# -----------------------------------------------------------------------
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS banks (
|
||||
bank_id VARCHAR2(256) NOT NULL,
|
||||
internal_id RAW(16) DEFAULT SYS_GUID() NOT NULL,
|
||||
name VARCHAR2(512),
|
||||
disposition CLOB DEFAULT '{"skepticism":3,"literalism":3,"empathy":3}' NOT NULL
|
||||
CONSTRAINT banks_disposition_json CHECK (disposition IS JSON),
|
||||
mission CLOB,
|
||||
personality CLOB DEFAULT '{}' NOT NULL
|
||||
CONSTRAINT banks_personality_json CHECK (personality IS JSON),
|
||||
config CLOB DEFAULT '{}' NOT NULL
|
||||
CONSTRAINT banks_config_json CHECK (config IS JSON),
|
||||
created_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,
|
||||
updated_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,
|
||||
CONSTRAINT pk_banks PRIMARY KEY (bank_id),
|
||||
CONSTRAINT banks_internal_id_unique UNIQUE (internal_id)
|
||||
)
|
||||
""",
|
||||
# -----------------------------------------------------------------------
|
||||
# 2. DOCUMENTS
|
||||
# -----------------------------------------------------------------------
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS documents (
|
||||
id VARCHAR2(512) NOT NULL,
|
||||
bank_id VARCHAR2(256) NOT NULL,
|
||||
original_text CLOB,
|
||||
content_hash VARCHAR2(128),
|
||||
metadata CLOB DEFAULT '{}' NOT NULL
|
||||
CONSTRAINT docs_metadata_json CHECK (metadata IS JSON),
|
||||
retain_params CLOB CONSTRAINT docs_retain_params_json CHECK (retain_params IS JSON OR retain_params IS NULL),
|
||||
file_storage_key VARCHAR2(512),
|
||||
file_original_name VARCHAR2(512),
|
||||
file_content_type VARCHAR2(256),
|
||||
tags CLOB DEFAULT '[]' NOT NULL,
|
||||
created_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,
|
||||
updated_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,
|
||||
CONSTRAINT pk_documents PRIMARY KEY (id, bank_id),
|
||||
CONSTRAINT fk_documents_bank FOREIGN KEY (bank_id) REFERENCES banks(bank_id) ON DELETE CASCADE
|
||||
)
|
||||
""",
|
||||
# -----------------------------------------------------------------------
|
||||
# 3. CHUNKS
|
||||
# -----------------------------------------------------------------------
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS chunks (
|
||||
chunk_id VARCHAR2(512) NOT NULL,
|
||||
document_id VARCHAR2(512) NOT NULL,
|
||||
bank_id VARCHAR2(256) NOT NULL,
|
||||
chunk_index NUMBER(10) NOT NULL,
|
||||
chunk_text CLOB NOT NULL,
|
||||
content_hash VARCHAR2(128),
|
||||
created_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,
|
||||
CONSTRAINT pk_chunks PRIMARY KEY (chunk_id),
|
||||
CONSTRAINT fk_chunks_document FOREIGN KEY (document_id, bank_id)
|
||||
REFERENCES documents(id, bank_id) ON DELETE CASCADE
|
||||
)
|
||||
""",
|
||||
# -----------------------------------------------------------------------
|
||||
# 4. MEMORY_UNITS
|
||||
# -----------------------------------------------------------------------
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS memory_units (
|
||||
id RAW(16) DEFAULT SYS_GUID() NOT NULL,
|
||||
bank_id VARCHAR2(256) NOT NULL,
|
||||
document_id VARCHAR2(512),
|
||||
chunk_id VARCHAR2(512),
|
||||
text CLOB NOT NULL,
|
||||
embedding VECTOR(384, FLOAT32),
|
||||
context CLOB,
|
||||
event_date TIMESTAMP WITH TIME ZONE NOT NULL,
|
||||
occurred_start TIMESTAMP WITH TIME ZONE,
|
||||
occurred_end TIMESTAMP WITH TIME ZONE,
|
||||
mentioned_at TIMESTAMP WITH TIME ZONE,
|
||||
fact_type VARCHAR2(64) DEFAULT 'world' NOT NULL,
|
||||
confidence_score BINARY_DOUBLE,
|
||||
access_count NUMBER(10) DEFAULT 0 NOT NULL,
|
||||
consolidated_at TIMESTAMP WITH TIME ZONE,
|
||||
observation_scopes CLOB CONSTRAINT mu_obs_scopes_json CHECK (observation_scopes IS JSON OR observation_scopes IS NULL),
|
||||
tags CLOB DEFAULT '[]' NOT NULL,
|
||||
metadata CLOB DEFAULT '{}' NOT NULL
|
||||
CONSTRAINT mu_metadata_json CHECK (metadata IS JSON),
|
||||
proof_count NUMBER(10) DEFAULT 1,
|
||||
source_memory_ids CLOB,
|
||||
history CLOB DEFAULT '[]'
|
||||
CONSTRAINT mu_history_json CHECK (history IS JSON OR history IS NULL),
|
||||
text_signals CLOB,
|
||||
consolidation_failed_at TIMESTAMP WITH TIME ZONE,
|
||||
search_vector CLOB,
|
||||
created_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,
|
||||
updated_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,
|
||||
CONSTRAINT pk_memory_units PRIMARY KEY (id),
|
||||
CONSTRAINT fk_mu_document FOREIGN KEY (document_id, bank_id)
|
||||
REFERENCES documents(id, bank_id) ON DELETE CASCADE,
|
||||
CONSTRAINT fk_mu_chunk FOREIGN KEY (chunk_id)
|
||||
REFERENCES chunks(chunk_id) ON DELETE SET NULL,
|
||||
CONSTRAINT chk_mu_fact_type CHECK (fact_type IN ('world', 'experience', 'observation')),
|
||||
CONSTRAINT chk_mu_confidence CHECK (
|
||||
confidence_score IS NULL
|
||||
OR (confidence_score >= 0.0 AND confidence_score <= 1.0)
|
||||
)
|
||||
)
|
||||
PARTITION BY LIST (bank_id) AUTOMATIC
|
||||
(PARTITION p_default VALUES ('__default__'))
|
||||
""",
|
||||
# -----------------------------------------------------------------------
|
||||
# 5. ENTITIES
|
||||
# -----------------------------------------------------------------------
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS entities (
|
||||
id RAW(16) DEFAULT SYS_GUID() NOT NULL,
|
||||
bank_id VARCHAR2(256) NOT NULL,
|
||||
canonical_name VARCHAR2(512) NOT NULL,
|
||||
metadata CLOB DEFAULT '{}' NOT NULL
|
||||
CONSTRAINT ent_metadata_json CHECK (metadata IS JSON),
|
||||
first_seen TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,
|
||||
last_seen TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,
|
||||
mention_count NUMBER(10) DEFAULT 1 NOT NULL,
|
||||
CONSTRAINT pk_entities PRIMARY KEY (id)
|
||||
)
|
||||
""",
|
||||
# -----------------------------------------------------------------------
|
||||
# 6. UNIT_ENTITIES (junction)
|
||||
# -----------------------------------------------------------------------
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS unit_entities (
|
||||
unit_id RAW(16) NOT NULL,
|
||||
entity_id RAW(16) NOT NULL,
|
||||
CONSTRAINT pk_unit_entities PRIMARY KEY (unit_id, entity_id),
|
||||
CONSTRAINT fk_ue_unit FOREIGN KEY (unit_id) REFERENCES memory_units(id) ON DELETE CASCADE,
|
||||
CONSTRAINT fk_ue_entity FOREIGN KEY (entity_id) REFERENCES entities(id) ON DELETE CASCADE
|
||||
)
|
||||
""",
|
||||
# -----------------------------------------------------------------------
|
||||
# 7. ENTITY_COOCCURRENCES
|
||||
# -----------------------------------------------------------------------
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS entity_cooccurrences (
|
||||
entity_id_1 RAW(16) NOT NULL,
|
||||
entity_id_2 RAW(16) NOT NULL,
|
||||
cooccurrence_count NUMBER(10) DEFAULT 1 NOT NULL,
|
||||
last_cooccurred TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,
|
||||
CONSTRAINT pk_entity_cooccurrences PRIMARY KEY (entity_id_1, entity_id_2),
|
||||
CONSTRAINT fk_ec_entity1 FOREIGN KEY (entity_id_1) REFERENCES entities(id) ON DELETE CASCADE,
|
||||
CONSTRAINT fk_ec_entity2 FOREIGN KEY (entity_id_2) REFERENCES entities(id) ON DELETE CASCADE
|
||||
)
|
||||
""",
|
||||
# -----------------------------------------------------------------------
|
||||
# 8. MEMORY_LINKS
|
||||
# -----------------------------------------------------------------------
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS memory_links (
|
||||
from_unit_id RAW(16) NOT NULL,
|
||||
to_unit_id RAW(16) NOT NULL,
|
||||
link_type VARCHAR2(64) NOT NULL,
|
||||
entity_id RAW(16),
|
||||
bank_id VARCHAR2(256),
|
||||
weight BINARY_DOUBLE DEFAULT 1.0 NOT NULL,
|
||||
source_memory_ids CLOB,
|
||||
created_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,
|
||||
CONSTRAINT fk_ml_from FOREIGN KEY (from_unit_id) REFERENCES memory_units(id) ON DELETE CASCADE,
|
||||
CONSTRAINT fk_ml_to FOREIGN KEY (to_unit_id) REFERENCES memory_units(id) ON DELETE CASCADE,
|
||||
CONSTRAINT fk_ml_entity FOREIGN KEY (entity_id) REFERENCES entities(id) ON DELETE CASCADE,
|
||||
CONSTRAINT chk_ml_link_type CHECK (
|
||||
link_type IN ('temporal', 'semantic', 'entity', 'causes', 'caused_by', 'enables', 'prevents')
|
||||
),
|
||||
CONSTRAINT chk_ml_weight CHECK (weight >= 0.0 AND weight <= 1.0)
|
||||
)
|
||||
""",
|
||||
# -----------------------------------------------------------------------
|
||||
# 9. MENTAL_MODELS
|
||||
# -----------------------------------------------------------------------
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS mental_models (
|
||||
id VARCHAR2(256) NOT NULL,
|
||||
bank_id VARCHAR2(256) NOT NULL,
|
||||
subtype VARCHAR2(32) NOT NULL,
|
||||
name VARCHAR2(256) NOT NULL,
|
||||
description CLOB NOT NULL,
|
||||
source_query CLOB,
|
||||
content CLOB,
|
||||
embedding VECTOR(384, FLOAT32),
|
||||
entity_id RAW(16),
|
||||
observations CLOB DEFAULT '{"observations":[]}' NOT NULL
|
||||
CONSTRAINT mm_obs_json CHECK (observations IS JSON),
|
||||
links CLOB,
|
||||
tags CLOB DEFAULT '[]' NOT NULL,
|
||||
max_tokens NUMBER(10) DEFAULT 2048 NOT NULL,
|
||||
"trigger" CLOB DEFAULT '{"refresh_after_consolidation":false}' NOT NULL
|
||||
CONSTRAINT mm_trigger_json CHECK ("trigger" IS JSON),
|
||||
structured_content CLOB CONSTRAINT mm_sc_json CHECK (structured_content IS JSON OR structured_content IS NULL),
|
||||
last_refreshed_source_query CLOB,
|
||||
reflect_response CLOB CONSTRAINT mm_reflect_resp_json CHECK (reflect_response IS JSON OR reflect_response IS NULL),
|
||||
history CLOB DEFAULT '[]' NOT NULL
|
||||
CONSTRAINT mm_history_json CHECK (history IS JSON),
|
||||
last_refreshed_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,
|
||||
last_updated TIMESTAMP WITH TIME ZONE,
|
||||
created_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,
|
||||
CONSTRAINT pk_mental_models PRIMARY KEY (id, bank_id),
|
||||
CONSTRAINT fk_mm_bank FOREIGN KEY (bank_id) REFERENCES banks(bank_id) ON DELETE CASCADE,
|
||||
CONSTRAINT fk_mm_entity FOREIGN KEY (entity_id) REFERENCES entities(id) ON DELETE SET NULL,
|
||||
CONSTRAINT chk_mm_subtype CHECK (subtype IN ('structural', 'emergent', 'pinned', 'learned'))
|
||||
)
|
||||
""",
|
||||
# -----------------------------------------------------------------------
|
||||
# 10. DIRECTIVES
|
||||
# -----------------------------------------------------------------------
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS directives (
|
||||
id RAW(16) DEFAULT SYS_GUID() NOT NULL,
|
||||
bank_id VARCHAR2(256) NOT NULL,
|
||||
name VARCHAR2(256) NOT NULL,
|
||||
content CLOB NOT NULL,
|
||||
priority NUMBER(10) DEFAULT 0 NOT NULL,
|
||||
is_active NUMBER(1) DEFAULT 1 NOT NULL,
|
||||
tags CLOB DEFAULT '[]' NOT NULL,
|
||||
created_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,
|
||||
updated_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,
|
||||
CONSTRAINT pk_directives PRIMARY KEY (id),
|
||||
CONSTRAINT fk_dir_bank FOREIGN KEY (bank_id) REFERENCES banks(bank_id) ON DELETE CASCADE
|
||||
)
|
||||
""",
|
||||
# -----------------------------------------------------------------------
|
||||
# 11. ASYNC_OPERATIONS
|
||||
# -----------------------------------------------------------------------
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS async_operations (
|
||||
operation_id RAW(16) DEFAULT SYS_GUID() NOT NULL,
|
||||
bank_id VARCHAR2(256) NOT NULL,
|
||||
operation_type VARCHAR2(128) NOT NULL,
|
||||
status VARCHAR2(32) DEFAULT 'pending' NOT NULL,
|
||||
worker_id VARCHAR2(256),
|
||||
claimed_at TIMESTAMP WITH TIME ZONE,
|
||||
retry_count NUMBER(10) DEFAULT 0 NOT NULL,
|
||||
next_retry_at TIMESTAMP WITH TIME ZONE,
|
||||
task_payload CLOB CONSTRAINT ao_payload_json CHECK (task_payload IS JSON OR task_payload IS NULL),
|
||||
result_metadata CLOB DEFAULT '{}' NOT NULL
|
||||
CONSTRAINT ao_result_json CHECK (result_metadata IS JSON),
|
||||
error_message CLOB,
|
||||
created_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,
|
||||
updated_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,
|
||||
completed_at TIMESTAMP WITH TIME ZONE,
|
||||
CONSTRAINT pk_async_operations PRIMARY KEY (operation_id),
|
||||
CONSTRAINT fk_ao_bank FOREIGN KEY (bank_id) REFERENCES banks(bank_id) ON DELETE CASCADE,
|
||||
CONSTRAINT chk_ao_status CHECK (status IN ('pending', 'processing', 'completed', 'failed'))
|
||||
)
|
||||
""",
|
||||
# -----------------------------------------------------------------------
|
||||
# 11. WEBHOOKS
|
||||
# -----------------------------------------------------------------------
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS webhooks (
|
||||
id RAW(16) DEFAULT SYS_GUID() NOT NULL,
|
||||
bank_id VARCHAR2(256) NOT NULL,
|
||||
url VARCHAR2(2048) NOT NULL,
|
||||
secret VARCHAR2(512),
|
||||
event_types CLOB DEFAULT '[]' NOT NULL,
|
||||
http_config CLOB DEFAULT '{}' NOT NULL
|
||||
CONSTRAINT wh_http_config_json CHECK (http_config IS JSON),
|
||||
enabled NUMBER(1) DEFAULT 1 NOT NULL,
|
||||
created_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,
|
||||
updated_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,
|
||||
CONSTRAINT pk_webhooks PRIMARY KEY (id),
|
||||
CONSTRAINT fk_wh_bank FOREIGN KEY (bank_id) REFERENCES banks(bank_id) ON DELETE CASCADE
|
||||
)
|
||||
""",
|
||||
# -----------------------------------------------------------------------
|
||||
# 12. FILE_STORAGE
|
||||
# -----------------------------------------------------------------------
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS file_storage (
|
||||
storage_key VARCHAR2(512) NOT NULL,
|
||||
data BLOB NOT NULL,
|
||||
CONSTRAINT pk_file_storage PRIMARY KEY (storage_key)
|
||||
)
|
||||
""",
|
||||
# -----------------------------------------------------------------------
|
||||
# 13. AUDIT_LOG
|
||||
# -----------------------------------------------------------------------
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS audit_log (
|
||||
id RAW(16) DEFAULT SYS_GUID() NOT NULL,
|
||||
action VARCHAR2(128) NOT NULL,
|
||||
transport VARCHAR2(64) NOT NULL,
|
||||
bank_id VARCHAR2(256),
|
||||
started_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,
|
||||
ended_at TIMESTAMP WITH TIME ZONE,
|
||||
request CLOB CONSTRAINT al_request_json CHECK (request IS JSON OR request IS NULL),
|
||||
response CLOB CONSTRAINT al_response_json CHECK (response IS JSON OR response IS NULL),
|
||||
metadata CLOB DEFAULT '{}' NOT NULL
|
||||
CONSTRAINT al_metadata_json CHECK (metadata IS JSON),
|
||||
CONSTRAINT pk_audit_log PRIMARY KEY (id)
|
||||
)
|
||||
""",
|
||||
# -----------------------------------------------------------------------
|
||||
# 11. OBSERVATION_SOURCES — junction table replacing source_memory_ids
|
||||
# column. Enables standard SQL joins instead of dialect-specific array
|
||||
# operators (PG unnest/&&) or JSON_TABLE (Oracle).
|
||||
# -----------------------------------------------------------------------
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS observation_sources (
|
||||
observation_id RAW(16) NOT NULL,
|
||||
source_id RAW(16) NOT NULL,
|
||||
CONSTRAINT pk_observation_sources PRIMARY KEY (observation_id, source_id),
|
||||
CONSTRAINT fk_obs_src_observation FOREIGN KEY (observation_id)
|
||||
REFERENCES memory_units(id) ON DELETE CASCADE
|
||||
)
|
||||
""",
|
||||
]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Indexes — created with IF NOT EXISTS where Oracle 23ai supports it,
|
||||
# otherwise guarded by PL/SQL exception handler.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _idx(name: str, ddl: str) -> str:
|
||||
"""Wrap CREATE INDEX in a PL/SQL block that silently ignores ORA-00955 (name already used)."""
|
||||
return f"""
|
||||
BEGIN
|
||||
EXECUTE IMMEDIATE '{ddl.strip().replace("'", "''")}';
|
||||
EXCEPTION
|
||||
WHEN OTHERS THEN
|
||||
IF SQLCODE = -955 THEN NULL; -- index already exists
|
||||
ELSE RAISE;
|
||||
END IF;
|
||||
END;
|
||||
"""
|
||||
|
||||
|
||||
_DDL_INDEXES = [
|
||||
# --- documents ---
|
||||
_idx("idx_docs_bank_id", "CREATE INDEX idx_docs_bank_id ON documents(bank_id)"),
|
||||
_idx("idx_docs_content_hash", "CREATE INDEX idx_docs_content_hash ON documents(content_hash)"),
|
||||
# --- chunks ---
|
||||
_idx("idx_chunks_document_id", "CREATE INDEX idx_chunks_document_id ON chunks(document_id)"),
|
||||
_idx("idx_chunks_bank_id", "CREATE INDEX idx_chunks_bank_id ON chunks(bank_id)"),
|
||||
# --- memory_units ---
|
||||
_idx("idx_mu_bank_id", "CREATE INDEX idx_mu_bank_id ON memory_units(bank_id)"),
|
||||
_idx("idx_mu_document_id", "CREATE INDEX idx_mu_document_id ON memory_units(document_id)"),
|
||||
_idx("idx_mu_chunk_id", "CREATE INDEX idx_mu_chunk_id ON memory_units(chunk_id)"),
|
||||
_idx("idx_mu_event_date", "CREATE INDEX idx_mu_event_date ON memory_units(event_date DESC)"),
|
||||
_idx("idx_mu_bank_date", "CREATE INDEX idx_mu_bank_date ON memory_units(bank_id, event_date DESC)"),
|
||||
_idx("idx_mu_access_count", "CREATE INDEX idx_mu_access_count ON memory_units(access_count DESC)"),
|
||||
_idx("idx_mu_fact_type", "CREATE INDEX idx_mu_fact_type ON memory_units(fact_type)"),
|
||||
_idx("idx_mu_bank_fact_type", "CREATE INDEX idx_mu_bank_fact_type ON memory_units(bank_id, fact_type)"),
|
||||
_idx(
|
||||
"idx_mu_bank_type_date",
|
||||
"CREATE INDEX idx_mu_bank_type_date ON memory_units(bank_id, fact_type, event_date DESC)",
|
||||
),
|
||||
# --- entities ---
|
||||
_idx("idx_ent_bank_id", "CREATE INDEX idx_ent_bank_id ON entities(bank_id)"),
|
||||
_idx("idx_ent_canonical_name", "CREATE INDEX idx_ent_canonical_name ON entities(canonical_name)"),
|
||||
_idx("idx_ent_bank_name", "CREATE INDEX idx_ent_bank_name ON entities(bank_id, canonical_name)"),
|
||||
_idx(
|
||||
"idx_ent_bank_lower_name",
|
||||
"CREATE UNIQUE INDEX idx_ent_bank_lower_name ON entities(bank_id, LOWER(canonical_name))",
|
||||
),
|
||||
# --- unit_entities ---
|
||||
_idx("idx_ue_unit", "CREATE INDEX idx_ue_unit ON unit_entities(unit_id)"),
|
||||
_idx("idx_ue_entity", "CREATE INDEX idx_ue_entity ON unit_entities(entity_id)"),
|
||||
# --- entity_cooccurrences ---
|
||||
_idx("idx_ec_entity1", "CREATE INDEX idx_ec_entity1 ON entity_cooccurrences(entity_id_1)"),
|
||||
_idx("idx_ec_entity2", "CREATE INDEX idx_ec_entity2 ON entity_cooccurrences(entity_id_2)"),
|
||||
_idx("idx_ec_count", "CREATE INDEX idx_ec_count ON entity_cooccurrences(cooccurrence_count DESC)"),
|
||||
# --- memory_links ---
|
||||
# Unique constraint matching PG's idx_memory_links_unique — required for ON CONFLICT DO NOTHING
|
||||
# duplicate suppression. Oracle function-based unique index uses NVL (Oracle equivalent of COALESCE)
|
||||
# with the nil UUID as raw bytes to handle nullable entity_id.
|
||||
_idx(
|
||||
"idx_memory_links_unique",
|
||||
"CREATE UNIQUE INDEX idx_memory_links_unique ON memory_links("
|
||||
"from_unit_id, to_unit_id, link_type, "
|
||||
"NVL(entity_id, HEXTORAW('00000000000000000000000000000000')))",
|
||||
),
|
||||
_idx("idx_ml_from_unit", "CREATE INDEX idx_ml_from_unit ON memory_links(from_unit_id)"),
|
||||
_idx("idx_ml_to_unit", "CREATE INDEX idx_ml_to_unit ON memory_links(to_unit_id)"),
|
||||
_idx("idx_ml_entity", "CREATE INDEX idx_ml_entity ON memory_links(entity_id)"),
|
||||
_idx("idx_ml_link_type", "CREATE INDEX idx_ml_link_type ON memory_links(link_type)"),
|
||||
_idx("idx_ml_bank_id", "CREATE INDEX idx_ml_bank_id ON memory_links(bank_id)"),
|
||||
# --- directives ---
|
||||
_idx("idx_dir_bank_id", "CREATE INDEX idx_dir_bank_id ON directives(bank_id)"),
|
||||
_idx("idx_dir_bank_active", "CREATE INDEX idx_dir_bank_active ON directives(bank_id, is_active)"),
|
||||
# --- mental_models ---
|
||||
_idx("idx_mm_bank_id", "CREATE INDEX idx_mm_bank_id ON mental_models(bank_id)"),
|
||||
_idx("idx_mm_subtype", "CREATE INDEX idx_mm_subtype ON mental_models(bank_id, subtype)"),
|
||||
_idx("idx_mm_entity_id", "CREATE INDEX idx_mm_entity_id ON mental_models(entity_id)"),
|
||||
# --- async_operations ---
|
||||
_idx("idx_ao_bank_id", "CREATE INDEX idx_ao_bank_id ON async_operations(bank_id)"),
|
||||
_idx("idx_ao_status", "CREATE INDEX idx_ao_status ON async_operations(status)"),
|
||||
_idx("idx_ao_bank_status", "CREATE INDEX idx_ao_bank_status ON async_operations(bank_id, status)"),
|
||||
_idx("idx_ao_status_retry", "CREATE INDEX idx_ao_status_retry ON async_operations(status, next_retry_at)"),
|
||||
# --- webhooks ---
|
||||
_idx("idx_wh_bank_id", "CREATE INDEX idx_wh_bank_id ON webhooks(bank_id)"),
|
||||
# --- audit_log ---
|
||||
_idx("idx_al_action_started", "CREATE INDEX idx_al_action_started ON audit_log(action, started_at DESC)"),
|
||||
_idx("idx_al_bank_started", "CREATE INDEX idx_al_bank_started ON audit_log(bank_id, started_at DESC)"),
|
||||
_idx("idx_al_started", "CREATE INDEX idx_al_started ON audit_log(started_at DESC)"),
|
||||
# --- observation_sources ---
|
||||
_idx(
|
||||
"idx_obs_sources_source_id",
|
||||
"CREATE INDEX idx_obs_sources_source_id ON observation_sources(source_id, observation_id)",
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Vector and text indexes (Oracle 23ai specific)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_DDL_VECTOR_INDEX = _idx(
|
||||
"idx_mu_embedding_hnsw",
|
||||
"CREATE VECTOR INDEX idx_mu_embedding_hnsw ON memory_units(embedding) "
|
||||
"ORGANIZATION NEIGHBOR PARTITIONS "
|
||||
"DISTANCE COSINE "
|
||||
"WITH TARGET ACCURACY 95",
|
||||
)
|
||||
|
||||
_DDL_TEXT_INDEX = """
|
||||
BEGIN
|
||||
EXECUTE IMMEDIATE '
|
||||
CREATE INDEX idx_mu_content_text ON memory_units(text)
|
||||
INDEXTYPE IS CTXSYS.CONTEXT
|
||||
PARAMETERS (''SYNC (ON COMMIT)'')
|
||||
';
|
||||
EXCEPTION
|
||||
WHEN OTHERS THEN
|
||||
IF SQLCODE = -955 THEN NULL;
|
||||
ELSE RAISE;
|
||||
END IF;
|
||||
END;
|
||||
"""
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Public API
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def run_oracle_migrations(dsn: str, *, schema: str | None = None) -> None:
|
||||
"""Run Oracle schema migrations.
|
||||
|
||||
Creates all tables, indexes, and constraints using idempotent DDL.
|
||||
Safe to call multiple times.
|
||||
|
||||
Args:
|
||||
dsn: Oracle connection string (oracle://user:pass@host:port/service)
|
||||
schema: Target schema (Oracle user). None uses the connecting user's default.
|
||||
"""
|
||||
try:
|
||||
import oracledb # type: ignore[import-not-found]
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"python-oracledb is required for Oracle migrations. Install with: pip install oracledb"
|
||||
) from None
|
||||
|
||||
oracledb.defaults.fetch_lobs = False
|
||||
|
||||
# Parse URL-format DSN
|
||||
parsed = urlparse(dsn)
|
||||
connect_kwargs: dict = {}
|
||||
if parsed.scheme in ("oracle", "oracle+oracledb"):
|
||||
connect_kwargs["user"] = parsed.username
|
||||
connect_kwargs["password"] = parsed.password
|
||||
host = parsed.hostname or "localhost"
|
||||
port = parsed.port or 1521
|
||||
service = parsed.path.lstrip("/") if parsed.path else "FREEPDB1"
|
||||
connect_kwargs["dsn"] = f"{host}:{port}/{service}"
|
||||
else:
|
||||
connect_kwargs["dsn"] = dsn
|
||||
|
||||
logger.info("Running Oracle schema migrations (dsn=%s, schema=%s)", connect_kwargs.get("dsn", dsn), schema)
|
||||
|
||||
conn = oracledb.connect(**connect_kwargs)
|
||||
cursor = conn.cursor()
|
||||
|
||||
try:
|
||||
# Wait up to 30s for DDL locks instead of failing immediately (ORA-00054)
|
||||
cursor.execute("ALTER SESSION SET DDL_LOCK_TIMEOUT = 30")
|
||||
|
||||
# Set schema if specified
|
||||
if schema:
|
||||
cursor.execute(f'ALTER SESSION SET CURRENT_SCHEMA = "{schema}"')
|
||||
|
||||
# Create tables
|
||||
for i, ddl in enumerate(_DDL_TABLES):
|
||||
try:
|
||||
cursor.execute(ddl.strip())
|
||||
conn.commit()
|
||||
except oracledb.DatabaseError as e:
|
||||
err = e.args[0]
|
||||
if hasattr(err, "code") and err.code == 955:
|
||||
# ORA-00955: name is already used by an existing object
|
||||
pass
|
||||
else:
|
||||
logger.error("Failed to create table (statement %d): %s", i, e)
|
||||
raise
|
||||
|
||||
# Convert memory_units to automatic list partitioning on bank_id.
|
||||
# New installs get this from CREATE TABLE; this handles existing installs.
|
||||
# Oracle 12.2+ supports online conversion via ALTER TABLE MODIFY.
|
||||
#
|
||||
# IMPORTANT: ALTER TABLE MODIFY PARTITION invalidates CTXSYS.CONTEXT
|
||||
# domain indexes (ORA-29861). We drop the text index before conversion
|
||||
# and recreate it afterward. The text index creation below handles both
|
||||
# fresh installs and this post-conversion recreation.
|
||||
try:
|
||||
# Drop text index first if it exists — it will be invalidated by partitioning.
|
||||
try:
|
||||
cursor.execute("DROP INDEX idx_mu_content_text FORCE")
|
||||
conn.commit()
|
||||
logger.debug("Dropped text index before partitioning conversion")
|
||||
except oracledb.DatabaseError:
|
||||
pass # Index doesn't exist yet (fresh install)
|
||||
|
||||
cursor.execute("""
|
||||
ALTER TABLE memory_units MODIFY
|
||||
PARTITION BY LIST (bank_id) AUTOMATIC
|
||||
(PARTITION p_default VALUES ('__default__'))
|
||||
""")
|
||||
conn.commit()
|
||||
logger.info("memory_units partitioned by bank_id (automatic list)")
|
||||
except oracledb.DatabaseError as e:
|
||||
err = e.args[0]
|
||||
# ORA-14504: table is already partitioned — safe to ignore
|
||||
if hasattr(err, "code") and err.code == 14504:
|
||||
logger.debug("memory_units already partitioned")
|
||||
else:
|
||||
logger.debug("Partitioning memory_units skipped: %s", e)
|
||||
|
||||
# Deduplicate memory_links before creating unique index.
|
||||
# Earlier versions lacked a unique constraint, so duplicate rows may exist.
|
||||
try:
|
||||
cursor.execute("""
|
||||
DELETE FROM memory_links WHERE ROWID IN (
|
||||
SELECT rid FROM (
|
||||
SELECT ROWID AS rid,
|
||||
ROW_NUMBER() OVER (
|
||||
PARTITION BY from_unit_id, to_unit_id, link_type,
|
||||
NVL(entity_id, HEXTORAW('00000000000000000000000000000000'))
|
||||
ORDER BY created_at
|
||||
) AS rn
|
||||
FROM memory_links
|
||||
) WHERE rn > 1
|
||||
)
|
||||
""")
|
||||
if cursor.rowcount > 0:
|
||||
logger.info("Deduplicated %d memory_links rows before unique index creation", cursor.rowcount)
|
||||
conn.commit()
|
||||
except oracledb.DatabaseError as e:
|
||||
logger.debug("memory_links dedup skipped (table may not exist yet): %s", e)
|
||||
|
||||
# Create B-tree indexes
|
||||
for idx_ddl in _DDL_INDEXES:
|
||||
try:
|
||||
cursor.execute(idx_ddl.strip())
|
||||
conn.commit()
|
||||
except oracledb.DatabaseError as e:
|
||||
logger.debug("Index creation (may already exist): %s", e)
|
||||
|
||||
# Create vector index
|
||||
try:
|
||||
cursor.execute(_DDL_VECTOR_INDEX.strip())
|
||||
conn.commit()
|
||||
except oracledb.DatabaseError as e:
|
||||
logger.debug("Vector index creation (may already exist or VECTOR not supported): %s", e)
|
||||
|
||||
# Create Oracle Text index
|
||||
try:
|
||||
cursor.execute(_DDL_TEXT_INDEX.strip())
|
||||
conn.commit()
|
||||
except oracledb.DatabaseError as e:
|
||||
logger.debug("Text index creation (may already exist): %s", e)
|
||||
|
||||
# Backfill observation_sources from source_memory_ids CLOB (JSON array).
|
||||
# Uses MERGE to be idempotent — safe to run multiple times.
|
||||
try:
|
||||
cursor.execute("""
|
||||
MERGE INTO observation_sources tgt
|
||||
USING (
|
||||
SELECT mu.id AS observation_id,
|
||||
HEXTORAW(jt.source_id) AS source_id
|
||||
FROM memory_units mu,
|
||||
JSON_TABLE(mu.source_memory_ids, '$[*]'
|
||||
COLUMNS (source_id VARCHAR2(36) PATH '$')
|
||||
) jt
|
||||
WHERE mu.fact_type = 'observation'
|
||||
AND mu.source_memory_ids IS NOT NULL
|
||||
) src
|
||||
ON (tgt.observation_id = src.observation_id AND tgt.source_id = src.source_id)
|
||||
WHEN NOT MATCHED THEN
|
||||
INSERT (observation_id, source_id) VALUES (src.observation_id, src.source_id)
|
||||
""")
|
||||
conn.commit()
|
||||
logger.info("observation_sources backfill completed")
|
||||
except oracledb.DatabaseError as e:
|
||||
logger.debug("observation_sources backfill (may be empty or already done): %s", e)
|
||||
|
||||
logger.info("Oracle schema migrations completed successfully")
|
||||
|
||||
finally:
|
||||
cursor.close()
|
||||
conn.close()
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user