Compare commits

..
Author SHA1 Message Date
DK09876andClaude Opus 4.6 b22f4bf258 fix(opencode): fix message parsing, shared state, and post-compaction retain
Three bugs fixed:
1. msg.role → msg.info.role: OpenCode SDK wraps role inside info, so all
   messages were silently filtered out, breaking retain and recall (#941)
2. Move PluginState to module level so it persists across sessions instead
   of being recreated per plugin instantiation
3. Reset lastRetainedTurn after compaction so idle-retain resumes when the
   message list shrinks

Co-Authored-By: Claude Opus 4.6 <[email protected]>
2026-04-13 09:04:11 -07:00
876 changed files with 13582 additions and 58960 deletions
+1 -1
View File
@@ -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:
-174
View File
@@ -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
View File
@@ -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;
-7
View File
@@ -1,7 +0,0 @@
{
"semi": true,
"singleQuote": false,
"tabWidth": 2,
"trailingComma": "es5",
"printWidth": 100
}
+1 -2
View File
@@ -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 -2
View File
@@ -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
View File
@@ -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 -1
View File
@@ -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",
+20 -24
View File
@@ -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']);
});
});
+3 -3
View File
@@ -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}`];
}
+6 -6
View File
@@ -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';
+15 -15
View File
@@ -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();
+43 -43
View File
@@ -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 -1
View File
@@ -1,4 +1,4 @@
import type { Logger } from "./logger.js";
import type { Logger } from './logger.js';
/**
* Options for {@link HindsightServer}.
+4 -4
View File
@@ -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,
+3 -3
View File
@@ -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',
},
});
+1 -1
View File
@@ -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"
+30 -41
View File
@@ -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:
+1 -1
View File
@@ -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"
-39
View File
@@ -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()
+1 -1
View File
@@ -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,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
@@ -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
@@ -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")
@@ -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
@@ -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")
@@ -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
@@ -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 '{{}}'")
@@ -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
@@ -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'))"
)
+7 -409
View File
@@ -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",
+10 -230
View File
@@ -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)
+88 -98
View File
@@ -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 = [
+1 -6
View File
@@ -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
+104 -475
View File
@@ -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.
+5 -7
View File
@@ -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"}]
-248
View File
@@ -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})"
)
+1 -1
View File
@@ -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)
-141
View File
@@ -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.
+29 -125
View File
@@ -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)
+1 -404
View File
@@ -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