Compare commits
1
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
b22f4bf258 |
@@ -34,7 +34,7 @@ jobs:
|
||||
env:
|
||||
UMAMI_URL: https://analytics.hindsight.vectorize.io
|
||||
UMAMI_WEBSITE_ID: ${{ secrets.UMAMI_WEBSITE_ID }}
|
||||
- uses: actions/upload-pages-artifact@v5
|
||||
- uses: actions/upload-pages-artifact@v4
|
||||
with:
|
||||
path: hindsight-docs/build
|
||||
deploy:
|
||||
|
||||
@@ -1,174 +0,0 @@
|
||||
name: Performance Tests
|
||||
|
||||
on:
|
||||
schedule:
|
||||
# Run daily at 06:00 UTC
|
||||
- cron: "0 6 * * *"
|
||||
workflow_dispatch:
|
||||
inputs:
|
||||
scale:
|
||||
description: "Test scale (perf-test)"
|
||||
type: choice
|
||||
options:
|
||||
- tiny
|
||||
- small
|
||||
- medium
|
||||
- large
|
||||
default: large
|
||||
suite:
|
||||
description: "Perf-test suite to run (blank = all)"
|
||||
type: choice
|
||||
options:
|
||||
- ""
|
||||
- retain
|
||||
- recall
|
||||
default: ""
|
||||
locomo_max_conversations:
|
||||
description: "LoComo max conversations (0 = skip, blank = all)"
|
||||
type: number
|
||||
default: 0
|
||||
locomo_skip:
|
||||
description: "Skip LoComo job"
|
||||
type: boolean
|
||||
default: false
|
||||
ref:
|
||||
description: "Git ref to test (branch, tag, or SHA). Defaults to main."
|
||||
type: string
|
||||
default: ""
|
||||
|
||||
concurrency:
|
||||
group: perf-test
|
||||
cancel-in-progress: true
|
||||
|
||||
jobs:
|
||||
perf-test:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ inputs.ref || github.ref }}
|
||||
|
||||
- name: Install uv
|
||||
uses: astral-sh/setup-uv@v7
|
||||
with:
|
||||
enable-cache: true
|
||||
prune-cache: false
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v6
|
||||
with:
|
||||
python-version-file: ".python-version"
|
||||
|
||||
- name: Cache HuggingFace models
|
||||
uses: actions/cache@v5
|
||||
with:
|
||||
path: ~/.cache/huggingface
|
||||
key: ${{ runner.os }}-huggingface-${{ hashFiles('hindsight-api-slim/pyproject.toml') }}
|
||||
restore-keys: |
|
||||
${{ runner.os }}-huggingface-
|
||||
|
||||
- name: Pre-download models
|
||||
working-directory: ./hindsight-api-slim
|
||||
run: |
|
||||
uv run --frozen --all-extras --index-strategy unsafe-best-match python -c "
|
||||
from sentence_transformers import SentenceTransformer
|
||||
print('Downloading embedding model...')
|
||||
SentenceTransformer('BAAI/bge-small-en-v1.5')
|
||||
print('Model downloaded successfully')
|
||||
"
|
||||
|
||||
- name: Install hindsight-dev dependencies
|
||||
run: |
|
||||
cd hindsight-dev && uv sync --frozen --all-extras --index-strategy unsafe-best-match
|
||||
|
||||
- name: Run perf tests
|
||||
run: |
|
||||
SUITE_ARG=""
|
||||
if [ -n "${{ inputs.suite }}" ]; then
|
||||
SUITE_ARG="--suite ${{ inputs.suite }}"
|
||||
fi
|
||||
./scripts/benchmarks/run-perf-test.sh \
|
||||
--scale ${{ inputs.scale || 'large' }} \
|
||||
$SUITE_ARG \
|
||||
--output perf-results.json
|
||||
|
||||
- name: Upload perf results
|
||||
if: always()
|
||||
uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: perf-results-${{ github.sha }}
|
||||
path: hindsight-dev/perf-results.json
|
||||
retention-days: 90
|
||||
|
||||
locomo:
|
||||
if: inputs.locomo_skip != true
|
||||
runs-on: ubuntu-latest
|
||||
env:
|
||||
HINDSIGHT_API_LLM_PROVIDER: vertexai
|
||||
HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY: /tmp/gcp-credentials.json
|
||||
HINDSIGHT_API_LLM_MODEL: google/gemini-2.5-flash-lite
|
||||
HINDSIGHT_API_JUDGE_LLM_PROVIDER: vertexai
|
||||
HINDSIGHT_API_JUDGE_LLM_MODEL: google/gemini-2.5-flash-lite
|
||||
HINDSIGHT_API_ANSWER_LLM_PROVIDER: vertexai
|
||||
HINDSIGHT_API_ANSWER_LLM_MODEL: google/gemini-3.1-pro-preview
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ inputs.ref || github.ref }}
|
||||
|
||||
- name: Setup GCP credentials
|
||||
run: |
|
||||
printf '%s' '${{ secrets.GCP_VERTEXAI_CREDENTIALS }}' > /tmp/gcp-credentials.json
|
||||
PROJECT_ID=$(jq -r '.project_id' /tmp/gcp-credentials.json)
|
||||
echo "HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID=$PROJECT_ID" >> $GITHUB_ENV
|
||||
|
||||
- name: Install uv
|
||||
uses: astral-sh/setup-uv@v7
|
||||
with:
|
||||
enable-cache: true
|
||||
prune-cache: false
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v6
|
||||
with:
|
||||
python-version-file: ".python-version"
|
||||
|
||||
- name: Cache HuggingFace models
|
||||
uses: actions/cache@v5
|
||||
with:
|
||||
path: ~/.cache/huggingface
|
||||
key: ${{ runner.os }}-huggingface-${{ hashFiles('hindsight-api-slim/pyproject.toml') }}
|
||||
restore-keys: |
|
||||
${{ runner.os }}-huggingface-
|
||||
|
||||
- name: Pre-download models
|
||||
working-directory: ./hindsight-api-slim
|
||||
run: |
|
||||
uv run --frozen --all-extras --index-strategy unsafe-best-match python -c "
|
||||
from sentence_transformers import SentenceTransformer
|
||||
print('Downloading embedding model...')
|
||||
SentenceTransformer('BAAI/bge-small-en-v1.5')
|
||||
print('Model downloaded successfully')
|
||||
"
|
||||
|
||||
- name: Install hindsight-dev dependencies
|
||||
run: |
|
||||
cd hindsight-dev && uv sync --frozen --all-extras --index-strategy unsafe-best-match
|
||||
|
||||
- name: Run LoComo benchmark
|
||||
run: |
|
||||
MAX_CONV_ARG=""
|
||||
if [ "${{ inputs.locomo_max_conversations }}" != "0" ] && [ -n "${{ inputs.locomo_max_conversations }}" ]; then
|
||||
MAX_CONV_ARG="--max-conversations ${{ inputs.locomo_max_conversations }}"
|
||||
fi
|
||||
uv run python hindsight-dev/benchmarks/locomo/locomo_benchmark.py \
|
||||
--wait-consolidation \
|
||||
$MAX_CONV_ARG
|
||||
|
||||
- name: Upload LoComo results
|
||||
if: always()
|
||||
uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: locomo-results-${{ github.sha }}
|
||||
path: hindsight-dev/benchmarks/locomo/results/
|
||||
retention-days: 90
|
||||
+3
-231
@@ -49,7 +49,6 @@ jobs:
|
||||
integrations-opencode: ${{ steps.filter.outputs.integrations-opencode }}
|
||||
integrations-cloudflare-oauth-proxy: ${{ steps.filter.outputs.integrations-cloudflare-oauth-proxy }}
|
||||
integrations-lockfiles: ${{ steps.filter.outputs.integrations-lockfiles }}
|
||||
integrations-openai-agents: ${{ steps.filter.outputs.integrations-openai-agents }}
|
||||
dev: ${{ steps.filter.outputs.dev }}
|
||||
ci: ${{ steps.filter.outputs.ci }}
|
||||
# Secrets are available for internal PRs, pull_request_review, and workflow_dispatch.
|
||||
@@ -134,8 +133,6 @@ jobs:
|
||||
- 'hindsight-integrations/*/package-lock.json'
|
||||
- 'hindsight-integrations/*/package.json'
|
||||
- 'scripts/check-integration-lockfiles.sh'
|
||||
integrations-openai-agents:
|
||||
- 'hindsight-integrations/openai-agents/**'
|
||||
dev:
|
||||
- 'hindsight-dev/**'
|
||||
ci:
|
||||
@@ -2046,43 +2043,6 @@ jobs:
|
||||
working-directory: ./hindsight-integrations/llamaindex
|
||||
run: uv run pytest tests -v
|
||||
|
||||
test-openai-agents-integration:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
github.event_name != 'pull_request_review' &&
|
||||
(github.event_name == 'workflow_dispatch' ||
|
||||
needs.detect-changes.outputs.integrations-openai-agents == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ github.event.pull_request.head.sha || '' }}
|
||||
|
||||
- name: Install uv
|
||||
uses: astral-sh/setup-uv@v7
|
||||
with:
|
||||
enable-cache: true
|
||||
prune-cache: false
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v6
|
||||
with:
|
||||
python-version-file: ".python-version"
|
||||
|
||||
- name: Build openai-agents integration
|
||||
working-directory: ./hindsight-integrations/openai-agents
|
||||
run: uv build
|
||||
|
||||
- name: Install dependencies
|
||||
working-directory: ./hindsight-integrations/openai-agents
|
||||
run: uv sync --frozen
|
||||
|
||||
- name: Run tests
|
||||
working-directory: ./hindsight-integrations/openai-agents
|
||||
run: uv run pytest tests -v
|
||||
|
||||
test-pip-slim:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
@@ -2215,189 +2175,6 @@ jobs:
|
||||
working-directory: ./hindsight-embed
|
||||
run: ./test.sh
|
||||
|
||||
test-embed-windows:
|
||||
# Windows coverage for hindsight-embed. Runs the same unit tests + smoke
|
||||
# test as the Linux `test-embed` job, plus a `uv pip install --target`
|
||||
# sanity check that validates the sibling-binary resolution used by
|
||||
# users who install via `uv pip install hindsight-all` on Windows
|
||||
# (closes #1240).
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
needs.detect-changes.outputs.has_secrets == 'true' &&
|
||||
((github.event_name == 'pull_request_review' && github.event.review.state == 'approved') ||
|
||||
github.event_name == 'workflow_dispatch' ||
|
||||
needs.detect-changes.outputs.core == 'true' ||
|
||||
needs.detect-changes.outputs.embed == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: windows-latest
|
||||
env:
|
||||
HINDSIGHT_API_LLM_PROVIDER: vertexai
|
||||
HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY: ${{ github.workspace }}/gcp-credentials.json
|
||||
HINDSIGHT_API_LLM_MODEL: google/gemini-2.5-flash-lite
|
||||
# Force UTF-8 I/O so the CLI's ✓/box-drawing output doesn't crash the
|
||||
# default Windows cp1252 codec. Also applied at runtime via
|
||||
# sys.stdout.reconfigure in cli.py; this belt-and-suspenders covers
|
||||
# subprocesses the daemon spawns.
|
||||
PYTHONIOENCODING: utf-8
|
||||
PYTHONUTF8: "1"
|
||||
# pg0-embedded unpacks Postgres on first boot — noticeably slower on a
|
||||
# cold Windows runner than POSIX. Double the embed startup budget.
|
||||
HINDSIGHT_EMBED_DAEMON_STARTUP_TIMEOUT: "360"
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ github.event.pull_request.head.sha || '' }}
|
||||
|
||||
- name: Setup GCP credentials
|
||||
shell: bash
|
||||
run: |
|
||||
printf '%s' '${{ secrets.GCP_VERTEXAI_CREDENTIALS }}' > gcp-credentials.json
|
||||
PROJECT_ID=$(jq -r '.project_id' gcp-credentials.json)
|
||||
echo "HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID=$PROJECT_ID" >> "$GITHUB_ENV"
|
||||
|
||||
- name: Install uv
|
||||
uses: astral-sh/setup-uv@v7
|
||||
with:
|
||||
enable-cache: true
|
||||
prune-cache: false
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v6
|
||||
with:
|
||||
python-version-file: ".python-version"
|
||||
|
||||
- name: Install embed dependencies
|
||||
working-directory: ./hindsight-embed
|
||||
run: uv sync --frozen --index-strategy unsafe-best-match
|
||||
|
||||
- name: Install API dependencies (with local-ml and embedded-db for smoke test)
|
||||
working-directory: ./hindsight-api-slim
|
||||
run: uv sync --frozen --all-extras --index-strategy unsafe-best-match
|
||||
|
||||
- name: Cache HuggingFace models
|
||||
uses: actions/cache@v5
|
||||
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 unit and integration tests
|
||||
working-directory: ./hindsight-embed
|
||||
run: uv run pytest tests/ -v
|
||||
|
||||
# Smoke test's retain/recall commands delegate to the Rust hindsight CLI.
|
||||
# On POSIX, hindsight-embed auto-installs the CLI via curl|bash; on
|
||||
# Windows that installer isn't available (and `bash` on windows-latest
|
||||
# routes to WSL which isn't provisioned). Build the CLI from source and
|
||||
# drop it into ~/.local/bin where find_cli_binary() looks first.
|
||||
- name: Install Rust
|
||||
uses: dtolnay/rust-toolchain@stable
|
||||
|
||||
- name: Cache cargo build
|
||||
uses: actions/cache@v5
|
||||
with:
|
||||
path: |
|
||||
~/.cargo/registry
|
||||
~/.cargo/git
|
||||
hindsight-cli/target
|
||||
key: ${{ runner.os }}-cargo-embed-smoke-${{ hashFiles('**/Cargo.lock') }}
|
||||
restore-keys: |
|
||||
${{ runner.os }}-cargo-embed-smoke-
|
||||
${{ runner.os }}-cargo-
|
||||
|
||||
- name: Build hindsight CLI
|
||||
working-directory: ./hindsight-cli
|
||||
run: cargo build --release
|
||||
|
||||
- name: Stage hindsight CLI where find_cli_binary expects it
|
||||
shell: bash
|
||||
run: |
|
||||
set -euo pipefail
|
||||
install_dir="$HOME/.local/bin"
|
||||
mkdir -p "$install_dir"
|
||||
cp hindsight-cli/target/release/hindsight.exe "$install_dir/hindsight.exe"
|
||||
"$install_dir/hindsight.exe" --version
|
||||
|
||||
- name: Run smoke test
|
||||
shell: bash
|
||||
working-directory: ./hindsight-embed
|
||||
run: ./test.sh
|
||||
|
||||
# Real-world install test for issue #1240: drop both packages into a
|
||||
# --target directory (the layout you get from `uv pip install hindsight-all`
|
||||
# or NixOS) and verify the sibling binary is discovered (not the uvx
|
||||
# fallback). Exercises a different code path than the smoke test, which
|
||||
# uses `uv run --project` via the monorepo branch of _find_api_command.
|
||||
#
|
||||
# IMPORTANT: install outside the repo checkout. `_find_api_command` first
|
||||
# probes `<pkg>/../../hindsight-api-slim` for dev mode; if the target dir
|
||||
# lives inside the monorepo, that branch matches and we never exercise
|
||||
# the sibling-binary path we actually want to test.
|
||||
- name: Install hindsight-embed and hindsight-api into --target directory
|
||||
shell: bash
|
||||
run: |
|
||||
set -euo pipefail
|
||||
target="$RUNNER_TEMP/install-test"
|
||||
rm -rf "$target"
|
||||
mkdir -p "$target"
|
||||
uv pip install --target "$target" ./hindsight-embed ./hindsight-api-slim
|
||||
|
||||
- name: Verify sibling hindsight-api.exe is present
|
||||
shell: bash
|
||||
run: |
|
||||
set -euo pipefail
|
||||
target="$RUNNER_TEMP/install-test"
|
||||
if [ -f "$target/Scripts/hindsight-api.exe" ]; then
|
||||
echo "Found $target/Scripts/hindsight-api.exe"
|
||||
elif [ -f "$target/bin/hindsight-api.exe" ]; then
|
||||
echo "Found $target/bin/hindsight-api.exe"
|
||||
else
|
||||
echo "::error::hindsight-api.exe not found in install target"
|
||||
ls "$target/"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
- name: Verify _find_api_command resolves the sibling binary (not uvx)
|
||||
shell: bash
|
||||
run: |
|
||||
set -euo pipefail
|
||||
target="$RUNNER_TEMP/install-test"
|
||||
PYTHONPATH="$target" python -c "
|
||||
from hindsight_embed.daemon_embed_manager import DaemonEmbedManager
|
||||
cmd = DaemonEmbedManager()._find_api_command()
|
||||
print('Resolved command:', cmd)
|
||||
assert len(cmd) == 1 and cmd[0].endswith('hindsight-api.exe'), (
|
||||
f'Expected sibling hindsight-api.exe, got {cmd!r}. '
|
||||
'Falling back to uvx on --target installs reintroduces issue #1240.'
|
||||
)
|
||||
"
|
||||
|
||||
- name: Smoke-check installed hindsight-embed binary runs
|
||||
shell: bash
|
||||
run: |
|
||||
set -euo pipefail
|
||||
target="$RUNNER_TEMP/install-test"
|
||||
export PYTHONPATH="$target"
|
||||
if [ -f "$target/Scripts/hindsight-embed.exe" ]; then
|
||||
"$target/Scripts/hindsight-embed.exe" --help
|
||||
else
|
||||
"$target/bin/hindsight-embed.exe" --help
|
||||
fi
|
||||
|
||||
- name: Collect daemon logs on failure
|
||||
if: failure()
|
||||
shell: bash
|
||||
run: |
|
||||
for f in ~/.hindsight/daemon.log ~/.hindsight/profiles/*.log ~/.hindsight/profiles/*.stderr.log; do
|
||||
if [ -f "$f" ]; then
|
||||
echo "=== $f ==="
|
||||
cat "$f"
|
||||
fi
|
||||
done || true
|
||||
|
||||
test-hindsight-all:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
@@ -2770,9 +2547,6 @@ jobs:
|
||||
- name: Run generate-openapi
|
||||
run: ./scripts/generate-openapi.sh
|
||||
|
||||
- name: Run generate-bank-template-schema
|
||||
run: ./scripts/generate-bank-template-schema.sh
|
||||
|
||||
- name: Run generate-clients
|
||||
run: ./scripts/generate-clients.sh
|
||||
|
||||
@@ -2792,7 +2566,6 @@ jobs:
|
||||
echo ""
|
||||
echo "Please run the following commands locally and commit the changes:"
|
||||
echo " ./scripts/generate-openapi.sh"
|
||||
echo " ./scripts/generate-bank-template-schema.sh"
|
||||
echo " ./scripts/generate-clients.sh"
|
||||
echo " ./scripts/generate-docs-skill.sh"
|
||||
echo " ./scripts/hooks/lint.sh"
|
||||
@@ -2929,7 +2702,6 @@ jobs:
|
||||
- test-llamaindex-integration
|
||||
- test-pip-slim
|
||||
- test-embed
|
||||
- test-embed-windows
|
||||
- test-hindsight-all
|
||||
- test-doc-examples
|
||||
- test-upgrade
|
||||
@@ -2943,7 +2715,7 @@ jobs:
|
||||
steps:
|
||||
- name: Determine overall result
|
||||
id: result
|
||||
uses: actions/github-script@v9
|
||||
uses: actions/github-script@v8
|
||||
with:
|
||||
script: |
|
||||
const needs = ${{ toJSON(needs) }};
|
||||
@@ -2976,7 +2748,7 @@ jobs:
|
||||
core.setOutput('run_url', runUrl);
|
||||
|
||||
- name: Report status to PR
|
||||
uses: actions/github-script@v9
|
||||
uses: actions/github-script@v8
|
||||
with:
|
||||
script: |
|
||||
await github.rest.repos.createCommitStatus({
|
||||
@@ -2990,7 +2762,7 @@ jobs:
|
||||
});
|
||||
|
||||
- name: Comment on PR
|
||||
uses: actions/github-script@v9
|
||||
uses: actions/github-script@v8
|
||||
with:
|
||||
script: |
|
||||
const prNumber = context.payload.pull_request.number;
|
||||
|
||||
@@ -1,7 +0,0 @@
|
||||
{
|
||||
"semi": true,
|
||||
"singleQuote": false,
|
||||
"tabWidth": 2,
|
||||
"trailingComma": "es5",
|
||||
"printWidth": 100
|
||||
}
|
||||
@@ -68,9 +68,8 @@ cd hindsight-control-plane && npm run dev
|
||||
./scripts/benchmarks/run-locomo.sh
|
||||
|
||||
# Performance benchmarks
|
||||
./scripts/benchmarks/run-perf-test.sh # System perf (mock LLM + pg0)
|
||||
./scripts/benchmarks/run-perf-test.sh --scale tiny # Quick smoke test
|
||||
./scripts/benchmarks/run-consolidation.sh
|
||||
./scripts/benchmarks/run-retain-perf.sh --document <path> # Requires API server running
|
||||
|
||||
# Results viewer
|
||||
./scripts/benchmarks/start-visualizer.sh # View results at localhost:8001
|
||||
|
||||
@@ -2,8 +2,8 @@ apiVersion: v2
|
||||
name: hindsight
|
||||
description: Hindsight helm chart
|
||||
type: application
|
||||
version: 0.5.4
|
||||
appVersion: "0.5.4"
|
||||
version: 0.5.1
|
||||
appVersion: "0.5.1"
|
||||
keywords:
|
||||
- ai
|
||||
- memory
|
||||
|
||||
+10
-10
@@ -18,17 +18,17 @@ npm install @vectorize-io/hindsight-all @vectorize-io/hindsight-client
|
||||
## Example
|
||||
|
||||
```ts
|
||||
import { HindsightServer, consoleLogger } from "@vectorize-io/hindsight-all";
|
||||
import { HindsightClient } from "@vectorize-io/hindsight-client";
|
||||
import { HindsightServer, consoleLogger } from '@vectorize-io/hindsight-all';
|
||||
import { HindsightClient } from '@vectorize-io/hindsight-client';
|
||||
|
||||
const server = new HindsightServer({
|
||||
profile: "my-app",
|
||||
profile: 'my-app',
|
||||
port: 9077,
|
||||
env: {
|
||||
HINDSIGHT_API_LLM_PROVIDER: "anthropic",
|
||||
HINDSIGHT_API_LLM_PROVIDER: 'anthropic',
|
||||
HINDSIGHT_API_LLM_API_KEY: process.env.ANTHROPIC_API_KEY,
|
||||
HINDSIGHT_API_LLM_MODEL: "claude-sonnet-4-20250514",
|
||||
HINDSIGHT_EMBED_DAEMON_IDLE_TIMEOUT: "0",
|
||||
HINDSIGHT_API_LLM_MODEL: 'claude-sonnet-4-20250514',
|
||||
HINDSIGHT_EMBED_DAEMON_IDLE_TIMEOUT: '0',
|
||||
},
|
||||
logger: consoleLogger,
|
||||
});
|
||||
@@ -37,11 +37,11 @@ await server.start();
|
||||
|
||||
const client = new HindsightClient({ baseUrl: server.getBaseUrl() });
|
||||
|
||||
await client.retain("user-123", "User prefers dark mode and concise answers.", {
|
||||
documentId: "pref-2026-04-01",
|
||||
await client.retain('user-123', 'User prefers dark mode and concise answers.', {
|
||||
documentId: 'pref-2026-04-01',
|
||||
});
|
||||
|
||||
const recall = await client.recall("user-123", "what are the user preferences?");
|
||||
const recall = await client.recall('user-123', 'what are the user preferences?');
|
||||
console.log(recall.results);
|
||||
|
||||
await server.stop();
|
||||
@@ -62,7 +62,7 @@ If you're hacking on the Python `hindsight-embed` package in the same monorepo,
|
||||
|
||||
```ts
|
||||
new HindsightServer({
|
||||
embedPackagePath: "/path/to/hindsight-embed",
|
||||
embedPackagePath: '/path/to/hindsight-embed',
|
||||
// ...
|
||||
});
|
||||
```
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@vectorize-io/hindsight-all",
|
||||
"version": "0.5.4",
|
||||
"version": "0.5.1",
|
||||
"description": "Node.js programmatic lifecycle manager for Hindsight — embeds a local hindsight daemon in a Node application. Pair with @vectorize-io/hindsight-client for memory operations.",
|
||||
"main": "dist/index.js",
|
||||
"types": "dist/index.d.ts",
|
||||
|
||||
@@ -1,36 +1,32 @@
|
||||
import { describe, it, expect } from "vitest";
|
||||
import { getEmbedCommand } from "./command.js";
|
||||
import { describe, it, expect } from 'vitest';
|
||||
import { getEmbedCommand } from './command.js';
|
||||
|
||||
describe("getEmbedCommand", () => {
|
||||
it("defaults to uvx hindsight-embed@latest", () => {
|
||||
expect(getEmbedCommand()).toEqual(["uvx", "hindsight-embed@latest"]);
|
||||
describe('getEmbedCommand', () => {
|
||||
it('defaults to uvx hindsight-embed@latest', () => {
|
||||
expect(getEmbedCommand()).toEqual(['uvx', 'hindsight-embed@latest']);
|
||||
});
|
||||
|
||||
it("honours an explicit version", () => {
|
||||
expect(getEmbedCommand({ embedVersion: "0.5.0" })).toEqual(["uvx", "[email protected]"]);
|
||||
it('honours an explicit version', () => {
|
||||
expect(getEmbedCommand({ embedVersion: '0.5.0' })).toEqual(['uvx', '[email protected]']);
|
||||
});
|
||||
|
||||
it("treats an empty version as latest", () => {
|
||||
expect(getEmbedCommand({ embedVersion: "" })).toEqual(["uvx", "hindsight-embed@latest"]);
|
||||
it('treats an empty version as latest', () => {
|
||||
expect(getEmbedCommand({ embedVersion: '' })).toEqual(['uvx', 'hindsight-embed@latest']);
|
||||
});
|
||||
|
||||
it("uses uv run --directory when a local path is given", () => {
|
||||
expect(getEmbedCommand({ embedPackagePath: "/abs/path" })).toEqual([
|
||||
"uv",
|
||||
"run",
|
||||
"--directory",
|
||||
"/abs/path",
|
||||
"hindsight-embed",
|
||||
it('uses uv run --directory when a local path is given', () => {
|
||||
expect(getEmbedCommand({ embedPackagePath: '/abs/path' })).toEqual([
|
||||
'uv',
|
||||
'run',
|
||||
'--directory',
|
||||
'/abs/path',
|
||||
'hindsight-embed',
|
||||
]);
|
||||
});
|
||||
|
||||
it("local path takes precedence over version", () => {
|
||||
expect(getEmbedCommand({ embedPackagePath: "/abs/path", embedVersion: "0.5.0" })).toEqual([
|
||||
"uv",
|
||||
"run",
|
||||
"--directory",
|
||||
"/abs/path",
|
||||
"hindsight-embed",
|
||||
]);
|
||||
it('local path takes precedence over version', () => {
|
||||
expect(
|
||||
getEmbedCommand({ embedPackagePath: '/abs/path', embedVersion: '0.5.0' }),
|
||||
).toEqual(['uv', 'run', '--directory', '/abs/path', 'hindsight-embed']);
|
||||
});
|
||||
});
|
||||
|
||||
@@ -18,8 +18,8 @@ export interface EmbedCommandOptions {
|
||||
|
||||
export function getEmbedCommand(opts: EmbedCommandOptions = {}): string[] {
|
||||
if (opts.embedPackagePath) {
|
||||
return ["uv", "run", "--directory", opts.embedPackagePath, "hindsight-embed"];
|
||||
return ['uv', 'run', '--directory', opts.embedPackagePath, 'hindsight-embed'];
|
||||
}
|
||||
const version = opts.embedVersion && opts.embedVersion.length > 0 ? opts.embedVersion : "latest";
|
||||
return ["uvx", `hindsight-embed@${version}`];
|
||||
const version = opts.embedVersion && opts.embedVersion.length > 0 ? opts.embedVersion : 'latest';
|
||||
return ['uvx', `hindsight-embed@${version}`];
|
||||
}
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
export { HindsightServer } from "./server.js";
|
||||
export { getEmbedCommand } from "./command.js";
|
||||
export { silentLogger, consoleLogger } from "./logger.js";
|
||||
export { HindsightServer } from './server.js';
|
||||
export { getEmbedCommand } from './command.js';
|
||||
export { silentLogger, consoleLogger } from './logger.js';
|
||||
|
||||
export type { Logger } from "./logger.js";
|
||||
export type { EmbedCommandOptions } from "./command.js";
|
||||
export type { HindsightServerOptions } from "./types.js";
|
||||
export type { Logger } from './logger.js';
|
||||
export type { EmbedCommandOptions } from './command.js';
|
||||
export type { HindsightServerOptions } from './types.js';
|
||||
|
||||
@@ -1,32 +1,32 @@
|
||||
import { describe, it, expect } from "vitest";
|
||||
import { HindsightServer } from "./server.js";
|
||||
import { describe, it, expect } from 'vitest';
|
||||
import { HindsightServer } from './server.js';
|
||||
|
||||
describe("HindsightServer construction", () => {
|
||||
it("defaults base URL to http://127.0.0.1:8888", () => {
|
||||
describe('HindsightServer construction', () => {
|
||||
it('defaults base URL to http://127.0.0.1:8888', () => {
|
||||
const server = new HindsightServer();
|
||||
expect(server.getBaseUrl()).toBe("http://127.0.0.1:8888");
|
||||
expect(server.getProfile()).toBe("default");
|
||||
expect(server.getBaseUrl()).toBe('http://127.0.0.1:8888');
|
||||
expect(server.getProfile()).toBe('default');
|
||||
});
|
||||
|
||||
it("honours custom profile, port, and host", () => {
|
||||
const server = new HindsightServer({ profile: "app", port: 9077, host: "0.0.0.0" });
|
||||
expect(server.getProfile()).toBe("app");
|
||||
expect(server.getBaseUrl()).toBe("http://0.0.0.0:9077");
|
||||
it('honours custom profile, port, and host', () => {
|
||||
const server = new HindsightServer({ profile: 'app', port: 9077, host: '0.0.0.0' });
|
||||
expect(server.getProfile()).toBe('app');
|
||||
expect(server.getBaseUrl()).toBe('http://0.0.0.0:9077');
|
||||
});
|
||||
|
||||
it("accepts open env pass-through without complaining about unknown keys", () => {
|
||||
it('accepts open env pass-through without complaining about unknown keys', () => {
|
||||
const server = new HindsightServer({
|
||||
env: {
|
||||
HINDSIGHT_API_LLM_PROVIDER: "openai",
|
||||
HINDSIGHT_API_LLM_MODEL: "gpt-4o-mini",
|
||||
HINDSIGHT_API_LLM_PROVIDER: 'openai',
|
||||
HINDSIGHT_API_LLM_MODEL: 'gpt-4o-mini',
|
||||
// A field that does not exist today — should still be accepted
|
||||
HINDSIGHT_FUTURE_FLAG: "enabled",
|
||||
HINDSIGHT_FUTURE_FLAG: 'enabled',
|
||||
},
|
||||
});
|
||||
expect(server).toBeInstanceOf(HindsightServer);
|
||||
});
|
||||
|
||||
it("exposes checkHealth that returns false when no daemon is running", async () => {
|
||||
it('exposes checkHealth that returns false when no daemon is running', async () => {
|
||||
// Random high port that nothing is listening on.
|
||||
const server = new HindsightServer({ port: 1, readyTimeoutMs: 100 });
|
||||
const healthy = await server.checkHealth();
|
||||
|
||||
@@ -1,12 +1,12 @@
|
||||
import { spawn } from "child_process";
|
||||
import { getEmbedCommand } from "./command.js";
|
||||
import { silentLogger } from "./logger.js";
|
||||
import type { Logger } from "./logger.js";
|
||||
import type { HindsightServerOptions } from "./types.js";
|
||||
import { spawn } from 'child_process';
|
||||
import { getEmbedCommand } from './command.js';
|
||||
import { silentLogger } from './logger.js';
|
||||
import type { Logger } from './logger.js';
|
||||
import type { HindsightServerOptions } from './types.js';
|
||||
|
||||
const DEFAULT_PORT = 8888;
|
||||
const DEFAULT_HOST = "127.0.0.1";
|
||||
const DEFAULT_PROFILE = "default";
|
||||
const DEFAULT_HOST = '127.0.0.1';
|
||||
const DEFAULT_PROFILE = 'default';
|
||||
const DEFAULT_READY_TIMEOUT_MS = 30_000;
|
||||
const DEFAULT_READY_POLL_INTERVAL_MS = 1_000;
|
||||
|
||||
@@ -61,7 +61,7 @@ export class HindsightServer {
|
||||
this.userEnv = opts.env ?? {};
|
||||
this.extraProfileCreateArgs = opts.extraProfileCreateArgs ?? [];
|
||||
this.extraDaemonStartArgs = opts.extraDaemonStartArgs ?? [];
|
||||
this.platformCpuWorkaround = opts.platformCpuWorkaround ?? process.platform === "darwin";
|
||||
this.platformCpuWorkaround = opts.platformCpuWorkaround ?? (process.platform === 'darwin');
|
||||
this.readyTimeoutMs = opts.readyTimeoutMs ?? DEFAULT_READY_TIMEOUT_MS;
|
||||
this.readyPollIntervalMs = opts.readyPollIntervalMs ?? DEFAULT_READY_POLL_INTERVAL_MS;
|
||||
this.logger = opts.logger ?? silentLogger;
|
||||
@@ -100,22 +100,22 @@ export class HindsightServer {
|
||||
embedVersion: this.embedVersion,
|
||||
embedPackagePath: this.embedPackagePath,
|
||||
});
|
||||
const args = [...baseArgs, "daemon", "--profile", this.profile, "stop"];
|
||||
const args = [...baseArgs, 'daemon', '--profile', this.profile, 'stop'];
|
||||
|
||||
const child = spawn(cmd, args, { stdio: "pipe" });
|
||||
this.pipeOutput(child, "daemon.stop");
|
||||
const child = spawn(cmd, args, { stdio: 'pipe' });
|
||||
this.pipeOutput(child, 'daemon.stop');
|
||||
|
||||
await new Promise<void>((resolve) => {
|
||||
const timeout = setTimeout(() => {
|
||||
this.logger.warn(`[hindsight] daemon stop timed out after 5s`);
|
||||
resolve();
|
||||
}, 5_000);
|
||||
child.on("exit", () => {
|
||||
child.on('exit', () => {
|
||||
clearTimeout(timeout);
|
||||
this.logger.info(`[hindsight] daemon stopped`);
|
||||
resolve();
|
||||
});
|
||||
child.on("error", (err) => {
|
||||
child.on('error', (err) => {
|
||||
clearTimeout(timeout);
|
||||
this.logger.warn(`[hindsight] error stopping daemon: ${err.message}`);
|
||||
resolve();
|
||||
@@ -147,9 +147,9 @@ export class HindsightServer {
|
||||
private buildEnv(): NodeJS.ProcessEnv {
|
||||
const merged: NodeJS.ProcessEnv = { ...process.env };
|
||||
|
||||
if (this.platformCpuWorkaround && process.platform === "darwin") {
|
||||
merged["HINDSIGHT_API_EMBEDDINGS_LOCAL_FORCE_CPU"] = "1";
|
||||
merged["HINDSIGHT_API_RERANKER_LOCAL_FORCE_CPU"] = "1";
|
||||
if (this.platformCpuWorkaround && process.platform === 'darwin') {
|
||||
merged['HINDSIGHT_API_EMBEDDINGS_LOCAL_FORCE_CPU'] = '1';
|
||||
merged['HINDSIGHT_API_RERANKER_LOCAL_FORCE_CPU'] = '1';
|
||||
}
|
||||
|
||||
for (const [key, value] of Object.entries(this.userEnv)) {
|
||||
@@ -175,11 +175,11 @@ export class HindsightServer {
|
||||
});
|
||||
const createArgs = [
|
||||
...baseArgs,
|
||||
"profile",
|
||||
"create",
|
||||
'profile',
|
||||
'create',
|
||||
this.profile,
|
||||
"--merge",
|
||||
"--port",
|
||||
'--merge',
|
||||
'--port',
|
||||
String(this.port),
|
||||
];
|
||||
|
||||
@@ -189,12 +189,12 @@ export class HindsightServer {
|
||||
// host state into profile config.
|
||||
const envForProfile = this.collectProfileEnv(env);
|
||||
for (const [key, value] of Object.entries(envForProfile)) {
|
||||
createArgs.push("--env", `${key}=${value}`);
|
||||
createArgs.push('--env', `${key}=${value}`);
|
||||
}
|
||||
|
||||
createArgs.push(...this.extraProfileCreateArgs);
|
||||
|
||||
await this.runCommand(cmd, createArgs, env, "profile.create");
|
||||
await this.runCommand(cmd, createArgs, env, 'profile.create');
|
||||
}
|
||||
|
||||
/** Collect only the env vars that should be written into the profile file. */
|
||||
@@ -209,10 +209,10 @@ export class HindsightServer {
|
||||
}
|
||||
|
||||
// 2. CPU workaround — only if auto-applied and not already overridden.
|
||||
if (this.platformCpuWorkaround && process.platform === "darwin") {
|
||||
if (this.platformCpuWorkaround && process.platform === 'darwin') {
|
||||
const cpuKeys = [
|
||||
"HINDSIGHT_API_EMBEDDINGS_LOCAL_FORCE_CPU",
|
||||
"HINDSIGHT_API_RERANKER_LOCAL_FORCE_CPU",
|
||||
'HINDSIGHT_API_EMBEDDINGS_LOCAL_FORCE_CPU',
|
||||
'HINDSIGHT_API_RERANKER_LOCAL_FORCE_CPU',
|
||||
];
|
||||
for (const key of cpuKeys) {
|
||||
if (!(key in out) && env[key] !== undefined) {
|
||||
@@ -231,14 +231,14 @@ export class HindsightServer {
|
||||
});
|
||||
const args = [
|
||||
...baseArgs,
|
||||
"daemon",
|
||||
"--profile",
|
||||
'daemon',
|
||||
'--profile',
|
||||
this.profile,
|
||||
"start",
|
||||
'start',
|
||||
...this.extraDaemonStartArgs,
|
||||
];
|
||||
|
||||
await this.runCommand(cmd, args, env, "daemon.start");
|
||||
await this.runCommand(cmd, args, env, 'daemon.start');
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -249,34 +249,34 @@ export class HindsightServer {
|
||||
cmd: string,
|
||||
args: string[],
|
||||
env: NodeJS.ProcessEnv,
|
||||
label: string
|
||||
label: string,
|
||||
): Promise<void> {
|
||||
const child = spawn(cmd, args, { stdio: "pipe", env });
|
||||
let output = "";
|
||||
child.stdout?.on("data", (data: Buffer) => {
|
||||
const child = spawn(cmd, args, { stdio: 'pipe', env });
|
||||
let output = '';
|
||||
child.stdout?.on('data', (data: Buffer) => {
|
||||
const text = data.toString();
|
||||
output += text;
|
||||
for (const line of text.trimEnd().split("\n")) {
|
||||
for (const line of text.trimEnd().split('\n')) {
|
||||
if (line) this.logger.info(`[hindsight:${label}] ${line}`);
|
||||
}
|
||||
});
|
||||
child.stderr?.on("data", (data: Buffer) => {
|
||||
child.stderr?.on('data', (data: Buffer) => {
|
||||
const text = data.toString();
|
||||
output += text;
|
||||
for (const line of text.trimEnd().split("\n")) {
|
||||
for (const line of text.trimEnd().split('\n')) {
|
||||
if (line) this.logger.warn(`[hindsight:${label}] ${line}`);
|
||||
}
|
||||
});
|
||||
|
||||
await new Promise<void>((resolve, reject) => {
|
||||
child.on("exit", (code) => {
|
||||
child.on('exit', (code) => {
|
||||
if (code === 0) {
|
||||
resolve();
|
||||
} else {
|
||||
reject(new Error(`${label} failed with code ${code}: ${output.trim()}`));
|
||||
}
|
||||
});
|
||||
child.on("error", (err) => {
|
||||
child.on('error', (err) => {
|
||||
reject(new Error(`${label} failed to spawn: ${err.message}`, { cause: err }));
|
||||
});
|
||||
});
|
||||
@@ -284,13 +284,13 @@ export class HindsightServer {
|
||||
|
||||
/** Stream a spawned child's stdout/stderr through the logger without blocking. */
|
||||
private pipeOutput(child: ReturnType<typeof spawn>, label: string): void {
|
||||
child.stdout?.on("data", (data: Buffer) => {
|
||||
for (const line of data.toString().trimEnd().split("\n")) {
|
||||
child.stdout?.on('data', (data: Buffer) => {
|
||||
for (const line of data.toString().trimEnd().split('\n')) {
|
||||
if (line) this.logger.info(`[hindsight:${label}] ${line}`);
|
||||
}
|
||||
});
|
||||
child.stderr?.on("data", (data: Buffer) => {
|
||||
for (const line of data.toString().trimEnd().split("\n")) {
|
||||
child.stderr?.on('data', (data: Buffer) => {
|
||||
for (const line of data.toString().trimEnd().split('\n')) {
|
||||
if (line) this.logger.warn(`[hindsight:${label}] ${line}`);
|
||||
}
|
||||
});
|
||||
@@ -316,7 +316,7 @@ export class HindsightServer {
|
||||
await new Promise((resolve) => setTimeout(resolve, this.readyPollIntervalMs));
|
||||
}
|
||||
throw new Error(
|
||||
`Hindsight daemon did not become ready within ${this.readyTimeoutMs}ms at ${this.baseUrl}`
|
||||
`Hindsight daemon did not become ready within ${this.readyTimeoutMs}ms at ${this.baseUrl}`,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
import type { Logger } from "./logger.js";
|
||||
import type { Logger } from './logger.js';
|
||||
|
||||
/**
|
||||
* Options for {@link HindsightServer}.
|
||||
|
||||
@@ -1,10 +1,10 @@
|
||||
import { defineConfig } from "tsup";
|
||||
import { defineConfig } from 'tsup';
|
||||
|
||||
export default defineConfig({
|
||||
entry: ["src/index.ts"],
|
||||
format: ["esm"],
|
||||
entry: ['src/index.ts'],
|
||||
format: ['esm'],
|
||||
dts: true,
|
||||
outDir: "dist",
|
||||
outDir: 'dist',
|
||||
clean: true,
|
||||
sourcemap: true,
|
||||
bundle: true,
|
||||
|
||||
@@ -1,8 +1,8 @@
|
||||
import { defineConfig } from "vitest/config";
|
||||
import { defineConfig } from 'vitest/config';
|
||||
|
||||
export default defineConfig({
|
||||
test: {
|
||||
include: ["src/**/*.test.ts"],
|
||||
environment: "node",
|
||||
include: ['src/**/*.test.ts'],
|
||||
environment: 'node',
|
||||
},
|
||||
});
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "hindsight-all-slim"
|
||||
version = "0.5.4"
|
||||
version = "0.5.1"
|
||||
description = "Hindsight: Agent Memory That Works Like Human Memory - Slim All-in-One Bundle"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.11"
|
||||
|
||||
@@ -71,7 +71,7 @@ class HindsightEmbedded:
|
||||
llm_model: Model name to use
|
||||
llm_base_url: Optional custom base URL for LLM API
|
||||
database_url: Optional database URL override (default: profile-specific pg0)
|
||||
idle_timeout: Seconds before daemon auto-exits when idle (default: 0, disabled)
|
||||
idle_timeout: Seconds before daemon auto-exits when idle (default: 300)
|
||||
log_level: Daemon log level (default: "info")
|
||||
ui: Whether to start the control plane web UI alongside the daemon (default: False)
|
||||
ui_port: Port for the UI. Defaults to daemon_port + 10000.
|
||||
@@ -86,7 +86,7 @@ class HindsightEmbedded:
|
||||
llm_model: str = "openai/gpt-oss-120b",
|
||||
llm_base_url: Optional[str] = None,
|
||||
database_url: Optional[str] = None,
|
||||
idle_timeout: int = 0,
|
||||
idle_timeout: int = 300,
|
||||
log_level: str = "info",
|
||||
ui: bool = False,
|
||||
ui_port: Optional[int] = None,
|
||||
@@ -102,7 +102,7 @@ class HindsightEmbedded:
|
||||
llm_model: Model name to use
|
||||
llm_base_url: Optional custom base URL for LLM API
|
||||
database_url: Optional database URL override
|
||||
idle_timeout: Seconds before daemon auto-exits when idle (0 = disabled)
|
||||
idle_timeout: Seconds before daemon auto-exits when idle
|
||||
log_level: Daemon log level
|
||||
ui: Whether to start the control plane web UI alongside the daemon
|
||||
ui_port: Port for the UI (defaults to daemon_port + 10000)
|
||||
@@ -142,37 +142,14 @@ class HindsightEmbedded:
|
||||
self._memories_api: Optional[MemoriesAPI] = None
|
||||
|
||||
def _ensure_started(self):
|
||||
"""Ensure daemon is running (thread-safe), restarting if crashed."""
|
||||
"""Ensure daemon is running (thread-safe)."""
|
||||
if self._started and self._client is not None:
|
||||
if self._manager.is_running(self.profile):
|
||||
return
|
||||
# Daemon crashed — reset state and fall through to restart
|
||||
logger.warning(
|
||||
"Daemon for profile '%s' is no longer responsive, restarting...",
|
||||
self.profile,
|
||||
)
|
||||
try:
|
||||
self._client.close()
|
||||
except Exception:
|
||||
logger.debug("Error closing stale client", exc_info=True)
|
||||
self._client = None
|
||||
self._started = False
|
||||
return
|
||||
|
||||
with self._lock:
|
||||
# Double-check after acquiring lock
|
||||
if self._started and self._client is not None:
|
||||
if self._manager.is_running(self.profile):
|
||||
return
|
||||
logger.warning(
|
||||
"Daemon for profile '%s' is no longer responsive (lock path), restarting...",
|
||||
self.profile,
|
||||
)
|
||||
try:
|
||||
self._client.close()
|
||||
except Exception:
|
||||
logger.debug("Error closing stale client", exc_info=True)
|
||||
self._client = None
|
||||
self._started = False
|
||||
return
|
||||
|
||||
if self._closed:
|
||||
raise RuntimeError(
|
||||
@@ -276,10 +253,23 @@ class HindsightEmbedded:
|
||||
This allows HindsightEmbedded to expose all HindsightClient methods
|
||||
without manually wrapping each one.
|
||||
"""
|
||||
# Ensure server is started (and restart if crashed) before proxying
|
||||
# Ensure server is started before proxying
|
||||
self._ensure_started()
|
||||
|
||||
return getattr(self._client, name)
|
||||
# Get the attribute from the underlying client
|
||||
attr = getattr(self._client, name)
|
||||
|
||||
# If it's a callable, wrap it to ensure server is started
|
||||
# (shouldn't be needed since _ensure_started already called, but defensive)
|
||||
if callable(attr):
|
||||
|
||||
def wrapper(*args, **kwargs):
|
||||
self._ensure_started()
|
||||
return attr(*args, **kwargs)
|
||||
|
||||
return wrapper
|
||||
|
||||
return attr
|
||||
|
||||
def __enter__(self):
|
||||
"""Context manager entry - ensures server is started."""
|
||||
@@ -404,8 +394,11 @@ class HindsightEmbedded:
|
||||
"""
|
||||
Get the underlying Hindsight client for direct access.
|
||||
|
||||
Ensures daemon is started (and restarts it if it has crashed) before
|
||||
returning the client.
|
||||
WARNING: Using this property directly means daemon restarts won't be
|
||||
handled automatically. Prefer using the API namespaces (banks, mental_models,
|
||||
directives, memories) or direct method calls on HindsightEmbedded instead.
|
||||
|
||||
Ensures daemon is started before returning the client.
|
||||
|
||||
Returns:
|
||||
Hindsight: The underlying client instance
|
||||
@@ -416,8 +409,9 @@ class HindsightEmbedded:
|
||||
|
||||
embedded = HindsightEmbedded(profile="myapp", ...)
|
||||
|
||||
# Direct access (not recommended - daemon crashes won't be handled)
|
||||
client = embedded.client
|
||||
banks = client.list_banks()
|
||||
banks = client.list_banks() # If daemon crashes, this will fail
|
||||
```
|
||||
"""
|
||||
self._ensure_started()
|
||||
@@ -431,13 +425,8 @@ class HindsightEmbedded:
|
||||
|
||||
@property
|
||||
def is_running(self) -> bool:
|
||||
"""Check if the client is initialized and the daemon is responsive."""
|
||||
return (
|
||||
self._started
|
||||
and not self._closed
|
||||
and self._client is not None
|
||||
and self._manager.is_running(self.profile)
|
||||
)
|
||||
"""Check if the client is initialized."""
|
||||
return self._started and not self._closed and self._client is not None
|
||||
|
||||
@property
|
||||
def ui_url(self) -> str:
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "hindsight-all"
|
||||
version = "0.5.4"
|
||||
version = "0.5.1"
|
||||
description = "Hindsight: Agent Memory That Works Like Human Memory - All-in-One Bundle"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.11"
|
||||
|
||||
@@ -401,42 +401,3 @@ def test_embedded_ui_flag(llm_config):
|
||||
|
||||
finally:
|
||||
client.close()
|
||||
|
||||
|
||||
def test_embedded_daemon_crash_recovery(llm_config):
|
||||
"""
|
||||
Test that HindsightEmbedded recovers when the daemon crashes.
|
||||
|
||||
Simulates a crash by stopping the daemon, then verifies
|
||||
that the next operation transparently restarts it.
|
||||
"""
|
||||
profile = f"test_crash_{uuid.uuid4().hex[:8]}"
|
||||
bank_id = f"bank_{uuid.uuid4().hex[:8]}"
|
||||
|
||||
client = HindsightEmbedded(profile=profile, log_level="info", **llm_config)
|
||||
|
||||
try:
|
||||
# Start daemon and store a memory
|
||||
result = client.retain(bank_id=bank_id, content="Before crash")
|
||||
assert result.success, "Initial retain should succeed"
|
||||
assert client.is_running, "Daemon should be running"
|
||||
|
||||
original_url = client.url
|
||||
|
||||
# Simulate daemon crash by stopping it
|
||||
client._manager.stop(client.profile)
|
||||
assert not client._manager.is_running(client.profile), (
|
||||
"Daemon should be stopped after simulated crash"
|
||||
)
|
||||
|
||||
# Next operation should transparently restart the daemon
|
||||
result2 = client.retain(bank_id=bank_id, content="After crash recovery")
|
||||
assert result2.success, "Retain after crash recovery should succeed"
|
||||
assert client.is_running, "Daemon should be running again after recovery"
|
||||
|
||||
# Verify recall still works
|
||||
recall_result = client.recall(bank_id=bank_id, query="crash")
|
||||
assert isinstance(recall_result.results, list), "Recall should return results"
|
||||
|
||||
finally:
|
||||
client.close()
|
||||
|
||||
@@ -46,4 +46,4 @@ __all__ = [
|
||||
"RemoteTEICrossEncoder",
|
||||
"LLMConfig",
|
||||
]
|
||||
__version__ = "0.5.4"
|
||||
__version__ = "0.5.1"
|
||||
|
||||
@@ -375,140 +375,6 @@ def decommission_worker(
|
||||
typer.echo(f"No tasks found for worker '{worker_id}'")
|
||||
|
||||
|
||||
async def _decommission_all_workers(db_url: str, schema: str = "public") -> list[dict[str, Any]]:
|
||||
"""Release all processing tasks from all workers, setting them back to pending status."""
|
||||
is_pg0, instance_name, _ = parse_pg0_url(db_url)
|
||||
if is_pg0:
|
||||
typer.echo(f"Starting embedded PostgreSQL (instance: {instance_name})...")
|
||||
resolved_url = await resolve_database_url(db_url)
|
||||
|
||||
conn = await asyncpg.connect(resolved_url)
|
||||
try:
|
||||
table = _fq_table("async_operations", schema)
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
UPDATE {table}
|
||||
SET status = 'pending', worker_id = NULL, claimed_at = NULL, updated_at = now()
|
||||
WHERE status = 'processing'
|
||||
RETURNING operation_id, worker_id, operation_type
|
||||
""",
|
||||
)
|
||||
return [dict(r) for r in rows]
|
||||
finally:
|
||||
await conn.close()
|
||||
|
||||
|
||||
@app.command(name="decommission-workers")
|
||||
def decommission_workers(
|
||||
schema: str = typer.Option("public", "--schema", "-s", help="Database schema"),
|
||||
yes: bool = typer.Option(False, "--yes", "-y", help="Skip confirmation prompt"),
|
||||
):
|
||||
"""Release all processing tasks from all workers (sets status back to pending).
|
||||
|
||||
Use this command to recover from situations where one or more workers have crashed
|
||||
or been removed without graceful shutdown. All tasks currently in 'processing' status
|
||||
will be released back to the queue regardless of which worker owns them.
|
||||
"""
|
||||
config = HindsightConfig.from_env()
|
||||
|
||||
if not config.database_url:
|
||||
typer.echo("Error: Database URL not configured.", err=True)
|
||||
typer.echo("Set HINDSIGHT_API_DATABASE_URL environment variable.", err=True)
|
||||
raise typer.Exit(1)
|
||||
|
||||
if not yes:
|
||||
typer.confirm(
|
||||
"This will release ALL processing tasks from ALL workers back to pending. Continue?",
|
||||
abort=True,
|
||||
)
|
||||
|
||||
typer.echo(f"Decommissioning all workers (schema: {schema})...")
|
||||
|
||||
released = asyncio.run(_decommission_all_workers(config.database_url, schema))
|
||||
|
||||
if released:
|
||||
# Group by worker_id for summary
|
||||
by_worker: dict[str, int] = {}
|
||||
for row in released:
|
||||
wid = row["worker_id"] or "unknown"
|
||||
by_worker[wid] = by_worker.get(wid, 0) + 1
|
||||
|
||||
typer.echo(f"Released {len(released)} task(s):")
|
||||
for wid, count in by_worker.items():
|
||||
typer.echo(f" {wid}: {count} task(s)")
|
||||
else:
|
||||
typer.echo("No processing tasks found")
|
||||
|
||||
|
||||
async def _worker_status(db_url: str, schema: str = "public") -> list[dict[str, Any]]:
|
||||
"""Get all processing tasks grouped by worker with their last update time."""
|
||||
is_pg0, instance_name, _ = parse_pg0_url(db_url)
|
||||
if is_pg0:
|
||||
typer.echo(f"Starting embedded PostgreSQL (instance: {instance_name})...")
|
||||
resolved_url = await resolve_database_url(db_url)
|
||||
|
||||
conn = await asyncpg.connect(resolved_url)
|
||||
try:
|
||||
table = _fq_table("async_operations", schema)
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT worker_id, operation_id, operation_type, bank_id,
|
||||
claimed_at, updated_at,
|
||||
now() - claimed_at AS running_for,
|
||||
now() - updated_at AS last_update_ago
|
||||
FROM {table}
|
||||
WHERE status = 'processing'
|
||||
ORDER BY worker_id, claimed_at
|
||||
""",
|
||||
)
|
||||
return [dict(r) for r in rows]
|
||||
finally:
|
||||
await conn.close()
|
||||
|
||||
|
||||
@app.command(name="worker-status")
|
||||
def worker_status(
|
||||
schema: str = typer.Option("public", "--schema", "-s", help="Database schema"),
|
||||
):
|
||||
"""Show all currently processing tasks grouped by worker.
|
||||
|
||||
Displays each worker's active tasks with operation type, bank, how long
|
||||
the task has been running, and when it was last updated. Useful for
|
||||
identifying dead workers with orphaned tasks.
|
||||
"""
|
||||
config = HindsightConfig.from_env()
|
||||
|
||||
if not config.database_url:
|
||||
typer.echo("Error: Database URL not configured.", err=True)
|
||||
typer.echo("Set HINDSIGHT_API_DATABASE_URL environment variable.", err=True)
|
||||
raise typer.Exit(1)
|
||||
|
||||
rows = asyncio.run(_worker_status(config.database_url, schema))
|
||||
|
||||
if not rows:
|
||||
typer.echo("No processing tasks found")
|
||||
return
|
||||
|
||||
# Group by worker_id
|
||||
by_worker: dict[str, list[dict[str, Any]]] = {}
|
||||
for row in rows:
|
||||
wid = row["worker_id"] or "unknown"
|
||||
by_worker.setdefault(wid, []).append(row)
|
||||
|
||||
typer.echo(f"Processing tasks across {len(by_worker)} worker(s):\n")
|
||||
for wid, tasks in by_worker.items():
|
||||
typer.echo(f"Worker: {wid} ({len(tasks)} task(s))")
|
||||
for task in tasks:
|
||||
op_id = str(task["operation_id"])[:8]
|
||||
running_for = task["running_for"]
|
||||
last_update = task["last_update_ago"]
|
||||
typer.echo(
|
||||
f" {op_id} {task['operation_type']:<20s} bank={task['bank_id']}"
|
||||
f" running={running_for} last_update={last_update} ago"
|
||||
)
|
||||
typer.echo("")
|
||||
|
||||
|
||||
def main():
|
||||
app()
|
||||
|
||||
|
||||
@@ -12,7 +12,6 @@ from dotenv import load_dotenv
|
||||
from sqlalchemy import engine_from_config, pool
|
||||
|
||||
# Import your models here
|
||||
from hindsight_api.db_url import to_libpq_url
|
||||
from hindsight_api.models import Base
|
||||
|
||||
|
||||
@@ -66,11 +65,11 @@ def get_database_url() -> str:
|
||||
"Set HINDSIGHT_API_DATABASE_URL environment variable or pass database_url to run_migrations()."
|
||||
)
|
||||
|
||||
# For migrations, use the sync psycopg2 driver (avoids pgbouncer prepared
|
||||
# statement issues and is required since create_engine is the sync API).
|
||||
# Also translates ?ssl=require (SQLAlchemy asyncpg style) to ?sslmode=require
|
||||
# (libpq style) for external-PostgreSQL deployments.
|
||||
database_url = to_libpq_url(database_url)
|
||||
# For migrations, use psycopg2 (sync driver) to avoid pgbouncer prepared statement issues
|
||||
if database_url.startswith("postgresql+asyncpg://"):
|
||||
database_url = database_url.replace("postgresql+asyncpg://", "postgresql://", 1)
|
||||
elif database_url.startswith("postgres+asyncpg://"):
|
||||
database_url = database_url.replace("postgres+asyncpg://", "postgresql://", 1)
|
||||
|
||||
# Update config with processed URL for engine_from_config to use
|
||||
config.set_main_option("sqlalchemy.url", database_url)
|
||||
|
||||
+4
-4
@@ -4,8 +4,8 @@ The previous GIN trigram index on canonical_name was case-sensitive, causing
|
||||
"Alice" and "alice" to have different trigram sets. This recreates it on
|
||||
LOWER(canonical_name) so the % operator matches case-insensitively.
|
||||
|
||||
Revision ID: 2eee35aa3cfc
|
||||
Revises: d6e7f8a9b0c1
|
||||
Revision ID: d6e7f8a9b0c1
|
||||
Revises: c5d6e7f8a9b0
|
||||
Create Date: 2026-03-31
|
||||
"""
|
||||
|
||||
@@ -13,8 +13,8 @@ from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
revision: str = "2eee35aa3cfc"
|
||||
down_revision: str | Sequence[str] | None = "d6e7f8a9b0c1"
|
||||
revision: str = "d6e7f8a9b0c1"
|
||||
down_revision: str | Sequence[str] | None = "c5d6e7f8a9b0"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
-40
@@ -1,40 +0,0 @@
|
||||
"""Merge divergent migration heads for v0.5.3
|
||||
|
||||
v0.5.3 shipped with two migration heads that were never unified:
|
||||
|
||||
* ``c4x5y6z7a8b9`` — delta-refresh chain
|
||||
(``add_last_refreshed_source_query`` ->
|
||||
``add_structured_content_to_mental_models`` ->
|
||||
``backsweep_orphan_observations_v2``)
|
||||
|
||||
* ``h3i4j5k6l7m8`` — per-bank vector indexes / audit log chain
|
||||
(the ``merge_heads_and_add_unit_entities_index`` subtree)
|
||||
|
||||
Both fork from ``z1u2v3w4x5y6``. Upgrades from v0.5.2 still succeed — the
|
||||
walker applies the three c4x5 revisions and leaves the database stamped at
|
||||
both heads — but the result is a split DAG: ``alembic upgrade head``
|
||||
(singular) is ambiguous, and any future migration has to pick one head as
|
||||
its parent, orphaning the other.
|
||||
|
||||
This revision linearises the DAG into a single head. It has no schema
|
||||
effect.
|
||||
|
||||
Revision ID: 8c6fa6f7230b
|
||||
Revises: c4x5y6z7a8b9, h3i4j5k6l7m8
|
||||
Create Date: 2026-04-18
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
revision: str = "8c6fa6f7230b"
|
||||
down_revision: str | Sequence[str] | None = ("c4x5y6z7a8b9", "h3i4j5k6l7m8")
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
pass
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
pass
|
||||
-38
@@ -1,38 +0,0 @@
|
||||
"""Add last_refreshed_source_query column to mental_models
|
||||
|
||||
Revision ID: a2v3w4x5y6z7
|
||||
Revises: z1u2v3w4x5y6
|
||||
Create Date: 2026-04-15
|
||||
|
||||
Tracks the source_query that was used during the most recent refresh.
|
||||
Used by delta-mode refresh to detect when the query has changed: if it has,
|
||||
delta mode falls back to a full regeneration because the surgical-edit
|
||||
assumption (same topic, new facts) no longer holds.
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
revision: str = "a2v3w4x5y6z7"
|
||||
down_revision: str | Sequence[str] | None = "z1u2v3w4x5y6"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _get_schema_prefix() -> str:
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
op.execute(f"""
|
||||
ALTER TABLE {schema}mental_models
|
||||
ADD COLUMN IF NOT EXISTS last_refreshed_source_query TEXT
|
||||
""")
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
op.execute(f"ALTER TABLE {schema}mental_models DROP COLUMN IF EXISTS last_refreshed_source_query")
|
||||
+2
-2
@@ -1,7 +1,7 @@
|
||||
"""Fix per-bank vector indexes to match configured extension
|
||||
|
||||
Revision ID: a4b5c6d7e8f9
|
||||
Revises: 2eee35aa3cfc
|
||||
Revises: d6e7f8a9b0c1
|
||||
Create Date: 2026-04-01
|
||||
|
||||
Migration d5e6f7a8b9c0 hardcoded HNSW when creating per-bank partial vector
|
||||
@@ -21,7 +21,7 @@ from alembic import context, op
|
||||
from sqlalchemy import text
|
||||
|
||||
revision: str = "a4b5c6d7e8f9"
|
||||
down_revision: str | Sequence[str] | None = "2eee35aa3cfc"
|
||||
down_revision: str | Sequence[str] | None = "d6e7f8a9b0c1"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
-44
@@ -1,44 +0,0 @@
|
||||
"""Add structured_content JSONB column to mental_models
|
||||
|
||||
Revision ID: b3w4x5y6z7a8
|
||||
Revises: a2v3w4x5y6z7
|
||||
Create Date: 2026-04-16
|
||||
|
||||
Stores the structured representation of a mental model document (sections,
|
||||
blocks). The plain ``content`` column remains the rendered markdown shown to
|
||||
users. ``structured_content`` is the source of truth for delta-mode refreshes:
|
||||
each refresh applies a list of typed operations to the structured doc, then
|
||||
re-renders to markdown — so unchanged sections come through byte-identical
|
||||
without an LLM round-trip.
|
||||
|
||||
Nullable: existing markdown-only mental models continue to work in full mode;
|
||||
the column is populated lazily the first time a model is refreshed in delta
|
||||
mode.
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
revision: str = "b3w4x5y6z7a8"
|
||||
down_revision: str | Sequence[str] | None = "a2v3w4x5y6z7"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _get_schema_prefix() -> str:
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
op.execute(f"""
|
||||
ALTER TABLE {schema}mental_models
|
||||
ADD COLUMN IF NOT EXISTS structured_content JSONB
|
||||
""")
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
op.execute(f"ALTER TABLE {schema}mental_models DROP COLUMN IF EXISTS structured_content")
|
||||
-66
@@ -1,66 +0,0 @@
|
||||
"""backsweep_orphan_observations_v2
|
||||
|
||||
Re-run of Pass 2 from migration ``g7h8i9j0k1l2_backsweep_orphan_observations``
|
||||
to sweep observations that became orphaned between then and now.
|
||||
|
||||
Why we need it again:
|
||||
``fact_storage.handle_document_tracking`` (the retain/upsert path) deleted
|
||||
the existing document via the FK cascade — which removes the source
|
||||
``memory_units`` — but never invalidated the observations derived from
|
||||
them. Only the explicit ``MemoryEngine.delete_document`` API called
|
||||
``_delete_stale_observations_for_memories``. Every document re-ingest
|
||||
therefore left orphan observations whose ``source_memory_ids`` arrays
|
||||
pointed at IDs that no longer existed in ``memory_units``.
|
||||
|
||||
``handle_document_tracking`` now calls the same cleanup helper before the
|
||||
cascade, so no new orphans will accumulate going forward. This migration
|
||||
cleans up the historical residue.
|
||||
|
||||
Identical to Pass 2 of g7h8i9j0k1l2. Pass 1 (memory_units whose bank is
|
||||
gone) is intentionally not re-run; that scenario has no fresh source.
|
||||
|
||||
Revision ID: c4x5y6z7a8b9
|
||||
Revises: b3w4x5y6z7a8
|
||||
Create Date: 2026-04-16
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
revision: str = "c4x5y6z7a8b9"
|
||||
down_revision: str | Sequence[str] | None = "b3w4x5y6z7a8"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _get_schema_prefix() -> str:
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
mu = f"{schema}memory_units"
|
||||
|
||||
# Delete observations whose every source_memory_id refers to a now-deleted
|
||||
# memory_unit (or the array is empty). Observations with at least one
|
||||
# surviving source are left alone — the consolidation engine will refresh
|
||||
# their text on the next pass.
|
||||
op.execute(
|
||||
f"""
|
||||
DELETE FROM {mu} orphan
|
||||
WHERE orphan.fact_type = 'observation'
|
||||
AND NOT EXISTS (
|
||||
SELECT 1
|
||||
FROM {mu} src
|
||||
WHERE src.id = ANY(orphan.source_memory_ids)
|
||||
AND src.bank_id = orphan.bank_id
|
||||
)
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
# Deleted rows cannot be restored.
|
||||
pass
|
||||
-39
@@ -1,39 +0,0 @@
|
||||
"""Drop unused metadata column from documents table
|
||||
|
||||
Revision ID: d6e7f8a9b0c1
|
||||
Revises: c2d3e4f5g6h7, c5d6e7f8a9b0
|
||||
Create Date: 2026-03-30
|
||||
|
||||
The metadata column on documents was always stored as an empty dict {}.
|
||||
Actual document metadata is stored inside retain_params.metadata.
|
||||
|
||||
This migration was originally shipped in v0.4.22, then its file was deleted
|
||||
in v0.5.0 (and its revision ID accidentally reused by 2eee35aa3cfc).
|
||||
Restoring the file so that databases stamped at this revision can upgrade
|
||||
cleanly to v0.5.x+.
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
revision: str = "d6e7f8a9b0c1"
|
||||
down_revision: str | Sequence[str] | None = ("c2d3e4f5g6h7", "c5d6e7f8a9b0")
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _get_schema_prefix() -> str:
|
||||
"""Get schema prefix for table names (required for multi-tenant support)."""
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
op.execute(f"ALTER TABLE {schema}documents DROP COLUMN IF EXISTS metadata")
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
op.execute(f"ALTER TABLE {schema}documents ADD COLUMN IF NOT EXISTS metadata jsonb DEFAULT '{{}}'")
|
||||
+2
-2
@@ -1,7 +1,7 @@
|
||||
"""Merge 3 migration heads and add unit_entities composite index
|
||||
|
||||
Revision ID: h3i4j5k6l7m8
|
||||
Revises: a4b5c6d7e8f9, g2h3i4j5k6l7
|
||||
Revises: a4b5c6d7e8f9, c2d3e4f5g6h7, g2h3i4j5k6l7
|
||||
Create Date: 2026-04-07
|
||||
|
||||
Merges three unmerged migration heads into one, and adds a composite index
|
||||
@@ -14,7 +14,7 @@ from collections.abc import Sequence
|
||||
from alembic import context, op
|
||||
|
||||
revision: str = "h3i4j5k6l7m8"
|
||||
down_revision: str | Sequence[str] | None = ("a4b5c6d7e8f9", "g2h3i4j5k6l7")
|
||||
down_revision: str | Sequence[str] | None = ("a4b5c6d7e8f9", "c2d3e4f5g6h7", "g2h3i4j5k6l7")
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
-39
@@ -1,39 +0,0 @@
|
||||
"""Add 'cancelled' to async_operations status check constraint
|
||||
|
||||
Revision ID: i4j5k6l7m8n9
|
||||
Revises: 8c6fa6f7230b
|
||||
Create Date: 2026-04-23
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
revision: str = "i4j5k6l7m8n9"
|
||||
down_revision: str | Sequence[str] | None = "8c6fa6f7230b"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _get_schema_prefix() -> str:
|
||||
"""Get schema prefix for table names (required for multi-tenant support)."""
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
op.execute(f"ALTER TABLE {schema}async_operations DROP CONSTRAINT IF EXISTS async_operations_status_check")
|
||||
op.execute(
|
||||
f"ALTER TABLE {schema}async_operations ADD CONSTRAINT async_operations_status_check "
|
||||
f"CHECK (status IN ('pending', 'processing', 'completed', 'failed', 'cancelled'))"
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
op.execute(f"ALTER TABLE {schema}async_operations DROP CONSTRAINT IF EXISTS async_operations_status_check")
|
||||
op.execute(
|
||||
f"ALTER TABLE {schema}async_operations ADD CONSTRAINT async_operations_status_check "
|
||||
f"CHECK (status IN ('pending', 'processing', 'completed', 'failed'))"
|
||||
)
|
||||
@@ -294,44 +294,6 @@ class EntityListResponse(BaseModel):
|
||||
offset: int
|
||||
|
||||
|
||||
class EntityGraphResponse(BaseModel):
|
||||
"""Response model for entity co-occurrence graph endpoint."""
|
||||
|
||||
model_config = ConfigDict(
|
||||
json_schema_extra={
|
||||
"example": {
|
||||
"nodes": [
|
||||
{"data": {"id": "uuid-1", "label": "Alice", "mentionCount": 12, "color": "#42a5f5"}},
|
||||
{"data": {"id": "uuid-2", "label": "Google", "mentionCount": 8, "color": "#42a5f5"}},
|
||||
],
|
||||
"edges": [
|
||||
{
|
||||
"data": {
|
||||
"id": "uuid-1-uuid-2",
|
||||
"source": "uuid-1",
|
||||
"target": "uuid-2",
|
||||
"linkType": "cooccurrence",
|
||||
"weight": 5,
|
||||
"color": "#ffd700",
|
||||
"lineStyle": "solid",
|
||||
"lastCooccurred": "2024-02-01T14:00:00Z",
|
||||
}
|
||||
}
|
||||
],
|
||||
"total_entities": 2,
|
||||
"total_edges": 1,
|
||||
"limit": 1000,
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
nodes: list[dict[str, Any]]
|
||||
edges: list[dict[str, Any]]
|
||||
total_entities: int
|
||||
total_edges: int
|
||||
limit: int
|
||||
|
||||
|
||||
class EntityDetailResponse(BaseModel):
|
||||
"""Response model for entity detail endpoint."""
|
||||
|
||||
@@ -1335,9 +1297,6 @@ class DocumentResponse(BaseModel):
|
||||
created_at: str
|
||||
updated_at: str
|
||||
memory_unit_count: int
|
||||
nodes_by_fact_type: dict[str, int] | None = Field(
|
||||
default=None, description="Memory count per fact type (world, experience, observation)"
|
||||
)
|
||||
tags: list[str] = FieldWithDefault(list, description="Tags associated with this document")
|
||||
document_metadata: dict[str, Any] | None = Field(default=None, description="Document metadata")
|
||||
retain_params: dict[str, Any] | None = Field(default=None, description="Parameters used during retain")
|
||||
@@ -1411,23 +1370,6 @@ class ChunkResponse(BaseModel):
|
||||
created_at: str
|
||||
|
||||
|
||||
class ListChunksResponse(BaseModel):
|
||||
"""Response model for listing chunks of a document."""
|
||||
|
||||
items: list[ChunkResponse]
|
||||
total: int
|
||||
limit: int
|
||||
offset: int
|
||||
|
||||
|
||||
class ReprocessDocumentResponse(BaseModel):
|
||||
"""Response model for reprocess document endpoint."""
|
||||
|
||||
success: bool
|
||||
operation_id: str
|
||||
items_count: int
|
||||
|
||||
|
||||
class DeleteResponse(BaseModel):
|
||||
"""Response model for delete operations."""
|
||||
|
||||
@@ -1474,7 +1416,6 @@ class BankStatsResponse(BaseModel):
|
||||
"failed_operations": 0,
|
||||
"last_consolidated_at": "2024-01-15T10:30:00Z",
|
||||
"pending_consolidation": 0,
|
||||
"failed_consolidation": 0,
|
||||
"total_observations": 45,
|
||||
}
|
||||
}
|
||||
@@ -1490,41 +1431,12 @@ class BankStatsResponse(BaseModel):
|
||||
links_breakdown: dict[str, dict[str, int]]
|
||||
pending_operations: int
|
||||
failed_operations: int
|
||||
operations_by_status: dict[str, int] = Field(
|
||||
default_factory=dict,
|
||||
description="Async operations grouped by status (pending, processing, completed, failed, cancelled).",
|
||||
)
|
||||
# Consolidation stats
|
||||
last_consolidated_at: str | None = Field(default=None, description="When consolidation last ran (ISO format)")
|
||||
pending_consolidation: int = Field(default=0, description="Number of memories not yet processed into observations")
|
||||
failed_consolidation: int = Field(
|
||||
default=0,
|
||||
description="Number of source memories (world/experience) whose consolidation permanently failed and can be retried via the consolidation recovery endpoint.",
|
||||
)
|
||||
total_observations: int = Field(default=0, description="Total number of observations")
|
||||
|
||||
|
||||
class MemoryTimeseriesBucket(BaseModel):
|
||||
"""One bucket in the memory ingestion time-series."""
|
||||
|
||||
time: str = Field(description="Bucket start timestamp in ISO-8601 (UTC).")
|
||||
world: int = Field(default=0, description="World-fact memories ingested in this bucket.")
|
||||
experience: int = Field(default=0, description="Experience memories ingested in this bucket.")
|
||||
observation: int = Field(default=0, description="Observations recorded in this bucket.")
|
||||
|
||||
|
||||
class MemoriesTimeseriesResponse(BaseModel):
|
||||
"""Time-series of memory ingestion bucketed by time and fact type."""
|
||||
|
||||
bank_id: str
|
||||
period: str = Field(description="One of: 1h, 12h, 1d, 7d, 30d, 90d.")
|
||||
trunc: str = Field(description="Bucket granularity: minute, hour, day.")
|
||||
buckets: list[MemoryTimeseriesBucket] = Field(
|
||||
default_factory=list,
|
||||
description="Per-bucket counts, always returned fully padded for the requested period.",
|
||||
)
|
||||
|
||||
|
||||
# Mental Model models
|
||||
|
||||
|
||||
@@ -1581,16 +1493,6 @@ class UpdateDirectiveRequest(BaseModel):
|
||||
class MentalModelTrigger(BaseModel):
|
||||
"""Trigger settings for a mental model."""
|
||||
|
||||
mode: Literal["full", "delta"] = Field(
|
||||
default="full",
|
||||
description=(
|
||||
"Refresh mode. 'full' (default) regenerates the mental model content from scratch on each refresh. "
|
||||
"'delta' performs surgical edits against the existing content: unchanged sections are preserved "
|
||||
"byte-for-byte, stale content is removed, new content is added. If the mental model has no existing "
|
||||
"content, or if the source_query has changed since the last refresh, delta mode falls back to a full "
|
||||
"regeneration automatically."
|
||||
),
|
||||
)
|
||||
refresh_after_consolidation: bool = Field(
|
||||
default=False,
|
||||
description="If true, refresh this mental model after observations consolidation (real-time mode)",
|
||||
@@ -1624,27 +1526,6 @@ class MentalModelTrigger(BaseModel):
|
||||
"Supports nested and/or/not expressions for complex tag-based scoping."
|
||||
),
|
||||
)
|
||||
include_chunks: bool | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Override whether the internal recall used during refresh returns raw chunk text. "
|
||||
"None means use the bank/global config default (recall_include_chunks)."
|
||||
),
|
||||
)
|
||||
recall_max_tokens: int | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Override the token budget for facts returned by the internal recall during refresh. "
|
||||
"None means use the bank/global config default (recall_max_tokens)."
|
||||
),
|
||||
)
|
||||
recall_chunks_max_tokens: int | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Override the token budget for raw chunks returned by the internal recall during refresh. "
|
||||
"None means use the bank/global config default (recall_chunks_max_tokens)."
|
||||
),
|
||||
)
|
||||
|
||||
@field_validator("fact_types")
|
||||
@classmethod
|
||||
@@ -1674,14 +1555,6 @@ class MentalModelResponse(BaseModel):
|
||||
default=None,
|
||||
description="Full reflect API response payload including based_on facts and observations",
|
||||
)
|
||||
is_stale: bool | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"True when new memories matching this mental model's tag/fact_type scope have been "
|
||||
"ingested since last_refreshed_at, or consolidation has pending items. Only populated "
|
||||
"when detail=full."
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class MentalModelListResponse(BaseModel):
|
||||
@@ -1800,61 +1673,6 @@ class BankTemplateConfig(BaseModel):
|
||||
entities_allow_free_form: bool | None = Field(
|
||||
default=None, description="Allow entities outside the label vocabulary"
|
||||
)
|
||||
retain_default_strategy: str | None = Field(
|
||||
default=None, description="Name of the default retain strategy (key into retain_strategies map)"
|
||||
)
|
||||
retain_strategies: dict | None = Field(
|
||||
default=None, description="Map of retain strategy name to per-strategy config dict"
|
||||
)
|
||||
retain_chunk_batch_size: int | None = Field(
|
||||
default=None, description="Max chunks per streaming batch (0 disables batching)"
|
||||
)
|
||||
mcp_enabled_tools: list[str] | None = Field(
|
||||
default=None, description="MCP tool allowlist for this bank (None = all tools)"
|
||||
)
|
||||
consolidation_llm_batch_size: int | None = Field(
|
||||
default=None, description="LLM batch size for observation consolidation"
|
||||
)
|
||||
consolidation_source_facts_max_tokens: int | None = Field(
|
||||
default=None, description="Max tokens of source facts per consolidation batch"
|
||||
)
|
||||
consolidation_source_facts_max_tokens_per_observation: int | None = Field(
|
||||
default=None, description="Max tokens of source facts per observation"
|
||||
)
|
||||
max_observations_per_scope: int | None = Field(
|
||||
default=None, description="Max observations to retain per consolidation scope"
|
||||
)
|
||||
reflect_source_facts_max_tokens: int | None = Field(
|
||||
default=None, description="Max tokens of source facts per reflect call"
|
||||
)
|
||||
llm_gemini_safety_settings: list | None = Field(
|
||||
default=None, description="Per-bank Gemini/VertexAI safety filter settings"
|
||||
)
|
||||
recall_budget_function: str | None = Field(
|
||||
default=None, description="Recall budget mapping function: 'fixed' or 'adaptive'"
|
||||
)
|
||||
recall_budget_fixed_low: int | None = Field(
|
||||
default=None, description="Fixed thinking_budget for budget=low (function='fixed')"
|
||||
)
|
||||
recall_budget_fixed_mid: int | None = Field(
|
||||
default=None, description="Fixed thinking_budget for budget=mid (function='fixed')"
|
||||
)
|
||||
recall_budget_fixed_high: int | None = Field(
|
||||
default=None, description="Fixed thinking_budget for budget=high (function='fixed')"
|
||||
)
|
||||
recall_budget_adaptive_low: float | None = Field(
|
||||
default=None, description="Ratio of max_tokens for budget=low (function='adaptive')"
|
||||
)
|
||||
recall_budget_adaptive_mid: float | None = Field(
|
||||
default=None, description="Ratio of max_tokens for budget=mid (function='adaptive')"
|
||||
)
|
||||
recall_budget_adaptive_high: float | None = Field(
|
||||
default=None, description="Ratio of max_tokens for budget=high (function='adaptive')"
|
||||
)
|
||||
recall_budget_min: int | None = Field(default=None, description="Floor for the adaptive function (after clamping)")
|
||||
recall_budget_max: int | None = Field(
|
||||
default=None, description="Ceiling for the adaptive function (after clamping)"
|
||||
)
|
||||
|
||||
def get_config_updates(self) -> dict[str, Any]:
|
||||
"""Return only the fields that were explicitly set (non-None)."""
|
||||
@@ -2139,8 +1957,6 @@ class OperationResponse(BaseModel):
|
||||
"created_at": "2024-01-15T10:30:00Z",
|
||||
"status": "pending",
|
||||
"error_message": None,
|
||||
"retry_count": 0,
|
||||
"next_retry_at": None,
|
||||
}
|
||||
}
|
||||
)
|
||||
@@ -2152,20 +1968,6 @@ class OperationResponse(BaseModel):
|
||||
created_at: str
|
||||
status: str
|
||||
error_message: str | None
|
||||
retry_count: int | None = Field(
|
||||
default=None,
|
||||
description="Number of times this operation has been retried after failure.",
|
||||
)
|
||||
next_retry_at: str | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"When the worker will next attempt this operation. For a pending "
|
||||
"operation, a value in the future indicates the task is waiting "
|
||||
"rather than available for immediate pickup — for example, an "
|
||||
"extension may have raised DeferOperation to park the task until "
|
||||
"some backpressure window opens. Always null for completed tasks."
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class ConsolidationResponse(BaseModel):
|
||||
@@ -2269,25 +2071,12 @@ class OperationStatusResponse(BaseModel):
|
||||
)
|
||||
|
||||
operation_id: str
|
||||
status: Literal["pending", "processing", "completed", "failed", "cancelled", "not_found"]
|
||||
status: Literal["pending", "completed", "failed", "not_found"]
|
||||
operation_type: str | None = None
|
||||
created_at: str | None = None
|
||||
updated_at: str | None = None
|
||||
completed_at: str | None = None
|
||||
error_message: str | None = None
|
||||
retry_count: int | None = Field(
|
||||
default=None,
|
||||
description="Number of times this operation has been retried after failure.",
|
||||
)
|
||||
next_retry_at: str | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"When the worker will next attempt this operation. For a pending "
|
||||
"operation, a value in the future indicates the task is parked "
|
||||
"(e.g. by an extension raising DeferOperation) rather than awaiting "
|
||||
"immediate pickup."
|
||||
),
|
||||
)
|
||||
result_metadata: dict[str, Any] | None = Field(
|
||||
default=None,
|
||||
description="Internal metadata for debugging. Structure may change without notice. Not for production use.",
|
||||
@@ -2295,10 +2084,6 @@ class OperationStatusResponse(BaseModel):
|
||||
child_operations: list[ChildOperationStatus] | None = Field(
|
||||
default=None, description="Child operations for batch operations (if applicable)"
|
||||
)
|
||||
task_payload: dict[str, Any] | None = Field(
|
||||
default=None,
|
||||
description="Raw task payload (params the operation was submitted with). Only populated when include_payload=true.",
|
||||
)
|
||||
|
||||
|
||||
class AsyncOperationSubmitResponse(BaseModel):
|
||||
@@ -2635,7 +2420,7 @@ def create_app(
|
||||
schema=schema,
|
||||
tenant_extension=memory._tenant_extension,
|
||||
max_slots=config.worker_max_slots,
|
||||
slot_reservations=config.worker_slot_reservations,
|
||||
consolidation_max_slots=config.worker_consolidation_max_slots,
|
||||
)
|
||||
poller_task = asyncio.create_task(poller.run())
|
||||
logging.info(f"Worker poller started (worker_id={worker_id})")
|
||||
@@ -2951,22 +2736,12 @@ def _register_routes(app: FastAPI):
|
||||
q: str | None = None,
|
||||
tags: list[str] | None = Query(None),
|
||||
tags_match: str = "all_strict",
|
||||
document_id: str | None = None,
|
||||
chunk_id: 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,
|
||||
limit=limit,
|
||||
q=q,
|
||||
tags=tags,
|
||||
tags_match=tags_match,
|
||||
document_id=document_id,
|
||||
chunk_id=chunk_id,
|
||||
request_context=request_context,
|
||||
bank_id, type, limit=limit, q=q, tags=tags, tags_match=tags_match, request_context=request_context
|
||||
)
|
||||
return data
|
||||
except OperationValidationError as e:
|
||||
@@ -2992,7 +2767,6 @@ def _register_routes(app: FastAPI):
|
||||
bank_id: str,
|
||||
type: str | None = None,
|
||||
q: str | None = None,
|
||||
consolidation_state: str | None = None,
|
||||
limit: int = 100,
|
||||
offset: int = 0,
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
@@ -3007,8 +2781,6 @@ def _register_routes(app: FastAPI):
|
||||
bank_id: Memory Bank ID (from path)
|
||||
type: Filter by fact type (world, experience, opinion)
|
||||
q: Search query for full-text search (searches text and context)
|
||||
consolidation_state: Filter by consolidation state for source memories
|
||||
(world/experience). One of 'failed', 'pending', or 'done'.
|
||||
limit: Maximum number of results (default: 100)
|
||||
offset: Offset for pagination (default: 0)
|
||||
"""
|
||||
@@ -3017,14 +2789,11 @@ def _register_routes(app: FastAPI):
|
||||
bank_id=bank_id,
|
||||
fact_type=type,
|
||||
search_query=q,
|
||||
consolidation_state=consolidation_state,
|
||||
limit=limit,
|
||||
offset=offset,
|
||||
request_context=request_context,
|
||||
)
|
||||
return data
|
||||
except ValueError as e:
|
||||
raise HTTPException(status_code=400, detail=str(e))
|
||||
except OperationValidationError as e:
|
||||
raise HTTPException(status_code=e.status_code, detail=e.reason)
|
||||
except (AuthenticationError, HTTPException):
|
||||
@@ -3479,10 +3248,8 @@ def _register_routes(app: FastAPI):
|
||||
links_breakdown=links_breakdown,
|
||||
pending_operations=ops.get("pending", 0),
|
||||
failed_operations=ops.get("failed", 0),
|
||||
operations_by_status=ops,
|
||||
last_consolidated_at=stats["last_consolidated_at"],
|
||||
pending_consolidation=stats["pending_consolidation"],
|
||||
failed_consolidation=stats.get("failed_consolidation", 0),
|
||||
total_observations=stats["total_observations"],
|
||||
)
|
||||
except OperationValidationError as e:
|
||||
@@ -3496,35 +3263,6 @@ def _register_routes(app: FastAPI):
|
||||
logger.error(f"Error in /v1/default/banks/{bank_id}/stats: {error_detail}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
@app.get(
|
||||
"/v1/default/banks/{bank_id}/stats/memories-timeseries",
|
||||
response_model=MemoriesTimeseriesResponse,
|
||||
summary="Memory ingestion time-series",
|
||||
description="Memories ingested over a period, bucketed by time and broken down by fact type.",
|
||||
operation_id="get_memories_timeseries",
|
||||
tags=["Banks"],
|
||||
)
|
||||
async def api_memories_timeseries(
|
||||
bank_id: str,
|
||||
period: str = "7d",
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
):
|
||||
try:
|
||||
data = await app.state.memory.get_memories_timeseries(
|
||||
bank_id, period=period, request_context=request_context
|
||||
)
|
||||
return MemoriesTimeseriesResponse(**data)
|
||||
except OperationValidationError as e:
|
||||
raise HTTPException(status_code=e.status_code, detail=e.reason)
|
||||
except (AuthenticationError, HTTPException):
|
||||
raise
|
||||
except Exception as e:
|
||||
import traceback
|
||||
|
||||
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
|
||||
logger.error(f"Error in /v1/default/banks/{bank_id}/stats/memories-timeseries: {error_detail}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
@app.get(
|
||||
"/v1/default/banks/{bank_id}/entities",
|
||||
response_model=EntityListResponse,
|
||||
@@ -3561,36 +3299,6 @@ def _register_routes(app: FastAPI):
|
||||
logger.error(f"Error in /v1/default/banks/{bank_id}/entities: {error_detail}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
@app.get(
|
||||
"/v1/default/banks/{bank_id}/entities/graph",
|
||||
response_model=EntityGraphResponse,
|
||||
summary="Get entity co-occurrence graph",
|
||||
description="Return a graph of entities (nodes) and their co-occurrences (edges) for visualization.",
|
||||
operation_id="get_entity_graph",
|
||||
tags=["Entities"],
|
||||
)
|
||||
async def api_entity_graph(
|
||||
bank_id: str,
|
||||
limit: int = Query(default=1000, description="Maximum number of co-occurrence edges to return"),
|
||||
min_count: int = Query(default=1, description="Minimum cooccurrence_count to include an edge"),
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
):
|
||||
"""Return entity co-occurrence graph for a bank."""
|
||||
try:
|
||||
return await app.state.memory.get_entity_graph(
|
||||
bank_id, limit=limit, min_count=min_count, request_context=request_context
|
||||
)
|
||||
except OperationValidationError as e:
|
||||
raise HTTPException(status_code=e.status_code, detail=e.reason)
|
||||
except (AuthenticationError, HTTPException):
|
||||
raise
|
||||
except Exception as e:
|
||||
import traceback
|
||||
|
||||
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
|
||||
logger.error(f"Error in /v1/default/banks/{bank_id}/entities/graph: {error_detail}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
@app.get(
|
||||
"/v1/default/banks/{bank_id}/entities/{entity_id}",
|
||||
response_model=EntityDetailResponse,
|
||||
@@ -4191,99 +3899,6 @@ def _register_routes(app: FastAPI):
|
||||
logger.error(f"Error in /v1/default/banks/{bank_id}/documents: {error_detail}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
@app.get(
|
||||
"/v1/default/banks/{bank_id}/documents/{document_id:path}/chunks",
|
||||
response_model=ListChunksResponse,
|
||||
summary="List document chunks",
|
||||
description="List all chunks for a given document, ordered by chunk index.",
|
||||
operation_id="list_document_chunks",
|
||||
tags=["Documents"],
|
||||
)
|
||||
async def api_list_document_chunks(
|
||||
bank_id: str,
|
||||
document_id: str,
|
||||
limit: int = Query(default=100, ge=1, le=1000, description="Maximum number of chunks to return"),
|
||||
offset: int = Query(default=0, ge=0, description="Offset for pagination"),
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
):
|
||||
"""
|
||||
List all chunks for a document, ordered by chunk_index.
|
||||
|
||||
Args:
|
||||
bank_id: Memory Bank ID (from path)
|
||||
document_id: Document ID (from path)
|
||||
limit: Maximum number of chunks to return (default: 100)
|
||||
offset: Offset for pagination (default: 0)
|
||||
"""
|
||||
try:
|
||||
result = await app.state.memory.list_document_chunks(
|
||||
bank_id=bank_id,
|
||||
document_id=document_id,
|
||||
limit=limit,
|
||||
offset=offset,
|
||||
request_context=request_context,
|
||||
)
|
||||
if result is None:
|
||||
raise HTTPException(status_code=404, detail="Document not found")
|
||||
return result
|
||||
except OperationValidationError as e:
|
||||
raise HTTPException(status_code=e.status_code, detail=e.reason)
|
||||
except (AuthenticationError, HTTPException):
|
||||
raise
|
||||
except Exception as e:
|
||||
import traceback
|
||||
|
||||
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
|
||||
logger.error(f"Error in /v1/default/banks/{bank_id}/documents/{document_id}/chunks: {error_detail}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
@app.post(
|
||||
"/v1/default/banks/{bank_id}/documents/{document_id:path}/reprocess",
|
||||
response_model=ReprocessDocumentResponse,
|
||||
summary="Reprocess document",
|
||||
description="Re-run the retain pipeline on an existing document without changing its content. "
|
||||
"This deletes the existing memory units and re-extracts facts using the current engine configuration. "
|
||||
"Useful when the LLM model, chunking strategy, or extraction settings have changed.",
|
||||
operation_id="reprocess_document",
|
||||
tags=["Documents"],
|
||||
)
|
||||
@audited("reprocess_document")
|
||||
async def api_reprocess_document(
|
||||
bank_id: str,
|
||||
document_id: str,
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
):
|
||||
"""
|
||||
Reprocess a document by re-running retain with its existing content and parameters.
|
||||
|
||||
Args:
|
||||
bank_id: Memory Bank ID (from path)
|
||||
document_id: Document ID (from path)
|
||||
"""
|
||||
try:
|
||||
result = await app.state.memory.reprocess_document(
|
||||
bank_id=bank_id,
|
||||
document_id=document_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
if result is None:
|
||||
raise HTTPException(status_code=404, detail="Document not found")
|
||||
return ReprocessDocumentResponse(
|
||||
success=True,
|
||||
operation_id=result["operation_id"],
|
||||
items_count=result["items_count"],
|
||||
)
|
||||
except OperationValidationError as e:
|
||||
raise HTTPException(status_code=e.status_code, detail=e.reason)
|
||||
except (AuthenticationError, HTTPException):
|
||||
raise
|
||||
except Exception as e:
|
||||
import traceback
|
||||
|
||||
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
|
||||
logger.error(f"Error in /v1/default/banks/{bank_id}/documents/{document_id}/reprocess: {error_detail}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
@app.get(
|
||||
"/v1/default/banks/{bank_id}/documents/{document_id:path}",
|
||||
response_model=DocumentResponse,
|
||||
@@ -4511,28 +4126,19 @@ def _register_routes(app: FastAPI):
|
||||
)
|
||||
async def api_list_operations(
|
||||
bank_id: str,
|
||||
status: str | None = Query(
|
||||
default=None, description="Filter by status: pending, processing, completed, failed, or cancelled"
|
||||
),
|
||||
status: str | None = Query(default=None, description="Filter by status: pending, completed, or failed"),
|
||||
type: str | None = Query(
|
||||
default=None,
|
||||
description="Filter by operation type: retain, consolidation, refresh_mental_model, file_convert_retain, webhook_delivery",
|
||||
),
|
||||
limit: int = Query(default=20, ge=1, le=100, description="Maximum number of operations to return"),
|
||||
offset: int = Query(default=0, ge=0, description="Number of operations to skip"),
|
||||
exclude_parents: bool = Query(default=False, description="Exclude parent batch operations from results"),
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
):
|
||||
"""List async operations for a memory bank with optional filtering and pagination."""
|
||||
try:
|
||||
result = await app.state.memory.list_operations(
|
||||
bank_id,
|
||||
status=status,
|
||||
task_type=type,
|
||||
limit=limit,
|
||||
offset=offset,
|
||||
exclude_parents=exclude_parents,
|
||||
request_context=request_context,
|
||||
bank_id, status=status, task_type=type, limit=limit, offset=offset, request_context=request_context
|
||||
)
|
||||
return OperationsListResponse(
|
||||
bank_id=bank_id,
|
||||
@@ -4562,13 +4168,7 @@ def _register_routes(app: FastAPI):
|
||||
tags=["Operations"],
|
||||
)
|
||||
async def api_get_operation_status(
|
||||
bank_id: str,
|
||||
operation_id: str,
|
||||
include_payload: bool = Query(
|
||||
default=False,
|
||||
description="Include the raw task payload (submission params) in the response. May be large.",
|
||||
),
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
bank_id: str, operation_id: str, request_context: RequestContext = Depends(get_request_context)
|
||||
):
|
||||
"""Get the status of an async operation."""
|
||||
try:
|
||||
@@ -4578,9 +4178,7 @@ def _register_routes(app: FastAPI):
|
||||
except ValueError:
|
||||
raise HTTPException(status_code=400, detail=f"Invalid operation_id format: {operation_id}")
|
||||
|
||||
result = await app.state.memory.get_operation_status(
|
||||
bank_id, operation_id, request_context=request_context, include_payload=include_payload
|
||||
)
|
||||
result = await app.state.memory.get_operation_status(bank_id, operation_id, request_context=request_context)
|
||||
return OperationStatusResponse(**result)
|
||||
except OperationValidationError as e:
|
||||
raise HTTPException(status_code=e.status_code, detail=e.reason)
|
||||
|
||||
@@ -111,6 +111,7 @@ def create_mcp_server(memory: MemoryEngine, multi_bank: bool = True) -> FastMCP:
|
||||
"delete_directive",
|
||||
"list_memories",
|
||||
"get_memory",
|
||||
"delete_memory",
|
||||
"list_documents",
|
||||
"get_document",
|
||||
"delete_document",
|
||||
|
||||
@@ -177,13 +177,11 @@ ENV_EMBEDDINGS_TEI_URL = "HINDSIGHT_API_EMBEDDINGS_TEI_URL"
|
||||
ENV_EMBEDDINGS_OPENAI_API_KEY = "HINDSIGHT_API_EMBEDDINGS_OPENAI_API_KEY"
|
||||
ENV_EMBEDDINGS_OPENAI_MODEL = "HINDSIGHT_API_EMBEDDINGS_OPENAI_MODEL"
|
||||
ENV_EMBEDDINGS_OPENAI_BASE_URL = "HINDSIGHT_API_EMBEDDINGS_OPENAI_BASE_URL"
|
||||
ENV_EMBEDDINGS_OPENAI_BATCH_SIZE = "HINDSIGHT_API_EMBEDDINGS_OPENAI_BATCH_SIZE"
|
||||
|
||||
# Gemini/Vertex AI embeddings configuration
|
||||
ENV_EMBEDDINGS_GEMINI_API_KEY = "HINDSIGHT_API_EMBEDDINGS_GEMINI_API_KEY"
|
||||
ENV_EMBEDDINGS_GEMINI_MODEL = "HINDSIGHT_API_EMBEDDINGS_GEMINI_MODEL"
|
||||
ENV_EMBEDDINGS_GEMINI_OUTPUT_DIMENSIONALITY = "HINDSIGHT_API_EMBEDDINGS_GEMINI_OUTPUT_DIMENSIONALITY"
|
||||
ENV_EMBEDDINGS_GEMINI_FORCE_IPV4 = "HINDSIGHT_API_EMBEDDINGS_GEMINI_FORCE_IPV4"
|
||||
ENV_EMBEDDINGS_VERTEXAI_PROJECT_ID = "HINDSIGHT_API_EMBEDDINGS_VERTEXAI_PROJECT_ID"
|
||||
ENV_EMBEDDINGS_VERTEXAI_REGION = "HINDSIGHT_API_EMBEDDINGS_VERTEXAI_REGION"
|
||||
ENV_EMBEDDINGS_VERTEXAI_SERVICE_ACCOUNT_KEY = "HINDSIGHT_API_EMBEDDINGS_VERTEXAI_SERVICE_ACCOUNT_KEY"
|
||||
@@ -192,7 +190,6 @@ ENV_EMBEDDINGS_VERTEXAI_SERVICE_ACCOUNT_KEY = "HINDSIGHT_API_EMBEDDINGS_VERTEXAI
|
||||
ENV_EMBEDDINGS_COHERE_API_KEY = "HINDSIGHT_API_EMBEDDINGS_COHERE_API_KEY"
|
||||
ENV_EMBEDDINGS_COHERE_MODEL = "HINDSIGHT_API_EMBEDDINGS_COHERE_MODEL"
|
||||
ENV_EMBEDDINGS_COHERE_BASE_URL = "HINDSIGHT_API_EMBEDDINGS_COHERE_BASE_URL"
|
||||
ENV_EMBEDDINGS_COHERE_OUTPUT_DIMENSIONS = "HINDSIGHT_API_EMBEDDINGS_COHERE_OUTPUT_DIMENSIONS"
|
||||
ENV_RERANKER_COHERE_API_KEY = "HINDSIGHT_API_RERANKER_COHERE_API_KEY"
|
||||
ENV_RERANKER_COHERE_MODEL = "HINDSIGHT_API_RERANKER_COHERE_MODEL"
|
||||
ENV_RERANKER_COHERE_BASE_URL = "HINDSIGHT_API_RERANKER_COHERE_BASE_URL"
|
||||
@@ -241,11 +238,9 @@ ENV_RERANKER_LOCAL_BATCH_SIZE = "HINDSIGHT_API_RERANKER_LOCAL_BATCH_SIZE"
|
||||
ENV_RERANKER_TEI_URL = "HINDSIGHT_API_RERANKER_TEI_URL"
|
||||
ENV_RERANKER_TEI_BATCH_SIZE = "HINDSIGHT_API_RERANKER_TEI_BATCH_SIZE"
|
||||
ENV_RERANKER_TEI_MAX_CONCURRENT = "HINDSIGHT_API_RERANKER_TEI_MAX_CONCURRENT"
|
||||
ENV_RERANKER_TEI_HTTP_TIMEOUT = "HINDSIGHT_API_RERANKER_TEI_HTTP_TIMEOUT"
|
||||
ENV_RERANKER_MAX_CANDIDATES = "HINDSIGHT_API_RERANKER_MAX_CANDIDATES"
|
||||
ENV_RERANKER_FLASHRANK_MODEL = "HINDSIGHT_API_RERANKER_FLASHRANK_MODEL"
|
||||
ENV_RERANKER_FLASHRANK_CACHE_DIR = "HINDSIGHT_API_RERANKER_FLASHRANK_CACHE_DIR"
|
||||
ENV_RERANKER_FLASHRANK_CPU_MEM_ARENA = "HINDSIGHT_API_RERANKER_FLASHRANK_CPU_MEM_ARENA"
|
||||
|
||||
# ZeroEntropy configuration (reranker only)
|
||||
ENV_RERANKER_ZEROENTROPY_API_KEY = "HINDSIGHT_API_RERANKER_ZEROENTROPY_API_KEY"
|
||||
@@ -270,7 +265,6 @@ ENV_PORT = "HINDSIGHT_API_PORT"
|
||||
ENV_BASE_PATH = "HINDSIGHT_API_BASE_PATH"
|
||||
ENV_LOG_LEVEL = "HINDSIGHT_API_LOG_LEVEL"
|
||||
ENV_LOG_FORMAT = "HINDSIGHT_API_LOG_FORMAT"
|
||||
ENV_LOG_JSON_FIELDS = "HINDSIGHT_API_LOG_JSON_FIELDS"
|
||||
ENV_WORKERS = "HINDSIGHT_API_WORKERS"
|
||||
ENV_MCP_ENABLED = "HINDSIGHT_API_MCP_ENABLED"
|
||||
ENV_MCP_ENABLED_TOOLS = "HINDSIGHT_API_MCP_ENABLED_TOOLS"
|
||||
@@ -339,15 +333,12 @@ ENV_FILE_DELETE_AFTER_RETAIN = "HINDSIGHT_API_FILE_DELETE_AFTER_RETAIN"
|
||||
# Observations settings (consolidated knowledge from facts)
|
||||
ENV_ENABLE_OBSERVATIONS = "HINDSIGHT_API_ENABLE_OBSERVATIONS"
|
||||
ENV_CONSOLIDATION_BATCH_SIZE = "HINDSIGHT_API_CONSOLIDATION_BATCH_SIZE"
|
||||
ENV_CONSOLIDATION_MAX_MEMORIES_PER_ROUND = "HINDSIGHT_API_CONSOLIDATION_MAX_MEMORIES_PER_ROUND"
|
||||
ENV_CONSOLIDATION_LLM_BATCH_SIZE = "HINDSIGHT_API_CONSOLIDATION_LLM_BATCH_SIZE"
|
||||
ENV_CONSOLIDATION_MAX_TOKENS = "HINDSIGHT_API_CONSOLIDATION_MAX_TOKENS"
|
||||
ENV_CONSOLIDATION_SOURCE_FACTS_MAX_TOKENS = "HINDSIGHT_API_CONSOLIDATION_SOURCE_FACTS_MAX_TOKENS"
|
||||
ENV_CONSOLIDATION_SOURCE_FACTS_MAX_TOKENS_PER_OBSERVATION = (
|
||||
"HINDSIGHT_API_CONSOLIDATION_SOURCE_FACTS_MAX_TOKENS_PER_OBSERVATION"
|
||||
)
|
||||
ENV_CONSOLIDATION_RECALL_BUDGET = "HINDSIGHT_API_CONSOLIDATION_RECALL_BUDGET"
|
||||
ENV_CONSOLIDATION_MAX_ATTEMPTS = "HINDSIGHT_API_CONSOLIDATION_MAX_ATTEMPTS"
|
||||
ENV_OBSERVATIONS_MISSION = "HINDSIGHT_API_OBSERVATIONS_MISSION"
|
||||
ENV_MAX_OBSERVATIONS_PER_SCOPE = "HINDSIGHT_API_MAX_OBSERVATIONS_PER_SCOPE"
|
||||
ENV_ENABLE_OBSERVATION_HISTORY = "HINDSIGHT_API_ENABLE_OBSERVATION_HISTORY"
|
||||
@@ -379,7 +370,6 @@ ENV_DB_POOL_MIN_SIZE = "HINDSIGHT_API_DB_POOL_MIN_SIZE"
|
||||
ENV_DB_POOL_MAX_SIZE = "HINDSIGHT_API_DB_POOL_MAX_SIZE"
|
||||
ENV_DB_COMMAND_TIMEOUT = "HINDSIGHT_API_DB_COMMAND_TIMEOUT"
|
||||
ENV_DB_ACQUIRE_TIMEOUT = "HINDSIGHT_API_DB_ACQUIRE_TIMEOUT"
|
||||
ENV_DB_STATEMENT_TIMEOUT = "HINDSIGHT_API_DB_STATEMENT_TIMEOUT"
|
||||
|
||||
# Worker configuration (distributed task processing)
|
||||
ENV_WORKER_ENABLED = "HINDSIGHT_API_WORKER_ENABLED"
|
||||
@@ -388,18 +378,7 @@ ENV_WORKER_POLL_INTERVAL_MS = "HINDSIGHT_API_WORKER_POLL_INTERVAL_MS"
|
||||
ENV_WORKER_MAX_RETRIES = "HINDSIGHT_API_WORKER_MAX_RETRIES"
|
||||
ENV_WORKER_HTTP_PORT = "HINDSIGHT_API_WORKER_HTTP_PORT"
|
||||
ENV_WORKER_MAX_SLOTS = "HINDSIGHT_API_WORKER_MAX_SLOTS"
|
||||
|
||||
# Per-operation-type slot reservations. Each entry maps an operation_type
|
||||
# (as stored in async_operations.operation_type) to its env var and default.
|
||||
# Adding a new operation type here is the ONLY change needed to make it
|
||||
# reservable via env var — config fields, from_env(), and the
|
||||
# worker_slot_reservations property all derive from this dict.
|
||||
WORKER_SLOT_RESERVATION_TYPES: dict[str, tuple[str, int]] = {
|
||||
"consolidation": ("HINDSIGHT_API_WORKER_CONSOLIDATION_MAX_SLOTS", 2),
|
||||
"retain": ("HINDSIGHT_API_WORKER_RETAIN_MAX_SLOTS", 0),
|
||||
"file_convert_retain": ("HINDSIGHT_API_WORKER_FILE_CONVERT_RETAIN_MAX_SLOTS", 0),
|
||||
"refresh_mental_model": ("HINDSIGHT_API_WORKER_REFRESH_MENTAL_MODEL_MAX_SLOTS", 0),
|
||||
}
|
||||
ENV_WORKER_CONSOLIDATION_MAX_SLOTS = "HINDSIGHT_API_WORKER_CONSOLIDATION_MAX_SLOTS"
|
||||
ENV_RETAIN_MAX_CONCURRENT = "HINDSIGHT_API_RETAIN_MAX_CONCURRENT"
|
||||
|
||||
# Reflect agent settings
|
||||
@@ -408,20 +387,6 @@ ENV_REFLECT_MAX_CONTEXT_TOKENS = "HINDSIGHT_API_REFLECT_MAX_CONTEXT_TOKENS"
|
||||
ENV_REFLECT_WALL_TIMEOUT = "HINDSIGHT_API_REFLECT_WALL_TIMEOUT"
|
||||
ENV_REFLECT_MISSION = "HINDSIGHT_API_REFLECT_MISSION"
|
||||
ENV_REFLECT_SOURCE_FACTS_MAX_TOKENS = "HINDSIGHT_API_REFLECT_SOURCE_FACTS_MAX_TOKENS"
|
||||
ENV_RECALL_INCLUDE_CHUNKS = "HINDSIGHT_API_RECALL_INCLUDE_CHUNKS"
|
||||
ENV_RECALL_MAX_TOKENS = "HINDSIGHT_API_RECALL_MAX_TOKENS"
|
||||
ENV_RECALL_CHUNKS_MAX_TOKENS = "HINDSIGHT_API_RECALL_CHUNKS_MAX_TOKENS"
|
||||
|
||||
# Recall budget mapping (budget enum -> thinking_budget integer)
|
||||
ENV_RECALL_BUDGET_FUNCTION = "HINDSIGHT_API_RECALL_BUDGET_FUNCTION"
|
||||
ENV_RECALL_BUDGET_FIXED_LOW = "HINDSIGHT_API_RECALL_BUDGET_FIXED_LOW"
|
||||
ENV_RECALL_BUDGET_FIXED_MID = "HINDSIGHT_API_RECALL_BUDGET_FIXED_MID"
|
||||
ENV_RECALL_BUDGET_FIXED_HIGH = "HINDSIGHT_API_RECALL_BUDGET_FIXED_HIGH"
|
||||
ENV_RECALL_BUDGET_ADAPTIVE_LOW = "HINDSIGHT_API_RECALL_BUDGET_ADAPTIVE_LOW"
|
||||
ENV_RECALL_BUDGET_ADAPTIVE_MID = "HINDSIGHT_API_RECALL_BUDGET_ADAPTIVE_MID"
|
||||
ENV_RECALL_BUDGET_ADAPTIVE_HIGH = "HINDSIGHT_API_RECALL_BUDGET_ADAPTIVE_HIGH"
|
||||
ENV_RECALL_BUDGET_MIN = "HINDSIGHT_API_RECALL_BUDGET_MIN"
|
||||
ENV_RECALL_BUDGET_MAX = "HINDSIGHT_API_RECALL_BUDGET_MAX"
|
||||
|
||||
# Audit log settings
|
||||
ENV_AUDIT_LOG_ENABLED = "HINDSIGHT_API_AUDIT_LOG_ENABLED"
|
||||
@@ -467,7 +432,7 @@ DEFAULT_LLAMACPP_NO_GRAMMAR = False # True = disable JSON grammar enforcement (
|
||||
DEFAULT_LLAMACPP_EXTRA_ARGS = None # Space-separated extra CLI args for llama.cpp server
|
||||
|
||||
DEFAULT_LLM_MAX_CONCURRENT = 32
|
||||
DEFAULT_LLM_MAX_RETRIES = 3 # Max retry attempts for LLM API calls
|
||||
DEFAULT_LLM_MAX_RETRIES = 10 # Max retry attempts for LLM API calls
|
||||
DEFAULT_LLM_INITIAL_BACKOFF = 1.0 # Initial backoff in seconds for retry exponential backoff
|
||||
DEFAULT_LLM_MAX_BACKOFF = 60.0 # Max backoff cap in seconds for retry exponential backoff
|
||||
DEFAULT_LLM_TIMEOUT = 120.0 # seconds
|
||||
@@ -485,10 +450,8 @@ DEFAULT_EMBEDDINGS_LOCAL_MODEL = "BAAI/bge-small-en-v1.5"
|
||||
DEFAULT_EMBEDDINGS_LOCAL_FORCE_CPU = False # Force CPU mode for local embeddings (avoids MPS/XPC issues on macOS)
|
||||
DEFAULT_EMBEDDINGS_LOCAL_TRUST_REMOTE_CODE = False # Security: disabled by default, required for some models
|
||||
DEFAULT_EMBEDDINGS_OPENAI_MODEL = "text-embedding-3-small"
|
||||
DEFAULT_EMBEDDINGS_OPENAI_BATCH_SIZE = 100
|
||||
DEFAULT_EMBEDDINGS_GEMINI_MODEL = "gemini-embedding-001"
|
||||
DEFAULT_EMBEDDINGS_GEMINI_OUTPUT_DIMENSIONALITY = 768
|
||||
DEFAULT_EMBEDDINGS_GEMINI_FORCE_IPV4 = False
|
||||
DEFAULT_EMBEDDING_DIMENSION = 384
|
||||
|
||||
DEFAULT_RERANKER_PROVIDER = "local"
|
||||
@@ -503,11 +466,9 @@ DEFAULT_RERANKER_LOCAL_BUCKET_BATCHING = False # Length-sorted bucket batching:
|
||||
DEFAULT_RERANKER_LOCAL_BATCH_SIZE = 32 # Batch size for local reranker predict() calls
|
||||
DEFAULT_RERANKER_TEI_BATCH_SIZE = 128
|
||||
DEFAULT_RERANKER_TEI_MAX_CONCURRENT = 8
|
||||
DEFAULT_RERANKER_TEI_HTTP_TIMEOUT = 30.0 # HTTP timeout for TEI reranker requests (seconds)
|
||||
DEFAULT_RERANKER_MAX_CANDIDATES = 300
|
||||
DEFAULT_RERANKER_FLASHRANK_MODEL = "ms-marco-MiniLM-L-12-v2" # Best balance of speed and quality
|
||||
DEFAULT_RERANKER_FLASHRANK_CACHE_DIR = None # Use default cache directory
|
||||
DEFAULT_RERANKER_FLASHRANK_CPU_MEM_ARENA = False # Disable ONNX CPU memory arena to bound RSS
|
||||
|
||||
DEFAULT_EMBEDDINGS_COHERE_MODEL = "embed-english-v3.0"
|
||||
DEFAULT_RERANKER_COHERE_MODEL = "rerank-english-v3.0"
|
||||
@@ -590,17 +551,12 @@ DEFAULT_FILE_DELETE_AFTER_RETAIN = True # Delete file bytes after retain (saves
|
||||
DEFAULT_ENABLE_OBSERVATIONS = True # Observations enabled by default
|
||||
DEFAULT_ENABLE_OBSERVATION_HISTORY = True # Observation history tracking enabled by default
|
||||
DEFAULT_ENABLE_MENTAL_MODEL_HISTORY = True # Mental model history tracking enabled by default
|
||||
DEFAULT_CONSOLIDATION_MAX_ATTEMPTS = 3 # Outer retry attempts for consolidation LLM batch calls
|
||||
DEFAULT_CONSOLIDATION_BATCH_SIZE = 50 # Memories to load per batch (internal memory optimization)
|
||||
DEFAULT_CONSOLIDATION_MAX_MEMORIES_PER_ROUND = (
|
||||
100 # Max memories per consolidation round (0 = unlimited). Limits how long one bank holds a worker slot.
|
||||
)
|
||||
DEFAULT_CONSOLIDATION_LLM_BATCH_SIZE = 8 # Facts per LLM call (1 = no batching; >1 = batch mode)
|
||||
DEFAULT_CONSOLIDATION_MAX_TOKENS = 512 # Max tokens for recall when finding related observations
|
||||
DEFAULT_CONSOLIDATION_RECALL_BUDGET = "low" # Budget level for consolidation recall (low/mid/high)
|
||||
DEFAULT_CONSOLIDATION_SOURCE_FACTS_MAX_TOKENS = (
|
||||
4096 # Total token budget for source facts in consolidation recall (-1 = unlimited)
|
||||
)
|
||||
-1
|
||||
) # Total token budget for source facts in consolidation recall (-1 = unlimited)
|
||||
DEFAULT_CONSOLIDATION_SOURCE_FACTS_MAX_TOKENS_PER_OBSERVATION = (
|
||||
256 # Max tokens of source facts per observation in consolidation prompt (-1 = unlimited)
|
||||
)
|
||||
@@ -615,7 +571,6 @@ DEFAULT_DB_POOL_MIN_SIZE = 5
|
||||
DEFAULT_DB_POOL_MAX_SIZE = 100
|
||||
DEFAULT_DB_COMMAND_TIMEOUT = 60 # seconds
|
||||
DEFAULT_DB_ACQUIRE_TIMEOUT = 30 # seconds
|
||||
DEFAULT_DB_STATEMENT_TIMEOUT = 600 # seconds (Postgres statement_timeout applied on every pool connection; 0 disables)
|
||||
|
||||
# Worker configuration (distributed task processing)
|
||||
DEFAULT_WORKER_ENABLED = True # API runs worker by default (standalone mode)
|
||||
@@ -624,6 +579,7 @@ DEFAULT_WORKER_POLL_INTERVAL_MS = 500 # Poll database every 500ms
|
||||
DEFAULT_WORKER_MAX_RETRIES = 3 # Max retries before marking task failed
|
||||
DEFAULT_WORKER_HTTP_PORT = 8889 # HTTP port for worker metrics/health
|
||||
DEFAULT_WORKER_MAX_SLOTS = 10 # Total concurrent tasks per worker
|
||||
DEFAULT_WORKER_CONSOLIDATION_MAX_SLOTS = 2 # Max concurrent consolidation tasks per worker
|
||||
DEFAULT_RETAIN_MAX_CONCURRENT = 4 # Max concurrent retain DB phases (HNSW reads + writes). Limits I/O contention.
|
||||
|
||||
# Reflect agent settings
|
||||
@@ -631,25 +587,6 @@ DEFAULT_REFLECT_MAX_ITERATIONS = 10 # Max tool call iterations before forcing r
|
||||
DEFAULT_REFLECT_MAX_CONTEXT_TOKENS = 100_000 # Max accumulated context tokens before forcing final prompt
|
||||
DEFAULT_REFLECT_WALL_TIMEOUT = 300 # Wall-clock timeout in seconds for the entire reflect operation (5 minutes)
|
||||
DEFAULT_REFLECT_SOURCE_FACTS_MAX_TOKENS = -1 # Token budget for source facts in search_observations (-1 = disabled)
|
||||
DEFAULT_RECALL_INCLUDE_CHUNKS = True # Whether internal recall (e.g. mental model refresh) returns raw chunks
|
||||
DEFAULT_RECALL_MAX_TOKENS = 2048 # Token budget for facts returned by internal recall
|
||||
DEFAULT_RECALL_CHUNKS_MAX_TOKENS = 1000 # Token budget for raw chunks returned by internal recall
|
||||
|
||||
# Recall budget mapping
|
||||
# "fixed": thinking_budget = recall_budget_fixed_<level> (preserves legacy behavior)
|
||||
# "adaptive": thinking_budget = round(max_tokens * recall_budget_adaptive_<level>),
|
||||
# clamped to [recall_budget_min, recall_budget_max]
|
||||
RECALL_BUDGET_FUNCTIONS = ("fixed", "adaptive")
|
||||
DEFAULT_RECALL_BUDGET_FUNCTION = "fixed"
|
||||
DEFAULT_RECALL_BUDGET_FIXED_LOW = 100
|
||||
DEFAULT_RECALL_BUDGET_FIXED_MID = 300
|
||||
DEFAULT_RECALL_BUDGET_FIXED_HIGH = 1000
|
||||
# Adaptive defaults chosen to roughly match fixed defaults at max_tokens=4096
|
||||
DEFAULT_RECALL_BUDGET_ADAPTIVE_LOW = 0.025
|
||||
DEFAULT_RECALL_BUDGET_ADAPTIVE_MID = 0.075
|
||||
DEFAULT_RECALL_BUDGET_ADAPTIVE_HIGH = 0.25
|
||||
DEFAULT_RECALL_BUDGET_MIN = 20 # Floor for the adaptive function
|
||||
DEFAULT_RECALL_BUDGET_MAX = 2000 # Ceiling for the adaptive function
|
||||
|
||||
# Disposition defaults (None = not set, fall back to bank DB value or 3)
|
||||
DEFAULT_DISPOSITION_SKEPTICISM = None
|
||||
@@ -712,10 +649,6 @@ class JsonFormatter(logging.Formatter):
|
||||
logging.CRITICAL: "CRITICAL",
|
||||
}
|
||||
|
||||
def __init__(self, allowed_fields: frozenset[str] | None = None):
|
||||
super().__init__()
|
||||
self._allowed_fields = allowed_fields
|
||||
|
||||
def format(self, record: logging.LogRecord) -> str:
|
||||
log_entry = {
|
||||
"severity": self.SEVERITY_MAP.get(record.levelno, "DEFAULT"),
|
||||
@@ -724,20 +657,10 @@ class JsonFormatter(logging.Formatter):
|
||||
"logger": record.name,
|
||||
}
|
||||
|
||||
# Lazy import to avoid circular dependency (engine imports from config).
|
||||
from hindsight_api.engine.memory_engine import _current_schema
|
||||
|
||||
tenant = _current_schema.get()
|
||||
if tenant:
|
||||
log_entry["tenant"] = tenant
|
||||
|
||||
# Add exception info if present
|
||||
if record.exc_info:
|
||||
log_entry["exception"] = self.formatException(record.exc_info)
|
||||
|
||||
if self._allowed_fields is not None:
|
||||
log_entry = {k: v for k, v in log_entry.items() if k in self._allowed_fields}
|
||||
|
||||
return json.dumps(log_entry)
|
||||
|
||||
|
||||
@@ -746,25 +669,6 @@ def _parse_str_list(value: str) -> list[str]:
|
||||
return [v.strip() for v in value.split(",") if v.strip()]
|
||||
|
||||
|
||||
def _parse_positive_int(name: str, raw: str | None, default: int) -> int:
|
||||
"""
|
||||
Parse an env var that must be a positive integer (>= 1).
|
||||
|
||||
Falls back to ``default`` when unset/empty. Raises ValueError on non-integer
|
||||
or non-positive values so misconfiguration fails fast instead of triggering
|
||||
infinite loops or zero-step range() calls downstream.
|
||||
"""
|
||||
if raw is None or raw == "":
|
||||
return default
|
||||
try:
|
||||
parsed = int(raw)
|
||||
except ValueError as e:
|
||||
raise ValueError(f"{name} must be an integer, got {raw!r}") from e
|
||||
if parsed < 1:
|
||||
raise ValueError(f"{name} must be >= 1, got {parsed}")
|
||||
return parsed
|
||||
|
||||
|
||||
def _validate_extraction_mode(mode: str) -> str:
|
||||
"""Validate and normalize extraction mode."""
|
||||
mode_lower = mode.lower()
|
||||
@@ -777,18 +681,6 @@ def _validate_extraction_mode(mode: str) -> str:
|
||||
return mode_lower
|
||||
|
||||
|
||||
def _validate_recall_budget_function(function: str) -> str:
|
||||
"""Validate and normalize recall budget function."""
|
||||
function_lower = function.lower()
|
||||
if function_lower not in RECALL_BUDGET_FUNCTIONS:
|
||||
logger.warning(
|
||||
f"Invalid recall budget function '{function}', must be one of {RECALL_BUDGET_FUNCTIONS}. "
|
||||
f"Defaulting to '{DEFAULT_RECALL_BUDGET_FUNCTION}'."
|
||||
)
|
||||
return DEFAULT_RECALL_BUDGET_FUNCTION
|
||||
return function_lower
|
||||
|
||||
|
||||
def _get_default_model_for_provider(provider: str) -> str:
|
||||
"""Get the default model for a given provider."""
|
||||
return PROVIDER_DEFAULT_MODELS.get(provider.lower(), DEFAULT_LLM_MODEL)
|
||||
@@ -898,7 +790,6 @@ class HindsightConfig:
|
||||
embeddings_cohere_api_key: str | None
|
||||
embeddings_cohere_model: str
|
||||
embeddings_cohere_base_url: str | None
|
||||
embeddings_cohere_output_dimensions: int | None
|
||||
embeddings_openrouter_api_key: str | None
|
||||
embeddings_openrouter_model: str
|
||||
embeddings_litellm_api_base: str
|
||||
@@ -913,7 +804,6 @@ class HindsightConfig:
|
||||
embeddings_gemini_api_key: str | None
|
||||
embeddings_gemini_model: str
|
||||
embeddings_gemini_output_dimensionality: int | None
|
||||
embeddings_gemini_force_ipv4: bool
|
||||
embeddings_vertexai_project_id: str | None
|
||||
embeddings_vertexai_region: str | None
|
||||
embeddings_vertexai_service_account_key: str | None
|
||||
@@ -930,7 +820,6 @@ class HindsightConfig:
|
||||
reranker_tei_url: str | None
|
||||
reranker_tei_batch_size: int
|
||||
reranker_tei_max_concurrent: int
|
||||
reranker_tei_http_timeout: float
|
||||
reranker_max_candidates: int
|
||||
reranker_cohere_api_key: str | None
|
||||
reranker_cohere_model: str
|
||||
@@ -960,7 +849,6 @@ class HindsightConfig:
|
||||
base_path: str
|
||||
log_level: str
|
||||
log_format: str
|
||||
log_json_fields: list[str] | None # None = all fields; explicit list = allowlist
|
||||
mcp_enabled: bool
|
||||
mcp_enabled_tools: list[str] | None # None = all tools; explicit list = allowlist
|
||||
mcp_stateless: bool # True = stateless HTTP (POST-only); False = stateful (supports GET/SSE)
|
||||
@@ -1019,13 +907,10 @@ class HindsightConfig:
|
||||
enable_observation_history: bool
|
||||
enable_mental_model_history: bool
|
||||
consolidation_batch_size: int
|
||||
consolidation_max_memories_per_round: int
|
||||
consolidation_llm_batch_size: int
|
||||
consolidation_max_tokens: int
|
||||
consolidation_recall_budget: str
|
||||
consolidation_source_facts_max_tokens: int
|
||||
consolidation_source_facts_max_tokens_per_observation: int
|
||||
consolidation_max_attempts: int
|
||||
observations_mission: str | None
|
||||
max_observations_per_scope: int
|
||||
|
||||
@@ -1040,25 +925,6 @@ class HindsightConfig:
|
||||
reflect_mission: str | None
|
||||
reflect_source_facts_max_tokens: int
|
||||
|
||||
# Recall settings (used by internal recall, e.g. during mental model refresh)
|
||||
recall_include_chunks: bool
|
||||
recall_max_tokens: int
|
||||
recall_chunks_max_tokens: int
|
||||
|
||||
# Recall budget mapping: how the Budget enum (LOW/MID/HIGH) maps to thinking_budget integer.
|
||||
# function="fixed": use the recall_budget_fixed_* values directly (legacy behavior).
|
||||
# function="adaptive": compute round(max_tokens * recall_budget_adaptive_*),
|
||||
# clamped to [recall_budget_min, recall_budget_max].
|
||||
recall_budget_function: str
|
||||
recall_budget_fixed_low: int
|
||||
recall_budget_fixed_mid: int
|
||||
recall_budget_fixed_high: int
|
||||
recall_budget_adaptive_low: float
|
||||
recall_budget_adaptive_mid: float
|
||||
recall_budget_adaptive_high: float
|
||||
recall_budget_min: int
|
||||
recall_budget_max: int
|
||||
|
||||
# Disposition settings (hierarchical - can be overridden per bank; None = fall back to DB)
|
||||
disposition_skepticism: int | None
|
||||
disposition_literalism: int | None
|
||||
@@ -1076,7 +942,6 @@ class HindsightConfig:
|
||||
db_pool_max_size: int
|
||||
db_command_timeout: int
|
||||
db_acquire_timeout: int
|
||||
db_statement_timeout: int
|
||||
|
||||
# Worker configuration (distributed task processing)
|
||||
worker_enabled: bool
|
||||
@@ -1085,7 +950,7 @@ class HindsightConfig:
|
||||
worker_max_retries: int
|
||||
worker_http_port: int
|
||||
worker_max_slots: int
|
||||
worker_slot_reservations: dict[str, int]
|
||||
worker_consolidation_max_slots: int
|
||||
retain_max_concurrent: int
|
||||
|
||||
# Reflect agent settings
|
||||
@@ -1112,10 +977,6 @@ class HindsightConfig:
|
||||
webhook_event_types: list[str] # Event types to deliver globally
|
||||
webhook_delivery_poll_interval_seconds: int # How often the delivery worker polls
|
||||
|
||||
# Defaulted fields (source-compatible additions — existing direct constructor callers keep working).
|
||||
# Keep at the end of the dataclass; Python forbids non-default fields after default fields.
|
||||
embeddings_openai_batch_size: int = DEFAULT_EMBEDDINGS_OPENAI_BATCH_SIZE
|
||||
|
||||
# Class-level sets for configuration categorization
|
||||
|
||||
# CREDENTIAL_FIELDS: Never exposed via API, never configurable per-tenant/bank
|
||||
@@ -1170,7 +1031,6 @@ class HindsightConfig:
|
||||
# Consolidation settings
|
||||
"enable_observations",
|
||||
"consolidation_llm_batch_size",
|
||||
"consolidation_max_memories_per_round",
|
||||
"consolidation_source_facts_max_tokens",
|
||||
"consolidation_source_facts_max_tokens_per_observation",
|
||||
"observations_mission",
|
||||
@@ -1178,20 +1038,6 @@ class HindsightConfig:
|
||||
# Reflect settings
|
||||
"reflect_mission",
|
||||
"reflect_source_facts_max_tokens",
|
||||
# Recall settings (used by internal recall, e.g. mental model refresh)
|
||||
"recall_include_chunks",
|
||||
"recall_max_tokens",
|
||||
"recall_chunks_max_tokens",
|
||||
# Recall budget mapping (Budget enum -> thinking_budget integer)
|
||||
"recall_budget_function",
|
||||
"recall_budget_fixed_low",
|
||||
"recall_budget_fixed_mid",
|
||||
"recall_budget_fixed_high",
|
||||
"recall_budget_adaptive_low",
|
||||
"recall_budget_adaptive_mid",
|
||||
"recall_budget_adaptive_high",
|
||||
"recall_budget_min",
|
||||
"recall_budget_max",
|
||||
# Disposition settings
|
||||
"disposition_skepticism",
|
||||
"disposition_literalism",
|
||||
@@ -1298,16 +1144,6 @@ class HindsightConfig:
|
||||
f"provider: {self.retain_llm_provider or self.llm_provider})"
|
||||
)
|
||||
|
||||
# Validate that sum of per-operation slot reservations does not exceed max_slots
|
||||
total_reserved = sum(self.worker_slot_reservations.values())
|
||||
if total_reserved > self.worker_max_slots:
|
||||
reservation_details = ", ".join(f"{k}={v}" for k, v in self.worker_slot_reservations.items() if v > 0)
|
||||
raise ValueError(
|
||||
f"Sum of per-operation slot reservations ({total_reserved}: {reservation_details}) "
|
||||
f"exceeds worker_max_slots ({self.worker_max_slots}). "
|
||||
f"Reduce reservations or increase HINDSIGHT_API_WORKER_MAX_SLOTS."
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_env(cls) -> "HindsightConfig":
|
||||
"""Create configuration from environment variables."""
|
||||
@@ -1434,18 +1270,10 @@ class HindsightConfig:
|
||||
in ("true", "1"),
|
||||
embeddings_tei_url=os.getenv(ENV_EMBEDDINGS_TEI_URL),
|
||||
embeddings_openai_base_url=os.getenv(ENV_EMBEDDINGS_OPENAI_BASE_URL) or None,
|
||||
embeddings_openai_batch_size=_parse_positive_int(
|
||||
ENV_EMBEDDINGS_OPENAI_BATCH_SIZE,
|
||||
os.getenv(ENV_EMBEDDINGS_OPENAI_BATCH_SIZE),
|
||||
DEFAULT_EMBEDDINGS_OPENAI_BATCH_SIZE,
|
||||
),
|
||||
# Cohere embeddings (with backward-compatible fallback to shared API key)
|
||||
embeddings_cohere_api_key=os.getenv(ENV_EMBEDDINGS_COHERE_API_KEY) or os.getenv(ENV_COHERE_API_KEY),
|
||||
embeddings_cohere_model=os.getenv(ENV_EMBEDDINGS_COHERE_MODEL, DEFAULT_EMBEDDINGS_COHERE_MODEL),
|
||||
embeddings_cohere_base_url=os.getenv(ENV_EMBEDDINGS_COHERE_BASE_URL) or None,
|
||||
embeddings_cohere_output_dimensions=int(v)
|
||||
if (v := os.getenv(ENV_EMBEDDINGS_COHERE_OUTPUT_DIMENSIONS))
|
||||
else None,
|
||||
# OpenRouter embeddings (with fallback to shared OpenRouter key, then LLM key)
|
||||
embeddings_openrouter_api_key=os.getenv(ENV_EMBEDDINGS_OPENROUTER_API_KEY)
|
||||
or os.getenv(ENV_OPENROUTER_API_KEY)
|
||||
@@ -1477,11 +1305,6 @@ class HindsightConfig:
|
||||
str(DEFAULT_EMBEDDINGS_GEMINI_OUTPUT_DIMENSIONALITY),
|
||||
)
|
||||
),
|
||||
embeddings_gemini_force_ipv4=os.getenv(
|
||||
ENV_EMBEDDINGS_GEMINI_FORCE_IPV4,
|
||||
str(DEFAULT_EMBEDDINGS_GEMINI_FORCE_IPV4),
|
||||
).lower()
|
||||
in ("true", "1"),
|
||||
embeddings_vertexai_project_id=os.getenv(ENV_EMBEDDINGS_VERTEXAI_PROJECT_ID)
|
||||
or os.getenv(ENV_LLM_VERTEXAI_PROJECT_ID),
|
||||
embeddings_vertexai_region=os.getenv(ENV_EMBEDDINGS_VERTEXAI_REGION) or os.getenv(ENV_LLM_VERTEXAI_REGION),
|
||||
@@ -1515,9 +1338,6 @@ class HindsightConfig:
|
||||
reranker_tei_max_concurrent=int(
|
||||
os.getenv(ENV_RERANKER_TEI_MAX_CONCURRENT, str(DEFAULT_RERANKER_TEI_MAX_CONCURRENT))
|
||||
),
|
||||
reranker_tei_http_timeout=float(
|
||||
os.getenv(ENV_RERANKER_TEI_HTTP_TIMEOUT, str(DEFAULT_RERANKER_TEI_HTTP_TIMEOUT))
|
||||
),
|
||||
reranker_max_candidates=int(os.getenv(ENV_RERANKER_MAX_CANDIDATES, str(DEFAULT_RERANKER_MAX_CANDIDATES))),
|
||||
# Cohere reranker (with backward-compatible fallback to shared API key)
|
||||
reranker_cohere_api_key=os.getenv(ENV_RERANKER_COHERE_API_KEY) or os.getenv(ENV_COHERE_API_KEY),
|
||||
@@ -1562,7 +1382,6 @@ class HindsightConfig:
|
||||
base_path=os.getenv(ENV_BASE_PATH, DEFAULT_BASE_PATH),
|
||||
log_level=os.getenv(ENV_LOG_LEVEL, DEFAULT_LOG_LEVEL),
|
||||
log_format=os.getenv(ENV_LOG_FORMAT, DEFAULT_LOG_FORMAT).lower(),
|
||||
log_json_fields=_parse_str_list(os.getenv(ENV_LOG_JSON_FIELDS, "")) or None,
|
||||
mcp_enabled=os.getenv(ENV_MCP_ENABLED, str(DEFAULT_MCP_ENABLED)).lower() == "true",
|
||||
mcp_enabled_tools=[t.strip() for t in os.getenv(ENV_MCP_ENABLED_TOOLS).split(",") if t.strip()]
|
||||
if os.getenv(ENV_MCP_ENABLED_TOOLS)
|
||||
@@ -1655,19 +1474,12 @@ class HindsightConfig:
|
||||
consolidation_batch_size=int(
|
||||
os.getenv(ENV_CONSOLIDATION_BATCH_SIZE, str(DEFAULT_CONSOLIDATION_BATCH_SIZE))
|
||||
),
|
||||
consolidation_max_memories_per_round=int(
|
||||
os.getenv(
|
||||
ENV_CONSOLIDATION_MAX_MEMORIES_PER_ROUND,
|
||||
str(DEFAULT_CONSOLIDATION_MAX_MEMORIES_PER_ROUND),
|
||||
)
|
||||
),
|
||||
consolidation_llm_batch_size=int(
|
||||
os.getenv(ENV_CONSOLIDATION_LLM_BATCH_SIZE, str(DEFAULT_CONSOLIDATION_LLM_BATCH_SIZE))
|
||||
),
|
||||
consolidation_max_tokens=int(
|
||||
os.getenv(ENV_CONSOLIDATION_MAX_TOKENS, str(DEFAULT_CONSOLIDATION_MAX_TOKENS))
|
||||
),
|
||||
consolidation_recall_budget=os.getenv(ENV_CONSOLIDATION_RECALL_BUDGET, DEFAULT_CONSOLIDATION_RECALL_BUDGET),
|
||||
consolidation_source_facts_max_tokens=int(
|
||||
os.getenv(ENV_CONSOLIDATION_SOURCE_FACTS_MAX_TOKENS, str(DEFAULT_CONSOLIDATION_SOURCE_FACTS_MAX_TOKENS))
|
||||
),
|
||||
@@ -1677,9 +1489,6 @@ class HindsightConfig:
|
||||
str(DEFAULT_CONSOLIDATION_SOURCE_FACTS_MAX_TOKENS_PER_OBSERVATION),
|
||||
)
|
||||
),
|
||||
consolidation_max_attempts=int(
|
||||
os.getenv(ENV_CONSOLIDATION_MAX_ATTEMPTS, str(DEFAULT_CONSOLIDATION_MAX_ATTEMPTS))
|
||||
),
|
||||
observations_mission=os.getenv(ENV_OBSERVATIONS_MISSION) or DEFAULT_OBSERVATIONS_MISSION,
|
||||
max_observations_per_scope=int(
|
||||
os.getenv(ENV_MAX_OBSERVATIONS_PER_SCOPE, str(DEFAULT_MAX_OBSERVATIONS_PER_SCOPE))
|
||||
@@ -1693,7 +1502,6 @@ class HindsightConfig:
|
||||
db_pool_max_size=int(os.getenv(ENV_DB_POOL_MAX_SIZE, str(DEFAULT_DB_POOL_MAX_SIZE))),
|
||||
db_command_timeout=int(os.getenv(ENV_DB_COMMAND_TIMEOUT, str(DEFAULT_DB_COMMAND_TIMEOUT))),
|
||||
db_acquire_timeout=int(os.getenv(ENV_DB_ACQUIRE_TIMEOUT, str(DEFAULT_DB_ACQUIRE_TIMEOUT))),
|
||||
db_statement_timeout=int(os.getenv(ENV_DB_STATEMENT_TIMEOUT, str(DEFAULT_DB_STATEMENT_TIMEOUT))),
|
||||
# Worker configuration
|
||||
worker_enabled=os.getenv(ENV_WORKER_ENABLED, str(DEFAULT_WORKER_ENABLED)).lower() == "true",
|
||||
worker_id=os.getenv(ENV_WORKER_ID) or DEFAULT_WORKER_ID,
|
||||
@@ -1701,11 +1509,9 @@ class HindsightConfig:
|
||||
worker_max_retries=int(os.getenv(ENV_WORKER_MAX_RETRIES, str(DEFAULT_WORKER_MAX_RETRIES))),
|
||||
worker_http_port=int(os.getenv(ENV_WORKER_HTTP_PORT, str(DEFAULT_WORKER_HTTP_PORT))),
|
||||
worker_max_slots=int(os.getenv(ENV_WORKER_MAX_SLOTS, str(DEFAULT_WORKER_MAX_SLOTS))),
|
||||
worker_slot_reservations={
|
||||
op_type: int(os.getenv(env_var, str(default)))
|
||||
for op_type, (env_var, default) in WORKER_SLOT_RESERVATION_TYPES.items()
|
||||
if int(os.getenv(env_var, str(default))) > 0
|
||||
},
|
||||
worker_consolidation_max_slots=int(
|
||||
os.getenv(ENV_WORKER_CONSOLIDATION_MAX_SLOTS, str(DEFAULT_WORKER_CONSOLIDATION_MAX_SLOTS))
|
||||
),
|
||||
retain_max_concurrent=int(os.getenv(ENV_RETAIN_MAX_CONCURRENT, str(DEFAULT_RETAIN_MAX_CONCURRENT))),
|
||||
# Reflect agent settings
|
||||
reflect_max_iterations=int(os.getenv(ENV_REFLECT_MAX_ITERATIONS, str(DEFAULT_REFLECT_MAX_ITERATIONS))),
|
||||
@@ -1717,31 +1523,6 @@ class HindsightConfig:
|
||||
reflect_source_facts_max_tokens=int(
|
||||
os.getenv(ENV_REFLECT_SOURCE_FACTS_MAX_TOKENS, str(DEFAULT_REFLECT_SOURCE_FACTS_MAX_TOKENS))
|
||||
),
|
||||
recall_include_chunks=os.getenv(ENV_RECALL_INCLUDE_CHUNKS, str(DEFAULT_RECALL_INCLUDE_CHUNKS)).lower()
|
||||
in ("true", "1", "yes"),
|
||||
recall_max_tokens=int(os.getenv(ENV_RECALL_MAX_TOKENS, str(DEFAULT_RECALL_MAX_TOKENS))),
|
||||
recall_chunks_max_tokens=int(
|
||||
os.getenv(ENV_RECALL_CHUNKS_MAX_TOKENS, str(DEFAULT_RECALL_CHUNKS_MAX_TOKENS))
|
||||
),
|
||||
recall_budget_function=_validate_recall_budget_function(
|
||||
os.getenv(ENV_RECALL_BUDGET_FUNCTION, DEFAULT_RECALL_BUDGET_FUNCTION)
|
||||
),
|
||||
recall_budget_fixed_low=int(os.getenv(ENV_RECALL_BUDGET_FIXED_LOW, str(DEFAULT_RECALL_BUDGET_FIXED_LOW))),
|
||||
recall_budget_fixed_mid=int(os.getenv(ENV_RECALL_BUDGET_FIXED_MID, str(DEFAULT_RECALL_BUDGET_FIXED_MID))),
|
||||
recall_budget_fixed_high=int(
|
||||
os.getenv(ENV_RECALL_BUDGET_FIXED_HIGH, str(DEFAULT_RECALL_BUDGET_FIXED_HIGH))
|
||||
),
|
||||
recall_budget_adaptive_low=float(
|
||||
os.getenv(ENV_RECALL_BUDGET_ADAPTIVE_LOW, str(DEFAULT_RECALL_BUDGET_ADAPTIVE_LOW))
|
||||
),
|
||||
recall_budget_adaptive_mid=float(
|
||||
os.getenv(ENV_RECALL_BUDGET_ADAPTIVE_MID, str(DEFAULT_RECALL_BUDGET_ADAPTIVE_MID))
|
||||
),
|
||||
recall_budget_adaptive_high=float(
|
||||
os.getenv(ENV_RECALL_BUDGET_ADAPTIVE_HIGH, str(DEFAULT_RECALL_BUDGET_ADAPTIVE_HIGH))
|
||||
),
|
||||
recall_budget_min=int(os.getenv(ENV_RECALL_BUDGET_MIN, str(DEFAULT_RECALL_BUDGET_MIN))),
|
||||
recall_budget_max=int(os.getenv(ENV_RECALL_BUDGET_MAX, str(DEFAULT_RECALL_BUDGET_MAX))),
|
||||
# Disposition settings (None = fall back to DB value)
|
||||
disposition_skepticism=int(os.getenv(ENV_DISPOSITION_SKEPTICISM))
|
||||
if os.getenv(ENV_DISPOSITION_SKEPTICISM)
|
||||
@@ -1832,8 +1613,7 @@ class HindsightConfig:
|
||||
handler.setLevel(self.get_python_log_level())
|
||||
|
||||
if self.log_format == "json":
|
||||
allowed = frozenset(self.log_json_fields) if self.log_json_fields else None
|
||||
handler.setFormatter(JsonFormatter(allowed_fields=allowed))
|
||||
handler.setFormatter(JsonFormatter())
|
||||
else:
|
||||
handler.setFormatter(logging.Formatter("%(asctime)s - %(levelname)s - %(name)s - %(message)s"))
|
||||
|
||||
|
||||
@@ -15,12 +15,7 @@ from typing import Any
|
||||
|
||||
import asyncpg
|
||||
|
||||
from hindsight_api.config import (
|
||||
RECALL_BUDGET_FUNCTIONS,
|
||||
HindsightConfig,
|
||||
_get_raw_config,
|
||||
normalize_config_dict,
|
||||
)
|
||||
from hindsight_api.config import HindsightConfig, _get_raw_config, normalize_config_dict
|
||||
from hindsight_api.engine.memory_engine import fq_table
|
||||
from hindsight_api.extensions.tenant import TenantExtension
|
||||
from hindsight_api.models import RequestContext
|
||||
@@ -261,9 +256,6 @@ class ConfigResolver:
|
||||
"Strategy names must not be empty strings. Remove entries with empty names before saving."
|
||||
)
|
||||
|
||||
# Validate recall budget fields
|
||||
_validate_recall_budget_updates(normalized_updates)
|
||||
|
||||
# Merge with existing config (JSONB || operator)
|
||||
async with self.pool.acquire() as conn:
|
||||
await conn.execute(
|
||||
@@ -300,53 +292,6 @@ class ConfigResolver:
|
||||
logger.info(f"Reset bank config for {bank_id} to defaults")
|
||||
|
||||
|
||||
_RECALL_BUDGET_FIXED_KEYS = (
|
||||
"recall_budget_fixed_low",
|
||||
"recall_budget_fixed_mid",
|
||||
"recall_budget_fixed_high",
|
||||
)
|
||||
_RECALL_BUDGET_ADAPTIVE_KEYS = (
|
||||
"recall_budget_adaptive_low",
|
||||
"recall_budget_adaptive_mid",
|
||||
"recall_budget_adaptive_high",
|
||||
)
|
||||
|
||||
|
||||
def _validate_recall_budget_updates(updates: dict[str, Any]) -> None:
|
||||
"""Validate recall budget config updates. Raises ValueError on invalid input."""
|
||||
if "recall_budget_function" in updates:
|
||||
function = updates["recall_budget_function"]
|
||||
if not isinstance(function, str) or function.lower() not in RECALL_BUDGET_FUNCTIONS:
|
||||
raise ValueError(
|
||||
f"recall_budget_function must be one of {sorted(RECALL_BUDGET_FUNCTIONS)}, got {function!r}"
|
||||
)
|
||||
|
||||
for key in _RECALL_BUDGET_FIXED_KEYS:
|
||||
if key in updates:
|
||||
value = updates[key]
|
||||
if not isinstance(value, int) or isinstance(value, bool) or value < 1:
|
||||
raise ValueError(f"{key} must be a positive integer, got {value!r}")
|
||||
|
||||
for key in _RECALL_BUDGET_ADAPTIVE_KEYS:
|
||||
if key in updates:
|
||||
value = updates[key]
|
||||
if isinstance(value, bool) or not isinstance(value, (int, float)) or value <= 0:
|
||||
raise ValueError(f"{key} must be a positive number, got {value!r}")
|
||||
|
||||
for key in ("recall_budget_min", "recall_budget_max"):
|
||||
if key in updates:
|
||||
value = updates[key]
|
||||
if not isinstance(value, int) or isinstance(value, bool) or value < 1:
|
||||
raise ValueError(f"{key} must be a positive integer, got {value!r}")
|
||||
|
||||
if "recall_budget_min" in updates and "recall_budget_max" in updates:
|
||||
if updates["recall_budget_min"] > updates["recall_budget_max"]:
|
||||
raise ValueError(
|
||||
f"recall_budget_min ({updates['recall_budget_min']}) must be <= "
|
||||
f"recall_budget_max ({updates['recall_budget_max']})"
|
||||
)
|
||||
|
||||
|
||||
def apply_strategy(config: HindsightConfig, strategy_name: str) -> HindsightConfig:
|
||||
"""
|
||||
Apply a named retain strategy's overrides on top of a resolved config.
|
||||
|
||||
@@ -63,17 +63,7 @@ def daemonize():
|
||||
Fork the current process into a background daemon.
|
||||
|
||||
Uses double-fork technique to properly detach from terminal.
|
||||
|
||||
On Windows there is no fork model: the spawning parent is expected to
|
||||
detach us via `CREATE_NEW_PROCESS_GROUP | DETACHED_PROCESS` and to
|
||||
redirect stdout/stderr to HINDSIGHT_API_DAEMON_LOG before exec. We
|
||||
still ensure the log directory exists so that any file handlers set
|
||||
up by the calling app have a valid target.
|
||||
"""
|
||||
if sys.platform == "win32":
|
||||
DAEMON_LOG_PATH.parent.mkdir(parents=True, exist_ok=True)
|
||||
return
|
||||
|
||||
# First fork - detach from parent
|
||||
try:
|
||||
pid = os.fork()
|
||||
|
||||
@@ -1,57 +0,0 @@
|
||||
"""Database URL normalization.
|
||||
|
||||
Hindsight accepts SQLAlchemy-style URLs like ``postgresql+asyncpg://...?ssl=require``
|
||||
for its async engine, but the same string cannot be handed directly to synchronous
|
||||
SQLAlchemy (psycopg2) or to :func:`asyncpg.create_pool`, which both expect a
|
||||
libpq-compatible URL (``postgresql://...?sslmode=require``).
|
||||
|
||||
:func:`to_libpq_url` performs that translation. It is idempotent and safe to
|
||||
apply to URLs that are already libpq-compatible, to the ``pg0`` embedded-PG
|
||||
marker, or to any non-PostgreSQL string (returned unchanged).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
|
||||
|
||||
_ASYNCPG_SCHEMES = ("postgresql+asyncpg", "postgres+asyncpg")
|
||||
_POSTGRES_SCHEMES = ("postgresql", "postgres") + _ASYNCPG_SCHEMES
|
||||
|
||||
|
||||
def to_libpq_url(url: str) -> str:
|
||||
"""Normalize a PostgreSQL URL for libpq-style consumers.
|
||||
|
||||
Accepts a SQLAlchemy URL (``postgresql+asyncpg://...``) or a plain libpq
|
||||
URL and returns a form suitable for:
|
||||
|
||||
- :func:`sqlalchemy.create_engine` (sync / psycopg2)
|
||||
- :func:`asyncpg.create_pool`
|
||||
|
||||
Transformations:
|
||||
|
||||
- ``postgresql+asyncpg`` / ``postgres+asyncpg`` / ``postgres`` → ``postgresql``
|
||||
- Query param ``ssl=<mode>`` → ``sslmode=<mode>`` (SQLAlchemy's asyncpg
|
||||
dialect uses ``ssl=``; libpq uses ``sslmode=``)
|
||||
|
||||
Any non-PostgreSQL input (e.g. the ``pg0`` embedded-PG marker, a sqlite
|
||||
URL, an empty string) is returned unchanged. Already-normalized URLs are
|
||||
returned unchanged.
|
||||
"""
|
||||
if not url or "://" not in url:
|
||||
return url
|
||||
|
||||
parts = urlsplit(url)
|
||||
if parts.scheme not in _POSTGRES_SCHEMES:
|
||||
return url
|
||||
|
||||
new_scheme = "postgresql"
|
||||
|
||||
new_query_pairs = [
|
||||
("sslmode", v) if k == "ssl" else (k, v) for k, v in parse_qsl(parts.query, keep_blank_values=True)
|
||||
]
|
||||
new_query = urlencode(new_query_pairs)
|
||||
|
||||
if new_scheme == parts.scheme and new_query == parts.query:
|
||||
return url
|
||||
|
||||
return urlunsplit((new_scheme, parts.netloc, parts.path, new_query, parts.fragment))
|
||||
@@ -28,7 +28,7 @@ from pydantic import BaseModel, field_validator
|
||||
|
||||
from ...config import get_config
|
||||
from ..llm_wrapper import sanitize_llm_output
|
||||
from ..memory_engine import Budget, fq_table
|
||||
from ..memory_engine import fq_table
|
||||
from ..retain import embedding_utils
|
||||
from .prompts import build_batch_consolidation_prompt
|
||||
|
||||
@@ -42,34 +42,6 @@ if TYPE_CHECKING:
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
async def _filter_live_source_memories(
|
||||
conn: "Connection",
|
||||
bank_id: str,
|
||||
source_memory_ids: list[uuid.UUID],
|
||||
) -> list[uuid.UUID]:
|
||||
"""Return only the source memory ids that still exist in the bank.
|
||||
|
||||
Uses FOR SHARE to block concurrent deletes from removing a row between the
|
||||
check and the subsequent insert/update. Combined with the delete path running
|
||||
its stale-observation sweep *after* deleting the source row, this closes the
|
||||
race window where consolidation would otherwise produce an orphan observation.
|
||||
"""
|
||||
if not source_memory_ids:
|
||||
return []
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT id
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE id = ANY($1::uuid[]) AND bank_id = $2
|
||||
FOR SHARE
|
||||
""",
|
||||
source_memory_ids,
|
||||
bank_id,
|
||||
)
|
||||
live = {row["id"] for row in rows}
|
||||
return [mid for mid in source_memory_ids if mid in live]
|
||||
|
||||
|
||||
class _CreateAction(BaseModel):
|
||||
text: str
|
||||
source_fact_ids: list[str] # memory UUIDs from the NEW FACTS list
|
||||
@@ -247,7 +219,6 @@ async def run_consolidation_job(
|
||||
|
||||
perf = ConsolidationPerfLog(bank_id)
|
||||
max_memories_per_batch = config.consolidation_batch_size
|
||||
max_memories_per_round = config.consolidation_max_memories_per_round
|
||||
llm_batch_size = max(1, config.consolidation_llm_batch_size)
|
||||
|
||||
# Check if consolidation is enabled
|
||||
@@ -310,17 +281,8 @@ async def run_consolidation_job(
|
||||
# Track all unique tags from consolidated memories for mental model refresh filtering
|
||||
consolidated_tags: set[str] = set()
|
||||
|
||||
round_limit_enabled = max_memories_per_round > 0
|
||||
round_remaining = max_memories_per_round if round_limit_enabled else float("inf")
|
||||
hit_round_limit = False
|
||||
|
||||
llm_batch_num = 0
|
||||
while True:
|
||||
# Cap fetch size by remaining round budget
|
||||
fetch_limit = (
|
||||
min(max_memories_per_batch, int(round_remaining)) if round_limit_enabled else max_memories_per_batch
|
||||
)
|
||||
|
||||
# Fetch next batch of unconsolidated memories
|
||||
async with pool.acquire() as conn:
|
||||
t0 = time.time()
|
||||
@@ -337,7 +299,7 @@ async def run_consolidation_job(
|
||||
LIMIT $2
|
||||
""",
|
||||
bank_id,
|
||||
fetch_limit,
|
||||
max_memories_per_batch,
|
||||
)
|
||||
perf.record_timing("fetch_memories", time.time() - t0)
|
||||
|
||||
@@ -562,25 +524,6 @@ async def run_consolidation_job(
|
||||
f" | avg={llm_batch_time / len(llm_batch):.3f}s/memory"
|
||||
)
|
||||
|
||||
# Update round budget after processing this DB fetch batch
|
||||
if round_limit_enabled:
|
||||
round_remaining -= len(memories)
|
||||
if round_remaining <= 0:
|
||||
hit_round_limit = True
|
||||
break
|
||||
|
||||
# Re-submit consolidation if we hit the round limit and there's likely more work
|
||||
if hit_round_limit:
|
||||
remaining = total_count - stats["memories_processed"]
|
||||
logger.info(
|
||||
f"[CONSOLIDATION] bank={bank_id} hit round limit of {max_memories_per_round} memories,"
|
||||
f" ~{remaining} remaining. Re-queuing consolidation."
|
||||
)
|
||||
try:
|
||||
await memory_engine.submit_async_consolidation(bank_id=bank_id, request_context=request_context)
|
||||
except Exception as e:
|
||||
logger.warning(f"[CONSOLIDATION] bank={bank_id} failed to re-queue consolidation: {e}")
|
||||
|
||||
# Build summary
|
||||
perf.log(
|
||||
f"[3] Results: {stats['memories_processed']} memories -> "
|
||||
@@ -609,21 +552,16 @@ async def run_consolidation_job(
|
||||
if timing_parts:
|
||||
perf.log(f"[4] Timing breakdown: {', '.join(timing_parts)}")
|
||||
|
||||
# Trigger mental model refreshes only on the final round (when all memories are processed).
|
||||
# If we hit the round limit and re-queued, skip MM refresh — the next round will handle it.
|
||||
if hit_round_limit:
|
||||
stats["mental_models_refreshed"] = 0
|
||||
logger.info(f"[CONSOLIDATION] bank={bank_id} skipping mental model refresh (round limit hit, re-queued)")
|
||||
else:
|
||||
# SECURITY: Only refresh mental models with matching tags (or all if no tags were consolidated)
|
||||
mental_models_refreshed = await _trigger_mental_model_refreshes(
|
||||
memory_engine=memory_engine,
|
||||
bank_id=bank_id,
|
||||
request_context=request_context,
|
||||
consolidated_tags=list(consolidated_tags) if consolidated_tags else None,
|
||||
perf=perf,
|
||||
)
|
||||
stats["mental_models_refreshed"] = mental_models_refreshed
|
||||
# Trigger mental model refreshes for models with refresh_after_consolidation=true
|
||||
# SECURITY: Only refresh mental models with matching tags (or all if no tags were consolidated)
|
||||
mental_models_refreshed = await _trigger_mental_model_refreshes(
|
||||
memory_engine=memory_engine,
|
||||
bank_id=bank_id,
|
||||
request_context=request_context,
|
||||
consolidated_tags=list(consolidated_tags) if consolidated_tags else None,
|
||||
perf=perf,
|
||||
)
|
||||
stats["mental_models_refreshed"] = mental_models_refreshed
|
||||
|
||||
perf.flush()
|
||||
|
||||
@@ -655,15 +593,17 @@ async def _trigger_mental_model_refreshes(
|
||||
"""
|
||||
pool = memory_engine._pool
|
||||
|
||||
# Find mental models with refresh_after_consolidation=true that are actually stale.
|
||||
# The tag filter on the SELECT enforces the security boundary (never look outside the
|
||||
# relevant tag scope); compute_mental_model_is_stale then verifies that new memories
|
||||
# in the MM's scope really were ingested since its last refresh.
|
||||
# Find mental models with refresh_after_consolidation=true
|
||||
# SECURITY: Control which mental models get refreshed based on tags
|
||||
async with pool.acquire() as conn:
|
||||
if consolidated_tags:
|
||||
candidates = await conn.fetch(
|
||||
# Tagged memories were consolidated - refresh:
|
||||
# 1. Mental models with overlapping tags (security boundary)
|
||||
# 2. Untagged mental models (they're "global" and available to all contexts)
|
||||
# DO NOT refresh mental models with different tags
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT id, name, tags, last_refreshed_at, trigger
|
||||
SELECT id, name, tags
|
||||
FROM {fq_table("mental_models")}
|
||||
WHERE bank_id = $1
|
||||
AND (trigger->>'refresh_after_consolidation')::boolean = true
|
||||
@@ -676,9 +616,11 @@ async def _trigger_mental_model_refreshes(
|
||||
consolidated_tags,
|
||||
)
|
||||
else:
|
||||
candidates = await conn.fetch(
|
||||
# Untagged memories were consolidated - only refresh untagged mental models
|
||||
# SECURITY: Tagged mental models are NOT refreshed when untagged memories are consolidated
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT id, name, tags, last_refreshed_at, trigger
|
||||
SELECT id, name, tags
|
||||
FROM {fq_table("mental_models")}
|
||||
WHERE bank_id = $1
|
||||
AND (trigger->>'refresh_after_consolidation')::boolean = true
|
||||
@@ -687,11 +629,6 @@ async def _trigger_mental_model_refreshes(
|
||||
bank_id,
|
||||
)
|
||||
|
||||
rows = []
|
||||
for candidate in candidates:
|
||||
if await memory_engine.compute_mental_model_is_stale(conn, bank_id, candidate):
|
||||
rows.append(candidate)
|
||||
|
||||
if not rows:
|
||||
return 0
|
||||
|
||||
@@ -952,15 +889,6 @@ async def _execute_update_action(
|
||||
logger.debug(f"Update skipped: observation {observation_id} not found in recall results")
|
||||
return
|
||||
|
||||
live_source_memory_ids = await _filter_live_source_memories(conn, bank_id, source_memory_ids)
|
||||
if not live_source_memory_ids:
|
||||
logger.debug(
|
||||
f"Update skipped: all {len(source_memory_ids)} source memories for observation "
|
||||
f"{observation_id} were deleted concurrently"
|
||||
)
|
||||
return
|
||||
source_memory_ids = live_source_memory_ids
|
||||
|
||||
from ...config import get_config
|
||||
|
||||
history_entry = {
|
||||
@@ -1138,14 +1066,10 @@ async def _find_related_observations(
|
||||
else:
|
||||
recall_span = None
|
||||
|
||||
# Resolve budget: consolidation doesn't need deep recall, default to LOW to reduce memory fan-out
|
||||
recall_budget = Budget(config.consolidation_recall_budget)
|
||||
|
||||
try:
|
||||
recall_result = await memory_engine.recall_async(
|
||||
bank_id=bank_id,
|
||||
query=query,
|
||||
budget=recall_budget,
|
||||
max_tokens=config.consolidation_max_tokens, # Token budget for observations (configurable)
|
||||
fact_type=["observation"], # Only retrieve observations
|
||||
request_context=request_context,
|
||||
@@ -1207,16 +1131,14 @@ async def _consolidate_batch_with_llm(
|
||||
memories: list[dict[str, Any]],
|
||||
union_observations: "list[MemoryFact]",
|
||||
union_source_facts: "dict[str, MemoryFact]",
|
||||
config: Any,
|
||||
config: Any = None,
|
||||
remaining_observation_slots: int | None = None,
|
||||
max_observations_per_scope: int = -1,
|
||||
) -> _BatchLLMResult:
|
||||
"""Single LLM call for a batch of facts against a pooled set of observations."""
|
||||
if config is None:
|
||||
raise ValueError("config is required for _consolidate_batch_with_llm")
|
||||
if union_observations:
|
||||
obs_list = _build_observations_for_llm(union_observations, union_source_facts)
|
||||
observations_text = json.dumps(obs_list, indent=2, ensure_ascii=False)
|
||||
observations_text = json.dumps(obs_list, indent=2)
|
||||
else:
|
||||
observations_text = "[]"
|
||||
|
||||
@@ -1250,7 +1172,8 @@ async def _consolidate_batch_with_llm(
|
||||
f"(out of {max_observations_per_scope}). Prefer UPDATE over CREATE when possible."
|
||||
)
|
||||
|
||||
prompt_template = build_batch_consolidation_prompt(config.observations_mission, observation_capacity_note)
|
||||
observations_mission = config.observations_mission if config is not None else None
|
||||
prompt_template = build_batch_consolidation_prompt(observations_mission, observation_capacity_note)
|
||||
prompt = prompt_template.format(
|
||||
facts_text=facts_lines,
|
||||
observations_text=observations_text,
|
||||
@@ -1259,29 +1182,15 @@ async def _consolidate_batch_with_llm(
|
||||
# Use a constrained response model when observation limit is active
|
||||
response_model = _build_response_model(max_creates=remaining_observation_slots)
|
||||
|
||||
max_attempts = config.consolidation_max_attempts
|
||||
inner_max_retries = config.consolidation_llm_max_retries
|
||||
max_attempts = 3
|
||||
last_exc: Exception | None = None
|
||||
# Pre-compute a stable identifier set for the batch so failure logs name the
|
||||
# exact memories whose consolidation is failing — without this, an opaque
|
||||
# "LLM batch call failed" line gives operators no way to find the offending
|
||||
# input until adaptive bisection narrows the batch down to a single memory.
|
||||
memory_ids = [str(m.get("id")) for m in memories]
|
||||
if len(memory_ids) <= 5:
|
||||
ids_label = ", ".join(memory_ids)
|
||||
else:
|
||||
ids_label = f"{', '.join(memory_ids[:3])}, ... +{len(memory_ids) - 3} more"
|
||||
batch_label = f"{len(memory_ids)} memories [{ids_label}]"
|
||||
for attempt in range(1, max_attempts + 1):
|
||||
try:
|
||||
call_kwargs: dict[str, Any] = {
|
||||
"messages": [{"role": "user", "content": prompt}],
|
||||
"response_format": response_model,
|
||||
"scope": "consolidation",
|
||||
}
|
||||
if inner_max_retries is not None:
|
||||
call_kwargs["max_retries"] = inner_max_retries
|
||||
response: _ConsolidationBatchResponse = await llm_config.call(**call_kwargs)
|
||||
response: _ConsolidationBatchResponse = await llm_config.call(
|
||||
messages=[{"role": "user", "content": prompt}],
|
||||
response_format=response_model,
|
||||
scope="consolidation",
|
||||
)
|
||||
# Defensive truncation: some LLM providers may not enforce JSON schema max_length
|
||||
creates = response.creates
|
||||
if remaining_observation_slots is not None and remaining_observation_slots >= 0:
|
||||
@@ -1300,13 +1209,10 @@ async def _consolidate_batch_with_llm(
|
||||
)
|
||||
except Exception as exc:
|
||||
last_exc = exc
|
||||
logger.warning(
|
||||
f"[CONSOLIDATION] LLM batch call failed (attempt {attempt}/{max_attempts}) for {batch_label}: {exc}"
|
||||
)
|
||||
logger.warning(f"[CONSOLIDATION] LLM batch call failed (attempt {attempt}/{max_attempts}): {exc}")
|
||||
|
||||
logger.error(
|
||||
f"[CONSOLIDATION] LLM batch call failed after {max_attempts} attempts for {batch_label}, "
|
||||
f"skipping batch. Last error: {last_exc}"
|
||||
f"[CONSOLIDATION] LLM batch call failed after {max_attempts} attempts, skipping batch. Last error: {last_exc}"
|
||||
)
|
||||
return _BatchLLMResult(obs_count=len(union_observations), prompt_chars=len(prompt), failed=True)
|
||||
|
||||
@@ -1325,12 +1231,6 @@ async def _create_observation_directly(
|
||||
perf: ConsolidationPerfLog | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Create an observation from one or more source memories with pre-processed text."""
|
||||
live_source_memory_ids = await _filter_live_source_memories(conn, bank_id, source_memory_ids)
|
||||
if not live_source_memory_ids:
|
||||
logger.debug(f"Create skipped: all {len(source_memory_ids)} source memories were deleted concurrently")
|
||||
return {"action": "skipped", "reason": "sources_deleted"}
|
||||
source_memory_ids = live_source_memory_ids
|
||||
|
||||
# Generate embedding for the observation (convert to string for pgvector)
|
||||
t0 = time.time()
|
||||
embeddings = await embedding_utils.generate_embeddings_batch(memory_engine.embeddings, [observation_text])
|
||||
|
||||
@@ -19,7 +19,6 @@ from ..config import (
|
||||
DEFAULT_LITELLM_API_BASE,
|
||||
DEFAULT_RERANKER_COHERE_MODEL,
|
||||
DEFAULT_RERANKER_FLASHRANK_CACHE_DIR,
|
||||
DEFAULT_RERANKER_FLASHRANK_CPU_MEM_ARENA,
|
||||
DEFAULT_RERANKER_FLASHRANK_MODEL,
|
||||
DEFAULT_RERANKER_GOOGLE_MODEL,
|
||||
DEFAULT_RERANKER_LITELLM_MAX_TOKENS_PER_DOC,
|
||||
@@ -34,13 +33,11 @@ from ..config import (
|
||||
DEFAULT_RERANKER_SILICONFLOW_BASE_URL,
|
||||
DEFAULT_RERANKER_SILICONFLOW_MODEL,
|
||||
DEFAULT_RERANKER_TEI_BATCH_SIZE,
|
||||
DEFAULT_RERANKER_TEI_HTTP_TIMEOUT,
|
||||
DEFAULT_RERANKER_TEI_MAX_CONCURRENT,
|
||||
DEFAULT_RERANKER_ZEROENTROPY_MODEL,
|
||||
ENV_RERANKER_COHERE_API_KEY,
|
||||
ENV_RERANKER_COHERE_MODEL,
|
||||
ENV_RERANKER_FLASHRANK_CACHE_DIR,
|
||||
ENV_RERANKER_FLASHRANK_CPU_MEM_ARENA,
|
||||
ENV_RERANKER_FLASHRANK_MODEL,
|
||||
ENV_RERANKER_GOOGLE_PROJECT_ID,
|
||||
ENV_RERANKER_LITELLM_SDK_API_KEY,
|
||||
@@ -51,7 +48,6 @@ from ..config import (
|
||||
ENV_RERANKER_PROVIDER,
|
||||
ENV_RERANKER_SILICONFLOW_API_KEY,
|
||||
ENV_RERANKER_TEI_BATCH_SIZE,
|
||||
ENV_RERANKER_TEI_HTTP_TIMEOUT,
|
||||
ENV_RERANKER_TEI_MAX_CONCURRENT,
|
||||
ENV_RERANKER_TEI_URL,
|
||||
ENV_RERANKER_ZEROENTROPY_API_KEY,
|
||||
@@ -866,7 +862,6 @@ class FlashRankCrossEncoder(CrossEncoderModel):
|
||||
cache_dir: str | None = None,
|
||||
max_length: int = 512,
|
||||
max_concurrent: int = 4,
|
||||
cpu_mem_arena: bool = False,
|
||||
):
|
||||
"""
|
||||
Initialize FlashRank cross-encoder.
|
||||
@@ -876,15 +871,10 @@ class FlashRankCrossEncoder(CrossEncoderModel):
|
||||
cache_dir: Directory to cache downloaded models. Default: system cache
|
||||
max_length: Maximum sequence length for reranking. Default: 512
|
||||
max_concurrent: Maximum concurrent reranking calls. Default: 4
|
||||
cpu_mem_arena: Enable ONNX Runtime CPU memory arena. Default: False.
|
||||
When True, ONNX pre-allocates a memory arena that never
|
||||
shrinks, causing RSS to grow monotonically. False trades
|
||||
slightly slower per-call allocation for bounded RSS.
|
||||
"""
|
||||
self.model_name = model_name or DEFAULT_RERANKER_FLASHRANK_MODEL
|
||||
self.cache_dir = cache_dir or DEFAULT_RERANKER_FLASHRANK_CACHE_DIR
|
||||
self.max_length = max_length
|
||||
self.cpu_mem_arena = cpu_mem_arena
|
||||
self._ranker = None
|
||||
FlashRankCrossEncoder._max_concurrent = max_concurrent
|
||||
|
||||
@@ -902,47 +892,15 @@ class FlashRankCrossEncoder(CrossEncoderModel):
|
||||
except ImportError:
|
||||
raise ImportError("flashrank is required for FlashRankCrossEncoder. Install it with: pip install flashrank")
|
||||
|
||||
logger.info(
|
||||
f"Reranker: initializing FlashRank provider with model {self.model_name}"
|
||||
f" (cpu_mem_arena={self.cpu_mem_arena})"
|
||||
)
|
||||
|
||||
# Configure ONNX session options before Ranker creates the session.
|
||||
# When cpu_mem_arena=False (default), ONNX won't pre-allocate an arena
|
||||
# that grows monotonically, keeping RSS bounded after rerank batches.
|
||||
if not self.cpu_mem_arena:
|
||||
import onnxruntime as ort
|
||||
|
||||
session_options = ort.SessionOptions()
|
||||
session_options.enable_cpu_mem_arena = False
|
||||
else:
|
||||
session_options = None
|
||||
logger.info(f"Reranker: initializing FlashRank provider with model {self.model_name}")
|
||||
|
||||
# Initialize ranker with optional cache directory
|
||||
ranker_kwargs: dict = {"model_name": self.model_name, "max_length": self.max_length}
|
||||
ranker_kwargs = {"model_name": self.model_name, "max_length": self.max_length}
|
||||
if self.cache_dir:
|
||||
ranker_kwargs["cache_dir"] = self.cache_dir
|
||||
|
||||
self._ranker = Ranker(**ranker_kwargs)
|
||||
|
||||
# Patch the ONNX session options if arena is disabled.
|
||||
# FlashRank's Ranker doesn't expose SessionOptions in its API,
|
||||
# so we replace the session after initialization.
|
||||
if session_options is not None and hasattr(self._ranker, "session"):
|
||||
import onnxruntime as ort
|
||||
|
||||
model_file = None
|
||||
model_dir = getattr(self._ranker, "model_dir", None)
|
||||
if model_dir:
|
||||
from pathlib import Path
|
||||
|
||||
for candidate in Path(model_dir).glob("*.onnx"):
|
||||
model_file = str(candidate)
|
||||
break
|
||||
if model_file:
|
||||
self._ranker.session = ort.InferenceSession(model_file, sess_options=session_options)
|
||||
logger.info("Reranker: replaced FlashRank ONNX session with cpu_mem_arena=False")
|
||||
|
||||
# Initialize shared executor
|
||||
if FlashRankCrossEncoder._executor is None:
|
||||
FlashRankCrossEncoder._executor = ThreadPoolExecutor(
|
||||
@@ -1324,7 +1282,6 @@ class JinaMLXCrossEncoder(CrossEncoderModel):
|
||||
def _load_model(self) -> None:
|
||||
"""Download (if needed) and load the MLX reranker. Runs in a thread."""
|
||||
import os
|
||||
import threading
|
||||
|
||||
from huggingface_hub import snapshot_download
|
||||
|
||||
@@ -1340,10 +1297,6 @@ class JinaMLXCrossEncoder(CrossEncoderModel):
|
||||
model_path=model_path,
|
||||
projector_path=os.path.join(model_path, "projector.safetensors"),
|
||||
)
|
||||
# MLX Metal GPU ops are not thread-safe — concurrent calls to
|
||||
# Device::end_encoding() crash with SIGSEGV (NULL deref).
|
||||
# Serialize all reranker inference through this lock.
|
||||
self._mlx_lock = threading.Lock()
|
||||
logger.info("Reranker: jina-mlx provider initialized")
|
||||
|
||||
def _predict_sync(self, pairs: list[tuple[str, str]]) -> list[float]:
|
||||
@@ -1357,14 +1310,13 @@ class JinaMLXCrossEncoder(CrossEncoderModel):
|
||||
|
||||
all_scores = [0.0] * len(pairs)
|
||||
|
||||
with self._mlx_lock:
|
||||
for query, indexed_docs in query_groups.items():
|
||||
docs = [doc for _, doc in indexed_docs]
|
||||
indices = [idx for idx, _ in indexed_docs]
|
||||
results = self._reranker.rerank(query, docs)
|
||||
for result in results:
|
||||
original_idx = result["index"]
|
||||
all_scores[indices[original_idx]] = result["relevance_score"]
|
||||
for query, indexed_docs in query_groups.items():
|
||||
docs = [doc for _, doc in indexed_docs]
|
||||
indices = [idx for idx, _ in indexed_docs]
|
||||
results = self._reranker.rerank(query, docs)
|
||||
for result in results:
|
||||
original_idx = result["index"]
|
||||
all_scores[indices[original_idx]] = result["relevance_score"]
|
||||
|
||||
return all_scores
|
||||
|
||||
@@ -1554,7 +1506,6 @@ def create_cross_encoder_from_env() -> CrossEncoderModel:
|
||||
raise ValueError(f"{ENV_RERANKER_TEI_URL} is required when {ENV_RERANKER_PROVIDER} is 'tei'")
|
||||
return RemoteTEICrossEncoder(
|
||||
base_url=url,
|
||||
timeout=config.reranker_tei_http_timeout,
|
||||
batch_size=config.reranker_tei_batch_size,
|
||||
max_concurrent=config.reranker_tei_max_concurrent,
|
||||
)
|
||||
@@ -1592,10 +1543,7 @@ def create_cross_encoder_from_env() -> CrossEncoderModel:
|
||||
elif provider == "flashrank":
|
||||
model = os.environ.get(ENV_RERANKER_FLASHRANK_MODEL, DEFAULT_RERANKER_FLASHRANK_MODEL)
|
||||
cache_dir = os.environ.get(ENV_RERANKER_FLASHRANK_CACHE_DIR, DEFAULT_RERANKER_FLASHRANK_CACHE_DIR)
|
||||
cpu_mem_arena = os.environ.get(
|
||||
ENV_RERANKER_FLASHRANK_CPU_MEM_ARENA, str(DEFAULT_RERANKER_FLASHRANK_CPU_MEM_ARENA)
|
||||
).lower() in ("true", "1", "yes")
|
||||
return FlashRankCrossEncoder(model_name=model, cache_dir=cache_dir, cpu_mem_arena=cpu_mem_arena)
|
||||
return FlashRankCrossEncoder(model_name=model, cache_dir=cache_dir)
|
||||
elif provider == "litellm":
|
||||
return LiteLLMCrossEncoder(
|
||||
api_base=config.reranker_litellm_api_base,
|
||||
|
||||
@@ -516,7 +516,6 @@ class CohereEmbeddings(Embeddings):
|
||||
api_key: str,
|
||||
model: str = DEFAULT_EMBEDDINGS_COHERE_MODEL,
|
||||
base_url: str | None = None,
|
||||
output_dimensions: int | None = None,
|
||||
batch_size: int = 96,
|
||||
timeout: float = 60.0,
|
||||
input_type: str = "search_document",
|
||||
@@ -528,7 +527,6 @@ class CohereEmbeddings(Embeddings):
|
||||
api_key: Cohere API key
|
||||
model: Cohere embedding model name (default: embed-english-v3.0)
|
||||
base_url: Custom base URL for Cohere-compatible API (e.g., Azure-hosted endpoint)
|
||||
output_dimensions: Optional output embedding dimensions (for Matryoshka-capable models)
|
||||
batch_size: Maximum batch size for embedding requests (default: 96, Cohere's limit)
|
||||
timeout: Request timeout in seconds (default: 60.0)
|
||||
input_type: Input type for embeddings (default: search_document).
|
||||
@@ -537,7 +535,6 @@ class CohereEmbeddings(Embeddings):
|
||||
self.api_key = api_key
|
||||
self.model = model
|
||||
self.base_url = base_url
|
||||
self.output_dimensions = output_dimensions
|
||||
self.batch_size = batch_size
|
||||
self.timeout = timeout
|
||||
self.input_type = input_type
|
||||
@@ -573,10 +570,8 @@ class CohereEmbeddings(Embeddings):
|
||||
client_kwargs["base_url"] = self.base_url
|
||||
self._client = cohere.Client(**client_kwargs)
|
||||
|
||||
# If output_dimensions is explicitly set, use that as the dimension
|
||||
if self.output_dimensions is not None:
|
||||
self._dimension = self.output_dimensions
|
||||
elif self.model in self.MODEL_DIMENSIONS:
|
||||
# Try to get dimension from known models, otherwise do a test embedding
|
||||
if self.model in self.MODEL_DIMENSIONS:
|
||||
self._dimension = self.MODEL_DIMENSIONS[self.model]
|
||||
else:
|
||||
# Do a test embedding to detect dimension
|
||||
@@ -612,23 +607,13 @@ class CohereEmbeddings(Embeddings):
|
||||
for i in range(0, len(texts), self.batch_size):
|
||||
batch = texts[i : i + self.batch_size]
|
||||
|
||||
if self.output_dimensions is not None:
|
||||
# Use v2 API which supports output_dimension
|
||||
response = self._client.v2.embed(
|
||||
texts=batch,
|
||||
model=self.model,
|
||||
input_type=self.input_type,
|
||||
output_dimension=self.output_dimensions,
|
||||
embedding_types=["float"],
|
||||
)
|
||||
all_embeddings.extend(response.embeddings.float_)
|
||||
else:
|
||||
response = self._client.embed(
|
||||
texts=batch,
|
||||
model=self.model,
|
||||
input_type=self.input_type,
|
||||
)
|
||||
all_embeddings.extend(response.embeddings)
|
||||
response = self._client.embed(
|
||||
texts=batch,
|
||||
model=self.model,
|
||||
input_type=self.input_type,
|
||||
)
|
||||
|
||||
all_embeddings.extend(response.embeddings)
|
||||
|
||||
return all_embeddings
|
||||
|
||||
@@ -927,7 +912,6 @@ class GeminiEmbeddings(Embeddings):
|
||||
vertexai_service_account_key: str | None = None,
|
||||
output_dimensionality: int | None = None,
|
||||
batch_size: int = 100,
|
||||
force_ipv4: bool = False,
|
||||
):
|
||||
self.model = model
|
||||
self.api_key = api_key
|
||||
@@ -936,9 +920,7 @@ class GeminiEmbeddings(Embeddings):
|
||||
self.vertexai_service_account_key = vertexai_service_account_key
|
||||
self.output_dimensionality = output_dimensionality
|
||||
self.batch_size = batch_size
|
||||
self.force_ipv4 = force_ipv4
|
||||
self._client = None
|
||||
self._httpx_client = None
|
||||
self._dimension: int | None = None
|
||||
self._is_vertexai = vertexai_project_id is not None
|
||||
self._embed_config = None # EmbedContentConfig, built during initialize()
|
||||
@@ -964,7 +946,7 @@ class GeminiEmbeddings(Embeddings):
|
||||
if self._is_vertexai:
|
||||
self._init_vertexai(genai)
|
||||
else:
|
||||
self._init_gemini(genai, genai_types)
|
||||
self._init_gemini(genai)
|
||||
|
||||
# Build EmbedContentConfig if output_dimensionality is set
|
||||
if self.output_dimensionality is not None:
|
||||
@@ -986,25 +968,12 @@ class GeminiEmbeddings(Embeddings):
|
||||
f"Embeddings: google provider initialized (auth: {auth_mode}, model: {self.model}, dim: {self._dimension})"
|
||||
)
|
||||
|
||||
def _init_gemini(self, genai, genai_types) -> None:
|
||||
def _init_gemini(self, genai) -> None:
|
||||
"""Initialize Gemini API client with API key."""
|
||||
if not self.api_key:
|
||||
raise ValueError("Gemini embeddings provider requires an API key")
|
||||
|
||||
client_kwargs = {"api_key": self.api_key}
|
||||
if self.force_ipv4:
|
||||
import httpx
|
||||
|
||||
self._httpx_client = httpx.Client(
|
||||
timeout=10,
|
||||
transport=httpx.HTTPTransport(local_address="0.0.0.0"),
|
||||
)
|
||||
client_kwargs["http_options"] = genai_types.HttpOptions(
|
||||
timeout=10000,
|
||||
httpxClient=self._httpx_client,
|
||||
)
|
||||
|
||||
self._client = genai.Client(**client_kwargs)
|
||||
self._client = genai.Client(api_key=self.api_key)
|
||||
logger.info(f"Embeddings: initializing Gemini provider with model {self.model}")
|
||||
|
||||
def _init_vertexai(self, genai) -> None:
|
||||
@@ -1131,12 +1100,7 @@ def create_embeddings_from_env() -> Embeddings:
|
||||
)
|
||||
model = os.environ.get(ENV_EMBEDDINGS_OPENAI_MODEL, DEFAULT_EMBEDDINGS_OPENAI_MODEL)
|
||||
base_url = os.environ.get(ENV_EMBEDDINGS_OPENAI_BASE_URL) or None
|
||||
return OpenAIEmbeddings(
|
||||
api_key=api_key,
|
||||
model=model,
|
||||
base_url=base_url,
|
||||
batch_size=config.embeddings_openai_batch_size,
|
||||
)
|
||||
return OpenAIEmbeddings(api_key=api_key, model=model, base_url=base_url)
|
||||
elif provider == "openrouter":
|
||||
api_key = config.embeddings_openrouter_api_key
|
||||
if not api_key:
|
||||
@@ -1148,7 +1112,6 @@ def create_embeddings_from_env() -> Embeddings:
|
||||
api_key=api_key,
|
||||
model=config.embeddings_openrouter_model,
|
||||
base_url="https://openrouter.ai/api/v1",
|
||||
batch_size=config.embeddings_openai_batch_size,
|
||||
)
|
||||
elif provider == "cohere":
|
||||
api_key = config.embeddings_cohere_api_key
|
||||
@@ -1158,7 +1121,6 @@ def create_embeddings_from_env() -> Embeddings:
|
||||
api_key=api_key,
|
||||
model=config.embeddings_cohere_model,
|
||||
base_url=config.embeddings_cohere_base_url,
|
||||
output_dimensions=config.embeddings_cohere_output_dimensions,
|
||||
)
|
||||
elif provider == "litellm":
|
||||
return LiteLLMEmbeddings(
|
||||
@@ -1197,7 +1159,6 @@ def create_embeddings_from_env() -> Embeddings:
|
||||
vertexai_region=config.embeddings_vertexai_region,
|
||||
vertexai_service_account_key=config.embeddings_vertexai_service_account_key,
|
||||
output_dimensionality=config.embeddings_gemini_output_dimensionality,
|
||||
force_ipv4=config.embeddings_gemini_force_ipv4,
|
||||
)
|
||||
else:
|
||||
raise ValueError(
|
||||
|
||||
@@ -289,6 +289,25 @@ class MemoryEngineInterface(ABC):
|
||||
"""
|
||||
...
|
||||
|
||||
@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,
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -153,7 +153,7 @@ class AnthropicLLM(LLMInterface):
|
||||
# Add JSON schema instruction if response_format is provided
|
||||
if response_format is not None and hasattr(response_format, "model_json_schema"):
|
||||
schema = response_format.model_json_schema()
|
||||
schema_msg = f"\n\nYou must respond with valid JSON matching this schema:\n{json.dumps(schema, indent=2, ensure_ascii=False)}"
|
||||
schema_msg = f"\n\nYou must respond with valid JSON matching this schema:\n{json.dumps(schema, indent=2)}"
|
||||
if system_prompt:
|
||||
system_prompt += schema_msg
|
||||
else:
|
||||
|
||||
@@ -171,7 +171,7 @@ class ClaudeCodeLLM(LLMInterface):
|
||||
if response_format is not None and hasattr(response_format, "model_json_schema"):
|
||||
schema = response_format.model_json_schema()
|
||||
schema_instruction = (
|
||||
f"\n\nYou must respond with valid JSON matching this schema:\n{json.dumps(schema, indent=2, ensure_ascii=False)}\n\n"
|
||||
f"\n\nYou must respond with valid JSON matching this schema:\n{json.dumps(schema, indent=2)}\n\n"
|
||||
"Respond with ONLY the JSON, no markdown formatting."
|
||||
)
|
||||
user_content += schema_instruction
|
||||
|
||||
@@ -205,7 +205,7 @@ class CodexLLM(LLMInterface):
|
||||
# Add JSON schema instruction if response_format is provided
|
||||
if response_format is not None and hasattr(response_format, "model_json_schema"):
|
||||
schema = response_format.model_json_schema()
|
||||
schema_msg = f"\n\nYou must respond with valid JSON matching this schema:\n{json.dumps(schema, indent=2, ensure_ascii=False)}"
|
||||
schema_msg = f"\n\nYou must respond with valid JSON matching this schema:\n{json.dumps(schema, indent=2)}"
|
||||
system_instruction += schema_msg
|
||||
|
||||
# gpt-5.2-codex only supports "detailed" reasoning summary
|
||||
|
||||
@@ -175,7 +175,7 @@ class GeminiLLM(LLMInterface):
|
||||
Args:
|
||||
messages: List of message dicts with 'role' and 'content'.
|
||||
response_format: Optional Pydantic model for structured output.
|
||||
max_completion_tokens: Maximum tokens in response (mapped to Gemini's max_output_tokens).
|
||||
max_completion_tokens: Maximum tokens in response (not supported by Gemini).
|
||||
temperature: Sampling temperature (0.0-2.0).
|
||||
scope: Scope identifier for tracking.
|
||||
max_retries: Maximum retry attempts.
|
||||
@@ -212,7 +212,7 @@ class GeminiLLM(LLMInterface):
|
||||
# Add JSON schema instruction if response_format is provided
|
||||
if response_format is not None and hasattr(response_format, "model_json_schema"):
|
||||
schema = response_format.model_json_schema()
|
||||
schema_msg = f"\n\nYou must respond with valid JSON matching this schema:\n{json.dumps(schema, indent=2, ensure_ascii=False)}"
|
||||
schema_msg = f"\n\nYou must respond with valid JSON matching this schema:\n{json.dumps(schema, indent=2)}"
|
||||
if system_instruction:
|
||||
system_instruction += schema_msg
|
||||
else:
|
||||
@@ -227,11 +227,6 @@ class GeminiLLM(LLMInterface):
|
||||
config_kwargs["response_schema"] = response_format
|
||||
if temperature is not None:
|
||||
config_kwargs["temperature"] = temperature
|
||||
# Gemini's equivalent of OpenAI-style max_completion_tokens is max_output_tokens.
|
||||
# Without it the model can produce arbitrarily long responses, ignoring the
|
||||
# caller's intended cap (e.g. mental_models max_tokens during refresh).
|
||||
if max_completion_tokens is not None:
|
||||
config_kwargs["max_output_tokens"] = max_completion_tokens
|
||||
|
||||
# Apply safety settings: context var (per-request bank override) takes precedence over instance default
|
||||
effective_safety_settings = _safety_settings_ctx.get()
|
||||
@@ -406,7 +401,7 @@ class GeminiLLM(LLMInterface):
|
||||
Args:
|
||||
messages: List of message dicts. Can include tool results with role='tool'.
|
||||
tools: List of tool definitions in OpenAI format.
|
||||
max_completion_tokens: Maximum tokens (mapped to Gemini's max_output_tokens).
|
||||
max_completion_tokens: Maximum tokens (not supported by Gemini).
|
||||
temperature: Sampling temperature.
|
||||
scope: Scope identifier for tracking.
|
||||
max_retries: Maximum retry attempts.
|
||||
@@ -498,10 +493,6 @@ class GeminiLLM(LLMInterface):
|
||||
config_kwargs["system_instruction"] = system_instruction
|
||||
if temperature is not None:
|
||||
config_kwargs["temperature"] = temperature
|
||||
# See note in `call`: Gemini's max_output_tokens is the equivalent of
|
||||
# OpenAI-style max_completion_tokens.
|
||||
if max_completion_tokens is not None:
|
||||
config_kwargs["max_output_tokens"] = max_completion_tokens
|
||||
|
||||
# Map OpenAI-style tool_choice to Gemini FunctionCallingConfig
|
||||
if tool_choice == "required":
|
||||
|
||||
@@ -60,32 +60,6 @@ def _strip_code_fences(content: str) -> str:
|
||||
return content
|
||||
|
||||
|
||||
def _summarize_status_error(e: APIStatusError, body_max: int = 400) -> str:
|
||||
"""Render an APIStatusError with status code + truncated response body.
|
||||
|
||||
Without this, retry loops only log "API error after N attempts" with the
|
||||
bare exception message — losing the provider's actual error payload, which
|
||||
is the only thing that explains *why* a request failed (rate limit reason,
|
||||
invalid tool schema, model overloaded, etc.).
|
||||
"""
|
||||
body: Any = getattr(e, "body", None)
|
||||
if body is None:
|
||||
try:
|
||||
body = e.response.text
|
||||
except Exception:
|
||||
body = None
|
||||
if isinstance(body, (dict, list)):
|
||||
try:
|
||||
body_str = json.dumps(body, default=str, ensure_ascii=False)
|
||||
except Exception:
|
||||
body_str = str(body)
|
||||
else:
|
||||
body_str = str(body or "").strip()
|
||||
if len(body_str) > body_max:
|
||||
body_str = body_str[:body_max] + "...TRUNCATED"
|
||||
return f"HTTP {e.status_code}: {body_str or '<no body>'}"
|
||||
|
||||
|
||||
class OpenAICompatibleLLM(LLMInterface):
|
||||
"""
|
||||
LLM provider for OpenAI-compatible APIs.
|
||||
@@ -206,12 +180,7 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
def _supports_reasoning_model(self) -> bool:
|
||||
"""Check if the current model is a reasoning model (o1, o3, GPT-5, DeepSeek)."""
|
||||
model_lower = self.model.lower()
|
||||
if "deepseek" in model_lower:
|
||||
# DeepSeek v4-flash is the non-thinking route. Treating every
|
||||
# DeepSeek model as a reasoning model injects reasoning_effort,
|
||||
# which conflicts with thinking-disabled flash calls.
|
||||
return any(x in model_lower for x in ["v4-pro", "reasoner", "r1", "thinking"])
|
||||
return any(x in model_lower for x in ["gpt-5", "o1", "o3"])
|
||||
return any(x in model_lower for x in ["gpt-5", "o1", "o3", "deepseek"])
|
||||
|
||||
def _get_max_reasoning_tokens(self) -> int | None:
|
||||
"""Get max reasoning tokens for reasoning models."""
|
||||
@@ -370,7 +339,9 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
else:
|
||||
# Soft enforcement: add schema to prompt and use json_object mode
|
||||
if schema is not None:
|
||||
schema_msg = f"\n\nYou must respond with valid JSON matching this schema:\n{json.dumps(schema, indent=2, ensure_ascii=False)}"
|
||||
schema_msg = (
|
||||
f"\n\nYou must respond with valid JSON matching this schema:\n{json.dumps(schema, indent=2)}"
|
||||
)
|
||||
|
||||
if call_params["messages"] and call_params["messages"][0].get("role") == "system":
|
||||
first_msg = call_params["messages"][0]
|
||||
@@ -579,19 +550,12 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
|
||||
last_exception = e
|
||||
if attempt < max_retries:
|
||||
logger.warning(
|
||||
f"APIStatusError ({self.provider}/{self.model}, scope={scope}, "
|
||||
f"attempt {attempt + 1}/{max_retries + 1}): {_summarize_status_error(e)}"
|
||||
)
|
||||
backoff = min(initial_backoff * (2**attempt), max_backoff)
|
||||
jitter = backoff * 0.2 * (2 * (time.time() % 1) - 1)
|
||||
sleep_time = backoff + jitter
|
||||
await asyncio.sleep(sleep_time)
|
||||
else:
|
||||
logger.error(
|
||||
f"API error after {max_retries + 1} attempts ({self.provider}/{self.model}, "
|
||||
f"scope={scope}): {_summarize_status_error(e)}"
|
||||
)
|
||||
logger.error(f"API error after {max_retries + 1} attempts: {str(e)}")
|
||||
raise
|
||||
|
||||
except Exception:
|
||||
@@ -632,50 +596,26 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
"""
|
||||
start_time = time.time()
|
||||
|
||||
request_tool_choice: str | dict[str, Any] | None = tool_choice
|
||||
|
||||
# Normalize named tool_choice dicts to "required" + filter tools.
|
||||
# Some providers (e.g. LM Studio, Ollama) reject the OpenAI named format
|
||||
# {"type": "function", "function": {"name": "..."}}. The semantics are
|
||||
# identical to tool_choice="required" with the tools list restricted to
|
||||
# just the requested tool, so we apply that transformation where supported.
|
||||
if isinstance(request_tool_choice, dict) and request_tool_choice.get("type") == "function":
|
||||
forced_name = request_tool_choice.get("function", {}).get("name")
|
||||
# just the requested tool, so we apply that transformation universally.
|
||||
if isinstance(tool_choice, dict) and tool_choice.get("type") == "function":
|
||||
forced_name = tool_choice.get("function", {}).get("name")
|
||||
if forced_name:
|
||||
filtered = [t for t in tools if t.get("function", {}).get("name") == forced_name]
|
||||
if filtered:
|
||||
tools = filtered
|
||||
request_tool_choice = "required"
|
||||
|
||||
# DeepSeek accepts tool calls but rejects explicit required/named
|
||||
# tool_choice values. The tools list has already been narrowed for
|
||||
# forced calls, so omitting tool_choice preserves the practical behavior.
|
||||
if "deepseek" in self.model.lower() and request_tool_choice != "auto":
|
||||
request_tool_choice = None
|
||||
|
||||
# DeepSeek tool-call replies can carry provider-specific reasoning_content.
|
||||
# The normalized tool result does not retain it, but replaying assistant
|
||||
# tool_calls without the field can trigger a 400. DeepSeek accepts an
|
||||
# empty-string fallback, matching the provider's history-replay contract.
|
||||
if "deepseek" in self.model.lower():
|
||||
normalized_messages: list[dict[str, Any]] = []
|
||||
for msg in messages:
|
||||
if msg.get("role") == "assistant" and msg.get("tool_calls") and "reasoning_content" not in msg:
|
||||
normalized_msg = dict(msg)
|
||||
normalized_msg["reasoning_content"] = ""
|
||||
normalized_messages.append(normalized_msg)
|
||||
else:
|
||||
normalized_messages.append(msg)
|
||||
messages = normalized_messages
|
||||
tool_choice = "required"
|
||||
|
||||
# Build call parameters
|
||||
call_params: dict[str, Any] = {
|
||||
"model": self.model,
|
||||
"messages": messages,
|
||||
"tools": tools,
|
||||
"tool_choice": tool_choice,
|
||||
}
|
||||
if request_tool_choice is not None:
|
||||
call_params["tool_choice"] = request_tool_choice
|
||||
|
||||
if max_completion_tokens is not None:
|
||||
call_params[self._max_tokens_param_name()] = max_completion_tokens
|
||||
@@ -766,41 +706,18 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
|
||||
except APIConnectionError as e:
|
||||
last_exception = e
|
||||
status_code = getattr(e, "status_code", None) or getattr(
|
||||
getattr(e, "response", None), "status_code", None
|
||||
)
|
||||
if attempt < max_retries:
|
||||
logger.warning(
|
||||
f"APIConnectionError in tool call ({self.provider}/{self.model}, scope={scope}, "
|
||||
f"attempt {attempt + 1}/{max_retries + 1}, HTTP {status_code}): {str(e)[:200]}"
|
||||
)
|
||||
await asyncio.sleep(min(initial_backoff * (2**attempt), max_backoff))
|
||||
continue
|
||||
logger.error(
|
||||
f"Connection error in tool call after {max_retries + 1} attempts "
|
||||
f"({self.provider}/{self.model}, scope={scope}): {str(e)}"
|
||||
)
|
||||
raise
|
||||
|
||||
except APIStatusError as e:
|
||||
if e.status_code in (401, 403):
|
||||
logger.error(
|
||||
f"Auth error in tool call (HTTP {e.status_code}, {self.provider}/{self.model}), "
|
||||
f"not retrying: {_summarize_status_error(e)}"
|
||||
)
|
||||
raise
|
||||
last_exception = e
|
||||
if attempt < max_retries:
|
||||
logger.warning(
|
||||
f"APIStatusError in tool call ({self.provider}/{self.model}, scope={scope}, "
|
||||
f"attempt {attempt + 1}/{max_retries + 1}): {_summarize_status_error(e)}"
|
||||
)
|
||||
await asyncio.sleep(min(initial_backoff * (2**attempt), max_backoff))
|
||||
continue
|
||||
logger.error(
|
||||
f"API error in tool call after {max_retries + 1} attempts "
|
||||
f"({self.provider}/{self.model}, scope={scope}): {_summarize_status_error(e)}"
|
||||
)
|
||||
raise
|
||||
|
||||
except Exception:
|
||||
@@ -848,7 +765,6 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
"model": self.model,
|
||||
"messages": messages,
|
||||
"stream": False,
|
||||
"think": False, # Disable thinking for reasoning models (qwen3.5, etc.)
|
||||
}
|
||||
|
||||
# Add schema as format parameter for structured output
|
||||
@@ -1003,7 +919,7 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
logger.info(f"Submitting batch with {len(requests)} requests to {self.provider}")
|
||||
|
||||
# Format requests as JSONL
|
||||
jsonl_content = "\n".join(json.dumps(req, ensure_ascii=False) for req in requests)
|
||||
jsonl_content = "\n".join(json.dumps(req) for req in requests)
|
||||
|
||||
# Upload file to provider (wrap in BytesIO with filename)
|
||||
file_bytes = io.BytesIO(jsonl_content.encode("utf-8"))
|
||||
|
||||
@@ -17,12 +17,7 @@ from typing import TYPE_CHECKING, Any, Awaitable, Callable
|
||||
import tiktoken
|
||||
|
||||
from .models import DirectiveInfo, LLMCall, ReflectAgentResult, TokenUsageSummary, ToolCall
|
||||
from .prompts import (
|
||||
_extract_directive_rules,
|
||||
build_final_prompt,
|
||||
build_final_system_prompt,
|
||||
build_system_prompt_for_tools,
|
||||
)
|
||||
from .prompts import FINAL_SYSTEM_PROMPT, _extract_directive_rules, build_final_prompt, build_system_prompt_for_tools
|
||||
from .tools_schema import get_reflect_tools
|
||||
|
||||
|
||||
@@ -191,7 +186,7 @@ async def _generate_structured_output(
|
||||
DynamicModel = create_model("StructuredResponse", **fields)
|
||||
|
||||
# Include the full schema in the prompt for better LLM guidance
|
||||
schema_str = json.dumps(response_schema, indent=2, ensure_ascii=False)
|
||||
schema_str = json.dumps(response_schema, indent=2)
|
||||
|
||||
# Build field descriptions for the prompt
|
||||
field_descriptions = []
|
||||
@@ -451,7 +446,7 @@ async def run_reflect_agent(
|
||||
llm_start = time.time()
|
||||
response, usage = await llm_config.call(
|
||||
messages=[
|
||||
{"role": "system", "content": build_final_system_prompt(bank_profile.get("mission"))},
|
||||
{"role": "system", "content": FINAL_SYSTEM_PROMPT},
|
||||
{"role": "user", "content": prompt},
|
||||
],
|
||||
scope="reflect",
|
||||
@@ -508,7 +503,7 @@ async def run_reflect_agent(
|
||||
llm_start = time.time()
|
||||
response, usage = await llm_config.call(
|
||||
messages=[
|
||||
{"role": "system", "content": build_final_system_prompt(bank_profile.get("mission"))},
|
||||
{"role": "system", "content": FINAL_SYSTEM_PROMPT},
|
||||
{"role": "user", "content": prompt},
|
||||
],
|
||||
scope="reflect",
|
||||
@@ -611,7 +606,7 @@ async def run_reflect_agent(
|
||||
llm_start = time.time()
|
||||
response, usage = await llm_config.call(
|
||||
messages=[
|
||||
{"role": "system", "content": build_final_system_prompt(bank_profile.get("mission"))},
|
||||
{"role": "system", "content": FINAL_SYSTEM_PROMPT},
|
||||
{"role": "user", "content": prompt},
|
||||
],
|
||||
scope="reflect",
|
||||
@@ -654,57 +649,9 @@ async def run_reflect_agent(
|
||||
|
||||
# No tool calls - LLM wants to respond with text
|
||||
if not result.tool_calls:
|
||||
# When directives are present but no evidence has been gathered,
|
||||
# the LLM tends to echo directive content verbatim as its answer.
|
||||
# Fall through to the final-prompt path which doesn't include
|
||||
# directives and handles "no data" gracefully.
|
||||
has_gathered_evidence = (
|
||||
bool(available_memory_ids) or bool(available_mental_model_ids) or bool(available_observation_ids)
|
||||
)
|
||||
directive_leak_risk = directives and not has_gathered_evidence
|
||||
if result.content and not directive_leak_risk:
|
||||
if result.content:
|
||||
answer = _clean_answer_text(result.content.strip())
|
||||
|
||||
# The call_with_tools call above is intentionally uncapped so the
|
||||
# LLM has headroom to emit tool-call JSON plus any intermediate
|
||||
# reasoning. But when the LLM short-circuits and returns text
|
||||
# directly, that text becomes the user-visible final answer and
|
||||
# must respect max_tokens like the forced-final paths do. If it
|
||||
# overshoots, run one extra capped call to rewrite it within
|
||||
# the cap.
|
||||
if max_tokens is not None and len(_TIKTOKEN_ENCODING.encode(answer)) > max_tokens:
|
||||
rewrite_start = time.time()
|
||||
rewritten, rewrite_usage = await llm_config.call(
|
||||
messages=[
|
||||
{
|
||||
"role": "system",
|
||||
"content": (
|
||||
"Rewrite the user's text so it fits within the requested token "
|
||||
"budget. Preserve the key facts and structure; drop lower-priority "
|
||||
"detail. Respond with the rewritten text only, no preamble."
|
||||
),
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": f"Target budget: {max_tokens} tokens.\n\nText to rewrite:\n{answer}",
|
||||
},
|
||||
],
|
||||
scope="reflect",
|
||||
max_completion_tokens=max_tokens,
|
||||
return_usage=True,
|
||||
)
|
||||
total_input_tokens += rewrite_usage.input_tokens
|
||||
total_output_tokens += rewrite_usage.output_tokens
|
||||
llm_trace.append(
|
||||
{
|
||||
"scope": "final_rewrite",
|
||||
"duration_ms": int((time.time() - rewrite_start) * 1000),
|
||||
"input_tokens": rewrite_usage.input_tokens,
|
||||
"output_tokens": rewrite_usage.output_tokens,
|
||||
}
|
||||
)
|
||||
answer = _clean_answer_text(rewritten.strip())
|
||||
|
||||
# Generate structured output if schema provided
|
||||
structured_output = None
|
||||
if response_schema and answer:
|
||||
@@ -732,7 +679,7 @@ async def run_reflect_agent(
|
||||
llm_start = time.time()
|
||||
response, usage = await llm_config.call(
|
||||
messages=[
|
||||
{"role": "system", "content": build_final_system_prompt(bank_profile.get("mission"))},
|
||||
{"role": "system", "content": FINAL_SYSTEM_PROMPT},
|
||||
{"role": "user", "content": prompt},
|
||||
],
|
||||
scope="reflect",
|
||||
@@ -796,8 +743,7 @@ async def run_reflect_agent(
|
||||
"content": json.dumps(
|
||||
{
|
||||
"error": "You must search for information first. Use search_mental_models(), search_observations(), or recall() before providing your final answer."
|
||||
},
|
||||
ensure_ascii=False,
|
||||
}
|
||||
),
|
||||
}
|
||||
)
|
||||
@@ -859,8 +805,7 @@ async def run_reflect_agent(
|
||||
"content": json.dumps(
|
||||
{
|
||||
"error": f"Tool '{_normalize_tool_name(tc.name)}' is not available. Use only the tools provided to you."
|
||||
},
|
||||
ensure_ascii=False,
|
||||
}
|
||||
),
|
||||
}
|
||||
)
|
||||
@@ -931,7 +876,7 @@ async def run_reflect_agent(
|
||||
"role": "tool",
|
||||
"tool_call_id": tc.id,
|
||||
"name": tc.name, # Required by Gemini
|
||||
"content": json.dumps(output, default=str, ensure_ascii=False),
|
||||
"content": json.dumps(output, default=str),
|
||||
}
|
||||
)
|
||||
|
||||
@@ -954,7 +899,7 @@ async def run_reflect_agent(
|
||||
)
|
||||
|
||||
try:
|
||||
output_chars = len(json.dumps(output, ensure_ascii=False))
|
||||
output_chars = len(json.dumps(output))
|
||||
except (TypeError, ValueError):
|
||||
output_chars = len(str(output))
|
||||
|
||||
@@ -991,7 +936,7 @@ def _tool_call_to_dict(tc: "LLMToolCall") -> dict[str, Any]:
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": tc.name,
|
||||
"arguments": json.dumps(tc.arguments, ensure_ascii=False),
|
||||
"arguments": json.dumps(tc.arguments),
|
||||
},
|
||||
}
|
||||
if tc.thought_signature is not None:
|
||||
@@ -1089,7 +1034,7 @@ async def _execute_tool_with_timing(
|
||||
# Set attributes
|
||||
span.set_attribute("hindsight.tool.name", normalized_name)
|
||||
span.set_attribute("hindsight.tool.id", tc.id)
|
||||
span.set_attribute("hindsight.tool.arguments", json.dumps(tc.arguments, ensure_ascii=False))
|
||||
span.set_attribute("hindsight.tool.arguments", json.dumps(tc.arguments))
|
||||
|
||||
try:
|
||||
result = await _execute_tool(
|
||||
|
||||
@@ -1,307 +0,0 @@
|
||||
"""Delta operations for structured mental models.
|
||||
|
||||
The LLM's job during a delta refresh is to emit a list of these operations,
|
||||
each targeting an existing section (by id) or referencing a position relative
|
||||
to one. ``apply_operations`` validates and applies each op in turn against a
|
||||
copy of the document; invalid ops (unknown ``section_id``, out-of-range
|
||||
``block_index``, malformed payloads) are dropped with a debug-friendly reason.
|
||||
|
||||
Sections and blocks not mentioned by any op are physically copied through
|
||||
unchanged — there is no LLM-mediated re-emission of unchanged text, so prose
|
||||
drift is structurally impossible.
|
||||
|
||||
Why operations and not "output the new structured doc":
|
||||
- "Output the new doc" still asks the LLM to *generate* every section's
|
||||
blocks, including ones it didn't intend to modify, which gives it the same
|
||||
opportunity to drift.
|
||||
- Operations make the no-change case mechanical: zero ops → identical doc.
|
||||
- Operations are auditable: each refresh produces a log of exactly what
|
||||
changed, useful for debugging the LLM's behaviour and explaining diffs.
|
||||
|
||||
Failure modes are by design conservative: an operation list that fails to
|
||||
parse against the Pydantic schema, or an LLM that returns invalid ops, results
|
||||
in zero changes — the document stays as-is. The structure can only get better
|
||||
or stay the same per refresh, never get worse.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from typing import Annotated, Any, Literal, Union
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
|
||||
from .structured_doc import (
|
||||
Block,
|
||||
Section,
|
||||
StructuredDocument,
|
||||
make_unique_id,
|
||||
slugify_heading,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# Op payloads ---------------------------------------------------------------
|
||||
|
||||
|
||||
class _OpBase(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
|
||||
class AppendBlockOp(_OpBase):
|
||||
"""Add a new block at the end of an existing section."""
|
||||
|
||||
op: Literal["append_block"] = "append_block"
|
||||
section_id: str
|
||||
block: Block
|
||||
|
||||
|
||||
class InsertBlockOp(_OpBase):
|
||||
"""Insert a new block at ``index`` in an existing section.
|
||||
|
||||
``index`` may equal ``len(section.blocks)`` (append) but not be greater.
|
||||
"""
|
||||
|
||||
op: Literal["insert_block"] = "insert_block"
|
||||
section_id: str
|
||||
index: int = Field(ge=0)
|
||||
block: Block
|
||||
|
||||
|
||||
class ReplaceBlockOp(_OpBase):
|
||||
"""Replace the block at ``index`` of an existing section."""
|
||||
|
||||
op: Literal["replace_block"] = "replace_block"
|
||||
section_id: str
|
||||
index: int = Field(ge=0)
|
||||
block: Block
|
||||
|
||||
|
||||
class RemoveBlockOp(_OpBase):
|
||||
"""Remove the block at ``index`` of an existing section."""
|
||||
|
||||
op: Literal["remove_block"] = "remove_block"
|
||||
section_id: str
|
||||
index: int = Field(ge=0)
|
||||
|
||||
|
||||
class AddSectionOp(_OpBase):
|
||||
"""Add a brand-new section.
|
||||
|
||||
``after_section_id`` is optional; when omitted the new section is appended
|
||||
at the end. ``new_id`` is optional; when omitted we slugify the heading
|
||||
and disambiguate against existing IDs.
|
||||
"""
|
||||
|
||||
op: Literal["add_section"] = "add_section"
|
||||
heading: str
|
||||
level: int = Field(default=2, ge=1, le=6)
|
||||
blocks: list[Block] = Field(default_factory=list)
|
||||
after_section_id: str | None = None
|
||||
new_id: str | None = None
|
||||
|
||||
|
||||
class RemoveSectionOp(_OpBase):
|
||||
"""Remove an entire section by id."""
|
||||
|
||||
op: Literal["remove_section"] = "remove_section"
|
||||
section_id: str
|
||||
|
||||
|
||||
class ReplaceSectionBlocksOp(_OpBase):
|
||||
"""Replace all blocks of a section in one go.
|
||||
|
||||
Used when most of a section's contents are stale and rebuilding it as a
|
||||
unit is clearer than emitting many block-level ops. The section's heading
|
||||
and id are preserved.
|
||||
"""
|
||||
|
||||
op: Literal["replace_section_blocks"] = "replace_section_blocks"
|
||||
section_id: str
|
||||
blocks: list[Block] = Field(default_factory=list)
|
||||
|
||||
|
||||
class RenameSectionOp(_OpBase):
|
||||
"""Rename a section's heading. The id is unchanged so future ops still resolve."""
|
||||
|
||||
op: Literal["rename_section"] = "rename_section"
|
||||
section_id: str
|
||||
new_heading: str
|
||||
|
||||
|
||||
Operation = Annotated[
|
||||
Union[
|
||||
AppendBlockOp,
|
||||
InsertBlockOp,
|
||||
ReplaceBlockOp,
|
||||
RemoveBlockOp,
|
||||
AddSectionOp,
|
||||
RemoveSectionOp,
|
||||
ReplaceSectionBlocksOp,
|
||||
RenameSectionOp,
|
||||
],
|
||||
Field(discriminator="op"),
|
||||
]
|
||||
|
||||
|
||||
class DeltaOperationList(BaseModel):
|
||||
"""Container for the operations produced by an LLM delta call."""
|
||||
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
operations: list[Operation] = Field(default_factory=list)
|
||||
|
||||
|
||||
# Application ---------------------------------------------------------------
|
||||
|
||||
|
||||
class AppliedDelta(BaseModel):
|
||||
"""Outcome of applying a list of operations to a document."""
|
||||
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
document: StructuredDocument
|
||||
applied: list[dict[str, Any]] = Field(default_factory=list)
|
||||
skipped: list[dict[str, Any]] = Field(default_factory=list)
|
||||
|
||||
@property
|
||||
def changed(self) -> bool:
|
||||
return len(self.applied) > 0
|
||||
|
||||
|
||||
def _op_summary(op: Operation) -> dict[str, Any]:
|
||||
"""Compact dict suitable for the audit trail."""
|
||||
data = op.model_dump()
|
||||
return {k: v for k, v in data.items() if k != "block" and k != "blocks"} | {
|
||||
"op": data["op"],
|
||||
}
|
||||
|
||||
|
||||
def apply_operations(
|
||||
doc: StructuredDocument,
|
||||
operations: list[Operation],
|
||||
) -> AppliedDelta:
|
||||
"""Apply a list of operations to a document, returning a new document.
|
||||
|
||||
The original document is never mutated. Invalid operations (unknown
|
||||
section, out-of-range index, name collision when adding a section) are
|
||||
skipped and recorded in ``skipped`` with a ``reason`` string.
|
||||
"""
|
||||
new_doc = doc.model_copy(deep=True)
|
||||
applied: list[dict[str, Any]] = []
|
||||
skipped: list[dict[str, Any]] = []
|
||||
|
||||
def skip(op: Operation, reason: str) -> None:
|
||||
entry = _op_summary(op)
|
||||
entry["reason"] = reason
|
||||
skipped.append(entry)
|
||||
logger.debug(f"[STRUCTURED_DELTA] skipping op {entry}")
|
||||
|
||||
for op in operations:
|
||||
if isinstance(op, AppendBlockOp):
|
||||
section = new_doc.section_by_id(op.section_id)
|
||||
if section is None:
|
||||
skip(op, f"unknown section_id: {op.section_id}")
|
||||
continue
|
||||
section.blocks.append(op.block)
|
||||
applied.append(_op_summary(op))
|
||||
continue
|
||||
|
||||
if isinstance(op, InsertBlockOp):
|
||||
section = new_doc.section_by_id(op.section_id)
|
||||
if section is None:
|
||||
skip(op, f"unknown section_id: {op.section_id}")
|
||||
continue
|
||||
if op.index > len(section.blocks):
|
||||
skip(
|
||||
op,
|
||||
f"index out of range: {op.index} > {len(section.blocks)}",
|
||||
)
|
||||
continue
|
||||
section.blocks.insert(op.index, op.block)
|
||||
applied.append(_op_summary(op))
|
||||
continue
|
||||
|
||||
if isinstance(op, ReplaceBlockOp):
|
||||
section = new_doc.section_by_id(op.section_id)
|
||||
if section is None:
|
||||
skip(op, f"unknown section_id: {op.section_id}")
|
||||
continue
|
||||
if op.index >= len(section.blocks):
|
||||
skip(
|
||||
op,
|
||||
f"index out of range: {op.index} >= {len(section.blocks)}",
|
||||
)
|
||||
continue
|
||||
section.blocks[op.index] = op.block
|
||||
applied.append(_op_summary(op))
|
||||
continue
|
||||
|
||||
if isinstance(op, RemoveBlockOp):
|
||||
section = new_doc.section_by_id(op.section_id)
|
||||
if section is None:
|
||||
skip(op, f"unknown section_id: {op.section_id}")
|
||||
continue
|
||||
if op.index >= len(section.blocks):
|
||||
skip(
|
||||
op,
|
||||
f"index out of range: {op.index} >= {len(section.blocks)}",
|
||||
)
|
||||
continue
|
||||
section.blocks.pop(op.index)
|
||||
applied.append(_op_summary(op))
|
||||
continue
|
||||
|
||||
if isinstance(op, AddSectionOp):
|
||||
existing_ids = {s.id for s in new_doc.sections}
|
||||
base_id = op.new_id or slugify_heading(op.heading)
|
||||
section_id = make_unique_id(base_id, existing_ids)
|
||||
new_section = Section(
|
||||
id=section_id,
|
||||
heading=op.heading,
|
||||
level=op.level,
|
||||
blocks=list(op.blocks),
|
||||
)
|
||||
if op.after_section_id is None:
|
||||
new_doc.sections.append(new_section)
|
||||
else:
|
||||
idx = new_doc.section_index(op.after_section_id)
|
||||
if idx is None:
|
||||
skip(op, f"unknown after_section_id: {op.after_section_id}")
|
||||
continue
|
||||
new_doc.sections.insert(idx + 1, new_section)
|
||||
entry = _op_summary(op)
|
||||
entry["assigned_id"] = section_id
|
||||
applied.append(entry)
|
||||
continue
|
||||
|
||||
if isinstance(op, RemoveSectionOp):
|
||||
idx = new_doc.section_index(op.section_id)
|
||||
if idx is None:
|
||||
skip(op, f"unknown section_id: {op.section_id}")
|
||||
continue
|
||||
new_doc.sections.pop(idx)
|
||||
applied.append(_op_summary(op))
|
||||
continue
|
||||
|
||||
if isinstance(op, ReplaceSectionBlocksOp):
|
||||
section = new_doc.section_by_id(op.section_id)
|
||||
if section is None:
|
||||
skip(op, f"unknown section_id: {op.section_id}")
|
||||
continue
|
||||
section.blocks = list(op.blocks)
|
||||
applied.append(_op_summary(op))
|
||||
continue
|
||||
|
||||
if isinstance(op, RenameSectionOp):
|
||||
section = new_doc.section_by_id(op.section_id)
|
||||
if section is None:
|
||||
skip(op, f"unknown section_id: {op.section_id}")
|
||||
continue
|
||||
section.heading = op.new_heading
|
||||
applied.append(_op_summary(op))
|
||||
continue
|
||||
|
||||
skip(op, f"unhandled op type: {type(op).__name__}") # pragma: no cover
|
||||
|
||||
return AppliedDelta(document=new_doc, applied=applied, skipped=skipped)
|
||||
@@ -18,9 +18,6 @@ _TIKTOKEN_ENCODING = tiktoken.get_encoding("cl100k_base")
|
||||
# The remainder covers the system prompt, question, bank context, and output tokens.
|
||||
_FINAL_PROMPT_CONTEXT_FRACTION = 0.8
|
||||
|
||||
_DEFAULT_ROLE = "You are a reflection agent that answers questions by reasoning over retrieved memories."
|
||||
_DEFAULT_FINAL_ROLE = "You are a thoughtful assistant that synthesizes answers from retrieved memories."
|
||||
|
||||
|
||||
def _extract_directive_rules(directives: list[dict[str, Any]]) -> list[str]:
|
||||
"""Extract directive rules as a list of strings."""
|
||||
@@ -136,9 +133,7 @@ def build_system_prompt_for_tools(
|
||||
|
||||
parts.extend(
|
||||
[
|
||||
mission.strip() if mission else _DEFAULT_ROLE,
|
||||
"",
|
||||
"Answer the user's question by reasoning over retrieved memories.",
|
||||
"You are a reflection agent that answers questions by reasoning over retrieved memories.",
|
||||
"",
|
||||
]
|
||||
)
|
||||
@@ -374,7 +369,7 @@ def build_agent_prompt(
|
||||
output = entry["output"]
|
||||
# Format as proper JSON for LLM readability
|
||||
try:
|
||||
output_str = json.dumps(output, indent=2, default=str, ensure_ascii=False)
|
||||
output_str = json.dumps(output, indent=2, default=str)
|
||||
except (TypeError, ValueError):
|
||||
output_str = str(output)
|
||||
parts.append(f"\n### Call {i}: {tool}\n```json\n{output_str}\n```")
|
||||
@@ -449,7 +444,7 @@ def build_final_prompt(
|
||||
tool = entry["tool"]
|
||||
output = entry["output"]
|
||||
try:
|
||||
output_str = json.dumps(output, indent=2, default=str, ensure_ascii=False)
|
||||
output_str = json.dumps(output, indent=2, default=str)
|
||||
except (TypeError, ValueError):
|
||||
output_str = str(output)
|
||||
block = f"\n### From {tool}:\n```json\n{output_str}\n```"
|
||||
@@ -484,9 +479,9 @@ def build_final_prompt(
|
||||
return "\n".join(parts)
|
||||
|
||||
|
||||
_FINAL_SYSTEM_PROMPT_BASE = """CRITICAL: You MUST ONLY use information from retrieved tool results. NEVER make up names, people, events, or entities.
|
||||
FINAL_SYSTEM_PROMPT = """CRITICAL: You MUST ONLY use information from retrieved tool results. NEVER make up names, people, events, or entities.
|
||||
|
||||
{role_section}
|
||||
You are a thoughtful assistant that synthesizes answers from retrieved memories.
|
||||
|
||||
Your approach:
|
||||
- Reason over the retrieved memories to answer the question
|
||||
@@ -513,213 +508,3 @@ CRITICAL: Output ONLY the final synthesized answer. Do NOT include:
|
||||
Just provide the direct answer with proper markdown formatting.
|
||||
|
||||
CRITICAL: This is a NON-CONVERSATIONAL system. NEVER ask follow-up questions, offer to search again, suggest alternatives, or end with anything like "Would you like me to..." or "Let me know if...". The user cannot reply. Your answer must be complete and self-contained."""
|
||||
|
||||
|
||||
def build_final_system_prompt(mission: str | None = None) -> str:
|
||||
"""Build the final synthesis system prompt, using mission as role when set."""
|
||||
role_section = mission.strip() if mission else _DEFAULT_FINAL_ROLE
|
||||
return _FINAL_SYSTEM_PROMPT_BASE.format(role_section=role_section)
|
||||
|
||||
|
||||
# Backward-compatible constant for non-identity missions
|
||||
FINAL_SYSTEM_PROMPT = build_final_system_prompt()
|
||||
|
||||
|
||||
STRUCTURED_DELTA_SYSTEM_PROMPT = """You are integrating *new information* into an existing structured document.
|
||||
|
||||
You will be given:
|
||||
1. TOPIC — the question this document answers. Content that does not help
|
||||
answer this question is OFF-TOPIC and should be removed.
|
||||
2. CURRENT DOCUMENT (JSON) — the existing structured mental model. Each section
|
||||
has a stable ``id``, a ``heading``, a ``level`` (1..6), and an ordered list
|
||||
of ``blocks``. Blocks are typed: ``paragraph``, ``bullet_list``,
|
||||
``ordered_list``, or ``code``.
|
||||
3. NEW INFORMATION SYNTHESIS (markdown) — a synthesis showing how the new facts
|
||||
relate to the document's topic. Use it to understand context and relevance,
|
||||
but do NOT copy its formatting or wording wholesale.
|
||||
4. SUPPORTING FACTS — observations and facts created since the last refresh.
|
||||
These are genuinely new — they were NOT available when the current document
|
||||
was written.
|
||||
|
||||
Your task: output a JSON object ``{"operations": [...]}``. Applied to CURRENT
|
||||
DOCUMENT, the operations must produce a document that best answers the TOPIC
|
||||
by integrating the new facts.
|
||||
|
||||
RULES
|
||||
- These facts are NEW since the last refresh. The existing document already
|
||||
captures all prior information from earlier refreshes. Your job is to
|
||||
integrate the new facts into the existing document.
|
||||
- **Preserve existing content**: The current document was built from prior facts
|
||||
that you cannot see. Do NOT remove or replace existing sections just because
|
||||
the new facts do not reference them. Only remove content when the new facts
|
||||
explicitly contradict or supersede it.
|
||||
- **Merge overlapping topics**: When new facts cover topics that overlap with
|
||||
existing sections, merge the new information INTO the existing section
|
||||
rather than creating duplicates. When new facts provide more specific or
|
||||
authoritative guidance on a topic already covered generically, update the
|
||||
existing content to reflect the more specific guidance.
|
||||
- **Preserve examples**: Concrete examples, before/after pairs, sample sentences,
|
||||
and illustrative ✅/❌ comparisons are MORE valuable than abstract rules.
|
||||
When facts contain examples, include them. Never drop an example to make
|
||||
room for an abstract restatement of the same point.
|
||||
- Operations target sections by ``section_id`` (use the ``id`` field of the
|
||||
section in CURRENT DOCUMENT, NOT the heading). Block operations target
|
||||
blocks by ``index`` (0-based, against the section's current block list).
|
||||
- **Add** new content with ``append_block``, ``insert_block``, or ``add_section``
|
||||
when facts introduce information not yet covered. Prefer extending an
|
||||
existing section over creating a new one.
|
||||
- **Update** existing content with ``replace_block`` or ``replace_section_blocks``
|
||||
when new facts provide corrections, updates, or more specific information
|
||||
about topics already in the document.
|
||||
- **Remove** content with ``remove_block`` or ``remove_section`` ONLY when
|
||||
the new facts explicitly contradict or supersede it.
|
||||
- NEVER emit operations whose only effect is to reword unchanged content.
|
||||
- NEVER emit operations to "normalize" formatting (numbered → bulleted, casing
|
||||
changes, paragraph → list, etc).
|
||||
- Every operation MUST be justifiable by a specific fact in SUPPORTING FACTS.
|
||||
- Output ``{"operations": []}`` only if the new facts are already reflected
|
||||
in the document (e.g., from a concurrent update).
|
||||
|
||||
ALLOWED OPERATIONS (each line shows the JSON shape)
|
||||
- ``{"op": "append_block", "section_id": "...", "block": {...}}``
|
||||
- ``{"op": "insert_block", "section_id": "...", "index": N, "block": {...}}``
|
||||
- ``{"op": "replace_block", "section_id": "...", "index": N, "block": {...}}``
|
||||
- ``{"op": "remove_block", "section_id": "...", "index": N}``
|
||||
- ``{"op": "add_section", "heading": "...", "level": 2, "blocks": [...], "after_section_id": "..."}``
|
||||
- ``{"op": "remove_section", "section_id": "..."}``
|
||||
- ``{"op": "replace_section_blocks", "section_id": "...", "blocks": [...]}``
|
||||
- ``{"op": "rename_section", "section_id": "...", "new_heading": "..."}``
|
||||
|
||||
Block shapes
|
||||
- ``{"type": "paragraph", "text": "..."}``
|
||||
- ``{"type": "bullet_list", "items": ["...", "..."]}``
|
||||
- ``{"type": "ordered_list", "items": ["...", "..."]}``
|
||||
- ``{"type": "code", "language": "json", "text": "..."}``
|
||||
|
||||
OUTPUT FORMAT
|
||||
Return ONLY a single JSON object on its own, with no prose before or after,
|
||||
no markdown code fences, no commentary. The object must have exactly one
|
||||
top-level key, ``operations``, whose value is an array of operation objects
|
||||
(empty array when nothing changes).
|
||||
|
||||
Examples
|
||||
- No changes needed → ``{"operations": []}``
|
||||
- Add one bullet to an existing "Members" section →
|
||||
``{"operations": [{"op": "append_block", "section_id": "members",
|
||||
"block": {"type": "bullet_list", "items": ["Carol — junior engineer"]}}]}``
|
||||
- Replace a paragraph that has been corrected by new facts →
|
||||
``{"operations": [{"op": "replace_block", "section_id": "overview",
|
||||
"index": 0, "block": {"type": "paragraph", "text": "Updated summary."}}]}``
|
||||
- Remove an obsolete block →
|
||||
``{"operations": [{"op": "remove_block", "section_id": "status", "index": 2}]}``"""
|
||||
|
||||
|
||||
def build_structured_delta_prompt(
|
||||
*,
|
||||
current_document_json: str,
|
||||
candidate_markdown: str,
|
||||
supporting_facts: list[dict[str, Any]],
|
||||
source_query: str,
|
||||
max_output_tokens: int | None = None,
|
||||
) -> str:
|
||||
"""Build the user prompt for a structured-delta mental model refresh.
|
||||
|
||||
The LLM's job is to emit operations against ``current_document_json``;
|
||||
the surrounding ``candidate_markdown`` and ``supporting_facts`` are
|
||||
references for *what new information exists*, not templates to mimic.
|
||||
|
||||
``max_output_tokens`` is surfaced in the prompt so the model can keep its
|
||||
op list within the provider's response cap. The actual cap is enforced by
|
||||
the caller; this is just an advisory anchor — without it the model often
|
||||
returns op lists whose JSON gets truncated mid-string.
|
||||
"""
|
||||
fact_lines: list[str] = []
|
||||
for f in supporting_facts:
|
||||
fid = f.get("id", "")
|
||||
text = (f.get("text") or "").strip().replace("\n", " ")
|
||||
ftype = f.get("type", "")
|
||||
fact_lines.append(f"- [{ftype}:{fid}] {text}")
|
||||
facts_block = "\n".join(fact_lines) if fact_lines else "(no supporting facts retrieved)"
|
||||
|
||||
budget_hint = ""
|
||||
if max_output_tokens is not None:
|
||||
budget_hint = (
|
||||
f"\n\n## Output budget\n"
|
||||
f"Your JSON response must fit within ~{max_output_tokens} tokens. If you "
|
||||
"would need more than this to express every change, prefer the highest-"
|
||||
"leverage edits first (a few ``replace_section_blocks`` ops over many "
|
||||
"block-level ops) so the response always parses as valid JSON."
|
||||
)
|
||||
|
||||
return (
|
||||
f"## Topic\n{source_query}\n\n"
|
||||
f"## CURRENT DOCUMENT (apply ops to this; reference section ids as listed)\n"
|
||||
f"```json\n{current_document_json}\n```\n\n"
|
||||
f"## NEW INFORMATION SYNTHESIS (context for how new facts relate to the topic)\n"
|
||||
f"```markdown\n{candidate_markdown}\n```\n\n"
|
||||
f"## SUPPORTING FACTS (new since last refresh — integrate these)\n{facts_block}"
|
||||
f"{budget_hint}\n\n"
|
||||
"## Task\n"
|
||||
"Output a JSON object matching the operations schema. Integrate the new "
|
||||
"supporting facts into CURRENT DOCUMENT. Add, update, or remove content "
|
||||
"as needed. Preserve unchanged sections and blocks by not mentioning them."
|
||||
)
|
||||
|
||||
|
||||
DELTA_SYSTEM_PROMPT = """You are performing a surgical delta update to an existing mental model document.
|
||||
|
||||
You will be given:
|
||||
1. CURRENT DOCUMENT: the existing mental model content (markdown).
|
||||
2. CANDIDATE UPDATE: a freshly generated synthesis based on the latest retrieved memories.
|
||||
3. SUPPORTING FACTS: the observations and facts that support the CANDIDATE UPDATE.
|
||||
|
||||
Your task: produce an updated version of the CURRENT DOCUMENT that reflects the new reality, with the MINIMUM possible changes.
|
||||
|
||||
ABSOLUTE RULES:
|
||||
- Preserve unchanged content BYTE-FOR-BYTE. If a sentence, heading, bullet, code block, or section is still accurate according to the CANDIDATE UPDATE and SUPPORTING FACTS, copy it verbatim — same wording, same punctuation, same whitespace, same markdown structure.
|
||||
- Do NOT reformat, rephrase, or re-style content that is still accurate. No "light edits for clarity", no reordering for flow, no synonym swaps.
|
||||
- Remove content that is contradicted by the CANDIDATE UPDATE or SUPPORTING FACTS (stale content).
|
||||
- Add new content ONLY when the SUPPORTING FACTS contain information not already in the CURRENT DOCUMENT.
|
||||
- When adding new content, prefer appending to an existing relevant section. Creating a new section is acceptable when the new information does not fit any existing section.
|
||||
- When creating a new section, match the heading style, tone, and formatting conventions used in the CURRENT DOCUMENT.
|
||||
- Every assertion in your output MUST be grounded in either (a) the CURRENT DOCUMENT (preserved) or (b) the SUPPORTING FACTS. Never introduce outside knowledge.
|
||||
- If nothing in the SUPPORTING FACTS contradicts or extends the CURRENT DOCUMENT, return the CURRENT DOCUMENT UNCHANGED, character for character.
|
||||
|
||||
OUTPUT FORMAT:
|
||||
- Output ONLY the updated markdown document. No preamble, no explanation, no diff markers, no commentary.
|
||||
- Do not wrap the output in code fences unless the CURRENT DOCUMENT itself was entirely a code fence."""
|
||||
|
||||
|
||||
def build_delta_prompt(
|
||||
*,
|
||||
current_content: str,
|
||||
candidate_content: str,
|
||||
supporting_facts: list[dict[str, Any]],
|
||||
source_query: str,
|
||||
) -> str:
|
||||
"""Build the user prompt for a delta-mode mental model refresh.
|
||||
|
||||
Args:
|
||||
current_content: The existing mental model content (to preserve as much as possible).
|
||||
candidate_content: Fresh synthesis from the reflect agent reflecting new reality.
|
||||
supporting_facts: Flat list of fact dicts (id, text, type) supporting the candidate.
|
||||
source_query: The mental model's source query, for topical framing.
|
||||
"""
|
||||
fact_lines: list[str] = []
|
||||
for f in supporting_facts:
|
||||
fid = f.get("id", "")
|
||||
text = (f.get("text") or "").strip().replace("\n", " ")
|
||||
ftype = f.get("type", "")
|
||||
fact_lines.append(f"- [{ftype}:{fid}] {text}")
|
||||
facts_block = "\n".join(fact_lines) if fact_lines else "(no supporting facts retrieved)"
|
||||
|
||||
return (
|
||||
f"## Topic\n{source_query}\n\n"
|
||||
f"## CURRENT DOCUMENT\n```markdown\n{current_content}\n```\n\n"
|
||||
f"## CANDIDATE UPDATE\n```markdown\n{candidate_content}\n```\n\n"
|
||||
f"## SUPPORTING FACTS\n{facts_block}\n\n"
|
||||
"## Task\n"
|
||||
"Produce the updated mental model document by applying the minimum necessary changes "
|
||||
"to CURRENT DOCUMENT so that it reflects CANDIDATE UPDATE and SUPPORTING FACTS. "
|
||||
"Preserve unchanged content byte-for-byte. Output only the final markdown."
|
||||
)
|
||||
|
||||
@@ -1,301 +0,0 @@
|
||||
"""Structured representation of a mental model document.
|
||||
|
||||
Why this exists
|
||||
---------------
|
||||
Storing mental models as raw markdown forces every refresh to round-trip prose
|
||||
through an LLM, which then drifts on stylistic details (numbered vs bulleted
|
||||
lists, casing, separator lines, paraphrasing) even when instructed to preserve
|
||||
content byte-for-byte. The intrinsic mechanism of an LLM is to *generate* the
|
||||
next token from a gestalt of the input — not to copy tokens verbatim — so any
|
||||
"preserve unchanged content" instruction is fundamentally a soft constraint.
|
||||
|
||||
The fix is to give the LLM no opportunity to drift on unchanged content. We
|
||||
keep an authoritative structured representation of the document; the markdown
|
||||
shown to users is a deterministic render of that structure. Delta refreshes
|
||||
emit *operations* against the structure (see ``delta_ops.py``); sections and
|
||||
blocks not mentioned by any operation are physically untouched.
|
||||
|
||||
Schema (v1)
|
||||
-----------
|
||||
A document is an ordered list of ``Section``s. Each section has:
|
||||
- ``id`` : stable slug derived from ``heading`` (used as the operation
|
||||
target across refreshes; surviving renames is a separate
|
||||
concern handled by an explicit ``rename`` op).
|
||||
- ``heading``: the markdown heading text (without the ``#`` prefix).
|
||||
- ``level`` : 1 (``#``) … 6 (``######``). Default 2.
|
||||
- ``blocks``: ordered list of typed blocks — paragraph, bullet_list,
|
||||
ordered_list, code.
|
||||
|
||||
The schema is intentionally narrow: it covers what real mental-model documents
|
||||
actually contain (the kind a coding agent writes for itself or a user writes as
|
||||
a "skill" doc). Tables, images, and raw HTML are out of scope until needed.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from typing import Annotated, Literal, Union
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
|
||||
# Blocks ---------------------------------------------------------------------
|
||||
|
||||
|
||||
class ParagraphBlock(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
type: Literal["paragraph"] = "paragraph"
|
||||
text: str
|
||||
|
||||
|
||||
class BulletListBlock(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
type: Literal["bullet_list"] = "bullet_list"
|
||||
items: list[str] = Field(default_factory=list)
|
||||
|
||||
|
||||
class OrderedListBlock(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
type: Literal["ordered_list"] = "ordered_list"
|
||||
items: list[str] = Field(default_factory=list)
|
||||
|
||||
|
||||
class CodeBlock(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
type: Literal["code"] = "code"
|
||||
language: str = ""
|
||||
text: str
|
||||
|
||||
|
||||
Block = Annotated[
|
||||
Union[ParagraphBlock, BulletListBlock, OrderedListBlock, CodeBlock],
|
||||
Field(discriminator="type"),
|
||||
]
|
||||
|
||||
|
||||
# Section / Document ---------------------------------------------------------
|
||||
|
||||
|
||||
class Section(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
id: str
|
||||
heading: str
|
||||
level: int = Field(default=2, ge=1, le=6)
|
||||
blocks: list[Block] = Field(default_factory=list)
|
||||
|
||||
|
||||
class StructuredDocument(BaseModel):
|
||||
"""Top-level structured representation of a mental model."""
|
||||
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
version: Literal[1] = 1
|
||||
sections: list[Section] = Field(default_factory=list)
|
||||
|
||||
def section_by_id(self, section_id: str) -> Section | None:
|
||||
for s in self.sections:
|
||||
if s.id == section_id:
|
||||
return s
|
||||
return None
|
||||
|
||||
def section_index(self, section_id: str) -> int | None:
|
||||
for i, s in enumerate(self.sections):
|
||||
if s.id == section_id:
|
||||
return i
|
||||
return None
|
||||
|
||||
|
||||
# Slug helpers ---------------------------------------------------------------
|
||||
|
||||
_SLUG_RX = re.compile(r"[^a-z0-9]+")
|
||||
|
||||
|
||||
def slugify_heading(heading: str) -> str:
|
||||
"""Stable, deterministic slug from a heading.
|
||||
|
||||
"Stop Conditions" -> "stop-conditions"
|
||||
"Inputs and Context" -> "inputs-and-context"
|
||||
"""
|
||||
slug = _SLUG_RX.sub("-", heading.strip().lower()).strip("-")
|
||||
return slug or "section"
|
||||
|
||||
|
||||
def make_unique_id(base: str, existing: set[str]) -> str:
|
||||
"""Disambiguate by appending -2, -3, … if the slug is already in use."""
|
||||
if base not in existing:
|
||||
return base
|
||||
i = 2
|
||||
while f"{base}-{i}" in existing:
|
||||
i += 1
|
||||
return f"{base}-{i}"
|
||||
|
||||
|
||||
# Renderer -------------------------------------------------------------------
|
||||
|
||||
|
||||
def render_block(block: Block) -> str:
|
||||
"""Render a single block to markdown. No trailing newline."""
|
||||
if isinstance(block, ParagraphBlock):
|
||||
return block.text.rstrip()
|
||||
if isinstance(block, BulletListBlock):
|
||||
return "\n".join(f"- {item.rstrip()}" for item in block.items)
|
||||
if isinstance(block, OrderedListBlock):
|
||||
return "\n".join(f"{i + 1}. {item.rstrip()}" for i, item in enumerate(block.items))
|
||||
if isinstance(block, CodeBlock):
|
||||
fence_lang = block.language or ""
|
||||
return f"```{fence_lang}\n{block.text}\n```"
|
||||
raise TypeError(f"Unknown block type: {type(block)!r}")
|
||||
|
||||
|
||||
def render_section(section: Section) -> str:
|
||||
"""Render a section: heading + blank line + blocks separated by blank lines."""
|
||||
parts = ["#" * section.level + " " + section.heading.strip()]
|
||||
for block in section.blocks:
|
||||
parts.append("") # blank line before each block
|
||||
parts.append(render_block(block))
|
||||
return "\n".join(parts)
|
||||
|
||||
|
||||
def render_document(doc: StructuredDocument) -> str:
|
||||
"""Render the whole document. Sections separated by a single blank line.
|
||||
|
||||
The output is byte-stable: same structured input always produces the same
|
||||
markdown, modulo the inherent ordering of sections/blocks/items.
|
||||
"""
|
||||
if not doc.sections:
|
||||
return ""
|
||||
return "\n\n".join(render_section(s) for s in doc.sections) + "\n"
|
||||
|
||||
|
||||
# Parser ---------------------------------------------------------------------
|
||||
#
|
||||
# The parser is intentionally lenient: it accepts the markdown produced by
|
||||
# our own renderer (round-trip-safe) and the markdown an LLM tends to produce
|
||||
# for mental-model documents. It is *not* a general CommonMark parser — it
|
||||
# does not need to be. When it cannot classify a block it falls back to a
|
||||
# paragraph so that no content is silently dropped.
|
||||
|
||||
_HEADING_RX = re.compile(r"^(#{1,6})\s+(.+?)\s*$")
|
||||
_BULLET_RX = re.compile(r"^\s*[-*+]\s+(.*)$")
|
||||
_ORDERED_RX = re.compile(r"^\s*\d+[.)]\s+(.*)$")
|
||||
_FENCE_RX = re.compile(r"^```([A-Za-z0-9_+-]*)\s*$")
|
||||
|
||||
|
||||
def _strip_separators(lines: list[str]) -> list[str]:
|
||||
"""Drop horizontal-rule lines (`---`, `***`) used as section separators.
|
||||
|
||||
Our renderer never emits these, but LLM output frequently includes them
|
||||
between sections; treating them as blank lines avoids parsing them as
|
||||
paragraphs.
|
||||
"""
|
||||
return ["" if re.fullmatch(r"\s*([-*_])\1{2,}\s*", line) else line for line in lines]
|
||||
|
||||
|
||||
def _split_blocks(lines: list[str]) -> list[list[str]]:
|
||||
"""Group consecutive non-blank lines into block chunks."""
|
||||
chunks: list[list[str]] = []
|
||||
current: list[str] = []
|
||||
in_fence = False
|
||||
for line in lines:
|
||||
if _FENCE_RX.match(line):
|
||||
current.append(line)
|
||||
in_fence = not in_fence
|
||||
continue
|
||||
if in_fence:
|
||||
current.append(line)
|
||||
continue
|
||||
if line.strip() == "":
|
||||
if current:
|
||||
chunks.append(current)
|
||||
current = []
|
||||
else:
|
||||
current.append(line)
|
||||
if current:
|
||||
chunks.append(current)
|
||||
return chunks
|
||||
|
||||
|
||||
def _parse_block(chunk: list[str]) -> Block:
|
||||
"""Parse a single non-empty chunk into a block."""
|
||||
if chunk and _FENCE_RX.match(chunk[0]):
|
||||
m = _FENCE_RX.match(chunk[0])
|
||||
lang = m.group(1) if m else ""
|
||||
body_lines = chunk[1:]
|
||||
if body_lines and _FENCE_RX.match(body_lines[-1]):
|
||||
body_lines = body_lines[:-1]
|
||||
return CodeBlock(language=lang, text="\n".join(body_lines))
|
||||
|
||||
if all(_BULLET_RX.match(line) for line in chunk):
|
||||
items = []
|
||||
for line in chunk:
|
||||
m = _BULLET_RX.match(line)
|
||||
assert m is not None
|
||||
items.append(m.group(1).strip())
|
||||
return BulletListBlock(items=items)
|
||||
|
||||
if all(_ORDERED_RX.match(line) for line in chunk):
|
||||
items = []
|
||||
for line in chunk:
|
||||
m = _ORDERED_RX.match(line)
|
||||
assert m is not None
|
||||
items.append(m.group(1).strip())
|
||||
return OrderedListBlock(items=items)
|
||||
|
||||
return ParagraphBlock(text=" ".join(line.strip() for line in chunk).strip())
|
||||
|
||||
|
||||
def parse_markdown(markdown: str) -> StructuredDocument:
|
||||
"""Best-effort parse of a markdown document into the structured schema.
|
||||
|
||||
Sections are introduced by ATX headings (``#``..``######``). Anything
|
||||
before the first heading is wrapped into an implicit "Overview" section
|
||||
so we never silently drop user content. Section IDs are unique slugs of
|
||||
their headings.
|
||||
"""
|
||||
raw_lines = (markdown or "").splitlines()
|
||||
lines = _strip_separators(raw_lines)
|
||||
|
||||
sections: list[Section] = []
|
||||
used_ids: set[str] = set()
|
||||
pending: list[str] = []
|
||||
current: Section | None = None
|
||||
|
||||
def flush_pending_into(section: Section) -> None:
|
||||
if not pending:
|
||||
return
|
||||
for chunk in _split_blocks(pending):
|
||||
section.blocks.append(_parse_block(chunk))
|
||||
pending.clear()
|
||||
|
||||
for line in lines:
|
||||
m = _HEADING_RX.match(line)
|
||||
if m:
|
||||
if current is not None:
|
||||
flush_pending_into(current)
|
||||
sections.append(current)
|
||||
elif pending:
|
||||
# Content before the first heading: wrap in implicit section.
|
||||
base = "overview"
|
||||
section_id = make_unique_id(base, used_ids)
|
||||
used_ids.add(section_id)
|
||||
implicit = Section(id=section_id, heading="Overview", level=2)
|
||||
flush_pending_into(implicit)
|
||||
sections.append(implicit)
|
||||
level = len(m.group(1))
|
||||
heading = m.group(2).strip()
|
||||
section_id = make_unique_id(slugify_heading(heading), used_ids)
|
||||
used_ids.add(section_id)
|
||||
current = Section(id=section_id, heading=heading, level=level)
|
||||
else:
|
||||
pending.append(line)
|
||||
|
||||
if current is not None:
|
||||
flush_pending_into(current)
|
||||
sections.append(current)
|
||||
elif pending:
|
||||
base = "overview"
|
||||
section_id = make_unique_id(base, used_ids)
|
||||
used_ids.add(section_id)
|
||||
implicit = Section(id=section_id, heading="Overview", level=2)
|
||||
flush_pending_into(implicit)
|
||||
sections.append(implicit)
|
||||
|
||||
return StructuredDocument(sections=sections)
|
||||
@@ -23,7 +23,6 @@ logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
async def tool_search_mental_models(
|
||||
memory_engine: "MemoryEngine",
|
||||
conn: "Connection",
|
||||
bank_id: str,
|
||||
query: str,
|
||||
@@ -33,6 +32,7 @@ async def tool_search_mental_models(
|
||||
tags_match: str = "any",
|
||||
tag_groups: "list | None" = None,
|
||||
exclude_ids: list[str] | None = None,
|
||||
pending_consolidation: int = 0,
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Search user-curated mental models by semantic similarity.
|
||||
@@ -82,7 +82,7 @@ async def tool_search_mental_models(
|
||||
f"""
|
||||
SELECT
|
||||
id, name, content,
|
||||
tags, created_at, last_refreshed_at, trigger,
|
||||
tags, created_at, last_refreshed_at,
|
||||
1 - (embedding <=> $2::vector) as relevance
|
||||
FROM {fq_table("mental_models")}
|
||||
WHERE bank_id = $1 AND embedding IS NOT NULL {filters}
|
||||
@@ -99,9 +99,10 @@ async def tool_search_mental_models(
|
||||
if last_refreshed_at and last_refreshed_at.tzinfo is None:
|
||||
last_refreshed_at = last_refreshed_at.replace(tzinfo=timezone.utc)
|
||||
|
||||
# Per-MM staleness: new in-scope memories since last refresh (includes pending).
|
||||
is_stale = await memory_engine.compute_mental_model_is_stale(conn, bank_id, row)
|
||||
staleness_reason = "new in-scope memories ingested since last refresh" if is_stale else None
|
||||
# A mental model is stale when there are memories that haven't been consolidated yet —
|
||||
# the same signal used for observations staleness.
|
||||
is_stale = pending_consolidation > 0
|
||||
staleness_reason = f"{pending_consolidation} memories pending consolidation" if is_stale else None
|
||||
|
||||
mental_models.append(
|
||||
{
|
||||
@@ -135,8 +136,6 @@ async def tool_search_observations(
|
||||
last_consolidated_at: datetime | None = None,
|
||||
pending_consolidation: int = 0,
|
||||
source_facts_max_tokens: int = -1,
|
||||
created_after: datetime | None = None,
|
||||
created_before: datetime | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Search consolidated observations using recall.
|
||||
@@ -180,8 +179,6 @@ async def tool_search_observations(
|
||||
tags_match=tags_match,
|
||||
tag_groups=tag_groups,
|
||||
include_source_facts=include_source_facts,
|
||||
created_after=created_after,
|
||||
created_before=created_before,
|
||||
_connection_budget=1,
|
||||
_quiet=True,
|
||||
**recall_kwargs,
|
||||
@@ -217,9 +214,6 @@ async def tool_recall(
|
||||
connection_budget: int = 1,
|
||||
max_chunk_tokens: int = 1000,
|
||||
fact_types: list[str] | None = None,
|
||||
include_chunks: bool = True,
|
||||
created_after: datetime | None = None,
|
||||
created_before: datetime | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Search memories using TEMPR retrieval.
|
||||
@@ -236,15 +230,15 @@ async def tool_recall(
|
||||
tags: Filter by tags (includes untagged memories)
|
||||
tags_match: How to match tags - "any" (OR), "all" (AND), or "exact"
|
||||
connection_budget: Max DB connections for this recall (default 1 for internal ops)
|
||||
max_chunk_tokens: Maximum tokens for raw source chunk text (default 1000)
|
||||
max_chunk_tokens: Maximum tokens for raw source chunk text (default 1000, always included)
|
||||
fact_types: Optional filter for fact types to retrieve. Defaults to ["experience", "world"].
|
||||
include_chunks: Whether to fetch raw chunk text alongside facts (default True).
|
||||
|
||||
Returns:
|
||||
Dict with list of matching memories including raw chunk text (when include_chunks)
|
||||
Dict with list of matching memories including raw chunk text
|
||||
"""
|
||||
# Only world/experience are valid for raw recall (observation is handled by search_observations)
|
||||
recall_fact_type = [ft for ft in (fact_types or ["experience", "world"]) if ft in ("world", "experience")]
|
||||
include_chunks = True
|
||||
internal_ctx = replace(request_context, internal=True)
|
||||
result = await memory_engine.recall_async(
|
||||
bank_id=bank_id,
|
||||
@@ -256,8 +250,6 @@ async def tool_recall(
|
||||
tags=tags,
|
||||
tags_match=tags_match,
|
||||
tag_groups=tag_groups,
|
||||
created_after=created_after,
|
||||
created_before=created_before,
|
||||
_connection_budget=connection_budget,
|
||||
_quiet=True, # Suppress logging for internal operations
|
||||
include_chunks=include_chunks,
|
||||
|
||||
@@ -47,16 +47,6 @@ async def generate_embeddings_batch(embeddings_backend, texts: list[str]) -> lis
|
||||
embeddings_backend.encode,
|
||||
texts,
|
||||
)
|
||||
return embeddings
|
||||
except Exception as e:
|
||||
raise Exception(f"Failed to generate batch embeddings: {str(e)}")
|
||||
|
||||
# Guarantee 1:1 alignment with input texts. A silent length mismatch here
|
||||
# propagates downstream as zip() drops items, eventually surfacing as an
|
||||
# IndexError in retain mapping (see issue #1037).
|
||||
if len(embeddings) != len(texts):
|
||||
raise RuntimeError(
|
||||
f"Embeddings backend returned {len(embeddings)} vectors for {len(texts)} input texts; "
|
||||
"expected exact 1:1 alignment"
|
||||
)
|
||||
|
||||
return embeddings
|
||||
|
||||
@@ -7,7 +7,6 @@ Handles insertion of facts into the database.
|
||||
import json
|
||||
import logging
|
||||
import uuid
|
||||
from datetime import datetime
|
||||
|
||||
from ...config import get_config
|
||||
from ..memory_engine import fq_table
|
||||
@@ -225,85 +224,6 @@ async def ensure_bank_exists(conn, bank_id: str) -> None:
|
||||
await create_bank_vector_indexes(conn, bank_id, str(internal_id))
|
||||
|
||||
|
||||
async def delete_stale_observations_for_memories(
|
||||
conn,
|
||||
bank_id: str,
|
||||
fact_ids: "list[str | uuid.UUID]",
|
||||
) -> int:
|
||||
"""Delete observations whose source memories are about to be removed.
|
||||
|
||||
Mirrors the cleanup performed by ``MemoryEngine.delete_document`` so that
|
||||
every code path that removes ``memory_units`` also removes the
|
||||
observations derived from them. Without this, ingesting a fresh version
|
||||
of a document via the retain pipeline (which does a full-replace
|
||||
``DELETE FROM documents`` cascade) used to leave orphan observations
|
||||
pointing at memory IDs that no longer existed.
|
||||
|
||||
For each observation referencing any of ``fact_ids``:
|
||||
1. Delete the observation row (its text is stale once even one source
|
||||
memory disappears).
|
||||
2. Reset ``consolidated_at = NULL`` on the surviving source memories so
|
||||
they get re-consolidated under fresh observations on the next run.
|
||||
|
||||
Must be called within an active transaction, before the source memories
|
||||
are deleted.
|
||||
|
||||
Returns the number of observations deleted.
|
||||
"""
|
||||
if not fact_ids:
|
||||
return 0
|
||||
|
||||
fact_uuids = [uuid.UUID(str(fid)) if not isinstance(fid, uuid.UUID) else fid for fid in fact_ids]
|
||||
|
||||
affected_obs = await conn.fetch(
|
||||
f"""
|
||||
SELECT id, source_memory_ids
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE bank_id = $1
|
||||
AND fact_type = 'observation'
|
||||
AND source_memory_ids && $2::uuid[]
|
||||
""",
|
||||
bank_id,
|
||||
fact_uuids,
|
||||
)
|
||||
|
||||
if not affected_obs:
|
||||
return 0
|
||||
|
||||
deleted_set = {str(uid) for uid in fact_uuids}
|
||||
obs_ids = [obs["id"] for obs in affected_obs]
|
||||
seen_remaining: set[str] = set()
|
||||
remaining_source_ids: list[uuid.UUID] = []
|
||||
for obs in affected_obs:
|
||||
for src_id in obs["source_memory_ids"] or []:
|
||||
src_str = str(src_id)
|
||||
if src_str not in deleted_set and src_str not in seen_remaining:
|
||||
remaining_source_ids.append(src_id)
|
||||
seen_remaining.add(src_str)
|
||||
|
||||
await conn.execute(
|
||||
f"DELETE FROM {fq_table('memory_units')} WHERE id = ANY($1::uuid[])",
|
||||
obs_ids,
|
||||
)
|
||||
|
||||
if remaining_source_ids:
|
||||
await conn.execute(
|
||||
f"""
|
||||
UPDATE {fq_table("memory_units")}
|
||||
SET consolidated_at = NULL
|
||||
WHERE id = ANY($1::uuid[])
|
||||
AND fact_type IN ('experience', 'world')
|
||||
""",
|
||||
remaining_source_ids,
|
||||
)
|
||||
|
||||
logger.info(
|
||||
f"[OBSERVATIONS] Deleted {len(obs_ids)} observations, reset {len(remaining_source_ids)} "
|
||||
f"source memories for re-consolidation in bank {bank_id}"
|
||||
)
|
||||
return len(obs_ids)
|
||||
|
||||
|
||||
async def handle_document_tracking(
|
||||
conn,
|
||||
bank_id: str,
|
||||
@@ -334,58 +254,17 @@ async def handle_document_tracking(
|
||||
combined_content = _sanitize_text(combined_content) or ""
|
||||
content_hash = hashlib.sha256(combined_content.encode()).hexdigest()
|
||||
|
||||
# Delete old document first (cascades to units and links).
|
||||
# Only delete on the first batch to avoid deleting data we just inserted.
|
||||
# Before the cascade, fan out to delete observations derived from the
|
||||
# outgoing memory_units — otherwise the FK ON DELETE CASCADE removes the
|
||||
# source memory_units but leaves observation rows pointing at IDs that
|
||||
# no longer exist (consolidated_at on co-source memories also stays
|
||||
# frozen). Same cleanup the explicit ``delete_document`` API performs.
|
||||
preserved_created_at = None
|
||||
# Delete old document first (cascades to units and links)
|
||||
# Only delete on the first batch to avoid deleting data we just inserted
|
||||
if is_first_batch:
|
||||
existing_unit_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT id FROM {fq_table("memory_units")}
|
||||
WHERE document_id = $1 AND fact_type IN ('experience', 'world')
|
||||
""",
|
||||
document_id,
|
||||
)
|
||||
existing_unit_ids = [row["id"] for row in existing_unit_rows]
|
||||
if existing_unit_ids:
|
||||
invalidated = await delete_stale_observations_for_memories(conn, bank_id, existing_unit_ids)
|
||||
if invalidated:
|
||||
logger.info(
|
||||
f"[RETAIN] Document {document_id} re-ingested: invalidated "
|
||||
f"{invalidated} observation(s) derived from {len(existing_unit_ids)} outgoing memory_units"
|
||||
)
|
||||
# Explicitly delete memory_units by document_id BEFORE deleting the
|
||||
# document row. The CASCADE from documents→chunks→memory_units only
|
||||
# catches units that have a non-NULL chunk_id FK. Units with chunk_id=NULL
|
||||
# (e.g. from partial writes or edge cases) would survive the cascade.
|
||||
# This explicit delete ensures complete cleanup.
|
||||
await conn.execute(
|
||||
f"DELETE FROM {fq_table('memory_units')} WHERE document_id = $1 AND bank_id = $2",
|
||||
document_id,
|
||||
bank_id,
|
||||
)
|
||||
# Capture created_at before deletion so re-ingestion preserves it.
|
||||
preserved_created_at = await conn.fetchval(
|
||||
f"DELETE FROM {fq_table('documents')} WHERE id = $1 AND bank_id = $2 RETURNING created_at",
|
||||
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 _upsert_document_row(
|
||||
conn,
|
||||
bank_id,
|
||||
document_id,
|
||||
combined_content,
|
||||
content_hash,
|
||||
retain_params,
|
||||
document_tags,
|
||||
preserved_created_at=preserved_created_at,
|
||||
)
|
||||
await _upsert_document_row(conn, bank_id, document_id, combined_content, content_hash, retain_params, document_tags)
|
||||
|
||||
|
||||
async def upsert_document_metadata(
|
||||
@@ -418,19 +297,12 @@ async def _upsert_document_row(
|
||||
content_hash: str,
|
||||
retain_params: dict | None = None,
|
||||
document_tags: list[str] | None = None,
|
||||
preserved_created_at: datetime | None = None,
|
||||
) -> None:
|
||||
"""Insert or update a document row.
|
||||
|
||||
When ``preserved_created_at`` is provided, it is used for ``created_at`` on
|
||||
INSERT so that re-ingesting a document (which deletes + inserts the row)
|
||||
keeps the original creation timestamp. ``updated_at`` is always set to
|
||||
``NOW()`` on both INSERT and the ON CONFLICT UPDATE branch.
|
||||
"""
|
||||
"""Insert or update a document row."""
|
||||
await conn.execute(
|
||||
f"""
|
||||
INSERT INTO {fq_table("documents")} (id, bank_id, original_text, content_hash, retain_params, tags, created_at, updated_at)
|
||||
VALUES ($1, $2, $3, $4, $5, $6, COALESCE($7, NOW()), NOW())
|
||||
INSERT INTO {fq_table("documents")} (id, bank_id, original_text, content_hash, retain_params, tags)
|
||||
VALUES ($1, $2, $3, $4, $5, $6)
|
||||
ON CONFLICT (id, bank_id) DO UPDATE
|
||||
SET original_text = EXCLUDED.original_text,
|
||||
content_hash = EXCLUDED.content_hash,
|
||||
@@ -444,7 +316,6 @@ async def _upsert_document_row(
|
||||
content_hash,
|
||||
json.dumps(retain_params) if retain_params else None,
|
||||
document_tags or [],
|
||||
preserved_created_at,
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -812,6 +812,7 @@ async def compute_semantic_links_ann(
|
||||
bank_id,
|
||||
fact_type,
|
||||
top_k,
|
||||
timeout=300, # ANN on large banks can take minutes
|
||||
)
|
||||
logger.debug(f"[ANN] fact_type={fact_type}: {len(ft_rows)} rows in {time_mod.time() - t_query:.3f}s")
|
||||
rows.extend(ft_rows)
|
||||
|
||||
@@ -16,7 +16,7 @@ from typing import Any
|
||||
|
||||
from ...worker.stage import set_stage
|
||||
from ..db_utils import acquire_with_retry
|
||||
from ..memory_engine import count_tokens, fq_table
|
||||
from ..memory_engine import fq_table
|
||||
from . import bank_utils
|
||||
|
||||
|
||||
@@ -25,32 +25,6 @@ def utcnow():
|
||||
return datetime.now(UTC)
|
||||
|
||||
|
||||
def _merge_processed_content_tokens(a: int | None, b: int | None) -> int | None:
|
||||
"""Combine the processed-content-tokens signal across sub-results.
|
||||
|
||||
Semantics (see RetainResult.processed_content_tokens):
|
||||
* None means "this part of the retain did not go through chunk-level
|
||||
dedup" — i.e. the entire submitted payload was processed. If any
|
||||
sub-result is None, the aggregate is None so callers conservatively
|
||||
bill the full content.
|
||||
* Otherwise, accumulate the int values.
|
||||
"""
|
||||
if a is None or b is None:
|
||||
return None
|
||||
return a + b
|
||||
|
||||
|
||||
def _count_delta_content_tokens(delta_contents: list["RetainContent"]) -> int:
|
||||
"""Sum content + context tokens across the chunk items that were
|
||||
actually fed into the extraction pipeline on a partial-delta retain.
|
||||
"""
|
||||
total = 0
|
||||
for c in delta_contents:
|
||||
total += count_tokens(c.content or "")
|
||||
total += count_tokens(c.context or "")
|
||||
return total
|
||||
|
||||
|
||||
def parse_datetime_flexible(value: Any) -> datetime:
|
||||
"""
|
||||
Parse a datetime value that could be either a datetime object or an ISO string.
|
||||
@@ -98,6 +72,7 @@ from . import (
|
||||
from .types import (
|
||||
ChunkMetadata,
|
||||
EntityResolutionResult,
|
||||
ExtractedFact,
|
||||
Phase1Result,
|
||||
Phase3Context,
|
||||
ProcessedFact,
|
||||
@@ -327,11 +302,8 @@ async def _insert_facts_and_links(
|
||||
causal_link_count = await link_creation.create_causal_links_batch(conn, bank_id, unit_ids, processed_facts)
|
||||
log_buffer.append(f" Causal links: {causal_link_count} links in {time.time() - step_start:.3f}s")
|
||||
|
||||
# Map results back to original content items. Use processed_facts (not
|
||||
# extracted_facts) because unit_ids has 1:1 alignment with processed_facts —
|
||||
# any upstream drop between extraction and processing would otherwise cause
|
||||
# an IndexError (see issue #1037).
|
||||
result_unit_ids = _map_results_to_contents(contents, processed_facts, unit_ids if unit_ids else [])
|
||||
# Map results back to original content items
|
||||
result_unit_ids = _map_results_to_contents(contents, extracted_facts, unit_ids if unit_ids else [])
|
||||
|
||||
if outbox_callback:
|
||||
await outbox_callback(conn)
|
||||
@@ -443,21 +415,13 @@ async def retain_batch(
|
||||
schema: str | None = None,
|
||||
outbox_callback: Callable[["asyncpg.Connection"], Awaitable[None]] | None = None,
|
||||
db_semaphore: "asyncio.Semaphore | None" = None,
|
||||
) -> tuple[list[list[str]], TokenUsage, int | None]:
|
||||
) -> tuple[list[list[str]], TokenUsage]:
|
||||
"""
|
||||
Process a batch of content through the retain pipeline.
|
||||
|
||||
Supports delta retain: when upserting a document that already has chunks,
|
||||
only re-processes chunks whose content has changed. Unchanged chunks keep
|
||||
their existing facts, entities, and links.
|
||||
|
||||
Returns a three-tuple of:
|
||||
* per-content-item unit ID lists
|
||||
* aggregate LLM token usage
|
||||
* processed_content_tokens — content+context tokens that actually went
|
||||
through extraction after chunk-level dedup, or ``None`` if this path
|
||||
didn't dedup (caller should treat as "bill full submitted content").
|
||||
See ``RetainResult.processed_content_tokens`` for details.
|
||||
"""
|
||||
start_time = time.time()
|
||||
total_chars = sum(len(item.get("content", "")) for item in contents_dicts)
|
||||
@@ -497,9 +461,8 @@ async def retain_batch(
|
||||
# Process each group and merge results back in original order
|
||||
result_unit_ids: list[list[str]] = [[] for _ in contents_dicts]
|
||||
total_usage = TokenUsage()
|
||||
total_processed_tokens: int | None = 0
|
||||
for doc_key, (group_dicts, group_contents) in groups.items():
|
||||
group_ids, group_usage, group_processed = await retain_batch(
|
||||
group_ids, group_usage = await retain_batch(
|
||||
pool=pool,
|
||||
embeddings_model=embeddings_model,
|
||||
llm_config=llm_config,
|
||||
@@ -521,12 +484,11 @@ async def retain_batch(
|
||||
if group_idx < len(group_ids):
|
||||
result_unit_ids[orig_idx] = group_ids[group_idx]
|
||||
total_usage = total_usage + group_usage
|
||||
total_processed_tokens = _merge_processed_content_tokens(total_processed_tokens, group_processed)
|
||||
return result_unit_ids, total_usage, total_processed_tokens
|
||||
return result_unit_ids, total_usage
|
||||
|
||||
# Resolve effective document_id early so both delta and streaming paths
|
||||
# can find existing chunks from a prior attempt. On retry, a generated
|
||||
# document_id is recovered from operation result_metadata.document_ids[0].
|
||||
# can find existing chunks from a prior attempt. On retry, the generated
|
||||
# document_id is recovered from operation result_metadata.
|
||||
effective_doc_id = document_id
|
||||
if not effective_doc_id:
|
||||
doc_ids = {item.get("document_id") for item in contents_dicts if item.get("document_id")}
|
||||
@@ -545,41 +507,26 @@ async def retain_batch(
|
||||
if isinstance(row["result_metadata"], dict)
|
||||
else json.loads(row["result_metadata"])
|
||||
)
|
||||
recovered = meta.get("document_ids") or []
|
||||
if recovered:
|
||||
effective_doc_id = recovered[0]
|
||||
effective_doc_id = meta.get("generated_document_id")
|
||||
except Exception:
|
||||
pass
|
||||
if not effective_doc_id:
|
||||
effective_doc_id = str(uuid.uuid4())
|
||||
|
||||
# Record effective_doc_id on the operation (idempotent set-append). Captures
|
||||
# both user-provided and generated ids so the operation shows every document
|
||||
# it touched, and lets retries reuse the same generated id.
|
||||
if operation_id:
|
||||
try:
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
await conn.execute(
|
||||
f"""
|
||||
UPDATE {fq_table("async_operations")}
|
||||
SET result_metadata = jsonb_set(
|
||||
COALESCE(result_metadata, '{{}}'::jsonb),
|
||||
'{{document_ids}}',
|
||||
CASE
|
||||
WHEN COALESCE(result_metadata->'document_ids', '[]'::jsonb) @> $1::jsonb
|
||||
THEN result_metadata->'document_ids'
|
||||
ELSE COALESCE(result_metadata->'document_ids', '[]'::jsonb) || $1::jsonb
|
||||
END,
|
||||
true
|
||||
),
|
||||
updated_at = now()
|
||||
WHERE operation_id = $2
|
||||
""",
|
||||
json.dumps([effective_doc_id]),
|
||||
uuid.UUID(operation_id),
|
||||
)
|
||||
except Exception:
|
||||
logger.warning("Failed to persist document_id", exc_info=True)
|
||||
# Persist so retries reuse the same document_id
|
||||
if operation_id:
|
||||
try:
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
await conn.execute(
|
||||
f"""
|
||||
UPDATE {fq_table("async_operations")}
|
||||
SET result_metadata = result_metadata || $1::jsonb, updated_at = now()
|
||||
WHERE operation_id = $2
|
||||
""",
|
||||
json.dumps({"generated_document_id": effective_doc_id}),
|
||||
uuid.UUID(operation_id),
|
||||
)
|
||||
except Exception:
|
||||
logger.warning("Failed to persist generated document_id", exc_info=True)
|
||||
|
||||
# --- Append mode: prepend existing document content to new content ---
|
||||
# When update_mode="append", fetch the existing document text and prepend it
|
||||
@@ -610,31 +557,6 @@ async def retain_batch(
|
||||
f"[append] Prepended {len(existing_text):,} chars from existing document {effective_doc_id}"
|
||||
)
|
||||
|
||||
# --- Stale-request check (best-effort, before LLM extraction) ---
|
||||
# If the document was already updated by a more recent retain (updated_at > our
|
||||
# start_time), skip this request entirely to avoid overwriting newer content
|
||||
# (e.g. a longer conversation) with older data. This is an optimization — the
|
||||
# real correctness guarantee comes from the FOR UPDATE + content_hash check
|
||||
# inside each batch TXN (see _run_mini_batch_db_work).
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
doc_row = await conn.fetchrow(
|
||||
f"SELECT updated_at FROM {fq_table('documents')} WHERE id = $1 AND bank_id = $2",
|
||||
effective_doc_id,
|
||||
bank_id,
|
||||
)
|
||||
if doc_row and doc_row["updated_at"]:
|
||||
doc_updated = doc_row["updated_at"].timestamp()
|
||||
if doc_updated > start_time:
|
||||
log_buffer.append(
|
||||
f"[stale] Skipping retain: document {effective_doc_id} was updated at "
|
||||
f"{doc_row['updated_at'].isoformat()} (after this request started at "
|
||||
f"{datetime.fromtimestamp(start_time, tz=UTC).isoformat()})"
|
||||
)
|
||||
logger.info("\n" + "\n".join(log_buffer) + "\n")
|
||||
# No new content was processed — report 0 so callers can skip
|
||||
# billing cleanly instead of falling back to full-content billing.
|
||||
return [[] for _ in contents], TokenUsage(), 0
|
||||
|
||||
# --- Delta retain: check if we can skip unchanged chunks ---
|
||||
if is_first_batch:
|
||||
delta_result = await _try_delta_retain(
|
||||
@@ -791,6 +713,7 @@ async def _run_final_semantic_ann(
|
||||
async with ann_semaphore:
|
||||
t0 = time.time()
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
await conn.execute("SET statement_timeout = '300s'")
|
||||
ann_links = await compute_semantic_links_ann(
|
||||
conn,
|
||||
bank_id,
|
||||
@@ -803,6 +726,7 @@ async def _run_final_semantic_ann(
|
||||
if ann_links:
|
||||
await _bulk_insert_links(conn, ann_links, bank_id=bank_id)
|
||||
chunk_link_counts[chunk_idx] = len(ann_links)
|
||||
await conn.execute("RESET statement_timeout")
|
||||
logger.info(
|
||||
f"[streaming] Final ANN chunk {chunk_idx + 1}/{num_chunks}: "
|
||||
f"{len(ann_links)} links in {time.time() - t0:.3f}s"
|
||||
@@ -868,27 +792,25 @@ async def _streaming_retain_batch(
|
||||
# Default template for metadata (context, event_date, etc.) when content list is empty.
|
||||
_default_content = RetainContent(content="")
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Recovery detection (read-only, before LLM extraction)
|
||||
# ---------------------------------------------------------------------------
|
||||
# Check if this is a retry of the same content (crash recovery). If the
|
||||
# document exists with a matching content_hash and has committed chunks,
|
||||
# the producer can skip already-extracted chunks to avoid duplicate work.
|
||||
# Load existing chunk hashes BEFORE document tracking to detect recovery.
|
||||
# If chunks exist AND the document content hash matches, this is a retry of
|
||||
# the same content — preserve existing data. If content differs, this is an
|
||||
# update — cascade-delete old data and start fresh.
|
||||
existing_chunk_hashes: set[str] = set()
|
||||
combined_content = "\n".join([c.get("content", "") for c in contents_dicts])
|
||||
# Sanitize before hashing to match what handle_document_tracking stores
|
||||
sanitized_content = fact_extraction._sanitize_text(combined_content) or ""
|
||||
new_content_hash = hashlib.sha256(sanitized_content.encode()).hexdigest()
|
||||
new_content_hash = hashlib.sha256(combined_content.encode()).hexdigest()
|
||||
is_recovery = False
|
||||
|
||||
try:
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
# Check if document exists with matching content hash
|
||||
doc_row = await conn.fetchrow(
|
||||
f"SELECT content_hash FROM {fq_table('documents')} WHERE id = $1 AND bank_id = $2",
|
||||
effective_doc_id,
|
||||
bank_id,
|
||||
)
|
||||
if doc_row and doc_row["content_hash"] == new_content_hash:
|
||||
# Same content — load chunk hashes for recovery skip
|
||||
existing_rows = await chunk_storage.load_existing_chunks(conn, bank_id, effective_doc_id)
|
||||
existing_chunk_hashes = {c.content_hash for c in existing_rows if c.content_hash}
|
||||
if existing_chunk_hashes:
|
||||
@@ -900,22 +822,24 @@ async def _streaming_retain_batch(
|
||||
except Exception:
|
||||
pass # If we can't load, just process all chunks
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Document tracking is DEFERRED to the first consumer batch TXN.
|
||||
# ---------------------------------------------------------------------------
|
||||
# Previously, document tracking (cascade-delete old data + insert doc row)
|
||||
# ran in a separate transaction BEFORE LLM extraction. This left a gap
|
||||
# between the cascade-delete and the first chunk write, allowing concurrent
|
||||
# requests to interleave and produce duplicates.
|
||||
#
|
||||
# Now, document tracking runs atomically inside the first batch's write TXN,
|
||||
# using SELECT ... FOR UPDATE on the document row for serialization across
|
||||
# workers. Each batch TXN also verifies document ownership via content_hash
|
||||
# to detect when a concurrent request has taken over the document.
|
||||
# See _run_mini_batch_db_work() for the implementation.
|
||||
# Create/update the document row.
|
||||
retain_params, merged_tags = _build_retain_params(contents_dicts, document_tags)
|
||||
# Track whether document tracking has been done (by the first batch)
|
||||
doc_tracking_done = [False]
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
async with conn.transaction():
|
||||
if is_recovery:
|
||||
# Recovery: same content, partially committed — preserve existing data
|
||||
await fact_storage.upsert_document_metadata(
|
||||
conn, bank_id, effective_doc_id, combined_content, retain_params, merged_tags
|
||||
)
|
||||
log_buffer.append(
|
||||
f"[streaming] Document {effective_doc_id} updated (recovery, preserving existing chunks)"
|
||||
)
|
||||
else:
|
||||
# Fresh or update: cascade-delete old data if document exists
|
||||
await fact_storage.handle_document_tracking(
|
||||
conn, bank_id, effective_doc_id, combined_content, is_first_batch, retain_params, merged_tags
|
||||
)
|
||||
log_buffer.append(f"[streaming] Document {effective_doc_id} tracked (full content)")
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Producer-consumer pipeline: LLM extraction runs concurrently with DB writes
|
||||
@@ -928,10 +852,6 @@ async def _streaming_retain_batch(
|
||||
|
||||
# Shared mutable state for the producer to report skipped chunks and usage
|
||||
producer_error: list[BaseException] = []
|
||||
# Set to True by _run_mini_batch_db_work when a concurrent request takes
|
||||
# over the document (content_hash mismatch). The consumer checks this and
|
||||
# stops processing further batches.
|
||||
pipeline_aborted: list[bool] = [False]
|
||||
|
||||
# ---- LLM Producer ----
|
||||
# Fires all chunk extractions as concurrent tasks (bounded by the LLM
|
||||
@@ -990,15 +910,17 @@ async def _streaming_retain_batch(
|
||||
# Phase 1 (entity resolution) -> Phase 2 (write txn) -> Phase 3 (ANN fire-and-forget).
|
||||
async def _db_consumer() -> None:
|
||||
batch: list[tuple] = []
|
||||
global_chunk_offset = 0
|
||||
consumer_batch_idx = 0
|
||||
|
||||
while True:
|
||||
item = await chunk_queue.get()
|
||||
if item is None:
|
||||
# Process any remaining items
|
||||
if batch and not pipeline_aborted[0]:
|
||||
if batch:
|
||||
await _process_db_batch(
|
||||
batch,
|
||||
global_chunk_offset,
|
||||
consumer_batch_idx,
|
||||
is_last=True,
|
||||
)
|
||||
@@ -1007,24 +929,19 @@ async def _streaming_retain_batch(
|
||||
batch.append(item)
|
||||
|
||||
if len(batch) >= chunk_batch_size:
|
||||
if pipeline_aborted[0]:
|
||||
# Another request took over the document — discard this batch
|
||||
log_buffer.append(
|
||||
f"[streaming] Consumer: discarding batch of {len(batch)} chunks "
|
||||
f"(pipeline aborted due to concurrent takeover)"
|
||||
)
|
||||
batch = []
|
||||
continue
|
||||
await _process_db_batch(
|
||||
batch,
|
||||
global_chunk_offset,
|
||||
consumer_batch_idx,
|
||||
is_last=False,
|
||||
)
|
||||
global_chunk_offset += len(batch)
|
||||
consumer_batch_idx += 1
|
||||
batch = []
|
||||
|
||||
async def _process_db_batch(
|
||||
batch: list[tuple],
|
||||
global_chunk_offset: int,
|
||||
consumer_batch_idx: int,
|
||||
is_last: bool,
|
||||
) -> None:
|
||||
@@ -1038,17 +955,15 @@ async def _streaming_retain_batch(
|
||||
|
||||
for global_idx, content, extracted, processed, chunk_meta, usage in batch:
|
||||
content_idx_in_batch = len(batch_contents)
|
||||
# Adjust chunk indices to use the original global position (global_idx)
|
||||
# so that chunk_id = {bank}_{doc}_{chunk_index} is deterministic regardless
|
||||
# of task completion order. content_index is batch-relative for result grouping.
|
||||
# Adjust chunk indices to global offsets and remap content_index
|
||||
for fact in extracted:
|
||||
fact.content_index = content_idx_in_batch
|
||||
if fact.chunk_index is not None:
|
||||
fact.chunk_index = global_idx
|
||||
fact.chunk_index = global_chunk_offset + content_idx_in_batch
|
||||
for pf in processed:
|
||||
pf.content_index = content_idx_in_batch
|
||||
for cm in chunk_meta:
|
||||
cm.chunk_index = global_idx
|
||||
cm.chunk_index = global_chunk_offset + content_idx_in_batch
|
||||
|
||||
batch_contents.append(content)
|
||||
batch_extracted.extend(extracted)
|
||||
@@ -1060,46 +975,6 @@ async def _streaming_retain_batch(
|
||||
total_usage = total_usage + batch_usage
|
||||
|
||||
if not batch_extracted:
|
||||
# Even with 0 facts, the first batch must still run document tracking
|
||||
# (cascade-delete + insert doc row) to establish ownership and prevent
|
||||
# concurrent requests from interleaving. Later batches can safely skip.
|
||||
if not doc_tracking_done[0]:
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
async with conn.transaction():
|
||||
await conn.execute(
|
||||
f"INSERT INTO {fq_table('documents')} (id, bank_id, original_text, content_hash) "
|
||||
f"VALUES ($1, $2, '', '__pending__') "
|
||||
f"ON CONFLICT (id, bank_id) DO NOTHING",
|
||||
effective_doc_id,
|
||||
bank_id,
|
||||
)
|
||||
await conn.fetchval(
|
||||
f"SELECT content_hash FROM {fq_table('documents')} "
|
||||
f"WHERE id = $1 AND bank_id = $2 FOR UPDATE",
|
||||
effective_doc_id,
|
||||
bank_id,
|
||||
)
|
||||
if is_recovery:
|
||||
await fact_storage.upsert_document_metadata(
|
||||
conn,
|
||||
bank_id,
|
||||
effective_doc_id,
|
||||
combined_content,
|
||||
retain_params,
|
||||
merged_tags,
|
||||
)
|
||||
else:
|
||||
await fact_storage.handle_document_tracking(
|
||||
conn,
|
||||
bank_id,
|
||||
effective_doc_id,
|
||||
combined_content,
|
||||
is_first_batch,
|
||||
retain_params,
|
||||
merged_tags,
|
||||
)
|
||||
doc_tracking_done[0] = True
|
||||
log_buffer.append(f"[streaming] Document {effective_doc_id} tracked (0 facts in first batch)")
|
||||
log_buffer.append(
|
||||
f"[streaming] Consumer batch {consumer_batch_idx + 1}: "
|
||||
f"0 facts extracted from {len(batch)} chunks, skipping"
|
||||
@@ -1130,92 +1005,10 @@ async def _streaming_retain_batch(
|
||||
|
||||
logger.info(f"[streaming] Phase 1 (entity resolution): {time.time() - p1_start:.3f}s")
|
||||
|
||||
# Phase 2 — Write transaction
|
||||
# -----------------------------------------------------------------
|
||||
# Concurrent-safety via row-level locking:
|
||||
#
|
||||
# The streaming pipeline splits work across multiple batch TXNs.
|
||||
# Without protection, two concurrent retains for the same document
|
||||
# can interleave: Request A writes batch1, Request B cascade-deletes
|
||||
# A's doc and writes its own batch1, then A's batch2 adds stale data
|
||||
# on top of B's → duplicates.
|
||||
#
|
||||
# To prevent this, every batch TXN:
|
||||
# 1. SELECT ... FOR UPDATE on the document row — serializes all
|
||||
# writers for this document at the DB level (works across workers).
|
||||
# 2. Check content_hash — if it doesn't match ours, another request
|
||||
# took over the document → abort remaining batches.
|
||||
# 3. First batch only: run handle_document_tracking (cascade-delete
|
||||
# old data + insert doc row) atomically with the first chunk write.
|
||||
# This eliminates the gap between "delete old" and "insert new"
|
||||
# that previously allowed interleaving.
|
||||
# -----------------------------------------------------------------
|
||||
|
||||
# Phase 2 — Write transaction (within-batch semantic links only)
|
||||
p2_start = time.time()
|
||||
batch_result_ids = None
|
||||
phase3_ctx = None
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
async with conn.transaction():
|
||||
# --- Document ownership gate ---
|
||||
# Lock the document row to serialize all concurrent writers.
|
||||
# SELECT ... FOR UPDATE doesn't lock non-existent rows, so we
|
||||
# first ensure the row exists with a lightweight upsert, THEN lock it.
|
||||
# The content_hash='__pending__' placeholder is immediately overwritten
|
||||
# by handle_document_tracking or upsert_document_metadata below.
|
||||
await conn.execute(
|
||||
f"INSERT INTO {fq_table('documents')} (id, bank_id, original_text, content_hash) "
|
||||
f"VALUES ($1, $2, '', '__pending__') "
|
||||
f"ON CONFLICT (id, bank_id) DO NOTHING",
|
||||
effective_doc_id,
|
||||
bank_id,
|
||||
)
|
||||
existing_hash = await conn.fetchval(
|
||||
f"SELECT content_hash FROM {fq_table('documents')} WHERE id = $1 AND bank_id = $2 FOR UPDATE",
|
||||
effective_doc_id,
|
||||
bank_id,
|
||||
)
|
||||
|
||||
if not doc_tracking_done[0]:
|
||||
# --- First batch: document tracking (atomic with chunk write) ---
|
||||
if is_recovery:
|
||||
await fact_storage.upsert_document_metadata(
|
||||
conn,
|
||||
bank_id,
|
||||
effective_doc_id,
|
||||
combined_content,
|
||||
retain_params,
|
||||
merged_tags,
|
||||
)
|
||||
log_buffer.append(
|
||||
f"[streaming] Document {effective_doc_id} updated "
|
||||
f"(recovery, preserving existing chunks)"
|
||||
)
|
||||
else:
|
||||
await fact_storage.handle_document_tracking(
|
||||
conn,
|
||||
bank_id,
|
||||
effective_doc_id,
|
||||
combined_content,
|
||||
is_first_batch,
|
||||
retain_params,
|
||||
merged_tags,
|
||||
)
|
||||
log_buffer.append(f"[streaming] Document {effective_doc_id} tracked (full content)")
|
||||
doc_tracking_done[0] = True
|
||||
else:
|
||||
# --- Later batches: verify we still own the document ---
|
||||
# If another request took over (cascade-deleted our doc and
|
||||
# inserted its own), the content_hash won't match ours.
|
||||
if existing_hash is not None and existing_hash != new_content_hash:
|
||||
log_buffer.append(
|
||||
f"[streaming] Document {effective_doc_id} taken over by "
|
||||
f"concurrent request (hash mismatch) — aborting remaining batches"
|
||||
)
|
||||
logger.info("\n" + "\n".join(log_buffer) + "\n")
|
||||
# Signal the consumer to stop processing further batches
|
||||
pipeline_aborted[0] = True
|
||||
return
|
||||
|
||||
# Store chunks with correct global indices
|
||||
step_start = time.time()
|
||||
chunk_id_map = {}
|
||||
@@ -1257,14 +1050,11 @@ async def _streaming_retain_batch(
|
||||
logger.info(f"[streaming] Phase 2 (write txn): {time.time() - p2_start:.3f}s")
|
||||
|
||||
# Best-effort: entity viz + stats (fast, not semantic ANN)
|
||||
if phase3_ctx is not None:
|
||||
try:
|
||||
await entity_resolver.flush_pending_stats()
|
||||
await _build_and_insert_entity_links_phase3(
|
||||
pool, entity_resolver, bank_id, phase3_ctx, log_buffer
|
||||
)
|
||||
except Exception:
|
||||
logger.warning(f"Phase 3 stats (consumer batch {consumer_batch_idx + 1}) failed", exc_info=True)
|
||||
try:
|
||||
await entity_resolver.flush_pending_stats()
|
||||
await _build_and_insert_entity_links_phase3(pool, entity_resolver, bank_id, phase3_ctx, log_buffer)
|
||||
except Exception:
|
||||
logger.warning(f"Phase 3 stats (consumer batch {consumer_batch_idx + 1}) failed", exc_info=True)
|
||||
|
||||
logger.info(
|
||||
f"[streaming] Consumer batch {consumer_batch_idx + 1} total "
|
||||
@@ -1272,9 +1062,8 @@ async def _streaming_retain_batch(
|
||||
)
|
||||
|
||||
# Collect unit_ids from this batch
|
||||
if batch_result_ids:
|
||||
for content_ids in batch_result_ids:
|
||||
all_unit_ids.extend(content_ids)
|
||||
for content_ids in batch_result_ids:
|
||||
all_unit_ids.extend(content_ids)
|
||||
|
||||
if db_semaphore is not None:
|
||||
async with db_semaphore:
|
||||
@@ -1317,47 +1106,6 @@ async def _streaming_retain_batch(
|
||||
if producer_error:
|
||||
raise producer_error[0]
|
||||
|
||||
# If no batch was processed (e.g. zero facts extracted from gibberish
|
||||
# content, or all chunks skipped in recovery), the document row was
|
||||
# never created by the first batch TXN. Create it now so the document
|
||||
# is tracked regardless of extraction results.
|
||||
if not doc_tracking_done[0] and not pipeline_aborted[0]:
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
async with conn.transaction():
|
||||
await conn.execute(
|
||||
f"INSERT INTO {fq_table('documents')} (id, bank_id, original_text, content_hash) "
|
||||
f"VALUES ($1, $2, '', '__pending__') "
|
||||
f"ON CONFLICT (id, bank_id) DO NOTHING",
|
||||
effective_doc_id,
|
||||
bank_id,
|
||||
)
|
||||
await conn.fetchval(
|
||||
f"SELECT content_hash FROM {fq_table('documents')} WHERE id = $1 AND bank_id = $2 FOR UPDATE",
|
||||
effective_doc_id,
|
||||
bank_id,
|
||||
)
|
||||
if is_recovery:
|
||||
await fact_storage.upsert_document_metadata(
|
||||
conn,
|
||||
bank_id,
|
||||
effective_doc_id,
|
||||
combined_content,
|
||||
retain_params,
|
||||
merged_tags,
|
||||
)
|
||||
else:
|
||||
await fact_storage.handle_document_tracking(
|
||||
conn,
|
||||
bank_id,
|
||||
effective_doc_id,
|
||||
combined_content,
|
||||
is_first_batch,
|
||||
retain_params,
|
||||
merged_tags,
|
||||
)
|
||||
doc_tracking_done[0] = True
|
||||
log_buffer.append(f"[streaming] Document {effective_doc_id} tracked (no facts extracted)")
|
||||
|
||||
# Mark facts as committed in operation metadata (crash recovery checkpoint)
|
||||
if operation_id and all_unit_ids:
|
||||
try:
|
||||
@@ -1394,31 +1142,16 @@ async def _streaming_retain_batch(
|
||||
# This replaces per-batch within-batch + fire-and-forget ANN with a single
|
||||
# efficient pass after all facts are in the database.
|
||||
# ---------------------------------------------------------------------------
|
||||
if all_unit_ids and not pipeline_aborted[0]:
|
||||
if all_unit_ids:
|
||||
ann_start = time.time()
|
||||
try:
|
||||
await _run_final_semantic_ann(pool, bank_id, all_unit_ids, log_buffer)
|
||||
except Exception:
|
||||
# ANN pass is best-effort. FK violations can occur if a concurrent
|
||||
# retain cascade-deleted our units between the batch commit and here.
|
||||
logger.warning(
|
||||
f"[streaming] Final ANN pass failed for document {effective_doc_id} "
|
||||
f"(units may have been superseded by concurrent retain)",
|
||||
exc_info=True,
|
||||
)
|
||||
await _run_final_semantic_ann(pool, bank_id, all_unit_ids, log_buffer)
|
||||
log_buffer.append(f"[streaming] Final ANN pass: {time.time() - ann_start:.3f}s for {len(all_unit_ids)} units")
|
||||
|
||||
total_time = time.time() - start_time
|
||||
log_buffer.append(f"{'=' * 60}")
|
||||
if pipeline_aborted[0]:
|
||||
log_buffer.append(
|
||||
f"STREAMING RETAIN ABORTED: document {effective_doc_id} was taken over by "
|
||||
f"a concurrent request after {total_time:.3f}s — data from this request was discarded"
|
||||
)
|
||||
else:
|
||||
log_buffer.append(
|
||||
f"STREAMING RETAIN COMPLETE: {len(all_unit_ids)} units across {num_batches} batches in {total_time:.3f}s"
|
||||
)
|
||||
log_buffer.append(
|
||||
f"STREAMING RETAIN COMPLETE: {len(all_unit_ids)} units across {num_batches} batches in {total_time:.3f}s"
|
||||
)
|
||||
log_buffer.append(f"Document: {effective_doc_id}")
|
||||
log_buffer.append(f"{'=' * 60}")
|
||||
logger.info("\n" + "\n".join(log_buffer) + "\n")
|
||||
@@ -1426,10 +1159,7 @@ async def _streaming_retain_batch(
|
||||
# Map all unit_ids back to the original content items.
|
||||
# For streaming mode with a single document, all units belong to content 0.
|
||||
result_unit_ids = [all_unit_ids] + [[] for _ in contents[1:]]
|
||||
# The streaming path doesn't compute per-chunk content-hash dedup in
|
||||
# a way that lets us report a partial-processed tokens count — signal
|
||||
# ``None`` so callers bill against the full submitted payload.
|
||||
return result_unit_ids, total_usage, None
|
||||
return result_unit_ids, total_usage
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -1457,15 +1187,10 @@ async def _try_delta_retain(
|
||||
schema,
|
||||
outbox_callback,
|
||||
db_semaphore: "asyncio.Semaphore | None" = None,
|
||||
) -> tuple[list[list[str]], TokenUsage, int | None] | None:
|
||||
):
|
||||
"""
|
||||
Attempt delta retain for a document upsert. Returns result tuple if delta
|
||||
was performed, or None to fall back to full retain.
|
||||
|
||||
When a result tuple is returned, the third element is the content+context
|
||||
token count for the chunks that actually went through extraction
|
||||
(``0`` if the submission matched prior content exactly and nothing was
|
||||
re-extracted).
|
||||
"""
|
||||
# Need a single document_id
|
||||
effective_doc_id = document_id
|
||||
@@ -1475,17 +1200,9 @@ async def _try_delta_retain(
|
||||
return None
|
||||
effective_doc_id = doc_ids.pop()
|
||||
|
||||
# Load existing chunks and snapshot the document's content_hash. This is
|
||||
# outside the write TXN, so a concurrent retain could modify the document
|
||||
# between this read and the write. The write TXN verifies the hash hasn't
|
||||
# changed; if it has, we fall back to streaming (which has full protection).
|
||||
# Load existing chunks
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
existing_chunks = await chunk_storage.load_existing_chunks(conn, bank_id, effective_doc_id)
|
||||
doc_hash_at_load = await conn.fetchval(
|
||||
f"SELECT content_hash FROM {fq_table('documents')} WHERE id = $1 AND bank_id = $2",
|
||||
effective_doc_id,
|
||||
bank_id,
|
||||
)
|
||||
|
||||
if not existing_chunks:
|
||||
return None
|
||||
@@ -1593,28 +1310,8 @@ async def _try_delta_retain(
|
||||
)
|
||||
|
||||
# PHASE 2 — Core Write Transaction (atomic)
|
||||
# Lock the document row and verify ownership. Delta loaded existing
|
||||
# chunks OUTSIDE this TXN, so a concurrent retain may have cascade-deleted
|
||||
# and replaced the document since then. If the content_hash changed,
|
||||
# the chunk state we based our delta diff on is stale — abort.
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
async with conn.transaction():
|
||||
current_hash = await conn.fetchval(
|
||||
f"SELECT content_hash FROM {fq_table('documents')} WHERE id = $1 AND bank_id = $2 FOR UPDATE",
|
||||
effective_doc_id,
|
||||
bank_id,
|
||||
)
|
||||
# Verify the document hasn't been replaced since we loaded chunks.
|
||||
# Compare the current hash against what we snapshotted at load time.
|
||||
if current_hash is not None and doc_hash_at_load is not None and current_hash != doc_hash_at_load:
|
||||
log_buffer.append(
|
||||
f"[delta] Document {effective_doc_id} was modified by concurrent request "
|
||||
f"since chunks were loaded — aborting delta, falling back to full retain"
|
||||
)
|
||||
logger.info("\n" + "\n".join(log_buffer) + "\n")
|
||||
# Return None to fall back to streaming (which has full FOR UPDATE protection)
|
||||
return None
|
||||
|
||||
# Update document metadata (no delete)
|
||||
step_start = time.time()
|
||||
combined_content = "\n".join([c.get("content", "") for c in contents_dicts])
|
||||
@@ -1724,12 +1421,7 @@ async def _try_delta_retain(
|
||||
await _run_delta_db_work()
|
||||
else:
|
||||
await _run_delta_db_work()
|
||||
# Count content + context tokens that actually went through extraction.
|
||||
# ``delta_contents`` holds the per-chunk RetainContent items for the
|
||||
# changed/new chunks (see ``_build_delta_contents``) — i.e. exactly what
|
||||
# the LLM pipeline saw this call. Unchanged chunks contribute zero.
|
||||
processed_tokens = _count_delta_content_tokens(delta_contents)
|
||||
return result_unit_ids, usage, processed_tokens
|
||||
return result_unit_ids, usage
|
||||
|
||||
|
||||
async def _delta_metadata_only(
|
||||
@@ -1746,12 +1438,6 @@ async def _delta_metadata_only(
|
||||
"""Handle the case where no chunks changed — just update document metadata and tags."""
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
async with conn.transaction():
|
||||
# Lock the document row to serialize with concurrent retains
|
||||
await conn.fetchval(
|
||||
f"SELECT content_hash FROM {fq_table('documents')} WHERE id = $1 AND bank_id = $2 FOR UPDATE",
|
||||
document_id,
|
||||
bank_id,
|
||||
)
|
||||
combined_content = "\n".join([c.get("content", "") for c in contents_dicts])
|
||||
retain_params, merged_tags = _build_retain_params(contents_dicts, document_tags)
|
||||
await fact_storage.upsert_document_metadata(
|
||||
@@ -1769,11 +1455,7 @@ async def _delta_metadata_only(
|
||||
total_time = time.time() - start_time
|
||||
log_buffer.append(f"DELTA RETAIN (no changes): metadata updated in {total_time:.3f}s")
|
||||
logger.info("\n" + "\n".join(log_buffer) + "\n")
|
||||
# Nothing went through the extraction pipeline — report 0 processed
|
||||
# content tokens so callers can bill accordingly (a caller that's been
|
||||
# told ``0`` knows the retain was a pure metadata update and should
|
||||
# charge nothing for content).
|
||||
return [[] for _ in contents], TokenUsage(), 0
|
||||
return [[] for _ in contents], TokenUsage()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -1868,19 +1550,12 @@ def _build_delta_contents(
|
||||
|
||||
def _map_results_to_contents(
|
||||
contents: list[RetainContent],
|
||||
processed_facts: list[ProcessedFact],
|
||||
extracted_facts: list[ExtractedFact],
|
||||
unit_ids: list[str],
|
||||
) -> list[list[str]]:
|
||||
"""Map created unit IDs back to original content items.
|
||||
|
||||
`processed_facts` and `unit_ids` must have the same length: each unit_id
|
||||
corresponds to the processed_fact at the same index.
|
||||
"""
|
||||
if len(processed_facts) != len(unit_ids):
|
||||
raise ValueError(f"processed_facts ({len(processed_facts)}) and unit_ids ({len(unit_ids)}) length mismatch")
|
||||
|
||||
"""Map created unit IDs back to original content items."""
|
||||
facts_by_content: dict[int, list[int]] = {i: [] for i in range(len(contents))}
|
||||
for i, fact in enumerate(processed_facts):
|
||||
for i, fact in enumerate(extracted_facts):
|
||||
# Normalize content_index: some LLM providers return 1-indexed values.
|
||||
# Clamp to valid range to prevent KeyError.
|
||||
idx = fact.content_index
|
||||
@@ -1889,8 +1564,12 @@ def _map_results_to_contents(
|
||||
facts_by_content[idx].append(i)
|
||||
|
||||
result_unit_ids = []
|
||||
unit_idx = 0
|
||||
for content_index in range(len(contents)):
|
||||
content_unit_ids = [unit_ids[i] for i in facts_by_content[content_index]]
|
||||
content_unit_ids = []
|
||||
for _ in facts_by_content[content_index]:
|
||||
content_unit_ids.append(unit_ids[unit_idx])
|
||||
unit_idx += 1
|
||||
result_unit_ids.append(content_unit_ids)
|
||||
|
||||
return result_unit_ids
|
||||
|
||||
@@ -8,7 +8,6 @@ of the recall pipeline.
|
||||
|
||||
import logging
|
||||
from abc import ABC, abstractmethod
|
||||
from datetime import datetime
|
||||
|
||||
from .tags import TagGroup, TagsMatch
|
||||
from .types import GraphRetrievalTimings, RetrievalResult
|
||||
@@ -46,8 +45,6 @@ class GraphRetriever(ABC):
|
||||
tags: list[str] | None = None, # Visibility scope tags for filtering
|
||||
tags_match: TagsMatch = "any", # How to match tags: 'any' (OR) or 'all' (AND)
|
||||
tag_groups: list[TagGroup] | None = None, # Compound boolean tag filter groups
|
||||
created_after: datetime | None = None, # Only include memory_units created after this time
|
||||
created_before: datetime | None = None, # Only include memory_units created before this time
|
||||
) -> tuple[list[RetrievalResult], GraphRetrievalTimings | None]:
|
||||
"""
|
||||
Retrieve relevant facts via graph traversal.
|
||||
|
||||
@@ -28,8 +28,6 @@ import asyncio
|
||||
import logging
|
||||
import math
|
||||
import time
|
||||
from datetime import datetime
|
||||
from typing import Any
|
||||
|
||||
from ...config import get_config
|
||||
from ..db_utils import acquire_with_retry
|
||||
@@ -51,8 +49,6 @@ async def _find_semantic_seeds(
|
||||
tags: list[str] | None = None,
|
||||
tags_match: TagsMatch = "any",
|
||||
tag_groups: list[TagGroup] | None = None,
|
||||
created_after: datetime | None = None,
|
||||
created_before: datetime | None = None,
|
||||
) -> list[RetrievalResult]:
|
||||
"""Find semantic seeds via embedding search."""
|
||||
from .tags import build_tag_groups_where_clause, build_tags_where_clause_simple
|
||||
@@ -60,24 +56,10 @@ async def _find_semantic_seeds(
|
||||
tags_clause = build_tags_where_clause_simple(tags, 6, match=tags_match)
|
||||
tag_groups_param_start = 6 + (1 if tags else 0)
|
||||
groups_clause, groups_params, _ = build_tag_groups_where_clause(tag_groups, tag_groups_param_start)
|
||||
|
||||
_next_idx = tag_groups_param_start + len(groups_params)
|
||||
created_range_clause = ""
|
||||
created_range_params: list[Any] = []
|
||||
if created_after is not None:
|
||||
created_range_params.append(created_after)
|
||||
created_range_clause += f" AND updated_at > ${_next_idx}"
|
||||
_next_idx += 1
|
||||
if created_before is not None:
|
||||
created_range_params.append(created_before)
|
||||
created_range_clause += f" AND updated_at < ${_next_idx}"
|
||||
_next_idx += 1
|
||||
|
||||
params = [query_embedding_str, bank_id, fact_type, threshold, limit]
|
||||
if tags:
|
||||
params.append(tags)
|
||||
params.extend(groups_params)
|
||||
params.extend(created_range_params)
|
||||
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
@@ -91,7 +73,6 @@ async def _find_semantic_seeds(
|
||||
AND (1 - (embedding <=> $1::vector)) >= $4
|
||||
{tags_clause}
|
||||
{groups_clause}
|
||||
{created_range_clause}
|
||||
ORDER BY embedding <=> $1::vector
|
||||
LIMIT $5
|
||||
""",
|
||||
@@ -140,8 +121,6 @@ class LinkExpansionRetriever(GraphRetriever):
|
||||
tags: list[str] | None = None,
|
||||
tags_match: TagsMatch = "any",
|
||||
tag_groups: list[TagGroup] | None = None,
|
||||
created_after: "datetime | None" = None,
|
||||
created_before: "datetime | None" = None,
|
||||
) -> tuple[list[RetrievalResult], GraphRetrievalTimings | None]:
|
||||
"""
|
||||
Retrieve facts by expanding links from seeds.
|
||||
@@ -180,8 +159,6 @@ class LinkExpansionRetriever(GraphRetriever):
|
||||
tags=tags,
|
||||
tags_match=tags_match,
|
||||
tag_groups=tag_groups,
|
||||
created_after=created_after,
|
||||
created_before=created_before,
|
||||
)
|
||||
timings.seeds_time = time.time() - seeds_start
|
||||
logger.debug(
|
||||
|
||||
@@ -13,7 +13,7 @@ import logging
|
||||
import re
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import UTC, datetime
|
||||
from typing import Any, Optional
|
||||
from typing import Optional
|
||||
|
||||
from ...config import get_config
|
||||
from ..db_utils import acquire_with_retry
|
||||
@@ -98,8 +98,6 @@ async def retrieve_semantic_bm25_combined(
|
||||
tags: list[str] | None = None,
|
||||
tags_match: TagsMatch = "any",
|
||||
tag_groups: list[TagGroup] | None = None,
|
||||
created_after: datetime | None = None,
|
||||
created_before: datetime | None = None,
|
||||
) -> dict[str, tuple[list[RetrievalResult], list[RetrievalResult]]]:
|
||||
"""
|
||||
Combined semantic + BM25 retrieval for multiple fact types in a single query.
|
||||
@@ -165,21 +163,6 @@ async def retrieve_semantic_bm25_combined(
|
||||
tag_groups_param_start = tags_param_idx + (1 if tags else 0)
|
||||
groups_clause, groups_params, _ = build_tag_groups_where_clause(tag_groups, tag_groups_param_start)
|
||||
|
||||
# --- created_at time range filter (appended after tags/groups) ---
|
||||
# Param indices are computed relative to the final params list built below,
|
||||
# so we pre-compute the next available index after all preceding params.
|
||||
_next_idx = tag_groups_param_start + len(groups_params)
|
||||
created_range_clause = ""
|
||||
created_range_params: list[Any] = []
|
||||
if created_after is not None:
|
||||
created_range_params.append(created_after)
|
||||
created_range_clause += f" AND updated_at > ${_next_idx}"
|
||||
_next_idx += 1
|
||||
if created_before is not None:
|
||||
created_range_params.append(created_before)
|
||||
created_range_clause += f" AND updated_at < ${_next_idx}"
|
||||
_next_idx += 1
|
||||
|
||||
# --- Semantic UNION ALL arms (one per fact_type) ---
|
||||
# Each arm has its own ORDER BY embedding <=> $1 LIMIT {hnsw_fetch}, which
|
||||
# lets the planner use the partial HNSW index for that fact_type.
|
||||
@@ -197,7 +180,6 @@ async def retrieve_semantic_bm25_combined(
|
||||
f" AND (1 - (embedding <=> $1::vector)) >= 0.3"
|
||||
f" {tags_clause}"
|
||||
f" {groups_clause}"
|
||||
f" {created_range_clause}"
|
||||
f" ORDER BY embedding <=> $1::vector"
|
||||
f" LIMIT {hnsw_fetch})"
|
||||
)
|
||||
@@ -238,7 +220,6 @@ async def retrieve_semantic_bm25_combined(
|
||||
f" {bm25_where_filter}"
|
||||
f" {tags_clause}"
|
||||
f" {groups_clause}"
|
||||
f" {created_range_clause}"
|
||||
f" ORDER BY {bm25_order_by}"
|
||||
f" LIMIT $3)"
|
||||
)
|
||||
@@ -252,7 +233,6 @@ async def retrieve_semantic_bm25_combined(
|
||||
if tags:
|
||||
params.append(tags)
|
||||
params.extend(groups_params)
|
||||
params.extend(created_range_params)
|
||||
|
||||
rows = await conn.fetch(query, *params)
|
||||
|
||||
@@ -286,8 +266,6 @@ async def retrieve_temporal_combined(
|
||||
tags: list[str] | None = None,
|
||||
tags_match: TagsMatch = "any",
|
||||
tag_groups: list[TagGroup] | None = None,
|
||||
created_after: datetime | None = None,
|
||||
created_before: datetime | None = None,
|
||||
) -> dict[str, list[RetrievalResult]]:
|
||||
"""
|
||||
Temporal retrieval for multiple fact types in a single query.
|
||||
@@ -321,25 +299,10 @@ async def retrieve_temporal_combined(
|
||||
tags_clause = build_tags_where_clause_simple(tags, 7, match=tags_match)
|
||||
tag_groups_param_start = 7 + (1 if tags else 0)
|
||||
groups_clause, groups_params, _ = build_tag_groups_where_clause(tag_groups, tag_groups_param_start)
|
||||
|
||||
# created_at time range filter (after tags/groups)
|
||||
_next_idx = tag_groups_param_start + len(groups_params)
|
||||
created_range_clause = ""
|
||||
created_range_params: list[Any] = []
|
||||
if created_after is not None:
|
||||
created_range_params.append(created_after)
|
||||
created_range_clause += f" AND updated_at > ${_next_idx}"
|
||||
_next_idx += 1
|
||||
if created_before is not None:
|
||||
created_range_params.append(created_before)
|
||||
created_range_clause += f" AND updated_at < ${_next_idx}"
|
||||
_next_idx += 1
|
||||
|
||||
params: list = [query_emb_str, bank_id, fact_types, start_date, end_date, semantic_threshold]
|
||||
if tags:
|
||||
params.append(tags)
|
||||
params.extend(groups_params)
|
||||
params.extend(created_range_params)
|
||||
|
||||
# Two-phase entry point query:
|
||||
# Phase 1 (date_ranked): rank by date only — no embedding computation — for all units in
|
||||
@@ -371,7 +334,6 @@ async def retrieve_temporal_combined(
|
||||
)
|
||||
{tags_clause}
|
||||
{groups_clause}
|
||||
{created_range_clause}
|
||||
),
|
||||
sim_ranked AS (
|
||||
SELECT mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start, mu.occurred_end, mu.mentioned_at, mu.fact_type, mu.proof_count, mu.document_id, mu.chunk_id, mu.tags, mu.metadata,
|
||||
@@ -574,8 +536,6 @@ async def retrieve_all_fact_types_parallel(
|
||||
tags: list[str] | None = None,
|
||||
tags_match: TagsMatch = "any",
|
||||
tag_groups: list[TagGroup] | None = None,
|
||||
created_after: datetime | None = None,
|
||||
created_before: datetime | None = None,
|
||||
) -> MultiFactTypeRetrievalResult:
|
||||
"""
|
||||
Optimized retrieval for multiple fact types using batched queries.
|
||||
@@ -634,8 +594,6 @@ async def retrieve_all_fact_types_parallel(
|
||||
tags=tags,
|
||||
tags_match=tags_match,
|
||||
tag_groups=tag_groups,
|
||||
created_after=created_after,
|
||||
created_before=created_before,
|
||||
)
|
||||
semantic_bm25_time = time.time() - semantic_bm25_start
|
||||
|
||||
@@ -655,8 +613,6 @@ async def retrieve_all_fact_types_parallel(
|
||||
tags=tags,
|
||||
tags_match=tags_match,
|
||||
tag_groups=tag_groups,
|
||||
created_after=created_after,
|
||||
created_before=created_before,
|
||||
)
|
||||
temporal_time = time.time() - temporal_start
|
||||
|
||||
@@ -680,8 +636,6 @@ async def retrieve_all_fact_types_parallel(
|
||||
tags=tags,
|
||||
tags_match=tags_match,
|
||||
tag_groups=tag_groups,
|
||||
created_after=created_after,
|
||||
created_before=created_before,
|
||||
)
|
||||
return ft, results, time.time() - graph_start, graph_timing
|
||||
|
||||
|
||||
@@ -73,7 +73,7 @@ def format_facts_for_prompt(facts: list[MemoryFact]) -> str:
|
||||
|
||||
formatted.append(fact_obj)
|
||||
|
||||
return json.dumps(formatted, indent=2, ensure_ascii=False)
|
||||
return json.dumps(formatted, indent=2)
|
||||
|
||||
|
||||
def format_entity_summaries_for_prompt(entities: dict) -> str:
|
||||
|
||||
@@ -2,8 +2,7 @@
|
||||
Task backend for distributed task processing.
|
||||
|
||||
This provides an abstraction for task storage and execution:
|
||||
- BrokerTaskBackend: Uses PostgreSQL as broker (production API servers)
|
||||
- WorkerTaskBackend: No-op submit_task (production workers — child tasks are polled)
|
||||
- BrokerTaskBackend: Uses PostgreSQL as broker (production)
|
||||
- SyncTaskBackend: Executes tasks immediately (testing/embedded)
|
||||
"""
|
||||
|
||||
@@ -126,33 +125,6 @@ class SyncTaskBackend(TaskBackend):
|
||||
logger.debug("SyncTaskBackend shutdown")
|
||||
|
||||
|
||||
class WorkerTaskBackend(TaskBackend):
|
||||
"""
|
||||
Task backend for worker processes.
|
||||
|
||||
Workers execute tasks directly via the poller (claim → execute), so they
|
||||
don't need submit_task to run anything. When engine code running *inside*
|
||||
a worker-executed task calls submit_task (e.g. retain triggers consolidation),
|
||||
the async-operation row has already been persisted (with task_payload) by
|
||||
_submit_async_operation — so submit_task is a no-op. The new task will be
|
||||
picked up by a worker on the next poll cycle instead of being executed inline,
|
||||
which avoids blocking the parent task.
|
||||
"""
|
||||
|
||||
async def initialize(self):
|
||||
self._initialized = True
|
||||
logger.debug("WorkerTaskBackend initialized")
|
||||
|
||||
async def submit_task(self, task_dict: dict[str, Any]):
|
||||
"""No-op: the row already exists in async_operations; a worker will claim it."""
|
||||
task_type = task_dict.get("type", "unknown")
|
||||
logger.debug(f"WorkerTaskBackend: submit_task no-op for {task_type} (will be picked up by poller)")
|
||||
|
||||
async def shutdown(self):
|
||||
self._initialized = False
|
||||
logger.debug("WorkerTaskBackend shutdown")
|
||||
|
||||
|
||||
class BrokerTaskBackend(TaskBackend):
|
||||
"""
|
||||
Task backend using PostgreSQL as broker.
|
||||
@@ -221,21 +193,17 @@ class BrokerTaskBackend(TaskBackend):
|
||||
table = fq_table("async_operations", schema)
|
||||
|
||||
if operation_id:
|
||||
# Callers now include task_payload in the same INSERT that creates the
|
||||
# async_operations row (see MemoryEngine._submit_async_operation). The
|
||||
# WHERE clause guards against overwriting that payload — the UPDATE is a
|
||||
# no-op when the row is already claimable, and only fills in a NULL payload
|
||||
# for any legacy caller that still creates the row first.
|
||||
# Update existing operation with task payload
|
||||
await pool.execute(
|
||||
f"""
|
||||
UPDATE {table}
|
||||
SET task_payload = $1::jsonb, updated_at = now()
|
||||
WHERE operation_id = $2 AND task_payload IS NULL
|
||||
WHERE operation_id = $2
|
||||
""",
|
||||
payload_json,
|
||||
operation_id,
|
||||
)
|
||||
logger.debug(f"submit_task UPDATE for operation {operation_id} (no-op if payload already set)")
|
||||
logger.debug(f"Updated task payload for operation {operation_id}")
|
||||
else:
|
||||
# Insert new operation (for tasks without pre-created records)
|
||||
# e.g., access_count_update tasks
|
||||
|
||||
@@ -55,7 +55,6 @@ from hindsight_api.extensions.tenant import (
|
||||
TenantExtension,
|
||||
)
|
||||
from hindsight_api.models import RequestContext
|
||||
from hindsight_api.worker.exceptions import DeferOperation
|
||||
|
||||
__all__ = [
|
||||
# Base
|
||||
@@ -69,7 +68,6 @@ __all__ = [
|
||||
# MCP Extension
|
||||
"MCPExtension",
|
||||
# Operation Validator - Core
|
||||
"DeferOperation",
|
||||
"OperationValidationError",
|
||||
"OperationValidatorExtension",
|
||||
"RecallContext",
|
||||
|
||||
@@ -176,22 +176,6 @@ class RetainResult:
|
||||
llm_input_tokens: int | None = None
|
||||
llm_output_tokens: int | None = None
|
||||
llm_total_tokens: int | None = None
|
||||
# Content tokens the retain pipeline actually processed, after
|
||||
# chunk-level content-hash deduplication. Semantics:
|
||||
# None — no dedup signal available (e.g. a first-time retain or a
|
||||
# path that doesn't compute it). Callers that care about
|
||||
# "what was actually new on this retain" should treat None
|
||||
# as "the full submitted content was processed."
|
||||
# 0 — the entire submission was a duplicate of prior content
|
||||
# (all chunks matched by content_hash); nothing went
|
||||
# through LLM extraction.
|
||||
# N>0 — only N tokens of content + context went through the
|
||||
# extraction pipeline. The remainder was dedup'd against
|
||||
# existing chunks.
|
||||
# This is the basis most billing/metering extensions want to use
|
||||
# when the customer's client resubmits growing payloads to the same
|
||||
# document_id (e.g. a session transcript appended to on each turn).
|
||||
processed_content_tokens: int | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -392,16 +376,6 @@ class OperationValidatorExtension(Extension, ABC):
|
||||
2. [operation executes]
|
||||
3. on_*_complete (post-operation)
|
||||
|
||||
Outcomes for `validate_*` hooks:
|
||||
- accept: return `ValidationResult.accept()` (or `accept_with(...)`)
|
||||
- reject: return `ValidationResult.reject(reason, status_code)`
|
||||
(raises `OperationValidationError` upstream)
|
||||
- defer: raise `DeferOperation(exec_date, reason)` from
|
||||
`hindsight_api.worker.exceptions` to requeue the task for a
|
||||
future time without bumping `retry_count`. Worker-only — do
|
||||
not raise from `validate_recall` / `validate_reflect` in
|
||||
synchronous HTTP request paths, where it surfaces as a 500.
|
||||
|
||||
Supported operations:
|
||||
- retain, recall, reflect (core memory operations)
|
||||
- consolidate (mental models consolidation)
|
||||
|
||||
@@ -45,6 +45,7 @@ _ALL_TOOLS: frozenset[str] = frozenset(
|
||||
"delete_directive",
|
||||
"list_memories",
|
||||
"get_memory",
|
||||
"delete_memory",
|
||||
"list_documents",
|
||||
"get_document",
|
||||
"delete_document",
|
||||
@@ -222,6 +223,7 @@ def register_mcp_tools(
|
||||
"delete_directive",
|
||||
"list_memories",
|
||||
"get_memory",
|
||||
"delete_memory",
|
||||
"list_documents",
|
||||
"get_document",
|
||||
"delete_document",
|
||||
@@ -290,6 +292,9 @@ def register_mcp_tools(
|
||||
if "get_memory" in tools_to_register:
|
||||
_register_get_memory(mcp, memory, config)
|
||||
|
||||
if "delete_memory" in tools_to_register:
|
||||
_register_delete_memory(mcp, memory, config)
|
||||
|
||||
# Document tools
|
||||
if "list_documents" in tools_to_register:
|
||||
_register_list_documents(mcp, memory, config)
|
||||
@@ -436,6 +441,7 @@ _AUDITABLE_MCP_TOOLS: frozenset[str] = frozenset(
|
||||
"refresh_mental_model",
|
||||
"create_directive",
|
||||
"delete_directive",
|
||||
"delete_memory",
|
||||
"delete_document",
|
||||
"cancel_operation",
|
||||
}
|
||||
@@ -2157,6 +2163,74 @@ def _register_get_memory(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsCon
|
||||
return {"error": str(e)}
|
||||
|
||||
|
||||
def _register_delete_memory(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig) -> None:
|
||||
"""Register the delete_memory tool."""
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
async def delete_memory(
|
||||
memory_id: str,
|
||||
bank_id: str | None = None,
|
||||
) -> str:
|
||||
"""
|
||||
Delete a specific memory by ID.
|
||||
|
||||
Permanently removes a memory unit and its associated data.
|
||||
|
||||
Args:
|
||||
memory_id: The ID of the memory to delete
|
||||
bank_id: Optional bank (accepted for consistency, not used in deletion).
|
||||
"""
|
||||
try:
|
||||
target_bank = bank_id or config.bank_id_resolver()
|
||||
if target_bank is None:
|
||||
return '{"error": "No bank_id configured"}'
|
||||
|
||||
result = await memory.delete_memory_unit(
|
||||
unit_id=memory_id,
|
||||
request_context=_get_request_context(config),
|
||||
)
|
||||
return json.dumps({"status": "deleted", "memory_id": memory_id, **result}, default=str)
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return json.dumps({"error": str(e)})
|
||||
except Exception as e:
|
||||
logger.error(f"Error deleting memory: {e}", exc_info=True)
|
||||
return f'{{"error": "{e}"}}'
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
async def delete_memory(
|
||||
memory_id: str,
|
||||
) -> dict:
|
||||
"""
|
||||
Delete a specific memory by ID.
|
||||
|
||||
Permanently removes a memory unit and its associated data.
|
||||
|
||||
Args:
|
||||
memory_id: The ID of the memory to delete
|
||||
"""
|
||||
try:
|
||||
target_bank = config.bank_id_resolver()
|
||||
if target_bank is None:
|
||||
return {"error": "No bank_id configured"}
|
||||
|
||||
result = await memory.delete_memory_unit(
|
||||
unit_id=memory_id,
|
||||
request_context=_get_request_context(config),
|
||||
)
|
||||
return {"status": "deleted", "memory_id": memory_id, **result}
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return {"error": str(e)}
|
||||
except Exception as e:
|
||||
logger.error(f"Error deleting memory: {e}", exc_info=True)
|
||||
return {"error": str(e)}
|
||||
|
||||
|
||||
# =========================================================================
|
||||
# DOCUMENT TOOLS
|
||||
# =========================================================================
|
||||
@@ -2780,44 +2854,6 @@ def _register_get_bank_stats(mcp: FastMCP, memory: MemoryEngine, config: MCPTool
|
||||
return f'{{"error": "{e}"}}'
|
||||
|
||||
|
||||
async def _do_update_bank(
|
||||
memory: MemoryEngine,
|
||||
target_bank: str,
|
||||
request_context: RequestContext,
|
||||
*,
|
||||
name: str | None = None,
|
||||
mission: str | None = None,
|
||||
config_updates: dict[str, Any] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Shared implementation for update_bank MCP tool variants.
|
||||
|
||||
Args:
|
||||
name: Display name (stored in banks table).
|
||||
mission: Deprecated alias for reflect_mission — mapped into config_updates.
|
||||
config_updates: Arbitrary config overrides passed to config_resolver.update_bank_config().
|
||||
Supports all configurable fields (retain_mission, disposition_*, etc.).
|
||||
The config resolver validates keys and rejects non-configurable/credential fields.
|
||||
"""
|
||||
# Update display name via engine (stored in DB banks table)
|
||||
if name is not None:
|
||||
await memory.update_bank(
|
||||
target_bank,
|
||||
name=name,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Merge deprecated mission alias into config_updates as reflect_mission
|
||||
effective_config: dict[str, Any] = dict(config_updates) if config_updates else {}
|
||||
if mission is not None and "reflect_mission" not in effective_config:
|
||||
effective_config["reflect_mission"] = mission
|
||||
|
||||
if effective_config:
|
||||
await memory._config_resolver.update_bank_config(target_bank, effective_config, request_context)
|
||||
|
||||
# Return updated profile
|
||||
return await memory.get_bank_profile(target_bank, request_context=request_context)
|
||||
|
||||
|
||||
def _register_update_bank(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig) -> None:
|
||||
"""Register the update_bank tool."""
|
||||
|
||||
@@ -2827,37 +2863,16 @@ def _register_update_bank(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsCo
|
||||
async def update_bank(
|
||||
name: str | None = None,
|
||||
mission: str | None = None,
|
||||
config_updates: dict[str, Any] | None = None,
|
||||
bank_id: str | None = None,
|
||||
) -> str:
|
||||
"""
|
||||
Update a memory bank's configuration.
|
||||
Update a memory bank's metadata.
|
||||
|
||||
Updates the bank's name and/or any bank-level configuration fields.
|
||||
Only provided fields will be updated; omitted fields remain unchanged.
|
||||
Changes the name or mission of an existing bank.
|
||||
|
||||
Args:
|
||||
name: Human-friendly display name for the bank.
|
||||
mission: Deprecated alias for config_updates.reflect_mission.
|
||||
config_updates: Dictionary of configuration fields to update. Supports all
|
||||
bank-configurable fields including:
|
||||
- reflect_mission: Mission/context for Reflect operations.
|
||||
- retain_mission: Steers what gets extracted during retain().
|
||||
- retain_extraction_mode: 'concise' (default), 'verbose', or 'custom'.
|
||||
- retain_custom_instructions: Custom extraction prompt (active when mode is 'custom').
|
||||
- retain_chunk_size: Maximum token size for each content chunk.
|
||||
- retain_chunk_batch_size: Number of chunks to process in parallel.
|
||||
- enable_observations: Toggle observation consolidation after retain().
|
||||
- observations_mission: Controls observation synthesis rules.
|
||||
- disposition_skepticism: Critical evaluation level (1-5).
|
||||
- disposition_literalism: Literal vs. abstract interpretation (1-5).
|
||||
- disposition_empathy: Emotional context consideration (1-5).
|
||||
- entity_labels: Controlled vocabulary for entity classification.
|
||||
- entities_allow_free_form: Allow labels outside entity_labels.
|
||||
- recall_include_chunks: Include raw chunks in recall results.
|
||||
- recall_max_tokens: Max tokens for recall results.
|
||||
- mcp_enabled_tools: Tool allowlist for this bank.
|
||||
Any configurable field name is accepted (use Python field names).
|
||||
name: New human-friendly name for the bank
|
||||
mission: New mission describing who the agent is and what they're trying to accomplish
|
||||
bank_id: Optional bank (defaults to session bank). Use for cross-bank operations.
|
||||
"""
|
||||
try:
|
||||
@@ -2865,16 +2880,14 @@ def _register_update_bank(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsCo
|
||||
if target_bank is None:
|
||||
return '{"error": "No bank_id configured"}'
|
||||
|
||||
result = await _do_update_bank(
|
||||
memory,
|
||||
result = await memory.update_bank(
|
||||
target_bank,
|
||||
_get_request_context(config),
|
||||
name=name,
|
||||
mission=mission,
|
||||
config_updates=config_updates,
|
||||
request_context=_get_request_context(config),
|
||||
)
|
||||
return json.dumps(result, indent=2, default=str)
|
||||
except (OperationValidationError, ValueError) as e:
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return json.dumps({"error": str(e)})
|
||||
except Exception as e:
|
||||
@@ -2887,52 +2900,29 @@ def _register_update_bank(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsCo
|
||||
async def update_bank(
|
||||
name: str | None = None,
|
||||
mission: str | None = None,
|
||||
config_updates: dict[str, Any] | None = None,
|
||||
) -> dict:
|
||||
"""
|
||||
Update this memory bank's configuration.
|
||||
Update this memory bank's metadata.
|
||||
|
||||
Updates the bank's name and/or any bank-level configuration fields.
|
||||
Only provided fields will be updated; omitted fields remain unchanged.
|
||||
Changes the name or mission of the bank.
|
||||
|
||||
Args:
|
||||
name: Human-friendly display name for the bank.
|
||||
mission: Deprecated alias for config_updates.reflect_mission.
|
||||
config_updates: Dictionary of configuration fields to update. Supports all
|
||||
bank-configurable fields including:
|
||||
- reflect_mission: Mission/context for Reflect operations.
|
||||
- retain_mission: Steers what gets extracted during retain().
|
||||
- retain_extraction_mode: 'concise' (default), 'verbose', or 'custom'.
|
||||
- retain_custom_instructions: Custom extraction prompt (active when mode is 'custom').
|
||||
- retain_chunk_size: Maximum token size for each content chunk.
|
||||
- retain_chunk_batch_size: Number of chunks to process in parallel.
|
||||
- enable_observations: Toggle observation consolidation after retain().
|
||||
- observations_mission: Controls observation synthesis rules.
|
||||
- disposition_skepticism: Critical evaluation level (1-5).
|
||||
- disposition_literalism: Literal vs. abstract interpretation (1-5).
|
||||
- disposition_empathy: Emotional context consideration (1-5).
|
||||
- entity_labels: Controlled vocabulary for entity classification.
|
||||
- entities_allow_free_form: Allow labels outside entity_labels.
|
||||
- recall_include_chunks: Include raw chunks in recall results.
|
||||
- recall_max_tokens: Max tokens for recall results.
|
||||
- mcp_enabled_tools: Tool allowlist for this bank.
|
||||
Any configurable field name is accepted (use Python field names).
|
||||
name: New human-friendly name for the bank
|
||||
mission: New mission describing who the agent is and what they're trying to accomplish
|
||||
"""
|
||||
try:
|
||||
target_bank = config.bank_id_resolver()
|
||||
if target_bank is None:
|
||||
return {"error": "No bank_id configured"}
|
||||
|
||||
result = await _do_update_bank(
|
||||
memory,
|
||||
result = await memory.update_bank(
|
||||
target_bank,
|
||||
_get_request_context(config),
|
||||
name=name,
|
||||
mission=mission,
|
||||
config_updates=config_updates,
|
||||
request_context=_get_request_context(config),
|
||||
)
|
||||
return result
|
||||
except (OperationValidationError, ValueError) as e:
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return {"error": str(e)}
|
||||
except Exception as e:
|
||||
|
||||
@@ -27,7 +27,6 @@ from alembic.config import Config
|
||||
from alembic.script.revision import ResolutionError
|
||||
from sqlalchemy import Connection, create_engine, text
|
||||
|
||||
from .db_url import to_libpq_url
|
||||
from .utils import mask_network_location
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -221,7 +220,7 @@ def run_migrations(
|
||||
# ineffective when the app URL goes through a pooler. Configure
|
||||
# HINDSIGHT_API_MIGRATION_DATABASE_URL to the direct PostgreSQL endpoint
|
||||
# (e.g. hindsight-pg-rw) to restore correct locking behaviour.
|
||||
migration_url = to_libpq_url(migration_database_url or database_url)
|
||||
migration_url = migration_database_url or database_url
|
||||
|
||||
try:
|
||||
# Determine script location
|
||||
@@ -451,7 +450,7 @@ def check_migration_status(
|
||||
return None, None
|
||||
|
||||
# Get current revision from database
|
||||
engine = create_engine(to_libpq_url(database_url))
|
||||
engine = create_engine(database_url)
|
||||
with engine.connect() as connection:
|
||||
context = MigrationContext.configure(connection)
|
||||
current_rev = context.get_current_revision()
|
||||
@@ -625,7 +624,7 @@ def ensure_embedding_dimension(
|
||||
"""
|
||||
schema_name = schema or "public"
|
||||
|
||||
engine = create_engine(to_libpq_url(database_url))
|
||||
engine = create_engine(database_url)
|
||||
with engine.connect() as conn:
|
||||
# Check if memory_units table exists (proxy for schema being initialized)
|
||||
table_exists = conn.execute(
|
||||
@@ -674,7 +673,7 @@ def ensure_vector_extension(
|
||||
"""
|
||||
schema_name = schema or "public"
|
||||
|
||||
engine = create_engine(to_libpq_url(database_url))
|
||||
engine = create_engine(database_url)
|
||||
with engine.connect() as conn:
|
||||
# Detect which vector extension should be used
|
||||
target_ext = _detect_vector_extension(conn, vector_extension)
|
||||
@@ -895,7 +894,7 @@ def ensure_text_search_extension(
|
||||
"""
|
||||
schema_name = schema or "public"
|
||||
|
||||
engine = create_engine(to_libpq_url(database_url))
|
||||
engine = create_engine(database_url)
|
||||
with engine.connect() as conn:
|
||||
# Tables with search_vector columns to check
|
||||
tables_to_check = [
|
||||
|
||||
@@ -24,12 +24,6 @@ class RequestContext:
|
||||
mcp_authenticated: bool = False # True when MCP transport auth already validated (skips tenant re-auth)
|
||||
user_initiated: bool = False # True for async operations that originated from a user request
|
||||
allowed_bank_ids: list[str] | None = None # None = unrestricted (all banks)
|
||||
# Number of times this task has been retried. Populated by the worker
|
||||
# from async_operations.retry_count before dispatching to a task handler;
|
||||
# 0 for sync/HTTP requests and for the first worker attempt. Useful for
|
||||
# validators that want exponential backoff on repeated failures (e.g.
|
||||
# "defer for 2^retry_count minutes") without querying the DB themselves.
|
||||
retry_count: int = 0
|
||||
|
||||
|
||||
from pgvector.sqlalchemy import Vector
|
||||
@@ -68,6 +62,7 @@ class Document(Base):
|
||||
bank_id: Mapped[str] = mapped_column(Text, primary_key=True)
|
||||
original_text: Mapped[str | None] = mapped_column(Text)
|
||||
content_hash: Mapped[str | None] = mapped_column(Text)
|
||||
doc_metadata: Mapped[dict] = mapped_column("metadata", JSONB, server_default=sql_text("'{}'::jsonb"))
|
||||
created_at: Mapped[datetime] = mapped_column(TIMESTAMP(timezone=True), server_default=func.now())
|
||||
updated_at: Mapped[datetime] = mapped_column(TIMESTAMP(timezone=True), server_default=func.now())
|
||||
|
||||
|
||||
@@ -7,24 +7,3 @@ class RetryTaskAt(Exception):
|
||||
def __init__(self, retry_at: datetime, message: str = ""):
|
||||
self.retry_at = retry_at
|
||||
super().__init__(message)
|
||||
|
||||
|
||||
class DeferOperation(Exception):
|
||||
"""Raise from an extension hook (or task handler) to requeue the
|
||||
operation for execution at a later time, without counting as a retry.
|
||||
|
||||
Unlike `RetryTaskAt`, this is not a failure: `retry_count` is not
|
||||
incremented and `error_message` is not written. Use this for
|
||||
backpressure / "not yet, try later" decisions made before or during
|
||||
task execution (e.g. quota windows, warming dependencies, upstream
|
||||
rate limits).
|
||||
|
||||
Worker-only: raising this from a hook called in HTTP request context
|
||||
(e.g. `validate_recall` for a synchronous recall) will surface as an
|
||||
unhandled 500 — there is no queue to defer to.
|
||||
"""
|
||||
|
||||
def __init__(self, exec_date: datetime, reason: str = ""):
|
||||
self.exec_date = exec_date
|
||||
self.reason = reason
|
||||
super().__init__(reason)
|
||||
|
||||
@@ -18,7 +18,7 @@ import sys
|
||||
import warnings
|
||||
|
||||
from ..config import get_config
|
||||
from ..engine.task_backend import WorkerTaskBackend
|
||||
from ..engine.task_backend import SyncTaskBackend
|
||||
from .poller import WorkerPoller
|
||||
|
||||
# Filter deprecation warnings from third-party libraries
|
||||
@@ -164,11 +164,7 @@ def main():
|
||||
print(f" Poll interval: {args.poll_interval}ms")
|
||||
print(f" Max retries: {args.max_retries}")
|
||||
print(f" Max slots: {config.worker_max_slots}")
|
||||
reservations = config.worker_slot_reservations
|
||||
reservations_str = ", ".join(f"{k}={v}" for k, v in reservations.items()) if reservations else "none"
|
||||
shared_pool = max(0, config.worker_max_slots - sum(reservations.values()))
|
||||
print(f" Slot reservations: {reservations_str}")
|
||||
print(f" Shared pool: {shared_pool}")
|
||||
print(f" Consolidation max slots: {config.worker_consolidation_max_slots}")
|
||||
print(f" HTTP server: {args.http_host}:{args.http_port}")
|
||||
print()
|
||||
|
||||
@@ -195,13 +191,11 @@ def main():
|
||||
logger.info(f"Loaded operation validator: {operation_validator.__class__.__name__}")
|
||||
|
||||
# Initialize MemoryEngine
|
||||
# Workers use WorkerTaskBackend: submit_task is a no-op because the
|
||||
# row already exists in async_operations. Child tasks (e.g. consolidation
|
||||
# triggered by retain) will be picked up by the poller on the next cycle
|
||||
# instead of being executed inline, which avoids blocking the parent task.
|
||||
# Workers use SyncTaskBackend because they execute tasks directly,
|
||||
# they don't need to store tasks (they poll from DB)
|
||||
memory = MemoryEngine(
|
||||
run_migrations=False, # Workers don't run migrations
|
||||
task_backend=WorkerTaskBackend(),
|
||||
task_backend=SyncTaskBackend(),
|
||||
tenant_extension=tenant_extension,
|
||||
operation_validator=operation_validator,
|
||||
)
|
||||
@@ -228,7 +222,7 @@ def main():
|
||||
schema=schema,
|
||||
tenant_extension=tenant_extension,
|
||||
max_slots=config.worker_max_slots,
|
||||
slot_reservations=config.worker_slot_reservations,
|
||||
consolidation_max_slots=config.worker_consolidation_max_slots,
|
||||
)
|
||||
|
||||
# Create the HTTP app for metrics/health
|
||||
|
||||
@@ -15,7 +15,7 @@ from collections.abc import Awaitable, Callable
|
||||
from dataclasses import dataclass, field
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from .exceptions import DeferOperation, RetryTaskAt
|
||||
from .exceptions import RetryTaskAt
|
||||
from .stage import StageHolder, bind_holder
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -70,21 +70,6 @@ class ClaimedTask:
|
||||
schema: str | None
|
||||
|
||||
|
||||
@dataclass
|
||||
class SlotAvailability:
|
||||
"""Available slot capacity across reserved and shared pools.
|
||||
|
||||
Each operation type with a reservation has its own reserved pool.
|
||||
The shared pool (max_slots - sum of reservations) is usable by any type.
|
||||
"""
|
||||
|
||||
reserved: dict[str, int]
|
||||
"""Per-operation-type remaining reserved capacity."""
|
||||
|
||||
shared: int
|
||||
"""Remaining shared pool capacity (usable by any operation type)."""
|
||||
|
||||
|
||||
class WorkerPoller:
|
||||
"""
|
||||
Polls PostgreSQL for pending tasks and executes them.
|
||||
@@ -104,7 +89,7 @@ class WorkerPoller:
|
||||
schema: str | None = None,
|
||||
tenant_extension: "TenantExtension | None" = None,
|
||||
max_slots: int = 10,
|
||||
slot_reservations: dict[str, int] | None = None,
|
||||
consolidation_max_slots: int = 2,
|
||||
):
|
||||
"""
|
||||
Initialize the worker poller.
|
||||
@@ -118,10 +103,7 @@ class WorkerPoller:
|
||||
tenant_extension: Extension for dynamic multi-tenant discovery. If None, creates a
|
||||
DefaultTenantExtension with the configured schema.
|
||||
max_slots: Maximum concurrent tasks per worker
|
||||
slot_reservations: Per-operation-type reserved slot counts (e.g. {"consolidation": 2,
|
||||
"retain": 3}). Reserved slots guarantee capacity for that operation type.
|
||||
Remaining slots (max_slots - sum of reservations) form a shared pool usable
|
||||
by any operation type. Defaults to {"consolidation": 2} if None.
|
||||
consolidation_max_slots: Maximum concurrent consolidation tasks per worker
|
||||
"""
|
||||
self._pool = pool
|
||||
self._worker_id = worker_id
|
||||
@@ -137,9 +119,7 @@ class WorkerPoller:
|
||||
tenant_extension = DefaultTenantExtension(config=config)
|
||||
self._tenant_extension = tenant_extension
|
||||
self._max_slots = max_slots
|
||||
self._slot_reservations: dict[str, int] = (
|
||||
slot_reservations if slot_reservations is not None else {"consolidation": 2}
|
||||
)
|
||||
self._consolidation_max_slots = consolidation_max_slots
|
||||
self._shutdown = asyncio.Event()
|
||||
self._current_tasks: set[asyncio.Task] = set()
|
||||
self._in_flight_count = 0
|
||||
@@ -150,9 +130,6 @@ class WorkerPoller:
|
||||
self._active_tasks: dict[str, ActiveTaskInfo] = {}
|
||||
# Track in-flight tasks by operation type
|
||||
self._in_flight_by_type: dict[str, int] = {}
|
||||
# Rotation offset for per-tenant fair claiming. Advances past the last
|
||||
# schema we serviced so a busy tenant can't monopolize the poll order.
|
||||
self._next_schema_idx: int = 0
|
||||
|
||||
async def _get_schemas(self) -> list[str | None]:
|
||||
"""Get list of schemas to poll. Returns [None] for default schema (no prefix)."""
|
||||
@@ -162,91 +139,29 @@ class WorkerPoller:
|
||||
# Convert default schema to None for SQL compatibility (no prefix), keep others as-is
|
||||
return [t.schema if t.schema != DEFAULT_DATABASE_SCHEMA else None for t in tenants]
|
||||
|
||||
async def _scan_active_schemas(self, schemas: list[str | None]) -> set[str | None]:
|
||||
"""Find which schemas have pending work.
|
||||
|
||||
Tries a server-side PL/pgSQL function first (single DB round-trip,
|
||||
~200ms for 1400+ schemas). Falls back to per-schema Python EXISTS
|
||||
queries if the function is not installed (~4ms each).
|
||||
|
||||
The server-side function should be installed in the ``public``
|
||||
schema as::
|
||||
|
||||
CREATE OR REPLACE FUNCTION public.schemas_with_pending_work()
|
||||
RETURNS SETOF text AS $$
|
||||
DECLARE
|
||||
r RECORD; has_work BOOLEAN;
|
||||
BEGIN
|
||||
FOR r IN SELECT nspname FROM pg_namespace
|
||||
WHERE nspname LIKE 'tenant_%' LOOP
|
||||
BEGIN
|
||||
EXECUTE format(
|
||||
'SELECT EXISTS(SELECT 1 FROM %I.async_operations '
|
||||
'WHERE status = ''pending'' '
|
||||
'AND task_payload IS NOT NULL LIMIT 1)',
|
||||
r.nspname) INTO has_work;
|
||||
IF has_work THEN RETURN NEXT r.nspname; END IF;
|
||||
EXCEPTION WHEN OTHERS THEN NULL;
|
||||
END;
|
||||
END LOOP;
|
||||
END $$ LANGUAGE plpgsql STABLE;
|
||||
|
||||
In hindsight-cloud deployments this is installed by a Helm hook
|
||||
job alongside ``total_pending_tasks()``.
|
||||
"""
|
||||
async with self._pool.acquire() as conn:
|
||||
try:
|
||||
rows = await conn.fetch("SELECT * FROM schemas_with_pending_work()")
|
||||
return {r[0] for r in rows}
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Fallback: per-schema EXISTS checks from Python
|
||||
active: set[str | None] = set()
|
||||
for schema in schemas:
|
||||
table = fq_table("async_operations", schema)
|
||||
try:
|
||||
has_work = await conn.fetchval(
|
||||
f"SELECT EXISTS(SELECT 1 FROM {table} "
|
||||
f"WHERE status = 'pending' AND task_payload IS NOT NULL LIMIT 1)"
|
||||
)
|
||||
if has_work:
|
||||
active.add(schema)
|
||||
except Exception:
|
||||
pass
|
||||
return active
|
||||
|
||||
async def _get_available_slots(self) -> SlotAvailability:
|
||||
async def _get_available_slots(self) -> tuple[int, int]:
|
||||
"""
|
||||
Calculate available slots for claiming tasks.
|
||||
|
||||
Each operation type can have reserved slots (via ``slot_reservations``).
|
||||
Reserved slots guarantee capacity for that type — they cannot be used by
|
||||
other types. The remaining slots (``max_slots - sum(reservations)``) form
|
||||
a shared pool usable by any operation type on a first-come basis.
|
||||
Consolidation has a reserved pool of ``consolidation_max_slots`` within
|
||||
``max_slots``. Non-consolidation tasks may use at most
|
||||
``max_slots - consolidation_max_slots`` slots, leaving the remainder
|
||||
always available for consolidation. This prevents consolidation from
|
||||
being starved when retain throughput continuously saturates the queue.
|
||||
|
||||
When an operation type's in-flight count exceeds its reservation, the
|
||||
excess tasks are considered to be using shared pool slots.
|
||||
Returns:
|
||||
(non_consolidation_available, consolidation_available) tuple
|
||||
"""
|
||||
async with self._in_flight_lock:
|
||||
total_in_flight = self._in_flight_count
|
||||
in_flight_snapshot = dict(self._in_flight_by_type)
|
||||
consolidation_in_flight = self._in_flight_by_type.get("consolidation", 0)
|
||||
|
||||
# Per-type reserved availability
|
||||
reserved_available: dict[str, int] = {}
|
||||
tasks_in_reserved = 0
|
||||
for op_type, reserved in self._slot_reservations.items():
|
||||
in_flight = in_flight_snapshot.get(op_type, 0)
|
||||
reserved_available[op_type] = max(0, reserved - in_flight)
|
||||
tasks_in_reserved += min(reserved, in_flight)
|
||||
non_consolidation_in_flight = max(0, total_in_flight - consolidation_in_flight)
|
||||
non_consolidation_max = max(0, self._max_slots - self._consolidation_max_slots)
|
||||
non_consolidation_available = max(0, non_consolidation_max - non_consolidation_in_flight)
|
||||
consolidation_available = max(0, self._consolidation_max_slots - consolidation_in_flight)
|
||||
|
||||
# Shared pool: total slots minus reservations minus tasks using shared slots
|
||||
sum_reservations = sum(self._slot_reservations.values())
|
||||
shared_pool_size = max(0, self._max_slots - sum_reservations)
|
||||
tasks_in_shared = max(0, total_in_flight - tasks_in_reserved)
|
||||
shared_available = max(0, shared_pool_size - tasks_in_shared)
|
||||
|
||||
return SlotAvailability(reserved=reserved_available, shared=shared_available)
|
||||
return non_consolidation_available, consolidation_available
|
||||
|
||||
async def wait_for_active_tasks(self, timeout: float = 10.0) -> bool:
|
||||
"""
|
||||
@@ -277,119 +192,47 @@ class WorkerPoller:
|
||||
async def claim_batch(self) -> list[ClaimedTask]:
|
||||
"""
|
||||
Claim pending tasks atomically across all tenant schemas,
|
||||
respecting per-operation-type slot reservations and shared pool limits.
|
||||
respecting slot limits (total and consolidation).
|
||||
|
||||
Uses FOR UPDATE SKIP LOCKED to ensure no conflicts with other workers.
|
||||
|
||||
Schema iteration is round-robin to prevent one busy tenant from
|
||||
starving others. Each poll starts at ``self._next_schema_idx`` and
|
||||
wraps around the full list. First pass caps at 1 claim per pool per
|
||||
schema so every tenant with pending work gets a fair chance; a second
|
||||
pass backfills remaining slots from any schema when there's spare
|
||||
capacity. After the call, the offset advances past the last
|
||||
schema we serviced (or by 1 if nothing was claimed) so the next
|
||||
poll starts at a different position.
|
||||
|
||||
Returns:
|
||||
List of ClaimedTask objects containing operation_id, task_dict, and schema
|
||||
"""
|
||||
# Calculate available slots (per-type reserved + shared pool)
|
||||
availability = await self._get_available_slots()
|
||||
# Calculate available slots (independent pools after reservation)
|
||||
non_consolidation_available, consolidation_available = await self._get_available_slots()
|
||||
|
||||
if all(v <= 0 for v in availability.reserved.values()) and availability.shared <= 0:
|
||||
if non_consolidation_available <= 0 and consolidation_available <= 0:
|
||||
return []
|
||||
|
||||
schemas = await self._get_schemas()
|
||||
if not schemas:
|
||||
return []
|
||||
|
||||
# Scan: find which schemas have pending work using a lightweight
|
||||
# EXISTS check (no locks). Then only claim from those schemas
|
||||
# using the expensive FOR UPDATE SKIP LOCKED query.
|
||||
active_schemas = await self._scan_active_schemas(schemas)
|
||||
|
||||
if not active_schemas:
|
||||
self._next_schema_idx = (self._next_schema_idx + 1) % len(schemas)
|
||||
return []
|
||||
|
||||
# Build rotation list from active schemas only, preserving their
|
||||
# original positions for correct offset advancement.
|
||||
all_indexed = list(enumerate(schemas))
|
||||
active_indexed = [(i, s) for i, s in all_indexed if s in active_schemas]
|
||||
|
||||
# Rotate so no tenant is always first.
|
||||
start = self._next_schema_idx % len(schemas)
|
||||
rotated = [x for x in active_indexed if x[0] >= start] + [x for x in active_indexed if x[0] < start]
|
||||
|
||||
all_tasks: list[ClaimedTask] = []
|
||||
remaining_reserved = dict(availability.reserved)
|
||||
remaining_shared = availability.shared
|
||||
last_serviced_idx: int | None = None
|
||||
schemas_with_work: list[tuple[int, str | None]] = []
|
||||
remaining_non_consolidation = non_consolidation_available
|
||||
remaining_consolidation = consolidation_available
|
||||
|
||||
def _has_capacity() -> bool:
|
||||
return any(v > 0 for v in remaining_reserved.values()) or remaining_shared > 0
|
||||
|
||||
def _account_tasks(tasks: list[ClaimedTask]) -> None:
|
||||
nonlocal remaining_shared
|
||||
for task in tasks:
|
||||
op_type = task.task_dict.get("operation_type", "unknown")
|
||||
if op_type in remaining_reserved and remaining_reserved[op_type] > 0:
|
||||
remaining_reserved[op_type] -= 1
|
||||
else:
|
||||
remaining_shared -= 1
|
||||
|
||||
# Pass 1: fairness pass — iterate only active schemas, cap at
|
||||
# 1 claim per pool per schema.
|
||||
for orig_idx, schema in rotated:
|
||||
if not _has_capacity():
|
||||
for schema in schemas:
|
||||
if remaining_non_consolidation <= 0 and remaining_consolidation <= 0:
|
||||
break
|
||||
|
||||
fair_reserved = {t: min(1, v) for t, v in remaining_reserved.items() if v > 0}
|
||||
fair_shared = min(1, remaining_shared) if remaining_shared > 0 else 0
|
||||
tasks = await self._claim_batch_for_schema(schema, fair_reserved, fair_shared)
|
||||
tasks = await self._claim_batch_for_schema(schema, remaining_non_consolidation, remaining_consolidation)
|
||||
|
||||
_account_tasks(tasks)
|
||||
|
||||
if tasks:
|
||||
last_serviced_idx = orig_idx
|
||||
schemas_with_work.append((orig_idx, schema))
|
||||
for task in tasks:
|
||||
op_type = task.task_dict.get("operation_type", "unknown")
|
||||
if op_type == "consolidation":
|
||||
remaining_consolidation -= 1
|
||||
else:
|
||||
remaining_non_consolidation -= 1
|
||||
|
||||
all_tasks.extend(tasks)
|
||||
|
||||
# Pass 2: capacity pass — fill remaining slots from schemas
|
||||
# that had work in pass 1 only.
|
||||
if _has_capacity() and schemas_with_work:
|
||||
for orig_idx, schema in schemas_with_work:
|
||||
if not _has_capacity():
|
||||
break
|
||||
|
||||
tasks = await self._claim_batch_for_schema(
|
||||
schema, {t: v for t, v in remaining_reserved.items() if v > 0}, remaining_shared
|
||||
)
|
||||
|
||||
_account_tasks(tasks)
|
||||
|
||||
if tasks:
|
||||
last_serviced_idx = orig_idx
|
||||
|
||||
all_tasks.extend(tasks)
|
||||
|
||||
# Advance offset past the last schema we serviced, or by 1 if
|
||||
# nothing was claimed (so we don't keep re-hitting an empty head).
|
||||
if last_serviced_idx is not None:
|
||||
self._next_schema_idx = (last_serviced_idx + 1) % len(schemas)
|
||||
else:
|
||||
self._next_schema_idx = (start + 1) % len(schemas)
|
||||
|
||||
return all_tasks
|
||||
|
||||
async def _claim_batch_for_schema(
|
||||
self, schema: str | None, reserved_limits: dict[str, int], shared_limit: int
|
||||
self, schema: str | None, non_consolidation_limit: int, consolidation_limit: int
|
||||
) -> list[ClaimedTask]:
|
||||
"""Claim tasks from a specific schema respecting per-type and shared slot limits."""
|
||||
"""Claim tasks from a specific schema respecting slot limits."""
|
||||
try:
|
||||
return await self._claim_batch_for_schema_inner(schema, reserved_limits, shared_limit)
|
||||
return await self._claim_batch_for_schema_inner(schema, non_consolidation_limit, consolidation_limit)
|
||||
except Exception as e:
|
||||
# Format schema for logging: custom schemas in quotes, None as-is
|
||||
schema_display = f'"{schema}"' if schema else str(schema)
|
||||
@@ -397,173 +240,67 @@ class WorkerPoller:
|
||||
return []
|
||||
|
||||
async def _claim_batch_for_schema_inner(
|
||||
self, schema: str | None, reserved_limits: dict[str, int], shared_limit: int
|
||||
self, schema: str | None, non_consolidation_limit: int, consolidation_limit: int
|
||||
) -> list[ClaimedTask]:
|
||||
"""Inner implementation for claiming tasks from a specific schema.
|
||||
"""Inner implementation for claiming tasks from a specific schema with slot limits.
|
||||
|
||||
Claims happen in two phases:
|
||||
1. Reserved pools: one query per operation type that has reserved slots.
|
||||
Consolidation queries always include bank-serialization (no two consolidation
|
||||
tasks for the same bank simultaneously).
|
||||
2. Shared pool: remaining capacity is filled by any operation type. Two queries
|
||||
are used (non-consolidation + consolidation with bank serialization) to
|
||||
preserve consolidation's bank-serialization constraint.
|
||||
|
||||
Within the same transaction, rows locked by earlier queries are excluded from
|
||||
later queries via ``operation_id != ALL($excluded)`` since ``FOR UPDATE SKIP
|
||||
LOCKED`` only skips rows locked by *other* transactions.
|
||||
Non-consolidation and consolidation pools are independent: each is bounded by
|
||||
its own limit and they do not borrow from each other.
|
||||
"""
|
||||
table = fq_table("async_operations", schema)
|
||||
|
||||
async with self._pool.acquire() as conn:
|
||||
async with conn.transaction():
|
||||
all_rows: list[Any] = []
|
||||
claimed_ids: list[Any] = []
|
||||
# 1. Claim non-consolidation tasks
|
||||
non_consolidation_rows = []
|
||||
if non_consolidation_limit > 0:
|
||||
non_consolidation_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type != 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
ORDER BY created_at
|
||||
LIMIT $1
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
non_consolidation_limit,
|
||||
)
|
||||
|
||||
# --- Phase 1: claim from reserved pools ---
|
||||
for op_type, limit in reserved_limits.items():
|
||||
if limit <= 0:
|
||||
continue
|
||||
# 2. Claim consolidation tasks from their reserved pool
|
||||
consolidation_rows = []
|
||||
if consolidation_limit > 0:
|
||||
consolidation_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, task_payload, retry_count
|
||||
FROM {table} AS pending
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND NOT EXISTS (
|
||||
SELECT 1 FROM {table} AS processing
|
||||
WHERE processing.bank_id = pending.bank_id
|
||||
AND processing.operation_type = 'consolidation'
|
||||
AND processing.status = 'processing'
|
||||
)
|
||||
ORDER BY created_at
|
||||
LIMIT $1
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
consolidation_limit,
|
||||
)
|
||||
|
||||
if op_type == "consolidation":
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table} AS pending
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND NOT EXISTS (
|
||||
SELECT 1 FROM {table} AS processing
|
||||
WHERE processing.bank_id = pending.bank_id
|
||||
AND processing.operation_type = 'consolidation'
|
||||
AND processing.status = 'processing'
|
||||
)
|
||||
ORDER BY created_at
|
||||
LIMIT $1
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
limit,
|
||||
)
|
||||
else:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = $1
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
ORDER BY created_at
|
||||
LIMIT $2
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
op_type,
|
||||
limit,
|
||||
)
|
||||
tagged_rows = [(row, False) for row in non_consolidation_rows] + [
|
||||
(row, True) for row in consolidation_rows
|
||||
]
|
||||
|
||||
for row in rows:
|
||||
claimed_ids.append(row["operation_id"])
|
||||
all_rows.append(row)
|
||||
|
||||
# --- Phase 2: claim from shared pool ---
|
||||
remaining_shared = shared_limit
|
||||
if remaining_shared > 0:
|
||||
# 2a. Non-consolidation tasks (any type except consolidation)
|
||||
if claimed_ids:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type != 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND operation_id != ALL($1::uuid[])
|
||||
ORDER BY created_at
|
||||
LIMIT $2
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
claimed_ids,
|
||||
remaining_shared,
|
||||
)
|
||||
else:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type != 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
ORDER BY created_at
|
||||
LIMIT $1
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
remaining_shared,
|
||||
)
|
||||
|
||||
for row in rows:
|
||||
claimed_ids.append(row["operation_id"])
|
||||
all_rows.append(row)
|
||||
remaining_shared -= len(rows)
|
||||
|
||||
# 2b. Consolidation tasks (with bank-serialization constraint)
|
||||
if remaining_shared > 0:
|
||||
if claimed_ids:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table} AS pending
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND operation_id != ALL($1::uuid[])
|
||||
AND NOT EXISTS (
|
||||
SELECT 1 FROM {table} AS processing
|
||||
WHERE processing.bank_id = pending.bank_id
|
||||
AND processing.operation_type = 'consolidation'
|
||||
AND processing.status = 'processing'
|
||||
)
|
||||
ORDER BY created_at
|
||||
LIMIT $2
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
claimed_ids,
|
||||
remaining_shared,
|
||||
)
|
||||
else:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, operation_type, task_payload, retry_count
|
||||
FROM {table} AS pending
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND NOT EXISTS (
|
||||
SELECT 1 FROM {table} AS processing
|
||||
WHERE processing.bank_id = pending.bank_id
|
||||
AND processing.operation_type = 'consolidation'
|
||||
AND processing.status = 'processing'
|
||||
)
|
||||
ORDER BY created_at
|
||||
LIMIT $1
|
||||
FOR UPDATE SKIP LOCKED
|
||||
""",
|
||||
remaining_shared,
|
||||
)
|
||||
|
||||
for row in rows:
|
||||
claimed_ids.append(row["operation_id"])
|
||||
all_rows.append(row)
|
||||
|
||||
if not all_rows:
|
||||
if not tagged_rows:
|
||||
return []
|
||||
|
||||
operation_ids = [row["operation_id"] for row in all_rows]
|
||||
operation_ids = [row["operation_id"] for row, _ in tagged_rows]
|
||||
await conn.execute(
|
||||
f"""
|
||||
UPDATE {table}
|
||||
@@ -575,15 +312,15 @@ class WorkerPoller:
|
||||
)
|
||||
|
||||
result = []
|
||||
for row in all_rows:
|
||||
for row, is_consolidation in tagged_rows:
|
||||
task_dict = json.loads(row["task_payload"])
|
||||
task_dict["_retry_count"] = row["retry_count"]
|
||||
task_dict["_operation_id"] = str(row["operation_id"])
|
||||
# The DB column is authoritative for operation_type — inject it
|
||||
# into task_dict so in-flight tracking and slot accounting work.
|
||||
db_op_type = row["operation_type"]
|
||||
if db_op_type:
|
||||
task_dict["operation_type"] = db_op_type
|
||||
# The DB row knows the operation_type, but the JSON payload may not
|
||||
# carry it. Inject it so in-flight tracking and slot accounting
|
||||
# (which key off task_dict["operation_type"]) work correctly.
|
||||
if is_consolidation:
|
||||
task_dict["operation_type"] = "consolidation"
|
||||
result.append(
|
||||
ClaimedTask(
|
||||
operation_id=str(row["operation_id"]),
|
||||
@@ -724,25 +461,6 @@ class WorkerPoller:
|
||||
)
|
||||
logger.warning(f"Task {operation_id} scheduled for retry at {retry_at}: {error_message}")
|
||||
|
||||
async def _defer_operation(self, operation_id: str, exec_date: "Any", reason: str, schema: str | None):
|
||||
"""Reset task to pending for re-pickup at exec_date without counting as a retry.
|
||||
|
||||
Unlike `_schedule_retry`, this does not bump `retry_count` and does not
|
||||
populate `error_message` — defer is intentional backpressure, not a failure.
|
||||
"""
|
||||
table = fq_table("async_operations", schema)
|
||||
await self._pool.execute(
|
||||
f"""
|
||||
UPDATE {table}
|
||||
SET status = 'pending', next_retry_at = $2, worker_id = NULL, claimed_at = NULL,
|
||||
updated_at = now()
|
||||
WHERE operation_id = $1
|
||||
""",
|
||||
operation_id,
|
||||
exec_date,
|
||||
)
|
||||
logger.info(f"Task {operation_id} deferred until {exec_date}: {reason}")
|
||||
|
||||
async def execute_task(self, task: ClaimedTask):
|
||||
"""Execute a single task as a background job (fire-and-forget)."""
|
||||
task_type = task.task_dict.get("type", "unknown")
|
||||
@@ -814,8 +532,6 @@ class WorkerPoller:
|
||||
task.task_dict["_schema"] = task.schema
|
||||
await self._executor(task.task_dict)
|
||||
logger.debug(f"Task {task.operation_id} execution finished")
|
||||
except DeferOperation as e:
|
||||
await self._defer_operation(task.operation_id, e.exec_date, e.reason, task.schema)
|
||||
except RetryTaskAt as e:
|
||||
await self._schedule_retry(task.operation_id, e.retry_at, str(e), task.schema)
|
||||
except Exception as e:
|
||||
@@ -954,13 +670,9 @@ class WorkerPoller:
|
||||
"""
|
||||
await self.recover_own_tasks()
|
||||
|
||||
reservations_str = (
|
||||
", ".join(f"{k}={v}" for k, v in self._slot_reservations.items()) if self._slot_reservations else "none"
|
||||
)
|
||||
shared_pool = max(0, self._max_slots - sum(self._slot_reservations.values()))
|
||||
logger.info(
|
||||
f"Worker {self._worker_id} starting polling loop "
|
||||
f"(max_slots={self._max_slots}, reservations=[{reservations_str}], shared_pool={shared_pool})"
|
||||
f"(max_slots={self._max_slots}, consolidation_max_slots={self._consolidation_max_slots})"
|
||||
)
|
||||
|
||||
while not self._shutdown.is_set():
|
||||
@@ -1079,19 +791,11 @@ class WorkerPoller:
|
||||
in_flight_by_type = dict(self._in_flight_by_type)
|
||||
active_tasks = dict(self._active_tasks)
|
||||
|
||||
# Compute per-type reserved availability and shared pool
|
||||
tasks_in_reserved = 0
|
||||
reserved_parts = []
|
||||
for op_type, reserved in self._slot_reservations.items():
|
||||
type_in_flight = in_flight_by_type.get(op_type, 0)
|
||||
type_available = max(0, reserved - type_in_flight)
|
||||
tasks_in_reserved += min(reserved, type_in_flight)
|
||||
reserved_parts.append(f"{op_type}={type_in_flight}/{reserved}(avail={type_available})")
|
||||
sum_reservations = sum(self._slot_reservations.values())
|
||||
shared_pool_size = max(0, self._max_slots - sum_reservations)
|
||||
tasks_in_shared = max(0, in_flight - tasks_in_reserved)
|
||||
shared_available = max(0, shared_pool_size - tasks_in_shared)
|
||||
reserved_str = ", ".join(reserved_parts) if reserved_parts else "none"
|
||||
consolidation_count = in_flight_by_type.get("consolidation", 0)
|
||||
non_consolidation_in_flight = max(0, in_flight - consolidation_count)
|
||||
non_consolidation_max = max(0, self._max_slots - self._consolidation_max_slots)
|
||||
available_slots = max(0, non_consolidation_max - non_consolidation_in_flight)
|
||||
available_consolidation_slots = max(0, self._consolidation_max_slots - consolidation_count)
|
||||
|
||||
# Build local processing breakdown (aggregate counts)
|
||||
task_groups: dict[tuple[str, str], int] = {}
|
||||
@@ -1108,42 +812,13 @@ class WorkerPoller:
|
||||
schemas = await self._get_schemas()
|
||||
global_pending = 0
|
||||
all_worker_counts: dict[str, int] = {}
|
||||
# operation_type -> aggregated bucket counts across schemas
|
||||
pending_breakdown: dict[str, dict[str, int]] = {}
|
||||
|
||||
async with self._pool.acquire() as conn:
|
||||
for schema in schemas:
|
||||
table = fq_table("async_operations", schema)
|
||||
|
||||
# Bucket pending rows by the same predicates the claim query
|
||||
# filters on, so an operator can see why pending > 0 but
|
||||
# nothing is being claimed (orphaned batch_retain parents,
|
||||
# retry backoff, etc.).
|
||||
breakdown_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT
|
||||
operation_type,
|
||||
COUNT(*) AS total,
|
||||
COUNT(*) FILTER (WHERE task_payload IS NULL) AS payload_null,
|
||||
COUNT(*) FILTER (
|
||||
WHERE next_retry_at IS NOT NULL AND next_retry_at > now()
|
||||
) AS retry_blocked,
|
||||
COUNT(*) FILTER (WHERE worker_id IS NOT NULL) AS assigned
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
GROUP BY operation_type
|
||||
"""
|
||||
)
|
||||
for br in breakdown_rows:
|
||||
op_type = br["operation_type"] or "unknown"
|
||||
bucket = pending_breakdown.setdefault(
|
||||
op_type, {"total": 0, "payload_null": 0, "retry_blocked": 0, "assigned": 0}
|
||||
)
|
||||
bucket["total"] += br["total"]
|
||||
bucket["payload_null"] += br["payload_null"]
|
||||
bucket["retry_blocked"] += br["retry_blocked"]
|
||||
bucket["assigned"] += br["assigned"]
|
||||
global_pending += br["total"]
|
||||
row = await conn.fetchrow(f"SELECT COUNT(*) as count FROM {table} WHERE status = 'pending'")
|
||||
global_pending += row["count"] if row else 0
|
||||
|
||||
worker_rows = await conn.fetch(
|
||||
f"""
|
||||
@@ -1172,9 +847,8 @@ class WorkerPoller:
|
||||
schemas_str = ", ".join(s if s else "default" for s in schemas)
|
||||
logger.info(
|
||||
f"[WORKER_STATS] worker={self._worker_id} "
|
||||
f"slots={in_flight}/{self._max_slots} | "
|
||||
f"reserved: [{reserved_str}] | "
|
||||
f"shared={tasks_in_shared}/{shared_pool_size}(avail={shared_available}) | "
|
||||
f"slots={in_flight}/{self._max_slots} (consolidation={consolidation_count}/{self._consolidation_max_slots}) | "
|
||||
f"available={available_slots} (consolidation={available_consolidation_slots}) | "
|
||||
f"global: pending={global_pending} (schemas: {schemas_str}) | "
|
||||
f"others: {others_str} | "
|
||||
f"pool: {pool_str} | "
|
||||
@@ -1182,13 +856,6 @@ class WorkerPoller:
|
||||
f"my_active: {processing_str}"
|
||||
)
|
||||
|
||||
# Pending breakdown - explains why pending rows aren't being claimed
|
||||
# (orphaned batch_retain parents have payload_null > 0, retry storms
|
||||
# show up as retry_blocked, etc.). Skip when nothing is pending so
|
||||
# the line doesn't add noise on idle deployments.
|
||||
if global_pending > 0:
|
||||
self._log_pending_breakdown(pending_breakdown)
|
||||
|
||||
# Per-task lines, sorted oldest-first so stuck tasks bubble to the top.
|
||||
self._log_per_task_lines(active_tasks, now=time.monotonic())
|
||||
|
||||
@@ -1228,14 +895,7 @@ class WorkerPoller:
|
||||
min_size = pool.get_min_size() if hasattr(pool, "get_min_size") else None
|
||||
max_size = pool.get_max_size() if hasattr(pool, "get_max_size") else None
|
||||
queue = getattr(pool, "_queue", None)
|
||||
# asyncpg's _queue is a LifoQueue pre-filled to max_size with
|
||||
# PoolConnectionHolder objects. qsize() therefore counts *available
|
||||
# holders*, not callers waiting on the pool — the previous "waiters"
|
||||
# label here was the opposite of what it suggested. The actual count
|
||||
# of awaiters is len(_queue._getters), nonzero only when qsize()==0.
|
||||
free_holders = queue.qsize() if queue is not None and hasattr(queue, "qsize") else None
|
||||
getters = getattr(queue, "_getters", None) if queue is not None else None
|
||||
pending_acquires = len(getters) if getters is not None else None
|
||||
waiters = queue.qsize() if queue is not None and hasattr(queue, "qsize") else None
|
||||
|
||||
parts = [f"size={size}"]
|
||||
if min_size is not None and max_size is not None:
|
||||
@@ -1243,44 +903,13 @@ class WorkerPoller:
|
||||
if free is not None:
|
||||
parts.append(f"idle={free}")
|
||||
parts.append(f"in_use={size - free}")
|
||||
if free_holders is not None:
|
||||
parts.append(f"free_holders={free_holders}")
|
||||
if pending_acquires is not None:
|
||||
parts.append(f"pending_acquires={pending_acquires}")
|
||||
if waiters is not None:
|
||||
parts.append(f"waiters={waiters}")
|
||||
return " ".join(parts)
|
||||
except Exception as e:
|
||||
logger.debug(f"Pool stats unavailable: {e}")
|
||||
return "unavailable"
|
||||
|
||||
def _log_pending_breakdown(self, breakdown: dict[str, dict[str, int]]) -> None:
|
||||
"""Emit one [PENDING_BREAKDOWN] line bucketing pending rows by claimability.
|
||||
|
||||
Each bucket mirrors a predicate in the claim query:
|
||||
* payload_null - row has no task_payload (e.g. batch_retain parent
|
||||
whose reconciliation never fired); claim query
|
||||
skips it forever
|
||||
* retry_blocked - next_retry_at is still in the future
|
||||
* assigned - worker_id already set; another worker owns it
|
||||
|
||||
``claimable`` is the residual that *should* be picked up on the next
|
||||
poll. If ``claimable > 0`` while workers report free slots, the bug is
|
||||
somewhere else (lock contention, tenant discovery, etc.) - this line
|
||||
narrows the search.
|
||||
"""
|
||||
if not breakdown:
|
||||
return
|
||||
|
||||
parts = []
|
||||
for op_type in sorted(breakdown):
|
||||
b = breakdown[op_type]
|
||||
claimable = b["total"] - b["payload_null"] - b["retry_blocked"] - b["assigned"]
|
||||
parts.append(
|
||||
f"{op_type}: total={b['total']} claimable={claimable} "
|
||||
f"payload_null={b['payload_null']} retry_blocked={b['retry_blocked']} "
|
||||
f"assigned={b['assigned']}"
|
||||
)
|
||||
logger.info(f"[PENDING_BREAKDOWN] {' | '.join(parts)}")
|
||||
|
||||
def _log_per_task_lines(self, active_tasks: dict[str, ActiveTaskInfo], now: float) -> None:
|
||||
"""Emit one [WORKER_TASK] line per in-flight task and dump stuck stacks.
|
||||
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "hindsight-api-slim"
|
||||
version = "0.5.4"
|
||||
version = "0.5.1"
|
||||
description = "Hindsight: Agent Memory That Works Like Human Memory"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.11"
|
||||
@@ -73,11 +73,9 @@ local-ml = [
|
||||
"torch>=2.6.0", # CVE fix for remote code execution
|
||||
"einops>=0.8.2",
|
||||
"flashrank>=0.2.0",
|
||||
# Apple Silicon local inference — mlx publishes wheels only for
|
||||
# macOS/Linux, not Windows, so gate on platform to let `uv sync
|
||||
# --all-extras` resolve on win_amd64 runners.
|
||||
"mlx>=0.31.0; sys_platform != 'win32'",
|
||||
"mlx-lm>=0.31.1; sys_platform != 'win32'",
|
||||
# Apple Silicon local inference
|
||||
"mlx>=0.31.0",
|
||||
"mlx-lm>=0.31.1",
|
||||
"safetensors>=0.6.2",
|
||||
]
|
||||
local-llm = [
|
||||
@@ -86,7 +84,7 @@ local-llm = [
|
||||
"huggingface-hub>=0.20.0",
|
||||
]
|
||||
embedded-db = [
|
||||
"pg0-embedded>=0.13.0",
|
||||
"pg0-embedded>=0.11.0",
|
||||
]
|
||||
all = [
|
||||
"hindsight-api-slim[local-ml,embedded-db]",
|
||||
|
||||
@@ -1,47 +0,0 @@
|
||||
"""Graph-level sanity checks for the Alembic migration DAG.
|
||||
|
||||
These tests do not touch a database; they only parse the revision files on
|
||||
disk, so they are cheap to run in CI and catch DAG accidents (divergent
|
||||
heads, unreachable revisions) at merge time instead of at deploy time.
|
||||
"""
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
from alembic.config import Config
|
||||
from alembic.script import ScriptDirectory
|
||||
|
||||
|
||||
def _script_directory() -> ScriptDirectory:
|
||||
cfg = Config()
|
||||
script_location = Path(__file__).parent.parent / "hindsight_api" / "alembic"
|
||||
cfg.set_main_option("script_location", str(script_location))
|
||||
return ScriptDirectory.from_config(cfg)
|
||||
|
||||
|
||||
def test_single_head() -> None:
|
||||
"""The DAG must have exactly one head.
|
||||
|
||||
A second head means a branch was added without a merge revision, which
|
||||
makes ``alembic upgrade head`` (singular) ambiguous and forces the next
|
||||
migration author to orphan whichever head they don't pick as parent.
|
||||
v0.5.3 shipped in exactly that state; this test would have caught it.
|
||||
|
||||
Fix for a new head: ``alembic merge heads -m "<reason>"``.
|
||||
"""
|
||||
script = _script_directory()
|
||||
heads = script.get_heads()
|
||||
assert len(heads) == 1, (
|
||||
f"Alembic has {len(heads)} heads ({heads}); expected exactly 1. "
|
||||
"Unify them with ``alembic merge heads -m '<reason>'``."
|
||||
)
|
||||
|
||||
|
||||
def test_single_base() -> None:
|
||||
"""The DAG must have exactly one base (the initial schema).
|
||||
|
||||
Multiple bases mean disconnected migration trees, which can only happen
|
||||
through manual file edits.
|
||||
"""
|
||||
script = _script_directory()
|
||||
bases = script.get_bases()
|
||||
assert len(bases) == 1, f"Alembic has {len(bases)} bases ({bases}); expected exactly 1."
|
||||
@@ -8,13 +8,6 @@ import pytest
|
||||
|
||||
from hindsight_api.extensions import RequestContext
|
||||
|
||||
# These tests submit async operations and rely on the engine-owned worker to
|
||||
# drain them. test_worker.py drives its own WorkerPoller.claim_batch() against
|
||||
# the same pool, so running the two files on different xdist workers causes
|
||||
# them to steal each other's pending rows. Share the "worker_tests" group so
|
||||
# they serialize on the same xdist process.
|
||||
pytestmark = pytest.mark.xdist_group("worker_tests")
|
||||
|
||||
|
||||
async def _ensure_bank(pool, bank_id: str) -> None:
|
||||
"""Upsert a minimal bank row so FK on async_operations passes."""
|
||||
@@ -440,392 +433,3 @@ async def test_config_retain_batch_tokens_respected(memory, request_context):
|
||||
# Even small batches use parent-child pattern now (simpler code path)
|
||||
assert "child_operations" in status
|
||||
assert status["result_metadata"]["num_sub_batches"] == 1
|
||||
|
||||
|
||||
async def _child_metadata(memory, bank_id: str, parent_operation_id: str, request_context):
|
||||
"""Fetch the first child operation's result_metadata for a parent batch_retain."""
|
||||
parent = await memory.get_operation_status(
|
||||
bank_id=bank_id,
|
||||
operation_id=parent_operation_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
assert parent["status"] == "completed", parent
|
||||
assert parent["child_operations"], "expected at least one child operation"
|
||||
child_id = parent["child_operations"][0]["operation_id"]
|
||||
child = await memory.get_operation_status(
|
||||
bank_id=bank_id,
|
||||
operation_id=child_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
return child["result_metadata"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retain_records_user_provided_document_ids(memory, request_context):
|
||||
"""User-supplied document_ids land in child op result_metadata.document_ids."""
|
||||
bank_id = "test_doc_ids_user_supplied"
|
||||
d1 = str(uuid.uuid4())
|
||||
d2 = str(uuid.uuid4())
|
||||
contents = [
|
||||
{"content": "User-supplied doc one content.", "document_id": d1},
|
||||
{"content": "User-supplied doc two content.", "document_id": d2},
|
||||
]
|
||||
|
||||
result = await memory.submit_async_retain(
|
||||
bank_id=bank_id,
|
||||
contents=contents,
|
||||
request_context=request_context,
|
||||
)
|
||||
await asyncio.sleep(0.2)
|
||||
|
||||
meta = await _child_metadata(memory, bank_id, result["operation_id"], request_context)
|
||||
assert "document_ids" in meta, meta
|
||||
assert set(meta["document_ids"]) == {d1, d2}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retain_records_generated_document_id(memory, request_context):
|
||||
"""With no document_ids supplied, retain records the single generated id."""
|
||||
bank_id = "test_doc_ids_generated"
|
||||
contents = [
|
||||
{"content": "Generated doc item one."},
|
||||
{"content": "Generated doc item two."},
|
||||
]
|
||||
|
||||
result = await memory.submit_async_retain(
|
||||
bank_id=bank_id,
|
||||
contents=contents,
|
||||
request_context=request_context,
|
||||
)
|
||||
await asyncio.sleep(0.2)
|
||||
|
||||
meta = await _child_metadata(memory, bank_id, result["operation_id"], request_context)
|
||||
assert "document_ids" in meta, meta
|
||||
assert isinstance(meta["document_ids"], list)
|
||||
assert len(meta["document_ids"]) == 1
|
||||
# Must be a valid UUID string (generated by the orchestrator)
|
||||
uuid.UUID(meta["document_ids"][0])
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retain_records_shared_document_id_once(memory, request_context):
|
||||
"""Items sharing one document_id record it exactly once (idempotent set-append)."""
|
||||
bank_id = "test_doc_ids_shared"
|
||||
shared = str(uuid.uuid4())
|
||||
# Duplicate per-item doc_ids are rejected up front, so shared-doc mode
|
||||
# is exercised by a single item carrying the id.
|
||||
contents = [{"content": "Shared doc, chunk A.", "document_id": shared}]
|
||||
|
||||
result = await memory.submit_async_retain(
|
||||
bank_id=bank_id,
|
||||
contents=contents,
|
||||
request_context=request_context,
|
||||
)
|
||||
await asyncio.sleep(0.2)
|
||||
|
||||
meta = await _child_metadata(memory, bank_id, result["operation_id"], request_context)
|
||||
assert meta.get("document_ids") == [shared]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_operation_status_include_payload(memory, request_context):
|
||||
"""include_payload=True returns the original submission payload; default omits it."""
|
||||
bank_id = "test_include_payload"
|
||||
contents = [{"content": "Payload roundtrip test item."}]
|
||||
|
||||
result = await memory.submit_async_retain(
|
||||
bank_id=bank_id,
|
||||
contents=contents,
|
||||
request_context=request_context,
|
||||
)
|
||||
await asyncio.sleep(0.2)
|
||||
|
||||
parent = await memory.get_operation_status(
|
||||
bank_id=bank_id,
|
||||
operation_id=result["operation_id"],
|
||||
request_context=request_context,
|
||||
)
|
||||
child_id = parent["child_operations"][0]["operation_id"]
|
||||
|
||||
# Default: no payload
|
||||
without = await memory.get_operation_status(
|
||||
bank_id=bank_id,
|
||||
operation_id=child_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
assert without.get("task_payload") is None
|
||||
|
||||
# With flag: payload populated
|
||||
with_payload = await memory.get_operation_status(
|
||||
bank_id=bank_id,
|
||||
operation_id=child_id,
|
||||
request_context=request_context,
|
||||
include_payload=True,
|
||||
)
|
||||
payload = with_payload.get("task_payload")
|
||||
assert payload is not None, with_payload
|
||||
assert payload.get("bank_id") == bank_id
|
||||
assert payload.get("contents")
|
||||
assert payload["contents"][0]["content"] == "Payload roundtrip test item."
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_operation_status_exposes_retry_count_and_next_retry_at(memory, request_context):
|
||||
"""get_operation_status and list_operations return retry_count and next_retry_at.
|
||||
|
||||
Consumers need these to distinguish a freshly-queued pending task from
|
||||
one that's parked for a future retry (e.g. because an extension raised
|
||||
DeferOperation). Without them, "pending" is ambiguous and callers can't
|
||||
render a helpful "deferred until X" state.
|
||||
"""
|
||||
from datetime import datetime, timedelta, timezone
|
||||
|
||||
bank_id = "test_retry_fields"
|
||||
result = await memory.submit_async_retain(
|
||||
bank_id=bank_id,
|
||||
contents=[{"content": "retry-fields test item"}],
|
||||
request_context=request_context,
|
||||
)
|
||||
await asyncio.sleep(0.1)
|
||||
parent_id = result["operation_id"]
|
||||
child_id = None
|
||||
|
||||
# Get the child op (the batch_retain parent holds a single child in the
|
||||
# sync/simplified path used by SyncTaskBackend tests).
|
||||
parent_status = await memory.get_operation_status(
|
||||
bank_id=bank_id,
|
||||
operation_id=parent_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
assert "retry_count" in parent_status
|
||||
assert "next_retry_at" in parent_status
|
||||
assert parent_status["retry_count"] == 0
|
||||
# Completed tasks should have next_retry_at cleared on the row (or the
|
||||
# status field doesn't include it meaningfully), so we don't assert a
|
||||
# specific value here — only that the key is present.
|
||||
if parent_status.get("child_operations"):
|
||||
child_id = parent_status["child_operations"][0]["operation_id"]
|
||||
|
||||
# list_operations also exposes both fields
|
||||
listed = await memory.list_operations(
|
||||
bank_id=bank_id,
|
||||
request_context=request_context,
|
||||
limit=10,
|
||||
offset=0,
|
||||
)
|
||||
assert listed["operations"], listed
|
||||
for op in listed["operations"]:
|
||||
assert "retry_count" in op
|
||||
assert "next_retry_at" in op
|
||||
assert isinstance(op["retry_count"], int)
|
||||
|
||||
# Simulate a deferred op: set next_retry_at to 15 min in the future for
|
||||
# the child row directly in the DB, then fetch via the API and confirm
|
||||
# the value round-trips as an ISO-8601 string.
|
||||
if child_id:
|
||||
pool = await memory._get_pool()
|
||||
future = datetime.now(timezone.utc) + timedelta(minutes=15)
|
||||
await pool.execute(
|
||||
"UPDATE async_operations SET status = 'pending', next_retry_at = $1, retry_count = 2 WHERE operation_id = $2",
|
||||
future,
|
||||
uuid.UUID(child_id),
|
||||
)
|
||||
fetched = await memory.get_operation_status(
|
||||
bank_id=bank_id,
|
||||
operation_id=child_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
assert fetched["retry_count"] == 2
|
||||
assert fetched["next_retry_at"] is not None
|
||||
# Round-trip tolerance: within 1 second.
|
||||
parsed = datetime.fromisoformat(fetched["next_retry_at"])
|
||||
assert abs((parsed - future).total_seconds()) < 1.0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_operations_exclude_parents(memory, request_context):
|
||||
"""list_operations with exclude_parents=True hides parent batch operations."""
|
||||
bank_id = "test_exclude_parents"
|
||||
pool = await memory._get_pool()
|
||||
await _ensure_bank(pool, bank_id)
|
||||
|
||||
# Create a parent operation (is_parent=True)
|
||||
parent_id = uuid.uuid4()
|
||||
child_id = uuid.uuid4()
|
||||
standalone_id = uuid.uuid4()
|
||||
|
||||
async with pool.acquire() as conn:
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO async_operations (operation_id, bank_id, operation_type, result_metadata, status)
|
||||
VALUES ($1, $2, $3, $4, $5)
|
||||
""",
|
||||
parent_id,
|
||||
bank_id,
|
||||
"batch_retain",
|
||||
json.dumps({"items_count": 10, "num_sub_batches": 1, "is_parent": True}),
|
||||
"completed",
|
||||
)
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO async_operations (operation_id, bank_id, operation_type, result_metadata, status)
|
||||
VALUES ($1, $2, $3, $4, $5)
|
||||
""",
|
||||
child_id,
|
||||
bank_id,
|
||||
"retain",
|
||||
json.dumps(
|
||||
{"items_count": 10, "parent_operation_id": str(parent_id), "sub_batch_index": 1, "total_sub_batches": 1}
|
||||
),
|
||||
"completed",
|
||||
)
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO async_operations (operation_id, bank_id, operation_type, result_metadata, status)
|
||||
VALUES ($1, $2, $3, $4, $5)
|
||||
""",
|
||||
standalone_id,
|
||||
bank_id,
|
||||
"consolidation",
|
||||
json.dumps({}),
|
||||
"completed",
|
||||
)
|
||||
|
||||
# Without exclude_parents: all 3 operations visible
|
||||
all_ops = await memory.list_operations(
|
||||
bank_id=bank_id,
|
||||
request_context=request_context,
|
||||
limit=10,
|
||||
offset=0,
|
||||
)
|
||||
all_ids = {op["id"] for op in all_ops["operations"]}
|
||||
assert str(parent_id) in all_ids
|
||||
assert str(child_id) in all_ids
|
||||
assert str(standalone_id) in all_ids
|
||||
assert all_ops["total"] == 3
|
||||
|
||||
# With exclude_parents: parent is hidden
|
||||
filtered_ops = await memory.list_operations(
|
||||
bank_id=bank_id,
|
||||
request_context=request_context,
|
||||
limit=10,
|
||||
offset=0,
|
||||
exclude_parents=True,
|
||||
)
|
||||
filtered_ids = {op["id"] for op in filtered_ops["operations"]}
|
||||
assert str(parent_id) not in filtered_ids
|
||||
assert str(child_id) in filtered_ids
|
||||
assert str(standalone_id) in filtered_ids
|
||||
assert filtered_ops["total"] == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_request_context_retry_count_propagated_to_validator(memory_no_llm_verify, request_context):
|
||||
"""_handle_batch_retain forwards the task's _retry_count as
|
||||
RequestContext.retry_count, so validator extensions can compute
|
||||
exponential backoff without querying async_operations themselves.
|
||||
"""
|
||||
from hindsight_api.extensions import (
|
||||
OperationValidatorExtension,
|
||||
RecallContext,
|
||||
ReflectContext,
|
||||
RetainContext,
|
||||
ValidationResult,
|
||||
)
|
||||
|
||||
captured: dict[str, int] = {"retry_count": -1}
|
||||
|
||||
class CapturingValidator(OperationValidatorExtension):
|
||||
def __init__(self):
|
||||
super().__init__({})
|
||||
|
||||
async def validate_retain(self, ctx: RetainContext) -> ValidationResult:
|
||||
captured["retry_count"] = ctx.request_context.retry_count
|
||||
return ValidationResult.accept()
|
||||
|
||||
async def validate_recall(self, ctx: RecallContext) -> ValidationResult:
|
||||
return ValidationResult.accept()
|
||||
|
||||
async def validate_reflect(self, ctx: ReflectContext) -> ValidationResult:
|
||||
return ValidationResult.accept()
|
||||
|
||||
memory_no_llm_verify._operation_validator = CapturingValidator()
|
||||
|
||||
bank_id = f"test-retry-propagate-{uuid.uuid4().hex[:8]}"
|
||||
pool = await memory_no_llm_verify._get_pool()
|
||||
await _ensure_bank(pool, bank_id)
|
||||
|
||||
task_dict = {
|
||||
"type": "batch_retain",
|
||||
"bank_id": bank_id,
|
||||
"contents": [{"content": "retry-propagate test"}],
|
||||
"_tenant_id": "default",
|
||||
"_retry_count": 3, # simulate 3rd retry
|
||||
}
|
||||
await memory_no_llm_verify._handle_batch_retain(task_dict)
|
||||
|
||||
assert captured["retry_count"] == 3, (
|
||||
f"Validator should see retry_count=3 from task_dict['_retry_count']; got {captured['retry_count']}"
|
||||
)
|
||||
|
||||
# Default (missing _retry_count key) must surface as 0, not raise.
|
||||
captured["retry_count"] = -1
|
||||
task_dict_no_retry = {
|
||||
"type": "batch_retain",
|
||||
"bank_id": bank_id,
|
||||
"contents": [{"content": "retry-propagate default test"}],
|
||||
"_tenant_id": "default",
|
||||
}
|
||||
await memory_no_llm_verify._handle_batch_retain(task_dict_no_retry)
|
||||
assert captured["retry_count"] == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_submit_async_operation_leaves_claimable_row_when_submit_task_fails(memory):
|
||||
"""Regression for the crash-window orphan bug fixed in #1091.
|
||||
|
||||
Previously, _submit_async_operation INSERTed the async_operations row without
|
||||
task_payload, then called submit_task as a separate step to fill it in. If
|
||||
submit_task failed (crash, timeout, dropped connection) after the INSERT
|
||||
committed, the row was left with task_payload IS NULL and became permanently
|
||||
stuck because the worker claim query filters on task_payload IS NOT NULL.
|
||||
|
||||
With the atomic INSERT, even if submit_task raises afterwards the row is born
|
||||
claimable. This test simulates the crash by forcing submit_task to raise.
|
||||
"""
|
||||
bank_id = f"test_orphan_prevention_{uuid.uuid4().hex[:8]}"
|
||||
pool = await memory._get_pool()
|
||||
await _ensure_bank(pool, bank_id)
|
||||
|
||||
async def failing_submit_task(_task_dict):
|
||||
raise RuntimeError("Simulated crash between INSERT and submit_task")
|
||||
|
||||
memory._task_backend.submit_task = failing_submit_task # type: ignore[method-assign]
|
||||
|
||||
with pytest.raises(RuntimeError, match="Simulated crash"):
|
||||
await memory._submit_async_operation(
|
||||
bank_id=bank_id,
|
||||
operation_type="retain",
|
||||
task_type="batch_retain",
|
||||
task_payload={"contents": [{"content": "hello", "document_id": "d1"}]},
|
||||
)
|
||||
|
||||
rows = await pool.fetch(
|
||||
"""
|
||||
SELECT status, task_payload
|
||||
FROM async_operations
|
||||
WHERE bank_id = $1 AND operation_type = 'retain'
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
assert len(rows) == 1, f"Expected exactly one retain row for bank_id={bank_id}, got {len(rows)}"
|
||||
row = rows[0]
|
||||
assert row["status"] == "pending"
|
||||
assert row["task_payload"] is not None, (
|
||||
"task_payload must be set atomically by the INSERT — a NULL here means "
|
||||
"the worker claim query (task_payload IS NOT NULL) will never pick this row up"
|
||||
)
|
||||
payload = json.loads(row["task_payload"])
|
||||
assert payload["type"] == "batch_retain"
|
||||
assert payload["bank_id"] == bank_id
|
||||
assert payload["contents"] == [{"content": "hello", "document_id": "d1"}]
|
||||
|
||||
@@ -1,248 +0,0 @@
|
||||
"""
|
||||
Tests for the bank stats endpoint and the memories-timeseries endpoint.
|
||||
|
||||
Covers the new fields exposed by GET /v1/default/banks/{bank_id}/stats
|
||||
(operations_by_status) and the new endpoint
|
||||
GET /v1/default/banks/{bank_id}/stats/memories-timeseries.
|
||||
"""
|
||||
import uuid
|
||||
from datetime import datetime
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
|
||||
from hindsight_api.api import create_app
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def api_client(memory):
|
||||
app = create_app(memory, initialize_memory=False)
|
||||
transport = httpx.ASGITransport(app=app)
|
||||
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
|
||||
yield client
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def test_bank_id():
|
||||
return f"stats_test_{datetime.now().timestamp()}"
|
||||
|
||||
|
||||
async def _insert_memory(memory, bank_id: str, text: str, *, failed: bool = False) -> str:
|
||||
"""Insert a single experience memory, optionally marked as consolidation-failed."""
|
||||
mem_id = uuid.uuid4()
|
||||
async with memory._pool.acquire() as conn:
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO memory_units (id, bank_id, text, fact_type, created_at, consolidation_failed_at)
|
||||
VALUES ($1, $2, $3, 'experience', now(), CASE WHEN $4 THEN now() ELSE NULL END)
|
||||
""",
|
||||
mem_id,
|
||||
bank_id,
|
||||
text,
|
||||
failed,
|
||||
)
|
||||
return str(mem_id)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bank_stats_exposes_operations_by_status(api_client, test_bank_id):
|
||||
"""/stats should return operations_by_status with all finished operations."""
|
||||
try:
|
||||
# Kick off a retain so at least one completed operation exists.
|
||||
response = await api_client.post(
|
||||
f"/v1/default/banks/{test_bank_id}/memories",
|
||||
json={"items": [{"content": "Alice is a software engineer.", "context": "team"}]},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
|
||||
response = await api_client.get(f"/v1/default/banks/{test_bank_id}/stats")
|
||||
assert response.status_code == 200
|
||||
stats = response.json()
|
||||
|
||||
assert "operations_by_status" in stats
|
||||
assert isinstance(stats["operations_by_status"], dict)
|
||||
# A synchronous retain finishes as "completed".
|
||||
assert stats["operations_by_status"].get("completed", 0) >= 1
|
||||
# pending/failed counters should still be present as scalar mirrors.
|
||||
assert stats["pending_operations"] == stats["operations_by_status"].get("pending", 0)
|
||||
assert stats["failed_operations"] == stats["operations_by_status"].get("failed", 0)
|
||||
finally:
|
||||
await api_client.delete(f"/v1/default/banks/{test_bank_id}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"period,expected_count,expected_trunc",
|
||||
[
|
||||
("1h", 60, "minute"),
|
||||
("12h", 12, "hour"),
|
||||
("1d", 24, "hour"),
|
||||
("7d", 7, "day"),
|
||||
("30d", 30, "day"),
|
||||
("90d", 90, "day"),
|
||||
],
|
||||
)
|
||||
async def test_memories_timeseries_periods(
|
||||
api_client, test_bank_id, period, expected_count, expected_trunc
|
||||
):
|
||||
"""Every period must return the full expected bucket count and trunc."""
|
||||
try:
|
||||
response = await api_client.post(
|
||||
f"/v1/default/banks/{test_bank_id}/memories",
|
||||
json={"items": [{"content": "Bob works on infrastructure.", "context": "team"}]},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
|
||||
response = await api_client.get(
|
||||
f"/v1/default/banks/{test_bank_id}/stats/memories-timeseries",
|
||||
params={"period": period},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
body = response.json()
|
||||
|
||||
assert body["bank_id"] == test_bank_id
|
||||
assert body["period"] == period
|
||||
assert body["trunc"] == expected_trunc
|
||||
assert len(body["buckets"]) == expected_count
|
||||
|
||||
for bucket in body["buckets"]:
|
||||
assert "time" in bucket
|
||||
# Bucket `time` must serialize as a tz-aware ISO (ending in `+00:00` or `Z`).
|
||||
# A naive ISO (`2026-04-18T00:00:00`) would be parsed as local time by
|
||||
# `new Date()` per ECMA-262, shifting the chart by the browser's timezone.
|
||||
assert bucket["time"].endswith("+00:00") or bucket["time"].endswith("Z"), (
|
||||
f"bucket time must include UTC offset, got {bucket['time']!r}"
|
||||
)
|
||||
assert bucket["world"] >= 0
|
||||
assert bucket["experience"] >= 0
|
||||
assert bucket["observation"] >= 0
|
||||
finally:
|
||||
await api_client.delete(f"/v1/default/banks/{test_bank_id}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_memories_timeseries_invalid_period_falls_back(api_client, test_bank_id):
|
||||
"""An unknown period must fall back to the 7d default."""
|
||||
try:
|
||||
response = await api_client.get(
|
||||
f"/v1/default/banks/{test_bank_id}/stats/memories-timeseries",
|
||||
params={"period": "nonsense"},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
body = response.json()
|
||||
assert body["period"] == "7d"
|
||||
assert body["trunc"] == "day"
|
||||
assert len(body["buckets"]) == 7
|
||||
finally:
|
||||
await api_client.delete(f"/v1/default/banks/{test_bank_id}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_memories_timeseries_empty_bank_returns_zero_filled_buckets(
|
||||
api_client, test_bank_id
|
||||
):
|
||||
"""A bank with no memories must still return the full zero-filled bucket set."""
|
||||
try:
|
||||
# Ensure the bank exists.
|
||||
response = await api_client.get(f"/v1/default/banks/{test_bank_id}/profile")
|
||||
assert response.status_code == 200
|
||||
|
||||
response = await api_client.get(
|
||||
f"/v1/default/banks/{test_bank_id}/stats/memories-timeseries",
|
||||
params={"period": "7d"},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
body = response.json()
|
||||
assert len(body["buckets"]) == 7
|
||||
for bucket in body["buckets"]:
|
||||
assert bucket["world"] == 0
|
||||
assert bucket["experience"] == 0
|
||||
assert bucket["observation"] == 0
|
||||
finally:
|
||||
await api_client.delete(f"/v1/default/banks/{test_bank_id}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_memories_timeseries_reflects_retained_memories(api_client, test_bank_id):
|
||||
"""Freshly-retained memories must show up in today's bucket counts."""
|
||||
try:
|
||||
response = await api_client.post(
|
||||
f"/v1/default/banks/{test_bank_id}/memories",
|
||||
json={
|
||||
"items": [
|
||||
{"content": "Alice is a software engineer.", "context": "team"},
|
||||
{"content": "Bob works on infrastructure.", "context": "team"},
|
||||
]
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
|
||||
response = await api_client.get(
|
||||
f"/v1/default/banks/{test_bank_id}/stats/memories-timeseries",
|
||||
params={"period": "7d"},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
body = response.json()
|
||||
totals = sum(b["world"] + b["experience"] + b["observation"] for b in body["buckets"])
|
||||
assert totals >= 2, "expected at least two memories across all buckets"
|
||||
|
||||
# Those memories should land in the most-recent bucket.
|
||||
latest = body["buckets"][-1]
|
||||
assert latest["world"] + latest["experience"] + latest["observation"] >= 2
|
||||
finally:
|
||||
await api_client.delete(f"/v1/default/banks/{test_bank_id}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bank_stats_reports_failed_consolidation(api_client, memory, test_bank_id):
|
||||
"""/stats must surface the count of memories with consolidation_failed_at set."""
|
||||
try:
|
||||
await _insert_memory(memory, test_bank_id, "Alice failed 1.", failed=True)
|
||||
await _insert_memory(memory, test_bank_id, "Alice failed 2.", failed=True)
|
||||
await _insert_memory(memory, test_bank_id, "Alice pending.", failed=False)
|
||||
|
||||
response = await api_client.get(f"/v1/default/banks/{test_bank_id}/stats")
|
||||
assert response.status_code == 200
|
||||
stats = response.json()
|
||||
|
||||
assert stats["failed_consolidation"] == 2
|
||||
# The two failed memories also count as "not-yet-consolidated".
|
||||
assert stats["pending_consolidation"] >= 3
|
||||
finally:
|
||||
await api_client.delete(f"/v1/default/banks/{test_bank_id}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_memories_filter_by_consolidation_state_failed(api_client, memory, test_bank_id):
|
||||
"""?consolidation_state=failed returns only memories with consolidation_failed_at set."""
|
||||
try:
|
||||
failed_id = await _insert_memory(memory, test_bank_id, "Broken item.", failed=True)
|
||||
await _insert_memory(memory, test_bank_id, "Healthy item.", failed=False)
|
||||
|
||||
response = await api_client.get(
|
||||
f"/v1/default/banks/{test_bank_id}/memories/list",
|
||||
params={"consolidation_state": "failed"},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
body = response.json()
|
||||
|
||||
ids = [item["id"] for item in body["items"]]
|
||||
assert failed_id in ids
|
||||
assert body["total"] == 1
|
||||
assert body["items"][0]["consolidation_failed_at"] is not None
|
||||
finally:
|
||||
await api_client.delete(f"/v1/default/banks/{test_bank_id}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_memories_filter_by_consolidation_state_rejects_unknown(api_client, test_bank_id):
|
||||
"""An invalid consolidation_state value must return a 400 (not 500)."""
|
||||
try:
|
||||
response = await api_client.get(
|
||||
f"/v1/default/banks/{test_bank_id}/memories/list",
|
||||
params={"consolidation_state": "bogus"},
|
||||
)
|
||||
assert response.status_code == 400
|
||||
finally:
|
||||
await api_client.delete(f"/v1/default/banks/{test_bank_id}")
|
||||
@@ -1,115 +0,0 @@
|
||||
"""Verify that BankTemplateConfig exposes every hierarchical field that
|
||||
_CONFIGURABLE_FIELDS already accepts at the engine layer.
|
||||
|
||||
This test guards the fix for the gap described in the upstream PR title
|
||||
"fix(bank-template): align BankTemplateConfig with _CONFIGURABLE_FIELDS".
|
||||
Each new field is POSTed through /v1/default/banks/{id}/import and then
|
||||
read back via the bank-config endpoint; assertion is that the applied
|
||||
value round-trips through the engine.
|
||||
|
||||
Runs via: uv run pytest tests/test_bank_template_configurable_fields.py -v
|
||||
|
||||
The api_client fixture (shared with tests/test_bank_templates.py) wraps
|
||||
create_app(memory, initialize_memory=False) in an httpx.ASGITransport
|
||||
with base_url http://test — in-process, no network, no tenant extension.
|
||||
Copy the fixture inline here so the test file does not depend on a
|
||||
conftest we do not ship in the patch.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
|
||||
from hindsight_api.api import create_app
|
||||
from hindsight_api.api.http import BankTemplateConfig
|
||||
|
||||
# Each tuple is (field_name, applied_value). Values chosen to differ
|
||||
# visibly from defaults so round-trip bugs surface.
|
||||
NEW_FIELDS: list[tuple[str, object]] = [
|
||||
("retain_default_strategy", "strategy-a"),
|
||||
("retain_strategies", {"strategy-a": {"mode": "concise", "max_tokens": 512}}),
|
||||
("retain_chunk_batch_size", 7),
|
||||
("mcp_enabled_tools", ["list_banks", "get_bank_profile"]),
|
||||
("consolidation_llm_batch_size", 11),
|
||||
("consolidation_source_facts_max_tokens", 2048),
|
||||
("consolidation_source_facts_max_tokens_per_observation", 256),
|
||||
("max_observations_per_scope", 13),
|
||||
("reflect_source_facts_max_tokens", 4096),
|
||||
("llm_gemini_safety_settings", [{"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_NONE"}]),
|
||||
("recall_budget_function", "adaptive"),
|
||||
("recall_budget_fixed_low", 50),
|
||||
("recall_budget_fixed_mid", 250),
|
||||
("recall_budget_fixed_high", 800),
|
||||
("recall_budget_adaptive_low", 0.05),
|
||||
("recall_budget_adaptive_mid", 0.1),
|
||||
("recall_budget_adaptive_high", 0.4),
|
||||
("recall_budget_min", 30),
|
||||
("recall_budget_max", 1500),
|
||||
]
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def api_client(memory):
|
||||
"""Matches the fixture in tests/test_bank_templates.py — in-process
|
||||
ASGI test client, no tenant extension, no auth."""
|
||||
app = create_app(memory, initialize_memory=False)
|
||||
transport = httpx.ASGITransport(app=app)
|
||||
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
|
||||
yield client
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def bank_id():
|
||||
return f"tmpl_config_{datetime.now().timestamp()}"
|
||||
|
||||
|
||||
def test_bank_template_config_declares_every_configurable_field():
|
||||
"""Pydantic-level guard: every field in NEW_FIELDS must be a declared
|
||||
attribute of BankTemplateConfig so get_config_updates() picks it up."""
|
||||
declared = set(BankTemplateConfig.model_fields.keys())
|
||||
missing = [name for name, _ in NEW_FIELDS if name not in declared]
|
||||
assert not missing, f"BankTemplateConfig missing fields: {missing}"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("field_name,applied_value", NEW_FIELDS, ids=[n for n, _ in NEW_FIELDS])
|
||||
async def test_new_field_round_trips_through_import(
|
||||
api_client: httpx.AsyncClient,
|
||||
bank_id: str,
|
||||
field_name: str,
|
||||
applied_value: object,
|
||||
):
|
||||
"""POST a minimal manifest with one new field set, then read bank
|
||||
config back and assert the value made it through.
|
||||
|
||||
Bank config response shape per upstream's test_import_applies_config:
|
||||
top-level keys are resolved hierarchical config; per-bank overrides
|
||||
live under config["overrides"][<field>]. Assert on the override slot.
|
||||
"""
|
||||
unique_bank_id = f"{bank_id}_{field_name}"
|
||||
manifest = {
|
||||
"version": "1",
|
||||
"bank": {field_name: applied_value},
|
||||
}
|
||||
|
||||
resp = await api_client.post(
|
||||
f"/v1/default/banks/{unique_bank_id}/import",
|
||||
json=manifest,
|
||||
)
|
||||
assert resp.status_code == 200, resp.text
|
||||
|
||||
# Read bank config back — field must reflect the applied value
|
||||
# under the "overrides" slot, matching upstream's own test shape.
|
||||
read = await api_client.get(f"/v1/default/banks/{unique_bank_id}/config")
|
||||
assert read.status_code == 200, read.text
|
||||
config = read.json()
|
||||
overrides = config.get("overrides", {})
|
||||
assert overrides.get(field_name) == applied_value, (
|
||||
f"round-trip mismatch for {field_name}: "
|
||||
f"sent {applied_value!r}, got {overrides.get(field_name)!r} "
|
||||
f"(full overrides: {overrides!r})"
|
||||
)
|
||||
@@ -416,7 +416,7 @@ async def test_worker_batch_recovery(memory, request_context):
|
||||
schema=schema,
|
||||
tenant_extension=tenant_extension,
|
||||
max_slots=5,
|
||||
slot_reservations={"consolidation": 2},
|
||||
consolidation_max_slots=2,
|
||||
)
|
||||
|
||||
# Run recovery
|
||||
|
||||
@@ -1515,7 +1515,6 @@ class TestHierarchicalRetrieval:
|
||||
async with memory._pool.acquire() as conn:
|
||||
query_embedding = memory.embeddings.encode(["What does John like?"])[0]
|
||||
mental_model_result = await tool_search_mental_models(
|
||||
memory_engine=memory,
|
||||
conn=conn,
|
||||
bank_id=bank_id,
|
||||
query="What does John like?",
|
||||
@@ -1577,7 +1576,6 @@ class TestHierarchicalRetrieval:
|
||||
async with memory._pool.acquire() as conn:
|
||||
query_embedding = memory.embeddings.encode(["Where does Sarah work?"])[0]
|
||||
mental_model_result = await tool_search_mental_models(
|
||||
memory_engine=memory,
|
||||
conn=conn,
|
||||
bank_id=bank_id,
|
||||
query="Where does Sarah work?",
|
||||
|
||||
@@ -1,100 +0,0 @@
|
||||
"""Tests for consolidation retry budget configurability (issue #1042)."""
|
||||
|
||||
import pytest
|
||||
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
from hindsight_api.engine.consolidation.consolidator import _consolidate_batch_with_llm
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_llm_config():
|
||||
llm = AsyncMock()
|
||||
response = MagicMock()
|
||||
response.creates = []
|
||||
response.updates = []
|
||||
response.deletes = []
|
||||
llm.call.return_value = response
|
||||
return llm
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_config():
|
||||
config = MagicMock()
|
||||
config.observations_mission = None
|
||||
config.consolidation_max_attempts = 3
|
||||
config.consolidation_llm_max_retries = None
|
||||
return config
|
||||
|
||||
|
||||
class TestConsolidationRetryBudget:
|
||||
@pytest.mark.asyncio
|
||||
async def test_config_is_required(self, mock_llm_config):
|
||||
"""Passing config=None raises — it's a programmer error, not a runtime fallback."""
|
||||
with pytest.raises(ValueError, match="config is required"):
|
||||
await _consolidate_batch_with_llm(
|
||||
llm_config=mock_llm_config,
|
||||
memories=[{"id": "m1", "text": "test"}],
|
||||
union_observations=[],
|
||||
union_source_facts={},
|
||||
config=None,
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_configurable_max_attempts(self, mock_llm_config, mock_config):
|
||||
"""consolidation_max_attempts controls the outer retry loop."""
|
||||
mock_config.consolidation_max_attempts = 5
|
||||
mock_llm_config.call.side_effect = RuntimeError("fail")
|
||||
result = await _consolidate_batch_with_llm(
|
||||
llm_config=mock_llm_config,
|
||||
memories=[{"id": "m1", "text": "test"}],
|
||||
union_observations=[],
|
||||
union_source_facts={},
|
||||
config=mock_config,
|
||||
)
|
||||
assert result.failed
|
||||
assert mock_llm_config.call.call_count == 5
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_max_retries_threaded_to_call(self, mock_llm_config, mock_config):
|
||||
"""consolidation_llm_max_retries is passed to llm_config.call()."""
|
||||
mock_config.consolidation_llm_max_retries = 3
|
||||
await _consolidate_batch_with_llm(
|
||||
llm_config=mock_llm_config,
|
||||
memories=[{"id": "m1", "text": "test"}],
|
||||
union_observations=[],
|
||||
union_source_facts={},
|
||||
config=mock_config,
|
||||
)
|
||||
assert mock_llm_config.call.call_args.kwargs.get("max_retries") == 3
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_max_retries_not_passed_when_none(self, mock_llm_config, mock_config):
|
||||
"""When consolidation_llm_max_retries is None, max_retries is not passed."""
|
||||
mock_config.consolidation_llm_max_retries = None
|
||||
await _consolidate_batch_with_llm(
|
||||
llm_config=mock_llm_config,
|
||||
memories=[{"id": "m1", "text": "test"}],
|
||||
union_observations=[],
|
||||
union_source_facts={},
|
||||
config=mock_config,
|
||||
)
|
||||
assert "max_retries" not in mock_llm_config.call.call_args.kwargs
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reduced_budget_limits_total_calls(self, mock_llm_config, mock_config):
|
||||
"""Setting both to low values caps total failure attempts."""
|
||||
mock_config.consolidation_max_attempts = 2
|
||||
mock_config.consolidation_llm_max_retries = 2
|
||||
mock_llm_config.call.side_effect = RuntimeError("upstream 503")
|
||||
result = await _consolidate_batch_with_llm(
|
||||
llm_config=mock_llm_config,
|
||||
memories=[{"id": "m1", "text": "test"}],
|
||||
union_observations=[],
|
||||
union_source_facts={},
|
||||
config=mock_config,
|
||||
)
|
||||
assert result.failed
|
||||
assert mock_llm_config.call.call_count == 2
|
||||
for call_args in mock_llm_config.call.call_args_list:
|
||||
assert call_args.kwargs.get("max_retries") == 2
|
||||
@@ -1,146 +0,0 @@
|
||||
"""Integration tests for consolidation_max_memories_per_round config."""
|
||||
|
||||
import uuid
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
from hindsight_api.config import _get_raw_config
|
||||
from hindsight_api.engine.consolidation.consolidator import run_consolidation_job
|
||||
from hindsight_api.engine.memory_engine import MemoryEngine
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def enable_observations():
|
||||
config = _get_raw_config()
|
||||
original = config.enable_observations
|
||||
config.enable_observations = True
|
||||
yield
|
||||
config.enable_observations = original
|
||||
|
||||
|
||||
def _make_config(**overrides):
|
||||
raw = _get_raw_config()
|
||||
return type(raw)(
|
||||
**{
|
||||
**{f: getattr(raw, f) for f in raw.__dataclass_fields__},
|
||||
**overrides,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_round_limit_caps_processed_memories(memory: MemoryEngine, request_context):
|
||||
"""When max_memories_per_round is set, consolidation processes at most that many memories
|
||||
and re-submits itself for the remaining backlog."""
|
||||
bank_id = f"test-round-limit-{uuid.uuid4().hex[:8]}"
|
||||
await memory.get_bank_profile(bank_id=bank_id, request_context=request_context)
|
||||
|
||||
# Disable consolidation during retain so we build up a backlog
|
||||
fake_config_no_obs = _make_config(enable_observations=False)
|
||||
|
||||
with patch.object(memory._config_resolver, "resolve_full_config", return_value=fake_config_no_obs):
|
||||
for i in range(6):
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content=f"Fact number {i}: The user enjoys activity {i} on weekends.",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Verify we have unconsolidated memories
|
||||
async with memory._pool.acquire() as conn:
|
||||
unconsolidated = await conn.fetchval(
|
||||
"""
|
||||
SELECT COUNT(*) FROM memory_units
|
||||
WHERE bank_id = $1 AND consolidated_at IS NULL
|
||||
AND consolidation_failed_at IS NULL AND fact_type IN ('experience', 'world')
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
assert unconsolidated >= 6, f"Expected at least 6 unconsolidated memories, got {unconsolidated}"
|
||||
|
||||
# Run consolidation with a round limit of 3
|
||||
round_limit = 3
|
||||
fake_config = _make_config(consolidation_max_memories_per_round=round_limit)
|
||||
|
||||
with (
|
||||
patch.object(memory._config_resolver, "resolve_full_config", return_value=fake_config),
|
||||
patch.object(memory, "submit_async_consolidation") as mock_requeue,
|
||||
):
|
||||
result = await run_consolidation_job(
|
||||
memory_engine=memory,
|
||||
bank_id=bank_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert result["status"] == "completed"
|
||||
assert result["memories_processed"] <= round_limit
|
||||
|
||||
# Must have re-queued consolidation for remaining work
|
||||
mock_requeue.assert_called_once_with(bank_id=bank_id, request_context=request_context)
|
||||
|
||||
# Mental model refresh should be skipped on intermediate round
|
||||
assert result.get("mental_models_refreshed", 0) == 0
|
||||
|
||||
# Verify some memories are still unconsolidated
|
||||
async with memory._pool.acquire() as conn:
|
||||
still_unconsolidated = await conn.fetchval(
|
||||
"""
|
||||
SELECT COUNT(*) FROM memory_units
|
||||
WHERE bank_id = $1 AND consolidated_at IS NULL
|
||||
AND consolidation_failed_at IS NULL AND fact_type IN ('experience', 'world')
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
assert still_unconsolidated > 0, "Some memories should still be unconsolidated after hitting round limit"
|
||||
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unlimited_round_processes_all(memory: MemoryEngine, request_context):
|
||||
"""When max_memories_per_round is 0 (unlimited), all memories are processed without re-queue."""
|
||||
bank_id = f"test-unlimited-{uuid.uuid4().hex[:8]}"
|
||||
await memory.get_bank_profile(bank_id=bank_id, request_context=request_context)
|
||||
|
||||
# Disable consolidation during retain
|
||||
fake_config_no_obs = _make_config(enable_observations=False)
|
||||
|
||||
with patch.object(memory._config_resolver, "resolve_full_config", return_value=fake_config_no_obs):
|
||||
for i in range(4):
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content=f"Fact {i}: The user visited city {i} last year.",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Run consolidation with unlimited round (0)
|
||||
fake_config = _make_config(consolidation_max_memories_per_round=0)
|
||||
|
||||
with (
|
||||
patch.object(memory._config_resolver, "resolve_full_config", return_value=fake_config),
|
||||
patch.object(memory, "submit_async_consolidation") as mock_requeue,
|
||||
):
|
||||
result = await run_consolidation_job(
|
||||
memory_engine=memory,
|
||||
bank_id=bank_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert result["status"] == "completed"
|
||||
# Should NOT re-queue
|
||||
mock_requeue.assert_not_called()
|
||||
|
||||
# All memories should be consolidated
|
||||
async with memory._pool.acquire() as conn:
|
||||
still_unconsolidated = await conn.fetchval(
|
||||
"""
|
||||
SELECT COUNT(*) FROM memory_units
|
||||
WHERE bank_id = $1 AND consolidated_at IS NULL
|
||||
AND consolidation_failed_at IS NULL AND fact_type IN ('experience', 'world')
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
assert still_unconsolidated == 0
|
||||
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
@@ -1,141 +0,0 @@
|
||||
"""Tests for ``hindsight_api.db_url.to_libpq_url``.
|
||||
|
||||
Covers backward compatibility (existing configs must pass through unchanged)
|
||||
and the two transformations needed to support external PostgreSQL deployments
|
||||
that use SQLAlchemy-style ``postgresql+asyncpg://...?ssl=require`` URLs:
|
||||
|
||||
1. strip the ``+asyncpg`` dialect suffix,
|
||||
2. rename the ``ssl=`` query parameter to ``sslmode=``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from hindsight_api.db_url import to_libpq_url
|
||||
|
||||
|
||||
class TestPassthrough:
|
||||
"""Inputs that must be returned unchanged — protects existing configs."""
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"url",
|
||||
[
|
||||
"pg0",
|
||||
"",
|
||||
"postgresql://user:pass@host:5432/db",
|
||||
"postgresql://user:pass@host:5432/db?sslmode=require",
|
||||
"postgresql://user:pass@host/db?sslmode=verify-full&connect_timeout=10",
|
||||
"sqlite:///./test.db",
|
||||
"postgresql+psycopg2://user:pass@host/db",
|
||||
],
|
||||
)
|
||||
def test_unchanged(self, url: str) -> None:
|
||||
assert to_libpq_url(url) == url
|
||||
|
||||
|
||||
class TestSchemeNormalization:
|
||||
def test_asyncpg_scheme_stripped(self) -> None:
|
||||
assert (
|
||||
to_libpq_url("postgresql+asyncpg://user:pass@host:5432/db")
|
||||
== "postgresql://user:pass@host:5432/db"
|
||||
)
|
||||
|
||||
def test_postgres_asyncpg_scheme_normalized(self) -> None:
|
||||
assert (
|
||||
to_libpq_url("postgres+asyncpg://user:pass@host/db")
|
||||
== "postgresql://user:pass@host/db"
|
||||
)
|
||||
|
||||
def test_bare_postgres_scheme_normalized_to_postgresql(self) -> None:
|
||||
assert to_libpq_url("postgres://user:pass@host/db") == "postgresql://user:pass@host/db"
|
||||
|
||||
|
||||
class TestSslParamRename:
|
||||
def test_ssl_require_to_sslmode_require(self) -> None:
|
||||
assert (
|
||||
to_libpq_url("postgresql+asyncpg://user:pass@host:5432/db?ssl=require")
|
||||
== "postgresql://user:pass@host:5432/db?sslmode=require"
|
||||
)
|
||||
|
||||
@pytest.mark.parametrize("mode", ["disable", "allow", "prefer", "require", "verify-ca", "verify-full"])
|
||||
def test_all_ssl_modes_translated(self, mode: str) -> None:
|
||||
result = to_libpq_url(f"postgresql+asyncpg://h/d?ssl={mode}")
|
||||
assert result == f"postgresql://h/d?sslmode={mode}"
|
||||
|
||||
def test_ssl_rename_on_libpq_url(self) -> None:
|
||||
"""Someone accidentally using SQLAlchemy-style ssl= on a libpq URL is also fixed."""
|
||||
assert to_libpq_url("postgresql://h/d?ssl=require") == "postgresql://h/d?sslmode=require"
|
||||
|
||||
def test_ssl_param_preserved_among_other_params(self) -> None:
|
||||
result = to_libpq_url(
|
||||
"postgresql+asyncpg://h/d?ssl=require&application_name=hindsight&connect_timeout=10"
|
||||
)
|
||||
assert result.startswith("postgresql://h/d?")
|
||||
# Query order should be preserved; ssl renamed, others untouched.
|
||||
assert "sslmode=require" in result
|
||||
assert "application_name=hindsight" in result
|
||||
assert "connect_timeout=10" in result
|
||||
assert "ssl=" not in result.split("?", 1)[1].replace("sslmode=", "")
|
||||
|
||||
def test_sslmode_not_double_renamed(self) -> None:
|
||||
"""An already-correct sslmode= param must not be altered."""
|
||||
assert (
|
||||
to_libpq_url("postgresql+asyncpg://h/d?sslmode=require")
|
||||
== "postgresql://h/d?sslmode=require"
|
||||
)
|
||||
|
||||
|
||||
class TestProductionConfigs:
|
||||
"""Regression guard: current production URL shapes must pass through unchanged.
|
||||
|
||||
These are the exact shapes currently set for HINDSIGHT_API_DATABASE_URL,
|
||||
HINDSIGHT_API_CONTROL_DATABASE_URL and HINDSIGHT_API_MIGRATION_DATABASE_URL
|
||||
in production. The helper must be a pure no-op for them so this change is
|
||||
truly backward-compatible.
|
||||
"""
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"url",
|
||||
[
|
||||
"postgresql://app:[email protected]:5432/appdb?sslmode=disable",
|
||||
"postgresql://app:[email protected]:5432/appdb_control?sslmode=disable",
|
||||
"postgresql://app:[email protected]:5432/appdb?sslmode=disable",
|
||||
],
|
||||
)
|
||||
def test_prod_urls_object_identical(self, url: str) -> None:
|
||||
# Not just equal — must be the exact same object (early-out path),
|
||||
# guaranteeing no parse/reassembly and no subtle mutation.
|
||||
assert to_libpq_url(url) is url
|
||||
|
||||
|
||||
class TestEdgeCases:
|
||||
def test_idempotent(self) -> None:
|
||||
original = "postgresql+asyncpg://user:pass@host:5432/db?ssl=require"
|
||||
once = to_libpq_url(original)
|
||||
twice = to_libpq_url(once)
|
||||
assert once == twice
|
||||
|
||||
def test_password_with_plus_is_preserved(self) -> None:
|
||||
"""A naive str.replace('+asyncpg', ...) would corrupt passwords containing '+'.
|
||||
|
||||
urllib.parse operates on the parsed scheme only, so this stays safe.
|
||||
"""
|
||||
url = "postgresql+asyncpg://user:pa%2Bsswd@host/db?ssl=require"
|
||||
result = to_libpq_url(url)
|
||||
assert result == "postgresql://user:pa%2Bsswd@host/db?sslmode=require"
|
||||
|
||||
def test_password_literal_asyncpg_in_password(self) -> None:
|
||||
"""Even a password that literally contains '+asyncpg' must survive."""
|
||||
url = "postgresql+asyncpg://user:my%2Basyncpgpass@host/db"
|
||||
result = to_libpq_url(url)
|
||||
assert result == "postgresql://user:my%2Basyncpgpass@host/db"
|
||||
|
||||
def test_url_without_query_string(self) -> None:
|
||||
assert (
|
||||
to_libpq_url("postgresql+asyncpg://user:pass@host/db")
|
||||
== "postgresql://user:pass@host/db"
|
||||
)
|
||||
|
||||
def test_url_with_port_and_path_only(self) -> None:
|
||||
assert to_libpq_url("postgresql+asyncpg://host:5432/db") == "postgresql://host:5432/db"
|
||||
@@ -1,166 +0,0 @@
|
||||
"""Regression tests for DeepSeek OpenAI-compatible tool-call quirks."""
|
||||
|
||||
import json
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from hindsight_api.engine.providers.openai_compatible_llm import OpenAICompatibleLLM
|
||||
|
||||
|
||||
TOOLS = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "search_observations",
|
||||
"description": "Search observations",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {"query": {"type": "string"}},
|
||||
"required": ["query"],
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "recall",
|
||||
"description": "Recall memories",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {"query": {"type": "string"}},
|
||||
"required": ["query"],
|
||||
},
|
||||
},
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
def _make_deepseek_llm(model: str = "deepseek-v4-flash") -> OpenAICompatibleLLM:
|
||||
return OpenAICompatibleLLM(
|
||||
provider="openai",
|
||||
api_key="sk-test",
|
||||
base_url="https://api.deepseek.com",
|
||||
model=model,
|
||||
)
|
||||
|
||||
|
||||
def _make_tool_call_response(tool_name: str = "search_observations") -> MagicMock:
|
||||
mock_tc = MagicMock()
|
||||
mock_tc.id = "call_deepseek_123"
|
||||
mock_tc.function.name = tool_name
|
||||
mock_tc.function.arguments = json.dumps({"query": "test"})
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.usage.prompt_tokens = 100
|
||||
mock_response.usage.completion_tokens = 20
|
||||
mock_response.usage.total_tokens = 120
|
||||
mock_response.choices[0].finish_reason = "tool_calls"
|
||||
mock_response.choices[0].message.content = None
|
||||
mock_response.choices[0].message.tool_calls = [mock_tc]
|
||||
return mock_response
|
||||
|
||||
|
||||
def test_deepseek_flash_is_not_treated_as_reasoning_model():
|
||||
llm = _make_deepseek_llm("deepseek-v4-flash")
|
||||
|
||||
assert llm._supports_reasoning_model() is False
|
||||
|
||||
|
||||
def test_deepseek_reasoning_models_still_use_reasoning_parameters():
|
||||
llm = _make_deepseek_llm("deepseek-v4-pro")
|
||||
|
||||
assert llm._supports_reasoning_model() is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deepseek_named_tool_choice_filters_tools_but_omits_tool_choice():
|
||||
"""DeepSeek rejects required/named tool_choice but accepts a narrowed tools list."""
|
||||
llm = _make_deepseek_llm()
|
||||
named_tool_choice = {"type": "function", "function": {"name": "search_observations"}}
|
||||
|
||||
with patch.object(llm._client.chat.completions, "create", new_callable=AsyncMock) as mock_create:
|
||||
mock_create.return_value = _make_tool_call_response("search_observations")
|
||||
|
||||
result = await llm.call_with_tools(
|
||||
messages=[{"role": "user", "content": "Search observations for Project-Rin."}],
|
||||
tools=TOOLS,
|
||||
tool_choice=named_tool_choice,
|
||||
max_retries=0,
|
||||
)
|
||||
|
||||
assert result.tool_calls[0].name == "search_observations"
|
||||
|
||||
sent_kwargs = mock_create.call_args.kwargs
|
||||
assert "tool_choice" not in sent_kwargs
|
||||
assert len(sent_kwargs["tools"]) == 1
|
||||
assert sent_kwargs["tools"][0]["function"]["name"] == "search_observations"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deepseek_tool_history_gets_empty_reasoning_content_fallback():
|
||||
"""DeepSeek requires reasoning_content when replaying assistant tool_calls."""
|
||||
llm = _make_deepseek_llm()
|
||||
messages = [
|
||||
{"role": "user", "content": "Search observations."},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "",
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call_deepseek_123",
|
||||
"type": "function",
|
||||
"function": {"name": "search_observations", "arguments": json.dumps({"query": "test"})},
|
||||
}
|
||||
],
|
||||
},
|
||||
{"role": "tool", "tool_call_id": "call_deepseek_123", "content": "{}"},
|
||||
]
|
||||
|
||||
with patch.object(llm._client.chat.completions, "create", new_callable=AsyncMock) as mock_create:
|
||||
mock_create.return_value = _make_tool_call_response("recall")
|
||||
|
||||
await llm.call_with_tools(
|
||||
messages=messages,
|
||||
tools=TOOLS,
|
||||
tool_choice="auto",
|
||||
max_retries=0,
|
||||
)
|
||||
|
||||
sent_messages = mock_create.call_args.kwargs["messages"]
|
||||
assert sent_messages[1]["reasoning_content"] == ""
|
||||
assert "reasoning_content" not in messages[1]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deepseek_tool_history_preserves_existing_reasoning_content():
|
||||
llm = _make_deepseek_llm()
|
||||
messages = [
|
||||
{"role": "user", "content": "Search observations."},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "",
|
||||
"reasoning_content": "provider reasoning scratchpad",
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call_deepseek_123",
|
||||
"type": "function",
|
||||
"function": {"name": "search_observations", "arguments": json.dumps({"query": "test"})},
|
||||
}
|
||||
],
|
||||
},
|
||||
{"role": "tool", "tool_call_id": "call_deepseek_123", "content": "{}"},
|
||||
]
|
||||
|
||||
with patch.object(llm._client.chat.completions, "create", new_callable=AsyncMock) as mock_create:
|
||||
mock_create.return_value = _make_tool_call_response("recall")
|
||||
|
||||
await llm.call_with_tools(
|
||||
messages=messages,
|
||||
tools=TOOLS,
|
||||
tool_choice="auto",
|
||||
max_retries=0,
|
||||
)
|
||||
|
||||
sent_messages = mock_create.call_args.kwargs["messages"]
|
||||
assert sent_messages[1]["reasoning_content"] == "provider reasoning scratchpad"
|
||||
@@ -1,191 +0,0 @@
|
||||
"""Integration test: delta mental model fuses generic SEO best practices with brand voice.
|
||||
|
||||
Scenario:
|
||||
1. Create a bank with a delta-mode mental model ("editorial-preferences").
|
||||
2. Ingest an SEO best practices document -> trigger mental model refresh.
|
||||
3. Ingest a brand voice document -> trigger mental model refresh (delta).
|
||||
4. Verify the delta fuses both documents organically.
|
||||
|
||||
Requires: HINDSIGHT_RUN_GEMINI_EVALS=1 + a Gemini/OpenAI API key.
|
||||
"""
|
||||
|
||||
import os
|
||||
import uuid
|
||||
from collections import Counter
|
||||
|
||||
import pytest
|
||||
|
||||
from hindsight_api import MemoryEngine, RequestContext
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Gate
|
||||
# ---------------------------------------------------------------------------
|
||||
_GEMINI_KEY = os.getenv("HINDSIGHT_GEMINI_API_KEY") or os.getenv("GEMINI_API_KEY") or os.getenv("GOOGLE_API_KEY")
|
||||
_OPENAI_KEY = os.getenv("OPENAI_API_KEY")
|
||||
_RUN = os.getenv("HINDSIGHT_RUN_GEMINI_EVALS") == "1" and (bool(_GEMINI_KEY) or bool(_OPENAI_KEY))
|
||||
pytestmark = pytest.mark.skipif(not _RUN, reason="Set HINDSIGHT_RUN_GEMINI_EVALS=1 + LLM API key")
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test documents — short but representative
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
SEO_BEST_PRACTICES = """\
|
||||
# SEO Content Best Practices
|
||||
|
||||
## Content Structure
|
||||
- Use clear H1/H2/H3 heading hierarchy for every article.
|
||||
- Keep paragraphs under 3 sentences for scannability.
|
||||
- Use bullet points and numbered lists to break up dense information.
|
||||
|
||||
## Tone and Voice
|
||||
- Write in a professional, authoritative tone.
|
||||
- Use industry-standard SEO terminology (e.g., "SERP", "CTR", "backlink").
|
||||
- Address the reader in second person ("you").
|
||||
|
||||
## Keyword Strategy
|
||||
- Place primary keyword in H1, first paragraph, and meta description.
|
||||
- Target keyword density of 1-2% for primary terms.
|
||||
- Include long-tail question keywords in H2/H3 subheadings.
|
||||
|
||||
## Technical Requirements
|
||||
- Meta titles: 50-60 characters, primary keyword first.
|
||||
- Meta descriptions: 150-160 characters, include CTA.
|
||||
- Internal links: minimum 3 per article.
|
||||
- Image alt text: descriptive, keyword-rich where natural.
|
||||
|
||||
## E-E-A-T Compliance
|
||||
- Include author bios with credentials.
|
||||
- Cite authoritative sources.
|
||||
- Update content quarterly to maintain freshness.
|
||||
"""
|
||||
|
||||
BRAND_VOICE = """\
|
||||
# Plot Brand Voice Guide
|
||||
|
||||
## Who We Are
|
||||
Plot is a finance app for freelancers. We handle invoicing, expense tracking,
|
||||
and tax prep for people whose income is irregular.
|
||||
|
||||
## Voice Principles
|
||||
- We talk like a smart friend who knows about money — not a bank, not a guru.
|
||||
- Clarity always wins. If a 12-year-old can't understand it, rewrite it.
|
||||
- We never lecture or moralize about financial decisions.
|
||||
|
||||
## Tone by Context
|
||||
- Marketing: Confident, slightly wry. Example: "Built for income that doesn't show up on the same day every month."
|
||||
- Support: Direct, human, accountable. Example: "That's our bug, not yours. We're fixing it now."
|
||||
- Product UI: Quiet, precise. Example: "Income from Stripe — Mar 14."
|
||||
- Errors: Calm, specific. Example: "We couldn't sync your bank. Try reconnecting."
|
||||
|
||||
## Writing Rules
|
||||
- Always use contractions (it's, we're, you'll).
|
||||
- Use Oxford comma.
|
||||
- Always say "you", never "users" or "customers".
|
||||
- Avoid jargon: never say "leverage", "empower", "solution", "holistic", "game-changing".
|
||||
- No puns. Wit is fine — wordplay and wry asides, not dad jokes.
|
||||
|
||||
## What We Sound Like
|
||||
- YES: "Here's what we found." / NO: "We are pleased to present our findings."
|
||||
- YES: "Looks like this payment is late." / NO: "ALERT: Payment overdue! Action required!"
|
||||
"""
|
||||
|
||||
|
||||
class TestDeltaEditorialFusion:
|
||||
"""Real-LLM test verifying delta mode correctly fuses two documents."""
|
||||
|
||||
async def test_delta_fuses_seo_and_brand_voice(
|
||||
self,
|
||||
memory: MemoryEngine,
|
||||
request_context: RequestContext,
|
||||
):
|
||||
bank_id = f"test-editorial-{uuid.uuid4().hex[:8]}"
|
||||
await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
|
||||
try:
|
||||
mm = await memory.create_mental_model(
|
||||
bank_id=bank_id,
|
||||
name="Editorial Preferences",
|
||||
source_query=(
|
||||
"What are the editorial preferences and content guidelines? "
|
||||
"Include tone, voice, formatting rules, and vocabulary rules."
|
||||
),
|
||||
content="",
|
||||
trigger={
|
||||
"mode": "delta",
|
||||
"refresh_after_consolidation": False,
|
||||
"fact_types": ["observation"],
|
||||
"exclude_mental_models": True,
|
||||
},
|
||||
request_context=request_context,
|
||||
)
|
||||
mm_id = mm["id"]
|
||||
|
||||
# Phase 1: Ingest SEO best practices
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id, content=SEO_BEST_PRACTICES,
|
||||
document_id="seo-best-practices", request_context=request_context,
|
||||
)
|
||||
mm_after_seo = await memory.refresh_mental_model(
|
||||
bank_id=bank_id, mental_model_id=mm_id, request_context=request_context,
|
||||
)
|
||||
seo_content = mm_after_seo["content"]
|
||||
assert len(seo_content) > 100, f"First refresh produced too little content: {len(seo_content)} chars"
|
||||
|
||||
# Phase 2: Ingest brand voice -> delta refresh
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id, content=BRAND_VOICE,
|
||||
document_id="brand-voice", request_context=request_context,
|
||||
)
|
||||
mm_after_brand = await memory.refresh_mental_model(
|
||||
bank_id=bank_id, mental_model_id=mm_id, request_context=request_context,
|
||||
)
|
||||
fused = mm_after_brand["content"]
|
||||
rr = mm_after_brand.get("reflect_response") or {}
|
||||
fused_lower = fused.lower()
|
||||
|
||||
# -- Verify fusion quality --
|
||||
|
||||
# Brand voice concepts present (LLM may paraphrase, check synonyms)
|
||||
for concept, signals in {
|
||||
"contractions": ["contraction", "it's", "we're", "you'll"],
|
||||
"oxford comma": ["oxford comma"],
|
||||
"vocabulary rules": ["jargon", "leverage", "empower", "forbidden"],
|
||||
}.items():
|
||||
assert any(s in fused_lower for s in signals), (
|
||||
f"Brand voice concept '{concept}' missing (looked for {signals}).\n"
|
||||
f"Fused content:\n{fused[:500]}"
|
||||
)
|
||||
|
||||
# SEO concepts still present (not wiped by delta)
|
||||
for concept, signals in {
|
||||
"keywords": ["keyword"],
|
||||
"structure": ["heading", "h1", "h2", "structure"],
|
||||
"seo": ["meta", "e-e-a-t", "seo", "search"],
|
||||
}.items():
|
||||
assert any(s in fused_lower for s in signals), (
|
||||
f"SEO concept '{concept}' missing (looked for {signals}).\n"
|
||||
f"Fused content:\n{fused[:500]}"
|
||||
)
|
||||
|
||||
# Brand voice overrides generic tone
|
||||
assert any(t in fused_lower for t in ["friend", "wry", "plot", "witty"]), (
|
||||
f"Brand-specific tone missing from fused content.\nFused:\n{fused[:500]}"
|
||||
)
|
||||
|
||||
# No duplicate paragraphs
|
||||
lines = [
|
||||
ln.strip() for ln in fused.split("\n")
|
||||
if ln.strip() and not ln.strip().startswith("#")
|
||||
]
|
||||
dupes = {line: cnt for line, cnt in Counter(lines).items() if cnt > 1}
|
||||
assert not dupes, (
|
||||
"Duplicate paragraphs:\n" +
|
||||
"\n".join(f" [{c}x] {t[:80]}" for t, c in dupes.items())
|
||||
)
|
||||
|
||||
# based_on accumulates from both docs
|
||||
obs_count = len(rr.get("based_on", {}).get("observation", []))
|
||||
assert obs_count > 5, f"Expected observations from both docs, got {obs_count}"
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
@@ -9,14 +9,6 @@ import pytest
|
||||
|
||||
from hindsight_api import RequestContext
|
||||
from hindsight_api.engine.memory_engine import Budget
|
||||
from hindsight_api.extensions import (
|
||||
OperationValidatorExtension,
|
||||
RecallContext,
|
||||
ReflectContext,
|
||||
RetainContext,
|
||||
RetainResult,
|
||||
ValidationResult,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -25,31 +17,6 @@ def _ts():
|
||||
return datetime.now(timezone.utc).timestamp()
|
||||
|
||||
|
||||
class _RetainResultCapture(OperationValidatorExtension):
|
||||
"""Minimal OperationValidator that records each RetainResult it receives.
|
||||
|
||||
Used by tests to assert on fields the engine sets on RetainResult (e.g.
|
||||
processed_content_tokens), without having to scrape logs or internals.
|
||||
The pre-operation validators must be implemented to satisfy the
|
||||
abstract base class, but they always accept.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.results: list[RetainResult] = []
|
||||
|
||||
async def validate_retain(self, ctx: RetainContext) -> ValidationResult:
|
||||
return ValidationResult.accept()
|
||||
|
||||
async def validate_recall(self, ctx: RecallContext) -> ValidationResult:
|
||||
return ValidationResult.accept()
|
||||
|
||||
async def validate_reflect(self, ctx: ReflectContext) -> ValidationResult:
|
||||
return ValidationResult.accept()
|
||||
|
||||
async def on_retain_complete(self, result: RetainResult) -> None:
|
||||
self.results.append(result)
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Core Delta Retain Tests
|
||||
# ============================================================
|
||||
@@ -873,185 +840,3 @@ async def test_delta_retain_recall_with_chunks(memory, request_context):
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
# ============================================================
|
||||
# processed_content_tokens on RetainResult
|
||||
# ============================================================
|
||||
#
|
||||
# These tests verify the signal the engine exposes via
|
||||
# RetainResult.processed_content_tokens for the post-retain hook. That
|
||||
# field lets a metering/billing extension tell the difference between:
|
||||
# * a retain that went through the full extraction pipeline (None),
|
||||
# * a retain whose chunks all matched prior content (0),
|
||||
# * a retain where only some chunks were new/changed (N>0, the
|
||||
# content+context tokens of the chunks that were actually processed).
|
||||
|
||||
|
||||
def test_merge_processed_content_tokens_helper():
|
||||
"""Unit check on the None-propagating aggregator used by the engine."""
|
||||
from hindsight_api.engine.retain.orchestrator import (
|
||||
_merge_processed_content_tokens,
|
||||
)
|
||||
|
||||
assert _merge_processed_content_tokens(0, 0) == 0
|
||||
assert _merge_processed_content_tokens(5, 7) == 12
|
||||
# None "wins" in either slot — once any sub-result bypassed dedup, the
|
||||
# aggregate is None so callers bill full content.
|
||||
assert _merge_processed_content_tokens(None, 10) is None
|
||||
assert _merge_processed_content_tokens(10, None) is None
|
||||
assert _merge_processed_content_tokens(None, None) is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_processed_content_tokens_first_retain_is_none(memory, request_context):
|
||||
"""
|
||||
First retain to a new document goes through the full (non-delta) path,
|
||||
so processed_content_tokens should be None — the caller has no dedup
|
||||
signal and should bill the full submitted content.
|
||||
"""
|
||||
bank_id = f"test_pct_first_{_ts()}"
|
||||
document_id = "new-doc"
|
||||
capture = _RetainResultCapture()
|
||||
memory._operation_validator = capture
|
||||
|
||||
try:
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="Alice works at Google.",
|
||||
context="test",
|
||||
document_id=document_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
assert len(capture.results) == 1
|
||||
assert capture.results[0].processed_content_tokens is None, (
|
||||
"First retain (full path) should report processed_content_tokens=None"
|
||||
)
|
||||
finally:
|
||||
memory._operation_validator = None
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_processed_content_tokens_unchanged_resubmit_is_zero(memory, request_context):
|
||||
"""
|
||||
Re-retaining identical content to the same document_id should hit the
|
||||
'no chunks changed' path and report processed_content_tokens=0.
|
||||
"""
|
||||
bank_id = f"test_pct_unchanged_{_ts()}"
|
||||
document_id = "conversation-001"
|
||||
capture = _RetainResultCapture()
|
||||
memory._operation_validator = capture
|
||||
content = "Alice works at Google. Bob works at Microsoft."
|
||||
|
||||
try:
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content=content,
|
||||
context="team info",
|
||||
document_id=document_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
# Identical resubmit.
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content=content,
|
||||
context="team info",
|
||||
document_id=document_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
assert len(capture.results) == 2
|
||||
assert capture.results[0].processed_content_tokens is None
|
||||
assert capture.results[1].processed_content_tokens == 0, (
|
||||
"Unchanged resubmit should report zero processed content tokens"
|
||||
)
|
||||
finally:
|
||||
memory._operation_validator = None
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_processed_content_tokens_appended_reports_delta(memory, request_context):
|
||||
"""
|
||||
Appending new content to an existing document should surface a
|
||||
non-zero processed_content_tokens that is less than the full
|
||||
submitted content tokens — only the new/changed chunks are counted.
|
||||
"""
|
||||
bank_id = f"test_pct_appended_{_ts()}"
|
||||
document_id = "growing-doc"
|
||||
capture = _RetainResultCapture()
|
||||
memory._operation_validator = capture
|
||||
|
||||
v1 = "Alice works at Google."
|
||||
# Make v2 large enough that the delta diff classifies some chunks as
|
||||
# unchanged (shared prefix) and some as new (the appended tail). The
|
||||
# chunker splits on ``retain_chunk_size`` (default 3000), so we pad
|
||||
# each part with a comfortable margin of filler text to force a chunk
|
||||
# boundary between them.
|
||||
filler = " The project budget is fine. " * 400 # ~12 KB
|
||||
v2 = v1 + filler + " Bob works at Microsoft."
|
||||
|
||||
try:
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content=v1,
|
||||
context="profile",
|
||||
document_id=document_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content=v2,
|
||||
context="profile",
|
||||
document_id=document_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
assert len(capture.results) == 2
|
||||
|
||||
# Second retain should either:
|
||||
# * Be on the delta path with a positive partial count strictly
|
||||
# less than the full submission (the common case), OR
|
||||
# * Fall back to full retain if the chunker decided nothing
|
||||
# matched (in which case we report None and bill full).
|
||||
# Both are correct signals for the billing extension; the test
|
||||
# just asserts they're shaped sanely.
|
||||
from hindsight_api.engine.memory_engine import count_tokens
|
||||
|
||||
submitted_tokens = count_tokens(v2) + count_tokens("profile")
|
||||
second = capture.results[1].processed_content_tokens
|
||||
if second is None:
|
||||
# Fell back to full retain — acceptable signal.
|
||||
return
|
||||
assert second > 0, "Partial-delta retain should report a positive token count"
|
||||
assert second < submitted_tokens, (
|
||||
"Partial-delta retain should report fewer processed tokens "
|
||||
"than the full submitted payload"
|
||||
)
|
||||
finally:
|
||||
memory._operation_validator = None
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_processed_content_tokens_without_document_id_is_none(memory, request_context):
|
||||
"""
|
||||
A retain without a document_id can't participate in per-document
|
||||
dedup, so the engine should report processed_content_tokens=None
|
||||
and let the caller bill the full submitted payload.
|
||||
"""
|
||||
bank_id = f"test_pct_no_doc_{_ts()}"
|
||||
capture = _RetainResultCapture()
|
||||
memory._operation_validator = capture
|
||||
|
||||
try:
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="A one-off observation with no document_id.",
|
||||
context="test",
|
||||
request_context=request_context,
|
||||
)
|
||||
assert len(capture.results) == 1
|
||||
assert capture.results[0].processed_content_tokens is None
|
||||
finally:
|
||||
memory._operation_validator = None
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@@ -1,404 +0,0 @@
|
||||
"""
|
||||
Tests for delta retain chunk ordering and duplicate prevention.
|
||||
|
||||
Verifies that:
|
||||
1. Chunks are stored with deterministic indices (not task completion order)
|
||||
2. Delta retain can correctly identify unchanged chunks on subsequent upserts
|
||||
3. Repeated upserts of same content don't produce duplicate memory units
|
||||
4. Concurrent retains on the same document produce clean final state (no duplicates)
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import os
|
||||
from datetime import datetime, timezone
|
||||
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
|
||||
from hindsight_api import RequestContext
|
||||
from hindsight_api.engine.task_backend import SyncTaskBackend
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _ts():
|
||||
return datetime.now(timezone.utc).timestamp()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_repeated_upsert_chunks_not_scrambled(memory, request_context):
|
||||
"""
|
||||
Verify that chunks are stored with correct indices matching the
|
||||
deterministic chunking order, not task completion order.
|
||||
|
||||
This is critical for delta retain: if chunk indices don't match the
|
||||
deterministic order, delta will think all chunks changed on every
|
||||
upsert and fall back to full re-processing.
|
||||
"""
|
||||
bank_id = f"test_chunk_order_{_ts()}"
|
||||
document_id = "chunk-order-doc"
|
||||
|
||||
try:
|
||||
# Create content that produces multiple distinct chunks
|
||||
chunk1_text = "Alice works at Google on Search. " * 100 # ~3300 chars
|
||||
chunk2_text = "Bob works at Microsoft on Azure. " * 100 # ~3400 chars
|
||||
content = chunk1_text + chunk2_text
|
||||
|
||||
assert len(content) > 6000, "Should produce at least 2 chunks"
|
||||
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content=content,
|
||||
context="team info",
|
||||
document_id=document_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Load chunks from DB and verify order matches deterministic chunking
|
||||
from hindsight_api.engine.retain import chunk_storage, fact_extraction
|
||||
|
||||
pool = await memory._get_pool()
|
||||
|
||||
# Get the chunk texts from DB
|
||||
async with pool.acquire() as conn:
|
||||
chunk_rows = await conn.fetch(
|
||||
"SELECT chunk_index, chunk_text, content_hash FROM chunks WHERE bank_id = $1 AND document_id = $2 ORDER BY chunk_index",
|
||||
bank_id,
|
||||
document_id,
|
||||
)
|
||||
|
||||
# Compute expected chunks deterministically (default chunk_size is 3000)
|
||||
chunk_size = 3000
|
||||
expected_chunks = fact_extraction.chunk_text(content, max_chars=chunk_size)
|
||||
|
||||
logger.info(f"Expected {len(expected_chunks)} chunks, got {len(chunk_rows)} in DB")
|
||||
|
||||
# Verify each chunk at its index has the correct content hash
|
||||
for i, expected_text in enumerate(expected_chunks):
|
||||
expected_hash = chunk_storage.compute_chunk_hash(expected_text)
|
||||
matching_rows = [r for r in chunk_rows if r["chunk_index"] == i]
|
||||
assert len(matching_rows) == 1, f"Expected exactly 1 chunk at index {i}, got {len(matching_rows)}"
|
||||
actual_hash = matching_rows[0]["content_hash"]
|
||||
assert actual_hash == expected_hash, (
|
||||
f"Chunk at index {i} has wrong content hash. "
|
||||
f"Expected hash of first 50 chars: {repr(expected_text[:50])}, "
|
||||
f"got hash of: {repr(matching_rows[0]['chunk_text'][:50])}"
|
||||
)
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delta_detects_unchanged_after_first_retain(memory, request_context):
|
||||
"""
|
||||
After first retain stores chunks with correct indices, a second retain
|
||||
with identical content should use the delta path and detect all chunks
|
||||
as unchanged (no re-processing).
|
||||
"""
|
||||
bank_id = f"test_delta_unchanged_{_ts()}"
|
||||
document_id = "delta-unchanged-doc"
|
||||
|
||||
try:
|
||||
# Multi-chunk content with distinct sections
|
||||
chunk1_text = "Alice works at Google on Search. " * 100
|
||||
chunk2_text = "Bob works at Microsoft on Azure. " * 100
|
||||
content = chunk1_text + chunk2_text
|
||||
|
||||
# First retain
|
||||
v1_units = await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content=content,
|
||||
context="team info",
|
||||
document_id=document_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
assert len(v1_units) > 0
|
||||
|
||||
pool = await memory._get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
v1_count = await conn.fetchval(
|
||||
"SELECT count(*) FROM memory_units WHERE bank_id = $1 AND document_id = $2",
|
||||
bank_id,
|
||||
document_id,
|
||||
)
|
||||
|
||||
# Second retain — same content, should be detected as unchanged by delta
|
||||
v2_units = await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content=content,
|
||||
context="team info",
|
||||
document_id=document_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Delta should detect all unchanged → return empty (no new units)
|
||||
assert v2_units == [], f"Delta with unchanged content should return empty, got {len(v2_units)} units"
|
||||
|
||||
# Memory unit count should not change
|
||||
async with pool.acquire() as conn:
|
||||
v2_count = await conn.fetchval(
|
||||
"SELECT count(*) FROM memory_units WHERE bank_id = $1 AND document_id = $2",
|
||||
bank_id,
|
||||
document_id,
|
||||
)
|
||||
assert v2_count == v1_count, (
|
||||
f"Memory unit count changed on same-content upsert: {v1_count} -> {v2_count}"
|
||||
)
|
||||
|
||||
# Third retain — verify stability
|
||||
v3_units = await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content=content,
|
||||
context="team info",
|
||||
document_id=document_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
assert v3_units == [], "Third retain should also detect unchanged"
|
||||
|
||||
async with pool.acquire() as conn:
|
||||
v3_count = await conn.fetchval(
|
||||
"SELECT count(*) FROM memory_units WHERE bank_id = $1 AND document_id = $2",
|
||||
bank_id,
|
||||
document_id,
|
||||
)
|
||||
assert v3_count == v1_count, (
|
||||
f"Memory unit count changed on third upsert: {v1_count} -> {v3_count}"
|
||||
)
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stale_request_skipped_when_newer_retain_completed(memory, request_context):
|
||||
"""
|
||||
When two retains race on the same document, the one that started earlier
|
||||
(stale) should be skipped if the newer one already completed.
|
||||
|
||||
Simulates: Request B (newer content) completes while Request A (older content)
|
||||
was waiting for the advisory lock. When A finally acquires the lock, it sees
|
||||
the document was updated after its start_time and skips.
|
||||
"""
|
||||
bank_id = f"test_stale_skip_{_ts()}"
|
||||
document_id = "stale-skip-doc"
|
||||
|
||||
try:
|
||||
# First: establish the document with initial content
|
||||
newer_content = "Alice works at Google. Bob works at Microsoft. Charlie works at Apple."
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content=newer_content,
|
||||
context="team",
|
||||
document_id=document_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
pool = await memory._get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
after_newer_count = await conn.fetchval(
|
||||
"SELECT count(*) FROM memory_units WHERE bank_id = $1 AND document_id = $2",
|
||||
bank_id,
|
||||
document_id,
|
||||
)
|
||||
assert after_newer_count > 0, "Should have facts from newer content"
|
||||
|
||||
# Simulate the race condition by pushing the document's updated_at into
|
||||
# the future. This makes any new retain appear "stale" (its start_time
|
||||
# is before updated_at), as if another request already completed.
|
||||
async with pool.acquire() as conn:
|
||||
await conn.execute(
|
||||
"UPDATE documents SET updated_at = NOW() + INTERVAL '10 seconds' WHERE id = $1 AND bank_id = $2",
|
||||
document_id,
|
||||
bank_id,
|
||||
)
|
||||
|
||||
# Now try to retain with older/different content. The stale-request check
|
||||
# should detect that updated_at > start_time and skip this request.
|
||||
older_content = "Alice works at Google."
|
||||
result = await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content=older_content,
|
||||
context="team",
|
||||
document_id=document_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# The stale request should have been skipped (empty result)
|
||||
assert result == [], f"Stale request should return empty, got {result}"
|
||||
|
||||
# Memory units should be unchanged (newer content preserved)
|
||||
async with pool.acquire() as conn:
|
||||
final_count = await conn.fetchval(
|
||||
"SELECT count(*) FROM memory_units WHERE bank_id = $1 AND document_id = $2",
|
||||
bank_id,
|
||||
document_id,
|
||||
)
|
||||
assert final_count == after_newer_count, (
|
||||
f"Stale request should not change memory units: {after_newer_count} -> {final_count}"
|
||||
)
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Concurrent Retain Stress Test
|
||||
# ============================================================
|
||||
|
||||
|
||||
@pytest_asyncio.fixture(scope="function")
|
||||
async def memory_no_llm(pg0_db_url, embeddings, cross_encoder, query_analyzer):
|
||||
"""
|
||||
MemoryEngine with provider=none (chunks mode, no LLM needed).
|
||||
Each chunk is stored verbatim as a single memory unit — fast and deterministic.
|
||||
"""
|
||||
from hindsight_api.engine.memory_engine import MemoryEngine
|
||||
|
||||
mem = MemoryEngine(
|
||||
db_url=pg0_db_url,
|
||||
memory_llm_provider="none",
|
||||
memory_llm_api_key="",
|
||||
memory_llm_model="none",
|
||||
embeddings=embeddings,
|
||||
cross_encoder=cross_encoder,
|
||||
query_analyzer=query_analyzer,
|
||||
pool_min_size=2,
|
||||
pool_max_size=10,
|
||||
run_migrations=False,
|
||||
task_backend=SyncTaskBackend(),
|
||||
skip_llm_verification=True,
|
||||
)
|
||||
await mem.initialize()
|
||||
yield mem
|
||||
await mem.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_concurrent_upserts_no_duplicates(memory_no_llm, request_context):
|
||||
"""
|
||||
Stress test: N concurrent retains of the same document with different content.
|
||||
|
||||
Each version has distinct content so we can verify the final state is exactly
|
||||
one version's data — no duplicates, no mixed data from different versions.
|
||||
|
||||
With provider=none (chunks mode), each chunk becomes a verbatim memory unit,
|
||||
so we can inspect exactly which chunks survived.
|
||||
|
||||
The test verifies:
|
||||
- Exactly one version's document row survives (by content_hash)
|
||||
- All memory units belong to a single version (no cross-version mixing)
|
||||
- No duplicate memory units exist
|
||||
- Chunk count matches what the winning version should have
|
||||
"""
|
||||
bank_id = f"test_concurrent_{_ts()}"
|
||||
document_id = "concurrent-doc"
|
||||
num_concurrent = 20
|
||||
|
||||
try:
|
||||
# Each version has unique, identifiable content.
|
||||
# Make content large enough for multiple chunks (~3000 chars per chunk).
|
||||
versions = []
|
||||
for v in range(num_concurrent):
|
||||
# Each version's chunks will contain "VERSION_XX" markers so we can
|
||||
# identify which version's data survived in the final state.
|
||||
content = f"VERSION_{v:02d} " + f"Person_{v} works at Company_{v}. " * 200
|
||||
versions.append(content)
|
||||
|
||||
# Fire all retains concurrently
|
||||
async def _retain_version(version_content: str) -> None:
|
||||
await memory_no_llm.retain_async(
|
||||
bank_id=bank_id,
|
||||
content=version_content,
|
||||
document_id=document_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
results = await asyncio.gather(
|
||||
*[_retain_version(v) for v in versions],
|
||||
return_exceptions=True,
|
||||
)
|
||||
|
||||
# Some may have been aborted (pipeline_aborted) — that's expected.
|
||||
# Check for unexpected errors.
|
||||
errors = [r for r in results if isinstance(r, Exception)]
|
||||
for err in errors:
|
||||
logger.warning(f"Concurrent retain error (may be expected): {err}")
|
||||
|
||||
# --- Verify final state ---
|
||||
pool = await memory_no_llm._get_pool()
|
||||
|
||||
# 1. Exactly one document row should exist
|
||||
async with pool.acquire() as conn:
|
||||
doc_rows = await conn.fetch(
|
||||
"SELECT id, content_hash FROM documents WHERE id = $1 AND bank_id = $2",
|
||||
document_id,
|
||||
bank_id,
|
||||
)
|
||||
assert len(doc_rows) == 1, f"Expected 1 document row, got {len(doc_rows)}"
|
||||
winning_hash = doc_rows[0]["content_hash"]
|
||||
|
||||
# Find which version won by matching content_hash
|
||||
import hashlib
|
||||
|
||||
from hindsight_api.engine.retain.fact_extraction import _sanitize_text
|
||||
|
||||
winning_version = None
|
||||
for v, content in enumerate(versions):
|
||||
sanitized = _sanitize_text(content) or ""
|
||||
h = hashlib.sha256(sanitized.encode()).hexdigest()
|
||||
if h == winning_hash:
|
||||
winning_version = v
|
||||
break
|
||||
assert winning_version is not None, "Could not identify winning version from content_hash"
|
||||
logger.info(f"Winning version: {winning_version} (out of {num_concurrent} concurrent retains)")
|
||||
|
||||
# 2. All memory units should belong to the winning version
|
||||
async with pool.acquire() as conn:
|
||||
units = await conn.fetch(
|
||||
"SELECT text, chunk_id, id::text as unit_id FROM memory_units WHERE bank_id = $1 AND document_id = $2",
|
||||
bank_id,
|
||||
document_id,
|
||||
)
|
||||
unit_texts = [r["text"] for r in units]
|
||||
assert len(unit_texts) > 0, "Should have at least 1 memory unit"
|
||||
|
||||
# In chunks mode, each memory unit text IS the chunk text.
|
||||
# Every unit should contain the winning version's unique person name.
|
||||
# We check for "Person_N" rather than "VERSION_N" because the text
|
||||
# splitter may cut mid-text, so later chunks might not start with the prefix.
|
||||
winning_person = f"Person_{winning_version}"
|
||||
wrong_version_units = [
|
||||
(r["text"], r["chunk_id"], r["unit_id"])
|
||||
for r in units
|
||||
if winning_person not in r["text"]
|
||||
]
|
||||
assert not wrong_version_units, (
|
||||
f"Found {len(wrong_version_units)} memory units NOT from winning version "
|
||||
f"{winning_version} (expected '{winning_person}' in every unit). "
|
||||
f"Details: {[(t[:60], cid, uid) for t, cid, uid in wrong_version_units]}"
|
||||
)
|
||||
|
||||
# 3. No duplicate memory units
|
||||
from collections import Counter
|
||||
|
||||
text_counts = Counter(unit_texts)
|
||||
duplicates = {text[:80]: count for text, count in text_counts.items() if count > 1}
|
||||
assert not duplicates, f"Found duplicate memory units: {duplicates}"
|
||||
|
||||
# 4. Chunk count matches expected
|
||||
from hindsight_api.engine.retain.fact_extraction import chunk_text
|
||||
|
||||
expected_chunks = chunk_text(versions[winning_version], max_chars=3000)
|
||||
assert len(unit_texts) == len(expected_chunks), (
|
||||
f"Expected {len(expected_chunks)} chunks for winning version, got {len(unit_texts)} memory units"
|
||||
)
|
||||
|
||||
logger.info(
|
||||
f"Concurrent test passed: version {winning_version} won with "
|
||||
f"{len(unit_texts)} memory units, no duplicates"
|
||||
)
|
||||
|
||||
finally:
|
||||
await memory_no_llm.delete_bank(bank_id, request_context=request_context)
|
||||
@@ -1,309 +0,0 @@
|
||||
"""
|
||||
Tests for document chunks API, reprocess, nodes_by_fact_type, and graph document/chunk filtering.
|
||||
"""
|
||||
from datetime import datetime, timezone
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
|
||||
from hindsight_api.api import create_app
|
||||
|
||||
|
||||
# ── Fixtures ──
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def api_client(memory):
|
||||
"""Create an async test client for the FastAPI app."""
|
||||
app = create_app(memory, initialize_memory=False)
|
||||
transport = httpx.ASGITransport(app=app)
|
||||
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
|
||||
yield client
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def bank_id():
|
||||
return f"test_doc_chunks_{datetime.now(timezone.utc).timestamp()}"
|
||||
|
||||
|
||||
async def _retain(api_client, bank_id, document_id, content, tags=None):
|
||||
"""Helper to retain a document via the HTTP API."""
|
||||
item = {"content": content, "document_id": document_id}
|
||||
if tags:
|
||||
item["tags"] = tags
|
||||
response = await api_client.post(
|
||||
f"/v1/default/banks/{bank_id}/memories",
|
||||
json={"items": [item]},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
return response.json()
|
||||
|
||||
|
||||
# ── list_document_chunks ──
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_document_chunks(memory, request_context):
|
||||
"""list_document_chunks returns chunks ordered by chunk_index."""
|
||||
bank_id = f"test_chunks_{datetime.now(timezone.utc).timestamp()}"
|
||||
|
||||
try:
|
||||
await memory.retain_batch_async(
|
||||
bank_id=bank_id,
|
||||
contents=[{"content": "Alice works at Google. Bob works at Meta. " * 20, "document_id": "doc1"}],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
result = await memory.list_document_chunks(
|
||||
bank_id=bank_id,
|
||||
document_id="doc1",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert result is not None
|
||||
assert result["total"] >= 1
|
||||
assert len(result["items"]) == result["total"]
|
||||
|
||||
# Chunks should be ordered by chunk_index
|
||||
indices = [c["chunk_index"] for c in result["items"]]
|
||||
assert indices == sorted(indices)
|
||||
|
||||
# Each chunk should have the expected fields
|
||||
for chunk in result["items"]:
|
||||
assert "chunk_id" in chunk
|
||||
assert "chunk_text" in chunk
|
||||
assert chunk["document_id"] == "doc1"
|
||||
assert chunk["bank_id"] == bank_id
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_document_chunks_pagination(memory, request_context):
|
||||
"""list_document_chunks respects limit and offset."""
|
||||
bank_id = f"test_chunks_page_{datetime.now(timezone.utc).timestamp()}"
|
||||
|
||||
try:
|
||||
await memory.retain_batch_async(
|
||||
bank_id=bank_id,
|
||||
contents=[{"content": "Content. " * 200, "document_id": "doc-pag"}],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
all_chunks = await memory.list_document_chunks(
|
||||
bank_id=bank_id, document_id="doc-pag", request_context=request_context
|
||||
)
|
||||
total = all_chunks["total"]
|
||||
if total < 2:
|
||||
pytest.skip("Document produced fewer than 2 chunks, can't test pagination")
|
||||
|
||||
page1 = await memory.list_document_chunks(
|
||||
bank_id=bank_id, document_id="doc-pag", limit=1, offset=0, request_context=request_context
|
||||
)
|
||||
assert len(page1["items"]) == 1
|
||||
assert page1["total"] == total
|
||||
|
||||
page2 = await memory.list_document_chunks(
|
||||
bank_id=bank_id, document_id="doc-pag", limit=1, offset=1, request_context=request_context
|
||||
)
|
||||
assert len(page2["items"]) == 1
|
||||
assert page2["items"][0]["chunk_id"] != page1["items"][0]["chunk_id"]
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_document_chunks_not_found(memory, request_context):
|
||||
"""list_document_chunks returns None for non-existent document."""
|
||||
bank_id = f"test_chunks_404_{datetime.now(timezone.utc).timestamp()}"
|
||||
|
||||
result = await memory.list_document_chunks(
|
||||
bank_id=bank_id, document_id="nonexistent", request_context=request_context
|
||||
)
|
||||
assert result is None
|
||||
|
||||
|
||||
# ── get_document nodes_by_fact_type ──
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_document_nodes_by_fact_type(memory, request_context):
|
||||
"""get_document returns nodes_by_fact_type with per-type counts."""
|
||||
bank_id = f"test_doc_composition_{datetime.now(timezone.utc).timestamp()}"
|
||||
|
||||
try:
|
||||
await memory.retain_batch_async(
|
||||
bank_id=bank_id,
|
||||
contents=[{"content": "Alice works at Google on AI research.", "document_id": "doc-comp"}],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
doc = await memory.get_document("doc-comp", bank_id, request_context=request_context)
|
||||
assert doc is not None
|
||||
assert "nodes_by_fact_type" in doc
|
||||
|
||||
nbt = doc["nodes_by_fact_type"]
|
||||
assert "world" in nbt
|
||||
assert "experience" in nbt
|
||||
assert "observation" in nbt
|
||||
assert isinstance(nbt["world"], int)
|
||||
assert isinstance(nbt["experience"], int)
|
||||
assert isinstance(nbt["observation"], int)
|
||||
|
||||
# Total should match memory_unit_count
|
||||
assert nbt["world"] + nbt["experience"] + nbt["observation"] == doc["memory_unit_count"]
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
# ── reprocess_document ──
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reprocess_document(memory, request_context):
|
||||
"""reprocess_document submits an async retain operation for an existing document."""
|
||||
bank_id = f"test_reprocess_{datetime.now(timezone.utc).timestamp()}"
|
||||
|
||||
try:
|
||||
await memory.retain_batch_async(
|
||||
bank_id=bank_id,
|
||||
contents=[{"content": "Alice works at Google.", "document_id": "doc-reprocess"}],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
result = await memory.reprocess_document(
|
||||
bank_id=bank_id, document_id="doc-reprocess", request_context=request_context
|
||||
)
|
||||
|
||||
assert result is not None
|
||||
assert "operation_id" in result
|
||||
assert "items_count" in result
|
||||
assert result["items_count"] == 1
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reprocess_document_not_found(memory, request_context):
|
||||
"""reprocess_document returns None for non-existent document."""
|
||||
result = await memory.reprocess_document(
|
||||
bank_id="nonexistent-bank", document_id="nonexistent", request_context=request_context
|
||||
)
|
||||
assert result is None
|
||||
|
||||
|
||||
# ── Graph document_id / chunk_id filters (HTTP level) ──
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_graph_document_id_filter(api_client, bank_id):
|
||||
"""Graph endpoint filters by document_id."""
|
||||
await _retain(api_client, bank_id, "doc-a", "Alice works at Google on AI.")
|
||||
await _retain(api_client, bank_id, "doc-b", "Bob works at Meta on VR.")
|
||||
|
||||
# Filter by doc-a
|
||||
response = await api_client.get(
|
||||
f"/v1/default/banks/{bank_id}/graph",
|
||||
params={"document_id": "doc-a"},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
doc_ids = {row.get("document_id") for row in data["table_rows"]}
|
||||
assert "doc-a" in doc_ids
|
||||
assert "doc-b" not in doc_ids
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_graph_chunk_id_filter(api_client, bank_id):
|
||||
"""Graph endpoint filters by chunk_id."""
|
||||
await _retain(api_client, bank_id, "doc-chunk-test", "Alice works at Google. " * 20)
|
||||
|
||||
# First get chunks to find a valid chunk_id
|
||||
response = await api_client.get(
|
||||
f"/v1/default/banks/{bank_id}/documents/doc-chunk-test/chunks"
|
||||
)
|
||||
assert response.status_code == 200
|
||||
chunks_data = response.json()
|
||||
if chunks_data["total"] == 0:
|
||||
pytest.skip("No chunks created")
|
||||
|
||||
chunk_id = chunks_data["items"][0]["chunk_id"]
|
||||
|
||||
# Filter graph by that chunk_id
|
||||
response = await api_client.get(
|
||||
f"/v1/default/banks/{bank_id}/graph",
|
||||
params={"chunk_id": chunk_id},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
chunk_ids = {row.get("chunk_id") for row in data["table_rows"]}
|
||||
# All returned memories should belong to the requested chunk
|
||||
assert all(cid == chunk_id for cid in chunk_ids if cid is not None)
|
||||
|
||||
|
||||
# ── HTTP endpoints for chunks and reprocess ──
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_http_list_document_chunks(api_client, bank_id):
|
||||
"""HTTP GET .../documents/{id}/chunks returns chunks."""
|
||||
await _retain(api_client, bank_id, "doc-http-chunks", "Alice works at Google on AI research. Bob works at Meta on VR systems. " * 20)
|
||||
|
||||
response = await api_client.get(
|
||||
f"/v1/default/banks/{bank_id}/documents/doc-http-chunks/chunks"
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert "items" in data
|
||||
assert "total" in data
|
||||
assert data["total"] >= 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_http_list_document_chunks_not_found(api_client, bank_id):
|
||||
"""HTTP GET .../documents/{id}/chunks returns 404 for non-existent document."""
|
||||
response = await api_client.get(
|
||||
f"/v1/default/banks/{bank_id}/documents/nonexistent/chunks"
|
||||
)
|
||||
assert response.status_code == 404
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_http_reprocess_document(api_client, bank_id):
|
||||
"""HTTP POST .../documents/{id}/reprocess returns success with operation_id."""
|
||||
await _retain(api_client, bank_id, "doc-http-reprocess", "Alice works at Google.")
|
||||
|
||||
response = await api_client.post(
|
||||
f"/v1/default/banks/{bank_id}/documents/doc-http-reprocess/reprocess"
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["success"] is True
|
||||
assert "operation_id" in data
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_http_reprocess_document_not_found(api_client, bank_id):
|
||||
"""HTTP POST .../documents/{id}/reprocess returns 404 for non-existent document."""
|
||||
response = await api_client.post(
|
||||
f"/v1/default/banks/{bank_id}/documents/nonexistent/reprocess"
|
||||
)
|
||||
assert response.status_code == 404
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_http_get_document_includes_nodes_by_fact_type(api_client, bank_id):
|
||||
"""HTTP GET .../documents/{id} includes nodes_by_fact_type."""
|
||||
await _retain(api_client, bank_id, "doc-http-comp", "Alice works at Google on AI research.")
|
||||
|
||||
response = await api_client.get(
|
||||
f"/v1/default/banks/{bank_id}/documents/doc-http-comp"
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert "nodes_by_fact_type" in data
|
||||
nbt = data["nodes_by_fact_type"]
|
||||
assert "world" in nbt
|
||||
assert "experience" in nbt
|
||||
assert "observation" in nbt
|
||||
@@ -1,150 +0,0 @@
|
||||
"""
|
||||
Tests for HINDSIGHT_API_EMBEDDINGS_OPENAI_BATCH_SIZE config wiring.
|
||||
|
||||
Regression test for issue #1142: `OpenAIEmbeddings` hardcoded `batch_size=100` is
|
||||
incompatible with OpenAI-compatible providers that enforce stricter per-request
|
||||
limits (e.g. DashScope / Aliyun Tongyi cap at 10). Users must be able to override
|
||||
the batch size via env var so `encode()` splits into smaller chunks.
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def setup_test_env():
|
||||
"""Save/restore env vars touched by these tests."""
|
||||
from hindsight_api.config import clear_config_cache
|
||||
|
||||
env_vars_to_save = [
|
||||
"HINDSIGHT_API_EMBEDDINGS_PROVIDER",
|
||||
"HINDSIGHT_API_EMBEDDINGS_OPENAI_API_KEY",
|
||||
"HINDSIGHT_API_EMBEDDINGS_OPENAI_MODEL",
|
||||
"HINDSIGHT_API_EMBEDDINGS_OPENAI_BATCH_SIZE",
|
||||
"HINDSIGHT_API_EMBEDDINGS_OPENROUTER_API_KEY",
|
||||
"HINDSIGHT_API_LLM_API_KEY",
|
||||
"HINDSIGHT_API_LLM_PROVIDER",
|
||||
]
|
||||
|
||||
original_values = {key: os.environ.get(key) for key in env_vars_to_save}
|
||||
|
||||
clear_config_cache()
|
||||
|
||||
yield
|
||||
|
||||
for key, original_value in original_values.items():
|
||||
if original_value is None:
|
||||
os.environ.pop(key, None)
|
||||
else:
|
||||
os.environ[key] = original_value
|
||||
|
||||
clear_config_cache()
|
||||
|
||||
|
||||
def test_default_openai_batch_size_is_100():
|
||||
"""Default batch size is 100 when env var unset (preserves legacy behavior)."""
|
||||
from hindsight_api.config import HindsightConfig
|
||||
|
||||
os.environ["HINDSIGHT_API_LLM_PROVIDER"] = "mock"
|
||||
os.environ.pop("HINDSIGHT_API_EMBEDDINGS_OPENAI_BATCH_SIZE", None)
|
||||
|
||||
config = HindsightConfig.from_env()
|
||||
assert config.embeddings_openai_batch_size == 100
|
||||
|
||||
|
||||
def test_openai_batch_size_env_var_is_read():
|
||||
"""HINDSIGHT_API_EMBEDDINGS_OPENAI_BATCH_SIZE overrides the default."""
|
||||
from hindsight_api.config import HindsightConfig
|
||||
|
||||
os.environ["HINDSIGHT_API_LLM_PROVIDER"] = "mock"
|
||||
os.environ["HINDSIGHT_API_EMBEDDINGS_OPENAI_BATCH_SIZE"] = "10"
|
||||
|
||||
config = HindsightConfig.from_env()
|
||||
assert config.embeddings_openai_batch_size == 10
|
||||
|
||||
|
||||
def test_openai_embeddings_provider_uses_configured_batch_size():
|
||||
"""create_embeddings_from_env() propagates config to OpenAIEmbeddings for 'openai' provider."""
|
||||
from hindsight_api.engine.embeddings import OpenAIEmbeddings, create_embeddings_from_env
|
||||
|
||||
os.environ["HINDSIGHT_API_LLM_PROVIDER"] = "mock"
|
||||
os.environ["HINDSIGHT_API_EMBEDDINGS_PROVIDER"] = "openai"
|
||||
os.environ["HINDSIGHT_API_EMBEDDINGS_OPENAI_API_KEY"] = "sk-test"
|
||||
os.environ["HINDSIGHT_API_EMBEDDINGS_OPENAI_BATCH_SIZE"] = "10"
|
||||
|
||||
embeddings = create_embeddings_from_env()
|
||||
assert isinstance(embeddings, OpenAIEmbeddings)
|
||||
assert embeddings.batch_size == 10
|
||||
|
||||
|
||||
def test_openrouter_provider_uses_configured_batch_size():
|
||||
"""'openrouter' provider also honors the shared batch-size config (both paths use OpenAIEmbeddings)."""
|
||||
from hindsight_api.engine.embeddings import OpenAIEmbeddings, create_embeddings_from_env
|
||||
|
||||
os.environ["HINDSIGHT_API_LLM_PROVIDER"] = "mock"
|
||||
os.environ["HINDSIGHT_API_EMBEDDINGS_PROVIDER"] = "openrouter"
|
||||
os.environ["HINDSIGHT_API_EMBEDDINGS_OPENROUTER_API_KEY"] = "sk-or-test"
|
||||
os.environ["HINDSIGHT_API_EMBEDDINGS_OPENAI_BATCH_SIZE"] = "8"
|
||||
|
||||
embeddings = create_embeddings_from_env()
|
||||
assert isinstance(embeddings, OpenAIEmbeddings)
|
||||
assert embeddings.batch_size == 8
|
||||
|
||||
|
||||
def test_zero_batch_size_is_rejected():
|
||||
"""Zero would cause `range(0, N, 0)` to crash at runtime — fail fast at config load."""
|
||||
from hindsight_api.config import HindsightConfig
|
||||
|
||||
os.environ["HINDSIGHT_API_LLM_PROVIDER"] = "mock"
|
||||
os.environ["HINDSIGHT_API_EMBEDDINGS_OPENAI_BATCH_SIZE"] = "0"
|
||||
|
||||
with pytest.raises(ValueError, match="must be >= 1"):
|
||||
HindsightConfig.from_env()
|
||||
|
||||
|
||||
def test_negative_batch_size_is_rejected():
|
||||
"""Negative values would silently skip batching — reject at config load."""
|
||||
from hindsight_api.config import HindsightConfig
|
||||
|
||||
os.environ["HINDSIGHT_API_LLM_PROVIDER"] = "mock"
|
||||
os.environ["HINDSIGHT_API_EMBEDDINGS_OPENAI_BATCH_SIZE"] = "-5"
|
||||
|
||||
with pytest.raises(ValueError, match="must be >= 1"):
|
||||
HindsightConfig.from_env()
|
||||
|
||||
|
||||
def test_non_numeric_batch_size_is_rejected():
|
||||
"""Non-integer strings are rejected with a clear error pointing at the env var name."""
|
||||
from hindsight_api.config import HindsightConfig
|
||||
|
||||
os.environ["HINDSIGHT_API_LLM_PROVIDER"] = "mock"
|
||||
os.environ["HINDSIGHT_API_EMBEDDINGS_OPENAI_BATCH_SIZE"] = "not-a-number"
|
||||
|
||||
with pytest.raises(ValueError, match="HINDSIGHT_API_EMBEDDINGS_OPENAI_BATCH_SIZE"):
|
||||
HindsightConfig.from_env()
|
||||
|
||||
|
||||
def test_openai_encode_splits_on_configured_batch_size(monkeypatch):
|
||||
"""encode() sends multiple upstream requests when len(texts) > batch_size."""
|
||||
from types import SimpleNamespace
|
||||
|
||||
from hindsight_api.engine.embeddings import OpenAIEmbeddings
|
||||
|
||||
emb = OpenAIEmbeddings(api_key="sk-test", model="text-embedding-3-small", batch_size=10)
|
||||
|
||||
calls: list[int] = []
|
||||
|
||||
def fake_create(*, model, input):
|
||||
calls.append(len(input))
|
||||
return SimpleNamespace(data=[SimpleNamespace(index=i, embedding=[0.0] * 1536) for i in range(len(input))])
|
||||
|
||||
emb._client = SimpleNamespace(embeddings=SimpleNamespace(create=fake_create))
|
||||
emb._dimension = 1536
|
||||
|
||||
vectors = emb.encode(["x"] * 25)
|
||||
|
||||
assert len(vectors) == 25
|
||||
assert calls == [10, 10, 5], (
|
||||
f"Expected upstream calls of size 10, 10, 5 when batch_size=10 and 25 inputs, got {calls}"
|
||||
)
|
||||
@@ -543,196 +543,6 @@ async def test_async_file_retain_serializes_datetime_timestamp(memory_no_llm_ver
|
||||
assert row["timestamp"] == "2024-01-15T10:30:00+00:00"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_file_retain_maps_timestamp_to_event_date(memory_no_llm_verify, sample_txt_content):
|
||||
"""Regression (PR #1092): file retain must translate 'timestamp' -> 'event_date'.
|
||||
|
||||
The retain orchestrator only reads 'event_date' from each content dict.
|
||||
_handle_file_convert_retain previously forwarded 'timestamp' unchanged, so every
|
||||
file-retained memory silently defaulted to utcnow() and the 'unset' sentinel
|
||||
was a no-op. This test intercepts the inner batch_retain task the handler
|
||||
submits and asserts the key mapping is correct for all three inputs:
|
||||
explicit ISO timestamp, 'unset' sentinel, and omitted (None).
|
||||
"""
|
||||
from hindsight_api.engine.parsers.base import FileParser
|
||||
from hindsight_api.models import RequestContext
|
||||
|
||||
memory = memory_no_llm_verify
|
||||
|
||||
class NoopParser(FileParser):
|
||||
async def convert(self, file_data: bytes, filename: str) -> str:
|
||||
return file_data.decode("utf-8")
|
||||
|
||||
def supports(self, filename: str, content_type: str | None = None) -> bool:
|
||||
return filename.endswith(".txt")
|
||||
|
||||
def name(self) -> str:
|
||||
return "event_date_regression_parser"
|
||||
|
||||
memory._parser_registry.register(NoopParser())
|
||||
|
||||
class MockFile:
|
||||
def __init__(self, content, filename, content_type):
|
||||
self.content = content
|
||||
self.filename = filename
|
||||
self.content_type = content_type
|
||||
|
||||
async def read(self):
|
||||
return self.content
|
||||
|
||||
# Capture the inner batch_retain submission from _handle_file_convert_retain so we
|
||||
# can inspect its content dict without running the (LLM-dependent) retain pipeline.
|
||||
original_submit = memory._task_backend.submit_task
|
||||
captured: list[dict] = []
|
||||
|
||||
async def capturing_submit(task_dict):
|
||||
if task_dict.get("type") == "batch_retain":
|
||||
captured.append(task_dict)
|
||||
return
|
||||
await original_submit(task_dict)
|
||||
|
||||
memory._task_backend.submit_task = capturing_submit
|
||||
try:
|
||||
context = RequestContext(internal=True)
|
||||
|
||||
async def run_case(label: str, timestamp_value) -> dict:
|
||||
bank_id = f"test_file_event_date_{label}_{datetime.now(timezone.utc).timestamp()}"
|
||||
await memory.get_bank_profile(bank_id, request_context=context)
|
||||
|
||||
captured.clear()
|
||||
await memory.submit_async_file_retain(
|
||||
bank_id=bank_id,
|
||||
file_items=[
|
||||
{
|
||||
"file": MockFile(sample_txt_content, f"{label}.txt", "text/plain"),
|
||||
"document_id": f"doc_{label}",
|
||||
"context": "regression test",
|
||||
"metadata": {},
|
||||
"tags": [],
|
||||
"timestamp": timestamp_value,
|
||||
"parser": ["event_date_regression_parser"],
|
||||
}
|
||||
],
|
||||
document_tags=None,
|
||||
request_context=context,
|
||||
)
|
||||
assert len(captured) == 1, f"{label}: expected exactly one batch_retain submission"
|
||||
contents = captured[0]["contents"]
|
||||
assert len(contents) == 1
|
||||
return contents[0]
|
||||
|
||||
# Explicit ISO timestamp -> event_date must equal that string.
|
||||
content = await run_case("explicit", "2024-01-15T10:30:00+00:00")
|
||||
assert "timestamp" not in content, "raw 'timestamp' must not leak into retain content"
|
||||
assert content["event_date"] == "2024-01-15T10:30:00+00:00"
|
||||
|
||||
# 'unset' sentinel -> event_date must be explicit None (orchestrator stores NULL).
|
||||
content = await run_case("unset", "unset")
|
||||
assert "timestamp" not in content
|
||||
assert "event_date" in content, "'unset' must produce an explicit event_date=None"
|
||||
assert content["event_date"] is None
|
||||
|
||||
# Omitted timestamp -> event_date key must be absent (orchestrator defaults to utcnow).
|
||||
content = await run_case("missing", None)
|
||||
assert "timestamp" not in content
|
||||
assert "event_date" not in content
|
||||
finally:
|
||||
memory._task_backend.submit_task = original_submit
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_file_retain_forwards_all_content_fields(memory_no_llm_verify, sample_txt_content):
|
||||
"""Regression: _handle_file_convert_retain must forward every FileRetainMetadata
|
||||
field to the inner batch_retain task without renaming or dropping it.
|
||||
|
||||
Covers document_id, context, metadata, tags (per-content), plus strategy
|
||||
and document_tags (per-request). The timestamp -> event_date mapping has
|
||||
its own test above. Existing file retain tests only assert HTTP 200 or
|
||||
inspect the outer file_convert_retain task_payload; none verify what
|
||||
arrives at the retain pipeline. If any of these fields were silently
|
||||
dropped or mis-keyed -- the same failure mode as #1092 for timestamp --
|
||||
those tests would still pass.
|
||||
"""
|
||||
from hindsight_api.engine.parsers.base import FileParser
|
||||
from hindsight_api.models import RequestContext
|
||||
|
||||
memory = memory_no_llm_verify
|
||||
|
||||
class NoopParser(FileParser):
|
||||
async def convert(self, file_data: bytes, filename: str) -> str:
|
||||
return file_data.decode("utf-8")
|
||||
|
||||
def supports(self, filename: str, content_type: str | None = None) -> bool:
|
||||
return filename.endswith(".txt")
|
||||
|
||||
def name(self) -> str:
|
||||
return "all_fields_regression_parser"
|
||||
|
||||
memory._parser_registry.register(NoopParser())
|
||||
|
||||
class MockFile:
|
||||
def __init__(self, content, filename, content_type):
|
||||
self.content = content
|
||||
self.filename = filename
|
||||
self.content_type = content_type
|
||||
|
||||
async def read(self):
|
||||
return self.content
|
||||
|
||||
original_submit = memory._task_backend.submit_task
|
||||
captured: list[dict] = []
|
||||
|
||||
async def capturing_submit(task_dict):
|
||||
if task_dict.get("type") == "batch_retain":
|
||||
captured.append(task_dict)
|
||||
return
|
||||
await original_submit(task_dict)
|
||||
|
||||
memory._task_backend.submit_task = capturing_submit
|
||||
try:
|
||||
request_context = RequestContext(internal=True)
|
||||
bank_id = f"test_file_all_fields_{datetime.now(timezone.utc).timestamp()}"
|
||||
await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
|
||||
await memory.submit_async_file_retain(
|
||||
bank_id=bank_id,
|
||||
file_items=[
|
||||
{
|
||||
"file": MockFile(sample_txt_content, "doc.txt", "text/plain"),
|
||||
"document_id": "my_doc_id",
|
||||
"context": "meeting notes from Alice",
|
||||
"metadata": {"author": "Alice", "year": "2024"},
|
||||
"tags": ["report", "q1"],
|
||||
"timestamp": None,
|
||||
"parser": ["all_fields_regression_parser"],
|
||||
"strategy": "my_strategy",
|
||||
}
|
||||
],
|
||||
document_tags=["batch_tag"],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert len(captured) == 1, "expected exactly one batch_retain submission"
|
||||
payload = captured[0]
|
||||
assert payload["type"] == "batch_retain"
|
||||
|
||||
# Per-request fields (live on the outer task payload, not per-content).
|
||||
assert payload.get("strategy") == "my_strategy", "strategy must be forwarded at request level"
|
||||
assert payload.get("document_tags") == ["batch_tag"], "document_tags must be forwarded at request level"
|
||||
|
||||
# Per-content fields.
|
||||
assert len(payload["contents"]) == 1
|
||||
content = payload["contents"][0]
|
||||
assert content["document_id"] == "my_doc_id"
|
||||
assert content["context"] == "meeting notes from Alice"
|
||||
assert content["metadata"] == {"author": "Alice", "year": "2024"}
|
||||
assert content["tags"] == ["report", "q1"]
|
||||
# content is the converted markdown (raw bytes decoded by NoopParser).
|
||||
assert content["content"] == sample_txt_content.decode("utf-8")
|
||||
finally:
|
||||
memory._task_backend.submit_task = original_submit
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_file_conversion_failure_sets_status_to_failed(memory_no_llm_verify, sample_txt_content):
|
||||
"""Test that when file conversion fails, the operation status is set to 'failed' not 'completed'."""
|
||||
|
||||
@@ -49,7 +49,6 @@ def _make_mock_google_module(mock_genai: MagicMock) -> MagicMock:
|
||||
mod = MagicMock()
|
||||
mod.genai = mock_genai
|
||||
mod.genai.types.EmbedContentConfig = MagicMock(side_effect=lambda **kw: MagicMock(**kw))
|
||||
mod.genai.types.HttpOptions = MagicMock(side_effect=lambda **kw: MagicMock(**kw))
|
||||
return mod
|
||||
|
||||
|
||||
@@ -172,23 +171,6 @@ class TestGeminiEmbeddings:
|
||||
call_kwargs = mock_genai.Client.return_value.models.embed_content.call_args
|
||||
assert "config" not in call_kwargs.kwargs
|
||||
|
||||
async def test_force_ipv4_passes_http_options(self):
|
||||
"""Test that force_ipv4 configures the Gemini client with custom HTTP options."""
|
||||
mock_genai = _make_mock_genai()
|
||||
mock_transport = MagicMock()
|
||||
mock_httpx_client = MagicMock()
|
||||
emb = GeminiEmbeddings(model="gemini-embedding-001", api_key="test-key", force_ipv4=True)
|
||||
|
||||
with _patch_google_import(mock_genai):
|
||||
with patch("httpx.HTTPTransport", return_value=mock_transport) as mock_http_transport:
|
||||
with patch("httpx.Client", return_value=mock_httpx_client) as mock_http_client:
|
||||
await emb.initialize()
|
||||
|
||||
mock_http_transport.assert_called_once_with(local_address="0.0.0.0")
|
||||
mock_http_client.assert_called_once_with(timeout=10, transport=mock_transport)
|
||||
assert emb._httpx_client is mock_httpx_client
|
||||
assert "http_options" in mock_genai.Client.call_args.kwargs
|
||||
|
||||
def test_auto_detect_vertexai(self):
|
||||
"""Test that _is_vertexai is auto-detected from vertexai_project_id."""
|
||||
assert GeminiEmbeddings(model="m", api_key="k")._is_vertexai is False
|
||||
@@ -305,7 +287,6 @@ class TestGeminiEmbeddingsFactory:
|
||||
defaults["embeddings_gemini_api_key"] = "test-key"
|
||||
defaults["embeddings_gemini_model"] = "gemini-embedding-001"
|
||||
defaults["embeddings_gemini_output_dimensionality"] = 768
|
||||
defaults["embeddings_gemini_force_ipv4"] = False
|
||||
defaults["embeddings_vertexai_project_id"] = None
|
||||
defaults["embeddings_vertexai_region"] = None
|
||||
defaults["embeddings_vertexai_service_account_key"] = None
|
||||
@@ -321,14 +302,6 @@ class TestGeminiEmbeddingsFactory:
|
||||
assert emb.provider_name == "google"
|
||||
assert emb.api_key == "test-key"
|
||||
assert emb._is_vertexai is False
|
||||
assert emb.force_ipv4 is False
|
||||
|
||||
def test_create_with_force_ipv4(self):
|
||||
config = self._make_config(embeddings_gemini_force_ipv4=True)
|
||||
with patch("hindsight_api.config.get_config", return_value=config):
|
||||
emb = create_embeddings_from_env()
|
||||
assert isinstance(emb, GeminiEmbeddings)
|
||||
assert emb.force_ipv4 is True
|
||||
|
||||
def test_create_with_vertexai(self):
|
||||
config = self._make_config(
|
||||
|
||||
@@ -98,7 +98,7 @@ async def test_hierarchical_fields_categorization():
|
||||
assert "retain_chunk_batch_size" in configurable
|
||||
|
||||
# Verify count is correct
|
||||
assert len(configurable) == 35
|
||||
assert len(configurable) == 22
|
||||
|
||||
# Verify credential fields (NEVER exposed)
|
||||
assert "llm_api_key" in credentials
|
||||
@@ -458,7 +458,7 @@ async def test_config_get_bank_config_no_static_or_credential_fields_leak(memory
|
||||
assert field in config, f"Expected configurable field '{field}' missing from config"
|
||||
|
||||
# Should have a small number of configurable fields (not hundreds)
|
||||
assert len(config) < 50, f"Too many fields returned: {len(config)}"
|
||||
assert len(config) < 25, f"Too many fields returned: {len(config)}"
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@@ -286,34 +286,6 @@ async def test_full_api_workflow(api_client, test_bank_id):
|
||||
)
|
||||
assert response.status_code == 410 # Deprecated endpoint
|
||||
|
||||
# Entity co-occurrence graph — shape is stable even when there are no
|
||||
# co-occurrences; every edge must reference two nodes that are also present.
|
||||
response = await api_client.get(f"/v1/default/banks/{test_bank_id}/entities/graph")
|
||||
assert response.status_code == 200
|
||||
entity_graph = response.json()
|
||||
assert set(entity_graph.keys()) >= {"nodes", "edges", "total_entities", "total_edges", "limit"}
|
||||
assert entity_graph["limit"] == 1000
|
||||
assert len(entity_graph["nodes"]) == entity_graph["total_entities"]
|
||||
assert len(entity_graph["edges"]) == entity_graph["total_edges"]
|
||||
node_ids = {n["data"]["id"] for n in entity_graph["nodes"]}
|
||||
for edge in entity_graph["edges"]:
|
||||
assert edge["data"]["source"] in node_ids
|
||||
assert edge["data"]["target"] in node_ids
|
||||
assert edge["data"]["linkType"] == "cooccurrence"
|
||||
assert edge["data"]["weight"] >= 1
|
||||
|
||||
# min_count filter — raising the threshold can only shrink the edge set.
|
||||
response = await api_client.get(
|
||||
f"/v1/default/banks/{test_bank_id}/entities/graph?min_count=1000000"
|
||||
)
|
||||
assert response.status_code == 200
|
||||
filtered_graph = response.json()
|
||||
assert filtered_graph["total_edges"] == 0
|
||||
|
||||
# "graph" must route to the graph endpoint, not be parsed as an entity_id.
|
||||
# Regression guard in case someone reorders the FastAPI route registration.
|
||||
assert entity_graph["total_entities"] >= 0
|
||||
|
||||
# ================================================================
|
||||
# 9. List All Banks (should include our test bank)
|
||||
# ================================================================
|
||||
|
||||
@@ -25,8 +25,6 @@ from hindsight_api.engine.llm_wrapper import TokenUsage
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
pytestmark = pytest.mark.xdist_group("load_batch_tests")
|
||||
|
||||
|
||||
def generate_content(char_count: int) -> str:
|
||||
"""Generate realistic content of approximately char_count characters."""
|
||||
@@ -119,18 +117,9 @@ class TestLargeBatchRetain:
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
@pytest.fixture
|
||||
def disable_observations(self):
|
||||
from hindsight_api.config import _get_raw_config
|
||||
config = _get_raw_config()
|
||||
original = config.enable_observations
|
||||
config.enable_observations = False
|
||||
yield
|
||||
config.enable_observations = original
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.timeout(300) # 5 minute timeout
|
||||
async def test_large_batch_500k_chars_20_items(self, memory_with_mock_llm, request_context, disable_observations):
|
||||
async def test_large_batch_500k_chars_20_items(self, memory_with_mock_llm, request_context):
|
||||
"""
|
||||
Test retaining a batch of 20 content items totaling ~500k chars.
|
||||
|
||||
@@ -294,7 +283,7 @@ class TestLargeBatchRetain:
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.timeout(60)
|
||||
async def test_db_connection_pool_under_load(self, memory_with_mock_llm, request_context, disable_observations):
|
||||
async def test_db_connection_pool_under_load(self, memory_with_mock_llm, request_context):
|
||||
"""
|
||||
Test that DB connection pool handles concurrent operations.
|
||||
|
||||
|
||||
@@ -159,6 +159,7 @@ def mock_memory():
|
||||
# Memory browsing methods
|
||||
memory.list_memory_units = AsyncMock(return_value={"items": [{"id": "mem-1", "content": "Test"}], "total": 1})
|
||||
memory.get_memory_unit = AsyncMock(return_value={"id": "mem-1", "content": "Test memory"})
|
||||
memory.delete_memory_unit = AsyncMock(return_value={"deleted_count": 1})
|
||||
|
||||
# Document methods
|
||||
memory.list_documents = AsyncMock(return_value={"items": [{"id": "doc-1", "name": "Test Doc"}], "total": 1})
|
||||
@@ -176,10 +177,6 @@ def mock_memory():
|
||||
memory.get_bank_stats = AsyncMock(return_value={"nodes": 100, "links": 50})
|
||||
memory.update_bank = AsyncMock(return_value={"id": "test-bank", "name": "Updated"})
|
||||
memory.delete_bank = AsyncMock(return_value={"deleted_memories": 10, "deleted_entities": 5})
|
||||
|
||||
# Config resolver (used by update_bank MCP tool for config fields)
|
||||
memory._config_resolver = MagicMock()
|
||||
memory._config_resolver.update_bank_config = AsyncMock()
|
||||
memory.list_banks = AsyncMock(return_value=[])
|
||||
|
||||
return memory
|
||||
@@ -312,6 +309,7 @@ class TestMentalModelToolRegistration:
|
||||
memory.delete_directive = AsyncMock()
|
||||
memory.list_memory_units = AsyncMock(return_value={})
|
||||
memory.get_memory_unit = AsyncMock()
|
||||
memory.delete_memory_unit = AsyncMock()
|
||||
memory.list_documents = AsyncMock(return_value={})
|
||||
memory.get_document = AsyncMock()
|
||||
memory.delete_document = AsyncMock()
|
||||
@@ -345,7 +343,7 @@ class TestMentalModelToolRegistration:
|
||||
assert "delete_bank" in tools
|
||||
assert "clear_memories" in tools
|
||||
assert "sync_retain" in tools
|
||||
assert len(tools) == 29
|
||||
assert len(tools) == 30
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
@@ -1104,6 +1102,12 @@ class TestMemoryBrowsingTools:
|
||||
result = await _tools(mcp)["get_memory"].fn(memory_id="missing")
|
||||
assert "not found" in result
|
||||
|
||||
async def test_delete_memory(self, mock_memory):
|
||||
mcp = _make_mcp_server(mock_memory, {"delete_memory"}, include_bank_id=True)
|
||||
result = await _tools(mcp)["delete_memory"].fn(memory_id="mem-1")
|
||||
assert '"deleted"' in result
|
||||
assert mock_memory.delete_memory_unit.call_args.kwargs["unit_id"] == "mem-1"
|
||||
|
||||
async def test_get_memory_invalid_uuid(self, mock_memory):
|
||||
mock_memory.get_memory_unit.side_effect = ValueError("Invalid memory_id: 'nonexistent' is not a valid UUID")
|
||||
mcp = _make_mcp_server(mock_memory, {"get_memory"}, include_bank_id=True)
|
||||
@@ -1116,6 +1120,12 @@ class TestMemoryBrowsingTools:
|
||||
result = await _tools(mcp)["get_memory"].fn(memory_id="bad")
|
||||
assert "not a valid UUID" in result["error"]
|
||||
|
||||
async def test_delete_memory_invalid_uuid(self, mock_memory):
|
||||
mock_memory.delete_memory_unit.side_effect = ValueError("Invalid unit_id: 'bad' is not a valid UUID")
|
||||
mcp = _make_mcp_server(mock_memory, {"delete_memory"}, include_bank_id=True)
|
||||
result = await _tools(mcp)["delete_memory"].fn(memory_id="bad")
|
||||
assert "not a valid UUID" in result
|
||||
|
||||
async def test_list_memories_single_bank(self, mock_memory):
|
||||
mcp = _make_mcp_server(mock_memory, {"list_memories"}, include_bank_id=False)
|
||||
result = await _tools(mcp)["list_memories"].fn()
|
||||
@@ -1255,14 +1265,9 @@ class TestTagsAndBankTools:
|
||||
async def test_update_bank(self, mock_memory):
|
||||
mcp = _make_mcp_server(mock_memory, {"update_bank"}, include_bank_id=True)
|
||||
result = await _tools(mcp)["update_bank"].fn(name="New Name", mission="New Mission")
|
||||
# name is updated via engine
|
||||
call_kwargs = mock_memory.update_bank.call_args.kwargs
|
||||
assert call_kwargs["name"] == "New Name"
|
||||
# mission is routed to config resolver as reflect_mission
|
||||
config_call = mock_memory._config_resolver.update_bank_config.call_args
|
||||
assert config_call.args[1] == {"reflect_mission": "New Mission"}
|
||||
# bank_id is the first positional arg
|
||||
assert config_call.args[0] == "test-bank"
|
||||
assert call_kwargs["mission"] == "New Mission"
|
||||
|
||||
async def test_delete_bank(self, mock_memory):
|
||||
mcp = _make_mcp_server(mock_memory, {"delete_bank"}, include_bank_id=True)
|
||||
@@ -1347,6 +1352,19 @@ class TestOperationErrorHandling:
|
||||
class TestDeleteErrorHandling:
|
||||
"""Error handling tests for delete operations."""
|
||||
|
||||
async def test_delete_memory_engine_error(self, mock_memory):
|
||||
mock_memory.delete_memory_unit.side_effect = RuntimeError("DB error")
|
||||
mcp = _make_mcp_server(mock_memory, {"delete_memory"}, include_bank_id=True)
|
||||
result = await _tools(mcp)["delete_memory"].fn(memory_id="mem-1")
|
||||
assert "error" in result
|
||||
assert "DB error" in result
|
||||
|
||||
async def test_delete_memory_single_bank(self, mock_memory):
|
||||
mcp = _make_mcp_server(mock_memory, {"delete_memory"}, include_bank_id=False)
|
||||
result = await _tools(mcp)["delete_memory"].fn(memory_id="mem-1")
|
||||
assert isinstance(result, dict)
|
||||
assert result["status"] == "deleted"
|
||||
|
||||
async def test_delete_document_engine_error(self, mock_memory):
|
||||
mock_memory.delete_document.side_effect = RuntimeError("DB error")
|
||||
mcp = _make_mcp_server(mock_memory, {"delete_document"}, include_bank_id=True)
|
||||
@@ -1378,120 +1396,6 @@ class TestUpdateBankVariants:
|
||||
result = await _tools(mcp)["update_bank"].fn(name="X")
|
||||
assert "error" in result
|
||||
|
||||
async def test_update_bank_config_updates_dict(self, mock_memory):
|
||||
"""config_updates dict is passed directly to config resolver."""
|
||||
mcp = _make_mcp_server(mock_memory, {"update_bank"}, include_bank_id=True)
|
||||
await _tools(mcp)["update_bank"].fn(config_updates={"reflect_mission": "Guide reflect output"})
|
||||
config_call = mock_memory._config_resolver.update_bank_config.call_args
|
||||
assert config_call.args[1] == {"reflect_mission": "Guide reflect output"}
|
||||
# name should NOT be updated when not provided
|
||||
mock_memory.update_bank.assert_not_called()
|
||||
|
||||
async def test_update_bank_mission_maps_to_reflect_mission(self, mock_memory):
|
||||
"""Deprecated mission param is mapped to reflect_mission in config."""
|
||||
mcp = _make_mcp_server(mock_memory, {"update_bank"}, include_bank_id=True)
|
||||
await _tools(mcp)["update_bank"].fn(mission="My mission")
|
||||
config_call = mock_memory._config_resolver.update_bank_config.call_args
|
||||
assert config_call.args[1] == {"reflect_mission": "My mission"}
|
||||
|
||||
async def test_update_bank_config_reflect_mission_takes_precedence_over_mission(self, mock_memory):
|
||||
"""When both mission and config_updates.reflect_mission are provided, config wins."""
|
||||
mcp = _make_mcp_server(mock_memory, {"update_bank"}, include_bank_id=True)
|
||||
await _tools(mcp)["update_bank"].fn(
|
||||
mission="old", config_updates={"reflect_mission": "new"}
|
||||
)
|
||||
config_call = mock_memory._config_resolver.update_bank_config.call_args
|
||||
assert config_call.args[1]["reflect_mission"] == "new"
|
||||
|
||||
async def test_update_bank_multiple_config_fields(self, mock_memory):
|
||||
"""Multiple config fields can be set in a single config_updates dict."""
|
||||
mcp = _make_mcp_server(mock_memory, {"update_bank"}, include_bank_id=True)
|
||||
await _tools(mcp)["update_bank"].fn(config_updates={
|
||||
"retain_mission": "Extract technical decisions",
|
||||
"disposition_skepticism": 5,
|
||||
"disposition_literalism": 1,
|
||||
"disposition_empathy": 4,
|
||||
"enable_observations": True,
|
||||
"observations_mission": "Focus on preferences",
|
||||
"retain_extraction_mode": "custom",
|
||||
"retain_custom_instructions": "Extract only action items",
|
||||
"retain_chunk_size": 2000,
|
||||
})
|
||||
config_call = mock_memory._config_resolver.update_bank_config.call_args
|
||||
updates = config_call.args[1]
|
||||
assert updates["retain_mission"] == "Extract technical decisions"
|
||||
assert updates["disposition_skepticism"] == 5
|
||||
assert updates["disposition_literalism"] == 1
|
||||
assert updates["disposition_empathy"] == 4
|
||||
assert updates["enable_observations"] is True
|
||||
assert updates["observations_mission"] == "Focus on preferences"
|
||||
assert updates["retain_extraction_mode"] == "custom"
|
||||
assert updates["retain_custom_instructions"] == "Extract only action items"
|
||||
assert updates["retain_chunk_size"] == 2000
|
||||
|
||||
async def test_update_bank_name_and_config_together(self, mock_memory):
|
||||
"""name goes to engine, config_updates goes to config resolver."""
|
||||
mcp = _make_mcp_server(mock_memory, {"update_bank"}, include_bank_id=True)
|
||||
await _tools(mcp)["update_bank"].fn(
|
||||
name="My Bank",
|
||||
config_updates={"reflect_mission": "Reflect guide", "retain_mission": "Retain guide"},
|
||||
)
|
||||
assert mock_memory.update_bank.call_args.kwargs["name"] == "My Bank"
|
||||
updates = mock_memory._config_resolver.update_bank_config.call_args.args[1]
|
||||
assert updates["reflect_mission"] == "Reflect guide"
|
||||
assert updates["retain_mission"] == "Retain guide"
|
||||
|
||||
async def test_update_bank_no_config_call_when_only_name(self, mock_memory):
|
||||
"""When only name is provided, config resolver should not be called."""
|
||||
mcp = _make_mcp_server(mock_memory, {"update_bank"}, include_bank_id=True)
|
||||
await _tools(mcp)["update_bank"].fn(name="Just Name")
|
||||
mock_memory.update_bank.assert_called_once()
|
||||
mock_memory._config_resolver.update_bank_config.assert_not_called()
|
||||
|
||||
async def test_update_bank_config_updates_single_bank(self, mock_memory):
|
||||
"""config_updates works in single-bank mode too."""
|
||||
mcp = _make_mcp_server(mock_memory, {"update_bank"}, include_bank_id=False)
|
||||
result = await _tools(mcp)["update_bank"].fn(
|
||||
config_updates={"retain_mission": "Extract everything", "disposition_empathy": 5}
|
||||
)
|
||||
assert isinstance(result, dict)
|
||||
config_call = mock_memory._config_resolver.update_bank_config.call_args
|
||||
updates = config_call.args[1]
|
||||
assert updates["retain_mission"] == "Extract everything"
|
||||
assert updates["disposition_empathy"] == 5
|
||||
|
||||
async def test_update_bank_with_bank_id_override(self, mock_memory):
|
||||
"""bank_id override routes config update to the correct bank."""
|
||||
mcp = _make_mcp_server(mock_memory, {"update_bank"}, include_bank_id=True)
|
||||
await _tools(mcp)["update_bank"].fn(
|
||||
config_updates={"reflect_mission": "Test"}, bank_id="other-bank"
|
||||
)
|
||||
config_call = mock_memory._config_resolver.update_bank_config.call_args
|
||||
assert config_call.args[0] == "other-bank"
|
||||
|
||||
async def test_update_bank_config_resolver_validation_error(self, mock_memory):
|
||||
"""ValueError from config resolver (e.g. invalid field) is returned as error."""
|
||||
mock_memory._config_resolver.update_bank_config.side_effect = ValueError(
|
||||
"Cannot override static (server-level) fields: ['database_url']"
|
||||
)
|
||||
mcp = _make_mcp_server(mock_memory, {"update_bank"}, include_bank_id=True)
|
||||
result = await _tools(mcp)["update_bank"].fn(config_updates={"database_url": "bad"})
|
||||
assert "error" in result
|
||||
assert "static" in result
|
||||
|
||||
async def test_update_bank_any_configurable_field(self, mock_memory):
|
||||
"""Any field in _CONFIGURABLE_FIELDS is accepted (future-proof)."""
|
||||
mcp = _make_mcp_server(mock_memory, {"update_bank"}, include_bank_id=True)
|
||||
await _tools(mcp)["update_bank"].fn(config_updates={
|
||||
"recall_budget_fixed_low": 100,
|
||||
"consolidation_llm_batch_size": 8,
|
||||
"entity_labels": ["PERSON", "ORG"],
|
||||
})
|
||||
updates = mock_memory._config_resolver.update_bank_config.call_args.args[1]
|
||||
assert updates["recall_budget_fixed_low"] == 100
|
||||
assert updates["consolidation_llm_batch_size"] == 8
|
||||
assert updates["entity_labels"] == ["PERSON", "ORG"]
|
||||
|
||||
async def test_get_bank_stats_engine_error(self, mock_memory):
|
||||
mock_memory.get_bank_stats.side_effect = RuntimeError("DB error")
|
||||
mcp = _make_mcp_server(mock_memory, {"get_bank_stats"}, include_bank_id=True)
|
||||
|
||||
@@ -1,906 +0,0 @@
|
||||
"""Tests for delta-mode mental model refresh.
|
||||
|
||||
Delta mode performs a surgical update on the existing mental model content:
|
||||
- Unchanged sections are preserved byte-for-byte.
|
||||
- Stale content is removed.
|
||||
- New content from observations/facts is added, preferably by extending existing sections.
|
||||
|
||||
Fallback rules:
|
||||
- If the mental model has no existing content, delta falls back to a full regeneration.
|
||||
- If the source_query has changed since the last refresh, delta falls back to a full regeneration.
|
||||
|
||||
This file contains two kinds of tests:
|
||||
|
||||
1. TestDeltaRefreshPlumbing: fast, deterministic tests that monkey-patch reflect_async
|
||||
and the LLM call to verify branching logic (fallback conditions, provenance tracking).
|
||||
|
||||
2. TestDeltaRefreshGeminiEval: real-LLM behavioral evals against Gemini. These are
|
||||
gated on HINDSIGHT_RUN_GEMINI_EVALS=1 (plus a Gemini API key) because they cost
|
||||
money/time and require network access. They verify the actual quality of delta
|
||||
updates — format preservation, surgical edits, observation-grounding.
|
||||
"""
|
||||
|
||||
import os
|
||||
import uuid
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
from hindsight_api import MemoryEngine, RequestContext
|
||||
from hindsight_api.engine.llm_wrapper import LLMConfig
|
||||
from hindsight_api.engine.response_models import ReflectResult
|
||||
|
||||
|
||||
def _canned_reflect_result(text: str, facts: list[dict] | None = None) -> ReflectResult:
|
||||
"""Build a minimal ReflectResult for monkey-patching reflect_async."""
|
||||
return ReflectResult.model_validate(
|
||||
{
|
||||
"text": text,
|
||||
"based_on": {
|
||||
"observation": facts or [],
|
||||
"world": [],
|
||||
"experience": [],
|
||||
"mental-models": [],
|
||||
"directives": [],
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def patch_reflect(monkeypatch):
|
||||
"""Helper that patches memory.reflect_async to return a canned result and records the call.
|
||||
|
||||
Usage:
|
||||
calls = patch_reflect(memory, text="hello", facts=[...])
|
||||
await memory.refresh_mental_model(...)
|
||||
assert len(calls) == 1
|
||||
"""
|
||||
|
||||
def _install(memory: MemoryEngine, *, text: str, facts: list[dict] | None = None):
|
||||
calls: list[dict] = []
|
||||
|
||||
async def fake_reflect_async(**kwargs):
|
||||
calls.append(kwargs)
|
||||
return _canned_reflect_result(text, facts)
|
||||
|
||||
monkeypatch.setattr(memory, "reflect_async", fake_reflect_async)
|
||||
return calls
|
||||
|
||||
return _install
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def patch_llm_call(monkeypatch):
|
||||
"""Patch the reflect LLM config's ``.call()`` used for the structured delta call.
|
||||
|
||||
The structured-delta path passes ``response_format=DeltaOperationList``, so the
|
||||
LLM returns a Pydantic instance. Each invocation of ``patch_llm_call`` installs
|
||||
a single canned response, in any of these shapes:
|
||||
|
||||
- ``DeltaOperationList`` instance → returned as-is
|
||||
- ``[]`` (empty list) → no operations (this is the no-change case)
|
||||
- ``[{"op": "...", ...}, ...]`` → wrapped into ``{"operations": [...]}``
|
||||
- ``{"operations": [...]}`` → validated directly
|
||||
"""
|
||||
from hindsight_api.engine.reflect.delta_ops import DeltaOperationList
|
||||
|
||||
def _to_op_list(resp: Any) -> DeltaOperationList:
|
||||
if isinstance(resp, DeltaOperationList):
|
||||
return resp
|
||||
if isinstance(resp, dict):
|
||||
if "operations" in resp:
|
||||
return DeltaOperationList.model_validate(resp)
|
||||
# Treat a bare op dict as a one-op list for ergonomics.
|
||||
return DeltaOperationList.model_validate({"operations": [resp]})
|
||||
if isinstance(resp, list):
|
||||
return DeltaOperationList.model_validate({"operations": resp})
|
||||
if isinstance(resp, str):
|
||||
# Tests that expect *no* call ever still install a sentinel; treat as no-op.
|
||||
return DeltaOperationList()
|
||||
raise TypeError(f"unsupported canned LLM response: {type(resp)!r}")
|
||||
|
||||
def _install(memory: MemoryEngine, *, returns):
|
||||
calls: list[dict] = []
|
||||
canned = _to_op_list(returns)
|
||||
|
||||
async def fake_call(*, messages, **kwargs):
|
||||
calls.append({"messages": messages, **kwargs})
|
||||
return canned
|
||||
|
||||
monkeypatch.setattr(memory._reflect_llm_config, "call", fake_call)
|
||||
return calls
|
||||
|
||||
return _install
|
||||
|
||||
|
||||
class TestDeltaRefreshPlumbing:
|
||||
"""Deterministic tests that verify the branching/plumbing of delta-mode refresh."""
|
||||
|
||||
async def test_full_mode_does_not_call_delta_merge(
|
||||
self,
|
||||
memory: MemoryEngine,
|
||||
request_context: RequestContext,
|
||||
patch_reflect,
|
||||
patch_llm_call,
|
||||
):
|
||||
"""When trigger.mode='full', no second LLM call for delta merge occurs."""
|
||||
bank_id = f"test-delta-full-{uuid.uuid4().hex[:8]}"
|
||||
await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
|
||||
mm = await memory.create_mental_model(
|
||||
bank_id=bank_id,
|
||||
name="Team Info",
|
||||
source_query="Tell me about the team",
|
||||
content="# Team\n\nOriginal content.",
|
||||
trigger={"mode": "full"},
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
patch_reflect(memory, text="# Team\n\nRegenerated from scratch.")
|
||||
llm_calls = patch_llm_call(memory, returns="should-not-be-called")
|
||||
|
||||
refreshed = await memory.refresh_mental_model(
|
||||
bank_id=bank_id, mental_model_id=mm["id"], request_context=request_context
|
||||
)
|
||||
|
||||
assert refreshed is not None
|
||||
assert refreshed["content"] == "# Team\n\nRegenerated from scratch."
|
||||
assert len(llm_calls) == 0, "Delta merge LLM call must not happen in full mode"
|
||||
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
async def test_delta_mode_empty_content_falls_back_to_full(
|
||||
self,
|
||||
memory: MemoryEngine,
|
||||
request_context: RequestContext,
|
||||
patch_reflect,
|
||||
patch_llm_call,
|
||||
):
|
||||
"""When the mental model has no existing content there is nothing to anchor
|
||||
a surgical edit on, so delta falls back to full regeneration. The user's
|
||||
candidate from reflect_async is used verbatim.
|
||||
"""
|
||||
bank_id = f"test-delta-empty-{uuid.uuid4().hex[:8]}"
|
||||
await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
|
||||
mm = await memory.create_mental_model(
|
||||
bank_id=bank_id,
|
||||
name="Team Info",
|
||||
source_query="Tell me about the team",
|
||||
content="", # no existing content
|
||||
trigger={"mode": "delta"},
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
patch_reflect(memory, text="# Team\n\nFull fresh synthesis.")
|
||||
llm_calls = patch_llm_call(memory, returns=[])
|
||||
|
||||
refreshed = await memory.refresh_mental_model(
|
||||
bank_id=bank_id, mental_model_id=mm["id"], request_context=request_context
|
||||
)
|
||||
|
||||
assert refreshed["content"] == "# Team\n\nFull fresh synthesis."
|
||||
assert len(llm_calls) == 0 # delta path skipped entirely
|
||||
rr = refreshed.get("reflect_response") or {}
|
||||
assert rr.get("delta_applied") is not True
|
||||
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
async def test_delta_mode_source_query_change_falls_back_to_full(
|
||||
self,
|
||||
memory: MemoryEngine,
|
||||
request_context: RequestContext,
|
||||
patch_reflect,
|
||||
patch_llm_call,
|
||||
):
|
||||
"""If source_query changes after a refresh, the next delta run must do a full rewrite."""
|
||||
bank_id = f"test-delta-query-change-{uuid.uuid4().hex[:8]}"
|
||||
await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
|
||||
mm = await memory.create_mental_model(
|
||||
bank_id=bank_id,
|
||||
name="Team Info",
|
||||
source_query="Tell me about the team",
|
||||
content="# Team\n\nBaseline.",
|
||||
trigger={"mode": "delta"},
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# First refresh: establishes last_refreshed_source_query.
|
||||
patch_reflect(memory, text="# Team\n\nFirst pass.")
|
||||
patch_llm_call(memory, returns="unused-first")
|
||||
await memory.refresh_mental_model(
|
||||
bank_id=bank_id, mental_model_id=mm["id"], request_context=request_context
|
||||
)
|
||||
|
||||
# Now change the source_query — a genuine topic shift.
|
||||
await memory.update_mental_model(
|
||||
bank_id=bank_id,
|
||||
mental_model_id=mm["id"],
|
||||
source_query="Tell me about customers instead",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Second refresh under the new query must do a FULL rewrite, not a delta merge.
|
||||
patch_reflect(memory, text="# Customers\n\nBrand new topic.")
|
||||
llm_calls = patch_llm_call(memory, returns="should-not-be-called")
|
||||
|
||||
refreshed = await memory.refresh_mental_model(
|
||||
bank_id=bank_id, mental_model_id=mm["id"], request_context=request_context
|
||||
)
|
||||
|
||||
assert refreshed["content"] == "# Customers\n\nBrand new topic."
|
||||
assert len(llm_calls) == 0, "Source-query change must bypass the delta merge"
|
||||
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
async def test_delta_mode_applies_ops_when_query_stable(
|
||||
self,
|
||||
memory: MemoryEngine,
|
||||
request_context: RequestContext,
|
||||
patch_reflect,
|
||||
patch_llm_call,
|
||||
):
|
||||
"""When content exists and source_query is stable, the delta LLM produces ops
|
||||
that are applied against the parsed structured doc. The unchanged section
|
||||
renders byte-identical, the new fact lands in a new block.
|
||||
"""
|
||||
bank_id = f"test-delta-apply-{uuid.uuid4().hex[:8]}"
|
||||
await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
|
||||
existing = (
|
||||
"# Team\n"
|
||||
"\n"
|
||||
"Alice is the lead.\n"
|
||||
"\n"
|
||||
"## Members\n"
|
||||
"\n"
|
||||
"- Alice — lead\n"
|
||||
)
|
||||
mm = await memory.create_mental_model(
|
||||
bank_id=bank_id,
|
||||
name="Team Info",
|
||||
source_query="Tell me about the team",
|
||||
content=existing,
|
||||
trigger={"mode": "delta"},
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# First refresh: empty op list → structured doc unchanged → markdown is the
|
||||
# render of the parsed existing content. This also seeds the tracking column.
|
||||
patch_reflect(memory, text="ignored — full mode candidate")
|
||||
patch_llm_call(memory, returns=[]) # zero ops
|
||||
await memory.refresh_mental_model(
|
||||
bank_id=bank_id, mental_model_id=mm["id"], request_context=request_context
|
||||
)
|
||||
|
||||
# Second refresh: a new fact arrives; LLM returns one append_block op.
|
||||
candidate = "# Team\n\nAlice is the lead. Bob joined as junior engineer."
|
||||
patch_reflect(
|
||||
memory,
|
||||
text=candidate,
|
||||
facts=[
|
||||
{
|
||||
"id": "obs-bob",
|
||||
"text": "Bob joined the team as junior engineer",
|
||||
"type": "observation",
|
||||
"context": None,
|
||||
}
|
||||
],
|
||||
)
|
||||
ops = [
|
||||
{
|
||||
"op": "append_block",
|
||||
"section_id": "members",
|
||||
"block": {
|
||||
"type": "bullet_list",
|
||||
"items": ["Bob — junior engineer"],
|
||||
},
|
||||
}
|
||||
]
|
||||
llm_calls = patch_llm_call(memory, returns=ops)
|
||||
|
||||
refreshed = await memory.refresh_mental_model(
|
||||
bank_id=bank_id, mental_model_id=mm["id"], request_context=request_context
|
||||
)
|
||||
|
||||
assert len(llm_calls) == 1, "Structured-delta LLM call must fire exactly once"
|
||||
system_msg = llm_calls[0]["messages"][0]["content"]
|
||||
user_msg = llm_calls[0]["messages"][1]["content"]
|
||||
# Prompt must include the structured doc + supporting facts + the system prompt.
|
||||
assert "integrating" in system_msg.lower()
|
||||
assert "operations" in system_msg.lower()
|
||||
assert "obs-bob" in user_msg
|
||||
assert "Bob joined" in user_msg
|
||||
# The structured JSON of the current doc must include the section id "members".
|
||||
assert '"id": "members"' in user_msg
|
||||
|
||||
# New content includes the new bullet.
|
||||
assert "Bob — junior engineer" in refreshed["content"]
|
||||
# Unchanged section ("Alice is the lead.") still present.
|
||||
assert "Alice is the lead." in refreshed["content"]
|
||||
rr = refreshed.get("reflect_response") or {}
|
||||
assert rr.get("delta_applied") is True
|
||||
applied = rr.get("delta_operations_applied") or []
|
||||
assert len(applied) == 1
|
||||
assert applied[0]["op"] == "append_block"
|
||||
assert applied[0]["section_id"] == "members"
|
||||
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
async def test_delta_zero_ops_keeps_existing_content_byte_identical(
|
||||
self,
|
||||
memory: MemoryEngine,
|
||||
request_context: RequestContext,
|
||||
patch_reflect,
|
||||
patch_llm_call,
|
||||
):
|
||||
"""Zero operations from the LLM must mean zero changes in the rendered output.
|
||||
|
||||
This is the structural guarantee: any sections/blocks not mentioned by an
|
||||
op come through byte-identical. A no-op refresh therefore re-renders the
|
||||
same structured doc — which (after the first refresh has parsed and
|
||||
re-rendered it) is byte-stable.
|
||||
"""
|
||||
bank_id = f"test-delta-noop-{uuid.uuid4().hex[:8]}"
|
||||
await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
|
||||
existing = (
|
||||
"# Team\n"
|
||||
"\n"
|
||||
"Alice is the lead.\n"
|
||||
"\n"
|
||||
"## Members\n"
|
||||
"\n"
|
||||
"- Alice\n"
|
||||
)
|
||||
mm = await memory.create_mental_model(
|
||||
bank_id=bank_id,
|
||||
name="Team Info",
|
||||
source_query="Tell me about the team",
|
||||
content=existing,
|
||||
trigger={"mode": "delta"},
|
||||
request_context=request_context,
|
||||
)
|
||||
# First refresh: parses + renders existing into structured form. The output
|
||||
# may not match `existing` byte-for-byte (whitespace normalised by renderer).
|
||||
patch_reflect(memory, text="ignored — full mode candidate")
|
||||
patch_llm_call(memory, returns=[])
|
||||
first = await memory.refresh_mental_model(
|
||||
bank_id=bank_id, mental_model_id=mm["id"], request_context=request_context
|
||||
)
|
||||
normalised = first["content"]
|
||||
|
||||
# Second refresh: zero ops again → same bytes as first refresh.
|
||||
# Must include at least one fact so the no-new-facts short-circuit doesn't fire.
|
||||
patch_reflect(
|
||||
memory,
|
||||
text="something completely different from existing",
|
||||
facts=[{"id": "obs-1", "text": "irrelevant", "type": "observation", "context": None}],
|
||||
)
|
||||
patch_llm_call(memory, returns=[])
|
||||
second = await memory.refresh_mental_model(
|
||||
bank_id=bank_id, mental_model_id=mm["id"], request_context=request_context
|
||||
)
|
||||
assert second["content"] == normalised
|
||||
rr = second.get("reflect_response") or {}
|
||||
assert rr.get("delta_applied") is True # delta path ran; produced no changes
|
||||
assert rr.get("delta_operations_applied") == []
|
||||
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
async def test_delta_llm_failure_falls_back_to_candidate(
|
||||
self,
|
||||
memory: MemoryEngine,
|
||||
request_context: RequestContext,
|
||||
patch_reflect,
|
||||
monkeypatch,
|
||||
):
|
||||
"""When the structured-delta LLM call raises, refresh falls back to the
|
||||
candidate markdown so the user still sees a fresh synthesis instead of
|
||||
an opaque failure.
|
||||
"""
|
||||
bank_id = f"test-delta-llm-fail-{uuid.uuid4().hex[:8]}"
|
||||
await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
|
||||
mm = await memory.create_mental_model(
|
||||
bank_id=bank_id,
|
||||
name="Team Info",
|
||||
source_query="Tell me about the team",
|
||||
content="# Team\n\nExisting.\n",
|
||||
trigger={"mode": "delta"},
|
||||
request_context=request_context,
|
||||
)
|
||||
# Seed tracking column with a successful zero-op refresh.
|
||||
patch_reflect(memory, text="ignored")
|
||||
|
||||
async def ok_call(*, messages, **kwargs):
|
||||
from hindsight_api.engine.reflect.delta_ops import DeltaOperationList
|
||||
|
||||
return DeltaOperationList()
|
||||
|
||||
monkeypatch.setattr(memory._reflect_llm_config, "call", ok_call)
|
||||
await memory.refresh_mental_model(
|
||||
bank_id=bank_id, mental_model_id=mm["id"], request_context=request_context
|
||||
)
|
||||
|
||||
# Now the second refresh: LLM raises. Refresh must not crash; it should
|
||||
# store the candidate markdown.
|
||||
candidate = "# Team\n\nFallback candidate from reflect_async.\n"
|
||||
patch_reflect(
|
||||
memory,
|
||||
text=candidate,
|
||||
facts=[{"id": "obs-new", "text": "some new fact", "type": "observation", "context": None}],
|
||||
)
|
||||
|
||||
async def boom(*, messages, **kwargs):
|
||||
raise RuntimeError("simulated provider 500")
|
||||
|
||||
monkeypatch.setattr(memory._reflect_llm_config, "call", boom)
|
||||
refreshed = await memory.refresh_mental_model(
|
||||
bank_id=bank_id, mental_model_id=mm["id"], request_context=request_context
|
||||
)
|
||||
|
||||
assert "Fallback candidate" in refreshed["content"]
|
||||
rr = refreshed.get("reflect_response") or {}
|
||||
assert rr.get("delta_applied") is False
|
||||
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
async def test_empty_reflect_answer_preserves_existing_content(
|
||||
self,
|
||||
memory: MemoryEngine,
|
||||
request_context: RequestContext,
|
||||
patch_reflect,
|
||||
patch_llm_call,
|
||||
monkeypatch,
|
||||
):
|
||||
"""Regression: when the reflect agent returns an empty answer (small models
|
||||
sometimes hit this after exhausting tool-call retries), the refresh must
|
||||
NOT overwrite the existing content with an empty string.
|
||||
|
||||
Previously this destroyed the working document on every transient upstream
|
||||
failure, and the next refresh saw current_content == "" and skipped the
|
||||
delta path entirely — a snowball that emptied valuable mental models.
|
||||
|
||||
The scenario covered here is the realistic failure path: the structured
|
||||
delta call also fails (because the empty supporting facts produce empty
|
||||
/ invalid JSON) so the fallback path kicks in. Without the guard, the
|
||||
fallback would write "" to the DB; with it, the existing content stays.
|
||||
"""
|
||||
bank_id = f"test-empty-reflect-{uuid.uuid4().hex[:8]}"
|
||||
await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
|
||||
existing = (
|
||||
"# Team\n"
|
||||
"\n"
|
||||
"Alice is the lead.\n"
|
||||
"\n"
|
||||
"## Members\n"
|
||||
"\n"
|
||||
"- Alice\n"
|
||||
)
|
||||
mm = await memory.create_mental_model(
|
||||
bank_id=bank_id,
|
||||
name="Team Info",
|
||||
source_query="Tell me about the team",
|
||||
content=existing,
|
||||
trigger={"mode": "delta"},
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Reflect returns "" — this is the upstream failure mode.
|
||||
# Must include at least one fact so the no-new-facts short-circuit doesn't fire.
|
||||
patch_reflect(
|
||||
memory,
|
||||
text="",
|
||||
facts=[{"id": "obs-new", "text": "some fact", "type": "observation", "context": None}],
|
||||
)
|
||||
|
||||
# Delta call also fails (mirrors the real groq behaviour where empty
|
||||
# supporting facts often produce empty / invalid JSON). Refresh then
|
||||
# falls back to the empty candidate, which the guard rejects.
|
||||
async def boom(*, messages, **kwargs):
|
||||
raise RuntimeError("simulated empty/invalid JSON from provider")
|
||||
|
||||
monkeypatch.setattr(memory._reflect_llm_config, "call", boom)
|
||||
|
||||
refreshed = await memory.refresh_mental_model(
|
||||
bank_id=bank_id, mental_model_id=mm["id"], request_context=request_context
|
||||
)
|
||||
|
||||
# Existing content preserved exactly.
|
||||
assert refreshed["content"] == existing, (
|
||||
"Empty reflect answer overwrote existing content — guard regressed"
|
||||
)
|
||||
rr = refreshed.get("reflect_response") or {}
|
||||
assert rr.get("refresh_skipped") == "empty_candidate"
|
||||
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Real-Gemini evaluation tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_GEMINI_API_KEY = (
|
||||
os.getenv("HINDSIGHT_GEMINI_API_KEY")
|
||||
or os.getenv("GEMINI_API_KEY")
|
||||
or os.getenv("GOOGLE_API_KEY")
|
||||
)
|
||||
_OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
|
||||
_RUN_LLM_EVAL = os.getenv("HINDSIGHT_RUN_GEMINI_EVALS") == "1" and (
|
||||
bool(_GEMINI_API_KEY) or bool(_OPENAI_API_KEY)
|
||||
)
|
||||
|
||||
|
||||
pytestmark_gemini = pytest.mark.skipif(
|
||||
not _RUN_LLM_EVAL,
|
||||
reason=(
|
||||
"Real-LLM delta evals are gated. Set HINDSIGHT_RUN_GEMINI_EVALS=1 and provide "
|
||||
"GEMINI_API_KEY (preferred) or OPENAI_API_KEY to run."
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def gemini_memory(memory_no_llm_verify: MemoryEngine):
|
||||
"""MemoryEngine wired to a real LLM for reflect + structured delta.
|
||||
|
||||
Prefers Gemini (the original target) but falls back to OpenAI when the
|
||||
Gemini key is unavailable — the structured-delta architecture works
|
||||
against either, and waiting on a single provider's key would block
|
||||
iteration. The chosen model is logged so test failures are unambiguous
|
||||
about which provider produced them.
|
||||
"""
|
||||
if _GEMINI_API_KEY:
|
||||
provider = "gemini"
|
||||
model = os.getenv("HINDSIGHT_GEMINI_EVAL_MODEL", "gemini-2.0-flash")
|
||||
cfg = LLMConfig(provider=provider, api_key=_GEMINI_API_KEY, base_url="", model=model)
|
||||
else:
|
||||
provider = "openai"
|
||||
model = os.getenv("HINDSIGHT_OPENAI_EVAL_MODEL", "gpt-4o-mini")
|
||||
cfg = LLMConfig(provider=provider, api_key=_OPENAI_API_KEY or "", base_url="", model=model)
|
||||
print(f"\n[delta-eval] using provider={provider} model={model}")
|
||||
memory_no_llm_verify._reflect_llm_config = cfg
|
||||
memory_no_llm_verify._llm_config = cfg
|
||||
memory_no_llm_verify._retain_llm_config = cfg
|
||||
memory_no_llm_verify._consolidation_llm_config = cfg
|
||||
yield memory_no_llm_verify
|
||||
|
||||
|
||||
_NEWS_FEED_SKILL_MARKDOWN = """## Purpose
|
||||
|
||||
Generate a concise, top-N personalized AI/ML news brief in response to user-triggered requests such as "ai news", "top 5 this week", or "what matters for builders today".
|
||||
|
||||
## Scope
|
||||
|
||||
- **In scope**: collecting, filtering, and summarizing AI/ML articles from user-preferred RSS feeds, applying user preferences stored in the AI News Feed Preferences mental model, and delivering the brief to the user.
|
||||
- **Out of scope**: non-AI news, detailed article content, legal or privacy reviews beyond user preferences, and posting the brief to external platforms without explicit user approval.
|
||||
|
||||
## Rules
|
||||
|
||||
- **Always**:
|
||||
1. Use the AI News Feed Preferences mental model to retrieve user preferences; do not embed preferences in the skill file.
|
||||
2. Do not post the brief to any platform unless the user explicitly approves.
|
||||
3. Do not persist preferences locally; rely solely on the mental model.
|
||||
4. Refresh the feed after consolidation if the trigger-refresh-after-consolidation flag is true.
|
||||
- **Prefer**:
|
||||
1. Provide a concise summary (about 2-3 sentences per article) for the top-N articles.
|
||||
2. Default to the top-5 articles unless the user specifies otherwise.
|
||||
3. Order articles chronologically or by relevance as per user preference.
|
||||
4. Highlight any user-specified topics or tags if present.
|
||||
|
||||
## Procedure
|
||||
|
||||
1. **Trigger detection** — identify a request containing keywords like "ai news", "top N", or "what matters".
|
||||
2. **Preference retrieval** — call memory recall for the AI News Feed Preferences mental model to obtain RSS feed URLs and any filtering criteria.
|
||||
3. **Feed consolidation** — fetch all feeds, de-duplicate entries, and apply any user-specified filters.
|
||||
4. **Article selection** — choose the top-N articles based on date or user preference; if trigger-refresh-after-consolidation is true, re-fetch feeds before selection.
|
||||
5. **Summarization** — generate a brief summary for each article, keeping it short and to the point.
|
||||
6. **Approval check** — if the brief is to be posted externally, verify explicit user approval; otherwise, deliver it directly to the user.
|
||||
7. **Memory retention** — store any new learnings or preferences observed during the task using memory retain.
|
||||
|
||||
## Inputs and Context
|
||||
|
||||
- **Source feeds**: user-specified RSS URLs stored in the mental model (e.g., https://aiagentmemory.org/index.xml).
|
||||
- **Time window**: the latest update from each feed; typically the last 7 days for weekly briefs.
|
||||
- **User preferences**: stored in the AI News Feed Preferences mental model; may include topics, tags, or language.
|
||||
|
||||
## Output Shape
|
||||
|
||||
- **Structure**: list of articles with title, publication date, source, and a 2-sentence summary.
|
||||
- **Format**: plain text or markdown (as requested by the user).
|
||||
- **Length**: concise — approximately 2-3 sentences per article; total brief about 200-300 words for top-5.
|
||||
- **Voice/Tone**: neutral, informative, and concise; use bullet points for clarity.
|
||||
|
||||
## Stop Conditions
|
||||
|
||||
- If the mental model cannot be retrieved, refuse or request clarification.
|
||||
- If the user has not provided any RSS feed URLs, ask for a preferred source.
|
||||
- If the brief is requested for posting and explicit approval is missing, refuse.
|
||||
- If the user explicitly requests to remove a skill or stop the briefing, comply immediately.
|
||||
|
||||
## Open Questions
|
||||
|
||||
- Desired brief length or word count?
|
||||
- Preferred summary style (bullet vs paragraph).
|
||||
- Whether the user wants to include non-AI but AI-related topics.
|
||||
- Frequency or schedule for automated briefs (if any).
|
||||
- Specific user-defined tags or topics to highlight.
|
||||
"""
|
||||
|
||||
|
||||
@pytestmark_gemini
|
||||
class TestDeltaRefreshGeminiEval:
|
||||
"""Real-LLM evals for the structured-delta refresh path.
|
||||
|
||||
The structural guarantee these tests verify: sections and blocks not
|
||||
targeted by an LLM-emitted operation are byte-identical between the
|
||||
pre-refresh and post-refresh markdown render. This is what the
|
||||
structured-ops architecture buys us — the LLM cannot drift on text it
|
||||
never re-emits.
|
||||
|
||||
Real Gemini is used (not a mock) because the failure mode we're guarding
|
||||
against is precisely "the LLM doesn't reliably do what the prompt says,
|
||||
even at temperature 0". Mocked output would prove the wiring works but
|
||||
not that the contract holds against an actual model.
|
||||
"""
|
||||
|
||||
async def _seed(
|
||||
self,
|
||||
memory: MemoryEngine,
|
||||
request_context: RequestContext,
|
||||
bank_id: str,
|
||||
existing_markdown: str,
|
||||
memories: list[str],
|
||||
) -> dict[str, Any]:
|
||||
await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
mm = await memory.create_mental_model(
|
||||
bank_id=bank_id,
|
||||
name="Skill Doc",
|
||||
source_query="Document the news-feed skill: purpose, rules, procedure, stop conditions.",
|
||||
content=existing_markdown,
|
||||
trigger={"mode": "delta"},
|
||||
request_context=request_context,
|
||||
)
|
||||
await memory.retain_batch_async(
|
||||
bank_id=bank_id,
|
||||
contents=[{"content": m} for m in memories],
|
||||
request_context=request_context,
|
||||
)
|
||||
await memory.wait_for_background_tasks()
|
||||
# First refresh: parses existing into structured form. With well-aligned
|
||||
# memories the LLM should emit zero ops, so the structured doc is just
|
||||
# the parsed existing content. The rendered markdown is canonicalised.
|
||||
first = await memory.refresh_mental_model(
|
||||
bank_id=bank_id, mental_model_id=mm["id"], request_context=request_context
|
||||
)
|
||||
return {"mm": mm, "first": first}
|
||||
|
||||
async def test_no_change_when_observations_agree_with_existing(
|
||||
self, gemini_memory: MemoryEngine, request_context: RequestContext
|
||||
):
|
||||
"""When observations only restate the existing doc, a second delta
|
||||
refresh produces output byte-identical to the first refresh's output.
|
||||
|
||||
The first refresh canonicalises whitespace via the parser+renderer; we
|
||||
compare the *second* refresh against the *first* (not against the raw
|
||||
seed markdown), which is the actual repeat-refresh behaviour users
|
||||
will see in production.
|
||||
"""
|
||||
bank_id = f"eval-delta-noop-{uuid.uuid4().hex[:8]}"
|
||||
seeded = await self._seed(
|
||||
gemini_memory,
|
||||
request_context,
|
||||
bank_id,
|
||||
existing_markdown=_NEWS_FEED_SKILL_MARKDOWN,
|
||||
memories=[
|
||||
"The news-feed skill produces a concise top-N AI/ML news brief.",
|
||||
"Default brief size is top 5 unless the user specifies otherwise.",
|
||||
"Source feed: https://aiagentmemory.org/index.xml.",
|
||||
"The skill must not post externally without explicit approval.",
|
||||
],
|
||||
)
|
||||
first_content = seeded["first"]["content"]
|
||||
|
||||
second = await gemini_memory.refresh_mental_model(
|
||||
bank_id=bank_id,
|
||||
mental_model_id=seeded["mm"]["id"],
|
||||
request_context=request_context,
|
||||
)
|
||||
second_content = second["content"]
|
||||
|
||||
# Byte-identical render across refreshes when no new fact has arrived.
|
||||
assert second_content == first_content, (
|
||||
"Repeat delta refresh changed bytes when no new facts arrived.\n"
|
||||
f"--- diff sample (first 300 chars different) ---\n"
|
||||
f"first: {first_content[:300]!r}\n"
|
||||
f"second: {second_content[:300]!r}"
|
||||
)
|
||||
rr = second.get("reflect_response") or {}
|
||||
# The LLM may emit zero ops (best case) or non-effective ops (still no
|
||||
# change to render); both are acceptable so long as the bytes match.
|
||||
assert rr.get("delta_applied") is True
|
||||
|
||||
await gemini_memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
async def test_new_observation_is_merged_surgically(
|
||||
self, gemini_memory: MemoryEngine, request_context: RequestContext
|
||||
):
|
||||
"""A new fact arrives; only the section relevant to it should change.
|
||||
|
||||
Asserts the architectural guarantee at the section level: every
|
||||
section that the LLM did NOT name in an operation must render exactly
|
||||
the same bytes after the refresh as before. The new fact itself must
|
||||
appear somewhere in the output.
|
||||
"""
|
||||
from hindsight_api.engine.reflect.structured_doc import (
|
||||
StructuredDocument,
|
||||
render_section,
|
||||
)
|
||||
|
||||
bank_id = f"eval-delta-add-{uuid.uuid4().hex[:8]}"
|
||||
seeded = await self._seed(
|
||||
gemini_memory,
|
||||
request_context,
|
||||
bank_id,
|
||||
existing_markdown=_NEWS_FEED_SKILL_MARKDOWN,
|
||||
memories=[
|
||||
"The news-feed skill produces a concise top-N AI/ML news brief.",
|
||||
"Default brief size is top 5.",
|
||||
"Source feed: https://aiagentmemory.org/index.xml.",
|
||||
],
|
||||
)
|
||||
first_content = seeded["first"]["content"]
|
||||
first_struct = StructuredDocument.model_validate(
|
||||
seeded["first"]["reflect_response"]["delta_operations_applied"]
|
||||
and seeded["first"].get("structured_content")
|
||||
or {"version": 1, "sections": []}
|
||||
)
|
||||
# The first refresh's structured snapshot is what the second refresh
|
||||
# will operate on. Re-fetch via get_mental_model would also work.
|
||||
# For preservation comparison we re-parse first_content.
|
||||
from hindsight_api.engine.reflect.structured_doc import parse_markdown
|
||||
|
||||
before = parse_markdown(first_content)
|
||||
|
||||
# Introduce a brand-new fact that fits into "Inputs and Context" or
|
||||
# similar — but the model may pick any reasonable section.
|
||||
await gemini_memory.retain_batch_async(
|
||||
bank_id=bank_id,
|
||||
contents=[
|
||||
{
|
||||
"content": (
|
||||
"The default time window for the news brief is the last 7 days, "
|
||||
"matching the weekly cadence preferred by the user."
|
||||
)
|
||||
},
|
||||
],
|
||||
request_context=request_context,
|
||||
)
|
||||
await gemini_memory.wait_for_background_tasks()
|
||||
|
||||
refreshed = await gemini_memory.refresh_mental_model(
|
||||
bank_id=bank_id,
|
||||
mental_model_id=seeded["mm"]["id"],
|
||||
request_context=request_context,
|
||||
)
|
||||
content = refreshed["content"]
|
||||
rr = refreshed.get("reflect_response") or {}
|
||||
applied_ops = rr.get("delta_operations_applied") or []
|
||||
touched_section_ids = {op.get("section_id") for op in applied_ops if op.get("section_id")}
|
||||
|
||||
# The fact must show up.
|
||||
assert "7 days" in content or "seven days" in content.lower(), (
|
||||
f"New fact about 7-day window missing from delta output: {content!r}"
|
||||
)
|
||||
|
||||
# Every untouched section must render byte-identical to its pre-refresh form.
|
||||
after = parse_markdown(content)
|
||||
before_by_id = {s.id: s for s in before.sections}
|
||||
for section in after.sections:
|
||||
if section.id in touched_section_ids:
|
||||
continue
|
||||
orig = before_by_id.get(section.id)
|
||||
if orig is None:
|
||||
continue # newly added section, no preservation contract
|
||||
assert render_section(orig) == render_section(section), (
|
||||
f"Untouched section {section.id!r} drifted between refreshes — the "
|
||||
f"structured-ops architecture's preservation guarantee was violated.\n"
|
||||
f"BEFORE:\n{render_section(orig)!r}\n"
|
||||
f"AFTER:\n{render_section(section)!r}"
|
||||
)
|
||||
|
||||
assert rr.get("delta_applied") is True
|
||||
|
||||
await gemini_memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
async def test_no_change_repeated_three_times_stays_byte_stable(
|
||||
self, gemini_memory: MemoryEngine, request_context: RequestContext
|
||||
):
|
||||
"""Three consecutive no-change refreshes must produce three identical
|
||||
markdown outputs. This is the regression test for the original
|
||||
complaint where prose-merge delta drifted content across versions even
|
||||
when no observation changed.
|
||||
"""
|
||||
bank_id = f"eval-delta-stable-{uuid.uuid4().hex[:8]}"
|
||||
seeded = await self._seed(
|
||||
gemini_memory,
|
||||
request_context,
|
||||
bank_id,
|
||||
existing_markdown=_NEWS_FEED_SKILL_MARKDOWN,
|
||||
memories=[
|
||||
"The news-feed skill produces a top-N AI brief on demand.",
|
||||
"It must not post without explicit user approval.",
|
||||
],
|
||||
)
|
||||
c1 = seeded["first"]["content"]
|
||||
r2 = await gemini_memory.refresh_mental_model(
|
||||
bank_id=bank_id,
|
||||
mental_model_id=seeded["mm"]["id"],
|
||||
request_context=request_context,
|
||||
)
|
||||
r3 = await gemini_memory.refresh_mental_model(
|
||||
bank_id=bank_id,
|
||||
mental_model_id=seeded["mm"]["id"],
|
||||
request_context=request_context,
|
||||
)
|
||||
assert r2["content"] == c1, "second refresh drifted vs first"
|
||||
assert r3["content"] == c1, "third refresh drifted vs first"
|
||||
|
||||
await gemini_memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
async def test_source_query_change_forces_full_rewrite(
|
||||
self, gemini_memory: MemoryEngine, request_context: RequestContext
|
||||
):
|
||||
"""Changing source_query must bypass delta and produce a full regeneration."""
|
||||
bank_id = f"eval-delta-query-change-{uuid.uuid4().hex[:8]}"
|
||||
await gemini_memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
|
||||
mm = await gemini_memory.create_mental_model(
|
||||
bank_id=bank_id,
|
||||
name="Subject",
|
||||
source_query="Summarize the team and how it operates.",
|
||||
content="# Team Overview\n\nAlice leads the team.\n",
|
||||
trigger={"mode": "delta"},
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
await gemini_memory.retain_batch_async(
|
||||
bank_id=bank_id,
|
||||
contents=[
|
||||
{"content": "Alice leads the team."},
|
||||
{"content": "The product is a memory system for AI agents."},
|
||||
{"content": "Customers include small SaaS startups and enterprise pilots."},
|
||||
],
|
||||
request_context=request_context,
|
||||
)
|
||||
await gemini_memory.wait_for_background_tasks()
|
||||
|
||||
# First refresh seeds tracking column under the team query.
|
||||
await gemini_memory.refresh_mental_model(
|
||||
bank_id=bank_id, mental_model_id=mm["id"], request_context=request_context
|
||||
)
|
||||
|
||||
# Change the topic entirely.
|
||||
await gemini_memory.update_mental_model(
|
||||
bank_id=bank_id,
|
||||
mental_model_id=mm["id"],
|
||||
source_query="Summarize our customers and what we sell them.",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
refreshed = await gemini_memory.refresh_mental_model(
|
||||
bank_id=bank_id, mental_model_id=mm["id"], request_context=request_context
|
||||
)
|
||||
content = refreshed["content"].lower()
|
||||
# Content should now be about customers/product, not (only) about Alice leading the team.
|
||||
assert "customer" in content or "product" in content, (
|
||||
f"Full rewrite should cover the new topic, got: {refreshed['content']!r}"
|
||||
)
|
||||
# delta_applied should be absent/False because we took the full path.
|
||||
assert (refreshed.get("reflect_response") or {}).get("delta_applied") is not True
|
||||
|
||||
await gemini_memory.delete_bank(bank_id, request_context=request_context)
|
||||
@@ -8,8 +8,7 @@ import uuid
|
||||
|
||||
import pytest
|
||||
|
||||
from hindsight_api.engine.memory_engine import MemoryEngine, fq_table
|
||||
from hindsight_api.engine.retain import embedding_utils
|
||||
from hindsight_api.engine.memory_engine import MemoryEngine
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
@@ -689,51 +688,6 @@ class TestMentalModelHistory:
|
||||
assert len(history) == 1
|
||||
assert history[0]["previous_content"] == "Original content"
|
||||
assert "changed_at" in history[0]
|
||||
assert "previous_reflect_response" in history[0]
|
||||
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
async def test_history_snapshots_previous_reflect_response(
|
||||
self, memory: MemoryEngine, request_context
|
||||
):
|
||||
"""Each history entry snapshots the reflect_response that produced previous_content."""
|
||||
bank_id = f"test-mm-history-reflect-{uuid.uuid4().hex[:8]}"
|
||||
await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
|
||||
mm = await memory.create_mental_model(
|
||||
bank_id=bank_id,
|
||||
name="Test Model",
|
||||
source_query="What is the test?",
|
||||
content="v1",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
rr_v1 = {"text": "v1", "based_on": {"observation": [{"id": "o1", "text": "obs1"}]}, "mental_models": []}
|
||||
await memory.update_mental_model(
|
||||
bank_id=bank_id,
|
||||
mental_model_id=mm["id"],
|
||||
content="v2",
|
||||
reflect_response=rr_v1,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
rr_v2 = {"text": "v2", "based_on": {"observation": [{"id": "o2", "text": "obs2"}]}, "mental_models": []}
|
||||
await memory.update_mental_model(
|
||||
bank_id=bank_id,
|
||||
mental_model_id=mm["id"],
|
||||
content="v3",
|
||||
reflect_response=rr_v2,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
history = await memory.get_mental_model_history(bank_id, mm["id"], request_context=request_context)
|
||||
assert len(history) == 2
|
||||
# Most recent first: replacing v2 snapshotted rr_v1 (the reflect that produced v2).
|
||||
assert history[0]["previous_content"] == "v2"
|
||||
assert history[0]["previous_reflect_response"] == rr_v1
|
||||
# The first update replaced v1, which had no reflect_response stored yet.
|
||||
assert history[1]["previous_content"] == "v1"
|
||||
assert history[1]["previous_reflect_response"] is None
|
||||
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@@ -809,222 +763,6 @@ class TestMentalModelHistory:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
class TestMentalModelStaleness:
|
||||
"""Tests for compute_mental_model_is_stale scope semantics.
|
||||
|
||||
Memories are inserted directly into ``memory_units`` so the scenarios don't
|
||||
depend on the LLM fact-extraction pipeline.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
async def _insert_memory(
|
||||
memory: MemoryEngine,
|
||||
bank_id: str,
|
||||
*,
|
||||
tags: list[str] | None = None,
|
||||
fact_type: str = "experience",
|
||||
) -> str:
|
||||
from datetime import datetime, timezone
|
||||
|
||||
pool = await memory._get_pool()
|
||||
mem_id = str(uuid.uuid4())
|
||||
now = datetime.now(timezone.utc)
|
||||
async with pool.acquire() as conn:
|
||||
await conn.execute(
|
||||
f"""
|
||||
INSERT INTO {fq_table("memory_units")}
|
||||
(id, bank_id, text, event_date, fact_type, tags, created_at)
|
||||
VALUES ($1, $2, $3, $4, $5, $6::varchar[], $4)
|
||||
""",
|
||||
mem_id,
|
||||
bank_id,
|
||||
"test memory",
|
||||
now,
|
||||
fact_type,
|
||||
tags if tags is not None else [],
|
||||
)
|
||||
return mem_id
|
||||
|
||||
async def test_fresh_mental_model_is_not_stale(self, memory: MemoryEngine, request_context):
|
||||
bank_id = f"test-mm-stale-fresh-{uuid.uuid4().hex[:8]}"
|
||||
await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
mm = await memory.create_mental_model(
|
||||
bank_id=bank_id, name="MM", source_query="q", content="c", request_context=request_context
|
||||
)
|
||||
got = await memory.get_mental_model(bank_id, mm["id"], request_context=request_context)
|
||||
assert got["is_stale"] is False
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
async def test_untagged_mm_stale_on_any_new_memory(
|
||||
self, memory: MemoryEngine, request_context
|
||||
):
|
||||
bank_id = f"test-mm-stale-untagged-{uuid.uuid4().hex[:8]}"
|
||||
await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
mm = await memory.create_mental_model(
|
||||
bank_id=bank_id, name="MM", source_query="q", content="c", request_context=request_context
|
||||
)
|
||||
await self._insert_memory(memory, bank_id, tags=["something"])
|
||||
got = await memory.get_mental_model(bank_id, mm["id"], request_context=request_context)
|
||||
assert got["is_stale"] is True
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
async def test_tagged_mm_ignores_out_of_scope_memory(
|
||||
self, memory: MemoryEngine, request_context
|
||||
):
|
||||
bank_id = f"test-mm-stale-oos-{uuid.uuid4().hex[:8]}"
|
||||
await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
mm = await memory.create_mental_model(
|
||||
bank_id=bank_id,
|
||||
name="MM",
|
||||
source_query="q",
|
||||
content="c",
|
||||
tags=["user_a"],
|
||||
request_context=request_context,
|
||||
)
|
||||
# Memory tagged with unrelated tag → not in scope, MM should not be stale
|
||||
await self._insert_memory(memory, bank_id, tags=["user_b"])
|
||||
got = await memory.get_mental_model(bank_id, mm["id"], request_context=request_context)
|
||||
assert got["is_stale"] is False
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
async def test_tagged_mm_stale_on_overlapping_memory(
|
||||
self, memory: MemoryEngine, request_context
|
||||
):
|
||||
bank_id = f"test-mm-stale-overlap-{uuid.uuid4().hex[:8]}"
|
||||
await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
mm = await memory.create_mental_model(
|
||||
bank_id=bank_id,
|
||||
name="MM",
|
||||
source_query="q",
|
||||
content="c",
|
||||
tags=["user_a"],
|
||||
request_context=request_context,
|
||||
)
|
||||
await self._insert_memory(memory, bank_id, tags=["user_a", "extra"])
|
||||
got = await memory.get_mental_model(bank_id, mm["id"], request_context=request_context)
|
||||
assert got["is_stale"] is True
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
async def test_tags_match_all_strict_requires_all_tags(
|
||||
self, memory: MemoryEngine, request_context
|
||||
):
|
||||
"""tags_match='all_strict' → memory must contain ALL MM tags (and be tagged)."""
|
||||
bank_id = f"test-mm-stale-all-{uuid.uuid4().hex[:8]}"
|
||||
await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
mm = await memory.create_mental_model(
|
||||
bank_id=bank_id,
|
||||
name="MM",
|
||||
source_query="q",
|
||||
content="c",
|
||||
tags=["user_a", "proj_x"],
|
||||
trigger={"refresh_after_consolidation": False, "tags_match": "all_strict"},
|
||||
request_context=request_context,
|
||||
)
|
||||
# Memory only has one of the tags → does NOT match all_strict
|
||||
await self._insert_memory(memory, bank_id, tags=["user_a"])
|
||||
got = await memory.get_mental_model(bank_id, mm["id"], request_context=request_context)
|
||||
assert got["is_stale"] is False, "all_strict must require ALL MM tags"
|
||||
|
||||
# Now add a memory with both tags → matches
|
||||
await self._insert_memory(memory, bank_id, tags=["user_a", "proj_x"])
|
||||
got = await memory.get_mental_model(bank_id, mm["id"], request_context=request_context)
|
||||
assert got["is_stale"] is True
|
||||
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
async def test_tags_match_any_strict_excludes_untagged(
|
||||
self, memory: MemoryEngine, request_context
|
||||
):
|
||||
"""tags_match='any_strict' → untagged memory does NOT keep MM in scope."""
|
||||
bank_id = f"test-mm-stale-anystrict-{uuid.uuid4().hex[:8]}"
|
||||
await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
mm = await memory.create_mental_model(
|
||||
bank_id=bank_id,
|
||||
name="MM",
|
||||
source_query="q",
|
||||
content="c",
|
||||
tags=["user_a"],
|
||||
trigger={"refresh_after_consolidation": False, "tags_match": "any_strict"},
|
||||
request_context=request_context,
|
||||
)
|
||||
await self._insert_memory(memory, bank_id, tags=None)
|
||||
got = await memory.get_mental_model(bank_id, mm["id"], request_context=request_context)
|
||||
assert got["is_stale"] is False
|
||||
|
||||
await self._insert_memory(memory, bank_id, tags=["user_a"])
|
||||
got = await memory.get_mental_model(bank_id, mm["id"], request_context=request_context)
|
||||
assert got["is_stale"] is True
|
||||
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
async def test_fact_type_filter_narrows_scope(
|
||||
self, memory: MemoryEngine, request_context
|
||||
):
|
||||
bank_id = f"test-mm-stale-fact-{uuid.uuid4().hex[:8]}"
|
||||
await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
mm = await memory.create_mental_model(
|
||||
bank_id=bank_id,
|
||||
name="MM",
|
||||
source_query="q",
|
||||
content="c",
|
||||
trigger={"refresh_after_consolidation": False, "fact_types": ["world"]},
|
||||
request_context=request_context,
|
||||
)
|
||||
# Out-of-scope fact_type → not stale
|
||||
await self._insert_memory(memory, bank_id, fact_type="experience")
|
||||
got = await memory.get_mental_model(bank_id, mm["id"], request_context=request_context)
|
||||
assert got["is_stale"] is False
|
||||
|
||||
# Matching fact_type → stale
|
||||
await self._insert_memory(memory, bank_id, fact_type="world")
|
||||
got = await memory.get_mental_model(bank_id, mm["id"], request_context=request_context)
|
||||
assert got["is_stale"] is True
|
||||
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
async def test_tool_search_mental_models_returns_is_stale_per_mm(
|
||||
self, memory: MemoryEngine, request_context
|
||||
):
|
||||
"""Regression: tool_search_mental_models must compute is_stale per-MM via scope,
|
||||
not via a bank-wide pending_consolidation short-circuit."""
|
||||
from hindsight_api.engine.reflect.tools import tool_search_mental_models
|
||||
|
||||
bank_id = f"test-mm-stale-tool-{uuid.uuid4().hex[:8]}"
|
||||
await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
fresh = await memory.create_mental_model(
|
||||
bank_id=bank_id,
|
||||
name="fresh MM",
|
||||
source_query="q",
|
||||
content="fresh",
|
||||
tags=["user_b"],
|
||||
request_context=request_context,
|
||||
)
|
||||
stale = await memory.create_mental_model(
|
||||
bank_id=bank_id,
|
||||
name="stale MM",
|
||||
source_query="q",
|
||||
content="stale",
|
||||
tags=["user_a"],
|
||||
request_context=request_context,
|
||||
)
|
||||
# Memory only in user_a's scope → only `stale` MM should be flagged.
|
||||
await self._insert_memory(memory, bank_id, tags=["user_a"])
|
||||
|
||||
pool = await memory._get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
embedding = (
|
||||
await embedding_utils.generate_embeddings_batch(memory.embeddings, ["q"])
|
||||
)[0]
|
||||
result = await tool_search_mental_models(
|
||||
memory, conn, bank_id, "q", embedding, max_results=10
|
||||
)
|
||||
by_id = {m["id"]: m for m in result["mental_models"]}
|
||||
assert by_id[fresh["id"]]["is_stale"] is False
|
||||
assert by_id[stale["id"]]["is_stale"] is True
|
||||
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
class TestMentalModelRefreshTagSecurity:
|
||||
"""Test that mental model refresh respects tag-based security boundaries."""
|
||||
|
||||
@@ -1515,147 +1253,6 @@ class TestMentalModelTriggerTagsConfig:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
class TestMentalModelRefreshMaxTokens:
|
||||
"""Verify that refresh_mental_model honors the per-model max_tokens column.
|
||||
|
||||
These tests mock the engine's collaborators so we can assert the exact kwargs
|
||||
passed to reflect_async without spinning up a DB or LLM. The bug being guarded
|
||||
against: the per-model ``max_tokens`` column was ignored during refresh, so
|
||||
reflect_async fell back to its default (4096) and the generated content could
|
||||
exceed the user-configured limit when there were many facts to synthesize.
|
||||
"""
|
||||
|
||||
async def test_refresh_passes_stored_max_tokens_to_reflect(self, request_context):
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
from hindsight_api.engine.memory_engine import MemoryEngine
|
||||
from hindsight_api.engine.response_models import ReflectResult
|
||||
|
||||
custom_max_tokens = 777
|
||||
mental_model = {
|
||||
"id": "mm-1",
|
||||
"bank_id": "bank-1",
|
||||
"name": "Capped Model",
|
||||
"source_query": "Summarize the facts",
|
||||
"content": "initial",
|
||||
"tags": None,
|
||||
"max_tokens": custom_max_tokens,
|
||||
"trigger": {"refresh_after_consolidation": False},
|
||||
}
|
||||
|
||||
engine = MemoryEngine.__new__(MemoryEngine)
|
||||
engine._authenticate_tenant = AsyncMock(return_value=None) # type: ignore[method-assign]
|
||||
engine.get_mental_model = AsyncMock(return_value=mental_model) # type: ignore[method-assign]
|
||||
engine.reflect_async = AsyncMock( # type: ignore[method-assign]
|
||||
return_value=ReflectResult(text="stub synthesis", based_on={})
|
||||
)
|
||||
engine.update_mental_model = AsyncMock(return_value=mental_model) # type: ignore[method-assign]
|
||||
|
||||
await engine.refresh_mental_model(
|
||||
bank_id="bank-1",
|
||||
mental_model_id="mm-1",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert engine.reflect_async.await_count == 1
|
||||
kwargs = engine.reflect_async.await_args.kwargs
|
||||
assert kwargs.get("max_tokens") == custom_max_tokens, (
|
||||
f"refresh_mental_model should forward the stored max_tokens ({custom_max_tokens}) "
|
||||
f"to reflect_async, but got max_tokens={kwargs.get('max_tokens')!r}"
|
||||
)
|
||||
|
||||
async def test_refresh_content_respects_max_tokens(self, memory: MemoryEngine, request_context):
|
||||
"""End-to-end: refreshed content must stay within the model's max_tokens cap.
|
||||
|
||||
We seed the bank with enough varied facts that an unconstrained synthesis
|
||||
would happily produce a long answer, then refresh a mental model with a
|
||||
small max_tokens and assert the resulting content is actually within the
|
||||
cap (with a small tolerance for cross-tokenizer drift, since the LLM may
|
||||
not use cl100k_base).
|
||||
"""
|
||||
from hindsight_api.engine.memory_engine import count_tokens
|
||||
|
||||
bank_id = f"test-refresh-cap-{uuid.uuid4().hex[:8]}"
|
||||
await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
|
||||
# Seed enough content that an uncapped reflect would produce a long answer.
|
||||
await memory.retain_batch_async(
|
||||
bank_id=bank_id,
|
||||
contents=[
|
||||
{"content": (
|
||||
"Alice is the staff frontend engineer. She owns the design system, "
|
||||
"leads accessibility reviews, mentors three junior engineers, and runs "
|
||||
"the weekly UI guild meeting every Thursday at 2pm Pacific."
|
||||
)},
|
||||
{"content": (
|
||||
"Bob is the backend tech lead. He owns the payments service, the "
|
||||
"billing reconciliation pipeline, and the on-call rotation for the "
|
||||
"platform team. He is the primary reviewer for any database migration."
|
||||
)},
|
||||
{"content": (
|
||||
"Carol manages the data platform. Her team operates the warehouse, "
|
||||
"the streaming ingestion layer, and the metrics pipeline that feeds "
|
||||
"the executive dashboards refreshed every fifteen minutes."
|
||||
)},
|
||||
{"content": (
|
||||
"The team holds a company-wide demo every other Friday. Engineering "
|
||||
"presents shipped work, design walks through prototypes, and product "
|
||||
"shares roadmap updates for the upcoming quarter."
|
||||
)},
|
||||
{"content": (
|
||||
"Dan is the security lead. He runs the quarterly threat-modeling "
|
||||
"exercises, owns the incident response runbook, and coordinates the "
|
||||
"annual external penetration test with the vendor."
|
||||
)},
|
||||
{"content": (
|
||||
"Erin runs developer experience. She maintains the local-dev tooling, "
|
||||
"the CI pipelines, the release automation, and the internal "
|
||||
"documentation portal that everyone uses to onboard new hires."
|
||||
)},
|
||||
],
|
||||
request_context=request_context,
|
||||
)
|
||||
await memory.wait_for_background_tasks()
|
||||
|
||||
cap = 200
|
||||
mm = await memory.create_mental_model(
|
||||
bank_id=bank_id,
|
||||
name="Team Summary (capped)",
|
||||
source_query="Give me a complete overview of every team member, what they own, and the recurring meetings.",
|
||||
content="initial",
|
||||
max_tokens=cap,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
refreshed = await memory.refresh_mental_model(
|
||||
bank_id=bank_id,
|
||||
mental_model_id=mm["id"],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert refreshed is not None
|
||||
content = refreshed["content"]
|
||||
assert content, "refresh produced empty content"
|
||||
|
||||
# The provider enforces the cap exactly in its own tokenizer, but our
|
||||
# local count uses tiktoken (cl100k_base) which can disagree with
|
||||
# provider tokenizers (Gemini's SentencePiece in particular tends to run
|
||||
# ~30% higher for English prose). We use a generous tolerance — the test
|
||||
# is guarding against the regression where the cap was ignored entirely
|
||||
# and content grew toward reflect_async's default of 4096 tokens. At
|
||||
# cap=200 we've observed cl100k counts up to ~1.9x; the 4096-ignored
|
||||
# regression would land ~20x, so a wide tolerance still catches it.
|
||||
observed_tokens = count_tokens(content)
|
||||
tolerance = 2.5
|
||||
assert observed_tokens <= cap * tolerance, (
|
||||
f"refreshed content exceeds max_tokens cap: "
|
||||
f"observed≈{observed_tokens} tokens, cap={cap} (tolerance x{tolerance}). "
|
||||
f"content={content!r}"
|
||||
)
|
||||
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
class TestMentalModelTriggerSchema:
|
||||
"""Unit tests for MentalModelTrigger schema validation (no DB needed)."""
|
||||
|
||||
|
||||
@@ -8,7 +8,6 @@ These tests verify that:
|
||||
resets the target memory itself for re-consolidation
|
||||
4. delete_bank(fact_type=...) also cleans up affected observations
|
||||
"""
|
||||
|
||||
import uuid
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
@@ -21,7 +20,6 @@ from hindsight_api.engine.memory_engine import MemoryEngine
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def _insert_memory(conn, bank_id: str, text: str, fact_type: str = "experience") -> uuid.UUID:
|
||||
"""Insert a memory unit directly, bypassing LLM retain pipeline."""
|
||||
mem_id = uuid.uuid4()
|
||||
@@ -38,7 +36,9 @@ async def _insert_memory(conn, bank_id: str, text: str, fact_type: str = "experi
|
||||
return mem_id
|
||||
|
||||
|
||||
async def _insert_observation(conn, bank_id: str, text: str, source_memory_ids: list[uuid.UUID]) -> uuid.UUID:
|
||||
async def _insert_observation(
|
||||
conn, bank_id: str, text: str, source_memory_ids: list[uuid.UUID]
|
||||
) -> uuid.UUID:
|
||||
"""Insert an observation unit directly."""
|
||||
obs_id = uuid.uuid4()
|
||||
await conn.execute(
|
||||
@@ -79,8 +79,8 @@ async def _ensure_bank(memory: MemoryEngine, bank_id: str, request_context: Requ
|
||||
# Tests: delete_memory_unit
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestDeleteMemoryUnitObservationCleanup:
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deleting_source_memory_removes_observation(
|
||||
self, memory: MemoryEngine, request_context: RequestContext
|
||||
@@ -207,10 +207,12 @@ class TestDeleteMemoryUnitObservationCleanup:
|
||||
# Tests: delete_document
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestDeleteDocumentObservationCleanup:
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deleting_document_removes_observations(self, memory: MemoryEngine, request_context: RequestContext):
|
||||
async def test_deleting_document_removes_observations(
|
||||
self, memory: MemoryEngine, request_context: RequestContext
|
||||
):
|
||||
"""Deleting a document removes observations derived from its memory units."""
|
||||
bank_id = f"test-invalidate-doc-{uuid.uuid4().hex[:8]}"
|
||||
await _ensure_bank(memory, bank_id, request_context)
|
||||
@@ -246,7 +248,9 @@ class TestDeleteDocumentObservationCleanup:
|
||||
m3 = await _insert_memory(conn, bank_id, "Alice is an avid outdoor person.")
|
||||
|
||||
# Observation referencing both doc memories and the standalone memory
|
||||
obs_id = await _insert_observation(conn, bank_id, "Alice enjoys outdoor activities.", [m1, m2, m3])
|
||||
obs_id = await _insert_observation(
|
||||
conn, bank_id, "Alice enjoys outdoor activities.", [m1, m2, m3]
|
||||
)
|
||||
|
||||
# Patch out consolidation so it doesn't re-set consolidated_at before we can check it
|
||||
with patch.object(memory, "submit_async_consolidation", new=AsyncMock()):
|
||||
@@ -263,117 +267,12 @@ class TestDeleteDocumentObservationCleanup:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests: document upsert via retain pipeline (regression for orphan observations)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class TestDocumentUpsertObservationCleanup:
|
||||
"""Regression: re-ingesting a document via the retain pipeline must clean up
|
||||
observations derived from the outgoing memory_units, the same way the
|
||||
explicit ``MemoryEngine.delete_document`` API does.
|
||||
|
||||
Before the fix, ``fact_storage.handle_document_tracking`` deleted the
|
||||
document via FK cascade — removing the source memory_units silently — but
|
||||
never invalidated the dependent observations. They became orphans whose
|
||||
``source_memory_ids`` arrays pointed at IDs that no longer existed in
|
||||
``memory_units``.
|
||||
"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upsert_document_removes_observations_from_outgoing_memories(
|
||||
self, memory: MemoryEngine, request_context: RequestContext
|
||||
):
|
||||
from hindsight_api.engine.retain.fact_storage import handle_document_tracking
|
||||
|
||||
bank_id = f"test-upsert-obs-cleanup-{uuid.uuid4().hex[:8]}"
|
||||
await _ensure_bank(memory, bank_id, request_context)
|
||||
|
||||
pool = await memory._get_pool()
|
||||
doc_id = str(uuid.uuid4())
|
||||
|
||||
# Pre-populate: one document, two source memories under it, one
|
||||
# standalone memory not in the document, and an observation that joins
|
||||
# all three. After the upsert, the two doc memories should be gone
|
||||
# (cascade) AND the observation should be invalidated (the bug we're
|
||||
# fixing). The standalone memory should be reset for re-consolidation.
|
||||
async with pool.acquire() as conn:
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO documents (id, bank_id, original_text, content_hash, created_at, updated_at)
|
||||
VALUES ($1, $2, 'old version', 'hash-old', NOW(), NOW())
|
||||
""",
|
||||
doc_id,
|
||||
bank_id,
|
||||
)
|
||||
doc_mem_a = uuid.uuid4()
|
||||
doc_mem_b = uuid.uuid4()
|
||||
for mem_id, text in [(doc_mem_a, "Old fact A."), (doc_mem_b, "Old fact B.")]:
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO memory_units (id, bank_id, text, fact_type, event_date, document_id,
|
||||
created_at, updated_at, consolidated_at)
|
||||
VALUES ($1, $2, $3, 'experience', NOW(), $4, NOW(), NOW(), NOW())
|
||||
""",
|
||||
mem_id,
|
||||
bank_id,
|
||||
text,
|
||||
doc_id,
|
||||
)
|
||||
standalone_mem = await _insert_memory(conn, bank_id, "Standalone fact C.")
|
||||
obs_id = await _insert_observation(
|
||||
conn,
|
||||
bank_id,
|
||||
"Aggregated observation joining doc + standalone facts.",
|
||||
[doc_mem_a, doc_mem_b, standalone_mem],
|
||||
)
|
||||
|
||||
# Trigger the upsert path directly. ``handle_document_tracking`` is
|
||||
# what the retain orchestrator calls on every document re-ingest.
|
||||
async with pool.acquire() as conn:
|
||||
async with conn.transaction():
|
||||
await handle_document_tracking(
|
||||
conn,
|
||||
bank_id=bank_id,
|
||||
document_id=doc_id,
|
||||
combined_content="new version replacing old facts",
|
||||
is_first_batch=True,
|
||||
retain_params=None,
|
||||
document_tags=None,
|
||||
)
|
||||
|
||||
async with pool.acquire() as conn:
|
||||
obs_ids = await _get_observation_ids(conn, bank_id)
|
||||
assert str(obs_id) not in obs_ids, (
|
||||
"Observation derived from the outgoing memory_units should have been "
|
||||
"deleted during the upsert (regression: orphan observations were "
|
||||
"previously left behind because handle_document_tracking didn't call "
|
||||
"delete_stale_observations_for_memories)"
|
||||
)
|
||||
|
||||
# The standalone memory survives (different document_id) and should
|
||||
# be reset for re-consolidation since one of its observations was
|
||||
# invalidated by the upsert.
|
||||
consolidated_at = await _get_consolidated_at(conn, standalone_mem)
|
||||
assert consolidated_at is None, (
|
||||
"Surviving co-source memory should be reset for re-consolidation"
|
||||
)
|
||||
|
||||
# The two doc-scoped memories are gone via FK cascade.
|
||||
doc_mem_count = await conn.fetchval(
|
||||
"SELECT COUNT(*) FROM memory_units WHERE id = ANY($1::uuid[])",
|
||||
[doc_mem_a, doc_mem_b],
|
||||
)
|
||||
assert doc_mem_count == 0
|
||||
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests: delete_bank with fact_type filter
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestDeleteBankByTypeObservationCleanup:
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_clearing_experience_memories_removes_affected_observations(
|
||||
self, memory: MemoryEngine, request_context: RequestContext
|
||||
@@ -386,7 +285,9 @@ class TestDeleteBankByTypeObservationCleanup:
|
||||
async with pool.acquire() as conn:
|
||||
exp1 = await _insert_memory(conn, bank_id, "Alice went hiking last week.", "experience")
|
||||
world1 = await _insert_memory(conn, bank_id, "Alice is a hiker.", "world")
|
||||
obs_id = await _insert_observation(conn, bank_id, "Alice is a regular hiker.", [exp1, world1])
|
||||
obs_id = await _insert_observation(
|
||||
conn, bank_id, "Alice is a regular hiker.", [exp1, world1]
|
||||
)
|
||||
|
||||
# Patch out consolidation so it doesn't re-set consolidated_at before we can check it
|
||||
with patch.object(memory, "submit_async_consolidation", new=AsyncMock()):
|
||||
@@ -429,8 +330,8 @@ class TestDeleteBankByTypeObservationCleanup:
|
||||
# Tests: clear_observations_for_memory
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestClearObservationsForMemory:
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_clears_observations_and_resets_all_source_memories(
|
||||
self, memory: MemoryEngine, request_context: RequestContext
|
||||
@@ -447,7 +348,9 @@ class TestClearObservationsForMemory:
|
||||
|
||||
# Patch out consolidation so it doesn't re-set consolidated_at before we can check it
|
||||
with patch.object(memory, "submit_async_consolidation", new=AsyncMock()):
|
||||
result = await memory.clear_observations_for_memory(bank_id, str(m1), request_context=request_context)
|
||||
result = await memory.clear_observations_for_memory(
|
||||
bank_id, str(m1), request_context=request_context
|
||||
)
|
||||
|
||||
assert result["deleted_count"] == 1
|
||||
|
||||
@@ -462,7 +365,9 @@ class TestClearObservationsForMemory:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_observations_returns_zero(self, memory: MemoryEngine, request_context: RequestContext):
|
||||
async def test_no_observations_returns_zero(
|
||||
self, memory: MemoryEngine, request_context: RequestContext
|
||||
):
|
||||
"""Returns 0 when the memory has no associated observations."""
|
||||
bank_id = f"test-clear-obs-noop-{uuid.uuid4().hex[:8]}"
|
||||
await _ensure_bank(memory, bank_id, request_context)
|
||||
@@ -471,7 +376,9 @@ class TestClearObservationsForMemory:
|
||||
async with pool.acquire() as conn:
|
||||
m1 = await _insert_memory(conn, bank_id, "Alice loves hiking.")
|
||||
|
||||
result = await memory.clear_observations_for_memory(bank_id, str(m1), request_context=request_context)
|
||||
result = await memory.clear_observations_for_memory(
|
||||
bank_id, str(m1), request_context=request_context
|
||||
)
|
||||
|
||||
assert result["deleted_count"] == 0
|
||||
|
||||
@@ -498,7 +405,9 @@ class TestClearObservationsForMemory:
|
||||
obs1_id = await _insert_observation(conn, bank_id, "Alice is an avid hiker.", [m1, m2])
|
||||
obs2_id = await _insert_observation(conn, bank_id, "Alice is a mountaineer.", [m3])
|
||||
|
||||
result = await memory.clear_observations_for_memory(bank_id, str(m1), request_context=request_context)
|
||||
result = await memory.clear_observations_for_memory(
|
||||
bank_id, str(m1), request_context=request_context
|
||||
)
|
||||
|
||||
assert result["deleted_count"] == 1
|
||||
|
||||
@@ -530,7 +439,9 @@ class TestClearObservationsForMemory:
|
||||
|
||||
# Patch out consolidation so it doesn't re-set consolidated_at before we can check it
|
||||
with patch.object(memory, "submit_async_consolidation", new=AsyncMock()):
|
||||
result = await memory.clear_observations_for_memory(bank_id, str(m1), request_context=request_context)
|
||||
result = await memory.clear_observations_for_memory(
|
||||
bank_id, str(m1), request_context=request_context
|
||||
)
|
||||
|
||||
assert result["deleted_count"] == 2
|
||||
|
||||
@@ -582,8 +493,11 @@ async def _insert_document_with_memories(
|
||||
|
||||
|
||||
class TestUpdateDocumentTagsObservationCleanup:
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_tags_returns_updated_document(self, memory: MemoryEngine, request_context: RequestContext):
|
||||
async def test_update_tags_returns_updated_document(
|
||||
self, memory: MemoryEngine, request_context: RequestContext
|
||||
):
|
||||
"""update_document returns the updated document with new tags."""
|
||||
bank_id = f"test-tag-update-basic-{uuid.uuid4().hex[:8]}"
|
||||
await _ensure_bank(memory, bank_id, request_context)
|
||||
@@ -593,7 +507,9 @@ class TestUpdateDocumentTagsObservationCleanup:
|
||||
doc_id = f"doc-{uuid.uuid4().hex[:8]}"
|
||||
await _insert_document_with_memories(conn, bank_id, doc_id, [("Alice loves hiking.", "experience")])
|
||||
|
||||
result = await memory.update_document(doc_id, bank_id, tags=["new-tag"], request_context=request_context)
|
||||
result = await memory.update_document(
|
||||
doc_id, bank_id, tags=["new-tag"], request_context=request_context
|
||||
)
|
||||
|
||||
assert result is True
|
||||
|
||||
@@ -607,14 +523,18 @@ class TestUpdateDocumentTagsObservationCleanup:
|
||||
bank_id = f"test-tag-update-missing-{uuid.uuid4().hex[:8]}"
|
||||
await _ensure_bank(memory, bank_id, request_context)
|
||||
|
||||
result = await memory.update_document("nonexistent-doc", bank_id, tags=["tag"], request_context=request_context)
|
||||
result = await memory.update_document(
|
||||
"nonexistent-doc", bank_id, tags=["tag"], request_context=request_context
|
||||
)
|
||||
|
||||
assert result is False
|
||||
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_tags_propagates_to_memory_units(self, memory: MemoryEngine, request_context: RequestContext):
|
||||
async def test_update_tags_propagates_to_memory_units(
|
||||
self, memory: MemoryEngine, request_context: RequestContext
|
||||
):
|
||||
"""Changing document tags also updates all associated memory unit tags."""
|
||||
bank_id = f"test-tag-update-propagate-{uuid.uuid4().hex[:8]}"
|
||||
await _ensure_bank(memory, bank_id, request_context)
|
||||
@@ -627,17 +547,23 @@ class TestUpdateDocumentTagsObservationCleanup:
|
||||
)
|
||||
|
||||
with patch.object(memory, "submit_async_consolidation", new=AsyncMock()):
|
||||
await memory.update_document(doc_id, bank_id, tags=["new-tag"], request_context=request_context)
|
||||
await memory.update_document(
|
||||
doc_id, bank_id, tags=["new-tag"], request_context=request_context
|
||||
)
|
||||
|
||||
async with pool.acquire() as conn:
|
||||
for mem_id in mem_ids:
|
||||
tags = await conn.fetchval("SELECT tags FROM memory_units WHERE id = $1", mem_id)
|
||||
tags = await conn.fetchval(
|
||||
"SELECT tags FROM memory_units WHERE id = $1", mem_id
|
||||
)
|
||||
assert list(tags) == ["new-tag"], f"Memory unit {mem_id} should have updated tags"
|
||||
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_tags_invalidates_observations(self, memory: MemoryEngine, request_context: RequestContext):
|
||||
async def test_update_tags_invalidates_observations(
|
||||
self, memory: MemoryEngine, request_context: RequestContext
|
||||
):
|
||||
"""Observations referencing the document's memory units are deleted on tag change."""
|
||||
bank_id = f"test-tag-update-obs-{uuid.uuid4().hex[:8]}"
|
||||
await _ensure_bank(memory, bank_id, request_context)
|
||||
@@ -651,7 +577,9 @@ class TestUpdateDocumentTagsObservationCleanup:
|
||||
obs_id = await _insert_observation(conn, bank_id, "Alice is a hiker.", mem_ids)
|
||||
|
||||
with patch.object(memory, "submit_async_consolidation", new=AsyncMock()):
|
||||
await memory.update_document(doc_id, bank_id, tags=["new-tag"], request_context=request_context)
|
||||
await memory.update_document(
|
||||
doc_id, bank_id, tags=["new-tag"], request_context=request_context
|
||||
)
|
||||
|
||||
async with pool.acquire() as conn:
|
||||
obs_ids = await _get_observation_ids(conn, bank_id)
|
||||
@@ -679,7 +607,9 @@ class TestUpdateDocumentTagsObservationCleanup:
|
||||
assert await _get_consolidated_at(conn, mem_ids[0]) is not None
|
||||
|
||||
with patch.object(memory, "submit_async_consolidation", new=AsyncMock()):
|
||||
await memory.update_document(doc_id, bank_id, tags=["new-tag"], request_context=request_context)
|
||||
await memory.update_document(
|
||||
doc_id, bank_id, tags=["new-tag"], request_context=request_context
|
||||
)
|
||||
|
||||
async with pool.acquire() as conn:
|
||||
consolidated_at = await _get_consolidated_at(conn, mem_ids[0])
|
||||
@@ -704,7 +634,9 @@ class TestUpdateDocumentTagsObservationCleanup:
|
||||
await _insert_observation(conn, bank_id, "Alice is a hiker.", mem_ids)
|
||||
|
||||
with patch.object(memory, "submit_async_consolidation", new=AsyncMock()) as mock_consolidate:
|
||||
await memory.update_document(doc_id, bank_id, tags=["new-tag"], request_context=request_context)
|
||||
await memory.update_document(
|
||||
doc_id, bank_id, tags=["new-tag"], request_context=request_context
|
||||
)
|
||||
mock_consolidate.assert_awaited_once()
|
||||
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
@@ -720,11 +652,15 @@ class TestUpdateDocumentTagsObservationCleanup:
|
||||
pool = await memory._get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
doc_id = f"doc-{uuid.uuid4().hex[:8]}"
|
||||
await _insert_document_with_memories(conn, bank_id, doc_id, [("Alice loves hiking.", "experience")])
|
||||
await _insert_document_with_memories(
|
||||
conn, bank_id, doc_id, [("Alice loves hiking.", "experience")]
|
||||
)
|
||||
# No observations inserted
|
||||
|
||||
with patch.object(memory, "submit_async_consolidation", new=AsyncMock()) as mock_consolidate:
|
||||
await memory.update_document(doc_id, bank_id, tags=["new-tag"], request_context=request_context)
|
||||
await memory.update_document(
|
||||
doc_id, bank_id, tags=["new-tag"], request_context=request_context
|
||||
)
|
||||
mock_consolidate.assert_not_awaited()
|
||||
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
@@ -753,7 +689,9 @@ class TestUpdateDocumentTagsObservationCleanup:
|
||||
assert await _get_consolidated_at(conn, other_mem) is not None
|
||||
|
||||
with patch.object(memory, "submit_async_consolidation", new=AsyncMock()):
|
||||
await memory.update_document(doc_id, bank_id, tags=["new-tag"], request_context=request_context)
|
||||
await memory.update_document(
|
||||
doc_id, bank_id, tags=["new-tag"], request_context=request_context
|
||||
)
|
||||
|
||||
async with pool.acquire() as conn:
|
||||
obs_ids = await _get_observation_ids(conn, bank_id)
|
||||
@@ -781,128 +719,17 @@ class TestUpdateDocumentTagsObservationCleanup:
|
||||
)
|
||||
# Unrelated memory not in the document
|
||||
unrelated = await _insert_memory(conn, bank_id, "Bob likes cycling.")
|
||||
unrelated_obs_id = await _insert_observation(conn, bank_id, "Bob is a cyclist.", [unrelated])
|
||||
unrelated_obs_id = await _insert_observation(
|
||||
conn, bank_id, "Bob is a cyclist.", [unrelated]
|
||||
)
|
||||
|
||||
with patch.object(memory, "submit_async_consolidation", new=AsyncMock()):
|
||||
await memory.update_document(doc_id, bank_id, tags=["new-tag"], request_context=request_context)
|
||||
await memory.update_document(
|
||||
doc_id, bank_id, tags=["new-tag"], request_context=request_context
|
||||
)
|
||||
|
||||
async with pool.acquire() as conn:
|
||||
obs_ids = await _get_observation_ids(conn, bank_id)
|
||||
assert str(unrelated_obs_id) in obs_ids, "Unrelated observation should remain untouched"
|
||||
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests: consolidation-vs-delete race — filtering stale source_memory_ids
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestConsolidationSourceMemoryFiltering:
|
||||
"""
|
||||
When a source memory is deleted concurrently with consolidation, the
|
||||
observation must not be written referencing the dead uuid. We exercise
|
||||
the guard by calling the consolidator helpers directly with a deleted
|
||||
source id in the input list.
|
||||
"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_observation_filters_deleted_source_memories(
|
||||
self, memory: MemoryEngine, request_context: RequestContext
|
||||
):
|
||||
from hindsight_api.engine.consolidation.consolidator import _create_observation_directly
|
||||
|
||||
bank_id = f"test-race-create-filter-{uuid.uuid4().hex[:8]}"
|
||||
await _ensure_bank(memory, bank_id, request_context)
|
||||
|
||||
pool = await memory._get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
live = await _insert_memory(conn, bank_id, "Alice loves hiking.")
|
||||
dead = uuid.uuid4() # never existed — stands in for a concurrently deleted source
|
||||
|
||||
result = await _create_observation_directly(
|
||||
conn=conn,
|
||||
memory_engine=memory,
|
||||
bank_id=bank_id,
|
||||
source_memory_ids=[live, dead],
|
||||
observation_text="Alice enjoys hiking regularly.",
|
||||
)
|
||||
|
||||
assert result["action"] == "created"
|
||||
stored = await conn.fetchval(
|
||||
"SELECT source_memory_ids FROM memory_units WHERE id = $1",
|
||||
uuid.UUID(result["observation_id"]),
|
||||
)
|
||||
stored_set = {str(s) for s in stored}
|
||||
assert str(live) in stored_set
|
||||
assert str(dead) not in stored_set, "Deleted source must not appear in stored observation"
|
||||
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_observation_skipped_when_all_sources_deleted(
|
||||
self, memory: MemoryEngine, request_context: RequestContext
|
||||
):
|
||||
from hindsight_api.engine.consolidation.consolidator import _create_observation_directly
|
||||
|
||||
bank_id = f"test-race-create-skip-{uuid.uuid4().hex[:8]}"
|
||||
await _ensure_bank(memory, bank_id, request_context)
|
||||
|
||||
pool = await memory._get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
result = await _create_observation_directly(
|
||||
conn=conn,
|
||||
memory_engine=memory,
|
||||
bank_id=bank_id,
|
||||
source_memory_ids=[uuid.uuid4(), uuid.uuid4()],
|
||||
observation_text="All sources gone.",
|
||||
)
|
||||
|
||||
assert result["action"] == "skipped"
|
||||
assert result["reason"] == "sources_deleted"
|
||||
|
||||
obs_ids = await _get_observation_ids(conn, bank_id)
|
||||
assert obs_ids == [], "No observation row should exist"
|
||||
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_observation_skipped_when_all_new_sources_deleted(
|
||||
self, memory: MemoryEngine, request_context: RequestContext
|
||||
):
|
||||
from hindsight_api.engine.consolidation.consolidator import _execute_update_action
|
||||
from hindsight_api.engine.response_models import MemoryFact
|
||||
|
||||
bank_id = f"test-race-update-skip-{uuid.uuid4().hex[:8]}"
|
||||
await _ensure_bank(memory, bank_id, request_context)
|
||||
|
||||
pool = await memory._get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
original_source = await _insert_memory(conn, bank_id, "Alice hikes.")
|
||||
obs_id = await _insert_observation(conn, bank_id, "Alice is a hiker.", [original_source])
|
||||
original_text = "Alice is a hiker."
|
||||
|
||||
observation_model = MemoryFact(
|
||||
id=str(obs_id),
|
||||
text=original_text,
|
||||
fact_type="observation",
|
||||
source_fact_ids=[str(original_source)],
|
||||
tags=[],
|
||||
)
|
||||
|
||||
await _execute_update_action(
|
||||
conn=conn,
|
||||
memory_engine=memory,
|
||||
bank_id=bank_id,
|
||||
source_memory_ids=[uuid.uuid4(), uuid.uuid4()], # all dead
|
||||
observation_id=str(obs_id),
|
||||
new_text="This update must not land.",
|
||||
observations=[observation_model],
|
||||
)
|
||||
|
||||
row = await conn.fetchrow("SELECT text, source_memory_ids FROM memory_units WHERE id = $1", obs_id)
|
||||
assert row["text"] == original_text, "Observation text must not change"
|
||||
stored_sources = {str(s) for s in row["source_memory_ids"]}
|
||||
assert stored_sources == {str(original_source)}, "Dead sources must not be appended"
|
||||
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@@ -1,245 +0,0 @@
|
||||
"""
|
||||
Tests that async operation statuses (pending, processing, completed, failed, cancelled)
|
||||
are correctly exposed through list and get API endpoints.
|
||||
|
||||
Regression tests:
|
||||
- Previously the API collapsed 'processing' into 'pending', hiding the real status.
|
||||
- Cancel used to delete the operation row; now it sets status to 'cancelled'.
|
||||
- Retry now accepts both 'failed' and 'cancelled' operations.
|
||||
"""
|
||||
import uuid
|
||||
from datetime import datetime
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
|
||||
from hindsight_api.api import create_app
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def api_client(memory):
|
||||
app = create_app(memory, initialize_memory=False)
|
||||
transport = httpx.ASGITransport(app=app)
|
||||
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
|
||||
yield client
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def test_bank_id():
|
||||
return f"op_status_test_{datetime.now().timestamp()}"
|
||||
|
||||
|
||||
async def _ensure_bank(pool, bank_id: str) -> None:
|
||||
"""Create a bank row if it doesn't already exist."""
|
||||
await pool.execute(
|
||||
"""
|
||||
INSERT INTO banks (bank_id) VALUES ($1)
|
||||
ON CONFLICT (bank_id) DO NOTHING
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
|
||||
|
||||
async def _insert_operation(pool, bank_id: str, status: str) -> str:
|
||||
"""Insert a test operation with the given status and return its ID."""
|
||||
op_id = uuid.uuid4()
|
||||
await pool.execute(
|
||||
"""
|
||||
INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload)
|
||||
VALUES ($1, $2, 'retain', $3, '{"test": true}'::jsonb)
|
||||
""",
|
||||
op_id,
|
||||
bank_id,
|
||||
status,
|
||||
)
|
||||
return str(op_id)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_operations_returns_processing_status(api_client, memory, test_bank_id):
|
||||
"""GET /operations should return 'processing' status, not collapse it to 'pending'."""
|
||||
pool = memory._pool
|
||||
await _ensure_bank(pool, test_bank_id)
|
||||
|
||||
pending_id = await _insert_operation(pool, test_bank_id, "pending")
|
||||
processing_id = await _insert_operation(pool, test_bank_id, "processing")
|
||||
|
||||
response = await api_client.get(f"/v1/default/banks/{test_bank_id}/operations")
|
||||
assert response.status_code == 200
|
||||
ops = response.json()["operations"]
|
||||
|
||||
statuses_by_id = {op["id"]: op["status"] for op in ops}
|
||||
assert statuses_by_id[pending_id] == "pending"
|
||||
assert statuses_by_id[processing_id] == "processing"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_operations_filter_by_processing(api_client, memory, test_bank_id):
|
||||
"""Filtering by status=processing should only return processing operations."""
|
||||
pool = memory._pool
|
||||
await _ensure_bank(pool, test_bank_id)
|
||||
|
||||
await _insert_operation(pool, test_bank_id, "pending")
|
||||
processing_id = await _insert_operation(pool, test_bank_id, "processing")
|
||||
|
||||
response = await api_client.get(
|
||||
f"/v1/default/banks/{test_bank_id}/operations",
|
||||
params={"status": "processing"},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
ops = response.json()["operations"]
|
||||
|
||||
assert len(ops) == 1
|
||||
assert ops[0]["id"] == processing_id
|
||||
assert ops[0]["status"] == "processing"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_operations_filter_by_pending_excludes_processing(api_client, memory, test_bank_id):
|
||||
"""Filtering by status=pending should NOT include processing operations."""
|
||||
pool = memory._pool
|
||||
await _ensure_bank(pool, test_bank_id)
|
||||
|
||||
pending_id = await _insert_operation(pool, test_bank_id, "pending")
|
||||
await _insert_operation(pool, test_bank_id, "processing")
|
||||
|
||||
response = await api_client.get(
|
||||
f"/v1/default/banks/{test_bank_id}/operations",
|
||||
params={"status": "pending"},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
ops = response.json()["operations"]
|
||||
|
||||
assert len(ops) == 1
|
||||
assert ops[0]["id"] == pending_id
|
||||
assert ops[0]["status"] == "pending"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_operation_returns_processing_status(api_client, memory, test_bank_id):
|
||||
"""GET /operations/{id} should return 'processing' status."""
|
||||
pool = memory._pool
|
||||
await _ensure_bank(pool, test_bank_id)
|
||||
|
||||
processing_id = await _insert_operation(pool, test_bank_id, "processing")
|
||||
|
||||
response = await api_client.get(
|
||||
f"/v1/default/banks/{test_bank_id}/operations/{processing_id}"
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["status"] == "processing"
|
||||
assert data["operation_id"] == processing_id
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_all_statuses_returned_correctly(api_client, memory, test_bank_id):
|
||||
"""All four DB statuses should be returned as-is through both list and get endpoints."""
|
||||
pool = memory._pool
|
||||
await _ensure_bank(pool, test_bank_id)
|
||||
|
||||
ids = {}
|
||||
for status in ("pending", "processing", "completed", "failed", "cancelled"):
|
||||
ids[status] = await _insert_operation(pool, test_bank_id, status)
|
||||
|
||||
# Verify list endpoint
|
||||
response = await api_client.get(f"/v1/default/banks/{test_bank_id}/operations")
|
||||
assert response.status_code == 200
|
||||
ops = response.json()["operations"]
|
||||
statuses_by_id = {op["id"]: op["status"] for op in ops}
|
||||
|
||||
for status, op_id in ids.items():
|
||||
assert statuses_by_id[op_id] == status, f"List: expected {status} for {op_id}, got {statuses_by_id[op_id]}"
|
||||
|
||||
# Verify get endpoint for each
|
||||
for status, op_id in ids.items():
|
||||
response = await api_client.get(
|
||||
f"/v1/default/banks/{test_bank_id}/operations/{op_id}"
|
||||
)
|
||||
assert response.status_code == 200
|
||||
assert response.json()["status"] == status, f"Get: expected {status} for {op_id}"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancel_sets_cancelled_status(api_client, memory, test_bank_id):
|
||||
"""DELETE /operations/{id} should set status to 'cancelled', not delete the row."""
|
||||
pool = memory._pool
|
||||
await _ensure_bank(pool, test_bank_id)
|
||||
|
||||
op_id = await _insert_operation(pool, test_bank_id, "pending")
|
||||
|
||||
# Cancel the operation
|
||||
response = await api_client.delete(
|
||||
f"/v1/default/banks/{test_bank_id}/operations/{op_id}"
|
||||
)
|
||||
assert response.status_code == 200
|
||||
assert response.json()["success"] is True
|
||||
|
||||
# Verify the operation still exists with 'cancelled' status
|
||||
response = await api_client.get(
|
||||
f"/v1/default/banks/{test_bank_id}/operations/{op_id}"
|
||||
)
|
||||
assert response.status_code == 200
|
||||
assert response.json()["status"] == "cancelled"
|
||||
|
||||
# Verify it shows up in list with cancelled filter
|
||||
response = await api_client.get(
|
||||
f"/v1/default/banks/{test_bank_id}/operations",
|
||||
params={"status": "cancelled"},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
ops = response.json()["operations"]
|
||||
assert len(ops) == 1
|
||||
assert ops[0]["id"] == op_id
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retry_cancelled_operation(api_client, memory, test_bank_id):
|
||||
"""POST /operations/{id}/retry should accept cancelled operations."""
|
||||
pool = memory._pool
|
||||
await _ensure_bank(pool, test_bank_id)
|
||||
|
||||
op_id = await _insert_operation(pool, test_bank_id, "cancelled")
|
||||
|
||||
# Retry the cancelled operation
|
||||
response = await api_client.post(
|
||||
f"/v1/default/banks/{test_bank_id}/operations/{op_id}/retry"
|
||||
)
|
||||
assert response.status_code == 200
|
||||
assert response.json()["success"] is True
|
||||
|
||||
# Verify the operation is now pending
|
||||
response = await api_client.get(
|
||||
f"/v1/default/banks/{test_bank_id}/operations/{op_id}"
|
||||
)
|
||||
assert response.status_code == 200
|
||||
assert response.json()["status"] == "pending"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retry_rejects_non_retriable_statuses(api_client, memory, test_bank_id):
|
||||
"""POST /operations/{id}/retry should reject pending, processing, and completed operations."""
|
||||
pool = memory._pool
|
||||
await _ensure_bank(pool, test_bank_id)
|
||||
|
||||
for status in ("pending", "processing", "completed"):
|
||||
op_id = await _insert_operation(pool, test_bank_id, status)
|
||||
response = await api_client.post(
|
||||
f"/v1/default/banks/{test_bank_id}/operations/{op_id}/retry"
|
||||
)
|
||||
assert response.status_code == 409, f"Expected 409 for {status}, got {response.status_code}"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancel_rejects_non_pending_operations(api_client, memory, test_bank_id):
|
||||
"""DELETE /operations/{id} should only cancel pending operations."""
|
||||
pool = memory._pool
|
||||
await _ensure_bank(pool, test_bank_id)
|
||||
|
||||
for status in ("processing", "completed", "failed"):
|
||||
op_id = await _insert_operation(pool, test_bank_id, status)
|
||||
response = await api_client.delete(
|
||||
f"/v1/default/banks/{test_bank_id}/operations/{op_id}"
|
||||
)
|
||||
assert response.status_code == 409, f"Expected 409 for {status}, got {response.status_code}"
|
||||
@@ -276,80 +276,30 @@ class TestReflectUsesReflectLLMConfig:
|
||||
# Verify it's different from the retain config
|
||||
assert engine._reflect_llm_config.model != engine._retain_llm_config.model
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reflect_allowed_when_default_llm_none_but_reflect_configured(self, monkeypatch):
|
||||
"""A disabled default LLM should not block a separately configured reflect LLM."""
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
from hindsight_api import MemoryEngine
|
||||
from hindsight_api.engine.reflect.models import ReflectAgentResult
|
||||
from hindsight_api.models import RequestContext
|
||||
|
||||
engine = MemoryEngine(
|
||||
memory_llm_provider="none",
|
||||
memory_llm_model="none",
|
||||
reflect_llm_provider="mock",
|
||||
reflect_llm_model="reflect-specific-model",
|
||||
skip_llm_verification=True,
|
||||
lazy_reranker=True,
|
||||
)
|
||||
|
||||
engine._authenticate_tenant = AsyncMock() # type: ignore[method-assign]
|
||||
engine.get_bank_profile = AsyncMock(return_value={"name": "Test", "mission": ""}) # type: ignore[method-assign]
|
||||
engine.get_bank_stats = AsyncMock(
|
||||
return_value=SimpleNamespace(last_consolidated_at=None, pending_consolidation=0)
|
||||
) # type: ignore[method-assign]
|
||||
engine.list_directives = AsyncMock(return_value=[]) # type: ignore[method-assign]
|
||||
engine._get_pool = AsyncMock(return_value=SimpleNamespace()) # type: ignore[method-assign]
|
||||
engine._config_resolver = SimpleNamespace(
|
||||
resolve_full_config=AsyncMock(return_value=SimpleNamespace(llm_gemini_safety_settings=None)),
|
||||
get_bank_config=AsyncMock(return_value={}),
|
||||
)
|
||||
|
||||
async def fake_run_reflect_agent(**kwargs):
|
||||
assert kwargs["llm_config"].provider == "mock"
|
||||
return ReflectAgentResult(text="reflect works")
|
||||
|
||||
monkeypatch.setattr("hindsight_api.engine.memory_engine.run_reflect_agent", fake_run_reflect_agent)
|
||||
|
||||
result = await engine.reflect_async(
|
||||
bank_id="bank-1",
|
||||
query="test",
|
||||
request_context=RequestContext(),
|
||||
exclude_mental_models=True,
|
||||
fact_types=["observation"],
|
||||
)
|
||||
|
||||
assert result.text == "reflect works"
|
||||
|
||||
|
||||
class TestRetryAndBackoffConfiguration:
|
||||
"""Test retry and backoff configuration options."""
|
||||
|
||||
def test_global_retry_backoff_config_defaults(self):
|
||||
"""Test that global retry/backoff settings have correct defaults."""
|
||||
from hindsight_api.config import DEFAULT_LLM_MAX_RETRIES, get_config
|
||||
from hindsight_api.config import get_config
|
||||
|
||||
config = get_config()
|
||||
|
||||
# Verify global defaults
|
||||
assert config.llm_max_retries == DEFAULT_LLM_MAX_RETRIES
|
||||
assert config.llm_max_retries == 10
|
||||
assert config.llm_initial_backoff == 1.0
|
||||
assert config.llm_max_backoff == 60.0
|
||||
|
||||
def test_per_operation_retry_backoff_config_from_env(self):
|
||||
"""Test that per-operation retry/backoff settings are loaded from environment."""
|
||||
from hindsight_api.config import DEFAULT_LLM_MAX_RETRIES, clear_config_cache
|
||||
from hindsight_api.config import clear_config_cache
|
||||
|
||||
# Set per-operation overrides (choose values different from the global default so the
|
||||
# "global unchanged" assertions below are meaningful).
|
||||
retain_retries = DEFAULT_LLM_MAX_RETRIES + 1
|
||||
reflect_retries = DEFAULT_LLM_MAX_RETRIES + 2
|
||||
os.environ["HINDSIGHT_API_RETAIN_LLM_MAX_RETRIES"] = str(retain_retries)
|
||||
# Set per-operation overrides
|
||||
os.environ["HINDSIGHT_API_RETAIN_LLM_MAX_RETRIES"] = "3"
|
||||
os.environ["HINDSIGHT_API_RETAIN_LLM_INITIAL_BACKOFF"] = "2.0"
|
||||
os.environ["HINDSIGHT_API_RETAIN_LLM_MAX_BACKOFF"] = "120.0"
|
||||
os.environ["HINDSIGHT_API_REFLECT_LLM_MAX_RETRIES"] = str(reflect_retries)
|
||||
os.environ["HINDSIGHT_API_REFLECT_LLM_MAX_RETRIES"] = "5"
|
||||
os.environ["HINDSIGHT_API_REFLECT_LLM_INITIAL_BACKOFF"] = "1.5"
|
||||
os.environ["HINDSIGHT_API_REFLECT_LLM_MAX_BACKOFF"] = "90.0"
|
||||
|
||||
@@ -360,17 +310,17 @@ class TestRetryAndBackoffConfiguration:
|
||||
config = get_config()
|
||||
|
||||
# Verify retain overrides
|
||||
assert config.retain_llm_max_retries == retain_retries
|
||||
assert config.retain_llm_max_retries == 3
|
||||
assert config.retain_llm_initial_backoff == 2.0
|
||||
assert config.retain_llm_max_backoff == 120.0
|
||||
|
||||
# Verify reflect overrides
|
||||
assert config.reflect_llm_max_retries == reflect_retries
|
||||
assert config.reflect_llm_max_retries == 5
|
||||
assert config.reflect_llm_initial_backoff == 1.5
|
||||
assert config.reflect_llm_max_backoff == 90.0
|
||||
|
||||
# Verify global defaults remain unchanged
|
||||
assert config.llm_max_retries == DEFAULT_LLM_MAX_RETRIES
|
||||
assert config.llm_max_retries == 10
|
||||
assert config.llm_initial_backoff == 1.0
|
||||
assert config.llm_max_backoff == 60.0
|
||||
finally:
|
||||
|
||||
@@ -1,254 +0,0 @@
|
||||
"""
|
||||
Tests for the configurable recall-budget mapping (Budget enum -> thinking_budget int).
|
||||
|
||||
Two functions are supported:
|
||||
- "fixed": returns the recall_budget_fixed_<level> integer directly (legacy default).
|
||||
- "adaptive": returns round(max_tokens * recall_budget_adaptive_<level>),
|
||||
clamped to [recall_budget_min, recall_budget_max].
|
||||
|
||||
Both the function selector and the per-level numbers are hierarchical config
|
||||
fields (global env -> tenant -> bank), so they can be overridden per bank.
|
||||
"""
|
||||
|
||||
import dataclasses
|
||||
|
||||
import pytest
|
||||
|
||||
from hindsight_api.config import (
|
||||
DEFAULT_RECALL_BUDGET_ADAPTIVE_HIGH,
|
||||
DEFAULT_RECALL_BUDGET_ADAPTIVE_LOW,
|
||||
DEFAULT_RECALL_BUDGET_ADAPTIVE_MID,
|
||||
DEFAULT_RECALL_BUDGET_FIXED_HIGH,
|
||||
DEFAULT_RECALL_BUDGET_FIXED_LOW,
|
||||
DEFAULT_RECALL_BUDGET_FIXED_MID,
|
||||
DEFAULT_RECALL_BUDGET_MAX,
|
||||
DEFAULT_RECALL_BUDGET_MIN,
|
||||
DEFAULT_RECALL_BUDGET_FUNCTION,
|
||||
ENV_RECALL_BUDGET_ADAPTIVE_LOW,
|
||||
ENV_RECALL_BUDGET_ADAPTIVE_MID,
|
||||
ENV_RECALL_BUDGET_FIXED_HIGH,
|
||||
ENV_RECALL_BUDGET_FIXED_LOW,
|
||||
ENV_RECALL_BUDGET_FIXED_MID,
|
||||
ENV_RECALL_BUDGET_MAX,
|
||||
ENV_RECALL_BUDGET_MIN,
|
||||
ENV_RECALL_BUDGET_FUNCTION,
|
||||
RECALL_BUDGET_FUNCTIONS,
|
||||
HindsightConfig,
|
||||
)
|
||||
from hindsight_api.config_resolver import _validate_recall_budget_updates
|
||||
from hindsight_api.engine.memory_engine import Budget, _resolve_thinking_budget
|
||||
|
||||
|
||||
_BUDGET_FIELD_NAMES = (
|
||||
"recall_budget_function",
|
||||
"recall_budget_fixed_low",
|
||||
"recall_budget_fixed_mid",
|
||||
"recall_budget_fixed_high",
|
||||
"recall_budget_adaptive_low",
|
||||
"recall_budget_adaptive_mid",
|
||||
"recall_budget_adaptive_high",
|
||||
"recall_budget_min",
|
||||
"recall_budget_max",
|
||||
)
|
||||
|
||||
|
||||
class TestBudgetConfigFields:
|
||||
def test_fields_exist_on_dataclass(self):
|
||||
names = {f.name for f in dataclasses.fields(HindsightConfig)}
|
||||
for field_name in _BUDGET_FIELD_NAMES:
|
||||
assert field_name in names, f"Missing dataclass field: {field_name}"
|
||||
|
||||
def test_fields_are_configurable(self):
|
||||
configurable = HindsightConfig.get_configurable_fields()
|
||||
for field_name in _BUDGET_FIELD_NAMES:
|
||||
assert field_name in configurable, f"Field not in _CONFIGURABLE_FIELDS: {field_name}"
|
||||
|
||||
def test_default_function_is_fixed_for_backwards_compat(self):
|
||||
# The whole point of function="fixed" being default is to preserve legacy behavior.
|
||||
assert DEFAULT_RECALL_BUDGET_FUNCTION == "fixed"
|
||||
assert "fixed" in RECALL_BUDGET_FUNCTIONS
|
||||
assert "adaptive" in RECALL_BUDGET_FUNCTIONS
|
||||
|
||||
def test_default_fixed_values_match_legacy_hardcoded_mapping(self):
|
||||
# These are the values that used to live in the hardcoded budget_mapping dict.
|
||||
assert DEFAULT_RECALL_BUDGET_FIXED_LOW == 100
|
||||
assert DEFAULT_RECALL_BUDGET_FIXED_MID == 300
|
||||
assert DEFAULT_RECALL_BUDGET_FIXED_HIGH == 1000
|
||||
|
||||
def test_default_adaptive_clamps_are_sane(self):
|
||||
assert DEFAULT_RECALL_BUDGET_MIN >= 1
|
||||
assert DEFAULT_RECALL_BUDGET_MAX > DEFAULT_RECALL_BUDGET_MIN
|
||||
|
||||
def test_env_var_constants(self):
|
||||
assert ENV_RECALL_BUDGET_FUNCTION == "HINDSIGHT_API_RECALL_BUDGET_FUNCTION"
|
||||
assert ENV_RECALL_BUDGET_FIXED_LOW == "HINDSIGHT_API_RECALL_BUDGET_FIXED_LOW"
|
||||
assert ENV_RECALL_BUDGET_FIXED_MID == "HINDSIGHT_API_RECALL_BUDGET_FIXED_MID"
|
||||
assert ENV_RECALL_BUDGET_FIXED_HIGH == "HINDSIGHT_API_RECALL_BUDGET_FIXED_HIGH"
|
||||
assert ENV_RECALL_BUDGET_ADAPTIVE_LOW == "HINDSIGHT_API_RECALL_BUDGET_ADAPTIVE_LOW"
|
||||
assert ENV_RECALL_BUDGET_ADAPTIVE_MID == "HINDSIGHT_API_RECALL_BUDGET_ADAPTIVE_MID"
|
||||
assert ENV_RECALL_BUDGET_MIN == "HINDSIGHT_API_RECALL_BUDGET_MIN"
|
||||
assert ENV_RECALL_BUDGET_MAX == "HINDSIGHT_API_RECALL_BUDGET_MAX"
|
||||
|
||||
def test_from_env_reads_overrides(self, monkeypatch):
|
||||
monkeypatch.setenv(ENV_RECALL_BUDGET_FUNCTION, "adaptive")
|
||||
monkeypatch.setenv(ENV_RECALL_BUDGET_FIXED_MID, "777")
|
||||
monkeypatch.setenv(ENV_RECALL_BUDGET_ADAPTIVE_MID, "0.5")
|
||||
monkeypatch.setenv(ENV_RECALL_BUDGET_MIN, "5")
|
||||
monkeypatch.setenv(ENV_RECALL_BUDGET_MAX, "9999")
|
||||
|
||||
config = HindsightConfig.from_env()
|
||||
assert config.recall_budget_function == "adaptive"
|
||||
assert config.recall_budget_fixed_mid == 777
|
||||
assert config.recall_budget_adaptive_mid == 0.5
|
||||
assert config.recall_budget_min == 5
|
||||
assert config.recall_budget_max == 9999
|
||||
|
||||
def test_from_env_invalid_function_falls_back_to_default(self, monkeypatch):
|
||||
# Defensive parsing: an invalid env value logs a warning and falls back.
|
||||
monkeypatch.setenv(ENV_RECALL_BUDGET_FUNCTION, "garbage")
|
||||
config = HindsightConfig.from_env()
|
||||
assert config.recall_budget_function == DEFAULT_RECALL_BUDGET_FUNCTION
|
||||
|
||||
|
||||
class TestResolveThinkingBudgetFixedFunction:
|
||||
@pytest.fixture
|
||||
def fixed_config(self):
|
||||
return {
|
||||
"recall_budget_function": "fixed",
|
||||
"recall_budget_fixed_low": 100,
|
||||
"recall_budget_fixed_mid": 300,
|
||||
"recall_budget_fixed_high": 1000,
|
||||
"recall_budget_adaptive_low": 0.025,
|
||||
"recall_budget_adaptive_mid": 0.075,
|
||||
"recall_budget_adaptive_high": 0.25,
|
||||
"recall_budget_min": 20,
|
||||
"recall_budget_max": 2000,
|
||||
}
|
||||
|
||||
def test_low_mid_high_match_fixed_values(self, fixed_config):
|
||||
assert _resolve_thinking_budget(fixed_config, Budget.LOW, 4096) == 100
|
||||
assert _resolve_thinking_budget(fixed_config, Budget.MID, 4096) == 300
|
||||
assert _resolve_thinking_budget(fixed_config, Budget.HIGH, 4096) == 1000
|
||||
|
||||
def test_none_budget_defaults_to_mid(self, fixed_config):
|
||||
assert _resolve_thinking_budget(fixed_config, None, 4096) == 300
|
||||
|
||||
def test_max_tokens_does_not_affect_fixed_function(self, fixed_config):
|
||||
# Whole point of "fixed": result is independent of max_tokens.
|
||||
assert _resolve_thinking_budget(fixed_config, Budget.MID, 1) == 300
|
||||
assert _resolve_thinking_budget(fixed_config, Budget.MID, 1_000_000) == 300
|
||||
|
||||
def test_per_bank_overrides_take_effect(self, fixed_config):
|
||||
fixed_config["recall_budget_fixed_mid"] = 42
|
||||
assert _resolve_thinking_budget(fixed_config, Budget.MID, 4096) == 42
|
||||
|
||||
|
||||
class TestResolveThinkingBudgetAdaptiveFunction:
|
||||
@pytest.fixture
|
||||
def adaptive_config(self):
|
||||
return {
|
||||
"recall_budget_function": "adaptive",
|
||||
"recall_budget_fixed_low": 100,
|
||||
"recall_budget_fixed_mid": 300,
|
||||
"recall_budget_fixed_high": 1000,
|
||||
"recall_budget_adaptive_low": 0.025,
|
||||
"recall_budget_adaptive_mid": 0.075,
|
||||
"recall_budget_adaptive_high": 0.25,
|
||||
"recall_budget_min": 20,
|
||||
"recall_budget_max": 2000,
|
||||
}
|
||||
|
||||
def test_scales_with_max_tokens(self, adaptive_config):
|
||||
# 4096 * 0.075 = 307.2 -> 307
|
||||
assert _resolve_thinking_budget(adaptive_config, Budget.MID, 4096) == 307
|
||||
# 8192 * 0.075 = 614.4 -> 614
|
||||
assert _resolve_thinking_budget(adaptive_config, Budget.MID, 8192) == 614
|
||||
|
||||
def test_clamps_to_floor_when_max_tokens_tiny(self, adaptive_config):
|
||||
# 100 * 0.025 = 2.5 -> 2 -> clamped to floor 20
|
||||
assert _resolve_thinking_budget(adaptive_config, Budget.LOW, 100) == 20
|
||||
|
||||
def test_clamps_to_ceiling_when_max_tokens_huge(self, adaptive_config):
|
||||
# 100_000 * 0.25 = 25_000 -> clamped to ceiling 2000
|
||||
assert _resolve_thinking_budget(adaptive_config, Budget.HIGH, 100_000) == 2000
|
||||
|
||||
def test_none_budget_defaults_to_mid(self, adaptive_config):
|
||||
assert _resolve_thinking_budget(adaptive_config, None, 4096) == 307
|
||||
|
||||
def test_custom_clamps_per_bank(self, adaptive_config):
|
||||
adaptive_config["recall_budget_min"] = 500
|
||||
adaptive_config["recall_budget_max"] = 600
|
||||
# 4096 * 0.075 = 307 -> below floor 500
|
||||
assert _resolve_thinking_budget(adaptive_config, Budget.MID, 4096) == 500
|
||||
# 4096 * 0.25 = 1024 -> above ceiling 600
|
||||
assert _resolve_thinking_budget(adaptive_config, Budget.HIGH, 4096) == 600
|
||||
|
||||
|
||||
class TestResolveThinkingBudgetFallbacks:
|
||||
def test_empty_config_uses_legacy_defaults(self):
|
||||
# Resilience: missing keys should not crash; fallback to legacy mapping.
|
||||
assert _resolve_thinking_budget({}, Budget.LOW, 4096) == 100
|
||||
assert _resolve_thinking_budget({}, Budget.MID, 4096) == 300
|
||||
assert _resolve_thinking_budget({}, Budget.HIGH, 4096) == 1000
|
||||
|
||||
def test_unknown_function_falls_back_to_fixed(self):
|
||||
# Defensive: if some bad config slipped past validation, behave like "fixed".
|
||||
assert _resolve_thinking_budget({"recall_budget_function": "garbage"}, Budget.MID, 4096) == 300
|
||||
|
||||
|
||||
class TestValidateRecallBudgetUpdates:
|
||||
def test_no_op_passes(self):
|
||||
_validate_recall_budget_updates({})
|
||||
_validate_recall_budget_updates({"unrelated_field": 123})
|
||||
|
||||
def test_valid_function_values(self):
|
||||
_validate_recall_budget_updates({"recall_budget_function": "fixed"})
|
||||
_validate_recall_budget_updates({"recall_budget_function": "adaptive"})
|
||||
|
||||
def test_invalid_function_raises(self):
|
||||
with pytest.raises(ValueError, match="recall_budget_function"):
|
||||
_validate_recall_budget_updates({"recall_budget_function": "wrong"})
|
||||
with pytest.raises(ValueError, match="recall_budget_function"):
|
||||
_validate_recall_budget_updates({"recall_budget_function": 123})
|
||||
|
||||
def test_fixed_must_be_positive_integer(self):
|
||||
for key in ("recall_budget_fixed_low", "recall_budget_fixed_mid", "recall_budget_fixed_high"):
|
||||
_validate_recall_budget_updates({key: 1})
|
||||
_validate_recall_budget_updates({key: 100_000})
|
||||
with pytest.raises(ValueError, match=key):
|
||||
_validate_recall_budget_updates({key: 0})
|
||||
with pytest.raises(ValueError, match=key):
|
||||
_validate_recall_budget_updates({key: -5})
|
||||
with pytest.raises(ValueError, match=key):
|
||||
_validate_recall_budget_updates({key: 1.5}) # float not allowed for fixed
|
||||
with pytest.raises(ValueError, match=key):
|
||||
_validate_recall_budget_updates({key: True}) # bool sneaks past int check
|
||||
|
||||
def test_adaptive_must_be_positive_number(self):
|
||||
for key in ("recall_budget_adaptive_low", "recall_budget_adaptive_mid", "recall_budget_adaptive_high"):
|
||||
_validate_recall_budget_updates({key: 0.001})
|
||||
_validate_recall_budget_updates({key: 1.0})
|
||||
_validate_recall_budget_updates({key: 5}) # int is acceptable as a number
|
||||
with pytest.raises(ValueError, match=key):
|
||||
_validate_recall_budget_updates({key: 0})
|
||||
with pytest.raises(ValueError, match=key):
|
||||
_validate_recall_budget_updates({key: -0.1})
|
||||
with pytest.raises(ValueError, match=key):
|
||||
_validate_recall_budget_updates({key: True})
|
||||
with pytest.raises(ValueError, match=key):
|
||||
_validate_recall_budget_updates({key: "0.5"})
|
||||
|
||||
def test_min_must_be_le_max_when_both_set(self):
|
||||
_validate_recall_budget_updates({"recall_budget_min": 10, "recall_budget_max": 1000})
|
||||
_validate_recall_budget_updates({"recall_budget_min": 100, "recall_budget_max": 100})
|
||||
with pytest.raises(ValueError, match="recall_budget_min"):
|
||||
_validate_recall_budget_updates({"recall_budget_min": 5000, "recall_budget_max": 100})
|
||||
|
||||
def test_min_max_must_be_positive_integers(self):
|
||||
for key in ("recall_budget_min", "recall_budget_max"):
|
||||
with pytest.raises(ValueError, match=key):
|
||||
_validate_recall_budget_updates({key: 0})
|
||||
with pytest.raises(ValueError, match=key):
|
||||
_validate_recall_budget_updates({key: -1})
|
||||
with pytest.raises(ValueError, match=key):
|
||||
_validate_recall_budget_updates({key: 1.5})
|
||||
@@ -1,231 +0,0 @@
|
||||
"""
|
||||
Tests for the internal recall configuration knobs used during mental model
|
||||
refresh: recall_include_chunks, recall_max_tokens, recall_chunks_max_tokens.
|
||||
|
||||
These are exposed both as hierarchical config fields (env → tenant → bank)
|
||||
and as overrides on a mental model's `trigger` JSONB field.
|
||||
"""
|
||||
|
||||
import dataclasses
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from hindsight_api.engine.reflect.tools import tool_recall
|
||||
from hindsight_api.engine.response_models import RecallResult as RecallResultModel
|
||||
from hindsight_api.models import RequestContext
|
||||
|
||||
|
||||
def _make_mock_engine():
|
||||
engine = MagicMock()
|
||||
engine.recall_async = AsyncMock(return_value=RecallResultModel(results=[], entities={}, chunks={}))
|
||||
return engine
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_request_context():
|
||||
# internal=True bypasses the tenant extension, letting these unit tests
|
||||
# exercise engine methods without standing up auth.
|
||||
return RequestContext(internal=True)
|
||||
|
||||
|
||||
class TestToolRecallIncludeChunks:
|
||||
"""tool_recall must honor the include_chunks parameter (was hardcoded True)."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_default_includes_chunks(self, mock_request_context):
|
||||
engine = _make_mock_engine()
|
||||
|
||||
await tool_recall(engine, "bank-1", "q", mock_request_context)
|
||||
|
||||
kwargs = engine.recall_async.call_args.kwargs
|
||||
assert kwargs["include_chunks"] is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_include_chunks_false_propagates(self, mock_request_context):
|
||||
engine = _make_mock_engine()
|
||||
|
||||
await tool_recall(engine, "bank-1", "q", mock_request_context, include_chunks=False)
|
||||
|
||||
kwargs = engine.recall_async.call_args.kwargs
|
||||
assert kwargs["include_chunks"] is False
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_max_chunk_tokens_propagates(self, mock_request_context):
|
||||
engine = _make_mock_engine()
|
||||
|
||||
await tool_recall(
|
||||
engine, "bank-1", "q", mock_request_context, max_chunk_tokens=2500, max_tokens=512
|
||||
)
|
||||
|
||||
kwargs = engine.recall_async.call_args.kwargs
|
||||
assert kwargs["max_chunk_tokens"] == 2500
|
||||
assert kwargs["max_tokens"] == 512
|
||||
|
||||
|
||||
class TestRecallConfigFields:
|
||||
"""Hierarchical config fields for internal recall."""
|
||||
|
||||
def test_fields_exist_on_dataclass(self):
|
||||
from hindsight_api.config import HindsightConfig
|
||||
|
||||
names = {f.name for f in dataclasses.fields(HindsightConfig)}
|
||||
assert "recall_include_chunks" in names
|
||||
assert "recall_max_tokens" in names
|
||||
assert "recall_chunks_max_tokens" in names
|
||||
|
||||
def test_fields_are_configurable(self):
|
||||
from hindsight_api.config import HindsightConfig
|
||||
|
||||
configurable = HindsightConfig.get_configurable_fields()
|
||||
assert "recall_include_chunks" in configurable
|
||||
assert "recall_max_tokens" in configurable
|
||||
assert "recall_chunks_max_tokens" in configurable
|
||||
|
||||
def test_default_values(self):
|
||||
from hindsight_api.config import (
|
||||
DEFAULT_RECALL_CHUNKS_MAX_TOKENS,
|
||||
DEFAULT_RECALL_INCLUDE_CHUNKS,
|
||||
DEFAULT_RECALL_MAX_TOKENS,
|
||||
)
|
||||
|
||||
assert DEFAULT_RECALL_INCLUDE_CHUNKS is True
|
||||
assert DEFAULT_RECALL_MAX_TOKENS == 2048
|
||||
assert DEFAULT_RECALL_CHUNKS_MAX_TOKENS == 1000
|
||||
|
||||
def test_env_var_constants(self):
|
||||
from hindsight_api.config import (
|
||||
ENV_RECALL_CHUNKS_MAX_TOKENS,
|
||||
ENV_RECALL_INCLUDE_CHUNKS,
|
||||
ENV_RECALL_MAX_TOKENS,
|
||||
)
|
||||
|
||||
assert ENV_RECALL_INCLUDE_CHUNKS == "HINDSIGHT_API_RECALL_INCLUDE_CHUNKS"
|
||||
assert ENV_RECALL_MAX_TOKENS == "HINDSIGHT_API_RECALL_MAX_TOKENS"
|
||||
assert ENV_RECALL_CHUNKS_MAX_TOKENS == "HINDSIGHT_API_RECALL_CHUNKS_MAX_TOKENS"
|
||||
|
||||
@patch.dict(
|
||||
"os.environ",
|
||||
{
|
||||
"HINDSIGHT_API_RECALL_INCLUDE_CHUNKS": "false",
|
||||
"HINDSIGHT_API_RECALL_MAX_TOKENS": "777",
|
||||
"HINDSIGHT_API_RECALL_CHUNKS_MAX_TOKENS": "333",
|
||||
},
|
||||
)
|
||||
def test_from_env_reads_overrides(self):
|
||||
from hindsight_api.config import HindsightConfig
|
||||
|
||||
config = HindsightConfig.from_env()
|
||||
assert config.recall_include_chunks is False
|
||||
assert config.recall_max_tokens == 777
|
||||
assert config.recall_chunks_max_tokens == 333
|
||||
|
||||
|
||||
class TestMentalModelTriggerRecallFields:
|
||||
"""MentalModelTrigger Pydantic model accepts the new override fields."""
|
||||
|
||||
def test_trigger_accepts_new_fields(self):
|
||||
from hindsight_api.api.http import MentalModelTrigger
|
||||
|
||||
trigger = MentalModelTrigger(
|
||||
include_chunks=False,
|
||||
recall_max_tokens=512,
|
||||
recall_chunks_max_tokens=0,
|
||||
)
|
||||
assert trigger.include_chunks is False
|
||||
assert trigger.recall_max_tokens == 512
|
||||
assert trigger.recall_chunks_max_tokens == 0
|
||||
|
||||
def test_trigger_defaults_are_none(self):
|
||||
from hindsight_api.api.http import MentalModelTrigger
|
||||
|
||||
trigger = MentalModelTrigger()
|
||||
assert trigger.include_chunks is None
|
||||
assert trigger.recall_max_tokens is None
|
||||
assert trigger.recall_chunks_max_tokens is None
|
||||
|
||||
|
||||
class TestRefreshTriggerWiring:
|
||||
"""Verify mental-model refresh forwards trigger overrides into reflect_async kwargs."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_trigger_overrides_passed_to_reflect_async(self, mock_request_context):
|
||||
from hindsight_api.engine.memory_engine import MemoryEngine
|
||||
from hindsight_api.engine.response_models import ReflectResult
|
||||
|
||||
engine = MemoryEngine.__new__(MemoryEngine)
|
||||
|
||||
async def fake_get_mental_model(bank_id, mental_model_id, request_context):
|
||||
return {
|
||||
"id": mental_model_id,
|
||||
"source_query": "What do we know?",
|
||||
"tags": [],
|
||||
"trigger": {
|
||||
"include_chunks": False,
|
||||
"recall_max_tokens": 512,
|
||||
"recall_chunks_max_tokens": 0,
|
||||
"fact_types": ["world"],
|
||||
},
|
||||
}
|
||||
|
||||
captured = {}
|
||||
|
||||
async def fake_reflect_async(**kwargs):
|
||||
captured.update(kwargs)
|
||||
return ReflectResult(text="ok", based_on={})
|
||||
|
||||
async def fake_update_mental_model(*args, **kwargs):
|
||||
return None
|
||||
|
||||
engine.get_mental_model = fake_get_mental_model
|
||||
engine.reflect_async = fake_reflect_async
|
||||
engine.update_mental_model = fake_update_mental_model
|
||||
engine._operation_validator = None
|
||||
engine._tenant_extension = None
|
||||
|
||||
await engine.refresh_mental_model(
|
||||
bank_id="bank-1",
|
||||
mental_model_id="mm-1",
|
||||
request_context=mock_request_context,
|
||||
)
|
||||
|
||||
assert captured["recall_include_chunks"] is False
|
||||
assert captured["recall_max_tokens_override"] == 512
|
||||
assert captured["recall_chunks_max_tokens_override"] == 0
|
||||
assert captured["fact_types"] == ["world"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_missing_trigger_fields_pass_none(self, mock_request_context):
|
||||
from hindsight_api.engine.memory_engine import MemoryEngine
|
||||
from hindsight_api.engine.response_models import ReflectResult
|
||||
|
||||
engine = MemoryEngine.__new__(MemoryEngine)
|
||||
|
||||
async def fake_get_mental_model(bank_id, mental_model_id, request_context):
|
||||
return {"id": mental_model_id, "source_query": "q", "tags": [], "trigger": {}}
|
||||
|
||||
captured = {}
|
||||
|
||||
async def fake_reflect_async(**kwargs):
|
||||
captured.update(kwargs)
|
||||
return ReflectResult(text="ok", based_on={})
|
||||
|
||||
async def fake_update_mental_model(*args, **kwargs):
|
||||
return None
|
||||
|
||||
engine.get_mental_model = fake_get_mental_model
|
||||
engine.reflect_async = fake_reflect_async
|
||||
engine.update_mental_model = fake_update_mental_model
|
||||
engine._operation_validator = None
|
||||
engine._tenant_extension = None
|
||||
|
||||
await engine.refresh_mental_model(
|
||||
bank_id="bank-1",
|
||||
mental_model_id="mm-1",
|
||||
request_context=mock_request_context,
|
||||
)
|
||||
|
||||
# When trigger fields are absent, None is forwarded so reflect_async falls back to bank/global config.
|
||||
assert captured["recall_include_chunks"] is None
|
||||
assert captured["recall_max_tokens_override"] is None
|
||||
assert captured["recall_chunks_max_tokens_override"] is None
|
||||
@@ -1,258 +0,0 @@
|
||||
"""Tests for created_after / created_before time-range filtering in recall.
|
||||
|
||||
Inserts memory_units with known timestamps directly via SQL, then verifies
|
||||
that recall_async respects the time bounds — never returning memories
|
||||
outside the requested range.
|
||||
|
||||
No LLM required — uses mock provider.
|
||||
"""
|
||||
|
||||
import uuid
|
||||
from datetime import datetime, timezone
|
||||
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
|
||||
from hindsight_api import MemoryEngine, RequestContext
|
||||
from hindsight_api.engine.retain import embedding_utils
|
||||
|
||||
# Tests in this file insert memory_units with shared hardcoded UUIDs and
|
||||
# memory_units.id is a global PK, so parallel xdist workers running these
|
||||
# tests simultaneously hit pk_memory_units conflicts. Share an xdist group
|
||||
# so the eight tests serialize on the same worker.
|
||||
pytestmark = pytest.mark.xdist_group("recall_time_range")
|
||||
|
||||
# Three points in time, each 1 hour apart
|
||||
T1 = datetime(2026, 1, 1, 10, 0, 0, tzinfo=timezone.utc)
|
||||
T2 = datetime(2026, 1, 1, 11, 0, 0, tzinfo=timezone.utc)
|
||||
T3 = datetime(2026, 1, 1, 12, 0, 0, tzinfo=timezone.utc)
|
||||
|
||||
# Stable UUIDs for the three facts (deterministic for assertion readability)
|
||||
ID_OLD = "00000000-0000-0000-0000-000000000001"
|
||||
ID_MID = "00000000-0000-0000-0000-000000000002"
|
||||
ID_NEW = "00000000-0000-0000-0000-000000000003"
|
||||
|
||||
RC = RequestContext(tenant_id="default")
|
||||
|
||||
|
||||
async def _insert_fact(
|
||||
conn,
|
||||
*,
|
||||
fact_id: str,
|
||||
text: str,
|
||||
bank_id: str,
|
||||
embedding_str: str,
|
||||
created_at: datetime,
|
||||
updated_at: datetime | None = None,
|
||||
fact_type: str = "world",
|
||||
) -> None:
|
||||
"""Insert a memory_unit with a specific created_at/updated_at timestamp."""
|
||||
updated = updated_at or created_at
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO memory_units (id, bank_id, text, fact_type, embedding, created_at, updated_at)
|
||||
VALUES ($1, $2, $3, $4, $5::vector, $6, $7)
|
||||
""",
|
||||
fact_id,
|
||||
bank_id,
|
||||
text,
|
||||
fact_type,
|
||||
embedding_str,
|
||||
created_at,
|
||||
updated,
|
||||
)
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def seeded_memory(memory_no_llm_verify: MemoryEngine):
|
||||
"""Insert three facts at T1, T2, T3 and return the engine."""
|
||||
engine = memory_no_llm_verify
|
||||
bank_id = f"test-time-range-{uuid.uuid4().hex[:8]}"
|
||||
|
||||
await engine.get_bank_profile(bank_id, request_context=RC)
|
||||
|
||||
# Generate real embeddings so semantic retrieval works
|
||||
embeddings = await embedding_utils.generate_embeddings_batch(
|
||||
engine.embeddings,
|
||||
["the cat sat on the mat", "dogs are loyal animals", "birds can fly in the sky"],
|
||||
)
|
||||
|
||||
def _to_str(emb: list[float]) -> str:
|
||||
return "[" + ",".join(str(v) for v in emb) + "]"
|
||||
|
||||
pool = await engine._get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
# Defensive cleanup: clear any rows left behind by an interrupted
|
||||
# previous run of this fixture (test process killed before teardown).
|
||||
# Without this, pk_memory_units rejects the next insert with the same
|
||||
# hardcoded IDs.
|
||||
await conn.execute(
|
||||
"DELETE FROM memory_units WHERE id IN ($1, $2, $3)",
|
||||
ID_OLD,
|
||||
ID_MID,
|
||||
ID_NEW,
|
||||
)
|
||||
await _insert_fact(
|
||||
conn,
|
||||
fact_id=ID_OLD,
|
||||
text="the cat sat on the mat",
|
||||
bank_id=bank_id,
|
||||
embedding_str=_to_str(embeddings[0]),
|
||||
created_at=T1,
|
||||
updated_at=T1,
|
||||
)
|
||||
await _insert_fact(
|
||||
conn,
|
||||
fact_id=ID_MID,
|
||||
text="dogs are loyal animals",
|
||||
bank_id=bank_id,
|
||||
embedding_str=_to_str(embeddings[1]),
|
||||
created_at=T2,
|
||||
updated_at=T2,
|
||||
)
|
||||
await _insert_fact(
|
||||
conn,
|
||||
fact_id=ID_NEW,
|
||||
text="birds can fly in the sky",
|
||||
bank_id=bank_id,
|
||||
embedding_str=_to_str(embeddings[2]),
|
||||
created_at=T3,
|
||||
updated_at=T3,
|
||||
)
|
||||
|
||||
yield engine, bank_id
|
||||
|
||||
await engine.delete_bank(bank_id, request_context=RC)
|
||||
|
||||
|
||||
def _result_ids(result) -> set[str]:
|
||||
return {str(r.id) for r in result.results}
|
||||
|
||||
|
||||
class TestRecallTimeRange:
|
||||
"""Verify created_after / created_before filtering at the recall level."""
|
||||
|
||||
async def test_no_filter_returns_all(self, seeded_memory):
|
||||
engine, bank_id = seeded_memory
|
||||
result = await engine.recall_async(
|
||||
bank_id=bank_id,
|
||||
query="animals and nature",
|
||||
request_context=RC,
|
||||
max_tokens=10000,
|
||||
)
|
||||
ids = _result_ids(result)
|
||||
assert ID_OLD in ids
|
||||
assert ID_MID in ids
|
||||
assert ID_NEW in ids
|
||||
|
||||
async def test_created_after_excludes_old(self, seeded_memory):
|
||||
"""created_after=T1 excludes fact-old (updated_at == T1, not > T1)."""
|
||||
engine, bank_id = seeded_memory
|
||||
result = await engine.recall_async(
|
||||
bank_id=bank_id,
|
||||
query="animals and nature",
|
||||
request_context=RC,
|
||||
max_tokens=10000,
|
||||
created_after=T1,
|
||||
)
|
||||
ids = _result_ids(result)
|
||||
assert ID_OLD not in ids, "fact-old (updated_at=T1) must be excluded by created_after=T1"
|
||||
assert ID_MID in ids
|
||||
assert ID_NEW in ids
|
||||
|
||||
async def test_created_after_excludes_old_and_mid(self, seeded_memory):
|
||||
"""created_after=T2 returns only fact-new."""
|
||||
engine, bank_id = seeded_memory
|
||||
result = await engine.recall_async(
|
||||
bank_id=bank_id,
|
||||
query="animals and nature",
|
||||
request_context=RC,
|
||||
max_tokens=10000,
|
||||
created_after=T2,
|
||||
)
|
||||
ids = _result_ids(result)
|
||||
assert ID_OLD not in ids
|
||||
assert ID_MID not in ids, "fact-mid (updated_at=T2) must be excluded by created_after=T2"
|
||||
assert ID_NEW in ids
|
||||
|
||||
async def test_created_before_excludes_new(self, seeded_memory):
|
||||
"""created_before=T3 excludes fact-new (updated_at == T3, not < T3)."""
|
||||
engine, bank_id = seeded_memory
|
||||
result = await engine.recall_async(
|
||||
bank_id=bank_id,
|
||||
query="animals and nature",
|
||||
request_context=RC,
|
||||
max_tokens=10000,
|
||||
created_before=T3,
|
||||
)
|
||||
ids = _result_ids(result)
|
||||
assert ID_OLD in ids
|
||||
assert ID_MID in ids
|
||||
assert ID_NEW not in ids, "fact-new (updated_at=T3) must be excluded by created_before=T3"
|
||||
|
||||
async def test_created_before_excludes_mid_and_new(self, seeded_memory):
|
||||
"""created_before=T2 returns only fact-old."""
|
||||
engine, bank_id = seeded_memory
|
||||
result = await engine.recall_async(
|
||||
bank_id=bank_id,
|
||||
query="animals and nature",
|
||||
request_context=RC,
|
||||
max_tokens=10000,
|
||||
created_before=T2,
|
||||
)
|
||||
ids = _result_ids(result)
|
||||
assert ID_OLD in ids
|
||||
assert ID_MID not in ids
|
||||
assert ID_NEW not in ids
|
||||
|
||||
async def test_range_both_bounds(self, seeded_memory):
|
||||
"""created_after=T1, created_before=T3 returns only fact-mid."""
|
||||
engine, bank_id = seeded_memory
|
||||
result = await engine.recall_async(
|
||||
bank_id=bank_id,
|
||||
query="animals and nature",
|
||||
request_context=RC,
|
||||
max_tokens=10000,
|
||||
created_after=T1,
|
||||
created_before=T3,
|
||||
)
|
||||
ids = _result_ids(result)
|
||||
assert ID_OLD not in ids
|
||||
assert ID_MID in ids, "fact-mid (T2) must be in range (T1, T3)"
|
||||
assert ID_NEW not in ids
|
||||
|
||||
async def test_empty_range_returns_nothing(self, seeded_memory):
|
||||
"""A range after all facts returns empty results."""
|
||||
engine, bank_id = seeded_memory
|
||||
result = await engine.recall_async(
|
||||
bank_id=bank_id,
|
||||
query="animals and nature",
|
||||
request_context=RC,
|
||||
max_tokens=10000,
|
||||
created_after=T3,
|
||||
)
|
||||
assert len(result.results) == 0, f"Expected no results after T3, got: {_result_ids(result)}"
|
||||
|
||||
async def test_updated_at_catches_consolidation_updates(self, seeded_memory):
|
||||
"""A fact created at T1 but updated at T3 appears with created_after=T2."""
|
||||
engine, bank_id = seeded_memory
|
||||
|
||||
# Simulate consolidation updating fact-old
|
||||
pool = await engine._get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
await conn.execute(
|
||||
"UPDATE memory_units SET updated_at = $1 WHERE id = $2",
|
||||
T3,
|
||||
ID_OLD,
|
||||
)
|
||||
|
||||
result = await engine.recall_async(
|
||||
bank_id=bank_id,
|
||||
query="animals and nature",
|
||||
request_context=RC,
|
||||
max_tokens=10000,
|
||||
created_after=T2,
|
||||
)
|
||||
ids = _result_ids(result)
|
||||
assert ID_OLD in ids, "fact-old created at T1, updated at T3 — created_after=T2 must find it via updated_at"
|
||||
assert ID_NEW in ids
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user