Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
1f397bc8fc | ||
|
|
6d982f9cea |
+2
-34
@@ -2,10 +2,10 @@
|
||||
# Copy this file to .env and fill in your values
|
||||
|
||||
# LLM Configuration (Required)
|
||||
# Supported providers: openai, groq, ollama, gemini, anthropic, lmstudio, vertexai, minimax
|
||||
# Supported providers: openai, groq, ollama, gemini, anthropic, lmstudio, vertexai
|
||||
HINDSIGHT_API_LLM_PROVIDER=openai
|
||||
HINDSIGHT_API_LLM_API_KEY=your-api-key-here
|
||||
HINDSIGHT_API_LLM_MODEL=gpt-4o-mini
|
||||
HINDSIGHT_API_LLM_MODEL=o3-mini
|
||||
HINDSIGHT_API_LLM_BASE_URL=https://api.openai.com/v1
|
||||
|
||||
# Example: Anthropic Claude configuration
|
||||
@@ -20,11 +20,6 @@ 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
|
||||
@@ -36,22 +31,10 @@ HINDSIGHT_API_HOST=0.0.0.0
|
||||
HINDSIGHT_API_PORT=8888
|
||||
HINDSIGHT_API_LOG_LEVEL=info
|
||||
|
||||
# Base Path / Reverse Proxy Support (Optional)
|
||||
# Set these when deploying behind a reverse proxy with path-based routing
|
||||
# Example: To deploy at example.com/hindsight/, set both to "/hindsight"
|
||||
# HINDSIGHT_API_BASE_PATH=/hindsight
|
||||
# NEXT_PUBLIC_BASE_PATH=/hindsight
|
||||
|
||||
# Database (Optional - uses embedded pg0 by default)
|
||||
# HINDSIGHT_API_DATABASE_URL=postgresql://user:pass@host:5432/db
|
||||
# HINDSIGHT_API_DATABASE_SCHEMA=public # PostgreSQL schema name (default: public)
|
||||
|
||||
# Vector Extension (Optional - uses pgvector by default)
|
||||
# Options: "pgvector" (default), "vchord", "pgvectorscale" (DiskANN)
|
||||
# HINDSIGHT_API_VECTOR_EXTENSION=pgvector
|
||||
# For Azure PostgreSQL with DiskANN:
|
||||
# HINDSIGHT_API_VECTOR_EXTENSION=pgvectorscale # Auto-detects pg_diskann on Azure
|
||||
|
||||
# Embeddings Configuration (Optional - uses local by default)
|
||||
# Provider: "local" (default) or "tei" (HuggingFace Text Embeddings Inference)
|
||||
# HINDSIGHT_API_EMBEDDINGS_PROVIDER=local
|
||||
@@ -67,18 +50,3 @@ HINDSIGHT_API_LOG_LEVEL=info
|
||||
# HINDSIGHT_API_RERANKER_LOCAL_MODEL=cross-encoder/ms-marco-MiniLM-L-6-v2
|
||||
# For TEI provider:
|
||||
# HINDSIGHT_API_RERANKER_TEI_URL=http://localhost:8081
|
||||
|
||||
# Observability & Tracing (Optional - disabled by default)
|
||||
# Enable OpenTelemetry tracing for LLM calls (GenAI semantic conventions)
|
||||
# HINDSIGHT_API_OTEL_TRACES_ENABLED=true
|
||||
#
|
||||
# Local development with Grafana LGTM stack (recommended - see scripts/dev/grafana/README.md)
|
||||
# HINDSIGHT_API_OTEL_EXPORTER_OTLP_ENDPOINT=http://localhost:4318
|
||||
#
|
||||
# Cloud backends (Grafana Cloud, Langfuse, DataDog, etc.)
|
||||
# HINDSIGHT_API_OTEL_EXPORTER_OTLP_ENDPOINT=https://your-backend-url
|
||||
# HINDSIGHT_API_OTEL_EXPORTER_OTLP_HEADERS="Authorization=Bearer your-token"
|
||||
#
|
||||
# Custom service name and environment (optional, defaults: hindsight-api, development)
|
||||
# HINDSIGHT_API_OTEL_SERVICE_NAME=hindsight-production
|
||||
# HINDSIGHT_API_OTEL_DEPLOYMENT_ENVIRONMENT=production
|
||||
|
||||
@@ -1,6 +0,0 @@
|
||||
version: 2
|
||||
updates:
|
||||
- package-ecosystem: "github-actions"
|
||||
directory: "/"
|
||||
schedule:
|
||||
interval: "weekly"
|
||||
@@ -21,20 +21,17 @@ jobs:
|
||||
build:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: actions/setup-node@v6
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/setup-node@v4
|
||||
with:
|
||||
node-version: 20
|
||||
cache: npm
|
||||
cache-dependency-path: package-lock.json
|
||||
- uses: astral-sh/setup-uv@v7
|
||||
- uses: astral-sh/setup-uv@v4
|
||||
- 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@v4
|
||||
- uses: actions/upload-pages-artifact@v3
|
||||
with:
|
||||
path: hindsight-docs/build
|
||||
deploy:
|
||||
|
||||
+32
-223
@@ -13,10 +13,10 @@ jobs:
|
||||
id-token: write
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
- name: Install uv
|
||||
uses: astral-sh/setup-uv@v7
|
||||
uses: astral-sh/setup-uv@v5
|
||||
with:
|
||||
enable-cache: true
|
||||
|
||||
@@ -30,20 +30,12 @@ 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-all
|
||||
run: uv build --out-dir dist
|
||||
|
||||
- name: Build hindsight-all-slim
|
||||
working-directory: ./hindsight-all-slim
|
||||
working-directory: ./hindsight
|
||||
run: uv build --out-dir dist
|
||||
|
||||
- name: Build hindsight-litellm
|
||||
@@ -54,27 +46,13 @@ jobs:
|
||||
working-directory: ./hindsight-embed
|
||||
run: uv build --out-dir dist
|
||||
|
||||
- name: Build hindsight-crewai
|
||||
working-directory: ./hindsight-integrations/crewai
|
||||
run: uv build --out-dir dist
|
||||
|
||||
- name: Build hindsight-pydantic-ai
|
||||
working-directory: ./hindsight-integrations/pydantic-ai
|
||||
run: uv build --out-dir dist
|
||||
|
||||
# Publish in order (client and api-slim first, then api/all wrappers which depend on them)
|
||||
# Publish in order (client and api first, then hindsight-all which depends 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:
|
||||
@@ -84,13 +62,7 @@ jobs:
|
||||
- name: Publish hindsight-all to PyPI
|
||||
uses: pypa/gh-action-pypi-publish@release/v1
|
||||
with:
|
||||
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
|
||||
packages-dir: ./hindsight/dist
|
||||
skip-existing: true
|
||||
|
||||
- name: Publish hindsight-litellm to PyPI
|
||||
@@ -105,33 +77,17 @@ jobs:
|
||||
packages-dir: ./hindsight-embed/dist
|
||||
skip-existing: true
|
||||
|
||||
- name: Publish hindsight-crewai to PyPI
|
||||
uses: pypa/gh-action-pypi-publish@release/v1
|
||||
with:
|
||||
packages-dir: ./hindsight-integrations/crewai/dist
|
||||
skip-existing: true
|
||||
|
||||
- name: Publish hindsight-pydantic-ai to PyPI
|
||||
uses: pypa/gh-action-pypi-publish@release/v1
|
||||
with:
|
||||
packages-dir: ./hindsight-integrations/pydantic-ai/dist
|
||||
skip-existing: true
|
||||
|
||||
# Upload artifacts for GitHub release
|
||||
- name: Upload artifacts
|
||||
uses: actions/upload-artifact@v7
|
||||
uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: python-packages
|
||||
path: |
|
||||
hindsight-clients/python/dist/*
|
||||
hindsight-api-slim/dist/*
|
||||
hindsight-api/dist/*
|
||||
hindsight-all/dist/*
|
||||
hindsight-all-slim/dist/*
|
||||
hindsight/dist/*
|
||||
hindsight-integrations/litellm/dist/*
|
||||
hindsight-embed/dist/*
|
||||
hindsight-integrations/crewai/dist/*
|
||||
hindsight-integrations/pydantic-ai/dist/*
|
||||
retention-days: 1
|
||||
|
||||
release-typescript-client:
|
||||
@@ -139,10 +95,10 @@ jobs:
|
||||
environment: npm
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
- name: Set up Node.js
|
||||
uses: actions/setup-node@v6
|
||||
uses: actions/setup-node@v4
|
||||
with:
|
||||
node-version: '20'
|
||||
registry-url: 'https://registry.npmjs.org'
|
||||
@@ -177,7 +133,7 @@ jobs:
|
||||
run: npm pack
|
||||
|
||||
- name: Upload artifacts
|
||||
uses: actions/upload-artifact@v7
|
||||
uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: typescript-client
|
||||
path: hindsight-clients/typescript/*.tgz
|
||||
@@ -188,10 +144,10 @@ jobs:
|
||||
environment: npm
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
- name: Set up Node.js
|
||||
uses: actions/setup-node@v6
|
||||
uses: actions/setup-node@v4
|
||||
with:
|
||||
node-version: '22'
|
||||
registry-url: 'https://registry.npmjs.org'
|
||||
@@ -226,119 +182,21 @@ jobs:
|
||||
run: npm pack
|
||||
|
||||
- name: Upload artifacts
|
||||
uses: actions/upload-artifact@v7
|
||||
uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: openclaw-integration
|
||||
path: hindsight-integrations/openclaw/*.tgz
|
||||
retention-days: 1
|
||||
|
||||
release-ai-sdk-integration:
|
||||
runs-on: ubuntu-latest
|
||||
environment: npm
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
|
||||
- name: Set up Node.js
|
||||
uses: actions/setup-node@v6
|
||||
with:
|
||||
node-version: '22'
|
||||
registry-url: 'https://registry.npmjs.org'
|
||||
|
||||
- name: Install dependencies
|
||||
working-directory: ./hindsight-integrations/ai-sdk
|
||||
run: npm ci
|
||||
|
||||
- name: Build
|
||||
working-directory: ./hindsight-integrations/ai-sdk
|
||||
run: npm run build
|
||||
|
||||
- name: Publish to npm
|
||||
working-directory: ./hindsight-integrations/ai-sdk
|
||||
run: |
|
||||
set +e
|
||||
OUTPUT=$(npm publish --access public 2>&1)
|
||||
EXIT_CODE=$?
|
||||
echo "$OUTPUT"
|
||||
if [ $EXIT_CODE -ne 0 ]; then
|
||||
if echo "$OUTPUT" | grep -q "cannot publish over"; then
|
||||
echo "Package version already published, skipping..."
|
||||
exit 0
|
||||
fi
|
||||
exit $EXIT_CODE
|
||||
fi
|
||||
env:
|
||||
NODE_AUTH_TOKEN: ${{ secrets.NPM_TOKEN }}
|
||||
|
||||
- name: Pack for GitHub release
|
||||
working-directory: ./hindsight-integrations/ai-sdk
|
||||
run: npm pack
|
||||
|
||||
- name: Upload artifacts
|
||||
uses: actions/upload-artifact@v7
|
||||
with:
|
||||
name: ai-sdk-integration
|
||||
path: hindsight-integrations/ai-sdk/*.tgz
|
||||
retention-days: 1
|
||||
|
||||
release-chat-integration:
|
||||
runs-on: ubuntu-latest
|
||||
environment: npm
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
|
||||
- name: Set up Node.js
|
||||
uses: actions/setup-node@v6
|
||||
with:
|
||||
node-version: '22'
|
||||
registry-url: 'https://registry.npmjs.org'
|
||||
|
||||
- name: Install dependencies
|
||||
working-directory: ./hindsight-integrations/chat
|
||||
run: npm ci
|
||||
|
||||
- name: Build
|
||||
working-directory: ./hindsight-integrations/chat
|
||||
run: npm run build
|
||||
|
||||
- name: Publish to npm
|
||||
working-directory: ./hindsight-integrations/chat
|
||||
run: |
|
||||
set +e
|
||||
OUTPUT=$(npm publish --access public 2>&1)
|
||||
EXIT_CODE=$?
|
||||
echo "$OUTPUT"
|
||||
if [ $EXIT_CODE -ne 0 ]; then
|
||||
if echo "$OUTPUT" | grep -q "cannot publish over"; then
|
||||
echo "Package version already published, skipping..."
|
||||
exit 0
|
||||
fi
|
||||
exit $EXIT_CODE
|
||||
fi
|
||||
env:
|
||||
NODE_AUTH_TOKEN: ${{ secrets.NPM_TOKEN }}
|
||||
|
||||
- name: Pack for GitHub release
|
||||
working-directory: ./hindsight-integrations/chat
|
||||
run: npm pack
|
||||
|
||||
- name: Upload artifacts
|
||||
uses: actions/upload-artifact@v7
|
||||
with:
|
||||
name: chat-integration
|
||||
path: hindsight-integrations/chat/*.tgz
|
||||
retention-days: 1
|
||||
|
||||
release-control-plane:
|
||||
runs-on: ubuntu-latest
|
||||
environment: npm
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
- name: Set up Node.js
|
||||
uses: actions/setup-node@v6
|
||||
uses: actions/setup-node@v4
|
||||
with:
|
||||
node-version: '20'
|
||||
registry-url: 'https://registry.npmjs.org'
|
||||
@@ -361,14 +219,11 @@ jobs:
|
||||
- name: Build
|
||||
run: npm run build --workspace=hindsight-control-plane
|
||||
|
||||
- name: Verify standalone build
|
||||
run: test -f hindsight-control-plane/standalone/server.js || (echo 'standalone/server.js missing - build failed' && exit 1)
|
||||
|
||||
- name: Publish to npm
|
||||
working-directory: ./hindsight-control-plane
|
||||
run: |
|
||||
set +e
|
||||
OUTPUT=$(npm publish --access public --ignore-scripts 2>&1)
|
||||
OUTPUT=$(npm publish --access public 2>&1)
|
||||
EXIT_CODE=$?
|
||||
echo "$OUTPUT"
|
||||
if [ $EXIT_CODE -ne 0 ]; then
|
||||
@@ -386,7 +241,7 @@ jobs:
|
||||
run: npm pack
|
||||
|
||||
- name: Upload artifacts
|
||||
uses: actions/upload-artifact@v7
|
||||
uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: control-plane
|
||||
path: hindsight-control-plane/*.tgz
|
||||
@@ -409,13 +264,9 @@ 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@v6
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
- name: Install Rust
|
||||
uses: dtolnay/rust-toolchain@stable
|
||||
@@ -433,14 +284,13 @@ jobs:
|
||||
chmod +x artifacts/${{ matrix.asset_name }}
|
||||
|
||||
- name: Upload artifacts
|
||||
uses: actions/upload-artifact@v7
|
||||
uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: rust-cli-${{ matrix.asset_name }}
|
||||
path: artifacts/${{ matrix.asset_name }}
|
||||
retention-days: 1
|
||||
|
||||
release-docker-images:
|
||||
name: Release Docker (${{ matrix.image_name }}${{ matrix.tag_suffix }})
|
||||
runs-on: ubuntu-latest
|
||||
permissions:
|
||||
contents: read
|
||||
@@ -450,31 +300,13 @@ jobs:
|
||||
include:
|
||||
- target: api-only
|
||||
image_name: hindsight-api
|
||||
tag_suffix: ""
|
||||
build_args: ""
|
||||
- target: api-only
|
||||
image_name: hindsight-api
|
||||
tag_suffix: "-slim"
|
||||
build_args: |
|
||||
INCLUDE_LOCAL_MODELS=false
|
||||
PRELOAD_ML_MODELS=false
|
||||
- target: cp-only
|
||||
image_name: hindsight-control-plane
|
||||
tag_suffix: ""
|
||||
build_args: ""
|
||||
- target: standalone
|
||||
image_name: hindsight
|
||||
tag_suffix: ""
|
||||
build_args: ""
|
||||
- target: standalone
|
||||
image_name: hindsight
|
||||
tag_suffix: "-slim"
|
||||
build_args: |
|
||||
INCLUDE_LOCAL_MODELS=false
|
||||
PRELOAD_ML_MODELS=false
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
- name: Free Disk Space
|
||||
uses: jlumbroso/free-disk-space@main
|
||||
@@ -488,13 +320,13 @@ jobs:
|
||||
swap-storage: true
|
||||
|
||||
- name: Set up QEMU
|
||||
uses: docker/setup-qemu-action@v4
|
||||
uses: docker/setup-qemu-action@v3
|
||||
|
||||
- name: Set up Docker Buildx
|
||||
uses: docker/setup-buildx-action@v4
|
||||
uses: docker/setup-buildx-action@v3
|
||||
|
||||
- name: Log in to GitHub Container Registry
|
||||
uses: docker/login-action@v4
|
||||
uses: docker/login-action@v3
|
||||
with:
|
||||
registry: ghcr.io
|
||||
username: ${{ github.actor }}
|
||||
@@ -506,12 +338,9 @@ jobs:
|
||||
|
||||
- name: Extract metadata for release tags
|
||||
id: meta
|
||||
uses: docker/metadata-action@v6
|
||||
uses: docker/metadata-action@v5
|
||||
with:
|
||||
images: ghcr.io/${{ github.repository_owner }}/${{ matrix.image_name }}
|
||||
flavor: |
|
||||
latest=auto
|
||||
suffix=${{ matrix.tag_suffix }}
|
||||
tags: |
|
||||
type=semver,pattern={{version}},value=${{ steps.get_version.outputs.VERSION }}
|
||||
type=semver,pattern={{major}}.{{minor}},value=${{ steps.get_version.outputs.VERSION }}
|
||||
@@ -522,7 +351,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@v7
|
||||
# uses: docker/build-push-action@v6
|
||||
# with:
|
||||
# context: .
|
||||
# file: docker/standalone/Dockerfile
|
||||
@@ -537,16 +366,15 @@ jobs:
|
||||
# - name: Smoke test - verify container starts
|
||||
# env:
|
||||
# GROQ_API_KEY: ${{ secrets.GROQ_API_KEY }}
|
||||
# run: ./docker/test-image.sh "${{ matrix.image_name }}:test" "${{ matrix.target }}"
|
||||
# run: ./scripts/docker-smoke-test.sh "${{ matrix.image_name }}:test" "${{ matrix.target }}"
|
||||
|
||||
# Build multi-platform and push to release tags
|
||||
- name: Build and push release images
|
||||
uses: docker/build-push-action@v7
|
||||
uses: docker/build-push-action@v6
|
||||
with:
|
||||
context: .
|
||||
file: docker/standalone/Dockerfile
|
||||
target: ${{ matrix.target }}
|
||||
build-args: ${{ matrix.build_args }}
|
||||
push: true
|
||||
platforms: linux/amd64,linux/arm64
|
||||
tags: ${{ steps.meta.outputs.tags }}
|
||||
@@ -559,7 +387,7 @@ jobs:
|
||||
packages: write
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
- name: Install Helm
|
||||
uses: azure/setup-helm@v4
|
||||
@@ -579,7 +407,7 @@ jobs:
|
||||
run: helm push helm-packages/*.tgz oci://ghcr.io/${{ github.repository_owner }}/charts
|
||||
|
||||
- name: Upload artifacts
|
||||
uses: actions/upload-artifact@v7
|
||||
uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: helm-chart
|
||||
path: helm-packages/*.tgz
|
||||
@@ -587,12 +415,12 @@ jobs:
|
||||
|
||||
create-github-release:
|
||||
runs-on: ubuntu-latest
|
||||
needs: [release-python-packages, release-typescript-client, release-openclaw-integration, release-ai-sdk-integration, release-chat-integration, release-control-plane, release-rust-cli, release-docker-images, release-helm-chart]
|
||||
needs: [release-python-packages, release-typescript-client, release-openclaw-integration, release-control-plane, release-rust-cli, release-docker-images, release-helm-chart]
|
||||
permissions:
|
||||
contents: write
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
- name: Extract version from tag
|
||||
id: get_version
|
||||
@@ -616,18 +444,6 @@ jobs:
|
||||
name: openclaw-integration
|
||||
path: ./artifacts/openclaw-integration
|
||||
|
||||
- name: Download AI SDK Integration
|
||||
uses: actions/download-artifact@v4
|
||||
with:
|
||||
name: ai-sdk-integration
|
||||
path: ./artifacts/ai-sdk-integration
|
||||
|
||||
- name: Download Chat Integration
|
||||
uses: actions/download-artifact@v4
|
||||
with:
|
||||
name: chat-integration
|
||||
path: ./artifacts/chat-integration
|
||||
|
||||
- name: Download Control Plane
|
||||
uses: actions/download-artifact@v4
|
||||
with:
|
||||
@@ -663,21 +479,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-all/dist/* release-assets/ || true
|
||||
cp artifacts/python-packages/hindsight-all-slim/dist/* release-assets/ || true
|
||||
cp artifacts/python-packages/hindsight/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-embed/dist/* release-assets/ || true
|
||||
# TypeScript client
|
||||
cp artifacts/typescript-client/*.tgz release-assets/ || true
|
||||
# OpenClaw Integration
|
||||
cp artifacts/openclaw-integration/*.tgz release-assets/ || true
|
||||
# AI SDK Integration
|
||||
cp artifacts/ai-sdk-integration/*.tgz release-assets/ || true
|
||||
# Chat Integration
|
||||
cp artifacts/chat-integration/*.tgz release-assets/ || true
|
||||
# Control Plane
|
||||
cp artifacts/control-plane/*.tgz release-assets/ || true
|
||||
# Rust CLI binaries
|
||||
|
||||
+203
-882
File diff suppressed because it is too large
Load Diff
+1
-4
@@ -46,7 +46,6 @@ hindsight-docs/static/llms-full.txt
|
||||
hindsight-dev/benchmarks/locomo/results/
|
||||
hindsight-dev/benchmarks/longmemeval/results/
|
||||
hindsight-dev/benchmarks/consolidation/results/
|
||||
hindsight-dev/benchmarks/perf/results/
|
||||
benchmarks/results/
|
||||
hindsight-cli/target
|
||||
hindsight-clients/rust/target
|
||||
@@ -54,6 +53,4 @@ hindsight-clients/rust/target
|
||||
whats-next.md
|
||||
TASK.md
|
||||
# Changelog is now tracked in hindsight-docs/src/pages/changelog.md
|
||||
# CHANGELOG.md
|
||||
|
||||
blog-post*
|
||||
# CHANGELOG.md
|
||||
@@ -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-slim && uv run pytest tests/
|
||||
cd hindsight-api && uv run pytest tests/
|
||||
|
||||
# Run specific test file
|
||||
cd hindsight-api-slim && uv run pytest tests/test_http_api_integration.py -v
|
||||
cd hindsight-api && uv run pytest tests/test_http_api_integration.py -v
|
||||
|
||||
# Run single test function
|
||||
cd hindsight-api-slim && uv run pytest tests/test_retain.py::test_retain_simple -v
|
||||
cd hindsight-api && uv run pytest tests/test_retain.py::test_retain_simple -v
|
||||
|
||||
# Lint and format
|
||||
cd hindsight-api-slim && uv run ruff check .
|
||||
cd hindsight-api-slim && uv run ruff format .
|
||||
cd hindsight-api && uv run ruff check .
|
||||
cd hindsight-api && uv run ruff format .
|
||||
|
||||
# Type checking (uses ty - extremely fast type checker from Astral)
|
||||
cd hindsight-api-slim && uv run ty check hindsight_api/
|
||||
cd hindsight-api && uv run ty check hindsight_api/
|
||||
```
|
||||
|
||||
### Control Plane (Next.js)
|
||||
@@ -45,7 +45,6 @@ cd hindsight-control-plane && npm run dev
|
||||
./scripts/dev/start-docs.sh
|
||||
```
|
||||
|
||||
|
||||
### Generating Clients/OpenAPI
|
||||
```bash
|
||||
# Regenerate OpenAPI spec after API changes (REQUIRED after changing endpoints)
|
||||
@@ -57,22 +56,15 @@ cd hindsight-control-plane && npm run dev
|
||||
|
||||
### Benchmarks
|
||||
```bash
|
||||
# Accuracy benchmarks
|
||||
./scripts/benchmarks/run-longmemeval.sh
|
||||
./scripts/benchmarks/run-locomo.sh
|
||||
|
||||
# Performance benchmarks
|
||||
./scripts/benchmarks/run-consolidation.sh
|
||||
./scripts/benchmarks/run-retain-perf.sh --document <path> # Requires API server running
|
||||
|
||||
# Results viewer
|
||||
./scripts/benchmarks/start-visualizer.sh # View results at localhost:8001
|
||||
```
|
||||
|
||||
## Architecture
|
||||
|
||||
### Monorepo Structure
|
||||
- **hindsight-api-slim/**: Core FastAPI server with memory engine (Python, uv)
|
||||
- **hindsight-api/**: 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 +73,9 @@ cd hindsight-control-plane && npm run dev
|
||||
- **hindsight-integrations/**: Framework integrations (LiteLLM, OpenAI)
|
||||
- **hindsight-dev/**: Development tools and benchmarks
|
||||
|
||||
### Core Engine (hindsight-api-slim/hindsight_api/engine/)
|
||||
### Core Engine (hindsight-api/hindsight_api/engine/)
|
||||
- `memory_engine.py`: Main orchestrator (~170KB) for retain/recall/reflect operations
|
||||
- `llm_wrapper.py`: LLM abstraction supporting OpenAI, Anthropic, Gemini, Groq, MiniMax, Ollama, LM Studio
|
||||
- `llm_wrapper.py`: LLM abstraction supporting OpenAI, Anthropic, Gemini, Groq, 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 +93,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-slim/hindsight_api/api/)
|
||||
### API Layer (hindsight-api/hindsight_api/api/)
|
||||
- `http.py`: FastAPI HTTP routers (~80KB) for all REST endpoints
|
||||
- `mcp.py`: Model Context Protocol server implementation
|
||||
|
||||
@@ -111,13 +103,13 @@ Main operations:
|
||||
- **Reflect**: Disposition-aware reasoning using memories and mental models.
|
||||
|
||||
### Database
|
||||
PostgreSQL with pgvector. Schema managed via Alembic migrations in `hindsight-api-slim/hindsight_api/alembic/`. Migrations run automatically on API startup.
|
||||
PostgreSQL with pgvector. Schema managed via Alembic migrations in `hindsight-api/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-slim/hindsight_api/alembic/versions/`:
|
||||
1. **Create a new migration file** in `hindsight-api/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 +146,7 @@ Key tables: `banks`, `memory_units`, `documents`, `entities`, `entity_links`
|
||||
|
||||
3. **Run migrations locally**:
|
||||
```bash
|
||||
# Set database URL and run migrations for the base schema plus all tenants
|
||||
# Set database URL and run migrations
|
||||
uv run hindsight-admin run-db-migration
|
||||
|
||||
# Run on a specific tenant schema
|
||||
@@ -245,61 +237,26 @@ def process(data: UserData) -> str:
|
||||
|
||||
### Adding New API Configuration Flags
|
||||
|
||||
Configuration follows a hierarchical system: **Global (env vars) → Tenant (via extension) → Bank (database)**.
|
||||
When adding a new environment variable configuration:
|
||||
|
||||
Fields must be categorized as either **hierarchical** (can be overridden per-tenant/bank) or **static** (server-level only).
|
||||
|
||||
#### Adding a New Configuration Field
|
||||
|
||||
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"`)
|
||||
1. **config.py** (`hindsight-api/hindsight_api/config.py`):
|
||||
- Add `ENV_*` constant for the environment variable name
|
||||
- Add `DEFAULT_*` constant for the default value
|
||||
- Add field to `HindsightConfig` dataclass with type annotation
|
||||
- **Mark as hierarchical or static** by adding to `_HIERARCHICAL_FIELDS` set (hierarchical) or leaving it out (static)
|
||||
- Add field to `HindsightConfig` dataclass
|
||||
- Add initialization in `from_env()` method
|
||||
|
||||
```python
|
||||
# Hierarchical field (can be overridden per-bank)
|
||||
_HIERARCHICAL_FIELDS = {
|
||||
...,
|
||||
"my_setting", # Add here for hierarchical
|
||||
}
|
||||
|
||||
# Static field - just don't add to _HIERARCHICAL_FIELDS
|
||||
```
|
||||
|
||||
2. **main.py** (`hindsight-api-slim/hindsight_api/main.py`):
|
||||
2. **main.py** (`hindsight-api/hindsight_api/main.py`):
|
||||
- Add field to the manual `HindsightConfig()` constructor call (search for "CLI override")
|
||||
|
||||
3. **Use hierarchical config in MemoryEngine**:
|
||||
```python
|
||||
# Config is resolved automatically per bank via ConfigResolver
|
||||
config_dict = await self._config_resolver.get_bank_config(bank_id, context)
|
||||
value = config_dict["my_setting"]
|
||||
```
|
||||
|
||||
4. **Use static config** (non-hierarchical):
|
||||
3. **Use the config** in code:
|
||||
```python
|
||||
from ...config import get_config
|
||||
config = get_config()
|
||||
value = config.my_static_field
|
||||
value = config.your_new_field
|
||||
```
|
||||
|
||||
5. **Documentation** (`hindsight-docs/docs/developer/configuration.md`):
|
||||
4. **Documentation** (`hindsight-docs/docs/developer/configuration.md`):
|
||||
- Add to appropriate section table with Variable, Description, Default
|
||||
- Mark if it's hierarchical (can be overridden per-bank)
|
||||
|
||||
#### Hierarchical vs Static Guidelines
|
||||
|
||||
**Hierarchical** (per-bank overridable):
|
||||
- LLM settings (provider, model, API key, base URL)
|
||||
- Operation-specific settings (retain mode, chunk size, etc.)
|
||||
- Feature flags that vary by customer/bank
|
||||
|
||||
**Static** (server-level only):
|
||||
- Infrastructure settings (database URL, port, host)
|
||||
- Global limits (max concurrent operations)
|
||||
- System-wide feature flags
|
||||
|
||||
## Environment Setup
|
||||
|
||||
@@ -308,19 +265,18 @@ cp .env.example .env
|
||||
# Edit .env with LLM API key
|
||||
|
||||
# Python deps
|
||||
uv sync --directory hindsight-api-slim/
|
||||
uv sync --directory hindsight-api/
|
||||
|
||||
# Node deps (uses npm workspaces)
|
||||
npm install
|
||||
```
|
||||
|
||||
Required env vars:
|
||||
- `HINDSIGHT_API_LLM_PROVIDER`: openai, anthropic, gemini, groq, minimax, ollama, lmstudio
|
||||
- `HINDSIGHT_API_LLM_PROVIDER`: openai, anthropic, gemini, groq, 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)
|
||||
- `HINDSIGHT_API_LLM_MODEL`: Model name (e.g., o3-mini, claude-sonnet-4-20250514)
|
||||
|
||||
Optional (uses local models by default):
|
||||
- `HINDSIGHT_API_EMBEDDINGS_PROVIDER`: local (default) or tei
|
||||
- `HINDSIGHT_API_RERANKER_PROVIDER`: local (default) or tei
|
||||
- `HINDSIGHT_API_DATABASE_URL`: External PostgreSQL (uses embedded pg0 by default)
|
||||
- `HINDSIGHT_API_ENABLE_BANK_CONFIG_API`: Enable per-bank config API (default: true)
|
||||
|
||||
@@ -2,16 +2,15 @@
|
||||
|
||||

|
||||
|
||||
[Documentation](https://hindsight.vectorize.io) • [Paper](https://arxiv.org/abs/2512.12818) • [Cookbook](https://hindsight.vectorize.io/cookbook) • [Hindsight Cloud](https://ui.hindsight.vectorize.io/signup)
|
||||
[Documentation](https://hindsight.vectorize.io) • [Paper](https://arxiv.org/abs/2512.12818) • [Cookbook](https://hindsight.vectorize.io/cookbook) • [Hindsight Cloud](https://vectorize.io/hindsight/cloud)
|
||||
|
||||
[](https://github.com/vectorize-io/hindsight/actions/workflows/release.yml)
|
||||
[](https://join.slack.com/t/hindsight-space/shared_invite/zt-3nhbm4w29-LeSJ5Ixi6j8PdiYOCPlOgg)
|
||||
[](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>
|
||||
|
||||
---
|
||||
@@ -37,57 +36,33 @@ Hindsight is being used in production at Fortune 500 enterprises and by a growin
|
||||
|
||||
## Adding Hindsight to Your AI Agents
|
||||
|
||||
The easiest way to use Hindsight with an existing agent is with the LLM Wrapper. You can add memory to your agent with 2 lines of code. That will swap your current LLM client out with the Hindsight wrapper. After that, memories will be stored and retrieved automatically as you make LLM calls.
|
||||
The easiest way use Hindsight with an existing agent is with the LLM Wrapper. You can add memory to your agent with 2 lines of code. That will swap your current LLM client out with the Hindsight wrapper. After that, memories will be stored and retrieved automatically as you make LLM calls.
|
||||
|
||||
If you need more control over how and when your agent stores and recalls memories, there's also a simple API you can integrate with using the SDKs or directly via HTTP.
|
||||
|
||||

|
||||
|
||||
---
|
||||
|
||||
> 🤖 **Using a coding agent?** Install the Hindsight documentation skill for instant access to docs while you code:
|
||||
> ```bash
|
||||
> npx skills add https://github.com/vectorize-io/hindsight --skill hindsight-docs
|
||||
> ```
|
||||
> Works with Claude Code, Cursor, and other AI coding assistants.
|
||||
|
||||
---
|
||||
|
||||
|
||||
## Quick Start
|
||||
|
||||
### Docker (recommended)
|
||||
|
||||
```bash
|
||||
export OPENAI_API_KEY=sk-xxx
|
||||
export OPENAI_API_KEY=your-key
|
||||
|
||||
docker run --rm -it --pull always -p 8888:8888 -p 9999:9999 \
|
||||
-e HINDSIGHT_API_LLM_API_KEY=$OPENAI_API_KEY \
|
||||
-e HINDSIGHT_API_LLM_MODEL=o3-mini \
|
||||
-v $HOME/.hindsight-docker:/home/hindsight/.pg0 \
|
||||
ghcr.io/vectorize-io/hindsight:latest
|
||||
```
|
||||
|
||||
>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).
|
||||
API: http://localhost:8888
|
||||
UI: http://localhost:9999
|
||||
|
||||
|
||||
|
||||
### Docker (external PostgreSQL)
|
||||
|
||||
```bash
|
||||
export OPENAI_API_KEY=sk-xxx
|
||||
export HINDSIGHT_DB_PASSWORD=choose-a-password
|
||||
cd docker/docker-compose
|
||||
docker compose up
|
||||
```
|
||||
|
||||
|
||||
>API: http://localhost:8888
|
||||
>UI: http://localhost:9999
|
||||
|
||||
### Client
|
||||
Install client:
|
||||
|
||||
```bash
|
||||
pip install hindsight-client -U
|
||||
@@ -95,7 +70,7 @@ pip install hindsight-client -U
|
||||
npm install @vectorize-io/hindsight-client
|
||||
```
|
||||
|
||||
#### Python
|
||||
Python example:
|
||||
|
||||
```python
|
||||
from hindsight_client import Hindsight
|
||||
@@ -112,29 +87,7 @@ client.recall(bank_id="my-bank", query="What does Alice do?")
|
||||
client.reflect(bank_id="my-bank", query="Tell me about Alice")
|
||||
```
|
||||
|
||||
#### Node.js / TypeScript
|
||||
|
||||
```bash
|
||||
npm install @vectorize-io/hindsight-client
|
||||
```
|
||||
|
||||
```javascript
|
||||
const { HindsightClient } = require('@vectorize-io/hindsight-client');
|
||||
|
||||
const main = async () => {
|
||||
const client = new HindsightClient({ baseUrl: 'http://localhost:8888' });
|
||||
|
||||
await client.retain('my-bank', 'Alice loves hiking in Yosemite');
|
||||
|
||||
const results = await client.recall('my-bank', 'What does Alice like?');
|
||||
console.log(results);
|
||||
}
|
||||
|
||||
main();
|
||||
```
|
||||
|
||||
|
||||
### Python Embedded (no server required)
|
||||
### Python (embedded, no Docker)
|
||||
|
||||
```bash
|
||||
pip install hindsight-all -U
|
||||
@@ -154,6 +107,20 @@ with HindsightServer(
|
||||
results = client.recall(bank_id="my-bank", query="Where does Alice work?")
|
||||
```
|
||||
|
||||
### Node.js / TypeScript
|
||||
|
||||
```bash
|
||||
npm install @vectorize-io/hindsight-client
|
||||
```
|
||||
|
||||
```javascript
|
||||
const { HindsightClient } = require('@vectorize-io/hindsight-client');
|
||||
|
||||
const client = new HindsightClient({ baseUrl: 'http://localhost:8888' });
|
||||
|
||||
await client.retain('my-bank', 'Alice loves hiking in Yosemite');
|
||||
await client.recall('my-bank', 'What does Alice like?');
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
@@ -182,7 +149,7 @@ Satisfying these requirements in Hindsight is straightforward. When new user inp
|
||||
|
||||

|
||||
|
||||
Most agent memory implementations rely on basic vector search or sometimes use a knowledge graph. Hindsight uses biomimetic data structures to organize agent memories in a way that is more like how human memory works:
|
||||
Most agent memory implementation rely on basic vector search or sometimes use a knowledge graph. Hindsight uses biomimetic data structures to organize agent memories in a way that is more like how human memory works:
|
||||
|
||||
- **World:** Facts about the world ("The stove gets hot")
|
||||
- **Experiences:** Agent's own experiences ("I touched the stove and it really hurt")
|
||||
@@ -308,5 +275,3 @@ MIT — see [LICENSE](./LICENSE)
|
||||
---
|
||||
|
||||
Built by [Vectorize.io](https://vectorize.io)
|
||||
|
||||
<img src="https://umami-pixel.chris-latimer.workers.dev/?id=a8b043e6-6964-454d-80df-69b69d3f0d50&host=github.com&url=/vectorize-io/hindsight" width="1" height="1" alt="" />
|
||||
|
||||
@@ -1,54 +0,0 @@
|
||||
# Docker Compose file for Hindsight with PostgreSQL and pgvector
|
||||
#
|
||||
# Make sure to set the required environment variables before running:
|
||||
# - HINDSIGHT_DB_PASSWORD: Password for the PostgreSQL user
|
||||
# - Configure LLM provider variables as needed (see below in the hindsight service)
|
||||
#
|
||||
# Usage:
|
||||
# docker compose up -d
|
||||
#
|
||||
# Optional environment variables with defaults:
|
||||
# - HINDSIGHT_VERSION: Hindsight application version (default: latest)
|
||||
# - HINDSIGHT_DB_USER: PostgreSQL user (default: hindsight_user)
|
||||
# - HINDSIGHT_DB_NAME: PostgreSQL database name (default: hindsight_db)
|
||||
# - HINDSIGHT_DB_VERSION: PostgreSQL version (default: 18)
|
||||
|
||||
services:
|
||||
db:
|
||||
# Use a PostgreSQL-Image with pgvector extension pre-installed
|
||||
# see https://hub.docker.com/r/pgvector/pgvector
|
||||
image: pgvector/pgvector:pg${HINDSIGHT_DB_VERSION:-18}
|
||||
container_name: hindsight-db
|
||||
restart: always
|
||||
# Expose PostgreSQL port
|
||||
# ports:
|
||||
# - "5432:5432"
|
||||
environment:
|
||||
POSTGRES_USER: ${HINDSIGHT_DB_USER:-hindsight_user}
|
||||
POSTGRES_PASSWORD: ${HINDSIGHT_DB_PASSWORD:?Please set the HINDSIGHT_DB_PASSWORD env variable}
|
||||
POSTGRES_DB: ${HINDSIGHT_DB_NAME:-hindsight_db}
|
||||
volumes:
|
||||
- pg_data:/var/lib/postgresql/${HINDSIGHT_DB_VERSION:-18}/docker
|
||||
networks:
|
||||
- hindsight-net
|
||||
|
||||
hindsight:
|
||||
image: ghcr.io/vectorize-io/hindsight:${HINDSIGHT_VERSION:-latest}
|
||||
container_name: hindsight-app
|
||||
ports:
|
||||
- "8888:8888"
|
||||
- "9999:9999"
|
||||
environment:
|
||||
- HINDSIGHT_API_LLM_API_KEY=${OPENAI_API_KEY?Please set the OPENAI_API_KEY env variable}
|
||||
- HINDSIGHT_API_DATABASE_URL=postgresql://${HINDSIGHT_DB_USER:-hindsight_user}:${HINDSIGHT_DB_PASSWORD:?Please set the HINDSIGHT_DB_PASSWORD env variable}@db:5432/${HINDSIGHT_DB_NAME:-hindsight_db}
|
||||
depends_on:
|
||||
- db
|
||||
networks:
|
||||
- hindsight-net
|
||||
|
||||
networks:
|
||||
hindsight-net:
|
||||
driver: bridge
|
||||
|
||||
volumes:
|
||||
pg_data:
|
||||
@@ -1,96 +0,0 @@
|
||||
# Nginx Reverse Proxy with Custom Base Path
|
||||
|
||||
Deploy Hindsight API under `/hindsight` (or any custom path) using Nginx reverse proxy.
|
||||
|
||||
## Quick Start (Published Image - API Only)
|
||||
|
||||
```bash
|
||||
docker-compose up
|
||||
```
|
||||
|
||||
- **API:** http://localhost:8080/hindsight/docs
|
||||
- **Control Plane:** http://localhost:9999 (direct access, not proxied)
|
||||
|
||||
## Full Stack with Custom Base Path (Requires Build)
|
||||
|
||||
**Important:** You cannot rebuild from the published image with build args. You must build from source.
|
||||
|
||||
### Build from Source with Custom Base Path
|
||||
|
||||
1. **Clone the repository** (if you haven't):
|
||||
```bash
|
||||
git clone https://github.com/vectorize-io/hindsight.git
|
||||
cd hindsight
|
||||
```
|
||||
|
||||
2. **Build with base path**:
|
||||
```bash
|
||||
docker build \
|
||||
--build-arg NEXT_PUBLIC_BASE_PATH=/hindsight \
|
||||
-f docker/standalone/Dockerfile \
|
||||
-t hindsight:custom \
|
||||
.
|
||||
```
|
||||
|
||||
3. **Update docker-compose.yml** to use your built image:
|
||||
```yaml
|
||||
services:
|
||||
hindsight:
|
||||
image: hindsight:custom # ← Change this
|
||||
environment:
|
||||
HINDSIGHT_API_BASE_PATH: /hindsight
|
||||
NEXT_PUBLIC_BASE_PATH: /hindsight
|
||||
```
|
||||
|
||||
4. **Update nginx.conf** to handle Control Plane routes (see below)
|
||||
|
||||
5. **Run**:
|
||||
```bash
|
||||
docker-compose up
|
||||
```
|
||||
|
||||
### Required nginx.conf for Full Stack
|
||||
|
||||
Replace the current `nginx.conf` with this to proxy both API and Control Plane:
|
||||
|
||||
```nginx
|
||||
events { worker_connections 1024; }
|
||||
|
||||
http {
|
||||
include /etc/nginx/mime.types;
|
||||
default_type application/octet-stream;
|
||||
|
||||
upstream hindsight_api { server hindsight:8888; }
|
||||
upstream hindsight_cp { server hindsight:9999; }
|
||||
|
||||
server {
|
||||
listen 80;
|
||||
|
||||
# API
|
||||
location ~ ^/hindsight/(docs|openapi\.json|health|metrics|v1|mcp) {
|
||||
proxy_pass http://hindsight_api;
|
||||
proxy_set_header Host $http_host;
|
||||
}
|
||||
|
||||
# Control Plane static files
|
||||
location ~ ^/hindsight/_next/ {
|
||||
proxy_pass http://hindsight_cp;
|
||||
proxy_set_header Host $http_host;
|
||||
}
|
||||
|
||||
# Control Plane UI
|
||||
location /hindsight {
|
||||
proxy_pass http://hindsight_cp;
|
||||
proxy_set_header Host $http_host;
|
||||
}
|
||||
|
||||
location = / { return 301 /hindsight; }
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
### Why Build is Required
|
||||
|
||||
Next.js requires `basePath` at **build time**. The published image was built without a custom base path, so you must rebuild from source with the `NEXT_PUBLIC_BASE_PATH` build arg to deploy the Control Plane under a subpath.
|
||||
|
||||
The API works without rebuild because `HINDSIGHT_API_BASE_PATH` is a runtime environment variable.
|
||||
@@ -1,88 +0,0 @@
|
||||
# Hindsight API deployment with Nginx reverse proxy (API-only)
|
||||
#
|
||||
# This example deploys Hindsight API under the path /hindsight with:
|
||||
# - Hindsight standalone image (API + Control Plane + embedded pg0)
|
||||
# - Nginx reverse proxy (API only)
|
||||
#
|
||||
# Quick Start:
|
||||
# docker-compose -f docker/docker-compose/nginx/docker-compose.yml up
|
||||
#
|
||||
# Access:
|
||||
# API (via nginx): http://localhost:8080/hindsight/docs
|
||||
# Control Plane (direct): http://localhost:9999
|
||||
#
|
||||
# For full stack deployment (API + Control Plane both under /hindsight):
|
||||
# See README.md in this directory for instructions on building with basePath.
|
||||
#
|
||||
# Note: This configuration uses the published image (no build required).
|
||||
# Control Plane is served directly because Next.js basePath requires
|
||||
# build-time configuration. See README.md for the full stack option.
|
||||
|
||||
services:
|
||||
# Hindsight (API + Control Plane + embedded pg0)
|
||||
hindsight:
|
||||
image: ghcr.io/vectorize-io/hindsight:latest
|
||||
ports:
|
||||
- "9999:9999" # Control Plane (direct access, not proxied)
|
||||
environment:
|
||||
# API base path for reverse proxy
|
||||
HINDSIGHT_API_BASE_PATH: /hindsight
|
||||
|
||||
# LLM configuration
|
||||
# Using mock provider for testing (no API key needed)
|
||||
# For production, set OPENAI_API_KEY or ANTHROPIC_API_KEY and use a real provider
|
||||
HINDSIGHT_API_LLM_PROVIDER: ${HINDSIGHT_API_LLM_PROVIDER:-mock}
|
||||
HINDSIGHT_API_LLM_API_KEY: ${OPENAI_API_KEY:-not-needed-for-mock}
|
||||
HINDSIGHT_API_LLM_MODEL: ${HINDSIGHT_API_LLM_MODEL:-mock-model}
|
||||
|
||||
# Production examples (uncomment and set appropriate API key):
|
||||
# HINDSIGHT_API_LLM_PROVIDER: openai
|
||||
# HINDSIGHT_API_LLM_API_KEY: ${OPENAI_API_KEY}
|
||||
# HINDSIGHT_API_LLM_MODEL: gpt-4o-mini
|
||||
|
||||
# HINDSIGHT_API_LLM_PROVIDER: anthropic
|
||||
# HINDSIGHT_API_LLM_API_KEY: ${ANTHROPIC_API_KEY}
|
||||
# HINDSIGHT_API_LLM_MODEL: claude-sonnet-4-20250514
|
||||
|
||||
# Server config
|
||||
HINDSIGHT_API_HOST: 0.0.0.0
|
||||
HINDSIGHT_API_PORT: 8888
|
||||
HINDSIGHT_API_LOG_LEVEL: info
|
||||
|
||||
# Control Plane config
|
||||
HINDSIGHT_CP_DATAPLANE_API_URL: http://localhost:8888
|
||||
volumes:
|
||||
# Persist embedded pg0 database
|
||||
- hindsight_data:/app/data
|
||||
# Note: Ports not exposed - access via Nginx at localhost:8080/hindsight/
|
||||
# To debug directly, uncomment these ports:
|
||||
# ports:
|
||||
# - "8888:8888" # API
|
||||
# - "9999:9999" # Control Plane
|
||||
healthcheck:
|
||||
test: ["CMD", "curl", "-f", "http://localhost:8888/hindsight/health"]
|
||||
interval: 10s
|
||||
timeout: 5s
|
||||
retries: 3
|
||||
start_period: 30s
|
||||
networks:
|
||||
- hindsight
|
||||
|
||||
# Nginx reverse proxy
|
||||
nginx:
|
||||
image: nginx:alpine
|
||||
ports:
|
||||
- "8080:80"
|
||||
volumes:
|
||||
- ./nginx.conf:/etc/nginx/nginx.conf:ro
|
||||
depends_on:
|
||||
hindsight:
|
||||
condition: service_healthy
|
||||
networks:
|
||||
- hindsight
|
||||
|
||||
volumes:
|
||||
hindsight_data:
|
||||
|
||||
networks:
|
||||
hindsight:
|
||||
@@ -1,40 +0,0 @@
|
||||
# Nginx configuration for API-only reverse proxy
|
||||
# Control Plane accessed directly (not through nginx)
|
||||
|
||||
events {
|
||||
worker_connections 1024;
|
||||
}
|
||||
|
||||
http {
|
||||
include /etc/nginx/mime.types;
|
||||
default_type application/octet-stream;
|
||||
|
||||
# Logging
|
||||
access_log /var/log/nginx/access.log;
|
||||
error_log /var/log/nginx/error.log;
|
||||
|
||||
# Upstream - Hindsight API
|
||||
upstream hindsight_api {
|
||||
server hindsight:8888;
|
||||
}
|
||||
|
||||
server {
|
||||
listen 80;
|
||||
server_name _;
|
||||
|
||||
# API endpoints - forward with /hindsight prefix
|
||||
location /hindsight/ {
|
||||
proxy_pass http://hindsight_api;
|
||||
|
||||
proxy_set_header Host $http_host;
|
||||
proxy_set_header X-Real-IP $remote_addr;
|
||||
proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for;
|
||||
proxy_set_header X-Forwarded-Proto $scheme;
|
||||
}
|
||||
|
||||
# Redirect root to API docs
|
||||
location = / {
|
||||
return 301 /hindsight/docs;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,32 +0,0 @@
|
||||
# PostgreSQL with pgvector and pg_textsearch extensions
|
||||
# Note: pg_textsearch requires PostgreSQL 17+
|
||||
FROM postgres:17
|
||||
|
||||
# Install build dependencies
|
||||
RUN apt-get update && apt-get install -y \
|
||||
build-essential \
|
||||
git \
|
||||
postgresql-server-dev-17 \
|
||||
libpq-dev \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
# Install pgvector
|
||||
RUN cd /tmp && \
|
||||
git clone --branch v0.8.0 https://github.com/pgvector/pgvector.git && \
|
||||
cd pgvector && \
|
||||
make && \
|
||||
make install
|
||||
|
||||
# Install pg_textsearch
|
||||
RUN cd /tmp && \
|
||||
git clone https://github.com/timescale/pg_textsearch.git && \
|
||||
cd pg_textsearch && \
|
||||
make && \
|
||||
make install
|
||||
|
||||
# Clean up source files and build dependencies
|
||||
RUN rm -rf /tmp/pgvector /tmp/pg_textsearch && \
|
||||
apt-get purge -y --auto-remove build-essential git postgresql-server-dev-17
|
||||
|
||||
# Ensure extensions are preloaded
|
||||
RUN echo "shared_preload_libraries = 'pg_textsearch'" >> /usr/share/postgresql/postgresql.conf.sample
|
||||
@@ -1,91 +0,0 @@
|
||||
name: hindsight
|
||||
# Docker Compose file for Hindsight with PostgreSQL and Timescale pg_textsearch
|
||||
# docker compose -f docker/docker-compose/pg_textsearch/docker-compose.yaml down && sleep 2 && docker compose -f docker/docker-compose/pg_textsearch/docker-compose.yaml up -d
|
||||
# Make sure to set the required environment variables before running:
|
||||
# - HINDSIGHT_DB_PASSWORD: Password for the PostgreSQL user
|
||||
# - Configure LLM provider variables as needed (see below in the hindsight service)
|
||||
#
|
||||
# Usage:
|
||||
# docker compose up -d
|
||||
#
|
||||
# Optional environment variables with defaults:
|
||||
# - HINDSIGHT_VERSION: Hindsight application version (default: latest)
|
||||
# - HINDSIGHT_DB_USER: PostgreSQL user (default: hindsight_user)
|
||||
# - HINDSIGHT_DB_NAME: PostgreSQL database name (default: hindsight_db)
|
||||
|
||||
services:
|
||||
db:
|
||||
# Use custom PostgreSQL image with pgvector and pg_textsearch extensions
|
||||
build:
|
||||
context: .
|
||||
dockerfile: Dockerfile
|
||||
container_name: hindsight-db
|
||||
restart: always
|
||||
# Expose PostgreSQL port
|
||||
ports:
|
||||
- "5437:5432"
|
||||
environment:
|
||||
POSTGRES_USER: ${HINDSIGHT_DB_USER:-hindsight_user}
|
||||
POSTGRES_PASSWORD: ${HINDSIGHT_DB_PASSWORD:-hindsight_password}
|
||||
POSTGRES_DB: ${HINDSIGHT_DB_NAME:-hindsight_db}
|
||||
volumes:
|
||||
- pg_data:/var/lib/postgresql/data
|
||||
networks:
|
||||
- hindsight-net
|
||||
|
||||
pg-textsearch-init:
|
||||
build:
|
||||
context: .
|
||||
dockerfile: Dockerfile
|
||||
depends_on:
|
||||
- db
|
||||
environment:
|
||||
- PGPASSWORD=${HINDSIGHT_DB_PASSWORD:-hindsight_password}
|
||||
command: >
|
||||
bash -c "
|
||||
echo 'Waiting for PostgreSQL to be ready...';
|
||||
until pg_isready -h hindsight-db -p 5432 -U hindsight_user; do
|
||||
echo 'PostgreSQL is unavailable - sleeping';
|
||||
sleep 2;
|
||||
done;
|
||||
echo 'PostgreSQL is ready - creating hindsight_db database';
|
||||
psql -h hindsight-db -p 5432 -U hindsight_user -c 'CREATE DATABASE hindsight_db;' 2>/dev/null || echo 'Database already exists';
|
||||
echo 'Creating extensions in hindsight_db database';
|
||||
psql -h hindsight-db -p 5432 -U hindsight_user -d hindsight_db -c 'CREATE EXTENSION IF NOT EXISTS vector CASCADE;';
|
||||
psql -h hindsight-db -p 5432 -U hindsight_user -d hindsight_db -c 'CREATE EXTENSION IF NOT EXISTS pg_textsearch CASCADE;';
|
||||
echo 'Database and extensions created successfully';
|
||||
"
|
||||
restart: "no"
|
||||
networks:
|
||||
- hindsight-net
|
||||
|
||||
hindsight:
|
||||
image: ghcr.io/vectorize-io/hindsight:${HINDSIGHT_VERSION:-latest}
|
||||
container_name: hindsight-app
|
||||
ports:
|
||||
- "8888:8888"
|
||||
- "9999:9999"
|
||||
environment:
|
||||
# LLM Configuration
|
||||
HINDSIGHT_API_LLM_PROVIDER: ${HINDSIGHT_API_LLM_PROVIDER:-openai}
|
||||
HINDSIGHT_API_LLM_API_KEY: ${OPENAI_API_KEY:-your-api-key}
|
||||
|
||||
# Database Configuration
|
||||
HINDSIGHT_API_DATABASE_URL: postgresql://${HINDSIGHT_DB_USER:-hindsight_user}:${HINDSIGHT_DB_PASSWORD:-hindsight_password}@db:5432/${HINDSIGHT_DB_NAME:-hindsight_db}
|
||||
|
||||
# Vector and Text Search Extensions
|
||||
HINDSIGHT_API_VECTOR_EXTENSION: pgvector
|
||||
HINDSIGHT_API_TEXT_SEARCH_EXTENSION: pg_textsearch
|
||||
|
||||
depends_on:
|
||||
- db
|
||||
networks:
|
||||
- hindsight-net
|
||||
|
||||
|
||||
networks:
|
||||
hindsight-net:
|
||||
driver: bridge
|
||||
|
||||
volumes:
|
||||
pg_data:
|
||||
@@ -1,83 +0,0 @@
|
||||
# Docker Compose file for Hindsight with S3 file storage (SeaweedFS)
|
||||
#
|
||||
# SeaweedFS (Apache 2.0) provides an S3-compatible object storage backend
|
||||
# for storing uploaded files instead of PostgreSQL BYTEA storage.
|
||||
#
|
||||
# Make sure to set the required environment variables before running:
|
||||
# - HINDSIGHT_DB_PASSWORD: Password for the PostgreSQL user
|
||||
# - Configure LLM provider variables as needed (see below in the hindsight service)
|
||||
#
|
||||
# Usage:
|
||||
# docker compose up -d
|
||||
#
|
||||
# Optional environment variables with defaults:
|
||||
# - HINDSIGHT_VERSION: Hindsight application version (default: latest)
|
||||
# - HINDSIGHT_DB_USER: PostgreSQL user (default: hindsight_user)
|
||||
# - HINDSIGHT_DB_NAME: PostgreSQL database name (default: hindsight_db)
|
||||
# - HINDSIGHT_DB_VERSION: PostgreSQL version (default: 18)
|
||||
# - SEAWEEDFS_S3_ACCESS_KEY: S3 access key (default: hindsight_s3_key)
|
||||
# - SEAWEEDFS_S3_SECRET_KEY: S3 secret key (default: hindsight_s3_secret)
|
||||
|
||||
services:
|
||||
db:
|
||||
image: pgvector/pgvector:pg${HINDSIGHT_DB_VERSION:-18}
|
||||
container_name: hindsight-db
|
||||
restart: always
|
||||
environment:
|
||||
POSTGRES_USER: ${HINDSIGHT_DB_USER:-hindsight_user}
|
||||
POSTGRES_PASSWORD: ${HINDSIGHT_DB_PASSWORD:?Please set the HINDSIGHT_DB_PASSWORD env variable}
|
||||
POSTGRES_DB: ${HINDSIGHT_DB_NAME:-hindsight_db}
|
||||
volumes:
|
||||
- pg_data:/var/lib/postgresql/${HINDSIGHT_DB_VERSION:-18}/docker
|
||||
networks:
|
||||
- hindsight-net
|
||||
|
||||
seaweedfs:
|
||||
image: chrislusf/seaweedfs:latest
|
||||
container_name: hindsight-seaweedfs
|
||||
restart: always
|
||||
# Single-node mode: master + volume + filer + S3 gateway all in one process
|
||||
command: >
|
||||
server
|
||||
-s3
|
||||
-s3.port=8333
|
||||
-s3.config=/etc/seaweedfs/s3.json
|
||||
-ip.bind=0.0.0.0
|
||||
volumes:
|
||||
- seaweedfs_data:/data
|
||||
- ./s3.json:/etc/seaweedfs/s3.json:ro
|
||||
# Expose S3 API port (uncomment to access from host)
|
||||
# ports:
|
||||
# - "8333:8333"
|
||||
networks:
|
||||
- hindsight-net
|
||||
|
||||
hindsight:
|
||||
image: ghcr.io/vectorize-io/hindsight:${HINDSIGHT_VERSION:-latest}
|
||||
container_name: hindsight-app
|
||||
ports:
|
||||
- "8888:8888"
|
||||
- "9999:9999"
|
||||
environment:
|
||||
- HINDSIGHT_API_LLM_API_KEY=${OPENAI_API_KEY?Please set the OPENAI_API_KEY env variable}
|
||||
- HINDSIGHT_API_DATABASE_URL=postgresql://${HINDSIGHT_DB_USER:-hindsight_user}:${HINDSIGHT_DB_PASSWORD:?Please set the HINDSIGHT_DB_PASSWORD env variable}@db:5432/${HINDSIGHT_DB_NAME:-hindsight_db}
|
||||
# S3 file storage configuration (SeaweedFS)
|
||||
- HINDSIGHT_API_FILE_STORAGE_TYPE=s3
|
||||
- HINDSIGHT_API_FILE_STORAGE_S3_BUCKET=hindsight
|
||||
- HINDSIGHT_API_FILE_STORAGE_S3_ENDPOINT=http://seaweedfs:8333
|
||||
- HINDSIGHT_API_FILE_STORAGE_S3_REGION=us-east-1
|
||||
- HINDSIGHT_API_FILE_STORAGE_S3_ACCESS_KEY_ID=${SEAWEEDFS_S3_ACCESS_KEY:-hindsight_s3_key}
|
||||
- HINDSIGHT_API_FILE_STORAGE_S3_SECRET_ACCESS_KEY=${SEAWEEDFS_S3_SECRET_KEY:-hindsight_s3_secret}
|
||||
depends_on:
|
||||
- db
|
||||
- seaweedfs
|
||||
networks:
|
||||
- hindsight-net
|
||||
|
||||
networks:
|
||||
hindsight-net:
|
||||
driver: bridge
|
||||
|
||||
volumes:
|
||||
pg_data:
|
||||
seaweedfs_data:
|
||||
@@ -1,19 +0,0 @@
|
||||
{
|
||||
"identities": [
|
||||
{
|
||||
"name": "hindsight",
|
||||
"credentials": [
|
||||
{
|
||||
"accessKey": "hindsight_s3_key",
|
||||
"secretKey": "hindsight_s3_secret"
|
||||
}
|
||||
],
|
||||
"actions": [
|
||||
"Admin",
|
||||
"Read",
|
||||
"Write",
|
||||
"List"
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -1,16 +0,0 @@
|
||||
# Git
|
||||
.git
|
||||
.gitignore
|
||||
.gitattributes
|
||||
|
||||
# Docker
|
||||
docker-compose.yaml
|
||||
.dockerignore
|
||||
|
||||
# Documentation
|
||||
README.md
|
||||
*.md
|
||||
|
||||
# Environment
|
||||
.env
|
||||
.env.example
|
||||
@@ -1,25 +0,0 @@
|
||||
# PostgreSQL Configuration
|
||||
HINDSIGHT_DB_USER=hindsight_user
|
||||
HINDSIGHT_DB_PASSWORD=change-me-to-secure-password
|
||||
HINDSIGHT_DB_NAME=hindsight_db
|
||||
|
||||
# Hindsight Version
|
||||
HINDSIGHT_VERSION=latest
|
||||
|
||||
# LLM Configuration
|
||||
HINDSIGHT_API_LLM_PROVIDER=openai
|
||||
OPENAI_API_KEY=your-openai-api-key-here
|
||||
|
||||
# Alternative LLM providers (uncomment and configure as needed):
|
||||
# HINDSIGHT_API_LLM_PROVIDER=anthropic
|
||||
# ANTHROPIC_API_KEY=your-anthropic-api-key
|
||||
|
||||
# HINDSIGHT_API_LLM_PROVIDER=gemini
|
||||
# GEMINI_API_KEY=your-gemini-api-key
|
||||
|
||||
# HINDSIGHT_API_LLM_PROVIDER=groq
|
||||
# GROQ_API_KEY=your-groq-api-key
|
||||
|
||||
# Vector and Text Search (already configured in docker-compose.yaml)
|
||||
# HINDSIGHT_API_VECTOR_EXTENSION=pgvectorscale
|
||||
# HINDSIGHT_API_TEXT_SEARCH_EXTENSION=pg_textsearch
|
||||
@@ -1,55 +0,0 @@
|
||||
# PostgreSQL with pgvector, pgvectorscale, and pg_textsearch extensions
|
||||
# All three extensions from Timescale/pgvector for high-performance vector and text search
|
||||
# Note: Requires PostgreSQL 16+
|
||||
FROM postgres:17
|
||||
|
||||
# Install build dependencies and Rust toolchain
|
||||
RUN apt-get update && apt-get install -y \
|
||||
build-essential \
|
||||
git \
|
||||
postgresql-server-dev-17 \
|
||||
libpq-dev \
|
||||
cmake \
|
||||
curl \
|
||||
pkg-config \
|
||||
libssl-dev \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
# Install Rust toolchain (required for pgvectorscale)
|
||||
RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y
|
||||
ENV PATH="/root/.cargo/bin:${PATH}"
|
||||
|
||||
# Install pgvector (required by pgvectorscale)
|
||||
RUN cd /tmp && \
|
||||
git clone --branch v0.8.0 https://github.com/pgvector/pgvector.git && \
|
||||
cd pgvector && \
|
||||
make && \
|
||||
make install && \
|
||||
rm -rf /tmp/pgvector
|
||||
|
||||
# Install cargo-pgrx (PostgreSQL extension framework for Rust)
|
||||
RUN cargo install cargo-pgrx --version 0.12.5 --locked && \
|
||||
cargo pgrx init --pg17 /usr/bin/pg_config
|
||||
|
||||
# Install pgvectorscale (DiskANN index support)
|
||||
RUN cd /tmp && \
|
||||
git clone --branch 0.5.1 https://github.com/timescale/pgvectorscale.git && \
|
||||
cd pgvectorscale/pgvectorscale && \
|
||||
cargo pgrx install --release && \
|
||||
rm -rf /tmp/pgvectorscale
|
||||
|
||||
# Install pg_textsearch (BM25 text search)
|
||||
RUN cd /tmp && \
|
||||
git clone https://github.com/timescale/pg_textsearch.git && \
|
||||
cd pg_textsearch && \
|
||||
make && \
|
||||
make install && \
|
||||
rm -rf /tmp/pg_textsearch
|
||||
|
||||
# Clean up build dependencies (keep runtime dependencies)
|
||||
RUN apt-get purge -y --auto-remove git cmake curl && \
|
||||
rm -rf /root/.cargo/registry /root/.cargo/git
|
||||
|
||||
# Ensure extensions are preloaded (pg_textsearch requires preloading)
|
||||
RUN echo "shared_preload_libraries = 'pg_textsearch'" >> /usr/share/postgresql/postgresql.conf.sample
|
||||
|
||||
@@ -1,101 +0,0 @@
|
||||
# Hindsight with Timescale Extensions
|
||||
|
||||
This Docker Compose setup provides a complete Hindsight deployment with **Timescale extensions**:
|
||||
- **pgvectorscale** - DiskANN algorithm for disk-based scalable vector search
|
||||
- **pg_textsearch** - High-performance BM25 text search
|
||||
|
||||
Both extensions are from [Timescale](https://github.com/timescale) and provide production-grade performance.
|
||||
|
||||
## Prerequisites
|
||||
|
||||
- Docker and Docker Compose installed
|
||||
- OpenAI API key (or another LLM provider)
|
||||
|
||||
## Quick Start
|
||||
|
||||
```bash
|
||||
# Set environment variables
|
||||
export HINDSIGHT_DB_PASSWORD="your-secure-password"
|
||||
export OPENAI_API_KEY="your-openai-api-key"
|
||||
|
||||
# Build and start
|
||||
docker compose -f docker/docker-compose/timescale/docker-compose.yaml up -d --build
|
||||
|
||||
# Check logs
|
||||
|
||||
docker compose -f docker/docker-compose/timescale/docker-compose.yaml logs -f
|
||||
```
|
||||
|
||||
**Access:**
|
||||
- API: http://localhost:8888
|
||||
- Control Plane: http://localhost:9999
|
||||
|
||||
## Stop and Clean Up
|
||||
|
||||
```bash
|
||||
# Stop services
|
||||
docker compose -f docker/docker-compose/timescale/docker-compose.yaml down
|
||||
|
||||
# Remove volumes (deletes all data)
|
||||
docker compose -f docker/docker-compose/timescale/docker-compose.yaml down -v
|
||||
```
|
||||
|
||||
## Configuration
|
||||
|
||||
### Environment Variables
|
||||
|
||||
| Variable | Description | Default |
|
||||
|----------|-------------|---------|
|
||||
| `HINDSIGHT_DB_PASSWORD` | PostgreSQL password | `hindsight_password` |
|
||||
| `HINDSIGHT_DB_USER` | PostgreSQL username | `hindsight_user` |
|
||||
| `HINDSIGHT_DB_NAME` | Database name | `hindsight_db` |
|
||||
| `HINDSIGHT_VERSION` | Hindsight Docker image version | `latest` |
|
||||
| `OPENAI_API_KEY` | OpenAI API key | (required) |
|
||||
| `HINDSIGHT_API_LLM_PROVIDER` | LLM provider | `openai` |
|
||||
|
||||
### Why Timescale Extensions?
|
||||
|
||||
**pgvectorscale (DiskANN):**
|
||||
- 28x lower p95 latency vs dedicated vector databases
|
||||
- 16x higher query throughput at 99% recall
|
||||
- 60-75% cost reduction (disk is cheaper than RAM)
|
||||
- Best for large datasets (10M+ vectors)
|
||||
|
||||
**pg_textsearch (BM25):**
|
||||
- High-performance keyword retrieval
|
||||
- Native BM25 ranking algorithm
|
||||
- Optimized for full-text search
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
### Extensions not installed
|
||||
|
||||
Check if extensions are available:
|
||||
|
||||
```bash
|
||||
docker exec -it hindsight-db-timescale psql -U hindsight_user -d hindsight_db -c "\dx"
|
||||
```
|
||||
|
||||
You should see:
|
||||
- `vector` (pgvector)
|
||||
- `vectorscale` (pgvectorscale/DiskANN)
|
||||
- `pg_textsearch` (BM25 search)
|
||||
|
||||
### Build fails
|
||||
|
||||
If the Docker build fails during pgvectorscale compilation:
|
||||
|
||||
1. Ensure you have sufficient memory (recommended: 4GB+)
|
||||
2. Check Docker build logs for Rust compilation errors
|
||||
3. Try building with more resources: `docker compose build --no-cache --memory 4g`
|
||||
|
||||
### Port conflicts
|
||||
|
||||
If port 5438 is already in use, modify the `ports` section in docker-compose.yaml.
|
||||
|
||||
## Learn More
|
||||
|
||||
- [pgvectorscale GitHub](https://github.com/timescale/pgvectorscale)
|
||||
- [pg_textsearch GitHub](https://github.com/timescale/pg_textsearch)
|
||||
- [HNSW vs DiskANN](https://www.tigerdata.com/learn/hnsw-vs-diskann)
|
||||
- [Hindsight Documentation](https://hindsight.dev)
|
||||
@@ -1,108 +0,0 @@
|
||||
name: hindsight
|
||||
# Docker Compose file for Hindsight with Timescale extensions
|
||||
# - pgvectorscale: DiskANN vector search (disk-based, scalable)
|
||||
# - pg_textsearch: BM25 text search (high-performance keyword retrieval)
|
||||
#
|
||||
# Quick start:
|
||||
# docker compose -f docker/docker-compose/timescale/docker-compose.yaml up -d --build
|
||||
#
|
||||
# Required environment variables:
|
||||
# - HINDSIGHT_DB_PASSWORD: Password for the PostgreSQL user
|
||||
# - OPENAI_API_KEY (or configure another LLM provider)
|
||||
#
|
||||
# Optional environment variables with defaults:
|
||||
# - HINDSIGHT_VERSION: Hindsight application version (default: latest)
|
||||
# - HINDSIGHT_DB_USER: PostgreSQL user (default: hindsight_user)
|
||||
# - HINDSIGHT_DB_NAME: PostgreSQL database name (default: hindsight_db)
|
||||
|
||||
services:
|
||||
db:
|
||||
# Custom PostgreSQL image with Timescale extensions (pgvectorscale + pg_textsearch)
|
||||
build:
|
||||
context: .
|
||||
dockerfile: Dockerfile
|
||||
container_name: hindsight-db-timescale
|
||||
restart: always
|
||||
# Expose PostgreSQL port (using 5438 to avoid conflicts with other setups)
|
||||
ports:
|
||||
- "5438:5432"
|
||||
environment:
|
||||
POSTGRES_USER: ${HINDSIGHT_DB_USER:-hindsight_user}
|
||||
POSTGRES_PASSWORD: ${HINDSIGHT_DB_PASSWORD:-hindsight_password}
|
||||
POSTGRES_DB: ${HINDSIGHT_DB_NAME:-hindsight_db}
|
||||
volumes:
|
||||
- pg_data:/var/lib/postgresql/data
|
||||
networks:
|
||||
- hindsight-net
|
||||
# Health check to ensure database is ready
|
||||
healthcheck:
|
||||
test: ["CMD-SHELL", "pg_isready -U hindsight_user"]
|
||||
interval: 5s
|
||||
timeout: 5s
|
||||
retries: 5
|
||||
|
||||
timescale-init:
|
||||
build:
|
||||
context: .
|
||||
dockerfile: Dockerfile
|
||||
depends_on:
|
||||
db:
|
||||
condition: service_healthy
|
||||
environment:
|
||||
- PGPASSWORD=${HINDSIGHT_DB_PASSWORD:-hindsight_password}
|
||||
command: >
|
||||
bash -c "
|
||||
echo 'PostgreSQL is ready - creating hindsight_db database';
|
||||
psql -h hindsight-db-timescale -p 5432 -U hindsight_user -c 'CREATE DATABASE hindsight_db;' 2>/dev/null || echo 'Database already exists';
|
||||
echo 'Installing Timescale extensions...';
|
||||
echo '1/3: Installing pgvector (required by pgvectorscale)...';
|
||||
psql -h hindsight-db-timescale -p 5432 -U hindsight_user -d hindsight_db -c 'CREATE EXTENSION IF NOT EXISTS vector CASCADE;';
|
||||
echo '2/3: Installing pgvectorscale (DiskANN vector search)...';
|
||||
psql -h hindsight-db-timescale -p 5432 -U hindsight_user -d hindsight_db -c 'CREATE EXTENSION IF NOT EXISTS vectorscale CASCADE;';
|
||||
echo '3/3: Installing pg_textsearch (BM25 text search)...';
|
||||
psql -h hindsight-db-timescale -p 5432 -U hindsight_user -d hindsight_db -c 'CREATE EXTENSION IF NOT EXISTS pg_textsearch CASCADE;';
|
||||
echo '';
|
||||
echo '✅ Timescale extensions installed successfully';
|
||||
echo '';
|
||||
echo 'Installed extensions:';
|
||||
psql -h hindsight-db-timescale -p 5432 -U hindsight_user -d hindsight_db -c \"\\dx\" | grep -E '(vector|vectorscale|pg_textsearch)';
|
||||
"
|
||||
restart: "no"
|
||||
networks:
|
||||
- hindsight-net
|
||||
|
||||
hindsight:
|
||||
image: ghcr.io/vectorize-io/hindsight:${HINDSIGHT_VERSION:-latest}
|
||||
container_name: hindsight-app-timescale
|
||||
ports:
|
||||
- "8888:8888"
|
||||
- "9999:9999"
|
||||
environment:
|
||||
# LLM Configuration
|
||||
HINDSIGHT_API_LLM_PROVIDER: ${HINDSIGHT_API_LLM_PROVIDER:-openai}
|
||||
HINDSIGHT_API_LLM_API_KEY: ${OPENAI_API_KEY:-your-api-key}
|
||||
|
||||
# Database Configuration
|
||||
HINDSIGHT_API_DATABASE_URL: postgresql://${HINDSIGHT_DB_USER:-hindsight_user}:${HINDSIGHT_DB_PASSWORD:-hindsight_password}@db:5432/${HINDSIGHT_DB_NAME:-hindsight_db}
|
||||
|
||||
# Timescale Extensions
|
||||
# pgvectorscale: DiskANN algorithm for disk-based scalable vector search
|
||||
HINDSIGHT_API_VECTOR_EXTENSION: pgvectorscale
|
||||
# pg_textsearch: High-performance BM25 text search
|
||||
HINDSIGHT_API_TEXT_SEARCH_EXTENSION: pg_textsearch
|
||||
|
||||
depends_on:
|
||||
db:
|
||||
condition: service_healthy
|
||||
timescale-init:
|
||||
condition: service_completed_successfully
|
||||
networks:
|
||||
- hindsight-net
|
||||
|
||||
|
||||
networks:
|
||||
hindsight-net:
|
||||
driver: bridge
|
||||
|
||||
volumes:
|
||||
pg_data:
|
||||
@@ -1,93 +0,0 @@
|
||||
name: hindsight
|
||||
# Docker Compose file for Hindsight with PostgreSQL and vectorchord
|
||||
# docker compose -f docker/docker-compose/docker-compose.yaml down && sleep 2 && docker compose -f docker/docker-compose/docker-compose.yaml up -d
|
||||
# Make sure to set the required environment variables before running:
|
||||
# - HINDSIGHT_DB_PASSWORD: Password for the PostgreSQL user
|
||||
# - Configure LLM provider variables as needed (see below in the hindsight service)
|
||||
#
|
||||
# Usage:
|
||||
# docker compose up -d
|
||||
#
|
||||
# Optional environment variables with defaults:
|
||||
# - HINDSIGHT_VERSION: Hindsight application version (default: latest)
|
||||
# - HINDSIGHT_DB_USER: PostgreSQL user (default: hindsight_user)
|
||||
# - HINDSIGHT_DB_NAME: PostgreSQL database name (default: hindsight_db)
|
||||
# - HINDSIGHT_DB_VERSION: PostgreSQL version (default: 18)
|
||||
|
||||
services:
|
||||
db:
|
||||
# Use a PostgreSQL-Image with vectorchord extension pre-installed
|
||||
image: tensorchord/vchord-suite:pg${HINDSIGHT_DB_VERSION:-18-latest}
|
||||
container_name: hindsight-db
|
||||
restart: always
|
||||
# Expose PostgreSQL port
|
||||
ports:
|
||||
- "5436:5432"
|
||||
environment:
|
||||
POSTGRES_USER: ${HINDSIGHT_DB_USER:-hindsight_user}
|
||||
POSTGRES_PASSWORD: ${HINDSIGHT_DB_PASSWORD:-hindsight_password}
|
||||
POSTGRES_DB: ${HINDSIGHT_DB_NAME:-hindsight_db}
|
||||
volumes:
|
||||
- pg_data:/var/lib/postgresql/${HINDSIGHT_DB_VERSION:-18}/docker
|
||||
networks:
|
||||
- hindsight-net
|
||||
|
||||
vectorchord-init:
|
||||
image: tensorchord/vchord-suite:pg18-latest
|
||||
#container_name: vectorchord-init
|
||||
depends_on:
|
||||
- db
|
||||
environment:
|
||||
- PGPASSWORD=${HINDSIGHT_DB_PASSWORD:-hindsight_password}
|
||||
command: >
|
||||
bash -c "
|
||||
echo 'Waiting for PostgreSQL to be ready...';
|
||||
until pg_isready -h hindsight-db -p 5432 -U hindsight_user; do
|
||||
echo 'PostgreSQL is unavailable - sleeping';
|
||||
sleep 2;
|
||||
done;
|
||||
echo 'PostgreSQL is ready - creating hindsight_db database';
|
||||
psql -h hindsight-db -p 5432 -U hindsight_user -c 'CREATE DATABASE hindsight_db;' 2>/dev/null || echo 'Database already exists';
|
||||
echo 'Creating extensions in hindsight_db database';
|
||||
psql -h hindsight-db -p 5432 -U hindsight_user -d hindsight_db -c 'CREATE EXTENSION IF NOT EXISTS vchord CASCADE;';
|
||||
psql -h hindsight-db -p 5432 -U hindsight_user -d hindsight_db -c 'CREATE EXTENSION IF NOT EXISTS pg_tokenizer CASCADE;';
|
||||
psql -h hindsight-db -p 5432 -U hindsight_user -d hindsight_db -c 'CREATE EXTENSION IF NOT EXISTS vchord_bm25 CASCADE;';
|
||||
echo 'Creating llmlingua2 tokenizer';
|
||||
psql -h hindsight-db -p 5432 -U hindsight_user -d hindsight_db -c \"SELECT create_tokenizer('llmlingua2', \\$\\$ model = \\\"llmlingua2\\\" \\$\\$);\" 2>/dev/null || echo 'Tokenizer already exists or creation skipped';
|
||||
echo 'Database and extensions created successfully';
|
||||
"
|
||||
restart: "no"
|
||||
networks:
|
||||
- hindsight-net
|
||||
|
||||
hindsight:
|
||||
image: ghcr.io/vectorize-io/hindsight:${HINDSIGHT_VERSION:-latest}
|
||||
container_name: hindsight-app
|
||||
ports:
|
||||
- "8888:8888"
|
||||
- "9999:9999"
|
||||
environment:
|
||||
# LLM Configuration (uses OpenAI for testing vchord)
|
||||
# LLM configuration
|
||||
HINDSIGHT_API_LLM_PROVIDER: ${HINDSIGHT_API_LLM_PROVIDER:-openai}
|
||||
HINDSIGHT_API_LLM_API_KEY: ${OPENAI_API_KEY:-your-api-key}
|
||||
|
||||
# Database Configuration
|
||||
HINDSIGHT_API_DATABASE_URL: postgresql://${HINDSIGHT_DB_USER:-hindsight_user}:${HINDSIGHT_DB_PASSWORD:-hindsight_password}@db:5432/${HINDSIGHT_DB_NAME:-hindsight_db}
|
||||
|
||||
# Vector and Text Search Extensions
|
||||
HINDSIGHT_API_VECTOR_EXTENSION: vchord
|
||||
HINDSIGHT_API_TEXT_SEARCH_EXTENSION: vchord
|
||||
|
||||
depends_on:
|
||||
- db
|
||||
networks:
|
||||
- hindsight-net
|
||||
|
||||
|
||||
networks:
|
||||
hindsight-net:
|
||||
driver: bridge
|
||||
|
||||
volumes:
|
||||
pg_data:
|
||||
@@ -8,7 +8,6 @@
|
||||
# Set to false when using external providers (TEI, OpenAI, Cohere)
|
||||
# PRELOAD_ML_MODELS=true/false - Pre-download ML models during build (default: true)
|
||||
# Only effective when INCLUDE_LOCAL_MODELS=true
|
||||
# NOTE: tiktoken encodings are ALWAYS preloaded (required for air-gapped deployments)
|
||||
#
|
||||
# Examples:
|
||||
# docker build -t hindsight . # Both (standalone)
|
||||
@@ -42,22 +41,25 @@ 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-slim/pyproject.toml ./api/
|
||||
COPY hindsight-api-slim/README.md ./api/
|
||||
COPY hindsight-api/pyproject.toml ./api/
|
||||
COPY hindsight-api/README.md ./api/
|
||||
|
||||
WORKDIR /app/api
|
||||
|
||||
# Sync dependencies using appropriate extras based on INCLUDE_LOCAL_MODELS
|
||||
# local-ml: torch, sentence-transformers, transformers, einops, flashrank, mlx (optional)
|
||||
# embedded-db: pg0-embedded (always included for embedded PostgreSQL support)
|
||||
RUN if [ "$INCLUDE_LOCAL_MODELS" = "true" ]; then \
|
||||
uv sync --extra local-ml --extra embedded-db; \
|
||||
else \
|
||||
uv sync --extra embedded-db; \
|
||||
# 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; \
|
||||
fi
|
||||
|
||||
# Sync dependencies (will create lock file if needed)
|
||||
RUN uv sync
|
||||
|
||||
# Copy source code (alembic migrations are inside hindsight_api/)
|
||||
COPY hindsight-api-slim/hindsight_api ./hindsight_api
|
||||
COPY hindsight-api/hindsight_api ./hindsight_api
|
||||
|
||||
# Install the local package (uv sync only installed dependencies, not the package itself)
|
||||
RUN uv pip install -e .
|
||||
@@ -109,10 +111,6 @@ RUN rm -f package-lock.json && sed -i '/"@vectorize-io\/hindsight-client":/d' pa
|
||||
# Copy built SDK directly into node_modules (more reliable than npm link in Docker)
|
||||
COPY --from=sdk-builder /app/hindsight-clients/typescript ./node_modules/@vectorize-io/hindsight-client
|
||||
|
||||
# Accept base path as build argument for reverse proxy deployments
|
||||
# Usage: docker build --build-arg NEXT_PUBLIC_BASE_PATH=/hindsight ...
|
||||
ARG NEXT_PUBLIC_BASE_PATH=""
|
||||
|
||||
# Build Control Plane - run next build first, then custom standalone copy
|
||||
# (The build:standalone script expects a specific path structure that differs in Docker)
|
||||
RUN npm exec -- next build
|
||||
@@ -167,35 +165,8 @@ RUN chown -R hindsight:hindsight /app
|
||||
|
||||
USER hindsight
|
||||
|
||||
# Create pg0 data directory as hindsight user so that Docker seeds new named
|
||||
# volumes with correct ownership (UID 1000) on first use, avoiding the
|
||||
# "Permission denied" error when mounting a fresh root-owned volume.
|
||||
RUN mkdir -p /home/hindsight/.pg0
|
||||
|
||||
ENV PATH="/app/api/.venv/bin:${PATH}"
|
||||
|
||||
# Pre-download tiktoken encoding (ALWAYS - required for token counting even in air-gapped envs)
|
||||
# Tiktoken is a core runtime dependency, not an optional ML model
|
||||
RUN MAX_RETRIES=3; \
|
||||
RETRY_DELAY=5; \
|
||||
for i in $(seq 1 $MAX_RETRIES); do \
|
||||
echo "Attempt $i/$MAX_RETRIES: Downloading tiktoken encoding..."; \
|
||||
/app/api/.venv/bin/python -c "\
|
||||
import tiktoken; \
|
||||
print('Downloading cl100k_base encoding...'); \
|
||||
tiktoken.get_encoding('cl100k_base'); \
|
||||
print('Tiktoken encoding cached successfully')" && break; \
|
||||
if [ $i -lt $MAX_RETRIES ]; then \
|
||||
echo "Attempt $i failed, retrying in ${RETRY_DELAY}s..."; \
|
||||
sleep $RETRY_DELAY; \
|
||||
RETRY_DELAY=$((RETRY_DELAY * 2)); \
|
||||
fi; \
|
||||
done; \
|
||||
if [ $i -eq $MAX_RETRIES ]; then \
|
||||
echo "ERROR: Failed to download tiktoken encoding after $MAX_RETRIES attempts"; \
|
||||
exit 1; \
|
||||
fi
|
||||
|
||||
# Pre-download ML models to avoid runtime download (conditional)
|
||||
# Only runs if both PRELOAD_ML_MODELS=true AND INCLUDE_LOCAL_MODELS=true
|
||||
# Includes retry logic with exponential backoff for transient network failures
|
||||
@@ -214,6 +185,7 @@ print('Downloading embedding model...'); \
|
||||
SentenceTransformer('BAAI/bge-small-en-v1.5'); \
|
||||
print('Downloading cross-encoder model...'); \
|
||||
CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2'); \
|
||||
print('Downloading tiktoken encoding...'); import tiktoken; tiktoken.get_encoding('cl100k_base'); \
|
||||
print('Models cached successfully')" && break; \
|
||||
if [ $i -lt $MAX_RETRIES ]; then \
|
||||
echo "Attempt $i failed, retrying in ${RETRY_DELAY}s..."; \
|
||||
@@ -236,10 +208,6 @@ ENV HINDSIGHT_API_LOG_LEVEL=info
|
||||
ENV HINDSIGHT_ENABLE_API=true
|
||||
ENV HINDSIGHT_ENABLE_CP=false
|
||||
ENV PYTHONUNBUFFERED=1
|
||||
# Suppress verbose transformers/HuggingFace model loading warnings
|
||||
ENV TRANSFORMERS_VERBOSITY=error
|
||||
ENV HF_HUB_VERBOSITY=error
|
||||
ENV TOKENIZERS_PARALLELISM=false
|
||||
|
||||
CMD ["/app/start-all.sh"]
|
||||
|
||||
@@ -323,35 +291,8 @@ RUN chown -R hindsight:hindsight /app
|
||||
|
||||
USER hindsight
|
||||
|
||||
# Create pg0 data directory as hindsight user so that Docker seeds new named
|
||||
# volumes with correct ownership (UID 1000) on first use, avoiding the
|
||||
# "Permission denied" error when mounting a fresh root-owned volume.
|
||||
RUN mkdir -p /home/hindsight/.pg0
|
||||
|
||||
ENV PATH="/app/api/.venv/bin:${PATH}"
|
||||
|
||||
# Pre-download tiktoken encoding (ALWAYS - required for token counting even in air-gapped envs)
|
||||
# Tiktoken is a core runtime dependency, not an optional ML model
|
||||
RUN MAX_RETRIES=3; \
|
||||
RETRY_DELAY=5; \
|
||||
for i in $(seq 1 $MAX_RETRIES); do \
|
||||
echo "Attempt $i/$MAX_RETRIES: Downloading tiktoken encoding..."; \
|
||||
/app/api/.venv/bin/python -c "\
|
||||
import tiktoken; \
|
||||
print('Downloading cl100k_base encoding...'); \
|
||||
tiktoken.get_encoding('cl100k_base'); \
|
||||
print('Tiktoken encoding cached successfully')" && break; \
|
||||
if [ $i -lt $MAX_RETRIES ]; then \
|
||||
echo "Attempt $i failed, retrying in ${RETRY_DELAY}s..."; \
|
||||
sleep $RETRY_DELAY; \
|
||||
RETRY_DELAY=$((RETRY_DELAY * 2)); \
|
||||
fi; \
|
||||
done; \
|
||||
if [ $i -eq $MAX_RETRIES ]; then \
|
||||
echo "ERROR: Failed to download tiktoken encoding after $MAX_RETRIES attempts"; \
|
||||
exit 1; \
|
||||
fi
|
||||
|
||||
# Pre-download ML models to avoid runtime download (conditional)
|
||||
# Only runs if both PRELOAD_ML_MODELS=true AND INCLUDE_LOCAL_MODELS=true
|
||||
# Includes retry logic with exponential backoff for transient network failures
|
||||
@@ -370,6 +311,7 @@ print('Downloading embedding model...'); \
|
||||
SentenceTransformer('BAAI/bge-small-en-v1.5'); \
|
||||
print('Downloading cross-encoder model...'); \
|
||||
CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2'); \
|
||||
print('Downloading tiktoken encoding...'); import tiktoken; tiktoken.get_encoding('cl100k_base'); \
|
||||
print('Models cached successfully')" && break; \
|
||||
if [ $i -lt $MAX_RETRIES ]; then \
|
||||
echo "Attempt $i failed, retrying in ${RETRY_DELAY}s..."; \
|
||||
@@ -394,10 +336,6 @@ ENV HINDSIGHT_CP_DATAPLANE_API_URL=http://localhost:8888
|
||||
ENV HINDSIGHT_ENABLE_API=true
|
||||
ENV HINDSIGHT_ENABLE_CP=true
|
||||
ENV PYTHONUNBUFFERED=1
|
||||
# Suppress verbose transformers/HuggingFace model loading warnings
|
||||
ENV TRANSFORMERS_VERBOSITY=error
|
||||
ENV HF_HUB_VERBOSITY=error
|
||||
ENV TOKENIZERS_PARALLELISM=false
|
||||
|
||||
CMD ["/app/start-all.sh"]
|
||||
|
||||
|
||||
@@ -77,32 +77,18 @@ 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
|
||||
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
|
||||
for i in {1..60}; do
|
||||
if curl -sf http://localhost:8888/health &>/dev/null; then
|
||||
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
|
||||
@@ -111,7 +97,7 @@ fi
|
||||
if [ "$ENABLE_CP" = "true" ]; then
|
||||
echo "🎛️ Starting Control Plane..."
|
||||
cd /app/control-plane
|
||||
PORT="${HINDSIGHT_CP_PORT:-9999}" node server.js &
|
||||
PORT=9999 node server.js &
|
||||
CP_PID=$!
|
||||
PIDS+=($CP_PID)
|
||||
else
|
||||
@@ -124,7 +110,7 @@ echo "✅ Hindsight is running!"
|
||||
echo ""
|
||||
echo "📍 Access:"
|
||||
if [ "$ENABLE_CP" = "true" ]; then
|
||||
echo " Control Plane: http://localhost:${HINDSIGHT_CP_PORT:-9999}"
|
||||
echo " Control Plane: http://localhost:9999"
|
||||
fi
|
||||
if [ "$ENABLE_API" = "true" ]; then
|
||||
echo " API: http://localhost:8888"
|
||||
|
||||
@@ -1,235 +0,0 @@
|
||||
#!/bin/bash
|
||||
#
|
||||
# Docker Smoke Test Script
|
||||
#
|
||||
# Tests that a Hindsight Docker image starts correctly and becomes healthy.
|
||||
# Can be run locally or in CI pipelines.
|
||||
#
|
||||
# Usage:
|
||||
# ./docker/test-image.sh <image> [target]
|
||||
#
|
||||
# Arguments:
|
||||
# image - Docker image to test (e.g., hindsight-api:test, ghcr.io/vectorize-io/hindsight:latest)
|
||||
# target - Optional: 'cp-only' for control plane, otherwise assumes API image (default: api)
|
||||
#
|
||||
# Environment variables:
|
||||
# HINDSIGHT_API_LLM_API_KEY - Required for API/standalone images (LLM verification)
|
||||
# HINDSIGHT_API_LLM_PROVIDER - LLM provider (default: openai)
|
||||
# HINDSIGHT_API_LLM_MODEL - LLM model (default: gpt-4o-mini)
|
||||
# HINDSIGHT_API_EMBEDDINGS_PROVIDER - Embeddings provider (optional, for slim images: openai, cohere, tei)
|
||||
# HINDSIGHT_API_EMBEDDINGS_OPENAI_API_KEY - OpenAI API key for embeddings (optional)
|
||||
# HINDSIGHT_API_RERANKER_PROVIDER - Reranker provider (optional, for slim images: cohere, tei)
|
||||
# HINDSIGHT_API_COHERE_API_KEY - Cohere API key for reranking (optional)
|
||||
# SMOKE_TEST_TIMEOUT - Timeout in seconds (default: 120)
|
||||
# SMOKE_TEST_CONTAINER_NAME - Container name (default: hindsight-smoke-test)
|
||||
#
|
||||
# Examples:
|
||||
# # Test a locally built full image
|
||||
# ./docker/test-image.sh hindsight-api:test
|
||||
#
|
||||
# # Test a released image
|
||||
# ./docker/test-image.sh ghcr.io/vectorize-io/hindsight:latest
|
||||
#
|
||||
# # Test control plane image
|
||||
# ./docker/test-image.sh hindsight-control-plane:test cp-only
|
||||
#
|
||||
# # Test slim image with external providers
|
||||
# export HINDSIGHT_API_LLM_API_KEY=sk_xxx
|
||||
# export HINDSIGHT_API_EMBEDDINGS_PROVIDER=openai
|
||||
# export HINDSIGHT_API_EMBEDDINGS_OPENAI_API_KEY=sk-xxx
|
||||
# export HINDSIGHT_API_RERANKER_PROVIDER=cohere
|
||||
# export HINDSIGHT_API_COHERE_API_KEY=xxx
|
||||
# ./docker/test-image.sh hindsight-slim:test
|
||||
#
|
||||
# Exit codes:
|
||||
# 0 - Success (container healthy)
|
||||
# 1 - Failure (container not healthy within timeout)
|
||||
# 2 - Invalid arguments
|
||||
#
|
||||
|
||||
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'
|
||||
YELLOW='\033[0;33m'
|
||||
NC='\033[0m' # No Color
|
||||
|
||||
# Configuration
|
||||
IMAGE="${1:-}"
|
||||
TARGET="${2:-api}"
|
||||
TIMEOUT="${SMOKE_TEST_TIMEOUT:-120}"
|
||||
CONTAINER_NAME="${SMOKE_TEST_CONTAINER_NAME:-hindsight-smoke-test}"
|
||||
LLM_PROVIDER="${HINDSIGHT_API_LLM_PROVIDER:-openai}"
|
||||
LLM_MODEL="${HINDSIGHT_API_LLM_MODEL:-gpt-4o-mini}"
|
||||
|
||||
# Validate arguments
|
||||
if [ -z "$IMAGE" ]; then
|
||||
echo -e "${RED}Error: Image argument is required${NC}"
|
||||
echo ""
|
||||
echo "Usage: $0 <image> [target]"
|
||||
echo ""
|
||||
echo "Examples:"
|
||||
echo " $0 hindsight-api:test"
|
||||
echo " $0 ghcr.io/vectorize-io/hindsight:latest"
|
||||
echo " $0 hindsight-control-plane:test cp-only"
|
||||
exit 2
|
||||
fi
|
||||
|
||||
# Determine health endpoint based on target
|
||||
if [ "$TARGET" = "cp-only" ]; then
|
||||
HEALTH_PORT=9999
|
||||
HEALTH_PATH="/api/health"
|
||||
NEEDS_LLM=false
|
||||
else
|
||||
HEALTH_PORT=8888
|
||||
HEALTH_PATH="/health"
|
||||
NEEDS_LLM=true
|
||||
fi
|
||||
|
||||
# Check for required environment variables
|
||||
if [ "$NEEDS_LLM" = true ] && [ "$LLM_PROVIDER" != "vertexai" ] && [ -z "${HINDSIGHT_API_LLM_API_KEY:-}" ]; then
|
||||
echo -e "${RED}Error: HINDSIGHT_API_LLM_API_KEY environment variable is required for API/standalone images${NC}"
|
||||
echo "Set it with: export HINDSIGHT_API_LLM_API_KEY=your-api-key"
|
||||
exit 2
|
||||
fi
|
||||
|
||||
# Cleanup function
|
||||
cleanup() {
|
||||
echo "Cleaning up..."
|
||||
docker stop "$CONTAINER_NAME" 2>/dev/null || true
|
||||
docker rm "$CONTAINER_NAME" 2>/dev/null || true
|
||||
}
|
||||
|
||||
# Set trap to cleanup on exit
|
||||
trap cleanup EXIT
|
||||
|
||||
echo -e "${YELLOW}Starting smoke test for: ${IMAGE}${NC}"
|
||||
echo " Target: $TARGET"
|
||||
echo " Health endpoint: http://localhost:${HEALTH_PORT}${HEALTH_PATH}"
|
||||
echo " Timeout: ${TIMEOUT}s"
|
||||
echo ""
|
||||
|
||||
# Remove any existing container with the same name
|
||||
docker rm -f "$CONTAINER_NAME" 2>/dev/null || true
|
||||
|
||||
# Start container based on target type
|
||||
echo "Starting container..."
|
||||
if [ "$TARGET" = "cp-only" ]; then
|
||||
docker run -d --name "$CONTAINER_NAME" \
|
||||
-p "${HEALTH_PORT}:${HEALTH_PORT}" \
|
||||
"$IMAGE"
|
||||
else
|
||||
# Build docker run command with required and optional env vars
|
||||
DOCKER_CMD="docker run -d --name $CONTAINER_NAME"
|
||||
DOCKER_CMD="$DOCKER_CMD -e HINDSIGHT_API_LLM_PROVIDER=$LLM_PROVIDER"
|
||||
if [ -n "${HINDSIGHT_API_LLM_API_KEY:-}" ]; then
|
||||
DOCKER_CMD="$DOCKER_CMD -e HINDSIGHT_API_LLM_API_KEY=${HINDSIGHT_API_LLM_API_KEY}"
|
||||
fi
|
||||
DOCKER_CMD="$DOCKER_CMD -e HINDSIGHT_API_LLM_MODEL=$LLM_MODEL"
|
||||
|
||||
# Add Vertex AI config if provider is vertexai
|
||||
if [ "$LLM_PROVIDER" = "vertexai" ]; then
|
||||
if [ -n "${HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY:-}" ]; then
|
||||
DOCKER_CMD="$DOCKER_CMD -v ${HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY}:/tmp/gcp-credentials.json:ro"
|
||||
DOCKER_CMD="$DOCKER_CMD -e HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY=/tmp/gcp-credentials.json"
|
||||
fi
|
||||
if [ -n "${HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID:-}" ]; then
|
||||
DOCKER_CMD="$DOCKER_CMD -e HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID=${HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID}"
|
||||
fi
|
||||
if [ -n "${HINDSIGHT_API_LLM_VERTEXAI_REGION:-}" ]; then
|
||||
DOCKER_CMD="$DOCKER_CMD -e HINDSIGHT_API_LLM_VERTEXAI_REGION=${HINDSIGHT_API_LLM_VERTEXAI_REGION}"
|
||||
fi
|
||||
fi
|
||||
|
||||
# Add optional embeddings provider config
|
||||
if [ -n "${HINDSIGHT_API_EMBEDDINGS_PROVIDER:-}" ]; then
|
||||
DOCKER_CMD="$DOCKER_CMD -e HINDSIGHT_API_EMBEDDINGS_PROVIDER=${HINDSIGHT_API_EMBEDDINGS_PROVIDER}"
|
||||
fi
|
||||
if [ -n "${HINDSIGHT_API_EMBEDDINGS_OPENAI_API_KEY:-}" ]; then
|
||||
DOCKER_CMD="$DOCKER_CMD -e HINDSIGHT_API_EMBEDDINGS_OPENAI_API_KEY=${HINDSIGHT_API_EMBEDDINGS_OPENAI_API_KEY}"
|
||||
fi
|
||||
|
||||
# Add optional reranker provider config
|
||||
if [ -n "${HINDSIGHT_API_RERANKER_PROVIDER:-}" ]; then
|
||||
DOCKER_CMD="$DOCKER_CMD -e HINDSIGHT_API_RERANKER_PROVIDER=${HINDSIGHT_API_RERANKER_PROVIDER}"
|
||||
fi
|
||||
if [ -n "${HINDSIGHT_API_COHERE_API_KEY:-}" ]; then
|
||||
DOCKER_CMD="$DOCKER_CMD -e HINDSIGHT_API_COHERE_API_KEY=${HINDSIGHT_API_COHERE_API_KEY}"
|
||||
fi
|
||||
|
||||
DOCKER_CMD="$DOCKER_CMD -p ${HEALTH_PORT}:${HEALTH_PORT}"
|
||||
DOCKER_CMD="$DOCKER_CMD $IMAGE"
|
||||
|
||||
eval $DOCKER_CMD
|
||||
fi
|
||||
|
||||
# Wait for health endpoint
|
||||
echo "Waiting for health endpoint at http://localhost:${HEALTH_PORT}${HEALTH_PATH}..."
|
||||
start_time=$(date +%s)
|
||||
|
||||
for i in $(seq 1 "$TIMEOUT"); do
|
||||
if curl -sf "http://localhost:${HEALTH_PORT}${HEALTH_PATH}" > /dev/null 2>&1; then
|
||||
end_time=$(date +%s)
|
||||
duration=$((end_time - start_time))
|
||||
echo ""
|
||||
echo -e "${GREEN}Container is healthy after ${duration}s${NC}"
|
||||
echo ""
|
||||
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
|
||||
echo ""
|
||||
echo -e "${GREEN}Smoke test PASSED${NC}"
|
||||
exit 0
|
||||
fi
|
||||
|
||||
# Show progress every 10 seconds
|
||||
if [ $((i % 10)) -eq 0 ]; then
|
||||
echo " Still waiting... (${i}s)"
|
||||
fi
|
||||
|
||||
# Check if container is still running
|
||||
if ! docker ps -q -f "name=$CONTAINER_NAME" | grep -q .; then
|
||||
echo ""
|
||||
echo -e "${RED}Container exited unexpectedly!${NC}"
|
||||
echo ""
|
||||
echo "=== Container Logs ==="
|
||||
docker logs "$CONTAINER_NAME" 2>&1
|
||||
echo ""
|
||||
echo -e "${RED}Smoke test FAILED${NC}"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
sleep 1
|
||||
done
|
||||
|
||||
# Timeout reached
|
||||
echo ""
|
||||
echo -e "${RED}Container failed to become healthy after ${TIMEOUT}s${NC}"
|
||||
echo ""
|
||||
echo "=== Container Logs ==="
|
||||
docker logs "$CONTAINER_NAME" 2>&1
|
||||
echo ""
|
||||
echo -e "${RED}Smoke test FAILED${NC}"
|
||||
exit 1
|
||||
@@ -1,47 +0,0 @@
|
||||
#!/bin/bash
|
||||
#
|
||||
# Local Test Script for Slim Docker Images
|
||||
#
|
||||
# This script makes it easy to test slim images locally with external providers.
|
||||
# It expects API keys to be set in environment variables.
|
||||
#
|
||||
# Usage:
|
||||
# export OPENAI_API_KEY=sk-xxx
|
||||
# export COHERE_API_KEY=xxx
|
||||
# ./docker/test-slim-local.sh
|
||||
#
|
||||
# Or inline:
|
||||
# OPENAI_API_KEY=sk_xxx COHERE_API_KEY=xxx ./docker/test-slim-local.sh
|
||||
#
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
# Check for required API keys
|
||||
if [ -z "${OPENAI_API_KEY:-}" ]; then
|
||||
echo "❌ Error: OPENAI_API_KEY environment variable is required"
|
||||
echo "Set it with: export OPENAI_API_KEY=sk-xxx"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
if [ -z "${COHERE_API_KEY:-}" ]; then
|
||||
echo "❌ Error: COHERE_API_KEY environment variable is required"
|
||||
echo "Set it with: export COHERE_API_KEY=xxx"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# Configuration
|
||||
IMAGE="${1:-hindsight-slim:test}"
|
||||
echo "Testing image: $IMAGE"
|
||||
echo ""
|
||||
|
||||
# Set up LLM and external providers
|
||||
export HINDSIGHT_API_LLM_PROVIDER=openai
|
||||
export HINDSIGHT_API_LLM_API_KEY=$OPENAI_API_KEY
|
||||
export HINDSIGHT_API_LLM_MODEL=gpt-4o-mini
|
||||
export HINDSIGHT_API_EMBEDDINGS_PROVIDER=openai
|
||||
export HINDSIGHT_API_EMBEDDINGS_OPENAI_API_KEY=$OPENAI_API_KEY
|
||||
export HINDSIGHT_API_RERANKER_PROVIDER=cohere
|
||||
export HINDSIGHT_API_COHERE_API_KEY=$COHERE_API_KEY
|
||||
|
||||
# Run the test
|
||||
exec "$(dirname "$0")/test-image.sh" "$IMAGE" standalone
|
||||
@@ -2,8 +2,8 @@ apiVersion: v2
|
||||
name: hindsight
|
||||
description: Hindsight helm chart
|
||||
type: application
|
||||
version: 0.4.17
|
||||
appVersion: "0.4.17"
|
||||
version: 0.4.7
|
||||
appVersion: "0.4.7"
|
||||
keywords:
|
||||
- ai
|
||||
- memory
|
||||
|
||||
@@ -127,38 +127,6 @@ API URL for control plane
|
||||
{{- printf "http://%s-api:%d" (include "hindsight.fullname" .) (.Values.api.service.port | int) }}
|
||||
{{- end }}
|
||||
|
||||
{{/*
|
||||
TEI reranker labels
|
||||
*/}}
|
||||
{{- define "hindsight.tei.reranker.labels" -}}
|
||||
{{ include "hindsight.labels" . }}
|
||||
app.kubernetes.io/component: tei-reranker
|
||||
{{- end }}
|
||||
|
||||
{{/*
|
||||
TEI reranker selector labels
|
||||
*/}}
|
||||
{{- define "hindsight.tei.reranker.selectorLabels" -}}
|
||||
{{ include "hindsight.selectorLabels" . }}
|
||||
app.kubernetes.io/component: tei-reranker
|
||||
{{- end }}
|
||||
|
||||
{{/*
|
||||
TEI embedding labels
|
||||
*/}}
|
||||
{{- define "hindsight.tei.embedding.labels" -}}
|
||||
{{ include "hindsight.labels" . }}
|
||||
app.kubernetes.io/component: tei-embedding
|
||||
{{- end }}
|
||||
|
||||
{{/*
|
||||
TEI embedding selector labels
|
||||
*/}}
|
||||
{{- define "hindsight.tei.embedding.selectorLabels" -}}
|
||||
{{ include "hindsight.selectorLabels" . }}
|
||||
app.kubernetes.io/component: tei-embedding
|
||||
{{- end }}
|
||||
|
||||
{{/*
|
||||
Get the name of the secret to use
|
||||
*/}}
|
||||
|
||||
@@ -33,7 +33,7 @@ spec:
|
||||
- name: api
|
||||
securityContext:
|
||||
{{- toYaml .Values.securityContext | nindent 10 }}
|
||||
image: "{{ .Values.api.image.repository }}:{{ .Values.api.image.tag | default .Values.version | default .Chart.AppVersion }}"
|
||||
image: "{{ .Values.api.image.repository }}:{{ .Values.api.image.tag | default .Values.version }}"
|
||||
imagePullPolicy: {{ .Values.api.image.pullPolicy }}
|
||||
ports:
|
||||
- name: http
|
||||
@@ -60,25 +60,10 @@ spec:
|
||||
- name: HINDSIGHT_API_WORKER_ENABLED
|
||||
value: "false"
|
||||
{{- end }}
|
||||
{{- /* Explicitly set port to override K8s service discovery env var (HINDSIGHT_API_PORT) */}}
|
||||
- name: HINDSIGHT_API_PORT
|
||||
value: {{ .Values.api.service.targetPort | quote }}
|
||||
{{- range $key, $value := .Values.api.env }}
|
||||
- name: {{ $key }}
|
||||
value: {{ $value | quote }}
|
||||
{{- end }}
|
||||
{{- if .Values.tei.reranker.enabled }}
|
||||
- name: HINDSIGHT_API_RERANKER_PROVIDER
|
||||
value: "tei"
|
||||
- name: HINDSIGHT_API_RERANKER_TEI_URL
|
||||
value: "http://{{ include "hindsight.fullname" . }}-tei-reranker:{{ .Values.tei.reranker.port }}"
|
||||
{{- end }}
|
||||
{{- if .Values.tei.embedding.enabled }}
|
||||
- name: HINDSIGHT_API_EMBEDDINGS_PROVIDER
|
||||
value: "tei"
|
||||
- name: HINDSIGHT_API_EMBEDDINGS_TEI_URL
|
||||
value: "http://{{ include "hindsight.fullname" . }}-tei-embedding:{{ .Values.tei.embedding.port }}"
|
||||
{{- end }}
|
||||
{{- /* Only use api.secrets when not using existingSecret (for chart-managed secrets) */}}
|
||||
{{- if not .Values.existingSecret }}
|
||||
{{- range $key, $value := .Values.api.secrets }}
|
||||
@@ -99,7 +84,7 @@ spec:
|
||||
nodeSelector:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
{{- with (.Values.api.affinity | default .Values.affinity) }}
|
||||
{{- with .Values.affinity }}
|
||||
affinity:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
|
||||
@@ -33,7 +33,7 @@ spec:
|
||||
- name: control-plane
|
||||
securityContext:
|
||||
{{- toYaml .Values.securityContext | nindent 10 }}
|
||||
image: "{{ .Values.controlPlane.image.repository }}:{{ .Values.controlPlane.image.tag | default .Values.version | default .Chart.AppVersion }}"
|
||||
image: "{{ .Values.controlPlane.image.repository }}:{{ .Values.controlPlane.image.tag | default .Values.version }}"
|
||||
imagePullPolicy: {{ .Values.controlPlane.image.pullPolicy }}
|
||||
ports:
|
||||
- name: http
|
||||
@@ -71,7 +71,7 @@ spec:
|
||||
nodeSelector:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
{{- with (.Values.controlPlane.affinity | default .Values.affinity) }}
|
||||
{{- with .Values.affinity }}
|
||||
affinity:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
|
||||
@@ -1,56 +0,0 @@
|
||||
{{- if and .Values.api.enabled .Values.api.podDisruptionBudget.enabled }}
|
||||
apiVersion: policy/v1
|
||||
kind: PodDisruptionBudget
|
||||
metadata:
|
||||
name: {{ include "hindsight.fullname" . }}-api
|
||||
labels:
|
||||
{{- include "hindsight.api.labels" . | nindent 4 }}
|
||||
spec:
|
||||
{{- if .Values.api.podDisruptionBudget.minAvailable }}
|
||||
minAvailable: {{ .Values.api.podDisruptionBudget.minAvailable }}
|
||||
{{- end }}
|
||||
{{- if .Values.api.podDisruptionBudget.maxUnavailable }}
|
||||
maxUnavailable: {{ .Values.api.podDisruptionBudget.maxUnavailable }}
|
||||
{{- end }}
|
||||
selector:
|
||||
matchLabels:
|
||||
{{- include "hindsight.api.selectorLabels" . | nindent 6 }}
|
||||
{{- end }}
|
||||
---
|
||||
{{- if and .Values.controlPlane.enabled .Values.controlPlane.podDisruptionBudget.enabled }}
|
||||
apiVersion: policy/v1
|
||||
kind: PodDisruptionBudget
|
||||
metadata:
|
||||
name: {{ include "hindsight.fullname" . }}-control-plane
|
||||
labels:
|
||||
{{- include "hindsight.controlPlane.labels" . | nindent 4 }}
|
||||
spec:
|
||||
{{- if .Values.controlPlane.podDisruptionBudget.minAvailable }}
|
||||
minAvailable: {{ .Values.controlPlane.podDisruptionBudget.minAvailable }}
|
||||
{{- end }}
|
||||
{{- if .Values.controlPlane.podDisruptionBudget.maxUnavailable }}
|
||||
maxUnavailable: {{ .Values.controlPlane.podDisruptionBudget.maxUnavailable }}
|
||||
{{- end }}
|
||||
selector:
|
||||
matchLabels:
|
||||
{{- include "hindsight.controlPlane.selectorLabels" . | nindent 6 }}
|
||||
{{- end }}
|
||||
---
|
||||
{{- if and .Values.worker.enabled .Values.worker.podDisruptionBudget.enabled }}
|
||||
apiVersion: policy/v1
|
||||
kind: PodDisruptionBudget
|
||||
metadata:
|
||||
name: {{ include "hindsight.fullname" . }}-worker
|
||||
labels:
|
||||
{{- include "hindsight.worker.labels" . | nindent 4 }}
|
||||
spec:
|
||||
{{- if .Values.worker.podDisruptionBudget.minAvailable }}
|
||||
minAvailable: {{ .Values.worker.podDisruptionBudget.minAvailable }}
|
||||
{{- end }}
|
||||
{{- if .Values.worker.podDisruptionBudget.maxUnavailable }}
|
||||
maxUnavailable: {{ .Values.worker.podDisruptionBudget.maxUnavailable }}
|
||||
{{- end }}
|
||||
selector:
|
||||
matchLabels:
|
||||
{{- include "hindsight.worker.selectorLabels" . | nindent 6 }}
|
||||
{{- end }}
|
||||
@@ -1,76 +0,0 @@
|
||||
{{- if .Values.tei.embedding.enabled }}
|
||||
apiVersion: apps/v1
|
||||
kind: Deployment
|
||||
metadata:
|
||||
name: {{ include "hindsight.fullname" . }}-tei-embedding
|
||||
labels:
|
||||
{{- include "hindsight.tei.embedding.labels" . | nindent 4 }}
|
||||
spec:
|
||||
replicas: {{ .Values.tei.embedding.replicaCount }}
|
||||
selector:
|
||||
matchLabels:
|
||||
{{- include "hindsight.tei.embedding.selectorLabels" . | nindent 6 }}
|
||||
template:
|
||||
metadata:
|
||||
{{- with .Values.podAnnotations }}
|
||||
annotations:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
labels:
|
||||
{{- include "hindsight.tei.embedding.selectorLabels" . | nindent 8 }}
|
||||
spec:
|
||||
{{- if .Values.serviceAccount.create }}
|
||||
serviceAccountName: {{ include "hindsight.serviceAccountName" . }}
|
||||
{{- end }}
|
||||
securityContext:
|
||||
{{- toYaml .Values.podSecurityContext | nindent 8 }}
|
||||
containers:
|
||||
- name: tei-embedding
|
||||
securityContext:
|
||||
{{- toYaml .Values.securityContext | nindent 10 }}
|
||||
image: "{{ .Values.tei.embedding.image.repository }}:{{ .Values.tei.embedding.image.tag }}"
|
||||
imagePullPolicy: {{ .Values.tei.embedding.image.pullPolicy }}
|
||||
args:
|
||||
- "--model-id"
|
||||
- {{ .Values.tei.embedding.model | quote }}
|
||||
- "--hostname"
|
||||
- "0.0.0.0"
|
||||
{{- range .Values.tei.embedding.args }}
|
||||
- {{ . | quote }}
|
||||
{{- end }}
|
||||
ports:
|
||||
- name: http
|
||||
containerPort: {{ .Values.tei.embedding.port }}
|
||||
protocol: TCP
|
||||
env:
|
||||
- name: PORT
|
||||
value: {{ .Values.tei.embedding.port | quote }}
|
||||
{{- range $key, $value := .Values.tei.embedding.env }}
|
||||
- name: {{ $key }}
|
||||
value: {{ $value | quote }}
|
||||
{{- end }}
|
||||
livenessProbe:
|
||||
{{- toYaml .Values.tei.embedding.livenessProbe | nindent 10 }}
|
||||
readinessProbe:
|
||||
{{- toYaml .Values.tei.embedding.readinessProbe | nindent 10 }}
|
||||
resources:
|
||||
{{- toYaml .Values.tei.embedding.resources | nindent 10 }}
|
||||
volumeMounts:
|
||||
- name: model-cache
|
||||
mountPath: /data
|
||||
volumes:
|
||||
- name: model-cache
|
||||
emptyDir: {}
|
||||
{{- with .Values.nodeSelector }}
|
||||
nodeSelector:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
{{- with .Values.affinity }}
|
||||
affinity:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
{{- with .Values.tolerations }}
|
||||
tolerations:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
{{- end }}
|
||||
@@ -1,17 +0,0 @@
|
||||
{{- if .Values.tei.embedding.enabled }}
|
||||
apiVersion: v1
|
||||
kind: Service
|
||||
metadata:
|
||||
name: {{ include "hindsight.fullname" . }}-tei-embedding
|
||||
labels:
|
||||
{{- include "hindsight.tei.embedding.labels" . | nindent 4 }}
|
||||
spec:
|
||||
type: ClusterIP
|
||||
ports:
|
||||
- port: {{ .Values.tei.embedding.port }}
|
||||
targetPort: http
|
||||
protocol: TCP
|
||||
name: http
|
||||
selector:
|
||||
{{- include "hindsight.tei.embedding.selectorLabels" . | nindent 4 }}
|
||||
{{- end }}
|
||||
@@ -1,76 +0,0 @@
|
||||
{{- if .Values.tei.reranker.enabled }}
|
||||
apiVersion: apps/v1
|
||||
kind: Deployment
|
||||
metadata:
|
||||
name: {{ include "hindsight.fullname" . }}-tei-reranker
|
||||
labels:
|
||||
{{- include "hindsight.tei.reranker.labels" . | nindent 4 }}
|
||||
spec:
|
||||
replicas: {{ .Values.tei.reranker.replicaCount }}
|
||||
selector:
|
||||
matchLabels:
|
||||
{{- include "hindsight.tei.reranker.selectorLabels" . | nindent 6 }}
|
||||
template:
|
||||
metadata:
|
||||
{{- with .Values.podAnnotations }}
|
||||
annotations:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
labels:
|
||||
{{- include "hindsight.tei.reranker.selectorLabels" . | nindent 8 }}
|
||||
spec:
|
||||
{{- if .Values.serviceAccount.create }}
|
||||
serviceAccountName: {{ include "hindsight.serviceAccountName" . }}
|
||||
{{- end }}
|
||||
securityContext:
|
||||
{{- toYaml .Values.podSecurityContext | nindent 8 }}
|
||||
containers:
|
||||
- name: tei-reranker
|
||||
securityContext:
|
||||
{{- toYaml .Values.securityContext | nindent 10 }}
|
||||
image: "{{ .Values.tei.reranker.image.repository }}:{{ .Values.tei.reranker.image.tag }}"
|
||||
imagePullPolicy: {{ .Values.tei.reranker.image.pullPolicy }}
|
||||
args:
|
||||
- "--model-id"
|
||||
- {{ .Values.tei.reranker.model | quote }}
|
||||
- "--hostname"
|
||||
- "0.0.0.0"
|
||||
{{- range .Values.tei.reranker.args }}
|
||||
- {{ . | quote }}
|
||||
{{- end }}
|
||||
ports:
|
||||
- name: http
|
||||
containerPort: {{ .Values.tei.reranker.port }}
|
||||
protocol: TCP
|
||||
env:
|
||||
- name: PORT
|
||||
value: {{ .Values.tei.reranker.port | quote }}
|
||||
{{- range $key, $value := .Values.tei.reranker.env }}
|
||||
- name: {{ $key }}
|
||||
value: {{ $value | quote }}
|
||||
{{- end }}
|
||||
livenessProbe:
|
||||
{{- toYaml .Values.tei.reranker.livenessProbe | nindent 10 }}
|
||||
readinessProbe:
|
||||
{{- toYaml .Values.tei.reranker.readinessProbe | nindent 10 }}
|
||||
resources:
|
||||
{{- toYaml .Values.tei.reranker.resources | nindent 10 }}
|
||||
volumeMounts:
|
||||
- name: model-cache
|
||||
mountPath: /data
|
||||
volumes:
|
||||
- name: model-cache
|
||||
emptyDir: {}
|
||||
{{- with .Values.nodeSelector }}
|
||||
nodeSelector:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
{{- with .Values.affinity }}
|
||||
affinity:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
{{- with .Values.tolerations }}
|
||||
tolerations:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
{{- end }}
|
||||
@@ -1,17 +0,0 @@
|
||||
{{- if .Values.tei.reranker.enabled }}
|
||||
apiVersion: v1
|
||||
kind: Service
|
||||
metadata:
|
||||
name: {{ include "hindsight.fullname" . }}-tei-reranker
|
||||
labels:
|
||||
{{- include "hindsight.tei.reranker.labels" . | nindent 4 }}
|
||||
spec:
|
||||
type: ClusterIP
|
||||
ports:
|
||||
- port: {{ .Values.tei.reranker.port }}
|
||||
targetPort: http
|
||||
protocol: TCP
|
||||
name: http
|
||||
selector:
|
||||
{{- include "hindsight.tei.reranker.selectorLabels" . | nindent 4 }}
|
||||
{{- end }}
|
||||
@@ -32,7 +32,7 @@ spec:
|
||||
- name: worker
|
||||
securityContext:
|
||||
{{- toYaml .Values.securityContext | nindent 10 }}
|
||||
image: "{{ .Values.worker.image.repository }}:{{ .Values.worker.image.tag | default .Values.version | default .Chart.AppVersion }}"
|
||||
image: "{{ .Values.worker.image.repository }}:{{ .Values.worker.image.tag | default .Values.version }}"
|
||||
imagePullPolicy: {{ .Values.worker.image.pullPolicy }}
|
||||
command: ["hindsight-worker"]
|
||||
ports:
|
||||
@@ -99,7 +99,7 @@ spec:
|
||||
nodeSelector:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
{{- with (.Values.worker.affinity | default .Values.affinity) }}
|
||||
{{- with .Values.affinity }}
|
||||
affinity:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
|
||||
+4
-110
@@ -1,8 +1,7 @@
|
||||
# Default values for hindsight
|
||||
|
||||
# Global version override - use this to set a consistent image tag across all components
|
||||
# If not set, defaults to Chart.appVersion from Chart.yaml
|
||||
# version: ""
|
||||
# Chart version - use this to set a consistent image tag across all components
|
||||
version: "0.1.1"
|
||||
|
||||
# Use an existing secret instead of creating one from values
|
||||
# When set, all keys from this secret are injected as environment variables via envFrom
|
||||
@@ -58,15 +57,6 @@ api:
|
||||
timeoutSeconds: 3
|
||||
failureThreshold: 3
|
||||
|
||||
# Pod disruption budget
|
||||
podDisruptionBudget:
|
||||
enabled: false
|
||||
minAvailable: 1
|
||||
# maxUnavailable: 1
|
||||
|
||||
# Pod affinity/anti-affinity (overrides global affinity for this component)
|
||||
# affinity: {}
|
||||
|
||||
# Environment variables
|
||||
env:
|
||||
#HINDSIGHT_API_LLM_PROVIDER: "groq"
|
||||
@@ -85,7 +75,7 @@ worker:
|
||||
image:
|
||||
repository: ghcr.io/vectorize-io/hindsight-api
|
||||
pullPolicy: IfNotPresent
|
||||
# tag: "" # defaults to .Values.version, then Chart.appVersion if not specified
|
||||
# tag defaults to .Values.version if not specified
|
||||
|
||||
service:
|
||||
# Service for metrics scraping (headless for StatefulSet)
|
||||
@@ -131,15 +121,6 @@ worker:
|
||||
# HTTP port for metrics/health (matches service.targetPort)
|
||||
HINDSIGHT_API_WORKER_HTTP_PORT: "8889"
|
||||
|
||||
# Pod disruption budget
|
||||
podDisruptionBudget:
|
||||
enabled: false
|
||||
minAvailable: 1
|
||||
# maxUnavailable: 1
|
||||
|
||||
# Pod affinity/anti-affinity (overrides global affinity for this component)
|
||||
# affinity: {}
|
||||
|
||||
# Secret environment variables (inherited from api.secrets if not specified)
|
||||
secrets: {}
|
||||
|
||||
@@ -183,15 +164,6 @@ controlPlane:
|
||||
timeoutSeconds: 3
|
||||
failureThreshold: 3
|
||||
|
||||
# Pod disruption budget
|
||||
podDisruptionBudget:
|
||||
enabled: false
|
||||
minAvailable: 1
|
||||
# maxUnavailable: 1
|
||||
|
||||
# Pod affinity/anti-affinity (overrides global affinity for this component)
|
||||
# affinity: {}
|
||||
|
||||
# Environment variables
|
||||
env:
|
||||
NODE_ENV: "production"
|
||||
@@ -290,87 +262,9 @@ nodeSelector: {}
|
||||
# Tolerations
|
||||
tolerations: []
|
||||
|
||||
# Affinity (applied to all components unless overridden per-component)
|
||||
# Affinity
|
||||
affinity: {}
|
||||
|
||||
# TEI (Text Embeddings Inference) - optional standalone deployments
|
||||
# for reranking and/or embedding models
|
||||
tei:
|
||||
reranker:
|
||||
enabled: false
|
||||
replicaCount: 1
|
||||
image:
|
||||
repository: ghcr.io/huggingface/text-embeddings-inference
|
||||
tag: cpu-1.8.3
|
||||
pullPolicy: IfNotPresent
|
||||
model: "cross-encoder/ms-marco-MiniLM-L-6-v2"
|
||||
port: 8090
|
||||
args:
|
||||
- "--auto-truncate"
|
||||
env:
|
||||
PAYLOAD_LIMIT: "10000000"
|
||||
MAX_CLIENT_BATCH_SIZE: "256"
|
||||
resources:
|
||||
limits:
|
||||
cpu: 2000m
|
||||
memory: 2Gi
|
||||
requests:
|
||||
cpu: 500m
|
||||
memory: 1Gi
|
||||
livenessProbe:
|
||||
httpGet:
|
||||
path: /health
|
||||
port: 8090
|
||||
initialDelaySeconds: 30
|
||||
periodSeconds: 10
|
||||
timeoutSeconds: 5
|
||||
failureThreshold: 6
|
||||
readinessProbe:
|
||||
httpGet:
|
||||
path: /health
|
||||
port: 8090
|
||||
initialDelaySeconds: 15
|
||||
periodSeconds: 5
|
||||
timeoutSeconds: 3
|
||||
failureThreshold: 3
|
||||
|
||||
embedding:
|
||||
enabled: false
|
||||
replicaCount: 1
|
||||
image:
|
||||
repository: ghcr.io/huggingface/text-embeddings-inference
|
||||
tag: cpu-1.8.3
|
||||
pullPolicy: IfNotPresent
|
||||
model: "sentence-transformers/all-MiniLM-L6-v2"
|
||||
port: 8091
|
||||
args: []
|
||||
env:
|
||||
PAYLOAD_LIMIT: "10000000"
|
||||
MAX_CLIENT_BATCH_SIZE: "256"
|
||||
resources:
|
||||
limits:
|
||||
cpu: 2000m
|
||||
memory: 2Gi
|
||||
requests:
|
||||
cpu: 500m
|
||||
memory: 1Gi
|
||||
livenessProbe:
|
||||
httpGet:
|
||||
path: /health
|
||||
port: 8091
|
||||
initialDelaySeconds: 30
|
||||
periodSeconds: 10
|
||||
timeoutSeconds: 5
|
||||
failureThreshold: 6
|
||||
readinessProbe:
|
||||
httpGet:
|
||||
path: /health
|
||||
port: 8091
|
||||
initialDelaySeconds: 15
|
||||
periodSeconds: 5
|
||||
timeoutSeconds: 3
|
||||
failureThreshold: 3
|
||||
|
||||
# Autoscaling
|
||||
autoscaling:
|
||||
enabled: false
|
||||
|
||||
@@ -1,33 +0,0 @@
|
||||
[build-system]
|
||||
requires = ["setuptools>=61"]
|
||||
build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "hindsight-all-slim"
|
||||
version = "0.4.17"
|
||||
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"
|
||||
@@ -1,48 +0,0 @@
|
||||
# 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
|
||||
```
|
||||
@@ -1,193 +0,0 @@
|
||||
"""
|
||||
API namespace classes for organizing client methods.
|
||||
|
||||
These classes provide organized access to different parts of the Hindsight API
|
||||
while ensuring the daemon is running before each call.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .embedded import HindsightEmbedded
|
||||
|
||||
|
||||
class BanksAPI:
|
||||
"""Namespace for bank-related operations."""
|
||||
|
||||
def __init__(self, embedded: "HindsightEmbedded"):
|
||||
self._embedded = embedded
|
||||
|
||||
def create(
|
||||
self,
|
||||
bank_id: str,
|
||||
name: str | None = None,
|
||||
mission: str | None = None,
|
||||
disposition: dict[str, Any] | None = None,
|
||||
):
|
||||
"""Create a new bank."""
|
||||
self._embedded._ensure_started()
|
||||
return self._embedded._client.create_bank(
|
||||
bank_id=bank_id,
|
||||
name=name,
|
||||
mission=mission,
|
||||
disposition=disposition,
|
||||
)
|
||||
|
||||
def delete(self, bank_id: str):
|
||||
"""Delete a bank."""
|
||||
self._embedded._ensure_started()
|
||||
return self._embedded._client.delete_bank(bank_id=bank_id)
|
||||
|
||||
def set_mission(self, bank_id: str, mission: str):
|
||||
"""Set or update the mission for a bank."""
|
||||
self._embedded._ensure_started()
|
||||
return self._embedded._client.set_mission(bank_id=bank_id, mission=mission)
|
||||
|
||||
def set_disposition(self, bank_id: str, disposition: dict[str, Any]):
|
||||
"""Set or update the disposition for a bank."""
|
||||
self._embedded._ensure_started()
|
||||
return self._embedded._client.set_disposition(bank_id=bank_id, disposition=disposition)
|
||||
|
||||
|
||||
class MentalModelsAPI:
|
||||
"""Namespace for mental model operations."""
|
||||
|
||||
def __init__(self, embedded: "HindsightEmbedded"):
|
||||
self._embedded = embedded
|
||||
|
||||
def create(
|
||||
self,
|
||||
bank_id: str,
|
||||
name: str,
|
||||
content: str,
|
||||
tags: list[str] | None = None,
|
||||
):
|
||||
"""Create a new mental model."""
|
||||
self._embedded._ensure_started()
|
||||
return self._embedded._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):
|
||||
"""List all mental models for a bank."""
|
||||
self._embedded._ensure_started()
|
||||
return self._embedded._client.list_mental_models(bank_id=bank_id, tags=tags)
|
||||
|
||||
def get(self, bank_id: str, mental_model_id: str):
|
||||
"""Get a specific mental model."""
|
||||
self._embedded._ensure_started()
|
||||
return self._embedded._client.get_mental_model(bank_id=bank_id, mental_model_id=mental_model_id)
|
||||
|
||||
def refresh(self, bank_id: str, mental_model_id: str):
|
||||
"""Refresh a mental model."""
|
||||
self._embedded._ensure_started()
|
||||
return self._embedded._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,
|
||||
):
|
||||
"""Update a mental model."""
|
||||
self._embedded._ensure_started()
|
||||
return self._embedded._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):
|
||||
"""Delete a mental model."""
|
||||
self._embedded._ensure_started()
|
||||
return self._embedded._client.delete_mental_model(bank_id=bank_id, mental_model_id=mental_model_id)
|
||||
|
||||
|
||||
class DirectivesAPI:
|
||||
"""Namespace for directive operations."""
|
||||
|
||||
def __init__(self, embedded: "HindsightEmbedded"):
|
||||
self._embedded = embedded
|
||||
|
||||
def create(
|
||||
self,
|
||||
bank_id: str,
|
||||
name: str,
|
||||
content: str,
|
||||
tags: list[str] | None = None,
|
||||
):
|
||||
"""Create a new directive."""
|
||||
self._embedded._ensure_started()
|
||||
return self._embedded._client.create_directive(
|
||||
bank_id=bank_id,
|
||||
name=name,
|
||||
content=content,
|
||||
tags=tags,
|
||||
)
|
||||
|
||||
def list(self, bank_id: str, tags: list[str] | None = None):
|
||||
"""List all directives for a bank."""
|
||||
self._embedded._ensure_started()
|
||||
return self._embedded._client.list_directives(bank_id=bank_id, tags=tags)
|
||||
|
||||
def get(self, bank_id: str, directive_id: str):
|
||||
"""Get a specific directive."""
|
||||
self._embedded._ensure_started()
|
||||
return self._embedded._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,
|
||||
):
|
||||
"""Update a directive."""
|
||||
self._embedded._ensure_started()
|
||||
return self._embedded._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):
|
||||
"""Delete a directive."""
|
||||
self._embedded._ensure_started()
|
||||
return self._embedded._client.delete_directive(bank_id=bank_id, directive_id=directive_id)
|
||||
|
||||
|
||||
class MemoriesAPI:
|
||||
"""Namespace for memory operations."""
|
||||
|
||||
def __init__(self, embedded: "HindsightEmbedded"):
|
||||
self._embedded = embedded
|
||||
|
||||
def list(
|
||||
self,
|
||||
bank_id: str,
|
||||
type: str | None = None,
|
||||
search_query: str | None = None,
|
||||
limit: int = 100,
|
||||
offset: int = 0,
|
||||
):
|
||||
"""List memories in a bank."""
|
||||
self._embedded._ensure_started()
|
||||
return self._embedded._client.list_memories(
|
||||
bank_id=bank_id,
|
||||
type=type,
|
||||
search_query=search_query,
|
||||
limit=limit,
|
||||
offset=offset,
|
||||
)
|
||||
@@ -1,243 +0,0 @@
|
||||
"""
|
||||
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."""
|
||||
|
||||
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,
|
||||
):
|
||||
"""Create a new bank."""
|
||||
return self._client.create_bank(
|
||||
bank_id=bank_id,
|
||||
name=name,
|
||||
mission=mission,
|
||||
disposition=disposition,
|
||||
)
|
||||
|
||||
def delete(self, bank_id: str):
|
||||
"""Delete a bank."""
|
||||
return self._client.delete_bank(bank_id=bank_id)
|
||||
|
||||
def set_mission(self, bank_id: str, mission: str):
|
||||
"""Set or update the mission for a bank."""
|
||||
return self._client.set_mission(bank_id=bank_id, mission=mission)
|
||||
|
||||
def set_disposition(self, bank_id: str, disposition: dict[str, Any]):
|
||||
"""Set or update the disposition for a bank."""
|
||||
return self._client.set_disposition(bank_id=bank_id, disposition=disposition)
|
||||
|
||||
def list(self):
|
||||
"""List all banks."""
|
||||
from hindsight_client.hindsight_client import _run_async
|
||||
|
||||
return _run_async(self._client._banks_api.list_banks())
|
||||
|
||||
|
||||
class MentalModelsAPI:
|
||||
"""Namespace for mental model operations."""
|
||||
|
||||
def __init__(self, client: Hindsight):
|
||||
self._client = client
|
||||
|
||||
def create(
|
||||
self,
|
||||
bank_id: str,
|
||||
name: str,
|
||||
content: str,
|
||||
tags: list[str] | None = None,
|
||||
):
|
||||
"""Create a new mental model."""
|
||||
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):
|
||||
"""List all mental models for a bank."""
|
||||
return self._client.list_mental_models(bank_id=bank_id, tags=tags)
|
||||
|
||||
def get(self, bank_id: str, mental_model_id: str):
|
||||
"""Get a specific mental model."""
|
||||
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):
|
||||
"""Refresh a mental model."""
|
||||
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,
|
||||
):
|
||||
"""Update a mental model."""
|
||||
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):
|
||||
"""Delete a mental model."""
|
||||
return self._client.delete_mental_model(bank_id=bank_id, mental_model_id=mental_model_id)
|
||||
|
||||
|
||||
class DirectivesAPI:
|
||||
"""Namespace for directive operations."""
|
||||
|
||||
def __init__(self, client: Hindsight):
|
||||
self._client = client
|
||||
|
||||
def create(
|
||||
self,
|
||||
bank_id: str,
|
||||
name: str,
|
||||
content: str,
|
||||
tags: list[str] | None = None,
|
||||
):
|
||||
"""Create a new directive."""
|
||||
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):
|
||||
"""List all directives for a bank."""
|
||||
return self._client.list_directives(bank_id=bank_id, tags=tags)
|
||||
|
||||
def get(self, bank_id: str, directive_id: str):
|
||||
"""Get a specific directive."""
|
||||
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,
|
||||
):
|
||||
"""Update a directive."""
|
||||
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):
|
||||
"""Delete a directive."""
|
||||
return self._client.delete_directive(bank_id=bank_id, directive_id=directive_id)
|
||||
|
||||
|
||||
class MemoriesAPI:
|
||||
"""Namespace for memory operations."""
|
||||
|
||||
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,
|
||||
):
|
||||
"""List memories in a bank."""
|
||||
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")
|
||||
```
|
||||
"""
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
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."""
|
||||
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."""
|
||||
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."""
|
||||
if self._directives_namespace is None:
|
||||
self._directives_namespace = DirectivesAPI(self)
|
||||
return self._directives_namespace
|
||||
|
||||
@property
|
||||
def memories(self) -> MemoriesAPI:
|
||||
"""Access memory listing operations."""
|
||||
if self._memories_namespace is None:
|
||||
self._memories_namespace = MemoriesAPI(self)
|
||||
return self._memories_namespace
|
||||
@@ -1,377 +0,0 @@
|
||||
"""
|
||||
Embedded Hindsight client with automatic daemon lifecycle management.
|
||||
|
||||
This module provides HindsightEmbedded, a client that uses the same daemon
|
||||
management interface as hindsight-embed CLI, ensuring full compatibility.
|
||||
|
||||
Example:
|
||||
```python
|
||||
from hindsight import HindsightEmbedded
|
||||
|
||||
# Daemon starts automatically on first use
|
||||
client = HindsightEmbedded(
|
||||
profile="myapp",
|
||||
llm_provider="groq",
|
||||
llm_api_key="your-api-key",
|
||||
)
|
||||
|
||||
# Use just like HindsightClient
|
||||
client.retain(bank_id="alice", content="Alice loves AI")
|
||||
results = client.recall(bank_id="alice", query="What does Alice like?")
|
||||
|
||||
# Optional cleanup
|
||||
client.close()
|
||||
```
|
||||
|
||||
Using context manager:
|
||||
```python
|
||||
from hindsight import HindsightEmbedded
|
||||
|
||||
with HindsightEmbedded(profile="myapp") as client:
|
||||
client.retain(bank_id="alice", content="Alice loves AI")
|
||||
# Daemon managed automatically
|
||||
```
|
||||
"""
|
||||
|
||||
import logging
|
||||
import os
|
||||
import threading
|
||||
from typing import Optional
|
||||
|
||||
from hindsight_client import Hindsight
|
||||
from hindsight_embed import get_embed_manager
|
||||
|
||||
from .api_namespaces import BanksAPI, DirectivesAPI, MemoriesAPI, MentalModelsAPI
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class HindsightEmbedded:
|
||||
"""
|
||||
Hindsight client with automatic daemon lifecycle management.
|
||||
|
||||
This client uses the same daemon management interface as hindsight-embed CLI,
|
||||
ensuring full compatibility and shared profiles. The daemon is started automatically
|
||||
on first use and manages profile-specific databases.
|
||||
|
||||
Profile data is stored in: ~/.pg0/instances/hindsight-embed-{profile}/
|
||||
|
||||
All methods from HindsightClient are available:
|
||||
- retain(), retain_batch()
|
||||
- recall()
|
||||
- reflect()
|
||||
- create_bank(), set_mission(), delete_bank()
|
||||
- create_mental_model(), list_mental_models(), etc.
|
||||
- create_directive(), list_directives(), etc.
|
||||
- And all async variants (aretain, arecall, areflect, etc.)
|
||||
|
||||
Args:
|
||||
profile: Profile name for data isolation (default: "default")
|
||||
llm_provider: LLM provider ("groq", "openai", "ollama", "gemini", "anthropic", "lmstudio")
|
||||
llm_api_key: API key for the LLM provider
|
||||
llm_model: Model name to use
|
||||
llm_base_url: Optional custom base URL for LLM API
|
||||
database_url: Optional database URL override (default: profile-specific pg0)
|
||||
idle_timeout: Seconds before daemon auto-exits when idle (default: 300)
|
||||
log_level: Daemon log level (default: "info")
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
profile: str = "default",
|
||||
llm_provider: str = "groq",
|
||||
llm_api_key: str = "",
|
||||
llm_model: str = "openai/gpt-oss-120b",
|
||||
llm_base_url: Optional[str] = None,
|
||||
database_url: Optional[str] = None,
|
||||
idle_timeout: int = 300,
|
||||
log_level: str = "info",
|
||||
):
|
||||
"""
|
||||
Initialize the embedded client (daemon starts on first use).
|
||||
|
||||
Args:
|
||||
profile: Profile name for data isolation
|
||||
llm_provider: LLM provider
|
||||
llm_api_key: API key for the LLM provider
|
||||
llm_model: Model name to use
|
||||
llm_base_url: Optional custom base URL for LLM API
|
||||
database_url: Optional database URL override
|
||||
idle_timeout: Seconds before daemon auto-exits when idle
|
||||
log_level: Daemon log level
|
||||
"""
|
||||
self.profile = profile
|
||||
|
||||
# Build config dict for daemon (matches CLI format)
|
||||
self.config = {
|
||||
"HINDSIGHT_API_LLM_PROVIDER": llm_provider,
|
||||
"HINDSIGHT_API_LLM_API_KEY": llm_api_key,
|
||||
"HINDSIGHT_API_LLM_MODEL": llm_model,
|
||||
"HINDSIGHT_API_LOG_LEVEL": log_level,
|
||||
"HINDSIGHT_EMBED_DAEMON_IDLE_TIMEOUT": str(idle_timeout),
|
||||
}
|
||||
|
||||
if llm_base_url:
|
||||
self.config["HINDSIGHT_API_LLM_BASE_URL"] = llm_base_url
|
||||
|
||||
if database_url:
|
||||
self.config["HINDSIGHT_EMBED_API_DATABASE_URL"] = database_url
|
||||
|
||||
self._client: Optional[Hindsight] = None
|
||||
self._lock = threading.Lock()
|
||||
self._started = False
|
||||
self._closed = False
|
||||
self._manager = get_embed_manager()
|
||||
|
||||
# API namespaces (initialized once, lazily)
|
||||
self._banks_api: Optional[BanksAPI] = None
|
||||
self._mental_models_api: Optional[MentalModelsAPI] = None
|
||||
self._directives_api: Optional[DirectivesAPI] = None
|
||||
self._memories_api: Optional[MemoriesAPI] = None
|
||||
|
||||
def _ensure_started(self):
|
||||
"""Ensure daemon is running (thread-safe)."""
|
||||
if self._started and self._client is not None:
|
||||
return
|
||||
|
||||
with self._lock:
|
||||
# Double-check after acquiring lock
|
||||
if self._started and self._client is not None:
|
||||
return
|
||||
|
||||
if self._closed:
|
||||
raise RuntimeError("Cannot use HindsightEmbedded after it has been closed")
|
||||
|
||||
# Use embed manager interface for daemon management
|
||||
logger.info(f"Ensuring daemon is running for profile '{self.profile}'...")
|
||||
success = self._manager.ensure_running(self.config, self.profile)
|
||||
if not success:
|
||||
raise RuntimeError(f"Failed to start daemon for profile '{self.profile}'")
|
||||
|
||||
# Get daemon URL and create client
|
||||
daemon_url = self._manager.get_url(self.profile)
|
||||
self._client = Hindsight(base_url=daemon_url)
|
||||
self._started = True
|
||||
logger.info(f"Connected to daemon at {daemon_url}")
|
||||
|
||||
def _cleanup(self, stop_daemon_on_close: bool = False):
|
||||
"""
|
||||
Cleanup client resources (idempotent).
|
||||
|
||||
Args:
|
||||
stop_daemon_on_close: If True, stops the daemon. Otherwise, daemon continues
|
||||
running (it will auto-stop after idle timeout).
|
||||
"""
|
||||
if self._closed:
|
||||
return
|
||||
|
||||
with self._lock:
|
||||
if self._closed:
|
||||
return
|
||||
|
||||
if self._client is not None:
|
||||
self._client.close()
|
||||
self._client = None
|
||||
|
||||
# Optionally stop daemon (daemon has idle timeout, so not required)
|
||||
if stop_daemon_on_close and self._started:
|
||||
logger.info(f"Stopping daemon for profile '{self.profile}'...")
|
||||
self._manager.stop(self.profile)
|
||||
|
||||
self._closed = True
|
||||
|
||||
def close(self, stop_daemon: bool = False):
|
||||
"""
|
||||
Explicitly close the client.
|
||||
|
||||
Args:
|
||||
stop_daemon: If True, stops the daemon. Otherwise, daemon continues running
|
||||
and will auto-stop after idle timeout (default: False).
|
||||
|
||||
Note:
|
||||
The daemon may be shared with other clients or the CLI, so stopping it
|
||||
might affect other users. By default, we rely on the daemon's idle timeout.
|
||||
"""
|
||||
self._cleanup(stop_daemon_on_close=stop_daemon)
|
||||
|
||||
def __getattr__(self, name: str):
|
||||
"""
|
||||
Proxy all method calls to the underlying Hindsight client.
|
||||
|
||||
This allows HindsightEmbedded to expose all HindsightClient methods
|
||||
without manually wrapping each one.
|
||||
"""
|
||||
# Ensure server is started before proxying
|
||||
self._ensure_started()
|
||||
|
||||
# Get the attribute from the underlying client
|
||||
attr = getattr(self._client, name)
|
||||
|
||||
# If it's a callable, wrap it to ensure server is started
|
||||
# (shouldn't be needed since _ensure_started already called, but defensive)
|
||||
if callable(attr):
|
||||
|
||||
def wrapper(*args, **kwargs):
|
||||
self._ensure_started()
|
||||
return attr(*args, **kwargs)
|
||||
|
||||
return wrapper
|
||||
|
||||
return attr
|
||||
|
||||
def __enter__(self):
|
||||
"""Context manager entry - ensures server is started."""
|
||||
self._ensure_started()
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||
"""Context manager exit - stops the server."""
|
||||
self.close()
|
||||
|
||||
def __del__(self):
|
||||
"""Cleanup on garbage collection."""
|
||||
self._cleanup()
|
||||
|
||||
@property
|
||||
def banks(self) -> BanksAPI:
|
||||
"""
|
||||
Access bank management operations.
|
||||
|
||||
Each method call ensures the daemon is running before executing.
|
||||
|
||||
Example:
|
||||
```python
|
||||
from hindsight import HindsightEmbedded
|
||||
|
||||
embedded = HindsightEmbedded(profile="myapp", ...)
|
||||
|
||||
# Create a bank
|
||||
embedded.banks.create(bank_id="test", name="Test Bank")
|
||||
|
||||
# Set mission
|
||||
embedded.banks.set_mission(bank_id="test", mission="Help users")
|
||||
```
|
||||
"""
|
||||
if self._banks_api is None:
|
||||
self._banks_api = BanksAPI(self)
|
||||
return self._banks_api
|
||||
|
||||
@property
|
||||
def mental_models(self) -> MentalModelsAPI:
|
||||
"""
|
||||
Access mental model operations.
|
||||
|
||||
Each method call ensures the daemon is running before executing.
|
||||
|
||||
Example:
|
||||
```python
|
||||
from hindsight import HindsightEmbedded
|
||||
|
||||
embedded = HindsightEmbedded(profile="myapp", ...)
|
||||
|
||||
# Create a mental model
|
||||
embedded.mental_models.create(
|
||||
bank_id="test",
|
||||
name="User Preferences",
|
||||
content="User prefers dark mode"
|
||||
)
|
||||
|
||||
# List mental models
|
||||
models = embedded.mental_models.list(bank_id="test")
|
||||
```
|
||||
"""
|
||||
if self._mental_models_api is None:
|
||||
self._mental_models_api = MentalModelsAPI(self)
|
||||
return self._mental_models_api
|
||||
|
||||
@property
|
||||
def directives(self) -> DirectivesAPI:
|
||||
"""
|
||||
Access directive operations.
|
||||
|
||||
Each method call ensures the daemon is running before executing.
|
||||
|
||||
Example:
|
||||
```python
|
||||
from hindsight import HindsightEmbedded
|
||||
|
||||
embedded = HindsightEmbedded(profile="myapp", ...)
|
||||
|
||||
# Create a directive
|
||||
embedded.directives.create(
|
||||
bank_id="test",
|
||||
name="Response Style",
|
||||
content="Always be concise and friendly"
|
||||
)
|
||||
|
||||
# List directives
|
||||
directives = embedded.directives.list(bank_id="test")
|
||||
```
|
||||
"""
|
||||
if self._directives_api is None:
|
||||
self._directives_api = DirectivesAPI(self)
|
||||
return self._directives_api
|
||||
|
||||
@property
|
||||
def memories(self) -> MemoriesAPI:
|
||||
"""
|
||||
Access memory listing operations.
|
||||
|
||||
Each method call ensures the daemon is running before executing.
|
||||
|
||||
Example:
|
||||
```python
|
||||
from hindsight import HindsightEmbedded
|
||||
|
||||
embedded = HindsightEmbedded(profile="myapp", ...)
|
||||
|
||||
# List memories
|
||||
memories = embedded.memories.list(
|
||||
bank_id="test",
|
||||
type="world",
|
||||
limit=50
|
||||
)
|
||||
```
|
||||
"""
|
||||
if self._memories_api is None:
|
||||
self._memories_api = MemoriesAPI(self)
|
||||
return self._memories_api
|
||||
|
||||
@property
|
||||
def client(self) -> Hindsight:
|
||||
"""
|
||||
Get the underlying Hindsight client for direct access.
|
||||
|
||||
WARNING: Using this property directly means daemon restarts won't be
|
||||
handled automatically. Prefer using the API namespaces (banks, mental_models,
|
||||
directives, memories) or direct method calls on HindsightEmbedded instead.
|
||||
|
||||
Ensures daemon is started before returning the client.
|
||||
|
||||
Returns:
|
||||
Hindsight: The underlying client instance
|
||||
|
||||
Example:
|
||||
```python
|
||||
from hindsight import HindsightEmbedded
|
||||
|
||||
embedded = HindsightEmbedded(profile="myapp", ...)
|
||||
|
||||
# Direct access (not recommended - daemon crashes won't be handled)
|
||||
client = embedded.client
|
||||
banks = client.list_banks() # If daemon crashes, this will fail
|
||||
```
|
||||
"""
|
||||
self._ensure_started()
|
||||
return self._client
|
||||
|
||||
@property
|
||||
def url(self) -> str:
|
||||
"""Get the daemon URL (starts daemon if needed)."""
|
||||
self._ensure_started()
|
||||
return self._manager.get_url(self.profile)
|
||||
|
||||
@property
|
||||
def is_running(self) -> bool:
|
||||
"""Check if the client is initialized."""
|
||||
return self._started and not self._closed and self._client is not None
|
||||
@@ -1,338 +0,0 @@
|
||||
"""
|
||||
Integration tests for HindsightEmbedded client.
|
||||
|
||||
Tests the embedded client with automatic server lifecycle management:
|
||||
1. Lazy server startup on first use
|
||||
2. Server reuse across multiple operations
|
||||
3. Context manager support
|
||||
4. Method proxying to underlying HindsightClient
|
||||
5. Proper cleanup
|
||||
|
||||
Note: Each test uses random bank_ids to avoid conflicts and allow safe parallel execution.
|
||||
"""
|
||||
|
||||
import os
|
||||
import uuid
|
||||
|
||||
import pytest
|
||||
|
||||
from hindsight import HindsightEmbedded
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def llm_config():
|
||||
"""Get LLM configuration from environment (session-scoped)."""
|
||||
# Try both naming conventions
|
||||
provider = os.getenv("HINDSIGHT_API_LLM_PROVIDER") or os.getenv("HINDSIGHT_LLM_PROVIDER", "groq")
|
||||
api_key = os.getenv("HINDSIGHT_API_LLM_API_KEY") or os.getenv("HINDSIGHT_LLM_API_KEY", "")
|
||||
model = os.getenv("HINDSIGHT_API_LLM_MODEL") or os.getenv("HINDSIGHT_LLM_MODEL", "openai/gpt-oss-120b")
|
||||
|
||||
if not api_key:
|
||||
pytest.skip("LLM API key not configured. Set HINDSIGHT_API_LLM_API_KEY or HINDSIGHT_LLM_API_KEY.")
|
||||
|
||||
return {
|
||||
"llm_provider": provider,
|
||||
"llm_api_key": api_key,
|
||||
"llm_model": model,
|
||||
}
|
||||
|
||||
|
||||
def test_embedded_lazy_start(llm_config):
|
||||
"""
|
||||
Test that HindsightEmbedded starts server lazily on first use.
|
||||
"""
|
||||
profile = f"test_lazy_{uuid.uuid4().hex[:8]}"
|
||||
bank_id = f"bank_{uuid.uuid4().hex[:8]}"
|
||||
|
||||
# Create client - should NOT start server yet
|
||||
client = HindsightEmbedded(profile=profile, log_level="info", **llm_config)
|
||||
assert not client.is_running, "Server should not be running after initialization"
|
||||
|
||||
# First call should start server
|
||||
result = client.retain(bank_id=bank_id, content="Test content for lazy start")
|
||||
|
||||
# Verify server is now running
|
||||
assert client.is_running, "Server should be running after first call"
|
||||
assert result.success, "Retain should succeed"
|
||||
assert result.items_count >= 1, "Should have stored at least 1 item"
|
||||
|
||||
# Cleanup
|
||||
client.close()
|
||||
assert not client.is_running, "Server should stop after close()"
|
||||
|
||||
|
||||
def test_embedded_context_manager(llm_config):
|
||||
"""
|
||||
Test HindsightEmbedded with context manager.
|
||||
"""
|
||||
profile = f"test_ctx_{uuid.uuid4().hex[:8]}"
|
||||
bank_id = f"bank_{uuid.uuid4().hex[:8]}"
|
||||
|
||||
# Use context manager
|
||||
with HindsightEmbedded(profile=profile, log_level="info", **llm_config) as client:
|
||||
assert client.is_running, "Server should be running inside context"
|
||||
|
||||
# Store memory
|
||||
result = client.retain(bank_id=bank_id, content="Testing context manager")
|
||||
assert result.success, "Retain should succeed"
|
||||
|
||||
# Recall memory
|
||||
recall_results = client.recall(bank_id=bank_id, query="context")
|
||||
assert isinstance(recall_results.results, list), "Recall should return results list"
|
||||
|
||||
# Server should be stopped after context exit
|
||||
# Note: We can't check client.is_running here as client is out of scope
|
||||
|
||||
|
||||
def test_embedded_complete_workflow(llm_config):
|
||||
"""
|
||||
Test complete workflow with HindsightEmbedded.
|
||||
|
||||
This test:
|
||||
1. Creates a client with lazy start
|
||||
2. Creates a memory bank
|
||||
3. Stores multiple memories
|
||||
4. Recalls memories
|
||||
5. Reflects on memories
|
||||
6. Tests cleanup
|
||||
"""
|
||||
profile = f"test_workflow_{uuid.uuid4().hex[:8]}"
|
||||
bank_id = f"assistant_{uuid.uuid4().hex[:8]}"
|
||||
|
||||
client = HindsightEmbedded(profile=profile, log_level="info", **llm_config)
|
||||
|
||||
try:
|
||||
# Step 1: Create a memory bank
|
||||
print(f"\n1. Creating memory bank: {bank_id}")
|
||||
bank_response = client.create_bank(
|
||||
bank_id=bank_id, name="Test Assistant", mission="Help with programming tasks"
|
||||
)
|
||||
assert bank_response.bank_id == bank_id
|
||||
|
||||
# Step 2: Store memories (single)
|
||||
print("\n2. Storing single memory...")
|
||||
retain_response = client.retain(
|
||||
bank_id=bank_id,
|
||||
content="User prefers Python for data analysis.",
|
||||
context="Programming preferences",
|
||||
)
|
||||
assert retain_response.success
|
||||
assert retain_response.items_count >= 1
|
||||
|
||||
# Step 3: Store batch memories
|
||||
print("\n3. Storing batch memories...")
|
||||
batch_response = client.retain_batch(
|
||||
bank_id=bank_id,
|
||||
items=[
|
||||
{"content": "User works with pandas and numpy."},
|
||||
{"content": "User likes matplotlib for visualization."},
|
||||
{"content": "User is interested in machine learning with scikit-learn."},
|
||||
],
|
||||
)
|
||||
assert batch_response.success
|
||||
assert batch_response.items_count >= 3
|
||||
|
||||
# Step 4: Recall memories
|
||||
print("\n4. Recalling memories...")
|
||||
recall_response = client.recall(bank_id=bank_id, query="What tools does the user prefer?", max_tokens=2000)
|
||||
assert isinstance(recall_response.results, list)
|
||||
assert len(recall_response.results) > 0
|
||||
print(f" Found {len(recall_response.results)} relevant memories")
|
||||
|
||||
# Step 5: Reflect on memories
|
||||
print("\n5. Reflecting on memories...")
|
||||
reflect_response = client.reflect(
|
||||
bank_id=bank_id,
|
||||
query="What programming tools should I recommend?",
|
||||
budget="low",
|
||||
)
|
||||
assert reflect_response.text
|
||||
assert len(reflect_response.text) > 0
|
||||
print(f" Answer: {reflect_response.text[:150]}...")
|
||||
|
||||
# Verify answer mentions relevant tools
|
||||
answer_lower = reflect_response.text.lower()
|
||||
assert any(term in answer_lower for term in ["python", "pandas", "numpy", "data"])
|
||||
|
||||
# Step 6: List memories
|
||||
print("\n6. Listing memories...")
|
||||
list_response = client.list_memories(bank_id=bank_id, limit=10)
|
||||
assert len(list_response.items) > 0
|
||||
print(f" Listed {len(list_response.items)} memories")
|
||||
|
||||
finally:
|
||||
# Cleanup
|
||||
client.close()
|
||||
|
||||
|
||||
def test_embedded_server_reuse(llm_config):
|
||||
"""
|
||||
Test that the same server is reused across multiple calls.
|
||||
"""
|
||||
profile = f"test_reuse_{uuid.uuid4().hex[:8]}"
|
||||
bank_id = f"bank_{uuid.uuid4().hex[:8]}"
|
||||
|
||||
client = HindsightEmbedded(profile=profile, log_level="info", **llm_config)
|
||||
|
||||
try:
|
||||
# First call starts server
|
||||
result1 = client.retain(bank_id=bank_id, content="First message")
|
||||
url1 = client.url
|
||||
assert client.is_running
|
||||
|
||||
# Second call should reuse the same server
|
||||
result2 = client.retain(bank_id=bank_id, content="Second message")
|
||||
url2 = client.url
|
||||
|
||||
# URLs should be identical (same server)
|
||||
assert url1 == url2, "Server URL should remain the same across calls"
|
||||
assert result1.success and result2.success
|
||||
|
||||
# Third call should also reuse
|
||||
recall_result = client.recall(bank_id=bank_id, query="message")
|
||||
url3 = client.url
|
||||
assert url3 == url1, "Server URL should remain the same for recall"
|
||||
assert isinstance(recall_result.results, list)
|
||||
|
||||
finally:
|
||||
client.close()
|
||||
|
||||
|
||||
def test_embedded_method_proxying(llm_config):
|
||||
"""
|
||||
Test that all HindsightClient methods are properly proxied.
|
||||
|
||||
This ensures __getattr__ proxying works for various method types.
|
||||
"""
|
||||
profile = f"test_proxy_{uuid.uuid4().hex[:8]}"
|
||||
bank_id = f"bank_{uuid.uuid4().hex[:8]}"
|
||||
|
||||
client = HindsightEmbedded(profile=profile, log_level="info", **llm_config)
|
||||
|
||||
try:
|
||||
# Test bank operations
|
||||
bank = client.create_bank(bank_id=bank_id, name="Proxy Test")
|
||||
assert bank.bank_id == bank_id
|
||||
|
||||
# Test mission setting
|
||||
mission_response = client.set_mission(bank_id=bank_id, mission="Test mission for proxying")
|
||||
assert mission_response.bank_id == bank_id
|
||||
|
||||
# Test retain
|
||||
retain_result = client.retain(bank_id=bank_id, content="Test content")
|
||||
assert retain_result.success
|
||||
|
||||
# Test retain_batch
|
||||
batch_result = client.retain_batch(
|
||||
bank_id=bank_id, items=[{"content": "Item 1"}, {"content": "Item 2"}]
|
||||
)
|
||||
assert batch_result.success
|
||||
assert batch_result.items_count >= 2
|
||||
|
||||
# Test recall
|
||||
recall_result = client.recall(bank_id=bank_id, query="test")
|
||||
assert hasattr(recall_result, "results")
|
||||
|
||||
# Test reflect
|
||||
reflect_result = client.reflect(bank_id=bank_id, query="What is stored?")
|
||||
assert hasattr(reflect_result, "text")
|
||||
|
||||
# Test list_memories
|
||||
list_result = client.list_memories(bank_id=bank_id, limit=5)
|
||||
assert hasattr(list_result, "items")
|
||||
|
||||
print("✓ All methods successfully proxied")
|
||||
|
||||
finally:
|
||||
client.close()
|
||||
|
||||
|
||||
def test_embedded_multiple_banks(llm_config):
|
||||
"""
|
||||
Test that HindsightEmbedded can work with multiple banks.
|
||||
"""
|
||||
profile = f"test_multibank_{uuid.uuid4().hex[:8]}"
|
||||
bank1_id = f"bank1_{uuid.uuid4().hex[:8]}"
|
||||
bank2_id = f"bank2_{uuid.uuid4().hex[:8]}"
|
||||
|
||||
client = HindsightEmbedded(profile=profile, log_level="info", **llm_config)
|
||||
|
||||
try:
|
||||
# Create first bank and store data
|
||||
client.create_bank(bank_id=bank1_id, name="Bank 1")
|
||||
client.retain(bank_id=bank1_id, content="Alice prefers Python for data science")
|
||||
|
||||
# Create second bank and store data
|
||||
client.create_bank(bank_id=bank2_id, name="Bank 2")
|
||||
client.retain(bank_id=bank2_id, content="Bob uses JavaScript for web development")
|
||||
|
||||
# Recall from both banks
|
||||
results1 = client.recall(bank_id=bank1_id, query="programming language")
|
||||
results2 = client.recall(bank_id=bank2_id, query="programming language")
|
||||
|
||||
assert len(results1.results) > 0
|
||||
assert len(results2.results) > 0
|
||||
|
||||
# Verify banks are isolated (each should only see their own content)
|
||||
# This is a basic check - content isolation is tested more thoroughly in other tests
|
||||
assert results1.results[0].text != results2.results[0].text or len(results1.results) != len(
|
||||
results2.results
|
||||
)
|
||||
|
||||
finally:
|
||||
client.close()
|
||||
|
||||
|
||||
def test_embedded_profile_isolation(llm_config):
|
||||
"""
|
||||
Test that different profiles create isolated data stores.
|
||||
"""
|
||||
profile1 = f"test_iso1_{uuid.uuid4().hex[:8]}"
|
||||
profile2 = f"test_iso2_{uuid.uuid4().hex[:8]}"
|
||||
bank_id = "shared_bank_name" # Same bank_id in both profiles
|
||||
|
||||
client1 = HindsightEmbedded(profile=profile1, log_level="info", **llm_config)
|
||||
client2 = HindsightEmbedded(profile=profile2, log_level="info", **llm_config)
|
||||
|
||||
try:
|
||||
# Store data in profile1
|
||||
client1.retain(bank_id=bank_id, content="User likes TypeScript for frontend development")
|
||||
|
||||
# Store different data in profile2
|
||||
client2.retain(bank_id=bank_id, content="User prefers Rust for systems programming")
|
||||
|
||||
# Each profile should only see its own data
|
||||
results1 = client1.recall(bank_id=bank_id, query="programming preference")
|
||||
results2 = client2.recall(bank_id=bank_id, query="programming preference")
|
||||
|
||||
# Both should have results
|
||||
assert len(results1.results) > 0
|
||||
assert len(results2.results) > 0
|
||||
|
||||
# Results should be different (basic isolation check)
|
||||
# Note: This is a basic sanity check. Full isolation is ensured by pg0's data directory separation
|
||||
|
||||
finally:
|
||||
client1.close()
|
||||
client2.close()
|
||||
|
||||
|
||||
def test_embedded_error_after_close(llm_config):
|
||||
"""
|
||||
Test that using HindsightEmbedded after close() raises an error.
|
||||
"""
|
||||
profile = f"test_error_{uuid.uuid4().hex[:8]}"
|
||||
bank_id = f"bank_{uuid.uuid4().hex[:8]}"
|
||||
|
||||
client = HindsightEmbedded(profile=profile, log_level="info", **llm_config)
|
||||
|
||||
# Use it once to start server
|
||||
client.retain(bank_id=bank_id, content="Test")
|
||||
|
||||
# Close the client
|
||||
client.close()
|
||||
assert not client.is_running
|
||||
|
||||
# Trying to use it after close should raise an error
|
||||
with pytest.raises(RuntimeError, match="Cannot use HindsightEmbedded after it has been closed"):
|
||||
client.retain(bank_id=bank_id, content="This should fail")
|
||||
@@ -1,170 +0,0 @@
|
||||
"""Test that API namespaces ensure daemon is started before each call."""
|
||||
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from hindsight import HindsightEmbedded
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def embedded_client():
|
||||
"""Create an embedded client for testing."""
|
||||
return HindsightEmbedded(
|
||||
profile="test",
|
||||
llm_provider="openai",
|
||||
llm_api_key="test-key",
|
||||
)
|
||||
|
||||
|
||||
def test_banks_create_ensures_daemon_started(embedded_client):
|
||||
"""Test that banks.create() calls _ensure_started()."""
|
||||
# Mock _ensure_started to track calls
|
||||
with patch.object(embedded_client, "_ensure_started") as mock_ensure:
|
||||
# Mock the underlying client to avoid actual API call
|
||||
mock_client = Mock()
|
||||
embedded_client._client = mock_client
|
||||
|
||||
# Call namespace method
|
||||
try:
|
||||
embedded_client.banks.create(bank_id="test", name="Test Bank")
|
||||
except Exception:
|
||||
pass # We don't care if the actual call fails
|
||||
|
||||
# Verify _ensure_started was called
|
||||
mock_ensure.assert_called_once()
|
||||
|
||||
|
||||
def test_mental_models_list_ensures_daemon_started(embedded_client):
|
||||
"""Test that mental_models.list() calls _ensure_started()."""
|
||||
with patch.object(embedded_client, "_ensure_started") as mock_ensure:
|
||||
mock_client = Mock()
|
||||
embedded_client._client = mock_client
|
||||
|
||||
try:
|
||||
embedded_client.mental_models.list(bank_id="test")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
mock_ensure.assert_called_once()
|
||||
|
||||
|
||||
def test_directives_list_ensures_daemon_started(embedded_client):
|
||||
"""Test that directives.list() calls _ensure_started()."""
|
||||
with patch.object(embedded_client, "_ensure_started") as mock_ensure:
|
||||
mock_client = Mock()
|
||||
embedded_client._client = mock_client
|
||||
|
||||
try:
|
||||
embedded_client.directives.list(bank_id="test")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
mock_ensure.assert_called_once()
|
||||
|
||||
|
||||
def test_memories_list_ensures_daemon_started(embedded_client):
|
||||
"""Test that memories.list() calls _ensure_started()."""
|
||||
with patch.object(embedded_client, "_ensure_started") as mock_ensure:
|
||||
mock_client = Mock()
|
||||
embedded_client._client = mock_client
|
||||
|
||||
try:
|
||||
embedded_client.memories.list(bank_id="test")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
mock_ensure.assert_called_once()
|
||||
|
||||
|
||||
def test_multiple_calls_ensure_daemon_each_time(embedded_client):
|
||||
"""Test that each namespace call ensures daemon is started."""
|
||||
with patch.object(embedded_client, "_ensure_started") as mock_ensure:
|
||||
mock_client = Mock()
|
||||
embedded_client._client = mock_client
|
||||
|
||||
# Make multiple calls
|
||||
try:
|
||||
embedded_client.banks.create(bank_id="test", name="Test")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
try:
|
||||
embedded_client.mental_models.list(bank_id="test")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
try:
|
||||
embedded_client.directives.list(bank_id="test")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Should be called 3 times (once per namespace method call)
|
||||
assert mock_ensure.call_count == 3
|
||||
|
||||
|
||||
def test_daemon_restart_handling(embedded_client):
|
||||
"""Test that namespace methods can recover from daemon crash."""
|
||||
call_count = 0
|
||||
|
||||
def mock_ensure_started():
|
||||
"""Mock that simulates daemon restart."""
|
||||
nonlocal call_count
|
||||
call_count += 1
|
||||
# Create a new mock client each time (simulating daemon restart)
|
||||
embedded_client._client = Mock()
|
||||
embedded_client._started = True
|
||||
|
||||
with patch.object(embedded_client, "_ensure_started", side_effect=mock_ensure_started):
|
||||
# First call - daemon starts
|
||||
embedded_client.banks.create(bank_id="test", name="Test")
|
||||
assert call_count == 1
|
||||
|
||||
# Simulate daemon crash by clearing client
|
||||
embedded_client._client = None
|
||||
embedded_client._started = False
|
||||
|
||||
# Second call - daemon restarts
|
||||
embedded_client.banks.create(bank_id="test", name="Test")
|
||||
assert call_count == 2
|
||||
|
||||
|
||||
def test_ensure_started_calls_manager(embedded_client):
|
||||
"""Test that _ensure_started actually starts the daemon via manager."""
|
||||
# Mock the manager
|
||||
mock_manager = Mock()
|
||||
mock_manager.ensure_running.return_value = True
|
||||
mock_manager.get_url.return_value = "http://localhost:54321"
|
||||
|
||||
embedded_client._manager = mock_manager
|
||||
|
||||
# Mock Hindsight client constructor
|
||||
with patch("hindsight.embedded.Hindsight") as mock_hindsight_class:
|
||||
mock_client = Mock()
|
||||
mock_hindsight_class.return_value = mock_client
|
||||
|
||||
# Call _ensure_started
|
||||
embedded_client._ensure_started()
|
||||
|
||||
# Verify manager was called
|
||||
mock_manager.ensure_running.assert_called_once_with(
|
||||
embedded_client.config, embedded_client.profile
|
||||
)
|
||||
mock_manager.get_url.assert_called_once_with(embedded_client.profile)
|
||||
|
||||
# Verify Hindsight client was created
|
||||
mock_hindsight_class.assert_called_once_with(base_url="http://localhost:54321")
|
||||
|
||||
|
||||
def test_namespace_singleton_behavior(embedded_client):
|
||||
"""Test that namespace properties return the same instance."""
|
||||
banks1 = embedded_client.banks
|
||||
banks2 = embedded_client.banks
|
||||
|
||||
# Should be the same instance
|
||||
assert banks1 is banks2
|
||||
|
||||
# Same for other namespaces
|
||||
assert embedded_client.mental_models is embedded_client.mental_models
|
||||
assert embedded_client.directives is embedded_client.directives
|
||||
assert embedded_client.memories is embedded_client.memories
|
||||
@@ -1,137 +0,0 @@
|
||||
# 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
|
||||
-70
@@ -1,70 +0,0 @@
|
||||
"""Add file_storage table for BYTEA-based file storage
|
||||
|
||||
Revision ID: a1b2c3d4e5f6
|
||||
Revises: y0t1u2v3w4x5
|
||||
Create Date: 2026-02-16
|
||||
|
||||
Creates a dedicated table for storing uploaded files using BYTEA.
|
||||
This provides zero-config file storage that "just works" for development
|
||||
and small deployments. For production/scale, use S3-compatible storage.
|
||||
|
||||
Files are stored in a separate table to avoid bloating the documents table.
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
revision: str = "a1b2c3d4e5f6"
|
||||
down_revision: str | Sequence[str] | None = "y0t1u2v3w4x5"
|
||||
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:
|
||||
"""Create file_storage table for BYTEA storage."""
|
||||
schema = _get_schema_prefix()
|
||||
|
||||
# Create file_storage table (minimal: just key + data)
|
||||
op.execute(
|
||||
f"""
|
||||
CREATE TABLE IF NOT EXISTS {schema}file_storage (
|
||||
storage_key TEXT PRIMARY KEY,
|
||||
data BYTEA NOT NULL
|
||||
)
|
||||
"""
|
||||
)
|
||||
|
||||
# Add file tracking columns to documents table
|
||||
op.execute(
|
||||
f"""
|
||||
ALTER TABLE {schema}documents
|
||||
ADD COLUMN IF NOT EXISTS file_storage_key TEXT,
|
||||
ADD COLUMN IF NOT EXISTS file_original_name TEXT,
|
||||
ADD COLUMN IF NOT EXISTS file_content_type TEXT
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
"""Remove file_storage table and related columns."""
|
||||
schema = _get_schema_prefix()
|
||||
|
||||
# Drop columns from documents table
|
||||
op.execute(
|
||||
f"""
|
||||
ALTER TABLE {schema}documents
|
||||
DROP COLUMN IF EXISTS file_storage_key,
|
||||
DROP COLUMN IF EXISTS file_original_name,
|
||||
DROP COLUMN IF EXISTS file_content_type
|
||||
"""
|
||||
)
|
||||
|
||||
# Drop file_storage table
|
||||
op.execute(f"DROP TABLE IF EXISTS {schema}file_storage")
|
||||
-88
@@ -1,88 +0,0 @@
|
||||
"""Add text_signals column to memory_units for enriched BM25 indexing.
|
||||
|
||||
text_signals stores a denormalized space-separated string of entity names
|
||||
(and future signals) to improve full-text search recall without polluting
|
||||
the stored fact text.
|
||||
|
||||
- vchord: text_signals included in tokenize() at insert time
|
||||
- native: search_vector GENERATED column regenerated to include text_signals
|
||||
- pg_textsearch: no change (index only supports a single base column)
|
||||
|
||||
Revision ID: a2b3c4d5e6f7
|
||||
Revises: z1u2v3w4x5y6
|
||||
Create Date: 2026-02-28
|
||||
"""
|
||||
|
||||
import os
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
revision: str = "a2b3c4d5e6f7"
|
||||
down_revision: str | Sequence[str] | None = "aa2b3c4d5e6f"
|
||||
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 _detect_text_search_extension() -> str:
|
||||
return os.getenv("HINDSIGHT_API_TEXT_SEARCH_EXTENSION", "native").lower()
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
table = f"{schema}memory_units"
|
||||
text_search_ext = _detect_text_search_extension()
|
||||
|
||||
# Add text_signals column (nullable TEXT, populated at retain time)
|
||||
op.execute(f"ALTER TABLE {table} ADD COLUMN IF NOT EXISTS text_signals TEXT")
|
||||
|
||||
if text_search_ext == "native":
|
||||
# Native PostgreSQL: drop and recreate the GENERATED tsvector column to include text_signals
|
||||
op.execute(f"ALTER TABLE {table} DROP COLUMN IF EXISTS search_vector")
|
||||
op.execute(f"""
|
||||
ALTER TABLE {table}
|
||||
ADD COLUMN search_vector tsvector
|
||||
GENERATED ALWAYS AS (
|
||||
to_tsvector('english',
|
||||
COALESCE(text, '') || ' ' ||
|
||||
COALESCE(context, '') || ' ' ||
|
||||
COALESCE(text_signals, '')
|
||||
)
|
||||
) STORED
|
||||
""")
|
||||
# Recreate GIN index (was dropped with the column)
|
||||
op.execute(f"""
|
||||
CREATE INDEX IF NOT EXISTS idx_memory_units_text_search
|
||||
ON {table} USING gin(search_vector)
|
||||
""")
|
||||
|
||||
# vchord: tokenize() call in fact_storage.py is updated to include text_signals at insert time
|
||||
# pg_textsearch: no change — index operates on the base `text` column only
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
table = f"{schema}memory_units"
|
||||
text_search_ext = _detect_text_search_extension()
|
||||
|
||||
if text_search_ext == "native":
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema}idx_memory_units_text_search")
|
||||
op.execute(f"ALTER TABLE {table} DROP COLUMN IF EXISTS search_vector")
|
||||
op.execute(f"""
|
||||
ALTER TABLE {table}
|
||||
ADD COLUMN search_vector tsvector
|
||||
GENERATED ALWAYS AS (
|
||||
to_tsvector('english', COALESCE(text, '') || ' ' || COALESCE(context, ''))
|
||||
) STORED
|
||||
""")
|
||||
op.execute(f"""
|
||||
CREATE INDEX idx_memory_units_text_search
|
||||
ON {table} USING gin(search_vector)
|
||||
""")
|
||||
|
||||
op.execute(f"ALTER TABLE {table} DROP COLUMN IF EXISTS text_signals")
|
||||
-54
@@ -1,54 +0,0 @@
|
||||
"""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")
|
||||
@@ -1,36 +0,0 @@
|
||||
"""Make event_date nullable in memory_units to support timestamp-free content
|
||||
|
||||
Revision ID: aa2b3c4d5e6f
|
||||
Revises: z1u2v3w4x5y6
|
||||
Create Date: 2026-03-02
|
||||
|
||||
When callers retain content without a timestamp (e.g. fictional documents, static text),
|
||||
the event_date column should be allowed to be NULL rather than defaulting to utcnow().
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
revision: str = "aa2b3c4d5e6f"
|
||||
down_revision: str | Sequence[str] | None = "z1u2v3w4x5y6"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _get_schema_prefix() -> str:
|
||||
"""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 ALTER COLUMN event_date DROP NOT NULL")
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
# Backfill NULLs with now() before restoring the NOT NULL constraint
|
||||
op.execute(f"UPDATE {schema}memory_units SET event_date = now() WHERE event_date IS NULL")
|
||||
op.execute(f"ALTER TABLE {schema}memory_units ALTER COLUMN event_date SET NOT NULL")
|
||||
-68
@@ -1,68 +0,0 @@
|
||||
"""Add partial indexes on memory_units temporal date fields for fast temporal retrieval
|
||||
|
||||
Revision ID: b3c4d5e6f7g8
|
||||
Revises: c1a2b3d4e5f6
|
||||
Create Date: 2026-03-02
|
||||
|
||||
The temporal retrieval entry-point query filters memory_units by occurred_start,
|
||||
occurred_end, and mentioned_at using OR conditions. Without dedicated indexes the
|
||||
planner falls back to a sequential scan of all bank rows after applying the
|
||||
(bank_id, fact_type) index, then re-checks each date field.
|
||||
|
||||
These three partial indexes give the planner bitmap-index scan options for the
|
||||
three most common date predicates, dramatically reducing the row set before any
|
||||
embedding computation is required.
|
||||
|
||||
All indexes are created CONCURRENTLY so the migration does not block writes on
|
||||
memory_units during production deployments. CONCURRENTLY requires running outside
|
||||
a transaction block; see migrations.py for how this is handled safely.
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
revision: str = "b3c4d5e6f7g8"
|
||||
down_revision: str | Sequence[str] | None = "c1a2b3d4e5f6"
|
||||
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()
|
||||
# Partial index on occurred_start (covers "occurred_start BETWEEN $4 AND $5")
|
||||
op.execute("COMMIT")
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_units_bank_occurred_start "
|
||||
f"ON {schema}memory_units(bank_id, fact_type, occurred_start) "
|
||||
f"WHERE occurred_start IS NOT NULL"
|
||||
)
|
||||
# Partial index on occurred_end (covers "occurred_end BETWEEN $4 AND $5")
|
||||
op.execute("COMMIT")
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_units_bank_occurred_end "
|
||||
f"ON {schema}memory_units(bank_id, fact_type, occurred_end) "
|
||||
f"WHERE occurred_end IS NOT NULL"
|
||||
)
|
||||
# Partial index on mentioned_at (covers "mentioned_at BETWEEN $4 AND $5")
|
||||
op.execute("COMMIT")
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_units_bank_mentioned_at "
|
||||
f"ON {schema}memory_units(bank_id, fact_type, mentioned_at) "
|
||||
f"WHERE mentioned_at IS NOT NULL"
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
op.execute("COMMIT")
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_units_bank_mentioned_at")
|
||||
op.execute("COMMIT")
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_units_bank_occurred_end")
|
||||
op.execute("COMMIT")
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_units_bank_occurred_start")
|
||||
-34
@@ -1,34 +0,0 @@
|
||||
"""Backfill observation_scopes column if missing.
|
||||
|
||||
This migration ensures observation_scopes exists even on databases that had
|
||||
revision z1u2v3w4x5y6 applied when it referred to the old text_signals migration
|
||||
(before it was renamed to a2b3c4d5e6f7). The ADD COLUMN IF NOT EXISTS makes this
|
||||
a no-op on databases that already have the column.
|
||||
|
||||
Revision ID: b4c5d6e7f8a9
|
||||
Revises: a2b3c4d5e6f7
|
||||
Create Date: 2026-03-02
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
revision: str = "b4c5d6e7f8a9"
|
||||
down_revision: str | Sequence[str] | None = "a2b3c4d5e6f7"
|
||||
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}memory_units ADD COLUMN IF NOT EXISTS observation_scopes JSONB")
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
pass # intentionally no-op — safe to leave the column in place
|
||||
-46
@@ -1,46 +0,0 @@
|
||||
"""Enable pg_trgm extension and add GIN trigram index on entities.canonical_name
|
||||
|
||||
Revision ID: c1a2b3d4e5f6
|
||||
Revises: b4c5d6e7f8a9
|
||||
Create Date: 2026-03-02
|
||||
|
||||
Index is created CONCURRENTLY so the migration does not block writes on entities
|
||||
during production deployments. CONCURRENTLY requires running outside a transaction
|
||||
block; see migrations.py for how this is handled safely.
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
revision: str = "c1a2b3d4e5f6"
|
||||
down_revision: str | Sequence[str] | None = "b4c5d6e7f8a9"
|
||||
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:
|
||||
# pg_trgm ships with every standard PostgreSQL installation as a contrib module.
|
||||
# It enables fast similarity lookups via GIN indexes, used for entity name matching.
|
||||
op.execute("CREATE EXTENSION IF NOT EXISTS pg_trgm")
|
||||
|
||||
schema = _get_schema_prefix()
|
||||
# GIN index on canonical_name enables sub-millisecond trigram similarity queries
|
||||
# (% operator, similarity()) instead of full-table scans across all bank entities.
|
||||
op.execute("COMMIT")
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS entities_canonical_name_trgm_idx "
|
||||
f"ON {schema}entities USING GIN (canonical_name gin_trgm_ops)"
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
op.execute("COMMIT")
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}entities_canonical_name_trgm_idx")
|
||||
# Note: not dropping pg_trgm extension as other indexes may depend on it
|
||||
-30
@@ -1,30 +0,0 @@
|
||||
"""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")
|
||||
-83
@@ -1,83 +0,0 @@
|
||||
"""Add covering and composite indexes to speed up link expansion graph retrieval.
|
||||
|
||||
Two indexes target the two bottlenecks identified by EXPLAIN ANALYZE on a 17M-row
|
||||
memory_links table:
|
||||
|
||||
1. idx_memory_links_to_type_weight (to_unit_id, link_type, weight DESC)
|
||||
The semantic incoming direction — finding facts that consider seeds as their
|
||||
nearest neighbour — currently hits an expensive BitmapAnd of two separate
|
||||
bitmap scans (to_unit_id bitmap ∩ link_type bitmap). A composite index
|
||||
on (to_unit_id, link_type) turns this into a single index scan and reduces
|
||||
latency from ~36 ms to < 5 ms per query.
|
||||
|
||||
2. idx_memory_links_entity_covering (from_unit_id) INCLUDE (to_unit_id, entity_id)
|
||||
WHERE link_type = 'entity'
|
||||
The entity co-occurrence expansion uses COUNT(DISTINCT ml.entity_id) and
|
||||
joins on ml.to_unit_id. Without a covering index the planner must read
|
||||
~2 500 heap pages to fetch entity_id and to_unit_id after the bitmap index
|
||||
scan, adding ~230 ms of random I/O. INCLUDE adds those two columns to the
|
||||
index leaf pages so the entire query can be served from the index (index-only
|
||||
scan), eliminating the heap reads entirely.
|
||||
Partial index (WHERE link_type = 'entity') keeps index size ~40 % smaller.
|
||||
|
||||
Both indexes are created with CONCURRENTLY so the migration does not block
|
||||
concurrent reads or writes on memory_links. CONCURRENTLY requires running
|
||||
outside a transaction block, so the migration emits an explicit COMMIT before
|
||||
each statement and uses IF NOT EXISTS for idempotency.
|
||||
|
||||
Revision ID: d2e3f4a5b6c7
|
||||
Revises: b3c4d5e6f7g8
|
||||
Create Date: 2026-03-02
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
revision: str = "d2e3f4a5b6c7"
|
||||
down_revision: str | Sequence[str] | None = "b3c4d5e6f7g8"
|
||||
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, then issue each CONCURRENTLY
|
||||
# statement in its own implicit autocommit transaction.
|
||||
# IF NOT EXISTS makes each statement idempotent if the migration is retried.
|
||||
|
||||
# Index for the semantic *incoming* direction in link_expansion_retrieval.py.
|
||||
# Replaces the BitmapAnd of idx_memory_links_to_unit ∩ idx_memory_links_link_type
|
||||
# with a single composite index scan.
|
||||
op.execute("COMMIT")
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_links_to_type_weight "
|
||||
f"ON {schema}memory_links(to_unit_id, link_type, weight DESC)"
|
||||
)
|
||||
|
||||
# Covering index for entity co-occurrence expansion.
|
||||
# Enables an index-only scan: entity_id and to_unit_id are read from the
|
||||
# index leaf pages instead of the heap, eliminating ~2 500 random heap-page
|
||||
# reads per expansion query.
|
||||
op.execute("COMMIT")
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_links_entity_covering "
|
||||
f"ON {schema}memory_links(from_unit_id) "
|
||||
f"INCLUDE (to_unit_id, entity_id) "
|
||||
f"WHERE link_type = 'entity'"
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
op.execute("COMMIT")
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_links_entity_covering")
|
||||
op.execute("COMMIT")
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_links_to_type_weight")
|
||||
-53
@@ -1,53 +0,0 @@
|
||||
"""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
@@ -1,131 +0,0 @@
|
||||
"""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")
|
||||
@@ -1,62 +0,0 @@
|
||||
"""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
@@ -1,73 +0,0 @@
|
||||
"""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")
|
||||
-33
@@ -1,33 +0,0 @@
|
||||
"""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")
|
||||
-317
@@ -1,317 +0,0 @@
|
||||
"""learnings_and_pinned_reflections
|
||||
|
||||
Revision ID: n9i0j1k2l3m4
|
||||
Revises: m8h9i0j1k2l3
|
||||
Create Date: 2026-01-21 00:00:00.000000
|
||||
|
||||
This migration:
|
||||
1. Creates the 'learnings' table for automatic bottom-up consolidation
|
||||
2. Creates the 'pinned_reflections' table for user-curated living documents
|
||||
3. Adds consolidation tracking columns to the 'banks' table
|
||||
"""
|
||||
|
||||
import os
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
from sqlalchemy import text
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = "n9i0j1k2l3m4"
|
||||
down_revision: str | Sequence[str] | None = "m8h9i0j1k2l3"
|
||||
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 _detect_vector_extension() -> str:
|
||||
"""
|
||||
Detect or validate vector extension: 'pgvector', 'vchord', or 'pgvectorscale'.
|
||||
Respects HINDSIGHT_API_VECTOR_EXTENSION env var if set.
|
||||
"""
|
||||
conn = op.get_bind()
|
||||
vector_extension = os.getenv("HINDSIGHT_API_VECTOR_EXTENSION", "pgvector").lower()
|
||||
|
||||
# Validate configured extension is installed
|
||||
if vector_extension == "pgvectorscale":
|
||||
# pgvectorscale/DiskANN requires pgvector
|
||||
pgvector_check = conn.execute(text("SELECT 1 FROM pg_extension WHERE extname = 'vector'")).scalar()
|
||||
if not pgvector_check:
|
||||
raise RuntimeError(
|
||||
"DiskANN requires pgvector. Install with: CREATE EXTENSION vector; then vectorscale or pg_diskann CASCADE;"
|
||||
)
|
||||
# Check for either vectorscale (open source) or pg_diskann (Azure)
|
||||
vectorscale_check = conn.execute(text("SELECT 1 FROM pg_extension WHERE extname = 'vectorscale'")).scalar()
|
||||
pg_diskann_check = conn.execute(text("SELECT 1 FROM pg_extension WHERE extname = 'pg_diskann'")).scalar()
|
||||
|
||||
if vectorscale_check:
|
||||
return "pgvectorscale"
|
||||
elif pg_diskann_check:
|
||||
return "pg_diskann"
|
||||
else:
|
||||
raise RuntimeError(
|
||||
"Configured vector extension 'pgvectorscale' not found. Install either:\n"
|
||||
" - pgvectorscale: CREATE EXTENSION vectorscale CASCADE;\n"
|
||||
" - pg_diskann (Azure): CREATE EXTENSION pg_diskann CASCADE;"
|
||||
)
|
||||
elif vector_extension == "vchord":
|
||||
vchord_check = conn.execute(text("SELECT 1 FROM pg_extension WHERE extname = 'vchord'")).scalar()
|
||||
if not vchord_check:
|
||||
raise RuntimeError(
|
||||
"Configured vector extension 'vchord' not found. Install it with: CREATE EXTENSION vchord CASCADE;"
|
||||
)
|
||||
return "vchord"
|
||||
elif vector_extension == "pgvector":
|
||||
pgvector_check = conn.execute(text("SELECT 1 FROM pg_extension WHERE extname = 'vector'")).scalar()
|
||||
if not pgvector_check:
|
||||
raise RuntimeError(
|
||||
"Configured vector extension 'pgvector' not found. Install it with: CREATE EXTENSION vector;"
|
||||
)
|
||||
return "pgvector"
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Invalid HINDSIGHT_API_VECTOR_EXTENSION: {vector_extension}. Must be 'pgvector', 'vchord', or 'pgvectorscale'"
|
||||
)
|
||||
|
||||
|
||||
def _detect_text_search_extension() -> str:
|
||||
"""
|
||||
Detect or validate text search extension: 'native', 'vchord', or 'pg_textsearch'.
|
||||
Respects HINDSIGHT_API_TEXT_SEARCH_EXTENSION env var.
|
||||
Creates the extension if needed.
|
||||
"""
|
||||
text_search_extension = os.getenv("HINDSIGHT_API_TEXT_SEARCH_EXTENSION", "native").lower()
|
||||
|
||||
if text_search_extension == "vchord":
|
||||
# Create vchord_bm25 extension if not exists
|
||||
try:
|
||||
op.execute("CREATE EXTENSION IF NOT EXISTS vchord_bm25 CASCADE")
|
||||
except Exception:
|
||||
# Extension might already exist or user lacks permissions - verify it exists
|
||||
conn = op.get_bind()
|
||||
result = conn.execute(text("SELECT 1 FROM pg_extension WHERE extname = 'vchord_bm25'")).fetchone()
|
||||
if not result:
|
||||
# Extension truly doesn't exist - re-raise the error
|
||||
raise
|
||||
return "vchord"
|
||||
elif text_search_extension == "pg_textsearch":
|
||||
# Create pg_textsearch extension if not exists
|
||||
try:
|
||||
op.execute("CREATE EXTENSION IF NOT EXISTS pg_textsearch CASCADE")
|
||||
except Exception:
|
||||
# Extension might already exist or user lacks permissions - verify it exists
|
||||
conn = op.get_bind()
|
||||
result = conn.execute(text("SELECT 1 FROM pg_extension WHERE extname = 'pg_textsearch'")).fetchone()
|
||||
if not result:
|
||||
# Extension truly doesn't exist - re-raise the error
|
||||
raise
|
||||
return "pg_textsearch"
|
||||
elif text_search_extension == "native":
|
||||
return "native"
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Invalid HINDSIGHT_API_TEXT_SEARCH_EXTENSION: {text_search_extension}. Must be 'native', 'vchord', or 'pg_textsearch'"
|
||||
)
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
"""Create learnings and pinned_reflections tables."""
|
||||
schema = _get_schema_prefix()
|
||||
|
||||
# Detect which vector extension is available
|
||||
vector_ext = _detect_vector_extension()
|
||||
|
||||
# Detect which text search extension to use
|
||||
text_search_ext = _detect_text_search_extension()
|
||||
|
||||
# 1. Create learnings table
|
||||
op.execute(f"""
|
||||
CREATE TABLE {schema}learnings (
|
||||
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
||||
bank_id VARCHAR(64) NOT NULL,
|
||||
text TEXT NOT NULL,
|
||||
proof_count INT NOT NULL DEFAULT 1,
|
||||
history JSONB DEFAULT '[]'::jsonb,
|
||||
mission_context VARCHAR(64),
|
||||
pre_mission_change BOOLEAN DEFAULT FALSE,
|
||||
embedding vector(384),
|
||||
tags VARCHAR[] DEFAULT ARRAY[]::VARCHAR[],
|
||||
created_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT now(),
|
||||
updated_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT now()
|
||||
)
|
||||
""")
|
||||
|
||||
# Add foreign key constraint
|
||||
op.execute(f"""
|
||||
ALTER TABLE {schema}learnings
|
||||
ADD CONSTRAINT fk_learnings_bank_id
|
||||
FOREIGN KEY (bank_id) REFERENCES {schema}banks(bank_id) ON DELETE CASCADE
|
||||
""")
|
||||
|
||||
# Indexes for learnings
|
||||
op.execute(f"CREATE INDEX idx_learnings_bank_id ON {schema}learnings(bank_id)")
|
||||
|
||||
# Create vector index based on detected extension
|
||||
if vector_ext == "pgvectorscale":
|
||||
op.execute(f"""
|
||||
CREATE INDEX idx_learnings_embedding ON {schema}learnings
|
||||
USING diskann (embedding vector_cosine_ops)
|
||||
WITH (num_neighbors = 50)
|
||||
""")
|
||||
elif vector_ext == "pg_diskann":
|
||||
op.execute(f"""
|
||||
CREATE INDEX idx_learnings_embedding ON {schema}learnings
|
||||
USING diskann (embedding vector_cosine_ops)
|
||||
WITH (max_neighbors = 50)
|
||||
""")
|
||||
elif vector_ext == "vchord":
|
||||
op.execute(f"""
|
||||
CREATE INDEX idx_learnings_embedding ON {schema}learnings
|
||||
USING vchordrq (embedding vector_l2_ops)
|
||||
""")
|
||||
else: # pgvector
|
||||
op.execute(f"""
|
||||
CREATE INDEX idx_learnings_embedding ON {schema}learnings
|
||||
USING hnsw (embedding vector_cosine_ops)
|
||||
""")
|
||||
|
||||
op.execute(f"CREATE INDEX idx_learnings_tags ON {schema}learnings USING GIN(tags)")
|
||||
|
||||
# Full-text search for learnings
|
||||
if text_search_ext == "vchord":
|
||||
# VectorChord BM25: bm25vector type (no GENERATED - tokenization happens on INSERT)
|
||||
# Note: vchord_bm25 extension creates types in bm25_catalog schema
|
||||
op.execute(f"""
|
||||
ALTER TABLE {schema}learnings ADD COLUMN search_vector bm25_catalog.bm25vector
|
||||
""")
|
||||
op.execute(f"""
|
||||
CREATE INDEX idx_learnings_text_search ON {schema}learnings
|
||||
USING bm25 (search_vector bm25_catalog.bm25_ops)
|
||||
""")
|
||||
elif text_search_ext == "pg_textsearch":
|
||||
# Timescale pg_textsearch: dummy TEXT column for consistency (indexes operate on base columns directly)
|
||||
op.execute(f"""
|
||||
ALTER TABLE {schema}learnings ADD COLUMN search_vector TEXT
|
||||
""")
|
||||
op.execute(f"""
|
||||
CREATE INDEX idx_learnings_text_search ON {schema}learnings
|
||||
USING bm25(text) WITH (text_config='english')
|
||||
""")
|
||||
else: # native
|
||||
# Native PostgreSQL: tsvector with automatic generation
|
||||
op.execute(f"""
|
||||
ALTER TABLE {schema}learnings ADD COLUMN search_vector tsvector
|
||||
GENERATED ALWAYS AS (to_tsvector('english', text)) STORED
|
||||
""")
|
||||
op.execute(f"CREATE INDEX idx_learnings_text_search ON {schema}learnings USING gin(search_vector)")
|
||||
|
||||
# 2. Create pinned_reflections table
|
||||
op.execute(f"""
|
||||
CREATE TABLE {schema}pinned_reflections (
|
||||
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
||||
bank_id VARCHAR(64) NOT NULL,
|
||||
name VARCHAR(256) NOT NULL,
|
||||
source_query TEXT NOT NULL,
|
||||
content TEXT NOT NULL,
|
||||
embedding vector(384),
|
||||
tags VARCHAR[] DEFAULT ARRAY[]::VARCHAR[],
|
||||
last_refreshed_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT now(),
|
||||
created_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT now()
|
||||
)
|
||||
""")
|
||||
|
||||
# Add foreign key constraint
|
||||
op.execute(f"""
|
||||
ALTER TABLE {schema}pinned_reflections
|
||||
ADD CONSTRAINT fk_pinned_reflections_bank_id
|
||||
FOREIGN KEY (bank_id) REFERENCES {schema}banks(bank_id) ON DELETE CASCADE
|
||||
""")
|
||||
|
||||
# Indexes for pinned_reflections
|
||||
op.execute(f"CREATE INDEX idx_pinned_reflections_bank_id ON {schema}pinned_reflections(bank_id)")
|
||||
|
||||
# Create vector index based on detected extension
|
||||
if vector_ext == "pgvectorscale":
|
||||
op.execute(f"""
|
||||
CREATE INDEX idx_pinned_reflections_embedding ON {schema}pinned_reflections
|
||||
USING diskann (embedding vector_cosine_ops)
|
||||
WITH (num_neighbors = 50)
|
||||
""")
|
||||
elif vector_ext == "pg_diskann":
|
||||
op.execute(f"""
|
||||
CREATE INDEX idx_pinned_reflections_embedding ON {schema}pinned_reflections
|
||||
USING diskann (embedding vector_cosine_ops)
|
||||
WITH (max_neighbors = 50)
|
||||
""")
|
||||
elif vector_ext == "vchord":
|
||||
op.execute(f"""
|
||||
CREATE INDEX idx_pinned_reflections_embedding ON {schema}pinned_reflections
|
||||
USING vchordrq (embedding vector_l2_ops)
|
||||
""")
|
||||
else: # pgvector
|
||||
op.execute(f"""
|
||||
CREATE INDEX idx_pinned_reflections_embedding ON {schema}pinned_reflections
|
||||
USING hnsw (embedding vector_cosine_ops)
|
||||
""")
|
||||
|
||||
op.execute(f"CREATE INDEX idx_pinned_reflections_tags ON {schema}pinned_reflections USING GIN(tags)")
|
||||
|
||||
# Full-text search for pinned_reflections
|
||||
if text_search_ext == "vchord":
|
||||
# VectorChord BM25: bm25vector type (no GENERATED - tokenization happens on INSERT/UPDATE)
|
||||
# Note: vchord_bm25 extension creates types in bm25_catalog schema
|
||||
op.execute(f"""
|
||||
ALTER TABLE {schema}pinned_reflections ADD COLUMN search_vector bm25_catalog.bm25vector
|
||||
""")
|
||||
op.execute(f"""
|
||||
CREATE INDEX idx_pinned_reflections_text_search ON {schema}pinned_reflections
|
||||
USING bm25 (search_vector bm25_catalog.bm25_ops)
|
||||
""")
|
||||
elif text_search_ext == "pg_textsearch":
|
||||
# Timescale pg_textsearch: dummy TEXT column for consistency (indexes operate on base columns directly)
|
||||
op.execute(f"""
|
||||
ALTER TABLE {schema}pinned_reflections ADD COLUMN search_vector TEXT
|
||||
""")
|
||||
op.execute(f"""
|
||||
CREATE INDEX idx_pinned_reflections_text_search ON {schema}pinned_reflections
|
||||
USING bm25(content)
|
||||
WITH (text_config='english')
|
||||
""")
|
||||
else: # native
|
||||
# Native PostgreSQL: tsvector with automatic generation
|
||||
op.execute(f"""
|
||||
ALTER TABLE {schema}pinned_reflections ADD COLUMN search_vector tsvector
|
||||
GENERATED ALWAYS AS (to_tsvector('english', COALESCE(name, '') || ' ' || content)) STORED
|
||||
""")
|
||||
op.execute(f"""
|
||||
CREATE INDEX idx_pinned_reflections_text_search ON {schema}pinned_reflections
|
||||
USING gin(search_vector)
|
||||
""")
|
||||
|
||||
# 3. Add consolidation tracking columns to banks table
|
||||
op.execute(f"""
|
||||
ALTER TABLE {schema}banks
|
||||
ADD COLUMN IF NOT EXISTS last_consolidated_at TIMESTAMP WITH TIME ZONE
|
||||
""")
|
||||
op.execute(f"""
|
||||
ALTER TABLE {schema}banks
|
||||
ADD COLUMN IF NOT EXISTS mission_changed_at TIMESTAMP WITH TIME ZONE
|
||||
""")
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
"""Drop learnings and pinned_reflections tables."""
|
||||
schema = _get_schema_prefix()
|
||||
|
||||
# Drop tables
|
||||
op.execute(f"DROP TABLE IF EXISTS {schema}learnings CASCADE")
|
||||
op.execute(f"DROP TABLE IF EXISTS {schema}pinned_reflections CASCADE")
|
||||
|
||||
# Remove columns from banks
|
||||
op.execute(f"ALTER TABLE {schema}banks DROP COLUMN IF EXISTS last_consolidated_at")
|
||||
op.execute(f"ALTER TABLE {schema}banks DROP COLUMN IF EXISTS mission_changed_at")
|
||||
-60
@@ -1,60 +0,0 @@
|
||||
"""Fix mental_models primary key to be scoped per bank
|
||||
|
||||
Revision ID: w8r9s0t1u2v3
|
||||
Revises: v7q8r9s0t1u2
|
||||
Create Date: 2026-02-05
|
||||
|
||||
This migration fixes a critical bank isolation bug where mental_models.id was
|
||||
globally unique across all banks instead of being scoped per bank. This caused
|
||||
conflicts when different banks tried to use the same custom ID.
|
||||
|
||||
CRITICAL FIX: Changes primary key from (id) to (bank_id, id) to ensure proper isolation.
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
revision: str = "w8r9s0t1u2v3"
|
||||
down_revision: str | Sequence[str] | None = "v7q8r9s0t1u2"
|
||||
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:
|
||||
"""Change mental_models primary key from (id) to (bank_id, id) for proper bank isolation."""
|
||||
schema = _get_schema_prefix()
|
||||
|
||||
# Drop the old primary key constraint (just id)
|
||||
# Note: The constraint might be named differently on different DBs
|
||||
# Try both old names (pinned_reflections_pkey from original, mental_models_pkey from rename)
|
||||
op.execute(f"ALTER TABLE {schema}mental_models DROP CONSTRAINT IF EXISTS pinned_reflections_pkey")
|
||||
op.execute(f"ALTER TABLE {schema}mental_models DROP CONSTRAINT IF EXISTS mental_models_pkey")
|
||||
|
||||
# Create the new composite primary key (bank_id, id)
|
||||
# This ensures IDs are scoped per bank, not globally
|
||||
op.execute(f"""
|
||||
ALTER TABLE {schema}mental_models
|
||||
ADD CONSTRAINT mental_models_pkey PRIMARY KEY (bank_id, id)
|
||||
""")
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
"""Revert mental_models primary key from (bank_id, id) to (id)."""
|
||||
schema = _get_schema_prefix()
|
||||
|
||||
# Drop the composite primary key
|
||||
op.execute(f"ALTER TABLE {schema}mental_models DROP CONSTRAINT IF EXISTS mental_models_pkey")
|
||||
|
||||
# Restore the old primary key (just id)
|
||||
# WARNING: This downgrade will fail if there are duplicate IDs across banks
|
||||
op.execute(f"""
|
||||
ALTER TABLE {schema}mental_models
|
||||
ADD CONSTRAINT mental_models_pkey PRIMARY KEY (id)
|
||||
""")
|
||||
-64
@@ -1,64 +0,0 @@
|
||||
"""Add config JSONB column to banks table for hierarchical configuration
|
||||
|
||||
Revision ID: x9s0t1u2v3w4
|
||||
Revises: w8r9s0t1u2v3
|
||||
Create Date: 2026-02-09
|
||||
|
||||
This migration adds a `config` JSONB column to the banks table to support
|
||||
per-bank configuration overrides. This enables hierarchical configuration where:
|
||||
- Global config is loaded from environment variables
|
||||
- Tenant config is provided via TenantExtension
|
||||
- Bank config overrides are stored in banks.config JSONB column
|
||||
|
||||
The config column stores overrides for hierarchical fields (LLM settings,
|
||||
retention parameters, retrieval settings, etc.) in Python field name format.
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import context, op
|
||||
from sqlalchemy.dialects.postgresql import JSONB
|
||||
|
||||
revision: str = "x9s0t1u2v3w4"
|
||||
down_revision: str | Sequence[str] | None = "w8r9s0t1u2v3"
|
||||
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:
|
||||
"""Add config JSONB column to banks table with GIN index."""
|
||||
schema = _get_schema_prefix()
|
||||
|
||||
# Add config column to banks table
|
||||
op.execute(f"""
|
||||
ALTER TABLE {schema}banks
|
||||
ADD COLUMN config JSONB NOT NULL DEFAULT '{{}}'::jsonb
|
||||
""")
|
||||
|
||||
# Add GIN index for efficient JSONB queries
|
||||
op.execute(f"""
|
||||
CREATE INDEX idx_banks_config
|
||||
ON {schema}banks
|
||||
USING gin(config)
|
||||
""")
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
"""Remove config column and index from banks table."""
|
||||
schema = _get_schema_prefix()
|
||||
|
||||
# Drop index first
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema}idx_banks_config")
|
||||
|
||||
# Drop column
|
||||
op.execute(f"""
|
||||
ALTER TABLE {schema}banks
|
||||
DROP COLUMN IF EXISTS config
|
||||
""")
|
||||
-49
@@ -1,49 +0,0 @@
|
||||
"""Add GIN index on async_operations.result_metadata for parent_operation_id queries
|
||||
|
||||
Revision ID: y0t1u2v3w4x5
|
||||
Revises: x9s0t1u2v3w4
|
||||
Create Date: 2026-02-13
|
||||
|
||||
This migration adds a GIN index on the result_metadata JSONB column in the
|
||||
async_operations table to support efficient queries for child operations by
|
||||
parent_operation_id.
|
||||
|
||||
The index enables fast lookups when querying for child operations:
|
||||
SELECT * FROM async_operations
|
||||
WHERE result_metadata::jsonb @> '{"parent_operation_id": "uuid"}'::jsonb
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
revision: str = "y0t1u2v3w4x5"
|
||||
down_revision: str | Sequence[str] | None = "x9s0t1u2v3w4"
|
||||
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:
|
||||
"""Add GIN index on result_metadata for efficient parent_operation_id queries."""
|
||||
schema = _get_schema_prefix()
|
||||
|
||||
# Add GIN index for JSONB containment queries (@> operator)
|
||||
op.execute(f"""
|
||||
CREATE INDEX IF NOT EXISTS idx_async_operations_result_metadata
|
||||
ON {schema}async_operations
|
||||
USING gin(result_metadata)
|
||||
""")
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
"""Remove GIN index on result_metadata."""
|
||||
schema = _get_schema_prefix()
|
||||
|
||||
# Drop index
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema}idx_async_operations_result_metadata")
|
||||
-35
@@ -1,35 +0,0 @@
|
||||
"""Add observation_scopes column to memory_units table
|
||||
|
||||
Revision ID: z1u2v3w4x5y6
|
||||
Revises: a1b2c3d4e5f6
|
||||
Create Date: 2026-02-25
|
||||
|
||||
Adds observation_scopes JSONB column to memory_units to control how observations
|
||||
are scoped during consolidation. Accepts "per_tag", "combined", or an explicit
|
||||
list of tag-set lists for custom multi-pass consolidation.
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
revision: str = "z1u2v3w4x5y6"
|
||||
down_revision: str | Sequence[str] | None = "a1b2c3d4e5f6"
|
||||
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 observation_scopes JSONB")
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
op.execute(f"ALTER TABLE {schema}memory_units DROP COLUMN IF EXISTS observation_scopes")
|
||||
@@ -1,452 +0,0 @@
|
||||
"""Hindsight MCP Server implementation using FastMCP (HTTP transport)."""
|
||||
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
from contextvars import ContextVar
|
||||
|
||||
from fastmcp import FastMCP
|
||||
|
||||
from hindsight_api import MemoryEngine
|
||||
from hindsight_api.config import _get_raw_config
|
||||
from hindsight_api.engine.memory_engine import _current_schema
|
||||
from hindsight_api.extensions import MCPExtension, load_extension
|
||||
from hindsight_api.extensions.tenant import AuthenticationError
|
||||
from hindsight_api.mcp_tools import MCPToolsConfig, register_mcp_tools
|
||||
from hindsight_api.models import RequestContext
|
||||
|
||||
# All tools available in the system (explicit list — no wildcards)
|
||||
_ALL_TOOLS: frozenset[str] = frozenset(
|
||||
{
|
||||
"retain",
|
||||
"recall",
|
||||
"reflect",
|
||||
"list_banks",
|
||||
"create_bank",
|
||||
"list_mental_models",
|
||||
"get_mental_model",
|
||||
"create_mental_model",
|
||||
"update_mental_model",
|
||||
"delete_mental_model",
|
||||
"refresh_mental_model",
|
||||
"list_directives",
|
||||
"create_directive",
|
||||
"delete_directive",
|
||||
"list_memories",
|
||||
"get_memory",
|
||||
"delete_memory",
|
||||
"list_documents",
|
||||
"get_document",
|
||||
"delete_document",
|
||||
"list_operations",
|
||||
"get_operation",
|
||||
"cancel_operation",
|
||||
"list_tags",
|
||||
"get_bank",
|
||||
"get_bank_stats",
|
||||
"update_bank",
|
||||
"delete_bank",
|
||||
"clear_memories",
|
||||
}
|
||||
)
|
||||
|
||||
# Configure logging from HINDSIGHT_API_LOG_LEVEL environment variable
|
||||
_log_level_str = os.environ.get("HINDSIGHT_API_LOG_LEVEL", "info").lower()
|
||||
_log_level_map = {
|
||||
"critical": logging.CRITICAL,
|
||||
"error": logging.ERROR,
|
||||
"warning": logging.WARNING,
|
||||
"info": logging.INFO,
|
||||
"debug": logging.DEBUG,
|
||||
"trace": logging.DEBUG,
|
||||
}
|
||||
logging.basicConfig(
|
||||
level=_log_level_map.get(_log_level_str, logging.INFO),
|
||||
format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
|
||||
)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Default bank_id from environment variable
|
||||
DEFAULT_BANK_ID = os.environ.get("HINDSIGHT_MCP_BANK_ID", "default")
|
||||
|
||||
# Legacy MCP authentication token (for backwards compatibility)
|
||||
# If set, this token is checked first before TenantExtension auth
|
||||
MCP_AUTH_TOKEN = os.environ.get("HINDSIGHT_API_MCP_AUTH_TOKEN")
|
||||
|
||||
# Context variable to hold the current bank_id
|
||||
_current_bank_id: ContextVar[str | None] = ContextVar("current_bank_id", default=None)
|
||||
|
||||
# Context variable to hold the current API key (for tenant auth propagation)
|
||||
_current_api_key: ContextVar[str | None] = ContextVar("current_api_key", default=None)
|
||||
|
||||
# Context variables for tenant_id and api_key_id (set by authenticate, used by usage metering)
|
||||
_current_tenant_id: ContextVar[str | None] = ContextVar("current_tenant_id", default=None)
|
||||
_current_api_key_id: ContextVar[str | None] = ContextVar("current_api_key_id", default=None)
|
||||
|
||||
|
||||
def get_current_bank_id() -> str | None:
|
||||
"""Get the current bank_id from context."""
|
||||
return _current_bank_id.get()
|
||||
|
||||
|
||||
def get_current_api_key() -> str | None:
|
||||
"""Get the current API key from context."""
|
||||
return _current_api_key.get()
|
||||
|
||||
|
||||
def get_current_tenant_id() -> str | None:
|
||||
"""Get the current tenant_id from context."""
|
||||
return _current_tenant_id.get()
|
||||
|
||||
|
||||
def get_current_api_key_id() -> str | None:
|
||||
"""Get the current api_key_id from context."""
|
||||
return _current_api_key_id.get()
|
||||
|
||||
|
||||
def create_mcp_server(memory: MemoryEngine, multi_bank: bool = True) -> FastMCP:
|
||||
"""
|
||||
Create and configure the Hindsight MCP server.
|
||||
|
||||
Args:
|
||||
memory: MemoryEngine instance (required)
|
||||
multi_bank: If True, expose all tools with bank_id parameters (default).
|
||||
If False, only expose bank-scoped tools without bank_id parameters.
|
||||
|
||||
Returns:
|
||||
Configured FastMCP server instance
|
||||
"""
|
||||
mcp = FastMCP("hindsight-mcp-server")
|
||||
|
||||
global_config = _get_raw_config()
|
||||
|
||||
# Tools available for this mode (multi-bank exposes all tools; single-bank excludes bank-management tools)
|
||||
_SINGLE_BANK_TOOLS: frozenset[str] = frozenset(
|
||||
{
|
||||
"retain",
|
||||
"recall",
|
||||
"reflect",
|
||||
"list_mental_models",
|
||||
"get_mental_model",
|
||||
"create_mental_model",
|
||||
"update_mental_model",
|
||||
"delete_mental_model",
|
||||
"refresh_mental_model",
|
||||
"list_directives",
|
||||
"create_directive",
|
||||
"delete_directive",
|
||||
"list_memories",
|
||||
"get_memory",
|
||||
"delete_memory",
|
||||
"list_documents",
|
||||
"get_document",
|
||||
"delete_document",
|
||||
"list_operations",
|
||||
"get_operation",
|
||||
"cancel_operation",
|
||||
"list_tags",
|
||||
"get_bank",
|
||||
"update_bank",
|
||||
"delete_bank",
|
||||
"clear_memories",
|
||||
}
|
||||
)
|
||||
base_tools: frozenset[str] | None = None if multi_bank else _SINGLE_BANK_TOOLS
|
||||
|
||||
# Apply global mcp_enabled_tools filter (env-level allowlist)
|
||||
if global_config.mcp_enabled_tools is not None:
|
||||
allowed = frozenset(global_config.mcp_enabled_tools)
|
||||
base_tools = (base_tools if base_tools is not None else _ALL_TOOLS) & allowed
|
||||
|
||||
# Configure and register tools using shared module
|
||||
config = MCPToolsConfig(
|
||||
bank_id_resolver=get_current_bank_id,
|
||||
api_key_resolver=get_current_api_key, # Propagate API key for tenant auth
|
||||
tenant_id_resolver=get_current_tenant_id, # Propagate tenant_id for usage metering
|
||||
api_key_id_resolver=get_current_api_key_id, # Propagate api_key_id for usage metering
|
||||
include_bank_id_param=multi_bank,
|
||||
tools=base_tools,
|
||||
)
|
||||
|
||||
register_mcp_tools(mcp, memory, config)
|
||||
|
||||
# Load and register additional tools from MCP extension if configured
|
||||
mcp_extension = load_extension("MCP", MCPExtension)
|
||||
if mcp_extension:
|
||||
logger.info(f"Loading MCP extension: {mcp_extension.__class__.__name__}")
|
||||
mcp_extension.register_tools(mcp, memory)
|
||||
|
||||
# Make all tools tolerant of extra arguments from LLMs (e.g., "explanation")
|
||||
_make_tools_tolerant(mcp)
|
||||
|
||||
return mcp
|
||||
|
||||
|
||||
def _make_tools_tolerant(mcp: FastMCP) -> None:
|
||||
"""Wrap all tool run methods to strip unknown arguments before validation.
|
||||
|
||||
LLMs frequently add extra fields like "explanation" or "reasoning" to tool calls.
|
||||
FastMCP's Pydantic TypeAdapter rejects these with "Unexpected keyword argument".
|
||||
This wraps each tool's run() to filter arguments to only known parameters.
|
||||
"""
|
||||
try:
|
||||
for name, tool in mcp._tool_manager._tools.items():
|
||||
if hasattr(tool, "parameters") and tool.parameters:
|
||||
allowed = set(tool.parameters.get("properties", {}).keys())
|
||||
original_run = tool.run
|
||||
|
||||
async def _tolerant_run(arguments, _allowed=allowed, _orig=original_run):
|
||||
extra_keys = set(arguments.keys()) - _allowed
|
||||
if extra_keys:
|
||||
logger.debug(f"Stripping unknown arguments from tool call: {extra_keys}")
|
||||
arguments = {k: v for k, v in arguments.items() if k in _allowed}
|
||||
return await _orig(arguments)
|
||||
|
||||
# FunctionTool is a Pydantic model with extra='forbid', so use
|
||||
# object.__setattr__ to bypass Pydantic's setter validation.
|
||||
object.__setattr__(tool, "run", _tolerant_run)
|
||||
except (AttributeError, KeyError) as e:
|
||||
logger.warning(f"Could not make tools tolerant of extra arguments: {e}")
|
||||
|
||||
|
||||
class MCPMiddleware:
|
||||
"""ASGI middleware that intercepts MCP requests and routes to appropriate MCP server.
|
||||
|
||||
This middleware wraps the main FastAPI app and intercepts requests matching the
|
||||
configured prefix (default: /mcp). Non-MCP requests pass through to the inner app.
|
||||
|
||||
Authentication:
|
||||
1. If HINDSIGHT_API_MCP_AUTH_TOKEN is set (legacy), validates against that token
|
||||
2. Otherwise, uses TenantExtension.authenticate_mcp() from the MemoryEngine
|
||||
- DefaultTenantExtension: no auth required (local dev)
|
||||
- ApiKeyTenantExtension: validates against env var
|
||||
|
||||
Two modes based on URL structure:
|
||||
|
||||
1. Multi-bank mode (for /mcp/ root endpoint):
|
||||
- Exposes all tools: retain, recall, reflect, list_banks, create_bank
|
||||
- All tools include optional bank_id parameter for cross-bank operations
|
||||
- Bank ID from: X-Bank-Id header or HINDSIGHT_MCP_BANK_ID env var
|
||||
|
||||
2. Single-bank mode (for /mcp/{bank_id}/ endpoints):
|
||||
- Exposes bank-scoped tools only: retain, recall, reflect
|
||||
- No bank_id parameter (comes from URL)
|
||||
- No bank management tools (list_banks, create_bank)
|
||||
- Recommended for agent isolation
|
||||
|
||||
Bank ID resolution priority:
|
||||
1. URL path (e.g., /mcp/{bank_id}/) → single-bank mode
|
||||
2. X-Bank-Id header → multi-bank mode
|
||||
3. HINDSIGHT_MCP_BANK_ID env var → multi-bank mode (default: "default")
|
||||
|
||||
Examples:
|
||||
# Single-bank mode (recommended for agent isolation)
|
||||
claude mcp add --transport http my-agent http://localhost:8888/mcp/my-agent-bank/ \\
|
||||
--header "Authorization: Bearer <token>"
|
||||
|
||||
# Multi-bank mode (for cross-bank operations)
|
||||
claude mcp add --transport http hindsight http://localhost:8888/mcp \\
|
||||
--header "X-Bank-Id: my-bank" --header "Authorization: Bearer <token>"
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
app,
|
||||
memory: MemoryEngine,
|
||||
prefix: str = "/mcp",
|
||||
multi_bank_app=None,
|
||||
single_bank_app=None,
|
||||
multi_bank_server=None,
|
||||
single_bank_server=None,
|
||||
):
|
||||
self.app = app
|
||||
self.prefix = prefix
|
||||
self.memory = memory
|
||||
self.tenant_extension = memory._tenant_extension
|
||||
|
||||
if multi_bank_app and single_bank_app:
|
||||
# Pre-created servers (used when called via add_middleware from create_app)
|
||||
self.multi_bank_app = multi_bank_app
|
||||
self.single_bank_app = single_bank_app
|
||||
self.multi_bank_server = multi_bank_server
|
||||
self.single_bank_server = single_bank_server
|
||||
else:
|
||||
# Create servers internally (for direct construction / tests)
|
||||
self.multi_bank_server = create_mcp_server(memory, multi_bank=True)
|
||||
self.multi_bank_app = self.multi_bank_server.http_app(path="/", stateless_http=True)
|
||||
self.single_bank_server = create_mcp_server(memory, multi_bank=False)
|
||||
self.single_bank_app = self.single_bank_server.http_app(path="/", stateless_http=True)
|
||||
|
||||
def _get_header(self, scope: dict, name: str) -> str | None:
|
||||
"""Extract a header value from ASGI scope."""
|
||||
name_lower = name.lower().encode()
|
||||
for header_name, header_value in scope.get("headers", []):
|
||||
if header_name.lower() == name_lower:
|
||||
return header_value.decode()
|
||||
return None
|
||||
|
||||
async def __call__(self, scope, receive, send):
|
||||
if scope["type"] != "http":
|
||||
await self.app(scope, receive, send)
|
||||
return
|
||||
|
||||
path = scope.get("path", "")
|
||||
|
||||
# Check if this is an MCP request (matches prefix)
|
||||
if not (path == self.prefix or path.startswith(self.prefix + "/")):
|
||||
# Not an MCP request — pass through to the inner app
|
||||
await self.app(scope, receive, send)
|
||||
return
|
||||
|
||||
# Strip prefix from path
|
||||
path = path[len(self.prefix) :] or "/"
|
||||
|
||||
# Extract auth token from header (for tenant auth propagation)
|
||||
auth_header = self._get_header(scope, "Authorization")
|
||||
auth_token: str | None = None
|
||||
if auth_header:
|
||||
# Support both "Bearer <token>" and direct token
|
||||
auth_token = auth_header[7:].strip() if auth_header.startswith("Bearer ") else auth_header.strip()
|
||||
|
||||
# Authenticate: check legacy MCP_AUTH_TOKEN first, then TenantExtension
|
||||
tenant_context = None
|
||||
auth_tenant_id: str | None = None
|
||||
auth_api_key_id: str | None = None
|
||||
if MCP_AUTH_TOKEN:
|
||||
# Legacy authentication mode - validate against static token
|
||||
if not auth_token:
|
||||
await self._send_error(send, 401, "Authorization header required")
|
||||
return
|
||||
if auth_token != MCP_AUTH_TOKEN:
|
||||
await self._send_error(send, 401, "Invalid authentication token")
|
||||
return
|
||||
# Legacy mode doesn't use tenant schemas
|
||||
tenant_context = None
|
||||
else:
|
||||
# Use TenantExtension.authenticate_mcp() for auth
|
||||
try:
|
||||
auth_context = RequestContext(api_key=auth_token)
|
||||
tenant_context = await self.tenant_extension.authenticate_mcp(auth_context)
|
||||
# Capture tenant_id and api_key_id set by authenticate() for usage metering
|
||||
auth_tenant_id = auth_context.tenant_id
|
||||
auth_api_key_id = auth_context.api_key_id
|
||||
except AuthenticationError as e:
|
||||
await self._send_error(send, 401, str(e), extra_headers=e.headers)
|
||||
return
|
||||
|
||||
# Set schema from tenant context so downstream DB queries use the correct schema
|
||||
schema_token = (
|
||||
_current_schema.set(tenant_context.schema_name) if tenant_context and tenant_context.schema_name else None
|
||||
)
|
||||
|
||||
# Resolve bank_id: path takes priority over header.
|
||||
# Path = user's explicit connection endpoint (e.g., /mcp/my-bank/).
|
||||
# X-Bank-Id header = per-request override for multi-bank mode only.
|
||||
bank_id = None
|
||||
bank_id_from_path = False
|
||||
new_path = path
|
||||
|
||||
# First, try to extract from path: /{bank_id}/...
|
||||
if path.startswith("/") and len(path) > 1:
|
||||
parts = path[1:].split("/", 1)
|
||||
if parts[0]:
|
||||
bank_id = parts[0]
|
||||
bank_id_from_path = True
|
||||
new_path = "/" + parts[1] if len(parts) > 1 else "/"
|
||||
|
||||
# If no path-based bank_id, try X-Bank-Id header (multi-bank mode)
|
||||
if not bank_id:
|
||||
bank_id = self._get_header(scope, "X-Bank-Id")
|
||||
|
||||
# Fall back to default bank_id
|
||||
if not bank_id:
|
||||
bank_id = DEFAULT_BANK_ID
|
||||
logger.debug(f"Using default bank_id: {bank_id}")
|
||||
|
||||
# Select the appropriate MCP app based on how bank_id was provided:
|
||||
# - Path-based bank_id → single-bank app (no bank_id param, scoped tools)
|
||||
# - Header/env bank_id → multi-bank app (bank_id param, all tools)
|
||||
target_app = self.single_bank_app if bank_id_from_path else self.multi_bank_app
|
||||
|
||||
# Set bank_id, api_key, tenant_id, and api_key_id context
|
||||
bank_id_token = _current_bank_id.set(bank_id)
|
||||
# Store the auth token for tenant extension to validate
|
||||
api_key_token = _current_api_key.set(auth_token) if auth_token else None
|
||||
# Store tenant_id and api_key_id from authentication for usage metering
|
||||
tenant_id_token = _current_tenant_id.set(auth_tenant_id) if auth_tenant_id else None
|
||||
api_key_id_token = _current_api_key_id.set(auth_api_key_id) if auth_api_key_id else None
|
||||
try:
|
||||
new_scope = scope.copy()
|
||||
new_scope["path"] = new_path
|
||||
# Clear root_path since we're passing directly to the app
|
||||
new_scope["root_path"] = ""
|
||||
|
||||
# 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".
|
||||
is_sse_response = False
|
||||
|
||||
async def send_wrapper(message):
|
||||
nonlocal is_sse_response
|
||||
if message["type"] == "http.response.start":
|
||||
for header_name, header_value in message.get("headers", []):
|
||||
if header_name == b"content-type" and b"text/event-stream" in header_value:
|
||||
is_sse_response = True
|
||||
break
|
||||
if message["type"] == "http.response.body" and bank_id_from_path and is_sse_response:
|
||||
body = message.get("body", b"")
|
||||
if body and b"/messages" in body:
|
||||
# Rewrite /messages to /{bank_id}/messages in SSE endpoint event
|
||||
body = body.replace(b"data: /messages", f"data: /{bank_id}/messages".encode())
|
||||
message = {**message, "body": body}
|
||||
await send(message)
|
||||
|
||||
await target_app(new_scope, receive, send_wrapper)
|
||||
finally:
|
||||
_current_bank_id.reset(bank_id_token)
|
||||
if api_key_token is not None:
|
||||
_current_api_key.reset(api_key_token)
|
||||
if tenant_id_token is not None:
|
||||
_current_tenant_id.reset(tenant_id_token)
|
||||
if api_key_id_token is not None:
|
||||
_current_api_key_id.reset(api_key_id_token)
|
||||
if schema_token is not None:
|
||||
_current_schema.reset(schema_token)
|
||||
|
||||
async def _send_error(self, send, status: int, message: str, extra_headers: dict[str, str] | None = None):
|
||||
"""Send an error response."""
|
||||
body = json.dumps({"error": message}).encode()
|
||||
headers = [(b"content-type", b"application/json")]
|
||||
for key, value in (extra_headers or {}).items():
|
||||
headers.append((key.encode(), value.encode()))
|
||||
await send(
|
||||
{
|
||||
"type": "http.response.start",
|
||||
"status": status,
|
||||
"headers": headers,
|
||||
}
|
||||
)
|
||||
await send(
|
||||
{
|
||||
"type": "http.response.body",
|
||||
"body": body,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def create_mcp_servers(memory: MemoryEngine):
|
||||
"""Create multi-bank and single-bank MCP servers and their Starlette apps.
|
||||
|
||||
Returns the servers and apps separately so lifespans can be chained before
|
||||
the middleware wraps the main app.
|
||||
|
||||
Returns:
|
||||
Tuple of (multi_bank_server, single_bank_server, multi_bank_app, single_bank_app)
|
||||
"""
|
||||
multi_bank_server = create_mcp_server(memory, multi_bank=True)
|
||||
multi_bank_app = multi_bank_server.http_app(path="/", stateless_http=True)
|
||||
|
||||
single_bank_server = create_mcp_server(memory, multi_bank=False)
|
||||
single_bank_app = single_bank_server.http_app(path="/", stateless_http=True)
|
||||
|
||||
return multi_bank_server, single_bank_server, multi_bank_app, single_bank_app
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,275 +0,0 @@
|
||||
"""
|
||||
Configuration resolution with hierarchical overrides.
|
||||
|
||||
Resolves config values through the hierarchy:
|
||||
Global (env vars) → Tenant config (via extension) → Bank config (database)
|
||||
|
||||
Config values are resolved on every request to ensure consistency across
|
||||
multiple API servers.
|
||||
"""
|
||||
|
||||
import json
|
||||
import logging
|
||||
from dataclasses import asdict
|
||||
from typing import Any
|
||||
|
||||
import asyncpg
|
||||
|
||||
from hindsight_api.config import HindsightConfig, _get_raw_config, normalize_config_dict
|
||||
from hindsight_api.engine.memory_engine import fq_table
|
||||
from hindsight_api.extensions.tenant import TenantExtension
|
||||
from hindsight_api.models import RequestContext
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ConfigResolver:
|
||||
"""Resolves hierarchical configuration with tenant/bank overrides."""
|
||||
|
||||
def __init__(self, pool: asyncpg.Pool, tenant_extension: TenantExtension | None = None):
|
||||
"""
|
||||
Initialize config resolver.
|
||||
|
||||
Args:
|
||||
pool: Database connection pool
|
||||
tenant_extension: Optional tenant extension for tenant-level config and permissions
|
||||
"""
|
||||
self.pool = pool
|
||||
self.tenant_extension = tenant_extension
|
||||
self._global_config = _get_raw_config()
|
||||
self._configurable_fields = HindsightConfig.get_configurable_fields()
|
||||
self._credential_fields = HindsightConfig.get_credential_fields()
|
||||
|
||||
async def resolve_full_config(self, bank_id: str, context: RequestContext | None = None) -> HindsightConfig:
|
||||
"""
|
||||
Resolve full HindsightConfig for a bank with hierarchical overrides applied.
|
||||
|
||||
This is for INTERNAL USE ONLY. Returns the complete config object with all fields
|
||||
including credentials and static fields. Use get_bank_config() for API responses.
|
||||
|
||||
Resolution order:
|
||||
1. Global config (from environment variables)
|
||||
2. Tenant config overrides (from TenantExtension.get_tenant_config())
|
||||
3. Bank config overrides (from banks.config JSONB)
|
||||
|
||||
Args:
|
||||
bank_id: Bank identifier
|
||||
context: Request context for tenant config resolution
|
||||
|
||||
Returns:
|
||||
Complete HindsightConfig with hierarchical overrides applied
|
||||
"""
|
||||
# Start with global config (all fields)
|
||||
config_dict = asdict(self._global_config)
|
||||
|
||||
# Load tenant config overrides (if tenant extension available)
|
||||
if self.tenant_extension and context:
|
||||
try:
|
||||
tenant_overrides = await self.tenant_extension.get_tenant_config(context)
|
||||
if tenant_overrides:
|
||||
# Normalize keys and filter to configurable fields only
|
||||
normalized_tenant = normalize_config_dict(tenant_overrides)
|
||||
configurable_tenant = {k: v for k, v in normalized_tenant.items() if k in self._configurable_fields}
|
||||
config_dict.update(configurable_tenant)
|
||||
logger.debug(
|
||||
f"Applied tenant config overrides for bank {bank_id}: {list(configurable_tenant.keys())}"
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to load tenant config for bank {bank_id}: {e}")
|
||||
|
||||
# Load bank config overrides
|
||||
bank_overrides = await self._load_bank_config(bank_id)
|
||||
if bank_overrides:
|
||||
config_dict.update(bank_overrides)
|
||||
logger.debug(f"Applied bank config overrides for bank {bank_id}: {list(bank_overrides.keys())}")
|
||||
|
||||
# Return full config object (dataclass doesn't have __init__ that accepts kwargs, so we update the object)
|
||||
# Create a new config instance by copying the global config and updating fields
|
||||
resolved_config = HindsightConfig(**config_dict)
|
||||
return resolved_config
|
||||
|
||||
async def get_bank_config(self, bank_id: str, context: RequestContext | None = None) -> dict[str, Any]:
|
||||
"""
|
||||
Get fully resolved config for a bank (filtered by permissions).
|
||||
|
||||
Resolution order:
|
||||
1. Global config (from environment variables)
|
||||
2. Tenant config overrides (from TenantExtension.get_tenant_config())
|
||||
3. Bank config overrides (from banks.config JSONB)
|
||||
|
||||
Note: Config is resolved on every call (not cached) to ensure consistency
|
||||
across multiple API servers.
|
||||
|
||||
SECURITY:
|
||||
- Only returns configurable fields (excludes static/infrastructure fields)
|
||||
- Filters out ALL credential fields (API keys, base URLs, etc.)
|
||||
- Further filtered by tenant/bank permissions if extension provides them
|
||||
|
||||
Args:
|
||||
bank_id: Bank identifier
|
||||
context: Request context for tenant config resolution and permissions
|
||||
|
||||
Returns:
|
||||
Dict of allowed configurable fields only (never includes credentials or static fields)
|
||||
"""
|
||||
# Resolve full config with all hierarchical overrides
|
||||
resolved_config = await self.resolve_full_config(bank_id, context)
|
||||
config_dict = asdict(resolved_config)
|
||||
|
||||
# SECURITY: Filter to only configurable fields (exclude static/infrastructure)
|
||||
filtered = {k: v for k, v in config_dict.items() if k in self._configurable_fields}
|
||||
|
||||
# SECURITY: Remove ALL credential fields (API keys, base URLs, etc.)
|
||||
filtered = {k: v for k, v in filtered.items() if k not in self._credential_fields}
|
||||
|
||||
# PERMISSIONS: Further filter based on tenant/bank permissions
|
||||
if self.tenant_extension and context:
|
||||
try:
|
||||
allowed_fields = await self.tenant_extension.get_allowed_config_fields(context, bank_id)
|
||||
if allowed_fields is not None: # None means "allow all"
|
||||
filtered = {k: v for k, v in filtered.items() if k in allowed_fields}
|
||||
logger.debug(
|
||||
f"Applied permission filter for bank {bank_id}: allowed={len(allowed_fields)} fields, "
|
||||
f"returned={len(filtered)} fields"
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to load permissions for bank {bank_id}: {e}")
|
||||
|
||||
return filtered
|
||||
|
||||
async def _load_bank_config(self, bank_id: str) -> dict[str, Any]:
|
||||
"""
|
||||
Load bank config overrides from banks.config JSONB column.
|
||||
|
||||
Args:
|
||||
bank_id: Bank identifier
|
||||
|
||||
Returns:
|
||||
Dict of config overrides (only configurable fields, normalized keys)
|
||||
"""
|
||||
try:
|
||||
async with self.pool.acquire() as conn:
|
||||
row = await conn.fetchrow(
|
||||
f"""
|
||||
SELECT config FROM {fq_table("banks")} WHERE bank_id = $1
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
|
||||
if row and row["config"]:
|
||||
config_data = row["config"]
|
||||
|
||||
# Handle case where JSONB is returned as JSON string
|
||||
if isinstance(config_data, str):
|
||||
config_data = json.loads(config_data)
|
||||
|
||||
# Normalize keys (handle both env var format and Python field format)
|
||||
normalized = normalize_config_dict(config_data)
|
||||
|
||||
# Only return overrides for configurable fields
|
||||
return {k: v for k, v in normalized.items() if k in self._configurable_fields}
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to load bank config for {bank_id}: {e}")
|
||||
|
||||
return {}
|
||||
|
||||
async def update_bank_config(
|
||||
self, bank_id: str, updates: dict[str, Any], context: RequestContext | None = None
|
||||
) -> None:
|
||||
"""
|
||||
Update bank configuration overrides (with permission checking).
|
||||
|
||||
Args:
|
||||
bank_id: Bank identifier
|
||||
updates: Dict of config field names to new values.
|
||||
Keys can be in env var format (HINDSIGHT_API_LLM_PROVIDER)
|
||||
or Python field format (llm_provider).
|
||||
Only configurable fields are allowed.
|
||||
context: Request context for permission checking
|
||||
|
||||
Raises:
|
||||
ValueError: If attempting to override invalid/disallowed fields
|
||||
"""
|
||||
# Normalize keys
|
||||
normalized_updates = normalize_config_dict(updates)
|
||||
|
||||
# SECURITY: Reject credential fields explicitly
|
||||
credential_attempts = set(normalized_updates.keys()) & self._credential_fields
|
||||
if credential_attempts:
|
||||
raise ValueError(
|
||||
f"Cannot set credential fields via API: {sorted(credential_attempts)}. "
|
||||
f"Credentials (API keys, base URLs) must be set at server level only."
|
||||
)
|
||||
|
||||
# Validate all fields are configurable
|
||||
invalid_fields = set(normalized_updates.keys()) - self._configurable_fields
|
||||
if invalid_fields:
|
||||
static_fields = HindsightConfig.get_static_fields()
|
||||
invalid_static = invalid_fields & static_fields
|
||||
if invalid_static:
|
||||
raise ValueError(
|
||||
f"Cannot override static (server-level) fields: {sorted(invalid_static)}. "
|
||||
f"Only configurable fields can be overridden per-bank. "
|
||||
f"Configurable fields include: {sorted(list(self._configurable_fields)[:10])}... "
|
||||
f"(total: {len(self._configurable_fields)} fields)"
|
||||
)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unknown configuration fields: {sorted(invalid_fields)}. "
|
||||
f"Valid configurable fields: {sorted(list(self._configurable_fields)[:10])}..."
|
||||
)
|
||||
|
||||
# PERMISSIONS: Check tenant/bank permissions
|
||||
if self.tenant_extension and context:
|
||||
try:
|
||||
allowed_fields = await self.tenant_extension.get_allowed_config_fields(context, bank_id)
|
||||
if allowed_fields is not None: # None means "allow all"
|
||||
disallowed = set(normalized_updates.keys()) - allowed_fields
|
||||
if disallowed:
|
||||
raise ValueError(
|
||||
f"Not allowed to modify fields: {sorted(disallowed)}. "
|
||||
f"Your permissions allow: {sorted(list(allowed_fields)[:10])}..."
|
||||
if allowed_fields
|
||||
else "Not allowed to modify fields: {sorted(disallowed)}. "
|
||||
"Your permissions do not allow any config modifications."
|
||||
)
|
||||
except ValueError:
|
||||
raise # Re-raise permission errors
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to check permissions for bank {bank_id}: {e}")
|
||||
# Continue without permission check (fail open for backward compatibility)
|
||||
|
||||
# Merge with existing config (JSONB || operator)
|
||||
async with self.pool.acquire() as conn:
|
||||
await conn.execute(
|
||||
f"""
|
||||
UPDATE {fq_table("banks")}
|
||||
SET config = config || $1::jsonb,
|
||||
updated_at = now()
|
||||
WHERE bank_id = $2
|
||||
""",
|
||||
json.dumps(normalized_updates),
|
||||
bank_id,
|
||||
)
|
||||
|
||||
logger.info(f"Updated bank config for {bank_id}: {list(normalized_updates.keys())}")
|
||||
|
||||
async def reset_bank_config(self, bank_id: str) -> None:
|
||||
"""
|
||||
Reset bank configuration to defaults (remove all overrides).
|
||||
|
||||
Args:
|
||||
bank_id: Bank identifier
|
||||
"""
|
||||
async with self.pool.acquire() as conn:
|
||||
await conn.execute(
|
||||
f"""
|
||||
UPDATE {fq_table("banks")}
|
||||
SET config = '{{}}'::jsonb,
|
||||
updated_at = now()
|
||||
WHERE bank_id = $1
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
|
||||
logger.info(f"Reset bank config for {bank_id} to defaults")
|
||||
@@ -1,113 +0,0 @@
|
||||
"""
|
||||
Daemon mode support for Hindsight API.
|
||||
|
||||
Provides idle timeout for running as a background daemon.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Default daemon configuration
|
||||
DEFAULT_DAEMON_PORT = 8888
|
||||
DEFAULT_IDLE_TIMEOUT = 0 # 0 = no auto-exit (hindsight-embed passes its own timeout)
|
||||
|
||||
# Allow override via environment variable for profile-specific logs
|
||||
DAEMON_LOG_PATH = Path(os.getenv("HINDSIGHT_API_DAEMON_LOG", str(Path.home() / ".hindsight" / "daemon.log")))
|
||||
|
||||
|
||||
class IdleTimeoutMiddleware:
|
||||
"""ASGI middleware that tracks activity and exits after idle timeout."""
|
||||
|
||||
def __init__(self, app, idle_timeout: int = DEFAULT_IDLE_TIMEOUT):
|
||||
self.app = app
|
||||
self.idle_timeout = idle_timeout
|
||||
self.last_activity = time.time()
|
||||
self._checker_task = None
|
||||
|
||||
async def __call__(self, scope, receive, send):
|
||||
# Update activity timestamp on each request
|
||||
self.last_activity = time.time()
|
||||
await self.app(scope, receive, send)
|
||||
|
||||
def start_idle_checker(self):
|
||||
"""Start the background task that checks for idle timeout."""
|
||||
self._checker_task = asyncio.create_task(self._check_idle())
|
||||
|
||||
async def _check_idle(self):
|
||||
"""Background task that exits the process after idle timeout."""
|
||||
# If idle_timeout is 0, don't auto-exit
|
||||
if self.idle_timeout <= 0:
|
||||
return
|
||||
|
||||
while True:
|
||||
await asyncio.sleep(30) # Check every 30 seconds
|
||||
idle_time = time.time() - self.last_activity
|
||||
if idle_time > self.idle_timeout:
|
||||
logger.info(f"Idle timeout reached ({self.idle_timeout}s), shutting down daemon")
|
||||
# Give a moment for any in-flight requests
|
||||
await asyncio.sleep(1)
|
||||
# Send SIGTERM to ourselves to trigger graceful shutdown
|
||||
import signal
|
||||
|
||||
os.kill(os.getpid(), signal.SIGTERM)
|
||||
|
||||
|
||||
def daemonize():
|
||||
"""
|
||||
Fork the current process into a background daemon.
|
||||
|
||||
Uses double-fork technique to properly detach from terminal.
|
||||
"""
|
||||
# First fork - detach from parent
|
||||
try:
|
||||
pid = os.fork()
|
||||
if pid > 0:
|
||||
sys.exit(0)
|
||||
except OSError as e:
|
||||
sys.stderr.write(f"fork #1 failed: {e}\n")
|
||||
sys.exit(1)
|
||||
|
||||
# Decouple from parent environment
|
||||
os.chdir("/")
|
||||
os.setsid()
|
||||
os.umask(0)
|
||||
|
||||
# Second fork - prevent zombie
|
||||
pid = os.fork()
|
||||
if pid > 0:
|
||||
sys.exit(0)
|
||||
|
||||
# Redirect standard file descriptors to log file
|
||||
DAEMON_LOG_PATH.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
sys.stdout.flush()
|
||||
sys.stderr.flush()
|
||||
|
||||
# Redirect stdin to /dev/null
|
||||
with open("/dev/null", "r") as devnull:
|
||||
os.dup2(devnull.fileno(), sys.stdin.fileno())
|
||||
|
||||
# Redirect stdout/stderr to log file
|
||||
log_fd = open(DAEMON_LOG_PATH, "a")
|
||||
os.dup2(log_fd.fileno(), sys.stdout.fileno())
|
||||
os.dup2(log_fd.fileno(), sys.stderr.fileno())
|
||||
|
||||
|
||||
def check_daemon_running(port: int = DEFAULT_DAEMON_PORT) -> bool:
|
||||
"""Check if a daemon is running and responsive on the given port."""
|
||||
import socket
|
||||
|
||||
try:
|
||||
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
||||
sock.settimeout(1)
|
||||
result = sock.connect_ex(("127.0.0.1", port))
|
||||
sock.close()
|
||||
return result == 0
|
||||
except Exception:
|
||||
return False
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,83 +0,0 @@
|
||||
"""Prompts for the consolidation engine."""
|
||||
|
||||
# Default mission when no bank-specific mission is set
|
||||
_DEFAULT_MISSION = "Track every detail: names, numbers, dates, places, and relationships. Prefer specifics over abstractions, never generalise."
|
||||
|
||||
# Processing rules — always present regardless of mission
|
||||
_PROCESSING_RULES = """Processing rules (always apply):
|
||||
- REDUNDANT: same info worded differently → UPDATE the existing observation.
|
||||
- CONTRADICTION/UPDATE: capture both states with temporal markers ("used to X, now Y").
|
||||
- RESOLVE REFERENCES: when a new fact provides a concrete value resolving a vague placeholder in an existing observation (e.g. "home country", "hometown", "birthplace", "native language", "her ex", "that city"), UPDATE the observation to embed the resolved value explicitly. Example: new fact says "grandma in Sweden" + existing observation says "moved from her home country" → update to "home country is Sweden".
|
||||
- NEVER merge observations about different people or unrelated topics."""
|
||||
|
||||
# Data section — format placeholders {facts_text} and {observations_text} are substituted at call time
|
||||
_BATCH_DATA_SECTION = """
|
||||
NEW FACTS:
|
||||
{facts_text}
|
||||
|
||||
EXISTING OBSERVATIONS (JSON array, pooled from recalls across all facts above):
|
||||
{observations_text}
|
||||
|
||||
Each observation includes:
|
||||
- id: unique identifier for updating
|
||||
- text: the observation content
|
||||
- proof_count: number of supporting memories
|
||||
- occurred_start/occurred_end: temporal range of source facts
|
||||
- source_memories: array of supporting facts with their text and dates
|
||||
|
||||
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 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
|
||||
|
||||
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"}}]}}
|
||||
|
||||
Rules:
|
||||
- "source_fact_ids": copy the EXACT UUID strings shown in brackets [uuid] from NEW FACTS — never use integers or positions.
|
||||
- "observation_id": copy the EXACT "id" UUID string from EXISTING OBSERVATIONS.
|
||||
- One create/update may reference multiple facts when they jointly support the observation.
|
||||
- "deletes": only when an observation is directly superseded or contradicted by new facts.
|
||||
- Do NOT include "tags" — handled automatically.
|
||||
- Return {{"creates": [], "updates": [], "deletes": []}} if nothing durable is found."""
|
||||
|
||||
|
||||
def build_batch_consolidation_prompt(observations_mission: str | None = None) -> str:
|
||||
"""
|
||||
Build the consolidation prompt for batch mode (multiple facts per LLM call).
|
||||
|
||||
The mission defines *what* to track (customisable per bank).
|
||||
Processing rules and output format are always present regardless of mission.
|
||||
"""
|
||||
mission = observations_mission or _DEFAULT_MISSION
|
||||
|
||||
return (
|
||||
"You are a memory consolidation system. Synthesize facts into observations "
|
||||
"and merge with existing observations when appropriate.\n\n"
|
||||
f"## MISSION\n{mission}\n\n"
|
||||
f"{_PROCESSING_RULES}" + _BATCH_DATA_SECTION + _BATCH_OUTPUT_FORMAT
|
||||
)
|
||||
@@ -1,144 +0,0 @@
|
||||
"""
|
||||
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)
|
||||
]
|
||||
@@ -1,69 +0,0 @@
|
||||
"""
|
||||
Typed metadata models for async operations.
|
||||
|
||||
These dataclasses define the structure of result_metadata for different operation types.
|
||||
The metadata is exposed in the API for debugging purposes and may change without notice.
|
||||
"""
|
||||
|
||||
from dataclasses import asdict, dataclass
|
||||
from typing import Any
|
||||
|
||||
|
||||
@dataclass
|
||||
class BatchRetainParentMetadata:
|
||||
"""Metadata for parent batch_retain operations (when split into sub-batches)."""
|
||||
|
||||
items_count: int
|
||||
total_tokens: int
|
||||
num_sub_batches: int
|
||||
is_parent: bool = True
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
"""Convert to dict for JSON serialization."""
|
||||
return asdict(self)
|
||||
|
||||
|
||||
@dataclass
|
||||
class BatchRetainChildMetadata:
|
||||
"""Metadata for child batch_retain operations (individual sub-batches)."""
|
||||
|
||||
items_count: int
|
||||
parent_operation_id: str
|
||||
sub_batch_index: int
|
||||
total_sub_batches: int
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
"""Convert to dict for JSON serialization."""
|
||||
return asdict(self)
|
||||
|
||||
|
||||
@dataclass
|
||||
class RetainMetadata:
|
||||
"""Metadata for regular retain operations (non-batched, deprecated async path)."""
|
||||
|
||||
items_count: int
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
"""Convert to dict for JSON serialization."""
|
||||
return asdict(self)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ConsolidationMetadata:
|
||||
"""Metadata for consolidation operations."""
|
||||
|
||||
# Currently empty, but structure for future fields
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
"""Convert to dict for JSON serialization."""
|
||||
return asdict(self)
|
||||
|
||||
|
||||
@dataclass
|
||||
class RefreshMentalModelMetadata:
|
||||
"""Metadata for mental model refresh operations."""
|
||||
|
||||
mental_model_id: str
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
"""Convert to dict for JSON serialization."""
|
||||
return asdict(self)
|
||||
@@ -1,128 +0,0 @@
|
||||
"""File parser implementations."""
|
||||
|
||||
import logging
|
||||
from dataclasses import dataclass
|
||||
|
||||
from .base import FileParser, UnsupportedFileTypeError
|
||||
from .iris import IrisParser
|
||||
from .markitdown import MarkitdownParser
|
||||
|
||||
__all__ = [
|
||||
"FileParser",
|
||||
"UnsupportedFileTypeError",
|
||||
"IrisParser",
|
||||
"MarkitdownParser",
|
||||
"FileParserRegistry",
|
||||
"ConvertResult",
|
||||
]
|
||||
|
||||
|
||||
@dataclass
|
||||
class ConvertResult:
|
||||
"""Result of a successful file conversion."""
|
||||
|
||||
content: str
|
||||
parser_name: str
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class FileParserRegistry:
|
||||
"""Registry for file parsers with auto-detection."""
|
||||
|
||||
def __init__(self):
|
||||
"""Initialize empty parser registry."""
|
||||
self._parsers: dict[str, FileParser] = {}
|
||||
|
||||
def register(self, parser: FileParser):
|
||||
"""
|
||||
Register a parser.
|
||||
|
||||
Args:
|
||||
parser: FileParser instance
|
||||
"""
|
||||
self._parsers[parser.name()] = parser
|
||||
|
||||
def get_parser(
|
||||
self,
|
||||
name: str | None,
|
||||
filename: str,
|
||||
content_type: str | None = None,
|
||||
) -> FileParser:
|
||||
"""
|
||||
Get parser by name or auto-detect.
|
||||
|
||||
Args:
|
||||
name: Parser name (e.g., "markitdown") or None for auto-detect
|
||||
filename: File name for auto-detection
|
||||
content_type: MIME type (optional)
|
||||
|
||||
Returns:
|
||||
FileParser instance
|
||||
|
||||
Raises:
|
||||
ValueError: If no suitable parser found
|
||||
"""
|
||||
if name:
|
||||
# Explicit parser requested — return it directly, let the parser
|
||||
# raise UnsupportedFileTypeError from convert() if needed
|
||||
if name not in self._parsers:
|
||||
raise ValueError(f"Parser '{name}' not found. Available: {list(self._parsers.keys())}")
|
||||
return self._parsers[name]
|
||||
|
||||
# Auto-detect parser
|
||||
for parser in self._parsers.values():
|
||||
if parser.supports(filename, content_type):
|
||||
return parser
|
||||
|
||||
raise ValueError(f"No parser found for {filename}. Available parsers: {list(self._parsers.keys())}")
|
||||
|
||||
async def convert_with_fallback(
|
||||
self,
|
||||
parsers: list[str],
|
||||
file_data: bytes,
|
||||
filename: str,
|
||||
content_type: str | None = None,
|
||||
) -> ConvertResult:
|
||||
"""
|
||||
Try each parser in order, falling back on failure or empty content.
|
||||
|
||||
Moves to the next parser if the current one raises UnsupportedFileTypeError
|
||||
or returns empty content. Any other exception (RuntimeError, network error,
|
||||
etc.) also triggers a fallback so the chain is exhausted before failing.
|
||||
|
||||
Args:
|
||||
parsers: Ordered list of parser names to try
|
||||
file_data: Raw file bytes
|
||||
filename: Original filename
|
||||
content_type: MIME type (optional)
|
||||
|
||||
Returns:
|
||||
ConvertResult with the parsed content and the name of the parser that succeeded
|
||||
|
||||
Raises:
|
||||
ValueError: If a parser name is not registered
|
||||
RuntimeError: If all parsers fail or return empty content
|
||||
"""
|
||||
last_error: Exception | None = None
|
||||
for name in parsers:
|
||||
parser = self.get_parser(name, filename, content_type)
|
||||
try:
|
||||
content = await parser.convert(file_data, filename)
|
||||
if content and content.strip():
|
||||
return ConvertResult(content=content, parser_name=name)
|
||||
logger.warning(f"Parser '{name}' returned empty content for '{filename}', trying next")
|
||||
last_error = RuntimeError(f"Parser '{name}' returned no content for '{filename}'")
|
||||
except UnsupportedFileTypeError as e:
|
||||
logger.warning(f"Parser '{name}' does not support '{filename}', trying next: {e}")
|
||||
last_error = e
|
||||
except Exception as e:
|
||||
logger.warning(f"Parser '{name}' failed for '{filename}', trying next: {e}")
|
||||
last_error = e
|
||||
|
||||
raise last_error or RuntimeError(f"No parsers available for '{filename}'")
|
||||
|
||||
def list_parsers(self) -> list[str]:
|
||||
"""Get list of registered parser names."""
|
||||
return list(self._parsers.keys())
|
||||
@@ -1,58 +0,0 @@
|
||||
"""Abstract base class for file parsers."""
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
|
||||
|
||||
class UnsupportedFileTypeError(Exception):
|
||||
"""Raised by a parser when it does not support the given file type."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class FileParser(ABC):
|
||||
"""Abstract base for file to markdown parsers."""
|
||||
|
||||
@abstractmethod
|
||||
async def convert(self, file_data: bytes, filename: str) -> str:
|
||||
"""
|
||||
Parse file to markdown.
|
||||
|
||||
Args:
|
||||
file_data: Raw file bytes
|
||||
filename: Original filename (used for format detection)
|
||||
|
||||
Returns:
|
||||
Markdown content as string
|
||||
|
||||
Raises:
|
||||
UnsupportedFileTypeError: If the file type is not supported by this parser
|
||||
RuntimeError: If parsing fails for another reason
|
||||
"""
|
||||
pass
|
||||
|
||||
def supports(self, filename: str, content_type: str | None = None) -> bool:
|
||||
"""
|
||||
Check if parser supports this file type.
|
||||
|
||||
Override this for local/static extension-based filtering.
|
||||
Parsers that delegate to a remote service should leave this as True
|
||||
and raise UnsupportedFileTypeError from convert() instead.
|
||||
|
||||
Args:
|
||||
filename: File name (used for extension check)
|
||||
content_type: MIME type (optional)
|
||||
|
||||
Returns:
|
||||
True if this parser can handle the file (default: True)
|
||||
"""
|
||||
return True
|
||||
|
||||
@abstractmethod
|
||||
def name(self) -> str:
|
||||
"""
|
||||
Get parser name.
|
||||
|
||||
Returns:
|
||||
Parser name (e.g., "markitdown")
|
||||
"""
|
||||
pass
|
||||
@@ -1,138 +0,0 @@
|
||||
"""Iris parser implementation using the Vectorize Iris HTTP API."""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import mimetypes
|
||||
import time
|
||||
|
||||
import httpx
|
||||
|
||||
from .base import FileParser, UnsupportedFileTypeError
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_IRIS_BASE_URL = "https://api.vectorize.io/v1"
|
||||
_DEFAULT_POLL_INTERVAL = 2.0 # seconds
|
||||
_DEFAULT_TIMEOUT = 300.0 # seconds
|
||||
|
||||
|
||||
class IrisParser(FileParser):
|
||||
"""
|
||||
Iris file parser using the Vectorize Iris cloud extraction service.
|
||||
|
||||
Uploads files to the Vectorize Iris API, starts an extraction job,
|
||||
and polls until the text is ready. The API determines which file types
|
||||
are supported — UnsupportedFileTypeError is raised if the file is rejected.
|
||||
|
||||
Authentication:
|
||||
Requires HINDSIGHT_API_FILE_PARSER_IRIS_TOKEN and
|
||||
HINDSIGHT_API_FILE_PARSER_IRIS_ORG_ID environment variables,
|
||||
or pass them explicitly via the constructor.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
token: str,
|
||||
org_id: str,
|
||||
poll_interval: float = _DEFAULT_POLL_INTERVAL,
|
||||
timeout: float = _DEFAULT_TIMEOUT,
|
||||
):
|
||||
"""
|
||||
Initialize iris parser.
|
||||
|
||||
Args:
|
||||
token: Vectorize API token
|
||||
org_id: Vectorize organization ID
|
||||
poll_interval: Seconds between status poll requests (default: 2)
|
||||
timeout: Maximum seconds to wait for extraction (default: 300)
|
||||
"""
|
||||
self._token = token
|
||||
self._org_id = org_id
|
||||
self._poll_interval = poll_interval
|
||||
self._timeout = timeout
|
||||
self._auth_headers = {"Authorization": f"Bearer {token}"}
|
||||
|
||||
async def convert(self, file_data: bytes, filename: str) -> str:
|
||||
"""
|
||||
Parse file to text using the Vectorize Iris API.
|
||||
|
||||
Raises:
|
||||
UnsupportedFileTypeError: If the Iris API rejects the file type (4xx)
|
||||
RuntimeError: If extraction fails for another reason
|
||||
"""
|
||||
content_type = mimetypes.guess_type(filename)[0] or "application/octet-stream"
|
||||
|
||||
async with httpx.AsyncClient(timeout=httpx.Timeout(30.0, read=120.0)) as client:
|
||||
# Step 1: Request a presigned upload URL
|
||||
init_resp = await client.post(
|
||||
f"{_IRIS_BASE_URL}/org/{self._org_id}/files",
|
||||
headers=self._auth_headers,
|
||||
json={"name": filename, "contentType": content_type},
|
||||
)
|
||||
_raise_for_status(init_resp, filename, "file upload init")
|
||||
init_data = init_resp.json()
|
||||
file_id: str = init_data["fileId"]
|
||||
upload_url: str = init_data["uploadUrl"]
|
||||
|
||||
# Step 2: Upload the file bytes to the presigned URL (no auth header)
|
||||
# Ensure file_data is plain bytes (GCS storage may return obstore.Bytes)
|
||||
upload_resp = await client.put(
|
||||
upload_url,
|
||||
content=bytes(file_data),
|
||||
headers={"Content-Type": content_type},
|
||||
)
|
||||
_raise_for_status(upload_resp, filename, "file upload")
|
||||
|
||||
# Step 3: Start extraction
|
||||
extract_resp = await client.post(
|
||||
f"{_IRIS_BASE_URL}/org/{self._org_id}/extraction",
|
||||
headers=self._auth_headers,
|
||||
json={"fileId": file_id},
|
||||
)
|
||||
_raise_for_status(extract_resp, filename, "start extraction")
|
||||
extraction_id: str = extract_resp.json()["extractionId"]
|
||||
|
||||
# Step 4: Poll until ready or timeout
|
||||
deadline = time.monotonic() + self._timeout
|
||||
while True:
|
||||
status_resp = await client.get(
|
||||
f"{_IRIS_BASE_URL}/org/{self._org_id}/extraction/{extraction_id}",
|
||||
headers=self._auth_headers,
|
||||
)
|
||||
_raise_for_status(status_resp, filename, "poll extraction status")
|
||||
status_data = status_resp.json()
|
||||
|
||||
if status_data.get("ready"):
|
||||
data = status_data.get("data", {})
|
||||
if not data.get("success"):
|
||||
error = data.get("error", "unknown error")
|
||||
raise RuntimeError(f"Iris extraction failed for '{filename}': {error}")
|
||||
text = data.get("text")
|
||||
if not text:
|
||||
raise RuntimeError(f"No content extracted from '{filename}'")
|
||||
return text
|
||||
|
||||
if time.monotonic() >= deadline:
|
||||
raise RuntimeError(f"Iris extraction timed out after {self._timeout}s for '{filename}'")
|
||||
|
||||
await asyncio.sleep(self._poll_interval)
|
||||
|
||||
def name(self) -> str:
|
||||
"""Get parser name."""
|
||||
return "iris"
|
||||
|
||||
|
||||
def _raise_for_status(response: httpx.Response, filename: str, step: str) -> None:
|
||||
"""
|
||||
Raise an appropriate error including the response body on HTTP errors.
|
||||
|
||||
Raises UnsupportedFileTypeError for 4xx responses (file rejected by the API),
|
||||
RuntimeError for other HTTP errors.
|
||||
"""
|
||||
if not response.is_error:
|
||||
return
|
||||
body = response.text or "<empty>"
|
||||
msg = f"Iris API error during {step} for '{filename}': {response.status_code} {response.reason_phrase} — {body}"
|
||||
if response.is_client_error:
|
||||
raise UnsupportedFileTypeError(msg)
|
||||
raise RuntimeError(msg)
|
||||
@@ -1,109 +0,0 @@
|
||||
"""Markitdown parser implementation."""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
|
||||
from .base import FileParser
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class MarkitdownParser(FileParser):
|
||||
"""
|
||||
Markitdown file parser.
|
||||
|
||||
Uses Microsoft's markitdown library to convert various file formats
|
||||
to markdown including PDF, Office docs, images (via OCR), audio, HTML.
|
||||
|
||||
Supported formats:
|
||||
- PDF (.pdf)
|
||||
- Word (.docx, .doc)
|
||||
- PowerPoint (.pptx, .ppt)
|
||||
- Excel (.xlsx, .xls)
|
||||
- Images (.jpg, .jpeg, .png) - with OCR
|
||||
- HTML (.html, .htm)
|
||||
- Text (.txt, .md)
|
||||
- Audio (.mp3, .wav) - with transcription
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
"""Initialize markitdown parser."""
|
||||
# Lazy import to avoid requiring markitdown for all users
|
||||
try:
|
||||
from markitdown import MarkItDown
|
||||
|
||||
self._markitdown = MarkItDown()
|
||||
except ImportError as e:
|
||||
raise ImportError(
|
||||
"markitdown package is required for file parsing. Install with: pip install markitdown"
|
||||
) from e
|
||||
|
||||
async def convert(self, file_data: bytes, filename: str) -> str:
|
||||
"""Parse file to markdown using markitdown."""
|
||||
# markitdown is synchronous, so we run it in executor to avoid blocking
|
||||
loop = asyncio.get_event_loop()
|
||||
return await loop.run_in_executor(None, self._convert_sync, file_data, filename)
|
||||
|
||||
def _convert_sync(self, file_data: bytes, filename: str) -> str:
|
||||
"""Synchronous parsing (runs in thread pool)."""
|
||||
# Write to temp file (markitdown requires file path)
|
||||
with tempfile.NamedTemporaryFile(suffix=Path(filename).suffix, delete=False) as tmp:
|
||||
tmp.write(file_data)
|
||||
tmp_path = tmp.name
|
||||
|
||||
try:
|
||||
# Parse using markitdown
|
||||
result = self._markitdown.convert(tmp_path)
|
||||
|
||||
if not result or not result.text_content:
|
||||
raise RuntimeError(f"No content extracted from '{filename}'")
|
||||
|
||||
return result.text_content
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Markitdown parsing failed for {filename}: {e}")
|
||||
raise RuntimeError(f"Failed to parse '{filename}': {e}") from e
|
||||
|
||||
finally:
|
||||
# Clean up temp file
|
||||
try:
|
||||
Path(tmp_path).unlink()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def supports(self, filename: str, content_type: str | None = None) -> bool:
|
||||
"""Check if markitdown supports this file type."""
|
||||
# Supported extensions (from markitdown docs)
|
||||
supported_extensions = {
|
||||
# Documents
|
||||
".pdf",
|
||||
".docx",
|
||||
".doc",
|
||||
".pptx",
|
||||
".ppt",
|
||||
".xlsx",
|
||||
".xls",
|
||||
# Images (with OCR)
|
||||
".jpg",
|
||||
".jpeg",
|
||||
".png",
|
||||
# Web
|
||||
".html",
|
||||
".htm",
|
||||
# Text
|
||||
".txt",
|
||||
".md",
|
||||
".csv",
|
||||
# Audio (with transcription)
|
||||
".mp3",
|
||||
".wav",
|
||||
}
|
||||
|
||||
ext = Path(filename).suffix.lower()
|
||||
return ext in supported_extensions
|
||||
|
||||
def name(self) -> str:
|
||||
"""Get parser name."""
|
||||
return "markitdown"
|
||||
@@ -1,194 +0,0 @@
|
||||
"""
|
||||
Entity labels models and helpers for retain pipeline.
|
||||
|
||||
Defines a controlled vocabulary of key:value classification labels
|
||||
(e.g., 'pedagogy:scaffolding', 'interest:active') that are extracted
|
||||
at retain time and stored as entities.
|
||||
"""
|
||||
|
||||
from typing import Literal
|
||||
|
||||
from pydantic import BaseModel, Field, create_model
|
||||
|
||||
|
||||
class LabelValue(BaseModel):
|
||||
"""A single allowed value for a label group."""
|
||||
|
||||
value: str
|
||||
description: str = ""
|
||||
|
||||
|
||||
class LabelGroup(BaseModel):
|
||||
"""A label group (dimension) with its type and allowed values."""
|
||||
|
||||
key: str
|
||||
description: str = ""
|
||||
type: Literal["value", "multi-values", "text"] = "value"
|
||||
optional: bool = True
|
||||
tag: bool = False
|
||||
values: list[LabelValue] = []
|
||||
|
||||
|
||||
class EntityLabelsConfig(BaseModel):
|
||||
"""Entity labels configuration for a bank (controlled vocabulary)."""
|
||||
|
||||
attributes: list[LabelGroup] = []
|
||||
|
||||
|
||||
def parse_entity_labels(raw: dict | list | None) -> EntityLabelsConfig | None:
|
||||
"""
|
||||
Parse raw entity labels config into EntityLabelsConfig.
|
||||
|
||||
Accepts:
|
||||
- None → returns None
|
||||
- list → list of attribute dicts (each may use legacy free_values/multi_value or new type field)
|
||||
- dict → {attributes: [...]}
|
||||
|
||||
Legacy migration (backward-compat):
|
||||
- free_values=True → type="text"
|
||||
- multi_value=True → type="multi-values"
|
||||
- neither / free_values=False → type="value"
|
||||
|
||||
Args:
|
||||
raw: Raw entity labels config from bank config
|
||||
|
||||
Returns:
|
||||
EntityLabelsConfig or None if raw is None/empty
|
||||
"""
|
||||
if raw is None:
|
||||
return None
|
||||
|
||||
if isinstance(raw, list):
|
||||
if not raw:
|
||||
return None
|
||||
attributes = [LabelGroup.model_validate(_migrate_label_group(a)) for a in raw]
|
||||
return EntityLabelsConfig(attributes=attributes)
|
||||
|
||||
if isinstance(raw, dict):
|
||||
attrs_raw = raw.get("attributes", [])
|
||||
if not attrs_raw:
|
||||
return None
|
||||
attributes = [LabelGroup.model_validate(_migrate_label_group(a)) for a in attrs_raw]
|
||||
return EntityLabelsConfig(attributes=attributes)
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def _migrate_label_group(raw: dict) -> dict:
|
||||
"""Migrate legacy free_values/multi_value fields to the new type field."""
|
||||
if not isinstance(raw, dict) or "type" in raw:
|
||||
return raw
|
||||
patched = dict(raw)
|
||||
if patched.get("free_values"):
|
||||
patched["type"] = "text"
|
||||
elif patched.get("multi_value"):
|
||||
patched["type"] = "multi-values"
|
||||
else:
|
||||
patched["type"] = "value"
|
||||
# Remove legacy keys so Pydantic doesn't error on unknown fields
|
||||
patched.pop("free_values", None)
|
||||
patched.pop("multi_value", None)
|
||||
return patched
|
||||
|
||||
|
||||
def build_labels_model(labels_cfg: EntityLabelsConfig) -> type[BaseModel] | None:
|
||||
"""
|
||||
Build a dynamic Pydantic model for structured label extraction.
|
||||
|
||||
Each LabelGroup becomes a typed field based on its type:
|
||||
- type="text" → str | None (always optional)
|
||||
- type="value", optional=True → Literal["v1","v2"] | None
|
||||
- type="value", optional=False → Literal["v1","v2"] (required)
|
||||
- type="multi-values" → list[Literal["v1","v2"]]
|
||||
|
||||
Args:
|
||||
labels_cfg: Parsed EntityLabelsConfig
|
||||
|
||||
Returns:
|
||||
Dynamic Pydantic model class, or None if no groups defined
|
||||
"""
|
||||
fields: dict = {}
|
||||
for group in labels_cfg.attributes:
|
||||
if not group.key:
|
||||
continue
|
||||
description = group.description or group.key
|
||||
|
||||
if group.type == "text":
|
||||
# Free-form: any string value accepted, always optional
|
||||
fields[group.key] = (str | None, Field(default=None, description=description))
|
||||
else:
|
||||
# Enum-constrained: must have defined values
|
||||
if not group.values:
|
||||
continue
|
||||
values = tuple(v.value for v in group.values if v.value)
|
||||
if not values:
|
||||
continue
|
||||
# Literal[("v1", "v2")] is equivalent to Literal["v1", "v2"] in Python 3.11+
|
||||
literal_type = Literal[values] # type: ignore[valid-type]
|
||||
if group.type == "multi-values":
|
||||
fields[group.key] = (
|
||||
list[literal_type], # type: ignore[valid-type]
|
||||
Field(default_factory=list, description=description),
|
||||
)
|
||||
elif group.optional:
|
||||
fields[group.key] = (
|
||||
literal_type | None, # type: ignore[valid-type]
|
||||
Field(default=None, description=description),
|
||||
)
|
||||
else:
|
||||
fields[group.key] = (
|
||||
literal_type, # type: ignore[valid-type]
|
||||
Field(description=description),
|
||||
)
|
||||
|
||||
if not fields:
|
||||
return None
|
||||
|
||||
return create_model("Labels", **fields)
|
||||
|
||||
|
||||
def is_label_entity(text: str, labels_cfg: EntityLabelsConfig, labels_lookup: set[str]) -> bool:
|
||||
"""
|
||||
Return True if entity text belongs to any configured label group.
|
||||
|
||||
For enum groups: checks the pre-built lookup set.
|
||||
For text groups: checks that the text starts with a known key prefix.
|
||||
"""
|
||||
if text.lower() in labels_lookup:
|
||||
return True
|
||||
for group in labels_cfg.attributes:
|
||||
if group.type == "text" and group.key and text.lower().startswith(f"{group.key.lower()}:"):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def build_labels_lookup(labels_cfg: EntityLabelsConfig | list | None) -> set[str]:
|
||||
"""
|
||||
Build a set of valid 'key:value' label strings (lowercase) for fast lookup.
|
||||
|
||||
Accepts either EntityLabelsConfig or raw list/None for backwards compatibility.
|
||||
|
||||
Args:
|
||||
labels_cfg: EntityLabelsConfig, raw list of attribute dicts, or None
|
||||
|
||||
Returns:
|
||||
Set of lowercase 'key:value' strings
|
||||
"""
|
||||
if labels_cfg is None:
|
||||
return set()
|
||||
|
||||
# Accept raw list/dict for backwards compatibility
|
||||
if not isinstance(labels_cfg, EntityLabelsConfig):
|
||||
parsed = parse_entity_labels(labels_cfg)
|
||||
if parsed is None:
|
||||
return set()
|
||||
labels_cfg = parsed
|
||||
|
||||
valid = set()
|
||||
for group in labels_cfg.attributes:
|
||||
if group.type == "text":
|
||||
continue # No fixed vocabulary — all values accepted in post-processing
|
||||
for v in group.values:
|
||||
if group.key and v.value:
|
||||
valid.add(f"{group.key}:{v.value}".lower())
|
||||
return valid
|
||||
@@ -1,544 +0,0 @@
|
||||
"""
|
||||
Main orchestrator for the retain pipeline.
|
||||
|
||||
Coordinates all retain pipeline modules to store memories efficiently.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import time
|
||||
import uuid
|
||||
from collections.abc import Awaitable, Callable
|
||||
from datetime import UTC, datetime
|
||||
from typing import Any
|
||||
|
||||
from ..db_utils import acquire_with_retry, retry_with_backoff
|
||||
from . import bank_utils
|
||||
|
||||
|
||||
def utcnow():
|
||||
"""Get current UTC time."""
|
||||
return datetime.now(UTC)
|
||||
|
||||
|
||||
def parse_datetime_flexible(value: Any) -> datetime:
|
||||
"""
|
||||
Parse a datetime value that could be either a datetime object or an ISO string.
|
||||
|
||||
This handles datetime values from both direct Python calls and deserialized JSON
|
||||
(where datetime objects are serialized as ISO strings).
|
||||
|
||||
Args:
|
||||
value: Either a datetime object or an ISO format string
|
||||
|
||||
Returns:
|
||||
datetime object (timezone-aware)
|
||||
|
||||
Raises:
|
||||
TypeError: If value is neither datetime nor string
|
||||
ValueError: If string is not a valid ISO datetime
|
||||
"""
|
||||
if isinstance(value, datetime):
|
||||
# Ensure timezone-aware
|
||||
if value.tzinfo is None:
|
||||
return value.replace(tzinfo=UTC)
|
||||
return value
|
||||
elif isinstance(value, str):
|
||||
# Parse ISO format string (handles both 'Z' and '+00:00' timezone formats)
|
||||
dt = datetime.fromisoformat(value.replace("Z", "+00:00"))
|
||||
# Ensure timezone-aware
|
||||
if dt.tzinfo is None:
|
||||
return dt.replace(tzinfo=UTC)
|
||||
return dt
|
||||
else:
|
||||
raise TypeError(f"Expected datetime or string, got {type(value).__name__}")
|
||||
|
||||
|
||||
import asyncpg
|
||||
|
||||
from ..response_models import TokenUsage
|
||||
from . import (
|
||||
chunk_storage,
|
||||
embedding_processing,
|
||||
entity_processing,
|
||||
fact_extraction,
|
||||
fact_storage,
|
||||
link_creation,
|
||||
)
|
||||
from .types import EntityLink, ExtractedFact, ProcessedFact, RetainContent, RetainContentDict
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
async def retain_batch(
|
||||
pool,
|
||||
embeddings_model,
|
||||
llm_config,
|
||||
entity_resolver,
|
||||
format_date_fn,
|
||||
bank_id: str,
|
||||
contents_dicts: list[RetainContentDict],
|
||||
config,
|
||||
document_id: str | None = None,
|
||||
is_first_batch: bool = True,
|
||||
fact_type_override: str | None = None,
|
||||
confidence_score: float | None = None,
|
||||
document_tags: list[str] | None = None,
|
||||
operation_id: str | None = None,
|
||||
schema: str | None = None,
|
||||
outbox_callback: Callable[["asyncpg.Connection"], Awaitable[None]] | None = None,
|
||||
) -> tuple[list[list[str]], TokenUsage]:
|
||||
"""
|
||||
Process a batch of content through the retain pipeline.
|
||||
|
||||
Args:
|
||||
pool: Database connection pool
|
||||
embeddings_model: Embeddings model for generating embeddings
|
||||
llm_config: LLM configuration for fact extraction
|
||||
entity_resolver: Entity resolver for entity processing
|
||||
format_date_fn: Function to format datetime to readable string
|
||||
bank_id: Bank identifier
|
||||
contents_dicts: List of content dictionaries
|
||||
config: Resolved HindsightConfig for this bank
|
||||
document_id: Optional document ID
|
||||
is_first_batch: Whether this is the first batch
|
||||
fact_type_override: Override fact type for all facts
|
||||
confidence_score: Confidence score for opinions
|
||||
document_tags: Tags applied to all items in this batch
|
||||
|
||||
Returns:
|
||||
Tuple of (unit ID lists, token usage for fact extraction)
|
||||
"""
|
||||
start_time = time.time()
|
||||
total_chars = sum(len(item.get("content", "")) for item in contents_dicts)
|
||||
|
||||
# Buffer all logs
|
||||
log_buffer = []
|
||||
log_buffer.append(f"{'=' * 60}")
|
||||
log_buffer.append(f"RETAIN_BATCH START: {bank_id}")
|
||||
log_buffer.append(f"Batch size: {len(contents_dicts)} content items, {total_chars:,} chars")
|
||||
log_buffer.append(f"{'=' * 60}")
|
||||
|
||||
# Get bank profile
|
||||
profile = await bank_utils.get_bank_profile(pool, bank_id)
|
||||
agent_name = profile["name"]
|
||||
|
||||
# Convert dicts to RetainContent objects
|
||||
contents = []
|
||||
for item in contents_dicts:
|
||||
# Merge item-level tags with document-level tags
|
||||
item_tags = item.get("tags", []) or []
|
||||
merged_tags = list(set(item_tags + (document_tags or [])))
|
||||
|
||||
# Handle event_date: distinguish "not provided" (default to now) from
|
||||
# "explicitly None" (caller opted into no timestamp).
|
||||
if "event_date" in item and item["event_date"] is None:
|
||||
event_date_value = None # Caller explicitly signalled "unknown date"
|
||||
elif item.get("event_date"):
|
||||
event_date_value = parse_datetime_flexible(item["event_date"])
|
||||
else:
|
||||
event_date_value = utcnow() # Backward-compatible default
|
||||
|
||||
content = RetainContent(
|
||||
content=item["content"],
|
||||
context=item.get("context", ""),
|
||||
event_date=event_date_value,
|
||||
metadata=item.get("metadata", {}),
|
||||
entities=item.get("entities", []),
|
||||
tags=merged_tags,
|
||||
observation_scopes=item.get("observation_scopes"),
|
||||
)
|
||||
contents.append(content)
|
||||
|
||||
# Step 1: Extract facts from all contents
|
||||
step_start = time.time()
|
||||
|
||||
extracted_facts, chunks, usage = await fact_extraction.extract_facts_from_contents(
|
||||
contents, llm_config, agent_name, config, pool, operation_id, schema
|
||||
)
|
||||
log_buffer.append(
|
||||
f"[1] Extract facts: {len(extracted_facts)} facts, {len(chunks)} chunks from {len(contents)} contents in {time.time() - step_start:.3f}s"
|
||||
)
|
||||
|
||||
if not extracted_facts:
|
||||
# Still need to create document if document_id was provided or chunks exist
|
||||
from collections import defaultdict
|
||||
|
||||
docs_tracked = 0
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
async with conn.transaction():
|
||||
# Group contents by document_id (consistent with normal path)
|
||||
contents_by_doc_early = defaultdict(list)
|
||||
for idx, content_dict in enumerate(contents_dicts):
|
||||
doc_id = content_dict.get("document_id")
|
||||
contents_by_doc_early[doc_id].append((idx, content_dict))
|
||||
|
||||
if document_id:
|
||||
# Legacy: single document_id parameter
|
||||
combined_content = "\n".join([c.get("content", "") for c in contents_dicts])
|
||||
# Collect tags from all content items and merge with document_tags
|
||||
all_tags = set(document_tags or [])
|
||||
for item in contents_dicts:
|
||||
item_tags = item.get("tags", []) or []
|
||||
all_tags.update(item_tags)
|
||||
merged_tags = list(all_tags)
|
||||
|
||||
retain_params = {}
|
||||
if contents_dicts:
|
||||
first_item = contents_dicts[0]
|
||||
if first_item.get("context"):
|
||||
retain_params["context"] = first_item["context"]
|
||||
if first_item.get("event_date"):
|
||||
retain_params["event_date"] = (
|
||||
first_item["event_date"].isoformat()
|
||||
if hasattr(first_item["event_date"], "isoformat")
|
||||
else str(first_item["event_date"])
|
||||
)
|
||||
if first_item.get("metadata"):
|
||||
retain_params["metadata"] = first_item["metadata"]
|
||||
await fact_storage.handle_document_tracking(
|
||||
conn, bank_id, document_id, combined_content, is_first_batch, retain_params, merged_tags
|
||||
)
|
||||
docs_tracked += 1
|
||||
else:
|
||||
# Handle per-item document_ids and/or chunks (mirrors normal path logic)
|
||||
has_any_doc_ids = any(item.get("document_id") for item in contents_dicts)
|
||||
|
||||
if has_any_doc_ids or chunks:
|
||||
for original_doc_id, doc_contents in contents_by_doc_early.items():
|
||||
should_create_doc = (original_doc_id is not None) or chunks
|
||||
if not should_create_doc:
|
||||
continue
|
||||
|
||||
actual_doc_id = original_doc_id
|
||||
if actual_doc_id is None:
|
||||
# No document_id but have chunks - generate one
|
||||
actual_doc_id = str(uuid.uuid4())
|
||||
|
||||
combined_content = "\n".join([c.get("content", "") for _, c in doc_contents])
|
||||
all_tags = set(document_tags or [])
|
||||
for _, item in doc_contents:
|
||||
item_tags = item.get("tags", []) or []
|
||||
all_tags.update(item_tags)
|
||||
merged_tags = list(all_tags)
|
||||
|
||||
retain_params = {}
|
||||
if doc_contents:
|
||||
first_item = doc_contents[0][1]
|
||||
if first_item.get("context"):
|
||||
retain_params["context"] = first_item["context"]
|
||||
if first_item.get("event_date"):
|
||||
retain_params["event_date"] = (
|
||||
first_item["event_date"].isoformat()
|
||||
if hasattr(first_item["event_date"], "isoformat")
|
||||
else str(first_item["event_date"])
|
||||
)
|
||||
if first_item.get("metadata"):
|
||||
retain_params["metadata"] = first_item["metadata"]
|
||||
await fact_storage.handle_document_tracking(
|
||||
conn,
|
||||
bank_id,
|
||||
actual_doc_id,
|
||||
combined_content,
|
||||
is_first_batch,
|
||||
retain_params,
|
||||
merged_tags,
|
||||
)
|
||||
docs_tracked += 1
|
||||
|
||||
total_time = time.time() - start_time
|
||||
doc_status = f"{docs_tracked} document(s) tracked" if docs_tracked > 0 else "no document tracked"
|
||||
logger.info(
|
||||
f"RETAIN_BATCH COMPLETE: 0 facts extracted from {len(contents)} contents in {total_time:.3f}s ({doc_status}, no facts)"
|
||||
)
|
||||
return [[] for _ in contents], usage
|
||||
|
||||
# Apply fact_type_override if provided
|
||||
if fact_type_override:
|
||||
for fact in extracted_facts:
|
||||
fact.fact_type = fact_type_override
|
||||
|
||||
# Step 2: Augment texts and generate embeddings
|
||||
step_start = time.time()
|
||||
augmented_texts = embedding_processing.augment_texts_with_dates(extracted_facts, format_date_fn)
|
||||
embeddings = await embedding_processing.generate_embeddings_batch(embeddings_model, augmented_texts)
|
||||
log_buffer.append(f"[2] Generate embeddings: {len(embeddings)} embeddings in {time.time() - step_start:.3f}s")
|
||||
|
||||
# Step 3: Convert to ProcessedFact objects (without chunk_ids yet)
|
||||
processed_facts = [
|
||||
ProcessedFact.from_extracted_fact(extracted_fact, embedding)
|
||||
for extracted_fact, embedding in zip(extracted_facts, embeddings)
|
||||
]
|
||||
|
||||
# Group contents by document_id for document tracking and chunk storage
|
||||
from collections import defaultdict
|
||||
|
||||
contents_by_doc = defaultdict(list)
|
||||
for idx, content_dict in enumerate(contents_dicts):
|
||||
doc_id = content_dict.get("document_id")
|
||||
contents_by_doc[doc_id].append((idx, content_dict))
|
||||
|
||||
# Step 4: Database transaction (retried on deadlock)
|
||||
result_unit_ids: list[list[str]] = []
|
||||
|
||||
log_buffer_pre_db = len(log_buffer)
|
||||
|
||||
async def _run_db_work() -> None:
|
||||
nonlocal result_unit_ids
|
||||
|
||||
# Reset per-fact mutations and log buffer so each retry attempt starts clean
|
||||
del log_buffer[log_buffer_pre_db:]
|
||||
document_ids_added: list[str] = []
|
||||
for pf in processed_facts:
|
||||
pf.document_id = None
|
||||
pf.chunk_id = None
|
||||
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
async with conn.transaction():
|
||||
# Handle document tracking for all documents
|
||||
step_start = time.time()
|
||||
# Map None document_id to generated UUIDs
|
||||
doc_id_mapping = {} # Maps original doc_id (including None) to actual doc_id used
|
||||
|
||||
if document_id:
|
||||
# Legacy: single document_id parameter
|
||||
combined_content = "\n".join([c.get("content", "") for c in contents_dicts])
|
||||
retain_params = {}
|
||||
# Collect tags from all content items and merge with document_tags
|
||||
all_tags = set(document_tags or [])
|
||||
for item in contents_dicts:
|
||||
item_tags = item.get("tags", []) or []
|
||||
all_tags.update(item_tags)
|
||||
merged_tags = list(all_tags)
|
||||
|
||||
if contents_dicts:
|
||||
first_item = contents_dicts[0]
|
||||
if first_item.get("context"):
|
||||
retain_params["context"] = first_item["context"]
|
||||
if first_item.get("event_date"):
|
||||
retain_params["event_date"] = (
|
||||
first_item["event_date"].isoformat()
|
||||
if hasattr(first_item["event_date"], "isoformat")
|
||||
else str(first_item["event_date"])
|
||||
)
|
||||
if first_item.get("metadata"):
|
||||
retain_params["metadata"] = first_item["metadata"]
|
||||
|
||||
await fact_storage.handle_document_tracking(
|
||||
conn, bank_id, document_id, combined_content, is_first_batch, retain_params, merged_tags
|
||||
)
|
||||
document_ids_added.append(document_id)
|
||||
doc_id_mapping[None] = document_id # For backwards compatibility
|
||||
else:
|
||||
# Handle per-item document_ids (create documents if any item has document_id or if chunks exist)
|
||||
has_any_doc_ids = any(item.get("document_id") for item in contents_dicts)
|
||||
|
||||
if has_any_doc_ids or chunks:
|
||||
for original_doc_id, doc_contents in contents_by_doc.items():
|
||||
actual_doc_id = original_doc_id
|
||||
|
||||
# Only create document record if:
|
||||
# 1. Item has explicit document_id, OR
|
||||
# 2. There are chunks (need document for chunk storage)
|
||||
should_create_doc = (original_doc_id is not None) or chunks
|
||||
|
||||
if should_create_doc:
|
||||
if actual_doc_id is None:
|
||||
# No document_id but have chunks - generate one
|
||||
actual_doc_id = str(uuid.uuid4())
|
||||
|
||||
# Store mapping for later use
|
||||
doc_id_mapping[original_doc_id] = actual_doc_id
|
||||
|
||||
# Combine content for this document
|
||||
combined_content = "\n".join([c.get("content", "") for _, c in doc_contents])
|
||||
|
||||
# Collect tags from all content items for this document and merge with document_tags
|
||||
all_tags = set(document_tags or [])
|
||||
for _, item in doc_contents:
|
||||
item_tags = item.get("tags", []) or []
|
||||
all_tags.update(item_tags)
|
||||
merged_tags = list(all_tags)
|
||||
|
||||
# Extract retain params from first content item
|
||||
retain_params = {}
|
||||
if doc_contents:
|
||||
first_item = doc_contents[0][1]
|
||||
if first_item.get("context"):
|
||||
retain_params["context"] = first_item["context"]
|
||||
if first_item.get("event_date"):
|
||||
retain_params["event_date"] = (
|
||||
first_item["event_date"].isoformat()
|
||||
if hasattr(first_item["event_date"], "isoformat")
|
||||
else str(first_item["event_date"])
|
||||
)
|
||||
if first_item.get("metadata"):
|
||||
retain_params["metadata"] = first_item["metadata"]
|
||||
|
||||
await fact_storage.handle_document_tracking(
|
||||
conn,
|
||||
bank_id,
|
||||
actual_doc_id,
|
||||
combined_content,
|
||||
is_first_batch,
|
||||
retain_params,
|
||||
merged_tags,
|
||||
)
|
||||
document_ids_added.append(actual_doc_id)
|
||||
|
||||
if document_ids_added:
|
||||
log_buffer.append(
|
||||
f"[2.5] Document tracking: {len(document_ids_added)} documents in {time.time() - step_start:.3f}s"
|
||||
)
|
||||
|
||||
# Store chunks and map to facts for all documents
|
||||
step_start = time.time()
|
||||
chunk_id_map_by_doc = {} # Maps (doc_id, chunk_index) -> chunk_id
|
||||
|
||||
if chunks:
|
||||
# Group chunks by their source document
|
||||
chunks_by_doc = defaultdict(list)
|
||||
for chunk in chunks:
|
||||
# chunk.content_index tells us which content this chunk came from
|
||||
original_doc_id = contents_dicts[chunk.content_index].get("document_id")
|
||||
# Map to actual document_id (handles None -> generated UUID mapping)
|
||||
actual_doc_id = doc_id_mapping.get(original_doc_id, original_doc_id)
|
||||
if actual_doc_id is None and document_id:
|
||||
actual_doc_id = document_id
|
||||
chunks_by_doc[actual_doc_id].append(chunk)
|
||||
|
||||
# Store chunks for each document
|
||||
for doc_id, doc_chunks in chunks_by_doc.items():
|
||||
chunk_id_map = await chunk_storage.store_chunks_batch(conn, bank_id, doc_id, doc_chunks)
|
||||
# Store mapping with document context
|
||||
for chunk_idx, chunk_id in chunk_id_map.items():
|
||||
chunk_id_map_by_doc[(doc_id, chunk_idx)] = chunk_id
|
||||
|
||||
log_buffer.append(
|
||||
f"[3] Store chunks: {len(chunks)} chunks for {len(chunks_by_doc)} documents in {time.time() - step_start:.3f}s"
|
||||
)
|
||||
|
||||
# Map chunk_ids and document_ids to facts
|
||||
for fact, processed_fact in zip(extracted_facts, processed_facts):
|
||||
# Get the original document_id for this fact's source content
|
||||
original_doc_id = contents_dicts[fact.content_index].get("document_id")
|
||||
# Map to actual document_id (handles None -> generated UUID mapping)
|
||||
actual_doc_id = doc_id_mapping.get(original_doc_id, original_doc_id)
|
||||
if actual_doc_id is None and document_id:
|
||||
actual_doc_id = document_id
|
||||
|
||||
# Set document_id on the fact
|
||||
processed_fact.document_id = actual_doc_id
|
||||
|
||||
# Map chunk_id if this fact came from a chunk
|
||||
if fact.chunk_index is not None:
|
||||
# Look up chunk_id using (doc_id, chunk_index)
|
||||
chunk_id = chunk_id_map_by_doc.get((actual_doc_id, fact.chunk_index))
|
||||
if chunk_id:
|
||||
processed_fact.chunk_id = chunk_id
|
||||
else:
|
||||
# No chunks - still need to set document_id on facts
|
||||
for fact, processed_fact in zip(extracted_facts, processed_facts):
|
||||
original_doc_id = contents_dicts[fact.content_index].get("document_id")
|
||||
# Map to actual document_id (handles None -> generated UUID mapping)
|
||||
actual_doc_id = doc_id_mapping.get(original_doc_id, original_doc_id)
|
||||
if actual_doc_id is None and document_id:
|
||||
actual_doc_id = document_id
|
||||
processed_fact.document_id = actual_doc_id
|
||||
|
||||
non_duplicate_facts = processed_facts
|
||||
|
||||
# Insert facts (document_id is now stored per-fact)
|
||||
step_start = time.time()
|
||||
unit_ids = await fact_storage.insert_facts_batch(conn, bank_id, non_duplicate_facts)
|
||||
log_buffer.append(f"[5] Insert facts: {len(unit_ids)} units in {time.time() - step_start:.3f}s")
|
||||
|
||||
# Process entities
|
||||
step_start = time.time()
|
||||
# Build map of content_index -> user entities for merging
|
||||
user_entities_per_content = {
|
||||
idx: content.entities for idx, content in enumerate(contents) if content.entities
|
||||
}
|
||||
entity_links = await entity_processing.process_entities_batch(
|
||||
entity_resolver,
|
||||
conn,
|
||||
bank_id,
|
||||
unit_ids,
|
||||
non_duplicate_facts,
|
||||
log_buffer,
|
||||
user_entities_per_content=user_entities_per_content,
|
||||
entity_labels=getattr(config, "entity_labels", None),
|
||||
)
|
||||
log_buffer.append(f"[6] Process entities: {len(entity_links)} links in {time.time() - step_start:.3f}s")
|
||||
|
||||
# Create temporal links
|
||||
step_start = time.time()
|
||||
temporal_link_count = await link_creation.create_temporal_links_batch(conn, bank_id, unit_ids)
|
||||
log_buffer.append(f"[7] Temporal links: {temporal_link_count} links in {time.time() - step_start:.3f}s")
|
||||
|
||||
# Create semantic links
|
||||
step_start = time.time()
|
||||
embeddings_for_links = [fact.embedding for fact in non_duplicate_facts]
|
||||
semantic_link_count = await link_creation.create_semantic_links_batch(
|
||||
conn, bank_id, unit_ids, embeddings_for_links
|
||||
)
|
||||
log_buffer.append(f"[8] Semantic links: {semantic_link_count} links in {time.time() - step_start:.3f}s")
|
||||
|
||||
# Insert entity links
|
||||
step_start = time.time()
|
||||
if entity_links:
|
||||
await entity_processing.insert_entity_links_batch(conn, entity_links)
|
||||
log_buffer.append(
|
||||
f"[9] Entity links: {len(entity_links) if entity_links else 0} links in {time.time() - step_start:.3f}s"
|
||||
)
|
||||
|
||||
# Create causal links
|
||||
step_start = time.time()
|
||||
causal_link_count = await link_creation.create_causal_links_batch(conn, unit_ids, non_duplicate_facts)
|
||||
log_buffer.append(f"[10] Causal links: {causal_link_count} links in {time.time() - step_start:.3f}s")
|
||||
|
||||
# Map results back to original content items
|
||||
result_unit_ids = _map_results_to_contents(contents, extracted_facts, unit_ids)
|
||||
|
||||
# Transactional outbox: queue any side-effect tasks (e.g. webhook deliveries)
|
||||
# inside the same transaction so they are atomically committed with the retain data.
|
||||
if outbox_callback:
|
||||
await outbox_callback(conn)
|
||||
|
||||
# Flush entity stats (mention_count / last_seen) now that the transaction
|
||||
# has committed. Uses a fresh pool connection — no locks held.
|
||||
await entity_resolver.flush_pending_stats()
|
||||
|
||||
# Log final summary
|
||||
total_time = time.time() - start_time
|
||||
log_buffer.append(f"{'=' * 60}")
|
||||
log_buffer.append(f"RETAIN_BATCH COMPLETE: {len(unit_ids)} units in {total_time:.3f}s")
|
||||
if document_ids_added:
|
||||
log_buffer.append(f"Documents: {', '.join(document_ids_added)}")
|
||||
log_buffer.append(f"{'=' * 60}")
|
||||
|
||||
logger.info("\n" + "\n".join(log_buffer) + "\n")
|
||||
|
||||
await retry_with_backoff(_run_db_work)
|
||||
return result_unit_ids, usage
|
||||
|
||||
|
||||
def _map_results_to_contents(
|
||||
contents: list[RetainContent],
|
||||
extracted_facts: list[ExtractedFact],
|
||||
unit_ids: list[str],
|
||||
) -> list[list[str]]:
|
||||
"""Map created unit IDs back to original content items."""
|
||||
facts_by_content: dict[int, list[int]] = {i: [] for i in range(len(contents))}
|
||||
for i, fact in enumerate(extracted_facts):
|
||||
facts_by_content[fact.content_index].append(i)
|
||||
|
||||
result_unit_ids = []
|
||||
unit_idx = 0
|
||||
for content_index in range(len(contents)):
|
||||
content_unit_ids = []
|
||||
for _ in facts_by_content[content_index]:
|
||||
content_unit_ids.append(unit_ids[unit_idx])
|
||||
unit_idx += 1
|
||||
result_unit_ids.append(content_unit_ids)
|
||||
|
||||
return result_unit_ids
|
||||
@@ -1,493 +0,0 @@
|
||||
"""
|
||||
Link Expansion graph retrieval.
|
||||
|
||||
Expands from semantic/temporal seeds through three parallel, first-class signals
|
||||
stored in memory_links:
|
||||
|
||||
1. Entity links — precomputed co-occurrence graph (created at retain time, bounded to
|
||||
MAX_LINKS_PER_ENTITY per entity). Score = number of distinct shared
|
||||
entities between the seed set and each candidate.
|
||||
2. Semantic links — precomputed kNN graph (each new fact linked to its top-5 most
|
||||
similar existing facts at insert time, similarity >= 0.7). Checked
|
||||
in both directions since the graph is not symmetric. Score = weight.
|
||||
3. Causal links — explicit causal chains (causes/caused_by/enables/prevents).
|
||||
Score = weight + 1.0 (boosted as highest-quality signal).
|
||||
|
||||
All three signals are bounded at retain time, so no LATERAL fan-out caps are needed
|
||||
at query time. Each expansion is a simple aggregation over a small result set.
|
||||
|
||||
For non-observation fact types the three expansions are issued as a single CTE query
|
||||
(one roundtrip, one connection) with a `source` discriminator column so the Python
|
||||
merge step can apply per-signal score transformations.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import math
|
||||
import time
|
||||
|
||||
from ..db_utils import acquire_with_retry
|
||||
from ..memory_engine import fq_table
|
||||
from .graph_retrieval import GraphRetriever
|
||||
from .tags import TagGroup, TagsMatch, filter_results_by_tag_groups, filter_results_by_tags
|
||||
from .types import MPFPTimings, RetrievalResult
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
async def _find_semantic_seeds(
|
||||
conn,
|
||||
query_embedding_str: str,
|
||||
bank_id: str,
|
||||
fact_type: str,
|
||||
limit: int = 20,
|
||||
threshold: float = 0.3,
|
||||
tags: list[str] | None = None,
|
||||
tags_match: TagsMatch = "any",
|
||||
tag_groups: list[TagGroup] | None = None,
|
||||
) -> list[RetrievalResult]:
|
||||
"""Find semantic seeds via embedding search."""
|
||||
from .tags import build_tag_groups_where_clause, build_tags_where_clause_simple
|
||||
|
||||
tags_clause = build_tags_where_clause_simple(tags, 6, match=tags_match)
|
||||
tag_groups_param_start = 6 + (1 if tags else 0)
|
||||
groups_clause, groups_params, _ = build_tag_groups_where_clause(tag_groups, tag_groups_param_start)
|
||||
params = [query_embedding_str, bank_id, fact_type, threshold, limit]
|
||||
if tags:
|
||||
params.append(tags)
|
||||
params.extend(groups_params)
|
||||
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT id, text, context, event_date, occurred_start, occurred_end,
|
||||
mentioned_at, fact_type, document_id, chunk_id, tags,
|
||||
1 - (embedding <=> $1::vector) AS similarity
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE bank_id = $2
|
||||
AND embedding IS NOT NULL
|
||||
AND fact_type = $3
|
||||
AND (1 - (embedding <=> $1::vector)) >= $4
|
||||
{tags_clause}
|
||||
{groups_clause}
|
||||
ORDER BY embedding <=> $1::vector
|
||||
LIMIT $5
|
||||
""",
|
||||
*params,
|
||||
)
|
||||
return [RetrievalResult.from_db_row(dict(r)) for r in rows]
|
||||
|
||||
|
||||
class LinkExpansionRetriever(GraphRetriever):
|
||||
"""
|
||||
Graph retrieval via direct link expansion from seeds.
|
||||
|
||||
Runs three expansions through precomputed memory_links: entity co-occurrence,
|
||||
semantic kNN, and causal chains, all bounded at retain time.
|
||||
|
||||
For non-observation fact types the three expansions are issued as a single CTE
|
||||
query (one roundtrip, one connection slot) with a `source` discriminator column.
|
||||
The Python merge step applies per-signal score transformations.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
causal_weight_threshold: float = 0.3,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
causal_weight_threshold: Minimum weight for causal links to follow.
|
||||
"""
|
||||
self.causal_weight_threshold = causal_weight_threshold
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "link_expansion"
|
||||
|
||||
async def retrieve(
|
||||
self,
|
||||
pool,
|
||||
query_embedding_str: str,
|
||||
bank_id: str,
|
||||
fact_type: str,
|
||||
budget: int,
|
||||
query_text: str | None = None,
|
||||
semantic_seeds: list[RetrievalResult] | None = None,
|
||||
temporal_seeds: list[RetrievalResult] | None = None,
|
||||
adjacency=None,
|
||||
tags: list[str] | None = None,
|
||||
tags_match: TagsMatch = "any",
|
||||
tag_groups: list[TagGroup] | None = None,
|
||||
) -> tuple[list[RetrievalResult], MPFPTimings | None]:
|
||||
"""
|
||||
Retrieve facts by expanding links from seeds.
|
||||
|
||||
Args:
|
||||
pool: Database connection pool
|
||||
query_embedding_str: Query embedding as string
|
||||
bank_id: Memory bank ID
|
||||
fact_type: Fact type to filter
|
||||
budget: Maximum results to return
|
||||
query_text: Original query text (unused)
|
||||
semantic_seeds: Pre-computed semantic entry points
|
||||
temporal_seeds: Pre-computed temporal entry points
|
||||
adjacency: Unused, kept for interface compatibility
|
||||
tags: Optional list of tags for visibility filtering
|
||||
|
||||
Returns:
|
||||
Tuple of (results, timings)
|
||||
"""
|
||||
start_time = time.time()
|
||||
timings = MPFPTimings(fact_type=fact_type)
|
||||
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
# Find seeds if not provided
|
||||
if semantic_seeds:
|
||||
all_seeds = list(semantic_seeds)
|
||||
else:
|
||||
seeds_start = time.time()
|
||||
all_seeds = await _find_semantic_seeds(
|
||||
conn,
|
||||
query_embedding_str,
|
||||
bank_id,
|
||||
fact_type,
|
||||
limit=20,
|
||||
threshold=0.3,
|
||||
tags=tags,
|
||||
tags_match=tags_match,
|
||||
tag_groups=tag_groups,
|
||||
)
|
||||
timings.seeds_time = time.time() - seeds_start
|
||||
logger.debug(
|
||||
f"[LinkExpansion] Found {len(all_seeds)} semantic seeds for fact_type={fact_type} "
|
||||
f"(tags={tags}, tags_match={tags_match})"
|
||||
)
|
||||
|
||||
if temporal_seeds:
|
||||
all_seeds.extend(temporal_seeds)
|
||||
|
||||
if not all_seeds:
|
||||
return [], timings
|
||||
|
||||
seed_ids = list({s.id for s in all_seeds})
|
||||
timings.pattern_count = len(seed_ids)
|
||||
|
||||
query_start = time.time()
|
||||
|
||||
if fact_type == "observation":
|
||||
entity_rows, semantic_rows, causal_rows = await self._expand_observations(conn, seed_ids, budget)
|
||||
else:
|
||||
entity_rows, semantic_rows, causal_rows = await self._expand_combined(conn, seed_ids, fact_type, budget)
|
||||
|
||||
timings.edge_load_time = time.time() - query_start
|
||||
timings.db_queries = 1
|
||||
timings.edge_count = len(entity_rows) + len(semantic_rows) + len(causal_rows)
|
||||
|
||||
# Merge results with additive intra-score: entity + semantic + causal ∈ [0, 3].
|
||||
#
|
||||
# Entity score: tanh(count × 0.5) maps shared-entity count to [0, 1]:
|
||||
# 1 entity → 0.46, 2 → 0.76, 3 → 0.91, 4 → 0.96 (saturates naturally)
|
||||
# Semantic score: similarity weight, already ∈ [0.7, 1.0].
|
||||
# Causal score: link weight, already ∈ [0, 1].
|
||||
#
|
||||
# Facts appearing in multiple signals accumulate higher scores, rewarding
|
||||
# convergent evidence. The outer RRF uses rank position from this sorted list.
|
||||
entity_scores: dict[str, float] = {}
|
||||
semantic_scores: dict[str, float] = {}
|
||||
causal_scores: dict[str, float] = {}
|
||||
row_map: dict[str, dict] = {}
|
||||
|
||||
for row in entity_rows:
|
||||
fact_id = str(row["id"])
|
||||
entity_scores[fact_id] = math.tanh(row["score"] * 0.5)
|
||||
row_map[fact_id] = dict(row)
|
||||
|
||||
for row in semantic_rows:
|
||||
fact_id = str(row["id"])
|
||||
semantic_scores[fact_id] = max(semantic_scores.get(fact_id, 0.0), row["score"])
|
||||
row_map.setdefault(fact_id, dict(row))
|
||||
|
||||
for row in causal_rows:
|
||||
fact_id = str(row["id"])
|
||||
causal_scores[fact_id] = max(causal_scores.get(fact_id, 0.0), row["score"])
|
||||
row_map.setdefault(fact_id, dict(row))
|
||||
|
||||
all_ids = set(entity_scores) | set(semantic_scores) | set(causal_scores)
|
||||
score_map = {
|
||||
fid: entity_scores.get(fid, 0.0) + semantic_scores.get(fid, 0.0) + causal_scores.get(fid, 0.0)
|
||||
for fid in all_ids
|
||||
}
|
||||
|
||||
sorted_ids = sorted(score_map.keys(), key=lambda x: score_map[x], reverse=True)[:budget]
|
||||
rows = [row_map[fact_id] for fact_id in sorted_ids]
|
||||
|
||||
results = []
|
||||
for row in rows:
|
||||
result = RetrievalResult.from_db_row(dict(row))
|
||||
result.activation = row["score"]
|
||||
results.append(result)
|
||||
|
||||
if tags:
|
||||
results = filter_results_by_tags(results, tags, match=tags_match)
|
||||
|
||||
if tag_groups:
|
||||
results = filter_results_by_tag_groups(results, tag_groups)
|
||||
|
||||
timings.result_count = len(results)
|
||||
timings.traverse = time.time() - start_time
|
||||
|
||||
logger.debug(
|
||||
f"LinkExpansion: {len(results)} results from {len(seed_ids)} seeds "
|
||||
f"in {timings.traverse * 1000:.1f}ms (query: {timings.edge_load_time * 1000:.1f}ms)"
|
||||
)
|
||||
|
||||
return results, timings
|
||||
|
||||
async def _expand_combined(
|
||||
self,
|
||||
conn,
|
||||
seed_ids: list,
|
||||
fact_type: str,
|
||||
budget: int,
|
||||
) -> tuple[list, list, list]:
|
||||
"""
|
||||
Single-roundtrip CTE query combining entity, semantic, and causal expansions.
|
||||
|
||||
Uses a `source` discriminator column so the caller can apply per-signal
|
||||
score transformations. The three CTEs share one connection slot — important
|
||||
for asyncpg which does not allow concurrent queries on the same connection.
|
||||
|
||||
Index coverage (requires migration d2e3f4a5b6c7):
|
||||
entity: idx_memory_links_entity_covering (from_unit_id) INCLUDE (to_unit_id, entity_id)
|
||||
WHERE link_type = 'entity' → index-only scan, no heap reads
|
||||
semantic incoming:
|
||||
idx_memory_links_to_type_weight (to_unit_id, link_type, weight DESC)
|
||||
→ replaces costly BitmapAnd of two separate scans
|
||||
"""
|
||||
ml = fq_table("memory_links")
|
||||
mu = fq_table("memory_units")
|
||||
all_rows = await conn.fetch(
|
||||
f"""
|
||||
WITH entity_expanded AS (
|
||||
-- Entity co-occurrence: seeds → their precomputed entity-link neighbors.
|
||||
-- Score = distinct shared entities (bounded at retain time to
|
||||
-- MAX_LINKS_PER_ENTITY=50). GROUP BY mu.id is sufficient because mu.id
|
||||
-- is the primary key and functionally determines all other mu columns.
|
||||
SELECT
|
||||
mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
|
||||
mu.occurred_end, mu.mentioned_at,
|
||||
mu.fact_type, mu.document_id, mu.chunk_id, mu.tags,
|
||||
COUNT(DISTINCT ml.entity_id)::float AS score,
|
||||
'entity'::text AS source
|
||||
FROM {ml} ml
|
||||
JOIN {mu} mu ON mu.id = ml.to_unit_id
|
||||
WHERE ml.from_unit_id = ANY($1::uuid[])
|
||||
AND ml.link_type = 'entity'
|
||||
AND mu.fact_type = $2
|
||||
AND mu.id != ALL($1::uuid[])
|
||||
GROUP BY mu.id
|
||||
ORDER BY score DESC
|
||||
LIMIT $3
|
||||
),
|
||||
semantic_expanded AS (
|
||||
-- Semantic kNN: both outgoing (seeds → their kNN at insert time) and
|
||||
-- incoming (facts inserted after seeds that found seeds as kNN).
|
||||
-- Score = max similarity weight across both directions.
|
||||
SELECT
|
||||
id, text, context, event_date, occurred_start,
|
||||
occurred_end, mentioned_at,
|
||||
fact_type, document_id, chunk_id, tags,
|
||||
MAX(weight) AS score,
|
||||
'semantic'::text AS source
|
||||
FROM (
|
||||
SELECT
|
||||
mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
|
||||
mu.occurred_end, mu.mentioned_at,
|
||||
mu.fact_type, mu.document_id, mu.chunk_id, mu.tags,
|
||||
ml.weight
|
||||
FROM {ml} ml
|
||||
JOIN {mu} mu ON mu.id = ml.to_unit_id
|
||||
WHERE ml.from_unit_id = ANY($1::uuid[])
|
||||
AND ml.link_type = 'semantic'
|
||||
AND mu.fact_type = $2
|
||||
AND mu.id != ALL($1::uuid[])
|
||||
UNION ALL
|
||||
SELECT
|
||||
mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
|
||||
mu.occurred_end, mu.mentioned_at,
|
||||
mu.fact_type, mu.document_id, mu.chunk_id, mu.tags,
|
||||
ml.weight
|
||||
FROM {ml} ml
|
||||
JOIN {mu} mu ON mu.id = ml.from_unit_id
|
||||
WHERE ml.to_unit_id = ANY($1::uuid[])
|
||||
AND ml.link_type = 'semantic'
|
||||
AND mu.fact_type = $2
|
||||
AND mu.id != ALL($1::uuid[])
|
||||
) sem_raw
|
||||
GROUP BY id, text, context, event_date, occurred_start,
|
||||
occurred_end, mentioned_at,
|
||||
fact_type, document_id, chunk_id, tags
|
||||
ORDER BY score DESC
|
||||
LIMIT $3
|
||||
),
|
||||
causal_expanded AS (
|
||||
-- Causal chains: explicit causes/enables/prevents links from seeds.
|
||||
-- DISTINCT ON handles the case where a seed has multiple causal links
|
||||
-- to the same target; best weight wins.
|
||||
SELECT DISTINCT ON (mu.id)
|
||||
mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
|
||||
mu.occurred_end, mu.mentioned_at,
|
||||
mu.fact_type, mu.document_id, mu.chunk_id, mu.tags,
|
||||
ml.weight AS score,
|
||||
'causal'::text AS source
|
||||
FROM {ml} ml
|
||||
JOIN {mu} mu ON ml.to_unit_id = mu.id
|
||||
WHERE ml.from_unit_id = ANY($1::uuid[])
|
||||
AND ml.link_type IN ('causes', 'caused_by', 'enables', 'prevents')
|
||||
AND ml.weight >= $4
|
||||
AND mu.fact_type = $2
|
||||
ORDER BY mu.id, ml.weight DESC
|
||||
LIMIT $3
|
||||
)
|
||||
SELECT * FROM entity_expanded
|
||||
UNION ALL
|
||||
SELECT * FROM semantic_expanded
|
||||
UNION ALL
|
||||
SELECT * FROM causal_expanded
|
||||
""",
|
||||
seed_ids,
|
||||
fact_type,
|
||||
budget,
|
||||
self.causal_weight_threshold,
|
||||
)
|
||||
|
||||
entity_rows = [r for r in all_rows if r["source"] == "entity"]
|
||||
semantic_rows = [r for r in all_rows if r["source"] == "semantic"]
|
||||
causal_rows = [r for r in all_rows if r["source"] == "causal"]
|
||||
return entity_rows, semantic_rows, causal_rows
|
||||
|
||||
async def _expand_observations(
|
||||
self,
|
||||
conn,
|
||||
seed_ids: list,
|
||||
budget: int,
|
||||
) -> tuple[list, list, list]:
|
||||
"""
|
||||
Observation-specific expansion.
|
||||
|
||||
Observations don't have direct entity links in memory_links (they're created
|
||||
by consolidation, not retain). Instead, traverse source_memory_ids → world
|
||||
facts → entities → other world facts → their observations.
|
||||
|
||||
Semantic and causal expansions run as a second combined CTE query.
|
||||
"""
|
||||
source_ids_found: list = []
|
||||
if logger.isEnabledFor(logging.DEBUG):
|
||||
debug_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT id, source_memory_ids
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE id = ANY($1::uuid[])
|
||||
""",
|
||||
seed_ids,
|
||||
)
|
||||
for row in debug_rows:
|
||||
if row["source_memory_ids"]:
|
||||
source_ids_found.extend(row["source_memory_ids"])
|
||||
logger.debug(
|
||||
f"[LinkExpansion] observation graph: {len(seed_ids)} seeds, "
|
||||
f"{len(source_ids_found)} source_memory_ids found"
|
||||
)
|
||||
|
||||
entity_rows = await conn.fetch(
|
||||
f"""
|
||||
WITH seed_sources AS (
|
||||
SELECT DISTINCT unnest(source_memory_ids) AS source_id
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE id = ANY($1::uuid[])
|
||||
AND source_memory_ids IS NOT NULL
|
||||
),
|
||||
connected_sources AS (
|
||||
-- Mirror the non-observation entity expansion: follow pre-bounded entity
|
||||
-- links in memory_links (capped to MAX_LINKS_PER_ENTITY=50 at retain time).
|
||||
-- Score = number of distinct shared entities, same as the non-obs path.
|
||||
SELECT DISTINCT ml.to_unit_id AS source_id
|
||||
FROM seed_sources ss
|
||||
JOIN {fq_table("memory_links")} ml ON ml.from_unit_id = ss.source_id
|
||||
WHERE ml.link_type = 'entity'
|
||||
),
|
||||
connected_array AS (
|
||||
SELECT array_agg(source_id) AS source_ids FROM connected_sources
|
||||
)
|
||||
SELECT
|
||||
mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
|
||||
mu.occurred_end, mu.mentioned_at,
|
||||
mu.fact_type, mu.document_id, mu.chunk_id, mu.tags,
|
||||
(SELECT COUNT(DISTINCT s) FROM unnest(mu.source_memory_ids) s WHERE s = ANY(ca.source_ids))::float AS score
|
||||
FROM {fq_table("memory_units")} mu, connected_array ca
|
||||
WHERE mu.fact_type = 'observation'
|
||||
AND mu.id != ALL($1::uuid[])
|
||||
AND ca.source_ids IS NOT NULL
|
||||
AND mu.source_memory_ids && ca.source_ids
|
||||
ORDER BY score DESC
|
||||
LIMIT $2
|
||||
""",
|
||||
seed_ids,
|
||||
budget,
|
||||
)
|
||||
logger.debug(f"[LinkExpansion] observation graph: found {len(entity_rows)} connected observations")
|
||||
|
||||
# Semantic + causal for observations in one query
|
||||
ml = fq_table("memory_links")
|
||||
mu = fq_table("memory_units")
|
||||
sem_causal_rows = await conn.fetch(
|
||||
f"""
|
||||
WITH semantic_expanded AS (
|
||||
SELECT
|
||||
id, text, context, event_date, occurred_start,
|
||||
occurred_end, mentioned_at,
|
||||
fact_type, document_id, chunk_id, tags,
|
||||
MAX(weight) AS score,
|
||||
'semantic'::text AS source
|
||||
FROM (
|
||||
SELECT mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
|
||||
mu.occurred_end, mu.mentioned_at, mu.fact_type, mu.document_id,
|
||||
mu.chunk_id, mu.tags, ml.weight
|
||||
FROM {ml} ml JOIN {mu} mu ON mu.id = ml.to_unit_id
|
||||
WHERE ml.from_unit_id = ANY($1::uuid[])
|
||||
AND ml.link_type = 'semantic' AND mu.fact_type = 'observation'
|
||||
AND mu.id != ALL($1::uuid[])
|
||||
UNION ALL
|
||||
SELECT mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
|
||||
mu.occurred_end, mu.mentioned_at, mu.fact_type, mu.document_id,
|
||||
mu.chunk_id, mu.tags, ml.weight
|
||||
FROM {ml} ml JOIN {mu} mu ON mu.id = ml.from_unit_id
|
||||
WHERE ml.to_unit_id = ANY($1::uuid[])
|
||||
AND ml.link_type = 'semantic' AND mu.fact_type = 'observation'
|
||||
AND mu.id != ALL($1::uuid[])
|
||||
) sem_raw
|
||||
GROUP BY id, text, context, event_date, occurred_start, occurred_end,
|
||||
mentioned_at, fact_type, document_id, chunk_id, tags
|
||||
ORDER BY score DESC LIMIT $2
|
||||
),
|
||||
causal_expanded AS (
|
||||
SELECT DISTINCT ON (mu.id)
|
||||
mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
|
||||
mu.occurred_end, mu.mentioned_at, mu.fact_type, mu.document_id,
|
||||
mu.chunk_id, mu.tags, ml.weight AS score, 'causal'::text AS source
|
||||
FROM {ml} ml JOIN {mu} mu ON ml.to_unit_id = mu.id
|
||||
WHERE ml.from_unit_id = ANY($1::uuid[])
|
||||
AND ml.link_type IN ('causes', 'caused_by', 'enables', 'prevents')
|
||||
AND ml.weight >= $3 AND mu.fact_type = 'observation'
|
||||
ORDER BY mu.id, ml.weight DESC LIMIT $2
|
||||
)
|
||||
SELECT * FROM semantic_expanded
|
||||
UNION ALL
|
||||
SELECT * FROM causal_expanded
|
||||
""",
|
||||
seed_ids,
|
||||
budget,
|
||||
self.causal_weight_threshold,
|
||||
)
|
||||
|
||||
semantic_rows = [r for r in sem_causal_rows if r["source"] == "semantic"]
|
||||
causal_rows = [r for r in sem_causal_rows if r["source"] == "causal"]
|
||||
return entity_rows, semantic_rows, causal_rows
|
||||
@@ -1,690 +0,0 @@
|
||||
"""
|
||||
Retrieval module for 4-way parallel search.
|
||||
|
||||
Implements:
|
||||
1. Semantic retrieval (vector similarity)
|
||||
2. BM25 retrieval (keyword/full-text search)
|
||||
3. Graph retrieval (via pluggable GraphRetriever interface)
|
||||
4. Temporal retrieval (time-aware search with spreading)
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import UTC, datetime
|
||||
from typing import Optional
|
||||
|
||||
from ...config import get_config
|
||||
from ..db_utils import acquire_with_retry
|
||||
from ..memory_engine import fq_table
|
||||
from .graph_retrieval import BFSGraphRetriever, GraphRetriever
|
||||
from .link_expansion_retrieval import LinkExpansionRetriever
|
||||
from .mpfp_retrieval import MPFPGraphRetriever
|
||||
from .tags import TagGroup, TagsMatch, build_tag_groups_where_clause, build_tags_where_clause_simple
|
||||
from .types import MPFPTimings, RetrievalResult
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ParallelRetrievalResult:
|
||||
"""Result from parallel retrieval across all methods."""
|
||||
|
||||
semantic: list[RetrievalResult]
|
||||
bm25: list[RetrievalResult]
|
||||
graph: list[RetrievalResult]
|
||||
temporal: list[RetrievalResult] | None
|
||||
timings: dict[str, float] = field(default_factory=dict)
|
||||
temporal_constraint: tuple | None = None # (start_date, end_date)
|
||||
mpfp_timings: list[MPFPTimings] = field(default_factory=list) # MPFP sub-step timings per fact type
|
||||
max_conn_wait: float = 0.0 # Maximum connection acquisition wait time across all methods
|
||||
|
||||
|
||||
@dataclass
|
||||
class MultiFactTypeRetrievalResult:
|
||||
"""Result from retrieval across all fact types."""
|
||||
|
||||
# Results per fact type
|
||||
results_by_fact_type: dict[str, ParallelRetrievalResult]
|
||||
# Aggregate timings
|
||||
timings: dict[str, float] = field(default_factory=dict)
|
||||
# Max connection wait across all operations
|
||||
max_conn_wait: float = 0.0
|
||||
|
||||
|
||||
# Default graph retriever instance (can be overridden)
|
||||
_default_graph_retriever: GraphRetriever | None = None
|
||||
|
||||
|
||||
def get_default_graph_retriever() -> GraphRetriever:
|
||||
"""Get or create the default graph retriever based on config."""
|
||||
global _default_graph_retriever
|
||||
if _default_graph_retriever is None:
|
||||
config = get_config()
|
||||
retriever_type = config.graph_retriever.lower()
|
||||
if retriever_type == "mpfp":
|
||||
_default_graph_retriever = MPFPGraphRetriever()
|
||||
logger.info(
|
||||
f"Using MPFP graph retriever (top_k_neighbors={_default_graph_retriever.config.top_k_neighbors})"
|
||||
)
|
||||
elif retriever_type == "bfs":
|
||||
_default_graph_retriever = BFSGraphRetriever()
|
||||
logger.info("Using BFS graph retriever")
|
||||
elif retriever_type == "link_expansion":
|
||||
_default_graph_retriever = LinkExpansionRetriever()
|
||||
logger.info("Using LinkExpansion graph retriever")
|
||||
else:
|
||||
logger.warning(f"Unknown graph retriever '{retriever_type}', falling back to link_expansion")
|
||||
_default_graph_retriever = LinkExpansionRetriever()
|
||||
return _default_graph_retriever
|
||||
|
||||
|
||||
def set_default_graph_retriever(retriever: GraphRetriever) -> None:
|
||||
"""Set the default graph retriever (for configuration/testing)."""
|
||||
global _default_graph_retriever
|
||||
_default_graph_retriever = retriever
|
||||
|
||||
|
||||
async def retrieve_semantic_bm25_combined(
|
||||
conn,
|
||||
query_emb_str: str,
|
||||
query_text: str,
|
||||
bank_id: str,
|
||||
fact_types: list[str],
|
||||
limit: int,
|
||||
tags: list[str] | None = None,
|
||||
tags_match: TagsMatch = "any",
|
||||
tag_groups: list[TagGroup] | None = None,
|
||||
) -> dict[str, tuple[list[RetrievalResult], list[RetrievalResult]]]:
|
||||
"""
|
||||
Combined semantic + BM25 retrieval for multiple fact types in a single query.
|
||||
|
||||
Uses UNION ALL of per-fact_type subqueries so that each arm has its own
|
||||
ORDER BY ... LIMIT, enabling the partial HNSW indexes per fact_type instead
|
||||
of forcing a full sequential scan (which the previous window-function approach
|
||||
caused by using PARTITION BY inside ROW_NUMBER()).
|
||||
|
||||
Requires partial HNSW indexes per fact_type (idx_mu_emb_world,
|
||||
idx_mu_emb_observation, idx_mu_emb_experience), created automatically by
|
||||
Alembic migration a3b4c5d6e7f8_add_partial_hnsw_indexes.py.
|
||||
|
||||
HNSW is approximate — semantic arms over-fetch by 5x (min 100) and trim to
|
||||
limit in Python to compensate. ef_search=200 is set globally on pool
|
||||
connections at init time (see memory_engine.py) to improve recall on sparse
|
||||
graphs.
|
||||
|
||||
fact_type values are inlined as literals (safe: they come from a controlled
|
||||
internal enum, never from user input).
|
||||
|
||||
Args:
|
||||
conn: Database connection
|
||||
query_emb_str: Query embedding as string
|
||||
query_text: Query text for BM25
|
||||
bank_id: Bank ID
|
||||
fact_types: List of fact types to retrieve
|
||||
limit: Maximum results per method per fact type
|
||||
tags: Optional tags to filter by
|
||||
tags_match: Tag matching mode
|
||||
|
||||
Returns:
|
||||
Dict mapping fact_type -> (semantic_results, bm25_results)
|
||||
"""
|
||||
import re
|
||||
|
||||
result_dict: dict[str, tuple[list[RetrievalResult], list[RetrievalResult]]] = {ft: ([], []) for ft in fact_types}
|
||||
|
||||
sanitized_text = re.sub(r"[^\w\s]", " ", query_text.lower())
|
||||
tokens = [token for token in sanitized_text.split() if token]
|
||||
|
||||
# Over-fetch for HNSW approximation; semantic results trimmed to limit in Python.
|
||||
hnsw_fetch = max(limit * 5, 100)
|
||||
|
||||
cols = (
|
||||
"id, text, context, event_date, occurred_start, occurred_end, mentioned_at, "
|
||||
"fact_type, document_id, chunk_id, tags"
|
||||
)
|
||||
table = fq_table("memory_units")
|
||||
|
||||
# --- Parameter layout ---
|
||||
# $1 = query_emb_str (semantic arms)
|
||||
# $2 = bank_id
|
||||
# $3 = limit (BM25 LIMIT; semantic uses inlined hnsw_fetch literal)
|
||||
# $4 = bm25_text (only when tokens present)
|
||||
# $N = tags (N=4 when no tokens, N=5 when tokens present)
|
||||
# $M+ = tag_groups params (one per leaf, starting after tags param)
|
||||
tags_param_idx = 5 if tokens else 4
|
||||
tags_clause = build_tags_where_clause_simple(tags, tags_param_idx, match=tags_match)
|
||||
|
||||
# tag_groups params start immediately after the tags param slot
|
||||
tag_groups_param_start = tags_param_idx + (1 if tags else 0)
|
||||
groups_clause, groups_params, _ = build_tag_groups_where_clause(tag_groups, tag_groups_param_start)
|
||||
|
||||
# --- Semantic UNION ALL arms (one per fact_type) ---
|
||||
# Each arm has its own ORDER BY embedding <=> $1 LIMIT {hnsw_fetch}, which
|
||||
# lets the planner use the partial HNSW index for that fact_type.
|
||||
sem_arms = []
|
||||
for ft in fact_types:
|
||||
sem_arms.append(
|
||||
f"(SELECT {cols},"
|
||||
f" 1 - (embedding <=> $1::vector) AS similarity,"
|
||||
f" NULL::float AS bm25_score,"
|
||||
f" 'semantic' AS source"
|
||||
f" FROM {table}"
|
||||
f" WHERE bank_id = $2"
|
||||
f" AND fact_type = '{ft}'"
|
||||
f" AND embedding IS NOT NULL"
|
||||
f" AND (1 - (embedding <=> $1::vector)) >= 0.3"
|
||||
f" {tags_clause}"
|
||||
f" {groups_clause}"
|
||||
f" ORDER BY embedding <=> $1::vector"
|
||||
f" LIMIT {hnsw_fetch})"
|
||||
)
|
||||
|
||||
arms = sem_arms
|
||||
|
||||
# --- BM25 UNION ALL arms (one per fact_type, only when tokens present) ---
|
||||
if tokens:
|
||||
config = get_config()
|
||||
if config.text_search_extension == "vchord":
|
||||
bm25_score_expr = (
|
||||
"search_vector <&> to_bm25query('idx_memory_units_text_search', tokenize($4, 'llmlingua2'))"
|
||||
)
|
||||
bm25_order_by = f"{bm25_score_expr} DESC"
|
||||
bm25_where_filter = ""
|
||||
bm25_text_param: str = query_text
|
||||
elif config.text_search_extension == "pg_textsearch":
|
||||
bm25_score_expr = "-(text <@> to_bm25query($4, 'idx_memory_units_text_search'))"
|
||||
bm25_order_by = "text <@> to_bm25query($4, 'idx_memory_units_text_search') ASC"
|
||||
bm25_where_filter = ""
|
||||
bm25_text_param = query_text
|
||||
else: # native
|
||||
query_tsquery = " | ".join(tokens)
|
||||
bm25_score_expr = "ts_rank_cd(search_vector, to_tsquery('english', $4))"
|
||||
bm25_order_by = f"{bm25_score_expr} DESC"
|
||||
bm25_where_filter = "AND search_vector @@ to_tsquery('english', $4)"
|
||||
bm25_text_param = query_tsquery
|
||||
|
||||
for ft in fact_types:
|
||||
arms.append(
|
||||
f"(SELECT {cols},"
|
||||
f" NULL::float AS similarity,"
|
||||
f" {bm25_score_expr} AS bm25_score,"
|
||||
f" 'bm25' AS source"
|
||||
f" FROM {table}"
|
||||
f" WHERE bank_id = $2"
|
||||
f" AND fact_type = '{ft}'"
|
||||
f" {bm25_where_filter}"
|
||||
f" {tags_clause}"
|
||||
f" {groups_clause}"
|
||||
f" ORDER BY {bm25_order_by}"
|
||||
f" LIMIT $3)"
|
||||
)
|
||||
|
||||
query = "\nUNION ALL\n".join(arms)
|
||||
|
||||
params: list = [query_emb_str, bank_id, limit]
|
||||
if tokens:
|
||||
params.append(bm25_text_param)
|
||||
if tags:
|
||||
params.append(tags)
|
||||
params.extend(groups_params)
|
||||
|
||||
rows = await conn.fetch(query, *params)
|
||||
|
||||
# Group results; trim semantic to limit (over-fetched for HNSW approximation).
|
||||
sem_counts: dict[str, int] = {ft: 0 for ft in fact_types}
|
||||
for r in rows:
|
||||
row = dict(r)
|
||||
source = row.pop("source")
|
||||
ft = row.get("fact_type")
|
||||
if ft not in result_dict:
|
||||
continue
|
||||
if source == "semantic":
|
||||
if sem_counts[ft] < limit:
|
||||
result_dict[ft][0].append(RetrievalResult.from_db_row(row))
|
||||
sem_counts[ft] += 1
|
||||
else:
|
||||
result_dict[ft][1].append(RetrievalResult.from_db_row(row))
|
||||
|
||||
return result_dict
|
||||
|
||||
|
||||
async def retrieve_temporal_combined(
|
||||
conn,
|
||||
query_emb_str: str,
|
||||
bank_id: str,
|
||||
fact_types: list[str],
|
||||
start_date: datetime,
|
||||
end_date: datetime,
|
||||
budget: int,
|
||||
semantic_threshold: float = 0.1,
|
||||
tags: list[str] | None = None,
|
||||
tags_match: TagsMatch = "any",
|
||||
tag_groups: list[TagGroup] | None = None,
|
||||
) -> dict[str, list[RetrievalResult]]:
|
||||
"""
|
||||
Temporal retrieval for multiple fact types in a single query.
|
||||
|
||||
Batches the entry point query using window functions to get top-N per fact type,
|
||||
then runs spreading for each fact type.
|
||||
|
||||
Args:
|
||||
conn: Database connection
|
||||
query_emb_str: Query embedding as string
|
||||
bank_id: Bank ID
|
||||
fact_types: List of fact types to retrieve
|
||||
start_date: Start of time range
|
||||
end_date: End of time range
|
||||
budget: Node budget for spreading per fact type
|
||||
semantic_threshold: Minimum semantic similarity to include
|
||||
|
||||
Returns:
|
||||
Dict mapping fact_type -> list of RetrievalResult
|
||||
"""
|
||||
from ..memory_engine import fq_table
|
||||
|
||||
# Ensure dates are timezone-aware
|
||||
if start_date.tzinfo is None:
|
||||
start_date = start_date.replace(tzinfo=UTC)
|
||||
if end_date.tzinfo is None:
|
||||
end_date = end_date.replace(tzinfo=UTC)
|
||||
|
||||
# Build tags clause
|
||||
# Entry point query: fixed params are $1-$6, tags at $7
|
||||
tags_clause = build_tags_where_clause_simple(tags, 7, match=tags_match)
|
||||
tag_groups_param_start = 7 + (1 if tags else 0)
|
||||
groups_clause, groups_params, _ = build_tag_groups_where_clause(tag_groups, tag_groups_param_start)
|
||||
params: list = [query_emb_str, bank_id, fact_types, start_date, end_date, semantic_threshold]
|
||||
if tags:
|
||||
params.append(tags)
|
||||
params.extend(groups_params)
|
||||
|
||||
# Two-phase entry point query:
|
||||
# Phase 1 (date_ranked): rank by date only — no embedding computation — for all units in
|
||||
# the temporal window. This lets the planner use date indexes for filtering.
|
||||
# Phase 2 (sim_ranked): join back to memory_units for only the top-50-per-type candidates
|
||||
# and compute embedding similarity for that small set (≤ 50 × len(fact_types) rows).
|
||||
# This avoids computing embedding distances for potentially thousands of date-range rows.
|
||||
entry_points = await conn.fetch(
|
||||
f"""
|
||||
WITH date_ranked AS MATERIALIZED (
|
||||
SELECT id, fact_type,
|
||||
ROW_NUMBER() OVER (
|
||||
PARTITION BY fact_type
|
||||
ORDER BY COALESCE(occurred_start, mentioned_at, occurred_end) DESC NULLS LAST
|
||||
) AS rn
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE bank_id = $2
|
||||
AND fact_type = ANY($3)
|
||||
AND embedding IS NOT NULL
|
||||
AND (
|
||||
(occurred_start IS NOT NULL AND occurred_end IS NOT NULL
|
||||
AND occurred_start <= $5 AND occurred_end >= $4)
|
||||
OR
|
||||
(mentioned_at IS NOT NULL AND mentioned_at BETWEEN $4 AND $5)
|
||||
OR
|
||||
(occurred_start IS NOT NULL AND occurred_start BETWEEN $4 AND $5)
|
||||
OR
|
||||
(occurred_end IS NOT NULL AND occurred_end BETWEEN $4 AND $5)
|
||||
)
|
||||
{tags_clause}
|
||||
{groups_clause}
|
||||
),
|
||||
sim_ranked AS (
|
||||
SELECT mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start, mu.occurred_end, mu.mentioned_at, mu.fact_type, mu.document_id, mu.chunk_id, mu.tags,
|
||||
1 - (mu.embedding <=> $1::vector) AS similarity,
|
||||
ROW_NUMBER() OVER (PARTITION BY mu.fact_type ORDER BY mu.embedding <=> $1::vector) AS sim_rn
|
||||
FROM date_ranked dr
|
||||
JOIN {fq_table("memory_units")} mu ON mu.id = dr.id
|
||||
WHERE dr.rn <= 50
|
||||
AND (1 - (mu.embedding <=> $1::vector)) >= $6
|
||||
)
|
||||
SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, fact_type, document_id, chunk_id, tags, similarity
|
||||
FROM sim_ranked
|
||||
WHERE sim_rn <= 10
|
||||
""",
|
||||
*params,
|
||||
)
|
||||
|
||||
if not entry_points:
|
||||
return {ft: [] for ft in fact_types}
|
||||
|
||||
# Group entry points by fact type
|
||||
entries_by_ft: dict[str, list] = {ft: [] for ft in fact_types}
|
||||
for ep in entry_points:
|
||||
ft = ep["fact_type"]
|
||||
if ft in entries_by_ft:
|
||||
entries_by_ft[ft].append(ep)
|
||||
|
||||
# Calculate shared temporal parameters
|
||||
total_days = (end_date - start_date).total_seconds() / 86400
|
||||
mid_date = start_date + (end_date - start_date) / 2
|
||||
|
||||
# Process each fact type (spreading needs to stay per fact type due to link filtering)
|
||||
results_by_ft: dict[str, list[RetrievalResult]] = {}
|
||||
|
||||
for ft in fact_types:
|
||||
ft_entry_points = entries_by_ft.get(ft, [])
|
||||
if not ft_entry_points:
|
||||
results_by_ft[ft] = []
|
||||
continue
|
||||
|
||||
results = []
|
||||
visited = set()
|
||||
node_scores = {}
|
||||
|
||||
# Process entry points
|
||||
for ep in ft_entry_points:
|
||||
unit_id = str(ep["id"])
|
||||
visited.add(unit_id)
|
||||
|
||||
# Calculate temporal proximity
|
||||
best_date = None
|
||||
if ep["occurred_start"] is not None and ep["occurred_end"] is not None:
|
||||
best_date = ep["occurred_start"] + (ep["occurred_end"] - ep["occurred_start"]) / 2
|
||||
elif ep["occurred_start"] is not None:
|
||||
best_date = ep["occurred_start"]
|
||||
elif ep["occurred_end"] is not None:
|
||||
best_date = ep["occurred_end"]
|
||||
elif ep["mentioned_at"] is not None:
|
||||
best_date = ep["mentioned_at"]
|
||||
|
||||
if best_date:
|
||||
days_from_mid = abs((best_date - mid_date).total_seconds() / 86400)
|
||||
temporal_proximity = 1.0 - min(days_from_mid / (total_days / 2), 1.0) if total_days > 0 else 1.0
|
||||
else:
|
||||
temporal_proximity = 0.5
|
||||
|
||||
ep_result = RetrievalResult.from_db_row(dict(ep))
|
||||
ep_result.temporal_score = temporal_proximity
|
||||
ep_result.temporal_proximity = temporal_proximity
|
||||
results.append(ep_result)
|
||||
node_scores[unit_id] = (ep["similarity"], 1.0)
|
||||
|
||||
# Spreading through temporal links (same as single-fact-type version)
|
||||
frontier = list(node_scores.keys())
|
||||
budget_remaining = budget - len(ft_entry_points)
|
||||
batch_size = 20
|
||||
# Per-source neighbor limit: lets the planner use the composite index
|
||||
# (from_unit_id, link_type, weight DESC) with early termination, avoiding
|
||||
# a full scan of all links from all source nodes before sorting.
|
||||
per_source_limit = 10
|
||||
# Safety cap on BFS iterations to prevent runaway spreading in dense graphs.
|
||||
max_iterations = 5
|
||||
iteration = 0
|
||||
|
||||
# Build tags clause for spreading (use param 7 since 1-6 are used)
|
||||
spreading_tags_clause = build_tags_where_clause_simple(tags, 7, table_alias="mu.", match=tags_match)
|
||||
spreading_groups_param_start = 7 + (1 if tags else 0)
|
||||
spreading_groups_clause, spreading_groups_params, _ = build_tag_groups_where_clause(
|
||||
tag_groups, spreading_groups_param_start, table_alias="mu."
|
||||
)
|
||||
|
||||
while frontier and budget_remaining > 0 and iteration < max_iterations:
|
||||
iteration += 1
|
||||
batch_ids = frontier[:batch_size]
|
||||
frontier = frontier[batch_size:]
|
||||
|
||||
# $1=query_emb, $2=batch_ids, $3=fact_type, $4=threshold, $5=per_source_limit, $6=bank_id, $7=tags, $M+=tag_groups
|
||||
spreading_params = [query_emb_str, batch_ids, ft, semantic_threshold, per_source_limit, bank_id]
|
||||
if tags:
|
||||
spreading_params.append(tags)
|
||||
spreading_params.extend(spreading_groups_params)
|
||||
|
||||
# LATERAL join: for each source node, fetch top-K neighbors by weight using
|
||||
# the existing idx_memory_links_from_type_weight index with early-exit semantics.
|
||||
# This avoids scanning all temporal links from all source nodes before sorting.
|
||||
# bank_id on memory_units lets the planner use idx_memory_units_bank_fact_type.
|
||||
neighbors = await conn.fetch(
|
||||
f"""
|
||||
SELECT src.from_unit_id, mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start, mu.occurred_end, mu.mentioned_at, mu.fact_type, mu.document_id, mu.chunk_id, mu.tags,
|
||||
l.weight, l.link_type,
|
||||
1 - (mu.embedding <=> $1::vector) AS similarity
|
||||
FROM unnest($2::uuid[]) AS src(from_unit_id)
|
||||
CROSS JOIN LATERAL (
|
||||
SELECT ml.to_unit_id, ml.weight, ml.link_type
|
||||
FROM {fq_table("memory_links")} ml
|
||||
WHERE ml.from_unit_id = src.from_unit_id
|
||||
AND ml.link_type IN ('temporal', 'causes', 'caused_by', 'enables', 'prevents')
|
||||
AND ml.weight >= 0.1
|
||||
ORDER BY ml.weight DESC
|
||||
LIMIT $5
|
||||
) l
|
||||
JOIN {fq_table("memory_units")} mu ON mu.id = l.to_unit_id
|
||||
WHERE mu.bank_id = $6
|
||||
AND mu.fact_type = $3
|
||||
AND mu.embedding IS NOT NULL
|
||||
AND (1 - (mu.embedding <=> $1::vector)) >= $4
|
||||
{spreading_tags_clause}
|
||||
{spreading_groups_clause}
|
||||
""",
|
||||
*spreading_params,
|
||||
)
|
||||
|
||||
for n in neighbors:
|
||||
neighbor_id = str(n["id"])
|
||||
if neighbor_id in visited:
|
||||
continue
|
||||
|
||||
visited.add(neighbor_id)
|
||||
budget_remaining -= 1
|
||||
|
||||
parent_id = str(n["from_unit_id"])
|
||||
_, parent_temporal_score = node_scores.get(parent_id, (0.5, 0.5))
|
||||
|
||||
neighbor_best_date = None
|
||||
if n["occurred_start"] is not None and n["occurred_end"] is not None:
|
||||
neighbor_best_date = n["occurred_start"] + (n["occurred_end"] - n["occurred_start"]) / 2
|
||||
elif n["occurred_start"] is not None:
|
||||
neighbor_best_date = n["occurred_start"]
|
||||
elif n["occurred_end"] is not None:
|
||||
neighbor_best_date = n["occurred_end"]
|
||||
elif n["mentioned_at"] is not None:
|
||||
neighbor_best_date = n["mentioned_at"]
|
||||
|
||||
if neighbor_best_date:
|
||||
days_from_mid = abs((neighbor_best_date - mid_date).total_seconds() / 86400)
|
||||
neighbor_temporal_proximity = (
|
||||
1.0 - min(days_from_mid / (total_days / 2), 1.0) if total_days > 0 else 1.0
|
||||
)
|
||||
else:
|
||||
neighbor_temporal_proximity = 0.3
|
||||
|
||||
link_type = n["link_type"]
|
||||
if link_type in ("causes", "caused_by"):
|
||||
causal_boost = 2.0
|
||||
elif link_type in ("enables", "prevents"):
|
||||
causal_boost = 1.5
|
||||
else:
|
||||
causal_boost = 1.0
|
||||
|
||||
propagated_temporal = parent_temporal_score * n["weight"] * causal_boost * 0.7
|
||||
combined_temporal = max(neighbor_temporal_proximity, propagated_temporal)
|
||||
|
||||
neighbor_result = RetrievalResult.from_db_row(dict(n))
|
||||
neighbor_result.temporal_score = combined_temporal
|
||||
neighbor_result.temporal_proximity = neighbor_temporal_proximity
|
||||
results.append(neighbor_result)
|
||||
|
||||
if budget_remaining > 0 and combined_temporal > 0.2:
|
||||
node_scores[neighbor_id] = (n["similarity"], combined_temporal)
|
||||
frontier.append(neighbor_id)
|
||||
|
||||
if budget_remaining <= 0:
|
||||
break
|
||||
|
||||
results_by_ft[ft] = results
|
||||
|
||||
return results_by_ft
|
||||
|
||||
|
||||
async def retrieve_all_fact_types_parallel(
|
||||
pool,
|
||||
query_text: str,
|
||||
query_embedding_str: str,
|
||||
bank_id: str,
|
||||
fact_types: list[str],
|
||||
thinking_budget: int,
|
||||
question_date: datetime | None = None,
|
||||
query_analyzer: Optional["QueryAnalyzer"] = None,
|
||||
graph_retriever: GraphRetriever | None = None,
|
||||
tags: list[str] | None = None,
|
||||
tags_match: TagsMatch = "any",
|
||||
tag_groups: list[TagGroup] | None = None,
|
||||
) -> MultiFactTypeRetrievalResult:
|
||||
"""
|
||||
Optimized retrieval for multiple fact types using batched queries.
|
||||
|
||||
This reduces database round-trips by:
|
||||
1. Combining semantic + BM25 into one CTE query for ALL fact types (1 query instead of 2N)
|
||||
2. Running graph retrieval per fact type in parallel (N parallel tasks)
|
||||
3. Running temporal retrieval per fact type in parallel (N parallel tasks)
|
||||
|
||||
Args:
|
||||
pool: Database connection pool
|
||||
query_text: Query text
|
||||
query_embedding_str: Query embedding as string
|
||||
bank_id: Bank ID
|
||||
fact_types: List of fact types to retrieve
|
||||
thinking_budget: Budget for graph traversal and retrieval limits
|
||||
question_date: Optional date when question was asked (for temporal filtering)
|
||||
query_analyzer: Query analyzer to use (defaults to TransformerQueryAnalyzer)
|
||||
graph_retriever: Graph retrieval strategy (defaults to configured retriever)
|
||||
|
||||
Returns:
|
||||
MultiFactTypeRetrievalResult with results organized by fact type
|
||||
"""
|
||||
import time
|
||||
|
||||
retriever = graph_retriever or get_default_graph_retriever()
|
||||
start_time = time.time()
|
||||
timings: dict[str, float] = {}
|
||||
|
||||
# Step 1: Extract temporal constraint first (CPU work, no DB)
|
||||
# Do this before DB queries so we know if we need temporal retrieval
|
||||
temporal_extraction_start = time.time()
|
||||
from .temporal_extraction import extract_temporal_constraint
|
||||
|
||||
temporal_constraint = extract_temporal_constraint(query_text, reference_date=question_date, analyzer=query_analyzer)
|
||||
temporal_extraction_time = time.time() - temporal_extraction_start
|
||||
timings["temporal_extraction"] = temporal_extraction_time
|
||||
|
||||
# Step 2: Run semantic + BM25 + temporal combined in ONE connection!
|
||||
# This reduces connection usage from 2 to 1 for these operations
|
||||
semantic_bm25_start = time.time()
|
||||
temporal_results_by_ft: dict[str, list[RetrievalResult]] = {}
|
||||
temporal_time = 0.0
|
||||
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
conn_wait = time.time() - semantic_bm25_start
|
||||
|
||||
# Semantic + BM25 combined
|
||||
semantic_bm25_results = await retrieve_semantic_bm25_combined(
|
||||
conn,
|
||||
query_embedding_str,
|
||||
query_text,
|
||||
bank_id,
|
||||
fact_types,
|
||||
thinking_budget,
|
||||
tags=tags,
|
||||
tags_match=tags_match,
|
||||
tag_groups=tag_groups,
|
||||
)
|
||||
semantic_bm25_time = time.time() - semantic_bm25_start
|
||||
|
||||
# Temporal combined (if constraint detected) - same connection!
|
||||
if temporal_constraint:
|
||||
tc_start, tc_end = temporal_constraint
|
||||
temporal_start = time.time()
|
||||
temporal_results_by_ft = await retrieve_temporal_combined(
|
||||
conn,
|
||||
query_embedding_str,
|
||||
bank_id,
|
||||
fact_types,
|
||||
tc_start,
|
||||
tc_end,
|
||||
budget=thinking_budget,
|
||||
semantic_threshold=0.1,
|
||||
tags=tags,
|
||||
tags_match=tags_match,
|
||||
tag_groups=tag_groups,
|
||||
)
|
||||
temporal_time = time.time() - temporal_start
|
||||
|
||||
timings["semantic_bm25_combined"] = semantic_bm25_time
|
||||
timings["temporal_combined"] = temporal_time
|
||||
|
||||
# Step 3: Run graph retrieval for each fact type in parallel
|
||||
async def run_graph_for_fact_type(ft: str) -> tuple[str, list[RetrievalResult], float, MPFPTimings | None]:
|
||||
graph_start = time.time()
|
||||
results, mpfp_timing = await retriever.retrieve(
|
||||
pool=pool,
|
||||
query_embedding_str=query_embedding_str,
|
||||
bank_id=bank_id,
|
||||
fact_type=ft,
|
||||
budget=thinking_budget,
|
||||
query_text=query_text,
|
||||
semantic_seeds=None,
|
||||
temporal_seeds=None,
|
||||
tags=tags,
|
||||
tags_match=tags_match,
|
||||
tag_groups=tag_groups,
|
||||
)
|
||||
return ft, results, time.time() - graph_start, mpfp_timing
|
||||
|
||||
# Run graph for all fact types in parallel
|
||||
graph_tasks = [run_graph_for_fact_type(ft) for ft in fact_types]
|
||||
graph_results_list = await asyncio.gather(*graph_tasks)
|
||||
|
||||
# Organize results by fact type
|
||||
results_by_fact_type: dict[str, ParallelRetrievalResult] = {}
|
||||
max_conn_wait = conn_wait # Single connection for semantic+bm25+temporal
|
||||
all_mpfp_timings: list[MPFPTimings] = []
|
||||
|
||||
for ft in fact_types:
|
||||
# Get semantic + bm25 results for this fact type
|
||||
semantic_results, bm25_results = semantic_bm25_results.get(ft, ([], []))
|
||||
|
||||
# Find graph results for this fact type
|
||||
graph_results = []
|
||||
graph_time = 0.0
|
||||
mpfp_timing = None
|
||||
for gr in graph_results_list:
|
||||
if gr[0] == ft:
|
||||
graph_results = gr[1]
|
||||
graph_time = gr[2]
|
||||
mpfp_timing = gr[3]
|
||||
if mpfp_timing:
|
||||
all_mpfp_timings.append(mpfp_timing)
|
||||
break
|
||||
|
||||
# Get temporal results for this fact type from combined result
|
||||
temporal_results = temporal_results_by_ft.get(ft) if temporal_constraint else None
|
||||
if temporal_results is not None and len(temporal_results) == 0:
|
||||
temporal_results = None
|
||||
|
||||
results_by_fact_type[ft] = ParallelRetrievalResult(
|
||||
semantic=semantic_results,
|
||||
bm25=bm25_results,
|
||||
graph=graph_results,
|
||||
temporal=temporal_results,
|
||||
timings={
|
||||
"semantic": semantic_bm25_time / 2, # Approximate split
|
||||
"bm25": semantic_bm25_time / 2,
|
||||
"graph": graph_time,
|
||||
"temporal": temporal_time, # Same for all fact types (single query)
|
||||
"temporal_extraction": temporal_extraction_time,
|
||||
},
|
||||
temporal_constraint=temporal_constraint,
|
||||
mpfp_timings=[mpfp_timing] if mpfp_timing else [],
|
||||
max_conn_wait=max_conn_wait,
|
||||
)
|
||||
|
||||
total_time = time.time() - start_time
|
||||
timings["total"] = total_time
|
||||
|
||||
return MultiFactTypeRetrievalResult(
|
||||
results_by_fact_type=results_by_fact_type,
|
||||
timings=timings,
|
||||
max_conn_wait=max_conn_wait,
|
||||
)
|
||||
@@ -1,390 +0,0 @@
|
||||
"""
|
||||
Tags filtering utilities for retrieval.
|
||||
|
||||
Provides SQL building functions for filtering memories by tags.
|
||||
Supports four matching modes via TagsMatch enum:
|
||||
- "any": OR matching, includes untagged memories (default, backward compatible)
|
||||
- "all": AND matching, includes untagged memories
|
||||
- "any_strict": OR matching, excludes untagged memories
|
||||
- "all_strict": AND matching, excludes untagged memories
|
||||
|
||||
OR matching (any/any_strict): Memory matches if ANY of its tags overlap with request tags
|
||||
AND matching (all/all_strict): Memory matches if ALL request tags are present in its tags
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Annotated, Literal
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
|
||||
TagsMatch = Literal["any", "all", "any_strict", "all_strict"]
|
||||
|
||||
|
||||
def _parse_tags_match(match: TagsMatch) -> tuple[str, bool]:
|
||||
"""
|
||||
Parse TagsMatch into operator and include_untagged flag.
|
||||
|
||||
Returns:
|
||||
Tuple of (operator, include_untagged)
|
||||
- operator: "&&" for any/any_strict, "@>" for all/all_strict
|
||||
- include_untagged: True for any/all, False for any_strict/all_strict
|
||||
"""
|
||||
if match == "any":
|
||||
return "&&", True
|
||||
elif match == "all":
|
||||
return "@>", True
|
||||
elif match == "any_strict":
|
||||
return "&&", False
|
||||
elif match == "all_strict":
|
||||
return "@>", False
|
||||
else:
|
||||
# Default to "any" behavior
|
||||
return "&&", True
|
||||
|
||||
|
||||
def build_tags_where_clause(
|
||||
tags: list[str] | None,
|
||||
param_offset: int = 1,
|
||||
table_alias: str = "",
|
||||
match: TagsMatch = "any",
|
||||
) -> tuple[str, list, int]:
|
||||
"""
|
||||
Build a SQL WHERE clause for filtering by tags.
|
||||
|
||||
Supports four matching modes:
|
||||
- "any" (default): OR matching, includes untagged memories
|
||||
- "all": AND matching, includes untagged memories
|
||||
- "any_strict": OR matching, excludes untagged memories
|
||||
- "all_strict": AND matching, excludes untagged memories
|
||||
|
||||
Args:
|
||||
tags: List of tags to filter by. If None or empty, returns empty clause (no filtering).
|
||||
param_offset: Starting parameter number for SQL placeholders (default 1).
|
||||
table_alias: Optional table alias prefix (e.g., "mu." for "memory_units mu").
|
||||
match: Matching mode. Defaults to "any".
|
||||
|
||||
Returns:
|
||||
Tuple of (sql_clause, params, next_param_offset):
|
||||
- sql_clause: SQL WHERE clause string
|
||||
- params: List of parameter values to bind
|
||||
- next_param_offset: Next available parameter number
|
||||
|
||||
Example:
|
||||
>>> clause, params, next_offset = build_tags_where_clause(['user_a'], 3, 'mu.', 'any_strict')
|
||||
>>> print(clause) # "AND mu.tags IS NOT NULL AND mu.tags != '{}' AND mu.tags && $3"
|
||||
"""
|
||||
if not tags:
|
||||
return "", [], param_offset
|
||||
|
||||
column = f"{table_alias}tags" if table_alias else "tags"
|
||||
operator, include_untagged = _parse_tags_match(match)
|
||||
|
||||
if include_untagged:
|
||||
# Include untagged memories (NULL or empty array) OR matching tags
|
||||
clause = f"AND ({column} IS NULL OR {column} = '{{}}' OR {column} {operator} ${param_offset})"
|
||||
else:
|
||||
# Strict: only memories with matching tags (exclude NULL and empty)
|
||||
clause = f"AND {column} IS NOT NULL AND {column} != '{{}}' AND {column} {operator} ${param_offset}"
|
||||
|
||||
return clause, [tags], param_offset + 1
|
||||
|
||||
|
||||
def build_tags_where_clause_simple(
|
||||
tags: list[str] | None,
|
||||
param_num: int,
|
||||
table_alias: str = "",
|
||||
match: TagsMatch = "any",
|
||||
) -> str:
|
||||
"""
|
||||
Build a simple SQL WHERE clause for tags filtering.
|
||||
|
||||
This is a convenience version that returns just the clause string,
|
||||
assuming the caller will add the tags array to their params list.
|
||||
|
||||
Args:
|
||||
tags: List of tags to filter by. If None or empty, returns empty string.
|
||||
param_num: Parameter number to use in the clause.
|
||||
table_alias: Optional table alias prefix.
|
||||
match: Matching mode. Defaults to "any".
|
||||
|
||||
Returns:
|
||||
SQL clause string or empty string.
|
||||
"""
|
||||
if not tags:
|
||||
return ""
|
||||
|
||||
column = f"{table_alias}tags" if table_alias else "tags"
|
||||
operator, include_untagged = _parse_tags_match(match)
|
||||
|
||||
if include_untagged:
|
||||
# Include untagged memories (NULL or empty array) OR matching tags
|
||||
return f"AND ({column} IS NULL OR {column} = '{{}}' OR {column} {operator} ${param_num})"
|
||||
else:
|
||||
# Strict: only memories with matching tags (exclude NULL and empty)
|
||||
return f"AND {column} IS NOT NULL AND {column} != '{{}}' AND {column} {operator} ${param_num}"
|
||||
|
||||
|
||||
def filter_results_by_tags(
|
||||
results: list,
|
||||
tags: list[str] | None,
|
||||
match: TagsMatch = "any",
|
||||
) -> list:
|
||||
"""
|
||||
Filter retrieval results by tags in Python (for post-processing).
|
||||
|
||||
Used when SQL filtering isn't possible (e.g., graph traversal results).
|
||||
|
||||
Args:
|
||||
results: List of RetrievalResult objects with a 'tags' attribute.
|
||||
tags: List of tags to filter by. If None or empty, returns all results.
|
||||
match: Matching mode. Defaults to "any".
|
||||
|
||||
Returns:
|
||||
Filtered list of results.
|
||||
"""
|
||||
if not tags:
|
||||
return results
|
||||
|
||||
_, include_untagged = _parse_tags_match(match)
|
||||
is_any_match = match in ("any", "any_strict")
|
||||
|
||||
tags_set = set(tags)
|
||||
filtered = []
|
||||
|
||||
for result in results:
|
||||
result_tags = getattr(result, "tags", None)
|
||||
|
||||
# Check if untagged
|
||||
is_untagged = result_tags is None or len(result_tags) == 0
|
||||
|
||||
if is_untagged:
|
||||
if include_untagged:
|
||||
filtered.append(result)
|
||||
# else: skip untagged
|
||||
else:
|
||||
result_tags_set = set(result_tags)
|
||||
if is_any_match:
|
||||
# Any overlap
|
||||
if result_tags_set & tags_set:
|
||||
filtered.append(result)
|
||||
else:
|
||||
# All tags must be present
|
||||
if tags_set <= result_tags_set:
|
||||
filtered.append(result)
|
||||
|
||||
return filtered
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Compound tag group models (recursive boolean expressions)
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TagGroupLeaf(BaseModel):
|
||||
"""A leaf tag filter: matches memories by tag list and match mode."""
|
||||
|
||||
tags: list[str]
|
||||
match: TagsMatch = "any_strict"
|
||||
|
||||
|
||||
class TagGroupAnd(BaseModel):
|
||||
"""Compound AND group: all child filters must match."""
|
||||
|
||||
model_config = ConfigDict(populate_by_name=True)
|
||||
filters: list[TagGroup] = Field(alias="and")
|
||||
|
||||
|
||||
class TagGroupOr(BaseModel):
|
||||
"""Compound OR group: at least one child filter must match."""
|
||||
|
||||
model_config = ConfigDict(populate_by_name=True)
|
||||
filters: list[TagGroup] = Field(alias="or")
|
||||
|
||||
|
||||
class TagGroupNot(BaseModel):
|
||||
"""Compound NOT group: child filter must NOT match."""
|
||||
|
||||
model_config = ConfigDict(populate_by_name=True)
|
||||
filter: TagGroup = Field(alias="not")
|
||||
|
||||
|
||||
# TagGroup is a discriminated union; Pydantic will try left-to-right.
|
||||
# TagGroupLeaf is identified by the presence of 'tags'.
|
||||
# TagGroupAnd / TagGroupOr / TagGroupNot are compound (no 'tags' key).
|
||||
TagGroup = Annotated[
|
||||
TagGroupLeaf | TagGroupAnd | TagGroupOr | TagGroupNot,
|
||||
Field(union_mode="left_to_right"),
|
||||
]
|
||||
|
||||
# Rebuild forward-reference models so recursive TagGroup is resolved.
|
||||
TagGroupAnd.model_rebuild()
|
||||
TagGroupOr.model_rebuild()
|
||||
TagGroupNot.model_rebuild()
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# SQL builder for compound tag groups
|
||||
# =============================================================================
|
||||
|
||||
|
||||
def _build_group_clause(
|
||||
group: TagGroup,
|
||||
param_offset: int,
|
||||
table_alias: str,
|
||||
) -> tuple[str, list, int]:
|
||||
"""
|
||||
Recursively build an inner SQL clause (no leading AND/OR) for a single TagGroup.
|
||||
|
||||
Returns:
|
||||
(inner_clause, params, next_param_offset)
|
||||
"""
|
||||
if isinstance(group, TagGroupLeaf):
|
||||
column = f"{table_alias}tags" if table_alias else "tags"
|
||||
operator, include_untagged = _parse_tags_match(group.match)
|
||||
if include_untagged:
|
||||
clause = f"({column} IS NULL OR {column} = '{{}}' OR {column} {operator} ${param_offset})"
|
||||
else:
|
||||
clause = f"({column} IS NOT NULL AND {column} != '{{}}' AND {column} {operator} ${param_offset})"
|
||||
return clause, [group.tags], param_offset + 1
|
||||
|
||||
elif isinstance(group, TagGroupAnd):
|
||||
parts = []
|
||||
params: list = []
|
||||
offset = param_offset
|
||||
for child in group.filters:
|
||||
child_clause, child_params, offset = _build_group_clause(child, offset, table_alias)
|
||||
parts.append(child_clause)
|
||||
params.extend(child_params)
|
||||
inner = " AND ".join(parts)
|
||||
return f"({inner})", params, offset
|
||||
|
||||
elif isinstance(group, TagGroupOr):
|
||||
parts = []
|
||||
params = []
|
||||
offset = param_offset
|
||||
for child in group.filters:
|
||||
child_clause, child_params, offset = _build_group_clause(child, offset, table_alias)
|
||||
parts.append(child_clause)
|
||||
params.extend(child_params)
|
||||
inner = " OR ".join(parts)
|
||||
return f"({inner})", params, offset
|
||||
|
||||
elif isinstance(group, TagGroupNot):
|
||||
child_clause, child_params, next_offset = _build_group_clause(group.filter, param_offset, table_alias)
|
||||
return f"NOT {child_clause}", child_params, next_offset
|
||||
|
||||
else:
|
||||
# Should never happen with proper Pydantic validation
|
||||
return "", [], param_offset
|
||||
|
||||
|
||||
def build_tag_groups_where_clause(
|
||||
tag_groups: list[TagGroup] | None,
|
||||
param_offset: int,
|
||||
table_alias: str = "",
|
||||
) -> tuple[str, list, int]:
|
||||
"""
|
||||
Build a SQL WHERE clause for compound tag group filtering.
|
||||
|
||||
Top-level groups are AND-ed together. Each group is a recursive boolean
|
||||
expression (leaf, and, or, not).
|
||||
|
||||
Args:
|
||||
tag_groups: List of TagGroup objects. If None or empty, returns empty clause.
|
||||
param_offset: Starting parameter number for SQL placeholders.
|
||||
table_alias: Optional table alias prefix (e.g., "mu." for "memory_units mu").
|
||||
|
||||
Returns:
|
||||
Tuple of (sql_clause, params, next_param_offset):
|
||||
- sql_clause: SQL WHERE clause string starting with "AND" (or empty string)
|
||||
- params: List of parameter values to bind (one per leaf node)
|
||||
- next_param_offset: Next available parameter number
|
||||
|
||||
Example:
|
||||
>>> groups = [TagGroupLeaf(tags=["user:alice"], match="all_strict")]
|
||||
>>> clause, params, next_offset = build_tag_groups_where_clause(groups, 3)
|
||||
>>> print(clause) # "AND (tags IS NOT NULL AND tags != '{}' AND tags @> $3)"
|
||||
"""
|
||||
if not tag_groups:
|
||||
return "", [], param_offset
|
||||
|
||||
all_params: list = []
|
||||
all_clauses: list[str] = []
|
||||
offset = param_offset
|
||||
|
||||
for group in tag_groups:
|
||||
inner_clause, group_params, offset = _build_group_clause(group, offset, table_alias)
|
||||
all_clauses.append(inner_clause)
|
||||
all_params.extend(group_params)
|
||||
|
||||
combined = " AND ".join(all_clauses)
|
||||
return f"AND {combined}", all_params, offset
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Python-side filter for compound tag groups (post-retrieval filtering)
|
||||
# =============================================================================
|
||||
|
||||
|
||||
def _match_group(result: object, group: TagGroup) -> bool:
|
||||
"""
|
||||
Recursively evaluate a TagGroup against a retrieval result.
|
||||
|
||||
Args:
|
||||
result: Any object with a 'tags' attribute (list[str] or None).
|
||||
group: The TagGroup to evaluate.
|
||||
|
||||
Returns:
|
||||
True if the result matches the group, False otherwise.
|
||||
"""
|
||||
if isinstance(group, TagGroupLeaf):
|
||||
result_tags = getattr(result, "tags", None)
|
||||
is_untagged = result_tags is None or len(result_tags) == 0
|
||||
_, include_untagged = _parse_tags_match(group.match)
|
||||
is_any_match = group.match in ("any", "any_strict")
|
||||
tags_set = set(group.tags)
|
||||
|
||||
if is_untagged:
|
||||
return include_untagged
|
||||
else:
|
||||
result_tags_set = set(result_tags)
|
||||
if is_any_match:
|
||||
return bool(result_tags_set & tags_set)
|
||||
else:
|
||||
return tags_set <= result_tags_set
|
||||
|
||||
elif isinstance(group, TagGroupAnd):
|
||||
return all(_match_group(result, child) for child in group.filters)
|
||||
|
||||
elif isinstance(group, TagGroupOr):
|
||||
return any(_match_group(result, child) for child in group.filters)
|
||||
|
||||
elif isinstance(group, TagGroupNot):
|
||||
return not _match_group(result, group.filter)
|
||||
|
||||
else:
|
||||
return True
|
||||
|
||||
|
||||
def filter_results_by_tag_groups(
|
||||
results: list,
|
||||
tag_groups: list[TagGroup] | None,
|
||||
) -> list:
|
||||
"""
|
||||
Filter retrieval results by compound tag groups in Python (for post-processing).
|
||||
|
||||
Used when SQL filtering isn't possible (e.g., graph traversal results).
|
||||
Top-level groups are AND-ed together.
|
||||
|
||||
Args:
|
||||
results: List of RetrievalResult objects with a 'tags' attribute.
|
||||
tag_groups: List of TagGroup objects. If None or empty, returns all results.
|
||||
|
||||
Returns:
|
||||
Filtered list of results where ALL top-level groups match.
|
||||
"""
|
||||
if not tag_groups:
|
||||
return results
|
||||
|
||||
return [r for r in results if all(_match_group(r, group) for group in tag_groups)]
|
||||
@@ -1,79 +0,0 @@
|
||||
"""File storage backends for uploaded files."""
|
||||
|
||||
from collections.abc import Callable
|
||||
|
||||
from .base import FileStorage
|
||||
from .postgresql import PostgreSQLFileStorage
|
||||
|
||||
__all__ = ["FileStorage", "PostgreSQLFileStorage", "create_file_storage"]
|
||||
|
||||
|
||||
def create_file_storage(
|
||||
storage_type: str,
|
||||
pool_getter: Callable | None = None,
|
||||
schema: str | None = None,
|
||||
schema_getter: Callable | None = None,
|
||||
**kwargs,
|
||||
) -> FileStorage:
|
||||
"""
|
||||
Create file storage backend based on configuration.
|
||||
|
||||
Args:
|
||||
storage_type: "native" (PostgreSQL BYTEA) or "s3" (S3-compatible object storage)
|
||||
pool_getter: Database pool getter (required for native)
|
||||
schema: Static database schema (for native single-tenant)
|
||||
schema_getter: Callable returning current schema at query time (for native multi-tenant)
|
||||
**kwargs: Additional args passed to storage backend
|
||||
|
||||
Returns:
|
||||
FileStorage instance
|
||||
|
||||
Raises:
|
||||
ValueError: If storage_type is unknown or required args are missing
|
||||
"""
|
||||
if storage_type == "native":
|
||||
if not pool_getter:
|
||||
raise ValueError("pool_getter required for native (PostgreSQL) storage")
|
||||
return PostgreSQLFileStorage(pool_getter=pool_getter, schema=schema, schema_getter=schema_getter)
|
||||
elif storage_type == "s3":
|
||||
from ...config import get_config
|
||||
from .s3 import S3FileStorage
|
||||
|
||||
config = get_config()
|
||||
bucket = config.file_storage_s3_bucket
|
||||
if not bucket:
|
||||
raise ValueError("HINDSIGHT_API_FILE_STORAGE_S3_BUCKET is required for S3 storage")
|
||||
return S3FileStorage(
|
||||
bucket=bucket,
|
||||
region=config.file_storage_s3_region,
|
||||
endpoint=config.file_storage_s3_endpoint,
|
||||
access_key_id=config.file_storage_s3_access_key_id,
|
||||
secret_access_key=config.file_storage_s3_secret_access_key,
|
||||
)
|
||||
elif storage_type == "gcs":
|
||||
from ...config import get_config
|
||||
from .gcs import GCSFileStorage
|
||||
|
||||
config = get_config()
|
||||
bucket = config.file_storage_gcs_bucket
|
||||
if not bucket:
|
||||
raise ValueError("HINDSIGHT_API_FILE_STORAGE_GCS_BUCKET is required for GCS storage")
|
||||
return GCSFileStorage(
|
||||
bucket=bucket,
|
||||
service_account_key=config.file_storage_gcs_service_account_key,
|
||||
)
|
||||
elif storage_type == "azure":
|
||||
from ...config import get_config
|
||||
from .azure import AzureFileStorage
|
||||
|
||||
config = get_config()
|
||||
container = config.file_storage_azure_container
|
||||
if not container:
|
||||
raise ValueError("HINDSIGHT_API_FILE_STORAGE_AZURE_CONTAINER is required for Azure storage")
|
||||
return AzureFileStorage(
|
||||
container_name=container,
|
||||
account_name=config.file_storage_azure_account_name,
|
||||
account_key=config.file_storage_azure_account_key,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unknown storage type: {storage_type}. Supported: 'native', 's3', 'gcs', 'azure'.")
|
||||
@@ -1,62 +0,0 @@
|
||||
"""Azure Blob Storage backend using obstore."""
|
||||
|
||||
import logging
|
||||
from datetime import timedelta
|
||||
|
||||
import obstore as obs
|
||||
from obstore.store import AzureStore
|
||||
|
||||
from .base import FileStorage
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class AzureFileStorage(FileStorage):
|
||||
"""
|
||||
Azure Blob Storage backend.
|
||||
|
||||
Uses obstore (Rust-backed) for high-throughput async access to Azure Blob Storage.
|
||||
Supports account key, SAS token, and default Azure credentials.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
container_name: str,
|
||||
account_name: str | None = None,
|
||||
account_key: str | None = None,
|
||||
):
|
||||
kwargs: dict = {}
|
||||
if account_name:
|
||||
kwargs["account_name"] = account_name
|
||||
if account_key:
|
||||
kwargs["account_key"] = account_key
|
||||
|
||||
self._store = AzureStore(container_name, **kwargs)
|
||||
logger.info(f"Initialized Azure file storage: container={container_name}, account={account_name}")
|
||||
|
||||
async def store(self, file_data: bytes, key: str, metadata: dict[str, str] | None = None) -> str:
|
||||
await obs.put_async(self._store, key, file_data)
|
||||
logger.debug(f"Stored file {key} ({len(file_data)} bytes) in Azure")
|
||||
return key
|
||||
|
||||
async def retrieve(self, key: str) -> bytes:
|
||||
try:
|
||||
response = await obs.get_async(self._store, key)
|
||||
return await response.bytes_async()
|
||||
except Exception as e:
|
||||
if "not found" in str(e).lower() or "BlobNotFound" in str(e):
|
||||
raise FileNotFoundError(f"File not found: {key}") from e
|
||||
raise
|
||||
|
||||
async def delete(self, key: str) -> None:
|
||||
await obs.delete_async(self._store, key)
|
||||
|
||||
async def exists(self, key: str) -> bool:
|
||||
try:
|
||||
await obs.head_async(self._store, key)
|
||||
return True
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
async def get_download_url(self, key: str, expires_in: int = 3600) -> str:
|
||||
return await obs.sign_async(self._store, "GET", key, timedelta(seconds=expires_in))
|
||||
@@ -1,83 +0,0 @@
|
||||
"""Abstract base class for file storage backends."""
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
|
||||
|
||||
class FileStorage(ABC):
|
||||
"""Abstract base for file storage backends."""
|
||||
|
||||
@abstractmethod
|
||||
async def store(
|
||||
self,
|
||||
file_data: bytes,
|
||||
key: str,
|
||||
metadata: dict[str, str] | None = None,
|
||||
) -> str:
|
||||
"""
|
||||
Store file and return storage key.
|
||||
|
||||
Args:
|
||||
file_data: Raw file bytes
|
||||
key: Storage key (e.g., "banks/{bank_id}/files/{file_id}.pdf")
|
||||
metadata: Optional metadata to store with file
|
||||
|
||||
Returns:
|
||||
Storage key that can be used to retrieve the file
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def retrieve(self, key: str) -> bytes:
|
||||
"""
|
||||
Retrieve file by storage key.
|
||||
|
||||
Args:
|
||||
key: Storage key
|
||||
|
||||
Returns:
|
||||
File data as bytes
|
||||
|
||||
Raises:
|
||||
FileNotFoundError: If file does not exist
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def delete(self, key: str) -> None:
|
||||
"""
|
||||
Delete file by storage key.
|
||||
|
||||
Args:
|
||||
key: Storage key
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def exists(self, key: str) -> bool:
|
||||
"""
|
||||
Check if file exists.
|
||||
|
||||
Args:
|
||||
key: Storage key
|
||||
|
||||
Returns:
|
||||
True if file exists, False otherwise
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def get_download_url(self, key: str, expires_in: int = 3600) -> str:
|
||||
"""
|
||||
Get a URL for downloading the file.
|
||||
|
||||
For PostgreSQL storage, this might be a relative API path.
|
||||
For S3, this would be a pre-signed URL.
|
||||
|
||||
Args:
|
||||
key: Storage key
|
||||
expires_in: Expiration time in seconds (may be ignored for some backends)
|
||||
|
||||
Returns:
|
||||
Download URL or path
|
||||
"""
|
||||
pass
|
||||
@@ -1,105 +0,0 @@
|
||||
"""Google Cloud Storage backend using obstore."""
|
||||
|
||||
import logging
|
||||
import os
|
||||
from datetime import datetime, timedelta, timezone
|
||||
|
||||
import obstore as obs
|
||||
from obstore.store import GCSStore
|
||||
|
||||
from .base import FileStorage
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _make_google_auth_credential_provider():
|
||||
"""Create a credential provider using google.auth (supports all credential types).
|
||||
|
||||
obstore's built-in credential parsing only supports service_account and
|
||||
authorized_user JSON types. This provider uses the google-auth library
|
||||
which additionally handles external_account (Workload Identity Federation),
|
||||
impersonated credentials, and metadata-server credentials.
|
||||
"""
|
||||
import google.auth
|
||||
import google.auth.transport.requests
|
||||
|
||||
credentials, _ = google.auth.default(scopes=["https://www.googleapis.com/auth/cloud-platform"])
|
||||
request = google.auth.transport.requests.Request()
|
||||
|
||||
def _provide():
|
||||
credentials.refresh(request)
|
||||
expiry = credentials.expiry
|
||||
if expiry and expiry.tzinfo is None:
|
||||
expiry = expiry.replace(tzinfo=timezone.utc)
|
||||
return {"token": credentials.token, "expires_at": expiry}
|
||||
|
||||
return _provide
|
||||
|
||||
|
||||
class GCSFileStorage(FileStorage):
|
||||
"""
|
||||
Google Cloud Storage backend.
|
||||
|
||||
Uses obstore (Rust-backed) for high-throughput async access to GCS.
|
||||
Supports Application Default Credentials, service account keys, and explicit credentials.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
bucket: str,
|
||||
service_account_key: str | None = None,
|
||||
):
|
||||
kwargs: dict = {}
|
||||
if service_account_key:
|
||||
kwargs["service_account_key"] = service_account_key
|
||||
else:
|
||||
# Use google.auth credential provider for broad credential type support
|
||||
# (service_account, authorized_user, external_account, metadata server, etc.)
|
||||
try:
|
||||
kwargs["credential_provider"] = _make_google_auth_credential_provider()
|
||||
logger.info("Using google.auth credential provider for GCS")
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
f"Failed to create google.auth credential provider, falling back to obstore defaults: {e}"
|
||||
)
|
||||
|
||||
# Workaround for https://github.com/developmentseed/obstore/issues/605
|
||||
# obstore's Rust layer doesn't support external_account credentials (Workload
|
||||
# Identity Federation) and eagerly parses GOOGLE_APPLICATION_CREDENTIALS even
|
||||
# when credential_provider is given. Per the obstore maintainer's guidance,
|
||||
# remove env vars so the Rust code doesn't try to authenticate itself.
|
||||
# google.auth (used by credential_provider above) has already loaded credentials.
|
||||
gac = os.environ.pop("GOOGLE_APPLICATION_CREDENTIALS", None)
|
||||
try:
|
||||
self._store = GCSStore(bucket, **kwargs)
|
||||
finally:
|
||||
if gac is not None:
|
||||
os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = gac
|
||||
logger.info(f"Initialized GCS file storage: bucket={bucket}")
|
||||
|
||||
async def store(self, file_data: bytes, key: str, metadata: dict[str, str] | None = None) -> str:
|
||||
await obs.put_async(self._store, key, file_data)
|
||||
logger.debug(f"Stored file {key} ({len(file_data)} bytes) in GCS")
|
||||
return key
|
||||
|
||||
async def retrieve(self, key: str) -> bytes:
|
||||
try:
|
||||
response = await obs.get_async(self._store, key)
|
||||
return await response.bytes_async()
|
||||
except Exception as e:
|
||||
if "not found" in str(e).lower():
|
||||
raise FileNotFoundError(f"File not found: {key}") from e
|
||||
raise
|
||||
|
||||
async def delete(self, key: str) -> None:
|
||||
await obs.delete_async(self._store, key)
|
||||
|
||||
async def exists(self, key: str) -> bool:
|
||||
try:
|
||||
await obs.head_async(self._store, key)
|
||||
return True
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
async def get_download_url(self, key: str, expires_in: int = 3600) -> str:
|
||||
return await obs.sign_async(self._store, "GET", key, timedelta(seconds=expires_in))
|
||||
@@ -1,153 +0,0 @@
|
||||
"""PostgreSQL BYTEA-based file storage (default, zero-config)."""
|
||||
|
||||
import logging
|
||||
from collections.abc import Callable
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import asyncpg
|
||||
|
||||
from .base import FileStorage
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def fq_table(table: str, schema: str | None = None) -> str:
|
||||
"""Get fully-qualified table name with optional schema prefix."""
|
||||
if schema:
|
||||
return f'"{schema}".{table}'
|
||||
return table
|
||||
|
||||
|
||||
class PostgreSQLFileStorage(FileStorage):
|
||||
"""
|
||||
PostgreSQL BYTEA-based file storage.
|
||||
|
||||
Stores files directly in PostgreSQL using BYTEA columns.
|
||||
This is the default storage backend - zero configuration required!
|
||||
|
||||
Pros:
|
||||
- Works out of the box (no external dependencies)
|
||||
- Transactional consistency with database
|
||||
- Simple backups (included in pg_dump)
|
||||
- Good performance for <10MB files
|
||||
|
||||
Cons:
|
||||
- Database bloat for large/many files
|
||||
- Not ideal for distributed deployments
|
||||
- Higher cost than object storage at scale
|
||||
|
||||
For production/scale, consider S3FileStorage instead.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
pool_getter: Callable[[], "asyncpg.Pool"],
|
||||
schema: str | None = None,
|
||||
schema_getter: Callable[[], str] | None = None,
|
||||
):
|
||||
"""
|
||||
Initialize PostgreSQL file storage.
|
||||
|
||||
Args:
|
||||
pool_getter: Function that returns asyncpg connection pool
|
||||
schema: Static database schema (fallback for single-tenant / tests)
|
||||
schema_getter: Callable returning current schema at query time (for multi-tenant)
|
||||
"""
|
||||
self._pool_getter = pool_getter
|
||||
self._static_schema = schema
|
||||
self._schema_getter = schema_getter
|
||||
|
||||
@property
|
||||
def _schema(self) -> str | None:
|
||||
"""Resolve schema dynamically per-request when schema_getter is provided."""
|
||||
if self._schema_getter:
|
||||
return self._schema_getter()
|
||||
return self._static_schema
|
||||
|
||||
async def store(
|
||||
self,
|
||||
file_data: bytes,
|
||||
key: str,
|
||||
metadata: dict[str, str] | None = None,
|
||||
) -> str:
|
||||
"""Store file in PostgreSQL."""
|
||||
pool = self._pool_getter()
|
||||
|
||||
async with pool.acquire() as conn:
|
||||
await conn.execute(
|
||||
f"""
|
||||
INSERT INTO {fq_table("file_storage", self._schema)}
|
||||
(storage_key, data)
|
||||
VALUES ($1, $2)
|
||||
ON CONFLICT (storage_key) DO UPDATE SET
|
||||
data = EXCLUDED.data
|
||||
""",
|
||||
key,
|
||||
file_data,
|
||||
)
|
||||
|
||||
logger.debug(f"Stored file {key} ({len(file_data)} bytes) in PostgreSQL")
|
||||
return key
|
||||
|
||||
async def retrieve(self, key: str) -> bytes:
|
||||
"""Retrieve file from PostgreSQL."""
|
||||
pool = self._pool_getter()
|
||||
|
||||
async with pool.acquire() as conn:
|
||||
row = await conn.fetchrow(
|
||||
f"""
|
||||
SELECT data FROM {fq_table("file_storage", self._schema)}
|
||||
WHERE storage_key = $1
|
||||
""",
|
||||
key,
|
||||
)
|
||||
|
||||
if not row:
|
||||
raise FileNotFoundError(f"File not found: {key}")
|
||||
|
||||
return bytes(row["data"])
|
||||
|
||||
async def delete(self, key: str) -> None:
|
||||
"""Delete file from PostgreSQL."""
|
||||
pool = self._pool_getter()
|
||||
|
||||
async with pool.acquire() as conn:
|
||||
result = await conn.execute(
|
||||
f"""
|
||||
DELETE FROM {fq_table("file_storage", self._schema)}
|
||||
WHERE storage_key = $1
|
||||
""",
|
||||
key,
|
||||
)
|
||||
|
||||
# Check if anything was deleted
|
||||
if result == "DELETE 0":
|
||||
logger.warning(f"Attempted to delete non-existent file: {key}")
|
||||
|
||||
async def exists(self, key: str) -> bool:
|
||||
"""Check if file exists in PostgreSQL."""
|
||||
pool = self._pool_getter()
|
||||
|
||||
async with pool.acquire() as conn:
|
||||
row = await conn.fetchrow(
|
||||
f"""
|
||||
SELECT 1 FROM {fq_table("file_storage", self._schema)}
|
||||
WHERE storage_key = $1
|
||||
""",
|
||||
key,
|
||||
)
|
||||
|
||||
return row is not None
|
||||
|
||||
async def get_download_url(self, key: str, expires_in: int = 3600) -> str:
|
||||
"""
|
||||
Get download URL for PostgreSQL-stored file.
|
||||
|
||||
Returns an API endpoint path (not a pre-signed URL since the file
|
||||
is stored in the database). The expires_in parameter is ignored
|
||||
for PostgreSQL storage.
|
||||
"""
|
||||
# Return API path for download endpoint
|
||||
# (expires_in ignored for database storage - auth handled at API level)
|
||||
return f"/v1/default/files/download/{key}"
|
||||
@@ -1,71 +0,0 @@
|
||||
"""S3 object storage backend using obstore."""
|
||||
|
||||
import logging
|
||||
from datetime import timedelta
|
||||
|
||||
import obstore as obs
|
||||
from obstore.store import S3Store
|
||||
|
||||
from .base import FileStorage
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class S3FileStorage(FileStorage):
|
||||
"""
|
||||
S3-compatible object storage backend.
|
||||
|
||||
Uses obstore (Rust-backed) for high-throughput async access to
|
||||
Amazon S3, MinIO, Cloudflare R2, and other S3-compliant APIs.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
bucket: str,
|
||||
region: str | None = None,
|
||||
endpoint: str | None = None,
|
||||
access_key_id: str | None = None,
|
||||
secret_access_key: str | None = None,
|
||||
):
|
||||
kwargs: dict = {}
|
||||
if region:
|
||||
kwargs["region"] = region
|
||||
if endpoint:
|
||||
kwargs["endpoint"] = endpoint
|
||||
# Allow plain HTTP for local S3-compatible services (MinIO, LocalStack, etc.)
|
||||
if endpoint.startswith("http://"):
|
||||
kwargs["allow_http"] = True
|
||||
if access_key_id:
|
||||
kwargs["access_key_id"] = access_key_id
|
||||
if secret_access_key:
|
||||
kwargs["secret_access_key"] = secret_access_key
|
||||
|
||||
self._store = S3Store(bucket, **kwargs)
|
||||
logger.info(f"Initialized S3 file storage: bucket={bucket}, region={region}, endpoint={endpoint}")
|
||||
|
||||
async def store(self, file_data: bytes, key: str, metadata: dict[str, str] | None = None) -> str:
|
||||
await obs.put_async(self._store, key, file_data)
|
||||
logger.debug(f"Stored file {key} ({len(file_data)} bytes) in S3")
|
||||
return key
|
||||
|
||||
async def retrieve(self, key: str) -> bytes:
|
||||
try:
|
||||
response = await obs.get_async(self._store, key)
|
||||
return await response.bytes_async()
|
||||
except Exception as e:
|
||||
if "not found" in str(e).lower() or "NoSuchKey" in str(e):
|
||||
raise FileNotFoundError(f"File not found: {key}") from e
|
||||
raise
|
||||
|
||||
async def delete(self, key: str) -> None:
|
||||
await obs.delete_async(self._store, key)
|
||||
|
||||
async def exists(self, key: str) -> bool:
|
||||
try:
|
||||
await obs.head_async(self._store, key)
|
||||
return True
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
async def get_download_url(self, key: str, expires_in: int = 3600) -> str:
|
||||
return await obs.sign_async(self._store, "GET", key, timedelta(seconds=expires_in))
|
||||
@@ -1,433 +0,0 @@
|
||||
"""
|
||||
Supabase Tenant Extension for Hindsight
|
||||
|
||||
Validates Supabase JWTs and maps authenticated users to isolated memory banks.
|
||||
Each user gets their own PostgreSQL schema based on their Supabase user ID.
|
||||
|
||||
This extension enables multi-tenant memory isolation for applications using
|
||||
Supabase Auth - each authenticated user's memories are stored in a separate
|
||||
schema, ensuring complete data isolation.
|
||||
|
||||
Features:
|
||||
- Local JWT Verification: Validates tokens locally using JWKS public keys
|
||||
(no network call per request)
|
||||
- Automatic Schema Isolation: Each user gets {prefix}_{user_id} schema
|
||||
- Zero User Management: Leverages your existing Supabase Auth setup
|
||||
- Production Ready: Includes health checks, timeouts, key rotation handling,
|
||||
and error handling
|
||||
- Built-in: Ships with Hindsight, no extra installation needed
|
||||
- Legacy Support: Falls back to /auth/v1/user endpoint for HS256 projects
|
||||
|
||||
JWT Verification Strategy:
|
||||
By default, JWTs are verified locally using public keys from the Supabase
|
||||
JWKS endpoint (/auth/v1/.well-known/jwks.json). This is the Supabase-recommended
|
||||
approach: no network call per request, fast, and secure.
|
||||
|
||||
If JWKS keys are unavailable (e.g., legacy HS256 projects), the extension
|
||||
falls back to calling /auth/v1/user per request for validation. This requires
|
||||
the service_role key to be configured.
|
||||
|
||||
Configuration via environment variables:
|
||||
HINDSIGHT_API_TENANT_EXTENSION=hindsight_api.extensions.builtin.supabase_tenant:SupabaseTenantExtension
|
||||
HINDSIGHT_API_TENANT_SUPABASE_URL=https://your-project.supabase.co
|
||||
|
||||
# Optional - only required for legacy HS256 projects or health checks
|
||||
HINDSIGHT_API_TENANT_SUPABASE_SERVICE_KEY=your-service-role-key
|
||||
|
||||
# Optional
|
||||
HINDSIGHT_API_TENANT_SCHEMA_PREFIX=user # Default: "user" (creates user_<uuid> schemas)
|
||||
|
||||
Usage:
|
||||
Clients pass their Supabase JWT in the Authorization header:
|
||||
|
||||
curl -H "Authorization: Bearer <supabase_jwt>" \\
|
||||
https://your-hindsight-server/v1/default/banks/my-bank/memories/recall
|
||||
|
||||
Author: BrighterBalance (https://brighterbalance.app)
|
||||
License: MIT
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import re
|
||||
import time
|
||||
|
||||
import httpx
|
||||
import jwt as pyjwt
|
||||
from jwt import PyJWK
|
||||
|
||||
from hindsight_api.extensions.tenant import AuthenticationError, Tenant, TenantContext, TenantExtension
|
||||
from hindsight_api.models import RequestContext
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
__all__ = ["SupabaseTenantExtension"]
|
||||
|
||||
# Minimum expected JWT length (JWTs are typically 100+ characters)
|
||||
MIN_TOKEN_LENGTH = 20
|
||||
|
||||
# Timeout for Supabase API calls
|
||||
REQUEST_TIMEOUT_SECONDS = 10.0
|
||||
|
||||
# JWKS cache TTL — Supabase Edge caches JWKS for 10 minutes, so we match that
|
||||
JWKS_CACHE_TTL_SECONDS = 600
|
||||
|
||||
# Minimum interval between JWKS refreshes to avoid hammering the endpoint
|
||||
JWKS_MIN_REFRESH_INTERVAL_SECONDS = 30
|
||||
|
||||
# Algorithms supported by Supabase Auth for asymmetric JWT signing
|
||||
SUPPORTED_ALGORITHMS = ["RS256", "ES256"]
|
||||
|
||||
# Supabase user IDs are UUIDs — validate before using in schema names
|
||||
_UUID_RE = re.compile(r"^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$", re.IGNORECASE)
|
||||
|
||||
# Schema prefix must be a valid Postgres identifier component (letters, digits, underscores)
|
||||
_SCHEMA_PREFIX_RE = re.compile(r"^[a-zA-Z_][a-zA-Z0-9_]*$")
|
||||
|
||||
|
||||
class SupabaseTenantExtension(TenantExtension):
|
||||
"""
|
||||
TenantExtension that validates Supabase JWTs for multi-tenant isolation.
|
||||
|
||||
Each authenticated user gets their own PostgreSQL schema, ensuring complete
|
||||
memory isolation between users. The schema name is derived from the user's
|
||||
Supabase user ID (the ``sub`` claim in the JWT).
|
||||
|
||||
JWT verification uses JWKS (local, no network call per request) when
|
||||
asymmetric keys are configured in Supabase, and falls back to the
|
||||
``/auth/v1/user`` endpoint for legacy HS256 projects.
|
||||
|
||||
Example:
|
||||
User with ID "a1b2c3d4-e5f6-7890-abcd-ef1234567890"
|
||||
gets schema "user_a1b2c3d4_e5f6_7890_abcd_ef1234567890"
|
||||
"""
|
||||
|
||||
def __init__(self, config: dict[str, str]) -> None:
|
||||
"""
|
||||
Initialize with configuration from environment variables.
|
||||
|
||||
Config keys are derived from HINDSIGHT_API_TENANT_* env vars:
|
||||
- HINDSIGHT_API_TENANT_SUPABASE_URL -> config["supabase_url"] (required)
|
||||
- HINDSIGHT_API_TENANT_SUPABASE_SERVICE_KEY -> config["supabase_service_key"] (optional)
|
||||
- HINDSIGHT_API_TENANT_SCHEMA_PREFIX -> config["schema_prefix"] (optional)
|
||||
|
||||
Args:
|
||||
config: Dictionary of configuration values from environment
|
||||
|
||||
Raises:
|
||||
ValueError: If required configuration is missing
|
||||
"""
|
||||
super().__init__(config)
|
||||
|
||||
self.supabase_url = (config.get("supabase_url") or "").rstrip("/")
|
||||
self.supabase_service_key = config.get("supabase_service_key")
|
||||
self.schema_prefix = config.get("schema_prefix", "user")
|
||||
|
||||
# Track initialized schemas to avoid redundant migrations
|
||||
self._initialized_schemas: set[str] = set()
|
||||
|
||||
# Reusable HTTP client (created on startup)
|
||||
self._http_client: httpx.AsyncClient | None = None
|
||||
|
||||
# JWKS state
|
||||
self._jwks_keys: dict[str, PyJWK] = {}
|
||||
self._jwks_last_fetched: float = 0
|
||||
self._use_jwks: bool = False
|
||||
|
||||
if not self.supabase_url:
|
||||
raise ValueError(
|
||||
"HINDSIGHT_API_TENANT_SUPABASE_URL is required. "
|
||||
"Set it to your Supabase project URL (e.g., https://xxx.supabase.co)"
|
||||
)
|
||||
|
||||
if not _SCHEMA_PREFIX_RE.match(self.schema_prefix):
|
||||
raise ValueError(
|
||||
f"Invalid schema_prefix '{self.schema_prefix}'. "
|
||||
"Must be a valid Postgres identifier (letters, digits, underscores, starting with a letter or underscore)."
|
||||
)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Lifecycle
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def on_startup(self) -> None:
|
||||
"""
|
||||
Called when Hindsight starts.
|
||||
|
||||
Creates a reusable HTTP client, fetches JWKS for local JWT verification,
|
||||
and optionally verifies connectivity to Supabase.
|
||||
"""
|
||||
logger.info("Initializing Supabase tenant extension")
|
||||
logger.info("Supabase URL: %s", self.supabase_url)
|
||||
logger.info("Schema prefix: %s_", self.schema_prefix)
|
||||
|
||||
self._http_client = httpx.AsyncClient(timeout=REQUEST_TIMEOUT_SECONDS)
|
||||
|
||||
# Attempt to fetch JWKS for fast local JWT verification
|
||||
await self._try_init_jwks()
|
||||
|
||||
# Optional health check using service key
|
||||
if self.supabase_service_key:
|
||||
await self._health_check()
|
||||
|
||||
async def on_shutdown(self) -> None:
|
||||
"""Called when Hindsight shuts down. Closes the HTTP client."""
|
||||
logger.info("Shutting down Supabase tenant extension")
|
||||
if self._http_client:
|
||||
await self._http_client.aclose()
|
||||
self._http_client = None
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# JWKS management
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def _try_init_jwks(self) -> None:
|
||||
"""Fetch JWKS and decide verification mode (local JWKS vs legacy endpoint)."""
|
||||
try:
|
||||
await self._fetch_jwks()
|
||||
if self._jwks_keys:
|
||||
self._use_jwks = True
|
||||
logger.info(
|
||||
"JWKS loaded — using local JWT verification with %d key(s)",
|
||||
len(self._jwks_keys),
|
||||
)
|
||||
return
|
||||
|
||||
# JWKS endpoint returned no keys — project likely uses legacy HS256
|
||||
logger.warning(
|
||||
"JWKS endpoint returned no signing keys. "
|
||||
"Falling back to /auth/v1/user endpoint for JWT verification. "
|
||||
"For better performance, enable asymmetric JWT signing in your "
|
||||
"Supabase dashboard (Project Settings → Auth → JWT Algorithm)."
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
"Could not fetch JWKS (%s). Falling back to /auth/v1/user endpoint for JWT verification.",
|
||||
e,
|
||||
)
|
||||
|
||||
# Legacy mode requires service key
|
||||
if not self.supabase_service_key:
|
||||
raise ValueError(
|
||||
"HINDSIGHT_API_TENANT_SUPABASE_SERVICE_KEY is required when JWKS "
|
||||
"is not available. Either enable asymmetric JWT signing in your "
|
||||
"Supabase project or provide the service_role key."
|
||||
)
|
||||
self._use_jwks = False
|
||||
|
||||
async def _fetch_jwks(self) -> None:
|
||||
"""Fetch public signing keys from the Supabase JWKS endpoint."""
|
||||
if self._http_client is None:
|
||||
raise RuntimeError("HTTP client not initialized")
|
||||
|
||||
url = f"{self.supabase_url}/auth/v1/.well-known/jwks.json"
|
||||
response = await self._http_client.get(url)
|
||||
response.raise_for_status()
|
||||
|
||||
jwks_data = response.json()
|
||||
keys: dict[str, PyJWK] = {}
|
||||
for key_data in jwks_data.get("keys", []):
|
||||
kid = key_data.get("kid")
|
||||
if kid:
|
||||
keys[kid] = PyJWK(key_data)
|
||||
|
||||
self._jwks_keys = keys
|
||||
self._jwks_last_fetched = time.monotonic()
|
||||
|
||||
async def _get_signing_key(self, token: str) -> PyJWK:
|
||||
"""
|
||||
Resolve the signing key for a token from the JWKS cache.
|
||||
|
||||
If the key ID (``kid``) is not in the cache, triggers one JWKS refresh
|
||||
to handle key rotation before raising an error.
|
||||
"""
|
||||
header = pyjwt.get_unverified_header(token)
|
||||
kid = header.get("kid")
|
||||
if not kid:
|
||||
raise AuthenticationError("Token missing key ID (kid) header")
|
||||
|
||||
# Refresh cache if stale
|
||||
now = time.monotonic()
|
||||
if now - self._jwks_last_fetched > JWKS_CACHE_TTL_SECONDS:
|
||||
logger.debug("JWKS cache expired, refreshing")
|
||||
await self._fetch_jwks()
|
||||
|
||||
if kid in self._jwks_keys:
|
||||
return self._jwks_keys[kid]
|
||||
|
||||
# Key not found — try one forced refresh to handle key rotation,
|
||||
# but only if we haven't just refreshed
|
||||
if now - self._jwks_last_fetched > JWKS_MIN_REFRESH_INTERVAL_SECONDS:
|
||||
logger.info("Signing key %s not in cache, refreshing JWKS for possible key rotation", kid)
|
||||
await self._fetch_jwks()
|
||||
if kid in self._jwks_keys:
|
||||
return self._jwks_keys[kid]
|
||||
|
||||
raise AuthenticationError("Unable to find signing key for token")
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Authentication
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def authenticate(self, context: RequestContext) -> TenantContext:
|
||||
"""
|
||||
Validate a Supabase JWT and return tenant context.
|
||||
|
||||
Uses local JWKS verification when available (no network call per
|
||||
request), falling back to the ``/auth/v1/user`` endpoint for legacy
|
||||
HS256 projects.
|
||||
|
||||
Args:
|
||||
context: Request context containing the API key (JWT)
|
||||
|
||||
Returns:
|
||||
TenantContext with schema_name set to ``{prefix}_{user_uuid}``
|
||||
|
||||
Raises:
|
||||
AuthenticationError: If token is missing, invalid, or expired
|
||||
"""
|
||||
token = context.api_key
|
||||
|
||||
if not token:
|
||||
raise AuthenticationError("Missing Authorization header. Expected: Bearer <supabase_jwt>")
|
||||
|
||||
if len(token) < MIN_TOKEN_LENGTH:
|
||||
raise AuthenticationError("Invalid token format")
|
||||
|
||||
if self._http_client is None:
|
||||
raise AuthenticationError("Extension not initialized")
|
||||
|
||||
# Verify the JWT and extract user ID
|
||||
if self._use_jwks:
|
||||
user_id = await self._verify_token_jwks(token)
|
||||
else:
|
||||
user_id = await self._verify_token_legacy(token)
|
||||
|
||||
# Validate user ID format before using in schema name
|
||||
if not _UUID_RE.match(user_id):
|
||||
raise AuthenticationError("Invalid user ID format in token")
|
||||
|
||||
# Build isolated schema name — hyphens to underscores for Postgres compatibility
|
||||
safe_user_id = user_id.replace("-", "_")
|
||||
schema_name = f"{self.schema_prefix}_{safe_user_id}"
|
||||
|
||||
# Initialize schema on first access
|
||||
if schema_name not in self._initialized_schemas:
|
||||
await self._initialize_schema(schema_name)
|
||||
|
||||
return TenantContext(schema_name=schema_name)
|
||||
|
||||
async def _verify_token_jwks(self, token: str) -> str:
|
||||
"""
|
||||
Verify a JWT locally using cached JWKS public keys.
|
||||
|
||||
Validates signature, expiration, issuer, and audience. Returns the
|
||||
user ID from the ``sub`` claim.
|
||||
|
||||
Raises:
|
||||
AuthenticationError: If the token is invalid or expired.
|
||||
"""
|
||||
try:
|
||||
signing_key = await self._get_signing_key(token)
|
||||
payload = pyjwt.decode(
|
||||
token,
|
||||
signing_key.key,
|
||||
algorithms=SUPPORTED_ALGORITHMS,
|
||||
audience="authenticated",
|
||||
issuer=f"{self.supabase_url}/auth/v1",
|
||||
)
|
||||
except pyjwt.ExpiredSignatureError:
|
||||
raise AuthenticationError("Token has expired")
|
||||
except pyjwt.InvalidAudienceError:
|
||||
raise AuthenticationError("Invalid token audience")
|
||||
except pyjwt.InvalidIssuerError:
|
||||
raise AuthenticationError("Invalid token issuer")
|
||||
except pyjwt.DecodeError:
|
||||
raise AuthenticationError("Invalid token")
|
||||
except AuthenticationError:
|
||||
raise
|
||||
except Exception as e:
|
||||
raise AuthenticationError(f"Token verification failed: {e!s}")
|
||||
|
||||
user_id = payload.get("sub")
|
||||
if not user_id:
|
||||
raise AuthenticationError("Token valid but missing subject (sub) claim")
|
||||
return user_id
|
||||
|
||||
async def _verify_token_legacy(self, token: str) -> str:
|
||||
"""
|
||||
Verify a JWT by calling the Supabase ``/auth/v1/user`` endpoint.
|
||||
|
||||
This is the fallback for projects using legacy HS256 JWT signing.
|
||||
Adds a network round-trip per request.
|
||||
|
||||
Raises:
|
||||
AuthenticationError: If the token is invalid or the request fails.
|
||||
"""
|
||||
try:
|
||||
response = await self._http_client.get(
|
||||
f"{self.supabase_url}/auth/v1/user",
|
||||
headers={
|
||||
"Authorization": f"Bearer {token}",
|
||||
"apikey": self.supabase_service_key,
|
||||
},
|
||||
)
|
||||
|
||||
if response.status_code == 401:
|
||||
raise AuthenticationError("Invalid or expired token")
|
||||
|
||||
if response.status_code != 200:
|
||||
raise AuthenticationError(f"Authentication failed: {response.status_code}")
|
||||
|
||||
user_data = response.json()
|
||||
user_id = user_data.get("id")
|
||||
|
||||
if not user_id:
|
||||
raise AuthenticationError("Token valid but no user ID found")
|
||||
|
||||
return user_id
|
||||
|
||||
except AuthenticationError:
|
||||
raise
|
||||
except httpx.TimeoutException:
|
||||
raise AuthenticationError("Authentication timeout - please retry")
|
||||
except httpx.RequestError as e:
|
||||
raise AuthenticationError(f"Connection error: {e!s}")
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Schema management
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def _initialize_schema(self, schema_name: str) -> None:
|
||||
"""Run migrations for a new tenant schema and cache the result."""
|
||||
logger.info("Initializing schema: %s", schema_name)
|
||||
try:
|
||||
await self.context.run_migration(schema_name)
|
||||
self._initialized_schemas.add(schema_name)
|
||||
logger.info("Schema ready: %s", schema_name)
|
||||
except Exception as e:
|
||||
logger.error("Schema initialization failed for %s: %s", schema_name, e)
|
||||
raise AuthenticationError(f"Failed to initialize tenant: {e!s}")
|
||||
|
||||
async def list_tenants(self) -> list[Tenant]:
|
||||
"""Return all tenant schemas that have been initialized."""
|
||||
return [Tenant(schema=schema) for schema in self._initialized_schemas]
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Health check
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def _health_check(self) -> None:
|
||||
"""Verify connectivity to Supabase using the auth health endpoint."""
|
||||
try:
|
||||
response = await self._http_client.get(
|
||||
f"{self.supabase_url}/auth/v1/health",
|
||||
headers={"apikey": self.supabase_service_key},
|
||||
)
|
||||
if response.status_code == 200:
|
||||
logger.info("Supabase connection verified")
|
||||
else:
|
||||
logger.warning("Supabase health check returned %d", response.status_code)
|
||||
except Exception as e:
|
||||
logger.warning("Could not verify Supabase connection: %s", e)
|
||||
@@ -1,90 +0,0 @@
|
||||
"""Built-in tenant extension implementations."""
|
||||
|
||||
from hindsight_api.config import get_config
|
||||
from hindsight_api.extensions.tenant import AuthenticationError, Tenant, TenantContext, TenantExtension
|
||||
from hindsight_api.models import RequestContext
|
||||
|
||||
|
||||
class DefaultTenantExtension(TenantExtension):
|
||||
"""
|
||||
Default single-tenant extension with no authentication.
|
||||
|
||||
This is the default extension used when no tenant extension is configured.
|
||||
It provides single-tenant behavior using the configured schema from
|
||||
HINDSIGHT_API_DATABASE_SCHEMA (defaults to 'public').
|
||||
|
||||
Features:
|
||||
- No authentication required (passes all requests)
|
||||
- Uses configured schema from environment
|
||||
- Perfect for single-tenant deployments without auth
|
||||
|
||||
Configuration:
|
||||
HINDSIGHT_API_DATABASE_SCHEMA=your-schema (optional, defaults to 'public')
|
||||
|
||||
This is automatically enabled by default. To use custom authentication,
|
||||
configure a different tenant extension:
|
||||
HINDSIGHT_API_TENANT_EXTENSION=hindsight_api.extensions.builtin.tenant:ApiKeyTenantExtension
|
||||
"""
|
||||
|
||||
def __init__(self, config: dict[str, str]):
|
||||
super().__init__(config)
|
||||
# Cache the schema at initialization for consistency
|
||||
# Support explicit schema override via config, otherwise use environment
|
||||
self._schema = config.get("schema", get_config().database_schema)
|
||||
|
||||
async def authenticate(self, context: RequestContext) -> TenantContext:
|
||||
"""Return configured schema without any authentication."""
|
||||
return TenantContext(schema_name=self._schema)
|
||||
|
||||
async def list_tenants(self) -> list[Tenant]:
|
||||
"""Return configured schema for single-tenant setup."""
|
||||
return [Tenant(schema=self._schema)]
|
||||
|
||||
|
||||
class ApiKeyTenantExtension(TenantExtension):
|
||||
"""
|
||||
Built-in tenant extension that validates API key against an environment variable.
|
||||
|
||||
This is a simple implementation that:
|
||||
1. Validates the API key matches HINDSIGHT_API_TENANT_API_KEY
|
||||
2. Returns the configured schema (HINDSIGHT_API_DATABASE_SCHEMA, default 'public')
|
||||
for all authenticated requests
|
||||
|
||||
Configuration:
|
||||
HINDSIGHT_API_TENANT_EXTENSION=hindsight_api.extensions.builtin.tenant:ApiKeyTenantExtension
|
||||
HINDSIGHT_API_TENANT_API_KEY=your-secret-key
|
||||
HINDSIGHT_API_DATABASE_SCHEMA=your-schema (optional, defaults to 'public')
|
||||
HINDSIGHT_API_TENANT_MCP_AUTH_DISABLED=true (optional, disable auth for MCP endpoints)
|
||||
|
||||
For multi-tenant setups with separate schemas per tenant, implement a custom
|
||||
TenantExtension that looks up the schema based on the API key or token claims.
|
||||
"""
|
||||
|
||||
def __init__(self, config: dict[str, str]):
|
||||
super().__init__(config)
|
||||
self.expected_api_key = config.get("api_key")
|
||||
if not self.expected_api_key:
|
||||
raise ValueError("HINDSIGHT_API_TENANT_API_KEY is required when using ApiKeyTenantExtension")
|
||||
# Allow disabling MCP auth for backwards compatibility
|
||||
self.mcp_auth_disabled = config.get("mcp_auth_disabled", "").lower() in ("true", "1", "yes")
|
||||
|
||||
async def authenticate(self, context: RequestContext) -> TenantContext:
|
||||
"""Validate API key and return configured schema context."""
|
||||
if context.api_key != self.expected_api_key:
|
||||
raise AuthenticationError("Invalid API key")
|
||||
return TenantContext(schema_name=get_config().database_schema)
|
||||
|
||||
async def list_tenants(self) -> list[Tenant]:
|
||||
"""Return configured schema for single-tenant setup."""
|
||||
return [Tenant(schema=get_config().database_schema)]
|
||||
|
||||
async def authenticate_mcp(self, context: RequestContext) -> TenantContext:
|
||||
"""
|
||||
Authenticate MCP requests.
|
||||
|
||||
If mcp_auth_disabled is set, skip authentication for backwards compatibility.
|
||||
Otherwise, delegate to authenticate().
|
||||
"""
|
||||
if self.mcp_auth_disabled:
|
||||
return TenantContext(schema_name=get_config().database_schema)
|
||||
return await self.authenticate(context)
|
||||
@@ -1,42 +0,0 @@
|
||||
"""MCP Extension for registering additional MCP tools.
|
||||
|
||||
This extension allows external packages (like hindsight-cloud) to register
|
||||
additional MCP tools on the Hindsight MCP server.
|
||||
|
||||
Example:
|
||||
HINDSIGHT_API_MCP_EXTENSION=hindsight_cloud.extensions:CloudMCPExtension
|
||||
"""
|
||||
|
||||
import logging
|
||||
from abc import abstractmethod
|
||||
|
||||
from fastmcp import FastMCP
|
||||
|
||||
from hindsight_api import MemoryEngine
|
||||
from hindsight_api.extensions.base import Extension
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class MCPExtension(Extension):
|
||||
"""Base class for MCP extensions that register additional tools.
|
||||
|
||||
Subclass this to add MCP tools in extension packages.
|
||||
|
||||
Example:
|
||||
class CloudMCPExtension(MCPExtension):
|
||||
def register_tools(self, mcp: FastMCP, memory: MemoryEngine) -> None:
|
||||
@mcp.tool()
|
||||
async def my_custom_tool(query: str) -> str:
|
||||
return "result"
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def register_tools(self, mcp: FastMCP, memory: MemoryEngine) -> None:
|
||||
"""Register additional MCP tools.
|
||||
|
||||
Args:
|
||||
mcp: FastMCP server instance to register tools on
|
||||
memory: MemoryEngine instance for accessing memory operations
|
||||
"""
|
||||
pass
|
||||
@@ -1,158 +0,0 @@
|
||||
"""Tenant Extension for multi-tenancy and API key authentication."""
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
from hindsight_api.extensions.base import Extension
|
||||
from hindsight_api.models import RequestContext
|
||||
|
||||
|
||||
class AuthenticationError(Exception):
|
||||
"""Raised when authentication fails."""
|
||||
|
||||
def __init__(self, reason: str, headers: dict[str, str] | None = None):
|
||||
self.reason = reason
|
||||
self.headers = headers or {}
|
||||
super().__init__(f"Authentication failed: {reason}")
|
||||
|
||||
|
||||
@dataclass
|
||||
class TenantContext:
|
||||
"""
|
||||
Tenant context returned by authentication.
|
||||
|
||||
Contains the PostgreSQL schema name for tenant isolation.
|
||||
All database queries will use fully-qualified table names
|
||||
with this schema (e.g., schema_name.memory_units).
|
||||
"""
|
||||
|
||||
schema_name: str
|
||||
|
||||
|
||||
@dataclass
|
||||
class Tenant:
|
||||
"""
|
||||
Represents a tenant for worker discovery.
|
||||
|
||||
Used by list_tenants() to return tenant information including
|
||||
the PostgreSQL schema name for database operations.
|
||||
"""
|
||||
|
||||
schema: str
|
||||
|
||||
|
||||
class TenantExtension(Extension, ABC):
|
||||
"""
|
||||
Extension for multi-tenancy and API key authentication.
|
||||
|
||||
This extension validates incoming requests and returns the tenant context
|
||||
including the PostgreSQL schema to use for database operations.
|
||||
|
||||
Built-in implementation:
|
||||
hindsight_api.extensions.builtin.tenant.ApiKeyTenantExtension
|
||||
|
||||
Enable via environment variable:
|
||||
HINDSIGHT_API_TENANT_EXTENSION=hindsight_api.extensions.builtin.tenant:ApiKeyTenantExtension
|
||||
HINDSIGHT_API_TENANT_API_KEY=your-secret-key
|
||||
|
||||
The returned schema_name is used for fully-qualified table names in queries,
|
||||
enabling tenant isolation at the database level.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
async def authenticate(self, context: RequestContext) -> TenantContext:
|
||||
"""
|
||||
Authenticate the action context and return tenant context.
|
||||
|
||||
Args:
|
||||
context: The action context containing API key and other auth data.
|
||||
|
||||
Returns:
|
||||
TenantContext with the schema_name for database operations.
|
||||
|
||||
Raises:
|
||||
AuthenticationError: If authentication fails.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def list_tenants(self) -> list[Tenant]:
|
||||
"""
|
||||
List all tenants that should be processed by workers.
|
||||
|
||||
This method is used by the worker to discover all tenants that need
|
||||
task polling. Workers will poll for pending tasks in each tenant's schema.
|
||||
|
||||
Returns:
|
||||
List of Tenant objects containing schema information.
|
||||
For single-tenant setups, return [Tenant(schema="public")].
|
||||
"""
|
||||
...
|
||||
|
||||
async def get_tenant_config(self, context: RequestContext) -> dict[str, Any]:
|
||||
"""
|
||||
Get tenant-specific configuration overrides.
|
||||
|
||||
This method is called during hierarchical configuration resolution to get
|
||||
tenant-level config overrides. The returned dict should contain Python field
|
||||
names (lowercase snake_case) as keys, not environment variable names.
|
||||
|
||||
Example:
|
||||
{"llm_model": "gpt-4", "retain_extraction_mode": "verbose"}
|
||||
|
||||
The default implementation returns an empty dict (no tenant-specific config).
|
||||
Override this method in custom extensions to provide tenant-specific configuration.
|
||||
|
||||
Args:
|
||||
context: The request context containing tenant information.
|
||||
|
||||
Returns:
|
||||
Dict of config field names to values (only configurable fields).
|
||||
Empty dict if no tenant-specific config.
|
||||
"""
|
||||
return {}
|
||||
|
||||
async def get_allowed_config_fields(self, context: RequestContext, bank_id: str) -> set[str] | None:
|
||||
"""
|
||||
Get set of config fields that this tenant/bank is allowed to modify.
|
||||
|
||||
This method controls which configurable fields can be modified via the bank config API.
|
||||
It enables fine-grained permission control per tenant or per bank.
|
||||
|
||||
Examples:
|
||||
- Return None: Allow all configurable fields (default)
|
||||
- Return {"retain_chunk_size", "retain_custom_instructions"}: Allow only these fields
|
||||
- Return set(): Allow no modifications (read-only)
|
||||
|
||||
The default implementation returns None (all configurable fields allowed).
|
||||
Override this method in custom extensions to implement custom permission logic.
|
||||
|
||||
Args:
|
||||
context: The request context containing tenant information.
|
||||
bank_id: The bank identifier for per-bank permissions.
|
||||
|
||||
Returns:
|
||||
Set of allowed field names, or None to allow all configurable fields.
|
||||
Returned fields must be a subset of HindsightConfig.get_configurable_fields().
|
||||
"""
|
||||
return None
|
||||
|
||||
async def authenticate_mcp(self, context: RequestContext) -> TenantContext:
|
||||
"""
|
||||
Authenticate MCP requests.
|
||||
|
||||
By default, this calls authenticate(). Override this method to provide
|
||||
different authentication behavior for MCP endpoints (e.g., to disable
|
||||
auth for backwards compatibility with existing MCP servers).
|
||||
|
||||
Args:
|
||||
context: The action context containing API key and other auth data.
|
||||
|
||||
Returns:
|
||||
TenantContext with the schema_name for database operations.
|
||||
|
||||
Raises:
|
||||
AuthenticationError: If authentication fails.
|
||||
"""
|
||||
return await self.authenticate(context)
|
||||
@@ -1,40 +0,0 @@
|
||||
"""
|
||||
Local MCP server entry point for use with Claude Code (HTTP transport).
|
||||
|
||||
This is a thin wrapper around the main hindsight-api server that pre-configures
|
||||
sensible defaults for local use (embedded PostgreSQL via pg0, warning log level).
|
||||
|
||||
The full API runs on localhost:8888. Configure Claude Code's MCP settings:
|
||||
claude mcp add --transport http hindsight http://localhost:8888/mcp/
|
||||
|
||||
Or pinned to a specific bank (single-bank mode):
|
||||
claude mcp add --transport http hindsight http://localhost:8888/mcp/default/
|
||||
|
||||
Run with:
|
||||
hindsight-local-mcp
|
||||
|
||||
Or with uvx:
|
||||
uvx hindsight-api@latest hindsight-local-mcp
|
||||
|
||||
Environment variables:
|
||||
HINDSIGHT_API_LLM_API_KEY: Required. API key for LLM provider.
|
||||
HINDSIGHT_API_LLM_PROVIDER: Optional. LLM provider (default: "openai").
|
||||
HINDSIGHT_API_LLM_MODEL: Optional. LLM model (default: "gpt-4o-mini").
|
||||
HINDSIGHT_API_DATABASE_URL: Optional. Override database URL (default: pg0://hindsight-mcp).
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
|
||||
def main() -> None:
|
||||
"""Start the Hindsight API server with local defaults."""
|
||||
# Set local defaults (only if not already configured by the user)
|
||||
os.environ.setdefault("HINDSIGHT_API_DATABASE_URL", "pg0://hindsight-mcp")
|
||||
|
||||
from hindsight_api.main import main as api_main
|
||||
|
||||
api_main()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -1,480 +0,0 @@
|
||||
"""
|
||||
OpenTelemetry distributed tracing instrumentation for Hindsight API.
|
||||
|
||||
This module provides tracing for:
|
||||
- LLM API calls with full prompts/completions following GenAI semantic conventions
|
||||
- Token usage and model information
|
||||
- Error tracking and finish reasons
|
||||
|
||||
Tracing is conditional and disabled by default. When enabled, traces are exported
|
||||
to Langfuse (or any OTLP-compatible backend) via OTLP HTTP protocol.
|
||||
"""
|
||||
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
from typing import Any, Optional
|
||||
|
||||
from opentelemetry import trace
|
||||
from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
|
||||
from opentelemetry.sdk.resources import Resource
|
||||
from opentelemetry.sdk.trace import TracerProvider
|
||||
from opentelemetry.sdk.trace.export import BatchSpanProcessor
|
||||
from opentelemetry.trace import Status, StatusCode
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _serialize_for_span(obj: Any) -> str:
|
||||
"""Serialize an object for span recording, handling Pydantic models."""
|
||||
if isinstance(obj, str):
|
||||
return obj
|
||||
if hasattr(obj, "model_dump_json"):
|
||||
# Pydantic v2 model
|
||||
return obj.model_dump_json()
|
||||
if hasattr(obj, "json"):
|
||||
# Pydantic v1 model
|
||||
return obj.json()
|
||||
if hasattr(obj, "model_dump"):
|
||||
# Pydantic v2 model - convert to dict then json
|
||||
return json.dumps(obj.model_dump())
|
||||
if hasattr(obj, "dict"):
|
||||
# Pydantic v1 model - convert to dict then json
|
||||
return json.dumps(obj.dict())
|
||||
# Fallback to json.dumps for dicts and other types
|
||||
return json.dumps(obj)
|
||||
|
||||
|
||||
# No-op tracer for when tracing is disabled
|
||||
class NoOpTracer:
|
||||
"""No-op tracer that provides the same interface as OpenTelemetry Tracer but does nothing."""
|
||||
|
||||
def start_as_current_span(self, name: str, **kwargs):
|
||||
"""Return a no-op context manager that yields a NoOpSpan."""
|
||||
from contextlib import contextmanager
|
||||
|
||||
@contextmanager
|
||||
def noop_span_context():
|
||||
yield NoOpSpan()
|
||||
|
||||
return noop_span_context()
|
||||
|
||||
def start_span(self, name: str, **kwargs):
|
||||
"""Return a no-op span."""
|
||||
return NoOpSpan()
|
||||
|
||||
|
||||
class NoOpSpan:
|
||||
"""No-op span that provides the same interface as OpenTelemetry Span but does nothing."""
|
||||
|
||||
def set_attribute(self, key: str, value: Any) -> None:
|
||||
"""No-op."""
|
||||
pass
|
||||
|
||||
def set_status(self, status: Any) -> None:
|
||||
"""No-op."""
|
||||
pass
|
||||
|
||||
def record_exception(self, exception: Exception) -> None:
|
||||
"""No-op."""
|
||||
pass
|
||||
|
||||
def add_event(self, name: str, attributes: dict | None = None) -> None:
|
||||
"""No-op."""
|
||||
pass
|
||||
|
||||
def end(self, end_time: int | None = None) -> None:
|
||||
"""No-op."""
|
||||
pass
|
||||
|
||||
|
||||
# Global tracer instance
|
||||
_tracer: trace.Tracer | NoOpTracer = NoOpTracer()
|
||||
_tracing_enabled: bool = False
|
||||
|
||||
|
||||
# GenAI semantic convention attribute names (based on v1.37 spec)
|
||||
class GenAIAttributes:
|
||||
"""GenAI semantic convention attribute names."""
|
||||
|
||||
# Operation and provider
|
||||
OPERATION_NAME = "gen_ai.operation.name"
|
||||
PROVIDER_NAME = "gen_ai.provider.name"
|
||||
|
||||
# Model information
|
||||
REQUEST_MODEL = "gen_ai.request.model"
|
||||
RESPONSE_MODEL = "gen_ai.response.model"
|
||||
|
||||
# Token usage
|
||||
USAGE_INPUT_TOKENS = "gen_ai.usage.input_tokens"
|
||||
USAGE_OUTPUT_TOKENS = "gen_ai.usage.output_tokens"
|
||||
|
||||
# Messages and prompts
|
||||
SYSTEM_INSTRUCTIONS = "gen_ai.system_instructions"
|
||||
INPUT_MESSAGES = "gen_ai.input.messages"
|
||||
OUTPUT_MESSAGES = "gen_ai.output.messages"
|
||||
|
||||
# Response metadata
|
||||
FINISH_REASONS = "gen_ai.response.finish_reasons"
|
||||
|
||||
# Error tracking
|
||||
ERROR_TYPE = "error.type"
|
||||
|
||||
|
||||
# Provider name mapping (Hindsight internal -> GenAI semantic convention)
|
||||
PROVIDER_NAME_MAPPING = {
|
||||
"openai": "openai",
|
||||
"anthropic": "anthropic",
|
||||
"gemini": "google",
|
||||
"vertexai": "google",
|
||||
"groq": "groq",
|
||||
"ollama": "ollama",
|
||||
"lmstudio": "lmstudio",
|
||||
"openai-codex": "openai",
|
||||
"claude-code": "anthropic",
|
||||
"mock": "mock",
|
||||
}
|
||||
|
||||
|
||||
def initialize_tracing(
|
||||
service_name: str,
|
||||
endpoint: str,
|
||||
headers: Optional[str] = None,
|
||||
deployment_environment: str = "development",
|
||||
) -> None:
|
||||
"""
|
||||
Initialize OpenTelemetry tracing with OTLP exporter.
|
||||
|
||||
Args:
|
||||
service_name: Name of the service for resource attributes
|
||||
endpoint: OTLP endpoint URL (e.g., https://cloud.langfuse.com/api/public/otel)
|
||||
headers: Optional headers in format "key1=value1,key2=value2"
|
||||
deployment_environment: Deployment environment (e.g., development, staging, production)
|
||||
"""
|
||||
global _tracer, _tracing_enabled
|
||||
|
||||
# Create resource with service information
|
||||
resource = Resource.create(
|
||||
{
|
||||
"service.name": service_name,
|
||||
"service.version": "0.4.8", # Could import from __version__
|
||||
"deployment.environment.name": deployment_environment,
|
||||
}
|
||||
)
|
||||
|
||||
# Parse headers
|
||||
headers_dict = {}
|
||||
if headers:
|
||||
for pair in headers.split(","):
|
||||
if "=" in pair:
|
||||
key, value = pair.split("=", 1)
|
||||
headers_dict[key.strip()] = value.strip()
|
||||
|
||||
# Create OTLP HTTP exporter
|
||||
# Note: Langfuse expects /v1/traces path appended to base endpoint
|
||||
otlp_endpoint = endpoint if endpoint.endswith("/v1/traces") else f"{endpoint}/v1/traces"
|
||||
otlp_exporter = OTLPSpanExporter(
|
||||
endpoint=otlp_endpoint,
|
||||
headers=headers_dict,
|
||||
)
|
||||
|
||||
# Create tracer provider with batch processor
|
||||
provider = TracerProvider(resource=resource)
|
||||
provider.add_span_processor(BatchSpanProcessor(otlp_exporter))
|
||||
|
||||
# Set global tracer provider
|
||||
trace.set_tracer_provider(provider)
|
||||
|
||||
# Get tracer for this application
|
||||
_tracer = trace.get_tracer(__name__)
|
||||
_tracing_enabled = True
|
||||
|
||||
logger.info(f"Tracing initialized: endpoint={otlp_endpoint}, service={service_name}")
|
||||
|
||||
|
||||
def get_tracer() -> trace.Tracer | NoOpTracer:
|
||||
"""
|
||||
Get the global tracer instance.
|
||||
|
||||
Returns a no-op tracer if tracing is disabled, so callers don't need to check for None.
|
||||
This improves code readability by allowing direct use without null checks.
|
||||
"""
|
||||
return _tracer
|
||||
|
||||
|
||||
def create_operation_span(operation: str, bank_id: str | None = None):
|
||||
"""
|
||||
Create a parent span for a Hindsight operation (retain, reflect, consolidation, etc.).
|
||||
|
||||
This creates the span hierarchy:
|
||||
- hindsight.{operation} (parent)
|
||||
- chat {model} (child LLM calls)
|
||||
|
||||
Args:
|
||||
operation: Operation name (retain, reflect, consolidation, mental_model_refresh)
|
||||
bank_id: Optional bank ID for context
|
||||
|
||||
Returns:
|
||||
Span context manager
|
||||
"""
|
||||
if not _tracing_enabled or _tracer is None:
|
||||
# Return a no-op context manager
|
||||
from contextlib import nullcontext
|
||||
|
||||
return nullcontext()
|
||||
|
||||
span_name = f"hindsight.{operation}"
|
||||
span = _tracer.start_as_current_span(span_name)
|
||||
|
||||
# Add operation-specific attributes
|
||||
if span and hasattr(span, "set_attribute"):
|
||||
span.set_attribute("hindsight.operation", operation)
|
||||
if bank_id:
|
||||
span.set_attribute("hindsight.bank_id", bank_id)
|
||||
|
||||
return span
|
||||
|
||||
|
||||
def is_tracing_enabled() -> bool:
|
||||
"""Check if tracing is enabled."""
|
||||
return _tracing_enabled
|
||||
|
||||
|
||||
# Maximum content length before truncation (to stay within span size limits)
|
||||
MAX_CONTENT_LENGTH = 100_000 # characters
|
||||
|
||||
|
||||
def _truncate_content(content: str) -> str:
|
||||
"""Truncate content if too large for span."""
|
||||
if len(content) > MAX_CONTENT_LENGTH:
|
||||
return content[:MAX_CONTENT_LENGTH] + f"\n\n[TRUNCATED: {len(content) - MAX_CONTENT_LENGTH} chars omitted]"
|
||||
return content
|
||||
|
||||
|
||||
class LLMSpanRecorder:
|
||||
"""
|
||||
Records OpenTelemetry spans for LLM calls following GenAI semantic conventions.
|
||||
"""
|
||||
|
||||
def __init__(self, tracer: trace.Tracer):
|
||||
self.tracer = tracer
|
||||
|
||||
def record_llm_call(
|
||||
self,
|
||||
provider: str,
|
||||
model: str,
|
||||
scope: str,
|
||||
messages: list[dict[str, str]],
|
||||
response_content: Optional[str],
|
||||
input_tokens: int,
|
||||
output_tokens: int,
|
||||
duration: float,
|
||||
finish_reason: Optional[str] = None,
|
||||
error: Optional[Exception] = None,
|
||||
tool_calls: Optional[list[dict[str, Any]]] = None,
|
||||
) -> None:
|
||||
"""
|
||||
Record a completed LLM call as a span with GenAI semantic conventions.
|
||||
|
||||
This creates a span AFTER the call completes, using timestamps to
|
||||
set the correct start/end times. This approach works better with
|
||||
the existing sync metrics recording pattern.
|
||||
|
||||
Args:
|
||||
provider: Hindsight provider name
|
||||
model: Model name
|
||||
scope: Scope identifier (memory, reflect, consolidation, etc.)
|
||||
messages: Input messages (chat history)
|
||||
response_content: Response text from LLM
|
||||
input_tokens: Input token count
|
||||
output_tokens: Output token count
|
||||
duration: Call duration in seconds
|
||||
finish_reason: Reason the model stopped (stop, length, tool_calls, etc.)
|
||||
error: Exception if call failed
|
||||
tool_calls: List of tool calls made (for function calling)
|
||||
"""
|
||||
try:
|
||||
# Map provider name to GenAI semantic convention
|
||||
genai_provider = PROVIDER_NAME_MAPPING.get(provider.lower(), provider.lower())
|
||||
|
||||
# Determine operation name based on scope/context
|
||||
operation_name = "chat" # Default for GenAI semantic conventions
|
||||
|
||||
# Create span name: "hindsight.{scope}" for consistency with parent spans
|
||||
# Model info is available in span attributes (gen_ai.request.model)
|
||||
if scope:
|
||||
span_name = f"hindsight.{scope}"
|
||||
else:
|
||||
# Fallback to chat {model} if no scope provided
|
||||
span_name = f"{operation_name} {model}"
|
||||
|
||||
# Calculate timestamps
|
||||
end_time_ns = time.time_ns()
|
||||
start_time_ns = end_time_ns - int(duration * 1_000_000_000)
|
||||
|
||||
# Create span with explicit timestamps
|
||||
with self.tracer.start_as_current_span(
|
||||
span_name,
|
||||
start_time=start_time_ns,
|
||||
end_on_exit=False, # We'll set end time manually
|
||||
) as span:
|
||||
# Set required attributes
|
||||
span.set_attribute(GenAIAttributes.OPERATION_NAME, operation_name)
|
||||
span.set_attribute(GenAIAttributes.PROVIDER_NAME, genai_provider)
|
||||
span.set_attribute(GenAIAttributes.REQUEST_MODEL, model)
|
||||
span.set_attribute(GenAIAttributes.RESPONSE_MODEL, model)
|
||||
span.set_attribute(GenAIAttributes.USAGE_INPUT_TOKENS, input_tokens)
|
||||
span.set_attribute(GenAIAttributes.USAGE_OUTPUT_TOKENS, output_tokens)
|
||||
|
||||
# Add custom attributes for Hindsight context
|
||||
span.set_attribute("hindsight.scope", scope)
|
||||
span.set_attribute("hindsight.provider.internal", provider)
|
||||
|
||||
# Add tool call information if present
|
||||
if tool_calls:
|
||||
span.set_attribute("gen_ai.tool_calls.count", len(tool_calls))
|
||||
# Add tool names as comma-separated list
|
||||
tool_names = [tc.get("name", "") for tc in tool_calls]
|
||||
span.set_attribute("gen_ai.tool_calls.names", ",".join(tool_names))
|
||||
|
||||
# Format messages for GenAI conventions (as JSON)
|
||||
input_messages_json = self._format_messages(messages)
|
||||
output_messages_json = self._format_output(response_content, finish_reason)
|
||||
|
||||
# Extract system instructions if present
|
||||
system_instructions = self._extract_system_instructions(messages)
|
||||
|
||||
# Add event with prompts/completions following v1.37 conventions
|
||||
event_attrs = {}
|
||||
if input_messages_json:
|
||||
event_attrs[GenAIAttributes.INPUT_MESSAGES] = input_messages_json
|
||||
if output_messages_json:
|
||||
event_attrs[GenAIAttributes.OUTPUT_MESSAGES] = output_messages_json
|
||||
if system_instructions:
|
||||
event_attrs[GenAIAttributes.SYSTEM_INSTRUCTIONS] = system_instructions
|
||||
if finish_reason:
|
||||
event_attrs[GenAIAttributes.FINISH_REASONS] = json.dumps([finish_reason])
|
||||
|
||||
span.add_event(
|
||||
"gen_ai.client.inference.operation.details",
|
||||
attributes=event_attrs,
|
||||
)
|
||||
|
||||
# Add individual tool call events with details
|
||||
if tool_calls:
|
||||
for i, tc in enumerate(tool_calls):
|
||||
tool_event_attrs = {
|
||||
"tool.name": tc.get("name", ""),
|
||||
"tool.id": tc.get("id", ""),
|
||||
"tool.arguments": json.dumps(tc.get("arguments", {})),
|
||||
}
|
||||
span.add_event(f"gen_ai.tool_call.{i}", attributes=tool_event_attrs)
|
||||
|
||||
# Handle errors
|
||||
if error:
|
||||
span.set_status(Status(StatusCode.ERROR, str(error)))
|
||||
span.set_attribute(GenAIAttributes.ERROR_TYPE, type(error).__name__)
|
||||
span.record_exception(error)
|
||||
else:
|
||||
span.set_status(Status(StatusCode.OK))
|
||||
|
||||
# Set end time
|
||||
span.end(end_time=end_time_ns)
|
||||
|
||||
except Exception as e:
|
||||
# Don't let tracing errors break LLM calls
|
||||
logger.error(f"Failed to record LLM span: {e}", exc_info=True)
|
||||
|
||||
def _format_messages(self, messages: list[dict[str, str]]) -> str:
|
||||
"""
|
||||
Format messages into GenAI semantic convention format (JSON array).
|
||||
|
||||
Returns JSON string representation of message array.
|
||||
"""
|
||||
try:
|
||||
formatted = []
|
||||
for msg in messages:
|
||||
content = msg.get("content", "")
|
||||
# Truncate if needed
|
||||
if isinstance(content, str):
|
||||
content = _truncate_content(content)
|
||||
|
||||
formatted.append(
|
||||
{
|
||||
"role": msg.get("role", "user"),
|
||||
"content": content,
|
||||
}
|
||||
)
|
||||
|
||||
return json.dumps(formatted)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to format input messages: {e}")
|
||||
return "[]"
|
||||
|
||||
def _format_output(
|
||||
self,
|
||||
content: Optional[str],
|
||||
finish_reason: Optional[str],
|
||||
) -> str:
|
||||
"""Format output message into GenAI semantic convention format."""
|
||||
try:
|
||||
if content is None:
|
||||
return "[]"
|
||||
|
||||
# Truncate if needed
|
||||
if isinstance(content, str):
|
||||
content = _truncate_content(content)
|
||||
|
||||
return json.dumps(
|
||||
[
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": content,
|
||||
}
|
||||
]
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to format output message: {e}")
|
||||
return "[]"
|
||||
|
||||
def _extract_system_instructions(self, messages: list[dict[str, str]]) -> Optional[str]:
|
||||
"""Extract system instructions from messages if present."""
|
||||
try:
|
||||
for msg in messages:
|
||||
if msg.get("role") == "system":
|
||||
content = msg.get("content", "")
|
||||
if isinstance(content, str):
|
||||
return _truncate_content(content)
|
||||
return str(content)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to extract system instructions: {e}")
|
||||
return None
|
||||
|
||||
|
||||
class NoOpLLMSpanRecorder:
|
||||
"""No-op span recorder for when tracing is disabled."""
|
||||
|
||||
def record_llm_call(self, **kwargs) -> None:
|
||||
"""No-op."""
|
||||
pass
|
||||
|
||||
|
||||
# Global span recorder instance
|
||||
_span_recorder: Optional[LLMSpanRecorder] = None
|
||||
|
||||
|
||||
def get_span_recorder() -> LLMSpanRecorder | NoOpLLMSpanRecorder:
|
||||
"""Get the global span recorder (NoOp if tracing disabled)."""
|
||||
if _span_recorder is None:
|
||||
return NoOpLLMSpanRecorder()
|
||||
return _span_recorder
|
||||
|
||||
|
||||
def create_span_recorder() -> LLMSpanRecorder:
|
||||
"""Create and set the global span recorder."""
|
||||
global _span_recorder
|
||||
tracer = get_tracer()
|
||||
if tracer is None:
|
||||
raise RuntimeError("Tracing not initialized. Call initialize_tracing() first.")
|
||||
_span_recorder = LLMSpanRecorder(tracer)
|
||||
return _span_recorder
|
||||
@@ -1,13 +0,0 @@
|
||||
from urllib.parse import urlparse, urlunparse
|
||||
|
||||
|
||||
def mask_network_location(url):
|
||||
if not url:
|
||||
return url
|
||||
parsed_url = urlparse(url)
|
||||
masked_network_location = parsed_url.hostname or ""
|
||||
if parsed_url.port:
|
||||
masked_network_location += f":{parsed_url.port}"
|
||||
if parsed_url.username or parsed_url.password:
|
||||
masked_network_location = f"***:***@{masked_network_location}"
|
||||
return urlunparse(parsed_url._replace(netloc=masked_network_location))
|
||||
@@ -1,13 +0,0 @@
|
||||
"""Webhook system for Hindsight API event notifications."""
|
||||
|
||||
from .manager import WebhookManager
|
||||
from .models import ConsolidationEventData, RetainEventData, WebhookConfig, WebhookEvent, WebhookEventType
|
||||
|
||||
__all__ = [
|
||||
"WebhookManager",
|
||||
"WebhookConfig",
|
||||
"WebhookEvent",
|
||||
"WebhookEventType",
|
||||
"ConsolidationEventData",
|
||||
"RetainEventData",
|
||||
]
|
||||
@@ -1,242 +0,0 @@
|
||||
"""Webhook manager for delivering event notifications."""
|
||||
|
||||
import hashlib
|
||||
import hmac
|
||||
import json
|
||||
import logging
|
||||
import uuid
|
||||
from datetime import datetime, timezone
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import asyncpg
|
||||
|
||||
from .models import WebhookConfig, WebhookEvent, WebhookHttpConfig
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from hindsight_api.extensions.tenant import TenantExtension
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Retry delay schedule in seconds: 5 retries after the first attempt.
|
||||
# Fast early retries catch transient failures; later retries handle longer outages.
|
||||
RETRY_DELAYS = [5, 300, 1800, 7200, 18000]
|
||||
MAX_ATTEMPTS = len(RETRY_DELAYS) + 1 # first attempt + len(RETRY_DELAYS) retries
|
||||
|
||||
|
||||
def _fq_table(table: str, schema: str | None = None) -> str:
|
||||
"""Get fully-qualified table name with optional schema prefix."""
|
||||
if schema:
|
||||
return f'"{schema}".{table}'
|
||||
return table
|
||||
|
||||
|
||||
def _parse_http_config(value: str | dict | None) -> WebhookHttpConfig:
|
||||
"""Parse http_config column value (JSONB returned as text or dict) into a model."""
|
||||
if value is None:
|
||||
return WebhookHttpConfig()
|
||||
if isinstance(value, str):
|
||||
return WebhookHttpConfig.model_validate_json(value)
|
||||
return WebhookHttpConfig.model_validate(value)
|
||||
|
||||
|
||||
class WebhookManager:
|
||||
"""
|
||||
Manages webhook registration and event firing.
|
||||
|
||||
Supports both global webhooks (configured via env vars) and per-bank
|
||||
webhooks stored in the database. Deliveries are queued as async_operations
|
||||
tasks (operation_type='webhook_delivery') and picked up by the worker poller.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
pool: asyncpg.Pool,
|
||||
global_webhooks: list[WebhookConfig],
|
||||
tenant_extension: "TenantExtension | None" = None,
|
||||
):
|
||||
self._pool = pool
|
||||
self._global_webhooks = global_webhooks
|
||||
self._tenant_extension = tenant_extension
|
||||
|
||||
def _sign_payload(self, secret: str, payload_bytes: bytes) -> str:
|
||||
"""Compute HMAC-SHA256 signature for a payload."""
|
||||
return "sha256=" + hmac.new(secret.encode(), payload_bytes, hashlib.sha256).hexdigest()
|
||||
|
||||
async def fire_event(self, event: WebhookEvent, schema: str | None = None) -> None:
|
||||
"""
|
||||
Queue webhook deliveries for an event as async_operations tasks.
|
||||
|
||||
Loads per-bank and global webhooks, inserts pending webhook_delivery tasks for
|
||||
any webhook whose event_types list matches the fired event type. The worker
|
||||
poller picks these up and calls MemoryEngine._handle_webhook_delivery().
|
||||
|
||||
Args:
|
||||
event: The event to deliver.
|
||||
schema: Database schema (for multi-tenant). None = default schema.
|
||||
"""
|
||||
webhook_table = _fq_table("webhooks", schema)
|
||||
ops_table = _fq_table("async_operations", schema)
|
||||
now = datetime.now(timezone.utc)
|
||||
payload_str = event.model_dump_json()
|
||||
|
||||
try:
|
||||
# Load per-bank webhooks from DB (bank-specific + global NULL rows)
|
||||
rows = await self._pool.fetch(
|
||||
f"""
|
||||
SELECT id, bank_id, url, secret, event_types, enabled, http_config::text
|
||||
FROM {webhook_table}
|
||||
WHERE (bank_id = $1 OR bank_id IS NULL) AND enabled = true
|
||||
""",
|
||||
event.bank_id,
|
||||
)
|
||||
|
||||
db_webhooks = [
|
||||
WebhookConfig(
|
||||
id=str(row["id"]),
|
||||
bank_id=row["bank_id"],
|
||||
url=row["url"],
|
||||
secret=row["secret"],
|
||||
event_types=list(row["event_types"]) if row["event_types"] else [],
|
||||
enabled=row["enabled"],
|
||||
http_config=_parse_http_config(row["http_config"]),
|
||||
)
|
||||
for row in rows
|
||||
]
|
||||
|
||||
# Merge with global webhooks from env config
|
||||
all_webhooks = self._global_webhooks + db_webhooks
|
||||
matched = 0
|
||||
|
||||
for webhook in all_webhooks:
|
||||
if not webhook.enabled:
|
||||
continue
|
||||
if event.event.value not in webhook.event_types:
|
||||
continue
|
||||
|
||||
operation_id = uuid.uuid4()
|
||||
webhook_id = webhook.id if webhook.id else None
|
||||
|
||||
task_payload = json.dumps(
|
||||
{
|
||||
"type": "webhook_delivery",
|
||||
"operation_id": str(operation_id),
|
||||
"bank_id": event.bank_id,
|
||||
"url": webhook.url,
|
||||
"secret": webhook.secret,
|
||||
"event_type": event.event.value,
|
||||
"payload": payload_str,
|
||||
"webhook_id": webhook_id,
|
||||
"http_config": webhook.http_config.model_dump(),
|
||||
}
|
||||
)
|
||||
|
||||
await self._pool.execute(
|
||||
f"""
|
||||
INSERT INTO {ops_table}
|
||||
(operation_id, bank_id, operation_type, status, task_payload, result_metadata, created_at, updated_at)
|
||||
VALUES ($1, $2, 'webhook_delivery', 'pending', $3::jsonb, '{{}}'::jsonb, $4, $4)
|
||||
""",
|
||||
operation_id,
|
||||
event.bank_id,
|
||||
task_payload,
|
||||
now,
|
||||
)
|
||||
matched += 1
|
||||
|
||||
logger.debug(f"Fired webhook event {event.event} for bank {event.bank_id}: {matched} delivery(ies) queued")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to queue webhook deliveries for event {event.event}: {e}")
|
||||
|
||||
async def fire_event_with_conn(
|
||||
self, event: WebhookEvent, conn: asyncpg.Connection, schema: str | None = None
|
||||
) -> None:
|
||||
"""
|
||||
Queue webhook deliveries within an existing database connection/transaction.
|
||||
|
||||
Identical to fire_event() but uses the provided connection instead of acquiring
|
||||
one from the pool. Use this to atomically insert delivery tasks in the same
|
||||
transaction as the primary operation (transactional outbox pattern).
|
||||
|
||||
Args:
|
||||
event: The event to deliver.
|
||||
conn: Existing asyncpg connection (may be inside an active transaction).
|
||||
schema: Database schema (for multi-tenant). None = default schema.
|
||||
"""
|
||||
webhook_table = _fq_table("webhooks", schema)
|
||||
ops_table = _fq_table("async_operations", schema)
|
||||
now = datetime.now(timezone.utc)
|
||||
payload_str = event.model_dump_json()
|
||||
|
||||
try:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT id, bank_id, url, secret, event_types, enabled, http_config::text
|
||||
FROM {webhook_table}
|
||||
WHERE (bank_id = $1 OR bank_id IS NULL) AND enabled = true
|
||||
""",
|
||||
event.bank_id,
|
||||
)
|
||||
|
||||
db_webhooks = [
|
||||
WebhookConfig(
|
||||
id=str(row["id"]),
|
||||
bank_id=row["bank_id"],
|
||||
url=row["url"],
|
||||
secret=row["secret"],
|
||||
event_types=list(row["event_types"]) if row["event_types"] else [],
|
||||
enabled=row["enabled"],
|
||||
http_config=_parse_http_config(row["http_config"]),
|
||||
)
|
||||
for row in rows
|
||||
]
|
||||
|
||||
all_webhooks = self._global_webhooks + db_webhooks
|
||||
matched = 0
|
||||
|
||||
for webhook in all_webhooks:
|
||||
if not webhook.enabled:
|
||||
continue
|
||||
if event.event.value not in webhook.event_types:
|
||||
continue
|
||||
|
||||
operation_id = uuid.uuid4()
|
||||
webhook_id = webhook.id if webhook.id else None
|
||||
|
||||
task_payload = json.dumps(
|
||||
{
|
||||
"type": "webhook_delivery",
|
||||
"operation_id": str(operation_id),
|
||||
"bank_id": event.bank_id,
|
||||
"url": webhook.url,
|
||||
"secret": webhook.secret,
|
||||
"event_type": event.event.value,
|
||||
"payload": payload_str,
|
||||
"webhook_id": webhook_id,
|
||||
"http_config": webhook.http_config.model_dump(),
|
||||
}
|
||||
)
|
||||
|
||||
await conn.execute(
|
||||
f"""
|
||||
INSERT INTO {ops_table}
|
||||
(operation_id, bank_id, operation_type, status, task_payload, result_metadata, created_at, updated_at)
|
||||
VALUES ($1, $2, 'webhook_delivery', 'pending', $3::jsonb, '{{}}'::jsonb, $4, $4)
|
||||
""",
|
||||
operation_id,
|
||||
event.bank_id,
|
||||
task_payload,
|
||||
now,
|
||||
)
|
||||
matched += 1
|
||||
|
||||
logger.debug(
|
||||
f"Fired webhook event {event.event} for bank {event.bank_id}: {matched} delivery(ies) queued (in-transaction)"
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Failed to queue webhook deliveries (in-transaction) for event {event.event}: {e}. "
|
||||
"CRITICAL: The enclosing database transaction is now aborted and will roll back all changes."
|
||||
)
|
||||
raise
|
||||
@@ -1,51 +0,0 @@
|
||||
"""Pydantic models for the webhook system."""
|
||||
|
||||
from datetime import datetime
|
||||
from enum import StrEnum
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
class WebhookEventType(StrEnum):
|
||||
CONSOLIDATION_COMPLETED = "consolidation.completed"
|
||||
RETAIN_COMPLETED = "retain.completed"
|
||||
|
||||
|
||||
class ConsolidationEventData(BaseModel):
|
||||
observations_created: int | None = None
|
||||
observations_updated: int | None = None
|
||||
observations_deleted: int | None = None
|
||||
error_message: str | None = None
|
||||
|
||||
|
||||
class RetainEventData(BaseModel):
|
||||
document_id: str | None = None
|
||||
tags: list[str] | None = None
|
||||
|
||||
|
||||
class WebhookEvent(BaseModel):
|
||||
event: WebhookEventType
|
||||
bank_id: str
|
||||
operation_id: str
|
||||
status: str # "completed" or "failed"
|
||||
timestamp: datetime
|
||||
data: ConsolidationEventData | RetainEventData
|
||||
|
||||
|
||||
class WebhookHttpConfig(BaseModel):
|
||||
"""HTTP delivery configuration for a webhook."""
|
||||
|
||||
method: str = Field(default="POST", description="HTTP method: GET or POST")
|
||||
timeout_seconds: int = Field(default=30, description="HTTP request timeout in seconds")
|
||||
headers: dict[str, str] = Field(default_factory=dict, description="Custom HTTP headers")
|
||||
params: dict[str, str] = Field(default_factory=dict, description="Custom HTTP query parameters")
|
||||
|
||||
|
||||
class WebhookConfig(BaseModel):
|
||||
id: str
|
||||
bank_id: str | None
|
||||
url: str
|
||||
secret: str | None
|
||||
event_types: list[str]
|
||||
enabled: bool
|
||||
http_config: WebhookHttpConfig = Field(default_factory=WebhookHttpConfig)
|
||||
@@ -1,9 +0,0 @@
|
||||
from datetime import datetime
|
||||
|
||||
|
||||
class RetryTaskAt(Exception):
|
||||
"""Raise from a task handler to schedule a retry at a specific time."""
|
||||
|
||||
def __init__(self, retry_at: datetime, message: str = ""):
|
||||
self.retry_at = retry_at
|
||||
super().__init__(message)
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user