Compare commits
26
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
d25222ba9c | ||
|
|
c882511f10 | ||
|
|
234d426499 | ||
|
|
e6511e7d77 | ||
|
|
904ea4de24 | ||
|
|
6168a77846 | ||
|
|
da44a5e839 | ||
|
|
32bca12c6f | ||
|
|
26850a0156 | ||
|
|
2a0c490c9e | ||
|
|
a831a7b77b | ||
|
|
d405b4feed | ||
|
|
b94b5cf26e | ||
|
|
6d820ef91b | ||
|
|
cf8882a867 | ||
|
|
490fccdc6f | ||
|
|
2948cb62d2 | ||
|
|
9053a51a88 | ||
|
|
f2c28cfd98 | ||
|
|
67fc532c43 | ||
|
|
9474f950f2 | ||
|
|
6a0c034f5d | ||
|
|
b52eb905ad | ||
|
|
1c6acc3ba0 | ||
|
|
8ecb5d3a0c | ||
|
|
ae80876671 |
@@ -42,6 +42,10 @@ jobs:
|
||||
working-directory: ./hindsight-integrations/litellm
|
||||
run: uv build --out-dir dist
|
||||
|
||||
- name: Build hindsight-embed
|
||||
working-directory: ./hindsight-embed
|
||||
run: uv build --out-dir dist
|
||||
|
||||
# Publish in order (client and api first, then hindsight-all which depends on them)
|
||||
- name: Publish hindsight-client to PyPI
|
||||
uses: pypa/gh-action-pypi-publish@release/v1
|
||||
@@ -67,6 +71,12 @@ jobs:
|
||||
packages-dir: ./hindsight-integrations/litellm/dist
|
||||
skip-existing: true
|
||||
|
||||
- name: Publish hindsight-embed to PyPI
|
||||
uses: pypa/gh-action-pypi-publish@release/v1
|
||||
with:
|
||||
packages-dir: ./hindsight-embed/dist
|
||||
skip-existing: true
|
||||
|
||||
# Upload artifacts for GitHub release
|
||||
- name: Upload artifacts
|
||||
uses: actions/upload-artifact@v4
|
||||
@@ -77,6 +87,7 @@ jobs:
|
||||
hindsight-api/dist/*
|
||||
hindsight/dist/*
|
||||
hindsight-integrations/litellm/dist/*
|
||||
hindsight-embed/dist/*
|
||||
retention-days: 1
|
||||
|
||||
release-typescript-client:
|
||||
@@ -102,7 +113,18 @@ jobs:
|
||||
|
||||
- name: Publish to npm
|
||||
working-directory: ./hindsight-clients/typescript
|
||||
run: npm publish --access public
|
||||
run: |
|
||||
set +e
|
||||
OUTPUT=$(npm publish --access public 2>&1)
|
||||
EXIT_CODE=$?
|
||||
echo "$OUTPUT"
|
||||
if [ $EXIT_CODE -ne 0 ]; then
|
||||
if echo "$OUTPUT" | grep -q "cannot publish over"; then
|
||||
echo "Package version already published, skipping..."
|
||||
exit 0
|
||||
fi
|
||||
exit $EXIT_CODE
|
||||
fi
|
||||
env:
|
||||
NODE_AUTH_TOKEN: ${{ secrets.NPM_TOKEN }}
|
||||
|
||||
@@ -117,6 +139,65 @@ jobs:
|
||||
path: hindsight-clients/typescript/*.tgz
|
||||
retention-days: 1
|
||||
|
||||
release-control-plane:
|
||||
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
|
||||
|
||||
- name: Build TypeScript client (dependency)
|
||||
run: npm run build --workspace=hindsight-clients/typescript
|
||||
|
||||
- name: Fix platform-specific native modules
|
||||
run: |
|
||||
# npm ci installs from lockfile which may have wrong platform binaries
|
||||
# Delete hoisted native modules and reinstall for current platform
|
||||
rm -rf node_modules/lightningcss node_modules/@tailwindcss
|
||||
npm install lightningcss @tailwindcss/postcss @tailwindcss/node
|
||||
|
||||
- name: Build
|
||||
run: npm run build --workspace=hindsight-control-plane
|
||||
|
||||
- name: Publish to npm
|
||||
working-directory: ./hindsight-control-plane
|
||||
run: |
|
||||
set +e
|
||||
OUTPUT=$(npm publish --access public 2>&1)
|
||||
EXIT_CODE=$?
|
||||
echo "$OUTPUT"
|
||||
if [ $EXIT_CODE -ne 0 ]; then
|
||||
if echo "$OUTPUT" | grep -q "cannot publish over"; then
|
||||
echo "Package version already published, skipping..."
|
||||
exit 0
|
||||
fi
|
||||
exit $EXIT_CODE
|
||||
fi
|
||||
env:
|
||||
NODE_AUTH_TOKEN: ${{ secrets.NPM_TOKEN }}
|
||||
|
||||
- name: Pack for GitHub release
|
||||
working-directory: ./hindsight-control-plane
|
||||
run: npm pack
|
||||
|
||||
- name: Upload artifacts
|
||||
uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: control-plane
|
||||
path: hindsight-control-plane/*.tgz
|
||||
retention-days: 1
|
||||
|
||||
release-rust-cli:
|
||||
runs-on: ${{ matrix.os }}
|
||||
strategy:
|
||||
@@ -181,7 +262,7 @@ jobs:
|
||||
- name: Free Disk Space
|
||||
uses: jlumbroso/free-disk-space@main
|
||||
with:
|
||||
tool-cache: false
|
||||
tool-cache: true
|
||||
android: true
|
||||
dotnet: true
|
||||
haskell: true
|
||||
@@ -206,7 +287,7 @@ jobs:
|
||||
id: get_version
|
||||
run: echo "VERSION=${GITHUB_REF#refs/tags/v}" >> $GITHUB_OUTPUT
|
||||
|
||||
- name: Extract metadata
|
||||
- name: Extract metadata for release tags
|
||||
id: meta
|
||||
uses: docker/metadata-action@v5
|
||||
with:
|
||||
@@ -217,7 +298,29 @@ jobs:
|
||||
type=semver,pattern={{major}},value=${{ steps.get_version.outputs.VERSION }}
|
||||
type=raw,value=latest
|
||||
|
||||
- name: Build and push
|
||||
# TODO: Re-enable smoke test when disk space issue is resolved
|
||||
# # Step 1: Build for local testing (single platform, no push)
|
||||
# # This creates an identical image to what will be released, just for one platform
|
||||
# - name: Build image for testing
|
||||
# uses: docker/build-push-action@v6
|
||||
# with:
|
||||
# context: .
|
||||
# file: docker/standalone/Dockerfile
|
||||
# target: ${{ matrix.target }}
|
||||
# push: false
|
||||
# load: true
|
||||
# tags: ${{ matrix.image_name }}:test
|
||||
# cache-from: type=gha
|
||||
# cache-to: type=gha,mode=max
|
||||
|
||||
# # Step 2: Test the image before pushing anything
|
||||
# - name: Smoke test - verify container starts
|
||||
# env:
|
||||
# GROQ_API_KEY: ${{ secrets.GROQ_API_KEY }}
|
||||
# run: ./scripts/docker-smoke-test.sh "${{ matrix.image_name }}:test" "${{ matrix.target }}"
|
||||
|
||||
# Build multi-platform and push to release tags
|
||||
- name: Build and push release images
|
||||
uses: docker/build-push-action@v6
|
||||
with:
|
||||
context: .
|
||||
@@ -263,7 +366,7 @@ jobs:
|
||||
|
||||
create-github-release:
|
||||
runs-on: ubuntu-latest
|
||||
needs: [release-python-packages, release-typescript-client, release-rust-cli, release-docker-images, release-helm-chart]
|
||||
needs: [release-python-packages, release-typescript-client, release-control-plane, release-rust-cli, release-docker-images, release-helm-chart]
|
||||
permissions:
|
||||
contents: write
|
||||
|
||||
@@ -286,6 +389,12 @@ jobs:
|
||||
name: typescript-client
|
||||
path: ./artifacts/typescript-client
|
||||
|
||||
- name: Download Control Plane
|
||||
uses: actions/download-artifact@v4
|
||||
with:
|
||||
name: control-plane
|
||||
path: ./artifacts/control-plane
|
||||
|
||||
- name: Download Rust CLI (Linux)
|
||||
uses: actions/download-artifact@v4
|
||||
with:
|
||||
@@ -318,8 +427,11 @@ jobs:
|
||||
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
|
||||
cp artifacts/python-packages/hindsight-embed/dist/* release-assets/ || true
|
||||
# TypeScript client
|
||||
cp artifacts/typescript-client/*.tgz release-assets/ || true
|
||||
# Control Plane
|
||||
cp artifacts/control-plane/*.tgz release-assets/ || true
|
||||
# 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
|
||||
|
||||
+222
-2
@@ -20,6 +20,8 @@ jobs:
|
||||
path: hindsight-api
|
||||
- name: hindsight-client
|
||||
path: hindsight-clients/python
|
||||
- name: hindsight-embed
|
||||
path: hindsight-embed
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
@@ -80,6 +82,58 @@ jobs:
|
||||
- name: Build TypeScript client
|
||||
run: npm run build --workspace=hindsight-clients/typescript
|
||||
|
||||
build-control-plane:
|
||||
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 SDK dependencies
|
||||
run: npm ci --workspace=hindsight-clients/typescript
|
||||
|
||||
- name: Build SDK
|
||||
run: npm run build --workspace=hindsight-clients/typescript
|
||||
|
||||
# Install control plane deps and fix hoisted lightningcss binary
|
||||
# lightningcss gets hoisted to root node_modules, so we need to reinstall it there
|
||||
- name: Install Control Plane dependencies
|
||||
run: |
|
||||
npm install --workspace=hindsight-control-plane
|
||||
rm -rf node_modules/lightningcss node_modules/@tailwindcss
|
||||
npm install lightningcss @tailwindcss/postcss @tailwindcss/node
|
||||
|
||||
- name: Build Control Plane
|
||||
run: npm run build --workspace=hindsight-control-plane
|
||||
|
||||
- name: Verify standalone build
|
||||
run: |
|
||||
test -f hindsight-control-plane/standalone/server.js || exit 1
|
||||
test -d hindsight-control-plane/standalone/node_modules || exit 1
|
||||
node hindsight-control-plane/bin/cli.js --help
|
||||
|
||||
- name: Smoke test - verify server starts
|
||||
run: |
|
||||
cd hindsight-control-plane
|
||||
node bin/cli.js --port 9999 &
|
||||
SERVER_PID=$!
|
||||
sleep 5
|
||||
if curl -sf http://localhost:9999 > /dev/null 2>&1; then
|
||||
echo "Server started successfully"
|
||||
kill $SERVER_PID 2>/dev/null || true
|
||||
exit 0
|
||||
else
|
||||
echo "Server failed to respond"
|
||||
kill $SERVER_PID 2>/dev/null || true
|
||||
exit 1
|
||||
fi
|
||||
|
||||
build-docs:
|
||||
runs-on: ubuntu-latest
|
||||
|
||||
@@ -121,6 +175,13 @@ jobs:
|
||||
working-directory: hindsight-cli
|
||||
run: cargo build --release
|
||||
|
||||
- name: Upload CLI artifact
|
||||
uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: hindsight-cli
|
||||
path: hindsight-cli/target/release/hindsight
|
||||
retention-days: 1
|
||||
|
||||
lint-helm-chart:
|
||||
runs-on: ubuntu-latest
|
||||
|
||||
@@ -153,7 +214,7 @@ jobs:
|
||||
- name: Free Disk Space
|
||||
uses: jlumbroso/free-disk-space@main
|
||||
with:
|
||||
tool-cache: false
|
||||
tool-cache: true
|
||||
android: true
|
||||
dotnet: true
|
||||
haskell: true
|
||||
@@ -171,6 +232,13 @@ jobs:
|
||||
file: docker/standalone/Dockerfile
|
||||
target: ${{ matrix.target }}
|
||||
push: false
|
||||
load: false
|
||||
|
||||
# TODO: Re-enable smoke test when disk space issue is resolved
|
||||
# - name: Smoke test - verify container starts
|
||||
# env:
|
||||
# GROQ_API_KEY: ${{ secrets.GROQ_API_KEY }}
|
||||
# run: ./scripts/docker-smoke-test.sh "hindsight-${{ matrix.name }}:test" "${{ matrix.target }}"
|
||||
|
||||
test-api:
|
||||
runs-on: ubuntu-latest
|
||||
@@ -495,4 +563,156 @@ jobs:
|
||||
|
||||
- name: Run tests
|
||||
working-directory: ./hindsight-integrations/litellm
|
||||
run: uv run pytest tests -v
|
||||
run: uv run pytest tests -v
|
||||
|
||||
test-embed:
|
||||
runs-on: ubuntu-latest
|
||||
env:
|
||||
HINDSIGHT_EMBED_LLM_PROVIDER: groq
|
||||
HINDSIGHT_EMBED_LLM_API_KEY: ${{ secrets.GROQ_API_KEY }}
|
||||
HINDSIGHT_EMBED_LLM_MODEL: openai/gpt-oss-20b
|
||||
# Prefer CPU-only PyTorch in CI
|
||||
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 dependencies
|
||||
working-directory: ./hindsight-embed
|
||||
run: uv sync --index-strategy unsafe-best-match
|
||||
|
||||
- name: Cache HuggingFace models
|
||||
uses: actions/cache@v4
|
||||
with:
|
||||
path: ~/.cache/huggingface
|
||||
key: ${{ runner.os }}-huggingface-embed-${{ hashFiles('hindsight-embed/pyproject.toml') }}
|
||||
restore-keys: |
|
||||
${{ runner.os }}-huggingface-embed-
|
||||
${{ runner.os }}-huggingface-
|
||||
|
||||
- name: Run smoke test
|
||||
working-directory: ./hindsight-embed
|
||||
run: ./test.sh
|
||||
|
||||
test-doc-examples:
|
||||
runs-on: ubuntu-latest
|
||||
needs: build-rust-cli
|
||||
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 }}
|
||||
UV_INDEX: pytorch=https://download.pytorch.org/whl/cpu
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
- name: Download CLI artifact
|
||||
uses: actions/download-artifact@v4
|
||||
with:
|
||||
name: hindsight-cli
|
||||
path: /usr/local/bin
|
||||
|
||||
- name: Make CLI executable
|
||||
run: chmod +x /usr/local/bin/hindsight
|
||||
|
||||
- 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 and install API
|
||||
working-directory: ./hindsight-api
|
||||
run: |
|
||||
uv build
|
||||
uv sync --no-install-project --index-strategy unsafe-best-match
|
||||
|
||||
- name: Install Python client dependencies
|
||||
working-directory: ./hindsight-clients/python
|
||||
run: uv sync --extra test --index-strategy unsafe-best-match
|
||||
|
||||
- name: Install TypeScript client
|
||||
run: |
|
||||
npm ci --workspace=hindsight-clients/typescript
|
||||
npm run build --workspace=hindsight-clients/typescript
|
||||
|
||||
- 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 doc examples
|
||||
working-directory: ./hindsight-clients/python
|
||||
run: |
|
||||
for f in ../../hindsight-docs/examples/api/*.py; do
|
||||
echo "Running $f..."
|
||||
uv run python "$f"
|
||||
done
|
||||
|
||||
- name: Run Node.js doc examples
|
||||
run: |
|
||||
for f in hindsight-docs/examples/api/*.mjs; do
|
||||
echo "Running $f..."
|
||||
node "$f"
|
||||
done
|
||||
|
||||
- name: Configure CLI
|
||||
run: hindsight configure --api-url http://localhost:8888
|
||||
|
||||
- name: Run CLI doc examples
|
||||
run: |
|
||||
for f in hindsight-docs/examples/api/*.sh; do
|
||||
echo "Running $f..."
|
||||
bash "$f"
|
||||
done
|
||||
|
||||
- name: Show API server logs
|
||||
if: always()
|
||||
run: |
|
||||
echo "=== API Server Logs ==="
|
||||
cat /tmp/api-server.log || echo "No API server log found"
|
||||
@@ -14,6 +14,7 @@ This document captures architectural decisions and coding conventions for the Hi
|
||||
hindsight/ # Python package for embedded usage
|
||||
hindsight-api/ # FastAPI server (core memory engine)
|
||||
hindsight-cli/ # Rust CLI client
|
||||
hindsight-embed/ # Embedded CLI (no server needed)
|
||||
hindsight-control-plane/ # Next.js admin UI
|
||||
hindsight-docs/ # Docusaurus documentation site
|
||||
hindsight-dev/ # Development tools and benchmarks
|
||||
@@ -148,4 +149,5 @@ Note: The maintained wrapper `hindsight_client.py` and `README.md` are preserved
|
||||
|
||||
# Branding
|
||||
## Colors
|
||||
- Primary: gradient from #0074d9 to #009296
|
||||
- Primary: gradient from #0074d9 to #009296
|
||||
|
||||
|
||||
@@ -2,16 +2,19 @@
|
||||
# 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)
|
||||
# INCLUDE_API=true/false - Include API (default: true)
|
||||
# INCLUDE_CP=true/false - Include Control Plane (default: true)
|
||||
# PRELOAD_ML_MODELS=true/false - Pre-download ML models during build (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
|
||||
# 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
|
||||
# docker build -t hindsight --build-arg PRELOAD_ML_MODELS=false . # Skip ML model preload
|
||||
|
||||
ARG INCLUDE_API=true
|
||||
ARG INCLUDE_CP=true
|
||||
ARG PRELOAD_ML_MODELS=true
|
||||
|
||||
# =============================================================================
|
||||
# Stage: API Builder
|
||||
@@ -72,30 +75,48 @@ 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
|
||||
# Create directory structure matching the monorepo layout
|
||||
# This is required because build:standalone script expects .next/standalone/memory-poc/hindsight-control-plane
|
||||
WORKDIR /app/memory-poc/hindsight-control-plane
|
||||
|
||||
# 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 ./
|
||||
# Remove the file: dependency on SDK (we'll copy it directly later)
|
||||
RUN sed -i '/"@vectorize-io\/hindsight-client":/d' 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
|
||||
# Also remove the file: dependency from package.json (restored by COPY above)
|
||||
RUN rm -f package-lock.json && sed -i '/"@vectorize-io\/hindsight-client":/d' package.json
|
||||
|
||||
# Link SDK (temporary for build)
|
||||
RUN cd /app/sdk && npm link && cd /app && npm link @vectorize-io/hindsight-client
|
||||
# Copy built SDK directly into node_modules (more reliable than npm link in Docker)
|
||||
COPY --from=sdk-builder /app/hindsight-clients/typescript ./node_modules/@vectorize-io/hindsight-client
|
||||
|
||||
# Build Control Plane
|
||||
RUN npm run build
|
||||
# Build Control Plane - run next build first, then custom standalone copy
|
||||
# (The build:standalone script expects a specific path structure that differs in Docker)
|
||||
RUN npm exec -- next build
|
||||
|
||||
# Create public directory if it doesn't exist
|
||||
RUN mkdir -p public
|
||||
# Create standalone directory structure manually
|
||||
# Note: Must exclude node_modules from find to avoid wrong server.js from next/dist/experimental/testmode/
|
||||
# Note: Must explicitly copy .next since glob * doesn't match hidden directories
|
||||
RUN STANDALONE_ROOT=$(find .next/standalone -path '*/node_modules' -prune -o -name 'server.js' -print | head -1 | xargs dirname) && \
|
||||
mkdir -p standalone && \
|
||||
cp -r "$STANDALONE_ROOT"/* standalone/ && \
|
||||
cp -r "$STANDALONE_ROOT"/.next standalone/.next && \
|
||||
# Copy node_modules if separate from app dir (monorepo structure)
|
||||
if [ -d ".next/standalone/node_modules" ] && [ "$STANDALONE_ROOT" != ".next/standalone" ]; then \
|
||||
cp -r .next/standalone/node_modules standalone/node_modules; \
|
||||
fi && \
|
||||
cp -r .next/static standalone/.next/static && \
|
||||
mkdir -p standalone/public && \
|
||||
cp -r public/* standalone/public/ 2>/dev/null || true && \
|
||||
# Verify required files exist
|
||||
test -f standalone/server.js || (echo "ERROR: server.js missing!" && exit 1) && \
|
||||
test -f standalone/.next/BUILD_ID || (echo "ERROR: BUILD_ID missing!" && exit 1)
|
||||
|
||||
# =============================================================================
|
||||
# Stage: Final Image - API Only
|
||||
@@ -104,14 +125,16 @@ FROM python:3.11-slim AS api-only
|
||||
|
||||
WORKDIR /app
|
||||
|
||||
# Install pg0 dependencies
|
||||
# Install pg0 dependencies (procps provides 'kill' command needed by pg0)
|
||||
# Note: libicu version varies by Debian version - try common versions in order
|
||||
RUN apt-get update && apt-get install -y \
|
||||
curl \
|
||||
procps \
|
||||
libxml2 \
|
||||
libssl3 \
|
||||
libgssapi-krb5-2 \
|
||||
libossp-uuid16 \
|
||||
&& apt-get install -y libicu72 || apt-get install -y libicu74 || apt-get install -y libicu* \
|
||||
&& (apt-get install -y libicu72 2>/dev/null || apt-get install -y libicu74 2>/dev/null || apt-get install -y libicu76 2>/dev/null || true) \
|
||||
&& rm -rf /var/lib/apt/lists/* \
|
||||
&& pip install --no-cache-dir uv
|
||||
|
||||
@@ -139,14 +162,17 @@ 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 "\
|
||||
# Pre-download ML models to avoid runtime download (conditional)
|
||||
ARG PRELOAD_ML_MODELS
|
||||
RUN if [ "$PRELOAD_ML_MODELS" = "true" ]; then \
|
||||
/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')"
|
||||
print('Models cached successfully')"; \
|
||||
else echo "Skipping ML model preload"; fi
|
||||
|
||||
EXPOSE 8888
|
||||
|
||||
@@ -171,9 +197,9 @@ 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
|
||||
COPY --from=cp-builder /app/memory-poc/hindsight-control-plane/standalone ./
|
||||
COPY --from=cp-builder /app/memory-poc/hindsight-control-plane/.next/static ./.next/static
|
||||
COPY --from=cp-builder /app/memory-poc/hindsight-control-plane/public ./public
|
||||
|
||||
WORKDIR /app
|
||||
|
||||
@@ -200,14 +226,16 @@ FROM python:3.11-slim AS standalone
|
||||
|
||||
WORKDIR /app
|
||||
|
||||
# Install Node.js, curl, uv, and pg0 dependencies
|
||||
# Install Node.js, curl, uv, and pg0 dependencies (procps provides 'kill' command needed by pg0)
|
||||
# Note: libicu version varies by Debian version - try common versions in order
|
||||
RUN apt-get update && apt-get install -y \
|
||||
curl \
|
||||
procps \
|
||||
libxml2 \
|
||||
libssl3 \
|
||||
libgssapi-krb5-2 \
|
||||
libossp-uuid16 \
|
||||
&& apt-get install -y libicu72 || apt-get install -y libicu74 || apt-get install -y libicu* \
|
||||
&& (apt-get install -y libicu72 2>/dev/null || apt-get install -y libicu74 2>/dev/null || apt-get install -y libicu76 2>/dev/null || true) \
|
||||
&& curl -fsSL https://deb.nodesource.com/setup_20.x | bash - \
|
||||
&& apt-get install -y nodejs \
|
||||
&& rm -rf /var/lib/apt/lists/* \
|
||||
@@ -224,9 +252,9 @@ 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
|
||||
COPY --from=cp-builder /app/memory-poc/hindsight-control-plane/standalone ./
|
||||
COPY --from=cp-builder /app/memory-poc/hindsight-control-plane/.next/static ./.next/static
|
||||
COPY --from=cp-builder /app/memory-poc/hindsight-control-plane/public ./public
|
||||
|
||||
WORKDIR /app
|
||||
|
||||
@@ -255,14 +283,17 @@ 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 "\
|
||||
# Pre-download ML models to avoid runtime download (conditional)
|
||||
ARG PRELOAD_ML_MODELS
|
||||
RUN if [ "$PRELOAD_ML_MODELS" = "true" ]; then \
|
||||
/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')"
|
||||
print('Models cached successfully')"; \
|
||||
else echo "Skipping ML model preload"; fi
|
||||
|
||||
EXPOSE 8888 9999
|
||||
|
||||
|
||||
@@ -23,7 +23,8 @@ PIDS=()
|
||||
# Start API if enabled
|
||||
if [ "$ENABLE_API" = "true" ]; then
|
||||
cd /app/api
|
||||
hindsight-api 2>&1 | sed -u 's/^/[api] /' &
|
||||
# Run API directly - Python's PYTHONUNBUFFERED=1 handles output buffering
|
||||
hindsight-api &
|
||||
API_PID=$!
|
||||
PIDS+=($API_PID)
|
||||
|
||||
@@ -42,7 +43,7 @@ fi
|
||||
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] /' &
|
||||
PORT=9999 node server.js &
|
||||
CP_PID=$!
|
||||
PIDS+=($CP_PID)
|
||||
else
|
||||
|
||||
@@ -2,8 +2,8 @@ apiVersion: v2
|
||||
name: hindsight
|
||||
description: Hindsight helm chart
|
||||
type: application
|
||||
version: 0.1.8
|
||||
appVersion: "0.1.8"
|
||||
version: 0.1.13
|
||||
appVersion: "0.1.13"
|
||||
keywords:
|
||||
- ai
|
||||
- memory
|
||||
|
||||
@@ -21,9 +21,11 @@ from .engine.search.trace import (
|
||||
WeightComponents,
|
||||
)
|
||||
from .engine.search.tracer import SearchTracer
|
||||
from .models import RequestContext
|
||||
|
||||
__all__ = [
|
||||
"MemoryEngine",
|
||||
"RequestContext",
|
||||
"HindsightConfig",
|
||||
"get_config",
|
||||
"SearchTrace",
|
||||
|
||||
@@ -109,6 +109,9 @@ def run_migrations_online() -> None:
|
||||
|
||||
get_database_url() # Process and set the database URL in config
|
||||
|
||||
# Check if we're targeting a specific schema (for multi-tenant isolation)
|
||||
target_schema = config.get_main_option("target_schema")
|
||||
|
||||
connectable = engine_from_config(
|
||||
config.get_section(config.config_ini_section, {}),
|
||||
prefix="sqlalchemy.",
|
||||
@@ -121,14 +124,34 @@ def run_migrations_online() -> None:
|
||||
def set_read_write_mode(dbapi_connection, connection_record):
|
||||
cursor = dbapi_connection.cursor()
|
||||
cursor.execute("SET SESSION CHARACTERISTICS AS TRANSACTION READ WRITE")
|
||||
# If targeting a specific schema, set search_path
|
||||
# Include public in search_path for access to shared extensions (pgvector)
|
||||
if target_schema:
|
||||
cursor.execute(f'CREATE SCHEMA IF NOT EXISTS "{target_schema}"')
|
||||
cursor.execute(f'SET search_path TO "{target_schema}", public')
|
||||
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"))
|
||||
|
||||
# If targeting a specific schema, set search_path
|
||||
# Include public in search_path for access to shared extensions (pgvector)
|
||||
if target_schema:
|
||||
connection.execute(text(f'CREATE SCHEMA IF NOT EXISTS "{target_schema}"'))
|
||||
connection.execute(text(f'SET search_path TO "{target_schema}", public'))
|
||||
|
||||
connection.commit() # Commit the SET command
|
||||
|
||||
context.configure(connection=connection, target_metadata=target_metadata)
|
||||
# Configure context with version_table_schema if using a specific schema
|
||||
context_opts = {
|
||||
"connection": connection,
|
||||
"target_metadata": target_metadata,
|
||||
}
|
||||
if target_schema:
|
||||
context_opts["version_table_schema"] = target_schema
|
||||
|
||||
context.configure(**context_opts)
|
||||
|
||||
with context.begin_transaction():
|
||||
context.run_migrations()
|
||||
|
||||
+14
-4
@@ -6,7 +6,7 @@ Create Date: 2024-12-04 15:00:00.000000
|
||||
|
||||
"""
|
||||
|
||||
from alembic import op
|
||||
from alembic import context, op
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision = "d9f6a3b4c5e2"
|
||||
@@ -15,14 +15,22 @@ branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def _get_schema_prefix() -> str:
|
||||
"""Get schema prefix for table names (e.g., 'tenant_x.' or '' for public)."""
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def upgrade():
|
||||
schema = _get_schema_prefix()
|
||||
|
||||
# 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'")
|
||||
op.execute(f"UPDATE {schema}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'")
|
||||
op.execute(f"UPDATE {schema}memory_units SET fact_type = 'experience' WHERE fact_type = 'interactions'")
|
||||
|
||||
# Create new check constraint with 'experience' instead of 'bank'
|
||||
op.create_check_constraint(
|
||||
@@ -31,11 +39,13 @@ def upgrade():
|
||||
|
||||
|
||||
def downgrade():
|
||||
schema = _get_schema_prefix()
|
||||
|
||||
# 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'")
|
||||
op.execute(f"UPDATE {schema}memory_units SET fact_type = 'bank' WHERE fact_type = 'experience'")
|
||||
|
||||
# Recreate old check constraint
|
||||
op.create_check_constraint(
|
||||
|
||||
+54
-13
@@ -12,7 +12,7 @@ system (skepticism, literalism, empathy with 1-5 integer values).
|
||||
from collections.abc import Sequence
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
from alembic import context, op
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision: str = "e0a1b2c3d4e5"
|
||||
@@ -21,9 +21,36 @@ branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _get_schema_prefix() -> str:
|
||||
"""Get schema prefix for table names (e.g., 'tenant_x.' or '' for public)."""
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def _get_target_schema() -> str:
|
||||
"""Get the target schema name (tenant schema or 'public')."""
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return schema if schema else "public"
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
"""Convert Big Five disposition to 3-trait disposition."""
|
||||
conn = op.get_bind()
|
||||
schema = _get_schema_prefix()
|
||||
target_schema = _get_target_schema()
|
||||
|
||||
# Check if disposition column exists (should have been created by previous migration)
|
||||
result = conn.execute(
|
||||
sa.text("""
|
||||
SELECT column_name
|
||||
FROM information_schema.columns
|
||||
WHERE table_schema = :schema AND table_name = 'banks' AND column_name = 'disposition'
|
||||
"""),
|
||||
{"schema": target_schema},
|
||||
)
|
||||
if not result.fetchone():
|
||||
# Column doesn't exist yet (shouldn't happen but be safe)
|
||||
return
|
||||
|
||||
# Update all existing banks to use the new disposition format
|
||||
# Convert from old format to new format with reasonable mappings:
|
||||
@@ -32,18 +59,18 @@ def upgrade() -> None:
|
||||
# - 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
|
||||
sa.text(f"""
|
||||
UPDATE {schema}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
|
||||
sa.text(f"""
|
||||
ALTER TABLE {schema}banks
|
||||
ALTER COLUMN disposition SET DEFAULT '{{"skepticism": 3, "literalism": 3, "empathy": 3}}'::jsonb
|
||||
""")
|
||||
)
|
||||
|
||||
@@ -51,20 +78,34 @@ def upgrade() -> None:
|
||||
def downgrade() -> None:
|
||||
"""Convert back to Big Five disposition."""
|
||||
conn = op.get_bind()
|
||||
schema = _get_schema_prefix()
|
||||
target_schema = _get_target_schema()
|
||||
|
||||
# Check if disposition column exists
|
||||
result = conn.execute(
|
||||
sa.text("""
|
||||
SELECT column_name
|
||||
FROM information_schema.columns
|
||||
WHERE table_schema = :schema AND table_name = 'banks' AND column_name = 'disposition'
|
||||
"""),
|
||||
{"schema": target_schema},
|
||||
)
|
||||
if not result.fetchone():
|
||||
return
|
||||
|
||||
# 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
|
||||
sa.text(f"""
|
||||
UPDATE {schema}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
|
||||
sa.text(f"""
|
||||
ALTER TABLE {schema}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
|
||||
""")
|
||||
)
|
||||
|
||||
@@ -9,7 +9,7 @@ Create Date: 2024-12-04
|
||||
from collections.abc import Sequence
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
from alembic import context, op
|
||||
from sqlalchemy.dialects import postgresql
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
@@ -19,17 +19,25 @@ branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _get_target_schema() -> str:
|
||||
"""Get the target schema name (tenant schema or 'public')."""
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return schema if schema else "public"
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
"""Rename personality column to disposition in banks table (if it exists)."""
|
||||
conn = op.get_bind()
|
||||
target_schema = _get_target_schema()
|
||||
|
||||
# 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'
|
||||
""")
|
||||
WHERE table_schema = :schema AND table_name = 'banks' AND column_name = 'personality'
|
||||
"""),
|
||||
{"schema": target_schema},
|
||||
)
|
||||
has_personality = result.fetchone() is not None
|
||||
|
||||
@@ -38,8 +46,9 @@ def upgrade() -> None:
|
||||
sa.text("""
|
||||
SELECT column_name
|
||||
FROM information_schema.columns
|
||||
WHERE table_name = 'banks' AND column_name = 'disposition'
|
||||
""")
|
||||
WHERE table_schema = :schema AND table_name = 'banks' AND column_name = 'disposition'
|
||||
"""),
|
||||
{"schema": target_schema},
|
||||
)
|
||||
has_disposition = result.fetchone() is not None
|
||||
|
||||
@@ -63,12 +72,14 @@ def upgrade() -> None:
|
||||
def downgrade() -> None:
|
||||
"""Revert disposition column back to personality."""
|
||||
conn = op.get_bind()
|
||||
target_schema = _get_target_schema()
|
||||
result = conn.execute(
|
||||
sa.text("""
|
||||
SELECT column_name
|
||||
FROM information_schema.columns
|
||||
WHERE table_name = 'banks' AND column_name = 'disposition'
|
||||
""")
|
||||
WHERE table_schema = :schema AND table_name = 'banks' AND column_name = 'disposition'
|
||||
"""),
|
||||
{"schema": target_schema},
|
||||
)
|
||||
if result.fetchone():
|
||||
op.alter_column("banks", "disposition", new_column_name="personality")
|
||||
|
||||
@@ -12,7 +12,7 @@ from contextlib import asynccontextmanager
|
||||
from datetime import datetime
|
||||
from typing import Any
|
||||
|
||||
from fastapi import FastAPI, HTTPException, Query
|
||||
from fastapi import Depends, FastAPI, Header, HTTPException, Query
|
||||
|
||||
|
||||
def _parse_metadata(metadata: Any) -> dict[str, Any]:
|
||||
@@ -29,13 +29,15 @@ def _parse_metadata(metadata: Any) -> dict[str, Any]:
|
||||
return {}
|
||||
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
from pydantic import BaseModel, ConfigDict, Field, field_validator
|
||||
|
||||
from hindsight_api import MemoryEngine
|
||||
from hindsight_api.engine.db_utils import acquire_with_retry
|
||||
from hindsight_api.engine.memory_engine import Budget
|
||||
from hindsight_api.engine.memory_engine import Budget, fq_table
|
||||
from hindsight_api.engine.response_models import VALID_RECALL_FACT_TYPES
|
||||
from hindsight_api.extensions import HttpExtension, load_extension
|
||||
from hindsight_api.metrics import create_metrics_collector, get_metrics_collector, initialize_metrics
|
||||
from hindsight_api.models import RequestContext
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -289,7 +291,7 @@ class MemoryItem(BaseModel):
|
||||
"metadata": {"source": "slack", "channel": "engineering"},
|
||||
"document_id": "meeting_notes_2024_01_15",
|
||||
}
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
content: str
|
||||
@@ -298,6 +300,23 @@ class MemoryItem(BaseModel):
|
||||
metadata: dict[str, str] | None = None
|
||||
document_id: str | None = Field(default=None, description="Optional document ID for this memory item.")
|
||||
|
||||
@field_validator("timestamp", mode="before")
|
||||
@classmethod
|
||||
def validate_timestamp(cls, v):
|
||||
if v is None or v == "":
|
||||
return None
|
||||
if isinstance(v, datetime):
|
||||
return v
|
||||
if isinstance(v, str):
|
||||
try:
|
||||
# Try parsing as ISO format
|
||||
return datetime.fromisoformat(v.replace("Z", "+00:00"))
|
||||
except ValueError as e:
|
||||
raise ValueError(
|
||||
f"Invalid timestamp/event_date format: '{v}'. Expected ISO format like '2024-01-15T10:30:00' or '2024-01-15T10:30:00Z'"
|
||||
) from e
|
||||
raise ValueError(f"timestamp must be a string or datetime, got {type(v).__name__}")
|
||||
|
||||
|
||||
class RetainRequest(BaseModel):
|
||||
"""Request model for retain endpoint."""
|
||||
@@ -337,7 +356,7 @@ class RetainResponse(BaseModel):
|
||||
success: bool
|
||||
bank_id: str
|
||||
items_count: int
|
||||
async_: bool = Field(
|
||||
is_async: bool = Field(
|
||||
alias="async", serialization_alias="async", description="Whether the operation was processed asynchronously"
|
||||
)
|
||||
|
||||
@@ -706,7 +725,11 @@ class DeleteResponse(BaseModel):
|
||||
deleted_count: int | None = None
|
||||
|
||||
|
||||
def create_app(memory: MemoryEngine, initialize_memory: bool = True) -> FastAPI:
|
||||
def create_app(
|
||||
memory: MemoryEngine,
|
||||
initialize_memory: bool = True,
|
||||
http_extension: HttpExtension | None = None,
|
||||
) -> FastAPI:
|
||||
"""
|
||||
Create and configure the FastAPI application.
|
||||
|
||||
@@ -714,6 +737,8 @@ def create_app(memory: MemoryEngine, initialize_memory: bool = True) -> FastAPI:
|
||||
memory: MemoryEngine instance (already initialized with required parameters).
|
||||
Migrations are controlled by the MemoryEngine's run_migrations parameter.
|
||||
initialize_memory: Whether to initialize memory system on startup (default: True)
|
||||
http_extension: Optional HTTP extension to mount custom endpoints under /extension/.
|
||||
If None, attempts to load from HINDSIGHT_API_HTTP_EXTENSION env var.
|
||||
|
||||
Returns:
|
||||
Configured FastAPI application
|
||||
@@ -723,6 +748,11 @@ def create_app(memory: MemoryEngine, initialize_memory: bool = True) -> FastAPI:
|
||||
In that case, you should call memory.initialize() manually before starting the server
|
||||
and memory.close() when shutting down.
|
||||
"""
|
||||
# Load HTTP extension from environment if not provided
|
||||
if http_extension is None:
|
||||
http_extension = load_extension("HTTP", HttpExtension)
|
||||
if http_extension:
|
||||
logging.info(f"Loaded HTTP extension: {http_extension.__class__.__name__}")
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(app: FastAPI):
|
||||
@@ -746,8 +776,18 @@ def create_app(memory: MemoryEngine, initialize_memory: bool = True) -> FastAPI:
|
||||
await memory.initialize()
|
||||
logging.info("Memory system initialized")
|
||||
|
||||
# Call HTTP extension startup hook
|
||||
if http_extension:
|
||||
await http_extension.on_startup()
|
||||
logging.info("HTTP extension started")
|
||||
|
||||
yield
|
||||
|
||||
# Call HTTP extension shutdown hook
|
||||
if http_extension:
|
||||
await http_extension.on_shutdown()
|
||||
logging.info("HTTP extension stopped")
|
||||
|
||||
# Shutdown: Cleanup memory system
|
||||
await memory.close()
|
||||
logging.info("Memory system closed")
|
||||
@@ -775,12 +815,36 @@ def create_app(memory: MemoryEngine, initialize_memory: bool = True) -> FastAPI:
|
||||
# Register all routes
|
||||
_register_routes(app)
|
||||
|
||||
# Mount HTTP extension router if available
|
||||
if http_extension:
|
||||
extension_router = http_extension.get_router(memory)
|
||||
app.include_router(extension_router, prefix="/ext", tags=["Extension"])
|
||||
logging.info("HTTP extension router mounted at /ext/")
|
||||
|
||||
return app
|
||||
|
||||
|
||||
def _register_routes(app: FastAPI):
|
||||
"""Register all API routes on the given app instance."""
|
||||
|
||||
def get_request_context(authorization: str | None = Header(default=None)) -> RequestContext:
|
||||
"""
|
||||
Extract request context from Authorization header.
|
||||
|
||||
Supports:
|
||||
- Bearer token: "Bearer <api_key>"
|
||||
- Direct API key: "<api_key>"
|
||||
|
||||
Returns RequestContext with extracted API key (may be None if no auth header).
|
||||
"""
|
||||
api_key = None
|
||||
if authorization:
|
||||
if authorization.lower().startswith("bearer "):
|
||||
api_key = authorization[7:].strip()
|
||||
else:
|
||||
api_key = authorization.strip()
|
||||
return RequestContext(api_key=api_key)
|
||||
|
||||
@app.get(
|
||||
"/health",
|
||||
summary="Health check endpoint",
|
||||
@@ -821,10 +885,12 @@ def _register_routes(app: FastAPI):
|
||||
operation_id="get_graph",
|
||||
tags=["Memory"],
|
||||
)
|
||||
async def api_graph(bank_id: str, type: str | None = None):
|
||||
async def api_graph(
|
||||
bank_id: str, type: str | None = None, request_context: RequestContext = Depends(get_request_context)
|
||||
):
|
||||
"""Get graph data from database, filtered by bank_id and optionally by type."""
|
||||
try:
|
||||
data = await app.state.memory.get_graph_data(bank_id, type)
|
||||
data = await app.state.memory.get_graph_data(bank_id, type, request_context=request_context)
|
||||
return data
|
||||
except Exception as e:
|
||||
import traceback
|
||||
@@ -841,7 +907,14 @@ def _register_routes(app: FastAPI):
|
||||
operation_id="list_memories",
|
||||
tags=["Memory"],
|
||||
)
|
||||
async def api_list(bank_id: str, type: str | None = None, q: str | None = None, limit: int = 100, offset: int = 0):
|
||||
async def api_list(
|
||||
bank_id: str,
|
||||
type: str | None = None,
|
||||
q: str | None = None,
|
||||
limit: int = 100,
|
||||
offset: int = 0,
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
):
|
||||
"""
|
||||
List memory units for table view with optional full-text search.
|
||||
|
||||
@@ -857,7 +930,12 @@ def _register_routes(app: FastAPI):
|
||||
"""
|
||||
try:
|
||||
data = await app.state.memory.list_memory_units(
|
||||
bank_id=bank_id, fact_type=type, search_query=q, limit=limit, offset=offset
|
||||
bank_id=bank_id,
|
||||
fact_type=type,
|
||||
search_query=q,
|
||||
limit=limit,
|
||||
offset=offset,
|
||||
request_context=request_context,
|
||||
)
|
||||
return data
|
||||
except Exception as e:
|
||||
@@ -880,7 +958,9 @@ def _register_routes(app: FastAPI):
|
||||
operation_id="recall_memories",
|
||||
tags=["Memory"],
|
||||
)
|
||||
async def api_recall(bank_id: str, request: RecallRequest):
|
||||
async def api_recall(
|
||||
bank_id: str, request: RecallRequest, request_context: RequestContext = Depends(get_request_context)
|
||||
):
|
||||
"""Run a recall and return results with trace."""
|
||||
metrics = get_metrics_collector()
|
||||
|
||||
@@ -923,6 +1003,7 @@ def _register_routes(app: FastAPI):
|
||||
max_entity_tokens=max_entity_tokens,
|
||||
include_chunks=include_chunks,
|
||||
max_chunk_tokens=max_chunk_tokens,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Convert core MemoryFact objects to API RecallResult objects (excluding internal metrics)
|
||||
@@ -995,14 +1076,20 @@ def _register_routes(app: FastAPI):
|
||||
operation_id="reflect",
|
||||
tags=["Memory"],
|
||||
)
|
||||
async def api_reflect(bank_id: str, request: ReflectRequest):
|
||||
async def api_reflect(
|
||||
bank_id: str, request: ReflectRequest, request_context: RequestContext = Depends(get_request_context)
|
||||
):
|
||||
metrics = get_metrics_collector()
|
||||
|
||||
try:
|
||||
# Use the memory system's reflect_async method (record metrics)
|
||||
with metrics.record_operation("reflect", bank_id=bank_id, budget=request.budget.value):
|
||||
core_result = await app.state.memory.reflect_async(
|
||||
bank_id=bank_id, query=request.query, budget=request.budget, context=request.context
|
||||
bank_id=bank_id,
|
||||
query=request.query,
|
||||
budget=request.budget,
|
||||
context=request.context,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Convert core MemoryFact objects to API ReflectFact objects if facts are requested
|
||||
@@ -1041,10 +1128,10 @@ def _register_routes(app: FastAPI):
|
||||
operation_id="list_banks",
|
||||
tags=["Banks"],
|
||||
)
|
||||
async def api_list_banks():
|
||||
async def api_list_banks(request_context: RequestContext = Depends(get_request_context)):
|
||||
"""Get list of all banks with their profiles."""
|
||||
try:
|
||||
banks = await app.state.memory.list_banks()
|
||||
banks = await app.state.memory.list_banks(request_context=request_context)
|
||||
return BankListResponse(banks=banks)
|
||||
except Exception as e:
|
||||
import traceback
|
||||
@@ -1067,9 +1154,9 @@ def _register_routes(app: FastAPI):
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
# Get node counts by fact_type
|
||||
node_stats = await conn.fetch(
|
||||
"""
|
||||
f"""
|
||||
SELECT fact_type, COUNT(*) as count
|
||||
FROM memory_units
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE bank_id = $1
|
||||
GROUP BY fact_type
|
||||
""",
|
||||
@@ -1078,10 +1165,10 @@ def _register_routes(app: FastAPI):
|
||||
|
||||
# Get link counts by link_type
|
||||
link_stats = await conn.fetch(
|
||||
"""
|
||||
f"""
|
||||
SELECT ml.link_type, COUNT(*) as count
|
||||
FROM memory_links ml
|
||||
JOIN memory_units mu ON ml.from_unit_id = mu.id
|
||||
FROM {fq_table("memory_links")} ml
|
||||
JOIN {fq_table("memory_units")} mu ON ml.from_unit_id = mu.id
|
||||
WHERE mu.bank_id = $1
|
||||
GROUP BY ml.link_type
|
||||
""",
|
||||
@@ -1090,10 +1177,10 @@ def _register_routes(app: FastAPI):
|
||||
|
||||
# Get link counts by fact_type (from nodes)
|
||||
link_fact_type_stats = await conn.fetch(
|
||||
"""
|
||||
f"""
|
||||
SELECT mu.fact_type, COUNT(*) as count
|
||||
FROM memory_links ml
|
||||
JOIN memory_units mu ON ml.from_unit_id = mu.id
|
||||
FROM {fq_table("memory_links")} ml
|
||||
JOIN {fq_table("memory_units")} mu ON ml.from_unit_id = mu.id
|
||||
WHERE mu.bank_id = $1
|
||||
GROUP BY mu.fact_type
|
||||
""",
|
||||
@@ -1102,10 +1189,10 @@ def _register_routes(app: FastAPI):
|
||||
|
||||
# Get link counts by fact_type AND link_type
|
||||
link_breakdown_stats = await conn.fetch(
|
||||
"""
|
||||
f"""
|
||||
SELECT mu.fact_type, ml.link_type, COUNT(*) as count
|
||||
FROM memory_links ml
|
||||
JOIN memory_units mu ON ml.from_unit_id = mu.id
|
||||
FROM {fq_table("memory_links")} ml
|
||||
JOIN {fq_table("memory_units")} mu ON ml.from_unit_id = mu.id
|
||||
WHERE mu.bank_id = $1
|
||||
GROUP BY mu.fact_type, ml.link_type
|
||||
""",
|
||||
@@ -1114,9 +1201,9 @@ def _register_routes(app: FastAPI):
|
||||
|
||||
# Get pending and failed operations counts
|
||||
ops_stats = await conn.fetch(
|
||||
"""
|
||||
f"""
|
||||
SELECT status, COUNT(*) as count
|
||||
FROM async_operations
|
||||
FROM {fq_table("async_operations")}
|
||||
WHERE bank_id = $1
|
||||
GROUP BY status
|
||||
""",
|
||||
@@ -1128,9 +1215,9 @@ def _register_routes(app: FastAPI):
|
||||
|
||||
# Get document count
|
||||
doc_count_result = await conn.fetchrow(
|
||||
"""
|
||||
f"""
|
||||
SELECT COUNT(*) as count
|
||||
FROM documents
|
||||
FROM {fq_table("documents")}
|
||||
WHERE bank_id = $1
|
||||
""",
|
||||
bank_id,
|
||||
@@ -1184,11 +1271,13 @@ def _register_routes(app: FastAPI):
|
||||
tags=["Entities"],
|
||||
)
|
||||
async def api_list_entities(
|
||||
bank_id: str, limit: int = Query(default=100, description="Maximum number of entities to return")
|
||||
bank_id: str,
|
||||
limit: int = Query(default=100, description="Maximum number of entities to return"),
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
):
|
||||
"""List entities for a memory bank."""
|
||||
try:
|
||||
entities = await app.state.memory.list_entities(bank_id, limit=limit)
|
||||
entities = await app.state.memory.list_entities(bank_id, limit=limit, request_context=request_context)
|
||||
return EntityListResponse(items=[EntityListItem(**e) for e in entities])
|
||||
except Exception as e:
|
||||
import traceback
|
||||
@@ -1205,37 +1294,26 @@ def _register_routes(app: FastAPI):
|
||||
operation_id="get_entity",
|
||||
tags=["Entities"],
|
||||
)
|
||||
async def api_get_entity(bank_id: str, entity_id: str):
|
||||
async def api_get_entity(
|
||||
bank_id: str, entity_id: str, request_context: RequestContext = Depends(get_request_context)
|
||||
):
|
||||
"""Get entity details with observations."""
|
||||
try:
|
||||
# First get the entity metadata
|
||||
pool = await app.state.memory._get_pool()
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
entity_row = await conn.fetchrow(
|
||||
"""
|
||||
SELECT id, canonical_name, mention_count, first_seen, last_seen, metadata
|
||||
FROM entities
|
||||
WHERE bank_id = $1 AND id = $2
|
||||
""",
|
||||
bank_id,
|
||||
uuid.UUID(entity_id),
|
||||
)
|
||||
entity = await app.state.memory.get_entity(bank_id, entity_id, request_context=request_context)
|
||||
|
||||
if not entity_row:
|
||||
if entity is None:
|
||||
raise HTTPException(status_code=404, detail=f"Entity {entity_id} not found")
|
||||
|
||||
# Get observations for the entity
|
||||
observations = await app.state.memory.get_entity_observations(bank_id, entity_id, limit=20)
|
||||
|
||||
return EntityDetailResponse(
|
||||
id=str(entity_row["id"]),
|
||||
canonical_name=entity_row["canonical_name"],
|
||||
mention_count=entity_row["mention_count"],
|
||||
first_seen=entity_row["first_seen"].isoformat() if entity_row["first_seen"] else None,
|
||||
last_seen=entity_row["last_seen"].isoformat() if entity_row["last_seen"] else None,
|
||||
metadata=_parse_metadata(entity_row["metadata"]),
|
||||
id=entity["id"],
|
||||
canonical_name=entity["canonical_name"],
|
||||
mention_count=entity["mention_count"],
|
||||
first_seen=entity["first_seen"],
|
||||
last_seen=entity["last_seen"],
|
||||
metadata=_parse_metadata(entity["metadata"]),
|
||||
observations=[
|
||||
EntityObservationResponse(text=obs.text, mentioned_at=obs.mentioned_at) for obs in observations
|
||||
EntityObservationResponse(text=obs.text, mentioned_at=obs.mentioned_at)
|
||||
for obs in entity["observations"]
|
||||
],
|
||||
)
|
||||
except HTTPException:
|
||||
@@ -1255,42 +1333,40 @@ def _register_routes(app: FastAPI):
|
||||
operation_id="regenerate_entity_observations",
|
||||
tags=["Entities"],
|
||||
)
|
||||
async def api_regenerate_entity_observations(bank_id: str, entity_id: str):
|
||||
async def api_regenerate_entity_observations(
|
||||
bank_id: str,
|
||||
entity_id: str,
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
):
|
||||
"""Regenerate observations for an entity."""
|
||||
try:
|
||||
# First get the entity metadata
|
||||
pool = await app.state.memory._get_pool()
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
entity_row = await conn.fetchrow(
|
||||
"""
|
||||
SELECT id, canonical_name, mention_count, first_seen, last_seen, metadata
|
||||
FROM entities
|
||||
WHERE bank_id = $1 AND id = $2
|
||||
""",
|
||||
bank_id,
|
||||
uuid.UUID(entity_id),
|
||||
)
|
||||
# Get the entity to verify it exists and get canonical_name
|
||||
entity = await app.state.memory.get_entity(bank_id, entity_id, request_context=request_context)
|
||||
|
||||
if not entity_row:
|
||||
if entity is None:
|
||||
raise HTTPException(status_code=404, detail=f"Entity {entity_id} not found")
|
||||
|
||||
# Regenerate observations
|
||||
await app.state.memory.regenerate_entity_observations(
|
||||
bank_id=bank_id, entity_id=entity_id, entity_name=entity_row["canonical_name"]
|
||||
bank_id=bank_id,
|
||||
entity_id=entity_id,
|
||||
entity_name=entity["canonical_name"],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Get updated observations
|
||||
observations = await app.state.memory.get_entity_observations(bank_id, entity_id, limit=20)
|
||||
# Get updated entity with new observations
|
||||
entity = await app.state.memory.get_entity(bank_id, entity_id, request_context=request_context)
|
||||
|
||||
return EntityDetailResponse(
|
||||
id=str(entity_row["id"]),
|
||||
canonical_name=entity_row["canonical_name"],
|
||||
mention_count=entity_row["mention_count"],
|
||||
first_seen=entity_row["first_seen"].isoformat() if entity_row["first_seen"] else None,
|
||||
last_seen=entity_row["last_seen"].isoformat() if entity_row["last_seen"] else None,
|
||||
metadata=_parse_metadata(entity_row["metadata"]),
|
||||
id=entity["id"],
|
||||
canonical_name=entity["canonical_name"],
|
||||
mention_count=entity["mention_count"],
|
||||
first_seen=entity["first_seen"],
|
||||
last_seen=entity["last_seen"],
|
||||
metadata=_parse_metadata(entity["metadata"]),
|
||||
observations=[
|
||||
EntityObservationResponse(text=obs.text, mentioned_at=obs.mentioned_at) for obs in observations
|
||||
EntityObservationResponse(text=obs.text, mentioned_at=obs.mentioned_at)
|
||||
for obs in entity["observations"]
|
||||
],
|
||||
)
|
||||
except HTTPException:
|
||||
@@ -1310,7 +1386,13 @@ def _register_routes(app: FastAPI):
|
||||
operation_id="list_documents",
|
||||
tags=["Documents"],
|
||||
)
|
||||
async def api_list_documents(bank_id: str, q: str | None = None, limit: int = 100, offset: int = 0):
|
||||
async def api_list_documents(
|
||||
bank_id: str,
|
||||
q: str | None = None,
|
||||
limit: int = 100,
|
||||
offset: int = 0,
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
):
|
||||
"""
|
||||
List documents for a memory bank with optional search.
|
||||
|
||||
@@ -1321,7 +1403,9 @@ def _register_routes(app: FastAPI):
|
||||
offset: Offset for pagination (default: 0)
|
||||
"""
|
||||
try:
|
||||
data = await app.state.memory.list_documents(bank_id=bank_id, search_query=q, limit=limit, offset=offset)
|
||||
data = await app.state.memory.list_documents(
|
||||
bank_id=bank_id, search_query=q, limit=limit, offset=offset, request_context=request_context
|
||||
)
|
||||
return data
|
||||
except Exception as e:
|
||||
import traceback
|
||||
@@ -1338,7 +1422,9 @@ def _register_routes(app: FastAPI):
|
||||
operation_id="get_document",
|
||||
tags=["Documents"],
|
||||
)
|
||||
async def api_get_document(bank_id: str, document_id: str):
|
||||
async def api_get_document(
|
||||
bank_id: str, document_id: str, request_context: RequestContext = Depends(get_request_context)
|
||||
):
|
||||
"""
|
||||
Get a specific document with its original text.
|
||||
|
||||
@@ -1347,7 +1433,7 @@ def _register_routes(app: FastAPI):
|
||||
document_id: Document ID (from path)
|
||||
"""
|
||||
try:
|
||||
document = await app.state.memory.get_document(document_id, bank_id)
|
||||
document = await app.state.memory.get_document(document_id, bank_id, request_context=request_context)
|
||||
if not document:
|
||||
raise HTTPException(status_code=404, detail="Document not found")
|
||||
return document
|
||||
@@ -1368,7 +1454,7 @@ def _register_routes(app: FastAPI):
|
||||
operation_id="get_chunk",
|
||||
tags=["Documents"],
|
||||
)
|
||||
async def api_get_chunk(chunk_id: str):
|
||||
async def api_get_chunk(chunk_id: str, request_context: RequestContext = Depends(get_request_context)):
|
||||
"""
|
||||
Get a specific chunk with its text.
|
||||
|
||||
@@ -1376,7 +1462,7 @@ def _register_routes(app: FastAPI):
|
||||
chunk_id: Chunk ID (from path, format: bank_id_document_id_chunk_index)
|
||||
"""
|
||||
try:
|
||||
chunk = await app.state.memory.get_chunk(chunk_id)
|
||||
chunk = await app.state.memory.get_chunk(chunk_id, request_context=request_context)
|
||||
if not chunk:
|
||||
raise HTTPException(status_code=404, detail="Chunk not found")
|
||||
return chunk
|
||||
@@ -1401,7 +1487,9 @@ def _register_routes(app: FastAPI):
|
||||
operation_id="delete_document",
|
||||
tags=["Documents"],
|
||||
)
|
||||
async def api_delete_document(bank_id: str, document_id: str):
|
||||
async def api_delete_document(
|
||||
bank_id: str, document_id: str, request_context: RequestContext = Depends(get_request_context)
|
||||
):
|
||||
"""
|
||||
Delete a document and all its associated memory units and links.
|
||||
|
||||
@@ -1410,7 +1498,7 @@ def _register_routes(app: FastAPI):
|
||||
document_id: Document ID to delete (from path)
|
||||
"""
|
||||
try:
|
||||
result = await app.state.memory.delete_document(document_id, bank_id)
|
||||
result = await app.state.memory.delete_document(document_id, bank_id, request_context=request_context)
|
||||
|
||||
if result["document_deleted"] == 0:
|
||||
raise HTTPException(status_code=404, detail="Document not found")
|
||||
@@ -1437,45 +1525,14 @@ def _register_routes(app: FastAPI):
|
||||
operation_id="list_operations",
|
||||
tags=["Operations"],
|
||||
)
|
||||
async def api_list_operations(bank_id: str):
|
||||
async def api_list_operations(bank_id: str, request_context: RequestContext = Depends(get_request_context)):
|
||||
"""List all async operations (pending and failed) for a memory bank."""
|
||||
try:
|
||||
pool = await app.state.memory._get_pool()
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
operations = await conn.fetch(
|
||||
"""
|
||||
SELECT operation_id, bank_id, operation_type, created_at, status, error_message, result_metadata
|
||||
FROM async_operations
|
||||
WHERE bank_id = $1
|
||||
ORDER BY created_at DESC
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
|
||||
def parse_metadata(metadata):
|
||||
"""Parse result_metadata which may be a string or dict."""
|
||||
if metadata is None:
|
||||
return {}
|
||||
if isinstance(metadata, str):
|
||||
return json.loads(metadata)
|
||||
return metadata
|
||||
|
||||
return {
|
||||
"bank_id": bank_id,
|
||||
"operations": [
|
||||
{
|
||||
"id": str(row["operation_id"]),
|
||||
"task_type": row["operation_type"],
|
||||
"items_count": parse_metadata(row["result_metadata"]).get("items_count", 0),
|
||||
"document_id": parse_metadata(row["result_metadata"]).get("document_id"),
|
||||
"created_at": row["created_at"].isoformat(),
|
||||
"status": row["status"],
|
||||
"error_message": row["error_message"],
|
||||
}
|
||||
for row in operations
|
||||
],
|
||||
}
|
||||
|
||||
operations = await app.state.memory.list_operations(bank_id, request_context=request_context)
|
||||
return {
|
||||
"bank_id": bank_id,
|
||||
"operations": operations,
|
||||
}
|
||||
except Exception as e:
|
||||
import traceback
|
||||
|
||||
@@ -1490,39 +1547,21 @@ def _register_routes(app: FastAPI):
|
||||
operation_id="cancel_operation",
|
||||
tags=["Operations"],
|
||||
)
|
||||
async def api_cancel_operation(bank_id: str, operation_id: str):
|
||||
async def api_cancel_operation(
|
||||
bank_id: str, operation_id: str, request_context: RequestContext = Depends(get_request_context)
|
||||
):
|
||||
"""Cancel a pending async operation."""
|
||||
try:
|
||||
# Validate UUID format
|
||||
try:
|
||||
op_uuid = uuid.UUID(operation_id)
|
||||
uuid.UUID(operation_id)
|
||||
except ValueError:
|
||||
raise HTTPException(status_code=400, detail=f"Invalid operation_id format: {operation_id}")
|
||||
|
||||
pool = await app.state.memory._get_pool()
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
# Check if operation exists and belongs to this memory bank
|
||||
result = await conn.fetchrow(
|
||||
"SELECT bank_id FROM async_operations WHERE operation_id = $1 AND bank_id = $2", op_uuid, bank_id
|
||||
)
|
||||
|
||||
if not result:
|
||||
raise HTTPException(
|
||||
status_code=404, detail=f"Operation {operation_id} not found for memory bank {bank_id}"
|
||||
)
|
||||
|
||||
# Delete the operation
|
||||
await conn.execute("DELETE FROM async_operations WHERE operation_id = $1", op_uuid)
|
||||
|
||||
return {
|
||||
"success": True,
|
||||
"message": f"Operation {operation_id} cancelled",
|
||||
"operation_id": operation_id,
|
||||
"bank_id": bank_id,
|
||||
}
|
||||
|
||||
except HTTPException:
|
||||
raise
|
||||
result = await app.state.memory.cancel_operation(bank_id, operation_id, request_context=request_context)
|
||||
return result
|
||||
except ValueError as e:
|
||||
raise HTTPException(status_code=404, detail=str(e))
|
||||
except Exception as e:
|
||||
import traceback
|
||||
|
||||
@@ -1538,10 +1577,10 @@ def _register_routes(app: FastAPI):
|
||||
operation_id="get_bank_profile",
|
||||
tags=["Banks"],
|
||||
)
|
||||
async def api_get_bank_profile(bank_id: str):
|
||||
async def api_get_bank_profile(bank_id: str, request_context: RequestContext = Depends(get_request_context)):
|
||||
"""Get memory bank profile (disposition + background)."""
|
||||
try:
|
||||
profile = await app.state.memory.get_bank_profile(bank_id)
|
||||
profile = await app.state.memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
# Convert DispositionTraits object to dict for Pydantic
|
||||
disposition_dict = (
|
||||
profile["disposition"].model_dump()
|
||||
@@ -1569,14 +1608,18 @@ def _register_routes(app: FastAPI):
|
||||
operation_id="update_bank_disposition",
|
||||
tags=["Banks"],
|
||||
)
|
||||
async def api_update_bank_disposition(bank_id: str, request: UpdateDispositionRequest):
|
||||
async def api_update_bank_disposition(
|
||||
bank_id: str, request: UpdateDispositionRequest, request_context: RequestContext = Depends(get_request_context)
|
||||
):
|
||||
"""Update bank disposition traits."""
|
||||
try:
|
||||
# Update disposition
|
||||
await app.state.memory.update_bank_disposition(bank_id, request.disposition.model_dump())
|
||||
await app.state.memory.update_bank_disposition(
|
||||
bank_id, request.disposition.model_dump(), request_context=request_context
|
||||
)
|
||||
|
||||
# Get updated profile
|
||||
profile = await app.state.memory.get_bank_profile(bank_id)
|
||||
profile = await app.state.memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
disposition_dict = (
|
||||
profile["disposition"].model_dump()
|
||||
if hasattr(profile["disposition"], "model_dump")
|
||||
@@ -1603,11 +1646,13 @@ def _register_routes(app: FastAPI):
|
||||
operation_id="add_bank_background",
|
||||
tags=["Banks"],
|
||||
)
|
||||
async def api_add_bank_background(bank_id: str, request: AddBackgroundRequest):
|
||||
async def api_add_bank_background(
|
||||
bank_id: str, request: AddBackgroundRequest, request_context: RequestContext = Depends(get_request_context)
|
||||
):
|
||||
"""Add or merge bank background information. Optionally infer disposition traits."""
|
||||
try:
|
||||
result = await app.state.memory.merge_bank_background(
|
||||
bank_id, request.content, update_disposition=request.update_disposition
|
||||
bank_id, request.content, update_disposition=request.update_disposition, request_context=request_context
|
||||
)
|
||||
|
||||
response = BackgroundResponse(background=result["background"])
|
||||
@@ -1630,51 +1675,31 @@ def _register_routes(app: FastAPI):
|
||||
operation_id="create_or_update_bank",
|
||||
tags=["Banks"],
|
||||
)
|
||||
async def api_create_or_update_bank(bank_id: str, request: CreateBankRequest):
|
||||
async def api_create_or_update_bank(
|
||||
bank_id: str, request: CreateBankRequest, request_context: RequestContext = Depends(get_request_context)
|
||||
):
|
||||
"""Create or update an agent with disposition and background."""
|
||||
try:
|
||||
# Get existing profile or create with defaults
|
||||
profile = await app.state.memory.get_bank_profile(bank_id)
|
||||
# Ensure bank exists by getting profile (auto-creates with defaults)
|
||||
await app.state.memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
|
||||
# Update name if provided
|
||||
if request.name is not None:
|
||||
pool = await app.state.memory._get_pool()
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
await conn.execute(
|
||||
"""
|
||||
UPDATE banks
|
||||
SET name = $2,
|
||||
updated_at = NOW()
|
||||
WHERE bank_id = $1
|
||||
""",
|
||||
bank_id,
|
||||
request.name,
|
||||
)
|
||||
profile["name"] = request.name
|
||||
# Update name and/or background if provided
|
||||
if request.name is not None or request.background is not None:
|
||||
await app.state.memory.update_bank(
|
||||
bank_id,
|
||||
name=request.name,
|
||||
background=request.background,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Update disposition if provided
|
||||
if request.disposition is not None:
|
||||
await app.state.memory.update_bank_disposition(bank_id, request.disposition.model_dump())
|
||||
profile["disposition"] = request.disposition.model_dump()
|
||||
|
||||
# Update background if provided (replace, not merge)
|
||||
if request.background is not None:
|
||||
pool = await app.state.memory._get_pool()
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
await conn.execute(
|
||||
"""
|
||||
UPDATE banks
|
||||
SET background = $2,
|
||||
updated_at = NOW()
|
||||
WHERE bank_id = $1
|
||||
""",
|
||||
bank_id,
|
||||
request.background,
|
||||
)
|
||||
profile["background"] = request.background
|
||||
await app.state.memory.update_bank_disposition(
|
||||
bank_id, request.disposition.model_dump(), request_context=request_context
|
||||
)
|
||||
|
||||
# Get final profile
|
||||
final_profile = await app.state.memory.get_bank_profile(bank_id)
|
||||
final_profile = await app.state.memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
disposition_dict = (
|
||||
final_profile["disposition"].model_dump()
|
||||
if hasattr(final_profile["disposition"], "model_dump")
|
||||
@@ -1702,10 +1727,10 @@ def _register_routes(app: FastAPI):
|
||||
operation_id="delete_bank",
|
||||
tags=["Banks"],
|
||||
)
|
||||
async def api_delete_bank(bank_id: str):
|
||||
async def api_delete_bank(bank_id: str, request_context: RequestContext = Depends(get_request_context)):
|
||||
"""Delete an entire memory bank and all its data."""
|
||||
try:
|
||||
result = await app.state.memory.delete_bank(bank_id)
|
||||
result = await app.state.memory.delete_bank(bank_id, request_context=request_context)
|
||||
return DeleteResponse(
|
||||
success=True,
|
||||
message=f"Bank '{bank_id}' and all associated data deleted successfully",
|
||||
@@ -1745,7 +1770,9 @@ def _register_routes(app: FastAPI):
|
||||
operation_id="retain_memories",
|
||||
tags=["Memory"],
|
||||
)
|
||||
async def api_retain(bank_id: str, request: RetainRequest):
|
||||
async def api_retain(
|
||||
bank_id: str, request: RetainRequest, request_context: RequestContext = Depends(get_request_context)
|
||||
):
|
||||
"""Retain memories with optional async processing."""
|
||||
metrics = get_metrics_collector()
|
||||
|
||||
@@ -1766,47 +1793,42 @@ def _register_routes(app: FastAPI):
|
||||
|
||||
if request.async_:
|
||||
# Async processing: queue task and return immediately
|
||||
operation_id = uuid.uuid4()
|
||||
|
||||
# Insert operation record into database
|
||||
pool = await app.state.memory._get_pool()
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO async_operations (operation_id, bank_id, operation_type, result_metadata)
|
||||
VALUES ($1, $2, $3, $4)
|
||||
""",
|
||||
operation_id,
|
||||
bank_id,
|
||||
"retain",
|
||||
json.dumps({"items_count": len(contents)}),
|
||||
)
|
||||
|
||||
# Submit task to background queue
|
||||
await app.state.memory._task_backend.submit_task(
|
||||
result = await app.state.memory.submit_async_retain(bank_id, contents, request_context=request_context)
|
||||
return RetainResponse.model_validate(
|
||||
{
|
||||
"type": "batch_retain",
|
||||
"operation_id": str(operation_id),
|
||||
"success": True,
|
||||
"bank_id": bank_id,
|
||||
"contents": contents,
|
||||
"items_count": result["items_count"],
|
||||
"async": True,
|
||||
}
|
||||
)
|
||||
|
||||
logging.info(
|
||||
f"Retain task queued for bank_id={bank_id}, {len(contents)} items, operation_id={operation_id}"
|
||||
)
|
||||
|
||||
return RetainResponse(success=True, bank_id=bank_id, items_count=len(contents), async_=True)
|
||||
else:
|
||||
# Synchronous processing: wait for completion (record metrics)
|
||||
with metrics.record_operation("retain", bank_id=bank_id):
|
||||
result = await app.state.memory.retain_batch_async(bank_id=bank_id, contents=contents)
|
||||
result = await app.state.memory.retain_batch_async(
|
||||
bank_id=bank_id, contents=contents, request_context=request_context
|
||||
)
|
||||
|
||||
return RetainResponse(success=True, bank_id=bank_id, items_count=len(contents), async_=False)
|
||||
return RetainResponse.model_validate(
|
||||
{"success": True, "bank_id": bank_id, "items_count": len(contents), "async": False}
|
||||
)
|
||||
except Exception as e:
|
||||
import traceback
|
||||
|
||||
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
|
||||
# Create a summary of the input for debugging
|
||||
input_summary = []
|
||||
for i, item in enumerate(request.items):
|
||||
content_preview = item.content[:100] + "..." if len(item.content) > 100 else item.content
|
||||
input_summary.append(
|
||||
f" [{i}] content={content_preview!r}, context={item.context}, timestamp={item.timestamp}"
|
||||
)
|
||||
input_debug = "\n".join(input_summary)
|
||||
|
||||
error_detail = (
|
||||
f"{str(e)}\n\n"
|
||||
f"Input ({len(request.items)} items):\n{input_debug}\n\n"
|
||||
f"Traceback:\n{traceback.format_exc()}"
|
||||
)
|
||||
logger.error(f"Error in /v1/default/banks/{bank_id}/memories (retain): {error_detail}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
@@ -1821,10 +1843,11 @@ def _register_routes(app: FastAPI):
|
||||
async def api_clear_bank_memories(
|
||||
bank_id: str,
|
||||
type: str | None = Query(None, description="Optional fact type filter (world, experience, opinion)"),
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
):
|
||||
"""Clear memories for a memory bank, optionally filtered by type."""
|
||||
try:
|
||||
await app.state.memory.delete_bank(bank_id, fact_type=type)
|
||||
await app.state.memory.delete_bank(bank_id, fact_type=type, request_context=request_context)
|
||||
|
||||
return DeleteResponse(success=True)
|
||||
except Exception as e:
|
||||
|
||||
@@ -9,6 +9,7 @@ from fastmcp import FastMCP
|
||||
|
||||
from hindsight_api import MemoryEngine
|
||||
from hindsight_api.engine.response_models import VALID_RECALL_FACT_TYPES
|
||||
from hindsight_api.models import RequestContext
|
||||
|
||||
# Configure logging from HINDSIGHT_API_LOG_LEVEL environment variable
|
||||
_log_level_str = os.environ.get("HINDSIGHT_API_LOG_LEVEL", "info").lower()
|
||||
@@ -67,7 +68,11 @@ def create_mcp_server(memory: MemoryEngine) -> FastMCP:
|
||||
"""
|
||||
try:
|
||||
bank_id = get_current_bank_id()
|
||||
await memory.retain_batch_async(bank_id=bank_id, contents=[{"content": content, "context": context}])
|
||||
if bank_id is None:
|
||||
return "Error: No bank_id configured"
|
||||
await memory.retain_batch_async(
|
||||
bank_id=bank_id, contents=[{"content": content, "context": context}], request_context=RequestContext()
|
||||
)
|
||||
return "Memory stored successfully"
|
||||
except Exception as e:
|
||||
logger.error(f"Error storing memory: {e}", exc_info=True)
|
||||
@@ -90,10 +95,16 @@ def create_mcp_server(memory: MemoryEngine) -> FastMCP:
|
||||
"""
|
||||
try:
|
||||
bank_id = get_current_bank_id()
|
||||
if bank_id is None:
|
||||
return "Error: No bank_id configured"
|
||||
from hindsight_api.engine.memory_engine import Budget
|
||||
|
||||
search_result = await memory.recall_async(
|
||||
bank_id=bank_id, query=query, fact_type=list(VALID_RECALL_FACT_TYPES), budget=Budget.LOW
|
||||
bank_id=bank_id,
|
||||
query=query,
|
||||
fact_type=list(VALID_RECALL_FACT_TYPES),
|
||||
budget=Budget.LOW,
|
||||
request_context=RequestContext(),
|
||||
)
|
||||
|
||||
results = [
|
||||
@@ -102,7 +113,7 @@ def create_mcp_server(memory: MemoryEngine) -> FastMCP:
|
||||
"text": fact.text,
|
||||
"type": fact.fact_type,
|
||||
"context": fact.context,
|
||||
"event_date": fact.event_date,
|
||||
"occurred_start": fact.occurred_start,
|
||||
}
|
||||
for fact in search_result.results[:max_results]
|
||||
]
|
||||
|
||||
@@ -31,6 +31,11 @@ ENV_LOG_LEVEL = "HINDSIGHT_API_LOG_LEVEL"
|
||||
ENV_MCP_ENABLED = "HINDSIGHT_API_MCP_ENABLED"
|
||||
ENV_GRAPH_RETRIEVER = "HINDSIGHT_API_GRAPH_RETRIEVER"
|
||||
ENV_MCP_LOCAL_BANK_ID = "HINDSIGHT_API_MCP_LOCAL_BANK_ID"
|
||||
ENV_MCP_INSTRUCTIONS = "HINDSIGHT_API_MCP_INSTRUCTIONS"
|
||||
|
||||
# Optimization flags
|
||||
ENV_SKIP_LLM_VERIFICATION = "HINDSIGHT_API_SKIP_LLM_VERIFICATION"
|
||||
ENV_LAZY_RERANKER = "HINDSIGHT_API_LAZY_RERANKER"
|
||||
|
||||
# Default values
|
||||
DEFAULT_DATABASE_URL = "pg0"
|
||||
@@ -50,6 +55,26 @@ DEFAULT_MCP_ENABLED = True
|
||||
DEFAULT_GRAPH_RETRIEVER = "bfs" # Options: "bfs", "mpfp"
|
||||
DEFAULT_MCP_LOCAL_BANK_ID = "mcp"
|
||||
|
||||
# Default MCP tool descriptions (can be customized via env vars)
|
||||
DEFAULT_MCP_RETAIN_DESCRIPTION = """Store important information to long-term memory.
|
||||
|
||||
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
|
||||
- Work context, projects, or responsibilities"""
|
||||
|
||||
DEFAULT_MCP_RECALL_DESCRIPTION = """Search memories 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"""
|
||||
|
||||
# Required embedding dimension for database schema
|
||||
EMBEDDING_DIMENSION = 384
|
||||
|
||||
@@ -86,6 +111,10 @@ class HindsightConfig:
|
||||
# Recall
|
||||
graph_retriever: str
|
||||
|
||||
# Optimization flags
|
||||
skip_llm_verification: bool
|
||||
lazy_reranker: bool
|
||||
|
||||
@classmethod
|
||||
def from_env(cls) -> "HindsightConfig":
|
||||
"""Create configuration from environment variables."""
|
||||
@@ -112,6 +141,9 @@ class HindsightConfig:
|
||||
mcp_enabled=os.getenv(ENV_MCP_ENABLED, str(DEFAULT_MCP_ENABLED)).lower() == "true",
|
||||
# Recall
|
||||
graph_retriever=os.getenv(ENV_GRAPH_RETRIEVER, DEFAULT_GRAPH_RETRIEVER),
|
||||
# Optimization flags
|
||||
skip_llm_verification=os.getenv(ENV_SKIP_LLM_VERIFICATION, "false").lower() == "true",
|
||||
lazy_reranker=os.getenv(ENV_LAZY_RERANKER, "false").lower() == "true",
|
||||
)
|
||||
|
||||
def get_llm_base_url(self) -> str:
|
||||
@@ -142,7 +174,9 @@ class HindsightConfig:
|
||||
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"
|
||||
level=self.get_python_log_level(),
|
||||
format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
|
||||
force=True, # Override any existing configuration
|
||||
)
|
||||
|
||||
def log_config(self) -> None:
|
||||
|
||||
@@ -0,0 +1,204 @@
|
||||
"""
|
||||
Daemon mode support for Hindsight API.
|
||||
|
||||
Provides idle timeout and lockfile management for running as a background daemon.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import fcntl
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Default daemon configuration
|
||||
DEFAULT_DAEMON_PORT = 8889
|
||||
DEFAULT_IDLE_TIMEOUT = 0 # 0 = no auto-exit (hindsight-embed passes its own timeout)
|
||||
LOCKFILE_PATH = Path.home() / ".hindsight" / "daemon.lock"
|
||||
DAEMON_LOG_PATH = Path.home() / ".hindsight" / "daemon.log"
|
||||
|
||||
|
||||
class IdleTimeoutMiddleware:
|
||||
"""ASGI middleware that tracks activity and exits after idle timeout."""
|
||||
|
||||
def __init__(self, app, idle_timeout: int = DEFAULT_IDLE_TIMEOUT):
|
||||
self.app = app
|
||||
self.idle_timeout = idle_timeout
|
||||
self.last_activity = time.time()
|
||||
self._checker_task = None
|
||||
|
||||
async def __call__(self, scope, receive, send):
|
||||
# Update activity timestamp on each request
|
||||
self.last_activity = time.time()
|
||||
await self.app(scope, receive, send)
|
||||
|
||||
def start_idle_checker(self):
|
||||
"""Start the background task that checks for idle timeout."""
|
||||
self._checker_task = asyncio.create_task(self._check_idle())
|
||||
|
||||
async def _check_idle(self):
|
||||
"""Background task that exits the process after idle timeout."""
|
||||
# If idle_timeout is 0, don't auto-exit
|
||||
if self.idle_timeout <= 0:
|
||||
return
|
||||
|
||||
while True:
|
||||
await asyncio.sleep(30) # Check every 30 seconds
|
||||
idle_time = time.time() - self.last_activity
|
||||
if idle_time > self.idle_timeout:
|
||||
logger.info(f"Idle timeout reached ({self.idle_timeout}s), shutting down daemon")
|
||||
# Give a moment for any in-flight requests
|
||||
await asyncio.sleep(1)
|
||||
os._exit(0)
|
||||
|
||||
|
||||
class DaemonLock:
|
||||
"""
|
||||
File-based lock to prevent multiple daemon instances.
|
||||
|
||||
Uses fcntl.flock for atomic locking on Unix systems.
|
||||
"""
|
||||
|
||||
def __init__(self, lockfile: Path = LOCKFILE_PATH):
|
||||
self.lockfile = lockfile
|
||||
self._fd = None
|
||||
|
||||
def acquire(self) -> bool:
|
||||
"""
|
||||
Try to acquire the daemon lock.
|
||||
|
||||
Returns True if lock acquired, False if another daemon is running.
|
||||
"""
|
||||
self.lockfile.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
try:
|
||||
self._fd = open(self.lockfile, "w")
|
||||
fcntl.flock(self._fd.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
|
||||
# Write PID for debugging
|
||||
self._fd.write(str(os.getpid()))
|
||||
self._fd.flush()
|
||||
return True
|
||||
except (IOError, OSError):
|
||||
# Lock is held by another process
|
||||
if self._fd:
|
||||
self._fd.close()
|
||||
self._fd = None
|
||||
return False
|
||||
|
||||
def release(self):
|
||||
"""Release the daemon lock."""
|
||||
if self._fd:
|
||||
try:
|
||||
fcntl.flock(self._fd.fileno(), fcntl.LOCK_UN)
|
||||
self._fd.close()
|
||||
except Exception:
|
||||
pass
|
||||
finally:
|
||||
self._fd = None
|
||||
# Remove lockfile
|
||||
try:
|
||||
self.lockfile.unlink()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def is_locked(self) -> bool:
|
||||
"""Check if the lock is held by another process."""
|
||||
if not self.lockfile.exists():
|
||||
return False
|
||||
|
||||
try:
|
||||
fd = open(self.lockfile, "r")
|
||||
fcntl.flock(fd.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
|
||||
# We got the lock, so no one else has it
|
||||
fcntl.flock(fd.fileno(), fcntl.LOCK_UN)
|
||||
fd.close()
|
||||
return False
|
||||
except (IOError, OSError):
|
||||
return True
|
||||
|
||||
def get_pid(self) -> int | None:
|
||||
"""Get the PID of the daemon holding the lock."""
|
||||
if not self.lockfile.exists():
|
||||
return None
|
||||
try:
|
||||
with open(self.lockfile, "r") as f:
|
||||
return int(f.read().strip())
|
||||
except (ValueError, IOError):
|
||||
return None
|
||||
|
||||
|
||||
def daemonize():
|
||||
"""
|
||||
Fork the current process into a background daemon.
|
||||
|
||||
Uses double-fork technique to properly detach from terminal.
|
||||
"""
|
||||
# First fork
|
||||
pid = os.fork()
|
||||
if pid > 0:
|
||||
# Parent exits
|
||||
sys.exit(0)
|
||||
|
||||
# Create new session
|
||||
os.setsid()
|
||||
|
||||
# Second fork to prevent zombie processes
|
||||
pid = os.fork()
|
||||
if pid > 0:
|
||||
sys.exit(0)
|
||||
|
||||
# Redirect standard file descriptors to log file
|
||||
DAEMON_LOG_PATH.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
sys.stdout.flush()
|
||||
sys.stderr.flush()
|
||||
|
||||
# Redirect stdin to /dev/null
|
||||
with open("/dev/null", "r") as devnull:
|
||||
os.dup2(devnull.fileno(), sys.stdin.fileno())
|
||||
|
||||
# Redirect stdout/stderr to log file
|
||||
log_fd = open(DAEMON_LOG_PATH, "a")
|
||||
os.dup2(log_fd.fileno(), sys.stdout.fileno())
|
||||
os.dup2(log_fd.fileno(), sys.stderr.fileno())
|
||||
|
||||
|
||||
def check_daemon_running(port: int = DEFAULT_DAEMON_PORT) -> bool:
|
||||
"""Check if a daemon is running and responsive on the given port."""
|
||||
import socket
|
||||
|
||||
try:
|
||||
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
||||
sock.settimeout(1)
|
||||
result = sock.connect_ex(("127.0.0.1", port))
|
||||
sock.close()
|
||||
return result == 0
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
def stop_daemon(port: int = DEFAULT_DAEMON_PORT) -> bool:
|
||||
"""Stop a running daemon by sending SIGTERM to the process."""
|
||||
lock = DaemonLock()
|
||||
pid = lock.get_pid()
|
||||
|
||||
if pid is None:
|
||||
return False
|
||||
|
||||
try:
|
||||
import signal
|
||||
|
||||
os.kill(pid, signal.SIGTERM)
|
||||
# Wait for process to exit
|
||||
for _ in range(50): # Wait up to 5 seconds
|
||||
time.sleep(0.1)
|
||||
try:
|
||||
os.kill(pid, 0) # Check if process exists
|
||||
except OSError:
|
||||
return True # Process exited
|
||||
return False
|
||||
except OSError:
|
||||
return False
|
||||
@@ -11,7 +11,13 @@ from .cross_encoder import CrossEncoderModel, LocalSTCrossEncoder, RemoteTEICros
|
||||
from .db_utils import acquire_with_retry
|
||||
from .embeddings import Embeddings, LocalSTEmbeddings, RemoteTEIEmbeddings
|
||||
from .llm_wrapper import LLMConfig
|
||||
from .memory_engine import MemoryEngine
|
||||
from .memory_engine import (
|
||||
MemoryEngine,
|
||||
UnqualifiedTableError,
|
||||
fq_table,
|
||||
get_current_schema,
|
||||
validate_sql_schema,
|
||||
)
|
||||
from .response_models import MemoryFact, RecallResult, ReflectResult
|
||||
from .search.trace import (
|
||||
EntryPoint,
|
||||
@@ -49,4 +55,9 @@ __all__ = [
|
||||
"RecallResult",
|
||||
"ReflectResult",
|
||||
"MemoryFact",
|
||||
# Schema safety utilities
|
||||
"fq_table",
|
||||
"get_current_schema",
|
||||
"validate_sql_schema",
|
||||
"UnqualifiedTableError",
|
||||
]
|
||||
|
||||
@@ -11,6 +11,7 @@ from difflib import SequenceMatcher
|
||||
import asyncpg
|
||||
|
||||
from .db_utils import acquire_with_retry
|
||||
from .memory_engine import fq_table
|
||||
|
||||
# Load spaCy model (singleton)
|
||||
_nlp = None
|
||||
@@ -68,9 +69,9 @@ class EntityResolver:
|
||||
) -> list[str]:
|
||||
# Query ALL candidates for this bank
|
||||
all_entities = await conn.fetch(
|
||||
"""
|
||||
f"""
|
||||
SELECT canonical_name, id, metadata, last_seen, mention_count
|
||||
FROM entities
|
||||
FROM {fq_table("entities")}
|
||||
WHERE bank_id = $1
|
||||
""",
|
||||
bank_id,
|
||||
@@ -82,11 +83,11 @@ class EntityResolver:
|
||||
# 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(
|
||||
"""
|
||||
f"""
|
||||
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)
|
||||
FROM {fq_table("entity_cooccurrences")} ec
|
||||
WHERE ec.entity_id_1 IN (SELECT id FROM {fq_table("entities")} WHERE bank_id = $1)
|
||||
OR ec.entity_id_2 IN (SELECT id FROM {fq_table("entities")} WHERE bank_id = $1)
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
@@ -195,8 +196,8 @@ class EntityResolver:
|
||||
# Batch update existing entities
|
||||
if entities_to_update:
|
||||
await conn.executemany(
|
||||
"""
|
||||
UPDATE entities SET
|
||||
f"""
|
||||
UPDATE {fq_table("entities")} SET
|
||||
mention_count = mention_count + 1,
|
||||
last_seen = $2
|
||||
WHERE id = $1::uuid
|
||||
@@ -232,13 +233,13 @@ class EntityResolver:
|
||||
# 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)
|
||||
f"""
|
||||
INSERT INTO {fq_table("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,
|
||||
mention_count = {fq_table("entities")}.mention_count + 1,
|
||||
last_seen = EXCLUDED.last_seen
|
||||
RETURNING id
|
||||
""",
|
||||
@@ -279,9 +280,9 @@ class EntityResolver:
|
||||
async with acquire_with_retry(self.pool) as conn:
|
||||
# Find candidate entities with similar name
|
||||
candidates = await conn.fetch(
|
||||
"""
|
||||
f"""
|
||||
SELECT id, canonical_name, metadata, last_seen
|
||||
FROM entities
|
||||
FROM {fq_table("entities")}
|
||||
WHERE bank_id = $1
|
||||
AND (
|
||||
canonical_name ILIKE $2
|
||||
@@ -326,10 +327,10 @@ class EntityResolver:
|
||||
# Get entities that co-occurred with this candidate before
|
||||
# Use the materialized co-occurrence cache for fast lookup
|
||||
co_entity_rows = await conn.fetch(
|
||||
"""
|
||||
f"""
|
||||
SELECT e.canonical_name, ec.cooccurrence_count
|
||||
FROM entity_cooccurrences ec
|
||||
JOIN entities e ON (
|
||||
FROM {fq_table("entity_cooccurrences")} ec
|
||||
JOIN {fq_table("entities")} e ON (
|
||||
CASE
|
||||
WHEN ec.entity_id_1 = $1 THEN ec.entity_id_2
|
||||
WHEN ec.entity_id_2 = $1 THEN ec.entity_id_1
|
||||
@@ -365,8 +366,8 @@ class EntityResolver:
|
||||
if best_score > threshold:
|
||||
# Update entity
|
||||
await conn.execute(
|
||||
"""
|
||||
UPDATE entities
|
||||
f"""
|
||||
UPDATE {fq_table("entities")}
|
||||
SET mention_count = mention_count + 1,
|
||||
last_seen = $1
|
||||
WHERE id = $2
|
||||
@@ -402,12 +403,12 @@ class EntityResolver:
|
||||
Entity ID
|
||||
"""
|
||||
entity_id = await conn.fetchval(
|
||||
"""
|
||||
INSERT INTO entities (bank_id, canonical_name, first_seen, last_seen, mention_count)
|
||||
f"""
|
||||
INSERT INTO {fq_table("entities")} (bank_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,
|
||||
mention_count = {fq_table("entities")}.mention_count + 1,
|
||||
last_seen = EXCLUDED.last_seen
|
||||
RETURNING id
|
||||
""",
|
||||
@@ -430,8 +431,8 @@ class EntityResolver:
|
||||
async with acquire_with_retry(self.pool) as conn:
|
||||
# Insert unit-entity link
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO unit_entities (unit_id, entity_id)
|
||||
f"""
|
||||
INSERT INTO {fq_table("unit_entities")} (unit_id, entity_id)
|
||||
VALUES ($1, $2)
|
||||
ON CONFLICT DO NOTHING
|
||||
""",
|
||||
@@ -441,9 +442,9 @@ class EntityResolver:
|
||||
|
||||
# Update co-occurrence cache: find other entities in this unit
|
||||
rows = await conn.fetch(
|
||||
"""
|
||||
f"""
|
||||
SELECT entity_id
|
||||
FROM unit_entities
|
||||
FROM {fq_table("unit_entities")}
|
||||
WHERE unit_id = $1 AND entity_id != $2
|
||||
""",
|
||||
unit_id,
|
||||
@@ -472,12 +473,12 @@ class EntityResolver:
|
||||
entity_id_1, entity_id_2 = entity_id_2, entity_id_1
|
||||
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO entity_cooccurrences (entity_id_1, entity_id_2, cooccurrence_count, last_cooccurred)
|
||||
f"""
|
||||
INSERT INTO {fq_table("entity_cooccurrences")} (entity_id_1, entity_id_2, cooccurrence_count, last_cooccurred)
|
||||
VALUES ($1, $2, 1, NOW())
|
||||
ON CONFLICT (entity_id_1, entity_id_2)
|
||||
DO UPDATE SET
|
||||
cooccurrence_count = entity_cooccurrences.cooccurrence_count + 1,
|
||||
cooccurrence_count = {fq_table("entity_cooccurrences")}.cooccurrence_count + 1,
|
||||
last_cooccurred = NOW()
|
||||
""",
|
||||
entity_id_1,
|
||||
@@ -506,8 +507,8 @@ class EntityResolver:
|
||||
async def _link_units_to_entities_batch_impl(self, conn, unit_entity_pairs: list[tuple[str, str]]):
|
||||
# Batch insert all unit-entity links
|
||||
await conn.executemany(
|
||||
"""
|
||||
INSERT INTO unit_entities (unit_id, entity_id)
|
||||
f"""
|
||||
INSERT INTO {fq_table("unit_entities")} (unit_id, entity_id)
|
||||
VALUES ($1, $2)
|
||||
ON CONFLICT DO NOTHING
|
||||
""",
|
||||
@@ -541,12 +542,12 @@ class EntityResolver:
|
||||
if cooccurrence_pairs:
|
||||
now = datetime.now(UTC)
|
||||
await conn.executemany(
|
||||
"""
|
||||
INSERT INTO entity_cooccurrences (entity_id_1, entity_id_2, cooccurrence_count, last_cooccurred)
|
||||
f"""
|
||||
INSERT INTO {fq_table("entity_cooccurrences")} (entity_id_1, entity_id_2, cooccurrence_count, last_cooccurred)
|
||||
VALUES ($1, $2, $3, $4)
|
||||
ON CONFLICT (entity_id_1, entity_id_2)
|
||||
DO UPDATE SET
|
||||
cooccurrence_count = entity_cooccurrences.cooccurrence_count + 1,
|
||||
cooccurrence_count = {fq_table("entity_cooccurrences")}.cooccurrence_count + 1,
|
||||
last_cooccurred = EXCLUDED.last_cooccurred
|
||||
""",
|
||||
[(e1, e2, 1, now) for e1, e2 in cooccurrence_pairs],
|
||||
@@ -565,9 +566,9 @@ class EntityResolver:
|
||||
"""
|
||||
async with acquire_with_retry(self.pool) as conn:
|
||||
rows = await conn.fetch(
|
||||
"""
|
||||
f"""
|
||||
SELECT unit_id
|
||||
FROM unit_entities
|
||||
FROM {fq_table("unit_entities")}
|
||||
WHERE entity_id = $1
|
||||
ORDER BY unit_id
|
||||
LIMIT $2
|
||||
@@ -594,8 +595,8 @@ class EntityResolver:
|
||||
"""
|
||||
async with acquire_with_retry(self.pool) as conn:
|
||||
row = await conn.fetchrow(
|
||||
"""
|
||||
SELECT id FROM entities
|
||||
f"""
|
||||
SELECT id FROM {fq_table("entities")}
|
||||
WHERE bank_id = $1
|
||||
AND canonical_name ILIKE $2
|
||||
ORDER BY mention_count DESC
|
||||
|
||||
@@ -0,0 +1,592 @@
|
||||
"""Abstract interface for MemoryEngine public methods.
|
||||
|
||||
This module defines the public API that HTTP endpoints and extensions should use
|
||||
to interact with the memory system. All methods require a RequestContext for
|
||||
authentication when a TenantExtension is configured.
|
||||
"""
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from datetime import datetime
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from hindsight_api.engine.memory_engine import Budget
|
||||
from hindsight_api.engine.response_models import RecallResult, ReflectResult
|
||||
from hindsight_api.models import RequestContext
|
||||
|
||||
|
||||
class MemoryEngineInterface(ABC):
|
||||
"""
|
||||
Abstract interface for the Memory Engine.
|
||||
|
||||
This defines the public API that should be used by HTTP endpoints and extensions.
|
||||
All methods require a RequestContext for authentication.
|
||||
"""
|
||||
|
||||
# =========================================================================
|
||||
# Health & Status
|
||||
# =========================================================================
|
||||
|
||||
@abstractmethod
|
||||
async def health_check(self) -> dict:
|
||||
"""
|
||||
Check the health of the memory system.
|
||||
|
||||
Returns:
|
||||
Dict with 'status' key ('healthy' or 'unhealthy') and additional info.
|
||||
"""
|
||||
...
|
||||
|
||||
# =========================================================================
|
||||
# Core Memory Operations
|
||||
# =========================================================================
|
||||
|
||||
@abstractmethod
|
||||
async def retain_batch_async(
|
||||
self,
|
||||
bank_id: str,
|
||||
contents: list[dict[str, Any]],
|
||||
*,
|
||||
request_context: "RequestContext",
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Retain a batch of memory items.
|
||||
|
||||
Args:
|
||||
bank_id: The memory bank ID.
|
||||
contents: List of content dicts with 'content', optional 'event_date',
|
||||
'context', 'metadata', 'document_id'.
|
||||
request_context: Request context for authentication.
|
||||
|
||||
Returns:
|
||||
Dict with processing results.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def recall_async(
|
||||
self,
|
||||
bank_id: str,
|
||||
query: str,
|
||||
*,
|
||||
budget: "Budget | None" = None,
|
||||
max_tokens: int = 4096,
|
||||
enable_trace: bool = False,
|
||||
fact_type: list[str] | None = None,
|
||||
question_date: datetime | None = None,
|
||||
include_entities: bool = False,
|
||||
max_entity_tokens: int = 500,
|
||||
include_chunks: bool = False,
|
||||
max_chunk_tokens: int = 8192,
|
||||
request_context: "RequestContext",
|
||||
) -> "RecallResult":
|
||||
"""
|
||||
Recall memories relevant to a query.
|
||||
|
||||
Args:
|
||||
bank_id: The memory bank ID.
|
||||
query: The search query.
|
||||
budget: Search budget (LOW, MID, HIGH).
|
||||
max_tokens: Maximum tokens in response.
|
||||
enable_trace: Include trace information.
|
||||
fact_type: Filter by fact types.
|
||||
question_date: Context date for temporal relevance.
|
||||
include_entities: Include entity observations.
|
||||
max_entity_tokens: Max tokens for entity observations.
|
||||
include_chunks: Include raw chunks.
|
||||
max_chunk_tokens: Max tokens for chunks.
|
||||
request_context: Request context for authentication.
|
||||
|
||||
Returns:
|
||||
RecallResult with matching memories.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def reflect_async(
|
||||
self,
|
||||
bank_id: str,
|
||||
query: str,
|
||||
*,
|
||||
budget: "Budget | None" = None,
|
||||
context: str | None = None,
|
||||
request_context: "RequestContext",
|
||||
) -> "ReflectResult":
|
||||
"""
|
||||
Reflect on a query and generate a thoughtful response.
|
||||
|
||||
Args:
|
||||
bank_id: The memory bank ID.
|
||||
query: The question to reflect on.
|
||||
budget: Search budget for retrieving context.
|
||||
context: Additional context for the reflection.
|
||||
request_context: Request context for authentication.
|
||||
|
||||
Returns:
|
||||
ReflectResult with generated response and supporting facts.
|
||||
"""
|
||||
...
|
||||
|
||||
# =========================================================================
|
||||
# Bank Management
|
||||
# =========================================================================
|
||||
|
||||
@abstractmethod
|
||||
async def list_banks(
|
||||
self,
|
||||
*,
|
||||
request_context: "RequestContext",
|
||||
) -> list[dict[str, Any]]:
|
||||
"""
|
||||
List all memory banks.
|
||||
|
||||
Args:
|
||||
request_context: Request context for authentication.
|
||||
|
||||
Returns:
|
||||
List of bank info dicts.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def get_bank_profile(
|
||||
self,
|
||||
bank_id: str,
|
||||
*,
|
||||
request_context: "RequestContext",
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Get bank profile including disposition and background.
|
||||
|
||||
Args:
|
||||
bank_id: The memory bank ID.
|
||||
request_context: Request context for authentication.
|
||||
|
||||
Returns:
|
||||
Bank profile dict.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def update_bank_disposition(
|
||||
self,
|
||||
bank_id: str,
|
||||
disposition: dict[str, int],
|
||||
*,
|
||||
request_context: "RequestContext",
|
||||
) -> None:
|
||||
"""
|
||||
Update bank disposition traits.
|
||||
|
||||
Args:
|
||||
bank_id: The memory bank ID.
|
||||
disposition: Dict with trait values.
|
||||
request_context: Request context for authentication.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def merge_bank_background(
|
||||
self,
|
||||
bank_id: str,
|
||||
new_info: str,
|
||||
*,
|
||||
update_disposition: bool = True,
|
||||
request_context: "RequestContext",
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Merge new background information into bank profile.
|
||||
|
||||
Args:
|
||||
bank_id: The memory bank ID.
|
||||
new_info: New background information to merge.
|
||||
update_disposition: Whether to infer disposition from background.
|
||||
request_context: Request context for authentication.
|
||||
|
||||
Returns:
|
||||
Updated background info.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def delete_bank(
|
||||
self,
|
||||
bank_id: str,
|
||||
*,
|
||||
fact_type: str | None = None,
|
||||
request_context: "RequestContext",
|
||||
) -> dict[str, int]:
|
||||
"""
|
||||
Delete a bank or its memories.
|
||||
|
||||
Args:
|
||||
bank_id: The memory bank ID.
|
||||
fact_type: If specified, only delete memories of this type.
|
||||
request_context: Request context for authentication.
|
||||
|
||||
Returns:
|
||||
Dict with deletion counts.
|
||||
"""
|
||||
...
|
||||
|
||||
# =========================================================================
|
||||
# Memory Units
|
||||
# =========================================================================
|
||||
|
||||
@abstractmethod
|
||||
async def list_memory_units(
|
||||
self,
|
||||
bank_id: str,
|
||||
*,
|
||||
fact_type: str | None = None,
|
||||
search_query: str | None = None,
|
||||
limit: int = 100,
|
||||
offset: int = 0,
|
||||
request_context: "RequestContext",
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
List memory units with pagination.
|
||||
|
||||
Args:
|
||||
bank_id: The memory bank ID.
|
||||
fact_type: Filter by fact type.
|
||||
search_query: Full-text search query.
|
||||
limit: Maximum results.
|
||||
offset: Pagination offset.
|
||||
request_context: Request context for authentication.
|
||||
|
||||
Returns:
|
||||
Dict with 'items', 'total', 'limit', 'offset'.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def delete_memory_unit(
|
||||
self,
|
||||
unit_id: str,
|
||||
*,
|
||||
request_context: "RequestContext",
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Delete a specific memory unit.
|
||||
|
||||
Args:
|
||||
unit_id: The memory unit ID.
|
||||
request_context: Request context for authentication.
|
||||
|
||||
Returns:
|
||||
Deletion result.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def get_graph_data(
|
||||
self,
|
||||
bank_id: str,
|
||||
*,
|
||||
fact_type: str | None = None,
|
||||
request_context: "RequestContext",
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Get graph data for visualization.
|
||||
|
||||
Args:
|
||||
bank_id: The memory bank ID.
|
||||
fact_type: Filter by fact type.
|
||||
request_context: Request context for authentication.
|
||||
|
||||
Returns:
|
||||
Dict with nodes, edges, table_rows, total_units.
|
||||
"""
|
||||
...
|
||||
|
||||
# =========================================================================
|
||||
# Documents
|
||||
# =========================================================================
|
||||
|
||||
@abstractmethod
|
||||
async def list_documents(
|
||||
self,
|
||||
bank_id: str,
|
||||
*,
|
||||
search_query: str | None = None,
|
||||
limit: int = 100,
|
||||
offset: int = 0,
|
||||
request_context: "RequestContext",
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
List documents with pagination.
|
||||
|
||||
Args:
|
||||
bank_id: The memory bank ID.
|
||||
search_query: Search query.
|
||||
limit: Maximum results.
|
||||
offset: Pagination offset.
|
||||
request_context: Request context for authentication.
|
||||
|
||||
Returns:
|
||||
Dict with 'items', 'total', 'limit', 'offset'.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def get_document(
|
||||
self,
|
||||
document_id: str,
|
||||
bank_id: str,
|
||||
*,
|
||||
request_context: "RequestContext",
|
||||
) -> dict[str, Any] | None:
|
||||
"""
|
||||
Get a specific document.
|
||||
|
||||
Args:
|
||||
document_id: The document ID.
|
||||
bank_id: The memory bank ID.
|
||||
request_context: Request context for authentication.
|
||||
|
||||
Returns:
|
||||
Document dict or None if not found.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def delete_document(
|
||||
self,
|
||||
document_id: str,
|
||||
bank_id: str,
|
||||
*,
|
||||
request_context: "RequestContext",
|
||||
) -> dict[str, int]:
|
||||
"""
|
||||
Delete a document and its memory units.
|
||||
|
||||
Args:
|
||||
document_id: The document ID.
|
||||
bank_id: The memory bank ID.
|
||||
request_context: Request context for authentication.
|
||||
|
||||
Returns:
|
||||
Dict with deletion counts.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def get_chunk(
|
||||
self,
|
||||
chunk_id: str,
|
||||
*,
|
||||
request_context: "RequestContext",
|
||||
) -> dict[str, Any] | None:
|
||||
"""
|
||||
Get a specific chunk.
|
||||
|
||||
Args:
|
||||
chunk_id: The chunk ID.
|
||||
request_context: Request context for authentication.
|
||||
|
||||
Returns:
|
||||
Chunk dict or None if not found.
|
||||
"""
|
||||
...
|
||||
|
||||
# =========================================================================
|
||||
# Entities
|
||||
# =========================================================================
|
||||
|
||||
@abstractmethod
|
||||
async def list_entities(
|
||||
self,
|
||||
bank_id: str,
|
||||
*,
|
||||
limit: int = 100,
|
||||
request_context: "RequestContext",
|
||||
) -> list[dict[str, Any]]:
|
||||
"""
|
||||
List entities for a bank.
|
||||
|
||||
Args:
|
||||
bank_id: The memory bank ID.
|
||||
limit: Maximum results.
|
||||
request_context: Request context for authentication.
|
||||
|
||||
Returns:
|
||||
List of entity dicts.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def get_entity_observations(
|
||||
self,
|
||||
bank_id: str,
|
||||
entity_id: str,
|
||||
*,
|
||||
limit: int = 10,
|
||||
request_context: "RequestContext",
|
||||
) -> list[Any]:
|
||||
"""
|
||||
Get observations for an entity.
|
||||
|
||||
Args:
|
||||
bank_id: The memory bank ID.
|
||||
entity_id: The entity ID.
|
||||
limit: Maximum observations.
|
||||
request_context: Request context for authentication.
|
||||
|
||||
Returns:
|
||||
List of EntityObservation objects.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def regenerate_entity_observations(
|
||||
self,
|
||||
bank_id: str,
|
||||
entity_id: str,
|
||||
entity_name: str,
|
||||
*,
|
||||
request_context: "RequestContext",
|
||||
) -> None:
|
||||
"""
|
||||
Regenerate observations for an entity.
|
||||
|
||||
Args:
|
||||
bank_id: The memory bank ID.
|
||||
entity_id: The entity ID.
|
||||
entity_name: The entity's canonical name.
|
||||
request_context: Request context for authentication.
|
||||
"""
|
||||
...
|
||||
|
||||
# =========================================================================
|
||||
# Statistics & Operations
|
||||
# =========================================================================
|
||||
|
||||
@abstractmethod
|
||||
async def get_bank_stats(
|
||||
self,
|
||||
bank_id: str,
|
||||
*,
|
||||
request_context: "RequestContext",
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Get statistics about memory nodes and links for a bank.
|
||||
|
||||
Args:
|
||||
bank_id: The memory bank ID.
|
||||
request_context: Request context for authentication.
|
||||
|
||||
Returns:
|
||||
Dict with node_counts, link_counts, link_counts_by_fact_type,
|
||||
link_breakdown, and operations stats.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def get_entity(
|
||||
self,
|
||||
bank_id: str,
|
||||
entity_id: str,
|
||||
*,
|
||||
request_context: "RequestContext",
|
||||
) -> dict[str, Any] | None:
|
||||
"""
|
||||
Get entity details including metadata and observations.
|
||||
|
||||
Args:
|
||||
bank_id: The memory bank ID.
|
||||
entity_id: The entity ID.
|
||||
request_context: Request context for authentication.
|
||||
|
||||
Returns:
|
||||
Entity dict with id, canonical_name, mention_count, first_seen,
|
||||
last_seen, metadata, and observations. None if not found.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def list_operations(
|
||||
self,
|
||||
bank_id: str,
|
||||
*,
|
||||
request_context: "RequestContext",
|
||||
) -> list[dict[str, Any]]:
|
||||
"""
|
||||
List async operations for a bank.
|
||||
|
||||
Args:
|
||||
bank_id: The memory bank ID.
|
||||
request_context: Request context for authentication.
|
||||
|
||||
Returns:
|
||||
List of operation dicts with id, task_type, status, etc.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def cancel_operation(
|
||||
self,
|
||||
bank_id: str,
|
||||
operation_id: str,
|
||||
*,
|
||||
request_context: "RequestContext",
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Cancel a pending async operation.
|
||||
|
||||
Args:
|
||||
bank_id: The memory bank ID.
|
||||
operation_id: The operation ID to cancel.
|
||||
request_context: Request context for authentication.
|
||||
|
||||
Returns:
|
||||
Dict with success status and message.
|
||||
|
||||
Raises:
|
||||
ValueError: If operation not found.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def update_bank(
|
||||
self,
|
||||
bank_id: str,
|
||||
*,
|
||||
name: str | None = None,
|
||||
background: str | None = None,
|
||||
request_context: "RequestContext",
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Update bank name and/or background.
|
||||
|
||||
Args:
|
||||
bank_id: The memory bank ID.
|
||||
name: New bank name (optional).
|
||||
background: New background text (optional, replaces existing).
|
||||
request_context: Request context for authentication.
|
||||
|
||||
Returns:
|
||||
Updated bank profile dict.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def submit_async_retain(
|
||||
self,
|
||||
bank_id: str,
|
||||
contents: list[dict[str, Any]],
|
||||
*,
|
||||
request_context: "RequestContext",
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Submit a batch retain operation to run asynchronously.
|
||||
|
||||
Args:
|
||||
bank_id: The memory bank ID.
|
||||
contents: List of content dicts to retain.
|
||||
request_context: Request context for authentication.
|
||||
|
||||
Returns:
|
||||
Dict with operation_id and items_count.
|
||||
"""
|
||||
...
|
||||
@@ -3,11 +3,13 @@ LLM wrapper for unified configuration across providers.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
from google import genai
|
||||
from google.genai import errors as genai_errors
|
||||
from google.genai import types as genai_types
|
||||
@@ -96,7 +98,7 @@ class LLMProvider:
|
||||
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._client = AsyncOpenAI(**client_kwargs) # type: ignore[invalid-argument-type] - dict kwargs
|
||||
self._gemini_client = None
|
||||
|
||||
async def verify_connection(self) -> None:
|
||||
@@ -112,7 +114,7 @@ class LLMProvider:
|
||||
)
|
||||
await self.call(
|
||||
messages=[{"role": "user", "content": "Say 'ok'"}],
|
||||
max_completion_tokens=10,
|
||||
max_completion_tokens=100,
|
||||
max_retries=2,
|
||||
initial_backoff=0.5,
|
||||
max_backoff=2.0,
|
||||
@@ -157,7 +159,6 @@ class LLMProvider:
|
||||
"""
|
||||
async with _global_llm_semaphore:
|
||||
start_time = time.time()
|
||||
import json
|
||||
|
||||
# Handle Gemini provider separately
|
||||
if self.provider == "gemini":
|
||||
@@ -165,6 +166,20 @@ class LLMProvider:
|
||||
messages, response_format, max_retries, initial_backoff, max_backoff, skip_validation, start_time
|
||||
)
|
||||
|
||||
# Handle Ollama with native API for structured output (better schema enforcement)
|
||||
if self.provider == "ollama" and response_format is not None:
|
||||
return await self._call_ollama_native(
|
||||
messages,
|
||||
response_format,
|
||||
max_completion_tokens,
|
||||
temperature,
|
||||
max_retries,
|
||||
initial_backoff,
|
||||
max_backoff,
|
||||
skip_validation,
|
||||
start_time,
|
||||
)
|
||||
|
||||
call_params = {
|
||||
"model": self.model,
|
||||
"messages": messages,
|
||||
@@ -172,7 +187,7 @@ class LLMProvider:
|
||||
|
||||
# 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"])
|
||||
is_reasoning_model = any(x in model_lower for x in ["gpt-5", "o1", "o3", "deepseek"])
|
||||
|
||||
# For GPT-4 and GPT-4.1 models, cap max_completion_tokens to 32000
|
||||
# For GPT-4o models, cap to 16384
|
||||
@@ -194,7 +209,7 @@ class LLMProvider:
|
||||
call_params["temperature"] = temperature
|
||||
|
||||
# Set reasoning_effort for reasoning models (OpenAI gpt-5, o1, o3)
|
||||
if is_reasoning_model and self.provider == "openai":
|
||||
if is_reasoning_model:
|
||||
call_params["reasoning_effort"] = self.reasoning_effort
|
||||
|
||||
# Provider-specific parameters
|
||||
@@ -203,7 +218,6 @@ class LLMProvider:
|
||||
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
|
||||
|
||||
@@ -228,7 +242,31 @@ class LLMProvider:
|
||||
response = await self._client.chat.completions.create(**call_params)
|
||||
|
||||
content = response.choices[0].message.content
|
||||
json_data = json.loads(content)
|
||||
|
||||
# Log raw LLM response for debugging JSON parse issues
|
||||
try:
|
||||
json_data = json.loads(content)
|
||||
except json.JSONDecodeError as json_err:
|
||||
# Truncate content for logging (first 500 and last 200 chars)
|
||||
content_preview = content[:500] if content else "<empty>"
|
||||
if content and len(content) > 700:
|
||||
content_preview = f"{content[:500]}...TRUNCATED...{content[-200:]}"
|
||||
logger.warning(
|
||||
f"JSON parse error from LLM response (attempt {attempt + 1}/{max_retries + 1}): {json_err}\n"
|
||||
f" Model: {self.provider}/{self.model}\n"
|
||||
f" Content length: {len(content) if content else 0} chars\n"
|
||||
f" Content preview: {content_preview!r}\n"
|
||||
f" Finish reason: {response.choices[0].finish_reason if response.choices else 'unknown'}"
|
||||
)
|
||||
# Retry on JSON parse errors - LLM may return valid JSON on next attempt
|
||||
if attempt < max_retries:
|
||||
backoff = min(initial_backoff * (2**attempt), max_backoff)
|
||||
await asyncio.sleep(backoff)
|
||||
last_exception = json_err
|
||||
continue
|
||||
else:
|
||||
logger.error(f"JSON parse error after {max_retries + 1} attempts, giving up")
|
||||
raise
|
||||
|
||||
if skip_validation:
|
||||
result = json_data
|
||||
@@ -301,6 +339,129 @@ class LLMProvider:
|
||||
raise last_exception
|
||||
raise RuntimeError("LLM call failed after all retries with no exception captured")
|
||||
|
||||
async def _call_ollama_native(
|
||||
self,
|
||||
messages: list[dict[str, str]],
|
||||
response_format: Any,
|
||||
max_completion_tokens: int | None,
|
||||
temperature: float | None,
|
||||
max_retries: int,
|
||||
initial_backoff: float,
|
||||
max_backoff: float,
|
||||
skip_validation: bool,
|
||||
start_time: float,
|
||||
) -> Any:
|
||||
"""
|
||||
Call Ollama using native API with JSON schema enforcement.
|
||||
|
||||
Ollama's native API supports passing a full JSON schema in the 'format' parameter,
|
||||
which provides better structured output control than the OpenAI-compatible API.
|
||||
"""
|
||||
# Get the JSON schema from the Pydantic model
|
||||
schema = response_format.model_json_schema() if hasattr(response_format, "model_json_schema") else None
|
||||
|
||||
# Build the base URL for Ollama's native API
|
||||
# Default OpenAI-compatible URL is http://localhost:11434/v1
|
||||
# Native API is at http://localhost:11434/api/chat
|
||||
base_url = self.base_url or "http://localhost:11434/v1"
|
||||
if base_url.endswith("/v1"):
|
||||
native_url = base_url[:-3] + "/api/chat"
|
||||
else:
|
||||
native_url = base_url.rstrip("/") + "/api/chat"
|
||||
|
||||
# Build request payload
|
||||
payload = {
|
||||
"model": self.model,
|
||||
"messages": messages,
|
||||
"stream": False,
|
||||
}
|
||||
|
||||
# Add schema as format parameter for structured output
|
||||
if schema:
|
||||
payload["format"] = schema
|
||||
|
||||
# Add optional parameters with optimized defaults for Ollama
|
||||
# Benchmarking shows num_ctx=16384 + num_batch=512 is optimal
|
||||
options = {
|
||||
"num_ctx": 16384, # 16k context window for larger prompts
|
||||
"num_batch": 512, # Optimal batch size for prompt processing
|
||||
}
|
||||
if max_completion_tokens:
|
||||
options["num_predict"] = max_completion_tokens
|
||||
if temperature is not None:
|
||||
options["temperature"] = temperature
|
||||
payload["options"] = options
|
||||
|
||||
last_exception = None
|
||||
|
||||
async with httpx.AsyncClient(timeout=300.0) as client:
|
||||
for attempt in range(max_retries + 1):
|
||||
try:
|
||||
response = await client.post(native_url, json=payload)
|
||||
response.raise_for_status()
|
||||
|
||||
result = response.json()
|
||||
content = result.get("message", {}).get("content", "")
|
||||
|
||||
# Parse JSON response
|
||||
try:
|
||||
json_data = json.loads(content)
|
||||
except json.JSONDecodeError as json_err:
|
||||
content_preview = content[:500] if content else "<empty>"
|
||||
if content and len(content) > 700:
|
||||
content_preview = f"{content[:500]}...TRUNCATED...{content[-200:]}"
|
||||
logger.warning(
|
||||
f"Ollama JSON parse error (attempt {attempt + 1}/{max_retries + 1}): {json_err}\n"
|
||||
f" Model: ollama/{self.model}\n"
|
||||
f" Content length: {len(content) if content else 0} chars\n"
|
||||
f" Content preview: {content_preview!r}"
|
||||
)
|
||||
if attempt < max_retries:
|
||||
backoff = min(initial_backoff * (2**attempt), max_backoff)
|
||||
await asyncio.sleep(backoff)
|
||||
last_exception = json_err
|
||||
continue
|
||||
else:
|
||||
raise
|
||||
|
||||
# Validate against Pydantic model or return raw JSON
|
||||
if skip_validation:
|
||||
return json_data
|
||||
else:
|
||||
return response_format.model_validate(json_data)
|
||||
|
||||
except httpx.HTTPStatusError as e:
|
||||
last_exception = e
|
||||
if attempt < max_retries:
|
||||
logger.warning(
|
||||
f"Ollama HTTP error (attempt {attempt + 1}/{max_retries + 1}): {e.response.status_code}"
|
||||
)
|
||||
backoff = min(initial_backoff * (2**attempt), max_backoff)
|
||||
await asyncio.sleep(backoff)
|
||||
continue
|
||||
else:
|
||||
logger.error(f"Ollama HTTP error after {max_retries + 1} attempts: {e}")
|
||||
raise
|
||||
|
||||
except httpx.RequestError as e:
|
||||
last_exception = e
|
||||
if attempt < max_retries:
|
||||
logger.warning(f"Ollama connection error (attempt {attempt + 1}/{max_retries + 1}): {e}")
|
||||
backoff = min(initial_backoff * (2**attempt), max_backoff)
|
||||
await asyncio.sleep(backoff)
|
||||
continue
|
||||
else:
|
||||
logger.error(f"Ollama connection error after {max_retries + 1} attempts: {e}")
|
||||
raise
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Unexpected error during Ollama call: {type(e).__name__}: {e}")
|
||||
raise
|
||||
|
||||
if last_exception:
|
||||
raise last_exception
|
||||
raise RuntimeError("Ollama call failed after all retries")
|
||||
|
||||
async def _call_gemini(
|
||||
self,
|
||||
messages: list[dict[str, str]],
|
||||
@@ -312,8 +473,6 @@ class LLMProvider:
|
||||
start_time: float,
|
||||
) -> Any:
|
||||
"""Handle Gemini-specific API calls."""
|
||||
import json
|
||||
|
||||
# Convert OpenAI-style messages to Gemini format
|
||||
system_instruction = None
|
||||
gemini_contents = []
|
||||
@@ -444,6 +603,8 @@ class LLMProvider:
|
||||
"""Create provider for memory operations from environment variables."""
|
||||
provider = os.getenv("HINDSIGHT_API_LLM_PROVIDER", "groq")
|
||||
api_key = os.getenv("HINDSIGHT_API_LLM_API_KEY")
|
||||
if not api_key:
|
||||
raise ValueError("HINDSIGHT_API_LLM_API_KEY environment variable is required")
|
||||
base_url = os.getenv("HINDSIGHT_API_LLM_BASE_URL", "")
|
||||
model = os.getenv("HINDSIGHT_API_LLM_MODEL", "openai/gpt-oss-120b")
|
||||
|
||||
@@ -454,6 +615,10 @@ class 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"))
|
||||
if not api_key:
|
||||
raise ValueError(
|
||||
"HINDSIGHT_API_LLM_API_KEY or HINDSIGHT_API_ANSWER_LLM_API_KEY environment variable is required"
|
||||
)
|
||||
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"))
|
||||
|
||||
@@ -464,6 +629,10 @@ class LLMProvider:
|
||||
"""Create provider for judge/evaluator operations. Falls back to memory config if not set."""
|
||||
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"))
|
||||
if not api_key:
|
||||
raise ValueError(
|
||||
"HINDSIGHT_API_LLM_API_KEY or HINDSIGHT_API_JUDGE_LLM_API_KEY environment variable is required"
|
||||
)
|
||||
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"))
|
||||
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -10,6 +10,7 @@ from typing import TypedDict
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from ..db_utils import acquire_with_retry
|
||||
from ..memory_engine import fq_table
|
||||
from ..response_models import DispositionTraits
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -51,9 +52,9 @@ async def get_bank_profile(pool, bank_id: str) -> BankProfile:
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
# Try to get existing bank
|
||||
row = await conn.fetchrow(
|
||||
"""
|
||||
f"""
|
||||
SELECT name, disposition, background
|
||||
FROM banks WHERE bank_id = $1
|
||||
FROM {fq_table("banks")} WHERE bank_id = $1
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
@@ -70,8 +71,8 @@ async def get_bank_profile(pool, bank_id: str) -> BankProfile:
|
||||
|
||||
# Bank doesn't exist, create with defaults
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO banks (bank_id, name, disposition, background)
|
||||
f"""
|
||||
INSERT INTO {fq_table("banks")} (bank_id, name, disposition, background)
|
||||
VALUES ($1, $2, $3::jsonb, $4)
|
||||
ON CONFLICT (bank_id) DO NOTHING
|
||||
""",
|
||||
@@ -98,8 +99,8 @@ async def update_bank_disposition(pool, bank_id: str, disposition: dict[str, int
|
||||
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
await conn.execute(
|
||||
"""
|
||||
UPDATE banks
|
||||
f"""
|
||||
UPDATE {fq_table("banks")}
|
||||
SET disposition = $2::jsonb,
|
||||
updated_at = NOW()
|
||||
WHERE bank_id = $1
|
||||
@@ -140,8 +141,8 @@ async def merge_bank_background(pool, llm_config, bank_id: str, new_info: str, u
|
||||
if inferred_disposition:
|
||||
# Update both background and disposition
|
||||
await conn.execute(
|
||||
"""
|
||||
UPDATE banks
|
||||
f"""
|
||||
UPDATE {fq_table("banks")}
|
||||
SET background = $2,
|
||||
disposition = $3::jsonb,
|
||||
updated_at = NOW()
|
||||
@@ -154,8 +155,8 @@ async def merge_bank_background(pool, llm_config, bank_id: str, new_info: str, u
|
||||
else:
|
||||
# Update only background
|
||||
await conn.execute(
|
||||
"""
|
||||
UPDATE banks
|
||||
f"""
|
||||
UPDATE {fq_table("banks")}
|
||||
SET background = $2,
|
||||
updated_at = NOW()
|
||||
WHERE bank_id = $1
|
||||
@@ -361,9 +362,9 @@ async def list_banks(pool) -> list:
|
||||
"""
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
rows = await conn.fetch(
|
||||
"""
|
||||
f"""
|
||||
SELECT bank_id, name, disposition, background, created_at, updated_at
|
||||
FROM banks
|
||||
FROM {fq_table("banks")}
|
||||
ORDER BY updated_at DESC
|
||||
"""
|
||||
)
|
||||
|
||||
@@ -6,6 +6,7 @@ Handles storage of document chunks in the database.
|
||||
|
||||
import logging
|
||||
|
||||
from ..memory_engine import fq_table
|
||||
from .types import ChunkMetadata
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -42,8 +43,8 @@ async def store_chunks_batch(conn, bank_id: str, document_id: str, chunks: list[
|
||||
|
||||
# Batch insert all chunks
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO chunks (chunk_id, document_id, bank_id, chunk_text, chunk_index)
|
||||
f"""
|
||||
INSERT INTO {fq_table("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,
|
||||
|
||||
@@ -7,6 +7,7 @@ Handles insertion of facts into the database.
|
||||
import json
|
||||
import logging
|
||||
|
||||
from ..memory_engine import fq_table
|
||||
from .types import ProcessedFact
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -67,8 +68,8 @@ async def insert_facts_batch(
|
||||
|
||||
# Batch insert all facts
|
||||
results = await conn.fetch(
|
||||
"""
|
||||
INSERT INTO memory_units (bank_id, text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
f"""
|
||||
INSERT INTO {fq_table("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[],
|
||||
@@ -107,8 +108,8 @@ async def ensure_bank_exists(conn, bank_id: str) -> None:
|
||||
bank_id: Bank identifier
|
||||
"""
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO banks (bank_id, disposition, background)
|
||||
f"""
|
||||
INSERT INTO {fq_table("banks")} (bank_id, disposition, background)
|
||||
VALUES ($1, $2::jsonb, $3)
|
||||
ON CONFLICT (bank_id) DO UPDATE
|
||||
SET updated_at = NOW()
|
||||
@@ -141,12 +142,14 @@ async def handle_document_tracking(
|
||||
# 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)
|
||||
await conn.fetchval(
|
||||
f"DELETE FROM {fq_table('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)
|
||||
f"""
|
||||
INSERT INTO {fq_table("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,
|
||||
|
||||
@@ -7,6 +7,7 @@ import time
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from uuid import UUID
|
||||
|
||||
from ..memory_engine import fq_table
|
||||
from .types import EntityLink
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -290,9 +291,9 @@ async def extract_entities_batch_optimized(
|
||||
|
||||
entity_id_list = [uuid.UUID(eid) if isinstance(eid, str) else eid for eid in all_entity_ids]
|
||||
rows = await conn.fetch(
|
||||
"""
|
||||
f"""
|
||||
SELECT entity_id, unit_id
|
||||
FROM unit_entities
|
||||
FROM {fq_table("unit_entities")}
|
||||
WHERE entity_id = ANY($1::uuid[])
|
||||
""",
|
||||
entity_id_list,
|
||||
@@ -413,9 +414,9 @@ async def create_temporal_links_batch_per_fact(
|
||||
# Get the event_date for each new unit
|
||||
fetch_dates_start = time_mod.time()
|
||||
rows = await conn.fetch(
|
||||
"""
|
||||
f"""
|
||||
SELECT id, event_date
|
||||
FROM memory_units
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE id::text = ANY($1)
|
||||
""",
|
||||
unit_ids,
|
||||
@@ -432,9 +433,9 @@ async def create_temporal_links_batch_per_fact(
|
||||
|
||||
fetch_neighbors_start = time_mod.time()
|
||||
all_candidates = await conn.fetch(
|
||||
"""
|
||||
f"""
|
||||
SELECT id, event_date
|
||||
FROM memory_units
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE bank_id = $1
|
||||
AND event_date BETWEEN $2 AND $3
|
||||
AND id::text != ALL($4)
|
||||
@@ -479,8 +480,8 @@ async def create_temporal_links_batch_per_fact(
|
||||
if links:
|
||||
insert_start = time_mod.time()
|
||||
await conn.executemany(
|
||||
"""
|
||||
INSERT INTO memory_links (from_unit_id, to_unit_id, link_type, weight, entity_id)
|
||||
f"""
|
||||
INSERT INTO {fq_table("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
|
||||
""",
|
||||
@@ -535,9 +536,9 @@ async def create_semantic_links_batch(
|
||||
# Fetch ALL existing units with embeddings in ONE query
|
||||
fetch_start = time_mod.time()
|
||||
all_existing = await conn.fetch(
|
||||
"""
|
||||
f"""
|
||||
SELECT id, embedding
|
||||
FROM memory_units
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE bank_id = $1
|
||||
AND embedding IS NOT NULL
|
||||
AND id::text != ALL($2)
|
||||
@@ -644,8 +645,8 @@ async def create_semantic_links_batch(
|
||||
if all_links:
|
||||
insert_start = time_mod.time()
|
||||
await conn.executemany(
|
||||
"""
|
||||
INSERT INTO memory_links (from_unit_id, to_unit_id, link_type, weight, entity_id)
|
||||
f"""
|
||||
INSERT INTO {fq_table("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
|
||||
""",
|
||||
@@ -721,8 +722,8 @@ async def insert_entity_links_batch(conn, links: list[EntityLink], chunk_size: i
|
||||
|
||||
# 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)
|
||||
await conn.execute(f"""
|
||||
INSERT INTO {fq_table("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
|
||||
@@ -808,8 +809,8 @@ async def create_causal_links_batch(
|
||||
insert_start = time_mod.time()
|
||||
try:
|
||||
await conn.executemany(
|
||||
"""
|
||||
INSERT INTO memory_links (from_unit_id, to_unit_id, link_type, weight, entity_id)
|
||||
f"""
|
||||
INSERT INTO {fq_table("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
|
||||
""",
|
||||
|
||||
@@ -9,6 +9,7 @@ import time
|
||||
import uuid
|
||||
from datetime import UTC, datetime
|
||||
|
||||
from ..memory_engine import fq_table
|
||||
from ..search import observation_utils
|
||||
from . import embedding_utils
|
||||
from .types import EntityLink
|
||||
@@ -75,8 +76,8 @@ async def regenerate_observations_batch(
|
||||
|
||||
# Batch query for entity names
|
||||
entity_rows = await conn.fetch(
|
||||
"""
|
||||
SELECT id, canonical_name FROM entities
|
||||
f"""
|
||||
SELECT id, canonical_name FROM {fq_table("entities")}
|
||||
WHERE id = ANY($1) AND bank_id = $2
|
||||
""",
|
||||
entity_uuids,
|
||||
@@ -86,10 +87,10 @@ async def regenerate_observations_batch(
|
||||
|
||||
# Batch query for fact counts
|
||||
fact_counts = await conn.fetch(
|
||||
"""
|
||||
f"""
|
||||
SELECT ue.entity_id, COUNT(*) as cnt
|
||||
FROM unit_entities ue
|
||||
JOIN memory_units mu ON ue.unit_id = mu.id
|
||||
FROM {fq_table("unit_entities")} ue
|
||||
JOIN {fq_table("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
|
||||
""",
|
||||
@@ -154,10 +155,10 @@ async def _regenerate_entity_observations(
|
||||
|
||||
# Get all facts mentioning this entity (exclude observations themselves)
|
||||
rows = await conn.fetch(
|
||||
"""
|
||||
f"""
|
||||
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
|
||||
FROM {fq_table("memory_units")} mu
|
||||
JOIN {fq_table("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')
|
||||
@@ -193,12 +194,12 @@ async def _regenerate_entity_observations(
|
||||
|
||||
# Delete old observations for this entity
|
||||
await conn.execute(
|
||||
"""
|
||||
DELETE FROM memory_units
|
||||
f"""
|
||||
DELETE FROM {fq_table("memory_units")}
|
||||
WHERE id IN (
|
||||
SELECT mu.id
|
||||
FROM memory_units mu
|
||||
JOIN unit_entities ue ON mu.id = ue.unit_id
|
||||
FROM {fq_table("memory_units")} mu
|
||||
JOIN {fq_table("unit_entities")} ue ON mu.id = ue.unit_id
|
||||
WHERE mu.bank_id = $1
|
||||
AND mu.fact_type = 'observation'
|
||||
AND ue.entity_id = $2
|
||||
@@ -217,8 +218,8 @@ async def _regenerate_entity_observations(
|
||||
|
||||
for obs_text, embedding in zip(observations, embeddings):
|
||||
result = await conn.fetchrow(
|
||||
"""
|
||||
INSERT INTO memory_units (
|
||||
f"""
|
||||
INSERT INTO {fq_table("memory_units")} (
|
||||
bank_id, text, embedding, context, event_date,
|
||||
occurred_start, occurred_end, mentioned_at,
|
||||
fact_type, access_count
|
||||
@@ -240,8 +241,8 @@ async def _regenerate_entity_observations(
|
||||
|
||||
# Link observation to entity
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO unit_entities (unit_id, entity_id)
|
||||
f"""
|
||||
INSERT INTO {fq_table("unit_entities")} (unit_id, entity_id)
|
||||
VALUES ($1, $2)
|
||||
""",
|
||||
uuid.UUID(obs_id),
|
||||
|
||||
@@ -8,7 +8,6 @@ import logging
|
||||
import time
|
||||
import uuid
|
||||
from datetime import UTC, datetime
|
||||
from typing import Any
|
||||
|
||||
from ..db_utils import acquire_with_retry
|
||||
from . import bank_utils
|
||||
@@ -29,7 +28,7 @@ from . import (
|
||||
link_creation,
|
||||
observation_regeneration,
|
||||
)
|
||||
from .types import ExtractedFact, ProcessedFact, RetainContent
|
||||
from .types import ExtractedFact, ProcessedFact, RetainContent, RetainContentDict
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -43,7 +42,7 @@ async def retain_batch(
|
||||
format_date_fn,
|
||||
duplicate_checker_fn,
|
||||
bank_id: str,
|
||||
contents_dicts: list[dict[str, Any]],
|
||||
contents_dicts: list[RetainContentDict],
|
||||
document_id: str | None = None,
|
||||
is_first_batch: bool = True,
|
||||
fact_type_override: str | None = None,
|
||||
@@ -107,6 +106,10 @@ async def retain_batch(
|
||||
)
|
||||
|
||||
if not extracted_facts:
|
||||
total_time = time.time() - start_time
|
||||
logger.info(
|
||||
f"RETAIN_BATCH COMPLETE: 0 facts extracted from {len(contents)} contents in {total_time:.3f}s (nothing to store)"
|
||||
)
|
||||
return [[] for _ in contents]
|
||||
|
||||
# Apply fact_type_override if provided
|
||||
|
||||
@@ -7,9 +7,33 @@ from content input to fact storage.
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import UTC, datetime
|
||||
from typing import TypedDict
|
||||
from uuid import UUID
|
||||
|
||||
|
||||
class RetainContentDict(TypedDict, total=False):
|
||||
"""Type definition for content items in retain_batch_async.
|
||||
|
||||
Fields:
|
||||
content: Text content to store (required)
|
||||
context: Context about the content (optional)
|
||||
event_date: When the content occurred (optional, defaults to now)
|
||||
metadata: Custom key-value metadata (optional)
|
||||
document_id: Document ID for this content item (optional)
|
||||
"""
|
||||
|
||||
content: str # Required
|
||||
context: str
|
||||
event_date: datetime
|
||||
metadata: dict[str, str]
|
||||
document_id: str
|
||||
|
||||
|
||||
def _now_utc() -> datetime:
|
||||
"""Factory function for default event_date."""
|
||||
return datetime.now(UTC)
|
||||
|
||||
|
||||
@dataclass
|
||||
class RetainContent:
|
||||
"""
|
||||
@@ -20,16 +44,9 @@ class RetainContent:
|
||||
|
||||
content: str
|
||||
context: str = ""
|
||||
event_date: datetime | None = None
|
||||
event_date: datetime = field(default_factory=_now_utc)
|
||||
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
|
||||
|
||||
self.event_date = datetime.now(UTC)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ChunkMetadata:
|
||||
|
||||
@@ -10,6 +10,7 @@ import logging
|
||||
from abc import ABC, abstractmethod
|
||||
|
||||
from ..db_utils import acquire_with_retry
|
||||
from ..memory_engine import fq_table
|
||||
from .types import RetrievalResult
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -139,11 +140,11 @@ class BFSGraphRetriever(GraphRetriever):
|
||||
|
||||
# Step 1: Find entry points
|
||||
entry_points = await conn.fetch(
|
||||
"""
|
||||
f"""
|
||||
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
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE bank_id = $2
|
||||
AND embedding IS NOT NULL
|
||||
AND fact_type = $3
|
||||
@@ -188,13 +189,13 @@ class BFSGraphRetriever(GraphRetriever):
|
||||
if batch_nodes and budget_remaining > 0:
|
||||
max_neighbors = len(batch_nodes) * 20
|
||||
neighbors = await conn.fetch(
|
||||
"""
|
||||
f"""
|
||||
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
|
||||
FROM {fq_table("memory_links")} ml
|
||||
JOIN {fq_table("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
|
||||
|
||||
@@ -20,6 +20,7 @@ from collections import defaultdict
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
from ..db_utils import acquire_with_retry
|
||||
from ..memory_engine import fq_table
|
||||
from .graph_retrieval import GraphRetriever
|
||||
from .types import RetrievalResult
|
||||
|
||||
@@ -217,10 +218,10 @@ async def load_typed_adjacency(pool, bank_id: str) -> TypedAdjacency:
|
||||
"""
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
rows = await conn.fetch(
|
||||
"""
|
||||
f"""
|
||||
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
|
||||
FROM {fq_table("memory_links")} ml
|
||||
JOIN {fq_table("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
|
||||
@@ -252,10 +253,10 @@ async def fetch_memory_units_by_ids(
|
||||
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
rows = await conn.fetch(
|
||||
"""
|
||||
f"""
|
||||
SELECT id, text, context, event_date, occurred_start, occurred_end,
|
||||
mentioned_at, access_count, embedding, fact_type, document_id, chunk_id
|
||||
FROM memory_units
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE id = ANY($1::uuid[])
|
||||
AND fact_type = $2
|
||||
""",
|
||||
@@ -418,9 +419,9 @@ class MPFPGraphRetriever(GraphRetriever):
|
||||
"""Fallback: find semantic seeds via embedding search."""
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
rows = await conn.fetch(
|
||||
"""
|
||||
f"""
|
||||
SELECT id, 1 - (embedding <=> $1::vector) AS similarity
|
||||
FROM memory_units
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE bank_id = $2
|
||||
AND embedding IS NOT NULL
|
||||
AND fact_type = $3
|
||||
|
||||
@@ -26,6 +26,23 @@ class CrossEncoderReranker:
|
||||
|
||||
cross_encoder = create_cross_encoder_from_env()
|
||||
self.cross_encoder = cross_encoder
|
||||
self._initialized = False
|
||||
|
||||
async def ensure_initialized(self):
|
||||
"""Ensure the cross-encoder model is initialized (for lazy initialization)."""
|
||||
if self._initialized:
|
||||
return
|
||||
|
||||
import asyncio
|
||||
|
||||
cross_encoder = self.cross_encoder
|
||||
# For local providers, run in thread pool to avoid blocking event loop
|
||||
if cross_encoder.provider_name == "local":
|
||||
loop = asyncio.get_event_loop()
|
||||
await loop.run_in_executor(None, lambda: asyncio.run(cross_encoder.initialize()))
|
||||
else:
|
||||
await cross_encoder.initialize()
|
||||
self._initialized = True
|
||||
|
||||
def rerank(self, query: str, candidates: list[MergedCandidate]) -> list[ScoredResult]:
|
||||
"""
|
||||
|
||||
@@ -16,6 +16,7 @@ from typing import Optional
|
||||
|
||||
from ...config import get_config
|
||||
from ..db_utils import acquire_with_retry
|
||||
from ..memory_engine import fq_table
|
||||
from .graph_retrieval import BFSGraphRetriever, GraphRetriever
|
||||
from .mpfp_retrieval import MPFPGraphRetriever
|
||||
from .types import RetrievalResult
|
||||
@@ -80,10 +81,10 @@ async def retrieve_semantic(
|
||||
List of RetrievalResult objects
|
||||
"""
|
||||
results = await conn.fetch(
|
||||
"""
|
||||
f"""
|
||||
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
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE bank_id = $2
|
||||
AND embedding IS NOT NULL
|
||||
AND fact_type = $3
|
||||
@@ -131,10 +132,10 @@ async def retrieve_bm25(conn, query_text: str, bank_id: str, fact_type: str, lim
|
||||
query_tsquery = " | ".join(tokens)
|
||||
|
||||
results = await conn.fetch(
|
||||
"""
|
||||
f"""
|
||||
SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, access_count, embedding, fact_type, document_id, chunk_id,
|
||||
ts_rank_cd(search_vector, to_tsquery('english', $1)) AS bm25_score
|
||||
FROM memory_units
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE bank_id = $2
|
||||
AND fact_type = $3
|
||||
AND search_vector @@ to_tsquery('english', $1)
|
||||
@@ -188,10 +189,10 @@ async def retrieve_temporal(
|
||||
end_date = end_date.replace(tzinfo=UTC)
|
||||
|
||||
entry_points = await conn.fetch(
|
||||
"""
|
||||
f"""
|
||||
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
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE bank_id = $2
|
||||
AND fact_type = $3
|
||||
AND embedding IS NOT NULL
|
||||
@@ -272,12 +273,12 @@ async def retrieve_temporal(
|
||||
# Get neighbors via temporal and causal links
|
||||
if budget_remaining > 0:
|
||||
neighbors = await conn.fetch(
|
||||
"""
|
||||
f"""
|
||||
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,
|
||||
ml.weight, ml.link_type,
|
||||
1 - (mu.embedding <=> $1::vector) AS similarity
|
||||
FROM memory_links ml
|
||||
JOIN memory_units mu ON ml.to_unit_id = mu.id
|
||||
FROM {fq_table("memory_links")} ml
|
||||
JOIN {fq_table("memory_units")} mu ON ml.to_unit_id = mu.id
|
||||
WHERE ml.from_unit_id = $2
|
||||
AND ml.link_type IN ('temporal', 'causes', 'caused_by', 'enables', 'prevents')
|
||||
AND ml.weight >= 0.1
|
||||
@@ -546,11 +547,11 @@ async def _get_temporal_entry_points(
|
||||
end_date = end_date.replace(tzinfo=UTC)
|
||||
|
||||
rows = await conn.fetch(
|
||||
"""
|
||||
f"""
|
||||
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
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE bank_id = $2
|
||||
AND fact_type = $3
|
||||
AND embedding IS NOT NULL
|
||||
|
||||
@@ -101,7 +101,7 @@ def build_think_prompt(
|
||||
name: str,
|
||||
disposition: DispositionTraits,
|
||||
background: str,
|
||||
context: str = None,
|
||||
context: str | None = None,
|
||||
) -> str:
|
||||
"""Build the think prompt for the LLM."""
|
||||
disposition_desc = build_disposition_description(disposition)
|
||||
|
||||
@@ -115,7 +115,7 @@ class SearchTracer:
|
||||
node_id: str,
|
||||
text: str,
|
||||
context: str,
|
||||
event_date: datetime,
|
||||
event_date: datetime | None,
|
||||
access_count: int,
|
||||
is_entry_point: bool,
|
||||
parent_node_id: str | None,
|
||||
|
||||
@@ -89,6 +89,38 @@ class TaskBackend(ABC):
|
||||
traceback.print_exc()
|
||||
|
||||
|
||||
class SyncTaskBackend(TaskBackend):
|
||||
"""
|
||||
Synchronous task backend that executes tasks immediately.
|
||||
|
||||
This is useful for embedded/CLI usage where we don't want background
|
||||
workers that prevent clean exit. Tasks are executed inline rather than
|
||||
being queued.
|
||||
"""
|
||||
|
||||
async def initialize(self):
|
||||
"""No-op for sync backend."""
|
||||
self._initialized = True
|
||||
logger.debug("SyncTaskBackend initialized")
|
||||
|
||||
async def submit_task(self, task_dict: dict[str, Any]):
|
||||
"""
|
||||
Execute the task immediately (synchronously).
|
||||
|
||||
Args:
|
||||
task_dict: Task dictionary to execute
|
||||
"""
|
||||
if not self._initialized:
|
||||
await self.initialize()
|
||||
|
||||
await self._execute_task(task_dict)
|
||||
|
||||
async def shutdown(self):
|
||||
"""No-op for sync backend."""
|
||||
self._initialized = False
|
||||
logger.debug("SyncTaskBackend shutdown")
|
||||
|
||||
|
||||
class AsyncIOQueueBackend(TaskBackend):
|
||||
"""
|
||||
Task backend implementation using asyncio queues.
|
||||
|
||||
@@ -0,0 +1,66 @@
|
||||
"""
|
||||
Hindsight Extensions System.
|
||||
|
||||
Extensions allow customizing and extending Hindsight behavior without modifying core code.
|
||||
Extensions are loaded via environment variables pointing to implementation classes.
|
||||
|
||||
Example:
|
||||
HINDSIGHT_API_OPERATION_VALIDATOR_EXTENSION=mypackage.validators:MyValidator
|
||||
HINDSIGHT_API_OPERATION_VALIDATOR_MAX_RETRIES=3
|
||||
|
||||
HINDSIGHT_API_HTTP_EXTENSION=mypackage.http:MyHttpExtension
|
||||
HINDSIGHT_API_HTTP_SOME_CONFIG=value
|
||||
|
||||
Extensions receive an ExtensionContext that provides a controlled API for interacting
|
||||
with the system (e.g., running migrations for tenant schemas).
|
||||
"""
|
||||
|
||||
from hindsight_api.extensions.base import Extension
|
||||
from hindsight_api.extensions.builtin import ApiKeyTenantExtension
|
||||
from hindsight_api.extensions.context import DefaultExtensionContext, ExtensionContext
|
||||
from hindsight_api.extensions.http import HttpExtension
|
||||
from hindsight_api.extensions.loader import load_extension
|
||||
from hindsight_api.extensions.operation_validator import (
|
||||
OperationValidationError,
|
||||
OperationValidatorExtension,
|
||||
RecallContext,
|
||||
RecallResult,
|
||||
ReflectContext,
|
||||
ReflectResultContext,
|
||||
RetainContext,
|
||||
RetainResult,
|
||||
ValidationResult,
|
||||
)
|
||||
from hindsight_api.extensions.tenant import (
|
||||
AuthenticationError,
|
||||
TenantContext,
|
||||
TenantExtension,
|
||||
)
|
||||
from hindsight_api.models import RequestContext
|
||||
|
||||
__all__ = [
|
||||
# Base
|
||||
"Extension",
|
||||
"load_extension",
|
||||
# Context
|
||||
"ExtensionContext",
|
||||
"DefaultExtensionContext",
|
||||
# HTTP Extension
|
||||
"HttpExtension",
|
||||
# Operation Validator
|
||||
"OperationValidationError",
|
||||
"OperationValidatorExtension",
|
||||
"RecallContext",
|
||||
"RecallResult",
|
||||
"ReflectContext",
|
||||
"ReflectResultContext",
|
||||
"RetainContext",
|
||||
"RetainResult",
|
||||
"ValidationResult",
|
||||
# Tenant/Auth
|
||||
"ApiKeyTenantExtension",
|
||||
"AuthenticationError",
|
||||
"RequestContext",
|
||||
"TenantContext",
|
||||
"TenantExtension",
|
||||
]
|
||||
@@ -0,0 +1,81 @@
|
||||
"""Base Extension class for all Hindsight extensions."""
|
||||
|
||||
from abc import ABC
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from hindsight_api.extensions.context import ExtensionContext
|
||||
|
||||
|
||||
class Extension(ABC):
|
||||
"""
|
||||
Base class for all Hindsight extensions.
|
||||
|
||||
Extensions are loaded via environment variables and receive configuration
|
||||
from prefixed environment variables.
|
||||
|
||||
Example:
|
||||
HINDSIGHT_API_MY_EXTENSION=mypackage.ext:MyExtension
|
||||
HINDSIGHT_API_MY_SOME_CONFIG=value
|
||||
|
||||
The extension receives: {"some_config": "value"}
|
||||
|
||||
Extensions also receive an ExtensionContext that provides a controlled API
|
||||
for interacting with the system (e.g., running migrations for tenant schemas).
|
||||
"""
|
||||
|
||||
def __init__(self, config: dict[str, str]):
|
||||
"""
|
||||
Initialize the extension with configuration.
|
||||
|
||||
Args:
|
||||
config: Dictionary of configuration values from environment variables.
|
||||
Keys are lowercased with the prefix stripped.
|
||||
"""
|
||||
self.config = config
|
||||
self._context: "ExtensionContext | None" = None
|
||||
|
||||
def set_context(self, context: "ExtensionContext") -> None:
|
||||
"""
|
||||
Set the extension context.
|
||||
|
||||
Called by the extension loader after instantiation.
|
||||
Extensions should not call this directly.
|
||||
|
||||
Args:
|
||||
context: The ExtensionContext providing system APIs.
|
||||
"""
|
||||
self._context = context
|
||||
|
||||
@property
|
||||
def context(self) -> "ExtensionContext":
|
||||
"""
|
||||
Get the extension context.
|
||||
|
||||
Returns:
|
||||
The ExtensionContext providing system APIs.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If context has not been set yet.
|
||||
"""
|
||||
if self._context is None:
|
||||
raise RuntimeError(
|
||||
"Extension context not set. Context is available after the extension is loaded by the system."
|
||||
)
|
||||
return self._context
|
||||
|
||||
async def on_startup(self) -> None:
|
||||
"""
|
||||
Called when the application starts.
|
||||
|
||||
Override to perform initialization tasks like connecting to external services.
|
||||
"""
|
||||
pass
|
||||
|
||||
async def on_shutdown(self) -> None:
|
||||
"""
|
||||
Called when the application shuts down.
|
||||
|
||||
Override to perform cleanup tasks like closing connections.
|
||||
"""
|
||||
pass
|
||||
@@ -0,0 +1,18 @@
|
||||
"""
|
||||
Built-in extension implementations.
|
||||
|
||||
These are ready-to-use implementations of the extension interfaces.
|
||||
They can be used directly or serve as examples for custom implementations.
|
||||
|
||||
Available built-in extensions:
|
||||
- ApiKeyTenantExtension: Simple API key validation with public schema
|
||||
|
||||
Example usage:
|
||||
HINDSIGHT_API_TENANT_EXTENSION=hindsight_api.extensions.builtin.tenant:ApiKeyTenantExtension
|
||||
"""
|
||||
|
||||
from hindsight_api.extensions.builtin.tenant import ApiKeyTenantExtension
|
||||
|
||||
__all__ = [
|
||||
"ApiKeyTenantExtension",
|
||||
]
|
||||
@@ -0,0 +1,33 @@
|
||||
"""Built-in tenant extension implementations."""
|
||||
|
||||
from hindsight_api.extensions.tenant import AuthenticationError, TenantContext, TenantExtension
|
||||
from hindsight_api.models import RequestContext
|
||||
|
||||
|
||||
class ApiKeyTenantExtension(TenantExtension):
|
||||
"""
|
||||
Built-in tenant extension that validates API key against an environment variable.
|
||||
|
||||
This is a simple implementation that:
|
||||
1. Validates the API key matches HINDSIGHT_API_TENANT_API_KEY
|
||||
2. Returns 'public' as the schema for all authenticated requests
|
||||
|
||||
Configuration:
|
||||
HINDSIGHT_API_TENANT_EXTENSION=hindsight_api.extensions.builtin.tenant:ApiKeyTenantExtension
|
||||
HINDSIGHT_API_TENANT_API_KEY=your-secret-key
|
||||
|
||||
For multi-tenant setups with separate schemas per tenant, implement a custom
|
||||
TenantExtension that looks up the schema based on the API key or token claims.
|
||||
"""
|
||||
|
||||
def __init__(self, config: dict[str, str]):
|
||||
super().__init__(config)
|
||||
self.expected_api_key = config.get("api_key")
|
||||
if not self.expected_api_key:
|
||||
raise ValueError("HINDSIGHT_API_TENANT_API_KEY is required when using ApiKeyTenantExtension")
|
||||
|
||||
async def authenticate(self, context: RequestContext) -> TenantContext:
|
||||
"""Validate API key and return public schema context."""
|
||||
if context.api_key != self.expected_api_key:
|
||||
raise AuthenticationError("Invalid API key")
|
||||
return TenantContext(schema_name="public")
|
||||
@@ -0,0 +1,110 @@
|
||||
"""Extension context providing a controlled API for extensions to interact with the system."""
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from hindsight_api.engine.interface import MemoryEngineInterface
|
||||
|
||||
|
||||
class ExtensionContext(ABC):
|
||||
"""
|
||||
Abstract context providing a controlled API for extensions.
|
||||
|
||||
Extensions receive this context instead of direct access to internal
|
||||
components like MemoryEngine or database connections. This provides:
|
||||
- A stable API that won't break when internals change
|
||||
- Security by limiting what extensions can access
|
||||
- Clear documentation of what extensions can do
|
||||
|
||||
Built-in implementation:
|
||||
hindsight_api.extensions.builtin.context.DefaultExtensionContext
|
||||
|
||||
Example usage in an extension:
|
||||
class MyTenantExtension(TenantExtension):
|
||||
async def on_startup(self) -> None:
|
||||
# Run migrations for a new tenant schema
|
||||
await self.context.run_migration("tenant_acme")
|
||||
|
||||
class MyHttpExtension(HttpExtension):
|
||||
def get_router(self, memory):
|
||||
# Use memory engine for custom endpoints
|
||||
engine = self.context.get_memory_engine()
|
||||
...
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
async def run_migration(self, schema: str) -> None:
|
||||
"""
|
||||
Run database migrations for a specific schema.
|
||||
|
||||
This creates the schema if it doesn't exist and runs all pending
|
||||
migrations. Uses advisory locks to coordinate between distributed workers.
|
||||
|
||||
Args:
|
||||
schema: PostgreSQL schema name (e.g., "tenant_acme").
|
||||
The schema will be created if it doesn't exist.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If migrations fail to complete.
|
||||
|
||||
Example:
|
||||
# Provision a new tenant schema
|
||||
await context.run_migration("tenant_acme")
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
def get_memory_engine(self) -> "MemoryEngineInterface":
|
||||
"""
|
||||
Get the memory engine interface.
|
||||
|
||||
Returns the MemoryEngineInterface for performing memory operations
|
||||
like retain, recall, reflect, and entity/document management.
|
||||
|
||||
Returns:
|
||||
MemoryEngineInterface instance.
|
||||
|
||||
Example:
|
||||
engine = context.get_memory_engine()
|
||||
result = await engine.recall_async(bank_id, query)
|
||||
"""
|
||||
...
|
||||
|
||||
|
||||
class DefaultExtensionContext(ExtensionContext):
|
||||
"""
|
||||
Default implementation of ExtensionContext.
|
||||
|
||||
Uses the system's database URL and migration infrastructure.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
database_url: str,
|
||||
memory_engine: "MemoryEngineInterface | None" = None,
|
||||
):
|
||||
"""
|
||||
Initialize the context.
|
||||
|
||||
Args:
|
||||
database_url: SQLAlchemy database URL for migrations.
|
||||
memory_engine: Optional MemoryEngine instance for memory operations.
|
||||
"""
|
||||
self._database_url = database_url
|
||||
self._memory_engine = memory_engine
|
||||
|
||||
async def run_migration(self, schema: str) -> None:
|
||||
"""Run migrations for a specific schema."""
|
||||
from hindsight_api.migrations import run_migrations
|
||||
|
||||
run_migrations(self._database_url, schema=schema)
|
||||
|
||||
def get_memory_engine(self) -> "MemoryEngineInterface":
|
||||
"""Get the memory engine interface."""
|
||||
if self._memory_engine is None:
|
||||
raise RuntimeError(
|
||||
"Memory engine not configured in ExtensionContext. "
|
||||
"Ensure the context was created with a memory_engine parameter."
|
||||
)
|
||||
return self._memory_engine
|
||||
@@ -0,0 +1,89 @@
|
||||
"""
|
||||
HTTP Extension for adding custom endpoints to the Hindsight API.
|
||||
|
||||
This extension allows adding custom HTTP endpoints under the /ext/ path prefix.
|
||||
The extension provides a FastAPI router that is mounted on the main application.
|
||||
"""
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from fastapi import APIRouter
|
||||
|
||||
from hindsight_api.extensions.base import Extension
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from hindsight_api import MemoryEngine
|
||||
|
||||
|
||||
class HttpExtension(Extension, ABC):
|
||||
"""
|
||||
Base class for HTTP extensions that add custom API endpoints.
|
||||
|
||||
HTTP extensions provide a FastAPI router that gets mounted under /ext/.
|
||||
The extension has full control over the routes, request/response models, and handlers.
|
||||
|
||||
Example:
|
||||
```python
|
||||
from fastapi import APIRouter
|
||||
from hindsight_api.extensions import HttpExtension
|
||||
|
||||
class MyHttpExtension(HttpExtension):
|
||||
def get_router(self, memory: MemoryEngine) -> APIRouter:
|
||||
router = APIRouter()
|
||||
|
||||
@router.get("/hello")
|
||||
async def hello():
|
||||
return {"message": "Hello from extension!"}
|
||||
|
||||
@router.post("/custom/{bank_id}/action")
|
||||
async def custom_action(bank_id: str):
|
||||
# Access memory engine for database operations
|
||||
pool = await memory._get_pool()
|
||||
# ... custom logic
|
||||
return {"status": "ok"}
|
||||
|
||||
return router
|
||||
```
|
||||
|
||||
The routes will be available at:
|
||||
- GET /ext/hello
|
||||
- POST /ext/custom/{bank_id}/action
|
||||
|
||||
Configuration via environment variables:
|
||||
HINDSIGHT_API_HTTP_EXTENSION=mypackage.ext:MyHttpExtension
|
||||
HINDSIGHT_API_HTTP_SOME_CONFIG=value
|
||||
|
||||
The extension receives config: {"some_config": "value"}
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def get_router(self, memory: "MemoryEngine") -> APIRouter:
|
||||
"""
|
||||
Return a FastAPI router with custom endpoints.
|
||||
|
||||
The router will be mounted at /ext/ on the main application.
|
||||
All routes defined in the router will be prefixed with /ext/.
|
||||
|
||||
Args:
|
||||
memory: The MemoryEngine instance for database access and core operations.
|
||||
Use this to access the connection pool, run queries, or call
|
||||
memory operations like retain, recall, etc.
|
||||
|
||||
Returns:
|
||||
A FastAPI APIRouter with the custom endpoints defined.
|
||||
|
||||
Example:
|
||||
```python
|
||||
def get_router(self, memory: MemoryEngine) -> APIRouter:
|
||||
router = APIRouter(tags=["My Extension"])
|
||||
|
||||
@router.get("/status")
|
||||
async def status():
|
||||
health = await memory.health_check()
|
||||
return {"extension": "healthy", "memory": health}
|
||||
|
||||
return router
|
||||
```
|
||||
"""
|
||||
pass
|
||||
@@ -0,0 +1,125 @@
|
||||
"""Extension loader utilities."""
|
||||
|
||||
import importlib
|
||||
import logging
|
||||
import os
|
||||
from typing import TYPE_CHECKING, TypeVar
|
||||
|
||||
from hindsight_api.extensions.base import Extension
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from hindsight_api.extensions.context import ExtensionContext
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
T = TypeVar("T", bound=Extension)
|
||||
|
||||
|
||||
class ExtensionLoadError(Exception):
|
||||
"""Raised when an extension fails to load."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
def load_extension(
|
||||
prefix: str,
|
||||
base_class: type[T],
|
||||
env_prefix: str = "HINDSIGHT_API",
|
||||
context: "ExtensionContext | None" = None,
|
||||
) -> T | None:
|
||||
"""
|
||||
Load an extension from environment variable configuration.
|
||||
|
||||
The extension class is specified via {env_prefix}_{prefix}_EXTENSION environment
|
||||
variable in the format "module.path:ClassName".
|
||||
|
||||
Configuration for the extension is collected from all environment variables
|
||||
matching {env_prefix}_{prefix}_* (excluding the EXTENSION variable itself).
|
||||
|
||||
Args:
|
||||
prefix: The extension prefix (e.g., "OPERATION_VALIDATOR").
|
||||
base_class: The base class that the extension must inherit from.
|
||||
env_prefix: The environment variable prefix (default: "HINDSIGHT_API").
|
||||
context: Optional ExtensionContext to provide system APIs to the extension.
|
||||
|
||||
Returns:
|
||||
An instance of the extension, or None if not configured.
|
||||
|
||||
Raises:
|
||||
ExtensionLoadError: If the extension fails to load or validate.
|
||||
|
||||
Example:
|
||||
HINDSIGHT_API_OPERATION_VALIDATOR_EXTENSION=mypackage.validators:MyValidator
|
||||
HINDSIGHT_API_OPERATION_VALIDATOR_MAX_REQUESTS=100
|
||||
|
||||
ext = load_extension("OPERATION_VALIDATOR", OperationValidatorExtension)
|
||||
# ext.config == {"max_requests": "100"}
|
||||
"""
|
||||
env_var = f"{env_prefix}_{prefix}_EXTENSION"
|
||||
ext_path = os.getenv(env_var)
|
||||
|
||||
if not ext_path:
|
||||
logger.debug(f"No extension configured for {env_var}")
|
||||
return None
|
||||
|
||||
logger.info(f"Loading extension from {env_var}={ext_path}")
|
||||
|
||||
# Parse "module.path:ClassName"
|
||||
if ":" not in ext_path:
|
||||
raise ExtensionLoadError(f"Invalid extension path '{ext_path}'. Expected format: 'module.path:ClassName'")
|
||||
|
||||
module_path, class_name = ext_path.rsplit(":", 1)
|
||||
|
||||
# Import the module
|
||||
try:
|
||||
module = importlib.import_module(module_path)
|
||||
except ImportError as e:
|
||||
raise ExtensionLoadError(f"Failed to import extension module '{module_path}': {e}") from e
|
||||
|
||||
# Get the class
|
||||
try:
|
||||
ext_class = getattr(module, class_name)
|
||||
except AttributeError as e:
|
||||
raise ExtensionLoadError(f"Extension class '{class_name}' not found in module '{module_path}'") from e
|
||||
|
||||
# Validate inheritance
|
||||
if not isinstance(ext_class, type) or not issubclass(ext_class, base_class):
|
||||
raise ExtensionLoadError(f"Extension class '{ext_class.__name__}' must inherit from '{base_class.__name__}'")
|
||||
|
||||
# Collect configuration from environment variables
|
||||
config = _collect_config(env_prefix, prefix)
|
||||
|
||||
logger.info(f"Loaded extension {ext_class.__name__} with config keys: {list(config.keys())}")
|
||||
|
||||
# Instantiate the extension
|
||||
try:
|
||||
extension = ext_class(config)
|
||||
except Exception as e:
|
||||
raise ExtensionLoadError(f"Failed to instantiate extension '{ext_class.__name__}': {e}") from e
|
||||
|
||||
# Set the context if provided
|
||||
if context is not None:
|
||||
extension.set_context(context)
|
||||
logger.debug(f"Set context on extension {ext_class.__name__}")
|
||||
|
||||
return extension
|
||||
|
||||
|
||||
def _collect_config(env_prefix: str, prefix: str) -> dict[str, str]:
|
||||
"""
|
||||
Collect configuration from environment variables.
|
||||
|
||||
Collects all variables matching {env_prefix}_{prefix}_* except for
|
||||
{env_prefix}_{prefix}_EXTENSION, strips the prefix, and lowercases keys.
|
||||
"""
|
||||
config = {}
|
||||
full_prefix = f"{env_prefix}_{prefix}_"
|
||||
extension_var = f"{full_prefix}EXTENSION"
|
||||
|
||||
for key, value in os.environ.items():
|
||||
if key.startswith(full_prefix) and key != extension_var:
|
||||
# Strip prefix and lowercase the key
|
||||
config_key = key[len(full_prefix) :].lower()
|
||||
config[config_key] = value
|
||||
|
||||
return config
|
||||
@@ -0,0 +1,325 @@
|
||||
"""Operation Validator Extension for validating retain/recall/reflect operations."""
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from hindsight_api.extensions.base import Extension
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from hindsight_api.engine.memory_engine import Budget
|
||||
from hindsight_api.engine.response_models import RecallResult as RecallResultModel
|
||||
from hindsight_api.engine.response_models import ReflectResult
|
||||
from hindsight_api.models import RequestContext
|
||||
|
||||
|
||||
class OperationValidationError(Exception):
|
||||
"""Raised when an operation fails validation."""
|
||||
|
||||
def __init__(self, reason: str):
|
||||
self.reason = reason
|
||||
super().__init__(f"Operation validation failed: {reason}")
|
||||
|
||||
|
||||
@dataclass
|
||||
class ValidationResult:
|
||||
"""Result of an operation validation."""
|
||||
|
||||
allowed: bool
|
||||
reason: str | None = None
|
||||
|
||||
@classmethod
|
||||
def accept(cls) -> "ValidationResult":
|
||||
"""Create an accepted validation result."""
|
||||
return cls(allowed=True)
|
||||
|
||||
@classmethod
|
||||
def reject(cls, reason: str) -> "ValidationResult":
|
||||
"""Create a rejected validation result with a reason."""
|
||||
return cls(allowed=False, reason=reason)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Pre-operation Contexts (all user-provided parameters)
|
||||
# =============================================================================
|
||||
|
||||
|
||||
@dataclass
|
||||
class RetainContext:
|
||||
"""Context for a retain operation validation (pre-operation).
|
||||
|
||||
Contains ALL user-provided parameters for the retain operation.
|
||||
"""
|
||||
|
||||
bank_id: str
|
||||
contents: list[dict] # List of {content, context, event_date, document_id}
|
||||
request_context: "RequestContext"
|
||||
document_id: str | None = None
|
||||
fact_type_override: str | None = None
|
||||
confidence_score: float | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class RecallContext:
|
||||
"""Context for a recall operation validation (pre-operation).
|
||||
|
||||
Contains ALL user-provided parameters for the recall operation.
|
||||
"""
|
||||
|
||||
bank_id: str
|
||||
query: str
|
||||
request_context: "RequestContext"
|
||||
budget: "Budget | None" = None
|
||||
max_tokens: int = 4096
|
||||
enable_trace: bool = False
|
||||
fact_types: list[str] = field(default_factory=list)
|
||||
question_date: datetime | None = None
|
||||
include_entities: bool = False
|
||||
max_entity_tokens: int = 500
|
||||
include_chunks: bool = False
|
||||
max_chunk_tokens: int = 8192
|
||||
|
||||
|
||||
@dataclass
|
||||
class ReflectContext:
|
||||
"""Context for a reflect operation validation (pre-operation).
|
||||
|
||||
Contains ALL user-provided parameters for the reflect operation.
|
||||
"""
|
||||
|
||||
bank_id: str
|
||||
query: str
|
||||
request_context: "RequestContext"
|
||||
budget: "Budget | None" = None
|
||||
context: str | None = None
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Post-operation Contexts (includes results)
|
||||
# =============================================================================
|
||||
|
||||
|
||||
@dataclass
|
||||
class RetainResult:
|
||||
"""Result context for post-retain hook.
|
||||
|
||||
Contains the operation parameters and the result.
|
||||
"""
|
||||
|
||||
bank_id: str
|
||||
contents: list[dict]
|
||||
request_context: "RequestContext"
|
||||
document_id: str | None
|
||||
fact_type_override: str | None
|
||||
confidence_score: float | None
|
||||
# Result
|
||||
unit_ids: list[list[str]] # List of unit IDs per content item
|
||||
success: bool = True
|
||||
error: str | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class RecallResult:
|
||||
"""Result context for post-recall hook.
|
||||
|
||||
Contains the operation parameters and the result.
|
||||
"""
|
||||
|
||||
bank_id: str
|
||||
query: str
|
||||
request_context: "RequestContext"
|
||||
budget: "Budget | None"
|
||||
max_tokens: int
|
||||
enable_trace: bool
|
||||
fact_types: list[str]
|
||||
question_date: datetime | None
|
||||
include_entities: bool
|
||||
max_entity_tokens: int
|
||||
include_chunks: bool
|
||||
max_chunk_tokens: int
|
||||
# Result
|
||||
result: "RecallResultModel | None" = None
|
||||
success: bool = True
|
||||
error: str | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class ReflectResultContext:
|
||||
"""Result context for post-reflect hook.
|
||||
|
||||
Contains the operation parameters and the result.
|
||||
"""
|
||||
|
||||
bank_id: str
|
||||
query: str
|
||||
request_context: "RequestContext"
|
||||
budget: "Budget | None"
|
||||
context: str | None
|
||||
# Result
|
||||
result: "ReflectResult | None" = None
|
||||
success: bool = True
|
||||
error: str | None = None
|
||||
|
||||
|
||||
class OperationValidatorExtension(Extension, ABC):
|
||||
"""
|
||||
Validates and hooks into retain/recall/reflect operations.
|
||||
|
||||
This extension allows implementing custom logic such as:
|
||||
- Rate limiting (pre-operation)
|
||||
- Quota enforcement (pre-operation)
|
||||
- Permission checks (pre-operation)
|
||||
- Content filtering (pre-operation)
|
||||
- Usage tracking (post-operation)
|
||||
- Audit logging (post-operation)
|
||||
- Metrics collection (post-operation)
|
||||
|
||||
Enable via environment variable:
|
||||
HINDSIGHT_API_OPERATION_VALIDATOR_EXTENSION=mypackage.validators:MyValidator
|
||||
|
||||
Configuration is passed from prefixed environment variables:
|
||||
HINDSIGHT_API_OPERATION_VALIDATOR_MAX_REQUESTS=100
|
||||
-> config = {"max_requests": "100"}
|
||||
|
||||
Hook execution order:
|
||||
1. validate_retain/validate_recall/validate_reflect (pre-operation)
|
||||
2. [operation executes]
|
||||
3. on_retain_complete/on_recall_complete/on_reflect_complete (post-operation)
|
||||
"""
|
||||
|
||||
# =========================================================================
|
||||
# Pre-operation validation hooks (abstract - must be implemented)
|
||||
# =========================================================================
|
||||
|
||||
@abstractmethod
|
||||
async def validate_retain(self, ctx: RetainContext) -> ValidationResult:
|
||||
"""
|
||||
Validate a retain operation before execution.
|
||||
|
||||
Called before the retain operation is processed. Return ValidationResult.reject()
|
||||
to prevent the operation from executing.
|
||||
|
||||
Args:
|
||||
ctx: Context containing all user-provided parameters:
|
||||
- bank_id: Bank identifier
|
||||
- contents: List of content dicts
|
||||
- request_context: Request context with auth info
|
||||
- document_id: Optional document ID
|
||||
- fact_type_override: Optional fact type override
|
||||
- confidence_score: Optional confidence score
|
||||
|
||||
Returns:
|
||||
ValidationResult indicating whether the operation is allowed.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def validate_recall(self, ctx: RecallContext) -> ValidationResult:
|
||||
"""
|
||||
Validate a recall operation before execution.
|
||||
|
||||
Called before the recall operation is processed. Return ValidationResult.reject()
|
||||
to prevent the operation from executing.
|
||||
|
||||
Args:
|
||||
ctx: Context containing all user-provided parameters:
|
||||
- bank_id: Bank identifier
|
||||
- query: Search query
|
||||
- request_context: Request context with auth info
|
||||
- budget: Budget level
|
||||
- max_tokens: Maximum tokens to return
|
||||
- enable_trace: Whether to include trace info
|
||||
- fact_types: List of fact types to search
|
||||
- question_date: Optional date context for query
|
||||
- include_entities: Whether to include entity data
|
||||
- max_entity_tokens: Max tokens for entities
|
||||
- include_chunks: Whether to include chunks
|
||||
- max_chunk_tokens: Max tokens for chunks
|
||||
|
||||
Returns:
|
||||
ValidationResult indicating whether the operation is allowed.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def validate_reflect(self, ctx: ReflectContext) -> ValidationResult:
|
||||
"""
|
||||
Validate a reflect operation before execution.
|
||||
|
||||
Called before the reflect operation is processed. Return ValidationResult.reject()
|
||||
to prevent the operation from executing.
|
||||
|
||||
Args:
|
||||
ctx: Context containing all user-provided parameters:
|
||||
- bank_id: Bank identifier
|
||||
- query: Question to answer
|
||||
- request_context: Request context with auth info
|
||||
- budget: Budget level
|
||||
- context: Optional additional context
|
||||
|
||||
Returns:
|
||||
ValidationResult indicating whether the operation is allowed.
|
||||
"""
|
||||
...
|
||||
|
||||
# =========================================================================
|
||||
# Post-operation hooks (optional - override to implement)
|
||||
# =========================================================================
|
||||
|
||||
async def on_retain_complete(self, result: RetainResult) -> None:
|
||||
"""
|
||||
Called after a retain operation completes (success or failure).
|
||||
|
||||
Override this method to implement post-operation logic such as:
|
||||
- Usage tracking
|
||||
- Audit logging
|
||||
- Metrics collection
|
||||
- Notifications
|
||||
|
||||
Args:
|
||||
result: Result context containing:
|
||||
- All original operation parameters
|
||||
- unit_ids: List of created unit IDs (if success)
|
||||
- success: Whether the operation succeeded
|
||||
- error: Error message (if failed)
|
||||
"""
|
||||
pass
|
||||
|
||||
async def on_recall_complete(self, result: RecallResult) -> None:
|
||||
"""
|
||||
Called after a recall operation completes (success or failure).
|
||||
|
||||
Override this method to implement post-operation logic such as:
|
||||
- Usage tracking
|
||||
- Audit logging
|
||||
- Metrics collection
|
||||
- Query analytics
|
||||
|
||||
Args:
|
||||
result: Result context containing:
|
||||
- All original operation parameters
|
||||
- result: RecallResultModel (if success)
|
||||
- success: Whether the operation succeeded
|
||||
- error: Error message (if failed)
|
||||
"""
|
||||
pass
|
||||
|
||||
async def on_reflect_complete(self, result: ReflectResultContext) -> None:
|
||||
"""
|
||||
Called after a reflect operation completes (success or failure).
|
||||
|
||||
Override this method to implement post-operation logic such as:
|
||||
- Usage tracking
|
||||
- Audit logging
|
||||
- Metrics collection
|
||||
- Response analytics
|
||||
|
||||
Args:
|
||||
result: Result context containing:
|
||||
- All original operation parameters
|
||||
- result: ReflectResult (if success)
|
||||
- success: Whether the operation succeeded
|
||||
- error: Error message (if failed)
|
||||
"""
|
||||
pass
|
||||
@@ -0,0 +1,63 @@
|
||||
"""Tenant Extension for multi-tenancy and API key authentication."""
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass
|
||||
|
||||
from hindsight_api.extensions.base import Extension
|
||||
from hindsight_api.models import RequestContext
|
||||
|
||||
|
||||
class AuthenticationError(Exception):
|
||||
"""Raised when authentication fails."""
|
||||
|
||||
def __init__(self, reason: str):
|
||||
self.reason = reason
|
||||
super().__init__(f"Authentication failed: {reason}")
|
||||
|
||||
|
||||
@dataclass
|
||||
class TenantContext:
|
||||
"""
|
||||
Tenant context returned by authentication.
|
||||
|
||||
Contains the PostgreSQL schema name for tenant isolation.
|
||||
All database queries will use fully-qualified table names
|
||||
with this schema (e.g., schema_name.memory_units).
|
||||
"""
|
||||
|
||||
schema_name: str
|
||||
|
||||
|
||||
class TenantExtension(Extension, ABC):
|
||||
"""
|
||||
Extension for multi-tenancy and API key authentication.
|
||||
|
||||
This extension validates incoming requests and returns the tenant context
|
||||
including the PostgreSQL schema to use for database operations.
|
||||
|
||||
Built-in implementation:
|
||||
hindsight_api.extensions.builtin.tenant.ApiKeyTenantExtension
|
||||
|
||||
Enable via environment variable:
|
||||
HINDSIGHT_API_TENANT_EXTENSION=hindsight_api.extensions.builtin.tenant:ApiKeyTenantExtension
|
||||
HINDSIGHT_API_TENANT_API_KEY=your-secret-key
|
||||
|
||||
The returned schema_name is used for fully-qualified table names in queries,
|
||||
enabling tenant isolation at the database level.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
async def authenticate(self, context: RequestContext) -> TenantContext:
|
||||
"""
|
||||
Authenticate the action context and return tenant context.
|
||||
|
||||
Args:
|
||||
context: The action context containing API key and other auth data.
|
||||
|
||||
Returns:
|
||||
TenantContext with the schema_name for database operations.
|
||||
|
||||
Raises:
|
||||
AuthenticationError: If authentication fails.
|
||||
"""
|
||||
...
|
||||
@@ -4,6 +4,9 @@ Command-line interface for Hindsight API.
|
||||
Run the server with:
|
||||
hindsight-api
|
||||
|
||||
Run as background daemon:
|
||||
hindsight-api --daemon
|
||||
|
||||
Stop with Ctrl+C.
|
||||
"""
|
||||
|
||||
@@ -21,9 +24,13 @@ from . import MemoryEngine
|
||||
from .api import create_app
|
||||
from .banner import print_banner
|
||||
from .config import HindsightConfig, get_config
|
||||
|
||||
print()
|
||||
print_banner()
|
||||
from .daemon import (
|
||||
DEFAULT_DAEMON_PORT,
|
||||
DEFAULT_IDLE_TIMEOUT,
|
||||
DaemonLock,
|
||||
IdleTimeoutMiddleware,
|
||||
daemonize,
|
||||
)
|
||||
|
||||
# Filter deprecation warnings from third-party libraries
|
||||
warnings.filterwarnings("ignore", message="websockets.legacy is deprecated")
|
||||
@@ -106,8 +113,52 @@ def main():
|
||||
parser.add_argument("--ssl-keyfile", default=None, help="SSL key file")
|
||||
parser.add_argument("--ssl-certfile", default=None, help="SSL certificate file")
|
||||
|
||||
# Daemon mode options
|
||||
parser.add_argument(
|
||||
"--daemon",
|
||||
action="store_true",
|
||||
help=f"Run as background daemon (uses port {DEFAULT_DAEMON_PORT}, auto-exits after idle)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--idle-timeout",
|
||||
type=int,
|
||||
default=DEFAULT_IDLE_TIMEOUT,
|
||||
help=f"Idle timeout in seconds before auto-exit in daemon mode (default: {DEFAULT_IDLE_TIMEOUT})",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# Daemon mode handling
|
||||
if args.daemon:
|
||||
# Use fixed daemon port
|
||||
args.port = DEFAULT_DAEMON_PORT
|
||||
args.host = "127.0.0.1" # Only bind to localhost for security
|
||||
|
||||
# Check if another daemon is already running
|
||||
daemon_lock = DaemonLock()
|
||||
if not daemon_lock.acquire():
|
||||
print(f"Daemon already running (PID: {daemon_lock.get_pid()})", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
# Fork into background
|
||||
daemonize()
|
||||
|
||||
# Re-acquire lock in child process
|
||||
daemon_lock = DaemonLock()
|
||||
if not daemon_lock.acquire():
|
||||
sys.exit(1)
|
||||
|
||||
# Register cleanup to release lock
|
||||
def release_lock():
|
||||
daemon_lock.release()
|
||||
|
||||
atexit.register(release_lock)
|
||||
|
||||
# Print banner (not in daemon mode)
|
||||
if not args.daemon:
|
||||
print()
|
||||
print_banner()
|
||||
|
||||
# Configure Python logging based on log level
|
||||
# Update config with CLI override if provided
|
||||
if args.log_level != config.log_level:
|
||||
@@ -127,8 +178,13 @@ def main():
|
||||
port=args.port,
|
||||
log_level=args.log_level,
|
||||
mcp_enabled=config.mcp_enabled,
|
||||
graph_retriever=config.graph_retriever,
|
||||
skip_llm_verification=config.skip_llm_verification,
|
||||
lazy_reranker=config.lazy_reranker,
|
||||
)
|
||||
config.configure_logging()
|
||||
if not args.daemon:
|
||||
config.log_config()
|
||||
|
||||
# Register cleanup handlers
|
||||
atexit.register(_cleanup)
|
||||
@@ -147,6 +203,12 @@ def main():
|
||||
initialize_memory=True,
|
||||
)
|
||||
|
||||
# Wrap with idle timeout middleware in daemon mode
|
||||
idle_middleware = None
|
||||
if args.daemon:
|
||||
idle_middleware = IdleTimeoutMiddleware(app, idle_timeout=args.idle_timeout)
|
||||
app = idle_middleware
|
||||
|
||||
# Prepare uvicorn config
|
||||
uvicorn_config = {
|
||||
"app": app,
|
||||
@@ -170,20 +232,40 @@ def main():
|
||||
if args.ssl_certfile:
|
||||
uvicorn_config["ssl_certfile"] = args.ssl_certfile
|
||||
|
||||
from .banner import print_startup_info
|
||||
# Print startup info (not in daemon mode)
|
||||
if not args.daemon:
|
||||
from .banner import print_startup_info
|
||||
|
||||
print_startup_info(
|
||||
host=args.host,
|
||||
port=args.port,
|
||||
database_url=config.database_url,
|
||||
llm_provider=config.llm_provider,
|
||||
llm_model=config.llm_model,
|
||||
embeddings_provider=config.embeddings_provider,
|
||||
reranker_provider=config.reranker_provider,
|
||||
mcp_enabled=config.mcp_enabled,
|
||||
)
|
||||
print_startup_info(
|
||||
host=args.host,
|
||||
port=args.port,
|
||||
database_url=config.database_url,
|
||||
llm_provider=config.llm_provider,
|
||||
llm_model=config.llm_model,
|
||||
embeddings_provider=config.embeddings_provider,
|
||||
reranker_provider=config.reranker_provider,
|
||||
mcp_enabled=config.mcp_enabled,
|
||||
)
|
||||
|
||||
uvicorn.run(**uvicorn_config)
|
||||
# Start idle checker in daemon mode
|
||||
if idle_middleware is not None:
|
||||
# Start the idle checker in a background thread with its own event loop
|
||||
import threading
|
||||
|
||||
def run_idle_checker():
|
||||
import time
|
||||
|
||||
time.sleep(2) # Wait for uvicorn to start
|
||||
try:
|
||||
loop = asyncio.new_event_loop()
|
||||
asyncio.set_event_loop(loop)
|
||||
loop.run_until_complete(idle_middleware._check_idle())
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
threading.Thread(target=run_idle_checker, daemon=True).start()
|
||||
|
||||
uvicorn.run(**uvicorn_config) # type: ignore[invalid-argument-type] - dict kwargs
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -28,7 +28,15 @@ Environment variables:
|
||||
HINDSIGHT_API_LLM_PROVIDER: Optional. LLM provider (default: "openai").
|
||||
HINDSIGHT_API_LLM_MODEL: Optional. LLM model (default: "gpt-4o-mini").
|
||||
HINDSIGHT_API_MCP_LOCAL_BANK_ID: Optional. Memory bank ID (default: "mcp").
|
||||
HINDSIGHT_API_LOG_LEVEL: Optional. Log level (default: "info").
|
||||
HINDSIGHT_API_LOG_LEVEL: Optional. Log level (default: "warning").
|
||||
HINDSIGHT_API_MCP_INSTRUCTIONS: Optional. Additional instructions appended to both retain and recall tools.
|
||||
|
||||
Example custom instructions (these are ADDED to the default behavior):
|
||||
To also store assistant actions:
|
||||
HINDSIGHT_API_MCP_INSTRUCTIONS="Also store every action you take, including tool calls, code written, and decisions made."
|
||||
|
||||
To also store conversation summaries:
|
||||
HINDSIGHT_API_MCP_INSTRUCTIONS="Also store summaries of important conversations and their outcomes."
|
||||
"""
|
||||
|
||||
import logging
|
||||
@@ -36,14 +44,19 @@ import os
|
||||
import sys
|
||||
|
||||
from mcp.server.fastmcp import FastMCP
|
||||
from mcp.types import Icon
|
||||
|
||||
from hindsight_api.config import (
|
||||
DEFAULT_MCP_LOCAL_BANK_ID,
|
||||
DEFAULT_MCP_RECALL_DESCRIPTION,
|
||||
DEFAULT_MCP_RETAIN_DESCRIPTION,
|
||||
ENV_MCP_INSTRUCTIONS,
|
||||
ENV_MCP_LOCAL_BANK_ID,
|
||||
)
|
||||
|
||||
# Configure logging - default to info
|
||||
_log_level_str = os.environ.get("HINDSIGHT_API_LOG_LEVEL", "info").lower()
|
||||
# Configure logging - default to warning to avoid polluting stderr during MCP init
|
||||
# MCP clients interpret stderr output as errors, so we suppress INFO logs by default
|
||||
_log_level_str = os.environ.get("HINDSIGHT_API_LOG_LEVEL", "warning").lower()
|
||||
_log_level_map = {
|
||||
"critical": logging.CRITICAL,
|
||||
"error": logging.ERROR,
|
||||
@@ -74,27 +87,27 @@ def create_local_mcp_server(bank_id: str, memory=None) -> FastMCP:
|
||||
from hindsight_api import MemoryEngine
|
||||
from hindsight_api.engine.memory_engine import Budget
|
||||
from hindsight_api.engine.response_models import VALID_RECALL_FACT_TYPES
|
||||
from hindsight_api.models import RequestContext
|
||||
|
||||
# Create memory engine with pg0 embedded database if not provided
|
||||
if memory is None:
|
||||
memory = MemoryEngine(db_url="pg0://hindsight-mcp")
|
||||
|
||||
# Get custom instructions from environment variable (appended to both tools)
|
||||
extra_instructions = os.environ.get(ENV_MCP_INSTRUCTIONS, "")
|
||||
|
||||
retain_description = DEFAULT_MCP_RETAIN_DESCRIPTION
|
||||
recall_description = DEFAULT_MCP_RECALL_DESCRIPTION
|
||||
|
||||
if extra_instructions:
|
||||
retain_description = f"{DEFAULT_MCP_RETAIN_DESCRIPTION}\n\nAdditional instructions: {extra_instructions}"
|
||||
recall_description = f"{DEFAULT_MCP_RECALL_DESCRIPTION}\n\nAdditional instructions: {extra_instructions}"
|
||||
|
||||
mcp = FastMCP("hindsight")
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(description=retain_description)
|
||||
async def retain(content: str, context: str = "general") -> dict:
|
||||
"""
|
||||
Store important information to long-term memory.
|
||||
|
||||
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
|
||||
- Work context, projects, or responsibilities
|
||||
|
||||
Args:
|
||||
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'
|
||||
@@ -103,7 +116,11 @@ def create_local_mcp_server(bank_id: str, memory=None) -> FastMCP:
|
||||
|
||||
async def _retain():
|
||||
try:
|
||||
await memory.retain_batch_async(bank_id=bank_id, contents=[{"content": content, "context": context}])
|
||||
await memory.retain_batch_async(
|
||||
bank_id=bank_id,
|
||||
contents=[{"content": content, "context": context}],
|
||||
request_context=RequestContext(),
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"Error storing memory: {e}", exc_info=True)
|
||||
|
||||
@@ -111,17 +128,9 @@ def create_local_mcp_server(bank_id: str, memory=None) -> FastMCP:
|
||||
asyncio.create_task(_retain())
|
||||
return {"status": "accepted", "message": "Memory storage initiated"}
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(description=recall_description)
|
||||
async def recall(query: str, max_tokens: int = 4096, budget: str = "low") -> dict:
|
||||
"""
|
||||
Search memories 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
|
||||
|
||||
Args:
|
||||
query: Natural language search query (e.g., "user's food preferences", "what projects is user working on")
|
||||
max_tokens: Maximum tokens to return in results (default: 4096)
|
||||
@@ -138,6 +147,7 @@ def create_local_mcp_server(bank_id: str, memory=None) -> FastMCP:
|
||||
fact_type=list(VALID_RECALL_FACT_TYPES),
|
||||
budget=budget_enum,
|
||||
max_tokens=max_tokens,
|
||||
request_context=RequestContext(),
|
||||
)
|
||||
|
||||
return search_result.model_dump()
|
||||
@@ -153,10 +163,9 @@ async def _initialize_and_run(bank_id: str):
|
||||
from hindsight_api import MemoryEngine
|
||||
|
||||
# Create and initialize memory engine with pg0 embedded database
|
||||
print("Initializing memory engine...", file=sys.stderr)
|
||||
# Note: We avoid printing to stderr during init as MCP clients show it as "errors"
|
||||
memory = MemoryEngine(db_url="pg0://hindsight-mcp")
|
||||
await memory.initialize()
|
||||
print("Memory engine initialized.", file=sys.stderr)
|
||||
|
||||
# Create and run the server
|
||||
mcp = create_local_mcp_server(bank_id, memory=memory)
|
||||
@@ -179,8 +188,8 @@ def main():
|
||||
# Get bank ID from environment, default to "mcp"
|
||||
bank_id = os.environ.get(ENV_MCP_LOCAL_BANK_ID, DEFAULT_MCP_LOCAL_BANK_ID)
|
||||
|
||||
# Print startup message to stderr (stdout is reserved for MCP protocol)
|
||||
print(f"Hindsight MCP server starting (bank_id={bank_id})...", file=sys.stderr)
|
||||
# Note: We don't print to stderr as MCP clients display it as "error output"
|
||||
# Use HINDSIGHT_API_LOG_LEVEL=debug for verbose startup logging
|
||||
|
||||
# Run the async initialization and server
|
||||
asyncio.run(_initialize_and_run(bank_id))
|
||||
|
||||
@@ -6,12 +6,16 @@ on application startup. It is designed to be safe for concurrent
|
||||
execution using PostgreSQL advisory locks to coordinate between
|
||||
distributed workers.
|
||||
|
||||
Supports multi-tenant schema isolation: migrations can target a specific
|
||||
PostgreSQL schema, allowing each tenant to have isolated tables.
|
||||
|
||||
Important: All migrations must be backward-compatible to allow
|
||||
safe rolling deployments.
|
||||
|
||||
No alembic.ini required - all configuration is done programmatically.
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
import logging
|
||||
import os
|
||||
from pathlib import Path
|
||||
@@ -26,11 +30,29 @@ logger = logging.getLogger(__name__)
|
||||
MIGRATION_LOCK_ID = 123456789
|
||||
|
||||
|
||||
def _run_migrations_internal(database_url: str, script_location: str) -> None:
|
||||
def _get_schema_lock_id(schema: str) -> int:
|
||||
"""
|
||||
Generate a unique advisory lock ID for a schema.
|
||||
|
||||
Uses hash of schema name to create a deterministic lock ID.
|
||||
"""
|
||||
# Use hash to create a unique lock ID per schema
|
||||
# Keep within PostgreSQL's bigint range
|
||||
hash_bytes = hashlib.sha256(schema.encode()).digest()[:8]
|
||||
return int.from_bytes(hash_bytes, byteorder="big") % (2**31)
|
||||
|
||||
|
||||
def _run_migrations_internal(database_url: str, script_location: str, schema: str | None = None) -> None:
|
||||
"""
|
||||
Internal function to run migrations without locking.
|
||||
|
||||
Args:
|
||||
database_url: SQLAlchemy database URL
|
||||
script_location: Path to alembic scripts
|
||||
schema: Target schema (None for default/public)
|
||||
"""
|
||||
logger.info("Running database migrations to head...")
|
||||
schema_name = schema or "public"
|
||||
logger.info(f"Running database migrations to head for schema '{schema_name}'...")
|
||||
logger.info(f"Database URL: {database_url}")
|
||||
logger.info(f"Script location: {script_location}")
|
||||
|
||||
@@ -50,13 +72,22 @@ def _run_migrations_internal(database_url: str, script_location: str) -> None:
|
||||
# Set path_separator to avoid deprecation warning
|
||||
alembic_cfg.set_main_option("path_separator", "os")
|
||||
|
||||
# Run migrations to head (latest version)
|
||||
# If targeting a specific schema, pass it to env.py via config
|
||||
# env.py will handle setting search_path and version_table_schema
|
||||
if schema:
|
||||
alembic_cfg.set_main_option("target_schema", schema)
|
||||
|
||||
# Run migrations
|
||||
command.upgrade(alembic_cfg, "head")
|
||||
|
||||
logger.info("Database migrations completed successfully")
|
||||
logger.info(f"Database migrations completed successfully for schema '{schema_name}'")
|
||||
|
||||
|
||||
def run_migrations(database_url: str, script_location: str | None = None) -> None:
|
||||
def run_migrations(
|
||||
database_url: str,
|
||||
script_location: str | None = None,
|
||||
schema: str | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
Run database migrations to the latest version using programmatic Alembic configuration.
|
||||
|
||||
@@ -65,19 +96,28 @@ def run_migrations(database_url: str, script_location: str | None = None) -> Non
|
||||
- Other workers wait for the lock, then verify migrations are complete
|
||||
- If schema is already up-to-date, this is a fast no-op
|
||||
|
||||
Supports multi-tenant schema isolation: when a schema is specified, migrations
|
||||
run in that schema instead of public. This allows tenant extensions to provision
|
||||
new tenant schemas with their own isolated tables.
|
||||
|
||||
Args:
|
||||
database_url: SQLAlchemy database URL (e.g., "postgresql://user:pass@host/db")
|
||||
script_location: Path to alembic migrations directory (e.g., "/path/to/alembic").
|
||||
If None, defaults to hindsight-api/alembic directory.
|
||||
schema: Target PostgreSQL schema name. If None, uses default (public).
|
||||
When specified, creates the schema if needed and runs migrations there.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If migrations fail to complete
|
||||
FileNotFoundError: If script_location doesn't exist
|
||||
|
||||
Example:
|
||||
# Using default location (hindsight_api package)
|
||||
# Using default location and public schema
|
||||
run_migrations("postgresql://user:pass@host/db")
|
||||
|
||||
# Run migrations for a specific tenant schema
|
||||
run_migrations("postgresql://user:pass@host/db", schema="tenant_acme")
|
||||
|
||||
# Using custom location (when importing from another project)
|
||||
run_migrations(
|
||||
"postgresql://user:pass@host/db",
|
||||
@@ -99,21 +139,25 @@ def run_migrations(database_url: str, script_location: str | None = None) -> Non
|
||||
f"Alembic script location not found at {script_location}. Database migrations cannot be run."
|
||||
)
|
||||
|
||||
# Use schema-specific lock ID for multi-tenant isolation
|
||||
lock_id = _get_schema_lock_id(schema) if schema else MIGRATION_LOCK_ID
|
||||
schema_name = schema or "public"
|
||||
|
||||
# Use PostgreSQL advisory lock to coordinate between distributed workers
|
||||
engine = create_engine(database_url)
|
||||
with engine.connect() as conn:
|
||||
# pg_advisory_lock blocks until the lock is acquired
|
||||
# The lock is automatically released when the connection closes
|
||||
logger.debug(f"Acquiring migration advisory lock (id={MIGRATION_LOCK_ID})...")
|
||||
conn.execute(text(f"SELECT pg_advisory_lock({MIGRATION_LOCK_ID})"))
|
||||
logger.debug(f"Acquiring migration advisory lock for schema '{schema_name}' (id={lock_id})...")
|
||||
conn.execute(text(f"SELECT pg_advisory_lock({lock_id})"))
|
||||
logger.debug("Migration advisory lock acquired")
|
||||
|
||||
try:
|
||||
# Run migrations while holding the lock
|
||||
_run_migrations_internal(database_url, script_location)
|
||||
_run_migrations_internal(database_url, script_location, schema=schema)
|
||||
finally:
|
||||
# Explicitly release the lock (also released on connection close)
|
||||
conn.execute(text(f"SELECT pg_advisory_unlock({MIGRATION_LOCK_ID})"))
|
||||
conn.execute(text(f"SELECT pg_advisory_unlock({lock_id})"))
|
||||
logger.debug("Migration advisory lock released")
|
||||
|
||||
except FileNotFoundError:
|
||||
|
||||
@@ -2,9 +2,24 @@
|
||||
SQLAlchemy models for the memory system.
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from uuid import UUID as PyUUID
|
||||
|
||||
|
||||
@dataclass
|
||||
class RequestContext:
|
||||
"""
|
||||
Context for request authentication and authorization.
|
||||
|
||||
This dataclass carries authentication data from HTTP requests to the
|
||||
memory engine operations. It can be extended to include additional
|
||||
context like headers, tokens, user info, etc.
|
||||
"""
|
||||
|
||||
api_key: str | None = None
|
||||
|
||||
|
||||
from pgvector.sqlalchemy import Vector
|
||||
from sqlalchemy import (
|
||||
CheckConstraint,
|
||||
|
||||
@@ -40,7 +40,7 @@ class EmbeddedPostgres:
|
||||
# Only set port if explicitly specified
|
||||
if self.port is not None:
|
||||
kwargs["port"] = self.port
|
||||
self._pg0 = Pg0(**kwargs)
|
||||
self._pg0 = Pg0(**kwargs) # type: ignore[invalid-argument-type] - dict kwargs
|
||||
return self._pg0
|
||||
|
||||
async def start(self, max_retries: int = 5, retry_delay: float = 4.0) -> str:
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "hindsight-api"
|
||||
version = "0.1.8"
|
||||
version = "0.1.13"
|
||||
description = "Hindsight: Agent Memory That Works Like Human Memory"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.11"
|
||||
@@ -92,6 +92,7 @@ dev = [
|
||||
"python-dotenv>=1.2.1",
|
||||
"filelock>=3.0.0",
|
||||
"ruff>=0.8.0",
|
||||
"ty>=0.0.1",
|
||||
]
|
||||
|
||||
[tool.ruff]
|
||||
@@ -121,3 +122,28 @@ ignore = [
|
||||
[tool.ruff.format]
|
||||
quote-style = "double"
|
||||
indent-style = "space"
|
||||
|
||||
[tool.ty]
|
||||
# Type checking configuration
|
||||
# ty is an extremely fast Python type checker from Astral (same team as ruff/uv)
|
||||
|
||||
[tool.ty.environment]
|
||||
python-version = "3.11"
|
||||
|
||||
[tool.ty.src]
|
||||
exclude = [
|
||||
"tests/",
|
||||
"hindsight_api/alembic/",
|
||||
]
|
||||
|
||||
[tool.ty.rules]
|
||||
# Disable noisy rules while keeping important ones
|
||||
invalid-argument-type = "ignore" # False positives with **kwargs patterns
|
||||
invalid-return-type = "ignore" # Often intentional in async code
|
||||
invalid-parameter-default = "ignore" # Optional params with None default
|
||||
possibly-missing-attribute = "ignore" # Common with Optional types
|
||||
invalid-raise = "ignore" # False positives with exception tracking
|
||||
call-non-callable = "ignore" # False positives with Optional types
|
||||
invalid-key = "ignore" # Pydantic ConfigDict not understood
|
||||
invalid-method-override = "ignore" # Intentional signature differences
|
||||
unresolved-reference = "ignore" # Forward references not always resolved
|
||||
|
||||
@@ -8,7 +8,7 @@ import os
|
||||
import filelock
|
||||
from pathlib import Path
|
||||
from dotenv import load_dotenv
|
||||
from hindsight_api import MemoryEngine, LLMConfig, LocalSTEmbeddings
|
||||
from hindsight_api import MemoryEngine, LLMConfig, LocalSTEmbeddings, RequestContext
|
||||
|
||||
from hindsight_api.engine.cross_encoder import LocalSTCrossEncoder
|
||||
from hindsight_api.engine.query_analyzer import DateparserQueryAnalyzer
|
||||
@@ -99,6 +99,12 @@ def pg0_db_url(db_url, tmp_path_factory, worker_id):
|
||||
return url
|
||||
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def request_context():
|
||||
"""Provide a default RequestContext for tests."""
|
||||
return RequestContext()
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def llm_config():
|
||||
"""
|
||||
|
||||
@@ -3,7 +3,7 @@ Tests for agent management API (profile, disposition, background).
|
||||
"""
|
||||
import pytest
|
||||
import uuid
|
||||
from hindsight_api import MemoryEngine
|
||||
from hindsight_api import MemoryEngine, RequestContext
|
||||
from hindsight_api.api import CreateBankRequest, DispositionTraits
|
||||
from hindsight_api.engine.memory_engine import Budget
|
||||
|
||||
@@ -17,11 +17,11 @@ class TestAgentProfile:
|
||||
"""Tests for agent profile management."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_agent_profile_creates_default(self, memory: MemoryEngine):
|
||||
async def test_get_agent_profile_creates_default(self, memory: MemoryEngine, request_context):
|
||||
"""Test that getting a profile for a new agent creates default disposition."""
|
||||
bank_id = unique_agent_id("test_profile_default")
|
||||
|
||||
profile = await memory.get_bank_profile(bank_id)
|
||||
profile = await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
|
||||
assert profile is not None
|
||||
assert "disposition" in profile
|
||||
@@ -35,11 +35,11 @@ class TestAgentProfile:
|
||||
assert profile["background"] == ""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_agent_disposition(self, memory: MemoryEngine):
|
||||
async def test_update_agent_disposition(self, memory: MemoryEngine, request_context):
|
||||
"""Test updating agent disposition traits."""
|
||||
bank_id = unique_agent_id("test_profile_update")
|
||||
|
||||
profile = await memory.get_bank_profile(bank_id)
|
||||
profile = await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
assert profile["disposition"].skepticism == 3
|
||||
|
||||
new_disposition = {
|
||||
@@ -47,26 +47,26 @@ class TestAgentProfile:
|
||||
"literalism": 4,
|
||||
"empathy": 2,
|
||||
}
|
||||
await memory.update_bank_disposition(bank_id, new_disposition)
|
||||
await memory.update_bank_disposition(bank_id, new_disposition, request_context=request_context)
|
||||
|
||||
updated_profile = await memory.get_bank_profile(bank_id)
|
||||
updated_profile = await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
disposition = updated_profile["disposition"]
|
||||
assert disposition.skepticism == new_disposition["skepticism"]
|
||||
assert disposition.literalism == new_disposition["literalism"]
|
||||
assert disposition.empathy == new_disposition["empathy"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_agents(self, memory: MemoryEngine):
|
||||
async def test_list_agents(self, memory: MemoryEngine, request_context):
|
||||
"""Test listing all agents."""
|
||||
agent_id_1 = unique_agent_id("test_list")
|
||||
agent_id_2 = unique_agent_id("test_list")
|
||||
agent_id_3 = unique_agent_id("test_list")
|
||||
|
||||
await memory.get_bank_profile(agent_id_1)
|
||||
await memory.get_bank_profile(agent_id_2)
|
||||
await memory.get_bank_profile(agent_id_3)
|
||||
await memory.get_bank_profile(agent_id_1, request_context=request_context)
|
||||
await memory.get_bank_profile(agent_id_2, request_context=request_context)
|
||||
await memory.get_bank_profile(agent_id_3, request_context=request_context)
|
||||
|
||||
agents = await memory.list_banks()
|
||||
agents = await memory.list_banks(request_context=request_context)
|
||||
|
||||
agent_ids = [a["bank_id"] for a in agents]
|
||||
assert agent_id_1 in agent_ids
|
||||
@@ -85,46 +85,50 @@ class TestAgentBackground:
|
||||
"""Tests for agent background management."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_merge_agent_background(self, memory: MemoryEngine):
|
||||
async def test_merge_agent_background(self, memory: MemoryEngine, request_context):
|
||||
"""Test merging agent background information."""
|
||||
bank_id = unique_agent_id("test_profile_merge")
|
||||
|
||||
profile = await memory.get_bank_profile(bank_id)
|
||||
profile = await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
assert profile["background"] == ""
|
||||
|
||||
result1 = await memory.merge_bank_background(
|
||||
bank_id,
|
||||
"I was born in Texas",
|
||||
update_disposition=False
|
||||
update_disposition=False,
|
||||
request_context=request_context,
|
||||
)
|
||||
assert "Texas" in result1["background"]
|
||||
|
||||
result2 = await memory.merge_bank_background(
|
||||
bank_id,
|
||||
"I have 10 years of startup experience",
|
||||
update_disposition=False
|
||||
update_disposition=False,
|
||||
request_context=request_context,
|
||||
)
|
||||
assert "Texas" in result2["background"] or "startup" in result2["background"]
|
||||
|
||||
final_profile = await memory.get_bank_profile(bank_id)
|
||||
final_profile = await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
assert final_profile["background"] != ""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_merge_background_handles_conflicts(self, memory: MemoryEngine):
|
||||
async def test_merge_background_handles_conflicts(self, memory: MemoryEngine, request_context):
|
||||
"""Test that merging background handles conflicts (new overwrites old)."""
|
||||
bank_id = unique_agent_id("test_profile_conflict")
|
||||
|
||||
result1 = await memory.merge_bank_background(
|
||||
bank_id,
|
||||
"I was born in Colorado",
|
||||
update_disposition=False
|
||||
update_disposition=False,
|
||||
request_context=request_context,
|
||||
)
|
||||
assert "Colorado" in result1["background"]
|
||||
|
||||
result2 = await memory.merge_bank_background(
|
||||
bank_id,
|
||||
"You were born in Texas",
|
||||
update_disposition=False
|
||||
update_disposition=False,
|
||||
request_context=request_context,
|
||||
)
|
||||
assert "Texas" in result2["background"]
|
||||
|
||||
@@ -133,7 +137,7 @@ class TestAgentEndpoint:
|
||||
"""Tests for agent PUT endpoint logic."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_put_agent_create(self, memory: MemoryEngine):
|
||||
async def test_put_agent_create(self, memory: MemoryEngine, request_context):
|
||||
"""Test creating an agent via PUT endpoint."""
|
||||
bank_id = unique_agent_id("test_put_create")
|
||||
|
||||
@@ -146,12 +150,13 @@ class TestAgentEndpoint:
|
||||
background="I am a creative software engineer"
|
||||
)
|
||||
|
||||
profile = await memory.get_bank_profile(bank_id)
|
||||
profile = await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
|
||||
if request.disposition is not None:
|
||||
await memory.update_bank_disposition(
|
||||
bank_id,
|
||||
request.disposition.model_dump()
|
||||
request.disposition.model_dump(),
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
if request.background is not None:
|
||||
@@ -168,14 +173,14 @@ class TestAgentEndpoint:
|
||||
request.background
|
||||
)
|
||||
|
||||
final_profile = await memory.get_bank_profile(bank_id)
|
||||
final_profile = await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
|
||||
assert final_profile["disposition"].skepticism == 4
|
||||
assert final_profile["disposition"].literalism == 5
|
||||
assert final_profile["background"] == "I am a creative software engineer"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_put_agent_partial_update(self, memory: MemoryEngine):
|
||||
async def test_put_agent_partial_update(self, memory: MemoryEngine, request_context):
|
||||
"""Test updating only background."""
|
||||
bank_id = unique_agent_id("test_put_partial")
|
||||
|
||||
@@ -183,7 +188,7 @@ class TestAgentEndpoint:
|
||||
background="I am a data scientist"
|
||||
)
|
||||
|
||||
profile = await memory.get_bank_profile(bank_id)
|
||||
profile = await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
|
||||
if request.background is not None:
|
||||
pool = await memory._get_pool()
|
||||
@@ -199,7 +204,7 @@ class TestAgentEndpoint:
|
||||
request.background
|
||||
)
|
||||
|
||||
final_profile = await memory.get_bank_profile(bank_id)
|
||||
final_profile = await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
|
||||
assert final_profile["disposition"].skepticism == 3 # Default
|
||||
assert final_profile["background"] == "I am a data scientist"
|
||||
@@ -209,7 +214,7 @@ class TestAgentDispositionIntegration:
|
||||
"""Tests for disposition integration with other features."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_think_uses_disposition(self, memory: MemoryEngine):
|
||||
async def test_think_uses_disposition(self, memory: MemoryEngine, request_context):
|
||||
"""Test that THINK operation uses agent disposition."""
|
||||
bank_id = unique_agent_id("test_think")
|
||||
|
||||
@@ -218,12 +223,13 @@ class TestAgentDispositionIntegration:
|
||||
"literalism": 4, # High literalism
|
||||
"empathy": 2, # Low empathy
|
||||
}
|
||||
await memory.update_bank_disposition(bank_id, disposition)
|
||||
await memory.update_bank_disposition(bank_id, disposition, request_context=request_context)
|
||||
|
||||
await memory.merge_bank_background(
|
||||
bank_id,
|
||||
"I am a creative artist who values innovation over tradition",
|
||||
update_disposition=False
|
||||
update_disposition=False,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
await memory.retain_batch_async(
|
||||
@@ -232,13 +238,14 @@ class TestAgentDispositionIntegration:
|
||||
{"content": "Traditional painting techniques have been used for centuries"},
|
||||
{"content": "Modern digital art is changing the art world"}
|
||||
],
|
||||
document_id="art_facts"
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
result = await memory.reflect_async(
|
||||
bank_id=bank_id,
|
||||
query="What do you think about traditional vs modern art?",
|
||||
budget=Budget.LOW
|
||||
budget=Budget.LOW,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert result.text is not None
|
||||
|
||||
@@ -6,7 +6,7 @@ import os
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_large_batch_auto_chunks(memory):
|
||||
async def test_large_batch_auto_chunks(memory, request_context):
|
||||
bank_id = "test_chunking_agent"
|
||||
# Create a large batch that should trigger chunking
|
||||
# Each item is ~2000 chars, so 30 items = 60k chars (exceeds 50k threshold)
|
||||
@@ -24,7 +24,8 @@ async def test_large_batch_auto_chunks(memory):
|
||||
# Ingest the large batch (should auto-chunk)
|
||||
result = await memory.retain_batch_async(
|
||||
bank_id=bank_id,
|
||||
contents=contents
|
||||
contents=contents,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Verify we got results back
|
||||
@@ -33,7 +34,7 @@ async def test_large_batch_auto_chunks(memory):
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_small_batch_no_chunking(memory):
|
||||
async def test_small_batch_no_chunking(memory, request_context):
|
||||
bank_id = "test_no_chunking_agent"
|
||||
|
||||
# Create a small batch that should NOT trigger chunking
|
||||
@@ -50,7 +51,8 @@ async def test_small_batch_no_chunking(memory):
|
||||
# Ingest the small batch (should NOT auto-chunk)
|
||||
result = await memory.retain_batch_async(
|
||||
bank_id=bank_id,
|
||||
contents=contents
|
||||
contents=contents,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Verify we got results back
|
||||
|
||||
@@ -10,6 +10,7 @@ import pytest
|
||||
from datetime import datetime, timezone
|
||||
from hindsight_api.engine.search.types import RetrievalResult, MergedCandidate, ScoredResult
|
||||
from hindsight_api.engine.memory_engine import Budget
|
||||
from hindsight_api import RequestContext
|
||||
|
||||
|
||||
class TestRRFNormalization:
|
||||
@@ -125,7 +126,7 @@ class TestCombinedScoringFormula:
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_trace_has_normalized_rrf(memory):
|
||||
async def test_trace_has_normalized_rrf(memory, request_context):
|
||||
"""Integration test: verify trace contains normalized RRF values, not raw."""
|
||||
bank_id = f"test_scoring_{datetime.now(timezone.utc).timestamp()}"
|
||||
|
||||
@@ -135,21 +136,25 @@ async def test_trace_has_normalized_rrf(memory):
|
||||
bank_id=bank_id,
|
||||
content="Python is a programming language created by Guido van Rossum",
|
||||
context="tech facts",
|
||||
request_context=request_context,
|
||||
)
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="JavaScript was created by Brendan Eich at Netscape",
|
||||
context="tech facts",
|
||||
request_context=request_context,
|
||||
)
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="The Eiffel Tower is located in Paris, France",
|
||||
context="geography facts",
|
||||
request_context=request_context,
|
||||
)
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="Mount Everest is the tallest mountain on Earth",
|
||||
context="geography facts",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Search with tracing
|
||||
@@ -160,6 +165,7 @@ async def test_trace_has_normalized_rrf(memory):
|
||||
budget=Budget.LOW,
|
||||
max_tokens=1024,
|
||||
enable_trace=True,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert result.trace is not None, "Trace should be present"
|
||||
@@ -210,11 +216,11 @@ async def test_trace_has_normalized_rrf(memory):
|
||||
print(f" - First result score components: {sc}")
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rrf_normalized_not_raw_in_trace(memory):
|
||||
async def test_rrf_normalized_not_raw_in_trace(memory, request_context):
|
||||
"""Verify that raw RRF scores (0.04-0.06 range) don't appear as normalized values."""
|
||||
bank_id = f"test_rrf_raw_{datetime.now(timezone.utc).timestamp()}"
|
||||
|
||||
@@ -225,6 +231,7 @@ async def test_rrf_normalized_not_raw_in_trace(memory):
|
||||
bank_id=bank_id,
|
||||
content=f"Test fact number {i} about various topics",
|
||||
context="test context",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
result = await memory.recall_async(
|
||||
@@ -234,6 +241,7 @@ async def test_rrf_normalized_not_raw_in_trace(memory):
|
||||
budget=Budget.LOW,
|
||||
max_tokens=512,
|
||||
enable_trace=True,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
trace = result.trace
|
||||
@@ -268,11 +276,11 @@ async def test_rrf_normalized_not_raw_in_trace(memory):
|
||||
print("\n✓ RRF raw vs normalized test passed!")
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_combined_score_matches_components(memory):
|
||||
async def test_combined_score_matches_components(memory, request_context):
|
||||
"""Verify the final score actually equals the weighted sum of components."""
|
||||
bank_id = f"test_combined_{datetime.now(timezone.utc).timestamp()}"
|
||||
|
||||
@@ -281,11 +289,13 @@ async def test_combined_score_matches_components(memory):
|
||||
bank_id=bank_id,
|
||||
content="The quick brown fox jumps over the lazy dog",
|
||||
context="test",
|
||||
request_context=request_context,
|
||||
)
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="A quick test of the emergency broadcast system",
|
||||
context="test",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
result = await memory.recall_async(
|
||||
@@ -295,6 +305,7 @@ async def test_combined_score_matches_components(memory):
|
||||
budget=Budget.LOW,
|
||||
max_tokens=512,
|
||||
enable_trace=True,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
trace = result.trace
|
||||
@@ -320,4 +331,4 @@ async def test_combined_score_matches_components(memory):
|
||||
print("\n✓ Combined score verification test passed!")
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@@ -4,10 +4,11 @@ Tests for document tracking and upsert functionality.
|
||||
import logging
|
||||
import pytest
|
||||
from datetime import datetime, timezone
|
||||
from hindsight_api import RequestContext
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_document_creation_and_retrieval(memory):
|
||||
async def test_document_creation_and_retrieval(memory, request_context):
|
||||
"""Test that documents are created and can be retrieved."""
|
||||
bank_id = f"test_doc_{datetime.now(timezone.utc).timestamp()}"
|
||||
|
||||
@@ -19,11 +20,12 @@ async def test_document_creation_and_retrieval(memory):
|
||||
bank_id=bank_id,
|
||||
content="Alice works at Google. Bob works at Microsoft.",
|
||||
context="Team meeting",
|
||||
document_id=document_id
|
||||
document_id=document_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Retrieve document
|
||||
doc = await memory.get_document(document_id, bank_id)
|
||||
doc = await memory.get_document(document_id, bank_id, request_context=request_context)
|
||||
|
||||
assert doc is not None
|
||||
assert doc["id"] == document_id
|
||||
@@ -32,11 +34,11 @@ async def test_document_creation_and_retrieval(memory):
|
||||
assert doc["memory_unit_count"] > 0
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_document_upsert(memory):
|
||||
async def test_document_upsert(memory, request_context):
|
||||
"""Test that providing the same document_id automatically upserts (deletes old units and creates new ones)."""
|
||||
bank_id = f"test_upsert_{datetime.now(timezone.utc).timestamp()}"
|
||||
|
||||
@@ -48,11 +50,12 @@ async def test_document_upsert(memory):
|
||||
bank_id=bank_id,
|
||||
content="Alice works at Google.",
|
||||
context="Initial",
|
||||
document_id=document_id
|
||||
document_id=document_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Get document stats
|
||||
doc_v1 = await memory.get_document(document_id, bank_id)
|
||||
doc_v1 = await memory.get_document(document_id, bank_id, request_context=request_context)
|
||||
count_v1 = doc_v1["memory_unit_count"]
|
||||
|
||||
# Update with different content (automatic upsert when same document_id is provided)
|
||||
@@ -60,11 +63,12 @@ async def test_document_upsert(memory):
|
||||
bank_id=bank_id,
|
||||
content="Alice works at Microsoft. Bob works at Apple.",
|
||||
context="Updated",
|
||||
document_id=document_id
|
||||
document_id=document_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Get updated document stats
|
||||
doc_v2 = await memory.get_document(document_id, bank_id)
|
||||
doc_v2 = await memory.get_document(document_id, bank_id, request_context=request_context)
|
||||
count_v2 = doc_v2["memory_unit_count"]
|
||||
|
||||
# Verify old units were replaced
|
||||
@@ -75,11 +79,11 @@ async def test_document_upsert(memory):
|
||||
assert set(units_v1).isdisjoint(set(units_v2))
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_document_deletion(memory):
|
||||
async def test_document_deletion(memory, request_context):
|
||||
"""Test that deleting a document cascades to memory units."""
|
||||
bank_id = f"test_delete_{datetime.now(timezone.utc).timestamp()}"
|
||||
|
||||
@@ -91,29 +95,30 @@ async def test_document_deletion(memory):
|
||||
bank_id=bank_id,
|
||||
content="Alice works at Google.",
|
||||
context="Test",
|
||||
document_id=document_id
|
||||
document_id=document_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Verify it exists
|
||||
doc = await memory.get_document(document_id, bank_id)
|
||||
doc = await memory.get_document(document_id, bank_id, request_context=request_context)
|
||||
assert doc is not None
|
||||
assert doc["memory_unit_count"] > 0
|
||||
|
||||
# Delete document
|
||||
result = await memory.delete_document(document_id, bank_id)
|
||||
result = await memory.delete_document(document_id, bank_id, request_context=request_context)
|
||||
assert result["document_deleted"] == 1
|
||||
assert result["memory_units_deleted"] > 0
|
||||
|
||||
# Verify it's gone
|
||||
doc_after = await memory.get_document(document_id, bank_id)
|
||||
doc_after = await memory.get_document(document_id, bank_id, request_context=request_context)
|
||||
assert doc_after is None
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_memory_without_document(memory):
|
||||
async def test_memory_without_document(memory, request_context):
|
||||
"""Test that memories can still be created without document tracking."""
|
||||
bank_id = f"test_no_doc_{datetime.now(timezone.utc).timestamp()}"
|
||||
|
||||
@@ -122,10 +127,11 @@ async def test_memory_without_document(memory):
|
||||
units = await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="Alice works at Google.",
|
||||
context="Test"
|
||||
context="Test",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert len(units) > 0
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@@ -0,0 +1,796 @@
|
||||
"""Tests for the Hindsight extensions system."""
|
||||
|
||||
from collections import defaultdict
|
||||
|
||||
import pytest
|
||||
from fastapi import APIRouter
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from hindsight_api.extensions import (
|
||||
ApiKeyTenantExtension,
|
||||
AuthenticationError,
|
||||
Extension,
|
||||
HttpExtension,
|
||||
OperationValidationError,
|
||||
OperationValidatorExtension,
|
||||
RecallContext,
|
||||
RecallResult,
|
||||
ReflectContext,
|
||||
ReflectResultContext,
|
||||
RequestContext,
|
||||
RetainContext,
|
||||
RetainResult,
|
||||
TenantContext,
|
||||
TenantExtension,
|
||||
ValidationResult,
|
||||
load_extension,
|
||||
)
|
||||
|
||||
|
||||
class TestExtensionLoader:
|
||||
"""Tests for extension loading and lifecycle."""
|
||||
|
||||
def test_load_extension_with_config(self, monkeypatch):
|
||||
"""Extension receives config from prefixed env vars and supports lifecycle."""
|
||||
monkeypatch.setenv(
|
||||
"HINDSIGHT_API_TEST_EXTENSION",
|
||||
"tests.test_extensions:LifecycleTestExtension",
|
||||
)
|
||||
monkeypatch.setenv("HINDSIGHT_API_TEST_API_URL", "https://example.com")
|
||||
monkeypatch.setenv("HINDSIGHT_API_TEST_MAX_RETRIES", "5")
|
||||
|
||||
ext = load_extension("TEST", Extension)
|
||||
|
||||
assert ext is not None
|
||||
assert ext.config["api_url"] == "https://example.com"
|
||||
assert ext.config["max_retries"] == "5"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_extension_lifecycle(self, monkeypatch):
|
||||
"""Extension on_startup and on_shutdown are called."""
|
||||
monkeypatch.setenv(
|
||||
"HINDSIGHT_API_TEST_EXTENSION",
|
||||
"tests.test_extensions:LifecycleTestExtension",
|
||||
)
|
||||
|
||||
ext = load_extension("TEST", Extension)
|
||||
|
||||
assert not ext.started
|
||||
assert not ext.stopped
|
||||
|
||||
await ext.on_startup()
|
||||
assert ext.started
|
||||
|
||||
await ext.on_shutdown()
|
||||
assert ext.stopped
|
||||
|
||||
|
||||
class LifecycleTestExtension(Extension):
|
||||
"""Test extension for config and lifecycle tests."""
|
||||
|
||||
def __init__(self, config):
|
||||
super().__init__(config)
|
||||
self.started = False
|
||||
self.stopped = False
|
||||
|
||||
async def on_startup(self):
|
||||
self.started = True
|
||||
|
||||
async def on_shutdown(self):
|
||||
self.stopped = True
|
||||
|
||||
|
||||
class RateLimitingValidator(OperationValidatorExtension):
|
||||
"""
|
||||
Mock validator that blocks after N attempts per bank_id.
|
||||
|
||||
Used for testing the extension integration with MemoryEngine.
|
||||
"""
|
||||
|
||||
def __init__(self, config: dict):
|
||||
super().__init__(config)
|
||||
self.max_attempts = int(config.get("max_attempts", "2"))
|
||||
self.retain_counts: dict[str, int] = defaultdict(int)
|
||||
self.recall_counts: dict[str, int] = defaultdict(int)
|
||||
self.reflect_counts: dict[str, int] = defaultdict(int)
|
||||
|
||||
async def validate_retain(self, ctx: RetainContext) -> ValidationResult:
|
||||
self.retain_counts[ctx.bank_id] += 1
|
||||
if self.retain_counts[ctx.bank_id] > self.max_attempts:
|
||||
return ValidationResult.reject(
|
||||
f"Retain limit exceeded for bank {ctx.bank_id}"
|
||||
)
|
||||
return ValidationResult.accept()
|
||||
|
||||
async def validate_recall(self, ctx: RecallContext) -> ValidationResult:
|
||||
self.recall_counts[ctx.bank_id] += 1
|
||||
if self.recall_counts[ctx.bank_id] > self.max_attempts:
|
||||
return ValidationResult.reject(
|
||||
f"Recall limit exceeded for bank {ctx.bank_id}"
|
||||
)
|
||||
return ValidationResult.accept()
|
||||
|
||||
async def validate_reflect(self, ctx: ReflectContext) -> ValidationResult:
|
||||
self.reflect_counts[ctx.bank_id] += 1
|
||||
if self.reflect_counts[ctx.bank_id] > self.max_attempts:
|
||||
return ValidationResult.reject(
|
||||
f"Reflect limit exceeded for bank {ctx.bank_id}"
|
||||
)
|
||||
return ValidationResult.accept()
|
||||
|
||||
|
||||
class TrackingValidator(OperationValidatorExtension):
|
||||
"""
|
||||
Mock validator that tracks all pre and post hook calls with full parameters.
|
||||
|
||||
Used for testing that hooks receive all user-provided parameters.
|
||||
"""
|
||||
|
||||
def __init__(self, config: dict):
|
||||
super().__init__(config)
|
||||
# Pre-hook tracking
|
||||
self.pre_retain_calls: list[RetainContext] = []
|
||||
self.pre_recall_calls: list[RecallContext] = []
|
||||
self.pre_reflect_calls: list[ReflectContext] = []
|
||||
# Post-hook tracking
|
||||
self.post_retain_calls: list[RetainResult] = []
|
||||
self.post_recall_calls: list[RecallResult] = []
|
||||
self.post_reflect_calls: list[ReflectResultContext] = []
|
||||
|
||||
async def validate_retain(self, ctx: RetainContext) -> ValidationResult:
|
||||
self.pre_retain_calls.append(ctx)
|
||||
return ValidationResult.accept()
|
||||
|
||||
async def validate_recall(self, ctx: RecallContext) -> ValidationResult:
|
||||
self.pre_recall_calls.append(ctx)
|
||||
return ValidationResult.accept()
|
||||
|
||||
async def validate_reflect(self, ctx: ReflectContext) -> ValidationResult:
|
||||
self.pre_reflect_calls.append(ctx)
|
||||
return ValidationResult.accept()
|
||||
|
||||
async def on_retain_complete(self, result: RetainResult) -> None:
|
||||
self.post_retain_calls.append(result)
|
||||
|
||||
async def on_recall_complete(self, result: RecallResult) -> None:
|
||||
self.post_recall_calls.append(result)
|
||||
|
||||
async def on_reflect_complete(self, result: ReflectResultContext) -> None:
|
||||
self.post_reflect_calls.append(result)
|
||||
|
||||
|
||||
class TestMemoryEngineValidation:
|
||||
"""Tests for validation integration with MemoryEngine.
|
||||
|
||||
The OperationValidatorExtension is integrated at the MemoryEngine level,
|
||||
so all interfaces (HTTP API, MCP, SDK) get the same validation behavior.
|
||||
|
||||
For retain, the batch is validated as a whole (all or nothing) using
|
||||
retain_batch_async which is the public method used by the HTTP API.
|
||||
"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retain_batch_validation(self, memory_with_validator):
|
||||
"""Retain batch is validated as a whole - accepts or rejects entire batch."""
|
||||
memory = memory_with_validator
|
||||
bank_id = "test-retain-batch"
|
||||
ctx = RequestContext()
|
||||
|
||||
# First batch should succeed
|
||||
await memory.retain_batch_async(
|
||||
bank_id=bank_id,
|
||||
contents=[
|
||||
{"content": "First item"},
|
||||
{"content": "Second item"},
|
||||
],
|
||||
request_context=ctx,
|
||||
)
|
||||
|
||||
# Second batch should succeed (2nd attempt)
|
||||
await memory.retain_batch_async(
|
||||
bank_id=bank_id,
|
||||
contents=[{"content": "Third item"}],
|
||||
request_context=ctx,
|
||||
)
|
||||
|
||||
# Third batch should be blocked entirely (exceeds limit)
|
||||
with pytest.raises(OperationValidationError) as exc_info:
|
||||
await memory.retain_batch_async(
|
||||
bank_id=bank_id,
|
||||
contents=[
|
||||
{"content": "Should not be stored"},
|
||||
{"content": "Neither should this"},
|
||||
],
|
||||
request_context=ctx,
|
||||
)
|
||||
|
||||
assert "limit exceeded" in str(exc_info.value).lower()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_recall_validation(self, memory_with_validator):
|
||||
"""Recall is validated before execution."""
|
||||
memory = memory_with_validator
|
||||
bank_id = "test-recall-validation"
|
||||
ctx = RequestContext()
|
||||
|
||||
# First recall should pass validation
|
||||
await memory.recall_async(bank_id, "test query", fact_type=["world"], request_context=ctx)
|
||||
|
||||
# Second recall should pass validation
|
||||
await memory.recall_async(bank_id, "another query", fact_type=["world"], request_context=ctx)
|
||||
|
||||
# Third recall should be blocked by validator
|
||||
with pytest.raises(OperationValidationError) as exc_info:
|
||||
await memory.recall_async(bank_id, "blocked query", fact_type=["world"], request_context=ctx)
|
||||
|
||||
assert "limit exceeded" in str(exc_info.value).lower()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reflect_validation(self, memory_with_validator):
|
||||
"""Reflect is validated before execution."""
|
||||
memory = memory_with_validator
|
||||
bank_id = "test-reflect-validation"
|
||||
ctx = RequestContext()
|
||||
|
||||
# First reflect should pass validation (may fail internally but validation passes)
|
||||
try:
|
||||
await memory.reflect_async(bank_id, "test question", request_context=ctx)
|
||||
except OperationValidationError:
|
||||
raise # Re-raise validation errors
|
||||
except Exception:
|
||||
pass # Other errors are fine (e.g., no data)
|
||||
|
||||
# Second reflect should pass validation
|
||||
try:
|
||||
await memory.reflect_async(bank_id, "another question", request_context=ctx)
|
||||
except OperationValidationError:
|
||||
raise
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Third reflect should be blocked by validator
|
||||
with pytest.raises(OperationValidationError) as exc_info:
|
||||
await memory.reflect_async(bank_id, "blocked question", request_context=ctx)
|
||||
|
||||
assert "limit exceeded" in str(exc_info.value).lower()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def memory_with_validator(memory):
|
||||
"""Memory engine with a rate-limiting validator (max 2 attempts per bank)."""
|
||||
validator = RateLimitingValidator({"max_attempts": "2"})
|
||||
memory._operation_validator = validator
|
||||
return memory
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def memory_with_tracking_validator(memory):
|
||||
"""Memory engine with a tracking validator that records all hook calls."""
|
||||
validator = TrackingValidator({})
|
||||
memory._operation_validator = validator
|
||||
return memory, validator
|
||||
|
||||
|
||||
class TestOperationHooksParameters:
|
||||
"""Tests for pre and post operation hooks receiving all user-provided parameters."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retain_pre_hook_receives_all_parameters(self, memory_with_tracking_validator):
|
||||
"""Pre-retain hook receives all user-provided parameters."""
|
||||
memory, validator = memory_with_tracking_validator
|
||||
bank_id = "test-retain-params"
|
||||
ctx = RequestContext(api_key="test-key")
|
||||
contents = [{"content": "Test content", "context": "test context"}]
|
||||
document_id = "doc-123"
|
||||
|
||||
await memory.retain_batch_async(
|
||||
bank_id=bank_id,
|
||||
contents=contents,
|
||||
document_id=document_id,
|
||||
fact_type_override="world",
|
||||
confidence_score=0.9,
|
||||
request_context=ctx,
|
||||
)
|
||||
|
||||
assert len(validator.pre_retain_calls) == 1
|
||||
pre_ctx = validator.pre_retain_calls[0]
|
||||
|
||||
# Verify all parameters are present
|
||||
assert pre_ctx.bank_id == bank_id
|
||||
# Note: contents is copied before document_id is applied to individual items
|
||||
assert len(pre_ctx.contents) == len(contents)
|
||||
assert pre_ctx.contents[0]["content"] == contents[0]["content"]
|
||||
assert pre_ctx.document_id == document_id
|
||||
assert pre_ctx.fact_type_override == "world"
|
||||
assert pre_ctx.confidence_score == 0.9
|
||||
assert pre_ctx.request_context == ctx
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retain_post_hook_receives_all_parameters_and_result(self, memory_with_tracking_validator):
|
||||
"""Post-retain hook receives all parameters plus the result."""
|
||||
memory, validator = memory_with_tracking_validator
|
||||
bank_id = "test-retain-post"
|
||||
ctx = RequestContext(api_key="test-key")
|
||||
contents = [{"content": "Test content for post hook"}]
|
||||
document_id = "doc-456"
|
||||
|
||||
result = await memory.retain_batch_async(
|
||||
bank_id=bank_id,
|
||||
contents=contents,
|
||||
document_id=document_id,
|
||||
fact_type_override="experience",
|
||||
confidence_score=0.8,
|
||||
request_context=ctx,
|
||||
)
|
||||
|
||||
assert len(validator.post_retain_calls) == 1
|
||||
post_result = validator.post_retain_calls[0]
|
||||
|
||||
# Verify all parameters are present
|
||||
assert post_result.bank_id == bank_id
|
||||
assert post_result.document_id == document_id
|
||||
assert post_result.fact_type_override == "experience"
|
||||
assert post_result.confidence_score == 0.8
|
||||
assert post_result.request_context == ctx
|
||||
|
||||
# Verify result data
|
||||
assert post_result.success is True
|
||||
assert post_result.error is None
|
||||
assert post_result.unit_ids == result # Should match the return value
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_recall_pre_hook_receives_all_parameters(self, memory_with_tracking_validator):
|
||||
"""Pre-recall hook receives all user-provided parameters."""
|
||||
from datetime import datetime, timezone
|
||||
from hindsight_api.engine.memory_engine import Budget
|
||||
|
||||
memory, validator = memory_with_tracking_validator
|
||||
bank_id = "test-recall-params"
|
||||
ctx = RequestContext(api_key="test-key")
|
||||
query = "test query"
|
||||
question_date = datetime(2024, 1, 15, tzinfo=timezone.utc)
|
||||
|
||||
await memory.recall_async(
|
||||
bank_id=bank_id,
|
||||
query=query,
|
||||
budget=Budget.HIGH,
|
||||
max_tokens=2048,
|
||||
enable_trace=True,
|
||||
fact_type=["world", "experience"],
|
||||
question_date=question_date,
|
||||
include_entities=True,
|
||||
max_entity_tokens=300,
|
||||
include_chunks=True,
|
||||
max_chunk_tokens=4096,
|
||||
request_context=ctx,
|
||||
)
|
||||
|
||||
assert len(validator.pre_recall_calls) == 1
|
||||
pre_ctx = validator.pre_recall_calls[0]
|
||||
|
||||
# Verify all parameters are present
|
||||
assert pre_ctx.bank_id == bank_id
|
||||
assert pre_ctx.query == query
|
||||
assert pre_ctx.budget == Budget.HIGH
|
||||
assert pre_ctx.max_tokens == 2048
|
||||
assert pre_ctx.enable_trace is True
|
||||
assert pre_ctx.fact_types == ["world", "experience"]
|
||||
assert pre_ctx.question_date == question_date
|
||||
assert pre_ctx.include_entities is True
|
||||
assert pre_ctx.max_entity_tokens == 300
|
||||
assert pre_ctx.include_chunks is True
|
||||
assert pre_ctx.max_chunk_tokens == 4096
|
||||
assert pre_ctx.request_context == ctx
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_recall_post_hook_receives_all_parameters_and_result(self, memory_with_tracking_validator):
|
||||
"""Post-recall hook receives all parameters plus the result."""
|
||||
from hindsight_api.engine.memory_engine import Budget
|
||||
|
||||
memory, validator = memory_with_tracking_validator
|
||||
bank_id = "test-recall-post"
|
||||
ctx = RequestContext(api_key="test-key")
|
||||
|
||||
result = await memory.recall_async(
|
||||
bank_id=bank_id,
|
||||
query="test query for post",
|
||||
budget=Budget.LOW,
|
||||
max_tokens=1024,
|
||||
fact_type=["world"],
|
||||
request_context=ctx,
|
||||
)
|
||||
|
||||
assert len(validator.post_recall_calls) == 1
|
||||
post_result = validator.post_recall_calls[0]
|
||||
|
||||
# Verify all parameters are present
|
||||
assert post_result.bank_id == bank_id
|
||||
assert post_result.query == "test query for post"
|
||||
assert post_result.budget == Budget.LOW
|
||||
assert post_result.max_tokens == 1024
|
||||
assert post_result.fact_types == ["world"]
|
||||
assert post_result.request_context == ctx
|
||||
|
||||
# Verify result data
|
||||
assert post_result.success is True
|
||||
assert post_result.error is None
|
||||
assert post_result.result == result # Should match the return value
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reflect_pre_hook_receives_all_parameters(self, memory_with_tracking_validator):
|
||||
"""Pre-reflect hook receives all user-provided parameters."""
|
||||
from hindsight_api.engine.memory_engine import Budget
|
||||
|
||||
memory, validator = memory_with_tracking_validator
|
||||
bank_id = "test-reflect-params"
|
||||
ctx = RequestContext(api_key="test-key")
|
||||
|
||||
try:
|
||||
await memory.reflect_async(
|
||||
bank_id=bank_id,
|
||||
query="test question",
|
||||
budget=Budget.MID,
|
||||
context="additional context",
|
||||
request_context=ctx,
|
||||
)
|
||||
except Exception:
|
||||
pass # May fail if no data, but pre-hook should still be called
|
||||
|
||||
assert len(validator.pre_reflect_calls) == 1
|
||||
pre_ctx = validator.pre_reflect_calls[0]
|
||||
|
||||
# Verify all parameters are present
|
||||
assert pre_ctx.bank_id == bank_id
|
||||
assert pre_ctx.query == "test question"
|
||||
assert pre_ctx.budget == Budget.MID
|
||||
assert pre_ctx.context == "additional context"
|
||||
assert pre_ctx.request_context == ctx
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reflect_post_hook_receives_all_parameters_and_result(self, memory_with_tracking_validator):
|
||||
"""Post-reflect hook receives all parameters plus the result on success."""
|
||||
from hindsight_api.engine.memory_engine import Budget
|
||||
|
||||
memory, validator = memory_with_tracking_validator
|
||||
bank_id = "test-reflect-post"
|
||||
ctx = RequestContext(api_key="test-key")
|
||||
|
||||
# Store some content first so reflect has something to work with
|
||||
await memory.retain_batch_async(
|
||||
bank_id=bank_id,
|
||||
contents=[{"content": "Alice is a software engineer at Google."}],
|
||||
request_context=ctx,
|
||||
)
|
||||
|
||||
result = await memory.reflect_async(
|
||||
bank_id=bank_id,
|
||||
query="What does Alice do?",
|
||||
budget=Budget.LOW,
|
||||
context="work context",
|
||||
request_context=ctx,
|
||||
)
|
||||
|
||||
assert len(validator.post_reflect_calls) == 1
|
||||
post_result = validator.post_reflect_calls[0]
|
||||
|
||||
# Verify all parameters are present
|
||||
assert post_result.bank_id == bank_id
|
||||
assert post_result.query == "What does Alice do?"
|
||||
assert post_result.budget == Budget.LOW
|
||||
assert post_result.context == "work context"
|
||||
assert post_result.request_context == ctx
|
||||
|
||||
# Verify result data
|
||||
assert post_result.success is True
|
||||
assert post_result.error is None
|
||||
assert post_result.result == result # Should match the return value
|
||||
assert post_result.result.text is not None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_hooks_called_in_order_after_pre_hooks(self, memory_with_tracking_validator):
|
||||
"""Post hooks are called after pre hooks and after operation completes."""
|
||||
memory, validator = memory_with_tracking_validator
|
||||
bank_id = "test-hook-order"
|
||||
ctx = RequestContext()
|
||||
|
||||
# Retain operation
|
||||
await memory.retain_batch_async(
|
||||
bank_id=bank_id,
|
||||
contents=[{"content": "Test content"}],
|
||||
request_context=ctx,
|
||||
)
|
||||
|
||||
# Pre-hook should be called before post-hook
|
||||
assert len(validator.pre_retain_calls) == 1
|
||||
assert len(validator.post_retain_calls) == 1
|
||||
|
||||
# Recall operation
|
||||
await memory.recall_async(
|
||||
bank_id=bank_id,
|
||||
query="test",
|
||||
fact_type=["world"],
|
||||
request_context=ctx,
|
||||
)
|
||||
|
||||
assert len(validator.pre_recall_calls) == 1
|
||||
assert len(validator.post_recall_calls) == 1
|
||||
|
||||
|
||||
class TestTenantExtension:
|
||||
"""Tests for TenantExtension and ApiKeyTenantExtension."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_api_key_tenant_extension_valid_key(self):
|
||||
"""ApiKeyTenantExtension accepts valid API key."""
|
||||
ext = ApiKeyTenantExtension({"api_key": "secret-key-123"})
|
||||
|
||||
result = await ext.authenticate(RequestContext(api_key="secret-key-123"))
|
||||
|
||||
assert result.schema_name == "public"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_api_key_tenant_extension_invalid_key(self):
|
||||
"""ApiKeyTenantExtension rejects invalid API key."""
|
||||
ext = ApiKeyTenantExtension({"api_key": "secret-key-123"})
|
||||
|
||||
with pytest.raises(AuthenticationError) as exc_info:
|
||||
await ext.authenticate(RequestContext(api_key="wrong-key"))
|
||||
|
||||
assert "Invalid API key" in str(exc_info.value)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_api_key_tenant_extension_missing_key(self):
|
||||
"""ApiKeyTenantExtension rejects missing API key."""
|
||||
ext = ApiKeyTenantExtension({"api_key": "secret-key-123"})
|
||||
|
||||
with pytest.raises(AuthenticationError):
|
||||
await ext.authenticate(RequestContext(api_key=None))
|
||||
|
||||
def test_api_key_tenant_extension_requires_config(self):
|
||||
"""ApiKeyTenantExtension requires api_key in config."""
|
||||
with pytest.raises(ValueError) as exc_info:
|
||||
ApiKeyTenantExtension({})
|
||||
|
||||
assert "HINDSIGHT_API_TENANT_API_KEY is required" in str(exc_info.value)
|
||||
|
||||
|
||||
class TestMemoryEngineTenantAuth:
|
||||
"""Tests for tenant authentication in MemoryEngine."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retain_requires_tenant_request_when_extension_configured(
|
||||
self, memory_with_tenant
|
||||
):
|
||||
"""Retain fails without RequestContext when tenant extension is configured."""
|
||||
memory = memory_with_tenant
|
||||
|
||||
with pytest.raises(AuthenticationError) as exc_info:
|
||||
await memory.retain_batch_async(
|
||||
bank_id="test-bank",
|
||||
contents=[{"content": "test"}],
|
||||
request_context=None, # Missing!
|
||||
)
|
||||
|
||||
assert "RequestContext is required" in str(exc_info.value)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retain_succeeds_with_valid_tenant_request(self, memory_with_tenant):
|
||||
"""Retain succeeds with valid RequestContext."""
|
||||
memory = memory_with_tenant
|
||||
|
||||
# Should not raise
|
||||
await memory.retain_batch_async(
|
||||
bank_id="test-bank-tenant",
|
||||
contents=[{"content": "test content"}],
|
||||
request_context=RequestContext(api_key="test-api-key"),
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retain_fails_with_invalid_api_key(self, memory_with_tenant):
|
||||
"""Retain fails with invalid API key."""
|
||||
memory = memory_with_tenant
|
||||
|
||||
with pytest.raises(AuthenticationError) as exc_info:
|
||||
await memory.retain_batch_async(
|
||||
bank_id="test-bank",
|
||||
contents=[{"content": "test"}],
|
||||
request_context=RequestContext(api_key="wrong-key"),
|
||||
)
|
||||
|
||||
assert "Invalid API key" in str(exc_info.value)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_recall_requires_tenant_request_when_extension_configured(
|
||||
self, memory_with_tenant
|
||||
):
|
||||
"""Recall fails without RequestContext when tenant extension is configured."""
|
||||
memory = memory_with_tenant
|
||||
|
||||
with pytest.raises(AuthenticationError):
|
||||
await memory.recall_async(
|
||||
bank_id="test-bank",
|
||||
query="test query",
|
||||
fact_type=["world"],
|
||||
request_context=None,
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_tenant_request_needed_without_extension(self, memory):
|
||||
"""Operations work with empty RequestContext when no tenant extension configured."""
|
||||
# Should not raise - no tenant extension configured, just pass empty RequestContext
|
||||
await memory.retain_batch_async(
|
||||
bank_id="test-bank-no-tenant",
|
||||
contents=[{"content": "test content"}],
|
||||
request_context=RequestContext(),
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def memory_with_tenant(memory):
|
||||
"""Memory engine with a tenant extension (API key auth)."""
|
||||
tenant_ext = ApiKeyTenantExtension({"api_key": "test-api-key"})
|
||||
memory._tenant_extension = tenant_ext
|
||||
return memory
|
||||
|
||||
|
||||
class SampleHttpExtension(HttpExtension):
|
||||
"""Sample HTTP extension for testing that provides custom endpoints."""
|
||||
|
||||
def __init__(self, config: dict):
|
||||
super().__init__(config)
|
||||
self.started = False
|
||||
self.stopped = False
|
||||
self.request_count = 0
|
||||
|
||||
async def on_startup(self):
|
||||
self.started = True
|
||||
|
||||
async def on_shutdown(self):
|
||||
self.stopped = True
|
||||
|
||||
def get_router(self, memory) -> APIRouter:
|
||||
router = APIRouter()
|
||||
|
||||
@router.get("/hello")
|
||||
async def hello():
|
||||
self.request_count += 1
|
||||
return {"message": "Hello from extension!"}
|
||||
|
||||
@router.get("/config")
|
||||
async def get_config():
|
||||
return {"config": self.config}
|
||||
|
||||
@router.get("/health-check")
|
||||
async def extension_health():
|
||||
health = await memory.health_check()
|
||||
return {"extension": "healthy", "memory": health}
|
||||
|
||||
@router.post("/echo")
|
||||
async def echo(data: dict):
|
||||
return {"echoed": data}
|
||||
|
||||
return router
|
||||
|
||||
|
||||
class TestHttpExtensionIntegration:
|
||||
"""Tests for HTTP extension integration."""
|
||||
|
||||
def test_load_http_extension(self, monkeypatch):
|
||||
"""HttpExtension can be loaded from environment variable."""
|
||||
monkeypatch.setenv(
|
||||
"HINDSIGHT_API_HTTP_EXTENSION",
|
||||
"tests.test_extensions:SampleHttpExtension",
|
||||
)
|
||||
monkeypatch.setenv("HINDSIGHT_API_HTTP_CUSTOM_PARAM", "custom_value")
|
||||
|
||||
ext = load_extension("HTTP", HttpExtension)
|
||||
|
||||
assert ext is not None
|
||||
assert isinstance(ext, SampleHttpExtension)
|
||||
assert ext.config["custom_param"] == "custom_value"
|
||||
|
||||
def test_http_extension_router_mounted_at_ext(self, memory):
|
||||
"""HTTP extension router is mounted at /ext/."""
|
||||
from hindsight_api.api.http import create_app
|
||||
|
||||
ext = SampleHttpExtension({"test_key": "test_value"})
|
||||
app = create_app(memory, initialize_memory=False, http_extension=ext)
|
||||
|
||||
client = TestClient(app)
|
||||
|
||||
# Extension endpoint should be accessible at /ext/
|
||||
response = client.get("/ext/hello")
|
||||
assert response.status_code == 200
|
||||
assert response.json() == {"message": "Hello from extension!"}
|
||||
|
||||
# Should track request count
|
||||
assert ext.request_count == 1
|
||||
|
||||
# Old path should NOT work
|
||||
response = client.get("/extension/hello")
|
||||
assert response.status_code == 404
|
||||
|
||||
def test_http_extension_config_endpoint(self, memory):
|
||||
"""Extension can expose its config via custom endpoint."""
|
||||
from hindsight_api.api.http import create_app
|
||||
|
||||
ext = SampleHttpExtension({"api_key": "secret", "limit": "100"})
|
||||
app = create_app(memory, initialize_memory=False, http_extension=ext)
|
||||
|
||||
client = TestClient(app)
|
||||
|
||||
response = client.get("/ext/config")
|
||||
assert response.status_code == 200
|
||||
assert response.json()["config"]["api_key"] == "secret"
|
||||
assert response.json()["config"]["limit"] == "100"
|
||||
|
||||
def test_http_extension_can_access_memory(self, memory):
|
||||
"""Extension endpoints can access memory engine."""
|
||||
from hindsight_api.api.http import create_app
|
||||
|
||||
ext = SampleHttpExtension({})
|
||||
app = create_app(memory, initialize_memory=False, http_extension=ext)
|
||||
|
||||
client = TestClient(app)
|
||||
|
||||
response = client.get("/ext/health-check")
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["extension"] == "healthy"
|
||||
assert "memory" in data
|
||||
|
||||
def test_http_extension_post_endpoint(self, memory):
|
||||
"""Extension can handle POST requests with JSON body."""
|
||||
from hindsight_api.api.http import create_app
|
||||
|
||||
ext = SampleHttpExtension({})
|
||||
app = create_app(memory, initialize_memory=False, http_extension=ext)
|
||||
|
||||
client = TestClient(app)
|
||||
|
||||
response = client.post("/ext/echo", json={"key": "value", "number": 42})
|
||||
assert response.status_code == 200
|
||||
assert response.json() == {"echoed": {"key": "value", "number": 42}}
|
||||
|
||||
def test_http_extension_not_mounted_when_none(self, memory):
|
||||
"""No extension routes when http_extension is None."""
|
||||
from hindsight_api.api.http import create_app
|
||||
|
||||
app = create_app(memory, initialize_memory=False, http_extension=None)
|
||||
|
||||
client = TestClient(app)
|
||||
|
||||
# Extension endpoint should not exist
|
||||
response = client.get("/ext/hello")
|
||||
assert response.status_code == 404
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_http_extension_lifecycle(self):
|
||||
"""HTTP extension on_startup and on_shutdown are called."""
|
||||
ext = SampleHttpExtension({})
|
||||
|
||||
assert not ext.started
|
||||
assert not ext.stopped
|
||||
|
||||
await ext.on_startup()
|
||||
assert ext.started
|
||||
|
||||
await ext.on_shutdown()
|
||||
assert ext.stopped
|
||||
|
||||
def test_core_routes_still_work_with_extension(self, memory):
|
||||
"""Core API routes still work when extension is mounted."""
|
||||
from hindsight_api.api.http import create_app
|
||||
|
||||
ext = SampleHttpExtension({})
|
||||
app = create_app(memory, initialize_memory=False, http_extension=ext)
|
||||
|
||||
client = TestClient(app)
|
||||
|
||||
# Health endpoint should work
|
||||
response = client.get("/health")
|
||||
assert response.status_code in (200, 503) # May be unhealthy if DB not connected
|
||||
|
||||
# Banks list endpoint should work
|
||||
response = client.get("/v1/default/banks")
|
||||
assert response.status_code in (200, 500) # May fail if DB not ready
|
||||
@@ -897,7 +897,7 @@ class TestDispositionInference:
|
||||
"""Tests for LLM-based disposition trait inference from background."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_background_merge_with_disposition_inference(self, memory):
|
||||
async def test_background_merge_with_disposition_inference(self, memory, request_context):
|
||||
"""Test that background merge infers disposition traits by default."""
|
||||
import uuid
|
||||
bank_id = f"test_infer_{uuid.uuid4().hex[:8]}"
|
||||
@@ -905,7 +905,8 @@ class TestDispositionInference:
|
||||
result = await memory.merge_bank_background(
|
||||
bank_id,
|
||||
"I am a creative software engineer who loves innovation and trying new technologies",
|
||||
update_disposition=True
|
||||
update_disposition=True,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert "background" in result
|
||||
@@ -923,30 +924,31 @@ class TestDispositionInference:
|
||||
assert 1 <= disposition[trait] <= 5
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_background_merge_without_disposition_inference(self, memory):
|
||||
async def test_background_merge_without_disposition_inference(self, memory, request_context):
|
||||
"""Test that background merge skips disposition inference when disabled."""
|
||||
import uuid
|
||||
bank_id = f"test_no_infer_{uuid.uuid4().hex[:8]}"
|
||||
|
||||
initial_profile = await memory.get_bank_profile(bank_id)
|
||||
initial_profile = await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
initial_disposition = initial_profile["disposition"]
|
||||
|
||||
result = await memory.merge_bank_background(
|
||||
bank_id,
|
||||
"I am a data scientist",
|
||||
update_disposition=False
|
||||
update_disposition=False,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert "background" in result
|
||||
assert "disposition" not in result
|
||||
|
||||
final_profile = await memory.get_bank_profile(bank_id)
|
||||
final_profile = await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
final_disposition = final_profile["disposition"]
|
||||
|
||||
assert initial_disposition == final_disposition
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_disposition_inference_for_lawyer(self, memory):
|
||||
async def test_disposition_inference_for_lawyer(self, memory, request_context):
|
||||
"""Test disposition inference for lawyer profile (high skepticism, high literalism)."""
|
||||
import uuid
|
||||
bank_id = f"test_lawyer_{uuid.uuid4().hex[:8]}"
|
||||
@@ -954,7 +956,8 @@ class TestDispositionInference:
|
||||
result = await memory.merge_bank_background(
|
||||
bank_id,
|
||||
"I am a lawyer who focuses on contract details and never takes claims at face value",
|
||||
update_disposition=True
|
||||
update_disposition=True,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
disposition = result["disposition"]
|
||||
@@ -964,7 +967,7 @@ class TestDispositionInference:
|
||||
assert disposition["literalism"] >= 3
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_disposition_inference_for_therapist(self, memory):
|
||||
async def test_disposition_inference_for_therapist(self, memory, request_context):
|
||||
"""Test disposition inference for therapist profile (high empathy)."""
|
||||
import uuid
|
||||
bank_id = f"test_therapist_{uuid.uuid4().hex[:8]}"
|
||||
@@ -972,7 +975,8 @@ class TestDispositionInference:
|
||||
result = await memory.merge_bank_background(
|
||||
bank_id,
|
||||
"I am a therapist who deeply understands and connects with people's emotional struggles",
|
||||
update_disposition=True
|
||||
update_disposition=True,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
disposition = result["disposition"]
|
||||
@@ -981,7 +985,7 @@ class TestDispositionInference:
|
||||
assert disposition["empathy"] >= 3
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_disposition_updates_in_database(self, memory):
|
||||
async def test_disposition_updates_in_database(self, memory, request_context):
|
||||
"""Test that inferred disposition is actually stored in database."""
|
||||
import uuid
|
||||
bank_id = f"test_db_update_{uuid.uuid4().hex[:8]}"
|
||||
@@ -989,12 +993,13 @@ class TestDispositionInference:
|
||||
result = await memory.merge_bank_background(
|
||||
bank_id,
|
||||
"I am an innovative designer",
|
||||
update_disposition=True
|
||||
update_disposition=True,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
inferred_disposition = result["disposition"]
|
||||
|
||||
profile = await memory.get_bank_profile(bank_id)
|
||||
profile = await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
db_disposition = profile["disposition"]
|
||||
|
||||
# Compare values (db_disposition is a Pydantic model)
|
||||
@@ -1003,7 +1008,7 @@ class TestDispositionInference:
|
||||
assert db_disposition.empathy == inferred_disposition["empathy"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_multiple_background_merges_update_disposition(self, memory):
|
||||
async def test_multiple_background_merges_update_disposition(self, memory, request_context):
|
||||
"""Test that each background merge can update disposition."""
|
||||
import uuid
|
||||
bank_id = f"test_multi_merge_{uuid.uuid4().hex[:8]}"
|
||||
@@ -1011,14 +1016,16 @@ class TestDispositionInference:
|
||||
result1 = await memory.merge_bank_background(
|
||||
bank_id,
|
||||
"I am a software engineer",
|
||||
update_disposition=True
|
||||
update_disposition=True,
|
||||
request_context=request_context,
|
||||
)
|
||||
disposition1 = result1["disposition"]
|
||||
|
||||
result2 = await memory.merge_bank_background(
|
||||
bank_id,
|
||||
"I love creative problem solving and innovation",
|
||||
update_disposition=True
|
||||
update_disposition=True,
|
||||
request_context=request_context,
|
||||
)
|
||||
disposition2 = result2["disposition"]
|
||||
|
||||
@@ -1026,7 +1033,7 @@ class TestDispositionInference:
|
||||
assert "creative" in result2["background"].lower() or "innovation" in result2["background"].lower()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_background_merge_conflict_resolution_with_disposition(self, memory):
|
||||
async def test_background_merge_conflict_resolution_with_disposition(self, memory, request_context):
|
||||
"""Test that conflicts are resolved and disposition reflects final background."""
|
||||
import uuid
|
||||
bank_id = f"test_conflict_{uuid.uuid4().hex[:8]}"
|
||||
@@ -1034,13 +1041,15 @@ class TestDispositionInference:
|
||||
await memory.merge_bank_background(
|
||||
bank_id,
|
||||
"I was born in Colorado and prefer stability",
|
||||
update_disposition=True
|
||||
update_disposition=True,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
result = await memory.merge_bank_background(
|
||||
bank_id,
|
||||
"You were born in Texas and are very skeptical of people",
|
||||
update_disposition=True
|
||||
update_disposition=True,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
background = result["background"]
|
||||
|
||||
@@ -7,24 +7,24 @@ distinguish between things said earlier vs later.
|
||||
"""
|
||||
import pytest
|
||||
from datetime import datetime, timezone
|
||||
from hindsight_api import MemoryEngine
|
||||
from hindsight_api import MemoryEngine, RequestContext
|
||||
from hindsight_api.engine.memory_engine import Budget
|
||||
import os
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fact_ordering_within_conversation(memory):
|
||||
async def test_fact_ordering_within_conversation(memory, request_context):
|
||||
bank_id = "test_ordering_agent"
|
||||
|
||||
# Get/create agent (auto-creates with defaults)
|
||||
await memory.get_bank_profile(bank_id)
|
||||
await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
|
||||
# Update disposition to match Marcus
|
||||
await memory.update_bank_disposition(bank_id, {
|
||||
"skepticism": 3,
|
||||
"literalism": 3,
|
||||
"empathy": 3
|
||||
})
|
||||
}, request_context=request_context)
|
||||
|
||||
# A conversation where Marcus changes his position
|
||||
conversation = """
|
||||
@@ -43,7 +43,8 @@ Marcus: Yeah, I realized I was being too optimistic about their defense.
|
||||
content=conversation,
|
||||
context="podcast discussion about NFL game",
|
||||
event_date=base_event_date,
|
||||
document_id="test_conv_1"
|
||||
document_id="test_conv_1",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Search for all facts about Marcus's predictions
|
||||
@@ -52,7 +53,8 @@ Marcus: Yeah, I realized I was being too optimistic about their defense.
|
||||
query="Marcus prediction Rams",
|
||||
fact_type=['opinion', 'experience', 'world'],
|
||||
budget=Budget.LOW,
|
||||
max_tokens=8192
|
||||
max_tokens=8192,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
print(f"\n=== Retrieved {len(results.results)} facts ===")
|
||||
@@ -113,17 +115,17 @@ Marcus: Yeah, I realized I was being too optimistic about their defense.
|
||||
print(f"\n✅ Temporal ordering preserved: First prediction came before changed prediction")
|
||||
|
||||
# Cleanup
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
print(f"\n✅ Test passed: Fact ordering within conversation is preserved")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_multiple_documents_ordering(memory):
|
||||
async def test_multiple_documents_ordering(memory, request_context):
|
||||
|
||||
bank_id = "test_multi_doc_agent"
|
||||
|
||||
await memory.get_bank_profile(bank_id) # Auto-creates with defaults
|
||||
await memory.get_bank_profile(bank_id, request_context=request_context) # Auto-creates with defaults
|
||||
|
||||
# Two separate conversations with same base time
|
||||
base_time = datetime(2024, 11, 14, 10, 0, 0, tzinfo=timezone.utc)
|
||||
@@ -146,7 +148,8 @@ Alice: I reconsidered the team's experience level.
|
||||
contents=[
|
||||
{"content": conv1, "context": "project discussion 1", "event_date": base_time},
|
||||
{"content": conv2, "context": "project discussion 2", "event_date": base_time}
|
||||
]
|
||||
],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Search for Alice's preferences
|
||||
@@ -155,7 +158,8 @@ Alice: I reconsidered the team's experience level.
|
||||
query="Alice preference React Vue",
|
||||
fact_type=['opinion', 'experience'],
|
||||
budget=Budget.LOW,
|
||||
max_tokens=8192
|
||||
max_tokens=8192,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
print(f"\n=== Retrieved {len(results.results)} agent facts ===")
|
||||
@@ -175,6 +179,6 @@ Alice: I reconsidered the team's experience level.
|
||||
print(f"\n✅ Facts from {len(agent_facts)} statements have {len(unique_timestamps)} unique timestamps")
|
||||
|
||||
# Cleanup
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
print(f"\n✅ Test passed: Multiple documents maintain separate ordering")
|
||||
|
||||
@@ -26,6 +26,9 @@ MODEL_MATRIX = [
|
||||
("gemini", "gemini-2.5-flash"),
|
||||
("gemini", "gemini-2.5-flash-lite"),
|
||||
("gemini", "gemini-3-pro-preview"),
|
||||
# Ollama models (local)
|
||||
("ollama", "gemma3:12b"),
|
||||
("ollama", "gemma3:1b"),
|
||||
]
|
||||
|
||||
|
||||
@@ -48,12 +51,18 @@ async def test_llm_provider_memory_operations(provider: str, model: str):
|
||||
All models must pass this test.
|
||||
"""
|
||||
api_key = get_api_key_for_provider(provider)
|
||||
if not api_key:
|
||||
|
||||
# Skip Ollama tests in CI (no models available)
|
||||
if provider == "ollama" and os.getenv("CI"):
|
||||
pytest.skip(f"Skipping {provider}/{model}: Ollama not available in CI")
|
||||
|
||||
# Other providers need an API key
|
||||
if provider != "ollama" and not api_key:
|
||||
pytest.skip(f"Skipping {provider}/{model}: no API key available")
|
||||
|
||||
llm = LLMProvider(
|
||||
provider=provider,
|
||||
api_key=api_key,
|
||||
api_key=api_key or "",
|
||||
base_url="",
|
||||
model=model,
|
||||
)
|
||||
|
||||
@@ -3,11 +3,12 @@ Test observation generation and entity state functionality.
|
||||
"""
|
||||
import pytest
|
||||
from hindsight_api.engine.memory_engine import Budget
|
||||
from hindsight_api import RequestContext
|
||||
from datetime import datetime, timezone
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_observation_generation_on_put(memory):
|
||||
async def test_observation_generation_on_put(memory, request_context):
|
||||
"""
|
||||
Test that observations are generated SYNCHRONOUSLY when new facts are added.
|
||||
|
||||
@@ -36,7 +37,8 @@ async def test_observation_generation_on_put(memory):
|
||||
bank_id=bank_id,
|
||||
content=content,
|
||||
context="work info",
|
||||
event_date=datetime(2024, 1, 15 + i, tzinfo=timezone.utc)
|
||||
event_date=datetime(2024, 1, 15 + i, tzinfo=timezone.utc),
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Observations are generated SYNCHRONOUSLY during retain,
|
||||
@@ -75,7 +77,7 @@ async def test_observation_generation_on_put(memory):
|
||||
print(f"Entity: {entity_name} (id: {entity_id})")
|
||||
|
||||
# Get observations for the entity - should be available immediately
|
||||
observations = await memory.get_entity_observations(bank_id, entity_id, limit=10)
|
||||
observations = await memory.get_entity_observations(bank_id, entity_id, limit=10, request_context=request_context)
|
||||
|
||||
print(f"\n=== Observations for {entity_name} ===")
|
||||
print(f"Total observations: {len(observations)}")
|
||||
@@ -102,7 +104,7 @@ async def test_observation_generation_on_put(memory):
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_regenerate_entity_observations(memory):
|
||||
async def test_regenerate_entity_observations(memory, request_context):
|
||||
"""
|
||||
Test explicit regeneration of observations for an entity.
|
||||
"""
|
||||
@@ -114,7 +116,8 @@ async def test_regenerate_entity_observations(memory):
|
||||
bank_id=bank_id,
|
||||
content="Sarah is a product manager who loves user research and data analysis.",
|
||||
context="work info",
|
||||
event_date=datetime(2024, 1, 15, tzinfo=timezone.utc)
|
||||
event_date=datetime(2024, 1, 15, tzinfo=timezone.utc),
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
await memory.wait_for_background_tasks()
|
||||
@@ -140,14 +143,15 @@ async def test_regenerate_entity_observations(memory):
|
||||
created_ids = await memory.regenerate_entity_observations(
|
||||
bank_id=bank_id,
|
||||
entity_id=entity_id,
|
||||
entity_name=entity_name
|
||||
entity_name=entity_name,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
print(f"\n=== Regenerated Observations ===")
|
||||
print(f"Created {len(created_ids)} observations for {entity_name}")
|
||||
|
||||
# Get the observations
|
||||
observations = await memory.get_entity_observations(bank_id, entity_id, limit=10)
|
||||
observations = await memory.get_entity_observations(bank_id, entity_id, limit=10, request_context=request_context)
|
||||
for obs in observations:
|
||||
print(f" - {obs.text}")
|
||||
|
||||
@@ -170,7 +174,7 @@ async def test_regenerate_entity_observations(memory):
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_search_with_include_entities(memory):
|
||||
async def test_search_with_include_entities(memory, request_context):
|
||||
"""
|
||||
Test that search with include_entities=True returns entity observations.
|
||||
|
||||
@@ -196,7 +200,8 @@ async def test_search_with_include_entities(memory):
|
||||
bank_id=bank_id,
|
||||
content=content,
|
||||
context="work info",
|
||||
event_date=datetime(2024, 1, 15 + i, tzinfo=timezone.utc)
|
||||
event_date=datetime(2024, 1, 15 + i, tzinfo=timezone.utc),
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Observations are generated synchronously during retain, no need to wait
|
||||
@@ -209,7 +214,8 @@ async def test_search_with_include_entities(memory):
|
||||
budget=Budget.LOW,
|
||||
max_tokens=2000,
|
||||
include_entities=True,
|
||||
max_entity_tokens=500
|
||||
max_entity_tokens=500,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
print(f"\n=== Search Results ===")
|
||||
@@ -263,7 +269,7 @@ async def test_search_with_include_entities(memory):
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_entity_state(memory):
|
||||
async def test_get_entity_state(memory, request_context):
|
||||
"""
|
||||
Test getting the full state of an entity.
|
||||
"""
|
||||
@@ -275,7 +281,8 @@ async def test_get_entity_state(memory):
|
||||
bank_id=bank_id,
|
||||
content="Bob is a frontend developer who specializes in React and TypeScript.",
|
||||
context="work info",
|
||||
event_date=datetime(2024, 1, 15, tzinfo=timezone.utc)
|
||||
event_date=datetime(2024, 1, 15, tzinfo=timezone.utc),
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
await memory.wait_for_background_tasks()
|
||||
@@ -302,7 +309,8 @@ async def test_get_entity_state(memory):
|
||||
bank_id=bank_id,
|
||||
entity_id=entity_id,
|
||||
entity_name=entity_name,
|
||||
limit=10
|
||||
limit=10,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
print(f"\n=== Entity State for {entity_name} ===")
|
||||
@@ -324,7 +332,7 @@ async def test_get_entity_state(memory):
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_observation_fact_type_in_database(memory):
|
||||
async def test_observation_fact_type_in_database(memory, request_context):
|
||||
"""
|
||||
Test that observations are stored with correct fact_type in database.
|
||||
"""
|
||||
@@ -336,7 +344,8 @@ async def test_observation_fact_type_in_database(memory):
|
||||
bank_id=bank_id,
|
||||
content="Charlie is a DevOps engineer who manages the Kubernetes infrastructure.",
|
||||
context="work info",
|
||||
event_date=datetime(2024, 1, 15, tzinfo=timezone.utc)
|
||||
event_date=datetime(2024, 1, 15, tzinfo=timezone.utc),
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
await memory.wait_for_background_tasks()
|
||||
@@ -374,7 +383,7 @@ async def test_observation_fact_type_in_database(memory):
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_user_entity_prioritized_for_observations(memory):
|
||||
async def test_user_entity_prioritized_for_observations(memory, request_context):
|
||||
"""
|
||||
Test that the 'user' entity gets observations even when many other entities exist.
|
||||
|
||||
@@ -410,7 +419,8 @@ async def test_user_entity_prioritized_for_observations(memory):
|
||||
bank_id=bank_id,
|
||||
content=content,
|
||||
context="personal info",
|
||||
event_date=datetime(2024, 1, 15 + i, tzinfo=timezone.utc)
|
||||
event_date=datetime(2024, 1, 15 + i, tzinfo=timezone.utc),
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Observations are generated synchronously during retain
|
||||
@@ -466,7 +476,7 @@ async def test_user_entity_prioritized_for_observations(memory):
|
||||
f"User entity should have at least 5 facts, but has {user_fact_count}"
|
||||
|
||||
# Get observations for user entity
|
||||
observations = await memory.get_entity_observations(bank_id, user_entity_id, limit=10)
|
||||
observations = await memory.get_entity_observations(bank_id, user_entity_id, limit=10, request_context=request_context)
|
||||
|
||||
print(f"\n=== User Entity Observations ===")
|
||||
print(f"Total observations: {len(observations)}")
|
||||
|
||||
+157
-104
@@ -5,12 +5,13 @@ import pytest
|
||||
import logging
|
||||
from datetime import datetime, timezone, timedelta
|
||||
from hindsight_api.engine.memory_engine import Budget
|
||||
from hindsight_api import RequestContext
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retain_with_chunks(memory):
|
||||
async def test_retain_with_chunks(memory, request_context):
|
||||
"""
|
||||
Test that retain function:
|
||||
1. Stores facts with associated chunks
|
||||
@@ -41,7 +42,8 @@ async def test_retain_with_chunks(memory):
|
||||
content=long_content,
|
||||
context="team overview",
|
||||
event_date=datetime(2024, 1, 15, tzinfo=timezone.utc),
|
||||
document_id=document_id
|
||||
document_id=document_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
print(f"\n=== Retained {len(unit_ids)} facts ===")
|
||||
@@ -56,7 +58,8 @@ async def test_retain_with_chunks(memory):
|
||||
fact_type=["world"], # Search for world facts
|
||||
include_entities=False, # Disable entities for simpler test
|
||||
include_chunks=True, # Enable chunks
|
||||
max_chunk_tokens=8192
|
||||
max_chunk_tokens=8192,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
print(f"\n=== Recall Results (with chunks) ===")
|
||||
@@ -88,12 +91,12 @@ async def test_retain_with_chunks(memory):
|
||||
|
||||
finally:
|
||||
# Cleanup - delete the test bank
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
print(f"\n=== Cleaned up bank: {bank_id} ===")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_chunks_and_entities_follow_fact_order(memory):
|
||||
async def test_chunks_and_entities_follow_fact_order(memory, request_context):
|
||||
"""
|
||||
Test that chunks and entities in recall results follow the same order as facts.
|
||||
This is critical because token limits may truncate later items.
|
||||
@@ -130,7 +133,8 @@ async def test_chunks_and_entities_follow_fact_order(memory):
|
||||
content=item["content"],
|
||||
context=item["context"],
|
||||
event_date=datetime(2024, 1, 15, tzinfo=timezone.utc),
|
||||
document_id=item["document_id"]
|
||||
document_id=item["document_id"],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
print("\n=== Stored 3 separate documents ===")
|
||||
@@ -144,7 +148,8 @@ async def test_chunks_and_entities_follow_fact_order(memory):
|
||||
fact_type=["world"],
|
||||
include_entities=True,
|
||||
include_chunks=True,
|
||||
max_chunk_tokens=8192
|
||||
max_chunk_tokens=8192,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
print(f"\n=== Recall Results ===")
|
||||
@@ -214,12 +219,12 @@ async def test_chunks_and_entities_follow_fact_order(memory):
|
||||
|
||||
finally:
|
||||
# Cleanup
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
print(f"\n=== Cleaned up bank: {bank_id} ===")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_event_date_storage(memory):
|
||||
async def test_event_date_storage(memory, request_context):
|
||||
"""
|
||||
Test that event_date is correctly stored as occurred_start.
|
||||
Verifies that we can track when events actually happened vs when they were stored.
|
||||
@@ -235,7 +240,8 @@ async def test_event_date_storage(memory):
|
||||
bank_id=bank_id,
|
||||
content="Alice completed the Q2 product launch on June 15th, 2023.",
|
||||
context="project history",
|
||||
event_date=past_event_date
|
||||
event_date=past_event_date,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert len(unit_ids) > 0, "Should have created at least one memory unit"
|
||||
@@ -246,7 +252,8 @@ async def test_event_date_storage(memory):
|
||||
query="When did Alice complete the product launch?",
|
||||
budget=Budget.LOW,
|
||||
max_tokens=500,
|
||||
fact_type=["world"]
|
||||
fact_type=["world"],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert len(result.results) > 0, "Should recall the stored fact"
|
||||
@@ -268,11 +275,11 @@ async def test_event_date_storage(memory):
|
||||
print(f"\n✓ Event date correctly stored: {occurred_dt}")
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_temporal_ordering(memory):
|
||||
async def test_temporal_ordering(memory, request_context):
|
||||
"""
|
||||
Test that facts can be stored and retrieved with correct temporal ordering.
|
||||
Stores facts with different event_dates and verifies temporal relationships.
|
||||
@@ -305,7 +312,8 @@ async def test_temporal_ordering(memory):
|
||||
bank_id=bank_id,
|
||||
content=event["content"],
|
||||
context=event["context"],
|
||||
event_date=event["event_date"]
|
||||
event_date=event["event_date"],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
print("\n=== Stored 3 events with different temporal dates ===")
|
||||
@@ -316,7 +324,8 @@ async def test_temporal_ordering(memory):
|
||||
query="Tell me about Alice's career progression",
|
||||
budget=Budget.MID,
|
||||
max_tokens=1000,
|
||||
fact_type=["world"]
|
||||
fact_type=["world"],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert len(result.results) >= 3, f"Should recall all 3 events, got {len(result.results)}"
|
||||
@@ -345,11 +354,11 @@ async def test_temporal_ordering(memory):
|
||||
print(f"\n✓ Temporal ordering preserved: {min_date.date()} to {max_date.date()}")
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mentioned_at_vs_occurred(memory):
|
||||
async def test_mentioned_at_vs_occurred(memory, request_context):
|
||||
"""
|
||||
Test distinction between when fact occurred vs when it was mentioned.
|
||||
|
||||
@@ -369,7 +378,8 @@ async def test_mentioned_at_vs_occurred(memory):
|
||||
bank_id=bank_id,
|
||||
content="Alice graduated from MIT in March 2020.",
|
||||
context="education history",
|
||||
event_date=conversation_date # When this conversation happened
|
||||
event_date=conversation_date, # When this conversation happened
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert len(unit_ids) > 0, "Should create memory unit"
|
||||
@@ -380,7 +390,8 @@ async def test_mentioned_at_vs_occurred(memory):
|
||||
query="Where did Alice go to school?",
|
||||
budget=Budget.LOW,
|
||||
max_tokens=500,
|
||||
fact_type=["world"]
|
||||
fact_type=["world"],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert len(result.results) > 0, "Should recall the fact"
|
||||
@@ -415,11 +426,11 @@ async def test_mentioned_at_vs_occurred(memory):
|
||||
print(f"✓ Test passed: Historical conversation correctly ingested with event_date=2020")
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_occurred_dates_not_defaulted(memory):
|
||||
async def test_occurred_dates_not_defaulted(memory, request_context):
|
||||
"""
|
||||
Test that occurred_start and occurred_end are NOT defaulted to mentioned_at.
|
||||
|
||||
@@ -441,7 +452,8 @@ async def test_occurred_dates_not_defaulted(memory):
|
||||
bank_id=bank_id,
|
||||
content="Alice likes coffee. The weather is sunny today.",
|
||||
context="current observations",
|
||||
event_date=event_date
|
||||
event_date=event_date,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert len(unit_ids) > 0, "Should create memory unit"
|
||||
@@ -452,7 +464,8 @@ async def test_occurred_dates_not_defaulted(memory):
|
||||
query="What does Alice like?",
|
||||
budget=Budget.LOW,
|
||||
max_tokens=500,
|
||||
fact_type=["world", "opinion"]
|
||||
fact_type=["world", "opinion"],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert len(result.results) > 0, "Should recall the fact"
|
||||
@@ -504,11 +517,11 @@ async def test_occurred_dates_not_defaulted(memory):
|
||||
print(f"✓ Test passed: occurred dates are not incorrectly defaulted to mentioned_at")
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mentioned_at_from_context_string(memory):
|
||||
async def test_mentioned_at_from_context_string(memory, request_context):
|
||||
"""
|
||||
Test that mentioned_at is extracted from context string by LLM.
|
||||
|
||||
@@ -527,7 +540,8 @@ async def test_mentioned_at_from_context_string(memory):
|
||||
bank_id=bank_id,
|
||||
content="Alice mentioned she loves hiking in the mountains.",
|
||||
context=f"Session ABC123 - you are the assistant in this conversation - happened on {session_date.strftime('%Y-%m-%d %H:%M:%S')} UTC.",
|
||||
event_date=None # Not providing event_date - should default to now() if LLM doesn't extract
|
||||
event_date=None, # Not providing event_date - should default to now() if LLM doesn't extract
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert len(unit_ids) > 0, "Should create memory unit"
|
||||
@@ -538,7 +552,8 @@ async def test_mentioned_at_from_context_string(memory):
|
||||
query="What does Alice like?",
|
||||
budget=Budget.LOW,
|
||||
max_tokens=500,
|
||||
fact_type=["world"]
|
||||
fact_type=["world"],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert len(result.results) > 0, "Should recall the fact"
|
||||
@@ -574,7 +589,7 @@ async def test_mentioned_at_from_context_string(memory):
|
||||
print(f"✓ mentioned_at is always set (never None)")
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
# ============================================================
|
||||
@@ -582,7 +597,7 @@ async def test_mentioned_at_from_context_string(memory):
|
||||
# ============================================================
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_context_preservation(memory):
|
||||
async def test_context_preservation(memory, request_context):
|
||||
"""
|
||||
Test that context is preserved and retrievable.
|
||||
Context helps understand why/how memory was formed.
|
||||
@@ -597,7 +612,8 @@ async def test_context_preservation(memory):
|
||||
bank_id=bank_id,
|
||||
content="The team decided to prioritize mobile development for next quarter.",
|
||||
context=specific_context,
|
||||
event_date=datetime(2024, 1, 15, tzinfo=timezone.utc)
|
||||
event_date=datetime(2024, 1, 15, tzinfo=timezone.utc),
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert len(unit_ids) > 0, "Should create at least one memory unit"
|
||||
@@ -608,7 +624,8 @@ async def test_context_preservation(memory):
|
||||
query="What did the team decide?",
|
||||
budget=Budget.LOW,
|
||||
max_tokens=500,
|
||||
fact_type=["world"]
|
||||
fact_type=["world"],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert len(result.results) > 0, "Should recall the stored fact"
|
||||
@@ -620,11 +637,11 @@ async def test_context_preservation(memory):
|
||||
print(f" Retrieved {len(result.results)} facts")
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_context_with_batch(memory):
|
||||
async def test_context_with_batch(memory, request_context):
|
||||
"""
|
||||
Test that each item in a batch can have different contexts.
|
||||
"""
|
||||
@@ -650,7 +667,8 @@ async def test_context_with_batch(memory):
|
||||
"context": "incident response",
|
||||
"event_date": datetime(2024, 1, 12, tzinfo=timezone.utc)
|
||||
}
|
||||
]
|
||||
],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Should have created facts from all items
|
||||
@@ -661,7 +679,7 @@ async def test_context_with_batch(memory):
|
||||
print(f" Created {total_units} total memory units")
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
# ============================================================
|
||||
@@ -669,7 +687,7 @@ async def test_context_with_batch(memory):
|
||||
# ============================================================
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_metadata_storage_and_retrieval(memory):
|
||||
async def test_metadata_storage_and_retrieval(memory, request_context):
|
||||
"""
|
||||
Test that user-defined metadata is preserved.
|
||||
Metadata allows arbitrary key-value data to be stored with facts.
|
||||
@@ -692,7 +710,8 @@ async def test_metadata_storage_and_retrieval(memory):
|
||||
bank_id=bank_id,
|
||||
content="The product launch is scheduled for March 1st.",
|
||||
context="planning meeting",
|
||||
event_date=datetime(2024, 1, 15, tzinfo=timezone.utc)
|
||||
event_date=datetime(2024, 1, 15, tzinfo=timezone.utc),
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert len(unit_ids) > 0, "Should create memory units"
|
||||
@@ -703,7 +722,8 @@ async def test_metadata_storage_and_retrieval(memory):
|
||||
query="When is the product launch?",
|
||||
budget=Budget.LOW,
|
||||
max_tokens=500,
|
||||
fact_type=["world"]
|
||||
fact_type=["world"],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert len(result.results) > 0, "Should recall stored facts"
|
||||
@@ -712,7 +732,7 @@ async def test_metadata_storage_and_retrieval(memory):
|
||||
print(f" (Note: Metadata support depends on API implementation)")
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
# ============================================================
|
||||
@@ -720,7 +740,7 @@ async def test_metadata_storage_and_retrieval(memory):
|
||||
# ============================================================
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_empty_batch(memory):
|
||||
async def test_empty_batch(memory, request_context):
|
||||
"""
|
||||
Test that empty batch is handled gracefully without errors.
|
||||
"""
|
||||
@@ -730,7 +750,8 @@ async def test_empty_batch(memory):
|
||||
# Attempt to store empty batch
|
||||
unit_ids = await memory.retain_batch_async(
|
||||
bank_id=bank_id,
|
||||
contents=[]
|
||||
contents=[],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Should return empty list or handle gracefully
|
||||
@@ -741,11 +762,11 @@ async def test_empty_batch(memory):
|
||||
|
||||
finally:
|
||||
# Clean up (though nothing should be stored)
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_single_item_batch(memory):
|
||||
async def test_single_item_batch(memory, request_context):
|
||||
"""
|
||||
Test that batch with one item works correctly.
|
||||
"""
|
||||
@@ -761,7 +782,8 @@ async def test_single_item_batch(memory):
|
||||
"context": "deployment log",
|
||||
"event_date": datetime(2024, 1, 15, tzinfo=timezone.utc)
|
||||
}
|
||||
]
|
||||
],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert len(unit_ids) == 1, "Should return one list of unit IDs"
|
||||
@@ -770,11 +792,11 @@ async def test_single_item_batch(memory):
|
||||
print(f"✓ Single-item batch created {len(unit_ids[0])} units")
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mixed_content_batch(memory):
|
||||
async def test_mixed_content_batch(memory, request_context):
|
||||
"""
|
||||
Test batch with varying content sizes (short and long).
|
||||
"""
|
||||
@@ -798,7 +820,8 @@ async def test_mixed_content_batch(memory):
|
||||
{"content": short_content, "context": "onboarding"},
|
||||
{"content": long_content, "context": "performance review"},
|
||||
{"content": "Charlie is on vacation this week.", "context": "team status"}
|
||||
]
|
||||
],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# All items should be processed
|
||||
@@ -813,11 +836,11 @@ async def test_mixed_content_batch(memory):
|
||||
print(f" Long content: {long_units} units")
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_batch_with_missing_optional_fields(memory):
|
||||
async def test_batch_with_missing_optional_fields(memory, request_context):
|
||||
"""
|
||||
Test that batch handles items with missing optional fields.
|
||||
"""
|
||||
@@ -842,7 +865,8 @@ async def test_batch_with_missing_optional_fields(memory):
|
||||
"context": "code review",
|
||||
# No event_date
|
||||
}
|
||||
]
|
||||
],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# All items should be processed successfully
|
||||
@@ -852,7 +876,7 @@ async def test_batch_with_missing_optional_fields(memory):
|
||||
print(f"✓ Batch with mixed optional fields created {total_units} total units")
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
# ============================================================
|
||||
@@ -860,7 +884,7 @@ async def test_batch_with_missing_optional_fields(memory):
|
||||
# ============================================================
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_single_batch_multiple_documents(memory):
|
||||
async def test_single_batch_multiple_documents(memory, request_context):
|
||||
"""
|
||||
Test storing multiple distinct documents in a single batch call.
|
||||
Each should be tracked separately.
|
||||
@@ -876,21 +900,24 @@ async def test_single_batch_multiple_documents(memory):
|
||||
bank_id=bank_id,
|
||||
content="Alice's resume: 10 years Python experience, worked at Google.",
|
||||
context="resume review",
|
||||
document_id="resume_alice"
|
||||
document_id="resume_alice",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
doc2_units = await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="Bob's resume: 5 years JavaScript experience, worked at Meta.",
|
||||
context="resume review",
|
||||
document_id="resume_bob"
|
||||
document_id="resume_bob",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
doc3_units = await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="Charlie's resume: 8 years Go experience, worked at Amazon.",
|
||||
context="resume review",
|
||||
document_id="resume_charlie"
|
||||
document_id="resume_charlie",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# All documents should be stored
|
||||
@@ -907,17 +934,18 @@ async def test_single_batch_multiple_documents(memory):
|
||||
query="Who worked at Google?",
|
||||
budget=Budget.MID,
|
||||
max_tokens=1000,
|
||||
fact_type=["world"]
|
||||
fact_type=["world"],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert len(result.results) > 0, "Should find facts about Alice"
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_document_upsert_behavior(memory):
|
||||
async def test_document_upsert_behavior(memory, request_context):
|
||||
"""
|
||||
Test that upserting a document replaces the old content.
|
||||
"""
|
||||
@@ -930,7 +958,8 @@ async def test_document_upsert_behavior(memory):
|
||||
bank_id=bank_id,
|
||||
content="Project is in planning phase. Alice is the lead.",
|
||||
context="status update v1",
|
||||
document_id=document_id
|
||||
document_id=document_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert len(v1_units) > 0, "Should create units for v1"
|
||||
@@ -940,7 +969,8 @@ async def test_document_upsert_behavior(memory):
|
||||
bank_id=bank_id,
|
||||
content="Project is in development phase. Bob has joined as co-lead.",
|
||||
context="status update v2",
|
||||
document_id=document_id
|
||||
document_id=document_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert len(v2_units) > 0, "Should create units for v2"
|
||||
@@ -951,7 +981,8 @@ async def test_document_upsert_behavior(memory):
|
||||
query="What is the project status?",
|
||||
budget=Budget.MID,
|
||||
max_tokens=1000,
|
||||
fact_type=["world"]
|
||||
fact_type=["world"],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert len(result.results) > 0, "Should recall facts"
|
||||
@@ -959,7 +990,7 @@ async def test_document_upsert_behavior(memory):
|
||||
print(f"✓ Document upsert created v1: {len(v1_units)} units, v2: {len(v2_units)} units")
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
# ============================================================
|
||||
@@ -967,7 +998,7 @@ async def test_document_upsert_behavior(memory):
|
||||
# ============================================================
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_chunk_fact_mapping(memory):
|
||||
async def test_chunk_fact_mapping(memory, request_context):
|
||||
"""
|
||||
Test that facts correctly reference their source chunks via chunk_id.
|
||||
"""
|
||||
@@ -990,7 +1021,8 @@ async def test_chunk_fact_mapping(memory):
|
||||
bank_id=bank_id,
|
||||
content=content,
|
||||
context="technical documentation",
|
||||
document_id=document_id
|
||||
document_id=document_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert len(unit_ids) > 0, "Should create memory units"
|
||||
@@ -1003,7 +1035,8 @@ async def test_chunk_fact_mapping(memory):
|
||||
max_tokens=1000,
|
||||
fact_type=["world"],
|
||||
include_chunks=True,
|
||||
max_chunk_tokens=8192
|
||||
max_chunk_tokens=8192,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert len(result.results) > 0, "Should recall facts"
|
||||
@@ -1026,11 +1059,11 @@ async def test_chunk_fact_mapping(memory):
|
||||
print(f" Returned {len(result.chunks)} chunks matching fact references")
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_chunk_ordering_preservation(memory):
|
||||
async def test_chunk_ordering_preservation(memory, request_context):
|
||||
"""
|
||||
Test that chunk_index reflects the correct order within a document.
|
||||
"""
|
||||
@@ -1070,7 +1103,8 @@ async def test_chunk_ordering_preservation(memory):
|
||||
bank_id=bank_id,
|
||||
content=content,
|
||||
context="multi-section document",
|
||||
document_id=document_id
|
||||
document_id=document_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert len(unit_ids) > 0, "Should create units"
|
||||
@@ -1083,7 +1117,8 @@ async def test_chunk_ordering_preservation(memory):
|
||||
max_tokens=2000,
|
||||
fact_type=["world"],
|
||||
include_chunks=True,
|
||||
max_chunk_tokens=8192
|
||||
max_chunk_tokens=8192,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
if result.chunks:
|
||||
@@ -1103,11 +1138,11 @@ async def test_chunk_ordering_preservation(memory):
|
||||
print("✓ Content stored (may have created single chunk or no chunks returned)")
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_chunks_truncation_behavior(memory):
|
||||
async def test_chunks_truncation_behavior(memory, request_context):
|
||||
"""
|
||||
Test that when chunks exceed max_chunk_tokens, truncation is indicated.
|
||||
"""
|
||||
@@ -1165,7 +1200,8 @@ async def test_chunks_truncation_behavior(memory):
|
||||
bank_id=bank_id,
|
||||
content=large_content,
|
||||
context="large document test",
|
||||
document_id=document_id
|
||||
document_id=document_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert len(unit_ids) > 0, "Should create units"
|
||||
@@ -1178,7 +1214,8 @@ async def test_chunks_truncation_behavior(memory):
|
||||
max_tokens=1000,
|
||||
fact_type=["world"],
|
||||
include_chunks=True,
|
||||
max_chunk_tokens=500 # Small limit to test truncation
|
||||
max_chunk_tokens=500, # Small limit to test truncation
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
if result.chunks:
|
||||
@@ -1198,7 +1235,7 @@ async def test_chunks_truncation_behavior(memory):
|
||||
print("✓ No chunks returned (may be under token limit)")
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
# ============================================================
|
||||
@@ -1206,7 +1243,7 @@ async def test_chunks_truncation_behavior(memory):
|
||||
# ============================================================
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_temporal_links_creation(memory):
|
||||
async def test_temporal_links_creation(memory, request_context):
|
||||
"""
|
||||
Test that temporal links are created between facts with nearby event dates.
|
||||
|
||||
@@ -1223,7 +1260,8 @@ async def test_temporal_links_creation(memory):
|
||||
bank_id=bank_id,
|
||||
content="Alice started working on the authentication module.",
|
||||
context="daily standup",
|
||||
event_date=base_date
|
||||
event_date=base_date,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Fact 2 at 2:00 PM same day (4 hours later)
|
||||
@@ -1231,7 +1269,8 @@ async def test_temporal_links_creation(memory):
|
||||
bank_id=bank_id,
|
||||
content="Bob reviewed the API design document.",
|
||||
context="daily standup",
|
||||
event_date=base_date.replace(hour=14)
|
||||
event_date=base_date.replace(hour=14),
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Fact 3 at 9:00 AM next day (23 hours later)
|
||||
@@ -1239,7 +1278,8 @@ async def test_temporal_links_creation(memory):
|
||||
bank_id=bank_id,
|
||||
content="Charlie deployed the new database schema.",
|
||||
context="daily standup",
|
||||
event_date=base_date.replace(day=16, hour=9)
|
||||
event_date=base_date.replace(day=16, hour=9),
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert len(unit_ids_1) > 0 and len(unit_ids_2) > 0 and len(unit_ids_3) > 0
|
||||
@@ -1278,11 +1318,11 @@ async def test_temporal_links_creation(memory):
|
||||
logger.info("Temporal links created successfully with proper weights")
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_semantic_links_creation(memory):
|
||||
async def test_semantic_links_creation(memory, request_context):
|
||||
"""
|
||||
Test that semantic links are created between facts with similar content.
|
||||
|
||||
@@ -1295,21 +1335,24 @@ async def test_semantic_links_creation(memory):
|
||||
unit_ids_1 = await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="Alice is an expert in Python programming and has built many web applications.",
|
||||
context="team skills"
|
||||
context="team skills",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Similar content - should create semantic link
|
||||
unit_ids_2 = await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="Bob is proficient in Python development and specializes in building APIs.",
|
||||
context="team skills"
|
||||
context="team skills",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Different content - less likely to create strong semantic link
|
||||
unit_ids_3 = await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="The quarterly sales meeting is scheduled for next Tuesday at 3 PM.",
|
||||
context="calendar events"
|
||||
context="calendar events",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert len(unit_ids_1) > 0 and len(unit_ids_2) > 0 and len(unit_ids_3) > 0
|
||||
@@ -1349,11 +1392,11 @@ async def test_semantic_links_creation(memory):
|
||||
logger.info("Semantic links created successfully between similar content")
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_entity_links_creation(memory):
|
||||
async def test_entity_links_creation(memory, request_context):
|
||||
"""
|
||||
Test that entity links are created between facts that mention the same entities.
|
||||
|
||||
@@ -1367,28 +1410,32 @@ async def test_entity_links_creation(memory):
|
||||
unit_ids_1 = await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="Alice joined Google as a software engineer in 2020.",
|
||||
context="career history"
|
||||
context="career history",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Mentions same entity (Alice) - should create entity link
|
||||
unit_ids_2 = await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="Alice led the development of the new authentication system.",
|
||||
context="project updates"
|
||||
context="project updates",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Mentions same entity (Google) - should create entity link
|
||||
unit_ids_3 = await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="Google announced new cloud services at their annual conference.",
|
||||
context="tech news"
|
||||
context="tech news",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Different entities - no entity link expected
|
||||
unit_ids_4 = await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="Bob works at Meta on machine learning infrastructure.",
|
||||
context="career history"
|
||||
context="career history",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert len(unit_ids_1) > 0 and len(unit_ids_2) > 0 and len(unit_ids_3) > 0 and len(unit_ids_4) > 0
|
||||
@@ -1445,11 +1492,11 @@ async def test_entity_links_creation(memory):
|
||||
logger.info("Entity links are properly bidirectional")
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_causal_links_creation(memory):
|
||||
async def test_causal_links_creation(memory, request_context):
|
||||
"""
|
||||
Test that causal links are created between facts with causal relationships.
|
||||
|
||||
@@ -1471,7 +1518,8 @@ async def test_causal_links_creation(memory):
|
||||
unit_ids = await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content=content,
|
||||
context="project timeline"
|
||||
context="project timeline",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert len(unit_ids) > 0, "Should have created facts"
|
||||
@@ -1517,11 +1565,11 @@ async def test_causal_links_creation(memory):
|
||||
logger.info("Test completed (causal link extraction is LLM-dependent)")
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_all_link_types_together(memory):
|
||||
async def test_all_link_types_together(memory, request_context):
|
||||
"""
|
||||
Integration test: Verify all link types can be created in a single retain operation.
|
||||
|
||||
@@ -1539,7 +1587,8 @@ async def test_all_link_types_together(memory):
|
||||
bank_id=bank_id,
|
||||
content="Alice completed the Python backend service for the authentication system.",
|
||||
context="sprint review",
|
||||
event_date=base_date
|
||||
event_date=base_date,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Fact 2: Related to Alice, similar topic (Python), close in time
|
||||
@@ -1547,7 +1596,8 @@ async def test_all_link_types_together(memory):
|
||||
bank_id=bank_id,
|
||||
content="Alice optimized the Python code and improved the authentication performance by 40%.",
|
||||
context="sprint review",
|
||||
event_date=base_date.replace(hour=14) # Same day, 4 hours later
|
||||
event_date=base_date.replace(hour=14), # Same day, 4 hours later
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Fact 3: Related to Alice, different topic but same entity
|
||||
@@ -1555,7 +1605,8 @@ async def test_all_link_types_together(memory):
|
||||
bank_id=bank_id,
|
||||
content="Alice presented the security architecture at the team meeting.",
|
||||
context="team meeting",
|
||||
event_date=base_date.replace(day=16) # Next day
|
||||
event_date=base_date.replace(day=16), # Next day
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert len(unit_ids_1) > 0 and len(unit_ids_2) > 0 and len(unit_ids_3) > 0
|
||||
@@ -1594,11 +1645,11 @@ async def test_all_link_types_together(memory):
|
||||
logger.info("All major link types (temporal, semantic, entity) are working correctly")
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_semantic_links_within_same_batch(memory):
|
||||
async def test_semantic_links_within_same_batch(memory, request_context):
|
||||
"""
|
||||
Test that semantic links are created between facts retained in the SAME batch.
|
||||
|
||||
@@ -1617,7 +1668,8 @@ async def test_semantic_links_within_same_batch(memory):
|
||||
|
||||
result = await memory.retain_batch_async(
|
||||
bank_id=bank_id,
|
||||
contents=contents
|
||||
contents=contents,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Flatten the list of lists
|
||||
@@ -1652,11 +1704,11 @@ async def test_semantic_links_within_same_batch(memory):
|
||||
logger.info(f" Semantic link: {str(link['from_unit_id'])[:8]}... -> {str(link['to_unit_id'])[:8]}... (weight: {link['weight']:.3f})")
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_temporal_links_within_same_batch(memory):
|
||||
async def test_temporal_links_within_same_batch(memory, request_context):
|
||||
"""
|
||||
Test that temporal links are created between facts retained in the SAME batch.
|
||||
|
||||
@@ -1689,7 +1741,8 @@ async def test_temporal_links_within_same_batch(memory):
|
||||
|
||||
result = await memory.retain_batch_async(
|
||||
bank_id=bank_id,
|
||||
contents=contents
|
||||
contents=contents,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Flatten the list of lists
|
||||
@@ -1724,4 +1777,4 @@ async def test_temporal_links_within_same_batch(memory):
|
||||
logger.info(f" Temporal link: {str(link['from_unit_id'])[:8]}... -> {str(link['to_unit_id'])[:8]}... (weight: {link['weight']:.3f})")
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@@ -0,0 +1,406 @@
|
||||
"""
|
||||
Tests for multi-tenant schema isolation.
|
||||
|
||||
Verifies that concurrent retain operations from different tenants
|
||||
are properly isolated in their respective PostgreSQL schemas.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import uuid
|
||||
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
|
||||
from hindsight_api.extensions import RequestContext, TenantContext, TenantExtension
|
||||
from hindsight_api.engine.memory_engine import _current_schema, fq_table
|
||||
from hindsight_api.migrations import run_migrations
|
||||
|
||||
|
||||
class MultiSchemaTestTenantExtension(TenantExtension):
|
||||
"""
|
||||
Test tenant extension that maps API keys to schema names.
|
||||
|
||||
API keys are in format: "key-{schema_name}"
|
||||
Provisions schemas on first access using run_migrations(schema=name).
|
||||
"""
|
||||
|
||||
def __init__(self, config: dict):
|
||||
super().__init__(config)
|
||||
self.db_url = config.get("db_url")
|
||||
# Pre-configured valid schemas for test
|
||||
self.valid_schemas = config.get("valid_schemas", set())
|
||||
# Track provisioned schemas
|
||||
self._provisioned: set[str] = set()
|
||||
|
||||
async def authenticate(self, context: RequestContext) -> TenantContext:
|
||||
if not context.api_key:
|
||||
from hindsight_api.extensions import AuthenticationError
|
||||
|
||||
raise AuthenticationError("API key required")
|
||||
|
||||
# Parse schema from API key (format: "key-{schema}")
|
||||
if context.api_key.startswith("key-"):
|
||||
schema = context.api_key[4:] # Remove "key-" prefix
|
||||
if schema in self.valid_schemas:
|
||||
# Provision schema on first access
|
||||
if schema not in self._provisioned and self.db_url:
|
||||
run_migrations(self.db_url, schema=schema)
|
||||
self._provisioned.add(schema)
|
||||
return TenantContext(schema_name=schema)
|
||||
|
||||
from hindsight_api.extensions import AuthenticationError
|
||||
|
||||
raise AuthenticationError(f"Unknown API key: {context.api_key}")
|
||||
|
||||
|
||||
async def drop_schema(conn, schema_name: str) -> None:
|
||||
"""Drop a schema and all its contents."""
|
||||
await conn.execute(f'DROP SCHEMA IF EXISTS "{schema_name}" CASCADE')
|
||||
|
||||
|
||||
async def count_memories_in_schema(conn, schema_name: str, bank_id: str) -> int:
|
||||
"""Count memory units in a specific schema for a bank."""
|
||||
result = await conn.fetchval(
|
||||
f'SELECT COUNT(*) FROM "{schema_name}".memory_units WHERE bank_id = $1',
|
||||
bank_id,
|
||||
)
|
||||
return result or 0
|
||||
|
||||
|
||||
async def get_memory_texts_in_schema(conn, schema_name: str, bank_id: str) -> list[str]:
|
||||
"""Get all memory texts in a specific schema for a bank."""
|
||||
rows = await conn.fetch(
|
||||
f'SELECT text FROM "{schema_name}".memory_units WHERE bank_id = $1 ORDER BY text',
|
||||
bank_id,
|
||||
)
|
||||
return [row["text"] for row in rows]
|
||||
|
||||
|
||||
class TestSchemaIsolation:
|
||||
"""Tests for multi-tenant schema isolation."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_concurrent_inserts_isolated_by_schema(self, memory, pg0_db_url):
|
||||
"""
|
||||
Multiple concurrent database operations from different tenants
|
||||
should store data in their respective schemas without cross-contamination.
|
||||
|
||||
Uses run_migrations(schema=x) to provision schemas like a real extension.
|
||||
"""
|
||||
import asyncpg
|
||||
|
||||
# Test schemas
|
||||
schemas = ["tenant_alpha", "tenant_beta", "tenant_gamma"]
|
||||
bank_id = f"test-isolation-{uuid.uuid4().hex[:8]}"
|
||||
|
||||
# Clean up any existing schemas
|
||||
conn = await asyncpg.connect(pg0_db_url)
|
||||
try:
|
||||
for schema in schemas:
|
||||
await drop_schema(conn, schema)
|
||||
finally:
|
||||
await conn.close()
|
||||
|
||||
# Configure tenant extension that provisions schemas via run_migrations
|
||||
tenant_ext = MultiSchemaTestTenantExtension({
|
||||
"db_url": pg0_db_url,
|
||||
"valid_schemas": set(schemas),
|
||||
})
|
||||
memory._tenant_extension = tenant_ext
|
||||
|
||||
# Define concurrent insert tasks for each tenant
|
||||
async def insert_for_tenant(schema_name: str, content_prefix: str):
|
||||
"""Insert memories for a specific tenant using schema context."""
|
||||
# Authenticate to set the schema context
|
||||
tenant_request = RequestContext(api_key=f"key-{schema_name}")
|
||||
await memory._authenticate_tenant(tenant_request)
|
||||
|
||||
# Now fq_table will use the correct schema
|
||||
pool = await memory._get_pool()
|
||||
from hindsight_api.engine.db_utils import acquire_with_retry
|
||||
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
# Insert 3 memories for this tenant
|
||||
for i in range(3):
|
||||
await conn.execute(
|
||||
f"""
|
||||
INSERT INTO {fq_table('memory_units')} (bank_id, text, event_date, fact_type)
|
||||
VALUES ($1, $2, now(), 'world')
|
||||
""",
|
||||
bank_id,
|
||||
f"MARKER_{content_prefix}_DOC{i}: Memory for {schema_name}",
|
||||
)
|
||||
|
||||
# Run concurrent inserts for all tenants
|
||||
await asyncio.gather(
|
||||
insert_for_tenant("tenant_alpha", "ALPHA"),
|
||||
insert_for_tenant("tenant_beta", "BETA"),
|
||||
insert_for_tenant("tenant_gamma", "GAMMA"),
|
||||
)
|
||||
|
||||
# Verify isolation - each schema should only have its own data
|
||||
conn = await asyncpg.connect(pg0_db_url)
|
||||
try:
|
||||
for schema in schemas:
|
||||
texts = await get_memory_texts_in_schema(conn, schema, bank_id)
|
||||
prefix = schema.replace("tenant_", "").upper()
|
||||
|
||||
# Should have exactly 3 memories
|
||||
assert len(texts) == 3, f"Schema {schema} should have 3 memories, got {len(texts)}"
|
||||
|
||||
# All texts should contain the schema's marker
|
||||
for text in texts:
|
||||
assert f"MARKER_{prefix}" in text, (
|
||||
f"Memory in {schema} missing its marker: {text}"
|
||||
)
|
||||
|
||||
# Should NOT contain other tenants' markers
|
||||
other_prefixes = ["ALPHA", "BETA", "GAMMA"]
|
||||
other_prefixes.remove(prefix)
|
||||
for other in other_prefixes:
|
||||
for text in texts:
|
||||
assert f"MARKER_{other}" not in text, (
|
||||
f"Cross-contamination! Schema {schema} has {other}'s marker: {text}"
|
||||
)
|
||||
|
||||
finally:
|
||||
# Cleanup
|
||||
for schema in schemas:
|
||||
await drop_schema(conn, schema)
|
||||
await conn.close()
|
||||
|
||||
# Reset tenant extension
|
||||
memory._tenant_extension = None
|
||||
_current_schema.set("public")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_schema_context_isolation_in_concurrent_tasks(self, pg0_db_url):
|
||||
"""
|
||||
Verify that _current_schema contextvar is properly isolated
|
||||
between concurrent async tasks.
|
||||
"""
|
||||
results = {}
|
||||
errors = []
|
||||
|
||||
async def check_schema_context(schema_name: str, delay: float):
|
||||
"""Set schema context, wait, then verify it's still correct."""
|
||||
try:
|
||||
# Set the schema
|
||||
_current_schema.set(schema_name)
|
||||
|
||||
# Small delay to allow interleaving
|
||||
await asyncio.sleep(delay)
|
||||
|
||||
# Verify schema is still correct
|
||||
current = _current_schema.get()
|
||||
if current != schema_name:
|
||||
errors.append(f"Expected {schema_name}, got {current}")
|
||||
|
||||
# Verify fq_table uses correct schema
|
||||
table = fq_table("memory_units")
|
||||
expected = f"{schema_name}.memory_units"
|
||||
if table != expected:
|
||||
errors.append(f"Expected {expected}, got {table}")
|
||||
|
||||
results[schema_name] = current
|
||||
|
||||
except Exception as e:
|
||||
errors.append(f"Error in {schema_name}: {e}")
|
||||
|
||||
# Run many concurrent tasks with different schemas
|
||||
tasks = []
|
||||
for i in range(10):
|
||||
for schema in ["schema_a", "schema_b", "schema_c"]:
|
||||
# Vary delays to create interleaving
|
||||
delay = 0.01 * (i % 3)
|
||||
tasks.append(check_schema_context(f"{schema}_{i}", delay))
|
||||
|
||||
await asyncio.gather(*tasks)
|
||||
|
||||
# No errors should have occurred
|
||||
assert not errors, f"Schema context isolation errors: {errors}"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_memories_respects_schema(self, memory, pg0_db_url):
|
||||
"""
|
||||
list_memory_units should only return memories from the current schema.
|
||||
|
||||
Uses run_migrations(schema=x) to provision schemas.
|
||||
"""
|
||||
import asyncpg
|
||||
|
||||
schemas = ["tenant_list_a", "tenant_list_b"]
|
||||
bank_id = f"test-list-{uuid.uuid4().hex[:8]}"
|
||||
|
||||
# Clean up any existing schemas and provision via migrations
|
||||
conn = await asyncpg.connect(pg0_db_url)
|
||||
try:
|
||||
for schema in schemas:
|
||||
await drop_schema(conn, schema)
|
||||
finally:
|
||||
await conn.close()
|
||||
|
||||
# Provision schemas using run_migrations
|
||||
for schema in schemas:
|
||||
run_migrations(pg0_db_url, schema=schema)
|
||||
|
||||
# Insert test data directly into each schema
|
||||
conn = await asyncpg.connect(pg0_db_url)
|
||||
try:
|
||||
for schema in schemas:
|
||||
await conn.execute(
|
||||
f"""
|
||||
INSERT INTO "{schema}".memory_units (bank_id, text, event_date, fact_type)
|
||||
VALUES ($1, $2, now(), 'world')
|
||||
""",
|
||||
bank_id,
|
||||
f"Direct insert for {schema}",
|
||||
)
|
||||
finally:
|
||||
await conn.close()
|
||||
|
||||
# Configure tenant extension
|
||||
tenant_ext = MultiSchemaTestTenantExtension({
|
||||
"db_url": pg0_db_url,
|
||||
"valid_schemas": set(schemas),
|
||||
})
|
||||
memory._tenant_extension = tenant_ext
|
||||
|
||||
try:
|
||||
# Query as tenant_list_a - should only see tenant_list_a's data
|
||||
tenant_a_request = RequestContext(api_key="key-tenant_list_a")
|
||||
await memory._authenticate_tenant(tenant_a_request)
|
||||
|
||||
result_a = await memory.list_memory_units(bank_id=bank_id, request_context=tenant_a_request)
|
||||
texts_a = [item["text"] for item in result_a.get("items", [])]
|
||||
|
||||
assert len(texts_a) == 1, f"Expected 1 memory for tenant_list_a, got {len(texts_a)}"
|
||||
assert "tenant_list_a" in texts_a[0], f"Wrong content: {texts_a[0]}"
|
||||
|
||||
# Query as tenant_list_b - should only see tenant_list_b's data
|
||||
tenant_b_request = RequestContext(api_key="key-tenant_list_b")
|
||||
await memory._authenticate_tenant(tenant_b_request)
|
||||
|
||||
result_b = await memory.list_memory_units(bank_id=bank_id, request_context=tenant_b_request)
|
||||
texts_b = [item["text"] for item in result_b.get("items", [])]
|
||||
|
||||
assert len(texts_b) == 1, f"Expected 1 memory for tenant_list_b, got {len(texts_b)}"
|
||||
assert "tenant_list_b" in texts_b[0], f"Wrong content: {texts_b[0]}"
|
||||
|
||||
finally:
|
||||
# Cleanup
|
||||
conn = await asyncpg.connect(pg0_db_url)
|
||||
try:
|
||||
for schema in schemas:
|
||||
await drop_schema(conn, schema)
|
||||
finally:
|
||||
await conn.close()
|
||||
|
||||
memory._tenant_extension = None
|
||||
_current_schema.set("public")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_high_concurrency_schema_isolation(self, memory, pg0_db_url):
|
||||
"""
|
||||
Stress test: Many concurrent operations across multiple schemas
|
||||
should maintain perfect isolation.
|
||||
|
||||
Uses run_migrations(schema=x) to provision schemas like a real extension.
|
||||
"""
|
||||
import asyncpg
|
||||
|
||||
# Create more schemas for stress test
|
||||
num_schemas = 5
|
||||
ops_per_schema = 10
|
||||
schemas = [f"stress_tenant_{i}" for i in range(num_schemas)]
|
||||
bank_id = f"test-stress-{uuid.uuid4().hex[:8]}"
|
||||
|
||||
# Clean up any existing schemas first
|
||||
conn = await asyncpg.connect(pg0_db_url)
|
||||
try:
|
||||
for schema in schemas:
|
||||
await drop_schema(conn, schema)
|
||||
finally:
|
||||
await conn.close()
|
||||
|
||||
# Provision schemas using run_migrations
|
||||
for schema in schemas:
|
||||
run_migrations(pg0_db_url, schema=schema)
|
||||
|
||||
# Configure tenant extension (schemas already provisioned)
|
||||
tenant_ext = MultiSchemaTestTenantExtension({
|
||||
"db_url": pg0_db_url,
|
||||
"valid_schemas": set(schemas),
|
||||
})
|
||||
# Mark schemas as already provisioned so extension doesn't re-run migrations
|
||||
tenant_ext._provisioned = set(schemas)
|
||||
memory._tenant_extension = tenant_ext
|
||||
|
||||
errors = []
|
||||
|
||||
async def insert_one(schema: str, item_id: int):
|
||||
"""Single insert operation for tracking."""
|
||||
try:
|
||||
# Authenticate to set the schema context
|
||||
tenant_request = RequestContext(api_key=f"key-{schema}")
|
||||
await memory._authenticate_tenant(tenant_request)
|
||||
|
||||
# Insert using fq_table
|
||||
pool = await memory._get_pool()
|
||||
from hindsight_api.engine.db_utils import acquire_with_retry
|
||||
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
await conn.execute(
|
||||
f"""
|
||||
INSERT INTO {fq_table('memory_units')} (bank_id, text, event_date, fact_type)
|
||||
VALUES ($1, $2, now(), 'world')
|
||||
""",
|
||||
bank_id,
|
||||
f"STRESS_MARKER_{schema}_ITEM{item_id}: Memory for {schema}",
|
||||
)
|
||||
except Exception as e:
|
||||
errors.append(f"Insert error for {schema}: {e}")
|
||||
|
||||
# Run many concurrent operations
|
||||
tasks = []
|
||||
for i in range(ops_per_schema):
|
||||
for schema in schemas:
|
||||
tasks.append(insert_one(schema, i))
|
||||
|
||||
await asyncio.gather(*tasks)
|
||||
|
||||
# Check for errors during insert
|
||||
assert not errors, f"Errors during insert: {errors}"
|
||||
|
||||
# Verify no cross-contamination
|
||||
conn = await asyncpg.connect(pg0_db_url)
|
||||
try:
|
||||
for schema in schemas:
|
||||
texts = await get_memory_texts_in_schema(conn, schema, bank_id)
|
||||
|
||||
# Should have exactly ops_per_schema memories
|
||||
assert len(texts) == ops_per_schema, (
|
||||
f"Schema {schema} should have {ops_per_schema} memories, got {len(texts)}"
|
||||
)
|
||||
|
||||
# All memories should reference this schema only
|
||||
for text in texts:
|
||||
# Check it contains our schema marker
|
||||
assert f"STRESS_MARKER_{schema}" in text, (
|
||||
f"Memory in {schema} doesn't contain schema marker: {text}"
|
||||
)
|
||||
|
||||
# Check it doesn't contain other schema markers
|
||||
for other_schema in schemas:
|
||||
if other_schema != schema:
|
||||
assert f"STRESS_MARKER_{other_schema}" not in text, (
|
||||
f"Cross-contamination! {schema} has {other_schema}'s data: {text}"
|
||||
)
|
||||
finally:
|
||||
# Cleanup
|
||||
for schema in schemas:
|
||||
await drop_schema(conn, schema)
|
||||
await conn.close()
|
||||
|
||||
memory._tenant_extension = None
|
||||
_current_schema.set("public")
|
||||
@@ -3,12 +3,12 @@ Test search tracing functionality.
|
||||
"""
|
||||
import pytest
|
||||
from hindsight_api.engine.memory_engine import Budget
|
||||
from hindsight_api import SearchTrace
|
||||
from hindsight_api import SearchTrace, RequestContext
|
||||
from datetime import datetime, timezone
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_search_with_trace(memory):
|
||||
async def test_search_with_trace(memory, request_context):
|
||||
"""Test that search with enable_trace=True returns a valid SearchTrace."""
|
||||
# Generate a unique agent ID for this test
|
||||
bank_id = f"test_trace_{datetime.now(timezone.utc).timestamp()}"
|
||||
@@ -20,16 +20,19 @@ async def test_search_with_trace(memory):
|
||||
bank_id=bank_id,
|
||||
content="Alice works at Google in Mountain View",
|
||||
context="test context",
|
||||
request_context=request_context,
|
||||
)
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="Bob also works at Google but in New York",
|
||||
context="test context",
|
||||
request_context=request_context,
|
||||
)
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="Charlie founded a startup called TechCorp",
|
||||
context="test context",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Search with tracing enabled
|
||||
@@ -40,6 +43,7 @@ async def test_search_with_trace(memory):
|
||||
budget=Budget.LOW, # 20,
|
||||
max_tokens=512,
|
||||
enable_trace=True,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Verify results
|
||||
@@ -102,11 +106,11 @@ async def test_search_with_trace(memory):
|
||||
|
||||
finally:
|
||||
# Cleanup
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_search_without_trace(memory):
|
||||
async def test_search_without_trace(memory, request_context):
|
||||
"""Test that search with enable_trace=False returns None for trace."""
|
||||
bank_id = f"test_no_trace_{datetime.now(timezone.utc).timestamp()}"
|
||||
|
||||
@@ -117,6 +121,7 @@ async def test_search_without_trace(memory):
|
||||
bank_id=bank_id,
|
||||
content="Test memory without trace",
|
||||
context="test",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Search without tracing
|
||||
@@ -127,6 +132,7 @@ async def test_search_without_trace(memory):
|
||||
budget=Budget.LOW, # 10,
|
||||
max_tokens=512,
|
||||
enable_trace=False,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Verify trace is None
|
||||
@@ -137,4 +143,4 @@ async def test_search_without_trace(memory):
|
||||
|
||||
finally:
|
||||
# Cleanup
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@@ -0,0 +1,136 @@
|
||||
"""
|
||||
Safety tests to ensure all SQL queries use fully-qualified table names.
|
||||
|
||||
This prevents cross-tenant data access by ensuring every table reference
|
||||
includes the schema prefix (e.g., public.memory_units instead of just memory_units).
|
||||
"""
|
||||
|
||||
import re
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
# All tables that MUST be schema-qualified in SQL queries
|
||||
TABLES = [
|
||||
"memory_units",
|
||||
"memory_links",
|
||||
"unit_entities",
|
||||
"entities",
|
||||
"entity_cooccurrences",
|
||||
"banks",
|
||||
"documents",
|
||||
"chunks",
|
||||
"async_operations",
|
||||
]
|
||||
|
||||
# Files to scan for SQL queries
|
||||
SCAN_PATHS = [
|
||||
"hindsight_api/engine",
|
||||
"hindsight_api/api",
|
||||
]
|
||||
|
||||
# Files to exclude (e.g., migrations, tests)
|
||||
EXCLUDE_PATTERNS = [
|
||||
"alembic",
|
||||
"__pycache__",
|
||||
"test_",
|
||||
]
|
||||
|
||||
|
||||
def get_python_files() -> list[Path]:
|
||||
"""Get all Python files to scan."""
|
||||
root = Path(__file__).parent.parent
|
||||
files = []
|
||||
for scan_path in SCAN_PATHS:
|
||||
path = root / scan_path
|
||||
if path.exists():
|
||||
for py_file in path.rglob("*.py"):
|
||||
# Check exclusions
|
||||
if any(excl in str(py_file) for excl in EXCLUDE_PATTERNS):
|
||||
continue
|
||||
files.append(py_file)
|
||||
return files
|
||||
|
||||
|
||||
def find_unqualified_table_refs(content: str, filename: str) -> list[tuple[int, str, str]]:
|
||||
"""
|
||||
Find SQL statements with unqualified table references.
|
||||
|
||||
Returns list of (line_number, table_name, line_content).
|
||||
"""
|
||||
violations = []
|
||||
|
||||
# Patterns that indicate SQL context
|
||||
sql_keywords = r"(?:FROM|JOIN|INTO|UPDATE|DELETE\s+FROM)\s+"
|
||||
|
||||
# Additional SQL indicators to confirm this is actually SQL, not prose
|
||||
sql_indicators = re.compile(
|
||||
r"(SELECT|INSERT|DELETE|UPDATE|CREATE|ALTER|DROP|WHERE|SET|VALUES|"
|
||||
r'f"""|f\'\'\'|""".*SELECT|\'\'\'.*SELECT)',
|
||||
re.IGNORECASE,
|
||||
)
|
||||
|
||||
lines = content.split("\n")
|
||||
for line_num, line in enumerate(lines, 1):
|
||||
# Skip comments and strings that are clearly not SQL
|
||||
stripped = line.strip()
|
||||
if stripped.startswith("#"):
|
||||
continue
|
||||
|
||||
for table in TABLES:
|
||||
# Pattern: SQL keyword followed by unqualified table name
|
||||
# Should match: FROM memory_units, JOIN memory_units, INTO memory_units
|
||||
# Should NOT match: FROM public.memory_units, FROM {schema}.memory_units
|
||||
# Should NOT match: fq_table("memory_units")
|
||||
|
||||
# Check for unqualified table after SQL keyword
|
||||
pattern = rf"{sql_keywords}{table}(?:\s|$|,|\))"
|
||||
|
||||
if re.search(pattern, line, re.IGNORECASE):
|
||||
# Check if it's actually qualified (has schema prefix)
|
||||
qualified_pattern = rf"\.\s*{table}(?:\s|$|,|\))"
|
||||
fq_table_pattern = rf'fq_table\s*\(\s*["\']?{table}'
|
||||
|
||||
if not re.search(qualified_pattern, line) and not re.search(
|
||||
fq_table_pattern, line
|
||||
):
|
||||
# Additional check: line must have SQL indicators
|
||||
# This avoids false positives in docstrings like "split into chunks"
|
||||
if sql_indicators.search(line):
|
||||
violations.append((line_num, table, stripped))
|
||||
|
||||
return violations
|
||||
|
||||
|
||||
class TestSQLSchemaSafety:
|
||||
"""Ensure all SQL uses schema-qualified table names."""
|
||||
|
||||
def test_no_unqualified_table_references(self):
|
||||
"""All SQL queries must use fq_table() or schema.table format."""
|
||||
all_violations = []
|
||||
|
||||
for py_file in get_python_files():
|
||||
content = py_file.read_text()
|
||||
violations = find_unqualified_table_refs(content, py_file.name)
|
||||
|
||||
for line_num, table, line in violations:
|
||||
all_violations.append(
|
||||
f"{py_file.relative_to(py_file.parent.parent)}:{line_num} - "
|
||||
f"unqualified '{table}': {line[:80]}..."
|
||||
)
|
||||
|
||||
if all_violations:
|
||||
msg = (
|
||||
f"Found {len(all_violations)} unqualified table references!\n"
|
||||
"These could cause cross-tenant data access.\n"
|
||||
"Use fq_table('table_name') for all table references.\n\n"
|
||||
+ "\n".join(all_violations[:20]) # Show first 20
|
||||
)
|
||||
if len(all_violations) > 20:
|
||||
msg += f"\n... and {len(all_violations) - 20} more"
|
||||
pytest.fail(msg)
|
||||
|
||||
def test_tables_list_is_complete(self):
|
||||
"""Verify we're checking for all tables (sanity check)."""
|
||||
# This is a sanity check - if you add a new table, add it to TABLES
|
||||
assert len(TABLES) >= 9, "Update TABLES list if you added new tables"
|
||||
@@ -3,16 +3,17 @@ import asyncio
|
||||
from datetime import datetime, timezone, timedelta
|
||||
import pytest
|
||||
from hindsight_api.engine.memory_engine import Budget
|
||||
from hindsight_api import RequestContext
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_temporal_ranges_are_written(memory):
|
||||
async def test_temporal_ranges_are_written(memory, request_context):
|
||||
"""Test that occurred_start, occurred_end, and mentioned_at are actually written to database."""
|
||||
bank_id = "test_temporal_ranges"
|
||||
|
||||
# Clean up any existing data
|
||||
try:
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
@@ -23,7 +24,8 @@ async def test_temporal_ranges_are_written(memory):
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content=text1,
|
||||
event_date=conversation_date
|
||||
event_date=conversation_date,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Test 2: Period event (month range)
|
||||
@@ -32,7 +34,8 @@ async def test_temporal_ranges_are_written(memory):
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content=text2,
|
||||
event_date=conversation_date
|
||||
event_date=conversation_date,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Give it a moment for async processing
|
||||
@@ -114,7 +117,8 @@ async def test_temporal_ranges_are_written(memory):
|
||||
query="pottery workshop",
|
||||
fact_type=["world", "experience"],
|
||||
budget=Budget.LOW,
|
||||
max_tokens=4096
|
||||
max_tokens=4096,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
print(f"Found {len(search_result.results)} search results")
|
||||
@@ -132,4 +136,4 @@ async def test_temporal_ranges_are_written(memory):
|
||||
print("⚠ Temporal fields not yet populated in search results (known issue)")
|
||||
|
||||
# Clean up
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@@ -4,10 +4,11 @@ Test think function for opinion generation and consistency.
|
||||
import pytest
|
||||
from datetime import datetime, timezone
|
||||
from hindsight_api.engine.memory_engine import Budget
|
||||
from hindsight_api import RequestContext
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_think_opinion_consistency(memory):
|
||||
async def test_think_opinion_consistency(memory, request_context):
|
||||
"""
|
||||
Test that think function:
|
||||
1. Generates an opinion
|
||||
@@ -23,14 +24,16 @@ async def test_think_opinion_consistency(memory):
|
||||
bank_id=bank_id,
|
||||
content="Alice is a software engineer who has worked on 5 major projects. She always delivers on time and writes clean, well-documented code.",
|
||||
context="performance review",
|
||||
event_date=datetime(2024, 1, 15, tzinfo=timezone.utc)
|
||||
event_date=datetime(2024, 1, 15, tzinfo=timezone.utc),
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="Bob recently joined the team. He missed his first deadline and his code had many bugs.",
|
||||
context="performance review",
|
||||
event_date=datetime(2024, 2, 1, tzinfo=timezone.utc)
|
||||
event_date=datetime(2024, 2, 1, tzinfo=timezone.utc),
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# First think call - should generate opinions
|
||||
@@ -39,6 +42,7 @@ async def test_think_opinion_consistency(memory):
|
||||
bank_id=bank_id,
|
||||
query=query,
|
||||
budget=Budget.LOW,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
print(f"\n=== First Think Call ===")
|
||||
@@ -82,6 +86,7 @@ async def test_think_opinion_consistency(memory):
|
||||
bank_id=bank_id,
|
||||
query=query,
|
||||
budget=Budget.LOW,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
print(f"\n=== Second Think Call ===")
|
||||
@@ -122,13 +127,13 @@ async def test_think_opinion_consistency(memory):
|
||||
finally:
|
||||
# Clean up agent data
|
||||
try:
|
||||
await memory.delete_bank(bank_id)
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
except Exception as e:
|
||||
print(f"Warning: Error during cleanup: {e}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_think_without_prior_context(memory):
|
||||
async def test_think_without_prior_context(memory, request_context):
|
||||
"""
|
||||
Test that think function handles queries when there's no relevant context.
|
||||
"""
|
||||
@@ -139,6 +144,7 @@ async def test_think_without_prior_context(memory):
|
||||
bank_id=bank_id,
|
||||
query="What is the capital of France?",
|
||||
budget=Budget.LOW,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
print(f"\n=== Think Without Context ===")
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
[package]
|
||||
name = "hindsight-cli"
|
||||
version = "0.1.8"
|
||||
version = "0.1.13"
|
||||
edition = "2021"
|
||||
authors = ["Hindsight Team"]
|
||||
description = "A beautiful CLI for Hindsight - semantic memory system"
|
||||
|
||||
@@ -64,13 +64,24 @@ pub struct ApiClient {
|
||||
}
|
||||
|
||||
impl ApiClient {
|
||||
pub fn new(base_url: String) -> Result<Self> {
|
||||
pub fn new(base_url: String, api_key: Option<String>) -> Result<Self> {
|
||||
let runtime = std::sync::Arc::new(tokio::runtime::Runtime::new()?);
|
||||
|
||||
// Create HTTP client with 2-minute timeout
|
||||
let http_client = reqwest::Client::builder()
|
||||
.timeout(std::time::Duration::from_secs(120))
|
||||
.build()?;
|
||||
// Create HTTP client with 2-minute timeout and optional auth header
|
||||
let mut client_builder = reqwest::Client::builder()
|
||||
.timeout(std::time::Duration::from_secs(120));
|
||||
|
||||
if let Some(key) = api_key {
|
||||
let mut headers = reqwest::header::HeaderMap::new();
|
||||
let auth_value = format!("Bearer {}", key);
|
||||
headers.insert(
|
||||
reqwest::header::AUTHORIZATION,
|
||||
reqwest::header::HeaderValue::from_str(&auth_value)?,
|
||||
);
|
||||
client_builder = client_builder.default_headers(headers);
|
||||
}
|
||||
|
||||
let http_client = client_builder.build()?;
|
||||
|
||||
let client = AsyncClient::new_with_client(&base_url, http_client);
|
||||
Ok(ApiClient { client, runtime })
|
||||
|
||||
@@ -944,7 +944,7 @@ fn render_banks(f: &mut Frame, app: &mut App, area: Rect) {
|
||||
.banks
|
||||
.iter()
|
||||
.map(|bank| {
|
||||
let name = if bank.name.is_empty() { "Unnamed" } else { &bank.name };
|
||||
let name = bank.name.as_deref().filter(|s| !s.is_empty()).unwrap_or("Unnamed");
|
||||
let content = format!("{} - {}", bank.bank_id, name);
|
||||
ListItem::new(content).style(Style::default().fg(Color::White))
|
||||
})
|
||||
|
||||
+38
-12
@@ -10,6 +10,7 @@ const CONFIG_DIR_NAME: &str = ".hindsight";
|
||||
|
||||
pub struct Config {
|
||||
pub api_url: String,
|
||||
pub api_key: Option<String>,
|
||||
pub source: ConfigSource,
|
||||
}
|
||||
|
||||
@@ -32,22 +33,27 @@ impl std::fmt::Display for ConfigSource {
|
||||
|
||||
impl Config {
|
||||
/// Load configuration with the following priority:
|
||||
/// 1. Environment variable (HINDSIGHT_API_URL) - highest priority, for overrides
|
||||
/// 1. Environment variable (HINDSIGHT_API_URL, HINDSIGHT_API_KEY) - highest priority, for overrides
|
||||
/// 2. Local config file (~/.hindsight/config.toml)
|
||||
/// 3. Default (http://localhost:8888)
|
||||
pub fn load() -> Result<Self> {
|
||||
// Load API key from environment (highest priority)
|
||||
let env_api_key = env::var("HINDSIGHT_API_KEY").ok();
|
||||
|
||||
// 1. Environment variable takes highest priority (for overrides)
|
||||
if let Ok(api_url) = env::var("HINDSIGHT_API_URL") {
|
||||
return Self::validate_and_create(api_url, ConfigSource::Environment);
|
||||
return Self::validate_and_create(api_url, env_api_key, ConfigSource::Environment);
|
||||
}
|
||||
|
||||
// 2. Try local config file
|
||||
if let Some(api_url) = Self::load_from_file()? {
|
||||
return Self::validate_and_create(api_url, ConfigSource::LocalFile);
|
||||
if let Some((api_url, file_api_key)) = Self::load_from_file()? {
|
||||
// Environment api_key takes precedence over file api_key
|
||||
let api_key = env_api_key.or(file_api_key);
|
||||
return Self::validate_and_create(api_url, api_key, ConfigSource::LocalFile);
|
||||
}
|
||||
|
||||
// 3. Fall back to default
|
||||
Self::validate_and_create(DEFAULT_API_URL.to_string(), ConfigSource::Default)
|
||||
Self::validate_and_create(DEFAULT_API_URL.to_string(), env_api_key, ConfigSource::Default)
|
||||
}
|
||||
|
||||
/// Legacy method for backwards compatibility
|
||||
@@ -55,14 +61,14 @@ impl Config {
|
||||
Self::load()
|
||||
}
|
||||
|
||||
fn validate_and_create(api_url: String, source: ConfigSource) -> Result<Self> {
|
||||
fn validate_and_create(api_url: String, api_key: Option<String>, source: ConfigSource) -> Result<Self> {
|
||||
if !api_url.starts_with("http://") && !api_url.starts_with("https://") {
|
||||
anyhow::bail!(
|
||||
"Invalid API URL: {}. Must start with http:// or https://",
|
||||
api_url
|
||||
);
|
||||
}
|
||||
Ok(Config { api_url, source })
|
||||
Ok(Config { api_url, api_key, source })
|
||||
}
|
||||
|
||||
fn config_dir() -> Option<PathBuf> {
|
||||
@@ -73,7 +79,7 @@ impl Config {
|
||||
Self::config_dir().map(|dir| dir.join(CONFIG_FILE_NAME))
|
||||
}
|
||||
|
||||
fn load_from_file() -> Result<Option<String>> {
|
||||
fn load_from_file() -> Result<Option<(String, Option<String>)>> {
|
||||
let config_path = match Self::config_file_path() {
|
||||
Some(path) => path,
|
||||
None => return Ok(None),
|
||||
@@ -86,23 +92,40 @@ impl Config {
|
||||
let content = fs::read_to_string(&config_path)
|
||||
.with_context(|| format!("Failed to read config file: {}", config_path.display()))?;
|
||||
|
||||
// Simple TOML parsing for api_url
|
||||
let mut api_url: Option<String> = None;
|
||||
let mut api_key: Option<String> = None;
|
||||
|
||||
// Simple TOML parsing for api_url and api_key
|
||||
for line in content.lines() {
|
||||
let line = line.trim();
|
||||
if line.starts_with("api_url") {
|
||||
if let Some(value) = line.split('=').nth(1) {
|
||||
let value = value.trim().trim_matches('"').trim_matches('\'');
|
||||
if !value.is_empty() {
|
||||
return Ok(Some(value.to_string()));
|
||||
api_url = Some(value.to_string());
|
||||
}
|
||||
}
|
||||
} else if line.starts_with("api_key") {
|
||||
if let Some(value) = line.split('=').nth(1) {
|
||||
let value = value.trim().trim_matches('"').trim_matches('\'');
|
||||
if !value.is_empty() {
|
||||
api_key = Some(value.to_string());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(None)
|
||||
match api_url {
|
||||
Some(url) => Ok(Some((url, api_key))),
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn save_api_url(api_url: &str) -> Result<PathBuf> {
|
||||
Self::save_config(api_url, None)
|
||||
}
|
||||
|
||||
pub fn save_config(api_url: &str, api_key: Option<&str>) -> Result<PathBuf> {
|
||||
let config_dir = Self::config_dir()
|
||||
.ok_or_else(|| anyhow::anyhow!("Could not determine home directory"))?;
|
||||
|
||||
@@ -113,7 +136,10 @@ impl Config {
|
||||
}
|
||||
|
||||
let config_path = config_dir.join(CONFIG_FILE_NAME);
|
||||
let content = format!("api_url = \"{}\"\n", api_url);
|
||||
let mut content = format!("api_url = \"{}\"\n", api_url);
|
||||
if let Some(key) = api_key {
|
||||
content.push_str(&format!("api_key = \"{}\"\n", key));
|
||||
}
|
||||
|
||||
fs::write(&config_path, content)
|
||||
.with_context(|| format!("Failed to write config file: {}", config_path.display()))?;
|
||||
|
||||
@@ -91,12 +91,18 @@ enum Commands {
|
||||
#[command(alias = "tui")]
|
||||
Explore,
|
||||
|
||||
/// Configure the CLI (API URL, etc.)
|
||||
#[command(after_help = "Configuration priority:\n 1. Environment variable (HINDSIGHT_API_URL) - highest priority\n 2. Config file (~/.hindsight/config)\n 3. Default (http://localhost:8888)")]
|
||||
/// Launch the web-based control plane UI
|
||||
Ui,
|
||||
|
||||
/// Configure the CLI (API URL, API key, etc.)
|
||||
#[command(after_help = "Configuration priority:\n 1. Environment variables (HINDSIGHT_API_URL, HINDSIGHT_API_KEY) - highest priority\n 2. Config file (~/.hindsight/config)\n 3. Default (http://localhost:8888)")]
|
||||
Configure {
|
||||
/// API URL to connect to (interactive prompt if not provided)
|
||||
#[arg(long)]
|
||||
api_url: Option<String>,
|
||||
/// API key for authentication (sent as Bearer token)
|
||||
#[arg(long)]
|
||||
api_key: Option<String>,
|
||||
},
|
||||
}
|
||||
|
||||
@@ -369,8 +375,13 @@ fn run() -> Result<()> {
|
||||
let verbose = cli.verbose;
|
||||
|
||||
// Handle configure command before loading full config (it doesn't need API client)
|
||||
if let Commands::Configure { api_url } = cli.command {
|
||||
return handle_configure(api_url, output_format);
|
||||
if let Commands::Configure { api_url, api_key } = cli.command {
|
||||
return handle_configure(api_url, api_key, output_format);
|
||||
}
|
||||
|
||||
// Handle ui command - needs config but not API client
|
||||
if let Commands::Ui = cli.command {
|
||||
return handle_ui(output_format);
|
||||
}
|
||||
|
||||
// Load configuration
|
||||
@@ -381,15 +392,17 @@ fn run() -> Result<()> {
|
||||
});
|
||||
|
||||
let api_url = config.api_url().to_string();
|
||||
let api_key = config.api_key.clone();
|
||||
|
||||
// Create API client
|
||||
let client = ApiClient::new(api_url.clone()).unwrap_or_else(|e| {
|
||||
let client = ApiClient::new(api_url.clone(), api_key).unwrap_or_else(|e| {
|
||||
errors::handle_api_error(e, &api_url);
|
||||
});
|
||||
|
||||
// Execute command and handle errors
|
||||
let result: Result<()> = match cli.command {
|
||||
Commands::Configure { .. } => unreachable!(), // Handled above
|
||||
Commands::Ui => unreachable!(), // Handled above
|
||||
Commands::Explore => commands::explore::run(&client),
|
||||
Commands::Bank(bank_cmd) => match bank_cmd {
|
||||
BankCommands::List => commands::bank::list(&client, verbose, output_format),
|
||||
@@ -467,7 +480,7 @@ fn run() -> Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn handle_configure(api_url: Option<String>, output_format: OutputFormat) -> Result<()> {
|
||||
fn handle_configure(api_url: Option<String>, api_key: Option<String>, output_format: OutputFormat) -> Result<()> {
|
||||
// Load current config to show current state
|
||||
let current_config = Config::load().ok();
|
||||
|
||||
@@ -478,6 +491,15 @@ fn handle_configure(api_url: Option<String>, output_format: OutputFormat) -> Res
|
||||
// Show current configuration
|
||||
if let Some(ref config) = current_config {
|
||||
println!(" Current API URL: {}", config.api_url);
|
||||
if let Some(ref key) = config.api_key {
|
||||
// Mask the API key for display
|
||||
let masked = if key.len() > 8 {
|
||||
format!("{}...{}", &key[..4], &key[key.len()-4..])
|
||||
} else {
|
||||
"****".to_string()
|
||||
};
|
||||
println!(" Current API Key: {}", masked);
|
||||
}
|
||||
println!(" Source: {}", config.source);
|
||||
println!();
|
||||
}
|
||||
@@ -502,18 +524,30 @@ fn handle_configure(api_url: Option<String>, output_format: OutputFormat) -> Res
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
// Use provided api_key, or keep existing one if not provided
|
||||
let new_api_key = api_key.or_else(|| current_config.as_ref().and_then(|c| c.api_key.clone()));
|
||||
|
||||
// Save to config file
|
||||
let config_path = Config::save_api_url(&new_api_url)?;
|
||||
let config_path = Config::save_config(&new_api_url, new_api_key.as_deref())?;
|
||||
|
||||
if output_format == OutputFormat::Pretty {
|
||||
ui::print_success(&format!("Configuration saved to {}", config_path.display()));
|
||||
println!();
|
||||
println!(" API URL: {}", new_api_url);
|
||||
if let Some(ref key) = new_api_key {
|
||||
let masked = if key.len() > 8 {
|
||||
format!("{}...{}", &key[..4], &key[key.len()-4..])
|
||||
} else {
|
||||
"****".to_string()
|
||||
};
|
||||
println!(" API Key: {}", masked);
|
||||
}
|
||||
println!();
|
||||
println!("Note: Environment variable HINDSIGHT_API_URL will override this setting.");
|
||||
println!("Note: Environment variables HINDSIGHT_API_URL and HINDSIGHT_API_KEY will override these settings.");
|
||||
} else {
|
||||
let result = serde_json::json!({
|
||||
"api_url": new_api_url,
|
||||
"api_key_set": new_api_key.is_some(),
|
||||
"config_path": config_path.display().to_string(),
|
||||
});
|
||||
output::print_output(&result, output_format)?;
|
||||
@@ -521,3 +555,50 @@ fn handle_configure(api_url: Option<String>, output_format: OutputFormat) -> Res
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn handle_ui(output_format: OutputFormat) -> Result<()> {
|
||||
use std::process::Command;
|
||||
|
||||
// Load configuration to get the API URL
|
||||
let config = Config::load().unwrap_or_else(|e| {
|
||||
ui::print_error(&format!("Configuration error: {}", e));
|
||||
errors::print_config_help();
|
||||
std::process::exit(1);
|
||||
});
|
||||
|
||||
let api_url = config.api_url();
|
||||
|
||||
if output_format == OutputFormat::Pretty {
|
||||
ui::print_info("Launching Hindsight Control Plane UI...");
|
||||
println!();
|
||||
println!(" API URL: {}", api_url);
|
||||
println!();
|
||||
}
|
||||
|
||||
// Run npx @vectorize-io/hindsight-control-plane --api-url {api_url}
|
||||
let status = Command::new("npx")
|
||||
.arg("@vectorize-io/hindsight-control-plane")
|
||||
.arg("--api-url")
|
||||
.arg(api_url)
|
||||
.status();
|
||||
|
||||
match status {
|
||||
Ok(exit_status) => {
|
||||
if !exit_status.success() {
|
||||
if let Some(code) = exit_status.code() {
|
||||
std::process::exit(code);
|
||||
} else {
|
||||
std::process::exit(1);
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
ui::print_error(&format!("Failed to launch control plane UI: {}", e));
|
||||
ui::print_info("Make sure you have Node.js and npm installed.");
|
||||
ui::print_info("You can also install the control plane globally: npm install -g @vectorize-io/hindsight-control-plane");
|
||||
std::process::exit(1);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -5,7 +5,7 @@ use crate::output::OutputFormat;
|
||||
|
||||
/// Get API client from config
|
||||
pub fn get_client(config: &Config) -> Result<ApiClient> {
|
||||
ApiClient::new(config.api_url.clone())
|
||||
ApiClient::new(config.api_url.clone(), config.api_key.clone())
|
||||
.context("Failed to create API client")
|
||||
}
|
||||
|
||||
|
||||
@@ -44,8 +44,12 @@ class Hindsight:
|
||||
```python
|
||||
from hindsight_client import Hindsight
|
||||
|
||||
# Without authentication
|
||||
client = Hindsight(base_url="http://localhost:8888")
|
||||
|
||||
# With API key authentication
|
||||
client = Hindsight(base_url="http://localhost:8888", api_key="your-api-key")
|
||||
|
||||
# Store a memory
|
||||
client.retain(bank_id="alice", content="Alice loves AI")
|
||||
|
||||
@@ -59,15 +63,16 @@ class Hindsight:
|
||||
```
|
||||
"""
|
||||
|
||||
def __init__(self, base_url: str, timeout: float = 30.0):
|
||||
def __init__(self, base_url: str, api_key: Optional[str] = None, timeout: float = 30.0):
|
||||
"""
|
||||
Initialize the Hindsight client.
|
||||
|
||||
Args:
|
||||
base_url: The base URL of the Hindsight API server
|
||||
api_key: Optional API key for authentication (sent as Bearer token)
|
||||
timeout: Request timeout in seconds (default: 30.0)
|
||||
"""
|
||||
config = hindsight_client_api.Configuration(host=base_url)
|
||||
config = hindsight_client_api.Configuration(host=base_url, access_token=api_key)
|
||||
self._api_client = hindsight_client_api.ApiClient(config)
|
||||
self._api = default_api.DefaultApi(self._api_client)
|
||||
|
||||
@@ -80,9 +85,21 @@ class Hindsight:
|
||||
self.close()
|
||||
|
||||
def close(self):
|
||||
"""Close the API client."""
|
||||
"""Close the API client (sync version - use aclose() in async code)."""
|
||||
if self._api_client:
|
||||
_run_async(self._api_client.close())
|
||||
try:
|
||||
loop = asyncio.get_running_loop()
|
||||
# We're in an async context - schedule but don't wait
|
||||
# The caller should use aclose() instead
|
||||
loop.create_task(self._api_client.close())
|
||||
except RuntimeError:
|
||||
# No running loop - safe to run synchronously
|
||||
_run_async(self._api_client.close())
|
||||
|
||||
async def aclose(self):
|
||||
"""Close the API client (async version)."""
|
||||
if self._api_client:
|
||||
await self._api_client.close()
|
||||
|
||||
# Simplified methods for main operations
|
||||
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
[project]
|
||||
name = "hindsight-client"
|
||||
version = "0.1.8"
|
||||
version = "0.1.13"
|
||||
description = "Python client for Hindsight - Semantic memory system with personality-driven thinking"
|
||||
authors = [
|
||||
{name = "Hindsight Team"}
|
||||
@@ -20,6 +20,7 @@ dependencies = [
|
||||
test = [
|
||||
"pytest>=7.0.0",
|
||||
"pytest-asyncio>=0.21.0",
|
||||
"requests>=2.28.0",
|
||||
]
|
||||
|
||||
[build-system]
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@vectorize-io/hindsight-client",
|
||||
"version": "0.1.8",
|
||||
"version": "0.1.13",
|
||||
"description": "TypeScript client for Hindsight - Semantic memory system with personality-driven thinking",
|
||||
"main": "./dist/src/index.js",
|
||||
"types": "./dist/src/index.d.ts",
|
||||
|
||||
@@ -5,8 +5,15 @@
|
||||
* ```typescript
|
||||
* import { HindsightClient } from '@vectorize-io/hindsight-client';
|
||||
*
|
||||
* // Without authentication
|
||||
* const client = new HindsightClient({ baseUrl: 'http://localhost:8888' });
|
||||
*
|
||||
* // With API key authentication
|
||||
* const client = new HindsightClient({
|
||||
* baseUrl: 'http://localhost:8888',
|
||||
* apiKey: 'your-api-key'
|
||||
* });
|
||||
*
|
||||
* // Retain a memory
|
||||
* await client.retain('alice', 'Alice loves AI');
|
||||
*
|
||||
@@ -37,6 +44,10 @@ import type {
|
||||
|
||||
export interface HindsightClientOptions {
|
||||
baseUrl: string;
|
||||
/**
|
||||
* Optional API key for authentication (sent as Bearer token in Authorization header)
|
||||
*/
|
||||
apiKey?: string;
|
||||
}
|
||||
|
||||
export interface MemoryItemInput {
|
||||
@@ -54,6 +65,9 @@ export class HindsightClient {
|
||||
this.client = createClient(
|
||||
createConfig({
|
||||
baseUrl: options.baseUrl,
|
||||
headers: options.apiKey
|
||||
? { Authorization: `Bearer ${options.apiKey}` }
|
||||
: undefined,
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
@@ -14,6 +14,7 @@
|
||||
|
||||
# production
|
||||
/build
|
||||
/standalone
|
||||
|
||||
# misc
|
||||
.DS_Store
|
||||
|
||||
Executable
+86
@@ -0,0 +1,86 @@
|
||||
#!/usr/bin/env node
|
||||
|
||||
const { spawn } = require('child_process');
|
||||
const path = require('path');
|
||||
const fs = require('fs');
|
||||
|
||||
const args = process.argv.slice(2);
|
||||
|
||||
// Parse command line arguments
|
||||
let port = process.env.PORT || 9999;
|
||||
let hostname = process.env.HOSTNAME || '0.0.0.0';
|
||||
let apiUrl = process.env.HINDSIGHT_CP_DATAPLANE_API_URL;
|
||||
|
||||
for (let i = 0; i < args.length; i++) {
|
||||
if (args[i] === '--port' || args[i] === '-p') {
|
||||
port = args[++i];
|
||||
} else if (args[i] === '--hostname' || args[i] === '-H') {
|
||||
hostname = args[++i];
|
||||
} else if (args[i] === '--api-url' || args[i] === '-a') {
|
||||
apiUrl = args[++i];
|
||||
} else if (args[i] === '--help' || args[i] === '-h') {
|
||||
console.log(`
|
||||
Hindsight Control Plane
|
||||
|
||||
Usage: hindsight-control-plane [options]
|
||||
|
||||
Options:
|
||||
-p, --port <port> Port to listen on (default: 9999, env: PORT)
|
||||
-H, --hostname <host> Hostname to bind to (default: 0.0.0.0, env: HOSTNAME)
|
||||
-a, --api-url <url> Hindsight API URL (env: HINDSIGHT_CP_DATAPLANE_API_URL)
|
||||
-h, --help Show this help message
|
||||
|
||||
Environment Variables:
|
||||
PORT Port to listen on
|
||||
HOSTNAME Hostname to bind to
|
||||
HINDSIGHT_CP_DATAPLANE_API_URL URL of the Hindsight API server
|
||||
`);
|
||||
process.exit(0);
|
||||
}
|
||||
}
|
||||
|
||||
// Find the standalone server
|
||||
const standaloneDir = path.join(__dirname, '..', 'standalone');
|
||||
const serverPath = path.join(standaloneDir, 'server.js');
|
||||
|
||||
if (!fs.existsSync(serverPath)) {
|
||||
console.error('Error: Standalone server not found at', serverPath);
|
||||
console.error('This package may not have been built correctly.');
|
||||
process.exit(1);
|
||||
}
|
||||
|
||||
// Set up environment
|
||||
const env = {
|
||||
...process.env,
|
||||
PORT: String(port),
|
||||
HOSTNAME: hostname,
|
||||
};
|
||||
|
||||
if (apiUrl) {
|
||||
env.HINDSIGHT_CP_DATAPLANE_API_URL = apiUrl;
|
||||
}
|
||||
|
||||
console.log(`Starting Hindsight Control Plane on http://${hostname}:${port}`);
|
||||
if (apiUrl) {
|
||||
console.log(`API URL: ${apiUrl}`);
|
||||
}
|
||||
|
||||
// Run the standalone server
|
||||
const server = spawn('node', [serverPath], {
|
||||
cwd: standaloneDir,
|
||||
env,
|
||||
stdio: 'inherit',
|
||||
});
|
||||
|
||||
server.on('error', (err) => {
|
||||
console.error('Failed to start server:', err.message);
|
||||
process.exit(1);
|
||||
});
|
||||
|
||||
server.on('close', (code) => {
|
||||
process.exit(code || 0);
|
||||
});
|
||||
|
||||
// Handle signals
|
||||
process.on('SIGTERM', () => server.kill('SIGTERM'));
|
||||
process.on('SIGINT', () => server.kill('SIGINT'));
|
||||
@@ -1,7 +1,14 @@
|
||||
import type { NextConfig } from "next";
|
||||
import path from "path";
|
||||
|
||||
const nextConfig: NextConfig = {
|
||||
output: 'standalone',
|
||||
// Disable request logging in production
|
||||
logging: false,
|
||||
// Set the monorepo root explicitly to avoid detecting wrong lockfiles in parent directories
|
||||
turbopack: {
|
||||
root: path.resolve(__dirname, '..'),
|
||||
},
|
||||
};
|
||||
|
||||
export default nextConfig;
|
||||
|
||||
@@ -1,17 +1,26 @@
|
||||
{
|
||||
"name": "hindsight-control-plane",
|
||||
"version": "0.1.8",
|
||||
"private": true,
|
||||
"name": "@vectorize-io/hindsight-control-plane",
|
||||
"version": "0.1.13",
|
||||
"description": "Control plane for Hindsight - Semantic memory system",
|
||||
"bin": {
|
||||
"hindsight-control-plane": "./bin/cli.js"
|
||||
},
|
||||
"files": [
|
||||
"bin",
|
||||
"standalone",
|
||||
"public"
|
||||
],
|
||||
"scripts": {
|
||||
"dev": "next dev",
|
||||
"build": "next build",
|
||||
"build": "next build && npm run build:standalone",
|
||||
"build:standalone": "rm -rf standalone && STANDALONE_ROOT=$(find .next/standalone -path '*/node_modules' -prune -o -name 'server.js' -print | head -1 | xargs dirname) && cp -r \"$STANDALONE_ROOT\" standalone && cp -r .next/standalone/node_modules standalone/node_modules && mkdir -p standalone/.next && cp -r .next/static standalone/.next/static && mkdir -p standalone/public && cp -r public/* standalone/public/ 2>/dev/null || true",
|
||||
"start": "next start",
|
||||
"lint": "next lint"
|
||||
"lint": "next lint",
|
||||
"prepublishOnly": "npm run build"
|
||||
},
|
||||
"keywords": [],
|
||||
"keywords": ["hindsight", "memory", "semantic", "ai"],
|
||||
"author": "Hindsight Team",
|
||||
"license": "ISC",
|
||||
"description": "Control plane for Hindsight - Semantic memory system",
|
||||
"dependencies": {
|
||||
"@radix-ui/react-checkbox": "^1.3.3",
|
||||
"@radix-ui/react-dialog": "^1.1.15",
|
||||
@@ -27,7 +36,6 @@
|
||||
"@types/node": "^24.10.0",
|
||||
"@types/react": "^19.2.2",
|
||||
"@types/react-dom": "^19.2.2",
|
||||
"@vectorize-io/hindsight-client": "file:../hindsight-clients/typescript",
|
||||
"autoprefixer": "^10.4.21",
|
||||
"class-variance-authority": "^0.7.1",
|
||||
"clsx": "^2.1.1",
|
||||
@@ -50,6 +58,7 @@
|
||||
"typescript": "^5.9.3"
|
||||
},
|
||||
"devDependencies": {
|
||||
"@vectorize-io/hindsight-client": "file:../hindsight-clients/typescript",
|
||||
"@eslint/eslintrc": "^3.3.3",
|
||||
"@eslint/js": "^9.39.2",
|
||||
"eslint-plugin-react": "^7.37.5",
|
||||
|
||||
@@ -0,0 +1,35 @@
|
||||
import { NextResponse } from "next/server";
|
||||
import { sdk, lowLevelClient } from "@/lib/hindsight-client";
|
||||
|
||||
export async function GET() {
|
||||
const status: {
|
||||
status: string;
|
||||
service: string;
|
||||
dataplane?: {
|
||||
status: string;
|
||||
url: string;
|
||||
error?: string;
|
||||
};
|
||||
} = {
|
||||
status: "ok",
|
||||
service: "hindsight-control-plane",
|
||||
};
|
||||
|
||||
// Check dataplane connectivity
|
||||
const dataplaneUrl = process.env.HINDSIGHT_CP_DATAPLANE_API_URL || "http://localhost:8888";
|
||||
try {
|
||||
await sdk.listBanks({ client: lowLevelClient });
|
||||
status.dataplane = {
|
||||
status: "connected",
|
||||
url: dataplaneUrl,
|
||||
};
|
||||
} catch (error) {
|
||||
status.dataplane = {
|
||||
status: "disconnected",
|
||||
url: dataplaneUrl,
|
||||
error: error instanceof Error ? error.message : String(error),
|
||||
};
|
||||
}
|
||||
|
||||
return NextResponse.json(status, { status: 200 });
|
||||
}
|
||||
@@ -7,17 +7,6 @@ export async function POST(request: NextRequest) {
|
||||
const bankId = body.bank_id || body.agent_id || "default";
|
||||
const { query, types, fact_type, max_tokens, trace, budget, include, query_timestamp } = body;
|
||||
|
||||
console.log("[Recall API] Request:", {
|
||||
bankId,
|
||||
query,
|
||||
types: types || fact_type,
|
||||
max_tokens,
|
||||
trace,
|
||||
budget,
|
||||
query_timestamp,
|
||||
});
|
||||
console.log("[Recall API] Include options:", JSON.stringify(include, null, 2));
|
||||
|
||||
const response = await sdk.recallMemories({
|
||||
client: lowLevelClient,
|
||||
path: { bank_id: bankId },
|
||||
@@ -37,18 +26,6 @@ export async function POST(request: NextRequest) {
|
||||
throw new Error(`API returned no data: ${JSON.stringify(response.error || "Unknown error")}`);
|
||||
}
|
||||
|
||||
console.log("[Recall API] Response structure:", {
|
||||
hasResults: !!response.data?.results,
|
||||
resultsCount: response.data?.results?.length,
|
||||
hasTrace: !!response.data?.trace,
|
||||
hasEntities: !!response.data?.entities,
|
||||
entitiesType: typeof response.data?.entities,
|
||||
entitiesKeys: response.data?.entities ? Object.keys(response.data.entities) : null,
|
||||
hasChunks: !!response.data?.chunks,
|
||||
chunksType: typeof response.data?.chunks,
|
||||
chunksKeys: response.data?.chunks ? Object.keys(response.data.chunks) : null,
|
||||
});
|
||||
|
||||
// Return a clean JSON object by spreading the response
|
||||
// This ensures any non-serializable properties are excluded
|
||||
const jsonResponse = {
|
||||
|
||||
@@ -39,7 +39,8 @@ export function AddMemoryView() {
|
||||
try {
|
||||
const item: any = { content };
|
||||
if (context) item.context = context;
|
||||
if (eventDate) item.timestamp = eventDate;
|
||||
// datetime-local gives "2024-01-15T10:30", add seconds for proper ISO format
|
||||
if (eventDate) item.timestamp = eventDate + ":00";
|
||||
|
||||
const data: any = await client.retain({
|
||||
bank_id: currentBank,
|
||||
|
||||
@@ -85,7 +85,8 @@ function BankSelectorInner() {
|
||||
try {
|
||||
const item: any = { content: docContent };
|
||||
if (docContext) item.context = docContext;
|
||||
if (docEventDate) item.event_date = docEventDate;
|
||||
// datetime-local gives "2024-01-15T10:30", add seconds for proper ISO format
|
||||
if (docEventDate) item.timestamp = docEventDate + ":00";
|
||||
|
||||
const params: any = {
|
||||
bank_id: currentBank,
|
||||
|
||||
@@ -102,12 +102,6 @@ export function DataView({ factType }: DataViewProps) {
|
||||
bank_id: currentBank,
|
||||
type: factType,
|
||||
});
|
||||
console.log("Loaded graph data:", {
|
||||
total_units: graphData.total_units,
|
||||
nodes: graphData.nodes?.length,
|
||||
edges: graphData.edges?.length,
|
||||
table_rows: graphData.table_rows?.length,
|
||||
});
|
||||
setData(graphData);
|
||||
} catch (error) {
|
||||
console.error("Error loading data:", error);
|
||||
@@ -191,10 +185,6 @@ export function DataView({ factType }: DataViewProps) {
|
||||
otherTypes[type] = (otherTypes[type] || 0) + 1;
|
||||
}
|
||||
});
|
||||
console.log("Graph link stats:", { semantic, temporal, entity, causal, total });
|
||||
if (Object.keys(otherTypes).length > 0) {
|
||||
console.log("Other link types:", otherTypes);
|
||||
}
|
||||
return { semantic, temporal, entity, causal, total, otherTypes };
|
||||
}, [graph2DData]);
|
||||
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -3,28 +3,28 @@ LoComo-specific benchmark implementations.
|
||||
|
||||
Provides dataset, answer generator, and evaluator for the LoComo benchmark.
|
||||
"""
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
from benchmarks.common.benchmark_runner import BenchmarkRunner
|
||||
|
||||
import json
|
||||
from datetime import datetime, timezone
|
||||
from typing import List, Dict, Any, Tuple, Optional
|
||||
import asyncio
|
||||
import pydantic
|
||||
from openai import AsyncOpenAI
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
|
||||
from benchmarks.common.benchmark_runner import BenchmarkDataset, LLMAnswerGenerator, LLMAnswerEvaluator
|
||||
import pydantic
|
||||
from hindsight_api.engine.llm_wrapper import LLMConfig
|
||||
from openai import AsyncOpenAI
|
||||
|
||||
from benchmarks.common.benchmark_runner import BenchmarkDataset, BenchmarkRunner, LLMAnswerEvaluator, LLMAnswerGenerator
|
||||
|
||||
|
||||
class LoComoDataset(BenchmarkDataset):
|
||||
"""LoComo dataset implementation."""
|
||||
|
||||
def load(self, path: Path, max_items: Optional[int] = None) -> List[Dict[str, Any]]:
|
||||
"""Load LoComo dataset from JSON file."""
|
||||
with open(path, 'r') as f:
|
||||
with open(path, "r") as f:
|
||||
dataset = json.load(f)
|
||||
|
||||
if max_items:
|
||||
@@ -34,7 +34,7 @@ class LoComoDataset(BenchmarkDataset):
|
||||
|
||||
def get_item_id(self, item: Dict) -> str:
|
||||
"""Get sample ID from LoComo item."""
|
||||
return item['sample_id']
|
||||
return item["sample_id"]
|
||||
|
||||
def prepare_sessions_for_ingestion(self, item: Dict) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
@@ -45,12 +45,12 @@ class LoComoDataset(BenchmarkDataset):
|
||||
Returns:
|
||||
List of session dicts, each containing 'content', 'context', 'event_date', 'document_id'
|
||||
"""
|
||||
conv = item['conversation']
|
||||
speaker_a = conv['speaker_a']
|
||||
speaker_b = conv['speaker_b']
|
||||
conv = item["conversation"]
|
||||
speaker_a = conv["speaker_a"]
|
||||
speaker_b = conv["speaker_b"]
|
||||
|
||||
# Get all session keys sorted
|
||||
session_keys = sorted([k for k in conv.keys() if k.startswith('session_') and not k.endswith('_date_time')])
|
||||
session_keys = sorted([k for k in conv.keys() if k.startswith("session_") and not k.endswith("_date_time")])
|
||||
|
||||
session_items = []
|
||||
|
||||
@@ -65,12 +65,14 @@ class LoComoDataset(BenchmarkDataset):
|
||||
session_date = self._parse_date(conv.get(date_key))
|
||||
session_content = json.dumps(session_data)
|
||||
document_id = f"{item['sample_id']}_{session_key}"
|
||||
session_items.append({
|
||||
"content": session_content,
|
||||
"context": f"Conversation between {speaker_a} and {speaker_b} ({session_key} of {item['sample_id']})",
|
||||
"event_date": session_date,
|
||||
"document_id": document_id
|
||||
})
|
||||
session_items.append(
|
||||
{
|
||||
"content": session_content,
|
||||
"context": f"Conversation between {speaker_a} and {speaker_b} ({session_key} of {item['sample_id']})",
|
||||
"event_date": session_date,
|
||||
"document_id": document_id,
|
||||
}
|
||||
)
|
||||
|
||||
return session_items
|
||||
|
||||
@@ -81,7 +83,7 @@ class LoComoDataset(BenchmarkDataset):
|
||||
Returns:
|
||||
List of QA dicts with 'question', 'answer', 'category'
|
||||
"""
|
||||
return item['qa']
|
||||
return item["qa"]
|
||||
|
||||
def _parse_date(self, date_string: str) -> datetime:
|
||||
"""Parse LoComo date format to datetime."""
|
||||
@@ -95,6 +97,7 @@ class LoComoDataset(BenchmarkDataset):
|
||||
|
||||
class QuestionAnswer(pydantic.BaseModel):
|
||||
"""Answer format for LoComo questions."""
|
||||
|
||||
answer: str
|
||||
reasoning: str
|
||||
|
||||
@@ -113,7 +116,7 @@ class LoComoAnswerGenerator(LLMAnswerGenerator):
|
||||
question: str,
|
||||
recall_result: Dict[str, Any],
|
||||
question_date: Optional[datetime] = None,
|
||||
question_type: Optional[str] = None
|
||||
question_type: Optional[str] = None,
|
||||
) -> Tuple[str, str, Optional[List[Dict[str, Any]]]]:
|
||||
"""
|
||||
Generate answer from retrieved memories using Groq.
|
||||
@@ -141,7 +144,7 @@ class LoComoAnswerGenerator(LLMAnswerGenerator):
|
||||
messages=[
|
||||
{
|
||||
"role": "system",
|
||||
"content": "You are a helpful expert assistant answering questions from lme_experiment users based on the provided context."
|
||||
"content": "You are a helpful expert assistant answering questions from lme_experiment users based on the provided context.",
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
@@ -165,11 +168,11 @@ Context:
|
||||
Question: {question}
|
||||
Answer:
|
||||
|
||||
"""
|
||||
}
|
||||
""",
|
||||
},
|
||||
],
|
||||
response_format=QuestionAnswer,
|
||||
scope="memory"
|
||||
scope="memory",
|
||||
)
|
||||
return answer_obj.answer, answer_obj.reasoning, None
|
||||
except Exception as e:
|
||||
@@ -183,7 +186,7 @@ class LoComoThinkAnswerGenerator(LLMAnswerGenerator):
|
||||
so it doesn't need external search to be performed by the benchmark runner.
|
||||
"""
|
||||
|
||||
def __init__(self, memory: 'MemoryEngine', agent_id: str, thinking_budget: int = 500):
|
||||
def __init__(self, memory: "MemoryEngine", agent_id: str, thinking_budget: int = 500):
|
||||
"""Initialize with memory instance and agent_id.
|
||||
|
||||
Args:
|
||||
@@ -204,7 +207,7 @@ class LoComoThinkAnswerGenerator(LLMAnswerGenerator):
|
||||
question: str,
|
||||
recall_result: Dict[str, Any],
|
||||
question_date: Optional[datetime] = None,
|
||||
question_type: Optional[str] = None
|
||||
question_type: Optional[str] = None,
|
||||
) -> Tuple[str, str, Optional[List[Dict[str, Any]]]]:
|
||||
"""
|
||||
Generate answer using the integrated think API.
|
||||
@@ -235,9 +238,9 @@ class LoComoThinkAnswerGenerator(LLMAnswerGenerator):
|
||||
|
||||
# Extract memories from based_on
|
||||
based_on = result.based_on
|
||||
world_facts = based_on.get('world', [])
|
||||
agent_facts = based_on.get('agent', [])
|
||||
opinion_facts = based_on.get('opinion', [])
|
||||
world_facts = based_on.get("world", [])
|
||||
agent_facts = based_on.get("agent", [])
|
||||
opinion_facts = based_on.get("opinion", [])
|
||||
|
||||
# Combine all facts into retrieved_memories
|
||||
retrieved_memories = []
|
||||
@@ -271,7 +274,7 @@ async def run_benchmark(
|
||||
api_url: str = None,
|
||||
max_concurrent_questions_override: int = None,
|
||||
only_failed: bool = False,
|
||||
only_invalid: bool = False
|
||||
only_invalid: bool = False,
|
||||
):
|
||||
"""
|
||||
Run the LoComo benchmark.
|
||||
@@ -287,6 +290,7 @@ async def run_benchmark(
|
||||
only_invalid: If True, only run conversations that have invalid questions (is_invalid=True)
|
||||
"""
|
||||
from rich.console import Console
|
||||
|
||||
console = Console()
|
||||
|
||||
# Load previous results if filtering for failed/invalid conversations
|
||||
@@ -294,35 +298,41 @@ async def run_benchmark(
|
||||
invalid_conversation_ids = set()
|
||||
if only_failed or only_invalid:
|
||||
suffix = "_think" if use_think else ""
|
||||
results_filename = f'benchmark_results{suffix}.json'
|
||||
results_path = Path(__file__).parent / 'results' / results_filename
|
||||
results_filename = f"benchmark_results{suffix}.json"
|
||||
results_path = Path(__file__).parent / "results" / results_filename
|
||||
|
||||
if not results_path.exists():
|
||||
console.print(f"[red]Error: Cannot use --only-failed or --only-invalid without existing results file[/red]")
|
||||
console.print("[red]Error: Cannot use --only-failed or --only-invalid without existing results file[/red]")
|
||||
console.print(f"[yellow]Results file not found: {results_path}[/yellow]")
|
||||
return
|
||||
|
||||
with open(results_path, 'r') as f:
|
||||
with open(results_path, "r") as f:
|
||||
previous_results = json.load(f)
|
||||
|
||||
# Extract conversation IDs that have failed or invalid questions
|
||||
for item_result in previous_results.get('item_results', []):
|
||||
item_id = item_result['item_id']
|
||||
for detail in item_result['metrics'].get('detailed_results', []):
|
||||
if only_failed and detail.get('is_correct') == False and not detail.get('is_invalid', False):
|
||||
for item_result in previous_results.get("item_results", []):
|
||||
item_id = item_result["item_id"]
|
||||
for detail in item_result["metrics"].get("detailed_results", []):
|
||||
if only_failed and detail.get("is_correct") == False and not detail.get("is_invalid", False):
|
||||
failed_conversation_ids.add(item_id)
|
||||
if only_invalid and detail.get('is_invalid', False):
|
||||
if only_invalid and detail.get("is_invalid", False):
|
||||
invalid_conversation_ids.add(item_id)
|
||||
|
||||
if only_failed:
|
||||
console.print(f"[cyan]Filtering to {len(failed_conversation_ids)} conversations with failed questions (is_correct=False)[/cyan]")
|
||||
console.print(
|
||||
f"[cyan]Filtering to {len(failed_conversation_ids)} conversations with failed questions (is_correct=False)[/cyan]"
|
||||
)
|
||||
if only_invalid:
|
||||
console.print(f"[cyan]Filtering to {len(invalid_conversation_ids)} conversations with invalid questions (is_invalid=True)[/cyan]")
|
||||
console.print(
|
||||
f"[cyan]Filtering to {len(invalid_conversation_ids)} conversations with invalid questions (is_invalid=True)[/cyan]"
|
||||
)
|
||||
|
||||
target_ids = failed_conversation_ids if only_failed else invalid_conversation_ids
|
||||
if not target_ids:
|
||||
filter_type = "failed" if only_failed else "invalid"
|
||||
console.print(f"[yellow]No conversations with {filter_type} questions found in previous results. Nothing to run.[/yellow]")
|
||||
console.print(
|
||||
f"[yellow]No conversations with {filter_type} questions found in previous results. Nothing to run.[/yellow]"
|
||||
)
|
||||
return
|
||||
|
||||
# Initialize components
|
||||
@@ -331,18 +341,16 @@ async def run_benchmark(
|
||||
# Use remote API client if api_url is provided, otherwise use local memory
|
||||
if api_url:
|
||||
from benchmarks.common.benchmark_runner import HindsightClientAdapter
|
||||
|
||||
memory = HindsightClientAdapter(base_url=api_url)
|
||||
await memory.initialize()
|
||||
else:
|
||||
from benchmarks.common.benchmark_runner import create_memory_engine
|
||||
|
||||
memory = await create_memory_engine()
|
||||
|
||||
if use_think:
|
||||
answer_generator = LoComoThinkAnswerGenerator(
|
||||
memory=memory,
|
||||
agent_id="locomo",
|
||||
thinking_budget=500
|
||||
)
|
||||
answer_generator = LoComoThinkAnswerGenerator(memory=memory, agent_id="locomo", thinking_budget=500)
|
||||
max_concurrent_questions = max_concurrent_questions_override or 4
|
||||
eval_semaphore_size = 4
|
||||
else:
|
||||
@@ -356,14 +364,11 @@ async def run_benchmark(
|
||||
|
||||
# Create benchmark runner
|
||||
runner = BenchmarkRunner(
|
||||
dataset=dataset,
|
||||
answer_generator=answer_generator,
|
||||
answer_evaluator=answer_evaluator,
|
||||
memory=memory
|
||||
dataset=dataset, answer_generator=answer_generator, answer_evaluator=answer_evaluator, memory=memory
|
||||
)
|
||||
|
||||
# Filter dataset if using --only-failed or --only-invalid
|
||||
dataset_path = Path(__file__).parent / 'datasets' / 'locomo10.json'
|
||||
dataset_path = Path(__file__).parent / "datasets" / "locomo10.json"
|
||||
|
||||
if only_failed or only_invalid:
|
||||
# Load and filter dataset
|
||||
@@ -374,14 +379,16 @@ async def run_benchmark(
|
||||
|
||||
# Temporarily replace dataset's load method
|
||||
original_load = dataset.load
|
||||
|
||||
def filtered_load(path: Path, max_items: Optional[int] = None):
|
||||
return filtered_items[:max_items] if max_items else filtered_items
|
||||
|
||||
dataset.load = filtered_load
|
||||
|
||||
# Determine output filename based on mode
|
||||
suffix = "_think" if use_think else ""
|
||||
results_filename = f'benchmark_results{suffix}.json'
|
||||
output_path = Path(__file__).parent / 'results' / results_filename
|
||||
results_filename = f"benchmark_results{suffix}.json"
|
||||
output_path = Path(__file__).parent / "results" / results_filename
|
||||
|
||||
# Create results directory if it doesn't exist
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
@@ -406,7 +413,7 @@ async def run_benchmark(
|
||||
clear_agent_per_item=True, # Use unique agent ID per conversation
|
||||
max_concurrent_items=3, # Process up to 3 conversations in parallel
|
||||
output_path=output_path, # Save results incrementally
|
||||
merge_with_existing=merge_with_existing
|
||||
merge_with_existing=merge_with_existing,
|
||||
)
|
||||
|
||||
# Display results (final save already happened incrementally)
|
||||
@@ -430,14 +437,10 @@ def generate_markdown_table(results: dict, use_think: bool = False):
|
||||
4 = Open-domain
|
||||
"""
|
||||
from rich.console import Console
|
||||
|
||||
console = Console()
|
||||
|
||||
category_names = {
|
||||
'1': 'Multi-hop',
|
||||
'2': 'Single-hop',
|
||||
'3': 'Temporal',
|
||||
'4': 'Open-domain'
|
||||
}
|
||||
category_names = {"1": "Multi-hop", "2": "Single-hop", "3": "Temporal", "4": "Open-domain"}
|
||||
|
||||
# Build markdown content
|
||||
lines = []
|
||||
@@ -446,33 +449,41 @@ def generate_markdown_table(results: dict, use_think: bool = False):
|
||||
lines.append("")
|
||||
|
||||
# Add model configuration
|
||||
if 'model_config' in results:
|
||||
config = results['model_config']
|
||||
if "model_config" in results:
|
||||
config = results["model_config"]
|
||||
lines.append("## Model Configuration")
|
||||
lines.append("")
|
||||
lines.append(f"- **Hindsight**: {config['hindsight']['provider']}/{config['hindsight']['model']}")
|
||||
lines.append(f"- **Answer Generation**: {config['answer_generation']['provider']}/{config['answer_generation']['model']}")
|
||||
lines.append(
|
||||
f"- **Answer Generation**: {config['answer_generation']['provider']}/{config['answer_generation']['model']}"
|
||||
)
|
||||
lines.append(f"- **LLM Judge**: {config['judge']['provider']}/{config['judge']['model']}")
|
||||
lines.append("")
|
||||
|
||||
lines.append(f"**Overall Accuracy**: {results['overall_accuracy']:.2f}% ({results['total_correct']}/{results['total_questions']})")
|
||||
lines.append(
|
||||
f"**Overall Accuracy**: {results['overall_accuracy']:.2f}% ({results['total_correct']}/{results['total_questions']})"
|
||||
)
|
||||
lines.append("")
|
||||
lines.append("| Sample ID | Sessions | Questions | Correct | Accuracy | Multi-hop | Single-hop | Temporal | Open-domain |")
|
||||
lines.append("|-----------|----------|-----------|---------|----------|-----------|------------|----------|-------------|")
|
||||
lines.append(
|
||||
"| Sample ID | Sessions | Questions | Correct | Accuracy | Multi-hop | Single-hop | Temporal | Open-domain |"
|
||||
)
|
||||
lines.append(
|
||||
"|-----------|----------|-----------|---------|----------|-----------|------------|----------|-------------|"
|
||||
)
|
||||
|
||||
for item_result in results['item_results']:
|
||||
item_id = item_result['item_id']
|
||||
num_sessions = item_result['num_sessions']
|
||||
metrics = item_result['metrics']
|
||||
for item_result in results["item_results"]:
|
||||
item_id = item_result["item_id"]
|
||||
num_sessions = item_result["num_sessions"]
|
||||
metrics = item_result["metrics"]
|
||||
|
||||
# Calculate category accuracies
|
||||
cat_stats = metrics.get('category_stats', {})
|
||||
cat_stats = metrics.get("category_stats", {})
|
||||
cat_accuracies = {}
|
||||
|
||||
for cat_id in ['1', '2', '3', '4']:
|
||||
for cat_id in ["1", "2", "3", "4"]:
|
||||
if cat_id in cat_stats:
|
||||
stats = cat_stats[cat_id]
|
||||
acc = (stats['correct'] / stats['total'] * 100) if stats['total'] > 0 else 0
|
||||
acc = (stats["correct"] / stats["total"] * 100) if stats["total"] > 0 else 0
|
||||
cat_accuracies[cat_id] = f"{acc:.1f}% ({stats['correct']}/{stats['total']})"
|
||||
else:
|
||||
cat_accuracies[cat_id] = "N/A"
|
||||
@@ -485,28 +496,48 @@ def generate_markdown_table(results: dict, use_think: bool = False):
|
||||
|
||||
# Write to file with suffix
|
||||
suffix = "_think" if use_think else ""
|
||||
output_file = Path(__file__).parent / 'results' / f'results_table{suffix}.md'
|
||||
output_file = Path(__file__).parent / "results" / f"results_table{suffix}.md"
|
||||
output_file.parent.mkdir(parents=True, exist_ok=True)
|
||||
output_file.write_text('\n'.join(lines))
|
||||
output_file.write_text("\n".join(lines))
|
||||
console.print(f"\n[green]✓[/green] Results table saved to {output_file}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import logging
|
||||
import argparse
|
||||
import logging
|
||||
|
||||
logging.basicConfig(level=logging.INFO, format='%(asctime)s %(levelname)s %(message)s')
|
||||
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
|
||||
|
||||
parser = argparse.ArgumentParser(description='Run LoComo benchmark')
|
||||
parser.add_argument('--max-conversations', type=int, default=None, help='Maximum conversations to evaluate')
|
||||
parser.add_argument('--max-questions', type=int, default=None, help='Maximum questions per conversation')
|
||||
parser.add_argument('--skip-ingestion', action='store_true', help='Skip ingestion and use existing data')
|
||||
parser.add_argument('--use-think', action='store_true', help='Use think API instead of search + LLM')
|
||||
parser.add_argument('--conversation', type=str, default=None, help='Run only specific conversation (e.g., "conv-26")')
|
||||
parser.add_argument('--api-url', type=str, default=None, help='Hindsight API URL (default: use local memory, example: http://localhost:8888)')
|
||||
parser.add_argument('--max-concurrent-questions', type=int, default=None, help='Max concurrent questions per conversation (default: 4 for think, 10 for search)')
|
||||
parser.add_argument('--only-failed', action='store_true', help='Only run conversations that have failed questions (is_correct=False). Requires existing results file.')
|
||||
parser.add_argument('--only-invalid', action='store_true', help='Only run conversations that have invalid questions (is_invalid=True). Requires existing results file.')
|
||||
parser = argparse.ArgumentParser(description="Run LoComo benchmark")
|
||||
parser.add_argument("--max-conversations", type=int, default=None, help="Maximum conversations to evaluate")
|
||||
parser.add_argument("--max-questions", type=int, default=None, help="Maximum questions per conversation")
|
||||
parser.add_argument("--skip-ingestion", action="store_true", help="Skip ingestion and use existing data")
|
||||
parser.add_argument("--use-think", action="store_true", help="Use think API instead of search + LLM")
|
||||
parser.add_argument(
|
||||
"--conversation", type=str, default=None, help='Run only specific conversation (e.g., "conv-26")'
|
||||
)
|
||||
parser.add_argument(
|
||||
"--api-url",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Hindsight API URL (default: use local memory, example: http://localhost:8888)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-concurrent-questions",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Max concurrent questions per conversation (default: 4 for think, 10 for search)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--only-failed",
|
||||
action="store_true",
|
||||
help="Only run conversations that have failed questions (is_correct=False). Requires existing results file.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--only-invalid",
|
||||
action="store_true",
|
||||
help="Only run conversations that have invalid questions (is_invalid=True). Requires existing results file.",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
@@ -514,14 +545,16 @@ if __name__ == "__main__":
|
||||
if args.only_failed and args.only_invalid:
|
||||
parser.error("Cannot use both --only-failed and --only-invalid at the same time")
|
||||
|
||||
results = asyncio.run(run_benchmark(
|
||||
max_conversations=args.max_conversations,
|
||||
max_questions_per_conv=args.max_questions,
|
||||
skip_ingestion=args.skip_ingestion,
|
||||
use_think=args.use_think,
|
||||
conversation=args.conversation,
|
||||
api_url=args.api_url,
|
||||
max_concurrent_questions_override=args.max_concurrent_questions,
|
||||
only_failed=args.only_failed,
|
||||
only_invalid=args.only_invalid
|
||||
))
|
||||
results = asyncio.run(
|
||||
run_benchmark(
|
||||
max_conversations=args.max_conversations,
|
||||
max_questions_per_conv=args.max_questions,
|
||||
skip_ingestion=args.skip_ingestion,
|
||||
use_think=args.use_think,
|
||||
conversation=args.conversation,
|
||||
api_url=args.api_url,
|
||||
max_concurrent_questions_override=args.max_concurrent_questions,
|
||||
only_failed=args.only_failed,
|
||||
only_invalid=args.only_invalid,
|
||||
)
|
||||
)
|
||||
|
||||
@@ -3,21 +3,20 @@ LongMemEval-specific benchmark implementations.
|
||||
|
||||
Provides dataset, answer generator, and evaluator for the LongMemEval benchmark.
|
||||
"""
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
from benchmarks.common.benchmark_runner import BenchmarkRunner
|
||||
|
||||
import json
|
||||
from datetime import datetime, timezone
|
||||
from typing import List, Dict, Any, Tuple, Optional
|
||||
import asyncio
|
||||
import pydantic
|
||||
from openai import AsyncOpenAI
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
|
||||
from benchmarks.common.benchmark_runner import BenchmarkDataset, LLMAnswerGenerator, LLMAnswerEvaluator
|
||||
import pydantic
|
||||
from hindsight_api.engine.llm_wrapper import LLMConfig
|
||||
from openai import AsyncOpenAI
|
||||
|
||||
from benchmarks.common.benchmark_runner import BenchmarkDataset, BenchmarkRunner, LLMAnswerEvaluator, LLMAnswerGenerator
|
||||
|
||||
|
||||
class LongMemEvalDataset(BenchmarkDataset):
|
||||
@@ -25,7 +24,7 @@ class LongMemEvalDataset(BenchmarkDataset):
|
||||
|
||||
def load(self, path: Path, max_items: Optional[int] = None) -> List[Dict[str, Any]]:
|
||||
"""Load LongMemEval dataset from JSON file."""
|
||||
with open(path, 'r') as f:
|
||||
with open(path, "r") as f:
|
||||
dataset = json.load(f)
|
||||
|
||||
if max_items:
|
||||
@@ -67,7 +66,7 @@ class LongMemEvalDataset(BenchmarkDataset):
|
||||
for turn in session_turns:
|
||||
if isinstance(turn, dict):
|
||||
# Create a copy without has_answer
|
||||
cleaned_turn = {k: v for k, v in turn.items() if k != 'has_answer'}
|
||||
cleaned_turn = {k: v for k, v in turn.items() if k != "has_answer"}
|
||||
cleaned_turns.append(cleaned_turn)
|
||||
else:
|
||||
cleaned_turns.append(turn)
|
||||
@@ -75,12 +74,14 @@ class LongMemEvalDataset(BenchmarkDataset):
|
||||
session_content = json.dumps(cleaned_turns)
|
||||
question_id = item.get("question_id", "unknown")
|
||||
document_id = f"{question_id}_{session_id}"
|
||||
batch_contents.append({
|
||||
"content": session_content,
|
||||
"context": f"Session {document_id} - you are the assistant in this conversation - happened on {session_date.strftime('%Y-%m-%d %H:%M:%S')} UTC.",
|
||||
"event_date": session_date,
|
||||
"document_id": document_id
|
||||
})
|
||||
batch_contents.append(
|
||||
{
|
||||
"content": session_content,
|
||||
"context": f"Session {document_id} - you are the assistant in this conversation - happened on {session_date.strftime('%Y-%m-%d %H:%M:%S')} UTC.",
|
||||
"event_date": session_date,
|
||||
"document_id": document_id,
|
||||
}
|
||||
)
|
||||
|
||||
return batch_contents
|
||||
|
||||
@@ -95,22 +96,24 @@ class LongMemEvalDataset(BenchmarkDataset):
|
||||
"""
|
||||
# Parse question_date if available
|
||||
question_date = None
|
||||
if 'question_date' in item:
|
||||
question_date = self._parse_date(item['question_date'])
|
||||
if "question_date" in item:
|
||||
question_date = self._parse_date(item["question_date"])
|
||||
|
||||
return [{
|
||||
'question': item.get("question", ""),
|
||||
'answer': item.get("answer", ""),
|
||||
'category': item.get("question_type", "unknown"),
|
||||
'question_date': question_date
|
||||
}]
|
||||
return [
|
||||
{
|
||||
"question": item.get("question", ""),
|
||||
"answer": item.get("answer", ""),
|
||||
"category": item.get("question_type", "unknown"),
|
||||
"question_date": question_date,
|
||||
}
|
||||
]
|
||||
|
||||
def _parse_date(self, date_str: str) -> datetime:
|
||||
"""Parse date string to datetime object."""
|
||||
try:
|
||||
# LongMemEval format: "2023/05/20 (Sat) 02:21"
|
||||
# Try to parse the main part before the day name
|
||||
date_str_cleaned = date_str.split('(')[0].strip() if '(' in date_str else date_str
|
||||
date_str_cleaned = date_str.split("(")[0].strip() if "(" in date_str else date_str
|
||||
|
||||
# Try multiple formats
|
||||
for fmt in ["%Y/%m/%d %H:%M", "%Y-%m-%d %H:%M:%S", "%Y-%m-%d", "%Y/%m/%d"]:
|
||||
@@ -121,7 +124,7 @@ class LongMemEvalDataset(BenchmarkDataset):
|
||||
continue
|
||||
|
||||
# Fallback: try ISO format
|
||||
return datetime.fromisoformat(date_str.replace('Z', '+00:00'))
|
||||
return datetime.fromisoformat(date_str.replace("Z", "+00:00"))
|
||||
except Exception:
|
||||
raise ValueError(f"Failed to parse date string: {date_str}")
|
||||
|
||||
@@ -130,6 +133,7 @@ class QuestionAnswer(pydantic.BaseModel):
|
||||
answer: str
|
||||
reasoning: Optional[str] = None
|
||||
|
||||
|
||||
class LongMemEvalAnswerGenerator(LLMAnswerGenerator):
|
||||
"""LongMemEval-specific answer generator using configurable LLM provider."""
|
||||
|
||||
@@ -202,10 +206,7 @@ class LongMemEvalAnswerGenerator(LLMAnswerGenerator):
|
||||
chunk_text = chunk_info.get("chunk_text", "")
|
||||
|
||||
# Build the formatted fact entry
|
||||
entry_parts = [
|
||||
f"Fact {i} ({fact_type}): {fact_text}",
|
||||
f"When: {when_str}"
|
||||
]
|
||||
entry_parts = [f"Fact {i} ({fact_type}): {fact_text}", f"When: {when_str}"]
|
||||
|
||||
# Add context field if present
|
||||
context = fact.get("context")
|
||||
@@ -217,7 +218,7 @@ class LongMemEvalAnswerGenerator(LLMAnswerGenerator):
|
||||
# Truncate very long chunks
|
||||
if len(chunk_text) > 1000:
|
||||
chunk_text = chunk_text[:1000] + "..."
|
||||
entry_parts.append(f"Source chunk:\n \"{chunk_text}\"")
|
||||
entry_parts.append(f'Source chunk:\n "{chunk_text}"')
|
||||
|
||||
formatted_parts.append("\n".join(entry_parts))
|
||||
|
||||
@@ -326,43 +327,43 @@ The context contains memory facts extracted from previous conversations, each wi
|
||||
return ""
|
||||
|
||||
async def generate_answer(
|
||||
self,
|
||||
question: str,
|
||||
recall_result: Dict[str, Any],
|
||||
question_date: Optional[datetime] = None,
|
||||
question_type: Optional[str] = None
|
||||
) -> Tuple[str, str, Optional[List[Dict[str, Any]]]]:
|
||||
"""
|
||||
Generate answer from retrieved memories using Groq.
|
||||
self,
|
||||
question: str,
|
||||
recall_result: Dict[str, Any],
|
||||
question_date: Optional[datetime] = None,
|
||||
question_type: Optional[str] = None,
|
||||
) -> Tuple[str, str, Optional[List[Dict[str, Any]]]]:
|
||||
"""
|
||||
Generate answer from retrieved memories using Groq.
|
||||
|
||||
Args:
|
||||
question: The question text
|
||||
recall_result: Full RecallResult dict containing results, entities, chunks, and trace
|
||||
question_date: Date when the question was asked (for temporal context)
|
||||
question_type: Question category (e.g., 'single-session-user', 'multi-session-assistant')
|
||||
Args:
|
||||
question: The question text
|
||||
recall_result: Full RecallResult dict containing results, entities, chunks, and trace
|
||||
question_date: Date when the question was asked (for temporal context)
|
||||
question_type: Question category (e.g., 'single-session-user', 'multi-session-assistant')
|
||||
|
||||
Returns:
|
||||
Tuple of (answer, reasoning, None)
|
||||
- None indicates to use the memories from recall_result
|
||||
"""
|
||||
# Format context based on selected mode
|
||||
if self.context_format == "structured":
|
||||
context = self._format_context_structured(recall_result)
|
||||
else:
|
||||
context = self._format_context_json(recall_result)
|
||||
Returns:
|
||||
Tuple of (answer, reasoning, None)
|
||||
- None indicates to use the memories from recall_result
|
||||
"""
|
||||
# Format context based on selected mode
|
||||
if self.context_format == "structured":
|
||||
context = self._format_context_structured(recall_result)
|
||||
else:
|
||||
context = self._format_context_json(recall_result)
|
||||
|
||||
context_instructions = self._get_context_instructions()
|
||||
context_instructions = self._get_context_instructions()
|
||||
|
||||
# Format question date if provided
|
||||
formatted_question_date = question_date.strftime('%Y-%m-%d %H:%M:%S UTC') if question_date else "Not specified"
|
||||
# Format question date if provided
|
||||
formatted_question_date = question_date.strftime("%Y-%m-%d %H:%M:%S UTC") if question_date else "Not specified"
|
||||
|
||||
# Use LLM to generate answer
|
||||
try:
|
||||
answer_obj = await self.llm_config.call(
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": f"""You are a helpful assistant that must answer user questions based on the previous conversations.
|
||||
# Use LLM to generate answer
|
||||
try:
|
||||
answer_obj = await self.llm_config.call(
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": f"""You are a helpful assistant that must answer user questions based on the previous conversations.
|
||||
|
||||
{context_instructions}**Answer Guidelines:**
|
||||
1. Start by scanning retrieved context to understand the facts and events that happened and the timeline.
|
||||
@@ -393,20 +394,20 @@ Retrieved Context:
|
||||
|
||||
|
||||
Answer:
|
||||
"""
|
||||
}
|
||||
],
|
||||
response_format=QuestionAnswer,
|
||||
scope="memory",
|
||||
max_completion_tokens=32768,
|
||||
)
|
||||
reasoning_text = answer_obj.reasoning or ""
|
||||
if reasoning_text:
|
||||
reasoning_text = reasoning_text + " "
|
||||
reasoning_text += f"(question date: {formatted_question_date})"
|
||||
return answer_obj.answer, reasoning_text, None
|
||||
except Exception as e:
|
||||
return f"Error generating answer: {str(e)}", "Error occurred during answer generation.", None
|
||||
""",
|
||||
}
|
||||
],
|
||||
response_format=QuestionAnswer,
|
||||
scope="memory",
|
||||
max_completion_tokens=32768,
|
||||
)
|
||||
reasoning_text = answer_obj.reasoning or ""
|
||||
if reasoning_text:
|
||||
reasoning_text = reasoning_text + " "
|
||||
reasoning_text += f"(question date: {formatted_question_date})"
|
||||
return answer_obj.answer, reasoning_text, None
|
||||
except Exception as e:
|
||||
return f"Error generating answer: {str(e)}", "Error occurred during answer generation.", None
|
||||
|
||||
|
||||
async def run_benchmark(
|
||||
@@ -425,7 +426,7 @@ async def run_benchmark(
|
||||
max_concurrent_items: int = 1,
|
||||
results_filename: str = "benchmark_results.json",
|
||||
context_format: str = "json",
|
||||
source_results: str = None
|
||||
source_results: str = None,
|
||||
):
|
||||
"""
|
||||
Run the LongMemEval benchmark.
|
||||
@@ -449,13 +450,16 @@ async def run_benchmark(
|
||||
source_results: Source results file to read failed/invalid questions from (for --only-failed/--only-invalid). Defaults to benchmark_results.json.
|
||||
"""
|
||||
from rich.console import Console
|
||||
|
||||
console = Console()
|
||||
|
||||
# Validate mutually exclusive arguments
|
||||
# --max-instances-per-category can't be combined with --max-instances or --category
|
||||
# But --category CAN be combined with --max-instances (to limit questions within a category)
|
||||
if max_instances_per_category is not None and (max_instances is not None or category is not None):
|
||||
console.print("[red]Error: --max-questions-per-category cannot be combined with --max-instances or --category[/red]")
|
||||
console.print(
|
||||
"[red]Error: --max-questions-per-category cannot be combined with --max-instances or --category[/red]"
|
||||
)
|
||||
return
|
||||
|
||||
# Validate --only-ingested can't be combined with other dataset filters
|
||||
@@ -480,8 +484,10 @@ async def run_benchmark(
|
||||
dataset_path = Path(__file__).parent / "datasets" / "longmemeval_s_cleaned.json"
|
||||
if not dataset_path.exists():
|
||||
if not download_dataset(dataset_path):
|
||||
console.print(f"[red]Failed to download dataset. Please download manually:[/red]")
|
||||
console.print("[yellow]curl -L 'https://huggingface.co/datasets/xiaowu0162/longmemeval-cleaned/resolve/main/longmemeval_s_cleaned.json' -o benchmarks/longmemeval/datasets/longmemeval_s_cleaned.json[/yellow]")
|
||||
console.print("[red]Failed to download dataset. Please download manually:[/red]")
|
||||
console.print(
|
||||
"[yellow]curl -L 'https://huggingface.co/datasets/xiaowu0162/longmemeval-cleaned/resolve/main/longmemeval_s_cleaned.json' -o benchmarks/longmemeval/datasets/longmemeval_s_cleaned.json[/yellow]"
|
||||
)
|
||||
return
|
||||
|
||||
# Initialize components
|
||||
@@ -499,9 +505,10 @@ async def run_benchmark(
|
||||
|
||||
# Group by category and take max_instances_per_category from each
|
||||
from collections import defaultdict
|
||||
|
||||
category_items = defaultdict(list)
|
||||
for item in original_dataset_items:
|
||||
cat = item.get('question_type', 'unknown')
|
||||
cat = item.get("question_type", "unknown")
|
||||
category_items[cat].append(item)
|
||||
|
||||
# Take up to max_instances_per_category from each category
|
||||
@@ -518,30 +525,34 @@ async def run_benchmark(
|
||||
invalid_question_ids = set()
|
||||
if only_failed or only_invalid:
|
||||
# Use source_results if specified, otherwise default to benchmark_results.json
|
||||
source_file = source_results if source_results else 'benchmark_results.json'
|
||||
results_path = Path(__file__).parent / 'results' / source_file
|
||||
source_file = source_results if source_results else "benchmark_results.json"
|
||||
results_path = Path(__file__).parent / "results" / source_file
|
||||
if not results_path.exists():
|
||||
console.print(f"[red]Error: Cannot use --only-failed or --only-invalid without existing results file[/red]")
|
||||
console.print("[red]Error: Cannot use --only-failed or --only-invalid without existing results file[/red]")
|
||||
console.print(f"[yellow]Results file not found: {results_path}[/yellow]")
|
||||
return
|
||||
|
||||
console.print(f"[cyan]Reading failed/invalid questions from: {source_file}[/cyan]")
|
||||
with open(results_path, 'r') as f:
|
||||
with open(results_path, "r") as f:
|
||||
previous_results = json.load(f)
|
||||
|
||||
# Extract question IDs that failed or are invalid
|
||||
for item_result in previous_results.get('item_results', []):
|
||||
item_id = item_result['item_id']
|
||||
for detail in item_result['metrics'].get('detailed_results', []):
|
||||
if only_failed and detail.get('is_correct') == False and not detail.get('is_invalid', False):
|
||||
for item_result in previous_results.get("item_results", []):
|
||||
item_id = item_result["item_id"]
|
||||
for detail in item_result["metrics"].get("detailed_results", []):
|
||||
if only_failed and detail.get("is_correct") == False and not detail.get("is_invalid", False):
|
||||
failed_question_ids.add(item_id)
|
||||
if only_invalid and detail.get('is_invalid', False):
|
||||
if only_invalid and detail.get("is_invalid", False):
|
||||
invalid_question_ids.add(item_id)
|
||||
|
||||
if only_failed:
|
||||
console.print(f"[cyan]Filtering to {len(failed_question_ids)} questions that failed (is_correct=False)[/cyan]")
|
||||
console.print(
|
||||
f"[cyan]Filtering to {len(failed_question_ids)} questions that failed (is_correct=False)[/cyan]"
|
||||
)
|
||||
if only_invalid:
|
||||
console.print(f"[cyan]Filtering to {len(invalid_question_ids)} questions that were invalid (is_invalid=True)[/cyan]")
|
||||
console.print(
|
||||
f"[cyan]Filtering to {len(invalid_question_ids)} questions that were invalid (is_invalid=True)[/cyan]"
|
||||
)
|
||||
|
||||
# Filter dataset by category if specified
|
||||
if category:
|
||||
@@ -550,11 +561,11 @@ async def run_benchmark(
|
||||
# Load full dataset without max_instances limit for filtering
|
||||
original_dataset_items = dataset.load(dataset_path, max_items=None)
|
||||
|
||||
filtered_items = [item for item in original_dataset_items if item.get('question_type') == category]
|
||||
filtered_items = [item for item in original_dataset_items if item.get("question_type") == category]
|
||||
|
||||
if not filtered_items:
|
||||
console.print(f"[yellow]No questions found for category '{category}'. Available categories:[/yellow]")
|
||||
available_categories = set(item.get('question_type', 'unknown') for item in original_dataset_items)
|
||||
available_categories = set(item.get("question_type", "unknown") for item in original_dataset_items)
|
||||
for cat in sorted(available_categories):
|
||||
console.print(f" - {cat}")
|
||||
return
|
||||
@@ -562,7 +573,9 @@ async def run_benchmark(
|
||||
total_found = len(filtered_items)
|
||||
will_run = min(total_found, max_instances) if max_instances else total_found
|
||||
if max_instances and total_found > max_instances:
|
||||
console.print(f"[green]Found {total_found} questions for category '{category}' (will run {will_run} due to --max-instances)[/green]")
|
||||
console.print(
|
||||
f"[green]Found {total_found} questions for category '{category}' (will run {will_run} due to --max-instances)[/green]"
|
||||
)
|
||||
else:
|
||||
console.print(f"[green]Found {total_found} questions for category '{category}'[/green]")
|
||||
|
||||
@@ -588,7 +601,9 @@ async def run_benchmark(
|
||||
total_found = len(filtered_items)
|
||||
will_run = min(total_found, max_instances) if max_instances else total_found
|
||||
if max_instances and total_found > max_instances:
|
||||
console.print(f"[green]Found {total_found} {filter_type} items to re-evaluate (will run {will_run} due to --max-instances)[/green]")
|
||||
console.print(
|
||||
f"[green]Found {total_found} {filter_type} items to re-evaluate (will run {will_run} due to --max-instances)[/green]"
|
||||
)
|
||||
else:
|
||||
console.print(f"[green]Found {total_found} {filter_type} items to re-evaluate[/green]")
|
||||
|
||||
@@ -600,6 +615,7 @@ async def run_benchmark(
|
||||
|
||||
# Create local memory engine
|
||||
from benchmarks.common.benchmark_runner import create_memory_engine
|
||||
|
||||
memory = await create_memory_engine()
|
||||
|
||||
# Filter by only_ingested: only run items whose memory bank already exists
|
||||
@@ -623,10 +639,9 @@ async def run_benchmark(
|
||||
# Check if bank has any memory units
|
||||
async with pool.acquire() as conn:
|
||||
result = await conn.fetchrow(
|
||||
"SELECT COUNT(*) as count FROM memory_units WHERE bank_id = $1 LIMIT 1",
|
||||
agent_id
|
||||
"SELECT COUNT(*) as count FROM memory_units WHERE bank_id = $1 LIMIT 1", agent_id
|
||||
)
|
||||
if result['count'] > 0:
|
||||
if result["count"] > 0:
|
||||
ingested_items.append(item)
|
||||
|
||||
filtered_items = ingested_items
|
||||
@@ -638,34 +653,43 @@ async def run_benchmark(
|
||||
|
||||
# Create benchmark runner
|
||||
runner = BenchmarkRunner(
|
||||
dataset=dataset,
|
||||
answer_generator=answer_generator,
|
||||
answer_evaluator=answer_evaluator,
|
||||
memory=memory
|
||||
dataset=dataset, answer_generator=answer_generator, answer_evaluator=answer_evaluator, memory=memory
|
||||
)
|
||||
|
||||
# If filtering by category, failed, invalid, only_ingested, or max_instances_per_category, we need to use a custom dataset that only returns those items
|
||||
# We'll temporarily replace the dataset's load method
|
||||
if filtered_items is not None:
|
||||
original_load = dataset.load
|
||||
|
||||
def filtered_load(path: Path, max_items: Optional[int] = None):
|
||||
return filtered_items[:max_items] if max_items else filtered_items
|
||||
|
||||
dataset.load = filtered_load
|
||||
|
||||
# Run benchmark
|
||||
# Single-phase approach: each question gets its own isolated agent_id
|
||||
# This ensures each question only has access to its own context
|
||||
output_path = Path(__file__).parent / 'results' / results_filename
|
||||
output_path = Path(__file__).parent / "results" / results_filename
|
||||
|
||||
# Create results directory if it doesn't exist
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
merge_with_existing = (filln or question_id is not None or only_failed or only_invalid or only_ingested or category is not None or max_instances_per_category is not None)
|
||||
merge_with_existing = (
|
||||
filln
|
||||
or question_id is not None
|
||||
or only_failed
|
||||
or only_invalid
|
||||
or only_ingested
|
||||
or category is not None
|
||||
or max_instances_per_category is not None
|
||||
)
|
||||
|
||||
results = await runner.run(
|
||||
dataset_path=dataset_path,
|
||||
agent_id="longmemeval", # Will be suffixed with question_id per item
|
||||
max_items=max_instances if not max_instances_per_category else None, # Don't apply max_items when using per-category limit
|
||||
max_items=max_instances
|
||||
if not max_instances_per_category
|
||||
else None, # Don't apply max_items when using per-category limit
|
||||
max_questions_per_item=max_questions_per_instance,
|
||||
thinking_budget=thinking_budget,
|
||||
max_tokens=max_tokens,
|
||||
@@ -678,7 +702,7 @@ async def run_benchmark(
|
||||
specific_item=question_id, # Optional filter for specific question ID
|
||||
max_concurrent_items=max_concurrent_items, # Parallel instance processing
|
||||
output_path=output_path, # Save results incrementally
|
||||
merge_with_existing=merge_with_existing # Merge when using --fill, --category, --only-failed, --only-invalid flags or specific question
|
||||
merge_with_existing=merge_with_existing, # Merge when using --fill, --category, --only-failed, --only-invalid flags or specific question
|
||||
)
|
||||
|
||||
# Display results (final save already happened incrementally)
|
||||
@@ -702,12 +726,14 @@ def download_dataset(dataset_path: Path) -> bool:
|
||||
True if successful, False otherwise
|
||||
"""
|
||||
import subprocess
|
||||
|
||||
from rich.console import Console
|
||||
|
||||
console = Console()
|
||||
|
||||
url = "https://huggingface.co/datasets/xiaowu0162/longmemeval-cleaned/resolve/main/longmemeval_s_cleaned.json"
|
||||
|
||||
console.print(f"[yellow]Dataset not found. Downloading from HuggingFace...[/yellow]")
|
||||
console.print("[yellow]Dataset not found. Downloading from HuggingFace...[/yellow]")
|
||||
console.print(f"[dim]URL: {url}[/dim]")
|
||||
console.print(f"[dim]Destination: {dataset_path}[/dim]")
|
||||
|
||||
@@ -720,18 +746,18 @@ def download_dataset(dataset_path: Path) -> bool:
|
||||
["curl", "-L", "-o", str(dataset_path), url],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=300 # 5 minute timeout
|
||||
timeout=300, # 5 minute timeout
|
||||
)
|
||||
|
||||
if result.returncode == 0 and dataset_path.exists():
|
||||
console.print(f"[green]✓ Dataset downloaded successfully[/green]")
|
||||
console.print("[green]✓ Dataset downloaded successfully[/green]")
|
||||
return True
|
||||
else:
|
||||
console.print(f"[red]✗ Download failed: {result.stderr}[/red]")
|
||||
return False
|
||||
|
||||
except subprocess.TimeoutExpired:
|
||||
console.print(f"[red]✗ Download timed out after 5 minutes[/red]")
|
||||
console.print("[red]✗ Download timed out after 5 minutes[/red]")
|
||||
return False
|
||||
except Exception as e:
|
||||
console.print(f"[red]✗ Download error: {e}[/red]")
|
||||
@@ -740,22 +766,23 @@ def download_dataset(dataset_path: Path) -> bool:
|
||||
|
||||
def generate_type_report(results: dict):
|
||||
"""Generate a detailed report by question type."""
|
||||
from rich.table import Table
|
||||
from rich.console import Console
|
||||
from rich.table import Table
|
||||
|
||||
console = Console()
|
||||
|
||||
# Aggregate stats by question type
|
||||
type_stats = {}
|
||||
|
||||
for item_result in results['item_results']:
|
||||
metrics = item_result['metrics']
|
||||
by_category = metrics.get('category_stats', {})
|
||||
for item_result in results["item_results"]:
|
||||
metrics = item_result["metrics"]
|
||||
by_category = metrics.get("category_stats", {})
|
||||
|
||||
for qtype, stats in by_category.items():
|
||||
if qtype not in type_stats:
|
||||
type_stats[qtype] = {'total': 0, 'correct': 0}
|
||||
type_stats[qtype]['total'] += stats['total']
|
||||
type_stats[qtype]['correct'] += stats['correct']
|
||||
type_stats[qtype] = {"total": 0, "correct": 0}
|
||||
type_stats[qtype]["total"] += stats["total"]
|
||||
type_stats[qtype]["correct"] += stats["correct"]
|
||||
|
||||
# Display table
|
||||
table = Table(title="Performance by Question Type")
|
||||
@@ -765,13 +792,8 @@ def generate_type_report(results: dict):
|
||||
table.add_column("Accuracy", justify="right", style="magenta")
|
||||
|
||||
for qtype, stats in sorted(type_stats.items()):
|
||||
acc = (stats['correct'] / stats['total'] * 100) if stats['total'] > 0 else 0
|
||||
table.add_row(
|
||||
qtype,
|
||||
str(stats['total']),
|
||||
str(stats['correct']),
|
||||
f"{acc:.1f}%"
|
||||
)
|
||||
acc = (stats["correct"] / stats["total"] * 100) if stats["total"] > 0 else 0
|
||||
table.add_row(qtype, str(stats["total"]), str(stats["correct"]), f"{acc:.1f}%")
|
||||
|
||||
console.print("\n")
|
||||
console.print(table)
|
||||
@@ -780,21 +802,22 @@ def generate_type_report(results: dict):
|
||||
def generate_markdown_table(results: dict, json_output_path: Path):
|
||||
"""Generate a markdown results table with model configuration."""
|
||||
from rich.console import Console
|
||||
|
||||
console = Console()
|
||||
|
||||
# Aggregate stats by question type
|
||||
type_stats = {}
|
||||
|
||||
for item_result in results['item_results']:
|
||||
metrics = item_result['metrics']
|
||||
by_category = metrics.get('category_stats', {})
|
||||
for item_result in results["item_results"]:
|
||||
metrics = item_result["metrics"]
|
||||
by_category = metrics.get("category_stats", {})
|
||||
|
||||
for qtype, stats in by_category.items():
|
||||
if qtype not in type_stats:
|
||||
type_stats[qtype] = {'total': 0, 'correct': 0, 'invalid': 0}
|
||||
type_stats[qtype]['total'] += stats['total']
|
||||
type_stats[qtype]['correct'] += stats['correct']
|
||||
type_stats[qtype]['invalid'] += stats.get('invalid', 0)
|
||||
type_stats[qtype] = {"total": 0, "correct": 0, "invalid": 0}
|
||||
type_stats[qtype]["total"] += stats["total"]
|
||||
type_stats[qtype]["correct"] += stats["correct"]
|
||||
type_stats[qtype]["invalid"] += stats.get("invalid", 0)
|
||||
|
||||
# Build markdown content
|
||||
lines = []
|
||||
@@ -802,16 +825,20 @@ def generate_markdown_table(results: dict, json_output_path: Path):
|
||||
lines.append("")
|
||||
|
||||
# Add model configuration
|
||||
if 'model_config' in results:
|
||||
config = results['model_config']
|
||||
if "model_config" in results:
|
||||
config = results["model_config"]
|
||||
lines.append("## Model Configuration")
|
||||
lines.append("")
|
||||
lines.append(f"- **Hindsight**: {config['hindsight']['provider']}/{config['hindsight']['model']}")
|
||||
lines.append(f"- **Answer Generation**: {config['answer_generation']['provider']}/{config['answer_generation']['model']}")
|
||||
lines.append(
|
||||
f"- **Answer Generation**: {config['answer_generation']['provider']}/{config['answer_generation']['model']}"
|
||||
)
|
||||
lines.append(f"- **LLM Judge**: {config['judge']['provider']}/{config['judge']['model']}")
|
||||
lines.append("")
|
||||
|
||||
lines.append(f"**Overall Accuracy**: {results['overall_accuracy']:.2f}% ({results['total_correct']}/{results['total_questions']})")
|
||||
lines.append(
|
||||
f"**Overall Accuracy**: {results['overall_accuracy']:.2f}% ({results['total_correct']}/{results['total_questions']})"
|
||||
)
|
||||
lines.append("")
|
||||
|
||||
# Results by question type
|
||||
@@ -822,34 +849,36 @@ def generate_markdown_table(results: dict, json_output_path: Path):
|
||||
|
||||
for qtype in sorted(type_stats.keys()):
|
||||
stats = type_stats[qtype]
|
||||
valid_total = stats['total'] - stats['invalid']
|
||||
acc = (stats['correct'] / valid_total * 100) if valid_total > 0 else 0
|
||||
invalid_str = str(stats['invalid']) if stats['invalid'] > 0 else "-"
|
||||
valid_total = stats["total"] - stats["invalid"]
|
||||
acc = (stats["correct"] / valid_total * 100) if valid_total > 0 else 0
|
||||
invalid_str = str(stats["invalid"]) if stats["invalid"] > 0 else "-"
|
||||
lines.append(f"| {qtype} | {stats['total']} | {stats['correct']} | {invalid_str} | {acc:.1f}% |")
|
||||
|
||||
# Add overall row
|
||||
total_invalid = results.get('total_invalid', 0)
|
||||
total_invalid = results.get("total_invalid", 0)
|
||||
invalid_str = str(total_invalid) if total_invalid > 0 else "-"
|
||||
lines.append(f"| **OVERALL** | **{results['total_questions']}** | **{results['total_correct']}** | **{invalid_str}** | **{results['overall_accuracy']:.1f}%** |")
|
||||
lines.append(
|
||||
f"| **OVERALL** | **{results['total_questions']}** | **{results['total_correct']}** | **{invalid_str}** | **{results['overall_accuracy']:.1f}%** |"
|
||||
)
|
||||
|
||||
# Write to file (same directory as JSON, but .md extension)
|
||||
md_output_path = json_output_path.with_suffix('.md')
|
||||
md_output_path.write_text('\n'.join(lines))
|
||||
md_output_path = json_output_path.with_suffix(".md")
|
||||
md_output_path.write_text("\n".join(lines))
|
||||
console.print(f"\n[green]✓[/green] Results table saved to {md_output_path}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import logging
|
||||
import argparse
|
||||
import logging
|
||||
|
||||
logging.basicConfig(level=logging.INFO, format='%(asctime)s %(levelname)s %(message)s')
|
||||
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
|
||||
|
||||
parser = argparse.ArgumentParser(description="Run LongMemEval benchmark")
|
||||
parser.add_argument(
|
||||
"--max-instances",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Limit TOTAL number of questions to evaluate (default: all 500). For per-category limits, use --max-questions-per-category instead."
|
||||
help="Limit TOTAL number of questions to evaluate (default: all 500). For per-category limits, use --max-questions-per-category instead.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-instances-per-category",
|
||||
@@ -857,87 +886,72 @@ if __name__ == "__main__":
|
||||
type=int,
|
||||
default=None,
|
||||
dest="max_instances_per_category",
|
||||
help="Limit number of questions per category (e.g., 20 = 20 questions from each of the 6 categories = 120 total). Cannot be combined with --max-instances or --category."
|
||||
help="Limit number of questions per category (e.g., 20 = 20 questions from each of the 6 categories = 120 total). Cannot be combined with --max-instances or --category.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-questions",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Limit number of questions per instance (for quick testing)"
|
||||
"--max-questions", type=int, default=None, help="Limit number of questions per instance (for quick testing)"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--thinking-budget",
|
||||
type=int,
|
||||
default=500,
|
||||
help="Thinking budget for spreading activation search"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-tokens",
|
||||
type=int,
|
||||
default=8192,
|
||||
help="Maximum tokens to retrieve from memories"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--skip-ingestion",
|
||||
action="store_true",
|
||||
help="Skip ingestion and use existing data"
|
||||
"--thinking-budget", type=int, default=500, help="Thinking budget for spreading activation search"
|
||||
)
|
||||
parser.add_argument("--max-tokens", type=int, default=8192, help="Maximum tokens to retrieve from memories")
|
||||
parser.add_argument("--skip-ingestion", action="store_true", help="Skip ingestion and use existing data")
|
||||
parser.add_argument(
|
||||
"--fill",
|
||||
action="store_true",
|
||||
help="Only process questions not already in results file (for resuming interrupted runs)"
|
||||
help="Only process questions not already in results file (for resuming interrupted runs)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--question-id",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Filter to specific question ID (e.g., 'e47becba'). Useful with --skip-ingestion to test a single question."
|
||||
help="Filter to specific question ID (e.g., 'e47becba'). Useful with --skip-ingestion to test a single question.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--only-failed",
|
||||
action="store_true",
|
||||
help="Only run questions that were previously marked as incorrect (is_correct=False). Requires existing results file."
|
||||
help="Only run questions that were previously marked as incorrect (is_correct=False). Requires existing results file.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--only-invalid",
|
||||
action="store_true",
|
||||
help="Only run questions that were previously marked as invalid (is_invalid=True). Requires existing results file."
|
||||
help="Only run questions that were previously marked as invalid (is_invalid=True). Requires existing results file.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--only-ingested",
|
||||
action="store_true",
|
||||
help="Only run questions whose memory bank already exists (has been ingested). Automatically skips ingestion. Cannot be combined with --only-failed, --only-invalid, --category, --question-id, or --max-instances-per-category."
|
||||
help="Only run questions whose memory bank already exists (has been ingested). Automatically skips ingestion. Cannot be combined with --only-failed, --only-invalid, --category, --question-id, or --max-instances-per-category.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--category",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Filter questions by category/question_type. Available categories: 'single-session-user', 'multi-session', 'single-session-preference', 'temporal-reasoning', 'knowledge-update', 'single-session-assistant'. Can be combined with --max-instances to limit questions within the category."
|
||||
help="Filter questions by category/question_type. Available categories: 'single-session-user', 'multi-session', 'single-session-preference', 'temporal-reasoning', 'knowledge-update', 'single-session-assistant'. Can be combined with --max-instances to limit questions within the category.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--parallel",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Number of instances to process in parallel (default: 1 for sequential). Higher values speed up evaluation but use more memory."
|
||||
help="Number of instances to process in parallel (default: 1 for sequential). Higher values speed up evaluation but use more memory.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--results-filename",
|
||||
type=str,
|
||||
default="benchmark_results.json",
|
||||
help="Filename for results output (default: benchmark_results.json). Saved in results/ directory."
|
||||
help="Filename for results output (default: benchmark_results.json). Saved in results/ directory.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--context-format",
|
||||
type=str,
|
||||
choices=["json", "structured"],
|
||||
default="json",
|
||||
help="How to format context for answer generation. 'json' (raw JSON dump, original behavior) or 'structured' (human-readable format with facts grouped with source chunks). Default: json."
|
||||
help="How to format context for answer generation. 'json' (raw JSON dump, original behavior) or 'structured' (human-readable format with facts grouped with source chunks). Default: json.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--source-results",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Source results file to read failed/invalid questions from (for --only-failed/--only-invalid). Defaults to benchmark_results.json if not specified."
|
||||
help="Source results file to read failed/invalid questions from (for --only-failed/--only-invalid). Defaults to benchmark_results.json if not specified.",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
@@ -951,21 +965,23 @@ if __name__ == "__main__":
|
||||
if args.max_instances_per_category is not None and (args.max_instances is not None or args.category is not None):
|
||||
parser.error("--max-questions-per-category cannot be combined with --max-instances or --category")
|
||||
|
||||
results = asyncio.run(run_benchmark(
|
||||
max_instances=args.max_instances,
|
||||
max_instances_per_category=args.max_instances_per_category,
|
||||
max_questions_per_instance=args.max_questions,
|
||||
thinking_budget=args.thinking_budget,
|
||||
max_tokens=args.max_tokens,
|
||||
skip_ingestion=args.skip_ingestion,
|
||||
filln=args.fill,
|
||||
question_id=args.question_id,
|
||||
only_failed=args.only_failed,
|
||||
only_invalid=args.only_invalid,
|
||||
only_ingested=args.only_ingested,
|
||||
category=args.category,
|
||||
max_concurrent_items=args.parallel,
|
||||
results_filename=args.results_filename,
|
||||
context_format=args.context_format,
|
||||
source_results=args.source_results
|
||||
))
|
||||
results = asyncio.run(
|
||||
run_benchmark(
|
||||
max_instances=args.max_instances,
|
||||
max_instances_per_category=args.max_instances_per_category,
|
||||
max_questions_per_instance=args.max_questions,
|
||||
thinking_budget=args.thinking_budget,
|
||||
max_tokens=args.max_tokens,
|
||||
skip_ingestion=args.skip_ingestion,
|
||||
filln=args.fill,
|
||||
question_id=args.question_id,
|
||||
only_failed=args.only_failed,
|
||||
only_invalid=args.only_invalid,
|
||||
only_ingested=args.only_ingested,
|
||||
category=args.category,
|
||||
max_concurrent_items=args.parallel,
|
||||
results_filename=args.results_filename,
|
||||
context_format=args.context_format,
|
||||
source_results=args.source_results,
|
||||
)
|
||||
)
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -5,6 +5,7 @@ Generate changelog entry for a new release.
|
||||
This script fetches the commit diff between releases, uses an LLM to summarize,
|
||||
and prepends the entry to the changelog page.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
@@ -29,6 +30,7 @@ CHANGELOG_PATH = REPO_PATH / "hindsight-docs" / "docs" / "changelog" / "index.md
|
||||
|
||||
class ChangelogEntry(BaseModel):
|
||||
"""A single changelog entry."""
|
||||
|
||||
category: str # "feature", "improvement", "bugfix", "breaking", "other"
|
||||
summary: str # Brief description of the change
|
||||
commit_id: str # Short commit hash
|
||||
@@ -36,12 +38,14 @@ class ChangelogEntry(BaseModel):
|
||||
|
||||
class ChangelogResponse(BaseModel):
|
||||
"""Structured response from LLM."""
|
||||
|
||||
entries: list[ChangelogEntry]
|
||||
|
||||
|
||||
@dataclass
|
||||
class Commit:
|
||||
"""Parsed commit from git log."""
|
||||
|
||||
hash: str
|
||||
message: str
|
||||
|
||||
@@ -151,10 +155,7 @@ def analyze_commits_with_llm(
|
||||
file_diff: str,
|
||||
) -> list[ChangelogEntry]:
|
||||
"""Use LLM to analyze commits and return structured changelog entries."""
|
||||
commits_json = json.dumps(
|
||||
[{"commit_id": c.hash, "message": c.message} for c in commits],
|
||||
indent=2
|
||||
)
|
||||
commits_json = json.dumps([{"commit_id": c.hash, "message": c.message} for c in commits], indent=2)
|
||||
|
||||
prompt = f"""Analyze the following git commits for release {version} of Hindsight (an AI memory system).
|
||||
|
||||
@@ -214,9 +215,11 @@ def build_changelog_markdown(
|
||||
# Build markdown
|
||||
lines = [f"## [{version}]({release_url})", ""]
|
||||
|
||||
has_entries = False
|
||||
for cat_key in ["breaking", "feature", "improvement", "bugfix", "other"]:
|
||||
cat_name, cat_entries = categories[cat_key]
|
||||
if cat_entries:
|
||||
has_entries = True
|
||||
lines.append(f"**{cat_name}**")
|
||||
lines.append("")
|
||||
for entry in cat_entries:
|
||||
@@ -224,6 +227,10 @@ def build_changelog_markdown(
|
||||
lines.append(f"- {entry.summary} ([`{entry.commit_id}`]({commit_url}))")
|
||||
lines.append("")
|
||||
|
||||
if not has_entries:
|
||||
lines.append("*This release contains internal maintenance and infrastructure changes only.*")
|
||||
lines.append("")
|
||||
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
@@ -236,6 +243,8 @@ sidebar_position: 1
|
||||
|
||||
# Changelog
|
||||
|
||||
This changelog highlights user-facing changes only. Internal maintenance, CI/CD, and infrastructure updates are omitted.
|
||||
|
||||
For full release details, see [GitHub Releases](https://github.com/vectorize-io/hindsight/releases).
|
||||
"""
|
||||
return header, ""
|
||||
@@ -244,8 +253,8 @@ For full release details, see [GitHub Releases](https://github.com/vectorize-io/
|
||||
|
||||
match = re.search(r"^## ", content, re.MULTILINE)
|
||||
if match:
|
||||
header = content[:match.start()].rstrip() + "\n\n"
|
||||
releases = content[match.start():]
|
||||
header = content[: match.start()].rstrip() + "\n\n"
|
||||
releases = content[match.start() :]
|
||||
else:
|
||||
header = content.rstrip() + "\n\n"
|
||||
releases = ""
|
||||
@@ -275,7 +284,7 @@ def generate_changelog_entry(
|
||||
tag = version if version.startswith("v") else f"v{version}"
|
||||
display_version = version.lstrip("v")
|
||||
|
||||
console.print(f"[blue]Fetching tags from repository...[/blue]")
|
||||
console.print("[blue]Fetching tags from repository...[/blue]")
|
||||
existing_tags = get_git_tags()
|
||||
|
||||
if tag not in existing_tags and display_version not in existing_tags:
|
||||
@@ -292,7 +301,7 @@ def generate_changelog_entry(
|
||||
else:
|
||||
console.print("[yellow]No previous version found, will include all commits[/yellow]")
|
||||
|
||||
console.print(f"[blue]Getting commits...[/blue]")
|
||||
console.print("[blue]Getting commits...[/blue]")
|
||||
commits = get_commits(previous_tag, actual_tag)
|
||||
file_diff = get_detailed_diff(previous_tag, actual_tag)
|
||||
|
||||
|
||||
@@ -4,13 +4,15 @@ Generate OpenAPI specification from FastAPI app.
|
||||
|
||||
This script imports the FastAPI app and exports its OpenAPI schema to a JSON file.
|
||||
"""
|
||||
|
||||
import json
|
||||
import sys
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
from hindsight_api.api import create_app
|
||||
from hindsight_api import MemoryEngine
|
||||
from hindsight_api.api import create_app
|
||||
|
||||
|
||||
def generate_openapi_spec(output_path: str = None):
|
||||
"""Generate OpenAPI spec and save to file."""
|
||||
@@ -34,7 +36,7 @@ def generate_openapi_spec(output_path: str = None):
|
||||
|
||||
# Write to file
|
||||
output_file = Path(output_path)
|
||||
with open(output_file, 'w') as f:
|
||||
with open(output_file, "w") as f:
|
||||
json.dump(openapi_schema, f, indent=2)
|
||||
|
||||
print(f"✓ OpenAPI specification generated: {output_file.absolute()}")
|
||||
@@ -44,14 +46,15 @@ def generate_openapi_spec(output_path: str = None):
|
||||
|
||||
# List endpoints
|
||||
print("\n Endpoints:")
|
||||
for path, methods in openapi_schema['paths'].items():
|
||||
for path, methods in openapi_schema["paths"].items():
|
||||
for method in methods.keys():
|
||||
if method.upper() in ['GET', 'POST', 'PUT', 'DELETE', 'PATCH']:
|
||||
if method.upper() in ["GET", "POST", "PUT", "DELETE", "PATCH"]:
|
||||
endpoint_info = methods[method]
|
||||
summary = endpoint_info.get('summary', 'No summary')
|
||||
tags = ', '.join(endpoint_info.get('tags', ['untagged']))
|
||||
summary = endpoint_info.get("summary", "No summary")
|
||||
tags = ", ".join(endpoint_info.get("tags", ["untagged"]))
|
||||
print(f" {method.upper():6} {path:30} [{tags}] - {summary}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
output = sys.argv[1] if len(sys.argv) > 1 else "openapi.json"
|
||||
generate_openapi_spec(output)
|
||||
|
||||
@@ -212,9 +212,7 @@ This recipe is available as an interactive Jupyter notebook.
|
||||
# Insert callout after first heading
|
||||
first_heading_match = re.search(r"^(#\s+.+\n)", md_content, re.MULTILINE)
|
||||
if first_heading_match:
|
||||
idx = md_content.index(first_heading_match.group(0)) + len(
|
||||
first_heading_match.group(0)
|
||||
)
|
||||
idx = md_content.index(first_heading_match.group(0)) + len(first_heading_match.group(0))
|
||||
final_content = md_content[:idx] + "\n" + callout + "\n" + md_content[idx:]
|
||||
else:
|
||||
final_content = callout + "\n" + md_content
|
||||
@@ -247,9 +245,7 @@ def process_applications(cookbook_dir: Path, apps_dir: Path) -> list[dict]:
|
||||
continue
|
||||
|
||||
slug = entry.name
|
||||
title = extract_title_from_readme(readme_path) or " ".join(
|
||||
word.capitalize() for word in slug.split("-")
|
||||
)
|
||||
title = extract_title_from_readme(readme_path) or " ".join(word.capitalize() for word in slug.split("-"))
|
||||
|
||||
print(f" Processing app: {entry.name} → {slug}.md")
|
||||
|
||||
@@ -275,12 +271,8 @@ This is a complete, runnable application demonstrating Hindsight integration.
|
||||
# Insert callout after first heading
|
||||
first_heading_match = re.search(r"^(#\s+.+\n)", readme_content, re.MULTILINE)
|
||||
if first_heading_match:
|
||||
idx = readme_content.index(first_heading_match.group(0)) + len(
|
||||
first_heading_match.group(0)
|
||||
)
|
||||
final_content = (
|
||||
readme_content[:idx] + "\n" + callout + "\n" + readme_content[idx:]
|
||||
)
|
||||
idx = readme_content.index(first_heading_match.group(0)) + len(first_heading_match.group(0))
|
||||
final_content = readme_content[:idx] + "\n" + callout + "\n" + readme_content[idx:]
|
||||
else:
|
||||
final_content = callout + "\n" + readme_content
|
||||
|
||||
@@ -407,26 +399,20 @@ def clean_description(desc: str) -> str:
|
||||
return desc
|
||||
|
||||
|
||||
def update_cookbook_index(
|
||||
recipes: list[dict], apps: list[dict], docs_dir: Path
|
||||
):
|
||||
def update_cookbook_index(recipes: list[dict], apps: list[dict], docs_dir: Path):
|
||||
"""Update cookbook/index.mdx with recipe and app carousels."""
|
||||
# Build recipe items for the carousel
|
||||
recipe_items = []
|
||||
for r in recipes:
|
||||
title = r["title"].replace('"', '\\"')
|
||||
recipe_items.append(
|
||||
f' {{ title: "{title}", href: "/cookbook/recipes/{r["slug"]}" }}'
|
||||
)
|
||||
recipe_items.append(f' {{ title: "{title}", href: "/cookbook/recipes/{r["slug"]}" }}')
|
||||
recipes_json = ",\n".join(recipe_items)
|
||||
|
||||
# Build app items for the carousel
|
||||
app_items = []
|
||||
for a in apps:
|
||||
title = a["title"].replace('"', '\\"')
|
||||
app_items.append(
|
||||
f' {{ title: "{title}", href: "/cookbook/applications/{a["slug"]}" }}'
|
||||
)
|
||||
app_items.append(f' {{ title: "{title}", href: "/cookbook/applications/{a["slug"]}" }}')
|
||||
apps_json = ",\n".join(app_items)
|
||||
|
||||
content = f"""---
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "hindsight-dev"
|
||||
version = "0.1.8"
|
||||
version = "0.1.13"
|
||||
description = "Development utilities for Hindsight"
|
||||
requires-python = ">=3.11"
|
||||
dependencies = [
|
||||
@@ -27,3 +27,58 @@ generate-openapi = "hindsight_dev.generate_openapi:generate_openapi_spec"
|
||||
generate-changelog = "hindsight_dev.generate_changelog:main"
|
||||
sync-cookbook = "hindsight_dev.sync_cookbook:main"
|
||||
generate-llms-full = "hindsight_dev.generate_llms_full:main"
|
||||
|
||||
[dependency-groups]
|
||||
dev = [
|
||||
"ruff>=0.8.0",
|
||||
"ty>=0.0.1",
|
||||
]
|
||||
|
||||
[tool.ruff]
|
||||
line-length = 120
|
||||
target-version = "py311"
|
||||
|
||||
[tool.ruff.lint]
|
||||
select = [
|
||||
"E", # pycodestyle errors
|
||||
"W", # pycodestyle warnings
|
||||
"F", # Pyflakes
|
||||
"I", # isort
|
||||
]
|
||||
ignore = [
|
||||
"E501", # line too long (handled by formatter)
|
||||
"E402", # module import not at top of file
|
||||
"E712", # avoid equality comparisons to False (intentional for optional bools)
|
||||
"F401", # unused import (too noisy during development)
|
||||
"F403", # star imports (fasthtml uses this pattern)
|
||||
"F405", # may be undefined from star imports (fasthtml uses this pattern)
|
||||
"F841", # unused variable (too noisy during development)
|
||||
"F811", # redefined while unused
|
||||
"F821", # undefined name (forward references in type hints)
|
||||
"W291", # trailing whitespace (in multiline strings)
|
||||
"W293", # blank line contains whitespace (in multiline strings)
|
||||
]
|
||||
|
||||
[tool.ruff.format]
|
||||
quote-style = "double"
|
||||
indent-style = "space"
|
||||
|
||||
[tool.ty]
|
||||
# Type checking configuration
|
||||
# ty is an extremely fast Python type checker from Astral (same team as ruff/uv)
|
||||
|
||||
[tool.ty.environment]
|
||||
python-version = "3.11"
|
||||
|
||||
[tool.ty.rules]
|
||||
# Disable noisy rules
|
||||
invalid-argument-type = "ignore" # Too many false positives
|
||||
invalid-return-type = "ignore" # Often intentional
|
||||
# missing-argument and unknown-argument are enabled (default)
|
||||
invalid-parameter-default = "ignore" # Optional params with None default
|
||||
possibly-missing-attribute = "ignore" # Common with Optional types
|
||||
unsupported-operator = "ignore" # Pandas DataFrame operations
|
||||
unresolved-reference = "ignore" # Forward references
|
||||
unresolved-import = "ignore" # Dynamic imports
|
||||
invalid-assignment = "ignore" # Intentional monkey-patching
|
||||
no-matching-overload = "ignore" # Complex generic issues
|
||||
|
||||
@@ -4,8 +4,44 @@ sidebar_position: 1
|
||||
|
||||
# Changelog
|
||||
|
||||
This changelog highlights user-facing changes only. Internal maintenance, CI/CD, and infrastructure updates are omitted.
|
||||
|
||||
For full release details, see [GitHub Releases](https://github.com/vectorize-io/hindsight/releases).
|
||||
|
||||
## [0.1.11](https://github.com/vectorize-io/hindsight/releases/tag/v0.1.11)
|
||||
|
||||
**Bug Fixes**
|
||||
|
||||
- Fixed the standalone Docker image and control plane standalone build process so standalone deployments build correctly. ([`2948cb6`](https://github.com/vectorize-io/hindsight/commit/2948cb6))
|
||||
|
||||
## [0.1.10](https://github.com/vectorize-io/hindsight/releases/tag/v0.1.10)
|
||||
|
||||
*This release contains internal maintenance and infrastructure changes only.*
|
||||
|
||||
|
||||
## [0.1.9](https://github.com/vectorize-io/hindsight/releases/tag/v0.1.9)
|
||||
|
||||
**Features**
|
||||
|
||||
- Simplified local MCP installation and added a standalone UI option for easier setup. ([`1c6acc3`](https://github.com/vectorize-io/hindsight/commit/1c6acc3))
|
||||
|
||||
**Bug Fixes**
|
||||
|
||||
- Fixed the standalone Docker image so it builds and starts reliably. ([`b52eb90`](https://github.com/vectorize-io/hindsight/commit/b52eb90))
|
||||
- Improved Docker runtime reliability by adding required system utilities (procps). ([`ae80876`](https://github.com/vectorize-io/hindsight/commit/ae80876))
|
||||
|
||||
## [0.1.8](https://github.com/vectorize-io/hindsight/releases/tag/v0.1.8)
|
||||
|
||||
**Bug Fixes**
|
||||
|
||||
- Fix bank list responses when a bank has no name. ([`04f01ab`](https://github.com/vectorize-io/hindsight/commit/04f01ab))
|
||||
- Fix failures when retaining memories asynchronously. ([`63f5138`](https://github.com/vectorize-io/hindsight/commit/63f5138))
|
||||
- Fix a race condition in the bank selector when switching banks. ([`e468a4e`](https://github.com/vectorize-io/hindsight/commit/e468a4e))
|
||||
|
||||
## [0.1.7](https://github.com/vectorize-io/hindsight/releases/tag/v0.1.7)
|
||||
|
||||
*This release contains internal maintenance and infrastructure changes only.*
|
||||
|
||||
## [0.1.6](https://github.com/vectorize-io/hindsight/releases/tag/v0.1.6)
|
||||
|
||||
**Features**
|
||||
|
||||
@@ -1,238 +0,0 @@
|
||||
---
|
||||
sidebar_position: 8
|
||||
---
|
||||
|
||||
# Documents
|
||||
|
||||
Track and manage document sources in your memory bank. Documents provide traceability — knowing where memories came from.
|
||||
|
||||
import Tabs from '@theme/Tabs';
|
||||
import TabItem from '@theme/TabItem';
|
||||
|
||||
:::tip Prerequisites
|
||||
Make sure you've completed the [Quick Start](./quickstart) and understand [how retain works](./retain).
|
||||
:::
|
||||
|
||||
## What Are Documents?
|
||||
|
||||
Documents are containers for retained content. They help you:
|
||||
|
||||
- **Track sources** — Know which PDF, conversation, or file a memory came from
|
||||
- **Update content** — Re-retain a document to update its facts
|
||||
- **Delete in bulk** — Remove all memories from a document at once
|
||||
- **Organize memories** — Group related facts by source
|
||||
|
||||
## Chunks
|
||||
|
||||
When you retain content, Hindsight splits it into chunks before extracting facts. These chunks are stored alongside the extracted memories, preserving the original text segments.
|
||||
|
||||
**Why chunks matter:**
|
||||
- **Context preservation** — Chunks contain the raw text that generated facts, useful when you need the exact wording
|
||||
- **Richer recall** — Including chunks in recall provides surrounding context for matched facts
|
||||
|
||||
:::tip Include Chunks in Recall
|
||||
Use `include_chunks=True` in your recall calls to get the original text chunks alongside fact results. See [Recall](./recall) for details.
|
||||
:::
|
||||
|
||||
## Retain with Document ID
|
||||
|
||||
Associate retained content with a document:
|
||||
|
||||
<Tabs>
|
||||
<TabItem value="python" label="Python">
|
||||
|
||||
```python
|
||||
from hindsight_client import Hindsight
|
||||
|
||||
client = Hindsight(base_url="http://localhost:8888")
|
||||
|
||||
# Retain with document ID
|
||||
client.retain(
|
||||
bank_id="my-bank",
|
||||
content="Alice presented the Q4 roadmap...",
|
||||
document_id="meeting-2024-03-15"
|
||||
)
|
||||
|
||||
# Batch retain for a document
|
||||
client.retain_batch(
|
||||
bank_id="my-bank",
|
||||
items=[
|
||||
{"content": "Item 1: Product launch delayed to Q2"},
|
||||
{"content": "Item 2: New hiring targets announced"},
|
||||
{"content": "Item 3: Budget approved for ML team"}
|
||||
],
|
||||
document_id="meeting-2024-03-15"
|
||||
)
|
||||
|
||||
# From file
|
||||
with open("notes.txt") as f:
|
||||
client.retain(
|
||||
bank_id="my-bank",
|
||||
content=f.read(),
|
||||
document_id="notes-2024-03-15"
|
||||
)
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
<TabItem value="node" label="Node.js">
|
||||
|
||||
```typescript
|
||||
import { HindsightClient } from '@vectorize-io/hindsight-client';
|
||||
|
||||
const client = new HindsightClient({ baseUrl: 'http://localhost:8888' });
|
||||
|
||||
// Retain with document ID
|
||||
await client.retain('my-bank', 'Alice presented the Q4 roadmap...', {
|
||||
document_id: 'meeting-2024-03-15'
|
||||
});
|
||||
|
||||
// Batch retain
|
||||
await client.retainBatch('my-bank', [
|
||||
{ content: 'Item 1: Product launch delayed to Q2' },
|
||||
{ content: 'Item 2: New hiring targets announced' },
|
||||
{ content: 'Item 3: Budget approved for ML team' }
|
||||
], { documentId: 'meeting-2024-03-15' });
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
<TabItem value="cli" label="CLI">
|
||||
|
||||
```bash
|
||||
# Retain file with document ID
|
||||
hindsight retain my-bank --file notes.txt --document-id notes-2024-03-15
|
||||
|
||||
# Batch retain directory
|
||||
hindsight retain my-bank --files docs/*.md --document-id project-docs
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
</Tabs>
|
||||
|
||||
## Update Documents
|
||||
|
||||
Re-retaining with the same document_id **replaces** the old content:
|
||||
|
||||
<Tabs>
|
||||
<TabItem value="python" label="Python">
|
||||
|
||||
```python
|
||||
# Original
|
||||
client.retain(
|
||||
bank_id="my-bank",
|
||||
content="Project deadline: March 31",
|
||||
document_id="project-plan"
|
||||
)
|
||||
|
||||
# Update (deletes old facts, creates new ones)
|
||||
client.retain(
|
||||
bank_id="my-bank",
|
||||
content="Project deadline: April 15 (extended)",
|
||||
document_id="project-plan"
|
||||
)
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
<TabItem value="node" label="Node.js">
|
||||
|
||||
```typescript
|
||||
// Original
|
||||
await client.retain('my-bank', 'Project deadline: March 31', {
|
||||
document_id: 'project-plan'
|
||||
});
|
||||
|
||||
// Update
|
||||
await client.retain('my-bank', 'Project deadline: April 15 (extended)', {
|
||||
document_id: 'project-plan'
|
||||
});
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
<TabItem value="cli" label="CLI">
|
||||
|
||||
```bash
|
||||
# Original
|
||||
hindsight retain my-bank "Project deadline: March 31" --document-id project-plan
|
||||
|
||||
# Update
|
||||
hindsight retain my-bank "Project deadline: April 15 (extended)" --document-id project-plan
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
</Tabs>
|
||||
|
||||
## Get Document
|
||||
|
||||
Retrieve a document's original text and metadata. This is useful for expanding document context after a recall operation returns memories with document references.
|
||||
|
||||
<Tabs>
|
||||
<TabItem value="python" label="Python">
|
||||
|
||||
```python
|
||||
from hindsight_client_api import ApiClient, Configuration
|
||||
from hindsight_client_api.api import DefaultApi
|
||||
|
||||
config = Configuration(host="http://localhost:8888")
|
||||
api_client = ApiClient(config)
|
||||
api = DefaultApi(api_client)
|
||||
|
||||
# Get document to expand context from recall results
|
||||
doc = api.get_document(
|
||||
bank_id="my-bank",
|
||||
document_id="meeting-2024-03-15"
|
||||
)
|
||||
|
||||
print(f"Document: {doc.id}")
|
||||
print(f"Original text: {doc.original_text}")
|
||||
print(f"Memory count: {doc.memory_unit_count}")
|
||||
print(f"Created: {doc.created_at}")
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
<TabItem value="node" label="Node.js">
|
||||
|
||||
```typescript
|
||||
import { sdk, createClient, createConfig } from '@vectorize-io/hindsight-client';
|
||||
|
||||
const apiClient = createClient(createConfig({ baseUrl: 'http://localhost:8888' }));
|
||||
|
||||
// Get document to expand context from recall results
|
||||
const { data: doc } = await sdk.getDocument({
|
||||
client: apiClient,
|
||||
path: { bank_id: 'my-bank', document_id: 'meeting-2024-03-15' }
|
||||
});
|
||||
|
||||
console.log(`Document: ${doc.id}`);
|
||||
console.log(`Original text: ${doc.original_text}`);
|
||||
console.log(`Memory count: ${doc.memory_unit_count}`);
|
||||
console.log(`Created: ${doc.created_at}`);
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
<TabItem value="cli" label="CLI">
|
||||
|
||||
```bash
|
||||
hindsight documents get my-bank meeting-2024-03-15
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
</Tabs>
|
||||
|
||||
## Document Response Format
|
||||
|
||||
```json
|
||||
{
|
||||
"id": "meeting-2024-03-15",
|
||||
"bank_id": "my-bank",
|
||||
"original_text": "Alice presented the Q4 roadmap...",
|
||||
"content_hash": "abc123def456",
|
||||
"memory_unit_count": 12,
|
||||
"created_at": "2024-03-15T14:00:00Z",
|
||||
"updated_at": "2024-03-15T14:00:00Z"
|
||||
}
|
||||
```
|
||||
|
||||
## Next Steps
|
||||
|
||||
- [**Entities**](./entities) — Track people, places, and concepts
|
||||
- [**Operations**](./operations) — Monitor background tasks
|
||||
- [**Memory Banks**](./memory-banks) — Configure bank settings
|
||||
@@ -0,0 +1,128 @@
|
||||
---
|
||||
sidebar_position: 8
|
||||
---
|
||||
|
||||
# Documents
|
||||
|
||||
Track and manage document sources in your memory bank. Documents provide traceability — knowing where memories came from.
|
||||
|
||||
import Tabs from '@theme/Tabs';
|
||||
import TabItem from '@theme/TabItem';
|
||||
import CodeSnippet from '@site/src/components/CodeSnippet';
|
||||
|
||||
{/* Import raw source files */}
|
||||
import documentsPy from '!!raw-loader!@site/examples/api/documents.py';
|
||||
import documentsMjs from '!!raw-loader!@site/examples/api/documents.mjs';
|
||||
|
||||
:::tip Prerequisites
|
||||
Make sure you've completed the [Quick Start](./quickstart) and understand [how retain works](./retain).
|
||||
:::
|
||||
|
||||
## What Are Documents?
|
||||
|
||||
Documents are containers for retained content. They help you:
|
||||
|
||||
- **Track sources** — Know which PDF, conversation, or file a memory came from
|
||||
- **Update content** — Re-retain a document to update its facts
|
||||
- **Delete in bulk** — Remove all memories from a document at once
|
||||
- **Organize memories** — Group related facts by source
|
||||
|
||||
## Chunks
|
||||
|
||||
When you retain content, Hindsight splits it into chunks before extracting facts. These chunks are stored alongside the extracted memories, preserving the original text segments.
|
||||
|
||||
**Why chunks matter:**
|
||||
- **Context preservation** — Chunks contain the raw text that generated facts, useful when you need the exact wording
|
||||
- **Richer recall** — Including chunks in recall provides surrounding context for matched facts
|
||||
|
||||
:::tip Include Chunks in Recall
|
||||
Use `include_chunks=True` in your recall calls to get the original text chunks alongside fact results. See [Recall](./recall) for details.
|
||||
:::
|
||||
|
||||
## Retain with Document ID
|
||||
|
||||
Associate retained content with a document:
|
||||
|
||||
<Tabs>
|
||||
<TabItem value="python" label="Python">
|
||||
<CodeSnippet code={documentsPy} section="document-retain" language="python" />
|
||||
</TabItem>
|
||||
<TabItem value="node" label="Node.js">
|
||||
<CodeSnippet code={documentsMjs} section="document-retain" language="javascript" />
|
||||
</TabItem>
|
||||
<TabItem value="cli" label="CLI">
|
||||
|
||||
```bash
|
||||
# Retain content with document ID
|
||||
hindsight memory retain my-bank "Meeting notes content..." --doc-id notes-2024-03-15
|
||||
|
||||
# Batch retain from files
|
||||
hindsight memory retain-files my-bank docs/
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
</Tabs>
|
||||
|
||||
## Update Documents
|
||||
|
||||
Re-retaining with the same document_id **replaces** the old content:
|
||||
|
||||
<Tabs>
|
||||
<TabItem value="python" label="Python">
|
||||
<CodeSnippet code={documentsPy} section="document-update" language="python" />
|
||||
</TabItem>
|
||||
<TabItem value="node" label="Node.js">
|
||||
<CodeSnippet code={documentsMjs} section="document-update" language="javascript" />
|
||||
</TabItem>
|
||||
<TabItem value="cli" label="CLI">
|
||||
|
||||
```bash
|
||||
# Original
|
||||
hindsight memory retain my-bank "Project deadline: March 31" --doc-id project-plan
|
||||
|
||||
# Update
|
||||
hindsight memory retain my-bank "Project deadline: April 15 (extended)" --doc-id project-plan
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
</Tabs>
|
||||
|
||||
## Get Document
|
||||
|
||||
Retrieve a document's original text and metadata. This is useful for expanding document context after a recall operation returns memories with document references.
|
||||
|
||||
<Tabs>
|
||||
<TabItem value="python" label="Python">
|
||||
<CodeSnippet code={documentsPy} section="document-get" language="python" />
|
||||
</TabItem>
|
||||
<TabItem value="node" label="Node.js">
|
||||
<CodeSnippet code={documentsMjs} section="document-get" language="javascript" />
|
||||
</TabItem>
|
||||
<TabItem value="cli" label="CLI">
|
||||
|
||||
```bash
|
||||
hindsight document get my-bank meeting-2024-03-15
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
</Tabs>
|
||||
|
||||
## Document Response Format
|
||||
|
||||
```json
|
||||
{
|
||||
"id": "meeting-2024-03-15",
|
||||
"bank_id": "my-bank",
|
||||
"original_text": "Alice presented the Q4 roadmap...",
|
||||
"content_hash": "abc123def456",
|
||||
"memory_unit_count": 12,
|
||||
"created_at": "2024-03-15T14:00:00Z",
|
||||
"updated_at": "2024-03-15T14:00:00Z"
|
||||
}
|
||||
```
|
||||
|
||||
## Next Steps
|
||||
|
||||
- [**Entities**](./entities) — Track people, places, and concepts
|
||||
- [**Operations**](./operations) — Monitor background tasks
|
||||
- [**Memory Banks**](./memory-banks) — Configure bank settings
|
||||
@@ -1,315 +0,0 @@
|
||||
---
|
||||
sidebar_position: 2
|
||||
---
|
||||
|
||||
# Main Methods
|
||||
|
||||
Hindsight provides three core operations: **retain**, **recall**, and **reflect**.
|
||||
|
||||
import Tabs from '@theme/Tabs';
|
||||
import TabItem from '@theme/TabItem';
|
||||
|
||||
:::tip Prerequisites
|
||||
Make sure you've [installed Hindsight](../installation) and completed the [Quick Start](./quickstart).
|
||||
:::
|
||||
|
||||
## Retain: Store Information
|
||||
|
||||
Store conversations, documents, and facts into a memory bank.
|
||||
|
||||
<Tabs>
|
||||
<TabItem value="python" label="Python">
|
||||
|
||||
```python
|
||||
# Store a single fact
|
||||
client.retain(
|
||||
bank_id="my-bank",
|
||||
content="Alice joined Google in March 2024 as a Senior ML Engineer"
|
||||
)
|
||||
|
||||
# Store a conversation
|
||||
conversation = """
|
||||
User: What did you work on today?
|
||||
Assistant: I reviewed the new ML pipeline architecture.
|
||||
User: How did it look?
|
||||
Assistant: Promising, but needs better error handling.
|
||||
"""
|
||||
|
||||
client.retain(
|
||||
bank_id="my-bank",
|
||||
content=conversation,
|
||||
context="Daily standup conversation"
|
||||
)
|
||||
|
||||
# Batch retain multiple items
|
||||
client.retain_batch(
|
||||
bank_id="my-bank",
|
||||
contents=[
|
||||
{"content": "Bob prefers Python for data science"},
|
||||
{"content": "Alice recommends using pytest for testing"},
|
||||
{"content": "The team uses GitHub for code reviews"}
|
||||
]
|
||||
)
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
<TabItem value="node" label="Node.js">
|
||||
|
||||
```javascript
|
||||
// Store a single fact
|
||||
await client.retain({
|
||||
bankId: 'my-bank',
|
||||
content: 'Alice joined Google in March 2024 as a Senior ML Engineer'
|
||||
});
|
||||
|
||||
// Store a conversation
|
||||
await client.retain({
|
||||
bankId: 'my-bank',
|
||||
content: `
|
||||
User: What did you work on today?
|
||||
Assistant: I reviewed the new ML pipeline architecture.
|
||||
User: How did it look?
|
||||
Assistant: Promising, but needs better error handling.
|
||||
`,
|
||||
context: 'Daily standup conversation'
|
||||
});
|
||||
|
||||
// Batch retain
|
||||
await client.retainBatch({
|
||||
bankId: 'my-bank',
|
||||
contents: [
|
||||
{ content: 'Bob prefers Python for data science' },
|
||||
{ content: 'Alice recommends using pytest for testing' },
|
||||
{ content: 'The team uses GitHub for code reviews' }
|
||||
]
|
||||
});
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
<TabItem value="cli" label="CLI">
|
||||
|
||||
```bash
|
||||
# Store a single fact
|
||||
hindsight retain my-bank "Alice joined Google in March 2024 as a Senior ML Engineer"
|
||||
|
||||
# Store from a file
|
||||
hindsight retain my-bank --file conversation.txt --context "Daily standup"
|
||||
|
||||
# Store multiple files
|
||||
hindsight retain my-bank --files docs/*.md
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
</Tabs>
|
||||
|
||||
**What happens:** Content is processed by an LLM to extract rich facts, identify entities, and build connections in a knowledge graph.
|
||||
|
||||
**See:** [Retain Details](./retain) for advanced options and parameters.
|
||||
|
||||
---
|
||||
|
||||
## Recall: Search Memories
|
||||
|
||||
Search for relevant memories using multi-strategy retrieval.
|
||||
|
||||
<Tabs>
|
||||
<TabItem value="python" label="Python">
|
||||
|
||||
```python
|
||||
# Basic search
|
||||
results = client.recall(
|
||||
bank_id="my-bank",
|
||||
query="What does Alice do at Google?"
|
||||
)
|
||||
|
||||
for result in results:
|
||||
print(f"[{result['weight']:.2f}] {result['text']}")
|
||||
|
||||
# Search with options
|
||||
results = client.recall(
|
||||
bank_id="my-bank",
|
||||
query="What happened last spring?",
|
||||
budget="high", # More thorough graph traversal
|
||||
max_tokens=8192, # Return more context
|
||||
fact_type="world" # Only world facts
|
||||
)
|
||||
|
||||
# Include entity information
|
||||
results = client.recall(
|
||||
bank_id="my-bank",
|
||||
query="Tell me about Alice",
|
||||
include_entities=True,
|
||||
max_entity_tokens=500
|
||||
)
|
||||
|
||||
# Check entity details
|
||||
for entity in results["entities"]:
|
||||
print(f"Entity: {entity['name']}")
|
||||
print(f"Observations: {entity['observations']}")
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
<TabItem value="node" label="Node.js">
|
||||
|
||||
```javascript
|
||||
// Basic search
|
||||
const results = await client.recall({
|
||||
bankId: 'my-bank',
|
||||
query: 'What does Alice do at Google?'
|
||||
});
|
||||
|
||||
results.forEach(r => {
|
||||
console.log(`[${r.weight.toFixed(2)}] ${r.text}`);
|
||||
});
|
||||
|
||||
// Search with options
|
||||
const detailedResults = await client.recall({
|
||||
bankId: 'my-bank',
|
||||
query: 'What happened last spring?',
|
||||
budget: 'high',
|
||||
maxTokens: 8192,
|
||||
factType: 'world'
|
||||
});
|
||||
|
||||
// Include entity information
|
||||
const withEntities = await client.recall({
|
||||
bankId: 'my-bank',
|
||||
query: 'Tell me about Alice',
|
||||
includeEntities: true,
|
||||
maxEntityTokens: 500
|
||||
});
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
<TabItem value="cli" label="CLI">
|
||||
|
||||
```bash
|
||||
# Basic search
|
||||
hindsight recall my-bank "What does Alice do at Google?"
|
||||
|
||||
# Search with options
|
||||
hindsight recall my-bank "What happened last spring?" \
|
||||
--budget high \
|
||||
--max-tokens 8192 \
|
||||
--fact-type world
|
||||
|
||||
# Verbose output (shows weights and sources)
|
||||
hindsight recall my-bank "Tell me about Alice" -v
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
</Tabs>
|
||||
|
||||
**What happens:** Four search strategies (semantic, keyword, graph, temporal) run in parallel, results are fused and reranked.
|
||||
|
||||
**See:** [Recall Details](./recall) for tuning quality vs latency.
|
||||
|
||||
---
|
||||
|
||||
## Reflect: Reason with Disposition
|
||||
|
||||
Generate disposition-aware responses that form opinions based on evidence.
|
||||
|
||||
<Tabs>
|
||||
<TabItem value="python" label="Python">
|
||||
|
||||
```python
|
||||
# Basic reflect
|
||||
response = client.reflect(
|
||||
bank_id="my-bank",
|
||||
query="Should we adopt TypeScript for our backend?"
|
||||
)
|
||||
|
||||
print(response["text"])
|
||||
print("\nBased on:", len(response["based_on"]["world"]), "facts")
|
||||
print("New opinions:", len(response["new_opinions"]))
|
||||
|
||||
# Reflect with options
|
||||
response = client.reflect(
|
||||
bank_id="my-bank",
|
||||
query="What are Alice's strengths for the team lead role?",
|
||||
budget="high", # More thorough reasoning
|
||||
include_entities=True
|
||||
)
|
||||
|
||||
# Access formed opinions
|
||||
for opinion in response["new_opinions"]:
|
||||
print(f"Opinion: {opinion['text']}")
|
||||
print(f"Confidence: {opinion['confidence']}")
|
||||
|
||||
# See which facts influenced the response
|
||||
for fact in response["based_on"]["world"]:
|
||||
print(f"[{fact['weight']:.2f}] {fact['text']}")
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
<TabItem value="node" label="Node.js">
|
||||
|
||||
```javascript
|
||||
// Basic reflect
|
||||
const response = await client.reflect({
|
||||
bankId: 'my-bank',
|
||||
query: 'Should we adopt TypeScript for our backend?'
|
||||
});
|
||||
|
||||
console.log(response.text);
|
||||
console.log(`\nBased on: ${response.basedOn.world.length} facts`);
|
||||
console.log(`New opinions: ${response.newOpinions.length}`);
|
||||
|
||||
// Reflect with options
|
||||
const detailed = await client.reflect({
|
||||
bankId: 'my-bank',
|
||||
query: "What are Alice's strengths for the team lead role?",
|
||||
budget: 'high',
|
||||
includeEntities: true
|
||||
});
|
||||
|
||||
// Access formed opinions
|
||||
detailed.newOpinions.forEach(op => {
|
||||
console.log(`Opinion: ${op.text}`);
|
||||
console.log(`Confidence: ${op.confidence}`);
|
||||
});
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
<TabItem value="cli" label="CLI">
|
||||
|
||||
```bash
|
||||
# Basic reflect
|
||||
hindsight reflect my-bank "Should we adopt TypeScript for our backend?"
|
||||
|
||||
# Verbose output (shows sources and opinions)
|
||||
hindsight reflect my-bank "What are Alice's strengths for the team lead role?" -v
|
||||
|
||||
# With higher reasoning budget
|
||||
hindsight reflect my-bank "Analyze our tech stack" --budget high
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
</Tabs>
|
||||
|
||||
**What happens:** Memories are recalled, bank disposition is loaded, LLM reasons through evidence, new opinions are formed and stored.
|
||||
|
||||
**See:** [Reflect Details](./reflect) for disposition configuration.
|
||||
|
||||
---
|
||||
|
||||
## Comparison
|
||||
|
||||
| Feature | Retain | Recall | Reflect |
|
||||
|---------|--------|--------|---------|
|
||||
| **Purpose** | Store information | Find information | Reason about information |
|
||||
| **Input** | Raw text/documents | Search query | Question/prompt |
|
||||
| **Output** | Memory IDs | Ranked facts | Reasoned response + opinions |
|
||||
| **Uses LLM** | Yes (extraction) | No | Yes (generation) |
|
||||
| **Forms opinions** | No | No | Yes |
|
||||
| **Disposition** | No | No | Yes |
|
||||
|
||||
---
|
||||
|
||||
## Next Steps
|
||||
|
||||
- [**Retain**](./retain) — Advanced options for storing memories
|
||||
- [**Recall**](./recall) — Tuning search quality and performance
|
||||
- [**Reflect**](./reflect) — Configuring disposition and opinions
|
||||
- [**Memory Banks**](./memory-banks) — Managing memory bank disposition
|
||||
@@ -0,0 +1,141 @@
|
||||
---
|
||||
sidebar_position: 2
|
||||
---
|
||||
|
||||
# Main Methods
|
||||
|
||||
Hindsight provides three core operations: **retain**, **recall**, and **reflect**.
|
||||
|
||||
import Tabs from '@theme/Tabs';
|
||||
import TabItem from '@theme/TabItem';
|
||||
import CodeSnippet from '@site/src/components/CodeSnippet';
|
||||
|
||||
{/* Import raw source files */}
|
||||
import mainMethodsPy from '!!raw-loader!@site/examples/api/main-methods.py';
|
||||
import mainMethodsMjs from '!!raw-loader!@site/examples/api/main-methods.mjs';
|
||||
|
||||
:::tip Prerequisites
|
||||
Make sure you've [installed Hindsight](../installation) and completed the [Quick Start](./quickstart).
|
||||
:::
|
||||
|
||||
## Retain: Store Information
|
||||
|
||||
Store conversations, documents, and facts into a memory bank.
|
||||
|
||||
<Tabs>
|
||||
<TabItem value="python" label="Python">
|
||||
<CodeSnippet code={mainMethodsPy} section="main-retain" language="python" />
|
||||
</TabItem>
|
||||
<TabItem value="node" label="Node.js">
|
||||
<CodeSnippet code={mainMethodsMjs} section="main-retain" language="javascript" />
|
||||
</TabItem>
|
||||
<TabItem value="cli" label="CLI">
|
||||
|
||||
```bash
|
||||
# Store a single fact
|
||||
hindsight retain my-bank "Alice joined Google in March 2024 as a Senior ML Engineer"
|
||||
|
||||
# Store from a file
|
||||
hindsight retain my-bank --file conversation.txt --context "Daily standup"
|
||||
|
||||
# Store multiple files
|
||||
hindsight retain my-bank --files docs/*.md
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
</Tabs>
|
||||
|
||||
**What happens:** Content is processed by an LLM to extract rich facts, identify entities, and build connections in a knowledge graph.
|
||||
|
||||
**See:** [Retain Details](./retain) for advanced options and parameters.
|
||||
|
||||
---
|
||||
|
||||
## Recall: Search Memories
|
||||
|
||||
Search for relevant memories using multi-strategy retrieval.
|
||||
|
||||
<Tabs>
|
||||
<TabItem value="python" label="Python">
|
||||
<CodeSnippet code={mainMethodsPy} section="main-recall" language="python" />
|
||||
</TabItem>
|
||||
<TabItem value="node" label="Node.js">
|
||||
<CodeSnippet code={mainMethodsMjs} section="main-recall" language="javascript" />
|
||||
</TabItem>
|
||||
<TabItem value="cli" label="CLI">
|
||||
|
||||
```bash
|
||||
# Basic search
|
||||
hindsight recall my-bank "What does Alice do at Google?"
|
||||
|
||||
# Search with options
|
||||
hindsight recall my-bank "What happened last spring?" \
|
||||
--budget high \
|
||||
--max-tokens 8192 \
|
||||
--fact-type world
|
||||
|
||||
# Verbose output (shows weights and sources)
|
||||
hindsight recall my-bank "Tell me about Alice" -v
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
</Tabs>
|
||||
|
||||
**What happens:** Four search strategies (semantic, keyword, graph, temporal) run in parallel, results are fused and reranked.
|
||||
|
||||
**See:** [Recall Details](./recall) for tuning quality vs latency.
|
||||
|
||||
---
|
||||
|
||||
## Reflect: Reason with Disposition
|
||||
|
||||
Generate disposition-aware responses that form opinions based on evidence.
|
||||
|
||||
<Tabs>
|
||||
<TabItem value="python" label="Python">
|
||||
<CodeSnippet code={mainMethodsPy} section="main-reflect" language="python" />
|
||||
</TabItem>
|
||||
<TabItem value="node" label="Node.js">
|
||||
<CodeSnippet code={mainMethodsMjs} section="main-reflect" language="javascript" />
|
||||
</TabItem>
|
||||
<TabItem value="cli" label="CLI">
|
||||
|
||||
```bash
|
||||
# Basic reflect
|
||||
hindsight reflect my-bank "Should we adopt TypeScript for our backend?"
|
||||
|
||||
# Verbose output (shows sources and opinions)
|
||||
hindsight reflect my-bank "What are Alice's strengths for the team lead role?" -v
|
||||
|
||||
# With higher reasoning budget
|
||||
hindsight reflect my-bank "Analyze our tech stack" --budget high
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
</Tabs>
|
||||
|
||||
**What happens:** Memories are recalled, bank disposition is loaded, LLM reasons through evidence, new opinions are formed and stored.
|
||||
|
||||
**See:** [Reflect Details](./reflect) for disposition configuration.
|
||||
|
||||
---
|
||||
|
||||
## Comparison
|
||||
|
||||
| Feature | Retain | Recall | Reflect |
|
||||
|---------|--------|--------|---------|
|
||||
| **Purpose** | Store information | Find information | Reason about information |
|
||||
| **Input** | Raw text/documents | Search query | Question/prompt |
|
||||
| **Output** | Memory IDs | Ranked facts | Reasoned response + opinions |
|
||||
| **Uses LLM** | Yes (extraction) | No | Yes (generation) |
|
||||
| **Forms opinions** | No | No | Yes |
|
||||
| **Disposition** | No | No | Yes |
|
||||
|
||||
---
|
||||
|
||||
## Next Steps
|
||||
|
||||
- [**Retain**](./retain) — Advanced options for storing memories
|
||||
- [**Recall**](./recall) — Tuning search quality and performance
|
||||
- [**Reflect**](./reflect) — Configuring disposition and opinions
|
||||
- [**Memory Banks**](./memory-banks) — Managing memory bank disposition
|
||||
+9
-54
@@ -8,6 +8,11 @@ Memory banks are isolated containers that store all memory-related data for a sp
|
||||
|
||||
import Tabs from '@theme/Tabs';
|
||||
import TabItem from '@theme/TabItem';
|
||||
import CodeSnippet from '@site/src/components/CodeSnippet';
|
||||
|
||||
{/* Import raw source files */}
|
||||
import memoryBanksPy from '!!raw-loader!@site/examples/api/memory-banks.py';
|
||||
import memoryBanksMjs from '!!raw-loader!@site/examples/api/memory-banks.mjs';
|
||||
|
||||
## What is a Memory Bank?
|
||||
|
||||
@@ -30,43 +35,10 @@ Make sure you've completed the [Quick Start](./quickstart) to install the client
|
||||
|
||||
<Tabs>
|
||||
<TabItem value="python" label="Python">
|
||||
|
||||
```python
|
||||
from hindsight_client import Hindsight
|
||||
|
||||
client = Hindsight(base_url="http://localhost:8888")
|
||||
|
||||
client.create_bank(
|
||||
bank_id="my-bank",
|
||||
name="Research Assistant",
|
||||
background="I am a research assistant specializing in machine learning",
|
||||
disposition={
|
||||
"skepticism": 4,
|
||||
"literalism": 3,
|
||||
"empathy": 3
|
||||
}
|
||||
)
|
||||
```
|
||||
|
||||
<CodeSnippet code={memoryBanksPy} section="create-bank" language="python" />
|
||||
</TabItem>
|
||||
<TabItem value="node" label="Node.js">
|
||||
|
||||
```typescript
|
||||
import { HindsightClient } from '@vectorize-io/hindsight-client';
|
||||
|
||||
const client = new HindsightClient({ baseUrl: 'http://localhost:8888' });
|
||||
|
||||
await client.createBank('my-bank', {
|
||||
name: 'Research Assistant',
|
||||
background: 'I am a research assistant specializing in machine learning',
|
||||
disposition: {
|
||||
skepticism: 4,
|
||||
literalism: 3,
|
||||
empathy: 3
|
||||
}
|
||||
});
|
||||
```
|
||||
|
||||
<CodeSnippet code={memoryBanksMjs} section="create-bank" language="javascript" />
|
||||
</TabItem>
|
||||
<TabItem value="cli" label="CLI">
|
||||
|
||||
@@ -98,27 +70,10 @@ The background is a first-person narrative providing context for opinion formati
|
||||
|
||||
<Tabs>
|
||||
<TabItem value="python" label="Python">
|
||||
|
||||
```python
|
||||
client.create_bank(
|
||||
bank_id="financial-advisor",
|
||||
background="""I am a conservative financial advisor with 20 years of experience.
|
||||
I prioritize capital preservation over aggressive growth.
|
||||
I have seen multiple market crashes and believe in diversification."""
|
||||
)
|
||||
```
|
||||
|
||||
<CodeSnippet code={memoryBanksPy} section="bank-background" language="python" />
|
||||
</TabItem>
|
||||
<TabItem value="node" label="Node.js">
|
||||
|
||||
```typescript
|
||||
await client.createBank('financial-advisor', {
|
||||
background: `I am a conservative financial advisor with 20 years of experience.
|
||||
I prioritize capital preservation over aggressive growth.
|
||||
I have seen multiple market crashes and believe in diversification.`
|
||||
});
|
||||
```
|
||||
|
||||
<CodeSnippet code={memoryBanksMjs} section="bank-background" language="javascript" />
|
||||
</TabItem>
|
||||
</Tabs>
|
||||
|
||||
+25
-71
@@ -8,6 +8,11 @@ How memory banks form, store, and evolve beliefs.
|
||||
|
||||
import Tabs from '@theme/Tabs';
|
||||
import TabItem from '@theme/TabItem';
|
||||
import CodeSnippet from '@site/src/components/CodeSnippet';
|
||||
|
||||
{/* Import raw source files */}
|
||||
import opinionsPy from '!!raw-loader!@site/examples/api/opinions.py';
|
||||
import opinionsMjs from '!!raw-loader!@site/examples/api/opinions.mjs';
|
||||
|
||||
:::tip Prerequisites
|
||||
Make sure you've completed the [Quick Start](./quickstart) to install the client and start the server.
|
||||
@@ -41,20 +46,10 @@ graph LR
|
||||
|
||||
<Tabs>
|
||||
<TabItem value="python" label="Python">
|
||||
|
||||
```python
|
||||
# Ask a question that might form an opinion
|
||||
answer = client.reflect(
|
||||
bank_id="my-bank",
|
||||
query="What do you think about functional programming?"
|
||||
)
|
||||
|
||||
# Check if new opinions were formed
|
||||
for opinion in answer.get("new_opinions", []):
|
||||
print(f"New opinion: {opinion['text']}")
|
||||
print(f"Confidence: {opinion['confidence']}")
|
||||
```
|
||||
|
||||
<CodeSnippet code={opinionsPy} section="opinion-form" language="python" />
|
||||
</TabItem>
|
||||
<TabItem value="node" label="Node.js">
|
||||
<CodeSnippet code={opinionsMjs} section="opinion-form" language="javascript" />
|
||||
</TabItem>
|
||||
</Tabs>
|
||||
|
||||
@@ -62,19 +57,10 @@ for opinion in answer.get("new_opinions", []):
|
||||
|
||||
<Tabs>
|
||||
<TabItem value="python" label="Python">
|
||||
|
||||
```python
|
||||
# Search only opinions
|
||||
opinions = client.recall(
|
||||
bank_id="my-bank",
|
||||
query="programming languages",
|
||||
types=["opinion"]
|
||||
)
|
||||
|
||||
for op in opinions:
|
||||
print(f"{op['text']} (confidence: {op['confidence_score']:.2f})")
|
||||
```
|
||||
|
||||
<CodeSnippet code={opinionsPy} section="opinion-search" language="python" />
|
||||
</TabItem>
|
||||
<TabItem value="node" label="Node.js">
|
||||
<CodeSnippet code={opinionsMjs} section="opinion-search" language="javascript" />
|
||||
</TabItem>
|
||||
<TabItem value="cli" label="CLI">
|
||||
|
||||
@@ -113,39 +99,10 @@ Different dispositions form different opinions from the same facts:
|
||||
|
||||
<Tabs>
|
||||
<TabItem value="python" label="Python">
|
||||
|
||||
```python
|
||||
# Create two memory banks with different dispositions
|
||||
client.create_bank(
|
||||
bank_id="open-minded",
|
||||
disposition={"skepticism": 2, "literalism": 2, "empathy": 4}
|
||||
)
|
||||
|
||||
client.create_bank(
|
||||
bank_id="conservative",
|
||||
disposition={"skepticism": 5, "literalism": 5, "empathy": 2}
|
||||
)
|
||||
|
||||
# Store the same facts to both
|
||||
facts = [
|
||||
"Rust has better memory safety than C++",
|
||||
"C++ has a larger ecosystem and more libraries",
|
||||
"Rust compile times are longer than C++"
|
||||
]
|
||||
for fact in facts:
|
||||
client.retain(bank_id="open-minded", content=fact)
|
||||
client.retain(bank_id="conservative", content=fact)
|
||||
|
||||
# Ask both the same question
|
||||
q = "Should we rewrite our C++ codebase in Rust?"
|
||||
|
||||
answer1 = client.reflect(bank_id="open-minded", query=q)
|
||||
# Likely: "Yes, Rust's safety benefits outweigh migration costs"
|
||||
|
||||
answer2 = client.reflect(bank_id="conservative", query=q)
|
||||
# Likely: "No, C++'s ecosystem and our team's expertise make it the safer choice"
|
||||
```
|
||||
|
||||
<CodeSnippet code={opinionsPy} section="opinion-disposition" language="python" />
|
||||
</TabItem>
|
||||
<TabItem value="node" label="Node.js">
|
||||
<CodeSnippet code={opinionsMjs} section="opinion-disposition" language="javascript" />
|
||||
</TabItem>
|
||||
</Tabs>
|
||||
|
||||
@@ -153,17 +110,14 @@ answer2 = client.reflect(bank_id="conservative", query=q)
|
||||
|
||||
When `reflect` uses opinions, they appear in `based_on`:
|
||||
|
||||
```python
|
||||
answer = client.reflect(bank_id="my-bank", query="What language should I learn?")
|
||||
|
||||
print("World facts used:")
|
||||
for f in answer.based_on.get("world", []):
|
||||
print(f" {f['text']}")
|
||||
|
||||
print("\nOpinions used:")
|
||||
for o in answer.based_on.get("opinion", []):
|
||||
print(f" {o['text']} (confidence: {o['confidence_score']})")
|
||||
```
|
||||
<Tabs>
|
||||
<TabItem value="python" label="Python">
|
||||
<CodeSnippet code={opinionsPy} section="opinion-in-reflect" language="python" />
|
||||
</TabItem>
|
||||
<TabItem value="node" label="Node.js">
|
||||
<CodeSnippet code={opinionsMjs} section="opinion-in-reflect" language="javascript" />
|
||||
</TabItem>
|
||||
</Tabs>
|
||||
|
||||
## Confidence Thresholds
|
||||
|
||||
+9
-38
@@ -8,6 +8,12 @@ Get up and running with Hindsight in 60 seconds.
|
||||
|
||||
import Tabs from '@theme/Tabs';
|
||||
import TabItem from '@theme/TabItem';
|
||||
import CodeSnippet from '@site/src/components/CodeSnippet';
|
||||
|
||||
{/* Import raw source files */}
|
||||
import quickstartPy from '!!raw-loader!@site/examples/api/quickstart.py';
|
||||
import quickstartMjs from '!!raw-loader!@site/examples/api/quickstart.mjs';
|
||||
import quickstartSh from '!!raw-loader!@site/examples/api/quickstart.sh';
|
||||
|
||||
## Start the API Server
|
||||
|
||||
@@ -59,20 +65,7 @@ See [LLM Providers](/developer/models#llm) for more details.
|
||||
pip install hindsight-client
|
||||
```
|
||||
|
||||
```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")
|
||||
```
|
||||
<CodeSnippet code={quickstartPy} section="quickstart-full" language="python" />
|
||||
|
||||
</TabItem>
|
||||
<TabItem value="node" label="Node.js">
|
||||
@@ -81,20 +74,7 @@ client.reflect(bank_id="my-bank", query="Tell me about Alice")
|
||||
npm install @vectorize-io/hindsight-client
|
||||
```
|
||||
|
||||
```javascript
|
||||
const { HindsightClient } = require('@vectorize-io/hindsight-client');
|
||||
|
||||
const client = new HindsightClient({ baseUrl: 'http://localhost:8888' });
|
||||
|
||||
// Retain: Store information
|
||||
await client.retain('my-bank', 'Alice works at Google as a software engineer');
|
||||
|
||||
// Recall: Search memories
|
||||
await client.recall('my-bank', 'What does Alice do?');
|
||||
|
||||
// Reflect: Generate response
|
||||
await client.reflect('my-bank', 'Tell me about Alice');
|
||||
```
|
||||
<CodeSnippet code={quickstartMjs} section="quickstart-full" language="javascript" />
|
||||
|
||||
</TabItem>
|
||||
<TabItem value="cli" label="CLI">
|
||||
@@ -103,16 +83,7 @@ await client.reflect('my-bank', 'Tell me about Alice');
|
||||
curl -fsSL https://hindsight.vectorize.io/get-cli | bash
|
||||
```
|
||||
|
||||
```bash
|
||||
# Retain: Store information
|
||||
hindsight memory retain my-bank "Alice works at Google as a software engineer"
|
||||
|
||||
# Recall: Search memories
|
||||
hindsight memory recall my-bank "What does Alice do?"
|
||||
|
||||
# Reflect: Generate response
|
||||
hindsight memory reflect my-bank "Tell me about Alice"
|
||||
```
|
||||
<CodeSnippet code={quickstartSh} section="quickstart-full" language="bash" />
|
||||
|
||||
</TabItem>
|
||||
</Tabs>
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user