Compare commits
6
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
3d66988c3b | ||
|
|
67dc160ba3 | ||
|
|
8291386387 | ||
|
|
3511062c51 | ||
|
|
a134d27137 | ||
|
|
31a53c0870 |
@@ -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
@@ -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
|
||||
|
||||
@@ -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
@@ -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
@@ -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,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
|
||||
@@ -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
|
||||
@@ -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.
|
||||
@@ -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.
|
||||
@@ -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.
|
||||
@@ -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*.
|
||||
@@ -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.
|
||||
@@ -1,238 +1,58 @@
|
||||
<div align="center">
|
||||
# Hindsight
|
||||
|
||||

|
||||
**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.
|
||||
|
||||
[](https://github.com/vectorize-io/hindsight/actions/workflows/release.yml)
|
||||
[](https://opensource.org/licenses/MIT)
|
||||
[](https://pypi.org/project/hindsight-api/)
|
||||
[](https://pypi.org/project/hindsight-client/)
|
||||
[](https://www.npmjs.com/package/@vectorize-io/hindsight-client)
|
||||
[](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
|
||||
|
||||

|
||||
|
||||
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.
|
||||
|
||||

|
||||
|
||||
### 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
|
||||
|
||||

|
||||
|
||||
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?")
|
||||
```
|
||||
|
||||

|
||||
|
||||
---
|
||||
|
||||
## 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
@@ -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
@@ -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).
|
||||
@@ -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
|
||||
@@ -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"]
|
||||
Executable
+27
@@ -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 ""
|
||||
@@ -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"]
|
||||
@@ -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
@@ -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
|
||||
@@ -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
|
||||
@@ -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 $?
|
||||
Executable
+65
@@ -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
@@ -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 ""
|
||||
@@ -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
|
||||
@@ -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"
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
```
|
||||
@@ -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
|
||||
|
||||
@@ -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 }}
|
||||
|
||||
@@ -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 }}
|
||||
|
||||
@@ -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 }}
|
||||
@@ -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
@@ -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
|
||||
+49
@@ -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,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')
|
||||
-39
@@ -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')
|
||||
-48
@@ -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')
|
||||
@@ -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
@@ -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)
|
||||
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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",
|
||||
]
|
||||
|
||||
+145
-132
@@ -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'"
|
||||
)
|
||||
|
||||
@@ -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
+108
-331
@@ -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)
|
||||
|
||||
@@ -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
|
||||
+28
-39
@@ -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
|
||||
|
||||
+2
-4
@@ -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")
|
||||
+12
-29
@@ -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
Reference in New Issue
Block a user