Compare commits

..
4 Commits
Author SHA1 Message Date
Ben a2039581d8 feat(nemoclaw): add @vectorize-io/hindsight-nemoclaw setup CLI (#622)
* feat(nemoclaw): add hindsight-nemoclaw setup CLI package

Automates the full NemoClaw sandbox setup:
- Installs @vectorize-io/hindsight-openclaw plugin
- Configures external API mode in ~/.openclaw/openclaw.json
- Reads current openshell sandbox policy, merges Hindsight egress rule, re-applies
- Restarts the OpenClaw gateway

Options: --dry-run, --skip-policy, --skip-plugin-install
36 unit tests passing

* docs: add NEMOCLAW.md setup guide

* feat(nemoclaw): add README, docs page, and release pipeline
2026-03-20 14:59:29 +01:00
Ben dd14454e38 docs: NemoClaw persistent memory guide and blog post (#621)
* docs(openclaw): add NemoClaw blog
2026-03-19 17:36:01 -04:00
Nicolò Boschi 6fb8c0570f test(openclaw): export stripMemoryTags/extractRecallQuery and add hook integration tests
- Extract stripMemoryTags and extractRecallQuery as exported pure functions
  from index.ts so hooks share one implementation and tests cover the real code
- Update before_agent_start to call extractRecallQuery; update agent_end to
  call stripMemoryTags instead of duplicating the regex inline
- Rewrite index.test.ts to import the real functions (no more local duplicate)
  and add 11 tests for extractRecallQuery covering all envelope-stripping cases
- Add tests/hooks.integration.test.ts: loads the plugin via mock MoltbotPluginAPI
  in HTTP mode, spies on client.recall/retain, and exercises all hook behaviours:
  excluded providers, short messages, memory injection format, tag stripping,
  transcript formatting, array content blocks, metadata, document_id derivation
2026-02-18 13:50:37 +01:00
Nicolò Boschi 587939f337 fix: improve openclaw test coverage 2026-02-18 13:33:15 +01:00
1010 changed files with 30649 additions and 108441 deletions
+2 -7
View File
@@ -2,10 +2,10 @@
# Copy this file to .env and fill in your values
# LLM Configuration (Required)
# Supported providers: openai, groq, ollama, gemini, anthropic, lmstudio, vertexai, minimax
# Supported providers: openai, groq, ollama, gemini, anthropic, lmstudio, vertexai
HINDSIGHT_API_LLM_PROVIDER=openai
HINDSIGHT_API_LLM_API_KEY=your-api-key-here
HINDSIGHT_API_LLM_MODEL=gpt-4o-mini
HINDSIGHT_API_LLM_MODEL=o3-mini
HINDSIGHT_API_LLM_BASE_URL=https://api.openai.com/v1
# Example: Anthropic Claude configuration
@@ -20,11 +20,6 @@ HINDSIGHT_API_LLM_BASE_URL=https://api.openai.com/v1
# HINDSIGHT_API_LLM_VERTEXAI_REGION=us-central1
# HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY=/path/to/service-account-key.json # Optional, uses ADC if not set
# Example: MiniMax configuration (1M context window)
# HINDSIGHT_API_LLM_PROVIDER=minimax
# HINDSIGHT_API_LLM_API_KEY=your-minimax-api-key
# HINDSIGHT_API_LLM_MODEL=MiniMax-M2.7
# Example: LM Studio local configuration (Qwen 2.5 32B recommended)
# HINDSIGHT_API_LLM_PROVIDER=lmstudio
# HINDSIGHT_API_LLM_API_KEY=lmstudio
-6
View File
@@ -1,6 +0,0 @@
version: 2
updates:
- package-ecosystem: "github-actions"
directory: "/"
schedule:
interval: "weekly"
+4 -7
View File
@@ -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:
-99
View File
@@ -1,99 +0,0 @@
name: Release Integration
on:
push:
tags:
- 'integrations/**'
jobs:
publish:
runs-on: ubuntu-latest
permissions:
id-token: write # for PyPI trusted publishing
steps:
- uses: actions/checkout@v6
- name: Extract integration info
id: info
run: |
# refs/tags/integrations/litellm/v0.1.0 → integration=litellm, version=0.1.0
TAG="${GITHUB_REF#refs/tags/}"
INTEGRATION=$(echo "$TAG" | cut -d'/' -f2)
VERSION=$(echo "$TAG" | cut -d'/' -f3 | sed 's/^v//')
echo "integration=$INTEGRATION" >> $GITHUB_OUTPUT
echo "version=$VERSION" >> $GITHUB_OUTPUT
echo "tag=$TAG" >> $GITHUB_OUTPUT
echo "Integration: $INTEGRATION, Version: $VERSION"
- name: Detect integration type
id: type
run: |
if [ -f "hindsight-integrations/${{ steps.info.outputs.integration }}/pyproject.toml" ]; then
echo "type=python" >> $GITHUB_OUTPUT
else
echo "type=typescript" >> $GITHUB_OUTPUT
fi
# ── Python integrations (litellm, pydantic-ai, crewai) ──────────────────
- name: Install uv
if: steps.type.outputs.type == 'python'
uses: astral-sh/setup-uv@v7
with:
enable-cache: true
- name: Set up Python
if: steps.type.outputs.type == 'python'
uses: actions/setup-python@v5
with:
python-version-file: ".python-version"
- name: Build Python package
if: steps.type.outputs.type == 'python'
working-directory: ./hindsight-integrations/${{ steps.info.outputs.integration }}
run: uv build --out-dir dist
- name: Publish Python package to PyPI
if: steps.type.outputs.type == 'python'
uses: pypa/gh-action-pypi-publish@release/v1
with:
packages-dir: ./hindsight-integrations/${{ steps.info.outputs.integration }}/dist
skip-existing: true
# ── TypeScript integrations (ai-sdk, chat, openclaw) ────────────────────
- name: Set up Node.js
if: steps.type.outputs.type == 'typescript'
uses: actions/setup-node@v6
with:
node-version: '22'
registry-url: 'https://registry.npmjs.org'
- name: Install dependencies
if: steps.type.outputs.type == 'typescript'
working-directory: ./hindsight-integrations/${{ steps.info.outputs.integration }}
run: npm ci
- name: Build TypeScript package
if: steps.type.outputs.type == 'typescript'
working-directory: ./hindsight-integrations/${{ steps.info.outputs.integration }}
run: npm run build
- name: Publish TypeScript package to npm
if: steps.type.outputs.type == 'typescript'
working-directory: ./hindsight-integrations/${{ steps.info.outputs.integration }}
run: |
set +e
OUTPUT=$(npm publish --access public 2>&1)
EXIT_CODE=$?
echo "$OUTPUT"
if [ $EXIT_CODE -ne 0 ]; then
if echo "$OUTPUT" | grep -q "cannot publish over"; then
echo "Package version already published, skipping..."
exit 0
fi
exit $EXIT_CODE
fi
env:
NODE_AUTH_TOKEN: ${{ secrets.NPM_TOKEN }}
+213 -61
View File
@@ -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,168 @@ 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-nemoclaw-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/nemoclaw
run: npm ci
- name: Build
working-directory: ./hindsight-integrations/nemoclaw
run: npm run build
- name: Publish to npm
working-directory: ./hindsight-integrations/nemoclaw
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/nemoclaw
run: npm pack
- name: Upload artifacts
uses: actions/upload-artifact@v4
with:
name: nemoclaw-integration
path: hindsight-integrations/nemoclaw/*.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 +317,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 +339,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 +362,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 +382,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 +423,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 +437,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 +455,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 +471,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 +490,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 +508,7 @@ jobs:
packages: write
steps:
- uses: actions/checkout@v6
- uses: actions/checkout@v4
- name: Install Helm
uses: azure/setup-helm@v4
@@ -399,7 +528,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 +536,73 @@ 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-nemoclaw-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 NemoClaw Integration
uses: actions/download-artifact@v4
with:
name: nemoclaw-integration
path: ./artifacts/nemoclaw-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 +612,18 @@ 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
# NemoClaw Integration
cp artifacts/nemoclaw-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
+229 -750
View File
File diff suppressed because it is too large Load Diff
+19 -19
View File
@@ -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)
@@ -72,7 +72,7 @@ cd hindsight-control-plane && npm run dev
## 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 +81,9 @@ cd hindsight-control-plane && npm run dev
- **hindsight-integrations/**: Framework integrations (LiteLLM, OpenAI)
- **hindsight-dev/**: Development tools and benchmarks
### Core Engine (hindsight-api-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 +101,7 @@ cd hindsight-control-plane && npm run dev
- `fusion.py`: Reciprocal rank fusion for combining results
- `reranking.py`: Cross-encoder reranking
### API Layer (hindsight-api-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 +111,13 @@ Main operations:
- **Reflect**: Disposition-aware reasoning using memories and mental models.
### Database
PostgreSQL with pgvector. Schema managed via Alembic migrations in `hindsight-api-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 +154,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 +251,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 +268,7 @@ Fields must be categorized as either **hierarchical** (can be overridden per-ten
# Static field - just don't add to _HIERARCHICAL_FIELDS
```
2. **main.py** (`hindsight-api-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 +308,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)
+4 -7
View File
@@ -9,9 +9,8 @@
[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
![PyPI - Downloads](https://img.shields.io/pypi/dm/hindsight-api?label=PyPI)
![NPM Downloads](https://img.shields.io/npm/dm/%40vectorize-io%2Fhindsight-client?logoColor=orange&label=NPM&color=blue&link=https%3A%2F%2Fwww.npmjs.com%2Fpackage%2F%40vectorize-io%2Fhindsight-client)
<br/>
<a href="https://trendshift.io/repositories/15603" target="_blank"><img src="https://trendshift.io/api/badge/repositories/15603" alt="vectorize-io%2Fhindsight | Trendshift" style="width: 250px; height: 55px;" width="250" height="55"/></a>
</div>
---
@@ -37,7 +36,7 @@ Hindsight is being used in production at Fortune 500 enterprises and by a growin
## Adding Hindsight to Your AI Agents
The easiest way to use Hindsight with an existing agent is with the LLM Wrapper. You can add memory to your agent with 2 lines of code. That will swap your current LLM client out with the Hindsight wrapper. After that, memories will be stored and retrieved automatically as you make LLM calls.
The easiest way use Hindsight with an existing agent is with the LLM Wrapper. You can add memory to your agent with 2 lines of code. That will swap your current LLM client out with the Hindsight wrapper. After that, memories will be stored and retrieved automatically as you make LLM calls.
If you need more control over how and when your agent stores and recalls memories, there's also a simple API you can integrate with using the SDKs or directly via HTTP.
@@ -70,7 +69,7 @@ docker run --rm -it --pull always -p 8888:8888 -p 9999:9999 \
>API: http://localhost:8888
>UI: http://localhost:9999
You can modify the LLM provider by setting `HINDSIGHT_API_LLM_PROVIDER`. Valid options are `openai`, `anthropic`, `gemini`, `groq`, `ollama`, `lmstudio`, and `minimax`. The documentation provides more details on [supported models](https://hindsight.vectorize.io/developer/models).
You can modify the LLM provider by setting `HINDSIGHT_API_LLM_PROVIDER`. Valid options are `openai`, `anthropic`, `gemini`, `groq`, `ollama`, and `lmstudio`. The documentation provides more details on [supported models](https://hindsight.vectorize.io/developer/models).
@@ -182,7 +181,7 @@ Satisfying these requirements in Hindsight is straightforward. When new user inp
![Overview](./hindsight-docs/static/img/hindsight-overview.webp)
Most agent memory implementations rely on basic vector search or sometimes use a knowledge graph. Hindsight uses biomimetic data structures to organize agent memories in a way that is more like how human memory works:
Most agent memory implementation rely on basic vector search or sometimes use a knowledge graph. Hindsight uses biomimetic data structures to organize agent memories in a way that is more like how human memory works:
- **World:** Facts about the world ("The stove gets hot")
- **Experiences:** Agent's own experiences ("I touched the stove and it really hurt")
@@ -308,5 +307,3 @@ MIT — see [LICENSE](./LICENSE)
---
Built by [Vectorize.io](https://vectorize.io)
<img src="https://umami-pixel.chris-latimer.workers.dev/?id=a8b043e6-6964-454d-80df-69b69d3f0d50&host=github.com&url=/vectorize-io/hindsight" width="1" height="1" alt="" />
Generated
-139
View File
@@ -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"
]
}
}
}
}
}
+13 -20
View File
@@ -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)
+4 -19
View File
@@ -77,32 +77,18 @@ PIDS=()
# Start API if enabled
if [ "$ENABLE_API" = "true" ]; then
cd /app/api
API_HEALTH_URL="${HINDSIGHT_API_HEALTH_URL:-http://localhost:8888/health}"
API_STARTUP_WAIT_SECONDS="${HINDSIGHT_API_STARTUP_WAIT_SECONDS:-300}"
# Run API directly - Python's PYTHONUNBUFFERED=1 handles output buffering
hindsight-api &
API_PID=$!
PIDS+=($API_PID)
# Wait for API to be ready
api_ready=false
for ((i=1; i<=API_STARTUP_WAIT_SECONDS; i++)); do
if ! kill -0 "$API_PID" 2>/dev/null; then
wait "$API_PID"
exit $?
fi
if curl -sf "$API_HEALTH_URL" &>/dev/null; then
api_ready=true
for i in {1..60}; do
if curl -sf http://localhost:8888/health &>/dev/null; then
break
fi
sleep 1
done
if [ "$api_ready" != "true" ]; then
echo "❌ API did not become healthy within ${API_STARTUP_WAIT_SECONDS}s"
exit 1
fi
else
echo "API disabled (HINDSIGHT_ENABLE_API=false)"
fi
@@ -111,8 +97,7 @@ fi
if [ "$ENABLE_CP" = "true" ]; then
echo "🎛️ Starting Control Plane..."
cd /app/control-plane
export HOSTNAME="${HINDSIGHT_CP_HOSTNAME:-0.0.0.0}"
PORT="${HINDSIGHT_CP_PORT:-9999}" node server.js &
PORT=9999 node server.js &
CP_PID=$!
PIDS+=($CP_PID)
else
@@ -125,7 +110,7 @@ echo "✅ Hindsight is running!"
echo ""
echo "📍 Access:"
if [ "$ENABLE_CP" = "true" ]; then
echo " Control Plane: http://localhost:${HINDSIGHT_CP_PORT:-9999}"
echo " Control Plane: http://localhost:9999"
fi
if [ "$ENABLE_API" = "true" ]; then
echo " API: http://localhost:8888"
+10 -44
View File
@@ -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
+9 -5
View File
@@ -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 -2
View File
@@ -2,8 +2,8 @@ apiVersion: v2
name: hindsight
description: Hindsight helm chart
type: application
version: 0.4.19
appVersion: "0.4.19"
version: 0.4.11
appVersion: "0.4.11"
keywords:
- ai
- memory
-33
View File
@@ -1,33 +0,0 @@
[build-system]
requires = ["setuptools>=61"]
build-backend = "setuptools.build_meta"
[project]
name = "hindsight-all-slim"
version = "0.4.19"
description = "Hindsight: Agent Memory That Works Like Human Memory - Slim All-in-One Bundle"
readme = "README.md"
requires-python = ">=3.11"
dependencies = [
"hindsight-api-slim>=0.4.17",
"hindsight-client>=0.0.7",
"hindsight-embed>=0.1.0",
]
[tool.uv.sources]
hindsight-api-slim = { workspace = true }
hindsight-client = { workspace = true }
hindsight-embed = { workspace = true }
[project.optional-dependencies]
test = [
"pytest>=7.0.0",
"pytest-asyncio>=0.21.0",
]
[tool.setuptools]
packages = []
[tool.pytest.ini_options]
asyncio_mode = "auto"
asyncio_default_fixture_loop_scope = "function"
-48
View File
@@ -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
```
-423
View File
@@ -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
-137
View File
@@ -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
@@ -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")
@@ -1,54 +0,0 @@
"""Add GIN index on source_memory_ids for observation lookup performance
Without this index, queries using the array overlap operator (&&) or array
containment (@>) on source_memory_ids require a full sequential scan over all
observation memory_units. At ~77k observations this was measured at 45ms per
query, becoming a bottleneck during consolidation recall (57-64s timeouts) and
user recall (18-27s average).
The GIN index reduces these queries to index scans: 45ms → 0.049ms (927x
speedup). Recall dropped from 18-27s to ~6s, and consolidation recall
stabilised from timeout to ~15s.
Created with CONCURRENTLY so the migration does not block reads or writes.
CONCURRENTLY requires running outside a transaction block, so the migration
emits an explicit COMMIT before the statement and uses IF NOT EXISTS for
idempotency.
Revision ID: a2b3c4d5e6f8
Revises: f7g8h9i0j1k2
Create Date: 2026-03-04
"""
from collections.abc import Sequence
from alembic import context, op
revision: str = "a2b3c4d5e6f8"
down_revision: str | Sequence[str] | None = "f7g8h9i0j1k2"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
def _get_schema_prefix() -> str:
schema = context.config.get_main_option("target_schema")
return f'"{schema}".' if schema else ""
def upgrade() -> None:
schema = _get_schema_prefix()
# CREATE INDEX CONCURRENTLY cannot run inside a transaction block.
# Commit the current Alembic transaction first.
op.execute("COMMIT")
op.execute(
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_units_source_memory_ids "
f"ON {schema}memory_units USING GIN (source_memory_ids) "
f"WHERE source_memory_ids IS NOT NULL"
)
def downgrade() -> None:
schema = _get_schema_prefix()
op.execute("COMMIT")
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_units_source_memory_ids")
@@ -1,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")
@@ -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")
@@ -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
@@ -1,46 +0,0 @@
"""Enable pg_trgm extension and add GIN trigram index on entities.canonical_name
Revision ID: c1a2b3d4e5f6
Revises: b4c5d6e7f8a9
Create Date: 2026-03-02
Index is created CONCURRENTLY so the migration does not block writes on entities
during production deployments. CONCURRENTLY requires running outside a transaction
block; see migrations.py for how this is handled safely.
"""
from collections.abc import Sequence
from alembic import context, op
revision: str = "c1a2b3d4e5f6"
down_revision: str | Sequence[str] | None = "b4c5d6e7f8a9"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
def _get_schema_prefix() -> str:
schema = context.config.get_main_option("target_schema")
return f'"{schema}".' if schema else ""
def upgrade() -> None:
# pg_trgm ships with every standard PostgreSQL installation as a contrib module.
# It enables fast similarity lookups via GIN indexes, used for entity name matching.
op.execute("CREATE EXTENSION IF NOT EXISTS pg_trgm")
schema = _get_schema_prefix()
# GIN index on canonical_name enables sub-millisecond trigram similarity queries
# (% operator, similarity()) instead of full-table scans across all bank entities.
op.execute("COMMIT")
op.execute(
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS entities_canonical_name_trgm_idx "
f"ON {schema}entities USING GIN (canonical_name gin_trgm_ops)"
)
def downgrade() -> None:
schema = _get_schema_prefix()
op.execute("COMMIT")
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}entities_canonical_name_trgm_idx")
# Note: not dropping pg_trgm extension as other indexes may depend on it
@@ -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")
@@ -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")
@@ -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"
)
@@ -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")
@@ -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")
@@ -1,38 +0,0 @@
"""chunk_fk_cascade_delete
Revision ID: f6g7h8i9j0k1
Revises: e5f6g7h8i9j0
Create Date: 2026-03-16 00:00:00.000000
"""
from collections.abc import Sequence
from alembic import op
# revision identifiers, used by Alembic.
revision: str = "f6g7h8i9j0k1"
down_revision: str | Sequence[str] | None = "e5f6g7h8i9j0"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
def upgrade() -> None:
"""Change memory_units.chunk_id FK from SET NULL to CASCADE.
When a document is deleted the CASCADE reaches chunks first; with SET NULL
the memory_units rows survived with chunk_id = NULL, leaving ghost records.
Switching to CASCADE ensures they are removed together with their chunk.
"""
op.drop_constraint("memory_units_chunk_fkey", "memory_units", type_="foreignkey")
op.create_foreign_key(
"memory_units_chunk_fkey", "memory_units", "chunks", ["chunk_id"], ["chunk_id"], ondelete="CASCADE"
)
def downgrade() -> None:
"""Revert to SET NULL behaviour."""
op.drop_constraint("memory_units_chunk_fkey", "memory_units", type_="foreignkey")
op.create_foreign_key(
"memory_units_chunk_fkey", "memory_units", "chunks", ["chunk_id"], ["chunk_id"], ondelete="SET NULL"
)
@@ -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")
@@ -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
@@ -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")
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,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,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,194 +0,0 @@
"""
Entity labels models and helpers for retain pipeline.
Defines a controlled vocabulary of key:value classification labels
(e.g., 'pedagogy:scaffolding', 'interest:active') that are extracted
at retain time and stored as entities.
"""
from typing import Literal
from pydantic import BaseModel, Field, create_model
class LabelValue(BaseModel):
"""A single allowed value for a label group."""
value: str
description: str = ""
class LabelGroup(BaseModel):
"""A label group (dimension) with its type and allowed values."""
key: str
description: str = ""
type: Literal["value", "multi-values", "text"] = "value"
optional: bool = True
tag: bool = False
values: list[LabelValue] = []
class EntityLabelsConfig(BaseModel):
"""Entity labels configuration for a bank (controlled vocabulary)."""
attributes: list[LabelGroup] = []
def parse_entity_labels(raw: dict | list | None) -> EntityLabelsConfig | None:
"""
Parse raw entity labels config into EntityLabelsConfig.
Accepts:
- None → returns None
- list → list of attribute dicts (each may use legacy free_values/multi_value or new type field)
- dict → {attributes: [...]}
Legacy migration (backward-compat):
- free_values=True → type="text"
- multi_value=True → type="multi-values"
- neither / free_values=False → type="value"
Args:
raw: Raw entity labels config from bank config
Returns:
EntityLabelsConfig or None if raw is None/empty
"""
if raw is None:
return None
if isinstance(raw, list):
if not raw:
return None
attributes = [LabelGroup.model_validate(_migrate_label_group(a)) for a in raw]
return EntityLabelsConfig(attributes=attributes)
if isinstance(raw, dict):
attrs_raw = raw.get("attributes", [])
if not attrs_raw:
return None
attributes = [LabelGroup.model_validate(_migrate_label_group(a)) for a in attrs_raw]
return EntityLabelsConfig(attributes=attributes)
return None
def _migrate_label_group(raw: dict) -> dict:
"""Migrate legacy free_values/multi_value fields to the new type field."""
if not isinstance(raw, dict) or "type" in raw:
return raw
patched = dict(raw)
if patched.get("free_values"):
patched["type"] = "text"
elif patched.get("multi_value"):
patched["type"] = "multi-values"
else:
patched["type"] = "value"
# Remove legacy keys so Pydantic doesn't error on unknown fields
patched.pop("free_values", None)
patched.pop("multi_value", None)
return patched
def build_labels_model(labels_cfg: EntityLabelsConfig) -> type[BaseModel] | None:
"""
Build a dynamic Pydantic model for structured label extraction.
Each LabelGroup becomes a typed field based on its type:
- type="text" → str | None (always optional)
- type="value", optional=True → Literal["v1","v2"] | None
- type="value", optional=False → Literal["v1","v2"] (required)
- type="multi-values" → list[Literal["v1","v2"]]
Args:
labels_cfg: Parsed EntityLabelsConfig
Returns:
Dynamic Pydantic model class, or None if no groups defined
"""
fields: dict = {}
for group in labels_cfg.attributes:
if not group.key:
continue
description = group.description or group.key
if group.type == "text":
# Free-form: any string value accepted, always optional
fields[group.key] = (str | None, Field(default=None, description=description))
else:
# Enum-constrained: must have defined values
if not group.values:
continue
values = tuple(v.value for v in group.values if v.value)
if not values:
continue
# Literal[("v1", "v2")] is equivalent to Literal["v1", "v2"] in Python 3.11+
literal_type = Literal[values] # type: ignore[valid-type]
if group.type == "multi-values":
fields[group.key] = (
list[literal_type], # type: ignore[valid-type]
Field(default_factory=list, description=description),
)
elif group.optional:
fields[group.key] = (
literal_type | None, # type: ignore[valid-type]
Field(default=None, description=description),
)
else:
fields[group.key] = (
literal_type, # type: ignore[valid-type]
Field(description=description),
)
if not fields:
return None
return create_model("Labels", **fields)
def is_label_entity(text: str, labels_cfg: EntityLabelsConfig, labels_lookup: set[str]) -> bool:
"""
Return True if entity text belongs to any configured label group.
For enum groups: checks the pre-built lookup set.
For text groups: checks that the text starts with a known key prefix.
"""
if text.lower() in labels_lookup:
return True
for group in labels_cfg.attributes:
if group.type == "text" and group.key and text.lower().startswith(f"{group.key.lower()}:"):
return True
return False
def build_labels_lookup(labels_cfg: EntityLabelsConfig | list | None) -> set[str]:
"""
Build a set of valid 'key:value' label strings (lowercase) for fast lookup.
Accepts either EntityLabelsConfig or raw list/None for backwards compatibility.
Args:
labels_cfg: EntityLabelsConfig, raw list of attribute dicts, or None
Returns:
Set of lowercase 'key:value' strings
"""
if labels_cfg is None:
return set()
# Accept raw list/dict for backwards compatibility
if not isinstance(labels_cfg, EntityLabelsConfig):
parsed = parse_entity_labels(labels_cfg)
if parsed is None:
return set()
labels_cfg = parsed
valid = set()
for group in labels_cfg.attributes:
if group.type == "text":
continue # No fixed vocabulary — all values accepted in post-processing
for v in group.values:
if group.key and v.value:
valid.add(f"{group.key}:{v.value}".lower())
return valid
@@ -1,544 +0,0 @@
"""
Main orchestrator for the retain pipeline.
Coordinates all retain pipeline modules to store memories efficiently.
"""
import logging
import time
import uuid
from collections.abc import Awaitable, Callable
from datetime import UTC, datetime
from typing import Any
from ..db_utils import acquire_with_retry, retry_with_backoff
from . import bank_utils
def utcnow():
"""Get current UTC time."""
return datetime.now(UTC)
def parse_datetime_flexible(value: Any) -> datetime:
"""
Parse a datetime value that could be either a datetime object or an ISO string.
This handles datetime values from both direct Python calls and deserialized JSON
(where datetime objects are serialized as ISO strings).
Args:
value: Either a datetime object or an ISO format string
Returns:
datetime object (timezone-aware)
Raises:
TypeError: If value is neither datetime nor string
ValueError: If string is not a valid ISO datetime
"""
if isinstance(value, datetime):
# Ensure timezone-aware
if value.tzinfo is None:
return value.replace(tzinfo=UTC)
return value
elif isinstance(value, str):
# Parse ISO format string (handles both 'Z' and '+00:00' timezone formats)
dt = datetime.fromisoformat(value.replace("Z", "+00:00"))
# Ensure timezone-aware
if dt.tzinfo is None:
return dt.replace(tzinfo=UTC)
return dt
else:
raise TypeError(f"Expected datetime or string, got {type(value).__name__}")
import asyncpg
from ..response_models import TokenUsage
from . import (
chunk_storage,
embedding_processing,
entity_processing,
fact_extraction,
fact_storage,
link_creation,
)
from .types import EntityLink, ExtractedFact, ProcessedFact, RetainContent, RetainContentDict
logger = logging.getLogger(__name__)
async def retain_batch(
pool,
embeddings_model,
llm_config,
entity_resolver,
format_date_fn,
bank_id: str,
contents_dicts: list[RetainContentDict],
config,
document_id: str | None = None,
is_first_batch: bool = True,
fact_type_override: str | None = None,
confidence_score: float | None = None,
document_tags: list[str] | None = None,
operation_id: str | None = None,
schema: str | None = None,
outbox_callback: Callable[["asyncpg.Connection"], Awaitable[None]] | None = None,
) -> tuple[list[list[str]], TokenUsage]:
"""
Process a batch of content through the retain pipeline.
Args:
pool: Database connection pool
embeddings_model: Embeddings model for generating embeddings
llm_config: LLM configuration for fact extraction
entity_resolver: Entity resolver for entity processing
format_date_fn: Function to format datetime to readable string
bank_id: Bank identifier
contents_dicts: List of content dictionaries
config: Resolved HindsightConfig for this bank
document_id: Optional document ID
is_first_batch: Whether this is the first batch
fact_type_override: Override fact type for all facts
confidence_score: Confidence score for opinions
document_tags: Tags applied to all items in this batch
Returns:
Tuple of (unit ID lists, token usage for fact extraction)
"""
start_time = time.time()
total_chars = sum(len(item.get("content", "")) for item in contents_dicts)
# Buffer all logs
log_buffer = []
log_buffer.append(f"{'=' * 60}")
log_buffer.append(f"RETAIN_BATCH START: {bank_id}")
log_buffer.append(f"Batch size: {len(contents_dicts)} content items, {total_chars:,} chars")
log_buffer.append(f"{'=' * 60}")
# Get bank profile
profile = await bank_utils.get_bank_profile(pool, bank_id)
agent_name = profile["name"]
# Convert dicts to RetainContent objects
contents = []
for item in contents_dicts:
# Merge item-level tags with document-level tags
item_tags = item.get("tags", []) or []
merged_tags = list(set(item_tags + (document_tags or [])))
# Handle event_date: distinguish "not provided" (default to now) from
# "explicitly None" (caller opted into no timestamp).
if "event_date" in item and item["event_date"] is None:
event_date_value = None # Caller explicitly signalled "unknown date"
elif item.get("event_date"):
event_date_value = parse_datetime_flexible(item["event_date"])
else:
event_date_value = utcnow() # Backward-compatible default
content = RetainContent(
content=item["content"],
context=item.get("context", ""),
event_date=event_date_value,
metadata=item.get("metadata", {}),
entities=item.get("entities", []),
tags=merged_tags,
observation_scopes=item.get("observation_scopes"),
)
contents.append(content)
# Step 1: Extract facts from all contents
step_start = time.time()
extracted_facts, chunks, usage = await fact_extraction.extract_facts_from_contents(
contents, llm_config, agent_name, config, pool, operation_id, schema
)
log_buffer.append(
f"[1] Extract facts: {len(extracted_facts)} facts, {len(chunks)} chunks from {len(contents)} contents in {time.time() - step_start:.3f}s"
)
if not extracted_facts:
# Still need to create document if document_id was provided or chunks exist
from collections import defaultdict
docs_tracked = 0
async with acquire_with_retry(pool) as conn:
async with conn.transaction():
# Group contents by document_id (consistent with normal path)
contents_by_doc_early = defaultdict(list)
for idx, content_dict in enumerate(contents_dicts):
doc_id = content_dict.get("document_id")
contents_by_doc_early[doc_id].append((idx, content_dict))
if document_id:
# Legacy: single document_id parameter
combined_content = "\n".join([c.get("content", "") for c in contents_dicts])
# Collect tags from all content items and merge with document_tags
all_tags = set(document_tags or [])
for item in contents_dicts:
item_tags = item.get("tags", []) or []
all_tags.update(item_tags)
merged_tags = list(all_tags)
retain_params = {}
if contents_dicts:
first_item = contents_dicts[0]
if first_item.get("context"):
retain_params["context"] = first_item["context"]
if first_item.get("event_date"):
retain_params["event_date"] = (
first_item["event_date"].isoformat()
if hasattr(first_item["event_date"], "isoformat")
else str(first_item["event_date"])
)
if first_item.get("metadata"):
retain_params["metadata"] = first_item["metadata"]
await fact_storage.handle_document_tracking(
conn, bank_id, document_id, combined_content, is_first_batch, retain_params, merged_tags
)
docs_tracked += 1
else:
# Handle per-item document_ids and/or chunks (mirrors normal path logic)
has_any_doc_ids = any(item.get("document_id") for item in contents_dicts)
if has_any_doc_ids or chunks:
for original_doc_id, doc_contents in contents_by_doc_early.items():
should_create_doc = (original_doc_id is not None) or chunks
if not should_create_doc:
continue
actual_doc_id = original_doc_id
if actual_doc_id is None:
# No document_id but have chunks - generate one
actual_doc_id = str(uuid.uuid4())
combined_content = "\n".join([c.get("content", "") for _, c in doc_contents])
all_tags = set(document_tags or [])
for _, item in doc_contents:
item_tags = item.get("tags", []) or []
all_tags.update(item_tags)
merged_tags = list(all_tags)
retain_params = {}
if doc_contents:
first_item = doc_contents[0][1]
if first_item.get("context"):
retain_params["context"] = first_item["context"]
if first_item.get("event_date"):
retain_params["event_date"] = (
first_item["event_date"].isoformat()
if hasattr(first_item["event_date"], "isoformat")
else str(first_item["event_date"])
)
if first_item.get("metadata"):
retain_params["metadata"] = first_item["metadata"]
await fact_storage.handle_document_tracking(
conn,
bank_id,
actual_doc_id,
combined_content,
is_first_batch,
retain_params,
merged_tags,
)
docs_tracked += 1
total_time = time.time() - start_time
doc_status = f"{docs_tracked} document(s) tracked" if docs_tracked > 0 else "no document tracked"
logger.info(
f"RETAIN_BATCH COMPLETE: 0 facts extracted from {len(contents)} contents in {total_time:.3f}s ({doc_status}, no facts)"
)
return [[] for _ in contents], usage
# Apply fact_type_override if provided
if fact_type_override:
for fact in extracted_facts:
fact.fact_type = fact_type_override
# Step 2: Augment texts and generate embeddings
step_start = time.time()
augmented_texts = embedding_processing.augment_texts_with_dates(extracted_facts, format_date_fn)
embeddings = await embedding_processing.generate_embeddings_batch(embeddings_model, augmented_texts)
log_buffer.append(f"[2] Generate embeddings: {len(embeddings)} embeddings in {time.time() - step_start:.3f}s")
# Step 3: Convert to ProcessedFact objects (without chunk_ids yet)
processed_facts = [
ProcessedFact.from_extracted_fact(extracted_fact, embedding)
for extracted_fact, embedding in zip(extracted_facts, embeddings)
]
# Group contents by document_id for document tracking and chunk storage
from collections import defaultdict
contents_by_doc = defaultdict(list)
for idx, content_dict in enumerate(contents_dicts):
doc_id = content_dict.get("document_id")
contents_by_doc[doc_id].append((idx, content_dict))
# Step 4: Database transaction (retried on deadlock)
result_unit_ids: list[list[str]] = []
log_buffer_pre_db = len(log_buffer)
async def _run_db_work() -> None:
nonlocal result_unit_ids
# Reset per-fact mutations and log buffer so each retry attempt starts clean
del log_buffer[log_buffer_pre_db:]
document_ids_added: list[str] = []
for pf in processed_facts:
pf.document_id = None
pf.chunk_id = None
async with acquire_with_retry(pool) as conn:
async with conn.transaction():
# Handle document tracking for all documents
step_start = time.time()
# Map None document_id to generated UUIDs
doc_id_mapping = {} # Maps original doc_id (including None) to actual doc_id used
if document_id:
# Legacy: single document_id parameter
combined_content = "\n".join([c.get("content", "") for c in contents_dicts])
retain_params = {}
# Collect tags from all content items and merge with document_tags
all_tags = set(document_tags or [])
for item in contents_dicts:
item_tags = item.get("tags", []) or []
all_tags.update(item_tags)
merged_tags = list(all_tags)
if contents_dicts:
first_item = contents_dicts[0]
if first_item.get("context"):
retain_params["context"] = first_item["context"]
if first_item.get("event_date"):
retain_params["event_date"] = (
first_item["event_date"].isoformat()
if hasattr(first_item["event_date"], "isoformat")
else str(first_item["event_date"])
)
if first_item.get("metadata"):
retain_params["metadata"] = first_item["metadata"]
await fact_storage.handle_document_tracking(
conn, bank_id, document_id, combined_content, is_first_batch, retain_params, merged_tags
)
document_ids_added.append(document_id)
doc_id_mapping[None] = document_id # For backwards compatibility
else:
# Handle per-item document_ids (create documents if any item has document_id or if chunks exist)
has_any_doc_ids = any(item.get("document_id") for item in contents_dicts)
if has_any_doc_ids or chunks:
for original_doc_id, doc_contents in contents_by_doc.items():
actual_doc_id = original_doc_id
# Only create document record if:
# 1. Item has explicit document_id, OR
# 2. There are chunks (need document for chunk storage)
should_create_doc = (original_doc_id is not None) or chunks
if should_create_doc:
if actual_doc_id is None:
# No document_id but have chunks - generate one
actual_doc_id = str(uuid.uuid4())
# Store mapping for later use
doc_id_mapping[original_doc_id] = actual_doc_id
# Combine content for this document
combined_content = "\n".join([c.get("content", "") for _, c in doc_contents])
# Collect tags from all content items for this document and merge with document_tags
all_tags = set(document_tags or [])
for _, item in doc_contents:
item_tags = item.get("tags", []) or []
all_tags.update(item_tags)
merged_tags = list(all_tags)
# Extract retain params from first content item
retain_params = {}
if doc_contents:
first_item = doc_contents[0][1]
if first_item.get("context"):
retain_params["context"] = first_item["context"]
if first_item.get("event_date"):
retain_params["event_date"] = (
first_item["event_date"].isoformat()
if hasattr(first_item["event_date"], "isoformat")
else str(first_item["event_date"])
)
if first_item.get("metadata"):
retain_params["metadata"] = first_item["metadata"]
await fact_storage.handle_document_tracking(
conn,
bank_id,
actual_doc_id,
combined_content,
is_first_batch,
retain_params,
merged_tags,
)
document_ids_added.append(actual_doc_id)
if document_ids_added:
log_buffer.append(
f"[2.5] Document tracking: {len(document_ids_added)} documents in {time.time() - step_start:.3f}s"
)
# Store chunks and map to facts for all documents
step_start = time.time()
chunk_id_map_by_doc = {} # Maps (doc_id, chunk_index) -> chunk_id
if chunks:
# Group chunks by their source document
chunks_by_doc = defaultdict(list)
for chunk in chunks:
# chunk.content_index tells us which content this chunk came from
original_doc_id = contents_dicts[chunk.content_index].get("document_id")
# Map to actual document_id (handles None -> generated UUID mapping)
actual_doc_id = doc_id_mapping.get(original_doc_id, original_doc_id)
if actual_doc_id is None and document_id:
actual_doc_id = document_id
chunks_by_doc[actual_doc_id].append(chunk)
# Store chunks for each document
for doc_id, doc_chunks in chunks_by_doc.items():
chunk_id_map = await chunk_storage.store_chunks_batch(conn, bank_id, doc_id, doc_chunks)
# Store mapping with document context
for chunk_idx, chunk_id in chunk_id_map.items():
chunk_id_map_by_doc[(doc_id, chunk_idx)] = chunk_id
log_buffer.append(
f"[3] Store chunks: {len(chunks)} chunks for {len(chunks_by_doc)} documents in {time.time() - step_start:.3f}s"
)
# Map chunk_ids and document_ids to facts
for fact, processed_fact in zip(extracted_facts, processed_facts):
# Get the original document_id for this fact's source content
original_doc_id = contents_dicts[fact.content_index].get("document_id")
# Map to actual document_id (handles None -> generated UUID mapping)
actual_doc_id = doc_id_mapping.get(original_doc_id, original_doc_id)
if actual_doc_id is None and document_id:
actual_doc_id = document_id
# Set document_id on the fact
processed_fact.document_id = actual_doc_id
# Map chunk_id if this fact came from a chunk
if fact.chunk_index is not None:
# Look up chunk_id using (doc_id, chunk_index)
chunk_id = chunk_id_map_by_doc.get((actual_doc_id, fact.chunk_index))
if chunk_id:
processed_fact.chunk_id = chunk_id
else:
# No chunks - still need to set document_id on facts
for fact, processed_fact in zip(extracted_facts, processed_facts):
original_doc_id = contents_dicts[fact.content_index].get("document_id")
# Map to actual document_id (handles None -> generated UUID mapping)
actual_doc_id = doc_id_mapping.get(original_doc_id, original_doc_id)
if actual_doc_id is None and document_id:
actual_doc_id = document_id
processed_fact.document_id = actual_doc_id
non_duplicate_facts = processed_facts
# Insert facts (document_id is now stored per-fact)
step_start = time.time()
unit_ids = await fact_storage.insert_facts_batch(conn, bank_id, non_duplicate_facts)
log_buffer.append(f"[5] Insert facts: {len(unit_ids)} units in {time.time() - step_start:.3f}s")
# Process entities
step_start = time.time()
# Build map of content_index -> user entities for merging
user_entities_per_content = {
idx: content.entities for idx, content in enumerate(contents) if content.entities
}
entity_links = await entity_processing.process_entities_batch(
entity_resolver,
conn,
bank_id,
unit_ids,
non_duplicate_facts,
log_buffer,
user_entities_per_content=user_entities_per_content,
entity_labels=getattr(config, "entity_labels", None),
)
log_buffer.append(f"[6] Process entities: {len(entity_links)} links in {time.time() - step_start:.3f}s")
# Create temporal links
step_start = time.time()
temporal_link_count = await link_creation.create_temporal_links_batch(conn, bank_id, unit_ids)
log_buffer.append(f"[7] Temporal links: {temporal_link_count} links in {time.time() - step_start:.3f}s")
# Create semantic links
step_start = time.time()
embeddings_for_links = [fact.embedding for fact in non_duplicate_facts]
semantic_link_count = await link_creation.create_semantic_links_batch(
conn, bank_id, unit_ids, embeddings_for_links
)
log_buffer.append(f"[8] Semantic links: {semantic_link_count} links in {time.time() - step_start:.3f}s")
# Insert entity links
step_start = time.time()
if entity_links:
await entity_processing.insert_entity_links_batch(conn, entity_links)
log_buffer.append(
f"[9] Entity links: {len(entity_links) if entity_links else 0} links in {time.time() - step_start:.3f}s"
)
# Create causal links
step_start = time.time()
causal_link_count = await link_creation.create_causal_links_batch(conn, unit_ids, non_duplicate_facts)
log_buffer.append(f"[10] Causal links: {causal_link_count} links in {time.time() - step_start:.3f}s")
# Map results back to original content items
result_unit_ids = _map_results_to_contents(contents, extracted_facts, unit_ids)
# Transactional outbox: queue any side-effect tasks (e.g. webhook deliveries)
# inside the same transaction so they are atomically committed with the retain data.
if outbox_callback:
await outbox_callback(conn)
# Flush entity stats (mention_count / last_seen) now that the transaction
# has committed. Uses a fresh pool connection — no locks held.
await entity_resolver.flush_pending_stats()
# Log final summary
total_time = time.time() - start_time
log_buffer.append(f"{'=' * 60}")
log_buffer.append(f"RETAIN_BATCH COMPLETE: {len(unit_ids)} units in {total_time:.3f}s")
if document_ids_added:
log_buffer.append(f"Documents: {', '.join(document_ids_added)}")
log_buffer.append(f"{'=' * 60}")
logger.info("\n" + "\n".join(log_buffer) + "\n")
await retry_with_backoff(_run_db_work)
return result_unit_ids, usage
def _map_results_to_contents(
contents: list[RetainContent],
extracted_facts: list[ExtractedFact],
unit_ids: list[str],
) -> list[list[str]]:
"""Map created unit IDs back to original content items."""
facts_by_content: dict[int, list[int]] = {i: [] for i in range(len(contents))}
for i, fact in enumerate(extracted_facts):
facts_by_content[fact.content_index].append(i)
result_unit_ids = []
unit_idx = 0
for content_index in range(len(contents)):
content_unit_ids = []
for _ in facts_by_content[content_index]:
content_unit_ids.append(unit_ids[unit_idx])
unit_idx += 1
result_unit_ids.append(content_unit_ids)
return result_unit_ids
@@ -1,493 +0,0 @@
"""
Link Expansion graph retrieval.
Expands from semantic/temporal seeds through three parallel, first-class signals
stored in memory_links:
1. Entity links — precomputed co-occurrence graph (created at retain time, bounded to
MAX_LINKS_PER_ENTITY per entity). Score = number of distinct shared
entities between the seed set and each candidate.
2. Semantic links — precomputed kNN graph (each new fact linked to its top-5 most
similar existing facts at insert time, similarity >= 0.7). Checked
in both directions since the graph is not symmetric. Score = weight.
3. Causal links — explicit causal chains (causes/caused_by/enables/prevents).
Score = weight + 1.0 (boosted as highest-quality signal).
All three signals are bounded at retain time, so no LATERAL fan-out caps are needed
at query time. Each expansion is a simple aggregation over a small result set.
For non-observation fact types the three expansions are issued as a single CTE query
(one roundtrip, one connection) with a `source` discriminator column so the Python
merge step can apply per-signal score transformations.
"""
import logging
import math
import time
from ..db_utils import acquire_with_retry
from ..memory_engine import fq_table
from .graph_retrieval import GraphRetriever
from .tags import TagGroup, TagsMatch, filter_results_by_tag_groups, filter_results_by_tags
from .types import MPFPTimings, RetrievalResult
logger = logging.getLogger(__name__)
async def _find_semantic_seeds(
conn,
query_embedding_str: str,
bank_id: str,
fact_type: str,
limit: int = 20,
threshold: float = 0.3,
tags: list[str] | None = None,
tags_match: TagsMatch = "any",
tag_groups: list[TagGroup] | None = None,
) -> list[RetrievalResult]:
"""Find semantic seeds via embedding search."""
from .tags import build_tag_groups_where_clause, build_tags_where_clause_simple
tags_clause = build_tags_where_clause_simple(tags, 6, match=tags_match)
tag_groups_param_start = 6 + (1 if tags else 0)
groups_clause, groups_params, _ = build_tag_groups_where_clause(tag_groups, tag_groups_param_start)
params = [query_embedding_str, bank_id, fact_type, threshold, limit]
if tags:
params.append(tags)
params.extend(groups_params)
rows = await conn.fetch(
f"""
SELECT id, text, context, event_date, occurred_start, occurred_end,
mentioned_at, fact_type, document_id, chunk_id, tags,
1 - (embedding <=> $1::vector) AS similarity
FROM {fq_table("memory_units")}
WHERE bank_id = $2
AND embedding IS NOT NULL
AND fact_type = $3
AND (1 - (embedding <=> $1::vector)) >= $4
{tags_clause}
{groups_clause}
ORDER BY embedding <=> $1::vector
LIMIT $5
""",
*params,
)
return [RetrievalResult.from_db_row(dict(r)) for r in rows]
class LinkExpansionRetriever(GraphRetriever):
"""
Graph retrieval via direct link expansion from seeds.
Runs three expansions through precomputed memory_links: entity co-occurrence,
semantic kNN, and causal chains, all bounded at retain time.
For non-observation fact types the three expansions are issued as a single CTE
query (one roundtrip, one connection slot) with a `source` discriminator column.
The Python merge step applies per-signal score transformations.
"""
def __init__(
self,
causal_weight_threshold: float = 0.3,
):
"""
Args:
causal_weight_threshold: Minimum weight for causal links to follow.
"""
self.causal_weight_threshold = causal_weight_threshold
@property
def name(self) -> str:
return "link_expansion"
async def retrieve(
self,
pool,
query_embedding_str: str,
bank_id: str,
fact_type: str,
budget: int,
query_text: str | None = None,
semantic_seeds: list[RetrievalResult] | None = None,
temporal_seeds: list[RetrievalResult] | None = None,
adjacency=None,
tags: list[str] | None = None,
tags_match: TagsMatch = "any",
tag_groups: list[TagGroup] | None = None,
) -> tuple[list[RetrievalResult], MPFPTimings | None]:
"""
Retrieve facts by expanding links from seeds.
Args:
pool: Database connection pool
query_embedding_str: Query embedding as string
bank_id: Memory bank ID
fact_type: Fact type to filter
budget: Maximum results to return
query_text: Original query text (unused)
semantic_seeds: Pre-computed semantic entry points
temporal_seeds: Pre-computed temporal entry points
adjacency: Unused, kept for interface compatibility
tags: Optional list of tags for visibility filtering
Returns:
Tuple of (results, timings)
"""
start_time = time.time()
timings = MPFPTimings(fact_type=fact_type)
async with acquire_with_retry(pool) as conn:
# Find seeds if not provided
if semantic_seeds:
all_seeds = list(semantic_seeds)
else:
seeds_start = time.time()
all_seeds = await _find_semantic_seeds(
conn,
query_embedding_str,
bank_id,
fact_type,
limit=20,
threshold=0.3,
tags=tags,
tags_match=tags_match,
tag_groups=tag_groups,
)
timings.seeds_time = time.time() - seeds_start
logger.debug(
f"[LinkExpansion] Found {len(all_seeds)} semantic seeds for fact_type={fact_type} "
f"(tags={tags}, tags_match={tags_match})"
)
if temporal_seeds:
all_seeds.extend(temporal_seeds)
if not all_seeds:
return [], timings
seed_ids = list({s.id for s in all_seeds})
timings.pattern_count = len(seed_ids)
query_start = time.time()
if fact_type == "observation":
entity_rows, semantic_rows, causal_rows = await self._expand_observations(conn, seed_ids, budget)
else:
entity_rows, semantic_rows, causal_rows = await self._expand_combined(conn, seed_ids, fact_type, budget)
timings.edge_load_time = time.time() - query_start
timings.db_queries = 1
timings.edge_count = len(entity_rows) + len(semantic_rows) + len(causal_rows)
# Merge results with additive intra-score: entity + semantic + causal ∈ [0, 3].
#
# Entity score: tanh(count × 0.5) maps shared-entity count to [0, 1]:
# 1 entity → 0.46, 2 → 0.76, 3 → 0.91, 4 → 0.96 (saturates naturally)
# Semantic score: similarity weight, already ∈ [0.7, 1.0].
# Causal score: link weight, already ∈ [0, 1].
#
# Facts appearing in multiple signals accumulate higher scores, rewarding
# convergent evidence. The outer RRF uses rank position from this sorted list.
entity_scores: dict[str, float] = {}
semantic_scores: dict[str, float] = {}
causal_scores: dict[str, float] = {}
row_map: dict[str, dict] = {}
for row in entity_rows:
fact_id = str(row["id"])
entity_scores[fact_id] = math.tanh(row["score"] * 0.5)
row_map[fact_id] = dict(row)
for row in semantic_rows:
fact_id = str(row["id"])
semantic_scores[fact_id] = max(semantic_scores.get(fact_id, 0.0), row["score"])
row_map.setdefault(fact_id, dict(row))
for row in causal_rows:
fact_id = str(row["id"])
causal_scores[fact_id] = max(causal_scores.get(fact_id, 0.0), row["score"])
row_map.setdefault(fact_id, dict(row))
all_ids = set(entity_scores) | set(semantic_scores) | set(causal_scores)
score_map = {
fid: entity_scores.get(fid, 0.0) + semantic_scores.get(fid, 0.0) + causal_scores.get(fid, 0.0)
for fid in all_ids
}
sorted_ids = sorted(score_map.keys(), key=lambda x: score_map[x], reverse=True)[:budget]
rows = [row_map[fact_id] for fact_id in sorted_ids]
results = []
for row in rows:
result = RetrievalResult.from_db_row(dict(row))
result.activation = row["score"]
results.append(result)
if tags:
results = filter_results_by_tags(results, tags, match=tags_match)
if tag_groups:
results = filter_results_by_tag_groups(results, tag_groups)
timings.result_count = len(results)
timings.traverse = time.time() - start_time
logger.debug(
f"LinkExpansion: {len(results)} results from {len(seed_ids)} seeds "
f"in {timings.traverse * 1000:.1f}ms (query: {timings.edge_load_time * 1000:.1f}ms)"
)
return results, timings
async def _expand_combined(
self,
conn,
seed_ids: list,
fact_type: str,
budget: int,
) -> tuple[list, list, list]:
"""
Single-roundtrip CTE query combining entity, semantic, and causal expansions.
Uses a `source` discriminator column so the caller can apply per-signal
score transformations. The three CTEs share one connection slot — important
for asyncpg which does not allow concurrent queries on the same connection.
Index coverage (requires migration d2e3f4a5b6c7):
entity: idx_memory_links_entity_covering (from_unit_id) INCLUDE (to_unit_id, entity_id)
WHERE link_type = 'entity' → index-only scan, no heap reads
semantic incoming:
idx_memory_links_to_type_weight (to_unit_id, link_type, weight DESC)
→ replaces costly BitmapAnd of two separate scans
"""
ml = fq_table("memory_links")
mu = fq_table("memory_units")
all_rows = await conn.fetch(
f"""
WITH entity_expanded AS (
-- Entity co-occurrence: seeds → their precomputed entity-link neighbors.
-- Score = distinct shared entities (bounded at retain time to
-- MAX_LINKS_PER_ENTITY=50). GROUP BY mu.id is sufficient because mu.id
-- is the primary key and functionally determines all other mu columns.
SELECT
mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
mu.occurred_end, mu.mentioned_at,
mu.fact_type, mu.document_id, mu.chunk_id, mu.tags,
COUNT(DISTINCT ml.entity_id)::float AS score,
'entity'::text AS source
FROM {ml} ml
JOIN {mu} mu ON mu.id = ml.to_unit_id
WHERE ml.from_unit_id = ANY($1::uuid[])
AND ml.link_type = 'entity'
AND mu.fact_type = $2
AND mu.id != ALL($1::uuid[])
GROUP BY mu.id
ORDER BY score DESC
LIMIT $3
),
semantic_expanded AS (
-- Semantic kNN: both outgoing (seeds → their kNN at insert time) and
-- incoming (facts inserted after seeds that found seeds as kNN).
-- Score = max similarity weight across both directions.
SELECT
id, text, context, event_date, occurred_start,
occurred_end, mentioned_at,
fact_type, document_id, chunk_id, tags,
MAX(weight) AS score,
'semantic'::text AS source
FROM (
SELECT
mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
mu.occurred_end, mu.mentioned_at,
mu.fact_type, mu.document_id, mu.chunk_id, mu.tags,
ml.weight
FROM {ml} ml
JOIN {mu} mu ON mu.id = ml.to_unit_id
WHERE ml.from_unit_id = ANY($1::uuid[])
AND ml.link_type = 'semantic'
AND mu.fact_type = $2
AND mu.id != ALL($1::uuid[])
UNION ALL
SELECT
mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
mu.occurred_end, mu.mentioned_at,
mu.fact_type, mu.document_id, mu.chunk_id, mu.tags,
ml.weight
FROM {ml} ml
JOIN {mu} mu ON mu.id = ml.from_unit_id
WHERE ml.to_unit_id = ANY($1::uuid[])
AND ml.link_type = 'semantic'
AND mu.fact_type = $2
AND mu.id != ALL($1::uuid[])
) sem_raw
GROUP BY id, text, context, event_date, occurred_start,
occurred_end, mentioned_at,
fact_type, document_id, chunk_id, tags
ORDER BY score DESC
LIMIT $3
),
causal_expanded AS (
-- Causal chains: explicit causes/enables/prevents links from seeds.
-- DISTINCT ON handles the case where a seed has multiple causal links
-- to the same target; best weight wins.
SELECT DISTINCT ON (mu.id)
mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
mu.occurred_end, mu.mentioned_at,
mu.fact_type, mu.document_id, mu.chunk_id, mu.tags,
ml.weight AS score,
'causal'::text AS source
FROM {ml} ml
JOIN {mu} mu ON ml.to_unit_id = mu.id
WHERE ml.from_unit_id = ANY($1::uuid[])
AND ml.link_type IN ('causes', 'caused_by', 'enables', 'prevents')
AND ml.weight >= $4
AND mu.fact_type = $2
ORDER BY mu.id, ml.weight DESC
LIMIT $3
)
SELECT * FROM entity_expanded
UNION ALL
SELECT * FROM semantic_expanded
UNION ALL
SELECT * FROM causal_expanded
""",
seed_ids,
fact_type,
budget,
self.causal_weight_threshold,
)
entity_rows = [r for r in all_rows if r["source"] == "entity"]
semantic_rows = [r for r in all_rows if r["source"] == "semantic"]
causal_rows = [r for r in all_rows if r["source"] == "causal"]
return entity_rows, semantic_rows, causal_rows
async def _expand_observations(
self,
conn,
seed_ids: list,
budget: int,
) -> tuple[list, list, list]:
"""
Observation-specific expansion.
Observations don't have direct entity links in memory_links (they're created
by consolidation, not retain). Instead, traverse source_memory_ids → world
facts → entities → other world facts → their observations.
Semantic and causal expansions run as a second combined CTE query.
"""
source_ids_found: list = []
if logger.isEnabledFor(logging.DEBUG):
debug_rows = await conn.fetch(
f"""
SELECT id, source_memory_ids
FROM {fq_table("memory_units")}
WHERE id = ANY($1::uuid[])
""",
seed_ids,
)
for row in debug_rows:
if row["source_memory_ids"]:
source_ids_found.extend(row["source_memory_ids"])
logger.debug(
f"[LinkExpansion] observation graph: {len(seed_ids)} seeds, "
f"{len(source_ids_found)} source_memory_ids found"
)
entity_rows = await conn.fetch(
f"""
WITH seed_sources AS (
SELECT DISTINCT unnest(source_memory_ids) AS source_id
FROM {fq_table("memory_units")}
WHERE id = ANY($1::uuid[])
AND source_memory_ids IS NOT NULL
),
connected_sources AS (
-- Mirror the non-observation entity expansion: follow pre-bounded entity
-- links in memory_links (capped to MAX_LINKS_PER_ENTITY=50 at retain time).
-- Score = number of distinct shared entities, same as the non-obs path.
SELECT DISTINCT ml.to_unit_id AS source_id
FROM seed_sources ss
JOIN {fq_table("memory_links")} ml ON ml.from_unit_id = ss.source_id
WHERE ml.link_type = 'entity'
),
connected_array AS (
SELECT array_agg(source_id) AS source_ids FROM connected_sources
)
SELECT
mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
mu.occurred_end, mu.mentioned_at,
mu.fact_type, mu.document_id, mu.chunk_id, mu.tags,
(SELECT COUNT(DISTINCT s) FROM unnest(mu.source_memory_ids) s WHERE s = ANY(ca.source_ids))::float AS score
FROM {fq_table("memory_units")} mu, connected_array ca
WHERE mu.fact_type = 'observation'
AND mu.id != ALL($1::uuid[])
AND ca.source_ids IS NOT NULL
AND mu.source_memory_ids && ca.source_ids
ORDER BY score DESC
LIMIT $2
""",
seed_ids,
budget,
)
logger.debug(f"[LinkExpansion] observation graph: found {len(entity_rows)} connected observations")
# Semantic + causal for observations in one query
ml = fq_table("memory_links")
mu = fq_table("memory_units")
sem_causal_rows = await conn.fetch(
f"""
WITH semantic_expanded AS (
SELECT
id, text, context, event_date, occurred_start,
occurred_end, mentioned_at,
fact_type, document_id, chunk_id, tags,
MAX(weight) AS score,
'semantic'::text AS source
FROM (
SELECT mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
mu.occurred_end, mu.mentioned_at, mu.fact_type, mu.document_id,
mu.chunk_id, mu.tags, ml.weight
FROM {ml} ml JOIN {mu} mu ON mu.id = ml.to_unit_id
WHERE ml.from_unit_id = ANY($1::uuid[])
AND ml.link_type = 'semantic' AND mu.fact_type = 'observation'
AND mu.id != ALL($1::uuid[])
UNION ALL
SELECT mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
mu.occurred_end, mu.mentioned_at, mu.fact_type, mu.document_id,
mu.chunk_id, mu.tags, ml.weight
FROM {ml} ml JOIN {mu} mu ON mu.id = ml.from_unit_id
WHERE ml.to_unit_id = ANY($1::uuid[])
AND ml.link_type = 'semantic' AND mu.fact_type = 'observation'
AND mu.id != ALL($1::uuid[])
) sem_raw
GROUP BY id, text, context, event_date, occurred_start, occurred_end,
mentioned_at, fact_type, document_id, chunk_id, tags
ORDER BY score DESC LIMIT $2
),
causal_expanded AS (
SELECT DISTINCT ON (mu.id)
mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
mu.occurred_end, mu.mentioned_at, mu.fact_type, mu.document_id,
mu.chunk_id, mu.tags, ml.weight AS score, 'causal'::text AS source
FROM {ml} ml JOIN {mu} mu ON ml.to_unit_id = mu.id
WHERE ml.from_unit_id = ANY($1::uuid[])
AND ml.link_type IN ('causes', 'caused_by', 'enables', 'prevents')
AND ml.weight >= $3 AND mu.fact_type = 'observation'
ORDER BY mu.id, ml.weight DESC LIMIT $2
)
SELECT * FROM semantic_expanded
UNION ALL
SELECT * FROM causal_expanded
""",
seed_ids,
budget,
self.causal_weight_threshold,
)
semantic_rows = [r for r in sem_causal_rows if r["source"] == "semantic"]
causal_rows = [r for r in sem_causal_rows if r["source"] == "causal"]
return entity_rows, semantic_rows, causal_rows
@@ -1,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,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,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)
-205
View File
@@ -1,205 +0,0 @@
[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"
[project]
name = "hindsight-api-slim"
version = "0.4.19"
description = "Hindsight: Agent Memory That Works Like Human Memory"
readme = "README.md"
requires-python = ">=3.11"
dependencies = [
"asyncpg>=0.29.0",
"python-dotenv>=1.0.0",
"openai>=1.0.0",
"pydantic>=2.0.0",
"rich>=13.0.0",
"langchain-text-splitters>=0.3.0",
"fastapi[standard]>=0.120.3",
"uvicorn>=0.38.0",
"wsproto>=1.0.0",
"sqlalchemy>=2.0.44",
"alembic>=1.17.1",
"pgvector>=0.4.1",
"greenlet>=3.2.4",
"psycopg2-binary>=2.9.11",
"tiktoken>=0.12.0",
"httpx>=0.27.0",
"PyJWT[crypto]>=2.8.0",
"fastmcp>=2.14.0", # CVE-2025-66416
"python-dateutil>=2.8.0",
"opentelemetry-api>=1.20.0",
"opentelemetry-sdk>=1.20.0",
"opentelemetry-instrumentation-fastapi>=0.41b0",
"opentelemetry-exporter-prometheus>=0.41b0",
"opentelemetry-exporter-otlp-proto-http>=1.20.0",
"opentelemetry-semantic-conventions>=0.41b0",
"dateparser>=1.2.2",
"google-genai>=1.0.0",
"google-auth>=2.0.0",
"anthropic>=0.40.0",
"typer>=0.9.0",
"cohere>=5.0.0",
"litellm>=1.0.0",
"markitdown[pdf,docx,pptx,xlsx,xls]>=0.1.4", # File to markdown conversion
"obstore>=0.4.0", # S3/GCS/Azure object storage client (Rust-backed)
"uvloop>=0.22.1",
# Transitive dependency security fixes
"pyasn1>=0.6.2", # 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.6", # Account takeover vulnerability fix
"aiohttp>=3.13.3", # Multiple DoS vulnerabilities
"claude-agent-sdk>=0.1.27; sys_platform == 'darwin'",
]
[project.optional-dependencies]
local-ml = [
# Local ML models for embeddings/reranking
"sentence-transformers>=3.3.0",
"transformers>=4.53.0", # Security fixes for ReDoS vulnerabilities
"torch>=2.6.0", # CVE fix for remote code execution
"einops>=0.8.2",
"flashrank>=0.2.0",
# Apple Silicon local inference
"mlx>=0.31.0",
"mlx-lm>=0.31.1",
"safetensors>=0.6.2",
]
embedded-db = [
"pg0-embedded>=0.11.0",
]
all = [
"hindsight-api-slim[local-ml,embedded-db]",
]
test = [
"pytest>=7.0.0",
"pytest-asyncio>=0.21.0",
"pytest-timeout>=2.4.0",
"pytest-xdist>=3.0.0",
"filelock>=3.20.1", # TOCTOU race condition fix
"testcontainers>=4.0.0",
]
[project.scripts]
hindsight-api = "hindsight_api.main:main"
hindsight-worker = "hindsight_api.worker.main:main"
hindsight-local-mcp = "hindsight_api.mcp_local:main"
hindsight-admin = "hindsight_api.admin.cli:main"
[tool.hatch.build.targets.wheel]
packages = ["hindsight_api"]
[tool.hatch.build.targets.wheel.sources]
"hindsight_api" = "hindsight_api"
[tool.hatch.build.targets.sdist]
include = [
"hindsight_api/**/*",
]
[tool.hatch.build]
include = [
"hindsight_api/**/*.py",
"hindsight_api/alembic/**/*",
]
[tool.pytest.ini_options]
log_cli = true
log_cli_level = "INFO"
log_cli_format = "%(asctime)s - %(levelname)s - %(name)s - %(message)s"
log_cli_date_format = "%Y-%m-%d %H:%M:%S"
addopts = "--timeout 300 -n 8 --dist loadgroup --durations=10 -v"
asyncio_mode = "auto"
asyncio_default_fixture_loop_scope = "function"
log_auto_indent = true
filterwarnings = [
"ignore:The @wait_container_is_ready decorator is deprecated:DeprecationWarning",
"ignore::RuntimeWarning:asyncio",
]
[dependency-groups]
dev = [
"pytest>=9.0.0",
"pytest-asyncio>=1.3.0",
"pytest-timeout>=2.4.0",
"pytest-xdist>=3.8.0",
"python-dotenv>=1.2.1",
"filelock>=3.20.1", # TOCTOU race condition fix
"ruff>=0.8.0",
"ty>=0.0.1",
"testcontainers>=4.0.0",
]
[tool.ruff]
line-length = 120
target-version = "py311"
exclude = [
"tests/",
"**/tests/",
]
[tool.ruff.lint]
select = [
"E", # pycodestyle errors
"W", # pycodestyle warnings
"F", # Pyflakes
"I", # isort
]
ignore = [
"E501", # line too long (handled by formatter)
"E402", # module import not at top of file
"F401", # unused import (too noisy during development)
"F841", # unused variable (too noisy during development)
"F811", # redefined while unused
"F821", # undefined name (forward references in type hints)
]
[tool.ruff.lint.isort]
known-third-party = ["alembic"]
[tool.ruff.format]
quote-style = "double"
indent-style = "space"
[tool.uv]
# Use explicit index for PyTorch to prevent the pytorch index from serving
# non-pytorch packages (e.g. markupsafe) with incompatible wheels
[[tool.uv.index]]
name = "pytorch-cpu"
url = "https://download.pytorch.org/whl/cpu"
explicit = true
[tool.uv.sources]
# Route torch to the CPU-only PyTorch index; everything else uses PyPI
torch = { index = "pytorch-cpu" }
[tool.ty]
# Type checking configuration
# ty is an extremely fast Python type checker from Astral (same team as ruff/uv)
[tool.ty.environment]
python-version = "3.11"
[tool.ty.src]
exclude = [
"tests/",
"hindsight_api/alembic/",
]
[tool.ty.rules]
# Disable noisy rules while keeping important ones
invalid-argument-type = "ignore" # False positives with **kwargs patterns
invalid-return-type = "ignore" # Often intentional in async code
invalid-parameter-default = "ignore" # Optional params with None default
possibly-missing-attribute = "ignore" # Common with Optional types
invalid-raise = "ignore" # False positives with exception tracking
call-non-callable = "ignore" # False positives with Optional types
invalid-key = "ignore" # Pydantic ConfigDict not understood
invalid-method-override = "ignore" # Intentional signature differences
unresolved-reference = "ignore" # Forward references not always resolved
@@ -1,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)
File diff suppressed because it is too large Load Diff
@@ -1,70 +0,0 @@
"""
Tests for EntityResolver edge cases.
"""
import uuid
from datetime import datetime, timezone
import asyncpg
import pytest
from hindsight_api.engine.entity_resolver import EntityResolver
from hindsight_api.pg0 import resolve_database_url
@pytest.mark.asyncio
async def test_resolve_entities_batch_handles_unicode_lower_conflicts(pg0_db_url):
"""
Existing entities with PostgreSQL/Python lowercase mismatches should resolve
to the conflicted row instead of leaving a missing entity_id.
"""
resolved_url = await resolve_database_url(pg0_db_url)
pool = await asyncpg.create_pool(resolved_url, min_size=1, max_size=2, command_timeout=30)
bank_id = f"test-entity-resolver-{uuid.uuid4().hex[:8]}"
event_date = datetime(2024, 1, 15, tzinfo=timezone.utc)
resolver = EntityResolver(pool=pool, entity_lookup="full")
try:
async with pool.acquire() as conn:
existing_entity_id = await conn.fetchval(
"""
INSERT INTO entities (bank_id, canonical_name, first_seen, last_seen, mention_count)
VALUES ($1, $2, $3, $3, 1)
RETURNING id
""",
bank_id,
"İstanbul",
event_date,
)
resolved_ids = await resolver.resolve_entities_batch(
bank_id=bank_id,
entities_data=[
{
"text": "istanbul",
"nearby_entities": [],
"event_date": event_date,
}
],
context="unicode case mismatch",
unit_event_date=event_date,
conn=conn,
)
entity_rows = await conn.fetch(
"""
SELECT id, canonical_name
FROM entities
WHERE bank_id = $1
ORDER BY canonical_name
""",
bank_id,
)
assert resolved_ids == [existing_entity_id]
assert len(entity_rows) == 1
assert entity_rows[0]["id"] == existing_entity_id
assert entity_rows[0]["canonical_name"] == "İstanbul"
finally:
await pool.execute("DELETE FROM entities WHERE bank_id = $1", bank_id)
await pool.close()
@@ -1,61 +0,0 @@
"""
Unit tests for metadata inclusion in fact extraction LLM prompt.
"""
from datetime import datetime
from hindsight_api.engine.retain.fact_extraction import _build_user_message
def test_build_user_message_includes_metadata():
"""Metadata key-value pairs should appear in the user message."""
event_date = datetime(2024, 6, 15, 12, 0, 0)
metadata = {"title": "Q2 Planning Doc", "source": "confluence", "author": "Alice"}
msg = _build_user_message(
chunk="Some content.",
chunk_index=0,
total_chunks=1,
event_date=event_date,
context="planning meeting",
metadata=metadata,
)
assert "title" in msg
assert "Q2 Planning Doc" in msg
assert "source" in msg
assert "confluence" in msg
assert "author" in msg
assert "Alice" in msg
def test_build_user_message_no_metadata():
"""When metadata is empty, the message should still be valid and not include a metadata section."""
event_date = datetime(2024, 6, 15, 12, 0, 0)
msg = _build_user_message(
chunk="Some content.",
chunk_index=0,
total_chunks=1,
event_date=event_date,
context="planning meeting",
metadata={},
)
assert "Some content." in msg
assert "Metadata:" not in msg
def test_build_user_message_without_metadata_arg():
"""Calling without metadata (default) should behave the same as empty metadata."""
event_date = datetime(2024, 6, 15, 12, 0, 0)
msg = _build_user_message(
chunk="Some content.",
chunk_index=0,
total_chunks=1,
event_date=event_date,
context="none",
)
assert "Some content." in msg
assert "Metadata:" not in msg
@@ -1,141 +0,0 @@
"""
Unit tests for fact extraction retry logic.
Tests the fix for the TypeError when LLM returns invalid JSON across all retries.
Previously, `raise last_error` would raise None (TypeError) because last_error was
only set in the BadRequestError handler, not when the LLM returned non-dict JSON.
"""
from datetime import datetime, timezone
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
def _make_config(llm_max_retries: int = 3, retain_llm_max_retries: int | None = None):
"""Build a minimal HindsightConfig for fact extraction tests."""
from hindsight_api.config import HindsightConfig
cfg = MagicMock(spec=HindsightConfig)
cfg.retain_llm_max_retries = retain_llm_max_retries
cfg.llm_max_retries = llm_max_retries
cfg.retain_llm_initial_backoff = None
cfg.llm_initial_backoff = 0.0
cfg.retain_llm_max_backoff = None
cfg.llm_max_backoff = 0.0
cfg.retain_max_completion_tokens = 8192
cfg.retain_extraction_mode = "concise"
cfg.retain_extract_causal_links = False
cfg.retain_mission = None
return cfg
def _make_llm_config(mock_response):
"""Build a mock LLMProvider that returns the given response."""
from hindsight_api.engine.llm_wrapper import LLMProvider
llm = MagicMock(spec=LLMProvider)
llm.provider = "mock"
token_usage = MagicMock()
token_usage.__add__ = lambda self, other: self
llm.call = AsyncMock(return_value=(mock_response, token_usage))
return llm
@pytest.mark.asyncio
async def test_non_dict_json_all_retries_returns_empty():
"""
When LLM returns non-dict JSON on every attempt, extraction should return []
without raising TypeError ('exceptions must derive from BaseException').
This was the bug: the loop ran range(2) times (hardcoded), but comparisons
used config.llm_max_retries (default 10). On the last loop iteration (attempt=1),
`attempt < 10 - 1` was True, so the code called `continue`, the loop
exhausted, and `raise last_error` raised None → TypeError.
"""
from hindsight_api.engine.retain.fact_extraction import _extract_facts_from_chunk
# llm_max_retries=3 ensures the bug triggers with the old code (3 != 2 hardcoded)
config = _make_config(llm_max_retries=3, retain_llm_max_retries=None)
# Mock: always returns a list (non-dict), which is invalid
llm_config = _make_llm_config(mock_response=[{"invalid": "response"}])
with patch(
"hindsight_api.engine.retain.fact_extraction._build_extraction_prompt_and_schema",
return_value=("system prompt", MagicMock()),
):
facts, usage = await _extract_facts_from_chunk(
chunk="Alice visited Paris in 2023.",
chunk_index=0,
total_chunks=1,
event_date=datetime(2023, 1, 1, tzinfo=timezone.utc),
context="travel notes",
llm_config=llm_config,
config=config,
agent_name="test-agent",
)
assert facts == []
@pytest.mark.asyncio
async def test_non_dict_json_with_default_max_retries_returns_empty():
"""
Same scenario with the default llm_max_retries=10 (matching real default config).
The old code ran range(2) but checked against 10, always continuing until
the loop exhausted, then raised None → TypeError.
"""
from hindsight_api.engine.retain.fact_extraction import _extract_facts_from_chunk
config = _make_config(llm_max_retries=10, retain_llm_max_retries=None)
llm_config = _make_llm_config(mock_response="not a dict at all")
with patch(
"hindsight_api.engine.retain.fact_extraction._build_extraction_prompt_and_schema",
return_value=("system prompt", MagicMock()),
):
facts, usage = await _extract_facts_from_chunk(
chunk="Some text.",
chunk_index=0,
total_chunks=1,
event_date=datetime(2023, 6, 1, tzinfo=timezone.utc),
context="",
llm_config=llm_config,
config=config,
agent_name="agent",
)
assert facts == []
@pytest.mark.asyncio
async def test_retain_llm_max_retries_overrides_global():
"""
When retain_llm_max_retries is set, it should be used for the loop range
and all comparisons (no shadowing bug).
"""
from hindsight_api.engine.retain.fact_extraction import _extract_facts_from_chunk
# retain_llm_max_retries=5 should override llm_max_retries=10
config = _make_config(llm_max_retries=10, retain_llm_max_retries=5)
llm_config = _make_llm_config(mock_response=42) # non-dict: integer
with patch(
"hindsight_api.engine.retain.fact_extraction._build_extraction_prompt_and_schema",
return_value=("system prompt", MagicMock()),
):
facts, usage = await _extract_facts_from_chunk(
chunk="Bob likes Python.",
chunk_index=0,
total_chunks=1,
event_date=datetime(2024, 1, 1, tzinfo=timezone.utc),
context="",
llm_config=llm_config,
config=config,
agent_name="agent",
)
assert facts == []
# Verify it retried exactly retain_llm_max_retries times
assert llm_config.call.call_count == 5
@@ -1,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"
@@ -1,174 +0,0 @@
"""
Tests for list_documents pagination and tags filtering.
"""
from datetime import datetime, timezone
import pytest
async def _retain_doc(memory, bank_id, document_id, tags, request_context):
"""Helper to retain a document with given tags. Uses gibberish content to avoid LLM
fact extraction (documents are persisted even with zero facts)."""
await memory.retain_batch_async(
bank_id=bank_id,
contents=[{"content": f"xyzabc123 !@# $$$ {document_id}"}],
document_id=document_id,
document_tags=tags or None,
request_context=request_context,
)
@pytest.mark.asyncio
async def test_list_documents_offset_pagination(memory, request_context):
"""offset parameter returns the correct slice of documents."""
bank_id = f"test_list_docs_offset_{datetime.now(timezone.utc).timestamp()}"
try:
for i in range(4):
await _retain_doc(memory, bank_id, f"doc-{i:02d}", [], request_context)
# All documents, ordered by created_at DESC → doc-03, doc-02, doc-01, doc-00
all_docs = await memory.list_documents(
bank_id=bank_id, limit=10, offset=0, request_context=request_context
)
assert all_docs["total"] == 4
assert len(all_docs["items"]) == 4
all_ids = [d["id"] for d in all_docs["items"]]
# offset=2 should skip the first two and return the remaining two
page2 = await memory.list_documents(
bank_id=bank_id, limit=10, offset=2, request_context=request_context
)
assert page2["total"] == 4 # total is always the full count
assert len(page2["items"]) == 2
assert [d["id"] for d in page2["items"]] == all_ids[2:]
# offset beyond total returns empty items but correct total
beyond = await memory.list_documents(
bank_id=bank_id, limit=10, offset=10, request_context=request_context
)
assert beyond["total"] == 4
assert beyond["items"] == []
finally:
await memory.delete_bank(bank_id, request_context=request_context)
@pytest.mark.asyncio
async def test_list_documents_tags_filter_any_strict(memory, request_context):
"""tags filter with any_strict returns only tagged documents that match."""
bank_id = f"test_list_docs_tags_{datetime.now(timezone.utc).timestamp()}"
try:
await _retain_doc(memory, bank_id, "doc-alpha", ["team-a"], request_context)
await _retain_doc(memory, bank_id, "doc-beta", ["team-b"], request_context)
await _retain_doc(memory, bank_id, "doc-both", ["team-a", "team-b"], request_context)
await _retain_doc(memory, bank_id, "doc-untagged", [], request_context)
# any_strict: only docs with at least one of the given tags, untagged excluded
result = await memory.list_documents(
bank_id=bank_id,
tags=["team-a"],
tags_match="any_strict",
request_context=request_context,
)
ids = {d["id"] for d in result["items"]}
assert ids == {"doc-alpha", "doc-both"}
assert result["total"] == 2
finally:
await memory.delete_bank(bank_id, request_context=request_context)
@pytest.mark.asyncio
async def test_list_documents_tags_filter_any_includes_untagged(memory, request_context):
"""tags filter with 'any' mode includes untagged documents."""
bank_id = f"test_list_docs_tags_any_{datetime.now(timezone.utc).timestamp()}"
try:
await _retain_doc(memory, bank_id, "doc-tagged", ["team-a"], request_context)
await _retain_doc(memory, bank_id, "doc-other", ["team-b"], request_context)
await _retain_doc(memory, bank_id, "doc-untagged", [], request_context)
result = await memory.list_documents(
bank_id=bank_id,
tags=["team-a"],
tags_match="any",
request_context=request_context,
)
ids = {d["id"] for d in result["items"]}
# "any" includes untagged + matching tagged
assert "doc-tagged" in ids
assert "doc-untagged" in ids
assert "doc-other" not in ids
finally:
await memory.delete_bank(bank_id, request_context=request_context)
@pytest.mark.asyncio
async def test_list_documents_tags_filter_all_strict(memory, request_context):
"""tags filter with all_strict returns only docs that have ALL the specified tags."""
bank_id = f"test_list_docs_tags_all_{datetime.now(timezone.utc).timestamp()}"
try:
await _retain_doc(memory, bank_id, "doc-a-only", ["team-a"], request_context)
await _retain_doc(memory, bank_id, "doc-a-and-b", ["team-a", "team-b"], request_context)
await _retain_doc(memory, bank_id, "doc-untagged", [], request_context)
result = await memory.list_documents(
bank_id=bank_id,
tags=["team-a", "team-b"],
tags_match="all_strict",
request_context=request_context,
)
ids = {d["id"] for d in result["items"]}
assert ids == {"doc-a-and-b"}
finally:
await memory.delete_bank(bank_id, request_context=request_context)
@pytest.mark.asyncio
async def test_list_documents_no_tags_filter_returns_all(memory, request_context):
"""When no tags filter is specified, all documents are returned."""
bank_id = f"test_list_docs_no_tags_{datetime.now(timezone.utc).timestamp()}"
try:
await _retain_doc(memory, bank_id, "doc-tagged", ["team-a"], request_context)
await _retain_doc(memory, bank_id, "doc-untagged", [], request_context)
result = await memory.list_documents(
bank_id=bank_id,
tags=None,
request_context=request_context,
)
ids = {d["id"] for d in result["items"]}
assert ids == {"doc-tagged", "doc-untagged"}
finally:
await memory.delete_bank(bank_id, request_context=request_context)
@pytest.mark.asyncio
async def test_list_documents_tags_and_search_query_combined(memory, request_context):
"""tags filter and q (search_query) can be combined."""
bank_id = f"test_list_docs_tags_q_{datetime.now(timezone.utc).timestamp()}"
try:
await _retain_doc(memory, bank_id, "report-2024", ["team-a"], request_context)
await _retain_doc(memory, bank_id, "report-2025", ["team-b"], request_context)
await _retain_doc(memory, bank_id, "summary-2024", ["team-a"], request_context)
result = await memory.list_documents(
bank_id=bank_id,
search_query="report",
tags=["team-a"],
tags_match="any_strict",
request_context=request_context,
)
ids = {d["id"] for d in result["items"]}
assert ids == {"report-2024"}
finally:
await memory.delete_bank(bank_id, request_context=request_context)
@@ -1,37 +0,0 @@
import pytest
from hindsight_api.engine.llm_wrapper import sanitize_llm_output
@pytest.mark.parametrize(
"input_text, expected",
[
# Null bytes stripped
("hello\x00world", "helloworld"),
("FIRST\u0000PAGE", "FIRSTPAGE"),
# Multiple null bytes
("\x00\x00text\x00", "text"),
# Other control characters stripped (non-whitespace)
("text\x01\x02\x03end", "textend"),
("text\x08end", "textend"), # backspace
("text\x0cend", "textend"), # form feed
("text\x0bend", "textend"), # vertical tab
("text\x1fend", "textend"), # unit separator
("text\x7fend", "textend"), # DEL
# Whitespace preserved
("hello\tworld", "hello\tworld"),
("hello\nworld", "hello\nworld"),
("hello\r\nworld", "hello\r\nworld"),
# Unicode surrogates stripped
("text\ud800end", "textend"),
("text\udfffend", "textend"),
# Clean text unchanged
("normal text", "normal text"),
("unicode: café naïve", "unicode: café naïve"),
# Edge cases
("", ""),
(None, None),
],
)
def test_sanitize_llm_output(input_text, expected):
assert sanitize_llm_output(input_text) == expected
@@ -1,321 +0,0 @@
"""
Reproduce issue #520: Reflect fails with LM Studio due to unsupported tool_choice format.
The reflect agent forces tool selection via named tool_choice dicts on the first few iterations:
{"type": "function", "function": {"name": "search_mental_models"}}
LM Studio (and Ollama) reject this format with HTTP 400:
"Tool choice of type 'function' is not supported. Use 'auto', 'none', or 'required'."
The fix should convert named tool_choice to "required" and filter the tools list
to only the requested tool for providers that don't support named tool_choice.
"""
import json
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from openai import APIStatusError
from hindsight_api.engine.providers.openai_compatible_llm import OpenAICompatibleLLM
# Reflect agent tools (subset matching what agent.py uses)
REFLECT_TOOLS = [
{
"type": "function",
"function": {
"name": "search_mental_models",
"description": "Search consolidated mental models",
"parameters": {
"type": "object",
"properties": {"query": {"type": "string"}},
"required": ["query"],
},
},
},
{
"type": "function",
"function": {
"name": "search_observations",
"description": "Search raw observations",
"parameters": {
"type": "object",
"properties": {"query": {"type": "string"}},
"required": ["query"],
},
},
},
{
"type": "function",
"function": {
"name": "recall",
"description": "Recall semantic memories",
"parameters": {
"type": "object",
"properties": {"query": {"type": "string"}},
"required": ["query"],
},
},
},
{
"type": "function",
"function": {
"name": "done",
"description": "Finish and return the answer",
"parameters": {
"type": "object",
"properties": {"answer": {"type": "string"}},
"required": ["answer"],
},
},
},
]
def _make_lmstudio_llm() -> OpenAICompatibleLLM:
return OpenAICompatibleLLM(
provider="lmstudio",
api_key="local",
base_url="http://localhost:1234/v1",
model="openai/gpt-oss-20b",
)
def _lmstudio_400_error(msg: str = "Tool choice of type 'function' is not supported. Use 'auto', 'none', or 'required'.") -> APIStatusError:
"""Simulate the HTTP 400 LM Studio returns for unsupported tool_choice format."""
mock_response = MagicMock()
mock_response.status_code = 400
mock_response.headers = {}
return APIStatusError(
message=msg,
response=mock_response,
body={"error": {"message": msg, "type": "invalid_request_error"}},
)
def _make_tool_call_response(tool_name: str, arguments: dict) -> MagicMock:
"""Build a mock successful tool call response from the LLM API."""
mock_tc = MagicMock()
mock_tc.id = "call_abc123"
mock_tc.function.name = tool_name
mock_tc.function.arguments = json.dumps(arguments)
mock_response = MagicMock()
mock_response.usage.prompt_tokens = 120
mock_response.usage.completion_tokens = 40
mock_response.usage.total_tokens = 160
mock_response.choices[0].finish_reason = "tool_calls"
mock_response.choices[0].message.content = None
mock_response.choices[0].message.tool_calls = [mock_tc]
return mock_response
class TestLMStudioNamedToolChoiceBug:
"""
Reproduces issue #520.
The reflect agent (agent.py lines 546-555) sets tool_choice to a named dict
on the first iterations to force sequential retrieval:
iteration=0, has_mental_models=True → {"type": "function", "function": {"name": "search_mental_models"}}
iteration=0, has_mental_models=False → {"type": "function", "function": {"name": "search_observations"}}
iteration=1, has_mental_models=True → {"type": "function", "function": {"name": "search_observations"}}
iteration=1 or (2 with models) → {"type": "function", "function": {"name": "recall"}}
LM Studio rejects these dict formats with HTTP 400.
"""
@pytest.mark.asyncio
async def test_lmstudio_named_tool_choice_no_longer_causes_400(self):
"""
Regression test for issue #520: named tool_choice dict is converted to
"required" + filtered tools before the API call, so LM Studio never
sees the unsupported format and the 400 error no longer occurs.
"""
llm = _make_lmstudio_llm()
named_tool_choice = {"type": "function", "function": {"name": "search_mental_models"}}
success_response = _make_tool_call_response("search_mental_models", {"query": "user name"})
with patch.object(llm._client.chat.completions, "create", new_callable=AsyncMock) as mock_create:
mock_create.return_value = success_response
# Should succeed — no 400 because the dict is converted before sending
result = await llm.call_with_tools(
messages=[{"role": "user", "content": "What is the user's name?"}],
tools=REFLECT_TOOLS,
tool_choice=named_tool_choice,
max_retries=0,
)
assert len(result.tool_calls) == 1
assert result.tool_calls[0].name == "search_mental_models"
sent_kwargs = mock_create.call_args.kwargs
assert sent_kwargs["tool_choice"] == "required"
assert len(sent_kwargs["tools"]) == 1
assert sent_kwargs["tools"][0]["function"]["name"] == "search_mental_models"
@pytest.mark.asyncio
@pytest.mark.parametrize(
"forced_tool_name",
["search_mental_models", "search_observations", "recall"],
)
async def test_all_reflect_forced_tools_fail_on_lmstudio(self, forced_tool_name: str):
"""
Each named tool_choice the reflect agent uses on iterations 0-2 triggers
the same 400 error on LM Studio.
"""
llm = _make_lmstudio_llm()
named_tool_choice = {"type": "function", "function": {"name": forced_tool_name}}
with patch.object(llm._client.chat.completions, "create", new_callable=AsyncMock) as mock_create:
mock_create.side_effect = _lmstudio_400_error()
with pytest.raises(APIStatusError) as exc_info:
await llm.call_with_tools(
messages=[{"role": "user", "content": "Test query"}],
tools=REFLECT_TOOLS,
tool_choice=named_tool_choice,
max_retries=0,
)
assert exc_info.value.status_code == 400
@pytest.mark.asyncio
async def test_lmstudio_string_tool_choice_works_fine(self):
"""
String tool_choice values ("auto", "none", "required") ARE supported by LM Studio.
Only the dict format {"type": "function", "function": {"name": "..."}} fails.
This test confirms the control case works.
"""
llm = _make_lmstudio_llm()
success_response = _make_tool_call_response("search_mental_models", {"query": "user name"})
with patch.object(llm._client.chat.completions, "create", new_callable=AsyncMock) as mock_create:
mock_create.return_value = success_response
result = await llm.call_with_tools(
messages=[{"role": "user", "content": "What is the user's name?"}],
tools=REFLECT_TOOLS,
tool_choice="required", # string form — LM Studio accepts this
max_retries=0,
)
assert len(result.tool_calls) == 1
assert result.tool_calls[0].name == "search_mental_models"
# Confirm "required" was sent, not a dict
sent_kwargs = mock_create.call_args.kwargs
assert sent_kwargs["tool_choice"] == "required"
class TestExpectedFixBehavior:
"""
Tests that document the EXPECTED behavior after the fix is applied.
For lmstudio (and ollama) providers, when tool_choice is a named dict:
{"type": "function", "function": {"name": "search_mental_models"}}
The fix should:
1. Convert tool_choice to "required"
2. Filter tools to only the requested tool
These tests currently FAIL (because the fix is not yet implemented).
After the fix is applied, they should PASS.
"""
@pytest.mark.asyncio
async def test_fix_converts_named_tool_choice_to_required(self):
"""
After fix: named tool_choice dict is converted to "required" for lmstudio.
The API receives tool_choice="required" instead of the unsupported dict.
"""
llm = _make_lmstudio_llm()
named_tool_choice = {"type": "function", "function": {"name": "search_mental_models"}}
success_response = _make_tool_call_response("search_mental_models", {"query": "user name"})
with patch.object(llm._client.chat.completions, "create", new_callable=AsyncMock) as mock_create:
mock_create.return_value = success_response
result = await llm.call_with_tools(
messages=[{"role": "user", "content": "What is the user's name?"}],
tools=REFLECT_TOOLS,
tool_choice=named_tool_choice,
max_retries=0,
)
assert len(result.tool_calls) == 1
assert result.tool_calls[0].name == "search_mental_models"
sent_kwargs = mock_create.call_args.kwargs
# Fix: dict was converted to "required"
assert sent_kwargs["tool_choice"] == "required", (
f"Expected tool_choice='required', got {sent_kwargs['tool_choice']!r}"
)
# Fix: tools filtered to just the requested one
assert len(sent_kwargs["tools"]) == 1
assert sent_kwargs["tools"][0]["function"]["name"] == "search_mental_models"
@pytest.mark.asyncio
@pytest.mark.parametrize(
"forced_tool_name",
["search_mental_models", "search_observations", "recall"],
)
async def test_fix_filters_tools_to_requested_tool(self, forced_tool_name: str):
"""
After fix: tools list is filtered to only the forced tool so the model
can only call that one tool (equivalent to the named tool_choice behavior).
"""
llm = _make_lmstudio_llm()
named_tool_choice = {"type": "function", "function": {"name": forced_tool_name}}
success_response = _make_tool_call_response(forced_tool_name, {"query": "test"})
with patch.object(llm._client.chat.completions, "create", new_callable=AsyncMock) as mock_create:
mock_create.return_value = success_response
await llm.call_with_tools(
messages=[{"role": "user", "content": "Test query"}],
tools=REFLECT_TOOLS,
tool_choice=named_tool_choice,
max_retries=0,
)
sent_kwargs = mock_create.call_args.kwargs
assert sent_kwargs["tool_choice"] == "required"
assert len(sent_kwargs["tools"]) == 1
assert sent_kwargs["tools"][0]["function"]["name"] == forced_tool_name
@pytest.mark.asyncio
async def test_fix_also_applies_to_openai_provider(self):
"""
The fix is generalized: all providers convert named tool_choice to
"required" + filtered tools. OpenAI natively supports the dict format
too, so the behaviour is semantically identical either way.
"""
from hindsight_api.engine.providers.openai_compatible_llm import OpenAICompatibleLLM
openai_llm = OpenAICompatibleLLM(
provider="openai",
api_key="sk-test",
base_url="",
model="gpt-4o-mini",
)
named_tool_choice = {"type": "function", "function": {"name": "search_mental_models"}}
success_response = _make_tool_call_response("search_mental_models", {"query": "test"})
with patch.object(openai_llm._client.chat.completions, "create", new_callable=AsyncMock) as mock_create:
mock_create.return_value = success_response
await openai_llm.call_with_tools(
messages=[{"role": "user", "content": "Test"}],
tools=REFLECT_TOOLS,
tool_choice=named_tool_choice,
max_retries=0,
)
sent_kwargs = mock_create.call_args.kwargs
# Generalized fix applies to OpenAI too
assert sent_kwargs["tool_choice"] == "required"
assert len(sent_kwargs["tools"]) == 1
assert sent_kwargs["tools"][0]["function"]["name"] == "search_mental_models"
File diff suppressed because it is too large Load Diff
@@ -1,154 +0,0 @@
"""Tests for migration g7h8i9j0k1l2 (backsweep orphaned memory_units).
Uses a dedicated pg0 instance (port 5562) so the test can control exactly
which migrations have run before inserting the orphan seed data.
"""
import asyncio
import uuid
from pathlib import Path
import pytest
from alembic import command
from alembic.config import Config
from sqlalchemy import create_engine, text
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
_SCRIPT_LOCATION = str(Path(__file__).parent.parent / "hindsight_api" / "alembic")
def _alembic_cfg(db_url: str) -> Config:
cfg = Config()
cfg.set_main_option("script_location", _SCRIPT_LOCATION)
cfg.set_main_option("sqlalchemy.url", db_url)
cfg.set_main_option("prepend_sys_path", ".")
cfg.set_main_option("path_separator", "os")
return cfg
def _upgrade(db_url: str, revision: str) -> None:
command.upgrade(_alembic_cfg(db_url), revision)
# ---------------------------------------------------------------------------
# Fixture: fresh database at the revision just before the backsweep
# ---------------------------------------------------------------------------
@pytest.fixture(scope="module")
def pre_backsweep_db_url():
"""
Spin up a dedicated pg0 instance and run all migrations up to (but not
including) the backsweep revision so each test can seed orphan data and
then apply the backsweep itself.
"""
from hindsight_api.pg0 import EmbeddedPostgres
pg0 = EmbeddedPostgres(name="hindsight-backsweep-test", port=5562)
loop = asyncio.new_event_loop()
try:
url = loop.run_until_complete(pg0.ensure_running())
finally:
loop.close()
# Migrate up to the revision just before the backsweep.
_upgrade(url, "f6g7h8i9j0k1")
return url
# ---------------------------------------------------------------------------
# The test
# ---------------------------------------------------------------------------
def test_backsweep_removes_orphans_and_preserves_legit_rows(pre_backsweep_db_url):
"""
Seed four kinds of rows then apply the backsweep migration and verify:
Rows that MUST be deleted
─────────────────────────
A. Any fact_type, bank_id missing from banks
→ Pass 1 deletes these regardless of fact_type or source links.
B. observation, bank exists, but ALL source_memory_ids are gone
→ Pass 2 deletes these.
Rows that MUST survive
──────────────────────
C. observation, bank exists, at least ONE source_memory_id still live
→ Pass 2 must not touch these.
D. Non-observation (world), bank exists, no sources (not relevant)
→ Pass 1 must not touch these (bank exists).
"""
db_url = pre_backsweep_db_url
engine = create_engine(db_url)
alive_bank = f"bank_{uuid.uuid4().hex[:8]}"
ghost_bank = f"bank_{uuid.uuid4().hex[:8]}" # never inserted into banks
# UUIDs for memory units
id_pass1_world = uuid.uuid4() # A: world unit, ghost bank
id_pass1_obs = uuid.uuid4() # A: observation, ghost bank
id_pass2_obs = uuid.uuid4() # B: observation, all sources gone
id_keep_obs = uuid.uuid4() # C: observation with one live source
id_keep_world = uuid.uuid4() # D: world unit, alive bank
id_live_source = uuid.uuid4() # live source for C
with engine.connect() as conn:
# --- banks ---
conn.execute(text("INSERT INTO banks (bank_id) VALUES (:b)"), {"b": alive_bank})
# --- seed memory_units ---
def insert_mu(uid, bank, fact_type, sources=None):
src_arr = "{" + ",".join(str(s) for s in (sources or [])) + "}"
conn.execute(
text(
"""
INSERT INTO memory_units
(id, bank_id, text, fact_type, source_memory_ids)
VALUES
(:id, :bank, :text, :ft, CAST(:src AS uuid[]))
"""
),
{"id": uid, "bank": bank, "text": "test", "ft": fact_type, "src": src_arr},
)
# A: ghost-bank rows (Pass 1 targets)
insert_mu(id_pass1_world, ghost_bank, "world")
insert_mu(id_pass1_obs, ghost_bank, "observation", sources=[uuid.uuid4()])
# B: observation with all-dead sources (Pass 2 target)
insert_mu(id_pass2_obs, alive_bank, "observation", sources=[uuid.uuid4(), uuid.uuid4()])
# C: observation with one live source (must survive)
insert_mu(id_live_source, alive_bank, "world")
insert_mu(id_keep_obs, alive_bank, "observation", sources=[id_live_source, uuid.uuid4()])
# D: world unit in alive bank (must survive)
insert_mu(id_keep_world, alive_bank, "world")
conn.commit()
# --- apply the backsweep ---
_upgrade(db_url, "g7h8i9j0k1l2")
# --- verify ---
with engine.connect() as conn:
def exists(uid):
return conn.execute(
text("SELECT 1 FROM memory_units WHERE id = :id"), {"id": uid}
).fetchone() is not None
# Must be gone
assert not exists(id_pass1_world), "Pass 1: world unit with ghost bank should be deleted"
assert not exists(id_pass1_obs), "Pass 1: observation with ghost bank should be deleted"
assert not exists(id_pass2_obs), "Pass 2: observation with all-dead sources should be deleted"
# Must survive
assert exists(id_keep_obs), "observation with a live source must not be deleted"
assert exists(id_keep_world), "world unit in alive bank must not be deleted"
assert exists(id_live_source), "live source memory unit must not be deleted"
engine.dispose()
@@ -1,48 +0,0 @@
import threading
import time
from hindsight_api import migrations
def test_run_migrations_internal_serializes_alembic_upgrade(monkeypatch):
max_concurrent_upgrades = 0
active_upgrades = 0
active_lock = threading.Lock()
start_barrier = threading.Barrier(2)
def fake_upgrade(_cfg, _revision):
nonlocal max_concurrent_upgrades, active_upgrades
with active_lock:
active_upgrades += 1
max_concurrent_upgrades = max(max_concurrent_upgrades, active_upgrades)
time.sleep(0.05)
with active_lock:
active_upgrades -= 1
monkeypatch.setattr(migrations.command, "upgrade", fake_upgrade)
errors = []
def run_in_thread(schema):
try:
start_barrier.wait()
migrations._run_migrations_internal(
"postgresql://user:pass@localhost/db",
"/tmp/alembic",
schema=schema,
)
except Exception as exc: # pragma: no cover - diagnostic path
errors.append(exc)
threads = [
threading.Thread(target=run_in_thread, args=("tenant_alpha",)),
threading.Thread(target=run_in_thread, args=("tenant_beta",)),
]
for thread in threads:
thread.start()
for thread in threads:
thread.join()
assert not errors
assert max_concurrent_upgrades == 1
@@ -1,735 +0,0 @@
"""
Tests for observation invalidation when source memories are deleted.
These tests verify that:
1. Observations are deleted (not just updated) when their source memories are removed
2. Remaining source memories are reset for re-consolidation (consolidated_at=NULL)
3. The clear_observations_for_memory method correctly clears observations and
resets the target memory itself for re-consolidation
4. delete_bank(fact_type=...) also cleans up affected observations
"""
import uuid
from unittest.mock import AsyncMock, patch
import pytest
from hindsight_api import RequestContext
from hindsight_api.engine.memory_engine import MemoryEngine
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
async def _insert_memory(conn, bank_id: str, text: str, fact_type: str = "experience") -> uuid.UUID:
"""Insert a memory unit directly, bypassing LLM retain pipeline."""
mem_id = uuid.uuid4()
await conn.execute(
"""
INSERT INTO memory_units (id, bank_id, text, fact_type, event_date, created_at, updated_at, consolidated_at)
VALUES ($1, $2, $3, $4, NOW(), NOW(), NOW(), NOW())
""",
mem_id,
bank_id,
text,
fact_type,
)
return mem_id
async def _insert_observation(
conn, bank_id: str, text: str, source_memory_ids: list[uuid.UUID]
) -> uuid.UUID:
"""Insert an observation unit directly."""
obs_id = uuid.uuid4()
await conn.execute(
"""
INSERT INTO memory_units (
id, bank_id, text, fact_type, event_date, source_memory_ids, proof_count, created_at, updated_at
) VALUES ($1, $2, $3, 'observation', NOW(), $4, $5, NOW(), NOW())
""",
obs_id,
bank_id,
text,
source_memory_ids,
len(source_memory_ids),
)
return obs_id
async def _get_observation_ids(conn, bank_id: str) -> list[str]:
rows = await conn.fetch(
"SELECT id FROM memory_units WHERE bank_id = $1 AND fact_type = 'observation'",
bank_id,
)
return [str(r["id"]) for r in rows]
async def _get_consolidated_at(conn, memory_id: uuid.UUID):
return await conn.fetchval(
"SELECT consolidated_at FROM memory_units WHERE id = $1",
memory_id,
)
async def _ensure_bank(memory: MemoryEngine, bank_id: str, request_context: RequestContext):
await memory.get_bank_profile(bank_id=bank_id, request_context=request_context)
# ---------------------------------------------------------------------------
# Tests: delete_memory_unit
# ---------------------------------------------------------------------------
class TestDeleteMemoryUnitObservationCleanup:
@pytest.mark.asyncio
async def test_deleting_source_memory_removes_observation(
self, memory: MemoryEngine, request_context: RequestContext
):
"""Deleting a source memory removes observations derived from it."""
bank_id = f"test-invalidate-del-{uuid.uuid4().hex[:8]}"
await _ensure_bank(memory, bank_id, request_context)
pool = await memory._get_pool()
async with pool.acquire() as conn:
m1 = await _insert_memory(conn, bank_id, "Alice loves hiking.")
m2 = await _insert_memory(conn, bank_id, "Alice goes hiking every weekend.")
obs_id = await _insert_observation(conn, bank_id, "Alice enjoys hiking regularly.", [m1, m2])
await memory.delete_memory_unit(str(m1), request_context=request_context)
async with pool.acquire() as conn:
obs_ids = await _get_observation_ids(conn, bank_id)
assert str(obs_id) not in obs_ids, "Observation should have been deleted"
await memory.delete_bank(bank_id, request_context=request_context)
@pytest.mark.asyncio
async def test_deleting_source_memory_resets_remaining_source_consolidated_at(
self, memory: MemoryEngine, request_context: RequestContext
):
"""After deleting a source memory, remaining source memories are reset for re-consolidation."""
bank_id = f"test-invalidate-reset-{uuid.uuid4().hex[:8]}"
await _ensure_bank(memory, bank_id, request_context)
pool = await memory._get_pool()
async with pool.acquire() as conn:
m1 = await _insert_memory(conn, bank_id, "Alice loves hiking.")
m2 = await _insert_memory(conn, bank_id, "Alice goes hiking every weekend.")
await _insert_observation(conn, bank_id, "Alice enjoys hiking regularly.", [m1, m2])
# Verify m2 starts with consolidated_at set
assert await _get_consolidated_at(conn, m2) is not None
# Patch out consolidation so it doesn't re-set consolidated_at before we can check it
with patch.object(memory, "submit_async_consolidation", new=AsyncMock()):
await memory.delete_memory_unit(str(m1), request_context=request_context)
async with pool.acquire() as conn:
# m2 should have consolidated_at reset to NULL
consolidated_at = await _get_consolidated_at(conn, m2)
assert consolidated_at is None, "Remaining source memory should be reset for re-consolidation"
await memory.delete_bank(bank_id, request_context=request_context)
@pytest.mark.asyncio
async def test_deleting_non_source_memory_leaves_observations_intact(
self, memory: MemoryEngine, request_context: RequestContext
):
"""Deleting a memory that is not a source of any observation leaves observations unchanged."""
bank_id = f"test-invalidate-noop-{uuid.uuid4().hex[:8]}"
await _ensure_bank(memory, bank_id, request_context)
pool = await memory._get_pool()
async with pool.acquire() as conn:
m1 = await _insert_memory(conn, bank_id, "Alice loves hiking.")
m2 = await _insert_memory(conn, bank_id, "Alice goes hiking every weekend.")
unrelated = await _insert_memory(conn, bank_id, "Bob likes cycling.")
obs_id = await _insert_observation(conn, bank_id, "Alice enjoys hiking regularly.", [m1, m2])
await memory.delete_memory_unit(str(unrelated), request_context=request_context)
async with pool.acquire() as conn:
obs_ids = await _get_observation_ids(conn, bank_id)
assert str(obs_id) in obs_ids, "Observation should remain untouched"
# m1 and m2 should still be consolidated
assert await _get_consolidated_at(conn, m1) is not None
assert await _get_consolidated_at(conn, m2) is not None
await memory.delete_bank(bank_id, request_context=request_context)
@pytest.mark.asyncio
async def test_deleting_sole_source_memory_removes_observation_no_remaining_reset(
self, memory: MemoryEngine, request_context: RequestContext
):
"""When an observation has only one source and it's deleted, observation is removed with no remaining memories to reset."""
bank_id = f"test-invalidate-sole-{uuid.uuid4().hex[:8]}"
await _ensure_bank(memory, bank_id, request_context)
pool = await memory._get_pool()
async with pool.acquire() as conn:
m1 = await _insert_memory(conn, bank_id, "Alice loves hiking.")
obs_id = await _insert_observation(conn, bank_id, "Alice enjoys hiking.", [m1])
await memory.delete_memory_unit(str(m1), request_context=request_context)
async with pool.acquire() as conn:
obs_ids = await _get_observation_ids(conn, bank_id)
assert str(obs_id) not in obs_ids, "Observation should have been deleted"
await memory.delete_bank(bank_id, request_context=request_context)
@pytest.mark.asyncio
async def test_deleting_observation_type_memory_does_not_trigger_invalidation(
self, memory: MemoryEngine, request_context: RequestContext
):
"""Deleting a memory with fact_type='observation' directly does not trigger invalidation logic."""
bank_id = f"test-invalidate-obstype-{uuid.uuid4().hex[:8]}"
await _ensure_bank(memory, bank_id, request_context)
pool = await memory._get_pool()
async with pool.acquire() as conn:
m1 = await _insert_memory(conn, bank_id, "Alice loves hiking.")
obs_id = await _insert_observation(conn, bank_id, "Alice enjoys hiking.", [m1])
# Delete the observation directly (not the source memory)
await memory.delete_memory_unit(str(obs_id), request_context=request_context)
async with pool.acquire() as conn:
# Source memory should still be consolidated (not reset)
assert await _get_consolidated_at(conn, m1) is not None
obs_ids = await _get_observation_ids(conn, bank_id)
assert str(obs_id) not in obs_ids
await memory.delete_bank(bank_id, request_context=request_context)
# ---------------------------------------------------------------------------
# Tests: delete_document
# ---------------------------------------------------------------------------
class TestDeleteDocumentObservationCleanup:
@pytest.mark.asyncio
async def test_deleting_document_removes_observations(
self, memory: MemoryEngine, request_context: RequestContext
):
"""Deleting a document removes observations derived from its memory units."""
bank_id = f"test-invalidate-doc-{uuid.uuid4().hex[:8]}"
await _ensure_bank(memory, bank_id, request_context)
pool = await memory._get_pool()
# Create a document and attach memories to it
async with pool.acquire() as conn:
doc_id = str(uuid.uuid4()) # documents.id is TEXT
await conn.execute(
"""
INSERT INTO documents (id, bank_id, original_text, content_hash, created_at, updated_at)
VALUES ($1, $2, 'some doc', 'hash123', NOW(), NOW())
""",
doc_id,
bank_id,
)
m1 = uuid.uuid4()
m2 = uuid.uuid4()
for mem_id, text in [(m1, "Alice loves hiking."), (m2, "Alice goes hiking every weekend.")]:
await conn.execute(
"""
INSERT INTO memory_units (id, bank_id, text, fact_type, event_date, document_id, created_at, updated_at, consolidated_at)
VALUES ($1, $2, $3, 'experience', NOW(), $4, NOW(), NOW(), NOW())
""",
mem_id,
bank_id,
text,
doc_id,
)
# Standalone memory (not in document)
m3 = await _insert_memory(conn, bank_id, "Alice is an avid outdoor person.")
# Observation referencing both doc memories and the standalone memory
obs_id = await _insert_observation(
conn, bank_id, "Alice enjoys outdoor activities.", [m1, m2, m3]
)
# Patch out consolidation so it doesn't re-set consolidated_at before we can check it
with patch.object(memory, "submit_async_consolidation", new=AsyncMock()):
await memory.delete_document(str(doc_id), bank_id, request_context=request_context)
async with pool.acquire() as conn:
obs_ids = await _get_observation_ids(conn, bank_id)
assert str(obs_id) not in obs_ids, "Observation should have been deleted"
# m3 (remaining source) should be reset for re-consolidation
consolidated_at = await _get_consolidated_at(conn, m3)
assert consolidated_at is None, "Remaining source memory should be reset"
await memory.delete_bank(bank_id, request_context=request_context)
# ---------------------------------------------------------------------------
# Tests: delete_bank with fact_type filter
# ---------------------------------------------------------------------------
class TestDeleteBankByTypeObservationCleanup:
@pytest.mark.asyncio
async def test_clearing_experience_memories_removes_affected_observations(
self, memory: MemoryEngine, request_context: RequestContext
):
"""Clearing all experience memories removes observations sourced from them."""
bank_id = f"test-invalidate-banktype-{uuid.uuid4().hex[:8]}"
await _ensure_bank(memory, bank_id, request_context)
pool = await memory._get_pool()
async with pool.acquire() as conn:
exp1 = await _insert_memory(conn, bank_id, "Alice went hiking last week.", "experience")
world1 = await _insert_memory(conn, bank_id, "Alice is a hiker.", "world")
obs_id = await _insert_observation(
conn, bank_id, "Alice is a regular hiker.", [exp1, world1]
)
# Patch out consolidation so it doesn't re-set consolidated_at before we can check it
with patch.object(memory, "submit_async_consolidation", new=AsyncMock()):
await memory.delete_bank(bank_id, fact_type="experience", request_context=request_context)
async with pool.acquire() as conn:
obs_ids = await _get_observation_ids(conn, bank_id)
assert str(obs_id) not in obs_ids, "Observation should have been deleted"
# world1 (remaining source) should be reset for re-consolidation
consolidated_at = await _get_consolidated_at(conn, world1)
assert consolidated_at is None, "World memory should be reset for re-consolidation"
await memory.delete_bank(bank_id, request_context=request_context)
@pytest.mark.asyncio
async def test_clearing_unrelated_type_leaves_observations_intact(
self, memory: MemoryEngine, request_context: RequestContext
):
"""Clearing memories of a type that is not a source of any observation leaves observations untouched."""
bank_id = f"test-invalidate-banktype-noop-{uuid.uuid4().hex[:8]}"
await _ensure_bank(memory, bank_id, request_context)
pool = await memory._get_pool()
async with pool.acquire() as conn:
world1 = await _insert_memory(conn, bank_id, "Alice is a hiker.", "world")
obs_id = await _insert_observation(conn, bank_id, "Alice is a regular hiker.", [world1])
# Deleting 'experience' type should not affect observations sourced only from 'world'
await memory.delete_bank(bank_id, fact_type="experience", request_context=request_context)
async with pool.acquire() as conn:
obs_ids = await _get_observation_ids(conn, bank_id)
assert str(obs_id) in obs_ids, "Observation should remain untouched"
await memory.delete_bank(bank_id, request_context=request_context)
# ---------------------------------------------------------------------------
# Tests: clear_observations_for_memory
# ---------------------------------------------------------------------------
class TestClearObservationsForMemory:
@pytest.mark.asyncio
async def test_clears_observations_and_resets_all_source_memories(
self, memory: MemoryEngine, request_context: RequestContext
):
"""Clearing observations for a memory deletes them and resets all related source memories."""
bank_id = f"test-clear-obs-mem-{uuid.uuid4().hex[:8]}"
await _ensure_bank(memory, bank_id, request_context)
pool = await memory._get_pool()
async with pool.acquire() as conn:
m1 = await _insert_memory(conn, bank_id, "Alice loves hiking.")
m2 = await _insert_memory(conn, bank_id, "Alice hikes every weekend.")
obs_id = await _insert_observation(conn, bank_id, "Alice is an avid hiker.", [m1, m2])
# Patch out consolidation so it doesn't re-set consolidated_at before we can check it
with patch.object(memory, "submit_async_consolidation", new=AsyncMock()):
result = await memory.clear_observations_for_memory(
bank_id, str(m1), request_context=request_context
)
assert result["deleted_count"] == 1
async with pool.acquire() as conn:
obs_ids = await _get_observation_ids(conn, bank_id)
assert str(obs_id) not in obs_ids, "Observation should be deleted"
# Both m1 (target) and m2 (remaining source) should be reset
assert await _get_consolidated_at(conn, m1) is None, "Target memory should be reset"
assert await _get_consolidated_at(conn, m2) is None, "Remaining source should be reset"
await memory.delete_bank(bank_id, request_context=request_context)
@pytest.mark.asyncio
async def test_no_observations_returns_zero(
self, memory: MemoryEngine, request_context: RequestContext
):
"""Returns 0 when the memory has no associated observations."""
bank_id = f"test-clear-obs-noop-{uuid.uuid4().hex[:8]}"
await _ensure_bank(memory, bank_id, request_context)
pool = await memory._get_pool()
async with pool.acquire() as conn:
m1 = await _insert_memory(conn, bank_id, "Alice loves hiking.")
result = await memory.clear_observations_for_memory(
bank_id, str(m1), request_context=request_context
)
assert result["deleted_count"] == 0
async with pool.acquire() as conn:
# Memory should still be consolidated (no observations were cleared)
assert await _get_consolidated_at(conn, m1) is not None
await memory.delete_bank(bank_id, request_context=request_context)
@pytest.mark.asyncio
async def test_only_clears_observations_referencing_target_memory(
self, memory: MemoryEngine, request_context: RequestContext
):
"""Clearing observations for m1 does not affect observations that only reference m2."""
bank_id = f"test-clear-obs-selective-{uuid.uuid4().hex[:8]}"
await _ensure_bank(memory, bank_id, request_context)
pool = await memory._get_pool()
async with pool.acquire() as conn:
m1 = await _insert_memory(conn, bank_id, "Alice loves hiking.")
m2 = await _insert_memory(conn, bank_id, "Alice hikes every weekend.")
m3 = await _insert_memory(conn, bank_id, "Alice climbed a mountain.")
obs1_id = await _insert_observation(conn, bank_id, "Alice is an avid hiker.", [m1, m2])
obs2_id = await _insert_observation(conn, bank_id, "Alice is a mountaineer.", [m3])
result = await memory.clear_observations_for_memory(
bank_id, str(m1), request_context=request_context
)
assert result["deleted_count"] == 1
async with pool.acquire() as conn:
obs_ids = await _get_observation_ids(conn, bank_id)
assert str(obs1_id) not in obs_ids, "obs1 (references m1) should be deleted"
assert str(obs2_id) in obs_ids, "obs2 (does not reference m1) should remain"
# m3 should still be consolidated
assert await _get_consolidated_at(conn, m3) is not None
await memory.delete_bank(bank_id, request_context=request_context)
@pytest.mark.asyncio
async def test_multiple_observations_for_same_memory_all_cleared(
self, memory: MemoryEngine, request_context: RequestContext
):
"""All observations referencing the target memory are cleared in one call."""
bank_id = f"test-clear-obs-multi-{uuid.uuid4().hex[:8]}"
await _ensure_bank(memory, bank_id, request_context)
pool = await memory._get_pool()
async with pool.acquire() as conn:
m1 = await _insert_memory(conn, bank_id, "Alice loves hiking.")
m2 = await _insert_memory(conn, bank_id, "Alice hikes every weekend.")
obs1_id = await _insert_observation(conn, bank_id, "Alice hikes often.", [m1])
obs2_id = await _insert_observation(conn, bank_id, "Alice is outdoorsy.", [m1, m2])
# Patch out consolidation so it doesn't re-set consolidated_at before we can check it
with patch.object(memory, "submit_async_consolidation", new=AsyncMock()):
result = await memory.clear_observations_for_memory(
bank_id, str(m1), request_context=request_context
)
assert result["deleted_count"] == 2
async with pool.acquire() as conn:
obs_ids = await _get_observation_ids(conn, bank_id)
assert str(obs1_id) not in obs_ids
assert str(obs2_id) not in obs_ids
# m1 and m2 should both be reset
assert await _get_consolidated_at(conn, m1) is None
assert await _get_consolidated_at(conn, m2) is None
await memory.delete_bank(bank_id, request_context=request_context)
# ---------------------------------------------------------------------------
# Tests: update_document
# ---------------------------------------------------------------------------
async def _insert_document_with_memories(
conn, bank_id: str, doc_id: str, memories: list[tuple[str, str]]
) -> list[uuid.UUID]:
"""Insert a document and attach memory units to it. Returns list of memory UUIDs."""
await conn.execute(
"""
INSERT INTO documents (id, bank_id, original_text, content_hash, created_at, updated_at)
VALUES ($1, $2, 'some doc', 'hash123', NOW(), NOW())
""",
doc_id,
bank_id,
)
mem_ids = []
for text, fact_type in memories:
mem_id = uuid.uuid4()
await conn.execute(
"""
INSERT INTO memory_units (id, bank_id, text, fact_type, event_date, document_id, created_at, updated_at, consolidated_at)
VALUES ($1, $2, $3, $4, NOW(), $5, NOW(), NOW(), NOW())
""",
mem_id,
bank_id,
text,
fact_type,
doc_id,
)
mem_ids.append(mem_id)
return mem_ids
class TestUpdateDocumentTagsObservationCleanup:
@pytest.mark.asyncio
async def test_update_tags_returns_updated_document(
self, memory: MemoryEngine, request_context: RequestContext
):
"""update_document returns the updated document with new tags."""
bank_id = f"test-tag-update-basic-{uuid.uuid4().hex[:8]}"
await _ensure_bank(memory, bank_id, request_context)
pool = await memory._get_pool()
async with pool.acquire() as conn:
doc_id = f"doc-{uuid.uuid4().hex[:8]}"
await _insert_document_with_memories(conn, bank_id, doc_id, [("Alice loves hiking.", "experience")])
result = await memory.update_document(
doc_id, bank_id, tags=["new-tag"], request_context=request_context
)
assert result is True
await memory.delete_bank(bank_id, request_context=request_context)
@pytest.mark.asyncio
async def test_update_tags_returns_none_for_missing_document(
self, memory: MemoryEngine, request_context: RequestContext
):
"""update_document returns False when document does not exist."""
bank_id = f"test-tag-update-missing-{uuid.uuid4().hex[:8]}"
await _ensure_bank(memory, bank_id, request_context)
result = await memory.update_document(
"nonexistent-doc", bank_id, tags=["tag"], request_context=request_context
)
assert result is False
await memory.delete_bank(bank_id, request_context=request_context)
@pytest.mark.asyncio
async def test_update_tags_propagates_to_memory_units(
self, memory: MemoryEngine, request_context: RequestContext
):
"""Changing document tags also updates all associated memory unit tags."""
bank_id = f"test-tag-update-propagate-{uuid.uuid4().hex[:8]}"
await _ensure_bank(memory, bank_id, request_context)
pool = await memory._get_pool()
async with pool.acquire() as conn:
doc_id = f"doc-{uuid.uuid4().hex[:8]}"
mem_ids = await _insert_document_with_memories(
conn, bank_id, doc_id, [("Alice loves hiking.", "experience"), ("Alice hikes weekly.", "world")]
)
with patch.object(memory, "submit_async_consolidation", new=AsyncMock()):
await memory.update_document(
doc_id, bank_id, tags=["new-tag"], request_context=request_context
)
async with pool.acquire() as conn:
for mem_id in mem_ids:
tags = await conn.fetchval(
"SELECT tags FROM memory_units WHERE id = $1", mem_id
)
assert list(tags) == ["new-tag"], f"Memory unit {mem_id} should have updated tags"
await memory.delete_bank(bank_id, request_context=request_context)
@pytest.mark.asyncio
async def test_update_tags_invalidates_observations(
self, memory: MemoryEngine, request_context: RequestContext
):
"""Observations referencing the document's memory units are deleted on tag change."""
bank_id = f"test-tag-update-obs-{uuid.uuid4().hex[:8]}"
await _ensure_bank(memory, bank_id, request_context)
pool = await memory._get_pool()
async with pool.acquire() as conn:
doc_id = f"doc-{uuid.uuid4().hex[:8]}"
mem_ids = await _insert_document_with_memories(
conn, bank_id, doc_id, [("Alice loves hiking.", "experience")]
)
obs_id = await _insert_observation(conn, bank_id, "Alice is a hiker.", mem_ids)
with patch.object(memory, "submit_async_consolidation", new=AsyncMock()):
await memory.update_document(
doc_id, bank_id, tags=["new-tag"], request_context=request_context
)
async with pool.acquire() as conn:
obs_ids = await _get_observation_ids(conn, bank_id)
assert str(obs_id) not in obs_ids, "Observation should have been invalidated"
await memory.delete_bank(bank_id, request_context=request_context)
@pytest.mark.asyncio
async def test_update_tags_resets_consolidated_at_on_affected_units(
self, memory: MemoryEngine, request_context: RequestContext
):
"""Affected memory units get consolidated_at reset for re-consolidation under new tags."""
bank_id = f"test-tag-update-reset-{uuid.uuid4().hex[:8]}"
await _ensure_bank(memory, bank_id, request_context)
pool = await memory._get_pool()
async with pool.acquire() as conn:
doc_id = f"doc-{uuid.uuid4().hex[:8]}"
mem_ids = await _insert_document_with_memories(
conn, bank_id, doc_id, [("Alice loves hiking.", "experience")]
)
obs_id = await _insert_observation(conn, bank_id, "Alice is a hiker.", mem_ids)
# Verify memory starts consolidated
assert await _get_consolidated_at(conn, mem_ids[0]) is not None
with patch.object(memory, "submit_async_consolidation", new=AsyncMock()):
await memory.update_document(
doc_id, bank_id, tags=["new-tag"], request_context=request_context
)
async with pool.acquire() as conn:
consolidated_at = await _get_consolidated_at(conn, mem_ids[0])
assert consolidated_at is None, "Memory unit should be reset for re-consolidation"
await memory.delete_bank(bank_id, request_context=request_context)
@pytest.mark.asyncio
async def test_update_tags_triggers_consolidation_when_observations_invalidated(
self, memory: MemoryEngine, request_context: RequestContext
):
"""submit_async_consolidation is called when observations are invalidated."""
bank_id = f"test-tag-update-cons-{uuid.uuid4().hex[:8]}"
await _ensure_bank(memory, bank_id, request_context)
pool = await memory._get_pool()
async with pool.acquire() as conn:
doc_id = f"doc-{uuid.uuid4().hex[:8]}"
mem_ids = await _insert_document_with_memories(
conn, bank_id, doc_id, [("Alice loves hiking.", "experience")]
)
await _insert_observation(conn, bank_id, "Alice is a hiker.", mem_ids)
with patch.object(memory, "submit_async_consolidation", new=AsyncMock()) as mock_consolidate:
await memory.update_document(
doc_id, bank_id, tags=["new-tag"], request_context=request_context
)
mock_consolidate.assert_awaited_once()
await memory.delete_bank(bank_id, request_context=request_context)
@pytest.mark.asyncio
async def test_update_tags_no_consolidation_when_no_observations(
self, memory: MemoryEngine, request_context: RequestContext
):
"""submit_async_consolidation is NOT called when no observations are invalidated."""
bank_id = f"test-tag-update-nocons-{uuid.uuid4().hex[:8]}"
await _ensure_bank(memory, bank_id, request_context)
pool = await memory._get_pool()
async with pool.acquire() as conn:
doc_id = f"doc-{uuid.uuid4().hex[:8]}"
await _insert_document_with_memories(
conn, bank_id, doc_id, [("Alice loves hiking.", "experience")]
)
# No observations inserted
with patch.object(memory, "submit_async_consolidation", new=AsyncMock()) as mock_consolidate:
await memory.update_document(
doc_id, bank_id, tags=["new-tag"], request_context=request_context
)
mock_consolidate.assert_not_awaited()
await memory.delete_bank(bank_id, request_context=request_context)
@pytest.mark.asyncio
async def test_update_tags_resets_co_source_memories_from_other_documents(
self, memory: MemoryEngine, request_context: RequestContext
):
"""Co-source memories from other documents that shared an invalidated observation are also reset."""
bank_id = f"test-tag-update-cosource-{uuid.uuid4().hex[:8]}"
await _ensure_bank(memory, bank_id, request_context)
pool = await memory._get_pool()
async with pool.acquire() as conn:
doc_id = f"doc-{uuid.uuid4().hex[:8]}"
doc_mem_ids = await _insert_document_with_memories(
conn, bank_id, doc_id, [("Alice loves hiking.", "experience")]
)
# Unrelated memory from another document — co-sourced in the same observation
other_mem = await _insert_memory(conn, bank_id, "Alice also rock-climbs.")
obs_id = await _insert_observation(
conn, bank_id, "Alice loves outdoor activities.", doc_mem_ids + [other_mem]
)
# Verify other_mem starts consolidated
assert await _get_consolidated_at(conn, other_mem) is not None
with patch.object(memory, "submit_async_consolidation", new=AsyncMock()):
await memory.update_document(
doc_id, bank_id, tags=["new-tag"], request_context=request_context
)
async with pool.acquire() as conn:
obs_ids = await _get_observation_ids(conn, bank_id)
assert str(obs_id) not in obs_ids, "Observation should have been invalidated"
# other_mem (co-source from another document) must also be reset
consolidated_at = await _get_consolidated_at(conn, other_mem)
assert consolidated_at is None, "Co-source memory from other document should be reset"
await memory.delete_bank(bank_id, request_context=request_context)
@pytest.mark.asyncio
async def test_update_tags_does_not_affect_unrelated_observations(
self, memory: MemoryEngine, request_context: RequestContext
):
"""Observations referencing memories from a different document are not affected."""
bank_id = f"test-tag-update-unrelated-{uuid.uuid4().hex[:8]}"
await _ensure_bank(memory, bank_id, request_context)
pool = await memory._get_pool()
async with pool.acquire() as conn:
doc_id = f"doc-{uuid.uuid4().hex[:8]}"
mem_ids = await _insert_document_with_memories(
conn, bank_id, doc_id, [("Alice loves hiking.", "experience")]
)
# Unrelated memory not in the document
unrelated = await _insert_memory(conn, bank_id, "Bob likes cycling.")
unrelated_obs_id = await _insert_observation(
conn, bank_id, "Bob is a cyclist.", [unrelated]
)
with patch.object(memory, "submit_async_consolidation", new=AsyncMock()):
await memory.update_document(
doc_id, bank_id, tags=["new-tag"], request_context=request_context
)
async with pool.acquire() as conn:
obs_ids = await _get_observation_ids(conn, bank_id)
assert str(unrelated_obs_id) in obs_ids, "Unrelated observation should remain untouched"
await memory.delete_bank(bank_id, request_context=request_context)
@@ -1,311 +0,0 @@
"""Tests for operation cancellation when a bank is deleted.
Covers:
- CASCADE DELETE: deleting a bank removes async_operations and webhooks rows
- _check_op_alive: returns True when op exists, False when deleted
- _mark_operation_completed / _mark_operation_failed: graceful no-op when row is gone
- Consolidation checkpoint: stops early after a batch commit if op was deleted
- Retain checkpoint: stops between sub-batches if op was deleted
"""
import uuid
from unittest.mock import AsyncMock, patch
import pytest
import pytest_asyncio
from hindsight_api.engine.memory_engine import MemoryEngine
pytestmark = pytest.mark.xdist_group("op_cancellation_tests")
_BANK_PREFIX = "test-op-cancel"
@pytest_asyncio.fixture
async def pool(pg0_db_url):
import asyncpg
from hindsight_api.pg0 import resolve_database_url
resolved_url = await resolve_database_url(pg0_db_url)
p = await asyncpg.create_pool(resolved_url, min_size=1, max_size=5, command_timeout=30)
yield p
await p.close()
@pytest_asyncio.fixture(autouse=True)
async def cleanup(pool):
"""Remove test rows before and after each test."""
await pool.execute(f"DELETE FROM banks WHERE bank_id LIKE '{_BANK_PREFIX}%'")
yield
await pool.execute(f"DELETE FROM banks WHERE bank_id LIKE '{_BANK_PREFIX}%'")
async def _insert_bank(pool, bank_id: str):
await pool.execute(
"INSERT INTO banks (bank_id, name) VALUES ($1, $2) ON CONFLICT DO NOTHING",
bank_id,
bank_id,
)
async def _insert_op(pool, bank_id: str, op_id: uuid.UUID | None = None) -> uuid.UUID:
op_id = op_id or uuid.uuid4()
await pool.execute(
"""
INSERT INTO async_operations (operation_id, bank_id, operation_type, status)
VALUES ($1, $2, 'consolidation', 'processing')
""",
op_id,
bank_id,
)
return op_id
# ---------------------------------------------------------------------------
# CASCADE DELETE tests
# ---------------------------------------------------------------------------
class TestCascadeDeleteOnBankDeletion:
@pytest.mark.asyncio
async def test_bank_deletion_cascades_to_async_operations(self, pool):
bank_id = f"{_BANK_PREFIX}-{uuid.uuid4().hex[:8]}"
await _insert_bank(pool, bank_id)
op_id = await _insert_op(pool, bank_id)
# Verify op exists
row = await pool.fetchrow("SELECT operation_id FROM async_operations WHERE operation_id = $1", op_id)
assert row is not None
# Delete the bank — should cascade to async_operations
await pool.execute("DELETE FROM banks WHERE bank_id = $1", bank_id)
row = await pool.fetchrow("SELECT operation_id FROM async_operations WHERE operation_id = $1", op_id)
assert row is None, "async_operations row should be deleted by CASCADE"
@pytest.mark.asyncio
async def test_bank_deletion_cascades_to_webhooks(self, pool):
bank_id = f"{_BANK_PREFIX}-{uuid.uuid4().hex[:8]}"
await _insert_bank(pool, bank_id)
webhook_id = uuid.uuid4()
await pool.execute(
"""
INSERT INTO webhooks (id, bank_id, url, event_types)
VALUES ($1, $2, 'https://example.com/hook', '{}')
""",
webhook_id,
bank_id,
)
row = await pool.fetchrow("SELECT id FROM webhooks WHERE id = $1", webhook_id)
assert row is not None
await pool.execute("DELETE FROM banks WHERE bank_id = $1", bank_id)
row = await pool.fetchrow("SELECT id FROM webhooks WHERE id = $1", webhook_id)
assert row is None, "webhooks row should be deleted by CASCADE"
# ---------------------------------------------------------------------------
# _check_op_alive tests
# ---------------------------------------------------------------------------
class TestCheckOpAlive:
@pytest.mark.asyncio
async def test_returns_true_when_op_exists(self, memory: MemoryEngine, request_context):
bank_id = f"{_BANK_PREFIX}-{uuid.uuid4().hex[:8]}"
await memory.get_bank_profile(bank_id=bank_id, request_context=request_context)
op_id = uuid.uuid4()
async with memory._pool.acquire() as conn:
await conn.execute(
"""
INSERT INTO async_operations (operation_id, bank_id, operation_type, status)
VALUES ($1, $2, 'consolidation', 'processing')
""",
op_id,
bank_id,
)
assert await memory._check_op_alive(str(op_id)) is True
@pytest.mark.asyncio
async def test_returns_false_when_op_deleted(self, memory: MemoryEngine, request_context):
bank_id = f"{_BANK_PREFIX}-{uuid.uuid4().hex[:8]}"
await memory.get_bank_profile(bank_id=bank_id, request_context=request_context)
op_id = uuid.uuid4()
async with memory._pool.acquire() as conn:
await conn.execute(
"""
INSERT INTO async_operations (operation_id, bank_id, operation_type, status)
VALUES ($1, $2, 'consolidation', 'processing')
""",
op_id,
bank_id,
)
await conn.execute("DELETE FROM async_operations WHERE operation_id = $1", op_id)
assert await memory._check_op_alive(str(op_id)) is False
@pytest.mark.asyncio
async def test_returns_false_after_bank_cascade_delete(self, memory: MemoryEngine, request_context):
bank_id = f"{_BANK_PREFIX}-{uuid.uuid4().hex[:8]}"
await memory.get_bank_profile(bank_id=bank_id, request_context=request_context)
op_id = uuid.uuid4()
async with memory._pool.acquire() as conn:
await conn.execute(
"""
INSERT INTO async_operations (operation_id, bank_id, operation_type, status)
VALUES ($1, $2, 'consolidation', 'processing')
""",
op_id,
bank_id,
)
# Delete the bank — cascades to the op row
await memory.delete_bank(bank_id=bank_id, request_context=request_context)
assert await memory._check_op_alive(str(op_id)) is False
# ---------------------------------------------------------------------------
# _mark_operation_completed / _mark_operation_failed graceful no-op
# ---------------------------------------------------------------------------
class TestMarkOperationGracefulOnMissingRow:
@pytest.mark.asyncio
async def test_mark_completed_does_not_raise_when_row_missing(self, memory: MemoryEngine):
# Row never existed — should log and return cleanly
missing_id = str(uuid.uuid4())
await memory._mark_operation_completed(missing_id) # no exception
@pytest.mark.asyncio
async def test_mark_failed_does_not_raise_when_row_missing(self, memory: MemoryEngine):
missing_id = str(uuid.uuid4())
await memory._mark_operation_failed(missing_id, "some error", "traceback here") # no exception
@pytest.mark.asyncio
async def test_mark_completed_and_fire_webhook_does_not_raise_when_row_missing(
self, memory: MemoryEngine
):
missing_id = str(uuid.uuid4())
await memory._mark_operation_completed_and_fire_webhook(
operation_id=missing_id,
bank_id="nonexistent-bank",
status="completed",
result=None,
) # no exception
# ---------------------------------------------------------------------------
# Consolidation checkpoint
# ---------------------------------------------------------------------------
class TestConsolidationCheckpoint:
@pytest.mark.asyncio
async def test_consolidation_stops_early_when_op_cancelled(self, memory: MemoryEngine, request_context):
"""Consolidation returns 'cancelled' status after the first batch if _check_op_alive is False."""
from hindsight_api.config import _get_raw_config
from hindsight_api.engine.consolidation.consolidator import run_consolidation_job
config = _get_raw_config()
original = config.enable_observations
config.enable_observations = True
try:
bank_id = f"{_BANK_PREFIX}-{uuid.uuid4().hex[:8]}"
await memory.get_bank_profile(bank_id=bank_id, request_context=request_context)
# Insert a few unconsolidated memories directly so we control the batch without LLM
async with memory._pool.acquire() as conn:
for i in range(3):
await conn.execute(
"""
INSERT INTO memory_units
(id, bank_id, text, fact_type, created_at, updated_at)
VALUES (gen_random_uuid(), $1, $2, 'experience', NOW(), NOW())
""",
bank_id,
f"Test memory {i} for cancellation test",
)
op_id = str(uuid.uuid4())
call_count = 0
async def _fake_check(operation_id: str) -> bool:
nonlocal call_count
call_count += 1
# Return False on the very first checkpoint call
return False
with patch.object(memory, "_check_op_alive", side_effect=_fake_check):
result = await run_consolidation_job(
memory_engine=memory,
bank_id=bank_id,
request_context=request_context,
operation_id=op_id,
)
assert result["status"] == "cancelled"
assert call_count >= 1
finally:
config.enable_observations = original
# ---------------------------------------------------------------------------
# Retain checkpoint
# ---------------------------------------------------------------------------
class TestRetainCheckpoint:
@pytest.mark.asyncio
async def test_retain_stops_between_sub_batches_when_cancelled(
self, memory: MemoryEngine, request_context
):
"""retain_batch_async returns partial results if _check_op_alive is False between sub-batches."""
from hindsight_api.config import _get_raw_config
bank_id = f"{_BANK_PREFIX}-{uuid.uuid4().hex[:8]}"
await memory.get_bank_profile(bank_id=bank_id, request_context=request_context)
# Force sub-batch splitting by temporarily lowering the token threshold
config = _get_raw_config()
original_tokens = config.retain_batch_tokens
# Set threshold very low so each item becomes its own sub-batch
config.retain_batch_tokens = 1
try:
op_id = str(uuid.uuid4())
check_calls = 0
async def _fake_check(operation_id: str) -> bool:
nonlocal check_calls
check_calls += 1
# Cancel after the first sub-batch completes
return check_calls <= 1
contents = [
{"content": f"Memory item {i} about something interesting."} for i in range(4)
]
with patch.object(memory, "_check_op_alive", side_effect=_fake_check):
result = await memory.retain_batch_async(
bank_id=bank_id,
contents=contents,
request_context=request_context,
operation_id=op_id,
)
# Should have stopped early: fewer results than total items
assert len(result) < len(contents), (
f"Expected early stop but got {len(result)}/{len(contents)} results"
)
assert check_calls >= 1
finally:
config.retain_batch_tokens = original_tokens
@@ -1,71 +0,0 @@
"""
Regression test for UnboundLocalError in recall when the reranker raises.
Before the fix, `scored_results` and `pre_filtered_count` were only assigned
inside the `try` block, but referenced in the `finally` block. If
`reranker_instance.rerank()` (or `ensure_initialized()`) raised, the `finally`
block crashed with `UnboundLocalError` instead of propagating the original
exception.
Fix: initialise both variables to safe defaults before the try/finally block.
"""
from datetime import datetime, timezone
from unittest.mock import AsyncMock, patch
import pytest
@pytest.mark.asyncio
async def test_recall_reranker_error_does_not_raise_unbound_local(memory, request_context):
"""Recall must propagate the reranker's exception, not an UnboundLocalError."""
bank_id = f"test_reranker_err_{datetime.now(timezone.utc).timestamp()}"
try:
await memory.retain_async(
bank_id=bank_id,
content="Paris is the capital of France",
request_context=request_context,
)
# Simulate a reranker failure (e.g. Cohere API error on empty/small candidate set)
rerank_mock = AsyncMock(side_effect=RuntimeError("reranker API error"))
memory._cross_encoder_reranker._initialized = True # skip ensure_initialized
with patch.object(memory._cross_encoder_reranker, "rerank", rerank_mock):
with pytest.raises(Exception, match="reranker API error"):
await memory.recall_async(
bank_id=bank_id,
query="capital of France",
request_context=request_context,
)
finally:
await memory.delete_bank(bank_id, request_context=request_context)
@pytest.mark.asyncio
async def test_recall_reranker_init_error_does_not_raise_unbound_local(memory, request_context):
"""Same regression when ensure_initialized() raises (before pre_filtered_count is set)."""
bank_id = f"test_reranker_init_err_{datetime.now(timezone.utc).timestamp()}"
try:
await memory.retain_async(
bank_id=bank_id,
content="Paris is the capital of France",
request_context=request_context,
)
init_mock = AsyncMock(side_effect=RuntimeError("reranker init failed"))
memory._cross_encoder_reranker._initialized = False
with patch.object(memory._cross_encoder_reranker, "ensure_initialized", init_mock):
with pytest.raises(Exception, match="reranker init failed"):
await memory.recall_async(
bank_id=bank_id,
query="capital of France",
request_context=request_context,
)
finally:
await memory.delete_bank(bank_id, request_context=request_context)
@@ -1,170 +0,0 @@
"""Tests for source_facts token limiting in recall.
Covers:
- max_source_facts_tokens: total token budget across all source facts
- max_source_facts_tokens_per_observation: per-observation cap
Both parameters are tested at the recall_async level and verified to produce
fewer source facts when the budget is tight vs. unlimited.
"""
import pytest
from hindsight_api.config import _get_raw_config
from hindsight_api.engine.memory_engine import Budget
@pytest.fixture(autouse=True)
def enable_observations():
config = _get_raw_config()
original = config.enable_observations
config.enable_observations = True
yield
config.enable_observations = original
async def _setup_bank_with_observations(memory, bank_id, request_context):
"""Retain several memories and trigger consolidation to produce observations with source facts."""
contents = [
"Alice is a software engineer who loves Python programming.",
"Alice has been working at TechCorp for 5 years.",
"Alice recently completed a machine learning certification course.",
"Alice mentors junior developers on the team.",
"Alice prefers functional programming patterns in her code.",
]
for content in contents:
await memory.retain_async(
bank_id=bank_id,
content=content,
request_context=request_context,
)
await memory.run_consolidation(bank_id=bank_id, request_context=request_context)
class TestRecallSourceFactsPerObservationCap:
@pytest.mark.asyncio
async def test_per_observation_cap_reduces_source_facts(self, memory, request_context):
"""A tight per-observation token cap should return fewer source facts than unlimited."""
bank_id = "test-sf-per-obs-cap"
try:
await _setup_bank_with_observations(memory, bank_id, request_context)
result_limited = await memory.recall_async(
bank_id=bank_id,
query="Alice engineer",
fact_type=["observation"],
max_tokens=4096,
include_source_facts=True,
max_source_facts_tokens_per_observation=1, # Effectively cuts all source facts
budget=Budget.MID,
request_context=request_context,
)
result_unlimited = await memory.recall_async(
bank_id=bank_id,
query="Alice engineer",
fact_type=["observation"],
max_tokens=4096,
include_source_facts=True,
max_source_facts_tokens_per_observation=-1,
budget=Budget.MID,
request_context=request_context,
)
unlimited_count = len(result_unlimited.source_facts) if result_unlimited.source_facts else 0
limited_count = len(result_limited.source_facts) if result_limited.source_facts else 0
if unlimited_count > 0:
assert limited_count <= unlimited_count, (
f"Per-observation cap should yield fewer source facts ({limited_count} <= {unlimited_count})"
)
finally:
await memory.delete_bank(bank_id, request_context=request_context)
@pytest.mark.asyncio
async def test_per_observation_cap_does_not_mix_between_observations(self, memory, request_context):
"""Each observation's source facts are capped independently — not as a shared pool."""
bank_id = "test-sf-per-obs-independent"
try:
await _setup_bank_with_observations(memory, bank_id, request_context)
# With a generous per-observation limit each observation can have facts;
# with a global limit of 1 token the first observation would consume the whole budget.
result_per_obs = await memory.recall_async(
bank_id=bank_id,
query="Alice engineer",
fact_type=["observation"],
max_tokens=4096,
include_source_facts=True,
max_source_facts_tokens=4096, # large global budget
max_source_facts_tokens_per_observation=512, # reasonable per-obs limit
budget=Budget.MID,
request_context=request_context,
)
# Should not raise; source_facts may be populated for multiple observations
assert result_per_obs.source_facts is not None or len(result_per_obs.results) == 0
finally:
await memory.delete_bank(bank_id, request_context=request_context)
class TestRecallSourceFactsTotalBudget:
@pytest.mark.asyncio
async def test_total_budget_limits_source_facts(self, memory, request_context):
"""A tight total token budget should return fewer source facts than unlimited."""
bank_id = "test-sf-total-budget"
try:
await _setup_bank_with_observations(memory, bank_id, request_context)
result_tight = await memory.recall_async(
bank_id=bank_id,
query="Alice engineer",
fact_type=["observation"],
max_tokens=4096,
include_source_facts=True,
max_source_facts_tokens=1, # Effectively cuts all source facts
budget=Budget.MID,
request_context=request_context,
)
result_unlimited = await memory.recall_async(
bank_id=bank_id,
query="Alice engineer",
fact_type=["observation"],
max_tokens=4096,
include_source_facts=True,
max_source_facts_tokens=-1,
budget=Budget.MID,
request_context=request_context,
)
unlimited_count = len(result_unlimited.source_facts) if result_unlimited.source_facts else 0
tight_count = len(result_tight.source_facts) if result_tight.source_facts else 0
if unlimited_count > 0:
assert tight_count <= unlimited_count, (
f"Total budget should yield fewer source facts ({tight_count} <= {unlimited_count})"
)
finally:
await memory.delete_bank(bank_id, request_context=request_context)
@pytest.mark.asyncio
async def test_no_source_facts_without_flag(self, memory, request_context):
"""source_facts should be None when include_source_facts is not set."""
bank_id = "test-sf-no-flag"
try:
await _setup_bank_with_observations(memory, bank_id, request_context)
result = await memory.recall_async(
bank_id=bank_id,
query="Alice engineer",
fact_type=["observation"],
max_tokens=4096,
include_source_facts=False, # default
budget=Budget.MID,
request_context=request_context,
)
assert result.source_facts is None or len(result.source_facts) == 0
finally:
await memory.delete_bank(bank_id, request_context=request_context)
-794
View File
@@ -1,794 +0,0 @@
"""Tests for the webhook system.
Covers:
- Unit tests for HMAC signing and retry constants (no DB required)
- Integration tests for fire_event() using a real DB (inserts into async_operations)
- Integration tests for _handle_webhook_delivery() on the memory engine
- HTTP API integration tests for CRUD and delivery listing endpoints
"""
import json
import uuid
from datetime import datetime, timezone
from unittest.mock import AsyncMock, MagicMock, patch
import httpx
import pytest
import pytest_asyncio
from hindsight_api.api import create_app
from hindsight_api.engine.memory_engine import MemoryEngine
from hindsight_api.webhooks.manager import MAX_ATTEMPTS, RETRY_DELAYS, WebhookManager
from hindsight_api.webhooks.models import (
ConsolidationEventData,
RetainEventData,
WebhookConfig,
WebhookEvent,
WebhookEventType,
)
from hindsight_api.worker.exceptions import RetryTaskAt
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def _make_event(bank_id: str = "bank-1") -> WebhookEvent:
return WebhookEvent(
event=WebhookEventType.CONSOLIDATION_COMPLETED,
bank_id=bank_id,
operation_id=uuid.uuid4().hex,
status="completed",
timestamp=datetime.now(timezone.utc),
data=ConsolidationEventData(observations_created=1),
)
def _make_delivery_task(
bank_id: str = "bank-1",
url: str = "https://example.com/hook",
retry_count: int = 0,
webhook_id: str | None = None,
) -> dict:
return {
"type": "webhook_delivery",
"bank_id": bank_id,
"url": url,
"secret": None,
"event_type": "consolidation.completed",
"payload": '{"event":"consolidation.completed"}',
"webhook_id": webhook_id,
"_retry_count": retry_count,
}
# ---------------------------------------------------------------------------
# Unit tests (no DB)
# ---------------------------------------------------------------------------
class TestHmacSigning:
"""Unit tests for WebhookManager._sign_payload()."""
def _make_manager(self) -> WebhookManager:
"""Create a WebhookManager with a dummy pool (not used for signing)."""
pool = MagicMock()
return WebhookManager(pool=pool, global_webhooks=[])
def test_hmac_signing_format(self):
"""_sign_payload should return a string starting with 'sha256='."""
manager = self._make_manager()
sig = manager._sign_payload("my-secret", b"hello world")
assert sig.startswith("sha256="), f"Expected 'sha256=' prefix, got: {sig!r}"
hex_part = sig[len("sha256="):]
# SHA-256 hex digest is always 64 characters
assert len(hex_part) == 64
# Hex characters only
assert all(c in "0123456789abcdef" for c in hex_part)
def test_hmac_signing_is_deterministic(self):
"""Same secret + payload always produces the same signature."""
manager = self._make_manager()
payload = b'{"event":"consolidation.completed"}'
sig1 = manager._sign_payload("secret-key", payload)
sig2 = manager._sign_payload("secret-key", payload)
assert sig1 == sig2
def test_hmac_signing_differs_with_different_secret(self):
"""Different secrets must produce different signatures."""
manager = self._make_manager()
payload = b"payload"
sig1 = manager._sign_payload("secret-a", payload)
sig2 = manager._sign_payload("secret-b", payload)
assert sig1 != sig2
def test_hmac_signing_differs_with_different_payload(self):
"""Different payloads must produce different signatures."""
manager = self._make_manager()
sig1 = manager._sign_payload("secret", b"payload-one")
sig2 = manager._sign_payload("secret", b"payload-two")
assert sig1 != sig2
class TestRetryConstants:
"""Unit tests to verify retry schedule constants."""
def test_retry_delays_values(self):
"""RETRY_DELAYS must match the documented schedule."""
assert RETRY_DELAYS == [5, 300, 1800, 7200, 18000]
def test_max_attempts(self):
"""MAX_ATTEMPTS should be len(RETRY_DELAYS) + 1."""
assert MAX_ATTEMPTS == 6
assert MAX_ATTEMPTS == len(RETRY_DELAYS) + 1
# ---------------------------------------------------------------------------
# DB integration tests
# ---------------------------------------------------------------------------
@pytest_asyncio.fixture
async def webhook_manager(memory: MemoryEngine) -> WebhookManager:
"""Return a WebhookManager backed by the test pool with no global webhooks."""
return WebhookManager(pool=memory._pool, global_webhooks=[])
async def _ensure_bank(pool, bank_id: str) -> None:
"""Upsert a minimal bank row so FK constraints on async_operations/webhooks pass."""
await pool.execute(
"INSERT INTO banks (bank_id, name) VALUES ($1, $2) ON CONFLICT DO NOTHING",
bank_id,
bank_id,
)
class TestFireEvent:
"""Integration tests for WebhookManager.fire_event()."""
@pytest.mark.asyncio
async def test_fire_event_creates_delivery(
self, memory: MemoryEngine, webhook_manager: WebhookManager
):
"""fire_event() inserts a pending webhook_delivery task in async_operations."""
bank_id = f"wh-test-{uuid.uuid4().hex[:8]}"
webhook_id = uuid.uuid4()
async with memory._pool.acquire() as conn:
await _ensure_bank(memory._pool, bank_id)
await conn.execute(
"""
INSERT INTO webhooks (id, bank_id, url, secret, event_types, enabled, created_at, updated_at)
VALUES ($1, $2, $3, NULL, $4, true, NOW(), NOW())
""",
webhook_id,
bank_id,
"https://example.com/hook",
["consolidation.completed"],
)
try:
event = _make_event(bank_id)
await webhook_manager.fire_event(event)
async with memory._pool.acquire() as conn:
rows = await conn.fetch(
"""
SELECT status, task_payload
FROM async_operations
WHERE operation_type = 'webhook_delivery'
AND bank_id = $1
AND task_payload->>'webhook_id' = $2
""",
bank_id,
str(webhook_id),
)
assert len(rows) == 1
assert rows[0]["status"] == "pending"
payload = rows[0]["task_payload"]
if isinstance(payload, str):
payload = json.loads(payload)
assert payload["event_type"] == "consolidation.completed"
finally:
async with memory._pool.acquire() as conn:
await conn.execute(
"DELETE FROM async_operations WHERE operation_type = 'webhook_delivery' AND bank_id = $1",
bank_id,
)
await conn.execute("DELETE FROM webhooks WHERE id = $1", webhook_id)
@pytest.mark.asyncio
async def test_fire_event_global_webhook(
self, memory: MemoryEngine
):
"""fire_event() also queues delivery tasks for global webhooks (not stored in DB)."""
bank_id = f"wh-global-{uuid.uuid4().hex[:8]}"
await _ensure_bank(memory._pool, bank_id)
global_webhook = WebhookConfig(
id="", # No DB row
bank_id=None,
url="https://global.example.com/hook",
secret=None,
event_types=["consolidation.completed"],
enabled=True,
)
manager = WebhookManager(pool=memory._pool, global_webhooks=[global_webhook])
event = _make_event(bank_id)
await manager.fire_event(event)
async with memory._pool.acquire() as conn:
rows = await conn.fetch(
"""
SELECT status, task_payload
FROM async_operations
WHERE operation_type = 'webhook_delivery'
AND bank_id = $1
AND task_payload->>'url' = 'https://global.example.com/hook'
ORDER BY created_at DESC
LIMIT 1
"""
,
bank_id,
)
assert len(rows) == 1
assert rows[0]["status"] == "pending"
payload = rows[0]["task_payload"]
if isinstance(payload, str):
payload = json.loads(payload)
assert payload["webhook_id"] is None # global webhook has no DB row
# Cleanup
async with memory._pool.acquire() as conn:
await conn.execute(
"DELETE FROM async_operations WHERE operation_type = 'webhook_delivery' AND bank_id = $1",
bank_id,
)
@pytest.mark.asyncio
async def test_fire_event_no_match_if_event_type_mismatch(
self, memory: MemoryEngine, webhook_manager: WebhookManager
):
"""Webhooks registered for a different event type receive no delivery task."""
bank_id = f"wh-mismatch-{uuid.uuid4().hex[:8]}"
webhook_id = uuid.uuid4()
await _ensure_bank(memory._pool, bank_id)
async with memory._pool.acquire() as conn:
await conn.execute(
"""
INSERT INTO webhooks (id, bank_id, url, secret, event_types, enabled, created_at, updated_at)
VALUES ($1, $2, $3, NULL, $4, true, NOW(), NOW())
""",
webhook_id,
bank_id,
"https://example.com/other-hook",
["other.event"],
)
try:
event = _make_event(bank_id)
await webhook_manager.fire_event(event)
async with memory._pool.acquire() as conn:
count = await conn.fetchval(
"""
SELECT COUNT(*) FROM async_operations
WHERE operation_type = 'webhook_delivery' AND bank_id = $1
""",
bank_id,
)
assert count == 0
finally:
async with memory._pool.acquire() as conn:
await conn.execute("DELETE FROM webhooks WHERE id = $1", webhook_id)
class TestHandleWebhookDelivery:
"""Integration tests for MemoryEngine._handle_webhook_delivery()."""
@pytest.mark.asyncio
async def test_deliver_success(self, memory: MemoryEngine):
"""A successful HTTP POST completes without raising."""
task_dict = _make_delivery_task(retry_count=0)
mock_response = MagicMock()
mock_response.raise_for_status = MagicMock()
with patch.object(memory._http_client, "post", new=AsyncMock(return_value=mock_response)):
# Should not raise
await memory._handle_webhook_delivery(task_dict)
@pytest.mark.asyncio
async def test_deliver_failure_raises_retry_task_at(self, memory: MemoryEngine):
"""A failed HTTP POST raises RetryTaskAt when retries remain."""
task_dict = _make_delivery_task(retry_count=0)
with patch.object(
memory._http_client, "post", new=AsyncMock(side_effect=Exception("connection refused"))
):
with pytest.raises(RetryTaskAt):
await memory._handle_webhook_delivery(task_dict)
@pytest.mark.asyncio
async def test_deliver_exhausted_retries_raises(self, memory: MemoryEngine):
"""When retry_count reaches MAX_ATTEMPTS-1, a failure raises the original exception."""
task_dict = _make_delivery_task(retry_count=MAX_ATTEMPTS - 1)
with patch.object(
memory._http_client, "post", new=AsyncMock(side_effect=Exception("server error"))
):
with pytest.raises(Exception, match="server error"):
await memory._handle_webhook_delivery(task_dict)
@pytest.mark.asyncio
async def test_deliver_retry_at_uses_delay_schedule(self, memory: MemoryEngine):
"""RetryTaskAt.retry_at is approximately now + RETRY_DELAYS[retry_count]."""
from datetime import timedelta
task_dict = _make_delivery_task(retry_count=1)
with patch.object(
memory._http_client, "post", new=AsyncMock(side_effect=Exception("fail"))
):
before = datetime.now(timezone.utc)
with pytest.raises(RetryTaskAt) as exc_info:
await memory._handle_webhook_delivery(task_dict)
after = datetime.now(timezone.utc)
retry_at = exc_info.value.retry_at
expected_delay = RETRY_DELAYS[1] # retry_count=1
assert retry_at >= before + timedelta(seconds=expected_delay - 2)
assert retry_at <= after + timedelta(seconds=expected_delay + 2)
@pytest.mark.asyncio
async def test_execute_task_marks_operation_completed(self, memory: MemoryEngine):
"""After a successful delivery, execute_task marks the async_operations row as completed."""
operation_id = str(uuid.uuid4())
bank_id = f"wh-exec-{uuid.uuid4().hex[:8]}"
await _ensure_bank(memory._pool, bank_id)
# Insert a real async_operations row so _mark_operation_completed has something to update
async with memory._pool.acquire() as conn:
await conn.execute(
"""
INSERT INTO async_operations
(operation_id, bank_id, operation_type, status, task_payload, result_metadata, created_at, updated_at)
VALUES ($1, $2, 'webhook_delivery', 'processing', '{}'::jsonb, '{}'::jsonb, NOW(), NOW())
""",
uuid.UUID(operation_id),
bank_id,
)
task_dict = {
**_make_delivery_task(bank_id=bank_id, retry_count=0),
"operation_id": operation_id,
}
mock_response = MagicMock()
mock_response.raise_for_status = MagicMock()
with patch.object(memory._http_client, "post", new=AsyncMock(return_value=mock_response)):
await memory.execute_task(task_dict)
async with memory._pool.acquire() as conn:
row = await conn.fetchrow(
"SELECT status FROM async_operations WHERE operation_id = $1",
uuid.UUID(operation_id),
)
assert row is not None
assert row["status"] == "completed", f"Expected 'completed', got '{row['status']}'"
# Cleanup
async with memory._pool.acquire() as conn:
await conn.execute(
"DELETE FROM async_operations WHERE operation_id = $1",
uuid.UUID(operation_id),
)
# ---------------------------------------------------------------------------
# HTTP API integration tests
# ---------------------------------------------------------------------------
@pytest_asyncio.fixture
async def api_client(memory: MemoryEngine):
"""Async HTTP test client wired to 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
class TestWebhookHttpApi:
"""HTTP API integration tests for webhook CRUD endpoints."""
@pytest.mark.asyncio
async def test_http_create_webhook(self, api_client: httpx.AsyncClient):
"""POST /webhooks returns 201 and an id."""
bank_id = f"http-wh-{uuid.uuid4().hex[:8]}"
response = await api_client.post(
f"/v1/default/banks/{bank_id}/webhooks",
json={
"url": "https://example.com/create",
"event_types": ["consolidation.completed"],
},
)
assert response.status_code == 201, response.text
data = response.json()
assert "id" in data
assert data["url"] == "https://example.com/create"
assert data["bank_id"] == bank_id
assert data["secret"] is None # secrets are never echoed back
# Cleanup
await api_client.delete(
f"/v1/default/banks/{bank_id}/webhooks/{data['id']}"
)
@pytest.mark.asyncio
async def test_http_list_webhooks(self, api_client: httpx.AsyncClient):
"""GET /webhooks returns the webhooks registered for a bank."""
bank_id = f"http-wh-{uuid.uuid4().hex[:8]}"
create_resp = await api_client.post(
f"/v1/default/banks/{bank_id}/webhooks",
json={"url": "https://example.com/list", "event_types": ["consolidation.completed"]},
)
assert create_resp.status_code == 201
webhook_id = create_resp.json()["id"]
list_resp = await api_client.get(f"/v1/default/banks/{bank_id}/webhooks")
assert list_resp.status_code == 200
items = list_resp.json()["items"]
assert any(item["id"] == webhook_id for item in items)
# Cleanup
await api_client.delete(f"/v1/default/banks/{bank_id}/webhooks/{webhook_id}")
@pytest.mark.asyncio
async def test_http_delete_webhook(self, api_client: httpx.AsyncClient):
"""DELETE /webhooks/{id} removes the webhook; subsequent list returns empty for that bank."""
bank_id = f"http-wh-{uuid.uuid4().hex[:8]}"
create_resp = await api_client.post(
f"/v1/default/banks/{bank_id}/webhooks",
json={"url": "https://example.com/delete", "event_types": ["consolidation.completed"]},
)
assert create_resp.status_code == 201
webhook_id = create_resp.json()["id"]
delete_resp = await api_client.delete(
f"/v1/default/banks/{bank_id}/webhooks/{webhook_id}"
)
assert delete_resp.status_code == 200
assert delete_resp.json()["success"] is True
list_resp = await api_client.get(f"/v1/default/banks/{bank_id}/webhooks")
assert list_resp.status_code == 200
ids = [item["id"] for item in list_resp.json()["items"]]
assert webhook_id not in ids
@pytest.mark.asyncio
async def test_http_delete_webhook_not_found(self, api_client: httpx.AsyncClient):
"""DELETE with a non-existent webhook id returns 404."""
bank_id = f"http-wh-{uuid.uuid4().hex[:8]}"
missing_id = str(uuid.uuid4())
response = await api_client.delete(
f"/v1/default/banks/{bank_id}/webhooks/{missing_id}"
)
assert response.status_code == 404
@pytest.mark.asyncio
async def test_http_list_deliveries(
self, memory: MemoryEngine, api_client: httpx.AsyncClient
):
"""GET /webhooks/{id}/deliveries returns delivery records for a webhook."""
bank_id = f"http-wh-{uuid.uuid4().hex[:8]}"
# Create webhook via HTTP API
create_resp = await api_client.post(
f"/v1/default/banks/{bank_id}/webhooks",
json={
"url": "https://example.com/deliveries",
"event_types": ["consolidation.completed"],
},
)
assert create_resp.status_code == 201
webhook_id = create_resp.json()["id"]
# Insert a delivery row directly into async_operations
delivery_id = uuid.uuid4()
now = datetime.now(timezone.utc)
task_payload = json.dumps(
{
"type": "webhook_delivery",
"bank_id": bank_id,
"url": "https://example.com/deliveries",
"secret": None,
"event_type": "consolidation.completed",
"payload": '{"event":"consolidation.completed"}',
"webhook_id": webhook_id,
}
)
async with memory._pool.acquire() as conn:
await conn.execute(
"""
INSERT INTO async_operations
(operation_id, bank_id, operation_type, status, retry_count, task_payload, result_metadata, created_at, updated_at)
VALUES ($1, $2, 'webhook_delivery', 'completed', 0, $3::jsonb, '{}'::jsonb, $4, $4)
""",
delivery_id,
bank_id,
task_payload,
now,
)
try:
deliveries_resp = await api_client.get(
f"/v1/default/banks/{bank_id}/webhooks/{webhook_id}/deliveries"
)
assert deliveries_resp.status_code == 200
items = deliveries_resp.json()["items"]
ids = [item["id"] for item in items]
assert str(delivery_id) in ids
# Verify shape of a delivery item
delivery = next(item for item in items if item["id"] == str(delivery_id))
assert delivery["status"] == "completed"
assert delivery["event_type"] == "consolidation.completed"
assert delivery["attempts"] == 1
finally:
async with memory._pool.acquire() as conn:
await conn.execute(
"DELETE FROM async_operations WHERE operation_id = $1", delivery_id
)
await api_client.delete(
f"/v1/default/banks/{bank_id}/webhooks/{webhook_id}"
)
@pytest.mark.asyncio
async def test_http_list_deliveries_webhook_not_found(self, api_client: httpx.AsyncClient):
"""GET /webhooks/{id}/deliveries for a non-existent webhook returns 404."""
bank_id = f"http-wh-{uuid.uuid4().hex[:8]}"
missing_id = str(uuid.uuid4())
response = await api_client.get(
f"/v1/default/banks/{bank_id}/webhooks/{missing_id}/deliveries"
)
assert response.status_code == 404
@pytest.mark.asyncio
async def test_http_update_webhook_url(self, api_client: httpx.AsyncClient):
"""PATCH /webhooks/{id} updates only the provided fields."""
bank_id = f"http-wh-{uuid.uuid4().hex[:8]}"
create_resp = await api_client.post(
f"/v1/default/banks/{bank_id}/webhooks",
json={"url": "https://example.com/original", "event_types": ["consolidation.completed"]},
)
assert create_resp.status_code == 201
webhook_id = create_resp.json()["id"]
patch_resp = await api_client.patch(
f"/v1/default/banks/{bank_id}/webhooks/{webhook_id}",
json={"url": "https://example.com/updated"},
)
assert patch_resp.status_code == 200
data = patch_resp.json()
assert data["url"] == "https://example.com/updated"
# event_types should be unchanged
assert "consolidation.completed" in data["event_types"]
# Cleanup
await api_client.delete(f"/v1/default/banks/{bank_id}/webhooks/{webhook_id}")
@pytest.mark.asyncio
async def test_http_update_webhook_event_types(self, api_client: httpx.AsyncClient):
"""PATCH /webhooks/{id} can update event_types."""
bank_id = f"http-wh-{uuid.uuid4().hex[:8]}"
create_resp = await api_client.post(
f"/v1/default/banks/{bank_id}/webhooks",
json={"url": "https://example.com/hook", "event_types": ["consolidation.completed"]},
)
assert create_resp.status_code == 201
webhook_id = create_resp.json()["id"]
patch_resp = await api_client.patch(
f"/v1/default/banks/{bank_id}/webhooks/{webhook_id}",
json={"event_types": ["retain.completed"]},
)
assert patch_resp.status_code == 200
data = patch_resp.json()
assert data["event_types"] == ["retain.completed"]
# Cleanup
await api_client.delete(f"/v1/default/banks/{bank_id}/webhooks/{webhook_id}")
@pytest.mark.asyncio
async def test_http_update_webhook_enabled(self, api_client: httpx.AsyncClient):
"""PATCH /webhooks/{id} can toggle enabled."""
bank_id = f"http-wh-{uuid.uuid4().hex[:8]}"
create_resp = await api_client.post(
f"/v1/default/banks/{bank_id}/webhooks",
json={"url": "https://example.com/hook", "event_types": ["consolidation.completed"]},
)
assert create_resp.status_code == 201
webhook_id = create_resp.json()["id"]
assert create_resp.json()["enabled"] is True
patch_resp = await api_client.patch(
f"/v1/default/banks/{bank_id}/webhooks/{webhook_id}",
json={"enabled": False},
)
assert patch_resp.status_code == 200
assert patch_resp.json()["enabled"] is False
# Cleanup
await api_client.delete(f"/v1/default/banks/{bank_id}/webhooks/{webhook_id}")
@pytest.mark.asyncio
async def test_http_update_webhook_http_config(self, api_client: httpx.AsyncClient):
"""PATCH /webhooks/{id} can update http_config."""
bank_id = f"http-wh-{uuid.uuid4().hex[:8]}"
create_resp = await api_client.post(
f"/v1/default/banks/{bank_id}/webhooks",
json={"url": "https://example.com/hook", "event_types": ["consolidation.completed"]},
)
assert create_resp.status_code == 201
webhook_id = create_resp.json()["id"]
patch_resp = await api_client.patch(
f"/v1/default/banks/{bank_id}/webhooks/{webhook_id}",
json={
"http_config": {
"method": "POST",
"timeout_seconds": 10,
"headers": {"X-Custom": "value"},
"params": {},
}
},
)
assert patch_resp.status_code == 200
data = patch_resp.json()
assert data["http_config"]["timeout_seconds"] == 10
assert data["http_config"]["headers"] == {"X-Custom": "value"}
# Cleanup
await api_client.delete(f"/v1/default/banks/{bank_id}/webhooks/{webhook_id}")
@pytest.mark.asyncio
async def test_http_update_webhook_not_found(self, api_client: httpx.AsyncClient):
"""PATCH /webhooks/{id} returns 404 for a non-existent webhook."""
bank_id = f"http-wh-{uuid.uuid4().hex[:8]}"
missing_id = str(uuid.uuid4())
response = await api_client.patch(
f"/v1/default/banks/{bank_id}/webhooks/{missing_id}",
json={"url": "https://example.com/new"},
)
assert response.status_code == 404
@pytest.mark.asyncio
async def test_http_update_webhook_no_fields(self, api_client: httpx.AsyncClient):
"""PATCH /webhooks/{id} with empty body returns 422."""
bank_id = f"http-wh-{uuid.uuid4().hex[:8]}"
create_resp = await api_client.post(
f"/v1/default/banks/{bank_id}/webhooks",
json={"url": "https://example.com/hook", "event_types": ["consolidation.completed"]},
)
assert create_resp.status_code == 201
webhook_id = create_resp.json()["id"]
patch_resp = await api_client.patch(
f"/v1/default/banks/{bank_id}/webhooks/{webhook_id}",
json={},
)
assert patch_resp.status_code == 422
# Cleanup
await api_client.delete(f"/v1/default/banks/{bank_id}/webhooks/{webhook_id}")
# ---------------------------------------------------------------------------
# retain.completed webhook tests
# ---------------------------------------------------------------------------
class TestRetainCompletedWebhook:
"""Tests for the retain.completed webhook event."""
def test_retain_event_data_model(self):
"""RetainEventData can be constructed with optional fields."""
data = RetainEventData(document_id="doc-123", tags=["tag1", "tag2"])
assert data.document_id == "doc-123"
assert data.tags == ["tag1", "tag2"]
empty = RetainEventData()
assert empty.document_id is None
assert empty.tags is None
def test_retain_event_type_value(self):
"""WebhookEventType.RETAIN_COMPLETED has the correct string value."""
assert WebhookEventType.RETAIN_COMPLETED == "retain.completed"
@pytest.mark.asyncio
async def test_fire_retain_webhook_queues_per_document(
self, memory: MemoryEngine, webhook_manager: WebhookManager
):
"""_fire_retain_webhook queues one delivery task per content item."""
bank_id = f"wh-retain-{uuid.uuid4().hex[:8]}"
webhook_id = uuid.uuid4()
await _ensure_bank(memory._pool, bank_id)
async with memory._pool.acquire() as conn:
await conn.execute(
"""
INSERT INTO webhooks (id, bank_id, url, secret, event_types, enabled, created_at, updated_at)
VALUES ($1, $2, $3, NULL, $4, true, NOW(), NOW())
""",
webhook_id,
bank_id,
"https://example.com/retain-hook",
["retain.completed"],
)
try:
contents = [
{"content": "Alice works at Google", "document_id": "doc-1"},
{"content": "Bob loves Python", "document_id": "doc-2"},
]
# Temporarily replace webhook manager on memory engine
original_manager = memory._webhook_manager
memory._webhook_manager = webhook_manager
try:
callback = memory._build_retain_outbox_callback(
bank_id=bank_id,
contents=contents,
operation_id="test-op-123",
)
assert callback is not None
async with memory._pool.acquire() as conn:
await callback(conn)
finally:
memory._webhook_manager = original_manager
async with memory._pool.acquire() as conn:
rows = await conn.fetch(
"""
SELECT task_payload
FROM async_operations
WHERE operation_type = 'webhook_delivery'
AND bank_id = $1
AND task_payload->>'event_type' = 'retain.completed'
ORDER BY created_at
""",
bank_id,
)
assert len(rows) == 2
payloads = []
for row in rows:
p = row["task_payload"]
if isinstance(p, str):
p = json.loads(p)
payloads.append(p)
doc_ids_in_payloads = [json.loads(p["payload"]).get("data", {}).get("document_id") for p in payloads]
assert "doc-1" in doc_ids_in_payloads
assert "doc-2" in doc_ids_in_payloads
finally:
async with memory._pool.acquire() as conn:
await conn.execute(
"DELETE FROM async_operations WHERE operation_type = 'webhook_delivery' AND bank_id = $1",
bank_id,
)
await conn.execute("DELETE FROM webhooks WHERE id = $1", webhook_id)
@@ -46,4 +46,4 @@ __all__ = [
"RemoteTEICrossEncoder",
"LLMConfig",
]
__version__ = "0.4.19"
__version__ = "0.4.11"
@@ -14,8 +14,7 @@ from typing import Any
import asyncpg
import typer
from ..config import DEFAULT_DATABASE_SCHEMA, HindsightConfig
from ..extensions import TenantExtension, load_extension
from ..config import HindsightConfig
from ..pg0 import parse_pg0_url, resolve_database_url
@@ -215,81 +214,20 @@ def restore(
typer.echo("Restore complete")
async def _run_migration(
db_url: str,
schema: str | None = None,
base_schema: str = DEFAULT_DATABASE_SCHEMA,
embedding_dimension: int | None = None,
) -> list[str]:
"""Resolve database URL and run migrations for one schema or all discovered schemas."""
from ..migrations import (
ensure_embedding_dimension,
ensure_text_search_extension,
ensure_vector_extension,
run_migrations,
)
async def _run_migration(db_url: str, schema: str = "public") -> None:
"""Resolve database URL and run migrations."""
from ..migrations import run_migrations
is_pg0, instance_name, _ = parse_pg0_url(db_url)
if is_pg0:
typer.echo(f"Starting embedded PostgreSQL (instance: {instance_name})...")
resolved_url = await resolve_database_url(db_url)
config = HindsightConfig.from_env()
if schema:
schemas = [schema]
else:
tenant_extension = load_extension("TENANT", TenantExtension)
schemas = [base_schema or DEFAULT_DATABASE_SCHEMA]
if tenant_extension:
tenants = await tenant_extension.list_tenants()
schemas.extend(tenant.schema for tenant in tenants if tenant.schema)
# Preserve order while removing duplicates.
schemas = list(dict.fromkeys(schemas))
for schema in schemas:
run_migrations(resolved_url, schema=schema)
if embedding_dimension is not None:
for schema in schemas:
ensure_embedding_dimension(
resolved_url,
embedding_dimension,
schema=schema,
vector_extension=config.vector_extension,
)
for schema in schemas:
ensure_vector_extension(
resolved_url,
vector_extension=config.vector_extension,
schema=schema,
)
for schema in schemas:
ensure_text_search_extension(
resolved_url,
text_search_extension=config.text_search_extension,
schema=schema,
)
return schemas
run_migrations(resolved_url, schema=schema)
@app.command(name="run-db-migration")
def run_db_migration(
schema: str | None = typer.Option(
None,
"--schema",
"-s",
help="Database schema to run migrations on. If omitted, migrate the base schema and all discovered tenant schemas.",
),
embedding_dimension: int | None = typer.Option(
None,
"--embedding-dimension",
help="Expected embedding dimension to enforce after migrations. Omit to skip dimension sync.",
),
schema: str = typer.Option("public", "--schema", "-s", help="Database schema to run migrations on"),
):
"""Run database migrations to the latest version."""
config = HindsightConfig.from_env()
@@ -299,21 +237,11 @@ def run_db_migration(
typer.echo("Set HINDSIGHT_API_DATABASE_URL environment variable.", err=True)
raise typer.Exit(1)
if schema:
typer.echo(f"Running database migrations for schema: {schema}...")
else:
typer.echo("Running database migrations for base schema and all discovered tenant schemas...")
typer.echo(f"Running database migrations (schema: {schema})...")
schemas = asyncio.run(
_run_migration(
config.database_url,
schema=schema,
base_schema=config.database_schema,
embedding_dimension=embedding_dimension,
)
)
asyncio.run(_run_migration(config.database_url, schema))
typer.echo(f"Database migrations completed successfully for {len(schemas)} schema(s)")
typer.echo("Database migrations completed successfully")
async def _decommission_worker(db_url: str, worker_id: str, schema: str = "public") -> int:
@@ -34,7 +34,7 @@ def upgrade() -> None:
# Create file_storage table (minimal: just key + data)
op.execute(
f"""
CREATE TABLE IF NOT EXISTS {schema}file_storage (
CREATE TABLE {schema}file_storage (
storage_key TEXT PRIMARY KEY,
data BYTEA NOT NULL
)

Some files were not shown because too many files have changed in this diff Show More