Compare commits
3
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
7972fd3906 | ||
|
|
86b698460e | ||
|
|
6f9cef674b |
@@ -1,15 +0,0 @@
|
||||
{
|
||||
"$schema": "https://anthropic.com/claude-code/marketplace.schema.json",
|
||||
"name": "hindsight",
|
||||
"description": "Official Hindsight integrations for Claude Code",
|
||||
"owner": {
|
||||
"name": "vectorize-io"
|
||||
},
|
||||
"plugins": [
|
||||
{
|
||||
"name": "hindsight-memory",
|
||||
"description": "Automatic long-term memory for Claude Code via Hindsight",
|
||||
"source": "./hindsight-integrations/claude-code"
|
||||
}
|
||||
]
|
||||
}
|
||||
+2
-14
@@ -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, volcano
|
||||
# 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
|
||||
@@ -44,15 +39,8 @@ HINDSIGHT_API_LOG_LEVEL=info
|
||||
|
||||
# Database (Optional - uses embedded pg0 by default)
|
||||
# HINDSIGHT_API_DATABASE_URL=postgresql://user:pass@host:5432/db
|
||||
# HINDSIGHT_API_MIGRATION_DATABASE_URL= # Direct PostgreSQL URL for migrations (bypasses PgBouncer). Falls back to DATABASE_URL.
|
||||
# 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,111 +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
|
||||
elif [ -f "hindsight-integrations/${{ steps.info.outputs.integration }}/package.json" ]; then
|
||||
echo "type=typescript" >> $GITHUB_OUTPUT
|
||||
else
|
||||
echo "type=plugin" >> $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@v6
|
||||
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) ────────────────────
|
||||
|
||||
# ── Plugin integrations (claude-code) — no package to publish ───────────
|
||||
|
||||
- name: Plugin release
|
||||
if: steps.type.outputs.type == 'plugin'
|
||||
run: |
|
||||
echo "Plugin integration ${{ steps.info.outputs.integration }} v${{ steps.info.outputs.version }} — no package to publish."
|
||||
echo "Users install via: claude plugin marketplace add vectorize-io/hindsight --sparse hindsight-integrations"
|
||||
|
||||
# ── 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
-1269
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
|
||||
@@ -251,7 +244,7 @@ Fields must be categorized as either **hierarchical** (can be overridden per-ten
|
||||
|
||||
#### Adding a New Configuration Field
|
||||
|
||||
1. **config.py** (`hindsight-api-slim/hindsight_api/config.py`):
|
||||
1. **config.py** (`hindsight-api/hindsight_api/config.py`):
|
||||
- Add `ENV_*` constant for the environment variable name (e.g., `ENV_MY_SETTING = "HINDSIGHT_API_MY_SETTING"`)
|
||||
- Add `DEFAULT_*` constant for the default value
|
||||
- Add field to `HindsightConfig` dataclass with type annotation
|
||||
@@ -268,7 +261,7 @@ Fields must be categorized as either **hierarchical** (can be overridden per-ten
|
||||
# Static field - just don't add to _HIERARCHICAL_FIELDS
|
||||
```
|
||||
|
||||
2. **main.py** (`hindsight-api-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**:
|
||||
@@ -308,19 +301,19 @@ 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)
|
||||
- `HINDSIGHT_API_ENABLE_BANK_CONFIG_API`: Enable per-bank config API (default: false, disabled for security)
|
||||
|
||||
@@ -7,12 +7,10 @@
|
||||
[](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)
|
||||
[](https://gitcgr.com/vectorize-io/hindsight)
|
||||

|
||||

|
||||
<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>
|
||||
|
||||
---
|
||||
@@ -38,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.
|
||||
|
||||
@@ -71,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).
|
||||
|
||||
|
||||
|
||||
@@ -183,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")
|
||||
@@ -309,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,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:
|
||||
@@ -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 .
|
||||
@@ -167,11 +170,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 +321,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)
|
||||
|
||||
@@ -1,28 +1,6 @@
|
||||
#!/bin/bash
|
||||
set -e
|
||||
|
||||
# =============================================================================
|
||||
# Embedded pg0 data integrity check (#675)
|
||||
#
|
||||
# When using embedded pg0, check if the data directory has existing PostgreSQL
|
||||
# data before starting. If the directory exists but appears empty/corrupt
|
||||
# (e.g., missing PG_VERSION file), log a warning. This helps diagnose data
|
||||
# loss scenarios where a container restart caused the data directory to be
|
||||
# wiped despite a volume mount being present.
|
||||
# =============================================================================
|
||||
PG0_DATA_DIR="${HOME}/.pg0"
|
||||
if [ -d "$PG0_DATA_DIR" ]; then
|
||||
# Look for actual PostgreSQL data directories (pg0 creates subdirs per instance)
|
||||
if compgen -G "$PG0_DATA_DIR"/*/PG_VERSION > /dev/null 2>&1; then
|
||||
echo "✅ Existing pg0 data directory detected at $PG0_DATA_DIR"
|
||||
elif [ "$(ls -A "$PG0_DATA_DIR" 2>/dev/null)" ]; then
|
||||
echo "⚠️ WARNING: pg0 data directory exists at $PG0_DATA_DIR but no PG_VERSION found."
|
||||
echo " This may indicate data corruption or an incomplete previous shutdown."
|
||||
echo " If you see all migrations running from scratch after this, your data may have been lost."
|
||||
echo " See: https://github.com/vectorize-io/hindsight/issues/675"
|
||||
fi
|
||||
fi
|
||||
|
||||
# Service flags (default to true if not set)
|
||||
ENABLE_API="${HINDSIGHT_ENABLE_API:-true}"
|
||||
ENABLE_CP="${HINDSIGHT_ENABLE_CP:-true}"
|
||||
@@ -93,95 +71,24 @@ if [ "${HINDSIGHT_WAIT_FOR_DEPS:-false}" = "true" ]; then
|
||||
done
|
||||
fi
|
||||
|
||||
# =============================================================================
|
||||
# Graceful shutdown handler (#675)
|
||||
#
|
||||
# Docker sends SIGTERM on `docker stop`/`docker restart`. Without a trap, child
|
||||
# processes (hindsight-api + pg0, control-plane) are killed abruptly. For the
|
||||
# embedded pg0 database this can cause data loss when the data directory is on
|
||||
# a Docker volume that gets remounted after restart.
|
||||
#
|
||||
# The trap forwards SIGTERM to all tracked child PIDs so that:
|
||||
# - hindsight-api receives the signal and can run its shutdown hooks
|
||||
# - pg0 gets a clean PostgreSQL shutdown (checkpoint + WAL flush)
|
||||
# - The control-plane Node.js process exits cleanly
|
||||
# =============================================================================
|
||||
# Guard against concurrent cleanup (e.g., child crash + SIGTERM arriving together)
|
||||
SHUTTING_DOWN=false
|
||||
|
||||
cleanup() {
|
||||
if $SHUTTING_DOWN; then return; fi
|
||||
SHUTTING_DOWN=true
|
||||
|
||||
echo ""
|
||||
echo "🛑 Received shutdown signal, stopping services gracefully..."
|
||||
for pid in "${PIDS[@]}"; do
|
||||
if kill -0 "$pid" 2>/dev/null; then
|
||||
kill -TERM "$pid" 2>/dev/null
|
||||
fi
|
||||
done
|
||||
# Give processes time to shut down cleanly (pg0 needs to flush WAL).
|
||||
# NOTE: Docker's default stop_grace_period is 10s. If you use the default,
|
||||
# either set stop_grace_period: 30s in your compose file / docker stop -t 30,
|
||||
# or Docker will SIGKILL the container before this timeout expires.
|
||||
local timeout=30
|
||||
for ((i=1; i<=timeout; i++)); do
|
||||
local all_stopped=true
|
||||
for pid in "${PIDS[@]}"; do
|
||||
if kill -0 "$pid" 2>/dev/null; then
|
||||
all_stopped=false
|
||||
break
|
||||
fi
|
||||
done
|
||||
if $all_stopped; then
|
||||
echo "✅ All services stopped cleanly"
|
||||
exit 0
|
||||
fi
|
||||
sleep 1
|
||||
done
|
||||
# Force kill if still running after timeout
|
||||
echo "⚠️ Timeout reached, forcing shutdown..."
|
||||
for pid in "${PIDS[@]}"; do
|
||||
if kill -0 "$pid" 2>/dev/null; then
|
||||
kill -9 "$pid" 2>/dev/null
|
||||
fi
|
||||
done
|
||||
exit 1
|
||||
}
|
||||
trap cleanup SIGTERM SIGINT
|
||||
|
||||
# Track PIDs for wait
|
||||
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
|
||||
@@ -190,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
|
||||
@@ -204,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"
|
||||
@@ -217,21 +123,8 @@ if [ ${#PIDS[@]} -eq 0 ]; then
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# Wait for any process to exit (use wait -n with trap-safe loop)
|
||||
while true; do
|
||||
# wait -n returns when any child exits; it also returns on signal delivery
|
||||
# (the trap handler will run and exit, so this loop is just for robustness).
|
||||
# `&& true` prevents `set -e` from killing the script when wait -n returns
|
||||
# non-zero (child exited with error or no backgrounded children remain).
|
||||
wait -n && true
|
||||
# Check if any tracked PID has exited
|
||||
for pid in "${PIDS[@]}"; do
|
||||
if ! kill -0 "$pid" 2>/dev/null; then
|
||||
wait "$pid" 2>/dev/null
|
||||
exit_code=$?
|
||||
echo "⚠️ Service (PID $pid) exited with code $exit_code"
|
||||
# Trigger cleanup for remaining services
|
||||
cleanup
|
||||
fi
|
||||
done
|
||||
done
|
||||
# Wait for any process to exit
|
||||
wait -n
|
||||
|
||||
# Exit with status of first exited process
|
||||
exit $?
|
||||
|
||||
+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.20
|
||||
appVersion: "0.4.20"
|
||||
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.20"
|
||||
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")
|
||||
-32
@@ -1,32 +0,0 @@
|
||||
"""add content_hash to chunks table for delta retain
|
||||
|
||||
Revision ID: b3c4d5e6f7a8
|
||||
Revises: a3b4c5d6e7f8
|
||||
Create Date: 2026-03-25
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
revision: str = "b3c4d5e6f7a8"
|
||||
down_revision: str | Sequence[str] | None = "a3b4c5d6e7f8"
|
||||
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()
|
||||
# Add content_hash column to chunks table for delta comparison
|
||||
op.execute(f"ALTER TABLE {schema}chunks ADD COLUMN IF NOT EXISTS content_hash TEXT")
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
op.execute(f"ALTER TABLE {schema}chunks DROP COLUMN IF EXISTS content_hash")
|
||||
-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
|
||||
-59
@@ -1,59 +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
|
||||
|
||||
import sqlalchemy as sa
|
||||
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 most PostgreSQL installations as a contrib module.
|
||||
# It enables fast similarity lookups via GIN indexes, used for entity name matching.
|
||||
# On managed services (e.g. Azure Flexible Server), the extension may not be
|
||||
# available or may require manual enablement. We gracefully skip the index
|
||||
# creation if the extension cannot be loaded — the entity resolver will
|
||||
# auto-detect and fall back to the "full" lookup strategy at runtime. See #626.
|
||||
conn = op.get_bind()
|
||||
try:
|
||||
conn.execute(sa.text("CREATE EXTENSION IF NOT EXISTS pg_trgm"))
|
||||
except Exception:
|
||||
# Extension not available (managed Postgres, insufficient privileges, etc.)
|
||||
# Roll back the failed statement and skip index creation.
|
||||
conn.execute(sa.text("ROLLBACK"))
|
||||
conn.execute(sa.text("BEGIN"))
|
||||
return
|
||||
|
||||
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
|
||||
@@ -1,61 +0,0 @@
|
||||
"""Add audit_log table for feature usage tracking.
|
||||
|
||||
Merge migration that combines the two existing heads (a3b4c5d6e7f8 + c8e5f2a3b4d1).
|
||||
|
||||
Stores raw request/response as JSONB for expandability without future migrations.
|
||||
The metadata JSONB column allows adding arbitrary fields in the future.
|
||||
|
||||
Revision ID: c2d3e4f5g6h7
|
||||
Revises: a3b4c5d6e7f8, c8e5f2a3b4d1
|
||||
Create Date: 2026-03-26
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
revision: str = "c2d3e4f5g6h7"
|
||||
down_revision: str | Sequence[str] | None = ("a3b4c5d6e7f8", "c8e5f2a3b4d1")
|
||||
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"""
|
||||
CREATE TABLE IF NOT EXISTS {schema}audit_log (
|
||||
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
||||
action TEXT NOT NULL,
|
||||
transport TEXT NOT NULL,
|
||||
bank_id TEXT,
|
||||
started_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
|
||||
ended_at TIMESTAMPTZ,
|
||||
request JSONB,
|
||||
response JSONB,
|
||||
metadata JSONB DEFAULT '{{}}'::jsonb
|
||||
)
|
||||
"""
|
||||
)
|
||||
|
||||
op.execute(
|
||||
f"CREATE INDEX IF NOT EXISTS idx_audit_log_action_started ON {schema}audit_log (action, started_at DESC)"
|
||||
)
|
||||
op.execute(f"CREATE INDEX IF NOT EXISTS idx_audit_log_bank_started ON {schema}audit_log (bank_id, started_at DESC)")
|
||||
op.execute(f"CREATE INDEX IF NOT EXISTS idx_audit_log_started ON {schema}audit_log (started_at DESC)")
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema}idx_audit_log_started")
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema}idx_audit_log_bank_started")
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema}idx_audit_log_action_started")
|
||||
op.execute(f"DROP TABLE IF EXISTS {schema}audit_log")
|
||||
-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")
|
||||
-49
@@ -1,49 +0,0 @@
|
||||
"""Add bank_id column to memory_links for direct filtering
|
||||
|
||||
The stats endpoint JOINs memory_links to memory_units just to filter by
|
||||
bank_id. With millions of links this takes 18+ seconds. Adding bank_id
|
||||
directly to memory_links lets Postgres push the filter down before the JOIN.
|
||||
|
||||
Revision ID: c5d6e7f8a9b0
|
||||
Revises: b3c4d5e6f7a8
|
||||
Create Date: 2026-03-26
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
revision: str = "c5d6e7f8a9b0"
|
||||
down_revision: str | Sequence[str] | None = "b3c4d5e6f7a8"
|
||||
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()
|
||||
|
||||
# 1. Add nullable column
|
||||
op.execute(f"ALTER TABLE {schema}memory_links ADD COLUMN IF NOT EXISTS bank_id TEXT")
|
||||
|
||||
# 2. Backfill from memory_units
|
||||
op.execute(f"""
|
||||
UPDATE {schema}memory_links ml
|
||||
SET bank_id = mu.bank_id
|
||||
FROM {schema}memory_units mu
|
||||
WHERE ml.from_unit_id = mu.id
|
||||
AND ml.bank_id IS NULL
|
||||
""")
|
||||
|
||||
# 3. Set NOT NULL
|
||||
op.execute(f"ALTER TABLE {schema}memory_links ALTER COLUMN bank_id SET NOT NULL")
|
||||
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
op.execute(f"ALTER TABLE {schema}memory_links DROP COLUMN IF EXISTS bank_id")
|
||||
-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")
|
||||
-57
@@ -1,57 +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.
|
||||
"""
|
||||
from alembic import context
|
||||
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
schema_prefix = f'"{schema}".' if schema else ""
|
||||
# Use raw SQL with IF EXISTS so this is safe on schemas where the FK was
|
||||
# already dropped or never existed under this name.
|
||||
op.execute(f"ALTER TABLE {schema_prefix}memory_units DROP CONSTRAINT IF EXISTS memory_units_chunk_fkey")
|
||||
# Use a DO block so the ADD is also idempotent: if the FK already exists (e.g.
|
||||
# the schema was provisioned after the base migration already added it) the
|
||||
# duplicate_object exception is swallowed rather than failing the migration.
|
||||
op.execute(
|
||||
f"""
|
||||
DO $$ BEGIN
|
||||
ALTER TABLE {schema_prefix}memory_units
|
||||
ADD CONSTRAINT memory_units_chunk_fkey
|
||||
FOREIGN KEY (chunk_id)
|
||||
REFERENCES {schema_prefix}chunks (chunk_id)
|
||||
ON DELETE CASCADE;
|
||||
EXCEPTION
|
||||
WHEN duplicate_object THEN NULL;
|
||||
END $$;
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
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
|
||||
-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,209 +0,0 @@
|
||||
"""Audit logging for feature usage tracking.
|
||||
|
||||
Provides fire-and-forget audit logging of all mutating and core operations
|
||||
(retain, recall, reflect, bank CRUD, etc.) across HTTP, MCP, and system transports.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import uuid
|
||||
from collections.abc import Callable
|
||||
from contextlib import asynccontextmanager
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any
|
||||
|
||||
import asyncpg
|
||||
|
||||
from ..engine.db_utils import acquire_with_retry
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class AuditEntry:
|
||||
"""A single audit log entry."""
|
||||
|
||||
action: str
|
||||
transport: str # "http", "mcp", "system"
|
||||
bank_id: str | None = None
|
||||
started_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
|
||||
ended_at: datetime | None = None
|
||||
request: dict[str, Any] | None = None
|
||||
response: dict[str, Any] | None = None
|
||||
metadata: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
|
||||
def _json_default(obj: Any) -> str:
|
||||
"""JSON serializer for objects not serializable by default."""
|
||||
if isinstance(obj, datetime):
|
||||
return obj.isoformat()
|
||||
if isinstance(obj, uuid.UUID):
|
||||
return str(obj)
|
||||
if isinstance(obj, bytes):
|
||||
return "<bytes>"
|
||||
if isinstance(obj, set):
|
||||
return list(obj)
|
||||
return str(obj)
|
||||
|
||||
|
||||
def _safe_json(data: Any) -> str | None:
|
||||
"""Serialize data to JSON string, returning None on failure."""
|
||||
if data is None:
|
||||
return None
|
||||
try:
|
||||
return json.dumps(data, default=_json_default)
|
||||
except Exception:
|
||||
logger.debug("Failed to serialize audit data", exc_info=True)
|
||||
return None
|
||||
|
||||
|
||||
_SWEEP_INTERVAL_SECONDS = 3600 # Run retention sweep every hour
|
||||
|
||||
|
||||
class AuditLogger:
|
||||
"""Fire-and-forget audit log writer with optional retention sweep."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
pool_getter: Callable[[], asyncpg.Pool | None],
|
||||
schema_getter: Callable[[], str],
|
||||
enabled: bool,
|
||||
allowed_actions: list[str],
|
||||
retention_days: int = -1,
|
||||
) -> None:
|
||||
self._pool_getter = pool_getter
|
||||
self._schema_getter = schema_getter
|
||||
self._enabled = enabled
|
||||
self._allowed_actions: frozenset[str] | None = frozenset(allowed_actions) if allowed_actions else None
|
||||
self._retention_days = retention_days
|
||||
self._sweep_task: asyncio.Task | None = None
|
||||
|
||||
def is_enabled(self, action: str) -> bool:
|
||||
"""Check if audit logging is enabled for this action."""
|
||||
if not self._enabled:
|
||||
return False
|
||||
if self._allowed_actions is not None:
|
||||
return action in self._allowed_actions
|
||||
return True
|
||||
|
||||
def log_fire_and_forget(self, entry: AuditEntry) -> None:
|
||||
"""Schedule an audit write as a background task."""
|
||||
if not self.is_enabled(entry.action):
|
||||
return
|
||||
try:
|
||||
asyncio.create_task(self._safe_log(entry))
|
||||
except RuntimeError:
|
||||
# No running event loop (e.g. during shutdown)
|
||||
logger.debug("Cannot schedule audit log write: no running event loop")
|
||||
|
||||
async def _safe_log(self, entry: AuditEntry) -> None:
|
||||
"""Write audit entry to DB. Errors are logged, never raised."""
|
||||
pool = self._pool_getter()
|
||||
if pool is None:
|
||||
logger.debug("Audit log skipped: pool not available")
|
||||
return
|
||||
try:
|
||||
schema = self._schema_getter()
|
||||
table = f"{schema}.audit_log"
|
||||
async with acquire_with_retry(pool, max_retries=1) as conn:
|
||||
await conn.execute(
|
||||
f"""
|
||||
INSERT INTO {table}
|
||||
(id, action, transport, bank_id, started_at, ended_at, request, response, metadata)
|
||||
VALUES
|
||||
($1, $2, $3, $4, $5, $6, $7::jsonb, $8::jsonb, $9::jsonb)
|
||||
""",
|
||||
uuid.uuid4(),
|
||||
entry.action,
|
||||
entry.transport,
|
||||
entry.bank_id,
|
||||
entry.started_at,
|
||||
entry.ended_at,
|
||||
_safe_json(entry.request),
|
||||
_safe_json(entry.response),
|
||||
_safe_json(entry.metadata) or "{}",
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"Audit log write failed for action={entry.action}: {e}")
|
||||
|
||||
def start_retention_sweep(self) -> None:
|
||||
"""Start the periodic retention sweep if retention is configured."""
|
||||
if self._retention_days <= 0 or not self._enabled:
|
||||
return
|
||||
try:
|
||||
self._sweep_task = asyncio.create_task(self._sweep_loop())
|
||||
except RuntimeError:
|
||||
logger.debug("Cannot start retention sweep: no running event loop")
|
||||
|
||||
async def stop_retention_sweep(self) -> None:
|
||||
"""Stop the periodic retention sweep."""
|
||||
if self._sweep_task and not self._sweep_task.done():
|
||||
self._sweep_task.cancel()
|
||||
try:
|
||||
await self._sweep_task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
self._sweep_task = None
|
||||
|
||||
async def _sweep_loop(self) -> None:
|
||||
"""Periodically delete audit log entries older than retention_days."""
|
||||
while True:
|
||||
await self._run_sweep()
|
||||
await asyncio.sleep(_SWEEP_INTERVAL_SECONDS)
|
||||
|
||||
async def _run_sweep(self) -> None:
|
||||
"""Delete expired audit log entries. Concurrent-safe via row-level deletes."""
|
||||
pool = self._pool_getter()
|
||||
if pool is None:
|
||||
return
|
||||
try:
|
||||
schema = self._schema_getter()
|
||||
table = f"{schema}.audit_log"
|
||||
async with acquire_with_retry(pool, max_retries=1) as conn:
|
||||
result = await conn.execute(
|
||||
f"DELETE FROM {table} WHERE started_at < NOW() - INTERVAL '{self._retention_days} days'"
|
||||
)
|
||||
if result and result != "DELETE 0":
|
||||
logger.info(f"Audit log retention sweep: {result}")
|
||||
except Exception as e:
|
||||
logger.warning(f"Audit log retention sweep failed: {e}")
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def audit_context(
|
||||
audit_logger: AuditLogger | None,
|
||||
action: str,
|
||||
transport: str,
|
||||
bank_id: str | None = None,
|
||||
request: dict[str, Any] | None = None,
|
||||
metadata: dict[str, Any] | None = None,
|
||||
):
|
||||
"""Async context manager that times the operation and writes audit on exit.
|
||||
|
||||
Usage:
|
||||
async with audit_context(logger, "retain", "http", bank_id, request_dict) as entry:
|
||||
result = await do_work()
|
||||
entry.response = result_dict
|
||||
"""
|
||||
if audit_logger is None or not audit_logger.is_enabled(action):
|
||||
entry = AuditEntry(action=action, transport=transport, bank_id=bank_id)
|
||||
yield entry
|
||||
return
|
||||
|
||||
entry = AuditEntry(
|
||||
action=action,
|
||||
transport=transport,
|
||||
bank_id=bank_id,
|
||||
started_at=datetime.now(timezone.utc),
|
||||
request=request,
|
||||
metadata=metadata or {},
|
||||
)
|
||||
try:
|
||||
yield entry
|
||||
finally:
|
||||
entry.ended_at = datetime.now(timezone.utc)
|
||||
audit_logger.log_fire_and_forget(entry)
|
||||
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,380 +0,0 @@
|
||||
"""
|
||||
LiteLLM LLM provider for universal model support.
|
||||
|
||||
This provider enables using 100+ LLM providers via the LiteLLM SDK, including:
|
||||
- AWS Bedrock (bedrock/anthropic.claude-3-5-sonnet-...)
|
||||
- Azure OpenAI (azure/gpt-4o)
|
||||
- Together AI (together_ai/meta-llama/...)
|
||||
- Any other LiteLLM-supported provider
|
||||
|
||||
Uses litellm.acompletion() for async chat completions.
|
||||
Authentication for cloud providers (e.g., AWS Bedrock via boto3 credential chain)
|
||||
is handled automatically by LiteLLM.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
from hindsight_api.engine.llm_interface import LLMInterface, OutputTooLongError
|
||||
from hindsight_api.engine.response_models import LLMToolCall, LLMToolCallResult, TokenUsage
|
||||
from hindsight_api.metrics import get_metrics_collector
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class LiteLLMLLM(LLMInterface):
|
||||
"""
|
||||
LLM provider using the LiteLLM SDK for universal model support.
|
||||
|
||||
Supports any model accessible via litellm.acompletion(), including AWS Bedrock,
|
||||
Azure OpenAI, Together AI, Fireworks AI, and more.
|
||||
|
||||
Model names follow LiteLLM conventions with provider prefixes:
|
||||
- bedrock/anthropic.claude-3-5-sonnet-20241022-v2:0
|
||||
- azure/gpt-4o
|
||||
- together_ai/meta-llama/Llama-3-70b-chat-hf
|
||||
- fireworks_ai/accounts/fireworks/models/llama-v3p1-70b-instruct
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
provider: str,
|
||||
api_key: str,
|
||||
base_url: str,
|
||||
model: str,
|
||||
reasoning_effort: str = "low",
|
||||
timeout: float = 300.0,
|
||||
**kwargs: Any,
|
||||
):
|
||||
super().__init__(provider, api_key, base_url, model, reasoning_effort, **kwargs)
|
||||
self.timeout = timeout
|
||||
self._litellm: Any = None
|
||||
|
||||
try:
|
||||
import litellm
|
||||
|
||||
self._litellm = litellm
|
||||
# Suppress LiteLLM's verbose logging
|
||||
litellm.suppress_debug_info = True # type: ignore[assignment]
|
||||
# Drop unsupported params instead of raising errors (e.g. tool_choice on some Bedrock models)
|
||||
litellm.drop_params = True # type: ignore[assignment]
|
||||
logging.getLogger("LiteLLM").setLevel(logging.WARNING)
|
||||
logger.info(f"LiteLLM SDK initialized for model: {self.model}")
|
||||
except ImportError as e:
|
||||
raise RuntimeError("LiteLLM SDK not installed. Run: uv add litellm or pip install litellm") from e
|
||||
|
||||
async def verify_connection(self) -> None:
|
||||
try:
|
||||
test_messages = [{"role": "user", "content": "test"}]
|
||||
await self.call(
|
||||
messages=test_messages,
|
||||
max_completion_tokens=50,
|
||||
temperature=0.0,
|
||||
scope="verification",
|
||||
max_retries=0,
|
||||
)
|
||||
logger.info("LiteLLM connection verified successfully")
|
||||
except OutputTooLongError:
|
||||
# Truncation is fine for verification — it means the connection works
|
||||
logger.info("LiteLLM connection verified successfully (response truncated)")
|
||||
except Exception as e:
|
||||
logger.error(f"LiteLLM connection verification failed: {e}")
|
||||
raise RuntimeError(f"Failed to verify LiteLLM connection: {e}") from e
|
||||
|
||||
def _build_common_kwargs(
|
||||
self,
|
||||
messages: list[dict[str, Any]],
|
||||
max_completion_tokens: int | None = None,
|
||||
temperature: float | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Build common kwargs for litellm calls."""
|
||||
kwargs: dict[str, Any] = {
|
||||
"model": self.model,
|
||||
"messages": messages,
|
||||
"timeout": self.timeout,
|
||||
}
|
||||
|
||||
if self.api_key:
|
||||
kwargs["api_key"] = self.api_key
|
||||
if self.base_url:
|
||||
kwargs["api_base"] = self.base_url
|
||||
if max_completion_tokens is not None:
|
||||
kwargs["max_completion_tokens"] = max_completion_tokens
|
||||
if temperature is not None:
|
||||
kwargs["temperature"] = temperature
|
||||
|
||||
return kwargs
|
||||
|
||||
async def call(
|
||||
self,
|
||||
messages: list[dict[str, str]],
|
||||
response_format: Any | None = None,
|
||||
max_completion_tokens: int | None = None,
|
||||
temperature: float | None = None,
|
||||
scope: str = "memory",
|
||||
max_retries: int = 10,
|
||||
initial_backoff: float = 1.0,
|
||||
max_backoff: float = 60.0,
|
||||
skip_validation: bool = False,
|
||||
strict_schema: bool = False,
|
||||
return_usage: bool = False,
|
||||
) -> Any:
|
||||
start_time = time.time()
|
||||
|
||||
call_kwargs = self._build_common_kwargs(messages, max_completion_tokens, temperature)
|
||||
|
||||
# Add JSON schema response format if provided
|
||||
if response_format is not None and hasattr(response_format, "model_json_schema"):
|
||||
schema = response_format.model_json_schema()
|
||||
call_kwargs["response_format"] = {
|
||||
"type": "json_schema",
|
||||
"json_schema": {
|
||||
"name": response_format.__name__ if hasattr(response_format, "__name__") else "response",
|
||||
"schema": schema,
|
||||
"strict": strict_schema,
|
||||
},
|
||||
}
|
||||
|
||||
last_exception = None
|
||||
|
||||
for attempt in range(max_retries + 1):
|
||||
try:
|
||||
response = await self._litellm.acompletion(**call_kwargs)
|
||||
|
||||
content = response.choices[0].message.content or ""
|
||||
finish_reason = response.choices[0].finish_reason
|
||||
|
||||
# Check for length-limited output
|
||||
if finish_reason == "length":
|
||||
raise OutputTooLongError("LiteLLM response was truncated due to token limit")
|
||||
|
||||
if response_format is not None:
|
||||
# Strip markdown code fences if present
|
||||
clean_content = content
|
||||
if "```json" in content:
|
||||
clean_content = content.split("```json")[1].split("```")[0].strip()
|
||||
elif "```" in content:
|
||||
clean_content = content.split("```")[1].split("```")[0].strip()
|
||||
|
||||
try:
|
||||
json_data = json.loads(clean_content)
|
||||
except json.JSONDecodeError:
|
||||
json_data = json.loads(content)
|
||||
|
||||
if skip_validation:
|
||||
result = json_data
|
||||
else:
|
||||
result = response_format.model_validate(json_data)
|
||||
else:
|
||||
result = content
|
||||
|
||||
# Extract usage
|
||||
input_tokens = getattr(response.usage, "prompt_tokens", 0) or 0
|
||||
output_tokens = getattr(response.usage, "completion_tokens", 0) or 0
|
||||
total_tokens = input_tokens + output_tokens
|
||||
|
||||
# Record metrics
|
||||
duration = time.time() - start_time
|
||||
metrics = get_metrics_collector()
|
||||
metrics.record_llm_call(
|
||||
provider=self.provider,
|
||||
model=self.model,
|
||||
scope=scope,
|
||||
duration=duration,
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
success=True,
|
||||
)
|
||||
|
||||
# Record trace span
|
||||
from hindsight_api.tracing import _serialize_for_span, get_span_recorder
|
||||
|
||||
span_recorder = get_span_recorder()
|
||||
span_recorder.record_llm_call(
|
||||
provider=self.provider,
|
||||
model=self.model,
|
||||
scope=scope,
|
||||
messages=messages,
|
||||
response_content=_serialize_for_span(result),
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
duration=duration,
|
||||
finish_reason=finish_reason,
|
||||
error=None,
|
||||
)
|
||||
|
||||
if duration > 10.0:
|
||||
logger.info(
|
||||
f"slow llm call: scope={scope}, model={self.provider}/{self.model}, "
|
||||
f"input_tokens={input_tokens}, output_tokens={output_tokens}, "
|
||||
f"time={duration:.3f}s"
|
||||
)
|
||||
|
||||
if return_usage:
|
||||
token_usage = TokenUsage(
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
total_tokens=total_tokens,
|
||||
)
|
||||
return result, token_usage
|
||||
return result
|
||||
|
||||
except OutputTooLongError:
|
||||
raise
|
||||
|
||||
except json.JSONDecodeError as e:
|
||||
last_exception = e
|
||||
if attempt < max_retries:
|
||||
logger.warning("LiteLLM returned invalid JSON, retrying...")
|
||||
backoff = min(initial_backoff * (2**attempt), max_backoff)
|
||||
await asyncio.sleep(backoff)
|
||||
continue
|
||||
else:
|
||||
logger.error(f"LiteLLM returned invalid JSON after {max_retries + 1} attempts")
|
||||
raise
|
||||
|
||||
except Exception as e:
|
||||
error_str = str(e).lower()
|
||||
# Fast fail on auth errors
|
||||
if "401" in error_str or "403" in error_str or "unauthorized" in error_str:
|
||||
logger.error(f"LiteLLM auth error, not retrying: {e}")
|
||||
raise
|
||||
|
||||
last_exception = e
|
||||
if attempt < max_retries:
|
||||
# Retry on rate limits, connection errors, server errors
|
||||
is_retryable = any(
|
||||
keyword in error_str
|
||||
for keyword in ("rate", "limit", "timeout", "connection", "500", "502", "503", "529")
|
||||
)
|
||||
if is_retryable:
|
||||
backoff = min(initial_backoff * (2**attempt), max_backoff)
|
||||
jitter = backoff * 0.2 * (2 * (time.time() % 1) - 1)
|
||||
await asyncio.sleep(backoff + jitter)
|
||||
continue
|
||||
|
||||
logger.error(f"LiteLLM API error after {attempt + 1} attempts: {e}")
|
||||
raise
|
||||
|
||||
if last_exception:
|
||||
raise last_exception
|
||||
raise RuntimeError("LiteLLM call failed after all retries")
|
||||
|
||||
async def call_with_tools(
|
||||
self,
|
||||
messages: list[dict[str, Any]],
|
||||
tools: list[dict[str, Any]],
|
||||
max_completion_tokens: int | None = None,
|
||||
temperature: float | None = None,
|
||||
scope: str = "tools",
|
||||
max_retries: int = 5,
|
||||
initial_backoff: float = 1.0,
|
||||
max_backoff: float = 30.0,
|
||||
tool_choice: str | dict[str, Any] = "auto",
|
||||
) -> LLMToolCallResult:
|
||||
start_time = time.time()
|
||||
|
||||
call_kwargs = self._build_common_kwargs(messages, max_completion_tokens, temperature)
|
||||
call_kwargs["tools"] = tools
|
||||
call_kwargs["tool_choice"] = tool_choice
|
||||
|
||||
last_exception = None
|
||||
for attempt in range(max_retries + 1):
|
||||
try:
|
||||
response = await self._litellm.acompletion(**call_kwargs)
|
||||
|
||||
message = response.choices[0].message
|
||||
content = message.content
|
||||
finish_reason = response.choices[0].finish_reason
|
||||
|
||||
# Extract tool calls
|
||||
tool_calls: list[LLMToolCall] = []
|
||||
if message.tool_calls:
|
||||
for tc in message.tool_calls:
|
||||
arguments = tc.function.arguments
|
||||
if isinstance(arguments, str):
|
||||
arguments = json.loads(arguments)
|
||||
tool_calls.append(
|
||||
LLMToolCall(
|
||||
id=tc.id,
|
||||
name=tc.function.name,
|
||||
arguments=arguments,
|
||||
)
|
||||
)
|
||||
|
||||
# Extract usage
|
||||
input_tokens = getattr(response.usage, "prompt_tokens", 0) or 0
|
||||
output_tokens = getattr(response.usage, "completion_tokens", 0) or 0
|
||||
|
||||
# Record metrics
|
||||
duration = time.time() - start_time
|
||||
metrics = get_metrics_collector()
|
||||
metrics.record_llm_call(
|
||||
provider=self.provider,
|
||||
model=self.model,
|
||||
scope=scope,
|
||||
duration=duration,
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
success=True,
|
||||
)
|
||||
|
||||
# Record trace span
|
||||
from hindsight_api.tracing import get_span_recorder
|
||||
|
||||
span_recorder = get_span_recorder()
|
||||
tool_calls_dict = (
|
||||
[{"id": tc.id, "name": tc.name, "arguments": tc.arguments} for tc in tool_calls]
|
||||
if tool_calls
|
||||
else None
|
||||
)
|
||||
span_recorder.record_llm_call(
|
||||
provider=self.provider,
|
||||
model=self.model,
|
||||
scope=scope,
|
||||
messages=messages,
|
||||
response_content=content,
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
duration=duration,
|
||||
finish_reason=finish_reason,
|
||||
error=None,
|
||||
tool_calls=tool_calls_dict,
|
||||
)
|
||||
|
||||
return LLMToolCallResult(
|
||||
content=content,
|
||||
tool_calls=tool_calls,
|
||||
finish_reason=finish_reason or ("tool_calls" if tool_calls else "stop"),
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
error_str = str(e).lower()
|
||||
if "401" in error_str or "403" in error_str or "unauthorized" in error_str:
|
||||
raise
|
||||
|
||||
last_exception = e
|
||||
if attempt < max_retries:
|
||||
is_retryable = any(
|
||||
keyword in error_str
|
||||
for keyword in ("rate", "limit", "timeout", "connection", "500", "502", "503", "529")
|
||||
)
|
||||
if is_retryable:
|
||||
await asyncio.sleep(min(initial_backoff * (2**attempt), max_backoff))
|
||||
continue
|
||||
|
||||
logger.error(f"LiteLLM tool call error after {attempt + 1} attempts: {e}")
|
||||
raise
|
||||
|
||||
if last_exception:
|
||||
raise last_exception
|
||||
raise RuntimeError("LiteLLM tool call failed after all retries")
|
||||
|
||||
async def cleanup(self) -> None:
|
||||
"""Clean up resources."""
|
||||
pass
|
||||
@@ -1,78 +0,0 @@
|
||||
"""
|
||||
No-op LLM provider for chunk-only storage mode.
|
||||
|
||||
When the LLM provider is set to "none", the system operates without any LLM dependency.
|
||||
Retain uses chunks mode (no fact extraction), and reflect/consolidation are disabled.
|
||||
This provider acts as a safety net — if any code path unexpectedly tries to call the LLM,
|
||||
it raises a clear error instead of a confusing connection failure.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import Any
|
||||
|
||||
from ..llm_interface import LLMInterface
|
||||
from ..response_models import LLMToolCallResult
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class LLMNotAvailableError(Exception):
|
||||
"""Raised when an operation requires an LLM but the provider is set to 'none'."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class NoneLLM(LLMInterface):
|
||||
"""
|
||||
No-op LLM provider that rejects all LLM calls.
|
||||
|
||||
Used when HINDSIGHT_API_LLM_PROVIDER=none to run Hindsight as a chunk store
|
||||
with semantic search but without LLM-based features (fact extraction, reflect,
|
||||
consolidation).
|
||||
"""
|
||||
|
||||
async def verify_connection(self) -> None:
|
||||
"""No-op — no LLM connection to verify."""
|
||||
logger.debug("NoneLLM: no LLM connection to verify (provider=none)")
|
||||
|
||||
async def call(
|
||||
self,
|
||||
messages: list[dict[str, str]],
|
||||
response_format: Any | None = None,
|
||||
max_completion_tokens: int | None = None,
|
||||
temperature: float | None = None,
|
||||
scope: str = "memory",
|
||||
max_retries: int = 10,
|
||||
initial_backoff: float = 1.0,
|
||||
max_backoff: float = 60.0,
|
||||
skip_validation: bool = False,
|
||||
strict_schema: bool = False,
|
||||
return_usage: bool = False,
|
||||
) -> Any:
|
||||
"""Raise LLMNotAvailableError — no LLM is configured."""
|
||||
raise LLMNotAvailableError(
|
||||
"LLM provider is set to 'none'. This operation requires an LLM. "
|
||||
"Set HINDSIGHT_API_LLM_PROVIDER to a real provider (e.g., openai, anthropic, gemini)."
|
||||
)
|
||||
|
||||
async def call_with_tools(
|
||||
self,
|
||||
messages: list[dict[str, Any]],
|
||||
tools: list[dict[str, Any]],
|
||||
max_completion_tokens: int | None = None,
|
||||
temperature: float | None = None,
|
||||
scope: str = "tools",
|
||||
max_retries: int = 5,
|
||||
initial_backoff: float = 1.0,
|
||||
max_backoff: float = 30.0,
|
||||
tool_choice: str | dict[str, Any] = "auto",
|
||||
) -> LLMToolCallResult:
|
||||
"""Raise LLMNotAvailableError — no LLM is configured."""
|
||||
raise LLMNotAvailableError(
|
||||
"LLM provider is set to 'none'. This operation requires an LLM. "
|
||||
"Set HINDSIGHT_API_LLM_PROVIDER to a real provider (e.g., openai, anthropic, gemini)."
|
||||
)
|
||||
|
||||
async def cleanup(self) -> None:
|
||||
"""No-op — nothing to clean up."""
|
||||
pass
|
||||
@@ -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,883 +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 import defaultdict
|
||||
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 ChunkMetadata, EntityLink, ExtractedFact, ProcessedFact, RetainContent, RetainContentDict
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _build_retain_params(contents_dicts, document_tags=None, doc_contents=None):
|
||||
"""Build retain_params and merged_tags from content dicts."""
|
||||
if doc_contents is not None:
|
||||
# Per-document mode: doc_contents is list of (idx, content_dict)
|
||||
items = [item for _, item in doc_contents]
|
||||
else:
|
||||
items = contents_dicts
|
||||
|
||||
all_tags = set(document_tags or [])
|
||||
for item in items:
|
||||
item_tags = item.get("tags", []) or []
|
||||
all_tags.update(item_tags)
|
||||
merged_tags = list(all_tags)
|
||||
|
||||
retain_params = {}
|
||||
if items:
|
||||
first_item = items[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"]
|
||||
|
||||
return retain_params, merged_tags
|
||||
|
||||
|
||||
async def _insert_facts_and_links(
|
||||
conn,
|
||||
entity_resolver,
|
||||
bank_id: str,
|
||||
contents: list[RetainContent],
|
||||
extracted_facts: list,
|
||||
processed_facts: list[ProcessedFact],
|
||||
config,
|
||||
log_buffer: list[str],
|
||||
outbox_callback=None,
|
||||
) -> list[list[str]]:
|
||||
"""
|
||||
Shared pipeline: insert facts, process entities, create all link types.
|
||||
|
||||
Used by both the full retain and delta retain paths.
|
||||
|
||||
Returns:
|
||||
List of unit ID lists mapped back to original content items.
|
||||
"""
|
||||
unit_ids = await fact_storage.insert_facts_batch(conn, bank_id, processed_facts)
|
||||
step_start = time.time()
|
||||
log_buffer.append(f" Insert facts: {len(unit_ids)} units in {time.time() - step_start:.3f}s")
|
||||
|
||||
if unit_ids:
|
||||
# Process entities
|
||||
step_start = time.time()
|
||||
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,
|
||||
processed_facts,
|
||||
log_buffer,
|
||||
user_entities_per_content=user_entities_per_content,
|
||||
entity_labels=getattr(config, "entity_labels", None),
|
||||
)
|
||||
log_buffer.append(f" 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" 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 processed_facts]
|
||||
semantic_link_count = await link_creation.create_semantic_links_batch(
|
||||
conn, bank_id, unit_ids, embeddings_for_links
|
||||
)
|
||||
log_buffer.append(f" 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, bank_id)
|
||||
log_buffer.append(
|
||||
f" 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, bank_id, unit_ids, processed_facts)
|
||||
log_buffer.append(f" Causal links: {causal_link_count} links in {time.time() - step_start:.3f}s")
|
||||
|
||||
# Map results back to original content items
|
||||
result_unit_ids = _map_results_to_contents(contents, extracted_facts, unit_ids if unit_ids else [])
|
||||
|
||||
if outbox_callback:
|
||||
await outbox_callback(conn)
|
||||
|
||||
return result_unit_ids
|
||||
|
||||
|
||||
async def _extract_and_embed(
|
||||
contents: list[RetainContent],
|
||||
llm_config,
|
||||
agent_name: str,
|
||||
config,
|
||||
embeddings_model,
|
||||
format_date_fn,
|
||||
fact_type_override: str | None,
|
||||
log_buffer: list[str],
|
||||
pool=None,
|
||||
operation_id: str | None = None,
|
||||
schema: str | None = None,
|
||||
) -> tuple[list, list[ProcessedFact], list[ChunkMetadata], TokenUsage]:
|
||||
"""
|
||||
Shared pipeline: extract facts from contents and generate embeddings.
|
||||
|
||||
Returns:
|
||||
Tuple of (extracted_facts, processed_facts, chunks_metadata, usage)
|
||||
"""
|
||||
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" Extract facts: {len(extracted_facts)} facts, {len(chunks)} chunks "
|
||||
f"from {len(contents)} contents in {time.time() - step_start:.3f}s"
|
||||
)
|
||||
|
||||
if not extracted_facts:
|
||||
return extracted_facts, [], chunks, usage
|
||||
|
||||
if fact_type_override:
|
||||
for fact in extracted_facts:
|
||||
fact.fact_type = fact_type_override
|
||||
|
||||
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" Generate embeddings: {len(embeddings)} embeddings in {time.time() - step_start:.3f}s")
|
||||
|
||||
processed_facts = [ProcessedFact.from_extracted_fact(ef, emb) for ef, emb in zip(extracted_facts, embeddings)]
|
||||
|
||||
return extracted_facts, processed_facts, chunks, usage
|
||||
|
||||
|
||||
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.
|
||||
|
||||
Supports delta retain: when upserting a document that already has chunks,
|
||||
only re-processes chunks whose content has changed. Unchanged chunks keep
|
||||
their existing facts, entities, and links.
|
||||
"""
|
||||
start_time = time.time()
|
||||
total_chars = sum(len(item.get("content", "")) for item in contents_dicts)
|
||||
|
||||
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 = _build_contents(contents_dicts, document_tags)
|
||||
|
||||
# --- Delta retain: check if we can skip unchanged chunks ---
|
||||
if is_first_batch:
|
||||
delta_result = await _try_delta_retain(
|
||||
pool,
|
||||
embeddings_model,
|
||||
llm_config,
|
||||
entity_resolver,
|
||||
format_date_fn,
|
||||
bank_id,
|
||||
contents_dicts,
|
||||
contents,
|
||||
config,
|
||||
document_id,
|
||||
fact_type_override,
|
||||
document_tags,
|
||||
agent_name,
|
||||
log_buffer,
|
||||
start_time,
|
||||
operation_id,
|
||||
schema,
|
||||
outbox_callback,
|
||||
)
|
||||
if delta_result is not None:
|
||||
return delta_result
|
||||
|
||||
# --- Full retain path ---
|
||||
extracted_facts, processed_facts, chunks, usage = await _extract_and_embed(
|
||||
contents,
|
||||
llm_config,
|
||||
agent_name,
|
||||
config,
|
||||
embeddings_model,
|
||||
format_date_fn,
|
||||
fact_type_override,
|
||||
log_buffer,
|
||||
pool,
|
||||
operation_id,
|
||||
schema,
|
||||
)
|
||||
|
||||
if not extracted_facts:
|
||||
await _handle_zero_facts_documents(
|
||||
pool,
|
||||
bank_id,
|
||||
contents_dicts,
|
||||
contents,
|
||||
config,
|
||||
document_id,
|
||||
is_first_batch,
|
||||
document_tags,
|
||||
chunks,
|
||||
log_buffer,
|
||||
start_time,
|
||||
)
|
||||
return [[] for _ in contents], usage
|
||||
|
||||
# Group contents by document_id
|
||||
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))
|
||||
|
||||
# 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
|
||||
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
|
||||
entity_resolver.discard_pending_stats()
|
||||
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
async with conn.transaction():
|
||||
# Handle document tracking
|
||||
step_start = time.time()
|
||||
doc_id_mapping = {}
|
||||
|
||||
if document_id:
|
||||
combined_content = "\n".join([c.get("content", "") for c in contents_dicts])
|
||||
retain_params, merged_tags = _build_retain_params(contents_dicts, document_tags)
|
||||
await fact_storage.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
|
||||
else:
|
||||
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
|
||||
should_create_doc = (original_doc_id is not None) or chunks
|
||||
if should_create_doc:
|
||||
if actual_doc_id is None:
|
||||
actual_doc_id = str(uuid.uuid4())
|
||||
doc_id_mapping[original_doc_id] = actual_doc_id
|
||||
combined_content = "\n".join([c.get("content", "") for _, c in doc_contents])
|
||||
retain_params, merged_tags = _build_retain_params(
|
||||
contents_dicts, document_tags, doc_contents=doc_contents
|
||||
)
|
||||
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" Document tracking: {len(document_ids_added)} documents in {time.time() - step_start:.3f}s"
|
||||
)
|
||||
|
||||
# Store chunks and map to facts
|
||||
step_start = time.time()
|
||||
chunk_id_map_by_doc = {}
|
||||
if chunks:
|
||||
chunks_by_doc = defaultdict(list)
|
||||
for chunk in chunks:
|
||||
original_doc_id = contents_dicts[chunk.content_index].get("document_id")
|
||||
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)
|
||||
|
||||
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)
|
||||
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" Store chunks: {len(chunks)} chunks for {len(chunks_by_doc)} documents "
|
||||
f"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):
|
||||
original_doc_id = contents_dicts[fact.content_index].get("document_id")
|
||||
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
|
||||
if chunks and fact.chunk_index is not None:
|
||||
chunk_id = chunk_id_map_by_doc.get((actual_doc_id, fact.chunk_index))
|
||||
if chunk_id:
|
||||
processed_fact.chunk_id = chunk_id
|
||||
|
||||
# Insert facts and create all links (shared pipeline)
|
||||
result_unit_ids = await _insert_facts_and_links(
|
||||
conn,
|
||||
entity_resolver,
|
||||
bank_id,
|
||||
contents,
|
||||
extracted_facts,
|
||||
processed_facts,
|
||||
config,
|
||||
log_buffer,
|
||||
outbox_callback,
|
||||
)
|
||||
|
||||
await entity_resolver.flush_pending_stats()
|
||||
|
||||
total_time = time.time() - start_time
|
||||
log_buffer.append(f"{'=' * 60}")
|
||||
log_buffer.append(f"RETAIN_BATCH COMPLETE: {len(processed_facts)} 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
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Delta retain
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def _try_delta_retain(
|
||||
pool,
|
||||
embeddings_model,
|
||||
llm_config,
|
||||
entity_resolver,
|
||||
format_date_fn,
|
||||
bank_id,
|
||||
contents_dicts,
|
||||
contents,
|
||||
config,
|
||||
document_id,
|
||||
fact_type_override,
|
||||
document_tags,
|
||||
agent_name,
|
||||
log_buffer,
|
||||
start_time,
|
||||
operation_id,
|
||||
schema,
|
||||
outbox_callback,
|
||||
):
|
||||
"""
|
||||
Attempt delta retain for a document upsert. Returns result tuple if delta
|
||||
was performed, or None to fall back to full retain.
|
||||
"""
|
||||
# Need a single document_id
|
||||
effective_doc_id = document_id
|
||||
if not effective_doc_id:
|
||||
doc_ids = {item.get("document_id") for item in contents_dicts if item.get("document_id")}
|
||||
if len(doc_ids) != 1:
|
||||
return None
|
||||
effective_doc_id = doc_ids.pop()
|
||||
|
||||
# Load existing chunks
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
existing_chunks = await chunk_storage.load_existing_chunks(conn, bank_id, effective_doc_id)
|
||||
|
||||
if not existing_chunks:
|
||||
return None
|
||||
|
||||
if any(c.content_hash is None for c in existing_chunks):
|
||||
logger.info(f"Delta retain skipped for {effective_doc_id}: existing chunks lack content_hash (pre-migration)")
|
||||
return None
|
||||
|
||||
# Chunk new content and classify changes
|
||||
step_start = time.time()
|
||||
new_chunks_with_contents = _chunk_contents_for_delta(contents, config)
|
||||
log_buffer.append(
|
||||
f"[delta] Chunked new content: {len(new_chunks_with_contents)} chunks in {time.time() - step_start:.3f}s"
|
||||
)
|
||||
|
||||
existing_by_index = {c.chunk_index: c for c in existing_chunks}
|
||||
new_hashes = {idx: chunk_storage.compute_chunk_hash(text) for idx, text in new_chunks_with_contents.items()}
|
||||
|
||||
unchanged_indices, changed_indices, new_indices, removed_indices = [], [], [], []
|
||||
for idx, new_hash in new_hashes.items():
|
||||
existing = existing_by_index.get(idx)
|
||||
if existing and existing.content_hash == new_hash:
|
||||
unchanged_indices.append(idx)
|
||||
elif existing:
|
||||
changed_indices.append(idx)
|
||||
else:
|
||||
new_indices.append(idx)
|
||||
for idx in existing_by_index:
|
||||
if idx not in new_hashes:
|
||||
removed_indices.append(idx)
|
||||
|
||||
log_buffer.append(
|
||||
f"[delta] Chunk diff: {len(unchanged_indices)} unchanged, "
|
||||
f"{len(changed_indices)} changed, {len(new_indices)} new, "
|
||||
f"{len(removed_indices)} removed"
|
||||
)
|
||||
|
||||
if not unchanged_indices:
|
||||
logger.info(f"Delta retain: no unchanged chunks for {effective_doc_id}, falling back to full retain")
|
||||
return None
|
||||
|
||||
chunks_to_process = changed_indices + new_indices
|
||||
|
||||
if not chunks_to_process and not removed_indices:
|
||||
# Nothing changed — just update document metadata/tags
|
||||
log_buffer.append("[delta] No chunk changes detected — updating document metadata only")
|
||||
return await _delta_metadata_only(
|
||||
pool,
|
||||
bank_id,
|
||||
contents_dicts,
|
||||
contents,
|
||||
effective_doc_id,
|
||||
document_tags,
|
||||
log_buffer,
|
||||
start_time,
|
||||
outbox_callback,
|
||||
)
|
||||
|
||||
# Build content items for only the changed/new chunks
|
||||
delta_contents, delta_chunk_map = _build_delta_contents(contents, new_chunks_with_contents, chunks_to_process)
|
||||
|
||||
if not delta_contents:
|
||||
return await _delta_metadata_only(
|
||||
pool,
|
||||
bank_id,
|
||||
contents_dicts,
|
||||
contents,
|
||||
effective_doc_id,
|
||||
document_tags,
|
||||
log_buffer,
|
||||
start_time,
|
||||
outbox_callback,
|
||||
)
|
||||
|
||||
# Extract facts and generate embeddings (shared pipeline)
|
||||
extracted_facts, processed_facts, new_chunk_metadata, usage = await _extract_and_embed(
|
||||
delta_contents,
|
||||
llm_config,
|
||||
agent_name,
|
||||
config,
|
||||
embeddings_model,
|
||||
format_date_fn,
|
||||
fact_type_override,
|
||||
log_buffer,
|
||||
pool,
|
||||
operation_id,
|
||||
schema,
|
||||
)
|
||||
|
||||
# Database transaction
|
||||
result_unit_ids: list[list[str]] = []
|
||||
log_buffer_pre_db = len(log_buffer)
|
||||
|
||||
async def _run_delta_db_work() -> None:
|
||||
nonlocal result_unit_ids
|
||||
del log_buffer[log_buffer_pre_db:]
|
||||
for pf in processed_facts:
|
||||
pf.document_id = None
|
||||
pf.chunk_id = None
|
||||
entity_resolver.discard_pending_stats()
|
||||
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
async with conn.transaction():
|
||||
# Update document metadata (no delete)
|
||||
step_start = time.time()
|
||||
combined_content = "\n".join([c.get("content", "") for c in contents_dicts])
|
||||
retain_params, merged_tags = _build_retain_params(contents_dicts, document_tags)
|
||||
await fact_storage.upsert_document_metadata(
|
||||
conn,
|
||||
bank_id,
|
||||
effective_doc_id,
|
||||
combined_content,
|
||||
retain_params,
|
||||
merged_tags,
|
||||
)
|
||||
log_buffer.append(f" Document metadata update in {time.time() - step_start:.3f}s")
|
||||
|
||||
# Delete changed and removed chunks (cascades to memory_units and links)
|
||||
step_start = time.time()
|
||||
chunks_to_delete = [
|
||||
existing_by_index[idx].chunk_id
|
||||
for idx in changed_indices + removed_indices
|
||||
if idx in existing_by_index
|
||||
]
|
||||
await chunk_storage.delete_chunks_by_ids(conn, chunks_to_delete)
|
||||
log_buffer.append(
|
||||
f" Deleted {len(chunks_to_delete)} chunks "
|
||||
f"({len(changed_indices)} changed + {len(removed_indices)} removed) "
|
||||
f"in {time.time() - step_start:.3f}s"
|
||||
)
|
||||
|
||||
# Update tags on unchanged chunks' memory units
|
||||
step_start = time.time()
|
||||
updated_count = await fact_storage.update_memory_units_tags(
|
||||
conn, bank_id, effective_doc_id, merged_tags
|
||||
)
|
||||
log_buffer.append(
|
||||
f" Updated tags on {updated_count} existing memory units in {time.time() - step_start:.3f}s"
|
||||
)
|
||||
|
||||
# Store new/changed chunks
|
||||
step_start = time.time()
|
||||
chunk_id_map_by_doc = {}
|
||||
if new_chunk_metadata:
|
||||
remapped_chunks = [
|
||||
ChunkMetadata(
|
||||
chunk_text=cm.chunk_text,
|
||||
fact_count=cm.fact_count,
|
||||
content_index=cm.content_index,
|
||||
chunk_index=delta_chunk_map.get(cm.chunk_index, cm.chunk_index),
|
||||
)
|
||||
for cm in new_chunk_metadata
|
||||
]
|
||||
chunk_id_map = await chunk_storage.store_chunks_batch(
|
||||
conn, bank_id, effective_doc_id, remapped_chunks
|
||||
)
|
||||
for chunk_idx, chunk_id in chunk_id_map.items():
|
||||
chunk_id_map_by_doc[(effective_doc_id, chunk_idx)] = chunk_id
|
||||
log_buffer.append(
|
||||
f" Stored {len(remapped_chunks)} new/changed chunks in {time.time() - step_start:.3f}s"
|
||||
)
|
||||
|
||||
# Map chunk_ids and document_ids to processed facts
|
||||
for ef, pf in zip(extracted_facts, processed_facts):
|
||||
pf.document_id = effective_doc_id
|
||||
if ef.chunk_index is not None:
|
||||
original_idx = delta_chunk_map.get(ef.chunk_index, ef.chunk_index)
|
||||
chunk_id = chunk_id_map_by_doc.get((effective_doc_id, original_idx))
|
||||
if chunk_id:
|
||||
pf.chunk_id = chunk_id
|
||||
|
||||
# Insert facts and create all links (shared pipeline)
|
||||
result_unit_ids = await _insert_facts_and_links(
|
||||
conn,
|
||||
entity_resolver,
|
||||
bank_id,
|
||||
contents,
|
||||
extracted_facts,
|
||||
processed_facts,
|
||||
config,
|
||||
log_buffer,
|
||||
outbox_callback,
|
||||
)
|
||||
|
||||
await entity_resolver.flush_pending_stats()
|
||||
|
||||
total_time = time.time() - start_time
|
||||
log_buffer.append(f"{'=' * 60}")
|
||||
log_buffer.append(
|
||||
f"DELTA RETAIN COMPLETE: {len(processed_facts)} new units, "
|
||||
f"{len(unchanged_indices)} chunks unchanged in {total_time:.3f}s"
|
||||
)
|
||||
log_buffer.append(f"Document: {effective_doc_id}")
|
||||
log_buffer.append(f"{'=' * 60}")
|
||||
logger.info("\n" + "\n".join(log_buffer) + "\n")
|
||||
|
||||
await retry_with_backoff(_run_delta_db_work)
|
||||
return result_unit_ids, usage
|
||||
|
||||
|
||||
async def _delta_metadata_only(
|
||||
pool,
|
||||
bank_id,
|
||||
contents_dicts,
|
||||
contents,
|
||||
document_id,
|
||||
document_tags,
|
||||
log_buffer,
|
||||
start_time,
|
||||
outbox_callback,
|
||||
):
|
||||
"""Handle the case where no chunks changed — just update document metadata and tags."""
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
async with conn.transaction():
|
||||
combined_content = "\n".join([c.get("content", "") for c in contents_dicts])
|
||||
retain_params, merged_tags = _build_retain_params(contents_dicts, document_tags)
|
||||
await fact_storage.upsert_document_metadata(
|
||||
conn,
|
||||
bank_id,
|
||||
document_id,
|
||||
combined_content,
|
||||
retain_params,
|
||||
merged_tags,
|
||||
)
|
||||
await fact_storage.update_memory_units_tags(conn, bank_id, document_id, merged_tags)
|
||||
if outbox_callback:
|
||||
await outbox_callback(conn)
|
||||
|
||||
total_time = time.time() - start_time
|
||||
log_buffer.append(f"DELTA RETAIN (no changes): metadata updated in {total_time:.3f}s")
|
||||
logger.info("\n" + "\n".join(log_buffer) + "\n")
|
||||
return [[] for _ in contents], TokenUsage()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _build_contents(contents_dicts: list[RetainContentDict], document_tags: list[str] | None) -> list[RetainContent]:
|
||||
"""Convert content dicts to RetainContent objects."""
|
||||
contents = []
|
||||
for item in contents_dicts:
|
||||
item_tags = item.get("tags", []) or []
|
||||
merged_tags = list(set(item_tags + (document_tags or [])))
|
||||
|
||||
if "event_date" in item and item["event_date"] is None:
|
||||
event_date_value = None
|
||||
elif item.get("event_date"):
|
||||
event_date_value = parse_datetime_flexible(item["event_date"])
|
||||
else:
|
||||
event_date_value = utcnow()
|
||||
|
||||
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)
|
||||
return contents
|
||||
|
||||
|
||||
async def _handle_zero_facts_documents(
|
||||
pool,
|
||||
bank_id,
|
||||
contents_dicts,
|
||||
contents,
|
||||
config,
|
||||
document_id,
|
||||
is_first_batch,
|
||||
document_tags,
|
||||
chunks,
|
||||
log_buffer,
|
||||
start_time,
|
||||
):
|
||||
"""Handle document tracking when zero facts were extracted."""
|
||||
docs_tracked = 0
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
async with conn.transaction():
|
||||
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))
|
||||
|
||||
if document_id:
|
||||
combined_content = "\n".join([c.get("content", "") for c in contents_dicts])
|
||||
retain_params, merged_tags = _build_retain_params(contents_dicts, document_tags)
|
||||
await fact_storage.handle_document_tracking(
|
||||
conn, bank_id, document_id, combined_content, is_first_batch, retain_params, merged_tags
|
||||
)
|
||||
docs_tracked += 1
|
||||
else:
|
||||
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():
|
||||
should_create_doc = (original_doc_id is not None) or chunks
|
||||
if not should_create_doc:
|
||||
continue
|
||||
actual_doc_id = original_doc_id or str(uuid.uuid4())
|
||||
combined_content = "\n".join([c.get("content", "") for _, c in doc_contents])
|
||||
retain_params, merged_tags = _build_retain_params(
|
||||
contents_dicts, document_tags, doc_contents=doc_contents
|
||||
)
|
||||
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 "
|
||||
f"in {total_time:.3f}s ({doc_status}, no facts)"
|
||||
)
|
||||
|
||||
|
||||
def _chunk_contents_for_delta(contents: list[RetainContent], config) -> dict[int, str]:
|
||||
"""
|
||||
Chunk contents the same way fact_extraction does, returning a map of
|
||||
global_chunk_index -> chunk_text.
|
||||
"""
|
||||
result = {}
|
||||
global_chunk_idx = 0
|
||||
for content in contents:
|
||||
chunk_size = getattr(config, "retain_chunk_size", 120000)
|
||||
chunks = fact_extraction.chunk_text(content.content, chunk_size)
|
||||
for chunk_text in chunks:
|
||||
result[global_chunk_idx] = chunk_text
|
||||
global_chunk_idx += 1
|
||||
return result
|
||||
|
||||
|
||||
def _build_delta_contents(
|
||||
original_contents: list[RetainContent],
|
||||
new_chunks_with_contents: dict[int, str],
|
||||
chunks_to_process: list[int],
|
||||
) -> tuple[list[RetainContent], dict[int, int]]:
|
||||
"""
|
||||
Build RetainContent items containing only the chunks that need processing.
|
||||
|
||||
Returns:
|
||||
- List of RetainContent items (one per chunk to process)
|
||||
- Map of delta_chunk_index -> original_chunk_index
|
||||
"""
|
||||
if not chunks_to_process or not original_contents:
|
||||
return [], {}
|
||||
|
||||
template_content = original_contents[0]
|
||||
delta_contents = []
|
||||
delta_chunk_map = {}
|
||||
|
||||
for original_chunk_idx in sorted(chunks_to_process):
|
||||
chunk_text = new_chunks_with_contents.get(original_chunk_idx)
|
||||
if not chunk_text:
|
||||
continue
|
||||
delta_content = RetainContent(
|
||||
content=chunk_text,
|
||||
context=template_content.context,
|
||||
event_date=template_content.event_date,
|
||||
metadata=template_content.metadata,
|
||||
entities=template_content.entities,
|
||||
tags=template_content.tags,
|
||||
observation_scopes=template_content.observation_scopes,
|
||||
)
|
||||
delta_contents.append(delta_content)
|
||||
delta_chunk_map[len(delta_contents) - 1] = original_chunk_idx
|
||||
|
||||
return delta_contents, delta_chunk_map
|
||||
|
||||
|
||||
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,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
@@ -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,211 +0,0 @@
|
||||
[build-system]
|
||||
requires = ["hatchling"]
|
||||
build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "hindsight-api-slim"
|
||||
version = "0.4.20"
|
||||
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,<=1.82.6", # 1.82.7+ contains a supply chain attack (malicious .pth credential stealer)
|
||||
"markitdown[pdf,docx,pptx,xlsx,xls]>=0.1.4", # File to markdown conversion
|
||||
"obstore>=0.4.0", # S3/GCS/Azure object storage client (Rust-backed)
|
||||
"winloop>=0.1.0; sys_platform == 'win32'",
|
||||
"uvloop>=0.22.1; sys_platform != 'win32'",
|
||||
# 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
|
||||
"python-multipart>=0.0.22", # Arbitrary file write via non-default configuration fix
|
||||
"tornado>=6.5.5", # DoS multipart/incomplete cookie validation fix
|
||||
"aiohttp>=3.13.3", # Multiple DoS vulnerabilities
|
||||
"claude-agent-sdk>=0.1.27",
|
||||
"boto3>=1.42.74",
|
||||
]
|
||||
|
||||
[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,449 +0,0 @@
|
||||
"""
|
||||
Tests for the audit log feature.
|
||||
|
||||
Tests the audit log list, stats, filtering, and pagination endpoints.
|
||||
Verifies that audit entries are created for operations when audit logging is enabled.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
|
||||
from hindsight_api.api import create_app
|
||||
from hindsight_api.config import get_config
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def audit_api_client(memory):
|
||||
"""Create a test client with audit logging enabled."""
|
||||
# Enable audit logging on the memory engine's audit logger
|
||||
memory._audit_logger._enabled = True
|
||||
memory._audit_logger._allowed_actions = None # All actions
|
||||
|
||||
app = create_app(memory, initialize_memory=False)
|
||||
transport = httpx.ASGITransport(app=app)
|
||||
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
|
||||
yield client
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def bank_id():
|
||||
"""Provide a unique bank ID for audit tests."""
|
||||
from datetime import datetime
|
||||
|
||||
return f"audit_test_{datetime.now().timestamp()}"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_audit_log_list_empty(audit_api_client, bank_id):
|
||||
"""Test listing audit logs for a bank with no entries returns empty."""
|
||||
# Create the bank first
|
||||
await audit_api_client.put(
|
||||
f"/v1/default/banks/{bank_id}",
|
||||
json={"name": "Audit Test Bank"},
|
||||
)
|
||||
|
||||
# Small delay for fire-and-forget audit writes
|
||||
await asyncio.sleep(0.5)
|
||||
|
||||
response = await audit_api_client.get(f"/v1/default/banks/{bank_id}/audit-logs")
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["bank_id"] == bank_id
|
||||
assert "total" in data
|
||||
assert "items" in data
|
||||
assert "limit" in data
|
||||
assert "offset" in data
|
||||
assert isinstance(data["items"], list)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_audit_log_created_for_retain(audit_api_client, bank_id):
|
||||
"""Test that a retain operation creates an audit log entry."""
|
||||
# Create bank
|
||||
await audit_api_client.put(
|
||||
f"/v1/default/banks/{bank_id}",
|
||||
json={"name": "Audit Test Bank"},
|
||||
)
|
||||
|
||||
# Perform a retain
|
||||
response = await audit_api_client.post(
|
||||
f"/v1/default/banks/{bank_id}/memories",
|
||||
json={
|
||||
"items": [{"content": "Alice likes cats", "context": "preferences"}],
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
|
||||
# Wait for fire-and-forget audit writes
|
||||
await asyncio.sleep(1.0)
|
||||
|
||||
# List audit logs - should have entries for create_bank and retain
|
||||
response = await audit_api_client.get(f"/v1/default/banks/{bank_id}/audit-logs")
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["total"] >= 1
|
||||
|
||||
actions = [item["action"] for item in data["items"]]
|
||||
assert "retain" in actions, f"Expected 'retain' in audit actions, got: {actions}"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_audit_log_entry_fields(audit_api_client, bank_id):
|
||||
"""Test that audit log entries have all expected fields."""
|
||||
# Create bank + recall to generate entries
|
||||
await audit_api_client.put(
|
||||
f"/v1/default/banks/{bank_id}",
|
||||
json={"name": "Audit Test Bank"},
|
||||
)
|
||||
|
||||
await audit_api_client.post(
|
||||
f"/v1/default/banks/{bank_id}/memories/recall",
|
||||
json={"query": "test query"},
|
||||
)
|
||||
|
||||
await asyncio.sleep(1.0)
|
||||
|
||||
response = await audit_api_client.get(f"/v1/default/banks/{bank_id}/audit-logs")
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["total"] >= 1
|
||||
|
||||
# Check the recall entry has all fields
|
||||
recall_entries = [item for item in data["items"] if item["action"] == "recall"]
|
||||
assert len(recall_entries) >= 1, f"Expected recall entry, got actions: {[i['action'] for i in data['items']]}"
|
||||
|
||||
entry = recall_entries[0]
|
||||
assert entry["id"] is not None
|
||||
assert entry["action"] == "recall"
|
||||
assert entry["transport"] == "http"
|
||||
assert entry["bank_id"] == bank_id
|
||||
assert entry["started_at"] is not None
|
||||
assert entry["ended_at"] is not None
|
||||
# Request should contain the recall parameters
|
||||
assert entry["request"] is not None
|
||||
assert "query" in entry["request"]
|
||||
# Response should contain the recall results
|
||||
assert entry["response"] is not None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_audit_log_filter_by_action(audit_api_client, bank_id):
|
||||
"""Test filtering audit logs by action type."""
|
||||
# Create bank and do retain + recall
|
||||
await audit_api_client.put(
|
||||
f"/v1/default/banks/{bank_id}",
|
||||
json={"name": "Audit Test Bank"},
|
||||
)
|
||||
await audit_api_client.post(
|
||||
f"/v1/default/banks/{bank_id}/memories",
|
||||
json={"items": [{"content": "test content", "context": "test"}]},
|
||||
)
|
||||
await audit_api_client.post(
|
||||
f"/v1/default/banks/{bank_id}/memories/recall",
|
||||
json={"query": "test"},
|
||||
)
|
||||
|
||||
await asyncio.sleep(1.0)
|
||||
|
||||
# Filter by retain only
|
||||
response = await audit_api_client.get(
|
||||
f"/v1/default/banks/{bank_id}/audit-logs",
|
||||
params={"action": "retain"},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
for item in data["items"]:
|
||||
assert item["action"] == "retain"
|
||||
|
||||
# Filter by recall only
|
||||
response = await audit_api_client.get(
|
||||
f"/v1/default/banks/{bank_id}/audit-logs",
|
||||
params={"action": "recall"},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
for item in data["items"]:
|
||||
assert item["action"] == "recall"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_audit_log_filter_by_transport(audit_api_client, bank_id):
|
||||
"""Test filtering audit logs by transport type."""
|
||||
await audit_api_client.put(
|
||||
f"/v1/default/banks/{bank_id}",
|
||||
json={"name": "Audit Test Bank"},
|
||||
)
|
||||
|
||||
await asyncio.sleep(0.5)
|
||||
|
||||
# Filter by http transport
|
||||
response = await audit_api_client.get(
|
||||
f"/v1/default/banks/{bank_id}/audit-logs",
|
||||
params={"transport": "http"},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
for item in data["items"]:
|
||||
assert item["transport"] == "http"
|
||||
|
||||
# Filter by mcp transport - should be empty (no MCP calls in this test)
|
||||
response = await audit_api_client.get(
|
||||
f"/v1/default/banks/{bank_id}/audit-logs",
|
||||
params={"transport": "mcp"},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["total"] == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_audit_log_filter_by_date_range(audit_api_client, bank_id):
|
||||
"""Test filtering audit logs by date range."""
|
||||
from datetime import datetime, timedelta, timezone
|
||||
|
||||
await audit_api_client.put(
|
||||
f"/v1/default/banks/{bank_id}",
|
||||
json={"name": "Audit Test Bank"},
|
||||
)
|
||||
|
||||
await asyncio.sleep(0.5)
|
||||
|
||||
now = datetime.now(timezone.utc)
|
||||
|
||||
# Filter with start_date in the past - should include entries
|
||||
past = (now - timedelta(hours=1)).isoformat()
|
||||
response = await audit_api_client.get(
|
||||
f"/v1/default/banks/{bank_id}/audit-logs",
|
||||
params={"start_date": past},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["total"] >= 1
|
||||
|
||||
# Filter with start_date in the future - should be empty
|
||||
future = (now + timedelta(hours=1)).isoformat()
|
||||
response = await audit_api_client.get(
|
||||
f"/v1/default/banks/{bank_id}/audit-logs",
|
||||
params={"start_date": future},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["total"] == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_audit_log_pagination(audit_api_client, bank_id):
|
||||
"""Test audit log pagination with limit and offset."""
|
||||
await audit_api_client.put(
|
||||
f"/v1/default/banks/{bank_id}",
|
||||
json={"name": "Audit Test Bank"},
|
||||
)
|
||||
|
||||
# Generate multiple audit entries
|
||||
for i in range(5):
|
||||
await audit_api_client.post(
|
||||
f"/v1/default/banks/{bank_id}/memories/recall",
|
||||
json={"query": f"test query {i}"},
|
||||
)
|
||||
|
||||
await asyncio.sleep(1.5)
|
||||
|
||||
# Get first page
|
||||
response = await audit_api_client.get(
|
||||
f"/v1/default/banks/{bank_id}/audit-logs",
|
||||
params={"limit": 2, "offset": 0},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
page1 = response.json()
|
||||
assert len(page1["items"]) == 2
|
||||
assert page1["limit"] == 2
|
||||
assert page1["offset"] == 0
|
||||
assert page1["total"] >= 5 # At least 5 recall + 1 create_bank
|
||||
|
||||
# Get second page
|
||||
response = await audit_api_client.get(
|
||||
f"/v1/default/banks/{bank_id}/audit-logs",
|
||||
params={"limit": 2, "offset": 2},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
page2 = response.json()
|
||||
assert len(page2["items"]) == 2
|
||||
assert page2["offset"] == 2
|
||||
|
||||
# Entries should be different between pages
|
||||
page1_ids = {item["id"] for item in page1["items"]}
|
||||
page2_ids = {item["id"] for item in page2["items"]}
|
||||
assert page1_ids.isdisjoint(page2_ids), "Pages should not overlap"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_audit_log_stats(audit_api_client, bank_id):
|
||||
"""Test the audit log stats endpoint returns correct structure."""
|
||||
await audit_api_client.put(
|
||||
f"/v1/default/banks/{bank_id}",
|
||||
json={"name": "Audit Test Bank"},
|
||||
)
|
||||
|
||||
await audit_api_client.post(
|
||||
f"/v1/default/banks/{bank_id}/memories/recall",
|
||||
json={"query": "stats test"},
|
||||
)
|
||||
|
||||
await asyncio.sleep(1.0)
|
||||
|
||||
# Get stats for last 24h
|
||||
response = await audit_api_client.get(
|
||||
f"/v1/default/banks/{bank_id}/audit-logs/stats",
|
||||
params={"period": "1d"},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["bank_id"] == bank_id
|
||||
assert data["period"] == "1d"
|
||||
assert data["trunc"] == "day"
|
||||
assert "buckets" in data
|
||||
assert isinstance(data["buckets"], list)
|
||||
|
||||
# Should have at least one bucket with our operations
|
||||
assert len(data["buckets"]) >= 1
|
||||
bucket = data["buckets"][0]
|
||||
assert "time" in bucket
|
||||
assert "actions" in bucket
|
||||
assert "total" in bucket
|
||||
assert bucket["total"] >= 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_audit_log_stats_filter_by_action(audit_api_client, bank_id):
|
||||
"""Test stats endpoint filters by action."""
|
||||
await audit_api_client.put(
|
||||
f"/v1/default/banks/{bank_id}",
|
||||
json={"name": "Audit Test Bank"},
|
||||
)
|
||||
|
||||
await audit_api_client.post(
|
||||
f"/v1/default/banks/{bank_id}/memories/recall",
|
||||
json={"query": "test"},
|
||||
)
|
||||
|
||||
await asyncio.sleep(1.0)
|
||||
|
||||
# Stats filtered by recall
|
||||
response = await audit_api_client.get(
|
||||
f"/v1/default/banks/{bank_id}/audit-logs/stats",
|
||||
params={"period": "1d", "action": "recall"},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
for bucket in data["buckets"]:
|
||||
# All actions in buckets should be "recall" only
|
||||
for action_name in bucket["actions"]:
|
||||
assert action_name == "recall"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_audit_log_stats_periods(audit_api_client, bank_id):
|
||||
"""Test stats endpoint supports different periods."""
|
||||
await audit_api_client.put(
|
||||
f"/v1/default/banks/{bank_id}",
|
||||
json={"name": "Audit Test Bank"},
|
||||
)
|
||||
|
||||
await asyncio.sleep(0.5)
|
||||
|
||||
for period, expected_trunc in [("1d", "day"), ("7d", "day"), ("30d", "day")]:
|
||||
response = await audit_api_client.get(
|
||||
f"/v1/default/banks/{bank_id}/audit-logs/stats",
|
||||
params={"period": period},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["period"] == period
|
||||
assert data["trunc"] == expected_trunc
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_audit_log_disabled(memory):
|
||||
"""Test that no audit logs are created when audit logging is disabled."""
|
||||
# Ensure audit logging is disabled
|
||||
memory._audit_logger._enabled = False
|
||||
|
||||
app = create_app(memory, initialize_memory=False)
|
||||
transport = httpx.ASGITransport(app=app)
|
||||
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
|
||||
from datetime import datetime
|
||||
|
||||
bid = f"audit_disabled_test_{datetime.now().timestamp()}"
|
||||
|
||||
await client.put(f"/v1/default/banks/{bid}", json={"name": "No Audit"})
|
||||
await client.post(
|
||||
f"/v1/default/banks/{bid}/memories/recall",
|
||||
json={"query": "test"},
|
||||
)
|
||||
|
||||
await asyncio.sleep(0.5)
|
||||
|
||||
response = await client.get(f"/v1/default/banks/{bid}/audit-logs")
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["total"] == 0, "No audit entries should exist when audit logging is disabled"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_audit_log_action_allowlist(memory):
|
||||
"""Test that only allowed actions are audited when allowlist is set."""
|
||||
memory._audit_logger._enabled = True
|
||||
memory._audit_logger._allowed_actions = frozenset({"recall"}) # Only audit recall
|
||||
|
||||
app = create_app(memory, initialize_memory=False)
|
||||
transport = httpx.ASGITransport(app=app)
|
||||
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
|
||||
from datetime import datetime
|
||||
|
||||
bid = f"audit_allowlist_test_{datetime.now().timestamp()}"
|
||||
|
||||
# create_bank should NOT be audited
|
||||
await client.put(f"/v1/default/banks/{bid}", json={"name": "Allowlist Test"})
|
||||
# recall should be audited
|
||||
await client.post(
|
||||
f"/v1/default/banks/{bid}/memories/recall",
|
||||
json={"query": "allowlist test"},
|
||||
)
|
||||
|
||||
await asyncio.sleep(1.0)
|
||||
|
||||
response = await client.get(f"/v1/default/banks/{bid}/audit-logs")
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
actions = [item["action"] for item in data["items"]]
|
||||
assert "recall" in actions, "recall should be audited"
|
||||
assert "create_bank" not in actions, "create_bank should NOT be audited (not in allowlist)"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_audit_log_ordered_by_most_recent(audit_api_client, bank_id):
|
||||
"""Test that audit logs are returned ordered by most recent first."""
|
||||
await audit_api_client.put(
|
||||
f"/v1/default/banks/{bank_id}",
|
||||
json={"name": "Order Test Bank"},
|
||||
)
|
||||
|
||||
for i in range(3):
|
||||
await audit_api_client.post(
|
||||
f"/v1/default/banks/{bank_id}/memories/recall",
|
||||
json={"query": f"order test {i}"},
|
||||
)
|
||||
await asyncio.sleep(0.2) # Small gap between requests
|
||||
|
||||
await asyncio.sleep(1.0)
|
||||
|
||||
response = await audit_api_client.get(f"/v1/default/banks/{bank_id}/audit-logs")
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
|
||||
# Check descending order by started_at
|
||||
timestamps = [item["started_at"] for item in data["items"] if item["started_at"]]
|
||||
assert timestamps == sorted(timestamps, reverse=True), "Audit logs should be ordered most recent first"
|
||||
@@ -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,842 +0,0 @@
|
||||
"""
|
||||
Tests for delta retain — upsert optimization that only re-processes changed chunks.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from datetime import datetime, timezone
|
||||
|
||||
import pytest
|
||||
|
||||
from hindsight_api import RequestContext
|
||||
from hindsight_api.engine.memory_engine import Budget
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _ts():
|
||||
return datetime.now(timezone.utc).timestamp()
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Core Delta Retain Tests
|
||||
# ============================================================
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delta_retain_unchanged_content_skips_llm(memory, request_context):
|
||||
"""
|
||||
When upserting a document with identical content, no new facts should be
|
||||
extracted (LLM is not called for unchanged chunks). The existing facts
|
||||
should be preserved.
|
||||
"""
|
||||
bank_id = f"test_delta_unchanged_{_ts()}"
|
||||
document_id = "conversation-001"
|
||||
|
||||
try:
|
||||
content = "Alice works at Google. Bob works at Microsoft."
|
||||
|
||||
# First retain — full processing
|
||||
v1_units = await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content=content,
|
||||
context="team info",
|
||||
document_id=document_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
assert len(v1_units) > 0, "v1 should create facts"
|
||||
|
||||
# Get v1 document state
|
||||
doc_v1 = await memory.get_document(document_id, bank_id, request_context=request_context)
|
||||
v1_unit_count = doc_v1["memory_unit_count"]
|
||||
|
||||
# Second retain — same content, should use delta path (no new facts)
|
||||
v2_units = await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content=content,
|
||||
context="team info",
|
||||
document_id=document_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# No new units should be returned (nothing changed)
|
||||
assert v2_units == [], "Delta retain with unchanged content should return empty unit list"
|
||||
|
||||
# Existing facts should still be there
|
||||
doc_v2 = await memory.get_document(document_id, bank_id, request_context=request_context)
|
||||
assert doc_v2["memory_unit_count"] == v1_unit_count, "Existing facts should be preserved"
|
||||
|
||||
# Verify recall still works
|
||||
result = await memory.recall_async(
|
||||
bank_id=bank_id,
|
||||
query="Where does Alice work?",
|
||||
budget=Budget.MID,
|
||||
max_tokens=1000,
|
||||
request_context=request_context,
|
||||
)
|
||||
assert len(result.results) > 0, "Should still recall facts after delta retain"
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delta_retain_appended_content(memory, request_context):
|
||||
"""
|
||||
When a conversation grows (new content appended), only new chunks should
|
||||
be processed. Facts from unchanged chunks should be preserved.
|
||||
"""
|
||||
bank_id = f"test_delta_append_{_ts()}"
|
||||
document_id = "growing-conversation"
|
||||
|
||||
try:
|
||||
# First version — short content (single chunk)
|
||||
v1_content = "Alice is a software engineer at Google. She works on search infrastructure."
|
||||
|
||||
v1_units = await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content=v1_content,
|
||||
context="profile",
|
||||
document_id=document_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
assert len(v1_units) > 0
|
||||
|
||||
# Get v1 facts via recall
|
||||
v1_recall = await memory.recall_async(
|
||||
bank_id=bank_id,
|
||||
query="What does Alice do?",
|
||||
budget=Budget.MID,
|
||||
max_tokens=2000,
|
||||
request_context=request_context,
|
||||
)
|
||||
v1_fact_texts = {r.text for r in v1_recall.results}
|
||||
|
||||
# Second version — original content + new content appended
|
||||
# This should preserve facts from the first chunk and add new ones
|
||||
v2_content = v1_content + "\n\nBob joined Google as a product manager in 2024. He previously worked at Meta on AR/VR products."
|
||||
|
||||
v2_units = await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content=v2_content,
|
||||
context="profile",
|
||||
document_id=document_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Should have facts about Bob from the new content
|
||||
v2_recall = await memory.recall_async(
|
||||
bank_id=bank_id,
|
||||
query="What does Bob do?",
|
||||
budget=Budget.MID,
|
||||
max_tokens=2000,
|
||||
request_context=request_context,
|
||||
)
|
||||
bob_facts = [r for r in v2_recall.results if "bob" in r.text.lower()]
|
||||
assert len(bob_facts) > 0, "Should have facts about Bob from appended content"
|
||||
|
||||
# Should still have facts about Alice from original content
|
||||
alice_recall = await memory.recall_async(
|
||||
bank_id=bank_id,
|
||||
query="What does Alice do?",
|
||||
budget=Budget.MID,
|
||||
max_tokens=2000,
|
||||
request_context=request_context,
|
||||
)
|
||||
assert len(alice_recall.results) > 0, "Should still have Alice facts from original content"
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delta_retain_modified_chunk(memory, request_context):
|
||||
"""
|
||||
When content in the middle changes, that chunk should be re-processed
|
||||
while other chunks are preserved.
|
||||
"""
|
||||
bank_id = f"test_delta_modified_{_ts()}"
|
||||
document_id = "changing-doc"
|
||||
|
||||
try:
|
||||
# v1: Alice works at Google
|
||||
v1_content = "Alice works at Google as a senior engineer."
|
||||
v1_units = await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content=v1_content,
|
||||
context="team",
|
||||
document_id=document_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
assert len(v1_units) > 0
|
||||
|
||||
# v2: Alice works at Microsoft (changed)
|
||||
v2_content = "Alice works at Microsoft as a principal engineer."
|
||||
v2_units = await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content=v2_content,
|
||||
context="team",
|
||||
document_id=document_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# New facts should reflect the updated content
|
||||
result = await memory.recall_async(
|
||||
bank_id=bank_id,
|
||||
query="Where does Alice work?",
|
||||
budget=Budget.MID,
|
||||
max_tokens=2000,
|
||||
request_context=request_context,
|
||||
)
|
||||
all_texts = " ".join(r.text.lower() for r in result.results)
|
||||
assert "microsoft" in all_texts, f"Should have updated fact about Microsoft, got: {all_texts}"
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Entity & Link Tests
|
||||
# ============================================================
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delta_retain_entities_preserved_for_unchanged_chunks(memory, request_context):
|
||||
"""
|
||||
Entities linked to unchanged chunks should be preserved after delta retain.
|
||||
"""
|
||||
bank_id = f"test_delta_entities_{_ts()}"
|
||||
document_id = "entity-doc"
|
||||
|
||||
try:
|
||||
v1_content = "Alice works at Google. She is a senior engineer in the Cloud division."
|
||||
v1_units = await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content=v1_content,
|
||||
context="team",
|
||||
document_id=document_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
assert len(v1_units) > 0
|
||||
|
||||
# Check entities exist
|
||||
pool = await memory._get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
v1_entities = await conn.fetch(
|
||||
"SELECT canonical_name FROM entities WHERE bank_id = $1",
|
||||
bank_id,
|
||||
)
|
||||
v1_entity_names = {e["canonical_name"].lower() for e in v1_entities}
|
||||
assert len(v1_entity_names) > 0, "Should have entities after v1 retain"
|
||||
|
||||
# Upsert with same content — entities should persist
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content=v1_content,
|
||||
context="team",
|
||||
document_id=document_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
async with pool.acquire() as conn:
|
||||
v2_entities = await conn.fetch(
|
||||
"SELECT canonical_name FROM entities WHERE bank_id = $1",
|
||||
bank_id,
|
||||
)
|
||||
v2_entity_names = {e["canonical_name"].lower() for e in v2_entities}
|
||||
|
||||
# All v1 entities should still exist
|
||||
assert v1_entity_names.issubset(v2_entity_names), (
|
||||
f"v1 entities {v1_entity_names} should be preserved, got {v2_entity_names}"
|
||||
)
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delta_retain_new_entities_created_for_new_chunks(memory, request_context):
|
||||
"""
|
||||
New entities should be created for newly added chunks during delta retain.
|
||||
"""
|
||||
bank_id = f"test_delta_new_entities_{_ts()}"
|
||||
document_id = "entity-growth-doc"
|
||||
|
||||
try:
|
||||
v1_content = "Alice works at Google."
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content=v1_content,
|
||||
context="team",
|
||||
document_id=document_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
pool = await memory._get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
v1_entities = await conn.fetch(
|
||||
"SELECT canonical_name FROM entities WHERE bank_id = $1",
|
||||
bank_id,
|
||||
)
|
||||
v1_entity_names = {e["canonical_name"].lower() for e in v1_entities}
|
||||
|
||||
# Append content mentioning new entities
|
||||
v2_content = v1_content + "\n\nBob joined Facebook. He works with Charlie on the Reality Labs project."
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content=v2_content,
|
||||
context="team",
|
||||
document_id=document_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
async with pool.acquire() as conn:
|
||||
v2_entities = await conn.fetch(
|
||||
"SELECT canonical_name FROM entities WHERE bank_id = $1",
|
||||
bank_id,
|
||||
)
|
||||
v2_entity_names = {e["canonical_name"].lower() for e in v2_entities}
|
||||
|
||||
# Should have more entities after adding content with new people/orgs
|
||||
assert len(v2_entity_names) > len(v1_entity_names), (
|
||||
f"Should have more entities after append: v1={v1_entity_names}, v2={v2_entity_names}"
|
||||
)
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delta_retain_links_preserved_for_unchanged_chunks(memory, request_context):
|
||||
"""
|
||||
Memory links (temporal, semantic, entity) for unchanged chunks should be preserved.
|
||||
"""
|
||||
bank_id = f"test_delta_links_{_ts()}"
|
||||
document_id = "links-doc"
|
||||
|
||||
try:
|
||||
content = "Alice is a senior engineer at Google Cloud. She mentors junior engineers and reviews their code."
|
||||
v1_units = await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content=content,
|
||||
context="team",
|
||||
document_id=document_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
assert len(v1_units) > 0
|
||||
|
||||
# Count links after v1
|
||||
pool = await memory._get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
v1_link_count = await conn.fetchval(
|
||||
"""SELECT COUNT(*) FROM memory_links ml
|
||||
JOIN memory_units mu ON ml.from_unit_id = mu.id
|
||||
WHERE mu.bank_id = $1 AND mu.document_id = $2""",
|
||||
bank_id,
|
||||
document_id,
|
||||
)
|
||||
|
||||
# Upsert with same content
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content=content,
|
||||
context="team",
|
||||
document_id=document_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Links should be preserved
|
||||
async with pool.acquire() as conn:
|
||||
v2_link_count = await conn.fetchval(
|
||||
"""SELECT COUNT(*) FROM memory_links ml
|
||||
JOIN memory_units mu ON ml.from_unit_id = mu.id
|
||||
WHERE mu.bank_id = $1 AND mu.document_id = $2""",
|
||||
bank_id,
|
||||
document_id,
|
||||
)
|
||||
|
||||
assert v2_link_count == v1_link_count, (
|
||||
f"Links should be preserved: v1={v1_link_count}, v2={v2_link_count}"
|
||||
)
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Document Metadata & Tags Tests
|
||||
# ============================================================
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delta_retain_document_metadata_updated(memory, request_context):
|
||||
"""
|
||||
Document metadata (retain_params, tags) should be updated even when
|
||||
chunk content hasn't changed.
|
||||
"""
|
||||
bank_id = f"test_delta_meta_{_ts()}"
|
||||
document_id = "metadata-doc"
|
||||
|
||||
try:
|
||||
content = "Alice works at Google."
|
||||
|
||||
# v1 with initial tags
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content=content,
|
||||
context="initial context",
|
||||
document_id=document_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
doc_v1 = await memory.get_document(document_id, bank_id, request_context=request_context)
|
||||
assert doc_v1 is not None
|
||||
|
||||
# v2 with updated context (same content — triggers delta path)
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content=content,
|
||||
context="updated context",
|
||||
document_id=document_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
doc_v2 = await memory.get_document(document_id, bank_id, request_context=request_context)
|
||||
assert doc_v2 is not None
|
||||
assert doc_v2["updated_at"] >= doc_v1["updated_at"], "Document should have updated timestamp"
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delta_retain_tags_propagated_to_existing_units(memory, request_context):
|
||||
"""
|
||||
When tags change during an upsert with unchanged content, the new tags
|
||||
should be propagated to all existing memory units.
|
||||
"""
|
||||
bank_id = f"test_delta_tags_{_ts()}"
|
||||
document_id = "tags-doc"
|
||||
|
||||
try:
|
||||
content = "Alice works at Google."
|
||||
|
||||
# v1 with tag "team-a"
|
||||
await memory.retain_batch_async(
|
||||
bank_id=bank_id,
|
||||
contents=[{
|
||||
"content": content,
|
||||
"document_id": document_id,
|
||||
"tags": ["team-a"],
|
||||
}],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
pool = await memory._get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
v1_tags = await conn.fetch(
|
||||
"SELECT tags FROM memory_units WHERE bank_id = $1 AND document_id = $2",
|
||||
bank_id,
|
||||
document_id,
|
||||
)
|
||||
assert all("team-a" in row["tags"] for row in v1_tags), "v1 units should have team-a tag"
|
||||
|
||||
# v2 with same content but different tags
|
||||
await memory.retain_batch_async(
|
||||
bank_id=bank_id,
|
||||
contents=[{
|
||||
"content": content,
|
||||
"document_id": document_id,
|
||||
"tags": ["team-b", "important"],
|
||||
}],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
async with pool.acquire() as conn:
|
||||
v2_tags = await conn.fetch(
|
||||
"SELECT tags FROM memory_units WHERE bank_id = $1 AND document_id = $2",
|
||||
bank_id,
|
||||
document_id,
|
||||
)
|
||||
for row in v2_tags:
|
||||
assert "team-b" in row["tags"], f"v2 units should have team-b tag, got {row['tags']}"
|
||||
assert "important" in row["tags"], f"v2 units should have important tag, got {row['tags']}"
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Chunk Management Tests
|
||||
# ============================================================
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delta_retain_removed_chunks_delete_facts(memory, request_context):
|
||||
"""
|
||||
When content is shortened (chunks removed), facts from the removed
|
||||
chunks should be deleted.
|
||||
"""
|
||||
bank_id = f"test_delta_removed_{_ts()}"
|
||||
document_id = "shrinking-doc"
|
||||
|
||||
try:
|
||||
# v1: longer content with facts about Alice and Bob
|
||||
v1_content = (
|
||||
"Alice is a senior engineer at Google Cloud. "
|
||||
"She leads the infrastructure team and has been there for 5 years.\n\n"
|
||||
"Bob is a product manager at Facebook Reality Labs. "
|
||||
"He previously worked at Amazon on Alexa voice products."
|
||||
)
|
||||
|
||||
v1_units = await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content=v1_content,
|
||||
context="profiles",
|
||||
document_id=document_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
assert len(v1_units) > 0
|
||||
|
||||
doc_v1 = await memory.get_document(document_id, bank_id, request_context=request_context)
|
||||
v1_count = doc_v1["memory_unit_count"]
|
||||
|
||||
# v2: Completely different content — all chunks change
|
||||
v2_content = "Charlie works at Netflix as a data scientist."
|
||||
v2_units = await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content=v2_content,
|
||||
context="profiles",
|
||||
document_id=document_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
doc_v2 = await memory.get_document(document_id, bank_id, request_context=request_context)
|
||||
assert doc_v2 is not None
|
||||
|
||||
# Should have facts about Charlie
|
||||
result = await memory.recall_async(
|
||||
bank_id=bank_id,
|
||||
query="Who works at Netflix?",
|
||||
budget=Budget.MID,
|
||||
max_tokens=2000,
|
||||
request_context=request_context,
|
||||
)
|
||||
all_texts = " ".join(r.text.lower() for r in result.results)
|
||||
assert "charlie" in all_texts or "netflix" in all_texts, (
|
||||
f"Should have facts about Charlie/Netflix after replacing content, got: {all_texts}"
|
||||
)
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delta_retain_chunks_have_content_hash(memory, request_context):
|
||||
"""
|
||||
After retain, chunks should have content_hash populated.
|
||||
"""
|
||||
bank_id = f"test_delta_hash_{_ts()}"
|
||||
document_id = "hash-doc"
|
||||
|
||||
try:
|
||||
content = "Alice works at Google as a software engineer."
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content=content,
|
||||
document_id=document_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
pool = await memory._get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
chunks = await conn.fetch(
|
||||
"SELECT chunk_id, content_hash FROM chunks WHERE document_id = $1 AND bank_id = $2",
|
||||
document_id,
|
||||
bank_id,
|
||||
)
|
||||
|
||||
assert len(chunks) > 0, "Should have stored chunks"
|
||||
for chunk in chunks:
|
||||
assert chunk["content_hash"] is not None, f"Chunk {chunk['chunk_id']} should have content_hash"
|
||||
assert len(chunk["content_hash"]) == 64, "content_hash should be SHA256 hex (64 chars)"
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Backward Compatibility Tests
|
||||
# ============================================================
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retain_without_document_id_still_works(memory, request_context):
|
||||
"""
|
||||
Retain without document_id should still work normally (no delta path).
|
||||
"""
|
||||
bank_id = f"test_no_docid_{_ts()}"
|
||||
|
||||
try:
|
||||
units = await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="Alice works at Google.",
|
||||
context="test",
|
||||
request_context=request_context,
|
||||
)
|
||||
assert len(units) > 0, "Should create facts without document_id"
|
||||
|
||||
result = await memory.recall_async(
|
||||
bank_id=bank_id,
|
||||
query="Where does Alice work?",
|
||||
budget=Budget.MID,
|
||||
max_tokens=1000,
|
||||
request_context=request_context,
|
||||
)
|
||||
assert len(result.results) > 0
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delta_retain_first_retain_full_path(memory, request_context):
|
||||
"""
|
||||
First retain of a new document should use the full path (no delta possible).
|
||||
"""
|
||||
bank_id = f"test_first_retain_{_ts()}"
|
||||
document_id = "new-doc"
|
||||
|
||||
try:
|
||||
units = await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="Alice works at Google.",
|
||||
context="test",
|
||||
document_id=document_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
assert len(units) > 0, "First retain should create facts via full path"
|
||||
|
||||
doc = await memory.get_document(document_id, bank_id, request_context=request_context)
|
||||
assert doc is not None
|
||||
assert doc["memory_unit_count"] > 0
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Edge Cases
|
||||
# ============================================================
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delta_retain_empty_to_content(memory, request_context):
|
||||
"""
|
||||
Going from gibberish (zero facts) to real content should work.
|
||||
"""
|
||||
bank_id = f"test_delta_empty_{_ts()}"
|
||||
document_id = "empty-to-content"
|
||||
|
||||
try:
|
||||
# v1: content that probably produces zero facts
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="!!!###$$$%%%",
|
||||
document_id=document_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
doc_v1 = await memory.get_document(document_id, bank_id, request_context=request_context)
|
||||
assert doc_v1 is not None
|
||||
|
||||
# v2: real content
|
||||
v2_units = await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="Alice works at Google as a senior engineer.",
|
||||
document_id=document_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
doc_v2 = await memory.get_document(document_id, bank_id, request_context=request_context)
|
||||
assert doc_v2 is not None
|
||||
assert doc_v2["memory_unit_count"] > 0 or len(v2_units) > 0, "Should have facts after updating with real content"
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delta_retain_multiple_upserts(memory, request_context):
|
||||
"""
|
||||
Multiple sequential upserts should work correctly, with delta optimization
|
||||
kicking in after the first retain.
|
||||
"""
|
||||
bank_id = f"test_delta_multi_{_ts()}"
|
||||
document_id = "multi-upsert"
|
||||
|
||||
try:
|
||||
# v1: initial
|
||||
v1_content = "Alice works at Google."
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content=v1_content,
|
||||
document_id=document_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# v2: same content (delta: no changes)
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content=v1_content,
|
||||
document_id=document_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# v3: append
|
||||
v3_content = v1_content + "\n\nBob works at Microsoft."
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content=v3_content,
|
||||
document_id=document_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# v4: same as v3 (delta: no changes again)
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content=v3_content,
|
||||
document_id=document_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Final check: should have facts about both Alice and Bob
|
||||
result = await memory.recall_async(
|
||||
bank_id=bank_id,
|
||||
query="Who works where?",
|
||||
budget=Budget.MID,
|
||||
max_tokens=2000,
|
||||
request_context=request_context,
|
||||
)
|
||||
all_texts = " ".join(r.text.lower() for r in result.results)
|
||||
assert "alice" in all_texts or "google" in all_texts, f"Should have Alice/Google facts, got: {all_texts}"
|
||||
|
||||
doc = await memory.get_document(document_id, bank_id, request_context=request_context)
|
||||
assert doc is not None
|
||||
assert doc["memory_unit_count"] > 0
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delta_retain_with_user_entities(memory, request_context):
|
||||
"""
|
||||
User-provided entities should work correctly with delta retain.
|
||||
"""
|
||||
bank_id = f"test_delta_user_entities_{_ts()}"
|
||||
document_id = "user-entity-doc"
|
||||
|
||||
try:
|
||||
content = "The project is going well."
|
||||
|
||||
# v1 with user entities
|
||||
await memory.retain_batch_async(
|
||||
bank_id=bank_id,
|
||||
contents=[{
|
||||
"content": content,
|
||||
"document_id": document_id,
|
||||
"entities": [{"text": "Project Alpha", "type": "PROJECT"}],
|
||||
}],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
pool = await memory._get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
v1_entities = await conn.fetch(
|
||||
"SELECT canonical_name FROM entities WHERE bank_id = $1",
|
||||
bank_id,
|
||||
)
|
||||
v1_names = {e["canonical_name"].lower() for e in v1_entities}
|
||||
|
||||
# v2 with additional entity, same content
|
||||
# Note: same content = delta path (no re-extraction)
|
||||
# The user entities for NEW chunks only get processed
|
||||
v2_content = content + "\n\nThe timeline is on track for Q2 delivery."
|
||||
await memory.retain_batch_async(
|
||||
bank_id=bank_id,
|
||||
contents=[{
|
||||
"content": v2_content,
|
||||
"document_id": document_id,
|
||||
"entities": [
|
||||
{"text": "Project Alpha", "type": "PROJECT"},
|
||||
{"text": "Q2 Deadline", "type": "MILESTONE"},
|
||||
],
|
||||
}],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Should have entities from both v1 and v2
|
||||
async with pool.acquire() as conn:
|
||||
v2_entities = await conn.fetch(
|
||||
"SELECT canonical_name FROM entities WHERE bank_id = $1",
|
||||
bank_id,
|
||||
)
|
||||
v2_names = {e["canonical_name"].lower() for e in v2_entities}
|
||||
|
||||
# v1 entities should be preserved
|
||||
assert v1_names.issubset(v2_names), f"v1 entities should be preserved: {v1_names} not in {v2_names}"
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delta_retain_recall_with_chunks(memory, request_context):
|
||||
"""
|
||||
After delta retain, recall with include_chunks should return correct chunk data.
|
||||
"""
|
||||
bank_id = f"test_delta_recall_chunks_{_ts()}"
|
||||
document_id = "recall-chunks-doc"
|
||||
|
||||
try:
|
||||
content = "Alice is a senior engineer at Google Cloud. She designs distributed systems."
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content=content,
|
||||
context="profile",
|
||||
document_id=document_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Upsert with same content (delta: no changes)
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content=content,
|
||||
context="profile",
|
||||
document_id=document_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Recall with chunks
|
||||
result = await memory.recall_async(
|
||||
bank_id=bank_id,
|
||||
query="What does Alice do?",
|
||||
budget=Budget.MID,
|
||||
max_tokens=2000,
|
||||
include_chunks=True,
|
||||
max_chunk_tokens=8192,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert len(result.results) > 0, "Should recall facts"
|
||||
|
||||
# Facts with chunk_ids should have corresponding chunks
|
||||
facts_with_chunks = [r for r in result.results if r.chunk_id]
|
||||
if facts_with_chunks and result.chunks:
|
||||
for fact in facts_with_chunks:
|
||||
assert fact.chunk_id in result.chunks, (
|
||||
f"Chunk {fact.chunk_id} should be in returned chunks"
|
||||
)
|
||||
|
||||
finally:
|
||||
await memory.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,114 +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
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Unit tests for discard_pending_stats() — no database required
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_discard_pending_stats_clears_both_dicts():
|
||||
"""discard_pending_stats() must remove entries for the current task from
|
||||
both _pending_stats and _pending_cooccurrences."""
|
||||
resolver = EntityResolver(pool=None) # type: ignore[arg-type]
|
||||
key = resolver._task_key()
|
||||
|
||||
resolver._pending_stats[key] = [object()] # type: ignore[list-item]
|
||||
resolver._pending_cooccurrences[key] = [object()] # type: ignore[list-item]
|
||||
|
||||
resolver.discard_pending_stats()
|
||||
|
||||
assert key not in resolver._pending_stats
|
||||
assert key not in resolver._pending_cooccurrences
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_discard_pending_stats_is_idempotent():
|
||||
"""Calling discard_pending_stats() when nothing is pending must not raise."""
|
||||
resolver = EntityResolver(pool=None) # type: ignore[arg-type]
|
||||
resolver.discard_pending_stats()
|
||||
resolver.discard_pending_stats() # second call — still safe
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_discard_pending_stats_does_not_affect_other_task_keys():
|
||||
"""discard_pending_stats() must only remove the current task's entries,
|
||||
leaving entries keyed under other task IDs untouched."""
|
||||
resolver = EntityResolver(pool=None) # type: ignore[arg-type]
|
||||
other_key = -1 # A fake key that can never be a real task id
|
||||
|
||||
resolver._pending_stats[other_key] = [object()] # type: ignore[list-item]
|
||||
resolver._pending_cooccurrences[other_key] = [object()] # type: ignore[list-item]
|
||||
|
||||
resolver.discard_pending_stats() # discards current task's key only
|
||||
|
||||
assert other_key in resolver._pending_stats, "other task's stats must be preserved"
|
||||
assert other_key in resolver._pending_cooccurrences, "other task's cooccurrences must be preserved"
|
||||
|
||||
|
||||
@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,163 +0,0 @@
|
||||
"""
|
||||
Unit tests for EntityResolver pg_trgm auto-detection (PR #626/#649).
|
||||
|
||||
These tests verify:
|
||||
1. When entity_lookup="trigram" and pg_trgm IS available, the trigram path is used.
|
||||
2. When entity_lookup="trigram" and pg_trgm is NOT available, the resolver falls back
|
||||
to entity_lookup="full" and uses the full-scan path.
|
||||
3. The pg_trgm check is only performed once (_pg_trgm_checked flag prevents re-checking).
|
||||
4. When entity_lookup="full" from the start, the trgm check is never performed.
|
||||
"""
|
||||
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from hindsight_api.engine.entity_resolver import EntityResolver
|
||||
|
||||
|
||||
def _make_conn(pg_trgm_available: bool) -> MagicMock:
|
||||
"""Create a minimal mock asyncpg connection for the pg_trgm availability check."""
|
||||
conn = MagicMock()
|
||||
conn.fetchval = AsyncMock(return_value=pg_trgm_available)
|
||||
conn.fetch = AsyncMock(return_value=[])
|
||||
conn.executemany = AsyncMock()
|
||||
conn.fetchrow = AsyncMock(return_value=None)
|
||||
return conn
|
||||
|
||||
|
||||
def _make_resolver(entity_lookup: str = "trigram") -> EntityResolver:
|
||||
"""Return an EntityResolver with a None pool (not needed for unit tests)."""
|
||||
return EntityResolver(pool=None, entity_lookup=entity_lookup) # type: ignore[arg-type]
|
||||
|
||||
|
||||
class TestPgTrgmAutoDetection:
|
||||
"""Unit tests for pg_trgm detection logic inside _resolve_entities_batch_impl."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_falls_back_to_full_when_pg_trgm_unavailable(self):
|
||||
"""When pg_trgm is absent the resolver switches to 'full' and calls the full-scan path."""
|
||||
resolver = _make_resolver(entity_lookup="trigram")
|
||||
conn = _make_conn(pg_trgm_available=False)
|
||||
|
||||
with (
|
||||
patch.object(resolver, "_resolve_entities_batch_full", new=AsyncMock(return_value=[])) as mock_full,
|
||||
patch.object(resolver, "_resolve_entities_batch_trigram", new=AsyncMock(return_value=[])) as mock_trgm,
|
||||
):
|
||||
await resolver._resolve_entities_batch_impl(
|
||||
conn=conn,
|
||||
bank_id="test-bank",
|
||||
entities_data=[],
|
||||
context="",
|
||||
unit_event_date=None,
|
||||
)
|
||||
|
||||
# Trigram path must NOT be called
|
||||
mock_trgm.assert_not_called()
|
||||
# Full-scan path must be called as the fallback
|
||||
mock_full.assert_called_once()
|
||||
# Strategy is permanently downgraded
|
||||
assert resolver.entity_lookup == "full"
|
||||
assert resolver._pg_trgm_checked is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_uses_trigram_when_pg_trgm_available(self):
|
||||
"""When pg_trgm is present the trigram path is used."""
|
||||
resolver = _make_resolver(entity_lookup="trigram")
|
||||
conn = _make_conn(pg_trgm_available=True)
|
||||
|
||||
with (
|
||||
patch.object(resolver, "_resolve_entities_batch_full", new=AsyncMock(return_value=[])) as mock_full,
|
||||
patch.object(resolver, "_resolve_entities_batch_trigram", new=AsyncMock(return_value=[])) as mock_trgm,
|
||||
):
|
||||
await resolver._resolve_entities_batch_impl(
|
||||
conn=conn,
|
||||
bank_id="test-bank",
|
||||
entities_data=[],
|
||||
context="",
|
||||
unit_event_date=None,
|
||||
)
|
||||
|
||||
mock_trgm.assert_called_once()
|
||||
mock_full.assert_not_called()
|
||||
assert resolver.entity_lookup == "trigram"
|
||||
assert resolver._pg_trgm_checked is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pg_trgm_check_performed_only_once(self):
|
||||
"""The fetchval check is only issued on the first call; subsequent calls skip it."""
|
||||
resolver = _make_resolver(entity_lookup="trigram")
|
||||
conn = _make_conn(pg_trgm_available=True)
|
||||
|
||||
with patch.object(resolver, "_resolve_entities_batch_trigram", new=AsyncMock(return_value=[])):
|
||||
# First call — check is issued
|
||||
await resolver._resolve_entities_batch_impl(
|
||||
conn=conn,
|
||||
bank_id="test-bank",
|
||||
entities_data=[],
|
||||
context="",
|
||||
unit_event_date=None,
|
||||
)
|
||||
# Second call — check must NOT be issued again
|
||||
await resolver._resolve_entities_batch_impl(
|
||||
conn=conn,
|
||||
bank_id="test-bank",
|
||||
entities_data=[],
|
||||
context="",
|
||||
unit_event_date=None,
|
||||
)
|
||||
|
||||
# fetchval (the pg_trgm availability query) should be called exactly once
|
||||
assert conn.fetchval.call_count == 1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_full_strategy_skips_pg_trgm_check(self):
|
||||
"""When entity_lookup='full' from the start, no pg_trgm check is ever issued."""
|
||||
resolver = _make_resolver(entity_lookup="full")
|
||||
conn = _make_conn(pg_trgm_available=False)
|
||||
|
||||
with patch.object(resolver, "_resolve_entities_batch_full", new=AsyncMock(return_value=[])):
|
||||
await resolver._resolve_entities_batch_impl(
|
||||
conn=conn,
|
||||
bank_id="test-bank",
|
||||
entities_data=[],
|
||||
context="",
|
||||
unit_event_date=None,
|
||||
)
|
||||
|
||||
# fetchval should never be called when entity_lookup is already "full"
|
||||
conn.fetchval.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fallback_is_sticky_across_calls(self):
|
||||
"""After falling back to 'full', subsequent calls also use the full path."""
|
||||
resolver = _make_resolver(entity_lookup="trigram")
|
||||
conn = _make_conn(pg_trgm_available=False)
|
||||
|
||||
with (
|
||||
patch.object(resolver, "_resolve_entities_batch_full", new=AsyncMock(return_value=[])) as mock_full,
|
||||
patch.object(resolver, "_resolve_entities_batch_trigram", new=AsyncMock(return_value=[])) as mock_trgm,
|
||||
):
|
||||
# First call triggers the fallback
|
||||
await resolver._resolve_entities_batch_impl(
|
||||
conn=conn,
|
||||
bank_id="b",
|
||||
entities_data=[],
|
||||
context="",
|
||||
unit_event_date=None,
|
||||
)
|
||||
# Second call — _pg_trgm_checked is True so no re-check; entity_lookup=="full"
|
||||
await resolver._resolve_entities_batch_impl(
|
||||
conn=conn,
|
||||
bank_id="b",
|
||||
entities_data=[],
|
||||
context="",
|
||||
unit_event_date=None,
|
||||
)
|
||||
|
||||
# Trigram path is never called
|
||||
mock_trgm.assert_not_called()
|
||||
# Full-scan path is called both times
|
||||
assert mock_full.call_count == 2
|
||||
# pg_trgm check was issued exactly once
|
||||
assert conn.fetchval.call_count == 1
|
||||
@@ -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,193 +0,0 @@
|
||||
"""
|
||||
Tests for per-bank HNSW index lifecycle and UNION ALL retrieval.
|
||||
|
||||
Covers:
|
||||
- _hnsw_index_name deterministic naming
|
||||
- Per-bank HNSW indexes created on bank creation (retain_async / ensure_bank_exists)
|
||||
- Per-bank HNSW indexes dropped on bank deletion
|
||||
- retrieve_semantic_bm25_combined groups results correctly by fact_type and source
|
||||
"""
|
||||
import uuid
|
||||
from datetime import datetime, timezone
|
||||
|
||||
import pytest
|
||||
|
||||
from hindsight_api.engine.retain.bank_utils import _HNSW_FACT_TYPES, _hnsw_index_name
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Unit tests — no DB required
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestHnswIndexName:
|
||||
def test_deterministic(self):
|
||||
uid = "550e8400-e29b-41d4-a716-446655440000"
|
||||
assert _hnsw_index_name("world", uid) == _hnsw_index_name("world", uid)
|
||||
|
||||
def test_strips_dashes(self):
|
||||
uid = "550e8400-e29b-41d4-a716-446655440000"
|
||||
name = _hnsw_index_name("world", uid)
|
||||
# uid16 should be hex chars only
|
||||
assert "-" not in name
|
||||
|
||||
def test_uses_first_16_hex_chars(self):
|
||||
uid = "550e8400-e29b-41d4-a716-446655440000"
|
||||
uid16 = uid.replace("-", "")[:16] # "550e8400e29b41d4"
|
||||
assert name_ends_with(name=_hnsw_index_name("world", uid), suffix=uid16)
|
||||
|
||||
def test_suffix_per_fact_type(self):
|
||||
uid = "550e8400-e29b-41d4-a716-446655440000"
|
||||
names = {ft: _hnsw_index_name(ft, uid) for ft in _HNSW_FACT_TYPES}
|
||||
# All three names must be distinct
|
||||
assert len(set(names.values())) == 3
|
||||
|
||||
def test_all_fact_types_covered(self):
|
||||
assert set(_HNSW_FACT_TYPES) == {"world", "experience", "observation"}
|
||||
|
||||
def test_fits_pg_identifier_limit(self):
|
||||
# PostgreSQL max identifier length is 63 chars
|
||||
uid = "f" * 32 # simulated UUID without dashes
|
||||
for ft in _HNSW_FACT_TYPES:
|
||||
assert len(_hnsw_index_name(ft, uid)) <= 63
|
||||
|
||||
|
||||
def name_ends_with(name: str, suffix: str) -> bool:
|
||||
return name.endswith(suffix)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Integration tests — require DB (memory fixture)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def _get_bank_hnsw_indexes(pool, bank_id: str) -> list[str]:
|
||||
"""Return index names for memory_units that match the per-bank pattern."""
|
||||
async with pool.acquire() as conn:
|
||||
rows = await conn.fetch(
|
||||
"""
|
||||
SELECT indexname
|
||||
FROM pg_indexes
|
||||
WHERE tablename = 'memory_units'
|
||||
AND indexname LIKE 'idx_mu_emb_%'
|
||||
AND indexdef LIKE $1
|
||||
ORDER BY indexname
|
||||
""",
|
||||
f"%bank_id = '{bank_id}'%",
|
||||
)
|
||||
return [row["indexname"] for row in rows]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retain_creates_per_bank_hnsw_indexes(memory, request_context):
|
||||
"""retain_async on a new bank must create 3 per-(bank, fact_type) HNSW indexes."""
|
||||
bank_id = f"test_hnsw_create_{uuid.uuid4().hex[:8]}"
|
||||
try:
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="Alice is a software engineer.",
|
||||
request_context=request_context,
|
||||
)
|
||||
indexes = await _get_bank_hnsw_indexes(memory._pool, bank_id)
|
||||
assert len(indexes) == 3, f"Expected 3 per-bank HNSW indexes, got: {indexes}"
|
||||
for ft_short in _HNSW_FACT_TYPES.values():
|
||||
assert any(ft_short in idx for idx in indexes), (
|
||||
f"Missing index for fact_type short '{ft_short}' in {indexes}"
|
||||
)
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_bank_drops_hnsw_indexes(memory, request_context):
|
||||
"""delete_bank must drop all per-bank HNSW indexes."""
|
||||
bank_id = f"test_hnsw_drop_{uuid.uuid4().hex[:8]}"
|
||||
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="Bob is a data scientist.",
|
||||
request_context=request_context,
|
||||
)
|
||||
# Verify indexes exist before deletion
|
||||
indexes_before = await _get_bank_hnsw_indexes(memory._pool, bank_id)
|
||||
assert len(indexes_before) == 3
|
||||
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
indexes_after = await _get_bank_hnsw_indexes(memory._pool, bank_id)
|
||||
assert indexes_after == [], f"Indexes should be dropped after bank deletion, got: {indexes_after}"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retain_idempotent_bank_creation(memory, request_context):
|
||||
"""Retaining into the same bank twice must not error and still have exactly 3 indexes."""
|
||||
bank_id = f"test_hnsw_idem_{uuid.uuid4().hex[:8]}"
|
||||
try:
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="Carol is a product manager.",
|
||||
request_context=request_context,
|
||||
)
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="Carol joined the company in 2022.",
|
||||
request_context=request_context,
|
||||
)
|
||||
indexes = await _get_bank_hnsw_indexes(memory._pool, bank_id)
|
||||
assert len(indexes) == 3
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retrieve_semantic_bm25_grouped_by_fact_type(memory, request_context):
|
||||
"""
|
||||
retrieve_semantic_bm25_combined must return a dict keyed by fact_type with
|
||||
(semantic_list, bm25_list) tuples. All returned facts must belong to their
|
||||
declared fact_type.
|
||||
"""
|
||||
from hindsight_api.engine.search.retrieval import retrieve_semantic_bm25_combined
|
||||
|
||||
bank_id = f"test_retrieval_{uuid.uuid4().hex[:8]}"
|
||||
try:
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content=(
|
||||
"Alice is a software engineer at TechCorp. "
|
||||
"She visited Paris in 2023 for a conference."
|
||||
),
|
||||
context="background",
|
||||
event_date=datetime(2023, 6, 1, tzinfo=timezone.utc),
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
query_emb = memory.embeddings.encode(["software engineer Alice"])
|
||||
query_emb_str = str(query_emb[0])
|
||||
|
||||
fact_types = ["world", "experience"]
|
||||
async with memory._pool.acquire() as conn:
|
||||
results = await retrieve_semantic_bm25_combined(
|
||||
conn=conn,
|
||||
query_emb_str=query_emb_str,
|
||||
query_text="software engineer Alice",
|
||||
bank_id=bank_id,
|
||||
fact_types=fact_types,
|
||||
limit=5,
|
||||
)
|
||||
|
||||
# Must return an entry for every requested fact_type
|
||||
assert set(results.keys()) == set(fact_types)
|
||||
|
||||
for ft, (sem, bm25) in results.items():
|
||||
# Semantic and BM25 lists must be lists
|
||||
assert isinstance(sem, list)
|
||||
assert isinstance(bm25, list)
|
||||
# All semantic results must declare the correct fact_type
|
||||
for r in sem:
|
||||
assert r.fact_type == ft, f"Semantic result has wrong fact_type: {r.fact_type}"
|
||||
# All BM25 results must declare the correct fact_type
|
||||
for r in bm25:
|
||||
assert r.fact_type == ft, f"BM25 result has wrong fact_type: {r.fact_type}"
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
@@ -1,72 +0,0 @@
|
||||
"""
|
||||
Integration tests for the Iris file parser.
|
||||
|
||||
Tests are skipped automatically if HINDSIGHT_API_FILE_PARSER_IRIS_TOKEN
|
||||
and HINDSIGHT_API_FILE_PARSER_IRIS_ORG_ID are not set in the environment.
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
import pytest
|
||||
|
||||
from hindsight_api.config import ENV_FILE_PARSER_IRIS_ORG_ID, ENV_FILE_PARSER_IRIS_TOKEN
|
||||
from hindsight_api.engine.parsers.iris import IrisParser
|
||||
|
||||
_token = os.getenv(ENV_FILE_PARSER_IRIS_TOKEN)
|
||||
_org_id = os.getenv(ENV_FILE_PARSER_IRIS_ORG_ID)
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
not (_token and _org_id),
|
||||
reason="HINDSIGHT_API_FILE_PARSER_IRIS_TOKEN and HINDSIGHT_API_FILE_PARSER_IRIS_ORG_ID not set",
|
||||
)
|
||||
|
||||
# Minimal valid PDF with the text "Hello from Hindsight"
|
||||
_SAMPLE_PDF = 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 (Hello from Hindsight) Tj ET
|
||||
endstream
|
||||
endobj
|
||||
xref
|
||||
0 5
|
||||
0000000000 65535 f
|
||||
0000000009 00000 n
|
||||
0000000058 00000 n
|
||||
0000000115 00000 n
|
||||
0000000274 00000 n
|
||||
trailer << /Size 5 /Root 1 0 R >>
|
||||
startxref
|
||||
369
|
||||
%%EOF"""
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def iris_parser() -> IrisParser:
|
||||
return IrisParser(token=_token, org_id=_org_id)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_iris_parser_converts_pdf(iris_parser: IrisParser):
|
||||
"""IrisParser should extract text from a valid PDF."""
|
||||
result = await iris_parser.convert(_SAMPLE_PDF, "sample.pdf")
|
||||
assert isinstance(result, str)
|
||||
assert len(result) > 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_iris_parser_name(iris_parser: IrisParser):
|
||||
"""IrisParser.name() should return 'iris'."""
|
||||
assert iris_parser.name() == "iris"
|
||||
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user