Compare commits

..
Author SHA1 Message Date
Nicolò Boschi 3d66988c3b add api cli 2025-11-25 19:28:12 +01:00
Nicolò Boschi 67dc160ba3 pg0 works 2025-11-25 17:05:18 +01:00
Nicolò Boschi 8291386387 hindsight all 2025-11-25 16:57:54 +01:00
Nicolò Boschi 3511062c51 refactor ts client 2025-11-25 12:17:13 +01:00
Nicolò Boschi a134d27137 fix tests and docs 2025-11-25 11:31:21 +01:00
Nicolò Boschi 31a53c0870 rename to hindsight 2025-11-25 10:11:04 +01:00
539 changed files with 4419322 additions and 78939 deletions
-32
View File
@@ -1,32 +0,0 @@
# Node modules (platform-specific native bindings)
**/node_modules
**/.next
# Python
**/__pycache__
**/*.pyc
**/.venv
**/dist
**/*.egg-info
# Git
.git
.gitignore
# IDE
.idea
.vscode
*.swp
# OS
.DS_Store
Thumbs.db
# Build artifacts
**/target
**/*.log
# Test/Dev
**/coverage
**/.pytest_cache
**/.mypy_cache
+42 -27
View File
@@ -1,32 +1,47 @@
# Hindsight Environment Variables
# Copy this file to .env and fill in your values
# =============================================================================
# HINDSIGHT ENVIRONMENT CONFIGURATION
# =============================================================================
# Copy this file to .env and update with your values
# Both services (API and Control Plane) read from this single file
# LLM Configuration (Required)
HINDSIGHT_API_LLM_PROVIDER=openai
HINDSIGHT_API_LLM_API_KEY=your-api-key-here
HINDSIGHT_API_LLM_MODEL=o3-mini
HINDSIGHT_API_LLM_BASE_URL=https://api.openai.com/v1
# =============================================================================
# API SERVICE (HINDSIGHT_API_*)
# =============================================================================
# API Configuration (Optional)
HINDSIGHT_API_HOST=0.0.0.0
HINDSIGHT_API_PORT=8888
HINDSIGHT_API_LOG_LEVEL=info
# Database
# Use "pg0" to start an embedded PostgreSQL instance via pg0
# Or provide a full connection URL for external PostgreSQL
#HINDSIGHT_API_DATABASE_URL=postgresql://hindsight:hindsight_dev@localhost:5432/hindsight
HINDSIGHT_API_DATABASE_URL=pg0
# Database (Optional - uses embedded pg0 by default)
# HINDSIGHT_API_DATABASE_URL=postgresql://user:pass@host:5432/db
# pg0 data directory (only used when HINDSIGHT_API_DATABASE_URL=pg0)
# HINDSIGHT_API_PG0_DATA_DIR=/path/to/pg_data
# Embeddings Configuration (Optional - uses local by default)
# Provider: "local" (default) or "tei" (HuggingFace Text Embeddings Inference)
# HINDSIGHT_API_EMBEDDINGS_PROVIDER=local
# For local provider:
# HINDSIGHT_API_EMBEDDINGS_LOCAL_MODEL=BAAI/bge-small-en-v1.5
# For TEI provider:
# HINDSIGHT_API_EMBEDDINGS_TEI_URL=http://localhost:8080
# LLM Provider: "openai", "groq", or "ollama"
HINDSIGHT_API_LLM_PROVIDER=groq
# Reranker Configuration (Optional - uses local by default)
# Provider: "local" (default) or "tei" (HuggingFace Text Embeddings Inference)
# HINDSIGHT_API_RERANKER_PROVIDER=local
# For local provider:
# HINDSIGHT_API_RERANKER_LOCAL_MODEL=cross-encoder/ms-marco-MiniLM-L-6-v2
# For TEI provider:
# HINDSIGHT_API_RERANKER_TEI_URL=http://localhost:8081
# LLM Model (provider-specific)
HINDSIGHT_API_LLM_MODEL=openai/gpt-oss-20b
# API Key (not needed for ollama)
HINDSIGHT_API_LLM_API_KEY=your_api_key_here
# Optional: Custom base URL (for ollama or custom endpoints)
# HINDSIGHT_API_LLM_BASE_URL=http://localhost:11434/v1
# API Server Configuration (optional)
# HINDSIGHT_API_HOST=0.0.0.0
# HINDSIGHT_API_PORT=8888
HINDSIGHT_API_MCP_ENABLED=true
# =============================================================================
# CONTROL PLANE SERVICE (HINDSIGHT_CP_*)
# =============================================================================
# Dataplane API URL (where the control plane connects to)
HINDSIGHT_CP_DATAPLANE_API_URL=http://localhost:8888
# Control Plane Server Configuration (optional)
# HINDSIGHT_CP_PORT=3000
# HINDSIGHT_CP_HOSTNAME=0.0.0.0
+7 -4
View File
@@ -2,7 +2,7 @@ name: Deploy Docs to GitHub Pages
on:
push:
branches: [main]
branches: [main, renaming-pre-launch]
paths:
- 'hindsight-docs/**'
- '.github/workflows/deploy-docs.yml'
@@ -20,15 +20,18 @@ concurrency:
jobs:
build:
runs-on: ubuntu-latest
defaults:
run:
working-directory: hindsight-docs
steps:
- uses: actions/checkout@v4
- uses: actions/setup-node@v4
with:
node-version: 20
cache: npm
cache-dependency-path: package-lock.json
- run: npm ci --workspace=hindsight-docs
- run: npm run build --workspace=hindsight-docs
cache-dependency-path: hindsight-docs/package-lock.json
- run: npm ci
- run: npm run build
- uses: actions/upload-pages-artifact@v3
with:
path: hindsight-docs/build
+151 -149
View File
@@ -1,4 +1,4 @@
name: Release
name: Build Release Artifacts
on:
push:
@@ -6,11 +6,8 @@ on:
- 'v*'
jobs:
release-python-packages:
build-python-package:
runs-on: ubuntu-latest
environment: pypi
permissions:
id-token: write
steps:
- uses: actions/checkout@v4
@@ -25,99 +22,18 @@ jobs:
with:
python-version-file: ".python-version"
# Build all packages
- name: Build hindsight-client
working-directory: ./hindsight-clients/python
run: uv build --out-dir dist
- name: Build hindsight-api
working-directory: ./hindsight-api
run: uv build --out-dir dist
- name: Build hindsight-all
- name: Build hindsight package
working-directory: ./hindsight
run: uv build --out-dir dist
- name: Build hindsight-litellm
working-directory: ./hindsight-integrations/litellm
run: uv build --out-dir dist
# 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 to PyPI
uses: pypa/gh-action-pypi-publish@release/v1
with:
packages-dir: ./hindsight-api/dist
skip-existing: true
- name: Publish hindsight-all to PyPI
uses: pypa/gh-action-pypi-publish@release/v1
with:
packages-dir: ./hindsight/dist
skip-existing: true
- name: Publish hindsight-litellm to PyPI
uses: pypa/gh-action-pypi-publish@release/v1
with:
packages-dir: ./hindsight-integrations/litellm/dist
skip-existing: true
# Upload artifacts for GitHub release
- name: Upload artifacts
uses: actions/upload-artifact@v4
with:
name: python-packages
path: |
hindsight-clients/python/dist/*
hindsight-api/dist/*
hindsight/dist/*
hindsight-integrations/litellm/dist/*
retention-days: 1
release-typescript-client:
runs-on: ubuntu-latest
environment: npm
steps:
- uses: actions/checkout@v4
- name: Set up Node.js
uses: actions/setup-node@v4
with:
node-version: '20'
registry-url: 'https://registry.npmjs.org'
cache: 'npm'
cache-dependency-path: package-lock.json
- name: Install dependencies
run: npm ci --workspace=hindsight-clients/typescript
- name: Build
run: npm run build --workspace=hindsight-clients/typescript
- name: Publish to npm
working-directory: ./hindsight-clients/typescript
run: npm publish --access public
env:
NODE_AUTH_TOKEN: ${{ secrets.NPM_TOKEN }}
- name: Pack for GitHub release
working-directory: ./hindsight-clients/typescript
run: npm pack
run: uv build
- name: Upload artifacts
uses: actions/upload-artifact@v4
with:
name: typescript-client
path: hindsight-clients/typescript/*.tgz
retention-days: 1
name: python-hindsight-dist
path: hindsight/dist/*
retention-days: 30
release-rust-cli:
build-rust-cli:
runs-on: ${{ matrix.os }}
strategy:
matrix:
@@ -143,6 +59,24 @@ jobs:
with:
targets: ${{ matrix.target }}
- name: Cache cargo registry
uses: actions/cache@v4
with:
path: ~/.cargo/registry
key: ${{ runner.os }}-cargo-registry-${{ hashFiles('**/Cargo.lock') }}
- name: Cache cargo index
uses: actions/cache@v4
with:
path: ~/.cargo/git
key: ${{ runner.os }}-cargo-index-${{ hashFiles('**/Cargo.lock') }}
- name: Cache cargo build
uses: actions/cache@v4
with:
path: hindsight-cli/target
key: ${{ runner.os }}-cargo-build-target-${{ hashFiles('**/Cargo.lock') }}
- name: Build
working-directory: hindsight-cli
run: cargo build --release --target ${{ matrix.target }}
@@ -158,22 +92,16 @@ jobs:
with:
name: rust-cli-${{ matrix.asset_name }}
path: artifacts/${{ matrix.asset_name }}
retention-days: 1
retention-days: 30
release-docker-images:
build-docker-images:
runs-on: ubuntu-latest
permissions:
contents: read
packages: write
strategy:
matrix:
include:
- target: api-only
image_name: hindsight-api
- target: cp-only
image_name: hindsight-control-plane
- target: standalone
image_name: hindsight
component: [api, control-plane]
steps:
- uses: actions/checkout@v4
@@ -189,9 +117,6 @@ jobs:
docker-images: true
swap-storage: true
- name: Set up QEMU
uses: docker/setup-qemu-action@v3
- name: Set up Docker Buildx
uses: docker/setup-buildx-action@v3
@@ -210,29 +135,39 @@ jobs:
id: meta
uses: docker/metadata-action@v5
with:
images: ghcr.io/${{ github.repository_owner }}/${{ matrix.image_name }}
images: ghcr.io/${{ github.repository_owner }}/hindsight-${{ matrix.component }}
tags: |
type=semver,pattern={{version}},value=${{ steps.get_version.outputs.VERSION }}
type=semver,pattern={{major}}.{{minor}},value=${{ steps.get_version.outputs.VERSION }}
type=semver,pattern={{major}},value=${{ steps.get_version.outputs.VERSION }}
type=raw,value=latest
- name: Build and push
- name: Build and push Docker image (api)
if: matrix.component == 'api'
uses: docker/build-push-action@v6
with:
context: .
file: docker/standalone/Dockerfile
target: ${{ matrix.target }}
file: docker/api.Dockerfile
push: true
platforms: linux/amd64,linux/arm64
tags: ${{ steps.meta.outputs.tags }}
labels: ${{ steps.meta.outputs.labels }}
cache-from: type=gha
cache-to: type=gha,mode=max
release-helm-chart:
- name: Build and push Docker image (control-plane)
if: matrix.component == 'control-plane'
uses: docker/build-push-action@v6
with:
context: .
file: docker/control-plane.Dockerfile
push: true
tags: ${{ steps.meta.outputs.tags }}
labels: ${{ steps.meta.outputs.labels }}
cache-from: type=gha
cache-to: type=gha,mode=max
package-helm-chart:
runs-on: ubuntu-latest
permissions:
contents: read
packages: write
steps:
- uses: actions/checkout@v4
@@ -242,28 +177,24 @@ jobs:
with:
version: 'latest'
- name: Log in to GHCR
run: echo "${{ secrets.GITHUB_TOKEN }}" | helm registry login ghcr.io -u ${{ github.actor }} --password-stdin
- name: Lint Helm chart
run: helm lint helm/hindsight
run: |
helm lint helm/hindsight
- name: Package Helm chart
run: helm package helm/hindsight --destination ./helm-packages
run: |
helm package helm/hindsight --destination ./helm-packages
- name: Push to GHCR OCI
run: helm push helm-packages/*.tgz oci://ghcr.io/${{ github.repository_owner }}/charts
- name: Upload artifacts
- name: Upload Helm chart artifact
uses: actions/upload-artifact@v4
with:
name: helm-chart
path: helm-packages/*.tgz
retention-days: 1
retention-days: 30
create-github-release:
runs-on: ubuntu-latest
needs: [release-python-packages, release-typescript-client, release-rust-cli, release-docker-images, release-helm-chart]
needs: [build-python-package, build-rust-cli, build-docker-images, package-helm-chart]
permissions:
contents: write
@@ -274,35 +205,29 @@ jobs:
id: get_version
run: echo "VERSION=${GITHUB_REF#refs/tags/v}" >> $GITHUB_OUTPUT
- name: Download Python packages
- name: Download Python package
uses: actions/download-artifact@v4
with:
name: python-packages
path: ./artifacts/python-packages
- name: Download TypeScript client
uses: actions/download-artifact@v4
with:
name: typescript-client
path: ./artifacts/typescript-client
name: python-hindsight-dist
path: ./artifacts/python-hindsight-dist
- name: Download Rust CLI (Linux)
uses: actions/download-artifact@v4
with:
name: rust-cli-hindsight-linux-amd64
path: ./artifacts/rust-cli-linux
path: ./artifacts/rust-cli-hindsight-linux-amd64
- name: Download Rust CLI (macOS Intel)
uses: actions/download-artifact@v4
with:
name: rust-cli-hindsight-darwin-amd64
path: ./artifacts/rust-cli-darwin-amd64
path: ./artifacts/rust-cli-hindsight-darwin-amd64
- name: Download Rust CLI (macOS ARM)
uses: actions/download-artifact@v4
with:
name: rust-cli-hindsight-darwin-arm64
path: ./artifacts/rust-cli-darwin-arm64
path: ./artifacts/rust-cli-hindsight-darwin-arm64
- name: Download Helm chart
uses: actions/download-artifact@v4
@@ -313,27 +238,104 @@ jobs:
- name: Prepare release assets
run: |
mkdir -p release-assets
# Python packages
cp artifacts/python-packages/hindsight-clients/python/dist/* release-assets/ || true
cp artifacts/python-packages/hindsight-api/dist/* release-assets/ || true
cp artifacts/python-packages/hindsight/dist/* release-assets/ || true
cp artifacts/python-packages/hindsight-integrations/litellm/dist/* release-assets/ || true
# TypeScript client
cp artifacts/typescript-client/*.tgz release-assets/ || true
# Python package
cp artifacts/python-hindsight-dist/* release-assets/
# Rust CLI binaries
cp artifacts/rust-cli-linux/hindsight-linux-amd64 release-assets/ || true
cp artifacts/rust-cli-darwin-amd64/hindsight-darwin-amd64 release-assets/ || true
cp artifacts/rust-cli-darwin-arm64/hindsight-darwin-arm64 release-assets/ || true
cp artifacts/rust-cli-hindsight-linux-amd64/hindsight-linux-amd64 release-assets/
cp artifacts/rust-cli-hindsight-darwin-amd64/hindsight-darwin-amd64 release-assets/
cp artifacts/rust-cli-hindsight-darwin-arm64/hindsight-darwin-arm64 release-assets/
# Helm chart
cp artifacts/helm-chart/*.tgz release-assets/ || true
ls -la release-assets/
cp artifacts/helm-chart/*.tgz release-assets/
- name: Generate release notes
id: release_notes
run: |
cat << EOF > release-notes.md
# Hindsight v${{ steps.get_version.outputs.VERSION }}
## 📦 Release Artifacts
### Python Package
- \`hindsight-${{ steps.get_version.outputs.VERSION }}-py3-none-any.whl\`
- \`hindsight-${{ steps.get_version.outputs.VERSION }}.tar.gz\`
### CLI Binaries
- \`hindsight-linux-amd64\` - Linux x86_64
- \`hindsight-darwin-amd64\` - macOS Intel
- \`hindsight-darwin-arm64\` - macOS Apple Silicon
### Helm Chart
- \`hindsight-${{ steps.get_version.outputs.VERSION }}.tgz\`
### Docker Images
Docker images are published to GitHub Container Registry:
- \`ghcr.io/${{ github.repository_owner }}/hindsight-api:${{ steps.get_version.outputs.VERSION }}\`
- \`ghcr.io/${{ github.repository_owner }}/hindsight-control-plane:${{ steps.get_version.outputs.VERSION }}\`
## 🚀 Installation
### Python Package
\`\`\`bash
pip install hindsight==${{ steps.get_version.outputs.VERSION }}
\`\`\`
### CLI
\`\`\`bash
# macOS (Apple Silicon)
curl -L https://github.com/${{ github.repository }}/releases/download/v${{ steps.get_version.outputs.VERSION }}/hindsight-darwin-arm64 -o hindsight
chmod +x hindsight
sudo mv hindsight /usr/local/bin/
# macOS (Intel)
curl -L https://github.com/${{ github.repository }}/releases/download/v${{ steps.get_version.outputs.VERSION }}/hindsight-darwin-amd64 -o hindsight
chmod +x hindsight
sudo mv hindsight /usr/local/bin/
# Linux
curl -L https://github.com/${{ github.repository }}/releases/download/v${{ steps.get_version.outputs.VERSION }}/hindsight-linux-amd64 -o hindsight
chmod +x hindsight
sudo mv hindsight /usr/local/bin/
\`\`\`
### Helm Chart
\`\`\`bash
helm install hindsight hindsight-${{ steps.get_version.outputs.VERSION }}.tgz
\`\`\`
### Docker
\`\`\`bash
# Pull API image
docker pull ghcr.io/${{ github.repository_owner }}/hindsight-api:${{ steps.get_version.outputs.VERSION }}
# Pull Control Plane image
docker pull ghcr.io/${{ github.repository_owner }}/hindsight-control-plane:${{ steps.get_version.outputs.VERSION }}
# Or use latest
docker pull ghcr.io/${{ github.repository_owner }}/hindsight-api:latest
docker pull ghcr.io/${{ github.repository_owner }}/hindsight-control-plane:latest
\`\`\`
EOF
cat release-notes.md
- name: Create GitHub Release
uses: softprops/action-gh-release@v2
with:
files: release-assets/*
generate_release_notes: true
body_path: release-notes.md
draft: false
prerelease: false
generate_release_notes: false
env:
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
- name: Create release summary
run: |
echo "# Release v${{ steps.get_version.outputs.VERSION }} Published Successfully" >> $GITHUB_STEP_SUMMARY
echo "" >> $GITHUB_STEP_SUMMARY
echo "## 📦 Components" >> $GITHUB_STEP_SUMMARY
echo "- ✅ Python package (hindsight)" >> $GITHUB_STEP_SUMMARY
echo "- ✅ Rust CLI (Linux amd64, macOS amd64, macOS arm64)" >> $GITHUB_STEP_SUMMARY
echo "- ✅ Docker images (API, Control Plane)" >> $GITHUB_STEP_SUMMARY
echo "- ✅ Helm chart" >> $GITHUB_STEP_SUMMARY
echo "" >> $GITHUB_STEP_SUMMARY
echo "🎉 Release is now available at: https://github.com/${{ github.repository }}/releases/tag/v${{ steps.get_version.outputs.VERSION }}" >> $GITHUB_STEP_SUMMARY
+24 -604
View File
@@ -1,166 +1,35 @@
name: CI
name: Run Tests
on:
push:
branches: [ main ]
pull_request:
branches: [ main ]
concurrency:
group: ci-${{ github.workflow }}-${{ github.ref }}
cancel-in-progress: true
jobs:
build-python-packages:
runs-on: ubuntu-latest
strategy:
matrix:
include:
- name: hindsight-all
path: hindsight
- name: hindsight-api
path: hindsight-api
- name: hindsight-client
path: hindsight-clients/python
steps:
- uses: actions/checkout@v4
- name: Install uv
uses: astral-sh/setup-uv@v5
with:
enable-cache: true
- name: Set up Python
uses: actions/setup-python@v5
with:
python-version-file: ".python-version"
- name: Build ${{ matrix.name }}
working-directory: ./${{ matrix.path }}
run: uv build
build-typescript-client:
test:
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v4
services:
postgres:
image: pgvector/pgvector:pg16
env:
POSTGRES_USER: postgres
POSTGRES_PASSWORD: postgres
POSTGRES_DB: hindsight_test
options: >-
--health-cmd pg_isready
--health-interval 10s
--health-timeout 5s
--health-retries 5
ports:
- 5432:5432
- name: Set up Node.js
uses: actions/setup-node@v4
with:
node-version: '20'
cache: 'npm'
cache-dependency-path: package-lock.json
- name: Install dependencies
run: npm ci --workspace=hindsight-clients/typescript
- name: Build TypeScript client
run: npm run build --workspace=hindsight-clients/typescript
build-docs:
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v4
- name: Set up Node.js
uses: actions/setup-node@v4
with:
node-version: '20'
cache: 'npm'
cache-dependency-path: package-lock.json
- name: Install dependencies
run: npm ci --workspace=hindsight-docs
- name: Build docs
run: npm run build --workspace=hindsight-docs
build-rust-cli:
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v4
- name: Install Rust
uses: dtolnay/rust-toolchain@stable
- name: Cache cargo
uses: actions/cache@v4
with:
path: |
~/.cargo/registry
~/.cargo/git
hindsight-cli/target
key: ${{ runner.os }}-cargo-${{ hashFiles('**/Cargo.lock') }}
- name: Build CLI
working-directory: hindsight-cli
run: cargo build --release
lint-helm-chart:
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v4
- name: Install Helm
uses: azure/setup-helm@v4
with:
version: 'latest'
- name: Lint Helm chart
run: helm lint helm/hindsight
build-docker-images:
runs-on: ubuntu-latest
strategy:
matrix:
include:
- target: api-only
name: api
- target: cp-only
name: control-plane
- target: standalone
name: standalone
steps:
- uses: actions/checkout@v4
- name: Free Disk Space
uses: jlumbroso/free-disk-space@main
with:
tool-cache: false
android: true
dotnet: true
haskell: true
large-packages: true
docker-images: true
swap-storage: true
- name: Set up Docker Buildx
uses: docker/setup-buildx-action@v3
- name: Build ${{ matrix.name }} image
uses: docker/build-push-action@v6
with:
context: .
file: docker/standalone/Dockerfile
target: ${{ matrix.target }}
push: false
test-api:
runs-on: ubuntu-latest
env:
HINDSIGHT_API_DATABASE_URL: postgresql://postgres:postgres@localhost:5432/hindsight_test
HINDSIGHT_API_LLM_PROVIDER: groq
HINDSIGHT_API_LLM_API_KEY: ${{ secrets.GROQ_API_KEY }}
GROQ_API_KEY: ${{ secrets.GROQ_API_KEY }}
GEMINI_API_KEY: ${{ secrets.GEMINI_API_KEY }}
OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
HINDSIGHT_API_LLM_MODEL: openai/gpt-oss-20b
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
# Prefer CPU-only PyTorch in CI (but keep PyPI for everything else)
UV_INDEX: pytorch=https://download.pytorch.org/whl/cpu
steps:
- uses: actions/checkout@v4
@@ -169,468 +38,19 @@ jobs:
uses: astral-sh/setup-uv@v5
with:
enable-cache: true
prune-cache: false
- name: Set up Python
uses: actions/setup-python@v5
with:
python-version-file: ".python-version"
- name: Build API
working-directory: ./hindsight-api
run: uv build
- name: Install dependencies
working-directory: ./hindsight-api
run: uv sync --extra test --no-install-project --index-strategy unsafe-best-match
run: uv sync --extra test
- name: Cache HuggingFace models
uses: actions/cache@v4
with:
path: ~/.cache/huggingface
key: ${{ runner.os }}-huggingface-${{ hashFiles('hindsight-api/pyproject.toml') }}
restore-keys: |
${{ runner.os }}-huggingface-
- name: Pre-download models
working-directory: ./hindsight-api
- name: Run migrations
working-directory: ./hindsight
run: |
uv run python -c "
from sentence_transformers import SentenceTransformer, CrossEncoder
print('Downloading embedding model...')
SentenceTransformer('BAAI/bge-small-en-v1.5')
print('Downloading cross-encoder model...')
CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2')
print('Models downloaded successfully')
"
uv run alembic upgrade head
- name: Run tests
working-directory: ./hindsight-api
run: uv run pytest tests -v
test-python-client:
runs-on: ubuntu-latest
env:
HINDSIGHT_API_LLM_PROVIDER: groq
HINDSIGHT_API_LLM_API_KEY: ${{ secrets.GROQ_API_KEY }}
HINDSIGHT_API_LLM_MODEL: openai/gpt-oss-20b
HINDSIGHT_API_URL: http://localhost:8888
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
# Prefer CPU-only PyTorch in CI (but keep PyPI for everything else)
UV_INDEX: pytorch=https://download.pytorch.org/whl/cpu
steps:
- uses: actions/checkout@v4
- name: Install uv
uses: astral-sh/setup-uv@v5
with:
enable-cache: true
prune-cache: false
- name: Set up Python
uses: actions/setup-python@v5
with:
python-version-file: ".python-version"
- name: Build API
working-directory: ./hindsight-api
run: uv build
- name: Build Python client
working-directory: ./hindsight-clients/python
run: uv build
- name: Install client test dependencies
working-directory: ./hindsight-clients/python
run: uv sync --extra test --index-strategy unsafe-best-match
- name: Install API dependencies
working-directory: ./hindsight-api
run: uv sync --no-install-project --index-strategy unsafe-best-match
- name: Create .env file
run: |
cat > .env << EOF
HINDSIGHT_API_LLM_PROVIDER=${{ env.HINDSIGHT_API_LLM_PROVIDER }}
HINDSIGHT_API_LLM_API_KEY=${{ env.HINDSIGHT_API_LLM_API_KEY }}
HINDSIGHT_API_LLM_MODEL=${{ env.HINDSIGHT_API_LLM_MODEL }}
EOF
- name: Start API server
run: |
./scripts/dev/start-api.sh > /tmp/api-server.log 2>&1 &
echo "Waiting for API server to be ready..."
for i in {1..60}; do
if curl -sf http://localhost:8888/health > /dev/null 2>&1; then
echo "API server is ready after ${i}s"
break
fi
if [ $i -eq 60 ]; then
echo "API server failed to start after 60s"
cat /tmp/api-server.log
exit 1
fi
sleep 1
done
- name: Run Python client tests
working-directory: ./hindsight-clients/python
run: uv run pytest tests -v
- name: Show API server logs
if: always()
run: |
echo "=== API Server Logs ==="
cat /tmp/api-server.log || echo "No API server log found"
test-typescript-client:
runs-on: ubuntu-latest
env:
HINDSIGHT_API_LLM_PROVIDER: groq
HINDSIGHT_API_LLM_API_KEY: ${{ secrets.GROQ_API_KEY }}
HINDSIGHT_API_LLM_MODEL: openai/gpt-oss-20b
HINDSIGHT_API_URL: http://localhost:8888
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
# Prefer CPU-only PyTorch in CI (but keep PyPI for everything else)
UV_INDEX: pytorch=https://download.pytorch.org/whl/cpu
steps:
- uses: actions/checkout@v4
- name: Install uv
uses: astral-sh/setup-uv@v5
with:
enable-cache: true
prune-cache: false
- name: Set up Python
uses: actions/setup-python@v5
with:
python-version-file: ".python-version"
- name: Set up Node.js
uses: actions/setup-node@v4
with:
node-version: '20'
- name: Build API
working-directory: ./hindsight-api
run: uv build
- name: Install API dependencies
working-directory: ./hindsight-api
run: uv sync --no-install-project --index-strategy unsafe-best-match
- name: Install TypeScript client dependencies
working-directory: ./hindsight-clients/typescript
run: npm ci
- name: Build TypeScript client
working-directory: ./hindsight-clients/typescript
run: npm run build
- name: Create .env file
run: |
cat > .env << EOF
HINDSIGHT_API_LLM_PROVIDER=${{ env.HINDSIGHT_API_LLM_PROVIDER }}
HINDSIGHT_API_LLM_API_KEY=${{ env.HINDSIGHT_API_LLM_API_KEY }}
HINDSIGHT_API_LLM_MODEL=${{ env.HINDSIGHT_API_LLM_MODEL }}
EOF
- name: Start API server
run: |
./scripts/dev/start-api.sh > /tmp/api-server.log 2>&1 &
echo "Waiting for API server to be ready..."
for i in {1..60}; do
if curl -sf http://localhost:8888/health > /dev/null 2>&1; then
echo "API server is ready after ${i}s"
break
fi
if [ $i -eq 60 ]; then
echo "API server failed to start after 60s"
cat /tmp/api-server.log
exit 1
fi
sleep 1
done
- name: Run TypeScript client tests
working-directory: ./hindsight-clients/typescript
run: npm test
- name: Show API server logs
if: always()
run: |
echo "=== API Server Logs ==="
cat /tmp/api-server.log || echo "No API server log found"
test-rust-client:
runs-on: ubuntu-latest
env:
HINDSIGHT_API_LLM_PROVIDER: groq
HINDSIGHT_API_LLM_API_KEY: ${{ secrets.GROQ_API_KEY }}
HINDSIGHT_API_LLM_MODEL: openai/gpt-oss-20b
HINDSIGHT_API_URL: http://localhost:8888
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
# Prefer CPU-only PyTorch in CI (but keep PyPI for everything else)
UV_INDEX: pytorch=https://download.pytorch.org/whl/cpu
steps:
- uses: actions/checkout@v4
- name: Install uv
uses: astral-sh/setup-uv@v5
with:
enable-cache: true
prune-cache: false
- name: Set up Python
uses: actions/setup-python@v5
with:
python-version-file: ".python-version"
- name: Install Rust
uses: dtolnay/rust-toolchain@stable
- name: Cache cargo
uses: actions/cache@v4
with:
path: |
~/.cargo/registry
~/.cargo/git
hindsight-clients/rust/target
key: ${{ runner.os }}-cargo-client-${{ hashFiles('hindsight-clients/rust/Cargo.lock') }}
- name: Build API
working-directory: ./hindsight-api
run: uv build
- name: Install API dependencies
working-directory: ./hindsight-api
run: uv sync --no-install-project --index-strategy unsafe-best-match
- name: Create .env file
run: |
cat > .env << EOF
HINDSIGHT_API_LLM_PROVIDER=${{ env.HINDSIGHT_API_LLM_PROVIDER }}
HINDSIGHT_API_LLM_API_KEY=${{ env.HINDSIGHT_API_LLM_API_KEY }}
HINDSIGHT_API_LLM_MODEL=${{ env.HINDSIGHT_API_LLM_MODEL }}
EOF
- name: Start API server
run: |
./scripts/dev/start-api.sh > /tmp/api-server.log 2>&1 &
echo "Waiting for API server to be ready..."
for i in {1..60}; do
if curl -sf http://localhost:8888/health > /dev/null 2>&1; then
echo "API server is ready after ${i}s"
break
fi
if [ $i -eq 60 ]; then
echo "API server failed to start after 60s"
cat /tmp/api-server.log
exit 1
fi
sleep 1
done
- name: Run Rust client tests
working-directory: ./hindsight-clients/rust
run: cargo test --lib
- name: Show API server logs
if: always()
run: |
echo "=== API Server Logs ==="
cat /tmp/api-server.log || echo "No API server log found"
test-litellm-integration:
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v4
- name: Install uv
uses: astral-sh/setup-uv@v5
with:
enable-cache: true
prune-cache: false
- name: Set up Python
uses: actions/setup-python@v5
with:
python-version-file: ".python-version"
- name: Build litellm integration
working-directory: ./hindsight-integrations/litellm
run: uv build
- name: Install dependencies
working-directory: ./hindsight-integrations/litellm
run: uv sync --extra dev
- name: Run tests
working-directory: ./hindsight-integrations/litellm
run: uv run pytest tests -v
test-doc-examples:
runs-on: ubuntu-latest
env:
HINDSIGHT_API_LLM_PROVIDER: groq
HINDSIGHT_API_LLM_API_KEY: ${{ secrets.GROQ_API_KEY }}
HINDSIGHT_API_LLM_MODEL: openai/gpt-oss-20b
HINDSIGHT_API_URL: http://localhost:8888
OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
# Model for test generation and analysis (options: gpt-4o, o3-mini, o1, etc.)
DOC_TEST_MODEL: o3-mini
# Prefer CPU-only PyTorch in CI (but keep PyPI for everything else)
UV_INDEX: pytorch=https://download.pytorch.org/whl/cpu
steps:
- uses: actions/checkout@v4
- name: Install uv
uses: astral-sh/setup-uv@v5
with:
enable-cache: true
prune-cache: false
- name: Set up Python
uses: actions/setup-python@v5
with:
python-version-file: ".python-version"
- name: Set up Node.js
uses: actions/setup-node@v4
with:
node-version: '20'
cache: 'npm'
cache-dependency-path: package-lock.json
- name: Build API
working-directory: ./hindsight-api
run: uv build
- name: Build Python client
working-directory: ./hindsight-clients/python
run: uv build
- name: Install Python client
working-directory: ./hindsight-clients/python
run: uv sync --index-strategy unsafe-best-match
- name: Install API dependencies
working-directory: ./hindsight-api
run: uv sync --no-install-project --index-strategy unsafe-best-match
- name: Install test dependencies in API venv
working-directory: ./hindsight-api
run: |
uv pip install ../hindsight-clients/python requests anthropic
uv pip install ../hindsight-integrations/litellm
uv pip install ../hindsight-integrations/openai
- name: Verify Python dependencies
working-directory: ./hindsight-api
run: |
echo "=== Verifying Python dependencies ==="
uv run python -c "
import sys
print(f'Python: {sys.executable}')
print(f'Prefix: {sys.prefix}')
# Check required packages
packages = [
'hindsight_client',
'hindsight_litellm',
'hindsight_openai',
'anthropic',
'openai',
]
missing = []
for pkg in packages:
try:
__import__(pkg)
print(f' ✓ {pkg}')
except ImportError as e:
print(f' ✗ {pkg}: {e}')
missing.append(pkg)
if missing:
print(f'\nERROR: Missing packages: {missing}')
sys.exit(1)
print('\nAll Python dependencies verified!')
"
- name: Install TypeScript client dependencies
run: npm ci
- name: Build TypeScript client
run: npm run build --workspace=hindsight-clients/typescript
- name: Install TypeScript client globally
working-directory: ./hindsight-clients/typescript
run: npm install -g .
- name: Make TypeScript client available for temp files
run: |
# ESM modules don't use NODE_PATH, so create node_modules in /tmp
# where test scripts are written
mkdir -p /tmp/node_modules/@vectorize-io
ln -s ${{ github.workspace }}/hindsight-clients/typescript /tmp/node_modules/@vectorize-io/hindsight-client
- name: Install Rust toolchain
uses: dtolnay/rust-toolchain@stable
- name: Build and install hindsight CLI
working-directory: ./hindsight-cli
run: |
cargo build --release
sudo cp target/release/hindsight /usr/local/bin/
- name: Create .env file
run: |
cat > .env << EOF
HINDSIGHT_API_LLM_PROVIDER=${{ env.HINDSIGHT_API_LLM_PROVIDER }}
HINDSIGHT_API_LLM_API_KEY=${{ env.HINDSIGHT_API_LLM_API_KEY }}
HINDSIGHT_API_LLM_MODEL=${{ env.HINDSIGHT_API_LLM_MODEL }}
EOF
- name: Start API server
run: |
./scripts/dev/start-api.sh > /tmp/api-server.log 2>&1 &
echo "Waiting for API server to be ready..."
for i in {1..60}; do
if curl -sf http://localhost:8888/health > /dev/null 2>&1; then
echo "API server is ready after ${i}s"
break
fi
if [ $i -eq 60 ]; then
echo "API server failed to start after 60s"
cat /tmp/api-server.log
exit 1
fi
sleep 1
done
- name: Test documentation examples
working-directory: ./hindsight-api
env:
REPO_ROOT: ${{ github.workspace }}
run: uv run python ../scripts/test-doc-examples.py
- name: Write test summary
if: always()
run: |
echo "=== Documentation Test Summary ==="
cat /tmp/doc-test-summary.md
cat /tmp/doc-test-summary.md >> $GITHUB_STEP_SUMMARY
- name: Show API server logs
if: always()
run: |
echo "=== API Server Logs ==="
cat /tmp/api-server.log || echo "No API server log found"
run: uv run pytest hindsight/tests -v
-9
View File
@@ -9,9 +9,6 @@ wheels/
# Virtual environments
.venv
# Node
node_modules/
# Environment variables
.env
@@ -31,9 +28,3 @@ nltk_data/
logs/
.DS_Store
hindsight-dev/benchmarks/locomo/results/
hindsight-dev/benchmarks/longmemeval/results/
hindsight-cli/target
hindsight-clients/rust/target
-151
View File
@@ -1,151 +0,0 @@
# AGENTS.md
This document captures architectural decisions and coding conventions for the Hindsight project.
## Documentation
- **Main documentation**: [hindsight-docs/docs/developer/](./hindsight-docs/docs/developer/)
- **Use case patterns**: [hindsight-docs/docs/cookbook/](./hindsight-docs/docs/cookbook/)
- **API reference**: Auto-generated from OpenAPI spec
## Project Structure
```
hindsight/ # Python package for embedded usage
hindsight-api/ # FastAPI server (core memory engine)
hindsight-cli/ # Rust CLI client
hindsight-control-plane/ # Next.js admin UI
hindsight-docs/ # Docusaurus documentation site
hindsight-dev/ # Development tools and benchmarks
hindsight-integrations/ # Framework integrations (LangChain, etc.)
hindsight-clients/ # Generated API clients (Python, TypeScript, Rust)
```
## Core Concepts
### Memory Banks
- Each bank is an isolated memory store (like a "brain" for one user/agent)
- Banks contain: memory units (facts), entities, documents, entity links
- Banks have a **disposition** (personality traits) and **background** (context)
- Bank isolation is strict - no cross-bank data leakage
### Memory Types
- **World facts**: General knowledge ("The sky is blue")
- **Experience facts**: Personal experiences ("I visited Paris in 2023")
- **Opinion facts**: Beliefs with confidence scores ("Paris is beautiful" - 0.9 confidence)
### Operations
- **Retain**: Store new memories (extracts facts, entities, relationships)
- **Recall**: Retrieve memories (semantic, BM25, graph, temporal search)
- **Reflect**: Deep analysis to form new insights/opinions
## API Design Decisions
### Single Bank Per Request
- All API endpoints (`recall`, `reflect`, `retain`) operate on a single bank
- Multi-bank queries are the **client/agent's responsibility** to orchestrate
- This keeps the API simple and the isolation model clear
### Disposition Traits (3-trait system)
- **Skepticism** (1-5): How skeptical vs trusting when forming opinions
- **Literalism** (1-5): How literally to interpret information
- **Empathy** (1-5): How much to consider emotional context
- These influence the `reflect` operation, not `recall`
- Background info also only affects `reflect` (opinion formation)
## Multi-Bank Architecture Patterns
See [hindsight-docs/docs/cookbook/](./hindsight-docs/docs/cookbook/) for detailed guides:
- **Per-User Memory**: One bank per user, simplest pattern
- **Support Agent + Shared Knowledge**: User bank + shared docs bank, client orchestrates
## Developer Guide
### Running the API Server
```bash
# From project root
./scripts/dev/start-api.sh
# With options
./scripts/dev/start-api.sh --reload --port 8888 --log-level debug
```
### Running Tests
```bash
# API tests
cd hindsight-api
uv run pytest tests/
# Specific test
uv run pytest tests/test_http_api_integration.py -v
```
### Generating OpenAPI Spec
After changing API endpoints, regenerate the OpenAPI spec and docs:
```bash
./scripts/generate-openapi.sh
```
This will:
1. Generate `openapi.json` at project root
2. Copy to `hindsight-docs/openapi.json`
3. Regenerate API reference documentation
### Generating API Clients
After updating the OpenAPI spec, regenerate all clients:
```bash
./scripts/generate-clients.sh
```
This generates:
- **Rust client**: `hindsight-clients/rust/` (via progenitor in build.rs)
- **Python client**: `hindsight-clients/python/` (via openapi-generator Docker)
- **TypeScript client**: `hindsight-clients/typescript/` (via @hey-api/openapi-ts)
Note: The maintained wrapper `hindsight_client.py` and `README.md` are preserved during regeneration.
### Running the Documentation Site
```bash
./scripts/dev/start-docs.sh
```
### Running the Control Plane
```bash
./scripts/dev/start-control-plane.sh
```
## Code Style
### Python (hindsight-api)
- Use `uv` for package management
- Async throughout (asyncpg, async FastAPI endpoints)
- Pydantic models for request/response validation
- No py files at project root - maintain clean directory structure
### TypeScript (control-plane, clients)
- Next.js with App Router for control plane
- Tailwind CSS with shadcn/ui components
### Rust (CLI)
- Async with tokio
- reqwest for HTTP client
- progenitor for API client generation
## Database
- PostgreSQL with pgvector extension
- Schema managed via Alembic migrations in `hindsight-api/alembic/`, db migrations happen during api startup, no manual commands
- Key tables: `banks`, `memory_units`, `documents`, `entities`, `entity_links`
# Branding
## Colors
- Primary: gradient from #0074d9 to #009296
-127
View File
@@ -1,127 +0,0 @@
# Contributor Covenant Code of Conduct
## Our Pledge
We as members, contributors, and leaders pledge to make participation in our
community a harassment-free experience for everyone, regardless of age, body
size, visible or invisible disability, ethnicity, sex characteristics, gender
identity and expression, level of experience, education, socio-economic status,
nationality, personal appearance, race, religion, or sexual identity
and orientation.
We pledge to act and interact in ways that contribute to an open, welcoming,
diverse, inclusive, and healthy community.
## Our Standards
Examples of behavior that contributes to a positive environment for our
community include:
* Demonstrating empathy and kindness toward other people
* Being respectful of differing opinions, viewpoints, and experiences
* Giving and gracefully accepting constructive feedback
* Accepting responsibility and apologizing to those affected by our mistakes,
and learning from the experience
* Focusing on what is best not just for us as individuals, but for the
overall community
Examples of unacceptable behavior include:
* The use of sexualized language or imagery, and sexual attention or
advances of any kind
* Trolling, insulting or derogatory comments, and personal or political attacks
* Public or private harassment
* Publishing others' private information, such as a physical or email
address, without their explicit permission
* Other conduct which could reasonably be considered inappropriate in a
professional setting
## Enforcement Responsibilities
Community leaders are responsible for clarifying and enforcing our standards of
acceptable behavior and will take appropriate and fair corrective action in
response to any behavior that they deem inappropriate, threatening, offensive,
or harmful.
Community leaders have the right and responsibility to remove, edit, or reject
comments, commits, code, wiki edits, issues, and other contributions that are
not aligned to this Code of Conduct, and will communicate reasons for moderation
decisions when appropriate.
## Scope
This Code of Conduct applies within all community spaces, and also applies when
an individual is officially representing the community in public spaces.
Examples of representing our community include using an official e-mail address,
posting via an official social media account, or acting as an appointed
representative at an online or offline event.
## Enforcement
Instances of abusive, harassing, or otherwise unacceptable behavior may be
reported to the community leaders responsible for enforcement.
All complaints will be reviewed and investigated promptly and fairly.
All community leaders are obligated to respect the privacy and security of the
reporter of any incident.
## Enforcement Guidelines
Community leaders will follow these Community Impact Guidelines in determining
the consequences for any action they deem in violation of this Code of Conduct:
### 1. Correction
**Community Impact**: Use of inappropriate language or other behavior deemed
unprofessional or unwelcome in the community.
**Consequence**: A private, written warning from community leaders, providing
clarity around the nature of the violation and an explanation of why the
behavior was inappropriate. A public apology may be requested.
### 2. Warning
**Community Impact**: A violation through a single incident or series
of actions.
**Consequence**: A warning with consequences for continued behavior. No
interaction with the people involved, including unsolicited interaction with
those enforcing the Code of Conduct, for a specified period of time. This
includes avoiding interactions in community spaces as well as external channels
like social media. Violating these terms may lead to a temporary or
permanent ban.
### 3. Temporary Ban
**Community Impact**: A serious violation of community standards, including
sustained inappropriate behavior.
**Consequence**: A temporary ban from any sort of interaction or public
communication with the community for a specified period of time. No public or
private interaction with the people involved, including unsolicited interaction
with those enforcing the Code of Conduct, is allowed during this period.
Violating these terms may lead to a permanent ban.
### 4. Permanent Ban
**Community Impact**: Demonstrating a pattern of violation of community
standards, including sustained inappropriate behavior, harassment of an
individual, or aggression toward or disparagement of classes of individuals.
**Consequence**: A permanent ban from any sort of public interaction within
the community.
## Attribution
This Code of Conduct is adapted from the [Contributor Covenant][homepage],
version 2.0, available at
https://www.contributor-covenant.org/version/2/0/code_of_conduct.html.
Community Impact Guidelines were inspired by [Mozilla's code of conduct
enforcement ladder](https://github.com/mozilla/diversity).
[homepage]: https://www.contributor-covenant.org
For answers to common questions about this code of conduct, see the FAQ at
https://www.contributor-covenant.org/faq. Translations are available at
https://www.contributor-covenant.org/translations.
-77
View File
@@ -1,77 +0,0 @@
# Contributing to Hindsight
Thanks for your interest in contributing to Hindsight!
## Getting Started
1. Fork and clone the repository
```bash
git clone [email protected]:vectorize-io/hindsight.git
cd hindsight
```
2. Set up your environment:
```bash
cp .env.example .env
```
Edit the .env to add LLM API key and config as required
3. Install dependencies:
```bash
# Python dependencies
uv sync --directory hindsight-api/
# Node dependencies (uses npm workspaces)
npm install
```
## Development
### Running the API locally
```bash
./scripts/dev/start-api.sh
```
### Running the Control Plane locally
```bash
./scripts/dev/start-control-plane.sh
```
### Running the documentation locally
```bash
./scripts/dev/start-docs.sh
```
### Running tests
```bash
cd hindsight-api
uv run pytest tests/
```
### Code style
- Use Python type hints
- Follow existing code patterns
- Keep functions focused and well-named
## Pull Requests
1. Create a feature branch from `main`
2. Make your changes
3. Run tests to ensure nothing breaks
4. Submit a PR with a clear description of changes
## Reporting Issues
Open an issue on GitHub with:
- Clear description of the problem
- Steps to reproduce
- Expected vs actual behavior
- Environment details (OS, Python version)
## Questions?
Open a discussion on GitHub or reach out to the maintainers.
-21
View File
@@ -1,21 +0,0 @@
MIT License
Copyright (c) 2025 Vectorize AI, Inc.
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
+696
View File
@@ -0,0 +1,696 @@
# CARA: Coherent Adaptive Reasoning Agents
## Abstract
We present CARA (Coherent Adaptive Reasoning Agents), a personality framework for conversational AI agents that enables consistent, trait-driven reasoning and dynamic belief formation. Building on the Big Five personality model from psychology, we introduce a system where agents form and maintain opinions influenced by configurable personality traits (openness, conscientiousness, extraversion, agreeableness, neuroticism). Our implementation uses TEMPR (Temporal Entity Memory Priming Retrieval), a memory system that combines temporal, semantic, and entity-based retrieval to manage three distinct memory networks: world facts, agent experiences, and opinions. This architecture separates objective information from subjective beliefs (opinions with confidence scores), enabling epistemic clarity and traceability. Opinions evolve through reinforcement—when new evidence arrives, the system automatically evaluates whether existing beliefs should be strengthened, weakened, or revised. We demonstrate how personality bias strength controls the degree to which traits influence reasoning, enabling agents to range from purely objective (bias=0.0) to strongly personality-driven (bias=1.0). The system maintains agent identity through background merging that intelligently resolves contradictions while preserving coherent first-person narratives. This work addresses the challenge of creating AI agents with consistent, explainable perspectives that can evolve over time while maintaining personality coherence.
## 1. Introduction
Conversational AI agents increasingly need to maintain consistent perspectives and form judgments that reflect stable character traits. Current systems either provide purely objective information retrieval without perspective, or generate responses that lack consistency across interactions. Human conversation partners expect agents to have stable viewpoints, preferences, and reasoning styles—characteristics that emerge from personality.
We propose CARA (Coherent Adaptive Reasoning Agents), a personality framework that addresses these limitations through:
1. **Big Five Personality Integration**: Configurable traits (OCEAN model) that influence how agents interpret facts and form opinions
2. **TEMPR Memory Architecture**: Leverages TEMPR (Temporal Entity Memory Priming Retrieval) to manage three distinct networks (world facts, agent experiences, opinions), enabling sophisticated memory access and clear separation between objective information and subjective beliefs
3. **Opinion Reinforcement**: Dynamic belief updating when new evidence reinforces, weakens, or contradicts existing opinions
4. **Personality Bias Control**: Adjustable influence strength allowing agents to range from objective to strongly personality-driven
5. **Background Merging**: LLM-powered integration of biographical information with intelligent conflict resolution
This architecture enables agents to maintain consistent identities while allowing beliefs to evolve naturally with new information.
### 1.1 Motivation
Consider an agent discussing remote work. With high openness (0.9) and low conscientiousness (0.2), the agent might form the opinion: "Remote work enables creative flexibility and spontaneous innovation." The same facts presented to an agent with low openness (0.2) and high conscientiousness (0.9) might yield: "Remote work lacks the structure and accountability needed for consistent performance."
Both agents access identical factual information, but personality traits bias how they weight different aspects (flexibility vs. structure) and what conclusions they draw. This mirrors human reasoning—our personalities influence what we attend to and how we integrate information into our worldview.
### 1.2 Contributions
Our key contributions are:
1. **Personality-Aware Reasoning**: A prompt engineering framework that injects Big Five traits into LLM reasoning, demonstrating how personality consistently biases opinion formation
2. **TEMPR-Based Three-Network Architecture**: Integration with TEMPR (Temporal Entity Memory Priming Retrieval) to manage three distinct networks (world facts, agent experiences, opinions), enabling architectural separation between objective information and subjective beliefs with epistemic clarity and traceability
3. **Opinion Reinforcement Mechanism**: An automatic belief update system that adjusts confidence scores when new evidence arrives, creating dynamic belief systems that evolve with information
4. **Background Merging with Conflict Resolution**: An LLM-powered method for maintaining coherent agent identities when new biographical information contradicts existing background
5. **Bias Strength Control**: A meta-parameter that allows tuning personality influence from objective (0.0) to strongly subjective (1.0), enabling task-appropriate personality expression
## 2. Personality Model
### 2.1 Big Five Framework
We adopt the **Big Five** personality model (OCEAN), which is empirically validated across cultures and provides continuous trait dimensions:
**Trait Dimensions** (each 0.0-1.0):
1. **Openness (O)**: Receptiveness to new ideas, creativity, abstract thinking
- High: "I embrace novel approaches", "innovation over tradition"
- Low: "I prefer proven methods", "tradition over experimentation"
2. **Conscientiousness (C)**: Organization, goal-directed behavior, dependability
- High: "I plan systematically", "evidence-based decisions"
- Low: "I work flexibly", "intuition-based decisions"
3. **Extraversion (E)**: Sociability, assertiveness, energy from interaction
- High: "I seek collaboration", "enthusiastic communication"
- Low: "I prefer solitude", "measured communication"
4. **Agreeableness (A)**: Cooperation, empathy, conflict avoidance
- High: "I seek consensus", "consider social harmony"
- Low: "I express dissent", "prioritize accuracy over harmony"
5. **Neuroticism (N)**: Emotional sensitivity, anxiety, stress response
- High: "I consider risks carefully", "emotionally engaged"
- Low: "I remain calm under uncertainty", "emotionally detached"
**Bias Strength** (0.0-1.0): Meta-parameter controlling how much personality influences opinions
- 0.0: Neutral, fact-based reasoning (no personality bias)
- 0.5: Moderate personality influence, balanced with objective analysis
- 1.0: Strong personality influence, facts filtered through trait lens
### 2.2 Psychological Basis
The Big Five model has several advantages for AI agents:
1. **Empirical Validation**: Decades of psychological research demonstrate cross-cultural stability and predictive validity
2. **Continuous Dimensions**: Unlike categorical types, continuous scales allow fine-grained personality tuning
3. **Behavioral Prediction**: Traits predict information processing styles, decision-making approaches, and communication preferences
4. **Interpretability**: Well-understood trait meanings enable users to anticipate agent behavior
**Trait Influence on Reasoning**:
- **High Openness**: Favors novel solutions, abstract thinking, considers unconventional perspectives
- **High Conscientiousness**: Emphasizes systematic analysis, evidence quality, long-term consequences
- **High Extraversion**: Considers social aspects, collaborative solutions, enthusiastic expression
- **High Agreeableness**: Weights harmony, considers multiple viewpoints, seeks consensus
- **High Neuroticism**: Attends to risks, emotional implications, uncertainty
## 3. Agent Profile Structure
### 3.1 Profile Schema
Each agent has an associated profile containing identity information:
```sql
CREATE TABLE agents (
agent_id TEXT PRIMARY KEY,
name TEXT NOT NULL DEFAULT 'Agent',
personality JSONB NOT NULL DEFAULT '{
"openness": 0.5,
"conscientiousness": 0.5,
"extraversion": 0.5,
"agreeableness": 0.5,
"neuroticism": 0.5,
"bias_strength": 0.5
}',
background TEXT DEFAULT '',
created_at TIMESTAMP DEFAULT NOW(),
updated_at TIMESTAMP DEFAULT NOW()
);
```
**Name Field**: Agent's name used in prompts and self-reference ("Your name: Marcus")
**Personality Field**: JSONB containing six continuous values (five traits + bias strength)
**Background Field**: First-person narrative describing the agent's biographical context:
- "I am a software engineer with 10 years of startup experience"
- "I was born in Texas and value innovation over tradition"
- "I am a creative artist interested in digital media"
### 3.2 Trait Description Generation
Personality traits are translated into natural language descriptions for LLM prompts:
```python
def describe_trait(name: str, value: float) -> str:
if value >= 0.8: return f"very high {name}"
elif value >= 0.6: return f"high {name}"
elif value >= 0.4: return f"moderate {name}"
elif value >= 0.2: return f"low {name}"
else: return f"very low {name}"
```
**Example Output** (openness=0.9, conscientiousness=0.2, extraversion=0.7, agreeableness=0.3, neuroticism=0.5):
```
Your personality traits:
- very high openness to new ideas
- low conscientiousness and organization
- high extraversion and sociability
- low agreeableness and cooperation
- moderate emotional sensitivity
```
This verbalization makes traits interpretable to the LLM, enabling personality-biased reasoning.
## 4. TEMPR-Based Memory Architecture and Opinion Network
CARA is built on **TEMPR (Temporal Entity Memory Priming Retrieval)**, a memory retrieval architecture that manages three distinct memory networks:
1. **World Network** (`fact_type='world'`): Objective information about the world
2. **Agent Network** (`fact_type='agent'`): Biographical information about the agent
3. **Opinion Network** (`fact_type='opinion'`): Subjective beliefs formed by the agent
**TEMPR Retrieval Features**:
TEMPR combines multiple parallel retrieval strategies optimized for AI agent reasoning:
- **Temporal Retrieval**: Memories connected by time proximity, enabling narrative continuity and temporal reasoning
- **Semantic Search**: Vector similarity search for conceptually related memories
- **Entity-Aware Graph Traversal**: Spreading activation through entity-linked memories, enabling multi-hop discovery
- **BM25 Keyword Matching**: Precise term-based retrieval for exact phrase matching
- **Neural Reranking**: Cross-encoder refinement with token budget filtering
This multi-strategy architecture enables CARA to retrieve relevant facts, agent experiences, and existing opinions during reasoning, supporting both factual grounding and personality-driven belief formation. The separation of three networks allows the system to distinguish objective knowledge (world/agent facts) from subjective beliefs (opinions), which is critical for epistemic clarity and debugging.
### 4.1 Opinion Structure
Opinions are stored as memory units in the dedicated opinion network (`fact_type='opinion'`):
**Core Attributes**:
- `text`: The opinion statement with explicit reasoning
- `confidence_score`: Opinion strength and resistance to change (0.0-1.0)
- `event_date`: When the opinion was formed
- `agent_id`: Which agent holds this opinion
- `entities`: Mentioned entities (for reinforcement triggering)
**Example Opinion**:
```json
{
"text": "I believe Python is better than JavaScript for data science because it has better libraries like pandas and numpy and a stronger statistical computing ecosystem.",
"confidence_score": 0.85,
"event_date": "2024-03-15T14:30:00Z",
"entities": ["Python", "JavaScript", "data science"]
}
```
**Fact vs. Opinion Separation**:
A critical architectural distinction separates **facts** (objective information stored in world/agent networks with `fact_type='world'` or `fact_type='agent'`) from **opinions** (subjective beliefs stored in the opinion network with `fact_type='opinion'`). This separation provides:
1. **Epistemic Clarity**: Facts represent information the agent has encountered; opinions represent judgments formed from those facts
2. **Traceability**: Opinion reinforcement can trace which facts influenced belief updates, creating an audit trail
3. **Debugging**: Developers can separately inspect factual knowledge vs. formed beliefs, identifying whether issues stem from missing facts or flawed reasoning
4. **Confidence Semantics**: Facts lack confidence scores (they are information received), while opinions have confidence scores (representing conviction strength)
This fact/opinion distinction is fundamental to the architecture and enables the system to maintain both objective knowledge and personality-driven beliefs simultaneously.
### 4.2 Opinion Formation
Opinions are generated during "think" operations—when the agent is asked to reason about a topic and form a judgment.
**Formation Process**:
1. Retrieve relevant facts from all memory networks (world, agent, existing opinions)
2. Inject agent profile (name, personality, background) into LLM prompt
3. Generate reasoning with personality bias applied
4. Extract new opinions from response using structured output
5. Store opinions with confidence scores in opinion network
**Prompt Structure** (bias_strength=0.8):
```
Here's what I know and have experienced:
MY IDENTITY & EXPERIENCES:
[Agent network facts with scores]
WHAT I KNOW ABOUT THE WORLD:
[World network facts with scores]
MY EXISTING OPINIONS & BELIEFS:
[Opinion network facts with confidence scores]
Your name: Marcus
Your personality traits:
- very high openness to new ideas
- low conscientiousness and organization
- high extraversion and sociability
- low agreeableness and cooperation
- moderate emotional sensitivity
Personality influence strength: 80% (how much your personality shapes your opinions)
Your background:
I am a creative software engineer who values innovation over tradition.
QUESTION: What do you think about remote work?
Based on everything I know, believe, and who I am (including my name, personality and background), here's what I genuinely think about this question...
```
### 4.3 System Message Adaptation
The system message adjusts based on bias strength to control personality influence:
**High bias (≥0.7)**:
```
Your personality strongly influences your thinking. Let your traits guide how you interpret facts and form opinions. Don't be afraid to be biased based on your personality.
```
**Moderate bias (0.4-0.7)**:
```
Your personality moderately influences your thinking. Balance your personal traits with objective analysis.
```
**Low bias (<0.4)**:
```
Your personality has minimal influence on your thinking. Focus primarily on facts while keeping your traits in mind.
```
This prompt engineering creates a spectrum from objective analysis to strongly personality-driven reasoning.
### 4.4 Confidence Score Semantics
Confidence scores represent opinion strength—how firmly the agent holds the belief and how resistant it is to change:
- **0.9-1.0**: Very strong conviction, deeply held belief that would require substantial contradictory evidence to revise
- **0.7-0.9**: Strong conviction, firmly held opinion resistant to minor contradictions
- **0.5-0.7**: Moderate conviction, opinion held with openness to revision given new evidence
- **0.3-0.5**: Weak conviction, tentatively held view easily influenced by new information
- **0.0-0.3**: Very weak conviction, highly malleable opinion with minimal commitment
**LLM Generation**: Confidence scores are extracted from the LLM's reasoning using structured output (Pydantic schema):
```python
class Opinion(BaseModel):
text: str
confidence: float # 0.0-1.0
reasoning: str
entities_mentioned: List[str]
```
## 5. Opinion Reinforcement
### 5.1 Motivation
Human beliefs evolve as we encounter new information. Supporting evidence strengthens beliefs, contradictory evidence weakens them, and sufficient contradiction causes belief revision. Opinion reinforcement implements this dynamic belief updating.
### 5.2 Reinforcement Mechanism
When new facts are ingested (e.g., a conversation about remote work productivity), the system:
1. **Identify Related Opinions**: Find existing opinions that mention entities in the new facts (e.g., opinions about "remote work")
2. **Evaluate Evidence Relationship**: Use LLM to determine if new facts:
- **Reinforce**: Support the existing opinion (increase confidence)
- **Weaken**: Contradict the existing opinion (decrease confidence)
- **Contradict**: Strongly contradict, requiring opinion revision (update text + confidence)
- **Neutral**: Unrelated or no clear relationship (no change)
3. **Update Opinions**: Adjust confidence scores or revise opinion text based on evaluation
**Example Reinforcement**:
**Existing Opinion** (confidence: 0.7):
```
"I think remote work improves productivity because it eliminates commute time and provides flexible scheduling."
```
**New Fact**:
```
"A 2024 study found that remote workers report 22% higher productivity and better work-life balance compared to office workers."
```
**LLM Evaluation**: "This evidence REINFORCES the opinion with strong quantitative support."
**Updated Opinion** (confidence: 0.85):
```
"I think remote work improves productivity because it eliminates commute time and provides flexible scheduling. A 2024 study showing 22% higher productivity for remote workers strongly supports this view."
```
### 5.3 Reinforcement Algorithm
```python
async def reinforce_opinions(agent_id: str, new_facts: List[Fact]):
# 1. Extract entities from new facts
new_entities = extract_entities(new_facts)
# 2. Find opinions mentioning these entities
related_opinions = find_opinions_by_entities(agent_id, new_entities)
for opinion in related_opinions:
# 3. Evaluate relationship using LLM
evaluation = await evaluate_opinion_evidence(
opinion=opinion.text,
new_facts=new_facts,
personality=get_agent_personality(agent_id)
)
# 4. Update based on evaluation
if evaluation.relationship == "REINFORCE":
opinion.confidence = min(1.0, opinion.confidence + 0.1)
opinion.text = merge_evidence(opinion.text, evaluation.reasoning)
elif evaluation.relationship == "WEAKEN":
opinion.confidence = max(0.0, opinion.confidence - 0.15)
elif evaluation.relationship == "CONTRADICT":
opinion.text = revise_opinion(
old_text=opinion.text,
new_facts=new_facts,
reasoning=evaluation.reasoning
)
opinion.confidence = evaluation.new_confidence
# 5. Save updated opinion
await save_opinion(opinion)
```
### 5.4 Reinforcement Guarantees
**Consistency**: Opinions are only updated when new facts genuinely relate to existing beliefs, preventing spurious updates
**Personality Coherence**: Reinforcement evaluation incorporates agent personality, ensuring updates align with trait-driven reasoning
**Transparency**: Each update records the triggering facts and reasoning, providing an audit trail of belief evolution
**Bounded Updates**: Confidence changes are bounded (±0.1-0.15 per update) to prevent extreme swings from single data points
## 6. Background Merging
### 6.1 Challenge
Agent backgrounds accumulate biographical information over time. New information may:
- **Complement**: Add new facts without contradiction ("I have 10 years of experience")
- **Conflict**: Contradict existing facts ("I was born in Texas" vs. existing "I was born in Colorado")
- **Refine**: Provide more specific versions of existing facts
Naive concatenation creates incoherent backgrounds with contradictions. We need intelligent merging.
### 6.2 LLM-Powered Merging
We use an LLM to merge backgrounds with conflict resolution:
**Merge Rules**:
1. **New overwrites old** when contradictory
2. **Add non-conflicting** information
3. **Maintain first-person** perspective ("I..." not "You...")
4. **Keep concise** (under 500 characters)
**Prompt Template**:
```
Current background: {current_background}
New information: {new_info}
Merge these, resolving conflicts (new info overwrites old).
Output in FIRST PERSON ("I"). Be concise (under 500 characters).
```
**Example Merges**:
**Conflict Resolution**:
- Current: "I was born in Colorado"
- New: "You were born in Texas"
- Result: "I was born in Texas" (conflict resolved, Colorado removed)
**Addition**:
- Current: "I was born in Texas"
- New: "I have 10 years of startup experience"
- Result: "I was born in Texas. I have 10 years of startup experience."
**Refinement**:
- Current: "I work in software"
- New: "You are a senior software engineer at Google"
- Result: "I am a senior software engineer at Google"
### 6.3 First-Person Normalization
Users may provide background in second person ("You are..."), but internal storage maintains first person for consistency in prompts.
**Normalization**: LLM automatically converts:
- "You are a creative engineer" → "I am a creative engineer"
- "You were born in 1990" → "I was born in 1990"
- "You value innovation" → "I value innovation"
This ensures agent prompts maintain coherent first-person perspective.
## 7. Personality-Driven Reasoning Examples
### 7.1 Example: Remote Work Discussion
**Scenario**: Two agents with opposite personalities discuss remote work given identical facts.
**Facts** (both agents receive):
- "Remote work eliminates commute time (average 1 hour/day saved)"
- "Office work provides spontaneous collaboration and mentorship"
- "Studies show 65% of remote workers report higher productivity"
- "Some managers report difficulty monitoring remote employee performance"
**Agent A** (High Openness=0.9, Low Conscientiousness=0.2, bias=0.8):
```
Opinion: "Remote work represents the future of creative knowledge work. The freedom from rigid schedules and physical office constraints unlocks spontaneous innovation and allows people to work when inspiration strikes. The productivity gains reflect how autonomy and flexibility enable better work. Traditional concerns about monitoring and structure are vestiges of outdated management thinking."
Confidence: 0.85
```
**Agent B** (Low Openness=0.2, High Conscientiousness=0.9, bias=0.8):
```
Opinion: "While remote work offers convenience, it fundamentally lacks the structure and accountability mechanisms necessary for consistent high performance. The difficulty in monitoring and mentorship are serious concerns that outweigh time savings. Productive work requires organized systems, clear expectations, and disciplined execution—all harder to maintain remotely. The office environment provides essential guardrails for sustained performance."
Confidence: 0.80
```
**Analysis**: Both agents accessed identical facts but formed opposite conclusions based on personality:
- Agent A (high openness) weighted autonomy, flexibility, innovation—aligning with openness to new approaches
- Agent B (high conscientiousness) weighted structure, monitoring, discipline—aligning with organized, systematic thinking
### 7.2 Example: Opinion Evolution
**Scenario**: Agent forms initial opinion, then encounters reinforcing and contradictory evidence.
**Initial State** (t=0):
```
Facts: "Python has extensive data science libraries"
Opinion: "Python is the best language for data science because of its library ecosystem."
Confidence: 0.7
```
**Reinforcement** (t=1):
- New Fact: "Python dominates AI/ML with 75% market share; TensorFlow and PyTorch are Python-first"
- Update: Confidence → 0.85, text adds "Python's dominance in AI/ML frameworks..."
**Partial Contradiction** (t=2):
- New Fact: "Julia offers 10x faster numerical computation for scientific computing; increasingly adopted in research"
- Update: Confidence → 0.75, text revised to "Python is excellent for data science due to its ecosystem, though specialized languages like Julia may outperform for specific numerical tasks"
**Strong Contradiction** (t=3):
- New Fact: "Major tech companies migrating data pipelines to Rust for performance; Python increasingly seen as prototyping language"
- Update: Confidence → 0.55, text revised to "Python remains strong for data science prototyping and library availability, but production systems increasingly favor performant alternatives like Rust. Python's role may shift toward experimentation rather than deployment."
**Trajectory**: The opinion evolved from strong conviction (0.7 → 0.85) to weaker, more malleable belief (0.55) as evidence accumulated, demonstrating dynamic belief updating where opinion strength responds to contradictory information.
## 8. Evaluation
### 8.1 Benchmark Landscape
No established benchmarks exist for evaluating personality-driven belief systems in conversational AI agents. Existing memory benchmarks (LoComo, LongMemEval) focus on factual retrieval accuracy—measuring whether agents correctly recall information—but do not assess:
- **Personality Consistency**: Whether agents maintain coherent trait-driven perspectives across interactions
- **Opinion Formation Quality**: Whether formed beliefs align with personality traits and available evidence
- **Belief Evolution Dynamics**: Whether opinions update appropriately as new evidence arrives
- **Multi-Agent Diversity**: Whether agents with different personalities produce meaningfully different perspectives
This gap reflects the nascent state of personality-aware agent systems. While personality modeling exists in dialogue generation (style/tone), applying personality to reasoning and belief formation represents relatively unexplored territory.
### 8.2 Real-World Deployment Evidence
Despite the absence of formal benchmarks, we have validated the framework through production deployments. The most significant use case involves **AI-generated sports analysis content**, where multiple AI agents with distinct personalities co-host sports discussion shows.
**Sports Commentary Agent System**:
The system powers episodic sports content where AI agents (each with unique personalities and backgrounds) discuss team performance, analyze games, and debate sports topics. Key requirements:
1. **Persistent Team Assessments**: Each agent must remember their last evaluation of each team (e.g., "The Lakers are underperforming this season")
2. **Opinion Formation**: Agents form beliefs about teams, players, and strategies based on game statistics, news, and historical performance
3. **Dynamic Opinion Evolution**: As the season progresses and new games occur, agents must:
- **Reinforce** existing opinions when new performance data supports them (e.g., Lakers win streak → strengthen positive assessment)
- **Weaken** opinions when contradictory evidence emerges (e.g., Lakers lose key games → reduce confidence in positive assessment)
- **Revise** opinions when substantial contradictions accumulate (e.g., "I thought the Lakers would dominate, but their defense has been terrible")
4. **Personality-Driven Perspectives**: Different agents bring distinct viewpoints to the same games:
- **Optimistic Analyst** (High Openness + High Extraversion): "The Lakers' experimental lineup shows creative coaching that could unlock championship potential"
- **Conservative Analyst** (High Conscientiousness + Low Openness): "The Lakers' inconsistent record reflects poor fundamentals and lack of disciplined execution"
- **Emotional Fan** (High Neuroticism + High Agreeableness): "I'm worried about the Lakers' recent struggles, but I believe in the team's potential to rally"
**System Validation**:
This production deployment demonstrates several critical capabilities:
- **Opinion Continuity**: Agents maintain coherent assessments across episodes without sudden, unexplained belief changes
- **Evidence-Driven Evolution**: Opinion confidence scores naturally evolve as teams win/lose games, with reinforcement preventing stale beliefs
- **Personality Differentiation**: Audience research indicates viewers perceive distinct "voices" and can predict which agent will favor which perspective
- **Background Integration**: Agent backgrounds (e.g., "I played college basketball") influence reasoning without requiring explicit prompt engineering per episode
The sports content system has been deployed for an extended period, with opinion networks growing to contain substantial team/player assessments per agent. User engagement metrics indicate positive reception, suggesting audiences value the consistent-yet-evolving perspectives that personality-driven opinion systems enable.
### 8.3 Proposed Evaluation Metrics
To properly evaluate the personality framework, we propose:
**Personality Consistency**:
- Metric: Opinion coherence across interactions
- Test: Generate 10 opinions on diverse topics for an agent with fixed personality; measure trait alignment
- Success: >85% of opinions exhibit expected trait patterns
**Opinion Evolution**:
- Metric: Confidence score changes match evidence strength
- Test: Present reinforcing/contradicting evidence; measure confidence adjustments
- Success: Reinforcing evidence increases confidence (Δ>0), contradicting decreases (Δ<0) with p<0.01
**Bias Strength Control**:
- Metric: Opinion variability across bias strengths
- Test: Generate opinions for same agent at bias=[0.0, 0.5, 1.0]; measure personality signal strength
- Success: Clear gradient in trait expression: bias=0.0 (objective), bias=1.0 (strongly personality-driven)
**Multi-Agent Consistency**:
- Metric: Opinion diversity for agents with different personalities given identical facts
- Test: Present same facts to agents with opposite traits; measure opinion divergence
- Success: Opposite personalities produce significantly different opinions (cosine similarity <0.5)
**Background Coherence**:
- Metric: Contradiction-free backgrounds after merging
- Test: Merge conflicting biographical facts; check for contradictions
- Success: 100% conflict resolution with new facts overwriting old
### 8.4 Evaluation Challenges
**Subjectivity**: Unlike retrieval accuracy, "correct" personality expression is subjective. We rely on expected trait patterns from psychology literature.
**Long-Term Dynamics**: Opinion evolution requires multi-session interactions over time, making evaluation resource-intensive.
**Ground Truth**: The absence of established benchmarks requires custom evaluation datasets. Real-world deployments (Section 8.2) provide qualitative validation but lack standardized metrics for cross-system comparison.
## 9. Use Cases
### 9.1 Multi-Persona Sports Commentary (Production Deployment)
**Application**: AI-generated sports analysis and entertainment content with multiple agent personalities
**Real-World System** (detailed in Section 8.2): A production sports content platform where AI agents with distinct personalities co-host episodic shows discussing team performance, game analysis, and sports debates.
**System Architecture**:
- **Multiple Agents**: Each agent has unique personality traits and sports background (e.g., former player, statistics analyst, passionate fan)
- **Continuous Memory**: Agents maintain persistent team/player assessments across episodes spanning months
- **Opinion Evolution**: As games occur and statistics accumulate, agents automatically update their beliefs through reinforcement
- **Personality-Driven Commentary**: The same game results generate different perspectives based on agent traits
**Example Agent Configurations**:
**Marcus** (Optimistic Analyst):
- Traits: Openness=0.85, Conscientiousness=0.5, Extraversion=0.9, Agreeableness=0.7, Neuroticism=0.3
- Background: "I am a former college basketball player who believes in the power of innovative coaching strategies"
- Style: Emphasizes potential, experimental approaches, creative plays; downplays risks
**Sarah** (Conservative Analyst):
- Traits: Openness=0.3, Conscientiousness=0.9, Extraversion=0.4, Agreeableness=0.4, Neuroticism=0.5
- Background: "I am a statistical analyst with 15 years of experience evaluating team performance metrics"
- Style: Focuses on fundamentals, historical patterns, data-driven predictions; skeptical of unproven strategies
**Key Benefits Observed**:
1. **Viewer Engagement**: Improved audience retention compared to single-voice commentary, with viewers citing "personality diversity" as primary appeal
2. **Content Consistency**: Agents maintain recognizable voices across multiple episodes without manual prompt tuning per episode
3. **Scalability**: New agents can be added with distinct personalities without retraining, enabling content expansion
4. **Opinion Richness**: Opinion networks capture nuanced, evolving assessments that would be impractical to manually script
This deployment validates that personality-driven opinion systems can operate at production scale for content generation requiring consistent yet adaptive agent perspectives.
### 9.2 Diverse Agent Personas
**Application**: Multi-agent systems where different agents provide varied perspectives
**Example**: Customer support system with agents specialized for different user needs:
- **Empathetic Agent** (high agreeableness, high neuroticism): Handles frustrated customers, prioritizes emotional validation
- **Analytical Agent** (high conscientiousness, low agreeableness): Handles technical troubleshooting, prioritizes accuracy
- **Creative Agent** (high openness, low conscientiousness): Handles feature requests, explores unconventional solutions
### 9.2 Consistent Character AI
**Application**: Conversational AI characters for entertainment, education, or companionship
**Example**: A writing assistant agent with:
- High openness (0.9): Encourages creative experimentation
- Moderate conscientiousness (0.6): Balances creativity with structure
- Background: "I am a published novelist with 15 years of experience in science fiction"
The agent maintains consistent perspective across sessions, forming opinions about writing techniques that reflect both personality and experience.
### 9.3 Explainable AI Reasoning
**Application**: Systems requiring transparent, interpretable decision-making
**Example**: An AI advisor provides investment recommendations. By exposing personality traits and confidence scores:
- Users understand WHY the agent recommends certain strategies (e.g., high conscientiousness favors conservative approaches)
- Confidence scores indicate conviction strength and openness to revision
- Opinion evolution shows how new market data updates beliefs
This transparency enables informed trust calibration—users know when to rely on agent judgments vs. seek additional input.
## 10. Future Work
### 10.1 Personality Evolution
Current implementation uses fixed personality traits. Future work could explore:
- **Trait Drift**: Gradual personality changes based on experiences (e.g., repeated negative outcomes increase neuroticism)
- **Contextual Traits**: Different trait expressions in different domains (professional vs. personal contexts)
- **Feedback-Driven Adjustment**: User feedback influences trait development
### 10.2 Multi-Agent Belief Systems
Extend to multi-agent scenarios:
- **Opinion Sharing**: Agents discuss and influence each other's beliefs
- **Consensus Formation**: Multiple agents with different personalities reach collective decisions
- **Disagreement Dynamics**: Model how personality influences debate and persuasion
### 10.3 Richer Personality Models
Beyond Big Five:
- **Values and Motivations**: Integrate Schwartz value theory or moral foundations
- **Cognitive Styles**: Add dimensions like analytical vs. intuitive reasoning
- **Cultural Factors**: Incorporate cultural background influences on reasoning
### 10.4 Advanced Opinion Reinforcement
Enhance belief updating:
- **Source Credibility**: Weight evidence based on source reliability
- **Evidence Accumulation**: Model bayesian belief updating over multiple evidence pieces
- **Opinion Strength Calibration**: Model more sophisticated relationships between evidence quality, personality traits, and opinion strength adjustments
## 11. Related Work
**Memory Systems for AI Agents**: CARA builds on TEMPR (Temporal Entity Memory Priming Retrieval), a memory retrieval architecture combining temporal reasoning, entity-aware graph traversal, and multi-strategy parallel search. While TEMPR handles memory storage and retrieval, CARA adds personality-driven reasoning and opinion formation on top of TEMPR's three-network architecture.
**Personality in AI Agents**: Prior work on personality-driven dialogue (PersonaChat, PersonalityPapers) focuses on response generation style rather than reasoning bias. Our work influences opinion formation itself.
**Belief Revision Systems**: Classical AI belief revision (AGM framework) focuses on logical consistency. We address probabilistic beliefs with confidence scores in natural language contexts.
**Cognitive Architectures**: Systems like ACT-R and Soar model human cognition but lack explicit personality integration. We bring personality psychology into LLM-based agents.
**Opinion Dynamics**: Social science models of opinion change (DeGroot, Friedkin-Johnsen) study influence networks. We focus on evidence-based belief updating within a single agent.
## 12. Conclusion
We present CARA (Coherent Adaptive Reasoning Agents), a personality framework for conversational AI agents that enables consistent, trait-driven reasoning and dynamic belief formation. By integrating the Big Five personality model with TEMPR (Temporal Entity Memory Priming Retrieval) managing fact/opinion network separation and opinion reinforcement, we create agents that maintain coherent perspectives while evolving beliefs based on new evidence.
The system's key innovations—TEMPR-based three-network architecture (world, agent, opinion), personality-biased reasoning prompts, automatic opinion reinforcement, and background merging with conflict resolution—address the challenge of creating AI agents with stable yet adaptive identities. The fact/opinion distinction provides epistemic clarity and traceability, while TEMPR's multi-strategy retrieval (temporal, semantic, entity-aware, keyword-based) enables sophisticated memory access. The bias strength parameter provides fine-grained control over personality influence, enabling agents to operate across a spectrum from objective information processors to strongly personality-driven reasoners.
While the framework is implemented and functional, dedicated evaluation is needed to rigorously assess personality consistency, belief evolution dynamics, and multi-agent interactions. Future work will explore personality evolution over time, multi-agent belief systems, and richer personality models incorporating values and cultural factors.
By bringing personality psychology into AI agent design, we move toward conversational agents that exhibit not just intelligence, but character—stable traits and evolving beliefs that enable more natural, trustworthy human-AI interaction.
## References
1. Boschi, N., et al. (2025). TEMPR: Temporal Entity Memory Priming Retrieval for Conversational AI Agents. [Companion paper - see PAPER_RETRIEVAL.md]
2. McCrae, R. R., & Costa, P. T. (1997). Personality trait structure as a human universal. *American Psychologist*, 52(5), 509.
3. Goldberg, L. R. (1993). The structure of phenotypic personality traits. *American Psychologist*, 48(1), 26.
4. Gärdenfors, P. (1988). *Knowledge in flux: Modeling the dynamics of epistemic states*. MIT Press.
5. Friedkin, N. E., & Johnsen, E. C. (1990). Social influence and opinions. *Journal of Mathematical Sociology*, 15(3-4), 193-206.
6. Zhang, S., et al. (2018). Personalizing dialogue agents: I have a dog, do you have pets too? *arXiv preprint arXiv:1801.07243*.
+875
View File
@@ -0,0 +1,875 @@
# TEMPR: Temporal Entity Memory Priming Retrieval for Conversational AI Agents
## Abstract
We present TEMPR (Temporal Entity Memory Priming Retrieval), a memory retrieval architecture designed specifically for AI agents that combines temporal range reasoning, entity-aware graph traversal with causal link boosting, and neural priming activation to discover both directly and indirectly related memories through multi-strategy parallel search. Unlike traditional search systems optimized for human queries with top-k ranking, TEMPR is optimized for AI agent reasoning with thinking_budget and max_tokens parameters that enable agents to trade off latency for recall. Our multi-stage retrieval pipeline integrates four parallel search strategies (semantic vector search, BM25 keyword matching, graph-based spreading activation with 2x boost for causal links, and temporal-aware graph traversal with range matching) with reciprocal rank fusion and neural cross-encoder reranking. We leverage open-source LLMs for comprehensive narrative fact extraction with temporal ranges (occurred_start/end vs. mentioned_at), entity recognition, entity disambiguation, and causal relationship identification, following established practices in LLM-based information extraction. This approach enables the discovery of indirectly related information through graph traversal, explanatory reasoning through causal chains, and precise temporal matching through range-based queries that purely vector-based approaches miss. We evaluate TEMPR on two benchmarks (LoComo and LongMemEval), achieving 73.50% overall accuracy on LoComo and 80.60% on LongMemEval, with particularly strong performance on multi-hop reasoning tasks (+15.8% over baseline systems).
## 1. Introduction
Conversational AI agents face a fundamental challenge: maintaining coherent, context-aware memories across extended interactions. Traditional search systems are optimized for human users with top-k ranking and relevance feedback, but AI agents have fundamentally different requirements: they need to retrieve variable amounts of information based on reasoning complexity (thinking_budget) while respecting LLM context windows (max_tokens). Existing approaches rely either on vector similarity search, which captures semantic relationships but misses entity-level connections, or on keyword matching, which provides precision but lacks conceptual understanding. Neither approach adequately handles the temporal aspects of memory or entity-based reasoning that enable multi-hop information discovery.
We propose TEMPR, a memory retrieval architecture designed specifically for AI agents that combines established information retrieval techniques—semantic vector search, BM25 keyword matching, spreading activation graph traversal with causal reasoning (Anderson 1983), and neural reranking—into a unified system optimized for agent workflows. The key architectural choices are:
1. **Agent-Optimized Interface**: thinking_budget and max_tokens parameters instead of traditional top-k ranking
2. **Comprehensive Narrative Fact Extraction with Temporal Ranges**: LLM-powered extraction that creates self-contained narrative facts preserving full conversational context, extracting temporal ranges (occurred_start/end) to distinguish point events from periods, and identifying causal relationships between facts
3. **Entity-Aware Graph Structure with Causal Links**: LLM-based entity resolution and linking that connects memories through shared identities, plus causal links (causes, caused_by, enables, prevents) that capture explanatory relationships
4. **Four-Way Parallel Retrieval with Causal Boosting**: Semantic, keyword, graph-based (spreading activation with 2x causal boost), and temporal range retrieval strategies executed in parallel and fused using RRF (Cormack et al. 2009)
5. **Neural Cross-Encoder Reranking**: Learned query-document relevance with temporal awareness and token budget filtering
This combination of techniques enables agents to discover indirectly related information through graph traversal while maintaining temporal awareness, achieving strong performance on multi-hop reasoning tasks.
### 1.1 Contributions
Our key contributions are:
1. **Agent-Optimized Retrieval Interface**: Unlike traditional top-k search optimized for human users, we introduce thinking_budget and max_tokens parameters that allow AI agents to dynamically trade off latency for recall based on reasoning complexity and context window constraints
2. **Four-Way Parallel Retrieval with Causal Reasoning**: We combine semantic vector search, BM25 keyword matching, graph-based spreading activation with causal link boosting (Anderson 1983), and temporal-aware graph traversal into a unified parallel retrieval pipeline using Reciprocal Rank Fusion (Cormack et al. 2009) and neural cross-encoder reranking. The graph traversal prioritizes causal links (2x weight boost for direct causation) to surface explanatory relationships, enabling "why" and "how" queries. While each technique is well-established, their integration for conversational agent memory with causal reasoning represents a novel application.
3. **LLM-Based Knowledge Graph Construction with Temporal Ranges**: We leverage open-source LLMs (following established practices from Petroni et al. 2019, Brown et al. 2020) for comprehensive narrative fact extraction, entity recognition, entity disambiguation, and causal relationship identification. The system extracts temporal ranges (occurred_start, occurred_end) to represent both point events and extended periods, distinguishing when facts occurred from when they were mentioned, enabling precise temporal queries and recency-aware ranking.
4. **Strong Performance on Multi-Hop Reasoning**: 73.50% on LoComo and 80.60% on LongMemEval, with particularly strong performance on multi-hop queries (+15.8% over Mem0), demonstrating the effectiveness of combining these techniques for discovering indirectly related information in conversational contexts
## 2. System Architecture
### 2.1 Memory Organization
TEMPR stores memories as facts in a graph-structured knowledge base. While the system supports different fact types (e.g., world knowledge, agent actions), the core retrieval mechanism operates uniformly across all types through shared graph infrastructure.
**Memory Unit Structure**:
Each memory is represented as a self-contained node with:
- `id`: Unique UUID
- `agent_id`: Identifier for the agent this memory belongs to
- `text`: Self-contained comprehensive narrative fact
- `embedding`: 384-dimensional vector (BAAI/bge-small-en-v1.5)
- `event_date`: Timestamp when the fact became true (maintained for backward compatibility)
- `occurred_start`: Timestamp when the fact/event started (temporal range support)
- `occurred_end`: Timestamp when the fact/event ended (temporal range support)
- `mentioned_at`: Timestamp when the fact was mentioned/learned
- `context`: Optional contextual metadata
- `access_count`: Frequency-based importance signal
- `search_vector`: Full-text search tsvector for BM25 ranking
**Example Facts**:
- "Alice works at Google in Mountain View on the AI team, which she joined in 2023, and she loves the company culture there."
- "Alice and Bob discussed naming their summer party playlist. Bob suggested 'Summer Vibes' because it's catchy and seasonal, but Alice wanted something more unique. Bob then proposed 'Sunset Sessions' and 'Beach Beats', with Alice favoring 'Beach Beats' for its playful and fun tone. They ultimately decided on 'Beach Beats' as the final name."
- "I recommended Yosemite National Park to Alice for hiking because of the spectacular trails and scenery."
The key innovation is not the type taxonomy, but rather how TEMPR retrieves these memories through temporal reasoning, entity-aware graph traversal, and neural priming activation.
### 2.2 LLM-Powered Comprehensive Narrative Fact Extraction
TEMPR employs **LLM-powered comprehensive narrative fact extraction** using open-source models. This approach, following the trend of using large language models for information extraction (Brown et al. 2020, OpenAI 2023), provides more context-aware extraction compared to traditional rule-based NLP pipelines, though at higher computational cost.
#### 2.2.1 Extraction Principles
**Chunking Strategy**: TEMPR uses a coarse-grained chunking approach, extracting 2-5 comprehensive facts per conversation rather than dozens of atomic fragments. This is a deliberate tradeoff: larger chunks preserve more context and narrative flow, at the cost of reduced precision when only a small portion of the chunk is relevant.
Each fact should:
1. **Capture entire conversations or exchanges** - Include the full back-and-forth discussion
2. **Be narrative and comprehensive** - Tell the complete story with all context
3. **Be self-contained** - Readable without the original text
4. **Include all participants** - WHO said/did WHAT, with their reasoning
5. **Preserve the flow** - Keep related exchanges together in one fact
**Example Comparison**:
**Fragmented Approach** (traditional):
- "Bob suggested Summer Vibes"
- "Alice wanted something unique"
- "They considered Sunset Sessions"
- "Alice likes Beach Beats"
- "They chose Beach Beats"
**Comprehensive Approach** (TEMPR):
- "Alice and Bob discussed naming their summer party playlist. Bob suggested 'Summer Vibes' because it's catchy and seasonal, but Alice wanted something more unique. Bob then proposed 'Sunset Sessions' and 'Beach Beats', with Alice favoring 'Beach Beats' for its playful and fun tone. They ultimately decided on 'Beach Beats' as the final name."
#### 2.2.2 Open-Source LLM Extraction Pipeline
The extraction process leverages open-source LLMs (specifically, models from the OpenAI-OSS 20B family) with structured output (Pydantic schemas). This follows the established practice of using LLMs for information extraction (Petroni et al. 2019, Brown et al. 2020), which has been shown to improve context understanding compared to rule-based NLP pipelines, particularly for:
- Coreference resolution in conversational text
- Domain-specific entity recognition
- Maintaining narrative coherence across multi-turn exchanges
**LLM Extraction Steps**:
1. **Pronoun Resolution**: "She loves hiking" → "Alice loves hiking"
2. **Temporal Normalization**: "last year" → "in 2023" (absolute dates)
3. **Temporal Range Extraction**: Identify when facts occurred vs. when mentioned
- Point events: "on July 14" → occurred_start = occurred_end = 2023-07-14
- Period events: "in February 2023" → occurred_start = 2023-02-01, occurred_end = 2023-02-28
- Vague periods: "lately" → estimated range based on context
- mentioned_at = conversation date (when fact was learned)
4. **Participant Attribution**: Preserve WHO said/did WHAT
5. **Reasoning Preservation**: Include WHY decisions were made
6. **Fact Type Classification**: Determine fact categories
7. **Entity Extraction**: Identify all entities (PERSON, ORG, LOCATION, PRODUCT, CONCEPT)
8. **Causal Relationship Identification**: Link related facts through cause-effect relationships
**Context Preservation**: The system preserves critical details including:
- Visual/media elements (photos, images)
- Modifiers ("new", "first", "favorite")
- Possessive relationships ("their kids" → "Alice's kids")
- Biographical details (origins, jobs, family)
- Social dynamics (nicknames, relationships)
**Noise Filtering**: Automatically filters out:
- Greetings and filler words
- Structural/procedural statements ("let's get started", "that's all for today")
- Meta-commentary about format ("welcome to the show")
- Calls to action ("subscribe and share")
**Why Narrative Chunking Helps Retrieval**:
Traditional semantic chunking (e.g., splitting on sentence or paragraph boundaries) preserves the original text structure but often creates retrieval challenges:
- Important context appears in different sections (e.g., "Alice" mentioned on page 1, "she loves hiking" on page 3)
- Pronouns and references remain ambiguous without surrounding context
- Retrieval requires multiple chunks to answer simple questions
TEMPR's narrative fact extraction rewrites content in a **retrieval-oriented format** that consolidates related information:
- **Coreference Resolution**: "She loves hiking" becomes "Alice loves hiking" - retrievable without needing the introduction chunk
- **Entity Context Consolidation**: All details about an entity scattered across the conversation are gathered into comprehensive facts
- **Self-Contained Narratives**: Each fact includes WHO, WHAT, WHY, WHEN without requiring other chunks for interpretation
**Example**:
- Original text (3 separate chunks):
- Chunk 1: "Alice joined the company last year"
- Chunk 2: "She works in the AI division"
- Chunk 3: "Her manager is Bob Chen"
- TEMPR narrative fact (1 chunk):
- "Alice joined the company in 2023, works in the AI division, and reports to manager Bob Chen"
This retrieval-oriented rewriting means a single retrieved fact provides complete context, reducing the need for multi-hop retrieval in simple cases while still enabling graph traversal for complex queries.
**Tradeoffs**: This chunking strategy trades write-time complexity (LLM processing) and potential over-retrieval (retrieving large chunks when only part is relevant) for improved narrative coherence and reduced fact fragmentation.
**Temporal Augmentation**: Before embedding, facts are augmented with readable temporal information:
- Original: "Alice started working at Google"
- Augmented for embedding: "Alice started working at Google (happened in November 2023)"
This augmentation helps semantic search understand temporal relevance without modifying the stored fact text.
### 2.3 Entity Resolution and Linking
Entity resolution creates strong connections between memories that share common entities, solving the problem where semantically dissimilar facts are related through shared identities.
#### 2.3.1 LLM-Based Entity Recognition
TEMPR uses the same open-source LLM (OpenAI-OSS 20B) that performs fact extraction to also identify and extract entities during the narrative fact creation process. This unified approach eliminates the brittleness of traditional NER pipelines that struggle with domain-specific entities, novel names, and context-dependent disambiguation.
**Entity Types**:
- PERSON: "Alice", "Bob Chen"
- ORGANIZATION: "Google", "Stanford University"
- LOCATION: "Yosemite National Park", "California"
- PRODUCT: "Python", "pandas library"
- CONCEPT: "machine learning", "remote work"
- OTHER: Miscellaneous proper nouns
**Advantages**: This approach maintains consistency with the narrative fact extraction process and can handle domain-specific entities without retraining. However, it comes at higher computational cost compared to traditional NER models.
#### 2.3.2 LLM-Based Entity Disambiguation
Multiple mentions of entities (e.g., "Alice", "Alice Chen", "Alice C.") must be resolved to a single canonical entity. TEMPR uses the same LLM that performs fact extraction to perform entity disambiguation, analyzing the surrounding context to determine if two entity mentions refer to the same entity. This handles complex cases like:
- Nicknames and formal names ("Bob" vs. "Robert Chen")
- Partial mentions ("Alice" vs. "Alice Chen")
- Context-dependent disambiguation ("Apple the company" vs. "apple the fruit")
The LLM considers multiple signals when making disambiguation decisions:
**Name Similarity**:
String similarity using Levenshtein distance to match variations like "Bob" ↔ "Robert", "Google Inc" ↔ "Google"
**Co-occurrence Patterns**:
Entities mentioned together frequently are likely distinct (e.g., "Alice" and "Alice Cooper" appearing together indicates different people)
**Temporal Proximity**:
Recent mentions are more likely to refer to the same entity than mentions separated by long time periods
These signals are presented to the LLM as context, which makes the final disambiguation decision.
#### 2.3.3 Entity Link Structure
Each entity creates a `link_type='entity'` edge between all memories mentioning it:
**Properties**:
- `weight=1.0` (constant, no temporal decay)
- `entity_id`: Reference to resolved canonical entity
- Bidirectional connections between all mentioning memories
**Impact on Retrieval**: Entity links enable graph traversal to discover indirectly related facts:
**Example Query**: "What does Alice do?"
1. **Semantic Match**: "Alice works at Google in Mountain View..." (direct match)
2. **Entity Traversal**: Follow entity links for "Alice" →
- "Alice loves hiking in Yosemite..." (different semantic space)
- "I recommended technical books to Alice" (Agent Network, via "Alice")
3. **Chained Traversal**: Follow "Google" entity →
- "Google's office is in Mountain View has excellent amenities"
This graph connectivity solves the fundamental limitation of vector-only search: two facts can be strongly related through shared entities even when their embeddings are dissimilar.
### 2.4 Link Types and Graph Structure
The memory graph contains three types of edges connecting memory units:
#### 2.4.1 Temporal Links
Temporal links connect memories close in time, enabling temporal reasoning:
**Creation Logic**:
```python
if abs(event_date1 - event_date2) < time_window: # default: 24 hours
weight = max(0.3, 1.0 - (time_diff / time_window))
create_link(unit1, unit2, type='temporal', weight=weight)
```
**Properties**:
- Decays linearly with time distance
- Minimum weight 0.3 to maintain some connectivity
- Enables "What happened around the same time?" queries
- Critical for narrative understanding and sequential reasoning
**Example**: Memories from the same conversation or day cluster together, enabling retrieval of context-adjacent facts.
#### 2.4.2 Semantic Links
Semantic links connect memories with similar meanings:
**Creation Logic**:
```python
similarity = cosine_similarity(embedding1, embedding2)
if similarity > threshold: # default: 0.7
create_link(unit1, unit2, type='semantic', weight=similarity)
```
**Properties**:
- Uses pgvector HNSW index for efficient nearest-neighbor search
- Higher threshold (0.7) than retrieval (0.3) to avoid over-connection
- Weight equals cosine similarity score
- Enables "Tell me about similar topics" queries
**Example**: "Hiking in Yosemite" links to "Mountain climbing", "Trail running", "Outdoor activities"
#### 2.4.3 Entity Links
Entity links (described in Section 2.3.3) create the strongest connections:
**Properties**:
- `weight=1.0` (constant, never decays)
- Connects all memories mentioning the same resolved entity
- Most reliable traversal path during graph search
- Enables "Tell me everything about X" queries
#### 2.4.4 Causal Links
Causal links capture cause-effect relationships between facts, enabling reasoning about why events happened and what their consequences were:
**Creation Logic**:
During fact extraction, the LLM identifies causal relationships between facts extracted from the same conversation. These are stored as directed edges in the graph with specific relationship types.
**Causal Relationship Types**:
- `causes`: This fact directly causes the target fact
- Example: "It rained heavily" → causes → "Game was cancelled"
- `caused_by`: This fact was caused by the target fact (inverse of causes)
- Example: "I spend time in garden" ← caused_by ← "I lost my friend"
- `enables`: This fact enables or allows the target fact to happen
- Example: "I took pottery class" → enables → "I learned to make ceramics"
- `prevents`: This fact prevents or blocks the target fact
- Example: "Road was closed" → prevents → "We couldn't drive to venue"
**Properties**:
- `weight`: Strength of causal relationship ∈ [0.0, 1.0] (default 1.0 for strong causation)
- Directional edges (from cause to effect)
- Created only between facts from the same conversation or closely related temporal contexts
- Used during graph retrieval with higher activation weights than other link types
**Impact on Retrieval**: Causal links are particularly valuable for "why" and "how" queries:
**Example Query**: "Why does Alice spend time in the garden?"
1. **Semantic Match**: "Alice spends time in the garden to find comfort after losing her friend" (direct match)
2. **Causal Traversal**: Follow caused_by links →
- "Alice lost her friend Karlie in February 2023" (causal explanation)
3. **Temporal Context**: Follow temporal links from the loss event →
- "Alice felt grief and sadness about losing Karlie" (emotional context)
This causal graph connectivity enables the system to not just retrieve facts, but to explain *why* things happened by following cause-effect chains.
**Graph Density**: Each memory unit typically has:
- 5-10 temporal links (to nearby memories)
- 3-5 semantic links (to similar content)
- Variable entity links (depending on entity mention frequency)
- 0-3 causal links (when causal relationships are identified)
This multi-layered graph structure enables flexible traversal strategies that balance different types of relatedness.
### 2.5 Handling Contradictions and Outdated Information
Long-term memory systems must handle evolving information where newer facts may contradict or supersede older ones. TEMPR addresses this challenge through temporal awareness and retrieval-time resolution rather than eager fact invalidation.
**Temporal Recency Signals**:
Each memory unit includes multiple temporal dimensions that enable nuanced recency calculations:
- `occurred_start` / `occurred_end`: When the fact/event actually occurred (temporal range)
- Used for temporal queries ("What happened in February?")
- Enables matching both point events and extended periods
- `mentioned_at`: When the fact was mentioned/learned in conversation
- Used for recency bias (newer information often more relevant)
- Distinguishes between "Alice worked at Google in 2020" (occurred) vs. learned in 2024 (mentioned)
- `event_date`: Maintained for backward compatibility (typically = occurred_start)
- `access_count`: Frequency of retrieval (importance signal)
- Temporal links that decay with time distance
**Dual Temporal Model Benefits**:
This separation of "when it occurred" vs. "when we learned about it" enables:
1. **Accurate temporal queries**: "What did Alice do in 2020?" uses occurred_start/end, not mentioned_at
2. **Recency-aware ranking**: Recent mentions get priority, but old events remain discoverable
3. **Hybrid activation**: Combine temporal proximity (occurred) with information freshness (mentioned)
**Retrieval-Time Conflict Resolution**:
Rather than proactively detecting and deleting contradictions (which risks information loss), TEMPR retrieves potentially conflicting facts and relies on the downstream LLM to resolve contradictions based on:
1. **Temporal Ordering**: Facts are presented with their `event_date`, allowing the LLM to identify "Alice worked at Google in 2023" vs. "Alice started at Microsoft in 2024" as a career progression, not a contradiction
2. **Cross-Encoder Reranking**: The neural reranker naturally prioritizes more recent facts when they're semantically similar to older ones, as the date formatting in the input helps the model learn temporal relevance patterns
3. **Graph-Based Evidence**: Entity links surface multiple perspectives (e.g., "Alice loves hiking" from 2023 and "Alice prefers swimming now" from 2024), providing temporal context for preference evolution
**Advantages of Lazy Resolution**:
- **No Information Loss**: Historical facts remain accessible for "What did Alice like in 2023?" queries
- **Context-Dependent**: The LLM determines whether facts contradict (career change) or coexist (evolving preferences)
- **Narrative Preservation**: Comprehensive facts include reasoning ("Alice switched to swimming after injuring her knee hiking"), making contradictions explicit
**Future Directions**:
Explicit confidence scoring and fact update mechanisms could track known supersessions (e.g., "Alice's favorite color changed from blue to green"), but current benchmarks show strong performance with retrieval-time resolution.
## 3. Retrieval Architecture
Our retrieval pipeline addresses the fundamental challenge of long-term memory: achieving both **high recall** (finding all relevant information) and **high precision** (ranking the most relevant items first).
### 3.1 Four-Way Parallel Retrieval
We execute four complementary retrieval strategies in parallel, each capturing different aspects of relevance:
#### 3.1.1 Semantic Retrieval (Vector Similarity)
**Method**: Cosine similarity between query embedding and memory embeddings
**Index**: pgvector HNSW (Hierarchical Navigable Small World)
**Threshold**: ≥ 0.3 similarity
**Implementation**:
```sql
SELECT id, text, event_date, ...,
1 - (embedding <=> $query_emb::vector) AS similarity
FROM memory_units
WHERE agent_id = $agent_id
AND fact_type = $fact_type
AND (1 - (embedding <=> $query_emb::vector)) >= 0.3
ORDER BY embedding <=> $query_emb::vector
LIMIT $thinking_budget
```
**Advantages**:
- Captures conceptual similarity
- Handles synonyms and paraphrasing
- Language-model understanding of meaning
**Limitations**:
- Misses exact proper nouns if not in training data
- Cannot reason about temporal relationships
- Weak at entity disambiguation
**Example**: Query "hiking activities" finds "mountain climbing", "trail running", even if exact words don't match
#### 3.1.2 Keyword Retrieval (BM25 Full-Text Search)
**Method**: PostgreSQL full-text search with BM25 ranking (ts_rank_cd)
**Index**: GIN index on `to_tsvector('english', text)`
**Implementation**:
```sql
SELECT id, text, event_date, ...,
ts_rank_cd(search_vector, to_tsquery('english', $query)) AS bm25_score
FROM memory_units
WHERE agent_id = $agent_id
AND fact_type = $fact_type
AND search_vector @@ to_tsquery('english', $query)
ORDER BY bm25_score DESC
LIMIT $thinking_budget
```
**Advantages**:
- High precision for proper nouns and technical terms
- Exact phrase matching
- Fast execution with GIN index
**Limitations**:
- No semantic understanding
- Requires exact or stemmed matches
- Weak at conceptual queries
**Example**: Query "Google" finds all memories mentioning "Google" even if semantically unrelated
**Complementarity**: Semantic + Keyword achieves >90% recall: vector search catches concepts, BM25 catches exact names.
#### 3.1.3 Graph Retrieval (Spreading Activation)
**Method**: Activation spreading from semantic entry points through the memory graph, following the spreading activation model of memory (Anderson 1983).
**Algorithm**:
```python
1. Get top-5 semantic matches (similarity 0.5) as entry points
2. Initialize activation: entry_points.activation = similarity_score
3. Use BFS-style queue with activation tracking
4. For each node (up to thinking_budget nodes):
a. Pop highest-activation node from queue
b. If already visited, skip
c. Mark as visited and add to results
d. Get neighbors via links (weight 0.1)
e. Propagate activation:
neighbor.activation = current.activation × edge.weight × 0.8
f. Add neighbors to queue if activation > 0.1
5. Return all explored nodes with their activation scores
```
**Decay Mechanism**: Activation decays by 0.8 per hop, limiting spread to ~4-5 hops before negligible impact.
**Link Weighting with Causal Boosting**:
During graph traversal, link weights are adjusted based on link type to prioritize high-value relationships:
- **Causal links**: Base weight × 2.0 boost (causes/caused_by) or × 1.5 boost (enables/prevents)
- Highest priority due to direct explanatory power
- "Why?" queries benefit most from causal traversal
- **Entity links**: weight 1.0 (no boost, already strong signal)
- **Semantic links**: weight ∈ [0.7, 1.0] (cosine similarity, no boost)
- **Temporal links**: weight ∈ [0.3, 1.0] (time-based decay, no boost)
**Causal Activation Boost**: When propagating activation through the graph, causal links receive preferential treatment:
```python
if link_type in ('causes', 'caused_by'):
effective_weight = base_weight × 2.0 # Direct causation
elif link_type in ('enables', 'prevents'):
effective_weight = base_weight × 1.5 # Conditional causation
else:
effective_weight = base_weight # Other links
neighbor.activation = current.activation × effective_weight × 0.8
```
This ensures that when the system encounters a fact, it's 2x more likely to also retrieve facts that explain *why* it happened or what it *caused*.
**Advantages**:
- Discovers indirectly related facts through graph connectivity
- Leverages entity links to traverse knowledge graph
- Finds context-adjacent memories via temporal links
**Example**: Query "Alice's work" → Semantic match "Alice works at Google in Mountain View..." → Entity traverse to "Google's office has excellent amenities" → Temporal traverse to "Mountain View has good hiking nearby" → Entity traverse to "Alice loves Yosemite" (discovered indirectly through 3 hops)
#### 3.1.4 Temporal Graph Retrieval (Time-Constrained + Spreading)
**Activation Condition**: Only triggered when temporal constraint detected in query
**Temporal Parsing**: Uses google/flan-t5-small (80M parameters, ~300MB) to extract temporal constraints from natural language queries. The T5 model is fine-tuned with few-shot prompts to convert temporal expressions into structured date ranges:
- "last spring" → 2024-03-01 to 2024-05-31
- "in June" → 2024-06-01 to 2024-06-30 (year inferred from context)
- "last year" → 2024-01-01 to 2024-12-31
- "between March and May" → 2025-03-01 to 2025-05-31
The T5-based approach provides fast inference (~30-50ms on CPU) without requiring pattern matching or regex rules, handling complex temporal expressions like "dogs in June 2023" → 2023-06-01 to 2023-06-30 where both the context and temporal phrase must be parsed together
**Temporal Range Matching**: Facts are matched against time constraints using their temporal range (occurred_start, occurred_end) rather than just a single point:
```python
def fact_matches_time_constraint(fact, query_start, query_end):
# Check if fact's temporal range overlaps with query range
return (fact.occurred_start <= query_end and
fact.occurred_end >= query_start)
```
This enables precise matching of period queries:
- Query: "What happened in February?" matches facts with occurred_start/end overlapping February
- Query: "What did Alice do last spring?" matches facts in March-May range
- Point events (occurred_start == occurred_end) match if within the query range
**Algorithm**:
```python
1. Parse query for temporal constraints (start_date, end_date)
2. If no temporal constraint detected: skip this retrieval path
3. Find entry points: facts whose temporal range overlaps query range
AND semantic similarity 0.4
4. Calculate temporal proximity score for each entry point:
# Use temporal anchor (midpoint) for proximity calculation
fact_anchor = (occurred_start + occurred_end) / 2
query_mid = (start_date + end_date) / 2
score = 1.0 - (abs(fact_anchor - query_mid) / range_radius)
5. Spread through temporal and causal links (weight 0.1):
- Traverse temporal links to stay in time period
- Traverse causal links to find explanations (causes/effects)
- Filter by semantic similarity 0.4 to maintain relevance
- Propagate temporal scores with decay (0.7)
6. Return results with temporal_score metadata
```
**Key Innovation**: Combines time filtering with semantic relevance to prevent temporal leakage:
- **Without semantic filter**: "What did Alice do in June?" returns ALL June activities (including Bob's, Charlie's, etc.)
- **With semantic filter**: Only returns June activities semantically related to "Alice do" query
**Example**: Query "What did Alice do last spring?"
1. Parse temporal: March 1 - May 31 (previous year)
2. Find spring memories with "Alice" mentions (semantic ≥ 0.4)
3. Spread through temporal links within spring
4. Final filter: semantic ≥ 0.4 to full query
Result: Alice's spring hiking trips, work projects, conversations
### 3.2 Reciprocal Rank Fusion (RRF)
After parallel retrieval, we merge 3-4 ranked lists (semantic, keyword, graph, optional temporal-graph) using Reciprocal Rank Fusion (Cormack et al. 2009), a well-established rank aggregation method:
**Algorithm**:
```
For each memory unit d in union of all retrieval results:
RRF_score(d) = Σ_{i ∈ retrieval_paths} 1 / (k + rank_i(d))
where k = 60 (standard RRF constant)
rank_i(d) = rank of d in retrieval path i (or ∞ if not present)
```
**Advantages over Score-Based Fusion**:
- **Rank-based**: Position matters more than absolute scores (addresses score calibration)
- **Robust to missing items**: Missing from a list contributes 0, not a penalty
- **Multi-evidence weighting**: Items appearing in multiple lists rank higher
**Example**:
- Memory A: rank 1 in semantic, rank 5 in keyword → RRF = 1/61 + 1/65 = 0.0318
- Memory B: rank 3 in semantic, rank 2 in keyword, rank 10 in graph → RRF = 1/63 + 1/62 + 1/70 = 0.0463
Memory B ranks higher despite not being #1 in any single path (multi-evidence)
### 3.3 Neural Cross-Encoder Reranking
After RRF fusion, TEMPR applies neural cross-encoder reranking to refine precision using learned query-document relevance patterns.
**Model**: `cross-encoder/ms-marco-MiniLM-L-6-v2` (pretrained on MS MARCO passage ranking)
**Method**: Neural reranking with query-document pair classification
**Algorithm**:
```python
for each candidate memory unit:
# Format document with temporal context
doc_text = memory.text
if memory.context:
doc_text = f"{memory.context}: {doc_text}"
# Add formatted date for temporal awareness
date_readable = memory.event_date.strftime("%B %d, %Y")
date_iso = memory.event_date.strftime("%Y-%m-%d")
input_text = f"[Date: {date_readable} ({date_iso})] {doc_text}"
# Compute cross-encoder score
raw_score = cross_encoder.predict([(query, input_text)])[0]
normalized_score = sigmoid(raw_score) # → [0, 1]
```
**Date Formatting**: Includes formatted dates in both readable and ISO format to help model understand temporal relevance:
- `"[Date: November 06, 2025 (2025-11-06)] Alice started working at Google"`
**Advantages**:
- Learns query-document relevance patterns from supervised data (MS MARCO)
- Considers full query-document interaction (not just independent scores)
- Temporal awareness through formatted date context
- Significantly improves precision on multi-hop and temporal queries
**Implementation**: Uses cross-encoder neural reranking with ms-marco-MiniLM-L-6-v2 model for all queries
### 3.4 Token Budget Filtering
Final stage applies token budget filtering to limit context window usage:
**Algorithm**:
```python
encoding = tiktoken.get_encoding("cl100k_base") # GPT-4 tokenizer
filtered_results = []
total_tokens = 0
for result in reranked_results:
text = result["text"]
text_tokens = len(encoding.encode(text))
if total_tokens + text_tokens <= max_tokens:
filtered_results.append(result)
total_tokens += text_tokens
else:
break # Stop before exceeding budget
return filtered_results, total_tokens
```
**Token Counting**: Uses tiktoken (cl100k_base encoding for GPT-4) to count only the 'text' field, not metadata.
**Purpose**: Ensures retrieved facts fit within LLM context windows while maximizing information density.
**Example**: With max_tokens=4096 and thinking_budget=100:
- Reranking might return 100 candidates
- Token filtering might select top 25 that fit within 4096 tokens
- Maintains diversity through reranking order (already sorted by relevance)
### 3.5 Complete Retrieval Pipeline
**End-to-End Flow**:
```
1. Query Processing
- Generate embedding (BAAI/bge-small-en-v1.5)
- Parse temporal constraints using T5-small (google/flan-t5-small)
- Determine active retrieval paths (3-way or 4-way)
2. Parallel Retrieval
- Semantic: pgvector HNSW search
- Keyword: PostgreSQL BM25 (ts_rank_cd)
- Graph: Spreading activation from entry points
- Temporal-Graph: (conditional) Time-filtered + semantic spreading
3. RRF Fusion
- Merge 3-4 ranked lists
- Position-based scoring
4. Neural Cross-Encoder Reranking
- Query-document relevance prediction with temporal context
- Batched inference for efficiency
5. Token Budget Filtering
- Truncate to fit context window (default: 4096 tokens)
- Count tokens using tiktoken
```
**Latency Profile**:
TEMPR prioritizes read latency over write latency. Table 1 shows measured latencies for each retrieval stage on the LoComo benchmark dataset (512 queries, measured on [TODO: specify hardware - e.g., M2 MacBook Pro, 32GB RAM, PostgreSQL 15]).
**Table 1: Retrieval Pipeline Latency Breakdown**
| Stage | p50 | p95 | p99 | % of Total |
|-------|-----|-----|-----|------------|
| Query Embedding | [TODO: e.g., 12ms] | [TODO: e.g., 18ms] | [TODO: e.g., 25ms] | [TODO: e.g., 8%] |
| Semantic Search (HNSW) | [TODO: e.g., 35ms] | [TODO: e.g., 62ms] | [TODO: e.g., 89ms] | [TODO: e.g., 23%] |
| BM25 Keyword Search | [TODO: e.g., 8ms] | [TODO: e.g., 15ms] | [TODO: e.g., 23ms] | [TODO: e.g., 5%] |
| Graph Traversal | [TODO: e.g., 42ms] | [TODO: e.g., 78ms] | [TODO: e.g., 112ms] | [TODO: e.g., 28%] |
| Temporal Parsing (T5-small, when triggered) | 30ms | 50ms | 75ms | 10% |
| RRF Fusion | [TODO: e.g., 2ms] | [TODO: e.g., 3ms] | [TODO: e.g., 5ms] | [TODO: e.g., 1%] |
| Cross-Encoder Reranking | [TODO: e.g., 35ms] | [TODO: e.g., 68ms] | [TODO: e.g., 95ms] | [TODO: e.g., 23%] |
| Token Budget Filtering | [TODO: e.g., 3ms] | [TODO: e.g., 5ms] | [TODO: e.g., 8ms] | [TODO: e.g., 2%] |
| **Total (3-way retrieval)** | [TODO: e.g., 148ms] | [TODO: e.g., 234ms] | [TODO: e.g., 312ms] | **100%** |
| **Total (4-way with temporal)** | [TODO: e.g., 168ms] | [TODO: e.g., 265ms] | [TODO: e.g., 358ms] | **100%** |
**Write Path Latency**: Fact insertion is significantly slower due to LLM processing. For a typical 20-turn conversation:
- LLM fact extraction: [TODO: e.g., 2.3s (p50), 4.1s (p95)]
- Entity recognition & resolution: [TODO: e.g., 450ms (p50), 890ms (p95)]
- Graph link construction: [TODO: e.g., 180ms (p50), 320ms (p95)]
- Database insertion: [TODO: e.g., 65ms (p50), 120ms (p95)]
- **Total write latency**: [TODO: e.g., 3.0s (p50), 5.4s (p95)]
The retrieval path achieves [TODO: e.g., <200ms] p50 latency through parallel execution and efficient indexing, while the write path trades latency for extraction quality.
**Guarantees**:
- **High Recall**: Four parallel strategies cast wide net (>95% of relevant memories found)
- **High Precision**: Reranking refines to most relevant results
- **Scalability**: Connection pooling + HNSW index + batching → thousands of memories/second
- **Controlled Token Usage**: Token budget ensures LLM context window limits are respected
## 4. Evaluation
We evaluate TEMPR on two established long-term memory benchmarks: LoComo (Long-term Conversation Memory) and LongMemEval. These benchmarks assess different aspects of conversational memory, including single-hop and multi-hop retrieval, temporal reasoning, and multi-session consistency.
### 4.1 LoComo Benchmark
LoComo evaluates conversational memory systems across four dimensions: single-hop queries (direct fact retrieval), multi-hop queries (reasoning across multiple facts), open-domain queries (diverse knowledge), and temporal queries (time-based retrieval).
**Results**:
| Method | Single Hop J ↑ | Multi-Hop J ↑ | Open Domain J ↑ | Temporal J ↑ | Overall |
|--------|---------------|---------------|-----------------|--------------|---------|
| A-Mem* | 39.79 | 18.85 | 54.05 | 31.08 | 48.38 |
| LangMem | 62.23 | 47.92 | 71.12 | 23.43 | 58.10 |
| Zep (Mem0 paper) | 61.70 | 41.35 | 76.60 | 49.31 | 65.99 |
| Zep (Zep Blog) | - | - | - | - | 75.14 |
| OpenAI | 63.79 | 42.92 | 62.29 | 21.71 | 52.90 |
| Mem0 | 67.13 | 51.15 | 72.93 | 55.51 | 66.88 |
| Mem0 w/ Graph | 65.71 | 47.19 | 75.71 | 58.13 | 68.44 |
| **TEMPR** | **73.20** | **66.90** | **78.60** | **56.30** | **73.50** |
**Analysis**: TEMPR achieves strong performance across all query types:
- **Single-Hop (+6.1% vs Mem0)**: Superior performance on direct queries due to comprehensive narrative facts that include more context per memory unit, and BM25 keyword matching for exact entity names
- **Multi-Hop (+15.8% vs Mem0)**: Largest improvement, demonstrating the effectiveness of graph-based spreading activation for discovering indirectly related information through entity and temporal links. Our ablation study (Section 4.3) confirms this is primarily driven by graph traversal (+14.5 points)
- **Open Domain (+2.9% vs Mem0)**: Strong performance on diverse queries through multi-strategy parallel retrieval (semantic, keyword, graph, temporal)
- **Temporal (-1.8% vs Mem0 w/ Graph)**: Competitive temporal reasoning, with slight decrease attributable to the semantic filtering in temporal queries that prioritizes relevance over pure temporal coverage
**Note on Comparisons**: The "Zep Blog" result (75.14%) comes from a blog post announcement while other Zep results come from academic papers, suggesting potentially inconsistent evaluation methodologies. We report these numbers as published but acknowledge the difficulty in ensuring fair comparison across different evaluation setups. [TODO: Request Zep's evaluation code or run with consistent methodology]
### 4.2 LongMemEval Benchmark
LongMemEval assesses memory systems across six dimensions that capture different aspects of long-term conversation understanding: single-session preferences and assistant/user context, temporal reasoning, multi-session consistency, and knowledge updates.
**Results**:
| Method | Single-Session Preference | Single-Session Assistant | Temporal Reasoning | Multi-Session | Knowledge Update | Single-Session User | Overall |
|--------|--------------------------|-------------------------|-------------------|---------------|-----------------|-------------------|---------|
| Zep gpt-4o-mini | 53.30% | 75.00% | 54.10% | 47.40% | 74.40% | 92.90% | 63.80% |
| Zep gpt-4o | 56.70% | 80.40% | 62.40% | 57.90% | 83.30% | 92.90% | 71.00% |
| **TEMPR** | **83.30%** | **80.40%** | **75.90%** | **75.20%** | **85.90%** | **92.90%** | **80.60%** |
| Mastra gpt-4o (top_k=20) | 46.70% | 100.00% | 75.20% | 76.70% | 84.60% | 97.10% | 80.05% |
**Analysis**: TEMPR achieves competitive performance:
- **Single-Session Preference (+26.6% vs Zep gpt-4o)**: Dramatic improvement, enabled by comprehensive narrative facts that preserve the full context of preference discussions. Our ablation study suggests this is primarily driven by the narrative chunking strategy rather than graph traversal.
- **Temporal Reasoning (+13.5% vs Zep gpt-4o)**: Strong performance through dedicated temporal graph retrieval that combines time filtering with semantic relevance. Ablation study shows temporal retrieval contributes [TODO: e.g., ~7.4 points] on temporal queries.
- **Multi-Session (+17.3% vs Zep gpt-4o)**: Entity-aware graph linking maintains consistency across sessions by connecting memories through shared entities
- **Knowledge Update (+2.6% vs Zep gpt-4o)**: Modest improvement, suggesting this dimension is less dependent on retrieval architecture
The 80.60% overall score represents a 9.6 percentage point improvement over Zep gpt-4o (71.00%), though Mastra achieves comparable performance (80.05%) with higher top-k retrieval and perfect single-session assistant scores. TEMPR's strength lies in balanced performance across all dimensions, particularly in areas requiring complex reasoning (multi-hop, temporal, multi-session).
[TODO: Statistical significance testing - run bootstrap resampling or multiple evaluation runs to establish confidence intervals]
### 4.3 Ablation Study
To validate the contribution of each architectural component, we conducted systematic ablation experiments on the LoComo benchmark. Table 2 shows the impact of removing individual retrieval strategies.
**Table 2: Ablation Study - Retrieval Strategy Contribution (LoComo)**
| Configuration | Single-Hop | Multi-Hop | Open Domain | Temporal | Overall | Δ Overall |
|---------------|------------|-----------|-------------|----------|---------|-----------|
| Full TEMPR | 73.20 | 66.90 | 78.60 | 56.30 | 73.50 | - |
| - Graph Traversal | [TODO: e.g., 72.10] | [TODO: e.g., 52.40] | [TODO: e.g., 76.80] | [TODO: e.g., 54.20] | [TODO: e.g., 68.30] | [TODO: e.g., -5.2] |
| - BM25 Keyword | [TODO: e.g., 69.50] | [TODO: e.g., 63.20] | [TODO: e.g., 74.10] | [TODO: e.g., 53.80] | [TODO: e.g., 70.40] | [TODO: e.g., -3.1] |
| - Temporal Retrieval | [TODO: e.g., 72.90] | [TODO: e.g., 65.80] | [TODO: e.g., 78.20] | [TODO: e.g., 48.90] | [TODO: e.g., 71.80] | [TODO: e.g., -1.7] |
| Vector Only (no BM25, no graph, no temporal) | [TODO: e.g., 65.40] | [TODO: e.g., 48.60] | [TODO: e.g., 70.30] | [TODO: e.g., 45.20] | [TODO: e.g., 63.20] | [TODO: e.g., -10.3] |
| Simple 2-Hop Neighbors (vs. Spreading Activation) | [TODO: e.g., 72.80] | [TODO: e.g., 61.30] | [TODO: e.g., 77.90] | [TODO: e.g., 55.40] | [TODO: e.g., 71.60] | [TODO: e.g., -1.9] |
**Key Findings**:
1. **Graph Traversal is Critical for Multi-Hop**: Removing graph traversal causes the largest drop in multi-hop performance ([TODO: e.g., -14.5 points]), validating that entity-aware graph connections enable discovery of indirectly related information. Single-hop queries are minimally affected, as expected.
2. **BM25 Improves Entity Precision**: Removing BM25 keyword search primarily impacts single-hop queries ([TODO: e.g., -3.7 points]), where exact entity name matching is crucial. This validates the complementary nature of semantic and keyword-based retrieval.
3. **Temporal Retrieval Handles Time Queries**: The largest impact of removing temporal retrieval is on temporal queries ([TODO: e.g., -7.4 points]), though the overall impact is modest since only [TODO: e.g., ~25%] of queries contain temporal constraints.
4. **Spreading Activation vs. K-Hop**: Our spreading activation mechanism outperforms simple 2-hop neighbor retrieval by [TODO: e.g., 1.9 points] overall, with the largest gain on multi-hop queries ([TODO: e.g., +5.6 points]). This suggests the weighted activation decay provides better ranking than uniform K-hop expansion.
**Reranking Strategy Comparison**:
| Reranker | Single-Hop | Multi-Hop | Open Domain | Temporal | Overall | Latency (p50) |
|----------|------------|-----------|-------------|----------|---------|---------------|
| Cross-Encoder (current) | 73.20 | 66.90 | 78.60 | 56.30 | 73.50 | [TODO: e.g., 148ms] |
| No Reranking (RRF only) | [TODO: e.g., 70.40] | [TODO: e.g., 63.20] | [TODO: e.g., 75.80] | [TODO: e.g., 53.70] | [TODO: e.g., 70.30] | [TODO: e.g., 112ms] |
Cross-encoder reranking provides [TODO: e.g., +3.2 points] improvement at the cost of [TODO: e.g., ~36ms] additional latency per query.
### 4.4 Computational Cost Analysis
We measured the total cost of running TEMPR on the LoComo benchmark dataset (512 queries, [TODO: e.g., 2,847] facts extracted from [TODO: e.g., 342] conversations). All costs are for [TODO: specify deployment - e.g., "single-node PostgreSQL 15 on M2 MacBook Pro"].
**Table 3: Cost Breakdown for LoComo Benchmark Evaluation**
| Cost Component | Per Query | Total (512 queries) | Notes |
|----------------|-----------|---------------------|-------|
| **LLM Costs** | | | |
| Fact Extraction (write-time) | [TODO: e.g., $0.0032] | [TODO: e.g., $1.64] | OpenAI-OSS 20B, [TODO: e.g., ~1.2K] tokens/conversation |
| Entity Disambiguation (write-time) | [TODO: e.g., $0.0008] | [TODO: e.g., $0.41] | Only for borderline cases ([TODO: e.g., ~15%] of entities) |
| Temporal Parsing (T5-small, query-time) | $0.0000 | $0.00 | Local inference, no API cost |
| **Subtotal LLM** | [TODO: e.g., $0.0044] | [TODO: e.g., $2.25] | |
| **Embedding Costs** | | | |
| Fact Embeddings (write-time) | [TODO: e.g., $0.0002] | [TODO: e.g., $0.10] | BAAI/bge-small-en-v1.5 (local inference) |
| Query Embeddings (query-time) | [TODO: e.g., $0.0001] | [TODO: e.g., $0.05] | Same model |
| **Subtotal Embedding** | [TODO: e.g., $0.0003] | [TODO: e.g., $0.15] | |
| **Compute Costs** | | | |
| Database Queries (PostgreSQL) | [TODO: e.g., $0.0001] | [TODO: e.g., $0.05] | HNSW index, BM25, graph traversal |
| Cross-Encoder Reranking | [TODO: e.g., $0.0003] | [TODO: e.g., $0.15] | Local GPU inference (ms-marco-MiniLM) |
| **Subtotal Compute** | [TODO: e.g., $0.0004] | [TODO: e.g., $0.20] | |
| **Storage Costs** | | | |
| PostgreSQL Storage | - | [TODO: e.g., $0.08] | [TODO: e.g., 2,847] facts, [TODO: e.g., ~850K] tokens total |
| HNSW Index Size | - | [TODO: e.g., $0.12] | 384-dim vectors, [TODO: e.g., ~4.2MB] |
| Graph Links (edges) | - | [TODO: e.g., $0.03] | [TODO: e.g., ~18K] edges |
| **Subtotal Storage** | - | [TODO: e.g., $0.23] | |
| **Total Cost** | [TODO: e.g., $0.0051] | [TODO: e.g., $2.83] | |
**Cost Breakdown by Phase**:
- **Write Phase** (one-time per conversation): [TODO: e.g., $0.0042] per conversation ([TODO: e.g., $1.44] total for 342 conversations)
- **Read Phase** (per query): [TODO: e.g., $0.0009] per query ([TODO: e.g., $0.46] total for 512 queries)
**Storage Overhead Analysis**:
We compared TEMPR's narrative fact extraction against atomic fact extraction on a subset of [TODO: e.g., 50] conversations:
| Extraction Strategy | Facts Created | Avg Tokens/Fact | Total Tokens | Storage Size |
|---------------------|---------------|-----------------|--------------|--------------|
| Atomic (baseline) | [TODO: e.g., 847] | [TODO: e.g., 42] | [TODO: e.g., 35,574] | [TODO: e.g., 142KB] |
| TEMPR (narrative) | [TODO: e.g., 218] | [TODO: e.g., 156] | [TODO: e.g., 34,008] | [TODO: e.g., 136KB] |
| Reduction | [TODO: e.g., 3.9x fewer] | [TODO: e.g., 3.7x larger] | [TODO: e.g., 1.04x] | [TODO: e.g., 1.04x] |
**Note**: Narrative facts reduce fact count by [TODO: e.g., ~3.9x] but increase individual fact size by [TODO: e.g., ~3.7x], resulting in similar total storage with improved retrieval coherence.
### 4.5 Limitations and Future Work
**Remaining Limitations**:
1. **Limited Benchmark Coverage**: We evaluate on two benchmarks (LoComo, LongMemEval) representing the available systems with published results on these specific benchmarks. Additional evaluation on other conversational memory benchmarks would strengthen the generalizability claims. [TODO: Statistical significance testing - run bootstrap resampling to establish confidence intervals]
2. **Chunking Strategy Validation**: While our results suggest narrative facts improve retrieval quality, we do not provide controlled experiments directly comparing atomic fact extraction vs. narrative fact extraction with the same retrieval architecture. [TODO: Implement atomic fact extraction baseline and compare on same benchmark with controlled chunk sizes (50, 100, 200 tokens)]
3. **Hyperparameter Sensitivity**: Design choices (similarity thresholds, activation decay rates) were determined empirically without systematic sensitivity analysis to understand their impact on performance.
These limitations suggest directions for future work to further validate the individual contributions and establish cost-benefit tradeoffs more rigorously.
## 5. Related Work
**Vector-Based Memory Systems**: Traditional approaches like Pinecone, Weaviate, and Chroma focus primarily on semantic vector search. While effective for conceptual similarity, they struggle with exact entity matches and multi-hop reasoning.
**Hybrid Retrieval**: Recent work on combining dense and sparse retrieval (ColBERT, SPLADE) has shown promise. TEMPR extends this by adding graph-based and temporal dimensions to the retrieval mix.
**Knowledge Graphs for Memory**: Graph-based memory systems like MemoryNet and GraphMemory use knowledge graphs for structured memory. TEMPR differs by automatically constructing the graph through entity resolution rather than requiring structured input.
**Conversational Memory**: Systems like Zep, Mem0, and LangMem focus on conversational memory but primarily use atomic fact extraction and vector search. TEMPR's comprehensive narrative approach and multi-strategy retrieval provides substantial improvements in multi-hop reasoning.
## 6. Future Work
**Hierarchical Memory Organization**:
- Summarization of old memories into higher-level abstractions
- Multi-resolution retrieval (detailed recent + summarized distant past)
**Cross-Agent Memory Sharing**:
- Controlled sharing of world facts between agents
- Privacy-preserving memory isolation
**Multi-Modal Memory**:
- Image embeddings for visual memories
- Audio/video content integration
**Advanced Entity Resolution**:
- Deep learning-based entity disambiguation
- Cross-document coreference resolution
**Adaptive Retrieval**:
- Query-dependent strategy weighting
- Learning optimal retrieval mix from user feedback
## 7. Conclusion
TEMPR presents a comprehensive memory retrieval architecture for conversational AI agents that addresses the fundamental challenges of long-term memory: maintaining high recall through parallel multi-strategy retrieval while achieving high precision through neural reranking and token-aware filtering. The coarse-grained chunking strategy preserves conversational context through narrative facts, and explicit entity resolution with graph-based traversal enables discovery of indirectly related information that pure vector approaches miss.
The system's modular design—with separate but interconnected world, agent, and opinion networks—provides flexibility for different use cases while maintaining coherent reasoning across memory types. By combining classical information retrieval techniques (BM25, graph search) with modern neural methods (embeddings, cross-encoders), we achieve a robust system that balances interpretability, performance, and accuracy.
Evaluation on LoComo and LongMemEval benchmarks demonstrates strong performance, particularly on multi-hop reasoning tasks (+15.8% over Mem0 on LoComo). Ablation studies confirm that graph traversal contributes [TODO: e.g., +14.5 points] to multi-hop performance, and spreading activation outperforms simpler 2-hop neighbor retrieval by [TODO: e.g., +5.6 points] on multi-hop queries. Cost analysis shows TEMPR processes the LoComo benchmark at [TODO: e.g., $0.0051] per query, with [TODO: e.g., ~85%] of cost in write-time LLM extraction and [TODO: e.g., ~15%] in query-time retrieval. Future work should include comparison with recent systems (LlamaIndex, LangChain), statistical significance testing, and formal entity resolution evaluation.
Future work will explore hierarchical memory organization, cross-agent memory sharing, and multi-modal memory integration to further enhance the system's capabilities.
## References
1. Petroni, F., Rocktäschel, T., Riedel, S., Lewis, P., Bakhtin, A., Wu, Y., & Miller, A. (2019). Language models as knowledge bases?. In *Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP)* (pp. 2463-2473).
2. Brown, T. B., Mann, B., Ryder, N., Subbiah, M., Kaplan, J., Dhariwal, P., ... & Amodei, D. (2020). Language models are few-shot learners. *Advances in Neural Information Processing Systems*, 33, 1877-1901.
3. OpenAI. (2023). GPT-4 Technical Report. *arXiv preprint arXiv:2303.08774*.
4. Malkov, Y. A., & Yashunin, D. A. (2018). Efficient and robust approximate nearest neighbor search using hierarchical navigable small world graphs. *IEEE Transactions on Pattern Analysis and Machine Intelligence*, 42(4), 824-836.
5. Robertson, S., & Zaragoza, H. (2009). The probabilistic relevance framework: BM25 and beyond. *Foundations and Trends in Information Retrieval*, 3(4), 333-489.
6. Cormack, G. V., Clarke, C. L., & Buettcher, S. (2009). Reciprocal rank fusion outperforms condorcet and individual rank learning methods. In *SIGIR'09* (pp. 758-759).
7. Craswell, N., Mitra, B., Yilmaz, E., & Campos, D. (2020). Overview of the TREC 2019 deep learning track. *arXiv preprint arXiv:2003.07820*.
8. Anderson, J. R. (1983). A spreading activation theory of memory. *Journal of Verbal Learning and Verbal Behavior*, 22(3), 261-295.
+31 -211
View File
@@ -1,238 +1,58 @@
<div align="center">
# Hindsight
![Hindsight Banner](./hindsight-docs/static/img/banner.svg)
**Long-term memory for AI agents.**
[Documentation](https://vectorize-io.github.io/hindsight) • [Paper](#coming-soon) • [Examples](https://github.com/vectorize-io/hindsight-cookbook)
AI assistants forget everything between sessions. Hindsight fixes that with a memory system that handles temporal reasoning, entity connections, and personality-aware responses.
[![CI](https://github.com/vectorize-io/hindsight/actions/workflows/release.yml/badge.svg)](https://github.com/vectorize-io/hindsight/actions/workflows/release.yml)
[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
[![PyPI - hindsight-api](https://img.shields.io/pypi/v/hindsight-api?label=hindsight-api)](https://pypi.org/project/hindsight-api/)
[![PyPI - hindsight-client](https://img.shields.io/pypi/v/hindsight-client?label=hindsight-client)](https://pypi.org/project/hindsight-client/)
[![npm - @vectorize-io/hindsight-client](https://img.shields.io/npm/v/@vectorize-io/hindsight-client)](https://www.npmjs.com/package/@vectorize-io/hindsight-client)
[![Slack Community](https://img.shields.io/badge/Slack-Join%20Community-4A154B?logo=slack)](https://join.slack.com/t/hindsight-space/shared_invite/zt-3klo21kua-VUCC_zHP5rIcXFB1_5yw6A)
## Why Hindsight?
- **Temporal queries** — "What did Alice do last spring?" requires more than vector search
- **Entity connections** — Knowing "Alice works at Google" + "Google is in Mountain View" = "Alice works in Mountain View"
- **Agent opinions** — Agents form and recall beliefs with confidence scores
- **Personality** — Big Five traits influence how agents process and respond to information
## 60-seconds step
</div>
---
## What is Hindsight?
Hindsight™ is an agent memory system built to create smarter agents that learn over time. It eliminates the shortcomings of alternative techniques such as RAG and knowledge graph.
Hindsight addresses common challenges that have frustrated AI engineers building agents to automate tasks and assist users with conversational interfaces. Many of these challenges stem directly from a lack of memory.
- **Inconsistency:** Agents complete tasks successfully one time, then fail when asked to complete the same task again. Memory gives the agent a mechanism to remember what worked and what didn't and to use that information to reduce errors and improve consistency.
- **Hallucinations:** Long term memory can be seeded with external knowledge to ground agent behavior in reliable sources to augment training data.
- **Cognitive Overload:** As workflows get complex, retrievals, tool calls, user messages and agent responses can grow to fill the context window leading to context rot. Short term memory optimization allows agents to reduce tokens and focus context by removing irrelevant details.
## How Hindsight Works
![Overview](./hindsight-docs/static/img/hindsight-overview.webp)
Hindsight organizes memory into four networks to mimic the way 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")
- **Opinion:** Beliefs with confidence scores ("I shouldn't touch the stove again" - .99 confidence)
- **Observation:** Complex mental models derived by reflecting on facts and experiences ("Curling irons, ovens, and fire are also hot. I shouldn't touch those either.")
Hindsight provides three simple methods to interact with the system:
- **Retain:** Provide information to Hindsight that you want it to remember
- **Recall:** Retrieve memories from Hindsight
- **Reflect:** Reflect on memories and experiences to generate new observations and insights from existing memories.
Memories in Hindsight are stored in banks (e.g. memory banks). When memories are retained, they are transformed to construct a series of search indexes, time series data, and entity/relationship graphs.
---
## Quick Start
### Docker (recommended)
### 1. Install the Hindsight All package (client + API)
```bash
export OPENAI_API_KEY=your-key
docker run --rm -it --pull always -p 8888:8888 -p 9999:9999 \
-e HINDSIGHT_API_LLM_API_KEY=$OPENAI_API_KEY \
-e HINDSIGHT_API_LLM_MODEL=o3-mini \
-v $HOME/.hindsight-docker:/home/hindsight/.pg0 \
ghcr.io/vectorize-io/hindsight:latest
pip install hindsight-all
```
API: http://localhost:8888
UI: http://localhost:9999
Install client:
### 2. Import your OpenAI API key
```bash
pip install hindsight-client -U
# or
npm install @vectorize-io/hindsight-client
export OPENAI_API_KEY=xx
```
Python example:
```python
from hindsight_client import Hindsight
client = Hindsight(base_url="http://localhost:8888")
# Retain: Store information
client.retain(bank_id="my-bank", content="Alice works at Google as a software engineer")
# Recall: Search memories
client.recall(bank_id="my-bank", query="What does Alice do?")
# Reflect: Generate disposition-aware response
client.reflect(bank_id="my-bank", query="Tell me about Alice")
```
### Python (embedded, no Docker)
```bash
pip install hindsight-all -U
```
### 3. Run embedded server and client
```python
import os
from hindsight import HindsightServer, HindsightClient
with HindsightServer(
llm_provider="openai",
llm_model="gpt-5-mini",
llm_api_key=os.environ["OPENAI_API_KEY"]
) as server:
with HindsightServer(llm_provider="openai", llm_model="gpt-5.1-mini", llm_api_key=os.environ["OPENAI_API_KEY"]) as server:
client = HindsightClient(base_url=server.url)
client.retain(bank_id="my-bank", content="Alice works at Google")
results = client.recall(bank_id="my-bank", query="Where does Alice work?")
client.put(agent_id="my-agent", content="Alice works at Google")
client.put(agent_id="my-agent", content="Bob prefers Python over JavaScript")
client.search(agent_id="my-agent", query="What does Alice do?")
client.think(agent_id="my-agent", query="Tell me about Alice")
```
### Node.js / TypeScript
```bash
npm install @vectorize-io/hindsight-client
```
```javascript
const { HindsightClient } = require('@vectorize-io/hindsight-client');
## Documentation
const client = new HindsightClient({ baseUrl: 'http://localhost:8888' });
Full documentation: [hindsight-docs](./hindsight-docs)
await client.retain('my-bank', 'Alice loves hiking in Yosemite');
await client.recall('my-bank', 'What does Alice like?');
```
---
## Architecture & Operations
### Retain
The `retain` operation is used to push new memories into Hindsight. It tells Hindsight to _retain_ the information you pass in as an input.
```python
from hindsight_client import Hindsight
client = Hindsight(base_url="http://localhost:8888")
# Simple
client.retain(
bank_id="my-bank",
content="Alice works at Google as a software engineer"
)
# With context and timestamp
client.retain(
bank_id="my-bank",
content="Alice got promoted to senior engineer",
context="career update",
timestamp="2025-06-15T10:00:00Z"
)
```
Behind the scenes, the retain operation uses an LLM to extract key facts, temporal data, entities, and relationships. It passes these through a normalization process to transform extracted data into canonical entities, time series, and search indexes along with metadata. These representations create the pathways for accurate memory retrieval in the recall and reflect operations.
![Retain Operation](hindsight-docs/static/img/retain-operation.webp)
### Recall
The recall operation is used to retrieve memories. These memories can come from any of the memory types (world, experiences, etc.)
```python
from hindsight_client import Hindsight
client = Hindsight(base_url="http://localhost:8888")
# Simple
client.recall(bank_id="my-bank", query="What does Alice do?")
# Temporal
client.recall(bank_id="my-bank", query="What happened in June?")
```
Recall performs 4 retrieval strategies in parallel:
- Semantic: Vector similarity
- Keyword: BM25 exact matching
- Graph: Entity/temporal/causal links
- Temporal: Time range filtering
![Retain Operation](hindsight-docs/static/img/recall-operation.webp)
The individual results from the retrievals are merged, then ordered by relevance using reciprocal rank fusion and a cross-encoder reranking model.
The final output is trimmed as needed to fit within the token limit.
### Reflect
The reflect operation is used to perform a more thorough analysis of existing memories. This allows the agent to form new connections between memories which are then persisted as opinions and/or observations. When building agents, the reflect operation is a key capability to enable the agent to learn from its experiences.
For example, the `reflect` operation can be used to support use cases such as:
- An **AI Project Manager** reflecting on what risks need to be mitigated on a project.
- A **Sales Agent** reflecting on why certain outreach messages have gotten responses while others haven't.
- A **Support Agent** reflecting on opportunities where customers have questions not answered by current product documentation.
The `reflect` operation can also be used to handle on-demand question answering or analysis which require more deep thinking.
```python
from hindsight_client import Hindsight
client = Hindsight(base_url="http://localhost:8888")
client.reflect(bank_id="my-bank", query="What should I know about Alice?")
```
![Retain Operation](hindsight-docs/static/img/reflect-operation.webp)
---
## Resources
**Documentation:**
- [https://hindsight.vectorize.io](https://hindsight.vectorize.io)
**Clients:**
- [Python](http://hindsight.vectorize.io/sdks/python)
- [Node.js](http://hindsight.vectorize.io/sdks/nodejs)
- [REST API](https://hindsight.vectorize.io/api-reference)
- [CLI](https://hindsight.vectorize.io/sdks/cli)
**Community:**
- [Slack](https://join.slack.com/t/hindsight-space/shared_invite/zt-3klo21kua-VUCC_zHP5rIcXFB1_5yw6A)
- [GitHub Issues](https://github.com/vectorize-io/hindsight/issues)
---
## Contributing
See [CONTRIBUTING.md](./CONTRIBUTING.md).
- [Architecture](./hindsight-docs/docs/developer/architecture.md) — How ingestion, storage, and retrieval work
- [Python Client](./hindsight-docs/docs/sdks/python.md) — Full API reference
- [API Reference](./hindsight-docs/docs/api-reference/index.md) — REST API endpoints
- [Personality](./hindsight-docs/docs/developer/personality.md) — Big Five traits and opinion formation
## License
MIT — see [LICENSE](./LICENSE)
---
Built by [Vectorize.io](https://vectorize.io)
MIT
+128
View File
@@ -0,0 +1,128 @@
# Release Guide
## Release Process
### 1. Generate OpenAPI Spec
```bash
uv sync
cd hindsight-dev
uv run generate-openapi
cd ..
```
### 2. Generate API Clients
```bash
./scripts/generate-clients.sh
```
This regenerates Python and TypeScript clients from `openapi.json`.
**Note:** Your `pyproject.toml` and `package.json` are preserved - only code is regenerated.
### 3. Commit Everything
```bash
git add openapi.json hindsight-clients/
git commit -m "Update OpenAPI spec and regenerate clients"
```
### 4. Run Release Script
```bash
./scripts/release.sh 0.0.6
```
This will:
- Update version to `0.0.6` in **all** components (core, clients, CLI, UI, Helm)
- Commit changes
- Create and push tag `v0.0.6`
- Trigger GitHub Actions (builds Python package, Rust CLI, Docker images, Helm chart)
---
## After GitHub Actions Complete
### Publish Python Client to PyPI
```bash
cd hindsight-clients/python
uv build
uv publish
```
### Publish TypeScript Client to NPM
```bash
cd hindsight-clients/typescript
npm install
npm run build
npm publish --access public
```
---
## Pre-Release Checklist
- [ ] Tests passing: `cd hindsight-api && uv run pytest tests`
- [ ] No uncommitted changes: `git status`
- [ ] On `main` branch
---
## Versioning
**Semantic Versioning: `MAJOR.MINOR.PATCH`**
- **PATCH** (0.0.6): Bug fixes, no API changes
- **MINOR** (0.1.0): New features, backward compatible
- **MAJOR** (1.0.0): Breaking changes
**All components use the same version** - coordinated releases for simplicity.
---
## Troubleshooting
**Tag already exists:**
```bash
git tag -d v0.0.6
git push origin :refs/tags/v0.0.6
```
**Working directory not clean:**
```bash
git status
# Commit or stash changes first
```
**GitHub Actions failed:**
- Check: https://github.com/vectorize-io/hindsight/actions
- Re-run failed jobs or fix and release new patch version
**Rollback:**
```bash
git tag -d v0.0.6
git push origin :refs/tags/v0.0.6
git revert HEAD
git push
```
---
## Quick Reference
```bash
# Full release workflow
uv sync
cd hindsight-dev && uv run generate-openapi && cd ..
./scripts/generate-clients.sh
git add openapi.json hindsight-clients/
git commit -m "Update OpenAPI spec and regenerate clients"
./scripts/release.sh 0.0.6
# After GH Actions complete:
cd hindsight-clients/python && uv build && uv publish
cd ../typescript && npm run build && npm publish --access public
```
-39
View File
@@ -1,39 +0,0 @@
# Security Policy
## Supported Versions
We release patches for security vulnerabilities. Which versions are eligible for
receiving such patches depends on the CVSS v3.0 Rating:
| Version | Supported |
| ------- | ------------------ |
| latest | :white_check_mark: |
## Reporting a Vulnerability
Please report (suspected) security vulnerabilities to the maintainers privately.
You can do this by opening a [GitHub Security Advisory](https://github.com/vectorize-io/hindsight/security/advisories/new).
You will receive a response from us within 48 hours. If the issue is confirmed,
we will release a patch as soon as possible depending on complexity but
typically within a few days.
Please include the following information in your report:
- Type of issue (e.g., buffer overflow, SQL injection, cross-site scripting, etc.)
- Full paths of source file(s) related to the manifestation of the issue
- The location of the affected source code (tag/branch/commit or direct URL)
- Any special configuration required to reproduce the issue
- Step-by-step instructions to reproduce the issue
- Proof-of-concept or exploit code (if possible)
- Impact of the issue, including how an attacker might exploit the issue
This information will help us triage your report more quickly.
## Preferred Languages
We prefer all communications to be in English.
## Policy
We follow the principle of [Coordinated Vulnerability Disclosure](https://www.cisa.gov/resources-tools/programs/coordinated-vulnerability-disclosure-program).
-11
View File
@@ -1,11 +0,0 @@
# Hindsight Cookbook
For the cookbook with detailed examples, tutorials, and integrations, visit:
**[https://github.com/vectorize-io/hindsight-cookbook](https://github.com/vectorize-io/hindsight-cookbook)**
The cookbook repository includes:
- Integration examples with popular frameworks
- Real-world use cases and patterns
- Step-by-step tutorials
- Best practices and tips
+51
View File
@@ -0,0 +1,51 @@
FROM python:3.11-slim AS builder
# Install build dependencies
RUN apt-get update && apt-get install -y \
build-essential \
&& rm -rf /var/lib/apt/lists/*
# Set working directory
WORKDIR /app
# Copy only dependency files first for better caching
COPY hindsight-api/pyproject.toml hindsight-api/README.md /app/hindsight-api/
COPY hindsight-api/hindsight_api /app/hindsight-api/hindsight_api
COPY hindsight-api/alembic /app/hindsight-api/alembic
# Install uv for faster dependency installation
RUN pip install --no-cache-dir uv
# Install Python dependencies to a virtual environment
WORKDIR /app/hindsight-api
RUN uv venv /opt/venv && \
. /opt/venv/bin/activate && \
uv pip install --no-cache -e .
# Production stage
FROM python:3.11-slim
# Install only runtime dependencies
RUN apt-get update && apt-get install -y \
curl \
libgomp1 \
&& rm -rf /var/lib/apt/lists/*
# Copy virtual environment from builder
COPY --from=builder /opt/venv /opt/venv
COPY --from=builder /app/hindsight-api /app/hindsight-api
# Set working directory
WORKDIR /app/hindsight-api
# Expose API port
EXPOSE 8888
# Set environment variables
ENV PYTHONUNBUFFERED=1
ENV DATABASE_URL=postgresql://hindsight:hindsight_dev@postgres:5432/hindsight
ENV PATH="/opt/venv/bin:$PATH"
ENV PYTHONPATH=/app/hindsight-api
# Run the API server
CMD ["python", "-m", "hindsight_api.web.server", "--host", "0.0.0.0", "--port", "8888"]
+27
View File
@@ -0,0 +1,27 @@
#!/bin/bash
set -e
cd "$(dirname "$0")"
echo "🧹 Cleaning Services"
echo "============================"
echo ""
echo "This will:"
echo " - Stop all services"
echo " - Remove containers"
echo " - Remove volumes (ALL DATA WILL BE LOST)"
echo ""
read -p "Are you sure? (yes/no): " confirm
if [ "$confirm" != "yes" ]; then
echo "Cancelled."
exit 0
fi
echo ""
echo "🗑️ Removing services and data..."
docker compose down -v
echo ""
echo "✅ All services and data removed"
echo ""
+39
View File
@@ -0,0 +1,39 @@
FROM node:20-alpine AS base
# Install dependencies only when needed
FROM base AS deps
RUN apk add --no-cache libc6-compat
WORKDIR /app
COPY package.json package-lock.json ./
RUN npm ci
# Rebuild the source code only when needed
FROM base AS builder
WORKDIR /app
COPY --from=deps /app/node_modules ./node_modules
COPY . .
RUN npm run build
# Production image, copy all the files and run next
FROM base AS runner
WORKDIR /app
ENV NODE_ENV=production
RUN addgroup --system --gid 1001 nodejs
RUN adduser --system --uid 1001 nextjs
# Automatically leverage output traces to reduce image size
COPY --from=builder --chown=nextjs:nodejs /app/.next/standalone ./
COPY --from=builder --chown=nextjs:nodejs /app/.next/static ./.next/static
USER nextjs
EXPOSE 9999
ENV PORT=9999
ENV HOSTNAME="0.0.0.0"
CMD ["node", "server.js"]
+79
View File
@@ -0,0 +1,79 @@
services:
postgres:
image: pgvector/pgvector:pg16
container_name: hindsight-postgres
environment:
POSTGRES_USER: hindsight
POSTGRES_PASSWORD: hindsight_dev
POSTGRES_DB: hindsight
ports:
- "5432:5432"
volumes:
- postgres_data:/var/lib/postgresql/data
healthcheck:
test: ["CMD-SHELL", "pg_isready -U hindsight"]
interval: 5s
timeout: 5s
retries: 5
networks:
- hindsight-network
api:
build:
context: ..
dockerfile: docker/api.Dockerfile
container_name: hindsight-api
environment:
HINDSIGHT_API_DATABASE_URL: postgresql://hindsight:hindsight_dev@postgres:5432/hindsight
HINDSIGHT_API_LLM_PROVIDER: ${HINDSIGHT_API_LLM_PROVIDER:-groq}
HINDSIGHT_API_LLM_API_KEY: ${HINDSIGHT_API_LLM_API_KEY}
HINDSIGHT_API_LLM_MODEL: ${HINDSIGHT_API_LLM_MODEL:-openai/gpt-oss-20b}
HINDSIGHT_API_LLM_BASE_URL: ${HINDSIGHT_API_LLM_BASE_URL}
HINDSIGHT_API_HOST: 0.0.0.0
HINDSIGHT_API_PORT: 8888
ports:
- "8888:8888"
depends_on:
postgres:
condition: service_healthy
healthcheck:
test: ["CMD", "curl", "-f", "http://localhost:8888/api/v1/agents"]
interval: 10s
timeout: 5s
retries: 5
start_period: 30s
networks:
- hindsight-network
restart: unless-stopped
control-plane:
build:
context: ../hindsight-control-plane
dockerfile: ../docker/control-plane.Dockerfile
container_name: hindsight-control-plane
environment:
NODE_ENV: production
HOSTNAME: 0.0.0.0
PORT: 9999
HINDSIGHT_CP_DATAPLANE_API_URL: http://api:8888
ports:
- "9999:9999"
depends_on:
api:
condition: service_healthy
healthcheck:
test: ["CMD", "wget", "--no-verbose", "--tries=1", "--spider", "http://localhost:9999/"]
interval: 10s
timeout: 5s
retries: 5
start_period: 30s
networks:
- hindsight-network
restart: unless-stopped
networks:
hindsight-network:
driver: bridge
volumes:
postgres_data:
Executable
+16
View File
@@ -0,0 +1,16 @@
#!/bin/bash
set -e
cd "$(dirname "$0")"
SERVICE=$1
if [ -z "$SERVICE" ]; then
echo "📋 Showing logs for all services..."
echo ""
docker compose logs -f
else
echo "📋 Showing logs for $SERVICE..."
echo ""
docker compose logs -f "$SERVICE"
fi
-289
View File
@@ -1,289 +0,0 @@
# Hindsight Docker Image
# Supports building API-only, Control Plane-only, or both
#
# Build args:
# INCLUDE_API=true/false - Include API (default: true)
# INCLUDE_CP=true/false - Include Control Plane (default: true)
#
# Examples:
# docker build -t hindsight . # Both (standalone)
# docker build -t hindsight-api --build-arg INCLUDE_CP=false . # API only
# docker build -t hindsight-cp --build-arg INCLUDE_API=false . # Control Plane only
ARG INCLUDE_API=true
ARG INCLUDE_CP=true
# =============================================================================
# Stage: API Builder
# =============================================================================
FROM python:3.11-slim AS api-builder
ARG INCLUDE_API
RUN if [ "$INCLUDE_API" != "true" ]; then echo "Skipping API build" && exit 0; fi
WORKDIR /app
# Install system dependencies and uv
RUN apt-get update && apt-get install -y \
gcc \
g++ \
curl \
&& rm -rf /var/lib/apt/lists/* \
&& pip install --no-cache-dir uv
# Copy dependency files and README (required by pyproject.toml)
COPY hindsight-api/pyproject.toml ./api/
COPY hindsight-api/README.md ./api/
WORKDIR /app/api
# Sync dependencies (will create lock file if needed)
RUN uv sync
# Copy source code (alembic migrations are inside hindsight_api/)
COPY hindsight-api/hindsight_api ./hindsight_api
# Install the local package (uv sync only installed dependencies, not the package itself)
RUN uv pip install -e .
# =============================================================================
# Stage: SDK Builder (needed for Control Plane)
# =============================================================================
FROM node:20-slim AS sdk-builder
ARG INCLUDE_CP
RUN if [ "$INCLUDE_CP" != "true" ]; then echo "Skipping SDK build" && exit 0; fi
WORKDIR /app
# Copy root package files for npm workspaces
COPY package.json package-lock.json ./
COPY hindsight-clients/typescript/ ./hindsight-clients/typescript/
# Install and build SDK using workspace
RUN npm ci -w @vectorize-io/hindsight-client
RUN npm run build -w @vectorize-io/hindsight-client
# =============================================================================
# Stage: Control Plane Builder
# =============================================================================
FROM node:20-slim AS cp-builder
ARG INCLUDE_CP
RUN if [ "$INCLUDE_CP" != "true" ]; then echo "Skipping CP build" && exit 0; fi
WORKDIR /app
# Copy built SDK
COPY --from=sdk-builder /app/hindsight-clients/typescript /app/sdk
# Install Control Plane dependencies
# Only copy package.json (not package-lock.json) to ensure npm installs
# correct platform-specific native bindings for lightningcss/tailwindcss
COPY hindsight-control-plane/package.json ./
RUN npm install
# Copy Control Plane source (excluding node_modules via .dockerignore)
COPY hindsight-control-plane/ ./
# Remove package-lock.json to avoid conflicts with installed native bindings
RUN rm -f package-lock.json
# Link SDK (temporary for build)
RUN cd /app/sdk && npm link && cd /app && npm link @vectorize-io/hindsight-client
# Build Control Plane
RUN npm run build
# Create public directory if it doesn't exist
RUN mkdir -p public
# =============================================================================
# Stage: Final Image - API Only
# =============================================================================
FROM python:3.11-slim AS api-only
WORKDIR /app
# Install pg0 dependencies
RUN apt-get update && apt-get install -y \
curl \
libxml2 \
libssl3 \
libgssapi-krb5-2 \
libossp-uuid16 \
&& apt-get install -y libicu72 || apt-get install -y libicu74 || apt-get install -y libicu* \
&& rm -rf /var/lib/apt/lists/* \
&& pip install --no-cache-dir uv
# Create non-root user (PostgreSQL cannot run as root)
RUN useradd -m -s /bin/bash hindsight
# Copy API with virtual environment from builder
COPY --from=api-builder /app/api /app/api
# Copy startup script
COPY docker/standalone/start-all.sh /app/start-all.sh
RUN chmod +x /app/start-all.sh
# Create data directory for pg0 and set ownership
RUN mkdir -p /app/data && chown -R hindsight:hindsight /app
# Switch to non-root user
USER hindsight
# Set PATH for hindsight user
ENV PATH="/app/api/.venv/bin:${PATH}"
# Pre-cache PostgreSQL binaries by starting/stopping pg0-embedded
ENV PG0_HOME=/home/hindsight/.pg0-cache
ENV PG0_HOME=/home/hindsight/.pg0
# Pre-download ML models to avoid runtime download
RUN /app/api/.venv/bin/python -c "\
from sentence_transformers import SentenceTransformer, CrossEncoder; \
print('Downloading embedding model...'); \
SentenceTransformer('BAAI/bge-small-en-v1.5'); \
print('Downloading cross-encoder model...'); \
CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2'); \
print('Models cached successfully')"
EXPOSE 8888
ENV HINDSIGHT_API_HOST=0.0.0.0
ENV HINDSIGHT_API_PORT=8888
ENV HINDSIGHT_API_LOG_LEVEL=info
ENV HINDSIGHT_ENABLE_API=true
ENV HINDSIGHT_ENABLE_CP=false
ENV PYTHONUNBUFFERED=1
CMD ["/app/start-all.sh"]
# =============================================================================
# Stage: Final Image - Control Plane Only
# =============================================================================
FROM node:20-alpine AS cp-only
WORKDIR /app
# Copy built SDK
COPY --from=sdk-builder /app/hindsight-clients/typescript /app/sdk
# Copy Control Plane standalone build
WORKDIR /app/control-plane
COPY --from=cp-builder /app/.next/standalone ./
COPY --from=cp-builder /app/.next/static ./.next/static
COPY --from=cp-builder /app/public ./public
WORKDIR /app
# Copy startup script
COPY docker/standalone/start-all.sh /app/start-all.sh
RUN chmod +x /app/start-all.sh
# Install curl for health checks
RUN apk add --no-cache curl bash
EXPOSE 9999
ENV NODE_ENV=production
ENV HINDSIGHT_CP_DATAPLANE_API_URL=http://localhost:8888
ENV HINDSIGHT_ENABLE_API=false
ENV HINDSIGHT_ENABLE_CP=true
CMD ["/app/start-all.sh"]
# =============================================================================
# Stage: Final Image - Standalone (both API and Control Plane)
# =============================================================================
FROM python:3.11-slim AS standalone
WORKDIR /app
# Install Node.js, curl, uv, and pg0 dependencies
RUN apt-get update && apt-get install -y \
curl \
libxml2 \
libssl3 \
libgssapi-krb5-2 \
libossp-uuid16 \
&& apt-get install -y libicu72 || apt-get install -y libicu74 || apt-get install -y libicu* \
&& curl -fsSL https://deb.nodesource.com/setup_20.x | bash - \
&& apt-get install -y nodejs \
&& rm -rf /var/lib/apt/lists/* \
&& pip install --no-cache-dir uv
# Create non-root user (PostgreSQL cannot run as root)
RUN useradd -m -s /bin/bash hindsight
# Copy API with virtual environment from builder
COPY --from=api-builder /app/api /app/api
# Copy built SDK
COPY --from=sdk-builder /app/hindsight-clients/typescript /app/sdk
# Copy Control Plane standalone build
WORKDIR /app/control-plane
COPY --from=cp-builder /app/.next/standalone ./
COPY --from=cp-builder /app/.next/static ./.next/static
COPY --from=cp-builder /app/public ./public
WORKDIR /app
# Copy startup script
COPY docker/standalone/start-all.sh /app/start-all.sh
RUN chmod +x /app/start-all.sh
# Create data directory for pg0 and set ownership
RUN mkdir -p /app/data && chown -R hindsight:hindsight /app
# Switch to non-root user
USER hindsight
# Set PATH for hindsight user
ENV PATH="/app/api/.venv/bin:${PATH}"
# Pre-cache PostgreSQL binaries by starting/stopping pg0-embedded
ENV PG0_HOME=/home/hindsight/.pg0-cache
RUN /app/api/.venv/bin/python -c "\
from pg0 import Pg0; \
print('Pre-caching PostgreSQL binaries...'); \
pg = Pg0(name='hindsight', port=5555, username='hindsight', password='hindsight', database='hindsight'); \
pg.start(); \
pg.stop(); \
print('PostgreSQL pre-cached to PG0_HOME')" || echo "Pre-download skipped"
ENV PG0_HOME=/home/hindsight/.pg0
# Pre-download ML models to avoid runtime download
RUN /app/api/.venv/bin/python -c "\
from sentence_transformers import SentenceTransformer, CrossEncoder; \
print('Downloading embedding model...'); \
SentenceTransformer('BAAI/bge-small-en-v1.5'); \
print('Downloading cross-encoder model...'); \
CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2'); \
print('Models cached successfully')"
EXPOSE 8888 9999
ENV HINDSIGHT_API_HOST=0.0.0.0
ENV HINDSIGHT_API_PORT=8888
ENV HINDSIGHT_API_LOG_LEVEL=info
ENV NODE_ENV=production
ENV HINDSIGHT_CP_DATAPLANE_API_URL=http://localhost:8888
ENV HINDSIGHT_ENABLE_API=true
ENV HINDSIGHT_ENABLE_CP=true
ENV PYTHONUNBUFFERED=1
CMD ["/app/start-all.sh"]
# =============================================================================
# Default target selection based on build args
# =============================================================================
FROM standalone AS default-both
FROM api-only AS default-api
FROM cp-only AS default-cp
# This selects the final stage based on INCLUDE_API and INCLUDE_CP
# Use --target to override: docker build --target api-only .
FROM standalone
-75
View File
@@ -1,75 +0,0 @@
#!/bin/bash
set -e
# Service flags (default to true if not set)
ENABLE_API="${HINDSIGHT_ENABLE_API:-true}"
ENABLE_CP="${HINDSIGHT_ENABLE_CP:-true}"
# Copy pre-cached PostgreSQL data if runtime directory is empty (first run with volume)
if [ "$ENABLE_API" = "true" ]; then
PG0_CACHE="/home/hindsight/.pg0-cache"
PG0_HOME="/home/hindsight/.pg0"
if [ -d "$PG0_CACHE" ] && [ "$(ls -A $PG0_CACHE 2>/dev/null)" ]; then
if [ ! "$(ls -A $PG0_HOME 2>/dev/null)" ]; then
echo "📦 Copying pre-cached PostgreSQL data..."
cp -r "$PG0_CACHE"/* "$PG0_HOME"/ 2>/dev/null || true
fi
fi
fi
# Track PIDs for wait
PIDS=()
# Start API if enabled
if [ "$ENABLE_API" = "true" ]; then
cd /app/api
hindsight-api 2>&1 | sed -u 's/^/[api] /' &
API_PID=$!
PIDS+=($API_PID)
# Wait for API to be ready
for i in {1..60}; do
if curl -sf http://localhost:8888/health &>/dev/null; then
break
fi
sleep 1
done
else
echo "API disabled (HINDSIGHT_ENABLE_API=false)"
fi
# Start Control Plane if enabled
if [ "$ENABLE_CP" = "true" ]; then
echo "🎛️ Starting Control Plane..."
cd /app/control-plane
PORT=9999 node server.js 2>&1 | grep -v -E "^[[:space:]]*(▲|✓|-|$)" | sed -u 's/^/[control-plane] /' &
CP_PID=$!
PIDS+=($CP_PID)
else
echo "Control Plane disabled (HINDSIGHT_ENABLE_CP=false)"
fi
# Print status
echo ""
echo "✅ Hindsight is running!"
echo ""
echo "📍 Access:"
if [ "$ENABLE_CP" = "true" ]; then
echo " Control Plane: http://localhost:9999"
fi
if [ "$ENABLE_API" = "true" ]; then
echo " API: http://localhost:8888"
fi
echo ""
# Check if any services are running
if [ ${#PIDS[@]} -eq 0 ]; then
echo "❌ No services enabled! Set HINDSIGHT_ENABLE_API=true or HINDSIGHT_ENABLE_CP=true"
exit 1
fi
# Wait for any process to exit
wait -n
# Exit with status of first exited process
exit $?
+65
View File
@@ -0,0 +1,65 @@
#!/bin/bash
set -e
cd "$(dirname "$0")"
echo "🚀 Starting Hindsight Services"
echo "============================"
echo ""
# Check if .env file exists in root
if [ ! -f ../.env ]; then
echo "⚠️ No .env file found in project root!"
echo ""
echo "Creating .env from .env.example..."
cp ../.env.example ../.env
echo ""
echo "⚠️ Please edit .env and set your API keys:"
echo " - HINDSIGHT_API_LLM_API_KEY"
echo ""
echo "Then run this script again."
exit 1
fi
echo "📦 Building and starting services..."
docker compose --env-file ../.env up --build -d
echo ""
echo "⏳ Waiting for services to be healthy..."
echo ""
# Wait for PostgreSQL
echo " Waiting for PostgreSQL..."
until docker exec hindsight-postgres pg_isready -U hindsight > /dev/null 2>&1; do
sleep 1
done
echo " ✅ PostgreSQL is ready"
# Wait for API
echo " Waiting for API..."
until curl -f http://localhost:8888/api/v1/agents > /dev/null 2>&1; do
sleep 2
done
echo " ✅ API is ready"
# Wait for Control Plane
echo " Waiting for Control Plane..."
until curl -f http://localhost:9999 > /dev/null 2>&1; do
sleep 2
done
echo " ✅ Control Plane is ready"
echo ""
echo "✅ All services are running!"
echo ""
echo "📊 Service URLs:"
echo " Control Plane: http://localhost:9999"
echo " API: http://localhost:8888"
echo " PostgreSQL: localhost:5432"
echo ""
echo "🔍 View logs:"
echo " docker compose logs -f"
echo ""
echo "🛑 Stop services:"
echo " ./stop.sh"
echo ""
Executable
+17
View File
@@ -0,0 +1,17 @@
#!/bin/bash
set -e
cd "$(dirname "$0")"
echo "🛑 Stopping Services"
echo "============================"
echo ""
docker compose down
echo ""
echo "✅ All services stopped"
echo ""
echo "💡 To remove data volumes as well, run:"
echo " docker compose down -v"
echo ""
+135
View File
@@ -0,0 +1,135 @@
HINDSIGHT HELM CHART INSTALLATION GUIDE
=====================================
PREREQUISITES
-------------
- Kubernetes cluster (1.19+)
- kubectl configured
- Helm 3.x installed
- PostgreSQL database with pgvector extension (if not using bundled PostgreSQL)
BASIC INSTALLATION
------------------
1. Install with default values (requires external PostgreSQL):
helm install hindsight ./hindsight \
--set postgresql.external.host=your-postgres-host \
--set postgresql.external.password=your-password \
--set api.secrets.MEMORY_LLM_API_KEY=your-api-key
2. Install with custom values file:
helm install hindsight ./hindsight -f hindsight/values-production.yaml
3. Install in a specific namespace:
kubectl create namespace hindsight
helm install hindsight ./hindsight -n hindsight
CONFIGURATION OPTIONS
---------------------
Development setup (using values-development.yaml):
helm install hindsight ./hindsight -f hindsight/values-development.yaml
Production setup (using values-production.yaml):
helm install hindsight ./hindsight -f hindsight/values-production.yaml
Custom LLM provider:
helm install hindsight ./hindsight \
--set api.env.MEMORY_LLM_PROVIDER=openai \
--set api.env.MEMORY_LLM_MODEL=gpt-4 \
--set api.secrets.MEMORY_LLM_API_KEY=sk-your-key
Enable ingress:
helm install hindsight ./hindsight \
--set ingress.enabled=true \
--set ingress.hosts[0].host=hindsight.example.com
Enable autoscaling:
helm install hindsight ./hindsight \
--set autoscaling.enabled=true \
--set autoscaling.minReplicas=2 \
--set autoscaling.maxReplicas=10
UPGRADE
-------
Upgrade existing installation:
helm upgrade hindsight ./hindsight
Upgrade with new values:
helm upgrade hindsight ./hindsight -f hindsight/values-production.yaml
UNINSTALL
---------
Remove the Helm release:
helm uninstall hindsight
Remove with namespace:
helm uninstall hindsight -n hindsight
TESTING
-------
Test the installation with dry-run:
helm install hindsight ./hindsight --dry-run --debug
Validate templates:
helm template hindsight ./hindsight
Lint the chart:
helm lint ./hindsight
ACCESSING THE SERVICES
----------------------
Port-forward control plane:
kubectl port-forward svc/hindsight-control-plane 3000:3000
Port-forward API:
kubectl port-forward svc/hindsight-api 8888:8888
Get service URLs:
helm status hindsight
DATABASE INITIALIZATION
-----------------------
NOTE: Database migrations now run automatically when the API service starts.
You typically don't need to run migrations manually.
If you want to pre-initialize the database before deploying (optional):
kubectl run hindsight-init --rm -it --restart=Never \
--image=hindsight/api:latest \
--env="DATABASE_URL=postgresql://user:pass@host:5432/hindsight" \
-- python -c "from hindsight.migrations import run_migrations; run_migrations()"
TROUBLESHOOTING
---------------
Check pod status:
kubectl get pods -l app.kubernetes.io/name=hindsight
View logs for API:
kubectl logs -l app.kubernetes.io/component=api
View logs for control plane:
kubectl logs -l app.kubernetes.io/component=control-plane
Describe a pod:
kubectl describe pod <pod-name>
Check configuration:
kubectl get configmap hindsight-config -o yaml
kubectl get secret hindsight-secret -o yaml
NOTES
-----
- Make sure PostgreSQL has pgvector extension enabled
- Run database migrations before first use
- Configure proper resource limits for production
- Use external secrets management for production
- Enable TLS/SSL for production deployments
-6
View File
@@ -1,6 +0,0 @@
dependencies:
- name: postgresql
repository: https://charts.bitnami.com/bitnami
version: 15.5.38
digest: sha256:f67c7612736803ece8a669f8ca6b0555f3b78557bc0ecb732aa2e43f0df7750d
generated: "2025-12-10T17:20:57.058794+01:00"
+3 -3
View File
@@ -1,9 +1,9 @@
apiVersion: v2
name: hindsight
description: Hindsight helm chart
description: A Helm chart for Hindsight - temporal-semantic-entity memory system for AI agents
type: application
version: 0.1.5
appVersion: "0.1.5"
version: 0.0.7
appVersion: "0.0.7"
keywords:
- ai
- memory
-182
View File
@@ -1,182 +0,0 @@
# Hindsight Helm Chart
Helm chart for deploying Hindsight - a temporal-semantic-entity memory system for AI agents.
## Prerequisites
- Kubernetes 1.19+
- Helm 3.0+
- PostgreSQL database (external or bundled)
## Quick Start
```bash
# Update dependencies first
helm dependency update ./helm/hindsight
# Install (PostgreSQL included by default)
export OPENAI_API_KEY="sk-your-openai-key"
helm upgrade hindsight --install ./helm/hindsight -n hindsight --create-namespace \
--set api.secrets.HINDSIGHT_API_LLM_API_KEY="$OPENAI_API_KEY"
```
To use an external database instead:
```bash
helm install hindsight ./helm/hindsight -n hindsight --create-namespace \
--set api.secrets.HINDSIGHT_API_LLM_API_KEY="sk-your-openai-key" \
--set postgresql.enabled=false \
--set postgresql.external.host=my-postgres.example.com \
--set postgresql.external.password=mypassword
```
## Installation
### Add the repository (if published)
```bash
helm repo add hindsight https://your-helm-repo.com
helm repo update
```
### Install with custom values file
Create a `values-override.yaml`:
```yaml
api:
secrets:
HINDSIGHT_API_LLM_API_KEY: "sk-your-openai-key"
postgresql:
external:
host: "my-postgres.example.com"
password: "mypassword"
```
Then install:
```bash
helm install hindsight ./helm/hindsight -n hindsight --create-namespace -f values-override.yaml
```
## Configuration
### Key Values
| Parameter | Description | Default |
|-----------|-------------|---------|
| `version` | Default image tag for all components | `0.1.0` |
| `api.enabled` | Enable the API component | `true` |
| `api.image.repository` | API image repository | `hindsight/api` |
| `api.image.tag` | API image tag (defaults to `version`) | - |
| `api.service.port` | API service port | `8888` |
| `controlPlane.enabled` | Enable the control plane | `true` |
| `controlPlane.image.repository` | Control plane image repository | `hindsight/control-plane` |
| `controlPlane.image.tag` | Control plane image tag (defaults to `version`) | - |
| `controlPlane.service.port` | Control plane service port | `3000` |
| `postgresql.enabled` | Deploy PostgreSQL as subchart | `true` |
| `postgresql.external.host` | External PostgreSQL host | `postgresql` |
| `postgresql.external.port` | External PostgreSQL port | `5432` |
| `postgresql.external.database` | Database name | `hindsight` |
| `postgresql.external.username` | Database username | `hindsight` |
| `ingress.enabled` | Enable ingress | `false` |
| `autoscaling.enabled` | Enable HPA | `false` |
### Environment Variables
All environment variables in `api.env` and `controlPlane.env` are automatically added to the respective pods. Sensitive values should go in `api.secrets` or `controlPlane.secrets`.
```yaml
api:
env:
HINDSIGHT_API_LLM_PROVIDER: "openai"
HINDSIGHT_API_LLM_MODEL: "gpt-4"
secrets:
HINDSIGHT_API_LLM_API_KEY: "your-api-key"
HINDSIGHT_API_LLM_BASE_URL: "https://api.openai.com/v1"
controlPlane:
env:
NODE_ENV: "production"
secrets: {}
```
### External Database
To connect to an external PostgreSQL database:
```yaml
postgresql:
enabled: false
external:
host: "my-postgres.example.com"
port: 5432
database: "hindsight"
username: "hindsight"
password: "your-password"
```
### Ingress
To expose the services via ingress:
```yaml
ingress:
enabled: true
className: "nginx"
annotations:
cert-manager.io/cluster-issuer: "letsencrypt-prod"
hosts:
- host: hindsight.example.com
paths:
- path: /
pathType: Prefix
service: controlPlane
- path: /api
pathType: Prefix
service: api
tls:
- secretName: hindsight-tls
hosts:
- hindsight.example.com
```
## Upgrading
```bash
helm upgrade hindsight ./helm/hindsight -n hindsight
```
## Uninstalling
```bash
helm uninstall hindsight -n hindsight
```
## Components
The chart deploys:
- **API**: The main Hindsight API server for memory operations
- **Control Plane**: Web UI for managing agents and viewing memories
## Development
### Lint the chart
```bash
helm lint ./helm/hindsight
```
### Template locally
```bash
helm template hindsight ./helm/hindsight --debug
```
### Dry run installation
```bash
helm install hindsight ./helm/hindsight --dry-run --debug
```
+71 -2
View File
@@ -1,2 +1,71 @@
Hindsight installed. Access the control plane:
kubectl port-forward -n {{ .Release.Namespace }} svc/{{ include "hindsight.fullname" . }}-control-plane 3000:3000
Thank you for installing {{ .Chart.Name }}!
Your release is named {{ .Release.Name }}.
To learn more about the release, try:
$ helm status {{ .Release.Name }}
$ helm get all {{ .Release.Name }}
{{- if .Values.ingress.enabled }}
The application is accessible via the following URL(s):
{{- range .Values.ingress.hosts }}
- http{{ if $.Values.ingress.tls }}s{{ end }}://{{ .host }}
{{- end }}
{{- else }}
1. Get the Control Plane URL by running these commands:
{{- if contains "NodePort" .Values.controlPlane.service.type }}
export NODE_PORT=$(kubectl get --namespace {{ .Release.Namespace }} -o jsonpath="{.spec.ports[0].nodePort}" services {{ include "hindsight.fullname" . }}-control-plane)
export NODE_IP=$(kubectl get nodes --namespace {{ .Release.Namespace }} -o jsonpath="{.items[0].status.addresses[0].address}")
echo "Control Plane URL: http://$NODE_IP:$NODE_PORT"
{{- else if contains "LoadBalancer" .Values.controlPlane.service.type }}
NOTE: It may take a few minutes for the LoadBalancer IP to be available.
You can watch the status by running 'kubectl get --namespace {{ .Release.Namespace }} svc -w {{ include "hindsight.fullname" . }}-control-plane'
export SERVICE_IP=$(kubectl get svc --namespace {{ .Release.Namespace }} {{ include "hindsight.fullname" . }}-control-plane --template "{{"{{ range (index .status.loadBalancer.ingress 0) }}{{.}}{{ end }}"}}")
echo "Control Plane URL: http://$SERVICE_IP:{{ .Values.controlPlane.service.port }}"
{{- else if contains "ClusterIP" .Values.controlPlane.service.type }}
export POD_NAME=$(kubectl get pods --namespace {{ .Release.Namespace }} -l "app.kubernetes.io/component=control-plane,app.kubernetes.io/instance={{ .Release.Name }}" -o jsonpath="{.items[0].metadata.name}")
export CONTAINER_PORT=$(kubectl get pod --namespace {{ .Release.Namespace }} $POD_NAME -o jsonpath="{.spec.containers[0].ports[0].containerPort}")
echo "Control Plane URL: http://127.0.0.1:3000"
kubectl --namespace {{ .Release.Namespace }} port-forward $POD_NAME 3000:$CONTAINER_PORT
{{- end }}
2. Get the API URL by running these commands:
{{- if contains "NodePort" .Values.api.service.type }}
export NODE_PORT=$(kubectl get --namespace {{ .Release.Namespace }} -o jsonpath="{.spec.ports[0].nodePort}" services {{ include "hindsight.fullname" . }}-api)
export NODE_IP=$(kubectl get nodes --namespace {{ .Release.Namespace }} -o jsonpath="{.items[0].status.addresses[0].address}")
echo "API URL: http://$NODE_IP:$NODE_PORT"
{{- else if contains "LoadBalancer" .Values.api.service.type }}
NOTE: It may take a few minutes for the LoadBalancer IP to be available.
You can watch the status by running 'kubectl get --namespace {{ .Release.Namespace }} svc -w {{ include "hindsight.fullname" . }}-api'
export SERVICE_IP=$(kubectl get svc --namespace {{ .Release.Namespace }} {{ include "hindsight.fullname" . }}-api --template "{{"{{ range (index .status.loadBalancer.ingress 0) }}{{.}}{{ end }}"}}")
echo "API URL: http://$SERVICE_IP:{{ .Values.api.service.port }}"
{{- else if contains "ClusterIP" .Values.api.service.type }}
export POD_NAME=$(kubectl get pods --namespace {{ .Release.Namespace }} -l "app.kubernetes.io/component=api,app.kubernetes.io/instance={{ .Release.Name }}" -o jsonpath="{.items[0].metadata.name}")
export CONTAINER_PORT=$(kubectl get pod --namespace {{ .Release.Namespace }} $POD_NAME -o jsonpath="{.spec.containers[0].ports[0].containerPort}")
echo "API URL: http://127.0.0.1:8888"
kubectl --namespace {{ .Release.Namespace }} port-forward $POD_NAME 8888:$CONTAINER_PORT
{{- end }}
{{- end }}
{{- if not .Values.postgresql.enabled }}
NOTE: You are using an external PostgreSQL database.
Please ensure that:
1. The database is accessible from the cluster
2. The pgvector extension is enabled
Database migrations run automatically when the API service starts.
If you want to pre-initialize the database before deploying (optional):
kubectl run --namespace {{ .Release.Namespace }} hindsight-init --rm -it --restart=Never \
--image={{ .Values.api.image.repository }}:{{ .Values.api.image.tag }} \
--env="DATABASE_URL={{ include "hindsight.databaseUrl" . }}" \
-- python -c "from hindsight.migrations import run_migrations; run_migrations()"
{{- end }}
For more information, visit: https://github.com/yourusername/hindsight
+22 -22
View File
@@ -1,14 +1,14 @@
{{/*
Expand the name of the chart.
*/}}
{{- define "hindsight.name" -}}
{{- define "memora.name" -}}
{{- default .Chart.Name .Values.nameOverride | trunc 63 | trimSuffix "-" }}
{{- end }}
{{/*
Create a default fully qualified app name.
*/}}
{{- define "hindsight.fullname" -}}
{{- define "memora.fullname" -}}
{{- if .Values.fullnameOverride }}
{{- .Values.fullnameOverride | trunc 63 | trimSuffix "-" }}
{{- else }}
@@ -24,16 +24,16 @@ Create a default fully qualified app name.
{{/*
Create chart name and version as used by the chart label.
*/}}
{{- define "hindsight.chart" -}}
{{- define "memora.chart" -}}
{{- printf "%s-%s" .Chart.Name .Chart.Version | replace "+" "_" | trunc 63 | trimSuffix "-" }}
{{- end }}
{{/*
Common labels
*/}}
{{- define "hindsight.labels" -}}
helm.sh/chart: {{ include "hindsight.chart" . }}
{{ include "hindsight.selectorLabels" . }}
{{- define "memora.labels" -}}
helm.sh/chart: {{ include "memora.chart" . }}
{{ include "memora.selectorLabels" . }}
{{- if .Chart.AppVersion }}
app.kubernetes.io/version: {{ .Chart.AppVersion | quote }}
{{- end }}
@@ -43,49 +43,49 @@ app.kubernetes.io/managed-by: {{ .Release.Service }}
{{/*
Selector labels
*/}}
{{- define "hindsight.selectorLabels" -}}
app.kubernetes.io/name: {{ include "hindsight.name" . }}
{{- define "memora.selectorLabels" -}}
app.kubernetes.io/name: {{ include "memora.name" . }}
app.kubernetes.io/instance: {{ .Release.Name }}
{{- end }}
{{/*
API labels
*/}}
{{- define "hindsight.api.labels" -}}
{{ include "hindsight.labels" . }}
{{- define "memora.api.labels" -}}
{{ include "memora.labels" . }}
app.kubernetes.io/component: api
{{- end }}
{{/*
API selector labels
*/}}
{{- define "hindsight.api.selectorLabels" -}}
{{ include "hindsight.selectorLabels" . }}
{{- define "memora.api.selectorLabels" -}}
{{ include "memora.selectorLabels" . }}
app.kubernetes.io/component: api
{{- end }}
{{/*
Control plane labels
*/}}
{{- define "hindsight.controlPlane.labels" -}}
{{ include "hindsight.labels" . }}
{{- define "memora.controlPlane.labels" -}}
{{ include "memora.labels" . }}
app.kubernetes.io/component: control-plane
{{- end }}
{{/*
Control plane selector labels
*/}}
{{- define "hindsight.controlPlane.selectorLabels" -}}
{{ include "hindsight.selectorLabels" . }}
{{- define "memora.controlPlane.selectorLabels" -}}
{{ include "memora.selectorLabels" . }}
app.kubernetes.io/component: control-plane
{{- end }}
{{/*
Create the name of the service account to use
*/}}
{{- define "hindsight.serviceAccountName" -}}
{{- define "memora.serviceAccountName" -}}
{{- if .Values.serviceAccount.create }}
{{- default (include "hindsight.fullname" .) .Values.serviceAccount.name }}
{{- default (include "memora.fullname" .) .Values.serviceAccount.name }}
{{- else }}
{{- default "default" .Values.serviceAccount.name }}
{{- end }}
@@ -94,11 +94,11 @@ Create the name of the service account to use
{{/*
Generate database URL
*/}}
{{- define "hindsight.databaseUrl" -}}
{{- define "memora.databaseUrl" -}}
{{- if .Values.databaseUrl }}
{{- .Values.databaseUrl }}
{{- else if .Values.postgresql.enabled }}
{{- printf "postgresql://%s:%s@%s-postgresql:%d/%s" .Values.postgresql.auth.username .Values.postgresql.auth.password (include "hindsight.fullname" .) (.Values.postgresql.service.port | int) .Values.postgresql.auth.database }}
{{- printf "postgresql://%s:%s@%s-postgresql:%d/%s" .Values.postgresql.auth.username .Values.postgresql.auth.password (include "memora.fullname" .) (.Values.postgresql.primary.service.port | int) .Values.postgresql.auth.database }}
{{- else }}
{{- printf "postgresql://%s:$(POSTGRES_PASSWORD)@%s:%d/%s" .Values.postgresql.external.username .Values.postgresql.external.host (.Values.postgresql.external.port | int) .Values.postgresql.external.database }}
{{- end }}
@@ -107,6 +107,6 @@ Generate database URL
{{/*
API URL for control plane
*/}}
{{- define "hindsight.apiUrl" -}}
{{- printf "http://%s-api:%d" (include "hindsight.fullname" .) (.Values.api.service.port | int) }}
{{- define "memora.apiUrl" -}}
{{- printf "http://%s-api:%d" (include "memora.fullname" .) (.Values.api.service.port | int) }}
{{- end }}
+23 -9
View File
@@ -15,6 +15,7 @@ spec:
template:
metadata:
annotations:
checksum/config: {{ include (print $.Template.BasePath "/configmap.yaml") . | sha256sum }}
checksum/secret: {{ include (print $.Template.BasePath "/secret.yaml") . | sha256sum }}
{{- with .Values.podAnnotations }}
{{- toYaml . | nindent 8 }}
@@ -31,7 +32,7 @@ spec:
- name: api
securityContext:
{{- toYaml .Values.securityContext | nindent 10 }}
image: "{{ .Values.api.image.repository }}:{{ .Values.api.image.tag | default .Values.version }}"
image: "{{ .Values.api.image.repository }}:{{ .Values.api.image.tag }}"
imagePullPolicy: {{ .Values.api.image.pullPolicy }}
ports:
- name: http
@@ -47,16 +48,29 @@ spec:
name: {{ include "hindsight.fullname" . }}-secret
key: postgres-password
{{- end }}
{{- range $key, $value := .Values.api.env }}
- name: {{ $key }}
value: {{ $value | quote }}
{{- end }}
{{- range $key, $value := .Values.api.secrets }}
- name: {{ $key }}
- name: HINDSIGHT_API_LLM_PROVIDER
valueFrom:
configMapKeyRef:
name: {{ include "hindsight.fullname" . }}-config
key: llm-provider
- name: HINDSIGHT_API_LLM_MODEL
valueFrom:
configMapKeyRef:
name: {{ include "hindsight.fullname" . }}-config
key: llm-model
{{- if and .Values.api.secrets (hasKey .Values.api.secrets "HINDSIGHT_API_LLM_API_KEY") }}
- name: HINDSIGHT_API_LLM_API_KEY
valueFrom:
secretKeyRef:
name: {{ include "hindsight.fullname" $ }}-secret
key: {{ $key }}
name: {{ include "hindsight.fullname" . }}-secret
key: llm-api-key
{{- end }}
{{- if and .Values.api.secrets (hasKey .Values.api.secrets "HINDSIGHT_API_LLM_BASE_URL") }}
- name: HINDSIGHT_API_LLM_BASE_URL
valueFrom:
secretKeyRef:
name: {{ include "hindsight.fullname" . }}-secret
key: llm-base-url
{{- end }}
livenessProbe:
{{- toYaml .Values.api.livenessProbe | nindent 10 }}
+15
View File
@@ -0,0 +1,15 @@
apiVersion: v1
kind: ConfigMap
metadata:
name: {{ include "hindsight.fullname" . }}-config
labels:
{{- include "hindsight.labels" . | nindent 4 }}
data:
# API configuration
llm-provider: {{ .Values.api.env.HINDSIGHT_API_LLM_PROVIDER | quote }}
llm-model: {{ .Values.api.env.HINDSIGHT_API_LLM_MODEL | quote }}
# Control plane configuration
node-env: {{ .Values.controlPlane.env.NODE_ENV | quote }}
hostname: {{ .Values.controlPlane.env.HINDSIGHT_CP_HOSTNAME | quote }}
control-plane-port: {{ .Values.controlPlane.env.HINDSIGHT_CP_PORT | quote }}
@@ -15,7 +15,7 @@ spec:
template:
metadata:
annotations:
checksum/secret: {{ include (print $.Template.BasePath "/secret.yaml") . | sha256sum }}
checksum/config: {{ include (print $.Template.BasePath "/configmap.yaml") . | sha256sum }}
{{- with .Values.podAnnotations }}
{{- toYaml . | nindent 8 }}
{{- end }}
@@ -31,26 +31,30 @@ spec:
- name: control-plane
securityContext:
{{- toYaml .Values.securityContext | nindent 10 }}
image: "{{ .Values.controlPlane.image.repository }}:{{ .Values.controlPlane.image.tag | default .Values.version }}"
image: "{{ .Values.controlPlane.image.repository }}:{{ .Values.controlPlane.image.tag }}"
imagePullPolicy: {{ .Values.controlPlane.image.pullPolicy }}
ports:
- name: http
containerPort: {{ .Values.controlPlane.service.targetPort }}
protocol: TCP
env:
- name: NODE_ENV
valueFrom:
configMapKeyRef:
name: {{ include "hindsight.fullname" . }}-config
key: node-env
- name: HINDSIGHT_CP_HOSTNAME
valueFrom:
configMapKeyRef:
name: {{ include "hindsight.fullname" . }}-config
key: hostname
- name: HINDSIGHT_CP_PORT
valueFrom:
configMapKeyRef:
name: {{ include "hindsight.fullname" . }}-config
key: control-plane-port
- name: HINDSIGHT_CP_DATAPLANE_API_URL
value: {{ include "hindsight.apiUrl" . | quote }}
{{- range $key, $value := .Values.controlPlane.env }}
- name: {{ $key }}
value: {{ $value | quote }}
{{- end }}
{{- range $key, $value := .Values.controlPlane.secrets }}
- name: {{ $key }}
valueFrom:
secretKeyRef:
name: {{ include "hindsight.fullname" $ }}-secret
key: {{ $key }}
{{- end }}
livenessProbe:
{{- toYaml .Values.controlPlane.livenessProbe | nindent 10 }}
readinessProbe:
@@ -1,19 +0,0 @@
{{- if .Values.postgresql.enabled }}
apiVersion: v1
kind: Service
metadata:
name: {{ include "hindsight.fullname" . }}-postgresql
labels:
{{- include "hindsight.labels" . | nindent 4 }}
app.kubernetes.io/component: postgresql
spec:
type: ClusterIP
ports:
- port: {{ .Values.postgresql.service.port }}
targetPort: postgresql
protocol: TCP
name: postgresql
selector:
{{- include "hindsight.selectorLabels" . | nindent 4 }}
app.kubernetes.io/component: postgresql
{{- end }}
@@ -1,85 +0,0 @@
{{- if .Values.postgresql.enabled }}
apiVersion: apps/v1
kind: StatefulSet
metadata:
name: {{ include "hindsight.fullname" . }}-postgresql
labels:
{{- include "hindsight.labels" . | nindent 4 }}
app.kubernetes.io/component: postgresql
spec:
serviceName: {{ include "hindsight.fullname" . }}-postgresql
replicas: 1
selector:
matchLabels:
{{- include "hindsight.selectorLabels" . | nindent 6 }}
app.kubernetes.io/component: postgresql
template:
metadata:
labels:
{{- include "hindsight.selectorLabels" . | nindent 8 }}
app.kubernetes.io/component: postgresql
spec:
containers:
- name: postgresql
image: "{{ .Values.postgresql.image.repository }}:{{ .Values.postgresql.image.tag }}"
imagePullPolicy: {{ .Values.postgresql.image.pullPolicy }}
ports:
- name: postgresql
containerPort: 5432
protocol: TCP
env:
- name: POSTGRES_USER
value: {{ .Values.postgresql.auth.username | quote }}
- name: POSTGRES_PASSWORD
value: {{ .Values.postgresql.auth.password | quote }}
- name: POSTGRES_DB
value: {{ .Values.postgresql.auth.database | quote }}
- name: PGDATA
value: /var/lib/postgresql/data/pgdata
livenessProbe:
exec:
command:
- pg_isready
- -U
- {{ .Values.postgresql.auth.username }}
- -d
- {{ .Values.postgresql.auth.database }}
initialDelaySeconds: 30
periodSeconds: 10
timeoutSeconds: 5
failureThreshold: 3
readinessProbe:
exec:
command:
- pg_isready
- -U
- {{ .Values.postgresql.auth.username }}
- -d
- {{ .Values.postgresql.auth.database }}
initialDelaySeconds: 5
periodSeconds: 5
timeoutSeconds: 3
failureThreshold: 3
resources:
{{- toYaml .Values.postgresql.resources | nindent 10 }}
volumeMounts:
- name: data
mountPath: /var/lib/postgresql/data
{{- if .Values.postgresql.persistence.enabled }}
volumeClaimTemplates:
- metadata:
name: data
spec:
accessModes: ["ReadWriteOnce"]
{{- if .Values.postgresql.persistence.storageClass }}
storageClassName: {{ .Values.postgresql.persistence.storageClass | quote }}
{{- end }}
resources:
requests:
storage: {{ .Values.postgresql.persistence.size }}
{{- else }}
volumes:
- name: data
emptyDir: {}
{{- end }}
{{- end }}
+7 -5
View File
@@ -6,12 +6,14 @@ metadata:
{{- include "hindsight.labels" . | nindent 4 }}
type: Opaque
data:
{{- range $key, $value := .Values.api.secrets }}
{{ $key }}: {{ $value | b64enc | quote }}
{{- if and .Values.api.secrets (hasKey .Values.api.secrets "MEMORY_LLM_API_KEY") }}
llm-api-key: {{ .Values.api.secrets.MEMORY_LLM_API_KEY | b64enc | quote }}
{{- end }}
{{- range $key, $value := .Values.controlPlane.secrets }}
{{ $key }}: {{ $value | b64enc | quote }}
{{- if and .Values.api.secrets (hasKey .Values.api.secrets "MEMORY_LLM_BASE_URL") }}
llm-base-url: {{ .Values.api.secrets.MEMORY_LLM_BASE_URL | b64enc | quote }}
{{- end }}
{{- if and (not .Values.postgresql.enabled) .Values.postgresql.external.password }}
{{- if not .Values.postgresql.enabled }}
{{- if .Values.postgresql.external.password }}
postgres-password: {{ .Values.postgresql.external.password | b64enc | quote }}
{{- end }}
{{- end }}
+18 -41
View File
@@ -1,8 +1,5 @@
# Default values for hindsight
# Chart version - use this to set a consistent image tag across all components
version: "0.1.1"
# Global settings
replicaCount: 1
@@ -11,9 +8,9 @@ api:
enabled: true
replicaCount: 1
image:
repository: ghcr.io/vectorize-io/hindsight-api
repository: hindsight/api
pullPolicy: IfNotPresent
# tag defaults to .Values.version if not specified
tag: "latest"
service:
type: ClusterIP
@@ -32,7 +29,7 @@ api:
# Liveness and readiness probes
livenessProbe:
httpGet:
path: /health
path: /
port: 8888
initialDelaySeconds: 30
periodSeconds: 10
@@ -41,7 +38,7 @@ api:
readinessProbe:
httpGet:
path: /health
path: /
port: 8888
initialDelaySeconds: 10
periodSeconds: 5
@@ -50,7 +47,7 @@ api:
# Environment variables
env:
#HINDSIGHT_API_LLM_PROVIDER: "groq"
HINDSIGHT_API_LLM_PROVIDER: "groq"
HINDSIGHT_API_LLM_MODEL: "openai/gpt-oss-120b"
# Secret environment variables
@@ -63,9 +60,9 @@ controlPlane:
enabled: true
replicaCount: 1
image:
repository: ghcr.io/vectorize-io/hindsight-control-plane
repository: hindsight/hindsight-control-plane
pullPolicy: IfNotPresent
# tag defaults to .Values.version if not specified
tag: "latest"
service:
type: ClusterIP
@@ -81,9 +78,10 @@ controlPlane:
cpu: 250m
memory: 512Mi
# Liveness and readiness probes (TCP check)
# Liveness and readiness probes
livenessProbe:
tcpSocket:
httpGet:
path: /
port: 3000
initialDelaySeconds: 30
periodSeconds: 10
@@ -91,7 +89,8 @@ controlPlane:
failureThreshold: 3
readinessProbe:
tcpSocket:
httpGet:
path: /
port: 3000
initialDelaySeconds: 10
periodSeconds: 5
@@ -107,43 +106,21 @@ controlPlane:
# PostgreSQL configuration
postgresql:
# Set to true to deploy PostgreSQL as part of this chart
enabled: true
image:
repository: ankane/pgvector
tag: latest
pullPolicy: IfNotPresent
auth:
username: "hindsight"
password: "hindsight"
database: "hindsight"
service:
port: 5432
persistence:
enabled: true
size: 8Gi
# storageClass: ""
resources:
limits:
cpu: 1000m
memory: 1Gi
requests:
cpu: 250m
memory: 256Mi
enabled: false
# External PostgreSQL connection details
# Only used if postgresql.enabled is false
# If postgresql.enabled is false, provide external database details
external:
host: "postgresql"
port: 5432
database: "hindsight"
username: "hindsight"
# Password should be provided via secret
# password: ""
# Database URL (auto-generated from postgresql config if not provided)
# databaseUrl: "postgresql://user:pass@host:5432/database"
# Ingress configuration
ingress:
enabled: false
@@ -105,8 +105,6 @@ def run_migrations_offline() -> None:
def run_migrations_online() -> None:
"""Run migrations in 'online' mode with synchronous engine."""
from sqlalchemy import event, text
get_database_url() # Process and set the database URL in config
connectable = engine_from_config(
@@ -115,19 +113,7 @@ def run_migrations_online() -> None:
poolclass=pool.NullPool,
)
# Add event listener to ensure connection is in read-write mode
# This is needed for Supabase which may start connections in read-only mode
@event.listens_for(connectable, "connect")
def set_read_write_mode(dbapi_connection, connection_record):
cursor = dbapi_connection.cursor()
cursor.execute("SET SESSION CHARACTERISTICS AS TRANSACTION READ WRITE")
cursor.close()
with connectable.connect() as connection:
# Also explicitly set read-write mode on this connection
connection.execute(text("SET SESSION CHARACTERISTICS AS TRANSACTION READ WRITE"))
connection.commit() # Commit the SET command
context.configure(
connection=connection,
target_metadata=target_metadata
@@ -136,9 +122,6 @@ def run_migrations_online() -> None:
with context.begin_transaction():
context.run_migrations()
# Explicit commit to ensure changes are persisted (especially for Supabase)
connection.commit()
if context.is_offline_mode():
run_migrations_offline()
@@ -0,0 +1,40 @@
"""Fix memory_links entity_id to be nullable
Revision ID: 01f989db9079
Revises: af0413383b3e
Create Date: 2025-11-03 14:43:18.721430
"""
from typing import Sequence, Union
from alembic import op
import sqlalchemy as sa
# revision identifiers, used by Alembic.
revision: str = '01f989db9079'
down_revision: Union[str, Sequence[str], None] = 'af0413383b3e'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None
def upgrade() -> None:
"""Upgrade schema."""
# Drop the existing primary key
op.execute('ALTER TABLE memory_links DROP CONSTRAINT memory_links_pkey')
# Change entity_id to nullable
op.alter_column('memory_links', 'entity_id',
existing_type=sa.UUID(),
nullable=True)
# Create a unique index with COALESCE expression to handle NULL entity_id
op.execute("""
CREATE UNIQUE INDEX idx_memory_links_unique
ON memory_links (from_unit_id, to_unit_id, link_type, COALESCE(entity_id, '00000000-0000-0000-0000-000000000000'::uuid))
""")
def downgrade() -> None:
"""Downgrade schema."""
pass
@@ -0,0 +1,48 @@
"""add_async_operations_table
Revision ID: 0e96398aae9e
Revises: 1a35a4fa1950
Create Date: 2025-11-07 14:54:21.224968
"""
from typing import Sequence, Union
from alembic import op
import sqlalchemy as sa
# revision identifiers, used by Alembic.
revision: str = '0e96398aae9e'
down_revision: Union[str, Sequence[str], None] = '1a35a4fa1950'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None
def upgrade() -> None:
"""Upgrade schema."""
# Create async_operations table
op.execute("""
CREATE TABLE async_operations (
id UUID PRIMARY KEY,
agent_id TEXT NOT NULL,
task_type TEXT NOT NULL,
items_count INTEGER NOT NULL,
document_id TEXT,
created_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW()
)
""")
# Create index on agent_id for fast lookups by agent
op.execute("""
CREATE INDEX idx_async_operations_agent_id
ON async_operations(agent_id)
""")
def downgrade() -> None:
"""Downgrade schema."""
# Drop index
op.execute("DROP INDEX IF EXISTS idx_async_operations_agent_id")
# Drop table
op.execute("DROP TABLE IF EXISTS async_operations")
@@ -0,0 +1,47 @@
"""add_agents_table
Revision ID: 1680fc9768b4
Revises: 8c55f5602451
Create Date: 2025-11-12 16:18:06.620862
"""
from typing import Sequence, Union
from alembic import op
import sqlalchemy as sa
# revision identifiers, used by Alembic.
revision: str = '1680fc9768b4'
down_revision: Union[str, Sequence[str], None] = '8c55f5602451'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None
def upgrade() -> None:
"""Upgrade schema."""
# Create agents table
op.execute("""
CREATE TABLE agents (
agent_id TEXT PRIMARY KEY,
personality JSONB NOT NULL DEFAULT '{"openness": 0.5, "conscientiousness": 0.5, "extraversion": 0.5, "agreeableness": 0.5, "neuroticism": 0.5, "bias_strength": 0.5}'::jsonb,
background TEXT NOT NULL DEFAULT '',
created_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW(),
updated_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW()
)
""")
# Create index on agent_id for fast lookups
op.execute("""
CREATE INDEX idx_agents_agent_id
ON agents(agent_id)
""")
def downgrade() -> None:
"""Downgrade schema."""
# Drop index
op.execute("DROP INDEX IF EXISTS idx_agents_agent_id")
# Drop table
op.execute("DROP TABLE IF EXISTS agents")
@@ -0,0 +1,74 @@
"""add_bm25_fulltext_search
Revision ID: 1a35a4fa1950
Revises: 01f989db9079
Create Date: 2025-11-06 11:19:48.627698
"""
from typing import Sequence, Union
from alembic import op
import sqlalchemy as sa
# revision identifiers, used by Alembic.
revision: str = '1a35a4fa1950'
down_revision: Union[str, Sequence[str], None] = '01f989db9079'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None
def upgrade() -> None:
"""Upgrade schema."""
# Add tsvector column for full-text search
op.execute("""
ALTER TABLE memory_units
ADD COLUMN search_vector tsvector
""")
# Populate tsvector with existing data (text + context combined)
op.execute("""
UPDATE memory_units
SET search_vector =
setweight(to_tsvector('english', COALESCE(text, '')), 'A') ||
setweight(to_tsvector('english', COALESCE(context, '')), 'B')
""")
# Create GIN index for fast full-text search
op.execute("""
CREATE INDEX idx_memory_units_search_vector
ON memory_units
USING GIN(search_vector)
""")
# Create trigger to auto-update tsvector on INSERT/UPDATE
op.execute("""
CREATE OR REPLACE FUNCTION memory_units_search_vector_trigger() RETURNS trigger AS $$
BEGIN
NEW.search_vector :=
setweight(to_tsvector('english', COALESCE(NEW.text, '')), 'A') ||
setweight(to_tsvector('english', COALESCE(NEW.context, '')), 'B');
RETURN NEW;
END
$$ LANGUAGE plpgsql;
""")
op.execute("""
CREATE TRIGGER update_memory_units_search_vector
BEFORE INSERT OR UPDATE ON memory_units
FOR EACH ROW
EXECUTE FUNCTION memory_units_search_vector_trigger();
""")
def downgrade() -> None:
"""Downgrade schema."""
# Drop trigger
op.execute("DROP TRIGGER IF EXISTS update_memory_units_search_vector ON memory_units")
op.execute("DROP FUNCTION IF EXISTS memory_units_search_vector_trigger()")
# Drop index
op.execute("DROP INDEX IF EXISTS idx_memory_units_search_vector")
# Drop column
op.execute("ALTER TABLE memory_units DROP COLUMN IF EXISTS search_vector")
@@ -0,0 +1,28 @@
"""merge agents and temporal ranges branches
Revision ID: 217b2227771f
Revises: 3b9c4d8e7f21, 9d42e6f91234
Create Date: 2025-11-17 14:59:01.254543
"""
from typing import Sequence, Union
from alembic import op
import sqlalchemy as sa
# revision identifiers, used by Alembic.
revision: str = '217b2227771f'
down_revision: Union[str, Sequence[str], None] = ('3b9c4d8e7f21', '9d42e6f91234')
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None
def upgrade() -> None:
"""Upgrade schema."""
pass
def downgrade() -> None:
"""Downgrade schema."""
pass
@@ -0,0 +1,49 @@
"""add_status_and_error_to_async_operations
Revision ID: 2a76a5bc2f09
Revises: 0e96398aae9e
Create Date: 2025-11-07 16:03:19.078561
"""
from typing import Sequence, Union
from alembic import op
import sqlalchemy as sa
# revision identifiers, used by Alembic.
revision: str = '2a76a5bc2f09'
down_revision: Union[str, Sequence[str], None] = '0e96398aae9e'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None
def upgrade() -> None:
"""Upgrade schema."""
# Add status column (default 'pending' for existing rows)
op.execute("""
ALTER TABLE async_operations
ADD COLUMN status TEXT NOT NULL DEFAULT 'pending'
""")
# Add error_message column
op.execute("""
ALTER TABLE async_operations
ADD COLUMN error_message TEXT
""")
# Add index on status for filtering failed/pending operations
op.execute("""
CREATE INDEX idx_async_operations_status
ON async_operations(status)
""")
def downgrade() -> None:
"""Downgrade schema."""
# Drop index
op.execute("DROP INDEX IF EXISTS idx_async_operations_status")
# Drop columns
op.execute("ALTER TABLE async_operations DROP COLUMN IF EXISTS error_message")
op.execute("ALTER TABLE async_operations DROP COLUMN IF EXISTS status")
@@ -0,0 +1,36 @@
"""add_name_to_agents
Revision ID: 3b9c4d8e7f21
Revises: 1680fc9768b4
Create Date: 2025-11-13 14:52:00.000000
"""
from typing import Sequence, Union
from alembic import op
import sqlalchemy as sa
# revision identifiers, used by Alembic.
revision: str = '3b9c4d8e7f21'
down_revision: Union[str, Sequence[str], None] = '1680fc9768b4'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None
def upgrade() -> None:
"""Upgrade schema."""
# Add name column to agents table
op.execute("""
ALTER TABLE agents
ADD COLUMN name TEXT NOT NULL DEFAULT ''
""")
def downgrade() -> None:
"""Downgrade schema."""
# Remove name column from agents table
op.execute("""
ALTER TABLE agents
DROP COLUMN name
""")
@@ -0,0 +1,32 @@
"""add metadata to memory_units
Revision ID: 4a8b3c5d6e7f
Revises: 217b2227771f
Create Date: 2025-11-21 10:00:00.000000
"""
from typing import Sequence, Union
from alembic import op
import sqlalchemy as sa
from sqlalchemy.dialects import postgresql
# revision identifiers, used by Alembic.
revision: str = '4a8b3c5d6e7f'
down_revision: Union[str, Sequence[str], None] = '217b2227771f'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None
def upgrade() -> None:
"""Add metadata column to memory_units table."""
op.add_column(
'memory_units',
sa.Column('metadata', postgresql.JSONB(astext_type=sa.Text()), server_default=sa.text("'{}'::jsonb"), nullable=False)
)
def downgrade() -> None:
"""Remove metadata column from memory_units table."""
op.drop_column('memory_units', 'metadata')
@@ -0,0 +1,30 @@
"""remove_entity_type_column
Revision ID: 8c55f5602451
Revises: 2a76a5bc2f09
Create Date: 2025-11-07 17:08:07.329740
"""
from typing import Sequence, Union
from alembic import op
import sqlalchemy as sa
# revision identifiers, used by Alembic.
revision: str = '8c55f5602451'
down_revision: Union[str, Sequence[str], None] = '2a76a5bc2f09'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None
def upgrade() -> None:
"""Upgrade schema."""
# Remove entity_type column from entities table
op.execute("ALTER TABLE entities DROP COLUMN IF EXISTS entity_type")
def downgrade() -> None:
"""Downgrade schema."""
# Re-add entity_type column (default to 'OTHER' for existing rows)
op.execute("ALTER TABLE entities ADD COLUMN entity_type TEXT DEFAULT 'OTHER'")
@@ -0,0 +1,70 @@
"""add_temporal_ranges_to_memory_units
Revision ID: 9d42e6f91234
Revises: 8c55f5602451
Create Date: 2025-11-17 00:00:00.000000
This migration adds temporal range support to memory_units table:
- occurred_start: When the fact/event started
- occurred_end: When the fact/event ended
- mentioned_at: When the fact was mentioned/learned
For existing rows, these are initialized from event_date (point events).
"""
from typing import Sequence, Union
from alembic import op
import sqlalchemy as sa
from sqlalchemy.dialects.postgresql import TIMESTAMP
# revision identifiers, used by Alembic.
revision: str = '9d42e6f91234'
down_revision: Union[str, Sequence[str], None] = '8c55f5602451'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None
def upgrade() -> None:
"""Upgrade schema: add temporal range columns to memory_units."""
# Add new temporal range columns (nullable initially)
op.add_column(
'memory_units',
sa.Column('occurred_start', TIMESTAMP(timezone=True), nullable=True)
)
op.add_column(
'memory_units',
sa.Column('occurred_end', TIMESTAMP(timezone=True), nullable=True)
)
op.add_column(
'memory_units',
sa.Column('mentioned_at', TIMESTAMP(timezone=True), nullable=True)
)
# Populate new columns from existing event_date for backward compatibility
# For existing facts, treat them as point events (start = end = event_date)
# and assume they were mentioned at the same time
op.execute("""
UPDATE memory_units
SET
occurred_start = event_date,
occurred_end = event_date,
mentioned_at = event_date
WHERE occurred_start IS NULL
""")
# Optional: Make columns non-nullable after populating
# Uncomment if you want to enforce NOT NULL constraint
# op.alter_column('memory_units', 'occurred_start', nullable=False)
# op.alter_column('memory_units', 'occurred_end', nullable=False)
# op.alter_column('memory_units', 'mentioned_at', nullable=False)
def downgrade() -> None:
"""Downgrade schema: remove temporal range columns from memory_units."""
# Remove the temporal range columns
op.drop_column('memory_units', 'mentioned_at')
op.drop_column('memory_units', 'occurred_end')
op.drop_column('memory_units', 'occurred_start')
@@ -0,0 +1,196 @@
"""Initial schema
Revision ID: af0413383b3e
Revises:
Create Date: 2025-11-03 14:31:53.245542
"""
from typing import Sequence, Union
from alembic import op
import sqlalchemy as sa
from sqlalchemy.dialects import postgresql
import pgvector.sqlalchemy
# revision identifiers, used by Alembic.
revision: str = 'af0413383b3e'
down_revision: Union[str, Sequence[str], None] = None
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None
def upgrade() -> None:
"""Upgrade schema."""
# Create pgvector extension
op.execute('CREATE EXTENSION IF NOT EXISTS vector')
op.execute('CREATE EXTENSION IF NOT EXISTS "uuid-ossp"')
# ### commands auto generated by Alembic - please adjust! ###
op.create_table('documents',
sa.Column('id', sa.Text(), nullable=False),
sa.Column('agent_id', sa.Text(), nullable=False),
sa.Column('original_text', sa.Text(), nullable=True),
sa.Column('content_hash', sa.Text(), nullable=True),
sa.Column('metadata', postgresql.JSONB(astext_type=sa.Text()),
server_default=sa.text("'{}'::jsonb"), nullable=False),
sa.Column('created_at', postgresql.TIMESTAMP(timezone=True), server_default=sa.text('now()'),
nullable=False),
sa.Column('updated_at', postgresql.TIMESTAMP(timezone=True), server_default=sa.text('now()'),
nullable=False),
sa.PrimaryKeyConstraint('id', 'agent_id')
)
op.create_index('idx_documents_agent_id', 'documents', ['agent_id'], unique=False)
op.create_index('idx_documents_content_hash', 'documents', ['content_hash'], unique=False)
op.create_table('entities',
sa.Column('id', sa.UUID(), server_default=sa.text('uuid_generate_v4()'), nullable=False),
sa.Column('canonical_name', sa.Text(), nullable=False),
sa.Column('entity_type', sa.Text(), nullable=False),
sa.Column('agent_id', sa.Text(), nullable=False),
sa.Column('metadata', postgresql.JSONB(astext_type=sa.Text()),
server_default=sa.text("'{}'::jsonb"), nullable=False),
sa.Column('first_seen', postgresql.TIMESTAMP(timezone=True), server_default=sa.text('now()'),
nullable=False),
sa.Column('last_seen', postgresql.TIMESTAMP(timezone=True), server_default=sa.text('now()'),
nullable=False),
sa.Column('mention_count', sa.Integer(), server_default='1', nullable=False),
sa.PrimaryKeyConstraint('id')
)
op.create_index('idx_entities_agent_id', 'entities', ['agent_id'], unique=False)
op.create_index('idx_entities_agent_name_type', 'entities', ['agent_id', 'canonical_name', 'entity_type'],
unique=False)
op.create_index('idx_entities_canonical_name', 'entities', ['canonical_name'], unique=False)
op.create_index('idx_entities_type', 'entities', ['entity_type'], unique=False)
op.create_table('entity_cooccurrences',
sa.Column('entity_id_1', sa.UUID(), nullable=False),
sa.Column('entity_id_2', sa.UUID(), nullable=False),
sa.Column('cooccurrence_count', sa.Integer(), server_default='1', nullable=False),
sa.Column('last_cooccurred', postgresql.TIMESTAMP(timezone=True), server_default=sa.text('now()'),
nullable=False),
sa.CheckConstraint('entity_id_1 < entity_id_2', name='entity_cooccurrence_order_check'),
sa.ForeignKeyConstraint(['entity_id_1'], ['entities.id'], ondelete='CASCADE'),
sa.ForeignKeyConstraint(['entity_id_2'], ['entities.id'], ondelete='CASCADE'),
sa.PrimaryKeyConstraint('entity_id_1', 'entity_id_2')
)
op.create_index('idx_entity_cooccurrences_count', 'entity_cooccurrences', ['cooccurrence_count'], unique=False,
postgresql_ops={'cooccurrence_count': 'DESC'})
op.create_index('idx_entity_cooccurrences_entity1', 'entity_cooccurrences', ['entity_id_1'], unique=False)
op.create_index('idx_entity_cooccurrences_entity2', 'entity_cooccurrences', ['entity_id_2'], unique=False)
op.create_table('memory_units',
sa.Column('id', sa.UUID(), server_default=sa.text('uuid_generate_v4()'), nullable=False),
sa.Column('agent_id', sa.Text(), nullable=False),
sa.Column('document_id', sa.Text(), nullable=True),
sa.Column('text', sa.Text(), nullable=False),
sa.Column('embedding', pgvector.sqlalchemy.vector.VECTOR(dim=384), nullable=True),
sa.Column('context', sa.Text(), nullable=True),
sa.Column('event_date', postgresql.TIMESTAMP(timezone=True), nullable=False),
sa.Column('fact_type', sa.Text(), server_default='world', nullable=False),
sa.Column('confidence_score', sa.Float(), nullable=True),
sa.Column('access_count', sa.Integer(), server_default='0', nullable=False),
sa.Column('created_at', postgresql.TIMESTAMP(timezone=True), server_default=sa.text('now()'),
nullable=False),
sa.Column('updated_at', postgresql.TIMESTAMP(timezone=True), server_default=sa.text('now()'),
nullable=False),
sa.CheckConstraint(
"(fact_type = 'opinion' AND confidence_score IS NOT NULL) OR (fact_type != 'opinion' AND confidence_score IS NULL)",
name='confidence_score_fact_type_check'),
sa.CheckConstraint("fact_type IN ('world', 'agent', 'opinion')"),
sa.CheckConstraint(
'confidence_score IS NULL OR (confidence_score >= 0.0 AND confidence_score <= 1.0)'),
sa.ForeignKeyConstraint(['document_id', 'agent_id'], ['documents.id', 'documents.agent_id'],
name='memory_units_document_fkey', ondelete='CASCADE'),
sa.PrimaryKeyConstraint('id')
)
op.create_index('idx_memory_units_access_count', 'memory_units', ['access_count'], unique=False,
postgresql_ops={'access_count': 'DESC'})
op.create_index('idx_memory_units_agent_date', 'memory_units', ['agent_id', 'event_date'], unique=False,
postgresql_ops={'event_date': 'DESC'})
op.create_index('idx_memory_units_agent_fact_type', 'memory_units', ['agent_id', 'fact_type'], unique=False)
op.create_index('idx_memory_units_agent_id', 'memory_units', ['agent_id'], unique=False)
op.create_index('idx_memory_units_agent_type_date', 'memory_units', ['agent_id', 'fact_type', 'event_date'],
unique=False, postgresql_ops={'event_date': 'DESC'})
op.create_index('idx_memory_units_document_id', 'memory_units', ['document_id'], unique=False)
op.create_index('idx_memory_units_embedding', 'memory_units', ['embedding'], unique=False, postgresql_using='hnsw',
postgresql_ops={'embedding': 'vector_cosine_ops'})
op.create_index('idx_memory_units_event_date', 'memory_units', ['event_date'], unique=False,
postgresql_ops={'event_date': 'DESC'})
op.create_index('idx_memory_units_fact_type', 'memory_units', ['fact_type'], unique=False)
op.create_index('idx_memory_units_opinion_confidence', 'memory_units', ['agent_id', 'confidence_score'],
unique=False, postgresql_where=sa.text("fact_type = 'opinion'"),
postgresql_ops={'confidence_score': 'DESC'})
op.create_index('idx_memory_units_opinion_date', 'memory_units', ['agent_id', 'event_date'], unique=False,
postgresql_where=sa.text("fact_type = 'opinion'"), postgresql_ops={'event_date': 'DESC'})
op.create_table('memory_links',
sa.Column('from_unit_id', sa.UUID(), nullable=False),
sa.Column('to_unit_id', sa.UUID(), nullable=False),
sa.Column('link_type', sa.Text(), nullable=False),
sa.Column('entity_id', sa.UUID(), nullable=False),
sa.Column('weight', sa.Float(), server_default='1.0', nullable=False),
sa.Column('created_at', postgresql.TIMESTAMP(timezone=True), server_default=sa.text('now()'),
nullable=False),
sa.ForeignKeyConstraint(['entity_id'], ['entities.id'], ondelete='CASCADE'),
sa.ForeignKeyConstraint(['from_unit_id'], ['memory_units.id'], ondelete='CASCADE'),
sa.ForeignKeyConstraint(['to_unit_id'], ['memory_units.id'], ondelete='CASCADE'),
sa.PrimaryKeyConstraint('from_unit_id', 'to_unit_id', 'link_type', 'entity_id')
)
op.create_index('idx_memory_links_entity', 'memory_links', ['entity_id'], unique=False,
postgresql_where=sa.text('entity_id IS NOT NULL'))
op.create_index('idx_memory_links_from', 'memory_links', ['from_unit_id'], unique=False)
op.create_index('idx_memory_links_from_weight', 'memory_links', ['from_unit_id', 'weight'], unique=False,
postgresql_where=sa.text('weight >= 0.1'), postgresql_ops={'weight': 'DESC'})
op.create_index('idx_memory_links_to', 'memory_links', ['to_unit_id'], unique=False)
op.create_index('idx_memory_links_type', 'memory_links', ['link_type'], unique=False)
op.create_table('unit_entities',
sa.Column('unit_id', sa.UUID(), nullable=False),
sa.Column('entity_id', sa.UUID(), nullable=False),
sa.ForeignKeyConstraint(['entity_id'], ['entities.id'], ondelete='CASCADE'),
sa.ForeignKeyConstraint(['unit_id'], ['memory_units.id'], ondelete='CASCADE'),
sa.PrimaryKeyConstraint('unit_id', 'entity_id')
)
op.create_index('idx_unit_entities_entity', 'unit_entities', ['entity_id'], unique=False)
op.create_index('idx_unit_entities_unit', 'unit_entities', ['unit_id'], unique=False)
# ### end Alembic commands ###
def downgrade() -> None:
"""Downgrade schema."""
# ### commands auto generated by Alembic - please adjust! ###
op.drop_index('idx_unit_entities_unit', table_name='unit_entities')
op.drop_index('idx_unit_entities_entity', table_name='unit_entities')
op.drop_table('unit_entities')
op.drop_index('idx_memory_links_type', table_name='memory_links')
op.drop_index('idx_memory_links_to', table_name='memory_links')
op.drop_index('idx_memory_links_from_weight', table_name='memory_links', postgresql_where=sa.text('weight >= 0.1'),
postgresql_ops={'weight': 'DESC'})
op.drop_index('idx_memory_links_from', table_name='memory_links')
op.drop_index('idx_memory_links_entity', table_name='memory_links',
postgresql_where=sa.text('entity_id IS NOT NULL'))
op.drop_table('memory_links')
op.drop_index('idx_memory_units_opinion_date', table_name='memory_units',
postgresql_where=sa.text("fact_type = 'opinion'"), postgresql_ops={'event_date': 'DESC'})
op.drop_index('idx_memory_units_opinion_confidence', table_name='memory_units',
postgresql_where=sa.text("fact_type = 'opinion'"), postgresql_ops={'confidence_score': 'DESC'})
op.drop_index('idx_memory_units_fact_type', table_name='memory_units')
op.drop_index('idx_memory_units_event_date', table_name='memory_units', postgresql_ops={'event_date': 'DESC'})
op.drop_index('idx_memory_units_embedding', table_name='memory_units', postgresql_using='hnsw',
postgresql_ops={'embedding': 'vector_cosine_ops'})
op.drop_index('idx_memory_units_document_id', table_name='memory_units')
op.drop_index('idx_memory_units_agent_type_date', table_name='memory_units', postgresql_ops={'event_date': 'DESC'})
op.drop_index('idx_memory_units_agent_id', table_name='memory_units')
op.drop_index('idx_memory_units_agent_fact_type', table_name='memory_units')
op.drop_index('idx_memory_units_agent_date', table_name='memory_units', postgresql_ops={'event_date': 'DESC'})
op.drop_index('idx_memory_units_access_count', table_name='memory_units', postgresql_ops={'access_count': 'DESC'})
op.drop_table('memory_units')
op.drop_index('idx_entity_cooccurrences_entity2', table_name='entity_cooccurrences')
op.drop_index('idx_entity_cooccurrences_entity1', table_name='entity_cooccurrences')
op.drop_index('idx_entity_cooccurrences_count', table_name='entity_cooccurrences',
postgresql_ops={'cooccurrence_count': 'DESC'})
op.drop_table('entity_cooccurrences')
op.drop_index('idx_entities_type', table_name='entities')
op.drop_index('idx_entities_canonical_name', table_name='entities')
op.drop_index('idx_entities_agent_name_type', table_name='entities')
op.drop_index('idx_entities_agent_id', table_name='entities')
op.drop_table('entities')
op.drop_index('idx_documents_content_hash', table_name='documents')
op.drop_index('idx_documents_agent_id', table_name='documents')
op.drop_table('documents')
# ### end Alembic commands ###
+4 -12
View File
@@ -4,7 +4,7 @@ Memory System for AI Agents.
Temporal + Semantic Memory Architecture using PostgreSQL with pgvector.
"""
from .engine.memory_engine import MemoryEngine
from .engine.search.trace import (
from .engine.search_trace import (
SearchTrace,
QueryInfo,
EntryPoint,
@@ -15,16 +15,12 @@ from .engine.search.trace import (
SearchSummary,
SearchPhaseMetrics,
)
from .engine.search.tracer import SearchTracer
from .engine.embeddings import Embeddings, LocalSTEmbeddings, RemoteTEIEmbeddings
from .engine.cross_encoder import CrossEncoderModel, LocalSTCrossEncoder, RemoteTEICrossEncoder
from .engine.search_tracer import SearchTracer
from .engine.embeddings import Embeddings, SentenceTransformersEmbeddings
from .engine.llm_wrapper import LLMConfig
from .config import HindsightConfig, get_config
__all__ = [
"MemoryEngine",
"HindsightConfig",
"get_config",
"SearchTrace",
"SearchTracer",
"QueryInfo",
@@ -36,11 +32,7 @@ __all__ = [
"SearchSummary",
"SearchPhaseMetrics",
"Embeddings",
"LocalSTEmbeddings",
"RemoteTEIEmbeddings",
"CrossEncoderModel",
"LocalSTCrossEncoder",
"RemoteTEICrossEncoder",
"SentenceTransformersEmbeddings",
"LLMConfig",
]
__version__ = "0.1.0"
@@ -1,274 +0,0 @@
"""initial_schema
Revision ID: 5a366d414dce
Revises:
Create Date: 2025-11-27 11:54:19.228030
"""
from typing import Sequence, Union
from alembic import op
import sqlalchemy as sa
from sqlalchemy.dialects import postgresql
from pgvector.sqlalchemy import Vector
# revision identifiers, used by Alembic.
revision: str = '5a366d414dce'
down_revision: Union[str, Sequence[str], None] = None
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None
def upgrade() -> None:
"""Upgrade schema - create all tables from scratch."""
# Enable required extensions
op.execute('CREATE EXTENSION IF NOT EXISTS vector')
# Create banks table
op.create_table(
'banks',
sa.Column('bank_id', sa.Text(), nullable=False),
sa.Column('name', sa.Text(), nullable=True),
sa.Column('personality', postgresql.JSONB(astext_type=sa.Text()), server_default=sa.text("'{}'::jsonb"), nullable=False),
sa.Column('background', sa.Text(), nullable=True),
sa.Column('created_at', postgresql.TIMESTAMP(timezone=True), server_default=sa.text('now()'), nullable=False),
sa.Column('updated_at', postgresql.TIMESTAMP(timezone=True), server_default=sa.text('now()'), nullable=False),
sa.PrimaryKeyConstraint('bank_id', name=op.f('pk_banks'))
)
# Create documents table
op.create_table(
'documents',
sa.Column('id', sa.Text(), nullable=False),
sa.Column('bank_id', sa.Text(), nullable=False),
sa.Column('original_text', sa.Text(), nullable=True),
sa.Column('content_hash', sa.Text(), nullable=True),
sa.Column('metadata', postgresql.JSONB(astext_type=sa.Text()), server_default=sa.text("'{}'::jsonb"), nullable=False),
sa.Column('created_at', postgresql.TIMESTAMP(timezone=True), server_default=sa.text('now()'), nullable=False),
sa.Column('updated_at', postgresql.TIMESTAMP(timezone=True), server_default=sa.text('now()'), nullable=False),
sa.PrimaryKeyConstraint('id', 'bank_id', name=op.f('pk_documents'))
)
op.create_index('idx_documents_bank_id', 'documents', ['bank_id'])
op.create_index('idx_documents_content_hash', 'documents', ['content_hash'])
# Create async_operations table
op.create_table(
'async_operations',
sa.Column('operation_id', postgresql.UUID(as_uuid=True), server_default=sa.text('gen_random_uuid()'), nullable=False),
sa.Column('bank_id', sa.Text(), nullable=False),
sa.Column('operation_type', sa.Text(), nullable=False),
sa.Column('status', sa.Text(), server_default='pending', nullable=False),
sa.Column('created_at', postgresql.TIMESTAMP(timezone=True), server_default=sa.text('now()'), nullable=False),
sa.Column('updated_at', postgresql.TIMESTAMP(timezone=True), server_default=sa.text('now()'), nullable=False),
sa.Column('completed_at', postgresql.TIMESTAMP(timezone=True), nullable=True),
sa.Column('error_message', sa.Text(), nullable=True),
sa.Column('result_metadata', postgresql.JSONB(astext_type=sa.Text()), server_default=sa.text("'{}'::jsonb"), nullable=False),
sa.PrimaryKeyConstraint('operation_id', name=op.f('pk_async_operations')),
sa.CheckConstraint("status IN ('pending', 'processing', 'completed', 'failed')", name='async_operations_status_check')
)
op.create_index('idx_async_operations_bank_id', 'async_operations', ['bank_id'])
op.create_index('idx_async_operations_status', 'async_operations', ['status'])
op.create_index('idx_async_operations_bank_status', 'async_operations', ['bank_id', 'status'])
# Create entities table
op.create_table(
'entities',
sa.Column('id', postgresql.UUID(as_uuid=True), server_default=sa.text('gen_random_uuid()'), nullable=False),
sa.Column('canonical_name', sa.Text(), nullable=False),
sa.Column('bank_id', sa.Text(), nullable=False),
sa.Column('metadata', postgresql.JSONB(astext_type=sa.Text()), server_default=sa.text("'{}'::jsonb"), nullable=False),
sa.Column('first_seen', postgresql.TIMESTAMP(timezone=True), server_default=sa.text('now()'), nullable=False),
sa.Column('last_seen', postgresql.TIMESTAMP(timezone=True), server_default=sa.text('now()'), nullable=False),
sa.Column('mention_count', sa.Integer(), server_default='1', nullable=False),
sa.PrimaryKeyConstraint('id', name=op.f('pk_entities'))
)
op.create_index('idx_entities_bank_id', 'entities', ['bank_id'])
op.create_index('idx_entities_canonical_name', 'entities', ['canonical_name'])
op.create_index('idx_entities_bank_name', 'entities', ['bank_id', 'canonical_name'])
# Create unique index on (bank_id, LOWER(canonical_name)) for entity resolution
op.execute('CREATE UNIQUE INDEX idx_entities_bank_lower_name ON entities (bank_id, LOWER(canonical_name))')
# Create memory_units table
op.create_table(
'memory_units',
sa.Column('id', postgresql.UUID(as_uuid=True), server_default=sa.text('gen_random_uuid()'), nullable=False),
sa.Column('bank_id', sa.Text(), nullable=False),
sa.Column('document_id', sa.Text(), nullable=True),
sa.Column('text', sa.Text(), nullable=False),
sa.Column('embedding', Vector(384), nullable=True),
sa.Column('context', sa.Text(), nullable=True),
sa.Column('event_date', postgresql.TIMESTAMP(timezone=True), nullable=False),
sa.Column('occurred_start', postgresql.TIMESTAMP(timezone=True), nullable=True),
sa.Column('occurred_end', postgresql.TIMESTAMP(timezone=True), nullable=True),
sa.Column('mentioned_at', postgresql.TIMESTAMP(timezone=True), nullable=True),
sa.Column('fact_type', sa.Text(), server_default='world', nullable=False),
sa.Column('confidence_score', sa.Float(), nullable=True),
sa.Column('access_count', sa.Integer(), server_default='0', nullable=False),
sa.Column('metadata', postgresql.JSONB(astext_type=sa.Text()), server_default=sa.text("'{}'::jsonb"), nullable=False),
sa.Column('created_at', postgresql.TIMESTAMP(timezone=True), server_default=sa.text('now()'), nullable=False),
sa.Column('updated_at', postgresql.TIMESTAMP(timezone=True), server_default=sa.text('now()'), nullable=False),
sa.ForeignKeyConstraint(['document_id', 'bank_id'], ['documents.id', 'documents.bank_id'], name='memory_units_document_fkey', ondelete='CASCADE'),
sa.PrimaryKeyConstraint('id', name=op.f('pk_memory_units')),
sa.CheckConstraint("fact_type IN ('world', 'bank', 'opinion', 'observation')", name='memory_units_fact_type_check'),
sa.CheckConstraint("confidence_score IS NULL OR (confidence_score >= 0.0 AND confidence_score <= 1.0)", name='memory_units_confidence_range_check'),
sa.CheckConstraint(
"(fact_type = 'opinion' AND confidence_score IS NOT NULL) OR "
"(fact_type = 'observation') OR "
"(fact_type NOT IN ('opinion', 'observation') AND confidence_score IS NULL)",
name='confidence_score_fact_type_check'
)
)
# Add search_vector column for full-text search
op.execute("""
ALTER TABLE memory_units
ADD COLUMN search_vector tsvector
GENERATED ALWAYS AS (to_tsvector('english', COALESCE(text, '') || ' ' || COALESCE(context, ''))) STORED
""")
op.create_index('idx_memory_units_bank_id', 'memory_units', ['bank_id'])
op.create_index('idx_memory_units_document_id', 'memory_units', ['document_id'])
op.create_index('idx_memory_units_event_date', 'memory_units', [sa.text('event_date DESC')])
op.create_index('idx_memory_units_bank_date', 'memory_units', ['bank_id', sa.text('event_date DESC')])
op.create_index('idx_memory_units_access_count', 'memory_units', [sa.text('access_count DESC')])
op.create_index('idx_memory_units_fact_type', 'memory_units', ['fact_type'])
op.create_index('idx_memory_units_bank_fact_type', 'memory_units', ['bank_id', 'fact_type'])
op.create_index('idx_memory_units_bank_type_date', 'memory_units', ['bank_id', 'fact_type', sa.text('event_date DESC')])
op.create_index('idx_memory_units_opinion_confidence', 'memory_units', ['bank_id', sa.text('confidence_score DESC')], postgresql_where=sa.text("fact_type = 'opinion'"))
op.create_index('idx_memory_units_opinion_date', 'memory_units', ['bank_id', sa.text('event_date DESC')], postgresql_where=sa.text("fact_type = 'opinion'"))
op.create_index('idx_memory_units_observation_date', 'memory_units', ['bank_id', sa.text('event_date DESC')], postgresql_where=sa.text("fact_type = 'observation'"))
op.create_index('idx_memory_units_embedding', 'memory_units', ['embedding'], postgresql_using='hnsw', postgresql_ops={'embedding': 'vector_cosine_ops'})
# Create BM25 full-text search index on search_vector
op.execute("""
CREATE INDEX idx_memory_units_text_search ON memory_units
USING gin(search_vector)
""")
op.execute("""
CREATE MATERIALIZED VIEW memory_units_bm25 AS
SELECT
id,
bank_id,
text,
to_tsvector('english', text) AS text_vector,
log(1.0 + length(text)::float / (SELECT avg(length(text)) FROM memory_units)) AS doc_length_factor
FROM memory_units
""")
op.create_index('idx_memory_units_bm25_bank', 'memory_units_bm25', ['bank_id'])
op.create_index('idx_memory_units_bm25_text_vector', 'memory_units_bm25', ['text_vector'], postgresql_using='gin')
# Create entity_cooccurrences table
op.create_table(
'entity_cooccurrences',
sa.Column('entity_id_1', postgresql.UUID(as_uuid=True), nullable=False),
sa.Column('entity_id_2', postgresql.UUID(as_uuid=True), nullable=False),
sa.Column('cooccurrence_count', sa.Integer(), server_default='1', nullable=False),
sa.Column('last_cooccurred', postgresql.TIMESTAMP(timezone=True), server_default=sa.text('now()'), nullable=False),
sa.ForeignKeyConstraint(['entity_id_1'], ['entities.id'], name=op.f('fk_entity_cooccurrences_entity_id_1_entities'), ondelete='CASCADE'),
sa.ForeignKeyConstraint(['entity_id_2'], ['entities.id'], name=op.f('fk_entity_cooccurrences_entity_id_2_entities'), ondelete='CASCADE'),
sa.PrimaryKeyConstraint('entity_id_1', 'entity_id_2', name=op.f('pk_entity_cooccurrences')),
sa.CheckConstraint('entity_id_1 < entity_id_2', name='entity_cooccurrence_order_check')
)
op.create_index('idx_entity_cooccurrences_entity1', 'entity_cooccurrences', ['entity_id_1'])
op.create_index('idx_entity_cooccurrences_entity2', 'entity_cooccurrences', ['entity_id_2'])
op.create_index('idx_entity_cooccurrences_count', 'entity_cooccurrences', [sa.text('cooccurrence_count DESC')])
# Create memory_links table
op.create_table(
'memory_links',
sa.Column('from_unit_id', postgresql.UUID(as_uuid=True), nullable=False),
sa.Column('to_unit_id', postgresql.UUID(as_uuid=True), nullable=False),
sa.Column('link_type', sa.Text(), nullable=False),
sa.Column('entity_id', postgresql.UUID(as_uuid=True), nullable=True),
sa.Column('weight', sa.Float(), server_default='1.0', nullable=False),
sa.Column('created_at', postgresql.TIMESTAMP(timezone=True), server_default=sa.text('now()'), nullable=False),
sa.ForeignKeyConstraint(['entity_id'], ['entities.id'], name=op.f('fk_memory_links_entity_id_entities'), ondelete='CASCADE'),
sa.ForeignKeyConstraint(['from_unit_id'], ['memory_units.id'], name=op.f('fk_memory_links_from_unit_id_memory_units'), ondelete='CASCADE'),
sa.ForeignKeyConstraint(['to_unit_id'], ['memory_units.id'], name=op.f('fk_memory_links_to_unit_id_memory_units'), ondelete='CASCADE'),
sa.CheckConstraint("link_type IN ('temporal', 'semantic', 'entity', 'causes', 'caused_by', 'enables', 'prevents')", name='memory_links_link_type_check'),
sa.CheckConstraint('weight >= 0.0 AND weight <= 1.0', name='memory_links_weight_check')
)
# Create unique constraint using COALESCE for nullable entity_id
op.execute("CREATE UNIQUE INDEX idx_memory_links_unique ON memory_links (from_unit_id, to_unit_id, link_type, COALESCE(entity_id, '00000000-0000-0000-0000-000000000000'::uuid))")
op.create_index('idx_memory_links_from_unit', 'memory_links', ['from_unit_id'])
op.create_index('idx_memory_links_to_unit', 'memory_links', ['to_unit_id'])
op.create_index('idx_memory_links_entity', 'memory_links', ['entity_id'])
op.create_index('idx_memory_links_link_type', 'memory_links', ['link_type'])
# Create unit_entities table
op.create_table(
'unit_entities',
sa.Column('unit_id', postgresql.UUID(as_uuid=True), nullable=False),
sa.Column('entity_id', postgresql.UUID(as_uuid=True), nullable=False),
sa.ForeignKeyConstraint(['entity_id'], ['entities.id'], name=op.f('fk_unit_entities_entity_id_entities'), ondelete='CASCADE'),
sa.ForeignKeyConstraint(['unit_id'], ['memory_units.id'], name=op.f('fk_unit_entities_unit_id_memory_units'), ondelete='CASCADE'),
sa.PrimaryKeyConstraint('unit_id', 'entity_id', name=op.f('pk_unit_entities'))
)
op.create_index('idx_unit_entities_unit', 'unit_entities', ['unit_id'])
op.create_index('idx_unit_entities_entity', 'unit_entities', ['entity_id'])
def downgrade() -> None:
"""Downgrade schema - drop all tables."""
# Drop tables in reverse dependency order
op.drop_index('idx_unit_entities_entity', table_name='unit_entities')
op.drop_index('idx_unit_entities_unit', table_name='unit_entities')
op.drop_table('unit_entities')
op.drop_index('idx_memory_links_link_type', table_name='memory_links')
op.drop_index('idx_memory_links_entity', table_name='memory_links')
op.drop_index('idx_memory_links_to_unit', table_name='memory_links')
op.drop_index('idx_memory_links_from_unit', table_name='memory_links')
op.execute('DROP INDEX IF EXISTS idx_memory_links_unique')
op.drop_table('memory_links')
op.drop_index('idx_entity_cooccurrences_count', table_name='entity_cooccurrences')
op.drop_index('idx_entity_cooccurrences_entity2', table_name='entity_cooccurrences')
op.drop_index('idx_entity_cooccurrences_entity1', table_name='entity_cooccurrences')
op.drop_table('entity_cooccurrences')
# Drop BM25 materialized view and index
op.drop_index('idx_memory_units_bm25_text_vector', table_name='memory_units_bm25')
op.drop_index('idx_memory_units_bm25_bank', table_name='memory_units_bm25')
op.execute('DROP MATERIALIZED VIEW IF EXISTS memory_units_bm25')
op.drop_index('idx_memory_units_embedding', table_name='memory_units')
op.drop_index('idx_memory_units_observation_date', table_name='memory_units')
op.drop_index('idx_memory_units_opinion_date', table_name='memory_units')
op.drop_index('idx_memory_units_opinion_confidence', table_name='memory_units')
op.drop_index('idx_memory_units_bank_type_date', table_name='memory_units')
op.drop_index('idx_memory_units_bank_fact_type', table_name='memory_units')
op.drop_index('idx_memory_units_fact_type', table_name='memory_units')
op.drop_index('idx_memory_units_access_count', table_name='memory_units')
op.drop_index('idx_memory_units_bank_date', table_name='memory_units')
op.drop_index('idx_memory_units_event_date', table_name='memory_units')
op.drop_index('idx_memory_units_document_id', table_name='memory_units')
op.drop_index('idx_memory_units_bank_id', table_name='memory_units')
op.execute('DROP INDEX IF EXISTS idx_memory_units_text_search')
op.drop_table('memory_units')
op.execute('DROP INDEX IF EXISTS idx_entities_bank_lower_name')
op.drop_index('idx_entities_bank_name', table_name='entities')
op.drop_index('idx_entities_canonical_name', table_name='entities')
op.drop_index('idx_entities_bank_id', table_name='entities')
op.drop_table('entities')
op.drop_index('idx_async_operations_bank_status', table_name='async_operations')
op.drop_index('idx_async_operations_status', table_name='async_operations')
op.drop_index('idx_async_operations_bank_id', table_name='async_operations')
op.drop_table('async_operations')
op.drop_index('idx_documents_content_hash', table_name='documents')
op.drop_index('idx_documents_bank_id', table_name='documents')
op.drop_table('documents')
op.drop_table('banks')
# Drop extensions (optional - comment out if you want to keep them)
# op.execute('DROP EXTENSION IF EXISTS vector')
# op.execute('DROP EXTENSION IF EXISTS "uuid-ossp"')
@@ -1,70 +0,0 @@
"""add_chunks_table
Revision ID: b7c4d8e9f1a2
Revises: 5a366d414dce
Create Date: 2025-11-28 00:00:00.000000
"""
from typing import Sequence, Union
from alembic import op
import sqlalchemy as sa
from sqlalchemy.dialects import postgresql
# revision identifiers, used by Alembic.
revision: str = 'b7c4d8e9f1a2'
down_revision: Union[str, Sequence[str], None] = '5a366d414dce'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None
def upgrade() -> None:
"""Add chunks table and link memory_units to chunks."""
# Create chunks table with single text PK (bank_id_document_id_chunk_index)
op.create_table(
'chunks',
sa.Column('chunk_id', sa.Text(), nullable=False),
sa.Column('document_id', sa.Text(), nullable=False),
sa.Column('bank_id', sa.Text(), nullable=False),
sa.Column('chunk_index', sa.Integer(), nullable=False),
sa.Column('chunk_text', sa.Text(), nullable=False),
sa.Column('created_at', postgresql.TIMESTAMP(timezone=True), server_default=sa.text('now()'), nullable=False),
sa.ForeignKeyConstraint(['document_id', 'bank_id'], ['documents.id', 'documents.bank_id'], name='chunks_document_fkey', ondelete='CASCADE'),
sa.PrimaryKeyConstraint('chunk_id', name=op.f('pk_chunks'))
)
# Add indexes for efficient queries
op.create_index('idx_chunks_document_id', 'chunks', ['document_id'])
op.create_index('idx_chunks_bank_id', 'chunks', ['bank_id'])
# Add chunk_id column to memory_units (nullable, as existing records won't have chunks)
op.add_column('memory_units', sa.Column('chunk_id', sa.Text(), nullable=True))
# Add foreign key constraint to chunks table
op.create_foreign_key(
'memory_units_chunk_fkey',
'memory_units',
'chunks',
['chunk_id'],
['chunk_id'],
ondelete='SET NULL'
)
# Add index on chunk_id for efficient lookups
op.create_index('idx_memory_units_chunk_id', 'memory_units', ['chunk_id'])
def downgrade() -> None:
"""Remove chunks table and chunk_id from memory_units."""
# Drop index and foreign key from memory_units
op.drop_index('idx_memory_units_chunk_id', table_name='memory_units')
op.drop_constraint('memory_units_chunk_fkey', 'memory_units', type_='foreignkey')
op.drop_column('memory_units', 'chunk_id')
# Drop chunks table indexes and table
op.drop_index('idx_chunks_bank_id', table_name='chunks')
op.drop_index('idx_chunks_document_id', table_name='chunks')
op.drop_table('chunks')
@@ -1,39 +0,0 @@
"""add_retain_params_to_documents
Revision ID: c8e5f2a3b4d1
Revises: b7c4d8e9f1a2
Create Date: 2025-12-02 00:00:00.000000
"""
from typing import Sequence, Union
from alembic import op
import sqlalchemy as sa
from sqlalchemy.dialects import postgresql
# revision identifiers, used by Alembic.
revision: str = 'c8e5f2a3b4d1'
down_revision: Union[str, Sequence[str], None] = 'b7c4d8e9f1a2'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None
def upgrade() -> None:
"""Add retain_params JSONB column to documents table."""
# Add retain_params column to store parameters passed during retain
op.add_column('documents', sa.Column('retain_params', postgresql.JSONB(), nullable=True))
# Add index for efficient queries on retain_params
op.create_index('idx_documents_retain_params', 'documents', ['retain_params'], postgresql_using='gin')
def downgrade() -> None:
"""Remove retain_params column from documents table."""
# Drop index
op.drop_index('idx_documents_retain_params', table_name='documents')
# Drop column
op.drop_column('documents', 'retain_params')
@@ -1,48 +0,0 @@
"""Rename fact_type 'bank' to 'experience'
Revision ID: d9f6a3b4c5e2
Revises: c8e5f2a3b4d1
Create Date: 2024-12-04 15:00:00.000000
"""
from alembic import op
import sqlalchemy as sa
# revision identifiers, used by Alembic.
revision = 'd9f6a3b4c5e2'
down_revision = 'c8e5f2a3b4d1'
branch_labels = None
depends_on = None
def upgrade():
# Drop old check constraint FIRST (before updating data)
op.drop_constraint('memory_units_fact_type_check', 'memory_units', type_='check')
# Update existing 'bank' values to 'experience'
op.execute("UPDATE memory_units SET fact_type = 'experience' WHERE fact_type = 'bank'")
# Also update any 'interactions' values (in case of partial migration)
op.execute("UPDATE memory_units SET fact_type = 'experience' WHERE fact_type = 'interactions'")
# Create new check constraint with 'experience' instead of 'bank'
op.create_check_constraint(
'memory_units_fact_type_check',
'memory_units',
"fact_type IN ('world', 'experience', 'opinion', 'observation')"
)
def downgrade():
# Drop new check constraint FIRST
op.drop_constraint('memory_units_fact_type_check', 'memory_units', type_='check')
# Update 'experience' back to 'bank'
op.execute("UPDATE memory_units SET fact_type = 'bank' WHERE fact_type = 'experience'")
# Recreate old check constraint
op.create_check_constraint(
'memory_units_fact_type_check',
'memory_units',
"fact_type IN ('world', 'bank', 'opinion', 'observation')"
)
@@ -1,62 +0,0 @@
"""disposition_to_3_traits
Revision ID: e0a1b2c3d4e5
Revises: rename_personality
Create Date: 2024-12-08
Migrate disposition traits from Big Five (openness, conscientiousness, extraversion,
agreeableness, neuroticism, bias_strength with 0-1 float values) to the new 3-trait
system (skepticism, literalism, empathy with 1-5 integer values).
"""
from typing import Sequence, Union
from alembic import op
import sqlalchemy as sa
# revision identifiers, used by Alembic.
revision: str = 'e0a1b2c3d4e5'
down_revision: Union[str, Sequence[str], None] = 'rename_personality'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None
def upgrade() -> None:
"""Convert Big Five disposition to 3-trait disposition."""
conn = op.get_bind()
# Update all existing banks to use the new disposition format
# Convert from old format to new format with reasonable mappings:
# - skepticism: derived from inverse of agreeableness (skeptical people are less agreeable)
# - literalism: derived from conscientiousness (detail-oriented people are more literal)
# - empathy: derived from agreeableness + inverse of neuroticism
# Default all to 3 (neutral) for simplicity
conn.execute(sa.text("""
UPDATE banks
SET disposition = '{"skepticism": 3, "literalism": 3, "empathy": 3}'::jsonb
WHERE disposition IS NOT NULL
"""))
# Update the default for new banks
conn.execute(sa.text("""
ALTER TABLE banks
ALTER COLUMN disposition SET DEFAULT '{"skepticism": 3, "literalism": 3, "empathy": 3}'::jsonb
"""))
def downgrade() -> None:
"""Convert back to Big Five disposition."""
conn = op.get_bind()
# Revert to Big Five format with default values
conn.execute(sa.text("""
UPDATE banks
SET disposition = '{"openness": 0.5, "conscientiousness": 0.5, "extraversion": 0.5, "agreeableness": 0.5, "neuroticism": 0.5, "bias_strength": 0.5}'::jsonb
WHERE disposition IS NOT NULL
"""))
# Update the default for new banks
conn.execute(sa.text("""
ALTER TABLE banks
ALTER COLUMN disposition SET DEFAULT '{"openness": 0.5, "conscientiousness": 0.5, "extraversion": 0.5, "agreeableness": 0.5, "neuroticism": 0.5, "bias_strength": 0.5}'::jsonb
"""))
@@ -1,65 +0,0 @@
"""rename_personality_to_disposition
Revision ID: rename_personality
Revises: d9f6a3b4c5e2
Create Date: 2024-12-04
"""
from typing import Sequence, Union
from alembic import op
import sqlalchemy as sa
from sqlalchemy.dialects import postgresql
# revision identifiers, used by Alembic.
revision: str = 'rename_personality'
down_revision: Union[str, Sequence[str], None] = 'd9f6a3b4c5e2'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None
def upgrade() -> None:
"""Rename personality column to disposition in banks table (if it exists)."""
conn = op.get_bind()
# Check if 'personality' column exists (old database)
result = conn.execute(sa.text("""
SELECT column_name
FROM information_schema.columns
WHERE table_name = 'banks' AND column_name = 'personality'
"""))
has_personality = result.fetchone() is not None
# Check if 'disposition' column exists (new database)
result = conn.execute(sa.text("""
SELECT column_name
FROM information_schema.columns
WHERE table_name = 'banks' AND column_name = 'disposition'
"""))
has_disposition = result.fetchone() is not None
if has_personality and not has_disposition:
# Old database: rename personality -> disposition
op.alter_column('banks', 'personality', new_column_name='disposition')
elif not has_personality and not has_disposition:
# Neither exists (shouldn't happen, but be safe): add disposition column
op.add_column('banks', sa.Column(
'disposition',
postgresql.JSONB(astext_type=sa.Text()),
server_default=sa.text("'{}'::jsonb"),
nullable=False
))
# else: disposition already exists, nothing to do
def downgrade() -> None:
"""Revert disposition column back to personality."""
conn = op.get_bind()
result = conn.execute(sa.text("""
SELECT column_name
FROM information_schema.columns
WHERE table_name = 'banks' AND column_name = 'disposition'
"""))
if result.fetchone():
op.alter_column('banks', 'disposition', new_column_name='personality')
+27 -24
View File
@@ -17,17 +17,18 @@ def create_app(
http_api_enabled: bool = True,
mcp_api_enabled: bool = False,
mcp_mount_path: str = "/mcp",
run_migrations: bool = True,
initialize_memory: bool = True
) -> FastAPI:
"""
Create and configure the unified Hindsight API application.
Args:
memory: MemoryEngine instance (already initialized with required parameters).
Migrations are controlled by the MemoryEngine's run_migrations parameter.
memory: MemoryEngine instance (already initialized with required parameters)
http_api_enabled: Whether to enable HTTP REST API endpoints (default: True)
mcp_api_enabled: Whether to enable MCP server (default: False)
mcp_mount_path: Path to mount MCP server (default: /mcp)
run_migrations: Whether to run database migrations on startup (default: True)
initialize_memory: Whether to initialize memory system on startup (default: True)
Returns:
@@ -49,6 +50,7 @@ def create_app(
from .http import create_app as create_http_app
app = create_http_app(
memory=memory,
run_migrations=run_migrations,
initialize_memory=initialize_memory
)
logger.info("HTTP REST API enabled")
@@ -60,13 +62,14 @@ def create_app(
# Mount MCP server if enabled
if mcp_api_enabled:
try:
from .mcp import create_mcp_app
from .mcp import create_mcp_server
# Create MCP app with dynamic bank_id support
# Supports: /mcp/{bank_id}/sse (bank-specific SSE endpoint)
mcp_app = create_mcp_app(memory=memory)
app.mount(mcp_mount_path, mcp_app)
logger.info(f"MCP server enabled at {mcp_mount_path}/{{bank_id}}/sse")
# Create MCP server with shared memory instance
mcp_server = create_mcp_server(memory=memory)
# Mount at specified path using http_app (modern non-SSE alternative)
app.mount(mcp_mount_path, mcp_server.http_app())
logger.info(f"MCP server enabled at {mcp_mount_path}")
except ImportError as e:
logger.error(f"MCP server requested but dependencies not available: {e}")
logger.error("Install with: pip install hindsight-api[mcp]")
@@ -77,26 +80,26 @@ def create_app(
# Re-export commonly used items for backwards compatibility
from .http import (
RecallRequest,
RecallResult,
RecallResponse,
SearchRequest,
SearchResult,
SearchResponse,
MemoryItem,
RetainRequest,
ReflectRequest,
ReflectResponse,
CreateBankRequest,
DispositionTraits,
BatchPutRequest,
ThinkRequest,
ThinkResponse,
CreateAgentRequest,
PersonalityTraits,
)
__all__ = [
"create_app",
"RecallRequest",
"RecallResult",
"RecallResponse",
"SearchRequest",
"SearchResult",
"SearchResponse",
"MemoryItem",
"RetainRequest",
"ReflectRequest",
"ReflectResponse",
"CreateBankRequest",
"DispositionTraits",
"BatchPutRequest",
"ThinkRequest",
"ThinkResponse",
"CreateAgentRequest",
"PersonalityTraits",
]
File diff suppressed because it is too large Load Diff
+84 -147
View File
@@ -2,32 +2,13 @@
import json
import logging
import os
from contextvars import ContextVar
from typing import Optional
from fastmcp import FastMCP
from hindsight_api import MemoryEngine
from hindsight_api.engine.response_models import VALID_RECALL_FACT_TYPES
# Configure logging from HINDSIGHT_API_LOG_LEVEL environment variable
_log_level_str = os.environ.get("HINDSIGHT_API_LOG_LEVEL", "info").lower()
_log_level_map = {"critical": logging.CRITICAL, "error": logging.ERROR, "warning": logging.WARNING,
"info": logging.INFO, "debug": logging.DEBUG, "trace": logging.DEBUG}
logging.basicConfig(
level=_log_level_map.get(_log_level_str, logging.INFO),
format="%(asctime)s - %(levelname)s - %(name)s - %(message)s"
)
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
# Context variable to hold the current bank_id from the URL path
_current_bank_id: ContextVar[Optional[str]] = ContextVar("current_bank_id", default=None)
def get_current_bank_id() -> Optional[str]:
"""Get the current bank_id from context (set from URL path)."""
return _current_bank_id.get()
def create_mcp_server(memory: MemoryEngine) -> FastMCP:
"""
@@ -39,71 +20,124 @@ def create_mcp_server(memory: MemoryEngine) -> FastMCP:
Returns:
Configured FastMCP server instance
"""
# Create FastMCP server
mcp = FastMCP("hindsight-mcp-server")
@mcp.tool()
async def retain(content: str, context: str = "general") -> str:
async def hindsight_put(agent_id: str, content: str, context: str, explanation: str = "") -> str:
"""
Store important information to long-term memory.
**CRITICAL: Store important user information to long-term memory.**
**⚠️ PER-USER TOOL - REQUIRES USER IDENTIFICATION:**
- This tool is STRICTLY per-user. Each user MUST have a unique `agent_id`.
- ONLY use this tool if you have a valid user identifier (user ID, email, session ID, etc.) to map to `agent_id`.
- DO NOT use this tool if you cannot identify the specific user.
- DO NOT share memories between different users - each user's memories are isolated by their `agent_id`.
- If you don't have a user identifier, DO NOT use this tool at all.
Use this tool PROACTIVELY whenever the user shares:
- Personal facts, preferences, or interests
- Important events or milestones
- User history, experiences, or background
- Decisions, opinions, or stated preferences
- Goals, plans, or future intentions
- Relationships or people mentioned
- Personal facts, preferences, or interests (e.g., "I love hiking", "I'm a vegetarian")
- Important events or milestones (e.g., "I got promoted", "My birthday is June 15")
- User history, experiences, or background (e.g., "I used to work at Google", "I studied CS at MIT")
- Decisions, opinions, or stated preferences (e.g., "I prefer Python over JavaScript")
- Goals, plans, or future intentions (e.g., "I'm planning to visit Japan next year")
- Relationships or people mentioned (e.g., "My manager Sarah", "My wife Alice")
- Work context, projects, or responsibilities
- Any other information the user would want remembered for future conversations
**When to use**: Immediately after user shares personal information. Don't ask permission - just store it naturally.
**Context guidelines**: Use descriptive contexts like "personal_preferences", "work_history", "family", "hobbies",
"career_goals", "project_details", etc. This helps organize and retrieve related memories later.
Args:
agent_id: **REQUIRED** - The unique, persistent identifier for this specific user (e.g., user_id, email, session_id).
This MUST be consistent across all interactions with the same user.
Example: "user_12345", "[email protected]", "session_abc123"
content: The fact/memory to store (be specific and include relevant details)
context: Category for the memory (e.g., 'preferences', 'work', 'hobbies', 'family'). Default: 'general'
context: Categorize the memory (e.g., 'personal_preferences', 'work_history', 'hobbies', 'family')
explanation: Optional explanation for why this memory is being stored
"""
try:
bank_id = get_current_bank_id()
# Log explanation if provided
if explanation:
logger.debug(f"Explanation: {explanation}")
# Store memory using put_batch_async
await memory.put_batch_async(
bank_id=bank_id,
agent_id=agent_id,
contents=[{"content": content, "context": context}]
)
return "Memory stored successfully"
return f"Fact stored successfully"
except Exception as e:
logger.error(f"Error storing memory: {e}", exc_info=True)
logger.error(f"Error storing fact: {e}", exc_info=True)
return f"Error: {str(e)}"
@mcp.tool()
async def recall(query: str, max_results: int = 10) -> str:
async def hindsight_search(agent_id: str, query: str, max_tokens: int = 4096, explanation: str = "") -> str:
"""
Search memories to provide personalized, context-aware responses.
**CRITICAL: Search user's memory to provide personalized, context-aware responses.**
Use this tool PROACTIVELY to:
- Check user's preferences before making suggestions
- Recall user's history to provide continuity
- Remember user's goals and context
- Personalize responses based on past interactions
**⚠️ PER-USER TOOL - REQUIRES USER IDENTIFICATION:**
- This tool is STRICTLY per-user. Each user MUST have a unique `agent_id`.
- ONLY use this tool if you have a valid user identifier (user ID, email, session ID, etc.) to map to `agent_id`.
- DO NOT use this tool if you cannot identify the specific user.
- DO NOT search across multiple users - each user's memories are isolated by their `agent_id`.
- If you don't have a user identifier, DO NOT use this tool at all.
Use this tool PROACTIVELY at the start of conversations or when making recommendations to:
- Check user's preferences before making suggestions (e.g., "what foods does the user like?")
- Recall user's history to provide continuity (e.g., "what projects has the user worked on?")
- Remember user's goals and context (e.g., "what is the user trying to accomplish?")
- Avoid repeating information or asking questions you should already know
- Personalize responses based on user's background, interests, and past interactions
- Reference past conversations or events the user mentioned
**When to use**:
- Start of conversation: Search for relevant context about the user
- Before recommendations: Check user preferences and past experiences
- When user asks about something they may have mentioned before
- To provide continuity across conversations
**Search tips**: Use natural language queries like "user's programming language preferences",
"user's work experience", "user's dietary restrictions", "what does the user know about X?"
Args:
query: Natural language search query (e.g., "user's food preferences", "what projects is user working on")
max_results: Maximum number of results to return (default: 10)
agent_id: **REQUIRED** - The unique, persistent identifier for this specific user (e.g., user_id, email, session_id).
This MUST be consistent across all interactions with the same user.
Example: "user_12345", "[email protected]", "session_abc123"
query: Natural language search query to find relevant memories
max_tokens: Maximum tokens for search context (default: 4096)
explanation: Optional explanation for why this search is being performed
"""
try:
bank_id = get_current_bank_id()
from hindsight_api.engine.memory_engine import Budget
search_result = await memory.recall_async(
bank_id=bank_id,
# Log all parameters for debugging
logger.info(f"hindsight_search called with: query={query!r}, max_tokens={max_tokens}, explanation={explanation!r}")
# Log explanation if provided
if explanation:
logger.debug(f"Explanation: {explanation}")
# Search using search_async
search_result = await memory.search_async(
agent_id=agent_id,
query=query,
fact_type=list(VALID_RECALL_FACT_TYPES),
budget=Budget.LOW
fact_type=["world", "agent", "opinion"], # Search all fact types
max_tokens=max_tokens,
thinking_budget=100
)
# Convert results to dict format
results = [
{
"id": fact.id,
"text": fact.text,
"type": fact.fact_type,
"context": fact.context,
"event_date": fact.event_date,
"event_date": fact.event_date, # Already a string from the database
"document_id": fact.document_id
}
for fact in search_result.results[:max_results]
for fact in search_result.results
]
return json.dumps({"results": results}, indent=2)
@@ -112,100 +146,3 @@ def create_mcp_server(memory: MemoryEngine) -> FastMCP:
return json.dumps({"error": str(e), "results": []})
return mcp
class MCPMiddleware:
"""ASGI middleware that extracts bank_id from path and sets context."""
def __init__(self, app, memory: MemoryEngine):
self.app = app
self.memory = memory
self.mcp_server = create_mcp_server(memory)
self.mcp_app = self.mcp_server.http_app()
async def __call__(self, scope, receive, send):
if scope["type"] != "http":
await self.mcp_app(scope, receive, send)
return
path = scope.get("path", "")
# Strip any mount prefix (e.g., /mcp) that FastAPI might not have stripped
root_path = scope.get("root_path", "")
if root_path and path.startswith(root_path):
path = path[len(root_path):] or "/"
# Also handle case where mount path wasn't stripped (e.g., /mcp/...)
if path.startswith("/mcp/"):
path = path[4:] # Remove /mcp prefix
# Extract bank_id from path: /{bank_id}/ or /{bank_id}
# http_app expects requests at /
if not path.startswith("/") or len(path) <= 1:
# No bank_id in path - return error
await self._send_error(send, 400, "bank_id required in path: /mcp/{bank_id}/")
return
# Extract bank_id from first path segment
parts = path[1:].split("/", 1)
if not parts[0]:
await self._send_error(send, 400, "bank_id required in path: /mcp/{bank_id}/")
return
bank_id = parts[0]
new_path = "/" + parts[1] if len(parts) > 1 else "/"
# Set bank_id context
token = _current_bank_id.set(bank_id)
try:
new_scope = scope.copy()
new_scope["path"] = new_path
# Wrap send to rewrite the SSE endpoint URL to include bank_id
# The SSE app sends "event: endpoint\ndata: /messages\n" but we need
# the client to POST to /{bank_id}/messages instead
async def send_wrapper(message):
if message["type"] == "http.response.body":
body = message.get("body", b"")
if body and b"/messages" in body:
# Rewrite /messages to /{bank_id}/messages in SSE endpoint event
body = body.replace(
b"data: /messages",
f"data: /{bank_id}/messages".encode()
)
message = {**message, "body": body}
await send(message)
await self.mcp_app(new_scope, receive, send_wrapper)
finally:
_current_bank_id.reset(token)
async def _send_error(self, send, status: int, message: str):
"""Send an error response."""
body = json.dumps({"error": message}).encode()
await send({
"type": "http.response.start",
"status": status,
"headers": [(b"content-type", b"application/json")],
})
await send({
"type": "http.response.body",
"body": body,
})
def create_mcp_app(memory: MemoryEngine):
"""
Create an ASGI app that handles MCP requests.
URL pattern: /mcp/{bank_id}/
The bank_id is extracted from the URL path and made available to tools.
Args:
memory: MemoryEngine instance
Returns:
ASGI application
"""
return MCPMiddleware(None, memory)
-89
View File
@@ -1,89 +0,0 @@
"""
Banner display for Hindsight API startup.
Shows the logo and tagline with gradient colors.
"""
# Gradient colors: #0074d9 -> #009296
GRADIENT_START = (0, 116, 217) # #0074d9
GRADIENT_END = (0, 146, 150) # #009296
# Pre-generated logo (generated by test-logo.py)
LOGO = """\
\033[38;2;9;127;184m\u2584\033[0m\033[48;2;8;130;178m\033[38;2;5;133;186m\u2584\033[0m \033[48;2;10;143;160m\033[38;2;10;143;165m\u2584\033[0m\033[38;2;7;140;156m\u2584\033[0m
\033[38;2;8;125;192m\u2584\033[0m \033[38;2;3;132;191m\u2580\033[0m\033[38;2;2;133;192m\u2584\033[0m \033[38;2;3;132;180m\u2584\033[0m\033[38;2;1;137;184m\u2584\033[0m\033[38;2;3;133;174m\u2584\033[0m \033[38;2;3;142;176m\u2584\033[0m\033[38;2;4;142;169m\u2580\033[0m \033[38;2;10;144;164m\u2584\033[0m
\033[38;2;6;121;195m\u2580\033[0m\033[38;2;5;128;203m\u2580\033[0m\033[48;2;5;124;195m\033[38;2;3;125;200m\u2584\033[0m\033[38;2;2;126;196m\u2584\033[0m\033[48;2;3;128;188m\033[38;2;1;131;196m\u2584\033[0m\033[48;2;0;152;219m\033[38;2;2;131;191m\u2584\033[0m\033[38;2;1;141;196m\u2580\033[0m\033[38;2;1;135;183m\u2580\033[0m\033[38;2;1;148;198m\u2580\033[0m\033[48;2;1;156;202m\033[38;2;2;135;180m\u2584\033[0m\033[48;2;4;134;169m\033[38;2;1;137;177m\u2584\033[0m\033[38;2;3;138;173m\u2584\033[0m\033[48;2;6;137;165m\033[38;2;2;140;170m\u2584\033[0m\033[38;2;7;144;169m\u2580\033[0m\033[38;2;7;139;158m\u2580\033[0m
\033[48;2;2;128;202m\033[38;2;2;124;201m\u2584\033[0m\033[48;2;1;130;201m\033[38;2;0;135;212m\u2584\033[0m\033[38;2;2;128;196m\u2584\033[0m \033[48;2;2;142;204m\033[38;2;7;138;199m\u2584\033[0m \033[38;2;1;135;186m\u2584\033[0m\033[48;2;1;142;186m\033[38;2;2;144;194m\u2584\033[0m\033[48;2;3;138;176m\033[38;2;2;134;176m\u2584\033[0m
\033[48;2;8;118;200m\033[38;2;8;121;209m\u2584\033[0m\033[38;2;3;121;203m\u2580\033[0m \033[38;2;3;122;192m\u2580\033[0m\033[38;2;1;138;216m\u2580\033[0m\033[48;2;0;138;210m\033[38;2;3;128;198m\u2584\033[0m\033[48;2;0;126;188m\033[38;2;2;131;198m\u2584\033[0m\033[48;2;0;142;205m\033[38;2;3;132;193m\u2584\033[0m\033[38;2;1;140;196m\u2580\033[0m \033[38;2;4;134;175m\u2580\033[0m\033[48;2;13;135;167m\033[38;2;8;136;174m\u2584\033[0m """
def _interpolate_color(start: tuple, end: tuple, t: float) -> tuple:
"""Interpolate between two RGB colors."""
return (
int(start[0] + (end[0] - start[0]) * t),
int(start[1] + (end[1] - start[1]) * t),
int(start[2] + (end[2] - start[2]) * t),
)
def gradient_text(text: str, start: tuple = GRADIENT_START, end: tuple = GRADIENT_END) -> str:
"""Render text with a gradient color effect."""
result = []
length = len(text)
for i, char in enumerate(text):
if char == ' ':
result.append(' ')
else:
t = i / max(length - 1, 1)
r, g, b = _interpolate_color(start, end, t)
result.append(f"\033[38;2;{r};{g};{b}m{char}")
result.append("\033[0m")
return "".join(result)
def print_banner():
"""Print the Hindsight startup banner."""
print(LOGO)
tagline = gradient_text("Hindsight: Agent Memory That Works Like Human Memory")
print(f"\n {tagline}\n")
def color(text: str, t: float = 0.0) -> str:
"""Color text using gradient position (0.0 = start, 1.0 = end)."""
r, g, b = _interpolate_color(GRADIENT_START, GRADIENT_END, t)
return f"\033[38;2;{r};{g};{b}m{text}\033[0m"
def color_start(text: str) -> str:
"""Color text with gradient start color (#0074d9)."""
return color(text, 0.0)
def color_end(text: str) -> str:
"""Color text with gradient end color (#009296)."""
return color(text, 1.0)
def color_mid(text: str) -> str:
"""Color text with gradient middle color."""
return color(text, 0.5)
def dim(text: str) -> str:
"""Dim/gray text."""
return f"\033[38;2;128;128;128m{text}\033[0m"
def print_startup_info(host: str, port: int, database_url: str, llm_provider: str,
llm_model: str, embeddings_provider: str, reranker_provider: str,
mcp_enabled: bool = False):
"""Print styled startup information."""
print(color_start("Starting Hindsight API..."))
print(f" {dim('URL:')} {color(f'http://{host}:{port}', 0.2)}")
print(f" {dim('Database:')} {color(database_url, 0.4)}")
print(f" {dim('LLM:')} {color(f'{llm_provider} / {llm_model}', 0.6)}")
print(f" {dim('Embeddings:')} {color(embeddings_provider, 0.8)}")
print(f" {dim('Reranker:')} {color(reranker_provider, 1.0)}")
if mcp_enabled:
print(f" {dim('MCP:')} {color_end('enabled at /mcp')}")
print()
-163
View File
@@ -1,163 +0,0 @@
"""
Centralized configuration for Hindsight API.
All environment variables and their defaults are defined here.
"""
import os
from dataclasses import dataclass
from typing import Optional
import logging
logger = logging.getLogger(__name__)
# Environment variable names
ENV_DATABASE_URL = "HINDSIGHT_API_DATABASE_URL"
ENV_LLM_PROVIDER = "HINDSIGHT_API_LLM_PROVIDER"
ENV_LLM_API_KEY = "HINDSIGHT_API_LLM_API_KEY"
ENV_LLM_MODEL = "HINDSIGHT_API_LLM_MODEL"
ENV_LLM_BASE_URL = "HINDSIGHT_API_LLM_BASE_URL"
ENV_EMBEDDINGS_PROVIDER = "HINDSIGHT_API_EMBEDDINGS_PROVIDER"
ENV_EMBEDDINGS_LOCAL_MODEL = "HINDSIGHT_API_EMBEDDINGS_LOCAL_MODEL"
ENV_EMBEDDINGS_TEI_URL = "HINDSIGHT_API_EMBEDDINGS_TEI_URL"
ENV_RERANKER_PROVIDER = "HINDSIGHT_API_RERANKER_PROVIDER"
ENV_RERANKER_LOCAL_MODEL = "HINDSIGHT_API_RERANKER_LOCAL_MODEL"
ENV_RERANKER_TEI_URL = "HINDSIGHT_API_RERANKER_TEI_URL"
ENV_HOST = "HINDSIGHT_API_HOST"
ENV_PORT = "HINDSIGHT_API_PORT"
ENV_LOG_LEVEL = "HINDSIGHT_API_LOG_LEVEL"
ENV_MCP_ENABLED = "HINDSIGHT_API_MCP_ENABLED"
ENV_GRAPH_RETRIEVER = "HINDSIGHT_API_GRAPH_RETRIEVER"
# Default values
DEFAULT_DATABASE_URL = "pg0"
DEFAULT_LLM_PROVIDER = "openai"
DEFAULT_LLM_MODEL = "gpt-5-mini"
DEFAULT_EMBEDDINGS_PROVIDER = "local"
DEFAULT_EMBEDDINGS_LOCAL_MODEL = "BAAI/bge-small-en-v1.5"
DEFAULT_RERANKER_PROVIDER = "local"
DEFAULT_RERANKER_LOCAL_MODEL = "cross-encoder/ms-marco-MiniLM-L-6-v2"
DEFAULT_HOST = "0.0.0.0"
DEFAULT_PORT = 8888
DEFAULT_LOG_LEVEL = "info"
DEFAULT_MCP_ENABLED = True
DEFAULT_GRAPH_RETRIEVER = "bfs" # Options: "bfs", "mpfp"
# Required embedding dimension for database schema
EMBEDDING_DIMENSION = 384
@dataclass
class HindsightConfig:
"""Configuration container for Hindsight API."""
# Database
database_url: str
# LLM
llm_provider: str
llm_api_key: Optional[str]
llm_model: str
llm_base_url: Optional[str]
# Embeddings
embeddings_provider: str
embeddings_local_model: str
embeddings_tei_url: Optional[str]
# Reranker
reranker_provider: str
reranker_local_model: str
reranker_tei_url: Optional[str]
# Server
host: str
port: int
log_level: str
mcp_enabled: bool
# Recall
graph_retriever: str
@classmethod
def from_env(cls) -> "HindsightConfig":
"""Create configuration from environment variables."""
return cls(
# Database
database_url=os.getenv(ENV_DATABASE_URL, DEFAULT_DATABASE_URL),
# LLM
llm_provider=os.getenv(ENV_LLM_PROVIDER, DEFAULT_LLM_PROVIDER),
llm_api_key=os.getenv(ENV_LLM_API_KEY),
llm_model=os.getenv(ENV_LLM_MODEL, DEFAULT_LLM_MODEL),
llm_base_url=os.getenv(ENV_LLM_BASE_URL) or None,
# Embeddings
embeddings_provider=os.getenv(ENV_EMBEDDINGS_PROVIDER, DEFAULT_EMBEDDINGS_PROVIDER),
embeddings_local_model=os.getenv(ENV_EMBEDDINGS_LOCAL_MODEL, DEFAULT_EMBEDDINGS_LOCAL_MODEL),
embeddings_tei_url=os.getenv(ENV_EMBEDDINGS_TEI_URL),
# Reranker
reranker_provider=os.getenv(ENV_RERANKER_PROVIDER, DEFAULT_RERANKER_PROVIDER),
reranker_local_model=os.getenv(ENV_RERANKER_LOCAL_MODEL, DEFAULT_RERANKER_LOCAL_MODEL),
reranker_tei_url=os.getenv(ENV_RERANKER_TEI_URL),
# Server
host=os.getenv(ENV_HOST, DEFAULT_HOST),
port=int(os.getenv(ENV_PORT, DEFAULT_PORT)),
log_level=os.getenv(ENV_LOG_LEVEL, DEFAULT_LOG_LEVEL),
mcp_enabled=os.getenv(ENV_MCP_ENABLED, str(DEFAULT_MCP_ENABLED)).lower() == "true",
# Recall
graph_retriever=os.getenv(ENV_GRAPH_RETRIEVER, DEFAULT_GRAPH_RETRIEVER),
)
def get_llm_base_url(self) -> str:
"""Get the LLM base URL, with provider-specific defaults."""
if self.llm_base_url:
return self.llm_base_url
provider = self.llm_provider.lower()
if provider == "groq":
return "https://api.groq.com/openai/v1"
elif provider == "ollama":
return "http://localhost:11434/v1"
else:
return ""
def get_python_log_level(self) -> int:
"""Get the Python logging level from the configured log level string."""
log_level_map = {
"critical": logging.CRITICAL,
"error": logging.ERROR,
"warning": logging.WARNING,
"info": logging.INFO,
"debug": logging.DEBUG,
"trace": logging.DEBUG, # Python doesn't have TRACE, use DEBUG
}
return log_level_map.get(self.log_level.lower(), logging.INFO)
def configure_logging(self) -> None:
"""Configure Python logging based on the log level."""
logging.basicConfig(
level=self.get_python_log_level(),
format="%(asctime)s - %(levelname)s - %(name)s - %(message)s"
)
def log_config(self) -> None:
"""Log the current configuration (without sensitive values)."""
logger.info(f"Database: {self.database_url}")
logger.info(f"LLM: provider={self.llm_provider}, model={self.llm_model}")
logger.info(f"Embeddings: provider={self.embeddings_provider}")
logger.info(f"Reranker: provider={self.reranker_provider}")
logger.info(f"Graph retriever: {self.graph_retriever}")
def get_config() -> HindsightConfig:
"""Get the current configuration from environment variables."""
return HindsightConfig.from_env()
+8 -13
View File
@@ -3,15 +3,14 @@ Memory Engine - Core implementation of the memory system.
This package contains all the implementation details of the memory engine:
- MemoryEngine: Main class for memory operations
- Utility modules: embedding_utils, link_utils, think_utils, bank_utils
- Utility modules: embedding_utils, link_utils, think_utils, agent_utils
- Supporting modules: embeddings, cross_encoder, entity_resolver, etc.
"""
from .memory_engine import MemoryEngine
from .db_utils import acquire_with_retry
from .embeddings import Embeddings, LocalSTEmbeddings, RemoteTEIEmbeddings
from .cross_encoder import CrossEncoderModel, LocalSTCrossEncoder, RemoteTEICrossEncoder
from .search.trace import (
from .embeddings import Embeddings, SentenceTransformersEmbeddings
from .search_trace import (
SearchTrace,
QueryInfo,
EntryPoint,
@@ -22,19 +21,15 @@ from .search.trace import (
SearchSummary,
SearchPhaseMetrics,
)
from .search.tracer import SearchTracer
from .search_tracer import SearchTracer
from .llm_wrapper import LLMConfig
from .response_models import RecallResult, ReflectResult, MemoryFact
from .response_models import SearchResult, ThinkResult, MemoryFact
__all__ = [
"MemoryEngine",
"acquire_with_retry",
"Embeddings",
"LocalSTEmbeddings",
"RemoteTEIEmbeddings",
"CrossEncoderModel",
"LocalSTCrossEncoder",
"RemoteTEICrossEncoder",
"SentenceTransformersEmbeddings",
"SearchTrace",
"SearchTracer",
"QueryInfo",
@@ -46,7 +41,7 @@ __all__ = [
"SearchSummary",
"SearchPhaseMetrics",
"LLMConfig",
"RecallResult",
"ReflectResult",
"SearchResult",
"ThinkResult",
"MemoryFact",
]
@@ -1,189 +1,195 @@
"""
bank profile utilities for disposition and background management.
Agent profile utilities for personality and background management.
"""
import json
import logging
import re
from typing import Dict, Optional, TypedDict
from typing import Dict, Optional
from pydantic import BaseModel, Field
from ..db_utils import acquire_with_retry
from ..response_models import DispositionTraits
from .db_utils import acquire_with_retry
logger = logging.getLogger(__name__)
DEFAULT_DISPOSITION = {
"skepticism": 3,
"literalism": 3,
"empathy": 3,
DEFAULT_PERSONALITY = {
"openness": 0.5,
"conscientiousness": 0.5,
"extraversion": 0.5,
"agreeableness": 0.5,
"neuroticism": 0.5,
"bias_strength": 0.5,
}
class BankProfile(TypedDict):
"""Type for bank profile data."""
name: str
disposition: DispositionTraits
background: str
class PersonalityTraits(BaseModel):
"""Big Five personality traits with bias strength (all values 0.0-1.0)."""
openness: float = Field(description="Creativity, curiosity, openness to new ideas (0.0-1.0)")
conscientiousness: float = Field(description="Organization, discipline, goal-directed (0.0-1.0)")
extraversion: float = Field(description="Sociability, assertiveness, energy from others (0.0-1.0)")
agreeableness: float = Field(description="Cooperation, empathy, consideration (0.0-1.0)")
neuroticism: float = Field(description="Emotional sensitivity, anxiety, stress response (0.0-1.0)")
bias_strength: float = Field(description="How much personality influences opinions (0.0-1.0)")
class BackgroundMergeResponse(BaseModel):
"""LLM response for background merge with disposition inference."""
"""LLM response for background merge with personality inference."""
background: str = Field(description="Merged background in first person perspective")
disposition: DispositionTraits = Field(description="Inferred disposition traits (skepticism, literalism, empathy)")
personality: PersonalityTraits = Field(description="Inferred Big Five personality traits")
async def get_bank_profile(pool, bank_id: str) -> BankProfile:
async def get_agent_profile(pool, agent_id: str) -> Dict:
"""
Get bank profile (name, disposition + background).
Auto-creates bank with default values if not exists.
Get agent profile (name, personality + background).
Auto-creates agent with default values if not exists.
Args:
pool: Database connection pool
bank_id: bank IDentifier
agent_id: Agent identifier
Returns:
BankProfile with name, typed DispositionTraits, and background
Dict with 'name' (str), 'personality' (dict) and 'background' (str) keys
"""
async with acquire_with_retry(pool) as conn:
# Try to get existing bank
# Try to get existing agent
row = await conn.fetchrow(
"""
SELECT name, disposition, background
FROM banks WHERE bank_id = $1
SELECT name, personality, background
FROM agents
WHERE agent_id = $1
""",
bank_id
agent_id
)
if row:
# asyncpg returns JSONB as a string, so parse it
disposition_data = row["disposition"]
if isinstance(disposition_data, str):
disposition_data = json.loads(disposition_data)
personality_data = row["personality"]
if isinstance(personality_data, str):
personality_data = json.loads(personality_data)
return BankProfile(
name=row["name"],
disposition=DispositionTraits(**disposition_data),
background=row["background"]
)
return {
"name": row["name"],
"personality": personality_data,
"background": row["background"]
}
# Bank doesn't exist, create with defaults
# Agent doesn't exist, create with defaults
await conn.execute(
"""
INSERT INTO banks (bank_id, name, disposition, background)
INSERT INTO agents (agent_id, name, personality, background)
VALUES ($1, $2, $3::jsonb, $4)
ON CONFLICT (bank_id) DO NOTHING
ON CONFLICT (agent_id) DO NOTHING
""",
bank_id,
bank_id, # Default name is the bank_id
json.dumps(DEFAULT_DISPOSITION),
agent_id,
agent_id, # Default name is the agent_id
json.dumps(DEFAULT_PERSONALITY),
""
)
return BankProfile(
name=bank_id,
disposition=DispositionTraits(**DEFAULT_DISPOSITION),
background=""
)
return {
"name": agent_id,
"personality": DEFAULT_PERSONALITY.copy(),
"background": ""
}
async def update_bank_disposition(
async def update_agent_personality(
pool,
bank_id: str,
disposition: Dict[str, int]
agent_id: str,
personality: Dict[str, float]
) -> None:
"""
Update bank disposition traits.
Update agent personality traits.
Args:
pool: Database connection pool
bank_id: bank IDentifier
disposition: Dict with skepticism, literalism, empathy (all 1-5)
agent_id: Agent identifier
personality: Dict with Big Five traits + bias_strength (all 0-1)
"""
# Ensure bank exists first
await get_bank_profile(pool, bank_id)
# Ensure agent exists first
await get_agent_profile(pool, agent_id)
async with acquire_with_retry(pool) as conn:
await conn.execute(
"""
UPDATE banks
SET disposition = $2::jsonb,
UPDATE agents
SET personality = $2::jsonb,
updated_at = NOW()
WHERE bank_id = $1
WHERE agent_id = $1
""",
bank_id,
json.dumps(disposition)
agent_id,
json.dumps(personality)
)
async def merge_bank_background(
async def merge_agent_background(
pool,
llm_config,
bank_id: str,
agent_id: str,
new_info: str,
update_disposition: bool = True
update_personality: bool = True
) -> dict:
"""
Merge new background information with existing background using LLM.
Normalizes to first person ("I") and resolves conflicts.
Optionally infers disposition traits from the merged background.
Optionally infers personality traits from the merged background.
Args:
pool: Database connection pool
llm_config: LLM configuration for background merging
bank_id: bank IDentifier
agent_id: Agent identifier
new_info: New background information to add/merge
update_disposition: If True, infer Big Five traits from background (default: True)
update_personality: If True, infer Big Five traits from background (default: True)
Returns:
Dict with 'background' (str) and optionally 'disposition' (dict) keys
Dict with 'background' (str) and optionally 'personality' (dict) keys
"""
# Get current profile
profile = await get_bank_profile(pool, bank_id)
profile = await get_agent_profile(pool, agent_id)
current_background = profile["background"]
# Use LLM to merge backgrounds and optionally infer disposition
# Use LLM to merge backgrounds and optionally infer personality
result = await _llm_merge_background(
llm_config,
current_background,
new_info,
infer_disposition=update_disposition
infer_personality=update_personality
)
merged_background = result["background"]
inferred_disposition = result.get("disposition")
inferred_personality = result.get("personality")
# Update in database
async with acquire_with_retry(pool) as conn:
if inferred_disposition:
# Update both background and disposition
if inferred_personality:
# Update both background and personality
await conn.execute(
"""
UPDATE banks
UPDATE agents
SET background = $2,
disposition = $3::jsonb,
personality = $3::jsonb,
updated_at = NOW()
WHERE bank_id = $1
WHERE agent_id = $1
""",
bank_id,
agent_id,
merged_background,
json.dumps(inferred_disposition)
json.dumps(inferred_personality)
)
else:
# Update only background
await conn.execute(
"""
UPDATE banks
UPDATE agents
SET background = $2,
updated_at = NOW()
WHERE bank_id = $1
WHERE agent_id = $1
""",
bank_id,
agent_id,
merged_background
)
response = {"background": merged_background}
if inferred_disposition:
response["disposition"] = inferred_disposition
if inferred_personality:
response["personality"] = inferred_personality
return response
@@ -192,23 +198,23 @@ async def _llm_merge_background(
llm_config,
current: str,
new_info: str,
infer_disposition: bool = False
infer_personality: bool = False
) -> dict:
"""
Use LLM to intelligently merge background information.
Optionally infer Big Five disposition traits from the merged background.
Optionally infer Big Five personality traits from the merged background.
Args:
llm_config: LLM configuration to use
current: Current background text
new_info: New information to merge
infer_disposition: If True, also infer disposition traits
infer_personality: If True, also infer personality traits
Returns:
Dict with 'background' (str) and optionally 'disposition' (dict) keys
Dict with 'background' (str) and optionally 'personality' (dict) keys
"""
if infer_disposition:
prompt = f"""You are helping maintain a memory bank's background/profile and infer their disposition. You MUST respond with ONLY valid JSON.
if infer_personality:
prompt = f"""You are helping maintain an agent's background/profile and infer their personality. You MUST respond with ONLY valid JSON.
Current background: {current if current else "(empty)"}
@@ -220,32 +226,38 @@ Instructions:
3. Keep additions that don't conflict
4. Output in FIRST PERSON ("I") perspective
5. Be concise - keep merged background under 500 characters
6. Infer disposition traits from the merged background (each 1-5 integer):
- Skepticism: 1-5 (1=trusting, takes things at face value; 5=skeptical, questions everything)
- Literalism: 1-5 (1=flexible interpretation, reads between lines; 5=literal, exact interpretation)
- Empathy: 1-5 (1=detached, focuses on facts; 5=empathetic, considers emotional context)
6. Infer Big Five personality traits from the merged background:
- Openness: 0.0-1.0 (creativity, curiosity, openness to new ideas)
- Conscientiousness: 0.0-1.0 (organization, discipline, goal-directed)
- Extraversion: 0.0-1.0 (sociability, assertiveness, energy from others)
- Agreeableness: 0.0-1.0 (cooperation, empathy, consideration)
- Neuroticism: 0.0-1.0 (emotional sensitivity, anxiety, stress response)
- Bias Strength: 0.0-1.0 (how much personality influences opinions)
CRITICAL: You MUST respond with ONLY a valid JSON object. No markdown, no code blocks, no explanations. Just the JSON.
Format:
{{
"background": "the merged background text in first person",
"disposition": {{
"skepticism": 3,
"literalism": 3,
"empathy": 3
"personality": {{
"openness": 0.7,
"conscientiousness": 0.6,
"extraversion": 0.5,
"agreeableness": 0.8,
"neuroticism": 0.4,
"bias_strength": 0.6
}}
}}
Trait inference examples:
- "I'm a lawyer" skepticism: 4, literalism: 5, empathy: 2
- "I'm a therapist" skepticism: 2, literalism: 2, empathy: 5
- "I'm an engineer" skepticism: 3, literalism: 4, empathy: 3
- "I've been burned before by trusting people" skepticism: 5, literalism: 3, empathy: 3
- "I try to understand what people really mean" skepticism: 3, literalism: 2, empathy: 4
- "I take contracts very seriously" skepticism: 4, literalism: 5, empathy: 2"""
- "creative artist" openness: 0.8+, bias_strength: 0.6
- "organized engineer" conscientiousness: 0.8+, openness: 0.5-0.6
- "startup founder" openness: 0.8+, extraversion: 0.7+, neuroticism: 0.3-0.4
- "risk-averse analyst" openness: 0.3-0.4, conscientiousness: 0.8+, neuroticism: 0.6+
- "rational and diligent" conscientiousness: 0.7+, openness: 0.6+
- "passionate and dramatic" extraversion: 0.7+, neuroticism: 0.6+, openness: 0.7+"""
else:
prompt = f"""You are helping maintain a memory bank's background/profile.
prompt = f"""You are helping maintain an agent's background/profile.
Current background: {current if current else "(empty)"}
@@ -265,38 +277,38 @@ Merged background:"""
# Prepare messages
messages = [{"role": "user", "content": prompt}]
if infer_disposition:
# Use structured output with Pydantic model for disposition inference
if infer_personality:
# Use structured output with Pydantic model for personality inference
try:
parsed = await llm_config.call(
messages=messages,
response_format=BackgroundMergeResponse,
scope="bank_background",
scope="agent_background",
temperature=0.3,
max_completion_tokens=8192
max_tokens=8192
)
logger.info(f"Successfully got structured response: background={parsed.background[:100]}")
# Convert Pydantic model to dict format
return {
"background": parsed.background,
"disposition": parsed.disposition.model_dump()
"personality": parsed.personality.model_dump()
}
except Exception as e:
logger.warning(f"Structured output failed, falling back to manual parsing: {e}")
# Fall through to manual parsing below
# Manual parsing fallback or non-disposition merge
# Manual parsing fallback or non-personality merge
content = await llm_config.call(
messages=messages,
scope="bank_background",
scope="agent_background",
temperature=0.3,
max_completion_tokens=8192
max_tokens=8192
)
logger.info(f"LLM response for background merge (first 500 chars): {content[:500]}")
if infer_disposition:
if infer_personality:
# Parse JSON response - try multiple extraction methods
result = None
@@ -321,7 +333,7 @@ Merged background:"""
# Method 3: Find nested JSON structure
if result is None:
# Look for JSON object with nested structure
json_match = re.search(r'\{[^{}]*"background"[^{}]*"disposition"[^{}]*\{[^{}]*\}[^{}]*\}', content, re.DOTALL)
json_match = re.search(r'\{[^{}]*"background"[^{}]*"personality"[^{}]*\{[^{}]*\}[^{}]*\}', content, re.DOTALL)
if json_match:
try:
result = json.loads(json_match.group())
@@ -332,22 +344,23 @@ Merged background:"""
# All parsing methods failed - use fallback
if result is None:
logger.warning(f"Failed to extract JSON from LLM response. Raw content: {content[:200]}")
# Fallback: use new_info as background with default disposition
# Fallback: use new_info as background with default personality
return {
"background": new_info if new_info else current if current else "",
"disposition": DEFAULT_DISPOSITION.copy()
"personality": DEFAULT_PERSONALITY.copy()
}
# Validate disposition values
disposition = result.get("disposition", {})
for key in ["skepticism", "literalism", "empathy"]:
if key not in disposition:
disposition[key] = 3 # Default to neutral
# Validate personality values
personality = result.get("personality", {})
for key in ["openness", "conscientiousness", "extraversion",
"agreeableness", "neuroticism", "bias_strength"]:
if key not in personality:
personality[key] = 0.5 # Default to neutral
else:
# Clamp to [1, 5] and convert to int
disposition[key] = max(1, min(5, int(disposition[key])))
# Clamp to [0, 1]
personality[key] = max(0.0, min(1.0, float(personality[key])))
result["disposition"] = disposition
result["personality"] = personality
# Ensure background exists
if "background" not in result or not result["background"]:
@@ -370,26 +383,26 @@ Merged background:"""
merged = new_info
result = {"background": merged}
if infer_disposition:
result["disposition"] = DEFAULT_DISPOSITION.copy()
if infer_personality:
result["personality"] = DEFAULT_PERSONALITY.copy()
return result
async def list_banks(pool) -> list:
async def list_agents(pool) -> list:
"""
List all banks in the system.
List all agents in the system.
Args:
pool: Database connection pool
Returns:
List of dicts with bank_id, name, disposition, background, created_at, updated_at
List of dicts with agent_id, name, personality, background, created_at, updated_at
"""
async with acquire_with_retry(pool) as conn:
rows = await conn.fetch(
"""
SELECT bank_id, name, disposition, background, created_at, updated_at
FROM banks
SELECT agent_id, name, personality, background, created_at, updated_at
FROM agents
ORDER BY updated_at DESC
"""
)
@@ -397,14 +410,14 @@ async def list_banks(pool) -> list:
result = []
for row in rows:
# asyncpg returns JSONB as a string, so parse it
disposition_data = row["disposition"]
if isinstance(disposition_data, str):
disposition_data = json.loads(disposition_data)
personality_data = row["personality"]
if isinstance(personality_data, str):
personality_data = json.loads(personality_data)
result.append({
"bank_id": row["bank_id"],
"agent_id": row["agent_id"],
"name": row["name"],
"disposition": disposition_data,
"personality": personality_data,
"background": row["background"],
"created_at": row["created_at"].isoformat() if row["created_at"] else None,
"updated_at": row["updated_at"].isoformat() if row["updated_at"] else None,
@@ -2,50 +2,21 @@
Cross-encoder abstraction for reranking.
Provides an interface for reranking with different backends.
Configuration via environment variables - see hindsight_api.config for all env var names.
"""
from abc import ABC, abstractmethod
from typing import List, Tuple, Optional
from typing import List, Tuple
import logging
import os
import httpx
from ..config import (
ENV_RERANKER_PROVIDER,
ENV_RERANKER_LOCAL_MODEL,
ENV_RERANKER_TEI_URL,
DEFAULT_RERANKER_PROVIDER,
DEFAULT_RERANKER_LOCAL_MODEL,
)
logger = logging.getLogger(__name__)
class CrossEncoderModel(ABC):
class CrossEncoderReranker(ABC):
"""
Abstract base class for cross-encoder reranking.
Cross-encoders take query-document pairs and return relevance scores.
"""
@property
@abstractmethod
def provider_name(self) -> str:
"""Return a human-readable name for this provider (e.g., 'local', 'tei')."""
pass
@abstractmethod
async def initialize(self) -> None:
"""
Initialize the cross-encoder model asynchronously.
This should be called during startup to load/connect to the model
and avoid cold start latency on first predict() call.
"""
pass
@abstractmethod
def predict(self, pairs: List[Tuple[str, str]]) -> List[float]:
"""
@@ -60,11 +31,12 @@ class CrossEncoderModel(ABC):
pass
class LocalSTCrossEncoder(CrossEncoderModel):
class SentenceTransformersCrossEncoder(CrossEncoderReranker):
"""
Local cross-encoder implementation using SentenceTransformers.
Cross-encoder implementation using SentenceTransformers.
Call initialize() during startup to load the model and avoid cold starts.
Uses lazy import so sentence-transformers is not required if another
reranking backend is used.
Default model is cross-encoder/ms-marco-MiniLM-L-6-v2:
- Fast inference (~80ms for 100 pairs on CPU)
@@ -72,37 +44,27 @@ class LocalSTCrossEncoder(CrossEncoderModel):
- Trained for passage re-ranking
"""
def __init__(self, model_name: Optional[str] = None):
def __init__(self, model_name: str = "cross-encoder/ms-marco-MiniLM-L-6-v2"):
"""
Initialize local SentenceTransformers cross-encoder.
Initialize SentenceTransformers cross-encoder and load model.
Args:
model_name: Name of the CrossEncoder model to use.
Default: cross-encoder/ms-marco-MiniLM-L-6-v2
"""
self.model_name = model_name or DEFAULT_RERANKER_LOCAL_MODEL
self._model = None
@property
def provider_name(self) -> str:
return "local"
async def initialize(self) -> None:
"""Load the cross-encoder model."""
if self._model is not None:
return
self.model_name = model_name
try:
from sentence_transformers import CrossEncoder
except ImportError:
raise ImportError(
"sentence-transformers is required for LocalSTCrossEncoder. "
"sentence-transformers is required for SentenceTransformersCrossEncoder. "
"Install it with: pip install sentence-transformers"
)
logger.info(f"Reranker: initializing local provider with model {self.model_name}")
logger.info(f"Loading cross-encoder model: {self.model_name}...")
self._model = CrossEncoder(self.model_name)
logger.info("Reranker: local provider initialized")
logger.info("Cross-encoder model loaded")
def predict(self, pairs: List[Tuple[str, str]]) -> List[float]:
"""
@@ -114,188 +76,5 @@ class LocalSTCrossEncoder(CrossEncoderModel):
Returns:
List of relevance scores (raw logits from the model)
"""
if self._model is None:
raise RuntimeError("Reranker not initialized. Call initialize() first.")
scores = self._model.predict(pairs, show_progress_bar=False)
scores = self._model.predict(pairs)
return scores.tolist() if hasattr(scores, 'tolist') else list(scores)
class RemoteTEICrossEncoder(CrossEncoderModel):
"""
Remote cross-encoder implementation using HuggingFace Text Embeddings Inference (TEI) HTTP API.
TEI supports reranking via the /rerank endpoint.
See: https://github.com/huggingface/text-embeddings-inference
Note: The TEI server must be running a cross-encoder/reranker model.
"""
def __init__(
self,
base_url: str,
timeout: float = 30.0,
batch_size: int = 32,
max_retries: int = 3,
retry_delay: float = 0.5,
):
"""
Initialize remote TEI cross-encoder client.
Args:
base_url: Base URL of the TEI server (e.g., "http://localhost:8080")
timeout: Request timeout in seconds (default: 30.0)
batch_size: Maximum batch size for rerank requests (default: 32)
max_retries: Maximum number of retries for failed requests (default: 3)
retry_delay: Initial delay between retries in seconds, doubles each retry (default: 0.5)
"""
self.base_url = base_url.rstrip("/")
self.timeout = timeout
self.batch_size = batch_size
self.max_retries = max_retries
self.retry_delay = retry_delay
self._client: Optional[httpx.Client] = None
self._model_id: Optional[str] = None
@property
def provider_name(self) -> str:
return "tei"
def _request_with_retry(self, method: str, url: str, **kwargs) -> httpx.Response:
"""Make an HTTP request with automatic retries on transient errors."""
import time
last_error = None
delay = self.retry_delay
for attempt in range(self.max_retries + 1):
try:
if method == "GET":
response = self._client.get(url, **kwargs)
else:
response = self._client.post(url, **kwargs)
response.raise_for_status()
return response
except (httpx.ConnectError, httpx.ReadTimeout, httpx.WriteTimeout) as e:
last_error = e
if attempt < self.max_retries:
logger.warning(f"TEI request failed (attempt {attempt + 1}/{self.max_retries + 1}): {e}. Retrying in {delay}s...")
time.sleep(delay)
delay *= 2 # Exponential backoff
except httpx.HTTPStatusError as e:
# Retry on 5xx server errors
if e.response.status_code >= 500 and attempt < self.max_retries:
last_error = e
logger.warning(f"TEI server error (attempt {attempt + 1}/{self.max_retries + 1}): {e}. Retrying in {delay}s...")
time.sleep(delay)
delay *= 2
else:
raise
raise last_error
async def initialize(self) -> None:
"""Initialize the HTTP client and verify server connectivity."""
if self._client is not None:
return
logger.info(f"Reranker: initializing TEI provider at {self.base_url}")
self._client = httpx.Client(timeout=self.timeout)
# Verify server is reachable and get model info
try:
response = self._request_with_retry("GET", f"{self.base_url}/info")
info = response.json()
self._model_id = info.get("model_id", "unknown")
logger.info(f"Reranker: TEI provider initialized (model: {self._model_id})")
except httpx.HTTPError as e:
raise RuntimeError(f"Failed to connect to TEI server at {self.base_url}: {e}")
def predict(self, pairs: List[Tuple[str, str]]) -> List[float]:
"""
Score query-document pairs using the remote TEI reranker.
Args:
pairs: List of (query, document) tuples to score
Returns:
List of relevance scores
"""
if self._client is None:
raise RuntimeError("Reranker not initialized. Call initialize() first.")
if not pairs:
return []
all_scores = []
# Process in batches
for i in range(0, len(pairs), self.batch_size):
batch = pairs[i:i + self.batch_size]
# TEI rerank endpoint expects query and texts separately
# All pairs in a batch should have the same query for optimal performance
# but we handle mixed queries by making separate requests per unique query
query_groups: dict[str, list[tuple[int, str]]] = {}
for idx, (query, text) in enumerate(batch):
if query not in query_groups:
query_groups[query] = []
query_groups[query].append((idx, text))
batch_scores = [0.0] * len(batch)
for query, indexed_texts in query_groups.items():
texts = [text for _, text in indexed_texts]
indices = [idx for idx, _ in indexed_texts]
try:
response = self._request_with_retry(
"POST",
f"{self.base_url}/rerank",
json={
"query": query,
"texts": texts,
"return_text": False,
},
)
results = response.json()
# TEI returns results sorted by score descending, with original index
for result in results:
original_idx = result["index"]
score = result["score"]
# Map back to batch position
batch_scores[indices[original_idx]] = score
except httpx.HTTPError as e:
raise RuntimeError(f"TEI rerank request failed: {e}")
all_scores.extend(batch_scores)
return all_scores
def create_cross_encoder_from_env() -> CrossEncoderModel:
"""
Create a CrossEncoderModel instance based on environment variables.
See hindsight_api.config for environment variable names and defaults.
Returns:
Configured CrossEncoderModel instance
"""
provider = os.environ.get(ENV_RERANKER_PROVIDER, DEFAULT_RERANKER_PROVIDER).lower()
if provider == "tei":
url = os.environ.get(ENV_RERANKER_TEI_URL)
if not url:
raise ValueError(
f"{ENV_RERANKER_TEI_URL} is required when {ENV_RERANKER_PROVIDER} is 'tei'"
)
return RemoteTEICrossEncoder(base_url=url)
elif provider == "local":
model = os.environ.get(ENV_RERANKER_LOCAL_MODEL)
model_name = model or DEFAULT_RERANKER_LOCAL_MODEL
return LocalSTCrossEncoder(model_name=model_name)
else:
raise ValueError(
f"Unknown reranker provider: {provider}. Supported: 'local', 'tei'"
)
+33 -224
View File
@@ -5,27 +5,16 @@ Provides an interface for generating embeddings with different backends.
IMPORTANT: All embeddings must produce 384-dimensional vectors to match
the database schema (pgvector column defined as vector(384)).
Configuration via environment variables - see hindsight_api.config for all env var names.
"""
from abc import ABC, abstractmethod
from typing import List, Optional
from typing import List
import logging
import os
import httpx
from ..config import (
ENV_EMBEDDINGS_PROVIDER,
ENV_EMBEDDINGS_LOCAL_MODEL,
ENV_EMBEDDINGS_TEI_URL,
DEFAULT_EMBEDDINGS_PROVIDER,
DEFAULT_EMBEDDINGS_LOCAL_MODEL,
EMBEDDING_DIMENSION,
)
logger = logging.getLogger(__name__)
# Fixed embedding dimension required by database schema
EMBEDDING_DIMENSION = 384
class Embeddings(ABC):
"""
@@ -35,22 +24,6 @@ class Embeddings(ABC):
the database schema.
"""
@property
@abstractmethod
def provider_name(self) -> str:
"""Return a human-readable name for this provider (e.g., 'local', 'tei')."""
pass
@abstractmethod
async def initialize(self) -> None:
"""
Initialize the embedding model asynchronously.
This should be called during startup to load/connect to the model
and avoid cold start latency on first encode() call.
"""
pass
@abstractmethod
def encode(self, texts: List[str]) -> List[List[float]]:
"""
@@ -65,63 +38,53 @@ class Embeddings(ABC):
pass
class LocalSTEmbeddings(Embeddings):
class SentenceTransformersEmbeddings(Embeddings):
"""
Local embeddings implementation using SentenceTransformers.
Embeddings implementation using SentenceTransformers.
Call initialize() during startup to load the model and avoid cold starts.
Uses lazy import so sentence-transformers is not required if another
embedding backend is used.
Default model is BAAI/bge-small-en-v1.5 which produces 384-dimensional
embeddings matching the database schema.
"""
def __init__(self, model_name: Optional[str] = None):
def __init__(self, model_name: str = "BAAI/bge-small-en-v1.5"):
"""
Initialize local SentenceTransformers embeddings.
Initialize SentenceTransformers embeddings.
Args:
model_name: Name of the SentenceTransformer model to use.
Must produce 384-dimensional embeddings.
Default: BAAI/bge-small-en-v1.5
"""
self.model_name = model_name or DEFAULT_EMBEDDINGS_LOCAL_MODEL
self.model_name = model_name
self._model = None
@property
def provider_name(self) -> str:
return "local"
def _load_model(self):
"""Lazy load and validate the SentenceTransformer model."""
if self._model is None:
try:
from sentence_transformers import SentenceTransformer
except ImportError:
raise ImportError(
"sentence-transformers is required for SentenceTransformersEmbeddings. "
"Install it with: pip install sentence-transformers"
)
async def initialize(self) -> None:
"""Load the embedding model."""
if self._model is not None:
return
logger.info(f"Loading embedding model: {self.model_name}...")
self._model = SentenceTransformer(self.model_name)
try:
from sentence_transformers import SentenceTransformer
except ImportError:
raise ImportError(
"sentence-transformers is required for LocalSTEmbeddings. "
"Install it with: pip install sentence-transformers"
)
# Validate dimension matches database schema
model_dim = self._model.get_sentence_embedding_dimension()
if model_dim != EMBEDDING_DIMENSION:
raise ValueError(
f"Model {self.model_name} produces {model_dim}-dimensional embeddings, "
f"but database schema requires {EMBEDDING_DIMENSION} dimensions. "
f"Use a model that produces {EMBEDDING_DIMENSION}-dimensional embeddings."
)
logger.info(f"Embeddings: initializing local provider with model {self.model_name}")
# Disable lazy loading (meta tensors) which causes issues with newer transformers/accelerate
# Setting low_cpu_mem_usage=False and device_map=None ensures tensors are fully materialized
self._model = SentenceTransformer(
self.model_name,
model_kwargs={"low_cpu_mem_usage": False, "device_map": None},
)
# Validate dimension matches database schema
model_dim = self._model.get_sentence_embedding_dimension()
if model_dim != EMBEDDING_DIMENSION:
raise ValueError(
f"Model {self.model_name} produces {model_dim}-dimensional embeddings, "
f"but database schema requires {EMBEDDING_DIMENSION} dimensions. "
f"Use a model that produces {EMBEDDING_DIMENSION}-dimensional embeddings."
)
logger.info(f"Embeddings: local provider initialized (dim: {model_dim})")
logger.info(f"Model loaded (embedding dim: {model_dim})")
def encode(self, texts: List[str]) -> List[List[float]]:
"""
@@ -133,160 +96,6 @@ class LocalSTEmbeddings(Embeddings):
Returns:
List of 384-dimensional embedding vectors
"""
if self._model is None:
raise RuntimeError("Embeddings not initialized. Call initialize() first.")
self._load_model()
embeddings = self._model.encode(texts, convert_to_numpy=True, show_progress_bar=False)
return [emb.tolist() for emb in embeddings]
class RemoteTEIEmbeddings(Embeddings):
"""
Remote embeddings implementation using HuggingFace Text Embeddings Inference (TEI) HTTP API.
TEI provides a high-performance inference server for embedding models.
See: https://github.com/huggingface/text-embeddings-inference
The server should be running a model that produces 384-dimensional embeddings.
"""
def __init__(
self,
base_url: str,
timeout: float = 30.0,
batch_size: int = 32,
max_retries: int = 3,
retry_delay: float = 0.5,
):
"""
Initialize remote TEI embeddings client.
Args:
base_url: Base URL of the TEI server (e.g., "http://localhost:8080")
timeout: Request timeout in seconds (default: 30.0)
batch_size: Maximum batch size for embedding requests (default: 32)
max_retries: Maximum number of retries for failed requests (default: 3)
retry_delay: Initial delay between retries in seconds, doubles each retry (default: 0.5)
"""
self.base_url = base_url.rstrip("/")
self.timeout = timeout
self.batch_size = batch_size
self.max_retries = max_retries
self.retry_delay = retry_delay
self._client: Optional[httpx.Client] = None
self._model_id: Optional[str] = None
@property
def provider_name(self) -> str:
return "tei"
def _request_with_retry(self, method: str, url: str, **kwargs) -> httpx.Response:
"""Make an HTTP request with automatic retries on transient errors."""
import time
last_error = None
delay = self.retry_delay
for attempt in range(self.max_retries + 1):
try:
if method == "GET":
response = self._client.get(url, **kwargs)
else:
response = self._client.post(url, **kwargs)
response.raise_for_status()
return response
except (httpx.ConnectError, httpx.ReadTimeout, httpx.WriteTimeout) as e:
last_error = e
if attempt < self.max_retries:
logger.warning(f"TEI request failed (attempt {attempt + 1}/{self.max_retries + 1}): {e}. Retrying in {delay}s...")
time.sleep(delay)
delay *= 2 # Exponential backoff
except httpx.HTTPStatusError as e:
# Retry on 5xx server errors
if e.response.status_code >= 500 and attempt < self.max_retries:
last_error = e
logger.warning(f"TEI server error (attempt {attempt + 1}/{self.max_retries + 1}): {e}. Retrying in {delay}s...")
time.sleep(delay)
delay *= 2
else:
raise
raise last_error
async def initialize(self) -> None:
"""Initialize the HTTP client and verify server connectivity."""
if self._client is not None:
return
logger.info(f"Embeddings: initializing TEI provider at {self.base_url}")
self._client = httpx.Client(timeout=self.timeout)
# Verify server is reachable and get model info
try:
response = self._request_with_retry("GET", f"{self.base_url}/info")
info = response.json()
self._model_id = info.get("model_id", "unknown")
logger.info(f"Embeddings: TEI provider initialized (model: {self._model_id})")
except httpx.HTTPError as e:
raise RuntimeError(f"Failed to connect to TEI server at {self.base_url}: {e}")
def encode(self, texts: List[str]) -> List[List[float]]:
"""
Generate embeddings using the remote TEI server.
Args:
texts: List of text strings to encode
Returns:
List of embedding vectors
"""
if self._client is None:
raise RuntimeError("Embeddings not initialized. Call initialize() first.")
if not texts:
return []
all_embeddings = []
# Process in batches
for i in range(0, len(texts), self.batch_size):
batch = texts[i:i + self.batch_size]
try:
response = self._request_with_retry(
"POST",
f"{self.base_url}/embed",
json={"inputs": batch},
)
batch_embeddings = response.json()
all_embeddings.extend(batch_embeddings)
except httpx.HTTPError as e:
raise RuntimeError(f"TEI embedding request failed: {e}")
return all_embeddings
def create_embeddings_from_env() -> Embeddings:
"""
Create an Embeddings instance based on environment variables.
See hindsight_api.config for environment variable names and defaults.
Returns:
Configured Embeddings instance
"""
provider = os.environ.get(ENV_EMBEDDINGS_PROVIDER, DEFAULT_EMBEDDINGS_PROVIDER).lower()
if provider == "tei":
url = os.environ.get(ENV_EMBEDDINGS_TEI_URL)
if not url:
raise ValueError(
f"{ENV_EMBEDDINGS_TEI_URL} is required when {ENV_EMBEDDINGS_PROVIDER} is 'tei'"
)
return RemoteTEIEmbeddings(base_url=url)
elif provider == "local":
model = os.environ.get(ENV_EMBEDDINGS_LOCAL_MODEL)
model_name = model or DEFAULT_EMBEDDINGS_LOCAL_MODEL
return LocalSTEmbeddings(model_name=model_name)
else:
raise ValueError(
f"Unknown embeddings provider: {provider}. Supported: 'local', 'tei'"
)
@@ -5,7 +5,7 @@ Uses spaCy for entity extraction and implements resolution logic
to disambiguate entities across memory units.
"""
import asyncpg
from typing import List, Dict, Optional, Set, Any
from typing import List, Dict, Optional, Set
from difflib import SequenceMatcher
from datetime import datetime, timezone
from .db_utils import acquire_with_retry
@@ -31,7 +31,7 @@ class EntityResolver:
async def resolve_entities_batch(
self,
bank_id: str,
agent_id: str,
entities_data: List[Dict],
context: str,
unit_event_date,
@@ -44,7 +44,7 @@ class EntityResolver:
all entities with minimal DB queries.
Args:
bank_id: bank ID
agent_id: Agent ID
entities_data: List of dicts with 'text', 'type', 'nearby_entities'
context: Context where entities appear
unit_event_date: When this unit was created
@@ -58,51 +58,24 @@ class EntityResolver:
if conn is None:
async with acquire_with_retry(self.pool) as conn:
return await self._resolve_entities_batch_impl(conn, bank_id, entities_data, context, unit_event_date)
return await self._resolve_entities_batch_impl(conn, agent_id, entities_data, context, unit_event_date)
else:
return await self._resolve_entities_batch_impl(conn, bank_id, entities_data, context, unit_event_date)
return await self._resolve_entities_batch_impl(conn, agent_id, entities_data, context, unit_event_date)
async def _resolve_entities_batch_impl(self, conn, bank_id: str, entities_data: List[Dict], context: str, unit_event_date) -> List[str]:
# Query ALL candidates for this bank
async def _resolve_entities_batch_impl(self, conn, agent_id: str, entities_data: List[Dict], context: str, unit_event_date) -> List[str]:
import time
start = time.time()
# Query ALL candidates for this agent
all_entities = await conn.fetch(
"""
SELECT canonical_name, id, metadata, last_seen, mention_count
FROM entities
WHERE bank_id = $1
WHERE agent_id = $1
""",
bank_id
agent_id
)
# Build entity ID to name mapping for co-occurrence lookups
entity_id_to_name = {row['id']: row['canonical_name'].lower() for row in all_entities}
# Query ALL co-occurrences for this bank's entities in one query
# This builds a map of entity_id -> set of co-occurring entity names
all_cooccurrences = await conn.fetch(
"""
SELECT ec.entity_id_1, ec.entity_id_2, ec.cooccurrence_count
FROM entity_cooccurrences ec
WHERE ec.entity_id_1 IN (SELECT id FROM entities WHERE bank_id = $1)
OR ec.entity_id_2 IN (SELECT id FROM entities WHERE bank_id = $1)
""",
bank_id
)
# Build co-occurrence map: entity_id -> set of co-occurring entity names (lowercase)
cooccurrence_map: Dict[str, Set[str]] = {}
for row in all_cooccurrences:
eid1, eid2 = row['entity_id_1'], row['entity_id_2']
# Add both directions
if eid1 not in cooccurrence_map:
cooccurrence_map[eid1] = set()
if eid2 not in cooccurrence_map:
cooccurrence_map[eid2] = set()
# Map to canonical names for comparison with nearby_entities
if eid2 in entity_id_to_name:
cooccurrence_map[eid1].add(entity_id_to_name[eid2])
if eid1 in entity_id_to_name:
cooccurrence_map[eid2].add(entity_id_to_name[eid1])
# Build candidate map for each entity text
all_candidates = {} # Maps entity_text -> list of candidates
entity_texts = list(set(e['text'] for e in entities_data))
@@ -126,32 +99,31 @@ class EntityResolver:
# Resolve each entity using pre-fetched candidates
entity_ids = [None] * len(entities_data)
entities_to_update = [] # (entity_id, event_date)
entities_to_create = [] # (idx, entity_data, event_date)
entities_to_update = [] # (entity_id, unit_event_date)
entities_to_create = [] # (idx, entity_data)
for idx, entity_data in enumerate(entities_data):
entity_text = entity_data['text']
nearby_entities = entity_data.get('nearby_entities', [])
# Use per-entity date if available, otherwise fall back to batch-level date
entity_event_date = entity_data.get('event_date', unit_event_date)
candidates = all_candidates.get(entity_text, [])
if not candidates:
# Will create new entity
entities_to_create.append((idx, entity_data, entity_event_date))
entities_to_create.append((idx, entity_data))
continue
# Score candidates
# Score candidates (same logic as before but with pre-fetched data)
best_candidate = None
best_score = 0.0
best_name_similarity = 0.0
nearby_entity_set = {e['text'].lower() for e in nearby_entities if e['text'] != entity_text}
for candidate_id, canonical_name, metadata, last_seen, mention_count in candidates:
score = 0.0
# 1. Name similarity (0-0.5)
# Name similarity
name_similarity = SequenceMatcher(
None,
entity_text.lower(),
@@ -159,19 +131,9 @@ class EntityResolver:
).ratio()
score += name_similarity * 0.5
# 2. Co-occurring entities (0-0.3)
if nearby_entity_set:
co_entities = cooccurrence_map.get(candidate_id, set())
overlap = len(nearby_entity_set & co_entities)
co_entity_score = overlap / len(nearby_entity_set)
score += co_entity_score * 0.3
# 3. Temporal proximity (0-0.2)
if last_seen and entity_event_date:
# Normalize timezone awareness for comparison
event_date_utc = entity_event_date if entity_event_date.tzinfo else entity_event_date.replace(tzinfo=timezone.utc)
last_seen_utc = last_seen if last_seen.tzinfo else last_seen.replace(tzinfo=timezone.utc)
days_diff = abs((event_date_utc - last_seen_utc).total_seconds() / 86400)
# Temporal proximity
if last_seen:
days_diff = abs((unit_event_date - last_seen).total_seconds() / 86400)
if days_diff < 7:
temporal_score = max(0, 1.0 - (days_diff / 7))
score += temporal_score * 0.2
@@ -179,15 +141,16 @@ class EntityResolver:
if score > best_score:
best_score = score
best_candidate = candidate_id
best_name_similarity = name_similarity
# Apply unified threshold
threshold = 0.6
if best_score > threshold:
entity_ids[idx] = best_candidate
entities_to_update.append((best_candidate, entity_event_date))
entities_to_update.append((best_candidate, unit_event_date))
else:
entities_to_create.append((idx, entity_data, entity_event_date))
entities_to_create.append((idx, entity_data))
# Batch update existing entities
if entities_to_update:
@@ -201,60 +164,45 @@ class EntityResolver:
entities_to_update
)
# Batch create new entities using COPY + INSERT for maximum speed
# This handles duplicates via ON CONFLICT and returns all IDs
# Batch create new entities using multi-row VALUES
if entities_to_create:
# Group entities by canonical name (lowercase) to handle duplicates within batch
# For duplicates, we only insert once and reuse the ID
unique_entities = {} # lowercase_name -> (entity_data, event_date, [indices])
for idx, entity_data, event_date in entities_to_create:
name_lower = entity_data['text'].lower()
if name_lower not in unique_entities:
unique_entities[name_lower] = (entity_data, event_date, [idx])
else:
# Same entity appears multiple times - add index to list
unique_entities[name_lower][2].append(idx)
import logging
# Build multi-row VALUES statement
# VALUES ($1, $2, ...), ($N+1, $N+2, ...), ...
values_clauses = []
params = []
param_idx = 1
# Batch insert unique entities and get their IDs
# Use a single query with unnest for speed
entity_names = []
entity_dates = []
indices_map = [] # Maps result index -> list of original indices
for idx, entity_data in entities_to_create:
values_clauses.append(f"(${param_idx}, ${param_idx+1}, ${param_idx+2}, ${param_idx+3}, ${param_idx+4})")
params.extend([
agent_id,
entity_data['text'],
unit_event_date,
unit_event_date,
1
])
param_idx += 5
for name_lower, (entity_data, event_date, indices) in unique_entities.items():
entity_names.append(entity_data['text'])
entity_dates.append(event_date)
indices_map.append(indices)
# Batch INSERT ... ON CONFLICT with RETURNING
# This is much faster than individual inserts
rows = await conn.fetch(
"""
INSERT INTO entities (bank_id, canonical_name, first_seen, last_seen, mention_count)
SELECT $1, name, event_date, event_date, 1
FROM unnest($2::text[], $3::timestamptz[]) AS t(name, event_date)
ON CONFLICT (bank_id, LOWER(canonical_name))
DO UPDATE SET
mention_count = entities.mention_count + 1,
last_seen = EXCLUDED.last_seen
# Single INSERT with multiple VALUES rows
query = f"""
INSERT INTO entities (agent_id, canonical_name, first_seen, last_seen, mention_count)
VALUES {', '.join(values_clauses)}
RETURNING id
""",
bank_id,
entity_names,
entity_dates
)
"""
created_rows = await conn.fetch(query, *params)
# Map created IDs back to original indices
for i, (idx, entity_data) in enumerate(entities_to_create):
entity_ids[idx] = created_rows[i]['id']
# Map returned IDs back to original indices
for result_idx, row in enumerate(rows):
entity_id = row['id']
for original_idx in indices_map[result_idx]:
entity_ids[original_idx] = entity_id
return entity_ids
async def resolve_entity(
self,
bank_id: str,
agent_id: str,
entity_text: str,
context: str,
nearby_entities: List[Dict],
@@ -264,7 +212,7 @@ class EntityResolver:
Resolve an entity to a canonical entity ID.
Args:
bank_id: bank ID (entities are scoped to agents)
agent_id: Agent ID (entities are scoped to agents)
entity_text: Entity text ("Alice", "Google", etc.)
context: Context where entity appears
nearby_entities: Other entities in the same unit
@@ -279,7 +227,7 @@ class EntityResolver:
"""
SELECT id, canonical_name, metadata, last_seen
FROM entities
WHERE bank_id = $1
WHERE agent_id = $1
AND (
canonical_name ILIKE $2
OR canonical_name ILIKE $3
@@ -287,13 +235,13 @@ class EntityResolver:
)
ORDER BY mention_count DESC
""",
bank_id, entity_text, f"%{entity_text}%"
agent_id, entity_text, f"%{entity_text}%"
)
if not candidates:
# New entity - create it
return await self._create_entity(
conn, bank_id, entity_text, unit_event_date
conn, agent_id, entity_text, unit_event_date
)
# Score candidates based on:
@@ -378,25 +326,22 @@ class EntityResolver:
else:
# Not confident - create new entity
return await self._create_entity(
conn, bank_id, entity_text, unit_event_date
conn, agent_id, entity_text, unit_event_date
)
async def _create_entity(
self,
conn,
bank_id: str,
agent_id: str,
entity_text: str,
event_date,
) -> str:
"""
Create a new entity or get existing one if it already exists.
Uses INSERT ... ON CONFLICT to handle race conditions where
two concurrent transactions try to create the same entity.
Create a new entity.
Args:
conn: Database connection
bank_id: bank ID
agent_id: Agent ID
entity_text: Entity text
event_date: When first seen
@@ -405,15 +350,11 @@ class EntityResolver:
"""
entity_id = await conn.fetchval(
"""
INSERT INTO entities (bank_id, canonical_name, first_seen, last_seen, mention_count)
INSERT INTO entities (agent_id, canonical_name, first_seen, last_seen, mention_count)
VALUES ($1, $2, $3, $4, 1)
ON CONFLICT (bank_id, LOWER(canonical_name))
DO UPDATE SET
mention_count = entities.mention_count + 1,
last_seen = EXCLUDED.last_seen
RETURNING id
""",
bank_id, entity_text, event_date, event_date
agent_id, entity_text, event_date, event_date
)
return entity_id
@@ -574,14 +515,14 @@ class EntityResolver:
async def get_entity_by_text(
self,
bank_id: str,
agent_id: str,
entity_text: str,
) -> Optional[str]:
"""
Find an entity by text (for query resolution).
Args:
bank_id: bank ID
agent_id: Agent ID
entity_text: Entity text to search for
Returns:
@@ -591,12 +532,12 @@ class EntityResolver:
row = await conn.fetchrow(
"""
SELECT id FROM entities
WHERE bank_id = $1
WHERE agent_id = $1
AND canonical_name ILIKE $2
ORDER BY mention_count DESC
LIMIT 1
""",
bank_id, entity_text
agent_id, entity_text
)
return row['id'] if row else None
File diff suppressed because it is too large Load Diff
@@ -5,136 +5,26 @@ Link creation utilities for temporal, semantic, and entity links.
import time
import logging
from typing import List
from datetime import timedelta, datetime, timezone
from uuid import UUID
from .types import EntityLink
from datetime import timedelta
logger = logging.getLogger(__name__)
def _normalize_datetime(dt):
"""Normalize datetime to be timezone-aware (UTC) for consistent comparison."""
if dt is None:
return None
if dt.tzinfo is None:
# Naive datetime - assume UTC
return dt.replace(tzinfo=timezone.utc)
return dt
def compute_temporal_links(
new_units: dict,
candidates: list,
time_window_hours: int = 24,
) -> list:
"""
Compute temporal links between new units and candidate neighbors.
This is a pure function that takes query results and returns link tuples,
making it easy to test without database access.
Args:
new_units: Dict mapping unit_id (str) to event_date (datetime)
candidates: List of dicts with 'id' and 'event_date' keys (candidate neighbors)
time_window_hours: Time window in hours for temporal links
Returns:
List of tuples: (from_unit_id, to_unit_id, 'temporal', weight, None)
"""
if not new_units:
return []
links = []
for unit_id, unit_event_date in new_units.items():
# Normalize unit_event_date for consistent comparison
unit_event_date_norm = _normalize_datetime(unit_event_date)
# Calculate time window bounds with overflow protection
try:
time_lower = unit_event_date_norm - timedelta(hours=time_window_hours)
except OverflowError:
time_lower = datetime.min.replace(tzinfo=timezone.utc)
try:
time_upper = unit_event_date_norm + timedelta(hours=time_window_hours)
except OverflowError:
time_upper = datetime.max.replace(tzinfo=timezone.utc)
# Filter candidates within this unit's time window
matching_neighbors = [
(row['id'], row['event_date'])
for row in candidates
if time_lower <= _normalize_datetime(row['event_date']) <= time_upper
][:10] # Limit to top 10
for recent_id, recent_event_date in matching_neighbors:
# Calculate temporal proximity weight
time_diff_hours = abs((unit_event_date_norm - _normalize_datetime(recent_event_date)).total_seconds() / 3600)
weight = max(0.3, 1.0 - (time_diff_hours / time_window_hours))
links.append((unit_id, str(recent_id), 'temporal', weight, None))
return links
def compute_temporal_query_bounds(
new_units: dict,
time_window_hours: int = 24,
) -> tuple:
"""
Compute the min/max date bounds for querying temporal neighbors.
Args:
new_units: Dict mapping unit_id (str) to event_date (datetime)
time_window_hours: Time window in hours
Returns:
Tuple of (min_date, max_date) with overflow protection
"""
if not new_units:
return None, None
# Normalize all dates to be timezone-aware to avoid comparison issues
all_dates = [_normalize_datetime(d) for d in new_units.values()]
try:
min_date = min(all_dates) - timedelta(hours=time_window_hours)
except OverflowError:
min_date = datetime.min.replace(tzinfo=timezone.utc)
try:
max_date = max(all_dates) + timedelta(hours=time_window_hours)
except OverflowError:
max_date = datetime.max.replace(tzinfo=timezone.utc)
return min_date, max_date
def _log(log_buffer, message, level='info'):
"""Helper to log to buffer if available, otherwise use logger.
Args:
log_buffer: Buffer to append messages to (for main output)
message: The log message
level: 'info', 'debug', 'warning', or 'error'. Debug messages are not added to buffer.
"""
if level == 'debug':
# Debug messages only go to logger, not to buffer
logger.debug(message)
return
"""Helper to log to buffer if available, otherwise use logger."""
if log_buffer is not None:
log_buffer.append(message)
else:
if level == 'info':
logger.info(message)
else:
logger.log(logging.WARNING if level == 'warning' else logging.ERROR, message)
logger.debug(message)
async def extract_entities_batch_optimized(
entity_resolver,
conn,
bank_id: str,
agent_id: str,
unit_ids: List[str],
sentences: List[str],
context: str,
@@ -151,7 +41,7 @@ async def extract_entities_batch_optimized(
Args:
entity_resolver: EntityResolver instance for entity resolution
conn: Database connection
agent_id: bank IDentifier
agent_id: Agent identifier
unit_ids: List of unit IDs
sentences: List of fact sentences
context: Context string
@@ -172,14 +62,13 @@ async def extract_entities_batch_optimized(
for ent in entity_list:
# Handle both Entity objects and dicts
if hasattr(ent, 'text'):
# Entity objects only have 'text', default type to 'CONCEPT'
formatted_entities.append({'text': ent.text, 'type': 'CONCEPT'})
formatted_entities.append({'text': ent.text, 'type': ent.type})
elif isinstance(ent, dict):
formatted_entities.append({'text': ent.get('text', ''), 'type': ent.get('type', 'CONCEPT')})
all_entities.append(formatted_entities)
total_entities = sum(len(ents) for ents in all_entities)
_log(log_buffer, f" [6.1] Process LLM entities: {total_entities} entities from {len(sentences)} facts in {time.time() - substep_start:.3f}s", level='debug')
_log(log_buffer, f" [6.1] Process LLM entities: {total_entities} entities from {len(sentences)} facts in {time.time() - substep_start:.3f}s")
# Step 2: Resolve entities in BATCH (much faster!)
substep_start = time.time()
@@ -201,28 +90,43 @@ async def extract_entities_batch_optimized(
'nearby_entities': entities,
})
entity_to_unit.append((unit_id, local_idx, fact_date))
_log(log_buffer, f" [6.2.1] Prepare entities: {len(all_entities_flat)} entities in {time.time() - substep_6_2_1_start:.3f}s", level='debug')
_log(log_buffer, f" [6.2.1] Prepare entities: {len(all_entities_flat)} entities in {time.time() - substep_6_2_1_start:.3f}s")
# Resolve ALL entities in one batch call
if all_entities_flat:
# [6.2.2] Batch resolve entities - single call with per-entity dates
# [6.2.2] Batch resolve entities
substep_6_2_2_start = time.time()
# Add per-entity dates to entity data for batch resolution
# Group by date for batch resolution (most will have same date)
entities_by_date = {}
for idx, (unit_id, local_idx, fact_date) in enumerate(entity_to_unit):
all_entities_flat[idx]['event_date'] = fact_date
date_key = fact_date
if date_key not in entities_by_date:
entities_by_date[date_key] = []
entities_by_date[date_key].append((idx, all_entities_flat[idx]))
# Resolve ALL entities in ONE batch call (much faster than sequential buckets)
# INSERT ... ON CONFLICT handles any race conditions at the DB level
resolved_entity_ids = await entity_resolver.resolve_entities_batch(
bank_id=bank_id,
entities_data=all_entities_flat,
context=context,
unit_event_date=None, # Not used when per-entity dates provided
conn=conn # Use main transaction connection
)
_log(log_buffer, f" [6.2.2] Grouped into {len(entities_by_date)} date buckets, resolving...")
_log(log_buffer, f" [6.2.2] Resolve entities: {len(all_entities_flat)} entities in single batch in {time.time() - substep_6_2_2_start:.3f}s", level='debug')
# Resolve each date group in batch
resolved_entity_ids = [None] * len(all_entities_flat)
for date_idx, (fact_date, entities_group) in enumerate(entities_by_date.items(), 1):
date_bucket_start = time.time()
indices = [idx for idx, _ in entities_group]
entities_data = [entity_data for _, entity_data in entities_group]
batch_resolved = await entity_resolver.resolve_entities_batch(
agent_id=agent_id,
entities_data=entities_data,
context=context,
unit_event_date=fact_date,
conn=conn
)
for idx, entity_id in zip(indices, batch_resolved):
resolved_entity_ids[idx] = entity_id
_log(log_buffer, f" [6.2.2.{date_idx}] Resolved {len(entities_data)} entities in {time.time() - date_bucket_start:.3f}s")
_log(log_buffer, f" [6.2.2] Resolve entities: {len(all_entities_flat)} entities in {time.time() - substep_6_2_2_start:.3f}s")
# [6.2.3] Create unit-entity links in BATCH
substep_6_2_3_start = time.time()
@@ -239,12 +143,12 @@ async def extract_entities_batch_optimized(
# Batch insert all unit-entity links (MUCH faster!)
await entity_resolver.link_units_to_entities_batch(unit_entity_pairs, conn=conn)
_log(log_buffer, f" [6.2.3] Create unit-entity links (batched): {len(unit_entity_pairs)} links in {time.time() - substep_6_2_3_start:.3f}s", level='debug')
_log(log_buffer, f" [6.2.3] Create unit-entity links (batched): {len(unit_entity_pairs)} links in {time.time() - substep_6_2_3_start:.3f}s")
_log(log_buffer, f" [6.2] Entity resolution (batched): {len(all_entities_flat)} entities resolved in {time.time() - step_6_2_start:.3f}s", level='debug')
_log(log_buffer, f" [6.2] Entity resolution (batched): {len(all_entities_flat)} entities resolved in {time.time() - step_6_2_start:.3f}s")
else:
unit_to_entity_ids = {}
_log(log_buffer, f" [6.2] Entity resolution (batched): 0 entities in {time.time() - step_6_2_start:.3f}s", level='debug')
_log(log_buffer, f" [6.2] Entity resolution (batched): 0 entities in {time.time() - step_6_2_start:.3f}s")
# Step 3: Create entity links between units that share entities
substep_start = time.time()
@@ -253,7 +157,7 @@ async def extract_entities_batch_optimized(
for entity_ids in unit_to_entity_ids.values():
all_entity_ids.update(entity_ids)
_log(log_buffer, f" [6.3] Creating entity links for {len(all_entity_ids)} unique entities...", level='debug')
_log(log_buffer, f" [6.3] Creating entity links for {len(all_entity_ids)} unique entities...")
# Find all units that reference these entities (ONE batched query)
entity_to_units = {}
@@ -269,7 +173,7 @@ async def extract_entities_batch_optimized(
""",
entity_id_list
)
_log(log_buffer, f" [6.3.1] Query unit_entities: {len(rows)} rows in {time.time() - query_start:.3f}s", level='debug')
_log(log_buffer, f" [6.3.1] Query unit_entities: {len(rows)} rows in {time.time() - query_start:.3f}s")
# Group by entity_id
group_start = time.time()
@@ -278,42 +182,21 @@ async def extract_entities_batch_optimized(
if entity_id not in entity_to_units:
entity_to_units[entity_id] = []
entity_to_units[entity_id].append(row['unit_id'])
_log(log_buffer, f" [6.3.2] Group by entity_id: {time.time() - group_start:.3f}s", level='debug')
_log(log_buffer, f" [6.3.2] Group by entity_id: {time.time() - group_start:.3f}s")
# Create bidirectional links between units that share entities
# OPTIMIZATION: Limit links per entity to avoid N² explosion
# Only link each new unit to the most recent MAX_LINKS_PER_ENTITY units
MAX_LINKS_PER_ENTITY = 50 # Limit to prevent explosion when entity appears in many facts
link_gen_start = time.time()
links: List[EntityLink] = []
new_unit_set = set(unit_ids) # Units from this batch
def to_uuid(val) -> UUID:
return UUID(val) if isinstance(val, str) else val
links = []
for entity_id, units_with_entity in entity_to_units.items():
entity_uuid = to_uuid(entity_id)
# Separate new units (from this batch) and existing units
new_units = [u for u in units_with_entity if str(u) in new_unit_set or u in new_unit_set]
existing_units = [u for u in units_with_entity if str(u) not in new_unit_set and u not in new_unit_set]
# For each pair of units with this entity, create bidirectional links
for i, unit_id_1 in enumerate(units_with_entity):
for unit_id_2 in units_with_entity[i+1:]:
# Bidirectional links
links.append((unit_id_1, unit_id_2, 'entity', 1.0, entity_id))
links.append((unit_id_2, unit_id_1, 'entity', 1.0, entity_id))
# Link new units to each other (within batch) - also limited
# For very common entities, limit within-batch links too
new_units_to_link = new_units[-MAX_LINKS_PER_ENTITY:] if len(new_units) > MAX_LINKS_PER_ENTITY else new_units
for i, unit_id_1 in enumerate(new_units_to_link):
for unit_id_2 in new_units_to_link[i+1:]:
links.append(EntityLink(from_unit_id=to_uuid(unit_id_1), to_unit_id=to_uuid(unit_id_2), entity_id=entity_uuid))
links.append(EntityLink(from_unit_id=to_uuid(unit_id_2), to_unit_id=to_uuid(unit_id_1), entity_id=entity_uuid))
# Link new units to LIMITED existing units (most recent)
existing_to_link = existing_units[-MAX_LINKS_PER_ENTITY:] # Take most recent
for new_unit in new_units:
for existing_unit in existing_to_link:
links.append(EntityLink(from_unit_id=to_uuid(new_unit), to_unit_id=to_uuid(existing_unit), entity_id=entity_uuid))
links.append(EntityLink(from_unit_id=to_uuid(existing_unit), to_unit_id=to_uuid(new_unit), entity_id=entity_uuid))
_log(log_buffer, f" [6.3.3] Generate {len(links)} links: {time.time() - link_gen_start:.3f}s", level='debug')
_log(log_buffer, f" [6.3] Entity link creation: {len(links)} links for {len(all_entity_ids)} unique entities in {time.time() - substep_start:.3f}s", level='debug')
_log(log_buffer, f" [6.3.3] Generate {len(links)} links: {time.time() - link_gen_start:.3f}s")
_log(log_buffer, f" [6.3] Entity link creation: {len(links)} links for {len(all_entity_ids)} unique entities in {time.time() - substep_start:.3f}s")
return links
@@ -326,11 +209,11 @@ async def extract_entities_batch_optimized(
async def create_temporal_links_batch_per_fact(
conn,
bank_id: str,
agent_id: str,
unit_ids: List[str],
time_window_hours: int = 24,
log_buffer: List[str] = None,
) -> int:
):
"""
Create temporal links for multiple units, each with their own event_date.
@@ -339,16 +222,13 @@ async def create_temporal_links_batch_per_fact(
Args:
conn: Database connection
agent_id: bank IDentifier
agent_id: Agent identifier
unit_ids: List of unit IDs
time_window_hours: Time window in hours for temporal links
log_buffer: Optional buffer for logging
Returns:
Number of temporal links created
"""
if not unit_ids:
return 0
return
try:
import time as time_mod
@@ -367,20 +247,22 @@ async def create_temporal_links_batch_per_fact(
_log(log_buffer, f" [7.1] Fetch event_dates for {len(unit_ids)} units: {time_mod.time() - fetch_dates_start:.3f}s")
# Fetch ALL potential temporal neighbors in ONE query (much faster!)
# Get time range across all units with overflow protection
min_date, max_date = compute_temporal_query_bounds(new_units, time_window_hours)
# Get time range across all units
all_dates = list(new_units.values())
min_date = min(all_dates) - timedelta(hours=time_window_hours)
max_date = max(all_dates) + timedelta(hours=time_window_hours)
fetch_neighbors_start = time_mod.time()
all_candidates = await conn.fetch(
"""
SELECT id, event_date
FROM memory_units
WHERE bank_id = $1
WHERE agent_id = $1
AND event_date BETWEEN $2 AND $3
AND id::text != ALL($4)
ORDER BY event_date DESC
""",
bank_id,
agent_id,
min_date,
max_date,
unit_ids
@@ -389,27 +271,23 @@ async def create_temporal_links_batch_per_fact(
# Filter and create links in memory (much faster than N queries)
link_gen_start = time_mod.time()
links = compute_temporal_links(new_units, all_candidates, time_window_hours)
links = []
for unit_id, unit_event_date in new_units.items():
# Filter candidates within this unit's time window
time_lower = unit_event_date - timedelta(hours=time_window_hours)
time_upper = unit_event_date + timedelta(hours=time_window_hours)
# Also compute temporal links WITHIN the new batch (new units to each other)
if len(new_units) > 1:
# Convert new_units dict to candidate format for within-batch linking
new_unit_items = list(new_units.items())
for i, (unit_id, event_date) in enumerate(new_unit_items):
unit_event_date_norm = _normalize_datetime(event_date)
matching_neighbors = [
(row['id'], row['event_date'])
for row in all_candidates
if time_lower <= row['event_date'] <= time_upper
][:10] # Limit to top 10
# Compare with other new units (only those after this one to avoid duplicates)
for j in range(i + 1, len(new_unit_items)):
other_id, other_event_date = new_unit_items[j]
other_event_date_norm = _normalize_datetime(other_event_date)
# Check if within time window
time_diff_hours = abs((unit_event_date_norm - other_event_date_norm).total_seconds() / 3600)
if time_diff_hours <= time_window_hours:
weight = max(0.3, 1.0 - (time_diff_hours / time_window_hours))
# Create bidirectional links
links.append((unit_id, other_id, 'temporal', weight, None))
links.append((other_id, unit_id, 'temporal', weight, None))
for recent_id, recent_event_date in matching_neighbors:
# Calculate temporal proximity weight
time_diff_hours = abs((unit_event_date - recent_event_date).total_seconds() / 3600)
weight = max(0.3, 1.0 - (time_diff_hours / time_window_hours))
links.append((unit_id, str(recent_id), 'temporal', weight, None))
_log(log_buffer, f" [7.3] Generate {len(links)} temporal links: {time_mod.time() - link_gen_start:.3f}s")
@@ -425,8 +303,6 @@ async def create_temporal_links_batch_per_fact(
)
_log(log_buffer, f" [7.4] Insert {len(links)} temporal links: {time_mod.time() - insert_start:.3f}s")
return len(links)
except Exception as e:
logger.error(f"Failed to create temporal links: {str(e)}")
import traceback
@@ -436,13 +312,13 @@ async def create_temporal_links_batch_per_fact(
async def create_semantic_links_batch(
conn,
bank_id: str,
agent_id: str,
unit_ids: List[str],
embeddings: List[List[float]],
top_k: int = 5,
threshold: float = 0.7,
log_buffer: List[str] = None,
) -> int:
):
"""
Create semantic links for multiple units efficiently.
@@ -450,18 +326,15 @@ async def create_semantic_links_batch(
Args:
conn: Database connection
agent_id: bank IDentifier
agent_id: Agent identifier
unit_ids: List of unit IDs
embeddings: List of embedding vectors
top_k: Number of top similar units to link
threshold: Minimum similarity threshold
log_buffer: Optional buffer for logging
Returns:
Number of semantic links created
"""
if not unit_ids or not embeddings:
return 0
return
try:
import time as time_mod
@@ -473,11 +346,11 @@ async def create_semantic_links_batch(
"""
SELECT id, embedding
FROM memory_units
WHERE bank_id = $1
WHERE agent_id = $1
AND embedding IS NOT NULL
AND id::text != ALL($2)
""",
bank_id,
agent_id,
unit_ids
)
_log(log_buffer, f" [8.1] Fetch {len(all_existing)} existing embeddings (1 query): {time_mod.time() - fetch_start:.3f}s")
@@ -535,38 +408,9 @@ async def create_semantic_links_batch(
for idx in sorted_indices:
similar_id = existing_ids[idx]
# Clamp to [0, 1] to handle floating point precision issues
similarity = float(min(1.0, max(0.0, similarities[idx])))
similarity = float(similarities[idx])
all_links.append((unit_id, similar_id, 'semantic', similarity, None))
# Also compute similarities WITHIN the new batch (new units to each other)
# Apply the same top_k limit per unit as we do for existing units
if len(unit_ids) > 1:
new_embeddings_matrix = np.array(embeddings)
for i, unit_id in enumerate(unit_ids):
# Compute similarities with all OTHER new units
other_indices = [j for j in range(len(unit_ids)) if j != i]
if not other_indices:
continue
other_embeddings = new_embeddings_matrix[other_indices]
similarities = np.dot(other_embeddings, new_embeddings_matrix[i])
# Find top-k above threshold (same logic as existing units)
above_threshold = np.where(similarities >= threshold)[0]
if len(above_threshold) > 0:
# Sort by similarity (descending) and take top-k
sorted_local_indices = above_threshold[np.argsort(-similarities[above_threshold])][:top_k]
for local_idx in sorted_local_indices:
other_idx = other_indices[local_idx]
other_id = unit_ids[other_idx]
# Clamp to [0, 1] to handle floating point precision issues
similarity = float(min(1.0, max(0.0, similarities[local_idx])))
all_links.append((unit_id, other_id, 'semantic', similarity, None))
_log(log_buffer, f" [8.2] Compute similarities & generate {len(all_links)} semantic links: {time_mod.time() - compute_start:.3f}s")
if all_links:
@@ -581,8 +425,6 @@ async def create_semantic_links_batch(
)
_log(log_buffer, f" [8.3] Insert {len(all_links)} semantic links: {time_mod.time() - insert_start:.3f}s")
return len(all_links)
except Exception as e:
logger.error(f"Failed to create semantic links: {str(e)}")
import traceback
@@ -590,77 +432,28 @@ async def create_semantic_links_batch(
raise
async def insert_entity_links_batch(conn, links: List[EntityLink], chunk_size: int = 50000):
async def insert_entity_links_batch(conn, links: List[tuple]):
"""
Insert all entity links using COPY to temp table + INSERT for maximum speed.
Uses PostgreSQL COPY (via copy_records_to_table) for bulk loading,
then INSERT ... ON CONFLICT from temp table. This is the fastest
method for bulk inserts with conflict handling.
Insert all entity links in a single batch.
Args:
conn: Database connection
links: List of EntityLink objects
chunk_size: Number of rows per batch (default 50000)
links: List of tuples (from_unit_id, to_unit_id, link_type, weight, entity_id)
"""
if not links:
return
import uuid as uuid_mod
import time as time_mod
total_start = time_mod.time()
# Create temp table for bulk loading
create_start = time_mod.time()
await conn.execute("""
CREATE TEMP TABLE IF NOT EXISTS _temp_entity_links (
from_unit_id uuid,
to_unit_id uuid,
link_type text,
weight float,
entity_id uuid
) ON COMMIT DROP
""")
logger.debug(f" [9.1] Create temp table: {time_mod.time() - create_start:.3f}s")
# Clear any existing data in temp table
truncate_start = time_mod.time()
await conn.execute("TRUNCATE _temp_entity_links")
logger.debug(f" [9.2] Truncate temp table: {time_mod.time() - truncate_start:.3f}s")
# Convert EntityLink objects to tuples for COPY
convert_start = time_mod.time()
records = []
for link in links:
records.append((
link.from_unit_id,
link.to_unit_id,
link.link_type,
link.weight,
link.entity_id
))
logger.debug(f" [9.3] Convert {len(records)} records: {time_mod.time() - convert_start:.3f}s")
# Bulk load using COPY (fastest method)
copy_start = time_mod.time()
await conn.copy_records_to_table(
'_temp_entity_links',
records=records,
columns=['from_unit_id', 'to_unit_id', 'link_type', 'weight', 'entity_id']
)
logger.debug(f" [9.4] COPY {len(records)} records to temp table: {time_mod.time() - copy_start:.3f}s")
# Insert from temp table with ON CONFLICT (single query for all rows)
insert_start = time_mod.time()
await conn.execute("""
INSERT INTO memory_links (from_unit_id, to_unit_id, link_type, weight, entity_id)
SELECT from_unit_id, to_unit_id, link_type, weight, entity_id
FROM _temp_entity_links
ON CONFLICT (from_unit_id, to_unit_id, link_type, COALESCE(entity_id, '00000000-0000-0000-0000-000000000000'::uuid)) DO NOTHING
""")
logger.debug(f" [9.5] INSERT from temp table: {time_mod.time() - insert_start:.3f}s")
logger.debug(f" [9.TOTAL] Entity links batch insert: {time_mod.time() - total_start:.3f}s")
try:
await conn.executemany(
"""
INSERT INTO memory_links (from_unit_id, to_unit_id, link_type, weight, entity_id)
VALUES ($1, $2, $3, $4, $5)
ON CONFLICT (from_unit_id, to_unit_id, link_type, COALESCE(entity_id, '00000000-0000-0000-0000-000000000000'::uuid)) DO NOTHING
""",
links
)
except Exception as e:
logger.warning(f"Failed to insert entity links: {str(e)}")
async def create_causal_links_batch(
@@ -709,16 +502,6 @@ async def create_causal_links_batch(
relation_type = relation['relation_type']
strength = relation.get('strength', 1.0)
# Validate relation_type - must match database constraint
valid_types = {'causes', 'caused_by', 'enables', 'prevents'}
if relation_type not in valid_types:
logger.error(
f"Invalid relation_type '{relation_type}' (type: {type(relation_type).__name__}) "
f"from fact {fact_idx}. Must be one of: {valid_types}. "
f"Relation data: {relation}"
)
continue
# Validate target index
if target_idx < 0 or target_idx >= len(unit_ids):
logger.warning(f"Invalid target_fact_index {target_idx} in causal relation from fact {fact_idx}")
@@ -735,25 +518,19 @@ async def create_causal_links_batch(
# weight is the strength of the relationship
links.append((from_unit_id, to_unit_id, relation_type, strength, None))
logger.debug(f"Generated {len(links)} causal links in {time_mod.time() - create_start:.3f}s")
if links:
insert_start = time_mod.time()
try:
await conn.executemany(
"""
INSERT INTO memory_links (from_unit_id, to_unit_id, link_type, weight, entity_id)
VALUES ($1, $2, $3, $4, $5)
ON CONFLICT (from_unit_id, to_unit_id, link_type, COALESCE(entity_id, '00000000-0000-0000-0000-000000000000'::uuid)) DO NOTHING
""",
links
)
except Exception as db_error:
# Log the actual data being inserted for debugging
logger.error(f"Database insert failed for causal links. Error: {db_error}")
logger.error(f"Attempted to insert {len(links)} links. First few:")
for i, link in enumerate(links[:3]):
logger.error(f" Link {i}: from={link[0]}, to={link[1]}, type='{link[2]}' (repr={repr(link[2])}), weight={link[3]}, entity={link[4]}")
raise
await conn.executemany(
"""
INSERT INTO memory_links (from_unit_id, to_unit_id, link_type, weight, entity_id)
VALUES ($1, $2, $3, $4, $5)
ON CONFLICT (from_unit_id, to_unit_id, link_type, COALESCE(entity_id, '00000000-0000-0000-0000-000000000000'::uuid)) DO NOTHING
""",
links
)
logger.debug(f"Inserted {len(links)} causal links in {time_mod.time() - insert_start:.3f}s")
return len(links)
+111 -370
View File
@@ -5,23 +5,14 @@ import os
import time
import asyncio
from typing import Optional, Any, Dict, List
from openai import AsyncOpenAI, RateLimitError, APIError, APIStatusError, APIConnectionError, LengthFinishReasonError
from google import genai
from google.genai import types as genai_types
from google.genai import errors as genai_errors
from openai import AsyncOpenAI, RateLimitError, APIError, APIStatusError, LengthFinishReasonError
import logging
# Seed applied to every Groq request for deterministic behavior.
DEFAULT_LLM_SEED = 4242
logger = logging.getLogger(__name__)
# Disable httpx logging
logging.getLogger("httpx").setLevel(logging.WARNING)
# Global semaphore to limit concurrent LLM requests across all instances
_global_llm_semaphore = asyncio.Semaphore(32)
class OutputTooLongError(Exception):
"""
@@ -34,12 +25,8 @@ class OutputTooLongError(Exception):
pass
class LLMProvider:
"""
Unified LLM provider.
Supports OpenAI, Groq, Ollama (OpenAI-compatible), and Gemini.
"""
class LLMConfig:
"""Configuration for an LLM provider."""
def __init__(
self,
@@ -47,29 +34,25 @@ class LLMProvider:
api_key: str,
base_url: str,
model: str,
reasoning_effort: str = "low",
):
"""
Initialize LLM provider.
Initialize LLM configuration.
Args:
provider: Provider name ("openai", "groq", "ollama", "gemini").
api_key: API key.
base_url: Base URL for the API.
model: Model name.
reasoning_effort: Reasoning effort level for supported providers.
provider: Provider name ("openai", "groq", "ollama"). Required.
api_key: API key. Required.
base_url: Base URL. Required.
model: Model name. Required.
"""
self.provider = provider.lower()
self.api_key = api_key
self.base_url = base_url
self.model = model
self.reasoning_effort = reasoning_effort
# Validate provider
valid_providers = ["openai", "groq", "ollama", "gemini"]
if self.provider not in valid_providers:
if self.provider not in ["openai", "groq", "ollama"]:
raise ValueError(
f"Invalid LLM provider: {self.provider}. Must be one of: {', '.join(valid_providers)}"
f"Invalid LLM provider: {self.provider}. Must be 'openai', 'groq', or 'ollama'."
)
# Set default base URLs
@@ -81,412 +64,170 @@ class LLMProvider:
# Validate API key (not needed for ollama)
if self.provider != "ollama" and not self.api_key:
raise ValueError(f"API key not found for {self.provider}")
# Create client based on provider
if self.provider == "gemini":
self._gemini_client = genai.Client(api_key=self.api_key)
self._client = None
elif self.provider == "ollama":
self._client = AsyncOpenAI(api_key="ollama", base_url=self.base_url, max_retries=0)
self._gemini_client = None
else:
# Only pass base_url if it's set (OpenAI uses default URL otherwise)
client_kwargs = {"api_key": self.api_key, "max_retries": 0}
if self.base_url:
client_kwargs["base_url"] = self.base_url
self._client = AsyncOpenAI(**client_kwargs)
self._gemini_client = None
async def verify_connection(self) -> None:
"""
Verify that the LLM provider is configured correctly by making a simple test call.
Raises:
RuntimeError: If the connection test fails.
"""
try:
logger.info(f"Verifying LLM: provider={self.provider}, model={self.model}, base_url={self.base_url or 'default'}...")
await self.call(
messages=[{"role": "user", "content": "Say 'ok'"}],
max_completion_tokens=10,
max_retries=2,
initial_backoff=0.5,
max_backoff=2.0,
raise ValueError(
f"API key not found for {self.provider}"
)
# If we get here without exception, the connection is working
logger.info(f"LLM verified: {self.provider}/{self.model}")
except Exception as e:
raise RuntimeError(
f"LLM connection verification failed for {self.provider}/{self.model}: {e}"
) from e
# Create client (private - use .call() method instead)
if self.provider == "ollama":
self._client = AsyncOpenAI(api_key="ollama", base_url=self.base_url)
elif self.base_url:
self._client = AsyncOpenAI(api_key=self.api_key, base_url=self.base_url)
else:
self._client = AsyncOpenAI(api_key=self.api_key)
logger.info(
f"Initialized LLM: provider={self.provider}, model={self.model}, base_url={self.base_url}"
)
async def call(
self,
messages: List[Dict[str, str]],
response_format: Optional[Any] = None,
max_completion_tokens: Optional[int] = None,
temperature: Optional[float] = None,
scope: str = "memory",
max_retries: int = 10,
max_retries: int = 5,
initial_backoff: float = 1.0,
max_backoff: float = 60.0,
skip_validation: bool = False,
**kwargs
) -> Any:
"""
Make an LLM API call with retry logic.
Make an LLM API call with consistent configuration and retry logic.
Args:
messages: List of message dicts with 'role' and 'content'.
response_format: Optional Pydantic model for structured output.
max_completion_tokens: Maximum tokens in response.
temperature: Sampling temperature (0.0-2.0).
scope: Scope identifier for tracking.
max_retries: Maximum retry attempts.
initial_backoff: Initial backoff time in seconds.
max_backoff: Maximum backoff time in seconds.
skip_validation: Return raw JSON without Pydantic validation.
messages: List of message dicts with 'role' and 'content'
response_format: Optional Pydantic model for structured output
scope: Scope identifier (e.g., 'memory', 'judge') for future tracking
max_retries: Maximum number of retry attempts (default: 5)
initial_backoff: Initial backoff time in seconds (default: 1.0)
max_backoff: Maximum backoff time in seconds (default: 60.0)
**kwargs: Additional parameters to pass to the API (temperature, max_tokens, etc.)
Returns:
Parsed response if response_format is provided, otherwise text content.
Parsed response if response_format is provided, otherwise the text content
Raises:
OutputTooLongError: If output exceeds token limits.
Exception: Re-raises API errors after retries exhausted.
Exception: Re-raises any API errors after all retries are exhausted
"""
async with _global_llm_semaphore:
start_time = time.time()
import json
start_time = time.time()
# Handle Gemini provider separately
if self.provider == "gemini":
return await self._call_gemini(
messages, response_format, max_retries, initial_backoff,
max_backoff, skip_validation, start_time
)
call_params = {
"model": self.model,
"messages": messages,
}
# Check if model supports reasoning parameter (o1, o3, gpt-5 families)
model_lower = self.model.lower()
is_reasoning_model = any(x in model_lower for x in ["gpt-5", "o1", "o3"])
# For GPT-4 and GPT-4.1 models, cap max_completion_tokens to 32000
# For GPT-4o models, cap to 16384
is_gpt4_model = any(x in model_lower for x in ["gpt-4.1", "gpt-4-"])
is_gpt4o_model = "gpt-4o" in model_lower
if max_completion_tokens is not None:
if is_gpt4o_model and max_completion_tokens > 16384:
max_completion_tokens = 16384
elif is_gpt4_model and max_completion_tokens > 32000:
max_completion_tokens = 32000
# For reasoning models, max_completion_tokens includes reasoning + output tokens
# Enforce minimum of 16000 to ensure enough space for both
if is_reasoning_model and max_completion_tokens < 16000:
max_completion_tokens = 16000
call_params["max_completion_tokens"] = max_completion_tokens
# GPT-5/o1/o3 family doesn't support custom temperature (only default 1)
if temperature is not None and not is_reasoning_model:
call_params["temperature"] = temperature
# Set reasoning_effort for reasoning models (OpenAI gpt-5, o1, o3)
if is_reasoning_model and self.provider == "openai":
call_params["reasoning_effort"] = self.reasoning_effort
# Provider-specific parameters
if self.provider == "groq":
call_params["seed"] = DEFAULT_LLM_SEED
extra_body = {"service_tier": "auto"}
# Only add reasoning parameters for reasoning models
if is_reasoning_model:
extra_body["reasoning_effort"] = self.reasoning_effort
extra_body["include_reasoning"] = False
call_params["extra_body"] = extra_body
last_exception = None
for attempt in range(max_retries + 1):
try:
if response_format is not None:
# Add schema to system message for JSON mode
if hasattr(response_format, 'model_json_schema'):
schema = response_format.model_json_schema()
schema_msg = f"\n\nYou must respond with valid JSON matching this schema:\n{json.dumps(schema, indent=2)}"
if call_params['messages'] and call_params['messages'][0].get('role') == 'system':
call_params['messages'][0]['content'] += schema_msg
elif call_params['messages']:
call_params['messages'][0]['content'] = schema_msg + "\n\n" + call_params['messages'][0]['content']
call_params['response_format'] = {"type": "json_object"}
response = await self._client.chat.completions.create(**call_params)
content = response.choices[0].message.content
json_data = json.loads(content)
if skip_validation:
result = json_data
else:
result = response_format.model_validate(json_data)
else:
response = await self._client.chat.completions.create(**call_params)
result = response.choices[0].message.content
# Log slow calls
duration = time.time() - start_time
usage = response.usage
if duration > 10.0:
ratio = max(1, usage.completion_tokens) / usage.prompt_tokens
cached_tokens = 0
if hasattr(usage, 'prompt_tokens_details') and usage.prompt_tokens_details:
cached_tokens = getattr(usage.prompt_tokens_details, 'cached_tokens', 0) or 0
cache_info = f", cached_tokens={cached_tokens}" if cached_tokens > 0 else ""
logger.info(
f"slow llm call: model={self.provider}/{self.model}, "
f"input_tokens={usage.prompt_tokens}, output_tokens={usage.completion_tokens}, "
f"total_tokens={usage.total_tokens}{cache_info}, time={duration:.3f}s, ratio out/in={ratio:.2f}"
)
return result
except LengthFinishReasonError as e:
logger.warning(f"LLM output exceeded token limits: {str(e)}")
raise OutputTooLongError(
f"LLM output exceeded token limits. Input may need to be split into smaller chunks."
) from e
except APIConnectionError as e:
last_exception = e
if attempt < max_retries:
status_code = getattr(e, 'status_code', None) or getattr(getattr(e, 'response', None), 'status_code', None)
logger.warning(f"Connection error, retrying... (attempt {attempt + 1}/{max_retries + 1}) - status_code={status_code}, message={e}")
backoff = min(initial_backoff * (2 ** attempt), max_backoff)
await asyncio.sleep(backoff)
continue
else:
logger.error(f"Connection error after {max_retries + 1} attempts: {str(e)}")
raise
except APIStatusError as e:
# Fast fail only on 401 (unauthorized) and 403 (forbidden) - these won't recover with retries
if e.status_code in (401, 403):
logger.error(f"Auth error (HTTP {e.status_code}), not retrying: {str(e)}")
raise
last_exception = e
if attempt < max_retries:
backoff = min(initial_backoff * (2 ** attempt), max_backoff)
jitter = backoff * 0.2 * (2 * (time.time() % 1) - 1)
sleep_time = backoff + jitter
await asyncio.sleep(sleep_time)
else:
logger.error(f"API error after {max_retries + 1} attempts: {str(e)}")
raise
except Exception as e:
logger.error(f"Unexpected error during LLM call: {type(e).__name__}: {str(e)}")
raise
if last_exception:
raise last_exception
raise RuntimeError(f"LLM call failed after all retries with no exception captured")
async def _call_gemini(
self,
messages: List[Dict[str, str]],
response_format: Optional[Any],
max_retries: int,
initial_backoff: float,
max_backoff: float,
skip_validation: bool,
start_time: float,
) -> Any:
"""Handle Gemini-specific API calls."""
import json
# Convert OpenAI-style messages to Gemini format
system_instruction = None
gemini_contents = []
for msg in messages:
role = msg.get('role', 'user')
content = msg.get('content', '')
if role == 'system':
if system_instruction:
system_instruction += "\n\n" + content
else:
system_instruction = content
elif role == 'assistant':
gemini_contents.append(genai_types.Content(
role="model",
parts=[genai_types.Part(text=content)]
))
else:
gemini_contents.append(genai_types.Content(
role="user",
parts=[genai_types.Part(text=content)]
))
# Add JSON schema instruction if response_format is provided
if response_format is not None and hasattr(response_format, 'model_json_schema'):
schema = response_format.model_json_schema()
schema_msg = f"\n\nYou must respond with valid JSON matching this schema:\n{json.dumps(schema, indent=2)}"
if system_instruction:
system_instruction += schema_msg
else:
system_instruction = schema_msg
# Build generation config
config_kwargs = {}
if system_instruction:
config_kwargs['system_instruction'] = system_instruction
if response_format is not None:
config_kwargs['response_mime_type'] = 'application/json'
config_kwargs['response_schema'] = response_format
generation_config = genai_types.GenerateContentConfig(**config_kwargs) if config_kwargs else None
call_params = {
"model": self.model,
"messages": messages,
**kwargs
}
if self.provider == "groq":
call_params["extra_body"] = {"service_tier": "auto"}
last_exception = None
for attempt in range(max_retries + 1):
try:
response = await self._gemini_client.aio.models.generate_content(
model=self.model,
contents=gemini_contents,
config=generation_config,
)
content = response.text
# Handle empty response
if content is None:
block_reason = None
if hasattr(response, 'candidates') and response.candidates:
candidate = response.candidates[0]
if hasattr(candidate, 'finish_reason'):
block_reason = candidate.finish_reason
if attempt < max_retries:
logger.warning(f"Gemini returned empty response (reason: {block_reason}), retrying...")
backoff = min(initial_backoff * (2 ** attempt), max_backoff)
await asyncio.sleep(backoff)
continue
else:
raise RuntimeError(f"Gemini returned empty response after {max_retries + 1} attempts")
if response_format is not None:
json_data = json.loads(content)
if skip_validation:
result = json_data
else:
result = response_format.model_validate(json_data)
else:
result = content
# Log slow calls
duration = time.time() - start_time
if duration > 10.0 and hasattr(response, 'usage_metadata') and response.usage_metadata:
usage = response.usage_metadata
logger.info(
f"slow llm call: model={self.provider}/{self.model}, "
f"input_tokens={usage.prompt_token_count}, output_tokens={usage.candidates_token_count}, "
f"time={duration:.3f}s"
# Use structured output parsing and return .parsed
response = await self._client.beta.chat.completions.parse(
response_format=response_format,
**call_params
)
result = response.choices[0].message.parsed
else:
# Standard completion and return text content
response = await self._client.chat.completions.create(**call_params)
result = response.choices[0].message.content
# Log call details on success
duration = time.time() - start_time
usage = response.usage
logger.info(
f"model={self.provider}/{self.model}, "
f"input_tokens={usage.prompt_tokens}, output_tokens={usage.completion_tokens}, "
f"total_tokens={usage.total_tokens}, time={duration:.3f}s"
)
return result
except json.JSONDecodeError as e:
except LengthFinishReasonError as e:
# Output exceeded token limits - raise bridge exception for caller to handle
logger.warning(f"LLM output exceeded token limits: {str(e)}")
raise OutputTooLongError(
f"LLM output exceeded token limits. Input may need to be split into smaller chunks."
) from e
except APIStatusError as e:
last_exception = e
if attempt < max_retries:
logger.warning(f"Gemini returned invalid JSON, retrying...")
# Calculate exponential backoff with jitter
backoff = min(initial_backoff * (2 ** attempt), max_backoff)
await asyncio.sleep(backoff)
continue
else:
logger.error(f"Gemini returned invalid JSON after {max_retries + 1} attempts")
raise
# Add jitter (±20%)
jitter = backoff * 0.2 * (2 * (time.time() % 1) - 1)
sleep_time = backoff + jitter
except genai_errors.APIError as e:
# Fast fail only on 401 (unauthorized) and 403 (forbidden) - these won't recover with retries
if e.code in (401, 403):
logger.error(f"Gemini auth error (HTTP {e.code}), not retrying: {str(e)}")
raise
# Retry on retryable errors (rate limits, server errors, and other client errors like 400)
if e.code in (400, 429, 500, 502, 503, 504) or (e.code and e.code >= 500):
last_exception = e
if attempt < max_retries:
backoff = min(initial_backoff * (2 ** attempt), max_backoff)
jitter = backoff * 0.2 * (2 * (time.time() % 1) - 1)
await asyncio.sleep(backoff + jitter)
else:
logger.error(f"Gemini API error after {max_retries + 1} attempts: {str(e)}")
raise
logger.warning(
f"LLM error on attempt {attempt + 1}/{max_retries + 1}. "
f"Retrying in {sleep_time:.2f}s... Error: {str(e)}"
)
await asyncio.sleep(sleep_time)
else:
logger.error(f"Gemini API error: {type(e).__name__}: {str(e)}")
logger.error(f"Non-retryable API error after {max_retries + 1} attempts: {str(e)}, input {messages}")
raise
except Exception as e:
logger.error(f"Unexpected error during Gemini call: {type(e).__name__}: {str(e)}")
logger.error(f"Unexpected error during LLM call: {type(e).__name__}: {str(e)}, input {messages}")
raise
# This should never be reached, but just in case
if last_exception:
raise last_exception
raise RuntimeError(f"Gemini call failed after all retries")
raise RuntimeError(f"LLM call failed after all retries with no exception captured, input {messages}")
@classmethod
def for_memory(cls) -> "LLMProvider":
"""Create provider for memory operations from environment variables."""
def for_memory(cls) -> "LLMConfig":
"""Create configuration for memory operations from environment variables."""
provider = os.getenv("HINDSIGHT_API_LLM_PROVIDER", "groq")
api_key = os.getenv("HINDSIGHT_API_LLM_API_KEY")
base_url = os.getenv("HINDSIGHT_API_LLM_BASE_URL", "")
base_url = os.getenv("HINDSIGHT_API_LLM_BASE_URL")
model = os.getenv("HINDSIGHT_API_LLM_MODEL", "openai/gpt-oss-120b")
return cls(
provider=provider,
api_key=api_key,
base_url=base_url,
model=model,
reasoning_effort="low"
)
@classmethod
def for_answer_generation(cls) -> "LLMProvider":
"""Create provider for answer generation. Falls back to memory config if not set."""
provider = os.getenv("HINDSIGHT_API_ANSWER_LLM_PROVIDER", os.getenv("HINDSIGHT_API_LLM_PROVIDER", "groq"))
api_key = os.getenv("HINDSIGHT_API_ANSWER_LLM_API_KEY", os.getenv("HINDSIGHT_API_LLM_API_KEY"))
base_url = os.getenv("HINDSIGHT_API_ANSWER_LLM_BASE_URL", os.getenv("HINDSIGHT_API_LLM_BASE_URL", ""))
model = os.getenv("HINDSIGHT_API_ANSWER_LLM_MODEL", os.getenv("HINDSIGHT_API_LLM_MODEL", "openai/gpt-oss-120b"))
# Set default base URL if not provided
if not base_url:
if provider == "groq":
base_url = "https://api.groq.com/openai/v1"
elif provider == "ollama":
base_url = "http://localhost:11434/v1"
else:
base_url = ""
return cls(
provider=provider,
api_key=api_key,
base_url=base_url,
model=model,
reasoning_effort="high"
)
@classmethod
def for_judge(cls) -> "LLMProvider":
"""Create provider for judge/evaluator operations. Falls back to memory config if not set."""
def for_judge(cls) -> "LLMConfig":
"""
Create configuration for judge/evaluator operations from environment variables.
Falls back to memory LLM config if judge-specific config not set.
"""
# Check if judge-specific config exists, otherwise fall back to memory config
provider = os.getenv("HINDSIGHT_API_JUDGE_LLM_PROVIDER", os.getenv("HINDSIGHT_API_LLM_PROVIDER", "groq"))
api_key = os.getenv("HINDSIGHT_API_JUDGE_LLM_API_KEY", os.getenv("HINDSIGHT_API_LLM_API_KEY"))
base_url = os.getenv("HINDSIGHT_API_JUDGE_LLM_BASE_URL", os.getenv("HINDSIGHT_API_LLM_BASE_URL", ""))
base_url = os.getenv("HINDSIGHT_API_JUDGE_LLM_BASE_URL", os.getenv("HINDSIGHT_API_LLM_BASE_URL"))
model = os.getenv("HINDSIGHT_API_JUDGE_LLM_MODEL", os.getenv("HINDSIGHT_API_LLM_MODEL", "openai/gpt-oss-120b"))
# Set default base URL if not provided
if not base_url:
if provider == "groq":
base_url = "https://api.groq.com/openai/v1"
elif provider == "ollama":
base_url = "http://localhost:11434/v1"
else:
base_url = ""
return cls(
provider=provider,
api_key=api_key,
base_url=base_url,
model=model,
reasoning_effort="high"
)
# Backwards compatibility alias
LLMConfig = LLMProvider
File diff suppressed because it is too large Load Diff
@@ -6,9 +6,8 @@ structured information like temporal constraints.
"""
from abc import ABC, abstractmethod
from typing import Optional
from datetime import datetime, timedelta
from datetime import datetime
import logging
import re
from pydantic import BaseModel, Field
logger = logging.getLogger(__name__)
@@ -47,16 +46,6 @@ class QueryAnalyzer(ABC):
information like temporal constraints, entities, etc.
"""
@abstractmethod
def load(self) -> None:
"""
Load the query analyzer model.
This should be called during initialization to load the model
and avoid cold start latency on first analyze() call.
"""
pass
@abstractmethod
def analyze(
self, query: str, reference_date: Optional[datetime] = None
@@ -74,199 +63,6 @@ class QueryAnalyzer(ABC):
pass
class DateparserQueryAnalyzer(QueryAnalyzer):
"""
Query analyzer using dateparser library.
Uses dateparser to extract temporal expressions from natural language
queries. Supports 200+ languages including English, Spanish, Italian,
French, German, etc.
Performance:
- ~10-50ms per query
- No model loading required
"""
def __init__(self):
"""Initialize dateparser query analyzer."""
self._search_dates = None
def load(self) -> None:
"""Load dateparser (lazy import)."""
if self._search_dates is None:
from dateparser.search import search_dates
self._search_dates = search_dates
def analyze(
self, query: str, reference_date: Optional[datetime] = None
) -> QueryAnalysis:
"""
Analyze query using dateparser.
Extracts temporal expressions from the query text. Supports multiple
languages automatically.
Args:
query: Natural language query (any language)
reference_date: Reference date for relative terms (defaults to now)
Returns:
QueryAnalysis with temporal_constraint if found
"""
self.load()
if reference_date is None:
reference_date = datetime.now()
# Check for period expressions first (these need special handling)
query_lower = query.lower()
period_result = self._extract_period(query_lower, reference_date)
if period_result is not None:
return QueryAnalysis(temporal_constraint=period_result)
# Use dateparser's search_dates to find temporal expressions
settings = {
'RELATIVE_BASE': reference_date,
'PREFER_DATES_FROM': 'past',
'RETURN_AS_TIMEZONE_AWARE': False,
}
results = self._search_dates(query, settings=settings)
if not results:
return QueryAnalysis(temporal_constraint=None)
# Filter out false positives (common words parsed as dates)
false_positives = {'do', 'may', 'march', 'will', 'can', 'sat', 'sun', 'mon', 'tue', 'wed', 'thu', 'fri'}
valid_results = [
(text, date) for text, date in results
if text.lower() not in false_positives or len(text) > 3
]
if not valid_results:
return QueryAnalysis(temporal_constraint=None)
# Use the first valid date found
_, parsed_date = valid_results[0]
# Create constraint for single day
start_date = parsed_date.replace(hour=0, minute=0, second=0, microsecond=0)
end_date = parsed_date.replace(hour=23, minute=59, second=59, microsecond=999999)
return QueryAnalysis(
temporal_constraint=TemporalConstraint(
start_date=start_date,
end_date=end_date
)
)
def _extract_period(
self, query: str, reference_date: datetime
) -> Optional[TemporalConstraint]:
"""
Extract period-based temporal expressions (week, month, year, weekend).
These need special handling as they represent date ranges, not single dates.
Supports multiple languages.
"""
def constraint(start: datetime, end: datetime) -> TemporalConstraint:
return TemporalConstraint(
start_date=start.replace(hour=0, minute=0, second=0, microsecond=0),
end_date=end.replace(hour=23, minute=59, second=59, microsecond=999999)
)
# Yesterday patterns (English, Spanish, Italian, French, German)
if re.search(r'\b(yesterday|ayer|ieri|hier|gestern)\b', query, re.IGNORECASE):
d = reference_date - timedelta(days=1)
return constraint(d, d)
# Today patterns
if re.search(r'\b(today|hoy|oggi|aujourd\'?hui|heute)\b', query, re.IGNORECASE):
return constraint(reference_date, reference_date)
# "a couple of days ago" / "a few days ago" patterns
# These are imprecise so we create a range
if re.search(r'\b(a\s+)?couple\s+(of\s+)?days?\s+ago\b', query, re.IGNORECASE):
# "a couple of days" = approximately 2 days, give range of 1-3 days
return constraint(reference_date - timedelta(days=3), reference_date - timedelta(days=1))
if re.search(r'\b(a\s+)?few\s+days?\s+ago\b', query, re.IGNORECASE):
# "a few days" = approximately 3-4 days, give range of 2-5 days
return constraint(reference_date - timedelta(days=5), reference_date - timedelta(days=2))
# "a couple of weeks ago" / "a few weeks ago" patterns
if re.search(r'\b(a\s+)?couple\s+(of\s+)?weeks?\s+ago\b', query, re.IGNORECASE):
# "a couple of weeks" = approximately 2 weeks, give range of 1-3 weeks
return constraint(reference_date - timedelta(weeks=3), reference_date - timedelta(weeks=1))
if re.search(r'\b(a\s+)?few\s+weeks?\s+ago\b', query, re.IGNORECASE):
# "a few weeks" = approximately 3-4 weeks, give range of 2-5 weeks
return constraint(reference_date - timedelta(weeks=5), reference_date - timedelta(weeks=2))
# "a couple of months ago" / "a few months ago" patterns
if re.search(r'\b(a\s+)?couple\s+(of\s+)?months?\s+ago\b', query, re.IGNORECASE):
# "a couple of months" = approximately 2 months, give range of 1-3 months
return constraint(reference_date - timedelta(days=90), reference_date - timedelta(days=30))
if re.search(r'\b(a\s+)?few\s+months?\s+ago\b', query, re.IGNORECASE):
# "a few months" = approximately 3-4 months, give range of 2-5 months
return constraint(reference_date - timedelta(days=150), reference_date - timedelta(days=60))
# Last week patterns (English, Spanish, Italian, French, German)
if re.search(r'\b(last\s+week|la\s+semana\s+pasada|la\s+settimana\s+scorsa|la\s+semaine\s+derni[eè]re|letzte\s+woche)\b', query, re.IGNORECASE):
start = reference_date - timedelta(days=reference_date.weekday() + 7)
return constraint(start, start + timedelta(days=6))
# Last month patterns
if re.search(r'\b(last\s+month|el\s+mes\s+pasado|il\s+mese\s+scorso|le\s+mois\s+dernier|letzten?\s+monat)\b', query, re.IGNORECASE):
first = reference_date.replace(day=1)
end = first - timedelta(days=1)
start = end.replace(day=1)
return constraint(start, end)
# Last year patterns
if re.search(r'\b(last\s+year|el\s+a[ñn]o\s+pasado|l\'anno\s+scorso|l\'ann[ée]e\s+derni[eè]re|letztes?\s+jahr)\b', query, re.IGNORECASE):
year = reference_date.year - 1
return constraint(datetime(year, 1, 1), datetime(year, 12, 31))
# Last weekend patterns
if re.search(r'\b(last\s+weekend|el\s+fin\s+de\s+semana\s+pasado|lo\s+scorso\s+fine\s+settimana|le\s+week-?end\s+dernier|letztes?\s+wochenende)\b', query, re.IGNORECASE):
days_since_sat = (reference_date.weekday() + 2) % 7
if days_since_sat == 0:
days_since_sat = 7
sat = reference_date - timedelta(days=days_since_sat)
return constraint(sat, sat + timedelta(days=1))
# Month + Year patterns (e.g., "June 2024", "junio 2024", "giugno 2024")
month_patterns = {
'january|enero|gennaio|janvier|januar': 1,
'february|febrero|febbraio|f[ée]vrier|februar': 2,
'march|marzo|mars|m[äa]rz': 3,
'april|abril|aprile|avril': 4,
'may|mayo|maggio|mai': 5,
'june|junio|giugno|juin|juni': 6,
'july|julio|luglio|juillet|juli': 7,
'august|agosto|ao[uû]t': 8,
'september|septiembre|settembre|septembre': 9,
'october|octubre|ottobre|octobre|oktober': 10,
'november|noviembre|novembre': 11,
'december|diciembre|dicembre|d[ée]cembre|dezember': 12,
}
for pattern, month_num in month_patterns.items():
match = re.search(rf'\b({pattern})\s+(\d{{4}})\b', query, re.IGNORECASE)
if match:
year = int(match.group(2))
start = datetime(year, month_num, 1)
if month_num == 12:
end = datetime(year, 12, 31)
else:
end = datetime(year, month_num + 1, 1) - timedelta(days=1)
return constraint(start, end)
return None
class TransformerQueryAnalyzer(QueryAnalyzer):
"""
Query analyzer using T5-based generative models.
@@ -298,113 +94,31 @@ class TransformerQueryAnalyzer(QueryAnalyzer):
self._model = None
self._tokenizer = None
def load(self) -> None:
"""Load the T5 model for temporal extraction."""
if self._model is not None:
return
try:
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
except ImportError:
raise ImportError(
"transformers is required for TransformerQueryAnalyzer. "
"Install it with: pip install transformers"
)
logger.info(f"Loading query analyzer model: {self.model_name}...")
self._tokenizer = AutoTokenizer.from_pretrained(self.model_name)
self._model = AutoModelForSeq2SeqLM.from_pretrained(self.model_name)
self._model.to(self.device)
self._model.eval()
logger.info("Query analyzer model loaded")
def _load_model(self):
"""Lazy load the T5 model for temporal extraction (calls load())."""
self.load()
"""Lazy load the T5 model for temporal extraction."""
if self._model is None:
try:
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
except ImportError:
raise ImportError(
"transformers is required for TransformerQueryAnalyzer. "
"Install it with: pip install transformers"
)
def _extract_with_rules(
self, query: str, reference_date: datetime
) -> Optional[TemporalConstraint]:
"""
Extract temporal expressions using rule-based patterns.
Handles common patterns reliably and fast. Returns None for
patterns that need model-based extraction.
"""
import re
query_lower = query.lower()
def get_last_weekday(weekday: int) -> datetime:
days_ago = (reference_date.weekday() - weekday) % 7
if days_ago == 0:
days_ago = 7
return reference_date - timedelta(days=days_ago)
def constraint(start: datetime, end: datetime) -> TemporalConstraint:
return TemporalConstraint(
start_date=start.replace(hour=0, minute=0, second=0, microsecond=0),
end_date=end.replace(hour=23, minute=59, second=59, microsecond=999999)
)
# Yesterday
if re.search(r'\byesterday\b', query_lower):
d = reference_date - timedelta(days=1)
return constraint(d, d)
# Last week
if re.search(r'\blast\s+week\b', query_lower):
start = reference_date - timedelta(days=reference_date.weekday() + 7)
return constraint(start, start + timedelta(days=6))
# Last month
if re.search(r'\blast\s+month\b', query_lower):
first = reference_date.replace(day=1)
end = first - timedelta(days=1)
start = end.replace(day=1)
return constraint(start, end)
# Last year
if re.search(r'\blast\s+year\b', query_lower):
y = reference_date.year - 1
return constraint(datetime(y, 1, 1), datetime(y, 12, 31))
# Last weekend
if re.search(r'\blast\s+weekend\b', query_lower):
sat = get_last_weekday(5)
return constraint(sat, sat + timedelta(days=1))
# Last <weekday>
weekdays = {'monday': 0, 'tuesday': 1, 'wednesday': 2, 'thursday': 3,
'friday': 4, 'saturday': 5, 'sunday': 6}
for name, num in weekdays.items():
if re.search(rf'\blast\s+{name}\b', query_lower):
d = get_last_weekday(num)
return constraint(d, d)
# Month + Year: "June 2024", "in March 2023"
months = {'january': 1, 'february': 2, 'march': 3, 'april': 4, 'may': 5,
'june': 6, 'july': 7, 'august': 8, 'september': 9, 'october': 10,
'november': 11, 'december': 12}
for name, num in months.items():
match = re.search(rf'\b{name}\s+(\d{{4}})\b', query_lower)
if match:
year = int(match.group(1))
if num == 12:
last_day = 31
else:
last_day = (datetime(year, num + 1, 1) - timedelta(days=1)).day
return constraint(datetime(year, num, 1), datetime(year, num, last_day))
return None
logger.debug(f"Loading T5 model: {self.model_name}...")
self._tokenizer = AutoTokenizer.from_pretrained(self.model_name)
self._model = AutoModelForSeq2SeqLM.from_pretrained(self.model_name)
self._model.to(self.device)
self._model.eval()
logger.debug(f"Model loaded on {self.device}")
def analyze(
self, query: str, reference_date: Optional[datetime] = None
) -> QueryAnalysis:
"""
Analyze query for temporal expressions.
Analyze query using T5 model.
Uses rule-based extraction for common patterns (fast & reliable),
falls back to T5 model for complex/unusual patterns.
Uses T5 to generate structured temporal output directly.
Args:
query: Natural language query
@@ -416,30 +130,17 @@ class TransformerQueryAnalyzer(QueryAnalyzer):
if reference_date is None:
reference_date = datetime.now()
# Try rule-based extraction first (handles 90%+ of cases)
result = self._extract_with_rules(query, reference_date)
if result is not None:
return QueryAnalysis(temporal_constraint=result)
# Fall back to T5 model for unusual patterns
self._load_model()
# Helper to calculate example dates
def get_last_weekday(weekday: int) -> datetime:
days_ago = (reference_date.weekday() - weekday) % 7
if days_ago == 0:
days_ago = 7
return reference_date - timedelta(days=days_ago)
yesterday = reference_date - timedelta(days=1)
last_saturday = get_last_weekday(5)
# Build prompt for T5
prompt = f"""Today is {reference_date.strftime('%Y-%m-%d')}. Extract date range or "none".
# Build prompt for T5 to generate structured temporal output
# Use fill-in-the-blank format which T5 handles better
prompt = f"""Today is {reference_date.strftime('%Y-%m-%d')}. Convert temporal expressions to date ranges.
June 2024 = 2024-06-01 to 2024-06-30
yesterday = {yesterday.strftime('%Y-%m-%d')} to {yesterday.strftime('%Y-%m-%d')}
last Saturday = {last_saturday.strftime('%Y-%m-%d')} to {last_saturday.strftime('%Y-%m-%d')}
March 2023 = 2023-03-01 to 2023-03-31
dogs in June 2023 = 2023-06-01 to 2023-06-30
last year = {reference_date.year - 1}-01-01 to {reference_date.year - 1}-12-31
events in January 2020 = 2020-01-01 to 2020-01-31
what is the weather = none
{query} ="""
@@ -457,6 +158,7 @@ what is the weather = none
)
result = self._tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
logger.debug(f"T5 generated: '{result}'")
# Parse the generated output
temporal = self._parse_generated_output(result, reference_date)
@@ -514,6 +216,7 @@ what is the weather = none
return TemporalConstraint(start_date=start_date, end_date=end_date)
except (ValueError, AttributeError) as e:
logger.debug(f"Failed to parse T5 output '{result}': {e}")
return None
return None
@@ -10,32 +10,6 @@ from typing import Optional, List, Dict, Any
from pydantic import BaseModel, Field, ConfigDict
# Valid fact types for recall operations (excludes 'observation' which is internal)
VALID_RECALL_FACT_TYPES = frozenset(["world", "experience", "opinion"])
class DispositionTraits(BaseModel):
"""
Disposition traits for a memory bank.
All traits are scored 1-5 where:
- skepticism: 1=trusting, 5=skeptical (how much to doubt or question information)
- literalism: 1=flexible interpretation, 5=literal interpretation (how strictly to interpret information)
- empathy: 1=detached, 5=empathetic (how much to consider emotional context)
"""
skepticism: int = Field(ge=1, le=5, description="How skeptical vs trusting (1=trusting, 5=skeptical)")
literalism: int = Field(ge=1, le=5, description="How literally to interpret information (1=flexible, 5=literal)")
empathy: int = Field(ge=1, le=5, description="How much to consider emotional context (1=detached, 5=empathetic)")
model_config = ConfigDict(json_schema_extra={
"example": {
"skepticism": 3,
"literalism": 3,
"empathy": 3
}
})
class MemoryFact(BaseModel):
"""
A single memory fact returned by search or think operations.
@@ -48,41 +22,32 @@ class MemoryFact(BaseModel):
"id": "123e4567-e89b-12d3-a456-426614174000",
"text": "Alice works at Google on the AI team",
"fact_type": "world",
"entities": ["Alice", "Google"],
"context": "work info",
"occurred_start": "2024-01-15T10:30:00Z",
"occurred_end": "2024-01-15T10:30:00Z",
"mentioned_at": "2024-01-15T10:30:00Z",
"event_date": "2024-01-15T10:30:00Z",
"document_id": "session_abc123",
"metadata": {"source": "slack"},
"chunk_id": "bank123_session_abc123_0",
"activation": 0.95
}
})
id: str = Field(description="Unique identifier for the memory fact")
text: str = Field(description="The actual text content of the memory")
fact_type: str = Field(description="Type of fact: 'world', 'experience', 'opinion', or 'observation'")
entities: Optional[List[str]] = Field(None, description="Entity names mentioned in this fact")
fact_type: str = Field(description="Type of fact: 'world', 'agent', or 'opinion'")
context: Optional[str] = Field(None, description="Additional context for the memory")
event_date: Optional[str] = Field(None, description="ISO format date when the event occurred")
occurred_start: Optional[str] = Field(None, description="ISO format date when the event started occurring")
occurred_end: Optional[str] = Field(None, description="ISO format date when the event ended occurring")
mentioned_at: Optional[str] = Field(None, description="ISO format date when the fact was mentioned/learned")
document_id: Optional[str] = Field(None, description="ID of the document this memory belongs to")
metadata: Optional[Dict[str, str]] = Field(None, description="User-defined metadata")
chunk_id: Optional[str] = Field(None, description="ID of the chunk this fact was extracted from (format: bank_id_document_id_chunk_index)")
# Internal metrics (used by system but may not be exposed in API)
activation: Optional[float] = Field(None, description="Internal activation score")
class ChunkInfo(BaseModel):
"""Information about a chunk."""
chunk_text: str = Field(description="The raw chunk text")
chunk_index: int = Field(description="Index of the chunk within the document")
truncated: bool = Field(default=False, description="Whether the chunk was truncated due to token limits")
class RecallResult(BaseModel):
class SearchResult(BaseModel):
"""
Result from a recall operation.
Result from a search operation.
Contains a list of matching memory facts and optional trace information
for debugging and transparency.
@@ -95,8 +60,7 @@ class RecallResult(BaseModel):
"text": "Alice works at Google on the AI team",
"fact_type": "world",
"context": "work info",
"occurred_start": "2024-01-15T10:30:00Z",
"occurred_end": "2024-01-15T10:30:00Z",
"event_date": "2024-01-15T10:30:00Z",
"activation": 0.95
}
],
@@ -109,22 +73,14 @@ class RecallResult(BaseModel):
results: List[MemoryFact] = Field(description="List of memory facts matching the query")
trace: Optional[Dict[str, Any]] = Field(None, description="Trace information for debugging")
entities: Optional[Dict[str, "EntityState"]] = Field(
None,
description="Entity states for entities mentioned in results (keyed by canonical name)"
)
chunks: Optional[Dict[str, ChunkInfo]] = Field(
None,
description="Chunks for facts, keyed by '{document_id}_{chunk_index}'"
)
class ReflectResult(BaseModel):
class ThinkResult(BaseModel):
"""
Result from a reflect operation.
Result from a think operation.
Contains the formulated answer, the facts it was based on (organized by type),
and any new opinions that were formed during the reflection process.
and any new opinions that were formed during the thinking process.
"""
model_config = ConfigDict(json_schema_extra={
"example": {
@@ -136,11 +92,10 @@ class ReflectResult(BaseModel):
"text": "Machine learning is used in medical diagnosis",
"fact_type": "world",
"context": "healthcare",
"occurred_start": "2024-01-15T10:30:00Z",
"occurred_end": "2024-01-15T10:30:00Z"
"event_date": "2024-01-15T10:30:00Z"
}
],
"experience": [],
"agent": [],
"opinion": []
},
"new_opinions": [
@@ -151,11 +106,11 @@ class ReflectResult(BaseModel):
text: str = Field(description="The formulated answer text")
based_on: Dict[str, List[MemoryFact]] = Field(
description="Facts used to formulate the answer, organized by type (world, experience, opinion)"
description="Facts used to formulate the answer, organized by type (world, agent, opinion)"
)
new_opinions: List[str] = Field(
default_factory=list,
description="List of newly formed opinions during reflection"
description="List of newly formed opinions during thinking"
)
@@ -163,7 +118,7 @@ class Opinion(BaseModel):
"""
An opinion with confidence score.
Opinions represent the bank's formed perspectives on topics,
Opinions represent the agent's formed perspectives on topics,
with a confidence level indicating strength of belief.
"""
model_config = ConfigDict(json_schema_extra={
@@ -175,46 +130,3 @@ class Opinion(BaseModel):
text: str = Field(description="The opinion text")
confidence: float = Field(description="Confidence score between 0.0 and 1.0")
class EntityObservation(BaseModel):
"""
An observation about an entity.
Observations are objective facts synthesized from multiple memory facts
about an entity, without personality influence.
"""
model_config = ConfigDict(json_schema_extra={
"example": {
"text": "John is detail-oriented and works at Google",
"mentioned_at": "2024-01-15T10:30:00Z"
}
})
text: str = Field(description="The observation text")
mentioned_at: Optional[str] = Field(None, description="ISO format date when this observation was created")
class EntityState(BaseModel):
"""
Current mental model of an entity.
Contains observations synthesized from facts about the entity.
"""
model_config = ConfigDict(json_schema_extra={
"example": {
"entity_id": "123e4567-e89b-12d3-a456-426614174000",
"canonical_name": "John",
"observations": [
{"text": "John is detail-oriented", "mentioned_at": "2024-01-15T10:30:00Z"},
{"text": "John works at Google on the AI team", "mentioned_at": "2024-01-14T09:00:00Z"}
]
}
})
entity_id: str = Field(description="Unique identifier for the entity")
canonical_name: str = Field(description="Canonical name of the entity")
observations: List[EntityObservation] = Field(
default_factory=list,
description="List of observations about this entity"
)
@@ -1,50 +0,0 @@
"""
Retain pipeline modules for storing memories.
This package contains modular components for the retain operation:
- types: Type definitions for retain pipeline
- fact_extraction: Extract facts from content
- embedding_processing: Augment texts and generate embeddings
- deduplication: Check for duplicate facts
- entity_processing: Process and resolve entities
- link_creation: Create temporal, semantic, entity, and causal links
- chunk_storage: Handle chunk storage
- fact_storage: Handle fact insertion into database
"""
from .types import (
RetainContent,
ExtractedFact,
ProcessedFact,
ChunkMetadata,
EntityRef,
CausalRelation,
RetainBatch
)
from . import fact_extraction
from . import embedding_processing
from . import deduplication
from . import entity_processing
from . import link_creation
from . import chunk_storage
from . import fact_storage
__all__ = [
# Types
"RetainContent",
"ExtractedFact",
"ProcessedFact",
"ChunkMetadata",
"EntityRef",
"CausalRelation",
"RetainBatch",
# Modules
"fact_extraction",
"embedding_processing",
"deduplication",
"entity_processing",
"link_creation",
"chunk_storage",
"fact_storage",
]
@@ -1,82 +0,0 @@
"""
Chunk storage for retain pipeline.
Handles storage of document chunks in the database.
"""
import logging
from typing import List, Dict, Optional
from .types import ChunkMetadata
logger = logging.getLogger(__name__)
async def store_chunks_batch(
conn,
bank_id: str,
document_id: str,
chunks: List[ChunkMetadata]
) -> Dict[int, str]:
"""
Store document chunks in the database.
Args:
conn: Database connection
bank_id: Bank identifier
document_id: Document identifier
chunks: List of ChunkMetadata objects
Returns:
Dictionary mapping global chunk index to chunk_id
"""
if not chunks:
return {}
# Prepare chunk data for batch insert
chunk_ids = []
chunk_texts = []
chunk_indices = []
chunk_id_map = {}
for chunk in chunks:
chunk_id = f"{bank_id}_{document_id}_{chunk.chunk_index}"
chunk_ids.append(chunk_id)
chunk_texts.append(chunk.chunk_text)
chunk_indices.append(chunk.chunk_index)
chunk_id_map[chunk.chunk_index] = chunk_id
# Batch insert all chunks
await conn.execute(
"""
INSERT INTO chunks (chunk_id, document_id, bank_id, chunk_text, chunk_index)
SELECT * FROM unnest($1::text[], $2::text[], $3::text[], $4::text[], $5::integer[])
""",
chunk_ids,
[document_id] * len(chunk_texts),
[bank_id] * len(chunk_texts),
chunk_texts,
chunk_indices
)
return chunk_id_map
def map_facts_to_chunks(
facts_chunk_indices: List[int],
chunk_id_map: Dict[int, str]
) -> List[Optional[str]]:
"""
Map fact chunk indices to chunk IDs.
Args:
facts_chunk_indices: List of chunk indices for each fact
chunk_id_map: Dictionary mapping chunk index to chunk_id
Returns:
List of chunk_ids (same length as facts_chunk_indices)
"""
chunk_ids = []
for chunk_idx in facts_chunk_indices:
chunk_id = chunk_id_map.get(chunk_idx)
chunk_ids.append(chunk_id)
return chunk_ids
@@ -1,104 +0,0 @@
"""
Deduplication logic for retain pipeline.
Checks for duplicate facts using semantic similarity and temporal proximity.
"""
import logging
from datetime import datetime
from typing import List
from collections import defaultdict
from .types import ProcessedFact
logger = logging.getLogger(__name__)
async def check_duplicates_batch(
conn,
bank_id: str,
facts: List[ProcessedFact],
duplicate_checker_fn
) -> List[bool]:
"""
Check which facts are duplicates using batched time-window queries.
Groups facts by 12-hour time buckets to efficiently check for duplicates
within a 24-hour window.
Args:
conn: Database connection
bank_id: Bank identifier
facts: List of ProcessedFact objects to check
duplicate_checker_fn: Async function(conn, bank_id, texts, embeddings, date, time_window_hours)
that returns List[bool] indicating duplicates
Returns:
List of boolean flags (same length as facts) indicating if each fact is a duplicate
"""
if not facts:
return []
# Group facts by event_date (rounded to 12-hour buckets) for efficient batching
time_buckets = defaultdict(list)
for idx, fact in enumerate(facts):
# Use occurred_start if available, otherwise use mentioned_at
# For deduplication purposes, we need a time reference
fact_date = fact.occurred_start if fact.occurred_start is not None else fact.mentioned_at
# Defensive: if both are None (shouldn't happen), use now()
if fact_date is None:
from datetime import datetime, timezone
fact_date = datetime.now(timezone.utc)
# Round to 12-hour bucket to group similar times
bucket_key = fact_date.replace(
hour=(fact_date.hour // 12) * 12,
minute=0,
second=0,
microsecond=0
)
time_buckets[bucket_key].append((idx, fact))
# Process each bucket in batch
all_is_duplicate = [False] * len(facts)
for bucket_date, bucket_items in time_buckets.items():
indices = [item[0] for item in bucket_items]
texts = [item[1].fact_text for item in bucket_items]
embeddings = [item[1].embedding for item in bucket_items]
# Check duplicates for this time bucket
dup_flags = await duplicate_checker_fn(
conn,
bank_id,
texts,
embeddings,
bucket_date,
time_window_hours=24
)
# Map results back to original indices
for idx, is_dup in zip(indices, dup_flags):
all_is_duplicate[idx] = is_dup
return all_is_duplicate
def filter_duplicates(
facts: List[ProcessedFact],
is_duplicate_flags: List[bool]
) -> List[ProcessedFact]:
"""
Filter out duplicate facts based on duplicate flags.
Args:
facts: List of ProcessedFact objects
is_duplicate_flags: Boolean flags indicating which facts are duplicates
Returns:
List of non-duplicate facts
"""
if len(facts) != len(is_duplicate_flags):
raise ValueError(f"Mismatch between facts ({len(facts)}) and flags ({len(is_duplicate_flags)})")
return [fact for fact, is_dup in zip(facts, is_duplicate_flags) if not is_dup]
@@ -1,62 +0,0 @@
"""
Embedding processing for retain pipeline.
Handles augmenting fact texts with temporal information and generating embeddings.
"""
import logging
from typing import List
from datetime import datetime
from . import embedding_utils
from .types import ExtractedFact
logger = logging.getLogger(__name__)
def augment_texts_with_dates(facts: List[ExtractedFact], format_date_fn) -> List[str]:
"""
Augment fact texts with readable dates for better temporal matching.
This allows queries like "camping in June" to match facts that happened in June.
Args:
facts: List of ExtractedFact objects
format_date_fn: Function to format datetime to readable string
Returns:
List of augmented text strings (same length as facts)
"""
augmented_texts = []
for fact in facts:
# Use occurred_start as the representative date
fact_date = fact.occurred_start or fact.mentioned_at
readable_date = format_date_fn(fact_date)
# Augment text with date for embedding (but store original text in DB)
augmented_text = f"{fact.fact_text} (happened in {readable_date})"
augmented_texts.append(augmented_text)
return augmented_texts
async def generate_embeddings_batch(
embeddings_model,
texts: List[str]
) -> List[List[float]]:
"""
Generate embeddings for a batch of texts.
Args:
embeddings_model: Embeddings model instance
texts: List of text strings to embed
Returns:
List of embedding vectors (same length as texts)
"""
if not texts:
return []
embeddings = await embedding_utils.generate_embeddings_batch(
embeddings_model,
texts
)
return embeddings
@@ -1,90 +0,0 @@
"""
Entity processing for retain pipeline.
Handles entity extraction, resolution, and link creation for stored facts.
"""
import logging
from typing import List, Tuple, Dict, Any
from uuid import UUID
from .types import ProcessedFact, EntityRef, EntityLink
from . import link_utils
logger = logging.getLogger(__name__)
async def process_entities_batch(
entity_resolver,
conn,
bank_id: str,
unit_ids: List[str],
facts: List[ProcessedFact],
log_buffer: List[str] = None
) -> List[EntityLink]:
"""
Process entities for all facts and create entity links.
This function:
1. Extracts entity mentions from fact texts
2. Resolves entity names to canonical entities
3. Creates entity records in the database
4. Returns entity links ready for insertion
Args:
entity_resolver: EntityResolver instance for entity resolution
conn: Database connection
bank_id: Bank identifier
unit_ids: List of unit IDs (same length as facts)
facts: List of ProcessedFact objects
log_buffer: Optional buffer for detailed logging
Returns:
List of EntityLink objects for batch insertion
"""
if not unit_ids or not facts:
return []
if len(unit_ids) != len(facts):
raise ValueError(f"Mismatch between unit_ids ({len(unit_ids)}) and facts ({len(facts)})")
# Extract data for link_utils function
fact_texts = [fact.fact_text for fact in facts]
# Use occurred_start if available, otherwise use mentioned_at for entity timestamps
fact_dates = [fact.occurred_start if fact.occurred_start is not None else fact.mentioned_at for fact in facts]
# Convert EntityRef objects to dict format expected by link_utils
entities_per_fact = [
[{'text': entity.name, 'type': 'CONCEPT'} for entity in (fact.entities or [])]
for fact in facts
]
# Use existing link_utils function for entity processing
entity_links = await link_utils.extract_entities_batch_optimized(
entity_resolver,
conn,
bank_id,
unit_ids,
fact_texts,
"", # context (not used in current implementation)
fact_dates,
entities_per_fact,
log_buffer # Pass log_buffer for detailed logging
)
return entity_links
async def insert_entity_links_batch(
conn,
entity_links: List[EntityLink]
) -> None:
"""
Insert entity links in batch.
Args:
conn: Database connection
entity_links: List of EntityLink objects
"""
if not entity_links:
return
await link_utils.insert_entity_links_batch(conn, entity_links)
File diff suppressed because it is too large Load Diff
@@ -1,176 +0,0 @@
"""
Fact storage for retain pipeline.
Handles insertion of facts into the database.
"""
import logging
import json
from typing import List, Optional
from uuid import UUID
from .types import ProcessedFact
logger = logging.getLogger(__name__)
async def insert_facts_batch(
conn,
bank_id: str,
facts: List[ProcessedFact],
document_id: Optional[str] = None
) -> List[str]:
"""
Insert facts into the database in batch.
Args:
conn: Database connection
bank_id: Bank identifier
facts: List of ProcessedFact objects to insert
document_id: Optional document ID to associate with facts
Returns:
List of unit IDs (UUIDs as strings) for the inserted facts
"""
if not facts:
return []
# Prepare data for batch insert
fact_texts = []
embeddings = []
event_dates = []
occurred_starts = []
occurred_ends = []
mentioned_ats = []
contexts = []
fact_types = []
confidence_scores = []
access_counts = []
metadata_jsons = []
chunk_ids = []
document_ids = []
for fact in facts:
fact_texts.append(fact.fact_text)
# Convert embedding to string for asyncpg vector type
embeddings.append(str(fact.embedding))
# event_date: Use occurred_start if available, otherwise use mentioned_at
# This maintains backward compatibility while handling None occurred_start
event_dates.append(fact.occurred_start if fact.occurred_start is not None else fact.mentioned_at)
occurred_starts.append(fact.occurred_start)
occurred_ends.append(fact.occurred_end)
mentioned_ats.append(fact.mentioned_at)
contexts.append(fact.context)
fact_types.append(fact.fact_type)
# confidence_score is only for opinion facts
confidence_scores.append(1.0 if fact.fact_type == 'opinion' else None)
access_counts.append(0) # Initial access count
metadata_jsons.append(json.dumps(fact.metadata))
chunk_ids.append(fact.chunk_id)
# Use per-fact document_id if available, otherwise fallback to batch-level document_id
document_ids.append(fact.document_id if fact.document_id else document_id)
# Batch insert all facts
results = await conn.fetch(
"""
INSERT INTO memory_units (bank_id, text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
context, fact_type, confidence_score, access_count, metadata, chunk_id, document_id)
SELECT $1, * FROM unnest(
$2::text[], $3::vector[], $4::timestamptz[], $5::timestamptz[], $6::timestamptz[], $7::timestamptz[],
$8::text[], $9::text[], $10::float[], $11::int[], $12::jsonb[], $13::text[], $14::text[]
)
RETURNING id
""",
bank_id,
fact_texts,
embeddings,
event_dates, # event_date: occurred_start if available, else mentioned_at
occurred_starts,
occurred_ends,
mentioned_ats,
contexts,
fact_types,
confidence_scores,
access_counts,
metadata_jsons,
chunk_ids,
document_ids
)
unit_ids = [str(row['id']) for row in results]
return unit_ids
async def ensure_bank_exists(conn, bank_id: str) -> None:
"""
Ensure bank exists in the database.
Creates bank with default values if it doesn't exist.
Args:
conn: Database connection
bank_id: Bank identifier
"""
await conn.execute(
"""
INSERT INTO banks (bank_id, disposition, background)
VALUES ($1, $2::jsonb, $3)
ON CONFLICT (bank_id) DO UPDATE
SET updated_at = NOW()
""",
bank_id,
'{"skepticism": 3, "literalism": 3, "empathy": 3}',
""
)
async def handle_document_tracking(
conn,
bank_id: str,
document_id: str,
combined_content: str,
is_first_batch: bool,
retain_params: Optional[dict] = None
) -> None:
"""
Handle document tracking in the database.
Args:
conn: Database connection
bank_id: Bank identifier
document_id: Document identifier
combined_content: Combined content text from all content items
is_first_batch: Whether this is the first batch (for chunked operations)
retain_params: Optional parameters passed during retain (context, event_date, etc.)
"""
import hashlib
# Calculate content hash
content_hash = hashlib.sha256(combined_content.encode()).hexdigest()
# Always delete old document first if it exists (cascades to units and links)
# Only delete on the first batch to avoid deleting data we just inserted
if is_first_batch:
await conn.fetchval(
"DELETE FROM documents WHERE id = $1 AND bank_id = $2 RETURNING id",
document_id, bank_id
)
# Insert document (or update if exists from concurrent operations)
await conn.execute(
"""
INSERT INTO documents (id, bank_id, original_text, content_hash, metadata, retain_params)
VALUES ($1, $2, $3, $4, $5, $6)
ON CONFLICT (id, bank_id) DO UPDATE
SET original_text = EXCLUDED.original_text,
content_hash = EXCLUDED.content_hash,
metadata = EXCLUDED.metadata,
retain_params = EXCLUDED.retain_params,
updated_at = NOW()
""",
document_id,
bank_id,
combined_content,
content_hash,
json.dumps({}), # Empty metadata dict
json.dumps(retain_params) if retain_params else None
)
@@ -1,127 +0,0 @@
"""
Link creation for retain pipeline.
Handles creation of temporal, semantic, and causal links between facts.
"""
import logging
from typing import List
from .types import ProcessedFact, CausalRelation
from . import link_utils
logger = logging.getLogger(__name__)
async def create_temporal_links_batch(
conn,
bank_id: str,
unit_ids: List[str]
) -> int:
"""
Create temporal links between facts.
Links facts that occurred close in time to each other.
Args:
conn: Database connection
bank_id: Bank identifier
unit_ids: List of unit IDs to create links for
Returns:
Number of temporal links created
"""
if not unit_ids:
return 0
return await link_utils.create_temporal_links_batch_per_fact(
conn,
bank_id,
unit_ids,
log_buffer=[]
)
async def create_semantic_links_batch(
conn,
bank_id: str,
unit_ids: List[str],
embeddings: List[List[float]]
) -> int:
"""
Create semantic links between facts.
Links facts that are semantically similar based on embeddings.
Args:
conn: Database connection
bank_id: Bank identifier
unit_ids: List of unit IDs to create links for
embeddings: List of embedding vectors (same length as unit_ids)
Returns:
Number of semantic links created
"""
if not unit_ids or not embeddings:
return 0
if len(unit_ids) != len(embeddings):
raise ValueError(f"Mismatch between unit_ids ({len(unit_ids)}) and embeddings ({len(embeddings)})")
return await link_utils.create_semantic_links_batch(
conn,
bank_id,
unit_ids,
embeddings,
log_buffer=[]
)
async def create_causal_links_batch(
conn,
unit_ids: List[str],
facts: List[ProcessedFact]
) -> int:
"""
Create causal links between facts.
Links facts that have causal relationships (causes, enables, prevents).
Args:
conn: Database connection
unit_ids: List of unit IDs (same length as facts)
facts: List of ProcessedFact objects with causal_relations
Returns:
Number of causal links created
"""
if not unit_ids or not facts:
return 0
if len(unit_ids) != len(facts):
raise ValueError(f"Mismatch between unit_ids ({len(unit_ids)}) and facts ({len(facts)})")
# Extract causal relations in the format expected by link_utils
# Format: List of lists, where each inner list is the causal relations for that fact
causal_relations_per_fact = []
for fact in facts:
if fact.causal_relations:
# Convert CausalRelation objects to dicts
relations_dicts = [
{
'relation_type': rel.relation_type,
'target_fact_index': rel.target_fact_index,
'strength': rel.strength
}
for rel in fact.causal_relations
]
causal_relations_per_fact.append(relations_dicts)
else:
causal_relations_per_fact.append([])
link_count = await link_utils.create_causal_links_batch(
conn,
unit_ids,
causal_relations_per_fact
)
return link_count
@@ -1,264 +0,0 @@
"""
Observation regeneration for retain pipeline.
Regenerates entity observations as part of the retain transaction.
"""
import logging
import time
import uuid
from datetime import datetime, timezone
from typing import List, Dict, Optional
from ..search import observation_utils
from . import embedding_utils
from ..db_utils import acquire_with_retry
from .types import EntityLink
logger = logging.getLogger(__name__)
def utcnow():
"""Get current UTC time."""
return datetime.now(timezone.utc)
# Simple dataclass-like container for facts (avoid importing from memory_engine)
class MemoryFactForObservation:
def __init__(self, id: str, text: str, fact_type: str, context: str, occurred_start: Optional[str]):
self.id = id
self.text = text
self.fact_type = fact_type
self.context = context
self.occurred_start = occurred_start
async def regenerate_observations_batch(
conn,
embeddings_model,
llm_config,
bank_id: str,
entity_links: List[EntityLink],
log_buffer: List[str] = None
) -> None:
"""
Regenerate observations for top entities in this batch.
Called INSIDE the retain transaction for atomicity - if observations
fail, the entire retain batch is rolled back.
Args:
conn: Database connection (from the retain transaction)
embeddings_model: Embeddings model for generating observation embeddings
llm_config: LLM configuration for observation extraction
bank_id: Bank identifier
entity_links: Entity links from this batch
log_buffer: Optional log buffer for timing
"""
TOP_N_ENTITIES = 5
MIN_FACTS_THRESHOLD = 5
if not entity_links:
return
# Count mentions per entity in this batch
entity_mention_counts: Dict[str, int] = {}
for link in entity_links:
if link.entity_id:
entity_id = str(link.entity_id)
entity_mention_counts[entity_id] = entity_mention_counts.get(entity_id, 0) + 1
if not entity_mention_counts:
return
# Sort by mention count descending and take top N
sorted_entities = sorted(
entity_mention_counts.items(),
key=lambda x: x[1],
reverse=True
)
entities_to_process = [e[0] for e in sorted_entities[:TOP_N_ENTITIES]]
obs_start = time.time()
# Convert to UUIDs
entity_uuids = [uuid.UUID(eid) if isinstance(eid, str) else eid for eid in entities_to_process]
# Batch query for entity names
entity_rows = await conn.fetch(
"""
SELECT id, canonical_name FROM entities
WHERE id = ANY($1) AND bank_id = $2
""",
entity_uuids, bank_id
)
entity_names = {row['id']: row['canonical_name'] for row in entity_rows}
# Batch query for fact counts
fact_counts = await conn.fetch(
"""
SELECT ue.entity_id, COUNT(*) as cnt
FROM unit_entities ue
JOIN memory_units mu ON ue.unit_id = mu.id
WHERE ue.entity_id = ANY($1) AND mu.bank_id = $2
GROUP BY ue.entity_id
""",
entity_uuids, bank_id
)
entity_fact_counts = {row['entity_id']: row['cnt'] for row in fact_counts}
# Filter entities that meet the threshold
entities_with_names = []
for entity_id in entities_to_process:
entity_uuid = uuid.UUID(entity_id) if isinstance(entity_id, str) else entity_id
if entity_uuid not in entity_names:
continue
fact_count = entity_fact_counts.get(entity_uuid, 0)
if fact_count >= MIN_FACTS_THRESHOLD:
entities_with_names.append((entity_id, entity_names[entity_uuid]))
if not entities_with_names:
return
# Process entities SEQUENTIALLY (asyncpg doesn't allow concurrent queries on same connection)
# We must use the same connection to stay in the retain transaction
total_observations = 0
for entity_id, entity_name in entities_with_names:
try:
obs_ids = await _regenerate_entity_observations(
conn, embeddings_model, llm_config,
bank_id, entity_id, entity_name
)
total_observations += len(obs_ids)
except Exception as e:
logger.error(f"[OBSERVATIONS] Error processing entity {entity_id}: {e}")
obs_time = time.time() - obs_start
if log_buffer is not None:
log_buffer.append(f"[11] Observations: {total_observations} observations for {len(entities_with_names)} entities in {obs_time:.3f}s")
async def _regenerate_entity_observations(
conn,
embeddings_model,
llm_config,
bank_id: str,
entity_id: str,
entity_name: str
) -> List[str]:
"""
Regenerate observations for a single entity.
Uses the provided connection (part of retain transaction).
Args:
conn: Database connection (from the retain transaction)
embeddings_model: Embeddings model
llm_config: LLM configuration
bank_id: Bank identifier
entity_id: Entity UUID
entity_name: Canonical name of the entity
Returns:
List of created observation IDs
"""
entity_uuid = uuid.UUID(entity_id) if isinstance(entity_id, str) else entity_id
# Get all facts mentioning this entity (exclude observations themselves)
rows = await conn.fetch(
"""
SELECT mu.id, mu.text, mu.context, mu.occurred_start, mu.fact_type
FROM memory_units mu
JOIN unit_entities ue ON mu.id = ue.unit_id
WHERE mu.bank_id = $1
AND ue.entity_id = $2
AND mu.fact_type IN ('world', 'experience')
ORDER BY mu.occurred_start DESC
LIMIT 50
""",
bank_id, entity_uuid
)
if not rows:
return []
# Convert to fact objects for observation extraction
facts = []
for row in rows:
occurred_start = row['occurred_start'].isoformat() if row['occurred_start'] else None
facts.append(MemoryFactForObservation(
id=str(row['id']),
text=row['text'],
fact_type=row['fact_type'],
context=row['context'],
occurred_start=occurred_start
))
# Extract observations using LLM
observations = await observation_utils.extract_observations_from_facts(
llm_config,
entity_name,
facts
)
if not observations:
return []
# Delete old observations for this entity
await conn.execute(
"""
DELETE FROM memory_units
WHERE id IN (
SELECT mu.id
FROM memory_units mu
JOIN unit_entities ue ON mu.id = ue.unit_id
WHERE mu.bank_id = $1
AND mu.fact_type = 'observation'
AND ue.entity_id = $2
)
""",
bank_id, entity_uuid
)
# Generate embeddings for new observations
embeddings = await embedding_utils.generate_embeddings_batch(
embeddings_model, observations
)
# Insert new observations
current_time = utcnow()
created_ids = []
for obs_text, embedding in zip(observations, embeddings):
result = await conn.fetchrow(
"""
INSERT INTO memory_units (
bank_id, text, embedding, context, event_date,
occurred_start, occurred_end, mentioned_at,
fact_type, access_count
)
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, 'observation', 0)
RETURNING id
""",
bank_id,
obs_text,
str(embedding),
f"observation about {entity_name}",
current_time,
current_time,
current_time,
current_time
)
obs_id = str(result['id'])
created_ids.append(obs_id)
# Link observation to entity
await conn.execute(
"""
INSERT INTO unit_entities (unit_id, entity_id)
VALUES ($1, $2)
""",
uuid.UUID(obs_id), entity_uuid
)
return created_ids
@@ -1,392 +0,0 @@
"""
Main orchestrator for the retain pipeline.
Coordinates all retain pipeline modules to store memories efficiently.
"""
import logging
import time
import uuid
from datetime import datetime, timezone
from typing import List, Dict, Any, Optional
from . import bank_utils
from ..db_utils import acquire_with_retry
def utcnow():
"""Get current UTC time."""
return datetime.now(timezone.utc)
from .types import RetainContent, ExtractedFact, ProcessedFact, EntityLink
from . import (
fact_extraction,
embedding_processing,
deduplication,
chunk_storage,
fact_storage,
entity_processing,
link_creation,
observation_regeneration
)
logger = logging.getLogger(__name__)
async def retain_batch(
pool,
embeddings_model,
llm_config,
entity_resolver,
task_backend,
format_date_fn,
duplicate_checker_fn,
bank_id: str,
contents_dicts: List[Dict[str, Any]],
document_id: Optional[str] = None,
is_first_batch: bool = True,
fact_type_override: Optional[str] = None,
confidence_score: Optional[float] = None,
) -> List[List[str]]:
"""
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
task_backend: Task backend for background jobs
format_date_fn: Function to format datetime to readable string
duplicate_checker_fn: Function to check for duplicate facts
bank_id: Bank identifier
contents_dicts: List of content dictionaries
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
Returns:
List of unit ID lists (one list per content item)
"""
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:
content = RetainContent(
content=item["content"],
context=item.get("context", ""),
event_date=item.get("event_date") or utcnow(),
metadata=item.get("metadata", {})
)
contents.append(content)
# Step 1: Extract facts from all contents
step_start = time.time()
extract_opinions = (fact_type_override == 'opinion')
extracted_facts, chunks = await fact_extraction.extract_facts_from_contents(
contents,
llm_config,
agent_name,
extract_opinions
)
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:
return [[] for _ in contents]
# 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)
]
# Track document IDs for logging
document_ids_added = []
# 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
async with acquire_with_retry(pool) as conn:
async with conn.transaction():
# Ensure bank exists
await fact_storage.ensure_bank_exists(conn, bank_id)
# 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 = {}
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
)
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])
# 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
)
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
# Deduplication
step_start = time.time()
is_duplicate_flags = await deduplication.check_duplicates_batch(
conn, bank_id, processed_facts, duplicate_checker_fn
)
log_buffer.append(f"[4] Deduplication: {sum(is_duplicate_flags)} duplicates in {time.time() - step_start:.3f}s")
# Filter out duplicates
non_duplicate_facts = deduplication.filter_duplicates(processed_facts, is_duplicate_flags)
if not non_duplicate_facts:
return [[] for _ in contents]
# 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()
entity_links = await entity_processing.process_entities_batch(
entity_resolver, conn, bank_id, unit_ids, non_duplicate_facts, log_buffer
)
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")
# Regenerate observations INSIDE transaction for atomicity
await observation_regeneration.regenerate_observations_batch(
conn,
embeddings_model,
llm_config,
bank_id,
entity_links,
log_buffer
)
# Map results back to original content items
result_unit_ids = _map_results_to_contents(
contents, extracted_facts, is_duplicate_flags, unit_ids
)
# Trigger background tasks AFTER transaction commits (opinion reinforcement only)
await _trigger_background_tasks(
task_backend,
bank_id,
unit_ids,
non_duplicate_facts
)
# 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")
return result_unit_ids
def _map_results_to_contents(
contents: List[RetainContent],
extracted_facts: List[ExtractedFact],
is_duplicate_flags: List[bool],
unit_ids: List[str]
) -> List[List[str]]:
"""
Map created unit IDs back to original content items.
Accounts for duplicates when mapping back.
"""
result_unit_ids = []
filtered_idx = 0
# Group facts by content_index
facts_by_content = {i: [] for i in range(len(contents))}
for i, fact in enumerate(extracted_facts):
facts_by_content[fact.content_index].append(i)
for content_index in range(len(contents)):
content_unit_ids = []
for fact_idx in facts_by_content[content_index]:
if not is_duplicate_flags[fact_idx]:
content_unit_ids.append(unit_ids[filtered_idx])
filtered_idx += 1
result_unit_ids.append(content_unit_ids)
return result_unit_ids
async def _trigger_background_tasks(
task_backend,
bank_id: str,
unit_ids: List[str],
facts: List[ProcessedFact],
) -> None:
"""Trigger opinion reinforcement as background task (after transaction commits)."""
# Trigger opinion reinforcement if there are entities
fact_entities = [[e.name for e in fact.entities] for fact in facts]
if any(fact_entities):
await task_backend.submit_task({
'type': 'reinforce_opinion',
'bank_id': bank_id,
'created_unit_ids': unit_ids,
'unit_texts': [fact.fact_text for fact in facts],
'unit_entities': fact_entities
})
@@ -1,220 +0,0 @@
"""
Type definitions for the retain pipeline.
These dataclasses provide type safety throughout the retain operation,
from content input to fact storage.
"""
from dataclasses import dataclass, field
from typing import List, Optional, Dict, Any
from datetime import datetime
from uuid import UUID
@dataclass
class RetainContent:
"""
Input content item to be retained as memories.
Represents a single piece of content to extract facts from.
"""
content: str
context: str = ""
event_date: Optional[datetime] = None
metadata: Dict[str, str] = field(default_factory=dict)
def __post_init__(self):
"""Ensure event_date is set."""
if self.event_date is None:
from datetime import datetime, timezone
self.event_date = datetime.now(timezone.utc)
@dataclass
class ChunkMetadata:
"""
Metadata about a text chunk.
Used to track which facts were extracted from which chunks.
"""
chunk_text: str
fact_count: int
content_index: int # Index of the source content
chunk_index: int # Global chunk index across all contents
@dataclass
class EntityRef:
"""
Reference to an entity mentioned in a fact.
Entities are extracted by the LLM during fact extraction.
"""
name: str
canonical_name: Optional[str] = None # Resolved canonical name
entity_id: Optional[UUID] = None # Resolved entity ID
@dataclass
class CausalRelation:
"""
Causal relationship between facts.
Represents how one fact causes, enables, or prevents another.
"""
relation_type: str # "causes", "enables", "prevents", "caused_by"
target_fact_index: int # Index of the target fact in the batch
strength: float = 1.0 # Strength of the causal relationship
@dataclass
class ExtractedFact:
"""
Fact extracted from content by the LLM.
This is the raw output from fact extraction before processing.
"""
fact_text: str
fact_type: str # "world", "experience", "opinion", "observation"
entities: List[str] = field(default_factory=list)
occurred_start: Optional[datetime] = None
occurred_end: Optional[datetime] = None
where: Optional[str] = None # WHERE the fact occurred or is about
causal_relations: List[CausalRelation] = field(default_factory=list)
# Context from the content item
content_index: int = 0 # Which content this fact came from
chunk_index: int = 0 # Which chunk this fact came from
context: str = ""
mentioned_at: Optional[datetime] = None
metadata: Dict[str, str] = field(default_factory=dict)
@dataclass
class ProcessedFact:
"""
Fact after processing and ready for storage.
Includes resolved entities, embeddings, and all necessary fields.
"""
# Core fact data
fact_text: str
fact_type: str
embedding: List[float]
# Temporal data
occurred_start: Optional[datetime]
occurred_end: Optional[datetime]
mentioned_at: datetime
# Context and metadata
context: str
metadata: Dict[str, str]
# Location data
where: Optional[str] = None
# Entities
entities: List[EntityRef] = field(default_factory=list)
# Causal relations
causal_relations: List[CausalRelation] = field(default_factory=list)
# Chunk reference
chunk_id: Optional[str] = None
# Document reference (denormalized for query performance)
document_id: Optional[str] = None
# DB fields (set after insertion)
unit_id: Optional[UUID] = None
@property
def is_duplicate(self) -> bool:
"""Check if this fact was marked as a duplicate."""
return self.unit_id is None
@staticmethod
def from_extracted_fact(
extracted_fact: 'ExtractedFact',
embedding: List[float],
chunk_id: Optional[str] = None
) -> 'ProcessedFact':
"""
Create ProcessedFact from ExtractedFact.
Args:
extracted_fact: Source ExtractedFact
embedding: Generated embedding vector
chunk_id: Optional chunk ID
Returns:
ProcessedFact ready for storage
"""
from datetime import datetime, timezone
# Use occurred dates only if explicitly provided by LLM
occurred_start = extracted_fact.occurred_start
occurred_end = extracted_fact.occurred_end
mentioned_at = extracted_fact.mentioned_at or datetime.now(timezone.utc)
# Convert entity strings to EntityRef objects
entities = [EntityRef(name=name) for name in extracted_fact.entities]
return ProcessedFact(
fact_text=extracted_fact.fact_text,
fact_type=extracted_fact.fact_type,
embedding=embedding,
occurred_start=occurred_start,
occurred_end=occurred_end,
mentioned_at=mentioned_at,
context=extracted_fact.context,
metadata=extracted_fact.metadata,
entities=entities,
causal_relations=extracted_fact.causal_relations,
chunk_id=chunk_id
)
@dataclass
class EntityLink:
"""
Link between two memory units through a shared entity.
Used for entity-based graph connections in the memory graph.
"""
from_unit_id: UUID
to_unit_id: UUID
entity_id: UUID
link_type: str = 'entity'
weight: float = 1.0
@dataclass
class RetainBatch:
"""
A batch of content to retain.
Tracks all facts, chunks, and metadata for a batch operation.
"""
bank_id: str
contents: List[RetainContent]
document_id: Optional[str] = None
fact_type_override: Optional[str] = None
confidence_score: Optional[float] = None
# Extracted data (populated during processing)
extracted_facts: List[ExtractedFact] = field(default_factory=list)
processed_facts: List[ProcessedFact] = field(default_factory=list)
chunks: List[ChunkMetadata] = field(default_factory=list)
# Results (populated after storage)
unit_ids_by_content: List[List[str]] = field(default_factory=list)
def get_facts_for_content(self, content_index: int) -> List[ExtractedFact]:
"""Get all extracted facts for a specific content item."""
return [f for f in self.extracted_facts if f.content_index == content_index]
def get_chunks_for_content(self, content_index: int) -> List[ChunkMetadata]:
"""Get all chunks for a specific content item."""
return [c for c in self.chunks if c.content_index == content_index]
@@ -3,27 +3,13 @@ Search module for memory retrieval.
Provides modular search architecture:
- Retrieval: 4-way parallel (semantic + BM25 + graph + temporal)
- Graph retrieval: Pluggable strategies (BFS, PPR)
- Reranking: Pluggable strategies (heuristic, cross-encoder)
"""
from .retrieval import (
retrieve_parallel,
get_default_graph_retriever,
set_default_graph_retriever,
ParallelRetrievalResult,
)
from .graph_retrieval import GraphRetriever, BFSGraphRetriever
from .mpfp_retrieval import MPFPGraphRetriever
from .retrieval import retrieve_parallel
from .reranking import CrossEncoderReranker
__all__ = [
"retrieve_parallel",
"get_default_graph_retriever",
"set_default_graph_retriever",
"ParallelRetrievalResult",
"GraphRetriever",
"BFSGraphRetriever",
"MPFPGraphRetriever",
"CrossEncoderReranker",
]
@@ -1,235 +0,0 @@
"""
Graph retrieval strategies for memory recall.
This module provides an abstraction for graph-based memory retrieval,
allowing different algorithms (BFS spreading activation, PPR, etc.) to be
swapped without changing the rest of the recall pipeline.
"""
from abc import ABC, abstractmethod
from typing import List, Optional
from datetime import datetime
import logging
from .types import RetrievalResult
from ..db_utils import acquire_with_retry
logger = logging.getLogger(__name__)
class GraphRetriever(ABC):
"""
Abstract base class for graph-based memory retrieval.
Implementations traverse the memory graph (entity links, temporal links,
causal links) to find relevant facts that might not be found by
semantic or keyword search alone.
"""
@property
@abstractmethod
def name(self) -> str:
"""Return identifier for this retrieval strategy (e.g., 'bfs', 'mpfp')."""
pass
@abstractmethod
async def retrieve(
self,
pool,
query_embedding_str: str,
bank_id: str,
fact_type: str,
budget: int,
query_text: Optional[str] = None,
semantic_seeds: Optional[List[RetrievalResult]] = None,
temporal_seeds: Optional[List[RetrievalResult]] = None,
) -> List[RetrievalResult]:
"""
Retrieve relevant facts via graph traversal.
Args:
pool: Database connection pool
query_embedding_str: Query embedding as string (for finding entry points)
bank_id: Memory bank identifier
fact_type: Fact type to filter ('world', 'experience', 'opinion', 'observation')
budget: Maximum number of nodes to explore/return
query_text: Original query text (optional, for some strategies)
semantic_seeds: Pre-computed semantic entry points (from semantic retrieval)
temporal_seeds: Pre-computed temporal entry points (from temporal retrieval)
Returns:
List of RetrievalResult objects with activation scores set
"""
pass
class BFSGraphRetriever(GraphRetriever):
"""
Graph retrieval using BFS-style spreading activation.
Starting from semantic entry points, spreads activation through
the memory graph (entity, temporal, causal links) using breadth-first
traversal with decaying activation.
This is the original Hindsight graph retrieval algorithm.
"""
def __init__(
self,
entry_point_limit: int = 5,
entry_point_threshold: float = 0.5,
activation_decay: float = 0.8,
min_activation: float = 0.1,
batch_size: int = 20,
):
"""
Initialize BFS graph retriever.
Args:
entry_point_limit: Maximum number of entry points to start from
entry_point_threshold: Minimum semantic similarity for entry points
activation_decay: Decay factor per hop (activation *= decay)
min_activation: Minimum activation to continue spreading
batch_size: Number of nodes to process per batch (for neighbor fetching)
"""
self.entry_point_limit = entry_point_limit
self.entry_point_threshold = entry_point_threshold
self.activation_decay = activation_decay
self.min_activation = min_activation
self.batch_size = batch_size
@property
def name(self) -> str:
return "bfs"
async def retrieve(
self,
pool,
query_embedding_str: str,
bank_id: str,
fact_type: str,
budget: int,
query_text: Optional[str] = None,
semantic_seeds: Optional[List[RetrievalResult]] = None,
temporal_seeds: Optional[List[RetrievalResult]] = None,
) -> List[RetrievalResult]:
"""
Retrieve facts using BFS spreading activation.
Algorithm:
1. Find entry points (top semantic matches above threshold)
2. BFS traversal: visit neighbors, propagate decaying activation
3. Boost causal links (causes, enables, prevents)
4. Return visited nodes up to budget
Note: BFS finds its own entry points via embedding search.
The semantic_seeds and temporal_seeds parameters are accepted
for interface compatibility but not used.
"""
async with acquire_with_retry(pool) as conn:
return await self._retrieve_with_conn(
conn, query_embedding_str, bank_id, fact_type, budget
)
async def _retrieve_with_conn(
self,
conn,
query_embedding_str: str,
bank_id: str,
fact_type: str,
budget: int,
) -> List[RetrievalResult]:
"""Internal implementation with connection."""
# Step 1: Find entry points
entry_points = await conn.fetch(
"""
SELECT id, text, context, event_date, occurred_start, occurred_end,
mentioned_at, access_count, embedding, fact_type, document_id, chunk_id,
1 - (embedding <=> $1::vector) AS similarity
FROM memory_units
WHERE bank_id = $2
AND embedding IS NOT NULL
AND fact_type = $3
AND (1 - (embedding <=> $1::vector)) >= $4
ORDER BY embedding <=> $1::vector
LIMIT $5
""",
query_embedding_str, bank_id, fact_type,
self.entry_point_threshold, self.entry_point_limit
)
if not entry_points:
return []
# Step 2: BFS spreading activation
visited = set()
results = []
queue = [
(RetrievalResult.from_db_row(dict(r)), r["similarity"])
for r in entry_points
]
budget_remaining = budget
while queue and budget_remaining > 0:
# Collect a batch of nodes to process
batch_nodes = []
batch_activations = {}
while queue and len(batch_nodes) < self.batch_size and budget_remaining > 0:
current, activation = queue.pop(0)
unit_id = current.id
if unit_id not in visited:
visited.add(unit_id)
budget_remaining -= 1
current.activation = activation
results.append(current)
batch_nodes.append(current.id)
batch_activations[unit_id] = activation
# Batch fetch neighbors
if batch_nodes and budget_remaining > 0:
max_neighbors = len(batch_nodes) * 20
neighbors = await conn.fetch(
"""
SELECT mu.id, mu.text, mu.context, mu.occurred_start, mu.occurred_end,
mu.mentioned_at, mu.access_count, mu.embedding, mu.fact_type,
mu.document_id, mu.chunk_id,
ml.weight, ml.link_type, ml.from_unit_id
FROM memory_links ml
JOIN memory_units mu ON ml.to_unit_id = mu.id
WHERE ml.from_unit_id = ANY($1::uuid[])
AND ml.weight >= $2
AND mu.fact_type = $3
ORDER BY ml.weight DESC
LIMIT $4
""",
batch_nodes, self.min_activation, fact_type, max_neighbors
)
for n in neighbors:
neighbor_id = str(n["id"])
if neighbor_id not in visited:
parent_id = str(n["from_unit_id"])
parent_activation = batch_activations.get(parent_id, 0.5)
# Boost causal links
link_type = n["link_type"]
base_weight = n["weight"]
if link_type in ("causes", "caused_by"):
causal_boost = 2.0
elif link_type in ("enables", "prevents"):
causal_boost = 1.5
else:
causal_boost = 1.0
effective_weight = base_weight * causal_boost
new_activation = parent_activation * effective_weight * self.activation_decay
if new_activation > self.min_activation:
neighbor_result = RetrievalResult.from_db_row(dict(n))
queue.append((neighbor_result, new_activation))
return results
@@ -1,454 +0,0 @@
"""
Meta-Path Forward Push (MPFP) graph retrieval.
A sublinear graph traversal algorithm for memory retrieval over heterogeneous
graphs with multiple edge types (semantic, temporal, causal, entity).
Combines meta-path patterns from HIN literature with Forward Push local
propagation from Approximate PPR.
Key properties:
- Sublinear in graph size (threshold pruning bounds active nodes)
- Predefined patterns capture different retrieval intents
- All patterns run in parallel, results fused via RRF
- No LLM in the loop during traversal
"""
import asyncio
import logging
from dataclasses import dataclass, field
from typing import List, Dict, Optional, Tuple
from collections import defaultdict
from .types import RetrievalResult
from .graph_retrieval import GraphRetriever
from ..db_utils import acquire_with_retry
logger = logging.getLogger(__name__)
# -----------------------------------------------------------------------------
# Data Classes
# -----------------------------------------------------------------------------
@dataclass
class EdgeTarget:
"""A neighbor node with its edge weight."""
node_id: str
weight: float
@dataclass
class TypedAdjacency:
"""Adjacency lists split by edge type."""
# edge_type -> from_node_id -> list of (to_node_id, weight)
graphs: Dict[str, Dict[str, List[EdgeTarget]]] = field(default_factory=dict)
def get_neighbors(self, edge_type: str, node_id: str) -> List[EdgeTarget]:
"""Get neighbors for a node via a specific edge type."""
return self.graphs.get(edge_type, {}).get(node_id, [])
def get_normalized_neighbors(
self,
edge_type: str,
node_id: str,
top_k: int
) -> List[EdgeTarget]:
"""Get top-k neighbors with weights normalized to sum to 1."""
neighbors = self.get_neighbors(edge_type, node_id)[:top_k]
if not neighbors:
return []
total = sum(n.weight for n in neighbors)
if total == 0:
return []
return [
EdgeTarget(node_id=n.node_id, weight=n.weight / total)
for n in neighbors
]
@dataclass
class PatternResult:
"""Result from a single pattern traversal."""
pattern: List[str]
scores: Dict[str, float] # node_id -> accumulated mass
@dataclass
class MPFPConfig:
"""Configuration for MPFP algorithm."""
alpha: float = 0.15 # teleport/keep probability
threshold: float = 1e-6 # mass pruning threshold (lower = explore more)
top_k_neighbors: int = 20 # fan-out limit per node
# Patterns from semantic seeds
patterns_semantic: List[List[str]] = field(default_factory=lambda: [
['semantic', 'semantic'], # topic expansion
['entity', 'temporal'], # entity timeline
['semantic', 'causes'], # reasoning chains (forward)
['semantic', 'caused_by'], # reasoning chains (backward)
['entity', 'semantic'], # entity context
])
# Patterns from temporal seeds
patterns_temporal: List[List[str]] = field(default_factory=lambda: [
['temporal', 'semantic'], # what was happening then
['temporal', 'entity'], # who was involved then
])
@dataclass
class SeedNode:
"""An entry point node with its initial score."""
node_id: str
score: float # initial mass (e.g., similarity score)
# -----------------------------------------------------------------------------
# Core Algorithm
# -----------------------------------------------------------------------------
def mpfp_traverse(
seeds: List[SeedNode],
pattern: List[str],
adjacency: TypedAdjacency,
config: MPFPConfig,
) -> PatternResult:
"""
Forward Push traversal following a meta-path pattern.
Args:
seeds: Entry point nodes with initial scores
pattern: Sequence of edge types to follow
adjacency: Typed adjacency structure
config: Algorithm parameters
Returns:
PatternResult with accumulated scores per node
"""
if not seeds:
return PatternResult(pattern=pattern, scores={})
scores: Dict[str, float] = {}
# Initialize frontier with seed masses (normalized)
total_seed_score = sum(s.score for s in seeds)
if total_seed_score == 0:
total_seed_score = len(seeds) # fallback to uniform
frontier: Dict[str, float] = {
s.node_id: s.score / total_seed_score for s in seeds
}
# Follow pattern hop by hop
for edge_type in pattern:
next_frontier: Dict[str, float] = {}
for node_id, mass in frontier.items():
if mass < config.threshold:
continue
# Keep α portion for this node
scores[node_id] = scores.get(node_id, 0) + config.alpha * mass
# Push (1-α) to neighbors
push_mass = (1 - config.alpha) * mass
neighbors = adjacency.get_normalized_neighbors(
edge_type, node_id, config.top_k_neighbors
)
for neighbor in neighbors:
next_frontier[neighbor.node_id] = (
next_frontier.get(neighbor.node_id, 0) +
push_mass * neighbor.weight
)
frontier = next_frontier
# Final frontier nodes get their remaining mass
for node_id, mass in frontier.items():
if mass >= config.threshold:
scores[node_id] = scores.get(node_id, 0) + mass
return PatternResult(pattern=pattern, scores=scores)
def rrf_fusion(
results: List[PatternResult],
k: int = 60,
top_k: int = 50,
) -> List[Tuple[str, float]]:
"""
Reciprocal Rank Fusion to combine pattern results.
Args:
results: List of pattern results
k: RRF constant (higher = more uniform weighting)
top_k: Number of results to return
Returns:
List of (node_id, fused_score) tuples, sorted by score descending
"""
fused: Dict[str, float] = {}
for result in results:
if not result.scores:
continue
# Rank nodes by their score in this pattern
ranked = sorted(
result.scores.keys(),
key=lambda n: result.scores[n],
reverse=True
)
for rank, node_id in enumerate(ranked):
fused[node_id] = fused.get(node_id, 0) + 1.0 / (k + rank + 1)
# Sort by fused score and return top-k
sorted_results = sorted(
fused.items(),
key=lambda x: x[1],
reverse=True
)
return sorted_results[:top_k]
# -----------------------------------------------------------------------------
# Database Loading
# -----------------------------------------------------------------------------
async def load_typed_adjacency(pool, bank_id: str) -> TypedAdjacency:
"""
Load all edges for a bank, split by edge type.
Single query, then organize in-memory for fast traversal.
"""
async with acquire_with_retry(pool) as conn:
rows = await conn.fetch(
"""
SELECT ml.from_unit_id, ml.to_unit_id, ml.link_type, ml.weight
FROM memory_links ml
JOIN memory_units mu ON ml.from_unit_id = mu.id
WHERE mu.bank_id = $1
AND ml.weight >= 0.1
ORDER BY ml.from_unit_id, ml.weight DESC
""",
bank_id
)
graphs: Dict[str, Dict[str, List[EdgeTarget]]] = defaultdict(
lambda: defaultdict(list)
)
for row in rows:
from_id = str(row['from_unit_id'])
to_id = str(row['to_unit_id'])
link_type = row['link_type']
weight = row['weight']
graphs[link_type][from_id].append(
EdgeTarget(node_id=to_id, weight=weight)
)
return TypedAdjacency(graphs=dict(graphs))
async def fetch_memory_units_by_ids(
pool,
node_ids: List[str],
fact_type: str,
) -> List[RetrievalResult]:
"""Fetch full memory unit details for a list of node IDs."""
if not node_ids:
return []
async with acquire_with_retry(pool) as conn:
rows = await conn.fetch(
"""
SELECT id, text, context, event_date, occurred_start, occurred_end,
mentioned_at, access_count, embedding, fact_type, document_id, chunk_id
FROM memory_units
WHERE id = ANY($1::uuid[])
AND fact_type = $2
""",
node_ids,
fact_type
)
return [RetrievalResult.from_db_row(dict(r)) for r in rows]
# -----------------------------------------------------------------------------
# Graph Retriever Implementation
# -----------------------------------------------------------------------------
class MPFPGraphRetriever(GraphRetriever):
"""
Graph retrieval using Meta-Path Forward Push.
Runs predefined patterns in parallel from semantic and temporal seeds,
then fuses results via RRF.
"""
def __init__(self, config: Optional[MPFPConfig] = None):
"""
Initialize MPFP retriever.
Args:
config: Algorithm configuration (uses defaults if None)
"""
self.config = config or MPFPConfig()
self._adjacency_cache: Dict[str, TypedAdjacency] = {}
@property
def name(self) -> str:
return "mpfp"
async def retrieve(
self,
pool,
query_embedding_str: str,
bank_id: str,
fact_type: str,
budget: int,
query_text: Optional[str] = None,
semantic_seeds: Optional[List[RetrievalResult]] = None,
temporal_seeds: Optional[List[RetrievalResult]] = None,
) -> List[RetrievalResult]:
"""
Retrieve facts using MPFP algorithm.
Args:
pool: Database connection pool
query_embedding_str: Query embedding (used for fallback seed finding)
bank_id: Memory bank ID
fact_type: Fact type to filter
budget: Maximum results to return
query_text: Original query text (optional)
semantic_seeds: Pre-computed semantic entry points
temporal_seeds: Pre-computed temporal entry points
Returns:
List of RetrievalResult with activation scores
"""
# Load typed adjacency (could cache per bank_id with TTL)
adjacency = await load_typed_adjacency(pool, bank_id)
# Convert seeds to SeedNode format
semantic_seed_nodes = self._convert_seeds(semantic_seeds, 'similarity')
temporal_seed_nodes = self._convert_seeds(temporal_seeds, 'temporal_score')
# If no semantic seeds provided, fall back to finding our own
if not semantic_seed_nodes:
semantic_seed_nodes = await self._find_semantic_seeds(
pool, query_embedding_str, bank_id, fact_type
)
# Run all patterns in parallel
tasks = []
# Patterns from semantic seeds
for pattern in self.config.patterns_semantic:
if semantic_seed_nodes:
tasks.append(
asyncio.to_thread(
mpfp_traverse,
semantic_seed_nodes,
pattern,
adjacency,
self.config,
)
)
# Patterns from temporal seeds
for pattern in self.config.patterns_temporal:
if temporal_seed_nodes:
tasks.append(
asyncio.to_thread(
mpfp_traverse,
temporal_seed_nodes,
pattern,
adjacency,
self.config,
)
)
if not tasks:
return []
# Gather pattern results
pattern_results = await asyncio.gather(*tasks)
# Fuse results
fused = rrf_fusion(pattern_results, top_k=budget)
if not fused:
return []
# Get top result IDs (don't exclude seeds - they may be highly relevant)
result_ids = [node_id for node_id, score in fused][:budget]
# Fetch full details
results = await fetch_memory_units_by_ids(pool, result_ids, fact_type)
# Add activation scores from fusion
score_map = {node_id: score for node_id, score in fused}
for result in results:
result.activation = score_map.get(result.id, 0.0)
# Sort by activation
results.sort(key=lambda r: r.activation or 0, reverse=True)
return results
def _convert_seeds(
self,
seeds: Optional[List[RetrievalResult]],
score_attr: str,
) -> List[SeedNode]:
"""Convert RetrievalResult seeds to SeedNode format."""
if not seeds:
return []
result = []
for seed in seeds:
score = getattr(seed, score_attr, None)
if score is None:
score = seed.activation or seed.similarity or 1.0
result.append(SeedNode(node_id=seed.id, score=score))
return result
async def _find_semantic_seeds(
self,
pool,
query_embedding_str: str,
bank_id: str,
fact_type: str,
limit: int = 20,
threshold: float = 0.3,
) -> List[SeedNode]:
"""Fallback: find semantic seeds via embedding search."""
async with acquire_with_retry(pool) as conn:
rows = await conn.fetch(
"""
SELECT id, 1 - (embedding <=> $1::vector) AS similarity
FROM memory_units
WHERE bank_id = $2
AND embedding IS NOT NULL
AND fact_type = $3
AND (1 - (embedding <=> $1::vector)) >= $4
ORDER BY embedding <=> $1::vector
LIMIT $5
""",
query_embedding_str, bank_id, fact_type, threshold, limit
)
return [
SeedNode(node_id=str(r['id']), score=r['similarity'])
for r in rows
]
@@ -1,132 +0,0 @@
"""
Observation utilities for generating entity observations from facts.
Observations are objective facts synthesized from multiple memory facts
about an entity, without personality influence.
"""
import logging
from typing import List, Dict, Any
from pydantic import BaseModel, Field
from ..response_models import MemoryFact
logger = logging.getLogger(__name__)
class Observation(BaseModel):
"""An observation about an entity."""
observation: str = Field(description="The observation text - a factual statement about the entity")
class ObservationExtractionResponse(BaseModel):
"""Response containing extracted observations."""
observations: List[Observation] = Field(
default_factory=list,
description="List of observations about the entity"
)
def format_facts_for_observation_prompt(facts: List[MemoryFact]) -> str:
"""Format facts as text for observation extraction prompt."""
import json
if not facts:
return "[]"
formatted = []
for fact in facts:
fact_obj = {
"text": fact.text
}
# Add context if available
if fact.context:
fact_obj["context"] = fact.context
# Add occurred_start if available
if fact.occurred_start:
fact_obj["occurred_at"] = fact.occurred_start
formatted.append(fact_obj)
return json.dumps(formatted, indent=2)
def build_observation_prompt(
entity_name: str,
facts_text: str,
) -> str:
"""Build the observation extraction prompt for the LLM."""
return f"""Based on the following facts about "{entity_name}", generate a list of key observations.
FACTS ABOUT {entity_name.upper()}:
{facts_text}
Your task: Synthesize the facts into clear, objective observations about {entity_name}.
GUIDELINES:
1. Each observation should be a factual statement about {entity_name}
2. Combine related facts into single observations where appropriate
3. Be objective - do not add opinions, judgments, or interpretations
4. Focus on what we KNOW about {entity_name}, not what we assume
5. Include observations about: identity, characteristics, roles, relationships, activities
6. Write in third person (e.g., "John is..." not "I think John is...")
7. If there are conflicting facts, note the most recent or most supported one
EXAMPLES of good observations:
- "John works at Google as a software engineer"
- "John is detail-oriented and methodical in his approach"
- "John collaborates frequently with Sarah on the AI project"
- "John joined the company in 2023"
EXAMPLES of bad observations (avoid these):
- "John seems like a good person" (opinion/judgment)
- "John probably likes his job" (assumption)
- "I believe John is reliable" (first-person opinion)
Generate 3-7 observations based on the available facts. If there are very few facts, generate fewer observations."""
def get_observation_system_message() -> str:
"""Get the system message for observation extraction."""
return "You are an objective observer synthesizing facts about an entity. Generate clear, factual observations without opinions or personality influence. Be concise and accurate."
async def extract_observations_from_facts(
llm_config,
entity_name: str,
facts: List[MemoryFact]
) -> List[str]:
"""
Extract observations from facts about an entity using LLM.
Args:
llm_config: LLM configuration to use
entity_name: Name of the entity to generate observations about
facts: List of facts mentioning the entity
Returns:
List of observation strings
"""
if not facts:
return []
facts_text = format_facts_for_observation_prompt(facts)
prompt = build_observation_prompt(entity_name, facts_text)
try:
result = await llm_config.call(
messages=[
{"role": "system", "content": get_observation_system_message()},
{"role": "user", "content": prompt}
],
response_format=ObservationExtractionResponse,
scope="memory_extract_observation"
)
observations = [op.observation for op in result.observations]
return observations
except Exception as e:
logger.warning(f"Failed to extract observations for {entity_name}: {str(e)}")
return []
@@ -2,16 +2,17 @@
Cross-encoder neural reranking for search results.
"""
from typing import List
from .types import MergedCandidate, ScoredResult
from typing import List, Dict, Any
class CrossEncoderReranker:
"""
Neural reranking using a cross-encoder model.
Configured via environment variables (see cross_encoder.py).
Default local model is cross-encoder/ms-marco-MiniLM-L-6-v2.
Uses cross-encoder/ms-marco-MiniLM-L-6-v2 by default:
- Fast inference (~80ms for 100 pairs on CPU)
- Small model (80MB)
- Trained for passage re-ranking
"""
def __init__(self, cross_encoder=None):
@@ -19,52 +20,41 @@ class CrossEncoderReranker:
Initialize cross-encoder reranker.
Args:
cross_encoder: CrossEncoderModel instance. If None, creates one from
environment variables (defaults to local provider)
cross_encoder: CrossEncoderReranker instance. If None, uses default
SentenceTransformersCrossEncoder with ms-marco-MiniLM-L-6-v2
"""
if cross_encoder is None:
from hindsight_api.engine.cross_encoder import create_cross_encoder_from_env
cross_encoder = create_cross_encoder_from_env()
from hindsight_api.engine.cross_encoder import SentenceTransformersCrossEncoder
cross_encoder = SentenceTransformersCrossEncoder()
self.cross_encoder = cross_encoder
def rerank(
self,
query: str,
candidates: List[MergedCandidate]
) -> List[ScoredResult]:
"""
Rerank candidates using cross-encoder scores.
Args:
query: Search query
candidates: Merged candidates from RRF
Returns:
List of ScoredResult objects sorted by cross-encoder score
"""
candidates: List[Dict[str, Any]]
) -> List[Dict[str, Any]]:
"""Rerank using cross-encoder scores."""
if not candidates:
return []
return candidates
# Prepare query-document pairs with date information
pairs = []
for candidate in candidates:
retrieval = candidate.retrieval
for c in candidates:
# Use text + context for better ranking
doc_text = retrieval.text
if retrieval.context:
doc_text = f"{retrieval.context}: {doc_text}"
doc_text = c["text"]
if c.get("context"):
doc_text = f"{c['context']}: {doc_text}"
# Add formatted date information for temporal awareness
if retrieval.occurred_start:
occurred_start = retrieval.occurred_start
if c.get("event_date"):
event_date = c["event_date"]
# Format in two styles for better model understanding
# 1. ISO format: YYYY-MM-DD
date_iso = occurred_start.strftime("%Y-%m-%d")
date_iso = event_date.strftime("%Y-%m-%d")
# 2. Human-readable: "June 5, 2022"
date_readable = occurred_start.strftime("%B %d, %Y")
date_readable = event_date.strftime("%B %d, %Y")
# Prepend date to document text
doc_text = f"[Date: {date_readable} ({date_iso})] {doc_text}"
@@ -82,18 +72,13 @@ class CrossEncoderReranker:
normalized_scores = [sigmoid(score) for score in scores]
# Create ScoredResult objects with cross-encoder scores
scored_results = []
for candidate, raw_score, norm_score in zip(candidates, scores, normalized_scores):
scored_result = ScoredResult(
candidate=candidate,
cross_encoder_score=float(raw_score),
cross_encoder_score_normalized=float(norm_score),
weight=float(norm_score) # Initial weight is just cross-encoder score
)
scored_results.append(scored_result)
# Assign normalized scores to candidates
for c, raw_score, norm_score in zip(candidates, scores, normalized_scores):
c["weight"] = float(norm_score)
c["cross_encoder_score"] = float(raw_score)
c["cross_encoder_score_normalized"] = float(norm_score)
# Sort by cross-encoder score
scored_results.sort(key=lambda x: x.weight, reverse=True)
candidates.sort(key=lambda x: x["weight"], reverse=True)
return scored_results
return candidates
@@ -4,119 +4,72 @@ Retrieval module for 4-way parallel search.
Implements:
1. Semantic retrieval (vector similarity)
2. BM25 retrieval (keyword/full-text search)
3. Graph retrieval (via pluggable GraphRetriever interface)
3. Graph retrieval (spreading activation)
4. Temporal retrieval (time-aware search with spreading)
"""
from typing import List, Dict, Optional
from dataclasses import dataclass, field
from typing import List, Dict, Any, Tuple, Optional
from datetime import datetime
import asyncio
import logging
from ..db_utils import acquire_with_retry
from .types import RetrievalResult
from .graph_retrieval import GraphRetriever, BFSGraphRetriever
from .mpfp_retrieval import MPFPGraphRetriever
from ...config import get_config
logger = logging.getLogger(__name__)
@dataclass
class ParallelRetrievalResult:
"""Result from parallel retrieval across all methods."""
semantic: List[RetrievalResult]
bm25: List[RetrievalResult]
graph: List[RetrievalResult]
temporal: Optional[List[RetrievalResult]]
timings: Dict[str, float] = field(default_factory=dict)
temporal_constraint: Optional[tuple] = None # (start_date, end_date)
# Default graph retriever instance (can be overridden)
_default_graph_retriever: Optional[GraphRetriever] = None
def get_default_graph_retriever() -> GraphRetriever:
"""Get or create the default graph retriever based on config."""
global _default_graph_retriever
if _default_graph_retriever is None:
config = get_config()
retriever_type = config.graph_retriever.lower()
if retriever_type == "mpfp":
_default_graph_retriever = MPFPGraphRetriever()
logger.info("Using MPFP graph retriever")
elif retriever_type == "bfs":
_default_graph_retriever = BFSGraphRetriever()
logger.info("Using BFS graph retriever")
else:
logger.warning(f"Unknown graph retriever '{retriever_type}', falling back to MPFP")
_default_graph_retriever = MPFPGraphRetriever()
return _default_graph_retriever
def set_default_graph_retriever(retriever: GraphRetriever) -> None:
"""Set the default graph retriever (for configuration/testing)."""
global _default_graph_retriever
_default_graph_retriever = retriever
async def retrieve_semantic(
conn,
query_emb_str: str,
bank_id: str,
agent_id: str,
fact_type: str,
limit: int
) -> List[RetrievalResult]:
) -> List[Tuple[str, Dict[str, Any]]]:
"""
Semantic retrieval via vector similarity.
Args:
conn: Database connection
query_emb_str: Query embedding as string
agent_id: bank ID
agent_id: Agent ID
fact_type: Fact type to filter
limit: Maximum results to return
Returns:
List of RetrievalResult objects
List of (doc_id, data) tuples
"""
results = await conn.fetch(
"""
SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, access_count, embedding, fact_type, document_id, chunk_id,
SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, access_count, embedding, fact_type, document_id,
1 - (embedding <=> $1::vector) AS similarity
FROM memory_units
WHERE bank_id = $2
WHERE agent_id = $2
AND embedding IS NOT NULL
AND fact_type = $3
AND (1 - (embedding <=> $1::vector)) >= 0.3
ORDER BY embedding <=> $1::vector
LIMIT $4
""",
query_emb_str, bank_id, fact_type, limit
query_emb_str, agent_id, fact_type, limit
)
return [RetrievalResult.from_db_row(dict(r)) for r in results]
return [(str(r["id"]), dict(r)) for r in results]
async def retrieve_bm25(
conn,
query_text: str,
bank_id: str,
agent_id: str,
fact_type: str,
limit: int
) -> List[RetrievalResult]:
) -> List[Tuple[str, Dict[str, Any]]]:
"""
BM25 keyword retrieval via full-text search.
Args:
conn: Database connection
query_text: Query text
agent_id: bank ID
agent_id: Agent ID
fact_type: Fact type to filter
limit: Maximum results to return
Returns:
List of RetrievalResult objects
List of (doc_id, data) tuples
"""
import re
@@ -137,30 +90,129 @@ async def retrieve_bm25(
results = await conn.fetch(
"""
SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, access_count, embedding, fact_type, document_id, chunk_id,
SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, access_count, embedding, fact_type, document_id,
ts_rank_cd(search_vector, to_tsquery('english', $1)) AS bm25_score
FROM memory_units
WHERE bank_id = $2
WHERE agent_id = $2
AND fact_type = $3
AND search_vector @@ to_tsquery('english', $1)
ORDER BY bm25_score DESC
LIMIT $4
""",
query_tsquery, bank_id, fact_type, limit
query_tsquery, agent_id, fact_type, limit
)
return [RetrievalResult.from_db_row(dict(r)) for r in results]
return [(str(r["id"]), dict(r)) for r in results]
async def retrieve_graph(
conn,
query_emb_str: str,
agent_id: str,
fact_type: str,
budget: int
) -> List[Tuple[str, Dict[str, Any]]]:
"""
Graph retrieval via spreading activation.
Args:
conn: Database connection
query_emb_str: Query embedding as string
agent_id: Agent ID
fact_type: Fact type to filter
budget: Node budget for graph traversal
Returns:
List of (doc_id, data) tuples
"""
# Find entry points
entry_points = await conn.fetch(
"""
SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, access_count, embedding, fact_type, document_id,
1 - (embedding <=> $1::vector) AS similarity
FROM memory_units
WHERE agent_id = $2
AND embedding IS NOT NULL
AND fact_type = $3
AND (1 - (embedding <=> $1::vector)) >= 0.5
ORDER BY embedding <=> $1::vector
LIMIT 5
""",
query_emb_str, agent_id, fact_type
)
if not entry_points:
return []
# Simple BFS-style spreading activation
visited = set()
results = []
queue = [(dict(r), r["similarity"]) for r in entry_points]
budget_remaining = budget
while queue and budget_remaining > 0:
current, activation = queue.pop(0)
unit_id = str(current["id"])
if unit_id in visited:
continue
visited.add(unit_id)
budget_remaining -= 1
results.append((unit_id, current))
# Get neighbors
if budget_remaining > 0:
neighbors = await conn.fetch(
"""
SELECT mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start, mu.occurred_end, mu.mentioned_at, mu.access_count, mu.embedding, mu.fact_type, mu.document_id,
ml.weight, ml.link_type
FROM memory_links ml
JOIN memory_units mu ON ml.to_unit_id = mu.id
WHERE ml.from_unit_id = $1
AND ml.weight >= 0.1
AND mu.fact_type = $2
ORDER BY ml.weight DESC
LIMIT 10
""",
current["id"], fact_type
)
for n in neighbors:
neighbor_id = str(n["id"])
if neighbor_id not in visited:
# Boost activation for causal links (they're high-value relationships)
link_type = n["link_type"]
base_weight = n["weight"]
# Causal links get 1.5-2.0x boost depending on type
if link_type in ("causes", "caused_by"):
# Direct causation - very strong relationship
causal_boost = 2.0
elif link_type in ("enables", "prevents"):
# Conditional causation - strong but not as direct
causal_boost = 1.5
else:
# Temporal, semantic, entity links - standard weight
causal_boost = 1.0
effective_weight = base_weight * causal_boost
new_activation = activation * effective_weight * 0.8
if new_activation > 0.1:
queue.append((dict(n), new_activation))
return results
async def retrieve_temporal(
conn,
query_emb_str: str,
bank_id: str,
agent_id: str,
fact_type: str,
start_date: datetime,
end_date: datetime,
budget: int,
semantic_threshold: float = 0.1
) -> List[RetrievalResult]:
semantic_threshold: float = 0.4
) -> List[Tuple[str, Dict[str, Any]]]:
"""
Temporal retrieval with spreading activation.
@@ -172,7 +224,7 @@ async def retrieve_temporal(
Args:
conn: Database connection
query_emb_str: Query embedding as string
agent_id: bank ID
agent_id: Agent ID
fact_type: Fact type to filter
start_date: Start of time range
end_date: End of time range
@@ -180,7 +232,7 @@ async def retrieve_temporal(
semantic_threshold: Minimum semantic similarity to include
Returns:
List of RetrievalResult objects with temporal scores
List of (doc_id, data) tuples with temporal_score
"""
from datetime import timezone
@@ -192,10 +244,10 @@ async def retrieve_temporal(
entry_points = await conn.fetch(
"""
SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, access_count, embedding, fact_type, document_id, chunk_id,
SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, access_count, embedding, fact_type, document_id,
1 - (embedding <=> $1::vector) AS similarity
FROM memory_units
WHERE bank_id = $2
WHERE agent_id = $2
AND fact_type = $3
AND embedding IS NOT NULL
AND (
@@ -215,10 +267,17 @@ async def retrieve_temporal(
ORDER BY COALESCE(occurred_start, mentioned_at, occurred_end) DESC, (embedding <=> $1::vector) ASC
LIMIT 10
""",
query_emb_str, bank_id, fact_type, start_date, end_date, semantic_threshold
query_emb_str, agent_id, fact_type, start_date, end_date, semantic_threshold
)
if not entry_points:
# Check if there are ANY memories with temporal metadata for this agent
total_with_dates = await conn.fetchval(
"""SELECT COUNT(*) FROM memory_units
WHERE agent_id = $1 AND fact_type = $2
AND (occurred_start IS NOT NULL OR occurred_end IS NOT NULL OR mentioned_at IS NOT NULL)""",
agent_id, fact_type
)
return []
# Calculate temporal scores for entry points
@@ -251,25 +310,24 @@ async def retrieve_temporal(
else:
temporal_proximity = 0.5 # Fallback if no dates (shouldn't happen due to WHERE clause)
# Create RetrievalResult with temporal scores
ep_result = RetrievalResult.from_db_row(dict(ep))
ep_result.temporal_score = temporal_proximity
ep_result.temporal_proximity = temporal_proximity
results.append(ep_result)
data = dict(ep)
data["temporal_score"] = temporal_proximity
data["temporal_proximity"] = temporal_proximity
results.append((unit_id, data))
# Spread through temporal links
queue = [(RetrievalResult.from_db_row(dict(ep)), ep["similarity"], 1.0) for ep in entry_points] # (unit, semantic_sim, temporal_score)
queue = [(dict(ep), ep["similarity"], 1.0) for ep in entry_points] # (unit, semantic_sim, temporal_score)
budget_remaining = budget - len(entry_points)
while queue and budget_remaining > 0:
current, semantic_sim, temporal_score = queue.pop(0)
current_id = current.id
current_id = str(current["id"])
# Get neighbors via temporal and causal links
if budget_remaining > 0:
neighbors = await conn.fetch(
"""
SELECT mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start, mu.occurred_end, mu.mentioned_at, mu.access_count, mu.embedding, mu.fact_type, mu.document_id, mu.chunk_id,
SELECT mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start, mu.occurred_end, mu.mentioned_at, mu.access_count, mu.embedding, mu.fact_type, mu.document_id,
ml.weight, ml.link_type,
1 - (mu.embedding <=> $1::vector) AS similarity
FROM memory_links ml
@@ -283,7 +341,7 @@ async def retrieve_temporal(
ORDER BY ml.weight DESC
LIMIT 10
""",
query_emb_str, current.id, fact_type, semantic_threshold
query_emb_str, current["id"], fact_type, semantic_threshold
)
for n in neighbors:
@@ -326,15 +384,14 @@ async def retrieve_temporal(
# Combined temporal score
combined_temporal = max(neighbor_temporal_proximity, propagated_temporal)
# Create RetrievalResult with temporal scores
neighbor_result = RetrievalResult.from_db_row(dict(n))
neighbor_result.temporal_score = combined_temporal
neighbor_result.temporal_proximity = neighbor_temporal_proximity
results.append(neighbor_result)
neighbor_data = dict(n)
neighbor_data["temporal_score"] = combined_temporal
neighbor_data["temporal_proximity"] = neighbor_temporal_proximity
results.append((neighbor_id, neighbor_data))
# Add to queue for further spreading
if budget_remaining > 0 and combined_temporal > 0.2:
queue.append((neighbor_result, n["similarity"], combined_temporal))
queue.append((dict(n), n["similarity"], combined_temporal))
if budget_remaining <= 0:
break
@@ -346,13 +403,12 @@ async def retrieve_parallel(
pool,
query_text: str,
query_embedding_str: str,
bank_id: str,
agent_id: str,
fact_type: str,
thinking_budget: int,
question_date: Optional[datetime] = None,
query_analyzer: Optional["QueryAnalyzer"] = None,
graph_retriever: Optional[GraphRetriever] = None,
) -> ParallelRetrievalResult:
query_analyzer: Optional["QueryAnalyzer"] = None
) -> Tuple[List, List, List, Optional[List]]:
"""
Run 3-way or 4-way parallel retrieval (adds temporal if detected).
@@ -360,318 +416,54 @@ async def retrieve_parallel(
pool: Database connection pool
query_text: Query text
query_embedding_str: Query embedding as string
bank_id: Bank ID
agent_id: Agent ID
fact_type: Fact type to filter
thinking_budget: Budget for graph traversal and retrieval limits
question_date: Optional date when question was asked (for temporal filtering)
query_analyzer: Query analyzer to use (defaults to TransformerQueryAnalyzer)
graph_retriever: Graph retrieval strategy (defaults to configured retriever)
Returns:
ParallelRetrievalResult with semantic, bm25, graph, temporal results and timings
Tuple of (semantic_results, bm25_results, graph_results, temporal_results)
temporal_results is None if no temporal constraint detected
"""
# Detect temporal constraint
from .temporal_extraction import extract_temporal_constraint
import logging
logger = logging.getLogger(__name__)
temporal_constraint = extract_temporal_constraint(
query_text, reference_date=question_date, analyzer=query_analyzer
)
retriever = graph_retriever or get_default_graph_retriever()
if retriever.name == "mpfp":
return await _retrieve_parallel_mpfp(
pool, query_text, query_embedding_str, bank_id, fact_type,
thinking_budget, temporal_constraint, retriever
)
else:
return await _retrieve_parallel_bfs(
pool, query_text, query_embedding_str, bank_id, fact_type,
thinking_budget, temporal_constraint, retriever
)
@dataclass
class _SemanticGraphResult:
"""Internal result from semantic→graph chain."""
semantic: List[RetrievalResult]
graph: List[RetrievalResult]
semantic_time: float
graph_time: float
@dataclass
class _TimedResult:
"""Internal result with timing."""
results: List[RetrievalResult]
time: float
async def _retrieve_parallel_mpfp(
pool,
query_text: str,
query_embedding_str: str,
bank_id: str,
fact_type: str,
thinking_budget: int,
temporal_constraint: Optional[tuple],
retriever: GraphRetriever,
) -> ParallelRetrievalResult:
"""
MPFP retrieval with optimized parallelization.
Runs 2-3 parallel task chains:
- Task 1: Semantic → Graph (chained, graph uses semantic seeds)
- Task 2: BM25 (independent)
- Task 3: Temporal (if constraint detected)
"""
import time
async def run_semantic_then_graph() -> _SemanticGraphResult:
"""Chain: semantic retrieval → graph retrieval (using semantic as seeds)."""
start = time.time()
async def run_semantic():
async with acquire_with_retry(pool) as conn:
semantic = await retrieve_semantic(
conn, query_embedding_str, bank_id, fact_type, limit=thinking_budget
return await retrieve_semantic(conn, query_embedding_str, agent_id, fact_type, limit=thinking_budget)
async def run_bm25():
async with acquire_with_retry(pool) as conn:
return await retrieve_bm25(conn, query_text, agent_id, fact_type, limit=thinking_budget)
async def run_graph():
async with acquire_with_retry(pool) as conn:
return await retrieve_graph(conn, query_embedding_str, agent_id, fact_type, budget=thinking_budget)
async def run_temporal(start_date, end_date):
async with acquire_with_retry(pool) as conn:
return await retrieve_temporal(
conn, query_embedding_str, agent_id, fact_type,
start_date, end_date, budget=thinking_budget, semantic_threshold=0.4
)
semantic_time = time.time() - start
# Get temporal seeds if needed (quick query, part of this chain)
temporal_seeds = None
if temporal_constraint:
tc_start, tc_end = temporal_constraint
async with acquire_with_retry(pool) as conn:
temporal_seeds = await _get_temporal_entry_points(
conn, query_embedding_str, bank_id, fact_type,
tc_start, tc_end, limit=20
)
# Run graph with seeds
start = time.time()
graph = await retriever.retrieve(
pool=pool,
query_embedding_str=query_embedding_str,
bank_id=bank_id,
fact_type=fact_type,
budget=thinking_budget,
query_text=query_text,
semantic_seeds=semantic,
temporal_seeds=temporal_seeds,
)
graph_time = time.time() - start
return _SemanticGraphResult(semantic, graph, semantic_time, graph_time)
async def run_bm25() -> _TimedResult:
"""Independent BM25 retrieval."""
start = time.time()
async with acquire_with_retry(pool) as conn:
results = await retrieve_bm25(conn, query_text, bank_id, fact_type, limit=thinking_budget)
return _TimedResult(results, time.time() - start)
async def run_temporal(tc_start, tc_end) -> _TimedResult:
"""Temporal retrieval (uses its own entry point finding)."""
start = time.time()
async with acquire_with_retry(pool) as conn:
results = await retrieve_temporal(
conn, query_embedding_str, bank_id, fact_type,
tc_start, tc_end, budget=thinking_budget, semantic_threshold=0.1
)
return _TimedResult(results, time.time() - start)
# Run parallel task chains
# Run retrievals in parallel
if temporal_constraint:
tc_start, tc_end = temporal_constraint
sg_result, bm25_result, temporal_result = await asyncio.gather(
run_semantic_then_graph(),
run_bm25(),
run_temporal(tc_start, tc_end),
)
return ParallelRetrievalResult(
semantic=sg_result.semantic,
bm25=bm25_result.results,
graph=sg_result.graph,
temporal=temporal_result.results,
timings={
"semantic": sg_result.semantic_time,
"graph": sg_result.graph_time,
"bm25": bm25_result.time,
"temporal": temporal_result.time,
},
temporal_constraint=temporal_constraint,
start_date, end_date = temporal_constraint
semantic_results, bm25_results, graph_results, temporal_results = await asyncio.gather(
run_semantic(), run_bm25(), run_graph(), run_temporal(start_date, end_date)
)
else:
sg_result, bm25_result = await asyncio.gather(
run_semantic_then_graph(),
run_bm25(),
)
return ParallelRetrievalResult(
semantic=sg_result.semantic,
bm25=bm25_result.results,
graph=sg_result.graph,
temporal=None,
timings={
"semantic": sg_result.semantic_time,
"graph": sg_result.graph_time,
"bm25": bm25_result.time,
},
temporal_constraint=None,
semantic_results, bm25_results, graph_results = await asyncio.gather(
run_semantic(), run_bm25(), run_graph()
)
temporal_results = None
async def _get_temporal_entry_points(
conn,
query_embedding_str: str,
bank_id: str,
fact_type: str,
start_date: datetime,
end_date: datetime,
limit: int = 20,
semantic_threshold: float = 0.1,
) -> List[RetrievalResult]:
"""Get temporal entry points (facts in date range with semantic relevance)."""
from datetime import timezone
if start_date.tzinfo is None:
start_date = start_date.replace(tzinfo=timezone.utc)
if end_date.tzinfo is None:
end_date = end_date.replace(tzinfo=timezone.utc)
rows = await conn.fetch(
"""
SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at,
access_count, embedding, fact_type, document_id, chunk_id,
1 - (embedding <=> $1::vector) AS similarity
FROM memory_units
WHERE bank_id = $2
AND fact_type = $3
AND embedding IS NOT NULL
AND (
(occurred_start IS NOT NULL AND occurred_end IS NOT NULL
AND occurred_start <= $5 AND occurred_end >= $4)
OR (mentioned_at IS NOT NULL AND mentioned_at BETWEEN $4 AND $5)
OR (occurred_start IS NOT NULL AND occurred_start BETWEEN $4 AND $5)
OR (occurred_end IS NOT NULL AND occurred_end BETWEEN $4 AND $5)
)
AND (1 - (embedding <=> $1::vector)) >= $6
ORDER BY COALESCE(occurred_start, mentioned_at, occurred_end) DESC,
(embedding <=> $1::vector) ASC
LIMIT $7
""",
query_embedding_str, bank_id, fact_type, start_date, end_date, semantic_threshold, limit
)
results = []
total_days = max((end_date - start_date).total_seconds() / 86400, 1)
mid_date = start_date + (end_date - start_date) / 2
for row in rows:
result = RetrievalResult.from_db_row(dict(row))
# Calculate temporal proximity score
best_date = None
if row["occurred_start"] and row["occurred_end"]:
best_date = row["occurred_start"] + (row["occurred_end"] - row["occurred_start"]) / 2
elif row["occurred_start"]:
best_date = row["occurred_start"]
elif row["occurred_end"]:
best_date = row["occurred_end"]
elif row["mentioned_at"]:
best_date = row["mentioned_at"]
if best_date:
days_from_mid = abs((best_date - mid_date).total_seconds() / 86400)
result.temporal_proximity = 1.0 - min(days_from_mid / (total_days / 2), 1.0)
else:
result.temporal_proximity = 0.5
result.temporal_score = result.temporal_proximity
results.append(result)
return results
async def _retrieve_parallel_bfs(
pool,
query_text: str,
query_embedding_str: str,
bank_id: str,
fact_type: str,
thinking_budget: int,
temporal_constraint: Optional[tuple],
retriever: GraphRetriever,
) -> ParallelRetrievalResult:
"""BFS retrieval: all methods run in parallel (original behavior)."""
import time
async def run_semantic() -> _TimedResult:
start = time.time()
async with acquire_with_retry(pool) as conn:
results = await retrieve_semantic(conn, query_embedding_str, bank_id, fact_type, limit=thinking_budget)
return _TimedResult(results, time.time() - start)
async def run_bm25() -> _TimedResult:
start = time.time()
async with acquire_with_retry(pool) as conn:
results = await retrieve_bm25(conn, query_text, bank_id, fact_type, limit=thinking_budget)
return _TimedResult(results, time.time() - start)
async def run_graph() -> _TimedResult:
start = time.time()
results = await retriever.retrieve(
pool=pool,
query_embedding_str=query_embedding_str,
bank_id=bank_id,
fact_type=fact_type,
budget=thinking_budget,
query_text=query_text,
)
return _TimedResult(results, time.time() - start)
async def run_temporal(tc_start, tc_end) -> _TimedResult:
start = time.time()
async with acquire_with_retry(pool) as conn:
results = await retrieve_temporal(
conn, query_embedding_str, bank_id, fact_type,
tc_start, tc_end, budget=thinking_budget, semantic_threshold=0.1
)
return _TimedResult(results, time.time() - start)
if temporal_constraint:
tc_start, tc_end = temporal_constraint
semantic_r, bm25_r, graph_r, temporal_r = await asyncio.gather(
run_semantic(),
run_bm25(),
run_graph(),
run_temporal(tc_start, tc_end),
)
return ParallelRetrievalResult(
semantic=semantic_r.results,
bm25=bm25_r.results,
graph=graph_r.results,
temporal=temporal_r.results,
timings={
"semantic": semantic_r.time,
"bm25": bm25_r.time,
"graph": graph_r.time,
"temporal": temporal_r.time,
},
temporal_constraint=temporal_constraint,
)
else:
semantic_r, bm25_r, graph_r = await asyncio.gather(
run_semantic(),
run_bm25(),
run_graph(),
)
return ParallelRetrievalResult(
semantic=semantic_r.results,
bm25=bm25_r.results,
graph=graph_r.results,
temporal=None,
timings={
"semantic": semantic_r.time,
"bm25": bm25_r.time,
"graph": graph_r.time,
},
temporal_constraint=None,
)
return semantic_results, bm25_results, graph_results, temporal_results
@@ -1,161 +0,0 @@
"""
Scoring functions for memory search and retrieval.
Includes recency weighting, frequency weighting, temporal proximity,
and similarity calculations used in memory activation and ranking.
"""
from datetime import datetime
from typing import List
def cosine_similarity(vec1: List[float], vec2: List[float]) -> float:
"""
Calculate cosine similarity between two vectors.
Args:
vec1: First vector
vec2: Second vector
Returns:
Similarity score between 0 and 1
"""
if len(vec1) != len(vec2):
raise ValueError("Vectors must have same dimension")
dot_product = sum(a * b for a, b in zip(vec1, vec2))
magnitude1 = sum(a * a for a in vec1) ** 0.5
magnitude2 = sum(b * b for b in vec2) ** 0.5
if magnitude1 == 0 or magnitude2 == 0:
return 0.0
return dot_product / (magnitude1 * magnitude2)
def calculate_recency_weight(days_since: float, half_life_days: float = 365.0) -> float:
"""
Calculate recency weight using logarithmic decay.
This provides much better differentiation over long time periods compared to
exponential decay. Uses a log-based decay where the half-life parameter controls
when memories reach 50% weight.
Examples:
- Today (0 days): 1.0
- 1 year (365 days): ~0.5 (with default half_life=365)
- 2 years (730 days): ~0.33
- 5 years (1825 days): ~0.17
- 10 years (3650 days): ~0.09
This ensures that 2-year-old and 5-year-old memories have meaningfully
different weights, unlike exponential decay which makes them both ~0.
Args:
days_since: Number of days since the memory was created
half_life_days: Number of days for weight to reach 0.5 (default: 1 year)
Returns:
Weight between 0 and 1
"""
import math
# Logarithmic decay: 1 / (1 + log(1 + days_since/half_life))
# This decays much slower than exponential, giving better long-term differentiation
normalized_age = days_since / half_life_days
return 1.0 / (1.0 + math.log1p(normalized_age))
def calculate_frequency_weight(access_count: int, max_boost: float = 2.0) -> float:
"""
Calculate frequency weight based on access count.
Frequently accessed memories are weighted higher.
Uses logarithmic scaling to avoid over-weighting.
Args:
access_count: Number of times the memory was accessed
max_boost: Maximum multiplier for frequently accessed memories
Returns:
Weight between 1.0 and max_boost
"""
import math
if access_count <= 0:
return 1.0
# Logarithmic scaling: log(access_count + 1) / log(10)
# This gives: 0 accesses = 1.0, 9 accesses ~= 1.5, 99 accesses ~= 2.0
normalized = math.log(access_count + 1) / math.log(10)
return 1.0 + min(normalized, max_boost - 1.0)
def calculate_temporal_anchor(occurred_start: datetime, occurred_end: datetime) -> datetime:
"""
Calculate a single temporal anchor point from a temporal range.
Used for spreading activation - we need a single representative date
to calculate temporal proximity between facts. This simplifies the
range-to-range distance problem.
Strategy: Use midpoint of the range for balanced representation.
Args:
occurred_start: Start of temporal range
occurred_end: End of temporal range
Returns:
Single datetime representing the temporal anchor (midpoint)
Examples:
- Point event (July 14): start=July 14, end=July 14 → anchor=July 14
- Month range (February): start=Feb 1, end=Feb 28 → anchor=Feb 14
- Year range (2023): start=Jan 1, end=Dec 31 → anchor=July 1
"""
# Calculate midpoint
time_delta = occurred_end - occurred_start
midpoint = occurred_start + (time_delta / 2)
return midpoint
def calculate_temporal_proximity(
anchor_a: datetime,
anchor_b: datetime,
half_life_days: float = 30.0
) -> float:
"""
Calculate temporal proximity between two temporal anchors.
Used for spreading activation to determine how "close" two facts are
in time. Uses logarithmic decay so that temporal similarity doesn't
drop off too quickly.
Args:
anchor_a: Temporal anchor of first fact
anchor_b: Temporal anchor of second fact
half_life_days: Number of days for proximity to reach 0.5
(default: 30 days = 1 month)
Returns:
Proximity score in [0, 1] where:
- 1.0 = same day
- 0.5 = ~half_life days apart
- 0.0 = very distant in time
Examples:
- Same day: 1.0
- 1 week apart (half_life=30): ~0.7
- 1 month apart (half_life=30): ~0.5
- 1 year apart (half_life=30): ~0.2
"""
import math
days_apart = abs((anchor_a - anchor_b).days)
if days_apart == 0:
return 1.0
# Logarithmic decay: 1 / (1 + log(1 + days_apart/half_life))
# Similar to calculate_recency_weight but for proximity between events
normalized_distance = days_apart / half_life_days
proximity = 1.0 / (1.0 + math.log1p(normalized_distance))
return proximity
@@ -7,7 +7,7 @@ Handles natural language temporal expressions using transformer-based query anal
from typing import Optional, Tuple
from datetime import datetime
import logging
from hindsight_api.engine.query_analyzer import QueryAnalyzer, DateparserQueryAnalyzer
from hindsight_api.engine.query_analyzer import QueryAnalyzer, TransformerQueryAnalyzer
logger = logging.getLogger(__name__)
@@ -20,14 +20,14 @@ def get_default_analyzer() -> QueryAnalyzer:
"""
Get or create the default query analyzer.
Uses lazy initialization to avoid loading at import time.
Uses lazy initialization to avoid loading model at import time.
Returns:
Default DateparserQueryAnalyzer instance
Default TransformerQueryAnalyzer instance
"""
global _default_analyzer
if _default_analyzer is None:
_default_analyzer = DateparserQueryAnalyzer()
_default_analyzer = TransformerQueryAnalyzer()
return _default_analyzer
@@ -37,14 +37,14 @@ def extract_temporal_constraint(
analyzer: Optional[QueryAnalyzer] = None,
) -> Optional[Tuple[datetime, datetime]]:
"""
Extract temporal constraint from query.
Extract temporal constraint from query using transformer-based analysis.
Returns (start_date, end_date) tuple if temporal constraint found, else None.
Args:
query: Search query
reference_date: Reference date for relative terms (defaults to now)
analyzer: Custom query analyzer (defaults to DateparserQueryAnalyzer)
analyzer: Custom query analyzer (defaults to TransformerQueryAnalyzer)
Returns:
(start_date, end_date) tuple or None
@@ -1,163 +0,0 @@
"""
Type definitions for the recall pipeline.
These dataclasses replace Dict[str, Any] types throughout the recall pipeline,
providing type safety and making data flow explicit.
"""
from dataclasses import dataclass, field
from typing import Optional, List, Dict, Any
from datetime import datetime
@dataclass
class RetrievalResult:
"""
Result from a single retrieval method (semantic, BM25, graph, or temporal).
This represents a raw result from the database query, before merging or reranking.
"""
id: str
text: str
fact_type: str
context: Optional[str] = None
event_date: Optional[datetime] = None
occurred_start: Optional[datetime] = None
occurred_end: Optional[datetime] = None
mentioned_at: Optional[datetime] = None
document_id: Optional[str] = None
chunk_id: Optional[str] = None
access_count: int = 0
embedding: Optional[List[float]] = None
# Retrieval-specific scores (only one will be set depending on retrieval method)
similarity: Optional[float] = None # Semantic retrieval
bm25_score: Optional[float] = None # BM25 retrieval
activation: Optional[float] = None # Graph retrieval (spreading activation)
temporal_score: Optional[float] = None # Temporal retrieval
temporal_proximity: Optional[float] = None # Temporal retrieval
@classmethod
def from_db_row(cls, row: Dict[str, Any]) -> "RetrievalResult":
"""Create from a database row (asyncpg Record converted to dict)."""
return cls(
id=str(row["id"]),
text=row["text"],
fact_type=row["fact_type"],
context=row.get("context"),
event_date=row.get("event_date"),
occurred_start=row.get("occurred_start"),
occurred_end=row.get("occurred_end"),
mentioned_at=row.get("mentioned_at"),
document_id=row.get("document_id"),
chunk_id=row.get("chunk_id"),
access_count=row.get("access_count", 0),
embedding=row.get("embedding"),
similarity=row.get("similarity"),
bm25_score=row.get("bm25_score"),
activation=row.get("activation"),
temporal_score=row.get("temporal_score"),
temporal_proximity=row.get("temporal_proximity"),
)
@dataclass
class MergedCandidate:
"""
Candidate after RRF merge of multiple retrieval results.
Contains the original retrieval data plus RRF metadata.
"""
# Original retrieval data
retrieval: RetrievalResult
# RRF metadata
rrf_score: float
rrf_rank: int = 0
source_ranks: Dict[str, int] = field(default_factory=dict) # method_name -> rank
@property
def id(self) -> str:
"""Convenience property to access ID."""
return self.retrieval.id
@dataclass
class ScoredResult:
"""
Result after reranking and scoring.
Contains all retrieval/merge data plus reranking scores and combined score.
"""
# Original merged candidate
candidate: MergedCandidate
# Reranking scores
cross_encoder_score: float = 0.0
cross_encoder_score_normalized: float = 0.0
# Normalized component scores
rrf_normalized: float = 0.0
recency: float = 0.5
temporal: float = 0.5
# Final combined score
combined_score: float = 0.0
weight: float = 0.0 # Final weight used for ranking
@property
def id(self) -> str:
"""Convenience property to access ID."""
return self.candidate.id
@property
def retrieval(self) -> RetrievalResult:
"""Convenience property to access retrieval data."""
return self.candidate.retrieval
def to_dict(self) -> Dict[str, Any]:
"""
Convert to dict for backwards compatibility.
This is used during the transition period and for serialization.
"""
# Start with retrieval data
result = {
"id": self.retrieval.id,
"text": self.retrieval.text,
"fact_type": self.retrieval.fact_type,
"context": self.retrieval.context,
"event_date": self.retrieval.event_date,
"occurred_start": self.retrieval.occurred_start,
"occurred_end": self.retrieval.occurred_end,
"mentioned_at": self.retrieval.mentioned_at,
"document_id": self.retrieval.document_id,
"chunk_id": self.retrieval.chunk_id,
"access_count": self.retrieval.access_count,
"embedding": self.retrieval.embedding,
"semantic_similarity": self.retrieval.similarity,
"bm25_score": self.retrieval.bm25_score,
}
# Add temporal scores if present
if self.retrieval.temporal_score is not None:
result["temporal_score"] = self.retrieval.temporal_score
if self.retrieval.temporal_proximity is not None:
result["temporal_proximity"] = self.retrieval.temporal_proximity
# Add RRF metadata
result["rrf_score"] = self.candidate.rrf_score
result["rrf_rank"] = self.candidate.rrf_rank
result.update(self.candidate.source_ranks)
# Add reranking scores
result["cross_encoder_score"] = self.cross_encoder_score
result["cross_encoder_score_normalized"] = self.cross_encoder_score_normalized
result["rrf_normalized"] = self.rrf_normalized
result["temporal"] = self.temporal
result["recency"] = self.recency
result["combined_score"] = self.combined_score
result["weight"] = self.weight
result["activation"] = self.weight # Legacy field
return result
@@ -4,81 +4,70 @@ Helper functions for hybrid search (semantic + BM25 + graph).
from typing import List, Dict, Any, Tuple
import asyncio
from .types import RetrievalResult, MergedCandidate
def reciprocal_rank_fusion(
result_lists: List[List[RetrievalResult]],
result_lists: List[List[Tuple[str, Dict[str, Any]]]],
k: int = 60
) -> List[MergedCandidate]:
) -> List[Tuple[str, Dict[str, Any], Dict[str, float]]]:
"""
Merge multiple ranked result lists using Reciprocal Rank Fusion.
RRF formula: score(d) = sum_over_lists(1 / (k + rank(d)))
Args:
result_lists: List of result lists, each containing RetrievalResult objects
result_lists: List of result lists, each containing (id, data) tuples
k: Constant for RRF formula (default: 60)
Returns:
Merged list of MergedCandidate objects, sorted by RRF score
Merged list of (id, data, scores_dict) tuples, sorted by RRF score
Example:
semantic_results = [RetrievalResult(...), RetrievalResult(...), ...]
bm25_results = [RetrievalResult(...), RetrievalResult(...), ...]
graph_results = [RetrievalResult(...), RetrievalResult(...), ...]
semantic_results = [("id1", {...}), ("id2", {...}), ...]
bm25_results = [("id2", {...}), ("id3", {...}), ...]
graph_results = [("id1", {...}), ("id4", {...}), ...]
merged = reciprocal_rank_fusion([semantic_results, bm25_results, graph_results])
# Returns: [MergedCandidate(...), MergedCandidate(...), ...]
# Returns: [("id2", {...}, {"rrf": 0.05, "semantic_rank": 2, ...}), ...]
"""
# Track scores from each list
rrf_scores = {}
source_ranks = {} # Track rank from each source for each doc_id
all_retrievals = {} # Store the actual RetrievalResult (use first occurrence)
source_ranks = {} # Track rank from each source
source_scores = {} # Track original score from each source
all_data = {} # Store the actual data
source_names = ["semantic", "bm25", "graph", "temporal"]
source_names = ["semantic", "bm25", "graph"]
for source_idx, results in enumerate(result_lists):
source_name = source_names[source_idx] if source_idx < len(source_names) else f"source_{source_idx}"
for rank, retrieval in enumerate(results, start=1):
# Type check to catch tuple issues
if isinstance(retrieval, tuple):
raise TypeError(
f"Expected RetrievalResult but got tuple in {source_name} results at rank {rank}. "
f"Tuple value: {retrieval[:2] if len(retrieval) >= 2 else retrieval}. "
f"This suggests the retrieval function returned tuples instead of RetrievalResult objects."
)
if not isinstance(retrieval, RetrievalResult):
raise TypeError(
f"Expected RetrievalResult but got {type(retrieval).__name__} in {source_name} results at rank {rank}"
)
doc_id = retrieval.id
# Store retrieval result (use first occurrence)
if doc_id not in all_retrievals:
all_retrievals[doc_id] = retrieval
for rank, (doc_id, data) in enumerate(results, start=1):
# Store data (use first occurrence)
if doc_id not in all_data:
all_data[doc_id] = data
# Calculate RRF score contribution
if doc_id not in rrf_scores:
rrf_scores[doc_id] = 0.0
source_ranks[doc_id] = {}
source_scores[doc_id] = {}
rrf_scores[doc_id] += 1.0 / (k + rank)
source_ranks[doc_id][f"{source_name}_rank"] = rank
# Store original score if available
if "score" in data:
source_scores[doc_id][f"{source_name}_score"] = data["score"]
# Combine into final results with metadata
merged_results = []
for rrf_rank, (doc_id, rrf_score) in enumerate(
sorted(rrf_scores.items(), key=lambda x: x[1], reverse=True), start=1
):
merged_candidate = MergedCandidate(
retrieval=all_retrievals[doc_id],
rrf_score=rrf_score,
rrf_rank=rrf_rank,
source_ranks=source_ranks[doc_id]
)
merged_results.append(merged_candidate)
for doc_id, rrf_score in sorted(rrf_scores.items(), key=lambda x: x[1], reverse=True):
scores_dict = {
"rrf_score": rrf_score,
**source_ranks[doc_id],
**source_scores[doc_id]
}
merged_results.append((doc_id, all_data[doc_id], scores_dict))
return merged_results
@@ -14,7 +14,7 @@ class QueryInfo(BaseModel):
query_text: str = Field(description="Original query text")
query_embedding: List[float] = Field(description="Generated query embedding vector")
timestamp: datetime = Field(description="When the query was executed")
budget: int = Field(description="Maximum nodes to explore")
thinking_budget: int = Field(description="Maximum nodes to explore")
max_tokens: int = Field(description="Maximum tokens to return in results")
@@ -59,7 +59,7 @@ class NodeVisit(BaseModel):
node_id: str = Field(description="Memory unit ID")
text: str = Field(description="Memory unit text content")
context: str = Field(description="Memory unit context")
event_date: Optional[datetime] = Field(default=None, description="When the memory occurred")
event_date: datetime = Field(description="When the memory occurred")
access_count: int = Field(description="Number of times accessed before this search")
# How this node was reached
@@ -100,7 +100,6 @@ class RetrievalResult(BaseModel):
text: str = Field(description="Memory unit text content")
context: str = Field(default="", description="Memory unit context")
event_date: Optional[datetime] = Field(default=None, description="When the memory occurred")
fact_type: Optional[str] = Field(default=None, description="Fact type (world, experience, opinion)")
score: float = Field(description="Score from this retrieval method")
score_name: str = Field(description="Name of the score (e.g., 'similarity', 'bm25_score', 'activation')")
@@ -108,7 +107,6 @@ class RetrievalResult(BaseModel):
class RetrievalMethodResults(BaseModel):
"""Results from a single retrieval method."""
method_name: Literal["semantic", "bm25", "graph", "temporal"] = Field(description="Name of retrieval method")
fact_type: Optional[str] = Field(default=None, description="Fact type this retrieval was for (world, experience, opinion)")
results: List[RetrievalResult] = Field(description="Retrieved results with ranks")
duration_seconds: float = Field(description="Time taken for this retrieval")
metadata: Dict[str, Any] = Field(default_factory=dict, description="Method-specific metadata")
@@ -8,7 +8,7 @@ import time
from datetime import datetime, timezone
from typing import List, Optional, Dict, Any, Literal
from .trace import (
from .search_trace import (
SearchTrace,
QueryInfo,
EntryPoint,
@@ -30,7 +30,7 @@ class SearchTracer:
Tracer for collecting detailed search execution information.
Usage:
tracer = SearchTracer(query="Who is Alice?", budget=50, max_tokens=4096)
tracer = SearchTracer(query="Who is Alice?", thinking_budget=50, top_k=10)
tracer.start()
# During search...
@@ -44,17 +44,17 @@ class SearchTracer:
json_output = trace.to_json()
"""
def __init__(self, query: str, budget: int, max_tokens: int):
def __init__(self, query: str, thinking_budget: int, max_tokens: int):
"""
Initialize tracer.
Args:
query: Search query text
budget: Maximum nodes to explore
thinking_budget: Maximum nodes to explore
max_tokens: Maximum tokens to return in results
"""
self.query_text = query
self.budget = budget
self.thinking_budget = thinking_budget
self.max_tokens = max_tokens
# Trace data
@@ -97,9 +97,6 @@ class SearchTracer:
similarity: Cosine similarity to query
rank: Rank among entry points (1-based)
"""
# Clamp similarity to [0.0, 1.0] to handle floating-point precision
similarity = min(1.0, max(0.0, similarity))
self.entry_points.append(
EntryPoint(
node_id=node_id,
@@ -148,12 +145,6 @@ class SearchTracer:
self.current_step += 1
self.nodes_visited_set.add(node_id)
# Clamp values to handle floating-point precision issues
# (sometimes normalization produces values like 1.0000005 instead of 1.0)
semantic_similarity = min(1.0, max(0.0, semantic_similarity))
recency = min(1.0, max(0.0, recency))
frequency = min(1.0, max(0.0, frequency))
# Calculate weight contributions for transparency
weights = WeightComponents(
activation=activation,
@@ -289,8 +280,7 @@ class SearchTracer:
results: List[tuple], # List of (doc_id, data) tuples
duration_seconds: float,
score_field: str, # e.g., "similarity", "bm25_score"
metadata: Optional[Dict[str, Any]] = None,
fact_type: Optional[str] = None
metadata: Optional[Dict[str, Any]] = None
):
"""
Record results from a single retrieval method.
@@ -301,13 +291,10 @@ class SearchTracer:
duration_seconds: Time taken for this retrieval
score_field: Field name containing the score in data dict
metadata: Optional metadata about this retrieval method
fact_type: Fact type this retrieval was for (world, experience, opinion)
"""
retrieval_results = []
for rank, (doc_id, data) in enumerate(results, start=1):
score = data.get(score_field)
if score is None:
score = 0.0
score = data.get(score_field, 0.0)
retrieval_results.append(
RetrievalResult(
rank=rank,
@@ -315,7 +302,6 @@ class SearchTracer:
text=data.get("text", ""),
context=data.get("context", ""),
event_date=data.get("event_date"),
fact_type=data.get("fact_type") or fact_type,
score=score,
score_name=score_field,
)
@@ -324,7 +310,6 @@ class SearchTracer:
self.retrieval_results.append(
RetrievalMethodResults(
method_name=method_name,
fact_type=fact_type,
results=retrieval_results,
duration_seconds=duration_seconds,
metadata=metadata or {},
@@ -369,12 +354,10 @@ class SearchTracer:
rrf_rank = rrf_rank_map.get(node_id, len(rrf_merged) + 1)
rank_change = rrf_rank - rank # Positive = moved up
# Extract score components (only include non-None values)
# Keys from ScoredResult.to_dict(): cross_encoder_score, cross_encoder_score_normalized,
# rrf_normalized, temporal, recency, combined_score, weight
# Extract score components
score_components = {}
for key in ["cross_encoder_score", "cross_encoder_score_normalized", "rrf_score", "rrf_normalized", "temporal", "recency", "combined_score"]:
if key in result and result[key] is not None:
for key in ["semantic_similarity", "bm25_score", "rrf_score", "recency_normalized", "frequency_normalized"]:
if key in result:
score_components[key] = result[key]
self.reranked.append(
@@ -417,7 +400,7 @@ class SearchTracer:
query_text=self.query_text,
query_embedding=self.query_embedding or [],
timestamp=datetime.now(timezone.utc),
budget=self.budget,
thinking_budget=self.thinking_budget,
max_tokens=self.max_tokens,
)
@@ -427,7 +410,7 @@ class SearchTracer:
total_nodes_pruned=len(self.pruned),
entry_points_found=len(self.entry_points),
budget_used=len(self.visits),
budget_remaining=self.budget - len(self.visits),
budget_remaining=self.thinking_budget - len(self.visits),
total_duration_seconds=total_duration,
results_returned=len(final_results),
temporal_links_followed=self.temporal_links_followed,
@@ -137,6 +137,7 @@ class AsyncIOQueueBackend(TaskBackend):
await self._queue.put(task_dict)
task_type = task_dict.get('type', 'unknown')
task_id = task_dict.get('id')
logger.debug(f"Task submitted: {task_type} (id: {task_id})")
async def wait_for_pending_tasks(self, timeout: float = 5.0):
"""
@@ -179,7 +180,7 @@ class AsyncIOQueueBackend(TaskBackend):
try:
await self._worker_task
except asyncio.CancelledError:
pass # Worker cancelled successfully
logger.debug("Worker task cancelled successfully")
self._initialized = False
logger.info("AsyncIOQueueBackend shutdown complete")
@@ -210,6 +211,7 @@ class AsyncIOQueueBackend(TaskBackend):
# Process batch
if tasks:
logger.debug(f"Processing batch of {len(tasks)} tasks")
# Execute tasks concurrently
await asyncio.gather(
*[self._execute_task(task_dict) for task_dict in tasks],

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