Compare commits
3
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
344ac8fae8 | ||
|
|
4b0c617ecf | ||
|
|
0a04770450 |
+2
-19
@@ -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 (1M context window)
|
||||
# HINDSIGHT_API_LLM_PROVIDER=minimax
|
||||
# HINDSIGHT_API_LLM_API_KEY=your-minimax-api-key
|
||||
# HINDSIGHT_API_LLM_MODEL=MiniMax-M2.7
|
||||
|
||||
# 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
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -1,99 +0,0 @@
|
||||
name: Release Integration
|
||||
|
||||
on:
|
||||
push:
|
||||
tags:
|
||||
- 'integrations/**'
|
||||
|
||||
jobs:
|
||||
publish:
|
||||
runs-on: ubuntu-latest
|
||||
permissions:
|
||||
id-token: write # for PyPI trusted publishing
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
|
||||
- name: Extract integration info
|
||||
id: info
|
||||
run: |
|
||||
# refs/tags/integrations/litellm/v0.1.0 → integration=litellm, version=0.1.0
|
||||
TAG="${GITHUB_REF#refs/tags/}"
|
||||
INTEGRATION=$(echo "$TAG" | cut -d'/' -f2)
|
||||
VERSION=$(echo "$TAG" | cut -d'/' -f3 | sed 's/^v//')
|
||||
echo "integration=$INTEGRATION" >> $GITHUB_OUTPUT
|
||||
echo "version=$VERSION" >> $GITHUB_OUTPUT
|
||||
echo "tag=$TAG" >> $GITHUB_OUTPUT
|
||||
echo "Integration: $INTEGRATION, Version: $VERSION"
|
||||
|
||||
- name: Detect integration type
|
||||
id: type
|
||||
run: |
|
||||
if [ -f "hindsight-integrations/${{ steps.info.outputs.integration }}/pyproject.toml" ]; then
|
||||
echo "type=python" >> $GITHUB_OUTPUT
|
||||
else
|
||||
echo "type=typescript" >> $GITHUB_OUTPUT
|
||||
fi
|
||||
|
||||
# ── Python integrations (litellm, pydantic-ai, crewai) ──────────────────
|
||||
|
||||
- name: Install uv
|
||||
if: steps.type.outputs.type == 'python'
|
||||
uses: astral-sh/setup-uv@v7
|
||||
with:
|
||||
enable-cache: true
|
||||
|
||||
- name: Set up Python
|
||||
if: steps.type.outputs.type == 'python'
|
||||
uses: actions/setup-python@v5
|
||||
with:
|
||||
python-version-file: ".python-version"
|
||||
|
||||
- name: Build Python package
|
||||
if: steps.type.outputs.type == 'python'
|
||||
working-directory: ./hindsight-integrations/${{ steps.info.outputs.integration }}
|
||||
run: uv build --out-dir dist
|
||||
|
||||
- name: Publish Python package to PyPI
|
||||
if: steps.type.outputs.type == 'python'
|
||||
uses: pypa/gh-action-pypi-publish@release/v1
|
||||
with:
|
||||
packages-dir: ./hindsight-integrations/${{ steps.info.outputs.integration }}/dist
|
||||
skip-existing: true
|
||||
|
||||
# ── TypeScript integrations (ai-sdk, chat, openclaw) ────────────────────
|
||||
|
||||
- name: Set up Node.js
|
||||
if: steps.type.outputs.type == 'typescript'
|
||||
uses: actions/setup-node@v6
|
||||
with:
|
||||
node-version: '22'
|
||||
registry-url: 'https://registry.npmjs.org'
|
||||
|
||||
- name: Install dependencies
|
||||
if: steps.type.outputs.type == 'typescript'
|
||||
working-directory: ./hindsight-integrations/${{ steps.info.outputs.integration }}
|
||||
run: npm ci
|
||||
|
||||
- name: Build TypeScript package
|
||||
if: steps.type.outputs.type == 'typescript'
|
||||
working-directory: ./hindsight-integrations/${{ steps.info.outputs.integration }}
|
||||
run: npm run build
|
||||
|
||||
- name: Publish TypeScript package to npm
|
||||
if: steps.type.outputs.type == 'typescript'
|
||||
working-directory: ./hindsight-integrations/${{ steps.info.outputs.integration }}
|
||||
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 }}
|
||||
+156
-61
@@ -13,15 +13,15 @@ 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
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v6
|
||||
uses: actions/setup-python@v5
|
||||
with:
|
||||
python-version-file: ".python-version"
|
||||
|
||||
@@ -30,39 +30,29 @@ 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
|
||||
working-directory: ./hindsight
|
||||
run: uv build --out-dir dist
|
||||
|
||||
- name: Build hindsight-all-slim
|
||||
working-directory: ./hindsight-all-slim
|
||||
- name: Build hindsight-litellm
|
||||
working-directory: ./hindsight-integrations/litellm
|
||||
run: uv build --out-dir dist
|
||||
|
||||
- name: Build hindsight-embed
|
||||
working-directory: ./hindsight-embed
|
||||
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:
|
||||
@@ -72,13 +62,13 @@ jobs:
|
||||
- name: Publish hindsight-all to PyPI
|
||||
uses: pypa/gh-action-pypi-publish@release/v1
|
||||
with:
|
||||
packages-dir: ./hindsight-all/dist
|
||||
packages-dir: ./hindsight/dist
|
||||
skip-existing: true
|
||||
|
||||
- name: Publish hindsight-all-slim to PyPI
|
||||
- name: Publish hindsight-litellm to PyPI
|
||||
uses: pypa/gh-action-pypi-publish@release/v1
|
||||
with:
|
||||
packages-dir: ./hindsight-all-slim/dist
|
||||
packages-dir: ./hindsight-integrations/litellm/dist
|
||||
skip-existing: true
|
||||
|
||||
- name: Publish hindsight-embed to PyPI
|
||||
@@ -89,15 +79,14 @@ jobs:
|
||||
|
||||
# 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/*
|
||||
retention-days: 1
|
||||
|
||||
@@ -106,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'
|
||||
@@ -144,21 +133,119 @@ 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
|
||||
retention-days: 1
|
||||
|
||||
release-openclaw-integration:
|
||||
runs-on: ubuntu-latest
|
||||
environment: npm
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
- name: Set up Node.js
|
||||
uses: actions/setup-node@v4
|
||||
with:
|
||||
node-version: '22'
|
||||
registry-url: 'https://registry.npmjs.org'
|
||||
|
||||
- name: Install dependencies
|
||||
working-directory: ./hindsight-integrations/openclaw
|
||||
run: npm ci
|
||||
|
||||
- name: Build
|
||||
working-directory: ./hindsight-integrations/openclaw
|
||||
run: npm run build
|
||||
|
||||
- name: Publish to npm
|
||||
working-directory: ./hindsight-integrations/openclaw
|
||||
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/openclaw
|
||||
run: npm pack
|
||||
|
||||
- name: Upload artifacts
|
||||
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@v4
|
||||
|
||||
- name: Set up Node.js
|
||||
uses: actions/setup-node@v4
|
||||
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@v4
|
||||
with:
|
||||
name: ai-sdk-integration
|
||||
path: hindsight-integrations/ai-sdk/*.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'
|
||||
@@ -181,14 +268,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
|
||||
@@ -206,7 +290,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
|
||||
@@ -229,13 +313,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
|
||||
@@ -253,7 +333,7 @@ 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 }}
|
||||
@@ -294,7 +374,7 @@ jobs:
|
||||
PRELOAD_ML_MODELS=false
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
- name: Free Disk Space
|
||||
uses: jlumbroso/free-disk-space@main
|
||||
@@ -308,13 +388,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 }}
|
||||
@@ -326,7 +406,7 @@ 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: |
|
||||
@@ -342,7 +422,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
|
||||
@@ -361,7 +441,7 @@ jobs:
|
||||
|
||||
# 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
|
||||
@@ -379,7 +459,7 @@ jobs:
|
||||
packages: write
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
- name: Install Helm
|
||||
uses: azure/setup-helm@v4
|
||||
@@ -399,7 +479,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
|
||||
@@ -407,55 +487,67 @@ jobs:
|
||||
|
||||
create-github-release:
|
||||
runs-on: ubuntu-latest
|
||||
needs: [release-python-packages, release-typescript-client, release-control-plane, release-rust-cli, release-docker-images, release-helm-chart]
|
||||
needs: [release-python-packages, release-typescript-client, release-openclaw-integration, release-ai-sdk-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
|
||||
run: echo "VERSION=${GITHUB_REF#refs/tags/v}" >> $GITHUB_OUTPUT
|
||||
|
||||
- name: Download Python packages
|
||||
uses: actions/download-artifact@v8
|
||||
uses: actions/download-artifact@v4
|
||||
with:
|
||||
name: python-packages
|
||||
path: ./artifacts/python-packages
|
||||
|
||||
- name: Download TypeScript client
|
||||
uses: actions/download-artifact@v8
|
||||
uses: actions/download-artifact@v4
|
||||
with:
|
||||
name: typescript-client
|
||||
path: ./artifacts/typescript-client
|
||||
|
||||
- name: Download OpenClaw Integration
|
||||
uses: actions/download-artifact@v4
|
||||
with:
|
||||
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 Control Plane
|
||||
uses: actions/download-artifact@v8
|
||||
uses: actions/download-artifact@v4
|
||||
with:
|
||||
name: control-plane
|
||||
path: ./artifacts/control-plane
|
||||
|
||||
- name: Download Rust CLI (Linux)
|
||||
uses: actions/download-artifact@v8
|
||||
uses: actions/download-artifact@v4
|
||||
with:
|
||||
name: rust-cli-hindsight-linux-amd64
|
||||
path: ./artifacts/rust-cli-linux
|
||||
|
||||
- name: Download Rust CLI (macOS Intel)
|
||||
uses: actions/download-artifact@v8
|
||||
uses: actions/download-artifact@v4
|
||||
with:
|
||||
name: rust-cli-hindsight-darwin-amd64
|
||||
path: ./artifacts/rust-cli-darwin-amd64
|
||||
|
||||
- name: Download Rust CLI (macOS ARM)
|
||||
uses: actions/download-artifact@v8
|
||||
uses: actions/download-artifact@v4
|
||||
with:
|
||||
name: rust-cli-hindsight-darwin-arm64
|
||||
path: ./artifacts/rust-cli-darwin-arm64
|
||||
|
||||
- name: Download Helm chart
|
||||
uses: actions/download-artifact@v8
|
||||
uses: actions/download-artifact@v4
|
||||
with:
|
||||
name: helm-chart
|
||||
path: ./artifacts/helm-chart
|
||||
@@ -465,13 +557,16 @@ 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-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
|
||||
# Control Plane
|
||||
cp artifacts/control-plane/*.tgz release-assets/ || true
|
||||
# Rust CLI binaries
|
||||
|
||||
+213
-894
File diff suppressed because it is too large
Load Diff
@@ -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
|
||||
|
||||
@@ -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)
|
||||
@@ -57,22 +57,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 +74,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 +94,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 +104,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 +147,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 +238,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 +266,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,7 +36,7 @@ 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.
|
||||
|
||||
@@ -70,7 +69,7 @@ docker run --rm -it --pull always -p 8888:8888 -p 9999:9999 \
|
||||
>API: http://localhost:8888
|
||||
>UI: http://localhost:9999
|
||||
|
||||
You can modify the LLM provider by setting `HINDSIGHT_API_LLM_PROVIDER`. Valid options are `openai`, `anthropic`, `gemini`, `groq`, `ollama`, `lmstudio`, and `minimax`. 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`, and `lmstudio`. The documentation provides more details on [supported models](https://hindsight.vectorize.io/developer/models).
|
||||
|
||||
|
||||
|
||||
@@ -182,7 +181,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 +307,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,139 +0,0 @@
|
||||
{
|
||||
"version": "5",
|
||||
"specifiers": {
|
||||
"jsr:@std/assert@^1.0.17": "1.0.19",
|
||||
"jsr:@std/assert@^1.0.19": "1.0.19",
|
||||
"jsr:@std/expect@*": "1.0.18",
|
||||
"jsr:@std/internal@^1.0.12": "1.0.12",
|
||||
"jsr:@std/path@^1.1.4": "1.1.4",
|
||||
"jsr:@std/testing@*": "1.0.17"
|
||||
},
|
||||
"jsr": {
|
||||
"@std/[email protected]": {
|
||||
"integrity": "eaada96ee120cb980bc47e040f82814d786fe8162ecc53c91d8df60b8755991e",
|
||||
"dependencies": [
|
||||
"jsr:@std/internal"
|
||||
]
|
||||
},
|
||||
"@std/[email protected]": {
|
||||
"integrity": "8566eab35200466f8609eb7e7aed062ed0db314e9a258d5d201b1b8997ce801a",
|
||||
"dependencies": [
|
||||
"jsr:@std/assert@^1.0.19",
|
||||
"jsr:@std/internal",
|
||||
"jsr:@std/path"
|
||||
]
|
||||
},
|
||||
"@std/[email protected]": {
|
||||
"integrity": "972a634fd5bc34b242024402972cd5143eac68d8dffaca5eaa4dba30ce17b027"
|
||||
},
|
||||
"@std/[email protected]": {
|
||||
"integrity": "1d2d43f39efb1b42f0b1882a25486647cb851481862dc7313390b2bb044314b5",
|
||||
"dependencies": [
|
||||
"jsr:@std/internal"
|
||||
]
|
||||
},
|
||||
"@std/[email protected]": {
|
||||
"integrity": "87bdc2700fa98249d48a17cd72413352d3d3680dcfbdb64947fd0982d6bbf681",
|
||||
"dependencies": [
|
||||
"jsr:@std/assert@^1.0.17",
|
||||
"jsr:@std/internal"
|
||||
]
|
||||
}
|
||||
},
|
||||
"workspace": {
|
||||
"members": {
|
||||
"hindsight-clients/typescript": {
|
||||
"packageJson": {
|
||||
"dependencies": [
|
||||
"npm:@hey-api/[email protected]",
|
||||
"npm:@types/jest@29",
|
||||
"npm:@types/node@20",
|
||||
"npm:jest@29",
|
||||
"npm:ts-jest@29",
|
||||
"npm:tsup@^8.5.1",
|
||||
"npm:typescript@5"
|
||||
]
|
||||
}
|
||||
},
|
||||
"hindsight-control-plane": {
|
||||
"packageJson": {
|
||||
"dependencies": [
|
||||
"npm:@eslint/eslintrc@^3.3.3",
|
||||
"npm:@eslint/js@^9.39.2",
|
||||
"npm:@radix-ui/react-alert-dialog@^1.1.15",
|
||||
"npm:@radix-ui/react-checkbox@^1.3.3",
|
||||
"npm:@radix-ui/react-dialog@^1.1.15",
|
||||
"npm:@radix-ui/react-dropdown-menu@^2.1.16",
|
||||
"npm:@radix-ui/react-label@^2.1.8",
|
||||
"npm:@radix-ui/react-popover@^1.1.15",
|
||||
"npm:@radix-ui/react-radio-group@^1.3.8",
|
||||
"npm:@radix-ui/react-select@^2.2.6",
|
||||
"npm:@radix-ui/react-slider@^1.3.6",
|
||||
"npm:@radix-ui/react-slot@^1.2.4",
|
||||
"npm:@radix-ui/react-switch@^1.2.6",
|
||||
"npm:@radix-ui/react-tabs@^1.1.13",
|
||||
"npm:@radix-ui/react-tooltip@^1.2.8",
|
||||
"npm:@tailwindcss/postcss@^4.1.17",
|
||||
"npm:@tailwindcss/typography@~0.5.19",
|
||||
"npm:@types/cytoscape@^3.21.9",
|
||||
"npm:@types/node@^24.10.0",
|
||||
"npm:@types/react-dom@^19.2.2",
|
||||
"npm:@types/react@^19.2.2",
|
||||
"npm:autoprefixer@^10.4.21",
|
||||
"npm:class-variance-authority@~0.7.1",
|
||||
"npm:clsx@^2.1.1",
|
||||
"npm:cmdk@^1.1.1",
|
||||
"npm:cytoscape-fcose@^2.2.0",
|
||||
"npm:cytoscape@^3.33.1",
|
||||
"npm:eslint-config-next@^16.0.1",
|
||||
"npm:eslint-plugin-react-hooks@^7.0.1",
|
||||
"npm:eslint-plugin-react@^7.37.5",
|
||||
"npm:eslint@^9.39.1",
|
||||
"npm:[email protected]",
|
||||
"npm:next-themes@~0.4.6",
|
||||
"npm:next@^16.1.6",
|
||||
"npm:postcss@^8.5.6",
|
||||
"npm:prettier@^3.7.4",
|
||||
"npm:react-chrono@^2.9.1",
|
||||
"npm:react-dom@^19.2.0",
|
||||
"npm:react-markdown@^10.1.0",
|
||||
"npm:react18-json-view@~0.2.9",
|
||||
"npm:react@^19.2.0",
|
||||
"npm:recharts@^3.5.1",
|
||||
"npm:remark-gfm@^4.0.1",
|
||||
"npm:sonner@^2.0.7",
|
||||
"npm:tailwind-merge@^3.4.0",
|
||||
"npm:tailwindcss-animate@^1.0.7",
|
||||
"npm:tailwindcss@^4.1.17",
|
||||
"npm:[email protected]",
|
||||
"npm:typescript-eslint@^8.50.0",
|
||||
"npm:typescript@^5.9.3"
|
||||
]
|
||||
}
|
||||
},
|
||||
"hindsight-docs": {
|
||||
"packageJson": {
|
||||
"dependencies": [
|
||||
"npm:@docusaurus/[email protected]",
|
||||
"npm:@docusaurus/[email protected]",
|
||||
"npm:@docusaurus/[email protected]",
|
||||
"npm:@docusaurus/theme-common@^3.9.2",
|
||||
"npm:@docusaurus/theme-mermaid@^3.9.2",
|
||||
"npm:@docusaurus/[email protected]",
|
||||
"npm:@docusaurus/[email protected]",
|
||||
"npm:@easyops-cn/docusaurus-search-local@~0.52.2",
|
||||
"npm:@mdx-js/react@3",
|
||||
"npm:clsx@2",
|
||||
"npm:prism-react-renderer@^2.3.0",
|
||||
"npm:raw-loader@^4.0.2",
|
||||
"npm:react-dom@19",
|
||||
"npm:react-icons@^5.6.0",
|
||||
"npm:react@19",
|
||||
"npm:redocusaurus@^2.5.0",
|
||||
"npm:typescript@~5.6.2"
|
||||
]
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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:
|
||||
@@ -42,22 +42,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 +112,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,11 +166,6 @@ 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)
|
||||
@@ -323,11 +317,6 @@ 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)
|
||||
|
||||
@@ -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,8 +97,7 @@ fi
|
||||
if [ "$ENABLE_CP" = "true" ]; then
|
||||
echo "🎛️ Starting Control Plane..."
|
||||
cd /app/control-plane
|
||||
export HOSTNAME="${HINDSIGHT_CP_HOSTNAME:-0.0.0.0}"
|
||||
PORT="${HINDSIGHT_CP_PORT:-9999}" node server.js &
|
||||
PORT=9999 node server.js &
|
||||
CP_PID=$!
|
||||
PIDS+=($CP_PID)
|
||||
else
|
||||
@@ -125,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"
|
||||
|
||||
+10
-44
@@ -13,9 +13,9 @@
|
||||
# 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)
|
||||
# GROQ_API_KEY - Required for API/standalone images (LLM verification)
|
||||
# HINDSIGHT_API_LLM_PROVIDER - LLM provider (default: groq)
|
||||
# HINDSIGHT_API_LLM_MODEL - LLM model (default: llama-3.3-70b-versatile)
|
||||
# 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)
|
||||
@@ -34,7 +34,7 @@
|
||||
# ./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 GROQ_API_KEY=gsk_xxx
|
||||
# export HINDSIGHT_API_EMBEDDINGS_PROVIDER=openai
|
||||
# export HINDSIGHT_API_EMBEDDINGS_OPENAI_API_KEY=sk-xxx
|
||||
# export HINDSIGHT_API_RERANKER_PROVIDER=cohere
|
||||
@@ -49,9 +49,6 @@
|
||||
|
||||
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'
|
||||
@@ -63,8 +60,8 @@ 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}"
|
||||
LLM_PROVIDER="${HINDSIGHT_API_LLM_PROVIDER:-groq}"
|
||||
LLM_MODEL="${HINDSIGHT_API_LLM_MODEL:-llama-3.3-70b-versatile}"
|
||||
|
||||
# Validate arguments
|
||||
if [ -z "$IMAGE" ]; then
|
||||
@@ -91,9 +88,9 @@ else
|
||||
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"
|
||||
if [ "$NEEDS_LLM" = true ] && [ -z "${GROQ_API_KEY:-}" ]; then
|
||||
echo -e "${RED}Error: GROQ_API_KEY environment variable is required for API/standalone images${NC}"
|
||||
echo "Set it with: export GROQ_API_KEY=your-api-key"
|
||||
exit 2
|
||||
fi
|
||||
|
||||
@@ -126,25 +123,9 @@ 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_API_KEY=${GROQ_API_KEY}"
|
||||
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}"
|
||||
@@ -181,21 +162,6 @@ for i in $(seq 1 "$TIMEOUT"); do
|
||||
echo "=== Health Response ==="
|
||||
curl -s "http://localhost:${HEALTH_PORT}${HEALTH_PATH}" | python3 -m json.tool 2>/dev/null || curl -s "http://localhost:${HEALTH_PORT}${HEALTH_PATH}"
|
||||
echo ""
|
||||
|
||||
# Run retain/recall smoke test for API targets
|
||||
if [ "$TARGET" != "cp-only" ]; then
|
||||
echo ""
|
||||
echo "=== Retain/Recall Smoke Test ==="
|
||||
if ! "$REPO_ROOT/scripts/smoke-test-slim.sh" "http://localhost:${HEALTH_PORT}"; then
|
||||
echo ""
|
||||
echo "=== Container Logs (last 50 lines) ==="
|
||||
docker logs "$CONTAINER_NAME" 2>&1 | tail -50
|
||||
echo ""
|
||||
echo -e "${RED}Smoke test FAILED${NC}"
|
||||
exit 1
|
||||
fi
|
||||
fi
|
||||
|
||||
echo ""
|
||||
echo "=== Container Logs (last 50 lines) ==="
|
||||
docker logs "$CONTAINER_NAME" 2>&1 | tail -50
|
||||
|
||||
@@ -6,17 +6,24 @@
|
||||
# It expects API keys to be set in environment variables.
|
||||
#
|
||||
# Usage:
|
||||
# export GROQ_API_KEY=gsk_xxx
|
||||
# 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
|
||||
# GROQ_API_KEY=gsk_xxx OPENAI_API_KEY=sk_xxx COHERE_API_KEY=xxx ./docker/test-slim-local.sh
|
||||
#
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
# Check for required API keys
|
||||
if [ -z "${GROQ_API_KEY:-}" ]; then
|
||||
echo "❌ Error: GROQ_API_KEY environment variable is required"
|
||||
echo "Set it with: export GROQ_API_KEY=gsk_xxx"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
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"
|
||||
@@ -34,10 +41,7 @@ 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
|
||||
# Set up external providers
|
||||
export HINDSIGHT_API_EMBEDDINGS_PROVIDER=openai
|
||||
export HINDSIGHT_API_EMBEDDINGS_OPENAI_API_KEY=$OPENAI_API_KEY
|
||||
export HINDSIGHT_API_RERANKER_PROVIDER=cohere
|
||||
|
||||
@@ -2,8 +2,8 @@ apiVersion: v2
|
||||
name: hindsight
|
||||
description: Hindsight helm chart
|
||||
type: application
|
||||
version: 0.4.19
|
||||
appVersion: "0.4.19"
|
||||
version: 0.4.10
|
||||
appVersion: "0.4.10"
|
||||
keywords:
|
||||
- ai
|
||||
- memory
|
||||
|
||||
@@ -1,33 +0,0 @@
|
||||
[build-system]
|
||||
requires = ["setuptools>=61"]
|
||||
build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "hindsight-all-slim"
|
||||
version = "0.4.19"
|
||||
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,423 +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.
|
||||
|
||||
Provides methods to create, delete, and manage memory banks.
|
||||
"""
|
||||
|
||||
def __init__(self, client: Hindsight):
|
||||
self._client = client
|
||||
|
||||
def create(
|
||||
self,
|
||||
bank_id: str,
|
||||
name: str | None = None,
|
||||
mission: str | None = None,
|
||||
disposition: dict[str, Any] | None = None,
|
||||
) -> Any:
|
||||
"""Create a new bank.
|
||||
|
||||
Args:
|
||||
bank_id: Unique identifier for the bank.
|
||||
name: Optional display name for the bank.
|
||||
mission: Optional mission statement for the bank.
|
||||
disposition: Optional disposition configuration dict.
|
||||
|
||||
Returns:
|
||||
Bank creation response from the API.
|
||||
"""
|
||||
return self._client.create_bank(
|
||||
bank_id=bank_id,
|
||||
name=name,
|
||||
mission=mission,
|
||||
disposition=disposition,
|
||||
)
|
||||
|
||||
def delete(self, bank_id: str) -> Any:
|
||||
"""Delete a bank.
|
||||
|
||||
Args:
|
||||
bank_id: The ID of the bank to delete.
|
||||
|
||||
Returns:
|
||||
Deletion response from the API.
|
||||
"""
|
||||
return self._client.delete_bank(bank_id=bank_id)
|
||||
|
||||
def set_mission(self, bank_id: str, mission: str) -> Any:
|
||||
"""Set or update the mission for a bank.
|
||||
|
||||
Args:
|
||||
bank_id: The ID of the bank.
|
||||
mission: The mission statement to set.
|
||||
|
||||
Returns:
|
||||
API response confirming the update.
|
||||
"""
|
||||
return self._client.set_mission(bank_id=bank_id, mission=mission)
|
||||
|
||||
def set_disposition(self, bank_id: str, disposition: dict[str, Any]) -> Any:
|
||||
"""Set or update the disposition for a bank.
|
||||
|
||||
Args:
|
||||
bank_id: The ID of the bank.
|
||||
disposition: The disposition configuration dict.
|
||||
|
||||
Returns:
|
||||
API response confirming the update.
|
||||
"""
|
||||
return self._client.set_disposition(bank_id=bank_id, disposition=disposition)
|
||||
|
||||
def list(self) -> Any:
|
||||
"""List all banks.
|
||||
|
||||
Returns:
|
||||
List of banks from the API.
|
||||
"""
|
||||
from hindsight_client.hindsight_client import _run_async
|
||||
|
||||
return _run_async(self._client._banks_api.list_banks())
|
||||
|
||||
|
||||
class MentalModelsAPI:
|
||||
"""Namespace for mental model operations.
|
||||
|
||||
Mental models are reusable knowledge structures that guide agent behavior.
|
||||
"""
|
||||
|
||||
def __init__(self, client: Hindsight):
|
||||
self._client = client
|
||||
|
||||
def create(
|
||||
self,
|
||||
bank_id: str,
|
||||
name: str,
|
||||
content: str,
|
||||
tags: list[str] | None = None,
|
||||
) -> Any:
|
||||
"""Create a new mental model.
|
||||
|
||||
Args:
|
||||
bank_id: The ID of the bank to add the model to.
|
||||
name: Name for the mental model.
|
||||
content: The content/instructions for the mental model.
|
||||
tags: Optional list of tags for categorization.
|
||||
|
||||
Returns:
|
||||
Creation response from the API.
|
||||
"""
|
||||
return self._client.create_mental_model(
|
||||
bank_id=bank_id,
|
||||
name=name,
|
||||
content=content,
|
||||
tags=tags,
|
||||
)
|
||||
|
||||
def list(self, bank_id: str, tags: list[str] | None = None) -> Any:
|
||||
"""List all mental models for a bank.
|
||||
|
||||
Args:
|
||||
bank_id: The ID of the bank.
|
||||
tags: Optional filter by tags.
|
||||
|
||||
Returns:
|
||||
List of mental models.
|
||||
"""
|
||||
return self._client.list_mental_models(bank_id=bank_id, tags=tags)
|
||||
|
||||
def get(self, bank_id: str, mental_model_id: str) -> Any:
|
||||
"""Get a specific mental model.
|
||||
|
||||
Args:
|
||||
bank_id: The ID of the bank.
|
||||
mental_model_id: The ID of the mental model.
|
||||
|
||||
Returns:
|
||||
The mental model details.
|
||||
"""
|
||||
return self._client.get_mental_model(bank_id=bank_id, mental_model_id=mental_model_id)
|
||||
|
||||
def refresh(self, bank_id: str, mental_model_id: str) -> Any:
|
||||
"""Refresh a mental model.
|
||||
|
||||
Args:
|
||||
bank_id: The ID of the bank.
|
||||
mental_model_id: The ID of the mental model to refresh.
|
||||
|
||||
Returns:
|
||||
Refresh response from the API.
|
||||
"""
|
||||
return self._client.refresh_mental_model(bank_id=bank_id, mental_model_id=mental_model_id)
|
||||
|
||||
def update(
|
||||
self,
|
||||
bank_id: str,
|
||||
mental_model_id: str,
|
||||
name: str | None = None,
|
||||
content: str | None = None,
|
||||
tags: list[str] | None = None,
|
||||
) -> Any:
|
||||
"""Update a mental model.
|
||||
|
||||
Args:
|
||||
bank_id: The ID of the bank.
|
||||
mental_model_id: The ID of the mental model to update.
|
||||
name: Optional new name.
|
||||
content: Optional new content.
|
||||
tags: Optional new tags list.
|
||||
|
||||
Returns:
|
||||
Update response from the API.
|
||||
"""
|
||||
return self._client.update_mental_model(
|
||||
bank_id=bank_id,
|
||||
mental_model_id=mental_model_id,
|
||||
name=name,
|
||||
content=content,
|
||||
tags=tags,
|
||||
)
|
||||
|
||||
def delete(self, bank_id: str, mental_model_id: str) -> Any:
|
||||
"""Delete a mental model.
|
||||
|
||||
Args:
|
||||
bank_id: The ID of the bank.
|
||||
mental_model_id: The ID of the mental model to delete.
|
||||
|
||||
Returns:
|
||||
Deletion response from the API.
|
||||
"""
|
||||
return self._client.delete_mental_model(bank_id=bank_id, mental_model_id=mental_model_id)
|
||||
|
||||
|
||||
class DirectivesAPI:
|
||||
"""Namespace for directive operations.
|
||||
|
||||
Directives are explicit instructions that guide agent behavior.
|
||||
"""
|
||||
|
||||
def __init__(self, client: Hindsight):
|
||||
self._client = client
|
||||
|
||||
def create(
|
||||
self,
|
||||
bank_id: str,
|
||||
name: str,
|
||||
content: str,
|
||||
tags: list[str] | None = None,
|
||||
) -> Any:
|
||||
"""Create a new directive.
|
||||
|
||||
Args:
|
||||
bank_id: The ID of the bank to add the directive to.
|
||||
name: Name for the directive.
|
||||
content: The directive content/instructions.
|
||||
tags: Optional list of tags for categorization.
|
||||
|
||||
Returns:
|
||||
Creation response from the API.
|
||||
"""
|
||||
return self._client.create_directive(
|
||||
bank_id=bank_id,
|
||||
name=name,
|
||||
content=content,
|
||||
tags=tags,
|
||||
)
|
||||
|
||||
def list(self, bank_id: str, tags: list[str] | None = None) -> Any:
|
||||
"""List all directives for a bank.
|
||||
|
||||
Args:
|
||||
bank_id: The ID of the bank.
|
||||
tags: Optional filter by tags.
|
||||
|
||||
Returns:
|
||||
List of directives.
|
||||
"""
|
||||
return self._client.list_directives(bank_id=bank_id, tags=tags)
|
||||
|
||||
def get(self, bank_id: str, directive_id: str) -> Any:
|
||||
"""Get a specific directive.
|
||||
|
||||
Args:
|
||||
bank_id: The ID of the bank.
|
||||
directive_id: The ID of the directive.
|
||||
|
||||
Returns:
|
||||
The directive details.
|
||||
"""
|
||||
return self._client.get_directive(bank_id=bank_id, directive_id=directive_id)
|
||||
|
||||
def update(
|
||||
self,
|
||||
bank_id: str,
|
||||
directive_id: str,
|
||||
name: str | None = None,
|
||||
content: str | None = None,
|
||||
tags: list[str] | None = None,
|
||||
) -> Any:
|
||||
"""Update a directive.
|
||||
|
||||
Args:
|
||||
bank_id: The ID of the bank.
|
||||
directive_id: The ID of the directive to update.
|
||||
name: Optional new name.
|
||||
content: Optional new content.
|
||||
tags: Optional new tags list.
|
||||
|
||||
Returns:
|
||||
Update response from the API.
|
||||
"""
|
||||
return self._client.update_directive(
|
||||
bank_id=bank_id,
|
||||
directive_id=directive_id,
|
||||
name=name,
|
||||
content=content,
|
||||
tags=tags,
|
||||
)
|
||||
|
||||
def delete(self, bank_id: str, directive_id: str) -> Any:
|
||||
"""Delete a directive.
|
||||
|
||||
Args:
|
||||
bank_id: The ID of the bank.
|
||||
directive_id: The ID of the directive to delete.
|
||||
|
||||
Returns:
|
||||
Deletion response from the API.
|
||||
"""
|
||||
return self._client.delete_directive(bank_id=bank_id, directive_id=directive_id)
|
||||
|
||||
|
||||
class MemoriesAPI:
|
||||
"""Namespace for memory operations.
|
||||
|
||||
Provides methods to query and retrieve stored memories.
|
||||
"""
|
||||
|
||||
def __init__(self, client: Hindsight):
|
||||
self._client = client
|
||||
|
||||
def list(
|
||||
self,
|
||||
bank_id: str,
|
||||
type: str | None = None,
|
||||
search_query: str | None = None,
|
||||
limit: int = 100,
|
||||
offset: int = 0,
|
||||
) -> Any:
|
||||
"""List memories in a bank.
|
||||
|
||||
Args:
|
||||
bank_id: The ID of the bank to query.
|
||||
type: Optional filter by memory type.
|
||||
search_query: Optional search query for filtering.
|
||||
limit: Maximum number of results to return (default: 100).
|
||||
offset: Number of results to skip for pagination (default: 0).
|
||||
|
||||
Returns:
|
||||
List of memories matching the criteria.
|
||||
"""
|
||||
return self._client.list_memories(
|
||||
bank_id=bank_id,
|
||||
type=type,
|
||||
search_query=search_query,
|
||||
limit=limit,
|
||||
offset=offset,
|
||||
)
|
||||
|
||||
|
||||
class HindsightClient(Hindsight):
|
||||
"""
|
||||
Enhanced Hindsight client with organized API namespaces.
|
||||
|
||||
This wrapper extends the auto-generated Hindsight client with organized
|
||||
access to different parts of the API through namespaces.
|
||||
|
||||
Example:
|
||||
```python
|
||||
from hindsight import HindsightClient
|
||||
|
||||
client = HindsightClient(base_url="http://localhost:8888")
|
||||
|
||||
# Core operations (inherited from Hindsight)
|
||||
client.retain(bank_id="test", content="Hello")
|
||||
results = client.recall(bank_id="test", query="Hello")
|
||||
|
||||
# Organized API access through namespaces
|
||||
client.banks.create(bank_id="test", name="Test Bank")
|
||||
models = client.mental_models.list(bank_id="test")
|
||||
directives = client.directives.list(bank_id="test")
|
||||
memories = client.memories.list(bank_id="test")
|
||||
```
|
||||
|
||||
Attributes:
|
||||
banks: Namespace for bank management operations.
|
||||
mental_models: Namespace for mental model operations.
|
||||
directives: Namespace for directive operations.
|
||||
memories: Namespace for memory listing operations.
|
||||
"""
|
||||
|
||||
def __init__(self, *args: Any, **kwargs: Any) -> None:
|
||||
super().__init__(*args, **kwargs)
|
||||
self._banks_namespace: BanksAPI | None = None
|
||||
self._mental_models_namespace: MentalModelsAPI | None = None
|
||||
self._directives_namespace: DirectivesAPI | None = None
|
||||
self._memories_namespace: MemoriesAPI | None = None
|
||||
|
||||
@property
|
||||
def banks(self) -> BanksAPI:
|
||||
"""Access bank management operations.
|
||||
|
||||
Returns:
|
||||
BanksAPI instance for bank operations.
|
||||
"""
|
||||
if self._banks_namespace is None:
|
||||
self._banks_namespace = BanksAPI(self)
|
||||
return self._banks_namespace
|
||||
|
||||
@property
|
||||
def mental_models(self) -> MentalModelsAPI:
|
||||
"""Access mental model operations.
|
||||
|
||||
Returns:
|
||||
MentalModelsAPI instance for mental model operations.
|
||||
"""
|
||||
if self._mental_models_namespace is None:
|
||||
self._mental_models_namespace = MentalModelsAPI(self)
|
||||
return self._mental_models_namespace
|
||||
|
||||
@property
|
||||
def directives(self) -> DirectivesAPI:
|
||||
"""Access directive operations.
|
||||
|
||||
Returns:
|
||||
DirectivesAPI instance for directive operations.
|
||||
"""
|
||||
if self._directives_namespace is None:
|
||||
self._directives_namespace = DirectivesAPI(self)
|
||||
return self._directives_namespace
|
||||
|
||||
@property
|
||||
def memories(self) -> MemoriesAPI:
|
||||
"""Access memory listing operations.
|
||||
|
||||
Returns:
|
||||
MemoriesAPI instance for memory operations.
|
||||
"""
|
||||
if self._memories_namespace is None:
|
||||
self._memories_namespace = MemoriesAPI(self)
|
||||
return self._memories_namespace
|
||||
@@ -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")
|
||||
-52
@@ -1,52 +0,0 @@
|
||||
"""Add consolidation_failed_at column to memory_units for tracking persistent LLM failures.
|
||||
|
||||
When all LLM retries are exhausted on a single-memory batch, the memory is marked
|
||||
with consolidation_failed_at instead of consolidated_at, so it is not silently lost
|
||||
and can be retried later via the API.
|
||||
|
||||
Revision ID: a3b4c5d6e7f8
|
||||
Revises: g7h8i9j0k1l2
|
||||
Create Date: 2026-03-17
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
revision: str = "a3b4c5d6e7f8"
|
||||
down_revision: str | Sequence[str] | None = "g7h8i9j0k1l2"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _get_schema_prefix() -> str:
|
||||
"""Get schema prefix for table names (required for multi-tenant support)."""
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
|
||||
op.execute(
|
||||
f"""
|
||||
ALTER TABLE {schema}memory_units
|
||||
ADD COLUMN IF NOT EXISTS consolidation_failed_at TIMESTAMPTZ DEFAULT NULL
|
||||
"""
|
||||
)
|
||||
|
||||
# Index to efficiently query memories that failed consolidation for a given bank
|
||||
op.execute(
|
||||
f"""
|
||||
CREATE INDEX IF NOT EXISTS idx_memory_units_consolidation_failed
|
||||
ON {schema}memory_units (bank_id, consolidation_failed_at)
|
||||
WHERE consolidation_failed_at IS NOT NULL AND fact_type IN ('experience', 'world')
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema}idx_memory_units_consolidation_failed")
|
||||
op.execute(f"ALTER TABLE {schema}memory_units DROP COLUMN IF EXISTS consolidation_failed_at")
|
||||
@@ -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")
|
||||
-38
@@ -1,38 +0,0 @@
|
||||
"""chunk_fk_cascade_delete
|
||||
|
||||
Revision ID: f6g7h8i9j0k1
|
||||
Revises: e5f6g7h8i9j0
|
||||
Create Date: 2026-03-16 00:00:00.000000
|
||||
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import op
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = "f6g7h8i9j0k1"
|
||||
down_revision: str | Sequence[str] | None = "e5f6g7h8i9j0"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
"""Change memory_units.chunk_id FK from SET NULL to CASCADE.
|
||||
|
||||
When a document is deleted the CASCADE reaches chunks first; with SET NULL
|
||||
the memory_units rows survived with chunk_id = NULL, leaving ghost records.
|
||||
Switching to CASCADE ensures they are removed together with their chunk.
|
||||
"""
|
||||
op.drop_constraint("memory_units_chunk_fkey", "memory_units", type_="foreignkey")
|
||||
op.create_foreign_key(
|
||||
"memory_units_chunk_fkey", "memory_units", "chunks", ["chunk_id"], ["chunk_id"], ondelete="CASCADE"
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
"""Revert to SET NULL behaviour."""
|
||||
op.drop_constraint("memory_units_chunk_fkey", "memory_units", type_="foreignkey")
|
||||
op.create_foreign_key(
|
||||
"memory_units_chunk_fkey", "memory_units", "chunks", ["chunk_id"], ["chunk_id"], ondelete="SET NULL"
|
||||
)
|
||||
-33
@@ -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")
|
||||
-71
@@ -1,71 +0,0 @@
|
||||
"""backsweep_orphan_memory_units
|
||||
|
||||
Two-pass cleanup of memory_units rows that were never removed by earlier bugs:
|
||||
|
||||
Pass 1 — any fact_type, bank gone:
|
||||
memory_units whose bank_id no longer exists in banks. These accumulate when
|
||||
a bank is deleted without a proper cascade (no FK from memory_units to banks
|
||||
exists in the schema).
|
||||
|
||||
Pass 2 — observations only, all sources gone:
|
||||
observation rows whose bank still exists but every source_memory_id points
|
||||
to a deleted memory unit. These were left behind before PR #580 fixed the
|
||||
chunk FK cascade and before delete_document() called
|
||||
_delete_stale_observations_for_memories.
|
||||
|
||||
Revision ID: g7h8i9j0k1l2
|
||||
Revises: f6g7h8i9j0k1
|
||||
Create Date: 2026-03-16
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
revision: str = "g7h8i9j0k1l2"
|
||||
down_revision: str | Sequence[str] | None = "f6g7h8i9j0k1"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _get_schema_prefix() -> str:
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
mu = f"{schema}memory_units"
|
||||
banks = f"{schema}banks"
|
||||
|
||||
# Pass 1: delete all memory_units (any fact_type) whose bank no longer exists.
|
||||
# There is no FK from memory_units to banks, so these never cascade away.
|
||||
op.execute(
|
||||
f"""
|
||||
DELETE FROM {mu}
|
||||
WHERE NOT EXISTS (
|
||||
SELECT 1 FROM {banks} b WHERE b.bank_id = {mu}.bank_id
|
||||
)
|
||||
"""
|
||||
)
|
||||
|
||||
# Pass 2: delete orphaned observations whose bank still exists but every
|
||||
# source_memory_id refers to a now-deleted memory unit (or the array is
|
||||
# empty). Observations with at least one surviving source are left alone.
|
||||
op.execute(
|
||||
f"""
|
||||
DELETE FROM {mu} orphan
|
||||
WHERE orphan.fact_type = 'observation'
|
||||
AND NOT EXISTS (
|
||||
SELECT 1
|
||||
FROM {mu} src
|
||||
WHERE src.id = ANY(orphan.source_memory_ids)
|
||||
AND src.bank_id = orphan.bank_id
|
||||
)
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
# Deleted rows cannot be restored.
|
||||
pass
|
||||
-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")
|
||||
-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,315 +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, replace
|
||||
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)
|
||||
|
||||
# Validate retain_strategies: reject empty string keys
|
||||
if "retain_strategies" in normalized_updates and normalized_updates["retain_strategies"]:
|
||||
empty_keys = [k for k in normalized_updates["retain_strategies"] if not str(k).strip()]
|
||||
if empty_keys:
|
||||
raise ValueError(
|
||||
"Strategy names must not be empty strings. Remove entries with empty names before saving."
|
||||
)
|
||||
|
||||
# Merge with existing config (JSONB || operator)
|
||||
async with self.pool.acquire() as conn:
|
||||
await conn.execute(
|
||||
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")
|
||||
|
||||
|
||||
def apply_strategy(config: HindsightConfig, strategy_name: str) -> HindsightConfig:
|
||||
"""
|
||||
Apply a named retain strategy's overrides on top of a resolved config.
|
||||
|
||||
A strategy is a named set of hierarchical field overrides stored in
|
||||
config.retain_strategies. Any field in _HIERARCHICAL_FIELDS can be
|
||||
overridden, including retain_extraction_mode, retain_chunk_size,
|
||||
entity_labels, entities_allow_free_form, etc.
|
||||
|
||||
Unknown strategy names log a warning and return config unchanged.
|
||||
Unknown or non-hierarchical fields in the strategy are silently ignored.
|
||||
"""
|
||||
strategies = config.retain_strategies or {}
|
||||
if strategy_name not in strategies:
|
||||
logger.warning(f"Unknown retain strategy '{strategy_name}', using resolved config as-is")
|
||||
return config
|
||||
|
||||
overrides = strategies[strategy_name]
|
||||
if not isinstance(overrides, dict):
|
||||
logger.warning(f"Retain strategy '{strategy_name}' is not a dict, skipping")
|
||||
return config
|
||||
|
||||
configurable = HindsightConfig.get_configurable_fields()
|
||||
filtered = {k: v for k, v in overrides.items() if k in configurable}
|
||||
|
||||
if not filtered:
|
||||
return config
|
||||
|
||||
logger.debug(f"Applying retain strategy '{strategy_name}': {list(filtered.keys())}")
|
||||
return replace(config, **filtered)
|
||||
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,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,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)
|
||||
@@ -1,208 +0,0 @@
|
||||
[build-system]
|
||||
requires = ["hatchling"]
|
||||
build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "hindsight-api-slim"
|
||||
version = "0.4.19"
|
||||
description = "Hindsight: Agent Memory That Works Like Human Memory"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.11"
|
||||
dependencies = [
|
||||
"asyncpg>=0.29.0",
|
||||
"python-dotenv>=1.0.0",
|
||||
"openai>=1.0.0",
|
||||
"pydantic>=2.0.0",
|
||||
"rich>=13.0.0",
|
||||
"langchain-text-splitters>=0.3.0",
|
||||
"fastapi[standard]>=0.120.3",
|
||||
"uvicorn>=0.38.0",
|
||||
"wsproto>=1.0.0",
|
||||
"sqlalchemy>=2.0.44",
|
||||
"alembic>=1.17.1",
|
||||
"pgvector>=0.4.1",
|
||||
"greenlet>=3.2.4",
|
||||
"psycopg2-binary>=2.9.11",
|
||||
"tiktoken>=0.12.0",
|
||||
"httpx>=0.27.0",
|
||||
"PyJWT[crypto]>=2.8.0",
|
||||
"fastmcp>=2.14.0", # CVE-2025-66416
|
||||
"python-dateutil>=2.8.0",
|
||||
"opentelemetry-api>=1.20.0",
|
||||
"opentelemetry-sdk>=1.20.0",
|
||||
"opentelemetry-instrumentation-fastapi>=0.41b0",
|
||||
"opentelemetry-exporter-prometheus>=0.41b0",
|
||||
"opentelemetry-exporter-otlp-proto-http>=1.20.0",
|
||||
"opentelemetry-semantic-conventions>=0.41b0",
|
||||
"dateparser>=1.2.2",
|
||||
"google-genai>=1.0.0",
|
||||
"google-auth>=2.0.0",
|
||||
"anthropic>=0.40.0",
|
||||
"typer>=0.9.0",
|
||||
"cohere>=5.0.0",
|
||||
"litellm>=1.0.0",
|
||||
"markitdown[pdf,docx,pptx,xlsx,xls]>=0.1.4", # File to markdown conversion
|
||||
"obstore>=0.4.0", # S3/GCS/Azure object storage client (Rust-backed)
|
||||
"uvloop>=0.22.1",
|
||||
# Transitive dependency security fixes
|
||||
"pyasn1>=0.6.3", # DoS vulnerability fix
|
||||
"urllib3>=2.6.3", # Decompression-bomb safeguards bypass fix
|
||||
"langchain-core>=1.2.11", # Serialization injection + SSRF vulnerability fix
|
||||
"langsmith>=0.6.3", # SSRF via tracing header injection fix
|
||||
"protobuf>=6.33.5", # JSON recursion depth bypass fix
|
||||
"pillow>=12.1.1", # Out-of-bounds write in PSD image loading fix
|
||||
"cryptography>=46.0.5", # Subgroup attack vulnerability fix
|
||||
"filelock>=3.20.1", # TOCTOU race condition fix
|
||||
"authlib>=1.6.9", # Account takeover/JWS header injection vulnerability fix
|
||||
"pyjwt>=2.12.0", # Accepts unknown crit header extensions fix
|
||||
"orjson>=3.11.6", # Unbounded recursion DoS fix
|
||||
"tornado>=6.5.5", # DoS multipart/incomplete cookie validation fix
|
||||
"aiohttp>=3.13.3", # Multiple DoS vulnerabilities
|
||||
"claude-agent-sdk>=0.1.27; sys_platform == 'darwin'",
|
||||
]
|
||||
|
||||
[project.optional-dependencies]
|
||||
local-ml = [
|
||||
# Local ML models for embeddings/reranking
|
||||
"sentence-transformers>=3.3.0",
|
||||
"transformers>=4.53.0", # Security fixes for ReDoS vulnerabilities
|
||||
"torch>=2.6.0", # CVE fix for remote code execution
|
||||
"einops>=0.8.2",
|
||||
"flashrank>=0.2.0",
|
||||
# Apple Silicon local inference
|
||||
"mlx>=0.31.0",
|
||||
"mlx-lm>=0.31.1",
|
||||
"safetensors>=0.6.2",
|
||||
]
|
||||
embedded-db = [
|
||||
"pg0-embedded>=0.11.0",
|
||||
]
|
||||
all = [
|
||||
"hindsight-api-slim[local-ml,embedded-db]",
|
||||
]
|
||||
test = [
|
||||
"pytest>=7.0.0",
|
||||
"pytest-asyncio>=0.21.0",
|
||||
"pytest-timeout>=2.4.0",
|
||||
"pytest-xdist>=3.0.0",
|
||||
"filelock>=3.20.1", # TOCTOU race condition fix
|
||||
"testcontainers>=4.0.0",
|
||||
]
|
||||
|
||||
[project.scripts]
|
||||
hindsight-api = "hindsight_api.main:main"
|
||||
hindsight-worker = "hindsight_api.worker.main:main"
|
||||
hindsight-local-mcp = "hindsight_api.mcp_local:main"
|
||||
hindsight-admin = "hindsight_api.admin.cli:main"
|
||||
|
||||
[tool.hatch.build.targets.wheel]
|
||||
packages = ["hindsight_api"]
|
||||
|
||||
[tool.hatch.build.targets.wheel.sources]
|
||||
"hindsight_api" = "hindsight_api"
|
||||
|
||||
[tool.hatch.build.targets.sdist]
|
||||
include = [
|
||||
"hindsight_api/**/*",
|
||||
]
|
||||
|
||||
[tool.hatch.build]
|
||||
include = [
|
||||
"hindsight_api/**/*.py",
|
||||
"hindsight_api/alembic/**/*",
|
||||
]
|
||||
|
||||
[tool.pytest.ini_options]
|
||||
log_cli = true
|
||||
log_cli_level = "INFO"
|
||||
log_cli_format = "%(asctime)s - %(levelname)s - %(name)s - %(message)s"
|
||||
log_cli_date_format = "%Y-%m-%d %H:%M:%S"
|
||||
addopts = "--timeout 300 -n 8 --dist loadgroup --durations=10 -v"
|
||||
asyncio_mode = "auto"
|
||||
asyncio_default_fixture_loop_scope = "function"
|
||||
log_auto_indent = true
|
||||
filterwarnings = [
|
||||
"ignore:The @wait_container_is_ready decorator is deprecated:DeprecationWarning",
|
||||
"ignore::RuntimeWarning:asyncio",
|
||||
]
|
||||
|
||||
[dependency-groups]
|
||||
dev = [
|
||||
"pytest>=9.0.0",
|
||||
"pytest-asyncio>=1.3.0",
|
||||
"pytest-timeout>=2.4.0",
|
||||
"pytest-xdist>=3.8.0",
|
||||
"python-dotenv>=1.2.1",
|
||||
"filelock>=3.20.1", # TOCTOU race condition fix
|
||||
"ruff>=0.8.0",
|
||||
"ty>=0.0.1",
|
||||
"testcontainers>=4.0.0",
|
||||
]
|
||||
|
||||
[tool.ruff]
|
||||
line-length = 120
|
||||
target-version = "py311"
|
||||
exclude = [
|
||||
"tests/",
|
||||
"**/tests/",
|
||||
]
|
||||
|
||||
[tool.ruff.lint]
|
||||
select = [
|
||||
"E", # pycodestyle errors
|
||||
"W", # pycodestyle warnings
|
||||
"F", # Pyflakes
|
||||
"I", # isort
|
||||
]
|
||||
ignore = [
|
||||
"E501", # line too long (handled by formatter)
|
||||
"E402", # module import not at top of file
|
||||
"F401", # unused import (too noisy during development)
|
||||
"F841", # unused variable (too noisy during development)
|
||||
"F811", # redefined while unused
|
||||
"F821", # undefined name (forward references in type hints)
|
||||
]
|
||||
|
||||
[tool.ruff.lint.isort]
|
||||
known-third-party = ["alembic"]
|
||||
|
||||
[tool.ruff.format]
|
||||
quote-style = "double"
|
||||
indent-style = "space"
|
||||
|
||||
[tool.uv]
|
||||
# Use explicit index for PyTorch to prevent the pytorch index from serving
|
||||
# non-pytorch packages (e.g. markupsafe) with incompatible wheels
|
||||
[[tool.uv.index]]
|
||||
name = "pytorch-cpu"
|
||||
url = "https://download.pytorch.org/whl/cpu"
|
||||
explicit = true
|
||||
|
||||
[tool.uv.sources]
|
||||
# Route torch to the CPU-only PyTorch index; everything else uses PyPI
|
||||
torch = { index = "pytorch-cpu" }
|
||||
|
||||
[tool.ty]
|
||||
# Type checking configuration
|
||||
# ty is an extremely fast Python type checker from Astral (same team as ruff/uv)
|
||||
|
||||
[tool.ty.environment]
|
||||
python-version = "3.11"
|
||||
|
||||
[tool.ty.src]
|
||||
exclude = [
|
||||
"tests/",
|
||||
"hindsight_api/alembic/",
|
||||
]
|
||||
|
||||
[tool.ty.rules]
|
||||
# Disable noisy rules while keeping important ones
|
||||
invalid-argument-type = "ignore" # False positives with **kwargs patterns
|
||||
invalid-return-type = "ignore" # Often intentional in async code
|
||||
invalid-parameter-default = "ignore" # Optional params with None default
|
||||
possibly-missing-attribute = "ignore" # Common with Optional types
|
||||
invalid-raise = "ignore" # False positives with exception tracking
|
||||
call-non-callable = "ignore" # False positives with Optional types
|
||||
invalid-key = "ignore" # Pydantic ConfigDict not understood
|
||||
invalid-method-override = "ignore" # Intentional signature differences
|
||||
unresolved-reference = "ignore" # Forward references not always resolved
|
||||
@@ -1,435 +0,0 @@
|
||||
"""Test async batch retain with smart batching and parent-child operations."""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import uuid
|
||||
|
||||
import pytest
|
||||
|
||||
from hindsight_api.extensions import RequestContext
|
||||
|
||||
|
||||
async def _ensure_bank(pool, bank_id: str) -> None:
|
||||
"""Upsert a minimal bank row so FK on async_operations passes."""
|
||||
await pool.execute(
|
||||
"INSERT INTO banks (bank_id, name) VALUES ($1, $2) ON CONFLICT DO NOTHING",
|
||||
bank_id,
|
||||
bank_id,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_duplicate_document_ids_rejected_async(memory, request_context):
|
||||
"""Test that async retain rejects batches with duplicate document_ids."""
|
||||
bank_id = "test_duplicate_async"
|
||||
contents = [
|
||||
{"content": "First item", "document_id": "doc1"},
|
||||
{"content": "Second item", "document_id": "doc2"},
|
||||
{"content": "Third item", "document_id": "doc1"}, # Duplicate!
|
||||
]
|
||||
|
||||
# Should raise ValueError due to duplicate document_ids
|
||||
with pytest.raises(ValueError, match="duplicate document_ids.*doc1"):
|
||||
await memory.submit_async_retain(
|
||||
bank_id=bank_id,
|
||||
contents=contents,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_duplicate_document_ids_rejected_sync(memory, request_context):
|
||||
"""Test that sync retain also rejects batches with duplicate document_ids."""
|
||||
bank_id = "test_duplicate_sync"
|
||||
contents = [
|
||||
{"content": "First item", "document_id": "doc1"},
|
||||
{"content": "Second item", "document_id": "doc1"}, # Duplicate!
|
||||
]
|
||||
|
||||
# Should raise ValueError due to duplicate document_ids
|
||||
with pytest.raises(ValueError, match="duplicate document_ids.*doc1"):
|
||||
await memory.retain_batch_async(
|
||||
bank_id=bank_id,
|
||||
contents=contents,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_small_async_batch_no_splitting(memory, request_context):
|
||||
"""Test that small async batches create parent with single child (simplified code path)."""
|
||||
bank_id = "test_small_async"
|
||||
contents = [{"content": "Alice works at Google", "document_id": f"doc{i}"} for i in range(5)]
|
||||
|
||||
# Calculate total chars (should be well under threshold)
|
||||
total_chars = sum(len(item["content"]) for item in contents)
|
||||
assert total_chars < 10_000, "Test batch should be small"
|
||||
|
||||
# Submit async retain
|
||||
result = await memory.submit_async_retain(
|
||||
bank_id=bank_id,
|
||||
contents=contents,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Verify we got an operation_id back
|
||||
assert "operation_id" in result
|
||||
assert "items_count" in result
|
||||
assert result["items_count"] == 5
|
||||
|
||||
operation_id = result["operation_id"]
|
||||
|
||||
# Wait for task to complete (SyncTaskBackend executes immediately)
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
# Check operation status
|
||||
status = await memory.get_operation_status(
|
||||
bank_id=bank_id,
|
||||
operation_id=operation_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Should be a parent operation with single child (simplified code path)
|
||||
assert status["status"] == "completed"
|
||||
assert status["operation_type"] == "batch_retain"
|
||||
assert "child_operations" in status
|
||||
assert status["result_metadata"]["num_sub_batches"] == 1 # Single sub-batch
|
||||
assert len(status["child_operations"]) == 1
|
||||
assert status["child_operations"][0]["status"] == "completed"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_large_async_batch_auto_splits(memory, request_context):
|
||||
"""Test that large async batches automatically split into sub-batches with parent operation."""
|
||||
from hindsight_api.engine.memory_engine import count_tokens
|
||||
|
||||
bank_id = "test_large_async"
|
||||
|
||||
# Create a large batch that exceeds the threshold (10k tokens default)
|
||||
# Repeating "A"s gets heavily compressed by tokenizer, use varied content
|
||||
# Use ~22k chars per item = ~5.5k tokens per item, 2 items = ~11k tokens total (exceeds 10k)
|
||||
large_content = "The quick brown fox jumps over the lazy dog. " * 500 # ~22k chars = ~5.5k tokens
|
||||
contents = [{"content": large_content + f" item {i}", "document_id": f"doc{i}"} for i in range(2)]
|
||||
|
||||
# Calculate total tokens (should exceed threshold)
|
||||
total_tokens = sum(count_tokens(item["content"]) for item in contents)
|
||||
assert total_tokens > 10_000, "Test batch should exceed threshold"
|
||||
|
||||
# Submit async retain
|
||||
result = await memory.submit_async_retain(
|
||||
bank_id=bank_id,
|
||||
contents=contents,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Verify we got an operation_id back
|
||||
assert "operation_id" in result
|
||||
assert "items_count" in result
|
||||
assert result["items_count"] == 2
|
||||
|
||||
parent_operation_id = result["operation_id"]
|
||||
|
||||
# Wait for tasks to complete
|
||||
await asyncio.sleep(0.5)
|
||||
|
||||
# Check parent operation status
|
||||
parent_status = await memory.get_operation_status(
|
||||
bank_id=bank_id,
|
||||
operation_id=parent_operation_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Should be a parent operation with children
|
||||
assert parent_status["operation_type"] == "batch_retain"
|
||||
assert "child_operations" in parent_status
|
||||
assert "num_sub_batches" in parent_status["result_metadata"]
|
||||
assert parent_status["result_metadata"]["num_sub_batches"] >= 2 # Should split into at least 2 batches
|
||||
assert parent_status["result_metadata"]["items_count"] == 2
|
||||
|
||||
# Verify child operations
|
||||
child_ops = parent_status["child_operations"]
|
||||
assert len(child_ops) >= 2, "Should have at least 2 child operations"
|
||||
|
||||
# All children should be completed (SyncTaskBackend executes immediately)
|
||||
for child in child_ops:
|
||||
assert child["status"] == "completed"
|
||||
assert child["sub_batch_index"] is not None
|
||||
assert child["items_count"] > 0
|
||||
|
||||
# Parent status should be aggregated as "completed"
|
||||
assert parent_status["status"] == "completed"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_parent_operation_status_aggregation_pending(memory, request_context):
|
||||
"""Test that parent operation shows 'pending' when children are pending."""
|
||||
bank_id = "test_parent_pending"
|
||||
pool = await memory._get_pool()
|
||||
await _ensure_bank(pool, bank_id)
|
||||
|
||||
# Manually create a parent operation
|
||||
parent_id = uuid.uuid4()
|
||||
async with pool.acquire() as conn:
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO async_operations (operation_id, bank_id, operation_type, result_metadata, status)
|
||||
VALUES ($1, $2, $3, $4, $5)
|
||||
""",
|
||||
parent_id,
|
||||
bank_id,
|
||||
"batch_retain",
|
||||
json.dumps({"items_count": 20, "num_sub_batches": 2, "is_parent": True}),
|
||||
"pending",
|
||||
)
|
||||
|
||||
# Create 2 child operations - one completed, one pending
|
||||
child1_id = uuid.uuid4()
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO async_operations (operation_id, bank_id, operation_type, result_metadata, status)
|
||||
VALUES ($1, $2, $3, $4, $5)
|
||||
""",
|
||||
child1_id,
|
||||
bank_id,
|
||||
"retain",
|
||||
json.dumps(
|
||||
{
|
||||
"items_count": 10,
|
||||
"parent_operation_id": str(parent_id),
|
||||
"sub_batch_index": 1,
|
||||
"total_sub_batches": 2,
|
||||
}
|
||||
),
|
||||
"completed",
|
||||
)
|
||||
|
||||
child2_id = uuid.uuid4()
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO async_operations (operation_id, bank_id, operation_type, result_metadata, status)
|
||||
VALUES ($1, $2, $3, $4, $5)
|
||||
""",
|
||||
child2_id,
|
||||
bank_id,
|
||||
"retain",
|
||||
json.dumps(
|
||||
{
|
||||
"items_count": 10,
|
||||
"parent_operation_id": str(parent_id),
|
||||
"sub_batch_index": 2,
|
||||
"total_sub_batches": 2,
|
||||
}
|
||||
),
|
||||
"pending",
|
||||
)
|
||||
|
||||
# Check parent status
|
||||
parent_status = await memory.get_operation_status(
|
||||
bank_id=bank_id,
|
||||
operation_id=str(parent_id),
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Parent should aggregate as "pending" since one child is still pending
|
||||
assert parent_status["status"] == "pending"
|
||||
assert len(parent_status["child_operations"]) == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_parent_operation_status_aggregation_failed(memory, request_context):
|
||||
"""Test that parent operation shows 'failed' when any child fails."""
|
||||
bank_id = "test_parent_failed"
|
||||
pool = await memory._get_pool()
|
||||
await _ensure_bank(pool, bank_id)
|
||||
|
||||
# Manually create a parent operation
|
||||
parent_id = uuid.uuid4()
|
||||
async with pool.acquire() as conn:
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO async_operations (operation_id, bank_id, operation_type, result_metadata, status)
|
||||
VALUES ($1, $2, $3, $4, $5)
|
||||
""",
|
||||
parent_id,
|
||||
bank_id,
|
||||
"batch_retain",
|
||||
json.dumps({"items_count": 20, "num_sub_batches": 2, "is_parent": True}),
|
||||
"pending",
|
||||
)
|
||||
|
||||
# Create 2 child operations - one completed, one failed
|
||||
child1_id = uuid.uuid4()
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO async_operations (operation_id, bank_id, operation_type, result_metadata, status)
|
||||
VALUES ($1, $2, $3, $4, $5)
|
||||
""",
|
||||
child1_id,
|
||||
bank_id,
|
||||
"retain",
|
||||
json.dumps(
|
||||
{
|
||||
"items_count": 10,
|
||||
"parent_operation_id": str(parent_id),
|
||||
"sub_batch_index": 1,
|
||||
"total_sub_batches": 2,
|
||||
}
|
||||
),
|
||||
"completed",
|
||||
)
|
||||
|
||||
child2_id = uuid.uuid4()
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO async_operations (operation_id, bank_id, operation_type, result_metadata, status, error_message)
|
||||
VALUES ($1, $2, $3, $4, $5, $6)
|
||||
""",
|
||||
child2_id,
|
||||
bank_id,
|
||||
"retain",
|
||||
json.dumps(
|
||||
{
|
||||
"items_count": 10,
|
||||
"parent_operation_id": str(parent_id),
|
||||
"sub_batch_index": 2,
|
||||
"total_sub_batches": 2,
|
||||
}
|
||||
),
|
||||
"failed",
|
||||
"Test error",
|
||||
)
|
||||
|
||||
# Check parent status
|
||||
parent_status = await memory.get_operation_status(
|
||||
bank_id=bank_id,
|
||||
operation_id=str(parent_id),
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Parent should aggregate as "failed" since one child failed
|
||||
assert parent_status["status"] == "failed"
|
||||
assert len(parent_status["child_operations"]) == 2
|
||||
|
||||
# Verify child with error is included
|
||||
failed_child = [c for c in parent_status["child_operations"] if c["status"] == "failed"][0]
|
||||
assert failed_child["error_message"] == "Test error"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_parent_operation_status_aggregation_completed(memory, request_context):
|
||||
"""Test that parent operation shows 'completed' when all children are completed."""
|
||||
bank_id = "test_parent_completed"
|
||||
pool = await memory._get_pool()
|
||||
await _ensure_bank(pool, bank_id)
|
||||
|
||||
# Manually create a parent operation
|
||||
parent_id = uuid.uuid4()
|
||||
async with pool.acquire() as conn:
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO async_operations (operation_id, bank_id, operation_type, result_metadata, status)
|
||||
VALUES ($1, $2, $3, $4, $5)
|
||||
""",
|
||||
parent_id,
|
||||
bank_id,
|
||||
"batch_retain",
|
||||
json.dumps({"items_count": 20, "num_sub_batches": 2, "is_parent": True}),
|
||||
"pending",
|
||||
)
|
||||
|
||||
# Create 2 child operations - both completed
|
||||
child1_id = uuid.uuid4()
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO async_operations (operation_id, bank_id, operation_type, result_metadata, status)
|
||||
VALUES ($1, $2, $3, $4, $5)
|
||||
""",
|
||||
child1_id,
|
||||
bank_id,
|
||||
"retain",
|
||||
json.dumps(
|
||||
{
|
||||
"items_count": 10,
|
||||
"parent_operation_id": str(parent_id),
|
||||
"sub_batch_index": 1,
|
||||
"total_sub_batches": 2,
|
||||
}
|
||||
),
|
||||
"completed",
|
||||
)
|
||||
|
||||
child2_id = uuid.uuid4()
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO async_operations (operation_id, bank_id, operation_type, result_metadata, status)
|
||||
VALUES ($1, $2, $3, $4, $5)
|
||||
""",
|
||||
child2_id,
|
||||
bank_id,
|
||||
"retain",
|
||||
json.dumps(
|
||||
{
|
||||
"items_count": 10,
|
||||
"parent_operation_id": str(parent_id),
|
||||
"sub_batch_index": 2,
|
||||
"total_sub_batches": 2,
|
||||
}
|
||||
),
|
||||
"completed",
|
||||
)
|
||||
|
||||
# Check parent status
|
||||
parent_status = await memory.get_operation_status(
|
||||
bank_id=bank_id,
|
||||
operation_id=str(parent_id),
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Parent should aggregate as "completed" since all children are completed
|
||||
assert parent_status["status"] == "completed"
|
||||
assert len(parent_status["child_operations"]) == 2
|
||||
assert all(c["status"] == "completed" for c in parent_status["child_operations"])
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_config_retain_batch_tokens_respected(memory, request_context):
|
||||
"""Test that the retain_batch_tokens config setting is respected."""
|
||||
from hindsight_api.config import get_config
|
||||
from hindsight_api.engine.memory_engine import count_tokens
|
||||
|
||||
bank_id = "test_config_batch_tokens"
|
||||
config = get_config()
|
||||
|
||||
# Check that config has the retain_batch_tokens setting
|
||||
assert hasattr(config, "retain_batch_tokens")
|
||||
assert config.retain_batch_tokens > 0
|
||||
|
||||
# Create a batch that's just under the threshold
|
||||
# Use content that produces roughly half the token limit per item
|
||||
content_size = config.retain_batch_tokens * 2 # chars (rough estimate: 1 token ~= 4 chars)
|
||||
contents = [{"content": "A" * content_size, "document_id": f"doc{i}"} for i in range(2)]
|
||||
|
||||
total_tokens = sum(count_tokens(item["content"]) for item in contents)
|
||||
# Should be equal to threshold (boundary case, no splitting since we use > not >=)
|
||||
assert total_tokens <= config.retain_batch_tokens
|
||||
|
||||
# Submit - should NOT split
|
||||
result = await memory.submit_async_retain(
|
||||
bank_id=bank_id,
|
||||
contents=contents,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Wait for completion
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
# Check status - should be a parent with single child (even for small batches)
|
||||
status = await memory.get_operation_status(
|
||||
bank_id=bank_id,
|
||||
operation_id=result["operation_id"],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Even small batches use parent-child pattern now (simpler code path)
|
||||
assert "child_operations" in status
|
||||
assert status["result_metadata"]["num_sub_batches"] == 1
|
||||
@@ -1,95 +0,0 @@
|
||||
"""Unit tests for async retain tag propagation."""
|
||||
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from hindsight_api.engine.memory_engine import MemoryEngine
|
||||
from hindsight_api.models import RequestContext
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_submit_async_retain_includes_document_tags_in_task_payload():
|
||||
"""submit_async_retain should include document_tags in queued task payload."""
|
||||
engine = MemoryEngine.__new__(MemoryEngine)
|
||||
engine._initialized = True
|
||||
engine._authenticate_tenant = AsyncMock()
|
||||
engine._operation_validator = None
|
||||
engine._submit_async_operation = AsyncMock(return_value={"operation_id": "op-1"})
|
||||
|
||||
# Mock the pool and connection for parent operation creation
|
||||
mock_conn = AsyncMock()
|
||||
mock_conn.execute = AsyncMock()
|
||||
mock_conn.transaction = MagicMock()
|
||||
mock_conn.transaction.return_value.__aenter__ = AsyncMock()
|
||||
mock_conn.transaction.return_value.__aexit__ = AsyncMock()
|
||||
|
||||
mock_pool = AsyncMock()
|
||||
mock_pool.acquire = AsyncMock(return_value=mock_conn)
|
||||
mock_pool.release = AsyncMock()
|
||||
|
||||
engine._get_pool = AsyncMock(return_value=mock_pool)
|
||||
|
||||
request_context = RequestContext(tenant_id="tenant-a", api_key_id="key-a")
|
||||
contents = [{"content": "Async retain payload test."}]
|
||||
document_tags = ["scope:tools", "user:alice"]
|
||||
|
||||
with patch("hindsight_api.engine.memory_engine.bank_utils.get_bank_profile", new_callable=AsyncMock):
|
||||
result = await MemoryEngine.submit_async_retain(
|
||||
engine,
|
||||
bank_id="bank-1",
|
||||
contents=contents,
|
||||
document_tags=document_tags,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Check result structure
|
||||
assert "operation_id" in result
|
||||
assert "items_count" in result
|
||||
assert result["items_count"] == 1
|
||||
|
||||
# Verify authentication was called
|
||||
engine._authenticate_tenant.assert_awaited_once_with(request_context)
|
||||
|
||||
# Verify child operation was submitted
|
||||
engine._submit_async_operation.assert_awaited_once()
|
||||
|
||||
# Verify child operation payload contains document_tags
|
||||
kwargs = engine._submit_async_operation.await_args.kwargs
|
||||
assert kwargs["bank_id"] == "bank-1"
|
||||
assert kwargs["operation_type"] == "retain"
|
||||
assert kwargs["task_type"] == "batch_retain"
|
||||
assert kwargs["task_payload"]["contents"] == contents
|
||||
assert kwargs["task_payload"]["document_tags"] == document_tags
|
||||
assert kwargs["task_payload"]["_tenant_id"] == "tenant-a"
|
||||
assert kwargs["task_payload"]["_api_key_id"] == "key-a"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_batch_retain_forwards_document_tags_to_retain_batch_async():
|
||||
"""Worker handler should forward document_tags from task payload."""
|
||||
engine = MemoryEngine.__new__(MemoryEngine)
|
||||
engine._initialized = True
|
||||
engine.retain_batch_async = AsyncMock(return_value={"items_count": 1})
|
||||
|
||||
task_dict = {
|
||||
"bank_id": "bank-1",
|
||||
"contents": [{"content": "Forward tags test."}],
|
||||
"document_tags": ["scope:client"],
|
||||
"_tenant_id": "tenant-a",
|
||||
"_api_key_id": "key-a",
|
||||
}
|
||||
|
||||
await MemoryEngine._handle_batch_retain(engine, task_dict)
|
||||
|
||||
engine.retain_batch_async.assert_awaited_once()
|
||||
kwargs = engine.retain_batch_async.await_args.kwargs
|
||||
assert kwargs["bank_id"] == "bank-1"
|
||||
assert kwargs["contents"] == task_dict["contents"]
|
||||
assert kwargs["document_tags"] == ["scope:client"]
|
||||
|
||||
request_context = kwargs["request_context"]
|
||||
assert request_context.internal is True
|
||||
assert request_context.user_initiated is True
|
||||
assert request_context.tenant_id == "tenant-a"
|
||||
assert request_context.api_key_id == "key-a"
|
||||
@@ -1,189 +0,0 @@
|
||||
"""
|
||||
Integration test for API base path support.
|
||||
|
||||
Tests that the API works correctly when deployed with a base path (e.g., /hindsight)
|
||||
for reverse proxy deployments.
|
||||
"""
|
||||
import os
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
import httpx
|
||||
from hindsight_api.api import create_app
|
||||
from hindsight_api.config import clear_config_cache
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def api_client_with_base_path(memory):
|
||||
"""Create an async test client for the FastAPI app with a base path."""
|
||||
# Set base path in environment
|
||||
base_path = "/hindsight"
|
||||
os.environ["HINDSIGHT_API_BASE_PATH"] = base_path
|
||||
|
||||
# Clear config cache to force reload with new base_path
|
||||
clear_config_cache()
|
||||
|
||||
# Memory is already initialized by the conftest fixture (with migrations)
|
||||
app = create_app(memory, initialize_memory=False)
|
||||
|
||||
# Use base_url with base path
|
||||
transport = httpx.ASGITransport(app=app)
|
||||
async with httpx.AsyncClient(
|
||||
transport=transport,
|
||||
base_url=f"http://test{base_path}"
|
||||
) as client:
|
||||
yield client
|
||||
|
||||
# Cleanup: unset base path
|
||||
os.environ.pop("HINDSIGHT_API_BASE_PATH", None)
|
||||
clear_config_cache()
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def api_client_without_base_path(memory):
|
||||
"""Create an async test client for the FastAPI app without a base path (root)."""
|
||||
# Ensure no base path is set
|
||||
os.environ.pop("HINDSIGHT_API_BASE_PATH", None)
|
||||
clear_config_cache()
|
||||
|
||||
app = create_app(memory, initialize_memory=False)
|
||||
transport = httpx.ASGITransport(app=app)
|
||||
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
|
||||
yield client
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_base_path_health_endpoint(api_client_with_base_path):
|
||||
"""Test that health endpoint works with base path."""
|
||||
# With base path set to /hindsight, health should be at /hindsight/health
|
||||
# But since our client base_url is already http://test/hindsight, we request /health
|
||||
response = await api_client_with_base_path.get("/health")
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert "status" in data
|
||||
assert data["status"] in ["ok", "healthy"] # Accept both formats
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_base_path_banks_endpoint(api_client_with_base_path):
|
||||
"""Test that banks endpoint works with base path."""
|
||||
response = await api_client_with_base_path.get("/v1/default/banks")
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert "banks" in data
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_base_path_openapi_schema(api_client_with_base_path):
|
||||
"""Test that OpenAPI schema includes correct base path in servers."""
|
||||
response = await api_client_with_base_path.get("/openapi.json")
|
||||
assert response.status_code == 200
|
||||
openapi_schema = response.json()
|
||||
|
||||
# Check that servers array includes base path
|
||||
assert "servers" in openapi_schema
|
||||
servers = openapi_schema["servers"]
|
||||
assert len(servers) > 0
|
||||
# FastAPI should set server URL to the root_path
|
||||
assert servers[0]["url"] == "/hindsight"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_base_path_docs_redirect(api_client_with_base_path):
|
||||
"""Test that /docs redirects correctly with base path."""
|
||||
# FastAPI docs endpoint should work
|
||||
response = await api_client_with_base_path.get("/docs", follow_redirects=False)
|
||||
# Should either return 200 (direct) or 307 (redirect to trailing slash)
|
||||
assert response.status_code in [200, 307]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_base_path_metrics(api_client_with_base_path):
|
||||
"""Test that metrics endpoint works with base path."""
|
||||
response = await api_client_with_base_path.get("/metrics")
|
||||
assert response.status_code == 200
|
||||
# Metrics should be in Prometheus format
|
||||
assert "# HELP" in response.text or "# TYPE" in response.text
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_base_path_full_workflow(api_client_with_base_path):
|
||||
"""
|
||||
Test a full retain/recall workflow with base path.
|
||||
|
||||
This ensures that all memory operations work correctly when the API
|
||||
is deployed with a base path.
|
||||
"""
|
||||
bank_id = "test_base_path_bank"
|
||||
|
||||
# 1. Create/get bank
|
||||
response = await api_client_with_base_path.get(f"/v1/default/banks/{bank_id}/profile")
|
||||
assert response.status_code == 200
|
||||
|
||||
# 2. Store a memory
|
||||
response = await api_client_with_base_path.post(
|
||||
f"/v1/default/banks/{bank_id}/memories",
|
||||
json={
|
||||
"items": [
|
||||
{
|
||||
"content": "The API supports base path deployment for reverse proxy use cases.",
|
||||
"context": "testing base path feature"
|
||||
}
|
||||
]
|
||||
}
|
||||
)
|
||||
assert response.status_code == 200
|
||||
result = response.json()
|
||||
assert result["success"] is True
|
||||
|
||||
# 3. Recall the memory
|
||||
response = await api_client_with_base_path.post(
|
||||
f"/v1/default/banks/{bank_id}/memories/recall",
|
||||
json={
|
||||
"query": "base path support"
|
||||
}
|
||||
)
|
||||
assert response.status_code == 200
|
||||
recall_result = response.json()
|
||||
# API returns "results" not "memories"
|
||||
assert "results" in recall_result
|
||||
assert len(recall_result["results"]) > 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_without_base_path_still_works(api_client_without_base_path):
|
||||
"""
|
||||
Regression test: ensure default behavior (no base path) still works.
|
||||
|
||||
This test verifies that when HINDSIGHT_API_BASE_PATH is not set,
|
||||
the API works at the root path as before.
|
||||
"""
|
||||
# Health check at root
|
||||
response = await api_client_without_base_path.get("/health")
|
||||
assert response.status_code == 200
|
||||
|
||||
# Banks endpoint at root
|
||||
response = await api_client_without_base_path.get("/v1/default/banks")
|
||||
assert response.status_code == 200
|
||||
|
||||
# OpenAPI schema should have empty or "/" server path
|
||||
response = await api_client_without_base_path.get("/openapi.json")
|
||||
assert response.status_code == 200
|
||||
openapi_schema = response.json()
|
||||
servers = openapi_schema.get("servers", [])
|
||||
if servers:
|
||||
# Server URL should be empty string (root) or "/"
|
||||
assert servers[0]["url"] in ["", "/"]
|
||||
|
||||
|
||||
@pytest.mark.skip(reason="MCP endpoint routing with base path needs investigation")
|
||||
@pytest.mark.asyncio
|
||||
async def test_base_path_mcp_endpoint(api_client_with_base_path):
|
||||
"""Test that MCP endpoint is accessible with base path."""
|
||||
bank_id = "test_mcp_bank"
|
||||
|
||||
# MCP endpoint should be mounted at /mcp/{bank_id}/
|
||||
# The MCP server uses a different protocol, so just check the root exists
|
||||
response = await api_client_with_base_path.get(f"/mcp/{bank_id}/")
|
||||
# MCP may return various status codes, but should not be 404 (not found)
|
||||
# Accept 405 (method not allowed), 400 (bad request), etc.
|
||||
assert response.status_code != 404, "MCP endpoint should exist"
|
||||
@@ -1,507 +0,0 @@
|
||||
"""
|
||||
Test OpenAI Batch API integration for retain fact extraction.
|
||||
|
||||
Tests cover:
|
||||
- Normal batch API flow (submit, poll, complete)
|
||||
- Crash recovery (resume from existing batch_id)
|
||||
- Provider fallback (when batch API not supported)
|
||||
- Worker recovery on restart
|
||||
"""
|
||||
import pytest
|
||||
import asyncio
|
||||
import logging
|
||||
import json
|
||||
import uuid
|
||||
from datetime import datetime, timezone
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
from hindsight_api import RequestContext
|
||||
from hindsight_api.engine.retain.fact_extraction import (
|
||||
extract_facts_from_contents_batch_api,
|
||||
extract_facts_from_contents,
|
||||
RetainContent,
|
||||
)
|
||||
from hindsight_api.config import HindsightConfig
|
||||
from hindsight_api.engine.llm_wrapper import create_llm_provider
|
||||
from hindsight_api.worker.poller import WorkerPoller
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_llm_config():
|
||||
"""Create a mock LLM config with batch API support."""
|
||||
mock = MagicMock()
|
||||
mock.provider = "openai"
|
||||
mock.model = "gpt-4o-mini"
|
||||
mock._provider_impl = AsyncMock()
|
||||
return mock
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def test_contents():
|
||||
"""Create test content for fact extraction."""
|
||||
return [
|
||||
RetainContent(
|
||||
content="Alice is a senior software engineer at TechCorp. She specializes in distributed systems.",
|
||||
event_date=datetime(2024, 1, 15, tzinfo=timezone.utc),
|
||||
context="team overview",
|
||||
),
|
||||
RetainContent(
|
||||
content="Bob joined the team last month as a junior developer. He is learning React.",
|
||||
event_date=datetime(2024, 1, 15, tzinfo=timezone.utc),
|
||||
context="team overview",
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def hindsight_config():
|
||||
"""Create test config with batch API enabled."""
|
||||
config = HindsightConfig.from_env()
|
||||
config.retain_batch_enabled = True
|
||||
config.retain_batch_poll_interval_seconds = 1 # Fast polling for tests
|
||||
config.retain_chunk_size = 4000
|
||||
config.retain_extraction_mode = "concise"
|
||||
config.retain_extract_causal_links = False
|
||||
return config
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_batch_api_normal_flow(mock_llm_config, test_contents, hindsight_config, memory, request_context):
|
||||
"""Test normal batch API flow: submit, poll, complete."""
|
||||
bank_id = f"test_batch_{datetime.now(timezone.utc).timestamp()}"
|
||||
|
||||
try:
|
||||
# Mock batch API responses
|
||||
batch_id = "batch_test123"
|
||||
|
||||
# Mock supports_batch_api
|
||||
mock_llm_config._provider_impl.supports_batch_api = AsyncMock(return_value=True)
|
||||
|
||||
# Mock submit_batch - returns batch metadata
|
||||
mock_llm_config._provider_impl.submit_batch = AsyncMock(
|
||||
return_value={
|
||||
"batch_id": batch_id,
|
||||
"status": "validating",
|
||||
"request_counts": {"total": 2, "completed": 0, "failed": 0},
|
||||
}
|
||||
)
|
||||
|
||||
# Mock get_batch_status - simulate polling sequence
|
||||
status_sequence = [
|
||||
{"status": "in_progress", "request_counts": {"total": 2, "completed": 1, "failed": 0}},
|
||||
{"status": "completed", "request_counts": {"total": 2, "completed": 2, "failed": 0}},
|
||||
]
|
||||
mock_llm_config._provider_impl.get_batch_status = AsyncMock(side_effect=status_sequence)
|
||||
|
||||
# Mock retrieve_batch_results - returns fact extraction results
|
||||
mock_results = [
|
||||
{
|
||||
"custom_id": "chunk_0",
|
||||
"response": {
|
||||
"body": {
|
||||
"choices": [
|
||||
{
|
||||
"message": {
|
||||
"content": json.dumps({
|
||||
"facts": [
|
||||
{
|
||||
"what": "Alice is a senior software engineer at TechCorp",
|
||||
"when": "present",
|
||||
"where": "TechCorp",
|
||||
"who": "Alice",
|
||||
"why": "Professional background information",
|
||||
"fact_type": "world",
|
||||
"fact_kind": "conversation",
|
||||
}
|
||||
]
|
||||
})
|
||||
}
|
||||
}
|
||||
],
|
||||
"usage": {"prompt_tokens": 100, "completion_tokens": 50, "total_tokens": 150},
|
||||
}
|
||||
},
|
||||
},
|
||||
{
|
||||
"custom_id": "chunk_1",
|
||||
"response": {
|
||||
"body": {
|
||||
"choices": [
|
||||
{
|
||||
"message": {
|
||||
"content": json.dumps({
|
||||
"facts": [
|
||||
{
|
||||
"what": "Bob joined the team last month as a junior developer",
|
||||
"when": "last month",
|
||||
"where": "team",
|
||||
"who": "Bob",
|
||||
"why": "New team member information",
|
||||
"fact_type": "world",
|
||||
"fact_kind": "conversation",
|
||||
}
|
||||
]
|
||||
})
|
||||
}
|
||||
}
|
||||
],
|
||||
"usage": {"prompt_tokens": 100, "completion_tokens": 50, "total_tokens": 150},
|
||||
}
|
||||
},
|
||||
},
|
||||
]
|
||||
mock_llm_config._provider_impl.retrieve_batch_results = AsyncMock(return_value=mock_results)
|
||||
|
||||
# Call batch API extraction
|
||||
facts, chunks, usage = await extract_facts_from_contents_batch_api(
|
||||
contents=test_contents,
|
||||
llm_config=mock_llm_config,
|
||||
agent_name="test_agent",
|
||||
config=hindsight_config,
|
||||
pool=None, # No DB pool for this test
|
||||
operation_id=None,
|
||||
schema=None,
|
||||
)
|
||||
|
||||
# Verify results
|
||||
assert len(facts) == 2, "Should extract 2 facts (one per chunk)"
|
||||
# Facts are ExtractedFact objects with .fact_text field
|
||||
assert "Alice" in facts[0].fact_text and "senior software engineer" in facts[0].fact_text
|
||||
assert "Bob" in facts[1].fact_text and "junior developer" in facts[1].fact_text
|
||||
|
||||
# Verify chunks metadata
|
||||
assert len(chunks) == 2, "Should have 2 chunks metadata"
|
||||
assert chunks[0].fact_count == 1
|
||||
assert chunks[1].fact_count == 1
|
||||
|
||||
# Verify token usage
|
||||
assert usage.input_tokens == 200 # 100 per chunk
|
||||
assert usage.output_tokens == 100 # 50 per chunk
|
||||
assert usage.total_tokens == 300
|
||||
|
||||
# Verify API calls
|
||||
mock_llm_config._provider_impl.submit_batch.assert_called_once()
|
||||
assert mock_llm_config._provider_impl.get_batch_status.call_count == 2
|
||||
mock_llm_config._provider_impl.retrieve_batch_results.assert_called_once_with(batch_id)
|
||||
|
||||
logger.info("✅ Normal batch API flow test passed")
|
||||
|
||||
finally:
|
||||
# Cleanup
|
||||
try:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_batch_api_crash_recovery(mock_llm_config, test_contents, hindsight_config, memory, request_context):
|
||||
"""Test crash recovery: resume polling from existing batch_id."""
|
||||
bank_id = f"test_crash_{datetime.now(timezone.utc).timestamp()}"
|
||||
operation_id = str(uuid.uuid4()) # Must be UUID for async_operations table
|
||||
|
||||
try:
|
||||
# Ensure bank exists
|
||||
await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
|
||||
# Setup: Store batch_id in async_operations table (simulates partial execution)
|
||||
batch_id = "batch_recovered_456"
|
||||
pool = memory._pool
|
||||
schema = request_context.tenant_id
|
||||
|
||||
from hindsight_api.engine.task_backend import fq_table
|
||||
table = fq_table("async_operations", schema)
|
||||
|
||||
# Create operation with batch_id already stored
|
||||
await pool.execute(
|
||||
f"""
|
||||
INSERT INTO {table} (operation_id, operation_type, bank_id, status, result_metadata)
|
||||
VALUES ($1, 'retain', $2, 'processing', $3::jsonb)
|
||||
""",
|
||||
operation_id,
|
||||
bank_id,
|
||||
json.dumps({
|
||||
"batch_id": batch_id,
|
||||
"batch_provider": "openai",
|
||||
"chunk_count": 2,
|
||||
}),
|
||||
)
|
||||
|
||||
# Mock batch API responses for resume scenario
|
||||
mock_llm_config._provider_impl.supports_batch_api = AsyncMock(return_value=True)
|
||||
|
||||
# Mock get_batch_status - batch already in progress
|
||||
mock_llm_config._provider_impl.get_batch_status = AsyncMock(
|
||||
return_value={
|
||||
"status": "completed",
|
||||
"request_counts": {"total": 2, "completed": 2, "failed": 0},
|
||||
}
|
||||
)
|
||||
|
||||
# Mock retrieve_batch_results
|
||||
mock_results = [
|
||||
{
|
||||
"custom_id": "chunk_0",
|
||||
"response": {
|
||||
"body": {
|
||||
"choices": [
|
||||
{
|
||||
"message": {
|
||||
"content": json.dumps({
|
||||
"facts": [
|
||||
{
|
||||
"what": "Alice is a senior software engineer",
|
||||
"when": "present",
|
||||
"where": "TechCorp",
|
||||
"who": "Alice",
|
||||
"why": "Background",
|
||||
"fact_type": "world",
|
||||
"fact_kind": "conversation",
|
||||
}
|
||||
]
|
||||
})
|
||||
}
|
||||
}
|
||||
],
|
||||
"usage": {"prompt_tokens": 100, "completion_tokens": 50, "total_tokens": 150},
|
||||
}
|
||||
},
|
||||
},
|
||||
{
|
||||
"custom_id": "chunk_1",
|
||||
"response": {
|
||||
"body": {
|
||||
"choices": [
|
||||
{
|
||||
"message": {
|
||||
"content": json.dumps({
|
||||
"facts": [
|
||||
{
|
||||
"what": "Bob is a junior developer",
|
||||
"when": "last month",
|
||||
"where": "team",
|
||||
"who": "Bob",
|
||||
"why": "New member",
|
||||
"fact_type": "world",
|
||||
"fact_kind": "conversation",
|
||||
}
|
||||
]
|
||||
})
|
||||
}
|
||||
}
|
||||
],
|
||||
"usage": {"prompt_tokens": 100, "completion_tokens": 50, "total_tokens": 150},
|
||||
}
|
||||
},
|
||||
},
|
||||
]
|
||||
mock_llm_config._provider_impl.retrieve_batch_results = AsyncMock(return_value=mock_results)
|
||||
|
||||
# Call batch API extraction with operation_id (crash recovery scenario)
|
||||
facts, chunks, usage = await extract_facts_from_contents_batch_api(
|
||||
contents=test_contents,
|
||||
llm_config=mock_llm_config,
|
||||
agent_name="test_agent",
|
||||
config=hindsight_config,
|
||||
pool=pool,
|
||||
operation_id=operation_id, # Provides crash recovery context
|
||||
schema=schema,
|
||||
)
|
||||
|
||||
# Verify results
|
||||
assert len(facts) == 2, "Should extract 2 facts after recovery"
|
||||
|
||||
# CRITICAL: Verify submit_batch was NOT called (because batch_id already exists)
|
||||
mock_llm_config._provider_impl.submit_batch.assert_not_called()
|
||||
|
||||
# Verify get_batch_status WAS called (polling resumed)
|
||||
mock_llm_config._provider_impl.get_batch_status.assert_called()
|
||||
|
||||
# Verify retrieve_batch_results was called with the recovered batch_id
|
||||
mock_llm_config._provider_impl.retrieve_batch_results.assert_called_once_with(batch_id)
|
||||
|
||||
logger.info("✅ Crash recovery test passed - resumed polling without re-submission")
|
||||
|
||||
finally:
|
||||
# Cleanup
|
||||
try:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_batch_api_fallback_unsupported_provider(mock_llm_config, test_contents, hindsight_config):
|
||||
"""Test fallback to sync mode when provider doesn't support batch API."""
|
||||
|
||||
# Mock provider that doesn't support batch API
|
||||
mock_llm_config._provider_impl.supports_batch_api = AsyncMock(return_value=False)
|
||||
mock_llm_config.provider = "groq" # Example of provider
|
||||
|
||||
# Patch the sync mode function to verify it's called
|
||||
with patch(
|
||||
"hindsight_api.engine.retain.fact_extraction.extract_facts_from_contents"
|
||||
) as mock_sync_extract:
|
||||
mock_sync_extract.return_value = ([], [], MagicMock())
|
||||
|
||||
# Call batch API extraction (should fallback to sync)
|
||||
await extract_facts_from_contents_batch_api(
|
||||
contents=test_contents,
|
||||
llm_config=mock_llm_config,
|
||||
agent_name="test_agent",
|
||||
config=hindsight_config,
|
||||
pool=None,
|
||||
operation_id=None,
|
||||
schema=None,
|
||||
)
|
||||
|
||||
# Verify fallback occurred
|
||||
mock_sync_extract.assert_called_once()
|
||||
|
||||
# Verify batch API methods were NOT called
|
||||
mock_llm_config._provider_impl.submit_batch.assert_not_called()
|
||||
|
||||
logger.info("✅ Fallback to sync mode test passed")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_worker_batch_recovery(memory, request_context):
|
||||
"""Test that WorkerPoller._recover_batch_operations finds and resets orphaned batches."""
|
||||
bank_id = f"test_worker_recovery_{datetime.now(timezone.utc).timestamp()}"
|
||||
operation_id = str(uuid.uuid4()) # Must be UUID for async_operations table
|
||||
|
||||
try:
|
||||
# Ensure bank exists
|
||||
await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
|
||||
pool = memory._pool
|
||||
schema = request_context.tenant_id
|
||||
|
||||
from hindsight_api.engine.task_backend import fq_table
|
||||
table = fq_table("async_operations", schema)
|
||||
|
||||
# Create orphaned batch operation (simulates worker crash during polling)
|
||||
batch_id = "batch_orphaned_999"
|
||||
task_payload = {
|
||||
"operation_type": "retain",
|
||||
"bank_id": bank_id,
|
||||
"contents": [{"content": "test", "event_date": "2024-01-15T00:00:00Z"}],
|
||||
}
|
||||
|
||||
await pool.execute(
|
||||
f"""
|
||||
INSERT INTO {table} (operation_id, operation_type, bank_id, status, worker_id, result_metadata, task_payload)
|
||||
VALUES ($1, 'retain', $2, 'processing', 'worker_crashed', $3::jsonb, $4::jsonb)
|
||||
""",
|
||||
operation_id,
|
||||
bank_id,
|
||||
json.dumps({
|
||||
"batch_id": batch_id,
|
||||
"batch_provider": "openai",
|
||||
"chunk_count": 1,
|
||||
}),
|
||||
json.dumps(task_payload),
|
||||
)
|
||||
|
||||
# Create WorkerPoller
|
||||
from hindsight_api.extensions.builtin.tenant import DefaultTenantExtension
|
||||
tenant_extension = DefaultTenantExtension(config={"schema": schema} if schema else {})
|
||||
|
||||
poller = WorkerPoller(
|
||||
pool=pool,
|
||||
worker_id="test_worker_recovery",
|
||||
executor=memory,
|
||||
poll_interval_ms=100,
|
||||
schema=schema,
|
||||
tenant_extension=tenant_extension,
|
||||
max_slots=5,
|
||||
consolidation_max_slots=2,
|
||||
)
|
||||
|
||||
# Run recovery
|
||||
recovered_count = await poller._recover_batch_operations(schema)
|
||||
|
||||
# Verify recovery
|
||||
assert recovered_count == 1, "Should recover 1 batch operation"
|
||||
|
||||
# Verify operation was reset to pending
|
||||
row = await pool.fetchrow(
|
||||
f"SELECT status, worker_id FROM {table} WHERE operation_id = $1",
|
||||
operation_id,
|
||||
)
|
||||
|
||||
assert row["status"] == "pending", "Operation should be reset to pending"
|
||||
assert row["worker_id"] is None, "Worker ID should be cleared"
|
||||
|
||||
logger.info("✅ Worker batch recovery test passed")
|
||||
|
||||
finally:
|
||||
# Cleanup
|
||||
try:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_batch_api_via_extract_facts_from_contents(
|
||||
mock_llm_config, test_contents, hindsight_config, memory, request_context
|
||||
):
|
||||
"""Test that extract_facts_from_contents routes to batch API when enabled."""
|
||||
bank_id = f"test_routing_{datetime.now(timezone.utc).timestamp()}"
|
||||
|
||||
try:
|
||||
# Enable batch API in config
|
||||
hindsight_config.retain_batch_enabled = True
|
||||
|
||||
# Mock batch API support
|
||||
mock_llm_config._provider_impl.supports_batch_api = AsyncMock(return_value=True)
|
||||
mock_llm_config._provider_impl.submit_batch = AsyncMock(
|
||||
return_value={"batch_id": "batch_123", "status": "validating", "request_counts": {}}
|
||||
)
|
||||
mock_llm_config._provider_impl.get_batch_status = AsyncMock(
|
||||
return_value={"status": "completed", "request_counts": {"total": 1, "completed": 1, "failed": 0}}
|
||||
)
|
||||
mock_llm_config._provider_impl.retrieve_batch_results = AsyncMock(
|
||||
return_value=[
|
||||
{
|
||||
"custom_id": "chunk_0",
|
||||
"response": {
|
||||
"body": {
|
||||
"choices": [
|
||||
{
|
||||
"message": {
|
||||
"content": json.dumps({"facts": []})
|
||||
}
|
||||
}
|
||||
],
|
||||
"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
|
||||
}
|
||||
},
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
# Call main extract_facts_from_contents (should route to batch API)
|
||||
facts, chunks, usage = await extract_facts_from_contents(
|
||||
contents=test_contents,
|
||||
llm_config=mock_llm_config,
|
||||
agent_name="test_agent",
|
||||
config=hindsight_config,
|
||||
pool=None,
|
||||
operation_id=None,
|
||||
schema=None,
|
||||
)
|
||||
|
||||
# Verify batch API was called
|
||||
mock_llm_config._provider_impl.submit_batch.assert_called_once()
|
||||
|
||||
logger.info("✅ Routing to batch API test passed")
|
||||
|
||||
finally:
|
||||
# Cleanup
|
||||
try:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
except Exception:
|
||||
pass
|
||||
@@ -1,263 +0,0 @@
|
||||
"""
|
||||
Real integration test for OpenAI Batch API.
|
||||
|
||||
This test makes REAL API calls to OpenAI and measures actual timing.
|
||||
It will be slow (minutes to hours) depending on OpenAI's queue.
|
||||
|
||||
To run:
|
||||
pytest tests/test_batch_api_integration.py -v -s
|
||||
|
||||
To skip in CI:
|
||||
Add @pytest.mark.skip at the test level
|
||||
"""
|
||||
import pytest
|
||||
import os
|
||||
import asyncio
|
||||
import logging
|
||||
import time
|
||||
from datetime import datetime, timezone
|
||||
from dotenv import load_dotenv
|
||||
from hindsight_api import RequestContext
|
||||
from hindsight_api.engine.retain.fact_extraction import (
|
||||
extract_facts_from_contents_batch_api,
|
||||
RetainContent,
|
||||
)
|
||||
from hindsight_api.config import HindsightConfig
|
||||
from hindsight_api.engine.llm_wrapper import LLMProvider
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Load .env file for API keys
|
||||
load_dotenv()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def openai_api_key():
|
||||
"""Get OpenAI API key from environment."""
|
||||
# Try both current and commented keys from .env
|
||||
api_key = os.getenv("HINDSIGHT_API_LLM_API_KEY")
|
||||
|
||||
# Check if it's an OpenAI key (starts with sk-proj- or sk-)
|
||||
if not api_key or not api_key.startswith("sk-"):
|
||||
# Try the OpenAI-specific env var (if set separately)
|
||||
api_key = os.getenv("OPENAI_API_KEY")
|
||||
|
||||
if not api_key or not api_key.startswith("sk-"):
|
||||
pytest.skip("OpenAI API key not found in environment. Set OPENAI_API_KEY or uncomment OpenAI config in .env")
|
||||
|
||||
return api_key
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def real_llm_config(openai_api_key):
|
||||
"""Create real LLM config for OpenAI."""
|
||||
# Create config with OpenAI settings
|
||||
config = HindsightConfig.from_env()
|
||||
|
||||
# Use LLMProvider wrapper (which creates _provider_impl internally)
|
||||
llm_config = LLMProvider(
|
||||
provider="openai",
|
||||
api_key=openai_api_key,
|
||||
base_url="https://api.openai.com/v1",
|
||||
model="gpt-4o-mini", # Fast, cheap model for testing
|
||||
reasoning_effort="medium", # Required parameter
|
||||
)
|
||||
|
||||
return llm_config
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def test_contents_real():
|
||||
"""Create realistic test content for fact extraction."""
|
||||
return [
|
||||
RetainContent(
|
||||
content="""
|
||||
Alice is a senior software engineer at TechCorp, where she has been working for 5 years.
|
||||
She specializes in distributed systems and microservices architecture. Alice graduated
|
||||
from MIT with a degree in Computer Science in 2015. She is known for writing clean,
|
||||
well-documented code and mentoring junior developers.
|
||||
""",
|
||||
event_date=datetime(2024, 1, 15, 10, 30, tzinfo=timezone.utc),
|
||||
context="team member profile",
|
||||
),
|
||||
RetainContent(
|
||||
content="""
|
||||
Bob joined TechCorp last month as a junior developer. He is learning React and Node.js
|
||||
and recently completed his first feature, which was a user authentication flow. Bob
|
||||
graduated from Berkeley with a degree in Computer Science in 2023. He is enthusiastic
|
||||
and asks great questions during code reviews.
|
||||
""",
|
||||
event_date=datetime(2024, 1, 15, 10, 30, tzinfo=timezone.utc),
|
||||
context="team member profile",
|
||||
),
|
||||
RetainContent(
|
||||
content="""
|
||||
The team uses Kubernetes for container orchestration and deploys to AWS. They follow
|
||||
agile methodologies with two-week sprints. Code reviews are mandatory before merging
|
||||
any pull request. The team meets every morning for a 15-minute standup to discuss
|
||||
progress and blockers.
|
||||
""",
|
||||
event_date=datetime(2024, 1, 15, 10, 30, tzinfo=timezone.utc),
|
||||
context="team processes",
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def integration_config():
|
||||
"""Create config for integration test."""
|
||||
config = HindsightConfig.from_env()
|
||||
config.retain_batch_enabled = True
|
||||
config.retain_batch_poll_interval_seconds = 30 # Poll every 30 seconds (reasonable for real API)
|
||||
config.retain_chunk_size = 4000
|
||||
config.retain_extraction_mode = "concise"
|
||||
config.retain_extract_causal_links = False
|
||||
return config
|
||||
|
||||
|
||||
@pytest.mark.skip(reason="Real API test - takes minutes and costs money. Run manually with: pytest tests/test_batch_api_integration.py::test_real_openai_batch_api -v -s")
|
||||
@pytest.mark.integration # Mark as integration test
|
||||
@pytest.mark.slow # Mark as slow test
|
||||
@pytest.mark.asyncio
|
||||
async def test_real_openai_batch_api(real_llm_config, test_contents_real, integration_config, memory, request_context):
|
||||
"""
|
||||
REAL integration test: Submit actual batch to OpenAI and measure timing.
|
||||
|
||||
WARNING: This test:
|
||||
- Makes real API calls to OpenAI
|
||||
- Will take minutes to hours to complete
|
||||
- Costs money (though very little with gpt-4o-mini)
|
||||
- Requires valid OpenAI API key
|
||||
|
||||
To skip this test:
|
||||
pytest tests/test_batch_api_integration.py --skip-integration
|
||||
"""
|
||||
bank_id = f"test_real_batch_{datetime.now(timezone.utc).timestamp()}"
|
||||
|
||||
logger.info("=" * 80)
|
||||
logger.info("STARTING REAL OPENAI BATCH API INTEGRATION TEST")
|
||||
logger.info("=" * 80)
|
||||
logger.info(f"Test contents: {len(test_contents_real)} items")
|
||||
logger.info(f"Poll interval: {integration_config.retain_batch_poll_interval_seconds}s")
|
||||
logger.info(f"Model: {real_llm_config.model}")
|
||||
logger.info("This may take several minutes to hours depending on OpenAI's queue...")
|
||||
logger.info("=" * 80)
|
||||
|
||||
try:
|
||||
# Ensure bank exists
|
||||
await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
|
||||
# Get database pool and schema for crash recovery testing
|
||||
pool = memory._pool
|
||||
schema = request_context.tenant_id
|
||||
|
||||
# Track overall timing
|
||||
test_start_time = time.time()
|
||||
|
||||
# Call REAL batch API extraction
|
||||
logger.info("\n📤 Submitting batch to OpenAI...")
|
||||
|
||||
facts, chunks, usage = await extract_facts_from_contents_batch_api(
|
||||
contents=test_contents_real,
|
||||
llm_config=real_llm_config,
|
||||
agent_name="test_agent",
|
||||
config=integration_config,
|
||||
pool=pool,
|
||||
operation_id=None, # No crash recovery for this test
|
||||
schema=schema,
|
||||
)
|
||||
|
||||
test_end_time = time.time()
|
||||
total_duration = test_end_time - test_start_time
|
||||
|
||||
# Log results
|
||||
logger.info("\n" + "=" * 80)
|
||||
logger.info("✅ BATCH COMPLETED SUCCESSFULLY")
|
||||
logger.info("=" * 80)
|
||||
logger.info(f"Total duration: {total_duration:.1f} seconds ({total_duration/60:.1f} minutes)")
|
||||
logger.info(f"Facts extracted: {len(facts)}")
|
||||
logger.info(f"Chunks processed: {len(chunks)}")
|
||||
logger.info(f"Token usage: {usage.input_tokens} input + {usage.output_tokens} output = {usage.total_tokens} total")
|
||||
logger.info(f"Estimated cost: ${(usage.input_tokens * 0.00015 / 1000 + usage.output_tokens * 0.0006 / 1000):.4f}")
|
||||
logger.info("=" * 80)
|
||||
|
||||
# Log sample facts
|
||||
logger.info("\n📋 Sample extracted facts:")
|
||||
for i, fact in enumerate(facts[:5]): # Show first 5 facts
|
||||
logger.info(f"\nFact {i+1}:")
|
||||
logger.info(f" Type: {fact.fact_type}")
|
||||
logger.info(f" Text: {fact.fact_text[:100]}...")
|
||||
logger.info(f" Entities: {fact.entities}")
|
||||
|
||||
# Verify results
|
||||
assert len(facts) > 0, "Should extract at least some facts"
|
||||
assert len(chunks) == len(test_contents_real), f"Should have {len(test_contents_real)} chunks"
|
||||
assert usage.total_tokens > 0, "Should have token usage"
|
||||
|
||||
# Verify fact structure
|
||||
for fact in facts:
|
||||
assert hasattr(fact, "fact_text"), "Fact should have fact_text"
|
||||
assert hasattr(fact, "fact_type"), "Fact should have fact_type"
|
||||
assert fact.fact_type in ["world", "experience", "opinion"], f"Invalid fact_type: {fact.fact_type}"
|
||||
|
||||
logger.info("\n✅ All assertions passed!")
|
||||
|
||||
# Write timing report to file for later analysis
|
||||
report_path = "/tmp/openai_batch_api_timing_report.txt"
|
||||
with open(report_path, "w") as f:
|
||||
f.write(f"OpenAI Batch API Integration Test Report\n")
|
||||
f.write(f"={'=' * 60}\n\n")
|
||||
f.write(f"Test Date: {datetime.now(timezone.utc).isoformat()}\n")
|
||||
f.write(f"Model: {real_llm_config.model}\n")
|
||||
f.write(f"Contents: {len(test_contents_real)} items\n")
|
||||
f.write(f"Poll Interval: {integration_config.retain_batch_poll_interval_seconds}s\n\n")
|
||||
f.write(f"Results:\n")
|
||||
f.write(f" Total Duration: {total_duration:.1f}s ({total_duration/60:.1f} min)\n")
|
||||
f.write(f" Facts Extracted: {len(facts)}\n")
|
||||
f.write(f" Chunks Processed: {len(chunks)}\n")
|
||||
f.write(f" Token Usage: {usage.total_tokens} ({usage.input_tokens} in + {usage.output_tokens} out)\n")
|
||||
f.write(f" Estimated Cost: ${(usage.input_tokens * 0.00015 / 1000 + usage.output_tokens * 0.0006 / 1000):.4f}\n")
|
||||
|
||||
logger.info(f"\n📄 Timing report written to: {report_path}")
|
||||
|
||||
finally:
|
||||
# Cleanup
|
||||
try:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
logger.info(f"\n🧹 Cleaned up test bank: {bank_id}")
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to cleanup bank: {e}")
|
||||
|
||||
|
||||
@pytest.mark.skip(reason="Real API test - requires Groq API key. Run manually if needed.")
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.asyncio
|
||||
async def test_real_batch_supports_groq(integration_config):
|
||||
"""
|
||||
Test that Groq also supports batch API (if configured).
|
||||
|
||||
Groq has the same batch API interface as OpenAI.
|
||||
"""
|
||||
groq_api_key = os.getenv("HINDSIGHT_API_LLM_API_KEY")
|
||||
|
||||
if not groq_api_key or not groq_api_key.startswith("gsk_"):
|
||||
pytest.skip("Groq API key not found in environment")
|
||||
|
||||
llm_config = LLMProvider(
|
||||
provider="groq",
|
||||
api_key=groq_api_key,
|
||||
base_url="https://api.groq.com/openai/v1",
|
||||
model="llama-3.1-8b-instant",
|
||||
reasoning_effort="medium",
|
||||
)
|
||||
|
||||
# Check if Groq supports batch API
|
||||
supports_batch = await llm_config._provider_impl.supports_batch_api()
|
||||
|
||||
logger.info(f"Groq batch API support: {supports_batch}")
|
||||
|
||||
# Groq should support batch API (same interface as OpenAI)
|
||||
assert supports_batch, "Groq should support batch API"
|
||||
|
||||
logger.info("✅ Groq batch API support confirmed")
|
||||
@@ -1,38 +0,0 @@
|
||||
"""
|
||||
Test validation for batch API + synchronous retain.
|
||||
|
||||
When HINDSIGHT_API_RETAIN_BATCH_ENABLED=true, synchronous retain operations
|
||||
should be rejected with a 400 error since they will timeout.
|
||||
"""
|
||||
|
||||
import os
|
||||
import pytest
|
||||
from hindsight_api.engine.memory_engine import MemoryEngine
|
||||
from hindsight_api.config import HindsightConfig
|
||||
from hindsight_api import RequestContext
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_batch_api_validation(memory, request_context):
|
||||
"""
|
||||
Test that attempting synchronous retain with batch API enabled
|
||||
raises an error at the HTTP layer.
|
||||
|
||||
This test verifies the validation logic exists - actual HTTP testing
|
||||
would require full FastAPI app setup.
|
||||
"""
|
||||
# Create config with batch API enabled
|
||||
config = HindsightConfig.from_env()
|
||||
config.retain_batch_enabled = True
|
||||
config.retain_batch_poll_interval_seconds = 1
|
||||
|
||||
# Verify the validation exists in memory engine
|
||||
# The actual HTTP validation happens in http.py api_retain()
|
||||
# This test documents the expected behavior
|
||||
|
||||
assert config.retain_batch_enabled is True
|
||||
assert config.retain_batch_poll_interval_seconds == 1
|
||||
|
||||
# When batch API is enabled and async=false, the HTTP endpoint
|
||||
# should return 400 with message:
|
||||
# "Batch API is enabled (HINDSIGHT_API_RETAIN_BATCH_ENABLED=true) but async=false"
|
||||
@@ -1,171 +0,0 @@
|
||||
"""
|
||||
Tests for combined scoring (apply_combined_scoring).
|
||||
|
||||
The function applies multiplicative recency/temporal boosts to the cross-encoder
|
||||
score so that the relative influence of these signals is proportional to the base
|
||||
relevance score, independent of the cross-encoder model's score calibration.
|
||||
"""
|
||||
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from hindsight_api.engine.search.reranking import apply_combined_scoring, _RECENCY_ALPHA, _TEMPORAL_ALPHA
|
||||
from hindsight_api.engine.search.types import MergedCandidate, RetrievalResult, ScoredResult
|
||||
|
||||
UTC = timezone.utc
|
||||
NOW = datetime(2024, 6, 1, tzinfo=UTC)
|
||||
|
||||
|
||||
def _make_result(
|
||||
ce_norm: float,
|
||||
occurred_start: datetime | None = None,
|
||||
temporal_proximity: float | None = None,
|
||||
) -> ScoredResult:
|
||||
retrieval = MagicMock(spec=RetrievalResult)
|
||||
retrieval.occurred_start = occurred_start
|
||||
retrieval.temporal_proximity = temporal_proximity
|
||||
|
||||
candidate = MagicMock(spec=MergedCandidate)
|
||||
candidate.retrieval = retrieval
|
||||
candidate.rrf_score = 0.05
|
||||
|
||||
return ScoredResult(
|
||||
candidate=candidate,
|
||||
cross_encoder_score=1.0,
|
||||
cross_encoder_score_normalized=ce_norm,
|
||||
weight=ce_norm,
|
||||
)
|
||||
|
||||
|
||||
class TestBoostFormula:
|
||||
def test_neutral_signals_leave_score_unchanged(self):
|
||||
"""recency=0.5 and temporal=0.5 both produce boost=1.0, so weight == ce."""
|
||||
sr = _make_result(ce_norm=0.6)
|
||||
apply_combined_scoring([sr], now=NOW)
|
||||
assert abs(sr.weight - 0.6) < 1e-9
|
||||
|
||||
def test_max_recency_boost(self):
|
||||
"""A memory from today (recency≈1.0) should boost by (1 + alpha*0.5)."""
|
||||
sr = _make_result(ce_norm=0.5, occurred_start=NOW)
|
||||
apply_combined_scoring([sr], now=NOW)
|
||||
expected = 0.5 * (1.0 + _RECENCY_ALPHA * 0.5) * 1.0 # temporal neutral
|
||||
assert abs(sr.weight - expected) < 1e-6
|
||||
|
||||
def test_min_recency_penalty(self):
|
||||
"""A memory from >365 days ago (recency=0.1) should penalise score."""
|
||||
old = NOW - timedelta(days=400)
|
||||
sr = _make_result(ce_norm=0.5, occurred_start=old)
|
||||
apply_combined_scoring([sr], now=NOW)
|
||||
expected = 0.5 * (1.0 + _RECENCY_ALPHA * (0.1 - 0.5)) * 1.0
|
||||
assert abs(sr.weight - expected) < 1e-6
|
||||
|
||||
def test_max_temporal_boost(self):
|
||||
"""temporal_proximity=1.0 should boost by (1 + alpha*0.5)."""
|
||||
sr = _make_result(ce_norm=0.5, temporal_proximity=1.0)
|
||||
apply_combined_scoring([sr], now=NOW)
|
||||
expected = 0.5 * 1.0 * (1.0 + _TEMPORAL_ALPHA * 0.5) # recency neutral
|
||||
assert abs(sr.weight - expected) < 1e-6
|
||||
|
||||
def test_temporal_none_is_neutral(self):
|
||||
"""temporal_proximity=None must be treated as 0.5 (no boost/penalty)."""
|
||||
sr_none = _make_result(ce_norm=0.5, temporal_proximity=None)
|
||||
sr_half = _make_result(ce_norm=0.5, temporal_proximity=0.5)
|
||||
apply_combined_scoring([sr_none], now=NOW)
|
||||
apply_combined_scoring([sr_half], now=NOW)
|
||||
assert abs(sr_none.weight - sr_half.weight) < 1e-9
|
||||
|
||||
def test_both_signals_combined(self):
|
||||
"""Both boosts are applied multiplicatively."""
|
||||
sr = _make_result(ce_norm=0.5, occurred_start=NOW, temporal_proximity=1.0)
|
||||
apply_combined_scoring([sr], now=NOW)
|
||||
recency_boost = 1.0 + _RECENCY_ALPHA * (1.0 - 0.5)
|
||||
temporal_boost = 1.0 + _TEMPORAL_ALPHA * (1.0 - 0.5)
|
||||
expected = 0.5 * recency_boost * temporal_boost
|
||||
assert abs(sr.weight - expected) < 1e-6
|
||||
|
||||
def test_boost_is_proportional_to_ce(self):
|
||||
"""The absolute boost from recency scales with the CE score."""
|
||||
sr_high = _make_result(ce_norm=0.9, occurred_start=NOW)
|
||||
sr_low = _make_result(ce_norm=0.3, occurred_start=NOW)
|
||||
apply_combined_scoring([sr_high, sr_low], now=NOW)
|
||||
|
||||
# Both get the same recency boost factor — absolute gain is proportional to CE
|
||||
boost_factor = 1.0 + _RECENCY_ALPHA * 0.5
|
||||
assert abs(sr_high.weight - 0.9 * boost_factor) < 1e-6
|
||||
assert abs(sr_low.weight - 0.3 * boost_factor) < 1e-6
|
||||
|
||||
def test_boost_capped(self):
|
||||
"""Max boost: recency=1.0 + temporal=1.0 gives ≤21% uplift on CE."""
|
||||
sr = _make_result(ce_norm=1.0, occurred_start=NOW, temporal_proximity=1.0)
|
||||
apply_combined_scoring([sr], now=NOW)
|
||||
assert sr.weight <= 1.0 * (1 + _RECENCY_ALPHA / 2) * (1 + _TEMPORAL_ALPHA / 2) + 1e-9
|
||||
|
||||
def test_rrf_normalized_always_zero(self):
|
||||
"""RRF is excluded from scoring; rrf_normalized is set to 0.0 for trace clarity."""
|
||||
sr = _make_result(ce_norm=0.5)
|
||||
apply_combined_scoring([sr], now=NOW)
|
||||
assert sr.rrf_normalized == 0.0
|
||||
|
||||
def test_combined_score_equals_weight(self):
|
||||
"""combined_score and weight must stay in sync."""
|
||||
sr = _make_result(ce_norm=0.7, occurred_start=NOW, temporal_proximity=0.8)
|
||||
apply_combined_scoring([sr], now=NOW)
|
||||
assert sr.combined_score == sr.weight
|
||||
|
||||
def test_model_calibration_independence(self):
|
||||
"""
|
||||
A low-calibration model (low CE scores) and a high-calibration model
|
||||
(high CE scores) should produce the same ranking for identical content.
|
||||
|
||||
With additive scoring the recency term would dominate for low-CE models;
|
||||
with multiplicative boosting the relative ranking is stable.
|
||||
"""
|
||||
recent = NOW - timedelta(days=10)
|
||||
old = NOW - timedelta(days=300)
|
||||
|
||||
# High-calibration model: clear winner is #1 (more relevant, slightly older)
|
||||
h_relevant = _make_result(ce_norm=0.85, occurred_start=old)
|
||||
h_recent = _make_result(ce_norm=0.60, occurred_start=recent)
|
||||
apply_combined_scoring([h_relevant, h_recent], now=NOW)
|
||||
assert h_relevant.weight > h_recent.weight, "High-CE model: relevance should win"
|
||||
|
||||
# Low-calibration model: same relative difference, just compressed scores
|
||||
l_relevant = _make_result(ce_norm=0.34, occurred_start=old)
|
||||
l_recent = _make_result(ce_norm=0.24, occurred_start=recent)
|
||||
apply_combined_scoring([l_relevant, l_recent], now=NOW)
|
||||
assert l_relevant.weight > l_recent.weight, "Low-CE model: relevance should still win"
|
||||
|
||||
def test_no_occurred_start_defaults_recency_neutral(self):
|
||||
"""Missing occurred_start → recency=0.5 → no boost/penalty."""
|
||||
sr = _make_result(ce_norm=0.5, occurred_start=None)
|
||||
apply_combined_scoring([sr], now=NOW)
|
||||
assert sr.recency == 0.5
|
||||
assert abs(sr.weight - 0.5) < 1e-9
|
||||
|
||||
def test_timezone_naive_occurred_start_handled(self):
|
||||
"""Naive datetimes in occurred_start should not raise."""
|
||||
naive_date = datetime(2024, 1, 1) # no tzinfo
|
||||
sr = _make_result(ce_norm=0.5, occurred_start=naive_date)
|
||||
apply_combined_scoring([sr], now=NOW) # must not raise
|
||||
assert 0.0 < sr.weight < 1.0
|
||||
|
||||
def test_custom_alpha_values(self):
|
||||
"""Custom alpha parameters are respected."""
|
||||
sr = _make_result(ce_norm=0.5, occurred_start=NOW)
|
||||
apply_combined_scoring([sr], now=NOW, recency_alpha=0.4, temporal_alpha=0.0)
|
||||
expected = 0.5 * (1.0 + 0.4 * 0.5) * 1.0
|
||||
assert abs(sr.weight - expected) < 1e-6
|
||||
|
||||
def test_future_event_recency_capped_at_one(self):
|
||||
"""Events in the future must not produce recency > 1.0, keeping boost within bounds."""
|
||||
future = NOW + timedelta(days=180)
|
||||
sr = _make_result(ce_norm=0.5, occurred_start=future)
|
||||
apply_combined_scoring([sr], now=NOW)
|
||||
assert sr.recency == 1.0
|
||||
expected_max_boost = 1.0 + _RECENCY_ALPHA * 0.5
|
||||
assert sr.weight <= 0.5 * expected_max_boost + 1e-9
|
||||
|
||||
def test_empty_list_is_noop(self):
|
||||
apply_combined_scoring([], now=NOW) # must not raise
|
||||
@@ -1,476 +0,0 @@
|
||||
"""Tests for consolidation failure handling: adaptive batch splitting, consolidation_failed_at,
|
||||
and the recovery API.
|
||||
|
||||
These tests use a mock LLM to simulate LLM failures deterministically, without making real
|
||||
API calls. All tests insert memories directly into the database to bypass retain's LLM calls
|
||||
and focus exclusively on the consolidation code paths.
|
||||
"""
|
||||
|
||||
import uuid
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
|
||||
from hindsight_api.engine.consolidation.consolidator import run_consolidation_job
|
||||
from hindsight_api.engine.memory_engine import MemoryEngine
|
||||
from hindsight_api.engine.providers.mock_llm import MockLLM
|
||||
from hindsight_api.engine.task_backend import SyncTaskBackend
|
||||
|
||||
|
||||
@pytest_asyncio.fixture(scope="function")
|
||||
async def memory_no_llm_verify(pg0_db_url, embeddings, cross_encoder, query_analyzer):
|
||||
"""MemoryEngine with mock LLM.
|
||||
|
||||
Migrations are already applied by the session-scoped pg0_db_url fixture, so
|
||||
run_migrations=False avoids advisory-lock serialization overhead per test.
|
||||
"""
|
||||
mem = MemoryEngine(
|
||||
db_url=pg0_db_url,
|
||||
memory_llm_provider="mock",
|
||||
memory_llm_api_key="",
|
||||
memory_llm_model="mock",
|
||||
embeddings=embeddings,
|
||||
cross_encoder=cross_encoder,
|
||||
query_analyzer=query_analyzer,
|
||||
pool_min_size=1,
|
||||
pool_max_size=5,
|
||||
run_migrations=False,
|
||||
task_backend=SyncTaskBackend(),
|
||||
skip_llm_verification=True,
|
||||
)
|
||||
await mem.initialize()
|
||||
yield mem
|
||||
try:
|
||||
if mem._pool and not mem._pool._closing:
|
||||
await mem.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def enable_observations():
|
||||
"""Enable observations for all tests in this module."""
|
||||
from hindsight_api.config import _get_raw_config
|
||||
|
||||
config = _get_raw_config()
|
||||
original = config.enable_observations
|
||||
config.enable_observations = True
|
||||
yield
|
||||
config.enable_observations = original
|
||||
|
||||
|
||||
def _make_failing_mock_llm(*, fail_first_n: int = 999) -> MockLLM:
|
||||
"""Return a MockLLM that raises ValueError for the first `fail_first_n` consolidation calls."""
|
||||
mock_llm = MockLLM(provider="mock", api_key="", base_url="", model="mock-model")
|
||||
call_count = 0
|
||||
|
||||
def callback(messages, scope):
|
||||
nonlocal call_count
|
||||
if scope == "consolidation":
|
||||
call_count += 1
|
||||
if call_count <= fail_first_n:
|
||||
raise ValueError(f"Simulated LLM failure (call {call_count})")
|
||||
# Return empty response — no creates/updates/deletes
|
||||
from hindsight_api.engine.consolidation.consolidator import _ConsolidationBatchResponse
|
||||
|
||||
return _ConsolidationBatchResponse()
|
||||
|
||||
mock_llm.set_response_callback(callback)
|
||||
return mock_llm
|
||||
|
||||
|
||||
def _make_always_success_mock_llm() -> MockLLM:
|
||||
"""Return a MockLLM that always succeeds with an empty consolidation response."""
|
||||
mock_llm = MockLLM(provider="mock", api_key="", base_url="", model="mock-model")
|
||||
|
||||
def callback(messages, scope):
|
||||
from hindsight_api.engine.consolidation.consolidator import _ConsolidationBatchResponse
|
||||
|
||||
return _ConsolidationBatchResponse()
|
||||
|
||||
mock_llm.set_response_callback(callback)
|
||||
return mock_llm
|
||||
|
||||
|
||||
def _inject_mock_llm(memory: MemoryEngine, mock_llm: MockLLM) -> None:
|
||||
"""Replace memory._consolidation_llm_config with a wrapper that returns mock_llm from with_config."""
|
||||
wrapper = MagicMock()
|
||||
wrapper.with_config.return_value = mock_llm
|
||||
memory._consolidation_llm_config = wrapper
|
||||
|
||||
|
||||
async def _insert_memories(conn, bank_id: str, texts: list[str]) -> list[uuid.UUID]:
|
||||
"""Insert experience memories directly, bypassing LLM-based retain."""
|
||||
ids = []
|
||||
for text in texts:
|
||||
mem_id = uuid.uuid4()
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO memory_units (id, bank_id, text, fact_type, created_at)
|
||||
VALUES ($1, $2, $3, 'experience', now())
|
||||
""",
|
||||
mem_id,
|
||||
bank_id,
|
||||
text,
|
||||
)
|
||||
ids.append(mem_id)
|
||||
return ids
|
||||
|
||||
|
||||
class TestAdaptiveBatchSplitting:
|
||||
"""Verify that a failing batch is halved and retried until batch_size=1 succeeds."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_splitting_recovers_all_memories(self, memory_no_llm_verify: MemoryEngine, request_context):
|
||||
"""When a batch of 2 fails, both are retried individually and succeed."""
|
||||
bank_id = f"test-split-recovery-{uuid.uuid4().hex[:8]}"
|
||||
await memory_no_llm_verify.get_bank_profile(bank_id=bank_id, request_context=request_context)
|
||||
|
||||
async with memory_no_llm_verify._pool.acquire() as conn:
|
||||
mem_ids = await _insert_memories(
|
||||
conn,
|
||||
bank_id,
|
||||
[
|
||||
"Alice runs marathons every spring.",
|
||||
"Alice trained for six months for her last race.",
|
||||
],
|
||||
)
|
||||
|
||||
# Exhaust all 3 retries for batch=2 (calls 1-3 fail), then each batch=1 succeeds (calls 4-5)
|
||||
mock_llm = _make_failing_mock_llm(fail_first_n=3)
|
||||
_inject_mock_llm(memory_no_llm_verify, mock_llm)
|
||||
|
||||
result = await run_consolidation_job(
|
||||
memory_engine=memory_no_llm_verify,
|
||||
bank_id=bank_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert result["status"] == "completed"
|
||||
assert result["memories_processed"] == 2
|
||||
assert result["memories_failed"] == 0
|
||||
|
||||
# Both memories must have consolidated_at set and consolidation_failed_at NULL
|
||||
async with memory_no_llm_verify._pool.acquire() as conn:
|
||||
rows = await conn.fetch(
|
||||
"""
|
||||
SELECT id, consolidated_at, consolidation_failed_at
|
||||
FROM memory_units
|
||||
WHERE bank_id = $1 AND fact_type = 'experience'
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
assert len(rows) == 2
|
||||
for row in rows:
|
||||
assert row["consolidated_at"] is not None, f"Memory {row['id']} should have consolidated_at set"
|
||||
assert row["consolidation_failed_at"] is None, (
|
||||
f"Memory {row['id']} should NOT have consolidation_failed_at set"
|
||||
)
|
||||
|
||||
# LLM called 5 times: 3 retries failed (batch=2) + 1 succeeded (batch=1) + 1 succeeded (batch=1)
|
||||
consolidation_calls = [c for c in mock_llm.get_mock_calls() if c["scope"] == "consolidation"]
|
||||
assert len(consolidation_calls) == 5
|
||||
|
||||
await memory_no_llm_verify.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_splitting_with_larger_batch(self, memory_no_llm_verify: MemoryEngine, request_context):
|
||||
"""A batch of 4 that always fails at size>1 resolves to 4 individual calls."""
|
||||
bank_id = f"test-split-large-{uuid.uuid4().hex[:8]}"
|
||||
await memory_no_llm_verify.get_bank_profile(bank_id=bank_id, request_context=request_context)
|
||||
|
||||
async with memory_no_llm_verify._pool.acquire() as conn:
|
||||
await _insert_memories(
|
||||
conn,
|
||||
bank_id,
|
||||
[
|
||||
"Bob plays chess competitively.",
|
||||
"Bob won a regional chess tournament.",
|
||||
"Bob practices tactics every morning.",
|
||||
"Bob coaches youth chess on weekends.",
|
||||
],
|
||||
)
|
||||
|
||||
# Exhaust all 3 retries for batch=4 (calls 1-3 fail), then both batch=2 halves succeed
|
||||
# (calls 4-5). This verifies that halving once is sufficient when batch=2 works.
|
||||
mock_llm = _make_failing_mock_llm(fail_first_n=3)
|
||||
_inject_mock_llm(memory_no_llm_verify, mock_llm)
|
||||
|
||||
result = await run_consolidation_job(
|
||||
memory_engine=memory_no_llm_verify,
|
||||
bank_id=bank_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert result["memories_processed"] == 4
|
||||
assert result["memories_failed"] == 0
|
||||
|
||||
async with memory_no_llm_verify._pool.acquire() as conn:
|
||||
rows = await conn.fetch(
|
||||
"SELECT consolidated_at, consolidation_failed_at FROM memory_units "
|
||||
"WHERE bank_id = $1 AND fact_type = 'experience'",
|
||||
bank_id,
|
||||
)
|
||||
assert all(r["consolidated_at"] is not None for r in rows)
|
||||
assert all(r["consolidation_failed_at"] is None for r in rows)
|
||||
|
||||
await memory_no_llm_verify.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
class TestConsolidationFailedAt:
|
||||
"""Verify that consolidation_failed_at is set — and consolidated_at is NOT — when all retries fail."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_single_memory_permanent_failure(self, memory_no_llm_verify: MemoryEngine, request_context):
|
||||
"""A single memory that exhausts all LLM retries gets consolidation_failed_at, not consolidated_at."""
|
||||
bank_id = f"test-perm-fail-{uuid.uuid4().hex[:8]}"
|
||||
await memory_no_llm_verify.get_bank_profile(bank_id=bank_id, request_context=request_context)
|
||||
|
||||
async with memory_no_llm_verify._pool.acquire() as conn:
|
||||
(mem_id,) = await _insert_memories(conn, bank_id, ["Carol enjoys painting watercolors."])
|
||||
|
||||
# Always fail
|
||||
mock_llm = _make_failing_mock_llm(fail_first_n=999)
|
||||
_inject_mock_llm(memory_no_llm_verify, mock_llm)
|
||||
|
||||
result = await run_consolidation_job(
|
||||
memory_engine=memory_no_llm_verify,
|
||||
bank_id=bank_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert result["memories_failed"] == 1
|
||||
assert result["memories_processed"] == 1
|
||||
|
||||
async with memory_no_llm_verify._pool.acquire() as conn:
|
||||
row = await conn.fetchrow(
|
||||
"SELECT consolidated_at, consolidation_failed_at FROM memory_units WHERE id = $1",
|
||||
mem_id,
|
||||
)
|
||||
|
||||
assert row["consolidated_at"] is None, "consolidated_at must NOT be set for a permanently failed memory"
|
||||
assert row["consolidation_failed_at"] is not None, "consolidation_failed_at must be set"
|
||||
|
||||
await memory_no_llm_verify.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_failed_memory_excluded_from_next_run(self, memory_no_llm_verify: MemoryEngine, request_context):
|
||||
"""A memory marked consolidation_failed_at is not re-processed on the next consolidation run."""
|
||||
bank_id = f"test-excluded-{uuid.uuid4().hex[:8]}"
|
||||
await memory_no_llm_verify.get_bank_profile(bank_id=bank_id, request_context=request_context)
|
||||
|
||||
async with memory_no_llm_verify._pool.acquire() as conn:
|
||||
(mem_id,) = await _insert_memories(conn, bank_id, ["Dave collects vinyl records."])
|
||||
# Manually stamp consolidation_failed_at to simulate a prior failed run
|
||||
await conn.execute(
|
||||
"UPDATE memory_units SET consolidation_failed_at = NOW() WHERE id = $1",
|
||||
mem_id,
|
||||
)
|
||||
|
||||
# Even with a healthy LLM, the memory should be skipped
|
||||
mock_llm = _make_always_success_mock_llm()
|
||||
_inject_mock_llm(memory_no_llm_verify, mock_llm)
|
||||
|
||||
result = await run_consolidation_job(
|
||||
memory_engine=memory_no_llm_verify,
|
||||
bank_id=bank_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# No unconsolidated memories to pick up (consolidation_failed_at ≠ NULL, consolidated_at = NULL
|
||||
# but the SELECT filters on consolidated_at IS NULL AND fact_type IN ('experience','world'))
|
||||
assert result["status"] in ("no_new_memories", "completed")
|
||||
if result["status"] == "completed":
|
||||
assert result["memories_processed"] == 0
|
||||
|
||||
# Memory still has consolidation_failed_at set and consolidated_at NULL
|
||||
async with memory_no_llm_verify._pool.acquire() as conn:
|
||||
row = await conn.fetchrow(
|
||||
"SELECT consolidated_at, consolidation_failed_at FROM memory_units WHERE id = $1",
|
||||
mem_id,
|
||||
)
|
||||
assert row["consolidated_at"] is None
|
||||
assert row["consolidation_failed_at"] is not None
|
||||
|
||||
await memory_no_llm_verify.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_partial_batch_failure(self, memory_no_llm_verify: MemoryEngine, request_context):
|
||||
"""In a batch of 2, if only the first individual retry fails, the second still succeeds."""
|
||||
bank_id = f"test-partial-fail-{uuid.uuid4().hex[:8]}"
|
||||
await memory_no_llm_verify.get_bank_profile(bank_id=bank_id, request_context=request_context)
|
||||
|
||||
async with memory_no_llm_verify._pool.acquire() as conn:
|
||||
mem_ids = await _insert_memories(
|
||||
conn,
|
||||
bank_id,
|
||||
[
|
||||
"Eve speaks three languages fluently.",
|
||||
"Eve learned Japanese in two years.",
|
||||
],
|
||||
)
|
||||
|
||||
# Exhaust 3 retries for batch=2 (calls 1-3), exhaust 3 retries for first batch=1 (calls 4-6),
|
||||
# second batch=1 succeeds (call 7)
|
||||
mock_llm = _make_failing_mock_llm(fail_first_n=6)
|
||||
_inject_mock_llm(memory_no_llm_verify, mock_llm)
|
||||
|
||||
result = await run_consolidation_job(
|
||||
memory_engine=memory_no_llm_verify,
|
||||
bank_id=bank_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert result["memories_processed"] == 2
|
||||
assert result["memories_failed"] == 1
|
||||
|
||||
async with memory_no_llm_verify._pool.acquire() as conn:
|
||||
rows = {
|
||||
str(r["id"]): r
|
||||
for r in await conn.fetch(
|
||||
"SELECT id, consolidated_at, consolidation_failed_at FROM memory_units "
|
||||
"WHERE bank_id = $1 AND fact_type = 'experience'",
|
||||
bank_id,
|
||||
)
|
||||
}
|
||||
|
||||
# One should have failed, one should have succeeded
|
||||
failed = [r for r in rows.values() if r["consolidation_failed_at"] is not None]
|
||||
succeeded = [r for r in rows.values() if r["consolidated_at"] is not None]
|
||||
assert len(failed) == 1
|
||||
assert len(succeeded) == 1
|
||||
# They must be different memories
|
||||
assert str(failed[0]["id"]) != str(succeeded[0]["id"])
|
||||
|
||||
await memory_no_llm_verify.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
class TestRecoverConsolidation:
|
||||
"""Verify the retry_failed_consolidation() method and the /consolidation/recover endpoint."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_recover_resets_failed_memories(self, memory_no_llm_verify: MemoryEngine, request_context):
|
||||
"""retry_failed_consolidation resets consolidation_failed_at and consolidated_at."""
|
||||
bank_id = f"test-recover-reset-{uuid.uuid4().hex[:8]}"
|
||||
await memory_no_llm_verify.get_bank_profile(bank_id=bank_id, request_context=request_context)
|
||||
|
||||
async with memory_no_llm_verify._pool.acquire() as conn:
|
||||
ids = await _insert_memories(
|
||||
conn,
|
||||
bank_id,
|
||||
[
|
||||
"Frank is a competitive cyclist.",
|
||||
"Frank completed the Tour de France route.",
|
||||
],
|
||||
)
|
||||
# Mark both as failed
|
||||
for mem_id in ids:
|
||||
await conn.execute(
|
||||
"UPDATE memory_units SET consolidation_failed_at = NOW() WHERE id = $1",
|
||||
mem_id,
|
||||
)
|
||||
|
||||
result = await memory_no_llm_verify.retry_failed_consolidation(
|
||||
bank_id, request_context=request_context
|
||||
)
|
||||
|
||||
assert result["retried_count"] == 2
|
||||
|
||||
async with memory_no_llm_verify._pool.acquire() as conn:
|
||||
rows = await conn.fetch(
|
||||
"SELECT consolidated_at, consolidation_failed_at FROM memory_units "
|
||||
"WHERE bank_id = $1 AND fact_type = 'experience'",
|
||||
bank_id,
|
||||
)
|
||||
assert all(r["consolidation_failed_at"] is None for r in rows), "consolidation_failed_at must be cleared"
|
||||
assert all(r["consolidated_at"] is None for r in rows), "consolidated_at must also be cleared"
|
||||
|
||||
await memory_no_llm_verify.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_recover_returns_zero_when_none_failed(self, memory_no_llm_verify: MemoryEngine, request_context):
|
||||
"""retry_failed_consolidation returns 0 when no memories have failed."""
|
||||
bank_id = f"test-recover-zero-{uuid.uuid4().hex[:8]}"
|
||||
await memory_no_llm_verify.get_bank_profile(bank_id=bank_id, request_context=request_context)
|
||||
|
||||
result = await memory_no_llm_verify.retry_failed_consolidation(
|
||||
bank_id, request_context=request_context
|
||||
)
|
||||
|
||||
assert result["retried_count"] == 0
|
||||
|
||||
await memory_no_llm_verify.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_recover_then_consolidate_succeeds(self, memory_no_llm_verify: MemoryEngine, request_context):
|
||||
"""After recovery, the memory is picked up by the next consolidation run."""
|
||||
bank_id = f"test-recover-consolidate-{uuid.uuid4().hex[:8]}"
|
||||
await memory_no_llm_verify.get_bank_profile(bank_id=bank_id, request_context=request_context)
|
||||
|
||||
async with memory_no_llm_verify._pool.acquire() as conn:
|
||||
(mem_id,) = await _insert_memories(conn, bank_id, ["Grace is an expert rock climber."])
|
||||
await conn.execute(
|
||||
"UPDATE memory_units SET consolidation_failed_at = NOW() WHERE id = $1", mem_id
|
||||
)
|
||||
|
||||
# Recover
|
||||
recover_result = await memory_no_llm_verify.retry_failed_consolidation(
|
||||
bank_id, request_context=request_context
|
||||
)
|
||||
assert recover_result["retried_count"] == 1
|
||||
|
||||
# Now consolidate with a healthy LLM
|
||||
mock_llm = _make_always_success_mock_llm()
|
||||
_inject_mock_llm(memory_no_llm_verify, mock_llm)
|
||||
|
||||
run_result = await run_consolidation_job(
|
||||
memory_engine=memory_no_llm_verify,
|
||||
bank_id=bank_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert run_result["memories_processed"] == 1
|
||||
assert run_result["memories_failed"] == 0
|
||||
|
||||
async with memory_no_llm_verify._pool.acquire() as conn:
|
||||
row = await conn.fetchrow(
|
||||
"SELECT consolidated_at, consolidation_failed_at FROM memory_units WHERE id = $1",
|
||||
mem_id,
|
||||
)
|
||||
assert row["consolidated_at"] is not None, "Memory should be consolidated after recovery"
|
||||
assert row["consolidation_failed_at"] is None
|
||||
|
||||
await memory_no_llm_verify.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_recover_endpoint_via_http(self, memory_no_llm_verify: MemoryEngine, request_context):
|
||||
"""The POST /consolidation/recover endpoint returns the correct retried_count."""
|
||||
import httpx
|
||||
|
||||
from hindsight_api.api.http import create_app
|
||||
|
||||
bank_id = f"test-recover-http-{uuid.uuid4().hex[:8]}"
|
||||
await memory_no_llm_verify.get_bank_profile(bank_id=bank_id, request_context=request_context)
|
||||
|
||||
async with memory_no_llm_verify._pool.acquire() as conn:
|
||||
ids = await _insert_memories(
|
||||
conn,
|
||||
bank_id,
|
||||
["Henry is a professional chef.", "Henry trained at Le Cordon Bleu."],
|
||||
)
|
||||
for mem_id in ids:
|
||||
await conn.execute(
|
||||
"UPDATE memory_units SET consolidation_failed_at = NOW() WHERE id = $1", mem_id
|
||||
)
|
||||
|
||||
app = create_app(memory_no_llm_verify, initialize_memory=False)
|
||||
transport = httpx.ASGITransport(app=app)
|
||||
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
|
||||
response = await client.post(f"/v1/default/banks/{bank_id}/consolidation/recover")
|
||||
|
||||
assert response.status_code == 200
|
||||
body = response.json()
|
||||
assert body["retried_count"] == 2
|
||||
|
||||
await memory_no_llm_verify.delete_bank(bank_id, request_context=request_context)
|
||||
@@ -1,366 +0,0 @@
|
||||
"""
|
||||
Tests for document tracking and upsert functionality.
|
||||
"""
|
||||
import logging
|
||||
from datetime import datetime, timezone
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
from hindsight_api import RequestContext
|
||||
from hindsight_api.engine.response_models import TokenUsage
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_document_creation_and_retrieval(memory, request_context):
|
||||
"""Test that documents are created and can be retrieved."""
|
||||
bank_id = f"test_doc_{datetime.now(timezone.utc).timestamp()}"
|
||||
|
||||
try:
|
||||
document_id = "meeting-001"
|
||||
|
||||
# Store memory with document tracking
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="Alice works at Google. Bob works at Microsoft.",
|
||||
context="Team meeting",
|
||||
document_id=document_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Retrieve document
|
||||
doc = await memory.get_document(document_id, bank_id, request_context=request_context)
|
||||
|
||||
assert doc is not None
|
||||
assert doc["id"] == document_id
|
||||
assert doc["bank_id"] == bank_id
|
||||
assert "Alice works at Google" in doc["original_text"]
|
||||
assert doc["memory_unit_count"] > 0
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_document_upsert(memory, request_context):
|
||||
"""Test that providing the same document_id automatically upserts (deletes old units and creates new ones)."""
|
||||
bank_id = f"test_upsert_{datetime.now(timezone.utc).timestamp()}"
|
||||
|
||||
try:
|
||||
document_id = "meeting-002"
|
||||
|
||||
# First version
|
||||
units_v1 = await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="Alice works at Google.",
|
||||
context="Initial",
|
||||
document_id=document_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Get document stats
|
||||
doc_v1 = await memory.get_document(document_id, bank_id, request_context=request_context)
|
||||
count_v1 = doc_v1["memory_unit_count"]
|
||||
|
||||
# Update with different content (automatic upsert when same document_id is provided)
|
||||
units_v2 = await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="Alice works at Microsoft. Bob works at Apple.",
|
||||
context="Updated",
|
||||
document_id=document_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Get updated document stats
|
||||
doc_v2 = await memory.get_document(document_id, bank_id, request_context=request_context)
|
||||
count_v2 = doc_v2["memory_unit_count"]
|
||||
|
||||
# Verify old units were replaced
|
||||
assert "Microsoft" in doc_v2["original_text"]
|
||||
assert doc_v2["updated_at"] > doc_v1["created_at"]
|
||||
|
||||
# Different unit IDs (old ones deleted, new ones created)
|
||||
assert set(units_v1).isdisjoint(set(units_v2))
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_document_deletion(memory, request_context):
|
||||
"""Test that deleting a document cascades to memory units."""
|
||||
bank_id = f"test_delete_{datetime.now(timezone.utc).timestamp()}"
|
||||
|
||||
try:
|
||||
document_id = "meeting-003"
|
||||
|
||||
# Create document
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="Alice works at Google.",
|
||||
context="Test",
|
||||
document_id=document_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Verify it exists
|
||||
doc = await memory.get_document(document_id, bank_id, request_context=request_context)
|
||||
assert doc is not None
|
||||
assert doc["memory_unit_count"] > 0
|
||||
|
||||
# Delete document
|
||||
result = await memory.delete_document(document_id, bank_id, request_context=request_context)
|
||||
assert result["document_deleted"] == 1
|
||||
assert result["memory_units_deleted"] > 0
|
||||
|
||||
# Verify it's gone
|
||||
doc_after = await memory.get_document(document_id, bank_id, request_context=request_context)
|
||||
assert doc_after is None
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_memory_without_document(memory, request_context):
|
||||
"""Test that memories can still be created without document tracking."""
|
||||
bank_id = f"test_no_doc_{datetime.now(timezone.utc).timestamp()}"
|
||||
|
||||
try:
|
||||
# Create memory without document_id (backward compatibility)
|
||||
units = await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="Alice works at Google.",
|
||||
context="Test",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert len(units) > 0
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_document_persisted_with_zero_facts(memory, request_context):
|
||||
"""
|
||||
Test that documents are persisted even when zero facts are extracted.
|
||||
|
||||
This is a regression test for issue #324 where documents with no extractable
|
||||
facts were reported as disappearing from the system.
|
||||
"""
|
||||
bank_id = f"test_zero_facts_{datetime.now(timezone.utc).timestamp()}"
|
||||
|
||||
try:
|
||||
document_id = "doc-zero-facts"
|
||||
|
||||
# Retain content that produces zero facts (gibberish/random characters)
|
||||
units = await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="xyzabc123 !!!### @@@ $$$", # Random characters unlikely to produce facts
|
||||
context="Test zero facts",
|
||||
document_id=document_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Should return empty unit list (no facts extracted)
|
||||
assert len(units) == 0, "Should extract zero facts from gibberish content"
|
||||
|
||||
# But document should still be persisted and retrievable
|
||||
doc = await memory.get_document(document_id, bank_id, request_context=request_context)
|
||||
assert doc is not None, "Document should be persisted even with zero facts"
|
||||
assert doc["id"] == document_id
|
||||
assert doc["bank_id"] == bank_id
|
||||
assert doc["memory_unit_count"] == 0, "Should have zero memory units"
|
||||
assert len(doc["original_text"]) > 0, "Should have non-zero text length"
|
||||
assert "xyzabc123" in doc["original_text"], "Should contain original content"
|
||||
|
||||
# Document should also appear in list
|
||||
docs_list = await memory.list_documents(
|
||||
bank_id=bank_id,
|
||||
search_query=None,
|
||||
limit=100,
|
||||
offset=0,
|
||||
request_context=request_context,
|
||||
)
|
||||
assert docs_list["total"] == 1, "Document should appear in list"
|
||||
assert any(d["id"] == document_id for d in docs_list["items"]), "Document should be in items"
|
||||
|
||||
listed_doc = next(d for d in docs_list["items"] if d["id"] == document_id)
|
||||
assert listed_doc["memory_unit_count"] == 0, "Listed document should show zero memory units"
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_document_persisted_with_zero_facts_batch(memory, request_context):
|
||||
"""
|
||||
Test that documents are persisted with zero facts in batch retain operations.
|
||||
|
||||
This tests the async batch code path to ensure it also handles zero facts correctly.
|
||||
"""
|
||||
bank_id = f"test_zero_facts_batch_{datetime.now(timezone.utc).timestamp()}"
|
||||
|
||||
try:
|
||||
# Mix of content: some produces facts, some produces zero facts
|
||||
contents = [
|
||||
{
|
||||
"content": "Alice works at Google",
|
||||
"document_id": "doc-with-facts",
|
||||
},
|
||||
{
|
||||
"content": "!@# $$$ %%% ^^^ &&& ***", # Gibberish - zero facts expected
|
||||
"document_id": "doc-zero-facts",
|
||||
},
|
||||
]
|
||||
|
||||
unit_ids = await memory.retain_batch_async(
|
||||
bank_id=bank_id,
|
||||
contents=contents,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# First content should produce facts, second should not
|
||||
assert len(unit_ids[0]) > 0, "First content should produce facts"
|
||||
assert len(unit_ids[1]) == 0, "Second content should produce zero facts"
|
||||
|
||||
# Both documents should be persisted
|
||||
doc_with_facts = await memory.get_document("doc-with-facts", bank_id, request_context=request_context)
|
||||
assert doc_with_facts is not None
|
||||
assert doc_with_facts["memory_unit_count"] > 0
|
||||
|
||||
doc_zero_facts = await memory.get_document("doc-zero-facts", bank_id, request_context=request_context)
|
||||
assert doc_zero_facts is not None, "Document with zero facts should be persisted"
|
||||
assert doc_zero_facts["memory_unit_count"] == 0, "Should have zero memory units"
|
||||
assert "!@#" in doc_zero_facts["original_text"]
|
||||
|
||||
# Both should appear in list
|
||||
docs_list = await memory.list_documents(
|
||||
bank_id=bank_id,
|
||||
search_query=None,
|
||||
limit=100,
|
||||
offset=0,
|
||||
request_context=request_context,
|
||||
)
|
||||
assert docs_list["total"] == 2, "Both documents should appear in list"
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_document_persisted_with_zero_facts_async_submit(memory, request_context):
|
||||
"""
|
||||
Test that documents are persisted with zero facts in fire-and-forget async retain.
|
||||
|
||||
This tests the submit_async_retain (background task) code path to ensure it also
|
||||
handles zero facts correctly.
|
||||
"""
|
||||
import asyncio
|
||||
|
||||
bank_id = f"test_zero_facts_async_{datetime.now(timezone.utc).timestamp()}"
|
||||
|
||||
try:
|
||||
# Submit async retain with gibberish content
|
||||
result = await memory.submit_async_retain(
|
||||
bank_id=bank_id,
|
||||
contents=[
|
||||
{
|
||||
"content": "!@# $$$ %%% ^^^ &&& ***", # Gibberish - zero facts expected
|
||||
"document_id": "doc-async-zero-facts",
|
||||
}
|
||||
],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
operation_id = result["operation_id"]
|
||||
assert operation_id is not None, "Should return operation_id"
|
||||
|
||||
# Wait for background task to complete
|
||||
max_wait = 60 # 60 seconds max
|
||||
wait_interval = 0.5
|
||||
elapsed = 0
|
||||
|
||||
while elapsed < max_wait:
|
||||
await asyncio.sleep(wait_interval)
|
||||
elapsed += wait_interval
|
||||
|
||||
# Check if document exists
|
||||
doc = await memory.get_document(
|
||||
"doc-async-zero-facts", bank_id, request_context=request_context
|
||||
)
|
||||
if doc is not None:
|
||||
break
|
||||
|
||||
# Document should be persisted even with zero facts
|
||||
assert doc is not None, "Document should be persisted after async task completes"
|
||||
assert doc["id"] == "doc-async-zero-facts"
|
||||
assert doc["memory_unit_count"] == 0, "Should have zero memory units"
|
||||
assert "!@#" in doc["original_text"]
|
||||
|
||||
# Document should appear in list
|
||||
docs_list = await memory.list_documents(
|
||||
bank_id=bank_id,
|
||||
search_query=None,
|
||||
limit=100,
|
||||
offset=0,
|
||||
request_context=request_context,
|
||||
)
|
||||
assert docs_list["total"] == 1, "Document should appear in list"
|
||||
assert any(d["id"] == "doc-async-zero-facts" for d in docs_list["items"])
|
||||
|
||||
listed_doc = next(d for d in docs_list["items"] if d["id"] == "doc-async-zero-facts")
|
||||
assert listed_doc["memory_unit_count"] == 0, "Listed document should show zero memory units"
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_document_stored_without_chunks_when_zero_facts(memory_no_llm_verify, request_context):
|
||||
"""
|
||||
Regression test: when 0 facts are extracted from chunked content, the document row
|
||||
must be stored but no chunk rows should be written.
|
||||
"""
|
||||
bank_id = f"test_zero_facts_no_chunks_{datetime.now(timezone.utc).timestamp()}"
|
||||
document_id = "doc-zero-facts-chunked"
|
||||
|
||||
# Content large enough to exceed default retain_chunk_size (3000 chars) so chunking is triggered
|
||||
content = "Alice works at Google. " * 200 # ~4600 chars
|
||||
|
||||
async def mock_llm_zero_facts(*args, **kwargs):
|
||||
response = {"facts": []}
|
||||
if kwargs.get("return_usage", False):
|
||||
return response, TokenUsage(input_tokens=10, output_tokens=2)
|
||||
return response
|
||||
|
||||
try:
|
||||
with patch("hindsight_api.engine.llm_wrapper.LLMProvider.call", new=mock_llm_zero_facts):
|
||||
units = await memory_no_llm_verify.retain_async(
|
||||
bank_id=bank_id,
|
||||
content=content,
|
||||
document_id=document_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert units == [], "Should return no memory units when LLM extracts zero facts"
|
||||
|
||||
# Document row must exist
|
||||
doc = await memory_no_llm_verify.get_document(document_id, bank_id, request_context=request_context)
|
||||
assert doc is not None, "Document row must be stored even when zero facts are extracted"
|
||||
assert doc["id"] == document_id
|
||||
assert doc["memory_unit_count"] == 0
|
||||
|
||||
# No chunk rows should be stored
|
||||
pool = await memory_no_llm_verify._get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
chunk_count = await conn.fetchval(
|
||||
"SELECT COUNT(*) FROM chunks WHERE document_id = $1 AND bank_id = $2",
|
||||
document_id,
|
||||
bank_id,
|
||||
)
|
||||
assert chunk_count == 0, "No chunk rows should be stored when zero facts are extracted"
|
||||
|
||||
finally:
|
||||
await memory_no_llm_verify.delete_bank(bank_id, request_context=request_context)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,70 +0,0 @@
|
||||
"""
|
||||
Tests for EntityResolver edge cases.
|
||||
"""
|
||||
|
||||
import uuid
|
||||
from datetime import datetime, timezone
|
||||
|
||||
import asyncpg
|
||||
import pytest
|
||||
|
||||
from hindsight_api.engine.entity_resolver import EntityResolver
|
||||
from hindsight_api.pg0 import resolve_database_url
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resolve_entities_batch_handles_unicode_lower_conflicts(pg0_db_url):
|
||||
"""
|
||||
Existing entities with PostgreSQL/Python lowercase mismatches should resolve
|
||||
to the conflicted row instead of leaving a missing entity_id.
|
||||
"""
|
||||
resolved_url = await resolve_database_url(pg0_db_url)
|
||||
pool = await asyncpg.create_pool(resolved_url, min_size=1, max_size=2, command_timeout=30)
|
||||
bank_id = f"test-entity-resolver-{uuid.uuid4().hex[:8]}"
|
||||
event_date = datetime(2024, 1, 15, tzinfo=timezone.utc)
|
||||
resolver = EntityResolver(pool=pool, entity_lookup="full")
|
||||
|
||||
try:
|
||||
async with pool.acquire() as conn:
|
||||
existing_entity_id = await conn.fetchval(
|
||||
"""
|
||||
INSERT INTO entities (bank_id, canonical_name, first_seen, last_seen, mention_count)
|
||||
VALUES ($1, $2, $3, $3, 1)
|
||||
RETURNING id
|
||||
""",
|
||||
bank_id,
|
||||
"İstanbul",
|
||||
event_date,
|
||||
)
|
||||
|
||||
resolved_ids = await resolver.resolve_entities_batch(
|
||||
bank_id=bank_id,
|
||||
entities_data=[
|
||||
{
|
||||
"text": "istanbul",
|
||||
"nearby_entities": [],
|
||||
"event_date": event_date,
|
||||
}
|
||||
],
|
||||
context="unicode case mismatch",
|
||||
unit_event_date=event_date,
|
||||
conn=conn,
|
||||
)
|
||||
|
||||
entity_rows = await conn.fetch(
|
||||
"""
|
||||
SELECT id, canonical_name
|
||||
FROM entities
|
||||
WHERE bank_id = $1
|
||||
ORDER BY canonical_name
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
|
||||
assert resolved_ids == [existing_entity_id]
|
||||
assert len(entity_rows) == 1
|
||||
assert entity_rows[0]["id"] == existing_entity_id
|
||||
assert entity_rows[0]["canonical_name"] == "İstanbul"
|
||||
finally:
|
||||
await pool.execute("DELETE FROM entities WHERE bank_id = $1", bank_id)
|
||||
await pool.close()
|
||||
@@ -1,61 +0,0 @@
|
||||
"""
|
||||
Unit tests for metadata inclusion in fact extraction LLM prompt.
|
||||
"""
|
||||
from datetime import datetime
|
||||
|
||||
from hindsight_api.engine.retain.fact_extraction import _build_user_message
|
||||
|
||||
|
||||
def test_build_user_message_includes_metadata():
|
||||
"""Metadata key-value pairs should appear in the user message."""
|
||||
event_date = datetime(2024, 6, 15, 12, 0, 0)
|
||||
metadata = {"title": "Q2 Planning Doc", "source": "confluence", "author": "Alice"}
|
||||
|
||||
msg = _build_user_message(
|
||||
chunk="Some content.",
|
||||
chunk_index=0,
|
||||
total_chunks=1,
|
||||
event_date=event_date,
|
||||
context="planning meeting",
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
assert "title" in msg
|
||||
assert "Q2 Planning Doc" in msg
|
||||
assert "source" in msg
|
||||
assert "confluence" in msg
|
||||
assert "author" in msg
|
||||
assert "Alice" in msg
|
||||
|
||||
|
||||
def test_build_user_message_no_metadata():
|
||||
"""When metadata is empty, the message should still be valid and not include a metadata section."""
|
||||
event_date = datetime(2024, 6, 15, 12, 0, 0)
|
||||
|
||||
msg = _build_user_message(
|
||||
chunk="Some content.",
|
||||
chunk_index=0,
|
||||
total_chunks=1,
|
||||
event_date=event_date,
|
||||
context="planning meeting",
|
||||
metadata={},
|
||||
)
|
||||
|
||||
assert "Some content." in msg
|
||||
assert "Metadata:" not in msg
|
||||
|
||||
|
||||
def test_build_user_message_without_metadata_arg():
|
||||
"""Calling without metadata (default) should behave the same as empty metadata."""
|
||||
event_date = datetime(2024, 6, 15, 12, 0, 0)
|
||||
|
||||
msg = _build_user_message(
|
||||
chunk="Some content.",
|
||||
chunk_index=0,
|
||||
total_chunks=1,
|
||||
event_date=event_date,
|
||||
context="none",
|
||||
)
|
||||
|
||||
assert "Some content." in msg
|
||||
assert "Metadata:" not in msg
|
||||
@@ -1,141 +0,0 @@
|
||||
"""
|
||||
Unit tests for fact extraction retry logic.
|
||||
|
||||
Tests the fix for the TypeError when LLM returns invalid JSON across all retries.
|
||||
Previously, `raise last_error` would raise None (TypeError) because last_error was
|
||||
only set in the BadRequestError handler, not when the LLM returned non-dict JSON.
|
||||
"""
|
||||
|
||||
from datetime import datetime, timezone
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
def _make_config(llm_max_retries: int = 3, retain_llm_max_retries: int | None = None):
|
||||
"""Build a minimal HindsightConfig for fact extraction tests."""
|
||||
from hindsight_api.config import HindsightConfig
|
||||
|
||||
cfg = MagicMock(spec=HindsightConfig)
|
||||
cfg.retain_llm_max_retries = retain_llm_max_retries
|
||||
cfg.llm_max_retries = llm_max_retries
|
||||
cfg.retain_llm_initial_backoff = None
|
||||
cfg.llm_initial_backoff = 0.0
|
||||
cfg.retain_llm_max_backoff = None
|
||||
cfg.llm_max_backoff = 0.0
|
||||
cfg.retain_max_completion_tokens = 8192
|
||||
cfg.retain_extraction_mode = "concise"
|
||||
cfg.retain_extract_causal_links = False
|
||||
cfg.retain_mission = None
|
||||
return cfg
|
||||
|
||||
|
||||
def _make_llm_config(mock_response):
|
||||
"""Build a mock LLMProvider that returns the given response."""
|
||||
from hindsight_api.engine.llm_wrapper import LLMProvider
|
||||
|
||||
llm = MagicMock(spec=LLMProvider)
|
||||
llm.provider = "mock"
|
||||
token_usage = MagicMock()
|
||||
token_usage.__add__ = lambda self, other: self
|
||||
llm.call = AsyncMock(return_value=(mock_response, token_usage))
|
||||
return llm
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_non_dict_json_all_retries_returns_empty():
|
||||
"""
|
||||
When LLM returns non-dict JSON on every attempt, extraction should return []
|
||||
without raising TypeError ('exceptions must derive from BaseException').
|
||||
|
||||
This was the bug: the loop ran range(2) times (hardcoded), but comparisons
|
||||
used config.llm_max_retries (default 10). On the last loop iteration (attempt=1),
|
||||
`attempt < 10 - 1` was True, so the code called `continue`, the loop
|
||||
exhausted, and `raise last_error` raised None → TypeError.
|
||||
"""
|
||||
from hindsight_api.engine.retain.fact_extraction import _extract_facts_from_chunk
|
||||
|
||||
# llm_max_retries=3 ensures the bug triggers with the old code (3 != 2 hardcoded)
|
||||
config = _make_config(llm_max_retries=3, retain_llm_max_retries=None)
|
||||
|
||||
# Mock: always returns a list (non-dict), which is invalid
|
||||
llm_config = _make_llm_config(mock_response=[{"invalid": "response"}])
|
||||
|
||||
with patch(
|
||||
"hindsight_api.engine.retain.fact_extraction._build_extraction_prompt_and_schema",
|
||||
return_value=("system prompt", MagicMock()),
|
||||
):
|
||||
facts, usage = await _extract_facts_from_chunk(
|
||||
chunk="Alice visited Paris in 2023.",
|
||||
chunk_index=0,
|
||||
total_chunks=1,
|
||||
event_date=datetime(2023, 1, 1, tzinfo=timezone.utc),
|
||||
context="travel notes",
|
||||
llm_config=llm_config,
|
||||
config=config,
|
||||
agent_name="test-agent",
|
||||
)
|
||||
|
||||
assert facts == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_non_dict_json_with_default_max_retries_returns_empty():
|
||||
"""
|
||||
Same scenario with the default llm_max_retries=10 (matching real default config).
|
||||
The old code ran range(2) but checked against 10, always continuing until
|
||||
the loop exhausted, then raised None → TypeError.
|
||||
"""
|
||||
from hindsight_api.engine.retain.fact_extraction import _extract_facts_from_chunk
|
||||
|
||||
config = _make_config(llm_max_retries=10, retain_llm_max_retries=None)
|
||||
llm_config = _make_llm_config(mock_response="not a dict at all")
|
||||
|
||||
with patch(
|
||||
"hindsight_api.engine.retain.fact_extraction._build_extraction_prompt_and_schema",
|
||||
return_value=("system prompt", MagicMock()),
|
||||
):
|
||||
facts, usage = await _extract_facts_from_chunk(
|
||||
chunk="Some text.",
|
||||
chunk_index=0,
|
||||
total_chunks=1,
|
||||
event_date=datetime(2023, 6, 1, tzinfo=timezone.utc),
|
||||
context="",
|
||||
llm_config=llm_config,
|
||||
config=config,
|
||||
agent_name="agent",
|
||||
)
|
||||
|
||||
assert facts == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retain_llm_max_retries_overrides_global():
|
||||
"""
|
||||
When retain_llm_max_retries is set, it should be used for the loop range
|
||||
and all comparisons (no shadowing bug).
|
||||
"""
|
||||
from hindsight_api.engine.retain.fact_extraction import _extract_facts_from_chunk
|
||||
|
||||
# retain_llm_max_retries=5 should override llm_max_retries=10
|
||||
config = _make_config(llm_max_retries=10, retain_llm_max_retries=5)
|
||||
llm_config = _make_llm_config(mock_response=42) # non-dict: integer
|
||||
|
||||
with patch(
|
||||
"hindsight_api.engine.retain.fact_extraction._build_extraction_prompt_and_schema",
|
||||
return_value=("system prompt", MagicMock()),
|
||||
):
|
||||
facts, usage = await _extract_facts_from_chunk(
|
||||
chunk="Bob likes Python.",
|
||||
chunk_index=0,
|
||||
total_chunks=1,
|
||||
event_date=datetime(2024, 1, 1, tzinfo=timezone.utc),
|
||||
context="",
|
||||
llm_config=llm_config,
|
||||
config=config,
|
||||
agent_name="agent",
|
||||
)
|
||||
|
||||
assert facts == []
|
||||
# Verify it retried exactly retain_llm_max_retries times
|
||||
assert llm_config.call.call_count == 5
|
||||
@@ -1,839 +0,0 @@
|
||||
"""
|
||||
End-to-end tests for file retain (upload, convert, retain) functionality.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import io
|
||||
import json
|
||||
from datetime import datetime, timezone
|
||||
|
||||
import pytest
|
||||
from httpx import ASGITransport, AsyncClient
|
||||
|
||||
from hindsight_api.extensions import FileConvertResult, OperationValidatorExtension, ValidationResult
|
||||
from hindsight_api.extensions.operation_validator import (
|
||||
RecallContext,
|
||||
RecallResult,
|
||||
ReflectContext,
|
||||
RetainContext,
|
||||
RetainResult,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_pdf_content():
|
||||
"""Create a simple PDF-like content for testing."""
|
||||
# This is a minimal PDF that markitdown can parse
|
||||
return b"""%PDF-1.4
|
||||
1 0 obj
|
||||
<<
|
||||
/Type /Catalog
|
||||
/Pages 2 0 R
|
||||
>>
|
||||
endobj
|
||||
2 0 obj
|
||||
<<
|
||||
/Type /Pages
|
||||
/Kids [3 0 R]
|
||||
/Count 1
|
||||
>>
|
||||
endobj
|
||||
3 0 obj
|
||||
<<
|
||||
/Type /Page
|
||||
/Parent 2 0 R
|
||||
/MediaBox [0 0 612 792]
|
||||
/Contents 4 0 R
|
||||
/Resources <<
|
||||
/Font <<
|
||||
/F1 <<
|
||||
/Type /Font
|
||||
/Subtype /Type1
|
||||
/BaseFont /Helvetica
|
||||
>>
|
||||
>>
|
||||
>>
|
||||
>>
|
||||
endobj
|
||||
4 0 obj
|
||||
<<
|
||||
/Length 44
|
||||
>>
|
||||
stream
|
||||
BT
|
||||
/F1 12 Tf
|
||||
100 700 Td
|
||||
(Test Document) Tj
|
||||
ET
|
||||
endstream
|
||||
endobj
|
||||
xref
|
||||
0 5
|
||||
0000000000 65535 f
|
||||
0000000009 00000 n
|
||||
0000000058 00000 n
|
||||
0000000115 00000 n
|
||||
0000000317 00000 n
|
||||
trailer
|
||||
<<
|
||||
/Size 5
|
||||
/Root 1 0 R
|
||||
>>
|
||||
startxref
|
||||
410
|
||||
%%EOF
|
||||
"""
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_txt_content():
|
||||
"""Create simple text content."""
|
||||
return b"This is a test document.\nIt contains some important information.\nAlice works at Google."
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_file_retain_basic(memory_no_llm_verify, sample_txt_content):
|
||||
"""Test basic file upload and conversion."""
|
||||
from hindsight_api.api.http import create_app
|
||||
|
||||
app = create_app(memory_no_llm_verify, initialize_memory=False)
|
||||
|
||||
async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client:
|
||||
# Create a bank first
|
||||
bank_response = await client.put("/v1/default/banks/test-file-bank", json={"name": "Test File Bank"})
|
||||
assert bank_response.status_code in (200, 201)
|
||||
|
||||
# Upload file
|
||||
request_data = {
|
||||
"document_tags": ["test"],
|
||||
"async": True,
|
||||
}
|
||||
|
||||
files = {"files": ("test.txt", sample_txt_content, "text/plain")}
|
||||
data = {"request": json.dumps(request_data)}
|
||||
|
||||
response = await client.post(
|
||||
"/v1/default/banks/test-file-bank/files/retain",
|
||||
files=files,
|
||||
data=data,
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
result = response.json()
|
||||
assert "operation_ids" in result
|
||||
assert len(result["operation_ids"]) == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_file_retain_with_metadata(memory_no_llm_verify, sample_txt_content):
|
||||
"""Test file upload with per-file metadata."""
|
||||
from hindsight_api.api.http import create_app
|
||||
|
||||
app = create_app(memory_no_llm_verify, initialize_memory=False)
|
||||
|
||||
async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client:
|
||||
# Create bank
|
||||
bank_response = await client.put("/v1/default/banks/test-file-meta-bank", json={"name": "Test Meta Bank"})
|
||||
assert bank_response.status_code in (200, 201)
|
||||
|
||||
# Upload file with metadata
|
||||
request_data = {
|
||||
"document_tags": ["work", "reports"],
|
||||
"async": True,
|
||||
"files_metadata": [
|
||||
{
|
||||
"document_id": "test_doc_123",
|
||||
"context": "quarterly report",
|
||||
"metadata": {"author": "Alice", "year": "2024"},
|
||||
"tags": ["Q1"],
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
files = {"files": ("report.txt", sample_txt_content, "text/plain")}
|
||||
data = {"request": json.dumps(request_data)}
|
||||
|
||||
response = await client.post(
|
||||
"/v1/default/banks/test-file-meta-bank/files/retain",
|
||||
files=files,
|
||||
data=data,
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
result = response.json()
|
||||
assert "operation_ids" in result
|
||||
assert len(result["operation_ids"]) == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_file_retain_multiple_files(memory_no_llm_verify, sample_txt_content):
|
||||
"""Test uploading multiple files at once."""
|
||||
from hindsight_api.api.http import create_app
|
||||
|
||||
app = create_app(memory_no_llm_verify, initialize_memory=False)
|
||||
|
||||
async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client:
|
||||
# Create bank
|
||||
bank_response = await client.put("/v1/default/banks/test-multi-file-bank", json={"name": "Test Multi Bank"})
|
||||
assert bank_response.status_code in (200, 201)
|
||||
|
||||
# Upload multiple files
|
||||
request_data = {
|
||||
"async": True,
|
||||
"files_metadata": [
|
||||
{"document_id": "doc1", "tags": ["file1"]},
|
||||
{"document_id": "doc2", "tags": ["file2"]},
|
||||
],
|
||||
}
|
||||
|
||||
content1 = b"First document content"
|
||||
content2 = b"Second document content"
|
||||
|
||||
files = [
|
||||
("files", ("file1.txt", content1, "text/plain")),
|
||||
("files", ("file2.txt", content2, "text/plain")),
|
||||
]
|
||||
data = {"request": json.dumps(request_data)}
|
||||
|
||||
response = await client.post(
|
||||
"/v1/default/banks/test-multi-file-bank/files/retain",
|
||||
files=files,
|
||||
data=data,
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
result = response.json()
|
||||
assert "operation_ids" in result
|
||||
assert len(result["operation_ids"]) == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_file_retain_validation_errors(memory_no_llm_verify):
|
||||
"""Test validation errors."""
|
||||
from hindsight_api.api.http import create_app
|
||||
|
||||
app = create_app(memory_no_llm_verify, initialize_memory=False)
|
||||
|
||||
async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client:
|
||||
# Create bank
|
||||
bank_response = await client.put("/v1/default/banks/test-validation-bank", json={"name": "Test Validation Bank"})
|
||||
assert bank_response.status_code in (200, 201)
|
||||
|
||||
# Test: metadata count mismatch
|
||||
request_data = {
|
||||
"async": True,
|
||||
"files_metadata": [
|
||||
{"document_id": "doc1"},
|
||||
{"document_id": "doc2"}, # 2 metadata entries
|
||||
],
|
||||
}
|
||||
|
||||
files = {"files": ("file1.txt", b"content", "text/plain")} # But only 1 file
|
||||
data = {"request": json.dumps(request_data)}
|
||||
|
||||
response = await client.post(
|
||||
"/v1/default/banks/test-validation-bank/files/retain",
|
||||
files=files,
|
||||
data=data,
|
||||
)
|
||||
|
||||
assert response.status_code == 400
|
||||
assert "files_metadata count" in response.json()["detail"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_file_retain_no_files(memory_no_llm_verify):
|
||||
"""Test error when no files provided."""
|
||||
from hindsight_api.api.http import create_app
|
||||
|
||||
app = create_app(memory_no_llm_verify, initialize_memory=False)
|
||||
|
||||
async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client:
|
||||
# Create bank
|
||||
bank_response = await client.put("/v1/default/banks/test-no-files-bank", json={"name": "Test No Files Bank"})
|
||||
assert bank_response.status_code in (200, 201)
|
||||
|
||||
request_data = {
|
||||
"async": True,
|
||||
}
|
||||
|
||||
# No files provided
|
||||
data = {"request": json.dumps(request_data)}
|
||||
|
||||
response = await client.post(
|
||||
"/v1/default/banks/test-no-files-bank/files/retain",
|
||||
data=data,
|
||||
)
|
||||
|
||||
# FastAPI will return 422 for missing required field
|
||||
assert response.status_code == 422
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_file_retain_sync_not_supported(memory_no_llm_verify, sample_txt_content):
|
||||
"""Test that file retain is always async (sync is not supported)."""
|
||||
from hindsight_api.api.http import create_app
|
||||
|
||||
app = create_app(memory_no_llm_verify, initialize_memory=False)
|
||||
|
||||
async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client:
|
||||
# Create bank
|
||||
bank_response = await client.put("/v1/default/banks/test-sync-bank", json={"name": "Test Sync Bank"})
|
||||
assert bank_response.status_code in (200, 201)
|
||||
|
||||
# File retain is always async - just verify it succeeds and returns operation_ids
|
||||
files = {"files": ("test.txt", sample_txt_content, "text/plain")}
|
||||
data = {"request": json.dumps({})}
|
||||
|
||||
response = await client.post(
|
||||
"/v1/default/banks/test-sync-bank/files/retain",
|
||||
files=files,
|
||||
data=data,
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
result = response.json()
|
||||
assert "operation_ids" in result
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_file_storage_postgresql(memory_no_llm_verify, sample_txt_content):
|
||||
"""Test file storage in PostgreSQL."""
|
||||
# Test that files are stored and retrieved correctly
|
||||
storage = memory_no_llm_verify._file_storage
|
||||
|
||||
# Store a file
|
||||
key = "test/file1.txt"
|
||||
stored_key = await storage.store(
|
||||
file_data=sample_txt_content,
|
||||
key=key,
|
||||
metadata={"content_type": "text/plain"},
|
||||
)
|
||||
|
||||
assert stored_key == key
|
||||
|
||||
# Retrieve the file
|
||||
retrieved = await storage.retrieve(key)
|
||||
assert retrieved == sample_txt_content
|
||||
|
||||
# Check if file exists
|
||||
exists = await storage.exists(key)
|
||||
assert exists is True
|
||||
|
||||
# Delete the file
|
||||
await storage.delete(key)
|
||||
|
||||
# Check file no longer exists
|
||||
exists_after = await storage.exists(key)
|
||||
assert exists_after is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_markitdown_converter():
|
||||
"""Test markitdown parser."""
|
||||
from hindsight_api.engine.parsers import MarkitdownParser
|
||||
|
||||
parser = MarkitdownParser()
|
||||
|
||||
# Test simple text file
|
||||
text_content = b"This is a test document.\nWith multiple lines."
|
||||
result = await parser.convert(text_content, "test.txt")
|
||||
|
||||
assert isinstance(result, str)
|
||||
assert len(result) > 0
|
||||
assert "test document" in result.lower() or "multiple lines" in result.lower()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_converter_registry():
|
||||
"""Test file parser registry."""
|
||||
from hindsight_api.engine.parsers import FileParserRegistry, MarkitdownParser
|
||||
|
||||
registry = FileParserRegistry()
|
||||
parser = MarkitdownParser()
|
||||
registry.register(parser)
|
||||
|
||||
# Test get by name
|
||||
retrieved = registry.get_parser("markitdown", "test.txt")
|
||||
assert retrieved is parser
|
||||
|
||||
# Test auto-detection
|
||||
auto = registry.get_parser(None, "test.pdf")
|
||||
assert auto is parser
|
||||
|
||||
# Test unsupported format
|
||||
with pytest.raises(ValueError, match="No parser found"):
|
||||
registry.get_parser(None, "test.xyz")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_file_conversion_creates_separate_retain_operation(memory_no_llm_verify, sample_txt_content):
|
||||
"""Test that file conversion and retain are two separate async operations.
|
||||
|
||||
The file_convert_retain task should:
|
||||
1. Convert the file to markdown
|
||||
2. In a single transaction: create a separate 'retain' operation AND mark itself as 'completed'
|
||||
3. Free the worker slot immediately after conversion
|
||||
|
||||
The retain then runs as its own task. This prevents deadlocks where file conversion
|
||||
tasks hold worker slots while waiting for inline retain to finish.
|
||||
"""
|
||||
from hindsight_api.models import RequestContext
|
||||
|
||||
bank_id = "test_file_two_phase_bank"
|
||||
|
||||
context = RequestContext(internal=True)
|
||||
await memory_no_llm_verify.get_bank_profile(bank_id, request_context=context)
|
||||
|
||||
class MockFile:
|
||||
def __init__(self, content, filename, content_type):
|
||||
self.content = content
|
||||
self.filename = filename
|
||||
self.content_type = content_type
|
||||
|
||||
async def read(self):
|
||||
return self.content
|
||||
|
||||
mock_file = MockFile(sample_txt_content, "test.txt", "text/plain")
|
||||
|
||||
file_items = [
|
||||
{
|
||||
"file": mock_file,
|
||||
"document_id": "test_doc_two_phase",
|
||||
"context": "test context",
|
||||
"metadata": {"source": "test"},
|
||||
"tags": ["test_tag"],
|
||||
"timestamp": None,
|
||||
"parser": ["markitdown"],
|
||||
}
|
||||
]
|
||||
|
||||
result = await memory_no_llm_verify.submit_async_file_retain(
|
||||
bank_id=bank_id,
|
||||
file_items=file_items,
|
||||
document_tags=["two_phase_test"],
|
||||
request_context=context,
|
||||
)
|
||||
|
||||
assert "operation_ids" in result
|
||||
assert len(result["operation_ids"]) == 1
|
||||
convert_operation_id = result["operation_ids"][0]
|
||||
|
||||
import asyncio
|
||||
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
pool = await memory_no_llm_verify._get_pool()
|
||||
from hindsight_api.engine.memory_engine import get_current_schema
|
||||
|
||||
schema = get_current_schema()
|
||||
|
||||
async with pool.acquire() as conn:
|
||||
# 1. The file_convert_retain operation must be completed
|
||||
convert_op = await conn.fetchrow(
|
||||
f"SELECT status, operation_type FROM {schema}.async_operations WHERE operation_id = $1",
|
||||
convert_operation_id,
|
||||
)
|
||||
assert convert_op is not None
|
||||
assert convert_op["operation_type"] == "file_convert_retain"
|
||||
assert convert_op["status"] == "completed", (
|
||||
f"file_convert_retain should be 'completed' after conversion, got '{convert_op['status']}'"
|
||||
)
|
||||
|
||||
# 2. A separate retain operation must have been created
|
||||
retain_op = await conn.fetchrow(
|
||||
f"""
|
||||
SELECT status, operation_type
|
||||
FROM {schema}.async_operations
|
||||
WHERE bank_id = $1 AND operation_type = 'retain' AND operation_id != $2
|
||||
""",
|
||||
bank_id,
|
||||
convert_operation_id,
|
||||
)
|
||||
assert retain_op is not None, "A separate 'retain' operation should have been created by file conversion"
|
||||
# With SyncTaskBackend the retain runs immediately, so it should be completed
|
||||
assert retain_op["status"] == "completed"
|
||||
|
||||
# 3. The document should exist with file metadata and retained content
|
||||
doc = await conn.fetchrow(
|
||||
f"""
|
||||
SELECT id, original_text, file_original_name, file_content_type
|
||||
FROM {schema}.documents
|
||||
WHERE id = $1 AND bank_id = $2
|
||||
""",
|
||||
"test_doc_two_phase",
|
||||
bank_id,
|
||||
)
|
||||
|
||||
assert doc is not None
|
||||
assert doc["file_original_name"] == "test.txt"
|
||||
assert doc["file_content_type"] == "text/plain"
|
||||
assert doc["original_text"] is not None
|
||||
assert len(doc["original_text"]) > 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_file_retain_serializes_datetime_timestamp(memory_no_llm_verify, sample_txt_content):
|
||||
"""Async file retain should accept Python datetimes in task payloads."""
|
||||
from hindsight_api.engine.parsers.base import FileParser
|
||||
from hindsight_api.models import RequestContext
|
||||
|
||||
bank_id = f"test_file_timestamp_bank_{datetime.now(timezone.utc).timestamp()}"
|
||||
timestamp = datetime(2024, 1, 15, 10, 30, tzinfo=timezone.utc)
|
||||
|
||||
context = RequestContext(internal=True)
|
||||
await memory_no_llm_verify.get_bank_profile(bank_id, request_context=context)
|
||||
|
||||
class MockFile:
|
||||
def __init__(self, content, filename, content_type):
|
||||
self.content = content
|
||||
self.filename = filename
|
||||
self.content_type = content_type
|
||||
|
||||
async def read(self):
|
||||
return self.content
|
||||
|
||||
class TimestampParser(FileParser):
|
||||
async def convert(self, file_data: bytes, filename: str) -> str:
|
||||
return file_data.decode("utf-8")
|
||||
|
||||
def supports(self, filename: str, content_type: str | None = None) -> bool:
|
||||
return filename.endswith(".txt")
|
||||
|
||||
def name(self) -> str:
|
||||
return "timestamp_parser"
|
||||
|
||||
memory_no_llm_verify._parser_registry.register(TimestampParser())
|
||||
|
||||
mock_file = MockFile(sample_txt_content, "timestamped.txt", "text/plain")
|
||||
|
||||
result = await memory_no_llm_verify.submit_async_file_retain(
|
||||
bank_id=bank_id,
|
||||
file_items=[
|
||||
{
|
||||
"file": mock_file,
|
||||
"document_id": "timestamped_doc",
|
||||
"context": "timestamp test",
|
||||
"metadata": {},
|
||||
"tags": [],
|
||||
"timestamp": timestamp,
|
||||
"parser": ["timestamp_parser"],
|
||||
}
|
||||
],
|
||||
document_tags=None,
|
||||
request_context=context,
|
||||
)
|
||||
|
||||
operation_id = result["operation_ids"][0]
|
||||
pool = await memory_no_llm_verify._get_pool()
|
||||
from hindsight_api.engine.memory_engine import get_current_schema
|
||||
|
||||
async with pool.acquire() as conn:
|
||||
row = await conn.fetchrow(
|
||||
f"""
|
||||
SELECT status, task_payload->>'timestamp' AS timestamp
|
||||
FROM {get_current_schema()}.async_operations
|
||||
WHERE operation_id = $1
|
||||
""",
|
||||
operation_id,
|
||||
)
|
||||
|
||||
assert row is not None
|
||||
assert row["status"] == "completed"
|
||||
assert row["timestamp"] == "2024-01-15T10:30:00+00:00"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_file_conversion_failure_sets_status_to_failed(memory_no_llm_verify, sample_txt_content):
|
||||
"""Test that when file conversion fails, the operation status is set to 'failed' not 'completed'."""
|
||||
from hindsight_api.engine.parsers.base import FileParser
|
||||
from hindsight_api.models import RequestContext
|
||||
|
||||
bank_id = "test_file_failure_bank"
|
||||
|
||||
# Create a mock parser that always fails
|
||||
class FailingParser(FileParser):
|
||||
"""Mock parser that raises an error."""
|
||||
|
||||
async def convert(self, file_data: bytes, filename: str) -> str:
|
||||
# Simulate conversion failure
|
||||
raise RuntimeError(f"Failed to convert '{filename}': Mock conversion error")
|
||||
|
||||
def supports(self, filename: str, content_type: str | None = None) -> bool:
|
||||
return filename.endswith(".fail")
|
||||
|
||||
def name(self) -> str:
|
||||
return "failing_converter"
|
||||
|
||||
# Register the failing parser
|
||||
failing_converter = FailingParser()
|
||||
memory_no_llm_verify._parser_registry.register(failing_converter)
|
||||
|
||||
# Create bank
|
||||
context = RequestContext(internal=True)
|
||||
await memory_no_llm_verify.get_bank_profile(bank_id, request_context=context)
|
||||
|
||||
# Create mock file
|
||||
class MockFile:
|
||||
def __init__(self, content, filename, content_type):
|
||||
self.content = content
|
||||
self.filename = filename
|
||||
self.content_type = content_type
|
||||
|
||||
async def read(self):
|
||||
return self.content
|
||||
|
||||
mock_file = MockFile(sample_txt_content, "test.fail", "application/octet-stream")
|
||||
|
||||
file_items = [
|
||||
{
|
||||
"file": mock_file,
|
||||
"document_id": "test_doc_fail",
|
||||
"context": None,
|
||||
"metadata": {},
|
||||
"tags": [],
|
||||
"timestamp": None,
|
||||
"parser": ["failing_converter"],
|
||||
}
|
||||
]
|
||||
|
||||
# Submit async file retain with failing parser
|
||||
result = await memory_no_llm_verify.submit_async_file_retain(
|
||||
bank_id=bank_id,
|
||||
file_items=file_items,
|
||||
document_tags=None,
|
||||
request_context=context,
|
||||
)
|
||||
|
||||
assert "operation_ids" in result
|
||||
assert len(result["operation_ids"]) == 1
|
||||
operation_id = result["operation_ids"][0]
|
||||
|
||||
# Wait for async processing (with SyncTaskBackend, this is immediate)
|
||||
import asyncio
|
||||
|
||||
await asyncio.sleep(0.2)
|
||||
|
||||
# Check operation status - should be 'failed' not 'completed'
|
||||
pool = await memory_no_llm_verify._get_pool()
|
||||
from hindsight_api.engine.memory_engine import get_current_schema
|
||||
|
||||
async with pool.acquire() as conn:
|
||||
operation = await conn.fetchrow(
|
||||
f"""
|
||||
SELECT status, error_message
|
||||
FROM {get_current_schema()}.async_operations
|
||||
WHERE operation_id = $1
|
||||
""",
|
||||
operation_id,
|
||||
)
|
||||
|
||||
assert operation is not None, f"Operation {operation_id} not found"
|
||||
assert operation["status"] == "failed", f"Expected status 'failed' but got '{operation['status']}'"
|
||||
assert operation["error_message"] is not None
|
||||
assert "Mock conversion error" in operation["error_message"]
|
||||
assert "test.fail" in operation["error_message"]
|
||||
|
||||
|
||||
class FileConvertTrackingValidator(OperationValidatorExtension):
|
||||
"""Validator that tracks on_file_convert_complete hook calls."""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__({})
|
||||
self.convert_calls: list[FileConvertResult] = []
|
||||
|
||||
async def validate_retain(self, ctx: RetainContext) -> ValidationResult:
|
||||
return ValidationResult.accept()
|
||||
|
||||
async def validate_recall(self, ctx: RecallContext) -> ValidationResult:
|
||||
return ValidationResult.accept()
|
||||
|
||||
async def validate_reflect(self, ctx: ReflectContext) -> ValidationResult:
|
||||
return ValidationResult.accept()
|
||||
|
||||
async def on_retain_complete(self, result: RetainResult) -> None:
|
||||
pass
|
||||
|
||||
async def on_recall_complete(self, result: RecallResult) -> None:
|
||||
pass
|
||||
|
||||
async def on_file_convert_complete(self, result: FileConvertResult) -> None:
|
||||
self.convert_calls.append(result)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_on_file_convert_complete_hook_called(memory_no_llm_verify, sample_txt_content):
|
||||
"""Test that on_file_convert_complete hook is called after file conversion with correct parameters."""
|
||||
from hindsight_api.models import RequestContext
|
||||
|
||||
bank_id = "test_file_convert_hook_bank"
|
||||
validator = FileConvertTrackingValidator()
|
||||
memory_no_llm_verify._operation_validator = validator
|
||||
|
||||
context = RequestContext(internal=True, api_key_id="test-key-id", tenant_id="test-tenant")
|
||||
await memory_no_llm_verify.get_bank_profile(bank_id, request_context=context)
|
||||
|
||||
class MockFile:
|
||||
def __init__(self, content, filename, content_type):
|
||||
self.content = content
|
||||
self.filename = filename
|
||||
self.content_type = content_type
|
||||
|
||||
async def read(self):
|
||||
return self.content
|
||||
|
||||
mock_file = MockFile(sample_txt_content, "report.txt", "text/plain")
|
||||
|
||||
file_items = [
|
||||
{
|
||||
"file": mock_file,
|
||||
"document_id": "hook_test_doc",
|
||||
"context": "test context",
|
||||
"metadata": {},
|
||||
"tags": [],
|
||||
"timestamp": None,
|
||||
"parser": ["markitdown"],
|
||||
}
|
||||
]
|
||||
|
||||
await memory_no_llm_verify.submit_async_file_retain(
|
||||
bank_id=bank_id,
|
||||
file_items=file_items,
|
||||
document_tags=None,
|
||||
request_context=context,
|
||||
)
|
||||
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
assert len(validator.convert_calls) == 1
|
||||
result = validator.convert_calls[0]
|
||||
assert result.bank_id == bank_id
|
||||
assert result.filename == "report.txt"
|
||||
assert result.parser_name == "markitdown"
|
||||
assert result.output_chars > 0
|
||||
assert result.output_text is not None
|
||||
assert len(result.output_text) == result.output_chars
|
||||
assert result.success is True
|
||||
assert result.error is None
|
||||
assert result.request_context is not None
|
||||
assert result.request_context.api_key_id == "test-key-id"
|
||||
assert result.request_context.tenant_id == "test-tenant"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_on_file_convert_complete_hook_called_for_each_file(memory_no_llm_verify, sample_txt_content):
|
||||
"""Test that on_file_convert_complete is called once per file when uploading multiple files."""
|
||||
from hindsight_api.models import RequestContext
|
||||
|
||||
bank_id = "test_file_convert_hook_multi_bank"
|
||||
validator = FileConvertTrackingValidator()
|
||||
memory_no_llm_verify._operation_validator = validator
|
||||
|
||||
context = RequestContext(internal=True)
|
||||
await memory_no_llm_verify.get_bank_profile(bank_id, request_context=context)
|
||||
|
||||
class MockFile:
|
||||
def __init__(self, content, filename, content_type):
|
||||
self.content = content
|
||||
self.filename = filename
|
||||
self.content_type = content_type
|
||||
|
||||
async def read(self):
|
||||
return self.content
|
||||
|
||||
file_items = [
|
||||
{
|
||||
"file": MockFile(b"First document content", "first.txt", "text/plain"),
|
||||
"document_id": "doc_1",
|
||||
"context": None,
|
||||
"metadata": {},
|
||||
"tags": [],
|
||||
"timestamp": None,
|
||||
"parser": ["markitdown"],
|
||||
},
|
||||
{
|
||||
"file": MockFile(b"Second document content", "second.txt", "text/plain"),
|
||||
"document_id": "doc_2",
|
||||
"context": None,
|
||||
"metadata": {},
|
||||
"tags": [],
|
||||
"timestamp": None,
|
||||
"parser": ["markitdown"],
|
||||
},
|
||||
]
|
||||
|
||||
await memory_no_llm_verify.submit_async_file_retain(
|
||||
bank_id=bank_id,
|
||||
file_items=file_items,
|
||||
document_tags=None,
|
||||
request_context=context,
|
||||
)
|
||||
|
||||
await asyncio.sleep(0.2)
|
||||
|
||||
assert len(validator.convert_calls) == 2
|
||||
filenames = {r.filename for r in validator.convert_calls}
|
||||
assert filenames == {"first.txt", "second.txt"}
|
||||
for result in validator.convert_calls:
|
||||
assert result.bank_id == bank_id
|
||||
assert result.parser_name == "markitdown"
|
||||
assert result.output_chars > 0
|
||||
assert result.success is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_on_file_convert_complete_hook_not_called_on_conversion_failure(memory_no_llm_verify, sample_txt_content):
|
||||
"""Test that on_file_convert_complete is NOT called when file conversion fails."""
|
||||
from hindsight_api.engine.parsers.base import FileParser
|
||||
from hindsight_api.models import RequestContext
|
||||
|
||||
bank_id = "test_file_convert_hook_fail_bank"
|
||||
validator = FileConvertTrackingValidator()
|
||||
memory_no_llm_verify._operation_validator = validator
|
||||
|
||||
class FailingParser(FileParser):
|
||||
async def convert(self, file_data: bytes, filename: str) -> str:
|
||||
raise RuntimeError("Mock conversion failure")
|
||||
|
||||
def supports(self, filename: str, content_type: str | None = None) -> bool:
|
||||
return filename.endswith(".hookfail")
|
||||
|
||||
def name(self) -> str:
|
||||
return "hookfail_parser"
|
||||
|
||||
memory_no_llm_verify._parser_registry.register(FailingParser())
|
||||
|
||||
context = RequestContext(internal=True)
|
||||
await memory_no_llm_verify.get_bank_profile(bank_id, request_context=context)
|
||||
|
||||
class MockFile:
|
||||
def __init__(self, content, filename, content_type):
|
||||
self.content = content
|
||||
self.filename = filename
|
||||
self.content_type = content_type
|
||||
|
||||
async def read(self):
|
||||
return self.content
|
||||
|
||||
file_items = [
|
||||
{
|
||||
"file": MockFile(sample_txt_content, "bad.hookfail", "application/octet-stream"),
|
||||
"document_id": "fail_hook_doc",
|
||||
"context": None,
|
||||
"metadata": {},
|
||||
"tags": [],
|
||||
"timestamp": None,
|
||||
"parser": ["hookfail_parser"],
|
||||
}
|
||||
]
|
||||
|
||||
await memory_no_llm_verify.submit_async_file_retain(
|
||||
bank_id=bank_id,
|
||||
file_items=file_items,
|
||||
document_tags=None,
|
||||
request_context=context,
|
||||
)
|
||||
|
||||
await asyncio.sleep(0.2)
|
||||
|
||||
assert len(validator.convert_calls) == 0
|
||||
@@ -1,257 +0,0 @@
|
||||
"""
|
||||
Integration tests for S3FileStorage against a SeaweedFS Docker container.
|
||||
|
||||
SeaweedFS (Apache 2.0) provides an S3-compatible API via `weed server -s3`.
|
||||
Requires Docker to be running. Tests are skipped automatically if Docker is unavailable.
|
||||
"""
|
||||
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import subprocess
|
||||
import tempfile
|
||||
import time
|
||||
import uuid
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from httpx import ASGITransport, AsyncClient
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
try:
|
||||
from testcontainers.core.container import DockerContainer
|
||||
|
||||
_has_testcontainers = True
|
||||
except ImportError:
|
||||
_has_testcontainers = False
|
||||
|
||||
_in_ci = os.getenv("CI") == "true"
|
||||
|
||||
pytestmark = [
|
||||
pytest.mark.skipif(not _has_testcontainers, reason="testcontainers not installed"),
|
||||
pytest.mark.skipif(_in_ci, reason="SeaweedFS Docker image pull too slow in CI"),
|
||||
pytest.mark.timeout(300),
|
||||
]
|
||||
|
||||
SEAWEEDFS_S3_PORT = 8333
|
||||
TEST_BUCKET = "hindsight-test"
|
||||
ACCESS_KEY = "test_access_key"
|
||||
SECRET_KEY = "test_secret_key"
|
||||
|
||||
# SeaweedFS S3 IAM config granting full access to our test credentials
|
||||
_S3_CONFIG = {
|
||||
"identities": [
|
||||
{
|
||||
"name": "test-user",
|
||||
"credentials": [{"accessKey": ACCESS_KEY, "secretKey": SECRET_KEY}],
|
||||
"actions": ["Admin", "Read", "Write", "List"],
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
|
||||
def _docker_available() -> bool:
|
||||
"""Check if Docker daemon is running."""
|
||||
try:
|
||||
result = subprocess.run(
|
||||
["docker", "info"],
|
||||
capture_output=True,
|
||||
timeout=5,
|
||||
)
|
||||
return result.returncode == 0
|
||||
except (FileNotFoundError, subprocess.TimeoutExpired):
|
||||
return False
|
||||
|
||||
|
||||
def _wait_for_seaweedfs(endpoint: str, timeout: int = 30) -> None:
|
||||
"""Poll SeaweedFS S3 endpoint until ready."""
|
||||
deadline = time.time() + timeout
|
||||
while time.time() < deadline:
|
||||
try:
|
||||
resp = httpx.get(endpoint, timeout=2)
|
||||
# 200 = no auth, 403 = auth enabled but gateway is up — either means ready
|
||||
if resp.status_code in (200, 403):
|
||||
logger.info("SeaweedFS S3 is ready at %s", endpoint)
|
||||
return
|
||||
except httpx.HTTPError:
|
||||
pass
|
||||
time.sleep(0.5)
|
||||
raise TimeoutError(f"SeaweedFS did not become ready at {endpoint} within {timeout}s")
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def seaweedfs_container():
|
||||
"""Start a SeaweedFS container for the test module, shared across all tests.
|
||||
|
||||
Mounts an s3.json config file to set up S3 credentials for the test user.
|
||||
"""
|
||||
if not _docker_available():
|
||||
pytest.skip("Docker is not available")
|
||||
|
||||
# Write S3 IAM config to a temp file that persists for the module scope
|
||||
s3_config_file = tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False)
|
||||
json.dump(_S3_CONFIG, s3_config_file)
|
||||
s3_config_file.flush()
|
||||
|
||||
container = (
|
||||
DockerContainer(image="chrislusf/seaweedfs:latest")
|
||||
.with_exposed_ports(SEAWEEDFS_S3_PORT)
|
||||
.with_volume_mapping(s3_config_file.name, "/etc/seaweedfs/s3.json", "ro")
|
||||
.with_command(
|
||||
f"server -s3 -s3.port={SEAWEEDFS_S3_PORT} -s3.config=/etc/seaweedfs/s3.json -ip.bind=0.0.0.0"
|
||||
)
|
||||
)
|
||||
|
||||
container.start()
|
||||
|
||||
try:
|
||||
host = container.get_container_host_ip()
|
||||
port = container.get_exposed_port(SEAWEEDFS_S3_PORT)
|
||||
endpoint = f"http://{host}:{port}"
|
||||
|
||||
_wait_for_seaweedfs(endpoint, timeout=240)
|
||||
|
||||
# Create test bucket using obstore (proper SigV4 signing)
|
||||
import obstore as obs
|
||||
from obstore.store import S3Store
|
||||
|
||||
admin_store = S3Store(
|
||||
TEST_BUCKET,
|
||||
endpoint=endpoint,
|
||||
region="us-east-1",
|
||||
access_key_id=ACCESS_KEY,
|
||||
secret_access_key=SECRET_KEY,
|
||||
allow_http=True,
|
||||
)
|
||||
# SeaweedFS auto-creates buckets on first write
|
||||
obs.put(admin_store, ".bucket-init", b"")
|
||||
obs.delete(admin_store, ".bucket-init")
|
||||
logger.info("Test bucket '%s' is ready", TEST_BUCKET)
|
||||
|
||||
yield {
|
||||
"endpoint": endpoint,
|
||||
"access_key": ACCESS_KEY,
|
||||
"secret_key": SECRET_KEY,
|
||||
"bucket": TEST_BUCKET,
|
||||
}
|
||||
finally:
|
||||
container.stop()
|
||||
import os
|
||||
|
||||
os.unlink(s3_config_file.name)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def s3_storage(seaweedfs_container):
|
||||
"""Create an S3FileStorage instance pointing at the SeaweedFS container."""
|
||||
from hindsight_api.engine.storage.s3 import S3FileStorage
|
||||
|
||||
return S3FileStorage(
|
||||
bucket=seaweedfs_container["bucket"],
|
||||
region="us-east-1",
|
||||
endpoint=seaweedfs_container["endpoint"],
|
||||
access_key_id=seaweedfs_container["access_key"],
|
||||
secret_access_key=seaweedfs_container["secret_key"],
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_s3_storage_store_and_retrieve(s3_storage):
|
||||
"""Store a file, retrieve it, verify bytes match."""
|
||||
content = b"Hello, SeaweedFS! This is a test file."
|
||||
key = f"test/{uuid.uuid4()}.txt"
|
||||
|
||||
stored_key = await s3_storage.store(
|
||||
file_data=content,
|
||||
key=key,
|
||||
metadata={"content_type": "text/plain"},
|
||||
)
|
||||
assert stored_key == key
|
||||
|
||||
retrieved = await s3_storage.retrieve(key)
|
||||
assert retrieved == content
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_s3_storage_exists_and_delete(s3_storage):
|
||||
"""Store, check exists=True, delete, check exists=False."""
|
||||
content = b"File to be deleted."
|
||||
key = f"test/{uuid.uuid4()}.txt"
|
||||
|
||||
await s3_storage.store(file_data=content, key=key)
|
||||
|
||||
assert await s3_storage.exists(key) is True
|
||||
|
||||
await s3_storage.delete(key)
|
||||
|
||||
assert await s3_storage.exists(key) is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_s3_storage_file_not_found(s3_storage):
|
||||
"""Retrieve a non-existent key, expect FileNotFoundError."""
|
||||
with pytest.raises(FileNotFoundError):
|
||||
await s3_storage.retrieve(f"nonexistent/{uuid.uuid4()}.txt")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_s3_storage_get_download_url(s3_storage):
|
||||
"""Store a file, get a presigned URL, verify it's a valid URL string."""
|
||||
content = b"Presigned URL test content."
|
||||
key = f"test/{uuid.uuid4()}.txt"
|
||||
|
||||
await s3_storage.store(file_data=content, key=key)
|
||||
|
||||
url = await s3_storage.get_download_url(key, expires_in=300)
|
||||
assert isinstance(url, str)
|
||||
assert url.startswith("http")
|
||||
assert key in url
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_s3_file_retain_api_end_to_end(seaweedfs_container, memory_no_llm_verify):
|
||||
"""Full HTTP API flow: upload file via /files/retain with S3 storage backend."""
|
||||
from hindsight_api.api.http import create_app
|
||||
from hindsight_api.engine.storage.s3 import S3FileStorage
|
||||
|
||||
# Swap the engine's file storage to use the SeaweedFS-backed S3 storage
|
||||
original_storage = memory_no_llm_verify._file_storage
|
||||
s3_storage = S3FileStorage(
|
||||
bucket=seaweedfs_container["bucket"],
|
||||
region="us-east-1",
|
||||
endpoint=seaweedfs_container["endpoint"],
|
||||
access_key_id=seaweedfs_container["access_key"],
|
||||
secret_access_key=seaweedfs_container["secret_key"],
|
||||
)
|
||||
memory_no_llm_verify._file_storage = s3_storage
|
||||
|
||||
try:
|
||||
app = create_app(memory_no_llm_verify, initialize_memory=False)
|
||||
|
||||
async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client:
|
||||
bank_id = f"test-s3-bank-{uuid.uuid4().hex[:8]}"
|
||||
bank_response = await client.put(f"/v1/default/banks/{bank_id}", json={"name": "S3 Test Bank"})
|
||||
assert bank_response.status_code in (200, 201)
|
||||
|
||||
txt_content = b"Alice works at Acme Corp. She joined in 2024."
|
||||
request_data = {
|
||||
"document_tags": ["s3-test"],
|
||||
"async": True,
|
||||
}
|
||||
|
||||
files = {"files": ("notes.txt", txt_content, "text/plain")}
|
||||
data = {"request": json.dumps(request_data)}
|
||||
|
||||
response = await client.post(
|
||||
f"/v1/default/banks/{bank_id}/files/retain",
|
||||
files=files,
|
||||
data=data,
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
result = response.json()
|
||||
assert "operation_ids" in result
|
||||
assert len(result["operation_ids"]) == 1
|
||||
finally:
|
||||
memory_no_llm_verify._file_storage = original_storage
|
||||
@@ -1,340 +0,0 @@
|
||||
"""
|
||||
Tests for Gemini safety settings feature.
|
||||
|
||||
Verifies that:
|
||||
- Safety settings are read from env var and stored on GeminiLLM instances
|
||||
- Settings are applied to GenerateContentConfig in call() and call_with_tools()
|
||||
- The context variable override allows per-bank settings at request time
|
||||
- None (unset) means Gemini's default safety settings are used (no override)
|
||||
"""
|
||||
|
||||
import os
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
pytest.importorskip("google.genai")
|
||||
|
||||
|
||||
SAMPLE_SAFETY_SETTINGS = [
|
||||
{"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_NONE"},
|
||||
{"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE"},
|
||||
{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "BLOCK_NONE"},
|
||||
{"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_NONE"},
|
||||
]
|
||||
|
||||
|
||||
# ─── Config / env var parsing ─────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_gemini_safety_settings_parsed_from_env():
|
||||
"""Safety settings JSON from env var is parsed into HindsightConfig."""
|
||||
import json
|
||||
|
||||
from hindsight_api.config import ENV_LLM_GEMINI_SAFETY_SETTINGS, HindsightConfig, clear_config_cache
|
||||
|
||||
settings_json = json.dumps(SAMPLE_SAFETY_SETTINGS)
|
||||
with patch.dict(os.environ, {ENV_LLM_GEMINI_SAFETY_SETTINGS: settings_json}, clear=False):
|
||||
clear_config_cache()
|
||||
config = HindsightConfig.from_env()
|
||||
assert config.llm_gemini_safety_settings == SAMPLE_SAFETY_SETTINGS
|
||||
clear_config_cache()
|
||||
|
||||
|
||||
def test_gemini_safety_settings_default_is_none():
|
||||
"""When env var is not set, llm_gemini_safety_settings defaults to None."""
|
||||
from hindsight_api.config import ENV_LLM_GEMINI_SAFETY_SETTINGS, HindsightConfig, clear_config_cache
|
||||
|
||||
env = {k: v for k, v in os.environ.items() if k != ENV_LLM_GEMINI_SAFETY_SETTINGS}
|
||||
with patch.dict(os.environ, env, clear=True):
|
||||
clear_config_cache()
|
||||
config = HindsightConfig.from_env()
|
||||
assert config.llm_gemini_safety_settings is None
|
||||
clear_config_cache()
|
||||
|
||||
|
||||
def test_gemini_safety_settings_is_configurable_field():
|
||||
"""llm_gemini_safety_settings appears in configurable (per-bank) fields."""
|
||||
from hindsight_api.config import HindsightConfig
|
||||
|
||||
assert "llm_gemini_safety_settings" in HindsightConfig.get_configurable_fields()
|
||||
|
||||
|
||||
def test_gemini_safety_settings_not_in_credential_fields():
|
||||
"""llm_gemini_safety_settings is NOT a credential — it is safe to expose via API."""
|
||||
from hindsight_api.config import HindsightConfig
|
||||
|
||||
assert "llm_gemini_safety_settings" not in HindsightConfig.get_credential_fields()
|
||||
|
||||
|
||||
# ─── GeminiLLM instance ───────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def _make_gemini_provider(safety_settings=None):
|
||||
"""Return a GeminiLLM instance with a mocked genai.Client."""
|
||||
with patch("google.genai.Client") as mock_client_cls:
|
||||
mock_client_cls.return_value = MagicMock()
|
||||
from hindsight_api.engine.providers.gemini_llm import GeminiLLM
|
||||
|
||||
provider = GeminiLLM(
|
||||
provider="gemini",
|
||||
api_key="fake-api-key",
|
||||
base_url="",
|
||||
model="gemini-2.5-flash",
|
||||
gemini_safety_settings=safety_settings,
|
||||
)
|
||||
# Replace client with a fresh mock so we can inspect calls
|
||||
provider._client = MagicMock()
|
||||
return provider
|
||||
|
||||
|
||||
def test_gemini_llm_stores_safety_settings():
|
||||
"""GeminiLLM stores safety settings passed at construction."""
|
||||
provider = _make_gemini_provider(safety_settings=SAMPLE_SAFETY_SETTINGS)
|
||||
assert provider._safety_settings == SAMPLE_SAFETY_SETTINGS
|
||||
|
||||
|
||||
def test_gemini_llm_no_safety_settings_is_none():
|
||||
"""GeminiLLM._safety_settings is None when not provided."""
|
||||
provider = _make_gemini_provider(safety_settings=None)
|
||||
assert provider._safety_settings is None
|
||||
|
||||
|
||||
# ─── call() applies safety settings ──────────────────────────────────────────
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_applies_safety_settings():
|
||||
"""call() includes safety_settings in GenerateContentConfig when configured."""
|
||||
from google.genai import types as genai_types
|
||||
|
||||
provider = _make_gemini_provider(safety_settings=SAMPLE_SAFETY_SETTINGS)
|
||||
|
||||
# Build a fake successful response
|
||||
fake_response = MagicMock()
|
||||
fake_response.text = "hello"
|
||||
fake_response.candidates = [MagicMock(finish_reason="STOP")]
|
||||
fake_response.usage_metadata = MagicMock(prompt_token_count=5, candidates_token_count=2)
|
||||
|
||||
provider._client.aio.models.generate_content = AsyncMock(return_value=fake_response)
|
||||
|
||||
await provider.call(
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
scope="test",
|
||||
)
|
||||
|
||||
# Inspect the config passed to generate_content
|
||||
call_args = provider._client.aio.models.generate_content.call_args
|
||||
config_arg = call_args.kwargs.get("config") or call_args.args[0] if call_args.args else None
|
||||
# config may be in kwargs or positional; grab from kwargs
|
||||
config_arg = call_args.kwargs.get("config")
|
||||
|
||||
assert config_arg is not None, "GenerateContentConfig should have been passed"
|
||||
assert hasattr(config_arg, "safety_settings"), "Config should have safety_settings"
|
||||
assert config_arg.safety_settings is not None
|
||||
|
||||
categories = [s.category.value if hasattr(s.category, "value") else str(s.category) for s in config_arg.safety_settings]
|
||||
assert "HARM_CATEGORY_HARASSMENT" in categories
|
||||
assert "HARM_CATEGORY_HATE_SPEECH" in categories
|
||||
assert "HARM_CATEGORY_SEXUALLY_EXPLICIT" in categories
|
||||
assert "HARM_CATEGORY_DANGEROUS_CONTENT" in categories
|
||||
|
||||
thresholds = [s.threshold.value if hasattr(s.threshold, "value") else str(s.threshold) for s in config_arg.safety_settings]
|
||||
assert all(t == "BLOCK_NONE" for t in thresholds)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_no_safety_settings_omits_key():
|
||||
"""call() does NOT add safety_settings to GenerateContentConfig when none configured."""
|
||||
provider = _make_gemini_provider(safety_settings=None)
|
||||
|
||||
fake_response = MagicMock()
|
||||
fake_response.text = "hello"
|
||||
fake_response.candidates = [MagicMock(finish_reason="STOP")]
|
||||
fake_response.usage_metadata = MagicMock(prompt_token_count=5, candidates_token_count=2)
|
||||
|
||||
provider._client.aio.models.generate_content = AsyncMock(return_value=fake_response)
|
||||
|
||||
await provider.call(
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
scope="test",
|
||||
)
|
||||
|
||||
call_args = provider._client.aio.models.generate_content.call_args
|
||||
config_arg = call_args.kwargs.get("config")
|
||||
|
||||
# When no safety settings, config is either None or lacks safety_settings
|
||||
if config_arg is not None:
|
||||
assert not hasattr(config_arg, "safety_settings") or config_arg.safety_settings is None
|
||||
|
||||
|
||||
# ─── call_with_tools() applies safety settings ────────────────────────────────
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_with_tools_applies_safety_settings():
|
||||
"""call_with_tools() includes safety_settings in GenerateContentConfig."""
|
||||
provider = _make_gemini_provider(safety_settings=SAMPLE_SAFETY_SETTINGS)
|
||||
|
||||
# Build a fake tool-use response (no tool calls, just text)
|
||||
fake_part = MagicMock()
|
||||
fake_part.text = "answer"
|
||||
fake_part.function_call = None
|
||||
|
||||
fake_candidate = MagicMock()
|
||||
fake_candidate.content = MagicMock(parts=[fake_part])
|
||||
|
||||
fake_response = MagicMock()
|
||||
fake_response.candidates = [fake_candidate]
|
||||
fake_response.usage_metadata = MagicMock(prompt_token_count=5, candidates_token_count=3)
|
||||
|
||||
provider._client.aio.models.generate_content = AsyncMock(return_value=fake_response)
|
||||
|
||||
tools = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "test_tool",
|
||||
"description": "A test tool",
|
||||
"parameters": {"type": "object", "properties": {}, "required": []},
|
||||
},
|
||||
}
|
||||
]
|
||||
|
||||
await provider.call_with_tools(
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
tools=tools,
|
||||
scope="test",
|
||||
)
|
||||
|
||||
call_args = provider._client.aio.models.generate_content.call_args
|
||||
config_arg = call_args.kwargs.get("config")
|
||||
|
||||
assert config_arg is not None
|
||||
assert config_arg.safety_settings is not None
|
||||
categories = [s.category.value if hasattr(s.category, "value") else str(s.category) for s in config_arg.safety_settings]
|
||||
assert "HARM_CATEGORY_HARASSMENT" in categories
|
||||
|
||||
|
||||
# ─── with_config() override ───────────────────────────────────────────────────
|
||||
|
||||
|
||||
def _make_llm_provider(safety_settings=None):
|
||||
"""Return an LLMProvider (wrapping GeminiLLM) with a mocked genai.Client."""
|
||||
with patch("google.genai.Client") as mock_client_cls:
|
||||
mock_client_cls.return_value = MagicMock()
|
||||
from hindsight_api.engine.llm_wrapper import LLMProvider
|
||||
|
||||
provider = LLMProvider(
|
||||
provider="gemini",
|
||||
api_key="fake-api-key",
|
||||
base_url="",
|
||||
model="gemini-2.5-flash",
|
||||
gemini_safety_settings=safety_settings,
|
||||
)
|
||||
# Replace the underlying Gemini client with a fresh mock
|
||||
provider._provider_impl._client = MagicMock()
|
||||
return provider
|
||||
|
||||
|
||||
def _fake_response():
|
||||
r = MagicMock()
|
||||
r.text = "hello"
|
||||
r.candidates = [MagicMock(finish_reason="STOP")]
|
||||
r.usage_metadata = MagicMock(prompt_token_count=5, candidates_token_count=2)
|
||||
return r
|
||||
|
||||
|
||||
def _make_config(safety_settings):
|
||||
"""Return a minimal config-like object with llm_gemini_safety_settings."""
|
||||
cfg = MagicMock()
|
||||
cfg.llm_gemini_safety_settings = safety_settings
|
||||
return cfg
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_with_config_overrides_instance_settings():
|
||||
"""with_config() settings take precedence over the provider instance defaults."""
|
||||
instance_settings = [{"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_ONLY_HIGH"}]
|
||||
override_settings = [{"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE"}]
|
||||
|
||||
provider = _make_llm_provider(safety_settings=instance_settings)
|
||||
provider._provider_impl._client.aio.models.generate_content = AsyncMock(return_value=_fake_response())
|
||||
|
||||
configured = provider.with_config(_make_config(override_settings))
|
||||
await configured.call(messages=[{"role": "user", "content": "hi"}], scope="test")
|
||||
|
||||
config_arg = provider._provider_impl._client.aio.models.generate_content.call_args.kwargs.get("config")
|
||||
assert config_arg is not None
|
||||
categories = [s.category.value if hasattr(s.category, "value") else str(s.category) for s in config_arg.safety_settings]
|
||||
# Should use override_settings (HATE_SPEECH), not instance_settings (HARASSMENT)
|
||||
assert "HARM_CATEGORY_HATE_SPEECH" in categories
|
||||
assert "HARM_CATEGORY_HARASSMENT" not in categories
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_with_config_none_falls_back_to_instance():
|
||||
"""When with_config() supplies None, the instance default is used."""
|
||||
instance_settings = [{"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_NONE"}]
|
||||
|
||||
provider = _make_llm_provider(safety_settings=instance_settings)
|
||||
provider._provider_impl._client.aio.models.generate_content = AsyncMock(return_value=_fake_response())
|
||||
|
||||
configured = provider.with_config(_make_config(None))
|
||||
await configured.call(messages=[{"role": "user", "content": "hi"}], scope="test")
|
||||
|
||||
config_arg = provider._provider_impl._client.aio.models.generate_content.call_args.kwargs.get("config")
|
||||
assert config_arg is not None
|
||||
categories = [s.category.value if hasattr(s.category, "value") else str(s.category) for s in config_arg.safety_settings]
|
||||
assert "HARM_CATEGORY_HARASSMENT" in categories
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_with_config_resets_after_call():
|
||||
"""The ContextVar is properly reset after a with_config() call (no leakage)."""
|
||||
from hindsight_api.engine.providers.gemini_llm import _safety_settings_ctx
|
||||
|
||||
settings = [{"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_NONE"}]
|
||||
provider = _make_llm_provider(safety_settings=None)
|
||||
provider._provider_impl._client.aio.models.generate_content = AsyncMock(return_value=_fake_response())
|
||||
|
||||
before = _safety_settings_ctx.get()
|
||||
configured = provider.with_config(_make_config(settings))
|
||||
await configured.call(messages=[{"role": "user", "content": "hi"}], scope="test")
|
||||
after = _safety_settings_ctx.get()
|
||||
|
||||
assert after == before # ContextVar restored to its original value
|
||||
|
||||
|
||||
# ─── LLMProvider reads safety settings from config ────────────────────────────
|
||||
|
||||
|
||||
def test_llm_provider_reads_safety_settings_from_config():
|
||||
"""LLMProvider reads llm_gemini_safety_settings from global config for Gemini provider."""
|
||||
import json
|
||||
|
||||
from hindsight_api.config import ENV_LLM_GEMINI_SAFETY_SETTINGS, clear_config_cache
|
||||
|
||||
settings_json = json.dumps(SAMPLE_SAFETY_SETTINGS)
|
||||
env_overrides = {
|
||||
"HINDSIGHT_API_LLM_PROVIDER": "gemini",
|
||||
"HINDSIGHT_API_LLM_API_KEY": "fake-key",
|
||||
ENV_LLM_GEMINI_SAFETY_SETTINGS: settings_json,
|
||||
}
|
||||
|
||||
with patch.dict(os.environ, env_overrides, clear=False):
|
||||
clear_config_cache()
|
||||
with patch("google.genai.Client") as mock_client_cls:
|
||||
mock_client_cls.return_value = MagicMock()
|
||||
from hindsight_api.engine.llm_wrapper import LLMProvider
|
||||
|
||||
provider = LLMProvider(
|
||||
provider="gemini",
|
||||
api_key="fake-key",
|
||||
base_url="",
|
||||
model="gemini-2.5-flash",
|
||||
)
|
||||
|
||||
assert provider.gemini_safety_settings == SAMPLE_SAFETY_SETTINGS
|
||||
|
||||
clear_config_cache()
|
||||
@@ -1,171 +0,0 @@
|
||||
"""
|
||||
Tests for server-side filtering in the graph API endpoint.
|
||||
|
||||
Verifies that q (text search) and tags filters work correctly
|
||||
when passed as query parameters to GET /v1/default/banks/{bank_id}/graph.
|
||||
"""
|
||||
from datetime import datetime
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
|
||||
from hindsight_api.api import create_app
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def api_client(memory):
|
||||
"""Create an async test client for the FastAPI app."""
|
||||
app = create_app(memory, initialize_memory=False)
|
||||
transport = httpx.ASGITransport(app=app)
|
||||
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
|
||||
yield client
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def test_bank_id():
|
||||
"""Provide a unique bank ID for this test run."""
|
||||
return f"graph_filter_test_{datetime.now().timestamp()}"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_graph_no_filter_returns_all(api_client, test_bank_id):
|
||||
"""Without filters the graph endpoint returns all memories."""
|
||||
response = await api_client.post(
|
||||
f"/v1/default/banks/{test_bank_id}/memories",
|
||||
json={
|
||||
"items": [
|
||||
{"content": "Alice loves hiking in the mountains.", "tags": ["user_alice"]},
|
||||
{"content": "Bob enjoys swimming at the beach.", "tags": ["user_bob"]},
|
||||
]
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
|
||||
response = await api_client.get(f"/v1/default/banks/{test_bank_id}/graph")
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert "table_rows" in data
|
||||
texts = [row["text"] for row in data["table_rows"]]
|
||||
assert any("Alice" in t for t in texts)
|
||||
assert any("Bob" in t for t in texts)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_graph_q_filter_returns_matching(api_client, test_bank_id):
|
||||
"""The q parameter filters memories by text content."""
|
||||
response = await api_client.post(
|
||||
f"/v1/default/banks/{test_bank_id}/memories",
|
||||
json={
|
||||
"items": [
|
||||
{"content": "Alice loves hiking in the mountains."},
|
||||
{"content": "Bob enjoys swimming at the beach."},
|
||||
]
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
|
||||
response = await api_client.get(f"/v1/default/banks/{test_bank_id}/graph", params={"q": "Alice"})
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
texts = [row["text"] for row in data["table_rows"]]
|
||||
assert all("Alice" in t or "alice" in t.lower() for t in texts), (
|
||||
f"Expected only Alice memories, got: {texts}"
|
||||
)
|
||||
assert not any("Bob" in t for t in texts)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_graph_q_filter_case_insensitive(api_client, test_bank_id):
|
||||
"""The q filter is case-insensitive."""
|
||||
response = await api_client.post(
|
||||
f"/v1/default/banks/{test_bank_id}/memories",
|
||||
json={
|
||||
"items": [
|
||||
{"content": "Alice loves hiking in the mountains."},
|
||||
{"content": "Bob enjoys swimming at the beach."},
|
||||
]
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
|
||||
response = await api_client.get(f"/v1/default/banks/{test_bank_id}/graph", params={"q": "alice"})
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
texts = [row["text"] for row in data["table_rows"]]
|
||||
assert any("Alice" in t for t in texts)
|
||||
assert not any("Bob" in t for t in texts)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_graph_tags_filter_returns_matching(api_client, test_bank_id):
|
||||
"""The tags parameter filters memories to only those with matching tags."""
|
||||
response = await api_client.post(
|
||||
f"/v1/default/banks/{test_bank_id}/memories",
|
||||
json={
|
||||
"items": [
|
||||
{"content": "Alice loves hiking.", "tags": ["user_alice"]},
|
||||
{"content": "Bob enjoys swimming.", "tags": ["user_bob"]},
|
||||
]
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
|
||||
response = await api_client.get(
|
||||
f"/v1/default/banks/{test_bank_id}/graph",
|
||||
params={"tags": "user_alice", "tags_match": "all_strict"},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
texts = [row["text"] for row in data["table_rows"]]
|
||||
assert any("Alice" in t for t in texts)
|
||||
assert not any("Bob" in t for t in texts)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_graph_q_and_tags_filter_combined(api_client, test_bank_id):
|
||||
"""Combining q and tags filters applies both server-side."""
|
||||
response = await api_client.post(
|
||||
f"/v1/default/banks/{test_bank_id}/memories",
|
||||
json={
|
||||
"items": [
|
||||
{"content": "Alice loves hiking.", "tags": ["user_alice"]},
|
||||
{"content": "Alice also loves coding.", "tags": ["user_alice"]},
|
||||
{"content": "Bob enjoys swimming.", "tags": ["user_bob"]},
|
||||
]
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
|
||||
response = await api_client.get(
|
||||
f"/v1/default/banks/{test_bank_id}/graph",
|
||||
params={"q": "hiking", "tags": "user_alice", "tags_match": "all_strict"},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
texts = [row["text"] for row in data["table_rows"]]
|
||||
assert any("hiking" in t.lower() for t in texts)
|
||||
assert not any("coding" in t.lower() for t in texts)
|
||||
assert not any("Bob" in t for t in texts)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_graph_q_filter_empty_results(api_client, test_bank_id):
|
||||
"""The q filter returns empty results when no memory matches."""
|
||||
response = await api_client.post(
|
||||
f"/v1/default/banks/{test_bank_id}/memories",
|
||||
json={
|
||||
"items": [
|
||||
{"content": "Alice loves hiking."},
|
||||
]
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
|
||||
response = await api_client.get(
|
||||
f"/v1/default/banks/{test_bank_id}/graph",
|
||||
params={"q": "zzznomatchzzz"},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["table_rows"] == []
|
||||
@@ -1,501 +0,0 @@
|
||||
"""
|
||||
Tests for hierarchical configuration system.
|
||||
|
||||
Tests config resolution hierarchy (global → tenant → bank),
|
||||
key normalization, API endpoints, validation, and caching.
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
import pytest
|
||||
|
||||
from hindsight_api import MemoryEngine
|
||||
from hindsight_api.config import HindsightConfig, normalize_config_dict, normalize_config_key
|
||||
from hindsight_api.config_resolver import ConfigResolver
|
||||
from hindsight_api.extensions.tenant import TenantExtension
|
||||
from hindsight_api.models import RequestContext
|
||||
|
||||
|
||||
class MockTenantExtension(TenantExtension):
|
||||
"""Mock tenant extension for testing tenant-level config."""
|
||||
|
||||
def __init__(self, tenant_config: dict):
|
||||
self.tenant_config = tenant_config
|
||||
|
||||
async def authenticate(self, context):
|
||||
from hindsight_api.extensions.tenant import TenantContext
|
||||
|
||||
return TenantContext(schema_name="public")
|
||||
|
||||
async def list_tenants(self):
|
||||
from hindsight_api.extensions.tenant import Tenant
|
||||
|
||||
return [Tenant(schema="public")]
|
||||
|
||||
async def get_tenant_config(self, context):
|
||||
"""Return mock tenant config."""
|
||||
return self.tenant_config
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_config_key_normalization():
|
||||
"""Test that env var keys are normalized to Python field names."""
|
||||
# Test basic normalization
|
||||
assert normalize_config_key("HINDSIGHT_API_LLM_PROVIDER") == "llm_provider"
|
||||
assert normalize_config_key("HINDSIGHT_API_LLM_MODEL") == "llm_model"
|
||||
assert normalize_config_key("HINDSIGHT_API_RETAIN_LLM_PROVIDER") == "retain_llm_provider"
|
||||
|
||||
# Test already normalized keys
|
||||
assert normalize_config_key("llm_provider") == "llm_provider"
|
||||
assert normalize_config_key("llm_model") == "llm_model"
|
||||
|
||||
# Test dict normalization
|
||||
input_dict = {
|
||||
"HINDSIGHT_API_LLM_PROVIDER": "openai",
|
||||
"HINDSIGHT_API_LLM_MODEL": "gpt-4",
|
||||
"llm_base_url": "https://api.openai.com",
|
||||
}
|
||||
expected = {"llm_provider": "openai", "llm_model": "gpt-4", "llm_base_url": "https://api.openai.com"}
|
||||
assert normalize_config_dict(input_dict) == expected
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_hierarchical_fields_categorization():
|
||||
"""Test that fields are correctly categorized as configurable, credentials, or static."""
|
||||
configurable = HindsightConfig.get_configurable_fields()
|
||||
credentials = HindsightConfig.get_credential_fields()
|
||||
static = HindsightConfig.get_static_fields()
|
||||
|
||||
# Verify no overlap between configurable and credentials
|
||||
assert len(configurable & credentials) == 0, "Configurable fields should not include credentials"
|
||||
|
||||
# Verify configurable fields include behavioral settings (safe to modify)
|
||||
assert "retain_extraction_mode" in configurable
|
||||
assert "retain_mission" in configurable
|
||||
assert "retain_custom_instructions" in configurable
|
||||
assert "retain_chunk_size" in configurable
|
||||
assert "enable_observations" in configurable
|
||||
assert "consolidation_llm_batch_size" in configurable
|
||||
assert "consolidation_source_facts_max_tokens" in configurable
|
||||
assert "consolidation_source_facts_max_tokens_per_observation" in configurable
|
||||
assert "observations_mission" in configurable
|
||||
assert "reflect_mission" in configurable
|
||||
assert "disposition_skepticism" in configurable
|
||||
assert "disposition_literalism" in configurable
|
||||
assert "disposition_empathy" in configurable
|
||||
|
||||
# Verify entity labels fields are included
|
||||
assert "entities_allow_free_form" in configurable
|
||||
assert "entity_labels" in configurable
|
||||
|
||||
# Verify count is correct
|
||||
assert len(configurable) == 19
|
||||
|
||||
# Verify credential fields (NEVER exposed)
|
||||
assert "llm_api_key" in credentials
|
||||
assert "llm_base_url" in credentials
|
||||
assert "retain_llm_api_key" in credentials
|
||||
assert "reflect_llm_api_key" in credentials
|
||||
|
||||
# Verify static fields include server settings AND non-configurable LLM fields
|
||||
assert "database_url" in static
|
||||
assert "port" in static
|
||||
assert "host" in static
|
||||
assert "embeddings_provider" in static
|
||||
assert "reranker_provider" in static
|
||||
assert "worker_enabled" in static
|
||||
assert "llm_provider" in static # Not configurable (needs presets)
|
||||
assert "llm_model" in static # Not configurable (needs presets)
|
||||
assert "graph_retriever" in static # Performance tuning, not configurable
|
||||
assert "llm_max_concurrent" in static # Performance tuning, not configurable
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_config_hierarchy_resolution(memory, request_context):
|
||||
"""Test that config resolution follows global → tenant → bank hierarchy."""
|
||||
bank_id = "test-hierarchy-bank"
|
||||
|
||||
try:
|
||||
# Ensure bank exists in database
|
||||
await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
|
||||
# Set up mock tenant extension with tenant-level config (use configurable fields only)
|
||||
tenant_config = {"retain_chunk_size": 5000, "retain_extraction_mode": "tenant-mode"}
|
||||
mock_tenant = MockTenantExtension(tenant_config)
|
||||
|
||||
# Create config resolver with mock tenant extension
|
||||
resolver = ConfigResolver(pool=memory._pool, tenant_extension=mock_tenant)
|
||||
|
||||
# Test 1: Global config only (no overrides)
|
||||
context = RequestContext(api_key=None, api_key_id=None, tenant_id=None, internal=False)
|
||||
config = await resolver.get_bank_config(bank_id, context)
|
||||
|
||||
# Should have configurable fields from global config (NOT credentials or llm_provider/model)
|
||||
assert "retain_chunk_size" in config # Configurable field
|
||||
assert "llm_api_key" not in config # Credential - never exposed
|
||||
assert "llm_provider" not in config # Not configurable (needs presets)
|
||||
|
||||
# Test 2: Add tenant-level overrides
|
||||
config = await resolver.get_bank_config(bank_id, context)
|
||||
|
||||
# Should apply tenant overrides (only configurable fields)
|
||||
assert config["retain_chunk_size"] == 5000 # Tenant override
|
||||
assert config["retain_extraction_mode"] == "tenant-mode" # Tenant override
|
||||
|
||||
# Test 3: Add bank-level overrides (should take precedence)
|
||||
await resolver.update_bank_config(
|
||||
bank_id,
|
||||
{"retain_chunk_size": 2000, "retain_extraction_mode": "bank-mode"}, # Override tenant settings
|
||||
context,
|
||||
)
|
||||
|
||||
# Config should reflect changes immediately (no caching)
|
||||
config = await resolver.get_bank_config(bank_id, context)
|
||||
|
||||
# Bank overrides should take precedence over tenant
|
||||
assert config["retain_chunk_size"] == 2000 # Bank override wins
|
||||
assert config["retain_extraction_mode"] == "bank-mode" # Bank override wins
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_config_validation_rejects_static_fields(memory, request_context):
|
||||
"""Test that attempting to override static fields raises ValueError."""
|
||||
bank_id = "test-validation-bank"
|
||||
|
||||
try:
|
||||
# Ensure bank exists in database
|
||||
await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
|
||||
resolver = ConfigResolver(pool=memory._pool)
|
||||
|
||||
# Test 1: Configurable fields should work
|
||||
await resolver.update_bank_config(bank_id, {"retain_chunk_size": 4000, "retain_extraction_mode": "verbose"})
|
||||
|
||||
# Test 2: Static fields should raise ValueError
|
||||
with pytest.raises(ValueError, match="Cannot override static"):
|
||||
await resolver.update_bank_config(bank_id, {"port": 9000})
|
||||
|
||||
with pytest.raises(ValueError, match="Cannot override static"):
|
||||
await resolver.update_bank_config(bank_id, {"database_url": "postgresql://fake"})
|
||||
|
||||
with pytest.raises(ValueError, match="Cannot override static"):
|
||||
await resolver.update_bank_config(bank_id, {"embeddings_provider": "openai"})
|
||||
|
||||
# Test 3: Credential fields should raise ValueError
|
||||
with pytest.raises(ValueError, match="Cannot set credential fields"):
|
||||
await resolver.update_bank_config(bank_id, {"llm_api_key": "sk-fake"})
|
||||
|
||||
# Test 4: Non-configurable LLM fields should raise ValueError (need presets)
|
||||
with pytest.raises(ValueError, match="Cannot override static"):
|
||||
await resolver.update_bank_config(bank_id, {"llm_model": "gpt-4"})
|
||||
|
||||
# Test 5: Mix of configurable and static should fail
|
||||
with pytest.raises(ValueError, match="Cannot override static"):
|
||||
await resolver.update_bank_config(bank_id, {"retain_chunk_size": 4000, "port": 9000})
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_config_freshness_across_updates(memory, request_context):
|
||||
"""Test that config changes are immediately visible (no stale cache)."""
|
||||
bank1 = "freshness-test-1"
|
||||
|
||||
try:
|
||||
# Ensure bank exists in database
|
||||
await memory.get_bank_profile(bank1, request_context=request_context)
|
||||
|
||||
resolver = ConfigResolver(pool=memory._pool)
|
||||
|
||||
# Test 1: Initial config reflects global defaults
|
||||
config1 = await resolver.get_bank_config(bank1, None)
|
||||
initial_chunk_size = config1["retain_chunk_size"]
|
||||
|
||||
# Test 2: Update config
|
||||
await resolver.update_bank_config(bank1, {"retain_chunk_size": 4000})
|
||||
|
||||
# Test 3: Next call should see updated value immediately (no stale cache)
|
||||
config2 = await resolver.get_bank_config(bank1, None)
|
||||
assert config2["retain_chunk_size"] == 4000
|
||||
|
||||
# Test 4: Multiple updates are all immediately visible
|
||||
await resolver.update_bank_config(bank1, {"retain_chunk_size": 4500})
|
||||
config3 = await resolver.get_bank_config(bank1, None)
|
||||
assert config3["retain_chunk_size"] == 4500
|
||||
|
||||
# Test 5: Reset restores global defaults immediately
|
||||
await resolver.reset_bank_config(bank1)
|
||||
config4 = await resolver.get_bank_config(bank1, None)
|
||||
assert config4["retain_chunk_size"] == initial_chunk_size # Back to global default
|
||||
|
||||
# Test 6: Each call returns a fresh config dict (not a cached reference)
|
||||
config5 = await resolver.get_bank_config(bank1, None)
|
||||
config6 = await resolver.get_bank_config(bank1, None)
|
||||
assert config5 is not config6 # Different object instances
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank1, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_config_reset_to_defaults(memory, request_context):
|
||||
"""Test that resetting config removes all bank-specific overrides."""
|
||||
bank_id = "test-reset-bank"
|
||||
|
||||
try:
|
||||
# Ensure bank exists in database
|
||||
await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
|
||||
resolver = ConfigResolver(pool=memory._pool)
|
||||
|
||||
# Add bank-specific overrides
|
||||
await resolver.update_bank_config(
|
||||
bank_id,
|
||||
{
|
||||
"retain_chunk_size": 5500,
|
||||
"retain_extraction_mode": "custom",
|
||||
"retain_custom_instructions": "Custom instructions",
|
||||
},
|
||||
)
|
||||
|
||||
# Verify overrides applied
|
||||
config = await resolver.get_bank_config(bank_id, None)
|
||||
assert config["retain_chunk_size"] == 5500
|
||||
assert config["retain_extraction_mode"] == "custom"
|
||||
assert config["retain_custom_instructions"] == "Custom instructions"
|
||||
|
||||
# Reset to defaults
|
||||
await resolver.reset_bank_config(bank_id)
|
||||
|
||||
# Verify overrides removed (back to global defaults)
|
||||
config_reset = await resolver.get_bank_config(bank_id, None)
|
||||
assert config_reset["retain_chunk_size"] != 5500 # Should be global default
|
||||
assert config_reset["retain_extraction_mode"] != "custom" # Should be global default
|
||||
|
||||
# Verify bank_config is empty
|
||||
bank_overrides = await resolver._load_bank_config(bank_id)
|
||||
assert bank_overrides == {}
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_config_supports_both_key_formats(memory, request_context):
|
||||
"""Test that API accepts both env var and Python field formats."""
|
||||
bank_id = "test-key-format-bank"
|
||||
|
||||
try:
|
||||
# Ensure bank exists in database
|
||||
await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
|
||||
resolver = ConfigResolver(pool=memory._pool)
|
||||
|
||||
# Test 1: Python field format
|
||||
await resolver.update_bank_config(bank_id, {"retain_chunk_size": 7000})
|
||||
|
||||
config = await resolver.get_bank_config(bank_id, None)
|
||||
assert config["retain_chunk_size"] == 7000
|
||||
|
||||
# Test 2: Env var format (should be normalized)
|
||||
await resolver.update_bank_config(bank_id, {"HINDSIGHT_API_RETAIN_CHUNK_SIZE": 8000})
|
||||
|
||||
config = await resolver.get_bank_config(bank_id, None)
|
||||
assert config["retain_chunk_size"] == 8000
|
||||
|
||||
# Test 3: Mixed format in same request
|
||||
await resolver.update_bank_config(
|
||||
bank_id,
|
||||
{
|
||||
"retain_chunk_size": 9000, # Python format
|
||||
"HINDSIGHT_API_RETAIN_EXTRACTION_MODE": "verbose", # Env format
|
||||
},
|
||||
)
|
||||
|
||||
config = await resolver.get_bank_config(bank_id, None)
|
||||
assert config["retain_chunk_size"] == 9000
|
||||
assert config["retain_extraction_mode"] == "verbose"
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_config_only_configurable_fields_stored(memory, request_context):
|
||||
"""Test that only configurable fields are stored in bank config."""
|
||||
bank_id = "test-filter-bank"
|
||||
|
||||
try:
|
||||
# Ensure bank exists in database
|
||||
await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
|
||||
resolver = ConfigResolver(pool=memory._pool)
|
||||
|
||||
# Add valid configurable field
|
||||
await resolver.update_bank_config(bank_id, {"retain_chunk_size": 3500})
|
||||
|
||||
# Load bank config and verify only configurable fields present
|
||||
bank_overrides = await resolver._load_bank_config(bank_id)
|
||||
|
||||
for key in bank_overrides.keys():
|
||||
assert key in HindsightConfig.get_configurable_fields(), f"Non-configurable field {key} in bank config"
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_config_get_bank_config_no_static_or_credential_fields_leak(memory, request_context):
|
||||
"""
|
||||
SECURITY TEST: Verify get_bank_config() only returns configurable fields (no static/credentials).
|
||||
|
||||
This prevents leaking sensitive system configuration like database URLs,
|
||||
API keys, LLM providers/models, worker counts, etc. when retrieving bank configuration.
|
||||
"""
|
||||
bank_id = "test-security-bank"
|
||||
|
||||
try:
|
||||
# Ensure bank exists in database
|
||||
await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
|
||||
resolver = ConfigResolver(pool=memory._pool)
|
||||
|
||||
# Get bank config
|
||||
config = await resolver.get_bank_config(bank_id, None)
|
||||
|
||||
# Get field categorizations
|
||||
configurable_fields = HindsightConfig.get_configurable_fields()
|
||||
credential_fields = HindsightConfig.get_credential_fields()
|
||||
static_fields = HindsightConfig.get_static_fields()
|
||||
|
||||
# SECURITY: Verify ONLY configurable fields are returned (NO static, NO credentials)
|
||||
for key in config.keys():
|
||||
assert key in configurable_fields, (
|
||||
f"SECURITY VIOLATION: Non-configurable field '{key}' returned by get_bank_config(). "
|
||||
f"Only configurable fields should be returned to prevent leaking system config."
|
||||
)
|
||||
assert key not in credential_fields, (
|
||||
f"SECURITY VIOLATION: Credential field '{key}' returned by get_bank_config(). "
|
||||
f"Credentials must NEVER be exposed via API."
|
||||
)
|
||||
|
||||
# SECURITY: Verify specific sensitive fields are NOT present
|
||||
sensitive_fields = [
|
||||
"database_url", "api_port", "host", "worker_count", # Infrastructure
|
||||
"llm_api_key", "llm_base_url", # Credentials
|
||||
"retain_llm_api_key", "reflect_llm_api_key", # More credentials
|
||||
"llm_provider", "llm_model", # Not configurable (need presets)
|
||||
]
|
||||
for field in sensitive_fields:
|
||||
assert field not in config, (
|
||||
f"SECURITY VIOLATION: Sensitive field '{field}' returned by get_bank_config(). "
|
||||
f"Must not be exposed via bank config API."
|
||||
)
|
||||
|
||||
# Verify we have the expected configurable fields (small set)
|
||||
expected_configurable = ["retain_chunk_size", "retain_extraction_mode", "enable_observations"]
|
||||
for field in expected_configurable:
|
||||
assert field in config, f"Expected configurable field '{field}' missing from config"
|
||||
|
||||
# Should have a small number of configurable fields (not hundreds)
|
||||
assert len(config) < 20, f"Too many fields returned: {len(config)}"
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_config_permissions_system(memory, request_context):
|
||||
"""
|
||||
Test that tenant extension can control which fields banks are allowed to modify.
|
||||
|
||||
Tests get_allowed_config_fields() permission system.
|
||||
"""
|
||||
bank_id = "test-permissions-bank"
|
||||
|
||||
class PermissionTenantExtension(TenantExtension):
|
||||
"""Mock tenant extension with configurable permissions."""
|
||||
|
||||
def __init__(self, allowed_fields: set[str] | None):
|
||||
self.allowed_fields = allowed_fields
|
||||
|
||||
async def authenticate(self, context):
|
||||
from hindsight_api.extensions.tenant import TenantContext
|
||||
|
||||
return TenantContext(schema_name="public")
|
||||
|
||||
async def list_tenants(self):
|
||||
from hindsight_api.extensions.tenant import Tenant
|
||||
|
||||
return [Tenant(schema="public")]
|
||||
|
||||
async def get_allowed_config_fields(self, context, bank_id):
|
||||
"""Return configured allowed fields."""
|
||||
return self.allowed_fields
|
||||
|
||||
try:
|
||||
# Ensure bank exists in database
|
||||
await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
|
||||
# Test 1: None = allow all configurable fields
|
||||
extension = PermissionTenantExtension(allowed_fields=None)
|
||||
resolver = ConfigResolver(pool=memory._pool, tenant_extension=extension)
|
||||
|
||||
await resolver.update_bank_config(
|
||||
bank_id, {"retain_chunk_size": 4000, "retain_extraction_mode": "verbose"}, request_context
|
||||
)
|
||||
config = await resolver.get_bank_config(bank_id, request_context)
|
||||
assert config["retain_chunk_size"] == 4000
|
||||
assert config["retain_extraction_mode"] == "verbose"
|
||||
|
||||
# Reset for next test
|
||||
await resolver.reset_bank_config(bank_id)
|
||||
|
||||
# Test 2: Specific set = only those fields allowed
|
||||
extension = PermissionTenantExtension(allowed_fields={"retain_chunk_size"})
|
||||
resolver = ConfigResolver(pool=memory._pool, tenant_extension=extension)
|
||||
|
||||
# Should allow retain_chunk_size
|
||||
await resolver.update_bank_config(bank_id, {"retain_chunk_size": 5000}, request_context)
|
||||
config = await resolver.get_bank_config(bank_id, request_context)
|
||||
assert config["retain_chunk_size"] == 5000
|
||||
|
||||
# Should reject retain_extraction_mode (not in allowed list)
|
||||
with pytest.raises(ValueError, match="Not allowed to modify fields"):
|
||||
await resolver.update_bank_config(bank_id, {"retain_extraction_mode": "verbose"}, request_context)
|
||||
|
||||
# Should reject mix of allowed and disallowed
|
||||
with pytest.raises(ValueError, match="Not allowed to modify fields"):
|
||||
await resolver.update_bank_config(
|
||||
bank_id, {"retain_chunk_size": 6000, "retain_extraction_mode": "verbose"}, request_context
|
||||
)
|
||||
|
||||
# Reset for next test
|
||||
await resolver.reset_bank_config(bank_id)
|
||||
|
||||
# Test 3: Empty set = no modifications allowed (read-only)
|
||||
extension = PermissionTenantExtension(allowed_fields=set())
|
||||
resolver = ConfigResolver(pool=memory._pool, tenant_extension=extension)
|
||||
|
||||
with pytest.raises(ValueError, match="Not allowed to modify fields"):
|
||||
await resolver.update_bank_config(bank_id, {"retain_chunk_size": 7000}, request_context)
|
||||
|
||||
# Test 4: get_bank_config should filter response based on permissions
|
||||
extension = PermissionTenantExtension(allowed_fields={"retain_chunk_size", "enable_observations"})
|
||||
resolver = ConfigResolver(pool=memory._pool, tenant_extension=extension)
|
||||
|
||||
config = await resolver.get_bank_config(bank_id, request_context)
|
||||
|
||||
# Should only return allowed fields
|
||||
assert "retain_chunk_size" in config
|
||||
assert "enable_observations" in config
|
||||
# Other configurable fields should be filtered out
|
||||
assert "retain_extraction_mode" not in config
|
||||
assert "retain_custom_instructions" not in config
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user