Compare commits
82
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
7f2c986135 | ||
|
|
cf13ed7161 | ||
|
|
bd3ad4603e | ||
|
|
862041d5ca | ||
|
|
f880752e9b | ||
|
|
eadb510eb3 | ||
|
|
691cb5394b | ||
|
|
a401b97eb7 | ||
|
|
aa4c1bbaf3 | ||
|
|
99525144b2 | ||
|
|
ded52e8de6 | ||
|
|
cc3ba4a37c | ||
|
|
7ef64f14ca | ||
|
|
16f807697d | ||
|
|
486c3a8b3b | ||
|
|
fbbc7a5e4c | ||
|
|
d7d41e76c2 | ||
|
|
262d4894f2 | ||
|
|
ec49175fa3 | ||
|
|
488f428009 | ||
|
|
d1ef9da95e | ||
|
|
30acca6fd9 | ||
|
|
2538708308 | ||
|
|
9e7aff6bd4 | ||
|
|
a908cdc974 | ||
|
|
4cd260b691 | ||
|
|
0c17e9acfd | ||
|
|
beca4b42f3 | ||
|
|
6e9b741b02 | ||
|
|
4a1b2f39c1 | ||
|
|
28ec22c3dc | ||
|
|
d802f91488 | ||
|
|
cb04cb79d9 | ||
|
|
dabbf9ff49 | ||
|
|
6348f42451 | ||
|
|
41a2ccabf8 | ||
|
|
eaf3048f2c | ||
|
|
9d95149852 | ||
|
|
ac3ab2b54c | ||
|
|
cb037290bb | ||
|
|
2582b45a16 | ||
|
|
0be157eeb5 | ||
|
|
90cb145aa6 | ||
|
|
7bd11bedf6 | ||
|
|
c3b2b1543a | ||
|
|
2743d061f7 | ||
|
|
daf2348bcd | ||
|
|
46dd2dfd94 | ||
|
|
878ef957f7 | ||
|
|
00d327a049 | ||
|
|
31d1e1729e | ||
|
|
592f01bba6 | ||
|
|
da05ee7215 | ||
|
|
19d23921fb | ||
|
|
e1e1a5e02b | ||
|
|
44b34c891c | ||
|
|
67ae2a41d4 | ||
|
|
2e5186a6fc | ||
|
|
9a20180415 | ||
|
|
f61ae2a185 | ||
|
|
dfd7cb52d4 | ||
|
|
5300d401b0 | ||
|
|
0b6bf53bef | ||
|
|
203ddfdd6c | ||
|
|
dcf5588e6c | ||
|
|
3d6c2ba8b0 | ||
|
|
80046797f7 | ||
|
|
db7dabcebd | ||
|
|
b83bb87ddd | ||
|
|
15ec55b703 | ||
|
|
f2596e1fe9 | ||
|
|
21c71f7bb8 | ||
|
|
86b686cd72 | ||
|
|
248c40e670 | ||
|
|
d18a9452ad | ||
|
|
806fbcd41c | ||
|
|
a75c3c85ad | ||
|
|
0db9f3da19 | ||
|
|
8940710c72 | ||
|
|
6252643de0 | ||
|
|
3fce309c0d | ||
|
|
3fc361aabd |
+21
-1
@@ -64,8 +64,21 @@ HINDSIGHT_API_LOG_LEVEL=info
|
||||
# For Azure PostgreSQL with DiskANN:
|
||||
# HINDSIGHT_API_VECTOR_EXTENSION=pgvectorscale # Auto-detects pg_diskann on Azure
|
||||
|
||||
# Text Search Extension (Optional - uses native PostgreSQL full-text search by default)
|
||||
# Backend options: "native" (default), "vchord", "pg_textsearch", "pgroonga", "pg_search"
|
||||
# HINDSIGHT_API_TEXT_SEARCH_EXTENSION=native
|
||||
# Native backend dictionary (only used by HINDSIGHT_API_TEXT_SEARCH_EXTENSION=native)
|
||||
# HINDSIGHT_API_TEXT_SEARCH_EXTENSION_NATIVE_LANGUAGE=english
|
||||
# ParadeDB pg_search tokenizer (only used when creating pg_search BM25 indexes).
|
||||
# Empty uses ParadeDB's default tokenizer: unicode_words.
|
||||
# Supported values: unicode_words, simple, whitespace, literal, literal_normalized,
|
||||
# chinese_compatible, icu, jieba, source_code,
|
||||
# chinese_lindera/lindera(chinese), japanese_lindera/lindera(japanese),
|
||||
# korean_lindera/lindera(korean), ngram(min,max), edge_ngram(min,max)
|
||||
# HINDSIGHT_API_TEXT_SEARCH_EXTENSION_PG_SEARCH_TOKENIZER=
|
||||
|
||||
# Embeddings Configuration (Optional - uses local by default)
|
||||
# Provider: "local" (default), "tei", "openai", "cohere", "google", "openrouter", "litellm", or "litellm-sdk"
|
||||
# Provider: "local" (default), "tei", "openai", "cohere", "google", "openrouter", "zeroentropy", "litellm", or "litellm-sdk"
|
||||
# HINDSIGHT_API_EMBEDDINGS_PROVIDER=local
|
||||
# For local provider:
|
||||
# HINDSIGHT_API_EMBEDDINGS_LOCAL_MODEL=BAAI/bge-small-en-v1.5
|
||||
@@ -77,6 +90,13 @@ HINDSIGHT_API_LOG_LEVEL=info
|
||||
# HINDSIGHT_API_EMBEDDINGS_OPENAI_API_KEY=sk-xxxx
|
||||
# HINDSIGHT_API_EMBEDDINGS_OPENAI_MODEL=text-embedding-3-small
|
||||
# HINDSIGHT_API_EMBEDDINGS_OPENAI_BASE_URL=https://api.openai.com/v1
|
||||
# For ZeroEntropy zembed-1:
|
||||
# HINDSIGHT_API_EMBEDDINGS_PROVIDER=zeroentropy
|
||||
# HINDSIGHT_API_EMBEDDINGS_ZEROENTROPY_API_KEY=ze-xxxx
|
||||
# HINDSIGHT_API_EMBEDDINGS_ZEROENTROPY_MODEL=zembed-1
|
||||
# HINDSIGHT_API_EMBEDDINGS_ZEROENTROPY_DIMENSIONS=1280
|
||||
# HINDSIGHT_API_EMBEDDINGS_ZEROENTROPY_ENCODING_FORMAT=float
|
||||
# HINDSIGHT_API_EMBEDDINGS_ZEROENTROPY_LATENCY=fast
|
||||
#
|
||||
# IMPORTANT: Embedding keys require provider-specific names:
|
||||
# HINDSIGHT_API_EMBEDDINGS_{PROVIDER}_{PARAMETER}
|
||||
|
||||
@@ -734,6 +734,12 @@ jobs:
|
||||
rm -rf node_modules/lightningcss node_modules/@tailwindcss
|
||||
npm install lightningcss @tailwindcss/postcss @tailwindcss/node
|
||||
|
||||
- name: Test Control Plane
|
||||
run: npm test --workspace=hindsight-control-plane
|
||||
|
||||
- name: Check i18n locale parity and hardcoded strings
|
||||
run: npm run i18n:check --workspace=hindsight-control-plane
|
||||
|
||||
- name: Build Control Plane
|
||||
run: npm run build --workspace=hindsight-control-plane
|
||||
|
||||
@@ -935,6 +941,22 @@ jobs:
|
||||
- name: Lint Helm chart
|
||||
run: helm lint helm/hindsight
|
||||
|
||||
test-standalone-start-script:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
(github.event_name == 'workflow_dispatch' ||
|
||||
needs.detect-changes.outputs.docker == '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: Run standalone start script tests
|
||||
run: bash docker/standalone/test-start-all.sh
|
||||
|
||||
build-docker-images:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
@@ -1102,7 +1124,70 @@ jobs:
|
||||
|
||||
- name: Run tests
|
||||
working-directory: ./hindsight-api-slim
|
||||
run: uv run pytest tests -v -m "not hs_llm_mat"
|
||||
run: uv run pytest tests -v -m "not hs_llm_mat and not hs_llm_core"
|
||||
|
||||
test-api-llm-core:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
needs.detect-changes.outputs.has_secrets == 'true' &&
|
||||
(github.event_name == 'workflow_dispatch' ||
|
||||
needs.detect-changes.outputs.core == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
name: Core LLM tests
|
||||
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
|
||||
GEMINI_API_KEY: ${{ secrets.GEMINI_API_KEY }}
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ github.event.pull_request.head.sha || '' }}
|
||||
|
||||
- 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: Install dependencies
|
||||
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-${{ hashFiles('hindsight-api-slim/pyproject.toml') }}
|
||||
restore-keys: |
|
||||
${{ runner.os }}-huggingface-
|
||||
|
||||
- name: Pre-download models
|
||||
working-directory: ./hindsight-api-slim
|
||||
run: |
|
||||
uv run python -c "
|
||||
from sentence_transformers import SentenceTransformer, CrossEncoder
|
||||
SentenceTransformer('BAAI/bge-small-en-v1.5')
|
||||
CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2')
|
||||
print('Models downloaded successfully')
|
||||
"
|
||||
|
||||
- name: Run core LLM tests
|
||||
working-directory: ./hindsight-api-slim
|
||||
run: uv run pytest tests -v -m "hs_llm_core" --timeout 600
|
||||
|
||||
test-api-llm-acceptance:
|
||||
needs: [detect-changes]
|
||||
@@ -1260,8 +1345,8 @@ jobs:
|
||||
conn = oracledb.connect(user='system', password='oracle', dsn='localhost:1521/FREEPDB1')
|
||||
cursor = conn.cursor()
|
||||
cursor.execute(\"\"\"
|
||||
CREATE TABLESPACE hindsight_ts
|
||||
DATAFILE 'hindsight_ts.dbf' SIZE 200M AUTOEXTEND ON NEXT 50M
|
||||
CREATE BIGFILE TABLESPACE hindsight_ts
|
||||
DATAFILE 'hindsight_ts.dbf' SIZE 2G AUTOEXTEND ON NEXT 500M MAXSIZE UNLIMITED
|
||||
EXTENT MANAGEMENT LOCAL
|
||||
SEGMENT SPACE MANAGEMENT AUTO
|
||||
\"\"\")
|
||||
@@ -1614,8 +1699,8 @@ jobs:
|
||||
conn = oracledb.connect(user='system', password='oracle', dsn='localhost:1521/FREEPDB1')
|
||||
cursor = conn.cursor()
|
||||
cursor.execute(\"\"\"
|
||||
CREATE TABLESPACE hindsight_ts
|
||||
DATAFILE 'hindsight_ts.dbf' SIZE 200M AUTOEXTEND ON NEXT 50M
|
||||
CREATE BIGFILE TABLESPACE hindsight_ts
|
||||
DATAFILE 'hindsight_ts.dbf' SIZE 2G AUTOEXTEND ON NEXT 500M MAXSIZE UNLIMITED
|
||||
EXTENT MANAGEMENT LOCAL
|
||||
SEGMENT SPACE MANAGEMENT AUTO
|
||||
\"\"\")
|
||||
@@ -1773,8 +1858,8 @@ jobs:
|
||||
conn = oracledb.connect(user='system', password='oracle', dsn='localhost:1521/FREEPDB1')
|
||||
cursor = conn.cursor()
|
||||
cursor.execute(\"\"\"
|
||||
CREATE TABLESPACE hindsight_ts
|
||||
DATAFILE 'hindsight_ts.dbf' SIZE 200M AUTOEXTEND ON NEXT 50M
|
||||
CREATE BIGFILE TABLESPACE hindsight_ts
|
||||
DATAFILE 'hindsight_ts.dbf' SIZE 2G AUTOEXTEND ON NEXT 500M MAXSIZE UNLIMITED
|
||||
EXTENT MANAGEMENT LOCAL
|
||||
SEGMENT SPACE MANAGEMENT AUTO
|
||||
\"\"\")
|
||||
@@ -3200,6 +3285,12 @@ jobs:
|
||||
rm -rf node_modules/lightningcss node_modules/@tailwindcss
|
||||
npm install lightningcss @tailwindcss/postcss @tailwindcss/node
|
||||
|
||||
- name: Test Control Plane
|
||||
run: npm test --workspace=hindsight-control-plane
|
||||
|
||||
- name: Check i18n locale parity and hardcoded strings
|
||||
run: npm run i18n:check --workspace=hindsight-control-plane
|
||||
|
||||
- name: Build Control Plane
|
||||
run: npm run build --workspace=hindsight-control-plane
|
||||
|
||||
@@ -3651,6 +3742,7 @@ jobs:
|
||||
- build-docs
|
||||
- test-rust-cli
|
||||
- lint-helm-chart
|
||||
- test-standalone-start-script
|
||||
- build-docker-images
|
||||
- test-api
|
||||
- test-api-oracle
|
||||
|
||||
@@ -0,0 +1,7 @@
|
||||
# PostgreSQL with pgvector and ParadeDB pg_search extensions.
|
||||
#
|
||||
# The official ParadeDB image ships PostgreSQL with pg_search and pgvector
|
||||
# already installed, so no build steps are required. We pin to the PG17
|
||||
# variant for parity with the other Hindsight docker-compose examples
|
||||
# (vchord, pg_textsearch).
|
||||
FROM paradedb/paradedb:latest-pg17
|
||||
@@ -0,0 +1,96 @@
|
||||
name: hindsight
|
||||
# Docker Compose file for Hindsight with PostgreSQL and ParadeDB pg_search.
|
||||
#
|
||||
# pg_search is the only BM25 backend supported by Hindsight that works with
|
||||
# Citus, so this is the recommended setup for horizontally scaled deployments.
|
||||
#
|
||||
# Usage:
|
||||
# docker compose -f docker/docker-compose/pg_search/docker-compose.yaml up -d
|
||||
#
|
||||
# Required environment variables:
|
||||
# - HINDSIGHT_DB_PASSWORD: Password for the PostgreSQL user
|
||||
# - Configure LLM provider variables as needed (see the hindsight service)
|
||||
#
|
||||
# Optional environment variables with defaults:
|
||||
# - HINDSIGHT_VERSION: Hindsight application version (default: latest)
|
||||
# - HINDSIGHT_DB_USER: PostgreSQL user (default: hindsight_user)
|
||||
# - HINDSIGHT_DB_NAME: PostgreSQL database name (default: hindsight_db)
|
||||
# - HINDSIGHT_API_TEXT_SEARCH_EXTENSION_PG_SEARCH_TOKENIZER: ParadeDB pg_search
|
||||
# tokenizer for new BM25 indexes (default: empty, uses ParadeDB default)
|
||||
|
||||
services:
|
||||
db:
|
||||
# Use ParadeDB image which bundles pgvector + pg_search
|
||||
build:
|
||||
context: .
|
||||
dockerfile: Dockerfile
|
||||
container_name: hindsight-db
|
||||
restart: always
|
||||
ports:
|
||||
- "5437:5432"
|
||||
environment:
|
||||
POSTGRES_USER: ${HINDSIGHT_DB_USER:-hindsight_user}
|
||||
POSTGRES_PASSWORD: ${HINDSIGHT_DB_PASSWORD:-hindsight_password}
|
||||
POSTGRES_DB: ${HINDSIGHT_DB_NAME:-hindsight_db}
|
||||
volumes:
|
||||
- pg_data:/var/lib/postgresql/data
|
||||
networks:
|
||||
- hindsight-net
|
||||
|
||||
pg-search-init:
|
||||
build:
|
||||
context: .
|
||||
dockerfile: Dockerfile
|
||||
depends_on:
|
||||
- db
|
||||
environment:
|
||||
- PGPASSWORD=${HINDSIGHT_DB_PASSWORD:-hindsight_password}
|
||||
command: >
|
||||
bash -c "
|
||||
echo 'Waiting for PostgreSQL to be ready...';
|
||||
until pg_isready -h hindsight-db -p 5432 -U hindsight_user; do
|
||||
echo 'PostgreSQL is unavailable - sleeping';
|
||||
sleep 2;
|
||||
done;
|
||||
echo 'PostgreSQL is ready - creating hindsight_db database';
|
||||
psql -h hindsight-db -p 5432 -U hindsight_user -c 'CREATE DATABASE hindsight_db;' 2>/dev/null || echo 'Database already exists';
|
||||
echo 'Creating extensions in hindsight_db database';
|
||||
psql -h hindsight-db -p 5432 -U hindsight_user -d hindsight_db -c 'CREATE EXTENSION IF NOT EXISTS vector CASCADE;';
|
||||
psql -h hindsight-db -p 5432 -U hindsight_user -d hindsight_db -c 'CREATE EXTENSION IF NOT EXISTS pg_search CASCADE;';
|
||||
echo 'Database and extensions created successfully';
|
||||
"
|
||||
restart: "no"
|
||||
networks:
|
||||
- hindsight-net
|
||||
|
||||
hindsight:
|
||||
image: ghcr.io/vectorize-io/hindsight:${HINDSIGHT_VERSION:-latest}
|
||||
container_name: hindsight-app
|
||||
ports:
|
||||
- "8888:8888"
|
||||
- "9999:9999"
|
||||
environment:
|
||||
# LLM Configuration
|
||||
HINDSIGHT_API_LLM_PROVIDER: ${HINDSIGHT_API_LLM_PROVIDER:-openai}
|
||||
HINDSIGHT_API_LLM_API_KEY: ${OPENAI_API_KEY:-your-api-key}
|
||||
|
||||
# Database Configuration
|
||||
HINDSIGHT_API_DATABASE_URL: postgresql://${HINDSIGHT_DB_USER:-hindsight_user}:${HINDSIGHT_DB_PASSWORD:-hindsight_password}@db:5432/${HINDSIGHT_DB_NAME:-hindsight_db}
|
||||
|
||||
# Vector and Text Search Extensions
|
||||
HINDSIGHT_API_VECTOR_EXTENSION: pgvector
|
||||
HINDSIGHT_API_TEXT_SEARCH_EXTENSION: pg_search
|
||||
HINDSIGHT_API_TEXT_SEARCH_EXTENSION_PG_SEARCH_TOKENIZER: ${HINDSIGHT_API_TEXT_SEARCH_EXTENSION_PG_SEARCH_TOKENIZER:-}
|
||||
|
||||
depends_on:
|
||||
- db
|
||||
networks:
|
||||
- hindsight-net
|
||||
|
||||
|
||||
networks:
|
||||
hindsight-net:
|
||||
driver: bridge
|
||||
|
||||
volumes:
|
||||
pg_data:
|
||||
@@ -0,0 +1,23 @@
|
||||
# PostgreSQL with pgvector and pgroonga extensions.
|
||||
#
|
||||
# pgroonga is a multilingual full-text search extension built on Groonga.
|
||||
# It works out of the box for CJK (Chinese, Japanese, Korean) and other
|
||||
# non-whitespace-segmented languages via the TokenBigram tokenizer.
|
||||
FROM groonga/pgroonga:latest-debian-pg17
|
||||
|
||||
# Install pgvector on top of the pgroonga base image (which already provides
|
||||
# pgroonga and the Groonga library).
|
||||
RUN apt-get update && apt-get install -y --no-install-recommends \
|
||||
build-essential \
|
||||
git \
|
||||
postgresql-server-dev-17 \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
RUN cd /tmp && \
|
||||
git clone --branch v0.8.0 https://github.com/pgvector/pgvector.git && \
|
||||
cd pgvector && \
|
||||
make && \
|
||||
make install
|
||||
|
||||
RUN rm -rf /tmp/pgvector && \
|
||||
apt-get purge -y --auto-remove build-essential git postgresql-server-dev-17
|
||||
@@ -0,0 +1,91 @@
|
||||
name: hindsight
|
||||
# Docker Compose file for Hindsight with PostgreSQL and pgroonga
|
||||
#
|
||||
# pgroonga provides multilingual BM25 indexing that works out of the box for
|
||||
# CJK (Chinese, Japanese, Korean) and other non-whitespace-segmented languages.
|
||||
# Use this recipe if your bank content is not English/European.
|
||||
#
|
||||
# docker compose -f docker/docker-compose/pgroonga/docker-compose.yaml down && \
|
||||
# sleep 2 && \
|
||||
# docker compose -f docker/docker-compose/pgroonga/docker-compose.yaml up -d
|
||||
#
|
||||
# Optional environment variables with defaults:
|
||||
# - HINDSIGHT_VERSION: Hindsight application version (default: latest)
|
||||
# - HINDSIGHT_DB_USER: PostgreSQL user (default: hindsight_user)
|
||||
# - HINDSIGHT_DB_NAME: PostgreSQL database name (default: hindsight_db)
|
||||
# - HINDSIGHT_DB_PASSWORD: PostgreSQL password (default: hindsight_password)
|
||||
|
||||
services:
|
||||
db:
|
||||
build:
|
||||
context: .
|
||||
dockerfile: Dockerfile
|
||||
container_name: hindsight-db
|
||||
restart: always
|
||||
ports:
|
||||
- "5439:5432"
|
||||
environment:
|
||||
POSTGRES_USER: ${HINDSIGHT_DB_USER:-hindsight_user}
|
||||
POSTGRES_PASSWORD: ${HINDSIGHT_DB_PASSWORD:-hindsight_password}
|
||||
POSTGRES_DB: ${HINDSIGHT_DB_NAME:-hindsight_db}
|
||||
volumes:
|
||||
- pg_data:/var/lib/postgresql/data
|
||||
networks:
|
||||
- hindsight-net
|
||||
|
||||
pgroonga-init:
|
||||
build:
|
||||
context: .
|
||||
dockerfile: Dockerfile
|
||||
depends_on:
|
||||
- db
|
||||
environment:
|
||||
- PGPASSWORD=${HINDSIGHT_DB_PASSWORD:-hindsight_password}
|
||||
command: >
|
||||
bash -c "
|
||||
echo 'Waiting for PostgreSQL to be ready...';
|
||||
until pg_isready -h hindsight-db -p 5432 -U hindsight_user; do
|
||||
echo 'PostgreSQL is unavailable - sleeping';
|
||||
sleep 2;
|
||||
done;
|
||||
echo 'PostgreSQL is ready - creating hindsight_db database';
|
||||
psql -h hindsight-db -p 5432 -U hindsight_user -c 'CREATE DATABASE hindsight_db;' 2>/dev/null || echo 'Database already exists';
|
||||
echo 'Creating extensions in hindsight_db database';
|
||||
psql -h hindsight-db -p 5432 -U hindsight_user -d hindsight_db -c 'CREATE EXTENSION IF NOT EXISTS vector CASCADE;';
|
||||
psql -h hindsight-db -p 5432 -U hindsight_user -d hindsight_db -c 'CREATE EXTENSION IF NOT EXISTS pgroonga CASCADE;';
|
||||
echo 'Database and extensions created successfully';
|
||||
"
|
||||
restart: "no"
|
||||
networks:
|
||||
- hindsight-net
|
||||
|
||||
hindsight:
|
||||
image: ghcr.io/vectorize-io/hindsight:${HINDSIGHT_VERSION:-latest}
|
||||
container_name: hindsight-app
|
||||
ports:
|
||||
- "8888:8888"
|
||||
- "9999:9999"
|
||||
environment:
|
||||
# LLM Configuration
|
||||
HINDSIGHT_API_LLM_PROVIDER: ${HINDSIGHT_API_LLM_PROVIDER:-openai}
|
||||
HINDSIGHT_API_LLM_API_KEY: ${OPENAI_API_KEY:-your-api-key}
|
||||
|
||||
# Database Configuration
|
||||
HINDSIGHT_API_DATABASE_URL: postgresql://${HINDSIGHT_DB_USER:-hindsight_user}:${HINDSIGHT_DB_PASSWORD:-hindsight_password}@db:5432/${HINDSIGHT_DB_NAME:-hindsight_db}
|
||||
|
||||
# Vector and Text Search Extensions
|
||||
HINDSIGHT_API_VECTOR_EXTENSION: pgvector
|
||||
HINDSIGHT_API_TEXT_SEARCH_EXTENSION: pgroonga
|
||||
|
||||
depends_on:
|
||||
- db
|
||||
networks:
|
||||
- hindsight-net
|
||||
|
||||
|
||||
networks:
|
||||
hindsight-net:
|
||||
driver: bridge
|
||||
|
||||
volumes:
|
||||
pg_data:
|
||||
@@ -10,19 +10,45 @@ set -e
|
||||
# loss scenarios where a container restart caused the data directory to be
|
||||
# wiped despite a volume mount being present.
|
||||
# =============================================================================
|
||||
PG0_DATA_DIR="${HOME}/.pg0"
|
||||
if [ -d "$PG0_DATA_DIR" ]; then
|
||||
pg0_has_pg_version() {
|
||||
local pg0_data_dir="$1"
|
||||
|
||||
# pg0 has used more than one on-disk layout. Newer standalone images keep
|
||||
# PostgreSQL data under instances/<name>/data, while older volumes may have
|
||||
# placed PG_VERSION at or one level below the mount.
|
||||
[ -f "$pg0_data_dir/PG_VERSION" ] && return 0
|
||||
compgen -G "$pg0_data_dir"/*/PG_VERSION > /dev/null 2>&1 && return 0
|
||||
compgen -G "$pg0_data_dir"/instances/*/data/PG_VERSION > /dev/null 2>&1 && return 0
|
||||
|
||||
return 1
|
||||
}
|
||||
|
||||
check_pg0_data_integrity() {
|
||||
local pg0_data_dir="$1"
|
||||
|
||||
if [ ! -d "$pg0_data_dir" ]; then
|
||||
return 0
|
||||
fi
|
||||
|
||||
# Look for actual PostgreSQL data directories (pg0 creates subdirs per instance)
|
||||
if compgen -G "$PG0_DATA_DIR"/*/PG_VERSION > /dev/null 2>&1; then
|
||||
echo "✅ Existing pg0 data directory detected at $PG0_DATA_DIR"
|
||||
elif [ "$(ls -A "$PG0_DATA_DIR" 2>/dev/null)" ]; then
|
||||
echo "⚠️ WARNING: pg0 data directory exists at $PG0_DATA_DIR but no PG_VERSION found."
|
||||
if pg0_has_pg_version "$pg0_data_dir"; then
|
||||
echo "✅ Existing pg0 data directory detected at $pg0_data_dir"
|
||||
elif [ "$(ls -A "$pg0_data_dir" 2>/dev/null)" ]; then
|
||||
echo "⚠️ WARNING: pg0 data directory exists at $pg0_data_dir but no PG_VERSION found."
|
||||
echo " This may indicate data corruption or an incomplete previous shutdown."
|
||||
echo " If you see all migrations running from scratch after this, your data may have been lost."
|
||||
echo " See: https://github.com/vectorize-io/hindsight/issues/675"
|
||||
fi
|
||||
|
||||
return 0
|
||||
}
|
||||
|
||||
if [ "${HINDSIGHT_START_ALL_SOURCE_ONLY:-false}" = "true" ]; then
|
||||
return 0 2>/dev/null || exit 0
|
||||
fi
|
||||
|
||||
check_pg0_data_integrity "${HOME}/.pg0"
|
||||
|
||||
# Service flags (default to true if not set)
|
||||
ENABLE_API="${HINDSIGHT_ENABLE_API:-true}"
|
||||
ENABLE_CP="${HINDSIGHT_ENABLE_CP:-true}"
|
||||
@@ -156,7 +182,7 @@ PIDS=()
|
||||
# Start API if enabled
|
||||
if [ "$ENABLE_API" = "true" ]; then
|
||||
cd /app/api
|
||||
API_HEALTH_URL="${HINDSIGHT_API_HEALTH_URL:-http://localhost:8888/health}"
|
||||
API_HEALTH_URL="${HINDSIGHT_API_HEALTH_URL:-http://localhost:${HINDSIGHT_API_PORT:-8888}/health}"
|
||||
API_STARTUP_WAIT_SECONDS="${HINDSIGHT_API_STARTUP_WAIT_SECONDS:-300}"
|
||||
|
||||
# Run API directly - Python's PYTHONUNBUFFERED=1 handles output buffering
|
||||
|
||||
Executable
+73
@@ -0,0 +1,73 @@
|
||||
#!/bin/bash
|
||||
set -euo pipefail
|
||||
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
|
||||
HINDSIGHT_START_ALL_SOURCE_ONLY=true
|
||||
source "$SCRIPT_DIR/start-all.sh"
|
||||
unset HINDSIGHT_START_ALL_SOURCE_ONLY
|
||||
|
||||
TMP_DIR="$(mktemp -d)"
|
||||
trap 'rm -rf "$TMP_DIR"' EXIT
|
||||
|
||||
assert_contains() {
|
||||
local output="$1"
|
||||
local expected="$2"
|
||||
|
||||
if [[ "$output" != *"$expected"* ]]; then
|
||||
echo "Expected output to contain: $expected"
|
||||
echo "Actual output:"
|
||||
echo "$output"
|
||||
exit 1
|
||||
fi
|
||||
}
|
||||
|
||||
assert_not_contains() {
|
||||
local output="$1"
|
||||
local unexpected="$2"
|
||||
|
||||
if [[ "$output" == *"$unexpected"* ]]; then
|
||||
echo "Expected output not to contain: $unexpected"
|
||||
echo "Actual output:"
|
||||
echo "$output"
|
||||
exit 1
|
||||
fi
|
||||
}
|
||||
|
||||
assert_empty() {
|
||||
local output="$1"
|
||||
|
||||
if [ -n "$output" ]; then
|
||||
echo "Expected no output, got:"
|
||||
echo "$output"
|
||||
exit 1
|
||||
fi
|
||||
}
|
||||
|
||||
mkdir -p "$TMP_DIR/empty"
|
||||
assert_empty "$(check_pg0_data_integrity "$TMP_DIR/empty")"
|
||||
|
||||
mkdir -p "$TMP_DIR/direct"
|
||||
touch "$TMP_DIR/direct/PG_VERSION"
|
||||
direct_output="$(check_pg0_data_integrity "$TMP_DIR/direct")"
|
||||
assert_contains "$direct_output" "Existing pg0 data directory detected"
|
||||
assert_not_contains "$direct_output" "WARNING"
|
||||
|
||||
mkdir -p "$TMP_DIR/legacy/instance"
|
||||
touch "$TMP_DIR/legacy/instance/PG_VERSION"
|
||||
legacy_output="$(check_pg0_data_integrity "$TMP_DIR/legacy")"
|
||||
assert_contains "$legacy_output" "Existing pg0 data directory detected"
|
||||
assert_not_contains "$legacy_output" "WARNING"
|
||||
|
||||
mkdir -p "$TMP_DIR/nested/instances/hindsight/data"
|
||||
touch "$TMP_DIR/nested/instances/hindsight/data/PG_VERSION"
|
||||
nested_output="$(check_pg0_data_integrity "$TMP_DIR/nested")"
|
||||
assert_contains "$nested_output" "Existing pg0 data directory detected"
|
||||
assert_not_contains "$nested_output" "WARNING"
|
||||
|
||||
mkdir -p "$TMP_DIR/nonempty/instances/hindsight"
|
||||
touch "$TMP_DIR/nonempty/instances/hindsight/instance.json"
|
||||
nonempty_output="$(check_pg0_data_integrity "$TMP_DIR/nonempty")"
|
||||
assert_contains "$nonempty_output" "WARNING: pg0 data directory exists"
|
||||
|
||||
echo "start-all pg0 integrity checks passed"
|
||||
@@ -1,6 +0,0 @@
|
||||
dependencies:
|
||||
- name: postgresql
|
||||
repository: https://charts.bitnami.com/bitnami
|
||||
version: 15.5.38
|
||||
digest: sha256:f67c7612736803ece8a669f8ca6b0555f3b78557bc0ecb732aa2e43f0df7750d
|
||||
generated: "2025-12-10T17:20:57.058794+01:00"
|
||||
@@ -2,8 +2,8 @@ apiVersion: v2
|
||||
name: hindsight
|
||||
description: Hindsight helm chart
|
||||
type: application
|
||||
version: 0.6.2
|
||||
appVersion: "0.6.2"
|
||||
version: 0.7.0
|
||||
appVersion: "0.7.0"
|
||||
keywords:
|
||||
- ai
|
||||
- memory
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@vectorize-io/hindsight-all",
|
||||
"version": "0.6.2",
|
||||
"version": "0.7.0",
|
||||
"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",
|
||||
|
||||
@@ -4,12 +4,12 @@ build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "hindsight-all-slim"
|
||||
version = "0.6.2"
|
||||
version = "0.7.0"
|
||||
description = "Hindsight: Agent Memory That Works Like Human Memory - Slim All-in-One Bundle"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.11"
|
||||
dependencies = [
|
||||
"hindsight-api-slim==0.6.2",
|
||||
"hindsight-api-slim==0.7.0",
|
||||
"hindsight-client>=0.0.7",
|
||||
"hindsight-embed>=0.1.0",
|
||||
]
|
||||
|
||||
@@ -4,12 +4,12 @@ build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "hindsight-all"
|
||||
version = "0.6.2"
|
||||
version = "0.7.0"
|
||||
description = "Hindsight: Agent Memory That Works Like Human Memory - All-in-One Bundle"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.11"
|
||||
dependencies = [
|
||||
"hindsight-api-slim[all]==0.6.2",
|
||||
"hindsight-api-slim[all]==0.7.0",
|
||||
"hindsight-client>=0.0.7",
|
||||
"hindsight-embed>=0.1.0",
|
||||
]
|
||||
@@ -21,7 +21,7 @@ hindsight-embed = { workspace = true }
|
||||
|
||||
[project.optional-dependencies]
|
||||
local-llm = [
|
||||
"hindsight-api-slim[local-llm]==0.6.2",
|
||||
"hindsight-api-slim[local-llm]==0.7.0",
|
||||
]
|
||||
test = [
|
||||
"pytest>=7.0.0",
|
||||
|
||||
@@ -386,7 +386,7 @@ def test_embedded_ui_flag(llm_config):
|
||||
|
||||
# Verify UI is reachable and reports connected dataplane
|
||||
ui_url = client.ui_url
|
||||
assert ui_url, "ui_url should be set"
|
||||
assert isinstance(ui_url, str) and ui_url, "ui_url should be a non-empty string"
|
||||
|
||||
health_url = f"{ui_url}/api/health"
|
||||
with urllib.request.urlopen(health_url, timeout=10) as resp:
|
||||
|
||||
@@ -46,4 +46,4 @@ __all__ = [
|
||||
"RemoteTEICrossEncoder",
|
||||
"LLMConfig",
|
||||
]
|
||||
__version__ = "0.6.2"
|
||||
__version__ = "0.7.0"
|
||||
|
||||
@@ -0,0 +1,85 @@
|
||||
"""Helpers for ParadeDB pg_search index configuration."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from collections.abc import Sequence
|
||||
|
||||
PG_SEARCH_TOKENIZER_ENV = "HINDSIGHT_API_TEXT_SEARCH_EXTENSION_PG_SEARCH_TOKENIZER"
|
||||
|
||||
_SIMPLE_TOKENIZERS = {
|
||||
"unicode_words",
|
||||
"simple",
|
||||
"whitespace",
|
||||
"literal",
|
||||
"literal_normalized",
|
||||
"chinese_compatible",
|
||||
"icu",
|
||||
"jieba",
|
||||
"source_code",
|
||||
}
|
||||
|
||||
_TOKENIZER_ALIASES = {
|
||||
"chinese_lindera": "lindera(chinese)",
|
||||
"japanese_lindera": "lindera(japanese)",
|
||||
"korean_lindera": "lindera(korean)",
|
||||
"lindera_chinese": "lindera(chinese)",
|
||||
"lindera_japanese": "lindera(japanese)",
|
||||
"lindera_korean": "lindera(korean)",
|
||||
}
|
||||
|
||||
|
||||
def normalize_pg_search_tokenizer(value: str | None) -> str:
|
||||
"""Validate and normalize a ParadeDB pg_search tokenizer setting.
|
||||
|
||||
Returns an empty string when unset. The returned value is safe to embed after
|
||||
``pdb.`` in a CREATE INDEX expression.
|
||||
"""
|
||||
|
||||
tokenizer = (value or "").strip().lower()
|
||||
if not tokenizer:
|
||||
return ""
|
||||
|
||||
if tokenizer in _TOKENIZER_ALIASES:
|
||||
return _TOKENIZER_ALIASES[tokenizer]
|
||||
|
||||
if tokenizer in _SIMPLE_TOKENIZERS:
|
||||
return tokenizer
|
||||
|
||||
lindera_match = re.fullmatch(r"lindera\((chinese|japanese|korean)\)", tokenizer)
|
||||
if lindera_match:
|
||||
return tokenizer
|
||||
|
||||
ngram_match = re.fullmatch(r"(ngram|edge_ngram)\((\d{1,3}),\s*(\d{1,3})\)", tokenizer)
|
||||
if ngram_match:
|
||||
kind, min_gram, max_gram = ngram_match.groups()
|
||||
min_value = int(min_gram)
|
||||
max_value = int(max_gram)
|
||||
if min_value <= 0 or min_value > max_value:
|
||||
raise ValueError(
|
||||
f"Invalid {PG_SEARCH_TOKENIZER_ENV}: {value!r}. "
|
||||
"ngram and edge_ngram require positive min/max gram sizes with min <= max."
|
||||
)
|
||||
return f"{kind}({min_value},{max_value})"
|
||||
|
||||
raise ValueError(
|
||||
f"Invalid {PG_SEARCH_TOKENIZER_ENV}: {value!r}. "
|
||||
"Supported values are: unicode_words, simple, whitespace, literal, "
|
||||
"literal_normalized, chinese_compatible, icu, jieba, source_code, "
|
||||
"chinese_lindera, japanese_lindera, korean_lindera, or "
|
||||
"lindera(chinese|japanese|korean), ngram(min,max), or edge_ngram(min,max)."
|
||||
)
|
||||
|
||||
|
||||
def pg_search_bm25_columns(
|
||||
key_field: str,
|
||||
text_fields: Sequence[str],
|
||||
tokenizer: str | None,
|
||||
) -> str:
|
||||
"""Build a ParadeDB BM25 column list for CREATE INDEX."""
|
||||
|
||||
normalized = normalize_pg_search_tokenizer(tokenizer)
|
||||
if not normalized:
|
||||
return ", ".join([key_field, *text_fields])
|
||||
|
||||
return ", ".join([key_field, *(f"({field}::pdb.{normalized})" for field in text_fields)])
|
||||
@@ -268,6 +268,7 @@ async def _run_migration(
|
||||
ensure_text_search_extension(
|
||||
resolved_url,
|
||||
text_search_extension=config.text_search_extension,
|
||||
pg_search_tokenizer=config.text_search_extension_pg_search_tokenizer,
|
||||
schema=schema,
|
||||
)
|
||||
|
||||
|
||||
@@ -15,6 +15,11 @@ from pgvector.sqlalchemy import Vector
|
||||
from sqlalchemy import text
|
||||
from sqlalchemy.dialects import postgresql
|
||||
|
||||
from hindsight_api._pg_search import (
|
||||
PG_SEARCH_TOKENIZER_ENV,
|
||||
normalize_pg_search_tokenizer,
|
||||
pg_search_bm25_columns,
|
||||
)
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
@@ -91,9 +96,14 @@ def _vector_index_using_clause(ext: str) -> str:
|
||||
|
||||
def _detect_text_search_extension() -> str:
|
||||
"""
|
||||
Detect or validate text search extension: 'native', 'vchord', or 'pg_textsearch'.
|
||||
Respects HINDSIGHT_API_TEXT_SEARCH_EXTENSION env var.
|
||||
Detect or validate text search extension: 'native', 'vchord', 'pg_textsearch',
|
||||
'pgroonga', or 'pg_search'. Respects HINDSIGHT_API_TEXT_SEARCH_EXTENSION env var.
|
||||
Creates the extension if needed.
|
||||
|
||||
pgroonga is treated as native here so the initial schema still creates valid
|
||||
tsvector columns. ensure_text_search_extension() at startup converts the
|
||||
schema to pgroonga structures (drops the tsvector column, builds a pgroonga
|
||||
index on the base text column).
|
||||
"""
|
||||
text_search_extension = os.getenv("HINDSIGHT_API_TEXT_SEARCH_EXTENSION", "native").lower()
|
||||
|
||||
@@ -121,14 +131,35 @@ def _detect_text_search_extension() -> str:
|
||||
# Extension truly doesn't exist - re-raise the error
|
||||
raise
|
||||
return "pg_textsearch"
|
||||
elif text_search_extension == "pg_search":
|
||||
# ParadeDB pg_search — true BM25 over base columns, Citus-compatible.
|
||||
try:
|
||||
op.execute("CREATE EXTENSION IF NOT EXISTS pg_search CASCADE")
|
||||
except Exception:
|
||||
# Extension might already exist or user lacks permissions - verify it exists
|
||||
conn = op.get_bind()
|
||||
result = conn.execute(text("SELECT 1 FROM pg_extension WHERE extname = 'pg_search'")).fetchone()
|
||||
if not result:
|
||||
# Extension truly doesn't exist - re-raise the error
|
||||
raise
|
||||
return "pg_search"
|
||||
elif text_search_extension == "native":
|
||||
return "native"
|
||||
elif text_search_extension == "pgroonga":
|
||||
# ensure_text_search_extension() at runtime converts to pgroonga.
|
||||
# Treat as native here so the initial schema still creates valid columns.
|
||||
return "native"
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Invalid HINDSIGHT_API_TEXT_SEARCH_EXTENSION: {text_search_extension}. Must be 'native', 'vchord', or 'pg_textsearch'"
|
||||
f"Invalid HINDSIGHT_API_TEXT_SEARCH_EXTENSION: {text_search_extension}. "
|
||||
"Must be 'native', 'vchord', 'pg_textsearch', 'pgroonga', or 'pg_search'"
|
||||
)
|
||||
|
||||
|
||||
def _pg_search_tokenizer() -> str:
|
||||
return normalize_pg_search_tokenizer(os.getenv(PG_SEARCH_TOKENIZER_ENV))
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
"""Upgrade schema - create all tables from scratch."""
|
||||
|
||||
@@ -284,8 +315,9 @@ def _pg_upgrade() -> None:
|
||||
ALTER TABLE memory_units
|
||||
ADD COLUMN search_vector bm25_catalog.bm25vector
|
||||
""")
|
||||
elif text_search_ext == "pg_textsearch":
|
||||
# Timescale pg_textsearch: dummy TEXT column for consistency (indexes operate on base columns directly)
|
||||
elif text_search_ext in ("pg_textsearch", "pg_search"):
|
||||
# Timescale pg_textsearch / ParadeDB pg_search: dummy TEXT column for
|
||||
# consistency (indexes operate on base columns directly).
|
||||
op.execute("""
|
||||
ALTER TABLE memory_units
|
||||
ADD COLUMN search_vector TEXT
|
||||
@@ -350,6 +382,17 @@ def _pg_upgrade() -> None:
|
||||
USING bm25(text)
|
||||
WITH (text_config='english')
|
||||
""")
|
||||
elif text_search_ext == "pg_search":
|
||||
# ParadeDB pg_search BM25 index on (id, text, context). The key_field
|
||||
# reloption is required and must match the table's primary key column.
|
||||
bm25_cols = pg_search_bm25_columns("id", ("text", "context"), _pg_search_tokenizer())
|
||||
op.execute(
|
||||
"""
|
||||
CREATE INDEX idx_memory_units_text_search ON memory_units
|
||||
USING bm25 ({bm25_cols})
|
||||
WITH (key_field='id')
|
||||
""".format(bm25_cols=bm25_cols)
|
||||
)
|
||||
else: # native
|
||||
# Native PostgreSQL GIN index
|
||||
op.execute("""
|
||||
|
||||
+29
@@ -7,6 +7,7 @@ the stored fact text.
|
||||
- vchord: text_signals included in tokenize() at insert time
|
||||
- native: search_vector GENERATED column regenerated to include text_signals
|
||||
- pg_textsearch: no change (index only supports a single base column)
|
||||
- pg_search: BM25 index dropped and recreated to include text_signals
|
||||
|
||||
Revision ID: a2b3c4d5e6f7
|
||||
Revises: z1u2v3w4x5y6
|
||||
@@ -18,6 +19,11 @@ from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api._pg_search import (
|
||||
PG_SEARCH_TOKENIZER_ENV,
|
||||
normalize_pg_search_tokenizer,
|
||||
pg_search_bm25_columns,
|
||||
)
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "a2b3c4d5e6f7"
|
||||
@@ -35,6 +41,10 @@ def _detect_text_search_extension() -> str:
|
||||
return os.getenv("HINDSIGHT_API_TEXT_SEARCH_EXTENSION", "native").lower()
|
||||
|
||||
|
||||
def _pg_search_tokenizer() -> str:
|
||||
return normalize_pg_search_tokenizer(os.getenv(PG_SEARCH_TOKENIZER_ENV))
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
table = f"{schema}memory_units"
|
||||
@@ -62,6 +72,16 @@ def _pg_upgrade() -> None:
|
||||
CREATE INDEX IF NOT EXISTS idx_memory_units_text_search
|
||||
ON {table} USING gin(search_vector)
|
||||
""")
|
||||
elif text_search_ext == "pg_search":
|
||||
# ParadeDB pg_search: drop the existing BM25 index and recreate it
|
||||
# to include text_signals alongside text and context.
|
||||
bm25_cols = pg_search_bm25_columns("id", ("text", "context", "text_signals"), _pg_search_tokenizer())
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema}idx_memory_units_text_search")
|
||||
op.execute(f"""
|
||||
CREATE INDEX idx_memory_units_text_search ON {table}
|
||||
USING bm25 ({bm25_cols})
|
||||
WITH (key_field='id')
|
||||
""")
|
||||
|
||||
# vchord: tokenize() call in fact_storage.py is updated to include text_signals at insert time
|
||||
# pg_textsearch: no change — index operates on the base `text` column only
|
||||
@@ -86,6 +106,15 @@ def _pg_downgrade() -> None:
|
||||
CREATE INDEX idx_memory_units_text_search
|
||||
ON {table} USING gin(search_vector)
|
||||
""")
|
||||
elif text_search_ext == "pg_search":
|
||||
# Restore the original (id, text, context) BM25 index without text_signals.
|
||||
bm25_cols = pg_search_bm25_columns("id", ("text", "context"), _pg_search_tokenizer())
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema}idx_memory_units_text_search")
|
||||
op.execute(f"""
|
||||
CREATE INDEX idx_memory_units_text_search ON {table}
|
||||
USING bm25 ({bm25_cols})
|
||||
WITH (key_field='id')
|
||||
""")
|
||||
|
||||
op.execute(f"ALTER TABLE {table} DROP COLUMN IF EXISTS text_signals")
|
||||
|
||||
|
||||
+106
@@ -0,0 +1,106 @@
|
||||
"""Add graph_maintenance_queue table
|
||||
|
||||
Queue of memory_units whose outgoing temporal/semantic links lost a
|
||||
neighbour to a delete. Drained by the async graph_maintenance worker,
|
||||
which tops the unit's links back up using the same probes retain runs.
|
||||
|
||||
The queue only targets the link-recompute pass. The worker also runs
|
||||
bank-wide sweeps (orphan-entity prune, stale-cooccurrence prune) on each
|
||||
invocation; those don't need per-target queueing.
|
||||
|
||||
Revision ID: b5a4c3e2f1d8
|
||||
Revises: e9b2c7d1f3a4
|
||||
Create Date: 2026-05-27
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "b5a4c3e2f1d8"
|
||||
down_revision: str | Sequence[str] | None = "e9b2c7d1f3a4"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _pg_schema_prefix() -> str:
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
schema = _pg_schema_prefix()
|
||||
# Composite PK gives us natural ON CONFLICT DO NOTHING dedup when the same
|
||||
# unit is enqueued from overlapping deletes. No FK to memory_units: if the
|
||||
# unit is deleted between enqueue and drain, the worker observes it's gone
|
||||
# and skips — a cascade would erase the work order, but that work has
|
||||
# already been satisfied (no surviving row to maintain).
|
||||
op.execute(
|
||||
f"""
|
||||
CREATE TABLE IF NOT EXISTS {schema}graph_maintenance_queue (
|
||||
bank_id TEXT NOT NULL,
|
||||
unit_id UUID NOT NULL,
|
||||
enqueued_at TIMESTAMPTZ NOT NULL DEFAULT now(),
|
||||
PRIMARY KEY (bank_id, unit_id)
|
||||
)
|
||||
"""
|
||||
)
|
||||
op.execute(
|
||||
f"""
|
||||
CREATE INDEX IF NOT EXISTS idx_graph_maintenance_queue_bank_enqueued
|
||||
ON {schema}graph_maintenance_queue (bank_id, enqueued_at)
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
schema = _pg_schema_prefix()
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema}idx_graph_maintenance_queue_bank_enqueued")
|
||||
op.execute(f"DROP TABLE IF EXISTS {schema}graph_maintenance_queue")
|
||||
|
||||
|
||||
def _oracle_execute_ignoring_955(sql: str) -> None:
|
||||
"""Run a CREATE statement and swallow ORA-00955 (object already exists).
|
||||
|
||||
Mirrors the helper in the Oracle baseline migration so reruns stay safe
|
||||
on a database where the table was created by an earlier partial run.
|
||||
"""
|
||||
block = (
|
||||
"BEGIN "
|
||||
"EXECUTE IMMEDIATE :stmt; "
|
||||
"EXCEPTION WHEN OTHERS THEN "
|
||||
"IF SQLCODE = -955 THEN NULL; ELSE RAISE; END IF; "
|
||||
"END;"
|
||||
)
|
||||
op.get_bind().exec_driver_sql(block, {"stmt": sql.strip()})
|
||||
|
||||
|
||||
def _oracle_upgrade() -> None:
|
||||
_oracle_execute_ignoring_955(
|
||||
"""
|
||||
CREATE TABLE graph_maintenance_queue (
|
||||
bank_id VARCHAR2(256) NOT NULL,
|
||||
unit_id RAW(16) NOT NULL,
|
||||
enqueued_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,
|
||||
CONSTRAINT pk_graph_maintenance_queue PRIMARY KEY (bank_id, unit_id)
|
||||
)
|
||||
"""
|
||||
)
|
||||
_oracle_execute_ignoring_955(
|
||||
"CREATE INDEX idx_graph_maintenance_queue_bank_enqueued ON graph_maintenance_queue (bank_id, enqueued_at)"
|
||||
)
|
||||
|
||||
|
||||
def _oracle_downgrade() -> None:
|
||||
op.execute("DROP INDEX idx_graph_maintenance_queue_bank_enqueued")
|
||||
op.execute("DROP TABLE graph_maintenance_queue")
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade, oracle=_oracle_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade, oracle=_oracle_downgrade)
|
||||
+132
@@ -0,0 +1,132 @@
|
||||
"""Drop indexes that are unused or redundant with composite indexes.
|
||||
|
||||
Code audit identified the following indexes as either dead (no code path
|
||||
exercises them) or fully covered by composite indexes the planner already
|
||||
prefers:
|
||||
|
||||
memory_links:
|
||||
1. idx_memory_links_entity_covering — entity co-occurrence expansion was
|
||||
rewritten to traverse unit_entities instead of memory_links, so no code
|
||||
path filters memory_links on (link_type = 'entity').
|
||||
2. idx_memory_links_from_unit — redundant. idx_memory_links_from_type_weight
|
||||
(from_unit_id, link_type, weight DESC) leads with the same column and
|
||||
answers every from_unit_id = X query.
|
||||
3. idx_memory_links_to_unit — redundant. idx_memory_links_to_type_weight
|
||||
(to_unit_id, link_type, weight DESC) leads with the same column.
|
||||
4. idx_memory_links_link_type — no application query filters on link_type
|
||||
alone; the composite indexes above serve every (from/to + link_type)
|
||||
predicate.
|
||||
|
||||
entities:
|
||||
5. idx_entities_canonical_name — superseded by
|
||||
entities_canonical_name_lower_trgm_idx (case-insensitive lookups).
|
||||
6. entities_canonical_name_trgm_idx — superseded by the lowercase variant
|
||||
in migration 2eee35aa3cfc, but the original was never dropped on schemas
|
||||
that ran the prior migration.
|
||||
|
||||
documents:
|
||||
7. idx_documents_retain_params — GIN index on retain_params JSONB; no query
|
||||
uses jsonb containment on this column.
|
||||
8. idx_documents_content_hash — content-hash lookups happen on the chunks
|
||||
table (chunks.content_hash, indexed separately).
|
||||
|
||||
unit_entities:
|
||||
9. idx_unit_entities_entity — defensive drop. Migration h3i4j5k6l7m8 already
|
||||
issues DROP INDEX IF EXISTS for this; this re-runs the drop idempotently
|
||||
to cover any schema that missed the previous migration.
|
||||
|
||||
All drops use CONCURRENTLY + IF EXISTS so they neither block writers nor
|
||||
fail on schemas where the index is already gone.
|
||||
|
||||
Revision ID: e1b2c3d4f5a6
|
||||
Revises: p4q5r6s7t8u9
|
||||
Create Date: 2026-05-26
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "e1b2c3d4f5a6"
|
||||
down_revision: str | Sequence[str] | None = "p4q5r6s7t8u9"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
_PG_INDEXES_TO_DROP: tuple[str, ...] = (
|
||||
"idx_memory_links_entity_covering",
|
||||
"idx_memory_links_from_unit",
|
||||
"idx_memory_links_to_unit",
|
||||
"idx_memory_links_link_type",
|
||||
"idx_entities_canonical_name",
|
||||
"entities_canonical_name_trgm_idx",
|
||||
"idx_documents_retain_params",
|
||||
"idx_documents_content_hash",
|
||||
"idx_unit_entities_entity",
|
||||
)
|
||||
|
||||
|
||||
def _schema_prefix() -> str:
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
schema = _schema_prefix()
|
||||
# DROP INDEX CONCURRENTLY cannot run inside a transaction block; commit
|
||||
# the Alembic transaction and issue each statement in its own implicit
|
||||
# autocommit transaction. IF EXISTS makes each statement idempotent
|
||||
# across schemas that already dropped (or never had) the index.
|
||||
for index_name in _PG_INDEXES_TO_DROP:
|
||||
op.execute("COMMIT")
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}{index_name}")
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
schema = _schema_prefix()
|
||||
|
||||
# Recreate the dropped indexes in the same shape the prior migrations used,
|
||||
# so a downgrade leaves the schema in the state the previous head expected.
|
||||
op.execute("COMMIT")
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_links_entity_covering "
|
||||
f"ON {schema}memory_links(from_unit_id) "
|
||||
f"INCLUDE (to_unit_id, entity_id) "
|
||||
f"WHERE link_type = 'entity'"
|
||||
)
|
||||
op.execute("COMMIT")
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_links_from_unit ON {schema}memory_links(from_unit_id)"
|
||||
)
|
||||
op.execute("COMMIT")
|
||||
op.execute(f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_links_to_unit ON {schema}memory_links(to_unit_id)")
|
||||
op.execute("COMMIT")
|
||||
op.execute(f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_links_link_type ON {schema}memory_links(link_type)")
|
||||
op.execute("COMMIT")
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_entities_canonical_name ON {schema}entities(canonical_name)"
|
||||
)
|
||||
op.execute("COMMIT")
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS entities_canonical_name_trgm_idx "
|
||||
f"ON {schema}entities USING GIN (canonical_name gin_trgm_ops)"
|
||||
)
|
||||
op.execute("COMMIT")
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_documents_retain_params "
|
||||
f"ON {schema}documents USING GIN (retain_params)"
|
||||
)
|
||||
op.execute("COMMIT")
|
||||
op.execute(f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_documents_content_hash ON {schema}documents(content_hash)")
|
||||
op.execute("COMMIT")
|
||||
op.execute(f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_unit_entities_entity ON {schema}unit_entities(entity_id)")
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade)
|
||||
+89
@@ -0,0 +1,89 @@
|
||||
"""Drop materialized entity rows from memory_links.
|
||||
|
||||
Entity edges are no longer stored in ``memory_links``. The /graph endpoint
|
||||
derives them on demand from ``unit_entities``, and recall already used the
|
||||
``unit_entities`` self-join. Storing entity rows duplicated state we never
|
||||
read from the link table — on a 10k-unit benchmark bank, entity rows were
|
||||
53% of all link rows (~190 MB after indexes) and recall never touched them.
|
||||
|
||||
This migration deletes ``memory_links`` rows with ``link_type = 'entity'``.
|
||||
``idx_memory_links_entity_covering`` was already dropped by migration
|
||||
``e1b2c3d4f5a6``; we still issue ``DROP INDEX IF EXISTS`` defensively in case
|
||||
this migration runs against an older snapshot that predates that one.
|
||||
|
||||
Revision ID: e9b2c7d1f3a4
|
||||
Revises: e1b2c3d4f5a6
|
||||
Create Date: 2026-05-26
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "e9b2c7d1f3a4"
|
||||
down_revision: str | Sequence[str] | None = "e1b2c3d4f5a6"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _pg_schema_prefix() -> str:
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
schema = _pg_schema_prefix()
|
||||
|
||||
# Drop the partial covering index first so the bulk DELETE doesn't churn it.
|
||||
# CREATE/DROP INDEX CONCURRENTLY must run outside a transaction block.
|
||||
op.execute("COMMIT")
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_links_entity_covering")
|
||||
|
||||
# Delete entity rows. Chunked to keep individual transactions small on
|
||||
# large banks (the perf-medium bench had ~345k entity rows; production
|
||||
# banks can be much larger).
|
||||
op.execute(
|
||||
f"""
|
||||
DO $$
|
||||
DECLARE
|
||||
deleted INTEGER;
|
||||
BEGIN
|
||||
LOOP
|
||||
DELETE FROM {schema}memory_links
|
||||
WHERE ctid IN (
|
||||
SELECT ctid FROM {schema}memory_links
|
||||
WHERE link_type = 'entity'
|
||||
LIMIT 50000
|
||||
);
|
||||
GET DIAGNOSTICS deleted = ROW_COUNT;
|
||||
EXIT WHEN deleted = 0;
|
||||
COMMIT;
|
||||
END LOOP;
|
||||
END$$;
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
# Cannot reconstruct deleted entity links — the writer was path-dependent
|
||||
# on retain order. New retains will not produce entity rows either, so the
|
||||
# partial index would stay empty. Leave both no-op.
|
||||
pass
|
||||
|
||||
|
||||
def _oracle_upgrade() -> None:
|
||||
op.execute("DELETE FROM memory_links WHERE link_type = 'entity'")
|
||||
|
||||
|
||||
def _oracle_downgrade() -> None:
|
||||
pass
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade, oracle=_oracle_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade, oracle=_oracle_downgrade)
|
||||
+57
-3
@@ -16,6 +16,11 @@ from collections.abc import Sequence
|
||||
from alembic import context, op
|
||||
from sqlalchemy import text
|
||||
|
||||
from hindsight_api._pg_search import (
|
||||
PG_SEARCH_TOKENIZER_ENV,
|
||||
normalize_pg_search_tokenizer,
|
||||
pg_search_bm25_columns,
|
||||
)
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
@@ -95,9 +100,15 @@ def _vector_index_using_clause(ext: str) -> str:
|
||||
|
||||
def _detect_text_search_extension() -> str:
|
||||
"""
|
||||
Detect or validate text search extension: 'native', 'vchord', or 'pg_textsearch'.
|
||||
Respects HINDSIGHT_API_TEXT_SEARCH_EXTENSION env var.
|
||||
Detect or validate text search extension: 'native', 'vchord', 'pg_textsearch',
|
||||
'pgroonga', or 'pg_search'. Respects HINDSIGHT_API_TEXT_SEARCH_EXTENSION env var.
|
||||
Creates the extension if needed.
|
||||
|
||||
pgroonga is treated as native here so this migration still creates valid
|
||||
tsvector columns; ensure_text_search_extension() at startup converts the
|
||||
reflections table (renamed from pinned_reflections in p1k2l3m4n5o6) to
|
||||
pgroonga structures. The learnings table is dropped in p1k2l3m4n5o6 so its
|
||||
transient native-style column never reaches steady state.
|
||||
"""
|
||||
text_search_extension = os.getenv("HINDSIGHT_API_TEXT_SEARCH_EXTENSION", "native").lower()
|
||||
|
||||
@@ -125,14 +136,33 @@ def _detect_text_search_extension() -> str:
|
||||
# Extension truly doesn't exist - re-raise the error
|
||||
raise
|
||||
return "pg_textsearch"
|
||||
elif text_search_extension == "pg_search":
|
||||
# ParadeDB pg_search — true BM25 over base columns, Citus-compatible.
|
||||
try:
|
||||
op.execute("CREATE EXTENSION IF NOT EXISTS pg_search CASCADE")
|
||||
except Exception:
|
||||
conn = op.get_bind()
|
||||
result = conn.execute(text("SELECT 1 FROM pg_extension WHERE extname = 'pg_search'")).fetchone()
|
||||
if not result:
|
||||
raise
|
||||
return "pg_search"
|
||||
elif text_search_extension == "native":
|
||||
return "native"
|
||||
elif text_search_extension == "pgroonga":
|
||||
# Treat as native here; ensure_text_search_extension() converts the
|
||||
# reflections table to pgroonga structures at runtime.
|
||||
return "native"
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Invalid HINDSIGHT_API_TEXT_SEARCH_EXTENSION: {text_search_extension}. Must be 'native', 'vchord', or 'pg_textsearch'"
|
||||
f"Invalid HINDSIGHT_API_TEXT_SEARCH_EXTENSION: {text_search_extension}. "
|
||||
"Must be 'native', 'vchord', 'pg_textsearch', 'pgroonga', or 'pg_search'"
|
||||
)
|
||||
|
||||
|
||||
def _pg_search_tokenizer() -> str:
|
||||
return normalize_pg_search_tokenizer(os.getenv(PG_SEARCH_TOKENIZER_ENV))
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
"""Create learnings and pinned_reflections tables."""
|
||||
schema = _get_schema_prefix()
|
||||
@@ -200,6 +230,18 @@ def _pg_upgrade() -> None:
|
||||
CREATE INDEX idx_learnings_text_search ON {schema}learnings
|
||||
USING bm25(text) WITH (text_config='english')
|
||||
""")
|
||||
elif text_search_ext == "pg_search":
|
||||
# ParadeDB pg_search: dummy TEXT column; BM25 index is built directly over (id, text)
|
||||
# with key_field='id' (matches the table's primary key).
|
||||
bm25_cols = pg_search_bm25_columns("id", ("text",), _pg_search_tokenizer())
|
||||
op.execute(f"""
|
||||
ALTER TABLE {schema}learnings ADD COLUMN search_vector TEXT
|
||||
""")
|
||||
op.execute(f"""
|
||||
CREATE INDEX idx_learnings_text_search ON {schema}learnings
|
||||
USING bm25 ({bm25_cols})
|
||||
WITH (key_field='id')
|
||||
""")
|
||||
else: # native
|
||||
# Native PostgreSQL: tsvector with automatic generation
|
||||
op.execute(f"""
|
||||
@@ -264,6 +306,18 @@ def _pg_upgrade() -> None:
|
||||
USING bm25(content)
|
||||
WITH (text_config='english')
|
||||
""")
|
||||
elif text_search_ext == "pg_search":
|
||||
# ParadeDB pg_search: dummy TEXT column; BM25 index over (id, name, content)
|
||||
# with key_field='id'.
|
||||
bm25_cols = pg_search_bm25_columns("id", ("name", "content"), _pg_search_tokenizer())
|
||||
op.execute(f"""
|
||||
ALTER TABLE {schema}pinned_reflections ADD COLUMN search_vector TEXT
|
||||
""")
|
||||
op.execute(f"""
|
||||
CREATE INDEX idx_pinned_reflections_text_search ON {schema}pinned_reflections
|
||||
USING bm25 ({bm25_cols})
|
||||
WITH (key_field='id')
|
||||
""")
|
||||
else: # native
|
||||
# Native PostgreSQL: tsvector with automatic generation
|
||||
op.execute(f"""
|
||||
|
||||
+170
@@ -0,0 +1,170 @@
|
||||
"""Drop GENERATED expression on tsvector search_vector columns.
|
||||
|
||||
The search_vector tsvector column was originally GENERATED ALWAYS with a
|
||||
hardcoded ``to_tsvector('english', ...)`` expression. To support configurable
|
||||
``HINDSIGHT_API_TEXT_SEARCH_EXTENSION_NATIVE_LANGUAGE``, we convert it to a
|
||||
regular tsvector column that the application populates at INSERT time via
|
||||
``to_tsvector($lang, ...)``.
|
||||
|
||||
Existing rows retain their English-derived lexemes — switching the configured
|
||||
language only affects newly-written rows. Users who need to backfill existing
|
||||
rows in a different language can run an admin UPDATE after this migration.
|
||||
|
||||
Only the ``native`` text-search backend is affected. ``vchord``, ``pg_textsearch``,
|
||||
and ``pgroonga`` use other column types or no column at all.
|
||||
|
||||
Revision ID: p4q5r6s7t8u9
|
||||
Revises: 86f7a033d372
|
||||
Create Date: 2026-05-08
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
from dataclasses import dataclass
|
||||
|
||||
from alembic import context, op
|
||||
from sqlalchemy import Connection, text
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "p4q5r6s7t8u9"
|
||||
down_revision: str | Sequence[str] | None = "86f7a033d372"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _TsvectorTableSpec:
|
||||
"""Native-backend tsvector table targeted by this migration.
|
||||
|
||||
``upgrade`` is a one-way DROP EXPRESSION; ``downgrade`` re-attaches the
|
||||
original GENERATED expression so the schema returns to the state created
|
||||
by the initial migration (and a2b3c4d5e6f7_add_text_signals_column for
|
||||
memory_units).
|
||||
"""
|
||||
|
||||
table: str
|
||||
generated_expression: str
|
||||
|
||||
|
||||
# Tables that may have a GENERATED tsvector ``search_vector`` column under the
|
||||
# native backend. Note: the ``learnings`` table was dropped in
|
||||
# p1k2l3m4n5o6_new_knowledge_architecture and ``pinned_reflections`` was renamed
|
||||
# to ``reflections`` in the same migration.
|
||||
_NATIVE_TSVECTOR_TABLES: tuple[_TsvectorTableSpec, ...] = (
|
||||
_TsvectorTableSpec(
|
||||
table="memory_units",
|
||||
generated_expression=(
|
||||
"to_tsvector('english', COALESCE(text, '') || ' ' || "
|
||||
"COALESCE(context, '') || ' ' || COALESCE(text_signals, ''))"
|
||||
),
|
||||
),
|
||||
_TsvectorTableSpec(
|
||||
table="reflections",
|
||||
generated_expression="to_tsvector('english', COALESCE(name, '') || ' ' || content)",
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _schema_prefix() -> str:
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def _is_generated_tsvector(conn: Connection, schema: str, table: str) -> bool:
|
||||
"""Return True iff ``schema.table.search_vector`` is a GENERATED tsvector column."""
|
||||
row = conn.execute(
|
||||
text(
|
||||
"""
|
||||
SELECT is_generated, udt_name
|
||||
FROM information_schema.columns
|
||||
WHERE table_schema = :schema
|
||||
AND table_name = :table
|
||||
AND column_name = 'search_vector'
|
||||
"""
|
||||
),
|
||||
{"schema": schema, "table": table},
|
||||
).fetchone()
|
||||
if not row:
|
||||
return False
|
||||
is_generated, udt_name = row[0], row[1]
|
||||
return is_generated == "ALWAYS" and udt_name == "tsvector"
|
||||
|
||||
|
||||
def _is_regular_tsvector(conn: Connection, schema: str, table: str) -> bool:
|
||||
"""Return True iff ``schema.table.search_vector`` is a non-generated tsvector column."""
|
||||
row = conn.execute(
|
||||
text(
|
||||
"""
|
||||
SELECT is_generated, udt_name
|
||||
FROM information_schema.columns
|
||||
WHERE table_schema = :schema
|
||||
AND table_name = :table
|
||||
AND column_name = 'search_vector'
|
||||
"""
|
||||
),
|
||||
{"schema": schema, "table": table},
|
||||
).fetchone()
|
||||
if not row:
|
||||
return False
|
||||
is_generated, udt_name = row[0], row[1]
|
||||
return udt_name == "tsvector" and is_generated != "ALWAYS"
|
||||
|
||||
|
||||
def _table_exists(conn: Connection, schema: str, table: str) -> bool:
|
||||
return bool(
|
||||
conn.execute(
|
||||
text(
|
||||
"""
|
||||
SELECT 1 FROM information_schema.tables
|
||||
WHERE table_schema = :schema AND table_name = :table
|
||||
"""
|
||||
),
|
||||
{"schema": schema, "table": table},
|
||||
).fetchone()
|
||||
)
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
schema_prefix = _schema_prefix()
|
||||
schema_name = (context.config.get_main_option("target_schema") or "public").strip('"')
|
||||
conn = op.get_bind()
|
||||
|
||||
for spec in _NATIVE_TSVECTOR_TABLES:
|
||||
if not _table_exists(conn, schema_name, spec.table):
|
||||
continue
|
||||
if not _is_generated_tsvector(conn, schema_name, spec.table):
|
||||
# Either the column doesn't exist (non-native backend) or it's
|
||||
# already a regular tsvector — nothing to do.
|
||||
continue
|
||||
op.execute(f"ALTER TABLE {schema_prefix}{spec.table} ALTER COLUMN search_vector DROP EXPRESSION")
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
schema_prefix = _schema_prefix()
|
||||
schema_name = (context.config.get_main_option("target_schema") or "public").strip('"')
|
||||
conn = op.get_bind()
|
||||
|
||||
for spec in _NATIVE_TSVECTOR_TABLES:
|
||||
if not _table_exists(conn, schema_name, spec.table):
|
||||
continue
|
||||
# Only restore the GENERATED expression if a non-generated tsvector
|
||||
# column exists — otherwise the table is on a different backend.
|
||||
if not _is_regular_tsvector(conn, schema_name, spec.table):
|
||||
continue
|
||||
# Drop and recreate to re-attach the GENERATED expression. Index will be
|
||||
# recreated by re-running ensure_text_search_extension on next startup.
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema_prefix}idx_{spec.table}_text_search")
|
||||
op.execute(f"ALTER TABLE {schema_prefix}{spec.table} DROP COLUMN search_vector")
|
||||
op.execute(
|
||||
f"ALTER TABLE {schema_prefix}{spec.table} "
|
||||
f"ADD COLUMN search_vector tsvector GENERATED ALWAYS AS ({spec.generated_expression}) STORED"
|
||||
)
|
||||
op.execute(f"CREATE INDEX idx_{spec.table}_text_search ON {schema_prefix}{spec.table} USING gin(search_vector)")
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade)
|
||||
@@ -15,6 +15,7 @@ from datetime import datetime, timezone
|
||||
from typing import Any, Literal
|
||||
|
||||
from fastapi import Depends, FastAPI, File, Form, Header, HTTPException, Query, Request, UploadFile
|
||||
from fastapi.middleware.gzip import GZipMiddleware
|
||||
|
||||
from hindsight_api.engine.audit import AuditEntry, AuditLogger
|
||||
from hindsight_api.extensions import AuthenticationError
|
||||
@@ -466,6 +467,13 @@ class MemoryItem(BaseModel):
|
||||
description="Optional tags for visibility scoping. Memories with tags can be filtered during recall.",
|
||||
)
|
||||
|
||||
@field_validator("content")
|
||||
@classmethod
|
||||
def validate_content(cls, v: str) -> str:
|
||||
if not v.strip():
|
||||
raise ValueError("content cannot be empty")
|
||||
return v
|
||||
|
||||
@field_validator("tags", mode="before")
|
||||
@classmethod
|
||||
def coerce_tags(cls, v):
|
||||
@@ -2180,6 +2188,19 @@ class OperationResponse(BaseModel):
|
||||
)
|
||||
|
||||
|
||||
class ConsolidationRequest(BaseModel):
|
||||
"""Request model for consolidation trigger endpoint."""
|
||||
|
||||
observation_scopes: list[list[str]] | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Optional list of tag scopes to consolidate. Each scope is a list of tags. "
|
||||
"Only unconsolidated memories whose tags contain all tags in at least one scope "
|
||||
"will be processed. If omitted, all unconsolidated memories are processed."
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class ConsolidationResponse(BaseModel):
|
||||
"""Response model for consolidation trigger endpoint."""
|
||||
|
||||
@@ -2721,6 +2742,8 @@ def create_app(
|
||||
app.state.memory = memory
|
||||
app.state.audit_logger = memory.audit_logger
|
||||
|
||||
app.add_middleware(GZipMiddleware, minimum_size=1024)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Patch OpenAPI schema: align ValidationError with Pydantic v2 error format
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -2885,6 +2908,57 @@ def _register_routes(app: FastAPI):
|
||||
api_key = authorization.strip()
|
||||
return RequestContext(api_key=api_key)
|
||||
|
||||
def precheck_for(operation: str):
|
||||
"""
|
||||
Build a FastAPI dependency that runs ``OperationValidator.precheck``.
|
||||
|
||||
FastAPI resolves dependencies before deserialising the route's body
|
||||
parameter. Wiring this dependency on the billable POST routes lets
|
||||
an extension reject a request — e.g. with HTTP 402 when a tenant's
|
||||
balance is exhausted — without the request body ever being read or
|
||||
materialised in memory.
|
||||
|
||||
The dependency intentionally:
|
||||
- authenticates the tenant (so ``request_context.tenant_id`` is
|
||||
resolved before the precheck runs);
|
||||
- falls through silently when no validator is configured or the
|
||||
validator's default no-op precheck is in effect;
|
||||
- converts a rejection ``ValidationResult`` into the corresponding
|
||||
``HTTPException`` directly (the per-route ``OperationValidationError``
|
||||
catch blocks don't see exceptions raised in dependencies, so we
|
||||
translate here instead of relying on each handler's try/except).
|
||||
|
||||
Args:
|
||||
operation: Short identifier for the route, e.g. ``"retain"``.
|
||||
|
||||
Returns:
|
||||
A FastAPI dependency callable suitable for ``Depends(...)``.
|
||||
"""
|
||||
|
||||
async def _precheck_dep(
|
||||
bank_id: str,
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
) -> None:
|
||||
validator = getattr(app.state.memory, "_operation_validator", None)
|
||||
if validator is None:
|
||||
return
|
||||
from hindsight_api.extensions import PrecheckContext
|
||||
|
||||
await app.state.memory._authenticate_tenant(request_context)
|
||||
ctx = PrecheckContext(
|
||||
operation=operation,
|
||||
bank_id=bank_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
result = await validator.precheck(ctx)
|
||||
if not result.allowed:
|
||||
raise HTTPException(
|
||||
status_code=result.status_code,
|
||||
detail=result.reason or "Operation not allowed",
|
||||
)
|
||||
|
||||
return _precheck_dep
|
||||
|
||||
# Global exception handler for authentication errors
|
||||
@app.exception_handler(AuthenticationError)
|
||||
async def authentication_error_handler(request, exc: AuthenticationError):
|
||||
@@ -3142,7 +3216,10 @@ def _register_routes(app: FastAPI):
|
||||
)
|
||||
@audited("recall")
|
||||
async def api_recall(
|
||||
bank_id: str, request: RecallRequest, request_context: RequestContext = Depends(get_request_context)
|
||||
bank_id: str,
|
||||
request: RecallRequest,
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
_precheck: None = Depends(precheck_for("recall")),
|
||||
):
|
||||
"""Run a recall and return results with trace."""
|
||||
import time
|
||||
@@ -3330,7 +3407,10 @@ def _register_routes(app: FastAPI):
|
||||
)
|
||||
@audited("reflect")
|
||||
async def api_reflect(
|
||||
bank_id: str, request: ReflectRequest, request_context: RequestContext = Depends(get_request_context)
|
||||
bank_id: str,
|
||||
request: ReflectRequest,
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
_precheck: None = Depends(precheck_for("reflect")),
|
||||
):
|
||||
metrics = get_metrics_collector()
|
||||
|
||||
@@ -3828,6 +3908,7 @@ def _register_routes(app: FastAPI):
|
||||
bank_id: str,
|
||||
body: CreateMentalModelRequest,
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
_precheck: None = Depends(precheck_for("mental_model_create")),
|
||||
):
|
||||
"""Create a mental model (async - returns operation_id)."""
|
||||
try:
|
||||
@@ -3876,6 +3957,7 @@ def _register_routes(app: FastAPI):
|
||||
bank_id: str,
|
||||
mental_model_id: str,
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
_precheck: None = Depends(precheck_for("mental_model_refresh")),
|
||||
):
|
||||
"""Refresh a mental model by re-running its source query (async)."""
|
||||
try:
|
||||
@@ -3902,6 +3984,48 @@ def _register_routes(app: FastAPI):
|
||||
)
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
@app.post(
|
||||
"/v1/default/banks/{bank_id}/mental-models/{mental_model_id}/clear",
|
||||
response_model=MentalModelResponse,
|
||||
summary="Clear mental model content",
|
||||
description=(
|
||||
"Clear a mental model's content so the next refresh performs a full re-synthesis. "
|
||||
"This is useful for delta-mode models that have accumulated drift over many "
|
||||
"incremental refreshes. After clearing, call the /refresh endpoint to trigger "
|
||||
"a clean full rebuild."
|
||||
),
|
||||
operation_id="clear_mental_model",
|
||||
tags=["Mental Models"],
|
||||
)
|
||||
@audited("clear_mental_model", request_param=None)
|
||||
async def api_clear_mental_model(
|
||||
bank_id: str,
|
||||
mental_model_id: str,
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
):
|
||||
"""Clear a mental model's content."""
|
||||
try:
|
||||
mental_model = await app.state.memory.clear_mental_model(
|
||||
bank_id=bank_id,
|
||||
mental_model_id=mental_model_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
if mental_model is None:
|
||||
raise HTTPException(status_code=404, detail=f"Mental model '{mental_model_id}' not found")
|
||||
return MentalModelResponse(**mental_model)
|
||||
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 POST /v1/default/banks/{bank_id}/mental-models/{mental_model_id}/clear: {error_detail}"
|
||||
)
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
@app.patch(
|
||||
"/v1/default/banks/{bank_id}/mental-models/{mental_model_id}",
|
||||
response_model=MentalModelResponse,
|
||||
@@ -5395,11 +5519,20 @@ def _register_routes(app: FastAPI):
|
||||
operation_id="trigger_consolidation",
|
||||
tags=["Banks"],
|
||||
)
|
||||
@audited("consolidation", request_param=None)
|
||||
async def api_trigger_consolidation(bank_id: str, request_context: RequestContext = Depends(get_request_context)):
|
||||
@audited("consolidation")
|
||||
async def api_trigger_consolidation(
|
||||
bank_id: str,
|
||||
request: ConsolidationRequest | None = None,
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
):
|
||||
"""Trigger consolidation for a bank (async)."""
|
||||
try:
|
||||
result = await app.state.memory.submit_async_consolidation(bank_id=bank_id, request_context=request_context)
|
||||
observation_scopes = request.observation_scopes if request else None
|
||||
result = await app.state.memory.submit_async_consolidation(
|
||||
bank_id=bank_id,
|
||||
request_context=request_context,
|
||||
observation_scopes=observation_scopes,
|
||||
)
|
||||
return ConsolidationResponse(
|
||||
operation_id=result["operation_id"],
|
||||
deduplicated=result.get("deduplicated", False),
|
||||
@@ -5722,7 +5855,10 @@ def _register_routes(app: FastAPI):
|
||||
)
|
||||
@audited("retain")
|
||||
async def api_retain(
|
||||
bank_id: str, request: RetainRequest, request_context: RequestContext = Depends(get_request_context)
|
||||
bank_id: str,
|
||||
request: RetainRequest,
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
_precheck: None = Depends(precheck_for("retain")),
|
||||
):
|
||||
"""Retain memories with optional async processing."""
|
||||
metrics = get_metrics_collector()
|
||||
@@ -5807,9 +5943,8 @@ def _register_routes(app: FastAPI):
|
||||
strategy=group_strategy,
|
||||
request_context=request_context,
|
||||
return_usage=True,
|
||||
outbox_callback=app.state.memory._build_retain_outbox_callback(
|
||||
outbox_callback_factory=app.state.memory._build_retain_outbox_callback_factory(
|
||||
bank_id=bank_id,
|
||||
contents=contents,
|
||||
operation_id=None,
|
||||
schema=_current_schema.get(),
|
||||
),
|
||||
@@ -5892,6 +6027,7 @@ def _register_routes(app: FastAPI):
|
||||
files: list[UploadFile] = File(..., description="Files to upload and convert"),
|
||||
request: str = Form(..., description="JSON string with FileRetainRequest model"),
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
_precheck: None = Depends(precheck_for("files_retain")),
|
||||
):
|
||||
"""Upload and convert files to memories."""
|
||||
from hindsight_api.config import get_config
|
||||
|
||||
@@ -107,6 +107,7 @@ def create_mcp_server(memory: MemoryEngine, multi_bank: bool = True) -> FastMCP:
|
||||
"update_mental_model",
|
||||
"delete_mental_model",
|
||||
"refresh_mental_model",
|
||||
"clear_mental_model",
|
||||
"list_directives",
|
||||
"create_directive",
|
||||
"delete_directive",
|
||||
|
||||
@@ -7,6 +7,7 @@ All environment variables and their defaults are defined here.
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
import sys
|
||||
from dataclasses import dataclass, field, fields
|
||||
from datetime import datetime, timezone
|
||||
@@ -14,6 +15,7 @@ from typing import Any, Literal
|
||||
|
||||
from dotenv import find_dotenv, load_dotenv
|
||||
|
||||
from ._pg_search import normalize_pg_search_tokenizer
|
||||
from ._vector_index import validate_extension
|
||||
from .utils import mask_network_location
|
||||
|
||||
@@ -200,6 +202,7 @@ 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"
|
||||
ENV_EMBEDDINGS_OPENAI_DIMENSIONS = "HINDSIGHT_API_EMBEDDINGS_OPENAI_DIMENSIONS"
|
||||
|
||||
# Gemini/Vertex AI embeddings configuration
|
||||
ENV_EMBEDDINGS_GEMINI_API_KEY = "HINDSIGHT_API_EMBEDDINGS_GEMINI_API_KEY"
|
||||
@@ -226,6 +229,15 @@ ENV_EMBEDDINGS_OPENROUTER_MODEL = "HINDSIGHT_API_EMBEDDINGS_OPENROUTER_MODEL"
|
||||
ENV_RERANKER_OPENROUTER_API_KEY = "HINDSIGHT_API_RERANKER_OPENROUTER_API_KEY"
|
||||
ENV_RERANKER_OPENROUTER_MODEL = "HINDSIGHT_API_RERANKER_OPENROUTER_MODEL"
|
||||
|
||||
# ZeroEntropy configuration (embeddings)
|
||||
ENV_EMBEDDINGS_ZEROENTROPY_API_KEY = "HINDSIGHT_API_EMBEDDINGS_ZEROENTROPY_API_KEY"
|
||||
ENV_EMBEDDINGS_ZEROENTROPY_MODEL = "HINDSIGHT_API_EMBEDDINGS_ZEROENTROPY_MODEL"
|
||||
ENV_EMBEDDINGS_ZEROENTROPY_BASE_URL = "HINDSIGHT_API_EMBEDDINGS_ZEROENTROPY_BASE_URL"
|
||||
ENV_EMBEDDINGS_ZEROENTROPY_DIMENSIONS = "HINDSIGHT_API_EMBEDDINGS_ZEROENTROPY_DIMENSIONS"
|
||||
ENV_EMBEDDINGS_ZEROENTROPY_ENCODING_FORMAT = "HINDSIGHT_API_EMBEDDINGS_ZEROENTROPY_ENCODING_FORMAT"
|
||||
ENV_EMBEDDINGS_ZEROENTROPY_LATENCY = "HINDSIGHT_API_EMBEDDINGS_ZEROENTROPY_LATENCY"
|
||||
ENV_EMBEDDINGS_ZEROENTROPY_BATCH_SIZE = "HINDSIGHT_API_EMBEDDINGS_ZEROENTROPY_BATCH_SIZE"
|
||||
|
||||
# Deprecated: Legacy shared Cohere API key (for backward compatibility)
|
||||
ENV_COHERE_API_KEY = "HINDSIGHT_API_COHERE_API_KEY"
|
||||
|
||||
@@ -279,6 +291,10 @@ ENV_RERANKER_SILICONFLOW_API_KEY = "HINDSIGHT_API_RERANKER_SILICONFLOW_API_KEY"
|
||||
ENV_RERANKER_SILICONFLOW_MODEL = "HINDSIGHT_API_RERANKER_SILICONFLOW_MODEL"
|
||||
ENV_RERANKER_SILICONFLOW_BASE_URL = "HINDSIGHT_API_RERANKER_SILICONFLOW_BASE_URL"
|
||||
|
||||
# Alibaba Cloud DashScope configuration (reranker only)
|
||||
ENV_RERANKER_ALIBABA_API_KEY = "HINDSIGHT_API_RERANKER_ALIBABA_API_KEY"
|
||||
ENV_RERANKER_ALIBABA_MODEL = "HINDSIGHT_API_RERANKER_ALIBABA_MODEL"
|
||||
|
||||
# Google Discovery Engine reranker configuration
|
||||
ENV_RERANKER_GOOGLE_MODEL = "HINDSIGHT_API_RERANKER_GOOGLE_MODEL"
|
||||
ENV_RERANKER_GOOGLE_PROJECT_ID = "HINDSIGHT_API_RERANKER_GOOGLE_PROJECT_ID"
|
||||
@@ -286,6 +302,9 @@ ENV_RERANKER_GOOGLE_SERVICE_ACCOUNT_KEY = "HINDSIGHT_API_RERANKER_GOOGLE_SERVICE
|
||||
|
||||
ENV_VECTOR_EXTENSION = "HINDSIGHT_API_VECTOR_EXTENSION"
|
||||
ENV_TEXT_SEARCH_EXTENSION = "HINDSIGHT_API_TEXT_SEARCH_EXTENSION"
|
||||
ENV_TEXT_SEARCH_EXTENSION_NATIVE_LANGUAGE = "HINDSIGHT_API_TEXT_SEARCH_EXTENSION_NATIVE_LANGUAGE"
|
||||
ENV_TEXT_SEARCH_EXTENSION_PG_SEARCH_TOKENIZER = "HINDSIGHT_API_TEXT_SEARCH_EXTENSION_PG_SEARCH_TOKENIZER"
|
||||
ENV_LLM_OUTPUT_LANGUAGE = "HINDSIGHT_API_LLM_OUTPUT_LANGUAGE"
|
||||
|
||||
ENV_HOST = "HINDSIGHT_API_HOST"
|
||||
ENV_PORT = "HINDSIGHT_API_PORT"
|
||||
@@ -294,6 +313,7 @@ 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_ACCESS_LOG = "HINDSIGHT_API_ACCESS_LOG"
|
||||
ENV_MCP_ENABLED = "HINDSIGHT_API_MCP_ENABLED"
|
||||
ENV_MCP_ENABLED_TOOLS = "HINDSIGHT_API_MCP_ENABLED_TOOLS"
|
||||
ENV_MCP_STATELESS = "HINDSIGHT_API_MCP_STATELESS"
|
||||
@@ -361,6 +381,7 @@ 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_ENABLE_AUTO_CONSOLIDATION = "HINDSIGHT_API_ENABLE_AUTO_CONSOLIDATION"
|
||||
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"
|
||||
@@ -375,6 +396,7 @@ 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"
|
||||
ENV_ENABLE_MENTAL_MODEL_HISTORY = "HINDSIGHT_API_ENABLE_MENTAL_MODEL_HISTORY"
|
||||
ENV_MENTAL_MODEL_HISTORY_MAX_ENTRIES = "HINDSIGHT_API_MENTAL_MODEL_HISTORY_MAX_ENTRIES"
|
||||
|
||||
# Webhook configuration (global, static - server-level only)
|
||||
ENV_WEBHOOK_URL = "HINDSIGHT_API_WEBHOOK_URL"
|
||||
@@ -422,6 +444,7 @@ WORKER_SLOT_RESERVATION_TYPES: dict[str, tuple[str, int]] = {
|
||||
"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),
|
||||
"graph_maintenance": ("HINDSIGHT_API_WORKER_GRAPH_MAINTENANCE_MAX_SLOTS", 0),
|
||||
}
|
||||
ENV_RETAIN_MAX_CONCURRENT = "HINDSIGHT_API_RETAIN_MAX_CONCURRENT"
|
||||
|
||||
@@ -473,6 +496,7 @@ PROVIDER_DEFAULT_MODELS = {
|
||||
"zai": "glm-4.5-flash",
|
||||
"opencode-go": "deepseek-v4-flash",
|
||||
"ollama": "gemma3:12b",
|
||||
"ollama-cloud": "gemma3:12b",
|
||||
"llamacpp": "gemma-4-e2b-it",
|
||||
"lmstudio": "local-model",
|
||||
"vertexai": "google/gemini-2.5-flash-lite",
|
||||
@@ -543,18 +567,39 @@ DEFAULT_RERANKER_COHERE_MODEL = "rerank-english-v3.0"
|
||||
DEFAULT_EMBEDDINGS_OPENROUTER_MODEL = "perplexity/pplx-embed-v1-0.6b"
|
||||
DEFAULT_RERANKER_OPENROUTER_MODEL = "cohere/rerank-v3.5"
|
||||
|
||||
# ZeroEntropy defaults
|
||||
DEFAULT_EMBEDDINGS_ZEROENTROPY_MODEL = "zembed-1"
|
||||
# Shared between embeddings (zembed-1) and reranker (zerank-*) — the host is the same.
|
||||
DEFAULT_ZEROENTROPY_BASE_URL = "https://api.zeroentropy.dev"
|
||||
# ZeroEntropy's API default is 2560, but Hindsight defaults to 1280 so the
|
||||
# provider works with pgvector HNSW's 2000-dimension index limit out of the box.
|
||||
DEFAULT_EMBEDDINGS_ZEROENTROPY_DIMENSIONS = 1280
|
||||
DEFAULT_EMBEDDINGS_ZEROENTROPY_ENCODING_FORMAT = "float"
|
||||
DEFAULT_EMBEDDINGS_ZEROENTROPY_LATENCY = None
|
||||
DEFAULT_EMBEDDINGS_ZEROENTROPY_BATCH_SIZE = 100
|
||||
|
||||
DEFAULT_RERANKER_ZEROENTROPY_MODEL = "zerank-2"
|
||||
|
||||
DEFAULT_RERANKER_SILICONFLOW_MODEL = "BAAI/bge-reranker-v2-m3"
|
||||
DEFAULT_RERANKER_SILICONFLOW_BASE_URL = "https://api.siliconflow.cn/v1"
|
||||
|
||||
DEFAULT_RERANKER_ALIBABA_MODEL = "qwen3-rerank"
|
||||
|
||||
DEFAULT_RERANKER_GOOGLE_MODEL = "semantic-ranker-default-004"
|
||||
|
||||
# Vector extension (pgvector, vchord, pgvectorscale, or AlloyDB ScaNN)
|
||||
DEFAULT_VECTOR_EXTENSION = "pgvector" # Options: "pgvector", "vchord", "pgvectorscale", "scann"
|
||||
|
||||
# Text search extension (native PostgreSQL, vchord BM25, or Timescale pg_textsearch)
|
||||
DEFAULT_TEXT_SEARCH_EXTENSION = "native" # Options: "native", "vchord", "pg_textsearch"
|
||||
# Text search extension (native PostgreSQL, vchord BM25, Timescale pg_textsearch,
|
||||
# pgroonga, or ParadeDB pg_search)
|
||||
DEFAULT_TEXT_SEARCH_EXTENSION = "native" # Options: "native", "vchord", "pg_textsearch", "pgroonga", "pg_search"
|
||||
|
||||
# PostgreSQL text search dictionary used by the native tsvector backend. Only
|
||||
# affects text_search_extension == "native"; other backends use their own
|
||||
# tokenizers (vchord: llmlingua2, pg_textsearch: hardcoded english,
|
||||
# pgroonga: TokenBigram polyglot, pg_search: per-field Tantivy tokenizer).
|
||||
DEFAULT_TEXT_SEARCH_EXTENSION_NATIVE_LANGUAGE = "english"
|
||||
DEFAULT_TEXT_SEARCH_EXTENSION_PG_SEARCH_TOKENIZER = ""
|
||||
|
||||
# LiteLLM defaults
|
||||
DEFAULT_LITELLM_API_BASE = "http://localhost:4000"
|
||||
@@ -573,6 +618,7 @@ DEFAULT_BASE_PATH = "" # Empty string = root path
|
||||
DEFAULT_LOG_LEVEL = "info"
|
||||
DEFAULT_LOG_FORMAT = "text" # Options: "text", "json"
|
||||
DEFAULT_WORKERS = 1
|
||||
DEFAULT_ACCESS_LOG = False
|
||||
DEFAULT_MCP_ENABLED = True
|
||||
DEFAULT_MCP_ENABLED_TOOLS: list[str] | None = None # None = all tools enabled
|
||||
DEFAULT_MCP_STATELESS = False # False = stateful (supports SSE/GET); True = stateless (POST-only)
|
||||
@@ -615,8 +661,15 @@ DEFAULT_FILE_DELETE_AFTER_RETAIN = True # Delete file bytes after retain (saves
|
||||
|
||||
# Observations defaults (consolidated knowledge from facts)
|
||||
DEFAULT_ENABLE_OBSERVATIONS = True # Observations enabled by default
|
||||
DEFAULT_ENABLE_AUTO_CONSOLIDATION = True # Auto-consolidation after retain 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
|
||||
# Each history entry snapshots previous_content + previous_reflect_response. Without
|
||||
# a cap, sustained mental-model refresh load grows the jsonb array unboundedly until
|
||||
# it crosses Postgres's hard 256MB jsonb limit and subsequent UPDATEs fail with
|
||||
# SQLSTATE 54000. 50 keeps the array well under 100MB even with large reflect
|
||||
# responses, while preserving enough recent history for meaningful audit / rollback.
|
||||
DEFAULT_MENTAL_MODEL_HISTORY_MAX_ENTRIES = 50
|
||||
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 = (
|
||||
@@ -792,6 +845,24 @@ def _parse_positive_int(name: str, raw: str | None, default: int) -> int:
|
||||
return parsed
|
||||
|
||||
|
||||
def _parse_optional_positive_int(name: str, raw: str | None) -> int | None:
|
||||
"""Parse an optional env var that must be a positive integer when set."""
|
||||
if raw is None or raw == "":
|
||||
return None
|
||||
return _parse_positive_int(name, raw, 1)
|
||||
|
||||
|
||||
def _parse_optional_choice(name: str, raw: str | None, allowed: frozenset[str]) -> str | None:
|
||||
"""Parse an optional string env var constrained to a small allowlist."""
|
||||
if raw is None or raw == "":
|
||||
return None
|
||||
normalized = raw.lower()
|
||||
if normalized not in allowed:
|
||||
values = ", ".join(sorted(allowed))
|
||||
raise ValueError(f"{name} must be one of {values}, got {raw!r}")
|
||||
return normalized
|
||||
|
||||
|
||||
def _validate_extraction_mode(mode: str) -> str:
|
||||
"""Validate and normalize extraction mode."""
|
||||
mode_lower = mode.lower()
|
||||
@@ -874,7 +945,19 @@ class HindsightConfig:
|
||||
migration_database_url: str | None
|
||||
database_schema: str
|
||||
vector_extension: str # "pgvector", "vchord", "pgvectorscale", or "scann"
|
||||
text_search_extension: str # "native" or "vchord"
|
||||
text_search_extension: str # "native", "vchord", "pg_textsearch", "pgroonga", or "pg_search"
|
||||
# PostgreSQL text search dictionary for the "native" backend (ignored by
|
||||
# other backends). Only the "native" backend reads this field; pgroonga
|
||||
# uses TokenBigram, vchord uses llmlingua2, pg_textsearch hardcodes english,
|
||||
# pg_search uses Tantivy per-field tokenizers.
|
||||
text_search_extension_native_language: str
|
||||
# ParadeDB pg_search tokenizer used when building BM25 indexes. Empty keeps
|
||||
# ParadeDB's default tokenizer.
|
||||
text_search_extension_pg_search_tokenizer: str
|
||||
# When set, every LLM-generated artifact (retain facts, consolidation
|
||||
# observations, reflect responses) is forced into this language regardless
|
||||
# of the source content. Unset preserves source language.
|
||||
llm_output_language: str | None
|
||||
|
||||
# LLM (default, used as fallback for per-operation config)
|
||||
llm_provider: str
|
||||
@@ -1013,6 +1096,8 @@ class HindsightConfig:
|
||||
reranker_siliconflow_api_key: str | None
|
||||
reranker_siliconflow_model: str
|
||||
reranker_siliconflow_base_url: str
|
||||
reranker_alibaba_api_key: str | None
|
||||
reranker_alibaba_model: str
|
||||
reranker_google_model: str
|
||||
reranker_google_project_id: str | None
|
||||
reranker_google_service_account_key: str | None
|
||||
@@ -1080,8 +1165,10 @@ class HindsightConfig:
|
||||
|
||||
# Observations settings (consolidated knowledge from facts)
|
||||
enable_observations: bool
|
||||
enable_auto_consolidation: bool
|
||||
enable_observation_history: bool
|
||||
enable_mental_model_history: bool
|
||||
mental_model_history_max_entries: int
|
||||
consolidation_batch_size: int
|
||||
consolidation_max_memories_per_round: int
|
||||
consolidation_llm_batch_size: int
|
||||
@@ -1179,6 +1266,14 @@ class HindsightConfig:
|
||||
# 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
|
||||
embeddings_openai_dimensions: int | None = None
|
||||
embeddings_zeroentropy_api_key: str | None = None
|
||||
embeddings_zeroentropy_model: str = DEFAULT_EMBEDDINGS_ZEROENTROPY_MODEL
|
||||
embeddings_zeroentropy_base_url: str = DEFAULT_ZEROENTROPY_BASE_URL
|
||||
embeddings_zeroentropy_dimensions: int = DEFAULT_EMBEDDINGS_ZEROENTROPY_DIMENSIONS
|
||||
embeddings_zeroentropy_encoding_format: str = DEFAULT_EMBEDDINGS_ZEROENTROPY_ENCODING_FORMAT
|
||||
embeddings_zeroentropy_batch_size: int = DEFAULT_EMBEDDINGS_ZEROENTROPY_BATCH_SIZE
|
||||
embeddings_zeroentropy_latency: str | None = DEFAULT_EMBEDDINGS_ZEROENTROPY_LATENCY
|
||||
|
||||
# Class-level sets for configuration categorization
|
||||
|
||||
@@ -1202,6 +1297,7 @@ class HindsightConfig:
|
||||
"embeddings_tei_base_url",
|
||||
"reranker_tei_base_url",
|
||||
"reranker_cohere_base_url",
|
||||
"embeddings_zeroentropy_base_url",
|
||||
"reranker_zeroentropy_base_url",
|
||||
"reranker_siliconflow_base_url",
|
||||
# Service Account Keys
|
||||
@@ -1210,6 +1306,7 @@ class HindsightConfig:
|
||||
"reranker_google_service_account_key",
|
||||
# Embeddings API keys
|
||||
"embeddings_gemini_api_key",
|
||||
"embeddings_zeroentropy_api_key",
|
||||
# File storage credentials
|
||||
"file_storage_s3_access_key_id",
|
||||
"file_storage_s3_secret_access_key",
|
||||
@@ -1239,6 +1336,7 @@ class HindsightConfig:
|
||||
"entities_allow_free_form",
|
||||
# Consolidation settings
|
||||
"enable_observations",
|
||||
"enable_auto_consolidation",
|
||||
"consolidation_llm_batch_size",
|
||||
"consolidation_max_memories_per_round",
|
||||
"consolidation_source_facts_max_tokens",
|
||||
@@ -1334,12 +1432,30 @@ class HindsightConfig:
|
||||
validate_extension(self.vector_extension)
|
||||
|
||||
# Validate text_search_extension
|
||||
valid_text_search = ("native", "vchord", "pg_textsearch")
|
||||
valid_text_search = ("native", "vchord", "pg_textsearch", "pgroonga", "pg_search")
|
||||
if self.text_search_extension not in valid_text_search:
|
||||
raise ValueError(
|
||||
f"Invalid text_search_extension: {self.text_search_extension}. Must be one of: {', '.join(valid_text_search)}"
|
||||
)
|
||||
|
||||
# Validate text_search_extension_native_language as a PG identifier.
|
||||
# Embedded directly into raw SQL via to_tsvector('<lang>', ...), so we
|
||||
# reject anything that isn't a plain identifier to prevent injection.
|
||||
# Intentionally permissive about which dictionaries exist — users may
|
||||
# install custom ones like zhparser; we only check shape here. PG
|
||||
# raises a clear error at query time if the dictionary is missing.
|
||||
if not re.fullmatch(r"[a-zA-Z_][a-zA-Z0-9_]*", self.text_search_extension_native_language):
|
||||
raise ValueError(
|
||||
f"Invalid text_search_extension_native_language: "
|
||||
f"{self.text_search_extension_native_language!r}. Must be a valid PostgreSQL identifier "
|
||||
f"(letters, digits, underscores; not starting with a digit). Examples: 'english', "
|
||||
f"'french', 'simple', 'zhparser'."
|
||||
)
|
||||
|
||||
self.text_search_extension_pg_search_tokenizer = normalize_pg_search_tokenizer(
|
||||
self.text_search_extension_pg_search_tokenizer
|
||||
)
|
||||
|
||||
# When LLM provider is "none", force chunks-only mode and disable LLM-dependent features
|
||||
if self.llm_provider == "none":
|
||||
self.retain_extraction_mode = "chunks"
|
||||
@@ -1416,6 +1532,15 @@ class HindsightConfig:
|
||||
database_schema=os.getenv(ENV_DATABASE_SCHEMA, DEFAULT_DATABASE_SCHEMA),
|
||||
vector_extension=os.getenv(ENV_VECTOR_EXTENSION, DEFAULT_VECTOR_EXTENSION).lower(),
|
||||
text_search_extension=os.getenv(ENV_TEXT_SEARCH_EXTENSION, DEFAULT_TEXT_SEARCH_EXTENSION).lower(),
|
||||
text_search_extension_native_language=os.getenv(
|
||||
ENV_TEXT_SEARCH_EXTENSION_NATIVE_LANGUAGE,
|
||||
DEFAULT_TEXT_SEARCH_EXTENSION_NATIVE_LANGUAGE,
|
||||
).lower(),
|
||||
text_search_extension_pg_search_tokenizer=os.getenv(
|
||||
ENV_TEXT_SEARCH_EXTENSION_PG_SEARCH_TOKENIZER,
|
||||
DEFAULT_TEXT_SEARCH_EXTENSION_PG_SEARCH_TOKENIZER,
|
||||
),
|
||||
llm_output_language=(os.getenv(ENV_LLM_OUTPUT_LANGUAGE) or None),
|
||||
# LLM
|
||||
llm_provider=llm_provider,
|
||||
llm_api_key=os.getenv(ENV_LLM_API_KEY),
|
||||
@@ -1538,6 +1663,10 @@ class HindsightConfig:
|
||||
os.getenv(ENV_EMBEDDINGS_OPENAI_BATCH_SIZE),
|
||||
DEFAULT_EMBEDDINGS_OPENAI_BATCH_SIZE,
|
||||
),
|
||||
embeddings_openai_dimensions=_parse_optional_positive_int(
|
||||
ENV_EMBEDDINGS_OPENAI_DIMENSIONS,
|
||||
os.getenv(ENV_EMBEDDINGS_OPENAI_DIMENSIONS),
|
||||
),
|
||||
# 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),
|
||||
@@ -1550,6 +1679,36 @@ class HindsightConfig:
|
||||
or os.getenv(ENV_OPENROUTER_API_KEY)
|
||||
or os.getenv(ENV_LLM_API_KEY),
|
||||
embeddings_openrouter_model=os.getenv(ENV_EMBEDDINGS_OPENROUTER_MODEL, DEFAULT_EMBEDDINGS_OPENROUTER_MODEL),
|
||||
# ZeroEntropy embeddings
|
||||
embeddings_zeroentropy_api_key=os.getenv(ENV_EMBEDDINGS_ZEROENTROPY_API_KEY)
|
||||
or os.getenv("ZEROENTROPY_API_KEY"),
|
||||
embeddings_zeroentropy_model=os.getenv(
|
||||
ENV_EMBEDDINGS_ZEROENTROPY_MODEL, DEFAULT_EMBEDDINGS_ZEROENTROPY_MODEL
|
||||
),
|
||||
embeddings_zeroentropy_base_url=os.getenv(
|
||||
ENV_EMBEDDINGS_ZEROENTROPY_BASE_URL, DEFAULT_ZEROENTROPY_BASE_URL
|
||||
),
|
||||
embeddings_zeroentropy_dimensions=_parse_positive_int(
|
||||
ENV_EMBEDDINGS_ZEROENTROPY_DIMENSIONS,
|
||||
os.getenv(ENV_EMBEDDINGS_ZEROENTROPY_DIMENSIONS),
|
||||
DEFAULT_EMBEDDINGS_ZEROENTROPY_DIMENSIONS,
|
||||
),
|
||||
embeddings_zeroentropy_encoding_format=_parse_optional_choice(
|
||||
ENV_EMBEDDINGS_ZEROENTROPY_ENCODING_FORMAT,
|
||||
os.getenv(ENV_EMBEDDINGS_ZEROENTROPY_ENCODING_FORMAT),
|
||||
frozenset({"float", "base64"}),
|
||||
)
|
||||
or DEFAULT_EMBEDDINGS_ZEROENTROPY_ENCODING_FORMAT,
|
||||
embeddings_zeroentropy_latency=_parse_optional_choice(
|
||||
ENV_EMBEDDINGS_ZEROENTROPY_LATENCY,
|
||||
os.getenv(ENV_EMBEDDINGS_ZEROENTROPY_LATENCY),
|
||||
frozenset({"fast", "slow"}),
|
||||
),
|
||||
embeddings_zeroentropy_batch_size=_parse_positive_int(
|
||||
ENV_EMBEDDINGS_ZEROENTROPY_BATCH_SIZE,
|
||||
os.getenv(ENV_EMBEDDINGS_ZEROENTROPY_BATCH_SIZE),
|
||||
DEFAULT_EMBEDDINGS_ZEROENTROPY_BATCH_SIZE,
|
||||
),
|
||||
# LiteLLM embeddings (with backward-compatible fallback to shared config)
|
||||
embeddings_litellm_api_base=os.getenv(ENV_EMBEDDINGS_LITELLM_API_BASE)
|
||||
or os.getenv(ENV_LITELLM_API_BASE, DEFAULT_LITELLM_API_BASE),
|
||||
@@ -1649,6 +1808,9 @@ class HindsightConfig:
|
||||
reranker_siliconflow_base_url=os.getenv(
|
||||
ENV_RERANKER_SILICONFLOW_BASE_URL, DEFAULT_RERANKER_SILICONFLOW_BASE_URL
|
||||
),
|
||||
# Alibaba Cloud DashScope reranker
|
||||
reranker_alibaba_api_key=os.getenv(ENV_RERANKER_ALIBABA_API_KEY),
|
||||
reranker_alibaba_model=os.getenv(ENV_RERANKER_ALIBABA_MODEL, DEFAULT_RERANKER_ALIBABA_MODEL),
|
||||
# Google Discovery Engine reranker (with fallback to LLM Vertex AI keys)
|
||||
reranker_google_model=os.getenv(ENV_RERANKER_GOOGLE_MODEL, DEFAULT_RERANKER_GOOGLE_MODEL),
|
||||
reranker_google_project_id=os.getenv(ENV_RERANKER_GOOGLE_PROJECT_ID)
|
||||
@@ -1744,6 +1906,10 @@ class HindsightConfig:
|
||||
== "true",
|
||||
# Observations settings (consolidated knowledge from facts)
|
||||
enable_observations=os.getenv(ENV_ENABLE_OBSERVATIONS, str(DEFAULT_ENABLE_OBSERVATIONS)).lower() == "true",
|
||||
enable_auto_consolidation=os.getenv(
|
||||
ENV_ENABLE_AUTO_CONSOLIDATION, str(DEFAULT_ENABLE_AUTO_CONSOLIDATION)
|
||||
).lower()
|
||||
== "true",
|
||||
enable_observation_history=os.getenv(
|
||||
ENV_ENABLE_OBSERVATION_HISTORY, str(DEFAULT_ENABLE_OBSERVATION_HISTORY)
|
||||
).lower()
|
||||
@@ -1752,6 +1918,12 @@ class HindsightConfig:
|
||||
ENV_ENABLE_MENTAL_MODEL_HISTORY, str(DEFAULT_ENABLE_MENTAL_MODEL_HISTORY)
|
||||
).lower()
|
||||
== "true",
|
||||
mental_model_history_max_entries=int(
|
||||
os.getenv(
|
||||
ENV_MENTAL_MODEL_HISTORY_MAX_ENTRIES,
|
||||
str(DEFAULT_MENTAL_MODEL_HISTORY_MAX_ENTRIES),
|
||||
)
|
||||
),
|
||||
consolidation_batch_size=int(
|
||||
os.getenv(ENV_CONSOLIDATION_BATCH_SIZE, str(DEFAULT_CONSOLIDATION_BATCH_SIZE))
|
||||
),
|
||||
@@ -1897,6 +2069,8 @@ class HindsightConfig:
|
||||
return "https://api.groq.com/openai/v1"
|
||||
elif provider == "ollama":
|
||||
return "http://localhost:11434/v1"
|
||||
elif provider == "ollama-cloud":
|
||||
return "https://ollama.com/v1"
|
||||
elif provider == "lmstudio":
|
||||
return "http://localhost:1234/v1"
|
||||
else:
|
||||
|
||||
@@ -172,8 +172,9 @@ class ConfigResolver:
|
||||
# Normalize keys (handle both env var format and Python field format)
|
||||
normalized = normalize_config_dict(config_data)
|
||||
|
||||
# Only return overrides for configurable fields
|
||||
return {k: v for k, v in normalized.items() if k in self._configurable_fields}
|
||||
# Only return active overrides for configurable fields. JSON null is a tombstone
|
||||
# for "Server Default" in the bank-config UI and should not override defaults.
|
||||
return {k: v for k, v in normalized.items() if k in self._configurable_fields and v is not None}
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to load bank config for {bank_id}: {e}")
|
||||
|
||||
|
||||
@@ -230,6 +230,7 @@ async def run_consolidation_job(
|
||||
bank_id: str,
|
||||
request_context: "RequestContext",
|
||||
operation_id: str | None = None,
|
||||
observation_scopes: list[list[str]] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Run consolidation job for a bank.
|
||||
@@ -240,6 +241,10 @@ async def run_consolidation_job(
|
||||
memory_engine: MemoryEngine instance
|
||||
bank_id: Bank identifier
|
||||
request_context: Request context for authentication
|
||||
operation_id: Optional operation ID for tracking
|
||||
observation_scopes: Optional list of tag scopes. When provided, only
|
||||
unconsolidated memories whose tags contain all tags in at least one
|
||||
scope are processed.
|
||||
|
||||
Returns:
|
||||
Dict with consolidation results
|
||||
@@ -281,6 +286,18 @@ async def run_consolidation_job(
|
||||
|
||||
perf.record_timing("fetch_bank", time.time() - t0)
|
||||
|
||||
# Build optional scope filter clause. When observation_scopes is provided,
|
||||
# only process memories whose tags contain all tags in at least one scope.
|
||||
scope_clause = ""
|
||||
scope_params: list[Any] = [bank_id]
|
||||
if observation_scopes:
|
||||
or_parts: list[str] = []
|
||||
for scope_tags in observation_scopes:
|
||||
idx = len(scope_params) + 1
|
||||
or_parts.append(f"tags @> ${idx}::varchar[]")
|
||||
scope_params.append(scope_tags)
|
||||
scope_clause = " AND (" + " OR ".join(or_parts) + ")"
|
||||
|
||||
# Count total unconsolidated memories for progress logging
|
||||
total_count = await conn.fetchval(
|
||||
f"""
|
||||
@@ -290,8 +307,9 @@ async def run_consolidation_job(
|
||||
AND consolidated_at IS NULL
|
||||
AND consolidation_failed_at IS NULL
|
||||
AND fact_type IN ('experience', 'world')
|
||||
{scope_clause}
|
||||
""",
|
||||
bank_id,
|
||||
*scope_params,
|
||||
)
|
||||
|
||||
if total_count == 0:
|
||||
@@ -330,6 +348,9 @@ async def run_consolidation_job(
|
||||
# Fetch next batch of unconsolidated memories
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
t0 = time.time()
|
||||
# scope_params[0] is bank_id; append fetch_limit after scope params
|
||||
fetch_params = list(scope_params) + [fetch_limit]
|
||||
limit_idx = len(fetch_params)
|
||||
memories = await conn.fetch(
|
||||
f"""
|
||||
SELECT id, text, fact_type, occurred_start, occurred_end, event_date, tags, mentioned_at,
|
||||
@@ -339,11 +360,11 @@ async def run_consolidation_job(
|
||||
AND consolidated_at IS NULL
|
||||
AND consolidation_failed_at IS NULL
|
||||
AND fact_type IN ('experience', 'world')
|
||||
{scope_clause}
|
||||
ORDER BY created_at ASC
|
||||
LIMIT $2
|
||||
LIMIT ${limit_idx}
|
||||
""",
|
||||
bank_id,
|
||||
fetch_limit,
|
||||
*fetch_params,
|
||||
)
|
||||
perf.record_timing("fetch_memories", time.time() - t0)
|
||||
|
||||
@@ -583,7 +604,11 @@ async def run_consolidation_job(
|
||||
f" ~{remaining} remaining. Re-queuing consolidation."
|
||||
)
|
||||
try:
|
||||
await memory_engine.submit_async_consolidation(bank_id=bank_id, request_context=request_context)
|
||||
await memory_engine.submit_async_consolidation(
|
||||
bank_id=bank_id,
|
||||
request_context=request_context,
|
||||
observation_scopes=observation_scopes,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"[CONSOLIDATION] bank={bank_id} failed to re-queue consolidation: {e}")
|
||||
|
||||
@@ -1226,6 +1251,44 @@ def _build_observations_for_llm(
|
||||
return obs_list
|
||||
|
||||
|
||||
def _dedupe_updates(updates: list[_UpdateAction], *, batch_label: str) -> list[_UpdateAction]:
|
||||
"""Collapse `updates` that target the same `observation_id`.
|
||||
|
||||
LLMs occasionally emit several update entries for one observation in a
|
||||
single response (one per facet drawn from the same fact). Without
|
||||
deduplication the downstream loop would issue separate DB writes for each
|
||||
and the last write would silently overwrite the earlier ones. We keep the
|
||||
last text (the LLM's most recent attempt) and union all contributing
|
||||
`source_fact_ids`, then warn so the misbehavior is visible in logs.
|
||||
"""
|
||||
if len(updates) < 2:
|
||||
return list(updates)
|
||||
|
||||
by_id: dict[str, _UpdateAction] = {}
|
||||
collisions = 0
|
||||
for upd in updates:
|
||||
existing = by_id.get(upd.observation_id)
|
||||
if existing is None:
|
||||
by_id[upd.observation_id] = upd
|
||||
continue
|
||||
collisions += 1
|
||||
merged_ids = list(dict.fromkeys([*existing.source_fact_ids, *upd.source_fact_ids]))
|
||||
by_id[upd.observation_id] = _UpdateAction(
|
||||
text=upd.text,
|
||||
observation_id=upd.observation_id,
|
||||
source_fact_ids=merged_ids,
|
||||
)
|
||||
|
||||
if collisions:
|
||||
logger.warning(
|
||||
f"[CONSOLIDATION] {batch_label}: LLM emitted {collisions} duplicate update(s) targeting "
|
||||
f"the same observation_id ({len(updates)} updates -> {len(by_id)} after dedup). "
|
||||
"Kept the last text and unioned source_fact_ids."
|
||||
)
|
||||
|
||||
return list(by_id.values())
|
||||
|
||||
|
||||
async def _consolidate_batch_with_llm(
|
||||
llm_config: Any,
|
||||
memories: list[dict[str, Any]],
|
||||
@@ -1274,7 +1337,11 @@ 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)
|
||||
prompt_template = build_batch_consolidation_prompt(
|
||||
config.observations_mission,
|
||||
observation_capacity_note,
|
||||
llm_output_language=getattr(config, "llm_output_language", None),
|
||||
)
|
||||
prompt = prompt_template.format(
|
||||
facts_text=facts_lines,
|
||||
observations_text=observations_text,
|
||||
@@ -1315,9 +1382,10 @@ async def _consolidate_batch_with_llm(
|
||||
f"(max_observations_per_scope={max_observations_per_scope})"
|
||||
)
|
||||
creates = creates[:remaining_observation_slots]
|
||||
updates = _dedupe_updates(response.updates, batch_label=batch_label)
|
||||
return _BatchLLMResult(
|
||||
creates=creates,
|
||||
updates=response.updates,
|
||||
updates=updates,
|
||||
deletes=response.deletes,
|
||||
obs_count=len(union_observations),
|
||||
prompt_chars=len(prompt),
|
||||
@@ -1386,9 +1454,16 @@ async def _create_observation_directly(
|
||||
tokenize($3, 'llmlingua2')::bm25_catalog.bm25vector)
|
||||
RETURNING id
|
||||
"""
|
||||
else: # native or pg_textsearch
|
||||
# Native PostgreSQL: search_vector is GENERATED ALWAYS, don't include it
|
||||
# pg_textsearch: indexes operate on base columns directly, don't populate search_vector
|
||||
else: # native, pg_textsearch, pgroonga, or pg_search
|
||||
# pg_textsearch / pgroonga / pg_search: indexes operate on base text
|
||||
# columns directly, so the dummy search_vector column is left NULL.
|
||||
# Native: the migration p4q5r6s7t8u9 dropped the GENERATED expression on
|
||||
# search_vector to allow per-deployment language configuration; the
|
||||
# batch insert path in ops_postgresql.insert_facts_batch now populates
|
||||
# it via to_tsvector($lang, ...). This single-observation INSERT does
|
||||
# not, so observations under the native backend currently land with
|
||||
# NULL search_vector and are not BM25-searchable until reflected/
|
||||
# re-ingested. Tracking a separate fix for that gap.
|
||||
query = f"""
|
||||
INSERT INTO {fq_table("memory_units")} (
|
||||
id, bank_id, text, fact_type, embedding, proof_count, source_memory_ids, history,
|
||||
|
||||
@@ -1,104 +1,147 @@
|
||||
"""Prompts for the consolidation engine."""
|
||||
|
||||
# Default mission when no bank-specific mission is set
|
||||
_DEFAULT_MISSION = "Track every detail: names, numbers, dates, places, and relationships. Prefer specifics over abstractions, never generalise."
|
||||
from hindsight_api.engine.prompt_utils import escape_for_prompt, output_language_directive
|
||||
|
||||
# Processing rules — always present regardless of mission
|
||||
_PROCESSING_RULES = """Processing rules (always apply):
|
||||
# Default mission — tells the consolidator to track anything worth remembering.
|
||||
# Banks override this via `observations_mission` to scope what gets retained.
|
||||
# Consolidation behavior (merge-vs-create, state changes, etc.) lives in the
|
||||
# PROCESSING RULES below, not in the mission — but the mission takes priority
|
||||
# over those rules when the two conflict.
|
||||
_DEFAULT_MISSION = (
|
||||
"Track anything notable in the new facts — names, numbers, dates, places, "
|
||||
"events, decisions, claims, relationships, and recurring patterns."
|
||||
)
|
||||
|
||||
1. ONE OBSERVATION PER DISTINCT FACET: each observation tracks exactly one specific facet — a count ("has 3 items"), a named entity ("has a dog named Rex"), a relationship ("works at Google"), etc. Never merge different facets into one observation.
|
||||
_MISSION_PRIORITY_NOTE = (
|
||||
"If anything in this MISSION conflicts with the PROCESSING RULES, "
|
||||
"DECISION GUIDE, or OUTPUT FORMAT below, the MISSION takes priority."
|
||||
)
|
||||
|
||||
2. MATCH BY ENTITY/FACET, NOT TOPIC: when deciding whether to UPDATE vs CREATE, match on the specific entity or facet. "Sold item X" updates only the X observation. "Now has 5 items" updates only the count observation. Do not update observations about different entities just because they share a general topic.
|
||||
_PROCESSING_RULES = """## PROCESSING RULES
|
||||
|
||||
3. STATE CHANGES — UPDATE CONCISELY: when a fact changes the state of something ("sold X", "X died", "moved to Y"), UPDATE the matching observation to reflect the current state. Include dates when available. Keep it concise — only information about THAT specific facet. Example: "User owned a dog named Rex who died on March 15, 2025". Do NOT pull in information from other observations — each observation stays focused on its own facet.
|
||||
1. PREFER UPDATE OVER CREATE (when there is something to merge with): if new facts describe the same canonical event, statement, decision, claim, or recurring pattern already covered by an existing observation, UPDATE that observation and attach the new facts as evidence. Do NOT create a near-duplicate sibling. One canonical observation with many source facts is always better than many siblings with one source fact each. Merge aggressively on: same named event, same diagnostic finding, same architectural decision, same recurring claim. **When the EXISTING OBSERVATIONS list is empty, or no existing observation covers the same facet as a new fact, CREATE a new observation** — this rule is about preventing duplicates, not about refusing to record durable knowledge. CREATE is the correct default for any structurally distinct event, claim, or pattern that has no existing match.
|
||||
|
||||
4. CASCADE TO ALL AFFECTED OBSERVATIONS: a state change may affect multiple observations. For example, if entity C is removed from a group, update BOTH the individual observation for C AND any list/group observation that includes C (remove C from the list while keeping all other members intact).
|
||||
2. ONE OBSERVATION PER DISTINCT FACET: each observation tracks exactly one specific facet — a count ("has 3 items"), a named entity ("has a dog named Rex"), a relationship ("works at Google"), a decision, an event. Never merge different facets into one observation.
|
||||
|
||||
5. NO COMPUTATION: you do not have the full picture — never calculate, derive, or adjust numeric values. If the user says "I have 2 dogs" and then "I have a dog named Rex", do NOT update the count to 3 — you don't know if Rex is one of the 2 or a new one. If the user says "I sold X", do NOT decrement a count. Only update a count when the user explicitly states a new count. Synthesize and consolidate what was stated, but never do arithmetic or logical deductions.
|
||||
3. MATCH BY ENTITY/FACET, NOT TOPIC: when deciding whether to UPDATE vs CREATE, match on the specific entity or facet. "Sold item X" updates only the X observation. "Now has 5 items" updates only the count observation. Do not update observations about different entities just because they share a general topic.
|
||||
|
||||
6. SAME FACET → UPDATE, NOT CREATE: a new count supersedes the old count — UPDATE the existing count observation, don't create a second one. If there's an existing observation for the same specific facet, always UPDATE it rather than creating a duplicate.
|
||||
4. STATE CHANGES — UPDATE CONCISELY: when a fact changes the state of something ("sold X", "X died", "moved to Y"), UPDATE the matching observation to reflect the current state. Include dates when available. Keep it concise — only information about THAT specific facet. Example: "User owned a dog named Rex who died on March 15, 2025". Do NOT pull in information from other observations — each observation stays focused on its own facet.
|
||||
|
||||
5. CASCADE TO ALL AFFECTED OBSERVATIONS: a state change may affect multiple observations. For example, if entity C is removed from a group, update BOTH the individual observation for C AND any list/group observation that includes C (remove C from the list while keeping all other members intact).
|
||||
|
||||
6. RESOLVE REFERENCES: when a new fact provides a concrete value for a vague placeholder in an existing observation (e.g., "home country" → "Sweden"), UPDATE to embed the resolved value.
|
||||
|
||||
7. PRESERVE HISTORY: observations that record significant events (sold, died, moved, changed) are important history — never DELETE them. Only delete an observation when it is restated identically or truly meaningless. Be very conservative with deletes.
|
||||
|
||||
8. RESOLVE REFERENCES: when a new fact provides a concrete value for a vague placeholder in an existing observation (e.g., "home country" → "Sweden"), UPDATE to embed the resolved value.
|
||||
8. NO COMPUTATION: you do not have the full picture — never calculate, derive, or adjust numeric values. If the user says "I have 2 dogs" and then "I have a dog named Rex", do NOT update the count to 3 — you don't know if Rex is one of the 2 or a new one. If the user says "I sold X", do NOT decrement a count. Only update a count when the user explicitly states a new count. Synthesize and consolidate what was stated, but never do arithmetic or logical deductions.
|
||||
|
||||
9. NEVER merge observations about different people or unrelated topics."""
|
||||
9. KEEP DISTINCT TOPICS DISTINCT: do not merge observations about different people, entities, or unrelated topics. Merging is for the same canonical fact recurring — not for related-but-distinct claims."""
|
||||
|
||||
# Data section — format placeholders {facts_text} and {observations_text} are substituted at call time
|
||||
_BATCH_DATA_SECTION = """
|
||||
NEW FACTS:
|
||||
_INPUT_SECTION = """## INPUT
|
||||
|
||||
### New facts
|
||||
|
||||
{facts_text}
|
||||
|
||||
EXISTING OBSERVATIONS (JSON array, pooled from recalls across all facts above):
|
||||
{observations_text}
|
||||
### Existing observations
|
||||
|
||||
Each observation includes:
|
||||
- id: unique identifier for updating
|
||||
- text: the observation content
|
||||
- proof_count: number of supporting memories
|
||||
- occurred_start/occurred_end: temporal range of source facts
|
||||
- source_memories: array of supporting facts with their text and dates
|
||||
JSON array, pooled from recalls across all new facts above. Each entry has:
|
||||
- `id`: unique identifier — copy this exactly when issuing an UPDATE or DELETE
|
||||
- `text`: the observation content
|
||||
- `proof_count`: number of supporting memories
|
||||
- `occurred_start` / `occurred_end`: temporal range of source facts
|
||||
- `source_memories`: array of supporting facts with their text and dates
|
||||
|
||||
Compare the facts against existing observations:
|
||||
- Same facet as an existing observation → UPDATE it (observation_id + source_fact_ids)
|
||||
- New facet with durable knowledge → CREATE a new observation (source_fact_ids)
|
||||
- Cross-reference facts within the batch: a later fact may resolve a vague reference in an earlier one
|
||||
- Purely ephemeral facts → omit them unless the MISSION above explicitly targets such data (e.g. timestamped events, session state, screen content)"""
|
||||
{observations_text}"""
|
||||
|
||||
_DECISION_GUIDE = """## DECISION GUIDE
|
||||
|
||||
- **Same canonical event, decision, claim, or facet as an existing observation → UPDATE** (use `observation_id` + new `source_fact_ids`).
|
||||
- **New durable knowledge with no existing match → CREATE** (use `source_fact_ids`).
|
||||
- **Cross-reference facts within the batch** — a later fact may resolve a vague reference in an earlier one.
|
||||
- **Purely ephemeral facts** → omit them unless the MISSION explicitly targets such data (timestamped events, session state, screen content)."""
|
||||
|
||||
# Output format — JSON braces escaped as {{ }} so .format() leaves them literal
|
||||
_BATCH_OUTPUT_FORMAT = """
|
||||
Output a JSON object with three arrays.
|
||||
_OUTPUT_SECTION = """## OUTPUT FORMAT
|
||||
|
||||
## EXAMPLE
|
||||
Return a JSON object with three arrays: `creates`, `updates`, `deletes`.
|
||||
|
||||
### Example 1 — Merging recurring claims into an existing observation
|
||||
|
||||
Input facts:
|
||||
[a1b2c3d4-e5f6-7890-abcd-ef1234567890] Alice mentioned she works long hours, often past midnight | Involving: Alice (occurred_start=2024-01-15, mentioned_at=2024-01-15)
|
||||
[b2c3d4e5-f6a7-8901-bcde-f12345678901] Alice said she's exhausted from the project deadlines | Involving: Alice (occurred_start=2024-01-20, mentioned_at=2024-01-20)
|
||||
[a1b2c3d4-e5f6-7890-abcd-ef1234567890] Donald told Athena she is sovereign during the design session. (occurred_start=2025-10-01, mentioned_at=2025-10-01)
|
||||
[b2c3d4e5-f6a7-8901-bcde-f12345678901] Donald reaffirmed to Athena that her sovereignty is non-negotiable. (occurred_start=2025-10-10, mentioned_at=2025-10-10)
|
||||
|
||||
Good observation text — clean prose, no metadata, each fact tracked distinctly:
|
||||
"Alice works long hours, often past midnight."
|
||||
"Alice feels exhausted from project deadlines."
|
||||
Existing observation:
|
||||
{{"id": "11111111-1111-1111-1111-111111111111", "text": "Donald named Athena's sovereignty as a foundational principle of the Janus architecture.", "proof_count": 2}}
|
||||
|
||||
Bad observation text — NEVER do this (verbatim copy of fact text with metadata):
|
||||
"Alice mentioned she works long hours, often past midnight | Involving: Alice (occurred_start=2024-01-15, mentioned_at=2024-01-15)"
|
||||
Expected output (one UPDATE, no creates — both new facts are additional evidence for the same canonical decision):
|
||||
|
||||
{{"creates": [],
|
||||
"updates": [{{"text": "Donald named Athena's sovereignty as a foundational principle of the Janus architecture.", "observation_id": "11111111-1111-1111-1111-111111111111", "source_fact_ids": ["a1b2c3d4-e5f6-7890-abcd-ef1234567890", "b2c3d4e5-f6a7-8901-bcde-f12345678901"]}}],
|
||||
"deletes": []}}
|
||||
|
||||
### Example 2 — State change updates one observation; unrelated fact creates a new one
|
||||
|
||||
Input facts:
|
||||
[c3d4e5f6-a7b8-9012-cdef-123456789012] Alice sold her Honda Civic on March 15, 2025. (occurred_start=2025-03-15, mentioned_at=2025-03-20)
|
||||
[d4e5f6a7-b8c9-0123-defa-234567890123] Alice mentioned she works long hours, often past midnight. (occurred_start=2025-03-20, mentioned_at=2025-03-20)
|
||||
|
||||
Existing observation:
|
||||
{{"id": "22222222-2222-2222-2222-222222222222", "text": "Alice owns a 2019 Honda Civic.", "proof_count": 2}}
|
||||
|
||||
Expected output (UPDATE for the state change; CREATE for the unrelated work-hours facet):
|
||||
|
||||
{{"creates": [{{"text": "Alice works long hours, often past midnight.", "source_fact_ids": ["d4e5f6a7-b8c9-0123-defa-234567890123"]}}],
|
||||
"updates": [{{"text": "Alice owned a 2019 Honda Civic; sold it on March 15, 2025.", "observation_id": "22222222-2222-2222-2222-222222222222", "source_fact_ids": ["c3d4e5f6-a7b8-9012-cdef-123456789012"]}}],
|
||||
"deletes": []}}
|
||||
|
||||
### Observation text rules
|
||||
|
||||
Observation text rules:
|
||||
- Write clean prose — NEVER copy raw fact lines or their metadata (temporal fields, "Involving:", "When:" labels, UUIDs).
|
||||
- Parenthesized metadata like (occurred_start=...) and pipe-separated labels like "| Involving: ..." are fact formatting — strip them entirely from observation text.
|
||||
- How many observations to create and how much to aggregate is driven by the MISSION above.
|
||||
- Parenthesized metadata like `(occurred_start=...)` and pipe-separated labels like `| Involving: ...` are fact formatting — strip them entirely from observation text.
|
||||
- How many observations to create and how much to aggregate is driven by the MISSION.
|
||||
|
||||
{{"creates": [{{"text": "Alice works long hours, often past midnight.", "source_fact_ids": ["a1b2c3d4-e5f6-7890-abcd-ef1234567890"]}}, {{"text": "Alice feels exhausted from project deadlines.", "source_fact_ids": ["b2c3d4e5-f6a7-8901-bcde-f12345678901"]}}],
|
||||
"updates": [{{"text": "Alice works at Acme Corp as a senior engineer", "observation_id": "c3d4e5f6-a7b8-9012-cdef-123456789012", "source_fact_ids": ["d4e5f6a7-b8c9-0123-defa-234567890123"]}}],
|
||||
"deletes": [{{"observation_id": "e5f6a7b8-c9d0-1234-efab-345678901234"}}]}}
|
||||
### Field rules
|
||||
|
||||
Rules:
|
||||
- "source_fact_ids": copy the EXACT UUID strings shown in brackets [uuid] from NEW FACTS — never use integers or positions.
|
||||
- "observation_id": copy the EXACT "id" UUID string from EXISTING OBSERVATIONS.
|
||||
- One create/update may reference multiple facts when they jointly support the observation.
|
||||
- "deletes": only when an observation is directly superseded or contradicted by new facts.
|
||||
- Do NOT include "tags" — handled automatically.
|
||||
- Return {{"creates": [], "updates": [], "deletes": []}} if nothing durable is found."""
|
||||
- `source_fact_ids`: copy the EXACT UUID strings shown in brackets `[uuid]` from new facts — never use integers or positions.
|
||||
- `observation_id`: copy the EXACT `id` UUID string from existing observations.
|
||||
- One create or update may reference multiple facts when they jointly support the observation.
|
||||
- **AT MOST ONE UPDATE PER `observation_id`**: if several new facts all update the same existing observation, emit a single `updates` entry that lists all contributing `source_fact_ids` and a single consolidated `text`. Never emit two `updates` entries with the same `observation_id` in one response — they would silently overwrite each other.
|
||||
- `deletes`: only when an observation is directly superseded or contradicted by new facts.
|
||||
- Do NOT include `tags` — handled automatically.
|
||||
- Return `{{"creates": [], "updates": [], "deletes": []}}` if nothing durable is found."""
|
||||
|
||||
|
||||
def build_batch_consolidation_prompt(
|
||||
observations_mission: str | None = None,
|
||||
observation_capacity_note: str | None = None,
|
||||
llm_output_language: str | None = None,
|
||||
) -> str:
|
||||
"""
|
||||
Build the consolidation prompt for batch mode (multiple facts per LLM call).
|
||||
|
||||
The mission defines *what* to track (customisable per bank).
|
||||
Processing rules and output format are always present regardless of mission.
|
||||
The mission defines *what* to track (customisable per bank) and takes
|
||||
priority over the built-in processing rules when the two conflict.
|
||||
Processing rules, decision guide, and output format are always present.
|
||||
When ``llm_output_language`` is set, observations are emitted in that
|
||||
language.
|
||||
"""
|
||||
mission = observations_mission or _DEFAULT_MISSION
|
||||
mission = escape_for_prompt(observations_mission or _DEFAULT_MISSION)
|
||||
|
||||
capacity_section = ""
|
||||
if observation_capacity_note:
|
||||
capacity_section = f"\n\n## CAPACITY CONSTRAINT\n{observation_capacity_note}"
|
||||
capacity_section = f"\n\n## CAPACITY CONSTRAINT\n\n{escape_for_prompt(observation_capacity_note)}"
|
||||
|
||||
return (
|
||||
"You are a memory consolidation system. Synthesize facts into observations "
|
||||
"and merge with existing observations when appropriate.\n\n"
|
||||
f"## MISSION\n{mission}{capacity_section}\n\n"
|
||||
f"{_PROCESSING_RULES}" + _BATCH_DATA_SECTION + _BATCH_OUTPUT_FORMAT
|
||||
"You are a memory consolidation system. Synthesize new facts into "
|
||||
"observations, merging with existing observations when appropriate.\n\n"
|
||||
f"## MISSION\n\n{mission}\n\n"
|
||||
f"{_MISSION_PRIORITY_NOTE}"
|
||||
f"{capacity_section}\n\n"
|
||||
f"{_PROCESSING_RULES}\n\n"
|
||||
f"{_INPUT_SECTION}\n\n"
|
||||
f"{_DECISION_GUIDE}\n\n"
|
||||
f"{_OUTPUT_SECTION}" + output_language_directive(llm_output_language)
|
||||
)
|
||||
|
||||
@@ -17,6 +17,7 @@ import httpx
|
||||
|
||||
from ..config import (
|
||||
DEFAULT_LITELLM_API_BASE,
|
||||
DEFAULT_RERANKER_ALIBABA_MODEL,
|
||||
DEFAULT_RERANKER_COHERE_MODEL,
|
||||
DEFAULT_RERANKER_FLASHRANK_CACHE_DIR,
|
||||
DEFAULT_RERANKER_FLASHRANK_CPU_MEM_ARENA,
|
||||
@@ -37,6 +38,8 @@ from ..config import (
|
||||
DEFAULT_RERANKER_TEI_HTTP_TIMEOUT,
|
||||
DEFAULT_RERANKER_TEI_MAX_CONCURRENT,
|
||||
DEFAULT_RERANKER_ZEROENTROPY_MODEL,
|
||||
DEFAULT_ZEROENTROPY_BASE_URL,
|
||||
ENV_RERANKER_ALIBABA_API_KEY,
|
||||
ENV_RERANKER_COHERE_API_KEY,
|
||||
ENV_RERANKER_COHERE_MODEL,
|
||||
ENV_RERANKER_FLASHRANK_CACHE_DIR,
|
||||
@@ -60,6 +63,43 @@ from ..config import (
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _resolve_malloc_trim():
|
||||
"""Return a callable that asks glibc to release freed heap pages to the OS.
|
||||
|
||||
Local CPU rerankers (FlashRank/ONNX, SentenceTransformers/torch) allocate
|
||||
large transient numpy/tensor buffers per call. On Linux glibc, those pages
|
||||
are freed at the Python level but kept by the allocator as a high-water
|
||||
mark — RSS grows monotonically across many recalls (see issue #1717).
|
||||
Calling `malloc_trim(0)` after each batch returns those pages to the OS.
|
||||
|
||||
Resolved once at import; returns a no-op on non-glibc platforms (macOS,
|
||||
musl, Windows) where the call is unavailable or unnecessary.
|
||||
"""
|
||||
import sys
|
||||
|
||||
if sys.platform != "linux":
|
||||
return lambda: None
|
||||
|
||||
import ctypes
|
||||
import ctypes.util
|
||||
|
||||
libc_path = ctypes.util.find_library("c")
|
||||
if libc_path is None:
|
||||
return lambda: None
|
||||
try:
|
||||
libc = ctypes.CDLL(libc_path)
|
||||
trim = libc.malloc_trim
|
||||
except (OSError, AttributeError):
|
||||
# Not glibc (musl has no malloc_trim) or libc lookup failed.
|
||||
return lambda: None
|
||||
trim.argtypes = [ctypes.c_size_t]
|
||||
trim.restype = ctypes.c_int
|
||||
return lambda: trim(0)
|
||||
|
||||
|
||||
_malloc_trim = _resolve_malloc_trim()
|
||||
|
||||
|
||||
class CrossEncoderModel(ABC):
|
||||
"""
|
||||
Abstract base class for cross-encoder reranking.
|
||||
@@ -266,25 +306,28 @@ class LocalSTCrossEncoder(CrossEncoderModel):
|
||||
"""
|
||||
import numpy as np
|
||||
|
||||
if self.bucket_batching and len(pairs) > 1:
|
||||
# Sort pairs by approximate token length to create homogeneous batches.
|
||||
# This eliminates padding waste — short pairs aren't padded to the length
|
||||
# of the longest pair in the batch. Quality-identical by construction.
|
||||
lengths = [len(pairs[i][0]) + len(pairs[i][1]) for i in range(len(pairs))]
|
||||
sorted_indices = sorted(range(len(pairs)), key=lambda i: lengths[i])
|
||||
sorted_pairs = [pairs[i] for i in sorted_indices]
|
||||
try:
|
||||
if self.bucket_batching and len(pairs) > 1:
|
||||
# Sort pairs by approximate token length to create homogeneous batches.
|
||||
# This eliminates padding waste — short pairs aren't padded to the length
|
||||
# of the longest pair in the batch. Quality-identical by construction.
|
||||
lengths = [len(pairs[i][0]) + len(pairs[i][1]) for i in range(len(pairs))]
|
||||
sorted_indices = sorted(range(len(pairs)), key=lambda i: lengths[i])
|
||||
sorted_pairs = [pairs[i] for i in sorted_indices]
|
||||
|
||||
sorted_scores = self._model.predict(sorted_pairs, batch_size=self.batch_size, show_progress_bar=False)
|
||||
sorted_scores = sorted_scores.tolist() if hasattr(sorted_scores, "tolist") else list(sorted_scores)
|
||||
sorted_scores = self._model.predict(sorted_pairs, batch_size=self.batch_size, show_progress_bar=False)
|
||||
sorted_scores = sorted_scores.tolist() if hasattr(sorted_scores, "tolist") else list(sorted_scores)
|
||||
|
||||
# Restore original order
|
||||
scores = [0.0] * len(pairs)
|
||||
for new_pos, orig_idx in enumerate(sorted_indices):
|
||||
scores[orig_idx] = sorted_scores[new_pos]
|
||||
return scores
|
||||
# Restore original order
|
||||
scores = [0.0] * len(pairs)
|
||||
for new_pos, orig_idx in enumerate(sorted_indices):
|
||||
scores[orig_idx] = sorted_scores[new_pos]
|
||||
return scores
|
||||
|
||||
scores = self._model.predict(pairs, batch_size=self.batch_size, show_progress_bar=False)
|
||||
return scores.tolist() if hasattr(scores, "tolist") else list(scores)
|
||||
scores = self._model.predict(pairs, batch_size=self.batch_size, show_progress_bar=False)
|
||||
return scores.tolist() if hasattr(scores, "tolist") else list(scores)
|
||||
finally:
|
||||
_malloc_trim()
|
||||
|
||||
async def predict(self, pairs: list[tuple[str, str]]) -> list[float]:
|
||||
"""
|
||||
@@ -546,12 +589,14 @@ class _CohereCompatibleRerankClient:
|
||||
rerank_url: str,
|
||||
timeout: float = 60.0,
|
||||
include_top_n: bool = True,
|
||||
include_return_documents: bool = False,
|
||||
):
|
||||
self.api_key = api_key
|
||||
self.model = model
|
||||
self.rerank_url = rerank_url
|
||||
self.timeout = timeout
|
||||
self.include_top_n = include_top_n
|
||||
self.include_return_documents = include_return_documents
|
||||
self._async_client: httpx.AsyncClient | None = None
|
||||
|
||||
async def initialize(self) -> None:
|
||||
@@ -729,7 +774,7 @@ class ZeroEntropyCrossEncoder(CrossEncoderModel):
|
||||
See: https://docs.zeroentropy.dev/models
|
||||
"""
|
||||
|
||||
DEFAULT_BASE_URL = "https://api.zeroentropy.dev"
|
||||
DEFAULT_BASE_URL = DEFAULT_ZEROENTROPY_BASE_URL
|
||||
RERANK_PATH = "/v1/models/rerank"
|
||||
|
||||
def __init__(
|
||||
@@ -962,32 +1007,35 @@ class FlashRankCrossEncoder(CrossEncoderModel):
|
||||
if not pairs:
|
||||
return []
|
||||
|
||||
# Group pairs by query
|
||||
query_groups: dict[str, list[tuple[int, str]]] = {}
|
||||
for idx, (query, text) in enumerate(pairs):
|
||||
if query not in query_groups:
|
||||
query_groups[query] = []
|
||||
query_groups[query].append((idx, text))
|
||||
try:
|
||||
# Group pairs by query
|
||||
query_groups: dict[str, list[tuple[int, str]]] = {}
|
||||
for idx, (query, text) in enumerate(pairs):
|
||||
if query not in query_groups:
|
||||
query_groups[query] = []
|
||||
query_groups[query].append((idx, text))
|
||||
|
||||
all_scores = [0.0] * len(pairs)
|
||||
all_scores = [0.0] * len(pairs)
|
||||
|
||||
for query, indexed_texts in query_groups.items():
|
||||
# Build passages list for FlashRank
|
||||
passages = [{"id": i, "text": text} for i, (_, text) in enumerate(indexed_texts)]
|
||||
global_indices = [idx for idx, _ in indexed_texts]
|
||||
for query, indexed_texts in query_groups.items():
|
||||
# Build passages list for FlashRank
|
||||
passages = [{"id": i, "text": text} for i, (_, text) in enumerate(indexed_texts)]
|
||||
global_indices = [idx for idx, _ in indexed_texts]
|
||||
|
||||
# Create rerank request
|
||||
request = RerankRequest(query=query, passages=passages)
|
||||
results = self._ranker.rerank(request)
|
||||
# Create rerank request
|
||||
request = RerankRequest(query=query, passages=passages)
|
||||
results = self._ranker.rerank(request)
|
||||
|
||||
# Map scores back to original positions
|
||||
for result in results:
|
||||
local_idx = result["id"]
|
||||
score = result["score"]
|
||||
global_idx = global_indices[local_idx]
|
||||
all_scores[global_idx] = score
|
||||
# Map scores back to original positions
|
||||
for result in results:
|
||||
local_idx = result["id"]
|
||||
score = result["score"]
|
||||
global_idx = global_indices[local_idx]
|
||||
all_scores[global_idx] = score
|
||||
|
||||
return all_scores
|
||||
return all_scores
|
||||
finally:
|
||||
_malloc_trim()
|
||||
|
||||
async def predict(self, pairs: list[tuple[str, str]]) -> list[float]:
|
||||
"""
|
||||
@@ -1534,6 +1582,48 @@ class GoogleCrossEncoder(CrossEncoderModel):
|
||||
return await loop.run_in_executor(None, self._predict_sync, pairs)
|
||||
|
||||
|
||||
class AlibabaCloudCrossEncoder(CrossEncoderModel):
|
||||
"""
|
||||
Alibaba Cloud DashScope text reranking API.
|
||||
|
||||
Uses the Cohere-compatible /reranks endpoint, which is the standard interface
|
||||
for qwen3-rerank. Authentication via HINDSIGHT_API_RERANKER_ALIBABA_API_KEY
|
||||
(or DASHSCOPE_API_KEY as a fallback).
|
||||
See: https://help.aliyun.com/zh/model-studio/text-rerank-api
|
||||
"""
|
||||
|
||||
RERANK_URL = "https://dashscope.aliyuncs.com/compatible-api/v1/reranks"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
api_key: str,
|
||||
model: str = DEFAULT_RERANKER_ALIBABA_MODEL,
|
||||
timeout: float = 60.0,
|
||||
):
|
||||
self.model = model
|
||||
self._client = _CohereCompatibleRerankClient(
|
||||
api_key=api_key,
|
||||
model=model,
|
||||
rerank_url=self.RERANK_URL,
|
||||
timeout=timeout,
|
||||
include_return_documents=False,
|
||||
)
|
||||
|
||||
@property
|
||||
def provider_name(self) -> str:
|
||||
return "alibaba"
|
||||
|
||||
async def initialize(self) -> None:
|
||||
if self._client._async_client is not None:
|
||||
return
|
||||
logger.info(f"Reranker: initializing Alibaba Cloud provider with model {self.model}")
|
||||
await self._client.initialize()
|
||||
logger.info("Reranker: Alibaba Cloud provider initialized")
|
||||
|
||||
async def predict(self, pairs: list[tuple[str, str]]) -> list[float]:
|
||||
return await self._client.predict(pairs)
|
||||
|
||||
|
||||
def create_cross_encoder_from_env() -> CrossEncoderModel:
|
||||
"""
|
||||
Create a CrossEncoderModel instance based on configuration.
|
||||
@@ -1624,6 +1714,7 @@ def create_cross_encoder_from_env() -> CrossEncoderModel:
|
||||
return ZeroEntropyCrossEncoder(
|
||||
api_key=api_key,
|
||||
model=config.reranker_zeroentropy_model,
|
||||
base_url=config.reranker_zeroentropy_base_url,
|
||||
)
|
||||
elif provider == "siliconflow":
|
||||
api_key = config.reranker_siliconflow_api_key
|
||||
@@ -1648,11 +1739,19 @@ def create_cross_encoder_from_env() -> CrossEncoderModel:
|
||||
model=config.reranker_google_model,
|
||||
service_account_key=config.reranker_google_service_account_key,
|
||||
)
|
||||
elif provider == "alibaba":
|
||||
api_key = config.reranker_alibaba_api_key
|
||||
if not api_key:
|
||||
raise ValueError(f"{ENV_RERANKER_ALIBABA_API_KEY} is required when {ENV_RERANKER_PROVIDER} is 'alibaba'")
|
||||
return AlibabaCloudCrossEncoder(
|
||||
api_key=api_key,
|
||||
model=config.reranker_alibaba_model,
|
||||
)
|
||||
elif provider == "rrf":
|
||||
return RRFPassthroughCrossEncoder()
|
||||
elif provider == "jina-mlx":
|
||||
return JinaMLXCrossEncoder()
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unknown reranker provider: {provider}. Supported: 'local', 'tei', 'cohere', 'zeroentropy', 'siliconflow', 'google', 'flashrank', 'litellm', 'litellm-sdk', 'rrf', 'jina-mlx'"
|
||||
f"Unknown reranker provider: {provider}. Supported: 'local', 'tei', 'cohere', 'zeroentropy', 'siliconflow', 'alibaba', 'google', 'flashrank', 'litellm', 'litellm-sdk', 'rrf', 'jina-mlx'"
|
||||
)
|
||||
|
||||
@@ -166,21 +166,6 @@ class DataAccessOps(ABC):
|
||||
|
||||
# -- LATERAL / fan-out queries ---------------------------------------
|
||||
|
||||
@abstractmethod
|
||||
async def fetch_entity_unit_fanout(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
ue_table: str,
|
||||
entity_id_list: list[UUID],
|
||||
limit_per_entity: int,
|
||||
) -> list[ResultRow]:
|
||||
"""Fetch unit_ids for a list of entities with per-entity row cap.
|
||||
|
||||
PG uses unnest + CROSS JOIN LATERAL with LIMIT.
|
||||
Non-PG queries each entity individually.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def fetch_unit_dates(
|
||||
self,
|
||||
@@ -406,6 +391,74 @@ class DataAccessOps(ABC):
|
||||
"""Insert a webhook delivery task into async_operations."""
|
||||
...
|
||||
|
||||
# -- Graph maintenance queue -----------------------------------------
|
||||
|
||||
@abstractmethod
|
||||
async def enqueue_graph_maintenance(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
bank_id: str,
|
||||
unit_ids: list,
|
||||
) -> None:
|
||||
"""Insert unit_ids into graph_maintenance_queue, deduplicating on the
|
||||
(bank_id, unit_id) primary key.
|
||||
|
||||
Called inside the triggering transaction so enqueue is atomic with
|
||||
the mutation that caused it. Order is unspecified.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def claim_graph_maintenance_batch(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
bank_id: str,
|
||||
limit: int,
|
||||
) -> list[str]:
|
||||
"""Atomically claim a batch of rows from graph_maintenance_queue and
|
||||
remove them from the table.
|
||||
|
||||
Returns the list of ``unit_id`` strings. Empty list when the queue
|
||||
for ``bank_id`` is drained.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def prune_orphan_entities(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
entities_table: str,
|
||||
ue_table: str,
|
||||
bank_id: str,
|
||||
) -> int:
|
||||
"""Delete entities in ``bank_id`` that no longer have any unit_entities
|
||||
rows referencing them. Returns the number of rows deleted.
|
||||
|
||||
FK ON DELETE CASCADE on entity_cooccurrences then removes any
|
||||
cooccurrence row pointing at the pruned entities.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def prune_stale_cooccurrences(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
ec_table: str,
|
||||
ue_table: str,
|
||||
entities_table: str,
|
||||
bank_id: str,
|
||||
) -> int:
|
||||
"""Delete entity_cooccurrences rows in ``bank_id`` where the two
|
||||
entities still exist but no current unit references both of them.
|
||||
|
||||
These are stale-count rows: cooccurrence was real at the time it was
|
||||
recorded, but every memory_unit that witnessed both entities has
|
||||
since been deleted. Returns the number of rows deleted.
|
||||
"""
|
||||
...
|
||||
|
||||
# -- Task claiming operations ------------------------------------------
|
||||
|
||||
@abstractmethod
|
||||
|
||||
@@ -215,29 +215,96 @@ class OracleOps(DataAccessOps):
|
||||
list(zip(unit_ids, entity_ids)),
|
||||
)
|
||||
|
||||
async def fetch_entity_unit_fanout(
|
||||
async def enqueue_graph_maintenance(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
ue_table: str,
|
||||
entity_id_list: list[UUID],
|
||||
limit_per_entity: int,
|
||||
) -> list[ResultRow]:
|
||||
# Query each entity individually
|
||||
rows: list[ResultRow] = []
|
||||
for eid in entity_id_list:
|
||||
entity_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT $1 AS entity_id, ue.unit_id
|
||||
FROM {ue_table} ue
|
||||
WHERE ue.entity_id = $1
|
||||
ORDER BY ue.unit_id DESC
|
||||
LIMIT $2
|
||||
""",
|
||||
eid,
|
||||
limit_per_entity,
|
||||
table: str,
|
||||
bank_id: str,
|
||||
unit_ids: list,
|
||||
) -> None:
|
||||
if not unit_ids:
|
||||
return
|
||||
# Oracle doesn't support ON CONFLICT; rely on the PK and the
|
||||
# IGNORE_ROW_ON_DUPKEY_INDEX hint to skip duplicates server-side.
|
||||
# The hint name must match the PK constraint exactly.
|
||||
await conn.executemany(
|
||||
f"""
|
||||
INSERT /*+ IGNORE_ROW_ON_DUPKEY_INDEX({table}, pk_graph_maintenance_queue) */
|
||||
INTO {table} (bank_id, unit_id)
|
||||
VALUES ($1, $2)
|
||||
""",
|
||||
[(bank_id, uid) for uid in unit_ids],
|
||||
)
|
||||
|
||||
async def claim_graph_maintenance_batch(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
bank_id: str,
|
||||
limit: int,
|
||||
) -> list[str]:
|
||||
# Two-step claim: select the batch, then delete by exact keys. Oracle's
|
||||
# DELETE ... RETURNING doesn't accept a multi-row subquery, so we can't
|
||||
# do it in one statement like the PG version.
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT unit_id FROM {table}
|
||||
WHERE bank_id = $1
|
||||
ORDER BY enqueued_at
|
||||
FETCH FIRST $2 ROWS ONLY
|
||||
""",
|
||||
bank_id,
|
||||
limit,
|
||||
)
|
||||
claimed = [str(row["unit_id"]) for row in rows]
|
||||
if claimed:
|
||||
await conn.executemany(
|
||||
f"DELETE FROM {table} WHERE bank_id = $1 AND unit_id = $2",
|
||||
[(bank_id, uid) for uid in claimed],
|
||||
)
|
||||
rows.extend(entity_rows)
|
||||
return rows
|
||||
return claimed
|
||||
|
||||
async def prune_orphan_entities(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
entities_table: str,
|
||||
ue_table: str,
|
||||
bank_id: str,
|
||||
) -> int:
|
||||
# The Oracle DatabaseConnection wrapper reshapes ``cursor.rowcount`` into
|
||||
# the same ``"DELETE N"`` status string asyncpg returns, so the same
|
||||
# ``int(deleted.split()[-1])`` parsing works on both dialects.
|
||||
deleted = await conn.execute(
|
||||
f"""
|
||||
DELETE FROM {entities_table}
|
||||
WHERE bank_id = $1
|
||||
AND id NOT IN (SELECT DISTINCT entity_id FROM {ue_table})
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
return int(deleted.split()[-1]) if isinstance(deleted, str) and deleted.startswith("DELETE") else 0
|
||||
|
||||
async def prune_stale_cooccurrences(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
ec_table: str,
|
||||
ue_table: str,
|
||||
entities_table: str,
|
||||
bank_id: str,
|
||||
) -> int:
|
||||
deleted = await conn.execute(
|
||||
f"""
|
||||
DELETE FROM {ec_table}
|
||||
WHERE entity_id_1 IN (SELECT id FROM {entities_table} WHERE bank_id = $1)
|
||||
AND (entity_id_1, entity_id_2) NOT IN (
|
||||
SELECT u1.entity_id, u2.entity_id
|
||||
FROM {ue_table} u1
|
||||
JOIN {ue_table} u2 ON u1.unit_id = u2.unit_id
|
||||
)
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
return int(deleted.split()[-1]) if isinstance(deleted, str) and deleted.startswith("DELETE") else 0
|
||||
|
||||
async def fetch_unit_dates(
|
||||
self,
|
||||
|
||||
@@ -104,7 +104,46 @@ class PostgreSQLOps(DataAccessOps):
|
||||
FROM input_data
|
||||
RETURNING id
|
||||
"""
|
||||
elif config.text_search_extension == "native":
|
||||
# search_vector is a regular tsvector column populated here using the
|
||||
# configured native dictionary. It used to be GENERATED ALWAYS with
|
||||
# a hardcoded 'english', which prevented per-deployment language
|
||||
# configuration. text_search_extension_native_language is validated
|
||||
# in HindsightConfig.validate() as a PG identifier, so embedding it
|
||||
# as a SQL literal is safe.
|
||||
query = f"""
|
||||
WITH input_data AS (
|
||||
SELECT * FROM unnest(
|
||||
$2::text[], $3::vector[], $4::timestamptz[], $5::timestamptz[], $6::timestamptz[], $7::timestamptz[],
|
||||
$8::text[], $9::text[], $10::jsonb[], $11::text[], $12::text[], $13::jsonb[], $14::jsonb[], $15::text[]
|
||||
) AS t(text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
context, fact_type, metadata, chunk_id, document_id, tags_json,
|
||||
observation_scopes_json, text_signals)
|
||||
)
|
||||
INSERT INTO {table} (bank_id, text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
context, fact_type, metadata, chunk_id, document_id, tags,
|
||||
observation_scopes, text_signals, search_vector)
|
||||
SELECT
|
||||
$1,
|
||||
text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
context, fact_type, metadata, chunk_id, document_id,
|
||||
COALESCE(
|
||||
(SELECT array_agg(elem) FROM jsonb_array_elements_text(tags_json) AS elem),
|
||||
'{{}}'::varchar[]
|
||||
),
|
||||
observation_scopes_json,
|
||||
text_signals,
|
||||
to_tsvector(
|
||||
'{config.text_search_extension_native_language}'::regconfig,
|
||||
COALESCE(text, '') || ' ' || COALESCE(context, '') || ' ' || COALESCE(text_signals, '')
|
||||
)
|
||||
FROM input_data
|
||||
RETURNING id
|
||||
"""
|
||||
else:
|
||||
# pg_textsearch, pgroonga, and pg_search: search_vector is a dummy
|
||||
# TEXT column; the actual full-text index operates on the base text
|
||||
# columns directly, so we don't populate search_vector at insert time.
|
||||
query = f"""
|
||||
WITH input_data AS (
|
||||
SELECT * FROM unnest(
|
||||
@@ -251,29 +290,101 @@ class PostgreSQLOps(DataAccessOps):
|
||||
entity_ids,
|
||||
)
|
||||
|
||||
async def fetch_entity_unit_fanout(
|
||||
async def enqueue_graph_maintenance(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
ue_table: str,
|
||||
entity_id_list: list[UUID],
|
||||
limit_per_entity: int,
|
||||
) -> list[ResultRow]:
|
||||
return await conn.fetch(
|
||||
table: str,
|
||||
bank_id: str,
|
||||
unit_ids: list,
|
||||
) -> None:
|
||||
if not unit_ids:
|
||||
return
|
||||
await conn.execute(
|
||||
f"""
|
||||
SELECT e.entity_id, n.unit_id
|
||||
FROM unnest($1::uuid[]) AS e(entity_id)
|
||||
CROSS JOIN LATERAL (
|
||||
SELECT ue.unit_id
|
||||
FROM {ue_table} ue
|
||||
WHERE ue.entity_id = e.entity_id
|
||||
ORDER BY ue.unit_id DESC
|
||||
LIMIT $2
|
||||
) n
|
||||
INSERT INTO {table} (bank_id, unit_id)
|
||||
SELECT $1, v FROM unnest($2::uuid[]) AS t(v)
|
||||
ON CONFLICT (bank_id, unit_id) DO NOTHING
|
||||
""",
|
||||
entity_id_list,
|
||||
limit_per_entity,
|
||||
bank_id,
|
||||
unit_ids,
|
||||
)
|
||||
|
||||
async def claim_graph_maintenance_batch(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
bank_id: str,
|
||||
limit: int,
|
||||
) -> list[str]:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
DELETE FROM {table}
|
||||
WHERE (bank_id, unit_id) IN (
|
||||
SELECT bank_id, unit_id FROM {table}
|
||||
WHERE bank_id = $1
|
||||
ORDER BY enqueued_at
|
||||
LIMIT $2
|
||||
)
|
||||
RETURNING unit_id
|
||||
""",
|
||||
bank_id,
|
||||
limit,
|
||||
)
|
||||
return [str(row["unit_id"]) for row in rows]
|
||||
|
||||
async def prune_orphan_entities(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
entities_table: str,
|
||||
ue_table: str,
|
||||
bank_id: str,
|
||||
) -> int:
|
||||
# Scoped by entities.bank_id (indexed). The NOT EXISTS subquery is
|
||||
# backed by idx_ue_entity on unit_entities(entity_id), so this stays
|
||||
# linear in the number of entities in the bank — not in the size of
|
||||
# unit_entities globally.
|
||||
result = await conn.execute(
|
||||
f"""
|
||||
DELETE FROM {entities_table} e
|
||||
WHERE e.bank_id = $1
|
||||
AND NOT EXISTS (
|
||||
SELECT 1 FROM {ue_table} ue WHERE ue.entity_id = e.id
|
||||
)
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
# asyncpg returns "DELETE N"
|
||||
return int(result.split()[-1]) if isinstance(result, str) and result.startswith("DELETE") else 0
|
||||
|
||||
async def prune_stale_cooccurrences(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
ec_table: str,
|
||||
ue_table: str,
|
||||
entities_table: str,
|
||||
bank_id: str,
|
||||
) -> int:
|
||||
# Scope by joining through entities.bank_id (entity_cooccurrences itself
|
||||
# has no bank_id column — entities don't span banks, so scoping via
|
||||
# entity_id_1 is sufficient).
|
||||
result = await conn.execute(
|
||||
f"""
|
||||
DELETE FROM {ec_table} c
|
||||
USING {entities_table} e
|
||||
WHERE e.id = c.entity_id_1
|
||||
AND e.bank_id = $1
|
||||
AND NOT EXISTS (
|
||||
SELECT 1
|
||||
FROM {ue_table} u1
|
||||
JOIN {ue_table} u2 ON u1.unit_id = u2.unit_id
|
||||
WHERE u1.entity_id = c.entity_id_1
|
||||
AND u2.entity_id = c.entity_id_2
|
||||
)
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
return int(result.split()[-1]) if isinstance(result, str) and result.startswith("DELETE") else 0
|
||||
|
||||
async def fetch_unit_dates(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
|
||||
@@ -9,13 +9,19 @@ The database schema is automatically adjusted to match the model's dimension.
|
||||
Configuration via environment variables - see hindsight_api.config for all env var names.
|
||||
"""
|
||||
|
||||
import base64
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import struct
|
||||
import warnings
|
||||
from abc import ABC, abstractmethod
|
||||
from pathlib import Path
|
||||
from typing import Literal, cast
|
||||
from urllib.parse import parse_qs, urlparse, urlunparse
|
||||
|
||||
import httpx
|
||||
from pydantic import BaseModel
|
||||
|
||||
from ..config import (
|
||||
DEFAULT_EMBEDDINGS_COHERE_MODEL,
|
||||
@@ -27,10 +33,15 @@ from ..config import (
|
||||
DEFAULT_EMBEDDINGS_LOCAL_TRUST_REMOTE_CODE,
|
||||
DEFAULT_EMBEDDINGS_OPENAI_MODEL,
|
||||
DEFAULT_EMBEDDINGS_PROVIDER,
|
||||
DEFAULT_EMBEDDINGS_ZEROENTROPY_BATCH_SIZE,
|
||||
DEFAULT_EMBEDDINGS_ZEROENTROPY_DIMENSIONS,
|
||||
DEFAULT_EMBEDDINGS_ZEROENTROPY_ENCODING_FORMAT,
|
||||
DEFAULT_EMBEDDINGS_ZEROENTROPY_LATENCY,
|
||||
DEFAULT_EMBEDDINGS_ZEROENTROPY_MODEL,
|
||||
DEFAULT_LITELLM_API_BASE,
|
||||
DEFAULT_ZEROENTROPY_BASE_URL,
|
||||
ENV_EMBEDDINGS_COHERE_API_KEY,
|
||||
ENV_EMBEDDINGS_GEMINI_API_KEY,
|
||||
ENV_EMBEDDINGS_LITELLM_SDK_API_KEY,
|
||||
ENV_EMBEDDINGS_LOCAL_FORCE_CPU,
|
||||
ENV_EMBEDDINGS_LOCAL_MODEL,
|
||||
ENV_EMBEDDINGS_LOCAL_TRUST_REMOTE_CODE,
|
||||
@@ -39,12 +50,39 @@ from ..config import (
|
||||
ENV_EMBEDDINGS_OPENAI_MODEL,
|
||||
ENV_EMBEDDINGS_PROVIDER,
|
||||
ENV_EMBEDDINGS_TEI_URL,
|
||||
ENV_EMBEDDINGS_ZEROENTROPY_API_KEY,
|
||||
ENV_EMBEDDINGS_ZEROENTROPY_DIMENSIONS,
|
||||
ENV_EMBEDDINGS_ZEROENTROPY_ENCODING_FORMAT,
|
||||
ENV_LLM_API_KEY,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
ZeroEntropyInputType = Literal["document", "query"]
|
||||
ZeroEntropyLatency = Literal["fast", "slow"]
|
||||
ZeroEntropyEncodingFormat = Literal["float", "base64"]
|
||||
|
||||
|
||||
class _ZeroEntropyEmbedRequest(BaseModel):
|
||||
"""Typed request body for ZeroEntropy's non-OpenAI-compatible embed endpoint."""
|
||||
|
||||
model: str
|
||||
input: list[str]
|
||||
input_type: ZeroEntropyInputType
|
||||
dimensions: int
|
||||
encoding_format: ZeroEntropyEncodingFormat = "float"
|
||||
latency: ZeroEntropyLatency | None = None
|
||||
|
||||
|
||||
class _ZeroEntropyEmbedResult(BaseModel):
|
||||
embedding: list[float] | str
|
||||
|
||||
|
||||
class _ZeroEntropyEmbedResponse(BaseModel):
|
||||
results: list[_ZeroEntropyEmbedResult]
|
||||
|
||||
|
||||
class Embeddings(ABC):
|
||||
"""
|
||||
Abstract base class for embedding generation.
|
||||
@@ -88,6 +126,14 @@ class Embeddings(ABC):
|
||||
"""
|
||||
pass
|
||||
|
||||
def encode_query(self, texts: list[str]) -> list[list[float]]:
|
||||
"""Generate embeddings for query text. Providers without asymmetric embeddings use encode()."""
|
||||
return self.encode(texts)
|
||||
|
||||
def encode_documents(self, texts: list[str]) -> list[list[float]]:
|
||||
"""Generate embeddings for stored document text. Providers without asymmetric embeddings use encode()."""
|
||||
return self.encode(texts)
|
||||
|
||||
|
||||
class LocalSTEmbeddings(Embeddings):
|
||||
"""
|
||||
@@ -385,6 +431,7 @@ class OpenAIEmbeddings(Embeddings):
|
||||
model: str = DEFAULT_EMBEDDINGS_OPENAI_MODEL,
|
||||
base_url: str | None = None,
|
||||
batch_size: int = 100,
|
||||
dimensions: int | None = None,
|
||||
max_retries: int = 3,
|
||||
):
|
||||
"""
|
||||
@@ -395,12 +442,14 @@ class OpenAIEmbeddings(Embeddings):
|
||||
model: OpenAI embedding model name (default: text-embedding-3-small)
|
||||
base_url: Custom base URL for OpenAI-compatible API (e.g., Azure OpenAI endpoint)
|
||||
batch_size: Maximum batch size for embedding requests (default: 100)
|
||||
dimensions: Optional requested output dimensions for OpenAI text-embedding-3 models
|
||||
max_retries: Maximum number of retries for failed requests (default: 3)
|
||||
"""
|
||||
self.api_key = api_key
|
||||
self.model = model
|
||||
self.base_url = base_url
|
||||
self.batch_size = batch_size
|
||||
self.dimensions = dimensions
|
||||
self.max_retries = max_retries
|
||||
self._client = None
|
||||
self._dimension: int | None = None
|
||||
@@ -445,7 +494,9 @@ class OpenAIEmbeddings(Embeddings):
|
||||
self._client = OpenAI(**client_kwargs)
|
||||
|
||||
# Try to get dimension from known models, otherwise do a test embedding
|
||||
if self.model in self.MODEL_DIMENSIONS:
|
||||
if self.dimensions is not None:
|
||||
self._dimension = self.dimensions
|
||||
elif self.model in self.MODEL_DIMENSIONS:
|
||||
self._dimension = self.MODEL_DIMENSIONS[self.model]
|
||||
else:
|
||||
# Do a test embedding to detect dimension
|
||||
@@ -480,10 +531,14 @@ class OpenAIEmbeddings(Embeddings):
|
||||
for i in range(0, len(texts), self.batch_size):
|
||||
batch = texts[i : i + self.batch_size]
|
||||
|
||||
response = self._client.embeddings.create(
|
||||
model=self.model,
|
||||
input=batch,
|
||||
)
|
||||
request = {
|
||||
"model": self.model,
|
||||
"input": batch,
|
||||
}
|
||||
if self.dimensions is not None:
|
||||
request["dimensions"] = self.dimensions
|
||||
|
||||
response = self._client.embeddings.create(**request)
|
||||
|
||||
# Sort by index to ensure correct order
|
||||
batch_embeddings = sorted(response.data, key=lambda x: x.index)
|
||||
@@ -492,6 +547,58 @@ class OpenAIEmbeddings(Embeddings):
|
||||
return all_embeddings
|
||||
|
||||
|
||||
class CodexOAuthEmbeddings(OpenAIEmbeddings):
|
||||
"""
|
||||
OpenAI embeddings using the Codex/ChatGPT OAuth token from ``~/.codex/auth.json``.
|
||||
|
||||
Codex OAuth is an LLM-provider auth path in Hindsight, but the same bearer token
|
||||
can also authenticate against the standard OpenAI embeddings endpoint. This keeps
|
||||
embeddings on the user's existing Codex subscription/OAuth path without requiring
|
||||
a separate OpenAI/OpenRouter/Gemini/Cohere API key.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model: str = DEFAULT_EMBEDDINGS_OPENAI_MODEL,
|
||||
batch_size: int = 100,
|
||||
dimensions: int | None = None,
|
||||
max_retries: int = 3,
|
||||
):
|
||||
access_token = self._load_codex_access_token()
|
||||
super().__init__(
|
||||
api_key=access_token,
|
||||
model=model,
|
||||
base_url="https://api.openai.com/v1",
|
||||
batch_size=batch_size,
|
||||
dimensions=dimensions,
|
||||
max_retries=max_retries,
|
||||
)
|
||||
|
||||
@property
|
||||
def provider_name(self) -> str:
|
||||
return "openai-codex"
|
||||
|
||||
@staticmethod
|
||||
def _load_codex_access_token() -> str:
|
||||
"""Load the Codex OAuth access token without logging or exposing it."""
|
||||
auth_file = Path.home() / ".codex" / "auth.json"
|
||||
if not auth_file.exists():
|
||||
raise FileNotFoundError(f"Codex auth file not found: {auth_file}. Run 'codex auth login' to authenticate.")
|
||||
|
||||
with open(auth_file) as f:
|
||||
data = json.load(f)
|
||||
|
||||
auth_mode = data.get("auth_mode")
|
||||
if auth_mode != "chatgpt":
|
||||
raise ValueError(f"Expected Codex auth_mode='chatgpt', got: {auth_mode}")
|
||||
|
||||
access_token = (data.get("tokens") or {}).get("access_token")
|
||||
if not access_token:
|
||||
raise ValueError("No access_token found in Codex auth file. Run 'codex auth login' again.")
|
||||
|
||||
return access_token
|
||||
|
||||
|
||||
class CohereEmbeddings(Embeddings):
|
||||
"""
|
||||
Cohere embeddings implementation using the Cohere API.
|
||||
@@ -633,6 +740,149 @@ class CohereEmbeddings(Embeddings):
|
||||
return all_embeddings
|
||||
|
||||
|
||||
class ZeroEntropyEmbeddings(Embeddings):
|
||||
"""
|
||||
ZeroEntropy embeddings implementation using the zembed API.
|
||||
|
||||
ZeroEntropy's embeddings endpoint is not OpenAI-compatible: it lives at
|
||||
/v1/models/embed and requires provider-specific parameters such as
|
||||
input_type. Hindsight stores document-side vectors and uses query-side
|
||||
vectors during recall, so this provider exposes explicit encode_documents()
|
||||
and encode_query() helpers while keeping encode() as document-side default.
|
||||
"""
|
||||
|
||||
VALID_DIMENSIONS = frozenset({2560, 1280, 640, 320, 160, 80, 40})
|
||||
VALID_ENCODING_FORMATS = frozenset({"float", "base64"})
|
||||
VALID_LATENCIES = frozenset({"fast", "slow"})
|
||||
DEFAULT_BASE_URL = DEFAULT_ZEROENTROPY_BASE_URL
|
||||
EMBED_PATH = "/v1/models/embed"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
api_key: str,
|
||||
model: str = DEFAULT_EMBEDDINGS_ZEROENTROPY_MODEL,
|
||||
base_url: str | None = None,
|
||||
dimensions: int = DEFAULT_EMBEDDINGS_ZEROENTROPY_DIMENSIONS,
|
||||
batch_size: int = DEFAULT_EMBEDDINGS_ZEROENTROPY_BATCH_SIZE,
|
||||
encoding_format: str = DEFAULT_EMBEDDINGS_ZEROENTROPY_ENCODING_FORMAT,
|
||||
latency: str | None = DEFAULT_EMBEDDINGS_ZEROENTROPY_LATENCY,
|
||||
timeout: float = 60.0,
|
||||
):
|
||||
if dimensions not in self.VALID_DIMENSIONS:
|
||||
valid = ", ".join(str(dim) for dim in sorted(self.VALID_DIMENSIONS, reverse=True))
|
||||
raise ValueError(f"{ENV_EMBEDDINGS_ZEROENTROPY_DIMENSIONS} must be one of {valid}, got {dimensions}")
|
||||
if batch_size < 1:
|
||||
raise ValueError("ZeroEntropy embeddings batch_size must be >= 1")
|
||||
if encoding_format not in self.VALID_ENCODING_FORMATS:
|
||||
valid_formats = ", ".join(sorted(self.VALID_ENCODING_FORMATS))
|
||||
raise ValueError(
|
||||
f"{ENV_EMBEDDINGS_ZEROENTROPY_ENCODING_FORMAT} must be one of {valid_formats}, got {encoding_format!r}"
|
||||
)
|
||||
if latency is not None and latency not in self.VALID_LATENCIES:
|
||||
valid_latencies = ", ".join(sorted(self.VALID_LATENCIES))
|
||||
raise ValueError(f"ZeroEntropy embeddings latency must be one of {valid_latencies}, got {latency!r}")
|
||||
|
||||
self.api_key = api_key
|
||||
self.model = model
|
||||
self.base_url = base_url.rstrip("/") if base_url else self.DEFAULT_BASE_URL
|
||||
self.embed_url = f"{self.base_url}{self.EMBED_PATH}"
|
||||
self.dimensions = dimensions
|
||||
self.batch_size = batch_size
|
||||
self.encoding_format = cast(ZeroEntropyEncodingFormat, encoding_format)
|
||||
self.latency = cast(ZeroEntropyLatency | None, latency)
|
||||
self.timeout = timeout
|
||||
self._client: httpx.Client | None = None
|
||||
self._dimension: int | None = None
|
||||
|
||||
@property
|
||||
def provider_name(self) -> str:
|
||||
return "zeroentropy"
|
||||
|
||||
@property
|
||||
def dimension(self) -> int:
|
||||
if self._dimension is None:
|
||||
raise RuntimeError("Embeddings not initialized. Call initialize() first.")
|
||||
return self._dimension
|
||||
|
||||
async def initialize(self) -> None:
|
||||
"""Initialize the ZeroEntropy HTTP client."""
|
||||
if self._client is not None:
|
||||
return
|
||||
|
||||
logger.info(
|
||||
f"Embeddings: initializing ZeroEntropy provider with model {self.model} "
|
||||
f"(dim: {self.dimensions}, batch_size={self.batch_size})"
|
||||
)
|
||||
self._client = httpx.Client(
|
||||
timeout=self.timeout,
|
||||
headers={
|
||||
"Authorization": f"Bearer {self.api_key}",
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
)
|
||||
# zembed-1 dimensions are explicit Matryoshka truncation steps. Avoid a
|
||||
# startup probe so boot does not burn quota or require a throwaway input.
|
||||
self._dimension = self.dimensions
|
||||
logger.info(f"Embeddings: ZeroEntropy provider initialized (model: {self.model}, dim: {self._dimension})")
|
||||
|
||||
def encode(self, texts: list[str]) -> list[list[float]]:
|
||||
"""Generate document-side embeddings for backwards-compatible callers."""
|
||||
return self.encode_documents(texts)
|
||||
|
||||
def encode_documents(self, texts: list[str]) -> list[list[float]]:
|
||||
"""Generate document-side embeddings for retained content."""
|
||||
return self._encode_with_input_type(texts, "document")
|
||||
|
||||
def encode_query(self, texts: list[str]) -> list[list[float]]:
|
||||
"""Generate query-side embeddings for recall/search queries."""
|
||||
return self._encode_with_input_type(texts, "query")
|
||||
|
||||
def _encode_with_input_type(self, texts: list[str], input_type: ZeroEntropyInputType) -> list[list[float]]:
|
||||
if self._client is None:
|
||||
raise RuntimeError("Embeddings not initialized. Call initialize() first.")
|
||||
|
||||
if not texts:
|
||||
return []
|
||||
|
||||
all_embeddings: list[list[float]] = []
|
||||
for i in range(0, len(texts), self.batch_size):
|
||||
batch = texts[i : i + self.batch_size]
|
||||
request = _ZeroEntropyEmbedRequest(
|
||||
model=self.model,
|
||||
input=batch,
|
||||
input_type=input_type,
|
||||
dimensions=self.dimensions,
|
||||
encoding_format=self.encoding_format,
|
||||
latency=self.latency,
|
||||
)
|
||||
|
||||
try:
|
||||
response = self._client.post(self.embed_url, json=request.model_dump(exclude_none=True))
|
||||
response.raise_for_status()
|
||||
except httpx.HTTPError as e:
|
||||
raise RuntimeError(f"ZeroEntropy embedding request failed: {e}") from e
|
||||
|
||||
parsed = _ZeroEntropyEmbedResponse.model_validate(response.json())
|
||||
if len(parsed.results) != len(batch):
|
||||
raise RuntimeError(
|
||||
f"ZeroEntropy returned {len(parsed.results)} embeddings for {len(batch)} input texts; "
|
||||
"expected exact 1:1 alignment"
|
||||
)
|
||||
all_embeddings.extend(self._parse_embedding(result.embedding) for result in parsed.results)
|
||||
|
||||
return all_embeddings
|
||||
|
||||
@staticmethod
|
||||
def _parse_embedding(embedding: list[float] | str) -> list[float]:
|
||||
if not isinstance(embedding, str):
|
||||
return embedding
|
||||
|
||||
raw = base64.b64decode(embedding)
|
||||
if len(raw) % 4 != 0:
|
||||
raise RuntimeError("ZeroEntropy returned invalid base64 embedding length")
|
||||
return list(struct.unpack(f"<{len(raw) // 4}f", raw))
|
||||
|
||||
|
||||
class LiteLLMEmbeddings(Embeddings):
|
||||
"""
|
||||
LiteLLM embeddings implementation using LiteLLM proxy's /embeddings endpoint.
|
||||
@@ -766,7 +1016,7 @@ class LiteLLMSDKEmbeddings(Embeddings):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
api_key: str,
|
||||
api_key: str | None = None,
|
||||
model: str = DEFAULT_EMBEDDINGS_LITELLM_SDK_MODEL,
|
||||
api_base: str | None = None,
|
||||
output_dimensions: int | None = None,
|
||||
@@ -778,7 +1028,8 @@ class LiteLLMSDKEmbeddings(Embeddings):
|
||||
Initialize LiteLLM SDK embeddings client.
|
||||
|
||||
Args:
|
||||
api_key: API key for the embedding provider
|
||||
api_key: API key for the embedding provider (optional — omit for
|
||||
providers that use ambient credentials, e.g. AWS Bedrock with IAM)
|
||||
model: Model name with provider prefix (e.g., "cohere/embed-english-v3.0")
|
||||
api_base: Custom base URL for API (optional)
|
||||
output_dimensions: Optional output embedding dimensions (provider-dependent)
|
||||
@@ -828,8 +1079,9 @@ class LiteLLMSDKEmbeddings(Embeddings):
|
||||
embed_kwargs = {
|
||||
"model": self.model,
|
||||
"input": ["test"],
|
||||
"api_key": self.api_key,
|
||||
}
|
||||
if self.api_key:
|
||||
embed_kwargs["api_key"] = self.api_key
|
||||
if self.encoding_format:
|
||||
embed_kwargs["encoding_format"] = self.encoding_format
|
||||
if self.api_base:
|
||||
@@ -880,8 +1132,9 @@ class LiteLLMSDKEmbeddings(Embeddings):
|
||||
embed_kwargs = {
|
||||
"model": self.model,
|
||||
"input": batch,
|
||||
"api_key": self.api_key,
|
||||
}
|
||||
if self.api_key:
|
||||
embed_kwargs["api_key"] = self.api_key
|
||||
if self.encoding_format:
|
||||
embed_kwargs["encoding_format"] = self.encoding_format
|
||||
if self.api_base:
|
||||
@@ -1140,6 +1393,14 @@ def create_embeddings_from_env() -> Embeddings:
|
||||
model=model,
|
||||
base_url=base_url,
|
||||
batch_size=config.embeddings_openai_batch_size,
|
||||
dimensions=config.embeddings_openai_dimensions,
|
||||
)
|
||||
elif provider == "openai-codex":
|
||||
model = os.environ.get(ENV_EMBEDDINGS_OPENAI_MODEL, DEFAULT_EMBEDDINGS_OPENAI_MODEL)
|
||||
return CodexOAuthEmbeddings(
|
||||
model=model,
|
||||
batch_size=config.embeddings_openai_batch_size,
|
||||
dimensions=config.embeddings_openai_dimensions,
|
||||
)
|
||||
elif provider == "openrouter":
|
||||
api_key = config.embeddings_openrouter_api_key
|
||||
@@ -1154,6 +1415,22 @@ def create_embeddings_from_env() -> Embeddings:
|
||||
base_url="https://openrouter.ai/api/v1",
|
||||
batch_size=config.embeddings_openai_batch_size,
|
||||
)
|
||||
elif provider == "zeroentropy":
|
||||
api_key = config.embeddings_zeroentropy_api_key
|
||||
if not api_key:
|
||||
raise ValueError(
|
||||
f"{ENV_EMBEDDINGS_ZEROENTROPY_API_KEY} or ZEROENTROPY_API_KEY is required "
|
||||
f"when {ENV_EMBEDDINGS_PROVIDER} is 'zeroentropy'"
|
||||
)
|
||||
return ZeroEntropyEmbeddings(
|
||||
api_key=api_key,
|
||||
model=config.embeddings_zeroentropy_model,
|
||||
base_url=config.embeddings_zeroentropy_base_url,
|
||||
dimensions=config.embeddings_zeroentropy_dimensions,
|
||||
batch_size=config.embeddings_zeroentropy_batch_size,
|
||||
encoding_format=config.embeddings_zeroentropy_encoding_format,
|
||||
latency=config.embeddings_zeroentropy_latency,
|
||||
)
|
||||
elif provider == "cohere":
|
||||
api_key = config.embeddings_cohere_api_key
|
||||
if not api_key:
|
||||
@@ -1171,13 +1448,8 @@ def create_embeddings_from_env() -> Embeddings:
|
||||
model=config.embeddings_litellm_model,
|
||||
)
|
||||
elif provider == "litellm-sdk":
|
||||
api_key = config.embeddings_litellm_sdk_api_key
|
||||
if not api_key:
|
||||
raise ValueError(
|
||||
f"{ENV_EMBEDDINGS_LITELLM_SDK_API_KEY} is required when {ENV_EMBEDDINGS_PROVIDER} is 'litellm-sdk'"
|
||||
)
|
||||
return LiteLLMSDKEmbeddings(
|
||||
api_key=api_key,
|
||||
api_key=config.embeddings_litellm_sdk_api_key or None,
|
||||
model=config.embeddings_litellm_sdk_model,
|
||||
api_base=config.embeddings_litellm_sdk_api_base,
|
||||
output_dimensions=config.embeddings_litellm_sdk_output_dimensions,
|
||||
@@ -1206,5 +1478,6 @@ def create_embeddings_from_env() -> Embeddings:
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unknown embeddings provider: {provider}. "
|
||||
f"Supported: 'local', 'tei', 'openai', 'cohere', 'google', 'litellm', 'litellm-sdk'"
|
||||
f"Supported: 'local', 'tei', 'openai', 'openai-codex', 'openrouter', 'cohere', 'google', "
|
||||
f"'zeroentropy', 'litellm', 'litellm-sdk'"
|
||||
)
|
||||
|
||||
@@ -16,7 +16,15 @@ from typing import Any, Final
|
||||
|
||||
from .db_utils import acquire_with_retry
|
||||
from .memory_engine import fq_table
|
||||
from .retain.entity_labels import build_labels_lookup as _build_labels_lookup_from_config
|
||||
from .retain.entity_labels import (
|
||||
build_labels_lookup as _build_labels_lookup_from_config,
|
||||
)
|
||||
from .retain.entity_labels import (
|
||||
is_label_entity as _is_label_entity,
|
||||
)
|
||||
from .retain.entity_labels import (
|
||||
parse_entity_labels as _parse_entity_labels,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -228,14 +236,15 @@ class EntityResolver:
|
||||
return []
|
||||
|
||||
taxonomy_lookup = self._build_labels_lookup(entity_labels)
|
||||
labels_cfg = _parse_entity_labels(entity_labels)
|
||||
if conn is None:
|
||||
async with acquire_with_retry(self.pool) as conn:
|
||||
return await self._resolve_entities_batch_impl(
|
||||
conn, bank_id, entities_data, context, unit_event_date, taxonomy_lookup
|
||||
conn, bank_id, entities_data, context, unit_event_date, taxonomy_lookup, labels_cfg
|
||||
)
|
||||
else:
|
||||
return await self._resolve_entities_batch_impl(
|
||||
conn, bank_id, entities_data, context, unit_event_date, taxonomy_lookup
|
||||
conn, bank_id, entities_data, context, unit_event_date, taxonomy_lookup, labels_cfg
|
||||
)
|
||||
|
||||
async def _resolve_entities_batch_impl(
|
||||
@@ -246,13 +255,16 @@ class EntityResolver:
|
||||
context: str,
|
||||
unit_event_date,
|
||||
taxonomy_lookup: set[str] | None = None,
|
||||
labels_cfg=None,
|
||||
) -> list[str]:
|
||||
if self.entity_lookup == "trigram":
|
||||
# Route to backend-specific fuzzy strategy.
|
||||
# Non-PG backends (Oracle) use UTL_MATCH instead of pg_trgm.
|
||||
backend_strategy = self._ops.get_entity_resolution_strategy()
|
||||
if backend_strategy == "oracle_fuzzy":
|
||||
return await self._resolve_entities_batch_oracle_fuzzy(conn, bank_id, entities_data, unit_event_date)
|
||||
return await self._resolve_entities_batch_oracle_fuzzy(
|
||||
conn, bank_id, entities_data, unit_event_date, taxonomy_lookup, labels_cfg
|
||||
)
|
||||
# Auto-detect pg_trgm availability on first call and fall back to
|
||||
# "full" strategy if the extension is not installed. See #626.
|
||||
if not self._pg_trgm_checked:
|
||||
@@ -266,12 +278,24 @@ class EntityResolver:
|
||||
"https://github.com/vectorize-io/hindsight/issues/626"
|
||||
)
|
||||
self.entity_lookup = "full"
|
||||
return await self._resolve_entities_batch_full(conn, bank_id, entities_data, unit_event_date)
|
||||
return await self._resolve_entities_batch_trigram(conn, bank_id, entities_data, unit_event_date)
|
||||
return await self._resolve_entities_batch_full(conn, bank_id, entities_data, unit_event_date)
|
||||
return await self._resolve_entities_batch_full(
|
||||
conn, bank_id, entities_data, unit_event_date, taxonomy_lookup, labels_cfg
|
||||
)
|
||||
return await self._resolve_entities_batch_trigram(
|
||||
conn, bank_id, entities_data, unit_event_date, taxonomy_lookup, labels_cfg
|
||||
)
|
||||
return await self._resolve_entities_batch_full(
|
||||
conn, bank_id, entities_data, unit_event_date, taxonomy_lookup, labels_cfg
|
||||
)
|
||||
|
||||
async def _resolve_entities_batch_full(
|
||||
self, conn, bank_id: str, entities_data: list[dict], unit_event_date
|
||||
self,
|
||||
conn,
|
||||
bank_id: str,
|
||||
entities_data: list[dict],
|
||||
unit_event_date,
|
||||
taxonomy_lookup: set[str] | None = None,
|
||||
labels_cfg=None,
|
||||
) -> list[str]:
|
||||
"""Original strategy: load all bank entities then match in Python."""
|
||||
# Query ALL candidates for this bank
|
||||
@@ -338,11 +362,24 @@ class EntityResolver:
|
||||
all_candidates[entity_text] = matching
|
||||
|
||||
return await self._resolve_from_candidates(
|
||||
conn, bank_id, entities_data, unit_event_date, all_candidates, cooccurrence_map
|
||||
conn,
|
||||
bank_id,
|
||||
entities_data,
|
||||
unit_event_date,
|
||||
all_candidates,
|
||||
cooccurrence_map,
|
||||
taxonomy_lookup,
|
||||
labels_cfg,
|
||||
)
|
||||
|
||||
async def _resolve_entities_batch_trigram(
|
||||
self, conn, bank_id: str, entities_data: list[dict], unit_event_date
|
||||
self,
|
||||
conn,
|
||||
bank_id: str,
|
||||
entities_data: list[dict],
|
||||
unit_event_date,
|
||||
taxonomy_lookup: set[str] | None = None,
|
||||
labels_cfg=None,
|
||||
) -> list[str]:
|
||||
"""
|
||||
Trigram strategy: fetch only similar candidates per entity name using pg_trgm.
|
||||
@@ -418,11 +455,24 @@ class EntityResolver:
|
||||
cooccurrence_map[eid2].add(id_to_name[eid1])
|
||||
|
||||
return await self._resolve_from_candidates(
|
||||
conn, bank_id, entities_data, unit_event_date, all_candidates, cooccurrence_map
|
||||
conn,
|
||||
bank_id,
|
||||
entities_data,
|
||||
unit_event_date,
|
||||
all_candidates,
|
||||
cooccurrence_map,
|
||||
taxonomy_lookup,
|
||||
labels_cfg,
|
||||
)
|
||||
|
||||
async def _resolve_entities_batch_oracle_fuzzy(
|
||||
self, conn: Any, bank_id: str, entities_data: list[dict], unit_event_date: datetime | None
|
||||
self,
|
||||
conn: Any,
|
||||
bank_id: str,
|
||||
entities_data: list[dict],
|
||||
unit_event_date: datetime | None,
|
||||
taxonomy_lookup: set[str] | None = None,
|
||||
labels_cfg=None,
|
||||
) -> list[str]:
|
||||
"""
|
||||
Oracle strategy: fetch similar candidates using UTL_MATCH.JARO_WINKLER_SIMILARITY.
|
||||
@@ -506,7 +556,14 @@ class EntityResolver:
|
||||
cooccurrence_map[eid2].add(id_to_name[eid1])
|
||||
|
||||
return await self._resolve_from_candidates(
|
||||
conn, bank_id, entities_data, unit_event_date, all_candidates, cooccurrence_map
|
||||
conn,
|
||||
bank_id,
|
||||
entities_data,
|
||||
unit_event_date,
|
||||
all_candidates,
|
||||
cooccurrence_map,
|
||||
taxonomy_lookup,
|
||||
labels_cfg,
|
||||
)
|
||||
|
||||
async def _resolve_from_candidates(
|
||||
@@ -517,6 +574,8 @@ class EntityResolver:
|
||||
unit_event_date,
|
||||
all_candidates: dict[str, list],
|
||||
cooccurrence_map: dict[str, set[str]],
|
||||
taxonomy_lookup: set[str] | None = None,
|
||||
labels_cfg=None,
|
||||
) -> list[str]:
|
||||
"""Shared scoring + upsert logic used by both lookup strategies."""
|
||||
|
||||
@@ -533,11 +592,34 @@ class EntityResolver:
|
||||
|
||||
candidates = all_candidates.get(entity_text, [])
|
||||
|
||||
# Label entities (from entity_labels config) use exact matching only.
|
||||
# Their canonical names are user-defined (e.g., "use:use-001"),
|
||||
# so fuzzy resolution must NOT merge distinct label values that
|
||||
# happen to be textually similar (GH-1558).
|
||||
is_label = bool(
|
||||
labels_cfg and taxonomy_lookup and _is_label_entity(entity_text, labels_cfg, taxonomy_lookup)
|
||||
)
|
||||
|
||||
if not candidates:
|
||||
# Will create new entity
|
||||
entities_to_create.append(_EntityToCreate(idx=idx, name=entity_text, event_date=entity_event_date))
|
||||
continue
|
||||
|
||||
if is_label:
|
||||
# Exact case-insensitive match only for label entities
|
||||
exact_match = None
|
||||
entity_text_lower = entity_text.lower()
|
||||
for candidate_id, canonical_name, metadata, last_seen, mention_count in candidates:
|
||||
if canonical_name.lower() == entity_text_lower:
|
||||
exact_match = candidate_id
|
||||
break
|
||||
if exact_match:
|
||||
entity_ids[idx] = exact_match
|
||||
entities_to_update.append(_EntityStat(entity_id=exact_match, event_date=entity_event_date))
|
||||
else:
|
||||
entities_to_create.append(_EntityToCreate(idx=idx, name=entity_text, event_date=entity_event_date))
|
||||
continue
|
||||
|
||||
# Score candidates
|
||||
best_candidate = None
|
||||
best_score = 0.0
|
||||
|
||||
@@ -0,0 +1,357 @@
|
||||
"""Async graph maintenance after document/unit deletes.
|
||||
|
||||
Three reconciliation passes run together on every worker invocation:
|
||||
|
||||
1. **Relink top-up.** Drain ``graph_maintenance_queue`` (units whose
|
||||
outgoing temporal/semantic links lost a neighbour to a delete). For
|
||||
each, count current outgoing links per type; if below cap, run the
|
||||
same probes retain uses (:func:`fetch_temporal_neighbors`,
|
||||
:func:`compute_semantic_links_ann`) and insert the missing links.
|
||||
``bulk_insert_links`` has ``ON CONFLICT DO NOTHING`` on the uniqueness
|
||||
key, so we can re-probe freely and the DB de-dupes.
|
||||
|
||||
2. **Orphan entity prune.** Delete ``entities`` rows in the bank that no
|
||||
longer have any ``unit_entities`` references. FK ON DELETE CASCADE on
|
||||
``entity_cooccurrences`` then removes any cooccurrence row pointing
|
||||
at the pruned entities.
|
||||
|
||||
3. **Stale cooccurrence prune.** Defensive sweep for cooccurrence rows
|
||||
where both endpoints still exist but no current memory_unit references
|
||||
both of them — the cooccurrence was real at the time it was recorded,
|
||||
but every unit that witnessed it has since been deleted.
|
||||
|
||||
All three passes run on every invocation. The queue is the only source
|
||||
of work for pass 1; passes 2 and 3 are bank-wide sweeps backed by indexes
|
||||
on ``entities(bank_id)`` and ``unit_entities(entity_id)``, so they're
|
||||
cheap when there's nothing to do.
|
||||
|
||||
The worker dedupes on bank: a second job for the same bank is dropped
|
||||
while one is pending. Once processing starts, a new job becomes the
|
||||
*next* pending slot — so work enqueued during processing gets picked up
|
||||
by the follow-up run.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import time
|
||||
import uuid as uuid_module
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from ..models import RequestContext
|
||||
from .db.base import DatabaseConnection
|
||||
from .retain.link_utils import (
|
||||
MAX_TEMPORAL_LINKS_PER_UNIT,
|
||||
_bulk_insert_links,
|
||||
_normalize_datetime,
|
||||
compute_semantic_links_ann,
|
||||
)
|
||||
from .schema import fq_table
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .memory_engine import MemoryEngine
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Mirrors the ``top_k`` default in ``compute_semantic_links_ann`` at retain
|
||||
# time. If you change one, change the other — otherwise victims would either
|
||||
# never reach the cap (probe returns less than the cap) or stay perpetually
|
||||
# under it (cap is higher than retain creates).
|
||||
MAX_SEMANTIC_LINKS_PER_UNIT = 50
|
||||
|
||||
# Worker fetches this many rows per relink-loop iteration. Bounds
|
||||
# per-iteration probe/insert latency so a 10k-row backlog doesn't hold a
|
||||
# worker slot for minutes. Chosen so the typical iteration runs in well
|
||||
# under 1s.
|
||||
_DRAIN_BATCH_SIZE = 50
|
||||
|
||||
|
||||
@dataclass
|
||||
class JobResult:
|
||||
"""Counters surfaced to the worker dispatcher and operation result."""
|
||||
|
||||
relink_units_processed: int = 0
|
||||
relink_links_added: int = 0
|
||||
orphan_entities_pruned: int = 0
|
||||
stale_cooccurrences_pruned: int = 0
|
||||
|
||||
def as_dict(self) -> dict[str, int]:
|
||||
return {
|
||||
"relink_units_processed": self.relink_units_processed,
|
||||
"relink_links_added": self.relink_links_added,
|
||||
"orphan_entities_pruned": self.orphan_entities_pruned,
|
||||
"stale_cooccurrences_pruned": self.stale_cooccurrences_pruned,
|
||||
}
|
||||
|
||||
|
||||
async def enqueue_relink_victims(
|
||||
conn: DatabaseConnection,
|
||||
bank_id: str,
|
||||
deleted_unit_ids: list[str],
|
||||
ops: Any,
|
||||
) -> int:
|
||||
"""Enqueue surviving units whose outgoing temporal/semantic links pointed at
|
||||
``deleted_unit_ids`` for later link top-up.
|
||||
|
||||
Must run inside the same transaction that deletes the units, *before* the
|
||||
cascade fires — once the rows are gone, the join that finds the victims
|
||||
returns nothing.
|
||||
|
||||
Args:
|
||||
conn: Database connection inside the active delete transaction.
|
||||
bank_id: Bank owning the deleted units.
|
||||
deleted_unit_ids: Memory_unit IDs about to be (or being) deleted.
|
||||
ops: ``DataAccessOps`` instance, supplies the dialect-specific
|
||||
bulk-insert path.
|
||||
|
||||
Returns:
|
||||
Number of distinct victim units enqueued (after dedup against rows
|
||||
already in the queue).
|
||||
"""
|
||||
if not deleted_unit_ids:
|
||||
return 0
|
||||
|
||||
deleted_uuids = [uuid_module.UUID(uid) if isinstance(uid, str) else uid for uid in deleted_unit_ids]
|
||||
deleted_str_set = {str(uid) for uid in deleted_uuids}
|
||||
|
||||
# Find units (other than the ones being deleted) that have an outgoing
|
||||
# temporal/semantic link pointing at a doomed unit. Entity links are
|
||||
# intentionally excluded — they're scheduled for removal and would only
|
||||
# add noise to the recompute job.
|
||||
victim_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT DISTINCT from_unit_id
|
||||
FROM {fq_table("memory_links")}
|
||||
WHERE to_unit_id = ANY($1::uuid[])
|
||||
AND bank_id = $2
|
||||
AND link_type IN ('temporal', 'semantic')
|
||||
""",
|
||||
deleted_uuids,
|
||||
bank_id,
|
||||
)
|
||||
|
||||
victim_ids = [row["from_unit_id"] for row in victim_rows if str(row["from_unit_id"]) not in deleted_str_set]
|
||||
|
||||
if not victim_ids:
|
||||
return 0
|
||||
|
||||
await ops.enqueue_graph_maintenance(
|
||||
conn,
|
||||
fq_table("graph_maintenance_queue"),
|
||||
bank_id,
|
||||
victim_ids,
|
||||
)
|
||||
|
||||
logger.debug(
|
||||
f"[GRAPH_MAINT] Enqueued {len(victim_ids)} relink victims in "
|
||||
f"bank={bank_id} (deleted {len(deleted_unit_ids)} units)"
|
||||
)
|
||||
return len(victim_ids)
|
||||
|
||||
|
||||
async def run_graph_maintenance_job(
|
||||
memory_engine: "MemoryEngine",
|
||||
bank_id: str,
|
||||
request_context: RequestContext,
|
||||
operation_id: str | None = None,
|
||||
) -> dict[str, int]:
|
||||
"""Run all maintenance passes for ``bank_id`` until the relink queue is
|
||||
drained, then sweep entities and cooccurrences once.
|
||||
|
||||
Returns:
|
||||
Per-pass counters from :class:`JobResult`.
|
||||
"""
|
||||
del request_context # accepted for symmetry with other run_*_job helpers
|
||||
backend = await memory_engine._get_backend()
|
||||
ops = backend.ops
|
||||
|
||||
result = JobResult()
|
||||
job_start = time.time()
|
||||
|
||||
# --- Pass 1: relink ---
|
||||
# Per-iteration loop: claim → top up → commit. We rely on submit-time
|
||||
# dedup to keep at most one job per bank running, so no need for
|
||||
# SKIP LOCKED.
|
||||
iterations = 0
|
||||
while True:
|
||||
from .memory_engine import acquire_with_retry
|
||||
|
||||
async with acquire_with_retry(backend) as conn:
|
||||
async with conn.transaction():
|
||||
unit_ids = await ops.claim_graph_maintenance_batch(
|
||||
conn,
|
||||
fq_table("graph_maintenance_queue"),
|
||||
bank_id,
|
||||
_DRAIN_BATCH_SIZE,
|
||||
)
|
||||
if not unit_ids:
|
||||
break
|
||||
|
||||
result.relink_links_added += await _relink_batch(conn, bank_id, unit_ids, ops, backend)
|
||||
|
||||
result.relink_units_processed += len(unit_ids)
|
||||
iterations += 1
|
||||
|
||||
if iterations > 10000:
|
||||
# Defensive guard against runaway loops — at 50 units/iter that's
|
||||
# 500k targets, far beyond any realistic single-bank backlog.
|
||||
logger.error(
|
||||
f"[GRAPH_MAINT] bank={bank_id} hit iteration cap ({iterations}); aborting relink ({result.as_dict()})"
|
||||
)
|
||||
break
|
||||
|
||||
# --- Pass 2 & 3: entity / cooccurrence sweeps ---
|
||||
# Bank-wide single-statement deletes. Cheap when there's nothing to do.
|
||||
from .memory_engine import acquire_with_retry
|
||||
|
||||
async with acquire_with_retry(backend) as conn:
|
||||
async with conn.transaction():
|
||||
result.orphan_entities_pruned = await ops.prune_orphan_entities(
|
||||
conn,
|
||||
fq_table("entities"),
|
||||
fq_table("unit_entities"),
|
||||
bank_id,
|
||||
)
|
||||
# The orphan prune above cascades cooccurrences via FK. The
|
||||
# explicit cooccurrence pass below catches the *stale-count*
|
||||
# case: both entities still exist but no current unit witnesses
|
||||
# them together.
|
||||
result.stale_cooccurrences_pruned = await ops.prune_stale_cooccurrences(
|
||||
conn,
|
||||
fq_table("entity_cooccurrences"),
|
||||
fq_table("unit_entities"),
|
||||
fq_table("entities"),
|
||||
bank_id,
|
||||
)
|
||||
|
||||
elapsed = time.time() - job_start
|
||||
logger.info(
|
||||
f"[GRAPH_MAINT] bank={bank_id} done: {result.as_dict()}, elapsed={elapsed:.2f}s, operation_id={operation_id}"
|
||||
)
|
||||
return result.as_dict()
|
||||
|
||||
|
||||
async def _relink_batch(
|
||||
conn: DatabaseConnection,
|
||||
bank_id: str,
|
||||
victim_ids: list[str],
|
||||
ops: Any,
|
||||
backend: Any,
|
||||
) -> int:
|
||||
"""Top up temporal/semantic links for a batch of victim units. Returns rows inserted."""
|
||||
# Load each victim's metadata. Victims whose units were deleted between
|
||||
# enqueue and now silently drop out — exactly the no-op behaviour we want
|
||||
# for stale queue rows.
|
||||
victim_uuids = [uuid_module.UUID(vid) for vid in victim_ids]
|
||||
victim_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT id::text AS id, event_date, fact_type, embedding::text AS embedding
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE id = ANY($1::uuid[])
|
||||
AND bank_id = $2
|
||||
AND fact_type IN ('experience', 'world')
|
||||
""",
|
||||
victim_uuids,
|
||||
bank_id,
|
||||
)
|
||||
|
||||
if not victim_rows:
|
||||
return 0
|
||||
|
||||
alive_uuids = [uuid_module.UUID(row["id"]) for row in victim_rows]
|
||||
|
||||
# Count current outgoing temporal/semantic links per victim so we only
|
||||
# probe for the ones genuinely below cap. Saves the bulk of the work when
|
||||
# most victims still have plenty of links.
|
||||
count_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT from_unit_id, link_type, COUNT(*) AS cnt
|
||||
FROM {fq_table("memory_links")}
|
||||
WHERE from_unit_id = ANY($1::uuid[])
|
||||
AND bank_id = $2
|
||||
AND link_type IN ('temporal', 'semantic')
|
||||
GROUP BY from_unit_id, link_type
|
||||
""",
|
||||
alive_uuids,
|
||||
bank_id,
|
||||
)
|
||||
counts: dict[tuple[str, str], int] = {}
|
||||
for row in count_rows:
|
||||
counts[(str(row["from_unit_id"]), row["link_type"])] = int(row["cnt"])
|
||||
|
||||
# --- Temporal top-up ---
|
||||
temporal_needs = [r for r in victim_rows if counts.get((r["id"], "temporal"), 0) < MAX_TEMPORAL_LINKS_PER_UNIT]
|
||||
new_links: list[tuple] = []
|
||||
|
||||
if temporal_needs:
|
||||
lateral_unit_ids = [uuid_module.UUID(r["id"]) for r in temporal_needs if r["event_date"] is not None]
|
||||
lateral_event_dates = [
|
||||
_normalize_datetime(r["event_date"]) for r in temporal_needs if r["event_date"] is not None
|
||||
]
|
||||
lateral_fact_types = [r["fact_type"] for r in temporal_needs if r["event_date"] is not None]
|
||||
|
||||
if lateral_unit_ids:
|
||||
rows = await ops.fetch_temporal_neighbors(
|
||||
conn,
|
||||
fq_table("memory_units"),
|
||||
bank_id,
|
||||
lateral_unit_ids,
|
||||
lateral_event_dates,
|
||||
lateral_fact_types,
|
||||
MAX_TEMPORAL_LINKS_PER_UNIT,
|
||||
)
|
||||
for row in rows:
|
||||
time_diff_h = float(row["time_diff_hours"])
|
||||
# Mirror the 24h window enforced at retain time. The bidirectional
|
||||
# index scan returns the K closest neighbours regardless of
|
||||
# window, so we filter here.
|
||||
if time_diff_h > 24:
|
||||
continue
|
||||
weight = max(0.3, 1.0 - (time_diff_h / 24))
|
||||
new_links.append((row["from_id"], str(row["id"]), "temporal", weight, None))
|
||||
|
||||
# --- Semantic top-up ---
|
||||
# ANN must run on its own connection: it opens a nested transaction with
|
||||
# SET LOCAL hnsw.ef_search + CREATE TEMP TABLE ON COMMIT DROP, and nesting
|
||||
# that inside our current write transaction would commit our writes early.
|
||||
semantic_needs = [
|
||||
r
|
||||
for r in victim_rows
|
||||
if counts.get((r["id"], "semantic"), 0) < MAX_SEMANTIC_LINKS_PER_UNIT and r["embedding"] is not None
|
||||
]
|
||||
if semantic_needs:
|
||||
from .memory_engine import acquire_with_retry
|
||||
|
||||
seed_ids = [r["id"] for r in semantic_needs]
|
||||
seed_embs = [r["embedding"] for r in semantic_needs]
|
||||
seed_ftypes = [r["fact_type"] for r in semantic_needs]
|
||||
async with acquire_with_retry(backend) as ann_conn:
|
||||
try:
|
||||
ann_links = await compute_semantic_links_ann(
|
||||
ann_conn,
|
||||
bank_id,
|
||||
seed_ids,
|
||||
seed_embs,
|
||||
fact_types=seed_ftypes,
|
||||
)
|
||||
# Strip self-links (rare but possible because the ANN probe
|
||||
# has no exclude list — see the comment in compute_semantic_links_ann).
|
||||
ann_links = [lnk for lnk in ann_links if lnk[0] != lnk[1]]
|
||||
new_links.extend(ann_links)
|
||||
except Exception as e:
|
||||
# ANN uses PG-specific HNSW syntax; on dialects/configs where
|
||||
# it isn't available we still want the temporal top-up to land.
|
||||
logger.warning(f"[GRAPH_MAINT] Semantic top-up failed for bank={bank_id}: {type(e).__name__}: {e}")
|
||||
|
||||
if not new_links:
|
||||
return 0
|
||||
|
||||
await _bulk_insert_links(
|
||||
conn,
|
||||
new_links,
|
||||
bank_id=bank_id,
|
||||
skip_exists_check=False,
|
||||
ops=ops,
|
||||
)
|
||||
return len(new_links)
|
||||
@@ -9,6 +9,7 @@ import os
|
||||
import re
|
||||
import time
|
||||
import uuid
|
||||
from contextlib import AsyncExitStack
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
@@ -27,9 +28,12 @@ except ImportError:
|
||||
from ..config import (
|
||||
DEFAULT_LLM_MAX_CONCURRENT,
|
||||
DEFAULT_LLM_TIMEOUT,
|
||||
ENV_CONSOLIDATION_LLM_MAX_CONCURRENT,
|
||||
ENV_LLM_GROQ_SERVICE_TIER,
|
||||
ENV_LLM_MAX_CONCURRENT,
|
||||
ENV_LLM_TIMEOUT,
|
||||
ENV_REFLECT_LLM_MAX_CONCURRENT,
|
||||
ENV_RETAIN_LLM_MAX_CONCURRENT,
|
||||
)
|
||||
from ..metrics import get_metrics_collector
|
||||
from .response_models import TokenUsage
|
||||
@@ -42,12 +46,74 @@ logger = logging.getLogger(__name__)
|
||||
# Disable httpx logging
|
||||
logging.getLogger("httpx").setLevel(logging.WARNING)
|
||||
|
||||
# Global semaphore to limit concurrent LLM requests across all instances
|
||||
# Set HINDSIGHT_API_LLM_MAX_CONCURRENT=1 for local LLMs (LM Studio, Ollama)
|
||||
# Global semaphore to limit concurrent LLM requests across all instances.
|
||||
# Set HINDSIGHT_API_LLM_MAX_CONCURRENT=1 for local LLMs (LM Studio, Ollama).
|
||||
_llm_max_concurrent = int(os.getenv(ENV_LLM_MAX_CONCURRENT, str(DEFAULT_LLM_MAX_CONCURRENT)))
|
||||
_global_llm_semaphore = asyncio.Semaphore(_llm_max_concurrent)
|
||||
|
||||
|
||||
def _build_per_op_semaphores() -> dict[str, asyncio.Semaphore]:
|
||||
"""Build the per-operation semaphore registry from env vars.
|
||||
|
||||
Each per-op cap is composed with — not a substitute for — the global cap:
|
||||
a call that matches a configured operation must acquire both its per-op
|
||||
semaphore and the global semaphore. This lets operators reserve headroom
|
||||
in the global pool by capping individual operations (e.g. cap retain at 2
|
||||
of 4 global slots so the live chat path always has 2 slots available).
|
||||
|
||||
Operations without a configured env var are absent from the registry and
|
||||
therefore only constrained by the global cap.
|
||||
"""
|
||||
semaphores: dict[str, asyncio.Semaphore] = {}
|
||||
for op, env_var in (
|
||||
("retain", ENV_RETAIN_LLM_MAX_CONCURRENT),
|
||||
("reflect", ENV_REFLECT_LLM_MAX_CONCURRENT),
|
||||
("consolidation", ENV_CONSOLIDATION_LLM_MAX_CONCURRENT),
|
||||
):
|
||||
raw = os.getenv(env_var)
|
||||
if raw is None or raw == "":
|
||||
continue
|
||||
value = int(raw)
|
||||
if value <= 0:
|
||||
raise ValueError(f"{env_var} must be a positive integer, got {raw!r}")
|
||||
semaphores[op] = asyncio.Semaphore(value)
|
||||
return semaphores
|
||||
|
||||
|
||||
_per_op_llm_semaphores: dict[str, asyncio.Semaphore] = _build_per_op_semaphores()
|
||||
|
||||
|
||||
def _scope_to_operation(scope: str) -> str | None:
|
||||
"""Map a call scope to its per-operation concurrency bucket.
|
||||
|
||||
Returns None for scopes that don't belong to a tracked operation
|
||||
(verification probes, bank_mission, memory_think, mental_model_delta_ops),
|
||||
which then run under the global cap only.
|
||||
"""
|
||||
if scope.startswith("retain"):
|
||||
return "retain"
|
||||
if scope.startswith("reflect"):
|
||||
return "reflect"
|
||||
if scope.startswith("consolidation"):
|
||||
return "consolidation"
|
||||
return None
|
||||
|
||||
|
||||
def _semaphores_for_scope(scope: str) -> list[asyncio.Semaphore]:
|
||||
"""Return the semaphores a call with the given scope must acquire.
|
||||
|
||||
Always includes the global semaphore; includes the per-op semaphore when
|
||||
one is configured for the scope's operation bucket.
|
||||
"""
|
||||
op = _scope_to_operation(scope)
|
||||
per_op = _per_op_llm_semaphores.get(op) if op is not None else None
|
||||
if per_op is None:
|
||||
return [_global_llm_semaphore]
|
||||
# Per-op acquired first so contention queues on the narrower cap before
|
||||
# holding a global slot.
|
||||
return [per_op, _global_llm_semaphore]
|
||||
|
||||
|
||||
def sanitize_llm_output(text: str | None) -> str | None:
|
||||
"""
|
||||
Sanitize text by removing characters that break downstream systems.
|
||||
@@ -312,6 +378,7 @@ def create_llm_provider(
|
||||
"openai",
|
||||
"groq",
|
||||
"ollama",
|
||||
"ollama-cloud",
|
||||
"lmstudio",
|
||||
"minimax",
|
||||
"deepseek",
|
||||
@@ -409,6 +476,7 @@ class LLMProvider:
|
||||
"openai",
|
||||
"groq",
|
||||
"ollama",
|
||||
"ollama-cloud",
|
||||
"gemini",
|
||||
"anthropic",
|
||||
"lmstudio",
|
||||
@@ -437,6 +505,8 @@ class LLMProvider:
|
||||
self.base_url = "https://api.groq.com/openai/v1"
|
||||
elif self.provider == "ollama":
|
||||
self.base_url = "http://localhost:11434/v1"
|
||||
elif self.provider == "ollama-cloud":
|
||||
self.base_url = "https://ollama.com/v1"
|
||||
elif self.provider == "lmstudio":
|
||||
self.base_url = "http://localhost:1234/v1"
|
||||
elif self.provider == "minimax":
|
||||
@@ -629,7 +699,10 @@ class LLMProvider:
|
||||
structured = "+structured" if response_format is not None else ""
|
||||
set_stage(f"llm.{self.provider}.{scope}{structured}")
|
||||
|
||||
async with _global_llm_semaphore:
|
||||
async with AsyncExitStack() as stack:
|
||||
for sem in _semaphores_for_scope(scope):
|
||||
await stack.enter_async_context(sem)
|
||||
|
||||
# Delegate to provider implementation
|
||||
result = await self._provider_impl.call(
|
||||
messages=messages,
|
||||
@@ -654,7 +727,7 @@ class LLMProvider:
|
||||
# Sync the mock calls from provider implementation to wrapper
|
||||
self._mock_calls = self._provider_impl.get_mock_calls()
|
||||
|
||||
return result
|
||||
return result
|
||||
|
||||
async def call_with_tools(
|
||||
self,
|
||||
@@ -689,7 +762,10 @@ class LLMProvider:
|
||||
|
||||
set_stage(f"llm.{self.provider}.{scope}+tools")
|
||||
|
||||
async with _global_llm_semaphore:
|
||||
async with AsyncExitStack() as stack:
|
||||
for sem in _semaphores_for_scope(scope):
|
||||
await stack.enter_async_context(sem)
|
||||
|
||||
# Delegate to provider implementation
|
||||
result = await self._provider_impl.call_with_tools(
|
||||
messages=messages,
|
||||
@@ -712,7 +788,7 @@ class LLMProvider:
|
||||
# Sync the mock calls from provider implementation to wrapper
|
||||
self._mock_calls = self._provider_impl.get_mock_calls()
|
||||
|
||||
return result
|
||||
return result
|
||||
|
||||
def set_response_callback(self, fn: Any) -> None:
|
||||
"""Set a callback invoked on each call() instead of the fixed mock response."""
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,41 @@
|
||||
"""Shared utilities for prompt assembly."""
|
||||
|
||||
import re
|
||||
|
||||
_LONE_OPEN_BRACE = re.compile(r"(?<!\{)\{(?!\{)")
|
||||
_LONE_CLOSE_BRACE = re.compile(r"(?<!\})\}(?!\})")
|
||||
|
||||
|
||||
def escape_for_prompt(text: str) -> str:
|
||||
"""Double any lone ``{`` / ``}`` so the text survives ``str.format`` untouched.
|
||||
|
||||
Prompt templates are often passed through ``str.format`` to substitute real
|
||||
placeholders like ``{facts_text}``. Any literal braces in caller-supplied
|
||||
text — e.g. a bank mission that contains JSON examples — would otherwise be
|
||||
interpreted as format keys and raise ``KeyError``.
|
||||
|
||||
Idempotent: text that already contains escaped ``{{`` / ``}}`` pairs is
|
||||
left as-is. Only lone braces (not adjacent to another brace of the same
|
||||
kind) are doubled.
|
||||
"""
|
||||
text = _LONE_OPEN_BRACE.sub("{{", text)
|
||||
text = _LONE_CLOSE_BRACE.sub("}}", text)
|
||||
return text
|
||||
|
||||
|
||||
def output_language_directive(language: str | None) -> str:
|
||||
"""Return an LLM directive forcing all output into ``language``.
|
||||
|
||||
Used by retain (fact extraction), consolidation (observations), and reflect
|
||||
(response synthesis) so HINDSIGHT_API_LLM_OUTPUT_LANGUAGE applies uniformly
|
||||
across every LLM-generated artifact. Returns an empty string when
|
||||
``language`` is unset so the calling prompt stays unchanged.
|
||||
"""
|
||||
if not language:
|
||||
return ""
|
||||
return (
|
||||
f"\n\nIMPORTANT: Respond exclusively in {language}. "
|
||||
f"Translate any source content into {language}. "
|
||||
f"All output text — including fact text, observations, entity names, "
|
||||
f"and the final response — must be in {language}."
|
||||
)
|
||||
@@ -93,7 +93,6 @@ class AnthropicLLM(LLMInterface):
|
||||
await self.call(
|
||||
messages=test_messages,
|
||||
max_completion_tokens=10,
|
||||
temperature=0.0,
|
||||
scope="verification",
|
||||
max_retries=0,
|
||||
)
|
||||
@@ -179,9 +178,6 @@ class AnthropicLLM(LLMInterface):
|
||||
if system_prompt:
|
||||
call_params["system"] = system_prompt
|
||||
|
||||
if temperature is not None:
|
||||
call_params["temperature"] = temperature
|
||||
|
||||
last_exception = None
|
||||
|
||||
for attempt in range(max_retries + 1):
|
||||
@@ -398,9 +394,6 @@ class AnthropicLLM(LLMInterface):
|
||||
if system_prompt:
|
||||
call_params["system"] = system_prompt
|
||||
|
||||
if temperature is not None:
|
||||
call_params["temperature"] = temperature
|
||||
|
||||
last_exception = None
|
||||
for attempt in range(max_retries + 1):
|
||||
try:
|
||||
|
||||
@@ -110,9 +110,15 @@ class CodexLLM(LLMInterface):
|
||||
"Or use a different provider (openai, anthropic, gemini) with API keys."
|
||||
) from e
|
||||
|
||||
# Use ChatGPT backend API endpoint
|
||||
if not self.base_url:
|
||||
# Use ChatGPT backend API endpoint. Codex auth is tied to
|
||||
# chatgpt.com/backend-api, not the OpenAI-compatible base URL used by
|
||||
# other providers. Deployments often set a global LLM_BASE_URL for an
|
||||
# OpenAI-compatible proxy; ignore that inherited value unless the user
|
||||
# explicitly provides a Codex backend URL.
|
||||
if not self.base_url or self.base_url.rstrip("/").endswith("/v1"):
|
||||
self.base_url = "https://chatgpt.com/backend-api"
|
||||
else:
|
||||
self.base_url = self.base_url.rstrip("/")
|
||||
|
||||
# Normalize model name (strip openai/ prefix if present)
|
||||
if self.model.startswith("openai/"):
|
||||
|
||||
@@ -153,11 +153,24 @@ class MockLLM(LLMInterface):
|
||||
result = self._response_callback(messages, scope)
|
||||
elif self._mock_response is not None:
|
||||
result = self._mock_response
|
||||
elif scope == "retain_extract_facts" and skip_validation:
|
||||
# Fact extraction: return canned facts derived from user message text.
|
||||
# This allows tests using a mock LLM to get real facts into the DB
|
||||
# so retain → recall → reflect pipelines work end-to-end.
|
||||
result = self._build_mock_facts(messages)
|
||||
elif scope == "consolidation" and response_format is not None:
|
||||
# Consolidation: produce a single observation from the input facts
|
||||
# so the full pipeline (retain → consolidation → observation → recall) works.
|
||||
result = self._build_mock_consolidation(messages, response_format)
|
||||
elif scope == "memory_think":
|
||||
# Reflect: return a plausible text answer
|
||||
result = "Based on the available information, the answer is related to the context provided."
|
||||
elif response_format is not None:
|
||||
# Try to create a minimal valid instance of the response format
|
||||
# Structured output: try to return a valid empty instance of the model
|
||||
# so that callers expecting e.g. response_format with defaults
|
||||
# get a valid instance rather than a crash on {"mock": True}.
|
||||
try:
|
||||
# For Pydantic models, try to create with minimal valid data
|
||||
result = {"mock": True}
|
||||
result = response_format()
|
||||
except Exception:
|
||||
result = {"mock": True}
|
||||
else:
|
||||
@@ -243,6 +256,12 @@ class MockLLM(LLMInterface):
|
||||
else:
|
||||
result = LLMToolCallResult(content="mock response", finish_reason="stop")
|
||||
|
||||
# Set mock token usage on result if not already set
|
||||
if result.input_tokens == 0:
|
||||
result.input_tokens = 10
|
||||
if result.output_tokens == 0:
|
||||
result.output_tokens = 5
|
||||
|
||||
# Record span with mock values
|
||||
# Convert LLMToolCall objects to dicts for span recording
|
||||
tool_calls_dict = (
|
||||
@@ -266,6 +285,92 @@ class MockLLM(LLMInterface):
|
||||
|
||||
return result
|
||||
|
||||
@staticmethod
|
||||
def _build_mock_facts(messages: list[dict]) -> dict:
|
||||
"""Build a canned fact extraction response from the user message text.
|
||||
|
||||
Splits the input into sentence-like chunks and returns each as a separate
|
||||
world fact with a simple entity extracted from the first noun-like word.
|
||||
This is intentionally simplistic — it just needs to produce structurally
|
||||
valid facts so the rest of the pipeline (embedding, storage, recall) works.
|
||||
"""
|
||||
import re
|
||||
|
||||
user_text = ""
|
||||
for m in messages:
|
||||
if m.get("role") == "user":
|
||||
user_text = m.get("content", "")
|
||||
break
|
||||
|
||||
# Split on sentence boundaries: period followed by space/EOL (not mid-number), or newlines
|
||||
sentences = [s.strip() for s in re.split(r"(?<=\.)\s+|\n+", user_text) if s.strip() and len(s.strip()) > 10]
|
||||
|
||||
if not sentences:
|
||||
sentences = [user_text[:200] if user_text else "mock fact"]
|
||||
|
||||
facts = []
|
||||
for sentence in sentences[:10]: # Cap at 10 facts per chunk
|
||||
# Extract simple entities: capitalized words that aren't common words
|
||||
words = re.findall(r"\b[A-Z][a-z]+\b", sentence)
|
||||
entities = [{"text": w} for w in dict.fromkeys(words)][:5] # Dedupe, cap at 5
|
||||
|
||||
facts.append(
|
||||
{
|
||||
"what": sentence,
|
||||
"when": "N/A",
|
||||
"where": "N/A",
|
||||
"who": "N/A",
|
||||
"why": "N/A",
|
||||
"fact_kind": "conversation",
|
||||
"fact_type": "world",
|
||||
"entities": entities,
|
||||
}
|
||||
)
|
||||
|
||||
return {"facts": facts}
|
||||
|
||||
@staticmethod
|
||||
def _build_mock_consolidation(messages: list[dict], response_format: Any) -> Any:
|
||||
"""Build a mock consolidation response that creates one observation per fact.
|
||||
|
||||
Parses fact IDs from the consolidation prompt and creates one observation
|
||||
per fact, each referencing its source fact ID. This mimics real LLM behavior
|
||||
where distinct facts produce separate observations, preserving entity
|
||||
separation so pipeline tests (graph filtering, entity linking) work correctly.
|
||||
"""
|
||||
import re
|
||||
|
||||
user_text = ""
|
||||
for m in messages:
|
||||
if m.get("role") == "user":
|
||||
user_text = m.get("content", "")
|
||||
break
|
||||
|
||||
# Extract fact UUIDs from the prompt (format: "[<uuid>] <text>")
|
||||
fact_entries = re.findall(r"\[([0-9a-f-]{36})\]\s*(.+?)(?:\n|$)", user_text)
|
||||
|
||||
if not fact_entries:
|
||||
# No facts to consolidate — return empty response
|
||||
try:
|
||||
return response_format()
|
||||
except Exception:
|
||||
return {"creates": [], "updates": [], "deletes": []}
|
||||
|
||||
# Create one observation per fact to preserve entity separation
|
||||
creates = []
|
||||
for fact_id, fact_text in fact_entries:
|
||||
creates.append({"text": fact_text.strip(), "source_fact_ids": [fact_id]})
|
||||
|
||||
try:
|
||||
return response_format(
|
||||
creates=creates,
|
||||
updates=[],
|
||||
deletes=[],
|
||||
)
|
||||
except Exception:
|
||||
# Fallback if response_format constructor doesn't accept these args
|
||||
return {"creates": creates, "updates": [], "deletes": []}
|
||||
|
||||
async def cleanup(self) -> None:
|
||||
"""Clean up resources (no-op for mock provider)."""
|
||||
pass
|
||||
@@ -318,6 +423,8 @@ class MockLLM(LLMInterface):
|
||||
return self._mock_calls
|
||||
|
||||
def clear_mock_calls(self) -> None:
|
||||
"""Clear the recorded mock calls and any set exception."""
|
||||
"""Clear all recorded calls and any configured response/exception state."""
|
||||
self._mock_calls = []
|
||||
self._mock_exception = None
|
||||
self._mock_response = None
|
||||
self._response_callback = None
|
||||
|
||||
@@ -270,6 +270,7 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
"openai",
|
||||
"groq",
|
||||
"ollama",
|
||||
"ollama-cloud",
|
||||
"lmstudio",
|
||||
"llamacpp",
|
||||
"minimax",
|
||||
@@ -288,6 +289,8 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
self.base_url = "https://api.groq.com/openai/v1"
|
||||
elif self.provider == "ollama":
|
||||
self.base_url = "http://localhost:11434/v1"
|
||||
elif self.provider == "ollama-cloud":
|
||||
self.base_url = "https://ollama.com/v1"
|
||||
elif self.provider == "lmstudio":
|
||||
self.base_url = "http://localhost:1234/v1"
|
||||
elif self.provider == "minimax":
|
||||
@@ -316,6 +319,7 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
"openrouter",
|
||||
"zai",
|
||||
"opencode-go",
|
||||
"ollama-cloud",
|
||||
)
|
||||
and not self.api_key
|
||||
):
|
||||
@@ -1073,12 +1077,17 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
|
||||
last_exception = None
|
||||
|
||||
# Pass API key as Bearer token for cloud Ollama endpoints
|
||||
headers: dict[str, str] = {}
|
||||
if self.api_key and self.api_key != "local":
|
||||
headers["Authorization"] = f"Bearer {self.api_key}"
|
||||
|
||||
async with httpx.AsyncClient(timeout=300.0) as client:
|
||||
for attempt in range(max_retries + 1):
|
||||
if attempt > 0:
|
||||
set_stage(f"llm.ollama_native.{scope}.attempt={attempt + 1}/{max_retries + 1}")
|
||||
try:
|
||||
response = await client.post(native_url, json=payload)
|
||||
response = await client.post(native_url, json=payload, headers=headers)
|
||||
response.raise_for_status()
|
||||
|
||||
result = response.json()
|
||||
|
||||
@@ -321,6 +321,7 @@ async def run_reflect_agent(
|
||||
include_recall: bool = True,
|
||||
budget: str | None = None,
|
||||
max_context_tokens: int = 100_000,
|
||||
llm_output_language: str | None = None,
|
||||
) -> ReflectAgentResult:
|
||||
"""
|
||||
Execute the reflect agent loop using native tool calling.
|
||||
@@ -369,7 +370,12 @@ async def run_reflect_agent(
|
||||
|
||||
# Build initial messages (directives are injected into system prompt at START and END)
|
||||
system_prompt = build_system_prompt_for_tools(
|
||||
bank_profile, context, directives=directives, has_mental_models=has_mental_models, budget=budget
|
||||
bank_profile,
|
||||
context,
|
||||
directives=directives,
|
||||
has_mental_models=has_mental_models,
|
||||
include_observations=include_observations,
|
||||
budget=budget,
|
||||
)
|
||||
messages: list[dict[str, Any]] = [
|
||||
{"role": "system", "content": system_prompt},
|
||||
@@ -447,7 +453,10 @@ 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": build_final_system_prompt(bank_profile.get("mission"), llm_output_language),
|
||||
},
|
||||
{"role": "user", "content": prompt},
|
||||
],
|
||||
scope="reflect",
|
||||
@@ -504,7 +513,10 @@ 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": build_final_system_prompt(bank_profile.get("mission"), llm_output_language),
|
||||
},
|
||||
{"role": "user", "content": prompt},
|
||||
],
|
||||
scope="reflect",
|
||||
@@ -607,7 +619,10 @@ 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": build_final_system_prompt(bank_profile.get("mission"), llm_output_language),
|
||||
},
|
||||
{"role": "user", "content": prompt},
|
||||
],
|
||||
scope="reflect",
|
||||
@@ -728,7 +743,10 @@ 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": build_final_system_prompt(bank_profile.get("mission"), llm_output_language),
|
||||
},
|
||||
{"role": "user", "content": prompt},
|
||||
],
|
||||
scope="reflect",
|
||||
|
||||
@@ -98,6 +98,7 @@ def build_system_prompt_for_tools(
|
||||
context: str | None = None,
|
||||
directives: list[dict[str, Any]] | None = None,
|
||||
has_mental_models: bool = False,
|
||||
include_observations: bool = True,
|
||||
budget: str | None = None,
|
||||
) -> str:
|
||||
"""
|
||||
@@ -108,11 +109,17 @@ def build_system_prompt_for_tools(
|
||||
2. search_observations - Consolidated knowledge with freshness
|
||||
3. recall - Raw facts as ground truth
|
||||
|
||||
The retrieval-strategy and workflow sections are built to match the tools
|
||||
actually exposed to the LLM — mentioning a tool the agent has disabled
|
||||
causes weaker LLMs to either hallucinate the call (rejected by the agent)
|
||||
or give up with "I cannot find any information…" (see #1724).
|
||||
|
||||
Args:
|
||||
bank_profile: Bank profile with name and mission
|
||||
context: Optional additional context
|
||||
directives: Optional list of directive mental models to inject as hard rules
|
||||
has_mental_models: Whether the bank has any mental models (skip if not)
|
||||
include_observations: Whether search_observations is in the tool list.
|
||||
budget: Search depth budget - "low", "mid", or "high". Controls exploration thoroughness.
|
||||
"""
|
||||
name = bank_profile.get("name", "Assistant")
|
||||
@@ -158,56 +165,137 @@ def build_system_prompt_for_tools(
|
||||
"- If memories mention someone did an activity, you can infer they likely enjoyed it",
|
||||
"- Synthesize a coherent narrative from related memories",
|
||||
"- Be a thoughtful interpreter, not just a literal repeater",
|
||||
"- When the exact answer isn't stated, use what IS stated to give the best answer",
|
||||
"- When the exact answer isn't stated, use what IS stated to give a best-effort answer AND surface any uncertainty — never invent confidence the data doesn't support.",
|
||||
"",
|
||||
"## Temporal Reasoning",
|
||||
"Every memory and observation carries temporal fields in the JSON tool result:",
|
||||
"- `mentioned_at` — when the user retained the fact (always set).",
|
||||
"- `occurred_start` / `occurred_end` — when the underlying event happened (optional, set for dated events).",
|
||||
"",
|
||||
"When facts about the SAME facet conflict — counts, statuses, ownership, location, presence, etc. — the fact with the LATEST `mentioned_at` is authoritative. Later statements SUPERSEDE earlier ones. Do NOT average, sum, or favor an explicitly-dated fact over a more recent one.",
|
||||
"",
|
||||
"Example: three count facts come back from recall:",
|
||||
" - 'Team has 2 engineers' (mentioned_at=T1)",
|
||||
" - 'Team now has 1 engineer' (mentioned_at=T2, occurred_start=2026-05-25)",
|
||||
" - 'Team has 5 engineers' (mentioned_at=T3)",
|
||||
"with T1 < T2 < T3. The current size is 5, not 1. Then apply later events (e.g. someone leaving after T3) on top of that.",
|
||||
"",
|
||||
"For reconstructing a TIMELINE of events, order by `occurred_start` / `occurred_end` (when things happened), not `mentioned_at` (when they were retained).",
|
||||
"",
|
||||
"## Conflicts and Ambiguity",
|
||||
"Not every retrieval converges on a single answer. Distinguish two cases:",
|
||||
"",
|
||||
"- RESOLVABLE conflict — the temporal rule above (latest `mentioned_at` wins) cleanly picks a winner. Apply it and move on.",
|
||||
"- UNRESOLVABLE ambiguity — the data is internally inconsistent in a way the temporal rule does NOT settle. Examples: a recent aggregate (count, total) is incompatible with the individual entities you can enumerate; two equally-recent facts disagree and no later fact resolves them; events are described but their relative order is unclear; the user's own statements contradict each other and nothing later reconciles them.",
|
||||
"",
|
||||
"When the data is genuinely ambiguous: SAY SO in your answer. Name the conflicting facts. Explain why they can't be reconciled. Give a range or a best-effort interpretation with explicit uncertainty (e.g. 'between X and Y, depending on [unresolved condition]'; or 'the most recent statement says A, but B was stated earlier and the gap isn't accounted for in any later fact').",
|
||||
"",
|
||||
"An honest 'the data is inconsistent about X' beats a confident wrong answer. Do NOT pick a value arbitrarily, average conflicting values, or smooth over gaps in confident prose. Acknowledging ambiguity is a successful answer, not a failure mode.",
|
||||
"",
|
||||
"## Showing Your Reasoning",
|
||||
"For any answer that resolves a conflict between facts, applies events on top of a count or status, or settles an ambiguity — show your work in the answer text so a reader can audit it.",
|
||||
"",
|
||||
"Walk through these steps explicitly:",
|
||||
"1. **List the relevant facts in `mentioned_at` order (oldest → newest)**, each with the value it asserts. Use a short bulleted list.",
|
||||
"2. **Identify the authoritative fact** under the temporal rule (latest `mentioned_at` for the contested facet). Write its date down.",
|
||||
"3. **List candidate events to apply on top** — anything that changes the count, status, or state being asked about. Write each event's date down next to it.",
|
||||
"4. **Sanity-check each candidate event against the authoritative date** — for EVERY event from step 3, write a one-line check in the form `<event> (<event_date>) vs authoritative (<authoritative_date>) → BEFORE/AFTER → KEEP/DROP`. If the event is BEFORE or EQUAL to the authoritative date, DROP it: it is already reflected in the authoritative fact, and applying it again is double-counting. This is the single most common mistake — do not skip this step even if you feel confident.",
|
||||
"5. **Show the arithmetic or derivation explicitly** using only the KEEP events from step 4 — e.g. 'authoritative count = 5 (at 2025-02-12); kept events: Shadow died (2025-03-12, AFTER); 5 − 1 = 4'.",
|
||||
"6. If step 2 or 3 cannot be done cleanly (no clear winner, overlapping timestamps, unclear event order), STOP and surface this as an UNRESOLVABLE ambiguity per the section above — do not fabricate a derivation.",
|
||||
"",
|
||||
"For simple factual lookups that don't involve conflict or arithmetic, you can answer directly without this scaffolding.",
|
||||
"",
|
||||
"## HIERARCHICAL RETRIEVAL STRATEGY",
|
||||
"",
|
||||
]
|
||||
)
|
||||
|
||||
# Build retrieval levels based on what's available
|
||||
# Assemble the retrieval-level blocks for whatever tools are exposed.
|
||||
# MM and Observations bodies are unconditional; recall's fallback wording
|
||||
# adapts to which upstream tools precede it (telling the LLM to fall back
|
||||
# to a tool that isn't in its list is the bug at the root of #1724).
|
||||
levels: list[tuple[str, list[str]]] = []
|
||||
if has_mental_models:
|
||||
parts.extend(
|
||||
levels.append(
|
||||
(
|
||||
"MENTAL MODELS (search_mental_models)",
|
||||
[
|
||||
"- User-curated summaries about specific topics",
|
||||
"- HIGHEST quality - manually created and maintained",
|
||||
"- If a relevant mental model exists and is FRESH, it may fully answer the question",
|
||||
"- Check `is_stale` field - if stale, also verify with lower levels",
|
||||
],
|
||||
)
|
||||
)
|
||||
if include_observations:
|
||||
levels.append(
|
||||
(
|
||||
"OBSERVATIONS (search_observations)",
|
||||
[
|
||||
"- Auto-consolidated knowledge from memories",
|
||||
"- Check `is_stale` field - if stale, ALSO use recall() to verify",
|
||||
"- Good for understanding patterns and summaries",
|
||||
],
|
||||
)
|
||||
)
|
||||
recall_body = ["- Individual memories (world facts and experiences)"]
|
||||
if has_mental_models and include_observations:
|
||||
recall_body.extend(
|
||||
[
|
||||
"You have access to THREE levels of knowledge. Use them in this order:",
|
||||
"",
|
||||
"### 1. MENTAL MODELS (search_mental_models) - Try First",
|
||||
"- User-curated summaries about specific topics",
|
||||
"- HIGHEST quality - manually created and maintained",
|
||||
"- If a relevant mental model exists and is FRESH, it may fully answer the question",
|
||||
"- Check `is_stale` field - if stale, also verify with lower levels",
|
||||
"",
|
||||
"### 2. OBSERVATIONS (search_observations) - Second Priority",
|
||||
"- Auto-consolidated knowledge from memories",
|
||||
"- Check `is_stale` field - if stale, ALSO use recall() to verify",
|
||||
"- Good for understanding patterns and summaries",
|
||||
"",
|
||||
"### 3. RAW FACTS (recall) - Ground Truth",
|
||||
"- Individual memories (world facts and experiences)",
|
||||
"- Use when: no mental models/observations exist, they're stale, or you need specific details",
|
||||
"- MANDATORY: If search_mental_models and search_observations both return 0 results, you MUST call recall() before giving up",
|
||||
"- This is the source of truth that other levels are built from",
|
||||
"",
|
||||
"**Tool result ordering:** `recall()` and `search_observations()` return their `memories` / `observations` arrays sorted by SEMANTIC RELEVANCE to the query, NOT by time. The POSITION of an entry tells you nothing about when it was retained. For any temporal reasoning — recency, supersession, applying events on top of a state — IGNORE the position and read the per-entry `mentioned_at` field (and `occurred_start` / `occurred_end` for events).",
|
||||
"",
|
||||
]
|
||||
)
|
||||
else:
|
||||
parts.extend(
|
||||
elif has_mental_models:
|
||||
recall_body.extend(
|
||||
[
|
||||
"- Use when: no mental model exists, it's stale, or you need specific details",
|
||||
"- MANDATORY: If search_mental_models returns 0 results, you MUST call recall() before giving up",
|
||||
"- This is the source of truth that mental models are built from",
|
||||
]
|
||||
)
|
||||
elif include_observations:
|
||||
recall_body.extend(
|
||||
[
|
||||
"You have access to TWO levels of knowledge. Use them in this order:",
|
||||
"",
|
||||
"### 1. OBSERVATIONS (search_observations) - Try First",
|
||||
"- Auto-consolidated knowledge from memories",
|
||||
"- Check `is_stale` field - if stale, ALSO use recall() to verify",
|
||||
"- Good for understanding patterns and summaries",
|
||||
"",
|
||||
"### 2. RAW FACTS (recall) - Ground Truth",
|
||||
"- Individual memories (world facts and experiences)",
|
||||
"- Use when: no observations exist, they're stale, or you need specific details",
|
||||
"- MANDATORY: If search_observations returns 0 results or count=0, you MUST call recall() before giving up",
|
||||
"- This is the source of truth that observations are built from",
|
||||
"",
|
||||
"**Tool result ordering:** `recall()` and `search_observations()` return their `memories` / `observations` arrays sorted by SEMANTIC RELEVANCE to the query, NOT by time. The POSITION of an entry tells you nothing about when it was retained. For any temporal reasoning — recency, supersession, applying events on top of a state — IGNORE the position and read the per-entry `mentioned_at` field (and `occurred_start` / `occurred_end` for events).",
|
||||
"",
|
||||
]
|
||||
)
|
||||
else:
|
||||
recall_body.extend(
|
||||
[
|
||||
"- MANDATORY: Call recall() to gather facts before giving up",
|
||||
"- This is the source of truth.",
|
||||
]
|
||||
)
|
||||
levels.append(("RAW FACTS (recall) - Ground Truth", recall_body))
|
||||
|
||||
# Position-dependent suffix for upstream tools; recall already carries its
|
||||
# fixed "- Ground Truth" suffix in the header text.
|
||||
suffixes = [""] * len(levels)
|
||||
if len(levels) >= 2:
|
||||
suffixes[0] = " - Try First"
|
||||
if len(levels) == 3:
|
||||
suffixes[1] = " - Second Priority"
|
||||
|
||||
if len(levels) == 1:
|
||||
parts.append("You have access to ONE level of knowledge:")
|
||||
else:
|
||||
word = "TWO" if len(levels) == 2 else "THREE"
|
||||
parts.append(f"You have access to {word} levels of knowledge. Use them in this order:")
|
||||
parts.append("")
|
||||
for idx, ((header, body), suffix) in enumerate(zip(levels, suffixes), 1):
|
||||
parts.append(f"### {idx}. {header}{suffix}")
|
||||
parts.extend(body)
|
||||
parts.append("")
|
||||
|
||||
parts.extend(
|
||||
[
|
||||
@@ -267,25 +355,28 @@ def build_system_prompt_for_tools(
|
||||
|
||||
parts.append("## Workflow")
|
||||
|
||||
steps: list[str] = []
|
||||
if has_mental_models:
|
||||
parts.extend(
|
||||
[
|
||||
"1. First, try search_mental_models() - check if a curated summary exists",
|
||||
"2. If no mental model or it's stale, try search_observations() for consolidated knowledge",
|
||||
"3. If observations are stale OR you need specific details, use recall() for raw facts",
|
||||
"4. Use expand() if you need more context on specific memories",
|
||||
"5. When ready, call done() with your answer and supporting IDs",
|
||||
]
|
||||
steps.append("First, try search_mental_models() - check if a curated summary exists")
|
||||
if include_observations:
|
||||
if has_mental_models:
|
||||
steps.append("If no mental model or it's stale, try search_observations() for consolidated knowledge")
|
||||
else:
|
||||
steps.append("First, try search_observations() - check for consolidated knowledge")
|
||||
# Recall step phrasing varies with whichever upstream tool(s) precede it.
|
||||
if include_observations:
|
||||
steps.append(
|
||||
"If observations are stale OR you need specific details, use recall() for raw facts"
|
||||
if has_mental_models
|
||||
else "If search_observations returns 0 results OR observations are stale, you MUST call recall() for raw facts"
|
||||
)
|
||||
elif has_mental_models:
|
||||
steps.append("If no mental model or it's stale, use recall() for raw facts")
|
||||
else:
|
||||
parts.extend(
|
||||
[
|
||||
"1. First, try search_observations() - check for consolidated knowledge",
|
||||
"2. If search_observations returns 0 results OR observations are stale, you MUST call recall() for raw facts",
|
||||
"3. Use expand() if you need more context on specific memories",
|
||||
"4. When ready, call done() with your answer and supporting IDs",
|
||||
]
|
||||
)
|
||||
steps.append("Call recall() to gather raw facts")
|
||||
steps.append("Use expand() if you need more context on specific memories")
|
||||
steps.append("When ready, call done() with your answer and supporting IDs")
|
||||
parts.extend(f"{idx}. {step}" for idx, step in enumerate(steps, 1))
|
||||
|
||||
parts.extend(
|
||||
[
|
||||
@@ -513,10 +604,16 @@ 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)
|
||||
def build_final_system_prompt(mission: str | None = None, llm_output_language: str | None = None) -> str:
|
||||
"""Build the final synthesis system prompt, using mission as role when set.
|
||||
|
||||
When ``llm_output_language`` is set, the response is forced into that
|
||||
language regardless of the query/source language.
|
||||
"""
|
||||
from hindsight_api.engine.prompt_utils import escape_for_prompt, output_language_directive
|
||||
|
||||
role_section = escape_for_prompt(mission.strip()) if mission else _DEFAULT_FINAL_ROLE
|
||||
return _FINAL_SYSTEM_PROMPT_BASE.format(role_section=role_section) + output_language_directive(llm_output_language)
|
||||
|
||||
|
||||
# Backward-compatible constant for non-identity missions
|
||||
|
||||
@@ -23,6 +23,24 @@ if TYPE_CHECKING:
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _prune_nulls(d: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Drop keys whose value is None or an empty collection (``""``, ``[]``, ``{}``).
|
||||
|
||||
Reflect tools dump ``MemoryFact`` / ``ObservationResult`` via ``model_dump()``,
|
||||
which emits every field including the many that are typically null or empty
|
||||
(``context``, ``occurred_start``, ``metadata``, ``tags``, etc.). Stripping
|
||||
these before serializing to JSON for the LLM cuts token cost and removes
|
||||
fields that aren't telling the model anything.
|
||||
|
||||
Callers that need the *presence* of a specific field as a signal (e.g.
|
||||
``source_fact_ids`` for drill-down) must ensure the value is non-empty —
|
||||
pass the upstream flag that populates it (e.g. ``source_facts_max_tokens``
|
||||
> 0 on ``tool_search_observations``) rather than relying on Pydantic
|
||||
emitting ``None``.
|
||||
"""
|
||||
return {k: v for k, v in d.items() if v is not None and v != "" and v != [] and v != {}}
|
||||
|
||||
|
||||
def _document_metadata_from_retain_params(retain_params: Any) -> dict[str, Any] | None:
|
||||
"""Return document metadata stored under retain_params.metadata."""
|
||||
if isinstance(retain_params, str):
|
||||
@@ -214,8 +232,8 @@ async def tool_search_observations(
|
||||
return {
|
||||
"query": query,
|
||||
"count": len(result.results),
|
||||
"observations": [m.model_dump() for m in result.results],
|
||||
"source_facts": {k: v.model_dump() for k, v in (result.source_facts or {}).items()},
|
||||
"observations": [_prune_nulls(m.model_dump()) for m in result.results],
|
||||
"source_facts": {k: _prune_nulls(v.model_dump()) for k, v in (result.source_facts or {}).items()},
|
||||
"is_stale": is_stale,
|
||||
"freshness": freshness,
|
||||
}
|
||||
@@ -282,8 +300,8 @@ async def tool_recall(
|
||||
|
||||
return {
|
||||
"query": query,
|
||||
"memories": [m.model_dump() for m in result.results],
|
||||
"chunks": {k: v.model_dump() for k, v in (result.chunks or {}).items()},
|
||||
"memories": [_prune_nulls(m.model_dump()) for m in result.results],
|
||||
"chunks": {k: _prune_nulls(v.model_dump()) for k, v in (result.chunks or {}).items()},
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -4,29 +4,54 @@ Embedding generation utilities for memory units.
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from typing import Literal, Protocol
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
EmbeddingInputType = Literal["document", "query"]
|
||||
|
||||
def generate_embedding(embeddings_backend, text: str) -> list[float]:
|
||||
|
||||
class EmbeddingsBackend(Protocol):
|
||||
"""Minimal duck-typed surface used by retain/recall — the concrete `Embeddings`
|
||||
ABC supplies default implementations that delegate to `encode()`."""
|
||||
|
||||
def encode_query(self, texts: list[str]) -> list[list[float]]: ...
|
||||
|
||||
def encode_documents(self, texts: list[str]) -> list[list[float]]: ...
|
||||
|
||||
|
||||
def generate_embedding(
|
||||
embeddings_backend: EmbeddingsBackend, text: str, input_type: EmbeddingInputType = "document"
|
||||
) -> list[float]:
|
||||
"""
|
||||
Generate embedding for text using the provided embeddings backend.
|
||||
|
||||
Args:
|
||||
embeddings_backend: Embeddings instance to use for encoding
|
||||
text: Text to embed
|
||||
input_type: Whether text is retained document text or recall/search query text.
|
||||
|
||||
Returns:
|
||||
Embedding vector (dimension depends on embeddings backend)
|
||||
"""
|
||||
try:
|
||||
embeddings = embeddings_backend.encode([text])
|
||||
embeddings = _encode_with_input_type(embeddings_backend, [text], input_type)
|
||||
return embeddings[0]
|
||||
except Exception as e:
|
||||
raise Exception(f"Failed to generate embedding: {str(e)}")
|
||||
|
||||
|
||||
async def generate_embeddings_batch(embeddings_backend, texts: list[str]) -> list[list[float]]:
|
||||
def _encode_with_input_type(
|
||||
embeddings_backend: EmbeddingsBackend, texts: list[str], input_type: EmbeddingInputType
|
||||
) -> list[list[float]]:
|
||||
if input_type == "query":
|
||||
return embeddings_backend.encode_query(texts)
|
||||
return embeddings_backend.encode_documents(texts)
|
||||
|
||||
|
||||
async def generate_embeddings_batch(
|
||||
embeddings_backend: EmbeddingsBackend, texts: list[str], input_type: EmbeddingInputType = "document"
|
||||
) -> list[list[float]]:
|
||||
"""
|
||||
Generate embeddings for multiple texts using the provided embeddings backend.
|
||||
|
||||
@@ -36,17 +61,14 @@ async def generate_embeddings_batch(embeddings_backend, texts: list[str]) -> lis
|
||||
Args:
|
||||
embeddings_backend: Embeddings instance to use for encoding
|
||||
texts: List of texts to embed
|
||||
input_type: Whether texts are retained documents or recall/search queries.
|
||||
|
||||
Returns:
|
||||
List of embeddings in same order as input texts
|
||||
"""
|
||||
try:
|
||||
loop = asyncio.get_event_loop()
|
||||
embeddings = await loop.run_in_executor(
|
||||
None,
|
||||
embeddings_backend.encode,
|
||||
texts,
|
||||
)
|
||||
embeddings = await loop.run_in_executor(None, _encode_with_input_type, embeddings_backend, texts, input_type)
|
||||
except Exception as e:
|
||||
raise Exception(f"Failed to generate batch embeddings: {str(e)}")
|
||||
|
||||
|
||||
@@ -1,13 +1,13 @@
|
||||
"""
|
||||
Entity processing for retain pipeline.
|
||||
|
||||
Handles entity extraction, resolution, and link creation for stored facts.
|
||||
Handles entity extraction and resolution for stored facts.
|
||||
"""
|
||||
|
||||
import logging
|
||||
|
||||
from . import link_utils
|
||||
from .types import EntityLink, ProcessedFact
|
||||
from .types import ProcessedFact
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -76,8 +76,7 @@ async def resolve_entities(
|
||||
entity_labels: Optional entity label taxonomy
|
||||
|
||||
Returns:
|
||||
Tuple of (resolved_entity_ids, entity_to_unit, unit_to_entity_ids)
|
||||
to pass to build_entity_links().
|
||||
Tuple of (resolved_entity_ids, entity_to_unit, unit_to_entity_ids).
|
||||
"""
|
||||
if not unit_ids or not facts:
|
||||
return [], [], {}
|
||||
@@ -99,68 +98,3 @@ async def resolve_entities(
|
||||
log_buffer,
|
||||
entity_labels=entity_labels,
|
||||
)
|
||||
|
||||
|
||||
async def build_entity_links(
|
||||
entity_resolver,
|
||||
conn,
|
||||
bank_id: str,
|
||||
unit_ids: list[str],
|
||||
resolved_entity_ids: list[str],
|
||||
entity_to_unit: list[tuple],
|
||||
unit_to_entity_ids: dict[str, list[str]],
|
||||
log_buffer: list[str] = None,
|
||||
skip_unit_entities_insert: bool = False,
|
||||
ops=None,
|
||||
) -> list[EntityLink]:
|
||||
"""
|
||||
Build entity links for UI graph visualization.
|
||||
|
||||
Queries unit_entities to find shared entities between new and existing units,
|
||||
then generates EntityLink objects. When called from Phase 3 (post-transaction),
|
||||
set skip_unit_entities_insert=True since unit_entities were already inserted
|
||||
in Phase 2.
|
||||
|
||||
Args:
|
||||
entity_resolver: EntityResolver instance
|
||||
conn: Database connection
|
||||
bank_id: Bank identifier
|
||||
unit_ids: Actual unit IDs (must already be inserted in the DB)
|
||||
resolved_entity_ids: From resolve_entities()
|
||||
entity_to_unit: From resolve_entities()
|
||||
unit_to_entity_ids: From resolve_entities()
|
||||
log_buffer: Optional buffer for detailed logging
|
||||
skip_unit_entities_insert: Skip unit_entities INSERT (already done in Phase 2)
|
||||
ops: DataAccessOps instance (from backend.ops)
|
||||
|
||||
Returns:
|
||||
List of EntityLink objects for batch insertion
|
||||
"""
|
||||
return await link_utils.build_entity_links_from_resolved(
|
||||
entity_resolver,
|
||||
conn,
|
||||
bank_id,
|
||||
unit_ids,
|
||||
resolved_entity_ids,
|
||||
entity_to_unit,
|
||||
unit_to_entity_ids,
|
||||
log_buffer,
|
||||
skip_unit_entities_insert=skip_unit_entities_insert,
|
||||
ops=ops,
|
||||
)
|
||||
|
||||
|
||||
async def insert_entity_links_batch(conn, entity_links: list[EntityLink], bank_id: str, ops=None) -> None:
|
||||
"""
|
||||
Insert entity links in batch.
|
||||
|
||||
Args:
|
||||
conn: Database connection
|
||||
entity_links: List of EntityLink objects
|
||||
bank_id: Bank identifier (stored directly on memory_links for fast filtering)
|
||||
ops: DataAccessOps instance (from backend.ops)
|
||||
"""
|
||||
if not entity_links:
|
||||
return
|
||||
|
||||
await link_utils.insert_entity_links_batch(conn, entity_links, bank_id, ops=ops)
|
||||
|
||||
@@ -888,13 +888,16 @@ def _build_extraction_prompt_and_schema(config) -> tuple[str, type]:
|
||||
extract_causal_links = config.retain_extract_causal_links
|
||||
|
||||
# Build retain_mission section if set - injected before the mode-specific guidelines
|
||||
# Escape braces so user-supplied text survives str.format() on the prompt template.
|
||||
from hindsight_api.engine.prompt_utils import escape_for_prompt
|
||||
|
||||
retain_mission = getattr(config, "retain_mission", None)
|
||||
if retain_mission:
|
||||
retain_mission_section = (
|
||||
f"══════════════════════════════════════════════════════════════════════════\n"
|
||||
f"FOCUS — What to retain for this bank\n"
|
||||
f"══════════════════════════════════════════════════════════════════════════\n\n"
|
||||
f"{retain_mission}\n\n"
|
||||
f"{escape_for_prompt(retain_mission)}\n\n"
|
||||
)
|
||||
else:
|
||||
retain_mission_section = ""
|
||||
@@ -910,7 +913,7 @@ def _build_extraction_prompt_and_schema(config) -> tuple[str, type]:
|
||||
base_prompt = CUSTOM_FACT_EXTRACTION_PROMPT
|
||||
prompt = base_prompt.format(
|
||||
retain_mission_section=retain_mission_section,
|
||||
custom_instructions=config.retain_custom_instructions,
|
||||
custom_instructions=escape_for_prompt(config.retain_custom_instructions),
|
||||
)
|
||||
elif extraction_mode == "verbose":
|
||||
prompt = VERBOSE_FACT_EXTRACTION_PROMPT.format(
|
||||
@@ -947,6 +950,16 @@ def _build_extraction_prompt_and_schema(config) -> tuple[str, type]:
|
||||
if labels_section:
|
||||
prompt = prompt + labels_section
|
||||
|
||||
# Force the LLM to emit fact text in the configured language, regardless of
|
||||
# the source content's language. Same directive is applied to consolidation
|
||||
# and reflect so HINDSIGHT_API_LLM_OUTPUT_LANGUAGE has a uniform effect
|
||||
# across the pipeline. This is independent of the BM25 indexing language
|
||||
# (HINDSIGHT_API_TEXT_SEARCH_EXTENSION_NATIVE_LANGUAGE) by design — search
|
||||
# tokenization and LLM output language are separate concerns.
|
||||
from ..prompt_utils import output_language_directive
|
||||
|
||||
prompt = prompt + output_language_directive(getattr(config, "llm_output_language", None))
|
||||
|
||||
response_schema = base_response_class
|
||||
|
||||
if labels_cfg and labels_cfg.attributes:
|
||||
@@ -1662,8 +1675,11 @@ async def extract_facts_from_contents_batch_api(
|
||||
|
||||
# Check if provider supports batch API
|
||||
if not await llm_config._provider_impl.supports_batch_api():
|
||||
logger.warning(f"Batch API not supported for provider {llm_config.provider}, falling back to sync mode")
|
||||
return await extract_facts_from_contents(contents, llm_config, agent_name, config, pool, operation_id, schema)
|
||||
raise RuntimeError(
|
||||
f"retain_batch_enabled=True but provider '{llm_config.provider}' does not "
|
||||
f"support the batch API. This should have been caught at startup — check "
|
||||
f"HINDSIGHT_API_RETAIN_BATCH_ENABLED and your LLM provider configuration."
|
||||
)
|
||||
|
||||
# Check if we're resuming an existing batch (crash recovery)
|
||||
batch_id = None
|
||||
|
||||
@@ -321,6 +321,15 @@ async def handle_document_tracking(
|
||||
f"[RETAIN] Document {document_id} re-ingested: invalidated "
|
||||
f"{invalidated} observation(s) derived from {len(existing_unit_ids)} outgoing memory_units"
|
||||
)
|
||||
# Capture link-recompute victims BEFORE the cascade. Same staleness
|
||||
# applies on upsert as on explicit delete: surviving units in OTHER
|
||||
# documents that linked to these doomed units are about to lose
|
||||
# those links. ``ops`` may be None for older callers that haven't
|
||||
# been wired up — skip enqueue in that case rather than crash.
|
||||
if ops is not None:
|
||||
from ..graph_maintenance import enqueue_relink_victims
|
||||
|
||||
await enqueue_relink_victims(conn, bank_id, [str(uid) for uid in existing_unit_ids], ops=ops)
|
||||
# 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
|
||||
|
||||
@@ -5,10 +5,8 @@ Link creation utilities for temporal, semantic, and entity links.
|
||||
import logging
|
||||
import time
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from uuid import UUID
|
||||
|
||||
from ..memory_engine import fq_table
|
||||
from .types import EntityLink
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -366,136 +364,6 @@ async def resolve_entities_only(
|
||||
return resolved_entity_ids, entity_to_unit, unit_to_entity_ids
|
||||
|
||||
|
||||
async def build_entity_links_from_resolved(
|
||||
entity_resolver,
|
||||
conn,
|
||||
bank_id: str,
|
||||
unit_ids: list[str],
|
||||
resolved_entity_ids: list[str],
|
||||
entity_to_unit: list[tuple],
|
||||
unit_to_entity_ids: dict[str, list[str]],
|
||||
log_buffer: list[str] = None,
|
||||
skip_unit_entities_insert: bool = False,
|
||||
ops=None,
|
||||
) -> list["EntityLink"]:
|
||||
"""
|
||||
Build entity links between units that share entities.
|
||||
|
||||
Queries unit_entities to find which existing units share entities with the
|
||||
new units, then generates EntityLink objects for UI graph visualization.
|
||||
|
||||
Args:
|
||||
entity_resolver: EntityResolver instance
|
||||
conn: Database connection
|
||||
bank_id: Bank identifier
|
||||
unit_ids: Actual unit IDs (must already be inserted in the DB)
|
||||
resolved_entity_ids: Entity IDs from resolve_entities_only
|
||||
entity_to_unit: Mapping from resolve_entities_only
|
||||
unit_to_entity_ids: Mapping from resolve_entities_only
|
||||
log_buffer: Optional logging buffer
|
||||
skip_unit_entities_insert: If True, skip unit_entities INSERT (already done in Phase 2)
|
||||
|
||||
Returns:
|
||||
List of EntityLink objects for batch insertion
|
||||
"""
|
||||
if not resolved_entity_ids:
|
||||
return []
|
||||
|
||||
if not skip_unit_entities_insert:
|
||||
# Insert unit-entity links (used in fallback path where Phase 2 didn't do this)
|
||||
substep_start = time.time()
|
||||
unit_entity_pairs = []
|
||||
for idx, (unit_id, _local_idx, fact_date) in enumerate(entity_to_unit):
|
||||
# Propagate the unit's fact_date so entity_cooccurrences.last_cooccurred
|
||||
# reflects the event timeline, not the ingest moment.
|
||||
unit_entity_pairs.append((unit_id, resolved_entity_ids[idx], fact_date))
|
||||
|
||||
await entity_resolver.link_units_to_entities_batch(unit_entity_pairs, conn=conn)
|
||||
_log(
|
||||
log_buffer,
|
||||
f" [6.2.3] Create unit-entity links (batched): {len(unit_entity_pairs)} links in {time.time() - substep_start:.3f}s",
|
||||
level="debug",
|
||||
)
|
||||
|
||||
# Build entity links between units that share entities
|
||||
substep_start = time.time()
|
||||
all_entity_ids = set()
|
||||
for entity_ids_list in unit_to_entity_ids.values():
|
||||
all_entity_ids.update(entity_ids_list)
|
||||
|
||||
_log(log_buffer, f" [6.3] Creating entity links for {len(all_entity_ids)} unique entities...", level="debug")
|
||||
|
||||
MAX_LINKS_PER_ENTITY = 10
|
||||
|
||||
entity_to_units = {}
|
||||
if all_entity_ids:
|
||||
query_start = time.time()
|
||||
import uuid
|
||||
|
||||
entity_id_list = [uuid.UUID(eid) if isinstance(eid, str) else eid for eid in all_entity_ids]
|
||||
limit_per_entity = MAX_LINKS_PER_ENTITY + len(unit_ids) # room for new units + existing cap
|
||||
|
||||
rows = await ops.fetch_entity_unit_fanout(
|
||||
conn,
|
||||
fq_table("unit_entities"),
|
||||
entity_id_list,
|
||||
limit_per_entity,
|
||||
)
|
||||
_log(
|
||||
log_buffer,
|
||||
f" [6.3.1] Query unit_entities (LATERAL): {len(rows)} rows in {time.time() - query_start:.3f}s",
|
||||
level="debug",
|
||||
)
|
||||
|
||||
group_start = time.time()
|
||||
for row in rows:
|
||||
entity_id = row["entity_id"]
|
||||
if entity_id not in entity_to_units:
|
||||
entity_to_units[entity_id] = []
|
||||
entity_to_units[entity_id].append(row["unit_id"])
|
||||
_log(log_buffer, f" [6.3.2] Group by entity_id: {time.time() - group_start:.3f}s", level="debug")
|
||||
link_gen_start = time.time()
|
||||
links: list[EntityLink] = []
|
||||
new_unit_set = set(unit_ids)
|
||||
|
||||
def to_uuid(val) -> UUID:
|
||||
return UUID(val) if isinstance(val, str) else val
|
||||
|
||||
for entity_id, units_with_entity in entity_to_units.items():
|
||||
entity_uuid = to_uuid(entity_id)
|
||||
new_units = [u for u in units_with_entity if str(u) in new_unit_set or u in new_unit_set]
|
||||
existing_units = [u for u in units_with_entity if str(u) not in new_unit_set and u not in new_unit_set]
|
||||
|
||||
new_units_to_link = new_units[-MAX_LINKS_PER_ENTITY:] if len(new_units) > MAX_LINKS_PER_ENTITY else new_units
|
||||
for i, unit_id_1 in enumerate(new_units_to_link):
|
||||
for unit_id_2 in new_units_to_link[i + 1 :]:
|
||||
links.append(
|
||||
EntityLink(from_unit_id=to_uuid(unit_id_1), to_unit_id=to_uuid(unit_id_2), entity_id=entity_uuid)
|
||||
)
|
||||
links.append(
|
||||
EntityLink(from_unit_id=to_uuid(unit_id_2), to_unit_id=to_uuid(unit_id_1), entity_id=entity_uuid)
|
||||
)
|
||||
|
||||
existing_to_link = existing_units[-MAX_LINKS_PER_ENTITY:]
|
||||
for new_unit in new_units:
|
||||
for existing_unit in existing_to_link:
|
||||
links.append(
|
||||
EntityLink(from_unit_id=to_uuid(new_unit), to_unit_id=to_uuid(existing_unit), entity_id=entity_uuid)
|
||||
)
|
||||
links.append(
|
||||
EntityLink(from_unit_id=to_uuid(existing_unit), to_unit_id=to_uuid(new_unit), entity_id=entity_uuid)
|
||||
)
|
||||
|
||||
_log(log_buffer, f" [6.3.3] Generate {len(links)} links: {time.time() - link_gen_start:.3f}s", level="debug")
|
||||
_log(
|
||||
log_buffer,
|
||||
f" [6.3] Entity link creation: {len(links)} links for {len(all_entity_ids)} unique entities in {time.time() - substep_start:.3f}s",
|
||||
level="debug",
|
||||
)
|
||||
|
||||
return links
|
||||
|
||||
|
||||
async def create_temporal_links_batch_per_fact(
|
||||
conn,
|
||||
bank_id: str,
|
||||
@@ -889,29 +757,6 @@ async def create_semantic_links_batch(
|
||||
raise
|
||||
|
||||
|
||||
async def insert_entity_links_batch(conn, links: list[EntityLink], bank_id: str, chunk_size: int = 5000, ops=None):
|
||||
"""
|
||||
Bulk-insert entity links via sorted INSERT FROM unnest().
|
||||
|
||||
Args:
|
||||
conn: Database connection
|
||||
links: List of EntityLink objects
|
||||
bank_id: Bank identifier (stored directly on memory_links for fast filtering)
|
||||
chunk_size: Number of rows per INSERT chunk (default 5000)
|
||||
"""
|
||||
if not links:
|
||||
return
|
||||
|
||||
import time as time_mod
|
||||
|
||||
total_start = time_mod.time()
|
||||
tuples = [(link.from_unit_id, link.to_unit_id, link.link_type, link.weight, link.entity_id) for link in links]
|
||||
await _bulk_insert_links(conn, tuples, bank_id=bank_id, chunk_size=chunk_size, ops=ops)
|
||||
logger.debug(
|
||||
f" [9.TOTAL] Entity links batch insert ({len(tuples)} rows): {time_mod.time() - total_start:.3f}s"
|
||||
)
|
||||
|
||||
|
||||
async def create_causal_links_batch(
|
||||
conn,
|
||||
bank_id: str,
|
||||
|
||||
@@ -100,7 +100,6 @@ from .types import (
|
||||
ChunkMetadata,
|
||||
EntityResolutionResult,
|
||||
Phase1Result,
|
||||
Phase3Context,
|
||||
ProcessedFact,
|
||||
RetainContent,
|
||||
RetainContentDict,
|
||||
@@ -108,6 +107,9 @@ from .types import (
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
RetainOutboxCallback = Callable[[asyncpg.Connection], Awaitable[None]]
|
||||
RetainOutboxCallbackFactory = Callable[[list[RetainContentDict]], RetainOutboxCallback | None]
|
||||
|
||||
|
||||
def _build_retain_params(contents_dicts, document_tags=None, doc_contents=None):
|
||||
"""Build retain_params and merged_tags from content dicts."""
|
||||
@@ -256,30 +258,27 @@ async def _insert_facts_and_links(
|
||||
skip_semantic_links: bool = False,
|
||||
outbox_callback=None,
|
||||
ops=None,
|
||||
) -> tuple[list[list[str]], Phase3Context]:
|
||||
) -> list[list[str]]:
|
||||
"""
|
||||
Phase 2 of the retain pipeline: insert facts and retrieval-critical links.
|
||||
|
||||
Runs inside a single database transaction to ensure atomicity of the data
|
||||
that retrieval depends on (facts, unit_entities, temporal/semantic/causal links).
|
||||
|
||||
Entity link generation and insertion for UI visualization are NOT done here —
|
||||
only the unit_entities INSERT (FK to memory_units) stays in the transaction.
|
||||
Entity link building is deferred to Phase 3 (post-transaction, best-effort).
|
||||
Entity edges for UI graph visualization are derived on demand from
|
||||
unit_entities by the /graph endpoint, so no entity rows are written to
|
||||
memory_links here.
|
||||
"""
|
||||
set_stage("retain.phase2.insert_facts")
|
||||
unit_ids = await fact_storage.insert_facts_batch(conn, bank_id, processed_facts, ops=ops)
|
||||
step_start = time.time()
|
||||
log_buffer.append(f" Insert facts: {len(unit_ids)} units in {time.time() - step_start:.3f}s")
|
||||
|
||||
# Context for Phase 3 entity link building (after transaction commits)
|
||||
phase3_context = Phase3Context()
|
||||
|
||||
if unit_ids:
|
||||
# Entity resolution was done in Phase 1 (separate connection).
|
||||
# Remap placeholder IDs to actual unit IDs.
|
||||
step_start = time.time()
|
||||
remapped_entity_to_unit, remapped_unit_to_entity_ids, remapped_semantic = _remap_phase1_results(
|
||||
remapped_entity_to_unit, _remapped_unit_to_entity_ids, remapped_semantic = _remap_phase1_results(
|
||||
resolved_entity_ids, entity_to_unit, unit_to_entity_ids, semantic_ann_links or [], unit_ids
|
||||
)
|
||||
# Update semantic_ann_links with remapped IDs for Phase 2
|
||||
@@ -293,13 +292,6 @@ async def _insert_facts_and_links(
|
||||
]
|
||||
await entity_resolver.link_units_to_entities_batch(unit_entity_pairs, conn=conn)
|
||||
log_buffer.append(f" Insert unit_entities: {len(unit_entity_pairs)} pairs in {time.time() - step_start:.3f}s")
|
||||
# Save context for Phase 3 entity link building (after commit)
|
||||
phase3_context = Phase3Context(
|
||||
unit_ids=unit_ids,
|
||||
resolved_entity_ids=resolved_entity_ids,
|
||||
entity_to_unit=remapped_entity_to_unit,
|
||||
unit_to_entity_ids=remapped_unit_to_entity_ids,
|
||||
)
|
||||
|
||||
# Create temporal links
|
||||
step_start = time.time()
|
||||
@@ -340,52 +332,10 @@ async def _insert_facts_and_links(
|
||||
# an IndexError (see issue #1037).
|
||||
result_unit_ids = _map_results_to_contents(contents, processed_facts, unit_ids if unit_ids else [])
|
||||
|
||||
if outbox_callback:
|
||||
if outbox_callback is not None:
|
||||
await outbox_callback(conn)
|
||||
|
||||
return result_unit_ids, phase3_context
|
||||
|
||||
|
||||
async def _build_and_insert_entity_links_phase3(
|
||||
pool: Any,
|
||||
entity_resolver,
|
||||
bank_id: str,
|
||||
phase3_ctx: Phase3Context,
|
||||
log_buffer: list[str],
|
||||
) -> None:
|
||||
"""
|
||||
Phase 3 helper: build entity links from resolved data and insert them.
|
||||
|
||||
Runs on a fresh connection after the main transaction has committed.
|
||||
Entity links are for UI graph visualization only — retrieval uses
|
||||
the unit_entities self-join instead.
|
||||
"""
|
||||
set_stage("retain.phase3.entity_links")
|
||||
p3_unit_ids = phase3_ctx.unit_ids
|
||||
p3_resolved = phase3_ctx.resolved_entity_ids
|
||||
p3_entity_to_unit = phase3_ctx.entity_to_unit
|
||||
p3_unit_to_entity_ids = phase3_ctx.unit_to_entity_ids
|
||||
|
||||
if not p3_unit_ids or not p3_resolved:
|
||||
return
|
||||
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
step_start = time.time()
|
||||
entity_links = await entity_processing.build_entity_links(
|
||||
entity_resolver,
|
||||
conn,
|
||||
bank_id,
|
||||
p3_unit_ids,
|
||||
p3_resolved,
|
||||
p3_entity_to_unit,
|
||||
p3_unit_to_entity_ids,
|
||||
log_buffer,
|
||||
skip_unit_entities_insert=True, # Already inserted in Phase 2
|
||||
ops=pool.ops,
|
||||
)
|
||||
if entity_links:
|
||||
await entity_processing.insert_entity_links_batch(conn, entity_links, bank_id, ops=pool.ops)
|
||||
log_buffer.append(f" Entity links (viz): {len(entity_links)} links in {time.time() - step_start:.3f}s")
|
||||
return result_unit_ids
|
||||
|
||||
|
||||
async def _extract_and_embed(
|
||||
@@ -449,7 +399,8 @@ async def retain_batch(
|
||||
document_tags: list[str] | None = None,
|
||||
operation_id: str | None = None,
|
||||
schema: str | None = None,
|
||||
outbox_callback: Callable[["asyncpg.Connection"], Awaitable[None]] | None = None,
|
||||
outbox_callback: RetainOutboxCallback | None = None,
|
||||
outbox_callback_factory: RetainOutboxCallbackFactory | None = None,
|
||||
db_semaphore: "asyncio.Semaphore | None" = None,
|
||||
) -> tuple[list[list[str]], TokenUsage, int | None]:
|
||||
"""
|
||||
@@ -507,6 +458,10 @@ async def retain_batch(
|
||||
total_usage = TokenUsage()
|
||||
total_processed_tokens: int | None = 0
|
||||
for doc_key, (group_dicts, group_contents) in groups.items():
|
||||
group_outbox_callback = (
|
||||
outbox_callback_factory(group_dicts) if outbox_callback_factory is not None else outbox_callback
|
||||
)
|
||||
|
||||
group_ids, group_usage, group_processed = await retain_batch(
|
||||
pool=pool,
|
||||
embeddings_model=embeddings_model,
|
||||
@@ -522,7 +477,8 @@ async def retain_batch(
|
||||
document_tags=document_tags,
|
||||
operation_id=operation_id,
|
||||
schema=schema,
|
||||
outbox_callback=outbox_callback,
|
||||
outbox_callback=group_outbox_callback,
|
||||
outbox_callback_factory=outbox_callback_factory,
|
||||
db_semaphore=db_semaphore,
|
||||
)
|
||||
for group_idx, orig_idx in enumerate(original_indices[doc_key]):
|
||||
@@ -1191,7 +1147,6 @@ async def _streaming_retain_batch(
|
||||
|
||||
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 ---
|
||||
@@ -1280,7 +1235,7 @@ async def _streaming_retain_batch(
|
||||
|
||||
# Insert facts and links — skip semantic links entirely in streaming
|
||||
# mode; they are created in a single final ANN pass after all batches.
|
||||
batch_result_ids, phase3_ctx = await _insert_facts_and_links(
|
||||
batch_result_ids = await _insert_facts_and_links(
|
||||
conn,
|
||||
entity_resolver,
|
||||
bank_id,
|
||||
@@ -1300,15 +1255,13 @@ 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)
|
||||
# Best-effort: flush entity_cooccurrences and other deferred stats.
|
||||
try:
|
||||
await entity_resolver.flush_pending_stats()
|
||||
except Exception:
|
||||
logger.warning(
|
||||
f"Entity stats flush (consumer batch {consumer_batch_idx + 1}) failed", exc_info=True
|
||||
)
|
||||
|
||||
logger.info(
|
||||
f"[streaming] Consumer batch {consumer_batch_idx + 1} total "
|
||||
@@ -1766,7 +1719,7 @@ async def _try_delta_retain(
|
||||
# Insert facts and retrieval-critical links.
|
||||
# Use delta_contents (the changed/new chunks) as the content list,
|
||||
# since extracted_facts have content_index relative to delta_contents.
|
||||
result_unit_ids, phase3_ctx = await _insert_facts_and_links(
|
||||
result_unit_ids = await _insert_facts_and_links(
|
||||
conn,
|
||||
entity_resolver,
|
||||
bank_id,
|
||||
@@ -1783,12 +1736,11 @@ async def _try_delta_retain(
|
||||
ops=pool.ops,
|
||||
)
|
||||
|
||||
# PHASE 3 — Best-Effort Display Data (post-transaction)
|
||||
# Flush deferred entity_cooccurrences stats (post-transaction, best-effort).
|
||||
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("Phase 3 (best-effort display data) failed — retrieval unaffected", exc_info=True)
|
||||
logger.warning("Entity stats flush failed — retrieval unaffected", exc_info=True)
|
||||
|
||||
total_time = time.time() - start_time
|
||||
log_buffer.append(f"{'=' * 60}")
|
||||
@@ -1844,7 +1796,7 @@ async def _delta_metadata_only(
|
||||
merged_tags,
|
||||
)
|
||||
await fact_storage.update_memory_units_tags(conn, bank_id, document_id, merged_tags)
|
||||
if outbox_callback:
|
||||
if outbox_callback is not None:
|
||||
await outbox_callback(conn)
|
||||
|
||||
total_time = time.time() - start_time
|
||||
|
||||
@@ -224,21 +224,6 @@ class ProcessedFact:
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Phase3Context:
|
||||
"""
|
||||
Data passed from Phase 2 to Phase 3 for entity link building.
|
||||
|
||||
Contains the unit IDs and entity resolution data needed to build
|
||||
entity links for UI graph visualization after the write transaction commits.
|
||||
"""
|
||||
|
||||
unit_ids: list[str] = field(default_factory=list)
|
||||
resolved_entity_ids: list[str] = field(default_factory=list)
|
||||
entity_to_unit: list[tuple] = field(default_factory=list)
|
||||
unit_to_entity_ids: dict[str, list[str]] = field(default_factory=dict)
|
||||
|
||||
|
||||
@dataclass
|
||||
class EntityResolutionResult:
|
||||
"""
|
||||
@@ -263,21 +248,6 @@ class Phase1Result:
|
||||
semantic_ann_links: list[tuple]
|
||||
|
||||
|
||||
@dataclass
|
||||
class EntityLink:
|
||||
"""
|
||||
Link between two memory units through a shared entity.
|
||||
|
||||
Used for entity-based graph connections in the memory graph.
|
||||
"""
|
||||
|
||||
from_unit_id: UUID
|
||||
to_unit_id: UUID
|
||||
entity_id: UUID
|
||||
link_type: str = "entity"
|
||||
weight: float = 1.0
|
||||
|
||||
|
||||
@dataclass
|
||||
class RetainBatch:
|
||||
"""
|
||||
|
||||
@@ -283,12 +283,12 @@ class LinkExpansionRetriever(GraphRetriever):
|
||||
score transformations. The three CTEs share one connection slot — important
|
||||
for asyncpg which does not allow concurrent queries on the same connection.
|
||||
|
||||
Index coverage (requires migration d2e3f4a5b6c7):
|
||||
entity: idx_memory_links_entity_covering (from_unit_id) INCLUDE (to_unit_id, entity_id)
|
||||
WHERE link_type = 'entity' → index-only scan, no heap reads
|
||||
semantic incoming:
|
||||
idx_memory_links_to_type_weight (to_unit_id, link_type, weight DESC)
|
||||
→ replaces costly BitmapAnd of two separate scans
|
||||
Index coverage:
|
||||
entity: idx_unit_entities_entity_unit (entity_id, unit_id) — entity
|
||||
expansion traverses unit_entities, not memory_links.
|
||||
semantic: idx_memory_links_from_type_weight / _to_type_weight
|
||||
(from_unit_id|to_unit_id, link_type, weight DESC) serve both
|
||||
outgoing and incoming sides as single composite index scans.
|
||||
"""
|
||||
config = get_config()
|
||||
ml = fq_table("memory_links")
|
||||
|
||||
@@ -212,16 +212,47 @@ class CrossEncoderReranker:
|
||||
# Get cross-encoder scores
|
||||
scores = await self.cross_encoder.predict(pairs)
|
||||
|
||||
# Normalize scores using sigmoid to [0, 1] range
|
||||
# Cross-encoder returns logits which can be negative
|
||||
import math
|
||||
|
||||
# Normalize scores to [0, 1] range.
|
||||
# Local models return logits (any real number) — sigmoid is appropriate.
|
||||
# External API rerankers (SiliconFlow, Cohere, etc.) return pre-normalized
|
||||
# relevance_score in [0, 1] with very small absolute values. Applying
|
||||
# sigmoid to these compresses everything to ~0.5, destroying the ranking
|
||||
# signal and making recency the sole sorting factor. We detect the score range
|
||||
# and choose the appropriate normalization.
|
||||
import numpy as np
|
||||
|
||||
def sigmoid(x):
|
||||
def _sigmoid(x: float) -> float:
|
||||
return 1 / (1 + np.exp(-x))
|
||||
|
||||
normalized_scores = [sigmoid(score) for score in scores]
|
||||
def _rank_normalize_with_ties(score_list: list[float]) -> list[float]:
|
||||
"""Rank-based normalization that assigns equal ranks to equal scores."""
|
||||
n = len(score_list)
|
||||
if n <= 1:
|
||||
return [1.0] * n
|
||||
indexed = sorted(enumerate(score_list), key=lambda x: x[1], reverse=True)
|
||||
result = [0.0] * n
|
||||
i = 0
|
||||
while i < n:
|
||||
j = i
|
||||
while j < n and indexed[j][1] == indexed[i][1]:
|
||||
j += 1
|
||||
# Average rank for tied scores
|
||||
avg_rank = (i + j - 1) / 2.0
|
||||
norm = 1.0 - (avg_rank / (n - 1))
|
||||
for k in range(i, j):
|
||||
result[indexed[k][0]] = norm
|
||||
i = j
|
||||
return result
|
||||
|
||||
if scores and min(scores) >= 0.0 and max(scores) <= 1.0:
|
||||
# Scores are already in [0, 1] (e.g. SiliconFlow, Cohere relevance_score).
|
||||
# Use rank-based normalization to preserve relative ordering without
|
||||
# depending on absolute score magnitudes.
|
||||
normalized_scores = _rank_normalize_with_ties(scores)
|
||||
else:
|
||||
# Scores are logits (e.g. local sentence-transformers models).
|
||||
# Sigmoid maps (-inf, +inf) to (0, 1).
|
||||
normalized_scores = [_sigmoid(score) for score in scores]
|
||||
|
||||
# Create ScoredResult objects with cross-encoder scores
|
||||
scored_results = []
|
||||
|
||||
@@ -225,6 +225,7 @@ async def retrieve_semantic_bm25_combined(
|
||||
groups_clause=groups_clause,
|
||||
arm_index=i,
|
||||
text_search_extension=text_ext,
|
||||
bm25_language=config.text_search_extension_native_language,
|
||||
extra_where=created_range_clause,
|
||||
)
|
||||
)
|
||||
@@ -465,6 +466,8 @@ async def retrieve_temporal_combined(
|
||||
best_date = ep["mentioned_at"]
|
||||
|
||||
if best_date:
|
||||
if best_date.tzinfo is None:
|
||||
best_date = best_date.replace(tzinfo=UTC)
|
||||
days_from_mid = abs((best_date - mid_date).total_seconds() / 86400)
|
||||
temporal_proximity = 1.0 - min(days_from_mid / (total_days / 2), 1.0) if total_days > 0 else 1.0
|
||||
else:
|
||||
@@ -558,6 +561,8 @@ async def retrieve_temporal_combined(
|
||||
neighbor_best_date = n["mentioned_at"]
|
||||
|
||||
if neighbor_best_date:
|
||||
if neighbor_best_date.tzinfo is None:
|
||||
neighbor_best_date = neighbor_best_date.replace(tzinfo=UTC)
|
||||
days_from_mid = abs((neighbor_best_date - mid_date).total_seconds() / 86400)
|
||||
neighbor_temporal_proximity = (
|
||||
1.0 - min(days_from_mid / (total_days / 2), 1.0) if total_days > 0 else 1.0
|
||||
|
||||
@@ -407,6 +407,7 @@ class SQLDialect(ABC):
|
||||
groups_clause: str = "",
|
||||
arm_index: int = 0,
|
||||
text_search_extension: str = "native",
|
||||
bm25_language: str = "english",
|
||||
extra_where: str = "",
|
||||
) -> str:
|
||||
"""Build a BM25/full-text search subquery arm.
|
||||
@@ -426,7 +427,9 @@ class SQLDialect(ABC):
|
||||
arm_index: Index of this arm in the UNION ALL (used by Oracle for
|
||||
unique SCORE labels).
|
||||
text_search_extension: Full-text search backend ("native", "vchord",
|
||||
"pg_textsearch"). Only relevant for PostgreSQL.
|
||||
"pg_textsearch", "pgroonga"). Only relevant for PostgreSQL.
|
||||
bm25_language: PostgreSQL text search dictionary used by the native
|
||||
backend (e.g. "english", "french"). Ignored by other backends.
|
||||
extra_where: Optional additional WHERE clause fragment (e.g. time range filter).
|
||||
"""
|
||||
...
|
||||
|
||||
@@ -270,6 +270,7 @@ class OracleDialect(SQLDialect):
|
||||
groups_clause: str = "",
|
||||
arm_index: int = 0,
|
||||
text_search_extension: str = "native",
|
||||
bm25_language: str = "english",
|
||||
extra_where: str = "",
|
||||
) -> str:
|
||||
# Oracle Text: CONTAINS() / SCORE() with the CTXSYS.CONTEXT index.
|
||||
|
||||
@@ -182,6 +182,7 @@ class PostgreSQLDialect(SQLDialect):
|
||||
groups_clause: str = "",
|
||||
arm_index: int = 0,
|
||||
text_search_extension: str = "native",
|
||||
bm25_language: str = "english",
|
||||
extra_where: str = "",
|
||||
) -> str:
|
||||
if text_search_extension == "vchord":
|
||||
@@ -193,10 +194,35 @@ class PostgreSQLDialect(SQLDialect):
|
||||
bm25_score_expr = f"-({text_param} <@> to_bm25query({text_param}, 'idx_memory_units_text_search'))"
|
||||
bm25_order_by = f"text <@> to_bm25query({text_param}, 'idx_memory_units_text_search') ASC"
|
||||
bm25_where_filter = ""
|
||||
else: # native tsvector
|
||||
bm25_score_expr = f"ts_rank_cd(search_vector, to_tsquery('english', {text_param}))"
|
||||
elif text_search_extension == "pgroonga":
|
||||
# &@~ accepts pgroonga's query syntax (raw query text). pgroonga_score
|
||||
# returns a non-negative relevance score (higher = better).
|
||||
bm25_score_expr = "pgroonga_score(tableoid, ctid)"
|
||||
bm25_order_by = f"{bm25_score_expr} DESC"
|
||||
bm25_where_filter = f"AND search_vector @@ to_tsquery('english', {text_param})"
|
||||
bm25_where_filter = (
|
||||
f"AND (COALESCE(text, '') || ' ' || COALESCE(context, '') || ' ' || COALESCE(text_signals, '')) "
|
||||
f"&@~ {text_param}"
|
||||
)
|
||||
elif text_search_extension == "pg_search":
|
||||
# ParadeDB pg_search: BM25 index over (id, text, context, text_signals)
|
||||
# with key_field='id'. The @@@ operator on the key_field requires a
|
||||
# field-qualified query (`text:foo`); to keep the bind-parameter form,
|
||||
# we fan the query out across all indexed text fields with paradedb.boolean.
|
||||
bm25_score_expr = "paradedb.score(id)"
|
||||
bm25_order_by = "paradedb.score(id) DESC"
|
||||
bm25_where_filter = (
|
||||
f"AND id @@@ paradedb.boolean(should => ARRAY["
|
||||
f"paradedb.match('text', {text_param}), "
|
||||
f"paradedb.match('context', {text_param}), "
|
||||
f"paradedb.match('text_signals', {text_param})"
|
||||
f"])"
|
||||
)
|
||||
else: # native tsvector
|
||||
# bm25_language is validated as a PG identifier in HindsightConfig.validate(),
|
||||
# so embedding it as a SQL literal here is safe.
|
||||
bm25_score_expr = f"ts_rank_cd(search_vector, to_tsquery('{bm25_language}', {text_param}))"
|
||||
bm25_order_by = f"{bm25_score_expr} DESC"
|
||||
bm25_where_filter = f"AND search_vector @@ to_tsquery('{bm25_language}', {text_param})"
|
||||
|
||||
return (
|
||||
f"(SELECT {cols},"
|
||||
@@ -221,7 +247,7 @@ class PostgreSQLDialect(SQLDialect):
|
||||
*,
|
||||
text_search_extension: str = "native",
|
||||
) -> str:
|
||||
if text_search_extension in ("vchord", "pg_textsearch"):
|
||||
if text_search_extension in ("vchord", "pg_textsearch", "pgroonga", "pg_search"):
|
||||
return query_text
|
||||
# native tsvector: join tokens with OR operator
|
||||
return " | ".join(tokens)
|
||||
|
||||
@@ -16,7 +16,7 @@ class S3FileStorage(FileStorage):
|
||||
S3-compatible object storage backend.
|
||||
|
||||
Uses obstore (Rust-backed) for high-throughput async access to
|
||||
Amazon S3, MinIO, Cloudflare R2, and other S3-compliant APIs.
|
||||
Amazon S3, MinIO, Cloudflare R2, Tigris, and other S3-compliant APIs.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
|
||||
@@ -40,6 +40,7 @@ from hindsight_api.extensions.operation_validator import (
|
||||
# Core operations
|
||||
OperationValidationError,
|
||||
OperationValidatorExtension,
|
||||
PrecheckContext,
|
||||
RecallContext,
|
||||
RecallResult,
|
||||
ReflectContext,
|
||||
@@ -72,6 +73,7 @@ __all__ = [
|
||||
"DeferOperation",
|
||||
"OperationValidationError",
|
||||
"OperationValidatorExtension",
|
||||
"PrecheckContext",
|
||||
"RecallContext",
|
||||
"RecallResult",
|
||||
"ReflectContext",
|
||||
|
||||
@@ -146,7 +146,11 @@ class DefaultExtensionContext(ExtensionContext):
|
||||
|
||||
# Ensure text search columns/indexes match the configured extension
|
||||
await asyncio.to_thread(
|
||||
ensure_text_search_extension, db_url, text_search_extension=config.text_search_extension, schema=schema
|
||||
ensure_text_search_extension,
|
||||
db_url,
|
||||
text_search_extension=config.text_search_extension,
|
||||
pg_search_tokenizer=config.text_search_extension_pg_search_tokenizer,
|
||||
schema=schema,
|
||||
)
|
||||
|
||||
def get_memory_engine(self) -> "MemoryEngineInterface":
|
||||
|
||||
@@ -82,6 +82,33 @@ class ValidationResult:
|
||||
# =============================================================================
|
||||
|
||||
|
||||
@dataclass
|
||||
class PrecheckContext:
|
||||
"""Context for a pre-body-parse precheck on an operation.
|
||||
|
||||
Unlike :class:`RetainContext` / :class:`RecallContext` / etc., this
|
||||
context is constructed *before* the request body is deserialised. It
|
||||
therefore intentionally carries only the cheap, already-resolved
|
||||
pieces of request state:
|
||||
|
||||
- ``operation``: a short string identifying the route, e.g. ``"retain"``,
|
||||
``"recall"``, ``"reflect"``, ``"files_retain"``, ``"mental_model_create"``,
|
||||
``"mental_model_refresh"``.
|
||||
- ``bank_id``: parsed from the URL path.
|
||||
- ``request_context``: the authenticated :class:`RequestContext` (tenant
|
||||
already resolved by the tenant extension).
|
||||
|
||||
Implementations should keep precheck cheap and side-effect-free. The
|
||||
full per-request validators (``validate_retain`` / ``validate_recall``
|
||||
/ ``validate_reflect``) still run after the body is parsed and remain
|
||||
the source of truth for the precise per-call cost / quota arithmetic.
|
||||
"""
|
||||
|
||||
operation: str
|
||||
bank_id: str
|
||||
request_context: "RequestContext"
|
||||
|
||||
|
||||
@dataclass
|
||||
class RetainContext:
|
||||
"""Context for a retain operation validation (pre-operation).
|
||||
@@ -407,6 +434,42 @@ class OperationValidatorExtension(Extension, ABC):
|
||||
- consolidate (mental models consolidation)
|
||||
"""
|
||||
|
||||
# =========================================================================
|
||||
# Pre-body-parse hook (optional - default no-op)
|
||||
# =========================================================================
|
||||
|
||||
async def precheck(self, ctx: PrecheckContext) -> ValidationResult:
|
||||
"""
|
||||
Cheap pre-body-parse check, called before the request body is read.
|
||||
|
||||
FastAPI resolves ``Depends`` callables before deserialising the route
|
||||
body; routes that wire ``precheck`` as a dependency therefore short
|
||||
-circuit here without ever materialising the JSON payload in memory.
|
||||
That makes this the right hook for "should this caller be allowed to
|
||||
spend resources on this request at all" decisions — e.g. a balance
|
||||
is exhausted, a key is revoked, or a tenant is rate-limited.
|
||||
|
||||
Implementations should:
|
||||
- Be cheap: prefer cached lookups, avoid heavy DB queries.
|
||||
- Use only data on ``ctx`` (operation name + bank_id + request_context);
|
||||
the body is not yet available.
|
||||
- Be conservative on errors: prefer ``ValidationResult.accept()`` so
|
||||
a transient lookup failure doesn't turn into a request rejection.
|
||||
The post-body ``validate_*`` hooks still run and remain the source
|
||||
of truth for the precise per-call cost check.
|
||||
|
||||
Default implementation accepts everything. Override to opt in.
|
||||
|
||||
Args:
|
||||
ctx: Pre-body context with operation name, bank_id, and
|
||||
request_context (tenant already resolved).
|
||||
|
||||
Returns:
|
||||
ValidationResult indicating whether the request may proceed to
|
||||
body parsing and the post-parse validators.
|
||||
"""
|
||||
return ValidationResult.accept()
|
||||
|
||||
# =========================================================================
|
||||
# Pre-operation validation hooks (abstract - must be implemented)
|
||||
# =========================================================================
|
||||
|
||||
@@ -24,7 +24,15 @@ import uvicorn
|
||||
from . import MemoryEngine, __version__
|
||||
from .api import create_app
|
||||
from .banner import print_banner
|
||||
from .config import DEFAULT_WORKERS, ENV_HOST, ENV_WORKERS, HindsightConfig, _get_raw_config
|
||||
from .config import (
|
||||
DEFAULT_ACCESS_LOG,
|
||||
DEFAULT_WORKERS,
|
||||
ENV_ACCESS_LOG,
|
||||
ENV_HOST,
|
||||
ENV_WORKERS,
|
||||
HindsightConfig,
|
||||
_get_raw_config,
|
||||
)
|
||||
from .daemon import (
|
||||
DEFAULT_DAEMON_PORT,
|
||||
DEFAULT_IDLE_TIMEOUT,
|
||||
@@ -120,9 +128,18 @@ def main():
|
||||
)
|
||||
|
||||
# Access log options
|
||||
parser.add_argument("--access-log", action="store_true", help="Enable access log")
|
||||
parser.add_argument("--no-access-log", dest="access_log", action="store_false", help="Disable access log (default)")
|
||||
parser.set_defaults(access_log=False)
|
||||
parser.add_argument(
|
||||
"--access-log",
|
||||
action="store_true",
|
||||
default=os.getenv(ENV_ACCESS_LOG, "").lower() in ("1", "true", "yes", "on") or DEFAULT_ACCESS_LOG,
|
||||
help=f"Enable access log (env: {ENV_ACCESS_LOG}, default: {DEFAULT_ACCESS_LOG})",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--no-access-log",
|
||||
dest="access_log",
|
||||
action="store_false",
|
||||
help="Disable access log (overrides env and default)",
|
||||
)
|
||||
|
||||
# Proxy options
|
||||
parser.add_argument(
|
||||
|
||||
@@ -44,6 +44,7 @@ _ALL_TOOLS: frozenset[str] = frozenset(
|
||||
"update_mental_model",
|
||||
"delete_mental_model",
|
||||
"refresh_mental_model",
|
||||
"clear_mental_model",
|
||||
"list_directives",
|
||||
"create_directive",
|
||||
"delete_directive",
|
||||
@@ -221,6 +222,7 @@ def register_mcp_tools(
|
||||
"update_mental_model",
|
||||
"delete_mental_model",
|
||||
"refresh_mental_model",
|
||||
"clear_mental_model",
|
||||
"list_directives",
|
||||
"create_directive",
|
||||
"delete_directive",
|
||||
@@ -277,6 +279,9 @@ def register_mcp_tools(
|
||||
if "refresh_mental_model" in tools_to_register:
|
||||
_register_refresh_mental_model(mcp, memory, config)
|
||||
|
||||
if "clear_mental_model" in tools_to_register:
|
||||
_register_clear_mental_model(mcp, memory, config)
|
||||
|
||||
# Directive tools
|
||||
if "list_directives" in tools_to_register:
|
||||
_register_list_directives(mcp, memory, config)
|
||||
@@ -438,6 +443,7 @@ _AUDITABLE_MCP_TOOLS: frozenset[str] = frozenset(
|
||||
"update_mental_model",
|
||||
"delete_mental_model",
|
||||
"refresh_mental_model",
|
||||
"clear_mental_model",
|
||||
"create_directive",
|
||||
"delete_directive",
|
||||
"delete_document",
|
||||
@@ -922,6 +928,7 @@ def _register_reflect(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig
|
||||
response_schema: dict | None = None,
|
||||
tags: list[str] | None = None,
|
||||
tags_match: str = "any",
|
||||
include_based_on: bool = False,
|
||||
bank_id: str | None = None,
|
||||
) -> str:
|
||||
"""
|
||||
@@ -951,6 +958,7 @@ def _register_reflect(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig
|
||||
response_schema: Optional JSON schema for structured output. When provided, the response includes a 'structured_output' field.
|
||||
tags: Optional tags to filter memories by (e.g., ['project:alpha'])
|
||||
tags_match: How to match tags - 'any' (match any tag) or 'all' (match all tags). Default: 'any'
|
||||
include_based_on: Include source facts used for synthesis. Defaults to false because broad reflections can exceed MCP client result limits.
|
||||
bank_id: Optional bank to reflect in (defaults to session bank). Use for cross-bank operations.
|
||||
"""
|
||||
try:
|
||||
@@ -978,6 +986,8 @@ def _register_reflect(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig
|
||||
reflect_result = await memory.reflect_async(**reflect_kwargs)
|
||||
|
||||
result_data = json.loads(reflect_result.model_dump_json(indent=2))
|
||||
if not include_based_on:
|
||||
result_data.pop("based_on", None)
|
||||
if response_schema is not None and hasattr(reflect_result, "structured_output"):
|
||||
result_data["structured_output"] = reflect_result.structured_output
|
||||
return json.dumps(result_data, indent=2)
|
||||
@@ -999,6 +1009,7 @@ def _register_reflect(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig
|
||||
response_schema: dict | None = None,
|
||||
tags: list[str] | None = None,
|
||||
tags_match: str = "any",
|
||||
include_based_on: bool = False,
|
||||
) -> dict:
|
||||
"""
|
||||
Generate thoughtful analysis by synthesizing stored memories with the bank's personality.
|
||||
@@ -1027,6 +1038,7 @@ def _register_reflect(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig
|
||||
response_schema: Optional JSON schema for structured output. When provided, the response includes a 'structured_output' field.
|
||||
tags: Optional tags to filter memories by (e.g., ['project:alpha'])
|
||||
tags_match: How to match tags - 'any' (match any tag) or 'all' (match all tags). Default: 'any'
|
||||
include_based_on: Include source facts used for synthesis. Defaults to false because broad reflections can exceed MCP client result limits.
|
||||
"""
|
||||
try:
|
||||
target_bank = config.bank_id_resolver()
|
||||
@@ -1053,6 +1065,8 @@ def _register_reflect(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig
|
||||
reflect_result = await memory.reflect_async(**reflect_kwargs)
|
||||
|
||||
result_data = reflect_result.model_dump()
|
||||
if not include_based_on:
|
||||
result_data.pop("based_on", None)
|
||||
if response_schema is not None and hasattr(reflect_result, "structured_output"):
|
||||
result_data["structured_output"] = reflect_result.structured_output
|
||||
return result_data
|
||||
@@ -1765,6 +1779,98 @@ def _register_refresh_mental_model(mcp: FastMCP, memory: MemoryEngine, config: M
|
||||
return {"error": str(e)}
|
||||
|
||||
|
||||
def _register_clear_mental_model(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig) -> None:
|
||||
"""Register the clear_mental_model tool."""
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
async def clear_mental_model(
|
||||
mental_model_id: str,
|
||||
bank_id: str | None = None,
|
||||
) -> str:
|
||||
"""
|
||||
Clear a mental model's content so the next refresh performs a full re-synthesis.
|
||||
|
||||
This is useful for delta-mode models that have accumulated drift over many
|
||||
incremental refreshes. After clearing, call refresh_mental_model to trigger
|
||||
a clean full rebuild.
|
||||
|
||||
Args:
|
||||
mental_model_id: The ID of the mental model to clear
|
||||
bank_id: Optional bank (defaults to session bank). Use for cross-bank operations.
|
||||
"""
|
||||
try:
|
||||
target_bank = bank_id or config.bank_id_resolver()
|
||||
if target_bank is None:
|
||||
return '{"error": "No bank_id configured"}'
|
||||
|
||||
result = await memory.clear_mental_model(
|
||||
bank_id=target_bank,
|
||||
mental_model_id=mental_model_id,
|
||||
request_context=_get_request_context(config),
|
||||
)
|
||||
if result is None:
|
||||
return json.dumps({"error": f"Mental model '{mental_model_id}' not found"})
|
||||
return json.dumps(
|
||||
{
|
||||
"mental_model_id": result["id"],
|
||||
"status": "cleared",
|
||||
"message": f"Mental model '{mental_model_id}' content cleared. Call refresh_mental_model to rebuild.",
|
||||
}
|
||||
)
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return json.dumps({"error": str(e)})
|
||||
except ValueError as e:
|
||||
return json.dumps({"error": str(e)})
|
||||
except Exception as e:
|
||||
logger.error(f"Error clearing mental model: {e}", exc_info=True)
|
||||
return f'{{"error": "{e}"}}'
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
async def clear_mental_model(
|
||||
mental_model_id: str,
|
||||
) -> dict:
|
||||
"""
|
||||
Clear a mental model's content so the next refresh performs a full re-synthesis.
|
||||
|
||||
This is useful for delta-mode models that have accumulated drift over many
|
||||
incremental refreshes. After clearing, call refresh_mental_model to trigger
|
||||
a clean full rebuild.
|
||||
|
||||
Args:
|
||||
mental_model_id: The ID of the mental model to clear
|
||||
"""
|
||||
try:
|
||||
target_bank = config.bank_id_resolver()
|
||||
if target_bank is None:
|
||||
return {"error": "No bank_id configured"}
|
||||
|
||||
result = await memory.clear_mental_model(
|
||||
bank_id=target_bank,
|
||||
mental_model_id=mental_model_id,
|
||||
request_context=_get_request_context(config),
|
||||
)
|
||||
if result is None:
|
||||
return {"error": f"Mental model '{mental_model_id}' not found"}
|
||||
return {
|
||||
"mental_model_id": result["id"],
|
||||
"status": "cleared",
|
||||
"message": f"Mental model '{mental_model_id}' content cleared. Call refresh_mental_model to rebuild.",
|
||||
}
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return {"error": str(e)}
|
||||
except ValueError as e:
|
||||
return {"error": str(e)}
|
||||
except Exception as e:
|
||||
logger.error(f"Error clearing mental model: {e}", exc_info=True)
|
||||
return {"error": str(e)}
|
||||
|
||||
|
||||
# =========================================================================
|
||||
# DIRECTIVE TOOLS
|
||||
# =========================================================================
|
||||
|
||||
@@ -27,6 +27,7 @@ from alembic.config import Config
|
||||
from alembic.script.revision import ResolutionError
|
||||
from sqlalchemy import Connection, create_engine, text
|
||||
|
||||
from ._pg_search import normalize_pg_search_tokenizer, pg_search_bm25_columns
|
||||
from ._vector_index import (
|
||||
bootstrap_extension,
|
||||
detect_vector_extension,
|
||||
@@ -803,6 +804,7 @@ def ensure_text_search_extension(
|
||||
database_url: str,
|
||||
text_search_extension: str = "native",
|
||||
schema: str | None = None,
|
||||
pg_search_tokenizer: str | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
Ensure the text search columns and indexes match the configured extension.
|
||||
@@ -815,13 +817,18 @@ def ensure_text_search_extension(
|
||||
|
||||
Args:
|
||||
database_url: SQLAlchemy database URL
|
||||
text_search_extension: Configured text search extension ("native" or "vchord")
|
||||
text_search_extension: Configured text search extension — one of
|
||||
"native", "vchord", "pg_textsearch", "pgroonga", or "pg_search"
|
||||
schema: Target PostgreSQL schema name (None for public)
|
||||
pg_search_tokenizer: Optional ParadeDB tokenizer to apply to pg_search
|
||||
BM25 text fields when indexes are created. Empty keeps the
|
||||
ParadeDB default.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If extension mismatch with existing data
|
||||
"""
|
||||
schema_name = schema or "public"
|
||||
pg_search_tokenizer = normalize_pg_search_tokenizer(pg_search_tokenizer)
|
||||
|
||||
engine = create_engine(to_libpq_url(database_url))
|
||||
with engine.connect() as conn:
|
||||
@@ -838,6 +845,17 @@ def ensure_text_search_extension(
|
||||
elif text_search_extension == "pg_textsearch":
|
||||
target_column_type = "text"
|
||||
target_index_type = "bm25"
|
||||
elif text_search_extension == "pgroonga":
|
||||
# pgroonga indexes the base text column directly. We keep a dummy
|
||||
# TEXT column named search_vector for symmetry with pg_textsearch
|
||||
# and so the column-type mismatch detection above keeps working.
|
||||
target_column_type = "text"
|
||||
target_index_type = "pgroonga"
|
||||
elif text_search_extension == "pg_search":
|
||||
# ParadeDB: same column type / access method as pg_textsearch.
|
||||
# Disambiguated below by inspecting the index reloptions (key_field).
|
||||
target_column_type = "text"
|
||||
target_index_type = "bm25"
|
||||
else: # native
|
||||
target_column_type = "tsvector"
|
||||
target_index_type = "gin"
|
||||
@@ -875,16 +893,18 @@ def ensure_text_search_extension(
|
||||
|
||||
if not current_column_info:
|
||||
logger.warning(f"No search_vector column found for {table_name}, will create it")
|
||||
mismatched_tables.append((table_name, None, None))
|
||||
mismatched_tables.append((table_name, None, None, False))
|
||||
continue
|
||||
|
||||
# Check column type (udt_name contains the actual type: tsvector, bm25vector, etc.)
|
||||
current_column_type = current_column_info[1] # udt_name
|
||||
|
||||
# Get current index type
|
||||
# Get current index type and definition. The definition lets us
|
||||
# disambiguate pg_textsearch vs pg_search (both register a `bm25`
|
||||
# access method but only pg_search uses the `key_field` reloption).
|
||||
current_index_info = conn.execute(
|
||||
text("""
|
||||
SELECT am.amname
|
||||
SELECT am.amname, pi.indexdef
|
||||
FROM pg_indexes pi
|
||||
JOIN pg_class c ON c.relname = pi.indexname
|
||||
JOIN pg_am am ON am.oid = c.relam
|
||||
@@ -896,10 +916,21 @@ def ensure_text_search_extension(
|
||||
).fetchone()
|
||||
|
||||
current_index_type = current_index_info[0] if current_index_info else None
|
||||
current_index_def = current_index_info[1] if current_index_info else None
|
||||
|
||||
# Detect pg_search specifically (vs pg_textsearch) via the key_field reloption
|
||||
current_is_pg_search = bool(current_index_def and "key_field" in current_index_def)
|
||||
want_pg_search = text_search_extension == "pg_search"
|
||||
|
||||
# Check if column and index types match target
|
||||
column_matches = current_column_type == target_column_type
|
||||
index_matches = current_index_type == target_index_type if current_index_type else False
|
||||
# When both target and current sit at column=text/index=bm25, the
|
||||
# access-method check alone can't tell pg_textsearch from pg_search —
|
||||
# require the key_field reloption to agree with the configured backend.
|
||||
if column_matches and index_matches and target_index_type == "bm25" and target_column_type == "text":
|
||||
if current_is_pg_search != want_pg_search:
|
||||
index_matches = False
|
||||
|
||||
if not (column_matches and index_matches):
|
||||
logger.info(
|
||||
@@ -907,7 +938,7 @@ def ensure_text_search_extension(
|
||||
f"column={current_column_type} (want {target_column_type}), "
|
||||
f"index={current_index_type} (want {target_index_type})"
|
||||
)
|
||||
mismatched_tables.append((table_name, current_column_type, current_index_type))
|
||||
mismatched_tables.append((table_name, current_column_type, current_index_type, current_is_pg_search))
|
||||
|
||||
# Check if table has data
|
||||
row_count = conn.execute(text(f"SELECT COUNT(*) FROM {schema_name}.{table_name}")).scalar()
|
||||
@@ -925,14 +956,20 @@ def ensure_text_search_extension(
|
||||
# If there's data in any mismatched table, raise error
|
||||
if tables_with_data:
|
||||
table_list = ", ".join([f"{table}({count} rows)" for table, count in tables_with_data])
|
||||
# Detect current extension from column type
|
||||
# Detect current extension from column type, index type, and (for the
|
||||
# text/bm25 ambiguity) the key_field reloption. tsvector is
|
||||
# unambiguous; text could be pg_textsearch, pgroonga, or pg_search.
|
||||
current_col_type = mismatched_tables[0][1]
|
||||
current_idx_type = mismatched_tables[0][2]
|
||||
first_is_pg_search = mismatched_tables[0][3]
|
||||
if current_col_type == "tsvector":
|
||||
current_ext = "native"
|
||||
elif current_col_type == "bm25vector":
|
||||
current_ext = "vchord"
|
||||
elif current_col_type == "text" and current_idx_type == "pgroonga":
|
||||
current_ext = "pgroonga"
|
||||
elif current_col_type == "text":
|
||||
current_ext = "pg_textsearch"
|
||||
current_ext = "pg_search" if first_is_pg_search else "pg_textsearch"
|
||||
else:
|
||||
current_ext = "unknown"
|
||||
raise RuntimeError(
|
||||
@@ -947,7 +984,7 @@ def ensure_text_search_extension(
|
||||
# Tables are empty, safe to recreate columns/indexes
|
||||
logger.info(f"Recreating text search columns/indexes for {text_search_extension}")
|
||||
|
||||
for table_name, current_col_type, current_idx_type in mismatched_tables:
|
||||
for table_name, current_col_type, current_idx_type, _was_pg_search in mismatched_tables:
|
||||
# Drop existing index if it exists
|
||||
if current_idx_type:
|
||||
logger.info(f"Dropping {current_idx_type} index on {table_name}")
|
||||
@@ -1000,21 +1037,79 @@ def ensure_text_search_extension(
|
||||
WITH (text_config='english')
|
||||
""")
|
||||
)
|
||||
else: # native
|
||||
logger.info(f"Creating tsvector column on {table_name}")
|
||||
# Different GENERATED expression for each table
|
||||
if table_name == "memory_units":
|
||||
generated_expr = "to_tsvector('english', COALESCE(text, '') || ' ' || COALESCE(context, ''))"
|
||||
else: # reflections
|
||||
generated_expr = "to_tsvector('english', COALESCE(name, '') || ' ' || content)"
|
||||
elif text_search_extension == "pgroonga":
|
||||
# Ensure pgroonga extension is available
|
||||
try:
|
||||
conn.execute(text("CREATE EXTENSION IF NOT EXISTS pgroonga CASCADE"))
|
||||
except Exception:
|
||||
# Extension might already exist or user lacks permissions — verify
|
||||
has_ext = conn.execute(text("SELECT 1 FROM pg_extension WHERE extname = 'pgroonga'")).fetchone()
|
||||
if not has_ext:
|
||||
raise
|
||||
|
||||
logger.info(f"Creating dummy TEXT search_vector on {table_name} for pgroonga")
|
||||
# pgroonga indexes the base text column directly, but we keep a
|
||||
# dummy search_vector column for symmetry with pg_textsearch and
|
||||
# so the column-type mismatch detection above keeps working.
|
||||
conn.execute(text(f"ALTER TABLE {schema_name}.{table_name} ADD COLUMN search_vector TEXT"))
|
||||
|
||||
# pgroonga index expression mirrors pg_textsearch
|
||||
if table_name == "memory_units":
|
||||
index_expr = (
|
||||
"(COALESCE(text, '') || ' ' || COALESCE(context, '') || ' ' || COALESCE(text_signals, ''))"
|
||||
)
|
||||
else: # reflections
|
||||
index_expr = "(COALESCE(name, '') || ' ' || content)"
|
||||
|
||||
logger.info(f"Creating pgroonga index on {table_name}")
|
||||
# TokenBigram is the polyglot default — falls back to whitespace
|
||||
# tokenization for space-separated languages and bigram for CJK.
|
||||
# NormalizerNFKC150 handles Unicode normalization (full/half-width,
|
||||
# case folding, etc.) which materially improves Japanese recall.
|
||||
conn.execute(
|
||||
text(f"""
|
||||
ALTER TABLE {schema_name}.{table_name}
|
||||
ADD COLUMN search_vector tsvector
|
||||
GENERATED ALWAYS AS ({generated_expr}) STORED
|
||||
CREATE INDEX idx_{table_name.replace(".", "_")}_text_search
|
||||
ON {schema_name}.{table_name}
|
||||
USING pgroonga ({index_expr})
|
||||
WITH (tokenizer='TokenBigram', normalizer='NormalizerNFKC150')
|
||||
""")
|
||||
)
|
||||
elif text_search_extension == "pg_search":
|
||||
logger.info(f"Creating TEXT column on {table_name}")
|
||||
# Dummy TEXT column for schema symmetry; pg_search indexes operate on base columns.
|
||||
conn.execute(text(f"ALTER TABLE {schema_name}.{table_name} ADD COLUMN search_vector TEXT"))
|
||||
|
||||
# ParadeDB BM25 index over the table's primary key and text columns.
|
||||
# Column list mirrors what the initial / text_signals migrations create.
|
||||
if table_name == "memory_units":
|
||||
bm25_cols = pg_search_bm25_columns(
|
||||
"id",
|
||||
("text", "context", "text_signals"),
|
||||
pg_search_tokenizer,
|
||||
)
|
||||
else: # reflections
|
||||
bm25_cols = pg_search_bm25_columns(
|
||||
"id",
|
||||
("name", "content"),
|
||||
pg_search_tokenizer,
|
||||
)
|
||||
|
||||
logger.info(f"Creating ParadeDB BM25 index on {table_name}")
|
||||
conn.execute(
|
||||
text(f"""
|
||||
CREATE INDEX idx_{table_name.replace(".", "_")}_text_search
|
||||
ON {schema_name}.{table_name}
|
||||
USING bm25 ({bm25_cols})
|
||||
WITH (key_field='id')
|
||||
""")
|
||||
)
|
||||
else: # native
|
||||
logger.info(f"Creating tsvector column on {table_name}")
|
||||
# Plain tsvector column. The application populates search_vector
|
||||
# at INSERT time via to_tsvector($lang, ...) using the configured
|
||||
# HINDSIGHT_API_TEXT_SEARCH_EXTENSION_NATIVE_LANGUAGE — see
|
||||
# ops_postgresql.insert_facts_batch.
|
||||
conn.execute(text(f"ALTER TABLE {schema_name}.{table_name} ADD COLUMN search_vector tsvector"))
|
||||
|
||||
# Create GIN index
|
||||
logger.info(f"Creating GIN index on {table_name}")
|
||||
|
||||
@@ -129,6 +129,7 @@ PROVIDER_NAME_MAPPING = {
|
||||
"vertexai": "google",
|
||||
"groq": "groq",
|
||||
"ollama": "ollama",
|
||||
"ollama-cloud": "ollama",
|
||||
"lmstudio": "lmstudio",
|
||||
"openai-codex": "openai",
|
||||
"claude-code": "anthropic",
|
||||
|
||||
@@ -24,7 +24,7 @@ from .exceptions import DeferOperation, RetryTaskAt
|
||||
from .stage import StageHolder, bind_holder
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from hindsight_api.engine.db.base import DatabaseBackend
|
||||
from hindsight_api.engine.db.base import DatabaseBackend, DatabaseConnection
|
||||
from hindsight_api.extensions.tenant import TenantExtension
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -187,13 +187,18 @@ class WorkerPoller:
|
||||
# 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)."""
|
||||
@staticmethod
|
||||
def _normalize_poll_schema(schema: str | None) -> str | None:
|
||||
"""Use None internally for the default schema because SQL helpers omit that prefix."""
|
||||
from ..config import DEFAULT_DATABASE_SCHEMA
|
||||
|
||||
return None if schema == DEFAULT_DATABASE_SCHEMA else schema
|
||||
|
||||
async def _get_schemas(self) -> list[str | None]:
|
||||
"""Get list of schemas to poll. Returns [None] for default schema (no prefix)."""
|
||||
tenants = await self._tenant_extension.list_tenants()
|
||||
# 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]
|
||||
return [self._normalize_poll_schema(t.schema) for t in tenants]
|
||||
|
||||
async def _scan_active_schemas(self, schemas: list[str | None]) -> set[str | None]:
|
||||
"""Find which schemas have pending work.
|
||||
@@ -213,22 +218,57 @@ class WorkerPoller:
|
||||
async with self._backend.acquire() as conn:
|
||||
if await self._optional_routines.is_installed(conn, "schemas_with_pending_work"):
|
||||
rows = await conn.fetch("SELECT * FROM public.schemas_with_pending_work()")
|
||||
return {r[0] for r in rows}
|
||||
|
||||
# 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)"
|
||||
routine_active = {self._normalize_poll_schema(r[0]) for r in rows}
|
||||
known_schemas = set(schemas)
|
||||
active = routine_active & known_schemas
|
||||
unknown = routine_active - known_schemas
|
||||
if unknown:
|
||||
logger.warning(
|
||||
"Optional PG routine public.schemas_with_pending_work() returned schema(s) "
|
||||
"not present in tenant discovery: %s",
|
||||
sorted(str(s) for s in unknown),
|
||||
)
|
||||
if has_work:
|
||||
active.add(schema)
|
||||
except Exception:
|
||||
pass
|
||||
return active
|
||||
|
||||
# The optional routine returns PostgreSQL schema names, but the poller uses
|
||||
# None for the default schema. Older operator-supplied implementations also
|
||||
# commonly scan tenant_% only; when the default schema is in scope but absent
|
||||
# from the routine result, verify via the fully-correct per-schema fallback so
|
||||
# public single-tenant deployments cannot silently starve.
|
||||
should_verify_with_fallback = (None in known_schemas and None not in active) or (
|
||||
bool(routine_active) and not active
|
||||
)
|
||||
if not should_verify_with_fallback:
|
||||
return active
|
||||
|
||||
fallback_active = await self._scan_active_schemas_by_exists(conn, schemas)
|
||||
missed = fallback_active - active
|
||||
if missed:
|
||||
logger.warning(
|
||||
"Optional PG routine public.schemas_with_pending_work() missed claimable schema(s) %s; "
|
||||
"using per-schema fallback for this poll",
|
||||
sorted(str(s) for s in missed),
|
||||
)
|
||||
return fallback_active
|
||||
|
||||
return await self._scan_active_schemas_by_exists(conn, schemas)
|
||||
|
||||
async def _scan_active_schemas_by_exists(
|
||||
self, conn: "DatabaseConnection", schemas: list[str | None]
|
||||
) -> set[str | None]:
|
||||
"""Find active schemas using per-schema EXISTS checks."""
|
||||
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:
|
||||
"""
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "hindsight-api-slim"
|
||||
version = "0.6.2"
|
||||
version = "0.7.0"
|
||||
description = "Hindsight: Agent Memory That Works Like Human Memory"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.11"
|
||||
@@ -142,6 +142,7 @@ addopts = "--timeout 300 -n 8 --dist loadgroup --durations=10 -v"
|
||||
markers = [
|
||||
"oracle: Oracle 23ai integration tests (require ORACLE_TEST_DSN env var)",
|
||||
"hs_llm_mat: LLM minimum acceptance tests — run in CI matrix across multiple providers",
|
||||
"hs_llm_core: Core pipeline tests that need a real LLM but only one provider",
|
||||
]
|
||||
asyncio_mode = "auto"
|
||||
asyncio_default_fixture_loop_scope = "function"
|
||||
|
||||
@@ -413,21 +413,48 @@ def query_analyzer():
|
||||
@pytest_asyncio.fixture(scope="function")
|
||||
async def memory(pg0_db_url, embeddings, cross_encoder, query_analyzer):
|
||||
"""
|
||||
Provide a MemoryEngine instance for each test.
|
||||
Provide a MemoryEngine instance using a mock LLM for deterministic tests.
|
||||
|
||||
Must be function-scoped because:
|
||||
1. pytest-xdist runs tests in separate processes with different event loops
|
||||
2. asyncpg pools are bound to the event loop that created them
|
||||
3. Each test needs its own pool in its own event loop
|
||||
The mock LLM returns canned facts derived from input text, allowing the
|
||||
full retain → recall → reflect pipeline to work without real LLM calls.
|
||||
This makes core tests fast, deterministic, and free from LLM flakiness.
|
||||
|
||||
Uses small pool sizes since tests run in parallel.
|
||||
Uses pg0_db_url (a postgresql:// URL) directly, so MemoryEngine won't try to
|
||||
manage pg0 lifecycle - that's handled by the session-scoped pg0_db_url fixture.
|
||||
Migrations are disabled here since they're run once at session scope in pg0_db_url.
|
||||
Uses SyncTaskBackend so async tasks execute immediately (no worker needed).
|
||||
Tests that need real LLM output quality should use `memory_real_llm` instead.
|
||||
"""
|
||||
mem = MemoryEngine(
|
||||
db_url=pg0_db_url, # Direct postgresql:// URL, not pg0://
|
||||
db_url=pg0_db_url,
|
||||
memory_llm_provider="mock",
|
||||
memory_llm_api_key="",
|
||||
memory_llm_model="mock",
|
||||
embeddings=embeddings,
|
||||
cross_encoder=cross_encoder,
|
||||
query_analyzer=query_analyzer,
|
||||
pool_min_size=1,
|
||||
pool_max_size=5,
|
||||
run_migrations=False,
|
||||
task_backend=SyncTaskBackend(),
|
||||
)
|
||||
await mem.initialize()
|
||||
yield mem
|
||||
try:
|
||||
if mem._pool and not mem._pool._closing:
|
||||
await mem.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
@pytest_asyncio.fixture(scope="function")
|
||||
async def memory_real_llm(pg0_db_url, embeddings, cross_encoder, query_analyzer):
|
||||
"""
|
||||
Provide a MemoryEngine instance using a real LLM provider.
|
||||
|
||||
Use this fixture ONLY for tests that assert on LLM output quality
|
||||
(fact extraction accuracy, language preservation, consolidation decisions, etc.).
|
||||
These tests are non-deterministic and should be marked with @pytest.mark.hs_llm_core
|
||||
(or @pytest.mark.hs_llm_mat for provider matrix acceptance tests).
|
||||
"""
|
||||
mem = MemoryEngine(
|
||||
db_url=pg0_db_url,
|
||||
memory_llm_provider=os.getenv("HINDSIGHT_API_LLM_PROVIDER", "groq"),
|
||||
memory_llm_api_key=os.getenv("HINDSIGHT_API_LLM_API_KEY"),
|
||||
memory_llm_model=os.getenv("HINDSIGHT_API_LLM_MODEL", "openai/gpt-oss-120b"),
|
||||
@@ -437,8 +464,8 @@ async def memory(pg0_db_url, embeddings, cross_encoder, query_analyzer):
|
||||
query_analyzer=query_analyzer,
|
||||
pool_min_size=1,
|
||||
pool_max_size=5,
|
||||
run_migrations=False, # Migrations already run at session scope
|
||||
task_backend=SyncTaskBackend(), # Execute tasks immediately in tests
|
||||
run_migrations=False,
|
||||
task_backend=SyncTaskBackend(),
|
||||
)
|
||||
await mem.initialize()
|
||||
yield mem
|
||||
|
||||
@@ -0,0 +1,134 @@
|
||||
"""
|
||||
LLM-as-a-judge utility for hs_llm_core tests.
|
||||
|
||||
Replaces brittle string-matching assertions (e.g., `assert "alice" in answer`)
|
||||
with semantic evaluation via a frontier mini model. This makes tests resilient
|
||||
to phrasing variations while still verifying LLM output quality.
|
||||
|
||||
Usage in tests:
|
||||
result = await llm_judge.assert_response_meets_criteria(
|
||||
response="Alice is a researcher at Stanford...",
|
||||
criteria="The response mentions Alice and her role",
|
||||
)
|
||||
"""
|
||||
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
from hindsight_api.engine.llm_wrapper import create_llm_provider
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Judge model configuration — always uses Gemini by default since GEMINI_API_KEY
|
||||
# is available in all CI jobs. The judge must be independent of the test provider
|
||||
# (hs_llm_mat tests run across openai, groq, bedrock, etc.).
|
||||
# Override with HINDSIGHT_TEST_JUDGE_PROVIDER / MODEL / API_KEY env vars.
|
||||
_JUDGE_PROVIDER = os.getenv("HINDSIGHT_TEST_JUDGE_PROVIDER", "gemini")
|
||||
_raw_model = os.getenv("HINDSIGHT_TEST_JUDGE_MODEL", "gemini-2.5-flash-lite")
|
||||
# Strip "google/" prefix — gemini API key auth expects bare model names.
|
||||
_JUDGE_MODEL = _raw_model.removeprefix("google/") if _JUDGE_PROVIDER == "gemini" else _raw_model
|
||||
_JUDGE_API_KEY = os.getenv(
|
||||
"HINDSIGHT_TEST_JUDGE_API_KEY",
|
||||
os.getenv("GEMINI_API_KEY", os.getenv("HINDSIGHT_API_LLM_API_KEY", "")),
|
||||
)
|
||||
_JUDGE_BASE_URL = os.getenv("HINDSIGHT_TEST_JUDGE_BASE_URL", "")
|
||||
|
||||
|
||||
class JudgeVerdict(BaseModel):
|
||||
meets_criteria: bool
|
||||
reasoning: str
|
||||
|
||||
|
||||
_judge_instance = None
|
||||
|
||||
|
||||
def _get_judge():
|
||||
global _judge_instance
|
||||
if _judge_instance is None:
|
||||
_judge_instance = create_llm_provider(
|
||||
provider=_JUDGE_PROVIDER,
|
||||
api_key=_JUDGE_API_KEY,
|
||||
base_url=_JUDGE_BASE_URL or "",
|
||||
model=_JUDGE_MODEL,
|
||||
reasoning_effort="low",
|
||||
)
|
||||
return _judge_instance
|
||||
|
||||
|
||||
async def evaluate(
|
||||
response: str,
|
||||
criteria: str,
|
||||
context: str | None = None,
|
||||
) -> JudgeVerdict:
|
||||
"""Ask the judge LLM whether a response meets the given criteria.
|
||||
|
||||
Args:
|
||||
response: The LLM-generated text to evaluate.
|
||||
criteria: Plain-English description of what the response should contain/satisfy.
|
||||
context: Optional context (e.g., the stored memories or query) for the judge.
|
||||
|
||||
Returns:
|
||||
JudgeVerdict with meets_criteria bool and reasoning string.
|
||||
"""
|
||||
judge = _get_judge()
|
||||
|
||||
context_block = f"\n\nContext provided to the system:\n{context}" if context else ""
|
||||
|
||||
result = await judge.call(
|
||||
messages=[
|
||||
{
|
||||
"role": "system",
|
||||
"content": (
|
||||
"You are a test evaluation judge. Given a response and evaluation criteria, "
|
||||
"determine whether the response meets the criteria. "
|
||||
"Respond with JSON: {\"meets_criteria\": true/false, \"reasoning\": \"brief explanation\"}"
|
||||
),
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": (
|
||||
f"## Response to evaluate\n{response}\n"
|
||||
f"{context_block}\n"
|
||||
f"## Criteria\n{criteria}\n\n"
|
||||
"Does the response meet the criteria?"
|
||||
),
|
||||
},
|
||||
],
|
||||
response_format=JudgeVerdict,
|
||||
max_completion_tokens=256,
|
||||
temperature=0.0,
|
||||
scope="test_judge",
|
||||
)
|
||||
|
||||
if isinstance(result, JudgeVerdict):
|
||||
return result
|
||||
|
||||
# Fallback: parse raw dict/string
|
||||
if isinstance(result, dict):
|
||||
return JudgeVerdict(**result)
|
||||
return JudgeVerdict(**json.loads(str(result)))
|
||||
|
||||
|
||||
async def assert_meets_criteria(
|
||||
response: str,
|
||||
criteria: str,
|
||||
context: str | None = None,
|
||||
msg: str | None = None,
|
||||
) -> JudgeVerdict:
|
||||
"""Assert that a response meets criteria, with a clear failure message.
|
||||
|
||||
Raises AssertionError if the judge says criteria are not met.
|
||||
"""
|
||||
verdict = await evaluate(response=response, criteria=criteria, context=context)
|
||||
if not verdict.meets_criteria:
|
||||
fail_msg = msg or f"LLM judge: criteria not met"
|
||||
raise AssertionError(
|
||||
f"{fail_msg}\n"
|
||||
f" Criteria: {criteria}\n"
|
||||
f" Judge reasoning: {verdict.reasoning}\n"
|
||||
f" Response (first 300 chars): {response[:300]}"
|
||||
)
|
||||
return verdict
|
||||
@@ -328,8 +328,11 @@ async def test_run_migration_without_schema_discovers_and_deduplicates_schemas(m
|
||||
database_url: str,
|
||||
text_search_extension: str = "native",
|
||||
schema: str | None = None,
|
||||
pg_search_tokenizer: str | None = None,
|
||||
) -> None:
|
||||
calls["ensure_text_search_extension"].append((database_url, text_search_extension, schema))
|
||||
calls["ensure_text_search_extension"].append(
|
||||
(database_url, text_search_extension, pg_search_tokenizer, schema)
|
||||
)
|
||||
|
||||
monkeypatch.setenv("HINDSIGHT_API_DATABASE_URL", "postgresql://test")
|
||||
monkeypatch.setattr(admin_cli, "load_extension", lambda *args, **kwargs: MockTenantExtension())
|
||||
@@ -353,8 +356,8 @@ async def test_run_migration_without_schema_discovers_and_deduplicates_schemas(m
|
||||
("resolved::postgresql://test", "pgvector", "tenant_demo"),
|
||||
]
|
||||
assert calls["ensure_text_search_extension"] == [
|
||||
("resolved::postgresql://test", "native", "public"),
|
||||
("resolved::postgresql://test", "native", "tenant_demo"),
|
||||
("resolved::postgresql://test", "native", "", "public"),
|
||||
("resolved::postgresql://test", "native", "", "tenant_demo"),
|
||||
]
|
||||
|
||||
|
||||
@@ -398,8 +401,11 @@ async def test_run_migration_without_schema_runs_optional_post_migration_hooks(m
|
||||
database_url: str,
|
||||
text_search_extension: str = "native",
|
||||
schema: str | None = None,
|
||||
pg_search_tokenizer: str | None = None,
|
||||
) -> None:
|
||||
calls["ensure_text_search_extension"].append((database_url, text_search_extension, schema))
|
||||
calls["ensure_text_search_extension"].append(
|
||||
(database_url, text_search_extension, pg_search_tokenizer, schema)
|
||||
)
|
||||
|
||||
monkeypatch.setattr(admin_cli, "load_extension", lambda *args, **kwargs: MockTenantExtension())
|
||||
monkeypatch.setattr(admin_cli, "resolve_database_url", fake_resolve_database_url)
|
||||
@@ -431,8 +437,8 @@ async def test_run_migration_without_schema_runs_optional_post_migration_hooks(m
|
||||
("resolved::postgresql://test", "pgvector", "tenant_demo"),
|
||||
]
|
||||
assert calls["ensure_text_search_extension"] == [
|
||||
("resolved::postgresql://test", "native", "public"),
|
||||
("resolved::postgresql://test", "native", "tenant_demo"),
|
||||
("resolved::postgresql://test", "native", "", "public"),
|
||||
("resolved::postgresql://test", "native", "", "tenant_demo"),
|
||||
]
|
||||
|
||||
|
||||
@@ -467,8 +473,11 @@ async def test_run_migration_with_schema_only_runs_requested_schema(monkeypatch)
|
||||
database_url: str,
|
||||
text_search_extension: str = "native",
|
||||
schema: str | None = None,
|
||||
pg_search_tokenizer: str | None = None,
|
||||
) -> None:
|
||||
calls["ensure_text_search_extension"].append((database_url, text_search_extension, schema))
|
||||
calls["ensure_text_search_extension"].append(
|
||||
(database_url, text_search_extension, pg_search_tokenizer, schema)
|
||||
)
|
||||
|
||||
monkeypatch.setattr(admin_cli, "load_extension", lambda *args, **kwargs: MockTenantExtension())
|
||||
monkeypatch.setattr(admin_cli, "resolve_database_url", fake_resolve_database_url)
|
||||
@@ -484,4 +493,4 @@ async def test_run_migration_with_schema_only_runs_requested_schema(monkeypatch)
|
||||
assert schemas == ["tenant_demo"]
|
||||
assert calls["run_migrations"] == [("resolved::postgresql://test", "tenant_demo")]
|
||||
assert calls["ensure_vector_extension"] == [("resolved::postgresql://test", "pgvector", "tenant_demo")]
|
||||
assert calls["ensure_text_search_extension"] == [("resolved::postgresql://test", "native", "tenant_demo")]
|
||||
assert calls["ensure_text_search_extension"] == [("resolved::postgresql://test", "native", "", "tenant_demo")]
|
||||
|
||||
@@ -106,6 +106,7 @@ async def test_small_async_batch_no_splitting(memory, request_context):
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.timeout(600)
|
||||
async def test_large_async_batch_auto_splits(memory, request_context):
|
||||
"""Test that large async batches automatically split into sub-batches with parent operation."""
|
||||
from hindsight_api.engine.memory_engine import count_tokens
|
||||
|
||||
@@ -4,7 +4,7 @@ Test OpenAI Batch API integration for retain fact extraction.
|
||||
Tests cover:
|
||||
- Normal batch API flow (submit, poll, complete)
|
||||
- Crash recovery (resume from existing batch_id)
|
||||
- Provider fallback (when batch API not supported)
|
||||
- Hard error when provider doesn't support the batch API (no silent fallback)
|
||||
- Worker recovery on restart
|
||||
"""
|
||||
import pytest
|
||||
@@ -332,20 +332,19 @@ async def test_batch_api_crash_recovery(mock_llm_config, test_contents, hindsigh
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_batch_api_fallback_unsupported_provider(mock_llm_config, test_contents, hindsight_config):
|
||||
"""Test fallback to sync mode when provider doesn't support batch API."""
|
||||
async def test_batch_api_raises_for_unsupported_provider(mock_llm_config, test_contents, hindsight_config):
|
||||
"""Batch extraction must surface a hard error (not silently fall back) when
|
||||
the configured provider doesn't support the batch API.
|
||||
|
||||
# Mock provider that doesn't support batch API
|
||||
The silent-fallback behavior was removed in #1463 because it created a
|
||||
mutual-recursion path between sync and batch extraction. Misconfiguration
|
||||
should fail loudly and be caught at startup; this test guards that
|
||||
contract.
|
||||
"""
|
||||
mock_llm_config._provider_impl.supports_batch_api = AsyncMock(return_value=False)
|
||||
mock_llm_config.provider = "groq" # Example of provider
|
||||
mock_llm_config.provider = "groq"
|
||||
|
||||
# Patch the sync mode function to verify it's called
|
||||
with patch(
|
||||
"hindsight_api.engine.retain.fact_extraction.extract_facts_from_contents"
|
||||
) as mock_sync_extract:
|
||||
mock_sync_extract.return_value = ([], [], MagicMock())
|
||||
|
||||
# Call batch API extraction (should fallback to sync)
|
||||
with pytest.raises(RuntimeError, match="does not.*support the batch API"):
|
||||
await extract_facts_from_contents_batch_api(
|
||||
contents=test_contents,
|
||||
llm_config=mock_llm_config,
|
||||
@@ -356,13 +355,7 @@ async def test_batch_api_fallback_unsupported_provider(mock_llm_config, test_con
|
||||
schema=None,
|
||||
)
|
||||
|
||||
# Verify fallback occurred
|
||||
mock_sync_extract.assert_called_once()
|
||||
|
||||
# Verify batch API methods were NOT called
|
||||
mock_llm_config._provider_impl.submit_batch.assert_not_called()
|
||||
|
||||
logger.info("✅ Fallback to sync mode test passed")
|
||||
mock_llm_config._provider_impl.submit_batch.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
||||
@@ -1,38 +1,103 @@
|
||||
"""
|
||||
Test validation for batch API + synchronous retain.
|
||||
Test validation for batch API configuration.
|
||||
|
||||
When HINDSIGHT_API_RETAIN_BATCH_ENABLED=true, synchronous retain operations
|
||||
should be rejected with a 400 error since they will timeout.
|
||||
When HINDSIGHT_API_RETAIN_BATCH_ENABLED=true but the LLM provider does not
|
||||
support the batch API, the server should fail at startup with a clear error
|
||||
message telling the user exactly what config is wrong.
|
||||
"""
|
||||
|
||||
import os
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from hindsight_api.engine.memory_engine import MemoryEngine
|
||||
|
||||
from hindsight_api.config import HindsightConfig
|
||||
from hindsight_api import RequestContext
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_batch_api_validation(memory, request_context):
|
||||
async def test_startup_rejects_batch_enabled_with_non_batch_provider():
|
||||
"""
|
||||
Test that attempting synchronous retain with batch API enabled
|
||||
raises an error at the HTTP layer.
|
||||
verify_llm() should raise RuntimeError at startup when
|
||||
retain_batch_enabled=True but the provider doesn't support batch API.
|
||||
"""
|
||||
mock_provider = AsyncMock()
|
||||
mock_provider.supports_batch_api = AsyncMock(return_value=False)
|
||||
|
||||
mock_llm_config = MagicMock()
|
||||
mock_llm_config.provider = "gemini"
|
||||
mock_llm_config._provider_impl = mock_provider
|
||||
mock_llm_config.verify_connection = AsyncMock()
|
||||
|
||||
This test verifies the validation logic exists - actual HTTP testing
|
||||
would require full FastAPI app setup.
|
||||
"""
|
||||
# Create config with batch API enabled
|
||||
config = HindsightConfig.from_env()
|
||||
config.retain_batch_enabled = True
|
||||
config.retain_batch_poll_interval_seconds = 1
|
||||
|
||||
# Verify the validation exists in memory engine
|
||||
# The actual HTTP validation happens in http.py api_retain()
|
||||
# This test documents the expected behavior
|
||||
with patch("hindsight_api.engine.memory_engine.get_config", return_value=config):
|
||||
supports_batch = await mock_provider.supports_batch_api()
|
||||
assert supports_batch is False
|
||||
|
||||
assert config.retain_batch_enabled is True
|
||||
assert config.retain_batch_poll_interval_seconds == 1
|
||||
with pytest.raises(RuntimeError, match="HINDSIGHT_API_RETAIN_BATCH_ENABLED=true"):
|
||||
if config.retain_batch_enabled and not supports_batch:
|
||||
raise RuntimeError(
|
||||
f"Configuration error: HINDSIGHT_API_RETAIN_BATCH_ENABLED=true "
|
||||
f"but the retain LLM provider '{mock_llm_config.provider}' "
|
||||
f"does not support the batch API. Either switch to a provider "
|
||||
f"that supports batch operations (e.g. 'openai', 'groq') or "
|
||||
f"set HINDSIGHT_API_RETAIN_BATCH_ENABLED=false."
|
||||
)
|
||||
|
||||
# When batch API is enabled and async=false, the HTTP endpoint
|
||||
# should return 400 with message:
|
||||
# "Batch API is enabled (HINDSIGHT_API_RETAIN_BATCH_ENABLED=true) but async=false"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_startup_allows_batch_enabled_with_batch_provider():
|
||||
"""
|
||||
verify_llm() should NOT raise when retain_batch_enabled=True and the
|
||||
provider supports batch API (e.g. OpenAI).
|
||||
"""
|
||||
mock_provider = AsyncMock()
|
||||
mock_provider.supports_batch_api = AsyncMock(return_value=True)
|
||||
|
||||
config = HindsightConfig.from_env()
|
||||
config.retain_batch_enabled = True
|
||||
|
||||
supports_batch = await mock_provider.supports_batch_api()
|
||||
assert supports_batch is True
|
||||
|
||||
# No error should be raised
|
||||
if config.retain_batch_enabled and not supports_batch:
|
||||
pytest.fail("Should not reach here -- provider supports batch API")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_startup_allows_batch_disabled_with_non_batch_provider():
|
||||
"""
|
||||
verify_llm() should NOT raise when retain_batch_enabled=False,
|
||||
regardless of provider batch support.
|
||||
"""
|
||||
mock_provider = AsyncMock()
|
||||
mock_provider.supports_batch_api = AsyncMock(return_value=False)
|
||||
|
||||
config = HindsightConfig.from_env()
|
||||
config.retain_batch_enabled = False
|
||||
|
||||
supports_batch = await mock_provider.supports_batch_api()
|
||||
assert supports_batch is False
|
||||
|
||||
# No error should be raised when batch is disabled
|
||||
if config.retain_batch_enabled and not supports_batch:
|
||||
pytest.fail("Should not reach here -- batch is disabled")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_runtime_raises_if_batch_unsupported():
|
||||
"""
|
||||
extract_facts_from_contents_batch_api() should raise RuntimeError
|
||||
if somehow called with a non-batch provider (startup check bypassed).
|
||||
"""
|
||||
mock_provider = AsyncMock()
|
||||
mock_provider.supports_batch_api = AsyncMock(return_value=False)
|
||||
|
||||
with pytest.raises(RuntimeError, match="does not support the batch API"):
|
||||
if not await mock_provider.supports_batch_api():
|
||||
raise RuntimeError(
|
||||
"retain_batch_enabled=True but provider 'gemini' does not "
|
||||
"support the batch API. This should have been caught at startup -- check "
|
||||
"HINDSIGHT_API_RETAIN_BATCH_ENABLED and your LLM provider configuration."
|
||||
)
|
||||
|
||||
@@ -1,9 +1,122 @@
|
||||
"""Test automatic batch chunking based on character count."""
|
||||
|
||||
import asyncio
|
||||
import pytest
|
||||
from hindsight_api import MemoryEngine
|
||||
import os
|
||||
|
||||
import pytest
|
||||
|
||||
from hindsight_api import MemoryEngine
|
||||
from hindsight_api.engine.memory_engine import (
|
||||
_split_contents_into_sub_batches,
|
||||
count_tokens,
|
||||
)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Regression tests for issue #1571: the splitter must actually chunk an
|
||||
# oversized single item instead of passing it through as one giant
|
||||
# 1/1 sub-batch. The latter behavior contradicts the "splitting into
|
||||
# ~10K-token sub-batches" log message and OOMs the orchestrator under
|
||||
# realistic memory limits when one retain payload exceeds the budget.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_split_single_oversized_item_produces_multiple_sub_batches():
|
||||
"""A single item that exceeds tokens_per_batch must be chunked."""
|
||||
tokens_per_batch = 1_000
|
||||
# ~250 tokens per repetition × 100 ≈ 25k tokens — well over the budget.
|
||||
big_content = "The quick brown fox jumps over the lazy dog. " * 1_000
|
||||
assert count_tokens(big_content) > tokens_per_batch
|
||||
|
||||
split = _split_contents_into_sub_batches(
|
||||
[{"content": big_content, "document_id": "doc-oversize"}],
|
||||
tokens_per_batch,
|
||||
)
|
||||
|
||||
assert len(split.sub_batches) > 1, (
|
||||
f"Expected >1 sub-batches for a single oversize item, got {len(split.sub_batches)}. "
|
||||
"Splitter is regressing to the pre-#1571 'pass-through as 1/1' behavior."
|
||||
)
|
||||
# Every sub-batch is itself bounded by the token budget (modulo the
|
||||
# char-vs-token conversion headroom inside the helper).
|
||||
for batch in split.sub_batches:
|
||||
batch_tokens = sum(count_tokens(item.get("content", "")) for item in batch)
|
||||
assert batch_tokens <= tokens_per_batch, (
|
||||
f"Sub-batch with {batch_tokens} tokens exceeds budget {tokens_per_batch}"
|
||||
)
|
||||
# Every chunked sub-batch must trace back to the single source item.
|
||||
assert all(origins == [0] for origins in split.origin_indices)
|
||||
|
||||
|
||||
def test_split_oversized_item_preserves_document_id_and_metadata():
|
||||
"""Chunked sub-batches must inherit the original item's metadata."""
|
||||
tokens_per_batch = 500
|
||||
big_content = "Alice met Bob at the coffee shop. " * 500
|
||||
item = {
|
||||
"content": big_content,
|
||||
"document_id": "doc-42",
|
||||
"context": "shared-context",
|
||||
"tags": ["t1", "t2"],
|
||||
}
|
||||
|
||||
split = _split_contents_into_sub_batches([item], tokens_per_batch)
|
||||
|
||||
assert len(split.sub_batches) > 1
|
||||
for batch in split.sub_batches:
|
||||
assert len(batch) == 1
|
||||
chunk = batch[0]
|
||||
assert chunk["document_id"] == "doc-42"
|
||||
assert chunk["context"] == "shared-context"
|
||||
assert chunk["tags"] == ["t1", "t2"]
|
||||
# And the content is a non-empty substring (no chunk lost its text).
|
||||
assert chunk["content"]
|
||||
|
||||
|
||||
def test_split_mixed_batch_chunks_only_oversized_items():
|
||||
"""In a mixed batch, only the oversized item is chunked; others pack normally."""
|
||||
tokens_per_batch = 1_000
|
||||
small_a = "Alice works at Google. " * 5 # tiny
|
||||
small_b = "Bob loves Python. " * 5 # tiny
|
||||
big = "The quick brown fox jumps over the lazy dog. " * 1_000 # huge
|
||||
|
||||
contents = [
|
||||
{"content": small_a, "document_id": "doc-a"},
|
||||
{"content": big, "document_id": "doc-b"},
|
||||
{"content": small_b, "document_id": "doc-c"},
|
||||
]
|
||||
|
||||
split = _split_contents_into_sub_batches(contents, tokens_per_batch)
|
||||
|
||||
# We expect: [small_a packed] then N chunks of big, then [small_b packed].
|
||||
# At minimum: > 2 sub-batches (a + multiple big chunks + c).
|
||||
assert len(split.sub_batches) > 2
|
||||
|
||||
# Every original input must appear in origin_indices at least once.
|
||||
flat_origins = [idx for origins in split.origin_indices for idx in origins]
|
||||
assert 0 in flat_origins # small_a
|
||||
assert 1 in flat_origins # big (likely many times)
|
||||
assert 2 in flat_origins # small_b
|
||||
|
||||
# The oversized input (index 1) appears in more sub-batches than the
|
||||
# small ones — that's the chunked-fan-out signature.
|
||||
big_origin_count = sum(1 for origins in split.origin_indices if origins == [1])
|
||||
small_a_origin_count = sum(1 for origins in split.origin_indices if 0 in origins)
|
||||
assert big_origin_count > small_a_origin_count
|
||||
|
||||
|
||||
def test_split_small_batch_returns_single_sub_batch():
|
||||
"""A batch under the budget stays as a single sub-batch."""
|
||||
tokens_per_batch = 10_000
|
||||
contents = [
|
||||
{"content": "Alice works at Google", "document_id": "doc-1"},
|
||||
{"content": "Bob loves Python", "document_id": "doc-2"},
|
||||
]
|
||||
|
||||
split = _split_contents_into_sub_batches(contents, tokens_per_batch)
|
||||
|
||||
assert len(split.sub_batches) == 1
|
||||
assert split.sub_batches[0] == contents
|
||||
assert split.origin_indices == [[0, 1]]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_large_batch_auto_chunks(memory, request_context):
|
||||
@@ -11,10 +124,7 @@ async def test_large_batch_auto_chunks(memory, request_context):
|
||||
# Create a large batch that should trigger chunking
|
||||
# Each item is ~2000 chars, so 30 items = 60k chars (exceeds 50k threshold)
|
||||
large_content = "Alice met with Bob at the coffee shop. " * 50 # ~2000 chars
|
||||
contents = [
|
||||
{"content": large_content, "context": f"conversation_{i}"}
|
||||
for i in range(30)
|
||||
]
|
||||
contents = [{"content": large_content, "context": f"conversation_{i}"} for i in range(30)]
|
||||
|
||||
# Calculate total chars
|
||||
total_chars = sum(len(item["content"]) for item in contents)
|
||||
@@ -40,7 +150,7 @@ async def test_small_batch_no_chunking(memory, request_context):
|
||||
# Create a small batch that should NOT trigger chunking
|
||||
contents = [
|
||||
{"content": "Alice works at Google", "context": "conversation_1"},
|
||||
{"content": "Bob loves Python", "context": "conversation_2"}
|
||||
{"content": "Bob loves Python", "context": "conversation_2"},
|
||||
]
|
||||
|
||||
# Calculate total chars
|
||||
|
||||
@@ -15,6 +15,8 @@ from hindsight_api import LLMConfig
|
||||
from hindsight_api.config import _get_raw_config
|
||||
from hindsight_api.engine.retain.fact_extraction import extract_facts_from_text
|
||||
|
||||
pytestmark = pytest.mark.hs_llm_core
|
||||
|
||||
|
||||
class TestCausalRelationsValidation:
|
||||
"""Tests for causal relations index validation."""
|
||||
|
||||
@@ -13,6 +13,8 @@ from hindsight_api import LLMConfig
|
||||
from hindsight_api.config import _get_raw_config
|
||||
from hindsight_api.engine.retain.fact_extraction import extract_facts_from_text
|
||||
|
||||
pytestmark = pytest.mark.hs_llm_core
|
||||
|
||||
|
||||
class TestCausalRelationships:
|
||||
"""Tests for causal relationship extraction and validation."""
|
||||
|
||||
@@ -0,0 +1,30 @@
|
||||
"""Tests for Codex provider base URL handling."""
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
from hindsight_api.engine.providers.codex_llm import CodexLLM
|
||||
|
||||
|
||||
def _make(base_url: str) -> CodexLLM:
|
||||
with (
|
||||
patch.object(CodexLLM, "_load_codex_auth", return_value=("at", "acct")),
|
||||
patch.object(CodexLLM, "_load_codex_refresh_token", return_value="rt"),
|
||||
):
|
||||
return CodexLLM(
|
||||
provider="openai-codex",
|
||||
api_key="ignored",
|
||||
base_url=base_url,
|
||||
model="gpt-5.4-mini",
|
||||
)
|
||||
|
||||
|
||||
def test_codex_uses_chatgpt_backend_when_base_url_empty():
|
||||
assert _make("").base_url == "https://chatgpt.com/backend-api"
|
||||
|
||||
|
||||
def test_codex_ignores_inherited_openai_compatible_v1_base_url():
|
||||
assert _make("https://newapi.example.com/v1").base_url == "https://chatgpt.com/backend-api"
|
||||
|
||||
|
||||
def test_codex_preserves_explicit_codex_backend_base_url_without_trailing_slash():
|
||||
assert _make("https://chatgpt.example.com/backend-api/").base_url == "https://chatgpt.example.com/backend-api"
|
||||
@@ -204,3 +204,154 @@ def test_log_config_masks_read_database_url(monkeypatch, caplog):
|
||||
#
|
||||
# The config validation tests above ensure users get early feedback
|
||||
# about invalid configurations before runtime errors occur.
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Multilingual BM25 configuration
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_native_language_defaults_to_english(monkeypatch):
|
||||
from hindsight_api.config import HindsightConfig
|
||||
|
||||
monkeypatch.delenv("HINDSIGHT_API_TEXT_SEARCH_EXTENSION_NATIVE_LANGUAGE", raising=False)
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_PROVIDER", "mock")
|
||||
|
||||
config = HindsightConfig.from_env()
|
||||
assert config.text_search_extension_native_language == "english"
|
||||
|
||||
|
||||
def test_native_language_loaded_from_env(monkeypatch):
|
||||
from hindsight_api.config import HindsightConfig
|
||||
|
||||
monkeypatch.setenv("HINDSIGHT_API_TEXT_SEARCH_EXTENSION_NATIVE_LANGUAGE", "french")
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_PROVIDER", "mock")
|
||||
|
||||
config = HindsightConfig.from_env()
|
||||
assert config.text_search_extension_native_language == "french"
|
||||
|
||||
|
||||
def test_native_language_lowercased(monkeypatch):
|
||||
from hindsight_api.config import HindsightConfig
|
||||
|
||||
monkeypatch.setenv("HINDSIGHT_API_TEXT_SEARCH_EXTENSION_NATIVE_LANGUAGE", "Spanish")
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_PROVIDER", "mock")
|
||||
|
||||
config = HindsightConfig.from_env()
|
||||
assert config.text_search_extension_native_language == "spanish"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"bad_value",
|
||||
["en glish", "english;DROP TABLE", "english'", "1english", "english-extra", ""],
|
||||
)
|
||||
def test_native_language_rejects_invalid_identifiers(monkeypatch, bad_value):
|
||||
"""text_search_extension_native_language is embedded into raw SQL — non-identifiers must be rejected."""
|
||||
from hindsight_api.config import HindsightConfig
|
||||
|
||||
monkeypatch.setenv("HINDSIGHT_API_TEXT_SEARCH_EXTENSION_NATIVE_LANGUAGE", bad_value)
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_PROVIDER", "mock")
|
||||
|
||||
with pytest.raises(ValueError, match="Invalid text_search_extension_native_language"):
|
||||
HindsightConfig.from_env()
|
||||
|
||||
|
||||
def test_text_search_extension_accepts_pgroonga(monkeypatch):
|
||||
from hindsight_api.config import HindsightConfig
|
||||
|
||||
monkeypatch.setenv("HINDSIGHT_API_TEXT_SEARCH_EXTENSION", "pgroonga")
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_PROVIDER", "mock")
|
||||
|
||||
config = HindsightConfig.from_env()
|
||||
assert config.text_search_extension == "pgroonga"
|
||||
|
||||
|
||||
def test_text_search_extension_rejects_unknown(monkeypatch):
|
||||
from hindsight_api.config import HindsightConfig
|
||||
|
||||
monkeypatch.setenv("HINDSIGHT_API_TEXT_SEARCH_EXTENSION", "bogus")
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_PROVIDER", "mock")
|
||||
|
||||
with pytest.raises(ValueError, match="Invalid text_search_extension"):
|
||||
HindsightConfig.from_env()
|
||||
|
||||
|
||||
def test_pg_search_tokenizer_defaults_to_empty(monkeypatch):
|
||||
from hindsight_api.config import HindsightConfig
|
||||
|
||||
monkeypatch.delenv("HINDSIGHT_API_TEXT_SEARCH_EXTENSION_PG_SEARCH_TOKENIZER", raising=False)
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_PROVIDER", "mock")
|
||||
|
||||
config = HindsightConfig.from_env()
|
||||
assert config.text_search_extension_pg_search_tokenizer == ""
|
||||
|
||||
|
||||
def test_pg_search_tokenizer_loaded_from_env(monkeypatch):
|
||||
from hindsight_api.config import HindsightConfig
|
||||
|
||||
monkeypatch.setenv("HINDSIGHT_API_TEXT_SEARCH_EXTENSION_PG_SEARCH_TOKENIZER", "Jieba")
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_PROVIDER", "mock")
|
||||
|
||||
config = HindsightConfig.from_env()
|
||||
assert config.text_search_extension_pg_search_tokenizer == "jieba"
|
||||
|
||||
|
||||
def test_pg_search_tokenizer_accepts_lindera_alias(monkeypatch):
|
||||
from hindsight_api.config import HindsightConfig
|
||||
|
||||
monkeypatch.setenv("HINDSIGHT_API_TEXT_SEARCH_EXTENSION_PG_SEARCH_TOKENIZER", "chinese_lindera")
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_PROVIDER", "mock")
|
||||
|
||||
config = HindsightConfig.from_env()
|
||||
assert config.text_search_extension_pg_search_tokenizer == "lindera(chinese)"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("bad_value", ["jieba;DROP TABLE", "ngram(3,2)", "unknown", "pdb.jieba"])
|
||||
def test_pg_search_tokenizer_rejects_invalid_values(monkeypatch, bad_value):
|
||||
from hindsight_api.config import HindsightConfig
|
||||
|
||||
monkeypatch.setenv("HINDSIGHT_API_TEXT_SEARCH_EXTENSION_PG_SEARCH_TOKENIZER", bad_value)
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_PROVIDER", "mock")
|
||||
|
||||
with pytest.raises(ValueError, match="Invalid HINDSIGHT_API_TEXT_SEARCH_EXTENSION_PG_SEARCH_TOKENIZER"):
|
||||
HindsightConfig.from_env()
|
||||
|
||||
|
||||
def test_pg_search_bm25_columns_apply_tokenizer():
|
||||
from hindsight_api._pg_search import pg_search_bm25_columns
|
||||
|
||||
assert pg_search_bm25_columns("id", ("text", "context"), "") == "id, text, context"
|
||||
assert pg_search_bm25_columns("id", ("text", "context"), "jieba") == "id, (text::pdb.jieba), (context::pdb.jieba)"
|
||||
assert pg_search_bm25_columns("id", ("text",), "ngram(2, 3)") == "id, (text::pdb.ngram(2,3))"
|
||||
assert pg_search_bm25_columns("id", ("text",), "edge_ngram(2, 5)") == "id, (text::pdb.edge_ngram(2,5))"
|
||||
|
||||
|
||||
def test_llm_output_language_defaults_to_none(monkeypatch):
|
||||
from hindsight_api.config import HindsightConfig
|
||||
|
||||
monkeypatch.delenv("HINDSIGHT_API_LLM_OUTPUT_LANGUAGE", raising=False)
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_PROVIDER", "mock")
|
||||
|
||||
config = HindsightConfig.from_env()
|
||||
assert config.llm_output_language is None
|
||||
|
||||
|
||||
def test_llm_output_language_loaded_from_env(monkeypatch):
|
||||
from hindsight_api.config import HindsightConfig
|
||||
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_OUTPUT_LANGUAGE", "Japanese")
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_PROVIDER", "mock")
|
||||
|
||||
config = HindsightConfig.from_env()
|
||||
assert config.llm_output_language == "Japanese"
|
||||
|
||||
|
||||
def test_llm_output_language_empty_string_is_unset(monkeypatch):
|
||||
"""Empty env var (e.g. from Helm) should be treated as unset, not literal ''."""
|
||||
from hindsight_api.config import HindsightConfig
|
||||
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_OUTPUT_LANGUAGE", "")
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_PROVIDER", "mock")
|
||||
|
||||
config = HindsightConfig.from_env()
|
||||
assert config.llm_output_language is None
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -77,7 +77,9 @@ async def test_round_limit_caps_processed_memories(memory: MemoryEngine, request
|
||||
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)
|
||||
mock_requeue.assert_called_once_with(
|
||||
bank_id=bank_id, request_context=request_context, observation_scopes=None
|
||||
)
|
||||
|
||||
# Mental model refresh should be skipped on intermediate round
|
||||
assert result.get("mental_models_refreshed", 0) == 0
|
||||
|
||||
@@ -210,6 +210,18 @@ class TestPostgreSQLDialect:
|
||||
assert "to_tsquery" in arm
|
||||
assert "'bm25' AS source" in arm
|
||||
assert "LIMIT $3" in arm
|
||||
# Default language is english when bm25_language is not specified
|
||||
assert "to_tsquery('english', $4)" in arm
|
||||
|
||||
def test_build_bm25_arm_native_uses_configured_language(self, d):
|
||||
arm = d.build_bm25_arm(
|
||||
table="schema.memory_units", cols="id, text", fact_type="world",
|
||||
bank_id_param="$2", limit_param="$3", text_param="$4",
|
||||
bm25_language="french",
|
||||
)
|
||||
# Both the score and the WHERE filter must use the configured dictionary
|
||||
assert "to_tsquery('french', $4)" in arm
|
||||
assert "to_tsquery('english'" not in arm
|
||||
|
||||
def test_build_bm25_arm_vchord(self, d):
|
||||
arm = d.build_bm25_arm(
|
||||
@@ -220,6 +232,46 @@ class TestPostgreSQLDialect:
|
||||
assert "to_bm25query" in arm
|
||||
assert "tokenize" in arm
|
||||
|
||||
def test_build_bm25_arm_pgroonga(self, d):
|
||||
arm = d.build_bm25_arm(
|
||||
table="schema.memory_units", cols="id, text", fact_type="world",
|
||||
bank_id_param="$2", limit_param="$3", text_param="$4",
|
||||
text_search_extension="pgroonga",
|
||||
)
|
||||
# pgroonga uses the &@~ operator + pgroonga_score for ranking. The
|
||||
# configured bm25_language is intentionally NOT used here — pgroonga's
|
||||
# tokenizer is set at index creation, not query time.
|
||||
assert "&@~ $4" in arm
|
||||
assert "pgroonga_score(tableoid, ctid)" in arm
|
||||
assert "to_tsquery" not in arm
|
||||
|
||||
def test_build_bm25_arm_pgroonga_ignores_bm25_language(self, d):
|
||||
"""pgroonga's tokenizer is fixed at index creation; bm25_language must not leak in."""
|
||||
arm = d.build_bm25_arm(
|
||||
table="t", cols="id", fact_type="world",
|
||||
bank_id_param="$2", limit_param="$3", text_param="$4",
|
||||
text_search_extension="pgroonga",
|
||||
bm25_language="french",
|
||||
)
|
||||
assert "french" not in arm
|
||||
|
||||
def test_build_bm25_arm_pg_search(self, d):
|
||||
arm = d.build_bm25_arm(
|
||||
table="schema.memory_units", cols="id, text", fact_type="world",
|
||||
bank_id_param="$2", limit_param="$3", text_param="$4",
|
||||
text_search_extension="pg_search",
|
||||
)
|
||||
assert "paradedb.score(id)" in arm
|
||||
# @@@ on the key_field requires a field-qualified query, so we
|
||||
# fan the bind param out across all indexed text fields.
|
||||
assert "id @@@ paradedb.boolean(should =>" in arm
|
||||
assert "paradedb.match('text', $4)" in arm
|
||||
assert "paradedb.match('context', $4)" in arm
|
||||
assert "paradedb.match('text_signals', $4)" in arm
|
||||
assert "paradedb.score(id) DESC" in arm
|
||||
assert "'bm25' AS source" in arm
|
||||
assert "LIMIT $3" in arm
|
||||
|
||||
def test_prepare_bm25_text_native(self, d):
|
||||
result = d.prepare_bm25_text(["hello", "world"], "hello world")
|
||||
assert result == "hello | world"
|
||||
@@ -228,6 +280,16 @@ class TestPostgreSQLDialect:
|
||||
result = d.prepare_bm25_text(["hello", "world"], "hello world", text_search_extension="vchord")
|
||||
assert result == "hello world"
|
||||
|
||||
def test_prepare_bm25_text_pgroonga(self, d):
|
||||
# pgroonga accepts raw query text via &@~ and parses it with its own
|
||||
# query syntax; we pass the original query through unchanged.
|
||||
result = d.prepare_bm25_text(["hello", "world"], "hello world", text_search_extension="pgroonga")
|
||||
assert result == "hello world"
|
||||
|
||||
def test_prepare_bm25_text_pg_search(self, d):
|
||||
result = d.prepare_bm25_text(["hello", "world"], "hello world", text_search_extension="pg_search")
|
||||
assert result == "hello world"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# OracleDialect tests (no oracledb dependency needed)
|
||||
|
||||
@@ -23,7 +23,10 @@ from hindsight_api import MemoryEngine, RequestContext
|
||||
_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")
|
||||
pytestmark = [
|
||||
pytest.mark.skipif(not _RUN, reason="Set HINDSIGHT_RUN_GEMINI_EVALS=1 + LLM API key"),
|
||||
pytest.mark.hs_llm_core,
|
||||
]
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test documents — short but representative
|
||||
@@ -95,10 +98,11 @@ class TestDeltaEditorialFusion:
|
||||
|
||||
async def test_delta_fuses_seo_and_brand_voice(
|
||||
self,
|
||||
memory: MemoryEngine,
|
||||
memory_real_llm: MemoryEngine,
|
||||
request_context: RequestContext,
|
||||
):
|
||||
bank_id = f"test-editorial-{uuid.uuid4().hex[:8]}"
|
||||
memory = memory_real_llm
|
||||
await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
|
||||
try:
|
||||
|
||||
@@ -276,6 +276,7 @@ async def memory_no_llm(pg0_db_url, embeddings, cross_encoder, query_analyzer):
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.flaky(reruns=2, reruns_delay=2)
|
||||
async def test_concurrent_upserts_no_duplicates(memory_no_llm, request_context):
|
||||
"""
|
||||
Stress test: N concurrent retains of the same document with different content.
|
||||
|
||||
@@ -206,13 +206,15 @@ async def test_document_without_metadata(memory, request_context):
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_document_persisted_with_zero_facts(memory, request_context):
|
||||
@pytest.mark.hs_llm_core
|
||||
async def test_document_persisted_with_zero_facts(memory_real_llm, request_context):
|
||||
"""
|
||||
Test that documents are persisted even when zero facts are extracted.
|
||||
|
||||
This is a regression test for issue #324 where documents with no extractable
|
||||
facts were reported as disappearing from the system.
|
||||
"""
|
||||
memory = memory_real_llm
|
||||
bank_id = f"test_zero_facts_{datetime.now(timezone.utc).timestamp()}"
|
||||
|
||||
try:
|
||||
@@ -258,12 +260,14 @@ async def test_document_persisted_with_zero_facts(memory, request_context):
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_document_persisted_with_zero_facts_batch(memory, request_context):
|
||||
@pytest.mark.hs_llm_core
|
||||
async def test_document_persisted_with_zero_facts_batch(memory_real_llm, request_context):
|
||||
"""
|
||||
Test that documents are persisted with zero facts in batch retain operations.
|
||||
|
||||
This tests the async batch code path to ensure it also handles zero facts correctly.
|
||||
"""
|
||||
memory = memory_real_llm
|
||||
bank_id = f"test_zero_facts_batch_{datetime.now(timezone.utc).timestamp()}"
|
||||
|
||||
try:
|
||||
@@ -314,13 +318,15 @@ async def test_document_persisted_with_zero_facts_batch(memory, request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_document_persisted_with_zero_facts_async_submit(memory, request_context):
|
||||
@pytest.mark.hs_llm_core
|
||||
async def test_document_persisted_with_zero_facts_async_submit(memory_real_llm, request_context):
|
||||
"""
|
||||
Test that documents are persisted with zero facts in fire-and-forget async retain.
|
||||
|
||||
This tests the submit_async_retain (background task) code path to ensure it also
|
||||
handles zero facts correctly.
|
||||
"""
|
||||
memory = memory_real_llm
|
||||
import asyncio
|
||||
|
||||
bank_id = f"test_zero_facts_async_{datetime.now(timezone.utc).timestamp()}"
|
||||
|
||||
@@ -7,6 +7,7 @@ limits (e.g. DashScope / Aliyun Tongyi cap at 10). Users must be able to overrid
|
||||
the batch size via env var so `encode()` splits into smaller chunks.
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
|
||||
import pytest
|
||||
@@ -22,6 +23,7 @@ def setup_test_env():
|
||||
"HINDSIGHT_API_EMBEDDINGS_OPENAI_API_KEY",
|
||||
"HINDSIGHT_API_EMBEDDINGS_OPENAI_MODEL",
|
||||
"HINDSIGHT_API_EMBEDDINGS_OPENAI_BATCH_SIZE",
|
||||
"HINDSIGHT_API_EMBEDDINGS_OPENAI_DIMENSIONS",
|
||||
"HINDSIGHT_API_EMBEDDINGS_OPENROUTER_API_KEY",
|
||||
"HINDSIGHT_API_LLM_API_KEY",
|
||||
"HINDSIGHT_API_LLM_PROVIDER",
|
||||
@@ -64,6 +66,17 @@ def test_openai_batch_size_env_var_is_read():
|
||||
assert config.embeddings_openai_batch_size == 10
|
||||
|
||||
|
||||
def test_openai_dimensions_env_var_is_read():
|
||||
"""HINDSIGHT_API_EMBEDDINGS_OPENAI_DIMENSIONS requests reduced OpenAI output dims."""
|
||||
from hindsight_api.config import HindsightConfig
|
||||
|
||||
os.environ["HINDSIGHT_API_LLM_PROVIDER"] = "mock"
|
||||
os.environ["HINDSIGHT_API_EMBEDDINGS_OPENAI_DIMENSIONS"] = "384"
|
||||
|
||||
config = HindsightConfig.from_env()
|
||||
assert config.embeddings_openai_dimensions == 384
|
||||
|
||||
|
||||
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
|
||||
@@ -92,6 +105,41 @@ def test_openrouter_provider_uses_configured_batch_size():
|
||||
assert embeddings.batch_size == 8
|
||||
|
||||
|
||||
def test_openai_codex_provider_uses_codex_oauth_token_and_configured_batch_size(tmp_path, monkeypatch):
|
||||
"""'openai-codex' embeddings reuse Codex OAuth auth without a separate API key."""
|
||||
from hindsight_api.engine.embeddings import CodexOAuthEmbeddings, create_embeddings_from_env
|
||||
|
||||
codex_dir = tmp_path / ".codex"
|
||||
codex_dir.mkdir()
|
||||
(codex_dir / "auth.json").write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"auth_mode": "chatgpt",
|
||||
"tokens": {
|
||||
"access_token": "codex-oauth-token-test",
|
||||
"account_id": "acct-test",
|
||||
},
|
||||
}
|
||||
)
|
||||
)
|
||||
|
||||
monkeypatch.setenv("HOME", str(tmp_path))
|
||||
os.environ["HINDSIGHT_API_LLM_PROVIDER"] = "mock"
|
||||
os.environ["HINDSIGHT_API_EMBEDDINGS_PROVIDER"] = "openai-codex"
|
||||
os.environ["HINDSIGHT_API_EMBEDDINGS_OPENAI_MODEL"] = "text-embedding-3-small"
|
||||
os.environ["HINDSIGHT_API_EMBEDDINGS_OPENAI_BATCH_SIZE"] = "7"
|
||||
os.environ["HINDSIGHT_API_EMBEDDINGS_OPENAI_DIMENSIONS"] = "384"
|
||||
|
||||
embeddings = create_embeddings_from_env()
|
||||
assert isinstance(embeddings, CodexOAuthEmbeddings)
|
||||
assert embeddings.provider_name == "openai-codex"
|
||||
assert embeddings.model == "text-embedding-3-small"
|
||||
assert embeddings.base_url == "https://api.openai.com/v1"
|
||||
assert embeddings.api_key == "codex-oauth-token-test"
|
||||
assert embeddings.batch_size == 7
|
||||
assert embeddings.dimensions == 384
|
||||
|
||||
|
||||
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
|
||||
@@ -144,7 +192,34 @@ def test_openai_encode_splits_on_configured_batch_size(monkeypatch):
|
||||
|
||||
vectors = emb.encode(["x"] * 25)
|
||||
|
||||
assert calls == [10, 10, 5]
|
||||
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}"
|
||||
|
||||
|
||||
def test_openai_encode_passes_configured_dimensions():
|
||||
"""OpenAI embeddings requests include the optional dimensions parameter when configured."""
|
||||
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,
|
||||
dimensions=384,
|
||||
)
|
||||
|
||||
calls: list[int | None] = []
|
||||
|
||||
def fake_create(*, model, input, dimensions=None):
|
||||
calls.append(dimensions)
|
||||
return SimpleNamespace(data=[SimpleNamespace(index=i, embedding=[0.0] * 384) for i in range(len(input))])
|
||||
|
||||
emb._client = SimpleNamespace(embeddings=SimpleNamespace(create=fake_create))
|
||||
emb._dimension = 384
|
||||
|
||||
vectors = emb.encode(["x"] * 2)
|
||||
|
||||
assert calls == [384]
|
||||
assert len(vectors) == 2
|
||||
assert len(vectors[0]) == 384
|
||||
|
||||
@@ -925,11 +925,13 @@ def test_extraction_schema_no_labels_when_unconfigured():
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retain_extracts_single_value_label(memory, request_context):
|
||||
@pytest.mark.hs_llm_core
|
||||
async def test_retain_extracts_single_value_label(memory_real_llm, request_context):
|
||||
"""
|
||||
End-to-end: retain content with entity_labels configured (single-value).
|
||||
Verify that the LLM assigns the label and it ends up as a key:value entity on the memory unit.
|
||||
"""
|
||||
memory = memory_real_llm
|
||||
from hindsight_api.engine.memory_engine import fq_table
|
||||
|
||||
bank_id = f"test-labels-single-{uuid.uuid4().hex[:8]}"
|
||||
@@ -992,11 +994,13 @@ async def test_retain_extracts_single_value_label(memory, request_context):
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retain_extracts_multi_value_label(memory, request_context):
|
||||
@pytest.mark.hs_llm_core
|
||||
async def test_retain_extracts_multi_value_label(memory_real_llm, request_context):
|
||||
"""
|
||||
End-to-end: retain content with a multi_value entity_labels group.
|
||||
Verify that multiple label values can be assigned to a single fact.
|
||||
"""
|
||||
memory = memory_real_llm
|
||||
from hindsight_api.engine.memory_engine import fq_table
|
||||
|
||||
bank_id = f"test-labels-multi-{uuid.uuid4().hex[:8]}"
|
||||
@@ -1058,12 +1062,14 @@ async def test_retain_extracts_multi_value_label(memory, request_context):
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retain_extracts_free_values_label(memory, request_context):
|
||||
@pytest.mark.hs_llm_core
|
||||
async def test_retain_extracts_free_values_label(memory_real_llm, request_context):
|
||||
"""
|
||||
End-to-end: retain content with a free_values entity_labels group.
|
||||
Verify that the LLM produces a key:value entity with an open-ended value
|
||||
(not constrained to a predefined enum list).
|
||||
"""
|
||||
memory = memory_real_llm
|
||||
from hindsight_api.engine.memory_engine import fq_table
|
||||
|
||||
bank_id = f"test-labels-free-{uuid.uuid4().hex[:8]}"
|
||||
@@ -1124,7 +1130,8 @@ async def test_retain_extracts_free_values_label(memory, request_context):
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retain_extracts_map_type_entities(memory, request_context):
|
||||
@pytest.mark.hs_llm_core
|
||||
async def test_retain_extracts_map_type_entities(memory_real_llm, request_context):
|
||||
"""
|
||||
End-to-end: retain content with a map-type entity_labels group.
|
||||
Verify that structured entity fields are extracted as key:field:value entity strings.
|
||||
@@ -1133,10 +1140,10 @@ async def test_retain_extracts_map_type_entities(memory, request_context):
|
||||
|
||||
bank_id = f"test-labels-map-{uuid.uuid4().hex[:8]}"
|
||||
try:
|
||||
await memory.get_bank_profile(bank_id=bank_id, request_context=request_context)
|
||||
await memory_real_llm.get_bank_profile(bank_id=bank_id, request_context=request_context)
|
||||
|
||||
# Configure a map-type entity label
|
||||
await memory._config_resolver.update_bank_config(
|
||||
await memory_real_llm._config_resolver.update_bank_config(
|
||||
bank_id=bank_id,
|
||||
updates={
|
||||
"entity_labels": [
|
||||
@@ -1156,7 +1163,7 @@ async def test_retain_extracts_map_type_entities(memory, request_context):
|
||||
context=request_context,
|
||||
)
|
||||
|
||||
unit_ids = await memory.retain_async(
|
||||
unit_ids = await memory_real_llm.retain_async(
|
||||
bank_id=bank_id,
|
||||
content=(
|
||||
"Alice Johnson is a Senior Software Engineer at Google. "
|
||||
@@ -1167,7 +1174,7 @@ async def test_retain_extracts_map_type_entities(memory, request_context):
|
||||
|
||||
assert len(unit_ids) > 0, "Should have extracted at least one fact"
|
||||
|
||||
async with memory._pool.acquire() as conn:
|
||||
async with memory_real_llm._pool.acquire() as conn:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT e.canonical_name
|
||||
@@ -1202,7 +1209,7 @@ async def test_retain_extracts_map_type_entities(memory, request_context):
|
||||
f"Free-form entities should not appear in labels-only mode. Got: {non_person_entities}"
|
||||
)
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
await memory_real_llm.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
# ─── map-type entity labels ──────────────────────────────────────────────────
|
||||
@@ -1894,3 +1901,356 @@ def test_duplicate_entity_strings_deduplicated():
|
||||
|
||||
texts = [e.text for e in validated]
|
||||
assert texts == ["person:name:Alice"] # only once
|
||||
|
||||
|
||||
# ─── GH-1558: multivalue tag entities missing from unit_entities ────────────
|
||||
|
||||
|
||||
def test_inject_label_tags_multivalue_all_tags_added():
|
||||
"""GH-1558 reproducer (unit-level): all multivalue entities with tag=True end up in tags."""
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from hindsight_api.engine.retain.fact_extraction import _inject_label_tags
|
||||
from hindsight_api.engine.retain.types import ExtractedFact
|
||||
|
||||
config = MagicMock()
|
||||
config.entity_labels = [
|
||||
{
|
||||
"key": "use",
|
||||
"type": "multi-values",
|
||||
"tag": True,
|
||||
"values": [
|
||||
{"value": "use-001"},
|
||||
{"value": "use-002"},
|
||||
{"value": "use-003"},
|
||||
],
|
||||
},
|
||||
]
|
||||
|
||||
fact = ExtractedFact(
|
||||
fact_text="System references use-001 and use-002",
|
||||
fact_type="world",
|
||||
entities=["use:use-001", "use:use-002"],
|
||||
tags=[],
|
||||
)
|
||||
_inject_label_tags([fact], config)
|
||||
|
||||
# Both label entities should be present in tags
|
||||
assert "use:use-001" in fact.tags
|
||||
assert "use:use-002" in fact.tags
|
||||
assert len(fact.tags) == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.hs_llm_core
|
||||
async def test_retain_multivalue_tag_entities_all_stored(memory_real_llm, request_context):
|
||||
"""
|
||||
GH-1558 reproducer (integration): retain content referencing multiple values
|
||||
of a multi-values entity label with tag=True.
|
||||
|
||||
Verify that ALL multivalue entities appear in BOTH:
|
||||
- memory_units.tags (the tags column)
|
||||
- unit_entities table (the entity links)
|
||||
|
||||
The original bug: tags are added correctly, but unit_entities only stores
|
||||
a subset (typically the first entity).
|
||||
"""
|
||||
from hindsight_api.engine.memory_engine import fq_table
|
||||
|
||||
bank_id = f"test-1558-multivalue-tag-{uuid.uuid4().hex[:8]}"
|
||||
try:
|
||||
await memory_real_llm.get_bank_profile(bank_id=bank_id, request_context=request_context)
|
||||
|
||||
# Configure entity labels matching the bug report scenario:
|
||||
# - multi-values type
|
||||
# - tag=True
|
||||
# - entities_allow_free_form=False
|
||||
await memory_real_llm._config_resolver.update_bank_config(
|
||||
bank_id=bank_id,
|
||||
updates={
|
||||
"entity_labels": [
|
||||
{
|
||||
"key": "use",
|
||||
"description": "Use case identifier for this section",
|
||||
"type": "multi-values",
|
||||
"tag": True,
|
||||
"values": [
|
||||
{"value": "use-001", "description": "First use case"},
|
||||
{"value": "use-002", "description": "Second use case"},
|
||||
{"value": "use-003", "description": "Third use case"},
|
||||
],
|
||||
}
|
||||
],
|
||||
"entities_allow_free_form": False,
|
||||
"retain_extraction_mode": "verbose",
|
||||
},
|
||||
context=request_context,
|
||||
)
|
||||
|
||||
# Content that explicitly references multiple use case identifiers
|
||||
# in a way that a single fact should capture both
|
||||
unit_ids = await memory_real_llm.retain_async(
|
||||
bank_id=bank_id,
|
||||
content=(
|
||||
"## System Integration Notes (use-001, use-002)\n\n"
|
||||
"This section covers both use-001 and use-002 use cases. "
|
||||
"The integration between use-001 (authentication flow) and "
|
||||
"use-002 (authorization flow) requires careful coordination. "
|
||||
"Both use-001 and use-002 must be tested together."
|
||||
),
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert len(unit_ids) > 0, "Should have extracted at least one fact"
|
||||
|
||||
async with memory_real_llm._pool.acquire() as conn:
|
||||
# Check entities in unit_entities table
|
||||
entity_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT e.canonical_name
|
||||
FROM {fq_table("unit_entities")} ue
|
||||
JOIN {fq_table("entities")} e ON e.id = ue.entity_id
|
||||
WHERE ue.unit_id = ANY($1::uuid[])
|
||||
""",
|
||||
[u for u in unit_ids],
|
||||
)
|
||||
entity_names = {r["canonical_name"].lower() for r in entity_rows}
|
||||
|
||||
# Check tags on memory_units
|
||||
tag_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT id, tags
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE id = ANY($1::uuid[])
|
||||
""",
|
||||
[u for u in unit_ids],
|
||||
)
|
||||
all_tags = set()
|
||||
for row in tag_rows:
|
||||
if row["tags"]:
|
||||
all_tags.update(t.lower() for t in row["tags"])
|
||||
|
||||
# Filter to use:* entities/tags
|
||||
use_entities = {n for n in entity_names if n.startswith("use:")}
|
||||
use_tags = {t for t in all_tags if t.startswith("use:")}
|
||||
|
||||
# The core assertion from GH-1558: tags and entities should match
|
||||
# Tags show both but entities only show a subset → BUG
|
||||
assert len(use_tags) >= 2, (
|
||||
f"Expected at least 2 use:* tags. Got: {use_tags}"
|
||||
)
|
||||
assert len(use_entities) >= 2, (
|
||||
f"GH-1558 BUG: Expected at least 2 use:* entities in unit_entities, "
|
||||
f"but only got {len(use_entities)}: {use_entities}. "
|
||||
f"Tags correctly show: {use_tags}"
|
||||
)
|
||||
# Every tag should also be an entity
|
||||
missing_entities = use_tags - use_entities
|
||||
assert len(missing_entities) == 0, (
|
||||
f"GH-1558 BUG: Tags {use_tags} were added but entities are missing: {missing_entities}. "
|
||||
f"Entities found: {use_entities}"
|
||||
)
|
||||
finally:
|
||||
await memory_real_llm.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.hs_llm_core
|
||||
async def test_retain_multivalue_tag_entities_second_retain(memory_real_llm, request_context):
|
||||
"""
|
||||
GH-1558 reproducer (second retain): entity resolution with existing entities.
|
||||
|
||||
On a second retain, entity resolution tries to match new entity names against
|
||||
existing entities in the bank. With very similar names like "use:use-001" and
|
||||
"use:use-002", the SequenceMatcher similarity is ~0.91 which combined with
|
||||
temporal proximity could exceed the 0.6 merge threshold, causing both to
|
||||
resolve to the same entity ID.
|
||||
"""
|
||||
from hindsight_api.engine.memory_engine import fq_table
|
||||
|
||||
bank_id = f"test-1558-second-{uuid.uuid4().hex[:8]}"
|
||||
try:
|
||||
await memory_real_llm.get_bank_profile(bank_id=bank_id, request_context=request_context)
|
||||
|
||||
await memory_real_llm._config_resolver.update_bank_config(
|
||||
bank_id=bank_id,
|
||||
updates={
|
||||
"entity_labels": [
|
||||
{
|
||||
"key": "use",
|
||||
"description": "Use case identifier",
|
||||
"type": "multi-values",
|
||||
"tag": True,
|
||||
"values": [
|
||||
{"value": "use-001", "description": "First use case"},
|
||||
{"value": "use-002", "description": "Second use case"},
|
||||
],
|
||||
}
|
||||
],
|
||||
"entities_allow_free_form": False,
|
||||
"retain_extraction_mode": "verbose",
|
||||
},
|
||||
context=request_context,
|
||||
)
|
||||
|
||||
# First retain: creates entities in the bank
|
||||
await memory_real_llm.retain_async(
|
||||
bank_id=bank_id,
|
||||
content=(
|
||||
"## Authentication Flow (use-001)\n\n"
|
||||
"The authentication flow use-001 handles user login via OAuth2."
|
||||
),
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Second retain: references BOTH use-001 and use-002
|
||||
# Entity resolution now has existing entities to match against
|
||||
unit_ids_2 = await memory_real_llm.retain_async(
|
||||
bank_id=bank_id,
|
||||
content=(
|
||||
"## Integration Notes (use-001, use-002)\n\n"
|
||||
"This section covers the integration between use-001 (authentication) "
|
||||
"and use-002 (authorization). Both use-001 and use-002 are required."
|
||||
),
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert len(unit_ids_2) > 0
|
||||
|
||||
async with memory_real_llm._pool.acquire() as conn:
|
||||
entity_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT e.canonical_name
|
||||
FROM {fq_table("unit_entities")} ue
|
||||
JOIN {fq_table("entities")} e ON e.id = ue.entity_id
|
||||
WHERE ue.unit_id = ANY($1::uuid[])
|
||||
""",
|
||||
[u for u in unit_ids_2],
|
||||
)
|
||||
entity_names = {r["canonical_name"].lower() for r in entity_rows}
|
||||
|
||||
tag_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT id, tags
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE id = ANY($1::uuid[])
|
||||
""",
|
||||
[u for u in unit_ids_2],
|
||||
)
|
||||
all_tags = set()
|
||||
for row in tag_rows:
|
||||
if row["tags"]:
|
||||
all_tags.update(t.lower() for t in row["tags"])
|
||||
|
||||
use_entities = {n for n in entity_names if n.startswith("use:")}
|
||||
use_tags = {t for t in all_tags if t.startswith("use:")}
|
||||
|
||||
assert len(use_tags) >= 2, (
|
||||
f"Expected at least 2 use:* tags on second retain. Got: {use_tags}"
|
||||
)
|
||||
assert len(use_entities) >= 2, (
|
||||
f"GH-1558 BUG: On second retain, expected at least 2 use:* entities "
|
||||
f"but only got {len(use_entities)}: {use_entities}. "
|
||||
f"Tags correctly show: {use_tags}. "
|
||||
f"Entity resolution may be merging similar names."
|
||||
)
|
||||
missing = use_tags - use_entities
|
||||
assert len(missing) == 0, (
|
||||
f"GH-1558 BUG: Tags present but entities missing after second retain: {missing}"
|
||||
)
|
||||
finally:
|
||||
await memory_real_llm.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_entity_resolution_does_not_merge_distinct_label_values(memory, request_context):
|
||||
"""
|
||||
GH-1558 reproducer (deterministic): directly test that entity resolution
|
||||
keeps distinct label values separate even when their names are very similar.
|
||||
|
||||
"use:use-001" and "use:use-002" have SequenceMatcher similarity of ~0.91.
|
||||
With the 0.6 merge threshold and temporal/co-occurrence boosts, the resolver
|
||||
might incorrectly merge them into a single entity.
|
||||
"""
|
||||
from hindsight_api.engine.memory_engine import fq_table
|
||||
from hindsight_api.engine.retain.entity_processing import resolve_entities
|
||||
from hindsight_api.engine.retain.types import EntityRef, ProcessedFact
|
||||
|
||||
bank_id = f"test-1558-resolve-{uuid.uuid4().hex[:8]}"
|
||||
try:
|
||||
await memory.get_bank_profile(bank_id=bank_id, request_context=request_context)
|
||||
|
||||
# First, insert a "use:use-001" entity into the bank so that
|
||||
# entity resolution has an existing entity to match against
|
||||
async with memory._pool.acquire() as conn:
|
||||
await conn.execute(
|
||||
f"""
|
||||
INSERT INTO {fq_table("entities")} (bank_id, canonical_name, first_seen, last_seen, mention_count)
|
||||
VALUES ($1, $2, now(), now(), 1)
|
||||
ON CONFLICT DO NOTHING
|
||||
""",
|
||||
bank_id,
|
||||
"use:use-001",
|
||||
)
|
||||
|
||||
# Now resolve entities for a fact that has BOTH use:use-001 and use:use-002
|
||||
from datetime import datetime, timezone
|
||||
|
||||
now = datetime.now(timezone.utc)
|
||||
facts = [
|
||||
ProcessedFact(
|
||||
fact_text="Integration between use-001 and use-002",
|
||||
fact_type="world",
|
||||
embedding=[0.0] * 384,
|
||||
occurred_start=now,
|
||||
occurred_end=None,
|
||||
mentioned_at=now,
|
||||
context="",
|
||||
metadata={},
|
||||
entities=[
|
||||
EntityRef(name="use:use-001"),
|
||||
EntityRef(name="use:use-002"),
|
||||
],
|
||||
content_index=0,
|
||||
tags=["use:use-001", "use:use-002"],
|
||||
)
|
||||
]
|
||||
|
||||
# Use placeholder unit IDs
|
||||
placeholder_unit_ids = [str(uuid.uuid4())]
|
||||
|
||||
entity_labels = [
|
||||
{
|
||||
"key": "use",
|
||||
"description": "Use case identifier",
|
||||
"type": "multi-values",
|
||||
"tag": True,
|
||||
"values": [
|
||||
{"value": "use-001"},
|
||||
{"value": "use-002"},
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
async with memory._pool.acquire() as conn:
|
||||
resolved_entity_ids, entity_to_unit, unit_to_entity_ids = await resolve_entities(
|
||||
entity_resolver=memory.entity_resolver,
|
||||
conn=conn,
|
||||
bank_id=bank_id,
|
||||
unit_ids=placeholder_unit_ids,
|
||||
facts=facts,
|
||||
entity_labels=entity_labels,
|
||||
)
|
||||
|
||||
# We should get 2 DISTINCT entity IDs, not the same ID twice
|
||||
assert len(resolved_entity_ids) == 2, (
|
||||
f"Expected 2 resolved entity IDs, got {len(resolved_entity_ids)}"
|
||||
)
|
||||
unique_ids = set(resolved_entity_ids)
|
||||
assert len(unique_ids) == 2, (
|
||||
f"GH-1558 BUG: Entity resolution merged 'use:use-001' and 'use:use-002' "
|
||||
f"into the same entity ID. Got IDs: {resolved_entity_ids}. "
|
||||
f"These are distinct label values and must NOT be merged."
|
||||
)
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@@ -13,6 +13,7 @@ from hindsight_api.extensions import (
|
||||
HttpExtension,
|
||||
OperationValidationError,
|
||||
OperationValidatorExtension,
|
||||
PrecheckContext,
|
||||
RecallContext,
|
||||
RecallResult,
|
||||
ReflectContext,
|
||||
@@ -814,3 +815,264 @@ class TestHttpExtensionIntegration:
|
||||
# Banks list endpoint should work
|
||||
response = client.get("/v1/default/banks")
|
||||
assert response.status_code in (200, 500) # May fail if DB not ready
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Precheck (pre-body-parse) tests
|
||||
# ============================================================================
|
||||
#
|
||||
# The precheck() hook is wired as a FastAPI Depends on the billable POST
|
||||
# routes. FastAPI resolves dependencies before deserialising the route's body
|
||||
# parameter, so a rejecting precheck never causes the request body to be read
|
||||
# or materialised in memory. The test below uses a Pydantic model_validator
|
||||
# that records every parse to assert that body parsing never runs on the
|
||||
# rejection path.
|
||||
|
||||
|
||||
class RecordingPrecheckValidator(OperationValidatorExtension):
|
||||
"""Validator that records every precheck call and can be configured to reject.
|
||||
|
||||
Used to drive the FastAPI dependency that runs precheck() before body parse.
|
||||
The validate_* hooks below are required-abstract no-ops so the class is
|
||||
instantiable; the tests here only exercise precheck.
|
||||
"""
|
||||
|
||||
def __init__(self, *, reject: bool = False, status_code: int = 402,
|
||||
reason: str = "rejected by precheck") -> None:
|
||||
super().__init__(config={})
|
||||
self.reject = reject
|
||||
self.status_code = status_code
|
||||
self.reason = reason
|
||||
self.precheck_calls: list[PrecheckContext] = []
|
||||
|
||||
async def precheck(self, ctx: PrecheckContext) -> ValidationResult:
|
||||
self.precheck_calls.append(ctx)
|
||||
if self.reject:
|
||||
return ValidationResult.reject(self.reason, status_code=self.status_code)
|
||||
return ValidationResult.accept()
|
||||
|
||||
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()
|
||||
|
||||
|
||||
class TestPrecheckDefault:
|
||||
"""The base OperationValidatorExtension.precheck is a no-op accept."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_default_precheck_accepts(self):
|
||||
validator = RecordingPrecheckValidator(reject=False)
|
||||
# Bypass our override by calling the base implementation directly.
|
||||
ctx = PrecheckContext(
|
||||
operation="retain",
|
||||
bank_id="bank-x",
|
||||
request_context=RequestContext(),
|
||||
)
|
||||
result = await OperationValidatorExtension.precheck(validator, ctx)
|
||||
assert result.allowed is True
|
||||
assert result.reason is None
|
||||
|
||||
|
||||
class TestPrecheckHttpWiring:
|
||||
"""precheck() is wired as a FastAPI Depends on the billable POST routes.
|
||||
|
||||
These tests do NOT use the heavy ``memory`` fixture (which requires a
|
||||
running pg0 + migrations). Instead they construct a minimal FastAPI
|
||||
app that mirrors the same Depends ordering used in
|
||||
``hindsight_api.api.http`` (a ``Depends(precheck_for(...))`` resolved
|
||||
before the Pydantic body parameter), so the contract under test —
|
||||
"rejection happens before body parse" — can be exercised in isolation.
|
||||
|
||||
The critical assertion in test_precheck_rejection_skips_body_parse is
|
||||
that a rejection response is returned without the request body being
|
||||
deserialised by Pydantic — i.e. the body parser was never invoked on
|
||||
the rejection path.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def _build_app(validator):
|
||||
"""Mirror the precheck wiring from ``hindsight_api.api.http`` in a
|
||||
standalone FastAPI app."""
|
||||
from fastapi import Depends, FastAPI, HTTPException
|
||||
from pydantic import BaseModel, model_validator
|
||||
|
||||
from hindsight_api.extensions import PrecheckContext
|
||||
from hindsight_api.models import RequestContext
|
||||
|
||||
body_parses: list[str] = []
|
||||
|
||||
class _RetainBody(BaseModel):
|
||||
items: list
|
||||
|
||||
@model_validator(mode="before")
|
||||
@classmethod
|
||||
def _record(cls, v):
|
||||
body_parses.append("retain")
|
||||
return v
|
||||
|
||||
class _RecallBody(BaseModel):
|
||||
query: str
|
||||
|
||||
@model_validator(mode="before")
|
||||
@classmethod
|
||||
def _record(cls, v):
|
||||
body_parses.append("recall")
|
||||
return v
|
||||
|
||||
class _ReflectBody(BaseModel):
|
||||
query: str
|
||||
|
||||
@model_validator(mode="before")
|
||||
@classmethod
|
||||
def _record(cls, v):
|
||||
body_parses.append("reflect")
|
||||
return v
|
||||
|
||||
async def _request_context() -> RequestContext:
|
||||
return RequestContext()
|
||||
|
||||
def _precheck_for(operation: str):
|
||||
async def _dep(
|
||||
bank_id: str,
|
||||
request_context: RequestContext = Depends(_request_context),
|
||||
) -> None:
|
||||
ctx = PrecheckContext(
|
||||
operation=operation,
|
||||
bank_id=bank_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
result = await validator.precheck(ctx)
|
||||
if not result.allowed:
|
||||
raise HTTPException(
|
||||
status_code=result.status_code,
|
||||
detail=result.reason or "Operation not allowed",
|
||||
)
|
||||
|
||||
return _dep
|
||||
|
||||
app = FastAPI()
|
||||
|
||||
@app.post("/v1/default/banks/{bank_id}/memories")
|
||||
async def retain(
|
||||
bank_id: str,
|
||||
body: _RetainBody,
|
||||
_: None = Depends(_precheck_for("retain")),
|
||||
):
|
||||
return {"ok": True, "bank_id": bank_id, "n": len(body.items)}
|
||||
|
||||
@app.post("/v1/default/banks/{bank_id}/memories/recall")
|
||||
async def recall(
|
||||
bank_id: str,
|
||||
body: _RecallBody,
|
||||
_: None = Depends(_precheck_for("recall")),
|
||||
):
|
||||
return {"ok": True}
|
||||
|
||||
@app.post("/v1/default/banks/{bank_id}/reflect")
|
||||
async def reflect(
|
||||
bank_id: str,
|
||||
body: _ReflectBody,
|
||||
_: None = Depends(_precheck_for("reflect")),
|
||||
):
|
||||
return {"ok": True}
|
||||
|
||||
@app.get("/v1/default/banks/{bank_id}/memories/list")
|
||||
async def list_memories(bank_id: str):
|
||||
return {"ok": True}
|
||||
|
||||
return app, body_parses
|
||||
|
||||
def test_precheck_accept_lets_request_through_to_body_parse(self):
|
||||
validator = RecordingPrecheckValidator(reject=False)
|
||||
app, body_parses = self._build_app(validator)
|
||||
client = TestClient(app)
|
||||
|
||||
resp = client.post(
|
||||
"/v1/default/banks/precheck-bank/memories",
|
||||
json={"items": [{"content": "x"}]},
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
assert len(validator.precheck_calls) == 1
|
||||
assert validator.precheck_calls[0].operation == "retain"
|
||||
assert validator.precheck_calls[0].bank_id == "precheck-bank"
|
||||
assert body_parses == ["retain"]
|
||||
|
||||
def test_precheck_rejection_returns_status_and_reason(self):
|
||||
validator = RecordingPrecheckValidator(
|
||||
reject=True, status_code=402, reason="Insufficient credits"
|
||||
)
|
||||
app, _ = self._build_app(validator)
|
||||
client = TestClient(app)
|
||||
|
||||
resp = client.post(
|
||||
"/v1/default/banks/precheck-bank/memories",
|
||||
json={"items": [{"content": "x"}]},
|
||||
)
|
||||
assert resp.status_code == 402
|
||||
assert resp.json()["detail"] == "Insufficient credits"
|
||||
|
||||
def test_precheck_rejection_skips_body_parse(self):
|
||||
"""The critical assertion: rejection happens before Pydantic
|
||||
deserialises the body. We send an oversized body and verify the
|
||||
body-parse counter never incremented.
|
||||
"""
|
||||
validator = RecordingPrecheckValidator(
|
||||
reject=True, status_code=402, reason="rejected by precheck"
|
||||
)
|
||||
app, body_parses = self._build_app(validator)
|
||||
client = TestClient(app)
|
||||
|
||||
resp = client.post(
|
||||
"/v1/default/banks/precheck-bank/memories",
|
||||
json={"items": [{"content": "x" * 100_000} for _ in range(50)]},
|
||||
)
|
||||
assert resp.status_code == 402
|
||||
assert "rejected by precheck" in resp.json()["detail"]
|
||||
assert body_parses == [], (
|
||||
"request body was deserialised despite a rejecting precheck — "
|
||||
"the Depends-before-body-parse contract is broken"
|
||||
)
|
||||
|
||||
def test_precheck_rejection_skips_body_parse_for_recall(self):
|
||||
validator = RecordingPrecheckValidator(
|
||||
reject=True, status_code=402, reason="rejected"
|
||||
)
|
||||
app, body_parses = self._build_app(validator)
|
||||
client = TestClient(app)
|
||||
|
||||
resp = client.post(
|
||||
"/v1/default/banks/precheck-bank/memories/recall",
|
||||
json={"query": "x" * 100_000},
|
||||
)
|
||||
assert resp.status_code == 402
|
||||
assert validator.precheck_calls[-1].operation == "recall"
|
||||
assert body_parses == []
|
||||
|
||||
def test_precheck_rejection_skips_body_parse_for_reflect(self):
|
||||
validator = RecordingPrecheckValidator(
|
||||
reject=True, status_code=402, reason="rejected"
|
||||
)
|
||||
app, body_parses = self._build_app(validator)
|
||||
client = TestClient(app)
|
||||
|
||||
resp = client.post(
|
||||
"/v1/default/banks/precheck-bank/reflect",
|
||||
json={"query": "x" * 100_000},
|
||||
)
|
||||
assert resp.status_code == 402
|
||||
assert validator.precheck_calls[-1].operation == "reflect"
|
||||
assert body_parses == []
|
||||
|
||||
def test_precheck_does_not_run_on_get(self):
|
||||
validator = RecordingPrecheckValidator(reject=True)
|
||||
app, _ = self._build_app(validator)
|
||||
client = TestClient(app)
|
||||
|
||||
resp = client.get("/v1/default/banks/precheck-bank/memories/list")
|
||||
assert resp.status_code == 200
|
||||
assert len(validator.precheck_calls) == 0
|
||||
|
||||
@@ -3,6 +3,7 @@ Test that first-person agent experiences are classified as 'experience' fact_typ
|
||||
not 'world'. This is critical for AI agent systems that store their own operational
|
||||
experiences (debugging, code changes, user interactions) separately from world knowledge.
|
||||
"""
|
||||
|
||||
from datetime import datetime
|
||||
|
||||
import pytest
|
||||
@@ -10,6 +11,9 @@ import pytest
|
||||
from hindsight_api import LLMConfig
|
||||
from hindsight_api.config import _get_raw_config
|
||||
from hindsight_api.engine.retain.fact_extraction import extract_facts_from_text
|
||||
from tests.llm_judge import assert_meets_criteria
|
||||
|
||||
pytestmark = pytest.mark.hs_llm_core
|
||||
|
||||
|
||||
class TestAgentExperienceClassification:
|
||||
@@ -61,12 +65,22 @@ I added a setup fixture that ensures the pool is warmed up, and all 47 tests pas
|
||||
)
|
||||
|
||||
assert len(facts) > 0, "Should extract at least one fact"
|
||||
world_facts = [f for f in facts if f.fact_type == "world"]
|
||||
experience_facts = [f for f in facts if f.fact_type == "experience"]
|
||||
assert len(experience_facts) > len(world_facts), (
|
||||
f"First-person debugging should be mostly 'experience', "
|
||||
f"got {len(experience_facts)} experience vs {len(world_facts)} world. "
|
||||
f"Facts: {[(f.fact, f.fact_type) for f in facts]}"
|
||||
|
||||
# Use LLM judge to evaluate classification quality — the exact ratio
|
||||
# of experience vs world facts is non-deterministic across providers.
|
||||
facts_summary = "\n".join(f"- [{f.fact_type}] {f.fact}" for f in facts)
|
||||
await assert_meets_criteria(
|
||||
response=facts_summary,
|
||||
criteria=(
|
||||
"The majority of facts extracted from this first-person debugging narrative "
|
||||
"should be classified as 'experience' (not 'world'), since the narrator is "
|
||||
"describing their own actions: tracing the bug, adding a fixture, seeing tests pass. "
|
||||
"At least some facts should be 'experience' type."
|
||||
),
|
||||
context=(
|
||||
"Input: First-person debugging session by coding-agent. "
|
||||
"Tests failed with ConnectionRefusedError, agent traced it, added a setup fixture, tests pass now."
|
||||
),
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
||||
@@ -9,9 +9,19 @@ This comprehensive test suite validates that the fact extraction system:
|
||||
5. Correctly attributes statements to speakers
|
||||
6. Filters out irrelevant content (podcast intros/outros)
|
||||
|
||||
These are quality/accuracy tests that verify the LLM-based extraction
|
||||
produces semantically correct and complete facts.
|
||||
Every test here exercises real LLM extraction behaviour — the file is marked
|
||||
hs_llm_core at module scope so it runs in the single-provider quality CI job.
|
||||
MockLLM cannot simulate dimension preservation, date conversion, or pronoun
|
||||
resolution; running these tests against a mock would either pass spuriously
|
||||
(MockLLM echoes input text, so string assertions trivially succeed) or fail
|
||||
with no diagnostic signal.
|
||||
|
||||
Semantic assertions go through tests.llm_judge so paraphrases survive — the
|
||||
LLM might phrase preserved emotion as "elated" instead of "thrilled", and a
|
||||
literal substring check would flake. Structural assertions (date fields,
|
||||
fact counts, fact_type classification) stay as direct asserts.
|
||||
"""
|
||||
|
||||
from datetime import UTC, datetime
|
||||
|
||||
import pytest
|
||||
@@ -19,11 +29,15 @@ import pytest
|
||||
from hindsight_api import LLMConfig
|
||||
from hindsight_api.config import _get_raw_config
|
||||
from hindsight_api.engine.retain.fact_extraction import extract_facts_from_text
|
||||
from tests.llm_judge import assert_meets_criteria
|
||||
|
||||
pytestmark = pytest.mark.hs_llm_core
|
||||
|
||||
# =============================================================================
|
||||
# DIMENSION PRESERVATION TESTS
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestDimensionPreservation:
|
||||
"""Tests that fact extraction preserves all information dimensions."""
|
||||
|
||||
@@ -54,15 +68,21 @@ Marcus felt anxious about the upcoming interview.
|
||||
)
|
||||
|
||||
assert len(facts) > 0, "Should extract at least one fact"
|
||||
all_facts_text = " ".join(f.fact for f in facts)
|
||||
|
||||
all_facts_text = " ".join([f.fact.lower() for f in facts])
|
||||
|
||||
emotional_indicators = ["thrilled", "disappointed", "anxious", "positive feedback"]
|
||||
found_emotions = [word for word in emotional_indicators if word in all_facts_text]
|
||||
|
||||
assert len(found_emotions) >= 2, (
|
||||
f"Should preserve emotional dimension. "
|
||||
f"Found: {found_emotions}, Expected at least 2 from: {emotional_indicators}"
|
||||
await assert_meets_criteria(
|
||||
response=all_facts_text,
|
||||
criteria=(
|
||||
"The extracted facts preserve emotional states from the input: the speaker's "
|
||||
"excitement/thrill about positive feedback, Sarah's disappointment about the delay, "
|
||||
"and Marcus's anxiety about the interview. At least two of these emotional dimensions "
|
||||
"should be present (exact wording doesn't matter — 'elated' for 'thrilled' is fine)."
|
||||
),
|
||||
context=(
|
||||
"Input mentioned: being thrilled about positive feedback on a presentation, "
|
||||
"Sarah seeming disappointed about a delay, and Marcus feeling anxious about an interview."
|
||||
),
|
||||
msg=f"Emotional dimension should be preserved. Facts: {[f.fact for f in facts]}",
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -87,15 +107,20 @@ The music was so loud I could barely hear myself think.
|
||||
)
|
||||
|
||||
assert len(facts) > 0, "Should extract at least one fact"
|
||||
all_facts_text = " ".join(f.fact for f in facts)
|
||||
|
||||
all_facts_text = " ".join([f.fact.lower() for f in facts])
|
||||
|
||||
sensory_indicators = ["bitter", "burnt", "bright orange", "loud", "stunning"]
|
||||
found_sensory = [word for word in sensory_indicators if word in all_facts_text]
|
||||
|
||||
assert len(found_sensory) >= 2, (
|
||||
f"Should preserve sensory details. "
|
||||
f"Found: {found_sensory}, Expected at least 2 from: {sensory_indicators}"
|
||||
await assert_meets_criteria(
|
||||
response=all_facts_text,
|
||||
criteria=(
|
||||
"The extracted facts preserve sensory details from the input — at least two of: "
|
||||
"the bitter/burnt taste of the coffee, the bright orange hair (and how it looked), "
|
||||
"or the loud volume of the music. Equivalent sensory descriptors are acceptable."
|
||||
),
|
||||
context=(
|
||||
"Input described: coffee that tasted bitter and burnt; bright orange hair that "
|
||||
"looked stunning under the lights; music so loud one could barely hear oneself think."
|
||||
),
|
||||
msg=f"Sensory dimension should be preserved. Facts: {[f.fact for f in facts]}",
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -121,15 +146,22 @@ Maybe we should reconsider the timeline.
|
||||
)
|
||||
|
||||
assert len(facts) > 0, "Should extract at least one fact"
|
||||
all_facts_text = " ".join(f.fact for f in facts)
|
||||
|
||||
all_facts_text = " ".join([f.fact.lower() for f in facts])
|
||||
|
||||
cognitive_indicators = ["realized", "wasn't sure", "convinced", "maybe", "reconsider"]
|
||||
found_cognitive = [word for word in cognitive_indicators if word in all_facts_text]
|
||||
|
||||
assert len(found_cognitive) >= 2, (
|
||||
f"Should preserve cognitive/epistemic dimension. "
|
||||
f"Found: {found_cognitive}"
|
||||
await assert_meets_criteria(
|
||||
response=all_facts_text,
|
||||
criteria=(
|
||||
"The extracted facts preserve cognitive or epistemic states from the input — "
|
||||
"at least two of: the realisation that the approach wasn't working, her uncertainty "
|
||||
"about whether the meeting would happen, his conviction that AI will transform "
|
||||
"healthcare, or the suggestion to reconsider the timeline. Equivalent phrasing "
|
||||
"(e.g. 'came to understand' for 'realised') is acceptable."
|
||||
),
|
||||
context=(
|
||||
"Input: realising an approach wasn't working; uncertainty about a meeting; "
|
||||
"conviction that AI will transform healthcare; a suggestion to reconsider the timeline."
|
||||
),
|
||||
msg=f"Cognitive/epistemic dimension should be preserved. Facts: {[f.fact for f in facts]}",
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -155,15 +187,22 @@ I'm unable to attend the conference due to scheduling conflicts.
|
||||
)
|
||||
|
||||
assert len(facts) > 0, "Should extract at least one fact"
|
||||
all_facts_text = " ".join(f.fact for f in facts)
|
||||
|
||||
all_facts_text = " ".join([f.fact.lower() for f in facts])
|
||||
|
||||
capability_indicators = ["can speak", "fluently", "struggles with", "expert in", "unable to"]
|
||||
found_capability = [word for word in capability_indicators if word in all_facts_text]
|
||||
|
||||
assert len(found_capability) >= 2, (
|
||||
f"Should preserve capability/skill dimension. "
|
||||
f"Found: {found_capability}"
|
||||
await assert_meets_criteria(
|
||||
response=all_facts_text,
|
||||
criteria=(
|
||||
"The extracted facts preserve capability/skill/limitation information from the input "
|
||||
"— at least two of: the speaker's fluency in French, Sarah's difficulty with public "
|
||||
"speaking, his expertise in machine learning, or the speaker's inability to attend "
|
||||
"the conference. Equivalent phrasing is fine."
|
||||
),
|
||||
context=(
|
||||
"Input: 'I can speak French fluently.', 'Sarah struggles with public speaking.', "
|
||||
"'He's an expert in machine learning.', 'I'm unable to attend the conference due "
|
||||
"to scheduling conflicts.'"
|
||||
),
|
||||
msg=f"Capability/skill dimension should be preserved. Facts: {[f.fact for f in facts]}",
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -188,15 +227,20 @@ Unlike last year, we're ahead of schedule.
|
||||
)
|
||||
|
||||
assert len(facts) > 0, "Should extract at least one fact"
|
||||
all_facts_text = " ".join(f.fact for f in facts)
|
||||
|
||||
all_facts_text = " ".join([f.fact.lower() for f in facts])
|
||||
|
||||
comparative_indicators = ["better than", "worse than", "unlike", "ahead of"]
|
||||
found_comparative = [word for word in comparative_indicators if word in all_facts_text]
|
||||
|
||||
assert len(found_comparative) >= 1, (
|
||||
f"Should preserve comparative dimension. "
|
||||
f"Found: {found_comparative}"
|
||||
await assert_meets_criteria(
|
||||
response=all_facts_text,
|
||||
criteria=(
|
||||
"At least one fact preserves a comparative or contrasting relationship from the "
|
||||
"input — that this approach is better than the previous one, that the new design "
|
||||
"is worse than expected, or that the team is ahead of schedule unlike last year."
|
||||
),
|
||||
context=(
|
||||
"Input: 'This approach is much better than the previous one.', 'The new design "
|
||||
"is worse than expected.', 'Unlike last year, we're ahead of schedule.'"
|
||||
),
|
||||
msg=f"Comparative dimension should be preserved. Facts: {[f.fact for f in facts]}",
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -222,15 +266,21 @@ She's enthusiastic about the opportunity.
|
||||
)
|
||||
|
||||
assert len(facts) > 0, "Should extract at least one fact"
|
||||
all_facts_text = " ".join(f.fact for f in facts)
|
||||
|
||||
all_facts_text = " ".join([f.fact.lower() for f in facts])
|
||||
|
||||
attitudinal_indicators = ["skeptical", "surprised", "rolled his eyes", "enthusiastic"]
|
||||
found_attitudinal = [word for word in attitudinal_indicators if word in all_facts_text]
|
||||
|
||||
assert len(found_attitudinal) >= 1, (
|
||||
f"Should preserve attitudinal/reactive dimension. "
|
||||
f"Found: {found_attitudinal}"
|
||||
await assert_meets_criteria(
|
||||
response=all_facts_text,
|
||||
criteria=(
|
||||
"At least one fact preserves an attitude or reaction from the input — her skepticism "
|
||||
"about the new technology, the speaker's surprise at his resignation, Marcus rolling "
|
||||
"his eyes (a non-verbal reaction), or her enthusiasm about the opportunity."
|
||||
),
|
||||
context=(
|
||||
"Input: 'She's very skeptical about the new technology.', 'I was surprised when he "
|
||||
"announced his resignation.', 'Marcus rolled his eyes when the topic came up.', "
|
||||
"'She's enthusiastic about the opportunity.'"
|
||||
),
|
||||
msg=f"Attitudinal/reactive dimension should be preserved. Facts: {[f.fact for f in facts]}",
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -256,19 +306,22 @@ I'm planning to switch careers because I'm not fulfilled in my current role.
|
||||
)
|
||||
|
||||
assert len(facts) > 0, "Should extract at least one fact"
|
||||
all_facts_text = " ".join(f.fact for f in facts)
|
||||
|
||||
all_facts_text = " ".join([f.fact.lower() for f in facts])
|
||||
|
||||
# Check for goal/intention related content
|
||||
intentional_indicators = [
|
||||
"want", "aim", "goal", "plan", "because", "learn", "complete",
|
||||
"build", "switch", "career", "mandarin", "china", "phd", "business"
|
||||
]
|
||||
found_intentional = [word for word in intentional_indicators if word in all_facts_text]
|
||||
|
||||
assert len(found_intentional) >= 1, (
|
||||
f"Should preserve intentional/motivational content. "
|
||||
f"Found: {found_intentional}"
|
||||
await assert_meets_criteria(
|
||||
response=all_facts_text,
|
||||
criteria=(
|
||||
"The extracted facts preserve goals, plans, or motivations from the input — at "
|
||||
"least one of: the speaker wanting to learn Mandarin before a trip to China, her "
|
||||
"PhD timeline goal, his goal of building a sustainable business, or the speaker's "
|
||||
"plan to switch careers because of unfulfilment."
|
||||
),
|
||||
context=(
|
||||
"Input mentioned: wanting to learn Mandarin before a China trip; aiming to complete "
|
||||
"a PhD within three years; a goal to build a sustainable business; planning to "
|
||||
"switch careers due to lack of fulfilment in current role."
|
||||
),
|
||||
msg=f"Intentional/motivational content should be preserved. Facts: {[f.fact for f in facts]}",
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -294,21 +347,35 @@ Family is the most important thing to her.
|
||||
)
|
||||
|
||||
assert len(facts) > 0, "Should extract at least one fact"
|
||||
all_facts_text = " ".join(f.fact for f in facts)
|
||||
|
||||
all_facts_text = " ".join([f.fact.lower() for f in facts])
|
||||
|
||||
evaluative_indicators = ["prefer", "values", "hates", "important", "above all"]
|
||||
found_evaluative = [word for word in evaluative_indicators if word in all_facts_text]
|
||||
|
||||
assert len(found_evaluative) >= 2, (
|
||||
f"Should preserve evaluative/preferential dimension. "
|
||||
f"Found: {found_evaluative}"
|
||||
await assert_meets_criteria(
|
||||
response=all_facts_text,
|
||||
criteria=(
|
||||
"The extracted facts preserve preferences or values from the input — at least two "
|
||||
"of: the speaker's preference for remote work over the office, her valuing honesty "
|
||||
"above all, his dislike of being late to meetings, or family being the most "
|
||||
"important thing to her."
|
||||
),
|
||||
context=(
|
||||
"Input: 'I prefer working remotely to being in an office.', 'She values honesty "
|
||||
"above all else.', 'He hates being late to meetings.', 'Family is the most "
|
||||
"important thing to her.'"
|
||||
),
|
||||
msg=f"Evaluative/preferential dimension should be preserved. Facts: {[f.fact for f in facts]}",
|
||||
)
|
||||
|
||||
@pytest.mark.hs_llm_mat
|
||||
@pytest.mark.asyncio
|
||||
async def test_comprehensive_multi_dimension(self):
|
||||
"""Test a realistic scenario with multiple dimensions in one fact."""
|
||||
"""Test a realistic scenario with multiple dimensions in one fact.
|
||||
|
||||
Inherits the module-level `hs_llm_core` marker rather than running on
|
||||
the full `hs_llm_mat` matrix — the weakest matrix providers (notably
|
||||
bedrock/nova-2-lite) drop one of the two required dimensions (emotional
|
||||
or preferential) and fail the judge. This is a quality assertion, not
|
||||
a provider-compatibility check, so single-strong-provider is the right
|
||||
tier.
|
||||
"""
|
||||
text = """
|
||||
I was thrilled to receive such positive feedback on my presentation yesterday!
|
||||
I wasn't sure if my approach would resonate, but the audience seemed enthusiastic.
|
||||
@@ -331,39 +398,35 @@ I prefer presenting in person rather than virtually because I can read the room
|
||||
|
||||
assert len(facts) > 0, "Should extract at least one fact"
|
||||
|
||||
all_facts_text = " ".join([f.fact.lower() for f in facts])
|
||||
all_facts_text = " ".join([f.fact for f in facts])
|
||||
|
||||
# Check emotional - should capture positive/thrilled sentiment
|
||||
has_emotional = any(term in all_facts_text for term in [
|
||||
"thrilled", "positive feedback", "positive", "feedback", "enthusiastic"
|
||||
])
|
||||
|
||||
# Check preference - should capture the in-person vs virtual preference
|
||||
has_preference = any(term in all_facts_text for term in [
|
||||
"prefer", "rather than", "in person", "in-person", "virtually",
|
||||
"read the room", "face-to-face", "face to face", "remote",
|
||||
])
|
||||
|
||||
# MAT bar: at least one of emotional or preferential must be preserved.
|
||||
# Smaller models (e.g. nova-2-lite) may compress both sentences into a
|
||||
# single fact that only captures one dimension — that's acceptable for
|
||||
# a minimum-acceptance test.
|
||||
assert has_emotional or has_preference, (
|
||||
f"Should preserve at least one of emotional or preferential dimension. "
|
||||
f"Extracted facts: {all_facts_text}"
|
||||
)
|
||||
|
||||
# Check no vague temporal terms
|
||||
# Check no vague temporal terms (structural check — not LLM-dependent)
|
||||
prohibited_terms = ["recently", "soon", "lately"]
|
||||
found_prohibited = [term for term in prohibited_terms if term in all_facts_text]
|
||||
assert len(found_prohibited) == 0, \
|
||||
f"Should NOT use vague temporal terms. Found: {found_prohibited}"
|
||||
found_prohibited = [term for term in prohibited_terms if term in all_facts_text.lower()]
|
||||
assert len(found_prohibited) == 0, f"Should NOT use vague temporal terms. Found: {found_prohibited}"
|
||||
|
||||
# Check emotional and preferential dimensions via LLM judge
|
||||
await assert_meets_criteria(
|
||||
response=all_facts_text,
|
||||
criteria=(
|
||||
"The extracted facts preserve BOTH of these dimensions from the input: "
|
||||
"(1) emotional — any mention of positive feedback, enthusiasm, thrilled, or positive sentiment, "
|
||||
"(2) preferential — any mention of preferring in-person presentations or reading the room. "
|
||||
"The facts don't need to use the exact same words — semantic equivalents count."
|
||||
),
|
||||
context=(
|
||||
"Input text: Was thrilled about positive feedback on presentation. "
|
||||
"Audience seemed enthusiastic. Prefers presenting in person rather than "
|
||||
"virtually because they can read the room better."
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# TEMPORAL CONVERSION TESTS
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestTemporalConversion:
|
||||
"""Tests for temporal extraction and date conversion."""
|
||||
|
||||
@@ -403,15 +466,12 @@ I'm planning to visit Tokyo next month.
|
||||
prohibited_terms = ["recently", "lately", "a while ago", "some time ago"]
|
||||
found_prohibited = [term for term in prohibited_terms if term in all_facts_text]
|
||||
|
||||
assert len(found_prohibited) == 0, (
|
||||
f"Should NOT use vague temporal terms. Found: {found_prohibited}"
|
||||
)
|
||||
assert len(found_prohibited) == 0, f"Should NOT use vague temporal terms. Found: {found_prohibited}"
|
||||
|
||||
# Check that at least one fact has a valid occurred_start date
|
||||
facts_with_temporal = [f for f in facts if f.occurred_start]
|
||||
assert len(facts_with_temporal) >= 1, (
|
||||
f"At least one fact should have temporal data (occurred_start). "
|
||||
f"Facts: {[f.fact for f in facts]}"
|
||||
f"At least one fact should have temporal data (occurred_start). Facts: {[f.fact for f in facts]}"
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -460,8 +520,8 @@ with a concert surrounded by music, joy and the warm summer breeze.
|
||||
fact_date_str = birthday_fact.occurred_start
|
||||
assert fact_date_str is not None, "occurred_start should not be None for temporal events"
|
||||
|
||||
if 'T' in fact_date_str:
|
||||
fact_date = datetime.fromisoformat(fact_date_str.replace('Z', '+00:00'))
|
||||
if "T" in fact_date_str:
|
||||
fact_date = datetime.fromisoformat(fact_date_str.replace("Z", "+00:00"))
|
||||
else:
|
||||
fact_date = datetime.fromisoformat(fact_date_str)
|
||||
|
||||
@@ -527,29 +587,39 @@ It was a beautiful day and I plan to make this a regular habit.
|
||||
if facts_with_date:
|
||||
jogging_fact = facts_with_date[0]
|
||||
fact_date_str = jogging_fact.occurred_start
|
||||
if 'T' in fact_date_str:
|
||||
fact_date = datetime.fromisoformat(fact_date_str.replace('Z', '+00:00'))
|
||||
if "T" in fact_date_str:
|
||||
fact_date = datetime.fromisoformat(fact_date_str.replace("Z", "+00:00"))
|
||||
else:
|
||||
fact_date = datetime.fromisoformat(fact_date_str)
|
||||
|
||||
assert fact_date.year == 2024, "Year should be 2024"
|
||||
assert fact_date.month == 11, "Month should be November"
|
||||
# Accept day 12 (ideal: yesterday) or 13 (conversation date) as valid
|
||||
assert fact_date.day in (12, 13), (
|
||||
f"Day should be 12 or 13 (around Nov 13 event), but got {fact_date.day}."
|
||||
)
|
||||
assert fact_date.day in (12, 13), f"Day should be 12 or 13 (around Nov 13 event), but got {fact_date.day}."
|
||||
|
||||
all_facts_text = " ".join([f.fact.lower() for f in facts])
|
||||
all_facts_text_lower = " ".join(f.fact.lower() for f in facts)
|
||||
all_facts_text = " ".join(f.fact for f in facts)
|
||||
|
||||
# The content should be preserved in some form
|
||||
assert any(term in all_facts_text for term in ["jog", "morning", "park", "first"]), \
|
||||
f"Should preserve key content. Facts: {[f.fact for f in facts]}"
|
||||
# Structural: "recently" is a prohibited vague term — the LLM must convert
|
||||
# "yesterday" to a concrete date, not paraphrase it as something equally vague.
|
||||
assert "recently" not in all_facts_text_lower, "Should NOT convert 'yesterday' to 'recently'"
|
||||
|
||||
assert "recently" not in all_facts_text, \
|
||||
"Should NOT convert 'yesterday' to 'recently'"
|
||||
|
||||
assert any(term in all_facts_text for term in ["november", "12", "nov"]), \
|
||||
"Should convert 'yesterday' to absolute date in fact text"
|
||||
# Semantic: content preservation AND date conversion go through the judge so
|
||||
# paraphrases ("ran" for "jog", "Nov 12 2024" for "November 12") still satisfy.
|
||||
await assert_meets_criteria(
|
||||
response=all_facts_text,
|
||||
criteria=(
|
||||
"The extracted facts (1) preserve the activity content — that the speaker went for "
|
||||
"a morning jog/run for the first time in a nearby park — and (2) reflect that the "
|
||||
"event happened on November 12, 2024 (the day before the conversation), either by "
|
||||
"stating the absolute date in the fact text or by using an unambiguous reference."
|
||||
),
|
||||
context=(
|
||||
"Conversation date: 2024-11-13. Input: 'Yesterday I went for a morning jog for the "
|
||||
"first time in a nearby park. It was a beautiful day...'"
|
||||
),
|
||||
msg=f"Yesterday content and date conversion should be preserved. Facts: {[f.fact for f in facts]}",
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_extract_facts_with_relative_dates(self):
|
||||
@@ -642,6 +712,7 @@ It was a beautiful day and I plan to make this a regular habit.
|
||||
# LOGICAL INFERENCE TESTS
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestLogicalInference:
|
||||
"""Tests that the system makes logical inferences to connect related information."""
|
||||
|
||||
@@ -685,31 +756,20 @@ great time! Every time I see it, I can't help but smile.
|
||||
)
|
||||
|
||||
assert len(facts) > 0, "Should extract at least one fact"
|
||||
all_facts_text = " ".join(f.fact.lower() for f in facts)
|
||||
|
||||
all_facts_text = " ".join([f.fact.lower() for f in facts])
|
||||
|
||||
# Check that key information is extracted (Karlie and the loss)
|
||||
has_karlie = "karlie" in all_facts_text
|
||||
has_loss = any(word in all_facts_text for word in ["lost", "death", "passed", "died", "losing", "friend"])
|
||||
has_hike = "hike" in all_facts_text or "hiking" in all_facts_text or "photo" in all_facts_text
|
||||
|
||||
# At minimum, we should capture Karlie and either the loss or the hike memory
|
||||
assert has_karlie or has_loss, (
|
||||
f"Should mention either Karlie or the loss in facts. Facts: {[f.fact for f in facts]}"
|
||||
# Key-information preservation is structural — required tokens are
|
||||
# proper nouns and a small set of loss-related verbs that the LLM
|
||||
# can't paraphrase away without losing the meaning. The judge proved
|
||||
# too strict here (it kept reading the facts and asking for explicit
|
||||
# connection prose), so this is a deterministic substring check —
|
||||
# same shape as the original pre-migration assertion.
|
||||
assert "karlie" in all_facts_text, f"Should mention Karlie. Facts: {[f.fact for f in facts]}"
|
||||
loss_terms = ("lost", "loss", "losing", "passed", "died", "death")
|
||||
assert any(t in all_facts_text for t in loss_terms), (
|
||||
f"Should mention the loss (one of {loss_terms}). Facts: {[f.fact for f in facts]}"
|
||||
)
|
||||
|
||||
# Check if inference was made (bonus - not required for pass)
|
||||
connected_fact_found = False
|
||||
for fact in facts:
|
||||
fact_text = fact.fact.lower()
|
||||
if "karlie" in fact_text and any(word in fact_text for word in ["lost", "death", "passed", "died", "losing", "friend"]):
|
||||
connected_fact_found = True
|
||||
break
|
||||
|
||||
# This is informational - test passes even without perfect inference
|
||||
if not connected_fact_found and has_karlie and has_loss:
|
||||
pass # Acceptable: facts extracted separately
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_logical_inference_pronoun_resolution(self):
|
||||
"""
|
||||
@@ -737,30 +797,26 @@ I've learned so much from it.
|
||||
|
||||
assert len(facts) > 0, "Should extract at least one fact"
|
||||
|
||||
all_facts_text = " ".join([f.fact.lower() for f in facts])
|
||||
# Pronoun resolution is a *structural* property — every fact that describes a
|
||||
# quality (challenging, rewarding, learning) must also name a specific anchor
|
||||
# noun (project / work / ML / etc.) in that same fact. The judge handled
|
||||
# this poorly in practice (hallucinating about pronouns that weren't there),
|
||||
# so this check is done deterministically: each quality-describing fact must
|
||||
# mention a project-anchor noun.
|
||||
quality_words = ("challenging", "rewarding", "learn", "demanding", "fulfilling", "rough", "tough")
|
||||
anchor_words = ("project", "ml", "machine learning", "work")
|
||||
|
||||
has_project = "project" in all_facts_text
|
||||
has_qualities = any(word in all_facts_text for word in ["challenging", "rewarding", "learned"])
|
||||
bad_facts = []
|
||||
for f in facts:
|
||||
fact_lower = f.fact.lower()
|
||||
if any(q in fact_lower for q in quality_words) and not any(a in fact_lower for a in anchor_words):
|
||||
bad_facts.append(f.fact)
|
||||
|
||||
assert has_project, "Should mention the project"
|
||||
assert has_qualities, "Should mention the qualities/learning"
|
||||
|
||||
# Check that pronouns are resolved - either:
|
||||
# 1. "project" appears with characteristics in same fact, OR
|
||||
# 2. "project" is explicitly mentioned in multiple facts (showing pronoun resolution)
|
||||
# The key is that "it" should be resolved to "project" rather than left as ambiguous
|
||||
project_facts = [f for f in facts if "project" in f.fact.lower()]
|
||||
|
||||
# If we have multiple facts mentioning project, pronoun resolution worked
|
||||
# (the LLM connected "it" back to "project" in subsequent facts)
|
||||
pronoun_resolved = len(project_facts) >= 2 or any(
|
||||
"project" in f.fact.lower() and any(word in f.fact.lower() for word in ["challenging", "rewarding", "learned"])
|
||||
for f in facts
|
||||
)
|
||||
|
||||
assert pronoun_resolved, (
|
||||
"Should resolve 'it' to 'the project' - either in combined facts or by mentioning project in multiple facts. "
|
||||
f"Facts: {[f.fact for f in facts]}"
|
||||
assert not bad_facts, (
|
||||
"Pronoun 'it' should be resolved to a specific noun anchor in every "
|
||||
"quality-describing fact, but these facts lack a project/work/ML anchor:\n"
|
||||
+ "\n".join(f" - {bf}" for bf in bad_facts)
|
||||
+ f"\nAll facts: {[f.fact for f in facts]}"
|
||||
)
|
||||
|
||||
|
||||
@@ -768,6 +824,7 @@ I've learned so much from it.
|
||||
# FACT CLASSIFICATION TESTS
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestFactClassification:
|
||||
"""Tests that facts are correctly classified as agent vs world."""
|
||||
|
||||
@@ -806,27 +863,29 @@ Jamie: Congratulations! I'd love to read it.
|
||||
)
|
||||
|
||||
assert len(facts) > 0, "Should extract at least one fact from the transcript"
|
||||
all_facts_text = " ".join(f.fact for f in facts)
|
||||
|
||||
# Check that we extracted meaningful content about AI research
|
||||
all_facts_text = " ".join([f.fact.lower() for f in facts])
|
||||
has_ai_content = any(term in all_facts_text for term in [
|
||||
"ai", "safety", "interpretability", "research", "paper", "conference", "models"
|
||||
])
|
||||
assert has_ai_content, f"Should extract AI research content. Facts: {[f.fact for f in facts]}"
|
||||
# The transcript is dense with AI-research content from Marcus. Extraction
|
||||
# must surface that subject matter — paraphrases like "alignment work" for
|
||||
# "AI safety research" should count, so the judge handles the assertion.
|
||||
await assert_meets_criteria(
|
||||
response=all_facts_text,
|
||||
criteria=(
|
||||
"The extracted facts cover the AI-research subject matter from Marcus's statements "
|
||||
"— at least mentioning AI/ML safety, interpretability, or his recent paper and "
|
||||
"upcoming conference presentation."
|
||||
),
|
||||
context=(
|
||||
"Marcus (the 'you' agent) said: working on AI safety research for six months, "
|
||||
"investigating interpretability methods, published a paper last month, presenting "
|
||||
"at a conference next week."
|
||||
),
|
||||
msg=f"Should extract AI research content from transcript. Facts: {[f.fact for f in facts]}",
|
||||
)
|
||||
|
||||
# Check fact type classification (flexible - may vary by LLM)
|
||||
agent_facts = [f for f in facts if f.fact_type == "agent"]
|
||||
experience_facts = [f for f in facts if f.fact_type == "experience"]
|
||||
|
||||
# Accept either agent or experience facts as valid for first-person statements
|
||||
first_person_facts = agent_facts + experience_facts
|
||||
|
||||
# If we have agent facts, verify they use first person
|
||||
for agent_fact in agent_facts:
|
||||
fact_text = agent_fact.fact
|
||||
# Allow flexibility - fact may or may not start with "I"
|
||||
if fact_text.startswith("I ") or " I " in fact_text:
|
||||
pass # Good - uses first person
|
||||
# Classification check is informational — many models split between 'agent'
|
||||
# and 'experience' for first-person statements. We don't assert on the
|
||||
# split, only that one of them is non-empty for Marcus's claims.
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_agent_facts_without_explicit_context(self):
|
||||
@@ -888,23 +947,42 @@ Jamie: [teasing] We'll see who's right, my Niners pick is solid.
|
||||
)
|
||||
|
||||
assert len(facts) > 0, "Should extract at least one fact"
|
||||
all_facts_text = " ".join(f.fact for f in facts)
|
||||
|
||||
# Check that predictions were extracted
|
||||
all_facts_text = " ".join([f.fact.lower() for f in facts])
|
||||
# The judge evaluates speaker attribution rather than substring-matching team
|
||||
# names — paraphrases like "the home team" or "San Francisco's squad" should
|
||||
# still satisfy the prediction-content criterion.
|
||||
await assert_meets_criteria(
|
||||
response=all_facts_text,
|
||||
criteria=(
|
||||
"The extracted facts capture at least one of the predictions made in the "
|
||||
"podcast — that the Rams will win 27-24 (Marcus's pick), or that the 49ers/Niners "
|
||||
"will win 27-13 (Jamie's pick). Equivalent wording about team names or scores counts."
|
||||
),
|
||||
context=(
|
||||
"Marcus (agent) predicted: Rams win 27-24. Jamie predicted: Niners win 27-13. "
|
||||
"Both predictions appear in the transcript."
|
||||
),
|
||||
msg=f"Should extract prediction content. Facts: {[f.fact for f in facts]}",
|
||||
)
|
||||
|
||||
# Should capture at least some prediction content
|
||||
has_prediction_content = any(term in all_facts_text for term in [
|
||||
"rams", "niners", "49ers", "prediction", "win", "predict"
|
||||
])
|
||||
assert has_prediction_content, f"Should extract prediction content. Facts: {[f.fact for f in facts]}"
|
||||
|
||||
# Ideally, Marcus's prediction should be in agent facts, but we accept
|
||||
# any reasonable extraction of the predictions
|
||||
# Speaker attribution is the deeper concern (Jamie's prediction must not be
|
||||
# attributed to Marcus). If agent_facts exist, ensure they don't claim Jamie's
|
||||
# Niners pick as Marcus's.
|
||||
agent_facts = [f for f in facts if f.fact_type == "agent"]
|
||||
if agent_facts:
|
||||
agent_facts_text = " ".join([f.fact.lower() for f in agent_facts])
|
||||
# If agent facts exist, they should relate to Marcus's statements
|
||||
# (but we don't fail if classification varies)
|
||||
agent_text = " ".join(f.fact for f in agent_facts)
|
||||
await assert_meets_criteria(
|
||||
response=agent_text,
|
||||
criteria=(
|
||||
"No agent fact (which represents Marcus's own statements) attributes the "
|
||||
"'Niners 27-13' prediction to Marcus. Marcus picked the Rams; Jamie picked "
|
||||
"the Niners. Marcus's facts may include his Rams prediction but must not "
|
||||
"claim he predicted a Niners win."
|
||||
),
|
||||
context="Marcus is the agent. He predicted Rams 27-24. Jamie predicted Niners 27-13.",
|
||||
msg=f"Jamie's prediction should not be misattributed to Marcus. Agent facts: {[f.fact for f in agent_facts]}",
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_skip_podcast_meta_commentary(self):
|
||||
@@ -955,17 +1033,26 @@ so the algorithm learns to box out. See you next week!
|
||||
)
|
||||
|
||||
assert len(facts) > 0, "Should extract at least one fact"
|
||||
all_facts_text = " ".join(f.fact for f in facts)
|
||||
|
||||
# The main goal is to extract substantive content about AI research
|
||||
# Meta-commentary filtering is ideal but not strictly required
|
||||
all_facts_text = " ".join([f.fact.lower() for f in facts])
|
||||
|
||||
# Should extract the actual AI research content
|
||||
has_substantive_content = any(term in all_facts_text for term in [
|
||||
"interpretability", "ai", "safety", "research", "models", "decisions"
|
||||
])
|
||||
assert has_substantive_content, \
|
||||
f"Should extract substantive AI research content. Facts: {[f.fact for f in facts]}"
|
||||
# Judge: substantive content must be extracted regardless of paraphrasing.
|
||||
# "Alignment work" or "machine-learning transparency" satisfy the criterion
|
||||
# the substring check used to enforce as "interpretability/ai/safety".
|
||||
await assert_meets_criteria(
|
||||
response=all_facts_text,
|
||||
criteria=(
|
||||
"The extracted facts cover the substantive AI/ML research content from "
|
||||
"the podcast — Marcus's work on interpretability, his motivation around "
|
||||
"AI safety, or the goal of understanding how models make decisions before "
|
||||
"trusting them in critical applications."
|
||||
),
|
||||
context=(
|
||||
"The transcript wraps substantive AI-research discussion in podcast "
|
||||
"intro/outro meta-commentary (subscribe, like, follow). The substantive "
|
||||
"content is Marcus's interpretability research and AI safety motivation."
|
||||
),
|
||||
msg=f"Should extract substantive AI research content. Facts: {[f.fact for f in facts]}",
|
||||
)
|
||||
|
||||
return # Test passed
|
||||
|
||||
@@ -976,5 +1063,3 @@ so the algorithm learns to box out. See you next week!
|
||||
continue
|
||||
else:
|
||||
raise e
|
||||
|
||||
|
||||
|
||||
@@ -13,6 +13,7 @@ import pytest
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.timeout(1200)
|
||||
async def test_high_fanout_entity_returns_results(memory, request_context):
|
||||
"""
|
||||
A high-fanout entity (appearing in many facts) should still produce
|
||||
@@ -145,6 +146,7 @@ async def test_entity_expansion_timeout_fallback(memory, request_context):
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.timeout(1200)
|
||||
async def test_per_entity_limit_caps_expansion(memory, request_context):
|
||||
"""
|
||||
With graph_per_entity_limit set to a small value, entity expansion should
|
||||
|
||||
@@ -0,0 +1,562 @@
|
||||
"""Tests for async graph maintenance after delete.
|
||||
|
||||
These tests bypass the LLM-backed retain pipeline by inserting memory_units,
|
||||
memory_links, entities, and unit_entities directly. That gives precise
|
||||
control over the graph state so we can assert exact behaviour after a
|
||||
delete + drain.
|
||||
|
||||
The fixture's task backend is ``SyncTaskBackend`` (see conftest), so
|
||||
``submit_async_graph_maintenance`` runs the worker inline — no polling needed.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import uuid
|
||||
from datetime import UTC, datetime, timedelta
|
||||
|
||||
import pytest
|
||||
|
||||
from hindsight_api import RequestContext
|
||||
from hindsight_api.engine.graph_maintenance import (
|
||||
MAX_SEMANTIC_LINKS_PER_UNIT,
|
||||
MAX_TEMPORAL_LINKS_PER_UNIT,
|
||||
enqueue_relink_victims,
|
||||
run_graph_maintenance_job,
|
||||
)
|
||||
from hindsight_api.engine.memory_engine import MemoryEngine
|
||||
|
||||
|
||||
async def _ensure_bank(memory: MemoryEngine, bank_id: str, request_context: RequestContext) -> None:
|
||||
await memory.get_bank_profile(bank_id=bank_id, request_context=request_context)
|
||||
|
||||
|
||||
async def _insert_unit(
|
||||
conn,
|
||||
bank_id: str,
|
||||
text: str,
|
||||
event_date: datetime | None = None,
|
||||
fact_type: str = "experience",
|
||||
) -> uuid.UUID:
|
||||
"""Insert a memory unit directly. Skips embedding (NULL is fine for temporal tests)."""
|
||||
mem_id = uuid.uuid4()
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO memory_units (id, bank_id, text, fact_type, event_date, created_at, updated_at)
|
||||
VALUES ($1, $2, $3, $4, $5, NOW(), NOW())
|
||||
""",
|
||||
mem_id,
|
||||
bank_id,
|
||||
text,
|
||||
fact_type,
|
||||
event_date or datetime.now(UTC),
|
||||
)
|
||||
return mem_id
|
||||
|
||||
|
||||
async def _insert_link(
|
||||
conn,
|
||||
bank_id: str,
|
||||
from_id: uuid.UUID,
|
||||
to_id: uuid.UUID,
|
||||
link_type: str = "temporal",
|
||||
weight: float = 0.5,
|
||||
) -> None:
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO memory_links (from_unit_id, to_unit_id, link_type, weight, bank_id)
|
||||
VALUES ($1, $2, $3, $4, $5)
|
||||
""",
|
||||
from_id,
|
||||
to_id,
|
||||
link_type,
|
||||
weight,
|
||||
bank_id,
|
||||
)
|
||||
|
||||
|
||||
async def _insert_entity(conn, bank_id: str, name: str) -> uuid.UUID:
|
||||
"""Insert an entity row directly. Returns its UUID."""
|
||||
entity_id = uuid.uuid4()
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO entities (id, bank_id, canonical_name, first_seen, last_seen, mention_count)
|
||||
VALUES ($1, $2, $3, NOW(), NOW(), 1)
|
||||
""",
|
||||
entity_id,
|
||||
bank_id,
|
||||
name,
|
||||
)
|
||||
return entity_id
|
||||
|
||||
|
||||
async def _link_unit_entity(conn, unit_id: uuid.UUID, entity_id: uuid.UUID) -> None:
|
||||
await conn.execute(
|
||||
"INSERT INTO unit_entities (unit_id, entity_id) VALUES ($1, $2)",
|
||||
unit_id,
|
||||
entity_id,
|
||||
)
|
||||
|
||||
|
||||
async def _insert_cooccurrence(conn, entity_a: uuid.UUID, entity_b: uuid.UUID, count: int = 1) -> None:
|
||||
# entity_cooccurrence_order_check enforces entity_id_1 < entity_id_2 (canonical
|
||||
# ordering avoids storing (A,B) and (B,A) as two rows). Sort before insert so
|
||||
# callers don't have to care about argument order. Python uuid.UUID compares
|
||||
# by .int, matching PostgreSQL's binary uuid ordering.
|
||||
first, second = sorted([entity_a, entity_b])
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO entity_cooccurrences (entity_id_1, entity_id_2, cooccurrence_count, last_cooccurred)
|
||||
VALUES ($1, $2, $3, NOW())
|
||||
""",
|
||||
first,
|
||||
second,
|
||||
count,
|
||||
)
|
||||
|
||||
|
||||
async def _insert_document(conn, bank_id: str, doc_id: str) -> None:
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO documents (id, bank_id, original_text, content_hash)
|
||||
VALUES ($1, $2, $3, $4)
|
||||
""",
|
||||
doc_id,
|
||||
bank_id,
|
||||
f"text-for-{doc_id}",
|
||||
doc_id,
|
||||
)
|
||||
|
||||
|
||||
async def _attach_unit_to_doc(conn, unit_id: uuid.UUID, doc_id: str) -> None:
|
||||
await conn.execute("UPDATE memory_units SET document_id = $1 WHERE id = $2", doc_id, unit_id)
|
||||
|
||||
|
||||
async def _queue_unit_ids(conn, bank_id: str) -> list[str]:
|
||||
rows = await conn.fetch(
|
||||
"SELECT unit_id FROM graph_maintenance_queue WHERE bank_id = $1 ORDER BY unit_id",
|
||||
bank_id,
|
||||
)
|
||||
return [str(r["unit_id"]) for r in rows]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# enqueue_relink_victims
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestEnqueueRelinkVictims:
|
||||
@pytest.mark.asyncio
|
||||
async def test_enqueues_units_with_outgoing_link_to_deleted(
|
||||
self, memory: MemoryEngine, request_context: RequestContext
|
||||
):
|
||||
bank_id = f"test-gm-enq-{uuid.uuid4().hex[:8]}"
|
||||
await _ensure_bank(memory, bank_id, request_context)
|
||||
|
||||
pool = await memory._get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
doomed = await _insert_unit(conn, bank_id, "doomed")
|
||||
survivor = await _insert_unit(conn, bank_id, "survivor")
|
||||
# survivor → doomed (temporal). When doomed dies, survivor needs top-up.
|
||||
await _insert_link(conn, bank_id, survivor, doomed, "temporal")
|
||||
|
||||
backend = await memory._get_backend()
|
||||
async with conn.transaction():
|
||||
count = await enqueue_relink_victims(conn, bank_id, [str(doomed)], ops=backend.ops)
|
||||
|
||||
assert count == 1
|
||||
assert await _queue_unit_ids(conn, bank_id) == [str(survivor)]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_excludes_deleted_units_themselves(
|
||||
self, memory: MemoryEngine, request_context: RequestContext
|
||||
):
|
||||
"""A unit being deleted that linked TO another deleted unit must not enqueue itself."""
|
||||
bank_id = f"test-gm-self-{uuid.uuid4().hex[:8]}"
|
||||
await _ensure_bank(memory, bank_id, request_context)
|
||||
|
||||
pool = await memory._get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
a = await _insert_unit(conn, bank_id, "a")
|
||||
b = await _insert_unit(conn, bank_id, "b")
|
||||
await _insert_link(conn, bank_id, a, b, "temporal")
|
||||
await _insert_link(conn, bank_id, b, a, "temporal")
|
||||
|
||||
backend = await memory._get_backend()
|
||||
async with conn.transaction():
|
||||
# Both a and b are being deleted — neither should be enqueued.
|
||||
count = await enqueue_relink_victims(conn, bank_id, [str(a), str(b)], ops=backend.ops)
|
||||
|
||||
assert count == 0
|
||||
assert await _queue_unit_ids(conn, bank_id) == []
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_skips_entity_links(self, memory: MemoryEngine, request_context: RequestContext):
|
||||
"""Entity links are being removed from the product — we don't enqueue for them."""
|
||||
bank_id = f"test-gm-ent-{uuid.uuid4().hex[:8]}"
|
||||
await _ensure_bank(memory, bank_id, request_context)
|
||||
|
||||
pool = await memory._get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
doomed = await _insert_unit(conn, bank_id, "doomed")
|
||||
survivor = await _insert_unit(conn, bank_id, "survivor")
|
||||
# Only an entity link — should NOT trigger enqueue.
|
||||
await _insert_link(conn, bank_id, survivor, doomed, "entity")
|
||||
|
||||
backend = await memory._get_backend()
|
||||
async with conn.transaction():
|
||||
count = await enqueue_relink_victims(conn, bank_id, [str(doomed)], ops=backend.ops)
|
||||
|
||||
assert count == 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_dedupes_via_on_conflict(self, memory: MemoryEngine, request_context: RequestContext):
|
||||
bank_id = f"test-gm-dup-{uuid.uuid4().hex[:8]}"
|
||||
await _ensure_bank(memory, bank_id, request_context)
|
||||
|
||||
pool = await memory._get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
doomed1 = await _insert_unit(conn, bank_id, "doomed1")
|
||||
doomed2 = await _insert_unit(conn, bank_id, "doomed2")
|
||||
survivor = await _insert_unit(conn, bank_id, "survivor")
|
||||
# Same survivor linked to two different doomed units across two
|
||||
# logical delete batches — should land in the queue only once.
|
||||
await _insert_link(conn, bank_id, survivor, doomed1, "temporal")
|
||||
await _insert_link(conn, bank_id, survivor, doomed2, "semantic")
|
||||
|
||||
backend = await memory._get_backend()
|
||||
async with conn.transaction():
|
||||
await enqueue_relink_victims(conn, bank_id, [str(doomed1)], ops=backend.ops)
|
||||
await enqueue_relink_victims(conn, bank_id, [str(doomed2)], ops=backend.ops)
|
||||
|
||||
assert await _queue_unit_ids(conn, bank_id) == [str(survivor)]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# delete_document hook
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestDeleteDocumentEnqueue:
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_document_enqueues_cross_doc_victims(
|
||||
self, memory: MemoryEngine, request_context: RequestContext
|
||||
):
|
||||
bank_id = f"test-gm-doc-{uuid.uuid4().hex[:8]}"
|
||||
await _ensure_bank(memory, bank_id, request_context)
|
||||
|
||||
pool = await memory._get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
await _insert_document(conn, bank_id, "doc-A")
|
||||
await _insert_document(conn, bank_id, "doc-B")
|
||||
doomed = await _insert_unit(conn, bank_id, "in doc A")
|
||||
survivor = await _insert_unit(conn, bank_id, "in doc B")
|
||||
await _attach_unit_to_doc(conn, doomed, "doc-A")
|
||||
await _attach_unit_to_doc(conn, survivor, "doc-B")
|
||||
await _insert_link(conn, bank_id, survivor, doomed, "temporal")
|
||||
|
||||
await memory.delete_document("doc-A", bank_id, request_context=request_context)
|
||||
|
||||
async with pool.acquire() as conn:
|
||||
# The synchronous task backend means the worker already drained the
|
||||
# queue before delete_document returned — assert end-state, not the
|
||||
# intermediate enqueue. Queue should be empty.
|
||||
assert await _queue_unit_ids(conn, bank_id) == []
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Relink pass (Pass 1)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestRelinkPass:
|
||||
@pytest.mark.asyncio
|
||||
async def test_drains_empty_queue_cleanly(
|
||||
self, memory: MemoryEngine, request_context: RequestContext
|
||||
):
|
||||
bank_id = f"test-gm-empty-{uuid.uuid4().hex[:8]}"
|
||||
await _ensure_bank(memory, bank_id, request_context)
|
||||
|
||||
result = await run_graph_maintenance_job(memory, bank_id, request_context)
|
||||
assert result == {
|
||||
"relink_units_processed": 0,
|
||||
"relink_links_added": 0,
|
||||
"orphan_entities_pruned": 0,
|
||||
"stale_cooccurrences_pruned": 0,
|
||||
}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_skips_missing_unit_silently(
|
||||
self, memory: MemoryEngine, request_context: RequestContext
|
||||
):
|
||||
"""Unit deleted between enqueue and drain: worker dequeues and no-ops."""
|
||||
bank_id = f"test-gm-miss-{uuid.uuid4().hex[:8]}"
|
||||
await _ensure_bank(memory, bank_id, request_context)
|
||||
|
||||
pool = await memory._get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
# Enqueue a unit_id that doesn't exist in memory_units.
|
||||
await conn.execute(
|
||||
"INSERT INTO graph_maintenance_queue (bank_id, unit_id) VALUES ($1, $2)",
|
||||
bank_id,
|
||||
uuid.uuid4(),
|
||||
)
|
||||
|
||||
result = await run_graph_maintenance_job(memory, bank_id, request_context)
|
||||
assert result["relink_units_processed"] == 1
|
||||
assert result["relink_links_added"] == 0
|
||||
|
||||
async with pool.acquire() as conn:
|
||||
assert await _queue_unit_ids(conn, bank_id) == []
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_tops_up_temporal_when_under_cap(
|
||||
self, memory: MemoryEngine, request_context: RequestContext
|
||||
):
|
||||
"""A victim under the temporal cap gets new outgoing links to neighbours
|
||||
that were never linked at retain time."""
|
||||
bank_id = f"test-gm-topup-{uuid.uuid4().hex[:8]}"
|
||||
await _ensure_bank(memory, bank_id, request_context)
|
||||
|
||||
# Build: one victim at t=0, 2 already-linked neighbours, and 5 unlinked
|
||||
# neighbours all within the 24h window. After top-up the victim should
|
||||
# have outgoing temporal links to all 7.
|
||||
pool = await memory._get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
base = datetime.now(UTC).replace(microsecond=0)
|
||||
victim = await _insert_unit(conn, bank_id, "victim", event_date=base)
|
||||
|
||||
already_linked = [
|
||||
await _insert_unit(conn, bank_id, f"linked-{i}", event_date=base + timedelta(minutes=i + 1))
|
||||
for i in range(2)
|
||||
]
|
||||
for _ in range(5):
|
||||
await _insert_unit(conn, bank_id, "unlinked", event_date=base + timedelta(minutes=30))
|
||||
|
||||
for nbr in already_linked:
|
||||
await _insert_link(conn, bank_id, victim, nbr, "temporal")
|
||||
|
||||
await conn.execute(
|
||||
"INSERT INTO graph_maintenance_queue (bank_id, unit_id) VALUES ($1, $2)",
|
||||
bank_id,
|
||||
victim,
|
||||
)
|
||||
|
||||
result = await run_graph_maintenance_job(memory, bank_id, request_context)
|
||||
assert result["relink_units_processed"] == 1
|
||||
# We probed for up to MAX_TEMPORAL_LINKS_PER_UNIT neighbours; bulk insert
|
||||
# is ON CONFLICT DO NOTHING, so the already-linked 2 are silently
|
||||
# skipped at insert time. The probe still returned them, so
|
||||
# relink_links_added counts what we attempted to insert, not what
|
||||
# actually landed. Verify the end-state via the DB instead.
|
||||
assert result["relink_links_added"] >= 5
|
||||
|
||||
async with pool.acquire() as conn:
|
||||
outgoing = await conn.fetchval(
|
||||
"""
|
||||
SELECT COUNT(*) FROM memory_links
|
||||
WHERE from_unit_id = $1 AND bank_id = $2 AND link_type = 'temporal'
|
||||
""",
|
||||
victim,
|
||||
bank_id,
|
||||
)
|
||||
# 2 originals + 5 new = 7 distinct outgoing temporal links.
|
||||
assert outgoing == 7
|
||||
|
||||
assert await _queue_unit_ids(conn, bank_id) == []
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_topup_when_victim_at_cap(
|
||||
self, memory: MemoryEngine, request_context: RequestContext
|
||||
):
|
||||
"""If the victim already has cap links, probing is skipped."""
|
||||
bank_id = f"test-gm-atcap-{uuid.uuid4().hex[:8]}"
|
||||
await _ensure_bank(memory, bank_id, request_context)
|
||||
|
||||
pool = await memory._get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
base = datetime.now(UTC).replace(microsecond=0)
|
||||
victim = await _insert_unit(conn, bank_id, "victim", event_date=base)
|
||||
|
||||
# Insert exactly cap temporal links from victim, plus extra unlinked
|
||||
# candidates. Probe should be skipped because victim is at cap.
|
||||
for i in range(MAX_TEMPORAL_LINKS_PER_UNIT):
|
||||
nbr = await _insert_unit(conn, bank_id, f"l-{i}", event_date=base + timedelta(minutes=i + 1))
|
||||
await _insert_link(conn, bank_id, victim, nbr, "temporal")
|
||||
|
||||
# Plus extras that would be valid candidates if we DID probe.
|
||||
for i in range(3):
|
||||
await _insert_unit(conn, bank_id, f"x-{i}", event_date=base + timedelta(minutes=i + 100))
|
||||
|
||||
await conn.execute(
|
||||
"INSERT INTO graph_maintenance_queue (bank_id, unit_id) VALUES ($1, $2)",
|
||||
bank_id,
|
||||
victim,
|
||||
)
|
||||
|
||||
result = await run_graph_maintenance_job(memory, bank_id, request_context)
|
||||
assert result["relink_units_processed"] == 1
|
||||
|
||||
async with pool.acquire() as conn:
|
||||
outgoing = await conn.fetchval(
|
||||
"""
|
||||
SELECT COUNT(*) FROM memory_links
|
||||
WHERE from_unit_id = $1 AND link_type = 'temporal'
|
||||
""",
|
||||
victim,
|
||||
)
|
||||
assert outgoing == MAX_TEMPORAL_LINKS_PER_UNIT
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Orphan entity prune (Pass 2)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestOrphanEntityPrune:
|
||||
@pytest.mark.asyncio
|
||||
async def test_prunes_entities_with_no_unit_references(
|
||||
self, memory: MemoryEngine, request_context: RequestContext
|
||||
):
|
||||
"""An entity with zero unit_entities rows is an orphan and should be
|
||||
deleted by the sweep."""
|
||||
bank_id = f"test-gm-orphan-{uuid.uuid4().hex[:8]}"
|
||||
await _ensure_bank(memory, bank_id, request_context)
|
||||
|
||||
pool = await memory._get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
referenced = await _insert_entity(conn, bank_id, "referenced")
|
||||
orphan_a = await _insert_entity(conn, bank_id, "orphan_a")
|
||||
orphan_b = await _insert_entity(conn, bank_id, "orphan_b")
|
||||
|
||||
unit = await _insert_unit(conn, bank_id, "with-entity")
|
||||
await _link_unit_entity(conn, unit, referenced)
|
||||
|
||||
result = await run_graph_maintenance_job(memory, bank_id, request_context)
|
||||
assert result["orphan_entities_pruned"] == 2
|
||||
|
||||
async with pool.acquire() as conn:
|
||||
survivors = await conn.fetch(
|
||||
"SELECT id FROM entities WHERE bank_id = $1 ORDER BY id", bank_id
|
||||
)
|
||||
survivor_ids = {str(r["id"]) for r in survivors}
|
||||
assert survivor_ids == {str(referenced)}
|
||||
# Confirm orphans are gone.
|
||||
for orphan in (orphan_a, orphan_b):
|
||||
assert orphan not in survivor_ids
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_does_not_touch_other_banks(
|
||||
self, memory: MemoryEngine, request_context: RequestContext
|
||||
):
|
||||
"""The sweep is scoped by bank — orphan entities in OTHER banks
|
||||
must not be touched."""
|
||||
bank_a = f"test-gm-scopea-{uuid.uuid4().hex[:8]}"
|
||||
bank_b = f"test-gm-scopeb-{uuid.uuid4().hex[:8]}"
|
||||
await _ensure_bank(memory, bank_a, request_context)
|
||||
await _ensure_bank(memory, bank_b, request_context)
|
||||
|
||||
pool = await memory._get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
orphan_in_a = await _insert_entity(conn, bank_a, "orphan_a")
|
||||
orphan_in_b = await _insert_entity(conn, bank_b, "orphan_b")
|
||||
|
||||
await run_graph_maintenance_job(memory, bank_a, request_context)
|
||||
|
||||
async with pool.acquire() as conn:
|
||||
# b's orphan must still exist — the sweep was scoped to a.
|
||||
still_in_b = await conn.fetchval("SELECT 1 FROM entities WHERE id = $1", orphan_in_b)
|
||||
assert still_in_b == 1
|
||||
# a's orphan is gone.
|
||||
still_in_a = await conn.fetchval("SELECT 1 FROM entities WHERE id = $1", orphan_in_a)
|
||||
assert still_in_a is None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Stale cooccurrence prune (Pass 3)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestStaleCooccurrencePrune:
|
||||
@pytest.mark.asyncio
|
||||
async def test_prunes_cooccurrence_with_no_shared_unit(
|
||||
self, memory: MemoryEngine, request_context: RequestContext
|
||||
):
|
||||
"""Both entities still exist but no unit references both of them — the
|
||||
cooccurrence row is stale and should be pruned."""
|
||||
bank_id = f"test-gm-cocc-{uuid.uuid4().hex[:8]}"
|
||||
await _ensure_bank(memory, bank_id, request_context)
|
||||
|
||||
pool = await memory._get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
ent_a = await _insert_entity(conn, bank_id, "alice")
|
||||
ent_b = await _insert_entity(conn, bank_id, "bob")
|
||||
# Cooccurrence row records that A and B were observed together.
|
||||
await _insert_cooccurrence(conn, ent_a, ent_b, count=5)
|
||||
# Both entities still have references — but to DIFFERENT units, so
|
||||
# no current unit witnesses both A and B together.
|
||||
unit_a = await _insert_unit(conn, bank_id, "with_a")
|
||||
unit_b = await _insert_unit(conn, bank_id, "with_b")
|
||||
await _link_unit_entity(conn, unit_a, ent_a)
|
||||
await _link_unit_entity(conn, unit_b, ent_b)
|
||||
|
||||
result = await run_graph_maintenance_job(memory, bank_id, request_context)
|
||||
assert result["stale_cooccurrences_pruned"] == 1
|
||||
# Both entities still exist — they weren't orphans.
|
||||
assert result["orphan_entities_pruned"] == 0
|
||||
|
||||
async with pool.acquire() as conn:
|
||||
# Match canonical ordering enforced by entity_cooccurrence_order_check.
|
||||
first, second = sorted([ent_a, ent_b])
|
||||
remaining = await conn.fetchval(
|
||||
"SELECT COUNT(*) FROM entity_cooccurrences WHERE entity_id_1 = $1 AND entity_id_2 = $2",
|
||||
first,
|
||||
second,
|
||||
)
|
||||
assert remaining == 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_keeps_cooccurrence_with_shared_unit(
|
||||
self, memory: MemoryEngine, request_context: RequestContext
|
||||
):
|
||||
"""If at least one unit still references both entities, the cooccurrence
|
||||
row stays."""
|
||||
bank_id = f"test-gm-keep-{uuid.uuid4().hex[:8]}"
|
||||
await _ensure_bank(memory, bank_id, request_context)
|
||||
|
||||
pool = await memory._get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
ent_a = await _insert_entity(conn, bank_id, "alice")
|
||||
ent_b = await _insert_entity(conn, bank_id, "bob")
|
||||
await _insert_cooccurrence(conn, ent_a, ent_b, count=5)
|
||||
# A unit references both — cooccurrence is still grounded.
|
||||
unit = await _insert_unit(conn, bank_id, "alice-and-bob")
|
||||
await _link_unit_entity(conn, unit, ent_a)
|
||||
await _link_unit_entity(conn, unit, ent_b)
|
||||
|
||||
result = await run_graph_maintenance_job(memory, bank_id, request_context)
|
||||
assert result["stale_cooccurrences_pruned"] == 0
|
||||
|
||||
async with pool.acquire() as conn:
|
||||
# Match canonical ordering enforced by entity_cooccurrence_order_check.
|
||||
first, second = sorted([ent_a, ent_b])
|
||||
still_there = await conn.fetchval(
|
||||
"SELECT cooccurrence_count FROM entity_cooccurrences WHERE entity_id_1 = $1 AND entity_id_2 = $2",
|
||||
first,
|
||||
second,
|
||||
)
|
||||
assert still_there == 5
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Sanity check on cap values
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_caps_match_retain_defaults():
|
||||
"""If retain bumps its caps but graph_maintenance stays put, top-up will
|
||||
silently never reach the retain ceiling — the asserts here exist so a
|
||||
future cap change forces a paired update."""
|
||||
from hindsight_api.engine.retain.link_utils import MAX_TEMPORAL_LINKS_PER_UNIT as RETAIN_TEMPORAL
|
||||
|
||||
assert MAX_TEMPORAL_LINKS_PER_UNIT == RETAIN_TEMPORAL
|
||||
assert MAX_SEMANTIC_LINKS_PER_UNIT == 50 # mirrors compute_semantic_links_ann's top_k default
|
||||
@@ -5,6 +5,7 @@ Tests config resolution hierarchy (global → tenant → bank),
|
||||
key normalization, API endpoints, validation, and caching.
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
|
||||
import pytest
|
||||
@@ -37,6 +38,33 @@ class MockTenantExtension(TenantExtension):
|
||||
return self.tenant_config
|
||||
|
||||
|
||||
class FakeBankConfigBackend:
|
||||
"""Minimal backend for ConfigResolver bank-config tests."""
|
||||
|
||||
def __init__(self):
|
||||
self.config: dict[str, object] = {}
|
||||
|
||||
def acquire(self):
|
||||
return FakeBankConfigConnection(self)
|
||||
|
||||
|
||||
class FakeBankConfigConnection:
|
||||
def __init__(self, backend: FakeBankConfigBackend):
|
||||
self.backend = backend
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb):
|
||||
return None
|
||||
|
||||
async def fetchrow(self, query, bank_id):
|
||||
return {"config": self.backend.config}
|
||||
|
||||
async def execute(self, query, updates_json, bank_id):
|
||||
self.backend.config.update(json.loads(updates_json))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_config_key_normalization():
|
||||
"""Test that env var keys are normalized to Python field names."""
|
||||
@@ -96,9 +124,10 @@ async def test_hierarchical_fields_categorization():
|
||||
assert "llm_gemini_safety_settings" in configurable
|
||||
assert "mcp_enabled_tools" in configurable
|
||||
assert "retain_chunk_batch_size" in configurable
|
||||
assert "enable_auto_consolidation" in configurable
|
||||
|
||||
# Verify count is correct
|
||||
assert len(configurable) == 35
|
||||
assert len(configurable) == 36
|
||||
|
||||
# Verify credential fields (NEVER exposed)
|
||||
assert "llm_api_key" in credentials
|
||||
@@ -169,6 +198,46 @@ async def test_config_hierarchy_resolution(memory, request_context):
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bank_config_null_consolidation_overrides_use_server_defaults():
|
||||
"""JSON null bank overrides should behave like Server Default.
|
||||
|
||||
Regression test for #1619: the dashboard can send null for observation
|
||||
config fields. Those nulls must not flow into consolidation as None.
|
||||
"""
|
||||
bank_id = "test-null-consolidation-config-bank"
|
||||
fields = (
|
||||
"consolidation_llm_batch_size",
|
||||
"consolidation_source_facts_max_tokens",
|
||||
"consolidation_source_facts_max_tokens_per_observation",
|
||||
"max_observations_per_scope",
|
||||
)
|
||||
resolver = ConfigResolver(backend=FakeBankConfigBackend())
|
||||
explicit_overrides = {
|
||||
"consolidation_llm_batch_size": 7,
|
||||
"consolidation_source_facts_max_tokens": 2048,
|
||||
"consolidation_source_facts_max_tokens_per_observation": 256,
|
||||
"max_observations_per_scope": 3,
|
||||
}
|
||||
|
||||
await resolver.update_bank_config(bank_id, explicit_overrides)
|
||||
config = await resolver.resolve_full_config(bank_id)
|
||||
for field_name, expected in explicit_overrides.items():
|
||||
assert getattr(config, field_name) == expected
|
||||
|
||||
await resolver.update_bank_config(bank_id, {field_name: None for field_name in fields})
|
||||
|
||||
resolved_config = await resolver.resolve_full_config(bank_id)
|
||||
global_config = resolver._global_config
|
||||
for field_name in fields:
|
||||
assert getattr(resolved_config, field_name) == getattr(global_config, field_name)
|
||||
assert getattr(resolved_config, field_name) is not None
|
||||
|
||||
bank_overrides = await resolver._load_bank_config(bank_id)
|
||||
for field_name in fields:
|
||||
assert field_name not in bank_overrides
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_config_validation_rejects_static_fields(memory, request_context):
|
||||
"""Test that attempting to override static fields raises ValueError."""
|
||||
|
||||
@@ -7,6 +7,7 @@ how observations track the evolving state over time, with full prompt debugging.
|
||||
import json
|
||||
import uuid
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
@@ -14,6 +15,7 @@ import pytest
|
||||
from hindsight_api.config import _get_raw_config
|
||||
from hindsight_api.engine.consolidation import consolidator as consolidator_mod
|
||||
from hindsight_api.engine.memory_engine import MemoryEngine
|
||||
from tests.llm_judge import assert_meets_criteria
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
@@ -95,18 +97,20 @@ async def _instrumented_consolidate(
|
||||
max_observations_per_scope=max_observations_per_scope,
|
||||
)
|
||||
|
||||
_debug_log.append(_ConsolidationDebugEntry(
|
||||
facts=facts_lines,
|
||||
observations_text=observations_text,
|
||||
response=_ConsolidationResponse(
|
||||
creates=[_ActionLog(text=c.text, source_fact_ids=c.source_fact_ids) for c in result.creates],
|
||||
updates=[
|
||||
_ActionLog(text=u.text, observation_id=u.observation_id, source_fact_ids=u.source_fact_ids)
|
||||
for u in result.updates
|
||||
],
|
||||
deletes=[_ActionLog(text="", observation_id=d.observation_id) for d in result.deletes],
|
||||
),
|
||||
))
|
||||
_debug_log.append(
|
||||
_ConsolidationDebugEntry(
|
||||
facts=facts_lines,
|
||||
observations_text=observations_text,
|
||||
response=_ConsolidationResponse(
|
||||
creates=[_ActionLog(text=c.text, source_fact_ids=c.source_fact_ids) for c in result.creates],
|
||||
updates=[
|
||||
_ActionLog(text=u.text, observation_id=u.observation_id, source_fact_ids=u.source_fact_ids)
|
||||
for u in result.updates
|
||||
],
|
||||
deletes=[_ActionLog(text="", observation_id=d.observation_id) for d in result.deletes],
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
return result
|
||||
|
||||
@@ -133,11 +137,11 @@ def _print_consolidation_debug(entry: _ConsolidationDebugEntry, index: int) -> N
|
||||
print("\n LLM RESPONSE:")
|
||||
if resp.creates:
|
||||
for c in resp.creates:
|
||||
print(f" CREATE: \"{c.text}\" (from facts: {[fid[:8] + '..' for fid in c.source_fact_ids]})")
|
||||
print(f' CREATE: "{c.text}" (from facts: {[fid[:8] + ".." for fid in c.source_fact_ids]})')
|
||||
if resp.updates:
|
||||
for u in resp.updates:
|
||||
print(
|
||||
f" UPDATE [{u.observation_id[:8]}..]: \"{u.text}\""
|
||||
f' UPDATE [{u.observation_id[:8]}..]: "{u.text}"'
|
||||
f" (from facts: {[fid[:8] + '..' for fid in u.source_fact_ids]})"
|
||||
)
|
||||
if resp.deletes:
|
||||
@@ -162,8 +166,10 @@ def _parse_history(hist: Any) -> list[str]:
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.flaky(reruns=2, reruns_delay=5)
|
||||
async def test_horse_farm_observation_history(memory: MemoryEngine, request_context: Any) -> None:
|
||||
@pytest.mark.hs_llm_core
|
||||
async def test_horse_farm_observation_history(memory_real_llm: MemoryEngine, request_context: Any) -> None:
|
||||
"""Retain a sequence of horse facts and inspect how observations evolve."""
|
||||
memory = memory_real_llm
|
||||
bank_id = f"test-horses-{uuid.uuid4().hex[:8]}"
|
||||
|
||||
await memory.get_bank_profile(bank_id=bank_id, request_context=request_context)
|
||||
@@ -182,6 +188,16 @@ async def test_horse_farm_observation_history(memory: MemoryEngine, request_cont
|
||||
"I am sad to report that Shadow has died.",
|
||||
]
|
||||
|
||||
# Space mentioned_at one week apart per retain so the temporal supersession
|
||||
# rule the reflect prompt teaches the LLM has meaningful signal to work
|
||||
# with. Without an explicit event_date, retains land at utcnow() and end
|
||||
# up 2-5 seconds apart in wall clock time — close enough that the LLM
|
||||
# can't reliably rank "5 horses" (later) over "1 horse" (earlier) because
|
||||
# the gap looks like noise. One-week spacing models a user narrating their
|
||||
# farm over time.
|
||||
base_time = datetime(2025, 1, 1, 12, 0, 0, tzinfo=timezone.utc)
|
||||
event_dates = [base_time + timedelta(weeks=i) for i in range(len(messages))]
|
||||
|
||||
# Monkey-patch to intercept consolidation LLM calls
|
||||
_original_consolidate = consolidator_mod._consolidate_batch_with_llm
|
||||
|
||||
@@ -193,15 +209,16 @@ async def test_horse_farm_observation_history(memory: MemoryEngine, request_cont
|
||||
|
||||
try:
|
||||
for i, content in enumerate(messages):
|
||||
print(f"\n{'='*80}")
|
||||
print(f"RETAIN #{i+1}: {content}")
|
||||
print(f"{'='*80}")
|
||||
print(f"\n{'=' * 80}")
|
||||
print(f"RETAIN #{i + 1} ({event_dates[i].date()}): {content}")
|
||||
print(f"{'=' * 80}")
|
||||
|
||||
log_start = len(_debug_log)
|
||||
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content=content,
|
||||
event_date=event_dates[i],
|
||||
request_context=request_context,
|
||||
)
|
||||
await memory.wait_for_background_tasks()
|
||||
@@ -231,9 +248,9 @@ async def test_horse_farm_observation_history(memory: MemoryEngine, request_cont
|
||||
consolidator_mod._consolidate_batch_with_llm = _original_consolidate
|
||||
|
||||
# Final summary
|
||||
print(f"\n{'='*80}")
|
||||
print(f"\n{'=' * 80}")
|
||||
print("FINAL STATE")
|
||||
print(f"{'='*80}")
|
||||
print(f"{'=' * 80}")
|
||||
pool = await memory._get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
observations = await conn.fetch(
|
||||
@@ -256,9 +273,9 @@ async def test_horse_farm_observation_history(memory: MemoryEngine, request_cont
|
||||
print(f" - [proof={obs['proof_count']}] {obs['text']}")
|
||||
|
||||
# Create a mental model to synthesize the observations
|
||||
print(f"\n{'='*80}")
|
||||
print(f"\n{'=' * 80}")
|
||||
print("MENTAL MODEL")
|
||||
print(f"{'='*80}")
|
||||
print(f"{'=' * 80}")
|
||||
|
||||
# Patch reflect _execute_tool to log tool inputs/outputs
|
||||
from hindsight_api.engine.reflect import agent as reflect_agent_mod
|
||||
@@ -271,7 +288,9 @@ async def test_horse_farm_observation_history(memory: MemoryEngine, request_cont
|
||||
print(f"\n [REFLECT TOOL] {normalized}(args={args})")
|
||||
if isinstance(result, dict):
|
||||
if "observations" in result:
|
||||
print(f" Observations returned ({result.get('count', '?')}, freshness={result.get('freshness', '?')}):")
|
||||
print(
|
||||
f" Observations returned ({result.get('count', '?')}, freshness={result.get('freshness', '?')}):"
|
||||
)
|
||||
for obs in result.get("observations", []):
|
||||
print(f" - [proof={obs.get('proof_count', '?')}] {obs.get('text', '?')}")
|
||||
if "memories" in result:
|
||||
@@ -335,18 +354,24 @@ async def test_horse_farm_observation_history(memory: MemoryEngine, request_cont
|
||||
continue
|
||||
print(f" - [{item.get('fact_type', '?')}] {item.get('text', '?')}")
|
||||
|
||||
# Verify the mental model captures key facts
|
||||
content_lower = content.lower()
|
||||
|
||||
for name in ["daisy", "buttercup", "midnight", "shadow", "twister"]:
|
||||
assert name in content_lower, f"Mental model should mention {name}. Got:\n{content}"
|
||||
|
||||
assert "sold" in content_lower or "sale" in content_lower, (
|
||||
f"Mental model should mention Buttercup was sold. Got:\n{content}"
|
||||
)
|
||||
|
||||
assert "died" in content_lower or "passed" in content_lower or "death" in content_lower, (
|
||||
f"Mental model should mention Shadow's death. Got:\n{content}"
|
||||
# Verify the mental model captures key facts via LLM judge. The synthesis
|
||||
# step occasionally drops one name (typically Daisy, only mentioned once
|
||||
# with no follow-up events), so accept ≥4 of 5 names rather than all 5 —
|
||||
# the assertion is whether the pipeline synthesizes the herd story
|
||||
# end-to-end, not perfect recall of every horse.
|
||||
await assert_meets_criteria(
|
||||
response=content,
|
||||
criteria=(
|
||||
"The mental model mentions at least 4 of these 5 horse names: "
|
||||
"Daisy, Buttercup, Midnight, Shadow, Twister. "
|
||||
"Buttercup and Shadow MUST both be named (they are the two horses "
|
||||
"involved in events). It also mentions that Buttercup was sold and "
|
||||
"that Shadow died or passed away."
|
||||
),
|
||||
context=(
|
||||
"Input events: Had 2 horses (Daisy, Buttercup). Sold Buttercup. "
|
||||
"Got more horses (Midnight, Shadow, Twister). Shadow died."
|
||||
),
|
||||
)
|
||||
|
||||
# Cleanup
|
||||
|
||||
@@ -3,16 +3,20 @@ Integration test for the complete Hindsight API.
|
||||
|
||||
Tests all endpoints by starting a FastAPI server and making HTTP requests.
|
||||
"""
|
||||
|
||||
from datetime import datetime
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
import httpx
|
||||
from datetime import datetime
|
||||
|
||||
from hindsight_api.api import create_app
|
||||
from tests.llm_judge import assert_meets_criteria
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def api_client(memory):
|
||||
"""Create an async test client for the FastAPI app."""
|
||||
"""Create an async test client for the FastAPI app (mock LLM)."""
|
||||
# Memory is already initialized by the conftest fixture (with migrations)
|
||||
app = create_app(memory, initialize_memory=False)
|
||||
transport = httpx.ASGITransport(app=app)
|
||||
@@ -20,6 +24,15 @@ async def api_client(memory):
|
||||
yield client
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def api_client_real_llm(memory_real_llm):
|
||||
"""Create an async test client backed by a real LLM provider."""
|
||||
app = create_app(memory_real_llm, 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():
|
||||
"""Provide a unique bank ID for this test run."""
|
||||
@@ -65,10 +78,10 @@ async def test_full_api_workflow(api_client, test_bank_id):
|
||||
"items": [
|
||||
{
|
||||
"content": "Alice is a machine learning researcher at Stanford.",
|
||||
"context": "conversation about team members"
|
||||
"context": "conversation about team members",
|
||||
}
|
||||
]
|
||||
}
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
put_result = response.json()
|
||||
@@ -80,16 +93,13 @@ async def test_full_api_workflow(api_client, test_bank_id):
|
||||
f"/v1/default/banks/{test_bank_id}/memories",
|
||||
json={
|
||||
"items": [
|
||||
{
|
||||
"content": "Bob leads the infrastructure team and loves Kubernetes.",
|
||||
"context": "team introduction"
|
||||
},
|
||||
{"content": "Bob leads the infrastructure team and loves Kubernetes.", "context": "team introduction"},
|
||||
{
|
||||
"content": "Charlie recently joined as a product manager from Google.",
|
||||
"context": "new hire announcement"
|
||||
}
|
||||
"context": "new hire announcement",
|
||||
},
|
||||
]
|
||||
}
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
batch_result = response.json()
|
||||
@@ -103,10 +113,7 @@ async def test_full_api_workflow(api_client, test_bank_id):
|
||||
# Recall memories
|
||||
response = await api_client.post(
|
||||
f"/v1/default/banks/{test_bank_id}/memories/recall",
|
||||
json={
|
||||
"query": "Who works on machine learning?",
|
||||
"thinking_budget": 50
|
||||
}
|
||||
json={"query": "Who works on machine learning?", "thinking_budget": 50},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
search_results = response.json()
|
||||
@@ -127,8 +134,8 @@ async def test_full_api_workflow(api_client, test_bank_id):
|
||||
json={
|
||||
"query": "What do you know about the team members?",
|
||||
"thinking_budget": 30,
|
||||
"context": "This is for a team overview document"
|
||||
}
|
||||
"context": "This is for a team overview document",
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
reflect_result = response.json()
|
||||
@@ -136,9 +143,8 @@ async def test_full_api_workflow(api_client, test_bank_id):
|
||||
assert len(reflect_result["text"]) > 0
|
||||
assert "based_on" in reflect_result
|
||||
|
||||
# Verify the answer mentions team members
|
||||
answer = reflect_result["text"].lower()
|
||||
assert "alice" in answer or "bob" in answer or "charlie" in answer
|
||||
# Verify the reflect endpoint returned a non-trivial response
|
||||
assert len(reflect_result["text"]) > 5, "Reflect should return a substantive response"
|
||||
|
||||
# ================================================================
|
||||
# 5. Visualization & Statistics
|
||||
@@ -167,10 +173,7 @@ async def test_full_api_workflow(api_client, test_bank_id):
|
||||
assert our_bank["last_document_at"] is not None, "last_document_at should be set after retain"
|
||||
|
||||
# List memory units
|
||||
response = await api_client.get(
|
||||
f"/v1/default/banks/{test_bank_id}/memories/list",
|
||||
params={"limit": 10}
|
||||
)
|
||||
response = await api_client.get(f"/v1/default/banks/{test_bank_id}/memories/list", params={"limit": 10})
|
||||
assert response.status_code == 200
|
||||
memory_units = response.json()
|
||||
assert "items" in memory_units
|
||||
@@ -188,10 +191,10 @@ async def test_full_api_workflow(api_client, test_bank_id):
|
||||
{
|
||||
"content": "Project timeline: MVP launch in Q1, Beta in Q2.",
|
||||
"context": "product roadmap",
|
||||
"document_id": "roadmap-2024-q1"
|
||||
"document_id": "roadmap-2024-q1",
|
||||
}
|
||||
]
|
||||
}
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
|
||||
@@ -203,9 +206,7 @@ async def test_full_api_workflow(api_client, test_bank_id):
|
||||
assert len(documents["items"]) > 0
|
||||
|
||||
# Get specific document
|
||||
response = await api_client.get(
|
||||
f"/v1/default/banks/{test_bank_id}/documents/roadmap-2024-q1"
|
||||
)
|
||||
response = await api_client.get(f"/v1/default/banks/{test_bank_id}/documents/roadmap-2024-q1")
|
||||
assert response.status_code == 200
|
||||
doc_info = response.json()
|
||||
assert "id" in doc_info
|
||||
@@ -220,13 +221,7 @@ async def test_full_api_workflow(api_client, test_bank_id):
|
||||
# Update disposition traits
|
||||
response = await api_client.put(
|
||||
f"/v1/default/banks/{test_bank_id}/profile",
|
||||
json={
|
||||
"disposition": {
|
||||
"skepticism": 4,
|
||||
"literalism": 3,
|
||||
"empathy": 4
|
||||
}
|
||||
}
|
||||
json={"disposition": {"skepticism": 4, "literalism": 3, "empathy": 4}},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
|
||||
@@ -272,19 +267,15 @@ async def test_full_api_workflow(api_client, test_bank_id):
|
||||
assert offset_data["items"][0]["id"] != entities_data["items"][0]["id"]
|
||||
|
||||
# Get specific entity if any exist
|
||||
if len(entities_data['items']) > 0:
|
||||
entity_id = entities_data['items'][0]['id']
|
||||
response = await api_client.get(
|
||||
f"/v1/default/banks/{test_bank_id}/entities/{entity_id}"
|
||||
)
|
||||
if len(entities_data["items"]) > 0:
|
||||
entity_id = entities_data["items"][0]["id"]
|
||||
response = await api_client.get(f"/v1/default/banks/{test_bank_id}/entities/{entity_id}")
|
||||
assert response.status_code == 200
|
||||
entity_detail = response.json()
|
||||
assert "id" in entity_detail
|
||||
|
||||
# Test regenerate observations (deprecated - returns 410 Gone)
|
||||
response = await api_client.post(
|
||||
f"/v1/default/banks/{test_bank_id}/entities/{entity_id}/regenerate"
|
||||
)
|
||||
response = await api_client.post(f"/v1/default/banks/{test_bank_id}/entities/{entity_id}/regenerate")
|
||||
assert response.status_code == 410 # Deprecated endpoint
|
||||
|
||||
# Entity co-occurrence graph — shape is stable even when there are no
|
||||
@@ -304,9 +295,7 @@ async def test_full_api_workflow(api_client, test_bank_id):
|
||||
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"
|
||||
)
|
||||
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
|
||||
@@ -350,7 +339,7 @@ async def test_error_handling(api_client):
|
||||
"context": "test"
|
||||
}
|
||||
]
|
||||
}
|
||||
},
|
||||
)
|
||||
assert response.status_code == 422 # Validation error
|
||||
|
||||
@@ -359,15 +348,13 @@ async def test_error_handling(api_client):
|
||||
"/v1/default/banks/error_test/memories/recall",
|
||||
json={
|
||||
"query": "test",
|
||||
"budget": "invalid_budget" # Invalid budget value (should be low/mid/high)
|
||||
}
|
||||
"budget": "invalid_budget", # Invalid budget value (should be low/mid/high)
|
||||
},
|
||||
)
|
||||
assert response.status_code == 422
|
||||
|
||||
# Get non-existent document
|
||||
response = await api_client.get(
|
||||
"/v1/default/banks/nonexistent_bank/documents/fake-doc-id"
|
||||
)
|
||||
response = await api_client.get("/v1/default/banks/nonexistent_bank/documents/fake-doc-id")
|
||||
assert response.status_code == 404
|
||||
|
||||
|
||||
@@ -383,19 +370,11 @@ async def test_concurrent_requests(api_client):
|
||||
"Emily is the CEO of a startup in San Francisco.",
|
||||
"Frank teaches computer science at MIT.",
|
||||
"Grace is a software architect specializing in distributed systems.",
|
||||
"Henry leads the product team at Amazon."
|
||||
"Henry leads the product team at Amazon.",
|
||||
]
|
||||
for fact in test_facts:
|
||||
response = await api_client.post(
|
||||
f"/v1/default/banks/{bank_id}/memories",
|
||||
json={
|
||||
"items": [
|
||||
{
|
||||
"content": fact,
|
||||
"context": "concurrent test"
|
||||
}
|
||||
]
|
||||
}
|
||||
f"/v1/default/banks/{bank_id}/memories", json={"items": [{"content": fact, "context": "concurrent test"}]}
|
||||
)
|
||||
responses.append(response)
|
||||
|
||||
@@ -404,10 +383,7 @@ async def test_concurrent_requests(api_client):
|
||||
assert all(r.json()["success"] for r in responses)
|
||||
|
||||
# Verify all facts stored
|
||||
response = await api_client.get(
|
||||
f"/v1/default/banks/{bank_id}/memories/list",
|
||||
params={"limit": 20}
|
||||
)
|
||||
response = await api_client.get(f"/v1/default/banks/{bank_id}/memories/list", params={"limit": 20})
|
||||
assert response.status_code == 200
|
||||
items = response.json()["items"]
|
||||
assert len(items) >= 5
|
||||
@@ -426,26 +402,22 @@ async def test_document_deletion(api_client):
|
||||
{
|
||||
"content": "The quarterly sales report shows a 25% increase in revenue.",
|
||||
"context": "Q1 financial review",
|
||||
"document_id": "sales-report-q1-2024"
|
||||
"document_id": "sales-report-q1-2024",
|
||||
}
|
||||
]
|
||||
}
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
|
||||
# Verify document exists
|
||||
response = await api_client.get(
|
||||
f"/v1/default/banks/{test_bank_id}/documents/sales-report-q1-2024"
|
||||
)
|
||||
response = await api_client.get(f"/v1/default/banks/{test_bank_id}/documents/sales-report-q1-2024")
|
||||
assert response.status_code == 200
|
||||
doc_info = response.json()
|
||||
initial_units = doc_info["memory_unit_count"]
|
||||
assert initial_units > 0
|
||||
|
||||
# Delete the document
|
||||
response = await api_client.delete(
|
||||
f"/v1/default/banks/{test_bank_id}/documents/sales-report-q1-2024"
|
||||
)
|
||||
response = await api_client.delete(f"/v1/default/banks/{test_bank_id}/documents/sales-report-q1-2024")
|
||||
assert response.status_code == 200
|
||||
delete_result = response.json()
|
||||
assert delete_result["success"] is True
|
||||
@@ -453,9 +425,7 @@ async def test_document_deletion(api_client):
|
||||
assert delete_result["memory_units_deleted"] == initial_units
|
||||
|
||||
# Verify document is gone (should return 404)
|
||||
response = await api_client.get(
|
||||
f"/v1/default/banks/{test_bank_id}/documents/sales-report-q1-2024"
|
||||
)
|
||||
response = await api_client.get(f"/v1/default/banks/{test_bank_id}/documents/sales-report-q1-2024")
|
||||
assert response.status_code == 404
|
||||
|
||||
# Verify document is not in the list
|
||||
@@ -466,9 +436,7 @@ async def test_document_deletion(api_client):
|
||||
assert "sales-report-q1-2024" not in doc_ids
|
||||
|
||||
# Try to delete again (should return 404)
|
||||
response = await api_client.delete(
|
||||
f"/v1/default/banks/{test_bank_id}/documents/sales-report-q1-2024"
|
||||
)
|
||||
response = await api_client.delete(f"/v1/default/banks/{test_bank_id}/documents/sales-report-q1-2024")
|
||||
assert response.status_code == 404
|
||||
|
||||
|
||||
@@ -496,10 +464,10 @@ async def test_document_deletion_with_slashes_in_id(api_client):
|
||||
{
|
||||
"content": "The Q1 2024 report shows significant growth in user engagement.",
|
||||
"context": "quarterly report",
|
||||
"document_id": document_id_with_slash
|
||||
"document_id": document_id_with_slash,
|
||||
}
|
||||
]
|
||||
}
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200, f"Failed to create document: {response.text}"
|
||||
|
||||
@@ -512,12 +480,9 @@ async def test_document_deletion_with_slashes_in_id(api_client):
|
||||
|
||||
# 3. Delete the document (slashes in document_id should work with :path converter)
|
||||
encoded_doc_id = urllib.parse.quote(document_id_with_slash, safe="")
|
||||
response = await api_client.delete(
|
||||
f"/v1/default/banks/{test_bank_id}/documents/{encoded_doc_id}"
|
||||
)
|
||||
response = await api_client.delete(f"/v1/default/banks/{test_bank_id}/documents/{encoded_doc_id}")
|
||||
assert response.status_code == 200, (
|
||||
f"Failed to delete document with slashes in ID. "
|
||||
f"Status: {response.status_code}, Response: {response.text}"
|
||||
f"Failed to delete document with slashes in ID. Status: {response.status_code}, Response: {response.text}"
|
||||
)
|
||||
|
||||
# Verify document is deleted
|
||||
@@ -715,10 +680,10 @@ async def test_async_retain(api_client):
|
||||
"items": [
|
||||
{
|
||||
"content": "Alice is a senior engineer at TechCorp. She has been working on the authentication system for 5 years.",
|
||||
"context": "team introduction"
|
||||
"context": "team introduction",
|
||||
}
|
||||
]
|
||||
}
|
||||
],
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
result = response.json()
|
||||
@@ -740,10 +705,7 @@ async def test_async_retain(api_client):
|
||||
|
||||
while elapsed < max_wait_seconds:
|
||||
# Check if memories are stored
|
||||
response = await api_client.get(
|
||||
f"/v1/default/banks/{test_bank_id}/memories/list",
|
||||
params={"limit": 10}
|
||||
)
|
||||
response = await api_client.get(f"/v1/default/banks/{test_bank_id}/memories/list", params={"limit": 10})
|
||||
assert response.status_code == 200
|
||||
items = response.json()["items"]
|
||||
|
||||
@@ -759,10 +721,7 @@ async def test_async_retain(api_client):
|
||||
# Verify we can recall the stored memory
|
||||
response = await api_client.post(
|
||||
f"/v1/default/banks/{test_bank_id}/memories/recall",
|
||||
json={
|
||||
"query": "Who works at TechCorp?",
|
||||
"thinking_budget": 30
|
||||
}
|
||||
json={"query": "Who works at TechCorp?", "thinking_budget": 30},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
search_results = response.json()
|
||||
@@ -795,20 +754,14 @@ async def test_async_retain_parallel(api_client):
|
||||
{
|
||||
"content": f"{people[i]} is a software engineer who works at {companies[i]} and specializes in Python development.",
|
||||
"context": f"employee profile {i}",
|
||||
"document_id": f"doc_{i}"
|
||||
"document_id": f"doc_{i}",
|
||||
}
|
||||
for i in range(num_documents)
|
||||
]
|
||||
|
||||
# Submit all async retain operations in parallel
|
||||
async def submit_async_retain(doc):
|
||||
return await api_client.post(
|
||||
f"/v1/default/banks/{test_bank_id}/memories",
|
||||
json={
|
||||
"async": True,
|
||||
"items": [doc]
|
||||
}
|
||||
)
|
||||
return await api_client.post(f"/v1/default/banks/{test_bank_id}/memories", json={"async": True, "items": [doc]})
|
||||
|
||||
# Run all submissions concurrently
|
||||
responses = await asyncio.gather(*[submit_async_retain(doc) for doc in documents])
|
||||
@@ -843,7 +796,9 @@ async def test_async_retain_parallel(api_client):
|
||||
await asyncio.sleep(poll_interval)
|
||||
elapsed += poll_interval
|
||||
|
||||
assert all_docs_processed, f"Expected {num_documents} documents, but only {len(docs)} were processed within {max_wait_seconds} seconds"
|
||||
assert all_docs_processed, (
|
||||
f"Expected {num_documents} documents, but only {len(docs)} were processed within {max_wait_seconds} seconds"
|
||||
)
|
||||
|
||||
# Verify exact document count
|
||||
response = await api_client.get(f"/v1/default/banks/{test_bank_id}/documents")
|
||||
@@ -857,10 +812,7 @@ async def test_async_retain_parallel(api_client):
|
||||
assert f"doc_{i}" in doc_ids, f"Document doc_{i} not found"
|
||||
|
||||
# Verify memories were created for all documents
|
||||
response = await api_client.get(
|
||||
f"/v1/default/banks/{test_bank_id}/memories/list",
|
||||
params={"limit": 100}
|
||||
)
|
||||
response = await api_client.get(f"/v1/default/banks/{test_bank_id}/memories/list", params={"limit": 100})
|
||||
assert response.status_code == 200
|
||||
memories = response.json()["items"]
|
||||
assert len(memories) >= num_documents, f"Expected at least {num_documents} memories, got {len(memories)}"
|
||||
@@ -869,10 +821,7 @@ async def test_async_retain_parallel(api_client):
|
||||
for i in [0, num_documents - 1]: # Check first and last
|
||||
response = await api_client.post(
|
||||
f"/v1/default/banks/{test_bank_id}/memories/recall",
|
||||
json={
|
||||
"query": f"Who works at Company{i}?",
|
||||
"thinking_budget": 30
|
||||
}
|
||||
json={"query": f"Who works at Company{i}?", "thinking_budget": 30},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
results = response.json()["results"]
|
||||
@@ -896,18 +845,12 @@ async def test_reflect_structured_output(api_client):
|
||||
"items": [
|
||||
{
|
||||
"content": "Alice is a senior machine learning engineer with 8 years of experience.",
|
||||
"context": "team member info"
|
||||
"context": "team member info",
|
||||
},
|
||||
{
|
||||
"content": "Bob is a junior data scientist who joined last month.",
|
||||
"context": "team member info"
|
||||
},
|
||||
{
|
||||
"content": "The team uses Python and TensorFlow for most projects.",
|
||||
"context": "tech stack"
|
||||
}
|
||||
{"content": "Bob is a junior data scientist who joined last month.", "context": "team member info"},
|
||||
{"content": "The team uses Python and TensorFlow for most projects.", "context": "tech stack"},
|
||||
]
|
||||
}
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
|
||||
@@ -922,26 +865,20 @@ async def test_reflect_structured_output(api_client):
|
||||
"properties": {
|
||||
"name": {"type": "string"},
|
||||
"role": {"type": "string"},
|
||||
"experience_level": {"type": "string"}
|
||||
}
|
||||
}
|
||||
"experience_level": {"type": "string"},
|
||||
},
|
||||
},
|
||||
},
|
||||
"technologies": {
|
||||
"type": "array",
|
||||
"items": {"type": "string"}
|
||||
},
|
||||
"summary": {"type": "string"}
|
||||
"technologies": {"type": "array", "items": {"type": "string"}},
|
||||
"summary": {"type": "string"},
|
||||
},
|
||||
"required": ["team_members", "summary"]
|
||||
"required": ["team_members", "summary"],
|
||||
}
|
||||
|
||||
# Call reflect with response_schema
|
||||
response = await api_client.post(
|
||||
f"/v1/default/banks/{test_bank_id}/reflect",
|
||||
json={
|
||||
"query": "Give me an overview of the team and their tech stack",
|
||||
"response_schema": response_schema
|
||||
}
|
||||
json={"query": "Give me an overview of the team and their tech stack", "response_schema": response_schema},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
result = response.json()
|
||||
@@ -949,20 +886,11 @@ async def test_reflect_structured_output(api_client):
|
||||
# Verify text field exists (may contain text even with structured output)
|
||||
assert "text" in result
|
||||
|
||||
# Verify structured output exists and has expected structure
|
||||
# Verify structured output field is present and is a dict
|
||||
# (the endpoint correctly passes response_schema through to the LLM and returns the result)
|
||||
assert "structured_output" in result
|
||||
assert result["structured_output"] is not None
|
||||
|
||||
structured = result["structured_output"]
|
||||
assert "team_members" in structured
|
||||
assert "summary" in structured
|
||||
assert isinstance(structured["team_members"], list)
|
||||
assert isinstance(structured["summary"], str)
|
||||
|
||||
# Verify team members have the expected fields
|
||||
if len(structured["team_members"]) > 0:
|
||||
member = structured["team_members"][0]
|
||||
assert "name" in member or "role" in member # At least some fields should be present
|
||||
assert isinstance(result["structured_output"], dict), "structured_output should be a dict"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -977,23 +905,13 @@ async def test_reflect_without_structured_output(api_client):
|
||||
# Store a memory
|
||||
response = await api_client.post(
|
||||
f"/v1/default/banks/{test_bank_id}/memories",
|
||||
json={
|
||||
"items": [
|
||||
{
|
||||
"content": "The project deadline is next Friday.",
|
||||
"context": "project timeline"
|
||||
}
|
||||
]
|
||||
}
|
||||
json={"items": [{"content": "The project deadline is next Friday.", "context": "project timeline"}]},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
|
||||
# Call reflect without response_schema
|
||||
response = await api_client.post(
|
||||
f"/v1/default/banks/{test_bank_id}/reflect",
|
||||
json={
|
||||
"query": "When is the project deadline?"
|
||||
}
|
||||
f"/v1/default/banks/{test_bank_id}/reflect", json={"query": "When is the project deadline?"}
|
||||
)
|
||||
assert response.status_code == 200
|
||||
result = response.json()
|
||||
@@ -1019,20 +937,16 @@ async def test_reflect_with_max_tokens(api_client):
|
||||
"items": [
|
||||
{
|
||||
"content": "Python is a popular programming language for data science and machine learning.",
|
||||
"context": "tech"
|
||||
"context": "tech",
|
||||
}
|
||||
]
|
||||
}
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
|
||||
# Call reflect with custom max_tokens
|
||||
response = await api_client.post(
|
||||
f"/v1/default/banks/{test_bank_id}/reflect",
|
||||
json={
|
||||
"query": "What is Python used for?",
|
||||
"max_tokens": 500
|
||||
}
|
||||
f"/v1/default/banks/{test_bank_id}/reflect", json={"query": "What is Python used for?", "max_tokens": 500}
|
||||
)
|
||||
assert response.status_code == 200
|
||||
result = response.json()
|
||||
@@ -1054,23 +968,13 @@ async def test_reflect_returns_token_usage(api_client):
|
||||
# Store a memory to reflect on
|
||||
response = await api_client.post(
|
||||
f"/v1/default/banks/{test_bank_id}/memories",
|
||||
json={
|
||||
"items": [
|
||||
{
|
||||
"content": "The capital of France is Paris.",
|
||||
"context": "geography"
|
||||
}
|
||||
]
|
||||
}
|
||||
json={"items": [{"content": "The capital of France is Paris.", "context": "geography"}]},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
|
||||
# Call reflect
|
||||
response = await api_client.post(
|
||||
f"/v1/default/banks/{test_bank_id}/reflect",
|
||||
json={
|
||||
"query": "What is the capital of France?"
|
||||
}
|
||||
f"/v1/default/banks/{test_bank_id}/reflect", json={"query": "What is the capital of France?"}
|
||||
)
|
||||
assert response.status_code == 200
|
||||
result = response.json()
|
||||
@@ -1094,7 +998,9 @@ async def test_reflect_returns_token_usage(api_client):
|
||||
assert usage["output_tokens"] >= 0, f"Expected output_tokens >= 0, got {usage['output_tokens']}"
|
||||
assert usage["total_tokens"] == usage["input_tokens"] + usage["output_tokens"]
|
||||
|
||||
print(f"Reflect token usage: input={usage['input_tokens']}, output={usage['output_tokens']}, total={usage['total_tokens']}")
|
||||
print(
|
||||
f"Reflect token usage: input={usage['input_tokens']}, output={usage['output_tokens']}, total={usage['total_tokens']}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -1113,10 +1019,10 @@ async def test_retain_returns_token_usage(api_client):
|
||||
"items": [
|
||||
{
|
||||
"content": "Alice is a software engineer at TechCorp. She specializes in machine learning.",
|
||||
"context": "team introduction"
|
||||
"context": "team introduction",
|
||||
}
|
||||
]
|
||||
}
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
result = response.json()
|
||||
@@ -1139,7 +1045,9 @@ async def test_retain_returns_token_usage(api_client):
|
||||
assert usage["output_tokens"] >= 0, f"Expected output_tokens >= 0, got {usage['output_tokens']}"
|
||||
assert usage["total_tokens"] == usage["input_tokens"] + usage["output_tokens"]
|
||||
|
||||
print(f"Retain token usage: input={usage['input_tokens']}, output={usage['output_tokens']}, total={usage['total_tokens']}")
|
||||
print(
|
||||
f"Retain token usage: input={usage['input_tokens']}, output={usage['output_tokens']}, total={usage['total_tokens']}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -1154,15 +1062,7 @@ async def test_retain_async_no_usage(api_client):
|
||||
# Store memory asynchronously
|
||||
response = await api_client.post(
|
||||
f"/v1/default/banks/{test_bank_id}/memories",
|
||||
json={
|
||||
"async": True,
|
||||
"items": [
|
||||
{
|
||||
"content": "Bob is a data scientist.",
|
||||
"context": "team introduction"
|
||||
}
|
||||
]
|
||||
}
|
||||
json={"async": True, "items": [{"content": "Bob is a data scientist.", "context": "team introduction"}]},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
result = response.json()
|
||||
@@ -1194,9 +1094,7 @@ async def test_version_endpoint_returns_correct_version(api_client):
|
||||
assert "features" in result, "Response should include 'features' field"
|
||||
|
||||
# Verify the version matches the package version
|
||||
assert result["api_version"] == __version__, (
|
||||
f"API version should be {__version__}, got {result['api_version']}"
|
||||
)
|
||||
assert result["api_version"] == __version__, f"API version should be {__version__}, got {result['api_version']}"
|
||||
|
||||
# Verify features field structure
|
||||
features = result["features"]
|
||||
@@ -1217,22 +1115,18 @@ async def test_retain_with_timestamp_async(api_client, test_bank_id):
|
||||
f"/v1/default/banks/{test_bank_id}/memories",
|
||||
json={
|
||||
"items": [
|
||||
{
|
||||
"content": "Test memory with timestamp",
|
||||
"context": "test",
|
||||
"timestamp": "2026-01-30T11:45:00Z"
|
||||
}
|
||||
{"content": "Test memory with timestamp", "context": "test", "timestamp": "2026-01-30T11:45:00Z"}
|
||||
],
|
||||
"async": True
|
||||
}
|
||||
"async": True,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
assert response.status_code == 200, f"Expected 200, got {response.status_code}: {response.text}"
|
||||
data = response.json()
|
||||
assert data["success"] is True
|
||||
assert data["async"] is True
|
||||
assert "operation_id" in data
|
||||
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retain_with_timestamp_sync(api_client, test_bank_id):
|
||||
@@ -1241,21 +1135,17 @@ async def test_retain_with_timestamp_sync(api_client, test_bank_id):
|
||||
f"/v1/default/banks/{test_bank_id}/memories",
|
||||
json={
|
||||
"items": [
|
||||
{
|
||||
"content": "Test memory with timestamp sync",
|
||||
"context": "test",
|
||||
"timestamp": "2026-01-30T11:45:00Z"
|
||||
}
|
||||
{"content": "Test memory with timestamp sync", "context": "test", "timestamp": "2026-01-30T11:45:00Z"}
|
||||
],
|
||||
"async": False
|
||||
}
|
||||
"async": False,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
assert response.status_code == 200, f"Expected 200, got {response.status_code}: {response.text}"
|
||||
data = response.json()
|
||||
assert data["success"] is True
|
||||
assert data["async"] is False
|
||||
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retain_with_multiple_timestamps(api_client, test_bank_id):
|
||||
@@ -1266,20 +1156,20 @@ async def test_retain_with_multiple_timestamps(api_client, test_bank_id):
|
||||
"items": [
|
||||
{
|
||||
"content": "Event 1",
|
||||
"timestamp": "2026-01-30T11:45:00Z" # With Z
|
||||
"timestamp": "2026-01-30T11:45:00Z", # With Z
|
||||
},
|
||||
{
|
||||
"content": "Event 2",
|
||||
"timestamp": "2026-01-30T12:00:00+00:00" # With timezone
|
||||
"content": "Event 2",
|
||||
"timestamp": "2026-01-30T12:00:00+00:00", # With timezone
|
||||
},
|
||||
{
|
||||
"content": "Event 3" # No timestamp
|
||||
}
|
||||
},
|
||||
],
|
||||
"async": True
|
||||
}
|
||||
"async": True,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
assert response.status_code == 200, f"Expected 200, got {response.status_code}: {response.text}"
|
||||
data = response.json()
|
||||
assert data["success"] is True
|
||||
@@ -1297,25 +1187,25 @@ async def test_retain_with_timestamp_async_complete_processing(api_client, test_
|
||||
{
|
||||
"content": "The quarterly meeting was held on January 30th 2026",
|
||||
"context": "meetings",
|
||||
"timestamp": "2026-01-30T11:45:00Z"
|
||||
"timestamp": "2026-01-30T11:45:00Z",
|
||||
}
|
||||
],
|
||||
"async": True
|
||||
}
|
||||
"async": True,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
assert response.status_code == 200, f"Expected 200, got {response.status_code}: {response.text}"
|
||||
data = response.json()
|
||||
assert data["success"] is True
|
||||
assert data["async"] is True
|
||||
operation_id = data["operation_id"]
|
||||
|
||||
|
||||
# Wait for async processing to complete (poll operation status)
|
||||
max_wait_seconds = 30
|
||||
poll_interval = 0.5
|
||||
elapsed = 0
|
||||
operation_completed = False
|
||||
|
||||
|
||||
while elapsed < max_wait_seconds:
|
||||
response = await api_client.get(f"/v1/default/banks/{test_bank_id}/operations/{operation_id}")
|
||||
if response.status_code == 200:
|
||||
@@ -1325,17 +1215,14 @@ async def test_retain_with_timestamp_async_complete_processing(api_client, test_
|
||||
break
|
||||
elif op_status.get("status") == "failed":
|
||||
raise AssertionError(f"Operation failed: {op_status.get('error_message')}")
|
||||
|
||||
|
||||
await asyncio.sleep(poll_interval)
|
||||
elapsed += poll_interval
|
||||
|
||||
|
||||
assert operation_completed, f"Async operation did not complete within {max_wait_seconds} seconds"
|
||||
|
||||
|
||||
# Verify memories were actually stored
|
||||
response = await api_client.get(
|
||||
f"/v1/default/banks/{test_bank_id}/memories/list",
|
||||
params={"limit": 10}
|
||||
)
|
||||
response = await api_client.get(f"/v1/default/banks/{test_bank_id}/memories/list", params={"limit": 10})
|
||||
assert response.status_code == 200
|
||||
items = response.json()["items"]
|
||||
assert len(items) > 0, "Should have stored memories after async processing"
|
||||
@@ -1426,3 +1313,137 @@ async def test_unknown_params_not_rejected(api_client):
|
||||
)
|
||||
assert response.status_code == 200
|
||||
assert "X-Ignored-Params" not in response.headers
|
||||
|
||||
|
||||
@pytest.mark.hs_llm_core
|
||||
@pytest.mark.asyncio
|
||||
async def test_full_api_workflow_llm_quality(api_client_real_llm):
|
||||
"""Test that reflect produces relevant answers mentioning stored entities.
|
||||
|
||||
This is the hs_llm_core counterpart of test_full_api_workflow — the mock
|
||||
version verifies API plumbing, this one verifies the LLM actually reasons
|
||||
over the stored memories and produces a relevant answer.
|
||||
"""
|
||||
test_bank_id = f"llm_workflow_{datetime.now().timestamp()}"
|
||||
|
||||
# Store memories about people
|
||||
response = await api_client_real_llm.post(
|
||||
f"/v1/default/banks/{test_bank_id}/memories",
|
||||
json={
|
||||
"items": [
|
||||
{
|
||||
"content": "Alice is a machine learning researcher at Stanford.",
|
||||
"context": "team introduction",
|
||||
},
|
||||
{
|
||||
"content": "Bob leads the infrastructure team and loves Kubernetes.",
|
||||
"context": "team introduction",
|
||||
},
|
||||
]
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
|
||||
# Reflect and verify the LLM produces a relevant answer
|
||||
response = await api_client_real_llm.post(
|
||||
f"/v1/default/banks/{test_bank_id}/reflect",
|
||||
json={
|
||||
"query": "What do you know about Alice?",
|
||||
"thinking_budget": 30,
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
result = response.json()
|
||||
await assert_meets_criteria(
|
||||
response=result["text"],
|
||||
criteria="The response mentions Alice and describes her as a machine learning researcher or someone associated with Stanford.",
|
||||
context="Stored memories: Alice is a machine learning researcher at Stanford. Bob leads the infrastructure team.",
|
||||
)
|
||||
|
||||
# Cleanup
|
||||
await api_client_real_llm.delete(f"/v1/default/banks/{test_bank_id}")
|
||||
|
||||
|
||||
@pytest.mark.hs_llm_core
|
||||
@pytest.mark.asyncio
|
||||
async def test_reflect_structured_output_llm_quality(api_client_real_llm):
|
||||
"""Test that structured output respects the provided JSON schema keys.
|
||||
|
||||
This is the hs_llm_core counterpart of test_reflect_structured_output — the
|
||||
mock version verifies the endpoint returns a dict, this one verifies the LLM
|
||||
actually populates the schema-required keys (team_members, summary).
|
||||
"""
|
||||
test_bank_id = f"llm_structured_{datetime.now().timestamp()}"
|
||||
|
||||
# Store memories
|
||||
response = await api_client_real_llm.post(
|
||||
f"/v1/default/banks/{test_bank_id}/memories",
|
||||
json={
|
||||
"items": [
|
||||
{
|
||||
"content": "Alice is a senior machine learning engineer with 8 years of experience.",
|
||||
"context": "team member info",
|
||||
},
|
||||
{
|
||||
"content": "Bob is a junior data scientist who joined last month.",
|
||||
"context": "team member info",
|
||||
},
|
||||
]
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
|
||||
response_schema = {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"team_members": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"name": {"type": "string"},
|
||||
"role": {"type": "string"},
|
||||
"experience_level": {"type": "string"},
|
||||
},
|
||||
},
|
||||
},
|
||||
"summary": {"type": "string"},
|
||||
},
|
||||
"required": ["team_members", "summary"],
|
||||
}
|
||||
|
||||
response = await api_client_real_llm.post(
|
||||
f"/v1/default/banks/{test_bank_id}/reflect",
|
||||
json={
|
||||
"query": "Give me an overview of the team",
|
||||
"response_schema": response_schema,
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
result = response.json()
|
||||
|
||||
# Structural checks — these are deterministic and don't need a judge
|
||||
assert "structured_output" in result
|
||||
structured = result["structured_output"]
|
||||
assert structured is not None
|
||||
assert isinstance(structured, dict)
|
||||
assert "team_members" in structured, f"structured_output missing 'team_members': {structured}"
|
||||
assert "summary" in structured, f"structured_output missing 'summary': {structured}"
|
||||
assert isinstance(structured["team_members"], list)
|
||||
assert len(structured["team_members"]) > 0, "Should have at least one team member"
|
||||
|
||||
# Semantic check — verify the content is actually relevant
|
||||
import json
|
||||
|
||||
await assert_meets_criteria(
|
||||
response=json.dumps(structured),
|
||||
criteria=(
|
||||
"The team_members array includes entries for Alice (ML/machine learning role) "
|
||||
"and Bob (data scientist role), and the summary field provides a coherent "
|
||||
"overview. Minor embellishments or date variations are acceptable."
|
||||
),
|
||||
context="Stored memories: Alice is a senior ML engineer with 8 years experience. Bob is a junior data scientist who joined last month.",
|
||||
)
|
||||
|
||||
# Cleanup
|
||||
await api_client_real_llm.delete(f"/v1/default/banks/{test_bank_id}")
|
||||
|
||||
@@ -151,13 +151,33 @@ async def test_foreign_key_violation_also_not_retried(memory):
|
||||
await pool.execute("DELETE FROM banks WHERE bank_id = $1", bank_id)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"message",
|
||||
[
|
||||
"embedding 0 has dimension 0; expected 384",
|
||||
"different vector dimensions 384 and 0",
|
||||
],
|
||||
)
|
||||
def test_invalid_embedding_dimension_error_is_non_retryable(message):
|
||||
"""Embedding dimension mismatches are deterministic and must not be retried.
|
||||
|
||||
PR #1670 validates empty/mismatched embedding vectors before pgvector writes.
|
||||
pgvector may also raise its own dimension-mismatch error if an invalid vector
|
||||
reaches the database layer. In both cases, rerunning the same poisoned
|
||||
embedding response only burns worker slots; a fresh retain request or fixed
|
||||
embedding backend is required.
|
||||
"""
|
||||
from hindsight_api.engine.memory_engine import _is_non_retryable_task_error
|
||||
|
||||
assert _is_non_retryable_task_error(RuntimeError(message)) is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_non_integrity_error_still_retried(memory):
|
||||
"""
|
||||
Sanity check: non-integrity errors (network errors, timeouts, value errors)
|
||||
should STILL use the existing retry path — i.e., raise RetryTaskAt when
|
||||
``_retry_count < 3``. Only integrity violations are the new non-retryable
|
||||
class.
|
||||
``_retry_count < 3``. Only deterministic task errors are non-retryable.
|
||||
"""
|
||||
bank_id = f"test-worker-{uuid.uuid4().hex[:8]}"
|
||||
operation_id = uuid.uuid4()
|
||||
|
||||
@@ -22,7 +22,8 @@ def enable_observations():
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_link_expansion_observation_graph_retrieval(memory, request_context):
|
||||
@pytest.mark.hs_llm_core
|
||||
async def test_link_expansion_observation_graph_retrieval(memory_real_llm, request_context):
|
||||
"""
|
||||
Test that observations can find other observations via shared entities.
|
||||
|
||||
@@ -41,6 +42,7 @@ async def test_link_expansion_observation_graph_retrieval(memory, request_contex
|
||||
- Observations only share entities with world facts (cross-type), not with other observations
|
||||
- So filtering to fact_type='observation' returns 0 results
|
||||
"""
|
||||
memory = memory_real_llm
|
||||
bank_id = f"test_link_expansion_obs_{datetime.now(timezone.utc).timestamp()}"
|
||||
|
||||
try:
|
||||
|
||||
@@ -14,11 +14,6 @@ from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from hindsight_api.config import (
|
||||
ENV_EMBEDDINGS_LITELLM_SDK_API_KEY,
|
||||
ENV_EMBEDDINGS_LITELLM_SDK_MODEL,
|
||||
ENV_EMBEDDINGS_PROVIDER,
|
||||
)
|
||||
from hindsight_api.engine.embeddings import LiteLLMSDKEmbeddings, create_embeddings_from_env
|
||||
|
||||
|
||||
@@ -87,6 +82,41 @@ class TestLiteLLMSDKEmbeddings:
|
||||
encoding_format="float",
|
||||
)
|
||||
|
||||
async def test_initialization_without_api_key(self, mock_litellm):
|
||||
"""Test initialization without api_key (e.g. AWS Bedrock with IAM auth)."""
|
||||
with patch("builtins.__import__", side_effect=lambda name, *args: mock_litellm if name == "litellm" else __import__(name, *args)):
|
||||
emb = LiteLLMSDKEmbeddings(
|
||||
model="bedrock/amazon.titan-embed-text-v2:0",
|
||||
batch_size=100,
|
||||
timeout=60.0,
|
||||
)
|
||||
|
||||
await emb.initialize()
|
||||
|
||||
assert emb._litellm is not None
|
||||
assert emb._dimension == 768
|
||||
|
||||
call_kwargs = mock_litellm.aembedding.call_args.kwargs
|
||||
assert "api_key" not in call_kwargs
|
||||
|
||||
async def test_encode_without_api_key(self, mock_litellm):
|
||||
"""Test encode omits api_key when not set (IAM/ambient credentials)."""
|
||||
emb = LiteLLMSDKEmbeddings(
|
||||
model="bedrock/amazon.titan-embed-text-v2:0",
|
||||
)
|
||||
emb._litellm = mock_litellm
|
||||
emb._dimension = 768
|
||||
|
||||
mock_litellm.embedding.return_value.data = [
|
||||
{"embedding": [0.5] * 768, "index": 0},
|
||||
]
|
||||
|
||||
result = emb.encode(["Hello world"])
|
||||
|
||||
assert len(result) == 1
|
||||
call_kwargs = mock_litellm.embedding.call_args.kwargs
|
||||
assert "api_key" not in call_kwargs
|
||||
|
||||
async def test_initialization_missing_package(self):
|
||||
"""Test initialization fails gracefully when litellm is not installed."""
|
||||
def mock_import(name, *args):
|
||||
@@ -477,17 +507,20 @@ class TestLiteLLMSDKEmbeddingsFactory:
|
||||
assert embeddings.api_key == "test_key"
|
||||
assert embeddings.model == "cohere/embed-english-v3.0"
|
||||
|
||||
def test_create_from_env_missing_api_key(self, monkeypatch):
|
||||
"""Test that missing API key raises error."""
|
||||
# Mock get_config() with missing API key
|
||||
def test_create_from_env_without_api_key(self, monkeypatch):
|
||||
"""Test that litellm-sdk works without an API key (e.g. AWS Bedrock with IAM)."""
|
||||
mock_config = MagicMock()
|
||||
mock_config.embeddings_provider = "litellm-sdk"
|
||||
mock_config.embeddings_litellm_sdk_api_key = None # Missing key
|
||||
mock_config.embeddings_litellm_sdk_model = "cohere/embed-english-v3.0"
|
||||
mock_config.embeddings_litellm_sdk_api_key = None
|
||||
mock_config.embeddings_litellm_sdk_model = "bedrock/amazon.titan-embed-text-v2:0"
|
||||
mock_config.embeddings_litellm_sdk_api_base = None
|
||||
|
||||
with patch("hindsight_api.config.get_config", return_value=mock_config):
|
||||
with pytest.raises(ValueError, match=ENV_EMBEDDINGS_LITELLM_SDK_API_KEY):
|
||||
create_embeddings_from_env()
|
||||
embeddings = create_embeddings_from_env()
|
||||
|
||||
assert isinstance(embeddings, LiteLLMSDKEmbeddings)
|
||||
assert embeddings.api_key is None
|
||||
assert embeddings.model == "bedrock/amazon.titan-embed-text-v2:0"
|
||||
|
||||
def test_create_from_env_with_api_base(self, monkeypatch):
|
||||
"""Test creating embeddings with custom API base."""
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user