Compare commits
56
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
f515b8207b | ||
|
|
0e3ccd038b | ||
|
|
e4dd654ec5 | ||
|
|
0ad8c2d09c | ||
|
|
1e40cd22a6 | ||
|
|
1e5aa7de4d | ||
|
|
58fdac44f7 | ||
|
|
891c33b1d7 | ||
|
|
7ed57fdd85 | ||
|
|
e0a2ac63e7 | ||
|
|
75b95106ba | ||
|
|
d425e93cb4 | ||
|
|
cb6d1c469c | ||
|
|
4c058b4b98 | ||
|
|
aa8e5475c4 | ||
|
|
ad2cf72aab | ||
|
|
3f2a6ec9ce | ||
|
|
f17406fdf0 | ||
|
|
66423b85f5 | ||
|
|
abbf874d84 | ||
|
|
51d2fc5309 | ||
|
|
ea27ef95ec | ||
|
|
edf60e0f3c | ||
|
|
719e79a4d9 | ||
|
|
3857a30491 | ||
|
|
3d87ef5cee | ||
|
|
5c3d3274d7 | ||
|
|
144e4c49d1 | ||
|
|
861295dd7c | ||
|
|
15f4b8769b | ||
|
|
61bf428ba9 | ||
|
|
73ef99e7b1 | ||
|
|
7942f181c2 | ||
|
|
5aff8e0c70 | ||
|
|
e407f4bc55 | ||
|
|
8138fa9002 | ||
|
|
1d70abfe85 | ||
|
|
ecf609c8aa | ||
|
|
cab5a40f3a | ||
|
|
ab70da1ead | ||
|
|
9b96becc5c | ||
|
|
f903948a26 | ||
|
|
77defd96e9 | ||
|
|
c2876490df | ||
|
|
eaeaa1f24d | ||
|
|
f6f1a7d889 | ||
|
|
ecb833f40d | ||
|
|
5270aa5a6e | ||
|
|
ad1660b313 | ||
|
|
55af468187 | ||
|
|
2b5fb10dab | ||
|
|
5443c18bfc | ||
|
|
1c21c0c1a6 | ||
|
|
145454533c | ||
|
|
9aaa78b9a6 | ||
|
|
bfa09d1685 |
@@ -50,6 +50,10 @@ jobs:
|
||||
working-directory: ./hindsight-integrations/crewai
|
||||
run: uv build --out-dir dist
|
||||
|
||||
- name: Build hindsight-pydantic-ai
|
||||
working-directory: ./hindsight-integrations/pydantic-ai
|
||||
run: uv build --out-dir dist
|
||||
|
||||
# Publish in order (client and api first, then hindsight-all which depends on them)
|
||||
- name: Publish hindsight-client to PyPI
|
||||
uses: pypa/gh-action-pypi-publish@release/v1
|
||||
@@ -87,6 +91,12 @@ jobs:
|
||||
packages-dir: ./hindsight-integrations/crewai/dist
|
||||
skip-existing: true
|
||||
|
||||
- name: Publish hindsight-pydantic-ai to PyPI
|
||||
uses: pypa/gh-action-pypi-publish@release/v1
|
||||
with:
|
||||
packages-dir: ./hindsight-integrations/pydantic-ai/dist
|
||||
skip-existing: true
|
||||
|
||||
# Upload artifacts for GitHub release
|
||||
- name: Upload artifacts
|
||||
uses: actions/upload-artifact@v4
|
||||
@@ -99,6 +109,7 @@ jobs:
|
||||
hindsight-integrations/litellm/dist/*
|
||||
hindsight-embed/dist/*
|
||||
hindsight-integrations/crewai/dist/*
|
||||
hindsight-integrations/pydantic-ai/dist/*
|
||||
retention-days: 1
|
||||
|
||||
release-typescript-client:
|
||||
@@ -376,6 +387,10 @@ jobs:
|
||||
target: aarch64-apple-darwin
|
||||
artifact_name: hindsight
|
||||
asset_name: hindsight-darwin-arm64
|
||||
- os: ubuntu-24.04-arm
|
||||
target: aarch64-unknown-linux-gnu
|
||||
artifact_name: hindsight
|
||||
asset_name: hindsight-linux-arm64
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
@@ -629,6 +644,7 @@ jobs:
|
||||
cp artifacts/python-packages/hindsight-api/dist/* release-assets/ || true
|
||||
cp artifacts/python-packages/hindsight/dist/* release-assets/ || true
|
||||
cp artifacts/python-packages/hindsight-integrations/litellm/dist/* release-assets/ || true
|
||||
cp artifacts/python-packages/hindsight-integrations/pydantic-ai/dist/* release-assets/ || true
|
||||
cp artifacts/python-packages/hindsight-embed/dist/* release-assets/ || true
|
||||
# TypeScript client
|
||||
cp artifacts/typescript-client/*.tgz release-assets/ || true
|
||||
|
||||
+83
-10
@@ -686,6 +686,30 @@ jobs:
|
||||
echo "=== API Server Logs ==="
|
||||
cat /tmp/api-server.log || echo "No API server log found"
|
||||
|
||||
build-rust-cli-arm64:
|
||||
runs-on: ubuntu-24.04-arm
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
- name: Install Rust
|
||||
uses: dtolnay/rust-toolchain@stable
|
||||
with:
|
||||
targets: aarch64-unknown-linux-gnu
|
||||
|
||||
- name: Cache cargo
|
||||
uses: actions/cache@v4
|
||||
with:
|
||||
path: |
|
||||
~/.cargo/registry
|
||||
~/.cargo/git
|
||||
hindsight-cli/target
|
||||
key: linux-arm64-cargo-${{ hashFiles('**/Cargo.lock') }}
|
||||
|
||||
- name: Build CLI
|
||||
working-directory: hindsight-cli
|
||||
run: cargo build --release --target aarch64-unknown-linux-gnu
|
||||
|
||||
test-rust-client:
|
||||
runs-on: ubuntu-latest
|
||||
env:
|
||||
@@ -1162,6 +1186,35 @@ jobs:
|
||||
working-directory: ./hindsight-integrations/litellm
|
||||
run: uv run pytest tests -v
|
||||
|
||||
test-pydantic-ai-integration:
|
||||
runs-on: ubuntu-latest
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
- name: Install uv
|
||||
uses: astral-sh/setup-uv@v5
|
||||
with:
|
||||
enable-cache: true
|
||||
prune-cache: false
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v5
|
||||
with:
|
||||
python-version-file: ".python-version"
|
||||
|
||||
- name: Build pydantic-ai integration
|
||||
working-directory: ./hindsight-integrations/pydantic-ai
|
||||
run: uv build
|
||||
|
||||
- name: Install dependencies
|
||||
working-directory: ./hindsight-integrations/pydantic-ai
|
||||
run: uv sync --frozen
|
||||
|
||||
- name: Run tests
|
||||
working-directory: ./hindsight-integrations/pydantic-ai
|
||||
run: uv run pytest tests -v
|
||||
|
||||
test-embed:
|
||||
runs-on: ubuntu-latest
|
||||
env:
|
||||
@@ -1268,7 +1321,11 @@ jobs:
|
||||
|
||||
test-doc-examples:
|
||||
runs-on: ubuntu-latest
|
||||
needs: test-rust-cli
|
||||
strategy:
|
||||
fail-fast: false
|
||||
matrix:
|
||||
language: [python, node, cli, go]
|
||||
name: test-doc-examples (${{ matrix.language }})
|
||||
env:
|
||||
HINDSIGHT_API_LLM_PROVIDER: vertexai
|
||||
HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY: /tmp/gcp-credentials.json
|
||||
@@ -1286,14 +1343,26 @@ jobs:
|
||||
PROJECT_ID=$(jq -r '.project_id' /tmp/gcp-credentials.json)
|
||||
echo "HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID=$PROJECT_ID" >> $GITHUB_ENV
|
||||
|
||||
- name: Download CLI artifact
|
||||
uses: actions/download-artifact@v4
|
||||
with:
|
||||
name: hindsight-cli
|
||||
path: /usr/local/bin
|
||||
- name: Install Rust
|
||||
if: matrix.language == 'cli'
|
||||
uses: dtolnay/rust-toolchain@stable
|
||||
|
||||
- name: Make CLI executable
|
||||
run: chmod +x /usr/local/bin/hindsight
|
||||
- name: Cache cargo
|
||||
if: matrix.language == 'cli'
|
||||
uses: actions/cache@v4
|
||||
with:
|
||||
path: |
|
||||
~/.cargo/registry
|
||||
~/.cargo/git
|
||||
hindsight-cli/target
|
||||
key: ${{ runner.os }}-cargo-${{ hashFiles('**/Cargo.lock') }}
|
||||
|
||||
- name: Build CLI
|
||||
if: matrix.language == 'cli'
|
||||
working-directory: hindsight-cli
|
||||
run: |
|
||||
cargo build --release
|
||||
cp target/release/hindsight /usr/local/bin/hindsight
|
||||
|
||||
- name: Install uv
|
||||
uses: astral-sh/setup-uv@v5
|
||||
@@ -1307,6 +1376,7 @@ jobs:
|
||||
python-version-file: ".python-version"
|
||||
|
||||
- name: Set up Node.js
|
||||
if: matrix.language == 'node'
|
||||
uses: actions/setup-node@v4
|
||||
with:
|
||||
node-version: '20'
|
||||
@@ -1320,10 +1390,12 @@ jobs:
|
||||
uv sync --frozen --no-install-project --index-strategy unsafe-best-match
|
||||
|
||||
- name: Install Python client dependencies
|
||||
if: matrix.language == 'python'
|
||||
working-directory: ./hindsight-clients/python
|
||||
run: uv sync --frozen --extra test --index-strategy unsafe-best-match
|
||||
|
||||
- name: Install TypeScript client
|
||||
if: matrix.language == 'node'
|
||||
run: |
|
||||
npm ci --workspace=hindsight-clients/typescript
|
||||
npm run build --workspace=hindsight-clients/typescript
|
||||
@@ -1375,10 +1447,11 @@ jobs:
|
||||
done
|
||||
|
||||
- name: Configure CLI
|
||||
if: matrix.language == 'cli'
|
||||
run: hindsight configure --api-url http://localhost:8888
|
||||
|
||||
- name: Run all doc examples
|
||||
run: ./scripts/test-doc-examples.sh
|
||||
- name: Run doc examples (${{ matrix.language }})
|
||||
run: ./scripts/test-doc-examples.sh --lang ${{ matrix.language }}
|
||||
|
||||
- name: Show API server logs
|
||||
if: always()
|
||||
|
||||
@@ -97,7 +97,7 @@ fi
|
||||
if [ "$ENABLE_CP" = "true" ]; then
|
||||
echo "🎛️ Starting Control Plane..."
|
||||
cd /app/control-plane
|
||||
PORT=9999 node server.js &
|
||||
PORT="${HINDSIGHT_CP_PORT:-9999}" node server.js &
|
||||
CP_PID=$!
|
||||
PIDS+=($CP_PID)
|
||||
else
|
||||
@@ -110,7 +110,7 @@ echo "✅ Hindsight is running!"
|
||||
echo ""
|
||||
echo "📍 Access:"
|
||||
if [ "$ENABLE_CP" = "true" ]; then
|
||||
echo " Control Plane: http://localhost:9999"
|
||||
echo " Control Plane: http://localhost:${HINDSIGHT_CP_PORT:-9999}"
|
||||
fi
|
||||
if [ "$ENABLE_API" = "true" ]; then
|
||||
echo " API: http://localhost:8888"
|
||||
|
||||
@@ -2,8 +2,8 @@ apiVersion: v2
|
||||
name: hindsight
|
||||
description: Hindsight helm chart
|
||||
type: application
|
||||
version: 0.4.13
|
||||
appVersion: "0.4.13"
|
||||
version: 0.4.16
|
||||
appVersion: "0.4.16"
|
||||
keywords:
|
||||
- ai
|
||||
- memory
|
||||
|
||||
@@ -46,4 +46,4 @@ __all__ = [
|
||||
"RemoteTEICrossEncoder",
|
||||
"LLMConfig",
|
||||
]
|
||||
__version__ = "0.4.13"
|
||||
__version__ = "0.4.16"
|
||||
|
||||
@@ -0,0 +1,88 @@
|
||||
"""Add text_signals column to memory_units for enriched BM25 indexing.
|
||||
|
||||
text_signals stores a denormalized space-separated string of entity names
|
||||
(and future signals) to improve full-text search recall without polluting
|
||||
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)
|
||||
|
||||
Revision ID: a2b3c4d5e6f7
|
||||
Revises: z1u2v3w4x5y6
|
||||
Create Date: 2026-02-28
|
||||
"""
|
||||
|
||||
import os
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
revision: str = "a2b3c4d5e6f7"
|
||||
down_revision: str | Sequence[str] | None = "aa2b3c4d5e6f"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _get_schema_prefix() -> str:
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def _detect_text_search_extension() -> str:
|
||||
return os.getenv("HINDSIGHT_API_TEXT_SEARCH_EXTENSION", "native").lower()
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
table = f"{schema}memory_units"
|
||||
text_search_ext = _detect_text_search_extension()
|
||||
|
||||
# Add text_signals column (nullable TEXT, populated at retain time)
|
||||
op.execute(f"ALTER TABLE {table} ADD COLUMN IF NOT EXISTS text_signals TEXT")
|
||||
|
||||
if text_search_ext == "native":
|
||||
# Native PostgreSQL: drop and recreate the GENERATED tsvector column to include text_signals
|
||||
op.execute(f"ALTER TABLE {table} DROP COLUMN IF EXISTS search_vector")
|
||||
op.execute(f"""
|
||||
ALTER TABLE {table}
|
||||
ADD COLUMN search_vector tsvector
|
||||
GENERATED ALWAYS AS (
|
||||
to_tsvector('english',
|
||||
COALESCE(text, '') || ' ' ||
|
||||
COALESCE(context, '') || ' ' ||
|
||||
COALESCE(text_signals, '')
|
||||
)
|
||||
) STORED
|
||||
""")
|
||||
# Recreate GIN index (was dropped with the column)
|
||||
op.execute(f"""
|
||||
CREATE INDEX IF NOT EXISTS idx_memory_units_text_search
|
||||
ON {table} USING gin(search_vector)
|
||||
""")
|
||||
|
||||
# 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
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
table = f"{schema}memory_units"
|
||||
text_search_ext = _detect_text_search_extension()
|
||||
|
||||
if text_search_ext == "native":
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema}idx_memory_units_text_search")
|
||||
op.execute(f"ALTER TABLE {table} DROP COLUMN IF EXISTS search_vector")
|
||||
op.execute(f"""
|
||||
ALTER TABLE {table}
|
||||
ADD COLUMN search_vector tsvector
|
||||
GENERATED ALWAYS AS (
|
||||
to_tsvector('english', COALESCE(text, '') || ' ' || COALESCE(context, ''))
|
||||
) STORED
|
||||
""")
|
||||
op.execute(f"""
|
||||
CREATE INDEX idx_memory_units_text_search
|
||||
ON {table} USING gin(search_vector)
|
||||
""")
|
||||
|
||||
op.execute(f"ALTER TABLE {table} DROP COLUMN IF EXISTS text_signals")
|
||||
+54
@@ -0,0 +1,54 @@
|
||||
"""Add GIN index on source_memory_ids for observation lookup performance
|
||||
|
||||
Without this index, queries using the array overlap operator (&&) or array
|
||||
containment (@>) on source_memory_ids require a full sequential scan over all
|
||||
observation memory_units. At ~77k observations this was measured at 45ms per
|
||||
query, becoming a bottleneck during consolidation recall (57-64s timeouts) and
|
||||
user recall (18-27s average).
|
||||
|
||||
The GIN index reduces these queries to index scans: 45ms → 0.049ms (927x
|
||||
speedup). Recall dropped from 18-27s to ~6s, and consolidation recall
|
||||
stabilised from timeout to ~15s.
|
||||
|
||||
Created with CONCURRENTLY so the migration does not block reads or writes.
|
||||
CONCURRENTLY requires running outside a transaction block, so the migration
|
||||
emits an explicit COMMIT before the statement and uses IF NOT EXISTS for
|
||||
idempotency.
|
||||
|
||||
Revision ID: a2b3c4d5e6f8
|
||||
Revises: f7g8h9i0j1k2
|
||||
Create Date: 2026-03-04
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
revision: str = "a2b3c4d5e6f8"
|
||||
down_revision: str | Sequence[str] | None = "f7g8h9i0j1k2"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _get_schema_prefix() -> str:
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
|
||||
# CREATE INDEX CONCURRENTLY cannot run inside a transaction block.
|
||||
# Commit the current Alembic transaction first.
|
||||
op.execute("COMMIT")
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_units_source_memory_ids "
|
||||
f"ON {schema}memory_units USING GIN (source_memory_ids) "
|
||||
f"WHERE source_memory_ids IS NOT NULL"
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
op.execute("COMMIT")
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_units_source_memory_ids")
|
||||
@@ -0,0 +1,36 @@
|
||||
"""Make event_date nullable in memory_units to support timestamp-free content
|
||||
|
||||
Revision ID: aa2b3c4d5e6f
|
||||
Revises: z1u2v3w4x5y6
|
||||
Create Date: 2026-03-02
|
||||
|
||||
When callers retain content without a timestamp (e.g. fictional documents, static text),
|
||||
the event_date column should be allowed to be NULL rather than defaulting to utcnow().
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
revision: str = "aa2b3c4d5e6f"
|
||||
down_revision: str | Sequence[str] | None = "z1u2v3w4x5y6"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _get_schema_prefix() -> str:
|
||||
"""Get schema prefix for table names (required for multi-tenant support)."""
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
op.execute(f"ALTER TABLE {schema}memory_units ALTER COLUMN event_date DROP NOT NULL")
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
# Backfill NULLs with now() before restoring the NOT NULL constraint
|
||||
op.execute(f"UPDATE {schema}memory_units SET event_date = now() WHERE event_date IS NULL")
|
||||
op.execute(f"ALTER TABLE {schema}memory_units ALTER COLUMN event_date SET NOT NULL")
|
||||
+68
@@ -0,0 +1,68 @@
|
||||
"""Add partial indexes on memory_units temporal date fields for fast temporal retrieval
|
||||
|
||||
Revision ID: b3c4d5e6f7g8
|
||||
Revises: c1a2b3d4e5f6
|
||||
Create Date: 2026-03-02
|
||||
|
||||
The temporal retrieval entry-point query filters memory_units by occurred_start,
|
||||
occurred_end, and mentioned_at using OR conditions. Without dedicated indexes the
|
||||
planner falls back to a sequential scan of all bank rows after applying the
|
||||
(bank_id, fact_type) index, then re-checks each date field.
|
||||
|
||||
These three partial indexes give the planner bitmap-index scan options for the
|
||||
three most common date predicates, dramatically reducing the row set before any
|
||||
embedding computation is required.
|
||||
|
||||
All indexes are created CONCURRENTLY so the migration does not block writes on
|
||||
memory_units during production deployments. CONCURRENTLY requires running outside
|
||||
a transaction block; see migrations.py for how this is handled safely.
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
revision: str = "b3c4d5e6f7g8"
|
||||
down_revision: str | Sequence[str] | None = "c1a2b3d4e5f6"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _get_schema_prefix() -> str:
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
# Partial index on occurred_start (covers "occurred_start BETWEEN $4 AND $5")
|
||||
op.execute("COMMIT")
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_units_bank_occurred_start "
|
||||
f"ON {schema}memory_units(bank_id, fact_type, occurred_start) "
|
||||
f"WHERE occurred_start IS NOT NULL"
|
||||
)
|
||||
# Partial index on occurred_end (covers "occurred_end BETWEEN $4 AND $5")
|
||||
op.execute("COMMIT")
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_units_bank_occurred_end "
|
||||
f"ON {schema}memory_units(bank_id, fact_type, occurred_end) "
|
||||
f"WHERE occurred_end IS NOT NULL"
|
||||
)
|
||||
# Partial index on mentioned_at (covers "mentioned_at BETWEEN $4 AND $5")
|
||||
op.execute("COMMIT")
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_units_bank_mentioned_at "
|
||||
f"ON {schema}memory_units(bank_id, fact_type, mentioned_at) "
|
||||
f"WHERE mentioned_at IS NOT NULL"
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
op.execute("COMMIT")
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_units_bank_mentioned_at")
|
||||
op.execute("COMMIT")
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_units_bank_occurred_end")
|
||||
op.execute("COMMIT")
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_units_bank_occurred_start")
|
||||
+34
@@ -0,0 +1,34 @@
|
||||
"""Backfill observation_scopes column if missing.
|
||||
|
||||
This migration ensures observation_scopes exists even on databases that had
|
||||
revision z1u2v3w4x5y6 applied when it referred to the old text_signals migration
|
||||
(before it was renamed to a2b3c4d5e6f7). The ADD COLUMN IF NOT EXISTS makes this
|
||||
a no-op on databases that already have the column.
|
||||
|
||||
Revision ID: b4c5d6e7f8a9
|
||||
Revises: a2b3c4d5e6f7
|
||||
Create Date: 2026-03-02
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
revision: str = "b4c5d6e7f8a9"
|
||||
down_revision: str | Sequence[str] | None = "a2b3c4d5e6f7"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _get_schema_prefix() -> str:
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
op.execute(f"ALTER TABLE {schema}memory_units ADD COLUMN IF NOT EXISTS observation_scopes JSONB")
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
pass # intentionally no-op — safe to leave the column in place
|
||||
+46
@@ -0,0 +1,46 @@
|
||||
"""Enable pg_trgm extension and add GIN trigram index on entities.canonical_name
|
||||
|
||||
Revision ID: c1a2b3d4e5f6
|
||||
Revises: b4c5d6e7f8a9
|
||||
Create Date: 2026-03-02
|
||||
|
||||
Index is created CONCURRENTLY so the migration does not block writes on entities
|
||||
during production deployments. CONCURRENTLY requires running outside a transaction
|
||||
block; see migrations.py for how this is handled safely.
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
revision: str = "c1a2b3d4e5f6"
|
||||
down_revision: str | Sequence[str] | None = "b4c5d6e7f8a9"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _get_schema_prefix() -> str:
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
# pg_trgm ships with every standard PostgreSQL installation as a contrib module.
|
||||
# It enables fast similarity lookups via GIN indexes, used for entity name matching.
|
||||
op.execute("CREATE EXTENSION IF NOT EXISTS pg_trgm")
|
||||
|
||||
schema = _get_schema_prefix()
|
||||
# GIN index on canonical_name enables sub-millisecond trigram similarity queries
|
||||
# (% operator, similarity()) instead of full-table scans across all bank entities.
|
||||
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)"
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
op.execute("COMMIT")
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}entities_canonical_name_trgm_idx")
|
||||
# Note: not dropping pg_trgm extension as other indexes may depend on it
|
||||
+83
@@ -0,0 +1,83 @@
|
||||
"""Add covering and composite indexes to speed up link expansion graph retrieval.
|
||||
|
||||
Two indexes target the two bottlenecks identified by EXPLAIN ANALYZE on a 17M-row
|
||||
memory_links table:
|
||||
|
||||
1. idx_memory_links_to_type_weight (to_unit_id, link_type, weight DESC)
|
||||
The semantic incoming direction — finding facts that consider seeds as their
|
||||
nearest neighbour — currently hits an expensive BitmapAnd of two separate
|
||||
bitmap scans (to_unit_id bitmap ∩ link_type bitmap). A composite index
|
||||
on (to_unit_id, link_type) turns this into a single index scan and reduces
|
||||
latency from ~36 ms to < 5 ms per query.
|
||||
|
||||
2. idx_memory_links_entity_covering (from_unit_id) INCLUDE (to_unit_id, entity_id)
|
||||
WHERE link_type = 'entity'
|
||||
The entity co-occurrence expansion uses COUNT(DISTINCT ml.entity_id) and
|
||||
joins on ml.to_unit_id. Without a covering index the planner must read
|
||||
~2 500 heap pages to fetch entity_id and to_unit_id after the bitmap index
|
||||
scan, adding ~230 ms of random I/O. INCLUDE adds those two columns to the
|
||||
index leaf pages so the entire query can be served from the index (index-only
|
||||
scan), eliminating the heap reads entirely.
|
||||
Partial index (WHERE link_type = 'entity') keeps index size ~40 % smaller.
|
||||
|
||||
Both indexes are created with CONCURRENTLY so the migration does not block
|
||||
concurrent reads or writes on memory_links. CONCURRENTLY requires running
|
||||
outside a transaction block, so the migration emits an explicit COMMIT before
|
||||
each statement and uses IF NOT EXISTS for idempotency.
|
||||
|
||||
Revision ID: d2e3f4a5b6c7
|
||||
Revises: b3c4d5e6f7g8
|
||||
Create Date: 2026-03-02
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
revision: str = "d2e3f4a5b6c7"
|
||||
down_revision: str | Sequence[str] | None = "b3c4d5e6f7g8"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _get_schema_prefix() -> str:
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
|
||||
# CREATE INDEX CONCURRENTLY cannot run inside a transaction block.
|
||||
# Commit the current Alembic transaction, then issue each CONCURRENTLY
|
||||
# statement in its own implicit autocommit transaction.
|
||||
# IF NOT EXISTS makes each statement idempotent if the migration is retried.
|
||||
|
||||
# Index for the semantic *incoming* direction in link_expansion_retrieval.py.
|
||||
# Replaces the BitmapAnd of idx_memory_links_to_unit ∩ idx_memory_links_link_type
|
||||
# with a single composite index scan.
|
||||
op.execute("COMMIT")
|
||||
op.execute(
|
||||
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_links_to_type_weight "
|
||||
f"ON {schema}memory_links(to_unit_id, link_type, weight DESC)"
|
||||
)
|
||||
|
||||
# Covering index for entity co-occurrence expansion.
|
||||
# Enables an index-only scan: entity_id and to_unit_id are read from the
|
||||
# index leaf pages instead of the heap, eliminating ~2 500 random heap-page
|
||||
# reads per expansion query.
|
||||
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'"
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
op.execute("COMMIT")
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_links_entity_covering")
|
||||
op.execute("COMMIT")
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_links_to_type_weight")
|
||||
@@ -0,0 +1,62 @@
|
||||
"""Add webhooks table and next_retry_at to async_operations.
|
||||
|
||||
Webhook deliveries are handled as async_operations tasks (operation_type='webhook_delivery')
|
||||
rather than a dedicated webhook_deliveries table.
|
||||
|
||||
Revision ID: e4f5a6b7c8d9
|
||||
Revises: d2e3f4a5b6c7
|
||||
Create Date: 2026-03-04
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
revision: str = "e4f5a6b7c8d9"
|
||||
down_revision: str | Sequence[str] | None = "d2e3f4a5b6c7"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _get_schema_prefix() -> str:
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
|
||||
op.execute(
|
||||
f"""
|
||||
CREATE TABLE IF NOT EXISTS {schema}webhooks (
|
||||
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
||||
bank_id TEXT,
|
||||
url TEXT NOT NULL,
|
||||
secret TEXT,
|
||||
event_types TEXT[] NOT NULL DEFAULT '{{}}',
|
||||
enabled BOOLEAN NOT NULL DEFAULT TRUE,
|
||||
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
|
||||
updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
|
||||
)
|
||||
"""
|
||||
)
|
||||
|
||||
# Index for bank-scoped webhook lookup
|
||||
op.execute(f"CREATE INDEX IF NOT EXISTS idx_webhooks_bank_id ON {schema}webhooks(bank_id)")
|
||||
|
||||
# Add next_retry_at to async_operations for task-owned retry scheduling
|
||||
op.execute(f"ALTER TABLE {schema}async_operations ADD COLUMN IF NOT EXISTS next_retry_at TIMESTAMPTZ NULL")
|
||||
|
||||
# Index for polling: status + next_retry_at
|
||||
op.execute(
|
||||
f"CREATE INDEX IF NOT EXISTS idx_async_operations_status_retry "
|
||||
f"ON {schema}async_operations(status, next_retry_at)"
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema}idx_async_operations_status_retry")
|
||||
op.execute(f"ALTER TABLE {schema}async_operations DROP COLUMN IF EXISTS next_retry_at")
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema}idx_webhooks_bank_id")
|
||||
op.execute(f"DROP TABLE IF EXISTS {schema}webhooks")
|
||||
@@ -0,0 +1,33 @@
|
||||
"""Add http_config JSONB column to webhooks table.
|
||||
|
||||
Stores HTTP delivery configuration (method, timeout, headers, params) as a
|
||||
single JSONB column rather than separate columns.
|
||||
|
||||
Revision ID: f7g8h9i0j1k2
|
||||
Revises: e4f5a6b7c8d9
|
||||
Create Date: 2026-03-04
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
revision: str = "f7g8h9i0j1k2"
|
||||
down_revision: str | Sequence[str] | None = "e4f5a6b7c8d9"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _get_schema_prefix() -> str:
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
op.execute(f"ALTER TABLE {schema}webhooks ADD COLUMN IF NOT EXISTS http_config JSONB NOT NULL DEFAULT '{{}}'")
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
op.execute(f"ALTER TABLE {schema}webhooks DROP COLUMN IF EXISTS http_config")
|
||||
+35
@@ -0,0 +1,35 @@
|
||||
"""Add observation_scopes column to memory_units table
|
||||
|
||||
Revision ID: z1u2v3w4x5y6
|
||||
Revises: a1b2c3d4e5f6
|
||||
Create Date: 2026-02-25
|
||||
|
||||
Adds observation_scopes JSONB column to memory_units to control how observations
|
||||
are scoped during consolidation. Accepts "per_tag", "combined", or an explicit
|
||||
list of tag-set lists for custom multi-pass consolidation.
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
revision: str = "z1u2v3w4x5y6"
|
||||
down_revision: str | Sequence[str] | None = "a1b2c3d4e5f6"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _get_schema_prefix() -> str:
|
||||
"""Get schema prefix for table names (required for multi-tenant support)."""
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
op.execute(f"ALTER TABLE {schema}memory_units ADD COLUMN IF NOT EXISTS observation_scopes JSONB")
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
schema = _get_schema_prefix()
|
||||
op.execute(f"ALTER TABLE {schema}memory_units DROP COLUMN IF EXISTS observation_scopes")
|
||||
File diff suppressed because it is too large
Load Diff
@@ -166,7 +166,6 @@ def create_mcp_server(memory: MemoryEngine, multi_bank: bool = True) -> FastMCP:
|
||||
api_key_id_resolver=get_current_api_key_id, # Propagate api_key_id for usage metering
|
||||
include_bank_id_param=multi_bank,
|
||||
tools=base_tools,
|
||||
retain_fire_and_forget=False, # HTTP MCP supports sync/async modes
|
||||
)
|
||||
|
||||
register_mcp_tools(mcp, memory, config)
|
||||
@@ -332,7 +331,7 @@ class MCPMiddleware:
|
||||
auth_tenant_id = auth_context.tenant_id
|
||||
auth_api_key_id = auth_context.api_key_id
|
||||
except AuthenticationError as e:
|
||||
await self._send_error(send, 401, str(e))
|
||||
await self._send_error(send, 401, str(e), extra_headers=e.headers)
|
||||
return
|
||||
|
||||
# Set schema from tenant context so downstream DB queries use the correct schema
|
||||
@@ -414,14 +413,17 @@ class MCPMiddleware:
|
||||
if schema_token is not None:
|
||||
_current_schema.reset(schema_token)
|
||||
|
||||
async def _send_error(self, send, status: int, message: str):
|
||||
async def _send_error(self, send, status: int, message: str, extra_headers: dict[str, str] | None = None):
|
||||
"""Send an error response."""
|
||||
body = json.dumps({"error": message}).encode()
|
||||
headers = [(b"content-type", b"application/json")]
|
||||
for key, value in (extra_headers or {}).items():
|
||||
headers.append((key.encode(), value.encode()))
|
||||
await send(
|
||||
{
|
||||
"type": "http.response.start",
|
||||
"status": status,
|
||||
"headers": [(b"content-type", b"application/json")],
|
||||
"headers": headers,
|
||||
}
|
||||
)
|
||||
await send(
|
||||
|
||||
@@ -252,6 +252,9 @@ ENV_LLM_VERTEXAI_PROJECT_ID = "HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID"
|
||||
ENV_LLM_VERTEXAI_REGION = "HINDSIGHT_API_LLM_VERTEXAI_REGION"
|
||||
ENV_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY = "HINDSIGHT_API_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY"
|
||||
|
||||
# Gemini safety settings
|
||||
ENV_LLM_GEMINI_SAFETY_SETTINGS = "HINDSIGHT_API_LLM_GEMINI_SAFETY_SETTINGS"
|
||||
|
||||
# Retain settings
|
||||
ENV_RETAIN_MAX_COMPLETION_TOKENS = "HINDSIGHT_API_RETAIN_MAX_COMPLETION_TOKENS"
|
||||
ENV_RETAIN_CHUNK_SIZE = "HINDSIGHT_API_RETAIN_CHUNK_SIZE"
|
||||
@@ -260,6 +263,7 @@ ENV_RETAIN_EXTRACTION_MODE = "HINDSIGHT_API_RETAIN_EXTRACTION_MODE"
|
||||
ENV_RETAIN_MISSION = "HINDSIGHT_API_RETAIN_MISSION"
|
||||
ENV_RETAIN_CUSTOM_INSTRUCTIONS = "HINDSIGHT_API_RETAIN_CUSTOM_INSTRUCTIONS"
|
||||
ENV_RETAIN_BATCH_TOKENS = "HINDSIGHT_API_RETAIN_BATCH_TOKENS"
|
||||
ENV_RETAIN_ENTITY_LOOKUP = "HINDSIGHT_API_RETAIN_ENTITY_LOOKUP"
|
||||
ENV_RETAIN_BATCH_ENABLED = "HINDSIGHT_API_RETAIN_BATCH_ENABLED"
|
||||
ENV_RETAIN_BATCH_POLL_INTERVAL_SECONDS = "HINDSIGHT_API_RETAIN_BATCH_POLL_INTERVAL_SECONDS"
|
||||
|
||||
@@ -290,6 +294,12 @@ ENV_CONSOLIDATION_LLM_BATCH_SIZE = "HINDSIGHT_API_CONSOLIDATION_LLM_BATCH_SIZE"
|
||||
ENV_CONSOLIDATION_MAX_TOKENS = "HINDSIGHT_API_CONSOLIDATION_MAX_TOKENS"
|
||||
ENV_OBSERVATIONS_MISSION = "HINDSIGHT_API_OBSERVATIONS_MISSION"
|
||||
|
||||
# Webhook configuration (global, static - server-level only)
|
||||
ENV_WEBHOOK_URL = "HINDSIGHT_API_WEBHOOK_URL"
|
||||
ENV_WEBHOOK_SECRET = "HINDSIGHT_API_WEBHOOK_SECRET"
|
||||
ENV_WEBHOOK_EVENT_TYPES = "HINDSIGHT_API_WEBHOOK_EVENT_TYPES"
|
||||
ENV_WEBHOOK_DELIVERY_POLL_INTERVAL_SECONDS = "HINDSIGHT_API_WEBHOOK_DELIVERY_POLL_INTERVAL_SECONDS"
|
||||
|
||||
# Optimization flags
|
||||
ENV_SKIP_LLM_VERIFICATION = "HINDSIGHT_API_SKIP_LLM_VERIFICATION"
|
||||
ENV_LAZY_RERANKER = "HINDSIGHT_API_LAZY_RERANKER"
|
||||
@@ -314,6 +324,7 @@ ENV_WORKER_CONSOLIDATION_MAX_SLOTS = "HINDSIGHT_API_WORKER_CONSOLIDATION_MAX_SLO
|
||||
|
||||
# Reflect agent settings
|
||||
ENV_REFLECT_MAX_ITERATIONS = "HINDSIGHT_API_REFLECT_MAX_ITERATIONS"
|
||||
ENV_REFLECT_MAX_CONTEXT_TOKENS = "HINDSIGHT_API_REFLECT_MAX_CONTEXT_TOKENS"
|
||||
ENV_REFLECT_MISSION = "HINDSIGHT_API_REFLECT_MISSION"
|
||||
|
||||
# Disposition settings
|
||||
@@ -351,6 +362,9 @@ DEFAULT_LLM_VERTEXAI_PROJECT_ID = None # Required for Vertex AI
|
||||
DEFAULT_LLM_VERTEXAI_REGION = "us-central1"
|
||||
DEFAULT_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY = None # Optional, uses ADC if not set
|
||||
|
||||
# Gemini safety settings defaults
|
||||
DEFAULT_LLM_GEMINI_SAFETY_SETTINGS = None # None = use Gemini default safety settings
|
||||
|
||||
DEFAULT_EMBEDDINGS_PROVIDER = "local"
|
||||
DEFAULT_EMBEDDINGS_LOCAL_MODEL = "BAAI/bge-small-en-v1.5"
|
||||
DEFAULT_EMBEDDINGS_LOCAL_FORCE_CPU = False # Force CPU mode for local embeddings (avoids MPS/XPC issues on macOS)
|
||||
@@ -415,6 +429,7 @@ RETAIN_EXTRACTION_MODES = ("concise", "verbose", "custom") # Allowed extraction
|
||||
DEFAULT_RETAIN_MISSION = None # Declarative spec of what to retain (injected into any extraction mode)
|
||||
DEFAULT_RETAIN_CUSTOM_INSTRUCTIONS = None # Custom extraction guidelines (only used when mode="custom")
|
||||
DEFAULT_RETAIN_BATCH_TOKENS = 10_000 # ~40KB of text # Max chars per sub-batch for async retain auto-splitting
|
||||
DEFAULT_RETAIN_ENTITY_LOOKUP = "trigram" # "full" or "trigram"
|
||||
DEFAULT_RETAIN_BATCH_ENABLED = False # Use LLM Batch API for fact extraction (only when async=True)
|
||||
DEFAULT_RETAIN_BATCH_POLL_INTERVAL_SECONDS = 60 # Batch API polling interval in seconds
|
||||
|
||||
@@ -453,6 +468,7 @@ DEFAULT_WORKER_CONSOLIDATION_MAX_SLOTS = 2 # Max concurrent consolidation tasks
|
||||
|
||||
# Reflect agent settings
|
||||
DEFAULT_REFLECT_MAX_ITERATIONS = 10 # Max tool call iterations before forcing response
|
||||
DEFAULT_REFLECT_MAX_CONTEXT_TOKENS = 100_000 # Max accumulated context tokens before forcing final prompt
|
||||
|
||||
# Disposition defaults (None = not set, fall back to bank DB value or 3)
|
||||
DEFAULT_DISPOSITION_SKEPTICISM = None
|
||||
@@ -487,6 +503,12 @@ Use this tool PROACTIVELY to:
|
||||
# Default embedding dimension (used by initial migration, adjusted at runtime)
|
||||
EMBEDDING_DIMENSION = DEFAULT_EMBEDDING_DIMENSION
|
||||
|
||||
# Webhook configuration defaults
|
||||
DEFAULT_WEBHOOK_URL = None # None = no global webhook configured
|
||||
DEFAULT_WEBHOOK_SECRET = None # None = no signing
|
||||
DEFAULT_WEBHOOK_EVENT_TYPES = "consolidation.completed" # Comma-separated; default = all supported events
|
||||
DEFAULT_WEBHOOK_DELIVERY_POLL_INTERVAL_SECONDS = 30 # How often to poll for pending deliveries
|
||||
|
||||
|
||||
class JsonFormatter(logging.Formatter):
|
||||
"""JSON formatter for structured logging.
|
||||
@@ -563,6 +585,9 @@ class HindsightConfig:
|
||||
llm_vertexai_region: str
|
||||
llm_vertexai_service_account_key: str | None
|
||||
|
||||
# Gemini safety settings (None = use Gemini defaults; list of dicts with category/threshold)
|
||||
llm_gemini_safety_settings: list | None
|
||||
|
||||
# Per-operation LLM configuration (None = use default LLM config)
|
||||
retain_llm_provider: str | None
|
||||
retain_llm_api_key: str | None
|
||||
@@ -660,6 +685,7 @@ class HindsightConfig:
|
||||
retain_batch_tokens: int
|
||||
retain_batch_enabled: bool
|
||||
retain_batch_poll_interval_seconds: int
|
||||
retain_entity_lookup: str # "full" or "trigram"
|
||||
|
||||
# File storage (static - server-level only)
|
||||
file_storage_type: str # "native" (PostgreSQL) or "s3" (S3-compatible)
|
||||
@@ -688,6 +714,13 @@ class HindsightConfig:
|
||||
consolidation_max_tokens: int
|
||||
observations_mission: str | None
|
||||
|
||||
# Entity labels (controlled vocabulary of key:value classification labels extracted at retain time)
|
||||
# List of label group dicts: [{key, description, type, optional, values: [{value, description}]}]
|
||||
entity_labels: list | None
|
||||
# Whether to extract regular named entities alongside entity labels (default: True)
|
||||
# When False: only label entities are extracted (or no entities at all if no labels configured)
|
||||
entities_allow_free_form: bool
|
||||
|
||||
# Reflect agent settings
|
||||
reflect_mission: str | None
|
||||
|
||||
@@ -720,6 +753,7 @@ class HindsightConfig:
|
||||
|
||||
# Reflect agent settings
|
||||
reflect_max_iterations: int
|
||||
reflect_max_context_tokens: int
|
||||
|
||||
# OpenTelemetry tracing configuration
|
||||
otel_traces_enabled: bool
|
||||
@@ -728,6 +762,12 @@ class HindsightConfig:
|
||||
otel_service_name: str
|
||||
otel_deployment_environment: str
|
||||
|
||||
# Webhook configuration (static - server-level only, not per-bank)
|
||||
webhook_url: str | None # Global webhook URL (None = disabled)
|
||||
webhook_secret: str | None # HMAC signing secret (None = unsigned)
|
||||
webhook_event_types: list[str] # Event types to deliver globally
|
||||
webhook_delivery_poll_interval_seconds: int # How often the delivery worker polls
|
||||
|
||||
# Class-level sets for configuration categorization
|
||||
|
||||
# CREDENTIAL_FIELDS: Never exposed via API, never configurable per-tenant/bank
|
||||
@@ -767,6 +807,9 @@ class HindsightConfig:
|
||||
"retain_extraction_mode",
|
||||
"retain_mission",
|
||||
"retain_custom_instructions",
|
||||
# Entity labels (controlled vocabulary for entity classification)
|
||||
"entity_labels",
|
||||
"entities_allow_free_form",
|
||||
# Consolidation settings
|
||||
"enable_observations",
|
||||
"observations_mission",
|
||||
@@ -776,6 +819,8 @@ class HindsightConfig:
|
||||
"disposition_skepticism",
|
||||
"disposition_literalism",
|
||||
"disposition_empathy",
|
||||
# Gemini safety settings (controls content filtering for Gemini/VertexAI providers)
|
||||
"llm_gemini_safety_settings",
|
||||
}
|
||||
|
||||
@property
|
||||
@@ -896,6 +941,8 @@ class HindsightConfig:
|
||||
llm_vertexai_region=os.getenv(ENV_LLM_VERTEXAI_REGION, DEFAULT_LLM_VERTEXAI_REGION),
|
||||
llm_vertexai_service_account_key=os.getenv(ENV_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY)
|
||||
or DEFAULT_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY,
|
||||
# Gemini safety settings (JSON-encoded list of {category, threshold} dicts)
|
||||
llm_gemini_safety_settings=json.loads(os.getenv(ENV_LLM_GEMINI_SAFETY_SETTINGS, "null")),
|
||||
# Per-operation LLM config (None = use default)
|
||||
retain_llm_provider=os.getenv(ENV_RETAIN_LLM_PROVIDER) or None,
|
||||
retain_llm_api_key=os.getenv(ENV_RETAIN_LLM_API_KEY) or None,
|
||||
@@ -1071,6 +1118,7 @@ class HindsightConfig:
|
||||
retain_mission=os.getenv(ENV_RETAIN_MISSION) or DEFAULT_RETAIN_MISSION,
|
||||
retain_custom_instructions=os.getenv(ENV_RETAIN_CUSTOM_INSTRUCTIONS) or DEFAULT_RETAIN_CUSTOM_INSTRUCTIONS,
|
||||
retain_batch_tokens=int(os.getenv(ENV_RETAIN_BATCH_TOKENS, str(DEFAULT_RETAIN_BATCH_TOKENS))),
|
||||
retain_entity_lookup=os.getenv(ENV_RETAIN_ENTITY_LOOKUP, DEFAULT_RETAIN_ENTITY_LOOKUP),
|
||||
retain_batch_enabled=os.getenv(ENV_RETAIN_BATCH_ENABLED, str(DEFAULT_RETAIN_BATCH_ENABLED)).lower()
|
||||
== "true",
|
||||
retain_batch_poll_interval_seconds=int(
|
||||
@@ -1115,6 +1163,8 @@ class HindsightConfig:
|
||||
os.getenv(ENV_CONSOLIDATION_MAX_TOKENS, str(DEFAULT_CONSOLIDATION_MAX_TOKENS))
|
||||
),
|
||||
observations_mission=os.getenv(ENV_OBSERVATIONS_MISSION) or DEFAULT_OBSERVATIONS_MISSION,
|
||||
entity_labels=None,
|
||||
entities_allow_free_form=True,
|
||||
# Database migrations
|
||||
run_migrations_on_startup=os.getenv(ENV_RUN_MIGRATIONS_ON_STARTUP, "true").lower() == "true",
|
||||
# Database connection pool
|
||||
@@ -1134,6 +1184,9 @@ class HindsightConfig:
|
||||
),
|
||||
# Reflect agent settings
|
||||
reflect_max_iterations=int(os.getenv(ENV_REFLECT_MAX_ITERATIONS, str(DEFAULT_REFLECT_MAX_ITERATIONS))),
|
||||
reflect_max_context_tokens=int(
|
||||
os.getenv(ENV_REFLECT_MAX_CONTEXT_TOKENS, str(DEFAULT_REFLECT_MAX_CONTEXT_TOKENS))
|
||||
),
|
||||
reflect_mission=os.getenv(ENV_REFLECT_MISSION) or None,
|
||||
# Disposition settings (None = fall back to DB value)
|
||||
disposition_skepticism=int(os.getenv(ENV_DISPOSITION_SKEPTICISM))
|
||||
@@ -1152,6 +1205,20 @@ class HindsightConfig:
|
||||
otel_exporter_otlp_headers=os.getenv(ENV_OTEL_EXPORTER_OTLP_HEADERS) or None,
|
||||
otel_service_name=os.getenv(ENV_OTEL_SERVICE_NAME, DEFAULT_OTEL_SERVICE_NAME),
|
||||
otel_deployment_environment=os.getenv(ENV_OTEL_DEPLOYMENT_ENVIRONMENT, DEFAULT_OTEL_DEPLOYMENT_ENVIRONMENT),
|
||||
# Webhook configuration (static, server-level only)
|
||||
webhook_url=os.getenv(ENV_WEBHOOK_URL) or DEFAULT_WEBHOOK_URL,
|
||||
webhook_secret=os.getenv(ENV_WEBHOOK_SECRET) or DEFAULT_WEBHOOK_SECRET,
|
||||
webhook_event_types=[
|
||||
t.strip()
|
||||
for t in os.getenv(ENV_WEBHOOK_EVENT_TYPES, DEFAULT_WEBHOOK_EVENT_TYPES).split(",")
|
||||
if t.strip()
|
||||
],
|
||||
webhook_delivery_poll_interval_seconds=int(
|
||||
os.getenv(
|
||||
ENV_WEBHOOK_DELIVERY_POLL_INTERVAL_SECONDS,
|
||||
str(DEFAULT_WEBHOOK_DELIVERY_POLL_INTERVAL_SECONDS),
|
||||
)
|
||||
),
|
||||
)
|
||||
config.validate()
|
||||
return config
|
||||
|
||||
@@ -9,6 +9,10 @@ Observations are stored in memory_units with fact_type='observation' and include
|
||||
- proof_count: Number of supporting memories
|
||||
- source_memory_ids: Array of memory UUIDs that contribute to this observation
|
||||
- history: JSONB tracking changes over time
|
||||
|
||||
NOTE: Observations are distinct from mental models (pinned reflections).
|
||||
- Observations: auto-generated bottom-up by this engine from raw facts (memory_units table, fact_type='observation')
|
||||
- Mental models: user-defined queries stored in the mental_models table, refreshed on demand via reflect
|
||||
"""
|
||||
|
||||
import json
|
||||
@@ -17,6 +21,7 @@ import time
|
||||
import uuid
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime, timezone
|
||||
from itertools import combinations
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from pydantic import BaseModel
|
||||
@@ -66,6 +71,42 @@ class _BatchLLMResult:
|
||||
prompt_chars: int = 0
|
||||
|
||||
|
||||
@dataclass
|
||||
class _SourceAggregation:
|
||||
"""Fields inherited by an observation from its source memories."""
|
||||
|
||||
event_date: datetime | None
|
||||
occurred_start: datetime | None
|
||||
occurred_end: datetime | None
|
||||
mentioned_at: datetime | None
|
||||
tags: list[str]
|
||||
|
||||
|
||||
def _aggregate_source_fields(source_mems: list[dict[str, Any]], tags: list[str] | None = None) -> _SourceAggregation:
|
||||
"""Compute the observation fields inherited from a set of source memories.
|
||||
|
||||
Temporal aggregation rules:
|
||||
- ``event_date`` — earliest across sources (min)
|
||||
- ``occurred_start`` — earliest across sources (min)
|
||||
- ``occurred_end`` — latest across sources (max)
|
||||
- ``mentioned_at`` — latest across sources (max)
|
||||
|
||||
Fields remain ``None`` when no source memory carries that information, so
|
||||
observations are never stamped with an artificial timestamp.
|
||||
|
||||
``tags`` defaults to those of the first source memory when not explicitly
|
||||
provided (all memories in a consolidation batch share the same tag set).
|
||||
"""
|
||||
effective_tags = tags if tags is not None else (source_mems[0].get("tags") or [] if source_mems else [])
|
||||
return _SourceAggregation(
|
||||
event_date=_min_date(m.get("event_date") for m in source_mems),
|
||||
occurred_start=_min_date(m.get("occurred_start") for m in source_mems),
|
||||
occurred_end=_max_date(m.get("occurred_end") for m in source_mems),
|
||||
mentioned_at=_max_date(m.get("mentioned_at") for m in source_mems),
|
||||
tags=effective_tags,
|
||||
)
|
||||
|
||||
|
||||
class ConsolidationPerfLog:
|
||||
"""Performance logging for consolidation operations."""
|
||||
|
||||
@@ -125,6 +166,11 @@ async def run_consolidation_job(
|
||||
"""
|
||||
# Resolve bank-specific config with hierarchical overrides
|
||||
config = await memory_engine._config_resolver.resolve_full_config(bank_id, request_context)
|
||||
|
||||
# Build a configured LLM wrapper that applies per-bank settings (e.g. safety settings)
|
||||
# to every call without leaking across operations.
|
||||
llm_config = memory_engine._consolidation_llm_config.with_config(config)
|
||||
|
||||
perf = ConsolidationPerfLog(bank_id)
|
||||
max_memories_per_batch = config.consolidation_batch_size
|
||||
llm_batch_size = max(1, config.consolidation_llm_batch_size)
|
||||
@@ -174,11 +220,12 @@ async def run_consolidation_job(
|
||||
perf.log(f"[1] Found {total_count} pending memories to consolidate")
|
||||
|
||||
# Process each memory with individual commits for crash recovery
|
||||
stats = {
|
||||
stats: dict[str, int] = {
|
||||
"memories_processed": 0,
|
||||
"observations_created": 0,
|
||||
"observations_updated": 0,
|
||||
"observations_merged": 0,
|
||||
"observations_deleted": 0,
|
||||
"actions_executed": 0,
|
||||
"skipped": 0,
|
||||
}
|
||||
@@ -193,7 +240,8 @@ async def run_consolidation_job(
|
||||
t0 = time.time()
|
||||
memories = await conn.fetch(
|
||||
f"""
|
||||
SELECT id, text, fact_type, occurred_start, occurred_end, event_date, tags, mentioned_at
|
||||
SELECT id, text, fact_type, occurred_start, occurred_end, event_date, tags, mentioned_at,
|
||||
observation_scopes
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE bank_id = $1
|
||||
AND consolidated_at IS NULL
|
||||
@@ -239,15 +287,89 @@ async def run_consolidation_job(
|
||||
consolidated_tags.update(memory_tags)
|
||||
|
||||
async with pool.acquire() as conn:
|
||||
results = await _process_memory_batch(
|
||||
conn=conn,
|
||||
memory_engine=memory_engine,
|
||||
bank_id=bank_id,
|
||||
memories=llm_batch,
|
||||
request_context=request_context,
|
||||
perf=perf,
|
||||
config=config,
|
||||
)
|
||||
# Determine observation_scopes for this batch. All memories in a batch share
|
||||
# the same tags (enforced by tag_groups), so we only check the first memory.
|
||||
# asyncpg returns JSONB columns as raw JSON strings, so parse if needed.
|
||||
_obs_raw = llm_batch[0].get("observation_scopes") if llm_batch else None
|
||||
_obs_parsed = json.loads(_obs_raw) if isinstance(_obs_raw, str) else _obs_raw
|
||||
|
||||
# Resolve the scope spec into a concrete list[list[str]] (or None for combined).
|
||||
if _obs_parsed == "per_tag":
|
||||
_memory_tags = llm_batch[0].get("tags") or []
|
||||
obs_tags_list = [[tag] for tag in _memory_tags] if _memory_tags else None
|
||||
elif _obs_parsed == "all_combinations":
|
||||
_memory_tags = llm_batch[0].get("tags") or []
|
||||
obs_tags_list = (
|
||||
[
|
||||
list(combo)
|
||||
for r in range(1, len(_memory_tags) + 1)
|
||||
for combo in combinations(_memory_tags, r)
|
||||
]
|
||||
if _memory_tags
|
||||
else None
|
||||
)
|
||||
elif _obs_parsed == "combined" or _obs_parsed is None:
|
||||
obs_tags_list = None # single combined pass (default behaviour)
|
||||
else:
|
||||
# explicit list[list[str]]
|
||||
obs_tags_list = _obs_parsed
|
||||
|
||||
batch_deleted: int = 0
|
||||
if obs_tags_list:
|
||||
# Multi-pass: run one observation consolidation pass per tag set
|
||||
results = []
|
||||
for obs_tags in obs_tags_list:
|
||||
pass_results, pass_deleted = await _process_memory_batch(
|
||||
conn=conn,
|
||||
memory_engine=memory_engine,
|
||||
llm_config=llm_config,
|
||||
bank_id=bank_id,
|
||||
memories=llm_batch,
|
||||
request_context=request_context,
|
||||
perf=perf,
|
||||
config=config,
|
||||
obs_tags_override=obs_tags,
|
||||
)
|
||||
batch_deleted += pass_deleted
|
||||
# Merge results: prefer non-skipped actions
|
||||
if not results:
|
||||
results = pass_results
|
||||
else:
|
||||
for i, (existing, new) in enumerate(zip(results, pass_results)):
|
||||
if existing.get("action") == "skipped" and new.get("action") != "skipped":
|
||||
results[i] = new
|
||||
elif existing.get("action") != "skipped" and new.get("action") != "skipped":
|
||||
# Both did something — combine into "multiple"
|
||||
existing_created = existing.get(
|
||||
"created", 1 if existing.get("action") == "created" else 0
|
||||
)
|
||||
existing_updated = existing.get(
|
||||
"updated", 1 if existing.get("action") == "updated" else 0
|
||||
)
|
||||
new_created = new.get("created", 1 if new.get("action") == "created" else 0)
|
||||
new_updated = new.get("updated", 1 if new.get("action") == "updated" else 0)
|
||||
total = existing_created + existing_updated + new_created + new_updated
|
||||
results[i] = {
|
||||
"action": "multiple",
|
||||
"created": existing_created + new_created,
|
||||
"updated": existing_updated + new_updated,
|
||||
"merged": 0,
|
||||
"total_actions": total,
|
||||
}
|
||||
else:
|
||||
# Normal single pass using the memory's own tags
|
||||
results, batch_deleted = await _process_memory_batch(
|
||||
conn=conn,
|
||||
memory_engine=memory_engine,
|
||||
llm_config=llm_config,
|
||||
bank_id=bank_id,
|
||||
memories=llm_batch,
|
||||
request_context=request_context,
|
||||
perf=perf,
|
||||
config=config,
|
||||
)
|
||||
stats["observations_deleted"] += batch_deleted
|
||||
|
||||
await conn.executemany(
|
||||
f"UPDATE {fq_table('memory_units')} SET consolidated_at = NOW() WHERE id = $1",
|
||||
[(m["id"],) for m in llm_batch],
|
||||
@@ -436,12 +558,14 @@ async def _trigger_mental_model_refreshes(
|
||||
async def _process_memory_batch(
|
||||
conn: "Connection",
|
||||
memory_engine: "MemoryEngine",
|
||||
llm_config: Any,
|
||||
bank_id: str,
|
||||
memories: list[dict[str, Any]],
|
||||
request_context: "RequestContext",
|
||||
perf: ConsolidationPerfLog | None = None,
|
||||
config: Any = None,
|
||||
) -> list[dict[str, Any]]:
|
||||
obs_tags_override: list[str] | None = None,
|
||||
) -> tuple[list[dict[str, Any]], int]:
|
||||
"""
|
||||
Process a batch of memories in a single LLM call.
|
||||
|
||||
@@ -455,18 +579,26 @@ async def _process_memory_batch(
|
||||
Per-fact security: action execution validates each learning_id against the
|
||||
observations that were recalled specifically for that fact, so cross-tag
|
||||
updates cannot occur.
|
||||
|
||||
Args:
|
||||
obs_tags_override: When set, use these tags for observation recall and
|
||||
create/update instead of the memory's own tags. This enables multi-pass
|
||||
consolidation where a single memory can contribute to observations
|
||||
scoped at different tag levels (e.g., user-level vs session-level).
|
||||
"""
|
||||
import asyncio
|
||||
|
||||
# 1. Parallel recalls — one per fact
|
||||
# When obs_tags_override is set, use it as the observation scope for all facts.
|
||||
t0 = time.time()
|
||||
observation_scope_tags = obs_tags_override if obs_tags_override is not None else None
|
||||
recall_tasks = [
|
||||
_find_related_observations(
|
||||
memory_engine=memory_engine,
|
||||
bank_id=bank_id,
|
||||
query=m["text"],
|
||||
request_context=request_context,
|
||||
tags=m.get("tags") or [],
|
||||
tags=observation_scope_tags if observation_scope_tags is not None else (m.get("tags") or []),
|
||||
)
|
||||
for m in memories
|
||||
]
|
||||
@@ -495,7 +627,7 @@ async def _process_memory_batch(
|
||||
# 3. Single LLM call
|
||||
t0 = time.time()
|
||||
llm_result = await _consolidate_batch_with_llm(
|
||||
memory_engine=memory_engine,
|
||||
llm_config=llm_config,
|
||||
memories=memories,
|
||||
union_observations=union_observations,
|
||||
union_source_facts=union_source_facts,
|
||||
@@ -510,8 +642,13 @@ async def _process_memory_batch(
|
||||
per_memory_created: set[str] = set()
|
||||
per_memory_updated: set[str] = set()
|
||||
|
||||
# All memories in the batch share the same tag set (enforced by batching)
|
||||
fact_tags = memories[0].get("tags") or [] if memories else []
|
||||
# Determine effective tag scope for observations.
|
||||
# When obs_tags_override is set, use it; otherwise use the memory's own tags.
|
||||
if obs_tags_override is not None:
|
||||
fact_tags = obs_tags_override
|
||||
else:
|
||||
# All memories in the batch share the same tag set (enforced by batching)
|
||||
fact_tags = memories[0].get("tags") or [] if memories else []
|
||||
|
||||
mem_by_id = {str(m["id"]): m for m in memories}
|
||||
|
||||
@@ -519,17 +656,18 @@ async def _process_memory_batch(
|
||||
source_mems = [mem_by_id[fid] for fid in create.source_fact_ids if fid in mem_by_id]
|
||||
if not source_mems:
|
||||
continue
|
||||
agg = _aggregate_source_fields(source_mems, tags=fact_tags)
|
||||
await _execute_create_action(
|
||||
conn=conn,
|
||||
memory_engine=memory_engine,
|
||||
bank_id=bank_id,
|
||||
source_memory_ids=[m["id"] for m in source_mems],
|
||||
text=create.text,
|
||||
source_fact_tags=fact_tags,
|
||||
event_date=_min_date(m.get("event_date") for m in source_mems),
|
||||
occurred_start=_min_date(m.get("occurred_start") for m in source_mems),
|
||||
occurred_end=_max_date(m.get("occurred_end") for m in source_mems),
|
||||
mentioned_at=_max_date(m.get("mentioned_at") for m in source_mems),
|
||||
source_fact_tags=agg.tags,
|
||||
event_date=agg.event_date,
|
||||
occurred_start=agg.occurred_start,
|
||||
occurred_end=agg.occurred_end,
|
||||
mentioned_at=agg.mentioned_at,
|
||||
perf=perf,
|
||||
)
|
||||
for m in source_mems:
|
||||
@@ -546,6 +684,7 @@ async def _process_memory_batch(
|
||||
f"not in any source fact's recall"
|
||||
)
|
||||
continue
|
||||
agg = _aggregate_source_fields(source_mems, tags=fact_tags)
|
||||
await _execute_update_action(
|
||||
conn=conn,
|
||||
memory_engine=memory_engine,
|
||||
@@ -554,15 +693,16 @@ async def _process_memory_batch(
|
||||
observation_id=update.observation_id,
|
||||
new_text=update.text,
|
||||
observations=union_observations,
|
||||
source_fact_tags=fact_tags,
|
||||
source_occurred_start=_min_date(m.get("occurred_start") for m in source_mems),
|
||||
source_occurred_end=_max_date(m.get("occurred_end") for m in source_mems),
|
||||
source_mentioned_at=_max_date(m.get("mentioned_at") for m in source_mems),
|
||||
source_fact_tags=agg.tags,
|
||||
source_occurred_start=agg.occurred_start,
|
||||
source_occurred_end=agg.occurred_end,
|
||||
source_mentioned_at=agg.mentioned_at,
|
||||
perf=perf,
|
||||
)
|
||||
for m in source_mems:
|
||||
per_memory_updated.add(str(m["id"]))
|
||||
|
||||
deleted_count = 0
|
||||
for delete in llm_result.deletes:
|
||||
# Security: the observation must be present in the unioned recall
|
||||
if not any(str(obs.id) == delete.observation_id for obs in union_observations):
|
||||
@@ -571,6 +711,7 @@ async def _process_memory_batch(
|
||||
)
|
||||
continue
|
||||
await _execute_delete_action(conn=conn, bank_id=bank_id, observation_id=delete.observation_id)
|
||||
deleted_count += 1
|
||||
|
||||
# Build per-memory result dicts for the stats tracker in the outer loop
|
||||
results: list[dict[str, Any]] = []
|
||||
@@ -587,7 +728,7 @@ async def _process_memory_batch(
|
||||
else:
|
||||
results.append({"action": "skipped", "reason": "no_durable_knowledge"})
|
||||
|
||||
return results
|
||||
return results, deleted_count
|
||||
|
||||
|
||||
def _min_date(dates: "Any") -> "datetime | None":
|
||||
@@ -854,7 +995,7 @@ def _build_observations_for_llm(
|
||||
|
||||
|
||||
async def _consolidate_batch_with_llm(
|
||||
memory_engine: "MemoryEngine",
|
||||
llm_config: Any,
|
||||
memories: list[dict[str, Any]],
|
||||
union_observations: "list[MemoryFact]",
|
||||
union_source_facts: "dict[str, MemoryFact]",
|
||||
@@ -890,7 +1031,7 @@ async def _consolidate_batch_with_llm(
|
||||
last_exc: Exception | None = None
|
||||
for attempt in range(1, max_attempts + 1):
|
||||
try:
|
||||
response: _ConsolidationBatchResponse = await memory_engine._consolidation_llm_config.call(
|
||||
response: _ConsolidationBatchResponse = await llm_config.call(
|
||||
messages=[{"role": "user", "content": prompt}],
|
||||
response_format=_ConsolidationBatchResponse,
|
||||
scope="consolidation",
|
||||
@@ -936,8 +1077,8 @@ async def _create_observation_directly(
|
||||
# Create the observation as a memory_unit
|
||||
now = datetime.now(timezone.utc)
|
||||
obs_event_date = event_date or now
|
||||
obs_occurred_start = occurred_start or now
|
||||
obs_occurred_end = occurred_end or now
|
||||
obs_occurred_start = occurred_start
|
||||
obs_occurred_end = occurred_end
|
||||
obs_mentioned_at = mentioned_at or now
|
||||
obs_tags = tags or []
|
||||
|
||||
|
||||
@@ -566,7 +566,7 @@ class ZeroEntropyCrossEncoder(CrossEncoderModel):
|
||||
See: https://docs.zeroentropy.dev/models
|
||||
"""
|
||||
|
||||
RERANK_URL = "https://api.zeroentropy.dev/models/rerank"
|
||||
RERANK_URL = "https://api.zeroentropy.dev/v1/models/rerank"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
|
||||
@@ -20,6 +20,7 @@ RETRYABLE_EXCEPTIONS = (
|
||||
asyncpg.exceptions.InterfaceError,
|
||||
asyncpg.exceptions.ConnectionDoesNotExistError,
|
||||
asyncpg.exceptions.TooManyConnectionsError,
|
||||
asyncpg.exceptions.DeadlockDetectedError,
|
||||
OSError,
|
||||
ConnectionError,
|
||||
asyncio.TimeoutError,
|
||||
|
||||
@@ -5,6 +5,10 @@ Uses spaCy for entity extraction and implements resolution logic
|
||||
to disambiguate entities across memory units.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from collections import defaultdict
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import UTC, datetime
|
||||
from difflib import SequenceMatcher
|
||||
|
||||
@@ -12,6 +16,43 @@ import asyncpg
|
||||
|
||||
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
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class _EntityToCreate:
|
||||
"""An entity that needs to be inserted (no matching candidate found)."""
|
||||
|
||||
idx: int
|
||||
name: str
|
||||
event_date: datetime | None
|
||||
|
||||
|
||||
@dataclass
|
||||
class _EntityStat:
|
||||
"""Stat accumulation entry for a resolved entity (post-transaction update)."""
|
||||
|
||||
entity_id: str
|
||||
event_date: datetime | None
|
||||
|
||||
|
||||
@dataclass
|
||||
class _EntityStatAgg:
|
||||
"""Aggregated stats used when flushing pending updates."""
|
||||
|
||||
count: int = 0
|
||||
max_date: datetime | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class _CooccurrencePair:
|
||||
"""A (entity_id_1, entity_id_2) pair observed in a retain batch (for post-txn flush)."""
|
||||
|
||||
entity_id_1: str
|
||||
entity_id_2: str
|
||||
|
||||
|
||||
# Load spaCy model (singleton)
|
||||
_nlp = None
|
||||
@@ -22,14 +63,95 @@ class EntityResolver:
|
||||
Resolves entities to canonical IDs with disambiguation.
|
||||
"""
|
||||
|
||||
def __init__(self, pool: asyncpg.Pool):
|
||||
def __init__(self, pool: asyncpg.Pool, entity_lookup: str = "full"):
|
||||
"""
|
||||
Initialize entity resolver.
|
||||
|
||||
Args:
|
||||
pool: asyncpg connection pool
|
||||
entity_lookup: Lookup strategy — "full" loads all bank entities then
|
||||
matches in Python; "trigram" uses pg_trgm GIN index to fetch only
|
||||
similar candidates per entity name (much faster for large banks).
|
||||
"""
|
||||
self.pool = pool
|
||||
self.entity_lookup = entity_lookup
|
||||
# Keyed by asyncio task id so concurrent retain batches never mix their
|
||||
# pending updates. flush_pending_stats() pops only the calling task's items.
|
||||
self._pending_stats: dict[int, list[_EntityStat]] = {}
|
||||
self._pending_cooccurrences: dict[int, list[_CooccurrencePair]] = {}
|
||||
|
||||
def _task_key(self) -> int:
|
||||
"""Return a unique key for the current asyncio task (or 0 for non-task context)."""
|
||||
task = asyncio.current_task()
|
||||
return id(task) if task is not None else 0
|
||||
|
||||
async def flush_pending_stats(self) -> None:
|
||||
"""
|
||||
Flush accumulated entity stats and co-occurrence counts for the current task.
|
||||
|
||||
Must be called AFTER the retain transaction commits. Pops only the items
|
||||
accumulated by the calling asyncio task so concurrent retain batches never
|
||||
flush each other's uncommitted entity IDs.
|
||||
"""
|
||||
if self.pool is None:
|
||||
return
|
||||
|
||||
key = self._task_key()
|
||||
stats = self._pending_stats.pop(key, [])
|
||||
cooccurrences = self._pending_cooccurrences.pop(key, [])
|
||||
|
||||
if not stats and not cooccurrences:
|
||||
return
|
||||
|
||||
async with acquire_with_retry(self.pool) as conn:
|
||||
if stats:
|
||||
# Aggregate: sum counts and find max date per entity_id.
|
||||
agg: dict[str, _EntityStatAgg] = defaultdict(_EntityStatAgg)
|
||||
for s in stats:
|
||||
entry = agg[s.entity_id]
|
||||
entry.count += 1
|
||||
if s.event_date is not None:
|
||||
entry.max_date = s.event_date if entry.max_date is None else max(entry.max_date, s.event_date)
|
||||
|
||||
# Sort by entity_id so all concurrent workers acquire row locks in
|
||||
# the same order — prevents circular lock dependencies (deadlocks).
|
||||
rows = sorted((eid, a.count, a.max_date) for eid, a in agg.items())
|
||||
await conn.executemany(
|
||||
f"""
|
||||
UPDATE {fq_table("entities")} SET
|
||||
mention_count = mention_count + $2,
|
||||
last_seen = GREATEST(last_seen, $3)
|
||||
WHERE id = $1::uuid
|
||||
""",
|
||||
rows,
|
||||
)
|
||||
|
||||
if cooccurrences:
|
||||
# Aggregate: count occurrences per (entity_id_1, entity_id_2) pair.
|
||||
coo_agg: dict[tuple[str, str], int] = {}
|
||||
for c in cooccurrences:
|
||||
pair = (c.entity_id_1, c.entity_id_2)
|
||||
coo_agg[pair] = coo_agg.get(pair, 0) + 1
|
||||
|
||||
now = datetime.now(UTC)
|
||||
# Sort by (entity_id_1, entity_id_2) for consistent lock ordering.
|
||||
await conn.executemany(
|
||||
f"""
|
||||
INSERT INTO {fq_table("entity_cooccurrences")}
|
||||
(entity_id_1, entity_id_2, cooccurrence_count, last_cooccurred)
|
||||
VALUES ($1, $2, $3, $4)
|
||||
ON CONFLICT (entity_id_1, entity_id_2)
|
||||
DO UPDATE SET
|
||||
cooccurrence_count = {fq_table("entity_cooccurrences")}.cooccurrence_count + EXCLUDED.cooccurrence_count,
|
||||
last_cooccurred = GREATEST({fq_table("entity_cooccurrences")}.last_cooccurred, EXCLUDED.last_cooccurred)
|
||||
""",
|
||||
sorted((e1, e2, count, now) for (e1, e2), count in coo_agg.items()),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _build_labels_lookup(entity_labels: list | None) -> set[str]:
|
||||
"""Build a set of valid 'key:value' entity label strings for fast lookup."""
|
||||
return _build_labels_lookup_from_config(entity_labels)
|
||||
|
||||
async def resolve_entities_batch(
|
||||
self,
|
||||
@@ -38,6 +160,7 @@ class EntityResolver:
|
||||
context: str,
|
||||
unit_event_date,
|
||||
conn=None,
|
||||
entity_labels: list | None = None,
|
||||
) -> list[str]:
|
||||
"""
|
||||
Resolve multiple entities in batch (MUCH faster than sequential).
|
||||
@@ -58,15 +181,34 @@ class EntityResolver:
|
||||
if not entities_data:
|
||||
return []
|
||||
|
||||
taxonomy_lookup = self._build_labels_lookup(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)
|
||||
return await self._resolve_entities_batch_impl(
|
||||
conn, bank_id, entities_data, context, unit_event_date, taxonomy_lookup
|
||||
)
|
||||
else:
|
||||
return await self._resolve_entities_batch_impl(conn, bank_id, entities_data, context, unit_event_date)
|
||||
return await self._resolve_entities_batch_impl(
|
||||
conn, bank_id, entities_data, context, unit_event_date, taxonomy_lookup
|
||||
)
|
||||
|
||||
async def _resolve_entities_batch_impl(
|
||||
self, conn, bank_id: str, entities_data: list[dict], context: str, unit_event_date
|
||||
self,
|
||||
conn,
|
||||
bank_id: str,
|
||||
entities_data: list[dict],
|
||||
context: str,
|
||||
unit_event_date,
|
||||
taxonomy_lookup: set[str] | None = None,
|
||||
) -> list[str]:
|
||||
if self.entity_lookup == "trigram":
|
||||
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)
|
||||
|
||||
async def _resolve_entities_batch_full(
|
||||
self, conn, bank_id: str, entities_data: list[dict], unit_event_date
|
||||
) -> list[str]:
|
||||
"""Original strategy: load all bank entities then match in Python."""
|
||||
# Query ALL candidates for this bank
|
||||
all_entities = await conn.fetch(
|
||||
f"""
|
||||
@@ -130,10 +272,103 @@ class EntityResolver:
|
||||
matching.append((ent_id, canonical_name, metadata, last_seen, mention_count))
|
||||
all_candidates[entity_text] = matching
|
||||
|
||||
return await self._resolve_from_candidates(
|
||||
conn, bank_id, entities_data, unit_event_date, all_candidates, cooccurrence_map
|
||||
)
|
||||
|
||||
async def _resolve_entities_batch_trigram(
|
||||
self, conn, bank_id: str, entities_data: list[dict], unit_event_date
|
||||
) -> list[str]:
|
||||
"""
|
||||
Trigram strategy: fetch only similar candidates per entity name using pg_trgm.
|
||||
|
||||
Instead of loading all bank entities (O(N)), uses a GIN trigram index to fetch
|
||||
only the small set of candidates that are textually similar to each input name.
|
||||
Reduces DB data transfer from 165K rows to ~5-20 rows per entity.
|
||||
"""
|
||||
entity_texts = list(set(e["text"] for e in entities_data))
|
||||
|
||||
# Fetch candidates for all unique entity texts in a single batched query.
|
||||
# The trigram % operator uses the GIN index; the substring conditions cover
|
||||
# exact prefix/suffix matches that trigrams might miss at low similarity.
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT DISTINCT ON (e.id)
|
||||
e.id, e.canonical_name, e.metadata, e.last_seen, e.mention_count,
|
||||
q.query_text
|
||||
FROM unnest($2::text[]) AS q(query_text)
|
||||
JOIN {fq_table("entities")} e ON (
|
||||
e.bank_id = $1
|
||||
AND (
|
||||
e.canonical_name % q.query_text
|
||||
OR LOWER(e.canonical_name) LIKE '%' || LOWER(q.query_text) || '%'
|
||||
OR LOWER(q.query_text) LIKE '%' || LOWER(e.canonical_name) || '%'
|
||||
)
|
||||
)
|
||||
""",
|
||||
bank_id,
|
||||
entity_texts,
|
||||
)
|
||||
|
||||
# Group candidates by query_text
|
||||
all_candidates: dict[str, list] = {t: [] for t in entity_texts}
|
||||
candidate_ids: set = set()
|
||||
for row in rows:
|
||||
query_text = row["query_text"]
|
||||
all_candidates[query_text].append(
|
||||
(row["id"], row["canonical_name"], row["metadata"], row["last_seen"], row["mention_count"])
|
||||
)
|
||||
candidate_ids.add(row["id"])
|
||||
|
||||
# Fetch co-occurrences only for the candidate entities (not all bank entities)
|
||||
cooccurrence_map: dict[str, set[str]] = {}
|
||||
if candidate_ids:
|
||||
candidate_id_list = list(candidate_ids)
|
||||
cooc_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT ec.entity_id_1, ec.entity_id_2
|
||||
FROM {fq_table("entity_cooccurrences")} ec
|
||||
WHERE ec.entity_id_1 = ANY($1::uuid[])
|
||||
OR ec.entity_id_2 = ANY($1::uuid[])
|
||||
""",
|
||||
candidate_id_list,
|
||||
)
|
||||
# Build name lookup for co-occurrence mapping
|
||||
id_to_name = {
|
||||
row["id"]: row["canonical_name"].lower()
|
||||
for cands in all_candidates.values()
|
||||
for row in [{"id": c[0], "canonical_name": c[1]} for c in cands]
|
||||
}
|
||||
for row in cooc_rows:
|
||||
eid1, eid2 = row["entity_id_1"], row["entity_id_2"]
|
||||
if eid1 not in cooccurrence_map:
|
||||
cooccurrence_map[eid1] = set()
|
||||
if eid2 not in cooccurrence_map:
|
||||
cooccurrence_map[eid2] = set()
|
||||
if eid2 in id_to_name:
|
||||
cooccurrence_map[eid1].add(id_to_name[eid2])
|
||||
if eid1 in id_to_name:
|
||||
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
|
||||
)
|
||||
|
||||
async def _resolve_from_candidates(
|
||||
self,
|
||||
conn,
|
||||
bank_id: str,
|
||||
entities_data: list[dict],
|
||||
unit_event_date,
|
||||
all_candidates: dict[str, list],
|
||||
cooccurrence_map: dict[str, set[str]],
|
||||
) -> list[str]:
|
||||
"""Shared scoring + upsert logic used by both lookup strategies."""
|
||||
|
||||
# Resolve each entity using pre-fetched candidates
|
||||
entity_ids = [None] * len(entities_data)
|
||||
entities_to_update = [] # (entity_id, event_date)
|
||||
entities_to_create = [] # (idx, entity_data, event_date)
|
||||
entities_to_update: list[_EntityStat] = []
|
||||
entities_to_create: list[_EntityToCreate] = []
|
||||
|
||||
for idx, entity_data in enumerate(entities_data):
|
||||
entity_text = entity_data["text"]
|
||||
@@ -145,7 +380,7 @@ class EntityResolver:
|
||||
|
||||
if not candidates:
|
||||
# Will create new entity
|
||||
entities_to_create.append((idx, entity_data, entity_event_date))
|
||||
entities_to_create.append(_EntityToCreate(idx=idx, name=entity_text, event_date=entity_event_date))
|
||||
continue
|
||||
|
||||
# Score candidates
|
||||
@@ -189,73 +424,83 @@ class EntityResolver:
|
||||
|
||||
if best_score > threshold:
|
||||
entity_ids[idx] = best_candidate
|
||||
entities_to_update.append((best_candidate, entity_event_date))
|
||||
entities_to_update.append(_EntityStat(entity_id=best_candidate, event_date=entity_event_date))
|
||||
else:
|
||||
entities_to_create.append((idx, entity_data, entity_event_date))
|
||||
entities_to_create.append(
|
||||
_EntityToCreate(idx=idx, name=entity_data["text"], event_date=entity_event_date)
|
||||
)
|
||||
|
||||
# Batch update existing entities
|
||||
if entities_to_update:
|
||||
await conn.executemany(
|
||||
f"""
|
||||
UPDATE {fq_table("entities")} SET
|
||||
mention_count = mention_count + 1,
|
||||
last_seen = $2
|
||||
WHERE id = $1::uuid
|
||||
""",
|
||||
entities_to_update,
|
||||
)
|
||||
# Existing entities: IDs already known from the candidate SELECT above.
|
||||
# No in-transaction UPDATE — mention_count/last_seen are stats deferred to
|
||||
# flush_pending_stats() which the orchestrator calls after the transaction.
|
||||
pending: list[_EntityStat] = list(entities_to_update)
|
||||
|
||||
# Batch create new entities using COPY + INSERT for maximum speed
|
||||
# This handles duplicates via ON CONFLICT and returns all IDs
|
||||
# New entities: INSERT with DO NOTHING to avoid row locks on concurrent races.
|
||||
# ON CONFLICT DO NOTHING returns nothing for rows that conflicted; we handle
|
||||
# that rare case with a fallback SELECT.
|
||||
if entities_to_create:
|
||||
# Group entities by canonical name (lowercase) to handle duplicates within batch
|
||||
# For duplicates, we only insert once and reuse the ID, but track the count
|
||||
unique_entities = {} # lowercase_name -> (entity_data, event_date, [indices])
|
||||
for idx, entity_data, event_date in entities_to_create:
|
||||
name_lower = entity_data["text"].lower()
|
||||
if name_lower not in unique_entities:
|
||||
unique_entities[name_lower] = (entity_data, event_date, [idx])
|
||||
else:
|
||||
# Same entity appears multiple times - add index to list
|
||||
unique_entities[name_lower][2].append(idx)
|
||||
# Group by lowercase name — deduplicate within the batch.
|
||||
@dataclass
|
||||
class _NameGroup:
|
||||
name: str
|
||||
event_date: datetime | None
|
||||
indices: list[int] = field(default_factory=list)
|
||||
|
||||
# Batch insert unique entities and get their IDs
|
||||
# Use a single query with unnest for speed
|
||||
entity_names = []
|
||||
entity_dates = []
|
||||
entity_counts = [] # Track how many times each entity appears in this batch
|
||||
indices_map = [] # Maps result index -> list of original indices
|
||||
groups: dict[str, _NameGroup] = {}
|
||||
for e in entities_to_create:
|
||||
name_lower = e.name.lower()
|
||||
if name_lower not in groups:
|
||||
groups[name_lower] = _NameGroup(name=e.name, event_date=e.event_date)
|
||||
groups[name_lower].indices.append(e.idx)
|
||||
|
||||
for name_lower, (entity_data, event_date, indices) in unique_entities.items():
|
||||
entity_names.append(entity_data["text"])
|
||||
entity_dates.append(event_date)
|
||||
entity_counts.append(len(indices)) # Count of occurrences in this batch
|
||||
indices_map.append(indices)
|
||||
# Sort by lowercase name for deterministic ordering.
|
||||
sorted_groups = sorted(groups.items())
|
||||
entity_names = [g.name for _, g in sorted_groups]
|
||||
entity_dates = [g.event_date for _, g in sorted_groups]
|
||||
|
||||
# Batch INSERT ... ON CONFLICT with RETURNING
|
||||
# Uses the batch count for mention_count instead of always 1
|
||||
rows = await conn.fetch(
|
||||
# INSERT ... ON CONFLICT DO NOTHING — no row lock on already-existing entities.
|
||||
inserted_rows = await conn.fetch(
|
||||
f"""
|
||||
INSERT INTO {fq_table("entities")} (bank_id, canonical_name, first_seen, last_seen, mention_count)
|
||||
SELECT $1, name, event_date, event_date, cnt
|
||||
FROM unnest($2::text[], $3::timestamptz[], $4::int[]) AS t(name, event_date, cnt)
|
||||
SELECT $1, name, COALESCE(event_date, now()), COALESCE(event_date, now()), 1
|
||||
FROM unnest($2::text[], $3::timestamptz[]) AS t(name, event_date)
|
||||
ON CONFLICT (bank_id, LOWER(canonical_name))
|
||||
DO UPDATE SET
|
||||
mention_count = {fq_table("entities")}.mention_count + EXCLUDED.mention_count,
|
||||
last_seen = EXCLUDED.last_seen
|
||||
RETURNING id
|
||||
DO NOTHING
|
||||
RETURNING id, LOWER(canonical_name) AS name_lower
|
||||
""",
|
||||
bank_id,
|
||||
entity_names,
|
||||
entity_dates,
|
||||
entity_counts,
|
||||
)
|
||||
id_by_name: dict[str, str] = {row["name_lower"]: row["id"] for row in inserted_rows}
|
||||
|
||||
# Map returned IDs back to original indices
|
||||
for result_idx, row in enumerate(rows):
|
||||
entity_id = row["id"]
|
||||
for original_idx in indices_map[result_idx]:
|
||||
entity_ids[original_idx] = entity_id
|
||||
# Fallback SELECT for names that conflicted (another worker won the race).
|
||||
missing = [n for n, _ in sorted_groups if n not in id_by_name]
|
||||
if missing:
|
||||
existing_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT id, LOWER(canonical_name) AS name_lower
|
||||
FROM {fq_table("entities")}
|
||||
WHERE bank_id = $1 AND LOWER(canonical_name) = ANY($2::text[])
|
||||
""",
|
||||
bank_id,
|
||||
missing,
|
||||
)
|
||||
for row in existing_rows:
|
||||
id_by_name[row["name_lower"]] = row["id"]
|
||||
|
||||
# Assign entity IDs back and queue for post-txn stats flush.
|
||||
for name_lower, g in sorted_groups:
|
||||
entity_id = id_by_name.get(name_lower)
|
||||
if entity_id:
|
||||
for original_idx in g.indices:
|
||||
entity_ids[original_idx] = entity_id
|
||||
pending.append(_EntityStat(entity_id=entity_id, event_date=g.event_date))
|
||||
|
||||
# Accumulate into the resolver's pending list; the orchestrator flushes
|
||||
# these with await entity_resolver.flush_pending_stats() after the txn.
|
||||
key = self._task_key()
|
||||
self._pending_stats.setdefault(key, []).extend(pending)
|
||||
|
||||
return entity_ids
|
||||
|
||||
@@ -408,7 +653,7 @@ class EntityResolver:
|
||||
entity_id = await conn.fetchval(
|
||||
f"""
|
||||
INSERT INTO {fq_table("entities")} (bank_id, canonical_name, first_seen, last_seen, mention_count)
|
||||
VALUES ($1, $2, $3, $4, 1)
|
||||
VALUES ($1, $2, COALESCE($3, now()), COALESCE($4, now()), 1)
|
||||
ON CONFLICT (bank_id, LOWER(canonical_name))
|
||||
DO UPDATE SET
|
||||
mention_count = {fq_table("entities")}.mention_count + 1,
|
||||
@@ -541,19 +786,14 @@ class EntityResolver:
|
||||
entity_id_1, entity_id_2 = entity_id_2, entity_id_1
|
||||
cooccurrence_pairs.add((entity_id_1, entity_id_2))
|
||||
|
||||
# Batch update co-occurrences
|
||||
# Accumulate co-occurrence pairs for post-transaction flush.
|
||||
# The actual INSERT/UPDATE is deferred to flush_pending_stats() to avoid
|
||||
# row-level lock contention (ON CONFLICT DO UPDATE inside a long transaction
|
||||
# serialises concurrent writers on popular entity pairs).
|
||||
if cooccurrence_pairs:
|
||||
now = datetime.now(UTC)
|
||||
await conn.executemany(
|
||||
f"""
|
||||
INSERT INTO {fq_table("entity_cooccurrences")} (entity_id_1, entity_id_2, cooccurrence_count, last_cooccurred)
|
||||
VALUES ($1, $2, $3, $4)
|
||||
ON CONFLICT (entity_id_1, entity_id_2)
|
||||
DO UPDATE SET
|
||||
cooccurrence_count = {fq_table("entity_cooccurrences")}.cooccurrence_count + 1,
|
||||
last_cooccurred = EXCLUDED.last_cooccurred
|
||||
""",
|
||||
[(e1, e2, 1, now) for e1, e2 in cooccurrence_pairs],
|
||||
key = self._task_key()
|
||||
self._pending_cooccurrences.setdefault(key, []).extend(
|
||||
_CooccurrencePair(entity_id_1=e1, entity_id_2=e2) for e1, e2 in cooccurrence_pairs
|
||||
)
|
||||
|
||||
async def get_units_by_entity(self, entity_id: str, limit: int = 100) -> list[str]:
|
||||
|
||||
@@ -12,6 +12,7 @@ from typing import TYPE_CHECKING, Any
|
||||
if TYPE_CHECKING:
|
||||
from hindsight_api.engine.memory_engine import Budget
|
||||
from hindsight_api.engine.response_models import RecallResult, ReflectResult
|
||||
from hindsight_api.engine.search.tags import TagsMatch
|
||||
from hindsight_api.models import RequestContext
|
||||
|
||||
|
||||
@@ -337,6 +338,8 @@ class MemoryEngineInterface(ABC):
|
||||
bank_id: str,
|
||||
*,
|
||||
search_query: str | None = None,
|
||||
tags: list[str] | None = None,
|
||||
tags_match: "TagsMatch" = "any_strict",
|
||||
limit: int = 100,
|
||||
offset: int = 0,
|
||||
request_context: "RequestContext",
|
||||
@@ -346,7 +349,9 @@ class MemoryEngineInterface(ABC):
|
||||
|
||||
Args:
|
||||
bank_id: The memory bank ID.
|
||||
search_query: Search query.
|
||||
search_query: Case-insensitive substring filter on document ID.
|
||||
tags: Filter by tags.
|
||||
tags_match: How to match tags (any, all, any_strict, all_strict).
|
||||
limit: Maximum results.
|
||||
offset: Pagination offset.
|
||||
request_context: Request context for authentication.
|
||||
|
||||
@@ -124,6 +124,7 @@ def create_llm_provider(
|
||||
vertexai_project_id: str | None = None,
|
||||
vertexai_region: str | None = None,
|
||||
vertexai_credentials: Any = None,
|
||||
gemini_safety_settings: list | None = None,
|
||||
) -> Any: # Returns LLMInterface
|
||||
"""
|
||||
Factory function to create the appropriate LLM provider implementation.
|
||||
@@ -192,6 +193,7 @@ def create_llm_provider(
|
||||
vertexai_project_id=vertexai_project_id,
|
||||
vertexai_region=vertexai_region,
|
||||
vertexai_credentials=vertexai_credentials,
|
||||
gemini_safety_settings=gemini_safety_settings,
|
||||
)
|
||||
|
||||
elif provider_lower == "anthropic":
|
||||
@@ -234,6 +236,7 @@ class LLMProvider:
|
||||
reasoning_effort: str = "low",
|
||||
groq_service_tier: str | None = None,
|
||||
openai_service_tier: str | None = None,
|
||||
gemini_safety_settings: list | None = None,
|
||||
):
|
||||
"""
|
||||
Initialize LLM provider.
|
||||
@@ -246,6 +249,7 @@ class LLMProvider:
|
||||
reasoning_effort: Reasoning effort level for supported providers.
|
||||
groq_service_tier: Groq service tier ("on_demand", "flex", "auto") - from config.
|
||||
openai_service_tier: OpenAI service tier (None or "flex") - from config.
|
||||
gemini_safety_settings: Safety settings for Gemini/VertexAI providers.
|
||||
"""
|
||||
self.provider = provider.lower()
|
||||
self.api_key = api_key
|
||||
@@ -255,6 +259,8 @@ class LLMProvider:
|
||||
# Service tiers from hierarchical config (not env vars)
|
||||
self.groq_service_tier = groq_service_tier
|
||||
self.openai_service_tier = openai_service_tier
|
||||
# Gemini safety settings (instance default; can be overridden per-request via context var)
|
||||
self.gemini_safety_settings = gemini_safety_settings
|
||||
|
||||
# Validate provider
|
||||
valid_providers = [
|
||||
@@ -323,6 +329,18 @@ class LLMProvider:
|
||||
f"model={self.model}, auth={'service_account' if service_account_key else 'ADC'}"
|
||||
)
|
||||
|
||||
# For Gemini/VertexAI providers: read safety settings from global config if not explicitly provided
|
||||
# Use _get_raw_config() to bypass StaticConfigProxy (which blocks configurable fields),
|
||||
# since LLMProvider initialization legitimately needs the server-level default.
|
||||
if self.provider in ("gemini", "vertexai") and self.gemini_safety_settings is None:
|
||||
from ..config import _get_raw_config
|
||||
|
||||
try:
|
||||
raw_config = _get_raw_config()
|
||||
self.gemini_safety_settings = raw_config.llm_gemini_safety_settings
|
||||
except Exception:
|
||||
pass # Config may not be initialized in test environments
|
||||
|
||||
# Create provider implementation using factory
|
||||
self._provider_impl = create_llm_provider(
|
||||
provider=self.provider,
|
||||
@@ -335,6 +353,7 @@ class LLMProvider:
|
||||
vertexai_project_id=vertexai_project_id,
|
||||
vertexai_region=vertexai_region,
|
||||
vertexai_credentials=vertexai_credentials,
|
||||
gemini_safety_settings=self.gemini_safety_settings,
|
||||
)
|
||||
|
||||
# Backward compatibility: Keep mock provider properties
|
||||
@@ -503,6 +522,14 @@ class LLMProvider:
|
||||
|
||||
return result
|
||||
|
||||
def set_response_callback(self, fn: Any) -> None:
|
||||
"""Set a callback invoked on each call() instead of the fixed mock response."""
|
||||
if self.provider == "mock":
|
||||
from .providers.mock_llm import MockLLM
|
||||
|
||||
if isinstance(self._provider_impl, MockLLM):
|
||||
self._provider_impl.set_response_callback(fn)
|
||||
|
||||
def set_mock_response(self, response: Any) -> None:
|
||||
"""Set the response to return from mock calls."""
|
||||
# Backward compatibility: Store in both wrapper and provider implementation
|
||||
@@ -595,6 +622,23 @@ class LLMProvider:
|
||||
# SDK will automatically check for authentication when first used
|
||||
# No need to verify here - let it fail gracefully on first call with helpful error
|
||||
|
||||
def with_config(self, config: Any) -> "ConfiguredLLMProvider":
|
||||
"""
|
||||
Return a configured wrapper for a specific bank operation.
|
||||
|
||||
The wrapper applies per-bank overrides (e.g. Gemini safety settings)
|
||||
to every ``call()`` / ``call_with_tools()`` invocation without
|
||||
changing the underlying provider or its long-lived client connection.
|
||||
|
||||
Args:
|
||||
config: Resolved ``HindsightConfig`` for the current bank/request.
|
||||
|
||||
Returns:
|
||||
A ``ConfiguredLLMProvider`` that delegates to this provider with
|
||||
the supplied config applied.
|
||||
"""
|
||||
return ConfiguredLLMProvider(self, config.llm_gemini_safety_settings)
|
||||
|
||||
async def cleanup(self) -> None:
|
||||
"""Clean up resources."""
|
||||
pass
|
||||
@@ -656,5 +700,58 @@ class LLMProvider:
|
||||
return cls(provider=provider, api_key=api_key, base_url=base_url, model=model, reasoning_effort="high")
|
||||
|
||||
|
||||
class ConfiguredLLMProvider:
|
||||
"""
|
||||
Thin wrapper around LLMProvider that applies bank-specific config to every call.
|
||||
|
||||
Obtained via ``LLMProvider.with_config(resolved_config)``. The wrapper
|
||||
sets any provider-specific overrides (currently Gemini safety settings)
|
||||
immediately before each call using a ContextVar token, then resets it
|
||||
afterwards — so nesting is safe and the configuration cannot leak across
|
||||
operations.
|
||||
|
||||
All attribute access falls through to the underlying provider so callers
|
||||
that read ``llm.provider``, ``llm.model``, etc. continue to work without
|
||||
any changes.
|
||||
"""
|
||||
|
||||
def __init__(self, provider: "LLMProvider", gemini_safety_settings: list | None) -> None:
|
||||
# Use object.__setattr__ to avoid triggering __getattr__
|
||||
object.__setattr__(self, "_provider", provider)
|
||||
object.__setattr__(self, "_gemini_safety_settings", gemini_safety_settings)
|
||||
|
||||
# ── attribute passthrough ──────────────────────────────────────────────────
|
||||
|
||||
def __getattr__(self, name: str) -> Any:
|
||||
return getattr(object.__getattribute__(self, "_provider"), name)
|
||||
|
||||
# ── overridden call methods ────────────────────────────────────────────────
|
||||
|
||||
async def call(self, messages: list[dict[str, Any]], **kwargs: Any) -> Any:
|
||||
from .providers.gemini_llm import _safety_settings_ctx
|
||||
|
||||
token = _safety_settings_ctx.set(object.__getattribute__(self, "_gemini_safety_settings"))
|
||||
try:
|
||||
return await object.__getattribute__(self, "_provider").call(messages=messages, **kwargs)
|
||||
finally:
|
||||
_safety_settings_ctx.reset(token)
|
||||
|
||||
async def call_with_tools(
|
||||
self,
|
||||
messages: list[dict[str, Any]],
|
||||
tools: list[dict[str, Any]],
|
||||
**kwargs: Any,
|
||||
) -> "LLMToolCallResult":
|
||||
from .providers.gemini_llm import _safety_settings_ctx
|
||||
|
||||
token = _safety_settings_ctx.set(object.__getattribute__(self, "_gemini_safety_settings"))
|
||||
try:
|
||||
return await object.__getattribute__(self, "_provider").call_with_tools(
|
||||
messages=messages, tools=tools, **kwargs
|
||||
)
|
||||
finally:
|
||||
_safety_settings_ctx.reset(token)
|
||||
|
||||
|
||||
# Backwards compatibility alias
|
||||
LLMConfig = LLMProvider
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,14 +0,0 @@
|
||||
"""
|
||||
Mental models module for Hindsight.
|
||||
|
||||
Mental models contain directives - hard rules that are injected into reflect prompts.
|
||||
Directives are user-defined and their observations are user-provided (not LLM-generated).
|
||||
|
||||
Other types of consolidated knowledge are handled by:
|
||||
- Learnings: Automatic bottom-up consolidation from facts
|
||||
- Pinned Reflections: User-curated living documents
|
||||
"""
|
||||
|
||||
from .models import MentalModel, MentalModelSubtype
|
||||
|
||||
__all__ = ["MentalModel", "MentalModelSubtype"]
|
||||
@@ -1,53 +0,0 @@
|
||||
"""
|
||||
Pydantic models for mental models.
|
||||
"""
|
||||
|
||||
from datetime import datetime, timezone
|
||||
from enum import Enum
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
class MentalModelSubtype(str, Enum):
|
||||
"""Subtype of mental model.
|
||||
|
||||
Currently only DIRECTIVE is supported. Other types of consolidated knowledge
|
||||
are handled by:
|
||||
- Learnings: Automatic bottom-up consolidation from facts
|
||||
- Pinned Reflections: User-curated living documents
|
||||
"""
|
||||
|
||||
DIRECTIVE = "directive" # User-defined hard rules, observations user-provided
|
||||
|
||||
|
||||
class MentalModel(BaseModel):
|
||||
"""
|
||||
A mental model representing synthesized understanding.
|
||||
|
||||
Mental models are the agent's consolidated knowledge. Unlike raw facts,
|
||||
mental models provide:
|
||||
- A one-liner description for quick scanning/retrieval
|
||||
- A full summary for deep understanding
|
||||
- Links to related mental models
|
||||
"""
|
||||
|
||||
id: str = Field(description="Unique identifier within the bank")
|
||||
bank_id: str = Field(description="Bank this mental model belongs to")
|
||||
subtype: MentalModelSubtype = Field(description="How this model was created")
|
||||
name: str = Field(description="Human-readable name")
|
||||
description: str = Field(description="One-liner for quick scanning and retrieval matching")
|
||||
summary: str | None = Field(default=None, description="Full synthesized understanding")
|
||||
|
||||
# References
|
||||
entity_id: str | None = Field(default=None, description="Reference to entities table when type=entity")
|
||||
source_facts: list[str] = Field(default_factory=list, description="Fact IDs used to generate summary")
|
||||
links: list[str] = Field(default_factory=list, description="Related mental model IDs")
|
||||
|
||||
# Tags for scoped visibility (similar to document tags)
|
||||
tags: list[str] = Field(default_factory=list, description="Tags for scoped visibility filtering")
|
||||
|
||||
# Timestamps
|
||||
last_updated: datetime | None = Field(default=None, description="When summary was last regenerated")
|
||||
created_at: datetime = Field(
|
||||
default_factory=lambda: datetime.now(timezone.utc), description="When this model was created"
|
||||
)
|
||||
@@ -238,21 +238,24 @@ class ClaudeCodeLLM(LLMInterface):
|
||||
)
|
||||
|
||||
# Record trace span
|
||||
from hindsight_api.tracing import get_span_recorder
|
||||
try:
|
||||
from hindsight_api.tracing import get_span_recorder
|
||||
|
||||
span_recorder = get_span_recorder()
|
||||
span_recorder.record_llm_call(
|
||||
provider=self.provider,
|
||||
model=self.model,
|
||||
scope=scope,
|
||||
messages=messages,
|
||||
response_content=result if isinstance(result, str) else json.dumps(result),
|
||||
input_tokens=estimated_input,
|
||||
output_tokens=estimated_output,
|
||||
duration=duration,
|
||||
finish_reason=None,
|
||||
error=None,
|
||||
)
|
||||
span_recorder = get_span_recorder()
|
||||
span_recorder.record_llm_call(
|
||||
provider=self.provider,
|
||||
model=self.model,
|
||||
scope=scope,
|
||||
messages=messages,
|
||||
response_content=result if isinstance(result, str) else result.model_dump_json(),
|
||||
input_tokens=estimated_input,
|
||||
output_tokens=estimated_output,
|
||||
duration=duration,
|
||||
finish_reason=None,
|
||||
error=None,
|
||||
)
|
||||
except Exception:
|
||||
pass # logging failure must never affect the operation
|
||||
|
||||
# Log slow calls
|
||||
if duration > 10.0:
|
||||
|
||||
@@ -11,6 +11,7 @@ import json
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
from contextvars import ContextVar
|
||||
from typing import Any
|
||||
|
||||
from google import genai
|
||||
@@ -24,6 +25,12 @@ from hindsight_api.metrics import get_metrics_collector
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Per-request Gemini safety settings override.
|
||||
# Set exclusively by ConfiguredLLMProvider.call() / call_with_tools() via token-based
|
||||
# set/reset, so it is properly scoped to each individual LLM call and never leaks.
|
||||
_safety_settings_ctx: ContextVar[list | None] = ContextVar("gemini_safety_settings", default=None)
|
||||
|
||||
|
||||
# Vertex AI imports (optional)
|
||||
try:
|
||||
import google.auth
|
||||
@@ -58,6 +65,9 @@ class GeminiLLM(LLMInterface):
|
||||
self._client = None
|
||||
self._is_vertexai = self.provider == "vertexai"
|
||||
|
||||
# Safety settings: None means use Gemini's defaults
|
||||
self._safety_settings: list | None = kwargs.get("gemini_safety_settings")
|
||||
|
||||
if self._is_vertexai:
|
||||
self._init_vertexai(**kwargs)
|
||||
else:
|
||||
@@ -216,6 +226,16 @@ class GeminiLLM(LLMInterface):
|
||||
if temperature is not None:
|
||||
config_kwargs["temperature"] = temperature
|
||||
|
||||
# Apply safety settings: context var (per-request bank override) takes precedence over instance default
|
||||
effective_safety_settings = _safety_settings_ctx.get()
|
||||
if effective_safety_settings is None:
|
||||
effective_safety_settings = self._safety_settings
|
||||
if effective_safety_settings is not None:
|
||||
config_kwargs["safety_settings"] = [
|
||||
genai_types.SafetySetting(category=s["category"], threshold=s["threshold"])
|
||||
for s in effective_safety_settings
|
||||
]
|
||||
|
||||
generation_config = genai_types.GenerateContentConfig(**config_kwargs) if config_kwargs else None
|
||||
|
||||
last_exception = None
|
||||
@@ -489,6 +509,16 @@ class GeminiLLM(LLMInterface):
|
||||
)
|
||||
# "auto" is the default (no tool_config needed)
|
||||
|
||||
# Apply safety settings: context var (per-request bank override) takes precedence over instance default
|
||||
effective_safety_settings = _safety_settings_ctx.get()
|
||||
if effective_safety_settings is None:
|
||||
effective_safety_settings = self._safety_settings
|
||||
if effective_safety_settings is not None:
|
||||
config_kwargs["safety_settings"] = [
|
||||
genai_types.SafetySetting(category=s["category"], threshold=s["threshold"])
|
||||
for s in effective_safety_settings
|
||||
]
|
||||
|
||||
config = genai_types.GenerateContentConfig(**config_kwargs)
|
||||
|
||||
last_exception = None
|
||||
|
||||
@@ -6,6 +6,7 @@ without making actual API calls to external LLM services.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from collections.abc import Callable
|
||||
from typing import Any
|
||||
|
||||
from ..llm_interface import LLMInterface
|
||||
@@ -66,6 +67,7 @@ class MockLLM(LLMInterface):
|
||||
self._mock_calls: list[dict] = []
|
||||
self._mock_response: Any = None
|
||||
self._mock_exception: Exception | None = None
|
||||
self._response_callback: Callable[[list[dict], str], Any] | None = None
|
||||
|
||||
async def verify_connection(self) -> None:
|
||||
"""
|
||||
@@ -147,7 +149,9 @@ class MockLLM(LLMInterface):
|
||||
)
|
||||
|
||||
# Return mock response
|
||||
if self._mock_response is not None:
|
||||
if self._response_callback is not None:
|
||||
result = self._response_callback(messages, scope)
|
||||
elif self._mock_response is not None:
|
||||
result = self._mock_response
|
||||
elif response_format is not None:
|
||||
# Try to create a minimal valid instance of the response format
|
||||
@@ -214,7 +218,15 @@ class MockLLM(LLMInterface):
|
||||
|
||||
span_recorder = get_span_recorder()
|
||||
|
||||
if self._mock_response is not None:
|
||||
if self._response_callback is not None:
|
||||
cb_result = self._response_callback(messages, scope)
|
||||
if isinstance(cb_result, LLMToolCallResult):
|
||||
result = cb_result
|
||||
else:
|
||||
result = LLMToolCallResult(
|
||||
content=str(cb_result) if cb_result is not None else "mock response", finish_reason="stop"
|
||||
)
|
||||
elif self._mock_response is not None:
|
||||
if isinstance(self._mock_response, LLMToolCallResult):
|
||||
result = self._mock_response
|
||||
elif isinstance(self._mock_response, list):
|
||||
@@ -258,6 +270,16 @@ class MockLLM(LLMInterface):
|
||||
"""Clean up resources (no-op for mock provider)."""
|
||||
pass
|
||||
|
||||
def set_response_callback(self, fn: Callable[[list[dict], str], Any]) -> None:
|
||||
"""
|
||||
Set a callback invoked on each call() instead of _mock_response.
|
||||
|
||||
The callback receives (messages, scope) and returns the response.
|
||||
Useful for returning different responses per call (e.g., cycling
|
||||
through a corpus in a benchmark).
|
||||
"""
|
||||
self._response_callback = fn
|
||||
|
||||
def set_mock_response(self, response: Any) -> None:
|
||||
"""
|
||||
Set the response to return from mock calls.
|
||||
|
||||
@@ -92,11 +92,17 @@ class DateparserQueryAnalyzer(QueryAnalyzer):
|
||||
self._search_dates = None
|
||||
|
||||
def load(self) -> None:
|
||||
"""Load dateparser (lazy import)."""
|
||||
"""Load dateparser and warm up internal data structures.
|
||||
|
||||
Triggers the real initialization cost (regex tables, timezone data) at
|
||||
load time so the first actual recall doesn't pay the cold-start penalty.
|
||||
"""
|
||||
if self._search_dates is None:
|
||||
from dateparser.search import search_dates
|
||||
|
||||
self._search_dates = search_dates
|
||||
# Warm up: fire a dummy call to trigger lazy-loaded internal tables.
|
||||
self._search_dates("today")
|
||||
|
||||
def analyze(self, query: str, reference_date: datetime | None = None) -> QueryAnalysis:
|
||||
"""
|
||||
|
||||
@@ -14,6 +14,8 @@ import re
|
||||
import time
|
||||
from typing import TYPE_CHECKING, Any, Awaitable, Callable
|
||||
|
||||
import tiktoken
|
||||
|
||||
from .models import DirectiveInfo, LLMCall, ReflectAgentResult, TokenUsageSummary, ToolCall
|
||||
from .prompts import FINAL_SYSTEM_PROMPT, _extract_directive_rules, build_final_prompt, build_system_prompt_for_tools
|
||||
from .tools_schema import get_reflect_tools
|
||||
@@ -259,6 +261,46 @@ OUTPUT:"""
|
||||
return None, 0, 0
|
||||
|
||||
|
||||
_TIKTOKEN_ENCODING = tiktoken.get_encoding("cl100k_base")
|
||||
|
||||
|
||||
def _count_messages_tokens(messages: list[dict[str, Any]]) -> int:
|
||||
"""Estimate the token count of the messages list using cl100k_base encoding."""
|
||||
total = 0
|
||||
for msg in messages:
|
||||
content = msg.get("content") or ""
|
||||
if isinstance(content, str):
|
||||
total += len(_TIKTOKEN_ENCODING.encode(content))
|
||||
elif isinstance(content, list):
|
||||
for part in content:
|
||||
if isinstance(part, dict) and isinstance(part.get("text"), str):
|
||||
total += len(_TIKTOKEN_ENCODING.encode(part["text"]))
|
||||
# Tool call arguments and results also count
|
||||
for tc in msg.get("tool_calls") or []:
|
||||
if isinstance(tc, dict):
|
||||
func = tc.get("function", {})
|
||||
total += len(_TIKTOKEN_ENCODING.encode(func.get("arguments", "")))
|
||||
return total
|
||||
|
||||
|
||||
def _is_context_overflow_error(exc: Exception) -> bool:
|
||||
"""Return True if the exception signals the LLM context window was exceeded."""
|
||||
msg = str(exc).lower()
|
||||
return any(
|
||||
phrase in msg
|
||||
for phrase in (
|
||||
"context_length_exceeded",
|
||||
"context length exceeded",
|
||||
"maximum context length",
|
||||
"prompt_too_long",
|
||||
"prompt is too long",
|
||||
"resource_exhausted",
|
||||
"input is too long",
|
||||
"too many tokens",
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
async def run_reflect_agent(
|
||||
llm_config: "LLMProvider",
|
||||
bank_id: str,
|
||||
@@ -275,6 +317,7 @@ async def run_reflect_agent(
|
||||
directives: list[dict[str, Any]] | None = None,
|
||||
has_mental_models: bool = False,
|
||||
budget: str | None = None,
|
||||
max_context_tokens: int = 100_000,
|
||||
) -> ReflectAgentResult:
|
||||
"""
|
||||
Execute the reflect agent loop using native tool calling.
|
||||
@@ -388,7 +431,9 @@ async def run_reflect_agent(
|
||||
|
||||
if is_last:
|
||||
# Force text response on last iteration - no tools
|
||||
prompt = build_final_prompt(query, context_history, bank_profile, context)
|
||||
prompt = build_final_prompt(
|
||||
query, context_history, bank_profile, context, max_context_tokens=max_context_tokens
|
||||
)
|
||||
llm_start = time.time()
|
||||
response, usage = await llm_config.call(
|
||||
messages=[
|
||||
@@ -433,6 +478,62 @@ async def run_reflect_agent(
|
||||
directives_applied=directives_applied,
|
||||
)
|
||||
|
||||
# Proactive context-window guard: if accumulated messages would exceed the
|
||||
# configured token budget, bail out early and synthesize from what we have.
|
||||
estimated_tokens = _count_messages_tokens(messages)
|
||||
if estimated_tokens >= max_context_tokens and (
|
||||
bool(available_memory_ids) or bool(available_mental_model_ids) or bool(available_observation_ids)
|
||||
):
|
||||
logger.warning(
|
||||
f"[REFLECT {reflect_id}] Context budget exceeded on iteration {iteration + 1}: "
|
||||
f"~{estimated_tokens} tokens >= {max_context_tokens} limit. Forcing final synthesis."
|
||||
)
|
||||
prompt = build_final_prompt(
|
||||
query, context_history, bank_profile, context, max_context_tokens=max_context_tokens
|
||||
)
|
||||
llm_start = time.time()
|
||||
response, usage = await llm_config.call(
|
||||
messages=[
|
||||
{"role": "system", "content": FINAL_SYSTEM_PROMPT},
|
||||
{"role": "user", "content": prompt},
|
||||
],
|
||||
scope="reflect",
|
||||
max_completion_tokens=max_tokens,
|
||||
return_usage=True,
|
||||
)
|
||||
llm_duration = int((time.time() - llm_start) * 1000)
|
||||
total_input_tokens += usage.input_tokens
|
||||
total_output_tokens += usage.output_tokens
|
||||
llm_trace.append(
|
||||
{
|
||||
"scope": "final",
|
||||
"duration_ms": llm_duration,
|
||||
"input_tokens": usage.input_tokens,
|
||||
"output_tokens": usage.output_tokens,
|
||||
}
|
||||
)
|
||||
answer = _clean_answer_text(response.strip())
|
||||
|
||||
structured_output = None
|
||||
if response_schema and answer:
|
||||
structured_output, struct_in, struct_out = await _generate_structured_output(
|
||||
answer, response_schema, llm_config, reflect_id
|
||||
)
|
||||
total_input_tokens += struct_in
|
||||
total_output_tokens += struct_out
|
||||
|
||||
_log_completion(answer, iteration + 1, forced=True)
|
||||
return ReflectAgentResult(
|
||||
text=answer,
|
||||
structured_output=structured_output,
|
||||
iterations=iteration + 1,
|
||||
tools_called=total_tools_called,
|
||||
tool_trace=tool_trace,
|
||||
llm_trace=_get_llm_trace(),
|
||||
usage=_get_usage(),
|
||||
directives_applied=directives_applied,
|
||||
)
|
||||
|
||||
# Call LLM with tools
|
||||
llm_start = time.time()
|
||||
|
||||
@@ -478,13 +579,22 @@ async def run_reflect_agent(
|
||||
consecutive_errors += 1
|
||||
logger.warning(f"[REFLECT {reflect_id}] LLM error on iteration {iteration + 1}: {e} ({err_duration}ms)")
|
||||
llm_trace.append({"scope": f"agent_{iteration + 1}_err", "duration_ms": err_duration})
|
||||
# Guardrail: If no evidence gathered yet, retry (but cap consecutive errors to avoid long hangs)
|
||||
has_gathered_evidence = (
|
||||
bool(available_memory_ids) or bool(available_mental_model_ids) or bool(available_observation_ids)
|
||||
)
|
||||
if not has_gathered_evidence and iteration < max_iterations - 1 and consecutive_errors < 2:
|
||||
# Context overflow errors must never be retried — retrying would only make them worse.
|
||||
# Skip straight to final synthesis with whatever evidence we have.
|
||||
if _is_context_overflow_error(e):
|
||||
logger.warning(
|
||||
f"[REFLECT {reflect_id}] Context window exceeded on iteration {iteration + 1}, "
|
||||
"forcing final synthesis from gathered evidence."
|
||||
)
|
||||
# For other errors: retry if no evidence yet (but cap consecutive errors to avoid long hangs)
|
||||
elif not has_gathered_evidence and iteration < max_iterations - 1 and consecutive_errors < 2:
|
||||
continue
|
||||
prompt = build_final_prompt(query, context_history, bank_profile, context)
|
||||
prompt = build_final_prompt(
|
||||
query, context_history, bank_profile, context, max_context_tokens=max_context_tokens
|
||||
)
|
||||
llm_start = time.time()
|
||||
response, usage = await llm_config.call(
|
||||
messages=[
|
||||
@@ -555,7 +665,9 @@ async def run_reflect_agent(
|
||||
directives_applied=directives_applied,
|
||||
)
|
||||
# Empty response, force final
|
||||
prompt = build_final_prompt(query, context_history, bank_profile, context)
|
||||
prompt = build_final_prompt(
|
||||
query, context_history, bank_profile, context, max_context_tokens=max_context_tokens
|
||||
)
|
||||
llm_start = time.time()
|
||||
response, usage = await llm_config.call(
|
||||
messages=[
|
||||
|
||||
@@ -10,6 +10,14 @@ The reflect agent uses hierarchical retrieval:
|
||||
import json
|
||||
from typing import Any
|
||||
|
||||
import tiktoken
|
||||
|
||||
_TIKTOKEN_ENCODING = tiktoken.get_encoding("cl100k_base")
|
||||
|
||||
# Fraction of max_context_tokens reserved for tool results in the final synthesis prompt.
|
||||
# The remainder covers the system prompt, question, bank context, and output tokens.
|
||||
_FINAL_PROMPT_CONTEXT_FRACTION = 0.8
|
||||
|
||||
|
||||
def _extract_directive_rules(directives: list[dict[str, Any]]) -> list[str]:
|
||||
"""Extract directive rules as a list of strings."""
|
||||
@@ -394,6 +402,7 @@ def build_final_prompt(
|
||||
context_history: list[dict],
|
||||
bank_profile: dict,
|
||||
additional_context: str | None = None,
|
||||
max_context_tokens: int = 100_000,
|
||||
) -> str:
|
||||
"""Build the final prompt when forcing a text response (no tools)."""
|
||||
parts = []
|
||||
@@ -423,18 +432,32 @@ def build_final_prompt(
|
||||
if additional_context:
|
||||
parts.append(f"\n## Additional Context\n{additional_context}")
|
||||
|
||||
# Tool call history
|
||||
# Tool call history — include as many entries as fit within the token budget,
|
||||
# preferring the most recent calls (they tend to be the most targeted).
|
||||
if context_history:
|
||||
parts.append("\n## Retrieved Data (synthesize and reason from this data)")
|
||||
for entry in context_history:
|
||||
token_budget = int(max_context_tokens * _FINAL_PROMPT_CONTEXT_FRACTION)
|
||||
# Render entries newest-first, then reverse so the prompt reads chronologically.
|
||||
rendered: list[str] = []
|
||||
truncated = False
|
||||
for entry in reversed(context_history):
|
||||
tool = entry["tool"]
|
||||
output = entry["output"]
|
||||
# Format as proper JSON for LLM readability
|
||||
try:
|
||||
output_str = json.dumps(output, indent=2, default=str)
|
||||
except (TypeError, ValueError):
|
||||
output_str = str(output)
|
||||
parts.append(f"\n### From {tool}:\n```json\n{output_str}\n```")
|
||||
block = f"\n### From {tool}:\n```json\n{output_str}\n```"
|
||||
block_tokens = len(_TIKTOKEN_ENCODING.encode(block))
|
||||
if block_tokens > token_budget:
|
||||
truncated = True
|
||||
break
|
||||
rendered.append(block)
|
||||
token_budget -= block_tokens
|
||||
for block in reversed(rendered):
|
||||
parts.append(block)
|
||||
if truncated:
|
||||
parts.append("\n*Note: Some earlier tool results were omitted to stay within the context window.*")
|
||||
else:
|
||||
parts.append("\n## Retrieved Data\nNo data was retrieved.")
|
||||
|
||||
|
||||
@@ -5,7 +5,6 @@ This package contains modular components for the retain operation:
|
||||
- types: Type definitions for retain pipeline
|
||||
- fact_extraction: Extract facts from content
|
||||
- embedding_processing: Augment texts and generate embeddings
|
||||
- deduplication: Check for duplicate facts
|
||||
- entity_processing: Process and resolve entities
|
||||
- link_creation: Create temporal, semantic, entity, and causal links
|
||||
- chunk_storage: Handle chunk storage
|
||||
@@ -14,7 +13,6 @@ This package contains modular components for the retain operation:
|
||||
|
||||
from . import (
|
||||
chunk_storage,
|
||||
deduplication,
|
||||
embedding_processing,
|
||||
entity_processing,
|
||||
fact_extraction,
|
||||
@@ -35,7 +33,6 @@ __all__ = [
|
||||
# Modules
|
||||
"fact_extraction",
|
||||
"embedding_processing",
|
||||
"deduplication",
|
||||
"entity_processing",
|
||||
"link_creation",
|
||||
"chunk_storage",
|
||||
|
||||
@@ -1,85 +0,0 @@
|
||||
"""
|
||||
Deduplication logic for retain pipeline.
|
||||
|
||||
Checks for duplicate facts using semantic similarity and temporal proximity.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from collections import defaultdict
|
||||
from datetime import UTC
|
||||
|
||||
from .types import ProcessedFact
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
async def check_duplicates_batch(conn, bank_id: str, facts: list[ProcessedFact], duplicate_checker_fn) -> list[bool]:
|
||||
"""
|
||||
Check which facts are duplicates using batched time-window queries.
|
||||
|
||||
Groups facts by 12-hour time buckets to efficiently check for duplicates
|
||||
within a 24-hour window.
|
||||
|
||||
Args:
|
||||
conn: Database connection
|
||||
bank_id: Bank identifier
|
||||
facts: List of ProcessedFact objects to check
|
||||
duplicate_checker_fn: Async function(conn, bank_id, texts, embeddings, date, time_window_hours)
|
||||
that returns List[bool] indicating duplicates
|
||||
|
||||
Returns:
|
||||
List of boolean flags (same length as facts) indicating if each fact is a duplicate
|
||||
"""
|
||||
if not facts:
|
||||
return []
|
||||
|
||||
# Group facts by event_date (rounded to 12-hour buckets) for efficient batching
|
||||
time_buckets = defaultdict(list)
|
||||
for idx, fact in enumerate(facts):
|
||||
# Use occurred_start if available, otherwise use mentioned_at
|
||||
# For deduplication purposes, we need a time reference
|
||||
fact_date = fact.occurred_start if fact.occurred_start is not None else fact.mentioned_at
|
||||
|
||||
# Defensive: if both are None (shouldn't happen), use now()
|
||||
if fact_date is None:
|
||||
from datetime import datetime
|
||||
|
||||
fact_date = datetime.now(UTC)
|
||||
|
||||
# Round to 12-hour bucket to group similar times
|
||||
bucket_key = fact_date.replace(hour=(fact_date.hour // 12) * 12, minute=0, second=0, microsecond=0)
|
||||
time_buckets[bucket_key].append((idx, fact))
|
||||
|
||||
# Process each bucket in batch
|
||||
all_is_duplicate = [False] * len(facts)
|
||||
|
||||
for bucket_date, bucket_items in time_buckets.items():
|
||||
indices = [item[0] for item in bucket_items]
|
||||
texts = [item[1].fact_text for item in bucket_items]
|
||||
embeddings = [item[1].embedding for item in bucket_items]
|
||||
|
||||
# Check duplicates for this time bucket
|
||||
dup_flags = await duplicate_checker_fn(conn, bank_id, texts, embeddings, bucket_date, time_window_hours=24)
|
||||
|
||||
# Map results back to original indices
|
||||
for idx, is_dup in zip(indices, dup_flags):
|
||||
all_is_duplicate[idx] = is_dup
|
||||
|
||||
return all_is_duplicate
|
||||
|
||||
|
||||
def filter_duplicates(facts: list[ProcessedFact], is_duplicate_flags: list[bool]) -> list[ProcessedFact]:
|
||||
"""
|
||||
Filter out duplicate facts based on duplicate flags.
|
||||
|
||||
Args:
|
||||
facts: List of ProcessedFact objects
|
||||
is_duplicate_flags: Boolean flags indicating which facts are duplicates
|
||||
|
||||
Returns:
|
||||
List of non-duplicate facts
|
||||
"""
|
||||
if len(facts) != len(is_duplicate_flags):
|
||||
raise ValueError(f"Mismatch between facts ({len(facts)}) and flags ({len(is_duplicate_flags)})")
|
||||
|
||||
return [fact for fact, is_dup in zip(facts, is_duplicate_flags) if not is_dup]
|
||||
@@ -27,11 +27,21 @@ def augment_texts_with_dates(facts: list[ExtractedFact], format_date_fn) -> list
|
||||
"""
|
||||
augmented_texts = []
|
||||
for fact in facts:
|
||||
# Use occurred_start as the representative date
|
||||
# Use occurred_start as the representative date, fall back to mentioned_at
|
||||
fact_date = fact.occurred_start or fact.mentioned_at
|
||||
readable_date = format_date_fn(fact_date)
|
||||
# Augment text with date for embedding (but store original text in DB)
|
||||
augmented_text = f"{fact.fact_text} (happened in {readable_date})"
|
||||
# Augment text with date and entity names for embedding (but store original text in DB)
|
||||
# Entity names (including key:value labels) improve retrieval without polluting stored content
|
||||
if fact_date is not None:
|
||||
readable_date = format_date_fn(fact_date)
|
||||
if fact.occurred_end and fact.occurred_end != fact.occurred_start:
|
||||
readable_end = format_date_fn(fact.occurred_end)
|
||||
augmented_text = f"{fact.fact_text} (happened from {readable_date} to {readable_end})"
|
||||
else:
|
||||
augmented_text = f"{fact.fact_text} (happened in {readable_date})"
|
||||
else:
|
||||
augmented_text = fact.fact_text
|
||||
if fact.entities:
|
||||
augmented_text = f"{augmented_text} [{', '.join(fact.entities)}]"
|
||||
augmented_texts.append(augmented_text)
|
||||
return augmented_texts
|
||||
|
||||
|
||||
@@ -41,10 +41,9 @@ async def generate_embeddings_batch(embeddings_backend, texts: list[str]) -> lis
|
||||
List of embeddings in same order as input texts
|
||||
"""
|
||||
try:
|
||||
# Run embeddings in thread pool to avoid blocking event loop
|
||||
loop = asyncio.get_event_loop()
|
||||
embeddings = await loop.run_in_executor(
|
||||
None, # Use default thread pool
|
||||
None,
|
||||
embeddings_backend.encode,
|
||||
texts,
|
||||
)
|
||||
|
||||
@@ -0,0 +1,194 @@
|
||||
"""
|
||||
Entity labels models and helpers for retain pipeline.
|
||||
|
||||
Defines a controlled vocabulary of key:value classification labels
|
||||
(e.g., 'pedagogy:scaffolding', 'interest:active') that are extracted
|
||||
at retain time and stored as entities.
|
||||
"""
|
||||
|
||||
from typing import Literal
|
||||
|
||||
from pydantic import BaseModel, Field, create_model
|
||||
|
||||
|
||||
class LabelValue(BaseModel):
|
||||
"""A single allowed value for a label group."""
|
||||
|
||||
value: str
|
||||
description: str = ""
|
||||
|
||||
|
||||
class LabelGroup(BaseModel):
|
||||
"""A label group (dimension) with its type and allowed values."""
|
||||
|
||||
key: str
|
||||
description: str = ""
|
||||
type: Literal["value", "multi-values", "text"] = "value"
|
||||
optional: bool = True
|
||||
tag: bool = False
|
||||
values: list[LabelValue] = []
|
||||
|
||||
|
||||
class EntityLabelsConfig(BaseModel):
|
||||
"""Entity labels configuration for a bank (controlled vocabulary)."""
|
||||
|
||||
attributes: list[LabelGroup] = []
|
||||
|
||||
|
||||
def parse_entity_labels(raw: dict | list | None) -> EntityLabelsConfig | None:
|
||||
"""
|
||||
Parse raw entity labels config into EntityLabelsConfig.
|
||||
|
||||
Accepts:
|
||||
- None → returns None
|
||||
- list → list of attribute dicts (each may use legacy free_values/multi_value or new type field)
|
||||
- dict → {attributes: [...]}
|
||||
|
||||
Legacy migration (backward-compat):
|
||||
- free_values=True → type="text"
|
||||
- multi_value=True → type="multi-values"
|
||||
- neither / free_values=False → type="value"
|
||||
|
||||
Args:
|
||||
raw: Raw entity labels config from bank config
|
||||
|
||||
Returns:
|
||||
EntityLabelsConfig or None if raw is None/empty
|
||||
"""
|
||||
if raw is None:
|
||||
return None
|
||||
|
||||
if isinstance(raw, list):
|
||||
if not raw:
|
||||
return None
|
||||
attributes = [LabelGroup.model_validate(_migrate_label_group(a)) for a in raw]
|
||||
return EntityLabelsConfig(attributes=attributes)
|
||||
|
||||
if isinstance(raw, dict):
|
||||
attrs_raw = raw.get("attributes", [])
|
||||
if not attrs_raw:
|
||||
return None
|
||||
attributes = [LabelGroup.model_validate(_migrate_label_group(a)) for a in attrs_raw]
|
||||
return EntityLabelsConfig(attributes=attributes)
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def _migrate_label_group(raw: dict) -> dict:
|
||||
"""Migrate legacy free_values/multi_value fields to the new type field."""
|
||||
if not isinstance(raw, dict) or "type" in raw:
|
||||
return raw
|
||||
patched = dict(raw)
|
||||
if patched.get("free_values"):
|
||||
patched["type"] = "text"
|
||||
elif patched.get("multi_value"):
|
||||
patched["type"] = "multi-values"
|
||||
else:
|
||||
patched["type"] = "value"
|
||||
# Remove legacy keys so Pydantic doesn't error on unknown fields
|
||||
patched.pop("free_values", None)
|
||||
patched.pop("multi_value", None)
|
||||
return patched
|
||||
|
||||
|
||||
def build_labels_model(labels_cfg: EntityLabelsConfig) -> type[BaseModel] | None:
|
||||
"""
|
||||
Build a dynamic Pydantic model for structured label extraction.
|
||||
|
||||
Each LabelGroup becomes a typed field based on its type:
|
||||
- type="text" → str | None (always optional)
|
||||
- type="value", optional=True → Literal["v1","v2"] | None
|
||||
- type="value", optional=False → Literal["v1","v2"] (required)
|
||||
- type="multi-values" → list[Literal["v1","v2"]]
|
||||
|
||||
Args:
|
||||
labels_cfg: Parsed EntityLabelsConfig
|
||||
|
||||
Returns:
|
||||
Dynamic Pydantic model class, or None if no groups defined
|
||||
"""
|
||||
fields: dict = {}
|
||||
for group in labels_cfg.attributes:
|
||||
if not group.key:
|
||||
continue
|
||||
description = group.description or group.key
|
||||
|
||||
if group.type == "text":
|
||||
# Free-form: any string value accepted, always optional
|
||||
fields[group.key] = (str | None, Field(default=None, description=description))
|
||||
else:
|
||||
# Enum-constrained: must have defined values
|
||||
if not group.values:
|
||||
continue
|
||||
values = tuple(v.value for v in group.values if v.value)
|
||||
if not values:
|
||||
continue
|
||||
# Literal[("v1", "v2")] is equivalent to Literal["v1", "v2"] in Python 3.11+
|
||||
literal_type = Literal[values] # type: ignore[valid-type]
|
||||
if group.type == "multi-values":
|
||||
fields[group.key] = (
|
||||
list[literal_type], # type: ignore[valid-type]
|
||||
Field(default_factory=list, description=description),
|
||||
)
|
||||
elif group.optional:
|
||||
fields[group.key] = (
|
||||
literal_type | None, # type: ignore[valid-type]
|
||||
Field(default=None, description=description),
|
||||
)
|
||||
else:
|
||||
fields[group.key] = (
|
||||
literal_type, # type: ignore[valid-type]
|
||||
Field(description=description),
|
||||
)
|
||||
|
||||
if not fields:
|
||||
return None
|
||||
|
||||
return create_model("Labels", **fields)
|
||||
|
||||
|
||||
def is_label_entity(text: str, labels_cfg: EntityLabelsConfig, labels_lookup: set[str]) -> bool:
|
||||
"""
|
||||
Return True if entity text belongs to any configured label group.
|
||||
|
||||
For enum groups: checks the pre-built lookup set.
|
||||
For text groups: checks that the text starts with a known key prefix.
|
||||
"""
|
||||
if text.lower() in labels_lookup:
|
||||
return True
|
||||
for group in labels_cfg.attributes:
|
||||
if group.type == "text" and group.key and text.lower().startswith(f"{group.key.lower()}:"):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def build_labels_lookup(labels_cfg: EntityLabelsConfig | list | None) -> set[str]:
|
||||
"""
|
||||
Build a set of valid 'key:value' label strings (lowercase) for fast lookup.
|
||||
|
||||
Accepts either EntityLabelsConfig or raw list/None for backwards compatibility.
|
||||
|
||||
Args:
|
||||
labels_cfg: EntityLabelsConfig, raw list of attribute dicts, or None
|
||||
|
||||
Returns:
|
||||
Set of lowercase 'key:value' strings
|
||||
"""
|
||||
if labels_cfg is None:
|
||||
return set()
|
||||
|
||||
# Accept raw list/dict for backwards compatibility
|
||||
if not isinstance(labels_cfg, EntityLabelsConfig):
|
||||
parsed = parse_entity_labels(labels_cfg)
|
||||
if parsed is None:
|
||||
return set()
|
||||
labels_cfg = parsed
|
||||
|
||||
valid = set()
|
||||
for group in labels_cfg.attributes:
|
||||
if group.type == "text":
|
||||
continue # No fixed vocabulary — all values accepted in post-processing
|
||||
for v in group.values:
|
||||
if group.key and v.value:
|
||||
valid.add(f"{group.key}:{v.value}".lower())
|
||||
return valid
|
||||
@@ -20,6 +20,7 @@ async def process_entities_batch(
|
||||
facts: list[ProcessedFact],
|
||||
log_buffer: list[str] = None,
|
||||
user_entities_per_content: dict[int, list[dict]] = None,
|
||||
entity_labels: list | None = None,
|
||||
) -> list[EntityLink]:
|
||||
"""
|
||||
Process entities for all facts and create entity links.
|
||||
@@ -90,6 +91,7 @@ async def process_entities_batch(
|
||||
fact_dates,
|
||||
entities_per_fact,
|
||||
log_buffer, # Pass log_buffer for detailed logging
|
||||
entity_labels=entity_labels,
|
||||
)
|
||||
|
||||
return entity_links
|
||||
|
||||
@@ -10,22 +10,32 @@ import json
|
||||
import logging
|
||||
import re
|
||||
from datetime import datetime, timedelta
|
||||
from typing import Literal
|
||||
from typing import Literal, cast
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field, field_validator
|
||||
from pydantic import BaseModel, ConfigDict, Field, create_model, field_validator
|
||||
|
||||
from ...config import get_config
|
||||
from ..llm_wrapper import LLMConfig, OutputTooLongError
|
||||
from ..response_models import TokenUsage
|
||||
from .entity_labels import (
|
||||
EntityLabelsConfig,
|
||||
build_labels_lookup,
|
||||
build_labels_model,
|
||||
is_label_entity,
|
||||
parse_entity_labels,
|
||||
)
|
||||
|
||||
|
||||
def _infer_temporal_date(fact_text: str, event_date: datetime) -> str | None:
|
||||
def _infer_temporal_date(fact_text: str, event_date: datetime | None) -> str | None:
|
||||
"""
|
||||
Infer a temporal date from fact text when LLM didn't provide occurred_start.
|
||||
|
||||
This is a fallback for when the LLM fails to extract temporal information
|
||||
from relative time expressions like "last night", "yesterday", etc.
|
||||
"""
|
||||
if event_date is None:
|
||||
return None
|
||||
|
||||
fact_lower = fact_text.lower()
|
||||
|
||||
# Map relative time expressions to day offsets
|
||||
@@ -100,7 +110,6 @@ class Fact(BaseModel):
|
||||
# Optional temporal fields
|
||||
occurred_start: str | None = None
|
||||
occurred_end: str | None = None
|
||||
mentioned_at: str | None = None
|
||||
|
||||
# Optional location field
|
||||
where: str | None = Field(
|
||||
@@ -690,10 +699,62 @@ Example: "Lost job → couldn't pay rent → moved apartment"
|
||||
- Fact 2: Moved apartment, causal_relations: [{target_index: 1, relation_type: "caused_by"}]"""
|
||||
|
||||
|
||||
def _build_labels_prompt_section(labels_cfg: EntityLabelsConfig | list | None, free_form_entities: bool = True) -> str:
|
||||
"""Build the entity labels classification section for the extraction prompt."""
|
||||
if labels_cfg is None:
|
||||
return ""
|
||||
|
||||
# Accept raw list for backwards compatibility
|
||||
if isinstance(labels_cfg, list):
|
||||
if not labels_cfg:
|
||||
return ""
|
||||
labels_cfg = parse_entity_labels(labels_cfg)
|
||||
if labels_cfg is None:
|
||||
return ""
|
||||
|
||||
if not labels_cfg.attributes:
|
||||
return ""
|
||||
|
||||
if free_form_entities:
|
||||
entities_instruction = "Classify each fact using the structured 'labels' field below. Continue extracting regular named entities in the 'entities' field."
|
||||
else:
|
||||
entities_instruction = "Classify each fact using the structured 'labels' field below. Do NOT add regular named entities — labels-only mode."
|
||||
|
||||
lines = [
|
||||
"\n\n══════════════════════════════════════════════════════════════════════════",
|
||||
"ENTITY LABELS - CLASSIFICATION ATTRIBUTES",
|
||||
"══════════════════════════════════════════════════════════════════════════",
|
||||
"",
|
||||
entities_instruction,
|
||||
"",
|
||||
"For each fact, fill the 'labels' object. Each field is a label group:",
|
||||
"",
|
||||
]
|
||||
|
||||
for attr in labels_cfg.attributes:
|
||||
if attr.type == "text":
|
||||
# Free-text: no predefined values — LLM writes any relevant string or null
|
||||
lines.append(f"- {attr.key} (free text or null): {attr.description}")
|
||||
else:
|
||||
mode = "multi-value (list)" if attr.type == "multi-values" else "single value or null"
|
||||
lines.append(f"- {attr.key} ({mode}): {attr.description}")
|
||||
for v in attr.values:
|
||||
desc = f" — {v.description}" if v.description else ""
|
||||
lines.append(f' • "{v.value}"{desc}')
|
||||
lines.append("")
|
||||
|
||||
lines.append("Only assign labels when clearly applicable. Leave null/empty if the fact does not match.")
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
def _build_extraction_prompt_and_schema(config) -> tuple[str, type]:
|
||||
"""
|
||||
Build extraction prompt and response schema based on config.
|
||||
|
||||
When a taxonomy is configured, dynamically builds a Pydantic model with a
|
||||
typed `taxonomy_entities` field using an Enum built from valid taxonomy values.
|
||||
This enables JSON schema enforcement for structured outputs.
|
||||
|
||||
Returns:
|
||||
Tuple of (prompt, response_schema)
|
||||
"""
|
||||
@@ -736,9 +797,53 @@ def _build_extraction_prompt_and_schema(config) -> tuple[str, type]:
|
||||
# Add causal relationships section if enabled
|
||||
if extract_causal_links:
|
||||
prompt = prompt + CAUSAL_RELATIONSHIPS_SECTION
|
||||
response_schema = FactExtractionResponseVerbose if extraction_mode == "verbose" else FactExtractionResponse
|
||||
base_fact_class = ExtractedFactVerbose if extraction_mode == "verbose" else ExtractedFact
|
||||
base_response_class = FactExtractionResponseVerbose if extraction_mode == "verbose" else FactExtractionResponse
|
||||
else:
|
||||
response_schema = FactExtractionResponseNoCausal
|
||||
base_fact_class = ExtractedFactNoCausal
|
||||
base_response_class = FactExtractionResponseNoCausal
|
||||
|
||||
# Add entity labels section if configured and build dynamic schema
|
||||
entity_labels_raw = getattr(config, "entity_labels", None)
|
||||
labels_cfg = parse_entity_labels(entity_labels_raw)
|
||||
free_form_entities = getattr(config, "entities_allow_free_form", True)
|
||||
labels_section = _build_labels_prompt_section(labels_cfg, free_form_entities)
|
||||
if labels_section:
|
||||
prompt = prompt + labels_section
|
||||
|
||||
response_schema = base_response_class
|
||||
|
||||
if labels_cfg and labels_cfg.attributes:
|
||||
LabelsModel = build_labels_model(labels_cfg)
|
||||
if LabelsModel is not None:
|
||||
dynamic_fields: dict = {
|
||||
"labels": (
|
||||
LabelsModel,
|
||||
Field(
|
||||
description="Classification labels for this fact. Fill each applicable field; leave others null/empty."
|
||||
),
|
||||
)
|
||||
}
|
||||
if not free_form_entities:
|
||||
dynamic_fields["entities"] = (
|
||||
list[Entity] | None,
|
||||
Field(default=None, description="Leave empty — labels-only mode"),
|
||||
)
|
||||
# Inherit parent's required fields and add 'labels' so it appears in the JSON schema
|
||||
# required array (the base class json_schema_extra overrides required entirely)
|
||||
base_extra = base_fact_class.model_config.get("json_schema_extra")
|
||||
base_required = cast(dict, base_extra).get("required", []) if isinstance(base_extra, dict) else []
|
||||
DynamicFact = create_model(
|
||||
"LabelsFact",
|
||||
__base__=base_fact_class,
|
||||
__config__=ConfigDict(
|
||||
json_schema_mode="validation",
|
||||
json_schema_extra={"required": [*base_required, "labels"]},
|
||||
),
|
||||
**dynamic_fields,
|
||||
)
|
||||
DynamicResponse = create_model("LabelsResponse", facts=(list[DynamicFact], ...)) # type: ignore[valid-type]
|
||||
response_schema = DynamicResponse
|
||||
|
||||
return prompt, response_schema
|
||||
|
||||
@@ -747,7 +852,7 @@ def _build_user_message(
|
||||
chunk: str,
|
||||
chunk_index: int,
|
||||
total_chunks: int,
|
||||
event_date: datetime,
|
||||
event_date: datetime | None,
|
||||
context: str,
|
||||
metadata: dict[str, str] | None = None,
|
||||
) -> str:
|
||||
@@ -756,8 +861,12 @@ def _build_user_message(
|
||||
|
||||
sanitized_chunk = _sanitize_text(chunk)
|
||||
sanitized_context = _sanitize_text(context) if context else "none"
|
||||
event_date = parse_datetime_flexible(event_date)
|
||||
event_date_formatted = event_date.strftime("%A, %B %d, %Y")
|
||||
|
||||
if event_date is not None:
|
||||
event_date = parse_datetime_flexible(event_date)
|
||||
event_date_str = f"{event_date.strftime('%A, %B %d, %Y')} ({event_date.isoformat()})"
|
||||
else:
|
||||
event_date_str = "Unknown"
|
||||
|
||||
metadata_section = ""
|
||||
if metadata:
|
||||
@@ -767,7 +876,7 @@ def _build_user_message(
|
||||
return f"""Extract facts from the following text chunk.
|
||||
|
||||
Chunk: {chunk_index + 1}/{total_chunks}
|
||||
Event Date: {event_date_formatted} ({event_date.isoformat()})
|
||||
Event Date: {event_date_str}
|
||||
Context: {sanitized_context}{metadata_section}
|
||||
|
||||
Text:
|
||||
@@ -805,7 +914,7 @@ async def _extract_facts_from_chunk(
|
||||
chunk: str,
|
||||
chunk_index: int,
|
||||
total_chunks: int,
|
||||
event_date: datetime,
|
||||
event_date: datetime | None,
|
||||
context: str,
|
||||
llm_config: "LLMConfig",
|
||||
config,
|
||||
@@ -835,16 +944,15 @@ async def _extract_facts_from_chunk(
|
||||
user_message = _build_user_message(chunk, chunk_index, total_chunks, event_date, context, metadata)
|
||||
|
||||
# Retry logic for JSON validation errors
|
||||
max_retries = 2
|
||||
last_error = None
|
||||
# Use retain-specific overrides if set, otherwise fall back to global LLM config
|
||||
llm_max_retries = (
|
||||
config.retain_llm_max_retries if config.retain_llm_max_retries is not None else config.llm_max_retries
|
||||
)
|
||||
last_error: Exception | None = None
|
||||
|
||||
usage = TokenUsage() # Track cumulative usage across retries
|
||||
for attempt in range(max_retries):
|
||||
for attempt in range(llm_max_retries):
|
||||
try:
|
||||
# Use retain-specific overrides if set, otherwise fall back to global LLM config
|
||||
max_retries = (
|
||||
config.retain_llm_max_retries if config.retain_llm_max_retries is not None else config.llm_max_retries
|
||||
)
|
||||
initial_backoff = (
|
||||
config.retain_llm_initial_backoff
|
||||
if config.retain_llm_initial_backoff is not None
|
||||
@@ -860,7 +968,7 @@ async def _extract_facts_from_chunk(
|
||||
scope="retain_extract_facts",
|
||||
temperature=0.1,
|
||||
max_completion_tokens=config.retain_max_completion_tokens,
|
||||
max_retries=max_retries,
|
||||
max_retries=llm_max_retries,
|
||||
initial_backoff=initial_backoff,
|
||||
max_backoff=max_backoff,
|
||||
skip_validation=True, # Get raw JSON, we'll validate leniently
|
||||
@@ -874,14 +982,14 @@ async def _extract_facts_from_chunk(
|
||||
|
||||
# Handle malformed LLM responses
|
||||
if not isinstance(extraction_response_json, dict):
|
||||
if attempt < max_retries - 1:
|
||||
if attempt < llm_max_retries - 1:
|
||||
logger.warning(
|
||||
f"LLM returned non-dict JSON on attempt {attempt + 1}/{max_retries}: {type(extraction_response_json).__name__}. Retrying..."
|
||||
f"LLM returned non-dict JSON on attempt {attempt + 1}/{llm_max_retries}: {type(extraction_response_json).__name__}. Retrying..."
|
||||
)
|
||||
continue
|
||||
else:
|
||||
logger.warning(
|
||||
f"LLM returned non-dict JSON after {max_retries} attempts: {type(extraction_response_json).__name__}. "
|
||||
f"LLM returned non-dict JSON after {llm_max_retries} attempts: {type(extraction_response_json).__name__}. "
|
||||
f"Raw: {str(extraction_response_json)[:500]}"
|
||||
)
|
||||
return [], usage
|
||||
@@ -991,9 +1099,9 @@ async def _extract_facts_from_chunk(
|
||||
# Add entities if present (validate as Entity objects)
|
||||
# LLM sometimes returns strings instead of {"text": "..."} format
|
||||
entities = get_value("entities")
|
||||
validated_entities = []
|
||||
if entities:
|
||||
# Validate and normalize each entity
|
||||
validated_entities = []
|
||||
for ent in entities:
|
||||
if isinstance(ent, str):
|
||||
# Normalize string to Entity object
|
||||
@@ -1003,8 +1111,48 @@ async def _extract_facts_from_chunk(
|
||||
validated_entities.append(Entity.model_validate(ent))
|
||||
except Exception as e:
|
||||
logger.warning(f"Invalid entity {ent}: {e}")
|
||||
if validated_entities:
|
||||
fact_data["entities"] = validated_entities
|
||||
|
||||
# Post-process label entities from structured labels object
|
||||
entity_labels_raw = getattr(config, "entity_labels", None)
|
||||
labels_cfg = parse_entity_labels(entity_labels_raw)
|
||||
free_form_entities = getattr(config, "entities_allow_free_form", True)
|
||||
if labels_cfg and labels_cfg.attributes:
|
||||
labels_lookup = build_labels_lookup(labels_cfg)
|
||||
labels_data = llm_fact.get("labels") or {}
|
||||
if isinstance(labels_data, dict):
|
||||
existing_texts_lower = {e.text.lower() for e in validated_entities}
|
||||
for group in labels_cfg.attributes:
|
||||
value = labels_data.get(group.key)
|
||||
if not value:
|
||||
continue
|
||||
values_list = value if isinstance(value, list) else [value]
|
||||
for v in values_list:
|
||||
if not isinstance(v, str) or not v.strip() or v.lower() in ("none", "null", "n/a"):
|
||||
continue
|
||||
label_str = f"{group.key}:{v.strip()}"
|
||||
if group.type == "text":
|
||||
if label_str.lower() not in existing_texts_lower:
|
||||
validated_entities.append(Entity(text=label_str))
|
||||
existing_texts_lower.add(label_str.lower())
|
||||
elif (
|
||||
label_str.lower() in labels_lookup and label_str.lower() not in existing_texts_lower
|
||||
):
|
||||
validated_entities.append(Entity(text=label_str))
|
||||
existing_texts_lower.add(label_str.lower())
|
||||
else:
|
||||
logger.warning(f"Label '{label_str}' not in valid label values, skipping")
|
||||
|
||||
# In labels-only mode, keep only label entities
|
||||
if not free_form_entities:
|
||||
validated_entities = [
|
||||
e for e in validated_entities if is_label_entity(e.text, labels_cfg, labels_lookup)
|
||||
]
|
||||
elif not free_form_entities:
|
||||
# No labels but free_form disabled: clear all entities
|
||||
validated_entities = []
|
||||
|
||||
if validated_entities:
|
||||
fact_data["entities"] = validated_entities
|
||||
|
||||
# Add per-fact causal relations (only if enabled in config)
|
||||
if extract_causal_links:
|
||||
@@ -1043,8 +1191,9 @@ async def _extract_facts_from_chunk(
|
||||
if validated_relations:
|
||||
fact_data["causal_relations"] = validated_relations
|
||||
|
||||
# Always set mentioned_at to the event_date (when the conversation/document occurred)
|
||||
fact_data["mentioned_at"] = event_date.isoformat()
|
||||
# Set mentioned_at to the event_date (when the conversation/document occurred),
|
||||
# or None when the caller opted into no timestamp.
|
||||
fact_data["mentioned_at"] = event_date.isoformat() if event_date is not None else None
|
||||
|
||||
# Build Fact model instance
|
||||
try:
|
||||
@@ -1056,9 +1205,9 @@ async def _extract_facts_from_chunk(
|
||||
continue
|
||||
|
||||
# If we got malformed facts and haven't exhausted retries, try again
|
||||
if has_malformed_facts and len(chunk_facts) < len(raw_facts) * 0.8 and attempt < max_retries - 1:
|
||||
if has_malformed_facts and len(chunk_facts) < len(raw_facts) * 0.8 and attempt < llm_max_retries - 1:
|
||||
logger.warning(
|
||||
f"Got {len(raw_facts) - len(chunk_facts)} malformed facts out of {len(raw_facts)} on attempt {attempt + 1}/{max_retries}. Retrying..."
|
||||
f"Got {len(raw_facts) - len(chunk_facts)} malformed facts out of {len(raw_facts)} on attempt {attempt + 1}/{llm_max_retries}. Retrying..."
|
||||
)
|
||||
continue
|
||||
|
||||
@@ -1091,23 +1240,25 @@ async def _extract_facts_from_chunk(
|
||||
|
||||
if "json_validate_failed" in str(e):
|
||||
logger.warning(
|
||||
f" [1.3.{chunk_index + 1}] Attempt {attempt + 1}/{max_retries} failed with JSON validation error: {e}"
|
||||
f" [1.3.{chunk_index + 1}] Attempt {attempt + 1}/{llm_max_retries} failed with JSON validation error: {e}"
|
||||
)
|
||||
if attempt < max_retries - 1:
|
||||
if attempt < llm_max_retries - 1:
|
||||
logger.info(f" [1.3.{chunk_index + 1}] Retrying...")
|
||||
continue
|
||||
# If it's not a JSON validation error or we're out of retries, re-raise
|
||||
raise
|
||||
|
||||
# If we exhausted all retries, raise the last error
|
||||
raise last_error
|
||||
# If we exhausted all retries, raise the last error or a descriptive fallback
|
||||
if last_error is not None:
|
||||
raise last_error
|
||||
raise RuntimeError(f"Fact extraction failed after {llm_max_retries} attempts: LLM did not return valid JSON")
|
||||
|
||||
|
||||
async def _extract_facts_with_auto_split(
|
||||
chunk: str,
|
||||
chunk_index: int,
|
||||
total_chunks: int,
|
||||
event_date: datetime,
|
||||
event_date: datetime | None,
|
||||
context: str,
|
||||
llm_config: LLMConfig,
|
||||
config,
|
||||
@@ -1226,7 +1377,7 @@ async def _extract_facts_with_auto_split(
|
||||
|
||||
async def extract_facts_from_text(
|
||||
text: str,
|
||||
event_date: datetime,
|
||||
event_date: datetime | None,
|
||||
llm_config: LLMConfig,
|
||||
agent_name: str,
|
||||
config,
|
||||
@@ -1599,8 +1750,8 @@ async def extract_facts_from_contents_batch_api(
|
||||
|
||||
# Entities
|
||||
entities = get_value("entities")
|
||||
validated_entities = []
|
||||
if entities:
|
||||
validated_entities = []
|
||||
for ent in entities:
|
||||
if isinstance(ent, str):
|
||||
validated_entities.append(Entity(text=ent))
|
||||
@@ -1609,8 +1760,45 @@ async def extract_facts_from_contents_batch_api(
|
||||
validated_entities.append(Entity.model_validate(ent))
|
||||
except Exception:
|
||||
pass
|
||||
if validated_entities:
|
||||
fact_data["entities"] = validated_entities
|
||||
|
||||
# Post-process label entities from structured labels object
|
||||
entity_labels_raw = getattr(config, "entity_labels", None)
|
||||
labels_cfg_batch = parse_entity_labels(entity_labels_raw)
|
||||
free_form_entities_batch = getattr(config, "entities_allow_free_form", True)
|
||||
if labels_cfg_batch and labels_cfg_batch.attributes:
|
||||
labels_lookup_batch = build_labels_lookup(labels_cfg_batch)
|
||||
labels_data = llm_fact.get("labels") or {}
|
||||
if isinstance(labels_data, dict):
|
||||
existing_texts_lower = {e.text.lower() for e in validated_entities}
|
||||
for group in labels_cfg_batch.attributes:
|
||||
value = labels_data.get(group.key)
|
||||
if not value:
|
||||
continue
|
||||
values_list = value if isinstance(value, list) else [value]
|
||||
for v in values_list:
|
||||
if not isinstance(v, str) or not v.strip() or v.lower() in ("none", "null", "n/a"):
|
||||
continue
|
||||
label_str = f"{group.key}:{v.strip()}"
|
||||
if group.type == "text":
|
||||
if label_str.lower() not in existing_texts_lower:
|
||||
validated_entities.append(Entity(text=label_str))
|
||||
existing_texts_lower.add(label_str.lower())
|
||||
elif (
|
||||
label_str.lower() in labels_lookup_batch
|
||||
and label_str.lower() not in existing_texts_lower
|
||||
):
|
||||
validated_entities.append(Entity(text=label_str))
|
||||
existing_texts_lower.add(label_str.lower())
|
||||
|
||||
if not free_form_entities_batch:
|
||||
validated_entities = [
|
||||
e for e in validated_entities if is_label_entity(e.text, labels_cfg_batch, labels_lookup_batch)
|
||||
]
|
||||
elif not free_form_entities_batch:
|
||||
validated_entities = []
|
||||
|
||||
if validated_entities:
|
||||
fact_data["entities"] = validated_entities
|
||||
|
||||
# Causal relations
|
||||
if extract_causal_links:
|
||||
@@ -1641,8 +1829,9 @@ async def extract_facts_from_contents_batch_api(
|
||||
if validated_relations:
|
||||
fact_data["causal_relations"] = validated_relations
|
||||
|
||||
# Always set mentioned_at
|
||||
fact_data["mentioned_at"] = event_date.isoformat()
|
||||
# Set mentioned_at to the event_date (when the conversation/document occurred),
|
||||
# or None when the caller opted into no timestamp.
|
||||
fact_data["mentioned_at"] = event_date.isoformat() if event_date is not None else None
|
||||
|
||||
try:
|
||||
fact = Fact(fact=combined_text, fact_type=fact_type, **fact_data)
|
||||
@@ -1701,6 +1890,7 @@ async def extract_facts_from_contents_batch_api(
|
||||
mentioned_at=content.event_date,
|
||||
metadata=content.metadata,
|
||||
tags=content.tags,
|
||||
observation_scopes=content.observation_scopes,
|
||||
)
|
||||
|
||||
extracted_facts.append(extracted_fact)
|
||||
@@ -1709,6 +1899,9 @@ async def extract_facts_from_contents_batch_api(
|
||||
# Step 7: Add temporal offsets
|
||||
_add_temporal_offsets(extracted_facts, contents)
|
||||
|
||||
# Step 8: Auto-tag facts from label groups with tag=True
|
||||
_inject_label_tags(extracted_facts, config)
|
||||
|
||||
logger.info(f"Batch API extracted {len(extracted_facts)} facts from {len(all_chunks_info)} chunks")
|
||||
|
||||
return extracted_facts, chunks_metadata, total_usage
|
||||
@@ -1831,6 +2024,7 @@ async def extract_facts_from_contents(
|
||||
mentioned_at=content.event_date,
|
||||
metadata=content.metadata,
|
||||
tags=content.tags,
|
||||
observation_scopes=content.observation_scopes,
|
||||
)
|
||||
|
||||
extracted_facts.append(extracted_fact)
|
||||
@@ -1840,6 +2034,9 @@ async def extract_facts_from_contents(
|
||||
# Step 4: Add time offsets to preserve ordering within each content
|
||||
_add_temporal_offsets(extracted_facts, contents)
|
||||
|
||||
# Step 5: Auto-tag facts from label groups with tag=True
|
||||
_inject_label_tags(extracted_facts, config)
|
||||
|
||||
return extracted_facts, chunks_metadata, total_usage
|
||||
|
||||
|
||||
@@ -1895,3 +2092,24 @@ def _add_temporal_offsets(facts: list[ExtractedFactType], contents: list[RetainC
|
||||
fact.occurred_end = parse_datetime_flexible(fact.occurred_end) + offset
|
||||
if fact.mentioned_at:
|
||||
fact.mentioned_at = parse_datetime_flexible(fact.mentioned_at) + offset
|
||||
|
||||
|
||||
def _inject_label_tags(facts: list[ExtractedFactType], config) -> None:
|
||||
"""
|
||||
For label groups with tag=True, add extracted key:value label entities
|
||||
to each fact's tags list. Modifies facts in place.
|
||||
|
||||
This lets entity labels double as tags, enabling filtering via the
|
||||
existing tags API without any extra query infrastructure.
|
||||
"""
|
||||
labels_cfg = parse_entity_labels(getattr(config, "entity_labels", None))
|
||||
if not labels_cfg:
|
||||
return
|
||||
tag_group_keys = {g.key.lower() for g in labels_cfg.attributes if g.tag}
|
||||
if not tag_group_keys:
|
||||
return
|
||||
for fact in facts:
|
||||
label_tags = [e for e in fact.entities if ":" in e and e.split(":", 1)[0].lower() in tag_group_keys]
|
||||
if label_tags:
|
||||
existing = set(fact.tags)
|
||||
fact.tags = fact.tags + [t for t in label_tags if t not in existing]
|
||||
|
||||
@@ -47,6 +47,8 @@ async def insert_facts_batch(
|
||||
chunk_ids = []
|
||||
document_ids = []
|
||||
tags_list = []
|
||||
observation_scopes_list = []
|
||||
text_signals_list = []
|
||||
|
||||
for fact in facts:
|
||||
fact_texts.append(_sanitize_text(fact.fact_text))
|
||||
@@ -68,6 +70,19 @@ async def insert_facts_batch(
|
||||
document_ids.append(fact.document_id if fact.document_id else document_id)
|
||||
# Convert tags to JSON string for proper batch insertion (PostgreSQL unnest doesn't handle 2D arrays well)
|
||||
tags_list.append(json.dumps(fact.tags if fact.tags else []))
|
||||
# observation_scopes: stored as JSONB (string or 2D array), None if not provided
|
||||
observation_scopes_list.append(
|
||||
json.dumps(fact.observation_scopes) if fact.observation_scopes is not None else None
|
||||
)
|
||||
# Build text_signals: entity names + date tokens for enriched BM25 indexing
|
||||
signal_parts = []
|
||||
if fact.entities:
|
||||
signal_parts.extend(e.name for e in fact.entities)
|
||||
if fact.occurred_start:
|
||||
signal_parts.append(fact.occurred_start.strftime("%B %-d %Y"))
|
||||
if fact.occurred_end and fact.occurred_end != fact.occurred_start:
|
||||
signal_parts.append(fact.occurred_end.strftime("%B %-d %Y"))
|
||||
text_signals_list.append(" ".join(signal_parts) if signal_parts else None)
|
||||
|
||||
# Batch insert all facts
|
||||
# Note: tags are passed as JSON strings and converted back to varchar[] via jsonb_array_elements_text + array_agg
|
||||
@@ -75,16 +90,19 @@ async def insert_facts_batch(
|
||||
config = get_config()
|
||||
if config.text_search_extension == "vchord":
|
||||
# VectorChord: manually tokenize and insert search_vector
|
||||
# text_signals (entity names etc.) are included in the tokenize input for enriched BM25
|
||||
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::float[], $11::jsonb[], $12::text[], $13::text[], $14::jsonb[]
|
||||
$8::text[], $9::text[], $10::float[], $11::jsonb[], $12::text[], $13::text[], $14::jsonb[], $15::jsonb[], $16::text[]
|
||||
) AS t(text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
context, fact_type, confidence_score, metadata, chunk_id, document_id, tags_json)
|
||||
context, fact_type, confidence_score, metadata, chunk_id, document_id, tags_json,
|
||||
observation_scopes_json, text_signals)
|
||||
)
|
||||
INSERT INTO {fq_table("memory_units")} (bank_id, text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
context, fact_type, confidence_score, metadata, chunk_id, document_id, tags, search_vector)
|
||||
context, fact_type, confidence_score, metadata, chunk_id, document_id, tags,
|
||||
observation_scopes, text_signals, search_vector)
|
||||
SELECT
|
||||
$1,
|
||||
text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
@@ -93,23 +111,30 @@ async def insert_facts_batch(
|
||||
(SELECT array_agg(elem) FROM jsonb_array_elements_text(tags_json) AS elem),
|
||||
'{{}}'::varchar[]
|
||||
),
|
||||
tokenize(COALESCE(text, '') || ' ' || COALESCE(context, ''), 'llmlingua2')::bm25_catalog.bm25vector
|
||||
observation_scopes_json,
|
||||
text_signals,
|
||||
tokenize(
|
||||
COALESCE(text, '') || ' ' || COALESCE(context, '') || ' ' || COALESCE(text_signals, ''),
|
||||
'llmlingua2'
|
||||
)::bm25_catalog.bm25vector
|
||||
FROM input_data
|
||||
RETURNING id
|
||||
"""
|
||||
else: # native or pg_textsearch
|
||||
# Native PostgreSQL: search_vector is GENERATED ALWAYS, don't include it
|
||||
# Native PostgreSQL: search_vector is GENERATED ALWAYS (expression includes text_signals), don't include it
|
||||
# pg_textsearch: indexes operate on base columns directly, don't populate search_vector
|
||||
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::float[], $11::jsonb[], $12::text[], $13::text[], $14::jsonb[]
|
||||
$8::text[], $9::text[], $10::float[], $11::jsonb[], $12::text[], $13::text[], $14::jsonb[], $15::jsonb[], $16::text[]
|
||||
) AS t(text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
context, fact_type, confidence_score, metadata, chunk_id, document_id, tags_json)
|
||||
context, fact_type, confidence_score, metadata, chunk_id, document_id, tags_json,
|
||||
observation_scopes_json, text_signals)
|
||||
)
|
||||
INSERT INTO {fq_table("memory_units")} (bank_id, text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
context, fact_type, confidence_score, metadata, chunk_id, document_id, tags)
|
||||
context, fact_type, confidence_score, metadata, chunk_id, document_id, tags,
|
||||
observation_scopes, text_signals)
|
||||
SELECT
|
||||
$1,
|
||||
text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
@@ -117,7 +142,9 @@ async def insert_facts_batch(
|
||||
COALESCE(
|
||||
(SELECT array_agg(elem) FROM jsonb_array_elements_text(tags_json) AS elem),
|
||||
'{{}}'::varchar[]
|
||||
)
|
||||
),
|
||||
observation_scopes_json,
|
||||
text_signals
|
||||
FROM input_data
|
||||
RETURNING id
|
||||
"""
|
||||
@@ -138,6 +165,8 @@ async def insert_facts_batch(
|
||||
chunk_ids,
|
||||
document_ids,
|
||||
tags_list,
|
||||
observation_scopes_list,
|
||||
text_signals_list,
|
||||
)
|
||||
|
||||
unit_ids = [str(row["id"]) for row in results]
|
||||
|
||||
@@ -47,6 +47,9 @@ def compute_temporal_links(
|
||||
|
||||
links = []
|
||||
for unit_id, unit_event_date in new_units.items():
|
||||
# Units without event_date can't form temporal links
|
||||
if unit_event_date is None:
|
||||
continue
|
||||
# Normalize unit_event_date for consistent comparison
|
||||
unit_event_date_norm = _normalize_datetime(unit_event_date)
|
||||
|
||||
@@ -96,7 +99,11 @@ def compute_temporal_query_bounds(
|
||||
return None, None
|
||||
|
||||
# Normalize all dates to be timezone-aware to avoid comparison issues
|
||||
all_dates = [_normalize_datetime(d) for d in new_units.values()]
|
||||
# Filter out None values — units without event_date can't form temporal links
|
||||
all_dates = [_normalize_datetime(d) for d in new_units.values() if d is not None]
|
||||
|
||||
if not all_dates:
|
||||
return None, None
|
||||
|
||||
try:
|
||||
min_date = min(all_dates) - timedelta(hours=time_window_hours)
|
||||
@@ -143,6 +150,7 @@ async def extract_entities_batch_optimized(
|
||||
fact_dates: list,
|
||||
llm_entities: list[list[dict]],
|
||||
log_buffer: list[str] = None,
|
||||
entity_labels: list | None = None,
|
||||
) -> list[tuple]:
|
||||
"""
|
||||
Process LLM-extracted entities for ALL facts in batch.
|
||||
@@ -232,6 +240,7 @@ async def extract_entities_batch_optimized(
|
||||
context=context,
|
||||
unit_event_date=None, # Not used when per-entity dates provided
|
||||
conn=conn, # Use main transaction connection
|
||||
entity_labels=entity_labels,
|
||||
)
|
||||
|
||||
_log(
|
||||
@@ -432,20 +441,23 @@ async def create_temporal_links_batch_per_fact(
|
||||
min_date, max_date = compute_temporal_query_bounds(new_units, time_window_hours)
|
||||
|
||||
fetch_neighbors_start = time_mod.time()
|
||||
all_candidates = await conn.fetch(
|
||||
f"""
|
||||
SELECT id, event_date
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE bank_id = $1
|
||||
AND event_date BETWEEN $2 AND $3
|
||||
AND id::text != ALL($4)
|
||||
ORDER BY event_date DESC
|
||||
""",
|
||||
bank_id,
|
||||
min_date,
|
||||
max_date,
|
||||
unit_ids,
|
||||
)
|
||||
if min_date is not None and max_date is not None:
|
||||
all_candidates = await conn.fetch(
|
||||
f"""
|
||||
SELECT id, event_date
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE bank_id = $1
|
||||
AND event_date BETWEEN $2 AND $3
|
||||
AND id::text != ALL($4)
|
||||
ORDER BY event_date DESC
|
||||
""",
|
||||
bank_id,
|
||||
min_date,
|
||||
max_date,
|
||||
unit_ids,
|
||||
)
|
||||
else:
|
||||
all_candidates = []
|
||||
_log(
|
||||
log_buffer,
|
||||
f" [7.2] Fetch {len(all_candidates)} candidate neighbors (1 query): {time_mod.time() - fetch_neighbors_start:.3f}s",
|
||||
@@ -460,11 +472,15 @@ async def create_temporal_links_batch_per_fact(
|
||||
# Convert new_units dict to candidate format for within-batch linking
|
||||
new_unit_items = list(new_units.items())
|
||||
for i, (unit_id, event_date) in enumerate(new_unit_items):
|
||||
if event_date is None:
|
||||
continue # Skip units without event_date for temporal linking
|
||||
unit_event_date_norm = _normalize_datetime(event_date)
|
||||
|
||||
# Compare with other new units (only those after this one to avoid duplicates)
|
||||
for j in range(i + 1, len(new_unit_items)):
|
||||
other_id, other_event_date = new_unit_items[j]
|
||||
if other_event_date is None:
|
||||
continue # Skip units without event_date
|
||||
other_event_date_norm = _normalize_datetime(other_event_date)
|
||||
|
||||
# Check if within time window
|
||||
@@ -482,14 +498,13 @@ async def create_temporal_links_batch_per_fact(
|
||||
# Batch inserts to avoid timeout on large batches
|
||||
BATCH_SIZE = 1000
|
||||
for batch_start in range(0, len(links), BATCH_SIZE):
|
||||
batch = links[batch_start : batch_start + BATCH_SIZE]
|
||||
await conn.executemany(
|
||||
f"""
|
||||
INSERT INTO {fq_table("memory_links")} (from_unit_id, to_unit_id, link_type, weight, entity_id)
|
||||
VALUES ($1, $2, $3, $4, $5)
|
||||
ON CONFLICT (from_unit_id, to_unit_id, link_type, COALESCE(entity_id, '00000000-0000-0000-0000-000000000000'::uuid)) DO NOTHING
|
||||
""",
|
||||
batch,
|
||||
links[batch_start : batch_start + BATCH_SIZE],
|
||||
)
|
||||
_log(log_buffer, f" [7.4] Insert {len(links)} temporal links: {time_mod.time() - insert_start:.3f}s")
|
||||
|
||||
@@ -537,81 +552,45 @@ async def create_semantic_links_batch(
|
||||
|
||||
import numpy as np
|
||||
|
||||
# Fetch ALL existing units with embeddings in ONE query
|
||||
fetch_start = time_mod.time()
|
||||
all_existing = await conn.fetch(
|
||||
f"""
|
||||
SELECT id, embedding
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE bank_id = $1
|
||||
AND embedding IS NOT NULL
|
||||
AND id::text != ALL($2)
|
||||
""",
|
||||
bank_id,
|
||||
unit_ids,
|
||||
)
|
||||
_log(
|
||||
log_buffer,
|
||||
f" [8.1] Fetch {len(all_existing)} existing embeddings (1 query): {time_mod.time() - fetch_start:.3f}s",
|
||||
)
|
||||
|
||||
# Convert to numpy for vectorized similarity computation
|
||||
compute_start = time_mod.time()
|
||||
# Use pgvector ANN search (HNSW index) for each new unit instead of fetching
|
||||
# all existing embeddings into Python. At large scale (100K+ units) the old
|
||||
# approach would transfer 100K × 384 floats (~150 MB) per retain call; the
|
||||
# ANN query completes in <5 ms and transfers only top_k rows.
|
||||
ann_start = time_mod.time()
|
||||
all_links = []
|
||||
|
||||
if all_existing:
|
||||
# Convert existing embeddings to numpy array
|
||||
existing_ids = [str(row["id"]) for row in all_existing]
|
||||
# Stack embeddings as 2D array: (num_embeddings, embedding_dim)
|
||||
embedding_arrays = []
|
||||
for row in all_existing:
|
||||
raw_emb = row["embedding"]
|
||||
# Handle different pgvector formats
|
||||
if isinstance(raw_emb, str):
|
||||
# Parse string format: "[1.0, 2.0, ...]"
|
||||
import json
|
||||
# Build UUID exclude list once for all ANN queries
|
||||
import uuid as uuid_mod
|
||||
|
||||
emb = np.array(json.loads(raw_emb), dtype=np.float32)
|
||||
elif isinstance(raw_emb, (list, tuple)):
|
||||
emb = np.array(raw_emb, dtype=np.float32)
|
||||
else:
|
||||
# Try direct conversion (works for numpy arrays, pgvector objects, etc.)
|
||||
emb = np.array(raw_emb, dtype=np.float32)
|
||||
exclude_uuids = [uuid_mod.UUID(uid) if isinstance(uid, str) else uid for uid in unit_ids]
|
||||
|
||||
# Ensure it's 1D
|
||||
if emb.ndim != 1:
|
||||
raise ValueError(f"Expected 1D embedding, got shape {emb.shape}")
|
||||
embedding_arrays.append(emb)
|
||||
for unit_id, new_embedding in zip(unit_ids, embeddings):
|
||||
emb_str = str(list(new_embedding) if not isinstance(new_embedding, list) else new_embedding)
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT id::text,
|
||||
1 - (embedding <=> $1::vector) AS similarity
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE bank_id = $2
|
||||
AND embedding IS NOT NULL
|
||||
AND id != ALL($3::uuid[])
|
||||
ORDER BY embedding <=> $1::vector
|
||||
LIMIT $4
|
||||
""",
|
||||
emb_str,
|
||||
bank_id,
|
||||
exclude_uuids,
|
||||
top_k,
|
||||
)
|
||||
for row in rows:
|
||||
sim = float(min(1.0, max(0.0, row["similarity"])))
|
||||
if sim >= threshold:
|
||||
all_links.append((unit_id, str(row["id"]), "semantic", sim, None))
|
||||
|
||||
if not embedding_arrays:
|
||||
existing_embeddings = np.array([])
|
||||
elif len(embedding_arrays) == 1:
|
||||
# Single embedding: reshape to (1, dim)
|
||||
existing_embeddings = embedding_arrays[0].reshape(1, -1)
|
||||
else:
|
||||
# Multiple embeddings: vstack
|
||||
existing_embeddings = np.vstack(embedding_arrays)
|
||||
|
||||
# For each new unit, compute similarities with ALL existing units
|
||||
for unit_id, new_embedding in zip(unit_ids, embeddings):
|
||||
new_emb_array = np.array(new_embedding)
|
||||
|
||||
# Compute cosine similarities (dot product for normalized vectors)
|
||||
similarities = np.dot(existing_embeddings, new_emb_array)
|
||||
|
||||
# Find top-k above threshold
|
||||
# Get indices of similarities above threshold
|
||||
above_threshold = np.where(similarities >= threshold)[0]
|
||||
|
||||
if len(above_threshold) > 0:
|
||||
# Sort by similarity (descending) and take top-k
|
||||
sorted_indices = above_threshold[np.argsort(-similarities[above_threshold])][:top_k]
|
||||
|
||||
for idx in sorted_indices:
|
||||
similar_id = existing_ids[idx]
|
||||
# Clamp to [0, 1] to handle floating point precision issues
|
||||
similarity = float(min(1.0, max(0.0, similarities[idx])))
|
||||
all_links.append((unit_id, similar_id, "semantic", similarity, None))
|
||||
_log(
|
||||
log_buffer,
|
||||
f" [8.1] ANN search for {len(unit_ids)} new units → {len(all_links)} candidate links: {time_mod.time() - ann_start:.3f}s",
|
||||
)
|
||||
|
||||
# Also compute similarities WITHIN the new batch (new units to each other)
|
||||
# Apply the same top_k limit per unit as we do for existing units
|
||||
@@ -643,7 +622,7 @@ async def create_semantic_links_batch(
|
||||
|
||||
_log(
|
||||
log_buffer,
|
||||
f" [8.2] Compute similarities & generate {len(all_links)} semantic links: {time_mod.time() - compute_start:.3f}s",
|
||||
f" [8.2] Within-batch similarities added {len(all_links)} total semantic links",
|
||||
)
|
||||
|
||||
if all_links:
|
||||
@@ -651,14 +630,13 @@ async def create_semantic_links_batch(
|
||||
# Batch inserts to avoid timeout on large batches
|
||||
BATCH_SIZE = 1000
|
||||
for batch_start in range(0, len(all_links), BATCH_SIZE):
|
||||
batch = all_links[batch_start : batch_start + BATCH_SIZE]
|
||||
await conn.executemany(
|
||||
f"""
|
||||
INSERT INTO {fq_table("memory_links")} (from_unit_id, to_unit_id, link_type, weight, entity_id)
|
||||
VALUES ($1, $2, $3, $4, $5)
|
||||
ON CONFLICT (from_unit_id, to_unit_id, link_type, COALESCE(entity_id, '00000000-0000-0000-0000-000000000000'::uuid)) DO NOTHING
|
||||
""",
|
||||
batch,
|
||||
all_links[batch_start : batch_start + BATCH_SIZE],
|
||||
)
|
||||
_log(
|
||||
log_buffer, f" [8.3] Insert {len(all_links)} semantic links: {time_mod.time() - insert_start:.3f}s"
|
||||
@@ -674,18 +652,18 @@ async def create_semantic_links_batch(
|
||||
raise
|
||||
|
||||
|
||||
async def insert_entity_links_batch(conn, links: list[EntityLink], chunk_size: int = 50000):
|
||||
async def insert_entity_links_batch(conn, links: list[EntityLink], chunk_size: int = 5000):
|
||||
"""
|
||||
Insert all entity links using COPY to temp table + INSERT for maximum speed.
|
||||
Insert all entity links using COPY to temp table + chunked INSERT for reliability.
|
||||
|
||||
Uses PostgreSQL COPY (via copy_records_to_table) for bulk loading,
|
||||
then INSERT ... ON CONFLICT from temp table. This is the fastest
|
||||
method for bulk inserts with conflict handling.
|
||||
Uses PostgreSQL COPY (via copy_records_to_table) for bulk loading into a
|
||||
temp table, then INSERT ... ON CONFLICT in chunks of chunk_size. Chunking
|
||||
prevents single-query timeouts on very large tables (100M+ rows).
|
||||
|
||||
Args:
|
||||
conn: Database connection
|
||||
links: List of EntityLink objects
|
||||
chunk_size: Number of rows per batch (default 50000)
|
||||
chunk_size: Number of rows per INSERT chunk (default 5000)
|
||||
"""
|
||||
if not links:
|
||||
return
|
||||
@@ -694,10 +672,11 @@ async def insert_entity_links_batch(conn, links: list[EntityLink], chunk_size: i
|
||||
|
||||
total_start = time_mod.time()
|
||||
|
||||
# Create temp table for bulk loading
|
||||
# Create temp table with serial for stable chunked access
|
||||
create_start = time_mod.time()
|
||||
await conn.execute("""
|
||||
CREATE TEMP TABLE IF NOT EXISTS _temp_entity_links (
|
||||
_row_num SERIAL,
|
||||
from_unit_id uuid,
|
||||
to_unit_id uuid,
|
||||
link_type text,
|
||||
@@ -714,9 +693,7 @@ async def insert_entity_links_batch(conn, links: list[EntityLink], chunk_size: i
|
||||
|
||||
# Convert EntityLink objects to tuples for COPY
|
||||
convert_start = time_mod.time()
|
||||
records = []
|
||||
for link in links:
|
||||
records.append((link.from_unit_id, link.to_unit_id, link.link_type, link.weight, link.entity_id))
|
||||
records = [(link.from_unit_id, link.to_unit_id, link.link_type, link.weight, link.entity_id) for link in links]
|
||||
logger.debug(f" [9.3] Convert {len(records)} records: {time_mod.time() - convert_start:.3f}s")
|
||||
|
||||
# Bulk load using COPY (fastest method)
|
||||
@@ -728,15 +705,25 @@ async def insert_entity_links_batch(conn, links: list[EntityLink], chunk_size: i
|
||||
)
|
||||
logger.debug(f" [9.4] COPY {len(records)} records to temp table: {time_mod.time() - copy_start:.3f}s")
|
||||
|
||||
# Insert from temp table with ON CONFLICT (single query for all rows)
|
||||
# Insert from temp table in chunks to avoid single-query timeouts on large tables
|
||||
insert_start = time_mod.time()
|
||||
await conn.execute(f"""
|
||||
INSERT INTO {fq_table("memory_links")} (from_unit_id, to_unit_id, link_type, weight, entity_id)
|
||||
SELECT from_unit_id, to_unit_id, link_type, weight, entity_id
|
||||
FROM _temp_entity_links
|
||||
ON CONFLICT (from_unit_id, to_unit_id, link_type, COALESCE(entity_id, '00000000-0000-0000-0000-000000000000'::uuid)) DO NOTHING
|
||||
""")
|
||||
logger.debug(f" [9.5] INSERT from temp table: {time_mod.time() - insert_start:.3f}s")
|
||||
total_rows = len(records)
|
||||
chunks = 0
|
||||
for chunk_start in range(0, total_rows, chunk_size):
|
||||
chunk_end = chunk_start + chunk_size
|
||||
await conn.execute(
|
||||
f"""
|
||||
INSERT INTO {fq_table("memory_links")} (from_unit_id, to_unit_id, link_type, weight, entity_id)
|
||||
SELECT from_unit_id, to_unit_id, link_type, weight, entity_id
|
||||
FROM _temp_entity_links
|
||||
WHERE _row_num > $1 AND _row_num <= $2
|
||||
ON CONFLICT (from_unit_id, to_unit_id, link_type, COALESCE(entity_id, '00000000-0000-0000-0000-000000000000'::uuid)) DO NOTHING
|
||||
""",
|
||||
chunk_start,
|
||||
chunk_end,
|
||||
)
|
||||
chunks += 1
|
||||
logger.debug(f" [9.5] INSERT {total_rows} rows in {chunks} chunks: {time_mod.time() - insert_start:.3f}s")
|
||||
logger.debug(f" [9.TOTAL] Entity links batch insert: {time_mod.time() - total_start:.3f}s")
|
||||
|
||||
|
||||
|
||||
@@ -7,6 +7,7 @@ Coordinates all retain pipeline modules to store memories efficiently.
|
||||
import logging
|
||||
import time
|
||||
import uuid
|
||||
from collections.abc import Awaitable, Callable
|
||||
from datetime import UTC, datetime
|
||||
from typing import Any
|
||||
|
||||
@@ -52,10 +53,11 @@ def parse_datetime_flexible(value: Any) -> datetime:
|
||||
raise TypeError(f"Expected datetime or string, got {type(value).__name__}")
|
||||
|
||||
|
||||
import asyncpg
|
||||
|
||||
from ..response_models import TokenUsage
|
||||
from . import (
|
||||
chunk_storage,
|
||||
deduplication,
|
||||
embedding_processing,
|
||||
entity_processing,
|
||||
fact_extraction,
|
||||
@@ -73,7 +75,6 @@ async def retain_batch(
|
||||
llm_config,
|
||||
entity_resolver,
|
||||
format_date_fn,
|
||||
duplicate_checker_fn,
|
||||
bank_id: str,
|
||||
contents_dicts: list[RetainContentDict],
|
||||
config,
|
||||
@@ -84,6 +85,7 @@ 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,
|
||||
) -> tuple[list[list[str]], TokenUsage]:
|
||||
"""
|
||||
Process a batch of content through the retain pipeline.
|
||||
@@ -94,7 +96,6 @@ async def retain_batch(
|
||||
llm_config: LLM configuration for fact extraction
|
||||
entity_resolver: Entity resolver for entity processing
|
||||
format_date_fn: Function to format datetime to readable string
|
||||
duplicate_checker_fn: Function to check for duplicate facts
|
||||
bank_id: Bank identifier
|
||||
contents_dicts: List of content dictionaries
|
||||
config: Resolved HindsightConfig for this bank
|
||||
@@ -128,12 +129,14 @@ async def retain_batch(
|
||||
item_tags = item.get("tags", []) or []
|
||||
merged_tags = list(set(item_tags + (document_tags or [])))
|
||||
|
||||
# Handle event_date: parse flexibly (handles both datetime objects and ISO strings)
|
||||
event_date_value = item.get("event_date")
|
||||
if event_date_value:
|
||||
event_date_value = parse_datetime_flexible(event_date_value)
|
||||
# Handle event_date: distinguish "not provided" (default to now) from
|
||||
# "explicitly None" (caller opted into no timestamp).
|
||||
if "event_date" in item and item["event_date"] is None:
|
||||
event_date_value = None # Caller explicitly signalled "unknown date"
|
||||
elif item.get("event_date"):
|
||||
event_date_value = parse_datetime_flexible(item["event_date"])
|
||||
else:
|
||||
event_date_value = utcnow()
|
||||
event_date_value = utcnow() # Backward-compatible default
|
||||
|
||||
content = RetainContent(
|
||||
content=item["content"],
|
||||
@@ -142,6 +145,7 @@ async def retain_batch(
|
||||
metadata=item.get("metadata", {}),
|
||||
entities=item.get("entities", []),
|
||||
tags=merged_tags,
|
||||
observation_scopes=item.get("observation_scopes"),
|
||||
)
|
||||
contents.append(content)
|
||||
|
||||
@@ -162,8 +166,6 @@ async def retain_batch(
|
||||
docs_tracked = 0
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
async with conn.transaction():
|
||||
await fact_storage.ensure_bank_exists(conn, bank_id)
|
||||
|
||||
# Group contents by document_id (consistent with normal path)
|
||||
contents_by_doc_early = defaultdict(list)
|
||||
for idx, content_dict in enumerate(contents_dicts):
|
||||
@@ -281,9 +283,6 @@ async def retain_batch(
|
||||
# Step 4: Database transaction
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
async with conn.transaction():
|
||||
# Ensure bank exists
|
||||
await fact_storage.ensure_bank_exists(conn, bank_id)
|
||||
|
||||
# Handle document tracking for all documents
|
||||
step_start = time.time()
|
||||
# Map None document_id to generated UUIDs
|
||||
@@ -435,20 +434,7 @@ async def retain_batch(
|
||||
actual_doc_id = document_id
|
||||
processed_fact.document_id = actual_doc_id
|
||||
|
||||
# Deduplication
|
||||
step_start = time.time()
|
||||
is_duplicate_flags = await deduplication.check_duplicates_batch(
|
||||
conn, bank_id, processed_facts, duplicate_checker_fn
|
||||
)
|
||||
log_buffer.append(
|
||||
f"[4] Deduplication: {sum(is_duplicate_flags)} duplicates in {time.time() - step_start:.3f}s"
|
||||
)
|
||||
|
||||
# Filter out duplicates
|
||||
non_duplicate_facts = deduplication.filter_duplicates(processed_facts, is_duplicate_flags)
|
||||
|
||||
if not non_duplicate_facts:
|
||||
return [[] for _ in contents], usage
|
||||
non_duplicate_facts = processed_facts
|
||||
|
||||
# Insert facts (document_id is now stored per-fact)
|
||||
step_start = time.time()
|
||||
@@ -469,6 +455,7 @@ async def retain_batch(
|
||||
non_duplicate_facts,
|
||||
log_buffer,
|
||||
user_entities_per_content=user_entities_per_content,
|
||||
entity_labels=getattr(config, "entity_labels", None),
|
||||
)
|
||||
log_buffer.append(f"[6] Process entities: {len(entity_links)} links in {time.time() - step_start:.3f}s")
|
||||
|
||||
@@ -499,7 +486,16 @@ async def retain_batch(
|
||||
log_buffer.append(f"[10] Causal links: {causal_link_count} links in {time.time() - step_start:.3f}s")
|
||||
|
||||
# Map results back to original content items
|
||||
result_unit_ids = _map_results_to_contents(contents, extracted_facts, is_duplicate_flags, unit_ids)
|
||||
result_unit_ids = _map_results_to_contents(contents, extracted_facts, unit_ids)
|
||||
|
||||
# Transactional outbox: queue any side-effect tasks (e.g. webhook deliveries)
|
||||
# inside the same transaction so they are atomically committed with the retain data.
|
||||
if outbox_callback:
|
||||
await outbox_callback(conn)
|
||||
|
||||
# Flush entity stats (mention_count / last_seen) now that the transaction
|
||||
# has committed. Uses a fresh pool connection — no locks held.
|
||||
await entity_resolver.flush_pending_stats()
|
||||
|
||||
# Log final summary
|
||||
total_time = time.time() - start_time
|
||||
@@ -517,28 +513,20 @@ async def retain_batch(
|
||||
def _map_results_to_contents(
|
||||
contents: list[RetainContent],
|
||||
extracted_facts: list[ExtractedFact],
|
||||
is_duplicate_flags: list[bool],
|
||||
unit_ids: list[str],
|
||||
) -> list[list[str]]:
|
||||
"""
|
||||
Map created unit IDs back to original content items.
|
||||
|
||||
Accounts for duplicates when mapping back.
|
||||
"""
|
||||
result_unit_ids = []
|
||||
filtered_idx = 0
|
||||
|
||||
# Group facts by content_index
|
||||
facts_by_content = {i: [] for i in range(len(contents))}
|
||||
"""Map created unit IDs back to original content items."""
|
||||
facts_by_content: dict[int, list[int]] = {i: [] for i in range(len(contents))}
|
||||
for i, fact in enumerate(extracted_facts):
|
||||
facts_by_content[fact.content_index].append(i)
|
||||
|
||||
result_unit_ids = []
|
||||
unit_idx = 0
|
||||
for content_index in range(len(contents)):
|
||||
content_unit_ids = []
|
||||
for fact_idx in facts_by_content[content_index]:
|
||||
if not is_duplicate_flags[fact_idx]:
|
||||
content_unit_ids.append(unit_ids[filtered_idx])
|
||||
filtered_idx += 1
|
||||
for _ in facts_by_content[content_index]:
|
||||
content_unit_ids.append(unit_ids[unit_idx])
|
||||
unit_idx += 1
|
||||
result_unit_ids.append(content_unit_ids)
|
||||
|
||||
return result_unit_ids
|
||||
|
||||
@@ -6,8 +6,8 @@ from content input to fact storage.
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import UTC, datetime
|
||||
from typing import TypedDict
|
||||
from datetime import datetime
|
||||
from typing import Literal, TypedDict
|
||||
from uuid import UUID
|
||||
|
||||
|
||||
@@ -22,20 +22,21 @@ class RetainContentDict(TypedDict, total=False):
|
||||
document_id: Document ID for this content item (optional)
|
||||
entities: User-provided entities to merge with extracted entities (optional)
|
||||
tags: Visibility scope tags for this content item (optional)
|
||||
observation_scopes: How to scope observations for consolidation (optional).
|
||||
"per_tag" runs one pass per individual tag; "combined" (default) runs a
|
||||
single pass with all tags; a list[list[str]] specifies exact passes.
|
||||
"""
|
||||
|
||||
content: str # Required
|
||||
context: str
|
||||
event_date: datetime
|
||||
event_date: datetime | None
|
||||
metadata: dict[str, str]
|
||||
document_id: str
|
||||
entities: list[dict[str, str]] # [{"text": "...", "type": "..."}]
|
||||
tags: list[str] # Visibility scope tags
|
||||
|
||||
|
||||
def _now_utc() -> datetime:
|
||||
"""Factory function for default event_date."""
|
||||
return datetime.now(UTC)
|
||||
observation_scopes: (
|
||||
Literal["per_tag", "combined", "all_combinations"] | list[list[str]]
|
||||
) # Observation scopes for consolidation
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -48,10 +49,13 @@ class RetainContent:
|
||||
|
||||
content: str
|
||||
context: str = ""
|
||||
event_date: datetime = field(default_factory=_now_utc)
|
||||
event_date: datetime | None = None
|
||||
metadata: dict[str, str] = field(default_factory=dict)
|
||||
entities: list[dict[str, str]] = field(default_factory=list) # User-provided entities
|
||||
tags: list[str] = field(default_factory=list) # Visibility scope tags
|
||||
observation_scopes: Literal["per_tag", "combined", "all_combinations"] | list[list[str]] | None = (
|
||||
None # Observation scopes
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -117,6 +121,9 @@ class ExtractedFact:
|
||||
mentioned_at: datetime | None = None
|
||||
metadata: dict[str, str] = field(default_factory=dict)
|
||||
tags: list[str] = field(default_factory=list) # Visibility scope tags
|
||||
observation_scopes: Literal["per_tag", "combined", "all_combinations"] | list[list[str]] | None = (
|
||||
None # Observation scopes
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -135,7 +142,7 @@ class ProcessedFact:
|
||||
# Temporal data
|
||||
occurred_start: datetime | None
|
||||
occurred_end: datetime | None
|
||||
mentioned_at: datetime
|
||||
mentioned_at: datetime | None
|
||||
|
||||
# Context and metadata
|
||||
context: str
|
||||
@@ -165,6 +172,9 @@ class ProcessedFact:
|
||||
# Visibility scope tags
|
||||
tags: list[str] = field(default_factory=list)
|
||||
|
||||
# Observation scopes for consolidation
|
||||
observation_scopes: Literal["per_tag", "combined", "all_combinations"] | list[list[str]] | None = None
|
||||
|
||||
@property
|
||||
def is_duplicate(self) -> bool:
|
||||
"""Check if this fact was marked as a duplicate."""
|
||||
@@ -185,12 +195,10 @@ class ProcessedFact:
|
||||
Returns:
|
||||
ProcessedFact ready for storage
|
||||
"""
|
||||
from datetime import datetime
|
||||
|
||||
# Use occurred dates only if explicitly provided by LLM
|
||||
occurred_start = extracted_fact.occurred_start
|
||||
occurred_end = extracted_fact.occurred_end
|
||||
mentioned_at = extracted_fact.mentioned_at or datetime.now(UTC)
|
||||
mentioned_at = extracted_fact.mentioned_at # May be None when caller opted into no timestamp
|
||||
|
||||
# Convert entity strings to EntityRef objects
|
||||
entities = [EntityRef(name=name) for name in extracted_fact.entities]
|
||||
@@ -209,6 +217,7 @@ class ProcessedFact:
|
||||
chunk_id=chunk_id,
|
||||
content_index=extracted_fact.content_index,
|
||||
tags=extracted_fact.tags,
|
||||
observation_scopes=extracted_fact.observation_scopes,
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -1,18 +1,28 @@
|
||||
"""
|
||||
Link Expansion graph retrieval.
|
||||
|
||||
A simple, fast graph retrieval that expands from seeds via:
|
||||
1. Entity links: Find facts sharing entities with seeds (filtered by entity frequency)
|
||||
2. Causal links: Find facts causally linked to seeds (top-k by weight)
|
||||
Expands from semantic/temporal seeds through three parallel, first-class signals
|
||||
stored in memory_links:
|
||||
|
||||
Characteristics:
|
||||
- 2-3 DB queries (seed finding + parallel entity/causal expansion)
|
||||
- Sublinear: only touches connected facts via indexes
|
||||
- No iteration, no propagation, no normalization
|
||||
- Target: <100ms
|
||||
1. Entity links — precomputed co-occurrence graph (created at retain time, bounded to
|
||||
MAX_LINKS_PER_ENTITY per entity). Score = number of distinct shared
|
||||
entities between the seed set and each candidate.
|
||||
2. Semantic links — precomputed kNN graph (each new fact linked to its top-5 most
|
||||
similar existing facts at insert time, similarity >= 0.7). Checked
|
||||
in both directions since the graph is not symmetric. Score = weight.
|
||||
3. Causal links — explicit causal chains (causes/caused_by/enables/prevents).
|
||||
Score = weight + 1.0 (boosted as highest-quality signal).
|
||||
|
||||
All three signals are bounded at retain time, so no LATERAL fan-out caps are needed
|
||||
at query time. Each expansion is a simple aggregation over a small result set.
|
||||
|
||||
For non-observation fact types the three expansions are issued as a single CTE query
|
||||
(one roundtrip, one connection) with a `source` discriminator column so the Python
|
||||
merge step can apply per-signal score transformations.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import math
|
||||
import time
|
||||
|
||||
from ..db_utils import acquire_with_retry
|
||||
@@ -65,27 +75,23 @@ class LinkExpansionRetriever(GraphRetriever):
|
||||
"""
|
||||
Graph retrieval via direct link expansion from seeds.
|
||||
|
||||
Expands through entity co-occurrence and causal links in a single query.
|
||||
Fast and simple alternative to MPFP.
|
||||
Runs three expansions through precomputed memory_links: entity co-occurrence,
|
||||
semantic kNN, and causal chains, all bounded at retain time.
|
||||
|
||||
For non-observation fact types the three expansions are issued as a single CTE
|
||||
query (one roundtrip, one connection slot) with a `source` discriminator column.
|
||||
The Python merge step applies per-signal score transformations.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
max_entity_frequency: int = 500,
|
||||
causal_weight_threshold: float = 0.3,
|
||||
causal_limit_per_seed: int = 10,
|
||||
):
|
||||
"""
|
||||
Initialize link expansion retriever.
|
||||
|
||||
Args:
|
||||
max_entity_frequency: Skip entities appearing in more than this many facts
|
||||
causal_weight_threshold: Minimum weight for causal links
|
||||
causal_limit_per_seed: Max causal links to follow per seed
|
||||
causal_weight_threshold: Minimum weight for causal links to follow.
|
||||
"""
|
||||
self.max_entity_frequency = max_entity_frequency
|
||||
self.causal_weight_threshold = causal_weight_threshold
|
||||
self.causal_limit_per_seed = causal_limit_per_seed
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
@@ -110,7 +116,7 @@ class LinkExpansionRetriever(GraphRetriever):
|
||||
|
||||
Args:
|
||||
pool: Database connection pool
|
||||
query_embedding_str: Query embedding (unused, kept for interface)
|
||||
query_embedding_str: Query embedding as string
|
||||
bank_id: Memory bank ID
|
||||
fact_type: Fact type to filter
|
||||
budget: Maximum results to return
|
||||
@@ -118,7 +124,7 @@ class LinkExpansionRetriever(GraphRetriever):
|
||||
semantic_seeds: Pre-computed semantic entry points
|
||||
temporal_seeds: Pre-computed temporal entry points
|
||||
adjacency: Unused, kept for interface compatibility
|
||||
tags: Optional list of tags for visibility filtering (OR matching)
|
||||
tags: Optional list of tags for visibility filtering
|
||||
|
||||
Returns:
|
||||
Tuple of (results, timings)
|
||||
@@ -126,8 +132,6 @@ class LinkExpansionRetriever(GraphRetriever):
|
||||
start_time = time.time()
|
||||
timings = MPFPTimings(fact_type=fact_type)
|
||||
|
||||
# Use single connection for all queries to reduce pool pressure
|
||||
# (queries are fast ~50ms each, connection acquisition is the bottleneck)
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
# Find seeds if not provided
|
||||
if semantic_seeds:
|
||||
@@ -150,7 +154,6 @@ class LinkExpansionRetriever(GraphRetriever):
|
||||
f"(tags={tags}, tags_match={tags_match})"
|
||||
)
|
||||
|
||||
# Add temporal seeds if provided
|
||||
if temporal_seeds:
|
||||
all_seeds.extend(temporal_seeds)
|
||||
|
||||
@@ -160,223 +163,61 @@ class LinkExpansionRetriever(GraphRetriever):
|
||||
seed_ids = list({s.id for s in all_seeds})
|
||||
timings.pattern_count = len(seed_ids)
|
||||
|
||||
# Run entity and causal expansion sequentially on same connection
|
||||
query_start = time.time()
|
||||
|
||||
# For observations, traverse through source_memory_ids to find entity connections.
|
||||
# Observations don't have direct unit_entities - they inherit entities via their
|
||||
# source world/experience facts.
|
||||
#
|
||||
# Path: observation → source_memory_ids → world fact → entities →
|
||||
# ALL world facts with those entities → their observations (excluding seeds)
|
||||
if fact_type == "observation":
|
||||
# Debug: Check what source_memory_ids exist on seed observations
|
||||
debug_sources = await conn.fetch(
|
||||
f"""
|
||||
SELECT id, source_memory_ids
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE id = ANY($1::uuid[])
|
||||
""",
|
||||
seed_ids,
|
||||
)
|
||||
source_ids_found = []
|
||||
for row in debug_sources:
|
||||
if row["source_memory_ids"]:
|
||||
source_ids_found.extend(row["source_memory_ids"])
|
||||
logger.debug(
|
||||
f"[LinkExpansion] observation graph: {len(seed_ids)} seeds, "
|
||||
f"{len(source_ids_found)} source_memory_ids found"
|
||||
)
|
||||
|
||||
entity_rows = await conn.fetch(
|
||||
f"""
|
||||
WITH seed_sources AS (
|
||||
-- Get source memory IDs from seed observations
|
||||
SELECT DISTINCT unnest(source_memory_ids) AS source_id
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE id = ANY($1::uuid[])
|
||||
AND source_memory_ids IS NOT NULL
|
||||
),
|
||||
source_entities AS (
|
||||
-- Get entities from those source memories (filtered by frequency)
|
||||
SELECT DISTINCT ue.entity_id
|
||||
FROM seed_sources ss
|
||||
JOIN {fq_table("unit_entities")} ue ON ss.source_id = ue.unit_id
|
||||
JOIN {fq_table("entities")} e ON ue.entity_id = e.id
|
||||
WHERE e.mention_count < $2
|
||||
),
|
||||
all_connected_sources AS (
|
||||
-- Find ALL world facts sharing those entities (don't exclude seed sources)
|
||||
-- The exclusion happens at the observation level, not the source level
|
||||
SELECT DISTINCT other_ue.unit_id AS source_id
|
||||
FROM source_entities se
|
||||
JOIN {fq_table("unit_entities")} other_ue ON se.entity_id = other_ue.entity_id
|
||||
)
|
||||
-- Find observations derived from connected source memories
|
||||
-- Only exclude the actual seed observations
|
||||
SELECT
|
||||
mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
|
||||
mu.occurred_end, mu.mentioned_at,
|
||||
mu.fact_type, mu.document_id, mu.chunk_id, mu.tags,
|
||||
COUNT(DISTINCT cs.source_id)::float AS score
|
||||
FROM all_connected_sources cs
|
||||
JOIN {fq_table("memory_units")} mu
|
||||
ON mu.source_memory_ids @> ARRAY[cs.source_id]
|
||||
WHERE mu.fact_type = 'observation'
|
||||
AND mu.id != ALL($1::uuid[])
|
||||
GROUP BY mu.id
|
||||
ORDER BY score DESC
|
||||
LIMIT $3
|
||||
""",
|
||||
seed_ids,
|
||||
self.max_entity_frequency,
|
||||
budget,
|
||||
)
|
||||
logger.debug(f"[LinkExpansion] observation graph: found {len(entity_rows)} connected observations")
|
||||
entity_rows, semantic_rows, causal_rows = await self._expand_observations(conn, seed_ids, budget)
|
||||
else:
|
||||
# For world/experience facts, use direct entity lookup
|
||||
entity_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT
|
||||
mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
|
||||
mu.occurred_end, mu.mentioned_at,
|
||||
mu.fact_type, mu.document_id, mu.chunk_id, mu.tags,
|
||||
COUNT(*)::float AS score
|
||||
FROM {fq_table("unit_entities")} seed_ue
|
||||
JOIN {fq_table("entities")} e ON seed_ue.entity_id = e.id
|
||||
JOIN {fq_table("unit_entities")} other_ue ON seed_ue.entity_id = other_ue.entity_id
|
||||
JOIN {fq_table("memory_units")} mu ON other_ue.unit_id = mu.id
|
||||
WHERE seed_ue.unit_id = ANY($1::uuid[])
|
||||
AND e.mention_count < $2
|
||||
AND mu.id != ALL($1::uuid[])
|
||||
AND mu.fact_type = $3
|
||||
GROUP BY mu.id
|
||||
ORDER BY score DESC
|
||||
LIMIT $4
|
||||
""",
|
||||
seed_ids,
|
||||
self.max_entity_frequency,
|
||||
fact_type,
|
||||
budget,
|
||||
)
|
||||
|
||||
causal_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT DISTINCT ON (mu.id)
|
||||
mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
|
||||
mu.occurred_end, mu.mentioned_at,
|
||||
mu.fact_type, mu.document_id, mu.chunk_id, mu.tags,
|
||||
ml.weight + 1.0 AS score
|
||||
FROM {fq_table("memory_links")} ml
|
||||
JOIN {fq_table("memory_units")} mu ON ml.to_unit_id = mu.id
|
||||
WHERE ml.from_unit_id = ANY($1::uuid[])
|
||||
AND ml.link_type IN ('causes', 'caused_by', 'enables', 'prevents')
|
||||
AND ml.weight >= $2
|
||||
AND mu.fact_type = $3
|
||||
ORDER BY mu.id, ml.weight DESC
|
||||
LIMIT $4
|
||||
""",
|
||||
seed_ids,
|
||||
self.causal_weight_threshold,
|
||||
fact_type,
|
||||
budget,
|
||||
)
|
||||
|
||||
# Fallback: semantic/temporal/entity links from memory_links table
|
||||
# These are secondary to entity links (via unit_entities) and causal links
|
||||
# Weight is halved (0.5x) to prioritize primary link types
|
||||
# Check both directions: seeds -> others AND others -> seeds
|
||||
fallback_rows = await conn.fetch(
|
||||
f"""
|
||||
WITH outgoing AS (
|
||||
-- Links FROM seeds TO other facts
|
||||
SELECT mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
|
||||
mu.occurred_end, mu.mentioned_at,
|
||||
mu.fact_type, mu.document_id, mu.chunk_id, mu.tags,
|
||||
ml.weight
|
||||
FROM {fq_table("memory_links")} ml
|
||||
JOIN {fq_table("memory_units")} mu ON ml.to_unit_id = mu.id
|
||||
WHERE ml.from_unit_id = ANY($1::uuid[])
|
||||
AND ml.link_type IN ('semantic', 'temporal', 'entity')
|
||||
AND ml.weight >= $2
|
||||
AND mu.fact_type = $3
|
||||
AND mu.id != ALL($1::uuid[])
|
||||
),
|
||||
incoming AS (
|
||||
-- Links FROM other facts TO seeds (reverse direction)
|
||||
SELECT mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
|
||||
mu.occurred_end, mu.mentioned_at,
|
||||
mu.fact_type, mu.document_id, mu.chunk_id, mu.tags,
|
||||
ml.weight
|
||||
FROM {fq_table("memory_links")} ml
|
||||
JOIN {fq_table("memory_units")} mu ON ml.from_unit_id = mu.id
|
||||
WHERE ml.to_unit_id = ANY($1::uuid[])
|
||||
AND ml.link_type IN ('semantic', 'temporal', 'entity')
|
||||
AND ml.weight >= $2
|
||||
AND mu.fact_type = $3
|
||||
AND mu.id != ALL($1::uuid[])
|
||||
),
|
||||
combined AS (
|
||||
SELECT * FROM outgoing
|
||||
UNION ALL
|
||||
SELECT * FROM incoming
|
||||
)
|
||||
SELECT DISTINCT ON (id)
|
||||
id, text, context, event_date, occurred_start,
|
||||
occurred_end, mentioned_at,
|
||||
fact_type, document_id, chunk_id, tags,
|
||||
(MAX(weight) * 0.5) AS score
|
||||
FROM combined
|
||||
GROUP BY id, text, context, event_date, occurred_start,
|
||||
occurred_end, mentioned_at,
|
||||
fact_type, document_id, chunk_id, tags
|
||||
ORDER BY id, score DESC
|
||||
LIMIT $4
|
||||
""",
|
||||
seed_ids,
|
||||
self.causal_weight_threshold,
|
||||
fact_type,
|
||||
budget,
|
||||
)
|
||||
entity_rows, semantic_rows, causal_rows = await self._expand_combined(conn, seed_ids, fact_type, budget)
|
||||
|
||||
timings.edge_load_time = time.time() - query_start
|
||||
timings.db_queries = 3
|
||||
timings.edge_count = len(entity_rows) + len(causal_rows) + len(fallback_rows)
|
||||
timings.db_queries = 1
|
||||
timings.edge_count = len(entity_rows) + len(semantic_rows) + len(causal_rows)
|
||||
|
||||
# Merge results, taking max score per fact
|
||||
# Priority: entity links (unit_entities) > causal links > fallback links
|
||||
score_map: dict[str, float] = {}
|
||||
# Merge results with additive intra-score: entity + semantic + causal ∈ [0, 3].
|
||||
#
|
||||
# Entity score: tanh(count × 0.5) maps shared-entity count to [0, 1]:
|
||||
# 1 entity → 0.46, 2 → 0.76, 3 → 0.91, 4 → 0.96 (saturates naturally)
|
||||
# Semantic score: similarity weight, already ∈ [0.7, 1.0].
|
||||
# Causal score: link weight, already ∈ [0, 1].
|
||||
#
|
||||
# Facts appearing in multiple signals accumulate higher scores, rewarding
|
||||
# convergent evidence. The outer RRF uses rank position from this sorted list.
|
||||
entity_scores: dict[str, float] = {}
|
||||
semantic_scores: dict[str, float] = {}
|
||||
causal_scores: dict[str, float] = {}
|
||||
row_map: dict[str, dict] = {}
|
||||
|
||||
for row in entity_rows:
|
||||
fact_id = str(row["id"])
|
||||
score_map[fact_id] = max(score_map.get(fact_id, 0), row["score"])
|
||||
entity_scores[fact_id] = math.tanh(row["score"] * 0.5)
|
||||
row_map[fact_id] = dict(row)
|
||||
|
||||
for row in semantic_rows:
|
||||
fact_id = str(row["id"])
|
||||
semantic_scores[fact_id] = max(semantic_scores.get(fact_id, 0.0), row["score"])
|
||||
row_map.setdefault(fact_id, dict(row))
|
||||
|
||||
for row in causal_rows:
|
||||
fact_id = str(row["id"])
|
||||
score_map[fact_id] = max(score_map.get(fact_id, 0), row["score"])
|
||||
if fact_id not in row_map:
|
||||
row_map[fact_id] = dict(row)
|
||||
causal_scores[fact_id] = max(causal_scores.get(fact_id, 0.0), row["score"])
|
||||
row_map.setdefault(fact_id, dict(row))
|
||||
|
||||
for row in fallback_rows:
|
||||
fact_id = str(row["id"])
|
||||
score_map[fact_id] = max(score_map.get(fact_id, 0), row["score"])
|
||||
if fact_id not in row_map:
|
||||
row_map[fact_id] = dict(row)
|
||||
all_ids = set(entity_scores) | set(semantic_scores) | set(causal_scores)
|
||||
score_map = {
|
||||
fid: entity_scores.get(fid, 0.0) + semantic_scores.get(fid, 0.0) + causal_scores.get(fid, 0.0)
|
||||
for fid in all_ids
|
||||
}
|
||||
|
||||
# Sort by score and limit
|
||||
sorted_ids = sorted(score_map.keys(), key=lambda x: score_map[x], reverse=True)[:budget]
|
||||
rows = [row_map[fact_id] for fact_id in sorted_ids]
|
||||
|
||||
# Convert to results
|
||||
results = []
|
||||
for row in rows:
|
||||
result = RetrievalResult.from_db_row(dict(row))
|
||||
result.activation = row["score"]
|
||||
results.append(result)
|
||||
|
||||
# Apply tags filtering (graph expansion may reach untagged memories)
|
||||
if tags:
|
||||
results = filter_results_by_tags(results, tags, match=tags_match)
|
||||
|
||||
@@ -389,3 +230,254 @@ class LinkExpansionRetriever(GraphRetriever):
|
||||
)
|
||||
|
||||
return results, timings
|
||||
|
||||
async def _expand_combined(
|
||||
self,
|
||||
conn,
|
||||
seed_ids: list,
|
||||
fact_type: str,
|
||||
budget: int,
|
||||
) -> tuple[list, list, list]:
|
||||
"""
|
||||
Single-roundtrip CTE query combining entity, semantic, and causal expansions.
|
||||
|
||||
Uses a `source` discriminator column so the caller can apply per-signal
|
||||
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
|
||||
"""
|
||||
ml = fq_table("memory_links")
|
||||
mu = fq_table("memory_units")
|
||||
all_rows = await conn.fetch(
|
||||
f"""
|
||||
WITH entity_expanded AS (
|
||||
-- Entity co-occurrence: seeds → their precomputed entity-link neighbors.
|
||||
-- Score = distinct shared entities (bounded at retain time to
|
||||
-- MAX_LINKS_PER_ENTITY=50). GROUP BY mu.id is sufficient because mu.id
|
||||
-- is the primary key and functionally determines all other mu columns.
|
||||
SELECT
|
||||
mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
|
||||
mu.occurred_end, mu.mentioned_at,
|
||||
mu.fact_type, mu.document_id, mu.chunk_id, mu.tags,
|
||||
COUNT(DISTINCT ml.entity_id)::float AS score,
|
||||
'entity'::text AS source
|
||||
FROM {ml} ml
|
||||
JOIN {mu} mu ON mu.id = ml.to_unit_id
|
||||
WHERE ml.from_unit_id = ANY($1::uuid[])
|
||||
AND ml.link_type = 'entity'
|
||||
AND mu.fact_type = $2
|
||||
AND mu.id != ALL($1::uuid[])
|
||||
GROUP BY mu.id
|
||||
ORDER BY score DESC
|
||||
LIMIT $3
|
||||
),
|
||||
semantic_expanded AS (
|
||||
-- Semantic kNN: both outgoing (seeds → their kNN at insert time) and
|
||||
-- incoming (facts inserted after seeds that found seeds as kNN).
|
||||
-- Score = max similarity weight across both directions.
|
||||
SELECT
|
||||
id, text, context, event_date, occurred_start,
|
||||
occurred_end, mentioned_at,
|
||||
fact_type, document_id, chunk_id, tags,
|
||||
MAX(weight) AS score,
|
||||
'semantic'::text AS source
|
||||
FROM (
|
||||
SELECT
|
||||
mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
|
||||
mu.occurred_end, mu.mentioned_at,
|
||||
mu.fact_type, mu.document_id, mu.chunk_id, mu.tags,
|
||||
ml.weight
|
||||
FROM {ml} ml
|
||||
JOIN {mu} mu ON mu.id = ml.to_unit_id
|
||||
WHERE ml.from_unit_id = ANY($1::uuid[])
|
||||
AND ml.link_type = 'semantic'
|
||||
AND mu.fact_type = $2
|
||||
AND mu.id != ALL($1::uuid[])
|
||||
UNION ALL
|
||||
SELECT
|
||||
mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
|
||||
mu.occurred_end, mu.mentioned_at,
|
||||
mu.fact_type, mu.document_id, mu.chunk_id, mu.tags,
|
||||
ml.weight
|
||||
FROM {ml} ml
|
||||
JOIN {mu} mu ON mu.id = ml.from_unit_id
|
||||
WHERE ml.to_unit_id = ANY($1::uuid[])
|
||||
AND ml.link_type = 'semantic'
|
||||
AND mu.fact_type = $2
|
||||
AND mu.id != ALL($1::uuid[])
|
||||
) sem_raw
|
||||
GROUP BY id, text, context, event_date, occurred_start,
|
||||
occurred_end, mentioned_at,
|
||||
fact_type, document_id, chunk_id, tags
|
||||
ORDER BY score DESC
|
||||
LIMIT $3
|
||||
),
|
||||
causal_expanded AS (
|
||||
-- Causal chains: explicit causes/enables/prevents links from seeds.
|
||||
-- DISTINCT ON handles the case where a seed has multiple causal links
|
||||
-- to the same target; best weight wins.
|
||||
SELECT DISTINCT ON (mu.id)
|
||||
mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
|
||||
mu.occurred_end, mu.mentioned_at,
|
||||
mu.fact_type, mu.document_id, mu.chunk_id, mu.tags,
|
||||
ml.weight AS score,
|
||||
'causal'::text AS source
|
||||
FROM {ml} ml
|
||||
JOIN {mu} mu ON ml.to_unit_id = mu.id
|
||||
WHERE ml.from_unit_id = ANY($1::uuid[])
|
||||
AND ml.link_type IN ('causes', 'caused_by', 'enables', 'prevents')
|
||||
AND ml.weight >= $4
|
||||
AND mu.fact_type = $2
|
||||
ORDER BY mu.id, ml.weight DESC
|
||||
LIMIT $3
|
||||
)
|
||||
SELECT * FROM entity_expanded
|
||||
UNION ALL
|
||||
SELECT * FROM semantic_expanded
|
||||
UNION ALL
|
||||
SELECT * FROM causal_expanded
|
||||
""",
|
||||
seed_ids,
|
||||
fact_type,
|
||||
budget,
|
||||
self.causal_weight_threshold,
|
||||
)
|
||||
|
||||
entity_rows = [r for r in all_rows if r["source"] == "entity"]
|
||||
semantic_rows = [r for r in all_rows if r["source"] == "semantic"]
|
||||
causal_rows = [r for r in all_rows if r["source"] == "causal"]
|
||||
return entity_rows, semantic_rows, causal_rows
|
||||
|
||||
async def _expand_observations(
|
||||
self,
|
||||
conn,
|
||||
seed_ids: list,
|
||||
budget: int,
|
||||
) -> tuple[list, list, list]:
|
||||
"""
|
||||
Observation-specific expansion.
|
||||
|
||||
Observations don't have direct entity links in memory_links (they're created
|
||||
by consolidation, not retain). Instead, traverse source_memory_ids → world
|
||||
facts → entities → other world facts → their observations.
|
||||
|
||||
Semantic and causal expansions run as a second combined CTE query.
|
||||
"""
|
||||
source_ids_found: list = []
|
||||
if logger.isEnabledFor(logging.DEBUG):
|
||||
debug_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT id, source_memory_ids
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE id = ANY($1::uuid[])
|
||||
""",
|
||||
seed_ids,
|
||||
)
|
||||
for row in debug_rows:
|
||||
if row["source_memory_ids"]:
|
||||
source_ids_found.extend(row["source_memory_ids"])
|
||||
logger.debug(
|
||||
f"[LinkExpansion] observation graph: {len(seed_ids)} seeds, "
|
||||
f"{len(source_ids_found)} source_memory_ids found"
|
||||
)
|
||||
|
||||
entity_rows = await conn.fetch(
|
||||
f"""
|
||||
WITH seed_sources AS (
|
||||
SELECT DISTINCT unnest(source_memory_ids) AS source_id
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE id = ANY($1::uuid[])
|
||||
AND source_memory_ids IS NOT NULL
|
||||
),
|
||||
connected_sources AS (
|
||||
-- Mirror the non-observation entity expansion: follow pre-bounded entity
|
||||
-- links in memory_links (capped to MAX_LINKS_PER_ENTITY=50 at retain time).
|
||||
-- Score = number of distinct shared entities, same as the non-obs path.
|
||||
SELECT DISTINCT ml.to_unit_id AS source_id
|
||||
FROM seed_sources ss
|
||||
JOIN {fq_table("memory_links")} ml ON ml.from_unit_id = ss.source_id
|
||||
WHERE ml.link_type = 'entity'
|
||||
),
|
||||
connected_array AS (
|
||||
SELECT array_agg(source_id) AS source_ids FROM connected_sources
|
||||
)
|
||||
SELECT
|
||||
mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
|
||||
mu.occurred_end, mu.mentioned_at,
|
||||
mu.fact_type, mu.document_id, mu.chunk_id, mu.tags,
|
||||
(SELECT COUNT(DISTINCT s) FROM unnest(mu.source_memory_ids) s WHERE s = ANY(ca.source_ids))::float AS score
|
||||
FROM {fq_table("memory_units")} mu, connected_array ca
|
||||
WHERE mu.fact_type = 'observation'
|
||||
AND mu.id != ALL($1::uuid[])
|
||||
AND ca.source_ids IS NOT NULL
|
||||
AND mu.source_memory_ids && ca.source_ids
|
||||
ORDER BY score DESC
|
||||
LIMIT $2
|
||||
""",
|
||||
seed_ids,
|
||||
budget,
|
||||
)
|
||||
logger.debug(f"[LinkExpansion] observation graph: found {len(entity_rows)} connected observations")
|
||||
|
||||
# Semantic + causal for observations in one query
|
||||
ml = fq_table("memory_links")
|
||||
mu = fq_table("memory_units")
|
||||
sem_causal_rows = await conn.fetch(
|
||||
f"""
|
||||
WITH semantic_expanded AS (
|
||||
SELECT
|
||||
id, text, context, event_date, occurred_start,
|
||||
occurred_end, mentioned_at,
|
||||
fact_type, document_id, chunk_id, tags,
|
||||
MAX(weight) AS score,
|
||||
'semantic'::text AS source
|
||||
FROM (
|
||||
SELECT mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
|
||||
mu.occurred_end, mu.mentioned_at, mu.fact_type, mu.document_id,
|
||||
mu.chunk_id, mu.tags, ml.weight
|
||||
FROM {ml} ml JOIN {mu} mu ON mu.id = ml.to_unit_id
|
||||
WHERE ml.from_unit_id = ANY($1::uuid[])
|
||||
AND ml.link_type = 'semantic' AND mu.fact_type = 'observation'
|
||||
AND mu.id != ALL($1::uuid[])
|
||||
UNION ALL
|
||||
SELECT mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
|
||||
mu.occurred_end, mu.mentioned_at, mu.fact_type, mu.document_id,
|
||||
mu.chunk_id, mu.tags, ml.weight
|
||||
FROM {ml} ml JOIN {mu} mu ON mu.id = ml.from_unit_id
|
||||
WHERE ml.to_unit_id = ANY($1::uuid[])
|
||||
AND ml.link_type = 'semantic' AND mu.fact_type = 'observation'
|
||||
AND mu.id != ALL($1::uuid[])
|
||||
) sem_raw
|
||||
GROUP BY id, text, context, event_date, occurred_start, occurred_end,
|
||||
mentioned_at, fact_type, document_id, chunk_id, tags
|
||||
ORDER BY score DESC LIMIT $2
|
||||
),
|
||||
causal_expanded AS (
|
||||
SELECT DISTINCT ON (mu.id)
|
||||
mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start,
|
||||
mu.occurred_end, mu.mentioned_at, mu.fact_type, mu.document_id,
|
||||
mu.chunk_id, mu.tags, ml.weight AS score, 'causal'::text AS source
|
||||
FROM {ml} ml JOIN {mu} mu ON ml.to_unit_id = mu.id
|
||||
WHERE ml.from_unit_id = ANY($1::uuid[])
|
||||
AND ml.link_type IN ('causes', 'caused_by', 'enables', 'prevents')
|
||||
AND ml.weight >= $3 AND mu.fact_type = 'observation'
|
||||
ORDER BY mu.id, ml.weight DESC LIMIT $2
|
||||
)
|
||||
SELECT * FROM semantic_expanded
|
||||
UNION ALL
|
||||
SELECT * FROM causal_expanded
|
||||
""",
|
||||
seed_ids,
|
||||
budget,
|
||||
self.causal_weight_threshold,
|
||||
)
|
||||
|
||||
semantic_rows = [r for r in sem_causal_rows if r["source"] == "semantic"]
|
||||
causal_rows = [r for r in sem_causal_rows if r["source"] == "causal"]
|
||||
return entity_rows, semantic_rows, causal_rows
|
||||
|
||||
@@ -2,8 +2,72 @@
|
||||
Cross-encoder neural reranking for search results.
|
||||
"""
|
||||
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from .types import MergedCandidate, ScoredResult
|
||||
|
||||
UTC = timezone.utc
|
||||
|
||||
# Multiplicative boost alphas for recency and temporal proximity.
|
||||
# Each signal contributes at most ±(alpha/2) relative adjustment to the base CE score,
|
||||
# so the max combined boost is (1 + alpha/2)^2 ≈ +21% and min is (1 - alpha/2)^2 ≈ -19%.
|
||||
_RECENCY_ALPHA: float = 0.2
|
||||
_TEMPORAL_ALPHA: float = 0.2
|
||||
|
||||
|
||||
def apply_combined_scoring(
|
||||
scored_results: list[ScoredResult],
|
||||
now: datetime,
|
||||
recency_alpha: float = _RECENCY_ALPHA,
|
||||
temporal_alpha: float = _TEMPORAL_ALPHA,
|
||||
) -> None:
|
||||
"""Apply combined scoring to a list of ScoredResults in-place.
|
||||
|
||||
Uses the cross-encoder score as the primary relevance signal, with recency
|
||||
and temporal proximity applied as multiplicative boosts. This ensures the
|
||||
influence of these secondary signals is always proportional to the base
|
||||
relevance score, regardless of the cross-encoder model's score calibration.
|
||||
|
||||
Formula::
|
||||
|
||||
recency_boost = 1 + recency_alpha * (recency - 0.5) # in [1-α/2, 1+α/2]
|
||||
temporal_boost = 1 + temporal_alpha * (temporal - 0.5) # in [1-α/2, 1+α/2]
|
||||
combined_score = cross_encoder_score_normalized * recency_boost * temporal_boost
|
||||
|
||||
Temporal proximity is treated as neutral (0.5) when not set by temporal retrieval,
|
||||
so temporal_boost collapses to 1.0 for non-temporal queries.
|
||||
|
||||
Args:
|
||||
scored_results: Results from the cross-encoder reranker. Mutated in place.
|
||||
now: Current UTC datetime for recency calculation.
|
||||
recency_alpha: Max relative recency adjustment (default 0.2 → ±10%).
|
||||
temporal_alpha: Max relative temporal adjustment (default 0.2 → ±10%).
|
||||
"""
|
||||
if now.tzinfo is None:
|
||||
now = now.replace(tzinfo=UTC)
|
||||
|
||||
for sr in scored_results:
|
||||
# Recency: linear decay over 365 days → [0.1, 1.0]; neutral 0.5 if no date.
|
||||
sr.recency = 0.5
|
||||
if sr.retrieval.occurred_start:
|
||||
occurred = sr.retrieval.occurred_start
|
||||
if occurred.tzinfo is None:
|
||||
occurred = occurred.replace(tzinfo=UTC)
|
||||
days_ago = (now - occurred).total_seconds() / 86400
|
||||
sr.recency = max(0.1, min(1.0, 1.0 - (days_ago / 365)))
|
||||
|
||||
# Temporal proximity: meaningful only for temporal queries; neutral otherwise.
|
||||
sr.temporal = sr.retrieval.temporal_proximity if sr.retrieval.temporal_proximity is not None else 0.5
|
||||
|
||||
# RRF: kept at 0.0 for trace continuity but excluded from scoring.
|
||||
# RRF is batch-relative (min-max normalised) and redundant after reranking.
|
||||
sr.rrf_normalized = 0.0
|
||||
|
||||
recency_boost = 1.0 + recency_alpha * (sr.recency - 0.5)
|
||||
temporal_boost = 1.0 + temporal_alpha * (sr.temporal - 0.5)
|
||||
sr.combined_score = sr.cross_encoder_score_normalized * recency_boost * temporal_boost
|
||||
sr.weight = sr.combined_score
|
||||
|
||||
|
||||
class CrossEncoderReranker:
|
||||
"""
|
||||
|
||||
@@ -297,13 +297,20 @@ async def retrieve_temporal_combined(
|
||||
if tags:
|
||||
params.append(tags)
|
||||
|
||||
# Batch query: Get entry points for ALL fact types at once with window function
|
||||
# Two-phase entry point query:
|
||||
# Phase 1 (date_ranked): rank by date only — no embedding computation — for all units in
|
||||
# the temporal window. This lets the planner use date indexes for filtering.
|
||||
# Phase 2 (sim_ranked): join back to memory_units for only the top-50-per-type candidates
|
||||
# and compute embedding similarity for that small set (≤ 50 × len(fact_types) rows).
|
||||
# This avoids computing embedding distances for potentially thousands of date-range rows.
|
||||
entry_points = await conn.fetch(
|
||||
f"""
|
||||
WITH ranked_entries AS (
|
||||
SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, fact_type, document_id, chunk_id, tags,
|
||||
1 - (embedding <=> $1::vector) AS similarity,
|
||||
ROW_NUMBER() OVER (PARTITION BY fact_type ORDER BY COALESCE(occurred_start, mentioned_at, occurred_end) DESC, embedding <=> $1::vector) AS rn
|
||||
WITH date_ranked AS MATERIALIZED (
|
||||
SELECT id, fact_type,
|
||||
ROW_NUMBER() OVER (
|
||||
PARTITION BY fact_type
|
||||
ORDER BY COALESCE(occurred_start, mentioned_at, occurred_end) DESC NULLS LAST
|
||||
) AS rn
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE bank_id = $2
|
||||
AND fact_type = ANY($3)
|
||||
@@ -318,12 +325,20 @@ async def retrieve_temporal_combined(
|
||||
OR
|
||||
(occurred_end IS NOT NULL AND occurred_end BETWEEN $4 AND $5)
|
||||
)
|
||||
AND (1 - (embedding <=> $1::vector)) >= $6
|
||||
{tags_clause}
|
||||
),
|
||||
sim_ranked AS (
|
||||
SELECT mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start, mu.occurred_end, mu.mentioned_at, mu.fact_type, mu.document_id, mu.chunk_id, mu.tags,
|
||||
1 - (mu.embedding <=> $1::vector) AS similarity,
|
||||
ROW_NUMBER() OVER (PARTITION BY mu.fact_type ORDER BY mu.embedding <=> $1::vector) AS sim_rn
|
||||
FROM date_ranked dr
|
||||
JOIN {fq_table("memory_units")} mu ON mu.id = dr.id
|
||||
WHERE dr.rn <= 50
|
||||
AND (1 - (mu.embedding <=> $1::vector)) >= $6
|
||||
)
|
||||
SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, fact_type, document_id, chunk_id, tags, similarity
|
||||
FROM ranked_entries
|
||||
WHERE rn <= 10
|
||||
FROM sim_ranked
|
||||
WHERE sim_rn <= 10
|
||||
""",
|
||||
*params,
|
||||
)
|
||||
@@ -387,34 +402,52 @@ async def retrieve_temporal_combined(
|
||||
frontier = list(node_scores.keys())
|
||||
budget_remaining = budget - len(ft_entry_points)
|
||||
batch_size = 20
|
||||
# Per-source neighbor limit: lets the planner use the composite index
|
||||
# (from_unit_id, link_type, weight DESC) with early termination, avoiding
|
||||
# a full scan of all links from all source nodes before sorting.
|
||||
per_source_limit = 10
|
||||
# Safety cap on BFS iterations to prevent runaway spreading in dense graphs.
|
||||
max_iterations = 5
|
||||
iteration = 0
|
||||
|
||||
# Build tags clause for spreading (use param 6 since 1-5 are used)
|
||||
spreading_tags_clause = build_tags_where_clause_simple(tags, 6, table_alias="mu.", match=tags_match)
|
||||
# Build tags clause for spreading (use param 7 since 1-6 are used)
|
||||
spreading_tags_clause = build_tags_where_clause_simple(tags, 7, table_alias="mu.", match=tags_match)
|
||||
|
||||
while frontier and budget_remaining > 0:
|
||||
while frontier and budget_remaining > 0 and iteration < max_iterations:
|
||||
iteration += 1
|
||||
batch_ids = frontier[:batch_size]
|
||||
frontier = frontier[batch_size:]
|
||||
|
||||
spreading_params = [query_emb_str, batch_ids, ft, semantic_threshold, batch_size * 10]
|
||||
# $1=query_emb, $2=batch_ids, $3=fact_type, $4=threshold, $5=per_source_limit, $6=bank_id, $7=tags
|
||||
spreading_params = [query_emb_str, batch_ids, ft, semantic_threshold, per_source_limit, bank_id]
|
||||
if tags:
|
||||
spreading_params.append(tags)
|
||||
|
||||
# LATERAL join: for each source node, fetch top-K neighbors by weight using
|
||||
# the existing idx_memory_links_from_type_weight index with early-exit semantics.
|
||||
# This avoids scanning all temporal links from all source nodes before sorting.
|
||||
# bank_id on memory_units lets the planner use idx_memory_units_bank_fact_type.
|
||||
neighbors = await conn.fetch(
|
||||
f"""
|
||||
SELECT mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start, mu.occurred_end, mu.mentioned_at, mu.fact_type, mu.document_id, mu.chunk_id, mu.tags,
|
||||
ml.weight, ml.link_type, ml.from_unit_id,
|
||||
SELECT src.from_unit_id, mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start, mu.occurred_end, mu.mentioned_at, mu.fact_type, mu.document_id, mu.chunk_id, mu.tags,
|
||||
l.weight, l.link_type,
|
||||
1 - (mu.embedding <=> $1::vector) AS similarity
|
||||
FROM {fq_table("memory_links")} ml
|
||||
JOIN {fq_table("memory_units")} mu ON ml.to_unit_id = mu.id
|
||||
WHERE ml.from_unit_id = ANY($2::uuid[])
|
||||
AND ml.link_type IN ('temporal', 'causes', 'caused_by', 'enables', 'prevents')
|
||||
AND ml.weight >= 0.1
|
||||
FROM unnest($2::uuid[]) AS src(from_unit_id)
|
||||
CROSS JOIN LATERAL (
|
||||
SELECT ml.to_unit_id, ml.weight, ml.link_type
|
||||
FROM {fq_table("memory_links")} ml
|
||||
WHERE ml.from_unit_id = src.from_unit_id
|
||||
AND ml.link_type IN ('temporal', 'causes', 'caused_by', 'enables', 'prevents')
|
||||
AND ml.weight >= 0.1
|
||||
ORDER BY ml.weight DESC
|
||||
LIMIT $5
|
||||
) l
|
||||
JOIN {fq_table("memory_units")} mu ON mu.id = l.to_unit_id
|
||||
WHERE mu.bank_id = $6
|
||||
AND mu.fact_type = $3
|
||||
AND mu.embedding IS NOT NULL
|
||||
AND (1 - (mu.embedding <=> $1::vector)) >= $4
|
||||
{spreading_tags_clause}
|
||||
ORDER BY ml.weight DESC
|
||||
LIMIT $5
|
||||
""",
|
||||
*spreading_params,
|
||||
)
|
||||
|
||||
@@ -22,6 +22,11 @@ from hindsight_api.extensions.http import HttpExtension
|
||||
from hindsight_api.extensions.loader import load_extension
|
||||
from hindsight_api.extensions.mcp import MCPExtension
|
||||
from hindsight_api.extensions.operation_validator import (
|
||||
# Bank Management operations
|
||||
BankListContext,
|
||||
BankListResult,
|
||||
BankReadContext,
|
||||
BankWriteContext,
|
||||
# Consolidation operation
|
||||
ConsolidateContext,
|
||||
ConsolidateResult,
|
||||
@@ -70,6 +75,11 @@ __all__ = [
|
||||
"RetainContext",
|
||||
"RetainResult",
|
||||
"ValidationResult",
|
||||
# Operation Validator - Bank Management
|
||||
"BankListContext",
|
||||
"BankListResult",
|
||||
"BankReadContext",
|
||||
"BankWriteContext",
|
||||
# Operation Validator - Consolidation
|
||||
"ConsolidateContext",
|
||||
"ConsolidateResult",
|
||||
|
||||
@@ -96,6 +96,8 @@ class DefaultExtensionContext(ExtensionContext):
|
||||
|
||||
async def run_migration(self, schema: str) -> None:
|
||||
"""Run migrations for a specific schema."""
|
||||
import asyncio
|
||||
|
||||
from hindsight_api.config import get_config
|
||||
from hindsight_api.migrations import (
|
||||
ensure_embedding_dimension,
|
||||
@@ -111,10 +113,14 @@ class DefaultExtensionContext(ExtensionContext):
|
||||
if engine_url:
|
||||
db_url = engine_url
|
||||
|
||||
run_migrations(db_url, schema=schema)
|
||||
|
||||
# Get config for vector extension setting
|
||||
# Run synchronous migration functions in a thread so the asyncio event loop
|
||||
# remains free. This is critical for single-machine deployments where the
|
||||
# worker runs in-process: if run_migrations() blocks the event loop, any
|
||||
# in-flight asyncpg transactions cannot flush their COMMIT, and
|
||||
# CREATE INDEX CONCURRENTLY inside the migration waits for those transactions
|
||||
# forever — a deadlock.
|
||||
config = get_config()
|
||||
await asyncio.to_thread(run_migrations, db_url, schema=schema)
|
||||
|
||||
# Ensure embedding column dimension matches the model's dimension
|
||||
# This is needed because migrations create columns with default dimension
|
||||
@@ -123,15 +129,23 @@ class DefaultExtensionContext(ExtensionContext):
|
||||
if embeddings is not None:
|
||||
dimension = getattr(embeddings, "dimension", None)
|
||||
if dimension is not None:
|
||||
ensure_embedding_dimension(
|
||||
db_url, dimension, schema=schema, vector_extension=config.vector_extension
|
||||
await asyncio.to_thread(
|
||||
ensure_embedding_dimension,
|
||||
db_url,
|
||||
dimension,
|
||||
schema=schema,
|
||||
vector_extension=config.vector_extension,
|
||||
)
|
||||
|
||||
# Ensure vector indexes match the configured extension
|
||||
ensure_vector_extension(db_url, vector_extension=config.vector_extension, schema=schema)
|
||||
await asyncio.to_thread(
|
||||
ensure_vector_extension, db_url, vector_extension=config.vector_extension, schema=schema
|
||||
)
|
||||
|
||||
# Ensure text search columns/indexes match the configured extension
|
||||
ensure_text_search_extension(db_url, text_search_extension=config.text_search_extension, schema=schema)
|
||||
await asyncio.to_thread(
|
||||
ensure_text_search_extension, db_url, text_search_extension=config.text_search_extension, schema=schema
|
||||
)
|
||||
|
||||
def get_memory_engine(self) -> "MemoryEngineInterface":
|
||||
"""Get the memory engine interface."""
|
||||
|
||||
@@ -87,3 +87,15 @@ class HttpExtension(Extension, ABC):
|
||||
```
|
||||
"""
|
||||
pass
|
||||
|
||||
def get_root_router(self, memory: "MemoryEngine") -> APIRouter | None:
|
||||
"""
|
||||
Return a FastAPI router with endpoints mounted at the app root.
|
||||
|
||||
Unlike get_router() which is mounted at /ext/, this router is mounted
|
||||
directly on the application root. Use for well-known endpoints or other
|
||||
paths that must be at specific locations.
|
||||
|
||||
Returns None by default (no root routes). Override to provide root-level routes.
|
||||
"""
|
||||
return None
|
||||
|
||||
@@ -200,6 +200,44 @@ class ConsolidateResult:
|
||||
error: str | None = None
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Bank Management Contexts
|
||||
# =============================================================================
|
||||
|
||||
|
||||
@dataclass
|
||||
class BankReadContext:
|
||||
"""Context for a bank read operation validation (pre-operation)."""
|
||||
|
||||
bank_id: str
|
||||
operation: str # "get_bank_profile", "get_bank_stats"
|
||||
request_context: "RequestContext"
|
||||
|
||||
|
||||
@dataclass
|
||||
class BankWriteContext:
|
||||
"""Context for a bank write operation validation (pre-operation)."""
|
||||
|
||||
bank_id: str
|
||||
operation: str # "delete_bank", "update_bank", "update_bank_disposition", "set_bank_mission", "merge_bank_mission", "clear_observations", "clear_observations_for_memory"
|
||||
request_context: "RequestContext"
|
||||
|
||||
|
||||
@dataclass
|
||||
class BankListContext:
|
||||
"""Context for filtering the bank list (post-query)."""
|
||||
|
||||
banks: list[dict]
|
||||
request_context: "RequestContext"
|
||||
|
||||
|
||||
@dataclass
|
||||
class BankListResult:
|
||||
"""Result of filtering the bank list."""
|
||||
|
||||
banks: list[dict]
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Mental Model Contexts
|
||||
# =============================================================================
|
||||
@@ -535,3 +573,63 @@ class OperationValidatorExtension(Extension, ABC):
|
||||
- error: Error message (if failed)
|
||||
"""
|
||||
pass
|
||||
|
||||
# =========================================================================
|
||||
# Bank Management - Validation hooks (optional - override to implement)
|
||||
# =========================================================================
|
||||
|
||||
async def validate_bank_read(self, ctx: BankReadContext) -> ValidationResult:
|
||||
"""
|
||||
Validate a bank read operation before execution.
|
||||
|
||||
Override to implement custom validation logic for bank reads
|
||||
(get_bank_profile, get_bank_stats).
|
||||
|
||||
Args:
|
||||
ctx: Context containing:
|
||||
- bank_id: Bank identifier
|
||||
- operation: Operation name
|
||||
- request_context: Request context with auth info
|
||||
|
||||
Returns:
|
||||
ValidationResult indicating whether the operation is allowed.
|
||||
"""
|
||||
return ValidationResult.accept()
|
||||
|
||||
async def validate_bank_write(self, ctx: BankWriteContext) -> ValidationResult:
|
||||
"""
|
||||
Validate a bank write operation before execution.
|
||||
|
||||
Override to implement custom validation logic for bank writes
|
||||
(delete_bank, update_bank, update_bank_disposition, set_bank_mission,
|
||||
merge_bank_mission, clear_observations, clear_observations_for_memory).
|
||||
|
||||
Args:
|
||||
ctx: Context containing:
|
||||
- bank_id: Bank identifier
|
||||
- operation: Operation name
|
||||
- request_context: Request context with auth info
|
||||
|
||||
Returns:
|
||||
ValidationResult indicating whether the operation is allowed.
|
||||
"""
|
||||
return ValidationResult.accept()
|
||||
|
||||
async def filter_bank_list(self, ctx: BankListContext) -> BankListResult:
|
||||
"""
|
||||
Filter the bank list after querying.
|
||||
|
||||
Unlike validate_* methods, this is a post-query filter that narrows results
|
||||
rather than a gate that blocks the operation.
|
||||
|
||||
Override to implement custom filtering (e.g., restrict to allowed banks).
|
||||
|
||||
Args:
|
||||
ctx: Context containing:
|
||||
- banks: List of bank dicts from the database
|
||||
- request_context: Request context with auth info
|
||||
|
||||
Returns:
|
||||
BankListResult with the filtered list of banks.
|
||||
"""
|
||||
return BankListResult(banks=ctx.banks)
|
||||
|
||||
@@ -11,8 +11,9 @@ from hindsight_api.models import RequestContext
|
||||
class AuthenticationError(Exception):
|
||||
"""Raised when authentication fails."""
|
||||
|
||||
def __init__(self, reason: str):
|
||||
def __init__(self, reason: str, headers: dict[str, str] | None = None):
|
||||
self.reason = reason
|
||||
self.headers = headers or {}
|
||||
super().__init__(f"Authentication failed: {reason}")
|
||||
|
||||
|
||||
|
||||
@@ -171,6 +171,7 @@ def main():
|
||||
llm_vertexai_project_id=config.llm_vertexai_project_id,
|
||||
llm_vertexai_region=config.llm_vertexai_region,
|
||||
llm_vertexai_service_account_key=config.llm_vertexai_service_account_key,
|
||||
llm_gemini_safety_settings=config.llm_gemini_safety_settings,
|
||||
retain_llm_provider=config.retain_llm_provider,
|
||||
retain_llm_api_key=config.retain_llm_api_key,
|
||||
retain_llm_model=config.retain_llm_model,
|
||||
@@ -252,6 +253,7 @@ def main():
|
||||
retain_mission=config.retain_mission,
|
||||
retain_custom_instructions=config.retain_custom_instructions,
|
||||
retain_batch_tokens=config.retain_batch_tokens,
|
||||
retain_entity_lookup=config.retain_entity_lookup,
|
||||
retain_batch_enabled=config.retain_batch_enabled,
|
||||
retain_batch_poll_interval_seconds=config.retain_batch_poll_interval_seconds,
|
||||
file_storage_type=config.file_storage_type,
|
||||
@@ -277,6 +279,8 @@ def main():
|
||||
consolidation_llm_batch_size=config.consolidation_llm_batch_size,
|
||||
consolidation_max_tokens=config.consolidation_max_tokens,
|
||||
observations_mission=config.observations_mission,
|
||||
entity_labels=config.entity_labels,
|
||||
entities_allow_free_form=config.entities_allow_free_form,
|
||||
skip_llm_verification=config.skip_llm_verification,
|
||||
lazy_reranker=config.lazy_reranker,
|
||||
run_migrations_on_startup=config.run_migrations_on_startup,
|
||||
@@ -292,6 +296,7 @@ def main():
|
||||
worker_max_slots=config.worker_max_slots,
|
||||
worker_consolidation_max_slots=config.worker_consolidation_max_slots,
|
||||
reflect_max_iterations=config.reflect_max_iterations,
|
||||
reflect_max_context_tokens=config.reflect_max_context_tokens,
|
||||
reflect_mission=config.reflect_mission,
|
||||
disposition_skepticism=config.disposition_skepticism,
|
||||
disposition_literalism=config.disposition_literalism,
|
||||
@@ -302,6 +307,10 @@ def main():
|
||||
otel_exporter_otlp_headers=config.otel_exporter_otlp_headers,
|
||||
otel_service_name=config.otel_service_name,
|
||||
otel_deployment_environment=config.otel_deployment_environment,
|
||||
webhook_url=config.webhook_url,
|
||||
webhook_secret=config.webhook_secret,
|
||||
webhook_event_types=config.webhook_event_types,
|
||||
webhook_delivery_poll_interval_seconds=config.webhook_delivery_poll_interval_seconds,
|
||||
)
|
||||
config.configure_logging()
|
||||
if not args.daemon:
|
||||
@@ -380,6 +389,7 @@ def main():
|
||||
"ws": "wsproto", # Use wsproto instead of websockets to avoid deprecation warnings
|
||||
"loop": loop_impl, # Explicitly set event loop implementation
|
||||
"timeout_keep_alive": 30, # Exceed aiohttp's 15s client timeout so the client always closes first
|
||||
"timeout_graceful_shutdown": 5, # Cap graceful shutdown at 5s; also enables force-kill on second Ctrl+C
|
||||
}
|
||||
|
||||
# Add optional parameters if provided
|
||||
|
||||
@@ -20,6 +20,7 @@ from hindsight_api.config import (
|
||||
)
|
||||
from hindsight_api.engine.memory_engine import Budget
|
||||
from hindsight_api.engine.response_models import VALID_RECALL_FACT_TYPES
|
||||
from hindsight_api.extensions import OperationValidationError
|
||||
from hindsight_api.models import RequestContext
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -52,7 +53,6 @@ class MCPToolsConfig:
|
||||
recall_description: str | None = None
|
||||
|
||||
# Retain behavior
|
||||
retain_fire_and_forget: bool = False # If True, use asyncio.create_task pattern
|
||||
|
||||
|
||||
def _get_request_context(config: MCPToolsConfig) -> RequestContext:
|
||||
@@ -271,48 +271,71 @@ def register_mcp_tools(
|
||||
def _apply_bank_tool_filtering(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig) -> None:
|
||||
"""Filter bank-level mcp_enabled_tools from both tools/list and tool invocation.
|
||||
|
||||
Wraps _tool_manager.get_tools() so that:
|
||||
- tools/list only returns permitted tools (they are hidden, not just blocked)
|
||||
- tools/call for a disabled tool raises NotFoundError (via the manager) before run()
|
||||
|
||||
tool.run wrappers are kept as defense-in-depth for any caller that bypasses the manager.
|
||||
Compatible with FastMCP 2.x (_tool_manager pattern) and 3.x (provider pattern).
|
||||
"""
|
||||
try:
|
||||
tool_manager = mcp._tool_manager
|
||||
original_get_tools = tool_manager.get_tools
|
||||
|
||||
async def _filtered_get_tools():
|
||||
all_tools = await original_get_tools()
|
||||
bank_id = config.bank_id_resolver()
|
||||
if not bank_id:
|
||||
return all_tools
|
||||
request_context = _get_request_context(config)
|
||||
bank_cfg = await memory._config_resolver.get_bank_config(bank_id, request_context)
|
||||
enabled: list[str] | None = bank_cfg.get("mcp_enabled_tools")
|
||||
if enabled is None:
|
||||
return all_tools
|
||||
enabled_set = set(enabled)
|
||||
return {k: v for k, v in all_tools.items() if k in enabled_set}
|
||||
async def _get_enabled_tools() -> set[str] | None:
|
||||
"""Return the enabled tool set for the current bank, or None if unrestricted."""
|
||||
bank_id = config.bank_id_resolver()
|
||||
if not bank_id:
|
||||
return None
|
||||
request_context = _get_request_context(config)
|
||||
bank_cfg = await memory._config_resolver.get_bank_config(bank_id, request_context)
|
||||
enabled: list[str] | None = bank_cfg.get("mcp_enabled_tools")
|
||||
if enabled is None:
|
||||
return None
|
||||
return set(enabled)
|
||||
|
||||
setattr(tool_manager, "get_tools", _filtered_get_tools)
|
||||
if hasattr(mcp, "list_tools"):
|
||||
# FastMCP 3.x: wrap list_tools() and get_tool() on the instance
|
||||
original_list_tools = mcp.list_tools
|
||||
original_get_tool = mcp.get_tool
|
||||
|
||||
# Defense-in-depth: also wrap tool.run for any direct caller that bypasses the manager
|
||||
for name, tool in tool_manager._tools.items():
|
||||
original_run = tool.run
|
||||
async def _filtered_list_tools(**kwargs):
|
||||
tools = await original_list_tools(**kwargs)
|
||||
enabled_set = await _get_enabled_tools()
|
||||
if enabled_set is None:
|
||||
return tools
|
||||
return [t for t in tools if t.name in enabled_set]
|
||||
|
||||
async def _filtered_run(arguments, _name=name, _orig=original_run):
|
||||
bank_id = config.bank_id_resolver()
|
||||
if bank_id:
|
||||
request_context = _get_request_context(config)
|
||||
bank_cfg = await memory._config_resolver.get_bank_config(bank_id, request_context)
|
||||
enabled: list[str] | None = bank_cfg.get("mcp_enabled_tools")
|
||||
if enabled is not None and _name not in enabled:
|
||||
raise ValueError(f"Tool '{_name}' is not enabled for bank '{bank_id}'")
|
||||
return await _orig(arguments)
|
||||
async def _filtered_get_tool(name, **kwargs):
|
||||
enabled_set = await _get_enabled_tools()
|
||||
if enabled_set is not None and name not in enabled_set:
|
||||
return None # FastMCP treats None as "not found" → raises NotFoundError
|
||||
return await original_get_tool(name, **kwargs)
|
||||
|
||||
object.__setattr__(tool, "run", _filtered_run)
|
||||
except (AttributeError, KeyError) as e:
|
||||
logger.warning(f"Could not apply bank tool filtering: {e}")
|
||||
object.__setattr__(mcp, "list_tools", _filtered_list_tools)
|
||||
object.__setattr__(mcp, "get_tool", _filtered_get_tool)
|
||||
|
||||
elif hasattr(mcp, "_tool_manager"):
|
||||
# FastMCP 2.x: wrap _tool_manager.get_tools() and tool.run()
|
||||
try:
|
||||
tool_manager = mcp._tool_manager
|
||||
original_get_tools = tool_manager.get_tools
|
||||
|
||||
async def _filtered_get_tools():
|
||||
all_tools = await original_get_tools()
|
||||
enabled_set = await _get_enabled_tools()
|
||||
if enabled_set is None:
|
||||
return all_tools
|
||||
return {k: v for k, v in all_tools.items() if k in enabled_set}
|
||||
|
||||
setattr(tool_manager, "get_tools", _filtered_get_tools)
|
||||
|
||||
for name, tool in tool_manager._tools.items():
|
||||
original_run = tool.run
|
||||
|
||||
async def _filtered_run(arguments, _name=name, _orig=original_run):
|
||||
enabled_set = await _get_enabled_tools()
|
||||
if enabled_set is not None and _name not in enabled_set:
|
||||
raise ValueError(f"Tool '{_name}' is not enabled for bank '{config.bank_id_resolver()}'")
|
||||
return await _orig(arguments)
|
||||
|
||||
object.__setattr__(tool, "run", _filtered_run)
|
||||
except (AttributeError, KeyError) as e:
|
||||
logger.warning(f"Could not apply bank tool filtering (v2): {e}")
|
||||
else:
|
||||
logger.warning("Could not apply bank tool filtering: unknown FastMCP version")
|
||||
|
||||
|
||||
def _register_retain(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig) -> None:
|
||||
@@ -320,106 +343,56 @@ def _register_retain(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig)
|
||||
description = config.retain_description or DEFAULT_MCP_RETAIN_DESCRIPTION
|
||||
|
||||
if config.include_bank_id_param:
|
||||
if config.retain_fire_and_forget:
|
||||
|
||||
@mcp.tool(description=description)
|
||||
async def retain(
|
||||
content: str,
|
||||
context: str = "general",
|
||||
timestamp: str | None = None,
|
||||
tags: list[str] | None = None,
|
||||
metadata: dict[str, str] | None = None,
|
||||
document_id: str | None = None,
|
||||
bank_id: str | None = None,
|
||||
) -> dict:
|
||||
"""
|
||||
Args:
|
||||
content: The fact/memory to store (be specific and include relevant details)
|
||||
context: Category for the memory (e.g., 'preferences', 'work', 'hobbies', 'family'). Default: 'general'
|
||||
timestamp: When this event/fact occurred (ISO format, e.g., '2024-01-15T10:30:00Z'). Useful for timeline tracking.
|
||||
tags: Optional tags for scoped visibility filtering (e.g., ['project:alpha', 'user:123'])
|
||||
metadata: Optional key-value metadata to attach (e.g., {'source': 'slack', 'channel': 'general'})
|
||||
document_id: Optional document ID to associate this memory with
|
||||
bank_id: Optional bank to store in (defaults to session bank). Use for cross-bank operations.
|
||||
"""
|
||||
import asyncio
|
||||
@mcp.tool(description=description)
|
||||
async def retain(
|
||||
content: str,
|
||||
context: str = "general",
|
||||
timestamp: str | None = None,
|
||||
tags: list[str] | None = None,
|
||||
metadata: dict[str, str] | None = None,
|
||||
document_id: str | None = None,
|
||||
bank_id: str | None = None,
|
||||
) -> dict:
|
||||
"""
|
||||
Args:
|
||||
content: The fact/memory to store (be specific and include relevant details)
|
||||
context: Category for the memory (e.g., 'preferences', 'work', 'hobbies', 'family'). Default: 'general'
|
||||
timestamp: When this event/fact occurred (ISO format, e.g., '2024-01-15T10:30:00Z'). Useful for timeline tracking.
|
||||
tags: Optional tags for scoped visibility filtering (e.g., ['project:alpha', 'user:123'])
|
||||
metadata: Optional key-value metadata to attach (e.g., {'source': 'slack', 'channel': 'general'})
|
||||
document_id: Optional document ID to associate this memory with
|
||||
bank_id: Optional bank to store in (defaults to session bank). Use for cross-bank operations.
|
||||
"""
|
||||
target_bank = bank_id or config.bank_id_resolver()
|
||||
if target_bank is None:
|
||||
return {"status": "error", "message": "No bank_id configured"}
|
||||
|
||||
target_bank = bank_id or config.bank_id_resolver()
|
||||
if target_bank is None:
|
||||
return {"status": "error", "message": "No bank_id configured"}
|
||||
content_dict, error = build_content_dict(content, context, timestamp, tags, metadata, document_id)
|
||||
if error:
|
||||
return {"status": "error", "message": error}
|
||||
|
||||
content_dict, error = build_content_dict(content, context, timestamp, tags, metadata, document_id)
|
||||
if error:
|
||||
return {"status": "error", "message": error}
|
||||
request_context = _get_request_context(config)
|
||||
|
||||
request_context = _get_request_context(config)
|
||||
|
||||
async def _retain():
|
||||
try:
|
||||
await memory.retain_batch_async(
|
||||
bank_id=target_bank,
|
||||
contents=[content_dict],
|
||||
request_context=request_context,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"Error storing memory: {e}", exc_info=True)
|
||||
|
||||
asyncio.create_task(_retain())
|
||||
return {"status": "accepted", "message": "Memory storage initiated"}
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool(description=description)
|
||||
async def retain(
|
||||
content: str,
|
||||
context: str = "general",
|
||||
timestamp: str | None = None,
|
||||
tags: list[str] | None = None,
|
||||
metadata: dict[str, str] | None = None,
|
||||
document_id: str | None = None,
|
||||
async_processing: bool = True,
|
||||
bank_id: str | None = None,
|
||||
) -> str:
|
||||
"""
|
||||
Args:
|
||||
content: The fact/memory to store (be specific and include relevant details)
|
||||
context: Category for the memory (e.g., 'preferences', 'work', 'hobbies', 'family'). Default: 'general'
|
||||
timestamp: When this event/fact occurred (ISO format, e.g., '2024-01-15T10:30:00Z'). Useful for timeline tracking.
|
||||
tags: Optional tags for scoped visibility filtering (e.g., ['project:alpha', 'user:123'])
|
||||
metadata: Optional key-value metadata to attach (e.g., {'source': 'slack', 'channel': 'general'})
|
||||
document_id: Optional document ID to associate this memory with
|
||||
async_processing: If True, queue for background processing and return immediately. If False, wait for completion. Default: True
|
||||
bank_id: Optional bank to store in (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"
|
||||
|
||||
content_dict, error = build_content_dict(content, context, timestamp, tags, metadata, document_id)
|
||||
if error:
|
||||
return f"Error: {error}"
|
||||
|
||||
contents = [content_dict]
|
||||
request_context = _get_request_context(config)
|
||||
if async_processing:
|
||||
result = await memory.submit_async_retain(
|
||||
bank_id=target_bank, contents=contents, request_context=request_context
|
||||
)
|
||||
return f"Memory queued for background processing (operation_id: {result.get('operation_id', 'N/A')})"
|
||||
else:
|
||||
await memory.retain_batch_async(
|
||||
bank_id=target_bank,
|
||||
contents=contents,
|
||||
request_context=request_context,
|
||||
)
|
||||
return f"Memory stored successfully in bank '{target_bank}'"
|
||||
except Exception as e:
|
||||
logger.error(f"Error storing memory: {e}", exc_info=True)
|
||||
return f"Error: {str(e)}"
|
||||
try:
|
||||
result = await memory.submit_async_retain(
|
||||
bank_id=target_bank,
|
||||
contents=[content_dict],
|
||||
request_context=request_context,
|
||||
)
|
||||
return {
|
||||
"status": "accepted",
|
||||
"message": "Memory storage initiated",
|
||||
"operation_id": result.get("operation_id"),
|
||||
}
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Retain rejected: {e}")
|
||||
return {"status": "error", "message": str(e)}
|
||||
except Exception as e:
|
||||
logger.error(f"Error storing memory: {e}", exc_info=True)
|
||||
return {"status": "error", "message": str(e)}
|
||||
|
||||
else:
|
||||
# No bank_id param - use fixed bank from resolver
|
||||
|
||||
@mcp.tool(description=description)
|
||||
async def retain(
|
||||
@@ -439,8 +412,6 @@ def _register_retain(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig)
|
||||
metadata: Optional key-value metadata to attach (e.g., {'source': 'slack', 'channel': 'general'})
|
||||
document_id: Optional document ID to associate this memory with
|
||||
"""
|
||||
import asyncio
|
||||
|
||||
target_bank = config.bank_id_resolver()
|
||||
if target_bank is None:
|
||||
return {"status": "error", "message": "No bank_id configured"}
|
||||
@@ -451,18 +422,23 @@ def _register_retain(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig)
|
||||
|
||||
request_context = _get_request_context(config)
|
||||
|
||||
async def _retain():
|
||||
try:
|
||||
await memory.retain_batch_async(
|
||||
bank_id=target_bank,
|
||||
contents=[content_dict],
|
||||
request_context=request_context,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"Error storing memory: {e}", exc_info=True)
|
||||
|
||||
asyncio.create_task(_retain())
|
||||
return {"status": "accepted", "message": "Memory storage initiated"}
|
||||
try:
|
||||
result = await memory.submit_async_retain(
|
||||
bank_id=target_bank,
|
||||
contents=[content_dict],
|
||||
request_context=request_context,
|
||||
)
|
||||
return {
|
||||
"status": "accepted",
|
||||
"message": "Memory storage initiated",
|
||||
"operation_id": result.get("operation_id"),
|
||||
}
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Retain rejected: {e}")
|
||||
return {"status": "error", "message": str(e)}
|
||||
except Exception as e:
|
||||
logger.error(f"Error storing memory: {e}", exc_info=True)
|
||||
return {"status": "error", "message": str(e)}
|
||||
|
||||
|
||||
def _register_recall(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig) -> None:
|
||||
@@ -519,6 +495,9 @@ def _register_recall(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig)
|
||||
recall_result = await memory.recall_async(**recall_kwargs)
|
||||
|
||||
return recall_result.model_dump_json(indent=2)
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Recall rejected: {e}")
|
||||
return json.dumps({"error": str(e), "results": []})
|
||||
except ValueError as e:
|
||||
return f'{{"error": "{e}", "results": []}}'
|
||||
except Exception as e:
|
||||
@@ -573,6 +552,9 @@ def _register_recall(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig)
|
||||
recall_result = await memory.recall_async(**recall_kwargs)
|
||||
|
||||
return recall_result.model_dump()
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Recall rejected: {e}")
|
||||
return {"error": str(e), "results": []}
|
||||
except ValueError as e:
|
||||
return {"error": str(e), "results": []}
|
||||
except Exception as e:
|
||||
@@ -653,6 +635,9 @@ def _register_reflect(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig
|
||||
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)
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Reflect rejected: {e}")
|
||||
return json.dumps({"error": str(e)})
|
||||
except Exception as e:
|
||||
logger.error(f"Error reflecting: {e}", exc_info=True)
|
||||
return f'{{"error": "{e}", "text": ""}}'
|
||||
@@ -725,6 +710,9 @@ def _register_reflect(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig
|
||||
if response_schema is not None and hasattr(reflect_result, "structured_output"):
|
||||
result_data["structured_output"] = reflect_result.structured_output
|
||||
return result_data
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Reflect rejected: {e}")
|
||||
return {"error": str(e)}
|
||||
except Exception as e:
|
||||
logger.error(f"Error reflecting: {e}", exc_info=True)
|
||||
return {"error": str(e), "text": ""}
|
||||
@@ -747,6 +735,9 @@ def _register_list_banks(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsCon
|
||||
try:
|
||||
banks = await memory.list_banks(request_context=_get_request_context(config))
|
||||
return json.dumps({"banks": banks}, indent=2)
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return json.dumps({"error": str(e), "banks": []})
|
||||
except Exception as e:
|
||||
logger.error(f"Error listing banks: {e}", exc_info=True)
|
||||
return f'{{"error": "{e}", "banks": []}}'
|
||||
@@ -788,6 +779,9 @@ def _register_create_bank(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsCo
|
||||
if "disposition" in profile and hasattr(profile["disposition"], "model_dump"):
|
||||
profile["disposition"] = profile["disposition"].model_dump()
|
||||
return json.dumps(profile, indent=2)
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return json.dumps({"error": str(e)})
|
||||
except Exception as e:
|
||||
logger.error(f"Error creating bank: {e}", exc_info=True)
|
||||
return f'{{"error": "{e}"}}'
|
||||
@@ -843,6 +837,9 @@ def _register_list_mental_models(mcp: FastMCP, memory: MemoryEngine, config: MCP
|
||||
request_context=_get_request_context(config),
|
||||
)
|
||||
return json.dumps({"items": models}, indent=2, default=str)
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return json.dumps({"error": str(e)})
|
||||
except Exception as e:
|
||||
logger.error(f"Error listing mental models: {e}", exc_info=True)
|
||||
return f'{{"error": "{e}", "items": []}}'
|
||||
@@ -874,6 +871,9 @@ def _register_list_mental_models(mcp: FastMCP, memory: MemoryEngine, config: MCP
|
||||
request_context=_get_request_context(config),
|
||||
)
|
||||
return {"items": models}
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return {"error": str(e)}
|
||||
except Exception as e:
|
||||
logger.error(f"Error listing mental models: {e}", exc_info=True)
|
||||
return {"error": str(e), "items": []}
|
||||
@@ -912,6 +912,9 @@ def _register_get_mental_model(mcp: FastMCP, memory: MemoryEngine, config: MCPTo
|
||||
if model is None:
|
||||
return json.dumps({"error": f"Mental model '{mental_model_id}' not found in bank '{target_bank}'"})
|
||||
return json.dumps(model, indent=2, default=str)
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return json.dumps({"error": str(e)})
|
||||
except Exception as e:
|
||||
logger.error(f"Error getting mental model: {e}", exc_info=True)
|
||||
return f'{{"error": "{e}"}}'
|
||||
@@ -944,6 +947,9 @@ def _register_get_mental_model(mcp: FastMCP, memory: MemoryEngine, config: MCPTo
|
||||
if model is None:
|
||||
return {"error": f"Mental model '{mental_model_id}' not found in bank '{target_bank}'"}
|
||||
return model
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return {"error": str(e)}
|
||||
except Exception as e:
|
||||
logger.error(f"Error getting mental model: {e}", exc_info=True)
|
||||
return {"error": str(e)}
|
||||
@@ -1027,6 +1033,9 @@ def _register_create_mental_model(mcp: FastMCP, memory: MemoryEngine, config: MC
|
||||
"message": f"Mental model '{name}' created. Content is being generated asynchronously.",
|
||||
}
|
||||
)
|
||||
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:
|
||||
@@ -1102,6 +1111,9 @@ def _register_create_mental_model(mcp: FastMCP, memory: MemoryEngine, config: MC
|
||||
"status": "created",
|
||||
"message": f"Mental model '{name}' created. Content is being generated asynchronously.",
|
||||
}
|
||||
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:
|
||||
@@ -1166,6 +1178,9 @@ def _register_update_mental_model(mcp: FastMCP, memory: MemoryEngine, config: MC
|
||||
if model is None:
|
||||
return json.dumps({"error": f"Mental model '{mental_model_id}' not found in bank '{target_bank}'"})
|
||||
return json.dumps(model, indent=2, default=str)
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return json.dumps({"error": str(e)})
|
||||
except Exception as e:
|
||||
logger.error(f"Error updating mental model: {e}", exc_info=True)
|
||||
return f'{{"error": "{e}"}}'
|
||||
@@ -1222,6 +1237,9 @@ def _register_update_mental_model(mcp: FastMCP, memory: MemoryEngine, config: MC
|
||||
if model is None:
|
||||
return {"error": f"Mental model '{mental_model_id}' not found in bank '{target_bank}'"}
|
||||
return model
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return {"error": str(e)}
|
||||
except Exception as e:
|
||||
logger.error(f"Error updating mental model: {e}", exc_info=True)
|
||||
return {"error": str(e)}
|
||||
@@ -1259,6 +1277,9 @@ def _register_delete_mental_model(mcp: FastMCP, memory: MemoryEngine, config: MC
|
||||
if not deleted:
|
||||
return json.dumps({"error": f"Mental model '{mental_model_id}' not found in bank '{target_bank}'"})
|
||||
return json.dumps({"status": "deleted", "mental_model_id": mental_model_id})
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return json.dumps({"error": str(e)})
|
||||
except Exception as e:
|
||||
logger.error(f"Error deleting mental model: {e}", exc_info=True)
|
||||
return f'{{"error": "{e}"}}'
|
||||
@@ -1290,6 +1311,9 @@ def _register_delete_mental_model(mcp: FastMCP, memory: MemoryEngine, config: MC
|
||||
if not deleted:
|
||||
return {"error": f"Mental model '{mental_model_id}' not found in bank '{target_bank}'"}
|
||||
return {"status": "deleted", "mental_model_id": mental_model_id}
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return {"error": str(e)}
|
||||
except Exception as e:
|
||||
logger.error(f"Error deleting mental model: {e}", exc_info=True)
|
||||
return {"error": str(e)}
|
||||
@@ -1333,6 +1357,9 @@ def _register_refresh_mental_model(mcp: FastMCP, memory: MemoryEngine, config: M
|
||||
"message": f"Refresh queued for mental model '{mental_model_id}'.",
|
||||
}
|
||||
)
|
||||
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:
|
||||
@@ -1370,6 +1397,9 @@ def _register_refresh_mental_model(mcp: FastMCP, memory: MemoryEngine, config: M
|
||||
"status": "queued",
|
||||
"message": f"Refresh queued for mental model '{mental_model_id}'.",
|
||||
}
|
||||
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:
|
||||
@@ -1416,6 +1446,9 @@ def _register_list_directives(mcp: FastMCP, memory: MemoryEngine, config: MCPToo
|
||||
request_context=_get_request_context(config),
|
||||
)
|
||||
return json.dumps({"items": directives}, indent=2, default=str)
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return json.dumps({"error": str(e)})
|
||||
except Exception as e:
|
||||
logger.error(f"Error listing directives: {e}", exc_info=True)
|
||||
return f'{{"error": "{e}", "items": []}}'
|
||||
@@ -1449,6 +1482,9 @@ def _register_list_directives(mcp: FastMCP, memory: MemoryEngine, config: MCPToo
|
||||
request_context=_get_request_context(config),
|
||||
)
|
||||
return {"items": directives}
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return {"error": str(e)}
|
||||
except Exception as e:
|
||||
logger.error(f"Error listing directives: {e}", exc_info=True)
|
||||
return {"error": str(e), "items": []}
|
||||
@@ -1496,6 +1532,9 @@ def _register_create_directive(mcp: FastMCP, memory: MemoryEngine, config: MCPTo
|
||||
request_context=_get_request_context(config),
|
||||
)
|
||||
return json.dumps(directive, indent=2, default=str)
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return json.dumps({"error": str(e)})
|
||||
except Exception as e:
|
||||
logger.error(f"Error creating directive: {e}", exc_info=True)
|
||||
return f'{{"error": "{e}"}}'
|
||||
@@ -1537,6 +1576,9 @@ def _register_create_directive(mcp: FastMCP, memory: MemoryEngine, config: MCPTo
|
||||
request_context=_get_request_context(config),
|
||||
)
|
||||
return directive
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return {"error": str(e)}
|
||||
except Exception as e:
|
||||
logger.error(f"Error creating directive: {e}", exc_info=True)
|
||||
return {"error": str(e)}
|
||||
@@ -1574,6 +1616,9 @@ def _register_delete_directive(mcp: FastMCP, memory: MemoryEngine, config: MCPTo
|
||||
if not deleted:
|
||||
return json.dumps({"error": f"Directive '{directive_id}' not found"})
|
||||
return json.dumps({"status": "deleted", "directive_id": directive_id})
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return json.dumps({"error": str(e)})
|
||||
except Exception as e:
|
||||
logger.error(f"Error deleting directive: {e}", exc_info=True)
|
||||
return f'{{"error": "{e}"}}'
|
||||
@@ -1605,6 +1650,9 @@ def _register_delete_directive(mcp: FastMCP, memory: MemoryEngine, config: MCPTo
|
||||
if not deleted:
|
||||
return {"error": f"Directive '{directive_id}' not found"}
|
||||
return {"status": "deleted", "directive_id": directive_id}
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return {"error": str(e)}
|
||||
except Exception as e:
|
||||
logger.error(f"Error deleting directive: {e}", exc_info=True)
|
||||
return {"error": str(e)}
|
||||
@@ -1655,6 +1703,9 @@ def _register_list_memories(mcp: FastMCP, memory: MemoryEngine, config: MCPTools
|
||||
request_context=_get_request_context(config),
|
||||
)
|
||||
return json.dumps(result, indent=2, default=str)
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return json.dumps({"error": str(e)})
|
||||
except Exception as e:
|
||||
logger.error(f"Error listing memories: {e}", exc_info=True)
|
||||
return f'{{"error": "{e}"}}'
|
||||
@@ -1694,6 +1745,9 @@ def _register_list_memories(mcp: FastMCP, memory: MemoryEngine, config: MCPTools
|
||||
request_context=_get_request_context(config),
|
||||
)
|
||||
return result
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return {"error": str(e)}
|
||||
except Exception as e:
|
||||
logger.error(f"Error listing memories: {e}", exc_info=True)
|
||||
return {"error": str(e)}
|
||||
@@ -1731,6 +1785,9 @@ def _register_get_memory(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsCon
|
||||
if result is None:
|
||||
return json.dumps({"error": f"Memory '{memory_id}' not found"})
|
||||
return json.dumps(result, indent=2, default=str)
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return json.dumps({"error": str(e)})
|
||||
except Exception as e:
|
||||
logger.error(f"Error getting memory: {e}", exc_info=True)
|
||||
return f'{{"error": "{e}"}}'
|
||||
@@ -1762,6 +1819,9 @@ def _register_get_memory(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsCon
|
||||
if result is None:
|
||||
return {"error": f"Memory '{memory_id}' not found"}
|
||||
return result
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return {"error": str(e)}
|
||||
except Exception as e:
|
||||
logger.error(f"Error getting memory: {e}", exc_info=True)
|
||||
return {"error": str(e)}
|
||||
@@ -1796,6 +1856,9 @@ def _register_delete_memory(mcp: FastMCP, memory: MemoryEngine, config: MCPTools
|
||||
request_context=_get_request_context(config),
|
||||
)
|
||||
return json.dumps({"status": "deleted", "memory_id": memory_id, **result}, default=str)
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return json.dumps({"error": str(e)})
|
||||
except Exception as e:
|
||||
logger.error(f"Error deleting memory: {e}", exc_info=True)
|
||||
return f'{{"error": "{e}"}}'
|
||||
@@ -1824,6 +1887,9 @@ def _register_delete_memory(mcp: FastMCP, memory: MemoryEngine, config: MCPTools
|
||||
request_context=_get_request_context(config),
|
||||
)
|
||||
return {"status": "deleted", "memory_id": memory_id, **result}
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return {"error": str(e)}
|
||||
except Exception as e:
|
||||
logger.error(f"Error deleting memory: {e}", exc_info=True)
|
||||
return {"error": str(e)}
|
||||
@@ -1868,6 +1934,9 @@ def _register_list_documents(mcp: FastMCP, memory: MemoryEngine, config: MCPTool
|
||||
request_context=_get_request_context(config),
|
||||
)
|
||||
return json.dumps(result, indent=2, default=str)
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return json.dumps({"error": str(e)})
|
||||
except Exception as e:
|
||||
logger.error(f"Error listing documents: {e}", exc_info=True)
|
||||
return f'{{"error": "{e}"}}'
|
||||
@@ -1901,6 +1970,9 @@ def _register_list_documents(mcp: FastMCP, memory: MemoryEngine, config: MCPTool
|
||||
request_context=_get_request_context(config),
|
||||
)
|
||||
return result
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return {"error": str(e)}
|
||||
except Exception as e:
|
||||
logger.error(f"Error listing documents: {e}", exc_info=True)
|
||||
return {"error": str(e)}
|
||||
@@ -1938,6 +2010,9 @@ def _register_get_document(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsC
|
||||
if result is None:
|
||||
return json.dumps({"error": f"Document '{document_id}' not found"})
|
||||
return json.dumps(result, indent=2, default=str)
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return json.dumps({"error": str(e)})
|
||||
except Exception as e:
|
||||
logger.error(f"Error getting document: {e}", exc_info=True)
|
||||
return f'{{"error": "{e}"}}'
|
||||
@@ -1969,6 +2044,9 @@ def _register_get_document(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsC
|
||||
if result is None:
|
||||
return {"error": f"Document '{document_id}' not found"}
|
||||
return result
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return {"error": str(e)}
|
||||
except Exception as e:
|
||||
logger.error(f"Error getting document: {e}", exc_info=True)
|
||||
return {"error": str(e)}
|
||||
@@ -2004,6 +2082,9 @@ def _register_delete_document(mcp: FastMCP, memory: MemoryEngine, config: MCPToo
|
||||
request_context=_get_request_context(config),
|
||||
)
|
||||
return json.dumps({"status": "deleted", "document_id": document_id, **result}, default=str)
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return json.dumps({"error": str(e)})
|
||||
except Exception as e:
|
||||
logger.error(f"Error deleting document: {e}", exc_info=True)
|
||||
return f'{{"error": "{e}"}}'
|
||||
@@ -2033,6 +2114,9 @@ def _register_delete_document(mcp: FastMCP, memory: MemoryEngine, config: MCPToo
|
||||
request_context=_get_request_context(config),
|
||||
)
|
||||
return {"status": "deleted", "document_id": document_id, **result}
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return {"error": str(e)}
|
||||
except Exception as e:
|
||||
logger.error(f"Error deleting document: {e}", exc_info=True)
|
||||
return {"error": str(e)}
|
||||
@@ -2076,6 +2160,9 @@ def _register_list_operations(mcp: FastMCP, memory: MemoryEngine, config: MCPToo
|
||||
request_context=_get_request_context(config),
|
||||
)
|
||||
return json.dumps(result, indent=2, default=str)
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return json.dumps({"error": str(e)})
|
||||
except Exception as e:
|
||||
logger.error(f"Error listing operations: {e}", exc_info=True)
|
||||
return f'{{"error": "{e}"}}'
|
||||
@@ -2108,6 +2195,9 @@ def _register_list_operations(mcp: FastMCP, memory: MemoryEngine, config: MCPToo
|
||||
request_context=_get_request_context(config),
|
||||
)
|
||||
return result
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return {"error": str(e)}
|
||||
except Exception as e:
|
||||
logger.error(f"Error listing operations: {e}", exc_info=True)
|
||||
return {"error": str(e)}
|
||||
@@ -2143,6 +2233,9 @@ def _register_get_operation(mcp: FastMCP, memory: MemoryEngine, config: MCPTools
|
||||
request_context=_get_request_context(config),
|
||||
)
|
||||
return json.dumps(result, indent=2, default=str)
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return json.dumps({"error": str(e)})
|
||||
except Exception as e:
|
||||
logger.error(f"Error getting operation: {e}", exc_info=True)
|
||||
return f'{{"error": "{e}"}}'
|
||||
@@ -2172,6 +2265,9 @@ def _register_get_operation(mcp: FastMCP, memory: MemoryEngine, config: MCPTools
|
||||
request_context=_get_request_context(config),
|
||||
)
|
||||
return result
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return {"error": str(e)}
|
||||
except Exception as e:
|
||||
logger.error(f"Error getting operation: {e}", exc_info=True)
|
||||
return {"error": str(e)}
|
||||
@@ -2205,6 +2301,9 @@ def _register_cancel_operation(mcp: FastMCP, memory: MemoryEngine, config: MCPTo
|
||||
request_context=_get_request_context(config),
|
||||
)
|
||||
return json.dumps(result, indent=2, default=str)
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return json.dumps({"error": str(e)})
|
||||
except Exception as e:
|
||||
logger.error(f"Error cancelling operation: {e}", exc_info=True)
|
||||
return f'{{"error": "{e}"}}'
|
||||
@@ -2232,6 +2331,9 @@ def _register_cancel_operation(mcp: FastMCP, memory: MemoryEngine, config: MCPTo
|
||||
request_context=_get_request_context(config),
|
||||
)
|
||||
return result
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return {"error": str(e)}
|
||||
except Exception as e:
|
||||
logger.error(f"Error cancelling operation: {e}", exc_info=True)
|
||||
return {"error": str(e)}
|
||||
@@ -2275,6 +2377,9 @@ def _register_list_tags(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConf
|
||||
request_context=_get_request_context(config),
|
||||
)
|
||||
return json.dumps(result, indent=2, default=str)
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return json.dumps({"error": str(e)})
|
||||
except Exception as e:
|
||||
logger.error(f"Error listing tags: {e}", exc_info=True)
|
||||
return f'{{"error": "{e}"}}'
|
||||
@@ -2307,6 +2412,9 @@ def _register_list_tags(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConf
|
||||
request_context=_get_request_context(config),
|
||||
)
|
||||
return result
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return {"error": str(e)}
|
||||
except Exception as e:
|
||||
logger.error(f"Error listing tags: {e}", exc_info=True)
|
||||
return {"error": str(e)}
|
||||
@@ -2341,6 +2449,9 @@ def _register_get_bank(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfi
|
||||
if "disposition" in profile and hasattr(profile["disposition"], "model_dump"):
|
||||
profile["disposition"] = profile["disposition"].model_dump()
|
||||
return json.dumps(profile, indent=2, default=str)
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return json.dumps({"error": str(e)})
|
||||
except Exception as e:
|
||||
logger.error(f"Error getting bank: {e}", exc_info=True)
|
||||
return f'{{"error": "{e}"}}'
|
||||
@@ -2366,6 +2477,9 @@ def _register_get_bank(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfi
|
||||
if "disposition" in profile and hasattr(profile["disposition"], "model_dump"):
|
||||
profile["disposition"] = profile["disposition"].model_dump()
|
||||
return profile
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return {"error": str(e)}
|
||||
except Exception as e:
|
||||
logger.error(f"Error getting bank: {e}", exc_info=True)
|
||||
return {"error": str(e)}
|
||||
@@ -2396,6 +2510,9 @@ def _register_get_bank_stats(mcp: FastMCP, memory: MemoryEngine, config: MCPTool
|
||||
request_context=_get_request_context(config),
|
||||
)
|
||||
return json.dumps(result, indent=2, default=str)
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return json.dumps({"error": str(e)})
|
||||
except Exception as e:
|
||||
logger.error(f"Error getting bank stats: {e}", exc_info=True)
|
||||
return f'{{"error": "{e}"}}'
|
||||
@@ -2434,6 +2551,9 @@ def _register_update_bank(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsCo
|
||||
request_context=_get_request_context(config),
|
||||
)
|
||||
return json.dumps(result, indent=2, default=str)
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return json.dumps({"error": str(e)})
|
||||
except Exception as e:
|
||||
logger.error(f"Error updating bank: {e}", exc_info=True)
|
||||
return f'{{"error": "{e}"}}'
|
||||
@@ -2466,6 +2586,9 @@ def _register_update_bank(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsCo
|
||||
request_context=_get_request_context(config),
|
||||
)
|
||||
return result
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return {"error": str(e)}
|
||||
except Exception as e:
|
||||
logger.error(f"Error updating bank: {e}", exc_info=True)
|
||||
return {"error": str(e)}
|
||||
@@ -2499,6 +2622,9 @@ def _register_delete_bank(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsCo
|
||||
request_context=_get_request_context(config),
|
||||
)
|
||||
return json.dumps({"status": "deleted", "bank_id": target_bank, **result}, default=str)
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return json.dumps({"error": str(e)})
|
||||
except Exception as e:
|
||||
logger.error(f"Error deleting bank: {e}", exc_info=True)
|
||||
return f'{{"error": "{e}"}}'
|
||||
@@ -2523,6 +2649,9 @@ def _register_delete_bank(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsCo
|
||||
request_context=_get_request_context(config),
|
||||
)
|
||||
return {"status": "deleted", "bank_id": target_bank, **result}
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return {"error": str(e)}
|
||||
except Exception as e:
|
||||
logger.error(f"Error deleting bank: {e}", exc_info=True)
|
||||
return {"error": str(e)}
|
||||
@@ -2558,6 +2687,9 @@ def _register_clear_memories(mcp: FastMCP, memory: MemoryEngine, config: MCPTool
|
||||
request_context=_get_request_context(config),
|
||||
)
|
||||
return json.dumps({"status": "cleared", "bank_id": target_bank, **result}, default=str)
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return json.dumps({"error": str(e)})
|
||||
except Exception as e:
|
||||
logger.error(f"Error clearing memories: {e}", exc_info=True)
|
||||
return f'{{"error": "{e}"}}'
|
||||
@@ -2587,6 +2719,9 @@ def _register_clear_memories(mcp: FastMCP, memory: MemoryEngine, config: MCPTool
|
||||
request_context=_get_request_context(config),
|
||||
)
|
||||
return {"status": "cleared", "bank_id": target_bank, **result}
|
||||
except OperationValidationError as e:
|
||||
logger.warning(f"Operation rejected: {e}")
|
||||
return {"error": str(e)}
|
||||
except Exception as e:
|
||||
logger.error(f"Error clearing memories: {e}", exc_info=True)
|
||||
return {"error": str(e)}
|
||||
|
||||
@@ -18,6 +18,7 @@ No alembic.ini required - all configuration is done programmatically.
|
||||
import hashlib
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
from alembic import command
|
||||
@@ -220,13 +221,40 @@ def run_migrations(
|
||||
lock_id = _get_schema_lock_id(schema) if schema else MIGRATION_LOCK_ID
|
||||
schema_name = schema or "public"
|
||||
|
||||
# Use PostgreSQL advisory lock to coordinate between distributed workers
|
||||
# Use PostgreSQL advisory lock to coordinate between distributed workers.
|
||||
#
|
||||
# IMPORTANT: We must avoid holding an open transaction on the advisory-lock
|
||||
# connection while CREATE INDEX CONCURRENTLY runs inside a migration.
|
||||
# CONCURRENTLY waits for ALL active transactions to finish before the index
|
||||
# becomes valid. If the advisory-lock connection (or any waiting worker's
|
||||
# connection) holds an open transaction, CONCURRENTLY deadlocks:
|
||||
# - migration worker waits for other workers' transactions to close
|
||||
# - other workers wait for the advisory lock to be released
|
||||
#
|
||||
# Fix:
|
||||
# 1. Use pg_try_advisory_lock (non-blocking) in a poll loop instead of
|
||||
# blocking pg_advisory_lock, so we can COMMIT the transaction between
|
||||
# retries. Between retries the connection holds no open transaction.
|
||||
# 2. After acquiring the lock, COMMIT the transaction on the advisory-lock
|
||||
# connection itself before running migrations. pg_advisory_lock is
|
||||
# session-level, so the lock survives the COMMIT.
|
||||
engine = create_engine(database_url)
|
||||
with engine.connect() as conn:
|
||||
# pg_advisory_lock blocks until the lock is acquired
|
||||
# The lock is automatically released when the connection closes
|
||||
logger.debug(f"Acquiring migration advisory lock for schema '{schema_name}' (id={lock_id})...")
|
||||
conn.execute(text(f"SELECT pg_advisory_lock({lock_id})"))
|
||||
while True:
|
||||
acquired = conn.execute(text(f"SELECT pg_try_advisory_lock({lock_id})")).scalar()
|
||||
if acquired:
|
||||
break
|
||||
# Commit the transaction so this connection holds no open snapshot
|
||||
# while waiting. This prevents blocking CREATE INDEX CONCURRENTLY
|
||||
# that may be running in the migration worker.
|
||||
conn.commit()
|
||||
time.sleep(0.5)
|
||||
|
||||
# Commit AFTER acquiring the lock too. pg_advisory_lock is session-level
|
||||
# and survives the COMMIT, but the open transaction on this connection
|
||||
# would otherwise block any CREATE INDEX CONCURRENTLY in the migration.
|
||||
conn.commit()
|
||||
logger.debug("Migration advisory lock acquired")
|
||||
|
||||
try:
|
||||
@@ -347,6 +375,13 @@ def run_migrations(
|
||||
"Please install it with: CREATE EXTENSION vectorscale CASCADE;"
|
||||
) from e
|
||||
|
||||
# Commit any pending transaction on the advisory-lock connection
|
||||
# before running migrations. Some code paths above (e.g., the
|
||||
# pgvector extension check) may have started a transaction via
|
||||
# SQLAlchemy's autobegin. If we leave it open, CREATE INDEX
|
||||
# CONCURRENTLY inside a migration will deadlock waiting for it.
|
||||
conn.commit()
|
||||
|
||||
# Run migrations while holding the lock
|
||||
_run_migrations_internal(database_url, script_location, schema=schema)
|
||||
finally:
|
||||
|
||||
@@ -22,6 +22,7 @@ class RequestContext:
|
||||
tenant_id: str | None = None # Tenant identifier (set by extension after auth)
|
||||
internal: bool = False # True for background/internal operations (skips extension auth)
|
||||
user_initiated: bool = False # True for async operations that originated from a user request
|
||||
allowed_bank_ids: list[str] | None = None # None = unrestricted (all banks)
|
||||
|
||||
|
||||
from pgvector.sqlalchemy import Vector
|
||||
|
||||
@@ -0,0 +1,13 @@
|
||||
"""Webhook system for Hindsight API event notifications."""
|
||||
|
||||
from .manager import WebhookManager
|
||||
from .models import ConsolidationEventData, RetainEventData, WebhookConfig, WebhookEvent, WebhookEventType
|
||||
|
||||
__all__ = [
|
||||
"WebhookManager",
|
||||
"WebhookConfig",
|
||||
"WebhookEvent",
|
||||
"WebhookEventType",
|
||||
"ConsolidationEventData",
|
||||
"RetainEventData",
|
||||
]
|
||||
@@ -0,0 +1,242 @@
|
||||
"""Webhook manager for delivering event notifications."""
|
||||
|
||||
import hashlib
|
||||
import hmac
|
||||
import json
|
||||
import logging
|
||||
import uuid
|
||||
from datetime import datetime, timezone
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import asyncpg
|
||||
|
||||
from .models import WebhookConfig, WebhookEvent, WebhookHttpConfig
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from hindsight_api.extensions.tenant import TenantExtension
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Retry delay schedule in seconds: 5 retries after the first attempt.
|
||||
# Fast early retries catch transient failures; later retries handle longer outages.
|
||||
RETRY_DELAYS = [5, 300, 1800, 7200, 18000]
|
||||
MAX_ATTEMPTS = len(RETRY_DELAYS) + 1 # first attempt + len(RETRY_DELAYS) retries
|
||||
|
||||
|
||||
def _fq_table(table: str, schema: str | None = None) -> str:
|
||||
"""Get fully-qualified table name with optional schema prefix."""
|
||||
if schema:
|
||||
return f'"{schema}".{table}'
|
||||
return table
|
||||
|
||||
|
||||
def _parse_http_config(value: str | dict | None) -> WebhookHttpConfig:
|
||||
"""Parse http_config column value (JSONB returned as text or dict) into a model."""
|
||||
if value is None:
|
||||
return WebhookHttpConfig()
|
||||
if isinstance(value, str):
|
||||
return WebhookHttpConfig.model_validate_json(value)
|
||||
return WebhookHttpConfig.model_validate(value)
|
||||
|
||||
|
||||
class WebhookManager:
|
||||
"""
|
||||
Manages webhook registration and event firing.
|
||||
|
||||
Supports both global webhooks (configured via env vars) and per-bank
|
||||
webhooks stored in the database. Deliveries are queued as async_operations
|
||||
tasks (operation_type='webhook_delivery') and picked up by the worker poller.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
pool: asyncpg.Pool,
|
||||
global_webhooks: list[WebhookConfig],
|
||||
tenant_extension: "TenantExtension | None" = None,
|
||||
):
|
||||
self._pool = pool
|
||||
self._global_webhooks = global_webhooks
|
||||
self._tenant_extension = tenant_extension
|
||||
|
||||
def _sign_payload(self, secret: str, payload_bytes: bytes) -> str:
|
||||
"""Compute HMAC-SHA256 signature for a payload."""
|
||||
return "sha256=" + hmac.new(secret.encode(), payload_bytes, hashlib.sha256).hexdigest()
|
||||
|
||||
async def fire_event(self, event: WebhookEvent, schema: str | None = None) -> None:
|
||||
"""
|
||||
Queue webhook deliveries for an event as async_operations tasks.
|
||||
|
||||
Loads per-bank and global webhooks, inserts pending webhook_delivery tasks for
|
||||
any webhook whose event_types list matches the fired event type. The worker
|
||||
poller picks these up and calls MemoryEngine._handle_webhook_delivery().
|
||||
|
||||
Args:
|
||||
event: The event to deliver.
|
||||
schema: Database schema (for multi-tenant). None = default schema.
|
||||
"""
|
||||
webhook_table = _fq_table("webhooks", schema)
|
||||
ops_table = _fq_table("async_operations", schema)
|
||||
now = datetime.now(timezone.utc)
|
||||
payload_str = event.model_dump_json()
|
||||
|
||||
try:
|
||||
# Load per-bank webhooks from DB (bank-specific + global NULL rows)
|
||||
rows = await self._pool.fetch(
|
||||
f"""
|
||||
SELECT id, bank_id, url, secret, event_types, enabled, http_config::text
|
||||
FROM {webhook_table}
|
||||
WHERE (bank_id = $1 OR bank_id IS NULL) AND enabled = true
|
||||
""",
|
||||
event.bank_id,
|
||||
)
|
||||
|
||||
db_webhooks = [
|
||||
WebhookConfig(
|
||||
id=str(row["id"]),
|
||||
bank_id=row["bank_id"],
|
||||
url=row["url"],
|
||||
secret=row["secret"],
|
||||
event_types=list(row["event_types"]) if row["event_types"] else [],
|
||||
enabled=row["enabled"],
|
||||
http_config=_parse_http_config(row["http_config"]),
|
||||
)
|
||||
for row in rows
|
||||
]
|
||||
|
||||
# Merge with global webhooks from env config
|
||||
all_webhooks = self._global_webhooks + db_webhooks
|
||||
matched = 0
|
||||
|
||||
for webhook in all_webhooks:
|
||||
if not webhook.enabled:
|
||||
continue
|
||||
if event.event.value not in webhook.event_types:
|
||||
continue
|
||||
|
||||
operation_id = uuid.uuid4()
|
||||
webhook_id = webhook.id if webhook.id else None
|
||||
|
||||
task_payload = json.dumps(
|
||||
{
|
||||
"type": "webhook_delivery",
|
||||
"operation_id": str(operation_id),
|
||||
"bank_id": event.bank_id,
|
||||
"url": webhook.url,
|
||||
"secret": webhook.secret,
|
||||
"event_type": event.event.value,
|
||||
"payload": payload_str,
|
||||
"webhook_id": webhook_id,
|
||||
"http_config": webhook.http_config.model_dump(),
|
||||
}
|
||||
)
|
||||
|
||||
await self._pool.execute(
|
||||
f"""
|
||||
INSERT INTO {ops_table}
|
||||
(operation_id, bank_id, operation_type, status, task_payload, result_metadata, created_at, updated_at)
|
||||
VALUES ($1, $2, 'webhook_delivery', 'pending', $3::jsonb, '{{}}'::jsonb, $4, $4)
|
||||
""",
|
||||
operation_id,
|
||||
event.bank_id,
|
||||
task_payload,
|
||||
now,
|
||||
)
|
||||
matched += 1
|
||||
|
||||
logger.debug(f"Fired webhook event {event.event} for bank {event.bank_id}: {matched} delivery(ies) queued")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to queue webhook deliveries for event {event.event}: {e}")
|
||||
|
||||
async def fire_event_with_conn(
|
||||
self, event: WebhookEvent, conn: asyncpg.Connection, schema: str | None = None
|
||||
) -> None:
|
||||
"""
|
||||
Queue webhook deliveries within an existing database connection/transaction.
|
||||
|
||||
Identical to fire_event() but uses the provided connection instead of acquiring
|
||||
one from the pool. Use this to atomically insert delivery tasks in the same
|
||||
transaction as the primary operation (transactional outbox pattern).
|
||||
|
||||
Args:
|
||||
event: The event to deliver.
|
||||
conn: Existing asyncpg connection (may be inside an active transaction).
|
||||
schema: Database schema (for multi-tenant). None = default schema.
|
||||
"""
|
||||
webhook_table = _fq_table("webhooks", schema)
|
||||
ops_table = _fq_table("async_operations", schema)
|
||||
now = datetime.now(timezone.utc)
|
||||
payload_str = event.model_dump_json()
|
||||
|
||||
try:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT id, bank_id, url, secret, event_types, enabled, http_config::text
|
||||
FROM {webhook_table}
|
||||
WHERE (bank_id = $1 OR bank_id IS NULL) AND enabled = true
|
||||
""",
|
||||
event.bank_id,
|
||||
)
|
||||
|
||||
db_webhooks = [
|
||||
WebhookConfig(
|
||||
id=str(row["id"]),
|
||||
bank_id=row["bank_id"],
|
||||
url=row["url"],
|
||||
secret=row["secret"],
|
||||
event_types=list(row["event_types"]) if row["event_types"] else [],
|
||||
enabled=row["enabled"],
|
||||
http_config=_parse_http_config(row["http_config"]),
|
||||
)
|
||||
for row in rows
|
||||
]
|
||||
|
||||
all_webhooks = self._global_webhooks + db_webhooks
|
||||
matched = 0
|
||||
|
||||
for webhook in all_webhooks:
|
||||
if not webhook.enabled:
|
||||
continue
|
||||
if event.event.value not in webhook.event_types:
|
||||
continue
|
||||
|
||||
operation_id = uuid.uuid4()
|
||||
webhook_id = webhook.id if webhook.id else None
|
||||
|
||||
task_payload = json.dumps(
|
||||
{
|
||||
"type": "webhook_delivery",
|
||||
"operation_id": str(operation_id),
|
||||
"bank_id": event.bank_id,
|
||||
"url": webhook.url,
|
||||
"secret": webhook.secret,
|
||||
"event_type": event.event.value,
|
||||
"payload": payload_str,
|
||||
"webhook_id": webhook_id,
|
||||
"http_config": webhook.http_config.model_dump(),
|
||||
}
|
||||
)
|
||||
|
||||
await conn.execute(
|
||||
f"""
|
||||
INSERT INTO {ops_table}
|
||||
(operation_id, bank_id, operation_type, status, task_payload, result_metadata, created_at, updated_at)
|
||||
VALUES ($1, $2, 'webhook_delivery', 'pending', $3::jsonb, '{{}}'::jsonb, $4, $4)
|
||||
""",
|
||||
operation_id,
|
||||
event.bank_id,
|
||||
task_payload,
|
||||
now,
|
||||
)
|
||||
matched += 1
|
||||
|
||||
logger.debug(
|
||||
f"Fired webhook event {event.event} for bank {event.bank_id}: {matched} delivery(ies) queued (in-transaction)"
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Failed to queue webhook deliveries (in-transaction) for event {event.event}: {e}. "
|
||||
"CRITICAL: The enclosing database transaction is now aborted and will roll back all changes."
|
||||
)
|
||||
raise
|
||||
@@ -0,0 +1,51 @@
|
||||
"""Pydantic models for the webhook system."""
|
||||
|
||||
from datetime import datetime
|
||||
from enum import StrEnum
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
class WebhookEventType(StrEnum):
|
||||
CONSOLIDATION_COMPLETED = "consolidation.completed"
|
||||
RETAIN_COMPLETED = "retain.completed"
|
||||
|
||||
|
||||
class ConsolidationEventData(BaseModel):
|
||||
observations_created: int | None = None
|
||||
observations_updated: int | None = None
|
||||
observations_deleted: int | None = None
|
||||
error_message: str | None = None
|
||||
|
||||
|
||||
class RetainEventData(BaseModel):
|
||||
document_id: str | None = None
|
||||
tags: list[str] | None = None
|
||||
|
||||
|
||||
class WebhookEvent(BaseModel):
|
||||
event: WebhookEventType
|
||||
bank_id: str
|
||||
operation_id: str
|
||||
status: str # "completed" or "failed"
|
||||
timestamp: datetime
|
||||
data: ConsolidationEventData | RetainEventData
|
||||
|
||||
|
||||
class WebhookHttpConfig(BaseModel):
|
||||
"""HTTP delivery configuration for a webhook."""
|
||||
|
||||
method: str = Field(default="POST", description="HTTP method: GET or POST")
|
||||
timeout_seconds: int = Field(default=30, description="HTTP request timeout in seconds")
|
||||
headers: dict[str, str] = Field(default_factory=dict, description="Custom HTTP headers")
|
||||
params: dict[str, str] = Field(default_factory=dict, description="Custom HTTP query parameters")
|
||||
|
||||
|
||||
class WebhookConfig(BaseModel):
|
||||
id: str
|
||||
bank_id: str | None
|
||||
url: str
|
||||
secret: str | None
|
||||
event_types: list[str]
|
||||
enabled: bool
|
||||
http_config: WebhookHttpConfig = Field(default_factory=WebhookHttpConfig)
|
||||
@@ -0,0 +1,9 @@
|
||||
from datetime import datetime
|
||||
|
||||
|
||||
class RetryTaskAt(Exception):
|
||||
"""Raise from a task handler to schedule a retry at a specific time."""
|
||||
|
||||
def __init__(self, retry_at: datetime, message: str = ""):
|
||||
self.retry_at = retry_at
|
||||
super().__init__(message)
|
||||
@@ -219,7 +219,6 @@ def main():
|
||||
worker_id=args.worker_id,
|
||||
executor=memory.execute_task,
|
||||
poll_interval_ms=args.poll_interval,
|
||||
max_retries=args.max_retries,
|
||||
schema=schema,
|
||||
tenant_extension=tenant_extension,
|
||||
max_slots=config.worker_max_slots,
|
||||
|
||||
@@ -14,6 +14,8 @@ from collections.abc import Awaitable, Callable
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from .exceptions import RetryTaskAt
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import asyncpg
|
||||
|
||||
@@ -57,7 +59,6 @@ class WorkerPoller:
|
||||
worker_id: str,
|
||||
executor: Callable[[dict[str, Any]], Awaitable[None]],
|
||||
poll_interval_ms: int = 500,
|
||||
max_retries: int = 3,
|
||||
schema: str | None = None,
|
||||
tenant_extension: "TenantExtension | None" = None,
|
||||
max_slots: int = 10,
|
||||
@@ -71,7 +72,6 @@ class WorkerPoller:
|
||||
worker_id: Unique identifier for this worker
|
||||
executor: Async function to execute tasks (typically MemoryEngine.execute_task)
|
||||
poll_interval_ms: Interval between polls when no tasks found (milliseconds)
|
||||
max_retries: Maximum retry attempts before marking task as failed
|
||||
schema: Database schema for single-tenant support (deprecated, use tenant_extension)
|
||||
tenant_extension: Extension for dynamic multi-tenant discovery. If None, creates a
|
||||
DefaultTenantExtension with the configured schema.
|
||||
@@ -82,7 +82,6 @@ class WorkerPoller:
|
||||
self._worker_id = worker_id
|
||||
self._executor = executor
|
||||
self._poll_interval_ms = poll_interval_ms
|
||||
self._max_retries = max_retries
|
||||
self._schema = schema
|
||||
# Always set tenant extension (use DefaultTenantExtension if none provided)
|
||||
if tenant_extension is None:
|
||||
@@ -218,11 +217,12 @@ class WorkerPoller:
|
||||
# 1. Claim non-consolidation tasks (up to limit)
|
||||
non_consolidation_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, task_payload
|
||||
SELECT operation_id, task_payload, retry_count
|
||||
FROM {table}
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type != 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
ORDER BY created_at
|
||||
LIMIT $1
|
||||
FOR UPDATE SKIP LOCKED
|
||||
@@ -238,11 +238,12 @@ class WorkerPoller:
|
||||
if consolidation_limit > 0 and remaining_limit > 0:
|
||||
consolidation_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT operation_id, task_payload
|
||||
SELECT operation_id, task_payload, retry_count
|
||||
FROM {table} AS pending
|
||||
WHERE status = 'pending'
|
||||
AND task_payload IS NOT NULL
|
||||
AND operation_type = 'consolidation'
|
||||
AND (next_retry_at IS NULL OR next_retry_at <= NOW())
|
||||
AND NOT EXISTS (
|
||||
SELECT 1 FROM {table} AS processing
|
||||
WHERE processing.bank_id = pending.bank_id
|
||||
@@ -274,14 +275,19 @@ class WorkerPoller:
|
||||
)
|
||||
|
||||
# Parse and return task payloads with schema context
|
||||
return [
|
||||
ClaimedTask(
|
||||
operation_id=str(row["operation_id"]),
|
||||
task_dict=json.loads(row["task_payload"]),
|
||||
schema=schema,
|
||||
result = []
|
||||
for row in all_rows:
|
||||
task_dict = json.loads(row["task_payload"])
|
||||
task_dict["_retry_count"] = row["retry_count"]
|
||||
task_dict["_operation_id"] = str(row["operation_id"])
|
||||
result.append(
|
||||
ClaimedTask(
|
||||
operation_id=str(row["operation_id"]),
|
||||
task_dict=task_dict,
|
||||
schema=schema,
|
||||
)
|
||||
)
|
||||
for row in all_rows
|
||||
]
|
||||
return result
|
||||
|
||||
async def _mark_completed(self, operation_id: str, schema: str | None):
|
||||
"""Mark a task as completed."""
|
||||
@@ -310,40 +316,22 @@ class WorkerPoller:
|
||||
error_message,
|
||||
)
|
||||
|
||||
async def _retry_or_fail(self, operation_id: str, error_message: str, schema: str | None):
|
||||
"""Increment retry count or mark as failed if max retries exceeded."""
|
||||
async def _schedule_retry(self, operation_id: str, retry_at: "Any", error_message: str, schema: str | None):
|
||||
"""Reset task to pending with a future retry timestamp."""
|
||||
table = fq_table("async_operations", schema)
|
||||
|
||||
# Get current retry count
|
||||
row = await self._pool.fetchrow(
|
||||
f"SELECT retry_count FROM {table} WHERE operation_id = $1",
|
||||
error_message = error_message[:5000] if len(error_message) > 5000 else error_message
|
||||
await self._pool.execute(
|
||||
f"""
|
||||
UPDATE {table}
|
||||
SET status = 'pending', next_retry_at = $2, worker_id = NULL, claimed_at = NULL,
|
||||
retry_count = retry_count + 1, error_message = $3, updated_at = now()
|
||||
WHERE operation_id = $1
|
||||
""",
|
||||
operation_id,
|
||||
retry_at,
|
||||
error_message,
|
||||
)
|
||||
|
||||
if row is None:
|
||||
logger.warning(f"Operation {operation_id} not found, cannot retry")
|
||||
return
|
||||
|
||||
retry_count = row["retry_count"]
|
||||
|
||||
if retry_count >= self._max_retries:
|
||||
# Max retries exceeded, mark as failed
|
||||
await self._mark_failed(
|
||||
operation_id, f"Max retries ({self._max_retries}) exceeded. Last error: {error_message}", schema
|
||||
)
|
||||
logger.error(f"Task {operation_id} failed after {retry_count} retries")
|
||||
else:
|
||||
# Increment retry and reset to pending
|
||||
await self._pool.execute(
|
||||
f"""
|
||||
UPDATE {table}
|
||||
SET status = 'pending', worker_id = NULL, claimed_at = NULL,
|
||||
retry_count = retry_count + 1, updated_at = now()
|
||||
WHERE operation_id = $1
|
||||
""",
|
||||
operation_id,
|
||||
)
|
||||
logger.warning(f"Task {operation_id} failed, will retry (attempt {retry_count + 1}/{self._max_retries})")
|
||||
logger.warning(f"Task {operation_id} scheduled for retry at {retry_at}: {error_message}")
|
||||
|
||||
async def execute_task(self, task: ClaimedTask):
|
||||
"""Execute a single task as a background job (fire-and-forget)."""
|
||||
@@ -376,11 +364,12 @@ class WorkerPoller:
|
||||
del self._in_flight_by_type[operation_type]
|
||||
|
||||
async def _execute_task_inner(self, task: ClaimedTask):
|
||||
"""Inner task execution with error handling.
|
||||
"""Inner task execution with retry/fail handling.
|
||||
|
||||
Note: The executor (MemoryEngine.execute_task) handles status marking internally
|
||||
(marking operations as completed/failed and handling retries). This method should
|
||||
NOT override those status updates.
|
||||
Tasks that want to be retried raise RetryTaskAt; the poller sets next_retry_at
|
||||
and resets status to 'pending'. All other exceptions are marked as failed immediately.
|
||||
Non-retryable failures (e.g., file_convert_retain) are handled by the executor
|
||||
internally — it marks the operation as failed and returns normally.
|
||||
"""
|
||||
task_type = task.task_dict.get("type", "unknown")
|
||||
bank_id = task.task_dict.get("bank_id", "unknown")
|
||||
@@ -392,11 +381,12 @@ class WorkerPoller:
|
||||
task.task_dict["_schema"] = task.schema
|
||||
await self._executor(task.task_dict)
|
||||
logger.debug(f"Task {task.operation_id} execution finished")
|
||||
except RetryTaskAt as e:
|
||||
await self._schedule_retry(task.operation_id, e.retry_at, str(e), task.schema)
|
||||
except Exception as e:
|
||||
# The executor should handle its own errors, but if an unexpected exception
|
||||
# propagates (e.g., from schema setup), log it as a warning
|
||||
logger.error(f"Task {task.operation_id} raised unexpected exception: {e}")
|
||||
logger.error(f"Task {task.operation_id} failed: {e}")
|
||||
traceback.print_exc()
|
||||
await self._mark_failed(task.operation_id, str(e), task.schema)
|
||||
|
||||
async def recover_own_tasks(self) -> int:
|
||||
"""
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "hindsight-api"
|
||||
version = "0.4.13"
|
||||
version = "0.4.16"
|
||||
description = "Hindsight: Agent Memory That Works Like Human Memory"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.11"
|
||||
@@ -43,8 +43,8 @@ dependencies = [
|
||||
"cohere>=5.0.0",
|
||||
"flashrank>=0.2.0",
|
||||
"litellm>=1.0.0",
|
||||
"markitdown[pdf,docx,pptx,xlsx,xls]>=0.1.4", # File to markdown conversion
|
||||
"obstore>=0.4.0", # S3/GCS/Azure object storage client (Rust-backed)
|
||||
"markitdown[pdf,docx,pptx,xlsx,xls]>=0.1.4", # File to markdown conversion
|
||||
"obstore>=0.4.0", # S3/GCS/Azure object storage client (Rust-backed)
|
||||
# Local ML models for embeddings/reranking - can be excluded in Docker with INCLUDE_LOCAL_MODELS=false
|
||||
"sentence-transformers>=3.3.0",
|
||||
"transformers>=4.53.0", # Security fixes for ReDoS vulnerabilities
|
||||
@@ -53,11 +53,16 @@ dependencies = [
|
||||
# Transitive dependency security fixes
|
||||
"pyasn1>=0.6.2", # DoS vulnerability fix
|
||||
"urllib3>=2.6.3", # Decompression-bomb safeguards bypass fix
|
||||
"langchain-core>=1.2.5", # Serialization injection vulnerability fix
|
||||
"langchain-core>=1.2.11", # Serialization injection + SSRF vulnerability fix
|
||||
"langsmith>=0.6.3", # SSRF via tracing header injection fix
|
||||
"protobuf>=6.33.5", # JSON recursion depth bypass fix
|
||||
"pillow>=12.1.1", # Out-of-bounds write in PSD image loading fix
|
||||
"cryptography>=46.0.5", # Subgroup attack vulnerability fix
|
||||
"filelock>=3.20.1", # TOCTOU race condition fix
|
||||
"authlib>=1.6.6", # Account takeover vulnerability fix
|
||||
"aiohttp>=3.13.3", # Multiple DoS vulnerabilities
|
||||
"claude-agent-sdk>=0.1.27",
|
||||
"einops>=0.8.2",
|
||||
]
|
||||
|
||||
[project.optional-dependencies]
|
||||
|
||||
@@ -14,6 +14,7 @@ async def test_submit_async_retain_includes_document_tags_in_task_payload():
|
||||
engine = MemoryEngine.__new__(MemoryEngine)
|
||||
engine._initialized = True
|
||||
engine._authenticate_tenant = AsyncMock()
|
||||
engine._operation_validator = None
|
||||
engine._submit_async_operation = AsyncMock(return_value={"operation_id": "op-1"})
|
||||
|
||||
# Mock the pool and connection for parent operation creation
|
||||
|
||||
@@ -413,7 +413,6 @@ async def test_worker_batch_recovery(memory, request_context):
|
||||
worker_id="test_worker_recovery",
|
||||
executor=memory,
|
||||
poll_interval_ms=100,
|
||||
max_retries=3,
|
||||
schema=schema,
|
||||
tenant_extension=tenant_extension,
|
||||
max_slots=5,
|
||||
|
||||
@@ -1,334 +1,171 @@
|
||||
"""
|
||||
Tests for combined scoring functionality.
|
||||
Tests for combined scoring (apply_combined_scoring).
|
||||
|
||||
Verifies that:
|
||||
1. RRF scores are properly normalized to [0, 1] range
|
||||
2. Combined scoring formula is applied correctly
|
||||
3. Tracer captures normalized values (not raw values)
|
||||
The function applies multiplicative recency/temporal boosts to the cross-encoder
|
||||
score so that the relative influence of these signals is proportional to the base
|
||||
relevance score, independent of the cross-encoder model's score calibration.
|
||||
"""
|
||||
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
from datetime import datetime, timezone
|
||||
from hindsight_api.engine.search.types import RetrievalResult, MergedCandidate, ScoredResult
|
||||
from hindsight_api.engine.memory_engine import Budget
|
||||
from hindsight_api import RequestContext
|
||||
|
||||
from hindsight_api.engine.search.reranking import apply_combined_scoring, _RECENCY_ALPHA, _TEMPORAL_ALPHA
|
||||
from hindsight_api.engine.search.types import MergedCandidate, RetrievalResult, ScoredResult
|
||||
|
||||
UTC = timezone.utc
|
||||
NOW = datetime(2024, 6, 1, tzinfo=UTC)
|
||||
|
||||
|
||||
class TestRRFNormalization:
|
||||
"""Test that RRF scores are properly normalized."""
|
||||
def _make_result(
|
||||
ce_norm: float,
|
||||
occurred_start: datetime | None = None,
|
||||
temporal_proximity: float | None = None,
|
||||
) -> ScoredResult:
|
||||
retrieval = MagicMock(spec=RetrievalResult)
|
||||
retrieval.occurred_start = occurred_start
|
||||
retrieval.temporal_proximity = temporal_proximity
|
||||
|
||||
def test_rrf_normalized_range(self):
|
||||
"""RRF normalized values should be in [0, 1] range, not raw [0.04, 0.06]."""
|
||||
# Simulate RRF scores like what we get from actual retrieval
|
||||
raw_rrf_scores = [0.0607, 0.0550, 0.0480, 0.0390]
|
||||
candidate = MagicMock(spec=MergedCandidate)
|
||||
candidate.retrieval = retrieval
|
||||
candidate.rrf_score = 0.05
|
||||
|
||||
max_rrf = max(raw_rrf_scores)
|
||||
min_rrf = min(raw_rrf_scores)
|
||||
rrf_range = max_rrf - min_rrf
|
||||
|
||||
normalized = []
|
||||
for score in raw_rrf_scores:
|
||||
if rrf_range > 0:
|
||||
norm = (score - min_rrf) / rrf_range
|
||||
else:
|
||||
norm = 0.5
|
||||
normalized.append(norm)
|
||||
|
||||
# Verify normalized values are in [0, 1]
|
||||
for i, norm in enumerate(normalized):
|
||||
assert 0.0 <= norm <= 1.0, f"Normalized RRF {norm} not in [0, 1] for raw {raw_rrf_scores[i]}"
|
||||
|
||||
# Highest raw should be 1.0
|
||||
assert normalized[0] == 1.0, f"Highest RRF should normalize to 1.0, got {normalized[0]}"
|
||||
|
||||
# Lowest raw should be 0.0
|
||||
assert normalized[-1] == 0.0, f"Lowest RRF should normalize to 0.0, got {normalized[-1]}"
|
||||
|
||||
def test_rrf_all_same_scores(self):
|
||||
"""When all RRF scores are the same, normalized should be 0.5 (neutral)."""
|
||||
raw_rrf_scores = [0.0500, 0.0500, 0.0500]
|
||||
|
||||
max_rrf = max(raw_rrf_scores)
|
||||
min_rrf = min(raw_rrf_scores)
|
||||
rrf_range = max_rrf - min_rrf
|
||||
|
||||
normalized = []
|
||||
for score in raw_rrf_scores:
|
||||
if rrf_range > 0:
|
||||
norm = (score - min_rrf) / rrf_range
|
||||
else:
|
||||
norm = 0.5 # Neutral value when all same
|
||||
normalized.append(norm)
|
||||
|
||||
# All should be 0.5 when scores are identical
|
||||
for norm in normalized:
|
||||
assert norm == 0.5, f"Expected 0.5 for identical scores, got {norm}"
|
||||
return ScoredResult(
|
||||
candidate=candidate,
|
||||
cross_encoder_score=1.0,
|
||||
cross_encoder_score_normalized=ce_norm,
|
||||
weight=ce_norm,
|
||||
)
|
||||
|
||||
|
||||
class TestCombinedScoringFormula:
|
||||
"""Test that the combined scoring formula is applied correctly."""
|
||||
class TestBoostFormula:
|
||||
def test_neutral_signals_leave_score_unchanged(self):
|
||||
"""recency=0.5 and temporal=0.5 both produce boost=1.0, so weight == ce."""
|
||||
sr = _make_result(ce_norm=0.6)
|
||||
apply_combined_scoring([sr], now=NOW)
|
||||
assert abs(sr.weight - 0.6) < 1e-9
|
||||
|
||||
def test_combined_score_calculation(self):
|
||||
"""Verify the weighted combination: 0.6*CE + 0.2*RRF + 0.1*temporal + 0.1*recency."""
|
||||
# Test case 1: All components at 1.0
|
||||
ce_norm = 1.0
|
||||
rrf_norm = 1.0
|
||||
temporal = 1.0
|
||||
recency = 1.0
|
||||
def test_max_recency_boost(self):
|
||||
"""A memory from today (recency≈1.0) should boost by (1 + alpha*0.5)."""
|
||||
sr = _make_result(ce_norm=0.5, occurred_start=NOW)
|
||||
apply_combined_scoring([sr], now=NOW)
|
||||
expected = 0.5 * (1.0 + _RECENCY_ALPHA * 0.5) * 1.0 # temporal neutral
|
||||
assert abs(sr.weight - expected) < 1e-6
|
||||
|
||||
expected = 0.6 * ce_norm + 0.2 * rrf_norm + 0.1 * temporal + 0.1 * recency
|
||||
assert expected == 1.0, f"All 1.0 should give 1.0, got {expected}"
|
||||
def test_min_recency_penalty(self):
|
||||
"""A memory from >365 days ago (recency=0.1) should penalise score."""
|
||||
old = NOW - timedelta(days=400)
|
||||
sr = _make_result(ce_norm=0.5, occurred_start=old)
|
||||
apply_combined_scoring([sr], now=NOW)
|
||||
expected = 0.5 * (1.0 + _RECENCY_ALPHA * (0.1 - 0.5)) * 1.0
|
||||
assert abs(sr.weight - expected) < 1e-6
|
||||
|
||||
# Test case 2: All components at 0.0
|
||||
ce_norm = 0.0
|
||||
rrf_norm = 0.0
|
||||
temporal = 0.0
|
||||
recency = 0.0
|
||||
def test_max_temporal_boost(self):
|
||||
"""temporal_proximity=1.0 should boost by (1 + alpha*0.5)."""
|
||||
sr = _make_result(ce_norm=0.5, temporal_proximity=1.0)
|
||||
apply_combined_scoring([sr], now=NOW)
|
||||
expected = 0.5 * 1.0 * (1.0 + _TEMPORAL_ALPHA * 0.5) # recency neutral
|
||||
assert abs(sr.weight - expected) < 1e-6
|
||||
|
||||
expected = 0.6 * ce_norm + 0.2 * rrf_norm + 0.1 * temporal + 0.1 * recency
|
||||
assert expected == 0.0, f"All 0.0 should give 0.0, got {expected}"
|
||||
def test_temporal_none_is_neutral(self):
|
||||
"""temporal_proximity=None must be treated as 0.5 (no boost/penalty)."""
|
||||
sr_none = _make_result(ce_norm=0.5, temporal_proximity=None)
|
||||
sr_half = _make_result(ce_norm=0.5, temporal_proximity=0.5)
|
||||
apply_combined_scoring([sr_none], now=NOW)
|
||||
apply_combined_scoring([sr_half], now=NOW)
|
||||
assert abs(sr_none.weight - sr_half.weight) < 1e-9
|
||||
|
||||
# Test case 3: High CE, low RRF (cross-encoder finds something retrieval missed)
|
||||
ce_norm = 0.999
|
||||
rrf_norm = 0.0 # Lowest in set
|
||||
temporal = 0.5
|
||||
recency = 0.5
|
||||
def test_both_signals_combined(self):
|
||||
"""Both boosts are applied multiplicatively."""
|
||||
sr = _make_result(ce_norm=0.5, occurred_start=NOW, temporal_proximity=1.0)
|
||||
apply_combined_scoring([sr], now=NOW)
|
||||
recency_boost = 1.0 + _RECENCY_ALPHA * (1.0 - 0.5)
|
||||
temporal_boost = 1.0 + _TEMPORAL_ALPHA * (1.0 - 0.5)
|
||||
expected = 0.5 * recency_boost * temporal_boost
|
||||
assert abs(sr.weight - expected) < 1e-6
|
||||
|
||||
expected = 0.6 * ce_norm + 0.2 * rrf_norm + 0.1 * temporal + 0.1 * recency
|
||||
# 0.5994 + 0.0 + 0.05 + 0.05 = 0.6994
|
||||
assert abs(expected - 0.6994) < 0.001, f"Expected ~0.6994, got {expected}"
|
||||
def test_boost_is_proportional_to_ce(self):
|
||||
"""The absolute boost from recency scales with the CE score."""
|
||||
sr_high = _make_result(ce_norm=0.9, occurred_start=NOW)
|
||||
sr_low = _make_result(ce_norm=0.3, occurred_start=NOW)
|
||||
apply_combined_scoring([sr_high, sr_low], now=NOW)
|
||||
|
||||
# Test case 4: Medium CE, high RRF (retrieval consensus)
|
||||
ce_norm = 0.8
|
||||
rrf_norm = 1.0 # Highest in set
|
||||
temporal = 0.5
|
||||
recency = 0.5
|
||||
# Both get the same recency boost factor — absolute gain is proportional to CE
|
||||
boost_factor = 1.0 + _RECENCY_ALPHA * 0.5
|
||||
assert abs(sr_high.weight - 0.9 * boost_factor) < 1e-6
|
||||
assert abs(sr_low.weight - 0.3 * boost_factor) < 1e-6
|
||||
|
||||
expected = 0.6 * ce_norm + 0.2 * rrf_norm + 0.1 * temporal + 0.1 * recency
|
||||
# 0.48 + 0.2 + 0.05 + 0.05 = 0.78
|
||||
assert abs(expected - 0.78) < 0.001, f"Expected ~0.78, got {expected}"
|
||||
def test_boost_capped(self):
|
||||
"""Max boost: recency=1.0 + temporal=1.0 gives ≤21% uplift on CE."""
|
||||
sr = _make_result(ce_norm=1.0, occurred_start=NOW, temporal_proximity=1.0)
|
||||
apply_combined_scoring([sr], now=NOW)
|
||||
assert sr.weight <= 1.0 * (1 + _RECENCY_ALPHA / 2) * (1 + _TEMPORAL_ALPHA / 2) + 1e-9
|
||||
|
||||
def test_rrf_contribution_is_significant(self):
|
||||
"""Verify RRF actually contributes to the final score (not negligible)."""
|
||||
# Same CE, different RRF
|
||||
ce_norm = 0.8
|
||||
temporal = 0.5
|
||||
recency = 0.5
|
||||
def test_rrf_normalized_always_zero(self):
|
||||
"""RRF is excluded from scoring; rrf_normalized is set to 0.0 for trace clarity."""
|
||||
sr = _make_result(ce_norm=0.5)
|
||||
apply_combined_scoring([sr], now=NOW)
|
||||
assert sr.rrf_normalized == 0.0
|
||||
|
||||
# Low RRF
|
||||
score_low_rrf = 0.6 * ce_norm + 0.2 * 0.0 + 0.1 * temporal + 0.1 * recency
|
||||
def test_combined_score_equals_weight(self):
|
||||
"""combined_score and weight must stay in sync."""
|
||||
sr = _make_result(ce_norm=0.7, occurred_start=NOW, temporal_proximity=0.8)
|
||||
apply_combined_scoring([sr], now=NOW)
|
||||
assert sr.combined_score == sr.weight
|
||||
|
||||
# High RRF
|
||||
score_high_rrf = 0.6 * ce_norm + 0.2 * 1.0 + 0.1 * temporal + 0.1 * recency
|
||||
def test_model_calibration_independence(self):
|
||||
"""
|
||||
A low-calibration model (low CE scores) and a high-calibration model
|
||||
(high CE scores) should produce the same ranking for identical content.
|
||||
|
||||
# Difference should be 0.2 (20% contribution)
|
||||
diff = score_high_rrf - score_low_rrf
|
||||
assert abs(diff - 0.2) < 0.001, f"RRF should contribute 0.2 difference, got {diff}"
|
||||
With additive scoring the recency term would dominate for low-CE models;
|
||||
with multiplicative boosting the relative ranking is stable.
|
||||
"""
|
||||
recent = NOW - timedelta(days=10)
|
||||
old = NOW - timedelta(days=300)
|
||||
|
||||
# High-calibration model: clear winner is #1 (more relevant, slightly older)
|
||||
h_relevant = _make_result(ce_norm=0.85, occurred_start=old)
|
||||
h_recent = _make_result(ce_norm=0.60, occurred_start=recent)
|
||||
apply_combined_scoring([h_relevant, h_recent], now=NOW)
|
||||
assert h_relevant.weight > h_recent.weight, "High-CE model: relevance should win"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_trace_has_normalized_rrf(memory, request_context):
|
||||
"""Integration test: verify trace contains normalized RRF values, not raw."""
|
||||
bank_id = f"test_scoring_{datetime.now(timezone.utc).timestamp()}"
|
||||
# Low-calibration model: same relative difference, just compressed scores
|
||||
l_relevant = _make_result(ce_norm=0.34, occurred_start=old)
|
||||
l_recent = _make_result(ce_norm=0.24, occurred_start=recent)
|
||||
apply_combined_scoring([l_relevant, l_recent], now=NOW)
|
||||
assert l_relevant.weight > l_recent.weight, "Low-CE model: relevance should still win"
|
||||
|
||||
try:
|
||||
# Store multiple memories to ensure different RRF scores
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="Python is a programming language created by Guido van Rossum",
|
||||
context="tech facts",
|
||||
request_context=request_context,
|
||||
)
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="JavaScript was created by Brendan Eich at Netscape",
|
||||
context="tech facts",
|
||||
request_context=request_context,
|
||||
)
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="The Eiffel Tower is located in Paris, France",
|
||||
context="geography facts",
|
||||
request_context=request_context,
|
||||
)
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="Mount Everest is the tallest mountain on Earth",
|
||||
context="geography facts",
|
||||
request_context=request_context,
|
||||
)
|
||||
def test_no_occurred_start_defaults_recency_neutral(self):
|
||||
"""Missing occurred_start → recency=0.5 → no boost/penalty."""
|
||||
sr = _make_result(ce_norm=0.5, occurred_start=None)
|
||||
apply_combined_scoring([sr], now=NOW)
|
||||
assert sr.recency == 0.5
|
||||
assert abs(sr.weight - 0.5) < 1e-9
|
||||
|
||||
# Search with tracing
|
||||
result = await memory.recall_async(
|
||||
bank_id=bank_id,
|
||||
query="programming languages",
|
||||
fact_type=["world"],
|
||||
budget=Budget.LOW,
|
||||
max_tokens=1024,
|
||||
enable_trace=True,
|
||||
request_context=request_context,
|
||||
)
|
||||
def test_timezone_naive_occurred_start_handled(self):
|
||||
"""Naive datetimes in occurred_start should not raise."""
|
||||
naive_date = datetime(2024, 1, 1) # no tzinfo
|
||||
sr = _make_result(ce_norm=0.5, occurred_start=naive_date)
|
||||
apply_combined_scoring([sr], now=NOW) # must not raise
|
||||
assert 0.0 < sr.weight < 1.0
|
||||
|
||||
assert result.trace is not None, "Trace should be present"
|
||||
trace = result.trace
|
||||
def test_custom_alpha_values(self):
|
||||
"""Custom alpha parameters are respected."""
|
||||
sr = _make_result(ce_norm=0.5, occurred_start=NOW)
|
||||
apply_combined_scoring([sr], now=NOW, recency_alpha=0.4, temporal_alpha=0.0)
|
||||
expected = 0.5 * (1.0 + 0.4 * 0.5) * 1.0
|
||||
assert abs(sr.weight - expected) < 1e-6
|
||||
|
||||
# Check reranked results have proper score_components
|
||||
assert "reranked" in trace, "Trace should have reranked results"
|
||||
assert len(trace["reranked"]) > 0, "Should have reranked results"
|
||||
def test_future_event_recency_capped_at_one(self):
|
||||
"""Events in the future must not produce recency > 1.0, keeping boost within bounds."""
|
||||
future = NOW + timedelta(days=180)
|
||||
sr = _make_result(ce_norm=0.5, occurred_start=future)
|
||||
apply_combined_scoring([sr], now=NOW)
|
||||
assert sr.recency == 1.0
|
||||
expected_max_boost = 1.0 + _RECENCY_ALPHA * 0.5
|
||||
assert sr.weight <= 0.5 * expected_max_boost + 1e-9
|
||||
|
||||
has_valid_rrf = False
|
||||
has_valid_temporal = False
|
||||
has_valid_recency = False
|
||||
|
||||
for r in trace["reranked"]:
|
||||
sc = r.get("score_components", {})
|
||||
|
||||
# Check RRF normalized is present and in valid range
|
||||
if "rrf_normalized" in sc:
|
||||
rrf_norm = sc["rrf_normalized"]
|
||||
assert 0.0 <= rrf_norm <= 1.0, f"rrf_normalized {rrf_norm} should be in [0, 1]"
|
||||
# Should NOT be raw RRF score (which would be ~0.04-0.06)
|
||||
# A normalized value of exactly 0.0 or 1.0 is valid (min/max of set)
|
||||
# But raw scores like 0.0607 should never appear as normalized
|
||||
if rrf_norm > 0.1: # Any value > 0.1 is likely properly normalized
|
||||
has_valid_rrf = True
|
||||
|
||||
# Check temporal is present and in valid range
|
||||
if "temporal" in sc:
|
||||
temporal = sc["temporal"]
|
||||
assert 0.0 <= temporal <= 1.0, f"temporal {temporal} should be in [0, 1]"
|
||||
has_valid_temporal = True
|
||||
|
||||
# Check recency is present and in valid range
|
||||
if "recency" in sc:
|
||||
recency = sc["recency"]
|
||||
assert 0.0 <= recency <= 1.0, f"recency {recency} should be in [0, 1]"
|
||||
has_valid_recency = True
|
||||
|
||||
# At least some results should have these components
|
||||
# (might not have rrf > 0.1 if all scores are same, which is fine)
|
||||
assert has_valid_temporal, "Should have temporal scores in trace"
|
||||
assert has_valid_recency, "Should have recency scores in trace"
|
||||
|
||||
print("\n✓ Combined scoring trace test passed!")
|
||||
print(f" - Reranked results: {len(trace['reranked'])}")
|
||||
if trace["reranked"]:
|
||||
sc = trace["reranked"][0].get("score_components", {})
|
||||
print(f" - First result score components: {sc}")
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rrf_normalized_not_raw_in_trace(memory, request_context):
|
||||
"""Verify that raw RRF scores (0.04-0.06 range) don't appear as normalized values."""
|
||||
bank_id = f"test_rrf_raw_{datetime.now(timezone.utc).timestamp()}"
|
||||
|
||||
try:
|
||||
# Store enough memories to get varied RRF scores
|
||||
for i in range(5):
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content=f"Test fact number {i} about various topics",
|
||||
context="test context",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
result = await memory.recall_async(
|
||||
bank_id=bank_id,
|
||||
query="test fact",
|
||||
fact_type=["world"],
|
||||
budget=Budget.LOW,
|
||||
max_tokens=512,
|
||||
enable_trace=True,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
trace = result.trace
|
||||
assert trace is not None
|
||||
|
||||
# Check that rrf_normalized values are NOT in the raw range
|
||||
raw_rrf_range = (0.01, 0.08) # Raw RRF scores are typically in this range
|
||||
|
||||
for r in trace.get("reranked", []):
|
||||
sc = r.get("score_components", {})
|
||||
|
||||
if "rrf_normalized" in sc and "rrf_score" in sc:
|
||||
rrf_norm = sc["rrf_normalized"]
|
||||
rrf_raw = sc["rrf_score"]
|
||||
|
||||
# Raw should be in the typical range
|
||||
assert raw_rrf_range[0] <= rrf_raw <= raw_rrf_range[1], \
|
||||
f"Raw RRF {rrf_raw} should be in typical range {raw_rrf_range}"
|
||||
|
||||
# Normalized should either be:
|
||||
# - 0.0 (min in set)
|
||||
# - 1.0 (max in set)
|
||||
# - 0.5 (all same)
|
||||
# - Something in between (0.0 to 1.0)
|
||||
# But NOT the same as raw (which would indicate no normalization)
|
||||
if len(trace["reranked"]) > 1:
|
||||
# If we have multiple results, normalized should differ from raw
|
||||
# (unless by coincidence, which is very unlikely)
|
||||
assert rrf_norm != rrf_raw, \
|
||||
f"Normalized RRF ({rrf_norm}) should differ from raw ({rrf_raw})"
|
||||
|
||||
print("\n✓ RRF raw vs normalized test passed!")
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_combined_score_matches_components(memory, request_context):
|
||||
"""Verify the final score actually equals the weighted sum of components."""
|
||||
bank_id = f"test_combined_{datetime.now(timezone.utc).timestamp()}"
|
||||
|
||||
try:
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="The quick brown fox jumps over the lazy dog",
|
||||
context="test",
|
||||
request_context=request_context,
|
||||
)
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="A quick test of the emergency broadcast system",
|
||||
context="test",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
result = await memory.recall_async(
|
||||
bank_id=bank_id,
|
||||
query="quick test",
|
||||
fact_type=["world"],
|
||||
budget=Budget.LOW,
|
||||
max_tokens=512,
|
||||
enable_trace=True,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
trace = result.trace
|
||||
assert trace is not None
|
||||
|
||||
for r in trace.get("reranked", []):
|
||||
sc = r.get("score_components", {})
|
||||
final_score = r.get("rerank_score", 0)
|
||||
|
||||
# Get components (use defaults if missing)
|
||||
ce = sc.get("cross_encoder_score_normalized", 0)
|
||||
rrf = sc.get("rrf_normalized", 0.5)
|
||||
tmp = sc.get("temporal", 0.5)
|
||||
rec = sc.get("recency", 0.5)
|
||||
|
||||
# Calculate expected score
|
||||
expected = 0.6 * ce + 0.2 * rrf + 0.1 * tmp + 0.1 * rec
|
||||
|
||||
# Allow small floating point difference
|
||||
assert abs(final_score - expected) < 0.01, \
|
||||
f"Final score {final_score} doesn't match expected {expected} from components"
|
||||
|
||||
print("\n✓ Combined score verification test passed!")
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
def test_empty_list_is_noop(self):
|
||||
apply_combined_scoring([], now=NOW) # must not raise
|
||||
|
||||
@@ -5,11 +5,15 @@ Note: Consolidation runs automatically after retain via SyncTaskBackend in tests
|
||||
"""
|
||||
|
||||
import uuid
|
||||
from datetime import datetime, timezone
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
from hindsight_api.engine.consolidation.consolidator import run_consolidation_job
|
||||
from hindsight_api.engine.consolidation.consolidator import (
|
||||
_aggregate_source_fields,
|
||||
run_consolidation_job,
|
||||
)
|
||||
from hindsight_api.engine.memory_engine import MemoryEngine
|
||||
from hindsight_api.engine.reflect.tools import (
|
||||
tool_recall,
|
||||
@@ -2084,3 +2088,333 @@ async def test_consolidation_with_observations_mission(memory: "MemoryEngine", r
|
||||
else:
|
||||
os.environ["HINDSIGHT_API_OBSERVATIONS_MISSION"] = original
|
||||
clear_config_cache()
|
||||
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_observation_scopes_explicit_multi_pass(memory: MemoryEngine, request_context):
|
||||
"""Test that observation_scopes with an explicit list triggers separate consolidation passes.
|
||||
|
||||
A single memory stored with observation_scopes=[["user:alice"], ["teacher:ben"]]
|
||||
must produce:
|
||||
- At least one observation with tags containing ONLY "user:alice" (not "teacher:ben")
|
||||
- At least one observation with tags containing ONLY "teacher:ben" (not "user:alice")
|
||||
|
||||
The two tag scopes must remain isolated — no observation should carry both tags,
|
||||
which would indicate the scopes were incorrectly merged.
|
||||
"""
|
||||
bank_id = f"test-obs-scopes-explicit-{uuid.uuid4().hex[:8]}"
|
||||
|
||||
await memory.get_bank_profile(bank_id=bank_id, request_context=request_context)
|
||||
|
||||
# Retain a memory with two explicit observation scopes
|
||||
await memory.retain_batch_async(
|
||||
bank_id=bank_id,
|
||||
contents=[
|
||||
{
|
||||
"content": "Alice, a student, worked hard in the lesson with teacher Ben.",
|
||||
"observation_scopes": [["user:alice"], ["teacher:ben"]],
|
||||
}
|
||||
],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
async with memory._pool.acquire() as conn:
|
||||
observations = await conn.fetch(
|
||||
"""
|
||||
SELECT id, text, tags
|
||||
FROM memory_units
|
||||
WHERE bank_id = $1 AND fact_type = 'observation'
|
||||
ORDER BY created_at
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
|
||||
try:
|
||||
# Must have at least 2 observations (one per tag scope)
|
||||
assert len(observations) >= 2, (
|
||||
f"Expected at least 2 observations (one per tag scope), got {len(observations)}: "
|
||||
+ str([dict(o) for o in observations])
|
||||
)
|
||||
|
||||
tag_sets = [set(obs["tags"] or []) for obs in observations]
|
||||
|
||||
# There must be at least one observation scoped to user:alice only
|
||||
alice_only = [ts for ts in tag_sets if "user:alice" in ts and "teacher:ben" not in ts]
|
||||
assert alice_only, (
|
||||
f"Expected an observation scoped to 'user:alice' only, got tag sets: {tag_sets}"
|
||||
)
|
||||
|
||||
# There must be at least one observation scoped to teacher:ben only
|
||||
ben_only = [ts for ts in tag_sets if "teacher:ben" in ts and "user:alice" not in ts]
|
||||
assert ben_only, (
|
||||
f"Expected an observation scoped to 'teacher:ben' only, got tag sets: {tag_sets}"
|
||||
)
|
||||
|
||||
# No observation should carry both tags (scopes must not be merged)
|
||||
both = [ts for ts in tag_sets if "user:alice" in ts and "teacher:ben" in ts]
|
||||
assert not both, (
|
||||
f"Found observation(s) with both tags — scopes were incorrectly merged: {both}"
|
||||
)
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_observation_scopes_per_tag(memory: MemoryEngine, request_context):
|
||||
"""Test that observation_scopes='per_tag' derives one pass per individual tag.
|
||||
|
||||
A memory with tags=["user:alice", "teacher:ben"] and observation_scopes="per_tag"
|
||||
must produce isolated observations — one scoped to "user:alice" and one to "teacher:ben".
|
||||
"""
|
||||
bank_id = f"test-obs-scopes-pertag-{uuid.uuid4().hex[:8]}"
|
||||
|
||||
await memory.get_bank_profile(bank_id=bank_id, request_context=request_context)
|
||||
|
||||
await memory.retain_batch_async(
|
||||
bank_id=bank_id,
|
||||
contents=[
|
||||
{
|
||||
"content": "Alice, a student, worked hard in the lesson with teacher Ben.",
|
||||
"tags": ["user:alice", "teacher:ben"],
|
||||
"observation_scopes": "per_tag",
|
||||
}
|
||||
],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
async with memory._pool.acquire() as conn:
|
||||
observations = await conn.fetch(
|
||||
"""
|
||||
SELECT id, text, tags
|
||||
FROM memory_units
|
||||
WHERE bank_id = $1 AND fact_type = 'observation'
|
||||
ORDER BY created_at
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
|
||||
try:
|
||||
assert len(observations) >= 2, (
|
||||
f"Expected at least 2 observations (one per tag), got {len(observations)}: "
|
||||
+ str([dict(o) for o in observations])
|
||||
)
|
||||
|
||||
tag_sets = [set(obs["tags"] or []) for obs in observations]
|
||||
|
||||
alice_only = [ts for ts in tag_sets if "user:alice" in ts and "teacher:ben" not in ts]
|
||||
assert alice_only, f"Expected an observation scoped to 'user:alice' only, got: {tag_sets}"
|
||||
|
||||
ben_only = [ts for ts in tag_sets if "teacher:ben" in ts and "user:alice" not in ts]
|
||||
assert ben_only, f"Expected an observation scoped to 'teacher:ben' only, got: {tag_sets}"
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_observation_scopes_combined(memory: MemoryEngine, request_context):
|
||||
"""Test that observation_scopes='combined' produces a single observation with all tags.
|
||||
|
||||
A memory with tags=["user:alice", "teacher:ben"] and observation_scopes="combined"
|
||||
must produce at least one observation that carries both tags together, and no
|
||||
observation scoped to only one of them.
|
||||
"""
|
||||
bank_id = f"test-obs-scopes-combined-{uuid.uuid4().hex[:8]}"
|
||||
|
||||
await memory.get_bank_profile(bank_id=bank_id, request_context=request_context)
|
||||
|
||||
await memory.retain_batch_async(
|
||||
bank_id=bank_id,
|
||||
contents=[
|
||||
{
|
||||
"content": "Alice, a student, worked hard in the lesson with teacher Ben.",
|
||||
"tags": ["user:alice", "teacher:ben"],
|
||||
"observation_scopes": "combined",
|
||||
}
|
||||
],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
async with memory._pool.acquire() as conn:
|
||||
observations = await conn.fetch(
|
||||
"""
|
||||
SELECT id, text, tags
|
||||
FROM memory_units
|
||||
WHERE bank_id = $1 AND fact_type = 'observation'
|
||||
ORDER BY created_at
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
|
||||
try:
|
||||
assert len(observations) >= 1, (
|
||||
"Expected at least 1 observation, got 0"
|
||||
)
|
||||
|
||||
tag_sets = [set(obs["tags"] or []) for obs in observations]
|
||||
|
||||
# All observations must carry both tags (combined scope)
|
||||
combined = [ts for ts in tag_sets if "user:alice" in ts and "teacher:ben" in ts]
|
||||
assert combined, f"Expected at least one observation with both tags, got: {tag_sets}"
|
||||
|
||||
# No observation should be scoped to only one tag
|
||||
alice_only = [ts for ts in tag_sets if "user:alice" in ts and "teacher:ben" not in ts]
|
||||
assert not alice_only, f"Expected no alice-only observation in combined mode, got: {tag_sets}"
|
||||
|
||||
ben_only = [ts for ts in tag_sets if "teacher:ben" in ts and "user:alice" not in ts]
|
||||
assert not ben_only, f"Expected no ben-only observation in combined mode, got: {tag_sets}"
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_observation_scopes_all_combinations(memory: MemoryEngine, request_context):
|
||||
"""Test that observation_scopes='all_combinations' generates passes for every tag subset.
|
||||
|
||||
A memory with tags=["user:alice", "teacher:ben"] and observation_scopes="all_combinations"
|
||||
must produce observations covering all subsets: ["user:alice"], ["teacher:ben"], and
|
||||
["user:alice", "teacher:ben"].
|
||||
"""
|
||||
bank_id = f"test-obs-scopes-allcombos-{uuid.uuid4().hex[:8]}"
|
||||
|
||||
await memory.get_bank_profile(bank_id=bank_id, request_context=request_context)
|
||||
|
||||
await memory.retain_batch_async(
|
||||
bank_id=bank_id,
|
||||
contents=[
|
||||
{
|
||||
"content": "Alice, a student, worked hard in the lesson with teacher Ben.",
|
||||
"tags": ["user:alice", "teacher:ben"],
|
||||
"observation_scopes": "all_combinations",
|
||||
}
|
||||
],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
async with memory._pool.acquire() as conn:
|
||||
observations = await conn.fetch(
|
||||
"""
|
||||
SELECT id, text, tags
|
||||
FROM memory_units
|
||||
WHERE bank_id = $1 AND fact_type = 'observation'
|
||||
ORDER BY created_at
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
|
||||
try:
|
||||
# With 2 tags there are 3 subsets: {alice}, {ben}, {alice, ben}
|
||||
assert len(observations) >= 3, (
|
||||
f"Expected at least 3 observations (one per subset), got {len(observations)}: "
|
||||
+ str([dict(o) for o in observations])
|
||||
)
|
||||
|
||||
tag_sets = [set(obs["tags"] or []) for obs in observations]
|
||||
|
||||
alice_only = [ts for ts in tag_sets if "user:alice" in ts and "teacher:ben" not in ts]
|
||||
assert alice_only, f"Expected an observation scoped to 'user:alice' only, got: {tag_sets}"
|
||||
|
||||
ben_only = [ts for ts in tag_sets if "teacher:ben" in ts and "user:alice" not in ts]
|
||||
assert ben_only, f"Expected an observation scoped to 'teacher:ben' only, got: {tag_sets}"
|
||||
|
||||
combined = [ts for ts in tag_sets if "user:alice" in ts and "teacher:ben" in ts]
|
||||
assert combined, f"Expected an observation scoped to both tags, got: {tag_sets}"
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
def _dt(year: int, month: int, day: int) -> datetime:
|
||||
return datetime(year, month, day, tzinfo=timezone.utc)
|
||||
|
||||
|
||||
class TestAggregateSourceFields:
|
||||
"""Unit tests for _aggregate_source_fields – no database required."""
|
||||
|
||||
def test_all_none_temporal_fields_stay_none(self):
|
||||
"""When source memories carry no temporal data, all fields must remain None."""
|
||||
source_mems = [
|
||||
{"tags": ["t1"], "event_date": None, "occurred_start": None, "occurred_end": None, "mentioned_at": None},
|
||||
{"tags": ["t1"], "event_date": None, "occurred_start": None, "occurred_end": None, "mentioned_at": None},
|
||||
]
|
||||
agg = _aggregate_source_fields(source_mems)
|
||||
assert agg.event_date is None
|
||||
assert agg.occurred_start is None
|
||||
assert agg.occurred_end is None
|
||||
assert agg.mentioned_at is None
|
||||
|
||||
def test_temporal_fields_aggregated_correctly(self):
|
||||
"""occurred_start and event_date are minimised; occurred_end and mentioned_at are maximised."""
|
||||
early = _dt(2023, 1, 1)
|
||||
late = _dt(2024, 6, 15)
|
||||
source_mems = [
|
||||
{
|
||||
"tags": [],
|
||||
"event_date": late,
|
||||
"occurred_start": late,
|
||||
"occurred_end": early,
|
||||
"mentioned_at": early,
|
||||
},
|
||||
{
|
||||
"tags": [],
|
||||
"event_date": early,
|
||||
"occurred_start": early,
|
||||
"occurred_end": late,
|
||||
"mentioned_at": late,
|
||||
},
|
||||
]
|
||||
agg = _aggregate_source_fields(source_mems)
|
||||
assert agg.event_date == early
|
||||
assert agg.occurred_start == early
|
||||
assert agg.occurred_end == late
|
||||
assert agg.mentioned_at == late
|
||||
|
||||
def test_partial_temporal_fields_ignored_when_none(self):
|
||||
"""None values in individual sources do not corrupt the min/max from sources that do have dates."""
|
||||
d = _dt(2023, 3, 10)
|
||||
source_mems = [
|
||||
{"tags": [], "event_date": None, "occurred_start": None, "occurred_end": None, "mentioned_at": None},
|
||||
{"tags": [], "event_date": d, "occurred_start": d, "occurred_end": d, "mentioned_at": d},
|
||||
]
|
||||
agg = _aggregate_source_fields(source_mems)
|
||||
assert agg.event_date == d
|
||||
assert agg.occurred_start == d
|
||||
assert agg.occurred_end == d
|
||||
assert agg.mentioned_at == d
|
||||
|
||||
def test_tags_inherited_from_first_source_memory(self):
|
||||
"""Tags default to those of the first source memory (batch invariant)."""
|
||||
source_mems = [
|
||||
{"tags": ["user:alice"], "event_date": None, "occurred_start": None, "occurred_end": None, "mentioned_at": None},
|
||||
{"tags": ["user:alice"], "event_date": None, "occurred_start": None, "occurred_end": None, "mentioned_at": None},
|
||||
]
|
||||
agg = _aggregate_source_fields(source_mems)
|
||||
assert agg.tags == ["user:alice"]
|
||||
|
||||
def test_tags_override_takes_precedence(self):
|
||||
"""Explicit tags parameter overrides the source-memory tags."""
|
||||
source_mems = [
|
||||
{"tags": ["user:alice"], "event_date": None, "occurred_start": None, "occurred_end": None, "mentioned_at": None},
|
||||
]
|
||||
agg = _aggregate_source_fields(source_mems, tags=["scope:override"])
|
||||
assert agg.tags == ["scope:override"]
|
||||
|
||||
def test_empty_tags_override_is_respected(self):
|
||||
"""An explicit empty list override must not fall back to source tags."""
|
||||
source_mems = [
|
||||
{"tags": ["user:alice"], "event_date": None, "occurred_start": None, "occurred_end": None, "mentioned_at": None},
|
||||
]
|
||||
agg = _aggregate_source_fields(source_mems, tags=[])
|
||||
assert agg.tags == []
|
||||
|
||||
def test_single_source_memory(self):
|
||||
"""Single-source aggregation should just pass through that memory's fields."""
|
||||
d = _dt(2024, 11, 5)
|
||||
source_mems = [
|
||||
{"tags": ["x"], "event_date": d, "occurred_start": d, "occurred_end": d, "mentioned_at": d},
|
||||
]
|
||||
agg = _aggregate_source_fields(source_mems)
|
||||
assert agg.event_date == d
|
||||
assert agg.occurred_start == d
|
||||
assert agg.occurred_end == d
|
||||
assert agg.mentioned_at == d
|
||||
assert agg.tags == ["x"]
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,141 @@
|
||||
"""
|
||||
Unit tests for fact extraction retry logic.
|
||||
|
||||
Tests the fix for the TypeError when LLM returns invalid JSON across all retries.
|
||||
Previously, `raise last_error` would raise None (TypeError) because last_error was
|
||||
only set in the BadRequestError handler, not when the LLM returned non-dict JSON.
|
||||
"""
|
||||
|
||||
from datetime import datetime, timezone
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
def _make_config(llm_max_retries: int = 3, retain_llm_max_retries: int | None = None):
|
||||
"""Build a minimal HindsightConfig for fact extraction tests."""
|
||||
from hindsight_api.config import HindsightConfig
|
||||
|
||||
cfg = MagicMock(spec=HindsightConfig)
|
||||
cfg.retain_llm_max_retries = retain_llm_max_retries
|
||||
cfg.llm_max_retries = llm_max_retries
|
||||
cfg.retain_llm_initial_backoff = None
|
||||
cfg.llm_initial_backoff = 0.0
|
||||
cfg.retain_llm_max_backoff = None
|
||||
cfg.llm_max_backoff = 0.0
|
||||
cfg.retain_max_completion_tokens = 8192
|
||||
cfg.retain_extraction_mode = "concise"
|
||||
cfg.retain_extract_causal_links = False
|
||||
cfg.retain_mission = None
|
||||
return cfg
|
||||
|
||||
|
||||
def _make_llm_config(mock_response):
|
||||
"""Build a mock LLMProvider that returns the given response."""
|
||||
from hindsight_api.engine.llm_wrapper import LLMProvider
|
||||
|
||||
llm = MagicMock(spec=LLMProvider)
|
||||
llm.provider = "mock"
|
||||
token_usage = MagicMock()
|
||||
token_usage.__add__ = lambda self, other: self
|
||||
llm.call = AsyncMock(return_value=(mock_response, token_usage))
|
||||
return llm
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_non_dict_json_all_retries_returns_empty():
|
||||
"""
|
||||
When LLM returns non-dict JSON on every attempt, extraction should return []
|
||||
without raising TypeError ('exceptions must derive from BaseException').
|
||||
|
||||
This was the bug: the loop ran range(2) times (hardcoded), but comparisons
|
||||
used config.llm_max_retries (default 10). On the last loop iteration (attempt=1),
|
||||
`attempt < 10 - 1` was True, so the code called `continue`, the loop
|
||||
exhausted, and `raise last_error` raised None → TypeError.
|
||||
"""
|
||||
from hindsight_api.engine.retain.fact_extraction import _extract_facts_from_chunk
|
||||
|
||||
# llm_max_retries=3 ensures the bug triggers with the old code (3 != 2 hardcoded)
|
||||
config = _make_config(llm_max_retries=3, retain_llm_max_retries=None)
|
||||
|
||||
# Mock: always returns a list (non-dict), which is invalid
|
||||
llm_config = _make_llm_config(mock_response=[{"invalid": "response"}])
|
||||
|
||||
with patch(
|
||||
"hindsight_api.engine.retain.fact_extraction._build_extraction_prompt_and_schema",
|
||||
return_value=("system prompt", MagicMock()),
|
||||
):
|
||||
facts, usage = await _extract_facts_from_chunk(
|
||||
chunk="Alice visited Paris in 2023.",
|
||||
chunk_index=0,
|
||||
total_chunks=1,
|
||||
event_date=datetime(2023, 1, 1, tzinfo=timezone.utc),
|
||||
context="travel notes",
|
||||
llm_config=llm_config,
|
||||
config=config,
|
||||
agent_name="test-agent",
|
||||
)
|
||||
|
||||
assert facts == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_non_dict_json_with_default_max_retries_returns_empty():
|
||||
"""
|
||||
Same scenario with the default llm_max_retries=10 (matching real default config).
|
||||
The old code ran range(2) but checked against 10, always continuing until
|
||||
the loop exhausted, then raised None → TypeError.
|
||||
"""
|
||||
from hindsight_api.engine.retain.fact_extraction import _extract_facts_from_chunk
|
||||
|
||||
config = _make_config(llm_max_retries=10, retain_llm_max_retries=None)
|
||||
llm_config = _make_llm_config(mock_response="not a dict at all")
|
||||
|
||||
with patch(
|
||||
"hindsight_api.engine.retain.fact_extraction._build_extraction_prompt_and_schema",
|
||||
return_value=("system prompt", MagicMock()),
|
||||
):
|
||||
facts, usage = await _extract_facts_from_chunk(
|
||||
chunk="Some text.",
|
||||
chunk_index=0,
|
||||
total_chunks=1,
|
||||
event_date=datetime(2023, 6, 1, tzinfo=timezone.utc),
|
||||
context="",
|
||||
llm_config=llm_config,
|
||||
config=config,
|
||||
agent_name="agent",
|
||||
)
|
||||
|
||||
assert facts == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retain_llm_max_retries_overrides_global():
|
||||
"""
|
||||
When retain_llm_max_retries is set, it should be used for the loop range
|
||||
and all comparisons (no shadowing bug).
|
||||
"""
|
||||
from hindsight_api.engine.retain.fact_extraction import _extract_facts_from_chunk
|
||||
|
||||
# retain_llm_max_retries=5 should override llm_max_retries=10
|
||||
config = _make_config(llm_max_retries=10, retain_llm_max_retries=5)
|
||||
llm_config = _make_llm_config(mock_response=42) # non-dict: integer
|
||||
|
||||
with patch(
|
||||
"hindsight_api.engine.retain.fact_extraction._build_extraction_prompt_and_schema",
|
||||
return_value=("system prompt", MagicMock()),
|
||||
):
|
||||
facts, usage = await _extract_facts_from_chunk(
|
||||
chunk="Bob likes Python.",
|
||||
chunk_index=0,
|
||||
total_chunks=1,
|
||||
event_date=datetime(2024, 1, 1, tzinfo=timezone.utc),
|
||||
context="",
|
||||
llm_config=llm_config,
|
||||
config=config,
|
||||
agent_name="agent",
|
||||
)
|
||||
|
||||
assert facts == []
|
||||
# Verify it retried exactly retain_llm_max_retries times
|
||||
assert llm_config.call.call_count == 5
|
||||
@@ -0,0 +1,340 @@
|
||||
"""
|
||||
Tests for Gemini safety settings feature.
|
||||
|
||||
Verifies that:
|
||||
- Safety settings are read from env var and stored on GeminiLLM instances
|
||||
- Settings are applied to GenerateContentConfig in call() and call_with_tools()
|
||||
- The context variable override allows per-bank settings at request time
|
||||
- None (unset) means Gemini's default safety settings are used (no override)
|
||||
"""
|
||||
|
||||
import os
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
pytest.importorskip("google.genai")
|
||||
|
||||
|
||||
SAMPLE_SAFETY_SETTINGS = [
|
||||
{"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_NONE"},
|
||||
{"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE"},
|
||||
{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "BLOCK_NONE"},
|
||||
{"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_NONE"},
|
||||
]
|
||||
|
||||
|
||||
# ─── Config / env var parsing ─────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_gemini_safety_settings_parsed_from_env():
|
||||
"""Safety settings JSON from env var is parsed into HindsightConfig."""
|
||||
import json
|
||||
|
||||
from hindsight_api.config import ENV_LLM_GEMINI_SAFETY_SETTINGS, HindsightConfig, clear_config_cache
|
||||
|
||||
settings_json = json.dumps(SAMPLE_SAFETY_SETTINGS)
|
||||
with patch.dict(os.environ, {ENV_LLM_GEMINI_SAFETY_SETTINGS: settings_json}, clear=False):
|
||||
clear_config_cache()
|
||||
config = HindsightConfig.from_env()
|
||||
assert config.llm_gemini_safety_settings == SAMPLE_SAFETY_SETTINGS
|
||||
clear_config_cache()
|
||||
|
||||
|
||||
def test_gemini_safety_settings_default_is_none():
|
||||
"""When env var is not set, llm_gemini_safety_settings defaults to None."""
|
||||
from hindsight_api.config import ENV_LLM_GEMINI_SAFETY_SETTINGS, HindsightConfig, clear_config_cache
|
||||
|
||||
env = {k: v for k, v in os.environ.items() if k != ENV_LLM_GEMINI_SAFETY_SETTINGS}
|
||||
with patch.dict(os.environ, env, clear=True):
|
||||
clear_config_cache()
|
||||
config = HindsightConfig.from_env()
|
||||
assert config.llm_gemini_safety_settings is None
|
||||
clear_config_cache()
|
||||
|
||||
|
||||
def test_gemini_safety_settings_is_configurable_field():
|
||||
"""llm_gemini_safety_settings appears in configurable (per-bank) fields."""
|
||||
from hindsight_api.config import HindsightConfig
|
||||
|
||||
assert "llm_gemini_safety_settings" in HindsightConfig.get_configurable_fields()
|
||||
|
||||
|
||||
def test_gemini_safety_settings_not_in_credential_fields():
|
||||
"""llm_gemini_safety_settings is NOT a credential — it is safe to expose via API."""
|
||||
from hindsight_api.config import HindsightConfig
|
||||
|
||||
assert "llm_gemini_safety_settings" not in HindsightConfig.get_credential_fields()
|
||||
|
||||
|
||||
# ─── GeminiLLM instance ───────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def _make_gemini_provider(safety_settings=None):
|
||||
"""Return a GeminiLLM instance with a mocked genai.Client."""
|
||||
with patch("google.genai.Client") as mock_client_cls:
|
||||
mock_client_cls.return_value = MagicMock()
|
||||
from hindsight_api.engine.providers.gemini_llm import GeminiLLM
|
||||
|
||||
provider = GeminiLLM(
|
||||
provider="gemini",
|
||||
api_key="fake-api-key",
|
||||
base_url="",
|
||||
model="gemini-2.5-flash",
|
||||
gemini_safety_settings=safety_settings,
|
||||
)
|
||||
# Replace client with a fresh mock so we can inspect calls
|
||||
provider._client = MagicMock()
|
||||
return provider
|
||||
|
||||
|
||||
def test_gemini_llm_stores_safety_settings():
|
||||
"""GeminiLLM stores safety settings passed at construction."""
|
||||
provider = _make_gemini_provider(safety_settings=SAMPLE_SAFETY_SETTINGS)
|
||||
assert provider._safety_settings == SAMPLE_SAFETY_SETTINGS
|
||||
|
||||
|
||||
def test_gemini_llm_no_safety_settings_is_none():
|
||||
"""GeminiLLM._safety_settings is None when not provided."""
|
||||
provider = _make_gemini_provider(safety_settings=None)
|
||||
assert provider._safety_settings is None
|
||||
|
||||
|
||||
# ─── call() applies safety settings ──────────────────────────────────────────
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_applies_safety_settings():
|
||||
"""call() includes safety_settings in GenerateContentConfig when configured."""
|
||||
from google.genai import types as genai_types
|
||||
|
||||
provider = _make_gemini_provider(safety_settings=SAMPLE_SAFETY_SETTINGS)
|
||||
|
||||
# Build a fake successful response
|
||||
fake_response = MagicMock()
|
||||
fake_response.text = "hello"
|
||||
fake_response.candidates = [MagicMock(finish_reason="STOP")]
|
||||
fake_response.usage_metadata = MagicMock(prompt_token_count=5, candidates_token_count=2)
|
||||
|
||||
provider._client.aio.models.generate_content = AsyncMock(return_value=fake_response)
|
||||
|
||||
await provider.call(
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
scope="test",
|
||||
)
|
||||
|
||||
# Inspect the config passed to generate_content
|
||||
call_args = provider._client.aio.models.generate_content.call_args
|
||||
config_arg = call_args.kwargs.get("config") or call_args.args[0] if call_args.args else None
|
||||
# config may be in kwargs or positional; grab from kwargs
|
||||
config_arg = call_args.kwargs.get("config")
|
||||
|
||||
assert config_arg is not None, "GenerateContentConfig should have been passed"
|
||||
assert hasattr(config_arg, "safety_settings"), "Config should have safety_settings"
|
||||
assert config_arg.safety_settings is not None
|
||||
|
||||
categories = [s.category.value if hasattr(s.category, "value") else str(s.category) for s in config_arg.safety_settings]
|
||||
assert "HARM_CATEGORY_HARASSMENT" in categories
|
||||
assert "HARM_CATEGORY_HATE_SPEECH" in categories
|
||||
assert "HARM_CATEGORY_SEXUALLY_EXPLICIT" in categories
|
||||
assert "HARM_CATEGORY_DANGEROUS_CONTENT" in categories
|
||||
|
||||
thresholds = [s.threshold.value if hasattr(s.threshold, "value") else str(s.threshold) for s in config_arg.safety_settings]
|
||||
assert all(t == "BLOCK_NONE" for t in thresholds)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_no_safety_settings_omits_key():
|
||||
"""call() does NOT add safety_settings to GenerateContentConfig when none configured."""
|
||||
provider = _make_gemini_provider(safety_settings=None)
|
||||
|
||||
fake_response = MagicMock()
|
||||
fake_response.text = "hello"
|
||||
fake_response.candidates = [MagicMock(finish_reason="STOP")]
|
||||
fake_response.usage_metadata = MagicMock(prompt_token_count=5, candidates_token_count=2)
|
||||
|
||||
provider._client.aio.models.generate_content = AsyncMock(return_value=fake_response)
|
||||
|
||||
await provider.call(
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
scope="test",
|
||||
)
|
||||
|
||||
call_args = provider._client.aio.models.generate_content.call_args
|
||||
config_arg = call_args.kwargs.get("config")
|
||||
|
||||
# When no safety settings, config is either None or lacks safety_settings
|
||||
if config_arg is not None:
|
||||
assert not hasattr(config_arg, "safety_settings") or config_arg.safety_settings is None
|
||||
|
||||
|
||||
# ─── call_with_tools() applies safety settings ────────────────────────────────
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_with_tools_applies_safety_settings():
|
||||
"""call_with_tools() includes safety_settings in GenerateContentConfig."""
|
||||
provider = _make_gemini_provider(safety_settings=SAMPLE_SAFETY_SETTINGS)
|
||||
|
||||
# Build a fake tool-use response (no tool calls, just text)
|
||||
fake_part = MagicMock()
|
||||
fake_part.text = "answer"
|
||||
fake_part.function_call = None
|
||||
|
||||
fake_candidate = MagicMock()
|
||||
fake_candidate.content = MagicMock(parts=[fake_part])
|
||||
|
||||
fake_response = MagicMock()
|
||||
fake_response.candidates = [fake_candidate]
|
||||
fake_response.usage_metadata = MagicMock(prompt_token_count=5, candidates_token_count=3)
|
||||
|
||||
provider._client.aio.models.generate_content = AsyncMock(return_value=fake_response)
|
||||
|
||||
tools = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "test_tool",
|
||||
"description": "A test tool",
|
||||
"parameters": {"type": "object", "properties": {}, "required": []},
|
||||
},
|
||||
}
|
||||
]
|
||||
|
||||
await provider.call_with_tools(
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
tools=tools,
|
||||
scope="test",
|
||||
)
|
||||
|
||||
call_args = provider._client.aio.models.generate_content.call_args
|
||||
config_arg = call_args.kwargs.get("config")
|
||||
|
||||
assert config_arg is not None
|
||||
assert config_arg.safety_settings is not None
|
||||
categories = [s.category.value if hasattr(s.category, "value") else str(s.category) for s in config_arg.safety_settings]
|
||||
assert "HARM_CATEGORY_HARASSMENT" in categories
|
||||
|
||||
|
||||
# ─── with_config() override ───────────────────────────────────────────────────
|
||||
|
||||
|
||||
def _make_llm_provider(safety_settings=None):
|
||||
"""Return an LLMProvider (wrapping GeminiLLM) with a mocked genai.Client."""
|
||||
with patch("google.genai.Client") as mock_client_cls:
|
||||
mock_client_cls.return_value = MagicMock()
|
||||
from hindsight_api.engine.llm_wrapper import LLMProvider
|
||||
|
||||
provider = LLMProvider(
|
||||
provider="gemini",
|
||||
api_key="fake-api-key",
|
||||
base_url="",
|
||||
model="gemini-2.5-flash",
|
||||
gemini_safety_settings=safety_settings,
|
||||
)
|
||||
# Replace the underlying Gemini client with a fresh mock
|
||||
provider._provider_impl._client = MagicMock()
|
||||
return provider
|
||||
|
||||
|
||||
def _fake_response():
|
||||
r = MagicMock()
|
||||
r.text = "hello"
|
||||
r.candidates = [MagicMock(finish_reason="STOP")]
|
||||
r.usage_metadata = MagicMock(prompt_token_count=5, candidates_token_count=2)
|
||||
return r
|
||||
|
||||
|
||||
def _make_config(safety_settings):
|
||||
"""Return a minimal config-like object with llm_gemini_safety_settings."""
|
||||
cfg = MagicMock()
|
||||
cfg.llm_gemini_safety_settings = safety_settings
|
||||
return cfg
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_with_config_overrides_instance_settings():
|
||||
"""with_config() settings take precedence over the provider instance defaults."""
|
||||
instance_settings = [{"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_ONLY_HIGH"}]
|
||||
override_settings = [{"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE"}]
|
||||
|
||||
provider = _make_llm_provider(safety_settings=instance_settings)
|
||||
provider._provider_impl._client.aio.models.generate_content = AsyncMock(return_value=_fake_response())
|
||||
|
||||
configured = provider.with_config(_make_config(override_settings))
|
||||
await configured.call(messages=[{"role": "user", "content": "hi"}], scope="test")
|
||||
|
||||
config_arg = provider._provider_impl._client.aio.models.generate_content.call_args.kwargs.get("config")
|
||||
assert config_arg is not None
|
||||
categories = [s.category.value if hasattr(s.category, "value") else str(s.category) for s in config_arg.safety_settings]
|
||||
# Should use override_settings (HATE_SPEECH), not instance_settings (HARASSMENT)
|
||||
assert "HARM_CATEGORY_HATE_SPEECH" in categories
|
||||
assert "HARM_CATEGORY_HARASSMENT" not in categories
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_with_config_none_falls_back_to_instance():
|
||||
"""When with_config() supplies None, the instance default is used."""
|
||||
instance_settings = [{"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_NONE"}]
|
||||
|
||||
provider = _make_llm_provider(safety_settings=instance_settings)
|
||||
provider._provider_impl._client.aio.models.generate_content = AsyncMock(return_value=_fake_response())
|
||||
|
||||
configured = provider.with_config(_make_config(None))
|
||||
await configured.call(messages=[{"role": "user", "content": "hi"}], scope="test")
|
||||
|
||||
config_arg = provider._provider_impl._client.aio.models.generate_content.call_args.kwargs.get("config")
|
||||
assert config_arg is not None
|
||||
categories = [s.category.value if hasattr(s.category, "value") else str(s.category) for s in config_arg.safety_settings]
|
||||
assert "HARM_CATEGORY_HARASSMENT" in categories
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_with_config_resets_after_call():
|
||||
"""The ContextVar is properly reset after a with_config() call (no leakage)."""
|
||||
from hindsight_api.engine.providers.gemini_llm import _safety_settings_ctx
|
||||
|
||||
settings = [{"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_NONE"}]
|
||||
provider = _make_llm_provider(safety_settings=None)
|
||||
provider._provider_impl._client.aio.models.generate_content = AsyncMock(return_value=_fake_response())
|
||||
|
||||
before = _safety_settings_ctx.get()
|
||||
configured = provider.with_config(_make_config(settings))
|
||||
await configured.call(messages=[{"role": "user", "content": "hi"}], scope="test")
|
||||
after = _safety_settings_ctx.get()
|
||||
|
||||
assert after == before # ContextVar restored to its original value
|
||||
|
||||
|
||||
# ─── LLMProvider reads safety settings from config ────────────────────────────
|
||||
|
||||
|
||||
def test_llm_provider_reads_safety_settings_from_config():
|
||||
"""LLMProvider reads llm_gemini_safety_settings from global config for Gemini provider."""
|
||||
import json
|
||||
|
||||
from hindsight_api.config import ENV_LLM_GEMINI_SAFETY_SETTINGS, clear_config_cache
|
||||
|
||||
settings_json = json.dumps(SAMPLE_SAFETY_SETTINGS)
|
||||
env_overrides = {
|
||||
"HINDSIGHT_API_LLM_PROVIDER": "gemini",
|
||||
"HINDSIGHT_API_LLM_API_KEY": "fake-key",
|
||||
ENV_LLM_GEMINI_SAFETY_SETTINGS: settings_json,
|
||||
}
|
||||
|
||||
with patch.dict(os.environ, env_overrides, clear=False):
|
||||
clear_config_cache()
|
||||
with patch("google.genai.Client") as mock_client_cls:
|
||||
mock_client_cls.return_value = MagicMock()
|
||||
from hindsight_api.engine.llm_wrapper import LLMProvider
|
||||
|
||||
provider = LLMProvider(
|
||||
provider="gemini",
|
||||
api_key="fake-key",
|
||||
base_url="",
|
||||
model="gemini-2.5-flash",
|
||||
)
|
||||
|
||||
assert provider.gemini_safety_settings == SAMPLE_SAFETY_SETTINGS
|
||||
|
||||
clear_config_cache()
|
||||
@@ -81,8 +81,12 @@ async def test_hierarchical_fields_categorization():
|
||||
assert "disposition_literalism" in configurable
|
||||
assert "disposition_empathy" in configurable
|
||||
|
||||
# Verify entity labels fields are included
|
||||
assert "entities_allow_free_form" in configurable
|
||||
assert "entity_labels" in configurable
|
||||
|
||||
# Verify count is correct
|
||||
assert len(configurable) == 11
|
||||
assert len(configurable) == 14
|
||||
|
||||
# Verify credential fields (NEVER exposed)
|
||||
assert "llm_api_key" in credentials
|
||||
|
||||
@@ -0,0 +1,174 @@
|
||||
"""
|
||||
Tests for list_documents pagination and tags filtering.
|
||||
"""
|
||||
from datetime import datetime, timezone
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
async def _retain_doc(memory, bank_id, document_id, tags, request_context):
|
||||
"""Helper to retain a document with given tags. Uses gibberish content to avoid LLM
|
||||
fact extraction (documents are persisted even with zero facts)."""
|
||||
await memory.retain_batch_async(
|
||||
bank_id=bank_id,
|
||||
contents=[{"content": f"xyzabc123 !@# $$$ {document_id}"}],
|
||||
document_id=document_id,
|
||||
document_tags=tags or None,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_documents_offset_pagination(memory, request_context):
|
||||
"""offset parameter returns the correct slice of documents."""
|
||||
bank_id = f"test_list_docs_offset_{datetime.now(timezone.utc).timestamp()}"
|
||||
|
||||
try:
|
||||
for i in range(4):
|
||||
await _retain_doc(memory, bank_id, f"doc-{i:02d}", [], request_context)
|
||||
|
||||
# All documents, ordered by created_at DESC → doc-03, doc-02, doc-01, doc-00
|
||||
all_docs = await memory.list_documents(
|
||||
bank_id=bank_id, limit=10, offset=0, request_context=request_context
|
||||
)
|
||||
assert all_docs["total"] == 4
|
||||
assert len(all_docs["items"]) == 4
|
||||
all_ids = [d["id"] for d in all_docs["items"]]
|
||||
|
||||
# offset=2 should skip the first two and return the remaining two
|
||||
page2 = await memory.list_documents(
|
||||
bank_id=bank_id, limit=10, offset=2, request_context=request_context
|
||||
)
|
||||
assert page2["total"] == 4 # total is always the full count
|
||||
assert len(page2["items"]) == 2
|
||||
assert [d["id"] for d in page2["items"]] == all_ids[2:]
|
||||
|
||||
# offset beyond total returns empty items but correct total
|
||||
beyond = await memory.list_documents(
|
||||
bank_id=bank_id, limit=10, offset=10, request_context=request_context
|
||||
)
|
||||
assert beyond["total"] == 4
|
||||
assert beyond["items"] == []
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_documents_tags_filter_any_strict(memory, request_context):
|
||||
"""tags filter with any_strict returns only tagged documents that match."""
|
||||
bank_id = f"test_list_docs_tags_{datetime.now(timezone.utc).timestamp()}"
|
||||
|
||||
try:
|
||||
await _retain_doc(memory, bank_id, "doc-alpha", ["team-a"], request_context)
|
||||
await _retain_doc(memory, bank_id, "doc-beta", ["team-b"], request_context)
|
||||
await _retain_doc(memory, bank_id, "doc-both", ["team-a", "team-b"], request_context)
|
||||
await _retain_doc(memory, bank_id, "doc-untagged", [], request_context)
|
||||
|
||||
# any_strict: only docs with at least one of the given tags, untagged excluded
|
||||
result = await memory.list_documents(
|
||||
bank_id=bank_id,
|
||||
tags=["team-a"],
|
||||
tags_match="any_strict",
|
||||
request_context=request_context,
|
||||
)
|
||||
ids = {d["id"] for d in result["items"]}
|
||||
assert ids == {"doc-alpha", "doc-both"}
|
||||
assert result["total"] == 2
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_documents_tags_filter_any_includes_untagged(memory, request_context):
|
||||
"""tags filter with 'any' mode includes untagged documents."""
|
||||
bank_id = f"test_list_docs_tags_any_{datetime.now(timezone.utc).timestamp()}"
|
||||
|
||||
try:
|
||||
await _retain_doc(memory, bank_id, "doc-tagged", ["team-a"], request_context)
|
||||
await _retain_doc(memory, bank_id, "doc-other", ["team-b"], request_context)
|
||||
await _retain_doc(memory, bank_id, "doc-untagged", [], request_context)
|
||||
|
||||
result = await memory.list_documents(
|
||||
bank_id=bank_id,
|
||||
tags=["team-a"],
|
||||
tags_match="any",
|
||||
request_context=request_context,
|
||||
)
|
||||
ids = {d["id"] for d in result["items"]}
|
||||
# "any" includes untagged + matching tagged
|
||||
assert "doc-tagged" in ids
|
||||
assert "doc-untagged" in ids
|
||||
assert "doc-other" not in ids
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_documents_tags_filter_all_strict(memory, request_context):
|
||||
"""tags filter with all_strict returns only docs that have ALL the specified tags."""
|
||||
bank_id = f"test_list_docs_tags_all_{datetime.now(timezone.utc).timestamp()}"
|
||||
|
||||
try:
|
||||
await _retain_doc(memory, bank_id, "doc-a-only", ["team-a"], request_context)
|
||||
await _retain_doc(memory, bank_id, "doc-a-and-b", ["team-a", "team-b"], request_context)
|
||||
await _retain_doc(memory, bank_id, "doc-untagged", [], request_context)
|
||||
|
||||
result = await memory.list_documents(
|
||||
bank_id=bank_id,
|
||||
tags=["team-a", "team-b"],
|
||||
tags_match="all_strict",
|
||||
request_context=request_context,
|
||||
)
|
||||
ids = {d["id"] for d in result["items"]}
|
||||
assert ids == {"doc-a-and-b"}
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_documents_no_tags_filter_returns_all(memory, request_context):
|
||||
"""When no tags filter is specified, all documents are returned."""
|
||||
bank_id = f"test_list_docs_no_tags_{datetime.now(timezone.utc).timestamp()}"
|
||||
|
||||
try:
|
||||
await _retain_doc(memory, bank_id, "doc-tagged", ["team-a"], request_context)
|
||||
await _retain_doc(memory, bank_id, "doc-untagged", [], request_context)
|
||||
|
||||
result = await memory.list_documents(
|
||||
bank_id=bank_id,
|
||||
tags=None,
|
||||
request_context=request_context,
|
||||
)
|
||||
ids = {d["id"] for d in result["items"]}
|
||||
assert ids == {"doc-tagged", "doc-untagged"}
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_documents_tags_and_search_query_combined(memory, request_context):
|
||||
"""tags filter and q (search_query) can be combined."""
|
||||
bank_id = f"test_list_docs_tags_q_{datetime.now(timezone.utc).timestamp()}"
|
||||
|
||||
try:
|
||||
await _retain_doc(memory, bank_id, "report-2024", ["team-a"], request_context)
|
||||
await _retain_doc(memory, bank_id, "report-2025", ["team-b"], request_context)
|
||||
await _retain_doc(memory, bank_id, "summary-2024", ["team-a"], request_context)
|
||||
|
||||
result = await memory.list_documents(
|
||||
bank_id=bank_id,
|
||||
search_query="report",
|
||||
tags=["team-a"],
|
||||
tags_match="any_strict",
|
||||
request_context=request_context,
|
||||
)
|
||||
ids = {d["id"] for d in result["items"]}
|
||||
assert ids == {"report-2024"}
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
@@ -332,8 +332,8 @@ async def test_llm_provider_consolidation(memory_no_llm_verify, request_context,
|
||||
test_bank_id = f"llm_test_consolidation_{provider}_{model}_{datetime.now().timestamp()}"
|
||||
|
||||
# Enable observations for this bank
|
||||
from hindsight_api.config import get_config
|
||||
config = get_config()
|
||||
from hindsight_api.config import _get_raw_config
|
||||
config = _get_raw_config()
|
||||
original_value = config.enable_observations
|
||||
config.enable_observations = True
|
||||
|
||||
|
||||
@@ -46,16 +46,15 @@ async def test_mcp_tools_use_context_bank_id(mock_memory):
|
||||
assert "retain" in tools
|
||||
assert "recall" in tools
|
||||
|
||||
# Test retain with bank_id from context (use async_processing=False for synchronous test)
|
||||
token = _current_bank_id.set("context-bank-id")
|
||||
try:
|
||||
retain_tool = tools["retain"]
|
||||
result = await retain_tool.fn(content="test content", context="test_context", async_processing=False)
|
||||
assert "successfully" in result.lower()
|
||||
result = await retain_tool.fn(content="test content", context="test_context")
|
||||
assert result["status"] == "accepted"
|
||||
|
||||
# Verify the memory was called with the context bank_id
|
||||
mock_memory.retain_batch_async.assert_called_once()
|
||||
call_kwargs = mock_memory.retain_batch_async.call_args.kwargs
|
||||
mock_memory.submit_async_retain.assert_called_once()
|
||||
call_kwargs = mock_memory.submit_async_retain.call_args.kwargs
|
||||
assert call_kwargs["bank_id"] == "context-bank-id"
|
||||
finally:
|
||||
_current_bank_id.reset(token)
|
||||
@@ -133,12 +132,12 @@ async def test_mcp_tools_propagate_api_key(mock_memory):
|
||||
api_key_token = _current_api_key.set("test-bearer-token")
|
||||
try:
|
||||
retain_tool = tools["retain"]
|
||||
result = await retain_tool.fn(content="test content", context="test_context", async_processing=False)
|
||||
assert "successfully" in result.lower()
|
||||
result = await retain_tool.fn(content="test content", context="test_context")
|
||||
assert result["status"] == "accepted"
|
||||
|
||||
# Verify the memory was called with request_context containing api_key
|
||||
mock_memory.retain_batch_async.assert_called_once()
|
||||
call_kwargs = mock_memory.retain_batch_async.call_args.kwargs
|
||||
mock_memory.submit_async_retain.assert_called_once()
|
||||
call_kwargs = mock_memory.submit_async_retain.call_args.kwargs
|
||||
assert call_kwargs["request_context"].api_key == "test-bearer-token"
|
||||
finally:
|
||||
_current_bank_id.reset(bank_token)
|
||||
@@ -200,11 +199,11 @@ async def test_mcp_tools_propagate_tenant_id_and_api_key_id(mock_memory):
|
||||
key_id_token = _current_api_key_id.set("key-uuid-456")
|
||||
try:
|
||||
retain_tool = tools["retain"]
|
||||
await retain_tool.fn(content="test content", context="test_context", async_processing=False)
|
||||
await retain_tool.fn(content="test content", context="test_context")
|
||||
|
||||
# Verify the RequestContext passed to memory engine has all auth fields
|
||||
mock_memory.retain_batch_async.assert_called_once()
|
||||
request_context = mock_memory.retain_batch_async.call_args.kwargs["request_context"]
|
||||
mock_memory.submit_async_retain.assert_called_once()
|
||||
request_context = mock_memory.submit_async_retain.call_args.kwargs["request_context"]
|
||||
assert request_context.api_key == "hsk_test_key"
|
||||
assert request_context.tenant_id == "org-billing-123"
|
||||
assert request_context.api_key_id == "key-uuid-456"
|
||||
|
||||
@@ -15,6 +15,10 @@ logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.xfail(
|
||||
strict=False,
|
||||
reason="Gemini sometimes consistently translates Chinese content to English despite instructions",
|
||||
)
|
||||
async def test_retain_chinese_content(memory, request_context):
|
||||
"""
|
||||
Test that retain correctly extracts facts from Chinese content
|
||||
@@ -24,70 +28,87 @@ async def test_retain_chinese_content(memory, request_context):
|
||||
1. Facts are extracted from Chinese text
|
||||
2. The extracted facts contain Chinese characters
|
||||
3. Entity names are preserved in Chinese
|
||||
|
||||
Note: LLM fact extraction is non-deterministic and may sometimes translate
|
||||
content to English despite instructions. We retry up to 3 times.
|
||||
"""
|
||||
bank_id = f"test_chinese_retain_{datetime.now(timezone.utc).timestamp()}"
|
||||
max_retries = 3
|
||||
last_error = None
|
||||
|
||||
try:
|
||||
# Chinese content about a person and their activities
|
||||
chinese_content = """
|
||||
张伟是一位资深软件工程师,在腾讯工作了五年。他专门研究分布式系统,
|
||||
并领导了公司微服务架构的开发。他以编写干净、文档完善的代码而闻名。
|
||||
for attempt in range(max_retries):
|
||||
bank_id = f"test_chinese_retain_{datetime.now(timezone.utc).timestamp()}_{attempt}"
|
||||
|
||||
李明上个月加入团队担任初级开发人员。他正在学习React和Node.js。
|
||||
李明很有热情,在代码审查中提出很好的问题。他最近完成了他的第一个功能,
|
||||
这是一个用户认证流程。
|
||||
try:
|
||||
# Chinese content about a person and their activities
|
||||
chinese_content = """
|
||||
张伟是一位资深软件工程师,在腾讯工作了五年。他专门研究分布式系统,
|
||||
并领导了公司微服务架构的开发。他以编写干净、文档完善的代码而闻名。
|
||||
|
||||
团队使用Kubernetes进行容器编排,并部署到阿里云。他们遵循敏捷方法论,
|
||||
采用两周冲刺周期。合并前必须进行代码审查。
|
||||
"""
|
||||
李明上个月加入团队担任初级开发人员。他正在学习React和Node.js。
|
||||
李明很有热情,在代码审查中提出很好的问题。他最近完成了他的第一个功能,
|
||||
这是一个用户认证流程。
|
||||
|
||||
# Retain the Chinese content
|
||||
unit_ids = await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content=chinese_content,
|
||||
context="团队概述", # Chinese context
|
||||
event_date=datetime(2024, 1, 15, tzinfo=timezone.utc),
|
||||
request_context=request_context,
|
||||
)
|
||||
团队使用Kubernetes进行容器编排,并部署到阿里云。他们遵循敏捷方法论,
|
||||
采用两周冲刺周期。合并前必须进行代码审查。
|
||||
"""
|
||||
|
||||
logger.info(f"Retained {len(unit_ids)} facts from Chinese content")
|
||||
assert len(unit_ids) > 0, "Should have extracted and stored facts from Chinese content"
|
||||
# Retain the Chinese content
|
||||
unit_ids = await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content=chinese_content,
|
||||
context="团队概述", # Chinese context
|
||||
event_date=datetime(2024, 1, 15, tzinfo=timezone.utc),
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Recall the facts with a Chinese query
|
||||
result = await memory.recall_async(
|
||||
bank_id=bank_id,
|
||||
query="告诉我关于张伟的信息", # "Tell me about Zhang Wei"
|
||||
budget=Budget.MID,
|
||||
max_tokens=1000,
|
||||
fact_type=["world"],
|
||||
request_context=request_context,
|
||||
)
|
||||
logger.info(f"Retained {len(unit_ids)} facts from Chinese content (attempt {attempt + 1})")
|
||||
assert len(unit_ids) > 0, "Should have extracted and stored facts from Chinese content"
|
||||
|
||||
logger.info(f"Recalled {len(result.results)} facts")
|
||||
assert len(result.results) > 0, "Should recall facts about Zhang Wei"
|
||||
# Recall the facts with a Chinese query
|
||||
result = await memory.recall_async(
|
||||
bank_id=bank_id,
|
||||
query="告诉我关于张伟的信息", # "Tell me about Zhang Wei"
|
||||
budget=Budget.MID,
|
||||
max_tokens=1000,
|
||||
fact_type=["world"],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Verify that the facts contain Chinese characters
|
||||
# At least one fact should mention 张伟 (Zhang Wei) or related Chinese content
|
||||
chinese_facts_found = 0
|
||||
for fact in result.results:
|
||||
logger.info(f"Fact: {fact.text[:100]}...")
|
||||
# Check for common Chinese characters or the name
|
||||
if any(
|
||||
char in fact.text
|
||||
for char in ["张", "伟", "腾讯", "软件", "工程师", "分布式", "系统", "代码"]
|
||||
):
|
||||
chinese_facts_found += 1
|
||||
logger.info(f"Recalled {len(result.results)} facts")
|
||||
assert len(result.results) > 0, "Should recall facts about Zhang Wei"
|
||||
|
||||
logger.info(f"Found {chinese_facts_found} facts with Chinese content")
|
||||
assert chinese_facts_found > 0, (
|
||||
f"Expected facts to contain Chinese characters, but none found. "
|
||||
f"Facts: {[f.text for f in result.results]}"
|
||||
)
|
||||
# Verify that the facts contain Chinese characters
|
||||
# At least one fact should mention 张伟 (Zhang Wei) or related Chinese content
|
||||
chinese_facts_found = 0
|
||||
for fact in result.results:
|
||||
logger.info(f"Fact: {fact.text[:100]}...")
|
||||
# Check for common Chinese characters or the name
|
||||
if any(
|
||||
char in fact.text
|
||||
for char in ["张", "伟", "腾讯", "软件", "工程师", "分布式", "系统", "代码"]
|
||||
):
|
||||
chinese_facts_found += 1
|
||||
|
||||
logger.info("Chinese retain test passed - facts preserved in Chinese")
|
||||
logger.info(f"Found {chinese_facts_found} facts with Chinese content")
|
||||
assert chinese_facts_found > 0, (
|
||||
f"Expected facts to contain Chinese characters, but none found. "
|
||||
f"Facts: {[f.text for f in result.results]}"
|
||||
)
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
logger.info("Chinese retain test passed - facts preserved in Chinese")
|
||||
return # Test passed
|
||||
|
||||
except AssertionError as e:
|
||||
last_error = e
|
||||
if attempt < max_retries - 1:
|
||||
logger.warning(f"Attempt {attempt + 1} failed: {e}. Retrying...")
|
||||
else:
|
||||
raise e
|
||||
finally:
|
||||
try:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
||||
@@ -231,6 +231,61 @@ async def test_recall_chunks_ordering_by_relevance(memory, request_context):
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_recall_chunks_for_observations(memory, request_context):
|
||||
"""
|
||||
Test that chunks are returned when recalling only observations.
|
||||
|
||||
Observations have no direct chunk_id (they are synthesized from source memories).
|
||||
When include_chunks=True, chunks should be resolved via source_memory_ids.
|
||||
"""
|
||||
bank_id = "test-chunks-observations"
|
||||
|
||||
try:
|
||||
# Retain content that will generate observations via consolidation
|
||||
test_content = """
|
||||
Alice is a senior software engineer at a large technology company.
|
||||
She specializes in distributed systems and has 10 years of experience.
|
||||
Alice leads a team of 8 engineers working on cloud infrastructure.
|
||||
She holds a PhD in computer science from Stanford University.
|
||||
Alice has published several papers on fault-tolerant distributed systems.
|
||||
""" * 8
|
||||
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content=test_content,
|
||||
context="profile notes",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Trigger consolidation explicitly to ensure observations exist
|
||||
await memory.run_consolidation(bank_id=bank_id, request_context=request_context)
|
||||
|
||||
# Recall observations only with chunks enabled
|
||||
result = await memory.recall_async(
|
||||
bank_id=bank_id,
|
||||
query="Alice software engineer",
|
||||
fact_type=["observation"],
|
||||
max_tokens=4096,
|
||||
include_chunks=True,
|
||||
max_chunk_tokens=2000,
|
||||
budget=Budget.MID,
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# If observations were created, chunks should be resolved from source memories
|
||||
if len(result.results) > 0:
|
||||
assert result.chunks is not None, "Should include chunks dict when observations are found"
|
||||
assert len(result.chunks) > 0, "Should return chunks resolved from observation source memories"
|
||||
|
||||
for chunk_id, chunk_info in result.chunks.items():
|
||||
assert len(chunk_info.chunk_text) > 0, "Chunks should contain text"
|
||||
assert chunk_info.chunk_index >= 0, "Chunk should have valid index"
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_recall_chunks_without_include_flag(memory, request_context):
|
||||
"""
|
||||
|
||||
@@ -7,14 +7,17 @@ These tests verify:
|
||||
3. Recovery from tool execution errors
|
||||
"""
|
||||
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
from hindsight_api.engine.reflect.agent import (
|
||||
_normalize_tool_name,
|
||||
_is_done_tool,
|
||||
_clean_answer_text,
|
||||
_clean_done_answer,
|
||||
_count_messages_tokens,
|
||||
_is_context_overflow_error,
|
||||
_is_done_tool,
|
||||
_normalize_tool_name,
|
||||
run_reflect_agent,
|
||||
)
|
||||
from hindsight_api.engine.response_models import LLMToolCall, LLMToolCallResult, TokenUsage
|
||||
@@ -412,3 +415,193 @@ class TestReflectAgentMocked:
|
||||
# Should have a result even if no memories found
|
||||
assert result is not None
|
||||
assert result.iterations == 3
|
||||
|
||||
|
||||
class TestContextOverflowHelpers:
|
||||
"""Unit tests for context-overflow detection helpers."""
|
||||
|
||||
def test_count_messages_tokens_basic(self):
|
||||
messages = [
|
||||
{"role": "system", "content": "You are a helpful assistant."},
|
||||
{"role": "user", "content": "What is the capital of France?"},
|
||||
]
|
||||
count = _count_messages_tokens(messages)
|
||||
assert count > 0
|
||||
# Rough sanity check: ~10 tokens for each message
|
||||
assert count < 100
|
||||
|
||||
def test_count_messages_tokens_with_tool_result(self):
|
||||
"""A large tool result should substantially increase the count."""
|
||||
small_messages = [{"role": "user", "content": "hi"}]
|
||||
large_messages = [
|
||||
{"role": "user", "content": "hi"},
|
||||
{
|
||||
"role": "tool",
|
||||
"tool_call_id": "x",
|
||||
"name": "recall",
|
||||
"content": '{"memories": [' + ', '.join([f'{{"id": "m{i}", "content": "A long memory fact about some topic that goes on and on."}}' for i in range(50)]) + ']}',
|
||||
},
|
||||
]
|
||||
small = _count_messages_tokens(small_messages)
|
||||
large = _count_messages_tokens(large_messages)
|
||||
assert large > small + 200
|
||||
|
||||
def test_is_context_overflow_error_openai(self):
|
||||
assert _is_context_overflow_error(Exception("context_length_exceeded: too many tokens"))
|
||||
assert _is_context_overflow_error(Exception("This model's maximum context length is 128000 tokens. However, your messages resulted in 142164 tokens."))
|
||||
|
||||
def test_is_context_overflow_error_anthropic(self):
|
||||
assert _is_context_overflow_error(Exception("prompt_too_long"))
|
||||
assert _is_context_overflow_error(Exception("prompt is too long for this model"))
|
||||
|
||||
def test_is_context_overflow_error_gemini(self):
|
||||
assert _is_context_overflow_error(Exception("RESOURCE_EXHAUSTED: quota exceeded"))
|
||||
|
||||
def test_is_context_overflow_error_generic(self):
|
||||
assert _is_context_overflow_error(Exception("input is too long to process"))
|
||||
assert _is_context_overflow_error(Exception("too many tokens in the request"))
|
||||
|
||||
def test_is_context_overflow_error_unrelated(self):
|
||||
assert not _is_context_overflow_error(Exception("connection timeout"))
|
||||
assert not _is_context_overflow_error(Exception("rate limit exceeded"))
|
||||
assert not _is_context_overflow_error(ValueError("invalid argument"))
|
||||
|
||||
|
||||
class TestContextOverflowBehavior:
|
||||
"""Test that the reflect agent handles context overflow gracefully."""
|
||||
|
||||
@pytest.fixture
|
||||
def mock_llm(self):
|
||||
llm = MagicMock()
|
||||
llm.call_with_tools = AsyncMock()
|
||||
llm.call = AsyncMock(
|
||||
return_value=("Synthesized answer from gathered evidence.", TokenUsage(input_tokens=50, output_tokens=20, total_tokens=70))
|
||||
)
|
||||
return llm
|
||||
|
||||
@pytest.fixture
|
||||
def mock_functions_with_large_output(self):
|
||||
"""Mock functions that return a large enough payload to exceed a tiny token budget."""
|
||||
large_memories = [
|
||||
{"id": f"mem-{i}", "content": f"Memory fact number {i}: " + "A" * 200}
|
||||
for i in range(20)
|
||||
]
|
||||
return {
|
||||
"search_mental_models_fn": AsyncMock(return_value={"mental_models": []}),
|
||||
"search_observations_fn": AsyncMock(return_value={"observations": []}),
|
||||
"recall_fn": AsyncMock(return_value={"memories": large_memories}),
|
||||
"expand_fn": AsyncMock(return_value={"memories": []}),
|
||||
}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_proactive_guard_fires_when_budget_exceeded(self, mock_llm, mock_functions_with_large_output):
|
||||
"""When token count exceeds max_context_tokens after a tool call, the agent
|
||||
should immediately synthesize from gathered evidence instead of making
|
||||
another LLM call that would overflow."""
|
||||
# First call: LLM calls recall (forced by iter 0 with no mental models)
|
||||
mock_llm.call_with_tools.return_value = LLMToolCallResult(
|
||||
tool_calls=[LLMToolCall(id="1", name="recall", arguments={"query": "test"})],
|
||||
finish_reason="tool_calls",
|
||||
)
|
||||
|
||||
# Set a tiny token budget — the recall result alone will blow past it
|
||||
result = await run_reflect_agent(
|
||||
llm_config=mock_llm,
|
||||
bank_id="test-bank",
|
||||
query="What do you know?",
|
||||
bank_profile={"name": "Test", "mission": "Testing"},
|
||||
max_context_tokens=100,
|
||||
**mock_functions_with_large_output,
|
||||
)
|
||||
|
||||
assert result.text == "Synthesized answer from gathered evidence."
|
||||
# call_with_tools was called once (for the forced recall), then the guard
|
||||
# kicked in — no further tool-call iterations
|
||||
assert mock_llm.call_with_tools.call_count == 1
|
||||
# llm.call() was invoked to generate the final synthesis
|
||||
mock_llm.call.assert_called_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_context_overflow_error_skips_retry(self, mock_llm, mock_functions_with_large_output):
|
||||
"""A context_length_exceeded error from the LLM should NOT be retried —
|
||||
it should immediately fall back to final synthesis."""
|
||||
mock_llm.call_with_tools.side_effect = Exception(
|
||||
"context_length_exceeded: messages resulted in 150000 tokens."
|
||||
)
|
||||
|
||||
result = await run_reflect_agent(
|
||||
llm_config=mock_llm,
|
||||
bank_id="test-bank",
|
||||
query="What do you know?",
|
||||
bank_profile={"name": "Test", "mission": "Testing"},
|
||||
max_iterations=5,
|
||||
**mock_functions_with_large_output,
|
||||
)
|
||||
|
||||
assert result is not None
|
||||
# Should have attempted only 1 iteration (no retry on overflow error)
|
||||
assert mock_llm.call_with_tools.call_count == 1
|
||||
# Final synthesis was called
|
||||
mock_llm.call.assert_called_once()
|
||||
|
||||
|
||||
class TestContextOverflowIntegration:
|
||||
"""Integration test: real LLM with a very small max_context_tokens.
|
||||
|
||||
The agent will make one real LLM call (forced tool choice), receive a large
|
||||
tool result that exceeds the tiny budget, then synthesize from it via a second
|
||||
real LLM call — all without raising a context_length_exceeded error.
|
||||
"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reflect_completes_with_tiny_context_budget(self, memory, request_context):
|
||||
"""End-to-end: reflect on a bank with max_context_tokens=1 (tiny budget).
|
||||
|
||||
Setting max_context_tokens=1 guarantees the proactive guard fires as soon
|
||||
as the first tool result is received and evidence is available.
|
||||
The result must be a non-empty string with no exception raised.
|
||||
"""
|
||||
import uuid
|
||||
from unittest.mock import patch
|
||||
|
||||
bank_id = f"test-ctx-overflow-{uuid.uuid4().hex[:8]}"
|
||||
try:
|
||||
# Retain a handful of facts so the recall tool has something to return
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="Alice is a software engineer who enjoys hiking on weekends.",
|
||||
request_context=request_context,
|
||||
)
|
||||
await memory.retain_async(
|
||||
bank_id=bank_id,
|
||||
content="Bob is a designer who loves cooking Italian food.",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
# Patch get_config where memory_engine uses it, injecting a tiny
|
||||
# max_context_tokens. Everything else delegates to the real config.
|
||||
real_config = memory._get_raw_config() if hasattr(memory, "_get_raw_config") else None
|
||||
from hindsight_api.config import get_config as _real_get_config
|
||||
|
||||
class _TinyContextProxy:
|
||||
"""Forwards all attribute access to the real config proxy except
|
||||
reflect_max_context_tokens which is forced to 1."""
|
||||
_real = _real_get_config()
|
||||
|
||||
def __getattr__(self, name: str):
|
||||
if name == "reflect_max_context_tokens":
|
||||
return 1
|
||||
return getattr(self._real, name)
|
||||
|
||||
with patch("hindsight_api.engine.memory_engine.get_config", return_value=_TinyContextProxy()):
|
||||
result = await memory.reflect_async(
|
||||
bank_id=bank_id,
|
||||
query="Tell me about the people you know.",
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert result.text, "reflect must return a non-empty answer"
|
||||
assert result.usage.total_tokens > 0
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
@@ -591,6 +591,125 @@ async def test_mentioned_at_from_context_string(memory, request_context):
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
# ============================================================
|
||||
# No Timestamp Tests
|
||||
# ============================================================
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retain_no_timestamp(memory, request_context):
|
||||
"""
|
||||
Test retaining content with explicit "no timestamp" sentinel.
|
||||
|
||||
When event_date=None is passed explicitly in the dict (i.e. caller opted into
|
||||
no timestamp), mentioned_at should be NULL in the DB rather than defaulting to now().
|
||||
"""
|
||||
bank_id = f"test_no_timestamp_{datetime.now(timezone.utc).timestamp()}"
|
||||
|
||||
try:
|
||||
# Use retain_batch_async with explicit event_date=None key to signal "no timestamp"
|
||||
unit_ids_list = await memory.retain_batch_async(
|
||||
bank_id=bank_id,
|
||||
contents=[
|
||||
{
|
||||
"content": "The capital of France is Paris. The Eiffel Tower is located in Paris.",
|
||||
"context": "general knowledge",
|
||||
"event_date": None, # Explicit sentinel: no timestamp
|
||||
}
|
||||
],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert len(unit_ids_list) > 0, "Should create at least one batch result"
|
||||
unit_ids = unit_ids_list[0]
|
||||
assert len(unit_ids) > 0, "Should have extracted and stored facts"
|
||||
|
||||
# Recall the facts
|
||||
result = await memory.recall_async(
|
||||
bank_id=bank_id,
|
||||
query="Where is the Eiffel Tower?",
|
||||
budget=Budget.LOW,
|
||||
max_tokens=500,
|
||||
fact_type=["world"],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert len(result.results) > 0, "Should recall the stored fact"
|
||||
|
||||
# All temporal fields should be None for temporally agnostic content
|
||||
for fact in result.results:
|
||||
assert fact.mentioned_at is None, (
|
||||
f"mentioned_at should be None for no-timestamp content, got {fact.mentioned_at}"
|
||||
)
|
||||
|
||||
print(f"\n✓ Test passed: mentioned_at is None for {len(result.results)} fact(s)")
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retain_omit_timestamp_defaults_to_now(memory, request_context):
|
||||
"""
|
||||
Backward-compatibility regression test: omitting event_date still stores a real datetime.
|
||||
|
||||
When event_date is absent from the content dict (key not present), the orchestrator
|
||||
should default to utcnow() — preserving existing behavior.
|
||||
"""
|
||||
bank_id = f"test_default_timestamp_{datetime.now(timezone.utc).timestamp()}"
|
||||
before = datetime.now(timezone.utc)
|
||||
|
||||
try:
|
||||
# Omit event_date entirely — should default to now()
|
||||
unit_ids_list = await memory.retain_batch_async(
|
||||
bank_id=bank_id,
|
||||
contents=[
|
||||
{
|
||||
"content": "Alice is a software engineer who loves Python.",
|
||||
"context": "profile",
|
||||
# event_date intentionally omitted
|
||||
}
|
||||
],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
after = datetime.now(timezone.utc)
|
||||
|
||||
assert len(unit_ids_list) > 0
|
||||
unit_ids = unit_ids_list[0]
|
||||
assert len(unit_ids) > 0, "Should have extracted and stored facts"
|
||||
|
||||
# Recall and verify mentioned_at is a real datetime close to now
|
||||
result = await memory.recall_async(
|
||||
bank_id=bank_id,
|
||||
query="Who is Alice?",
|
||||
budget=Budget.LOW,
|
||||
max_tokens=500,
|
||||
fact_type=["world"],
|
||||
request_context=request_context,
|
||||
)
|
||||
|
||||
assert len(result.results) > 0, "Should recall the fact"
|
||||
fact = result.results[0]
|
||||
|
||||
assert fact.mentioned_at is not None, "mentioned_at should be set when event_date is omitted"
|
||||
|
||||
if isinstance(fact.mentioned_at, str):
|
||||
mentioned_dt = datetime.fromisoformat(fact.mentioned_at.replace("Z", "+00:00"))
|
||||
else:
|
||||
mentioned_dt = fact.mentioned_at
|
||||
|
||||
# Should be within 60s of when we ran the test
|
||||
assert before <= mentioned_dt <= after + timedelta(seconds=60), (
|
||||
f"mentioned_at {mentioned_dt} should be close to now ({before} – {after})"
|
||||
)
|
||||
|
||||
print(f"\n✓ Test passed: mentioned_at={mentioned_dt} is a real datetime (backward compat)")
|
||||
|
||||
finally:
|
||||
await memory.delete_bank(bank_id, request_context=request_context)
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Context Tracking Tests
|
||||
# ============================================================
|
||||
|
||||
@@ -233,6 +233,35 @@ class TestFilterResultsByTags:
|
||||
assert len(filtered) == 1
|
||||
assert filtered[0].tags == ["a", "b", "c"] # Has a, b, AND c
|
||||
|
||||
def test_all_strict_superset_observation_matches_incoming_memory_tags(self):
|
||||
"""
|
||||
Consolidation scenario: an incoming memory with tags ['user:bob', 'session:id1']
|
||||
uses all_strict matching to find existing observations.
|
||||
|
||||
An observation tagged ['user:bob', 'session:id1', 'place:online'] IS matched
|
||||
because it contains all of the incoming memory's tags (superset).
|
||||
This is NOT exact matching — an observation with extra tags is still a valid match.
|
||||
"""
|
||||
# Incoming memory tags (e.g. from a new retain call)
|
||||
incoming_tags = ["user:bob", "session:id1"]
|
||||
|
||||
# Candidate observations with different tag sets
|
||||
exact_match = MockResult(["user:bob", "session:id1"])
|
||||
superset_match = MockResult(["session:id1", "user:bob", "place:online"])
|
||||
different_user = MockResult(["user:alice", "session:id1"])
|
||||
missing_session = MockResult(["user:bob"])
|
||||
|
||||
results = [exact_match, superset_match, different_user, missing_session]
|
||||
filtered = filter_results_by_tags(results, incoming_tags, match="all_strict")
|
||||
|
||||
# Both exact_match and superset_match have all incoming tags → both match
|
||||
assert len(filtered) == 2
|
||||
assert exact_match in filtered
|
||||
assert superset_match in filtered
|
||||
# different_user and missing_session are excluded because they lack at least one tag
|
||||
assert different_user not in filtered
|
||||
assert missing_session not in filtered
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Integration Tests for tags in retain/recall/reflect
|
||||
|
||||
@@ -0,0 +1,780 @@
|
||||
"""Tests for the webhook system.
|
||||
|
||||
Covers:
|
||||
- Unit tests for HMAC signing and retry constants (no DB required)
|
||||
- Integration tests for fire_event() using a real DB (inserts into async_operations)
|
||||
- Integration tests for _handle_webhook_delivery() on the memory engine
|
||||
- HTTP API integration tests for CRUD and delivery listing endpoints
|
||||
"""
|
||||
|
||||
import json
|
||||
import uuid
|
||||
from datetime import datetime, timezone
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
|
||||
from hindsight_api.api import create_app
|
||||
from hindsight_api.engine.memory_engine import MemoryEngine
|
||||
from hindsight_api.webhooks.manager import MAX_ATTEMPTS, RETRY_DELAYS, WebhookManager
|
||||
from hindsight_api.webhooks.models import (
|
||||
ConsolidationEventData,
|
||||
RetainEventData,
|
||||
WebhookConfig,
|
||||
WebhookEvent,
|
||||
WebhookEventType,
|
||||
)
|
||||
from hindsight_api.worker.exceptions import RetryTaskAt
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _make_event(bank_id: str = "bank-1") -> WebhookEvent:
|
||||
return WebhookEvent(
|
||||
event=WebhookEventType.CONSOLIDATION_COMPLETED,
|
||||
bank_id=bank_id,
|
||||
operation_id=uuid.uuid4().hex,
|
||||
status="completed",
|
||||
timestamp=datetime.now(timezone.utc),
|
||||
data=ConsolidationEventData(observations_created=1),
|
||||
)
|
||||
|
||||
|
||||
def _make_delivery_task(
|
||||
bank_id: str = "bank-1",
|
||||
url: str = "https://example.com/hook",
|
||||
retry_count: int = 0,
|
||||
webhook_id: str | None = None,
|
||||
) -> dict:
|
||||
return {
|
||||
"type": "webhook_delivery",
|
||||
"bank_id": bank_id,
|
||||
"url": url,
|
||||
"secret": None,
|
||||
"event_type": "consolidation.completed",
|
||||
"payload": '{"event":"consolidation.completed"}',
|
||||
"webhook_id": webhook_id,
|
||||
"_retry_count": retry_count,
|
||||
}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Unit tests (no DB)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestHmacSigning:
|
||||
"""Unit tests for WebhookManager._sign_payload()."""
|
||||
|
||||
def _make_manager(self) -> WebhookManager:
|
||||
"""Create a WebhookManager with a dummy pool (not used for signing)."""
|
||||
pool = MagicMock()
|
||||
return WebhookManager(pool=pool, global_webhooks=[])
|
||||
|
||||
def test_hmac_signing_format(self):
|
||||
"""_sign_payload should return a string starting with 'sha256='."""
|
||||
manager = self._make_manager()
|
||||
sig = manager._sign_payload("my-secret", b"hello world")
|
||||
assert sig.startswith("sha256="), f"Expected 'sha256=' prefix, got: {sig!r}"
|
||||
hex_part = sig[len("sha256="):]
|
||||
# SHA-256 hex digest is always 64 characters
|
||||
assert len(hex_part) == 64
|
||||
# Hex characters only
|
||||
assert all(c in "0123456789abcdef" for c in hex_part)
|
||||
|
||||
def test_hmac_signing_is_deterministic(self):
|
||||
"""Same secret + payload always produces the same signature."""
|
||||
manager = self._make_manager()
|
||||
payload = b'{"event":"consolidation.completed"}'
|
||||
sig1 = manager._sign_payload("secret-key", payload)
|
||||
sig2 = manager._sign_payload("secret-key", payload)
|
||||
assert sig1 == sig2
|
||||
|
||||
def test_hmac_signing_differs_with_different_secret(self):
|
||||
"""Different secrets must produce different signatures."""
|
||||
manager = self._make_manager()
|
||||
payload = b"payload"
|
||||
sig1 = manager._sign_payload("secret-a", payload)
|
||||
sig2 = manager._sign_payload("secret-b", payload)
|
||||
assert sig1 != sig2
|
||||
|
||||
def test_hmac_signing_differs_with_different_payload(self):
|
||||
"""Different payloads must produce different signatures."""
|
||||
manager = self._make_manager()
|
||||
sig1 = manager._sign_payload("secret", b"payload-one")
|
||||
sig2 = manager._sign_payload("secret", b"payload-two")
|
||||
assert sig1 != sig2
|
||||
|
||||
|
||||
class TestRetryConstants:
|
||||
"""Unit tests to verify retry schedule constants."""
|
||||
|
||||
def test_retry_delays_values(self):
|
||||
"""RETRY_DELAYS must match the documented schedule."""
|
||||
assert RETRY_DELAYS == [5, 300, 1800, 7200, 18000]
|
||||
|
||||
def test_max_attempts(self):
|
||||
"""MAX_ATTEMPTS should be len(RETRY_DELAYS) + 1."""
|
||||
assert MAX_ATTEMPTS == 6
|
||||
assert MAX_ATTEMPTS == len(RETRY_DELAYS) + 1
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# DB integration tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def webhook_manager(memory: MemoryEngine) -> WebhookManager:
|
||||
"""Return a WebhookManager backed by the test pool with no global webhooks."""
|
||||
return WebhookManager(pool=memory._pool, global_webhooks=[])
|
||||
|
||||
|
||||
class TestFireEvent:
|
||||
"""Integration tests for WebhookManager.fire_event()."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fire_event_creates_delivery(
|
||||
self, memory: MemoryEngine, webhook_manager: WebhookManager
|
||||
):
|
||||
"""fire_event() inserts a pending webhook_delivery task in async_operations."""
|
||||
bank_id = f"wh-test-{uuid.uuid4().hex[:8]}"
|
||||
webhook_id = uuid.uuid4()
|
||||
|
||||
async with memory._pool.acquire() as conn:
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO webhooks (id, bank_id, url, secret, event_types, enabled, created_at, updated_at)
|
||||
VALUES ($1, $2, $3, NULL, $4, true, NOW(), NOW())
|
||||
""",
|
||||
webhook_id,
|
||||
bank_id,
|
||||
"https://example.com/hook",
|
||||
["consolidation.completed"],
|
||||
)
|
||||
|
||||
try:
|
||||
event = _make_event(bank_id)
|
||||
await webhook_manager.fire_event(event)
|
||||
|
||||
async with memory._pool.acquire() as conn:
|
||||
rows = await conn.fetch(
|
||||
"""
|
||||
SELECT status, task_payload
|
||||
FROM async_operations
|
||||
WHERE operation_type = 'webhook_delivery'
|
||||
AND bank_id = $1
|
||||
AND task_payload->>'webhook_id' = $2
|
||||
""",
|
||||
bank_id,
|
||||
str(webhook_id),
|
||||
)
|
||||
|
||||
assert len(rows) == 1
|
||||
assert rows[0]["status"] == "pending"
|
||||
payload = rows[0]["task_payload"]
|
||||
if isinstance(payload, str):
|
||||
payload = json.loads(payload)
|
||||
assert payload["event_type"] == "consolidation.completed"
|
||||
finally:
|
||||
async with memory._pool.acquire() as conn:
|
||||
await conn.execute(
|
||||
"DELETE FROM async_operations WHERE operation_type = 'webhook_delivery' AND bank_id = $1",
|
||||
bank_id,
|
||||
)
|
||||
await conn.execute("DELETE FROM webhooks WHERE id = $1", webhook_id)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fire_event_global_webhook(
|
||||
self, memory: MemoryEngine
|
||||
):
|
||||
"""fire_event() also queues delivery tasks for global webhooks (not stored in DB)."""
|
||||
bank_id = f"wh-global-{uuid.uuid4().hex[:8]}"
|
||||
global_webhook = WebhookConfig(
|
||||
id="", # No DB row
|
||||
bank_id=None,
|
||||
url="https://global.example.com/hook",
|
||||
secret=None,
|
||||
event_types=["consolidation.completed"],
|
||||
enabled=True,
|
||||
)
|
||||
manager = WebhookManager(pool=memory._pool, global_webhooks=[global_webhook])
|
||||
|
||||
event = _make_event(bank_id)
|
||||
await manager.fire_event(event)
|
||||
|
||||
async with memory._pool.acquire() as conn:
|
||||
rows = await conn.fetch(
|
||||
"""
|
||||
SELECT status, task_payload
|
||||
FROM async_operations
|
||||
WHERE operation_type = 'webhook_delivery'
|
||||
AND bank_id = $1
|
||||
AND task_payload->>'url' = 'https://global.example.com/hook'
|
||||
ORDER BY created_at DESC
|
||||
LIMIT 1
|
||||
"""
|
||||
,
|
||||
bank_id,
|
||||
)
|
||||
|
||||
assert len(rows) == 1
|
||||
assert rows[0]["status"] == "pending"
|
||||
payload = rows[0]["task_payload"]
|
||||
if isinstance(payload, str):
|
||||
payload = json.loads(payload)
|
||||
assert payload["webhook_id"] is None # global webhook has no DB row
|
||||
|
||||
# Cleanup
|
||||
async with memory._pool.acquire() as conn:
|
||||
await conn.execute(
|
||||
"DELETE FROM async_operations WHERE operation_type = 'webhook_delivery' AND bank_id = $1",
|
||||
bank_id,
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fire_event_no_match_if_event_type_mismatch(
|
||||
self, memory: MemoryEngine, webhook_manager: WebhookManager
|
||||
):
|
||||
"""Webhooks registered for a different event type receive no delivery task."""
|
||||
bank_id = f"wh-mismatch-{uuid.uuid4().hex[:8]}"
|
||||
webhook_id = uuid.uuid4()
|
||||
|
||||
async with memory._pool.acquire() as conn:
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO webhooks (id, bank_id, url, secret, event_types, enabled, created_at, updated_at)
|
||||
VALUES ($1, $2, $3, NULL, $4, true, NOW(), NOW())
|
||||
""",
|
||||
webhook_id,
|
||||
bank_id,
|
||||
"https://example.com/other-hook",
|
||||
["other.event"],
|
||||
)
|
||||
|
||||
try:
|
||||
event = _make_event(bank_id)
|
||||
await webhook_manager.fire_event(event)
|
||||
|
||||
async with memory._pool.acquire() as conn:
|
||||
count = await conn.fetchval(
|
||||
"""
|
||||
SELECT COUNT(*) FROM async_operations
|
||||
WHERE operation_type = 'webhook_delivery' AND bank_id = $1
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
|
||||
assert count == 0
|
||||
finally:
|
||||
async with memory._pool.acquire() as conn:
|
||||
await conn.execute("DELETE FROM webhooks WHERE id = $1", webhook_id)
|
||||
|
||||
|
||||
class TestHandleWebhookDelivery:
|
||||
"""Integration tests for MemoryEngine._handle_webhook_delivery()."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deliver_success(self, memory: MemoryEngine):
|
||||
"""A successful HTTP POST completes without raising."""
|
||||
task_dict = _make_delivery_task(retry_count=0)
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.raise_for_status = MagicMock()
|
||||
|
||||
with patch.object(memory._http_client, "post", new=AsyncMock(return_value=mock_response)):
|
||||
# Should not raise
|
||||
await memory._handle_webhook_delivery(task_dict)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deliver_failure_raises_retry_task_at(self, memory: MemoryEngine):
|
||||
"""A failed HTTP POST raises RetryTaskAt when retries remain."""
|
||||
task_dict = _make_delivery_task(retry_count=0)
|
||||
|
||||
with patch.object(
|
||||
memory._http_client, "post", new=AsyncMock(side_effect=Exception("connection refused"))
|
||||
):
|
||||
with pytest.raises(RetryTaskAt):
|
||||
await memory._handle_webhook_delivery(task_dict)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deliver_exhausted_retries_raises(self, memory: MemoryEngine):
|
||||
"""When retry_count reaches MAX_ATTEMPTS-1, a failure raises the original exception."""
|
||||
task_dict = _make_delivery_task(retry_count=MAX_ATTEMPTS - 1)
|
||||
|
||||
with patch.object(
|
||||
memory._http_client, "post", new=AsyncMock(side_effect=Exception("server error"))
|
||||
):
|
||||
with pytest.raises(Exception, match="server error"):
|
||||
await memory._handle_webhook_delivery(task_dict)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deliver_retry_at_uses_delay_schedule(self, memory: MemoryEngine):
|
||||
"""RetryTaskAt.retry_at is approximately now + RETRY_DELAYS[retry_count]."""
|
||||
from datetime import timedelta
|
||||
|
||||
task_dict = _make_delivery_task(retry_count=1)
|
||||
|
||||
with patch.object(
|
||||
memory._http_client, "post", new=AsyncMock(side_effect=Exception("fail"))
|
||||
):
|
||||
before = datetime.now(timezone.utc)
|
||||
with pytest.raises(RetryTaskAt) as exc_info:
|
||||
await memory._handle_webhook_delivery(task_dict)
|
||||
after = datetime.now(timezone.utc)
|
||||
|
||||
retry_at = exc_info.value.retry_at
|
||||
expected_delay = RETRY_DELAYS[1] # retry_count=1
|
||||
assert retry_at >= before + timedelta(seconds=expected_delay - 2)
|
||||
assert retry_at <= after + timedelta(seconds=expected_delay + 2)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_task_marks_operation_completed(self, memory: MemoryEngine):
|
||||
"""After a successful delivery, execute_task marks the async_operations row as completed."""
|
||||
operation_id = str(uuid.uuid4())
|
||||
bank_id = f"wh-exec-{uuid.uuid4().hex[:8]}"
|
||||
|
||||
# Insert a real async_operations row so _mark_operation_completed has something to update
|
||||
async with memory._pool.acquire() as conn:
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO async_operations
|
||||
(operation_id, bank_id, operation_type, status, task_payload, result_metadata, created_at, updated_at)
|
||||
VALUES ($1, $2, 'webhook_delivery', 'processing', '{}'::jsonb, '{}'::jsonb, NOW(), NOW())
|
||||
""",
|
||||
uuid.UUID(operation_id),
|
||||
bank_id,
|
||||
)
|
||||
|
||||
task_dict = {
|
||||
**_make_delivery_task(bank_id=bank_id, retry_count=0),
|
||||
"operation_id": operation_id,
|
||||
}
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.raise_for_status = MagicMock()
|
||||
|
||||
with patch.object(memory._http_client, "post", new=AsyncMock(return_value=mock_response)):
|
||||
await memory.execute_task(task_dict)
|
||||
|
||||
async with memory._pool.acquire() as conn:
|
||||
row = await conn.fetchrow(
|
||||
"SELECT status FROM async_operations WHERE operation_id = $1",
|
||||
uuid.UUID(operation_id),
|
||||
)
|
||||
|
||||
assert row is not None
|
||||
assert row["status"] == "completed", f"Expected 'completed', got '{row['status']}'"
|
||||
|
||||
# Cleanup
|
||||
async with memory._pool.acquire() as conn:
|
||||
await conn.execute(
|
||||
"DELETE FROM async_operations WHERE operation_id = $1",
|
||||
uuid.UUID(operation_id),
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# HTTP API integration tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def api_client(memory: MemoryEngine):
|
||||
"""Async HTTP test client wired to the FastAPI app."""
|
||||
app = create_app(memory, initialize_memory=False)
|
||||
transport = httpx.ASGITransport(app=app)
|
||||
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
|
||||
yield client
|
||||
|
||||
|
||||
class TestWebhookHttpApi:
|
||||
"""HTTP API integration tests for webhook CRUD endpoints."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_http_create_webhook(self, api_client: httpx.AsyncClient):
|
||||
"""POST /webhooks returns 201 and an id."""
|
||||
bank_id = f"http-wh-{uuid.uuid4().hex[:8]}"
|
||||
response = await api_client.post(
|
||||
f"/v1/default/banks/{bank_id}/webhooks",
|
||||
json={
|
||||
"url": "https://example.com/create",
|
||||
"event_types": ["consolidation.completed"],
|
||||
},
|
||||
)
|
||||
assert response.status_code == 201, response.text
|
||||
data = response.json()
|
||||
assert "id" in data
|
||||
assert data["url"] == "https://example.com/create"
|
||||
assert data["bank_id"] == bank_id
|
||||
assert data["secret"] is None # secrets are never echoed back
|
||||
|
||||
# Cleanup
|
||||
await api_client.delete(
|
||||
f"/v1/default/banks/{bank_id}/webhooks/{data['id']}"
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_http_list_webhooks(self, api_client: httpx.AsyncClient):
|
||||
"""GET /webhooks returns the webhooks registered for a bank."""
|
||||
bank_id = f"http-wh-{uuid.uuid4().hex[:8]}"
|
||||
|
||||
create_resp = await api_client.post(
|
||||
f"/v1/default/banks/{bank_id}/webhooks",
|
||||
json={"url": "https://example.com/list", "event_types": ["consolidation.completed"]},
|
||||
)
|
||||
assert create_resp.status_code == 201
|
||||
webhook_id = create_resp.json()["id"]
|
||||
|
||||
list_resp = await api_client.get(f"/v1/default/banks/{bank_id}/webhooks")
|
||||
assert list_resp.status_code == 200
|
||||
items = list_resp.json()["items"]
|
||||
assert any(item["id"] == webhook_id for item in items)
|
||||
|
||||
# Cleanup
|
||||
await api_client.delete(f"/v1/default/banks/{bank_id}/webhooks/{webhook_id}")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_http_delete_webhook(self, api_client: httpx.AsyncClient):
|
||||
"""DELETE /webhooks/{id} removes the webhook; subsequent list returns empty for that bank."""
|
||||
bank_id = f"http-wh-{uuid.uuid4().hex[:8]}"
|
||||
|
||||
create_resp = await api_client.post(
|
||||
f"/v1/default/banks/{bank_id}/webhooks",
|
||||
json={"url": "https://example.com/delete", "event_types": ["consolidation.completed"]},
|
||||
)
|
||||
assert create_resp.status_code == 201
|
||||
webhook_id = create_resp.json()["id"]
|
||||
|
||||
delete_resp = await api_client.delete(
|
||||
f"/v1/default/banks/{bank_id}/webhooks/{webhook_id}"
|
||||
)
|
||||
assert delete_resp.status_code == 200
|
||||
assert delete_resp.json()["success"] is True
|
||||
|
||||
list_resp = await api_client.get(f"/v1/default/banks/{bank_id}/webhooks")
|
||||
assert list_resp.status_code == 200
|
||||
ids = [item["id"] for item in list_resp.json()["items"]]
|
||||
assert webhook_id not in ids
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_http_delete_webhook_not_found(self, api_client: httpx.AsyncClient):
|
||||
"""DELETE with a non-existent webhook id returns 404."""
|
||||
bank_id = f"http-wh-{uuid.uuid4().hex[:8]}"
|
||||
missing_id = str(uuid.uuid4())
|
||||
response = await api_client.delete(
|
||||
f"/v1/default/banks/{bank_id}/webhooks/{missing_id}"
|
||||
)
|
||||
assert response.status_code == 404
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_http_list_deliveries(
|
||||
self, memory: MemoryEngine, api_client: httpx.AsyncClient
|
||||
):
|
||||
"""GET /webhooks/{id}/deliveries returns delivery records for a webhook."""
|
||||
bank_id = f"http-wh-{uuid.uuid4().hex[:8]}"
|
||||
|
||||
# Create webhook via HTTP API
|
||||
create_resp = await api_client.post(
|
||||
f"/v1/default/banks/{bank_id}/webhooks",
|
||||
json={
|
||||
"url": "https://example.com/deliveries",
|
||||
"event_types": ["consolidation.completed"],
|
||||
},
|
||||
)
|
||||
assert create_resp.status_code == 201
|
||||
webhook_id = create_resp.json()["id"]
|
||||
|
||||
# Insert a delivery row directly into async_operations
|
||||
delivery_id = uuid.uuid4()
|
||||
now = datetime.now(timezone.utc)
|
||||
task_payload = json.dumps(
|
||||
{
|
||||
"type": "webhook_delivery",
|
||||
"bank_id": bank_id,
|
||||
"url": "https://example.com/deliveries",
|
||||
"secret": None,
|
||||
"event_type": "consolidation.completed",
|
||||
"payload": '{"event":"consolidation.completed"}',
|
||||
"webhook_id": webhook_id,
|
||||
}
|
||||
)
|
||||
async with memory._pool.acquire() as conn:
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO async_operations
|
||||
(operation_id, bank_id, operation_type, status, retry_count, task_payload, result_metadata, created_at, updated_at)
|
||||
VALUES ($1, $2, 'webhook_delivery', 'completed', 0, $3::jsonb, '{}'::jsonb, $4, $4)
|
||||
""",
|
||||
delivery_id,
|
||||
bank_id,
|
||||
task_payload,
|
||||
now,
|
||||
)
|
||||
|
||||
try:
|
||||
deliveries_resp = await api_client.get(
|
||||
f"/v1/default/banks/{bank_id}/webhooks/{webhook_id}/deliveries"
|
||||
)
|
||||
assert deliveries_resp.status_code == 200
|
||||
items = deliveries_resp.json()["items"]
|
||||
ids = [item["id"] for item in items]
|
||||
assert str(delivery_id) in ids
|
||||
|
||||
# Verify shape of a delivery item
|
||||
delivery = next(item for item in items if item["id"] == str(delivery_id))
|
||||
assert delivery["status"] == "completed"
|
||||
assert delivery["event_type"] == "consolidation.completed"
|
||||
assert delivery["attempts"] == 1
|
||||
finally:
|
||||
async with memory._pool.acquire() as conn:
|
||||
await conn.execute(
|
||||
"DELETE FROM async_operations WHERE operation_id = $1", delivery_id
|
||||
)
|
||||
await api_client.delete(
|
||||
f"/v1/default/banks/{bank_id}/webhooks/{webhook_id}"
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_http_list_deliveries_webhook_not_found(self, api_client: httpx.AsyncClient):
|
||||
"""GET /webhooks/{id}/deliveries for a non-existent webhook returns 404."""
|
||||
bank_id = f"http-wh-{uuid.uuid4().hex[:8]}"
|
||||
missing_id = str(uuid.uuid4())
|
||||
response = await api_client.get(
|
||||
f"/v1/default/banks/{bank_id}/webhooks/{missing_id}/deliveries"
|
||||
)
|
||||
assert response.status_code == 404
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_http_update_webhook_url(self, api_client: httpx.AsyncClient):
|
||||
"""PATCH /webhooks/{id} updates only the provided fields."""
|
||||
bank_id = f"http-wh-{uuid.uuid4().hex[:8]}"
|
||||
|
||||
create_resp = await api_client.post(
|
||||
f"/v1/default/banks/{bank_id}/webhooks",
|
||||
json={"url": "https://example.com/original", "event_types": ["consolidation.completed"]},
|
||||
)
|
||||
assert create_resp.status_code == 201
|
||||
webhook_id = create_resp.json()["id"]
|
||||
|
||||
patch_resp = await api_client.patch(
|
||||
f"/v1/default/banks/{bank_id}/webhooks/{webhook_id}",
|
||||
json={"url": "https://example.com/updated"},
|
||||
)
|
||||
assert patch_resp.status_code == 200
|
||||
data = patch_resp.json()
|
||||
assert data["url"] == "https://example.com/updated"
|
||||
# event_types should be unchanged
|
||||
assert "consolidation.completed" in data["event_types"]
|
||||
|
||||
# Cleanup
|
||||
await api_client.delete(f"/v1/default/banks/{bank_id}/webhooks/{webhook_id}")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_http_update_webhook_event_types(self, api_client: httpx.AsyncClient):
|
||||
"""PATCH /webhooks/{id} can update event_types."""
|
||||
bank_id = f"http-wh-{uuid.uuid4().hex[:8]}"
|
||||
|
||||
create_resp = await api_client.post(
|
||||
f"/v1/default/banks/{bank_id}/webhooks",
|
||||
json={"url": "https://example.com/hook", "event_types": ["consolidation.completed"]},
|
||||
)
|
||||
assert create_resp.status_code == 201
|
||||
webhook_id = create_resp.json()["id"]
|
||||
|
||||
patch_resp = await api_client.patch(
|
||||
f"/v1/default/banks/{bank_id}/webhooks/{webhook_id}",
|
||||
json={"event_types": ["retain.completed"]},
|
||||
)
|
||||
assert patch_resp.status_code == 200
|
||||
data = patch_resp.json()
|
||||
assert data["event_types"] == ["retain.completed"]
|
||||
|
||||
# Cleanup
|
||||
await api_client.delete(f"/v1/default/banks/{bank_id}/webhooks/{webhook_id}")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_http_update_webhook_enabled(self, api_client: httpx.AsyncClient):
|
||||
"""PATCH /webhooks/{id} can toggle enabled."""
|
||||
bank_id = f"http-wh-{uuid.uuid4().hex[:8]}"
|
||||
|
||||
create_resp = await api_client.post(
|
||||
f"/v1/default/banks/{bank_id}/webhooks",
|
||||
json={"url": "https://example.com/hook", "event_types": ["consolidation.completed"]},
|
||||
)
|
||||
assert create_resp.status_code == 201
|
||||
webhook_id = create_resp.json()["id"]
|
||||
assert create_resp.json()["enabled"] is True
|
||||
|
||||
patch_resp = await api_client.patch(
|
||||
f"/v1/default/banks/{bank_id}/webhooks/{webhook_id}",
|
||||
json={"enabled": False},
|
||||
)
|
||||
assert patch_resp.status_code == 200
|
||||
assert patch_resp.json()["enabled"] is False
|
||||
|
||||
# Cleanup
|
||||
await api_client.delete(f"/v1/default/banks/{bank_id}/webhooks/{webhook_id}")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_http_update_webhook_http_config(self, api_client: httpx.AsyncClient):
|
||||
"""PATCH /webhooks/{id} can update http_config."""
|
||||
bank_id = f"http-wh-{uuid.uuid4().hex[:8]}"
|
||||
|
||||
create_resp = await api_client.post(
|
||||
f"/v1/default/banks/{bank_id}/webhooks",
|
||||
json={"url": "https://example.com/hook", "event_types": ["consolidation.completed"]},
|
||||
)
|
||||
assert create_resp.status_code == 201
|
||||
webhook_id = create_resp.json()["id"]
|
||||
|
||||
patch_resp = await api_client.patch(
|
||||
f"/v1/default/banks/{bank_id}/webhooks/{webhook_id}",
|
||||
json={
|
||||
"http_config": {
|
||||
"method": "POST",
|
||||
"timeout_seconds": 10,
|
||||
"headers": {"X-Custom": "value"},
|
||||
"params": {},
|
||||
}
|
||||
},
|
||||
)
|
||||
assert patch_resp.status_code == 200
|
||||
data = patch_resp.json()
|
||||
assert data["http_config"]["timeout_seconds"] == 10
|
||||
assert data["http_config"]["headers"] == {"X-Custom": "value"}
|
||||
|
||||
# Cleanup
|
||||
await api_client.delete(f"/v1/default/banks/{bank_id}/webhooks/{webhook_id}")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_http_update_webhook_not_found(self, api_client: httpx.AsyncClient):
|
||||
"""PATCH /webhooks/{id} returns 404 for a non-existent webhook."""
|
||||
bank_id = f"http-wh-{uuid.uuid4().hex[:8]}"
|
||||
missing_id = str(uuid.uuid4())
|
||||
response = await api_client.patch(
|
||||
f"/v1/default/banks/{bank_id}/webhooks/{missing_id}",
|
||||
json={"url": "https://example.com/new"},
|
||||
)
|
||||
assert response.status_code == 404
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_http_update_webhook_no_fields(self, api_client: httpx.AsyncClient):
|
||||
"""PATCH /webhooks/{id} with empty body returns 422."""
|
||||
bank_id = f"http-wh-{uuid.uuid4().hex[:8]}"
|
||||
|
||||
create_resp = await api_client.post(
|
||||
f"/v1/default/banks/{bank_id}/webhooks",
|
||||
json={"url": "https://example.com/hook", "event_types": ["consolidation.completed"]},
|
||||
)
|
||||
assert create_resp.status_code == 201
|
||||
webhook_id = create_resp.json()["id"]
|
||||
|
||||
patch_resp = await api_client.patch(
|
||||
f"/v1/default/banks/{bank_id}/webhooks/{webhook_id}",
|
||||
json={},
|
||||
)
|
||||
assert patch_resp.status_code == 422
|
||||
|
||||
# Cleanup
|
||||
await api_client.delete(f"/v1/default/banks/{bank_id}/webhooks/{webhook_id}")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# retain.completed webhook tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestRetainCompletedWebhook:
|
||||
"""Tests for the retain.completed webhook event."""
|
||||
|
||||
def test_retain_event_data_model(self):
|
||||
"""RetainEventData can be constructed with optional fields."""
|
||||
data = RetainEventData(document_id="doc-123", tags=["tag1", "tag2"])
|
||||
assert data.document_id == "doc-123"
|
||||
assert data.tags == ["tag1", "tag2"]
|
||||
|
||||
empty = RetainEventData()
|
||||
assert empty.document_id is None
|
||||
assert empty.tags is None
|
||||
|
||||
def test_retain_event_type_value(self):
|
||||
"""WebhookEventType.RETAIN_COMPLETED has the correct string value."""
|
||||
assert WebhookEventType.RETAIN_COMPLETED == "retain.completed"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fire_retain_webhook_queues_per_document(
|
||||
self, memory: MemoryEngine, webhook_manager: WebhookManager
|
||||
):
|
||||
"""_fire_retain_webhook queues one delivery task per content item."""
|
||||
bank_id = f"wh-retain-{uuid.uuid4().hex[:8]}"
|
||||
webhook_id = uuid.uuid4()
|
||||
|
||||
async with memory._pool.acquire() as conn:
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO webhooks (id, bank_id, url, secret, event_types, enabled, created_at, updated_at)
|
||||
VALUES ($1, $2, $3, NULL, $4, true, NOW(), NOW())
|
||||
""",
|
||||
webhook_id,
|
||||
bank_id,
|
||||
"https://example.com/retain-hook",
|
||||
["retain.completed"],
|
||||
)
|
||||
|
||||
try:
|
||||
contents = [
|
||||
{"content": "Alice works at Google", "document_id": "doc-1"},
|
||||
{"content": "Bob loves Python", "document_id": "doc-2"},
|
||||
]
|
||||
# Temporarily replace webhook manager on memory engine
|
||||
original_manager = memory._webhook_manager
|
||||
memory._webhook_manager = webhook_manager
|
||||
try:
|
||||
callback = memory._build_retain_outbox_callback(
|
||||
bank_id=bank_id,
|
||||
contents=contents,
|
||||
operation_id="test-op-123",
|
||||
)
|
||||
assert callback is not None
|
||||
async with memory._pool.acquire() as conn:
|
||||
await callback(conn)
|
||||
finally:
|
||||
memory._webhook_manager = original_manager
|
||||
|
||||
async with memory._pool.acquire() as conn:
|
||||
rows = await conn.fetch(
|
||||
"""
|
||||
SELECT task_payload
|
||||
FROM async_operations
|
||||
WHERE operation_type = 'webhook_delivery'
|
||||
AND bank_id = $1
|
||||
AND task_payload->>'event_type' = 'retain.completed'
|
||||
ORDER BY created_at
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
|
||||
assert len(rows) == 2
|
||||
payloads = []
|
||||
for row in rows:
|
||||
p = row["task_payload"]
|
||||
if isinstance(p, str):
|
||||
p = json.loads(p)
|
||||
payloads.append(p)
|
||||
|
||||
doc_ids_in_payloads = [json.loads(p["payload"]).get("data", {}).get("document_id") for p in payloads]
|
||||
assert "doc-1" in doc_ids_in_payloads
|
||||
assert "doc-2" in doc_ids_in_payloads
|
||||
finally:
|
||||
async with memory._pool.acquire() as conn:
|
||||
await conn.execute(
|
||||
"DELETE FROM async_operations WHERE operation_type = 'webhook_delivery' AND bank_id = $1",
|
||||
bank_id,
|
||||
)
|
||||
await conn.execute("DELETE FROM webhooks WHERE id = $1", webhook_id)
|
||||
@@ -48,11 +48,11 @@ async def pool(pg0_db_url):
|
||||
@pytest_asyncio.fixture
|
||||
async def clean_operations(pool):
|
||||
"""Clean up async_operations table before and after tests."""
|
||||
# Clean before test
|
||||
await pool.execute("DELETE FROM async_operations WHERE bank_id LIKE 'test-worker-%'")
|
||||
# Clean before test - covers both 'test-worker-' and 'test_worker_recovery' patterns
|
||||
await pool.execute("DELETE FROM async_operations WHERE bank_id LIKE 'test-worker-%' OR bank_id LIKE 'test_worker_%'")
|
||||
yield
|
||||
# Clean after test
|
||||
await pool.execute("DELETE FROM async_operations WHERE bank_id LIKE 'test-worker-%'")
|
||||
await pool.execute("DELETE FROM async_operations WHERE bank_id LIKE 'test-worker-%' OR bank_id LIKE 'test_worker_%'")
|
||||
|
||||
|
||||
class TestBrokerTaskBackend:
|
||||
@@ -268,32 +268,38 @@ class TestWorkerPoller:
|
||||
assert row["completed_at"] is not None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_executor_exception_does_not_crash_poller(self, pool, clean_operations):
|
||||
"""Test that unexpected exceptions from executor are caught and don't crash the poller.
|
||||
async def test_executor_exception_triggers_retry(self, pool, clean_operations):
|
||||
"""Test that exceptions from the executor trigger _retry_or_fail (not a crash).
|
||||
|
||||
If the executor raises an unexpected exception (which MemoryEngine.execute_task should NOT do,
|
||||
but could happen from schema setup or other infrastructure issues), the poller should catch it
|
||||
gracefully. Status remains 'processing' since neither executor nor poller handled it.
|
||||
When the executor re-raises an exception (as MemoryEngine.execute_task does for
|
||||
retryable task failures), the poller calls _retry_or_fail, which resets the task
|
||||
back to 'pending' and increments retry_count so it can be reclaimed.
|
||||
|
||||
This is the fix for the consolidation deadlock: previously submit_task was called
|
||||
with only a task_payload update, leaving status='processing' forever.
|
||||
"""
|
||||
from hindsight_api.worker import WorkerPoller
|
||||
from hindsight_api.worker.poller import ClaimedTask
|
||||
|
||||
# Create a pending task
|
||||
bank_id = f"test-worker-{uuid.uuid4().hex[:8]}"
|
||||
op_id = uuid.uuid4()
|
||||
payload = json.dumps({"type": "test_task", "operation_id": str(op_id), "bank_id": bank_id})
|
||||
payload = json.dumps({"type": "consolidation", "operation_id": str(op_id), "bank_id": bank_id})
|
||||
await pool.execute(
|
||||
"""
|
||||
INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload, worker_id)
|
||||
VALUES ($1, $2, 'test', 'processing', $3::jsonb, 'test-worker-1')
|
||||
INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload, worker_id, claimed_at)
|
||||
VALUES ($1, $2, 'consolidation', 'processing', $3::jsonb, 'test-worker-1', now())
|
||||
""",
|
||||
op_id,
|
||||
bank_id,
|
||||
payload,
|
||||
)
|
||||
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from hindsight_api.worker.exceptions import RetryTaskAt
|
||||
|
||||
async def failing_executor(task_dict):
|
||||
raise ValueError("Unexpected infrastructure failure")
|
||||
raise RetryTaskAt(retry_at=datetime.now(timezone.utc), message="TimeoutError during recall")
|
||||
|
||||
poller = WorkerPoller(
|
||||
pool=pool,
|
||||
@@ -301,35 +307,86 @@ class TestWorkerPoller:
|
||||
executor=failing_executor,
|
||||
)
|
||||
|
||||
# Execute - should catch exception without crashing
|
||||
task_dict = json.loads(payload)
|
||||
claimed_task = ClaimedTask(operation_id=str(op_id), task_dict=task_dict, schema=None)
|
||||
await poller.execute_task(claimed_task)
|
||||
|
||||
# Wait for background task to complete
|
||||
completed = await poller.wait_for_active_tasks(timeout=5.0)
|
||||
assert completed, "Task did not complete within timeout"
|
||||
|
||||
# Status stays 'processing' since the poller no longer manages status
|
||||
# Task must be reset to 'pending' with worker_id/claimed_at cleared — not left as
|
||||
# 'processing', which would cause a permanent deadlock via the NOT EXISTS guard.
|
||||
row = await pool.fetchrow(
|
||||
"SELECT status FROM async_operations WHERE operation_id = $1",
|
||||
"SELECT status, worker_id, claimed_at, retry_count FROM async_operations WHERE operation_id = $1",
|
||||
op_id,
|
||||
)
|
||||
assert row["status"] == "processing"
|
||||
assert row["status"] == "pending", (
|
||||
f"REGRESSION: Task status is '{row['status']}' instead of 'pending'. "
|
||||
"A task stuck in 'processing' after a retry causes a consolidation deadlock."
|
||||
)
|
||||
assert row["worker_id"] is None, "worker_id must be cleared on retry"
|
||||
assert row["claimed_at"] is None, "claimed_at must be cleared on retry"
|
||||
assert row["retry_count"] == 1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_executor_exception_marks_failed_immediately(self, pool, clean_operations):
|
||||
"""Test that a plain exception (not RetryTaskAt) permanently marks a task as 'failed'.
|
||||
|
||||
With the task-owned retry model, plain exceptions are non-retryable — the poller
|
||||
marks them as 'failed' immediately. Tasks that want to be retried must raise RetryTaskAt.
|
||||
"""
|
||||
from hindsight_api.worker import WorkerPoller
|
||||
from hindsight_api.worker.poller import ClaimedTask
|
||||
|
||||
bank_id = f"test-worker-{uuid.uuid4().hex[:8]}"
|
||||
op_id = uuid.uuid4()
|
||||
payload = json.dumps({"type": "consolidation", "operation_id": str(op_id), "bank_id": bank_id})
|
||||
await pool.execute(
|
||||
"""
|
||||
INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload, worker_id, claimed_at, retry_count)
|
||||
VALUES ($1, $2, 'consolidation', 'processing', $3::jsonb, 'test-worker-1', now(), 0)
|
||||
""",
|
||||
op_id,
|
||||
bank_id,
|
||||
payload,
|
||||
)
|
||||
|
||||
async def failing_executor(task_dict):
|
||||
raise ValueError("Non-retryable error")
|
||||
|
||||
poller = WorkerPoller(
|
||||
pool=pool,
|
||||
worker_id="test-worker-1",
|
||||
executor=failing_executor,
|
||||
)
|
||||
|
||||
task_dict = json.loads(payload)
|
||||
claimed_task = ClaimedTask(operation_id=str(op_id), task_dict=task_dict, schema=None)
|
||||
await poller.execute_task(claimed_task)
|
||||
|
||||
completed = await poller.wait_for_active_tasks(timeout=5.0)
|
||||
assert completed, "Task did not complete within timeout"
|
||||
|
||||
row = await pool.fetchrow(
|
||||
"SELECT status, error_message, retry_count FROM async_operations WHERE operation_id = $1",
|
||||
op_id,
|
||||
)
|
||||
assert row["status"] == "failed", (
|
||||
f"Expected 'failed' for plain exception, got '{row['status']}'"
|
||||
)
|
||||
assert row["error_message"] is not None
|
||||
assert row["retry_count"] == 0 # not incremented; plain exception = immediate fail
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_executor_failed_status_not_overridden(self, pool, clean_operations):
|
||||
"""REGRESSION TEST: Verify poller does NOT overwrite executor's 'failed' status to 'completed'.
|
||||
|
||||
This test catches the bug where the poller always called _mark_completed() after executor
|
||||
returned, overwriting the 'failed' status that the executor had already set.
|
||||
|
||||
Scenario:
|
||||
1. Executor catches an internal error and marks the operation as 'failed' in the DB
|
||||
2. Executor returns normally (does NOT re-raise) - this is how MemoryEngine.execute_task works
|
||||
This test covers the non-retryable failure path (e.g., file_convert_retain):
|
||||
1. Executor catches an internal error, marks the operation as 'failed' in the DB
|
||||
2. Executor returns normally (does NOT re-raise) — so no exception reaches the poller
|
||||
3. The poller must NOT overwrite the 'failed' status to 'completed'
|
||||
|
||||
With the old buggy code, this test would FAIL (status would be 'completed').
|
||||
Retryable failures re-raise instead (see test_executor_exception_triggers_retry).
|
||||
"""
|
||||
from hindsight_api.worker import WorkerPoller
|
||||
from hindsight_api.worker.poller import ClaimedTask
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
[package]
|
||||
name = "hindsight-cli"
|
||||
version = "0.4.13"
|
||||
version = "0.4.16"
|
||||
edition = "2021"
|
||||
authors = ["Hindsight Team"]
|
||||
description = "A beautiful CLI for Hindsight - semantic memory system"
|
||||
|
||||
@@ -300,6 +300,8 @@ impl ApiClient {
|
||||
offset.map(|o| o as i64),
|
||||
q,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
).await?;
|
||||
Ok(response.into_inner())
|
||||
})
|
||||
|
||||
@@ -387,6 +387,7 @@ pub fn retain(
|
||||
document_id: Some(doc_id.clone()),
|
||||
entities: None,
|
||||
tags: None,
|
||||
observation_scopes: None,
|
||||
};
|
||||
|
||||
let request = RetainRequest {
|
||||
|
||||
@@ -7,7 +7,7 @@ info:
|
||||
name: Apache 2.0
|
||||
url: https://www.apache.org/licenses/LICENSE-2.0.html
|
||||
title: Hindsight HTTP API
|
||||
version: 0.4.13
|
||||
version: 0.4.16
|
||||
servers:
|
||||
- url: /
|
||||
paths:
|
||||
@@ -1173,7 +1173,9 @@ paths:
|
||||
title: Bank Id
|
||||
type: string
|
||||
style: simple
|
||||
- explode: true
|
||||
- description: Case-insensitive substring filter on document ID (e.g. 'report'
|
||||
matches 'report-2024')
|
||||
explode: true
|
||||
in: query
|
||||
name: q
|
||||
required: false
|
||||
@@ -1181,6 +1183,28 @@ paths:
|
||||
nullable: true
|
||||
type: string
|
||||
style: form
|
||||
- description: Filter documents by tags
|
||||
explode: true
|
||||
in: query
|
||||
name: tags
|
||||
required: false
|
||||
schema:
|
||||
items:
|
||||
type: string
|
||||
nullable: true
|
||||
type: array
|
||||
style: form
|
||||
- description: "How to match tags: 'any', 'all', 'any_strict', 'all_strict'"
|
||||
explode: true
|
||||
in: query
|
||||
name: tags_match
|
||||
required: false
|
||||
schema:
|
||||
default: any_strict
|
||||
description: "How to match tags: 'any', 'all', 'any_strict', 'all_strict'"
|
||||
title: Tags Match
|
||||
type: string
|
||||
style: form
|
||||
- explode: true
|
||||
in: query
|
||||
name: limit
|
||||
@@ -2087,6 +2111,248 @@ paths:
|
||||
summary: Trigger consolidation
|
||||
tags:
|
||||
- Banks
|
||||
/v1/default/banks/{bank_id}/webhooks:
|
||||
get:
|
||||
description: List all webhooks registered for a bank.
|
||||
operationId: list_webhooks
|
||||
parameters:
|
||||
- explode: false
|
||||
in: path
|
||||
name: bank_id
|
||||
required: true
|
||||
schema:
|
||||
title: Bank Id
|
||||
type: string
|
||||
style: simple
|
||||
- explode: false
|
||||
in: header
|
||||
name: authorization
|
||||
required: false
|
||||
schema:
|
||||
nullable: true
|
||||
type: string
|
||||
style: simple
|
||||
responses:
|
||||
"200":
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
$ref: '#/components/schemas/WebhookListResponse'
|
||||
description: Successful Response
|
||||
"422":
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
$ref: '#/components/schemas/HTTPValidationError'
|
||||
description: Validation Error
|
||||
summary: List webhooks
|
||||
tags:
|
||||
- Webhooks
|
||||
post:
|
||||
description: Register a webhook endpoint to receive event notifications for
|
||||
this bank.
|
||||
operationId: create_webhook
|
||||
parameters:
|
||||
- explode: false
|
||||
in: path
|
||||
name: bank_id
|
||||
required: true
|
||||
schema:
|
||||
title: Bank Id
|
||||
type: string
|
||||
style: simple
|
||||
- explode: false
|
||||
in: header
|
||||
name: authorization
|
||||
required: false
|
||||
schema:
|
||||
nullable: true
|
||||
type: string
|
||||
style: simple
|
||||
requestBody:
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
$ref: '#/components/schemas/CreateWebhookRequest'
|
||||
required: true
|
||||
responses:
|
||||
"201":
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
$ref: '#/components/schemas/WebhookResponse'
|
||||
description: Successful Response
|
||||
"422":
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
$ref: '#/components/schemas/HTTPValidationError'
|
||||
description: Validation Error
|
||||
summary: Register webhook
|
||||
tags:
|
||||
- Webhooks
|
||||
/v1/default/banks/{bank_id}/webhooks/{webhook_id}:
|
||||
delete:
|
||||
description: Remove a registered webhook.
|
||||
operationId: delete_webhook
|
||||
parameters:
|
||||
- explode: false
|
||||
in: path
|
||||
name: bank_id
|
||||
required: true
|
||||
schema:
|
||||
title: Bank Id
|
||||
type: string
|
||||
style: simple
|
||||
- explode: false
|
||||
in: path
|
||||
name: webhook_id
|
||||
required: true
|
||||
schema:
|
||||
title: Webhook Id
|
||||
type: string
|
||||
style: simple
|
||||
- explode: false
|
||||
in: header
|
||||
name: authorization
|
||||
required: false
|
||||
schema:
|
||||
nullable: true
|
||||
type: string
|
||||
style: simple
|
||||
responses:
|
||||
"200":
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
$ref: '#/components/schemas/DeleteResponse'
|
||||
description: Successful Response
|
||||
"422":
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
$ref: '#/components/schemas/HTTPValidationError'
|
||||
description: Validation Error
|
||||
summary: Delete webhook
|
||||
tags:
|
||||
- Webhooks
|
||||
patch:
|
||||
description: Update one or more fields of a registered webhook. Only provided
|
||||
fields are changed.
|
||||
operationId: update_webhook
|
||||
parameters:
|
||||
- explode: false
|
||||
in: path
|
||||
name: bank_id
|
||||
required: true
|
||||
schema:
|
||||
title: Bank Id
|
||||
type: string
|
||||
style: simple
|
||||
- explode: false
|
||||
in: path
|
||||
name: webhook_id
|
||||
required: true
|
||||
schema:
|
||||
title: Webhook Id
|
||||
type: string
|
||||
style: simple
|
||||
- explode: false
|
||||
in: header
|
||||
name: authorization
|
||||
required: false
|
||||
schema:
|
||||
nullable: true
|
||||
type: string
|
||||
style: simple
|
||||
requestBody:
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
$ref: '#/components/schemas/UpdateWebhookRequest'
|
||||
required: true
|
||||
responses:
|
||||
"200":
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
$ref: '#/components/schemas/WebhookResponse'
|
||||
description: Successful Response
|
||||
"422":
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
$ref: '#/components/schemas/HTTPValidationError'
|
||||
description: Validation Error
|
||||
summary: Update webhook
|
||||
tags:
|
||||
- Webhooks
|
||||
/v1/default/banks/{bank_id}/webhooks/{webhook_id}/deliveries:
|
||||
get:
|
||||
description: Inspect delivery history for a webhook (useful for debugging).
|
||||
operationId: list_webhook_deliveries
|
||||
parameters:
|
||||
- explode: false
|
||||
in: path
|
||||
name: bank_id
|
||||
required: true
|
||||
schema:
|
||||
title: Bank Id
|
||||
type: string
|
||||
style: simple
|
||||
- explode: false
|
||||
in: path
|
||||
name: webhook_id
|
||||
required: true
|
||||
schema:
|
||||
title: Webhook Id
|
||||
type: string
|
||||
style: simple
|
||||
- description: Maximum number of deliveries to return
|
||||
explode: true
|
||||
in: query
|
||||
name: limit
|
||||
required: false
|
||||
schema:
|
||||
default: 50
|
||||
description: Maximum number of deliveries to return
|
||||
maximum: 200
|
||||
title: Limit
|
||||
type: integer
|
||||
style: form
|
||||
- description: Pagination cursor (created_at of last item)
|
||||
explode: true
|
||||
in: query
|
||||
name: cursor
|
||||
required: false
|
||||
schema:
|
||||
nullable: true
|
||||
type: string
|
||||
style: form
|
||||
- explode: false
|
||||
in: header
|
||||
name: authorization
|
||||
required: false
|
||||
schema:
|
||||
nullable: true
|
||||
type: string
|
||||
style: simple
|
||||
responses:
|
||||
"200":
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
$ref: '#/components/schemas/WebhookDeliveryListResponse'
|
||||
description: Successful Response
|
||||
"422":
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
$ref: '#/components/schemas/HTTPValidationError'
|
||||
description: Validation Error
|
||||
summary: List webhook deliveries
|
||||
tags:
|
||||
- Webhooks
|
||||
/v1/default/banks/{bank_id}/memories:
|
||||
delete:
|
||||
description: "Delete memory units for a memory bank. Optionally filter by type\
|
||||
@@ -2855,6 +3121,47 @@ components:
|
||||
required:
|
||||
- operation_id
|
||||
title: CreateMentalModelResponse
|
||||
CreateWebhookRequest:
|
||||
description: Request model for registering a webhook.
|
||||
example:
|
||||
event_types:
|
||||
- event_types
|
||||
- event_types
|
||||
secret: secret
|
||||
http_config:
|
||||
headers:
|
||||
key: headers
|
||||
method: POST
|
||||
timeout_seconds: 0
|
||||
params:
|
||||
key: params
|
||||
url: url
|
||||
enabled: true
|
||||
properties:
|
||||
url:
|
||||
description: HTTP(S) endpoint URL to deliver events to
|
||||
title: Url
|
||||
type: string
|
||||
secret:
|
||||
nullable: true
|
||||
type: string
|
||||
event_types:
|
||||
default:
|
||||
- consolidation.completed
|
||||
description: "List of event types to deliver. Currently supported: 'consolidation.completed'"
|
||||
items:
|
||||
type: string
|
||||
type: array
|
||||
enabled:
|
||||
default: true
|
||||
description: Whether this webhook is active
|
||||
title: Enabled
|
||||
type: boolean
|
||||
http_config:
|
||||
$ref: '#/components/schemas/WebhookHttpConfig'
|
||||
required:
|
||||
- url
|
||||
title: CreateWebhookRequest
|
||||
DeleteDocumentResponse:
|
||||
description: Response model for delete document endpoint.
|
||||
example:
|
||||
@@ -3477,9 +3784,7 @@ components:
|
||||
title: Content
|
||||
type: string
|
||||
timestamp:
|
||||
format: date-time
|
||||
nullable: true
|
||||
type: string
|
||||
$ref: '#/components/schemas/Timestamp'
|
||||
context:
|
||||
nullable: true
|
||||
type: string
|
||||
@@ -3500,6 +3805,8 @@ components:
|
||||
type: string
|
||||
nullable: true
|
||||
type: array
|
||||
observation_scopes:
|
||||
$ref: '#/components/schemas/ObservationScopes'
|
||||
required:
|
||||
- content
|
||||
title: MemoryItem
|
||||
@@ -4413,6 +4720,41 @@ components:
|
||||
trigger:
|
||||
$ref: '#/components/schemas/MentalModelTrigger'
|
||||
title: UpdateMentalModelRequest
|
||||
UpdateWebhookRequest:
|
||||
description: Request model for updating a webhook. Only provided fields are
|
||||
updated.
|
||||
example:
|
||||
event_types:
|
||||
- event_types
|
||||
- event_types
|
||||
secret: secret
|
||||
http_config:
|
||||
headers:
|
||||
key: headers
|
||||
method: POST
|
||||
timeout_seconds: 0
|
||||
params:
|
||||
key: params
|
||||
url: url
|
||||
enabled: true
|
||||
properties:
|
||||
url:
|
||||
nullable: true
|
||||
type: string
|
||||
secret:
|
||||
nullable: true
|
||||
type: string
|
||||
event_types:
|
||||
items:
|
||||
type: string
|
||||
nullable: true
|
||||
type: array
|
||||
enabled:
|
||||
nullable: true
|
||||
type: boolean
|
||||
http_config:
|
||||
$ref: '#/components/schemas/WebhookHttpConfig'
|
||||
title: UpdateWebhookRequest
|
||||
ValidationError:
|
||||
example:
|
||||
msg: msg
|
||||
@@ -4457,6 +4799,274 @@ components:
|
||||
- api_version
|
||||
- features
|
||||
title: VersionResponse
|
||||
WebhookDeliveryListResponse:
|
||||
description: Response model for listing webhook deliveries.
|
||||
example:
|
||||
next_cursor: next_cursor
|
||||
items:
|
||||
- last_response_body: last_response_body
|
||||
last_attempt_at: last_attempt_at
|
||||
created_at: created_at
|
||||
last_response_status: 6
|
||||
url: url
|
||||
event_type: event_type
|
||||
updated_at: updated_at
|
||||
webhook_id: webhook_id
|
||||
next_retry_at: next_retry_at
|
||||
id: id
|
||||
last_error: last_error
|
||||
status: status
|
||||
attempts: 0
|
||||
- last_response_body: last_response_body
|
||||
last_attempt_at: last_attempt_at
|
||||
created_at: created_at
|
||||
last_response_status: 6
|
||||
url: url
|
||||
event_type: event_type
|
||||
updated_at: updated_at
|
||||
webhook_id: webhook_id
|
||||
next_retry_at: next_retry_at
|
||||
id: id
|
||||
last_error: last_error
|
||||
status: status
|
||||
attempts: 0
|
||||
properties:
|
||||
items:
|
||||
items:
|
||||
$ref: '#/components/schemas/WebhookDeliveryResponse'
|
||||
type: array
|
||||
next_cursor:
|
||||
nullable: true
|
||||
type: string
|
||||
required:
|
||||
- items
|
||||
title: WebhookDeliveryListResponse
|
||||
WebhookDeliveryResponse:
|
||||
description: Response model for a webhook delivery record.
|
||||
example:
|
||||
last_response_body: last_response_body
|
||||
last_attempt_at: last_attempt_at
|
||||
created_at: created_at
|
||||
last_response_status: 6
|
||||
url: url
|
||||
event_type: event_type
|
||||
updated_at: updated_at
|
||||
webhook_id: webhook_id
|
||||
next_retry_at: next_retry_at
|
||||
id: id
|
||||
last_error: last_error
|
||||
status: status
|
||||
attempts: 0
|
||||
properties:
|
||||
id:
|
||||
title: Id
|
||||
type: string
|
||||
webhook_id:
|
||||
nullable: true
|
||||
type: string
|
||||
url:
|
||||
title: Url
|
||||
type: string
|
||||
event_type:
|
||||
title: Event Type
|
||||
type: string
|
||||
status:
|
||||
title: Status
|
||||
type: string
|
||||
attempts:
|
||||
title: Attempts
|
||||
type: integer
|
||||
next_retry_at:
|
||||
nullable: true
|
||||
type: string
|
||||
last_error:
|
||||
nullable: true
|
||||
type: string
|
||||
last_response_status:
|
||||
nullable: true
|
||||
type: integer
|
||||
last_response_body:
|
||||
nullable: true
|
||||
type: string
|
||||
last_attempt_at:
|
||||
nullable: true
|
||||
type: string
|
||||
created_at:
|
||||
nullable: true
|
||||
type: string
|
||||
updated_at:
|
||||
nullable: true
|
||||
type: string
|
||||
required:
|
||||
- attempts
|
||||
- event_type
|
||||
- id
|
||||
- status
|
||||
- url
|
||||
- webhook_id
|
||||
title: WebhookDeliveryResponse
|
||||
WebhookHttpConfig:
|
||||
description: HTTP delivery configuration for a webhook.
|
||||
example:
|
||||
headers:
|
||||
key: headers
|
||||
method: POST
|
||||
timeout_seconds: 0
|
||||
params:
|
||||
key: params
|
||||
properties:
|
||||
method:
|
||||
default: POST
|
||||
description: "HTTP method: GET or POST"
|
||||
title: Method
|
||||
type: string
|
||||
timeout_seconds:
|
||||
default: 30
|
||||
description: HTTP request timeout in seconds
|
||||
title: Timeout Seconds
|
||||
type: integer
|
||||
headers:
|
||||
additionalProperties:
|
||||
type: string
|
||||
description: Custom HTTP headers
|
||||
title: Headers
|
||||
params:
|
||||
additionalProperties:
|
||||
type: string
|
||||
description: Custom HTTP query parameters
|
||||
title: Params
|
||||
title: WebhookHttpConfig
|
||||
WebhookListResponse:
|
||||
description: Response model for listing webhooks.
|
||||
example:
|
||||
items:
|
||||
- event_types:
|
||||
- event_types
|
||||
- event_types
|
||||
updated_at: updated_at
|
||||
bank_id: bank_id
|
||||
created_at: created_at
|
||||
id: id
|
||||
secret: secret
|
||||
http_config:
|
||||
headers:
|
||||
key: headers
|
||||
method: POST
|
||||
timeout_seconds: 0
|
||||
params:
|
||||
key: params
|
||||
url: url
|
||||
enabled: true
|
||||
- event_types:
|
||||
- event_types
|
||||
- event_types
|
||||
updated_at: updated_at
|
||||
bank_id: bank_id
|
||||
created_at: created_at
|
||||
id: id
|
||||
secret: secret
|
||||
http_config:
|
||||
headers:
|
||||
key: headers
|
||||
method: POST
|
||||
timeout_seconds: 0
|
||||
params:
|
||||
key: params
|
||||
url: url
|
||||
enabled: true
|
||||
properties:
|
||||
items:
|
||||
items:
|
||||
$ref: '#/components/schemas/WebhookResponse'
|
||||
type: array
|
||||
required:
|
||||
- items
|
||||
title: WebhookListResponse
|
||||
WebhookResponse:
|
||||
description: Response model for a webhook.
|
||||
example:
|
||||
event_types:
|
||||
- event_types
|
||||
- event_types
|
||||
updated_at: updated_at
|
||||
bank_id: bank_id
|
||||
created_at: created_at
|
||||
id: id
|
||||
secret: secret
|
||||
http_config:
|
||||
headers:
|
||||
key: headers
|
||||
method: POST
|
||||
timeout_seconds: 0
|
||||
params:
|
||||
key: params
|
||||
url: url
|
||||
enabled: true
|
||||
properties:
|
||||
id:
|
||||
title: Id
|
||||
type: string
|
||||
bank_id:
|
||||
nullable: true
|
||||
type: string
|
||||
url:
|
||||
title: Url
|
||||
type: string
|
||||
secret:
|
||||
nullable: true
|
||||
type: string
|
||||
event_types:
|
||||
items:
|
||||
type: string
|
||||
type: array
|
||||
enabled:
|
||||
title: Enabled
|
||||
type: boolean
|
||||
http_config:
|
||||
$ref: '#/components/schemas/WebhookHttpConfig'
|
||||
created_at:
|
||||
nullable: true
|
||||
type: string
|
||||
updated_at:
|
||||
nullable: true
|
||||
type: string
|
||||
required:
|
||||
- bank_id
|
||||
- enabled
|
||||
- event_types
|
||||
- id
|
||||
- url
|
||||
title: WebhookResponse
|
||||
Timestamp:
|
||||
anyOf:
|
||||
- format: date-time
|
||||
type: string
|
||||
- type: string
|
||||
description: "When the content occurred. Accepts an ISO 8601 datetime string\
|
||||
\ (e.g. '2024-01-15T10:30:00Z'), null/omitted (defaults to now), or the special\
|
||||
\ string 'unset' to explicitly store without any timestamp (use this for timeless\
|
||||
\ content such as fictional documents or static reference material)."
|
||||
nullable: true
|
||||
title: Timestamp
|
||||
ObservationScopes:
|
||||
anyOf:
|
||||
- enum:
|
||||
- per_tag
|
||||
- combined
|
||||
- all_combinations
|
||||
type: string
|
||||
- items:
|
||||
items:
|
||||
type: string
|
||||
type: array
|
||||
type: array
|
||||
description: "How to scope observations during consolidation. 'per_tag' runs\
|
||||
\ one consolidation pass per individual tag, creating separate observations\
|
||||
\ for each tag. 'combined' (default) runs a single pass with all tags together.\
|
||||
\ A list of tag lists runs one pass per inner list, giving full control over\
|
||||
\ which combinations to use."
|
||||
nullable: true
|
||||
title: ObservationScopes
|
||||
ValidationError_loc_inner:
|
||||
anyOf:
|
||||
- type: string
|
||||
|
||||
@@ -3,7 +3,7 @@ Hindsight HTTP API
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
API version: 0.4.13
|
||||
API version: 0.4.16
|
||||
*/
|
||||
|
||||
// Code generated by OpenAPI Generator (https://openapi-generator.tech); DO NOT EDIT.
|
||||
|
||||
@@ -3,7 +3,7 @@ Hindsight HTTP API
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
API version: 0.4.13
|
||||
API version: 0.4.16
|
||||
*/
|
||||
|
||||
// Code generated by OpenAPI Generator (https://openapi-generator.tech); DO NOT EDIT.
|
||||
|
||||
@@ -3,7 +3,7 @@ Hindsight HTTP API
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
API version: 0.4.13
|
||||
API version: 0.4.16
|
||||
*/
|
||||
|
||||
// Code generated by OpenAPI Generator (https://openapi-generator.tech); DO NOT EDIT.
|
||||
@@ -17,6 +17,7 @@ import (
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
"reflect"
|
||||
)
|
||||
|
||||
|
||||
@@ -409,16 +410,31 @@ type ApiListDocumentsRequest struct {
|
||||
ApiService *DocumentsAPIService
|
||||
bankId string
|
||||
q *string
|
||||
tags *[]string
|
||||
tagsMatch *string
|
||||
limit *int32
|
||||
offset *int32
|
||||
authorization *string
|
||||
}
|
||||
|
||||
// Case-insensitive substring filter on document ID (e.g. 'report' matches 'report-2024')
|
||||
func (r ApiListDocumentsRequest) Q(q string) ApiListDocumentsRequest {
|
||||
r.q = &q
|
||||
return r
|
||||
}
|
||||
|
||||
// Filter documents by tags
|
||||
func (r ApiListDocumentsRequest) Tags(tags []string) ApiListDocumentsRequest {
|
||||
r.tags = &tags
|
||||
return r
|
||||
}
|
||||
|
||||
// How to match tags: 'any', 'all', 'any_strict', 'all_strict'
|
||||
func (r ApiListDocumentsRequest) TagsMatch(tagsMatch string) ApiListDocumentsRequest {
|
||||
r.tagsMatch = &tagsMatch
|
||||
return r
|
||||
}
|
||||
|
||||
func (r ApiListDocumentsRequest) Limit(limit int32) ApiListDocumentsRequest {
|
||||
r.limit = &limit
|
||||
return r
|
||||
@@ -480,6 +496,23 @@ func (a *DocumentsAPIService) ListDocumentsExecute(r ApiListDocumentsRequest) (*
|
||||
if r.q != nil {
|
||||
parameterAddToHeaderOrQuery(localVarQueryParams, "q", r.q, "form", "")
|
||||
}
|
||||
if r.tags != nil {
|
||||
t := *r.tags
|
||||
if reflect.TypeOf(t).Kind() == reflect.Slice {
|
||||
s := reflect.ValueOf(t)
|
||||
for i := 0; i < s.Len(); i++ {
|
||||
parameterAddToHeaderOrQuery(localVarQueryParams, "tags", s.Index(i).Interface(), "form", "multi")
|
||||
}
|
||||
} else {
|
||||
parameterAddToHeaderOrQuery(localVarQueryParams, "tags", t, "form", "multi")
|
||||
}
|
||||
}
|
||||
if r.tagsMatch != nil {
|
||||
parameterAddToHeaderOrQuery(localVarQueryParams, "tags_match", r.tagsMatch, "form", "")
|
||||
} else {
|
||||
var defaultValue string = "any_strict"
|
||||
r.tagsMatch = &defaultValue
|
||||
}
|
||||
if r.limit != nil {
|
||||
parameterAddToHeaderOrQuery(localVarQueryParams, "limit", r.limit, "form", "")
|
||||
} else {
|
||||
|
||||
@@ -3,7 +3,7 @@ Hindsight HTTP API
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
API version: 0.4.13
|
||||
API version: 0.4.16
|
||||
*/
|
||||
|
||||
// Code generated by OpenAPI Generator (https://openapi-generator.tech); DO NOT EDIT.
|
||||
|
||||
@@ -3,7 +3,7 @@ Hindsight HTTP API
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
API version: 0.4.13
|
||||
API version: 0.4.16
|
||||
*/
|
||||
|
||||
// Code generated by OpenAPI Generator (https://openapi-generator.tech); DO NOT EDIT.
|
||||
|
||||
@@ -3,7 +3,7 @@ Hindsight HTTP API
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
API version: 0.4.13
|
||||
API version: 0.4.16
|
||||
*/
|
||||
|
||||
// Code generated by OpenAPI Generator (https://openapi-generator.tech); DO NOT EDIT.
|
||||
|
||||
@@ -3,7 +3,7 @@ Hindsight HTTP API
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
API version: 0.4.13
|
||||
API version: 0.4.16
|
||||
*/
|
||||
|
||||
// Code generated by OpenAPI Generator (https://openapi-generator.tech); DO NOT EDIT.
|
||||
|
||||
@@ -3,7 +3,7 @@ Hindsight HTTP API
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
API version: 0.4.13
|
||||
API version: 0.4.16
|
||||
*/
|
||||
|
||||
// Code generated by OpenAPI Generator (https://openapi-generator.tech); DO NOT EDIT.
|
||||
|
||||
@@ -3,7 +3,7 @@ Hindsight HTTP API
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
API version: 0.4.13
|
||||
API version: 0.4.16
|
||||
*/
|
||||
|
||||
// Code generated by OpenAPI Generator (https://openapi-generator.tech); DO NOT EDIT.
|
||||
|
||||
@@ -0,0 +1,691 @@
|
||||
/*
|
||||
Hindsight HTTP API
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
API version: 0.4.16
|
||||
*/
|
||||
|
||||
// Code generated by OpenAPI Generator (https://openapi-generator.tech); DO NOT EDIT.
|
||||
|
||||
package hindsight
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
)
|
||||
|
||||
|
||||
// WebhooksAPIService WebhooksAPI service
|
||||
type WebhooksAPIService service
|
||||
|
||||
type ApiCreateWebhookRequest struct {
|
||||
ctx context.Context
|
||||
ApiService *WebhooksAPIService
|
||||
bankId string
|
||||
createWebhookRequest *CreateWebhookRequest
|
||||
authorization *string
|
||||
}
|
||||
|
||||
func (r ApiCreateWebhookRequest) CreateWebhookRequest(createWebhookRequest CreateWebhookRequest) ApiCreateWebhookRequest {
|
||||
r.createWebhookRequest = &createWebhookRequest
|
||||
return r
|
||||
}
|
||||
|
||||
func (r ApiCreateWebhookRequest) Authorization(authorization string) ApiCreateWebhookRequest {
|
||||
r.authorization = &authorization
|
||||
return r
|
||||
}
|
||||
|
||||
func (r ApiCreateWebhookRequest) Execute() (*WebhookResponse, *http.Response, error) {
|
||||
return r.ApiService.CreateWebhookExecute(r)
|
||||
}
|
||||
|
||||
/*
|
||||
CreateWebhook Register webhook
|
||||
|
||||
Register a webhook endpoint to receive event notifications for this bank.
|
||||
|
||||
@param ctx context.Context - for authentication, logging, cancellation, deadlines, tracing, etc. Passed from http.Request or context.Background().
|
||||
@param bankId
|
||||
@return ApiCreateWebhookRequest
|
||||
*/
|
||||
func (a *WebhooksAPIService) CreateWebhook(ctx context.Context, bankId string) ApiCreateWebhookRequest {
|
||||
return ApiCreateWebhookRequest{
|
||||
ApiService: a,
|
||||
ctx: ctx,
|
||||
bankId: bankId,
|
||||
}
|
||||
}
|
||||
|
||||
// Execute executes the request
|
||||
// @return WebhookResponse
|
||||
func (a *WebhooksAPIService) CreateWebhookExecute(r ApiCreateWebhookRequest) (*WebhookResponse, *http.Response, error) {
|
||||
var (
|
||||
localVarHTTPMethod = http.MethodPost
|
||||
localVarPostBody interface{}
|
||||
formFiles []formFile
|
||||
localVarReturnValue *WebhookResponse
|
||||
)
|
||||
|
||||
localBasePath, err := a.client.cfg.ServerURLWithContext(r.ctx, "WebhooksAPIService.CreateWebhook")
|
||||
if err != nil {
|
||||
return localVarReturnValue, nil, &GenericOpenAPIError{error: err.Error()}
|
||||
}
|
||||
|
||||
localVarPath := localBasePath + "/v1/default/banks/{bank_id}/webhooks"
|
||||
localVarPath = strings.Replace(localVarPath, "{"+"bank_id"+"}", url.PathEscape(parameterValueToString(r.bankId, "bankId")), -1)
|
||||
|
||||
localVarHeaderParams := make(map[string]string)
|
||||
localVarQueryParams := url.Values{}
|
||||
localVarFormParams := url.Values{}
|
||||
if r.createWebhookRequest == nil {
|
||||
return localVarReturnValue, nil, reportError("createWebhookRequest is required and must be specified")
|
||||
}
|
||||
|
||||
// to determine the Content-Type header
|
||||
localVarHTTPContentTypes := []string{"application/json"}
|
||||
|
||||
// set Content-Type header
|
||||
localVarHTTPContentType := selectHeaderContentType(localVarHTTPContentTypes)
|
||||
if localVarHTTPContentType != "" {
|
||||
localVarHeaderParams["Content-Type"] = localVarHTTPContentType
|
||||
}
|
||||
|
||||
// to determine the Accept header
|
||||
localVarHTTPHeaderAccepts := []string{"application/json"}
|
||||
|
||||
// set Accept header
|
||||
localVarHTTPHeaderAccept := selectHeaderAccept(localVarHTTPHeaderAccepts)
|
||||
if localVarHTTPHeaderAccept != "" {
|
||||
localVarHeaderParams["Accept"] = localVarHTTPHeaderAccept
|
||||
}
|
||||
if r.authorization != nil {
|
||||
parameterAddToHeaderOrQuery(localVarHeaderParams, "authorization", r.authorization, "simple", "")
|
||||
}
|
||||
// body params
|
||||
localVarPostBody = r.createWebhookRequest
|
||||
req, err := a.client.prepareRequest(r.ctx, localVarPath, localVarHTTPMethod, localVarPostBody, localVarHeaderParams, localVarQueryParams, localVarFormParams, formFiles)
|
||||
if err != nil {
|
||||
return localVarReturnValue, nil, err
|
||||
}
|
||||
|
||||
localVarHTTPResponse, err := a.client.callAPI(req)
|
||||
if err != nil || localVarHTTPResponse == nil {
|
||||
return localVarReturnValue, localVarHTTPResponse, err
|
||||
}
|
||||
|
||||
localVarBody, err := io.ReadAll(localVarHTTPResponse.Body)
|
||||
localVarHTTPResponse.Body.Close()
|
||||
localVarHTTPResponse.Body = io.NopCloser(bytes.NewBuffer(localVarBody))
|
||||
if err != nil {
|
||||
return localVarReturnValue, localVarHTTPResponse, err
|
||||
}
|
||||
|
||||
if localVarHTTPResponse.StatusCode >= 300 {
|
||||
newErr := &GenericOpenAPIError{
|
||||
body: localVarBody,
|
||||
error: localVarHTTPResponse.Status,
|
||||
}
|
||||
if localVarHTTPResponse.StatusCode == 422 {
|
||||
var v HTTPValidationError
|
||||
err = a.client.decode(&v, localVarBody, localVarHTTPResponse.Header.Get("Content-Type"))
|
||||
if err != nil {
|
||||
newErr.error = err.Error()
|
||||
return localVarReturnValue, localVarHTTPResponse, newErr
|
||||
}
|
||||
newErr.error = formatErrorMessage(localVarHTTPResponse.Status, &v)
|
||||
newErr.model = v
|
||||
}
|
||||
return localVarReturnValue, localVarHTTPResponse, newErr
|
||||
}
|
||||
|
||||
err = a.client.decode(&localVarReturnValue, localVarBody, localVarHTTPResponse.Header.Get("Content-Type"))
|
||||
if err != nil {
|
||||
newErr := &GenericOpenAPIError{
|
||||
body: localVarBody,
|
||||
error: err.Error(),
|
||||
}
|
||||
return localVarReturnValue, localVarHTTPResponse, newErr
|
||||
}
|
||||
|
||||
return localVarReturnValue, localVarHTTPResponse, nil
|
||||
}
|
||||
|
||||
type ApiDeleteWebhookRequest struct {
|
||||
ctx context.Context
|
||||
ApiService *WebhooksAPIService
|
||||
bankId string
|
||||
webhookId string
|
||||
authorization *string
|
||||
}
|
||||
|
||||
func (r ApiDeleteWebhookRequest) Authorization(authorization string) ApiDeleteWebhookRequest {
|
||||
r.authorization = &authorization
|
||||
return r
|
||||
}
|
||||
|
||||
func (r ApiDeleteWebhookRequest) Execute() (*DeleteResponse, *http.Response, error) {
|
||||
return r.ApiService.DeleteWebhookExecute(r)
|
||||
}
|
||||
|
||||
/*
|
||||
DeleteWebhook Delete webhook
|
||||
|
||||
Remove a registered webhook.
|
||||
|
||||
@param ctx context.Context - for authentication, logging, cancellation, deadlines, tracing, etc. Passed from http.Request or context.Background().
|
||||
@param bankId
|
||||
@param webhookId
|
||||
@return ApiDeleteWebhookRequest
|
||||
*/
|
||||
func (a *WebhooksAPIService) DeleteWebhook(ctx context.Context, bankId string, webhookId string) ApiDeleteWebhookRequest {
|
||||
return ApiDeleteWebhookRequest{
|
||||
ApiService: a,
|
||||
ctx: ctx,
|
||||
bankId: bankId,
|
||||
webhookId: webhookId,
|
||||
}
|
||||
}
|
||||
|
||||
// Execute executes the request
|
||||
// @return DeleteResponse
|
||||
func (a *WebhooksAPIService) DeleteWebhookExecute(r ApiDeleteWebhookRequest) (*DeleteResponse, *http.Response, error) {
|
||||
var (
|
||||
localVarHTTPMethod = http.MethodDelete
|
||||
localVarPostBody interface{}
|
||||
formFiles []formFile
|
||||
localVarReturnValue *DeleteResponse
|
||||
)
|
||||
|
||||
localBasePath, err := a.client.cfg.ServerURLWithContext(r.ctx, "WebhooksAPIService.DeleteWebhook")
|
||||
if err != nil {
|
||||
return localVarReturnValue, nil, &GenericOpenAPIError{error: err.Error()}
|
||||
}
|
||||
|
||||
localVarPath := localBasePath + "/v1/default/banks/{bank_id}/webhooks/{webhook_id}"
|
||||
localVarPath = strings.Replace(localVarPath, "{"+"bank_id"+"}", url.PathEscape(parameterValueToString(r.bankId, "bankId")), -1)
|
||||
localVarPath = strings.Replace(localVarPath, "{"+"webhook_id"+"}", url.PathEscape(parameterValueToString(r.webhookId, "webhookId")), -1)
|
||||
|
||||
localVarHeaderParams := make(map[string]string)
|
||||
localVarQueryParams := url.Values{}
|
||||
localVarFormParams := url.Values{}
|
||||
|
||||
// to determine the Content-Type header
|
||||
localVarHTTPContentTypes := []string{}
|
||||
|
||||
// set Content-Type header
|
||||
localVarHTTPContentType := selectHeaderContentType(localVarHTTPContentTypes)
|
||||
if localVarHTTPContentType != "" {
|
||||
localVarHeaderParams["Content-Type"] = localVarHTTPContentType
|
||||
}
|
||||
|
||||
// to determine the Accept header
|
||||
localVarHTTPHeaderAccepts := []string{"application/json"}
|
||||
|
||||
// set Accept header
|
||||
localVarHTTPHeaderAccept := selectHeaderAccept(localVarHTTPHeaderAccepts)
|
||||
if localVarHTTPHeaderAccept != "" {
|
||||
localVarHeaderParams["Accept"] = localVarHTTPHeaderAccept
|
||||
}
|
||||
if r.authorization != nil {
|
||||
parameterAddToHeaderOrQuery(localVarHeaderParams, "authorization", r.authorization, "simple", "")
|
||||
}
|
||||
req, err := a.client.prepareRequest(r.ctx, localVarPath, localVarHTTPMethod, localVarPostBody, localVarHeaderParams, localVarQueryParams, localVarFormParams, formFiles)
|
||||
if err != nil {
|
||||
return localVarReturnValue, nil, err
|
||||
}
|
||||
|
||||
localVarHTTPResponse, err := a.client.callAPI(req)
|
||||
if err != nil || localVarHTTPResponse == nil {
|
||||
return localVarReturnValue, localVarHTTPResponse, err
|
||||
}
|
||||
|
||||
localVarBody, err := io.ReadAll(localVarHTTPResponse.Body)
|
||||
localVarHTTPResponse.Body.Close()
|
||||
localVarHTTPResponse.Body = io.NopCloser(bytes.NewBuffer(localVarBody))
|
||||
if err != nil {
|
||||
return localVarReturnValue, localVarHTTPResponse, err
|
||||
}
|
||||
|
||||
if localVarHTTPResponse.StatusCode >= 300 {
|
||||
newErr := &GenericOpenAPIError{
|
||||
body: localVarBody,
|
||||
error: localVarHTTPResponse.Status,
|
||||
}
|
||||
if localVarHTTPResponse.StatusCode == 422 {
|
||||
var v HTTPValidationError
|
||||
err = a.client.decode(&v, localVarBody, localVarHTTPResponse.Header.Get("Content-Type"))
|
||||
if err != nil {
|
||||
newErr.error = err.Error()
|
||||
return localVarReturnValue, localVarHTTPResponse, newErr
|
||||
}
|
||||
newErr.error = formatErrorMessage(localVarHTTPResponse.Status, &v)
|
||||
newErr.model = v
|
||||
}
|
||||
return localVarReturnValue, localVarHTTPResponse, newErr
|
||||
}
|
||||
|
||||
err = a.client.decode(&localVarReturnValue, localVarBody, localVarHTTPResponse.Header.Get("Content-Type"))
|
||||
if err != nil {
|
||||
newErr := &GenericOpenAPIError{
|
||||
body: localVarBody,
|
||||
error: err.Error(),
|
||||
}
|
||||
return localVarReturnValue, localVarHTTPResponse, newErr
|
||||
}
|
||||
|
||||
return localVarReturnValue, localVarHTTPResponse, nil
|
||||
}
|
||||
|
||||
type ApiListWebhookDeliveriesRequest struct {
|
||||
ctx context.Context
|
||||
ApiService *WebhooksAPIService
|
||||
bankId string
|
||||
webhookId string
|
||||
limit *int32
|
||||
cursor *string
|
||||
authorization *string
|
||||
}
|
||||
|
||||
// Maximum number of deliveries to return
|
||||
func (r ApiListWebhookDeliveriesRequest) Limit(limit int32) ApiListWebhookDeliveriesRequest {
|
||||
r.limit = &limit
|
||||
return r
|
||||
}
|
||||
|
||||
// Pagination cursor (created_at of last item)
|
||||
func (r ApiListWebhookDeliveriesRequest) Cursor(cursor string) ApiListWebhookDeliveriesRequest {
|
||||
r.cursor = &cursor
|
||||
return r
|
||||
}
|
||||
|
||||
func (r ApiListWebhookDeliveriesRequest) Authorization(authorization string) ApiListWebhookDeliveriesRequest {
|
||||
r.authorization = &authorization
|
||||
return r
|
||||
}
|
||||
|
||||
func (r ApiListWebhookDeliveriesRequest) Execute() (*WebhookDeliveryListResponse, *http.Response, error) {
|
||||
return r.ApiService.ListWebhookDeliveriesExecute(r)
|
||||
}
|
||||
|
||||
/*
|
||||
ListWebhookDeliveries List webhook deliveries
|
||||
|
||||
Inspect delivery history for a webhook (useful for debugging).
|
||||
|
||||
@param ctx context.Context - for authentication, logging, cancellation, deadlines, tracing, etc. Passed from http.Request or context.Background().
|
||||
@param bankId
|
||||
@param webhookId
|
||||
@return ApiListWebhookDeliveriesRequest
|
||||
*/
|
||||
func (a *WebhooksAPIService) ListWebhookDeliveries(ctx context.Context, bankId string, webhookId string) ApiListWebhookDeliveriesRequest {
|
||||
return ApiListWebhookDeliveriesRequest{
|
||||
ApiService: a,
|
||||
ctx: ctx,
|
||||
bankId: bankId,
|
||||
webhookId: webhookId,
|
||||
}
|
||||
}
|
||||
|
||||
// Execute executes the request
|
||||
// @return WebhookDeliveryListResponse
|
||||
func (a *WebhooksAPIService) ListWebhookDeliveriesExecute(r ApiListWebhookDeliveriesRequest) (*WebhookDeliveryListResponse, *http.Response, error) {
|
||||
var (
|
||||
localVarHTTPMethod = http.MethodGet
|
||||
localVarPostBody interface{}
|
||||
formFiles []formFile
|
||||
localVarReturnValue *WebhookDeliveryListResponse
|
||||
)
|
||||
|
||||
localBasePath, err := a.client.cfg.ServerURLWithContext(r.ctx, "WebhooksAPIService.ListWebhookDeliveries")
|
||||
if err != nil {
|
||||
return localVarReturnValue, nil, &GenericOpenAPIError{error: err.Error()}
|
||||
}
|
||||
|
||||
localVarPath := localBasePath + "/v1/default/banks/{bank_id}/webhooks/{webhook_id}/deliveries"
|
||||
localVarPath = strings.Replace(localVarPath, "{"+"bank_id"+"}", url.PathEscape(parameterValueToString(r.bankId, "bankId")), -1)
|
||||
localVarPath = strings.Replace(localVarPath, "{"+"webhook_id"+"}", url.PathEscape(parameterValueToString(r.webhookId, "webhookId")), -1)
|
||||
|
||||
localVarHeaderParams := make(map[string]string)
|
||||
localVarQueryParams := url.Values{}
|
||||
localVarFormParams := url.Values{}
|
||||
|
||||
if r.limit != nil {
|
||||
parameterAddToHeaderOrQuery(localVarQueryParams, "limit", r.limit, "form", "")
|
||||
} else {
|
||||
var defaultValue int32 = 50
|
||||
r.limit = &defaultValue
|
||||
}
|
||||
if r.cursor != nil {
|
||||
parameterAddToHeaderOrQuery(localVarQueryParams, "cursor", r.cursor, "form", "")
|
||||
}
|
||||
// to determine the Content-Type header
|
||||
localVarHTTPContentTypes := []string{}
|
||||
|
||||
// set Content-Type header
|
||||
localVarHTTPContentType := selectHeaderContentType(localVarHTTPContentTypes)
|
||||
if localVarHTTPContentType != "" {
|
||||
localVarHeaderParams["Content-Type"] = localVarHTTPContentType
|
||||
}
|
||||
|
||||
// to determine the Accept header
|
||||
localVarHTTPHeaderAccepts := []string{"application/json"}
|
||||
|
||||
// set Accept header
|
||||
localVarHTTPHeaderAccept := selectHeaderAccept(localVarHTTPHeaderAccepts)
|
||||
if localVarHTTPHeaderAccept != "" {
|
||||
localVarHeaderParams["Accept"] = localVarHTTPHeaderAccept
|
||||
}
|
||||
if r.authorization != nil {
|
||||
parameterAddToHeaderOrQuery(localVarHeaderParams, "authorization", r.authorization, "simple", "")
|
||||
}
|
||||
req, err := a.client.prepareRequest(r.ctx, localVarPath, localVarHTTPMethod, localVarPostBody, localVarHeaderParams, localVarQueryParams, localVarFormParams, formFiles)
|
||||
if err != nil {
|
||||
return localVarReturnValue, nil, err
|
||||
}
|
||||
|
||||
localVarHTTPResponse, err := a.client.callAPI(req)
|
||||
if err != nil || localVarHTTPResponse == nil {
|
||||
return localVarReturnValue, localVarHTTPResponse, err
|
||||
}
|
||||
|
||||
localVarBody, err := io.ReadAll(localVarHTTPResponse.Body)
|
||||
localVarHTTPResponse.Body.Close()
|
||||
localVarHTTPResponse.Body = io.NopCloser(bytes.NewBuffer(localVarBody))
|
||||
if err != nil {
|
||||
return localVarReturnValue, localVarHTTPResponse, err
|
||||
}
|
||||
|
||||
if localVarHTTPResponse.StatusCode >= 300 {
|
||||
newErr := &GenericOpenAPIError{
|
||||
body: localVarBody,
|
||||
error: localVarHTTPResponse.Status,
|
||||
}
|
||||
if localVarHTTPResponse.StatusCode == 422 {
|
||||
var v HTTPValidationError
|
||||
err = a.client.decode(&v, localVarBody, localVarHTTPResponse.Header.Get("Content-Type"))
|
||||
if err != nil {
|
||||
newErr.error = err.Error()
|
||||
return localVarReturnValue, localVarHTTPResponse, newErr
|
||||
}
|
||||
newErr.error = formatErrorMessage(localVarHTTPResponse.Status, &v)
|
||||
newErr.model = v
|
||||
}
|
||||
return localVarReturnValue, localVarHTTPResponse, newErr
|
||||
}
|
||||
|
||||
err = a.client.decode(&localVarReturnValue, localVarBody, localVarHTTPResponse.Header.Get("Content-Type"))
|
||||
if err != nil {
|
||||
newErr := &GenericOpenAPIError{
|
||||
body: localVarBody,
|
||||
error: err.Error(),
|
||||
}
|
||||
return localVarReturnValue, localVarHTTPResponse, newErr
|
||||
}
|
||||
|
||||
return localVarReturnValue, localVarHTTPResponse, nil
|
||||
}
|
||||
|
||||
type ApiListWebhooksRequest struct {
|
||||
ctx context.Context
|
||||
ApiService *WebhooksAPIService
|
||||
bankId string
|
||||
authorization *string
|
||||
}
|
||||
|
||||
func (r ApiListWebhooksRequest) Authorization(authorization string) ApiListWebhooksRequest {
|
||||
r.authorization = &authorization
|
||||
return r
|
||||
}
|
||||
|
||||
func (r ApiListWebhooksRequest) Execute() (*WebhookListResponse, *http.Response, error) {
|
||||
return r.ApiService.ListWebhooksExecute(r)
|
||||
}
|
||||
|
||||
/*
|
||||
ListWebhooks List webhooks
|
||||
|
||||
List all webhooks registered for a bank.
|
||||
|
||||
@param ctx context.Context - for authentication, logging, cancellation, deadlines, tracing, etc. Passed from http.Request or context.Background().
|
||||
@param bankId
|
||||
@return ApiListWebhooksRequest
|
||||
*/
|
||||
func (a *WebhooksAPIService) ListWebhooks(ctx context.Context, bankId string) ApiListWebhooksRequest {
|
||||
return ApiListWebhooksRequest{
|
||||
ApiService: a,
|
||||
ctx: ctx,
|
||||
bankId: bankId,
|
||||
}
|
||||
}
|
||||
|
||||
// Execute executes the request
|
||||
// @return WebhookListResponse
|
||||
func (a *WebhooksAPIService) ListWebhooksExecute(r ApiListWebhooksRequest) (*WebhookListResponse, *http.Response, error) {
|
||||
var (
|
||||
localVarHTTPMethod = http.MethodGet
|
||||
localVarPostBody interface{}
|
||||
formFiles []formFile
|
||||
localVarReturnValue *WebhookListResponse
|
||||
)
|
||||
|
||||
localBasePath, err := a.client.cfg.ServerURLWithContext(r.ctx, "WebhooksAPIService.ListWebhooks")
|
||||
if err != nil {
|
||||
return localVarReturnValue, nil, &GenericOpenAPIError{error: err.Error()}
|
||||
}
|
||||
|
||||
localVarPath := localBasePath + "/v1/default/banks/{bank_id}/webhooks"
|
||||
localVarPath = strings.Replace(localVarPath, "{"+"bank_id"+"}", url.PathEscape(parameterValueToString(r.bankId, "bankId")), -1)
|
||||
|
||||
localVarHeaderParams := make(map[string]string)
|
||||
localVarQueryParams := url.Values{}
|
||||
localVarFormParams := url.Values{}
|
||||
|
||||
// to determine the Content-Type header
|
||||
localVarHTTPContentTypes := []string{}
|
||||
|
||||
// set Content-Type header
|
||||
localVarHTTPContentType := selectHeaderContentType(localVarHTTPContentTypes)
|
||||
if localVarHTTPContentType != "" {
|
||||
localVarHeaderParams["Content-Type"] = localVarHTTPContentType
|
||||
}
|
||||
|
||||
// to determine the Accept header
|
||||
localVarHTTPHeaderAccepts := []string{"application/json"}
|
||||
|
||||
// set Accept header
|
||||
localVarHTTPHeaderAccept := selectHeaderAccept(localVarHTTPHeaderAccepts)
|
||||
if localVarHTTPHeaderAccept != "" {
|
||||
localVarHeaderParams["Accept"] = localVarHTTPHeaderAccept
|
||||
}
|
||||
if r.authorization != nil {
|
||||
parameterAddToHeaderOrQuery(localVarHeaderParams, "authorization", r.authorization, "simple", "")
|
||||
}
|
||||
req, err := a.client.prepareRequest(r.ctx, localVarPath, localVarHTTPMethod, localVarPostBody, localVarHeaderParams, localVarQueryParams, localVarFormParams, formFiles)
|
||||
if err != nil {
|
||||
return localVarReturnValue, nil, err
|
||||
}
|
||||
|
||||
localVarHTTPResponse, err := a.client.callAPI(req)
|
||||
if err != nil || localVarHTTPResponse == nil {
|
||||
return localVarReturnValue, localVarHTTPResponse, err
|
||||
}
|
||||
|
||||
localVarBody, err := io.ReadAll(localVarHTTPResponse.Body)
|
||||
localVarHTTPResponse.Body.Close()
|
||||
localVarHTTPResponse.Body = io.NopCloser(bytes.NewBuffer(localVarBody))
|
||||
if err != nil {
|
||||
return localVarReturnValue, localVarHTTPResponse, err
|
||||
}
|
||||
|
||||
if localVarHTTPResponse.StatusCode >= 300 {
|
||||
newErr := &GenericOpenAPIError{
|
||||
body: localVarBody,
|
||||
error: localVarHTTPResponse.Status,
|
||||
}
|
||||
if localVarHTTPResponse.StatusCode == 422 {
|
||||
var v HTTPValidationError
|
||||
err = a.client.decode(&v, localVarBody, localVarHTTPResponse.Header.Get("Content-Type"))
|
||||
if err != nil {
|
||||
newErr.error = err.Error()
|
||||
return localVarReturnValue, localVarHTTPResponse, newErr
|
||||
}
|
||||
newErr.error = formatErrorMessage(localVarHTTPResponse.Status, &v)
|
||||
newErr.model = v
|
||||
}
|
||||
return localVarReturnValue, localVarHTTPResponse, newErr
|
||||
}
|
||||
|
||||
err = a.client.decode(&localVarReturnValue, localVarBody, localVarHTTPResponse.Header.Get("Content-Type"))
|
||||
if err != nil {
|
||||
newErr := &GenericOpenAPIError{
|
||||
body: localVarBody,
|
||||
error: err.Error(),
|
||||
}
|
||||
return localVarReturnValue, localVarHTTPResponse, newErr
|
||||
}
|
||||
|
||||
return localVarReturnValue, localVarHTTPResponse, nil
|
||||
}
|
||||
|
||||
type ApiUpdateWebhookRequest struct {
|
||||
ctx context.Context
|
||||
ApiService *WebhooksAPIService
|
||||
bankId string
|
||||
webhookId string
|
||||
updateWebhookRequest *UpdateWebhookRequest
|
||||
authorization *string
|
||||
}
|
||||
|
||||
func (r ApiUpdateWebhookRequest) UpdateWebhookRequest(updateWebhookRequest UpdateWebhookRequest) ApiUpdateWebhookRequest {
|
||||
r.updateWebhookRequest = &updateWebhookRequest
|
||||
return r
|
||||
}
|
||||
|
||||
func (r ApiUpdateWebhookRequest) Authorization(authorization string) ApiUpdateWebhookRequest {
|
||||
r.authorization = &authorization
|
||||
return r
|
||||
}
|
||||
|
||||
func (r ApiUpdateWebhookRequest) Execute() (*WebhookResponse, *http.Response, error) {
|
||||
return r.ApiService.UpdateWebhookExecute(r)
|
||||
}
|
||||
|
||||
/*
|
||||
UpdateWebhook Update webhook
|
||||
|
||||
Update one or more fields of a registered webhook. Only provided fields are changed.
|
||||
|
||||
@param ctx context.Context - for authentication, logging, cancellation, deadlines, tracing, etc. Passed from http.Request or context.Background().
|
||||
@param bankId
|
||||
@param webhookId
|
||||
@return ApiUpdateWebhookRequest
|
||||
*/
|
||||
func (a *WebhooksAPIService) UpdateWebhook(ctx context.Context, bankId string, webhookId string) ApiUpdateWebhookRequest {
|
||||
return ApiUpdateWebhookRequest{
|
||||
ApiService: a,
|
||||
ctx: ctx,
|
||||
bankId: bankId,
|
||||
webhookId: webhookId,
|
||||
}
|
||||
}
|
||||
|
||||
// Execute executes the request
|
||||
// @return WebhookResponse
|
||||
func (a *WebhooksAPIService) UpdateWebhookExecute(r ApiUpdateWebhookRequest) (*WebhookResponse, *http.Response, error) {
|
||||
var (
|
||||
localVarHTTPMethod = http.MethodPatch
|
||||
localVarPostBody interface{}
|
||||
formFiles []formFile
|
||||
localVarReturnValue *WebhookResponse
|
||||
)
|
||||
|
||||
localBasePath, err := a.client.cfg.ServerURLWithContext(r.ctx, "WebhooksAPIService.UpdateWebhook")
|
||||
if err != nil {
|
||||
return localVarReturnValue, nil, &GenericOpenAPIError{error: err.Error()}
|
||||
}
|
||||
|
||||
localVarPath := localBasePath + "/v1/default/banks/{bank_id}/webhooks/{webhook_id}"
|
||||
localVarPath = strings.Replace(localVarPath, "{"+"bank_id"+"}", url.PathEscape(parameterValueToString(r.bankId, "bankId")), -1)
|
||||
localVarPath = strings.Replace(localVarPath, "{"+"webhook_id"+"}", url.PathEscape(parameterValueToString(r.webhookId, "webhookId")), -1)
|
||||
|
||||
localVarHeaderParams := make(map[string]string)
|
||||
localVarQueryParams := url.Values{}
|
||||
localVarFormParams := url.Values{}
|
||||
if r.updateWebhookRequest == nil {
|
||||
return localVarReturnValue, nil, reportError("updateWebhookRequest is required and must be specified")
|
||||
}
|
||||
|
||||
// to determine the Content-Type header
|
||||
localVarHTTPContentTypes := []string{"application/json"}
|
||||
|
||||
// set Content-Type header
|
||||
localVarHTTPContentType := selectHeaderContentType(localVarHTTPContentTypes)
|
||||
if localVarHTTPContentType != "" {
|
||||
localVarHeaderParams["Content-Type"] = localVarHTTPContentType
|
||||
}
|
||||
|
||||
// to determine the Accept header
|
||||
localVarHTTPHeaderAccepts := []string{"application/json"}
|
||||
|
||||
// set Accept header
|
||||
localVarHTTPHeaderAccept := selectHeaderAccept(localVarHTTPHeaderAccepts)
|
||||
if localVarHTTPHeaderAccept != "" {
|
||||
localVarHeaderParams["Accept"] = localVarHTTPHeaderAccept
|
||||
}
|
||||
if r.authorization != nil {
|
||||
parameterAddToHeaderOrQuery(localVarHeaderParams, "authorization", r.authorization, "simple", "")
|
||||
}
|
||||
// body params
|
||||
localVarPostBody = r.updateWebhookRequest
|
||||
req, err := a.client.prepareRequest(r.ctx, localVarPath, localVarHTTPMethod, localVarPostBody, localVarHeaderParams, localVarQueryParams, localVarFormParams, formFiles)
|
||||
if err != nil {
|
||||
return localVarReturnValue, nil, err
|
||||
}
|
||||
|
||||
localVarHTTPResponse, err := a.client.callAPI(req)
|
||||
if err != nil || localVarHTTPResponse == nil {
|
||||
return localVarReturnValue, localVarHTTPResponse, err
|
||||
}
|
||||
|
||||
localVarBody, err := io.ReadAll(localVarHTTPResponse.Body)
|
||||
localVarHTTPResponse.Body.Close()
|
||||
localVarHTTPResponse.Body = io.NopCloser(bytes.NewBuffer(localVarBody))
|
||||
if err != nil {
|
||||
return localVarReturnValue, localVarHTTPResponse, err
|
||||
}
|
||||
|
||||
if localVarHTTPResponse.StatusCode >= 300 {
|
||||
newErr := &GenericOpenAPIError{
|
||||
body: localVarBody,
|
||||
error: localVarHTTPResponse.Status,
|
||||
}
|
||||
if localVarHTTPResponse.StatusCode == 422 {
|
||||
var v HTTPValidationError
|
||||
err = a.client.decode(&v, localVarBody, localVarHTTPResponse.Header.Get("Content-Type"))
|
||||
if err != nil {
|
||||
newErr.error = err.Error()
|
||||
return localVarReturnValue, localVarHTTPResponse, newErr
|
||||
}
|
||||
newErr.error = formatErrorMessage(localVarHTTPResponse.Status, &v)
|
||||
newErr.model = v
|
||||
}
|
||||
return localVarReturnValue, localVarHTTPResponse, newErr
|
||||
}
|
||||
|
||||
err = a.client.decode(&localVarReturnValue, localVarBody, localVarHTTPResponse.Header.Get("Content-Type"))
|
||||
if err != nil {
|
||||
newErr := &GenericOpenAPIError{
|
||||
body: localVarBody,
|
||||
error: err.Error(),
|
||||
}
|
||||
return localVarReturnValue, localVarHTTPResponse, newErr
|
||||
}
|
||||
|
||||
return localVarReturnValue, localVarHTTPResponse, nil
|
||||
}
|
||||
@@ -3,7 +3,7 @@ Hindsight HTTP API
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
API version: 0.4.13
|
||||
API version: 0.4.16
|
||||
*/
|
||||
|
||||
// Code generated by OpenAPI Generator (https://openapi-generator.tech); DO NOT EDIT.
|
||||
@@ -41,7 +41,7 @@ var (
|
||||
queryDescape = strings.NewReplacer( "%5B", "[", "%5D", "]" )
|
||||
)
|
||||
|
||||
// APIClient manages communication with the Hindsight HTTP API API v0.4.13
|
||||
// APIClient manages communication with the Hindsight HTTP API API v0.4.16
|
||||
// In most cases there should be only one, shared, APIClient.
|
||||
type APIClient struct {
|
||||
cfg *Configuration
|
||||
@@ -66,6 +66,8 @@ type APIClient struct {
|
||||
MonitoringAPI *MonitoringAPIService
|
||||
|
||||
OperationsAPI *OperationsAPIService
|
||||
|
||||
WebhooksAPI *WebhooksAPIService
|
||||
}
|
||||
|
||||
type service struct {
|
||||
@@ -93,6 +95,7 @@ func NewAPIClient(cfg *Configuration) *APIClient {
|
||||
c.MentalModelsAPI = (*MentalModelsAPIService)(&c.common)
|
||||
c.MonitoringAPI = (*MonitoringAPIService)(&c.common)
|
||||
c.OperationsAPI = (*OperationsAPIService)(&c.common)
|
||||
c.WebhooksAPI = (*WebhooksAPIService)(&c.common)
|
||||
|
||||
return c
|
||||
}
|
||||
|
||||
@@ -3,7 +3,7 @@ Hindsight HTTP API
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
API version: 0.4.13
|
||||
API version: 0.4.16
|
||||
*/
|
||||
|
||||
// Code generated by OpenAPI Generator (https://openapi-generator.tech); DO NOT EDIT.
|
||||
|
||||
@@ -67,7 +67,7 @@ func TestRetainWithContext(t *testing.T) {
|
||||
Items: []MemoryItem{
|
||||
{
|
||||
Content: "Bob went hiking in the mountains",
|
||||
Timestamp: *NewNullableTime(PtrTime(timestamp)),
|
||||
Timestamp: *NewNullableTimestamp(&Timestamp{TimeTime: ×tamp}),
|
||||
Context: *NewNullableString(PtrString("outdoor activities")),
|
||||
},
|
||||
},
|
||||
|
||||
@@ -3,7 +3,7 @@ Hindsight HTTP API
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
API version: 0.4.13
|
||||
API version: 0.4.16
|
||||
*/
|
||||
|
||||
// Code generated by OpenAPI Generator (https://openapi-generator.tech); DO NOT EDIT.
|
||||
|
||||
@@ -3,7 +3,7 @@ Hindsight HTTP API
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
API version: 0.4.13
|
||||
API version: 0.4.16
|
||||
*/
|
||||
|
||||
// Code generated by OpenAPI Generator (https://openapi-generator.tech); DO NOT EDIT.
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user