Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
d796735d7d | ||
|
|
1fc5fefde3 | ||
|
|
df4636d623 | ||
|
|
721fb5a645 | ||
|
|
ae2dbd9689 | ||
|
|
8c378b981a | ||
|
|
e2b19d3b38 | ||
|
|
210a40665d | ||
|
|
28dac7c7f8 | ||
|
|
f88f0a3b26 | ||
|
|
e4f8a157c2 | ||
|
|
ef90842f87 | ||
|
|
f68e2e2851 | ||
|
|
61b01cc040 | ||
|
|
bbcfe2f5ab | ||
|
|
d2bfa84bca | ||
|
|
8a64dc8db6 | ||
|
|
e7da7d0e4f | ||
|
|
f09ad9deac | ||
|
|
f27bd95382 | ||
|
|
7eabe5e168 | ||
|
|
33565a8236 | ||
|
|
51f05365a1 | ||
|
|
1e6cb15e99 | ||
|
|
bd6348aa08 | ||
|
|
836fd81e19 | ||
|
|
32b00cea4f | ||
|
|
21f9f46ca3 | ||
|
|
c7db770281 | ||
|
|
5fdb0e863f | ||
|
|
4a69a422a0 | ||
|
|
26472df166 | ||
|
|
5de793eec7 | ||
|
|
06200f1752 | ||
|
|
d4131f88fa | ||
|
|
94598fbd25 | ||
|
|
15ea23d5d6 | ||
|
|
720d42c576 | ||
|
|
8cbad0ef7a | ||
|
|
3d2c62ef09 | ||
|
|
56462a30cd | ||
|
|
067acf1ba5 | ||
|
|
4eaa2f3566 | ||
|
|
eeb938fc65 | ||
|
|
2344484f77 | ||
|
|
a01bb18bc4 | ||
|
|
b17f338e17 | ||
|
|
e210953d05 | ||
|
|
06b0f74a48 | ||
|
|
0560f6260d | ||
|
|
220851e6f4 | ||
|
|
5b360c83d2 | ||
|
|
1caf5ec9ee | ||
|
|
66dedb8d41 | ||
|
|
43b3efc494 | ||
|
|
00ac3d8834 | ||
|
|
2191654b1f | ||
|
|
dcaacbe407 | ||
|
|
32a4882a10 | ||
|
|
cd3a6a227b | ||
|
|
28308a14d6 | ||
|
|
fc71664b5f | ||
|
|
9a694f64b8 | ||
|
|
7bcf26097c | ||
|
|
1cdfb7c2e2 | ||
|
|
3e967add78 | ||
|
|
00ccf0b218 | ||
|
|
f7a60f898d | ||
|
|
7accac94b2 | ||
|
|
fa3501d448 | ||
|
|
f88b50a45e | ||
|
|
1b4ad7f435 | ||
|
|
d2504ac5ed | ||
|
|
d9d7021a49 | ||
|
|
e2baca8bfe | ||
|
|
576473b6aa | ||
|
|
99220d0527 | ||
|
|
8540c33236 | ||
|
|
5d05962db0 | ||
|
|
1d17dea2f1 | ||
|
|
928dc696e8 | ||
|
|
e4dd654ec5 | ||
|
|
0ad8c2d09c | ||
|
|
1e40cd22a6 | ||
|
|
1e5aa7de4d | ||
|
|
58fdac44f7 | ||
|
|
891c33b1d7 | ||
|
|
7ed57fdd85 | ||
|
|
e0a2ac63e7 | ||
|
|
75b95106ba | ||
|
|
d425e93cb4 | ||
|
|
cb6d1c469c | ||
|
|
4c058b4b98 | ||
|
|
aa8e5475c4 | ||
|
|
ad2cf72aab | ||
|
|
3f2a6ec9ce | ||
|
|
f17406fdf0 | ||
|
|
66423b85f5 | ||
|
|
abbf874d84 | ||
|
|
51d2fc5309 | ||
|
|
ea27ef95ec | ||
|
|
edf60e0f3c | ||
|
|
719e79a4d9 | ||
|
|
3857a30491 | ||
|
|
3d87ef5cee | ||
|
|
5c3d3274d7 |
+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
|
||||
# Supported providers: openai, groq, ollama, gemini, anthropic, lmstudio, vertexai, minimax
|
||||
HINDSIGHT_API_LLM_PROVIDER=openai
|
||||
HINDSIGHT_API_LLM_API_KEY=your-api-key-here
|
||||
HINDSIGHT_API_LLM_MODEL=gpt-4o-mini
|
||||
@@ -20,6 +20,11 @@ HINDSIGHT_API_LLM_BASE_URL=https://api.openai.com/v1
|
||||
# HINDSIGHT_API_LLM_VERTEXAI_REGION=us-central1
|
||||
# HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY=/path/to/service-account-key.json # Optional, uses ADC if not set
|
||||
|
||||
# Example: MiniMax configuration (204K context window)
|
||||
# HINDSIGHT_API_LLM_PROVIDER=minimax
|
||||
# HINDSIGHT_API_LLM_API_KEY=your-minimax-api-key
|
||||
# HINDSIGHT_API_LLM_MODEL=MiniMax-M2.5
|
||||
|
||||
# Example: LM Studio local configuration (Qwen 2.5 32B recommended)
|
||||
# HINDSIGHT_API_LLM_PROVIDER=lmstudio
|
||||
# HINDSIGHT_API_LLM_API_KEY=lmstudio
|
||||
|
||||
@@ -0,0 +1,6 @@
|
||||
version: 2
|
||||
updates:
|
||||
- package-ecosystem: "github-actions"
|
||||
directory: "/"
|
||||
schedule:
|
||||
interval: "weekly"
|
||||
@@ -21,20 +21,20 @@ jobs:
|
||||
build:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/setup-node@v4
|
||||
- uses: actions/checkout@v6
|
||||
- uses: actions/setup-node@v6
|
||||
with:
|
||||
node-version: 20
|
||||
cache: npm
|
||||
cache-dependency-path: package-lock.json
|
||||
- uses: astral-sh/setup-uv@v4
|
||||
- uses: astral-sh/setup-uv@v7
|
||||
- run: npm ci --workspace=hindsight-docs
|
||||
- run: uv run generate-llms-full
|
||||
- run: npm run build --workspace=hindsight-docs
|
||||
env:
|
||||
UMAMI_URL: https://analytics.hindsight.vectorize.io
|
||||
UMAMI_WEBSITE_ID: ${{ secrets.UMAMI_WEBSITE_ID }}
|
||||
- uses: actions/upload-pages-artifact@v3
|
||||
- uses: actions/upload-pages-artifact@v4
|
||||
with:
|
||||
path: hindsight-docs/build
|
||||
deploy:
|
||||
|
||||
@@ -13,15 +13,15 @@ jobs:
|
||||
id-token: write
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v6
|
||||
|
||||
- name: Install uv
|
||||
uses: astral-sh/setup-uv@v5
|
||||
uses: astral-sh/setup-uv@v7
|
||||
with:
|
||||
enable-cache: true
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v5
|
||||
uses: actions/setup-python@v6
|
||||
with:
|
||||
python-version-file: ".python-version"
|
||||
|
||||
@@ -30,12 +30,20 @@ jobs:
|
||||
working-directory: ./hindsight-clients/python
|
||||
run: uv build --out-dir dist
|
||||
|
||||
- name: Build hindsight-api-slim
|
||||
working-directory: ./hindsight-api-slim
|
||||
run: uv build --out-dir dist
|
||||
|
||||
- name: Build hindsight-api
|
||||
working-directory: ./hindsight-api
|
||||
run: uv build --out-dir dist
|
||||
|
||||
- name: Build hindsight-all
|
||||
working-directory: ./hindsight
|
||||
working-directory: ./hindsight-all
|
||||
run: uv build --out-dir dist
|
||||
|
||||
- name: Build hindsight-all-slim
|
||||
working-directory: ./hindsight-all-slim
|
||||
run: uv build --out-dir dist
|
||||
|
||||
- name: Build hindsight-litellm
|
||||
@@ -54,13 +62,27 @@ jobs:
|
||||
working-directory: ./hindsight-integrations/pydantic-ai
|
||||
run: uv build --out-dir dist
|
||||
|
||||
# Publish in order (client and api first, then hindsight-all which depends on them)
|
||||
- name: Build hindsight-hermes
|
||||
working-directory: ./hindsight-integrations/hermes
|
||||
run: uv build --out-dir dist
|
||||
|
||||
- name: Build hindsight-agno
|
||||
working-directory: ./hindsight-integrations/agno
|
||||
run: uv build --out-dir dist
|
||||
|
||||
# Publish in order (client and api-slim first, then api/all wrappers which depend on them)
|
||||
- name: Publish hindsight-client to PyPI
|
||||
uses: pypa/gh-action-pypi-publish@release/v1
|
||||
with:
|
||||
packages-dir: ./hindsight-clients/python/dist
|
||||
skip-existing: true
|
||||
|
||||
- name: Publish hindsight-api-slim to PyPI
|
||||
uses: pypa/gh-action-pypi-publish@release/v1
|
||||
with:
|
||||
packages-dir: ./hindsight-api-slim/dist
|
||||
skip-existing: true
|
||||
|
||||
- name: Publish hindsight-api to PyPI
|
||||
uses: pypa/gh-action-pypi-publish@release/v1
|
||||
with:
|
||||
@@ -70,7 +92,13 @@ jobs:
|
||||
- name: Publish hindsight-all to PyPI
|
||||
uses: pypa/gh-action-pypi-publish@release/v1
|
||||
with:
|
||||
packages-dir: ./hindsight/dist
|
||||
packages-dir: ./hindsight-all/dist
|
||||
skip-existing: true
|
||||
|
||||
- name: Publish hindsight-all-slim to PyPI
|
||||
uses: pypa/gh-action-pypi-publish@release/v1
|
||||
with:
|
||||
packages-dir: ./hindsight-all-slim/dist
|
||||
skip-existing: true
|
||||
|
||||
- name: Publish hindsight-litellm to PyPI
|
||||
@@ -97,19 +125,35 @@ jobs:
|
||||
packages-dir: ./hindsight-integrations/pydantic-ai/dist
|
||||
skip-existing: true
|
||||
|
||||
- name: Publish hindsight-hermes to PyPI
|
||||
uses: pypa/gh-action-pypi-publish@release/v1
|
||||
with:
|
||||
packages-dir: ./hindsight-integrations/hermes/dist
|
||||
skip-existing: true
|
||||
|
||||
- name: Publish hindsight-agno to PyPI
|
||||
uses: pypa/gh-action-pypi-publish@release/v1
|
||||
with:
|
||||
packages-dir: ./hindsight-integrations/agno/dist
|
||||
skip-existing: true
|
||||
|
||||
# Upload artifacts for GitHub release
|
||||
- name: Upload artifacts
|
||||
uses: actions/upload-artifact@v4
|
||||
uses: actions/upload-artifact@v7
|
||||
with:
|
||||
name: python-packages
|
||||
path: |
|
||||
hindsight-clients/python/dist/*
|
||||
hindsight-api-slim/dist/*
|
||||
hindsight-api/dist/*
|
||||
hindsight/dist/*
|
||||
hindsight-all/dist/*
|
||||
hindsight-all-slim/dist/*
|
||||
hindsight-integrations/litellm/dist/*
|
||||
hindsight-embed/dist/*
|
||||
hindsight-integrations/crewai/dist/*
|
||||
hindsight-integrations/pydantic-ai/dist/*
|
||||
hindsight-integrations/hermes/dist/*
|
||||
hindsight-integrations/agno/dist/*
|
||||
retention-days: 1
|
||||
|
||||
release-typescript-client:
|
||||
@@ -117,10 +161,10 @@ jobs:
|
||||
environment: npm
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v6
|
||||
|
||||
- name: Set up Node.js
|
||||
uses: actions/setup-node@v4
|
||||
uses: actions/setup-node@v6
|
||||
with:
|
||||
node-version: '20'
|
||||
registry-url: 'https://registry.npmjs.org'
|
||||
@@ -155,7 +199,7 @@ jobs:
|
||||
run: npm pack
|
||||
|
||||
- name: Upload artifacts
|
||||
uses: actions/upload-artifact@v4
|
||||
uses: actions/upload-artifact@v7
|
||||
with:
|
||||
name: typescript-client
|
||||
path: hindsight-clients/typescript/*.tgz
|
||||
@@ -166,10 +210,10 @@ jobs:
|
||||
environment: npm
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v6
|
||||
|
||||
- name: Set up Node.js
|
||||
uses: actions/setup-node@v4
|
||||
uses: actions/setup-node@v6
|
||||
with:
|
||||
node-version: '22'
|
||||
registry-url: 'https://registry.npmjs.org'
|
||||
@@ -204,7 +248,7 @@ jobs:
|
||||
run: npm pack
|
||||
|
||||
- name: Upload artifacts
|
||||
uses: actions/upload-artifact@v4
|
||||
uses: actions/upload-artifact@v7
|
||||
with:
|
||||
name: openclaw-integration
|
||||
path: hindsight-integrations/openclaw/*.tgz
|
||||
@@ -215,10 +259,10 @@ jobs:
|
||||
environment: npm
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v6
|
||||
|
||||
- name: Set up Node.js
|
||||
uses: actions/setup-node@v4
|
||||
uses: actions/setup-node@v6
|
||||
with:
|
||||
node-version: '22'
|
||||
registry-url: 'https://registry.npmjs.org'
|
||||
@@ -253,7 +297,7 @@ jobs:
|
||||
run: npm pack
|
||||
|
||||
- name: Upload artifacts
|
||||
uses: actions/upload-artifact@v4
|
||||
uses: actions/upload-artifact@v7
|
||||
with:
|
||||
name: ai-sdk-integration
|
||||
path: hindsight-integrations/ai-sdk/*.tgz
|
||||
@@ -264,10 +308,10 @@ jobs:
|
||||
environment: npm
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v6
|
||||
|
||||
- name: Set up Node.js
|
||||
uses: actions/setup-node@v4
|
||||
uses: actions/setup-node@v6
|
||||
with:
|
||||
node-version: '22'
|
||||
registry-url: 'https://registry.npmjs.org'
|
||||
@@ -302,7 +346,7 @@ jobs:
|
||||
run: npm pack
|
||||
|
||||
- name: Upload artifacts
|
||||
uses: actions/upload-artifact@v4
|
||||
uses: actions/upload-artifact@v7
|
||||
with:
|
||||
name: chat-integration
|
||||
path: hindsight-integrations/chat/*.tgz
|
||||
@@ -313,10 +357,10 @@ jobs:
|
||||
environment: npm
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v6
|
||||
|
||||
- name: Set up Node.js
|
||||
uses: actions/setup-node@v4
|
||||
uses: actions/setup-node@v6
|
||||
with:
|
||||
node-version: '20'
|
||||
registry-url: 'https://registry.npmjs.org'
|
||||
@@ -364,7 +408,7 @@ jobs:
|
||||
run: npm pack
|
||||
|
||||
- name: Upload artifacts
|
||||
uses: actions/upload-artifact@v4
|
||||
uses: actions/upload-artifact@v7
|
||||
with:
|
||||
name: control-plane
|
||||
path: hindsight-control-plane/*.tgz
|
||||
@@ -387,9 +431,13 @@ jobs:
|
||||
target: aarch64-apple-darwin
|
||||
artifact_name: hindsight
|
||||
asset_name: hindsight-darwin-arm64
|
||||
- os: ubuntu-24.04-arm
|
||||
target: aarch64-unknown-linux-gnu
|
||||
artifact_name: hindsight
|
||||
asset_name: hindsight-linux-arm64
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v6
|
||||
|
||||
- name: Install Rust
|
||||
uses: dtolnay/rust-toolchain@stable
|
||||
@@ -407,7 +455,7 @@ jobs:
|
||||
chmod +x artifacts/${{ matrix.asset_name }}
|
||||
|
||||
- name: Upload artifacts
|
||||
uses: actions/upload-artifact@v4
|
||||
uses: actions/upload-artifact@v7
|
||||
with:
|
||||
name: rust-cli-${{ matrix.asset_name }}
|
||||
path: artifacts/${{ matrix.asset_name }}
|
||||
@@ -448,7 +496,7 @@ jobs:
|
||||
PRELOAD_ML_MODELS=false
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v6
|
||||
|
||||
- name: Free Disk Space
|
||||
uses: jlumbroso/free-disk-space@main
|
||||
@@ -462,13 +510,13 @@ jobs:
|
||||
swap-storage: true
|
||||
|
||||
- name: Set up QEMU
|
||||
uses: docker/setup-qemu-action@v3
|
||||
uses: docker/setup-qemu-action@v4
|
||||
|
||||
- name: Set up Docker Buildx
|
||||
uses: docker/setup-buildx-action@v3
|
||||
uses: docker/setup-buildx-action@v4
|
||||
|
||||
- name: Log in to GitHub Container Registry
|
||||
uses: docker/login-action@v3
|
||||
uses: docker/login-action@v4
|
||||
with:
|
||||
registry: ghcr.io
|
||||
username: ${{ github.actor }}
|
||||
@@ -480,7 +528,7 @@ jobs:
|
||||
|
||||
- name: Extract metadata for release tags
|
||||
id: meta
|
||||
uses: docker/metadata-action@v5
|
||||
uses: docker/metadata-action@v6
|
||||
with:
|
||||
images: ghcr.io/${{ github.repository_owner }}/${{ matrix.image_name }}
|
||||
flavor: |
|
||||
@@ -496,7 +544,7 @@ jobs:
|
||||
# # Step 1: Build for local testing (single platform, no push)
|
||||
# # This creates an identical image to what will be released, just for one platform
|
||||
# - name: Build image for testing
|
||||
# uses: docker/build-push-action@v6
|
||||
# uses: docker/build-push-action@v7
|
||||
# with:
|
||||
# context: .
|
||||
# file: docker/standalone/Dockerfile
|
||||
@@ -515,7 +563,7 @@ jobs:
|
||||
|
||||
# Build multi-platform and push to release tags
|
||||
- name: Build and push release images
|
||||
uses: docker/build-push-action@v6
|
||||
uses: docker/build-push-action@v7
|
||||
with:
|
||||
context: .
|
||||
file: docker/standalone/Dockerfile
|
||||
@@ -533,7 +581,7 @@ jobs:
|
||||
packages: write
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v6
|
||||
|
||||
- name: Install Helm
|
||||
uses: azure/setup-helm@v4
|
||||
@@ -553,7 +601,7 @@ jobs:
|
||||
run: helm push helm-packages/*.tgz oci://ghcr.io/${{ github.repository_owner }}/charts
|
||||
|
||||
- name: Upload artifacts
|
||||
uses: actions/upload-artifact@v4
|
||||
uses: actions/upload-artifact@v7
|
||||
with:
|
||||
name: helm-chart
|
||||
path: helm-packages/*.tgz
|
||||
@@ -566,68 +614,68 @@ jobs:
|
||||
contents: write
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v6
|
||||
|
||||
- name: Extract version from tag
|
||||
id: get_version
|
||||
run: echo "VERSION=${GITHUB_REF#refs/tags/v}" >> $GITHUB_OUTPUT
|
||||
|
||||
- name: Download Python packages
|
||||
uses: actions/download-artifact@v4
|
||||
uses: actions/download-artifact@v8
|
||||
with:
|
||||
name: python-packages
|
||||
path: ./artifacts/python-packages
|
||||
|
||||
- name: Download TypeScript client
|
||||
uses: actions/download-artifact@v4
|
||||
uses: actions/download-artifact@v8
|
||||
with:
|
||||
name: typescript-client
|
||||
path: ./artifacts/typescript-client
|
||||
|
||||
- name: Download OpenClaw Integration
|
||||
uses: actions/download-artifact@v4
|
||||
uses: actions/download-artifact@v8
|
||||
with:
|
||||
name: openclaw-integration
|
||||
path: ./artifacts/openclaw-integration
|
||||
|
||||
- name: Download AI SDK Integration
|
||||
uses: actions/download-artifact@v4
|
||||
uses: actions/download-artifact@v8
|
||||
with:
|
||||
name: ai-sdk-integration
|
||||
path: ./artifacts/ai-sdk-integration
|
||||
|
||||
- name: Download Chat Integration
|
||||
uses: actions/download-artifact@v4
|
||||
uses: actions/download-artifact@v8
|
||||
with:
|
||||
name: chat-integration
|
||||
path: ./artifacts/chat-integration
|
||||
|
||||
- name: Download Control Plane
|
||||
uses: actions/download-artifact@v4
|
||||
uses: actions/download-artifact@v8
|
||||
with:
|
||||
name: control-plane
|
||||
path: ./artifacts/control-plane
|
||||
|
||||
- name: Download Rust CLI (Linux)
|
||||
uses: actions/download-artifact@v4
|
||||
uses: actions/download-artifact@v8
|
||||
with:
|
||||
name: rust-cli-hindsight-linux-amd64
|
||||
path: ./artifacts/rust-cli-linux
|
||||
|
||||
- name: Download Rust CLI (macOS Intel)
|
||||
uses: actions/download-artifact@v4
|
||||
uses: actions/download-artifact@v8
|
||||
with:
|
||||
name: rust-cli-hindsight-darwin-amd64
|
||||
path: ./artifacts/rust-cli-darwin-amd64
|
||||
|
||||
- name: Download Rust CLI (macOS ARM)
|
||||
uses: actions/download-artifact@v4
|
||||
uses: actions/download-artifact@v8
|
||||
with:
|
||||
name: rust-cli-hindsight-darwin-arm64
|
||||
path: ./artifacts/rust-cli-darwin-arm64
|
||||
|
||||
- name: Download Helm chart
|
||||
uses: actions/download-artifact@v4
|
||||
uses: actions/download-artifact@v8
|
||||
with:
|
||||
name: helm-chart
|
||||
path: ./artifacts/helm-chart
|
||||
@@ -637,10 +685,14 @@ jobs:
|
||||
mkdir -p release-assets
|
||||
# Python packages
|
||||
cp artifacts/python-packages/hindsight-clients/python/dist/* release-assets/ || true
|
||||
cp artifacts/python-packages/hindsight-api-slim/dist/* release-assets/ || true
|
||||
cp artifacts/python-packages/hindsight-api/dist/* release-assets/ || true
|
||||
cp artifacts/python-packages/hindsight/dist/* release-assets/ || true
|
||||
cp artifacts/python-packages/hindsight-all/dist/* release-assets/ || true
|
||||
cp artifacts/python-packages/hindsight-all-slim/dist/* release-assets/ || true
|
||||
cp artifacts/python-packages/hindsight-integrations/litellm/dist/* release-assets/ || true
|
||||
cp artifacts/python-packages/hindsight-integrations/pydantic-ai/dist/* release-assets/ || true
|
||||
cp artifacts/python-packages/hindsight-integrations/hermes/dist/* release-assets/ || true
|
||||
cp artifacts/python-packages/hindsight-integrations/agno/dist/* release-assets/ || true
|
||||
cp artifacts/python-packages/hindsight-embed/dist/* release-assets/ || true
|
||||
# TypeScript client
|
||||
cp artifacts/typescript-client/*.tgz release-assets/ || true
|
||||
|
||||
+411
-183
File diff suppressed because it is too large
Load Diff
@@ -17,20 +17,20 @@ Hindsight is an agent memory system that provides long-term memory for AI agents
|
||||
./scripts/dev/start-api.sh
|
||||
|
||||
# Run all tests (parallelized with pytest-xdist)
|
||||
cd hindsight-api && uv run pytest tests/
|
||||
cd hindsight-api-slim && uv run pytest tests/
|
||||
|
||||
# Run specific test file
|
||||
cd hindsight-api && uv run pytest tests/test_http_api_integration.py -v
|
||||
cd hindsight-api-slim && uv run pytest tests/test_http_api_integration.py -v
|
||||
|
||||
# Run single test function
|
||||
cd hindsight-api && uv run pytest tests/test_retain.py::test_retain_simple -v
|
||||
cd hindsight-api-slim && uv run pytest tests/test_retain.py::test_retain_simple -v
|
||||
|
||||
# Lint and format
|
||||
cd hindsight-api && uv run ruff check .
|
||||
cd hindsight-api && uv run ruff format .
|
||||
cd hindsight-api-slim && uv run ruff check .
|
||||
cd hindsight-api-slim && uv run ruff format .
|
||||
|
||||
# Type checking (uses ty - extremely fast type checker from Astral)
|
||||
cd hindsight-api && uv run ty check hindsight_api/
|
||||
cd hindsight-api-slim && uv run ty check hindsight_api/
|
||||
```
|
||||
|
||||
### Control Plane (Next.js)
|
||||
@@ -72,7 +72,7 @@ cd hindsight-control-plane && npm run dev
|
||||
## Architecture
|
||||
|
||||
### Monorepo Structure
|
||||
- **hindsight-api/**: Core FastAPI server with memory engine (Python, uv)
|
||||
- **hindsight-api-slim/**: Core FastAPI server with memory engine (Python, uv)
|
||||
- **hindsight/**: Embedded Python bundle (hindsight-all package)
|
||||
- **hindsight-control-plane/**: Admin UI (Next.js, npm)
|
||||
- **hindsight-cli/**: CLI tool (Rust, cargo, uses progenitor for API client)
|
||||
@@ -81,9 +81,9 @@ cd hindsight-control-plane && npm run dev
|
||||
- **hindsight-integrations/**: Framework integrations (LiteLLM, OpenAI)
|
||||
- **hindsight-dev/**: Development tools and benchmarks
|
||||
|
||||
### Core Engine (hindsight-api/hindsight_api/engine/)
|
||||
### Core Engine (hindsight-api-slim/hindsight_api/engine/)
|
||||
- `memory_engine.py`: Main orchestrator (~170KB) for retain/recall/reflect operations
|
||||
- `llm_wrapper.py`: LLM abstraction supporting OpenAI, Anthropic, Gemini, Groq, Ollama, LM Studio
|
||||
- `llm_wrapper.py`: LLM abstraction supporting OpenAI, Anthropic, Gemini, Groq, MiniMax, Ollama, LM Studio
|
||||
- `embeddings.py`: Embedding generation (local sentence-transformers or TEI)
|
||||
- `cross_encoder.py`: Reranking (local or TEI)
|
||||
- `entity_resolver.py`: Entity extraction and normalization
|
||||
@@ -101,7 +101,7 @@ cd hindsight-control-plane && npm run dev
|
||||
- `fusion.py`: Reciprocal rank fusion for combining results
|
||||
- `reranking.py`: Cross-encoder reranking
|
||||
|
||||
### API Layer (hindsight-api/hindsight_api/api/)
|
||||
### API Layer (hindsight-api-slim/hindsight_api/api/)
|
||||
- `http.py`: FastAPI HTTP routers (~80KB) for all REST endpoints
|
||||
- `mcp.py`: Model Context Protocol server implementation
|
||||
|
||||
@@ -111,13 +111,13 @@ Main operations:
|
||||
- **Reflect**: Disposition-aware reasoning using memories and mental models.
|
||||
|
||||
### Database
|
||||
PostgreSQL with pgvector. Schema managed via Alembic migrations in `hindsight-api/hindsight_api/alembic/`. Migrations run automatically on API startup.
|
||||
PostgreSQL with pgvector. Schema managed via Alembic migrations in `hindsight-api-slim/hindsight_api/alembic/`. Migrations run automatically on API startup.
|
||||
|
||||
Key tables: `banks`, `memory_units`, `documents`, `entities`, `entity_links`
|
||||
|
||||
### Adding Database Migrations
|
||||
|
||||
1. **Create a new migration file** in `hindsight-api/hindsight_api/alembic/versions/`:
|
||||
1. **Create a new migration file** in `hindsight-api-slim/hindsight_api/alembic/versions/`:
|
||||
- File name format: `<revision_id>_<description>.py` (e.g., `f1a2b3c4d5e6_add_new_index.py`)
|
||||
- Use a unique hex revision ID (12 chars)
|
||||
- Set `down_revision` to the previous migration's revision ID
|
||||
@@ -154,7 +154,7 @@ Key tables: `banks`, `memory_units`, `documents`, `entities`, `entity_links`
|
||||
|
||||
3. **Run migrations locally**:
|
||||
```bash
|
||||
# Set database URL and run migrations
|
||||
# Set database URL and run migrations for the base schema plus all tenants
|
||||
uv run hindsight-admin run-db-migration
|
||||
|
||||
# Run on a specific tenant schema
|
||||
@@ -251,7 +251,7 @@ Fields must be categorized as either **hierarchical** (can be overridden per-ten
|
||||
|
||||
#### Adding a New Configuration Field
|
||||
|
||||
1. **config.py** (`hindsight-api/hindsight_api/config.py`):
|
||||
1. **config.py** (`hindsight-api-slim/hindsight_api/config.py`):
|
||||
- Add `ENV_*` constant for the environment variable name (e.g., `ENV_MY_SETTING = "HINDSIGHT_API_MY_SETTING"`)
|
||||
- Add `DEFAULT_*` constant for the default value
|
||||
- Add field to `HindsightConfig` dataclass with type annotation
|
||||
@@ -268,7 +268,7 @@ Fields must be categorized as either **hierarchical** (can be overridden per-ten
|
||||
# Static field - just don't add to _HIERARCHICAL_FIELDS
|
||||
```
|
||||
|
||||
2. **main.py** (`hindsight-api/hindsight_api/main.py`):
|
||||
2. **main.py** (`hindsight-api-slim/hindsight_api/main.py`):
|
||||
- Add field to the manual `HindsightConfig()` constructor call (search for "CLI override")
|
||||
|
||||
3. **Use hierarchical config in MemoryEngine**:
|
||||
@@ -308,14 +308,14 @@ cp .env.example .env
|
||||
# Edit .env with LLM API key
|
||||
|
||||
# Python deps
|
||||
uv sync --directory hindsight-api/
|
||||
uv sync --directory hindsight-api-slim/
|
||||
|
||||
# Node deps (uses npm workspaces)
|
||||
npm install
|
||||
```
|
||||
|
||||
Required env vars:
|
||||
- `HINDSIGHT_API_LLM_PROVIDER`: openai, anthropic, gemini, groq, ollama, lmstudio
|
||||
- `HINDSIGHT_API_LLM_PROVIDER`: openai, anthropic, gemini, groq, minimax, ollama, lmstudio
|
||||
- `HINDSIGHT_API_LLM_API_KEY`: Your API key
|
||||
- `HINDSIGHT_API_LLM_MODEL`: Model name (e.g., gpt-4o-mini, claude-sonnet-4-20250514)
|
||||
|
||||
|
||||
@@ -9,8 +9,9 @@
|
||||
[](https://opensource.org/licenses/MIT)
|
||||

|
||||

|
||||
<br/>
|
||||
|
||||
|
||||
<a href="https://trendshift.io/repositories/15603" target="_blank"><img src="https://trendshift.io/api/badge/repositories/15603" alt="vectorize-io%2Fhindsight | Trendshift" style="width: 250px; height: 55px;" width="250" height="55"/></a>
|
||||
</div>
|
||||
|
||||
---
|
||||
@@ -69,7 +70,7 @@ docker run --rm -it --pull always -p 8888:8888 -p 9999:9999 \
|
||||
>API: http://localhost:8888
|
||||
>UI: http://localhost:9999
|
||||
|
||||
You can modify the LLM provider by setting `HINDSIGHT_API_LLM_PROVIDER`. Valid options are `openai`, `anthropic`, `gemini`, `groq`, `ollama`, and `lmstudio`. The documentation provides more details on [supported models](https://hindsight.vectorize.io/developer/models).
|
||||
You can modify the LLM provider by setting `HINDSIGHT_API_LLM_PROVIDER`. Valid options are `openai`, `anthropic`, `gemini`, `groq`, `ollama`, `lmstudio`, and `minimax`. The documentation provides more details on [supported models](https://hindsight.vectorize.io/developer/models).
|
||||
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,139 @@
|
||||
{
|
||||
"version": "5",
|
||||
"specifiers": {
|
||||
"jsr:@std/assert@^1.0.17": "1.0.19",
|
||||
"jsr:@std/assert@^1.0.19": "1.0.19",
|
||||
"jsr:@std/expect@*": "1.0.18",
|
||||
"jsr:@std/internal@^1.0.12": "1.0.12",
|
||||
"jsr:@std/path@^1.1.4": "1.1.4",
|
||||
"jsr:@std/testing@*": "1.0.17"
|
||||
},
|
||||
"jsr": {
|
||||
"@std/[email protected]": {
|
||||
"integrity": "eaada96ee120cb980bc47e040f82814d786fe8162ecc53c91d8df60b8755991e",
|
||||
"dependencies": [
|
||||
"jsr:@std/internal"
|
||||
]
|
||||
},
|
||||
"@std/[email protected]": {
|
||||
"integrity": "8566eab35200466f8609eb7e7aed062ed0db314e9a258d5d201b1b8997ce801a",
|
||||
"dependencies": [
|
||||
"jsr:@std/assert@^1.0.19",
|
||||
"jsr:@std/internal",
|
||||
"jsr:@std/path"
|
||||
]
|
||||
},
|
||||
"@std/[email protected]": {
|
||||
"integrity": "972a634fd5bc34b242024402972cd5143eac68d8dffaca5eaa4dba30ce17b027"
|
||||
},
|
||||
"@std/[email protected]": {
|
||||
"integrity": "1d2d43f39efb1b42f0b1882a25486647cb851481862dc7313390b2bb044314b5",
|
||||
"dependencies": [
|
||||
"jsr:@std/internal"
|
||||
]
|
||||
},
|
||||
"@std/[email protected]": {
|
||||
"integrity": "87bdc2700fa98249d48a17cd72413352d3d3680dcfbdb64947fd0982d6bbf681",
|
||||
"dependencies": [
|
||||
"jsr:@std/assert@^1.0.17",
|
||||
"jsr:@std/internal"
|
||||
]
|
||||
}
|
||||
},
|
||||
"workspace": {
|
||||
"members": {
|
||||
"hindsight-clients/typescript": {
|
||||
"packageJson": {
|
||||
"dependencies": [
|
||||
"npm:@hey-api/[email protected]",
|
||||
"npm:@types/jest@29",
|
||||
"npm:@types/node@20",
|
||||
"npm:jest@29",
|
||||
"npm:ts-jest@29",
|
||||
"npm:tsup@^8.5.1",
|
||||
"npm:typescript@5"
|
||||
]
|
||||
}
|
||||
},
|
||||
"hindsight-control-plane": {
|
||||
"packageJson": {
|
||||
"dependencies": [
|
||||
"npm:@eslint/eslintrc@^3.3.3",
|
||||
"npm:@eslint/js@^9.39.2",
|
||||
"npm:@radix-ui/react-alert-dialog@^1.1.15",
|
||||
"npm:@radix-ui/react-checkbox@^1.3.3",
|
||||
"npm:@radix-ui/react-dialog@^1.1.15",
|
||||
"npm:@radix-ui/react-dropdown-menu@^2.1.16",
|
||||
"npm:@radix-ui/react-label@^2.1.8",
|
||||
"npm:@radix-ui/react-popover@^1.1.15",
|
||||
"npm:@radix-ui/react-radio-group@^1.3.8",
|
||||
"npm:@radix-ui/react-select@^2.2.6",
|
||||
"npm:@radix-ui/react-slider@^1.3.6",
|
||||
"npm:@radix-ui/react-slot@^1.2.4",
|
||||
"npm:@radix-ui/react-switch@^1.2.6",
|
||||
"npm:@radix-ui/react-tabs@^1.1.13",
|
||||
"npm:@radix-ui/react-tooltip@^1.2.8",
|
||||
"npm:@tailwindcss/postcss@^4.1.17",
|
||||
"npm:@tailwindcss/typography@~0.5.19",
|
||||
"npm:@types/cytoscape@^3.21.9",
|
||||
"npm:@types/node@^24.10.0",
|
||||
"npm:@types/react-dom@^19.2.2",
|
||||
"npm:@types/react@^19.2.2",
|
||||
"npm:autoprefixer@^10.4.21",
|
||||
"npm:class-variance-authority@~0.7.1",
|
||||
"npm:clsx@^2.1.1",
|
||||
"npm:cmdk@^1.1.1",
|
||||
"npm:cytoscape-fcose@^2.2.0",
|
||||
"npm:cytoscape@^3.33.1",
|
||||
"npm:eslint-config-next@^16.0.1",
|
||||
"npm:eslint-plugin-react-hooks@^7.0.1",
|
||||
"npm:eslint-plugin-react@^7.37.5",
|
||||
"npm:eslint@^9.39.1",
|
||||
"npm:[email protected]",
|
||||
"npm:next-themes@~0.4.6",
|
||||
"npm:next@^16.1.6",
|
||||
"npm:postcss@^8.5.6",
|
||||
"npm:prettier@^3.7.4",
|
||||
"npm:react-chrono@^2.9.1",
|
||||
"npm:react-dom@^19.2.0",
|
||||
"npm:react-markdown@^10.1.0",
|
||||
"npm:react18-json-view@~0.2.9",
|
||||
"npm:react@^19.2.0",
|
||||
"npm:recharts@^3.5.1",
|
||||
"npm:remark-gfm@^4.0.1",
|
||||
"npm:sonner@^2.0.7",
|
||||
"npm:tailwind-merge@^3.4.0",
|
||||
"npm:tailwindcss-animate@^1.0.7",
|
||||
"npm:tailwindcss@^4.1.17",
|
||||
"npm:[email protected]",
|
||||
"npm:typescript-eslint@^8.50.0",
|
||||
"npm:typescript@^5.9.3"
|
||||
]
|
||||
}
|
||||
},
|
||||
"hindsight-docs": {
|
||||
"packageJson": {
|
||||
"dependencies": [
|
||||
"npm:@docusaurus/[email protected]",
|
||||
"npm:@docusaurus/[email protected]",
|
||||
"npm:@docusaurus/[email protected]",
|
||||
"npm:@docusaurus/theme-common@^3.9.2",
|
||||
"npm:@docusaurus/theme-mermaid@^3.9.2",
|
||||
"npm:@docusaurus/[email protected]",
|
||||
"npm:@docusaurus/[email protected]",
|
||||
"npm:@easyops-cn/docusaurus-search-local@~0.52.2",
|
||||
"npm:@mdx-js/react@3",
|
||||
"npm:clsx@2",
|
||||
"npm:prism-react-renderer@^2.3.0",
|
||||
"npm:raw-loader@^4.0.2",
|
||||
"npm:react-dom@19",
|
||||
"npm:react-icons@^5.6.0",
|
||||
"npm:react@19",
|
||||
"npm:redocusaurus@^2.5.0",
|
||||
"npm:typescript@~5.6.2"
|
||||
]
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -42,25 +42,22 @@ RUN apt-get update && apt-get install -y \
|
||||
&& pip install --no-cache-dir uv
|
||||
|
||||
# Copy dependency files and README (required by pyproject.toml)
|
||||
COPY hindsight-api/pyproject.toml ./api/
|
||||
COPY hindsight-api/README.md ./api/
|
||||
COPY hindsight-api-slim/pyproject.toml ./api/
|
||||
COPY hindsight-api-slim/README.md ./api/
|
||||
|
||||
WORKDIR /app/api
|
||||
|
||||
# Remove local ML model dependencies if INCLUDE_LOCAL_MODELS=false
|
||||
# This creates a smaller image when using external providers (TEI, OpenAI, Cohere)
|
||||
RUN if [ "$INCLUDE_LOCAL_MODELS" != "true" ]; then \
|
||||
echo "Removing local-models dependencies (sentence-transformers, torch, transformers)..." && \
|
||||
sed -i '/"sentence-transformers/d' pyproject.toml && \
|
||||
sed -i '/"transformers/d' pyproject.toml && \
|
||||
sed -i '/"torch/d' pyproject.toml; \
|
||||
# Sync dependencies using appropriate extras based on INCLUDE_LOCAL_MODELS
|
||||
# local-ml: torch, sentence-transformers, transformers, einops, flashrank, mlx (optional)
|
||||
# embedded-db: pg0-embedded (always included for embedded PostgreSQL support)
|
||||
RUN if [ "$INCLUDE_LOCAL_MODELS" = "true" ]; then \
|
||||
uv sync --extra local-ml --extra embedded-db; \
|
||||
else \
|
||||
uv sync --extra embedded-db; \
|
||||
fi
|
||||
|
||||
# Sync dependencies (will create lock file if needed)
|
||||
RUN uv sync
|
||||
|
||||
# Copy source code (alembic migrations are inside hindsight_api/)
|
||||
COPY hindsight-api/hindsight_api ./hindsight_api
|
||||
COPY hindsight-api-slim/hindsight_api ./hindsight_api
|
||||
|
||||
# Install the local package (uv sync only installed dependencies, not the package itself)
|
||||
RUN uv pip install -e .
|
||||
|
||||
@@ -77,18 +77,32 @@ PIDS=()
|
||||
# Start API if enabled
|
||||
if [ "$ENABLE_API" = "true" ]; then
|
||||
cd /app/api
|
||||
API_HEALTH_URL="${HINDSIGHT_API_HEALTH_URL:-http://localhost:8888/health}"
|
||||
API_STARTUP_WAIT_SECONDS="${HINDSIGHT_API_STARTUP_WAIT_SECONDS:-300}"
|
||||
|
||||
# Run API directly - Python's PYTHONUNBUFFERED=1 handles output buffering
|
||||
hindsight-api &
|
||||
API_PID=$!
|
||||
PIDS+=($API_PID)
|
||||
|
||||
# Wait for API to be ready
|
||||
for i in {1..60}; do
|
||||
if curl -sf http://localhost:8888/health &>/dev/null; then
|
||||
api_ready=false
|
||||
for ((i=1; i<=API_STARTUP_WAIT_SECONDS; i++)); do
|
||||
if ! kill -0 "$API_PID" 2>/dev/null; then
|
||||
wait "$API_PID"
|
||||
exit $?
|
||||
fi
|
||||
if curl -sf "$API_HEALTH_URL" &>/dev/null; then
|
||||
api_ready=true
|
||||
break
|
||||
fi
|
||||
sleep 1
|
||||
done
|
||||
|
||||
if [ "$api_ready" != "true" ]; then
|
||||
echo "❌ API did not become healthy within ${API_STARTUP_WAIT_SECONDS}s"
|
||||
exit 1
|
||||
fi
|
||||
else
|
||||
echo "API disabled (HINDSIGHT_ENABLE_API=false)"
|
||||
fi
|
||||
@@ -97,6 +111,7 @@ fi
|
||||
if [ "$ENABLE_CP" = "true" ]; then
|
||||
echo "🎛️ Starting Control Plane..."
|
||||
cd /app/control-plane
|
||||
export HOSTNAME="${HINDSIGHT_CP_HOSTNAME:-0.0.0.0}"
|
||||
PORT="${HINDSIGHT_CP_PORT:-9999}" node server.js &
|
||||
CP_PID=$!
|
||||
PIDS+=($CP_PID)
|
||||
|
||||
@@ -49,6 +49,9 @@
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
REPO_ROOT="$(dirname "$SCRIPT_DIR")"
|
||||
|
||||
# Colors for output
|
||||
RED='\033[0;31m'
|
||||
GREEN='\033[0;32m'
|
||||
@@ -178,6 +181,21 @@ for i in $(seq 1 "$TIMEOUT"); do
|
||||
echo "=== Health Response ==="
|
||||
curl -s "http://localhost:${HEALTH_PORT}${HEALTH_PATH}" | python3 -m json.tool 2>/dev/null || curl -s "http://localhost:${HEALTH_PORT}${HEALTH_PATH}"
|
||||
echo ""
|
||||
|
||||
# Run retain/recall smoke test for API targets
|
||||
if [ "$TARGET" != "cp-only" ]; then
|
||||
echo ""
|
||||
echo "=== Retain/Recall Smoke Test ==="
|
||||
if ! "$REPO_ROOT/scripts/smoke-test-slim.sh" "http://localhost:${HEALTH_PORT}"; then
|
||||
echo ""
|
||||
echo "=== Container Logs (last 50 lines) ==="
|
||||
docker logs "$CONTAINER_NAME" 2>&1 | tail -50
|
||||
echo ""
|
||||
echo -e "${RED}Smoke test FAILED${NC}"
|
||||
exit 1
|
||||
fi
|
||||
fi
|
||||
|
||||
echo ""
|
||||
echo "=== Container Logs (last 50 lines) ==="
|
||||
docker logs "$CONTAINER_NAME" 2>&1 | tail -50
|
||||
|
||||
@@ -2,8 +2,8 @@ apiVersion: v2
|
||||
name: hindsight
|
||||
description: Hindsight helm chart
|
||||
type: application
|
||||
version: 0.4.15
|
||||
appVersion: "0.4.15"
|
||||
version: 0.4.18
|
||||
appVersion: "0.4.18"
|
||||
keywords:
|
||||
- ai
|
||||
- memory
|
||||
|
||||
@@ -0,0 +1,33 @@
|
||||
[build-system]
|
||||
requires = ["setuptools>=61"]
|
||||
build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "hindsight-all-slim"
|
||||
version = "0.4.18"
|
||||
description = "Hindsight: Agent Memory That Works Like Human Memory - Slim All-in-One Bundle"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.11"
|
||||
dependencies = [
|
||||
"hindsight-api-slim>=0.4.17",
|
||||
"hindsight-client>=0.0.7",
|
||||
"hindsight-embed>=0.1.0",
|
||||
]
|
||||
|
||||
[tool.uv.sources]
|
||||
hindsight-api-slim = { workspace = true }
|
||||
hindsight-client = { workspace = true }
|
||||
hindsight-embed = { workspace = true }
|
||||
|
||||
[project.optional-dependencies]
|
||||
test = [
|
||||
"pytest>=7.0.0",
|
||||
"pytest-asyncio>=0.21.0",
|
||||
]
|
||||
|
||||
[tool.setuptools]
|
||||
packages = []
|
||||
|
||||
[tool.pytest.ini_options]
|
||||
asyncio_mode = "auto"
|
||||
asyncio_default_fixture_loop_scope = "function"
|
||||
@@ -0,0 +1,48 @@
|
||||
# hindsight-all
|
||||
|
||||
All-in-one package for Hindsight - Agent Memory That Works Like Human Memory
|
||||
|
||||
## Quick Start
|
||||
|
||||
```python
|
||||
from hindsight import start_server, HindsightClient
|
||||
|
||||
# Start server with embedded PostgreSQL
|
||||
server = start_server(
|
||||
llm_provider="groq",
|
||||
llm_api_key="your-api-key",
|
||||
llm_model="openai/gpt-oss-120b"
|
||||
)
|
||||
|
||||
# Create client
|
||||
client = HindsightClient(base_url=server.url)
|
||||
|
||||
# Store memories
|
||||
client.put(agent_id="assistant", content="User prefers Python for data analysis")
|
||||
|
||||
# Search memories
|
||||
results = client.search(agent_id="assistant", query="programming preferences")
|
||||
|
||||
# Generate contextual response
|
||||
response = client.think(agent_id="assistant", query="What languages should I recommend?")
|
||||
|
||||
# Stop server when done
|
||||
server.stop()
|
||||
```
|
||||
|
||||
## Using Context Manager
|
||||
|
||||
```python
|
||||
from hindsight import HindsightServer, HindsightClient
|
||||
|
||||
with HindsightServer(llm_provider="groq", llm_api_key="...") as server:
|
||||
client = HindsightClient(base_url=server.url)
|
||||
# ... use client ...
|
||||
# Server automatically stops
|
||||
```
|
||||
|
||||
## Installation
|
||||
|
||||
```bash
|
||||
pip install hindsight-all
|
||||
```
|
||||
@@ -0,0 +1,423 @@
|
||||
"""
|
||||
Wrapper for Hindsight client that adds API namespaces.
|
||||
|
||||
Provides organized access to different parts of the Hindsight API through
|
||||
namespaces like .banks, .mental_models, etc.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from hindsight_client import Hindsight
|
||||
|
||||
|
||||
class BanksAPI:
|
||||
"""Namespace for bank-related operations.
|
||||
|
||||
Provides methods to create, delete, and manage memory banks.
|
||||
"""
|
||||
|
||||
def __init__(self, client: Hindsight):
|
||||
self._client = client
|
||||
|
||||
def create(
|
||||
self,
|
||||
bank_id: str,
|
||||
name: str | None = None,
|
||||
mission: str | None = None,
|
||||
disposition: dict[str, Any] | None = None,
|
||||
) -> Any:
|
||||
"""Create a new bank.
|
||||
|
||||
Args:
|
||||
bank_id: Unique identifier for the bank.
|
||||
name: Optional display name for the bank.
|
||||
mission: Optional mission statement for the bank.
|
||||
disposition: Optional disposition configuration dict.
|
||||
|
||||
Returns:
|
||||
Bank creation response from the API.
|
||||
"""
|
||||
return self._client.create_bank(
|
||||
bank_id=bank_id,
|
||||
name=name,
|
||||
mission=mission,
|
||||
disposition=disposition,
|
||||
)
|
||||
|
||||
def delete(self, bank_id: str) -> Any:
|
||||
"""Delete a bank.
|
||||
|
||||
Args:
|
||||
bank_id: The ID of the bank to delete.
|
||||
|
||||
Returns:
|
||||
Deletion response from the API.
|
||||
"""
|
||||
return self._client.delete_bank(bank_id=bank_id)
|
||||
|
||||
def set_mission(self, bank_id: str, mission: str) -> Any:
|
||||
"""Set or update the mission for a bank.
|
||||
|
||||
Args:
|
||||
bank_id: The ID of the bank.
|
||||
mission: The mission statement to set.
|
||||
|
||||
Returns:
|
||||
API response confirming the update.
|
||||
"""
|
||||
return self._client.set_mission(bank_id=bank_id, mission=mission)
|
||||
|
||||
def set_disposition(self, bank_id: str, disposition: dict[str, Any]) -> Any:
|
||||
"""Set or update the disposition for a bank.
|
||||
|
||||
Args:
|
||||
bank_id: The ID of the bank.
|
||||
disposition: The disposition configuration dict.
|
||||
|
||||
Returns:
|
||||
API response confirming the update.
|
||||
"""
|
||||
return self._client.set_disposition(bank_id=bank_id, disposition=disposition)
|
||||
|
||||
def list(self) -> Any:
|
||||
"""List all banks.
|
||||
|
||||
Returns:
|
||||
List of banks from the API.
|
||||
"""
|
||||
from hindsight_client.hindsight_client import _run_async
|
||||
|
||||
return _run_async(self._client._banks_api.list_banks())
|
||||
|
||||
|
||||
class MentalModelsAPI:
|
||||
"""Namespace for mental model operations.
|
||||
|
||||
Mental models are reusable knowledge structures that guide agent behavior.
|
||||
"""
|
||||
|
||||
def __init__(self, client: Hindsight):
|
||||
self._client = client
|
||||
|
||||
def create(
|
||||
self,
|
||||
bank_id: str,
|
||||
name: str,
|
||||
content: str,
|
||||
tags: list[str] | None = None,
|
||||
) -> Any:
|
||||
"""Create a new mental model.
|
||||
|
||||
Args:
|
||||
bank_id: The ID of the bank to add the model to.
|
||||
name: Name for the mental model.
|
||||
content: The content/instructions for the mental model.
|
||||
tags: Optional list of tags for categorization.
|
||||
|
||||
Returns:
|
||||
Creation response from the API.
|
||||
"""
|
||||
return self._client.create_mental_model(
|
||||
bank_id=bank_id,
|
||||
name=name,
|
||||
content=content,
|
||||
tags=tags,
|
||||
)
|
||||
|
||||
def list(self, bank_id: str, tags: list[str] | None = None) -> Any:
|
||||
"""List all mental models for a bank.
|
||||
|
||||
Args:
|
||||
bank_id: The ID of the bank.
|
||||
tags: Optional filter by tags.
|
||||
|
||||
Returns:
|
||||
List of mental models.
|
||||
"""
|
||||
return self._client.list_mental_models(bank_id=bank_id, tags=tags)
|
||||
|
||||
def get(self, bank_id: str, mental_model_id: str) -> Any:
|
||||
"""Get a specific mental model.
|
||||
|
||||
Args:
|
||||
bank_id: The ID of the bank.
|
||||
mental_model_id: The ID of the mental model.
|
||||
|
||||
Returns:
|
||||
The mental model details.
|
||||
"""
|
||||
return self._client.get_mental_model(bank_id=bank_id, mental_model_id=mental_model_id)
|
||||
|
||||
def refresh(self, bank_id: str, mental_model_id: str) -> Any:
|
||||
"""Refresh a mental model.
|
||||
|
||||
Args:
|
||||
bank_id: The ID of the bank.
|
||||
mental_model_id: The ID of the mental model to refresh.
|
||||
|
||||
Returns:
|
||||
Refresh response from the API.
|
||||
"""
|
||||
return self._client.refresh_mental_model(bank_id=bank_id, mental_model_id=mental_model_id)
|
||||
|
||||
def update(
|
||||
self,
|
||||
bank_id: str,
|
||||
mental_model_id: str,
|
||||
name: str | None = None,
|
||||
content: str | None = None,
|
||||
tags: list[str] | None = None,
|
||||
) -> Any:
|
||||
"""Update a mental model.
|
||||
|
||||
Args:
|
||||
bank_id: The ID of the bank.
|
||||
mental_model_id: The ID of the mental model to update.
|
||||
name: Optional new name.
|
||||
content: Optional new content.
|
||||
tags: Optional new tags list.
|
||||
|
||||
Returns:
|
||||
Update response from the API.
|
||||
"""
|
||||
return self._client.update_mental_model(
|
||||
bank_id=bank_id,
|
||||
mental_model_id=mental_model_id,
|
||||
name=name,
|
||||
content=content,
|
||||
tags=tags,
|
||||
)
|
||||
|
||||
def delete(self, bank_id: str, mental_model_id: str) -> Any:
|
||||
"""Delete a mental model.
|
||||
|
||||
Args:
|
||||
bank_id: The ID of the bank.
|
||||
mental_model_id: The ID of the mental model to delete.
|
||||
|
||||
Returns:
|
||||
Deletion response from the API.
|
||||
"""
|
||||
return self._client.delete_mental_model(bank_id=bank_id, mental_model_id=mental_model_id)
|
||||
|
||||
|
||||
class DirectivesAPI:
|
||||
"""Namespace for directive operations.
|
||||
|
||||
Directives are explicit instructions that guide agent behavior.
|
||||
"""
|
||||
|
||||
def __init__(self, client: Hindsight):
|
||||
self._client = client
|
||||
|
||||
def create(
|
||||
self,
|
||||
bank_id: str,
|
||||
name: str,
|
||||
content: str,
|
||||
tags: list[str] | None = None,
|
||||
) -> Any:
|
||||
"""Create a new directive.
|
||||
|
||||
Args:
|
||||
bank_id: The ID of the bank to add the directive to.
|
||||
name: Name for the directive.
|
||||
content: The directive content/instructions.
|
||||
tags: Optional list of tags for categorization.
|
||||
|
||||
Returns:
|
||||
Creation response from the API.
|
||||
"""
|
||||
return self._client.create_directive(
|
||||
bank_id=bank_id,
|
||||
name=name,
|
||||
content=content,
|
||||
tags=tags,
|
||||
)
|
||||
|
||||
def list(self, bank_id: str, tags: list[str] | None = None) -> Any:
|
||||
"""List all directives for a bank.
|
||||
|
||||
Args:
|
||||
bank_id: The ID of the bank.
|
||||
tags: Optional filter by tags.
|
||||
|
||||
Returns:
|
||||
List of directives.
|
||||
"""
|
||||
return self._client.list_directives(bank_id=bank_id, tags=tags)
|
||||
|
||||
def get(self, bank_id: str, directive_id: str) -> Any:
|
||||
"""Get a specific directive.
|
||||
|
||||
Args:
|
||||
bank_id: The ID of the bank.
|
||||
directive_id: The ID of the directive.
|
||||
|
||||
Returns:
|
||||
The directive details.
|
||||
"""
|
||||
return self._client.get_directive(bank_id=bank_id, directive_id=directive_id)
|
||||
|
||||
def update(
|
||||
self,
|
||||
bank_id: str,
|
||||
directive_id: str,
|
||||
name: str | None = None,
|
||||
content: str | None = None,
|
||||
tags: list[str] | None = None,
|
||||
) -> Any:
|
||||
"""Update a directive.
|
||||
|
||||
Args:
|
||||
bank_id: The ID of the bank.
|
||||
directive_id: The ID of the directive to update.
|
||||
name: Optional new name.
|
||||
content: Optional new content.
|
||||
tags: Optional new tags list.
|
||||
|
||||
Returns:
|
||||
Update response from the API.
|
||||
"""
|
||||
return self._client.update_directive(
|
||||
bank_id=bank_id,
|
||||
directive_id=directive_id,
|
||||
name=name,
|
||||
content=content,
|
||||
tags=tags,
|
||||
)
|
||||
|
||||
def delete(self, bank_id: str, directive_id: str) -> Any:
|
||||
"""Delete a directive.
|
||||
|
||||
Args:
|
||||
bank_id: The ID of the bank.
|
||||
directive_id: The ID of the directive to delete.
|
||||
|
||||
Returns:
|
||||
Deletion response from the API.
|
||||
"""
|
||||
return self._client.delete_directive(bank_id=bank_id, directive_id=directive_id)
|
||||
|
||||
|
||||
class MemoriesAPI:
|
||||
"""Namespace for memory operations.
|
||||
|
||||
Provides methods to query and retrieve stored memories.
|
||||
"""
|
||||
|
||||
def __init__(self, client: Hindsight):
|
||||
self._client = client
|
||||
|
||||
def list(
|
||||
self,
|
||||
bank_id: str,
|
||||
type: str | None = None,
|
||||
search_query: str | None = None,
|
||||
limit: int = 100,
|
||||
offset: int = 0,
|
||||
) -> Any:
|
||||
"""List memories in a bank.
|
||||
|
||||
Args:
|
||||
bank_id: The ID of the bank to query.
|
||||
type: Optional filter by memory type.
|
||||
search_query: Optional search query for filtering.
|
||||
limit: Maximum number of results to return (default: 100).
|
||||
offset: Number of results to skip for pagination (default: 0).
|
||||
|
||||
Returns:
|
||||
List of memories matching the criteria.
|
||||
"""
|
||||
return self._client.list_memories(
|
||||
bank_id=bank_id,
|
||||
type=type,
|
||||
search_query=search_query,
|
||||
limit=limit,
|
||||
offset=offset,
|
||||
)
|
||||
|
||||
|
||||
class HindsightClient(Hindsight):
|
||||
"""
|
||||
Enhanced Hindsight client with organized API namespaces.
|
||||
|
||||
This wrapper extends the auto-generated Hindsight client with organized
|
||||
access to different parts of the API through namespaces.
|
||||
|
||||
Example:
|
||||
```python
|
||||
from hindsight import HindsightClient
|
||||
|
||||
client = HindsightClient(base_url="http://localhost:8888")
|
||||
|
||||
# Core operations (inherited from Hindsight)
|
||||
client.retain(bank_id="test", content="Hello")
|
||||
results = client.recall(bank_id="test", query="Hello")
|
||||
|
||||
# Organized API access through namespaces
|
||||
client.banks.create(bank_id="test", name="Test Bank")
|
||||
models = client.mental_models.list(bank_id="test")
|
||||
directives = client.directives.list(bank_id="test")
|
||||
memories = client.memories.list(bank_id="test")
|
||||
```
|
||||
|
||||
Attributes:
|
||||
banks: Namespace for bank management operations.
|
||||
mental_models: Namespace for mental model operations.
|
||||
directives: Namespace for directive operations.
|
||||
memories: Namespace for memory listing operations.
|
||||
"""
|
||||
|
||||
def __init__(self, *args: Any, **kwargs: Any) -> None:
|
||||
super().__init__(*args, **kwargs)
|
||||
self._banks_namespace: BanksAPI | None = None
|
||||
self._mental_models_namespace: MentalModelsAPI | None = None
|
||||
self._directives_namespace: DirectivesAPI | None = None
|
||||
self._memories_namespace: MemoriesAPI | None = None
|
||||
|
||||
@property
|
||||
def banks(self) -> BanksAPI:
|
||||
"""Access bank management operations.
|
||||
|
||||
Returns:
|
||||
BanksAPI instance for bank operations.
|
||||
"""
|
||||
if self._banks_namespace is None:
|
||||
self._banks_namespace = BanksAPI(self)
|
||||
return self._banks_namespace
|
||||
|
||||
@property
|
||||
def mental_models(self) -> MentalModelsAPI:
|
||||
"""Access mental model operations.
|
||||
|
||||
Returns:
|
||||
MentalModelsAPI instance for mental model operations.
|
||||
"""
|
||||
if self._mental_models_namespace is None:
|
||||
self._mental_models_namespace = MentalModelsAPI(self)
|
||||
return self._mental_models_namespace
|
||||
|
||||
@property
|
||||
def directives(self) -> DirectivesAPI:
|
||||
"""Access directive operations.
|
||||
|
||||
Returns:
|
||||
DirectivesAPI instance for directive operations.
|
||||
"""
|
||||
if self._directives_namespace is None:
|
||||
self._directives_namespace = DirectivesAPI(self)
|
||||
return self._directives_namespace
|
||||
|
||||
@property
|
||||
def memories(self) -> MemoriesAPI:
|
||||
"""Access memory listing operations.
|
||||
|
||||
Returns:
|
||||
MemoriesAPI instance for memory operations.
|
||||
"""
|
||||
if self._memories_namespace is None:
|
||||
self._memories_namespace = MemoriesAPI(self)
|
||||
return self._memories_namespace
|
||||
@@ -4,18 +4,18 @@ build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "hindsight-all"
|
||||
version = "0.4.15"
|
||||
version = "0.4.18"
|
||||
description = "Hindsight: Agent Memory That Works Like Human Memory - All-in-One Bundle"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.11"
|
||||
dependencies = [
|
||||
"hindsight-api>=0.0.7",
|
||||
"hindsight-api-slim[all]>=0.4.17",
|
||||
"hindsight-client>=0.0.7",
|
||||
"hindsight-embed>=0.1.0",
|
||||
]
|
||||
|
||||
[tool.uv.sources]
|
||||
hindsight-api = { workspace = true }
|
||||
hindsight-api-slim = { workspace = true }
|
||||
hindsight-client = { workspace = true }
|
||||
hindsight-embed = { workspace = true }
|
||||
|
||||
@@ -0,0 +1,137 @@
|
||||
# Hindsight API
|
||||
|
||||
**Memory System for AI Agents** — Temporal + Semantic + Entity Memory Architecture using PostgreSQL with pgvector.
|
||||
|
||||
Hindsight gives AI agents persistent memory that works like human memory: it stores facts, tracks entities and relationships, handles temporal reasoning ("what happened last spring?"), and forms opinions based on configurable disposition traits.
|
||||
|
||||
## Installation
|
||||
|
||||
```bash
|
||||
pip install hindsight-api
|
||||
```
|
||||
|
||||
## Quick Start
|
||||
|
||||
### Run the Server
|
||||
|
||||
```bash
|
||||
# Set your LLM provider
|
||||
export HINDSIGHT_API_LLM_PROVIDER=openai
|
||||
export HINDSIGHT_API_LLM_API_KEY=sk-xxxxxxxxxxxx
|
||||
|
||||
# Start the server (uses embedded PostgreSQL by default)
|
||||
hindsight-api
|
||||
```
|
||||
|
||||
The server starts at http://localhost:8888 with:
|
||||
- REST API for memory operations
|
||||
- MCP server at `/mcp` for tool-use integration
|
||||
|
||||
### Use the Python API
|
||||
|
||||
```python
|
||||
from hindsight_api import MemoryEngine
|
||||
|
||||
# Create and initialize the memory engine
|
||||
memory = MemoryEngine()
|
||||
await memory.initialize()
|
||||
|
||||
# Create a memory bank for your agent
|
||||
bank = await memory.create_memory_bank(
|
||||
name="my-assistant",
|
||||
background="A helpful coding assistant"
|
||||
)
|
||||
|
||||
# Store a memory
|
||||
await memory.retain(
|
||||
memory_bank_id=bank.id,
|
||||
content="The user prefers Python for data science projects"
|
||||
)
|
||||
|
||||
# Recall memories
|
||||
results = await memory.recall(
|
||||
memory_bank_id=bank.id,
|
||||
query="What programming language does the user prefer?"
|
||||
)
|
||||
|
||||
# Reflect with reasoning
|
||||
response = await memory.reflect(
|
||||
memory_bank_id=bank.id,
|
||||
query="Should I recommend Python or R for this ML project?"
|
||||
)
|
||||
```
|
||||
|
||||
## CLI Options
|
||||
|
||||
```bash
|
||||
hindsight-api --help
|
||||
|
||||
# Common options
|
||||
hindsight-api --port 9000 # Custom port (default: 8888)
|
||||
hindsight-api --host 127.0.0.1 # Bind to localhost only
|
||||
hindsight-api --workers 4 # Multiple worker processes
|
||||
hindsight-api --log-level debug # Verbose logging
|
||||
```
|
||||
|
||||
## Configuration
|
||||
|
||||
Configure via environment variables:
|
||||
|
||||
| Variable | Description | Default |
|
||||
|----------|-------------|---------|
|
||||
| `HINDSIGHT_API_DATABASE_URL` | PostgreSQL connection string | `pg0` (embedded) |
|
||||
| `HINDSIGHT_API_LLM_PROVIDER` | `openai`, `anthropic`, `gemini`, `groq`, `ollama`, `lmstudio` | `openai` |
|
||||
| `HINDSIGHT_API_LLM_API_KEY` | API key for LLM provider | - |
|
||||
| `HINDSIGHT_API_LLM_MODEL` | Model name | `gpt-4o-mini` |
|
||||
| `HINDSIGHT_API_HOST` | Server bind address | `0.0.0.0` |
|
||||
| `HINDSIGHT_API_PORT` | Server port | `8888` |
|
||||
|
||||
### Example with External PostgreSQL
|
||||
|
||||
```bash
|
||||
export HINDSIGHT_API_DATABASE_URL=postgresql://user:pass@localhost:5432/hindsight
|
||||
export HINDSIGHT_API_LLM_PROVIDER=groq
|
||||
export HINDSIGHT_API_LLM_API_KEY=gsk_xxxxxxxxxxxx
|
||||
|
||||
hindsight-api
|
||||
```
|
||||
|
||||
## Docker
|
||||
|
||||
```bash
|
||||
docker run --rm -it -p 8888:8888 \
|
||||
-e HINDSIGHT_API_LLM_API_KEY=$OPENAI_API_KEY \
|
||||
-v $HOME/.hindsight-docker:/home/hindsight/.pg0 \
|
||||
ghcr.io/vectorize-io/hindsight:latest
|
||||
```
|
||||
|
||||
## MCP Server
|
||||
|
||||
For local MCP integration without running the full API server:
|
||||
|
||||
```bash
|
||||
hindsight-local-mcp
|
||||
```
|
||||
|
||||
This runs a stdio-based MCP server that can be used directly with MCP-compatible clients.
|
||||
|
||||
## Key Features
|
||||
|
||||
- **Multi-Strategy Retrieval (TEMPR)** — Semantic, keyword, graph, and temporal search combined with RRF fusion
|
||||
- **Entity Graph** — Automatic entity extraction and relationship tracking
|
||||
- **Temporal Reasoning** — Native support for time-based queries
|
||||
- **Disposition Traits** — Configurable skepticism, literalism, and empathy influence opinion formation
|
||||
- **Three Memory Types** — World facts, bank actions, and formed opinions with confidence scores
|
||||
|
||||
## Documentation
|
||||
|
||||
Full documentation: [https://hindsight.vectorize.io](https://hindsight.vectorize.io)
|
||||
|
||||
- [Installation Guide](https://hindsight.vectorize.io/developer/installation)
|
||||
- [Configuration Reference](https://hindsight.vectorize.io/developer/configuration)
|
||||
- [API Reference](https://hindsight.vectorize.io/api-reference)
|
||||
- [Python SDK](https://hindsight.vectorize.io/sdks/python)
|
||||
|
||||
## License
|
||||
|
||||
Apache 2.0
|
||||
+1
-1
@@ -46,4 +46,4 @@ __all__ = [
|
||||
"RemoteTEICrossEncoder",
|
||||
"LLMConfig",
|
||||
]
|
||||
__version__ = "0.4.15"
|
||||
__version__ = "0.4.18"
|
||||
+81
-9
@@ -14,7 +14,8 @@ from typing import Any
|
||||
import asyncpg
|
||||
import typer
|
||||
|
||||
from ..config import HindsightConfig
|
||||
from ..config import DEFAULT_DATABASE_SCHEMA, HindsightConfig
|
||||
from ..extensions import TenantExtension, load_extension
|
||||
from ..pg0 import parse_pg0_url, resolve_database_url
|
||||
|
||||
|
||||
@@ -214,20 +215,81 @@ def restore(
|
||||
typer.echo("Restore complete")
|
||||
|
||||
|
||||
async def _run_migration(db_url: str, schema: str = "public") -> None:
|
||||
"""Resolve database URL and run migrations."""
|
||||
from ..migrations import run_migrations
|
||||
async def _run_migration(
|
||||
db_url: str,
|
||||
schema: str | None = None,
|
||||
base_schema: str = DEFAULT_DATABASE_SCHEMA,
|
||||
embedding_dimension: int | None = None,
|
||||
) -> list[str]:
|
||||
"""Resolve database URL and run migrations for one schema or all discovered schemas."""
|
||||
from ..migrations import (
|
||||
ensure_embedding_dimension,
|
||||
ensure_text_search_extension,
|
||||
ensure_vector_extension,
|
||||
run_migrations,
|
||||
)
|
||||
|
||||
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)
|
||||
run_migrations(resolved_url, schema=schema)
|
||||
|
||||
config = HindsightConfig.from_env()
|
||||
if schema:
|
||||
schemas = [schema]
|
||||
else:
|
||||
tenant_extension = load_extension("TENANT", TenantExtension)
|
||||
|
||||
schemas = [base_schema or DEFAULT_DATABASE_SCHEMA]
|
||||
if tenant_extension:
|
||||
tenants = await tenant_extension.list_tenants()
|
||||
schemas.extend(tenant.schema for tenant in tenants if tenant.schema)
|
||||
|
||||
# Preserve order while removing duplicates.
|
||||
schemas = list(dict.fromkeys(schemas))
|
||||
|
||||
for schema in schemas:
|
||||
run_migrations(resolved_url, schema=schema)
|
||||
|
||||
if embedding_dimension is not None:
|
||||
for schema in schemas:
|
||||
ensure_embedding_dimension(
|
||||
resolved_url,
|
||||
embedding_dimension,
|
||||
schema=schema,
|
||||
vector_extension=config.vector_extension,
|
||||
)
|
||||
|
||||
for schema in schemas:
|
||||
ensure_vector_extension(
|
||||
resolved_url,
|
||||
vector_extension=config.vector_extension,
|
||||
schema=schema,
|
||||
)
|
||||
|
||||
for schema in schemas:
|
||||
ensure_text_search_extension(
|
||||
resolved_url,
|
||||
text_search_extension=config.text_search_extension,
|
||||
schema=schema,
|
||||
)
|
||||
|
||||
return schemas
|
||||
|
||||
|
||||
@app.command(name="run-db-migration")
|
||||
def run_db_migration(
|
||||
schema: str = typer.Option("public", "--schema", "-s", help="Database schema to run migrations on"),
|
||||
schema: str | None = typer.Option(
|
||||
None,
|
||||
"--schema",
|
||||
"-s",
|
||||
help="Database schema to run migrations on. If omitted, migrate the base schema and all discovered tenant schemas.",
|
||||
),
|
||||
embedding_dimension: int | None = typer.Option(
|
||||
None,
|
||||
"--embedding-dimension",
|
||||
help="Expected embedding dimension to enforce after migrations. Omit to skip dimension sync.",
|
||||
),
|
||||
):
|
||||
"""Run database migrations to the latest version."""
|
||||
config = HindsightConfig.from_env()
|
||||
@@ -237,11 +299,21 @@ def run_db_migration(
|
||||
typer.echo("Set HINDSIGHT_API_DATABASE_URL environment variable.", err=True)
|
||||
raise typer.Exit(1)
|
||||
|
||||
typer.echo(f"Running database migrations (schema: {schema})...")
|
||||
if schema:
|
||||
typer.echo(f"Running database migrations for schema: {schema}...")
|
||||
else:
|
||||
typer.echo("Running database migrations for base schema and all discovered tenant schemas...")
|
||||
|
||||
asyncio.run(_run_migration(config.database_url, schema))
|
||||
schemas = asyncio.run(
|
||||
_run_migration(
|
||||
config.database_url,
|
||||
schema=schema,
|
||||
base_schema=config.database_schema,
|
||||
embedding_dimension=embedding_dimension,
|
||||
)
|
||||
)
|
||||
|
||||
typer.echo("Database migrations completed successfully")
|
||||
typer.echo(f"Database migrations completed successfully for {len(schemas)} schema(s)")
|
||||
|
||||
|
||||
async def _decommission_worker(db_url: str, worker_id: str, schema: str = "public") -> int:
|
||||
+1
-1
@@ -34,7 +34,7 @@ def upgrade() -> None:
|
||||
# Create file_storage table (minimal: just key + data)
|
||||
op.execute(
|
||||
f"""
|
||||
CREATE TABLE {schema}file_storage (
|
||||
CREATE TABLE IF NOT EXISTS {schema}file_storage (
|
||||
storage_key TEXT PRIMARY KEY,
|
||||
data BYTEA NOT NULL
|
||||
)
|
||||
+54
@@ -0,0 +1,54 @@
|
||||
"""Add GIN index on source_memory_ids for observation lookup performance
|
||||
|
||||
Without this index, queries using the array overlap operator (&&) or array
|
||||
containment (@>) on source_memory_ids require a full sequential scan over all
|
||||
observation memory_units. At ~77k observations this was measured at 45ms per
|
||||
query, becoming a bottleneck during consolidation recall (57-64s timeouts) and
|
||||
user recall (18-27s average).
|
||||
|
||||
The GIN index reduces these queries to index scans: 45ms → 0.049ms (927x
|
||||
speedup). Recall dropped from 18-27s to ~6s, and consolidation recall
|
||||
stabilised from timeout to ~15s.
|
||||
|
||||
Created with CONCURRENTLY so the migration does not block reads or writes.
|
||||
CONCURRENTLY requires running outside a transaction block, so the migration
|
||||
emits an explicit COMMIT before the statement and uses IF NOT EXISTS for
|
||||
idempotency.
|
||||
|
||||
Revision ID: a2b3c4d5e6f8
|
||||
Revises: f7g8h9i0j1k2
|
||||
Create Date: 2026-03-04
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
revision: str = "a2b3c4d5e6f8"
|
||||
down_revision: str | Sequence[str] | None = "f7g8h9i0j1k2"
|
||||
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()
|
||||
|
||||
# CREATE INDEX CONCURRENTLY cannot run inside a transaction block.
|
||||
# Commit the current Alembic transaction first.
|
||||
op.execute("COMMIT")
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_units_source_memory_ids "
|
||||
f"ON {schema}memory_units USING GIN (source_memory_ids) "
|
||||
f"WHERE source_memory_ids IS NOT NULL"
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
op.execute("COMMIT")
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_units_source_memory_ids")
|
||||
+52
@@ -0,0 +1,52 @@
|
||||
"""Add consolidation_failed_at column to memory_units for tracking persistent LLM failures.
|
||||
|
||||
When all LLM retries are exhausted on a single-memory batch, the memory is marked
|
||||
with consolidation_failed_at instead of consolidated_at, so it is not silently lost
|
||||
and can be retried later via the API.
|
||||
|
||||
Revision ID: a3b4c5d6e7f8
|
||||
Revises: g7h8i9j0k1l2
|
||||
Create Date: 2026-03-17
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
revision: str = "a3b4c5d6e7f8"
|
||||
down_revision: str | Sequence[str] | None = "g7h8i9j0k1l2"
|
||||
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}memory_units
|
||||
ADD COLUMN IF NOT EXISTS consolidation_failed_at TIMESTAMPTZ DEFAULT NULL
|
||||
"""
|
||||
)
|
||||
|
||||
# Index to efficiently query memories that failed consolidation for a given bank
|
||||
op.execute(
|
||||
f"""
|
||||
CREATE INDEX IF NOT EXISTS idx_memory_units_consolidation_failed
|
||||
ON {schema}memory_units (bank_id, consolidation_failed_at)
|
||||
WHERE consolidation_failed_at IS NOT NULL AND fact_type IN ('experience', 'world')
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema}idx_memory_units_consolidation_failed")
|
||||
op.execute(f"ALTER TABLE {schema}memory_units DROP COLUMN IF EXISTS consolidation_failed_at")
|
||||
+30
@@ -0,0 +1,30 @@
|
||||
"""Add history column to mental_models
|
||||
|
||||
Revision ID: c3d4e5f6g7h8
|
||||
Revises: a2b3c4d5e6f7, a2b3c4d5e6f8
|
||||
Create Date: 2026-03-06
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
revision: str = "c3d4e5f6g7h8"
|
||||
down_revision: str | Sequence[str] | None = ("a2b3c4d5e6f7", "a2b3c4d5e6f8")
|
||||
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 history JSONB DEFAULT '[]'::jsonb")
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
op.execute(f"ALTER TABLE {schema}mental_models DROP COLUMN IF EXISTS history")
|
||||
+53
@@ -0,0 +1,53 @@
|
||||
"""Recreate idx_memory_units_source_memory_ids GIN index with fastupdate=off
|
||||
|
||||
GIN indexes use a "fastupdate" pending list by default: small writes are
|
||||
buffered there and flushed to the main GIN tree in bulk. Flushing requires
|
||||
AccessExclusiveLock on the index. Under high insert concurrency (e.g. 8
|
||||
parallel pytest-xdist workers all calling retain_async) two transactions can
|
||||
each trigger a flush simultaneously and deadlock.
|
||||
|
||||
Disabling fastupdate makes every insert write directly to the GIN tree
|
||||
(slightly slower per insert, but no pending-list lock cycles).
|
||||
|
||||
Revision ID: d4e5f6g7h8i9
|
||||
Revises: d5e6f7a8b9c0
|
||||
Create Date: 2026-03-11
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
revision: str = "d4e5f6g7h8i9"
|
||||
down_revision: str | Sequence[str] | None = "d5e6f7a8b9c0"
|
||||
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()
|
||||
# DROP + CREATE CONCURRENTLY must run outside a transaction block.
|
||||
op.execute("COMMIT")
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_units_source_memory_ids")
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_units_source_memory_ids "
|
||||
f"ON {schema}memory_units USING GIN (source_memory_ids) "
|
||||
f"WITH (fastupdate=off) "
|
||||
f"WHERE source_memory_ids IS NOT NULL"
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
op.execute("COMMIT")
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_units_source_memory_ids")
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_units_source_memory_ids "
|
||||
f"ON {schema}memory_units USING GIN (source_memory_ids) "
|
||||
f"WHERE source_memory_ids IS NOT NULL"
|
||||
)
|
||||
+131
@@ -0,0 +1,131 @@
|
||||
"""Add internal_id to banks and per-(bank, fact_type) partial HNSW indexes
|
||||
|
||||
Revision ID: d5e6f7a8b9c0
|
||||
Revises: a3b4c5d6e7f8
|
||||
Create Date: 2026-03-11
|
||||
|
||||
This migration:
|
||||
1. Adds internal_id UUID column to banks (stable identifier for index naming)
|
||||
2. Drops the global HNSW index (competes with per-bank partial indexes)
|
||||
3. Creates per-(bank_id, fact_type) partial HNSW indexes for all existing banks
|
||||
(new banks get indexes created at bank-creation time via bank_utils.create_bank_hnsw_indexes)
|
||||
|
||||
Why per-(bank, fact_type) indexes:
|
||||
- fact_type-only partial indexes are never chosen by the planner when bank_id is in the WHERE
|
||||
clause, because the idx_memory_units_bank_id B-tree index always wins at planning time.
|
||||
- Per-(bank, fact_type) partial indexes have both predicates matching → planner selects them.
|
||||
- The global HNSW index competes for larger partitions (world, observation) and must be dropped.
|
||||
|
||||
For large deployments, create indexes CONCURRENTLY before running this migration:
|
||||
SELECT internal_id, bank_id FROM banks;
|
||||
-- for each bank and each fact_type in (world, experience, observation):
|
||||
CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_mu_emb_{ft}_{uid16}
|
||||
ON memory_units USING hnsw (embedding vector_cosine_ops)
|
||||
WHERE fact_type = '{ft}' AND bank_id = '{bank_id}';
|
||||
DROP INDEX CONCURRENTLY IF EXISTS idx_memory_units_embedding;
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
from sqlalchemy import text
|
||||
|
||||
revision: str = "d5e6f7a8b9c0"
|
||||
down_revision: str | Sequence[str] | None = "c3d4e5f6g7h8"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
_HNSW_FACT_TYPES: dict[str, str] = {
|
||||
"world": "worl",
|
||||
"experience": "expr",
|
||||
"observation": "obsv",
|
||||
}
|
||||
|
||||
|
||||
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()
|
||||
|
||||
# 1. Add internal_id column to banks
|
||||
op.execute(
|
||||
f"ALTER TABLE {schema}banks ADD COLUMN IF NOT EXISTS internal_id UUID DEFAULT gen_random_uuid() NOT NULL"
|
||||
)
|
||||
op.execute(f"ALTER TABLE {schema}banks ADD CONSTRAINT banks_internal_id_unique UNIQUE (internal_id)")
|
||||
|
||||
# 2. Drop any fact_type-only partial HNSW indexes that may exist from prior migrations
|
||||
# (bank_id B-tree always wins over them when bank_id is in the WHERE clause)
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema}idx_mu_emb_world")
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema}idx_mu_emb_observation")
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema}idx_mu_emb_experience")
|
||||
|
||||
# 4. Drop global HNSW index (competes with per-bank partial indexes)
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema}idx_memory_units_embedding")
|
||||
|
||||
# 5. Create per-(bank, fact_type) partial HNSW indexes for all existing banks
|
||||
bind = op.get_bind()
|
||||
schema_name = context.config.get_main_option("target_schema")
|
||||
table_ref = f'"{schema_name}".memory_units' if schema_name else "memory_units"
|
||||
banks_ref = f'"{schema_name}".banks' if schema_name else "banks"
|
||||
|
||||
rows = bind.execute(text(f"SELECT bank_id, internal_id FROM {banks_ref}")).fetchall() # noqa: S608
|
||||
for row in rows:
|
||||
bank_id = row[0]
|
||||
internal_id = str(row[1]).replace("-", "")[:16]
|
||||
escaped_bank_id = bank_id.replace("'", "''")
|
||||
for ft, ft_short in _HNSW_FACT_TYPES.items():
|
||||
idx_name = f"idx_mu_emb_{ft_short}_{internal_id}"
|
||||
# Index name is schema-unqualified (indexes live in the schema of their table)
|
||||
bind.execute(
|
||||
text(
|
||||
f"CREATE INDEX IF NOT EXISTS {idx_name} "
|
||||
f"ON {table_ref} USING hnsw (embedding vector_cosine_ops) "
|
||||
f"WHERE fact_type = '{ft}' AND bank_id = '{escaped_bank_id}'"
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
|
||||
# Drop per-bank HNSW indexes (iterate existing banks)
|
||||
bind = op.get_bind()
|
||||
schema_name = context.config.get_main_option("target_schema")
|
||||
banks_ref = f'"{schema_name}".banks' if schema_name else "banks"
|
||||
|
||||
rows = bind.execute(text(f"SELECT internal_id FROM {banks_ref}")).fetchall() # noqa: S608
|
||||
for row in rows:
|
||||
internal_id = str(row[0]).replace("-", "")[:16]
|
||||
for ft_short in _HNSW_FACT_TYPES.values():
|
||||
idx_name = f"idx_mu_emb_{ft_short}_{internal_id}"
|
||||
bind.execute(text(f"DROP INDEX IF EXISTS {schema}{idx_name}"))
|
||||
|
||||
# Restore the global HNSW index
|
||||
table_ref = f'"{schema_name}".memory_units' if schema_name else "memory_units"
|
||||
op.execute(
|
||||
f"CREATE INDEX IF NOT EXISTS idx_memory_units_embedding ON {table_ref} USING hnsw (embedding vector_cosine_ops)"
|
||||
)
|
||||
|
||||
# Restore old fact_type-only partial indexes
|
||||
op.execute(
|
||||
f"CREATE INDEX IF NOT EXISTS idx_mu_emb_world "
|
||||
f"ON {table_ref} USING hnsw (embedding vector_cosine_ops) "
|
||||
f"WHERE fact_type = 'world'"
|
||||
)
|
||||
op.execute(
|
||||
f"CREATE INDEX IF NOT EXISTS idx_mu_emb_observation "
|
||||
f"ON {table_ref} USING hnsw (embedding vector_cosine_ops) "
|
||||
f"WHERE fact_type = 'observation'"
|
||||
)
|
||||
op.execute(
|
||||
f"CREATE INDEX IF NOT EXISTS idx_mu_emb_experience "
|
||||
f"ON {table_ref} USING hnsw (embedding vector_cosine_ops) "
|
||||
f"WHERE fact_type = 'experience'"
|
||||
)
|
||||
|
||||
# Drop internal_id column
|
||||
op.execute(f"ALTER TABLE {schema}banks DROP CONSTRAINT IF EXISTS banks_internal_id_unique")
|
||||
op.execute(f"ALTER TABLE {schema}banks DROP COLUMN IF EXISTS internal_id")
|
||||
@@ -0,0 +1,62 @@
|
||||
"""Add webhooks table and next_retry_at to async_operations.
|
||||
|
||||
Webhook deliveries are handled as async_operations tasks (operation_type='webhook_delivery')
|
||||
rather than a dedicated webhook_deliveries table.
|
||||
|
||||
Revision ID: e4f5a6b7c8d9
|
||||
Revises: d2e3f4a5b6c7
|
||||
Create Date: 2026-03-04
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
revision: str = "e4f5a6b7c8d9"
|
||||
down_revision: str | Sequence[str] | None = "d2e3f4a5b6c7"
|
||||
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"""
|
||||
CREATE TABLE IF NOT EXISTS {schema}webhooks (
|
||||
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
||||
bank_id TEXT,
|
||||
url TEXT NOT NULL,
|
||||
secret TEXT,
|
||||
event_types TEXT[] NOT NULL DEFAULT '{{}}',
|
||||
enabled BOOLEAN NOT NULL DEFAULT TRUE,
|
||||
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
|
||||
updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
|
||||
)
|
||||
"""
|
||||
)
|
||||
|
||||
# Index for bank-scoped webhook lookup
|
||||
op.execute(f"CREATE INDEX IF NOT EXISTS idx_webhooks_bank_id ON {schema}webhooks(bank_id)")
|
||||
|
||||
# Add next_retry_at to async_operations for task-owned retry scheduling
|
||||
op.execute(f"ALTER TABLE {schema}async_operations ADD COLUMN IF NOT EXISTS next_retry_at TIMESTAMPTZ NULL")
|
||||
|
||||
# Index for polling: status + next_retry_at
|
||||
op.execute(
|
||||
f"CREATE INDEX IF NOT EXISTS idx_async_operations_status_retry "
|
||||
f"ON {schema}async_operations(status, next_retry_at)"
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema}idx_async_operations_status_retry")
|
||||
op.execute(f"ALTER TABLE {schema}async_operations DROP COLUMN IF EXISTS next_retry_at")
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema}idx_webhooks_bank_id")
|
||||
op.execute(f"DROP TABLE IF EXISTS {schema}webhooks")
|
||||
+73
@@ -0,0 +1,73 @@
|
||||
"""Add CASCADE DELETE FK from async_operations and webhooks to banks.
|
||||
|
||||
When a bank is deleted, all its async_operations and webhooks rows are
|
||||
automatically deleted by the database. This ensures that any in-flight
|
||||
worker tasks detect the deletion via _check_op_alive() and abort early.
|
||||
|
||||
Revision ID: e5f6g7h8i9j0
|
||||
Revises: d4e5f6g7h8i9
|
||||
Create Date: 2026-03-11
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
revision: str = "e5f6g7h8i9j0"
|
||||
down_revision: str | Sequence[str] | None = "d4e5f6g7h8i9"
|
||||
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()
|
||||
|
||||
# Remove orphaned async_operations rows whose bank no longer exists
|
||||
# (can happen because there was no FK before this migration).
|
||||
op.execute(
|
||||
f"""
|
||||
DELETE FROM {schema}async_operations
|
||||
WHERE bank_id IS NOT NULL
|
||||
AND bank_id NOT IN (SELECT bank_id FROM {schema}banks)
|
||||
"""
|
||||
)
|
||||
|
||||
# Remove orphaned webhooks rows whose bank no longer exists.
|
||||
op.execute(
|
||||
f"""
|
||||
DELETE FROM {schema}webhooks
|
||||
WHERE bank_id IS NOT NULL
|
||||
AND bank_id NOT IN (SELECT bank_id FROM {schema}banks)
|
||||
"""
|
||||
)
|
||||
|
||||
# Add FK with ON DELETE CASCADE so that deleting a bank automatically
|
||||
# cleans up all its pending/processing operations and webhook configs.
|
||||
op.execute(
|
||||
f"""
|
||||
ALTER TABLE {schema}async_operations
|
||||
ADD CONSTRAINT fk_async_operations_bank_id
|
||||
FOREIGN KEY (bank_id) REFERENCES {schema}banks(bank_id)
|
||||
ON DELETE CASCADE
|
||||
"""
|
||||
)
|
||||
|
||||
op.execute(
|
||||
f"""
|
||||
ALTER TABLE {schema}webhooks
|
||||
ADD CONSTRAINT fk_webhooks_bank_id
|
||||
FOREIGN KEY (bank_id) REFERENCES {schema}banks(bank_id)
|
||||
ON DELETE CASCADE
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
op.execute(f"ALTER TABLE {schema}async_operations DROP CONSTRAINT IF EXISTS fk_async_operations_bank_id")
|
||||
op.execute(f"ALTER TABLE {schema}webhooks DROP CONSTRAINT IF EXISTS fk_webhooks_bank_id")
|
||||
+38
@@ -0,0 +1,38 @@
|
||||
"""chunk_fk_cascade_delete
|
||||
|
||||
Revision ID: f6g7h8i9j0k1
|
||||
Revises: e5f6g7h8i9j0
|
||||
Create Date: 2026-03-16 00:00:00.000000
|
||||
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import op
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = "f6g7h8i9j0k1"
|
||||
down_revision: str | Sequence[str] | None = "e5f6g7h8i9j0"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
"""Change memory_units.chunk_id FK from SET NULL to CASCADE.
|
||||
|
||||
When a document is deleted the CASCADE reaches chunks first; with SET NULL
|
||||
the memory_units rows survived with chunk_id = NULL, leaving ghost records.
|
||||
Switching to CASCADE ensures they are removed together with their chunk.
|
||||
"""
|
||||
op.drop_constraint("memory_units_chunk_fkey", "memory_units", type_="foreignkey")
|
||||
op.create_foreign_key(
|
||||
"memory_units_chunk_fkey", "memory_units", "chunks", ["chunk_id"], ["chunk_id"], ondelete="CASCADE"
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
"""Revert to SET NULL behaviour."""
|
||||
op.drop_constraint("memory_units_chunk_fkey", "memory_units", type_="foreignkey")
|
||||
op.create_foreign_key(
|
||||
"memory_units_chunk_fkey", "memory_units", "chunks", ["chunk_id"], ["chunk_id"], ondelete="SET NULL"
|
||||
)
|
||||
+33
@@ -0,0 +1,33 @@
|
||||
"""Add http_config JSONB column to webhooks table.
|
||||
|
||||
Stores HTTP delivery configuration (method, timeout, headers, params) as a
|
||||
single JSONB column rather than separate columns.
|
||||
|
||||
Revision ID: f7g8h9i0j1k2
|
||||
Revises: e4f5a6b7c8d9
|
||||
Create Date: 2026-03-04
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
revision: str = "f7g8h9i0j1k2"
|
||||
down_revision: str | Sequence[str] | None = "e4f5a6b7c8d9"
|
||||
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}webhooks ADD COLUMN IF NOT EXISTS http_config JSONB NOT NULL DEFAULT '{{}}'")
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
op.execute(f"ALTER TABLE {schema}webhooks DROP COLUMN IF EXISTS http_config")
|
||||
+71
@@ -0,0 +1,71 @@
|
||||
"""backsweep_orphan_memory_units
|
||||
|
||||
Two-pass cleanup of memory_units rows that were never removed by earlier bugs:
|
||||
|
||||
Pass 1 — any fact_type, bank gone:
|
||||
memory_units whose bank_id no longer exists in banks. These accumulate when
|
||||
a bank is deleted without a proper cascade (no FK from memory_units to banks
|
||||
exists in the schema).
|
||||
|
||||
Pass 2 — observations only, all sources gone:
|
||||
observation rows whose bank still exists but every source_memory_id points
|
||||
to a deleted memory unit. These were left behind before PR #580 fixed the
|
||||
chunk FK cascade and before delete_document() called
|
||||
_delete_stale_observations_for_memories.
|
||||
|
||||
Revision ID: g7h8i9j0k1l2
|
||||
Revises: f6g7h8i9j0k1
|
||||
Create Date: 2026-03-16
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
revision: str = "g7h8i9j0k1l2"
|
||||
down_revision: str | Sequence[str] | None = "f6g7h8i9j0k1"
|
||||
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"
|
||||
banks = f"{schema}banks"
|
||||
|
||||
# Pass 1: delete all memory_units (any fact_type) whose bank no longer exists.
|
||||
# There is no FK from memory_units to banks, so these never cascade away.
|
||||
op.execute(
|
||||
f"""
|
||||
DELETE FROM {mu}
|
||||
WHERE NOT EXISTS (
|
||||
SELECT 1 FROM {banks} b WHERE b.bank_id = {mu}.bank_id
|
||||
)
|
||||
"""
|
||||
)
|
||||
|
||||
# Pass 2: delete orphaned observations whose bank still exists but 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.
|
||||
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
|
||||
+1
-1
@@ -35,7 +35,7 @@ def upgrade() -> None:
|
||||
|
||||
# Add GIN index for JSONB containment queries (@> operator)
|
||||
op.execute(f"""
|
||||
CREATE INDEX idx_async_operations_result_metadata
|
||||
CREATE INDEX IF NOT EXISTS idx_async_operations_result_metadata
|
||||
ON {schema}async_operations
|
||||
USING gin(result_metadata)
|
||||
""")
|
||||
+834
-64
File diff suppressed because it is too large
Load Diff
@@ -381,6 +381,15 @@ class MCPMiddleware:
|
||||
# Clear root_path since we're passing directly to the app
|
||||
new_scope["root_path"] = ""
|
||||
|
||||
# Ensure Accept header includes required MIME types for MCP SDK.
|
||||
# Some clients (e.g., Claude Code) don't send Accept, causing
|
||||
# the SDK to reject with 406 Not Acceptable.
|
||||
accept_header = self._get_header(new_scope, "accept")
|
||||
if not accept_header or "text/event-stream" not in accept_header:
|
||||
headers = [(k, v) for k, v in new_scope.get("headers", []) if k.lower() != b"accept"]
|
||||
headers.append((b"accept", b"application/json, text/event-stream"))
|
||||
new_scope["headers"] = headers
|
||||
|
||||
# Wrap send to rewrite the SSE endpoint URL to include bank_id if using path-based routing.
|
||||
# Only rewrite SSE (text/event-stream) responses to avoid corrupting tool results
|
||||
# that might contain the literal string "data: /messages".
|
||||
+123
-4
@@ -193,6 +193,7 @@ ENV_EMBEDDINGS_LITELLM_MODEL = "HINDSIGHT_API_EMBEDDINGS_LITELLM_MODEL"
|
||||
ENV_RERANKER_LITELLM_API_BASE = "HINDSIGHT_API_RERANKER_LITELLM_API_BASE"
|
||||
ENV_RERANKER_LITELLM_API_KEY = "HINDSIGHT_API_RERANKER_LITELLM_API_KEY"
|
||||
ENV_RERANKER_LITELLM_MODEL = "HINDSIGHT_API_RERANKER_LITELLM_MODEL"
|
||||
ENV_RERANKER_LITELLM_MAX_TOKENS_PER_DOC = "HINDSIGHT_API_RERANKER_LITELLM_MAX_TOKENS_PER_DOC"
|
||||
|
||||
# LiteLLM SDK configuration (direct API access, no proxy needed)
|
||||
ENV_EMBEDDINGS_LITELLM_SDK_API_KEY = "HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_API_KEY"
|
||||
@@ -211,6 +212,9 @@ ENV_RERANKER_LOCAL_MODEL = "HINDSIGHT_API_RERANKER_LOCAL_MODEL"
|
||||
ENV_RERANKER_LOCAL_FORCE_CPU = "HINDSIGHT_API_RERANKER_LOCAL_FORCE_CPU"
|
||||
ENV_RERANKER_LOCAL_MAX_CONCURRENT = "HINDSIGHT_API_RERANKER_LOCAL_MAX_CONCURRENT"
|
||||
ENV_RERANKER_LOCAL_TRUST_REMOTE_CODE = "HINDSIGHT_API_RERANKER_LOCAL_TRUST_REMOTE_CODE"
|
||||
ENV_RERANKER_LOCAL_FP16 = "HINDSIGHT_API_RERANKER_LOCAL_FP16"
|
||||
ENV_RERANKER_LOCAL_BUCKET_BATCHING = "HINDSIGHT_API_RERANKER_LOCAL_BUCKET_BATCHING"
|
||||
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"
|
||||
@@ -238,6 +242,7 @@ ENV_GRAPH_RETRIEVER = "HINDSIGHT_API_GRAPH_RETRIEVER"
|
||||
ENV_MPFP_TOP_K_NEIGHBORS = "HINDSIGHT_API_MPFP_TOP_K_NEIGHBORS"
|
||||
ENV_RECALL_MAX_CONCURRENT = "HINDSIGHT_API_RECALL_MAX_CONCURRENT"
|
||||
ENV_RECALL_CONNECTION_BUDGET = "HINDSIGHT_API_RECALL_CONNECTION_BUDGET"
|
||||
ENV_RECALL_MAX_QUERY_TOKENS = "HINDSIGHT_API_RECALL_MAX_QUERY_TOKENS"
|
||||
ENV_MENTAL_MODEL_REFRESH_CONCURRENCY = "HINDSIGHT_API_MENTAL_MODEL_REFRESH_CONCURRENCY"
|
||||
|
||||
# OpenTelemetry tracing configuration
|
||||
@@ -262,6 +267,7 @@ ENV_RETAIN_EXTRACT_CAUSAL_LINKS = "HINDSIGHT_API_RETAIN_EXTRACT_CAUSAL_LINKS"
|
||||
ENV_RETAIN_EXTRACTION_MODE = "HINDSIGHT_API_RETAIN_EXTRACTION_MODE"
|
||||
ENV_RETAIN_MISSION = "HINDSIGHT_API_RETAIN_MISSION"
|
||||
ENV_RETAIN_CUSTOM_INSTRUCTIONS = "HINDSIGHT_API_RETAIN_CUSTOM_INSTRUCTIONS"
|
||||
ENV_RETAIN_DEFAULT_STRATEGY = "HINDSIGHT_API_RETAIN_DEFAULT_STRATEGY"
|
||||
ENV_RETAIN_BATCH_TOKENS = "HINDSIGHT_API_RETAIN_BATCH_TOKENS"
|
||||
ENV_RETAIN_ENTITY_LOOKUP = "HINDSIGHT_API_RETAIN_ENTITY_LOOKUP"
|
||||
ENV_RETAIN_BATCH_ENABLED = "HINDSIGHT_API_RETAIN_BATCH_ENABLED"
|
||||
@@ -280,6 +286,7 @@ ENV_FILE_STORAGE_AZURE_CONTAINER = "HINDSIGHT_API_FILE_STORAGE_AZURE_CONTAINER"
|
||||
ENV_FILE_STORAGE_AZURE_ACCOUNT_NAME = "HINDSIGHT_API_FILE_STORAGE_AZURE_ACCOUNT_NAME"
|
||||
ENV_FILE_STORAGE_AZURE_ACCOUNT_KEY = "HINDSIGHT_API_FILE_STORAGE_AZURE_ACCOUNT_KEY"
|
||||
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_CONVERSION_MAX_BATCH_SIZE_MB = "HINDSIGHT_API_FILE_CONVERSION_MAX_BATCH_SIZE_MB"
|
||||
@@ -292,7 +299,19 @@ ENV_ENABLE_OBSERVATIONS = "HINDSIGHT_API_ENABLE_OBSERVATIONS"
|
||||
ENV_CONSOLIDATION_BATCH_SIZE = "HINDSIGHT_API_CONSOLIDATION_BATCH_SIZE"
|
||||
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_OBSERVATIONS_MISSION = "HINDSIGHT_API_OBSERVATIONS_MISSION"
|
||||
ENV_ENABLE_OBSERVATION_HISTORY = "HINDSIGHT_API_ENABLE_OBSERVATION_HISTORY"
|
||||
ENV_ENABLE_MENTAL_MODEL_HISTORY = "HINDSIGHT_API_ENABLE_MENTAL_MODEL_HISTORY"
|
||||
|
||||
# Webhook configuration (global, static - server-level only)
|
||||
ENV_WEBHOOK_URL = "HINDSIGHT_API_WEBHOOK_URL"
|
||||
ENV_WEBHOOK_SECRET = "HINDSIGHT_API_WEBHOOK_SECRET"
|
||||
ENV_WEBHOOK_EVENT_TYPES = "HINDSIGHT_API_WEBHOOK_EVENT_TYPES"
|
||||
ENV_WEBHOOK_DELIVERY_POLL_INTERVAL_SECONDS = "HINDSIGHT_API_WEBHOOK_DELIVERY_POLL_INTERVAL_SECONDS"
|
||||
|
||||
# Optimization flags
|
||||
ENV_SKIP_LLM_VERIFICATION = "HINDSIGHT_API_SKIP_LLM_VERIFICATION"
|
||||
@@ -337,6 +356,7 @@ PROVIDER_DEFAULT_MODELS = {
|
||||
"anthropic": "claude-haiku-4-5-20251001",
|
||||
"gemini": "gemini-2.5-flash",
|
||||
"groq": "openai/gpt-oss-120b",
|
||||
"minimax": "MiniMax-M2.5",
|
||||
"ollama": "gemma3:12b",
|
||||
"lmstudio": "local-model",
|
||||
"vertexai": "google/gemini-2.5-flash-lite",
|
||||
@@ -373,6 +393,9 @@ DEFAULT_RERANKER_LOCAL_MAX_CONCURRENT = 4 # Limit concurrent CPU-bound rerankin
|
||||
DEFAULT_RERANKER_LOCAL_TRUST_REMOTE_CODE = (
|
||||
False # Security: disabled by default, required for some models like jina-reranker-v2
|
||||
)
|
||||
DEFAULT_RERANKER_LOCAL_FP16 = False # FP16 inference: opt-in, faster on MPS/CUDA (not CPU)
|
||||
DEFAULT_RERANKER_LOCAL_BUCKET_BATCHING = False # Length-sorted bucket batching: opt-in, 36-54% speedup
|
||||
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_MAX_CANDIDATES = 300
|
||||
@@ -394,6 +417,7 @@ DEFAULT_TEXT_SEARCH_EXTENSION = "native" # Options: "native", "vchord", "pg_tex
|
||||
DEFAULT_LITELLM_API_BASE = "http://localhost:4000"
|
||||
DEFAULT_EMBEDDINGS_LITELLM_MODEL = "text-embedding-3-small"
|
||||
DEFAULT_RERANKER_LITELLM_MODEL = "cohere/rerank-english-v3.0"
|
||||
DEFAULT_RERANKER_LITELLM_MAX_TOKENS_PER_DOC: int | None = None
|
||||
|
||||
# LiteLLM SDK defaults
|
||||
DEFAULT_EMBEDDINGS_LITELLM_SDK_MODEL = "cohere/embed-english-v3.0"
|
||||
@@ -412,6 +436,7 @@ DEFAULT_GRAPH_RETRIEVER = "link_expansion" # Options: "link_expansion", "mpfp",
|
||||
DEFAULT_MPFP_TOP_K_NEIGHBORS = 20 # Fan-out limit per node in MPFP graph traversal
|
||||
DEFAULT_RECALL_MAX_CONCURRENT = 32 # Max concurrent recall operations per worker
|
||||
DEFAULT_RECALL_CONNECTION_BUDGET = 4 # Max concurrent DB connections per recall operation
|
||||
DEFAULT_RECALL_MAX_QUERY_TOKENS = 500 # Maximum tokens allowed in recall query
|
||||
DEFAULT_MENTAL_MODEL_REFRESH_CONCURRENCY = 8 # Max concurrent mental model refreshes
|
||||
|
||||
# Retain settings
|
||||
@@ -419,9 +444,11 @@ DEFAULT_RETAIN_MAX_COMPLETION_TOKENS = 64000 # Max tokens for fact extraction L
|
||||
DEFAULT_RETAIN_CHUNK_SIZE = 3000 # Max chars per chunk for fact extraction
|
||||
DEFAULT_RETAIN_EXTRACT_CAUSAL_LINKS = True # Extract causal links between facts
|
||||
DEFAULT_RETAIN_EXTRACTION_MODE = "concise" # Extraction mode: "concise", "verbose", or "custom"
|
||||
RETAIN_EXTRACTION_MODES = ("concise", "verbose", "custom") # Allowed extraction modes
|
||||
RETAIN_EXTRACTION_MODES = ("concise", "verbose", "custom", "verbatim", "chunks") # Allowed extraction modes
|
||||
DEFAULT_RETAIN_MISSION = None # Declarative spec of what to retain (injected into any extraction mode)
|
||||
DEFAULT_RETAIN_CUSTOM_INSTRUCTIONS = None # Custom extraction guidelines (only used when mode="custom")
|
||||
DEFAULT_RETAIN_DEFAULT_STRATEGY = None # Default strategy name (None = no strategy override)
|
||||
DEFAULT_RETAIN_STRATEGIES: dict | None = None # Named retain strategies (dict of name → config overrides)
|
||||
DEFAULT_RETAIN_BATCH_TOKENS = 10_000 # ~40KB of text # Max chars per sub-batch for async retain auto-splitting
|
||||
DEFAULT_RETAIN_ENTITY_LOOKUP = "trigram" # "full" or "trigram"
|
||||
DEFAULT_RETAIN_BATCH_ENABLED = False # Use LLM Batch API for fact extraction (only when async=True)
|
||||
@@ -429,7 +456,8 @@ DEFAULT_RETAIN_BATCH_POLL_INTERVAL_SECONDS = 60 # Batch API polling interval in
|
||||
|
||||
# File storage defaults
|
||||
DEFAULT_FILE_STORAGE_TYPE = "native" # PostgreSQL BYTEA storage
|
||||
DEFAULT_FILE_PARSER = "markitdown" # File parser to use (markitdown is the only supported parser)
|
||||
DEFAULT_FILE_PARSER = "markitdown" # Default parser fallback chain (comma-separated, e.g. "iris,markitdown")
|
||||
DEFAULT_FILE_PARSER_ALLOWLIST = None # Allowlist of parsers clients may request (None = all registered parsers)
|
||||
DEFAULT_FILE_CONVERSION_MAX_BATCH_SIZE_MB = 100 # Max total batch size in MB (all files combined)
|
||||
DEFAULT_FILE_CONVERSION_MAX_BATCH_SIZE = 10 # Max files per batch upload
|
||||
DEFAULT_ENABLE_FILE_UPLOAD_API = True # Enable file upload endpoint
|
||||
@@ -437,9 +465,17 @@ DEFAULT_FILE_DELETE_AFTER_RETAIN = True # Delete file bytes after retain (saves
|
||||
|
||||
# Observations defaults (consolidated knowledge from facts)
|
||||
DEFAULT_ENABLE_OBSERVATIONS = True # Observations enabled by default
|
||||
DEFAULT_ENABLE_OBSERVATION_HISTORY = True # Observation history tracking enabled by default
|
||||
DEFAULT_ENABLE_MENTAL_MODEL_HISTORY = True # Mental model history tracking enabled by default
|
||||
DEFAULT_CONSOLIDATION_BATCH_SIZE = 50 # Memories to load per batch (internal memory optimization)
|
||||
DEFAULT_CONSOLIDATION_LLM_BATCH_SIZE = 8 # Facts per LLM call (1 = no batching; >1 = batch mode)
|
||||
DEFAULT_CONSOLIDATION_MAX_TOKENS = 512 # Max tokens for recall when finding related observations
|
||||
DEFAULT_CONSOLIDATION_SOURCE_FACTS_MAX_TOKENS = (
|
||||
-1
|
||||
) # 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)
|
||||
)
|
||||
DEFAULT_OBSERVATIONS_MISSION = None # Declarative spec of what observations are for this bank
|
||||
|
||||
# Database migrations
|
||||
@@ -497,6 +533,12 @@ Use this tool PROACTIVELY to:
|
||||
# Default embedding dimension (used by initial migration, adjusted at runtime)
|
||||
EMBEDDING_DIMENSION = DEFAULT_EMBEDDING_DIMENSION
|
||||
|
||||
# Webhook configuration defaults
|
||||
DEFAULT_WEBHOOK_URL = None # None = no global webhook configured
|
||||
DEFAULT_WEBHOOK_SECRET = None # None = no signing
|
||||
DEFAULT_WEBHOOK_EVENT_TYPES = "consolidation.completed" # Comma-separated; default = all supported events
|
||||
DEFAULT_WEBHOOK_DELIVERY_POLL_INTERVAL_SECONDS = 30 # How often to poll for pending deliveries
|
||||
|
||||
|
||||
class JsonFormatter(logging.Formatter):
|
||||
"""JSON formatter for structured logging.
|
||||
@@ -528,6 +570,11 @@ class JsonFormatter(logging.Formatter):
|
||||
return json.dumps(log_entry)
|
||||
|
||||
|
||||
def _parse_str_list(value: str) -> list[str]:
|
||||
"""Parse a comma-separated string into a non-empty list of stripped tokens."""
|
||||
return [v.strip() for v in value.split(",") if v.strip()]
|
||||
|
||||
|
||||
def _validate_extraction_mode(mode: str) -> str:
|
||||
"""Validate and normalize extraction mode."""
|
||||
mode_lower = mode.lower()
|
||||
@@ -630,6 +677,9 @@ class HindsightConfig:
|
||||
reranker_local_force_cpu: bool
|
||||
reranker_local_max_concurrent: int
|
||||
reranker_local_trust_remote_code: bool
|
||||
reranker_local_fp16: bool
|
||||
reranker_local_bucket_batching: bool
|
||||
reranker_local_batch_size: int
|
||||
reranker_tei_url: str | None
|
||||
reranker_tei_batch_size: int
|
||||
reranker_tei_max_concurrent: int
|
||||
@@ -640,6 +690,7 @@ class HindsightConfig:
|
||||
reranker_litellm_api_base: str
|
||||
reranker_litellm_api_key: str | None
|
||||
reranker_litellm_model: str
|
||||
reranker_litellm_max_tokens_per_doc: int | None
|
||||
reranker_litellm_sdk_api_key: str | None
|
||||
reranker_litellm_sdk_model: str
|
||||
reranker_litellm_sdk_api_base: str | None
|
||||
@@ -661,6 +712,7 @@ class HindsightConfig:
|
||||
mpfp_top_k_neighbors: int
|
||||
recall_max_concurrent: int
|
||||
recall_connection_budget: int
|
||||
recall_max_query_tokens: int
|
||||
mental_model_refresh_concurrency: int
|
||||
|
||||
# Retain settings
|
||||
@@ -670,6 +722,8 @@ class HindsightConfig:
|
||||
retain_extraction_mode: str
|
||||
retain_mission: str | None
|
||||
retain_custom_instructions: str | None
|
||||
retain_default_strategy: str | None
|
||||
retain_strategies: dict | None
|
||||
retain_batch_tokens: int
|
||||
retain_batch_enabled: bool
|
||||
retain_batch_poll_interval_seconds: int
|
||||
@@ -687,7 +741,8 @@ class HindsightConfig:
|
||||
file_storage_azure_container: str | None # Azure container name (required for azure storage)
|
||||
file_storage_azure_account_name: str | None # Azure storage account name
|
||||
file_storage_azure_account_key: str | None # Azure storage account key
|
||||
file_parser: str # File parser to use (e.g., "markitdown", "iris")
|
||||
file_parser: list[str] # Ordered fallback chain of parsers (e.g. ["iris", "markitdown"])
|
||||
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_conversion_max_batch_size_mb: int # Max total batch size in MB (all files combined)
|
||||
@@ -697,9 +752,13 @@ class HindsightConfig:
|
||||
|
||||
# Observations settings (consolidated knowledge from facts)
|
||||
enable_observations: bool
|
||||
enable_observation_history: bool
|
||||
enable_mental_model_history: bool
|
||||
consolidation_batch_size: int
|
||||
consolidation_llm_batch_size: int
|
||||
consolidation_max_tokens: int
|
||||
consolidation_source_facts_max_tokens: int
|
||||
consolidation_source_facts_max_tokens_per_observation: int
|
||||
observations_mission: str | None
|
||||
|
||||
# Entity labels (controlled vocabulary of key:value classification labels extracted at retain time)
|
||||
@@ -750,6 +809,12 @@ class HindsightConfig:
|
||||
otel_service_name: str
|
||||
otel_deployment_environment: str
|
||||
|
||||
# Webhook configuration (static - server-level only, not per-bank)
|
||||
webhook_url: str | None # Global webhook URL (None = disabled)
|
||||
webhook_secret: str | None # HMAC signing secret (None = unsigned)
|
||||
webhook_event_types: list[str] # Event types to deliver globally
|
||||
webhook_delivery_poll_interval_seconds: int # How often the delivery worker polls
|
||||
|
||||
# Class-level sets for configuration categorization
|
||||
|
||||
# CREDENTIAL_FIELDS: Never exposed via API, never configurable per-tenant/bank
|
||||
@@ -789,11 +854,16 @@ class HindsightConfig:
|
||||
"retain_extraction_mode",
|
||||
"retain_mission",
|
||||
"retain_custom_instructions",
|
||||
"retain_default_strategy",
|
||||
"retain_strategies",
|
||||
# Entity labels (controlled vocabulary for entity classification)
|
||||
"entity_labels",
|
||||
"entities_allow_free_form",
|
||||
# Consolidation settings
|
||||
"enable_observations",
|
||||
"consolidation_llm_batch_size",
|
||||
"consolidation_source_facts_max_tokens",
|
||||
"consolidation_source_facts_max_tokens_per_observation",
|
||||
"observations_mission",
|
||||
# Reflect settings
|
||||
"reflect_mission",
|
||||
@@ -1038,6 +1108,15 @@ class HindsightConfig:
|
||||
ENV_RERANKER_LOCAL_TRUST_REMOTE_CODE, str(DEFAULT_RERANKER_LOCAL_TRUST_REMOTE_CODE)
|
||||
).lower()
|
||||
in ("true", "1"),
|
||||
reranker_local_fp16=os.getenv(ENV_RERANKER_LOCAL_FP16, str(DEFAULT_RERANKER_LOCAL_FP16)).lower()
|
||||
in ("true", "1"),
|
||||
reranker_local_bucket_batching=os.getenv(
|
||||
ENV_RERANKER_LOCAL_BUCKET_BATCHING, str(DEFAULT_RERANKER_LOCAL_BUCKET_BATCHING)
|
||||
).lower()
|
||||
in ("true", "1"),
|
||||
reranker_local_batch_size=int(
|
||||
os.getenv(ENV_RERANKER_LOCAL_BATCH_SIZE, str(DEFAULT_RERANKER_LOCAL_BATCH_SIZE))
|
||||
),
|
||||
reranker_tei_url=os.getenv(ENV_RERANKER_TEI_URL),
|
||||
reranker_tei_batch_size=int(os.getenv(ENV_RERANKER_TEI_BATCH_SIZE, str(DEFAULT_RERANKER_TEI_BATCH_SIZE))),
|
||||
reranker_tei_max_concurrent=int(
|
||||
@@ -1053,6 +1132,9 @@ class HindsightConfig:
|
||||
or os.getenv(ENV_LITELLM_API_BASE, DEFAULT_LITELLM_API_BASE),
|
||||
reranker_litellm_api_key=os.getenv(ENV_RERANKER_LITELLM_API_KEY) or os.getenv(ENV_LITELLM_API_KEY),
|
||||
reranker_litellm_model=os.getenv(ENV_RERANKER_LITELLM_MODEL, DEFAULT_RERANKER_LITELLM_MODEL),
|
||||
reranker_litellm_max_tokens_per_doc=int(v)
|
||||
if (v := os.getenv(ENV_RERANKER_LITELLM_MAX_TOKENS_PER_DOC))
|
||||
else DEFAULT_RERANKER_LITELLM_MAX_TOKENS_PER_DOC,
|
||||
# LiteLLM SDK reranker (direct API access)
|
||||
reranker_litellm_sdk_api_key=os.getenv(ENV_RERANKER_LITELLM_SDK_API_KEY),
|
||||
reranker_litellm_sdk_model=os.getenv(ENV_RERANKER_LITELLM_SDK_MODEL, DEFAULT_RERANKER_LITELLM_SDK_MODEL),
|
||||
@@ -1079,6 +1161,7 @@ class HindsightConfig:
|
||||
recall_connection_budget=int(
|
||||
os.getenv(ENV_RECALL_CONNECTION_BUDGET, str(DEFAULT_RECALL_CONNECTION_BUDGET))
|
||||
),
|
||||
recall_max_query_tokens=int(os.getenv(ENV_RECALL_MAX_QUERY_TOKENS, str(DEFAULT_RECALL_MAX_QUERY_TOKENS))),
|
||||
mental_model_refresh_concurrency=int(
|
||||
os.getenv(ENV_MENTAL_MODEL_REFRESH_CONCURRENCY, str(DEFAULT_MENTAL_MODEL_REFRESH_CONCURRENCY))
|
||||
),
|
||||
@@ -1099,6 +1182,8 @@ class HindsightConfig:
|
||||
),
|
||||
retain_mission=os.getenv(ENV_RETAIN_MISSION) or DEFAULT_RETAIN_MISSION,
|
||||
retain_custom_instructions=os.getenv(ENV_RETAIN_CUSTOM_INSTRUCTIONS) or DEFAULT_RETAIN_CUSTOM_INSTRUCTIONS,
|
||||
retain_default_strategy=os.getenv(ENV_RETAIN_DEFAULT_STRATEGY) or DEFAULT_RETAIN_DEFAULT_STRATEGY,
|
||||
retain_strategies=DEFAULT_RETAIN_STRATEGIES,
|
||||
retain_batch_tokens=int(os.getenv(ENV_RETAIN_BATCH_TOKENS, str(DEFAULT_RETAIN_BATCH_TOKENS))),
|
||||
retain_entity_lookup=os.getenv(ENV_RETAIN_ENTITY_LOOKUP, DEFAULT_RETAIN_ENTITY_LOOKUP),
|
||||
retain_batch_enabled=os.getenv(ENV_RETAIN_BATCH_ENABLED, str(DEFAULT_RETAIN_BATCH_ENABLED)).lower()
|
||||
@@ -1118,7 +1203,10 @@ class HindsightConfig:
|
||||
file_storage_azure_container=os.getenv(ENV_FILE_STORAGE_AZURE_CONTAINER) or None,
|
||||
file_storage_azure_account_name=os.getenv(ENV_FILE_STORAGE_AZURE_ACCOUNT_NAME) or None,
|
||||
file_storage_azure_account_key=os.getenv(ENV_FILE_STORAGE_AZURE_ACCOUNT_KEY) or None,
|
||||
file_parser=os.getenv(ENV_FILE_PARSER, DEFAULT_FILE_PARSER),
|
||||
file_parser=_parse_str_list(os.getenv(ENV_FILE_PARSER, DEFAULT_FILE_PARSER)),
|
||||
file_parser_allowlist=_parse_str_list(os.getenv(ENV_FILE_PARSER_ALLOWLIST))
|
||||
if os.getenv(ENV_FILE_PARSER_ALLOWLIST)
|
||||
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_conversion_max_batch_size_mb=int(
|
||||
@@ -1135,6 +1223,14 @@ class HindsightConfig:
|
||||
== "true",
|
||||
# Observations settings (consolidated knowledge from facts)
|
||||
enable_observations=os.getenv(ENV_ENABLE_OBSERVATIONS, str(DEFAULT_ENABLE_OBSERVATIONS)).lower() == "true",
|
||||
enable_observation_history=os.getenv(
|
||||
ENV_ENABLE_OBSERVATION_HISTORY, str(DEFAULT_ENABLE_OBSERVATION_HISTORY)
|
||||
).lower()
|
||||
== "true",
|
||||
enable_mental_model_history=os.getenv(
|
||||
ENV_ENABLE_MENTAL_MODEL_HISTORY, str(DEFAULT_ENABLE_MENTAL_MODEL_HISTORY)
|
||||
).lower()
|
||||
== "true",
|
||||
consolidation_batch_size=int(
|
||||
os.getenv(ENV_CONSOLIDATION_BATCH_SIZE, str(DEFAULT_CONSOLIDATION_BATCH_SIZE))
|
||||
),
|
||||
@@ -1144,6 +1240,15 @@ class HindsightConfig:
|
||||
consolidation_max_tokens=int(
|
||||
os.getenv(ENV_CONSOLIDATION_MAX_TOKENS, str(DEFAULT_CONSOLIDATION_MAX_TOKENS))
|
||||
),
|
||||
consolidation_source_facts_max_tokens=int(
|
||||
os.getenv(ENV_CONSOLIDATION_SOURCE_FACTS_MAX_TOKENS, str(DEFAULT_CONSOLIDATION_SOURCE_FACTS_MAX_TOKENS))
|
||||
),
|
||||
consolidation_source_facts_max_tokens_per_observation=int(
|
||||
os.getenv(
|
||||
ENV_CONSOLIDATION_SOURCE_FACTS_MAX_TOKENS_PER_OBSERVATION,
|
||||
str(DEFAULT_CONSOLIDATION_SOURCE_FACTS_MAX_TOKENS_PER_OBSERVATION),
|
||||
)
|
||||
),
|
||||
observations_mission=os.getenv(ENV_OBSERVATIONS_MISSION) or DEFAULT_OBSERVATIONS_MISSION,
|
||||
entity_labels=None,
|
||||
entities_allow_free_form=True,
|
||||
@@ -1187,6 +1292,20 @@ class HindsightConfig:
|
||||
otel_exporter_otlp_headers=os.getenv(ENV_OTEL_EXPORTER_OTLP_HEADERS) or None,
|
||||
otel_service_name=os.getenv(ENV_OTEL_SERVICE_NAME, DEFAULT_OTEL_SERVICE_NAME),
|
||||
otel_deployment_environment=os.getenv(ENV_OTEL_DEPLOYMENT_ENVIRONMENT, DEFAULT_OTEL_DEPLOYMENT_ENVIRONMENT),
|
||||
# Webhook configuration (static, server-level only)
|
||||
webhook_url=os.getenv(ENV_WEBHOOK_URL) or DEFAULT_WEBHOOK_URL,
|
||||
webhook_secret=os.getenv(ENV_WEBHOOK_SECRET) or DEFAULT_WEBHOOK_SECRET,
|
||||
webhook_event_types=[
|
||||
t.strip()
|
||||
for t in os.getenv(ENV_WEBHOOK_EVENT_TYPES, DEFAULT_WEBHOOK_EVENT_TYPES).split(",")
|
||||
if t.strip()
|
||||
],
|
||||
webhook_delivery_poll_interval_seconds=int(
|
||||
os.getenv(
|
||||
ENV_WEBHOOK_DELIVERY_POLL_INTERVAL_SECONDS,
|
||||
str(DEFAULT_WEBHOOK_DELIVERY_POLL_INTERVAL_SECONDS),
|
||||
)
|
||||
),
|
||||
)
|
||||
config.validate()
|
||||
return config
|
||||
+41
-1
@@ -10,7 +10,7 @@ multiple API servers.
|
||||
|
||||
import json
|
||||
import logging
|
||||
from dataclasses import asdict
|
||||
from dataclasses import asdict, replace
|
||||
from typing import Any
|
||||
|
||||
import asyncpg
|
||||
@@ -239,6 +239,14 @@ class ConfigResolver:
|
||||
logger.warning(f"Failed to check permissions for bank {bank_id}: {e}")
|
||||
# Continue without permission check (fail open for backward compatibility)
|
||||
|
||||
# Validate retain_strategies: reject empty string keys
|
||||
if "retain_strategies" in normalized_updates and normalized_updates["retain_strategies"]:
|
||||
empty_keys = [k for k in normalized_updates["retain_strategies"] if not str(k).strip()]
|
||||
if empty_keys:
|
||||
raise ValueError(
|
||||
"Strategy names must not be empty strings. Remove entries with empty names before saving."
|
||||
)
|
||||
|
||||
# Merge with existing config (JSONB || operator)
|
||||
async with self.pool.acquire() as conn:
|
||||
await conn.execute(
|
||||
@@ -273,3 +281,35 @@ class ConfigResolver:
|
||||
)
|
||||
|
||||
logger.info(f"Reset bank config for {bank_id} to defaults")
|
||||
|
||||
|
||||
def apply_strategy(config: HindsightConfig, strategy_name: str) -> HindsightConfig:
|
||||
"""
|
||||
Apply a named retain strategy's overrides on top of a resolved config.
|
||||
|
||||
A strategy is a named set of hierarchical field overrides stored in
|
||||
config.retain_strategies. Any field in _HIERARCHICAL_FIELDS can be
|
||||
overridden, including retain_extraction_mode, retain_chunk_size,
|
||||
entity_labels, entities_allow_free_form, etc.
|
||||
|
||||
Unknown strategy names log a warning and return config unchanged.
|
||||
Unknown or non-hierarchical fields in the strategy are silently ignored.
|
||||
"""
|
||||
strategies = config.retain_strategies or {}
|
||||
if strategy_name not in strategies:
|
||||
logger.warning(f"Unknown retain strategy '{strategy_name}', using resolved config as-is")
|
||||
return config
|
||||
|
||||
overrides = strategies[strategy_name]
|
||||
if not isinstance(overrides, dict):
|
||||
logger.warning(f"Retain strategy '{strategy_name}' is not a dict, skipping")
|
||||
return config
|
||||
|
||||
configurable = HindsightConfig.get_configurable_fields()
|
||||
filtered = {k: v for k, v in overrides.items() if k in configurable}
|
||||
|
||||
if not filtered:
|
||||
return config
|
||||
|
||||
logger.debug(f"Applying retain strategy '{strategy_name}': {list(filtered.keys())}")
|
||||
return replace(config, **filtered)
|
||||
+239
-105
@@ -9,6 +9,10 @@ Observations are stored in memory_units with fact_type='observation' and include
|
||||
- proof_count: Number of supporting memories
|
||||
- source_memory_ids: Array of memory UUIDs that contribute to this observation
|
||||
- history: JSONB tracking changes over time
|
||||
|
||||
NOTE: Observations are distinct from mental models (pinned reflections).
|
||||
- Observations: auto-generated bottom-up by this engine from raw facts (memory_units table, fact_type='observation')
|
||||
- Mental models: user-defined queries stored in the mental_models table, refreshed on demand via reflect
|
||||
"""
|
||||
|
||||
import json
|
||||
@@ -20,9 +24,10 @@ from datetime import datetime, timezone
|
||||
from itertools import combinations
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from pydantic import BaseModel
|
||||
from pydantic import BaseModel, field_validator
|
||||
|
||||
from ...config import get_config
|
||||
from ..llm_wrapper import sanitize_llm_output
|
||||
from ..memory_engine import fq_table
|
||||
from ..retain import embedding_utils
|
||||
from .prompts import build_batch_consolidation_prompt
|
||||
@@ -41,12 +46,22 @@ class _CreateAction(BaseModel):
|
||||
text: str
|
||||
source_fact_ids: list[str] # memory UUIDs from the NEW FACTS list
|
||||
|
||||
@field_validator("text", mode="before")
|
||||
@classmethod
|
||||
def sanitize_text(cls, v: str) -> str:
|
||||
return sanitize_llm_output(v) or ""
|
||||
|
||||
|
||||
class _UpdateAction(BaseModel):
|
||||
text: str
|
||||
observation_id: str # UUID of the existing observation to update
|
||||
source_fact_ids: list[str] # memory UUIDs from the NEW FACTS list
|
||||
|
||||
@field_validator("text", mode="before")
|
||||
@classmethod
|
||||
def sanitize_text(cls, v: str) -> str:
|
||||
return sanitize_llm_output(v) or ""
|
||||
|
||||
|
||||
class _DeleteAction(BaseModel):
|
||||
observation_id: str # UUID of the observation to remove
|
||||
@@ -65,6 +80,43 @@ class _BatchLLMResult:
|
||||
deletes: list[_DeleteAction] = field(default_factory=list)
|
||||
obs_count: int = 0
|
||||
prompt_chars: int = 0
|
||||
failed: bool = False
|
||||
|
||||
|
||||
@dataclass
|
||||
class _SourceAggregation:
|
||||
"""Fields inherited by an observation from its source memories."""
|
||||
|
||||
event_date: datetime | None
|
||||
occurred_start: datetime | None
|
||||
occurred_end: datetime | None
|
||||
mentioned_at: datetime | None
|
||||
tags: list[str]
|
||||
|
||||
|
||||
def _aggregate_source_fields(source_mems: list[dict[str, Any]], tags: list[str] | None = None) -> _SourceAggregation:
|
||||
"""Compute the observation fields inherited from a set of source memories.
|
||||
|
||||
Temporal aggregation rules:
|
||||
- ``event_date`` — earliest across sources (min)
|
||||
- ``occurred_start`` — earliest across sources (min)
|
||||
- ``occurred_end`` — latest across sources (max)
|
||||
- ``mentioned_at`` — latest across sources (max)
|
||||
|
||||
Fields remain ``None`` when no source memory carries that information, so
|
||||
observations are never stamped with an artificial timestamp.
|
||||
|
||||
``tags`` defaults to those of the first source memory when not explicitly
|
||||
provided (all memories in a consolidation batch share the same tag set).
|
||||
"""
|
||||
effective_tags = tags if tags is not None else (source_mems[0].get("tags") or [] if source_mems else [])
|
||||
return _SourceAggregation(
|
||||
event_date=_min_date(m.get("event_date") for m in source_mems),
|
||||
occurred_start=_min_date(m.get("occurred_start") for m in source_mems),
|
||||
occurred_end=_max_date(m.get("occurred_end") for m in source_mems),
|
||||
mentioned_at=_max_date(m.get("mentioned_at") for m in source_mems),
|
||||
tags=effective_tags,
|
||||
)
|
||||
|
||||
|
||||
class ConsolidationPerfLog:
|
||||
@@ -110,6 +162,7 @@ async def run_consolidation_job(
|
||||
memory_engine: "MemoryEngine",
|
||||
bank_id: str,
|
||||
request_context: "RequestContext",
|
||||
operation_id: str | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Run consolidation job for a bank.
|
||||
@@ -167,6 +220,7 @@ async def run_consolidation_job(
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE bank_id = $1
|
||||
AND consolidated_at IS NULL
|
||||
AND consolidation_failed_at IS NULL
|
||||
AND fact_type IN ('experience', 'world')
|
||||
""",
|
||||
bank_id,
|
||||
@@ -180,13 +234,15 @@ async def run_consolidation_job(
|
||||
perf.log(f"[1] Found {total_count} pending memories to consolidate")
|
||||
|
||||
# Process each memory with individual commits for crash recovery
|
||||
stats = {
|
||||
stats: dict[str, int] = {
|
||||
"memories_processed": 0,
|
||||
"observations_created": 0,
|
||||
"observations_updated": 0,
|
||||
"observations_merged": 0,
|
||||
"observations_deleted": 0,
|
||||
"actions_executed": 0,
|
||||
"skipped": 0,
|
||||
"memories_failed": 0,
|
||||
}
|
||||
|
||||
# Track all unique tags from consolidated memories for mental model refresh filtering
|
||||
@@ -204,6 +260,7 @@ async def run_consolidation_job(
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE bank_id = $1
|
||||
AND consolidated_at IS NULL
|
||||
AND consolidation_failed_at IS NULL
|
||||
AND fact_type IN ('experience', 'world')
|
||||
ORDER BY created_at ASC
|
||||
LIMIT $2
|
||||
@@ -245,91 +302,148 @@ async def run_consolidation_job(
|
||||
if memory_tags:
|
||||
consolidated_tags.update(memory_tags)
|
||||
|
||||
async with pool.acquire() as conn:
|
||||
# Determine observation_scopes for this batch. All memories in a batch share
|
||||
# the same tags (enforced by tag_groups), so we only check the first memory.
|
||||
# asyncpg returns JSONB columns as raw JSON strings, so parse if needed.
|
||||
_obs_raw = llm_batch[0].get("observation_scopes") if llm_batch else None
|
||||
_obs_parsed = json.loads(_obs_raw) if isinstance(_obs_raw, str) else _obs_raw
|
||||
# Process llm_batch with adaptive splitting: on LLM failure, halve the sub-batch
|
||||
# and retry, down to batch_size=1. Only if a single-memory batch still fails is
|
||||
# the memory marked with consolidation_failed_at and excluded from future runs
|
||||
# until explicitly retried via the API.
|
||||
all_results: list[dict[str, Any]] = []
|
||||
all_deleted = 0
|
||||
succeeded_ids: list[Any] = []
|
||||
failed_ids: list[Any] = []
|
||||
|
||||
# Resolve the scope spec into a concrete list[list[str]] (or None for combined).
|
||||
if _obs_parsed == "per_tag":
|
||||
_memory_tags = llm_batch[0].get("tags") or []
|
||||
obs_tags_list = [[tag] for tag in _memory_tags] if _memory_tags else None
|
||||
elif _obs_parsed == "all_combinations":
|
||||
_memory_tags = llm_batch[0].get("tags") or []
|
||||
obs_tags_list = (
|
||||
[
|
||||
list(combo)
|
||||
for r in range(1, len(_memory_tags) + 1)
|
||||
for combo in combinations(_memory_tags, r)
|
||||
]
|
||||
if _memory_tags
|
||||
else None
|
||||
)
|
||||
elif _obs_parsed == "combined" or _obs_parsed is None:
|
||||
obs_tags_list = None # single combined pass (default behaviour)
|
||||
else:
|
||||
# explicit list[list[str]]
|
||||
obs_tags_list = _obs_parsed
|
||||
pending: list[list[dict[str, Any]]] = [llm_batch]
|
||||
while pending:
|
||||
sub_batch = pending.pop(0)
|
||||
|
||||
if obs_tags_list:
|
||||
# Multi-pass: run one observation consolidation pass per tag set
|
||||
results = []
|
||||
for obs_tags in obs_tags_list:
|
||||
pass_results = await _process_memory_batch(
|
||||
async with pool.acquire() 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.
|
||||
_obs_raw = sub_batch[0].get("observation_scopes") if sub_batch else None
|
||||
_obs_parsed = json.loads(_obs_raw) if isinstance(_obs_raw, str) else _obs_raw
|
||||
|
||||
# Resolve the scope spec into a concrete list[list[str]] (or None for combined).
|
||||
if _obs_parsed == "per_tag":
|
||||
_memory_tags = sub_batch[0].get("tags") or []
|
||||
obs_tags_list = [[tag] for tag in _memory_tags] if _memory_tags else None
|
||||
elif _obs_parsed == "all_combinations":
|
||||
_memory_tags = sub_batch[0].get("tags") or []
|
||||
obs_tags_list = (
|
||||
[
|
||||
list(combo)
|
||||
for r in range(1, len(_memory_tags) + 1)
|
||||
for combo in combinations(_memory_tags, r)
|
||||
]
|
||||
if _memory_tags
|
||||
else None
|
||||
)
|
||||
elif _obs_parsed == "combined" or _obs_parsed is None:
|
||||
obs_tags_list = None # single combined pass (default behaviour)
|
||||
else:
|
||||
# explicit list[list[str]]
|
||||
obs_tags_list = _obs_parsed
|
||||
|
||||
sub_deleted: int = 0
|
||||
sub_llm_failed = False
|
||||
if obs_tags_list:
|
||||
# Multi-pass: run one observation consolidation pass per tag set
|
||||
sub_results: list[dict[str, Any]] = []
|
||||
for obs_tags in obs_tags_list:
|
||||
pass_results, pass_deleted, pass_failed = await _process_memory_batch(
|
||||
conn=conn,
|
||||
memory_engine=memory_engine,
|
||||
llm_config=llm_config,
|
||||
bank_id=bank_id,
|
||||
memories=sub_batch,
|
||||
request_context=request_context,
|
||||
perf=perf,
|
||||
config=config,
|
||||
obs_tags_override=obs_tags,
|
||||
)
|
||||
sub_deleted += pass_deleted
|
||||
sub_llm_failed = sub_llm_failed or pass_failed
|
||||
# Merge results: prefer non-skipped actions
|
||||
if not sub_results:
|
||||
sub_results = pass_results
|
||||
else:
|
||||
for i, (existing, new) in enumerate(zip(sub_results, pass_results)):
|
||||
if existing.get("action") == "skipped" and new.get("action") != "skipped":
|
||||
sub_results[i] = new
|
||||
elif existing.get("action") != "skipped" and new.get("action") != "skipped":
|
||||
# Both did something — combine into "multiple"
|
||||
existing_created = existing.get(
|
||||
"created", 1 if existing.get("action") == "created" else 0
|
||||
)
|
||||
existing_updated = existing.get(
|
||||
"updated", 1 if existing.get("action") == "updated" else 0
|
||||
)
|
||||
new_created = new.get("created", 1 if new.get("action") == "created" else 0)
|
||||
new_updated = new.get("updated", 1 if new.get("action") == "updated" else 0)
|
||||
total = existing_created + existing_updated + new_created + new_updated
|
||||
sub_results[i] = {
|
||||
"action": "multiple",
|
||||
"created": existing_created + new_created,
|
||||
"updated": existing_updated + new_updated,
|
||||
"merged": 0,
|
||||
"total_actions": total,
|
||||
}
|
||||
else:
|
||||
# Normal single pass using the memory's own tags
|
||||
sub_results, sub_deleted, sub_llm_failed = await _process_memory_batch(
|
||||
conn=conn,
|
||||
memory_engine=memory_engine,
|
||||
llm_config=llm_config,
|
||||
bank_id=bank_id,
|
||||
memories=llm_batch,
|
||||
memories=sub_batch,
|
||||
request_context=request_context,
|
||||
perf=perf,
|
||||
config=config,
|
||||
obs_tags_override=obs_tags,
|
||||
)
|
||||
# Merge results: prefer non-skipped actions
|
||||
if not results:
|
||||
results = pass_results
|
||||
else:
|
||||
for i, (existing, new) in enumerate(zip(results, pass_results)):
|
||||
if existing.get("action") == "skipped" and new.get("action") != "skipped":
|
||||
results[i] = new
|
||||
elif existing.get("action") != "skipped" and new.get("action") != "skipped":
|
||||
# Both did something — combine into "multiple"
|
||||
existing_created = existing.get(
|
||||
"created", 1 if existing.get("action") == "created" else 0
|
||||
)
|
||||
existing_updated = existing.get(
|
||||
"updated", 1 if existing.get("action") == "updated" else 0
|
||||
)
|
||||
new_created = new.get("created", 1 if new.get("action") == "created" else 0)
|
||||
new_updated = new.get("updated", 1 if new.get("action") == "updated" else 0)
|
||||
total = existing_created + existing_updated + new_created + new_updated
|
||||
results[i] = {
|
||||
"action": "multiple",
|
||||
"created": existing_created + new_created,
|
||||
"updated": existing_updated + new_updated,
|
||||
"merged": 0,
|
||||
"total_actions": total,
|
||||
}
|
||||
|
||||
all_deleted += sub_deleted
|
||||
|
||||
if sub_llm_failed and len(sub_batch) > 1:
|
||||
# Split and retry with smaller batches
|
||||
mid = len(sub_batch) // 2
|
||||
logger.warning(
|
||||
f"[CONSOLIDATION] bank={bank_id} LLM failed for sub-batch of {len(sub_batch)},"
|
||||
f" splitting into {mid}/{len(sub_batch) - mid}"
|
||||
)
|
||||
pending[0:0] = [sub_batch[:mid], sub_batch[mid:]]
|
||||
elif sub_llm_failed:
|
||||
# batch_size=1 and still failing — mark as permanently failed for now
|
||||
failed_ids.append(sub_batch[0]["id"])
|
||||
all_results.append({"action": "failed"})
|
||||
logger.warning(
|
||||
f"[CONSOLIDATION] bank={bank_id} LLM failed for single memory"
|
||||
f" {sub_batch[0]['id']}, marking consolidation_failed_at"
|
||||
)
|
||||
else:
|
||||
# Normal single pass using the memory's own tags
|
||||
results = await _process_memory_batch(
|
||||
conn=conn,
|
||||
memory_engine=memory_engine,
|
||||
llm_config=llm_config,
|
||||
bank_id=bank_id,
|
||||
memories=llm_batch,
|
||||
request_context=request_context,
|
||||
perf=perf,
|
||||
config=config,
|
||||
succeeded_ids.extend(m["id"] for m in sub_batch)
|
||||
all_results.extend(sub_results)
|
||||
|
||||
# Commit consolidated_at / consolidation_failed_at in a single DB round-trip
|
||||
async with pool.acquire() as conn:
|
||||
if succeeded_ids:
|
||||
await conn.executemany(
|
||||
f"UPDATE {fq_table('memory_units')} SET consolidated_at = NOW() WHERE id = $1",
|
||||
[(mem_id,) for mem_id in succeeded_ids],
|
||||
)
|
||||
if failed_ids:
|
||||
await conn.executemany(
|
||||
f"UPDATE {fq_table('memory_units')} SET consolidation_failed_at = NOW() WHERE id = $1",
|
||||
[(mem_id,) for mem_id in failed_ids],
|
||||
)
|
||||
|
||||
await conn.executemany(
|
||||
f"UPDATE {fq_table('memory_units')} SET consolidated_at = NOW() WHERE id = $1",
|
||||
[(m["id"],) for m in llm_batch],
|
||||
stats["observations_deleted"] += all_deleted
|
||||
results = all_results
|
||||
|
||||
# Checkpoint: abort if the operation (and thus the bank) was deleted mid-run.
|
||||
if operation_id and not await memory_engine._check_op_alive(operation_id):
|
||||
logger.info(
|
||||
f"[CONSOLIDATION] bank={bank_id} operation {operation_id} cancelled (bank deleted), stopping early"
|
||||
)
|
||||
return {"status": "cancelled", "bank_id": bank_id, **stats}
|
||||
|
||||
for result in results:
|
||||
stats["memories_processed"] += 1
|
||||
@@ -350,6 +464,8 @@ async def run_consolidation_job(
|
||||
stats["actions_executed"] += result.get("total_actions", 0)
|
||||
elif action == "skipped":
|
||||
stats["skipped"] += 1
|
||||
elif action == "failed":
|
||||
stats["memories_failed"] += 1
|
||||
|
||||
# Per-LLM-batch log
|
||||
llm_batch_time = time.time() - llm_batch_start
|
||||
@@ -362,6 +478,7 @@ async def run_consolidation_job(
|
||||
batch_created = stats["observations_created"] - snap_stats["observations_created"]
|
||||
batch_updated = stats["observations_updated"] - snap_stats["observations_updated"]
|
||||
batch_skipped = stats["skipped"] - snap_stats["skipped"]
|
||||
batch_failed = stats["memories_failed"] - snap_stats["memories_failed"]
|
||||
llm_calls_made = perf.llm_calls - snap_llm_calls
|
||||
logger.info(
|
||||
f"[CONSOLIDATION] bank={bank_id} llm_batch #{llm_batch_num}"
|
||||
@@ -369,7 +486,8 @@ async def run_consolidation_job(
|
||||
f" | {stats['memories_processed']}/{total_count} processed"
|
||||
f" | {', '.join(timing_parts)}"
|
||||
f" | created={batch_created} updated={batch_updated} skipped={batch_skipped}"
|
||||
f" | input_tokens=~{input_tokens}"
|
||||
+ (f" failed={batch_failed}" if batch_failed else "")
|
||||
+ f" | input_tokens=~{input_tokens}"
|
||||
f" | avg={llm_batch_time / len(llm_batch):.3f}s/memory"
|
||||
)
|
||||
|
||||
@@ -521,7 +639,7 @@ async def _process_memory_batch(
|
||||
perf: ConsolidationPerfLog | None = None,
|
||||
config: Any = None,
|
||||
obs_tags_override: list[str] | None = None,
|
||||
) -> list[dict[str, Any]]:
|
||||
) -> tuple[list[dict[str, Any]], int, bool]:
|
||||
"""
|
||||
Process a batch of memories in a single LLM call.
|
||||
|
||||
@@ -612,17 +730,18 @@ async def _process_memory_batch(
|
||||
source_mems = [mem_by_id[fid] for fid in create.source_fact_ids if fid in mem_by_id]
|
||||
if not source_mems:
|
||||
continue
|
||||
agg = _aggregate_source_fields(source_mems, tags=fact_tags)
|
||||
await _execute_create_action(
|
||||
conn=conn,
|
||||
memory_engine=memory_engine,
|
||||
bank_id=bank_id,
|
||||
source_memory_ids=[m["id"] for m in source_mems],
|
||||
text=create.text,
|
||||
source_fact_tags=fact_tags,
|
||||
event_date=_min_date(m.get("event_date") for m in source_mems),
|
||||
occurred_start=_min_date(m.get("occurred_start") for m in source_mems),
|
||||
occurred_end=_max_date(m.get("occurred_end") for m in source_mems),
|
||||
mentioned_at=_max_date(m.get("mentioned_at") for m in source_mems),
|
||||
source_fact_tags=agg.tags,
|
||||
event_date=agg.event_date,
|
||||
occurred_start=agg.occurred_start,
|
||||
occurred_end=agg.occurred_end,
|
||||
mentioned_at=agg.mentioned_at,
|
||||
perf=perf,
|
||||
)
|
||||
for m in source_mems:
|
||||
@@ -639,6 +758,7 @@ async def _process_memory_batch(
|
||||
f"not in any source fact's recall"
|
||||
)
|
||||
continue
|
||||
agg = _aggregate_source_fields(source_mems, tags=fact_tags)
|
||||
await _execute_update_action(
|
||||
conn=conn,
|
||||
memory_engine=memory_engine,
|
||||
@@ -647,15 +767,16 @@ async def _process_memory_batch(
|
||||
observation_id=update.observation_id,
|
||||
new_text=update.text,
|
||||
observations=union_observations,
|
||||
source_fact_tags=fact_tags,
|
||||
source_occurred_start=_min_date(m.get("occurred_start") for m in source_mems),
|
||||
source_occurred_end=_max_date(m.get("occurred_end") for m in source_mems),
|
||||
source_mentioned_at=_max_date(m.get("mentioned_at") for m in source_mems),
|
||||
source_fact_tags=agg.tags,
|
||||
source_occurred_start=agg.occurred_start,
|
||||
source_occurred_end=agg.occurred_end,
|
||||
source_mentioned_at=agg.mentioned_at,
|
||||
perf=perf,
|
||||
)
|
||||
for m in source_mems:
|
||||
per_memory_updated.add(str(m["id"]))
|
||||
|
||||
deleted_count = 0
|
||||
for delete in llm_result.deletes:
|
||||
# Security: the observation must be present in the unioned recall
|
||||
if not any(str(obs.id) == delete.observation_id for obs in union_observations):
|
||||
@@ -664,6 +785,7 @@ async def _process_memory_batch(
|
||||
)
|
||||
continue
|
||||
await _execute_delete_action(conn=conn, bank_id=bank_id, observation_id=delete.observation_id)
|
||||
deleted_count += 1
|
||||
|
||||
# Build per-memory result dicts for the stats tracker in the outer loop
|
||||
results: list[dict[str, Any]] = []
|
||||
@@ -680,7 +802,7 @@ async def _process_memory_batch(
|
||||
else:
|
||||
results.append({"action": "skipped", "reason": "no_durable_knowledge"})
|
||||
|
||||
return results
|
||||
return results, deleted_count, llm_result.failed
|
||||
|
||||
|
||||
def _min_date(dates: "Any") -> "datetime | None":
|
||||
@@ -718,13 +840,17 @@ async def _execute_update_action(
|
||||
logger.debug(f"Update skipped: observation {observation_id} not found in recall results")
|
||||
return
|
||||
|
||||
history = [
|
||||
{
|
||||
"previous_text": model.text,
|
||||
"changed_at": datetime.now(timezone.utc).isoformat(),
|
||||
"source_memory_ids": [str(mid) for mid in source_memory_ids],
|
||||
}
|
||||
]
|
||||
from ...config import get_config
|
||||
|
||||
history_entry = {
|
||||
"previous_text": model.text,
|
||||
"previous_tags": list(model.tags or []),
|
||||
"previous_occurred_start": model.occurred_start,
|
||||
"previous_occurred_end": model.occurred_end,
|
||||
"previous_mentioned_at": model.mentioned_at,
|
||||
"changed_at": datetime.now(timezone.utc).isoformat(),
|
||||
"new_source_memory_ids": [str(mid) for mid in source_memory_ids],
|
||||
}
|
||||
|
||||
source_ids = list(model.source_fact_ids or []) + source_memory_ids
|
||||
|
||||
@@ -739,13 +865,18 @@ async def _execute_update_action(
|
||||
if perf:
|
||||
perf.record_timing("embedding", time.time() - t0)
|
||||
|
||||
config = get_config()
|
||||
history_clause = (
|
||||
"history = COALESCE(history, '[]'::jsonb) || $3::jsonb," if config.enable_observation_history else ""
|
||||
)
|
||||
|
||||
t0 = time.time()
|
||||
await conn.execute(
|
||||
f"""
|
||||
UPDATE {fq_table("memory_units")}
|
||||
SET text = $1,
|
||||
embedding = $2::vector,
|
||||
history = $3,
|
||||
{history_clause}
|
||||
source_memory_ids = $4,
|
||||
proof_count = $5,
|
||||
tags = $10,
|
||||
@@ -757,7 +888,7 @@ async def _execute_update_action(
|
||||
""",
|
||||
new_text,
|
||||
embedding_str,
|
||||
json.dumps(history),
|
||||
json.dumps([history_entry]),
|
||||
source_ids,
|
||||
len(source_ids),
|
||||
uuid.UUID(observation_id),
|
||||
@@ -869,10 +1000,9 @@ async def _find_related_observations(
|
||||
"""
|
||||
# Use recall to find related observations with token budget
|
||||
# max_tokens naturally limits how many observations are returned
|
||||
from ...config import get_config
|
||||
from ...tracing import get_tracer, is_tracing_enabled
|
||||
|
||||
config = get_config()
|
||||
config = await memory_engine._config_resolver.resolve_full_config(bank_id, request_context)
|
||||
|
||||
# SECURITY: Use all_strict matching if tags provided to prevent cross-scope consolidation
|
||||
tags_match = "all_strict" if tags else "any"
|
||||
@@ -897,7 +1027,8 @@ async def _find_related_observations(
|
||||
tags=tags, # Filter by source memory's tags
|
||||
tags_match=tags_match, # Use strict matching for security
|
||||
include_source_facts=True, # Embed source facts so we avoid a separate DB fetch
|
||||
max_source_facts_tokens=-1, # No token limit — we need all source facts for consolidation
|
||||
max_source_facts_tokens=config.consolidation_source_facts_max_tokens,
|
||||
max_source_facts_tokens_per_observation=config.consolidation_source_facts_max_tokens_per_observation,
|
||||
_quiet=True, # Suppress logging
|
||||
)
|
||||
finally:
|
||||
@@ -961,14 +1092,17 @@ async def _consolidate_batch_with_llm(
|
||||
observations_text = "[]"
|
||||
|
||||
def _fact_line(m: dict[str, Any]) -> str:
|
||||
parts = [f"[{m['id']}] {m['text']}"]
|
||||
text = f"[{m['id']}] {m['text']}"
|
||||
temporal_parts = []
|
||||
if m.get("occurred_start"):
|
||||
parts.append(f"occurred_start={m['occurred_start']}")
|
||||
temporal_parts.append(f"occurred_start={m['occurred_start']}")
|
||||
if m.get("occurred_end"):
|
||||
parts.append(f"occurred_end={m['occurred_end']}")
|
||||
temporal_parts.append(f"occurred_end={m['occurred_end']}")
|
||||
if m.get("mentioned_at"):
|
||||
parts.append(f"mentioned_at={m['mentioned_at']}")
|
||||
return " | ".join(parts)
|
||||
temporal_parts.append(f"mentioned_at={m['mentioned_at']}")
|
||||
if temporal_parts:
|
||||
text += f" ({', '.join(temporal_parts)})"
|
||||
return text
|
||||
|
||||
facts_lines = "\n".join(_fact_line(m) for m in memories)
|
||||
|
||||
@@ -1002,7 +1136,7 @@ async def _consolidate_batch_with_llm(
|
||||
logger.error(
|
||||
f"[CONSOLIDATION] LLM batch call failed after {max_attempts} attempts, skipping batch. Last error: {last_exc}"
|
||||
)
|
||||
return _BatchLLMResult(obs_count=len(union_observations), prompt_chars=len(prompt))
|
||||
return _BatchLLMResult(obs_count=len(union_observations), prompt_chars=len(prompt), failed=True)
|
||||
|
||||
|
||||
async def _create_observation_directly(
|
||||
@@ -1029,8 +1163,8 @@ async def _create_observation_directly(
|
||||
# Create the observation as a memory_unit
|
||||
now = datetime.now(timezone.utc)
|
||||
obs_event_date = event_date or now
|
||||
obs_occurred_start = occurred_start or now
|
||||
obs_occurred_end = occurred_end or now
|
||||
obs_occurred_start = occurred_start
|
||||
obs_occurred_end = occurred_end
|
||||
obs_mentioned_at = mentioned_at or now
|
||||
obs_tags = tags or []
|
||||
|
||||
+20
-3
@@ -29,14 +29,31 @@ Compare the facts against existing observations:
|
||||
- Same topic as an existing observation → UPDATE it (observation_id + source_fact_ids)
|
||||
- New topic with durable knowledge → CREATE a new observation (source_fact_ids)
|
||||
- Cross-reference facts within the batch: a later fact may resolve a vague reference in an earlier one
|
||||
- Purely ephemeral facts → omit them (no create/update needed)"""
|
||||
- Purely ephemeral facts → omit them unless the MISSION above explicitly targets such data (e.g. timestamped events, session state, screen content)"""
|
||||
|
||||
# Output format — JSON braces escaped as {{ }} so .format() leaves them literal
|
||||
_BATCH_OUTPUT_FORMAT = """
|
||||
Output a JSON object with three arrays.
|
||||
|
||||
Example (showing the required UUID format for all IDs):
|
||||
{{"creates": [{{"text": "Alice lives in Berlin", "source_fact_ids": ["a1b2c3d4-e5f6-7890-abcd-ef1234567890", "b2c3d4e5-f6a7-8901-bcde-f12345678901"]}}],
|
||||
## EXAMPLE
|
||||
|
||||
Input facts:
|
||||
[a1b2c3d4-e5f6-7890-abcd-ef1234567890] Alice mentioned she works long hours, often past midnight | Involving: Alice (occurred_start=2024-01-15, mentioned_at=2024-01-15)
|
||||
[b2c3d4e5-f6a7-8901-bcde-f12345678901] Alice said she's exhausted from the project deadlines | Involving: Alice (occurred_start=2024-01-20, mentioned_at=2024-01-20)
|
||||
|
||||
Good observation text — clean prose, no metadata, each fact tracked distinctly:
|
||||
"Alice works long hours, often past midnight."
|
||||
"Alice feels exhausted from project deadlines."
|
||||
|
||||
Bad observation text — NEVER do this (verbatim copy of fact text with metadata):
|
||||
"Alice mentioned she works long hours, often past midnight | Involving: Alice (occurred_start=2024-01-15, mentioned_at=2024-01-15)"
|
||||
|
||||
Observation text rules:
|
||||
- Write clean prose — NEVER copy raw fact lines or their metadata (temporal fields, "Involving:", "When:" labels, UUIDs).
|
||||
- Parenthesized metadata like (occurred_start=...) and pipe-separated labels like "| Involving: ..." are fact formatting — strip them entirely from observation text.
|
||||
- How many observations to create and how much to aggregate is driven by the MISSION above.
|
||||
|
||||
{{"creates": [{{"text": "Alice works long hours, often past midnight.", "source_fact_ids": ["a1b2c3d4-e5f6-7890-abcd-ef1234567890"]}}, {{"text": "Alice feels exhausted from project deadlines.", "source_fact_ids": ["b2c3d4e5-f6a7-8901-bcde-f12345678901"]}}],
|
||||
"updates": [{{"text": "Alice works at Acme Corp as a senior engineer", "observation_id": "c3d4e5f6-a7b8-9012-cdef-123456789012", "source_fact_ids": ["d4e5f6a7-b8c9-0123-defa-234567890123"]}}],
|
||||
"deletes": [{{"observation_id": "e5f6a7b8-c9d0-1234-efab-345678901234"}}]}}
|
||||
|
||||
+189
-3
@@ -20,8 +20,10 @@ from ..config import (
|
||||
DEFAULT_RERANKER_COHERE_MODEL,
|
||||
DEFAULT_RERANKER_FLASHRANK_CACHE_DIR,
|
||||
DEFAULT_RERANKER_FLASHRANK_MODEL,
|
||||
DEFAULT_RERANKER_LITELLM_MAX_TOKENS_PER_DOC,
|
||||
DEFAULT_RERANKER_LITELLM_MODEL,
|
||||
DEFAULT_RERANKER_LITELLM_SDK_MODEL,
|
||||
DEFAULT_RERANKER_LOCAL_BATCH_SIZE,
|
||||
DEFAULT_RERANKER_LOCAL_FORCE_CPU,
|
||||
DEFAULT_RERANKER_LOCAL_MAX_CONCURRENT,
|
||||
DEFAULT_RERANKER_LOCAL_MODEL,
|
||||
@@ -110,6 +112,9 @@ class LocalSTCrossEncoder(CrossEncoderModel):
|
||||
max_concurrent: int = 4,
|
||||
force_cpu: bool = False,
|
||||
trust_remote_code: bool = False,
|
||||
fp16: bool = False,
|
||||
bucket_batching: bool = False,
|
||||
batch_size: int = DEFAULT_RERANKER_LOCAL_BATCH_SIZE,
|
||||
):
|
||||
"""
|
||||
Initialize local SentenceTransformers cross-encoder.
|
||||
@@ -124,10 +129,20 @@ class LocalSTCrossEncoder(CrossEncoderModel):
|
||||
trust_remote_code: Allow loading models with custom code (security risk).
|
||||
Required for some models like jina-reranker-v2-base-multilingual.
|
||||
Default: False (disabled for security)
|
||||
fp16: Use FP16 (half precision) inference. Faster on MPS and CUDA,
|
||||
may be slower on CPU. Default: False (opt-in via env var).
|
||||
bucket_batching: Sort pairs by token length before batching to reduce
|
||||
padding waste. 36-54% speedup, quality-identical.
|
||||
Default: False (opt-in via env var).
|
||||
batch_size: Batch size for predict() calls. Optimal values vary by
|
||||
hardware and model (MPS: 32, CUDA: 128+). Default: 32.
|
||||
"""
|
||||
self.model_name = model_name or DEFAULT_RERANKER_LOCAL_MODEL
|
||||
self.force_cpu = force_cpu
|
||||
self.trust_remote_code = trust_remote_code
|
||||
self.fp16 = fp16
|
||||
self.bucket_batching = bucket_batching
|
||||
self.batch_size = batch_size
|
||||
self._model = None
|
||||
LocalSTCrossEncoder._max_concurrent = max_concurrent
|
||||
|
||||
@@ -175,6 +190,24 @@ class LocalSTCrossEncoder(CrossEncoderModel):
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to detect GPU/MPS, falling back to CPU: {e}")
|
||||
|
||||
# Patch transformers 5.x compatibility for models using XLM-RoBERTa
|
||||
# (e.g., jina-reranker-v2-base-multilingual). transformers 5.x removed
|
||||
# create_position_ids_from_input_ids as a module-level function; the custom
|
||||
# code in these models still references it. This monkey-patch restores it.
|
||||
try:
|
||||
import transformers.models.xlm_roberta.modeling_xlm_roberta as xlm_module
|
||||
from transformers.models.xlm_roberta.modeling_xlm_roberta import XLMRobertaEmbeddings
|
||||
|
||||
if not hasattr(xlm_module, "create_position_ids_from_input_ids"):
|
||||
setattr(
|
||||
xlm_module,
|
||||
"create_position_ids_from_input_ids",
|
||||
XLMRobertaEmbeddings.create_position_ids_from_input_ids,
|
||||
)
|
||||
logger.info("Reranker: applied transformers 5.x compatibility patch for XLM-RoBERTa")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Suppress verbose transformers warnings during model loading
|
||||
# This suppresses the "UNEXPECTED" warnings from CrossEncoder which are harmless
|
||||
# but look alarming to users (e.g., "embeddings.position_ids | UNEXPECTED")
|
||||
@@ -199,6 +232,12 @@ class LocalSTCrossEncoder(CrossEncoderModel):
|
||||
# Restore original logging level
|
||||
transformers_logger.setLevel(original_level)
|
||||
|
||||
# FP16 inference: convert model weights to half precision.
|
||||
# Empirically validated: 27-36% faster on MPS, quality-identical (20/20 overlap).
|
||||
if self.fp16 and device != "cpu":
|
||||
self._model.model.half()
|
||||
logger.info("Reranker: FP16 inference enabled")
|
||||
|
||||
# Initialize shared executor (limited workers naturally limits concurrency)
|
||||
if LocalSTCrossEncoder._executor is None:
|
||||
LocalSTCrossEncoder._executor = ThreadPoolExecutor(
|
||||
@@ -210,8 +249,32 @@ class LocalSTCrossEncoder(CrossEncoderModel):
|
||||
logger.info("Reranker: local provider initialized (using existing executor)")
|
||||
|
||||
def _predict_sync(self, pairs: list[tuple[str, str]]) -> list[float]:
|
||||
"""Synchronous prediction wrapper for thread pool execution."""
|
||||
scores = self._model.predict(pairs, show_progress_bar=False)
|
||||
"""Synchronous prediction wrapper for thread pool execution.
|
||||
|
||||
Supports two optimizations (controlled via .env):
|
||||
- bucket_batching: sort pairs by token length to reduce padding waste (36-54% speedup)
|
||||
- batch_size: explicit batch size for predict() calls (MPS optimal: 32)
|
||||
"""
|
||||
import numpy as np
|
||||
|
||||
if self.bucket_batching and len(pairs) > 1:
|
||||
# Sort pairs by approximate token length to create homogeneous batches.
|
||||
# This eliminates padding waste — short pairs aren't padded to the length
|
||||
# of the longest pair in the batch. Quality-identical by construction.
|
||||
lengths = [len(pairs[i][0]) + len(pairs[i][1]) for i in range(len(pairs))]
|
||||
sorted_indices = sorted(range(len(pairs)), key=lambda i: lengths[i])
|
||||
sorted_pairs = [pairs[i] for i in sorted_indices]
|
||||
|
||||
sorted_scores = self._model.predict(sorted_pairs, batch_size=self.batch_size, show_progress_bar=False)
|
||||
sorted_scores = sorted_scores.tolist() if hasattr(sorted_scores, "tolist") else list(sorted_scores)
|
||||
|
||||
# Restore original order
|
||||
scores = [0.0] * len(pairs)
|
||||
for new_pos, orig_idx in enumerate(sorted_indices):
|
||||
scores[orig_idx] = sorted_scores[new_pos]
|
||||
return scores
|
||||
|
||||
scores = self._model.predict(pairs, batch_size=self.batch_size, show_progress_bar=False)
|
||||
return scores.tolist() if hasattr(scores, "tolist") else list(scores)
|
||||
|
||||
async def predict(self, pairs: list[tuple[str, str]]) -> list[float]:
|
||||
@@ -820,6 +883,17 @@ class FlashRankCrossEncoder(CrossEncoderModel):
|
||||
return await loop.run_in_executor(FlashRankCrossEncoder._executor, self._predict_sync, pairs)
|
||||
|
||||
|
||||
def _truncate_to_tokens(text: str, max_tokens: int) -> str:
|
||||
"""Truncate text to at most max_tokens using the shared tiktoken encoder."""
|
||||
from .memory_engine import _get_tiktoken_encoding
|
||||
|
||||
enc = _get_tiktoken_encoding()
|
||||
tokens = enc.encode(text)
|
||||
if len(tokens) <= max_tokens:
|
||||
return text
|
||||
return enc.decode(tokens[:max_tokens])
|
||||
|
||||
|
||||
class LiteLLMCrossEncoder(CrossEncoderModel):
|
||||
"""
|
||||
LiteLLM cross-encoder implementation using LiteLLM proxy's /rerank endpoint.
|
||||
@@ -843,6 +917,7 @@ class LiteLLMCrossEncoder(CrossEncoderModel):
|
||||
api_key: str | None = None,
|
||||
model: str = DEFAULT_RERANKER_LITELLM_MODEL,
|
||||
timeout: float = 60.0,
|
||||
max_tokens_per_doc: int | None = DEFAULT_RERANKER_LITELLM_MAX_TOKENS_PER_DOC,
|
||||
):
|
||||
"""
|
||||
Initialize LiteLLM cross-encoder client.
|
||||
@@ -853,11 +928,15 @@ class LiteLLMCrossEncoder(CrossEncoderModel):
|
||||
model: Reranking model name (default: cohere/rerank-english-v3.0)
|
||||
Use provider prefix (e.g., cohere/, together_ai/, voyage/)
|
||||
timeout: Request timeout in seconds (default: 60.0)
|
||||
max_tokens_per_doc: If set, truncate each document to this many tokens before
|
||||
sending to the reranker (uses tiktoken cl100k_base encoding).
|
||||
Useful for models with small context windows (e.g. 1024 tokens).
|
||||
"""
|
||||
self.api_base = api_base.rstrip("/")
|
||||
self.api_key = api_key
|
||||
self.model = model
|
||||
self.timeout = timeout
|
||||
self.max_tokens_per_doc = max_tokens_per_doc
|
||||
self._async_client: httpx.AsyncClient | None = None
|
||||
|
||||
@property
|
||||
@@ -905,6 +984,8 @@ class LiteLLMCrossEncoder(CrossEncoderModel):
|
||||
|
||||
for query, indexed_texts in query_groups.items():
|
||||
texts = [text for _, text in indexed_texts]
|
||||
if self.max_tokens_per_doc is not None:
|
||||
texts = [_truncate_to_tokens(t, self.max_tokens_per_doc) for t in texts]
|
||||
indices = [idx for idx, _ in indexed_texts]
|
||||
|
||||
# LiteLLM /rerank follows Cohere API format
|
||||
@@ -950,6 +1031,7 @@ class LiteLLMSDKCrossEncoder(CrossEncoderModel):
|
||||
model: str = DEFAULT_RERANKER_LITELLM_SDK_MODEL,
|
||||
api_base: str | None = None,
|
||||
timeout: float = 60.0,
|
||||
max_tokens_per_doc: int | None = DEFAULT_RERANKER_LITELLM_MAX_TOKENS_PER_DOC,
|
||||
):
|
||||
"""
|
||||
Initialize LiteLLM SDK cross-encoder client.
|
||||
@@ -959,11 +1041,15 @@ class LiteLLMSDKCrossEncoder(CrossEncoderModel):
|
||||
model: Model name with provider prefix (e.g., "deepinfra/Qwen3-reranker-8B")
|
||||
api_base: Custom base URL for API (optional)
|
||||
timeout: Request timeout in seconds (default: 60.0)
|
||||
max_tokens_per_doc: If set, truncate each document to this many tokens before
|
||||
sending to the reranker (uses tiktoken cl100k_base encoding).
|
||||
Useful for models with small context windows (e.g. 1024 tokens).
|
||||
"""
|
||||
self.api_key = api_key
|
||||
self.model = model
|
||||
self.api_base = api_base
|
||||
self.timeout = timeout
|
||||
self.max_tokens_per_doc = max_tokens_per_doc
|
||||
self._initialized = False
|
||||
self._litellm = None # Will be set during initialization
|
||||
|
||||
@@ -1017,6 +1103,8 @@ class LiteLLMSDKCrossEncoder(CrossEncoderModel):
|
||||
|
||||
for query, indexed_texts in query_groups.items():
|
||||
texts = [text for _, text in indexed_texts]
|
||||
if self.max_tokens_per_doc is not None:
|
||||
texts = [_truncate_to_tokens(t, self.max_tokens_per_doc) for t in texts]
|
||||
indices = [idx for idx, _ in indexed_texts]
|
||||
|
||||
# Build kwargs for rerank call
|
||||
@@ -1050,6 +1138,97 @@ class LiteLLMSDKCrossEncoder(CrossEncoderModel):
|
||||
return all_scores
|
||||
|
||||
|
||||
class JinaMLXCrossEncoder(CrossEncoderModel):
|
||||
"""
|
||||
Jina Reranker v3 MLX implementation for Apple Silicon.
|
||||
|
||||
Uses jinaai/jina-reranker-v3-mlx — a 0.6B parameter multilingual listwise reranker
|
||||
optimized for Apple Silicon via the MLX framework. No transformers/PyTorch dependency.
|
||||
|
||||
The model is downloaded automatically from HuggingFace Hub on first use.
|
||||
Requires: mlx>=0.31.0, mlx-lm>=0.31.1, safetensors>=0.6.2
|
||||
"""
|
||||
|
||||
HF_REPO_ID = "jinaai/jina-reranker-v3-mlx"
|
||||
|
||||
def __init__(self, model_path: str | None = None):
|
||||
"""
|
||||
Args:
|
||||
model_path: Local path to the downloaded model directory.
|
||||
If None, the model is downloaded from HuggingFace Hub.
|
||||
"""
|
||||
self.model_path = model_path
|
||||
self._reranker = None
|
||||
|
||||
@property
|
||||
def provider_name(self) -> str:
|
||||
return "jina-mlx"
|
||||
|
||||
async def initialize(self) -> None:
|
||||
if self._reranker is not None:
|
||||
return
|
||||
|
||||
try:
|
||||
import mlx.core # noqa: F401
|
||||
import mlx_lm # noqa: F401
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"mlx and mlx-lm are required for JinaMLXCrossEncoder. "
|
||||
"Install with: pip install mlx>=0.31.0 mlx-lm>=0.31.1 safetensors>=0.6.2"
|
||||
)
|
||||
|
||||
loop = asyncio.get_event_loop()
|
||||
await loop.run_in_executor(None, self._load_model)
|
||||
|
||||
def _load_model(self) -> None:
|
||||
"""Download (if needed) and load the MLX reranker. Runs in a thread."""
|
||||
import os
|
||||
|
||||
from huggingface_hub import snapshot_download
|
||||
|
||||
from .jina_mlx_reranker import MLXReranker
|
||||
|
||||
model_path = self.model_path
|
||||
if model_path is None:
|
||||
logger.info(f"Reranker: downloading {self.HF_REPO_ID} from HuggingFace Hub...")
|
||||
model_path = snapshot_download(repo_id=self.HF_REPO_ID)
|
||||
|
||||
logger.info(f"Reranker: loading jina-reranker-v3-mlx from {model_path}")
|
||||
self._reranker = MLXReranker(
|
||||
model_path=model_path,
|
||||
projector_path=os.path.join(model_path, "projector.safetensors"),
|
||||
)
|
||||
logger.info("Reranker: jina-mlx provider initialized")
|
||||
|
||||
def _predict_sync(self, pairs: list[tuple[str, str]]) -> list[float]:
|
||||
"""Score pairs grouped by query. Runs in a thread."""
|
||||
if not pairs:
|
||||
return []
|
||||
|
||||
query_groups: dict[str, list[tuple[int, str]]] = {}
|
||||
for idx, (query, doc) in enumerate(pairs):
|
||||
query_groups.setdefault(query, []).append((idx, doc))
|
||||
|
||||
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"]
|
||||
|
||||
return all_scores
|
||||
|
||||
async def predict(self, pairs: list[tuple[str, str]]) -> list[float]:
|
||||
if self._reranker is None:
|
||||
raise RuntimeError("Reranker not initialized. Call initialize() first.")
|
||||
|
||||
loop = asyncio.get_event_loop()
|
||||
return await loop.run_in_executor(None, self._predict_sync, pairs)
|
||||
|
||||
|
||||
def create_cross_encoder_from_env() -> CrossEncoderModel:
|
||||
"""
|
||||
Create a CrossEncoderModel instance based on configuration.
|
||||
@@ -1079,6 +1258,9 @@ def create_cross_encoder_from_env() -> CrossEncoderModel:
|
||||
max_concurrent=config.reranker_local_max_concurrent,
|
||||
force_cpu=config.reranker_local_force_cpu,
|
||||
trust_remote_code=config.reranker_local_trust_remote_code,
|
||||
fp16=config.reranker_local_fp16,
|
||||
bucket_batching=config.reranker_local_bucket_batching,
|
||||
batch_size=config.reranker_local_batch_size,
|
||||
)
|
||||
elif provider == "cohere":
|
||||
api_key = config.reranker_cohere_api_key
|
||||
@@ -1098,6 +1280,7 @@ def create_cross_encoder_from_env() -> CrossEncoderModel:
|
||||
api_base=config.reranker_litellm_api_base,
|
||||
api_key=config.reranker_litellm_api_key,
|
||||
model=config.reranker_litellm_model,
|
||||
max_tokens_per_doc=config.reranker_litellm_max_tokens_per_doc,
|
||||
)
|
||||
elif provider == "litellm-sdk":
|
||||
api_key = config.reranker_litellm_sdk_api_key
|
||||
@@ -1109,6 +1292,7 @@ def create_cross_encoder_from_env() -> CrossEncoderModel:
|
||||
api_key=api_key,
|
||||
model=config.reranker_litellm_sdk_model,
|
||||
api_base=config.reranker_litellm_sdk_api_base,
|
||||
max_tokens_per_doc=config.reranker_litellm_max_tokens_per_doc,
|
||||
)
|
||||
elif provider == "zeroentropy":
|
||||
api_key = config.reranker_zeroentropy_api_key
|
||||
@@ -1122,7 +1306,9 @@ def create_cross_encoder_from_env() -> CrossEncoderModel:
|
||||
)
|
||||
elif provider == "rrf":
|
||||
return RRFPassthroughCrossEncoder()
|
||||
elif provider == "jina-mlx":
|
||||
return JinaMLXCrossEncoder()
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unknown reranker provider: {provider}. Supported: 'local', 'tei', 'cohere', 'zeroentropy', 'flashrank', 'litellm', 'litellm-sdk', 'rrf'"
|
||||
f"Unknown reranker provider: {provider}. Supported: 'local', 'tei', 'cohere', 'zeroentropy', 'flashrank', 'litellm', 'litellm-sdk', 'rrf', 'jina-mlx'"
|
||||
)
|
||||
+10
-4
@@ -58,10 +58,16 @@ async def retry_with_backoff(
|
||||
last_exception = e
|
||||
if attempt < max_retries:
|
||||
delay = min(base_delay * (2**attempt), max_delay)
|
||||
logger.warning(
|
||||
f"Database operation failed (attempt {attempt + 1}/{max_retries + 1}): {e}. "
|
||||
f"Retrying in {delay:.1f}s..."
|
||||
)
|
||||
if isinstance(e, asyncpg.exceptions.DeadlockDetectedError):
|
||||
logger.warning(
|
||||
f"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:
|
||||
logger.warning(
|
||||
f"Database operation failed (attempt {attempt + 1}/{max_retries + 1}): {e}. "
|
||||
f"Retrying in {delay:.1f}s..."
|
||||
)
|
||||
await asyncio.sleep(delay)
|
||||
else:
|
||||
logger.error(f"Database operation failed after {max_retries + 1} attempts: {e}")
|
||||
+7
-3
@@ -459,10 +459,12 @@ class EntityResolver:
|
||||
entity_dates = [g.event_date for _, g in sorted_groups]
|
||||
|
||||
# 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()), 1
|
||||
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
|
||||
@@ -489,13 +491,15 @@ class EntityResolver:
|
||||
for row in existing_rows:
|
||||
id_by_name[row["name_lower"]] = row["id"]
|
||||
|
||||
# Assign entity IDs back and queue for post-txn stats flush.
|
||||
# Assign entity IDs back and queue one stat per original mention so that
|
||||
# flush_pending_stats() increments mention_count by the true mention count,
|
||||
# not just 1 per unique name.
|
||||
for name_lower, g in sorted_groups:
|
||||
entity_id = id_by_name.get(name_lower)
|
||||
if entity_id:
|
||||
for original_idx in g.indices:
|
||||
entity_ids[original_idx] = entity_id
|
||||
pending.append(_EntityStat(entity_id=entity_id, event_date=g.event_date))
|
||||
pending.append(_EntityStat(entity_id=entity_id, event_date=g.event_date))
|
||||
|
||||
# Accumulate into the resolver's pending list; the orchestrator flushes
|
||||
# these with await entity_resolver.flush_pending_stats() after the txn.
|
||||
@@ -0,0 +1,144 @@
|
||||
"""
|
||||
MLX implementation of jina-reranker-v3 for Apple Silicon.
|
||||
|
||||
This file is adapted from the official model repository:
|
||||
https://huggingface.co/jinaai/jina-reranker-v3-mlx/blob/main/rerank.py
|
||||
|
||||
License: CC BY-NC 4.0 (contact Jina AI for commercial usage)
|
||||
|
||||
Changes from upstream:
|
||||
- Removed the __main__ example block
|
||||
- Type annotations added to public methods
|
||||
- top_n parameter added to rerank() (upstream only exposed it implicitly)
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
class _MLPProjector:
|
||||
def __init__(self):
|
||||
import mlx.nn as nn
|
||||
|
||||
self.linear1 = nn.Linear(1024, 512, bias=False)
|
||||
self.linear2 = nn.Linear(512, 512, bias=False)
|
||||
|
||||
def __call__(self, x):
|
||||
import mlx.nn as nn
|
||||
|
||||
x = self.linear1(x)
|
||||
x = nn.relu(x)
|
||||
x = self.linear2(x)
|
||||
return x
|
||||
|
||||
|
||||
def _load_projector(projector_path: str) -> _MLPProjector:
|
||||
import mlx.core as mx
|
||||
from safetensors import safe_open
|
||||
|
||||
projector = _MLPProjector()
|
||||
with safe_open(projector_path, framework="numpy") as f:
|
||||
projector.linear1.weight = mx.array(f.get_tensor("linear1.weight"))
|
||||
projector.linear2.weight = mx.array(f.get_tensor("linear2.weight"))
|
||||
return projector
|
||||
|
||||
|
||||
def _sanitize(text: str, special_tokens: dict[str, str]) -> str:
|
||||
for token in special_tokens.values():
|
||||
text = text.replace(token, "")
|
||||
return text
|
||||
|
||||
|
||||
def _format_prompt(query: str, docs: list[str], special_tokens: dict[str, str]) -> str:
|
||||
query = _sanitize(query, special_tokens)
|
||||
docs = [_sanitize(d, special_tokens) for d in docs]
|
||||
|
||||
doc_token = special_tokens["doc_embed_token"]
|
||||
query_token = special_tokens["query_embed_token"]
|
||||
|
||||
prefix = (
|
||||
"<|im_start|>system\n"
|
||||
"You are a search relevance expert who can determine a ranking of the passages based on how relevant they are to the query. "
|
||||
"If the query is a question, how relevant a passage is depends on how well it answers the question. "
|
||||
"If not, try to analyze the intent of the query and assess how well each passage satisfies the intent. "
|
||||
"If an instruction is provided, you should follow the instruction when determining the ranking."
|
||||
"<|im_end|>\n<|im_start|>user\n"
|
||||
)
|
||||
suffix = "<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n"
|
||||
|
||||
body = (
|
||||
f"I will provide you with {len(docs)} passages, each indicated by a numerical identifier. "
|
||||
f"Rank the passages based on their relevance to query: {query}\n"
|
||||
)
|
||||
body += "\n".join(f'<passage id="{i}">\n{doc}{doc_token}\n</passage>' for i, doc in enumerate(docs))
|
||||
body += f"\n<query>\n{query}{query_token}\n</query>"
|
||||
return prefix + body + suffix
|
||||
|
||||
|
||||
class MLXReranker:
|
||||
"""
|
||||
MLX-accelerated jina-reranker-v3 for Apple Silicon.
|
||||
|
||||
Loads the model from a local directory (use huggingface_hub.snapshot_download
|
||||
to fetch jinaai/jina-reranker-v3-mlx if you don't have it already).
|
||||
"""
|
||||
|
||||
_SPECIAL_TOKENS = {
|
||||
"query_embed_token": "<|rerank_token|>",
|
||||
"doc_embed_token": "<|embed_token|>",
|
||||
}
|
||||
_DOC_TOKEN_ID = 151670
|
||||
_QUERY_TOKEN_ID = 151671
|
||||
|
||||
def __init__(self, model_path: str, projector_path: str):
|
||||
from mlx_lm import load
|
||||
|
||||
self.model, self.tokenizer = load(model_path)
|
||||
self.model.eval()
|
||||
self.projector = _load_projector(projector_path)
|
||||
|
||||
def rerank(self, query: str, documents: list[str], top_n: int | None = None) -> list[dict]:
|
||||
"""
|
||||
Rank documents by relevance to a query.
|
||||
|
||||
Returns a list of dicts with keys: document, relevance_score, index.
|
||||
Sorted by descending relevance_score.
|
||||
"""
|
||||
import mlx.core as mx
|
||||
|
||||
prompt = _format_prompt(query, documents, self._SPECIAL_TOKENS)
|
||||
input_ids = self.tokenizer.encode(prompt)
|
||||
hidden_states = self.model.model([input_ids])[0] # [seq_len, hidden_size]
|
||||
|
||||
input_ids_np = np.array(input_ids)
|
||||
query_positions = np.where(input_ids_np == self._QUERY_TOKEN_ID)[0]
|
||||
doc_positions = np.where(input_ids_np == self._DOC_TOKEN_ID)[0]
|
||||
|
||||
if len(query_positions) == 0:
|
||||
raise ValueError("Query embed token not found in prompt")
|
||||
if len(doc_positions) == 0:
|
||||
raise ValueError("Document embed tokens not found in prompt")
|
||||
|
||||
query_hidden = mx.expand_dims(hidden_states[int(query_positions[0])], axis=0)
|
||||
doc_hidden = mx.stack([hidden_states[int(p)] for p in doc_positions])
|
||||
|
||||
query_emb = self.projector(query_hidden) # [1, 512]
|
||||
doc_emb = self.projector(doc_hidden) # [num_docs, 512]
|
||||
|
||||
query_exp = mx.broadcast_to(mx.expand_dims(query_emb, 0), (1, len(documents), 512))
|
||||
doc_exp = mx.expand_dims(doc_emb, 0)
|
||||
|
||||
scores = mx.sum(doc_exp * query_exp, axis=-1) / (
|
||||
mx.sqrt(mx.sum(doc_exp * doc_exp, axis=-1)) * mx.sqrt(mx.sum(query_exp * query_exp, axis=-1))
|
||||
) # [1, num_docs]
|
||||
scores_np = np.array(scores[0])
|
||||
|
||||
order = np.argsort(scores_np)[::-1]
|
||||
n = min(top_n, len(documents)) if top_n is not None else len(documents)
|
||||
return [
|
||||
{
|
||||
"document": documents[order[i]],
|
||||
"relevance_score": float(scores_np[order[i]]),
|
||||
"index": int(order[i]),
|
||||
}
|
||||
for i in range(n)
|
||||
]
|
||||
+27
-2
@@ -48,6 +48,28 @@ _llm_max_concurrent = int(os.getenv(ENV_LLM_MAX_CONCURRENT, str(DEFAULT_LLM_MAX_
|
||||
_global_llm_semaphore = asyncio.Semaphore(_llm_max_concurrent)
|
||||
|
||||
|
||||
def sanitize_llm_output(text: str | None) -> str | None:
|
||||
"""
|
||||
Sanitize text by removing characters that break downstream systems.
|
||||
|
||||
Removes:
|
||||
- ASCII control characters (0x00-0x08, 0x0B-0x0C, 0x0E-0x1F, 0x7F): break
|
||||
json.loads and PostgreSQL UTF-8 encoding; tab (0x09), newline (0x0A), and
|
||||
carriage return (0x0D) are preserved as they are valid in text and JSON.
|
||||
- Unicode surrogates (U+D800-U+DFFF): Invalid in UTF-8, break LLM APIs
|
||||
|
||||
Surrogate characters are used in UTF-16 encoding but cannot be encoded
|
||||
in UTF-8. They can appear in Python strings from improperly decoded data
|
||||
(e.g., from JavaScript or broken files). Control characters commonly appear
|
||||
in LLM output embedded inside JSON string values.
|
||||
"""
|
||||
if text is None:
|
||||
return None
|
||||
if not text:
|
||||
return text
|
||||
return re.sub(r"[\x00-\x08\x0b\x0c\x0e-\x1f\x7f\ud800-\udfff]", "", text)
|
||||
|
||||
|
||||
class OutputTooLongError(Exception):
|
||||
"""
|
||||
Bridge exception raised when LLM output exceeds token limits.
|
||||
@@ -205,7 +227,7 @@ def create_llm_provider(
|
||||
reasoning_effort=reasoning_effort,
|
||||
)
|
||||
|
||||
elif provider_lower in ("openai", "groq", "ollama", "lmstudio"):
|
||||
elif provider_lower in ("openai", "groq", "ollama", "lmstudio", "minimax"):
|
||||
return OpenAICompatibleLLM(
|
||||
provider=provider,
|
||||
api_key=api_key,
|
||||
@@ -274,6 +296,7 @@ class LLMProvider:
|
||||
"openai-codex",
|
||||
"claude-code",
|
||||
"mock",
|
||||
"minimax",
|
||||
]
|
||||
if self.provider not in valid_providers:
|
||||
raise ValueError(f"Invalid LLM provider: {self.provider}. Must be one of: {', '.join(valid_providers)}")
|
||||
@@ -286,6 +309,8 @@ class LLMProvider:
|
||||
self.base_url = "http://localhost:11434/v1"
|
||||
elif self.provider == "lmstudio":
|
||||
self.base_url = "http://localhost:1234/v1"
|
||||
elif self.provider == "minimax":
|
||||
self.base_url = "https://api.minimax.io/v1"
|
||||
|
||||
# Prepare Vertex AI config (if applicable)
|
||||
vertexai_project_id = None
|
||||
@@ -608,7 +633,7 @@ class LLMProvider:
|
||||
# Reduce Claude Agent SDK logging verbosity
|
||||
import logging as sdk_logging
|
||||
|
||||
from claude_agent_sdk import query # noqa: F401
|
||||
from claude_agent_sdk import query # noqa: F401 # type: ignore[unresolved-import]
|
||||
|
||||
sdk_logging.getLogger("claude_agent_sdk").setLevel(sdk_logging.WARNING)
|
||||
sdk_logging.getLogger("claude_agent_sdk._internal").setLevel(sdk_logging.WARNING)
|
||||
+938
-96
File diff suppressed because it is too large
Load Diff
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user