Compare commits
155
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
56cf3b1def | ||
|
|
d61e8ffff6 | ||
|
|
c20e08fecc | ||
|
|
2142d43f6f | ||
|
|
0c38d46ee9 | ||
|
|
375ec091f3 | ||
|
|
eb5b29f067 | ||
|
|
000fb9ddbe | ||
|
|
dfa02c8b61 | ||
|
|
d28b852732 | ||
|
|
6e03dd2d4c | ||
|
|
b11e053323 | ||
|
|
11154d48b7 | ||
|
|
a483682da6 | ||
|
|
8a7a70b828 | ||
|
|
dddd571a99 | ||
|
|
8bd9ce194b | ||
|
|
0e0fd14ed4 | ||
|
|
946a80bfb8 | ||
|
|
07af5b4a37 | ||
|
|
59b008a461 | ||
|
|
188f8fcb32 | ||
|
|
d7059d2840 | ||
|
|
5adaf60a9f | ||
|
|
3fb33b9873 | ||
|
|
ca97f947d9 | ||
|
|
347b9c23c4 | ||
|
|
ebe438b25b | ||
|
|
c05ca6529f | ||
|
|
6a6d4f2261 | ||
|
|
cf7aece729 | ||
|
|
36e94454e6 | ||
|
|
5bfef3caa4 | ||
|
|
c1908a0205 | ||
|
|
54354e4735 | ||
|
|
4df4b398f5 | ||
|
|
81aa4979b3 | ||
|
|
0108cd7019 | ||
|
|
aad0af9756 | ||
|
|
bd49f6a7c7 | ||
|
|
263eba1342 | ||
|
|
418524051d | ||
|
|
327aa05e80 | ||
|
|
9bf0023163 | ||
|
|
d64921221e | ||
|
|
7bb3d1925b | ||
|
|
eca0fd5a29 | ||
|
|
44398633bb | ||
|
|
52b893b93b | ||
|
|
52c216c1fb | ||
|
|
d9bc612a3c | ||
|
|
1fe43ec3fb | ||
|
|
ca87e29891 | ||
|
|
d2b14e51ee | ||
|
|
9676fc1699 | ||
|
|
73a5b576c9 | ||
|
|
685e50b9af | ||
|
|
c27fafb298 | ||
|
|
37fa0adf93 | ||
|
|
b95055ba28 | ||
|
|
4b78761d20 | ||
|
|
1549987015 | ||
|
|
b0e5d103d8 | ||
|
|
5ab6bdc9b6 | ||
|
|
ed1083803b | ||
|
|
ec3b415c42 | ||
|
|
3a188971b9 | ||
|
|
b3e32ce5b6 | ||
|
|
d0bbfa1015 | ||
|
|
043a68fa44 | ||
|
|
7542035e44 | ||
|
|
bee6f5d114 | ||
|
|
e20b1815fc | ||
|
|
395823f7b6 | ||
|
|
a910fd8a0b | ||
|
|
64ee029a18 | ||
|
|
25df91ca53 | ||
|
|
8987fb8267 | ||
|
|
86ff344c93 | ||
|
|
b6c7b2a2e9 | ||
|
|
9efce6a470 | ||
|
|
6cc7484b78 | ||
|
|
920afcce20 | ||
|
|
6d82c8a554 | ||
|
|
4b52b10e2e | ||
|
|
ac06df1ade | ||
|
|
5f1a867650 | ||
|
|
d284119246 | ||
|
|
84b9aa56ce | ||
|
|
f58ecee0b9 | ||
|
|
b9d16fe86f | ||
|
|
2ccde7a5cd | ||
|
|
d2ca26afaf | ||
|
|
b52feb305e | ||
|
|
558b2f8b67 | ||
|
|
383f0caa16 | ||
|
|
16e8c4e216 | ||
|
|
c6a1b507ba | ||
|
|
cb4fe70b63 | ||
|
|
408d7c34c8 | ||
|
|
7f2df54e01 | ||
|
|
1758d8510b | ||
|
|
7c3a5619f9 | ||
|
|
2b5b47f97d | ||
|
|
b2508bf2d6 | ||
|
|
b0fb1111ec | ||
|
|
4bf126bf52 | ||
|
|
e4449326e6 | ||
|
|
7bab4db28d | ||
|
|
e29ee58603 | ||
|
|
82e67315a4 | ||
|
|
2e82edee14 | ||
|
|
0338ab4d73 | ||
|
|
f4106fff55 | ||
|
|
d4f700ed22 | ||
|
|
041f0f13f4 | ||
|
|
9dee1c594b | ||
|
|
a1ebb2d9c6 | ||
|
|
06ddf041e5 | ||
|
|
d251fcb7d2 | ||
|
|
e839c65537 | ||
|
|
0cde79b831 | ||
|
|
8767a518db | ||
|
|
10ed288d80 | ||
|
|
7143684a81 | ||
|
|
11da432db2 | ||
|
|
73d3231bbd | ||
|
|
59d825dfca | ||
|
|
ae7099fd02 | ||
|
|
56db6d7cf6 | ||
|
|
0eb52762ae | ||
|
|
e93c560288 | ||
|
|
1f213d00f7 | ||
|
|
8f51f99dde | ||
|
|
5cc1482a72 | ||
|
|
639d84ad32 | ||
|
|
1c74f795a6 | ||
|
|
b4f9fbe1b5 | ||
|
|
b992ba996d | ||
|
|
f00d3c7f66 | ||
|
|
e97b615547 | ||
|
|
29cc1d7fdc | ||
|
|
016b5f0363 | ||
|
|
dd7e252452 | ||
|
|
fda1a77f70 | ||
|
|
0accef8e98 | ||
|
|
c77e2368de | ||
|
|
38ef0247c2 | ||
|
|
a158b819f3 | ||
|
|
381963c28a | ||
|
|
ba158c9cdb | ||
|
|
767a2c0061 | ||
|
|
36334f27a1 | ||
|
|
6a479dddb9 | ||
|
|
7058d1aad7 |
@@ -1,7 +1,7 @@
|
||||
{
|
||||
"$schema": "https://anthropic.com/claude-code/marketplace.schema.json",
|
||||
"name": "hindsight",
|
||||
"version": "0.7.2",
|
||||
"version": "0.7.5",
|
||||
"description": "Official Hindsight integrations for Claude Code",
|
||||
"owner": {
|
||||
"name": "vectorize-io"
|
||||
@@ -11,6 +11,11 @@
|
||||
"name": "hindsight-memory",
|
||||
"description": "Automatic long-term memory for Claude Code via Hindsight",
|
||||
"source": "./hindsight-integrations/claude-code"
|
||||
},
|
||||
{
|
||||
"name": "hindsight-zcode",
|
||||
"description": "No-MCP long-term memory for ZCode via Hindsight hooks",
|
||||
"source": "./hindsight-integrations/zcode"
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
@@ -78,6 +78,11 @@ results = await asyncio.gather(*tasks, return_exceptions=True)
|
||||
- **Authentication/tenancy is enforced inside each engine method, not assumed by the handler.** Every engine method that touches bank-scoped data must authenticate via `request_context` — typically `await self._authenticate_tenant(request_context)` (often indirectly through `get_bank_profile(...)`) — so the correct tenant schema is resolved before any query runs. Handlers must thread `request_context` through to the engine method; never query a tenant-scoped table assuming the schema is already set.
|
||||
- Engine methods return typed models (Pydantic/dataclass), not raw dicts (see Type Safety).
|
||||
|
||||
### Database Locking
|
||||
- **Never use PostgreSQL advisory locks** (`pg_advisory_lock`, `pg_try_advisory_lock`, `pg_advisory_xact_lock`, `pg_advisory_unlock`, …) in migrations, engine code, or anything else. Hindsight runs against connection poolers and managed/PG-compatible services where advisory locks are unreliable or unsupported: session-level locks silently leak or vanish when a pooler hands the session to another client, and callers can block forever on a lock the server never grants. Reject any new occurrence, including ones that look "safe" because they are transaction-scoped.
|
||||
- The pre-existing usage in `hindsight_api/migrations.py` is grandfathered, not a precedent — it is tracked for removal. Don't copy it.
|
||||
- Design the concurrency out instead of locking around it: give each process its own object to write (e.g. per-schema DDL rather than a shared `public.` object), make the operation idempotent, or use a real row/table constraint (`INSERT ... ON CONFLICT`, `SELECT ... FOR UPDATE` in a fixed order). See #2690 for a migration that reached for `pg_advisory_xact_lock` and had to be reverted.
|
||||
|
||||
### Branch Hygiene
|
||||
- **Always start new feature branches from `origin/main`** — rebase to ensure a clean base.
|
||||
- **Only include commits relevant to the PR/branch/feature** — no unrelated changes. If the branch contains commits that don't belong, they must be removed before merging.
|
||||
@@ -204,6 +209,14 @@ in `hindsight-api-slim/hindsight_api/config.py`):
|
||||
The `test_bundled_template_matches_repo_root` sync test fails on drift; if the
|
||||
root file changed without re-copying, flag it as a **must fix**.
|
||||
|
||||
### 11c. Check for advisory locks
|
||||
|
||||
Grep the diff for `advisory` (`git diff main...HEAD | grep -in advisory`). Any new
|
||||
`pg_advisory_lock` / `pg_try_advisory_lock` / `pg_advisory_xact_lock` /
|
||||
`pg_advisory_unlock` call is a **must fix** — see Database Locking above. Point the
|
||||
author at the alternatives (per-process objects, idempotent DDL, row-level
|
||||
constraints) rather than just asking them to drop the lock.
|
||||
|
||||
### 12. Review against other coding standards
|
||||
|
||||
Check the diff for violations of the standards listed above:
|
||||
|
||||
@@ -59,6 +59,15 @@ HINDSIGHT_API_LLM_BASE_URL=https://api.openai.com/v1
|
||||
# HINDSIGHT_API_LLM_BASE_URL=http://localhost:1234/v1
|
||||
# HINDSIGHT_API_LLM_MODEL=qwen2.5-32b-instruct
|
||||
|
||||
# Example: Ollama local configuration (native provider)
|
||||
# HINDSIGHT_API_LLM_PROVIDER=ollama
|
||||
# HINDSIGHT_API_LLM_BASE_URL=http://localhost:11434/v1
|
||||
# HINDSIGHT_API_LLM_MODEL=gemma3:12b
|
||||
# Native Ollama context-window override (num_ctx). Leave unset to let Ollama use
|
||||
# the model Modelfile / server default; set a positive integer only to force a
|
||||
# specific context size (e.g. 16384 to keep the previous request behavior).
|
||||
# HINDSIGHT_API_LLM_OLLAMA_NUM_CTX=16384
|
||||
|
||||
# Multi-LLM strategies: configure extra LLMs by index alongside the primary above,
|
||||
# then pick a routing strategy. Unset = single primary LLM (default). Members are
|
||||
# numbered from 1; indices must be contiguous. Each operation can override with a
|
||||
@@ -80,6 +89,10 @@ HINDSIGHT_API_LOG_LEVEL=info
|
||||
# Unset uses HINDSIGHT_API_RETAIN_CHUNK_SIZE as the structured-chunk limit.
|
||||
# HINDSIGHT_API_RETAIN_STRUCTURED_CHUNK_SIZE=
|
||||
|
||||
# When true, a retain operation that hit any fact-extraction errors is marked
|
||||
# 'failed' (not 'completed'), surfacing silently-dropped facts. Default false.
|
||||
# HINDSIGHT_API_FAIL_ON_EXTRACTION_ERRORS=false
|
||||
|
||||
# Dry-run extraction preview endpoint (POST /memories/dry-run-extract). Enabled by default; it makes
|
||||
# a real LLM call but stores nothing. Set to false to remove the endpoint (returns 404).
|
||||
# HINDSIGHT_API_ENABLE_DRY_RUN_EXTRACT=true
|
||||
@@ -95,7 +108,10 @@ HINDSIGHT_API_LOG_LEVEL=info
|
||||
# HINDSIGHT_API_READ_DATABASE_URL= # Optional read-replica URL. When set, recall queries (semantic, BM25, graph, temporal) flow through a separate pool against this URL, offloading the primary. Typically points to a read-only endpoint (CNPG's <cluster>-ro service or Aurora reader endpoint).
|
||||
# HINDSIGHT_API_MIGRATION_DATABASE_URL= # Direct PostgreSQL URL for migrations (bypasses PgBouncer). Falls back to DATABASE_URL.
|
||||
# HINDSIGHT_API_DATABASE_SCHEMA=public # PostgreSQL schema name (default: public)
|
||||
# HINDSIGHT_API_DB_MAX_PARALLEL_WORKERS_PER_GATHER= # Optional cap on Postgres planner parallelism for this process's pool connections. Unset leaves the server default; 0 makes background/bulk queries run serially (useful on worker processes sharing a primary with latency-sensitive traffic).
|
||||
# HINDSIGHT_API_MIGRATION_CONCURRENCY=1 # Tenant schemas to migrate concurrently (PG only, each in its own process; per-schema work stays sequential). Each worker has ~1-2s startup cost + uses ~3 DB connections, so it only pays off with many schemas (tens+) or slow migrations; keep concurrency*3 <= spare max_connections. Default: 1 (sequential).
|
||||
# HINDSIGHT_API_OPERATION_RETENTION_DAYS=30 # Prune terminal operation rows, payloads, and metadata after this many days; 0 (the default) keeps them forever.
|
||||
# HINDSIGHT_API_OPERATION_CLEANUP_BATCH_SIZE=1000 # Maximum expired terminal rows deleted per tenant schema in each cleanup cycle; must be positive.
|
||||
|
||||
# Vector Extension (Optional - uses pgvector by default)
|
||||
# Options: "pgvector" (default), "vchord", "pgvectorscale" (DiskANN)
|
||||
@@ -115,6 +131,11 @@ HINDSIGHT_API_LOG_LEVEL=info
|
||||
# chinese_lindera/lindera(chinese), japanese_lindera/lindera(japanese),
|
||||
# korean_lindera/lindera(korean), ngram(min,max), edge_ngram(min,max)
|
||||
# HINDSIGHT_API_TEXT_SEARCH_EXTENSION_PG_SEARCH_TOKENIZER=
|
||||
# Optional cap on the number of terms in the native PostgreSQL BM25 tsquery.
|
||||
# Long queries OR-join every normalized token, which can match too much of a
|
||||
# large bank. 0 (default) keeps the historical uncapped behavior; a positive
|
||||
# value bounds only the native backend (other BM25 backends get the raw query).
|
||||
# HINDSIGHT_API_BM25_MAX_QUERY_TERMS=0
|
||||
|
||||
# File Parser (Optional - uses markitdown by default)
|
||||
# HINDSIGHT_API_FILE_PARSER=markitdown
|
||||
@@ -133,6 +154,8 @@ HINDSIGHT_API_LOG_LEVEL=info
|
||||
# HINDSIGHT_API_EMBEDDINGS_PROVIDER=local
|
||||
# For local provider:
|
||||
# HINDSIGHT_API_EMBEDDINGS_LOCAL_MODEL=BAAI/bge-small-en-v1.5
|
||||
# Force CPU if local embeddings hit MPS/XPC instability on macOS:
|
||||
# HINDSIGHT_API_EMBEDDINGS_LOCAL_FORCE_CPU=false
|
||||
# For ONNX provider (local CPU embeddings without an Ollama/TEI sidecar):
|
||||
# HINDSIGHT_API_EMBEDDINGS_PROVIDER=onnx
|
||||
# HINDSIGHT_API_EMBEDDINGS_ONNX_MODEL_ID=intfloat/multilingual-e5-small
|
||||
@@ -172,8 +195,13 @@ HINDSIGHT_API_LOG_LEVEL=info
|
||||
# Reranker Configuration (Optional - uses local by default)
|
||||
# Provider: "local" (default) or "tei" (HuggingFace Text Embeddings Inference)
|
||||
# HINDSIGHT_API_RERANKER_PROVIDER=local
|
||||
# Trusted gateway attribution (disabled by default). When enabled, remote
|
||||
# reranker requests include X-Hindsight-Bank-Id with the current bank ID.
|
||||
# HINDSIGHT_API_RERANKER_SEND_BANK_AS_HEADER=false
|
||||
# For local provider:
|
||||
# HINDSIGHT_API_RERANKER_LOCAL_MODEL=cross-encoder/ms-marco-MiniLM-L-6-v2
|
||||
# Force CPU if the local reranker hits MPS/XPC instability on macOS:
|
||||
# HINDSIGHT_API_RERANKER_LOCAL_FORCE_CPU=false
|
||||
# For TEI provider:
|
||||
# HINDSIGHT_API_RERANKER_TEI_URL=http://localhost:8081
|
||||
|
||||
|
||||
+54
-17
@@ -41,6 +41,7 @@ jobs:
|
||||
integrations-github-copilot: ${{ steps.filter.outputs.integrations-github-copilot }}
|
||||
integrations-continue: ${{ steps.filter.outputs.integrations-continue }}
|
||||
integrations-cursor-cli: ${{ steps.filter.outputs.integrations-cursor-cli }}
|
||||
integrations-zcode: ${{ steps.filter.outputs.integrations-zcode }}
|
||||
integrations-crewai: ${{ steps.filter.outputs.integrations-crewai }}
|
||||
integrations-litellm: ${{ steps.filter.outputs.integrations-litellm }}
|
||||
integrations-pydantic-ai: ${{ steps.filter.outputs.integrations-pydantic-ai }}
|
||||
@@ -180,6 +181,8 @@ jobs:
|
||||
- 'hindsight-integrations/cursor/**'
|
||||
integrations-zed:
|
||||
- 'hindsight-integrations/zed/**'
|
||||
integrations-zcode:
|
||||
- 'hindsight-integrations/zcode/**'
|
||||
integrations-n8n:
|
||||
- 'hindsight-integrations/n8n/**'
|
||||
integrations-zapier:
|
||||
@@ -520,22 +523,17 @@ jobs:
|
||||
with:
|
||||
ref: ${{ github.event.pull_request.head.sha || '' }}
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v6
|
||||
- name: Set up Node.js
|
||||
uses: actions/setup-node@v6
|
||||
with:
|
||||
python-version: '3.11'
|
||||
|
||||
- name: Install package and pytest
|
||||
working-directory: ./hindsight-integrations/zed
|
||||
# Installs the package (incl. the zstandard runtime dep) so the threads.db
|
||||
# reader tests can decompress Zed's zstd blobs.
|
||||
run: pip install -e . pytest
|
||||
node-version: '22'
|
||||
|
||||
- name: Run tests
|
||||
working-directory: ./hindsight-integrations/zed
|
||||
# PR CI runs only the deterministic bucket; the real-LLM E2E bucket
|
||||
# (requires_real_llm) needs a live Hindsight server and runs separately.
|
||||
run: python -m pytest tests/ -v -m "not requires_real_llm"
|
||||
# Config-only integration with no dependencies — it uses Node's built-in
|
||||
# test runner. The runtime MCP bridge is `npx mcp-remote` (Node), so this
|
||||
# integration requires only Node.js (no Python).
|
||||
run: npm test
|
||||
|
||||
test-omo-integration:
|
||||
needs: [detect-changes]
|
||||
@@ -702,6 +700,43 @@ jobs:
|
||||
working-directory: ./hindsight-integrations/cursor-cli
|
||||
run: uv run pytest tests -v
|
||||
|
||||
test-zcode-integration:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
(github.event_name == 'workflow_dispatch' ||
|
||||
needs.detect-changes.outputs.integrations-zcode == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
with:
|
||||
ref: ${{ github.event.pull_request.head.sha || '' }}
|
||||
|
||||
- name: Install uv
|
||||
uses: astral-sh/setup-uv@v7
|
||||
with:
|
||||
enable-cache: true
|
||||
prune-cache: false
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v6
|
||||
with:
|
||||
python-version-file: ".python-version"
|
||||
|
||||
- name: Build zcode integration
|
||||
working-directory: ./hindsight-integrations/zcode
|
||||
run: uv build
|
||||
|
||||
- name: Install dependencies
|
||||
working-directory: ./hindsight-integrations/zcode
|
||||
run: uv sync --frozen
|
||||
|
||||
- name: Run tests
|
||||
working-directory: ./hindsight-integrations/zcode
|
||||
run: uv run pytest tests -v
|
||||
|
||||
build-ai-sdk-integration:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
@@ -1229,10 +1264,10 @@ jobs:
|
||||
|
||||
build-docs:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
(github.event_name == 'workflow_dispatch' ||
|
||||
needs.detect-changes.outputs.docs == 'true' ||
|
||||
needs.detect-changes.outputs.ci == 'true')
|
||||
# Keep the production docs build as an unconditional PR check. OpenAPI
|
||||
# generation used to build the site again inside verify-generated-files;
|
||||
# running the existing job for every PR preserves that coverage without
|
||||
# serializing two full Docusaurus builds in the generated-files check.
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 30
|
||||
|
||||
@@ -4771,7 +4806,8 @@ jobs:
|
||||
cd ../hindsight-embed && uv sync --frozen --index-strategy unsafe-best-match
|
||||
|
||||
- name: Run generate-openapi
|
||||
run: ./scripts/generate-openapi.sh
|
||||
working-directory: hindsight-dev
|
||||
run: uv run generate-openapi
|
||||
|
||||
- name: Run generate-bank-template-schema
|
||||
run: ./scripts/generate-bank-template-schema.sh
|
||||
@@ -4909,6 +4945,7 @@ jobs:
|
||||
- test-github-copilot-integration
|
||||
- test-codex-integration
|
||||
- test-cursor-cli-integration
|
||||
- test-zcode-integration
|
||||
- build-ai-sdk-integration
|
||||
- test-ai-sdk-integration-deno
|
||||
- test-opencode-integration
|
||||
|
||||
@@ -41,6 +41,7 @@ nltk_data/
|
||||
logs/
|
||||
|
||||
.DS_Store
|
||||
.sesskey
|
||||
|
||||
# Generated docs files
|
||||
hindsight-docs/static/llms-full.txt
|
||||
|
||||
@@ -56,7 +56,6 @@ BACKUP_TABLES = [
|
||||
"observation_history",
|
||||
"mental_models",
|
||||
"mental_model_history",
|
||||
"knowledge_pages",
|
||||
"directives",
|
||||
"async_operations",
|
||||
"webhooks",
|
||||
@@ -77,7 +76,8 @@ async def _admin_connect(db_url: str) -> asyncpg.Connection:
|
||||
is the only step needed to connect. JSON codecs are registered so ``jsonb``
|
||||
columns decode to Python objects (used by the export row dumps).
|
||||
"""
|
||||
is_pg0, instance_name, _ = parse_pg0_url(db_url)
|
||||
_pg0 = parse_pg0_url(db_url)
|
||||
is_pg0, instance_name = _pg0.is_pg0, _pg0.instance_name
|
||||
if is_pg0:
|
||||
typer.echo(f"Starting embedded PostgreSQL (instance: {instance_name})...")
|
||||
conn = await asyncpg.connect(await resolve_database_url(db_url))
|
||||
@@ -178,7 +178,8 @@ async def _restore(database_url: str, input_path: Path, schema: str = "public")
|
||||
|
||||
async def _run_backup(db_url: str, output: Path, schema: str = "public") -> dict[str, Any]:
|
||||
"""Resolve database URL and run backup."""
|
||||
is_pg0, instance_name, _ = parse_pg0_url(db_url)
|
||||
_pg0 = parse_pg0_url(db_url)
|
||||
is_pg0, instance_name = _pg0.is_pg0, _pg0.instance_name
|
||||
if is_pg0:
|
||||
typer.echo(f"Starting embedded PostgreSQL (instance: {instance_name})...")
|
||||
resolved_url = await resolve_database_url(db_url)
|
||||
@@ -187,7 +188,8 @@ async def _run_backup(db_url: str, output: Path, schema: str = "public") -> dict
|
||||
|
||||
async def _run_restore(db_url: str, input_file: Path, schema: str = "public") -> dict[str, Any]:
|
||||
"""Resolve database URL and run restore."""
|
||||
is_pg0, instance_name, _ = parse_pg0_url(db_url)
|
||||
_pg0 = parse_pg0_url(db_url)
|
||||
is_pg0, instance_name = _pg0.is_pg0, _pg0.instance_name
|
||||
if is_pg0:
|
||||
typer.echo(f"Starting embedded PostgreSQL (instance: {instance_name})...")
|
||||
resolved_url = await resolve_database_url(db_url)
|
||||
@@ -262,7 +264,8 @@ async def _run_migration(
|
||||
"""Resolve database URL and run migrations for one schema or all discovered schemas."""
|
||||
from ..migrations import run_migrations_for_schemas
|
||||
|
||||
is_pg0, instance_name, _ = parse_pg0_url(db_url)
|
||||
_pg0 = parse_pg0_url(db_url)
|
||||
is_pg0, instance_name = _pg0.is_pg0, _pg0.instance_name
|
||||
if is_pg0:
|
||||
typer.echo(f"Starting embedded PostgreSQL (instance: {instance_name})...")
|
||||
resolved_url = await resolve_database_url(db_url)
|
||||
@@ -473,7 +476,8 @@ def import_bank_command(
|
||||
|
||||
async def _decommission_worker(db_url: str, worker_id: str, schema: str = "public") -> int:
|
||||
"""Release all tasks owned by a worker, setting them back to pending status."""
|
||||
is_pg0, instance_name, _ = parse_pg0_url(db_url)
|
||||
_pg0 = parse_pg0_url(db_url)
|
||||
is_pg0, instance_name = _pg0.is_pg0, _pg0.instance_name
|
||||
if is_pg0:
|
||||
typer.echo(f"Starting embedded PostgreSQL (instance: {instance_name})...")
|
||||
resolved_url = await resolve_database_url(db_url)
|
||||
@@ -532,7 +536,8 @@ def decommission_worker(
|
||||
|
||||
async def _decommission_all_workers(db_url: str, schema: str = "public") -> list[dict[str, Any]]:
|
||||
"""Release all processing tasks from all workers, setting them back to pending status."""
|
||||
is_pg0, instance_name, _ = parse_pg0_url(db_url)
|
||||
_pg0 = parse_pg0_url(db_url)
|
||||
is_pg0, instance_name = _pg0.is_pg0, _pg0.instance_name
|
||||
if is_pg0:
|
||||
typer.echo(f"Starting embedded PostgreSQL (instance: {instance_name})...")
|
||||
resolved_url = await resolve_database_url(db_url)
|
||||
@@ -597,7 +602,8 @@ def decommission_workers(
|
||||
|
||||
async def _worker_status(db_url: str, schema: str = "public") -> list[dict[str, Any]]:
|
||||
"""Get all processing tasks grouped by worker with their last update time."""
|
||||
is_pg0, instance_name, _ = parse_pg0_url(db_url)
|
||||
_pg0 = parse_pg0_url(db_url)
|
||||
is_pg0, instance_name = _pg0.is_pg0, _pg0.instance_name
|
||||
if is_pg0:
|
||||
typer.echo(f"Starting embedded PostgreSQL (instance: {instance_name})...")
|
||||
resolved_url = await resolve_database_url(db_url)
|
||||
|
||||
@@ -96,7 +96,8 @@ def get_database_url() -> str:
|
||||
# for the sync engine used during migrations.
|
||||
database_url = to_libpq_url(database_url)
|
||||
|
||||
config.set_main_option("sqlalchemy.url", database_url)
|
||||
# Alembic stores options through ConfigParser, where '%' is interpolation.
|
||||
config.set_main_option("sqlalchemy.url", database_url.replace("%", "%%"))
|
||||
return database_url
|
||||
|
||||
|
||||
|
||||
-52
@@ -1,52 +0,0 @@
|
||||
"""Add managed flag to knowledge_pages.
|
||||
|
||||
The knowledge base is managed by clients (CRUD over folders/pages). ``managed``
|
||||
lets a client tag a node as system-owned vs. hand-authored; it carries no
|
||||
server-side behaviour.
|
||||
|
||||
Revision ID: a5b6c7d8e9f0
|
||||
Revises: a9b8c7d6e5f4
|
||||
Create Date: 2026-06-26
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "a5b6c7d8e9f0"
|
||||
down_revision: str | Sequence[str] | None = "a9b8c7d6e5f4"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _pg_schema_prefix() -> str:
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
schema = _pg_schema_prefix()
|
||||
op.execute(f"ALTER TABLE {schema}knowledge_pages ADD COLUMN IF NOT EXISTS managed BOOLEAN NOT NULL DEFAULT false")
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
schema = _pg_schema_prefix()
|
||||
op.execute(f"ALTER TABLE {schema}knowledge_pages DROP COLUMN IF EXISTS managed")
|
||||
|
||||
|
||||
def _oracle_upgrade() -> None:
|
||||
op.execute("ALTER TABLE knowledge_pages ADD (managed NUMBER(1) DEFAULT 0 NOT NULL)")
|
||||
|
||||
|
||||
def _oracle_downgrade() -> None:
|
||||
op.execute("ALTER TABLE knowledge_pages DROP COLUMN managed")
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade, oracle=_oracle_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade, oracle=_oracle_downgrade)
|
||||
+82
@@ -0,0 +1,82 @@
|
||||
"""Add indexes for terminal cleanup and newest-first operation listing.
|
||||
|
||||
Revision ID: a8c1e4f7b0d3
|
||||
Revises: e7c3a9f1b2d5
|
||||
Create Date: 2026-07-14
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "a8c1e4f7b0d3"
|
||||
down_revision: str | Sequence[str] | None = "e7c3a9f1b2d5"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _pg_schema_prefix() -> str:
|
||||
"""Schema-qualifier for PostgreSQL multi-tenant migration runs."""
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
schema = _pg_schema_prefix()
|
||||
# These can be large tables in long-running installations. Concurrent DDL
|
||||
# keeps operation submission, polling, and status reads available.
|
||||
with op.get_context().autocommit_block():
|
||||
op.execute(
|
||||
"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_async_operations_terminal_cleanup "
|
||||
f"ON {schema}async_operations (updated_at, operation_id) "
|
||||
"WHERE status IN ('completed', 'failed', 'cancelled')"
|
||||
)
|
||||
op.execute(
|
||||
"CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_async_operations_bank_created_desc "
|
||||
f"ON {schema}async_operations (bank_id, created_at DESC)"
|
||||
)
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
schema = _pg_schema_prefix()
|
||||
with op.get_context().autocommit_block():
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_async_operations_bank_created_desc")
|
||||
op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_async_operations_terminal_cleanup")
|
||||
|
||||
|
||||
def _oracle_create_index(sql: str) -> None:
|
||||
"""Create an index idempotently for rerun-safe Oracle migrations."""
|
||||
block = (
|
||||
"BEGIN "
|
||||
"EXECUTE IMMEDIATE :stmt; "
|
||||
"EXCEPTION WHEN OTHERS THEN "
|
||||
"IF SQLCODE = -955 THEN NULL; ELSE RAISE; END IF; "
|
||||
"END;"
|
||||
)
|
||||
op.get_bind().exec_driver_sql(block, {"stmt": sql})
|
||||
|
||||
|
||||
def _oracle_upgrade() -> None:
|
||||
# Oracle migrations run with CURRENT_SCHEMA set to each tenant, so table
|
||||
# and index names intentionally remain unqualified here.
|
||||
_oracle_create_index(
|
||||
"CREATE INDEX idx_async_operations_terminal_cleanup ON async_operations (updated_at, operation_id, status)"
|
||||
)
|
||||
_oracle_create_index(
|
||||
"CREATE INDEX idx_async_operations_bank_created_desc ON async_operations (bank_id, created_at DESC)"
|
||||
)
|
||||
|
||||
|
||||
def _oracle_downgrade() -> None:
|
||||
op.execute("DROP INDEX idx_async_operations_bank_created_desc")
|
||||
op.execute("DROP INDEX idx_async_operations_terminal_cleanup")
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade, oracle=_oracle_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade, oracle=_oracle_downgrade)
|
||||
-110
@@ -1,110 +0,0 @@
|
||||
"""Add knowledge_pages table (knowledge-base hierarchy).
|
||||
|
||||
The knowledge base organizes synthesized mental models into a navigable tree of
|
||||
**folders** and **pages**. A page references the mental model that holds its
|
||||
content (``mental_model_id``); a folder is a pure container (``mental_model_id``
|
||||
NULL). Hierarchy is a single self-referential ``parent_id`` so folders can nest
|
||||
arbitrarily. Content stays in ``mental_models`` — this table is metadata + tree
|
||||
structure only.
|
||||
|
||||
Revision ID: a9b8c7d6e5f4
|
||||
Revises: b57a7c9e0d13
|
||||
Create Date: 2026-06-25
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "a9b8c7d6e5f4"
|
||||
down_revision: str | Sequence[str] | None = "b57a7c9e0d13"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _pg_schema_prefix() -> str:
|
||||
"""Schema-qualifier for raw SQL on PG (multi-tenant search_path)."""
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
schema = _pg_schema_prefix()
|
||||
# parent_id self-FK cascades so deleting a folder row removes its whole
|
||||
# subtree of rows in one shot. The mental_model FK is composite (matches the
|
||||
# mental_models (id, bank_id) PK) and cascades too, so deleting a page's
|
||||
# mental model removes the page row — folders skip the FK because a NULL
|
||||
# column in a composite FK is not enforced (MATCH SIMPLE).
|
||||
op.execute(
|
||||
f"""
|
||||
CREATE TABLE IF NOT EXISTS {schema}knowledge_pages (
|
||||
id VARCHAR(64) NOT NULL,
|
||||
bank_id TEXT NOT NULL,
|
||||
parent_id VARCHAR(64),
|
||||
kind VARCHAR(16) NOT NULL,
|
||||
name TEXT NOT NULL,
|
||||
mental_model_id VARCHAR(64),
|
||||
sort_order INTEGER NOT NULL DEFAULT 0,
|
||||
created_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT now(),
|
||||
updated_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT now(),
|
||||
CONSTRAINT pk_knowledge_pages PRIMARY KEY (id),
|
||||
CONSTRAINT ck_knowledge_pages_kind CHECK (kind IN ('folder', 'page')),
|
||||
CONSTRAINT fk_kp_bank FOREIGN KEY (bank_id)
|
||||
REFERENCES {schema}banks(bank_id) ON DELETE CASCADE,
|
||||
CONSTRAINT fk_kp_parent FOREIGN KEY (parent_id)
|
||||
REFERENCES {schema}knowledge_pages(id) ON DELETE CASCADE,
|
||||
CONSTRAINT fk_kp_mm FOREIGN KEY (mental_model_id, bank_id)
|
||||
REFERENCES {schema}mental_models(id, bank_id) ON DELETE CASCADE
|
||||
)
|
||||
"""
|
||||
)
|
||||
op.execute(
|
||||
f"CREATE INDEX IF NOT EXISTS idx_kp_bank_parent ON {schema}knowledge_pages (bank_id, parent_id, sort_order)"
|
||||
)
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
schema = _pg_schema_prefix()
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema}idx_kp_bank_parent")
|
||||
op.execute(f"DROP TABLE IF EXISTS {schema}knowledge_pages")
|
||||
|
||||
|
||||
def _oracle_upgrade() -> None:
|
||||
op.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS knowledge_pages (
|
||||
id VARCHAR2(64) NOT NULL,
|
||||
bank_id VARCHAR2(256) NOT NULL,
|
||||
parent_id VARCHAR2(64),
|
||||
kind VARCHAR2(16) NOT NULL,
|
||||
name CLOB NOT NULL,
|
||||
mental_model_id VARCHAR2(64),
|
||||
sort_order NUMBER DEFAULT 0 NOT NULL,
|
||||
created_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,
|
||||
updated_at TIMESTAMP WITH TIME ZONE DEFAULT SYSTIMESTAMP NOT NULL,
|
||||
CONSTRAINT pk_knowledge_pages PRIMARY KEY (id),
|
||||
CONSTRAINT ck_knowledge_pages_kind CHECK (kind IN ('folder', 'page')),
|
||||
CONSTRAINT fk_kp_bank FOREIGN KEY (bank_id)
|
||||
REFERENCES banks(bank_id) ON DELETE CASCADE,
|
||||
CONSTRAINT fk_kp_parent FOREIGN KEY (parent_id)
|
||||
REFERENCES knowledge_pages(id) ON DELETE CASCADE,
|
||||
CONSTRAINT fk_kp_mm FOREIGN KEY (mental_model_id, bank_id)
|
||||
REFERENCES mental_models(id, bank_id) ON DELETE CASCADE
|
||||
)
|
||||
"""
|
||||
)
|
||||
op.execute("CREATE INDEX idx_kp_bank_parent ON knowledge_pages (bank_id, parent_id, sort_order)")
|
||||
|
||||
|
||||
def _oracle_downgrade() -> None:
|
||||
op.execute("DROP TABLE knowledge_pages CASCADE CONSTRAINTS")
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade, oracle=_oracle_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade, oracle=_oracle_downgrade)
|
||||
+259
@@ -0,0 +1,259 @@
|
||||
"""Install the maintenance discovery routines into the configured schema.
|
||||
|
||||
The three discovery routines driving the background maintenance loop —
|
||||
``banks_needing_consolidation()``, ``schemas_with_expired_rows(...)`` and
|
||||
``mental_models_with_cron()`` — were installed into ``public`` and gated on the
|
||||
run being the base run (no ``target_schema``) or an explicit
|
||||
``target_schema='public'`` run (``e5f6a7b8c9d0`` → ``b2d4f6a8c1e3`` →
|
||||
``c7e9f1a3b5d2``, ``f4d1c2b3a5e6``).
|
||||
|
||||
That leaves a **single-tenant deployment migrated into a dedicated, non-**
|
||||
``public`` **schema** (``HINDSIGHT_API_DATABASE_SCHEMA=<non-public>``) with no
|
||||
routines at all: the runtime migrates only that one schema, so ``target_schema``
|
||||
is never falsy or ``public``, the gate never opens, and the maintenance loop
|
||||
logs, forever::
|
||||
|
||||
function public.banks_needing_consolidation() does not exist
|
||||
function public.schemas_with_expired_rows(...) does not exist
|
||||
|
||||
The revision is stamped applied, so redeploying the same version does not help
|
||||
(issue #2638; #2056 only fixed the ``public``/base-run case).
|
||||
|
||||
**The bug was the hardcoded literal, not the gating.** These routines are
|
||||
database-global — each enumerates ``pg_class`` across every schema and dispatches
|
||||
per schema — so exactly one copy should exist, and the maintenance loop calls the
|
||||
one in ``get_config().database_schema`` (see ``fq_routine``). The old gate
|
||||
installed into whichever schema was named ``public`` instead of whichever schema
|
||||
the deployment is actually configured to use. Comparing ``target_schema`` against
|
||||
the configured schema instead of the literal fixes #2638 at the source.
|
||||
|
||||
That also keeps the property the gate existed for: exactly one migration run
|
||||
satisfies the predicate, so concurrent per-schema runs never issue competing
|
||||
``CREATE OR REPLACE`` against the same ``pg_proc`` row and cannot hit
|
||||
``tuple concurrently updated``. No cross-process coordination is required — in
|
||||
particular no advisory lock, which is unusable here because Hindsight runs behind
|
||||
connection poolers and managed PG services (see #2817).
|
||||
|
||||
Runs targeting any *other* schema drop the routines from that schema rather than
|
||||
merely skipping. An earlier revision of this migration installed a copy into
|
||||
every schema it touched, which left one dead duplicate per tenant on any database
|
||||
that ran it; the drop makes the next migration pass clean those up instead of
|
||||
leaving them behind forever.
|
||||
|
||||
PostgreSQL only: the maintenance loop and worker poller are PG-only, so the
|
||||
Oracle slot is intentionally absent (mirrors ``e5f6a7b8c9d0``).
|
||||
|
||||
Revision ID: b6d2f8a4c1e7
|
||||
Revises: a8c1e4f7b0d3
|
||||
Create Date: 2026-07-20
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
from hindsight_api.config import get_config
|
||||
|
||||
revision: str = "b6d2f8a4c1e7"
|
||||
down_revision: str | Sequence[str] | None = "a8c1e4f7b0d3"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _configured_schema() -> str:
|
||||
"""The one schema this deployment's routines live in and are called from."""
|
||||
return get_config().database_schema or "public"
|
||||
|
||||
|
||||
def _target_schema() -> str | None:
|
||||
return context.config.get_main_option("target_schema")
|
||||
|
||||
|
||||
def _is_install_run() -> bool:
|
||||
"""True for the single run that owns the routines.
|
||||
|
||||
The base run (no ``target_schema``) and the run targeting the configured
|
||||
schema are the same deployment-level run; every other target is a tenant
|
||||
schema that must not carry its own copy.
|
||||
"""
|
||||
target = _target_schema()
|
||||
return not target or target == _configured_schema()
|
||||
|
||||
|
||||
def _prefix(schema: str | None) -> str:
|
||||
"""Qualifier for ``schema``, or ``""`` to fall back to ``search_path``."""
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
if not _is_install_run():
|
||||
_drop_stray_copies()
|
||||
return
|
||||
schema = _prefix(_target_schema())
|
||||
|
||||
op.execute(
|
||||
f"""
|
||||
CREATE OR REPLACE FUNCTION {schema}banks_needing_consolidation()
|
||||
RETURNS TABLE(schema_name text, bank_id text)
|
||||
LANGUAGE plpgsql STABLE
|
||||
AS $fn$
|
||||
DECLARE
|
||||
sch text;
|
||||
BEGIN
|
||||
FOR sch IN
|
||||
SELECT n.nspname
|
||||
FROM pg_class c
|
||||
JOIN pg_namespace n ON n.oid = c.relnamespace
|
||||
WHERE c.relname = 'memory_units' AND c.relkind = 'r'
|
||||
LOOP
|
||||
BEGIN
|
||||
RETURN QUERY EXECUTE format($q$
|
||||
SELECT %1$L::text, m.bank_id
|
||||
FROM %1$I.memory_units m
|
||||
JOIN %1$I.banks b ON b.bank_id = m.bank_id
|
||||
WHERE m.consolidated_at IS NULL
|
||||
AND m.consolidation_failed_at IS NULL
|
||||
AND m.fact_type IN ('experience', 'world')
|
||||
AND COALESCE(b.config -> 'enable_auto_consolidation', 'true'::jsonb) <> 'false'::jsonb
|
||||
AND NOT EXISTS (
|
||||
SELECT 1 FROM %1$I.async_operations o
|
||||
WHERE o.bank_id = m.bank_id
|
||||
AND o.operation_type = 'consolidation'
|
||||
AND o.status IN ('pending', 'processing')
|
||||
)
|
||||
GROUP BY m.bank_id
|
||||
$q$, sch);
|
||||
EXCEPTION
|
||||
-- Schema or its tables vanished between the pg_class
|
||||
-- snapshot and this query (tenant dropped or migrating).
|
||||
WHEN undefined_table OR invalid_schema_name OR undefined_column THEN
|
||||
CONTINUE;
|
||||
END;
|
||||
END LOOP;
|
||||
END;
|
||||
$fn$;
|
||||
"""
|
||||
)
|
||||
|
||||
op.execute(
|
||||
f"""
|
||||
CREATE OR REPLACE FUNCTION {schema}schemas_with_expired_rows(
|
||||
p_table text, p_ts_col text, p_days int
|
||||
)
|
||||
RETURNS SETOF text
|
||||
LANGUAGE plpgsql STABLE
|
||||
AS $fn$
|
||||
DECLARE
|
||||
sch text;
|
||||
has_expired boolean;
|
||||
BEGIN
|
||||
IF p_days IS NULL OR p_days <= 0 THEN
|
||||
RETURN;
|
||||
END IF;
|
||||
FOR sch IN
|
||||
SELECT n.nspname
|
||||
FROM pg_class c
|
||||
JOIN pg_namespace n ON n.oid = c.relnamespace
|
||||
WHERE c.relname = p_table AND c.relkind = 'r'
|
||||
LOOP
|
||||
BEGIN
|
||||
EXECUTE format(
|
||||
'SELECT EXISTS (SELECT 1 FROM %I.%I WHERE %I < NOW() - make_interval(days => $1))',
|
||||
sch, p_table, p_ts_col
|
||||
) INTO has_expired USING p_days;
|
||||
EXCEPTION
|
||||
-- Schema or its table vanished mid-scan; skip it.
|
||||
WHEN undefined_table OR invalid_schema_name OR undefined_column THEN
|
||||
CONTINUE;
|
||||
END;
|
||||
IF has_expired THEN
|
||||
RETURN NEXT sch;
|
||||
END IF;
|
||||
END LOOP;
|
||||
END;
|
||||
$fn$;
|
||||
"""
|
||||
)
|
||||
|
||||
op.execute(
|
||||
f"""
|
||||
CREATE OR REPLACE FUNCTION {schema}mental_models_with_cron()
|
||||
RETURNS TABLE(schema_name text, bank_id text, mental_model_id text,
|
||||
refresh_cron text, last_refreshed_at timestamptz)
|
||||
LANGUAGE plpgsql STABLE
|
||||
AS $fn$
|
||||
DECLARE
|
||||
sch text;
|
||||
BEGIN
|
||||
FOR sch IN
|
||||
SELECT n.nspname
|
||||
FROM pg_class c
|
||||
JOIN pg_namespace n ON n.oid = c.relnamespace
|
||||
WHERE c.relname = 'mental_models' AND c.relkind = 'r'
|
||||
LOOP
|
||||
BEGIN
|
||||
RETURN QUERY EXECUTE format($q$
|
||||
SELECT %1$L::text, mm.bank_id::text, mm.id::text,
|
||||
mm.trigger->>'refresh_cron', mm.last_refreshed_at
|
||||
FROM %1$I.mental_models mm
|
||||
WHERE COALESCE(mm.trigger->>'refresh_cron', '') <> ''
|
||||
AND NOT EXISTS (
|
||||
SELECT 1 FROM %1$I.async_operations o
|
||||
WHERE o.bank_id = mm.bank_id
|
||||
AND o.operation_type = 'refresh_mental_model'
|
||||
AND o.status IN ('pending', 'processing')
|
||||
AND o.task_payload->>'mental_model_id' = mm.id::text
|
||||
)
|
||||
$q$, sch);
|
||||
EXCEPTION
|
||||
-- Schema or its tables vanished between the pg_class
|
||||
-- snapshot and this query (tenant dropped or migrating).
|
||||
WHEN undefined_table OR invalid_schema_name OR undefined_column THEN
|
||||
CONTINUE;
|
||||
END;
|
||||
END LOOP;
|
||||
END;
|
||||
$fn$;
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def _drop_routines(schema: str | None) -> None:
|
||||
prefix = _prefix(schema)
|
||||
op.execute(f"DROP FUNCTION IF EXISTS {prefix}mental_models_with_cron()")
|
||||
op.execute(f"DROP FUNCTION IF EXISTS {prefix}schemas_with_expired_rows(text, text, int)")
|
||||
op.execute(f"DROP FUNCTION IF EXISTS {prefix}banks_needing_consolidation()")
|
||||
|
||||
|
||||
def _drop_stray_copies() -> None:
|
||||
"""Remove per-tenant duplicates left by the first cut of this migration.
|
||||
|
||||
That version installed a copy into every schema it touched, so a database
|
||||
that ran it carries one dead duplicate per tenant — only the copy in the
|
||||
configured schema is ever called. Dropping here means the next migration pass
|
||||
cleans them up; without it they would persist for the life of the database.
|
||||
|
||||
Safe on a database that never had them: ``DROP FUNCTION IF EXISTS`` is a
|
||||
no-op, and this branch never runs for the configured schema.
|
||||
"""
|
||||
_drop_routines(_target_schema())
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
# Only drop what this migration uniquely owns. When the configured schema is
|
||||
# ``public`` the copies there belong to e5f6a7b8c9d0 / f4d1c2b3a5e6, which are
|
||||
# still applied at this point and drop them on their own downgrade — removing
|
||||
# them here would strand those migrations without the functions they claim to
|
||||
# have installed.
|
||||
if not _is_install_run() or _configured_schema() == "public":
|
||||
return
|
||||
_drop_routines(_target_schema())
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade)
|
||||
-71
@@ -1,71 +0,0 @@
|
||||
"""Unique page name per folder in knowledge_pages.
|
||||
|
||||
The folder curator can fire concurrently (folder-create trigger + the
|
||||
post-consolidation sweep), and an in-process lock can't serialize runs that
|
||||
execute in different threads/loops. A partial unique index on
|
||||
(bank_id, parent, lower(name)) for pages makes duplicate-named pages in the same
|
||||
folder impossible at the DB level — the second concurrent insert fails and the
|
||||
curator treats it as "already exists".
|
||||
|
||||
PostgreSQL only: the Oracle ``name`` column is a CLOB and cannot back a
|
||||
functional unique index; Oracle relies on the in-process serialization instead.
|
||||
|
||||
Revision ID: c3d4e5f6a7b8
|
||||
Revises: a5b6c7d8e9f0
|
||||
Create Date: 2026-06-26
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "c3d4e5f6a7b8"
|
||||
down_revision: str | Sequence[str] | None = "a5b6c7d8e9f0"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _pg_schema_prefix() -> str:
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
schema = _pg_schema_prefix()
|
||||
# First drop any pre-existing duplicate pages (created by the racy curator
|
||||
# before this guard existed), keeping the earliest row of each duplicate set,
|
||||
# so the unique index can be built. Their backing mental models are left in
|
||||
# place (harmless orphans).
|
||||
op.execute(
|
||||
f"""
|
||||
DELETE FROM {schema}knowledge_pages a
|
||||
USING {schema}knowledge_pages b
|
||||
WHERE a.kind = 'page' AND b.kind = 'page'
|
||||
AND a.bank_id = b.bank_id
|
||||
AND COALESCE(a.parent_id, '') = COALESCE(b.parent_id, '')
|
||||
AND lower(a.name) = lower(b.name)
|
||||
AND a.ctid > b.ctid
|
||||
"""
|
||||
)
|
||||
# COALESCE(parent_id, '') so root-level pages (NULL parent) are also unique by
|
||||
# name — NULLs would otherwise compare distinct and allow duplicates.
|
||||
op.execute(
|
||||
"CREATE UNIQUE INDEX IF NOT EXISTS uq_kp_folder_pagename "
|
||||
f"ON {schema}knowledge_pages (bank_id, COALESCE(parent_id, ''), lower(name)) "
|
||||
"WHERE kind = 'page'"
|
||||
)
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
schema = _pg_schema_prefix()
|
||||
op.execute(f"DROP INDEX IF EXISTS {schema}uq_kp_folder_pagename")
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade) # oracle slot intentionally absent (CLOB name)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade)
|
||||
+150
@@ -0,0 +1,150 @@
|
||||
"""Add the ``schemas_with_expired_operations`` cross-tenant discovery routine.
|
||||
|
||||
The worker's terminal-operation cleanup (``a8c1e4f7b0d3``) opens a connection
|
||||
and a prune transaction against *every* tenant schema on every cleanup cycle,
|
||||
whether or not that tenant has anything to prune. At thousands of tenants that
|
||||
is a per-cycle query storm whose cost is paid entirely by idle schemas.
|
||||
|
||||
This is the same problem ``public.schemas_with_expired_rows`` already solves for
|
||||
the ``audit_log`` / ``llm_requests`` retention sweeps (``e5f6a7b8c9d0``): one
|
||||
round-trip returns just the schemas that actually hold expired rows, and the
|
||||
caller then does real work only there. ``async_operations`` needs its own
|
||||
routine rather than reusing that one because eligibility is not "row older than
|
||||
N days" — pending and processing rows are never prunable, so the status filter
|
||||
has to be part of the predicate.
|
||||
|
||||
Install policy mirrors ``b6d2f8a4c1e7`` (#2638/#2824), the current behaviour for
|
||||
the sibling routines: the routine is database-global — it enumerates ``pg_class``
|
||||
across every schema and dispatches per schema — so exactly one copy should exist,
|
||||
installed into the schema this deployment is *configured* to use and called from
|
||||
there via ``fq_routine``. Gating on the literal ``"public"`` instead of the
|
||||
configured schema is what left single-tenant deployments in a dedicated
|
||||
non-``public`` schema without the routine (#2638).
|
||||
|
||||
Exactly one migration run satisfies that predicate, so concurrent per-schema runs
|
||||
never issue competing ``CREATE OR REPLACE`` against the same ``pg_proc`` row and
|
||||
cannot hit ``tuple concurrently updated``. No cross-process coordination is
|
||||
required — in particular no advisory lock, which is unusable here because
|
||||
Hindsight runs behind connection poolers and managed PG services (see #2817).
|
||||
|
||||
Each per-schema probe runs in its own ``BEGIN ... EXCEPTION`` block so a tenant
|
||||
dropped mid-scan is skipped instead of aborting the sweep (see ``c7e9f1a3b5d2``).
|
||||
|
||||
Revision ID: d7b2f8a1c934
|
||||
Revises: b6d2f8a4c1e7
|
||||
Create Date: 2026-07-20
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
from hindsight_api.config import get_config
|
||||
|
||||
revision: str = "d7b2f8a1c934"
|
||||
down_revision: str | Sequence[str] | None = "b6d2f8a4c1e7"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _configured_schema() -> str:
|
||||
"""The one schema this deployment's routines live in and are called from."""
|
||||
return get_config().database_schema or "public"
|
||||
|
||||
|
||||
def _target_schema() -> str | None:
|
||||
return context.config.get_main_option("target_schema")
|
||||
|
||||
|
||||
def _is_install_run() -> bool:
|
||||
"""True for the single run that owns the routine (mirrors b6d2f8a4c1e7)."""
|
||||
target = _target_schema()
|
||||
return not target or target == _configured_schema()
|
||||
|
||||
|
||||
def _prefix(schema: str | None) -> str:
|
||||
"""Qualifier for ``schema``, or ``""`` to fall back to ``search_path``."""
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def _drop_routine(schema: str | None) -> None:
|
||||
op.execute(f"DROP FUNCTION IF EXISTS {_prefix(schema)}schemas_with_expired_operations(int)")
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
if not _is_install_run():
|
||||
# Tenant schemas must not carry their own copy: the routine is
|
||||
# database-global and only the configured schema's copy is ever called.
|
||||
# Dropping (rather than skipping) also cleans up after any interim build
|
||||
# of this branch that installed per-schema copies.
|
||||
_drop_routine(_target_schema())
|
||||
return
|
||||
schema = _prefix(_target_schema())
|
||||
op.execute(
|
||||
f"""
|
||||
CREATE OR REPLACE FUNCTION {schema}schemas_with_expired_operations(p_days int)
|
||||
RETURNS SETOF text
|
||||
LANGUAGE plpgsql STABLE
|
||||
AS $fn$
|
||||
DECLARE
|
||||
sch text;
|
||||
has_expired boolean;
|
||||
BEGIN
|
||||
-- Zero (or negative) retention means "keep forever": report nothing
|
||||
-- so the caller skips the sweep entirely.
|
||||
IF p_days IS NULL OR p_days <= 0 THEN
|
||||
RETURN;
|
||||
END IF;
|
||||
FOR sch IN
|
||||
SELECT n.nspname
|
||||
FROM pg_class c
|
||||
JOIN pg_namespace n ON n.oid = c.relnamespace
|
||||
WHERE c.relname = 'async_operations' AND c.relkind = 'r'
|
||||
LOOP
|
||||
BEGIN
|
||||
-- Matches the worker's prune predicate: only terminal rows
|
||||
-- are eligible, so a schema holding nothing but pending or
|
||||
-- processing work is correctly reported as having nothing
|
||||
-- to prune. Uses idx_async_operations_terminal_cleanup.
|
||||
EXECUTE format(
|
||||
'SELECT EXISTS ('
|
||||
' SELECT 1 FROM %I.async_operations'
|
||||
' WHERE status IN (''completed'', ''failed'', ''cancelled'')'
|
||||
' AND updated_at < NOW() - make_interval(days => $1)'
|
||||
')',
|
||||
sch
|
||||
) INTO has_expired USING p_days;
|
||||
EXCEPTION
|
||||
-- Schema or its table vanished between the pg_class
|
||||
-- snapshot and this probe (tenant dropped or migrating).
|
||||
WHEN undefined_table OR invalid_schema_name OR undefined_column THEN
|
||||
CONTINUE;
|
||||
END;
|
||||
IF has_expired THEN
|
||||
RETURN NEXT sch;
|
||||
END IF;
|
||||
END LOOP;
|
||||
END;
|
||||
$fn$;
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
# This migration is the sole creator of this routine — no older migration
|
||||
# owns a copy the way e5f6a7b8c9d0 owns the public sibling routines — so the
|
||||
# install run's own copy is always ours to drop.
|
||||
if not _is_install_run():
|
||||
return
|
||||
_drop_routine(_target_schema())
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
# Oracle slot intentionally absent: this mirrors the PostgreSQL-only
|
||||
# maintenance routines, and the Oracle worker keeps its per-schema sweep.
|
||||
run_for_dialect(pg=_pg_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade)
|
||||
+96
@@ -0,0 +1,96 @@
|
||||
"""Drop the search_vector column from the curation archive (invalidated_memory_units).
|
||||
|
||||
The archive is cold storage, never a recall surface, and carries no text-search
|
||||
index. Like ``embedding`` (dropped in d4f6a8c2e1b3), ``search_vector`` is a
|
||||
recall-surface column whose type follows the configured text-search backend, so
|
||||
it has no business living on the archive. Earlier curation code copied the live
|
||||
row's ``search_vector`` into ``invalidated_memory_units`` on invalidate; the
|
||||
engine now leaves it out on invalidate and recomputes it on revert, so the
|
||||
column is dead weight.
|
||||
|
||||
Dropping it removes a latent failure mode (#2503): under a non-native backend
|
||||
(pgroonga / pg_textsearch / pg_search / vchord) ``ensure_text_search_extension``
|
||||
reconciles ``memory_units.search_vector`` to ``text`` / ``bm25vector`` but never
|
||||
touched the archive, which the ``LIKE memory_units`` clone (c9a1b2d3e4f5) created
|
||||
as ``tsvector``. The type mismatch then broke the curation INSERT … SELECT
|
||||
round-trip:
|
||||
|
||||
column "search_vector" is of type tsvector but expression is of type text
|
||||
|
||||
With no column at all, there is nothing to mismatch. Unlike ``embedding`` (whose
|
||||
creation sites already omit it), the ``LIKE`` clone still adds ``search_vector``,
|
||||
so this migration does real work on both fresh and existing PostgreSQL databases.
|
||||
|
||||
DROP COLUMN is a metadata-only operation on both PostgreSQL and Oracle 23ai (no
|
||||
table rewrite), so it is cheap even across many tenant schemas. The downgrade
|
||||
re-adds an empty ``tsvector`` column (its original creation type).
|
||||
|
||||
Revision ID: e7c3a9f1b2d5
|
||||
Revises: b57a7c9e0d13
|
||||
Create Date: 2026-07-02
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "e7c3a9f1b2d5"
|
||||
down_revision: str | Sequence[str] | None = "b57a7c9e0d13"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def _pg_schema_prefix() -> str:
|
||||
"""Schema-qualifier for raw SQL on PG (multi-tenant search_path)."""
|
||||
schema = context.config.get_main_option("target_schema")
|
||||
return f'"{schema}".' if schema else ""
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
schema = _pg_schema_prefix()
|
||||
op.execute(f"ALTER TABLE {schema}invalidated_memory_units DROP COLUMN IF EXISTS search_vector")
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
schema = _pg_schema_prefix()
|
||||
# Re-add as the original tsvector creation type; comes back empty regardless.
|
||||
op.execute(f"ALTER TABLE {schema}invalidated_memory_units ADD COLUMN IF NOT EXISTS search_vector tsvector")
|
||||
|
||||
|
||||
def _oracle_upgrade() -> None:
|
||||
# Oracle has no `DROP COLUMN IF EXISTS`; swallow ORA-00904 (column does not
|
||||
# exist) so the migration is idempotent and safe on a schema whose baseline
|
||||
# may already omit the column.
|
||||
op.execute(
|
||||
"""
|
||||
BEGIN
|
||||
EXECUTE IMMEDIATE 'ALTER TABLE invalidated_memory_units DROP COLUMN search_vector';
|
||||
EXCEPTION WHEN OTHERS THEN
|
||||
IF SQLCODE != -904 THEN RAISE; END IF;
|
||||
END;
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def _oracle_downgrade() -> None:
|
||||
# Swallow ORA-01430 (column already exists) for idempotency. Oracle stores
|
||||
# search_vector as CLOB (see the Oracle baseline), so re-add it as CLOB.
|
||||
op.execute(
|
||||
"""
|
||||
BEGIN
|
||||
EXECUTE IMMEDIATE 'ALTER TABLE invalidated_memory_units ADD (search_vector CLOB)';
|
||||
EXCEPTION WHEN OTHERS THEN
|
||||
IF SQLCODE != -1430 THEN RAISE; END IF;
|
||||
END;
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
run_for_dialect(pg=_pg_upgrade, oracle=_oracle_upgrade)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade, oracle=_oracle_downgrade)
|
||||
@@ -18,7 +18,6 @@ from typing import Any, Literal, TypeVar
|
||||
from fastapi import Depends, FastAPI, File, Form, Header, HTTPException, Query, Request, UploadFile
|
||||
from fastapi.middleware.gzip import GZipMiddleware
|
||||
|
||||
from hindsight_api.api import okf
|
||||
from hindsight_api.api.disconnect import ClientDisconnectCancellationMiddleware, get_scope_cancellation_token
|
||||
from hindsight_api.cancellation import OperationCancelledError
|
||||
from hindsight_api.engine.audit import (
|
||||
@@ -53,6 +52,7 @@ from fastapi.routing import APIRoute
|
||||
from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
|
||||
|
||||
from hindsight_api import MemoryEngine
|
||||
from hindsight_api.config import RETAIN_EXTRACTION_MODES
|
||||
|
||||
|
||||
def _annotation_is_nullable(annotation: Any) -> bool:
|
||||
@@ -1246,7 +1246,7 @@ class CreateBankRequest(BaseModel):
|
||||
)
|
||||
retain_extraction_mode: str | None = Field(
|
||||
default=None,
|
||||
description="Fact extraction mode: 'concise' (default), 'verbose', or 'custom'.",
|
||||
description="Fact extraction mode: 'concise' (default), 'verbose', 'custom', 'verbatim', or 'chunks'.",
|
||||
)
|
||||
retain_custom_instructions: str | None = Field(
|
||||
default=None,
|
||||
@@ -1434,6 +1434,7 @@ class ListMemoryUnitsResponse(BaseModel):
|
||||
"date": "2024-01-15T10:30:00Z",
|
||||
"type": "world",
|
||||
"entities": "Alice (PERSON), Google (ORGANIZATION)",
|
||||
"metadata": {"source": "slack", "channel": "engineering"},
|
||||
}
|
||||
],
|
||||
"total": 150,
|
||||
@@ -1667,8 +1668,8 @@ class UpdateMemoryRequest(BaseModel):
|
||||
|
||||
@model_validator(mode="after")
|
||||
def _require_an_edit(self) -> "UpdateMemoryRequest":
|
||||
if all(
|
||||
v is None
|
||||
has_value_edit = any(
|
||||
v is not None
|
||||
for v in (
|
||||
self.text,
|
||||
self.context,
|
||||
@@ -1678,7 +1679,9 @@ class UpdateMemoryRequest(BaseModel):
|
||||
self.entities,
|
||||
self.state,
|
||||
)
|
||||
):
|
||||
)
|
||||
has_date_clear = bool({"occurred_start", "occurred_end"} & self.model_fields_set)
|
||||
if not has_value_edit and not has_date_clear:
|
||||
raise ValueError("Provide at least one field to update.")
|
||||
if self.state is not None and self.state not in ("valid", "invalidated"):
|
||||
raise ValueError("state must be 'valid' or 'invalidated'.")
|
||||
@@ -2111,150 +2114,6 @@ class MentalModelListResponse(BaseModel):
|
||||
items: list[MentalModelResponse]
|
||||
|
||||
|
||||
# =========================================================================
|
||||
# KNOWLEDGE BASE (folders + pages over mental models, projected to OKF)
|
||||
# =========================================================================
|
||||
|
||||
|
||||
class KnowledgeNode(BaseModel):
|
||||
"""A node in the knowledge-base tree — a folder or a page.
|
||||
|
||||
Pages carry ``description``/``tags`` from their backing mental model. The
|
||||
knowledge base is client-managed (CRUD); ``managed`` lets a client tag a node
|
||||
as system-owned vs. hand-authored.
|
||||
"""
|
||||
|
||||
id: str
|
||||
kind: Literal["folder", "page"]
|
||||
name: str
|
||||
parent_id: str | None = None
|
||||
mental_model_id: str | None = Field(default=None, description="Backing mental model id (pages only).")
|
||||
managed: bool = Field(default=False, description="Client-set flag: true = system-owned, false = hand-authored.")
|
||||
description: str | None = Field(default=None, description="Page source query (OKF `description`).")
|
||||
tags: list[str] = FieldWithDefault(list)
|
||||
timestamp: str | None = Field(default=None, description="Last refresh (page) or last update (folder).")
|
||||
children: list["KnowledgeNode"] = FieldWithDefault(list)
|
||||
|
||||
|
||||
class KnowledgeTreeResponse(BaseModel):
|
||||
"""The knowledge base as a nested folder/page tree."""
|
||||
|
||||
roots: list[KnowledgeNode]
|
||||
|
||||
|
||||
class CreateFolderRequest(BaseModel):
|
||||
"""Create a folder under an optional parent folder."""
|
||||
|
||||
name: str
|
||||
parent_id: str | None = None
|
||||
|
||||
|
||||
class CreatePageRequest(BaseModel):
|
||||
"""Create a page (a mental model + tree node) under an optional parent folder."""
|
||||
|
||||
name: str
|
||||
source_query: str
|
||||
parent_id: str | None = None
|
||||
tags: list[str] | None = None
|
||||
max_tokens: int | None = None
|
||||
trigger: MentalModelTrigger | None = None
|
||||
|
||||
|
||||
class UpdateNodeRequest(BaseModel):
|
||||
"""Rename and/or move a node. Each field applies only when present."""
|
||||
|
||||
name: str | None = None
|
||||
parent_id: str | None = None
|
||||
|
||||
|
||||
class CreateKnowledgePageResponse(BaseModel):
|
||||
"""Result of creating a page: the node id, its mental model, and the refresh op."""
|
||||
|
||||
page_id: str
|
||||
mental_model_id: str
|
||||
operation_id: str | None = None
|
||||
|
||||
|
||||
class KnowledgePageResponse(BaseModel):
|
||||
"""A knowledge page rendered as an OKF document."""
|
||||
|
||||
id: str
|
||||
name: str
|
||||
type: str = Field(description="OKF document type — from a `type:<x>` tag, else 'knowledge-page'.")
|
||||
description: str | None = Field(default=None, description="The source query that rebuilds the page.")
|
||||
tags: list[str] = FieldWithDefault(list)
|
||||
timestamp: str | None = Field(default=None, description="Last refresh time (falls back to creation).")
|
||||
body: str | None = Field(default=None, description="The page's synthesized markdown body.")
|
||||
markdown: str = Field(description="The full OKF document: YAML frontmatter + markdown body.")
|
||||
|
||||
|
||||
class KnowledgePageGraphResponse(BaseModel):
|
||||
"""Constellation graph of knowledge pages linked by shared tags."""
|
||||
|
||||
nodes: list[dict[str, Any]]
|
||||
edges: list[dict[str, Any]]
|
||||
total_pages: int
|
||||
total_edges: int
|
||||
|
||||
|
||||
class KnowledgePageBundleFile(BaseModel):
|
||||
"""One file in a portable OKF bundle."""
|
||||
|
||||
path: str
|
||||
content: str
|
||||
|
||||
|
||||
class KnowledgePageBundleResponse(BaseModel):
|
||||
"""A portable OKF bundle — a flat set of markdown files (index + pages + logs)."""
|
||||
|
||||
files: list[KnowledgePageBundleFile]
|
||||
|
||||
|
||||
def _knowledge_node_model(node: dict[str, Any]) -> KnowledgeNode:
|
||||
"""Project an engine node dict into a (childless) KnowledgeNode."""
|
||||
is_page = node.get("kind") == "page"
|
||||
return KnowledgeNode(
|
||||
id=node["id"],
|
||||
kind=node["kind"],
|
||||
name=node["name"],
|
||||
parent_id=node.get("parent_id"),
|
||||
mental_model_id=node.get("mental_model_id"),
|
||||
managed=bool(node.get("managed")),
|
||||
description=node.get("source_query") if is_page else None,
|
||||
tags=list(node.get("tags") or []) if is_page else [],
|
||||
timestamp=(node.get("last_refreshed_at") if is_page else node.get("updated_at")),
|
||||
)
|
||||
|
||||
|
||||
def _build_knowledge_tree(nodes: list[dict[str, Any]]) -> list[KnowledgeNode]:
|
||||
"""Assemble the flat node list into a nested tree of roots."""
|
||||
models = {n["id"]: _knowledge_node_model(n) for n in nodes}
|
||||
roots: list[KnowledgeNode] = []
|
||||
for node in nodes:
|
||||
model = models[node["id"]]
|
||||
parent_id = node.get("parent_id")
|
||||
if parent_id and parent_id in models:
|
||||
models[parent_id].children.append(model)
|
||||
else:
|
||||
roots.append(model)
|
||||
return roots
|
||||
|
||||
|
||||
def _knowledge_page_response(node: dict[str, Any]) -> KnowledgePageResponse:
|
||||
"""Project a page node (with merged mental-model content) into an OKF document."""
|
||||
page = okf.page_type(node.get("tags"))
|
||||
return KnowledgePageResponse(
|
||||
id=node["id"],
|
||||
name=node["name"],
|
||||
type=page.type,
|
||||
description=node.get("source_query"),
|
||||
tags=page.display_tags,
|
||||
timestamp=node.get("last_refreshed_at") or node.get("created_at"),
|
||||
body=node.get("content"),
|
||||
markdown=okf.render_document(node),
|
||||
)
|
||||
|
||||
|
||||
class CreateMentalModelRequest(BaseModel):
|
||||
"""Request model for creating a mental model."""
|
||||
|
||||
@@ -2348,7 +2207,8 @@ class BankTemplateConfig(BaseModel):
|
||||
reflect_mission: str | None = Field(default=None, description="Mission/context for Reflect operations")
|
||||
retain_mission: str | None = Field(default=None, description="Steers what gets extracted during retain")
|
||||
retain_extraction_mode: str | None = Field(
|
||||
default=None, description="Fact extraction mode: 'concise' (default), 'verbose', or 'custom'"
|
||||
default=None,
|
||||
description="Fact extraction mode: 'concise' (default), 'verbose', 'custom', 'verbatim', or 'chunks'",
|
||||
)
|
||||
retain_custom_instructions: str | None = Field(
|
||||
default=None, description="Custom extraction prompt (when mode='custom')"
|
||||
@@ -2574,10 +2434,10 @@ def validate_bank_template(manifest: "BankTemplateManifest") -> list[str]:
|
||||
if manifest.bank:
|
||||
bank = manifest.bank
|
||||
if bank.retain_extraction_mode is not None:
|
||||
valid_modes = ("concise", "verbose", "custom", "chunks")
|
||||
if bank.retain_extraction_mode not in valid_modes:
|
||||
if bank.retain_extraction_mode not in RETAIN_EXTRACTION_MODES:
|
||||
errors.append(
|
||||
f"bank.retain_extraction_mode: must be one of {valid_modes}, got '{bank.retain_extraction_mode}'"
|
||||
"bank.retain_extraction_mode: "
|
||||
f"must be one of {RETAIN_EXTRACTION_MODES}, got '{bank.retain_extraction_mode}'"
|
||||
)
|
||||
if bank.retain_custom_instructions and bank.retain_extraction_mode != "custom":
|
||||
errors.append("bank.retain_custom_instructions: requires retain_extraction_mode='custom'")
|
||||
@@ -3887,6 +3747,7 @@ def _register_routes(app: FastAPI):
|
||||
operation_id="update_memory",
|
||||
tags=["Memory"],
|
||||
)
|
||||
@audited("update_memory")
|
||||
async def api_update_memory(
|
||||
bank_id: str,
|
||||
memory_id: str,
|
||||
@@ -3895,13 +3756,23 @@ def _register_routes(app: FastAPI):
|
||||
):
|
||||
"""Curate a single memory unit (edit text / invalidate / revert)."""
|
||||
try:
|
||||
occurred_start = (
|
||||
""
|
||||
if "occurred_start" in request.model_fields_set and request.occurred_start is None
|
||||
else request.occurred_start
|
||||
)
|
||||
occurred_end = (
|
||||
""
|
||||
if "occurred_end" in request.model_fields_set and request.occurred_end is None
|
||||
else request.occurred_end
|
||||
)
|
||||
data = await app.state.memory.update_memory_unit(
|
||||
bank_id=bank_id,
|
||||
memory_id=memory_id,
|
||||
text=request.text,
|
||||
context=request.context,
|
||||
occurred_start=request.occurred_start,
|
||||
occurred_end=request.occurred_end,
|
||||
occurred_start=occurred_start,
|
||||
occurred_end=occurred_end,
|
||||
new_fact_type=request.fact_type,
|
||||
entities=request.entities,
|
||||
state=request.state,
|
||||
@@ -4926,333 +4797,6 @@ def _register_routes(app: FastAPI):
|
||||
logger.error(f"Error in DELETE /v1/default/banks/{bank_id}/mental-models/{mental_model_id}: {error_detail}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
# =========================================================================
|
||||
# KNOWLEDGE BASE ENDPOINTS (folders + pages, Open Knowledge Format)
|
||||
# =========================================================================
|
||||
# A hierarchy of folders and pages over mental models. Pages project to OKF
|
||||
# documents (markdown body + YAML frontmatter); see api/okf.py. The static
|
||||
# sub-paths (/tree, /folders, /pages, /graph, /export) are declared before
|
||||
# the /pages/{id} and /nodes/{id} path-parameter routes so they win.
|
||||
|
||||
@app.get(
|
||||
"/v1/default/banks/{bank_id}/knowledge-base/tree",
|
||||
response_model=KnowledgeTreeResponse,
|
||||
summary="Get the knowledge-base tree",
|
||||
description="Return the knowledge base as a nested tree of folders and pages.",
|
||||
operation_id="get_knowledge_base_tree",
|
||||
tags=["Knowledge Base"],
|
||||
)
|
||||
async def api_knowledge_base_tree(
|
||||
bank_id: str,
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
):
|
||||
"""Return the folder/page tree for a bank."""
|
||||
try:
|
||||
nodes = await app.state.memory.list_knowledge_nodes(bank_id=bank_id, request_context=request_context)
|
||||
return KnowledgeTreeResponse(roots=_build_knowledge_tree(nodes))
|
||||
except OperationValidationError as e:
|
||||
raise HTTPException(status_code=e.status_code, detail=e.reason)
|
||||
except (AuthenticationError, HTTPException):
|
||||
raise
|
||||
except Exception as e:
|
||||
import traceback
|
||||
|
||||
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
|
||||
logger.error(f"Error in GET /v1/default/banks/{bank_id}/knowledge-base/tree: {error_detail}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
@app.post(
|
||||
"/v1/default/banks/{bank_id}/knowledge-base/folders",
|
||||
response_model=KnowledgeNode,
|
||||
status_code=201,
|
||||
summary="Create a knowledge-base folder",
|
||||
description="Create a folder, optionally nested under a parent folder.",
|
||||
operation_id="create_knowledge_folder",
|
||||
tags=["Knowledge Base"],
|
||||
)
|
||||
async def api_create_knowledge_folder(
|
||||
bank_id: str,
|
||||
body: CreateFolderRequest,
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
):
|
||||
"""Create a folder node."""
|
||||
try:
|
||||
node = await app.state.memory.create_knowledge_folder(
|
||||
bank_id=bank_id,
|
||||
name=body.name,
|
||||
parent_id=body.parent_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
return _knowledge_node_model(node)
|
||||
except ValueError as e:
|
||||
raise HTTPException(status_code=400, detail=str(e))
|
||||
except OperationValidationError as e:
|
||||
raise HTTPException(status_code=e.status_code, detail=e.reason)
|
||||
except (AuthenticationError, HTTPException):
|
||||
raise
|
||||
except Exception as e:
|
||||
import traceback
|
||||
|
||||
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
|
||||
logger.error(f"Error in POST /v1/default/banks/{bank_id}/knowledge-base/folders: {error_detail}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
@app.post(
|
||||
"/v1/default/banks/{bank_id}/knowledge-base/pages",
|
||||
response_model=CreateKnowledgePageResponse,
|
||||
status_code=201,
|
||||
summary="Create a knowledge-base page",
|
||||
description="Create a page (a mental model + tree node). Content is generated asynchronously; "
|
||||
"use the returned operation_id to track completion.",
|
||||
operation_id="create_knowledge_page",
|
||||
tags=["Knowledge Base"],
|
||||
)
|
||||
async def api_create_knowledge_page(
|
||||
bank_id: str,
|
||||
body: CreatePageRequest,
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
):
|
||||
"""Create a page node (async content generation)."""
|
||||
try:
|
||||
node = await app.state.memory.create_knowledge_page(
|
||||
bank_id=bank_id,
|
||||
name=body.name,
|
||||
source_query=body.source_query,
|
||||
content="Generating content...",
|
||||
parent_id=body.parent_id,
|
||||
tags=body.tags if body.tags else None,
|
||||
max_tokens=body.max_tokens,
|
||||
trigger=body.trigger.model_dump() if body.trigger else None,
|
||||
request_context=request_context,
|
||||
)
|
||||
if node is None:
|
||||
raise HTTPException(status_code=409, detail=f"A page named '{body.name}' already exists in this folder")
|
||||
result = await app.state.memory.submit_async_refresh_mental_model(
|
||||
bank_id=bank_id,
|
||||
mental_model_id=node["mental_model_id"],
|
||||
request_context=request_context,
|
||||
)
|
||||
return CreateKnowledgePageResponse(
|
||||
page_id=node["id"],
|
||||
mental_model_id=node["mental_model_id"],
|
||||
operation_id=result["operation_id"],
|
||||
)
|
||||
except ValueError as e:
|
||||
raise HTTPException(status_code=400, detail=str(e))
|
||||
except OperationValidationError as e:
|
||||
raise HTTPException(status_code=e.status_code, detail=e.reason)
|
||||
except (AuthenticationError, HTTPException):
|
||||
raise
|
||||
except Exception as e:
|
||||
import traceback
|
||||
|
||||
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
|
||||
logger.error(f"Error in POST /v1/default/banks/{bank_id}/knowledge-base/pages: {error_detail}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
@app.get(
|
||||
"/v1/default/banks/{bank_id}/knowledge-base/graph",
|
||||
response_model=KnowledgePageGraphResponse,
|
||||
summary="Knowledge-base constellation graph",
|
||||
description="Return pages as nodes linked by shared tags, for the constellation view.",
|
||||
operation_id="get_knowledge_base_graph",
|
||||
tags=["Knowledge Base"],
|
||||
)
|
||||
async def api_knowledge_base_graph(
|
||||
bank_id: str,
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
):
|
||||
"""Return the shared-tag constellation graph for a bank's pages."""
|
||||
try:
|
||||
nodes = await app.state.memory.list_knowledge_nodes(bank_id=bank_id, request_context=request_context)
|
||||
pages = [n for n in nodes if n.get("kind") == "page"]
|
||||
# Cluster the constellation by parent folder (the knowledge base's own
|
||||
# structure) rather than by the retired type: tag.
|
||||
folder_names = {n["id"]: n["name"] for n in nodes if n.get("kind") == "folder"}
|
||||
graph = okf.knowledge_graph(pages, cluster_for=lambda p: folder_names.get(p.get("parent_id"), "Ungrouped"))
|
||||
return KnowledgePageGraphResponse(
|
||||
nodes=graph.nodes,
|
||||
edges=graph.edges,
|
||||
total_pages=len(graph.nodes),
|
||||
total_edges=len(graph.edges),
|
||||
)
|
||||
except OperationValidationError as e:
|
||||
raise HTTPException(status_code=e.status_code, detail=e.reason)
|
||||
except (AuthenticationError, HTTPException):
|
||||
raise
|
||||
except Exception as e:
|
||||
import traceback
|
||||
|
||||
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
|
||||
logger.error(f"Error in GET /v1/default/banks/{bank_id}/knowledge-base/graph: {error_detail}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
@app.get(
|
||||
"/v1/default/banks/{bank_id}/knowledge-base/export",
|
||||
response_model=KnowledgePageBundleResponse,
|
||||
summary="Export the knowledge base as an OKF bundle",
|
||||
description="Return a portable OKF bundle: a nested index.md, one <id>.md per page, and history logs.",
|
||||
operation_id="export_knowledge_base",
|
||||
tags=["Knowledge Base"],
|
||||
)
|
||||
async def api_export_knowledge_base(
|
||||
bank_id: str,
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
):
|
||||
"""Export a bank's knowledge base as a flat OKF markdown bundle."""
|
||||
try:
|
||||
nodes = await app.state.memory.list_knowledge_nodes(bank_id=bank_id, request_context=request_context)
|
||||
files = [KnowledgePageBundleFile(path=okf.INDEX_FILENAME, content=okf.render_index(nodes))]
|
||||
for node in nodes:
|
||||
if node.get("kind") != "page":
|
||||
continue
|
||||
page = await app.state.memory.get_knowledge_page(
|
||||
bank_id=bank_id, page_id=node["id"], request_context=request_context
|
||||
)
|
||||
if page is None:
|
||||
continue
|
||||
files.append(
|
||||
KnowledgePageBundleFile(path=okf.page_filename(node["id"]), content=okf.render_document(page))
|
||||
)
|
||||
if node.get("mental_model_id"):
|
||||
history = (
|
||||
await app.state.memory.get_mental_model_history(
|
||||
bank_id=bank_id,
|
||||
mental_model_id=node["mental_model_id"],
|
||||
request_context=request_context,
|
||||
)
|
||||
or []
|
||||
)
|
||||
if history:
|
||||
files.append(
|
||||
KnowledgePageBundleFile(
|
||||
path=okf.log_filename(node["id"]), content=okf.render_log(page, history)
|
||||
)
|
||||
)
|
||||
return KnowledgePageBundleResponse(files=files)
|
||||
except OperationValidationError as e:
|
||||
raise HTTPException(status_code=e.status_code, detail=e.reason)
|
||||
except (AuthenticationError, HTTPException):
|
||||
raise
|
||||
except Exception as e:
|
||||
import traceback
|
||||
|
||||
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
|
||||
logger.error(f"Error in GET /v1/default/banks/{bank_id}/knowledge-base/export: {error_detail}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
@app.get(
|
||||
"/v1/default/banks/{bank_id}/knowledge-base/pages/{page_id}",
|
||||
response_model=KnowledgePageResponse,
|
||||
summary="Get a knowledge-base page",
|
||||
description="Return a single page as an OKF document (frontmatter + markdown body).",
|
||||
operation_id="get_knowledge_page",
|
||||
tags=["Knowledge Base"],
|
||||
)
|
||||
async def api_get_knowledge_page(
|
||||
bank_id: str,
|
||||
page_id: str,
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
):
|
||||
"""Get a single page as an OKF document."""
|
||||
try:
|
||||
node = await app.state.memory.get_knowledge_page(
|
||||
bank_id=bank_id, page_id=page_id, request_context=request_context
|
||||
)
|
||||
if node is None:
|
||||
raise HTTPException(status_code=404, detail=f"Knowledge page '{page_id}' not found")
|
||||
return _knowledge_page_response(node)
|
||||
except (AuthenticationError, HTTPException):
|
||||
raise
|
||||
except OperationValidationError as e:
|
||||
raise HTTPException(status_code=e.status_code, detail=e.reason)
|
||||
except Exception as e:
|
||||
import traceback
|
||||
|
||||
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
|
||||
logger.error(f"Error in GET /v1/default/banks/{bank_id}/knowledge-base/pages/{page_id}: {error_detail}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
@app.patch(
|
||||
"/v1/default/banks/{bank_id}/knowledge-base/nodes/{node_id}",
|
||||
response_model=KnowledgeNode,
|
||||
summary="Rename or move a knowledge-base node",
|
||||
description="Rename a node (set `name`) and/or move it under another folder (set `parent_id`, "
|
||||
"null for the root).",
|
||||
operation_id="update_knowledge_node",
|
||||
tags=["Knowledge Base"],
|
||||
)
|
||||
async def api_update_knowledge_node(
|
||||
bank_id: str,
|
||||
node_id: str,
|
||||
body: UpdateNodeRequest,
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
):
|
||||
"""Rename and/or move a node."""
|
||||
try:
|
||||
updated: dict[str, Any] | None = None
|
||||
did_change = False
|
||||
if body.name is not None:
|
||||
did_change = True
|
||||
updated = await app.state.memory.rename_knowledge_node(
|
||||
bank_id=bank_id, node_id=node_id, name=body.name, request_context=request_context
|
||||
)
|
||||
# parent_id is applied only when present in the body, so passing null
|
||||
# moves the node to the root (distinct from "not provided").
|
||||
if "parent_id" in body.model_fields_set:
|
||||
did_change = True
|
||||
updated = await app.state.memory.move_knowledge_node(
|
||||
bank_id=bank_id, node_id=node_id, new_parent_id=body.parent_id, request_context=request_context
|
||||
)
|
||||
if not did_change:
|
||||
raise HTTPException(status_code=400, detail="Provide name and/or parent_id to update")
|
||||
if updated is None:
|
||||
raise HTTPException(status_code=404, detail=f"Knowledge node '{node_id}' not found")
|
||||
return _knowledge_node_model(updated)
|
||||
except ValueError as e:
|
||||
raise HTTPException(status_code=400, detail=str(e))
|
||||
except (AuthenticationError, HTTPException):
|
||||
raise
|
||||
except OperationValidationError as e:
|
||||
raise HTTPException(status_code=e.status_code, detail=e.reason)
|
||||
except Exception as e:
|
||||
import traceback
|
||||
|
||||
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
|
||||
logger.error(f"Error in PATCH /v1/default/banks/{bank_id}/knowledge-base/nodes/{node_id}: {error_detail}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
@app.delete(
|
||||
"/v1/default/banks/{bank_id}/knowledge-base/nodes/{node_id}",
|
||||
summary="Delete a knowledge-base node",
|
||||
description="Delete a folder or page and its whole subtree (pages' mental models are removed too).",
|
||||
operation_id="delete_knowledge_node",
|
||||
tags=["Knowledge Base"],
|
||||
)
|
||||
async def api_delete_knowledge_node(
|
||||
bank_id: str,
|
||||
node_id: str,
|
||||
request_context: RequestContext = Depends(get_request_context),
|
||||
):
|
||||
"""Delete a node and its subtree."""
|
||||
try:
|
||||
deleted = await app.state.memory.delete_knowledge_node(
|
||||
bank_id=bank_id, node_id=node_id, request_context=request_context
|
||||
)
|
||||
if not deleted:
|
||||
raise HTTPException(status_code=404, detail=f"Knowledge node '{node_id}' not found")
|
||||
return {"status": "deleted"}
|
||||
except (AuthenticationError, HTTPException):
|
||||
raise
|
||||
except OperationValidationError as e:
|
||||
raise HTTPException(status_code=e.status_code, detail=e.reason)
|
||||
except Exception as e:
|
||||
import traceback
|
||||
|
||||
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
|
||||
logger.error(f"Error in DELETE /v1/default/banks/{bank_id}/knowledge-base/nodes/{node_id}: {error_detail}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
# =========================================================================
|
||||
# DIRECTIVES ENDPOINTS
|
||||
# =========================================================================
|
||||
@@ -5879,8 +5423,9 @@ def _register_routes(app: FastAPI):
|
||||
"/v1/default/banks/{bank_id}/operations/{operation_id}",
|
||||
response_model=OperationStatusResponse,
|
||||
summary="Get operation status",
|
||||
description="Get the status of a specific async operation. Returns 'pending', 'completed', or 'failed'. "
|
||||
"Completed operations are removed from storage, so 'completed' means the operation finished successfully.",
|
||||
description="Get the status of a specific async operation. Returns 'pending', 'processing', 'completed', "
|
||||
"'failed', or 'cancelled'. Completed operations remain queryable with their payload for the configured "
|
||||
"retention window and are pruned afterward.",
|
||||
operation_id="get_operation_status",
|
||||
tags=["Operations"],
|
||||
)
|
||||
@@ -6120,8 +5665,12 @@ def _register_routes(app: FastAPI):
|
||||
):
|
||||
"""Create or update an agent with disposition and mission."""
|
||||
try:
|
||||
# Ensure bank exists by getting profile (auto-creates with defaults)
|
||||
await app.state.memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
# Ensure bank exists, validating create_bank only when this call
|
||||
# actually creates a missing bank.
|
||||
await app.state.memory._ensure_bank_exists(
|
||||
bank_id,
|
||||
request_context,
|
||||
)
|
||||
|
||||
# Update name if provided (stored in DB for display only, deprecated)
|
||||
if request.name is not None:
|
||||
@@ -6313,8 +5862,12 @@ def _register_routes(app: FastAPI):
|
||||
dry_run=True,
|
||||
)
|
||||
|
||||
# Ensure bank exists (auto-creates with defaults if needed)
|
||||
await app.state.memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
# Ensure bank exists, validating create_bank only when this import
|
||||
# actually creates a missing target bank.
|
||||
await app.state.memory._ensure_bank_exists(
|
||||
bank_id,
|
||||
request_context,
|
||||
)
|
||||
|
||||
return await apply_bank_template_manifest(
|
||||
memory=app.state.memory,
|
||||
|
||||
@@ -1,263 +0,0 @@
|
||||
"""Open Knowledge Format (OKF) projection for knowledge pages.
|
||||
|
||||
Knowledge pages are a *read-only* OKF view over the existing mental models: each
|
||||
mental model is projected into an OKF document — a markdown body with YAML
|
||||
frontmatter (``type`` required; ``title``/``description``/``tags``/``timestamp``
|
||||
optional) — and pages are linked into a constellation graph via shared tags.
|
||||
|
||||
See the Open Knowledge Format spec:
|
||||
https://github.com/GoogleCloudPlatform/knowledge-catalog/tree/main/okf
|
||||
|
||||
This module is intentionally pure: every function transforms the mental-model
|
||||
dicts returned by ``MemoryEngine.list_mental_models`` / ``get_mental_model`` and
|
||||
never touches the database. That keeps the OKF contract unit-testable without a
|
||||
DB or LLM and lets the HTTP layer stay a thin wrapper.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
# OKF requires exactly one frontmatter field — ``type``. We default to this when
|
||||
# a page does not declare one via a ``type:<x>`` tag.
|
||||
DEFAULT_PAGE_TYPE = "knowledge-page"
|
||||
|
||||
# A page declares its OKF ``type`` through a tag of the form ``type:runbook``.
|
||||
# This keeps the projection schema-free (no new mental_models column): the type
|
||||
# is lifted from the existing tags array.
|
||||
TYPE_TAG_PREFIX = "type:"
|
||||
|
||||
INDEX_FILENAME = "index.md"
|
||||
|
||||
# Deterministic, colour-blind-friendly palette. Type → colour is stable across
|
||||
# requests so the constellation keeps the same colours between reloads.
|
||||
_PALETTE = (
|
||||
"#0074d9", # blue
|
||||
"#2ecc40", # green
|
||||
"#b10dc9", # purple
|
||||
"#ff851b", # orange
|
||||
"#39cccc", # teal
|
||||
"#f012be", # magenta
|
||||
"#3d9970", # olive
|
||||
"#ff4136", # red
|
||||
)
|
||||
|
||||
_EDGE_COLOR = "#9aa5b1"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class PageType:
|
||||
"""A page's OKF ``type`` and the tags that remain after the type tag is split off."""
|
||||
|
||||
type: str
|
||||
display_tags: list[str]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class KnowledgeGraph:
|
||||
"""Cytoscape-style node/edge graph of knowledge pages linked by shared tags."""
|
||||
|
||||
nodes: list[dict[str, Any]] = field(default_factory=list)
|
||||
edges: list[dict[str, Any]] = field(default_factory=list)
|
||||
|
||||
|
||||
def _color_for(key: str) -> str:
|
||||
"""Stable colour for a string key (FNV-ish hash into the fixed palette)."""
|
||||
h = 0
|
||||
for ch in key:
|
||||
h = (h * 31 + ord(ch)) & 0xFFFFFFFF
|
||||
return _PALETTE[h % len(_PALETTE)]
|
||||
|
||||
|
||||
def _scalar(value: Any) -> str:
|
||||
"""Emit a YAML-safe double-quoted scalar.
|
||||
|
||||
We always double-quote so arbitrary page names / source queries can't be
|
||||
misread as YAML special forms (``true``, ``2026-01-01``, ``- x``, etc.).
|
||||
"""
|
||||
text = str(value)
|
||||
escaped = text.replace("\\", "\\\\").replace('"', '\\"').replace("\n", "\\n").replace("\r", "")
|
||||
return f'"{escaped}"'
|
||||
|
||||
|
||||
def page_type(tags: list[str] | None) -> PageType:
|
||||
"""Split an OKF ``type`` out of the tag list.
|
||||
|
||||
The first ``type:<x>`` tag wins; all ``type:`` tags are removed from the
|
||||
returned ``display_tags`` so they don't pollute the constellation's
|
||||
shared-tag edges. Falls back to :data:`DEFAULT_PAGE_TYPE`.
|
||||
"""
|
||||
resolved = DEFAULT_PAGE_TYPE
|
||||
display: list[str] = []
|
||||
for tag in tags or []:
|
||||
if tag.startswith(TYPE_TAG_PREFIX):
|
||||
suffix = tag[len(TYPE_TAG_PREFIX) :].strip()
|
||||
if suffix and resolved == DEFAULT_PAGE_TYPE:
|
||||
resolved = suffix
|
||||
continue
|
||||
display.append(tag)
|
||||
return PageType(type=resolved, display_tags=display)
|
||||
|
||||
|
||||
def _timestamp(mm: dict[str, Any]) -> str | None:
|
||||
return mm.get("last_refreshed_at") or mm.get("created_at")
|
||||
|
||||
|
||||
def frontmatter(mm: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Build the ordered OKF frontmatter mapping for a mental model.
|
||||
|
||||
``None``/empty values are dropped by :func:`render_frontmatter`.
|
||||
"""
|
||||
pt = page_type(mm.get("tags"))
|
||||
return {
|
||||
"id": mm.get("id"),
|
||||
"type": pt.type,
|
||||
"title": mm.get("name"),
|
||||
"description": mm.get("source_query"),
|
||||
"tags": pt.display_tags,
|
||||
"timestamp": _timestamp(mm),
|
||||
}
|
||||
|
||||
|
||||
def render_frontmatter(fm: dict[str, Any]) -> str:
|
||||
"""Render a frontmatter mapping into a ``---`` fenced YAML block."""
|
||||
lines = ["---"]
|
||||
for key, value in fm.items():
|
||||
if value is None:
|
||||
continue
|
||||
if isinstance(value, list):
|
||||
if not value:
|
||||
continue
|
||||
lines.append(f"{key}:")
|
||||
lines.extend(f" - {_scalar(item)}" for item in value)
|
||||
else:
|
||||
lines.append(f"{key}: {_scalar(value)}")
|
||||
lines.append("---")
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
def render_document(mm: dict[str, Any]) -> str:
|
||||
"""Render a full OKF document: frontmatter block + markdown body."""
|
||||
body = (mm.get("content") or "").strip()
|
||||
return f"{render_frontmatter(frontmatter(mm))}\n\n{body}\n" if body else f"{render_frontmatter(frontmatter(mm))}\n"
|
||||
|
||||
|
||||
def page_filename(page_id: str) -> str:
|
||||
"""OKF bundle filename for a page id."""
|
||||
return f"{page_id}.md"
|
||||
|
||||
|
||||
def log_filename(page_id: str) -> str:
|
||||
"""OKF reserved per-page history filename."""
|
||||
return f"{page_id}.log.md"
|
||||
|
||||
|
||||
def render_index(nodes: list[dict[str, Any]]) -> str:
|
||||
"""Render the reserved ``index.md`` — nested OKF navigation over the tree.
|
||||
|
||||
``nodes`` is the flat folder/page list (each with ``id``, ``kind``, ``name``,
|
||||
``parent_id``); folders nest their children, pages link to their ``.md``.
|
||||
"""
|
||||
fm = render_frontmatter({"type": "index", "title": "Knowledge base"})
|
||||
lines = [fm, "", "# Knowledge base", ""]
|
||||
|
||||
children: dict[Any, list[dict[str, Any]]] = {}
|
||||
for node in nodes:
|
||||
children.setdefault(node.get("parent_id"), []).append(node)
|
||||
|
||||
def walk(parent: Any, depth: int) -> None:
|
||||
ordered = sorted(children.get(parent, []), key=lambda n: (n.get("sort_order", 0), n.get("name") or ""))
|
||||
for node in ordered:
|
||||
indent = " " * depth
|
||||
if node.get("kind") == "folder":
|
||||
lines.append(f"{indent}- **{node['name']}/**")
|
||||
walk(node["id"], depth + 1)
|
||||
else:
|
||||
description = node.get("source_query") or node.get("description")
|
||||
link = f"{indent}- [{node['name']}](./{page_filename(node['id'])})"
|
||||
lines.append(f"{link} — {description}" if description else link)
|
||||
|
||||
walk(None, 0)
|
||||
if len(lines) == 4:
|
||||
lines.append("_No knowledge pages yet._")
|
||||
return "\n".join(lines) + "\n"
|
||||
|
||||
|
||||
def render_log(mm: dict[str, Any], history: list[dict[str, Any]]) -> str:
|
||||
"""Render the reserved per-page ``log.md`` from refresh history.
|
||||
|
||||
Each history entry is ``{previous_content, previous_reflect_response,
|
||||
changed_at}`` (newest first), capturing the content *before* a refresh.
|
||||
"""
|
||||
name = mm.get("name") or mm.get("id")
|
||||
fm = render_frontmatter({"type": "log", "title": f"{name} — history"})
|
||||
lines = [fm, "", f"# {name} — history", ""]
|
||||
if not history:
|
||||
lines.append("_No refresh history._")
|
||||
return "\n".join(lines) + "\n"
|
||||
for entry in history:
|
||||
changed_at = entry.get("changed_at") or "unknown"
|
||||
previous = (entry.get("previous_content") or "").strip()
|
||||
lines.append(f"## {changed_at}")
|
||||
lines.append("")
|
||||
lines.append(previous if previous else "_(empty)_")
|
||||
lines.append("")
|
||||
return "\n".join(lines).rstrip() + "\n"
|
||||
|
||||
|
||||
def knowledge_graph(
|
||||
pages: list[dict[str, Any]],
|
||||
cluster_for: "Callable[[dict[str, Any]], str] | None" = None,
|
||||
) -> KnowledgeGraph:
|
||||
"""Derive the constellation graph: pages as nodes, shared tags as edges.
|
||||
|
||||
Two pages are linked when they share at least one (non-``type:``) tag; the
|
||||
edge weight is the number of shared tags. Each node's cluster (``type`` field
|
||||
+ colour) comes from ``cluster_for(page)`` — the knowledge base groups by
|
||||
parent folder; the default groups by OKF ``type``.
|
||||
"""
|
||||
nodes: list[dict[str, Any]] = []
|
||||
tag_sets: list[tuple[str, frozenset[str]]] = []
|
||||
for mm in pages:
|
||||
page_id = mm["id"]
|
||||
pt = page_type(mm.get("tags"))
|
||||
cluster = cluster_for(mm) if cluster_for else pt.type
|
||||
tag_sets.append((page_id, frozenset(pt.display_tags)))
|
||||
nodes.append(
|
||||
{
|
||||
"data": {
|
||||
"id": page_id,
|
||||
"label": mm.get("name") or page_id,
|
||||
"type": cluster,
|
||||
"tagCount": len(pt.display_tags),
|
||||
"color": _color_for(cluster),
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
edges: list[dict[str, Any]] = []
|
||||
for i in range(len(tag_sets)):
|
||||
source_id, source_tags = tag_sets[i]
|
||||
if not source_tags:
|
||||
continue
|
||||
for j in range(i + 1, len(tag_sets)):
|
||||
target_id, target_tags = tag_sets[j]
|
||||
shared = source_tags & target_tags
|
||||
if not shared:
|
||||
continue
|
||||
edges.append(
|
||||
{
|
||||
"data": {
|
||||
"id": f"{source_id}--{target_id}",
|
||||
"source": source_id,
|
||||
"target": target_id,
|
||||
"sharedTags": sorted(shared),
|
||||
"weight": len(shared),
|
||||
"color": _EDGE_COLOR,
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
return KnowledgeGraph(nodes=nodes, edges=edges)
|
||||
@@ -147,6 +147,7 @@ ENV_LLM_EXTRA_BODY = "HINDSIGHT_API_LLM_EXTRA_BODY"
|
||||
ENV_LLM_DEFAULT_HEADERS = "HINDSIGHT_API_LLM_DEFAULT_HEADERS"
|
||||
ENV_LLM_STRICT_SCHEMA = "HINDSIGHT_API_LLM_STRICT_SCHEMA"
|
||||
ENV_LLM_SEND_BANK_AS_USER = "HINDSIGHT_API_LLM_SEND_BANK_AS_USER"
|
||||
ENV_LLM_OLLAMA_NUM_CTX = "HINDSIGHT_API_LLM_OLLAMA_NUM_CTX"
|
||||
|
||||
# Per-operation sampling temperature. Each internal LLM call uses a temperature
|
||||
# tuned for its task (deterministic extraction vs. creative reflection). These
|
||||
@@ -373,6 +374,7 @@ ENV_EMBEDDINGS_LITELLM_SDK_MODEL = "HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_MODEL"
|
||||
ENV_EMBEDDINGS_LITELLM_SDK_API_BASE = "HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_API_BASE"
|
||||
ENV_EMBEDDINGS_LITELLM_SDK_OUTPUT_DIMENSIONS = "HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_OUTPUT_DIMENSIONS"
|
||||
ENV_EMBEDDINGS_LITELLM_SDK_ENCODING_FORMAT = "HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_ENCODING_FORMAT"
|
||||
ENV_EMBEDDINGS_LITELLM_SDK_MAX_INPUT_TOKENS = "HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_MAX_INPUT_TOKENS"
|
||||
ENV_RERANKER_LITELLM_SDK_API_KEY = "HINDSIGHT_API_RERANKER_LITELLM_SDK_API_KEY"
|
||||
ENV_RERANKER_LITELLM_SDK_MODEL = "HINDSIGHT_API_RERANKER_LITELLM_SDK_MODEL"
|
||||
ENV_RERANKER_LITELLM_SDK_API_BASE = "HINDSIGHT_API_RERANKER_LITELLM_SDK_API_BASE"
|
||||
@@ -382,6 +384,7 @@ ENV_LITELLM_API_BASE = "HINDSIGHT_API_LITELLM_API_BASE"
|
||||
ENV_LITELLM_API_KEY = "HINDSIGHT_API_LITELLM_API_KEY"
|
||||
|
||||
ENV_RERANKER_PROVIDER = "HINDSIGHT_API_RERANKER_PROVIDER"
|
||||
ENV_RERANKER_SEND_BANK_AS_HEADER = "HINDSIGHT_API_RERANKER_SEND_BANK_AS_HEADER"
|
||||
ENV_RERANKER_LOCAL_MODEL = "HINDSIGHT_API_RERANKER_LOCAL_MODEL"
|
||||
ENV_RERANKER_LOCAL_FORCE_CPU = "HINDSIGHT_API_RERANKER_LOCAL_FORCE_CPU"
|
||||
ENV_RERANKER_LOCAL_MAX_CONCURRENT = "HINDSIGHT_API_RERANKER_LOCAL_MAX_CONCURRENT"
|
||||
@@ -584,6 +587,7 @@ ENV_DB_POOL_MAX_SIZE = "HINDSIGHT_API_DB_POOL_MAX_SIZE"
|
||||
ENV_DB_COMMAND_TIMEOUT = "HINDSIGHT_API_DB_COMMAND_TIMEOUT"
|
||||
ENV_DB_ACQUIRE_TIMEOUT = "HINDSIGHT_API_DB_ACQUIRE_TIMEOUT"
|
||||
ENV_DB_STATEMENT_TIMEOUT = "HINDSIGHT_API_DB_STATEMENT_TIMEOUT"
|
||||
ENV_DB_MAX_PARALLEL_WORKERS_PER_GATHER = "HINDSIGHT_API_DB_MAX_PARALLEL_WORKERS_PER_GATHER"
|
||||
|
||||
# Wall-clock cap on model/connection initialization at startup. If embeddings,
|
||||
# cross-encoder, or LLM verification hang (e.g. an offline HuggingFace download
|
||||
@@ -598,6 +602,8 @@ ENV_WORKER_MAX_RETRIES = "HINDSIGHT_API_WORKER_MAX_RETRIES"
|
||||
ENV_WORKER_TASK_RETRY_BACKOFF_SECONDS = "HINDSIGHT_API_WORKER_TASK_RETRY_BACKOFF_SECONDS"
|
||||
ENV_WORKER_HTTP_PORT = "HINDSIGHT_API_WORKER_HTTP_PORT"
|
||||
ENV_WORKER_MAX_SLOTS = "HINDSIGHT_API_WORKER_MAX_SLOTS"
|
||||
ENV_OPERATION_RETENTION_DAYS = "HINDSIGHT_API_OPERATION_RETENTION_DAYS"
|
||||
ENV_OPERATION_CLEANUP_BATCH_SIZE = "HINDSIGHT_API_OPERATION_CLEANUP_BATCH_SIZE"
|
||||
|
||||
# Per-operation-type slot reservations. Each entry maps an operation_type
|
||||
# (as stored in async_operations.operation_type) to its env var and default.
|
||||
@@ -617,6 +623,7 @@ ENV_RETAIN_MAX_CONCURRENT = "HINDSIGHT_API_RETAIN_MAX_CONCURRENT"
|
||||
|
||||
# Reflect agent settings
|
||||
ENV_REFLECT_MAX_ITERATIONS = "HINDSIGHT_API_REFLECT_MAX_ITERATIONS"
|
||||
ENV_REFLECT_PROMPT_CACHE_ENABLED = "HINDSIGHT_API_REFLECT_PROMPT_CACHE_ENABLED"
|
||||
ENV_REFLECT_MAX_CONTEXT_TOKENS = "HINDSIGHT_API_REFLECT_MAX_CONTEXT_TOKENS"
|
||||
ENV_REFLECT_WALL_TIMEOUT = "HINDSIGHT_API_REFLECT_WALL_TIMEOUT"
|
||||
ENV_REFLECT_MISSION = "HINDSIGHT_API_REFLECT_MISSION"
|
||||
@@ -638,6 +645,7 @@ ENV_RECALL_BUDGET_MAX = "HINDSIGHT_API_RECALL_BUDGET_MAX"
|
||||
|
||||
# Recall candidate gating (per-source cap + BM25 score floor)
|
||||
ENV_BM25_MIN_SCORE = "HINDSIGHT_API_BM25_MIN_SCORE"
|
||||
ENV_BM25_MAX_QUERY_TERMS = "HINDSIGHT_API_BM25_MAX_QUERY_TERMS"
|
||||
ENV_RECALL_MAX_CANDIDATES_PER_SOURCE = "HINDSIGHT_API_RECALL_MAX_CANDIDATES_PER_SOURCE"
|
||||
# Per-strategy recall boost. Prioritises specific retrieval arms (semantic,
|
||||
# bm25, graph, temporal) on recall via a human priority level — e.g.
|
||||
@@ -661,6 +669,9 @@ ENV_AUDIT_LOG_ENABLED = "HINDSIGHT_API_AUDIT_LOG_ENABLED"
|
||||
ENV_AUDIT_LOG_ACTIONS = "HINDSIGHT_API_AUDIT_LOG_ACTIONS"
|
||||
ENV_AUDIT_LOG_RETENTION_DAYS = "HINDSIGHT_API_AUDIT_LOG_RETENTION_DAYS"
|
||||
|
||||
# Retain reliability settings
|
||||
ENV_FAIL_ON_EXTRACTION_ERRORS = "HINDSIGHT_API_FAIL_ON_EXTRACTION_ERRORS"
|
||||
|
||||
# LLM request tracing settings
|
||||
ENV_LLM_TRACE_ENABLED = "HINDSIGHT_API_LLM_TRACE_ENABLED"
|
||||
ENV_LLM_TRACE_SCOPES = "HINDSIGHT_API_LLM_TRACE_SCOPES"
|
||||
@@ -761,6 +772,7 @@ DEFAULT_EMBEDDINGS_GEMINI_FORCE_IPV4 = False
|
||||
DEFAULT_EMBEDDING_DIMENSION = 384
|
||||
|
||||
DEFAULT_RERANKER_PROVIDER = "local"
|
||||
DEFAULT_RERANKER_SEND_BANK_AS_HEADER = False
|
||||
DEFAULT_RERANKER_LOCAL_MODEL = "cross-encoder/ms-marco-MiniLM-L-6-v2"
|
||||
DEFAULT_RERANKER_LOCAL_FORCE_CPU = False # Force CPU mode for local reranker
|
||||
DEFAULT_RERANKER_LOCAL_MAX_CONCURRENT = 4 # Limit concurrent CPU-bound reranking to prevent thrashing
|
||||
@@ -789,6 +801,9 @@ DEFAULT_SEMANTIC_MIN_SIMILARITY = 0.3
|
||||
# zero-score (non-matching) rows on backends — notably VectorChord — whose
|
||||
# operator ranks every document rather than pre-filtering to term matches.
|
||||
DEFAULT_BM25_MIN_SCORE = 0.0
|
||||
# Native tsvector BM25 can optionally cap the OR tsquery built from normalized
|
||||
# query tokens. 0 preserves the historical uncapped behavior.
|
||||
DEFAULT_BM25_MAX_QUERY_TERMS = 0
|
||||
# Per-source candidate cap applied to each retrieval arm (semantic, BM25, graph,
|
||||
# temporal) before RRF, so a single over-expanding backend cannot fill the
|
||||
# reranker's global candidate budget on its own. 0 disables the cap.
|
||||
@@ -909,6 +924,10 @@ DEFAULT_RERANKER_LITELLM_MAX_TOKENS_PER_DOC: int | None = None
|
||||
# LiteLLM SDK defaults
|
||||
DEFAULT_EMBEDDINGS_LITELLM_SDK_MODEL = "cohere/embed-english-v3.0"
|
||||
DEFAULT_EMBEDDINGS_LITELLM_SDK_ENCODING_FORMAT = "float"
|
||||
# Opt-in per-text input truncation (tiktoken cl100k_base tokens). Off by default;
|
||||
# set to the embedding model's real input limit (e.g. 8192 for Bedrock Titan V2)
|
||||
# to keep oversized content from permanently failing the embed call. See #2501.
|
||||
DEFAULT_EMBEDDINGS_LITELLM_SDK_MAX_INPUT_TOKENS: int | None = None
|
||||
DEFAULT_RERANKER_LITELLM_SDK_MODEL = "cohere/rerank-english-v3.0"
|
||||
|
||||
DEFAULT_HOST = "0.0.0.0"
|
||||
@@ -1040,6 +1059,14 @@ DEFAULT_DB_POOL_MAX_SIZE = 100
|
||||
DEFAULT_DB_COMMAND_TIMEOUT = 60 # seconds
|
||||
DEFAULT_DB_ACQUIRE_TIMEOUT = 30 # seconds
|
||||
DEFAULT_DB_STATEMENT_TIMEOUT = 600 # seconds (Postgres statement_timeout applied on every pool connection; 0 disables)
|
||||
# Optional cap on Postgres planner parallelism for this process's pool
|
||||
# connections (SET max_parallel_workers_per_gather). None leaves the server
|
||||
# default untouched. Setting 0 on background-worker processes keeps bulk
|
||||
# maintenance queries (consolidation, graph upkeep) from fanning out across
|
||||
# cores that latency-sensitive foreground traffic is sharing — parallel
|
||||
# workers buy latency, which background work doesn't need, at the cost of
|
||||
# concurrent CPU footprint, which multi-tenant primaries do care about.
|
||||
DEFAULT_DB_MAX_PARALLEL_WORKERS_PER_GATHER: int | None = None
|
||||
DEFAULT_MODEL_INIT_TIMEOUT = 300 # seconds (cap on startup model/connection init; covers first-time downloads)
|
||||
|
||||
# Worker configuration (distributed task processing)
|
||||
@@ -1050,10 +1077,20 @@ DEFAULT_WORKER_MAX_RETRIES = 3 # Max retries before marking task failed
|
||||
DEFAULT_WORKER_TASK_RETRY_BACKOFF_SECONDS = 60 # Seconds between retries on transient task failure
|
||||
DEFAULT_WORKER_HTTP_PORT = 8889 # HTTP port for worker metrics/health
|
||||
DEFAULT_WORKER_MAX_SLOTS = 10 # Total concurrent tasks per worker
|
||||
# Terminal rows keep their payload and metadata for one coherent debug/retry TTL.
|
||||
# Zero retention days disables automatic pruning entirely, and is the default:
|
||||
# operation history is a user-visible audit trail, so bounding it is an opt-in
|
||||
# policy decision rather than something an upgrade silently applies.
|
||||
DEFAULT_OPERATION_RETENTION_DAYS = 0
|
||||
DEFAULT_OPERATION_CLEANUP_BATCH_SIZE = 1000
|
||||
DEFAULT_RETAIN_MAX_CONCURRENT = 4 # Max concurrent retain DB phases (HNSW reads + writes). Limits I/O contention.
|
||||
|
||||
# Reflect agent settings
|
||||
DEFAULT_REFLECT_MAX_ITERATIONS = 10 # Max tool call iterations before forcing response
|
||||
# Step-by-step context caching for the reflect tool loop (Gemini). On by default;
|
||||
# requires the global prompt cache (HINDSIGHT_API_LLM_PROMPT_CACHE_ENABLED) to also
|
||||
# be on. Set false to force reflect to run uncached even when prompt caching is on.
|
||||
DEFAULT_REFLECT_PROMPT_CACHE_ENABLED = True
|
||||
DEFAULT_REFLECT_MAX_CONTEXT_TOKENS = 100_000 # Max accumulated context tokens before forcing final prompt
|
||||
DEFAULT_REFLECT_WALL_TIMEOUT = 300 # Wall-clock timeout in seconds for the entire reflect operation (5 minutes)
|
||||
DEFAULT_REFLECT_SOURCE_FACTS_MAX_TOKENS = -1 # Token budget for source facts in search_observations (-1 = disabled)
|
||||
@@ -1094,6 +1131,11 @@ DEFAULT_AUDIT_LOG_ENABLED = False # Disabled by default
|
||||
DEFAULT_AUDIT_LOG_ACTIONS = "" # Empty = audit all eligible actions
|
||||
DEFAULT_AUDIT_LOG_RETENTION_DAYS = -1 # -1 = keep forever
|
||||
|
||||
# Retain reliability defaults
|
||||
DEFAULT_FAIL_ON_EXTRACTION_ERRORS = (
|
||||
False # Preserve existing behavior: retain completes even if some chunks fail extraction
|
||||
)
|
||||
|
||||
# LLM request tracing defaults
|
||||
DEFAULT_LLM_TRACE_ENABLED = True # Enabled by default
|
||||
DEFAULT_LLM_TRACE_SCOPES = "" # Empty = trace all call scopes
|
||||
@@ -1209,6 +1251,19 @@ def _parse_positive_int(name: str, raw: str | None, default: int) -> int:
|
||||
return parsed
|
||||
|
||||
|
||||
def _parse_non_negative_int(name: str, raw: str | None, default: int) -> int:
|
||||
"""Parse an env var that must be an integer >= 0."""
|
||||
if raw is None or raw == "":
|
||||
return default
|
||||
try:
|
||||
parsed = int(raw)
|
||||
except ValueError as e:
|
||||
raise ValueError(f"{name} must be an integer, got {raw!r}") from e
|
||||
if parsed < 0:
|
||||
raise ValueError(f"{name} must be >= 0, got {parsed}")
|
||||
return parsed
|
||||
|
||||
|
||||
def _parse_optional_positive_int(name: str, raw: str | None) -> int | None:
|
||||
"""Parse an optional env var that must be a positive integer when set."""
|
||||
if raw is None or raw == "":
|
||||
@@ -1216,6 +1271,25 @@ def _parse_optional_positive_int(name: str, raw: str | None) -> int | None:
|
||||
return _parse_positive_int(name, raw, 1)
|
||||
|
||||
|
||||
def _parse_optional_non_negative_int(name: str, raw: str | None) -> int | None:
|
||||
"""
|
||||
Parse an optional env var that must be a non-negative integer when set.
|
||||
|
||||
Unlike ``_parse_optional_positive_int``, 0 is a meaningful value here —
|
||||
e.g. ``max_parallel_workers_per_gather = 0`` disables planner parallelism
|
||||
entirely. Unset/empty means "no opinion" (None).
|
||||
"""
|
||||
if raw is None or raw == "":
|
||||
return None
|
||||
try:
|
||||
parsed = int(raw)
|
||||
except ValueError as e:
|
||||
raise ValueError(f"{name} must be an integer, got {raw!r}") from e
|
||||
if parsed < 0:
|
||||
raise ValueError(f"{name} must be >= 0, got {parsed}")
|
||||
return parsed
|
||||
|
||||
|
||||
def _validate_retain_chunking_int(name: str, value: Any) -> int:
|
||||
if isinstance(value, bool) or not isinstance(value, int):
|
||||
raise ValueError(f"{name} must be an integer, got {value!r}")
|
||||
@@ -1577,6 +1651,9 @@ class HindsightConfig:
|
||||
# LiteLLM, Helicone) key attribution on the OpenAI `user` field. Opt-in; never
|
||||
# overrides a `user` the caller already set.
|
||||
llm_send_bank_as_user: bool
|
||||
# Optional native Ollama context window override. Unset lets Ollama use the
|
||||
# model/server default instead of forcing a Hindsight-wide value.
|
||||
llm_ollama_num_ctx: int | None = field(default=None, kw_only=True)
|
||||
|
||||
# Per-operation sampling temperature. None means the temperature parameter is
|
||||
# omitted from the call (for models that reject explicit temperatures). See
|
||||
@@ -1685,6 +1762,7 @@ class HindsightConfig:
|
||||
embeddings_litellm_sdk_api_base: str | None
|
||||
embeddings_litellm_sdk_output_dimensions: int | None
|
||||
embeddings_litellm_sdk_encoding_format: str | None
|
||||
embeddings_litellm_sdk_max_input_tokens: int | None
|
||||
# Gemini/Vertex AI embeddings
|
||||
embeddings_gemini_api_key: str | None
|
||||
embeddings_gemini_model: str
|
||||
@@ -1696,6 +1774,7 @@ class HindsightConfig:
|
||||
|
||||
# Reranker
|
||||
reranker_provider: str
|
||||
reranker_send_bank_as_header: bool
|
||||
reranker_local_model: str
|
||||
reranker_local_force_cpu: bool
|
||||
reranker_local_max_concurrent: int
|
||||
@@ -1896,6 +1975,7 @@ class HindsightConfig:
|
||||
db_command_timeout: int
|
||||
db_acquire_timeout: int
|
||||
db_statement_timeout: int
|
||||
db_max_parallel_workers_per_gather: int | None
|
||||
model_init_timeout: float
|
||||
|
||||
# Worker configuration (distributed task processing)
|
||||
@@ -1908,12 +1988,15 @@ class HindsightConfig:
|
||||
worker_max_slots: int
|
||||
worker_slot_reservations: dict[str, int]
|
||||
worker_consolidation_bank_priority: dict[str, int]
|
||||
operation_retention_days: int
|
||||
operation_cleanup_batch_size: int
|
||||
retain_max_concurrent: int
|
||||
|
||||
# Reflect agent settings
|
||||
reflect_max_iterations: int
|
||||
reflect_max_context_tokens: int
|
||||
reflect_wall_timeout: int
|
||||
reflect_prompt_cache_enabled: bool
|
||||
|
||||
# OpenTelemetry tracing configuration
|
||||
otel_traces_enabled: bool
|
||||
@@ -1929,6 +2012,11 @@ class HindsightConfig:
|
||||
audit_log_actions: list[str] # Allowlist of action types (empty = all)
|
||||
audit_log_retention_days: int # -1 = keep forever, >0 = delete after N days
|
||||
|
||||
# Retain reliability configuration (static - server-level only)
|
||||
# When True, a retain operation that accumulated any fact-extraction errors is
|
||||
# marked 'failed' instead of 'completed', surfacing silent fact loss to clients.
|
||||
fail_on_extraction_errors: bool
|
||||
|
||||
# LLM request tracing configuration (static - server-level only)
|
||||
llm_trace_enabled: bool # Master switch for per-bank LLM request tracing
|
||||
llm_trace_scopes: list[str] # Allowlist of call scopes to trace (empty = all)
|
||||
@@ -1979,6 +2067,7 @@ class HindsightConfig:
|
||||
reflect_llm_strategy: LLMStrategyConfig | None = None
|
||||
consolidation_llm_members: list[LLMMemberConfig] = field(default_factory=list)
|
||||
consolidation_llm_strategy: LLMStrategyConfig | None = None
|
||||
bm25_max_query_terms: int = DEFAULT_BM25_MAX_QUERY_TERMS
|
||||
|
||||
# Class-level sets for configuration categorization
|
||||
|
||||
@@ -2179,6 +2268,9 @@ class HindsightConfig:
|
||||
f"Invalid semantic_min_similarity: {self.semantic_min_similarity}. Must be between 0.0 and 1.0"
|
||||
)
|
||||
|
||||
if self.bm25_max_query_terms < 0:
|
||||
raise ValueError(f"Invalid bm25_max_query_terms: {self.bm25_max_query_terms}. Must be >= 0")
|
||||
|
||||
# Validate bedrock_service_tier
|
||||
valid_bedrock_tiers = (None, "flex", "priority", "reserved")
|
||||
if self.llm_bedrock_service_tier not in valid_bedrock_tiers:
|
||||
@@ -2268,6 +2360,13 @@ class HindsightConfig:
|
||||
f"Reduce reservations or increase HINDSIGHT_API_WORKER_MAX_SLOTS."
|
||||
)
|
||||
|
||||
if self.operation_retention_days < 0:
|
||||
raise ValueError(f"{ENV_OPERATION_RETENTION_DAYS} must be >= 0, got {self.operation_retention_days}")
|
||||
if self.operation_cleanup_batch_size < 1:
|
||||
raise ValueError(
|
||||
f"{ENV_OPERATION_CLEANUP_BATCH_SIZE} must be >= 1, got {self.operation_cleanup_batch_size}"
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_env(cls) -> "HindsightConfig":
|
||||
"""Create configuration from environment variables."""
|
||||
@@ -2319,6 +2418,10 @@ class HindsightConfig:
|
||||
llm_strict_schema=os.getenv(ENV_LLM_STRICT_SCHEMA, str(DEFAULT_LLM_STRICT_SCHEMA)).lower() in ("true", "1"),
|
||||
llm_send_bank_as_user=os.getenv(ENV_LLM_SEND_BANK_AS_USER, str(DEFAULT_LLM_SEND_BANK_AS_USER)).lower()
|
||||
in ("true", "1"),
|
||||
llm_ollama_num_ctx=_parse_optional_positive_int(
|
||||
ENV_LLM_OLLAMA_NUM_CTX,
|
||||
os.getenv(ENV_LLM_OLLAMA_NUM_CTX),
|
||||
),
|
||||
llm_temperature_verification=_resolve_operation_temperature(
|
||||
ENV_LLM_TEMPERATURE_VERIFICATION, DEFAULT_LLM_TEMPERATURE_VERIFICATION
|
||||
),
|
||||
@@ -2555,6 +2658,9 @@ class HindsightConfig:
|
||||
embeddings_litellm_sdk_encoding_format=os.getenv(
|
||||
ENV_EMBEDDINGS_LITELLM_SDK_ENCODING_FORMAT, DEFAULT_EMBEDDINGS_LITELLM_SDK_ENCODING_FORMAT
|
||||
),
|
||||
embeddings_litellm_sdk_max_input_tokens=int(v)
|
||||
if (v := os.getenv(ENV_EMBEDDINGS_LITELLM_SDK_MAX_INPUT_TOKENS))
|
||||
else DEFAULT_EMBEDDINGS_LITELLM_SDK_MAX_INPUT_TOKENS,
|
||||
# Gemini/Vertex AI embeddings (with fallback to LLM keys)
|
||||
embeddings_gemini_api_key=os.getenv(ENV_EMBEDDINGS_GEMINI_API_KEY) or os.getenv(ENV_LLM_API_KEY),
|
||||
embeddings_gemini_model=os.getenv(ENV_EMBEDDINGS_GEMINI_MODEL, DEFAULT_EMBEDDINGS_GEMINI_MODEL),
|
||||
@@ -2576,6 +2682,11 @@ class HindsightConfig:
|
||||
or os.getenv(ENV_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY),
|
||||
# Reranker
|
||||
reranker_provider=os.getenv(ENV_RERANKER_PROVIDER, DEFAULT_RERANKER_PROVIDER),
|
||||
reranker_send_bank_as_header=os.getenv(
|
||||
ENV_RERANKER_SEND_BANK_AS_HEADER,
|
||||
str(DEFAULT_RERANKER_SEND_BANK_AS_HEADER),
|
||||
).lower()
|
||||
in ("true", "1"),
|
||||
reranker_local_model=os.getenv(ENV_RERANKER_LOCAL_MODEL, DEFAULT_RERANKER_LOCAL_MODEL),
|
||||
reranker_local_force_cpu=os.getenv(
|
||||
ENV_RERANKER_LOCAL_FORCE_CPU, str(DEFAULT_RERANKER_LOCAL_FORCE_CPU)
|
||||
@@ -2608,6 +2719,11 @@ class HindsightConfig:
|
||||
reranker_max_candidates=int(os.getenv(ENV_RERANKER_MAX_CANDIDATES, str(DEFAULT_RERANKER_MAX_CANDIDATES))),
|
||||
semantic_min_similarity=float(os.getenv(ENV_SEMANTIC_MIN_SIMILARITY, str(DEFAULT_SEMANTIC_MIN_SIMILARITY))),
|
||||
bm25_min_score=float(os.getenv(ENV_BM25_MIN_SCORE, str(DEFAULT_BM25_MIN_SCORE))),
|
||||
bm25_max_query_terms=_parse_non_negative_int(
|
||||
ENV_BM25_MAX_QUERY_TERMS,
|
||||
os.getenv(ENV_BM25_MAX_QUERY_TERMS),
|
||||
DEFAULT_BM25_MAX_QUERY_TERMS,
|
||||
),
|
||||
recall_max_candidates_per_source=int(
|
||||
os.getenv(ENV_RECALL_MAX_CANDIDATES_PER_SOURCE, str(DEFAULT_RECALL_MAX_CANDIDATES_PER_SOURCE))
|
||||
),
|
||||
@@ -2902,6 +3018,10 @@ class HindsightConfig:
|
||||
db_command_timeout=int(os.getenv(ENV_DB_COMMAND_TIMEOUT, str(DEFAULT_DB_COMMAND_TIMEOUT))),
|
||||
db_acquire_timeout=int(os.getenv(ENV_DB_ACQUIRE_TIMEOUT, str(DEFAULT_DB_ACQUIRE_TIMEOUT))),
|
||||
db_statement_timeout=int(os.getenv(ENV_DB_STATEMENT_TIMEOUT, str(DEFAULT_DB_STATEMENT_TIMEOUT))),
|
||||
db_max_parallel_workers_per_gather=_parse_optional_non_negative_int(
|
||||
ENV_DB_MAX_PARALLEL_WORKERS_PER_GATHER,
|
||||
os.getenv(ENV_DB_MAX_PARALLEL_WORKERS_PER_GATHER),
|
||||
),
|
||||
model_init_timeout=float(os.getenv(ENV_MODEL_INIT_TIMEOUT, str(DEFAULT_MODEL_INIT_TIMEOUT))),
|
||||
# Worker configuration
|
||||
worker_enabled=os.getenv(ENV_WORKER_ENABLED, str(DEFAULT_WORKER_ENABLED)).lower() == "true",
|
||||
@@ -2924,9 +3044,23 @@ class HindsightConfig:
|
||||
worker_consolidation_bank_priority=_parse_bank_priority(
|
||||
os.getenv(ENV_WORKER_CONSOLIDATION_BANK_PRIORITY, "")
|
||||
),
|
||||
operation_retention_days=_parse_non_negative_int(
|
||||
ENV_OPERATION_RETENTION_DAYS,
|
||||
os.getenv(ENV_OPERATION_RETENTION_DAYS),
|
||||
DEFAULT_OPERATION_RETENTION_DAYS,
|
||||
),
|
||||
operation_cleanup_batch_size=_parse_positive_int(
|
||||
ENV_OPERATION_CLEANUP_BATCH_SIZE,
|
||||
os.getenv(ENV_OPERATION_CLEANUP_BATCH_SIZE),
|
||||
DEFAULT_OPERATION_CLEANUP_BATCH_SIZE,
|
||||
),
|
||||
retain_max_concurrent=int(os.getenv(ENV_RETAIN_MAX_CONCURRENT, str(DEFAULT_RETAIN_MAX_CONCURRENT))),
|
||||
# Reflect agent settings
|
||||
reflect_max_iterations=int(os.getenv(ENV_REFLECT_MAX_ITERATIONS, str(DEFAULT_REFLECT_MAX_ITERATIONS))),
|
||||
reflect_prompt_cache_enabled=os.getenv(
|
||||
ENV_REFLECT_PROMPT_CACHE_ENABLED, str(DEFAULT_REFLECT_PROMPT_CACHE_ENABLED)
|
||||
).lower()
|
||||
in ("1", "true", "yes", "on"),
|
||||
reflect_max_context_tokens=int(
|
||||
os.getenv(ENV_REFLECT_MAX_CONTEXT_TOKENS, str(DEFAULT_REFLECT_MAX_CONTEXT_TOKENS))
|
||||
),
|
||||
@@ -2989,6 +3123,11 @@ class HindsightConfig:
|
||||
audit_log_retention_days=int(
|
||||
os.getenv(ENV_AUDIT_LOG_RETENTION_DAYS, str(DEFAULT_AUDIT_LOG_RETENTION_DAYS))
|
||||
),
|
||||
# Retain reliability configuration (static, server-level only)
|
||||
fail_on_extraction_errors=os.getenv(
|
||||
ENV_FAIL_ON_EXTRACTION_ERRORS, str(DEFAULT_FAIL_ON_EXTRACTION_ERRORS)
|
||||
).lower()
|
||||
== "true",
|
||||
# LLM request tracing configuration (static, server-level only)
|
||||
llm_trace_enabled=os.getenv(ENV_LLM_TRACE_ENABLED, str(DEFAULT_LLM_TRACE_ENABLED)).lower() == "true",
|
||||
llm_trace_scopes=[
|
||||
|
||||
@@ -13,6 +13,8 @@ but operators should opt in with that in mind.
|
||||
|
||||
from typing import Any
|
||||
|
||||
RERANKER_BANK_ID_HEADER = "X-Hindsight-Bank-Id"
|
||||
|
||||
|
||||
def apply_bank_attribution(request: dict[str, Any]) -> None:
|
||||
"""Tag ``request`` with ``user=<bank_id>`` for per-bank cost attribution.
|
||||
@@ -32,3 +34,14 @@ def apply_bank_attribution(request: dict[str, Any]) -> None:
|
||||
bank_id = get_current_bank_id()
|
||||
if bank_id:
|
||||
request["user"] = bank_id
|
||||
|
||||
|
||||
def reranker_bank_attribution_headers() -> dict[str, str]:
|
||||
"""Return the fixed per-bank header for trusted remote reranker endpoints."""
|
||||
from ..config import get_config
|
||||
from .memory_engine import get_current_bank_id
|
||||
|
||||
if not get_config().reranker_send_bank_as_header:
|
||||
return {}
|
||||
bank_id = get_current_bank_id()
|
||||
return {RERANKER_BANK_ID_HEADER: bank_id} if bank_id else {}
|
||||
|
||||
@@ -0,0 +1,13 @@
|
||||
"""Shared causal-link taxonomy.
|
||||
|
||||
Retain writes only the canonical relationship. Transfer import/export also
|
||||
preserves historical relationship types so existing banks keep their graph
|
||||
semantics without allowing new retain output to create those types.
|
||||
"""
|
||||
|
||||
CANONICAL_CAUSAL_LINK_TYPE = "caused_by"
|
||||
LEGACY_CAUSAL_LINK_TYPE_NAMES = ("causes", "enables", "prevents")
|
||||
|
||||
CANONICAL_CAUSAL_LINK_TYPES = frozenset({CANONICAL_CAUSAL_LINK_TYPE})
|
||||
LEGACY_CAUSAL_LINK_TYPES = frozenset(LEGACY_CAUSAL_LINK_TYPE_NAMES)
|
||||
CAUSAL_LINK_TYPES = (CANONICAL_CAUSAL_LINK_TYPE, *LEGACY_CAUSAL_LINK_TYPE_NAMES)
|
||||
@@ -109,6 +109,11 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
|
||||
end.replace(hour=23, minute=59, second=59, microsecond=999999),
|
||||
)
|
||||
|
||||
def safe_constraint(start: datetime | None, end: datetime | None) -> DateRange | NoTemporalConstraintSentinel:
|
||||
if start is None or end is None:
|
||||
return NO_TEMPORAL_CONSTRAINT
|
||||
return constraint(start, end)
|
||||
|
||||
def subtract_months(months: int) -> datetime:
|
||||
month_index = reference_date.month - months - 1
|
||||
year = reference_date.year + month_index // 12
|
||||
@@ -126,11 +131,21 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
|
||||
day = min(base_date.day, calendar.monthrange(year, month)[1])
|
||||
return base_date.replace(year=year, month=month, day=day)
|
||||
|
||||
def add_years(base_date: datetime, years: int) -> datetime:
|
||||
def add_years(base_date: datetime, years: int) -> datetime | None:
|
||||
year = base_date.year + years
|
||||
if year < datetime.min.year or year > datetime.max.year:
|
||||
return None
|
||||
day = min(base_date.day, calendar.monthrange(year, base_date.month)[1])
|
||||
return base_date.replace(year=year, day=day)
|
||||
|
||||
def add_days(base_date: datetime | None, days: int) -> datetime | None:
|
||||
if base_date is None:
|
||||
return None
|
||||
try:
|
||||
return base_date + timedelta(days=days)
|
||||
except OverflowError:
|
||||
return None
|
||||
|
||||
def has_chinese_temporal_context(match: re.Match[str]) -> bool:
|
||||
if match.end() >= len(query):
|
||||
return True
|
||||
@@ -438,6 +453,11 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
|
||||
return NO_TEMPORAL_CONSTRAINT
|
||||
return constraint(start, reference_date)
|
||||
|
||||
def safe_since_constraint(start: datetime | None) -> DateRange | NoTemporalConstraintSentinel:
|
||||
if start is None:
|
||||
return NO_TEMPORAL_CONSTRAINT
|
||||
return since_constraint(start)
|
||||
|
||||
def since_from_period(
|
||||
period: DateRange | None,
|
||||
) -> DateRange | NoTemporalConstraintSentinel | None:
|
||||
@@ -450,7 +470,7 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
|
||||
return None
|
||||
return since_constraint(day)
|
||||
|
||||
def relative_offset_datetime(amount: int, unit: str, direction: int) -> datetime:
|
||||
def relative_offset_datetime(amount: int, unit: str, direction: int) -> datetime | None:
|
||||
if unit in ("天", "日"):
|
||||
return reference_date + timedelta(days=direction * amount)
|
||||
if unit in ("周", "星期", "礼拜"):
|
||||
@@ -459,15 +479,15 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
|
||||
return add_months(reference_date, direction * amount)
|
||||
return add_years(reference_date, direction * amount)
|
||||
|
||||
def point_constraint_at_offset(amount: int, unit: str, direction: int) -> DateRange:
|
||||
def point_constraint_at_offset(amount: int, unit: str, direction: int) -> DateRange | NoTemporalConstraintSentinel:
|
||||
d = relative_offset_datetime(amount, unit, direction)
|
||||
return constraint(d, d)
|
||||
return safe_constraint(d, d)
|
||||
|
||||
def window_to_reference(amount: int, unit: str) -> DateRange:
|
||||
return constraint(relative_offset_datetime(amount, unit, -1), reference_date)
|
||||
def window_to_reference(amount: int, unit: str) -> DateRange | NoTemporalConstraintSentinel:
|
||||
return safe_constraint(relative_offset_datetime(amount, unit, -1), reference_date)
|
||||
|
||||
def window_from_reference(amount: int, unit: str) -> DateRange:
|
||||
return constraint(reference_date, relative_offset_datetime(amount, unit, 1))
|
||||
def window_from_reference(amount: int, unit: str) -> DateRange | NoTemporalConstraintSentinel:
|
||||
return safe_constraint(reference_date, relative_offset_datetime(amount, unit, 1))
|
||||
|
||||
# Chinese rule guide
|
||||
#
|
||||
@@ -781,8 +801,8 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
|
||||
if relative_year_fixed_day_since_match:
|
||||
year = relative_year_number(relative_year_fixed_day_since_match.group(1))
|
||||
base = add_years(reference_date, year - reference_date.year)
|
||||
d = base + timedelta(days=fixed_day_offset(relative_year_fixed_day_since_match.group(2)))
|
||||
return since_constraint(d)
|
||||
d = add_days(base, fixed_day_offset(relative_year_fixed_day_since_match.group(2)))
|
||||
return safe_since_constraint(d)
|
||||
|
||||
fixed_day_since_match = chinese_search(
|
||||
rf"(大大后天|大后天|后天|明天|明日|今天|今日|本日|当日|当天|昨天|昨日|大大前天|大前天|前天)"
|
||||
@@ -799,7 +819,7 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
|
||||
amount = parse_chinese_number(exact_relative_since_match.group(1))
|
||||
unit = exact_relative_since_match.group(2)
|
||||
if amount is not None:
|
||||
return since_constraint(relative_offset_datetime(amount, unit, -1))
|
||||
return safe_since_constraint(relative_offset_datetime(amount, unit, -1))
|
||||
|
||||
weekend_since_match = chinese_search(
|
||||
rf"(?<![上下大小每个各隔])"
|
||||
@@ -899,8 +919,8 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
|
||||
if relative_year_daypart_since_match:
|
||||
year = relative_year_number(relative_year_daypart_since_match.group(1))
|
||||
base = add_years(reference_date, year - reference_date.year)
|
||||
d = base + timedelta(days=daypart_day_offset(relative_year_daypart_since_match.group(2)))
|
||||
return since_constraint(d)
|
||||
d = add_days(base, daypart_day_offset(relative_year_daypart_since_match.group(2)))
|
||||
return safe_since_constraint(d)
|
||||
|
||||
daypart_since_match = chinese_search(
|
||||
rf"(昨晚|昨夜|前晚|前夜|今晚|今早|今晨|明早|明晚|明夜){chinese_since_suffix_pattern}"
|
||||
@@ -915,17 +935,17 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
|
||||
if relative_year_daypart_match:
|
||||
year = relative_year_number(relative_year_daypart_match.group(1))
|
||||
base = add_years(reference_date, year - reference_date.year)
|
||||
d = base + timedelta(days=daypart_day_offset(relative_year_daypart_match.group(2)))
|
||||
return constraint(d, d)
|
||||
d = add_days(base, daypart_day_offset(relative_year_daypart_match.group(2)))
|
||||
return safe_constraint(d, d)
|
||||
|
||||
# Day-part abbreviations still resolve only to date granularity.
|
||||
if chinese_search(r"昨晚|昨夜"):
|
||||
d = reference_date + timedelta(days=daypart_day_offset("昨晚"))
|
||||
return constraint(d, d)
|
||||
d = add_days(reference_date, daypart_day_offset("昨晚"))
|
||||
return safe_constraint(d, d)
|
||||
|
||||
if chinese_search(r"前晚|前夜"):
|
||||
d = reference_date + timedelta(days=daypart_day_offset("前晚"))
|
||||
return constraint(d, d)
|
||||
d = add_days(reference_date, daypart_day_offset("前晚"))
|
||||
return safe_constraint(d, d)
|
||||
|
||||
if chinese_search(r"今晚|今早|今晨"):
|
||||
return constraint(reference_date, reference_date)
|
||||
@@ -941,8 +961,8 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
|
||||
if relative_year_fixed_day_match:
|
||||
year = relative_year_number(relative_year_fixed_day_match.group(1))
|
||||
base = add_years(reference_date, year - reference_date.year)
|
||||
d = base + timedelta(days=fixed_day_offset(relative_year_fixed_day_match.group(2)))
|
||||
return constraint(d, d)
|
||||
d = add_days(base, fixed_day_offset(relative_year_fixed_day_match.group(2)))
|
||||
return safe_constraint(d, d)
|
||||
|
||||
if chinese_search(r"昨天|昨日"):
|
||||
d = reference_date - timedelta(days=1)
|
||||
@@ -1085,7 +1105,7 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
|
||||
end_amount = parse_chinese_number(amount_text[-1])
|
||||
unit = adjacent_fuzzy_future_match.group(2)
|
||||
if start_amount is not None and end_amount is not None:
|
||||
return constraint(
|
||||
return safe_constraint(
|
||||
relative_offset_datetime(start_amount, unit, 1),
|
||||
relative_offset_datetime(end_amount, unit, 1),
|
||||
)
|
||||
@@ -1093,7 +1113,7 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
|
||||
few_future_match = chinese_search(rf"[几数]个?(天|日|周|星期|礼拜|月|年){chinese_relative_future_suffix_pattern}")
|
||||
if few_future_match:
|
||||
unit = few_future_match.group(1)
|
||||
return constraint(relative_offset_datetime(2, unit, 1), relative_offset_datetime(5, unit, 1))
|
||||
return safe_constraint(relative_offset_datetime(2, unit, 1), relative_offset_datetime(5, unit, 1))
|
||||
|
||||
exact_future_match = chinese_search(
|
||||
rf"(?<![{_CHINESE_NUMERAL_PREFIX_CHARS}])([0-9]+|[{_CHINESE_NUMERAL_CHARS}]+)个?(天|日|周|星期|礼拜|月|年){chinese_relative_future_suffix_pattern}"
|
||||
@@ -1113,7 +1133,7 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
|
||||
second_amount = parse_chinese_number(adjacent_fuzzy_past_match.group(2))
|
||||
unit = adjacent_fuzzy_past_match.group(3)
|
||||
if first_amount is not None and second_amount is not None and second_amount == first_amount + 1:
|
||||
return constraint(
|
||||
return safe_constraint(
|
||||
relative_offset_datetime(second_amount, unit, -1),
|
||||
relative_offset_datetime(first_amount, unit, -1),
|
||||
)
|
||||
@@ -1144,7 +1164,7 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
|
||||
return constraint(reference_date - timedelta(days=150), reference_date - timedelta(days=60))
|
||||
|
||||
if chinese_search(r"一两年前|[两二]三年前|三两年前"):
|
||||
return constraint(add_years(reference_date, -3), add_years(reference_date, -1))
|
||||
return safe_constraint(add_years(reference_date, -3), add_years(reference_date, -1))
|
||||
|
||||
rolling_this_adjacent_match = chinese_search(
|
||||
r"这(一两|[两二]三|三两|三四|四五|五六|六七|七八|八九|九十)个?(天|日|周|星期|礼拜|月|年)"
|
||||
@@ -1154,7 +1174,7 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
|
||||
end_amount = 3 if amount_text in ("一两", "三两") else parse_chinese_number(amount_text[-1])
|
||||
unit = rolling_this_adjacent_match.group(2)
|
||||
if end_amount is not None:
|
||||
return constraint(relative_offset_datetime(end_amount, unit, -1), reference_date)
|
||||
return safe_constraint(relative_offset_datetime(end_amount, unit, -1), reference_date)
|
||||
|
||||
rolling_this_count_match = chinese_search(rf"这([0-9]+|[{_CHINESE_NUMERAL_CHARS}]+)个?(天|日|周|星期|礼拜|月|年)")
|
||||
if rolling_this_count_match:
|
||||
@@ -1191,7 +1211,7 @@ def extract_chinese_period(query: str, reference_date: datetime) -> DateRange |
|
||||
end_amount = 3 if amount_text in ("一两", "三两") else parse_chinese_number(amount_text[-1])
|
||||
unit = rolling_past_adjacent_match.group(3)
|
||||
if end_amount is not None:
|
||||
return constraint(relative_offset_datetime(end_amount, unit, -1), reference_date)
|
||||
return safe_constraint(relative_offset_datetime(end_amount, unit, -1), reference_date)
|
||||
|
||||
rolling_past_few_match = chinese_search(r"(过去|近|最近)几个?(天|日|周|星期|礼拜|月|年)")
|
||||
if rolling_past_few_match:
|
||||
|
||||
@@ -28,6 +28,7 @@ from fnmatch import fnmatchcase
|
||||
from itertools import combinations
|
||||
from typing import TYPE_CHECKING, Any, Literal
|
||||
|
||||
import asyncpg
|
||||
from pydantic import BaseModel, field_validator
|
||||
|
||||
from ...config import get_config
|
||||
@@ -102,6 +103,29 @@ class _DedupDecision(BaseModel):
|
||||
text: str = "" # the synthesized merged observation (when action == "merge")
|
||||
reason: str = ""
|
||||
|
||||
@field_validator("action", mode="before")
|
||||
@classmethod
|
||||
def _normalize_action(cls, value: object) -> str:
|
||||
if isinstance(value, str):
|
||||
normalized = value.strip().lower()
|
||||
if normalized in {"merge", "keep"}:
|
||||
return normalized
|
||||
|
||||
logger.warning("Invalid consolidation dedup action %r; defaulting to keep", value)
|
||||
return "keep"
|
||||
|
||||
|
||||
def _dedup_decision_from_response(raw: Any) -> _DedupDecision:
|
||||
try:
|
||||
if isinstance(raw, _DedupDecision):
|
||||
return raw
|
||||
if isinstance(raw, str):
|
||||
return _DedupDecision.model_validate_json(raw)
|
||||
return _DedupDecision.model_validate(raw)
|
||||
except ValueError as exc:
|
||||
logger.warning("Invalid consolidation dedup response %r; defaulting to keep: %s", raw, exc)
|
||||
return _DedupDecision(action="keep", reason="invalid structured response")
|
||||
|
||||
|
||||
_DEDUP_PROMPT = """You reconcile long-term memory observations. A NEW observation is about to be \
|
||||
stored, and it is highly similar to an EXISTING one:
|
||||
@@ -109,9 +133,20 @@ stored, and it is highly similar to an EXISTING one:
|
||||
[NEW] {new}
|
||||
[EXISTING] {existing}
|
||||
|
||||
If they assert the SAME fact (wording aside), respond action="merge" and provide `text`: a single \
|
||||
observation that preserves EVERY detail from both. If they differ in ANY important detail — a \
|
||||
number/quantity, a named entity or language, a negation, or a condition — respond action="keep"."""
|
||||
Respond with ONLY one valid JSON object matching one of these shapes:
|
||||
|
||||
For duplicate facts:
|
||||
{{"action": "merge", "text": "...", "reason": "..."}}
|
||||
|
||||
For distinct facts:
|
||||
{{"action": "keep", "text": "", "reason": "..."}}
|
||||
|
||||
Do NOT use key=value lines, markdown fences, or any text outside the JSON object.
|
||||
|
||||
If they assert the SAME fact (wording aside), set "action" to "merge" and provide "text": a \
|
||||
single observation that preserves EVERY detail from both. If they differ in ANY important detail \
|
||||
— a number/quantity, a named entity or language, a negation, or a condition — set "action" to \
|
||||
"keep" and "text" to an empty string."""
|
||||
|
||||
|
||||
def _dedup_active(config: Any) -> bool:
|
||||
@@ -189,10 +224,12 @@ async def _dedup_adjudicate(
|
||||
if best_id is None:
|
||||
return _DedupOutcome(best_id=None, merged_text="", should_merge=False)
|
||||
|
||||
decision: _DedupDecision = await dedup_llm_config.call(
|
||||
messages=[{"role": "user", "content": _DEDUP_PROMPT.format(new=anchor_text, existing=best_text)}],
|
||||
response_format=_DedupDecision,
|
||||
scope="consolidation_dedup",
|
||||
decision = _dedup_decision_from_response(
|
||||
await dedup_llm_config.call(
|
||||
messages=[{"role": "user", "content": _DEDUP_PROMPT.format(new=anchor_text, existing=best_text)}],
|
||||
response_format=_DedupDecision,
|
||||
scope="consolidation_dedup",
|
||||
)
|
||||
)
|
||||
if decision.action != "merge":
|
||||
return _DedupOutcome(best_id=best_id, merged_text="", should_merge=False)
|
||||
@@ -1766,15 +1803,22 @@ async def _append_observation_history(
|
||||
history from growing without bound.
|
||||
"""
|
||||
obs_uuid = uuid.UUID(observation_id)
|
||||
await conn.execute(
|
||||
f"""
|
||||
try:
|
||||
await conn.execute(
|
||||
f"""
|
||||
INSERT INTO {fq_table("observation_history")} (observation_id, bank_id, content, changed_at)
|
||||
VALUES ($1, $2, $3::jsonb, now())
|
||||
""",
|
||||
obs_uuid,
|
||||
bank_id,
|
||||
json.dumps(asdict(snapshot)),
|
||||
)
|
||||
obs_uuid,
|
||||
bank_id,
|
||||
json.dumps(asdict(snapshot)),
|
||||
)
|
||||
except asyncpg.exceptions.ForeignKeyViolationError:
|
||||
logger.warning(
|
||||
f"FK violation writing observation_history for {observation_id}: "
|
||||
"observation was removed before history could be written (race with parallel consolidation). Skipping."
|
||||
)
|
||||
return
|
||||
if max_entries and max_entries > 0:
|
||||
await conn.execute(
|
||||
f"""
|
||||
|
||||
@@ -7,11 +7,13 @@ Configuration via environment variables - see hindsight_api.config for all env v
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import gc
|
||||
import logging
|
||||
import os
|
||||
import warnings
|
||||
from abc import ABC, abstractmethod
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
@@ -45,6 +47,7 @@ from ..config import (
|
||||
ENV_RERANKER_TEI_URL,
|
||||
ENV_RERANKER_ZEROENTROPY_API_KEY,
|
||||
)
|
||||
from .bank_attribution import reranker_bank_attribution_headers
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -86,6 +89,12 @@ def _resolve_malloc_trim():
|
||||
_malloc_trim = _resolve_malloc_trim()
|
||||
|
||||
|
||||
def _release_rerank_heap() -> None:
|
||||
"""Release transient Python and native heap memory after local reranking."""
|
||||
gc.collect()
|
||||
_malloc_trim()
|
||||
|
||||
|
||||
class CrossEncoderModel(ABC):
|
||||
"""
|
||||
Abstract base class for cross-encoder reranking.
|
||||
@@ -315,7 +324,7 @@ class LocalSTCrossEncoder(CrossEncoderModel):
|
||||
scores = self._model.predict(pairs, batch_size=self.batch_size, show_progress_bar=False)
|
||||
return scores.tolist() if hasattr(scores, "tolist") else list(scores)
|
||||
finally:
|
||||
_malloc_trim()
|
||||
_release_rerank_heap()
|
||||
|
||||
async def predict(self, pairs: list[tuple[str, str]]) -> list[float]:
|
||||
"""
|
||||
@@ -484,6 +493,7 @@ class RemoteTEICrossEncoder(CrossEncoderModel):
|
||||
semaphore,
|
||||
"POST",
|
||||
f"{self.base_url}/rerank",
|
||||
headers=reranker_bank_attribution_headers(),
|
||||
json={
|
||||
"query": query,
|
||||
"texts": texts,
|
||||
@@ -624,7 +634,11 @@ class _CohereCompatibleRerankClient:
|
||||
if self.include_top_n:
|
||||
body["top_n"] = len(texts)
|
||||
|
||||
response = await self._async_client.post(self.rerank_url, json=body)
|
||||
response = await self._async_client.post(
|
||||
self.rerank_url,
|
||||
headers=reranker_bank_attribution_headers(),
|
||||
json=body,
|
||||
)
|
||||
response.raise_for_status()
|
||||
result = response.json()
|
||||
|
||||
@@ -990,11 +1004,11 @@ class FlashRankCrossEncoder(CrossEncoderModel):
|
||||
|
||||
def _predict_sync(self, pairs: list[tuple[str, str]]) -> list[float]:
|
||||
"""Synchronous predict - processes each query group."""
|
||||
from flashrank import RerankRequest
|
||||
|
||||
if not pairs:
|
||||
return []
|
||||
|
||||
from flashrank import RerankRequest
|
||||
|
||||
try:
|
||||
# Group pairs by query
|
||||
query_groups: dict[str, list[tuple[int, str]]] = {}
|
||||
@@ -1023,7 +1037,7 @@ class FlashRankCrossEncoder(CrossEncoderModel):
|
||||
|
||||
return all_scores
|
||||
finally:
|
||||
_malloc_trim()
|
||||
_release_rerank_heap()
|
||||
|
||||
async def predict(self, pairs: list[tuple[str, str]]) -> list[float]:
|
||||
"""
|
||||
@@ -1151,6 +1165,7 @@ class LiteLLMCrossEncoder(CrossEncoderModel):
|
||||
# LiteLLM /rerank follows Cohere API format
|
||||
response = await self._async_client.post(
|
||||
f"{self.api_base}/rerank",
|
||||
headers=reranker_bank_attribution_headers(),
|
||||
json={
|
||||
"model": self.model,
|
||||
"query": query,
|
||||
@@ -1269,10 +1284,11 @@ class LiteLLMSDKCrossEncoder(CrossEncoderModel):
|
||||
indices = [idx for idx, _ in indexed_texts]
|
||||
|
||||
# Build kwargs for rerank call
|
||||
rerank_kwargs = {
|
||||
rerank_kwargs: dict[str, Any] = {
|
||||
"model": self.model,
|
||||
"query": query,
|
||||
"documents": texts,
|
||||
"headers": reranker_bank_attribution_headers(),
|
||||
}
|
||||
if self.api_key:
|
||||
rerank_kwargs["api_key"] = self.api_key
|
||||
@@ -1281,21 +1297,9 @@ class LiteLLMSDKCrossEncoder(CrossEncoderModel):
|
||||
|
||||
response = await self._litellm.arerank(**rerank_kwargs)
|
||||
|
||||
# Map scores back to original positions
|
||||
# Response format: RerankResponse with results list
|
||||
# Each result is a TypedDict with "index" and "relevance_score"
|
||||
if hasattr(response, "results") and response.results:
|
||||
for result in response.results:
|
||||
# Results are TypedDicts, use dict-style access
|
||||
original_idx = result["index"]
|
||||
score = result.get("relevance_score", result.get("score", 0.0))
|
||||
all_scores[indices[original_idx]] = score
|
||||
elif isinstance(response, list):
|
||||
# Direct list of scores (unlikely but defensive)
|
||||
for i, score in enumerate(response):
|
||||
all_scores[indices[i]] = score
|
||||
else:
|
||||
logger.warning(f"Unexpected response format from LiteLLM rerank: {type(response)}")
|
||||
for result in response.results:
|
||||
original_idx = result["index"]
|
||||
all_scores[indices[original_idx]] = result["relevance_score"]
|
||||
|
||||
return all_scores
|
||||
|
||||
|
||||
@@ -307,6 +307,17 @@ class DatabaseBackend(ABC):
|
||||
"""Close the connection pool and release all resources."""
|
||||
...
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def is_ready(self) -> bool:
|
||||
"""Whether the pool exists and can serve connections.
|
||||
|
||||
False before :meth:`initialize` and after :meth:`shutdown`. Best-effort
|
||||
callers (tracing, auditing) check this to skip work during those windows
|
||||
instead of acquiring and interpreting the resulting error.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
@asynccontextmanager
|
||||
async def acquire(self) -> AsyncIterator[DatabaseConnection]:
|
||||
|
||||
@@ -18,6 +18,7 @@ and mirrors Django's ``DatabaseOperations`` architecture.
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from typing import Any
|
||||
|
||||
from .base import DatabaseConnection
|
||||
@@ -484,6 +485,23 @@ class DataAccessOps(ABC):
|
||||
|
||||
# -- Task claiming operations ------------------------------------------
|
||||
|
||||
@abstractmethod
|
||||
async def prune_terminal_operations(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
cutoff: datetime,
|
||||
*,
|
||||
batch_size: int,
|
||||
) -> int:
|
||||
"""Delete one deterministic batch of terminal operations older than ``cutoff``.
|
||||
|
||||
Implementations must lock candidates without waiting on rows another
|
||||
worker is pruning, never select pending/processing rows, and return the
|
||||
number deleted. The caller provides a transaction around this method.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def claim_tasks(
|
||||
self,
|
||||
|
||||
@@ -13,6 +13,8 @@ from .base import DatabaseConnection
|
||||
from .ops import DataAccessOps, TagListingParts
|
||||
from .result import DictResultRow as ResultRow
|
||||
|
||||
ORACLE_IN_LIST_LIMIT = 1000
|
||||
|
||||
|
||||
class OracleOps(DataAccessOps):
|
||||
"""Oracle-specific data access operations."""
|
||||
@@ -329,6 +331,12 @@ class OracleOps(DataAccessOps):
|
||||
entities_table: str,
|
||||
bank_id: str,
|
||||
) -> int:
|
||||
# NB: the Postgres path additionally selects victims FOR UPDATE in sorted
|
||||
# (entity_id_1, entity_id_2) order to prevent the #2529 deadlock against
|
||||
# retain's sorted cooccurrence upsert. Oracle's DELETE can't carry that
|
||||
# ordered-lock CTE the same way, so here we rely on the Pass 2/3 retry
|
||||
# wrap in run_graph_maintenance_job (retry_with_backoff is ORA-00060
|
||||
# deadlock-aware) to recover instead. Deliberate dialect asymmetry.
|
||||
deleted = await conn.execute(
|
||||
f"""
|
||||
DELETE FROM {ec_table}
|
||||
@@ -447,6 +455,14 @@ class OracleOps(DataAccessOps):
|
||||
FROM {ue_table} ue_target
|
||||
WHERE ue_target.entity_id = se.entity_id
|
||||
AND ue_target.unit_id != ALL($1::uuid[])
|
||||
-- Filter before applying the cap: candidates from other fact
|
||||
-- types must not consume this entity's bounded fan-out.
|
||||
AND EXISTS (
|
||||
SELECT 1
|
||||
FROM {mu_table} mu_target
|
||||
WHERE mu_target.id = ue_target.unit_id
|
||||
AND mu_target.fact_type = $2
|
||||
)
|
||||
ORDER BY ue_target.unit_id DESC
|
||||
FETCH FIRST {per_entity_limit} ROWS ONLY
|
||||
) t
|
||||
@@ -459,7 +475,6 @@ class OracleOps(DataAccessOps):
|
||||
es.score, 'entity' AS source
|
||||
FROM entity_scores es
|
||||
JOIN {mu_table} mu ON mu.id = es.unit_id
|
||||
WHERE mu.fact_type = $2
|
||||
ORDER BY es.score DESC
|
||||
FETCH FIRST $3 ROWS ONLY
|
||||
)"""
|
||||
@@ -824,6 +839,157 @@ class OracleOps(DataAccessOps):
|
||||
|
||||
# -- Task claiming operations ------------------------------------------
|
||||
|
||||
async def prune_terminal_operations(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
cutoff: datetime,
|
||||
*,
|
||||
batch_size: int,
|
||||
) -> int:
|
||||
# Oracle rejects a row-limited SELECT ... FOR UPDATE (ORA-02014). Pick
|
||||
# the deterministic bounded IDs first, then lock only that candidate
|
||||
# set and re-check eligibility before deleting in the same transaction.
|
||||
# Clamp to Oracle's 1000-expression IN-list limit because the adapter
|
||||
# expands the candidate UUID list into individual bind variables.
|
||||
# Cancelled children cannot complete parent aggregation, so retain the
|
||||
# parent guard only for completed/failed children. Before removing a
|
||||
# cancelled child, preserve its signal by cancelling a pending parent
|
||||
# in this transaction and refreshing the parent's retention window.
|
||||
# Validate metadata before HEXTORAW: CASE makes malformed UUIDs yield
|
||||
# NULL while keeping the indexed RAW parent.operation_id key unwrapped.
|
||||
effective_batch_size = min(batch_size, ORACLE_IN_LIST_LIMIT)
|
||||
candidates = await conn.fetch(
|
||||
f"""
|
||||
SELECT candidate_operation.operation_id
|
||||
FROM {table} candidate_operation
|
||||
WHERE candidate_operation.status IN ('completed', 'failed', 'cancelled')
|
||||
AND candidate_operation.updated_at < $1
|
||||
AND (
|
||||
candidate_operation.status = 'cancelled'
|
||||
OR NOT EXISTS (
|
||||
SELECT 1
|
||||
FROM {table} parent
|
||||
WHERE parent.operation_id = CASE
|
||||
WHEN REGEXP_LIKE(
|
||||
JSON_VALUE(
|
||||
candidate_operation.result_metadata,
|
||||
'$.parent_operation_id' RETURNING VARCHAR2(36) NULL ON ERROR
|
||||
),
|
||||
'^[0-9A-Fa-f]{{8}}-[0-9A-Fa-f]{{4}}-[0-9A-Fa-f]{{4}}-[0-9A-Fa-f]{{4}}-[0-9A-Fa-f]{{12}}$'
|
||||
)
|
||||
THEN HEXTORAW(REPLACE(
|
||||
JSON_VALUE(
|
||||
candidate_operation.result_metadata,
|
||||
'$.parent_operation_id' RETURNING VARCHAR2(36) NULL ON ERROR
|
||||
),
|
||||
'-',
|
||||
''
|
||||
))
|
||||
ELSE NULL
|
||||
END
|
||||
AND parent.bank_id = candidate_operation.bank_id
|
||||
)
|
||||
)
|
||||
ORDER BY candidate_operation.updated_at, candidate_operation.operation_id
|
||||
LIMIT $2
|
||||
""",
|
||||
cutoff,
|
||||
effective_batch_size,
|
||||
)
|
||||
if not candidates:
|
||||
return 0
|
||||
|
||||
candidate_ids = [row["operation_id"] for row in candidates]
|
||||
locked = await conn.fetch(
|
||||
f"""
|
||||
SELECT candidate_operation.operation_id
|
||||
FROM {table} candidate_operation
|
||||
WHERE candidate_operation.operation_id = ANY($1)
|
||||
AND candidate_operation.status IN ('completed', 'failed', 'cancelled')
|
||||
AND candidate_operation.updated_at < $2
|
||||
AND (
|
||||
candidate_operation.status = 'cancelled'
|
||||
OR NOT EXISTS (
|
||||
SELECT 1
|
||||
FROM {table} parent
|
||||
WHERE parent.operation_id = CASE
|
||||
WHEN REGEXP_LIKE(
|
||||
JSON_VALUE(
|
||||
candidate_operation.result_metadata,
|
||||
'$.parent_operation_id' RETURNING VARCHAR2(36) NULL ON ERROR
|
||||
),
|
||||
'^[0-9A-Fa-f]{{8}}-[0-9A-Fa-f]{{4}}-[0-9A-Fa-f]{{4}}-[0-9A-Fa-f]{{4}}-[0-9A-Fa-f]{{12}}$'
|
||||
)
|
||||
THEN HEXTORAW(REPLACE(
|
||||
JSON_VALUE(
|
||||
candidate_operation.result_metadata,
|
||||
'$.parent_operation_id' RETURNING VARCHAR2(36) NULL ON ERROR
|
||||
),
|
||||
'-',
|
||||
''
|
||||
))
|
||||
ELSE NULL
|
||||
END
|
||||
AND parent.bank_id = candidate_operation.bank_id
|
||||
)
|
||||
)
|
||||
ORDER BY candidate_operation.updated_at, candidate_operation.operation_id
|
||||
FOR UPDATE OF candidate_operation.operation_id SKIP LOCKED
|
||||
""",
|
||||
candidate_ids,
|
||||
cutoff,
|
||||
)
|
||||
if not locked:
|
||||
return 0
|
||||
operation_ids = [row["operation_id"] for row in locked]
|
||||
await conn.execute(
|
||||
f"""
|
||||
UPDATE {table} parent
|
||||
SET status = 'cancelled',
|
||||
updated_at = now(),
|
||||
completed_at = COALESCE(parent.completed_at, now()),
|
||||
error_message = COALESCE(
|
||||
parent.error_message,
|
||||
'Cancelled because a child operation was cancelled'
|
||||
)
|
||||
WHERE parent.status = 'pending'
|
||||
AND EXISTS (
|
||||
SELECT 1
|
||||
FROM {table} candidate_operation
|
||||
WHERE candidate_operation.operation_id = ANY($1)
|
||||
AND candidate_operation.status = 'cancelled'
|
||||
AND candidate_operation.updated_at < $2
|
||||
AND candidate_operation.bank_id = parent.bank_id
|
||||
AND parent.operation_id = CASE
|
||||
WHEN REGEXP_LIKE(
|
||||
JSON_VALUE(
|
||||
candidate_operation.result_metadata,
|
||||
'$.parent_operation_id' RETURNING VARCHAR2(36) NULL ON ERROR
|
||||
),
|
||||
'^[0-9A-Fa-f]{{8}}-[0-9A-Fa-f]{{4}}-[0-9A-Fa-f]{{4}}-[0-9A-Fa-f]{{4}}-[0-9A-Fa-f]{{12}}$'
|
||||
)
|
||||
THEN HEXTORAW(REPLACE(
|
||||
JSON_VALUE(
|
||||
candidate_operation.result_metadata,
|
||||
'$.parent_operation_id' RETURNING VARCHAR2(36) NULL ON ERROR
|
||||
),
|
||||
'-',
|
||||
''
|
||||
))
|
||||
ELSE NULL
|
||||
END
|
||||
)
|
||||
""",
|
||||
operation_ids,
|
||||
cutoff,
|
||||
)
|
||||
await conn.execute(
|
||||
f"DELETE FROM {table} WHERE operation_id = ANY($1)",
|
||||
operation_ids,
|
||||
)
|
||||
return len(operation_ids)
|
||||
|
||||
async def _claim_consolidation_tasks(
|
||||
self,
|
||||
conn,
|
||||
|
||||
@@ -4,11 +4,40 @@ Uses unnest(), LATERAL, DISTINCT ON, and native array operations for
|
||||
efficient batch operations.
|
||||
"""
|
||||
|
||||
from datetime import datetime
|
||||
|
||||
from .base import DatabaseConnection
|
||||
from .ops import DataAccessOps, TagListingParts
|
||||
from .result import ResultRow
|
||||
|
||||
|
||||
def pg_search_vector_expr(
|
||||
config,
|
||||
*,
|
||||
text_col: str = "text",
|
||||
context_col: str = "context",
|
||||
signals_col: str = "text_signals",
|
||||
) -> str | None:
|
||||
"""SQL expression that builds ``search_vector`` for the configured PG text-search backend.
|
||||
|
||||
Single source of truth shared by the batch insert (over the ``input_data``
|
||||
CTE columns) and the curation revert recompute (over a ``memory_units`` row),
|
||||
so the two can never drift. Returns ``None`` for backends that leave
|
||||
``search_vector`` unpopulated — pgroonga / pg_textsearch / pg_search index the
|
||||
base text columns directly and keep only a dummy column, so there is nothing
|
||||
to build.
|
||||
|
||||
``text_search_extension_native_language`` is validated as a PG identifier in
|
||||
``HindsightConfig.validate()``, so embedding it as a SQL literal is safe.
|
||||
"""
|
||||
combined = f"COALESCE({text_col}, '') || ' ' || COALESCE({context_col}, '') || ' ' || COALESCE({signals_col}, '')"
|
||||
if config.text_search_extension == "vchord":
|
||||
return f"tokenize({combined}, 'llmlingua2')::bm25_catalog.bm25vector"
|
||||
if config.text_search_extension == "native":
|
||||
return f"to_tsvector('{config.text_search_extension_native_language}'::regconfig, {combined})"
|
||||
return None
|
||||
|
||||
|
||||
class PostgreSQLOps(DataAccessOps):
|
||||
"""PostgreSQL-specific data access operations using unnest and LATERAL."""
|
||||
|
||||
@@ -93,101 +122,39 @@ class PostgreSQLOps(DataAccessOps):
|
||||
config = get_config()
|
||||
table = self._get_mu_table()
|
||||
|
||||
if config.text_search_extension == "vchord":
|
||||
query = f"""
|
||||
WITH input_data AS (
|
||||
SELECT * FROM unnest(
|
||||
$2::text[], $3::vector[], $4::timestamptz[], $5::timestamptz[], $6::timestamptz[], $7::timestamptz[],
|
||||
$8::text[], $9::text[], $10::jsonb[], $11::text[], $12::text[], $13::jsonb[], $14::jsonb[], $15::text[]
|
||||
) AS t(text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
context, fact_type, metadata, chunk_id, document_id, tags_json,
|
||||
observation_scopes_json, text_signals)
|
||||
)
|
||||
INSERT INTO {table} (bank_id, text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
context, fact_type, metadata, chunk_id, document_id, tags,
|
||||
observation_scopes, text_signals, search_vector)
|
||||
SELECT
|
||||
$1,
|
||||
text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
context, fact_type, metadata, chunk_id, document_id,
|
||||
COALESCE(
|
||||
(SELECT array_agg(elem) FROM jsonb_array_elements_text(tags_json) AS elem),
|
||||
'{{}}'::varchar[]
|
||||
),
|
||||
observation_scopes_json,
|
||||
text_signals,
|
||||
tokenize(
|
||||
COALESCE(text, '') || ' ' || COALESCE(context, '') || ' ' || COALESCE(text_signals, ''),
|
||||
'llmlingua2'
|
||||
)::bm25_catalog.bm25vector
|
||||
FROM input_data
|
||||
RETURNING id
|
||||
"""
|
||||
elif config.text_search_extension == "native":
|
||||
# search_vector is a regular tsvector column populated here using the
|
||||
# configured native dictionary. It used to be GENERATED ALWAYS with
|
||||
# a hardcoded 'english', which prevented per-deployment language
|
||||
# configuration. text_search_extension_native_language is validated
|
||||
# in HindsightConfig.validate() as a PG identifier, so embedding it
|
||||
# as a SQL literal is safe.
|
||||
query = f"""
|
||||
WITH input_data AS (
|
||||
SELECT * FROM unnest(
|
||||
$2::text[], $3::vector[], $4::timestamptz[], $5::timestamptz[], $6::timestamptz[], $7::timestamptz[],
|
||||
$8::text[], $9::text[], $10::jsonb[], $11::text[], $12::text[], $13::jsonb[], $14::jsonb[], $15::text[]
|
||||
) AS t(text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
context, fact_type, metadata, chunk_id, document_id, tags_json,
|
||||
observation_scopes_json, text_signals)
|
||||
)
|
||||
INSERT INTO {table} (bank_id, text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
context, fact_type, metadata, chunk_id, document_id, tags,
|
||||
observation_scopes, text_signals, search_vector)
|
||||
SELECT
|
||||
$1,
|
||||
text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
context, fact_type, metadata, chunk_id, document_id,
|
||||
COALESCE(
|
||||
(SELECT array_agg(elem) FROM jsonb_array_elements_text(tags_json) AS elem),
|
||||
'{{}}'::varchar[]
|
||||
),
|
||||
observation_scopes_json,
|
||||
text_signals,
|
||||
to_tsvector(
|
||||
'{config.text_search_extension_native_language}'::regconfig,
|
||||
COALESCE(text, '') || ' ' || COALESCE(context, '') || ' ' || COALESCE(text_signals, '')
|
||||
)
|
||||
FROM input_data
|
||||
RETURNING id
|
||||
"""
|
||||
else:
|
||||
# pg_textsearch, pgroonga, and pg_search: search_vector is a dummy
|
||||
# TEXT column; the actual full-text index operates on the base text
|
||||
# columns directly, so we don't populate search_vector at insert time.
|
||||
query = f"""
|
||||
WITH input_data AS (
|
||||
SELECT * FROM unnest(
|
||||
$2::text[], $3::vector[], $4::timestamptz[], $5::timestamptz[], $6::timestamptz[], $7::timestamptz[],
|
||||
$8::text[], $9::text[], $10::jsonb[], $11::text[], $12::text[], $13::jsonb[], $14::jsonb[], $15::text[]
|
||||
) AS t(text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
context, fact_type, metadata, chunk_id, document_id, tags_json,
|
||||
observation_scopes_json, text_signals)
|
||||
)
|
||||
INSERT INTO {table} (bank_id, text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
context, fact_type, metadata, chunk_id, document_id, tags,
|
||||
observation_scopes, text_signals)
|
||||
SELECT
|
||||
$1,
|
||||
text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
context, fact_type, metadata, chunk_id, document_id,
|
||||
COALESCE(
|
||||
(SELECT array_agg(elem) FROM jsonb_array_elements_text(tags_json) AS elem),
|
||||
'{{}}'::varchar[]
|
||||
),
|
||||
observation_scopes_json,
|
||||
text_signals
|
||||
FROM input_data
|
||||
RETURNING id
|
||||
"""
|
||||
# search_vector is populated inline for backends that store a real vector
|
||||
# (native tsvector, vchord bm25vector). pgroonga / pg_textsearch / pg_search
|
||||
# index the base text columns directly and keep only a dummy column, so the
|
||||
# expression is None and the column is left out of the insert entirely.
|
||||
# Same expression is reused by curation revert (see pg_search_vector_expr).
|
||||
sv_expr = pg_search_vector_expr(config)
|
||||
sv_insert_col = ", search_vector" if sv_expr else ""
|
||||
sv_select_val = f",\n {sv_expr}" if sv_expr else ""
|
||||
query = f"""
|
||||
WITH input_data AS (
|
||||
SELECT * FROM unnest(
|
||||
$2::text[], $3::vector[], $4::timestamptz[], $5::timestamptz[], $6::timestamptz[], $7::timestamptz[],
|
||||
$8::text[], $9::text[], $10::jsonb[], $11::text[], $12::text[], $13::jsonb[], $14::jsonb[], $15::text[]
|
||||
) AS t(text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
context, fact_type, metadata, chunk_id, document_id, tags_json,
|
||||
observation_scopes_json, text_signals)
|
||||
)
|
||||
INSERT INTO {table} (bank_id, text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
context, fact_type, metadata, chunk_id, document_id, tags,
|
||||
observation_scopes, text_signals{sv_insert_col})
|
||||
SELECT
|
||||
$1,
|
||||
text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
context, fact_type, metadata, chunk_id, document_id,
|
||||
COALESCE(
|
||||
(SELECT array_agg(elem) FROM jsonb_array_elements_text(tags_json) AS elem),
|
||||
'{{}}'::varchar[]
|
||||
),
|
||||
observation_scopes_json,
|
||||
text_signals{sv_select_val}
|
||||
FROM input_data
|
||||
RETURNING id
|
||||
"""
|
||||
|
||||
results = await conn.fetch(
|
||||
query,
|
||||
@@ -425,19 +392,48 @@ class PostgreSQLOps(DataAccessOps):
|
||||
# Scope by joining through entities.bank_id (entity_cooccurrences itself
|
||||
# has no bank_id column — entities don't span banks, so scoping via
|
||||
# entity_id_1 is sufficient).
|
||||
#
|
||||
# Ordered locking (deadlock avoidance, #2529): retain's concurrent
|
||||
# cooccurrence upsert (entity_resolver._flush_pending) locks rows in
|
||||
# sorted (entity_id_1, entity_id_2) order — sorted specifically to give
|
||||
# every writer one consistent lock-acquisition order. A plain
|
||||
# `DELETE ... USING` scans/locks in whatever order the join plan picks,
|
||||
# so it could lock the same rows in the opposite order and cycle. We
|
||||
# instead select the victims in that same sorted order `FOR UPDATE`
|
||||
# first — the locking clause materialises the CTE and places LockRows
|
||||
# above the Sort, so locks are acquired ascending, matching the upsert —
|
||||
# then delete the already-locked rows. Same order on both sides ⇒ no
|
||||
# cycle (the deadlock is prevented, not merely retried). The Pass 2/3
|
||||
# retry wrap in run_graph_maintenance_job stays as a backstop for the
|
||||
# residual paths (FK cascade from prune_orphan_entities, Oracle).
|
||||
#
|
||||
# The staleness predicate is an INTERSECT of the two entities' unit sets
|
||||
# rather than the equivalent `unit_entities u1 JOIN u2 ON u1.unit_id =
|
||||
# u2.unit_id` self-join (#2473): both INTERSECT branches resolve as Index
|
||||
# Only Scans on idx_unit_entities_entity_unit (entity_id, unit_id), so the
|
||||
# per-pair cost is bounded by the two entities' degrees. The self-join let
|
||||
# the planner pick an anti-join that rescanned a high-degree hub entity's
|
||||
# membership set for every pair — 28-30min on a bank with a ~100K-membership
|
||||
# hub, even when zero rows were stale. Don't "simplify" it back.
|
||||
result = await conn.execute(
|
||||
f"""
|
||||
WITH victims AS (
|
||||
SELECT c.entity_id_1, c.entity_id_2
|
||||
FROM {ec_table} c
|
||||
JOIN {entities_table} e ON e.id = c.entity_id_1
|
||||
WHERE e.bank_id = $1
|
||||
AND NOT EXISTS (
|
||||
SELECT unit_id FROM {ue_table} WHERE entity_id = c.entity_id_1
|
||||
INTERSECT
|
||||
SELECT unit_id FROM {ue_table} WHERE entity_id = c.entity_id_2
|
||||
)
|
||||
ORDER BY c.entity_id_1, c.entity_id_2
|
||||
FOR UPDATE OF c
|
||||
)
|
||||
DELETE FROM {ec_table} c
|
||||
USING {entities_table} e
|
||||
WHERE e.id = c.entity_id_1
|
||||
AND e.bank_id = $1
|
||||
AND NOT EXISTS (
|
||||
SELECT 1
|
||||
FROM {ue_table} u1
|
||||
JOIN {ue_table} u2 ON u1.unit_id = u2.unit_id
|
||||
WHERE u1.entity_id = c.entity_id_1
|
||||
AND u2.entity_id = c.entity_id_2
|
||||
)
|
||||
USING victims v
|
||||
WHERE c.entity_id_1 = v.entity_id_1
|
||||
AND c.entity_id_2 = v.entity_id_2
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
@@ -541,11 +537,18 @@ class PostgreSQLOps(DataAccessOps):
|
||||
FROM {ue_table} ue_target
|
||||
WHERE ue_target.entity_id = se.entity_id
|
||||
AND ue_target.unit_id != ALL($1::uuid[])
|
||||
-- Filter before applying the cap: candidates from other fact
|
||||
-- types must not consume this entity's bounded fan-out.
|
||||
AND EXISTS (
|
||||
SELECT 1
|
||||
FROM {mu_table} mu_target
|
||||
WHERE mu_target.id = ue_target.unit_id
|
||||
AND mu_target.fact_type = $2
|
||||
)
|
||||
ORDER BY ue_target.unit_id DESC
|
||||
LIMIT {per_entity_limit}
|
||||
) t
|
||||
JOIN {mu_table} mu ON mu.id = t.unit_id
|
||||
WHERE mu.fact_type = $2
|
||||
GROUP BY mu.id
|
||||
ORDER BY score DESC
|
||||
LIMIT $3
|
||||
@@ -894,6 +897,93 @@ class PostgreSQLOps(DataAccessOps):
|
||||
|
||||
# -- Task claiming operations ------------------------------------------
|
||||
|
||||
async def prune_terminal_operations(
|
||||
self,
|
||||
conn: DatabaseConnection,
|
||||
table: str,
|
||||
cutoff: datetime,
|
||||
*,
|
||||
batch_size: int,
|
||||
) -> int:
|
||||
# Lock only the bounded candidate set. SKIP LOCKED lets multiple
|
||||
# workers prune disjoint batches without waiting or double-deleting.
|
||||
# Cancelled children cannot complete parent aggregation, so retain the
|
||||
# parent guard only for completed/failed children. Before removing a
|
||||
# cancelled child, preserve its signal by cancelling a pending parent
|
||||
# in this transaction and refreshing the parent's retention window.
|
||||
candidates = await conn.fetch(
|
||||
f"""
|
||||
SELECT candidate_operation.operation_id
|
||||
FROM {table} candidate_operation
|
||||
WHERE candidate_operation.status IN ('completed', 'failed', 'cancelled')
|
||||
AND candidate_operation.updated_at < $1
|
||||
AND (
|
||||
candidate_operation.status = 'cancelled'
|
||||
OR NOT EXISTS (
|
||||
SELECT 1
|
||||
FROM {table} parent
|
||||
WHERE parent.operation_id = CASE
|
||||
WHEN candidate_operation.result_metadata->>'parent_operation_id'
|
||||
~* '^[0-9a-f]{{8}}-[0-9a-f]{{4}}-[0-9a-f]{{4}}-[0-9a-f]{{4}}-[0-9a-f]{{12}}$'
|
||||
THEN (candidate_operation.result_metadata->>'parent_operation_id')::uuid
|
||||
ELSE NULL
|
||||
END
|
||||
AND parent.bank_id = candidate_operation.bank_id
|
||||
)
|
||||
)
|
||||
ORDER BY candidate_operation.updated_at, candidate_operation.operation_id
|
||||
LIMIT $2
|
||||
FOR UPDATE OF candidate_operation SKIP LOCKED
|
||||
""",
|
||||
cutoff,
|
||||
batch_size,
|
||||
)
|
||||
if not candidates:
|
||||
return 0
|
||||
|
||||
candidate_ids = [row["operation_id"] for row in candidates]
|
||||
await conn.execute(
|
||||
f"""
|
||||
UPDATE {table} parent
|
||||
SET status = 'cancelled',
|
||||
updated_at = now(),
|
||||
completed_at = COALESCE(parent.completed_at, now()),
|
||||
error_message = COALESCE(
|
||||
parent.error_message,
|
||||
'Cancelled because a child operation was cancelled'
|
||||
)
|
||||
WHERE parent.status = 'pending'
|
||||
AND EXISTS (
|
||||
SELECT 1
|
||||
FROM {table} candidate_operation
|
||||
WHERE candidate_operation.operation_id = ANY($1)
|
||||
AND candidate_operation.status = 'cancelled'
|
||||
AND candidate_operation.updated_at < $2
|
||||
AND candidate_operation.bank_id = parent.bank_id
|
||||
AND parent.operation_id = CASE
|
||||
WHEN candidate_operation.result_metadata->>'parent_operation_id'
|
||||
~* '^[0-9a-f]{{8}}-[0-9a-f]{{4}}-[0-9a-f]{{4}}-[0-9a-f]{{4}}-[0-9a-f]{{12}}$'
|
||||
THEN (candidate_operation.result_metadata->>'parent_operation_id')::uuid
|
||||
ELSE NULL
|
||||
END
|
||||
)
|
||||
""",
|
||||
candidate_ids,
|
||||
cutoff,
|
||||
)
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
DELETE FROM {table}
|
||||
WHERE operation_id = ANY($1)
|
||||
AND status IN ('completed', 'failed', 'cancelled')
|
||||
AND updated_at < $2
|
||||
RETURNING operation_id
|
||||
""",
|
||||
candidate_ids,
|
||||
cutoff,
|
||||
)
|
||||
return len(rows)
|
||||
|
||||
async def _claim_consolidation_tasks(
|
||||
self,
|
||||
conn,
|
||||
|
||||
@@ -1242,6 +1242,10 @@ class OracleBackend(DatabaseBackend):
|
||||
def __init__(self) -> None:
|
||||
self._pool: Any = None
|
||||
self._oracledb: Any = None
|
||||
# Oracle pooled sessions retain CURRENT_SCHEMA across checkouts. Cache
|
||||
# SESSION_USER so default-schema acquisitions can explicitly reset a
|
||||
# connection that was previously used for a tenant schema.
|
||||
self._default_schema: str | None = None
|
||||
|
||||
async def initialize(
|
||||
self,
|
||||
@@ -1277,11 +1281,17 @@ class OracleBackend(DatabaseBackend):
|
||||
logger.info(f"Oracle pool created (min={min_size}, max={max_size})")
|
||||
|
||||
async def shutdown(self) -> None:
|
||||
if self._pool is not None:
|
||||
await self._pool.close(force=True)
|
||||
self._pool = None
|
||||
# Drop the reference before awaiting close() so is_ready flips False for
|
||||
# the whole teardown, not just after it completes (see PostgreSQLBackend).
|
||||
pool, self._pool = self._pool, None
|
||||
if pool is not None:
|
||||
await pool.close(force=True)
|
||||
logger.info("Oracle pool closed")
|
||||
|
||||
@property
|
||||
def is_ready(self) -> bool:
|
||||
return self._pool is not None
|
||||
|
||||
async def _set_session_schema(self, conn: Any) -> None:
|
||||
"""Set the session schema on an Oracle connection.
|
||||
|
||||
@@ -1294,10 +1304,23 @@ class OracleBackend(DatabaseBackend):
|
||||
from ..memory_engine import get_current_schema
|
||||
|
||||
schema = get_current_schema()
|
||||
if schema and schema != "public":
|
||||
cursor = conn.cursor()
|
||||
await cursor.execute(f'ALTER SESSION SET CURRENT_SCHEMA = "{schema}"')
|
||||
await cursor.close()
|
||||
cursor = conn.cursor()
|
||||
try:
|
||||
if self._default_schema is None:
|
||||
await cursor.execute("SELECT SYS_CONTEXT('USERENV', 'SESSION_USER') FROM DUAL")
|
||||
row = await cursor.fetchone()
|
||||
if not row or not row[0]:
|
||||
raise RuntimeError("Oracle did not return SESSION_USER while resetting CURRENT_SCHEMA")
|
||||
self._default_schema = str(row[0])
|
||||
|
||||
target_schema = self._default_schema if not schema or schema == "public" else schema
|
||||
safe_schema = target_schema.replace('"', '""')
|
||||
await cursor.execute(f'ALTER SESSION SET CURRENT_SCHEMA = "{safe_schema}"')
|
||||
finally:
|
||||
# oracledb's AsyncCursor.close() is synchronous (not a coroutine);
|
||||
# awaiting it raises "object NoneType can't be used in 'await'
|
||||
# expression" and aborts every acquire().
|
||||
cursor.close()
|
||||
|
||||
@asynccontextmanager
|
||||
async def acquire(self) -> AsyncIterator[OracleConnection]:
|
||||
|
||||
@@ -95,7 +95,12 @@ class PostgreSQLBackend(DatabaseBackend):
|
||||
command_timeout=command_timeout,
|
||||
statement_cache_size=statement_cache_size,
|
||||
timeout=acquire_timeout,
|
||||
# init runs once per new connection; setup runs on every acquire,
|
||||
# after asyncpg's release-time RESET ALL. Passing init_callback as
|
||||
# both keeps the per-connection session GUCs (hnsw.ef_search, etc.)
|
||||
# applied after a connection is reused, not just on first creation.
|
||||
init=init_callback,
|
||||
setup=init_callback,
|
||||
)
|
||||
logger.info(
|
||||
f"PostgreSQL pool created (min={min_size}, max={max_size}, "
|
||||
@@ -103,11 +108,19 @@ class PostgreSQLBackend(DatabaseBackend):
|
||||
)
|
||||
|
||||
async def shutdown(self) -> None:
|
||||
if self._pool is not None:
|
||||
await self._pool.close()
|
||||
self._pool = None
|
||||
# Drop the reference *before* awaiting close(): closing is not
|
||||
# instantaneous, and anything acquiring during that window would
|
||||
# otherwise get an asyncpg "pool is closing" error rather than seeing
|
||||
# is_ready False.
|
||||
pool, self._pool = self._pool, None
|
||||
if pool is not None:
|
||||
await pool.close()
|
||||
logger.info("PostgreSQL pool closed")
|
||||
|
||||
@property
|
||||
def is_ready(self) -> bool:
|
||||
return self._pool is not None
|
||||
|
||||
@asynccontextmanager
|
||||
async def acquire(self) -> AsyncIterator[PostgresConnection]:
|
||||
pool = self._ensure_pool()
|
||||
|
||||
@@ -4,6 +4,7 @@ Database utility functions for connection management with retry logic.
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import random
|
||||
import time
|
||||
from collections.abc import AsyncIterator
|
||||
from contextlib import AsyncExitStack, asynccontextmanager
|
||||
@@ -16,6 +17,20 @@ DEFAULT_MAX_RETRIES = 3
|
||||
DEFAULT_BASE_DELAY = 0.5 # seconds
|
||||
DEFAULT_MAX_DELAY = 5.0 # seconds
|
||||
|
||||
|
||||
def _backoff_delay(attempt: int, base_delay: float, max_delay: float) -> float:
|
||||
"""Exponential backoff with equal jitter.
|
||||
|
||||
Deterministic backoff makes concurrent retriers wake in lock-step and
|
||||
re-collide on the very same rows, re-triggering the deadlock they just
|
||||
backed off from. "Equal jitter" — half the window fixed, half random —
|
||||
keeps a floor (so we don't hot-spin) while decorrelating the wake-ups, so
|
||||
two contenders that deadlocked together are very unlikely to retry in sync.
|
||||
"""
|
||||
ceil = min(base_delay * (2**attempt), max_delay)
|
||||
return ceil / 2 + random.uniform(0, ceil / 2)
|
||||
|
||||
|
||||
# Retryable exception types (checked by class name to avoid hard imports)
|
||||
_RETRYABLE_EXCEPTION_NAMES = frozenset(
|
||||
{
|
||||
@@ -78,7 +93,7 @@ async def retry_with_backoff(
|
||||
raise
|
||||
last_exception = e
|
||||
if attempt < max_retries:
|
||||
delay = min(base_delay * (2**attempt), max_delay)
|
||||
delay = _backoff_delay(attempt, base_delay, max_delay)
|
||||
if type(e).__name__ == "DeadlockDetectedError" or _is_oracle_deadlock(e):
|
||||
logger.warning(
|
||||
"Deadlock detected during parallel document processing — "
|
||||
@@ -136,7 +151,7 @@ async def acquire_with_retry(backend_or_pool: Any, max_retries: int = DEFAULT_MA
|
||||
if not _is_retryable(e):
|
||||
raise
|
||||
if attempt < max_retries:
|
||||
delay = min(DEFAULT_BASE_DELAY * (2**attempt), DEFAULT_MAX_DELAY)
|
||||
delay = _backoff_delay(attempt, DEFAULT_BASE_DELAY, DEFAULT_MAX_DELAY)
|
||||
logger.warning(
|
||||
f"Database acquire failed (attempt {attempt + 1}/{max_retries + 1}): {e}. "
|
||||
f"Retrying in {delay:.1f}s..."
|
||||
|
||||
@@ -76,6 +76,25 @@ class _ZeroEntropyEmbedResponse(BaseModel):
|
||||
results: list[_ZeroEntropyEmbedResult]
|
||||
|
||||
|
||||
def _truncate_to_tokens(text: str, max_tokens: int) -> tuple[str, int]:
|
||||
"""Truncate ``text`` to at most ``max_tokens`` cl100k_base tokens.
|
||||
|
||||
tiktoken is an approximation of any given provider's tokenizer, so set
|
||||
``max_tokens`` with a little headroom below the model's real limit.
|
||||
|
||||
Returns the (possibly truncated) text and the original token count (so the
|
||||
caller can report how much was dropped); the count equals ``len(tokens)``
|
||||
whether or not truncation occurred.
|
||||
"""
|
||||
from .token_encoding import get_token_encoding
|
||||
|
||||
enc = get_token_encoding()
|
||||
tokens = enc.encode(text)
|
||||
if len(tokens) <= max_tokens:
|
||||
return text, len(tokens)
|
||||
return enc.decode(tokens[:max_tokens]), len(tokens)
|
||||
|
||||
|
||||
class Embeddings(ABC):
|
||||
"""
|
||||
Abstract base class for embedding generation.
|
||||
@@ -1202,6 +1221,7 @@ class LiteLLMSDKEmbeddings(Embeddings):
|
||||
batch_size: int = 100,
|
||||
timeout: float = 60.0,
|
||||
encoding_format: str | None = "float",
|
||||
max_input_tokens: int | None = None,
|
||||
):
|
||||
"""
|
||||
Initialize LiteLLM SDK embeddings client.
|
||||
@@ -1216,6 +1236,10 @@ class LiteLLMSDKEmbeddings(Embeddings):
|
||||
timeout: Request timeout in seconds (default: 60.0)
|
||||
encoding_format: Encoding format for embeddings (default: "float").
|
||||
Set to None or empty string to omit (needed for Voyage AI, Gemini).
|
||||
max_input_tokens: If set, truncate each input text to this many tokens
|
||||
(tiktoken cl100k_base) before embedding. Needed for models with a
|
||||
fixed input-token limit (e.g. Bedrock Titan V2's hard 8192 cap),
|
||||
where an oversized text would otherwise fail permanently (#2501).
|
||||
"""
|
||||
self.api_key = api_key
|
||||
self.model = model
|
||||
@@ -1224,6 +1248,7 @@ class LiteLLMSDKEmbeddings(Embeddings):
|
||||
self.batch_size = batch_size
|
||||
self.timeout = timeout
|
||||
self.encoding_format = encoding_format or None
|
||||
self.max_input_tokens = max_input_tokens
|
||||
self._litellm = None # Will be set during initialization
|
||||
self._dimension: int | None = None
|
||||
|
||||
@@ -1300,6 +1325,33 @@ class LiteLLMSDKEmbeddings(Embeddings):
|
||||
if not texts:
|
||||
return []
|
||||
|
||||
# Truncate oversized inputs before hitting the provider. Models with a
|
||||
# fixed input-token limit (e.g. Bedrock Titan V2, 8192) reject an
|
||||
# oversized text with a permanent error rather than truncating it
|
||||
# server-side, which strands the caller (e.g. a delta mental model whose
|
||||
# content grew past the cap) with no recovery path. See #2501.
|
||||
if self.max_input_tokens is not None:
|
||||
truncated_texts = []
|
||||
original_token_counts = []
|
||||
for t in texts:
|
||||
new_text, original_tokens = _truncate_to_tokens(t, self.max_input_tokens)
|
||||
truncated_texts.append(new_text)
|
||||
if original_tokens > self.max_input_tokens:
|
||||
original_token_counts.append(original_tokens)
|
||||
texts = truncated_texts
|
||||
if original_token_counts:
|
||||
logger.warning(
|
||||
"Embeddings: truncated %d of %d input(s) to %d tokens for model %s "
|
||||
"(largest was ~%d tokens); embedded content is incomplete. "
|
||||
"This usually means a mental model's content has grown past the model's "
|
||||
"input limit — see issue #2501.",
|
||||
len(original_token_counts),
|
||||
len(texts),
|
||||
self.max_input_tokens,
|
||||
self.model,
|
||||
max(original_token_counts),
|
||||
)
|
||||
|
||||
all_embeddings = []
|
||||
|
||||
# Process in batches
|
||||
@@ -1691,6 +1743,7 @@ def create_embeddings_from_env() -> Embeddings:
|
||||
api_base=config.embeddings_litellm_sdk_api_base,
|
||||
output_dimensions=config.embeddings_litellm_sdk_output_dimensions,
|
||||
encoding_format=config.embeddings_litellm_sdk_encoding_format,
|
||||
max_input_tokens=config.embeddings_litellm_sdk_max_input_tokens,
|
||||
)
|
||||
elif provider == "google":
|
||||
vertexai_project_id = config.embeddings_vertexai_project_id
|
||||
|
||||
@@ -9,6 +9,7 @@ import asyncio
|
||||
import json
|
||||
import logging
|
||||
from collections import defaultdict
|
||||
from collections.abc import Iterator
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import UTC, datetime
|
||||
from difflib import SequenceMatcher
|
||||
@@ -75,6 +76,22 @@ def _later_date(a: datetime | None, b: datetime | None) -> datetime | None:
|
||||
return a if a > b else b
|
||||
|
||||
|
||||
def _canonical_cooccurrence_pairs(entity_list: list[str]) -> Iterator[tuple[str, str]]:
|
||||
"""Yield each distinct pair of ``entity_list`` as ``(a, b)`` with ``a < b``.
|
||||
|
||||
Canonical ordering matches the entity_cooccurrences PK and check constraint.
|
||||
The pair is ordered into fresh locals rather than by swapping the loop
|
||||
variables: ``entity_id_1`` is the outer iterate, so swapping it would leak
|
||||
into the remaining inner iterations and build later pairs off the wrong
|
||||
element.
|
||||
"""
|
||||
for i, entity_id_1 in enumerate(entity_list):
|
||||
for entity_id_2 in entity_list[i + 1 :]:
|
||||
if entity_id_1 == entity_id_2:
|
||||
continue
|
||||
yield (entity_id_1, entity_id_2) if entity_id_1 < entity_id_2 else (entity_id_2, entity_id_1)
|
||||
|
||||
|
||||
@dataclass
|
||||
class _CooccurrencePair:
|
||||
"""A (entity_id_1, entity_id_2) pair observed in a retain batch (for post-txn flush)."""
|
||||
@@ -853,20 +870,12 @@ class EntityResolver:
|
||||
for unit_id, entity_ids in unit_to_entities.items():
|
||||
entity_list = list(entity_ids)
|
||||
event_date = unit_event_date.get(unit_id)
|
||||
for i, entity_id_1 in enumerate(entity_list):
|
||||
for entity_id_2 in entity_list[i + 1 :]:
|
||||
if entity_id_1 == entity_id_2:
|
||||
continue
|
||||
# Canonical ordering (entity_id_1 < entity_id_2) matches the
|
||||
# entity_cooccurrences PK and check constraint.
|
||||
if entity_id_1 > entity_id_2:
|
||||
entity_id_1, entity_id_2 = entity_id_2, entity_id_1
|
||||
key = (entity_id_1, entity_id_2)
|
||||
prev = cooccurrence_pairs.get(key, _SENTINEL_MISSING)
|
||||
if prev is _SENTINEL_MISSING:
|
||||
cooccurrence_pairs[key] = event_date
|
||||
else:
|
||||
cooccurrence_pairs[key] = _later_date(prev, event_date)
|
||||
for key in _canonical_cooccurrence_pairs(entity_list):
|
||||
prev = cooccurrence_pairs.get(key, _SENTINEL_MISSING)
|
||||
if prev is _SENTINEL_MISSING:
|
||||
cooccurrence_pairs[key] = event_date
|
||||
else:
|
||||
cooccurrence_pairs[key] = _later_date(prev, event_date)
|
||||
|
||||
# Accumulate co-occurrence pairs for post-transaction flush.
|
||||
# The actual INSERT/UPDATE is deferred to flush_pending_stats() to avoid
|
||||
|
||||
@@ -66,6 +66,20 @@ MAX_SEMANTIC_LINKS_PER_UNIT = 50
|
||||
# under 1s.
|
||||
_DRAIN_BATCH_SIZE = 50
|
||||
|
||||
# Retry budget for the idempotent Pass 2/3 entity/cooccurrence sweep. Higher
|
||||
# than db_utils' default (3) because the sweep has no client waiting on it and
|
||||
# is safe to rerun, so we'd rather spend a longer jittered-backoff tail than
|
||||
# drop a maintenance pass and leak stale graph rows (see run_graph_maintenance_job).
|
||||
_SWEEP_MAX_RETRIES = 8
|
||||
|
||||
|
||||
@dataclass
|
||||
class _SweepCounts:
|
||||
"""Prune counts returned by the Pass 2/3 sweep (avoids a bare tuple return)."""
|
||||
|
||||
orphan_entities_pruned: int
|
||||
stale_cooccurrences_pruned: int
|
||||
|
||||
|
||||
@dataclass
|
||||
class JobResult:
|
||||
@@ -203,27 +217,51 @@ async def run_graph_maintenance_job(
|
||||
|
||||
# --- Pass 2 & 3: entity / cooccurrence sweeps ---
|
||||
# Bank-wide single-statement deletes. Cheap when there's nothing to do.
|
||||
#
|
||||
# Unlike Pass 1's queue claim, these DELETEs aren't protected by any
|
||||
# consistent lock-ordering guarantee: prune_stale_cooccurrences scans
|
||||
# entity_cooccurrences via a join/NOT EXISTS plan, while retain's
|
||||
# concurrent cooccurrence upserts (entity_resolver._flush_pending) lock
|
||||
# the same rows in sorted (entity_id_1, entity_id_2) order. When a sweep
|
||||
# and a concurrent upsert touch overlapping rows in opposite orders,
|
||||
# Postgres detects a genuine circular wait and aborts one side with
|
||||
# DeadlockDetectedError. Both prunes are idempotent bank-wide sweeps —
|
||||
# rerunning only deletes what's still stale — so retrying the whole
|
||||
# transaction on deadlock is safe.
|
||||
from .db_utils import retry_with_backoff
|
||||
from .memory_engine import acquire_with_retry
|
||||
|
||||
async with acquire_with_retry(backend) as conn:
|
||||
async with conn.transaction():
|
||||
result.orphan_entities_pruned = await ops.prune_orphan_entities(
|
||||
conn,
|
||||
fq_table("entities"),
|
||||
fq_table("unit_entities"),
|
||||
bank_id,
|
||||
)
|
||||
# The orphan prune above cascades cooccurrences via FK. The
|
||||
# explicit cooccurrence pass below catches the *stale-count*
|
||||
# case: both entities still exist but no current unit witnesses
|
||||
# them together.
|
||||
result.stale_cooccurrences_pruned = await ops.prune_stale_cooccurrences(
|
||||
conn,
|
||||
fq_table("entity_cooccurrences"),
|
||||
fq_table("unit_entities"),
|
||||
fq_table("entities"),
|
||||
bank_id,
|
||||
)
|
||||
async def _run_sweep() -> _SweepCounts:
|
||||
async with acquire_with_retry(backend) as conn:
|
||||
async with conn.transaction():
|
||||
orphan_pruned = await ops.prune_orphan_entities(
|
||||
conn,
|
||||
fq_table("entities"),
|
||||
fq_table("unit_entities"),
|
||||
bank_id,
|
||||
)
|
||||
# The orphan prune above cascades cooccurrences via FK. The
|
||||
# explicit cooccurrence pass below catches the *stale-count*
|
||||
# case: both entities still exist but no current unit
|
||||
# witnesses them together.
|
||||
stale_pruned = await ops.prune_stale_cooccurrences(
|
||||
conn,
|
||||
fq_table("entity_cooccurrences"),
|
||||
fq_table("unit_entities"),
|
||||
fq_table("entities"),
|
||||
bank_id,
|
||||
)
|
||||
return _SweepCounts(orphan_entities_pruned=orphan_pruned, stale_cooccurrences_pruned=stale_pruned)
|
||||
|
||||
# A larger retry budget than the default (3): this is idempotent background
|
||||
# maintenance with no client waiting on it, so a longer retry tail costs
|
||||
# nothing, whereas a dropped sweep silently leaks orphan entities / stale
|
||||
# cooccurrences until the next run. With jittered backoff a single sweep
|
||||
# contending against continuous retain upserts effectively never exhausts
|
||||
# this budget (each retry independently clears with high probability).
|
||||
sweep = await retry_with_backoff(_run_sweep, max_retries=_SWEEP_MAX_RETRIES)
|
||||
result.orphan_entities_pruned = sweep.orphan_entities_pruned
|
||||
result.stale_cooccurrences_pruned = sweep.stale_cooccurrences_pruned
|
||||
|
||||
elapsed = time.time() - job_start
|
||||
logger.info(
|
||||
|
||||
@@ -116,6 +116,7 @@ class LLMInterface(ABC):
|
||||
max_backoff: float = 30.0,
|
||||
tool_choice: str | dict[str, Any] = "auto",
|
||||
cached_prefix: str | None = None,
|
||||
cached_prefix_message_count: int = 0,
|
||||
) -> LLMToolCallResult:
|
||||
"""
|
||||
Make an LLM API call with tool/function calling support.
|
||||
@@ -185,6 +186,45 @@ class LLMInterface(ABC):
|
||||
"""
|
||||
return None
|
||||
|
||||
# ── Step-by-step incremental prompt caching (optional) ─────────────────────
|
||||
#
|
||||
# For agentic loops (reflect) the dominant cost is the conversation prefix
|
||||
# re-sent every turn, not the static system prefix. Providers that can cache
|
||||
# a *growing* prefix implement these: the caller rolls one cache per step
|
||||
# (each covering the previous step's full input), passes its handle plus the
|
||||
# message count it covers to ``call_with_tools`` so only the new turns are
|
||||
# sent fresh, and tears the caches down when the loop ends. Default no-ops so
|
||||
# non-supporting providers transparently run uncached.
|
||||
|
||||
def supports_incremental_prompt_cache(self) -> bool:
|
||||
"""Whether this provider can cache a growing multi-turn conversation prefix."""
|
||||
return False
|
||||
|
||||
async def create_incremental_cache(
|
||||
self,
|
||||
*,
|
||||
session_id: str,
|
||||
messages: list[dict[str, Any]],
|
||||
tools: list[dict[str, Any]] | None = None,
|
||||
) -> str | None:
|
||||
"""Cache ``system + tools + messages`` and return an opaque handle, or None.
|
||||
|
||||
The handle is passed back to ``call_with_tools(cached_prefix=...,
|
||||
cached_prefix_message_count=len(messages))``. Caches are grouped under
|
||||
``session_id`` for teardown via ``delete_cache_session``. Returns None
|
||||
when caching is unavailable or the prefix is too small — caller falls
|
||||
back to an uncached call.
|
||||
"""
|
||||
return None
|
||||
|
||||
async def delete_cached_prefix(self, name: str) -> None:
|
||||
"""Best-effort delete of a single cache handle (a superseded step)."""
|
||||
return None
|
||||
|
||||
async def delete_cache_session(self, session_id: str) -> None:
|
||||
"""Best-effort teardown of every cache created under ``session_id``."""
|
||||
return None
|
||||
|
||||
async def submit_batch(
|
||||
self,
|
||||
requests: list[dict[str, Any]],
|
||||
|
||||
@@ -376,6 +376,26 @@ class LLMTraceRecorder:
|
||||
# INSERTs it patches — but it must not block on unrelated operations).
|
||||
self._pending: dict[str | None, set[asyncio.Task]] = {}
|
||||
|
||||
def _writable(self) -> Any | None:
|
||||
"""Return the pool to write through, or None if writing isn't possible.
|
||||
|
||||
Covers the two lifecycle windows in which best-effort trace writes must
|
||||
be skipped rather than attempted: before the backend pool is created
|
||||
(``initialize()`` verifies the LLM before the DB is up) and during/after
|
||||
shutdown. Writes already in flight need no handling — the pools close
|
||||
gracefully, waiting for their connections to be released.
|
||||
"""
|
||||
pool = self._pool_getter()
|
||||
if pool is None:
|
||||
return None
|
||||
# Backends declare readiness explicitly; a raw pool (some callers pass
|
||||
# one directly) has no lifecycle flag and is assumed usable.
|
||||
from .db.base import DatabaseBackend
|
||||
|
||||
if isinstance(pool, DatabaseBackend) and not pool.is_ready:
|
||||
return None
|
||||
return pool
|
||||
|
||||
def is_enabled(self, scope: str) -> bool:
|
||||
"""Whether tracing is active for the given call scope."""
|
||||
if not self._enabled:
|
||||
@@ -473,7 +493,7 @@ class LLMTraceRecorder:
|
||||
|
||||
async def _safe_write(self, record: LLMRequestRecord) -> None:
|
||||
"""Write a trace row. Errors are logged, never raised."""
|
||||
pool = self._pool_getter()
|
||||
pool = self._writable()
|
||||
if pool is None:
|
||||
logger.debug("LLM trace skipped: pool not available")
|
||||
return
|
||||
@@ -568,8 +588,9 @@ class LLMTraceRecorder:
|
||||
# so the UPDATE patches rows that already exist rather than racing ahead
|
||||
# of them (without blocking on unrelated operations' pending writes).
|
||||
await self._flush_pending(trace_id)
|
||||
pool = self._pool_getter()
|
||||
pool = self._writable()
|
||||
if pool is None:
|
||||
logger.debug("LLM trace memory_id attach skipped: pool not available")
|
||||
return
|
||||
try:
|
||||
schema = self._schema_getter()
|
||||
|
||||
@@ -31,9 +31,6 @@ from ..config import (
|
||||
if TYPE_CHECKING:
|
||||
from .response_models import LLMToolCallResult
|
||||
|
||||
# Seed applied to every Groq request for deterministic behavior.
|
||||
DEFAULT_LLM_SEED = 4242
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Disable httpx logging
|
||||
@@ -235,6 +232,17 @@ def requires_api_key(provider: str) -> bool:
|
||||
return provider.lower() not in _PROVIDERS_WITHOUT_API_KEY
|
||||
|
||||
|
||||
def _validate_ollama_num_ctx(value: Any) -> int | None:
|
||||
"""Validate a native Ollama context-window override."""
|
||||
if value is None:
|
||||
return None
|
||||
if isinstance(value, bool) or not isinstance(value, int):
|
||||
raise ValueError(f"ollama_num_ctx must be a positive integer, got {value!r}")
|
||||
if value < 1:
|
||||
raise ValueError(f"ollama_num_ctx must be >= 1, got {value}")
|
||||
return value
|
||||
|
||||
|
||||
def create_llm_provider(
|
||||
provider: str,
|
||||
api_key: str,
|
||||
@@ -254,6 +262,7 @@ def create_llm_provider(
|
||||
litellmrouter_config: dict[str, Any] | None = None,
|
||||
gemini_service_tier: str | None = None,
|
||||
timeout: float | None = None,
|
||||
ollama_num_ctx: int | None = None,
|
||||
) -> Any: # Returns LLMInterface
|
||||
"""
|
||||
Factory function to create the appropriate LLM provider implementation.
|
||||
@@ -268,6 +277,8 @@ def create_llm_provider(
|
||||
openai_service_tier: OpenAI service tier (for OpenAI provider) - None (default) or "flex" (50% cheaper).
|
||||
bedrock_service_tier: Bedrock service tier (for Bedrock provider) - None (default), "flex", "priority", or "reserved".
|
||||
gemini_service_tier: Gemini service tier (for Gemini provider) - None (default) or "flex" (50% cheaper).
|
||||
ollama_num_ctx: Native Ollama context window override. None lets Ollama use the
|
||||
model/server default.
|
||||
extra_body: Extra request-body params merged into the provider's native
|
||||
call. Threaded into OpenAI-compatible, Fireworks, Anthropic, Gemini/
|
||||
VertexAI and LiteLLM providers (each merges them in its own parameter
|
||||
@@ -291,6 +302,8 @@ def create_llm_provider(
|
||||
Returns:
|
||||
LLMInterface implementation for the specified provider.
|
||||
"""
|
||||
ollama_num_ctx = _validate_ollama_num_ctx(ollama_num_ctx)
|
||||
|
||||
from .providers import (
|
||||
AnthropicLLM,
|
||||
ClaudeCodeLLM,
|
||||
@@ -494,6 +507,7 @@ def create_llm_provider(
|
||||
groq_service_tier=groq_service_tier,
|
||||
openai_service_tier=openai_service_tier,
|
||||
extra_body=extra_body,
|
||||
ollama_num_ctx=ollama_num_ctx,
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
@@ -531,6 +545,7 @@ class LLMProvider:
|
||||
max_retries: int | None = None,
|
||||
initial_backoff: float | None = None,
|
||||
max_backoff: float | None = None,
|
||||
ollama_num_ctx: int | None = None,
|
||||
):
|
||||
"""
|
||||
Initialize LLM provider.
|
||||
@@ -545,6 +560,8 @@ class LLMProvider:
|
||||
openai_service_tier: OpenAI service tier (None or "flex") - from config.
|
||||
bedrock_service_tier: Bedrock service tier (None, "flex", "priority", "reserved") - from config.
|
||||
gemini_service_tier: Gemini service tier (None or "flex") - from config.
|
||||
ollama_num_ctx: Native Ollama context window override. ``None`` lets Ollama
|
||||
use the model/server default.
|
||||
gemini_safety_settings: Safety settings for Gemini/VertexAI providers.
|
||||
extra_body: Extra request-body params merged into the provider's native call
|
||||
(OpenAI-compatible, Fireworks, Anthropic, Gemini/VertexAI, LiteLLM).
|
||||
@@ -598,6 +615,7 @@ class LLMProvider:
|
||||
self.openai_service_tier = openai_service_tier
|
||||
self.bedrock_service_tier = bedrock_service_tier
|
||||
self.gemini_service_tier = gemini_service_tier
|
||||
self.ollama_num_ctx = _validate_ollama_num_ctx(ollama_num_ctx)
|
||||
# Gemini safety settings (instance default; can be overridden per-request via context var)
|
||||
self.gemini_safety_settings = gemini_safety_settings
|
||||
# Gemini prompt caching: when True, retain extraction (and any future
|
||||
@@ -742,6 +760,7 @@ class LLMProvider:
|
||||
gemini_safety_settings=self.gemini_safety_settings,
|
||||
prompt_cache_enabled=self.prompt_cache_enabled,
|
||||
litellmrouter_config=router_config,
|
||||
ollama_num_ctx=self.ollama_num_ctx,
|
||||
timeout=self.timeout,
|
||||
)
|
||||
|
||||
@@ -967,6 +986,7 @@ class LLMProvider:
|
||||
max_backoff: float | None = None,
|
||||
tool_choice: str | dict[str, Any] = "auto",
|
||||
cached_prefix: str | None = None,
|
||||
cached_prefix_message_count: int = 0,
|
||||
) -> "LLMToolCallResult":
|
||||
"""
|
||||
Make an LLM API call with tool/function calling support.
|
||||
@@ -1034,9 +1054,14 @@ class LLMProvider:
|
||||
await stack.enter_async_context(sem)
|
||||
|
||||
# cached_prefix is only set for providers that returned a handle
|
||||
# from get_or_create_cached_prefix(); forward it only when present
|
||||
# so non-caching providers keep their signature (same as call()).
|
||||
cache_kwarg = {"cached_prefix": cached_prefix} if cached_prefix is not None else {}
|
||||
# from get_or_create_cached_prefix() / create_incremental_cache();
|
||||
# forward it (plus how many leading messages it covers) only when
|
||||
# present so non-caching providers keep their signature.
|
||||
cache_kwarg = (
|
||||
{"cached_prefix": cached_prefix, "cached_prefix_message_count": cached_prefix_message_count}
|
||||
if cached_prefix is not None
|
||||
else {}
|
||||
)
|
||||
try:
|
||||
# Delegate to provider implementation
|
||||
result = await self._provider_impl.call_with_tools(
|
||||
@@ -1260,6 +1285,7 @@ class LLMProvider:
|
||||
ENV_LLM_GROQ_SERVICE_TIER,
|
||||
ENV_LLM_LITELLMROUTER_CONFIG,
|
||||
ENV_LLM_MODEL,
|
||||
ENV_LLM_OLLAMA_NUM_CTX,
|
||||
ENV_LLM_OPENAI_SERVICE_TIER,
|
||||
ENV_LLM_PROMPT_CACHE_ENABLED,
|
||||
ENV_LLM_PROVIDER,
|
||||
@@ -1270,6 +1296,7 @@ class LLMProvider:
|
||||
ENV_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY,
|
||||
_get_default_model_for_provider,
|
||||
_parse_llm_router_config,
|
||||
_parse_optional_positive_int,
|
||||
parse_gemini_service_tier,
|
||||
)
|
||||
|
||||
@@ -1314,6 +1341,7 @@ class LLMProvider:
|
||||
),
|
||||
gemini_safety_settings=json.loads(os.getenv(ENV_LLM_GEMINI_SAFETY_SETTINGS, "null")),
|
||||
prompt_cache_enabled=prompt_cache_enabled,
|
||||
ollama_num_ctx=_parse_optional_positive_int(ENV_LLM_OLLAMA_NUM_CTX, os.getenv(ENV_LLM_OLLAMA_NUM_CTX)),
|
||||
litellmrouter_config=_parse_llm_router_config(ENV_LLM_LITELLMROUTER_CONFIG),
|
||||
vertexai_project_id=os.getenv(ENV_LLM_VERTEXAI_PROJECT_ID) or None,
|
||||
vertexai_region=os.getenv(ENV_LLM_VERTEXAI_REGION) or None,
|
||||
|
||||
@@ -21,9 +21,10 @@ from one place, so we don't spawn a separate ``asyncio`` task per concern:
|
||||
The loop wakes on a short fixed tick and runs each job when its own
|
||||
``last_run + interval`` is due (run-at-start, then on interval), so adding jobs
|
||||
with different cadences doesn't burst CPU. Cross-tenant discovery goes through
|
||||
server-side PL/pgSQL routines (``public.schemas_with_expired_rows`` and
|
||||
``public.banks_needing_consolidation``) — one round-trip each — instead of a
|
||||
per-schema query storm, which matters at thousands of tenants.
|
||||
server-side PL/pgSQL routines (``schemas_with_expired_rows`` and
|
||||
``banks_needing_consolidation``, in the configured schema — see ``fq_routine``) —
|
||||
one round-trip each — instead of a per-schema query storm, which matters at
|
||||
thousands of tenants.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -32,13 +33,13 @@ import asyncio
|
||||
import logging
|
||||
import time
|
||||
from collections.abc import Coroutine
|
||||
from datetime import datetime, timezone
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from ..config import HindsightConfig, get_config
|
||||
from ..models import RequestContext
|
||||
from .db_utils import acquire_with_retry
|
||||
from .schema import _is_oracle, fq_table
|
||||
from .schema import _is_oracle, fq_routine, fq_table, fq_table_explicit
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .memory_engine import MemoryEngine
|
||||
@@ -49,6 +50,10 @@ logger = logging.getLogger(__name__)
|
||||
_TICK_SECONDS = 60
|
||||
# Retention sweeps are not time-sensitive; hourly matches the previous per-sweep cadence.
|
||||
_RETENTION_INTERVAL_SECONDS = 3600
|
||||
# Operation cleanup deletes one bounded batch per schema per run, so its cadence
|
||||
# sets the drain rate for a backlog. Kept at one-per-tick (the value it used while
|
||||
# it rode the worker's poll loop) so throughput is unchanged by the move.
|
||||
_OPERATION_CLEANUP_INTERVAL_SECONDS = 60
|
||||
|
||||
|
||||
class MaintenanceLoop:
|
||||
@@ -100,7 +105,8 @@ class MaintenanceLoop:
|
||||
audit_on = cfg.audit_log_enabled and cfg.audit_log_retention_days > 0
|
||||
llm_on = cfg.llm_trace_enabled and cfg.llm_trace_retention_days > 0
|
||||
mm_refresh_on = cfg.mental_model_refresh_tick_seconds > 0
|
||||
return reconcile_on or audit_on or llm_on or mm_refresh_on
|
||||
op_cleanup_on = cfg.operation_retention_days > 0
|
||||
return reconcile_on or audit_on or llm_on or mm_refresh_on or op_cleanup_on
|
||||
|
||||
# ── loop ───────────────────────────────────────────────────────────────
|
||||
|
||||
@@ -134,6 +140,8 @@ class MaintenanceLoop:
|
||||
mm_interval = cfg.mental_model_refresh_tick_seconds
|
||||
if mm_interval > 0 and self._is_due("mm_refresh", mm_interval):
|
||||
await self._run_timed("scheduled mental model refresh", self._run_scheduled_mm_refresh())
|
||||
if cfg.operation_retention_days > 0 and self._is_due("operation_cleanup", _OPERATION_CLEANUP_INTERVAL_SECONDS):
|
||||
await self._run_timed("operation cleanup", self._run_operation_cleanup(cfg))
|
||||
|
||||
async def _run_timed(self, name: str, coro: Coroutine[Any, Any, None]) -> None:
|
||||
"""Run a maintenance job and emit one timing line for it.
|
||||
@@ -163,7 +171,7 @@ class MaintenanceLoop:
|
||||
try:
|
||||
async with acquire_with_retry(backend, max_retries=1) as conn:
|
||||
rows = await conn.fetch(
|
||||
"SELECT * FROM public.schemas_with_expired_rows($1, $2, $3)", table, ts_col, days
|
||||
f"SELECT * FROM {fq_routine('schemas_with_expired_rows')}($1, $2, $3)", table, ts_col, days
|
||||
)
|
||||
for row in rows:
|
||||
schema = row[0]
|
||||
@@ -178,6 +186,73 @@ class MaintenanceLoop:
|
||||
except Exception as e:
|
||||
logger.warning(f"Retention sweep failed for {table}: {e}")
|
||||
|
||||
# ── terminal operation cleanup ─────────────────────────────────────────
|
||||
|
||||
async def _run_operation_cleanup(self, cfg: HindsightConfig) -> None:
|
||||
"""Prune one bounded batch of expired terminal operations per tenant schema.
|
||||
|
||||
Previously this rode the worker's task-claiming loop, so it only fired
|
||||
when that loop happened to iterate and was interleaved with claiming. It
|
||||
is a periodic housekeeping sweep like the retention jobs above, so it
|
||||
belongs on the same schedule.
|
||||
|
||||
Discovery is one cross-tenant round-trip (``schemas_with_expired_operations``)
|
||||
rather than a connection + prune transaction per tenant; pending and
|
||||
processing rows are never prunable, so a schema holding only in-flight
|
||||
work is correctly reported as having nothing to do.
|
||||
"""
|
||||
engine = self._engine
|
||||
backend = engine._backend
|
||||
try:
|
||||
async with acquire_with_retry(backend, max_retries=1) as conn:
|
||||
rows = await conn.fetch(
|
||||
f"SELECT * FROM {fq_routine('schemas_with_expired_operations')}($1)",
|
||||
cfg.operation_retention_days,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"Operation cleanup discovery failed: {e}")
|
||||
return
|
||||
if not rows:
|
||||
return
|
||||
|
||||
# Prune only schemas the deployment actually serves. The routine reports
|
||||
# every schema owning an async_operations table, including ones tenant
|
||||
# discovery doesn't claim.
|
||||
try:
|
||||
tenants = await engine._tenant_extension.list_tenants()
|
||||
except Exception as e:
|
||||
logger.warning(f"Operation cleanup tenant discovery failed: {e}")
|
||||
return
|
||||
known = {t.schema for t in tenants} | {get_config().database_schema}
|
||||
|
||||
from .memory_engine import _current_schema
|
||||
|
||||
cutoff = datetime.now(timezone.utc) - timedelta(days=cfg.operation_retention_days)
|
||||
pruned = 0
|
||||
for row in rows:
|
||||
schema = row[0]
|
||||
if schema not in known:
|
||||
continue
|
||||
# Oracle resolves unqualified names from a context-bound session
|
||||
# schema; on PostgreSQL this is harmless and fq_table stays explicit.
|
||||
token = _current_schema.set(schema)
|
||||
try:
|
||||
table = fq_table_explicit("async_operations", schema)
|
||||
async with acquire_with_retry(backend, max_retries=1) as conn:
|
||||
async with conn.transaction():
|
||||
deleted = await backend.ops.prune_terminal_operations(
|
||||
conn, table, cutoff, batch_size=cfg.operation_cleanup_batch_size
|
||||
)
|
||||
if deleted:
|
||||
pruned += deleted
|
||||
logger.info(f"Operation cleanup pruned {deleted} expired terminal operations from {schema}")
|
||||
except Exception as e:
|
||||
logger.warning(f"Operation cleanup failed for schema {schema}: {e}")
|
||||
finally:
|
||||
_current_schema.reset(token)
|
||||
if pruned:
|
||||
logger.info(f"Operation cleanup: pruned {pruned} operation(s) total")
|
||||
|
||||
# ── consolidation reconcile ──────────────────────────────────────────────
|
||||
|
||||
async def _run_reconcile(self) -> None:
|
||||
@@ -185,7 +260,9 @@ class MaintenanceLoop:
|
||||
engine = self._engine
|
||||
try:
|
||||
async with acquire_with_retry(engine._backend, max_retries=1) as conn:
|
||||
rows = await conn.fetch("SELECT schema_name, bank_id FROM public.banks_needing_consolidation()")
|
||||
rows = await conn.fetch(
|
||||
f"SELECT schema_name, bank_id FROM {fq_routine('banks_needing_consolidation')}()"
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"Consolidation reconcile discovery failed: {e}")
|
||||
return
|
||||
@@ -244,7 +321,7 @@ class MaintenanceLoop:
|
||||
|
||||
Discovery (the set of cron-scheduled models, minus any with an in-flight
|
||||
refresh) is one cross-tenant round-trip via
|
||||
``public.mental_models_with_cron()``. Cron *due-ness* is evaluated here in
|
||||
``mental_models_with_cron()``. Cron *due-ness* is evaluated here in
|
||||
Python — a scheduled fire has elapsed when the most recent cron boundary at
|
||||
or before now is later than ``last_refreshed_at`` — because cron arithmetic
|
||||
isn't expressible in plain SQL. Each due model is refreshed only when it is
|
||||
@@ -256,7 +333,7 @@ class MaintenanceLoop:
|
||||
async with acquire_with_retry(engine._backend, max_retries=1) as conn:
|
||||
rows = await conn.fetch(
|
||||
"SELECT schema_name, bank_id, mental_model_id, refresh_cron, last_refreshed_at "
|
||||
"FROM public.mental_models_with_cron()"
|
||||
f"FROM {fq_routine('mental_models_with_cron')}()"
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"Scheduled mental model refresh discovery failed: {e}")
|
||||
|
||||
@@ -60,6 +60,7 @@ from .llm_trace import (
|
||||
from .operation_metadata import (
|
||||
BatchRetainChildMetadata,
|
||||
BatchRetainParentMetadata,
|
||||
RefreshMentalModelOutcomeMetadata,
|
||||
RetainExtractionErrors,
|
||||
RetainOutcomeAggregate,
|
||||
RetainOutcomeMetadata,
|
||||
@@ -71,7 +72,7 @@ from .sql import SQLDialect, create_sql_dialect
|
||||
_current_schema: contextvars.ContextVar[str | None] = contextvars.ContextVar("current_schema", default=None)
|
||||
|
||||
# Context variable for the bank an operation runs for (async-safe, per-task isolation).
|
||||
# Set by the engine wherever it learns the bank (recall/retain/batch/task execution) so
|
||||
# Set by the engine wherever it learns the bank (recall/retain/batch/reflect/task execution) so
|
||||
# downstream provider calls can attribute spend per bank — e.g. tagging the OpenAI `user`
|
||||
# field for cost gateways. None outside a bank-scoped operation.
|
||||
_current_bank_id: contextvars.ContextVar[str | None] = contextvars.ContextVar("current_bank_id", default=None)
|
||||
@@ -438,6 +439,7 @@ def _member_to_llm(member: "LLMMemberConfig", config: HindsightConfig, defaults:
|
||||
reasoning_effort=member.reasoning_effort or config.llm_reasoning_effort,
|
||||
extra_body=member.extra_body,
|
||||
default_headers=member.default_headers or config.llm_default_headers,
|
||||
ollama_num_ctx=config.llm_ollama_num_ctx,
|
||||
bedrock_service_tier=member.bedrock_service_tier,
|
||||
gemini_service_tier=member.gemini_service_tier or config.llm_gemini_service_tier,
|
||||
gemini_safety_settings=_get_raw_config().llm_gemini_safety_settings,
|
||||
@@ -993,7 +995,10 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
# Initialize PostgreSQL connection URL
|
||||
# The actual URL will be set during initialize() after starting the server
|
||||
# Supports: "pg0" (default instance), "pg0://instance-name" (named instance), or regular postgresql:// URL
|
||||
self._use_pg0, self._pg0_instance_name, self._pg0_port = parse_pg0_url(db_url)
|
||||
_parsed_pg0 = parse_pg0_url(db_url)
|
||||
self._use_pg0 = _parsed_pg0.is_pg0
|
||||
self._pg0_instance_name = _parsed_pg0.instance_name
|
||||
self._pg0_port = _parsed_pg0.port
|
||||
if self._use_pg0:
|
||||
self.db_url = None
|
||||
else:
|
||||
@@ -1028,6 +1033,7 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
self._db_command_timeout = db_command_timeout if db_command_timeout is not None else config.db_command_timeout
|
||||
self._db_acquire_timeout = db_acquire_timeout if db_acquire_timeout is not None else config.db_acquire_timeout
|
||||
self._db_statement_timeout = config.db_statement_timeout
|
||||
self._db_max_parallel_workers_per_gather = config.db_max_parallel_workers_per_gather
|
||||
self._run_migrations = run_migrations
|
||||
self._retain_entity_lookup = config.retain_entity_lookup
|
||||
self._retain_entity_resolution_batch_size = config.retain_entity_resolution_batch_size
|
||||
@@ -1085,6 +1091,7 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
reasoning_effort=config.llm_reasoning_effort,
|
||||
extra_body=config.llm_extra_body,
|
||||
default_headers=config.llm_default_headers,
|
||||
ollama_num_ctx=config.llm_ollama_num_ctx,
|
||||
litellmrouter_config=config.llm_litellmrouter_config,
|
||||
bedrock_service_tier=config.llm_bedrock_service_tier,
|
||||
gemini_service_tier=config.llm_gemini_service_tier,
|
||||
@@ -1129,6 +1136,7 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
reasoning_effort=config.llm_reasoning_effort,
|
||||
extra_body=config.llm_extra_body,
|
||||
default_headers=config.llm_default_headers,
|
||||
ollama_num_ctx=config.llm_ollama_num_ctx,
|
||||
litellmrouter_config=config.retain_llm_litellmrouter_config or config.llm_litellmrouter_config,
|
||||
bedrock_service_tier=config.llm_bedrock_service_tier,
|
||||
gemini_service_tier=config.llm_gemini_service_tier,
|
||||
@@ -1167,6 +1175,7 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
reasoning_effort=config.llm_reasoning_effort,
|
||||
extra_body=config.llm_extra_body,
|
||||
default_headers=config.llm_default_headers,
|
||||
ollama_num_ctx=config.llm_ollama_num_ctx,
|
||||
litellmrouter_config=config.reflect_llm_litellmrouter_config or config.llm_litellmrouter_config,
|
||||
bedrock_service_tier=config.llm_bedrock_service_tier,
|
||||
gemini_service_tier=config.llm_gemini_service_tier,
|
||||
@@ -1205,6 +1214,7 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
reasoning_effort=config.llm_reasoning_effort,
|
||||
extra_body=config.llm_extra_body,
|
||||
default_headers=config.llm_default_headers,
|
||||
ollama_num_ctx=config.llm_ollama_num_ctx,
|
||||
litellmrouter_config=config.consolidation_llm_litellmrouter_config or config.llm_litellmrouter_config,
|
||||
bedrock_service_tier=config.llm_bedrock_service_tier,
|
||||
gemini_service_tier=config.llm_gemini_service_tier,
|
||||
@@ -1828,6 +1838,10 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
if refreshed is None:
|
||||
raise ValueError(f"Mental model {mental_model_id} not found in bank {bank_id}")
|
||||
|
||||
# Enrich the submit-time result_metadata with the semantic outcome
|
||||
# before the worker marks the operation completed (#2605).
|
||||
await self._write_refresh_outcome_metadata(task_dict.get("operation_id"), refreshed)
|
||||
|
||||
# Compute facts/mental_models counts for the post-op validator hook.
|
||||
# refresh_mental_model already persisted everything; the hook only needs
|
||||
# tallies that derive from the stored reflect_response payload.
|
||||
@@ -2369,25 +2383,70 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
|
||||
Also checks if this is a child operation and updates the parent if all siblings are done.
|
||||
Uses a single transaction to avoid race conditions when multiple children complete simultaneously.
|
||||
|
||||
Opt-in escape hatch: when ``HINDSIGHT_API_FAIL_ON_EXTRACTION_ERRORS`` is set and the
|
||||
operation's ``result_metadata`` recorded a non-zero ``extraction_errors_count`` (written by
|
||||
``_write_retain_outcome_metadata`` before this call), the operation is marked ``failed``
|
||||
instead of ``completed``. This surfaces silently-dropped facts as a hard failure rather
|
||||
than a clean success. Default is off, so existing behavior is unchanged (see issue #2700).
|
||||
"""
|
||||
try:
|
||||
backend = await self._get_backend()
|
||||
async with acquire_with_retry(backend) as conn:
|
||||
async with conn.transaction():
|
||||
# Mark this operation as completed
|
||||
# Read the accumulated extraction-error count (persisted by
|
||||
# _write_retain_outcome_metadata) to decide the terminal status.
|
||||
meta_row = await conn.fetchrow(
|
||||
f"SELECT result_metadata FROM {fq_table('async_operations')} WHERE operation_id = $1",
|
||||
uuid.UUID(operation_id),
|
||||
)
|
||||
extraction_errors_count = 0
|
||||
if meta_row is not None:
|
||||
metadata = conn.parse_json(meta_row["result_metadata"]) or {}
|
||||
extraction_errors_count = int(metadata.get("extraction_errors_count") or 0)
|
||||
|
||||
fail_on_errors = get_config().fail_on_extraction_errors
|
||||
if fail_on_errors and extraction_errors_count > 0:
|
||||
error_message = (
|
||||
f"Retain completed with {extraction_errors_count} fact extraction error(s); "
|
||||
"marked failed because HINDSIGHT_API_FAIL_ON_EXTRACTION_ERRORS is enabled. "
|
||||
"See result_metadata.extraction_errors_sample for details."
|
||||
)
|
||||
row = await conn.fetchrow(
|
||||
f"""
|
||||
UPDATE {fq_table("async_operations")}
|
||||
SET status = 'failed', error_message = $2, updated_at = NOW(), completed_at = NOW()
|
||||
WHERE operation_id = $1 AND status NOT IN ('completed', 'failed', 'cancelled')
|
||||
RETURNING operation_id
|
||||
""",
|
||||
uuid.UUID(operation_id),
|
||||
error_message,
|
||||
)
|
||||
if row is None:
|
||||
logger.info(f"Operation {operation_id} already terminal or deleted, skipping mark-failed")
|
||||
return
|
||||
logger.warning(
|
||||
f"Marked async operation as failed due to {extraction_errors_count} "
|
||||
f"extraction error(s): {operation_id}"
|
||||
)
|
||||
await self._maybe_update_parent_operation(operation_id, conn)
|
||||
return
|
||||
|
||||
# Mark this operation as completed. Guarded so an already-terminal
|
||||
# row is never re-terminalized: this keeps the engine idempotent
|
||||
# with the worker poller's completion backstop (PR #2608) and never
|
||||
# re-runs parent aggregation on a row that is already done.
|
||||
row = await conn.fetchrow(
|
||||
f"""
|
||||
UPDATE {fq_table("async_operations")}
|
||||
SET status = 'completed', updated_at = NOW(), completed_at = NOW()
|
||||
WHERE operation_id = $1
|
||||
WHERE operation_id = $1 AND status NOT IN ('completed', 'failed', 'cancelled')
|
||||
RETURNING operation_id
|
||||
""",
|
||||
uuid.UUID(operation_id),
|
||||
)
|
||||
if row is None:
|
||||
logger.info(
|
||||
f"Operation {operation_id} no longer exists (bank deleted), skipping mark-completed"
|
||||
)
|
||||
logger.info(f"Operation {operation_id} already terminal or deleted, skipping mark-completed")
|
||||
return
|
||||
logger.info(f"Marked async operation as completed: {operation_id}")
|
||||
|
||||
@@ -2438,6 +2497,47 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
# write silently regresses them to the ambiguous pre-fix behaviour.
|
||||
logger.warning(f"Failed to write retain outcome metadata for {operation_id}: {e}")
|
||||
|
||||
async def _write_refresh_outcome_metadata(self, operation_id: str | None, refreshed: dict[str, Any]) -> None:
|
||||
"""Persist completed refresh outcome fields before the operation is marked completed.
|
||||
|
||||
Refresh parity with ``_write_retain_outcome_metadata`` (#2605): merges the
|
||||
outcome into the submit-time ``{mental_model_id, name}`` metadata rather
|
||||
than replacing it, so consumers joining on those keys keep working.
|
||||
"""
|
||||
if not operation_id:
|
||||
return
|
||||
|
||||
from .reflect.agent import NO_ANSWER_TEXT
|
||||
|
||||
content = refreshed.get("content") or ""
|
||||
stripped = content.strip()
|
||||
based_on = (refreshed.get("reflect_response") or {}).get("based_on") or {}
|
||||
outcome = RefreshMentalModelOutcomeMetadata(
|
||||
content_len=len(content),
|
||||
# The no-answer stub and the pending placeholder complete
|
||||
# wire-successful but carry no real synthesis — a length check
|
||||
# alone would read them as populated.
|
||||
populated_content=bool(stripped) and stripped not in (MENTAL_MODEL_PENDING_CONTENT, NO_ANSWER_TEXT),
|
||||
based_on_counts={fact_type: len(facts or []) for fact_type, facts in based_on.items()},
|
||||
)
|
||||
try:
|
||||
backend = await self._get_backend()
|
||||
async with acquire_with_retry(backend) as conn:
|
||||
await conn.execute(
|
||||
f"""
|
||||
UPDATE {fq_table("async_operations")}
|
||||
SET result_metadata = COALESCE(result_metadata, '{{}}'::jsonb) || $2::jsonb,
|
||||
updated_at = now()
|
||||
WHERE operation_id = $1
|
||||
""",
|
||||
uuid.UUID(operation_id),
|
||||
json.dumps(outcome.to_dict()),
|
||||
)
|
||||
except Exception as e:
|
||||
# Best-effort, but log loudly: a missing write regresses clients to
|
||||
# fetch-and-measure health checks (the pre-#2605 behaviour).
|
||||
logger.warning(f"Failed to write refresh outcome metadata for {operation_id}: {e}")
|
||||
|
||||
async def _mark_operation_completed_and_fire_webhook(
|
||||
self,
|
||||
operation_id: str,
|
||||
@@ -2447,11 +2547,23 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
schema: str | None = None,
|
||||
error_message: str | None = None,
|
||||
) -> None:
|
||||
"""Mark an operation as completed and queue webhook deliveries in a single transaction.
|
||||
"""Mark an operation as completed and queue its consolidation webhook.
|
||||
|
||||
Uses the transactional outbox pattern: the webhook delivery row is inserted in the
|
||||
same database transaction as the status update. This guarantees at-least-once delivery
|
||||
even if the process crashes immediately after committing.
|
||||
Happy path uses the transactional outbox pattern: the webhook delivery row is
|
||||
inserted in the *same* transaction as the ``status = 'completed'`` update, which
|
||||
guarantees at-least-once delivery even if the process crashes right after commit.
|
||||
|
||||
The critical property is that a failure in the best-effort side-effects (webhook
|
||||
outbox insert, parent aggregation) must never roll back the completion with it.
|
||||
The original code wrapped everything in one transaction and swallowed the
|
||||
exception, so any hiccup left the operation stuck in ``processing`` forever while
|
||||
the log already said the work was done (issue #2601). If the combined transaction
|
||||
fails we therefore fall back to committing the completion on its own and fire the
|
||||
webhook best-effort (non-transactional) instead of dropping both.
|
||||
|
||||
The UPDATE only fires on a non-terminal row, so it is idempotent with the worker
|
||||
poller's completion backstop (PR #2608): whichever path runs second sees an
|
||||
already-terminal row, updates nothing, and does not re-run parent aggregation.
|
||||
"""
|
||||
from ..webhooks.models import ConsolidationEventData, WebhookEvent, WebhookEventType
|
||||
|
||||
@@ -2463,15 +2575,13 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
f"""
|
||||
UPDATE {fq_table("async_operations")}
|
||||
SET status = 'completed', updated_at = NOW(), completed_at = NOW()
|
||||
WHERE operation_id = $1
|
||||
WHERE operation_id = $1 AND status NOT IN ('completed', 'failed', 'cancelled')
|
||||
RETURNING operation_id
|
||||
""",
|
||||
uuid.UUID(operation_id),
|
||||
)
|
||||
if row is None:
|
||||
logger.info(
|
||||
f"Operation {operation_id} no longer exists (bank deleted), skipping mark-completed"
|
||||
)
|
||||
logger.info(f"Operation {operation_id} already terminal or deleted, skipping mark-completed")
|
||||
return
|
||||
logger.info(f"Marked async operation as completed: {operation_id}")
|
||||
await self._maybe_update_parent_operation(operation_id, conn)
|
||||
@@ -2493,8 +2603,51 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
data=data,
|
||||
)
|
||||
await self._webhook_manager.fire_event_with_conn(event, conn, schema=schema)
|
||||
return
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to mark operation completed and fire webhook {operation_id}: {e}")
|
||||
logger.error(
|
||||
f"Atomic complete+webhook failed for {operation_id}: {e}. "
|
||||
"Falling back to a completion-only commit so the operation is not left unfinished."
|
||||
)
|
||||
|
||||
# Fallback: the combined transaction above rolled back (atomically), so the row is
|
||||
# still non-terminal. Commit the terminal state on its own, then deliver the webhook
|
||||
# best-effort. Losing at-least-once atomicity for a single notification is far better
|
||||
# than leaving the operation stuck. We only re-fire the webhook when this fallback
|
||||
# actually transitioned the row: if the row is already terminal the happy-path
|
||||
# transaction had already committed (status + outbox together), so re-firing would
|
||||
# duplicate the delivery.
|
||||
completed_in_fallback = False
|
||||
try:
|
||||
backend = await self._get_backend()
|
||||
async with acquire_with_retry(backend) as conn:
|
||||
async with conn.transaction():
|
||||
row = await conn.fetchrow(
|
||||
f"""
|
||||
UPDATE {fq_table("async_operations")}
|
||||
SET status = 'completed', updated_at = NOW(), completed_at = NOW()
|
||||
WHERE operation_id = $1 AND status NOT IN ('completed', 'failed', 'cancelled')
|
||||
RETURNING operation_id
|
||||
""",
|
||||
uuid.UUID(operation_id),
|
||||
)
|
||||
if row is not None:
|
||||
completed_in_fallback = True
|
||||
await self._maybe_update_parent_operation(operation_id, conn)
|
||||
except Exception as e:
|
||||
# Last-resort: the worker poller's post-executor backstop (PR #2608) still
|
||||
# marks the row completed after this returns.
|
||||
logger.error(f"Fallback completion commit failed for {operation_id}: {e}")
|
||||
|
||||
if completed_in_fallback:
|
||||
await self._fire_consolidation_webhook(
|
||||
bank_id=bank_id,
|
||||
operation_id=operation_id,
|
||||
status=status,
|
||||
result=result,
|
||||
error_message=error_message,
|
||||
schema=schema,
|
||||
)
|
||||
|
||||
async def _maybe_update_parent_operation(self, child_operation_id: str, conn):
|
||||
"""Check if this is a child operation and update parent status if all siblings are done.
|
||||
@@ -2860,6 +3013,7 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
self._dialect = create_sql_dialect(self._database_backend_type)
|
||||
|
||||
stmt_timeout_s = self._db_statement_timeout
|
||||
max_parallel_gather = self._db_max_parallel_workers_per_gather
|
||||
text_search_extension = get_config().text_search_extension
|
||||
|
||||
# Per-connection initialization callback (PostgreSQL-specific for now)
|
||||
@@ -2897,6 +3051,18 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
if stmt_timeout_s > 0:
|
||||
await conn.execute(f"SET statement_timeout = '{stmt_timeout_s}s'")
|
||||
|
||||
# Optional cap on planner parallelism for this process's
|
||||
# connections. Deployments that run background workers against a
|
||||
# database shared with latency-sensitive traffic can set this to 0
|
||||
# on the worker process: bulk maintenance queries (consolidation,
|
||||
# graph upkeep) then run serially instead of fanning out across
|
||||
# parallel workers — parallelism buys latency, which background
|
||||
# work doesn't need, at the cost of concurrent CPU footprint,
|
||||
# which shared primaries do care about. None (default) leaves the
|
||||
# server setting untouched.
|
||||
if max_parallel_gather is not None:
|
||||
await conn.execute(f"SET max_parallel_workers_per_gather = {max_parallel_gather}")
|
||||
|
||||
await self._backend.initialize(
|
||||
self.db_url,
|
||||
min_size=self._pool_min_size,
|
||||
@@ -3337,6 +3503,8 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
if result and result.contents is not None:
|
||||
contents = cast(list[RetainContentDict], result.contents)
|
||||
|
||||
await self._ensure_bank_exists(bank_id, request_context)
|
||||
|
||||
# Engine-owned copy: the orchestrator clears per-item "content" strings
|
||||
# after building the document's combined text (memory pressure
|
||||
# optimization, see retain/orchestrator.py). Without an internal copy
|
||||
@@ -3823,6 +3991,14 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
# target bank's config before the restore.
|
||||
parsed = parse_bank_archive(archive_bytes)
|
||||
bank_id = target_bank_id or parsed.manifest.source_bank_id
|
||||
if self._operation_validator and await bank_utils.get_bank_profile_if_exists(backend, bank_id) is None:
|
||||
from hindsight_api.extensions import CreateBankContext
|
||||
|
||||
ctx = CreateBankContext(
|
||||
bank_id=bank_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
await self._validate_operation(self._operation_validator.validate_create_bank(ctx))
|
||||
resolved_config = await self._config_resolver.resolve_full_config(bank_id, request_context)
|
||||
return await import_bank(
|
||||
backend=backend,
|
||||
@@ -6570,13 +6746,18 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
)
|
||||
|
||||
collist = await self._memory_unit_columns(conn)
|
||||
# The archive is cold storage, never a recall surface, so the schema gives it
|
||||
# no `embedding` column at all (dropped in d4f6a8c2e1b3). The move in/out is
|
||||
# therefore over every memory_units column EXCEPT embedding; on revert the
|
||||
# embedding is recomputed from the unit's text/dates/entities below. This makes
|
||||
# a model switch (which re-dimensions memory_units) structurally unable to trip
|
||||
# a vector-dimension mismatch on the INSERT … SELECT round-trip (#2209).
|
||||
arch_cols = ", ".join(c for c in (s.strip() for s in collist.split(",")) if c != '"embedding"')
|
||||
# The archive is cold storage, never a recall surface and carries no index,
|
||||
# so the schema gives it neither the `embedding` (dropped in d4f6a8c2e1b3)
|
||||
# nor the `search_vector` column (dropped in e7c3a9f1b2d5). Both are
|
||||
# recall-surface columns whose type/shape follows server
|
||||
# config, so the move in/out is over every memory_units column EXCEPT those
|
||||
# two; on revert each is recomputed from the unit's text/dates/entities below.
|
||||
# This makes a model switch (which re-dimensions memory_units) structurally
|
||||
# unable to trip a vector-dimension mismatch (#2209), and a text-search backend
|
||||
# switch unable to trip a search_vector type mismatch (#2503), on the
|
||||
# INSERT … SELECT round-trip.
|
||||
_archive_omitted = ('"embedding"', '"search_vector"')
|
||||
arch_cols = ", ".join(c for c in (s.strip() for s in collist.split(",")) if c not in _archive_omitted)
|
||||
|
||||
# --- Edit fields (live rows only): text / context / dates / fact_type / entities ---
|
||||
doing_edit = any(
|
||||
@@ -6634,6 +6815,17 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
mentioned_at=live["mentioned_at"],
|
||||
entities=[r["canonical_name"] for r in ent_rows],
|
||||
)
|
||||
# Keep the stored text-search vector in sync with curated
|
||||
# text/context edits. Use the incoming parameters here:
|
||||
# PostgreSQL evaluates UPDATE RHS expressions before the
|
||||
# sibling SET assignments take effect, so column references
|
||||
# would see the pre-edit text/context.
|
||||
from .db.ops_postgresql import pg_search_vector_expr
|
||||
|
||||
sv_expr = pg_search_vector_expr(get_config(), text_col="$3", context_col="$4")
|
||||
search_vector_clause = (
|
||||
f",\n search_vector = {sv_expr}" if sv_expr else ""
|
||||
)
|
||||
await enqueue_relink_victims(conn, bank_id, [memory_id], ops=backend.ops)
|
||||
await conn.execute(
|
||||
f"""
|
||||
@@ -6641,7 +6833,7 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
SET text = $3, context = $4, fact_type = $5, occurred_start = $6,
|
||||
occurred_end = $7, event_date = $8, embedding = $9::vector,
|
||||
consolidated_at = NULL, consolidation_failed_at = NULL,
|
||||
edited_at = now(), updated_at = now()
|
||||
edited_at = now(), updated_at = now(){search_vector_clause}
|
||||
WHERE id = $1 AND bank_id = $2
|
||||
""",
|
||||
str(memory_uuid),
|
||||
@@ -6695,14 +6887,29 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
arch_row = await conn.fetchrow(
|
||||
f"SELECT entity_ids FROM {arch} WHERE id = $1 AND bank_id = $2", str(memory_uuid), bank_id
|
||||
)
|
||||
# The archive has no embedding column (see arch_cols above), so the live
|
||||
# row's embedding defaults to NULL on the way back and is recomputed below
|
||||
# once entities are restored.
|
||||
# The archive keeps neither embedding nor search_vector (see arch_cols
|
||||
# above), so both default to NULL on the way back and are recomputed here:
|
||||
# the embedding below once entities are restored, the search_vector now
|
||||
# from the row's own text/context/text_signals.
|
||||
await conn.execute(
|
||||
f"INSERT INTO {mu} ({arch_cols}) SELECT {arch_cols} FROM {arch} WHERE id = $1 AND bank_id = $2",
|
||||
str(memory_uuid),
|
||||
bank_id,
|
||||
)
|
||||
# Rebuild search_vector using the *current* text-search backend, so the
|
||||
# reverted unit is keyword-searchable again (more correct than carrying a
|
||||
# verbatim copy that could be stale/wrong-type if the backend changed while
|
||||
# the fact sat archived). None = pgroonga/pg_textsearch/pg_search, which
|
||||
# index base columns directly and leave search_vector empty (#2503).
|
||||
from .db.ops_postgresql import pg_search_vector_expr
|
||||
|
||||
sv_expr = pg_search_vector_expr(get_config())
|
||||
if sv_expr is not None:
|
||||
await conn.execute(
|
||||
f"UPDATE {mu} SET search_vector = {sv_expr} WHERE id = $1 AND bank_id = $2",
|
||||
str(memory_uuid),
|
||||
bank_id,
|
||||
)
|
||||
# Re-consolidate from scratch; links are rebuilt by graph maintenance.
|
||||
await conn.execute(
|
||||
f"UPDATE {mu} SET consolidated_at = NULL, consolidation_failed_at = NULL, updated_at = now() "
|
||||
@@ -7446,7 +7653,7 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
f"""
|
||||
SELECT id, text, event_date, context, fact_type, document_id,
|
||||
mentioned_at, occurred_start, occurred_end, chunk_id, proof_count,
|
||||
tags, consolidated_at, consolidation_failed_at, edited_at, {curation_cols}
|
||||
tags, metadata, consolidated_at, consolidation_failed_at, edited_at, {curation_cols}
|
||||
FROM {source_table}
|
||||
{where_clause}
|
||||
ORDER BY mentioned_at DESC NULLS LAST, created_at DESC
|
||||
@@ -7501,6 +7708,7 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
"chunk_id": row["chunk_id"] if row["chunk_id"] else None,
|
||||
"proof_count": row["proof_count"] if row["proof_count"] is not None else 1,
|
||||
"tags": list(row["tags"]) if row["tags"] else [],
|
||||
"metadata": conn.parse_json(row["metadata"]) if row["metadata"] is not None else {},
|
||||
"consolidated_at": row["consolidated_at"].isoformat() if row["consolidated_at"] else None,
|
||||
"consolidation_failed_at": (
|
||||
row["consolidation_failed_at"].isoformat() if row["consolidation_failed_at"] else None
|
||||
@@ -7553,7 +7761,7 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
# back to the archive (with its invalidation bookkeeping) on a miss.
|
||||
select_cols = (
|
||||
"id, text, context, event_date, occurred_start, occurred_end, "
|
||||
"mentioned_at, fact_type, document_id, chunk_id, tags, source_memory_ids, "
|
||||
"mentioned_at, fact_type, document_id, chunk_id, tags, metadata, source_memory_ids, "
|
||||
"observation_scopes, edited_at"
|
||||
)
|
||||
row = await conn.fetchrow(
|
||||
@@ -7597,7 +7805,10 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
"document_id": row["document_id"] if row["document_id"] else None,
|
||||
"chunk_id": str(row["chunk_id"]) if row["chunk_id"] else None,
|
||||
"tags": row["tags"] if row["tags"] else [],
|
||||
"observation_scopes": row["observation_scopes"] if row["observation_scopes"] else None,
|
||||
"metadata": conn.parse_json(row["metadata"]) if row["metadata"] is not None else {},
|
||||
"observation_scopes": (
|
||||
conn.parse_json(row["observation_scopes"]) if row["observation_scopes"] is not None else None
|
||||
),
|
||||
"state": unit_state,
|
||||
"invalidation_reason": row["invalidation_reason"],
|
||||
"invalidated_at": row["invalidated_at"].isoformat() if row["invalidated_at"] else None,
|
||||
@@ -8603,16 +8814,12 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
existing = await bank_utils.get_bank_profile_if_exists(backend, bank_id)
|
||||
if existing is None:
|
||||
return None
|
||||
profile, created = existing, False
|
||||
profile = existing
|
||||
else:
|
||||
result = await bank_utils.get_or_create_bank_profile(backend, bank_id)
|
||||
profile, created = result.profile, result.created
|
||||
|
||||
# Apply HINDSIGHT_API_DEFAULT_BANK_TEMPLATE to freshly-created banks. Done
|
||||
# before reading the resolved config below so the template's overrides
|
||||
# (e.g. reflect_mission, dispositions) are visible on this very call.
|
||||
if created:
|
||||
await self._apply_default_bank_template(bank_id, request_context)
|
||||
await self._ensure_bank_exists(bank_id, request_context)
|
||||
profile = await bank_utils.get_bank_profile_if_exists(backend, bank_id)
|
||||
if profile is None:
|
||||
raise RuntimeError(f"Bank '{bank_id}' was not found after ensuring it exists")
|
||||
|
||||
# reflect_mission and disposition in config take precedence over the legacy DB columns
|
||||
config_dict = await self._config_resolver.get_bank_config(bank_id, request_context)
|
||||
@@ -8668,6 +8875,20 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
True if the bank was freshly created on this call.
|
||||
"""
|
||||
backend = await self._get_backend()
|
||||
if self._operation_validator:
|
||||
if conn is not None:
|
||||
exists = await conn.fetchval(f"SELECT 1 FROM {fq_table('banks')} WHERE bank_id = $1", bank_id)
|
||||
else:
|
||||
exists = await bank_utils.get_bank_profile_if_exists(backend, bank_id)
|
||||
if not exists:
|
||||
from hindsight_api.extensions import CreateBankContext
|
||||
|
||||
ctx = CreateBankContext(
|
||||
bank_id=bank_id,
|
||||
request_context=request_context,
|
||||
)
|
||||
await self._validate_operation(self._operation_validator.validate_create_bank(ctx))
|
||||
|
||||
if conn is not None:
|
||||
result = await bank_utils.get_or_create_bank_profile_on_conn(conn, bank_id, ops=backend.ops)
|
||||
return result.created
|
||||
@@ -8858,6 +9079,7 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
|
||||
# ==================== Reflect Methods ====================
|
||||
|
||||
@_bind_bank_id()
|
||||
async def reflect_async(
|
||||
self,
|
||||
bank_id: str,
|
||||
@@ -10412,7 +10634,11 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
# outlives a mental-model insert that ultimately fails.
|
||||
async with acquire_with_retry(backend) as conn:
|
||||
async with conn.transaction():
|
||||
created = await self._ensure_bank_exists(bank_id, request_context, conn=conn)
|
||||
created = await self._ensure_bank_exists(
|
||||
bank_id,
|
||||
request_context,
|
||||
conn=conn,
|
||||
)
|
||||
if mental_model_id:
|
||||
row = await conn.fetchrow(
|
||||
f"""
|
||||
@@ -11119,335 +11345,6 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
|
||||
return result == "DELETE 1"
|
||||
|
||||
# =====================================================================
|
||||
# KNOWLEDGE BASE (folders + pages over mental models)
|
||||
# =====================================================================
|
||||
# The knowledge base is a tree of folders and pages stored in
|
||||
# ``knowledge_pages``. A page references the mental model holding its content
|
||||
# (``mental_model_id``); a folder is a container (``mental_model_id`` NULL).
|
||||
# Content lives in ``mental_models`` — this layer owns only tree structure.
|
||||
|
||||
# Default trigger for a knowledge page: a living document synthesized from the
|
||||
# bank's consolidated **observations** (not raw facts), refreshed incrementally
|
||||
# (delta) after each consolidation, and excluding other mental models so a page
|
||||
# never reflects on sibling pages. Applied when the client doesn't pass its own
|
||||
# ``trigger`` on create; a client can override any of these.
|
||||
KNOWLEDGE_PAGE_DEFAULT_TRIGGER = {
|
||||
"mode": "delta",
|
||||
"fact_types": ["observation"],
|
||||
"exclude_mental_models": True,
|
||||
"refresh_after_consolidation": True,
|
||||
}
|
||||
|
||||
# Knowledge pages default to a larger budget than a plain mental model (2048)
|
||||
# since they're meant to read as full documents. Applied when the client
|
||||
# doesn't pass ``max_tokens`` on create.
|
||||
KNOWLEDGE_PAGE_DEFAULT_MAX_TOKENS = 4096
|
||||
|
||||
@staticmethod
|
||||
def _row_to_knowledge_node(row) -> dict[str, Any]:
|
||||
"""Project a knowledge_pages row (optionally joined to its mental model)."""
|
||||
node: dict[str, Any] = {
|
||||
"id": row["id"],
|
||||
"bank_id": row["bank_id"],
|
||||
"parent_id": row["parent_id"],
|
||||
"kind": row["kind"],
|
||||
"name": row["name"],
|
||||
"mental_model_id": row["mental_model_id"],
|
||||
"sort_order": row["sort_order"],
|
||||
"managed": (row["managed"] if "managed" in row else False),
|
||||
"created_at": row["created_at"].isoformat() if row["created_at"] else None,
|
||||
"updated_at": row["updated_at"].isoformat() if row["updated_at"] else None,
|
||||
}
|
||||
# Page rows are returned LEFT JOINed to mental_models so the OKF
|
||||
# projection (type/tags/description) needs no second round-trip.
|
||||
if "mm_tags" in row:
|
||||
node["tags"] = list(row["mm_tags"] or [])
|
||||
node["source_query"] = row["mm_source_query"]
|
||||
node["last_refreshed_at"] = row["mm_last_refreshed_at"].isoformat() if row["mm_last_refreshed_at"] else None
|
||||
return node
|
||||
|
||||
# Column list for plain (non-joined) knowledge_pages reads/RETURNING.
|
||||
_KP_COLUMNS = "id, bank_id, parent_id, kind, name, mental_model_id, sort_order, managed, created_at, updated_at"
|
||||
|
||||
_KP_PAGE_SELECT = (
|
||||
"kp.id, kp.bank_id, kp.parent_id, kp.kind, kp.name, kp.mental_model_id, "
|
||||
"kp.sort_order, kp.managed, kp.created_at, kp.updated_at, "
|
||||
"mm.tags AS mm_tags, mm.source_query AS mm_source_query, "
|
||||
"mm.last_refreshed_at AS mm_last_refreshed_at"
|
||||
)
|
||||
|
||||
def _kp_join(self) -> str:
|
||||
kp = fq_table("knowledge_pages")
|
||||
mm = fq_table("mental_models")
|
||||
return f"{kp} kp LEFT JOIN {mm} mm ON mm.id = kp.mental_model_id AND mm.bank_id = kp.bank_id"
|
||||
|
||||
async def _kp_assert_folder_parent(self, conn, bank_id: str, parent_id: str | None) -> None:
|
||||
"""A non-null parent must be an existing folder in this bank."""
|
||||
if parent_id is None:
|
||||
return
|
||||
row = await conn.fetchrow(
|
||||
f"SELECT kind FROM {fq_table('knowledge_pages')} WHERE bank_id = $1 AND id = $2",
|
||||
bank_id,
|
||||
parent_id,
|
||||
)
|
||||
if row is None:
|
||||
raise ValueError(f"Parent folder '{parent_id}' not found")
|
||||
if row["kind"] != "folder":
|
||||
raise ValueError(f"Parent '{parent_id}' is not a folder")
|
||||
|
||||
async def create_knowledge_folder(
|
||||
self,
|
||||
bank_id: str,
|
||||
name: str,
|
||||
*,
|
||||
parent_id: str | None = None,
|
||||
managed: bool = False,
|
||||
request_context: "RequestContext",
|
||||
) -> dict[str, Any]:
|
||||
"""Create a folder (a container node) in the knowledge base.
|
||||
|
||||
The knowledge base is managed by clients (CRUD over folders/pages);
|
||||
``managed`` lets a client tag a node as system-owned vs. hand-authored.
|
||||
"""
|
||||
await self._authenticate_tenant(request_context)
|
||||
backend = await self._get_backend()
|
||||
folder_id = f"kf-{uuid.uuid4().hex}"
|
||||
async with acquire_with_retry(backend) as conn:
|
||||
async with conn.transaction():
|
||||
await self._ensure_bank_exists(bank_id, request_context, conn=conn)
|
||||
await self._kp_assert_folder_parent(conn, bank_id, parent_id)
|
||||
row = await conn.fetchrow(
|
||||
f"""
|
||||
INSERT INTO {fq_table("knowledge_pages")} (id, bank_id, parent_id, kind, name, managed)
|
||||
VALUES ($1, $2, $3, 'folder', $4, $5)
|
||||
RETURNING {self._KP_COLUMNS}
|
||||
""",
|
||||
folder_id,
|
||||
bank_id,
|
||||
parent_id,
|
||||
name,
|
||||
managed,
|
||||
)
|
||||
return self._row_to_knowledge_node(row)
|
||||
|
||||
async def create_knowledge_page(
|
||||
self,
|
||||
bank_id: str,
|
||||
name: str,
|
||||
source_query: str,
|
||||
content: str,
|
||||
*,
|
||||
parent_id: str | None = None,
|
||||
tags: list[str] | None = None,
|
||||
max_tokens: int | None = None,
|
||||
trigger: dict[str, Any] | None = None,
|
||||
mental_model_id: str | None = None,
|
||||
managed: bool = False,
|
||||
request_context: "RequestContext",
|
||||
) -> dict[str, Any] | None:
|
||||
"""Create a page: a backing mental model plus the tree node that refs it.
|
||||
|
||||
``managed`` lets a client tag the page as system-owned vs. hand-authored.
|
||||
When ``trigger`` is omitted the page uses ``KNOWLEDGE_PAGE_DEFAULT_TRIGGER``
|
||||
(observation-only, delta, auto-refresh) so a knowledge page is a living
|
||||
document by default.
|
||||
|
||||
Returns ``None`` when a page with the same name already exists in the same
|
||||
folder (a uniqueness violation) — the caller should treat that as
|
||||
"already exists" (surfaced by the API as a 409).
|
||||
"""
|
||||
await self._authenticate_tenant(request_context)
|
||||
# The mental model carries the content (and is created+validated by the
|
||||
# existing path, including lazy bank creation); the node only refs it.
|
||||
mm = await self.create_mental_model(
|
||||
bank_id=bank_id,
|
||||
name=name,
|
||||
source_query=source_query,
|
||||
content=content,
|
||||
mental_model_id=mental_model_id,
|
||||
tags=tags,
|
||||
max_tokens=max_tokens if max_tokens is not None else self.KNOWLEDGE_PAGE_DEFAULT_MAX_TOKENS,
|
||||
trigger=trigger if trigger is not None else dict(self.KNOWLEDGE_PAGE_DEFAULT_TRIGGER),
|
||||
request_context=request_context,
|
||||
)
|
||||
backend = await self._get_backend()
|
||||
page_id = f"kp-{uuid.uuid4().hex}"
|
||||
try:
|
||||
async with acquire_with_retry(backend) as conn:
|
||||
async with conn.transaction():
|
||||
await self._kp_assert_folder_parent(conn, bank_id, parent_id)
|
||||
row = await conn.fetchrow(
|
||||
f"""
|
||||
INSERT INTO {fq_table("knowledge_pages")}
|
||||
(id, bank_id, parent_id, kind, name, mental_model_id, managed)
|
||||
VALUES ($1, $2, $3, 'page', $4, $5, $6)
|
||||
RETURNING {self._KP_COLUMNS}
|
||||
""",
|
||||
page_id,
|
||||
bank_id,
|
||||
parent_id,
|
||||
name,
|
||||
mm["id"],
|
||||
managed,
|
||||
)
|
||||
except asyncpg.UniqueViolationError:
|
||||
# Duplicate page name in this folder (uq_kp_folder_pagename). Roll back
|
||||
# by deleting the orphan mental model we just created, then signal the
|
||||
# caller that the page already exists.
|
||||
await self.delete_mental_model(bank_id, mm["id"], request_context=request_context)
|
||||
return None
|
||||
node = self._row_to_knowledge_node(row)
|
||||
# Surface the mental-model metadata so the caller can render OKF or
|
||||
# schedule a content refresh without a second fetch.
|
||||
node["tags"] = list(mm.get("tags") or [])
|
||||
node["source_query"] = mm.get("source_query")
|
||||
node["last_refreshed_at"] = mm.get("last_refreshed_at")
|
||||
return node
|
||||
|
||||
async def list_knowledge_nodes(self, bank_id: str, *, request_context: "RequestContext") -> list[dict[str, Any]]:
|
||||
"""Return every folder/page node in the bank (flat; caller builds the tree)."""
|
||||
await self._authenticate_tenant(request_context)
|
||||
backend = await self._get_backend()
|
||||
async with acquire_with_retry(backend) as conn:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT {self._KP_PAGE_SELECT}
|
||||
FROM {self._kp_join()}
|
||||
WHERE kp.bank_id = $1
|
||||
ORDER BY kp.sort_order, kp.name
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
return [self._row_to_knowledge_node(r) for r in rows]
|
||||
|
||||
async def get_knowledge_page(
|
||||
self, bank_id: str, page_id: str, *, request_context: "RequestContext"
|
||||
) -> dict[str, Any] | None:
|
||||
"""Return a page node merged with its mental model's content (for OKF)."""
|
||||
await self._authenticate_tenant(request_context)
|
||||
backend = await self._get_backend()
|
||||
async with acquire_with_retry(backend) as conn:
|
||||
row = await conn.fetchrow(
|
||||
f"""
|
||||
SELECT {self._KP_PAGE_SELECT}, mm.content AS mm_content
|
||||
FROM {self._kp_join()}
|
||||
WHERE kp.bank_id = $1 AND kp.id = $2 AND kp.kind = 'page'
|
||||
""",
|
||||
bank_id,
|
||||
page_id,
|
||||
)
|
||||
if row is None:
|
||||
return None
|
||||
node = self._row_to_knowledge_node(row)
|
||||
node["content"] = row["mm_content"]
|
||||
return node
|
||||
|
||||
async def rename_knowledge_node(
|
||||
self, bank_id: str, node_id: str, name: str, *, request_context: "RequestContext"
|
||||
) -> dict[str, Any] | None:
|
||||
"""Rename a folder or page node."""
|
||||
await self._authenticate_tenant(request_context)
|
||||
backend = await self._get_backend()
|
||||
async with acquire_with_retry(backend) as conn:
|
||||
row = await conn.fetchrow(
|
||||
f"""
|
||||
UPDATE {fq_table("knowledge_pages")}
|
||||
SET name = $3, updated_at = now()
|
||||
WHERE bank_id = $1 AND id = $2
|
||||
RETURNING {self._KP_COLUMNS}
|
||||
""",
|
||||
bank_id,
|
||||
node_id,
|
||||
name,
|
||||
)
|
||||
return self._row_to_knowledge_node(row) if row else None
|
||||
|
||||
async def move_knowledge_node(
|
||||
self, bank_id: str, node_id: str, new_parent_id: str | None, *, request_context: "RequestContext"
|
||||
) -> dict[str, Any] | None:
|
||||
"""Re-parent a node, rejecting self-parenting and cycles."""
|
||||
await self._authenticate_tenant(request_context)
|
||||
if new_parent_id == node_id:
|
||||
raise ValueError("A node cannot be its own parent")
|
||||
backend = await self._get_backend()
|
||||
async with acquire_with_retry(backend) as conn:
|
||||
async with conn.transaction():
|
||||
await self._kp_assert_folder_parent(conn, bank_id, new_parent_id)
|
||||
# Cycle guard: walk up from the new parent; if we reach node_id,
|
||||
# the move would create a loop. Done in Python so the check stays
|
||||
# dialect-agnostic (no recursive CTE).
|
||||
if new_parent_id is not None:
|
||||
parents = {
|
||||
r["id"]: r["parent_id"]
|
||||
for r in await conn.fetch(
|
||||
f"SELECT id, parent_id FROM {fq_table('knowledge_pages')} WHERE bank_id = $1",
|
||||
bank_id,
|
||||
)
|
||||
}
|
||||
cursor: str | None = new_parent_id
|
||||
while cursor is not None:
|
||||
if cursor == node_id:
|
||||
raise ValueError("Cannot move a node into its own subtree")
|
||||
cursor = parents.get(cursor)
|
||||
row = await conn.fetchrow(
|
||||
f"""
|
||||
UPDATE {fq_table("knowledge_pages")}
|
||||
SET parent_id = $3, updated_at = now()
|
||||
WHERE bank_id = $1 AND id = $2
|
||||
RETURNING {self._KP_COLUMNS}
|
||||
""",
|
||||
bank_id,
|
||||
node_id,
|
||||
new_parent_id,
|
||||
)
|
||||
return self._row_to_knowledge_node(row) if row else None
|
||||
|
||||
async def delete_knowledge_node(self, bank_id: str, node_id: str, *, request_context: "RequestContext") -> bool:
|
||||
"""Delete a node and its whole subtree, including each page's mental model.
|
||||
|
||||
Deleting the mental models cascades their page rows away (FK ON DELETE
|
||||
CASCADE); deleting the node then cascades any remaining descendant folder
|
||||
rows. The subtree is gathered in Python so the logic is dialect-agnostic.
|
||||
"""
|
||||
await self._authenticate_tenant(request_context)
|
||||
backend = await self._get_backend()
|
||||
async with acquire_with_retry(backend) as conn:
|
||||
async with conn.transaction():
|
||||
all_rows = await conn.fetch(
|
||||
f"SELECT id, parent_id, mental_model_id FROM {fq_table('knowledge_pages')} WHERE bank_id = $1",
|
||||
bank_id,
|
||||
)
|
||||
by_parent: dict[str | None, list] = {}
|
||||
for r in all_rows:
|
||||
by_parent.setdefault(r["parent_id"], []).append(r)
|
||||
if not any(r["id"] == node_id for r in all_rows):
|
||||
return False
|
||||
# BFS the subtree rooted at node_id, collecting page mental models.
|
||||
stack = [node_id]
|
||||
mm_ids: list[str] = []
|
||||
while stack:
|
||||
current = stack.pop()
|
||||
for child in by_parent.get(current, []):
|
||||
stack.append(child["id"])
|
||||
node_row = next((r for r in all_rows if r["id"] == current), None)
|
||||
if node_row and node_row["mental_model_id"]:
|
||||
mm_ids.append(node_row["mental_model_id"])
|
||||
# Delete each backing mental model individually (the subtree is
|
||||
# small) to keep the SQL dialect-neutral — no PG array casts.
|
||||
for mm_id in mm_ids:
|
||||
await conn.execute(
|
||||
f"DELETE FROM {fq_table('mental_models')} WHERE bank_id = $1 AND id = $2",
|
||||
bank_id,
|
||||
mm_id,
|
||||
)
|
||||
await conn.execute(
|
||||
f"DELETE FROM {fq_table('knowledge_pages')} WHERE bank_id = $1 AND id = $2",
|
||||
bank_id,
|
||||
node_id,
|
||||
)
|
||||
return True
|
||||
|
||||
async def compute_mental_model_is_stale(
|
||||
self,
|
||||
conn,
|
||||
@@ -12230,22 +12127,11 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
op_uuid = uuid.UUID(operation_id)
|
||||
|
||||
async with acquire_with_retry(backend) as conn:
|
||||
row = await conn.fetchrow(
|
||||
f"SELECT bank_id, status FROM {fq_table('async_operations')} WHERE operation_id = $1 AND bank_id = $2",
|
||||
op_uuid,
|
||||
bank_id,
|
||||
)
|
||||
|
||||
if not row:
|
||||
raise ValueError(f"Operation {operation_id} not found for bank {bank_id}")
|
||||
|
||||
if row["status"] not in ("failed", "cancelled"):
|
||||
raise OperationValidationError(
|
||||
f"Operation {operation_id} cannot be retried: status is '{row['status']}', expected 'failed' or 'cancelled'",
|
||||
409,
|
||||
)
|
||||
|
||||
await conn.execute(
|
||||
# Make the retry transition a single conditional write. This
|
||||
# coordinates with retention cleanup's row locks: either retry wins
|
||||
# and the row becomes nonterminal, or pruning wins and this call
|
||||
# returns not-found instead of falsely acknowledging a vanished job.
|
||||
updated = await conn.fetchrow(
|
||||
f"""
|
||||
UPDATE {fq_table("async_operations")}
|
||||
SET status = 'pending',
|
||||
@@ -12257,10 +12143,27 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
retry_count = 0,
|
||||
updated_at = NOW()
|
||||
WHERE operation_id = $1
|
||||
AND bank_id = $2
|
||||
AND status IN ('failed', 'cancelled')
|
||||
RETURNING operation_id
|
||||
""",
|
||||
op_uuid,
|
||||
bank_id,
|
||||
)
|
||||
|
||||
if updated is None:
|
||||
row = await conn.fetchrow(
|
||||
f"SELECT status FROM {fq_table('async_operations')} WHERE operation_id = $1 AND bank_id = $2",
|
||||
op_uuid,
|
||||
bank_id,
|
||||
)
|
||||
if not row:
|
||||
raise ValueError(f"Operation {operation_id} not found for bank {bank_id}")
|
||||
raise OperationValidationError(
|
||||
f"Operation {operation_id} cannot be retried: status is '{row['status']}', expected 'failed' or 'cancelled'",
|
||||
409,
|
||||
)
|
||||
|
||||
return {
|
||||
"success": True,
|
||||
"message": f"Operation {operation_id} queued for retry",
|
||||
@@ -12362,7 +12265,11 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
# commit (or roll back) atomically.
|
||||
async with acquire_with_retry(backend) as conn:
|
||||
async with conn.transaction():
|
||||
created = await self._ensure_bank_exists(bank_id, request_context, conn=conn)
|
||||
created = await self._ensure_bank_exists(
|
||||
bank_id,
|
||||
request_context,
|
||||
conn=conn,
|
||||
)
|
||||
row = await backend.ops.create_webhook(
|
||||
conn,
|
||||
fq_table("webhooks"),
|
||||
@@ -12811,7 +12718,11 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
# async_operations.bank_id has a FK to banks. Create the bank
|
||||
# lazily inside this same transaction so it is atomic with the
|
||||
# parent + child operation rows.
|
||||
created = await self._ensure_bank_exists(bank_id, request_context, conn=conn)
|
||||
created = await self._ensure_bank_exists(
|
||||
bank_id,
|
||||
request_context,
|
||||
conn=conn,
|
||||
)
|
||||
await conn.execute(
|
||||
f"""
|
||||
INSERT INTO {fq_table("async_operations")} (operation_id, bank_id, operation_type, result_metadata, status)
|
||||
@@ -13176,4 +13087,3 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
result_metadata={"mental_model_id": mental_model_id, "name": mental_model["name"]},
|
||||
dedupe_by_bank=False,
|
||||
)
|
||||
|
||||
|
||||
@@ -142,3 +142,21 @@ class RefreshMentalModelMetadata:
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
"""Convert to dict for JSON serialization."""
|
||||
return asdict(self)
|
||||
|
||||
|
||||
@dataclass
|
||||
class RefreshMentalModelOutcomeMetadata:
|
||||
"""Machine-readable outcome metadata for a completed refresh_mental_model operation.
|
||||
|
||||
Refresh parity with RetainOutcomeMetadata (#2605): lets a monitoring layer
|
||||
distinguish "refreshed with real content" from "refreshed empty" by reading
|
||||
result_metadata alone, without a follow-up content fetch.
|
||||
"""
|
||||
|
||||
content_len: int
|
||||
populated_content: bool
|
||||
based_on_counts: dict[str, int] = field(default_factory=dict)
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
"""Convert to dict for JSON serialization."""
|
||||
return asdict(self)
|
||||
|
||||
@@ -186,7 +186,11 @@ class MarkitdownParser(FileParser):
|
||||
if Path(filename).suffix.lower() not in _TEXT_EXTENSIONS:
|
||||
return None
|
||||
try:
|
||||
file_data.decode("utf-8")
|
||||
# file_data may arrive as a non-``bytes`` buffer (e.g. a memoryview or
|
||||
# a native/Rust-backed buffer object) that has no ``.decode``; coerce
|
||||
# through the buffer protocol before the UTF-8 probe. The ``tmp.write``
|
||||
# in the caller already relies only on the same buffer protocol.
|
||||
bytes(file_data).decode("utf-8")
|
||||
except UnicodeDecodeError:
|
||||
return None
|
||||
from markitdown import StreamInfo
|
||||
|
||||
@@ -34,6 +34,43 @@ def _usage_from_anthropic_response(response: Any) -> LLMResponseUsage:
|
||||
)
|
||||
|
||||
|
||||
_EPHEMERAL_CACHE = {"type": "ephemeral"}
|
||||
|
||||
|
||||
def _cached_system_blocks(system_prompt: str) -> list[dict[str, Any]]:
|
||||
"""Render the system prompt as a block list with a cache_control marker.
|
||||
|
||||
Anthropic prompt caching is a prefix match: marking the (single) system
|
||||
block caches tools + system together. The system prompt is stable per
|
||||
scope — fact extraction reuses it across every chunk, reflect and
|
||||
consolidation keep their stable instructions there — so repeat calls read
|
||||
it at ~10% of the base input price. Markers below the model's minimum
|
||||
cacheable prefix are silently ignored (no write premium), so marking is
|
||||
safe unconditionally. This is the "inline-marker provider" strategy that
|
||||
``LLMInterface.get_or_create_cached_prefix`` documents for Anthropic.
|
||||
"""
|
||||
return [{"type": "text", "text": system_prompt, "cache_control": _EPHEMERAL_CACHE}]
|
||||
|
||||
|
||||
def _mark_last_message_for_caching(messages: list[dict[str, Any]]) -> None:
|
||||
"""Add a cache_control marker to the final content block, in place.
|
||||
|
||||
Used on the multi-turn (tool-calling) path: the reflect agent loop resends
|
||||
the entire growing conversation each iteration, so this request's
|
||||
end-marker becomes the next iteration's cache read point. Together with
|
||||
the system marker this uses 2 of the 4 allowed breakpoints.
|
||||
"""
|
||||
if not messages:
|
||||
return
|
||||
last = messages[-1]
|
||||
content = last.get("content")
|
||||
if isinstance(content, str):
|
||||
if content.strip(): # the API rejects empty text blocks
|
||||
last["content"] = [{"type": "text", "text": content, "cache_control": _EPHEMERAL_CACHE}]
|
||||
elif isinstance(content, list) and content and isinstance(content[-1], dict):
|
||||
content[-1]["cache_control"] = _EPHEMERAL_CACHE
|
||||
|
||||
|
||||
class AnthropicLLM(LLMInterface):
|
||||
"""
|
||||
LLM provider using Anthropic's Claude models.
|
||||
@@ -206,7 +243,9 @@ class AnthropicLLM(LLMInterface):
|
||||
}
|
||||
|
||||
if system_prompt:
|
||||
call_params["system"] = system_prompt
|
||||
# One-shot calls share only the system prompt with each other, so
|
||||
# that is the sole cache breakpoint on this path.
|
||||
call_params["system"] = _cached_system_blocks(system_prompt)
|
||||
|
||||
if use_forced_tool:
|
||||
# Single tool whose input_schema IS the response schema; force the model to
|
||||
@@ -450,6 +489,11 @@ class AnthropicLLM(LLMInterface):
|
||||
else:
|
||||
anthropic_messages.append({"role": role, "content": content})
|
||||
|
||||
# Multi-turn tool loop: cache the stable prefix (tools + system) via
|
||||
# the system marker, and the growing conversation via an end-marker
|
||||
# that the next iteration reads back.
|
||||
_mark_last_message_for_caching(anthropic_messages)
|
||||
|
||||
call_params: dict[str, Any] = {
|
||||
"model": self.model,
|
||||
"messages": anthropic_messages,
|
||||
@@ -457,7 +501,7 @@ class AnthropicLLM(LLMInterface):
|
||||
"max_tokens": max_completion_tokens or 4096,
|
||||
}
|
||||
if system_prompt:
|
||||
call_params["system"] = system_prompt
|
||||
call_params["system"] = _cached_system_blocks(system_prompt)
|
||||
|
||||
if self._extra_body:
|
||||
call_params["extra_body"] = self._extra_body
|
||||
@@ -543,6 +587,217 @@ class AnthropicLLM(LLMInterface):
|
||||
raise last_exception
|
||||
raise RuntimeError("Anthropic tool call failed")
|
||||
|
||||
# ── Message Batches API (50% token discount) ─────────────────────────────
|
||||
|
||||
_BATCH_TOOL_NAME = "structured_response"
|
||||
|
||||
async def supports_batch_api(self) -> bool:
|
||||
"""Anthropic supports batch operations via the Message Batches API."""
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
def _map_batch_status(processing_status: str) -> str:
|
||||
"""Map Anthropic ``processing_status`` onto the OpenAI vocabulary.
|
||||
|
||||
The engine's poll loop breaks on "completed" and hard-fails on
|
||||
"failed"/"expired"/"cancelled"; anything else keeps polling. Anthropic
|
||||
batches only end as "ended" (per-request failures surface in the
|
||||
results, mirroring OpenAI's "completed"-with-errors semantics), so
|
||||
"ended" maps to "completed" and the non-terminal states pass through.
|
||||
"""
|
||||
return "completed" if processing_status == "ended" else processing_status
|
||||
|
||||
def _translate_batch_body(self, body: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Translate one OpenAI-shaped request body into Messages API params.
|
||||
|
||||
Mirrors the conversion rules of ``call()``: system messages fold into
|
||||
the ``system`` param; ``max_completion_tokens`` becomes ``max_tokens``
|
||||
(default 4096); ``temperature`` is dropped (the sync path never sends
|
||||
it either — current Claude models reject non-default sampling params);
|
||||
an OpenAI ``response_format`` json_schema becomes a single forced
|
||||
tool_use tool when strict (native constrained decoding, issue #1002),
|
||||
else the schema is injected into the system prompt.
|
||||
|
||||
The system prompt carries the same cache_control marker as the sync
|
||||
one-shot path (its sole breakpoint): every request in a retain batch
|
||||
shares the fact-extraction system prompt, so the first item's cache
|
||||
write serves the remaining items as best-effort reads — and the
|
||||
cache-read discount stacks with the 50% batch discount.
|
||||
"""
|
||||
system_prompt: str | None = None
|
||||
messages: list[dict[str, Any]] = []
|
||||
for msg in body.get("messages", []):
|
||||
role = msg.get("role", "user")
|
||||
content = msg.get("content", "")
|
||||
if role == "system":
|
||||
system_prompt = (system_prompt + "\n\n" + content) if system_prompt else content
|
||||
else:
|
||||
messages.append({"role": role, "content": content})
|
||||
|
||||
params: dict[str, Any] = {
|
||||
"model": body.get("model") or self.model,
|
||||
"messages": messages,
|
||||
"max_tokens": body.get("max_completion_tokens") or 4096,
|
||||
}
|
||||
|
||||
json_schema = (body.get("response_format") or {}).get("json_schema") or {}
|
||||
schema = json_schema.get("schema")
|
||||
if schema is not None:
|
||||
if json_schema.get("strict"):
|
||||
params["tools"] = [
|
||||
{
|
||||
"name": self._BATCH_TOOL_NAME,
|
||||
"description": "Return the structured response.",
|
||||
"input_schema": schema,
|
||||
}
|
||||
]
|
||||
params["tool_choice"] = {"type": "tool", "name": self._BATCH_TOOL_NAME}
|
||||
else:
|
||||
schema_msg = "\n\nYou must respond with valid JSON matching this schema:\n" + json.dumps(
|
||||
schema, indent=2, ensure_ascii=False
|
||||
)
|
||||
system_prompt = (system_prompt + schema_msg) if system_prompt else schema_msg
|
||||
|
||||
if system_prompt:
|
||||
params["system"] = _cached_system_blocks(system_prompt)
|
||||
|
||||
# Batch params ARE the raw Messages body, so operator-configured extra
|
||||
# body params merge directly (the sync path routes them through the
|
||||
# SDK's extra_body, which does the same merge server-side).
|
||||
if self._extra_body:
|
||||
params.update(self._extra_body)
|
||||
|
||||
return params
|
||||
|
||||
def _translate_batch_message(self, message: Any) -> dict[str, Any]:
|
||||
"""Render an Anthropic Message as the OpenAI response body the engine parses.
|
||||
|
||||
The engine reads ``choices[0].message.content`` (json.loads'ing it when
|
||||
a schema was requested) and sums ``usage`` under the OpenAI key names.
|
||||
Forced-tool responses carry their JSON in the tool_use block's input,
|
||||
so that is re-serialized as the content string.
|
||||
"""
|
||||
content = ""
|
||||
tool_input = None
|
||||
for block in message.content:
|
||||
if block.type == "tool_use" and block.name == self._BATCH_TOOL_NAME:
|
||||
tool_input = block.input or {}
|
||||
elif block.type == "text":
|
||||
content += block.text
|
||||
if tool_input is not None:
|
||||
content = json.dumps(tool_input, ensure_ascii=False)
|
||||
|
||||
usage = getattr(message, "usage", None)
|
||||
input_tokens = (usage.input_tokens or 0) if usage else 0
|
||||
output_tokens = (usage.output_tokens or 0) if usage else 0
|
||||
|
||||
return {
|
||||
"choices": [
|
||||
{
|
||||
"message": {"role": "assistant", "content": content},
|
||||
"finish_reason": getattr(message, "stop_reason", None),
|
||||
}
|
||||
],
|
||||
"usage": {
|
||||
"prompt_tokens": input_tokens,
|
||||
"completion_tokens": output_tokens,
|
||||
"total_tokens": input_tokens + output_tokens,
|
||||
},
|
||||
}
|
||||
|
||||
async def submit_batch(
|
||||
self,
|
||||
requests: list[dict[str, Any]],
|
||||
endpoint: str = "/v1/chat/completions",
|
||||
completion_window: str = "24h",
|
||||
) -> dict[str, Any]:
|
||||
"""Submit a batch of requests to the Message Batches API.
|
||||
|
||||
Accepts the engine's OpenAI-JSONL-shaped entries. ``endpoint`` and
|
||||
``completion_window`` belong to that shared shape and have no Anthropic
|
||||
equivalent (batches always resolve within 24 hours); both are ignored.
|
||||
"""
|
||||
batch_requests = [
|
||||
{
|
||||
"custom_id": req["custom_id"],
|
||||
"params": self._translate_batch_body(req.get("body") or {}),
|
||||
}
|
||||
for req in requests
|
||||
]
|
||||
|
||||
logger.info(f"Submitting Anthropic message batch with {len(batch_requests)} requests")
|
||||
batch = await self._client.messages.batches.create(requests=batch_requests)
|
||||
logger.info(f"Anthropic batch submitted: {batch.id}, status={batch.processing_status}")
|
||||
|
||||
return {
|
||||
"batch_id": batch.id,
|
||||
"status": self._map_batch_status(batch.processing_status),
|
||||
"created_at": batch.created_at,
|
||||
"request_count": len(batch_requests),
|
||||
}
|
||||
|
||||
async def get_batch_status(self, batch_id: str) -> dict[str, Any]:
|
||||
"""Get batch status in the shape the engine's poll loop expects."""
|
||||
batch = await self._client.messages.batches.retrieve(batch_id)
|
||||
|
||||
counts = batch.request_counts
|
||||
processing = getattr(counts, "processing", 0) or 0
|
||||
succeeded = getattr(counts, "succeeded", 0) or 0
|
||||
errored = getattr(counts, "errored", 0) or 0
|
||||
canceled = getattr(counts, "canceled", 0) or 0
|
||||
expired = getattr(counts, "expired", 0) or 0
|
||||
resolved = succeeded + errored + canceled + expired
|
||||
|
||||
result: dict[str, Any] = {
|
||||
"batch_id": batch.id,
|
||||
"status": self._map_batch_status(batch.processing_status),
|
||||
"created_at": batch.created_at,
|
||||
"request_counts": {
|
||||
"total": processing + resolved,
|
||||
"completed": resolved,
|
||||
"failed": errored,
|
||||
},
|
||||
}
|
||||
|
||||
ended_at = getattr(batch, "ended_at", None)
|
||||
if ended_at:
|
||||
result["completed_at"] = ended_at
|
||||
|
||||
return result
|
||||
|
||||
async def retrieve_batch_results(self, batch_id: str) -> list[dict[str, Any]]:
|
||||
"""Retrieve completed batch results, translated to the OpenAI shape.
|
||||
|
||||
Succeeded entries become ``{"custom_id", "response": {"body": ...}}``;
|
||||
errored/canceled/expired entries become ``{"custom_id", "error": ...}``
|
||||
so the engine's per-result error handling applies unchanged.
|
||||
"""
|
||||
batch = await self._client.messages.batches.retrieve(batch_id)
|
||||
if batch.processing_status != "ended":
|
||||
raise ValueError(f"Batch {batch_id} is not completed yet (status: {batch.processing_status})")
|
||||
|
||||
decoder = await self._client.messages.batches.results(batch_id)
|
||||
results: list[dict[str, Any]] = []
|
||||
async for entry in decoder:
|
||||
outcome = entry.result
|
||||
if outcome.type == "succeeded":
|
||||
results.append(
|
||||
{
|
||||
"custom_id": entry.custom_id,
|
||||
"response": {"body": self._translate_batch_message(outcome.message)},
|
||||
}
|
||||
)
|
||||
else:
|
||||
error = getattr(outcome, "error", None)
|
||||
if error is not None:
|
||||
detail = f"{getattr(error, 'type', 'error')}: {getattr(error, 'message', error)}"
|
||||
else:
|
||||
detail = f"batch request {outcome.type}"
|
||||
results.append({"custom_id": entry.custom_id, "error": detail})
|
||||
|
||||
logger.info(f"Retrieved {len(results)} results for Anthropic batch {batch_id}")
|
||||
return results
|
||||
|
||||
async def cleanup(self) -> None:
|
||||
"""Clean up resources (close Anthropic client connections)."""
|
||||
if hasattr(self, "_client") and self._client:
|
||||
|
||||
@@ -49,6 +49,20 @@ def _get_isolated_claude_env() -> dict[str, str]:
|
||||
return _isolated_claude_env
|
||||
|
||||
|
||||
def _result_error_detail(message: Any) -> str:
|
||||
"""Build an actionable error string from an ``is_error`` ResultMessage.
|
||||
|
||||
The CLI can report a failure with ``is_error=True`` while ``subtype``
|
||||
still reads ``"success"``, putting the real detail in ``result`` (e.g.
|
||||
quota exhaustion: ``You've hit your weekly limit · resets ...`` with
|
||||
``api_error_status: 429``). The SDK's own fallback exception surfaces
|
||||
only the subtype, producing the misleading "Claude Code returned an
|
||||
error result: success" (issue #2702) — so prefer ``result``.
|
||||
"""
|
||||
detail = (message.result or "").strip() or message.subtype or "unknown error"
|
||||
return f"Claude Code reported an error: {detail}"
|
||||
|
||||
|
||||
class ClaudeCodeLLM(LLMInterface):
|
||||
"""
|
||||
LLM provider using Claude Code authentication.
|
||||
@@ -176,6 +190,7 @@ class ClaudeCodeLLM(LLMInterface):
|
||||
from claude_agent_sdk import ( # type: ignore[unresolved-import]
|
||||
AssistantMessage,
|
||||
ClaudeAgentOptions,
|
||||
ResultMessage,
|
||||
TextBlock,
|
||||
query,
|
||||
)
|
||||
@@ -228,6 +243,11 @@ class ClaudeCodeLLM(LLMInterface):
|
||||
for block in message.content:
|
||||
if isinstance(block, TextBlock):
|
||||
full_text += block.text
|
||||
elif isinstance(message, ResultMessage) and message.is_error:
|
||||
# Surface the CLI's actual error text (e.g. quota
|
||||
# exhaustion) instead of the SDK's subtype-based
|
||||
# fallback exception (issue #2702).
|
||||
raise RuntimeError(_result_error_detail(message))
|
||||
|
||||
# The Claude Agent SDK doesn't report exact counts; stash the same
|
||||
# char/4 estimate the success path traces so a later parse/validate
|
||||
@@ -393,6 +413,7 @@ class ClaudeCodeLLM(LLMInterface):
|
||||
AssistantMessage,
|
||||
ClaudeAgentOptions,
|
||||
ClaudeSDKClient,
|
||||
ResultMessage,
|
||||
SdkMcpTool,
|
||||
TextBlock,
|
||||
ToolUseBlock,
|
||||
@@ -532,6 +553,9 @@ class ClaudeCodeLLM(LLMInterface):
|
||||
|
||||
# Receive response
|
||||
async for message in client.receive_response():
|
||||
if isinstance(message, ResultMessage) and message.is_error:
|
||||
# Surface the CLI's actual error text (issue #2702).
|
||||
raise RuntimeError(_result_error_detail(message))
|
||||
if isinstance(message, AssistantMessage):
|
||||
for block in message.content:
|
||||
if isinstance(block, TextBlock):
|
||||
|
||||
@@ -19,6 +19,7 @@ from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import binascii
|
||||
import contextlib
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
@@ -31,6 +32,11 @@ from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
try:
|
||||
import fcntl
|
||||
except ImportError: # pragma: no cover - Windows
|
||||
fcntl = None # type: ignore[assignment]
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@@ -58,6 +64,9 @@ _CODEX_TOKEN_REFRESH_SKEW_SECONDS = 60
|
||||
_CODEX_TERMINAL_REFRESH_ERROR_CODES = frozenset(
|
||||
{"refresh_token_expired", "refresh_token_reused", "refresh_token_invalidated"}
|
||||
)
|
||||
_CODEX_AUTH_LOCK_TIMEOUT_SECONDS = 20.0
|
||||
_CODEX_AUTH_LOCKS_GUARD = threading.Lock()
|
||||
_CODEX_AUTH_LOCKS: dict[Path, threading.Lock] = {}
|
||||
|
||||
|
||||
def default_codex_auth_file() -> Path:
|
||||
@@ -76,6 +85,44 @@ def default_codex_auth_file() -> Path:
|
||||
return Path.home() / ".codex" / "auth.json"
|
||||
|
||||
|
||||
def _path_scoped_lock(auth_file: Path) -> threading.Lock:
|
||||
key = auth_file.expanduser().resolve(strict=False)
|
||||
with _CODEX_AUTH_LOCKS_GUARD:
|
||||
lock = _CODEX_AUTH_LOCKS.get(key)
|
||||
if lock is None:
|
||||
lock = threading.Lock()
|
||||
_CODEX_AUTH_LOCKS[key] = lock
|
||||
return lock
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _codex_auth_lock(auth_file: Path, timeout_seconds: float = _CODEX_AUTH_LOCK_TIMEOUT_SECONDS):
|
||||
"""Cross-process advisory lock for one Codex auth store."""
|
||||
with _path_scoped_lock(auth_file):
|
||||
if fcntl is None: # pragma: no cover - Windows
|
||||
logger.debug("fcntl unavailable; Codex refresh proceeds without a cross-process lock.")
|
||||
yield
|
||||
return
|
||||
|
||||
lock_path = auth_file.with_suffix(".lock")
|
||||
lock_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
with open(lock_path, "a+") as lock_file:
|
||||
deadline = time.monotonic() + max(1.0, timeout_seconds)
|
||||
while True:
|
||||
try:
|
||||
fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
|
||||
break
|
||||
except (BlockingIOError, OSError):
|
||||
if time.monotonic() >= deadline:
|
||||
raise TimeoutError("Timed out waiting for the Codex auth store lock") from None
|
||||
time.sleep(0.05)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
with contextlib.suppress(OSError):
|
||||
fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN)
|
||||
|
||||
|
||||
class CodexRefreshExpiredError(RuntimeError):
|
||||
"""Raised when the Codex refresh_token itself is no longer valid.
|
||||
|
||||
@@ -192,6 +239,34 @@ class CodexAuthManager:
|
||||
return None
|
||||
return data.get("tokens", {}).get("refresh_token")
|
||||
|
||||
@staticmethod
|
||||
def _load_tokens_from_file(auth_file: Path) -> dict[str, Any] | None:
|
||||
try:
|
||||
with open(auth_file) as f:
|
||||
data = json.load(f)
|
||||
except (OSError, json.JSONDecodeError):
|
||||
return None
|
||||
tokens = data.get("tokens")
|
||||
return tokens if isinstance(tokens, dict) else None
|
||||
|
||||
def _adopt_tokens(self, tokens: dict[str, Any]) -> bool:
|
||||
"""Adopt a newer on-disk Codex token set if present."""
|
||||
access_token = tokens.get("access_token")
|
||||
refresh_token = tokens.get("refresh_token")
|
||||
account_id = tokens.get("account_id")
|
||||
|
||||
changed = False
|
||||
if isinstance(access_token, str) and access_token and access_token != self.access_token:
|
||||
self.access_token = access_token
|
||||
changed = True
|
||||
if isinstance(refresh_token, str) and refresh_token and refresh_token != self.refresh_token:
|
||||
self.refresh_token = refresh_token
|
||||
changed = True
|
||||
if isinstance(account_id, str) and account_id and account_id != self.account_id:
|
||||
self.account_id = account_id
|
||||
changed = True
|
||||
return changed
|
||||
|
||||
@staticmethod
|
||||
def _decode_jwt_exp_unixtime(token: str) -> int | None:
|
||||
"""Return the JWT ``exp`` claim as a unix timestamp, or None on parse failure.
|
||||
@@ -225,6 +300,11 @@ class CodexAuthManager:
|
||||
return False
|
||||
return exp <= int(time.time()) + skew_seconds
|
||||
|
||||
def _token_is_fresh_with_known_expiry(self, skew_seconds: int = _CODEX_TOKEN_REFRESH_SKEW_SECONDS) -> bool:
|
||||
"""True only when the cached token has a known expiry outside the skew window."""
|
||||
exp = self._decode_jwt_exp_unixtime(self.access_token)
|
||||
return exp is not None and exp > int(time.time()) + skew_seconds
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Persistence
|
||||
# ------------------------------------------------------------------
|
||||
@@ -339,78 +419,93 @@ class CodexAuthManager:
|
||||
if not self._token_is_stale():
|
||||
return
|
||||
|
||||
if not self.refresh_token:
|
||||
raise RuntimeError(
|
||||
"Codex access_token is expired but no refresh_token is available. "
|
||||
"Run 'codex auth login' to re-authenticate."
|
||||
)
|
||||
with _codex_auth_lock(self._auth_file):
|
||||
disk_tokens = self._load_tokens_from_file(self._auth_file)
|
||||
if disk_tokens and self._adopt_tokens(disk_tokens):
|
||||
if force or self._token_is_fresh_with_known_expiry():
|
||||
return
|
||||
|
||||
log_reason = f" ({reason})" if reason else ""
|
||||
logger.info(f"Refreshing Codex OAuth access_token{log_reason}")
|
||||
|
||||
request_body = {
|
||||
"client_id": _CODEX_CLIENT_ID,
|
||||
"grant_type": "refresh_token",
|
||||
"refresh_token": self.refresh_token,
|
||||
}
|
||||
try:
|
||||
response = self._http_client.post(
|
||||
_CODEX_REFRESH_TOKEN_URL,
|
||||
json=request_body,
|
||||
headers={"Content-Type": "application/json"},
|
||||
timeout=30.0,
|
||||
)
|
||||
except httpx.RequestError as e:
|
||||
raise RuntimeError(f"Codex OAuth refresh network error: {type(e).__name__}") from e
|
||||
|
||||
if response.status_code == 401:
|
||||
error_code = self._extract_oauth_error_code(response)
|
||||
if error_code in _CODEX_TERMINAL_REFRESH_ERROR_CODES:
|
||||
raise CodexRefreshExpiredError(
|
||||
f"Codex refresh_token is permanently invalid (error.code={error_code}). "
|
||||
if not self.refresh_token:
|
||||
raise RuntimeError(
|
||||
"Codex access_token is expired but no refresh_token is available. "
|
||||
"Run 'codex auth login' to re-authenticate."
|
||||
)
|
||||
raise CodexRefreshExpiredError(
|
||||
f"Codex OAuth refresh returned 401 with unrecognized error code "
|
||||
f"({error_code or 'none'}). Run 'codex auth login' to re-authenticate."
|
||||
)
|
||||
|
||||
if response.status_code >= 400:
|
||||
raise RuntimeError(f"Codex OAuth refresh failed with HTTP {response.status_code}")
|
||||
log_reason = f" ({reason})" if reason else ""
|
||||
logger.info(f"Refreshing Codex OAuth access_token{log_reason}")
|
||||
|
||||
try:
|
||||
body = response.json()
|
||||
except json.JSONDecodeError as e:
|
||||
raise RuntimeError(f"Codex OAuth refresh returned non-JSON body: {e}") from e
|
||||
request_access_token = self.access_token
|
||||
request_refresh_token = self.refresh_token
|
||||
request_body = {
|
||||
"client_id": _CODEX_CLIENT_ID,
|
||||
"grant_type": "refresh_token",
|
||||
"refresh_token": request_refresh_token,
|
||||
}
|
||||
try:
|
||||
response = self._http_client.post(
|
||||
_CODEX_REFRESH_TOKEN_URL,
|
||||
json=request_body,
|
||||
headers={"Content-Type": "application/json"},
|
||||
timeout=30.0,
|
||||
)
|
||||
except httpx.RequestError as e:
|
||||
raise RuntimeError(f"Codex OAuth refresh network error: {type(e).__name__}") from e
|
||||
|
||||
new_access = body.get("access_token")
|
||||
if not new_access:
|
||||
raise RuntimeError("Codex OAuth refresh returned no access_token")
|
||||
if response.status_code == 401:
|
||||
error_code = self._extract_oauth_error_code(response)
|
||||
disk_tokens = self._load_tokens_from_file(self._auth_file)
|
||||
if disk_tokens and (
|
||||
disk_tokens.get("access_token") != request_access_token
|
||||
or disk_tokens.get("refresh_token") != request_refresh_token
|
||||
):
|
||||
self._adopt_tokens(disk_tokens)
|
||||
return
|
||||
if error_code in _CODEX_TERMINAL_REFRESH_ERROR_CODES:
|
||||
raise CodexRefreshExpiredError(
|
||||
f"Codex refresh_token is permanently invalid (error.code={error_code}). "
|
||||
"Run 'codex auth login' to re-authenticate."
|
||||
)
|
||||
raise CodexRefreshExpiredError(
|
||||
f"Codex OAuth refresh returned 401 with unrecognized error code "
|
||||
f"({error_code or 'none'}). Run 'codex auth login' to re-authenticate."
|
||||
)
|
||||
|
||||
new_refresh = body.get("refresh_token") or self.refresh_token
|
||||
new_id_token = body.get("id_token")
|
||||
if response.status_code >= 400:
|
||||
raise RuntimeError(f"Codex OAuth refresh failed with HTTP {response.status_code}")
|
||||
|
||||
# Update in-memory state first so waiters see fresh credentials
|
||||
# immediately, even if disk write fails.
|
||||
self.access_token = new_access
|
||||
self.refresh_token = new_refresh
|
||||
try:
|
||||
body = response.json()
|
||||
except json.JSONDecodeError as e:
|
||||
raise RuntimeError(f"Codex OAuth refresh returned non-JSON body: {e}") from e
|
||||
|
||||
persisted: dict[str, Any] = {
|
||||
"access_token": new_access,
|
||||
"refresh_token": new_refresh,
|
||||
}
|
||||
if new_id_token:
|
||||
persisted["id_token"] = new_id_token
|
||||
new_access = body.get("access_token")
|
||||
if not new_access:
|
||||
raise RuntimeError("Codex OAuth refresh returned no access_token")
|
||||
|
||||
try:
|
||||
self._persist_auth_atomic(persisted)
|
||||
except OSError as e:
|
||||
logger.warning(
|
||||
f"Codex OAuth refresh succeeded but persisting auth.json failed: {type(e).__name__}. "
|
||||
"In-memory credentials are up to date; on-disk file is stale."
|
||||
)
|
||||
new_refresh = body.get("refresh_token") or self.refresh_token
|
||||
new_id_token = body.get("id_token")
|
||||
|
||||
logger.info("Codex OAuth access_token refreshed successfully")
|
||||
# Update in-memory state first so waiters see fresh credentials
|
||||
# immediately, even if disk write fails.
|
||||
self.access_token = new_access
|
||||
self.refresh_token = new_refresh
|
||||
|
||||
persisted: dict[str, Any] = {
|
||||
"access_token": new_access,
|
||||
"refresh_token": new_refresh,
|
||||
}
|
||||
if new_id_token:
|
||||
persisted["id_token"] = new_id_token
|
||||
|
||||
try:
|
||||
self._persist_auth_atomic(persisted)
|
||||
except OSError as e:
|
||||
logger.warning(
|
||||
f"Codex OAuth refresh succeeded but persisting auth.json failed: {type(e).__name__}. "
|
||||
"In-memory credentials are up to date; on-disk file is stale."
|
||||
)
|
||||
|
||||
logger.info("Codex OAuth access_token refreshed successfully")
|
||||
|
||||
def ensure_fresh_token(self) -> None:
|
||||
"""Proactively refresh the access_token if it is near or past expiry.
|
||||
|
||||
@@ -53,6 +53,59 @@ __all__ = [
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Newer Codex models are gated on the first-party client identity; the previous
|
||||
# browser-shaped User-Agent returned "Model not found" for Luna (#2643).
|
||||
# Use a neutral version because Hindsight must not claim a specific Codex release.
|
||||
_CODEX_ORIGINATOR = "codex_cli_rs"
|
||||
_CODEX_USER_AGENT = "codex_cli_rs/0.0.0 (Hindsight)"
|
||||
|
||||
# Name of the single forced function tool used to carry structured output when
|
||||
# strict_schema is on. The Codex backend speaks the OpenAI Responses API, so a
|
||||
# forced function call gives us constrained decoding straight into the response
|
||||
# schema — no prompt-injected schema, no raw json.loads on free-form model text,
|
||||
# no invalid-\escape retry storm (issue #2504, same class as #1002 / #2339).
|
||||
_STRUCTURED_TOOL_NAME = "structured_response"
|
||||
|
||||
# Valid JSON string escape characters (the char that may follow a backslash).
|
||||
_VALID_JSON_ESCAPE_CHARS = set('"\\/bfnrtu')
|
||||
|
||||
|
||||
def _repair_invalid_json_escapes(text: str) -> str:
|
||||
"""Best-effort repair of invalid ``\\escape`` sequences in a JSON string.
|
||||
|
||||
Escape-heavy content (code, serial/CLI commands, Windows paths, regexes)
|
||||
makes weaker models emit backslashes that aren't valid JSON escapes (e.g.
|
||||
``\\d``, ``\\s``, ``C:\\Users``), so ``json.loads`` fails deterministically
|
||||
and every retry re-fails the same way (issue #2504). This doubles any
|
||||
backslash that isn't part of a valid escape so the payload parses. It is a
|
||||
lenient fallback only — the strict_schema forced-tool path is the real fix.
|
||||
"""
|
||||
result: list[str] = []
|
||||
i = 0
|
||||
n = len(text)
|
||||
while i < n:
|
||||
ch = text[i]
|
||||
if ch == "\\" and i + 1 < n:
|
||||
nxt = text[i + 1]
|
||||
if nxt in _VALID_JSON_ESCAPE_CHARS:
|
||||
# Preserve the valid escape (both chars) verbatim.
|
||||
result.append(ch)
|
||||
result.append(nxt)
|
||||
i += 2
|
||||
continue
|
||||
# Invalid escape: escape the lone backslash so JSON parses.
|
||||
result.append("\\\\")
|
||||
i += 1
|
||||
continue
|
||||
if ch == "\\" and i + 1 == n:
|
||||
# Trailing lone backslash — escape it.
|
||||
result.append("\\\\")
|
||||
i += 1
|
||||
continue
|
||||
result.append(ch)
|
||||
i += 1
|
||||
return "".join(result)
|
||||
|
||||
|
||||
class CodexLLM(LLMInterface):
|
||||
"""
|
||||
@@ -140,6 +193,18 @@ class CodexLLM(LLMInterface):
|
||||
def account_id(self) -> str:
|
||||
return self._auth_manager.account_id
|
||||
|
||||
def _build_request_headers(self) -> httpx.Headers:
|
||||
return httpx.Headers(
|
||||
{
|
||||
"Authorization": f"Bearer {self.access_token}",
|
||||
"Content-Type": "application/json",
|
||||
"OpenAI-Account-ID": self.account_id,
|
||||
"User-Agent": _CODEX_USER_AGENT,
|
||||
"Origin": "https://chatgpt.com",
|
||||
"originator": _CODEX_ORIGINATOR,
|
||||
}
|
||||
)
|
||||
|
||||
@property
|
||||
def refresh_token(self) -> str | None:
|
||||
return self._auth_manager.refresh_token
|
||||
@@ -336,7 +401,18 @@ class CodexLLM(LLMInterface):
|
||||
strict_schema: bool = False,
|
||||
return_usage: bool = False,
|
||||
) -> Any:
|
||||
"""Make API call to Codex backend with SSE streaming."""
|
||||
"""Make API call to Codex backend with SSE streaming.
|
||||
|
||||
Args:
|
||||
strict_schema: Route structured output through a single forced
|
||||
function tool (constrained decoding) instead of prompt-injecting
|
||||
the schema and parsing free-form text. The Codex backend speaks
|
||||
the OpenAI Responses API, so the forced function call emits the
|
||||
response schema directly as tool arguments — eliminating the
|
||||
invalid-``\\escape`` retry storm (issue #2504). When False, falls
|
||||
back to schema-in-prompt + JSON parse, now hardened with a lenient
|
||||
invalid-escape repair before giving up.
|
||||
"""
|
||||
start_time = time.time()
|
||||
|
||||
# Proactively refresh the OAuth access_token if it's near expiry.
|
||||
@@ -361,11 +437,22 @@ class CodexLLM(LLMInterface):
|
||||
else:
|
||||
user_messages.append(msg)
|
||||
|
||||
# Add JSON schema instruction if response_format is provided
|
||||
# Structured output: prefer a single forced function tool (constrained
|
||||
# decoding) over text-injecting the schema and parsing the reply. The
|
||||
# forced tool guarantees schema-shaped JSON in the tool arguments,
|
||||
# eliminating the invalid-\escape retry storm (issue #2504). When
|
||||
# strict_schema is off we keep the schema-in-prompt + json.loads
|
||||
# fallback (now hardened with a lenient escape repair) for callers that
|
||||
# can't force tools.
|
||||
schema = None
|
||||
use_forced_tool = False
|
||||
if response_format is not None and hasattr(response_format, "model_json_schema"):
|
||||
schema = response_format.model_json_schema()
|
||||
schema_msg = f"\n\nYou must respond with valid JSON matching this schema:\n{json.dumps(schema, indent=2, ensure_ascii=False)}"
|
||||
system_instruction += schema_msg
|
||||
if strict_schema:
|
||||
use_forced_tool = True
|
||||
else:
|
||||
schema_msg = f"\n\nYou must respond with valid JSON matching this schema:\n{json.dumps(schema, indent=2, ensure_ascii=False)}"
|
||||
system_instruction += schema_msg
|
||||
|
||||
# gpt-5.2-codex only supports "detailed" reasoning summary
|
||||
reasoning_summary = "detailed" if "5.2" in self.model else self.reasoning_summary
|
||||
@@ -392,13 +479,21 @@ class CodexLLM(LLMInterface):
|
||||
"prompt_cache_key": str(uuid.uuid4()),
|
||||
}
|
||||
|
||||
headers = {
|
||||
"Authorization": f"Bearer {self.access_token}",
|
||||
"Content-Type": "application/json",
|
||||
"OpenAI-Account-ID": self.account_id,
|
||||
"User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7)",
|
||||
"Origin": "https://chatgpt.com",
|
||||
}
|
||||
if use_forced_tool and schema is not None:
|
||||
# Single function tool whose parameters ARE the response schema;
|
||||
# force it via tool_choice so the backend does constrained decoding.
|
||||
payload["tools"] = [
|
||||
{
|
||||
"type": "function",
|
||||
"name": _STRUCTURED_TOOL_NAME,
|
||||
"description": "Return the structured response.",
|
||||
"parameters": schema,
|
||||
}
|
||||
]
|
||||
payload["tool_choice"] = {"type": "function", "name": _STRUCTURED_TOOL_NAME}
|
||||
payload["parallel_tool_calls"] = False
|
||||
|
||||
headers = self._build_request_headers()
|
||||
|
||||
url = f"{self.base_url}/codex/responses"
|
||||
|
||||
@@ -412,8 +507,15 @@ class CodexLLM(LLMInterface):
|
||||
response = await self._client.post(url, json=payload, headers=headers, timeout=120.0)
|
||||
response.raise_for_status()
|
||||
|
||||
# Parse SSE stream
|
||||
content = await self._parse_sse_stream(response)
|
||||
# Forced-tool path: read structured output from the function-call
|
||||
# arguments (already a JSON string in a dedicated channel) rather
|
||||
# than from free-form assistant text.
|
||||
if use_forced_tool:
|
||||
text_content, tool_calls = await self._parse_sse_tool_stream(response)
|
||||
content = text_content or ""
|
||||
else:
|
||||
tool_calls = []
|
||||
content = await self._parse_sse_stream(response)
|
||||
|
||||
# Codex SSE carries no usage block; stash the same char/4 estimate
|
||||
# the success path traces so a later parse/validate failure records
|
||||
@@ -426,7 +528,28 @@ class CodexLLM(LLMInterface):
|
||||
)
|
||||
|
||||
# Handle structured output
|
||||
if response_format is not None:
|
||||
if use_forced_tool:
|
||||
tool_input = None
|
||||
for tc in tool_calls:
|
||||
if tc.name == _STRUCTURED_TOOL_NAME:
|
||||
tool_input = tc.arguments if isinstance(tc.arguments, dict) else None
|
||||
break
|
||||
if tool_input is None:
|
||||
# Model ignored the forced tool (rare — e.g. a gateway that
|
||||
# drops tool_choice). Retry so we don't hard-fail.
|
||||
logger.warning(
|
||||
f"Codex forced structured tool missing from response "
|
||||
f"(attempt {attempt + 1}/{max_retries + 1})"
|
||||
)
|
||||
if attempt < max_retries:
|
||||
backoff = min(initial_backoff * (2**attempt), max_backoff)
|
||||
await asyncio.sleep(backoff)
|
||||
attempt += 1
|
||||
continue
|
||||
raise RuntimeError("Codex did not return the forced structured_response tool call")
|
||||
content = json.dumps(tool_input)
|
||||
result = tool_input if skip_validation else response_format.model_validate(tool_input)
|
||||
elif response_format is not None:
|
||||
# Models may wrap JSON in markdown
|
||||
clean_content = content
|
||||
if "```json" in content:
|
||||
@@ -437,13 +560,20 @@ class CodexLLM(LLMInterface):
|
||||
try:
|
||||
json_data = json.loads(clean_content)
|
||||
except json.JSONDecodeError as e:
|
||||
logger.warning(f"Codex JSON parse error (attempt {attempt + 1}/{max_retries + 1}): {e}")
|
||||
if attempt < max_retries:
|
||||
backoff = min(initial_backoff * (2**attempt), max_backoff)
|
||||
await asyncio.sleep(backoff)
|
||||
attempt += 1
|
||||
continue
|
||||
raise
|
||||
# Escape-heavy content deterministically re-fails every
|
||||
# retry (issue #2504). Try a lenient invalid-escape repair
|
||||
# before burning a retry / re-raising.
|
||||
try:
|
||||
json_data = json.loads(_repair_invalid_json_escapes(clean_content))
|
||||
logger.info("Codex JSON parsed after repairing invalid escape sequences")
|
||||
except json.JSONDecodeError:
|
||||
logger.warning(f"Codex JSON parse error (attempt {attempt + 1}/{max_retries + 1}): {e}")
|
||||
if attempt < max_retries:
|
||||
backoff = min(initial_backoff * (2**attempt), max_backoff)
|
||||
await asyncio.sleep(backoff)
|
||||
attempt += 1
|
||||
continue
|
||||
raise
|
||||
|
||||
if skip_validation:
|
||||
result = json_data
|
||||
@@ -725,13 +855,7 @@ class CodexLLM(LLMInterface):
|
||||
"prompt_cache_key": str(uuid.uuid4()),
|
||||
}
|
||||
|
||||
headers = {
|
||||
"Authorization": f"Bearer {self.access_token}",
|
||||
"Content-Type": "application/json",
|
||||
"OpenAI-Account-ID": self.account_id,
|
||||
"User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7)",
|
||||
"Origin": "https://chatgpt.com",
|
||||
}
|
||||
headers = self._build_request_headers()
|
||||
|
||||
url = f"{self.base_url}/codex/responses"
|
||||
|
||||
@@ -872,8 +996,13 @@ class CodexLLM(LLMInterface):
|
||||
try:
|
||||
arguments = json.loads(arguments_str)
|
||||
except json.JSONDecodeError:
|
||||
logger.warning(f"Failed to parse tool arguments: {arguments_str}")
|
||||
arguments = {}
|
||||
# Escape-heavy content can emit invalid \escape
|
||||
# sequences (issue #2504); repair before giving up.
|
||||
try:
|
||||
arguments = json.loads(_repair_invalid_json_escapes(arguments_str))
|
||||
except json.JSONDecodeError:
|
||||
logger.warning(f"Failed to parse tool arguments: {arguments_str}")
|
||||
arguments = {}
|
||||
|
||||
tool_calls.append(
|
||||
LLMToolCall(
|
||||
|
||||
@@ -56,6 +56,14 @@ _DEFAULT_REFRESH_MARGIN_SECONDS = 5 * 60
|
||||
# to None and callers proceed uncached, rather than stalling the whole batch.
|
||||
_DEFAULT_CREATE_TIMEOUT_SECONDS = 30.0
|
||||
|
||||
# TTL for the per-step reflect caches created by ``create_incremental``. These
|
||||
# live only for the duration of one reflect (seconds), so the TTL is just a
|
||||
# storage backstop in case the explicit ``delete_session`` at reflect end is
|
||||
# missed (crash / event-loop teardown). Short so orphaned caches age out fast —
|
||||
# storage is billed per token-hour, so a 5-minute cap keeps the cost of a leaked
|
||||
# cache negligible.
|
||||
_DEFAULT_INCREMENTAL_TTL_SECONDS = 5 * 60
|
||||
|
||||
|
||||
@dataclass
|
||||
class _CacheEntry:
|
||||
@@ -92,6 +100,10 @@ class GeminiCacheManager:
|
||||
self._create_timeout_seconds = create_timeout_seconds
|
||||
self._entries: dict[str, _CacheEntry] = {}
|
||||
self._lock = asyncio.Lock()
|
||||
# session_id -> CachedContent names created via ``create_incremental``.
|
||||
# A reflect creates a fresh rolling cache per step under one session id;
|
||||
# ``delete_session`` tears them all down when the reflect finishes.
|
||||
self._sessions: dict[str, list[str]] = {}
|
||||
|
||||
@staticmethod
|
||||
def fingerprint(
|
||||
@@ -229,18 +241,99 @@ class GeminiCacheManager:
|
||||
if entry.name == name:
|
||||
self._entries.pop(key, None)
|
||||
|
||||
async def create_incremental(
|
||||
self,
|
||||
*,
|
||||
session_id: str,
|
||||
model: str,
|
||||
system_instruction: str,
|
||||
contents: list[Any],
|
||||
tools: list[dict[str, Any]] | None = None,
|
||||
) -> str | None:
|
||||
"""Create a fresh CachedContent holding ``system + tools + contents`` and
|
||||
track it under ``session_id`` for later teardown.
|
||||
|
||||
Unlike ``get_or_create``, this does NOT deduplicate by fingerprint: each
|
||||
step of a reflect grows the conversation prefix, so every call is a
|
||||
distinct, single-use cache. The reflect loop creates one per step (each
|
||||
covering the previous step's full input) and reuses it for exactly the
|
||||
next model turn, then supersedes it. All caches for the session are
|
||||
deleted by ``delete_session`` when the reflect ends; the short TTL is
|
||||
only a backstop.
|
||||
|
||||
Returns the cache resource name, or ``None`` when caching is disabled,
|
||||
the prefix is below the model minimum, or the create otherwise fails —
|
||||
callers MUST fall back to an uncached call in that case.
|
||||
"""
|
||||
try:
|
||||
name = await self._create_cache(
|
||||
model=model,
|
||||
system_instruction=system_instruction,
|
||||
tools=tools,
|
||||
contents=contents,
|
||||
ttl_seconds=_DEFAULT_INCREMENTAL_TTL_SECONDS,
|
||||
)
|
||||
except _CacheNotEligible as e:
|
||||
logger.debug(
|
||||
"GeminiCacheManager: incremental prefix not eligible (model=%s, reason=%s) — caller falls back",
|
||||
model,
|
||||
e,
|
||||
)
|
||||
return None
|
||||
except Exception:
|
||||
logger.exception(
|
||||
"GeminiCacheManager: failed to create incremental cache (model=%s); caller falls back",
|
||||
model,
|
||||
)
|
||||
return None
|
||||
if name is not None:
|
||||
self._sessions.setdefault(session_id, []).append(name)
|
||||
return name
|
||||
|
||||
async def delete(self, name: str) -> None:
|
||||
"""Best-effort server-side delete of a single CachedContent.
|
||||
|
||||
Swallows all errors: a failed delete just means the cache ages out on
|
||||
its TTL. Also drops any matching in-process entry.
|
||||
"""
|
||||
self.invalidate(name)
|
||||
try:
|
||||
await self._client.aio.caches.delete(name=name)
|
||||
except Exception:
|
||||
logger.debug("GeminiCacheManager: delete of cache %s failed (will age out on TTL)", name, exc_info=True)
|
||||
|
||||
async def delete_session(self, session_id: str) -> None:
|
||||
"""Delete every CachedContent created for ``session_id`` (reflect teardown).
|
||||
|
||||
Deletes concurrently and best-effort — a reflect must never fail because
|
||||
a cache couldn't be torn down; the short TTL is the backstop.
|
||||
"""
|
||||
names = self._sessions.pop(session_id, [])
|
||||
if not names:
|
||||
return
|
||||
await asyncio.gather(*(self.delete(n) for n in names), return_exceptions=True)
|
||||
|
||||
async def _create_cache(
|
||||
self,
|
||||
*,
|
||||
model: str,
|
||||
system_instruction: str,
|
||||
tools: list[dict[str, Any]] | None = None,
|
||||
contents: list[Any] | None = None,
|
||||
ttl_seconds: int | None = None,
|
||||
) -> str | None:
|
||||
"""Wrap ``client.aio.caches.create`` with the config we want.
|
||||
|
||||
The SDK surface differs slightly across google-genai versions;
|
||||
this implementation targets the >=1.0.0 line where caches live
|
||||
under ``client.aio.caches``.
|
||||
|
||||
``contents`` (already-converted ``genai_types.Content`` turns) is
|
||||
appended after the system_instruction/tools so the cache can hold a
|
||||
growing multi-turn conversation prefix, not just the static prefix —
|
||||
this is what the step-by-step reflect cache relies on. ``ttl_seconds``
|
||||
overrides the manager default (used to give per-step reflect caches a
|
||||
short backstop TTL).
|
||||
"""
|
||||
# Lazy import so this module doesn't require the SDK at import time.
|
||||
from google.genai import types as genai_types
|
||||
@@ -254,8 +347,10 @@ class GeminiCacheManager:
|
||||
# still part of the fingerprint so a schema change keys a fresh cache.
|
||||
config_kwargs: dict[str, Any] = {
|
||||
"system_instruction": system_instruction,
|
||||
"ttl": f"{self._ttl_seconds}s",
|
||||
"ttl": f"{ttl_seconds if ttl_seconds is not None else self._ttl_seconds}s",
|
||||
}
|
||||
if contents:
|
||||
config_kwargs["contents"] = contents
|
||||
if tools:
|
||||
# OpenAI-style {"function": {...}} entries must be converted to
|
||||
# Gemini's Tool/FunctionDeclaration shape before caching.
|
||||
|
||||
@@ -13,6 +13,7 @@ import json
|
||||
import logging
|
||||
import time
|
||||
from contextvars import ContextVar
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
from google import genai
|
||||
@@ -63,6 +64,77 @@ def _usage_from_gemini_response(response: Any) -> LLMResponseUsage:
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _GeminiConversation:
|
||||
"""A message list converted to Gemini's request shape."""
|
||||
|
||||
system_instruction: str | None
|
||||
contents: list["genai_types.Content"]
|
||||
|
||||
|
||||
def _convert_messages_to_gemini(msg_list: list[dict[str, Any]]) -> _GeminiConversation:
|
||||
"""Convert OpenAI-style messages to a Gemini (system_instruction, contents) pair.
|
||||
|
||||
Shared by ``call_with_tools`` (request body) and the incremental cache
|
||||
builder so a cached prefix and the live request serialise turns identically —
|
||||
any drift would fingerprint differently and defeat the cache. Consecutive
|
||||
``role="tool"`` messages are grouped into a single ``user`` Content with
|
||||
multiple FunctionResponse parts, matching Gemini's multi-turn requirement.
|
||||
"""
|
||||
system_instruction: str | None = None
|
||||
gemini_contents: list[genai_types.Content] = []
|
||||
i = 0
|
||||
while i < len(msg_list):
|
||||
msg = msg_list[i]
|
||||
role = msg.get("role", "user")
|
||||
content = msg.get("content", "")
|
||||
|
||||
if role == "system":
|
||||
system_instruction = (system_instruction + "\n\n" + content) if system_instruction else content
|
||||
i += 1
|
||||
elif role == "tool":
|
||||
parts = []
|
||||
while i < len(msg_list) and msg_list[i].get("role") == "tool":
|
||||
tool_msg = msg_list[i]
|
||||
tool_content = tool_msg.get("content", "")
|
||||
parts.append(
|
||||
genai_types.Part(
|
||||
function_response=genai_types.FunctionResponse(
|
||||
name=tool_msg.get("name", ""),
|
||||
response={"result": tool_content},
|
||||
)
|
||||
)
|
||||
)
|
||||
i += 1
|
||||
gemini_contents.append(genai_types.Content(role="user", parts=parts))
|
||||
elif role == "assistant":
|
||||
tool_calls_in_msg = msg.get("tool_calls", [])
|
||||
if tool_calls_in_msg:
|
||||
parts = []
|
||||
if content:
|
||||
parts.append(genai_types.Part(text=content))
|
||||
for tc in tool_calls_in_msg:
|
||||
fn = tc.get("function", {})
|
||||
fn_name = fn.get("name", "")
|
||||
fn_args_str = fn.get("arguments", "{}")
|
||||
fn_args = parse_llm_json(fn_args_str)
|
||||
thought_signature = tc.get("thought_signature")
|
||||
fc_kwargs: dict[str, Any] = {"name": fn_name, "args": fn_args}
|
||||
part_kwargs: dict[str, Any] = {"function_call": genai_types.FunctionCall(**fc_kwargs)}
|
||||
if thought_signature:
|
||||
part_kwargs["thought_signature"] = base64.b64decode(thought_signature)
|
||||
parts.append(genai_types.Part(**part_kwargs))
|
||||
gemini_contents.append(genai_types.Content(role="model", parts=parts))
|
||||
else:
|
||||
gemini_contents.append(genai_types.Content(role="model", parts=[genai_types.Part(text=content)]))
|
||||
i += 1
|
||||
else:
|
||||
gemini_contents.append(genai_types.Content(role="user", parts=[genai_types.Part(text=content)]))
|
||||
i += 1
|
||||
|
||||
return _GeminiConversation(system_instruction=system_instruction, contents=gemini_contents)
|
||||
|
||||
|
||||
class GeminiLLM(LLMInterface):
|
||||
"""
|
||||
LLM provider for Google Gemini and Vertex AI.
|
||||
@@ -527,6 +599,7 @@ class GeminiLLM(LLMInterface):
|
||||
max_backoff: float = 30.0,
|
||||
tool_choice: str | dict[str, Any] = "auto",
|
||||
cached_prefix: str | None = None,
|
||||
cached_prefix_message_count: int = 0,
|
||||
) -> LLMToolCallResult:
|
||||
"""
|
||||
Make a Gemini/VertexAI API call with tool/function calling support.
|
||||
@@ -542,13 +615,20 @@ class GeminiLLM(LLMInterface):
|
||||
max_backoff: Maximum backoff time in seconds.
|
||||
tool_choice: How to choose tools (Gemini uses "auto" only).
|
||||
cached_prefix: Optional CachedContent resource name (from
|
||||
``GeminiCacheManager.get_or_create`` with ``tools=...``). When
|
||||
set, the system_instruction and tool definitions are assumed
|
||||
``GeminiCacheManager.get_or_create`` or ``create_incremental``).
|
||||
When set, the system_instruction and tool definitions are assumed
|
||||
to live in the cache; this call will skip resending them and
|
||||
the cached prefix is billed at the cached-input rate. The
|
||||
``tools`` argument is still required (the caller may pass
|
||||
an empty list when the cache holds them) so existing call
|
||||
sites don't break.
|
||||
cached_prefix_message_count: Number of leading ``messages`` already
|
||||
baked into ``cached_prefix`` (the step-by-step reflect cache holds
|
||||
a growing conversation prefix, not just system+tools). Only the
|
||||
messages AFTER this index are sent as request contents — the rest
|
||||
come from the cache and bill at the cached rate. 0 means the cache
|
||||
holds only the static prefix (system+tools), so the full
|
||||
conversation is still sent (legacy behaviour).
|
||||
|
||||
Returns:
|
||||
LLMToolCallResult with content and/or tool_calls.
|
||||
@@ -556,86 +636,45 @@ class GeminiLLM(LLMInterface):
|
||||
start_time = time.time()
|
||||
using_cache = cached_prefix is not None
|
||||
|
||||
# Convert tools to Gemini format. When the cache is in use, the
|
||||
# tool definitions are baked into the CachedContent at create time
|
||||
# and the SDK rejects re-sending them alongside ``cached_content``.
|
||||
# Convert tools to Gemini format. While the cache is in use the tool
|
||||
# definitions live in the CachedContent and the SDK rejects re-sending
|
||||
# them alongside ``cached_content`` (see ``_build_tools_config``), but we
|
||||
# still build them unconditionally so the cached-call-failed fallback —
|
||||
# which drops the cache and re-sends prefix + tools inline — has real
|
||||
# tools to send rather than an empty list.
|
||||
gemini_tools = []
|
||||
if not using_cache:
|
||||
for tool in tools:
|
||||
func = tool.get("function", {})
|
||||
gemini_tools.append(
|
||||
genai_types.Tool(
|
||||
function_declarations=[
|
||||
genai_types.FunctionDeclaration(
|
||||
name=func.get("name", ""),
|
||||
description=func.get("description", ""),
|
||||
parameters=func.get("parameters"),
|
||||
)
|
||||
]
|
||||
)
|
||||
)
|
||||
|
||||
# Convert messages
|
||||
system_instruction = None
|
||||
gemini_contents = []
|
||||
msg_list = list(messages)
|
||||
i = 0
|
||||
while i < len(msg_list):
|
||||
msg = msg_list[i]
|
||||
role = msg.get("role", "user")
|
||||
content = msg.get("content", "")
|
||||
|
||||
if role == "system":
|
||||
# Always capture system_instruction. _build_tools_config omits it
|
||||
# (and tools) from the request while the cache carries the prefix,
|
||||
# but it must be available so the cached-call-failed safety net can
|
||||
# re-send the prefix + tools inline.
|
||||
system_instruction = (system_instruction + "\n\n" + content) if system_instruction else content
|
||||
i += 1
|
||||
elif role == "tool":
|
||||
# Gemini requires ALL tool responses for a given model turn to be grouped
|
||||
# into a single Content with multiple FunctionResponse parts.
|
||||
# Consecutive role="tool" messages correspond to one model turn's tool calls.
|
||||
parts = []
|
||||
while i < len(msg_list) and msg_list[i].get("role") == "tool":
|
||||
tool_msg = msg_list[i]
|
||||
tool_content = tool_msg.get("content", "")
|
||||
parts.append(
|
||||
genai_types.Part(
|
||||
function_response=genai_types.FunctionResponse(
|
||||
name=tool_msg.get("name", ""),
|
||||
response={"result": tool_content},
|
||||
)
|
||||
for tool in tools:
|
||||
func = tool.get("function", {})
|
||||
gemini_tools.append(
|
||||
genai_types.Tool(
|
||||
function_declarations=[
|
||||
genai_types.FunctionDeclaration(
|
||||
name=func.get("name", ""),
|
||||
description=func.get("description", ""),
|
||||
parameters=func.get("parameters"),
|
||||
)
|
||||
)
|
||||
i += 1
|
||||
gemini_contents.append(genai_types.Content(role="user", parts=parts))
|
||||
elif role == "assistant":
|
||||
tool_calls_in_msg = msg.get("tool_calls", [])
|
||||
if tool_calls_in_msg:
|
||||
# Convert OpenAI-style tool_calls to Gemini function_call parts
|
||||
# This is required for proper multi-turn conversation history
|
||||
parts = []
|
||||
if content:
|
||||
parts.append(genai_types.Part(text=content))
|
||||
for tc in tool_calls_in_msg:
|
||||
fn = tc.get("function", {})
|
||||
fn_name = fn.get("name", "")
|
||||
fn_args_str = fn.get("arguments", "{}")
|
||||
fn_args = parse_llm_json(fn_args_str)
|
||||
thought_signature = tc.get("thought_signature")
|
||||
fc_kwargs: dict[str, Any] = {"name": fn_name, "args": fn_args}
|
||||
part_kwargs: dict[str, Any] = {"function_call": genai_types.FunctionCall(**fc_kwargs)}
|
||||
if thought_signature:
|
||||
part_kwargs["thought_signature"] = base64.b64decode(thought_signature)
|
||||
parts.append(genai_types.Part(**part_kwargs))
|
||||
gemini_contents.append(genai_types.Content(role="model", parts=parts))
|
||||
else:
|
||||
gemini_contents.append(genai_types.Content(role="model", parts=[genai_types.Part(text=content)]))
|
||||
i += 1
|
||||
else:
|
||||
gemini_contents.append(genai_types.Content(role="user", parts=[genai_types.Part(text=content)]))
|
||||
i += 1
|
||||
]
|
||||
)
|
||||
)
|
||||
|
||||
# Convert messages. ``system_instruction`` and the FULL contents are always
|
||||
# computed: _build_tools_config omits system/tools from the request while
|
||||
# the cache carries the prefix, but the cached-call-failed safety net must
|
||||
# be able to re-send the whole prefix + tools inline.
|
||||
converted = _convert_messages_to_gemini(list(messages))
|
||||
system_instruction = converted.system_instruction
|
||||
full_contents = converted.contents
|
||||
|
||||
# Step-by-step reflect cache: when the cache already holds the first
|
||||
# ``cached_prefix_message_count`` messages, send ONLY the newer turns as
|
||||
# request contents — the cached prefix supplies the rest at the cached
|
||||
# rate. The split is always at a whole-turn boundary (the reflect loop
|
||||
# advances the cache one completed turn at a time), so slicing the raw
|
||||
# messages before conversion never splits a grouped tool turn.
|
||||
if using_cache and cached_prefix_message_count > 0:
|
||||
delta_contents = _convert_messages_to_gemini(list(messages)[cached_prefix_message_count:]).contents
|
||||
else:
|
||||
delta_contents = full_contents
|
||||
|
||||
# Apply safety settings: context var (per-request bank override) takes precedence over instance default
|
||||
effective_safety_settings = _safety_settings_ctx.get()
|
||||
@@ -701,10 +740,14 @@ class GeminiLLM(LLMInterface):
|
||||
if attempt > 0:
|
||||
set_stage(f"llm.gemini.tools.attempt={attempt + 1}/{max_retries + 1}")
|
||||
try:
|
||||
# With the cache active, send only the un-cached tail (delta);
|
||||
# on the uncached fallback path send the full conversation so the
|
||||
# re-inlined system+tools prefix has its whole context.
|
||||
active_contents = delta_contents if cache_active else full_contents
|
||||
response = await asyncio.wait_for(
|
||||
self._client.aio.models.generate_content(
|
||||
model=self.model,
|
||||
contents=gemini_contents,
|
||||
contents=active_contents,
|
||||
config=config,
|
||||
),
|
||||
timeout=90.0, # Safety net for network hangs; valid slow responses are <90s
|
||||
@@ -883,6 +926,56 @@ class GeminiLLM(LLMInterface):
|
||||
tools=tools,
|
||||
)
|
||||
|
||||
# ── Step-by-step incremental prompt caching (reflect tool loop) ──────────
|
||||
|
||||
def supports_incremental_prompt_cache(self) -> bool:
|
||||
"""True when explicit caching is on — the reflect loop can then roll a
|
||||
per-step CachedContent that grows with the conversation."""
|
||||
return self._prompt_cache_enabled
|
||||
|
||||
def _ensure_cache_manager(self) -> Any:
|
||||
if self._cache_manager is None:
|
||||
from hindsight_api.engine.providers.gemini_cache import GeminiCacheManager
|
||||
|
||||
self._cache_manager = GeminiCacheManager(self._client)
|
||||
return self._cache_manager
|
||||
|
||||
async def create_incremental_cache(
|
||||
self,
|
||||
*,
|
||||
session_id: str,
|
||||
messages: list[dict[str, Any]],
|
||||
tools: list[dict[str, Any]] | None = None,
|
||||
) -> str | None:
|
||||
"""Cache ``system + tools + messages`` as a conversation prefix and return
|
||||
its resource name (or ``None`` — caller falls back to an uncached call).
|
||||
|
||||
The reflect loop calls this once per step with the growing message list so
|
||||
each step's cache entirely contains the previous step's input; the next
|
||||
model turn then references it and re-sends only its own delta. Caches are
|
||||
tracked under ``session_id`` and torn down by ``delete_cache_session``.
|
||||
"""
|
||||
if not self._prompt_cache_enabled or self._client is None:
|
||||
return None
|
||||
converted = _convert_messages_to_gemini(list(messages))
|
||||
return await self._ensure_cache_manager().create_incremental(
|
||||
session_id=session_id,
|
||||
model=self.model,
|
||||
system_instruction=converted.system_instruction or "",
|
||||
contents=converted.contents,
|
||||
tools=tools,
|
||||
)
|
||||
|
||||
async def delete_cached_prefix(self, name: str) -> None:
|
||||
"""Best-effort delete of a single CachedContent (superseded reflect step)."""
|
||||
if self._cache_manager is not None:
|
||||
await self._cache_manager.delete(name)
|
||||
|
||||
async def delete_cache_session(self, session_id: str) -> None:
|
||||
"""Tear down every CachedContent created for a reflect session."""
|
||||
if self._cache_manager is not None:
|
||||
await self._cache_manager.delete_session(session_id)
|
||||
|
||||
# ── Batch API (Gemini API only — not Vertex AI) ─────────────────────────
|
||||
#
|
||||
# Google's Gemini Batch API gives a flat 50% discount on input + output
|
||||
@@ -1030,7 +1123,7 @@ class GeminiLLM(LLMInterface):
|
||||
Mirrors the synchronous ``call`` path: system messages become
|
||||
``systemInstruction``; a ``response_format`` json_schema forces JSON
|
||||
output (``responseMimeType``), appends the schema as a textual hint, and
|
||||
grammar-enforces via ``responseJsonSchema`` when ``strict`` is set.
|
||||
grammar-enforces via ``responseJsonSchema`` whenever a schema is present.
|
||||
"""
|
||||
system_texts: list[str] = []
|
||||
contents: list[dict[str, Any]] = []
|
||||
@@ -1059,8 +1152,13 @@ class GeminiLLM(LLMInterface):
|
||||
system_texts.append(
|
||||
"You must respond with valid JSON matching this schema:\n" + json.dumps(schema, ensure_ascii=False)
|
||||
)
|
||||
if json_schema.get("strict"):
|
||||
generation_config["responseJsonSchema"] = schema
|
||||
# #2699: Gemini always grammar-enforces structured output via its native
|
||||
# response_schema (``strict`` is an OpenAI concept, meaningless here). Set
|
||||
# the native schema whenever one is present so the batch path mirrors the
|
||||
# interactive path; otherwise batch requests at default config
|
||||
# (HINDSIGHT_API_LLM_STRICT_SCHEMA=False) get only a textual hint and
|
||||
# intermittently emit malformed JSON, losing every fact in the chunk.
|
||||
generation_config["responseJsonSchema"] = schema
|
||||
|
||||
request: dict[str, Any] = {"contents": contents}
|
||||
if system_texts:
|
||||
|
||||
@@ -47,15 +47,31 @@ logger = logging.getLogger(__name__)
|
||||
# Seed applied to every Groq request for deterministic behavior
|
||||
DEFAULT_LLM_SEED = 4242
|
||||
JSON_MODE_USER_HINT = "Return valid json only."
|
||||
DEFAULT_VERIFICATION_MAX_COMPLETION_TOKENS = 512
|
||||
|
||||
# Self-hosted OpenAI-compatible servers that advertise tool_choice="required"
|
||||
|
||||
def _validate_ollama_num_ctx(value: Any) -> int | None:
|
||||
"""Validate a native Ollama context-window override."""
|
||||
if value is None:
|
||||
return None
|
||||
if isinstance(value, bool) or not isinstance(value, int):
|
||||
raise ValueError(f"ollama_num_ctx must be a positive integer, got {value!r}")
|
||||
if value < 1:
|
||||
raise ValueError(f"ollama_num_ctx must be >= 1, got {value}")
|
||||
return value
|
||||
|
||||
|
||||
# Provider implementations that advertise tool_choice="required"
|
||||
# but silently ignore it: instead of forcing a tool call they return
|
||||
# finish_reason "stop"/"tool_calls" with an EMPTY tool_calls array and no error.
|
||||
# Reflect's agent loop then sees no tool call, runs synthesis with no retrieval,
|
||||
# and answers "I don't have information" even when the bank holds the answer.
|
||||
# See issues #1563 (LM Studio), #1179 (LM Studio + Qwen), #1877 (vLLM with
|
||||
# --enable-auto-tool-choice). llama-server (the "llamacpp" provider) honors
|
||||
# "required" correctly and is intentionally excluded (#1179).
|
||||
# --enable-auto-tool-choice). The generic OpenAI provider is intentionally not
|
||||
# inferred from its URL: custom OpenAI-compatible endpoints can implement the
|
||||
# required-tool contract, and silently downgrading them changes request semantics.
|
||||
# llama-server (the "llamacpp" provider) honors "required" correctly and is
|
||||
# intentionally excluded (#1179).
|
||||
_TOOL_CHOICE_REQUIRED_UNSUPPORTED_PROVIDERS = frozenset({"lmstudio", "ollama"})
|
||||
|
||||
|
||||
@@ -67,23 +83,68 @@ class ProviderResponseError(RuntimeError):
|
||||
self.retryable = retryable
|
||||
|
||||
|
||||
def _is_json(text: str) -> bool:
|
||||
"""True if ``text`` parses as a JSON value."""
|
||||
try:
|
||||
json.loads(text)
|
||||
except (json.JSONDecodeError, ValueError):
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def _outer_json_span(content: str) -> str | None:
|
||||
"""Return the outermost ``{...}`` / ``[...]`` span if it parses as JSON, else None.
|
||||
|
||||
Fallback for responses where fences are partial/absent or the model wrapped
|
||||
the JSON in surrounding prose. Only returned when it is valid JSON so callers
|
||||
never receive a worse candidate than the raw content.
|
||||
"""
|
||||
starts = [i for i in (content.find("{"), content.find("[")) if i >= 0]
|
||||
ends = [i for i in (content.rfind("}"), content.rfind("]")) if i >= 0]
|
||||
if not starts or not ends:
|
||||
return None
|
||||
start, end = min(starts), max(ends)
|
||||
if end <= start:
|
||||
return None
|
||||
candidate = content[start : end + 1].strip()
|
||||
return candidate if _is_json(candidate) else None
|
||||
|
||||
|
||||
def _strip_code_fences(content: str) -> str:
|
||||
"""Strip markdown code fences from LLM response if present.
|
||||
|
||||
Many LLM providers (MiniMax, some Ollama models, Claude via proxies)
|
||||
wrap JSON responses in ```json ... ``` fences even when json_object
|
||||
response format is requested. This strips the fences while preserving
|
||||
the JSON content inside. Returns the original content unchanged if
|
||||
no fences are detected.
|
||||
response format is requested. Fences are detected by line (a closing
|
||||
``` must sit alone on its line) so triple-backticks *inside* JSON string
|
||||
values do not truncate the payload. When the stripped candidate is not
|
||||
valid JSON (partial fence, prose-wrapped output, truncated response), fall
|
||||
back to the outermost parseable JSON span. Returns the original content
|
||||
unchanged if no better candidate is found.
|
||||
"""
|
||||
if "```" not in content:
|
||||
return content
|
||||
try:
|
||||
if "```json" in content:
|
||||
return content.split("```json")[1].split("```")[0].strip()
|
||||
return content.split("```")[1].split("```")[0].strip()
|
||||
except (IndexError, ValueError):
|
||||
return content
|
||||
candidate = content
|
||||
if "```" in content:
|
||||
lines = content.split("\n")
|
||||
# Find first line that starts a code fence (``` optionally followed by language)
|
||||
fence_start = next((i for i, line in enumerate(lines) if line.startswith("```")), None)
|
||||
if fence_start is not None:
|
||||
# Find matching closing fence (``` alone or with trailing whitespace)
|
||||
fence_end = next(
|
||||
(j for j in range(fence_start + 1, len(lines)) if lines[j].strip() == "```"),
|
||||
None,
|
||||
)
|
||||
if fence_end is not None:
|
||||
candidate = "\n".join(lines[fence_start + 1 : fence_end]).strip()
|
||||
|
||||
if _is_json(candidate):
|
||||
return candidate
|
||||
|
||||
# Fence stripping did not yield valid JSON — try to recover the outer JSON span.
|
||||
span = _outer_json_span(content)
|
||||
if span is not None:
|
||||
return span
|
||||
|
||||
return candidate
|
||||
|
||||
|
||||
# Reasoning/thinking tags emitted by extended-thinking models. Some providers
|
||||
@@ -435,6 +496,8 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
timeout: float | None = None,
|
||||
groq_service_tier: str | None = None,
|
||||
extra_body: dict[str, Any] | None = None,
|
||||
*,
|
||||
ollama_num_ctx: int | None = None,
|
||||
**kwargs: Any,
|
||||
):
|
||||
"""
|
||||
@@ -449,6 +512,8 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
timeout: Request timeout in seconds (uses env var or 120s default).
|
||||
groq_service_tier: Groq service tier ("on_demand", "flex", "auto").
|
||||
extra_body: Extra body params merged into every API call.
|
||||
ollama_num_ctx: Native Ollama context window override. None lets Ollama use
|
||||
the model/server default.
|
||||
**kwargs: Additional provider-specific parameters.
|
||||
"""
|
||||
super().__init__(provider, api_key, base_url, model, reasoning_effort, **kwargs)
|
||||
@@ -529,6 +594,7 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
# Service tier configuration (from config, not env vars)
|
||||
self.groq_service_tier = groq_service_tier
|
||||
self.openai_service_tier = kwargs.get("openai_service_tier")
|
||||
self.ollama_num_ctx = _validate_ollama_num_ctx(ollama_num_ctx)
|
||||
# User-configured extra body params (merged into every API call)
|
||||
self._config_extra_body = extra_body or {}
|
||||
|
||||
@@ -559,17 +625,17 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
def _drops_tool_choice_required(self) -> bool:
|
||||
"""Whether this endpoint silently ignores ``tool_choice="required"``.
|
||||
|
||||
True for self-hosted OpenAI-compatible servers known to return an empty
|
||||
tool_calls array for "required" instead of forcing a call (#1563/#1179/
|
||||
#1877). Covers LM Studio / Ollama directly, plus any server reached via
|
||||
the generic "openai" provider with a custom ``base_url`` (e.g. a local
|
||||
vLLM endpoint). The real OpenAI API (no base_url override) honors
|
||||
"required", and cloud providers keep their own default base_urls, so both
|
||||
are left untouched.
|
||||
Only explicitly identified provider implementations are classified as
|
||||
unsupported. A custom base URL does not identify endpoint capabilities:
|
||||
an OpenAI-compatible endpoint may correctly enforce required tool calls,
|
||||
and replacing ``required`` with ``auto`` would violate the caller's named
|
||||
tool choice after the tools list has been narrowed.
|
||||
"""
|
||||
if self.provider in _TOOL_CHOICE_REQUIRED_UNSUPPORTED_PROVIDERS:
|
||||
return True
|
||||
return self.provider == "openai" and bool(self.base_url)
|
||||
return self.provider in _TOOL_CHOICE_REQUIRED_UNSUPPORTED_PROVIDERS
|
||||
|
||||
def _verification_max_completion_tokens(self) -> int:
|
||||
"""Return the startup verification budget for OpenAI-compatible gateways."""
|
||||
return DEFAULT_VERIFICATION_MAX_COMPLETION_TOKENS
|
||||
|
||||
async def verify_connection(self) -> None:
|
||||
"""
|
||||
@@ -582,7 +648,7 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
logger.info(f"Verifying connection: {self.provider}/{self.model}")
|
||||
await self.call(
|
||||
messages=[{"role": "user", "content": "Say 'ok'"}],
|
||||
max_completion_tokens=100,
|
||||
max_completion_tokens=self._verification_max_completion_tokens(),
|
||||
max_retries=2,
|
||||
initial_backoff=0.5,
|
||||
max_backoff=2.0,
|
||||
@@ -644,6 +710,11 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
# use the widely-supported max_tokens
|
||||
return "max_tokens"
|
||||
|
||||
def _apply_provider_extra_body_defaults(self, extra_body: dict[str, Any]) -> None:
|
||||
"""Apply provider-specific extra_body defaults while preserving user overrides."""
|
||||
if self.provider == "minimax":
|
||||
extra_body.setdefault("thinking", {"type": "disabled"})
|
||||
|
||||
async def call(
|
||||
self,
|
||||
messages: list[dict[str, str]],
|
||||
@@ -731,6 +802,7 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
|
||||
# Provider-specific parameters
|
||||
extra_body: dict[str, Any] = {**self._config_extra_body}
|
||||
self._apply_provider_extra_body_defaults(extra_body)
|
||||
if self.provider == "groq":
|
||||
call_params["seed"] = DEFAULT_LLM_SEED
|
||||
# Add service_tier if configured
|
||||
@@ -886,7 +958,9 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
output_tokens = max(0, output_tokens - thoughts_tokens)
|
||||
total_tokens = max(0, total_tokens - thoughts_tokens)
|
||||
|
||||
# Record LLM metrics
|
||||
# Record LLM metrics. ``output_tokens`` is visible-only by now, so
|
||||
# ``thoughts_tokens`` has to be recorded alongside it or the reasoning
|
||||
# half of the billed output reaches no counter at all.
|
||||
metrics = get_metrics_collector()
|
||||
metrics.record_llm_call(
|
||||
provider=self.provider,
|
||||
@@ -896,6 +970,8 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
success=True,
|
||||
cached_input_tokens=cached_tokens,
|
||||
thoughts_tokens=thoughts_tokens,
|
||||
)
|
||||
|
||||
# Record trace span
|
||||
@@ -1086,8 +1162,12 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
forced_name = request_tool_choice.get("function", {}).get("name")
|
||||
if forced_name:
|
||||
filtered = [t for t in tools if t.get("function", {}).get("name") == forced_name]
|
||||
if filtered:
|
||||
tools = filtered
|
||||
if len(filtered) != 1:
|
||||
raise ValueError(
|
||||
f"Named tool_choice must reference exactly one declared tool; "
|
||||
f"found {len(filtered)} definitions for {forced_name!r}"
|
||||
)
|
||||
tools = filtered
|
||||
request_tool_choice = "required"
|
||||
|
||||
# DeepSeek accepts tool calls but rejects explicit required/named
|
||||
@@ -1105,13 +1185,13 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
if request_tool_choice == "auto":
|
||||
request_tool_choice = None
|
||||
|
||||
# vLLM (--enable-auto-tool-choice), LM Studio, Ollama and similar
|
||||
# self-hosted servers silently drop tool_choice="required", returning an
|
||||
# empty tool_calls array instead of forcing a call (#1563/#1179/#1877).
|
||||
# LM Studio and Ollama silently drop tool_choice="required", returning an
|
||||
# empty tool_calls array instead of forcing a call (#1563/#1179).
|
||||
# Downgrade to auto (None) so the model still gets to call a tool. Named
|
||||
# tool_choice dicts were already normalized to "required" + a single
|
||||
# filtered tool above, so the call stays practically forced even under
|
||||
# auto. The real OpenAI API honors "required" and is left untouched.
|
||||
# auto. Generic OpenAI-compatible endpoints retain the canonical
|
||||
# ``required`` contract regardless of whether they use a custom base URL.
|
||||
if request_tool_choice == "required" and self._drops_tool_choice_required():
|
||||
request_tool_choice = None
|
||||
|
||||
@@ -1149,6 +1229,7 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
|
||||
# Provider-specific parameters
|
||||
extra_body: dict[str, Any] = {**self._config_extra_body}
|
||||
self._apply_provider_extra_body_defaults(extra_body)
|
||||
if self.provider == "groq":
|
||||
call_params["seed"] = DEFAULT_LLM_SEED
|
||||
if extra_body:
|
||||
@@ -1196,6 +1277,8 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
if thoughts_tokens:
|
||||
output_tokens = max(0, output_tokens - thoughts_tokens)
|
||||
|
||||
# See ``call()``: record the reasoning and cached counts too, so no
|
||||
# billed token is dropped from the metrics counters.
|
||||
metrics = get_metrics_collector()
|
||||
metrics.record_llm_call(
|
||||
provider=self.provider,
|
||||
@@ -1205,6 +1288,8 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
success=True,
|
||||
cached_input_tokens=cached_tokens,
|
||||
thoughts_tokens=thoughts_tokens,
|
||||
)
|
||||
|
||||
# Record OpenTelemetry span
|
||||
@@ -1337,9 +1422,10 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
|
||||
# Add optional parameters with optimized defaults for Ollama
|
||||
options: dict[str, Any] = {
|
||||
"num_ctx": 16384, # 16k context window for larger prompts
|
||||
"num_batch": 512, # Optimal batch size for prompt processing
|
||||
}
|
||||
if self.ollama_num_ctx is not None:
|
||||
options["num_ctx"] = self.ollama_num_ctx
|
||||
if max_completion_tokens:
|
||||
options["num_predict"] = max_completion_tokens
|
||||
if temperature is not None:
|
||||
|
||||
@@ -6,6 +6,7 @@ structured information like temporal constraints.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import re
|
||||
from abc import ABC, abstractmethod
|
||||
from datetime import datetime, timedelta
|
||||
|
||||
@@ -19,6 +20,103 @@ from hindsight_api.engine.temporal_periods import (
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# dateparser.search_dates over-matches: short common words that happen to be
|
||||
# weekday/month abbreviations in *some* language ("we"/"me"/"did" -> a weekday,
|
||||
# "do" -> Sunday) come back as bogus dates. When such a false positive appears
|
||||
# *before* the real date in the query, taking the first match (or a hard-coded
|
||||
# blacklist of such words) silently produces a wrong temporal window — worse
|
||||
# than none, because the constraint is non-null so nothing downstream can tell
|
||||
# extraction failed. See issue #2768.
|
||||
#
|
||||
# Instead of blacklisting words one at a time (a moving target — every short
|
||||
# word dateparser resolves is a new instance of the same bug), we score each
|
||||
# match by the date signal it actually carries and keep only matches with a
|
||||
# real signal, preferring the strongest. A bare weekday abbreviation carries no
|
||||
# day/month/year and scores zero, so it is rejected regardless of language or
|
||||
# dateparser version.
|
||||
_TOKEN_RE = re.compile(r"[a-z0-9]+")
|
||||
_MONTH_WORDS = {
|
||||
"january",
|
||||
"february",
|
||||
"march",
|
||||
"april",
|
||||
"may",
|
||||
"june",
|
||||
"july",
|
||||
"august",
|
||||
"september",
|
||||
"october",
|
||||
"november",
|
||||
"december",
|
||||
}
|
||||
_RELATIVE_WORDS = {"today", "yesterday", "tomorrow", "tonight", "now"}
|
||||
_WEEKDAY_WORDS = {
|
||||
"monday",
|
||||
"tuesday",
|
||||
"wednesday",
|
||||
"thursday",
|
||||
"friday",
|
||||
"saturday",
|
||||
"sunday",
|
||||
}
|
||||
_PERIOD_WORDS = {
|
||||
"last",
|
||||
"next",
|
||||
"this",
|
||||
"past",
|
||||
"coming",
|
||||
"ago",
|
||||
"week",
|
||||
"weeks",
|
||||
"month",
|
||||
"months",
|
||||
"year",
|
||||
"years",
|
||||
"day",
|
||||
"days",
|
||||
"hour",
|
||||
"hours",
|
||||
"minute",
|
||||
"minutes",
|
||||
"quarter",
|
||||
"decade",
|
||||
"century",
|
||||
"weekend",
|
||||
"morning",
|
||||
"afternoon",
|
||||
"evening",
|
||||
"night",
|
||||
"noon",
|
||||
"midnight",
|
||||
}
|
||||
|
||||
|
||||
def _date_match_score(text: str) -> int:
|
||||
"""Score how strong a temporal signal a matched span carries.
|
||||
|
||||
A score of 0 means the span is a bare token with no explicit date content
|
||||
(the false-positive class from issue #2768) and should be rejected. Higher
|
||||
scores mean a stronger, less ambiguous date reference. A digit is the
|
||||
strongest signal (day/year/ISO date); an explicit English month/relative
|
||||
word next; weekday names and period words weakest but still explicit.
|
||||
"""
|
||||
tokens = _TOKEN_RE.findall(text.lower())
|
||||
if not tokens:
|
||||
return 0
|
||||
score = 0
|
||||
if any(any(ch.isdigit() for ch in tok) for tok in tokens):
|
||||
score += 100
|
||||
token_set = set(tokens)
|
||||
if token_set & _MONTH_WORDS:
|
||||
score += 50
|
||||
if token_set & _RELATIVE_WORDS:
|
||||
score += 50
|
||||
if token_set & _WEEKDAY_WORDS:
|
||||
score += 30
|
||||
if token_set & _PERIOD_WORDS:
|
||||
score += 20
|
||||
return score
|
||||
|
||||
|
||||
class TemporalConstraint(BaseModel):
|
||||
"""
|
||||
@@ -164,20 +262,23 @@ class DateparserQueryAnalyzer(QueryAnalyzer):
|
||||
if not results:
|
||||
return QueryAnalysis(temporal_constraint=None)
|
||||
|
||||
# Filter out false positives (common words parsed as dates)
|
||||
false_positives = {"do", "may", "march", "will", "can", "sat", "sun", "mon", "tue", "wed", "thu", "fri"}
|
||||
valid_results = [
|
||||
(text, date)
|
||||
# Score each match by the date signal it carries and keep only those
|
||||
# with a real signal, rejecting bare weekday/month-abbreviation false
|
||||
# positives ("we"/"me"/"did"). Prefer the strongest match, breaking ties
|
||||
# by longest span, so an explicit date ("in May", "2026-06-10") always
|
||||
# beats an earlier weak word regardless of position. See issue #2768.
|
||||
scored_results = [
|
||||
(_date_match_score(text), len(text), date)
|
||||
for text, date in results
|
||||
if (text.lower() not in false_positives or len(text) > 3)
|
||||
and not is_embedded_cjk_dateparser_match(query, text)
|
||||
if not is_embedded_cjk_dateparser_match(query, text)
|
||||
]
|
||||
scored_results = [entry for entry in scored_results if entry[0] > 0]
|
||||
|
||||
if not valid_results:
|
||||
if not scored_results:
|
||||
return QueryAnalysis(temporal_constraint=None)
|
||||
|
||||
# Use the first valid date found
|
||||
_, parsed_date = valid_results[0]
|
||||
# Highest signal score wins; ties broken by the longest matched span.
|
||||
_, _, parsed_date = max(scored_results, key=lambda entry: (entry[0], entry[1]))
|
||||
|
||||
# Create constraint for single day
|
||||
start_date = parsed_date.replace(hour=0, minute=0, second=0, microsecond=0)
|
||||
|
||||
@@ -49,6 +49,11 @@ logger = logging.getLogger(__name__)
|
||||
|
||||
DEFAULT_MAX_ITERATIONS = 10
|
||||
|
||||
# Fallback answer when the LLM returns nothing usable. Consumers that need to
|
||||
# tell a real answer from this placeholder (e.g. refresh outcome metadata's
|
||||
# populated_content) compare against this constant rather than the literal.
|
||||
NO_ANSWER_TEXT = "No answer provided."
|
||||
|
||||
|
||||
def _normalize_tool_name(name: str) -> str:
|
||||
"""Normalize tool name from various LLM output formats.
|
||||
@@ -224,6 +229,7 @@ async def _generate_structured_output(
|
||||
response_schema: dict,
|
||||
llm_config: "LLMProvider",
|
||||
reflect_id: str,
|
||||
max_tokens: int | None = None,
|
||||
) -> StructuredOutputResult:
|
||||
"""Generate structured output from an answer using the provided JSON schema.
|
||||
|
||||
@@ -232,6 +238,10 @@ async def _generate_structured_output(
|
||||
response_schema: JSON Schema for the expected output structure
|
||||
llm_config: LLM provider for making the extraction call
|
||||
reflect_id: Reflect ID for logging
|
||||
max_tokens: Output-token budget for the extraction call, mirroring the
|
||||
plain reflect calls (omitted when None); without it, reasoning /
|
||||
preamble models can exhaust the provider default before emitting any
|
||||
JSON (finish_reason=length, empty content -> issue #2431)
|
||||
|
||||
Returns:
|
||||
A StructuredOutputResult carrying the structured output (None if
|
||||
@@ -322,6 +332,7 @@ OUTPUT:"""
|
||||
],
|
||||
response_format=DynamicModel,
|
||||
scope="reflect_structured",
|
||||
max_completion_tokens=max_tokens,
|
||||
max_retries=1,
|
||||
initial_backoff=0.25,
|
||||
max_backoff=1.0,
|
||||
@@ -413,7 +424,104 @@ def _all_mental_models_are_usable_and_fresh(tool_output: dict[str, Any]) -> bool
|
||||
return True
|
||||
|
||||
|
||||
# Detached cache-teardown tasks. asyncio holds only weak references to tasks, so
|
||||
# a fire-and-forget task can be garbage-collected mid-flight — keep a strong
|
||||
# reference here until it finishes.
|
||||
_cache_cleanup_tasks: set[asyncio.Task] = set()
|
||||
|
||||
|
||||
def _spawn_cache_cleanup(
|
||||
provider_impl: Any,
|
||||
session_id: str,
|
||||
cache_tasks: list[asyncio.Task],
|
||||
reflect_id: str,
|
||||
) -> None:
|
||||
"""Delete a reflect's ephemeral context caches in the background.
|
||||
|
||||
The per-reflect caches are dead the moment the reflect returns — nothing ever
|
||||
reuses them — so the caller must not wait on teardown: draining the in-flight
|
||||
create plus the delete round-trips would add latency to every single answer.
|
||||
Detach it instead. The short cache TTL is the backstop if the process dies
|
||||
before the task runs.
|
||||
"""
|
||||
|
||||
async def _cleanup() -> None:
|
||||
try:
|
||||
# Let any overlapped create land first, so its cache is registered in
|
||||
# the session and actually gets deleted rather than lingering to TTL.
|
||||
if cache_tasks:
|
||||
await asyncio.gather(*cache_tasks, return_exceptions=True)
|
||||
await provider_impl.delete_cache_session(session_id)
|
||||
except Exception:
|
||||
logger.debug("[REFLECT %s] cache session teardown failed (will age out on TTL)", reflect_id)
|
||||
|
||||
try:
|
||||
task = asyncio.create_task(_cleanup())
|
||||
except RuntimeError:
|
||||
# No running loop to detach onto (not expected in the server); TTL cleans up.
|
||||
return
|
||||
_cache_cleanup_tasks.add(task)
|
||||
task.add_done_callback(_cache_cleanup_tasks.discard)
|
||||
|
||||
|
||||
async def run_reflect_agent(
|
||||
llm_config: "LLMProvider",
|
||||
bank_id: str,
|
||||
query: str,
|
||||
bank_profile: dict[str, Any],
|
||||
search_mental_models_fn: Callable[[str, int], Awaitable[dict[str, Any]]],
|
||||
search_observations_fn: Callable[[str, int], Awaitable[dict[str, Any]]],
|
||||
recall_fn: Callable[[str, int, int], Awaitable[dict[str, Any]]],
|
||||
expand_fn: Callable[[list[str], str], Awaitable[dict[str, Any]]],
|
||||
**kwargs: Any,
|
||||
) -> ReflectAgentResult:
|
||||
"""Public entrypoint: runs the agent loop and tears down any per-step context
|
||||
caches it created.
|
||||
|
||||
The step-by-step caches (Gemini ``CachedContent``) are ephemeral — scoped to
|
||||
exactly one reflect and never reused after it — so teardown is scheduled on
|
||||
every exit path (answer, error, cancellation) but runs **detached**: the
|
||||
caller gets its answer without waiting on the delete round-trips. The short
|
||||
cache TTL is the backstop if the teardown never runs; the delete is
|
||||
best-effort and never allowed to fail a reflect.
|
||||
"""
|
||||
reflect_id = f"{bank_id[:8]}-{int(time.time() * 1000) % 100000}"
|
||||
provider_impl = getattr(llm_config, "_provider_impl", None)
|
||||
# Reflect step-by-step caching needs the provider to support it AND the
|
||||
# dedicated reflect flag (on by default; distinct from the global prompt-cache
|
||||
# switch so it can be turned off for reflect alone).
|
||||
incremental_caching = (
|
||||
provider_impl is not None
|
||||
and provider_impl.supports_incremental_prompt_cache()
|
||||
and get_config().reflect_prompt_cache_enabled
|
||||
)
|
||||
cache_session_id = f"reflect:{reflect_id}"
|
||||
# In-flight cache-create tasks (scheduled to overlap tool execution). Awaited
|
||||
# before teardown so every created cache is tracked and deleted — no orphans.
|
||||
cache_tasks: list[asyncio.Task] = []
|
||||
try:
|
||||
return await _run_reflect_agent_inner(
|
||||
llm_config,
|
||||
bank_id,
|
||||
query,
|
||||
bank_profile,
|
||||
search_mental_models_fn,
|
||||
search_observations_fn,
|
||||
recall_fn,
|
||||
expand_fn,
|
||||
reflect_id=reflect_id,
|
||||
provider_impl=provider_impl,
|
||||
incremental_caching=incremental_caching,
|
||||
cache_session_id=cache_session_id,
|
||||
cache_tasks=cache_tasks,
|
||||
**kwargs,
|
||||
)
|
||||
finally:
|
||||
if incremental_caching and provider_impl is not None:
|
||||
_spawn_cache_cleanup(provider_impl, cache_session_id, cache_tasks, reflect_id)
|
||||
|
||||
|
||||
async def _run_reflect_agent_inner(
|
||||
llm_config: "LLMProvider",
|
||||
bank_id: str,
|
||||
query: str,
|
||||
@@ -434,6 +542,12 @@ async def run_reflect_agent(
|
||||
max_context_tokens: int = 100_000,
|
||||
llm_output_language: str | None = None,
|
||||
cancel_check: Callable[[], None] | None = None,
|
||||
*,
|
||||
reflect_id: str,
|
||||
provider_impl: Any,
|
||||
incremental_caching: bool,
|
||||
cache_session_id: str,
|
||||
cache_tasks: list[asyncio.Task],
|
||||
) -> ReflectAgentResult:
|
||||
"""
|
||||
Execute the reflect agent loop using native tool calling.
|
||||
@@ -461,7 +575,6 @@ async def run_reflect_agent(
|
||||
Returns:
|
||||
ReflectAgentResult with final answer and metadata
|
||||
"""
|
||||
reflect_id = f"{bank_id[:8]}-{int(time.time() * 1000) % 100000}"
|
||||
start_time = time.time()
|
||||
|
||||
# Build directives_applied for the trace
|
||||
@@ -498,27 +611,68 @@ async def run_reflect_agent(
|
||||
{"role": "user", "content": query},
|
||||
]
|
||||
|
||||
# Opt into context caching for the agentic tool loop. The system
|
||||
# prompt and tool definitions are stable for the duration of this
|
||||
# reflect call (and across reflects against the same bank), so
|
||||
# caching them once and reusing across every iteration of the loop
|
||||
# collapses the dominant input cost — the prefix repeated on every
|
||||
# turn. ``get_or_create_cached_prefix`` returns None when caching is
|
||||
# disabled, unsupported, or the prefix is too small; the
|
||||
# ``call_with_tools`` invocation below transparently falls back to
|
||||
# the uncached path in that case.
|
||||
cached_prefix_name: str | None = None
|
||||
provider_impl = getattr(llm_config, "_provider_impl", None)
|
||||
if provider_impl is not None and provider_impl.supports_prompt_caching():
|
||||
# Step-by-step context caching for the agentic tool loop.
|
||||
#
|
||||
# Caching only the static system+tools prefix wins little here: it's dwarfed
|
||||
# by the tool results (recall/observations) that get re-sent on every turn.
|
||||
# Instead we roll a cache forward one step at a time — after each turn the
|
||||
# cache is extended to cover that turn's FULL input, so the next ``auto`` turn
|
||||
# reuses the entire prior conversation at the cached rate and sends only its
|
||||
# own new tool results as the delta. Each new tool payload is therefore billed
|
||||
# at full price exactly once (the turn it's produced), then cached thereafter.
|
||||
#
|
||||
# The cache create for turn N+1 covers turn N's input, which is fully known the
|
||||
# moment turn N's LLM call returns — so we kick it off as a background task that
|
||||
# runs CONCURRENTLY with turn N's tool execution (``_schedule_cache``) and only
|
||||
# await it (``_resolve_pending_cache``) right before the next ``auto`` call,
|
||||
# hiding the create latency behind work we'd do anyway.
|
||||
#
|
||||
# ``rolling_cache_boundary`` is the number of leading ``messages`` baked into
|
||||
# the adopted ``rolling_cache_name``. ``incremental_caching`` is False for
|
||||
# providers/config without explicit caching, so every branch below is a no-op.
|
||||
rolling_cache_name: str | None = None
|
||||
rolling_cache_boundary = 0
|
||||
pending_cache_task: asyncio.Task | None = None
|
||||
pending_cache_boundary = 0
|
||||
|
||||
async def _resolve_pending_cache() -> None:
|
||||
"""Adopt the overlapped next-cache once it's ready as the rolling cache.
|
||||
|
||||
Best-effort: a failed/``None`` create just leaves the previous (smaller)
|
||||
cache in place, so the next call sends a larger delta but stays correct.
|
||||
"""
|
||||
nonlocal rolling_cache_name, rolling_cache_boundary, pending_cache_task
|
||||
if pending_cache_task is None:
|
||||
return
|
||||
task = pending_cache_task
|
||||
pending_cache_task = None
|
||||
try:
|
||||
cached_prefix_name = await provider_impl.get_or_create_cached_prefix(
|
||||
system_instruction=system_prompt,
|
||||
tools=tools,
|
||||
new_name = await task
|
||||
except Exception:
|
||||
new_name = None
|
||||
if new_name is not None:
|
||||
rolling_cache_name = new_name
|
||||
rolling_cache_boundary = pending_cache_boundary
|
||||
|
||||
def _schedule_cache(upto: int) -> None:
|
||||
"""Start building the cache covering ``messages[:upto]`` in the background
|
||||
so it overlaps the tool execution that follows this turn."""
|
||||
nonlocal pending_cache_task, pending_cache_boundary
|
||||
# ``messages[:upto]`` is snapshotted now, so appends during tool execution
|
||||
# can't change what gets cached. ``ensure_future`` raises if the provider
|
||||
# didn't return a coroutine (e.g. a test double) — caching is a soft
|
||||
# optimisation and must never break a reflect, so swallow and skip.
|
||||
try:
|
||||
task = asyncio.ensure_future(
|
||||
provider_impl.create_incremental_cache(
|
||||
session_id=cache_session_id, messages=messages[:upto], tools=tools
|
||||
)
|
||||
)
|
||||
except Exception:
|
||||
# Caching is a soft optimisation; never let a cache-side
|
||||
# error block a reflect.
|
||||
cached_prefix_name = None
|
||||
return
|
||||
pending_cache_boundary = upto
|
||||
pending_cache_task = task
|
||||
cache_tasks.append(task)
|
||||
|
||||
# Tracking
|
||||
total_tools_called = 0
|
||||
@@ -640,7 +794,7 @@ async def run_reflect_agent(
|
||||
# Generate structured output if schema provided
|
||||
structured_output = None
|
||||
if response_schema and answer:
|
||||
struct = await _generate_structured_output(answer, response_schema, llm_config, reflect_id)
|
||||
struct = await _generate_structured_output(answer, response_schema, llm_config, reflect_id, max_tokens)
|
||||
structured_output = struct.structured_output
|
||||
total_input_tokens += struct.input_tokens
|
||||
total_output_tokens += struct.output_tokens
|
||||
@@ -704,7 +858,7 @@ async def run_reflect_agent(
|
||||
|
||||
structured_output = None
|
||||
if response_schema and answer:
|
||||
struct = await _generate_structured_output(answer, response_schema, llm_config, reflect_id)
|
||||
struct = await _generate_structured_output(answer, response_schema, llm_config, reflect_id, max_tokens)
|
||||
structured_output = struct.structured_output
|
||||
total_input_tokens += struct.input_tokens
|
||||
total_output_tokens += struct.output_tokens
|
||||
@@ -745,6 +899,29 @@ async def run_reflect_agent(
|
||||
else:
|
||||
iter_tool_choice = "auto"
|
||||
|
||||
# Will the NEXT turn be an ``auto`` turn (the only kind that references a
|
||||
# cache)? The cache we schedule this turn covers this turn's input and is
|
||||
# used by the next turn, so we only bother building it when the next turn
|
||||
# can use it — skipping the wasted creates between two forced turns.
|
||||
next_iter = iteration + 1
|
||||
if stop_forcing_from_iteration is not None and next_iter >= stop_forcing_from_iteration:
|
||||
next_is_auto = True
|
||||
elif next_iter < len(forced_sequence):
|
||||
next_is_auto = False
|
||||
else:
|
||||
next_is_auto = True
|
||||
|
||||
# Before an ``auto`` turn, adopt the cache that was being built in the
|
||||
# background during the previous turn's tool execution. It covers that
|
||||
# turn's full input, so THIS call reuses the entire prior conversation at
|
||||
# the cached rate and sends only the turns appended since. Forced turns
|
||||
# can't use a cache (Gemini rejects ``cached_content`` + ``tool_config``),
|
||||
# but the cache still advances underneath them, so the first ``auto`` turn
|
||||
# inherits a cache covering all the forced results.
|
||||
if incremental_caching and iter_tool_choice == "auto":
|
||||
await _resolve_pending_cache()
|
||||
|
||||
call_msg_count = len(messages)
|
||||
try:
|
||||
ct_kwargs: dict[str, Any] = dict(
|
||||
messages=messages,
|
||||
@@ -752,15 +929,9 @@ async def run_reflect_agent(
|
||||
scope="reflect_tool_call",
|
||||
tool_choice=iter_tool_choice,
|
||||
)
|
||||
# Gemini rejects ``cached_content`` alongside a per-request
|
||||
# ``tool_config`` (forced tool choice): "CachedContent can not be used
|
||||
# with GenerateContent request setting system_instruction, tools or
|
||||
# tool_config." The forced-sequence iterations set tool_config, so only
|
||||
# the ``auto`` iterations can reference the cache; forced iterations send
|
||||
# the prefix inline. The cache (tools + system prompt) is identical
|
||||
# either way, so this just limits *which* iterations are billed cached.
|
||||
if cached_prefix_name is not None and iter_tool_choice == "auto":
|
||||
ct_kwargs["cached_prefix"] = cached_prefix_name
|
||||
if incremental_caching and iter_tool_choice == "auto" and rolling_cache_name is not None:
|
||||
ct_kwargs["cached_prefix"] = rolling_cache_name
|
||||
ct_kwargs["cached_prefix_message_count"] = rolling_cache_boundary
|
||||
result = await llm_config.call_with_tools(**ct_kwargs)
|
||||
llm_duration = int((time.time() - llm_start) * 1000)
|
||||
consecutive_errors = 0
|
||||
@@ -831,7 +1002,7 @@ async def run_reflect_agent(
|
||||
# Generate structured output if schema provided
|
||||
structured_output = None
|
||||
if response_schema and answer:
|
||||
struct = await _generate_structured_output(answer, response_schema, llm_config, reflect_id)
|
||||
struct = await _generate_structured_output(answer, response_schema, llm_config, reflect_id, max_tokens)
|
||||
structured_output = struct.structured_output
|
||||
total_input_tokens += struct.input_tokens
|
||||
total_output_tokens += struct.output_tokens
|
||||
@@ -908,7 +1079,9 @@ async def run_reflect_agent(
|
||||
# Generate structured output if schema provided
|
||||
structured_output = None
|
||||
if response_schema and answer:
|
||||
struct = await _generate_structured_output(answer, response_schema, llm_config, reflect_id)
|
||||
struct = await _generate_structured_output(
|
||||
answer, response_schema, llm_config, reflect_id, max_tokens
|
||||
)
|
||||
structured_output = struct.structured_output
|
||||
total_input_tokens += struct.input_tokens
|
||||
total_output_tokens += struct.output_tokens
|
||||
@@ -963,7 +1136,7 @@ async def run_reflect_agent(
|
||||
# Generate structured output if schema provided
|
||||
structured_output = None
|
||||
if response_schema and answer:
|
||||
struct = await _generate_structured_output(answer, response_schema, llm_config, reflect_id)
|
||||
struct = await _generate_structured_output(answer, response_schema, llm_config, reflect_id, max_tokens)
|
||||
structured_output = struct.structured_output
|
||||
total_input_tokens += struct.input_tokens
|
||||
total_output_tokens += struct.output_tokens
|
||||
@@ -1035,6 +1208,7 @@ async def run_reflect_agent(
|
||||
directives_applied=directives_applied,
|
||||
llm_config=llm_config,
|
||||
response_schema=response_schema,
|
||||
max_tokens=max_tokens,
|
||||
)
|
||||
|
||||
# Execute other tools in parallel (exclude done tool in all its format variants)
|
||||
@@ -1078,6 +1252,16 @@ async def run_reflect_agent(
|
||||
|
||||
other_tools = allowed_tools
|
||||
|
||||
# Kick off the next-turn cache (covering THIS call's input) so it
|
||||
# builds concurrently with the tool execution below — hiding the
|
||||
# create latency. Only schedule when the next turn is ``auto`` (the
|
||||
# only kind that references it); the next turn's pre-call resolve then
|
||||
# adopts it. Resolve any prior in-flight create first so we don't drop
|
||||
# its handle.
|
||||
if incremental_caching and next_is_auto:
|
||||
await _resolve_pending_cache()
|
||||
_schedule_cache(call_msg_count)
|
||||
|
||||
# Execute tools in parallel
|
||||
tool_tasks = [
|
||||
_execute_tool_with_timing(
|
||||
@@ -1244,6 +1428,7 @@ async def _process_done_tool(
|
||||
directives_applied: list[DirectiveInfo],
|
||||
llm_config: "LLMProvider | None" = None,
|
||||
response_schema: dict | None = None,
|
||||
max_tokens: int | None = None,
|
||||
) -> ReflectAgentResult:
|
||||
"""Process the done tool call and return the result."""
|
||||
args = done_call.arguments
|
||||
@@ -1252,7 +1437,46 @@ async def _process_done_tool(
|
||||
raw_answer = args.get("answer", "").strip()
|
||||
answer = _clean_done_answer(raw_answer) if raw_answer else ""
|
||||
if not answer:
|
||||
answer = "No answer provided."
|
||||
answer = NO_ANSWER_TEXT
|
||||
|
||||
final_usage = usage
|
||||
if llm_config and max_tokens is not None and count_cl100k_tokens(answer) > max_tokens:
|
||||
rewrite_start = time.time()
|
||||
rewritten, rewrite_usage = await llm_config.call(
|
||||
messages=[
|
||||
{
|
||||
"role": "system",
|
||||
"content": (
|
||||
"Rewrite the user's text so it fits within the requested token budget. "
|
||||
"Preserve the key facts and structure; drop lower-priority detail. "
|
||||
"Respond with the rewritten text only, no preamble."
|
||||
),
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": f"Target budget: {max_tokens} tokens.\n\nText to rewrite:\n{answer}",
|
||||
},
|
||||
],
|
||||
scope="reflect",
|
||||
max_completion_tokens=max_tokens,
|
||||
return_usage=True,
|
||||
)
|
||||
answer = _clean_answer_text(rewritten.strip())
|
||||
final_usage = TokenUsageSummary(
|
||||
input_tokens=usage.input_tokens + rewrite_usage.input_tokens,
|
||||
output_tokens=usage.output_tokens + rewrite_usage.output_tokens,
|
||||
total_tokens=usage.total_tokens + rewrite_usage.input_tokens + rewrite_usage.output_tokens,
|
||||
cached_tokens=usage.cached_tokens + (getattr(rewrite_usage, "cached_tokens", 0) or 0),
|
||||
thoughts_tokens=usage.thoughts_tokens + (getattr(rewrite_usage, "thoughts_tokens", 0) or 0),
|
||||
)
|
||||
llm_trace.append(
|
||||
LLMCall(
|
||||
scope="final_rewrite",
|
||||
duration_ms=int((time.time() - rewrite_start) * 1000),
|
||||
input_tokens=rewrite_usage.input_tokens,
|
||||
output_tokens=rewrite_usage.output_tokens,
|
||||
)
|
||||
)
|
||||
|
||||
# Validate IDs (only include IDs that were actually retrieved)
|
||||
used_memory_ids = [mid for mid in (args.get("memory_ids") or []) if mid in available_memory_ids]
|
||||
@@ -1261,17 +1485,16 @@ async def _process_done_tool(
|
||||
|
||||
# Generate structured output if schema provided
|
||||
structured_output = None
|
||||
final_usage = usage
|
||||
if response_schema and llm_config and answer:
|
||||
struct = await _generate_structured_output(answer, response_schema, llm_config, reflect_id)
|
||||
struct = await _generate_structured_output(answer, response_schema, llm_config, reflect_id, max_tokens)
|
||||
structured_output = struct.structured_output
|
||||
# Add structured output tokens to usage
|
||||
final_usage = TokenUsageSummary(
|
||||
input_tokens=usage.input_tokens + struct.input_tokens,
|
||||
output_tokens=usage.output_tokens + struct.output_tokens,
|
||||
total_tokens=usage.total_tokens + struct.input_tokens + struct.output_tokens,
|
||||
cached_tokens=usage.cached_tokens + struct.cached_tokens,
|
||||
thoughts_tokens=usage.thoughts_tokens + struct.thoughts_tokens,
|
||||
input_tokens=final_usage.input_tokens + struct.input_tokens,
|
||||
output_tokens=final_usage.output_tokens + struct.output_tokens,
|
||||
total_tokens=final_usage.total_tokens + struct.input_tokens + struct.output_tokens,
|
||||
cached_tokens=final_usage.cached_tokens + struct.cached_tokens,
|
||||
thoughts_tokens=final_usage.thoughts_tokens + struct.thoughts_tokens,
|
||||
)
|
||||
|
||||
log_completion(answer, iterations)
|
||||
@@ -1386,22 +1609,35 @@ async def _execute_tool(
|
||||
query = args.get("query")
|
||||
if not query:
|
||||
return {"error": "search_mental_models requires a query parameter"}
|
||||
max_results = int(args.get("max_results") or 5)
|
||||
max_results, error = _parse_tool_int_arg_or_error(args, "max_results", default=5)
|
||||
if error:
|
||||
return {"error": error}
|
||||
return await search_mental_models_fn(query, max_results)
|
||||
|
||||
elif tool_name == "search_observations":
|
||||
query = args.get("query")
|
||||
if not query:
|
||||
return {"error": "search_observations requires a query parameter"}
|
||||
max_tokens = max(int(args.get("max_tokens") or 5000), 1000) # Default 5000, min 1000
|
||||
max_tokens, error = _parse_tool_int_arg_or_error(args, "max_tokens", default=5000, minimum=1000)
|
||||
if error:
|
||||
return {"error": error}
|
||||
return await search_observations_fn(query, max_tokens)
|
||||
|
||||
elif tool_name == "recall":
|
||||
query = args.get("query")
|
||||
if not query:
|
||||
return {"error": "recall requires a query parameter"}
|
||||
max_tokens = max(int(args.get("max_tokens") or 2048), 1000) # Default 2048, min 1000
|
||||
max_chunk_tokens = max(int(args.get("max_chunk_tokens") or 1000), 1000) # Always enabled, min 1000
|
||||
max_tokens, error = _parse_tool_int_arg_or_error(args, "max_tokens", default=2048, minimum=1000)
|
||||
if error:
|
||||
return {"error": error}
|
||||
max_chunk_tokens, error = _parse_tool_int_arg_or_error(
|
||||
args,
|
||||
"max_chunk_tokens",
|
||||
default=1000,
|
||||
minimum=1000,
|
||||
)
|
||||
if error:
|
||||
return {"error": error}
|
||||
return await recall_fn(query, max_tokens, max_chunk_tokens)
|
||||
|
||||
elif tool_name == "expand":
|
||||
@@ -1415,23 +1651,63 @@ async def _execute_tool(
|
||||
return {"error": f"Unknown tool: {tool_name}"}
|
||||
|
||||
|
||||
_NULLISH_TOOL_INT_STRINGS = {"", "none", "null"}
|
||||
|
||||
|
||||
def _parse_tool_int_arg(args: dict[str, Any], key: str, *, default: int, minimum: int | None = None) -> int:
|
||||
raw_value = args.get(key)
|
||||
if not raw_value:
|
||||
value = default
|
||||
elif isinstance(raw_value, str) and raw_value.strip().lower() in _NULLISH_TOOL_INT_STRINGS:
|
||||
value = default
|
||||
else:
|
||||
value = int(raw_value)
|
||||
if minimum is None:
|
||||
return value
|
||||
return max(value, minimum)
|
||||
|
||||
|
||||
def _parse_tool_int_arg_or_error(
|
||||
args: dict[str, Any],
|
||||
key: str,
|
||||
*,
|
||||
default: int,
|
||||
minimum: int | None = None,
|
||||
) -> tuple[int, str | None]:
|
||||
try:
|
||||
return _parse_tool_int_arg(args, key, default=default, minimum=minimum), None
|
||||
except (OverflowError, TypeError, ValueError):
|
||||
return default, f"{key} must be an integer or null-like value"
|
||||
|
||||
|
||||
def _summarize_tool_int_arg(args: dict[str, Any], key: str, *, default: int, minimum: int | None = None) -> str:
|
||||
try:
|
||||
return str(_parse_tool_int_arg(args, key, default=default, minimum=minimum))
|
||||
except (OverflowError, TypeError, ValueError):
|
||||
return f"invalid:{args.get(key)!r}"
|
||||
|
||||
|
||||
def _summarize_tool_query(args: dict[str, Any]) -> str:
|
||||
query = args.get("query") or ""
|
||||
if not isinstance(query, str):
|
||||
query = str(query)
|
||||
return f"'{query[:30]}...'" if len(query) > 30 else f"'{query}'"
|
||||
|
||||
|
||||
def _summarize_input(tool_name: str, args: dict[str, Any]) -> str:
|
||||
"""Create a summary of tool input for logging, showing all params."""
|
||||
if tool_name == "search_mental_models":
|
||||
query = args.get("query", "")
|
||||
query_preview = f"'{query[:30]}...'" if len(query) > 30 else f"'{query}'"
|
||||
max_results = int(args.get("max_results") or 5)
|
||||
query_preview = _summarize_tool_query(args)
|
||||
max_results = _summarize_tool_int_arg(args, "max_results", default=5)
|
||||
return f"(query={query_preview}, max_results={max_results})"
|
||||
elif tool_name == "search_observations":
|
||||
query = args.get("query", "")
|
||||
query_preview = f"'{query[:30]}...'" if len(query) > 30 else f"'{query}'"
|
||||
max_tokens = max(int(args.get("max_tokens") or 5000), 1000)
|
||||
query_preview = _summarize_tool_query(args)
|
||||
max_tokens = _summarize_tool_int_arg(args, "max_tokens", default=5000, minimum=1000)
|
||||
return f"(query={query_preview}, max_tokens={max_tokens})"
|
||||
elif tool_name == "recall":
|
||||
query = args.get("query", "")
|
||||
query_preview = f"'{query[:30]}...'" if len(query) > 30 else f"'{query}'"
|
||||
max_tokens = max(int(args.get("max_tokens") or 2048), 1000)
|
||||
max_chunk_tokens = max(int(args.get("max_chunk_tokens") or 1000), 1000)
|
||||
query_preview = _summarize_tool_query(args)
|
||||
max_tokens = _summarize_tool_int_arg(args, "max_tokens", default=2048, minimum=1000)
|
||||
max_chunk_tokens = _summarize_tool_int_arg(args, "max_chunk_tokens", default=1000, minimum=1000)
|
||||
return f"(query={query_preview}, max_tokens={max_tokens}, max_chunk_tokens={max_chunk_tokens})"
|
||||
elif tool_name == "expand":
|
||||
memory_ids = args.get("memory_ids", [])
|
||||
|
||||
@@ -177,20 +177,17 @@ _HEADING_RX = re.compile(r"^(#{1,6})\s+(.+?)\s*$")
|
||||
_BULLET_RX = re.compile(r"^\s*[-*+]\s+(.*)$")
|
||||
_ORDERED_RX = re.compile(r"^\s*\d+[.)]\s+(.*)$")
|
||||
_FENCE_RX = re.compile(r"^```([A-Za-z0-9_+-]*)\s*$")
|
||||
|
||||
|
||||
def _strip_separators(lines: list[str]) -> list[str]:
|
||||
"""Drop horizontal-rule lines (`---`, `***`) used as section separators.
|
||||
|
||||
Our renderer never emits these, but LLM output frequently includes them
|
||||
between sections; treating them as blank lines avoids parsing them as
|
||||
paragraphs.
|
||||
"""
|
||||
return ["" if re.fullmatch(r"\s*([-*_])\1{2,}\s*", line) else line for line in lines]
|
||||
_SEPARATOR_RX = re.compile(r"\s*([-*_])\1{2,}\s*")
|
||||
|
||||
|
||||
def _split_blocks(lines: list[str]) -> list[list[str]]:
|
||||
"""Group consecutive non-blank lines into block chunks."""
|
||||
"""Group consecutive non-blank lines into block chunks.
|
||||
|
||||
Horizontal-rule lines (`---`, `***`) count as blank. Our renderer never
|
||||
emits these, but LLM output frequently includes them between sections;
|
||||
treating them as blank avoids parsing them as paragraphs. Inside a fence
|
||||
they are code, not a separator, so they are kept verbatim.
|
||||
"""
|
||||
chunks: list[list[str]] = []
|
||||
current: list[str] = []
|
||||
in_fence = False
|
||||
@@ -202,7 +199,7 @@ def _split_blocks(lines: list[str]) -> list[list[str]]:
|
||||
if in_fence:
|
||||
current.append(line)
|
||||
continue
|
||||
if line.strip() == "":
|
||||
if line.strip() == "" or _SEPARATOR_RX.fullmatch(line):
|
||||
if current:
|
||||
chunks.append(current)
|
||||
current = []
|
||||
@@ -250,8 +247,7 @@ def parse_markdown(markdown: str) -> StructuredDocument:
|
||||
so we never silently drop user content. Section IDs are unique slugs of
|
||||
their headings.
|
||||
"""
|
||||
raw_lines = (markdown or "").splitlines()
|
||||
lines = _strip_separators(raw_lines)
|
||||
lines = (markdown or "").splitlines()
|
||||
|
||||
sections: list[Section] = []
|
||||
used_ids: set[str] = set()
|
||||
|
||||
@@ -328,18 +328,21 @@ async def tool_expand(
|
||||
if not memory_ids:
|
||||
return {"error": "memory_ids is required and must not be empty"}
|
||||
|
||||
# Validate and convert UUIDs
|
||||
valid_uuids: list[uuid.UUID] = []
|
||||
# Validate and convert UUIDs. Each id keeps a handle on its own UUID: a list of
|
||||
# only the valid ones no longer lines up with memory_ids once one id is invalid.
|
||||
uuid_by_id: dict[str, uuid.UUID] = {}
|
||||
errors: dict[str, str] = {}
|
||||
for mid in memory_ids:
|
||||
try:
|
||||
valid_uuids.append(uuid.UUID(mid))
|
||||
uuid_by_id[mid] = uuid.UUID(mid)
|
||||
except ValueError:
|
||||
errors[mid] = f"Invalid memory_id format: {mid}"
|
||||
|
||||
if not valid_uuids:
|
||||
if not uuid_by_id:
|
||||
return {"error": "No valid memory IDs provided", "details": errors}
|
||||
|
||||
valid_uuids = list(uuid_by_id.values())
|
||||
|
||||
# Batch fetch all memory units
|
||||
memories = await conn.fetch(
|
||||
f"""
|
||||
@@ -395,12 +398,12 @@ async def tool_expand(
|
||||
|
||||
# Build results
|
||||
results: list[dict[str, Any]] = []
|
||||
for mid, mem_uuid in zip(memory_ids, valid_uuids):
|
||||
for mid in memory_ids:
|
||||
if mid in errors:
|
||||
results.append({"memory_id": mid, "error": errors[mid]})
|
||||
continue
|
||||
|
||||
memory = memory_map.get(mem_uuid)
|
||||
memory = memory_map.get(uuid_by_id[mid])
|
||||
if not memory:
|
||||
results.append({"memory_id": mid, "error": f"Memory not found: {mid}"})
|
||||
continue
|
||||
|
||||
@@ -255,13 +255,20 @@ class MemoryFact(BaseModel):
|
||||
@field_validator("metadata", mode="before")
|
||||
@classmethod
|
||||
def parse_metadata(cls, v: Any) -> dict[str, str] | None:
|
||||
"""Parse metadata from JSON string if needed (asyncpg may return JSONB as str)."""
|
||||
"""Parse metadata from JSON string if needed (asyncpg may return JSONB as str).
|
||||
|
||||
Also coerces non-string dict values (e.g., integer IDs stored in JSONB)
|
||||
to strings, preventing ValidationError when consolidation encounters
|
||||
metadata like {"original_id": 348} instead of {"original_id": "348"}.
|
||||
"""
|
||||
if v is None:
|
||||
return None
|
||||
if isinstance(v, str):
|
||||
import json
|
||||
|
||||
return json.loads(v)
|
||||
v = json.loads(v)
|
||||
if isinstance(v, dict):
|
||||
return {str(k): str(val) for k, val in v.items()}
|
||||
return v
|
||||
|
||||
chunk_id: str | None = Field(
|
||||
|
||||
@@ -189,7 +189,8 @@ async def get_or_create_bank_profile(pool, bank_id: str) -> BankProfileResult:
|
||||
``get_or_create_bank_profile_on_conn`` instead.
|
||||
"""
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
return await get_or_create_bank_profile_on_conn(conn, bank_id, ops=pool.ops)
|
||||
async with conn.transaction():
|
||||
return await get_or_create_bank_profile_on_conn(conn, bank_id, ops=pool.ops)
|
||||
|
||||
|
||||
async def get_or_create_bank_profile_on_conn(conn, bank_id: str, *, ops) -> BankProfileResult:
|
||||
|
||||
@@ -64,8 +64,53 @@ async def delete_chunks_by_ids(conn, chunk_ids: list[str]) -> None:
|
||||
"""
|
||||
if not chunk_ids:
|
||||
return
|
||||
|
||||
# PostgreSQL's FK cascade deletes child memory_links in executor-chosen
|
||||
# order. Concurrent chunk deletes for the same bank can then lock overlapping
|
||||
# memory_links in opposite orders and deadlock. Delete links explicitly in a
|
||||
# total order before deleting chunks so every writer takes row locks the same
|
||||
# way; the FK cascade still handles anything inserted later in this txn.
|
||||
await conn.execute(
|
||||
f"DELETE FROM {fq_table('chunks')} WHERE chunk_id = ANY($1::text[])",
|
||||
f"""
|
||||
WITH target_units AS MATERIALIZED (
|
||||
SELECT id
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE chunk_id = ANY($1::text[])
|
||||
),
|
||||
ordered_links AS MATERIALIZED (
|
||||
SELECT ml.ctid
|
||||
FROM {fq_table("memory_links")} ml
|
||||
WHERE EXISTS (
|
||||
SELECT 1
|
||||
FROM target_units tu
|
||||
WHERE tu.id = ml.from_unit_id OR tu.id = ml.to_unit_id
|
||||
)
|
||||
ORDER BY
|
||||
LEAST(ml.from_unit_id, ml.to_unit_id),
|
||||
GREATEST(ml.from_unit_id, ml.to_unit_id),
|
||||
ml.link_type,
|
||||
COALESCE(ml.entity_id, '00000000-0000-0000-0000-000000000000'::uuid)
|
||||
FOR UPDATE OF ml
|
||||
)
|
||||
DELETE FROM {fq_table("memory_links")} ml
|
||||
USING ordered_links ol
|
||||
WHERE ml.ctid = ol.ctid
|
||||
""",
|
||||
chunk_ids,
|
||||
)
|
||||
await conn.execute(
|
||||
f"""
|
||||
WITH ordered_chunks AS MATERIALIZED (
|
||||
SELECT chunk_id
|
||||
FROM {fq_table("chunks")}
|
||||
WHERE chunk_id = ANY($1::text[])
|
||||
ORDER BY chunk_id
|
||||
FOR UPDATE
|
||||
)
|
||||
DELETE FROM {fq_table("chunks")} c
|
||||
USING ordered_chunks oc
|
||||
WHERE c.chunk_id = oc.chunk_id
|
||||
""",
|
||||
chunk_ids,
|
||||
)
|
||||
|
||||
|
||||
@@ -15,7 +15,7 @@ from typing import Any, Literal, cast
|
||||
from pydantic import BaseModel, ConfigDict, Field, create_model, field_validator
|
||||
|
||||
from ..llm_interface import ProviderRateLimitResetError
|
||||
from ..llm_wrapper import LLMConfig, OutputTooLongError, sanitize_llm_output
|
||||
from ..llm_wrapper import LLMConfig, OutputTooLongError, parse_llm_json, sanitize_llm_output
|
||||
from ..operation_metadata import RetainExtractionErrors
|
||||
from ..response_models import TokenUsage
|
||||
from .entity_labels import (
|
||||
@@ -32,7 +32,7 @@ def _extract_map_entities(
|
||||
entity_obj: dict,
|
||||
fields: dict[str, MapField],
|
||||
prefix: str,
|
||||
validated_entities: "list[Entity]",
|
||||
validated_entities: list[str],
|
||||
existing_texts_lower: set[str],
|
||||
) -> None:
|
||||
"""Recursively extract key:field:value entity strings from a map entity dict."""
|
||||
@@ -59,7 +59,7 @@ def _extract_map_entities(
|
||||
continue
|
||||
label_str = f"{prefix}{field_name}:{v.strip()}"
|
||||
if label_str.lower() not in existing_texts_lower:
|
||||
validated_entities.append(Entity(text=label_str))
|
||||
validated_entities.append(label_str)
|
||||
existing_texts_lower.add(label_str.lower())
|
||||
else:
|
||||
# text or value — single string
|
||||
@@ -67,7 +67,7 @@ def _extract_map_entities(
|
||||
continue
|
||||
label_str = f"{prefix}{field_name}:{field_val.strip()}"
|
||||
if label_str.lower() not in existing_texts_lower:
|
||||
validated_entities.append(Entity(text=label_str))
|
||||
validated_entities.append(label_str)
|
||||
existing_texts_lower.add(label_str.lower())
|
||||
|
||||
|
||||
@@ -114,12 +114,33 @@ def _sanitize_text(text: str | None) -> str | None:
|
||||
return sanitize_llm_output(text)
|
||||
|
||||
|
||||
class Entity(BaseModel):
|
||||
"""An entity extracted from text."""
|
||||
def _coerce_entity_strings(v: Any) -> Any:
|
||||
"""
|
||||
Normalize the LLM's `entities` field to a plain list of strings.
|
||||
|
||||
text: str = Field(
|
||||
description="The specific, named entity as it appears in the fact. Must be a proper noun or specific identifier."
|
||||
)
|
||||
The schema previously asked for `Entity` objects ({"text": "..."}) while the
|
||||
prompt's few-shot examples taught a flat string array. Models that followed
|
||||
the examples literally returned strings, and the entities were silently
|
||||
dropped — none were ever persisted (#2749). The `Entity` wrapper carried no
|
||||
information beyond the string, so it was removed rather than taught to the
|
||||
prompt; the object form is still unwrapped here for models that learned it
|
||||
and for in-flight batch jobs.
|
||||
|
||||
Returns non-list input untouched so pydantic reports the type error itself.
|
||||
"""
|
||||
if v is None:
|
||||
return []
|
||||
if not isinstance(v, list):
|
||||
return v
|
||||
coerced = []
|
||||
for item in v:
|
||||
if isinstance(item, dict):
|
||||
text = item.get("text")
|
||||
if isinstance(text, str):
|
||||
coerced.append(text)
|
||||
else:
|
||||
coerced.append(item)
|
||||
return coerced
|
||||
|
||||
|
||||
class Fact(BaseModel):
|
||||
@@ -144,7 +165,7 @@ class Fact(BaseModel):
|
||||
)
|
||||
|
||||
# Optional structured data
|
||||
entities: list[Entity] | None = None
|
||||
entities: list[str] | None = None
|
||||
causal_relations: list["CausalRelation"] | None = None
|
||||
|
||||
|
||||
@@ -195,7 +216,9 @@ class ExtractedFact(BaseModel):
|
||||
fact_type: Literal["world", "assistant"] = Field(
|
||||
description="'world' = objective/external facts, including user preferences, rules, corrections, and constraints even when stated during a conversation. 'assistant' = actions, experiences, or observations the assistant/agent actually performed."
|
||||
)
|
||||
entities: list[Entity] | None = Field(default=None, description="People, places, concepts")
|
||||
entities: list[str] = Field(
|
||||
default_factory=list, description='People, places, concepts - plain strings, e.g. ["Alice", "Kubernetes"]'
|
||||
)
|
||||
causal_relations: list[FactCausalRelation] | None = Field(
|
||||
default=None, description="Links to previous facts (target_index < this fact's index)"
|
||||
)
|
||||
@@ -203,10 +226,7 @@ class ExtractedFact(BaseModel):
|
||||
@field_validator("entities", mode="before")
|
||||
@classmethod
|
||||
def ensure_entities_list(cls, v):
|
||||
"""Ensure entities is always a list (convert None to empty list)."""
|
||||
if v is None:
|
||||
return []
|
||||
return v
|
||||
return _coerce_entity_strings(v)
|
||||
|
||||
def build_fact_text(self) -> str:
|
||||
"""Combine all dimensions into a single comprehensive fact string."""
|
||||
@@ -232,6 +252,55 @@ class FactExtractionResponse(BaseModel):
|
||||
facts: list[ExtractedFact] = Field(description="List of extracted factual statements")
|
||||
|
||||
|
||||
def _split_chunk_for_output_retry(chunk: str) -> tuple[str, str] | None:
|
||||
"""Split an oversized extraction chunk without corrupting structured input."""
|
||||
stripped = chunk.strip()
|
||||
if len(stripped) <= 1:
|
||||
return None
|
||||
|
||||
try:
|
||||
parsed = json.loads(stripped)
|
||||
except (TypeError, ValueError, json.JSONDecodeError):
|
||||
parsed = None
|
||||
|
||||
if isinstance(parsed, list):
|
||||
if len(parsed) >= 2:
|
||||
mid = len(parsed) // 2
|
||||
return json.dumps(parsed[:mid]), json.dumps(parsed[mid:])
|
||||
|
||||
if len(parsed) == 1 and isinstance(parsed[0], dict):
|
||||
turn = parsed[0]
|
||||
content = turn.get("content")
|
||||
if isinstance(content, str) and len(content) > 1:
|
||||
cut = len(content) // 2
|
||||
first_turn = dict(turn)
|
||||
second_turn = dict(turn)
|
||||
first_turn["content"] = content[:cut]
|
||||
second_turn["content"] = content[cut:]
|
||||
return json.dumps([first_turn]), json.dumps([second_turn])
|
||||
|
||||
return None
|
||||
|
||||
# Split plain text at the midpoint, preferring sentence boundaries nearby.
|
||||
mid_point = len(stripped) // 2
|
||||
search_range = int(len(stripped) * 0.2)
|
||||
search_start = max(0, mid_point - search_range)
|
||||
search_end = min(len(stripped), mid_point + search_range)
|
||||
|
||||
best_split = mid_point
|
||||
for ending in [". ", "! ", "? ", "\n\n"]:
|
||||
pos = stripped.rfind(ending, search_start, search_end)
|
||||
if pos != -1:
|
||||
best_split = pos + len(ending)
|
||||
break
|
||||
|
||||
first_half = stripped[:best_split].strip()
|
||||
second_half = stripped[best_split:].strip()
|
||||
if not first_half or not second_half or first_half == stripped or second_half == stripped:
|
||||
return None
|
||||
return first_half, second_half
|
||||
|
||||
|
||||
class ExtractedFactVerbose(BaseModel):
|
||||
"""A single extracted fact with verbose field descriptions for detailed extraction."""
|
||||
|
||||
@@ -299,9 +368,9 @@ class ExtractedFactVerbose(BaseModel):
|
||||
description="'world' = objective/external facts about the user, other people, events, general knowledge, preferences, rules, corrections, or constraints. 'assistant' = actions, experiences, or observations the assistant/agent actually performed (e.g., 'I changed X', 'I discovered Y')."
|
||||
)
|
||||
|
||||
entities: list[Entity] | None = Field(
|
||||
default=None,
|
||||
description="Named entities, objects, AND abstract concepts from the fact. Include: people names, organizations, places, significant objects (e.g., 'coffee maker', 'car'), AND abstract concepts/themes (e.g., 'friendship', 'career growth', 'loss', 'celebration'). Extract anything that could help link related facts together.",
|
||||
entities: list[str] = Field(
|
||||
default_factory=list,
|
||||
description="Named entities, objects, AND abstract concepts from the fact, as plain strings (e.g. [\"Alice\", \"friendship\"]). Include: people names, organizations, places, significant objects (e.g., 'coffee maker', 'car'), AND abstract concepts/themes (e.g., 'friendship', 'career growth', 'loss', 'celebration'). Extract anything that could help link related facts together.",
|
||||
)
|
||||
|
||||
causal_relations: list[FactCausalRelation] | None = Field(
|
||||
@@ -313,9 +382,7 @@ class ExtractedFactVerbose(BaseModel):
|
||||
@field_validator("entities", mode="before")
|
||||
@classmethod
|
||||
def ensure_entities_list(cls, v):
|
||||
if v is None:
|
||||
return []
|
||||
return v
|
||||
return _coerce_entity_strings(v)
|
||||
|
||||
|
||||
class FactExtractionResponseVerbose(BaseModel):
|
||||
@@ -348,17 +415,15 @@ class ExtractedFactNoCausal(BaseModel):
|
||||
fact_type: Literal["world", "assistant"] = Field(
|
||||
description="'world' = about the user/others, including user preferences, rules, corrections, and constraints. 'assistant' = actions or experiences the assistant/agent actually performed."
|
||||
)
|
||||
entities: list[Entity] | None = Field(
|
||||
default=None,
|
||||
description="Named entities, objects, and concepts from the fact.",
|
||||
entities: list[str] = Field(
|
||||
default_factory=list,
|
||||
description='Named entities, objects, and concepts from the fact, as plain strings (e.g. ["Alice", "Kubernetes"]).',
|
||||
)
|
||||
|
||||
@field_validator("entities", mode="before")
|
||||
@classmethod
|
||||
def ensure_entities_list(cls, v):
|
||||
if v is None:
|
||||
return []
|
||||
return v
|
||||
return _coerce_entity_strings(v)
|
||||
|
||||
|
||||
class FactExtractionResponseNoCausal(BaseModel):
|
||||
@@ -390,14 +455,14 @@ class VerbatimExtractedFact(BaseModel):
|
||||
fact_type: Literal["world", "assistant"] = Field(
|
||||
description="'world' = objective/external facts. 'assistant' = first-person actions, experiences, or observations by the speaker."
|
||||
)
|
||||
entities: list[Entity] | None = Field(default=None, description="People, places, concepts")
|
||||
entities: list[str] = Field(
|
||||
default_factory=list, description='People, places, concepts - plain strings, e.g. ["Alice", "Kubernetes"]'
|
||||
)
|
||||
|
||||
@field_validator("entities", mode="before")
|
||||
@classmethod
|
||||
def ensure_entities_list(cls, v):
|
||||
if v is None:
|
||||
return []
|
||||
return v
|
||||
return _coerce_entity_strings(v)
|
||||
|
||||
|
||||
class VerbatimFactExtractionResponse(BaseModel):
|
||||
@@ -681,6 +746,11 @@ Use "Event Date" from input as reference for relative dates.
|
||||
ENTITIES
|
||||
══════════════════════════════════════════════════════════════════════════
|
||||
|
||||
ALWAYS return "entities" as an array of plain strings — never objects, never null.
|
||||
Correct: entities=["Alice", "Kubernetes", "CKA"]
|
||||
Wrong: entities as an array of objects with a "text" key ← never use this form
|
||||
Use an empty array [] only when the fact truly names nothing.
|
||||
|
||||
Include: people names, organizations, places, key objects, abstract concepts (career, friendship, etc.)
|
||||
Always include "user" when fact is about the user.{examples}"""
|
||||
|
||||
@@ -1098,8 +1168,8 @@ def _build_extraction_prompt_and_schema(config) -> tuple[str, type]:
|
||||
}
|
||||
if not free_form_entities:
|
||||
dynamic_fields["entities"] = (
|
||||
list[Entity] | None,
|
||||
Field(default=None, description="Leave empty — labels-only mode"),
|
||||
list[str],
|
||||
Field(default_factory=list, 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)
|
||||
@@ -1236,6 +1306,15 @@ def _build_request_body(llm_config, config, prompt: str, user_message: str, resp
|
||||
return request_body
|
||||
|
||||
|
||||
def _coerce_fact_response(response: Any) -> dict[str, Any] | None:
|
||||
"""Accept the schema wrapper, or a recoverable top-level facts array."""
|
||||
if isinstance(response, dict):
|
||||
return response
|
||||
if isinstance(response, list) and all(isinstance(item, dict) for item in response):
|
||||
return {"facts": response}
|
||||
return None
|
||||
|
||||
|
||||
async def _extract_facts_from_chunk(
|
||||
chunk: str,
|
||||
chunk_index: int,
|
||||
@@ -1341,7 +1420,8 @@ async def _extract_facts_from_chunk(
|
||||
has_malformed_facts = False
|
||||
|
||||
# Handle malformed LLM responses
|
||||
if not isinstance(extraction_response_json, dict):
|
||||
coerced_response_json = _coerce_fact_response(extraction_response_json)
|
||||
if coerced_response_json is None:
|
||||
if attempt < llm_max_retries - 1:
|
||||
logger.warning(
|
||||
f"LLM returned non-dict JSON on attempt {attempt + 1}/{llm_max_retries}: {type(extraction_response_json).__name__}. Retrying..."
|
||||
@@ -1356,6 +1436,7 @@ async def _extract_facts_from_chunk(
|
||||
f"Fact extraction failed: LLM returned non-dict JSON after {llm_max_retries} attempts "
|
||||
f"({type(extraction_response_json).__name__}). Raw: {str(extraction_response_json)[:500]}"
|
||||
)
|
||||
extraction_response_json = coerced_response_json
|
||||
|
||||
raw_facts = extraction_response_json.get("facts", [])
|
||||
|
||||
@@ -1451,21 +1532,9 @@ async def _extract_facts_from_chunk(
|
||||
elif fact_data.get("occurred_start"):
|
||||
fact_data["occurred_end"] = fact_data["occurred_start"]
|
||||
|
||||
# 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
|
||||
for ent in entities:
|
||||
if isinstance(ent, str):
|
||||
# Normalize string to Entity object
|
||||
validated_entities.append(Entity(text=ent))
|
||||
elif isinstance(ent, dict) and "text" in ent:
|
||||
try:
|
||||
validated_entities.append(Entity.model_validate(ent))
|
||||
except Exception as e:
|
||||
logger.warning(f"Invalid entity {ent}: {e}")
|
||||
# Entities are plain strings. Older prompts taught a {"text": ...}
|
||||
# object form, so keep unwrapping it for models that still emit it.
|
||||
validated_entities = _coerce_entity_strings(get_value("entities"))
|
||||
|
||||
# Post-process label entities from structured labels object
|
||||
entity_labels_raw = getattr(config, "entity_labels", None)
|
||||
@@ -1475,7 +1544,7 @@ async def _extract_facts_from_chunk(
|
||||
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}
|
||||
existing_texts_lower = {e.lower() for e in validated_entities}
|
||||
for group in labels_cfg.attributes:
|
||||
value = labels_data.get(group.key)
|
||||
if not value:
|
||||
@@ -1500,12 +1569,12 @@ async def _extract_facts_from_chunk(
|
||||
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))
|
||||
validated_entities.append(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))
|
||||
validated_entities.append(label_str)
|
||||
existing_texts_lower.add(label_str.lower())
|
||||
else:
|
||||
logger.warning(f"Label '{label_str}' not in valid label values, skipping")
|
||||
@@ -1513,7 +1582,7 @@ async def _extract_facts_from_chunk(
|
||||
# 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)
|
||||
e for e in validated_entities if is_label_entity(e, labels_cfg, labels_lookup)
|
||||
]
|
||||
elif not free_form_entities:
|
||||
# No labels but free_form disabled: clear all entities
|
||||
@@ -1664,33 +1733,22 @@ async def _extract_facts_with_auto_split(
|
||||
metadata=metadata,
|
||||
)
|
||||
except OutputTooLongError:
|
||||
# Output exceeded token limits - split the chunk in half and retry
|
||||
# Output exceeded token limits - split the chunk and retry. Conversation
|
||||
# chunks are JSON arrays, so preserve array/turn boundaries when possible.
|
||||
logger.warning(
|
||||
f"Output too long for chunk {chunk_index + 1}/{total_chunks} "
|
||||
f"({len(chunk)} chars). Splitting in half and retrying..."
|
||||
f"({len(chunk)} chars). Splitting and retrying..."
|
||||
)
|
||||
|
||||
# Split at the midpoint, preferring sentence boundaries
|
||||
mid_point = len(chunk) // 2
|
||||
split_chunks = _split_chunk_for_output_retry(chunk)
|
||||
if split_chunks is None:
|
||||
logger.warning(
|
||||
f"Cannot make progress splitting chunk {chunk_index + 1}/{total_chunks} "
|
||||
f"({len(chunk)} chars); dropping this sub-chunk."
|
||||
)
|
||||
return [], TokenUsage()
|
||||
|
||||
# Try to find a sentence boundary near the midpoint
|
||||
# Look for ". ", "! ", "? " within 20% of midpoint
|
||||
search_range = int(len(chunk) * 0.2)
|
||||
search_start = max(0, mid_point - search_range)
|
||||
search_end = min(len(chunk), mid_point + search_range)
|
||||
|
||||
sentence_endings = [". ", "! ", "? ", "\n\n"]
|
||||
best_split = mid_point
|
||||
|
||||
for ending in sentence_endings:
|
||||
pos = chunk.rfind(ending, search_start, search_end)
|
||||
if pos != -1:
|
||||
best_split = pos + len(ending)
|
||||
break
|
||||
|
||||
# Split the chunk
|
||||
first_half = chunk[:best_split].strip()
|
||||
second_half = chunk[best_split:].strip()
|
||||
first_half, second_half = split_chunks
|
||||
|
||||
logger.info(
|
||||
f"Split chunk {chunk_index + 1} into two sub-chunks: {len(first_half)} chars and {len(second_half)} chars"
|
||||
@@ -2132,7 +2190,10 @@ async def extract_facts_from_contents_batch_api(
|
||||
content_str = message.get("content", "{}")
|
||||
|
||||
try:
|
||||
extraction_response_json = json.loads(content_str)
|
||||
# #2701: use the lenient parser (strips markdown fences, scrubs
|
||||
# embedded control chars) so recoverable batch responses — e.g.
|
||||
# transient Gemini quirks — aren't dropped along with all their facts.
|
||||
extraction_response_json = parse_llm_json(content_str)
|
||||
except json.JSONDecodeError as e:
|
||||
message = f"{custom_id}: failed to parse JSON: {e}"
|
||||
logger.error(message)
|
||||
@@ -2144,6 +2205,19 @@ async def extract_facts_from_contents_batch_api(
|
||||
)
|
||||
continue
|
||||
|
||||
response_type_name = type(extraction_response_json).__name__
|
||||
extraction_response_json = _coerce_fact_response(extraction_response_json)
|
||||
if extraction_response_json is None:
|
||||
message = f"{custom_id}: LLM returned non-dict JSON ({response_type_name})"
|
||||
logger.error(message)
|
||||
extraction_errors.add(message)
|
||||
chunks_metadata.append(
|
||||
ChunkMetadata(
|
||||
chunk_text=chunk_content, fact_count=0, content_index=content_index, chunk_index=chunk_idx
|
||||
)
|
||||
)
|
||||
continue
|
||||
|
||||
# Parse facts (reuse existing logic from _extract_facts_from_chunk)
|
||||
raw_facts = extraction_response_json.get("facts", [])
|
||||
chunk_facts = []
|
||||
@@ -2210,18 +2284,9 @@ async def extract_facts_from_contents_batch_api(
|
||||
elif fact_data.get("occurred_start"):
|
||||
fact_data["occurred_end"] = fact_data["occurred_start"]
|
||||
|
||||
# Entities
|
||||
entities = get_value("entities")
|
||||
validated_entities = []
|
||||
if entities:
|
||||
for ent in entities:
|
||||
if isinstance(ent, str):
|
||||
validated_entities.append(Entity(text=ent))
|
||||
elif isinstance(ent, dict) and "text" in ent:
|
||||
try:
|
||||
validated_entities.append(Entity.model_validate(ent))
|
||||
except Exception:
|
||||
pass
|
||||
# Entities are plain strings. Older prompts taught a {"text": ...}
|
||||
# object form, so keep unwrapping it for models that still emit it.
|
||||
validated_entities = _coerce_entity_strings(get_value("entities"))
|
||||
|
||||
# Post-process label entities from structured labels object
|
||||
entity_labels_raw = getattr(config, "entity_labels", None)
|
||||
@@ -2231,7 +2296,7 @@ async def extract_facts_from_contents_batch_api(
|
||||
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}
|
||||
existing_texts_lower = {e.lower() for e in validated_entities}
|
||||
for group in labels_cfg_batch.attributes:
|
||||
value = labels_data.get(group.key)
|
||||
if not value:
|
||||
@@ -2256,18 +2321,18 @@ async def extract_facts_from_contents_batch_api(
|
||||
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))
|
||||
validated_entities.append(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))
|
||||
validated_entities.append(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)
|
||||
e for e in validated_entities if is_label_entity(e, labels_cfg_batch, labels_lookup_batch)
|
||||
]
|
||||
elif not free_form_entities_batch:
|
||||
validated_entities = []
|
||||
@@ -2354,7 +2419,7 @@ async def extract_facts_from_contents_batch_api(
|
||||
extracted_fact = ExtractedFactType(
|
||||
fact_text=fact_from_llm.fact,
|
||||
fact_type=fact_from_llm.fact_type,
|
||||
entities=[e.text for e in (fact_from_llm.entities or [])],
|
||||
entities=list(fact_from_llm.entities or []),
|
||||
occurred_start=_parse_datetime(fact_from_llm.occurred_start) if fact_from_llm.occurred_start else None,
|
||||
occurred_end=_parse_datetime(fact_from_llm.occurred_end) if fact_from_llm.occurred_end else None,
|
||||
causal_relations=_convert_causal_relations(fact_from_llm.causal_relations or [], global_fact_idx),
|
||||
@@ -2551,7 +2616,7 @@ async def extract_facts_from_contents(
|
||||
extracted_fact = ExtractedFactType(
|
||||
fact_text=fact_from_llm.fact,
|
||||
fact_type=fact_from_llm.fact_type,
|
||||
entities=[e.text for e in (fact_from_llm.entities or [])],
|
||||
entities=list(fact_from_llm.entities or []),
|
||||
# occurred_start/end: from LLM only, leave None if not provided
|
||||
occurred_start=_parse_datetime(fact_from_llm.occurred_start)
|
||||
if fact_from_llm.occurred_start
|
||||
|
||||
@@ -74,7 +74,9 @@ async def create_causal_links_batch(
|
||||
"""
|
||||
Create causal links between facts.
|
||||
|
||||
Links facts that have causal relationships (causes, enables, prevents).
|
||||
Retain writes the canonical ``caused_by`` relationship only. The database and
|
||||
retrieval paths also recognize historical causal types so imported and
|
||||
pre-existing memories remain traversable.
|
||||
|
||||
Args:
|
||||
conn: Database connection
|
||||
@@ -90,22 +92,7 @@ async def create_causal_links_batch(
|
||||
if len(unit_ids) != len(facts):
|
||||
raise ValueError(f"Mismatch between unit_ids ({len(unit_ids)}) and facts ({len(facts)})")
|
||||
|
||||
# Extract causal relations in the format expected by link_utils
|
||||
# Format: List of lists, where each inner list is the causal relations for that fact
|
||||
causal_relations_per_fact = []
|
||||
for fact in facts:
|
||||
if fact.causal_relations:
|
||||
# Convert CausalRelation objects to dicts
|
||||
relations_dicts = [
|
||||
{
|
||||
"relation_type": rel.relation_type,
|
||||
"target_fact_index": rel.target_fact_index,
|
||||
}
|
||||
for rel in fact.causal_relations
|
||||
]
|
||||
causal_relations_per_fact.append(relations_dicts)
|
||||
else:
|
||||
causal_relations_per_fact.append([])
|
||||
causal_relations_per_fact = [fact.causal_relations or [] for fact in facts]
|
||||
|
||||
link_count = await link_utils.create_causal_links_batch(conn, bank_id, unit_ids, causal_relations_per_fact, ops=ops)
|
||||
|
||||
|
||||
@@ -7,7 +7,11 @@ import time
|
||||
from datetime import UTC, datetime, timedelta
|
||||
|
||||
from ..._vector_index import ann_search_tuning_settings, configured_vector_extension
|
||||
from ..causal_links import CANONICAL_CAUSAL_LINK_TYPES, LEGACY_CAUSAL_LINK_TYPES
|
||||
from ..db.base import DatabaseConnection
|
||||
from ..db.ops import DataAccessOps
|
||||
from ..memory_engine import fq_table
|
||||
from .types import CausalRelation
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -771,28 +775,61 @@ async def create_semantic_links_batch(
|
||||
|
||||
|
||||
async def create_causal_links_batch(
|
||||
conn,
|
||||
conn: DatabaseConnection,
|
||||
bank_id: str,
|
||||
unit_ids: list[str],
|
||||
causal_relations_per_fact: list[list[dict]],
|
||||
ops=None,
|
||||
causal_relations_per_fact: list[list[CausalRelation]],
|
||||
ops: DataAccessOps | None = None,
|
||||
) -> int:
|
||||
"""
|
||||
Create causal links between facts based on LLM-extracted causal relationships.
|
||||
"""Create canonical causal links for the retain pipeline.
|
||||
|
||||
Args:
|
||||
conn: Database connection
|
||||
unit_ids: List of unit IDs (in same order as causal_relations_per_fact)
|
||||
causal_relations_per_fact: List of causal relations for each fact.
|
||||
Each element is a list of dicts with:
|
||||
- target_fact_index: Index into unit_ids for the target fact
|
||||
- relation_type: "caused_by"
|
||||
Retain must only create the backward-looking ``caused_by`` form. Historical
|
||||
types are restored exclusively through ``restore_legacy_causal_links_batch``.
|
||||
"""
|
||||
return await _write_causal_links_batch(
|
||||
conn,
|
||||
bank_id,
|
||||
unit_ids,
|
||||
causal_relations_per_fact,
|
||||
CANONICAL_CAUSAL_LINK_TYPES,
|
||||
ops=ops,
|
||||
)
|
||||
|
||||
|
||||
async def restore_legacy_causal_links_batch(
|
||||
conn: DatabaseConnection,
|
||||
bank_id: str,
|
||||
unit_ids: list[str],
|
||||
causal_relations_per_fact: list[list[CausalRelation]],
|
||||
ops: DataAccessOps | None = None,
|
||||
) -> int:
|
||||
"""Restore historical causal links while importing a transfer archive.
|
||||
|
||||
This is deliberately separate from the retain writer: retrieval continues
|
||||
reading historical types, but only transfer import may create them.
|
||||
"""
|
||||
return await _write_causal_links_batch(
|
||||
conn,
|
||||
bank_id,
|
||||
unit_ids,
|
||||
causal_relations_per_fact,
|
||||
LEGACY_CAUSAL_LINK_TYPES,
|
||||
ops=ops,
|
||||
)
|
||||
|
||||
|
||||
async def _write_causal_links_batch(
|
||||
conn: DatabaseConnection,
|
||||
bank_id: str,
|
||||
unit_ids: list[str],
|
||||
causal_relations_per_fact: list[list[CausalRelation]],
|
||||
allowed_relation_types: frozenset[str],
|
||||
ops: DataAccessOps | None = None,
|
||||
) -> int:
|
||||
"""Write causal links after the caller has selected its allowed taxonomy.
|
||||
|
||||
Returns:
|
||||
Number of causal links created
|
||||
|
||||
Causal link type:
|
||||
- "caused_by": This fact was caused by the target fact
|
||||
"""
|
||||
if not unit_ids or not causal_relations_per_fact:
|
||||
return 0
|
||||
@@ -809,15 +846,13 @@ async def create_causal_links_batch(
|
||||
from_unit_id = unit_ids[fact_idx]
|
||||
|
||||
for relation in causal_relations:
|
||||
target_idx = relation["target_fact_index"]
|
||||
relation_type = relation["relation_type"]
|
||||
target_idx = relation.target_fact_index
|
||||
relation_type = relation.relation_type
|
||||
|
||||
# Validate relation_type - only "caused_by" is supported (DB constraint)
|
||||
valid_types = {"caused_by"}
|
||||
if relation_type not in valid_types:
|
||||
if relation_type not in allowed_relation_types:
|
||||
logger.error(
|
||||
f"Invalid relation_type '{relation_type}' (type: {type(relation_type).__name__}) "
|
||||
f"from fact {fact_idx}. Must be one of: {valid_types}. "
|
||||
f"from fact {fact_idx}. Must be one of: {allowed_relation_types}. "
|
||||
f"Relation data: {relation}"
|
||||
)
|
||||
continue
|
||||
|
||||
@@ -549,7 +549,11 @@ async def _extract_and_embed(
|
||||
embeddings = await embedding_processing.generate_embeddings_batch(embeddings_model, augmented_texts)
|
||||
log_buffer.append(f" Generate embeddings: {len(embeddings)} embeddings in {time.time() - step_start:.3f}s")
|
||||
|
||||
processed_facts = [ProcessedFact.from_extracted_fact(ef, emb) for ef, emb in zip(extracted_facts, embeddings)]
|
||||
processed_facts = [
|
||||
pf
|
||||
for ef, emb in zip(extracted_facts, embeddings)
|
||||
if (pf := ProcessedFact.from_extracted_fact(ef, emb)) is not None
|
||||
]
|
||||
|
||||
return extracted_facts, processed_facts, chunks, usage
|
||||
|
||||
@@ -831,6 +835,12 @@ async def retain_batch(
|
||||
first = contents_dicts[0]
|
||||
if first.get("context"):
|
||||
existing_content["context"] = first["context"]
|
||||
if first.get("event_date"):
|
||||
existing_content["event_date"] = first["event_date"]
|
||||
if first.get("metadata"):
|
||||
existing_content["metadata"] = first["metadata"]
|
||||
if first.get("observation_scopes") is not None:
|
||||
existing_content["observation_scopes"] = first["observation_scopes"]
|
||||
if first.get("tags"):
|
||||
existing_content["tags"] = first["tags"]
|
||||
contents_dicts = [existing_content, *contents_dicts]
|
||||
@@ -852,6 +862,12 @@ async def retain_batch(
|
||||
contents_dicts = [{"content": json.dumps(_merged, ensure_ascii=False)}]
|
||||
if first.get("context"):
|
||||
contents_dicts[0]["context"] = first["context"]
|
||||
if first.get("event_date"):
|
||||
contents_dicts[0]["event_date"] = first["event_date"]
|
||||
if first.get("metadata"):
|
||||
contents_dicts[0]["metadata"] = first["metadata"]
|
||||
if first.get("observation_scopes") is not None:
|
||||
contents_dicts[0]["observation_scopes"] = first["observation_scopes"]
|
||||
if first.get("tags"):
|
||||
contents_dicts[0]["tags"] = first["tags"]
|
||||
except (json.JSONDecodeError, ValueError, TypeError):
|
||||
|
||||
@@ -5,11 +5,14 @@ These dataclasses provide type safety throughout the retain operation,
|
||||
from content input to fact storage.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime
|
||||
from typing import Literal, TypedDict
|
||||
from uuid import UUID
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class RetainContentDict(TypedDict, total=False):
|
||||
"""Type definition for content items in retain_batch_async.
|
||||
@@ -96,10 +99,12 @@ class CausalRelation:
|
||||
"""
|
||||
Causal relationship between facts.
|
||||
|
||||
Represents how one fact was caused by another.
|
||||
Retain emits only the backward-looking ``caused_by`` form. Transfer import
|
||||
reuses this structure to restore historical causal types without allowing
|
||||
normal retain writes to create them.
|
||||
"""
|
||||
|
||||
relation_type: str # "caused_by"
|
||||
relation_type: str # ``caused_by`` for retain; legacy types for transfer restore
|
||||
target_fact_index: int # Index of the target fact in the batch
|
||||
|
||||
|
||||
@@ -185,10 +190,47 @@ class ProcessedFact:
|
||||
"""Check if this fact was marked as a duplicate."""
|
||||
return self.unit_id is None
|
||||
|
||||
@staticmethod
|
||||
def _is_degenerate_text(text: str) -> bool:
|
||||
"""Check if fact text has zero information content.
|
||||
|
||||
Rejects empty strings, whitespace-only, single punctuation marks,
|
||||
and common LLM hallucination patterns that carry no semantic meaning.
|
||||
"""
|
||||
stripped = (text or "").strip()
|
||||
if not stripped:
|
||||
return True
|
||||
# Single or repeated punctuation patterns with no semantic content
|
||||
degenerate_patterns = {
|
||||
"...",
|
||||
"…",
|
||||
"-",
|
||||
"--",
|
||||
"---",
|
||||
".",
|
||||
"..",
|
||||
"•",
|
||||
"·",
|
||||
"*",
|
||||
"**",
|
||||
"***",
|
||||
"_,_",
|
||||
"_, _, _",
|
||||
}
|
||||
if stripped in degenerate_patterns:
|
||||
return True
|
||||
# Strings composed entirely of punctuation and whitespace
|
||||
if all(c in ".,;:!?-–—…\"'`´ \t\n\r" for c in stripped):
|
||||
return True
|
||||
# Very short text (<= 2 chars) that is only punctuation
|
||||
if len(stripped) <= 2 and all(not c.isalnum() for c in stripped):
|
||||
return True
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
def from_extracted_fact(
|
||||
extracted_fact: "ExtractedFact", embedding: list[float], chunk_id: str | None = None
|
||||
) -> "ProcessedFact":
|
||||
) -> "ProcessedFact | None":
|
||||
"""
|
||||
Create ProcessedFact from ExtractedFact.
|
||||
|
||||
@@ -198,8 +240,17 @@ class ProcessedFact:
|
||||
chunk_id: Optional chunk ID
|
||||
|
||||
Returns:
|
||||
ProcessedFact ready for storage
|
||||
ProcessedFact ready for storage, or None if the fact text is degenerate
|
||||
(zero information content — punctuation-only, empty, etc.)
|
||||
"""
|
||||
fact_text = extracted_fact.fact_text or ""
|
||||
if ProcessedFact._is_degenerate_text(fact_text):
|
||||
logger.warning(
|
||||
f"Rejected degenerate fact text: type={extracted_fact.fact_type}, "
|
||||
f"text={fact_text[:80]!r}, entities={extracted_fact.entities}"
|
||||
)
|
||||
return None
|
||||
|
||||
# Use occurred dates only if explicitly provided by LLM
|
||||
occurred_start = extracted_fact.occurred_start
|
||||
occurred_end = extracted_fact.occurred_end
|
||||
@@ -209,7 +260,7 @@ class ProcessedFact:
|
||||
entities = [EntityRef(name=name) for name in extracted_fact.entities]
|
||||
|
||||
return ProcessedFact(
|
||||
fact_text=extracted_fact.fact_text,
|
||||
fact_text=fact_text,
|
||||
fact_type=extracted_fact.fact_type,
|
||||
embedding=embedding,
|
||||
occurred_start=occurred_start,
|
||||
|
||||
@@ -27,6 +27,23 @@ def fq_table(table_name: str) -> str:
|
||||
return f"{get_current_schema()}.{table_name}"
|
||||
|
||||
|
||||
def fq_routine(name: str) -> str:
|
||||
"""Schema-qualified name of a cross-tenant discovery routine.
|
||||
|
||||
These routines are database-global — each enumerates ``pg_class`` across every
|
||||
schema and dispatches per schema — so exactly one copy exists, installed into
|
||||
the configured schema by ``b6d2f8a4c1e7``. Calling it through the configured
|
||||
schema rather than a hardcoded ``public.`` is what makes a deployment living
|
||||
in a dedicated non-``public`` schema work (#2638).
|
||||
|
||||
Unlike :func:`fq_table` this ignores the per-request schema contextvar: the
|
||||
routines are deliberately cross-tenant, called from background loops that have
|
||||
no request context.
|
||||
"""
|
||||
schema = get_config().database_schema or "public"
|
||||
return '"' + schema.replace('"', '""') + '".' + name
|
||||
|
||||
|
||||
def fq_table_explicit(table: str, schema: str | None = None) -> str:
|
||||
"""Get fully-qualified table name with an explicit schema override.
|
||||
|
||||
|
||||
@@ -40,8 +40,6 @@ class GraphRetriever(ABC):
|
||||
fact_type: str,
|
||||
budget: int,
|
||||
query_text: str | None = None,
|
||||
semantic_seeds: list[RetrievalResult] | None = None,
|
||||
temporal_seeds: list[RetrievalResult] | None = None,
|
||||
adjacency=None, # TypedAdjacency, optional pre-loaded graph
|
||||
tags: list[str] | None = None, # Visibility scope tags for filtering
|
||||
tags_match: TagsMatch = "any", # How to match tags: 'any' (OR) or 'all' (AND)
|
||||
@@ -59,8 +57,6 @@ class GraphRetriever(ABC):
|
||||
fact_type: Fact type to filter ('world', 'experience', 'observation')
|
||||
budget: Maximum number of nodes to explore/return
|
||||
query_text: Original query text (optional, for some strategies)
|
||||
semantic_seeds: Pre-computed semantic entry points (from semantic retrieval)
|
||||
temporal_seeds: Pre-computed temporal entry points (from temporal retrieval)
|
||||
adjacency: Pre-loaded typed adjacency graph (optional)
|
||||
tags: Optional list of tags for visibility filtering (OR matching)
|
||||
|
||||
|
||||
@@ -1,8 +1,8 @@
|
||||
"""
|
||||
Link Expansion graph retrieval.
|
||||
|
||||
Expands from semantic/temporal seeds through three parallel, first-class signals
|
||||
stored in memory_links:
|
||||
Selects bounded semantic seeds, then expands through three parallel,
|
||||
first-class signals stored in memory_links:
|
||||
|
||||
1. Entity links — query-time self-join through unit_entities. Score = number of distinct
|
||||
shared entities between the seed set and each candidate, computed via
|
||||
@@ -127,8 +127,6 @@ class LinkExpansionRetriever(GraphRetriever):
|
||||
fact_type: str,
|
||||
budget: int,
|
||||
query_text: str | None = None,
|
||||
semantic_seeds: list[RetrievalResult] | None = None,
|
||||
temporal_seeds: list[RetrievalResult] | None = None,
|
||||
adjacency=None,
|
||||
tags: list[str] | None = None,
|
||||
tags_match: TagsMatch = "any",
|
||||
@@ -146,8 +144,6 @@ class LinkExpansionRetriever(GraphRetriever):
|
||||
fact_type: Fact type to filter
|
||||
budget: Maximum results to return
|
||||
query_text: Original query text (unused)
|
||||
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
|
||||
|
||||
@@ -158,32 +154,28 @@ class LinkExpansionRetriever(GraphRetriever):
|
||||
timings = GraphRetrievalTimings(fact_type=fact_type)
|
||||
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
# Find seeds if not provided
|
||||
if semantic_seeds:
|
||||
all_seeds = list(semantic_seeds)
|
||||
else:
|
||||
seeds_start = time.time()
|
||||
all_seeds = await _find_semantic_seeds(
|
||||
conn,
|
||||
query_embedding_str,
|
||||
bank_id,
|
||||
fact_type,
|
||||
limit=20,
|
||||
threshold=0.3,
|
||||
tags=tags,
|
||||
tags_match=tags_match,
|
||||
tag_groups=tag_groups,
|
||||
created_after=created_after,
|
||||
created_before=created_before,
|
||||
)
|
||||
timings.seeds_time = time.time() - seeds_start
|
||||
logger.debug(
|
||||
f"[LinkExpansion] Found {len(all_seeds)} semantic seeds for fact_type={fact_type} "
|
||||
f"(tags={tags}, tags_match={tags_match})"
|
||||
)
|
||||
|
||||
if temporal_seeds:
|
||||
all_seeds.extend(temporal_seeds)
|
||||
# Graph traversal deliberately chooses its own bounded seeds. The semantic and temporal
|
||||
# retrieval arms have independent candidate limits and thresholds, so reusing their
|
||||
# results would silently change graph-retrieval recall behavior.
|
||||
seeds_start = time.time()
|
||||
all_seeds = await _find_semantic_seeds(
|
||||
conn,
|
||||
query_embedding_str,
|
||||
bank_id,
|
||||
fact_type,
|
||||
limit=20,
|
||||
threshold=0.3,
|
||||
tags=tags,
|
||||
tags_match=tags_match,
|
||||
tag_groups=tag_groups,
|
||||
created_after=created_after,
|
||||
created_before=created_before,
|
||||
)
|
||||
timings.seeds_time = time.time() - seeds_start
|
||||
logger.debug(
|
||||
f"[LinkExpansion] Found {len(all_seeds)} semantic seeds for fact_type={fact_type} "
|
||||
f"(tags={tags}, tags_match={tags_match})"
|
||||
)
|
||||
|
||||
if not all_seeds:
|
||||
return [], timings
|
||||
@@ -243,12 +235,16 @@ class LinkExpansionRetriever(GraphRetriever):
|
||||
}
|
||||
|
||||
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]
|
||||
|
||||
results = []
|
||||
for row in rows:
|
||||
for fact_id in sorted_ids:
|
||||
row = row_map[fact_id]
|
||||
result = RetrievalResult.from_db_row(dict(row))
|
||||
result.activation = row["score"]
|
||||
# ``activation`` is used to re-sort graph results after fact types are
|
||||
# combined. It must retain the final additive score rather than the
|
||||
# raw score from one signal, which would otherwise discard the other
|
||||
# signals and make the cross-fact-type order disagree with this order.
|
||||
result.activation = score_map[fact_id]
|
||||
results.append(result)
|
||||
|
||||
# filter_results_by_tags is a no-op when no filter applies (tags falsy and not
|
||||
|
||||
@@ -15,7 +15,7 @@ from dataclasses import dataclass, field
|
||||
from datetime import UTC, datetime
|
||||
from typing import TYPE_CHECKING, Any, Optional
|
||||
|
||||
from ...config import get_config
|
||||
from ...config import DEFAULT_BM25_MAX_QUERY_TERMS, get_config
|
||||
from ..db_utils import acquire_with_retry
|
||||
from ..memory_engine import fq_table
|
||||
from ..sql import create_sql_dialect
|
||||
@@ -222,7 +222,12 @@ async def retrieve_semantic_bm25_combined(
|
||||
# --- BM25 UNION ALL arms (one per fact_type, only when tokens present) ---
|
||||
if _include_bm25:
|
||||
text_ext = config.text_search_extension
|
||||
bm25_text_param: str = dialect.prepare_bm25_text(tokens, query_text, text_search_extension=text_ext)
|
||||
bm25_text_param: str = dialect.prepare_bm25_text(
|
||||
tokens,
|
||||
query_text,
|
||||
text_search_extension=text_ext,
|
||||
max_query_terms=getattr(config, "bm25_max_query_terms", DEFAULT_BM25_MAX_QUERY_TERMS),
|
||||
)
|
||||
for i, ft in enumerate(fact_types):
|
||||
arms.append(
|
||||
dialect.build_bm25_arm(
|
||||
@@ -616,7 +621,7 @@ async def retrieve_temporal_combined(
|
||||
# bank_id on memory_units lets the planner use idx_memory_units_bank_fact_type.
|
||||
neighbors = await conn.fetch(
|
||||
f"""
|
||||
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, mu.metadata,
|
||||
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, mu.metadata, mu.proof_count,
|
||||
l.weight, l.link_type,
|
||||
1 - (mu.embedding <=> $1::vector) AS similarity
|
||||
FROM unnest($2::uuid[]) AS src(from_unit_id)
|
||||
@@ -793,7 +798,7 @@ async def retrieve_all_fact_types_parallel(
|
||||
tc_start,
|
||||
tc_end,
|
||||
budget=thinking_budget,
|
||||
semantic_threshold=min_semantic if min_semantic is not None else 0.1,
|
||||
semantic_threshold=0.1,
|
||||
tags=tags,
|
||||
tags_match=tags_match,
|
||||
tag_groups=tag_groups,
|
||||
@@ -817,8 +822,6 @@ async def retrieve_all_fact_types_parallel(
|
||||
fact_type=ft,
|
||||
budget=thinking_budget,
|
||||
query_text=query_text,
|
||||
semantic_seeds=None,
|
||||
temporal_seeds=None,
|
||||
tags=tags,
|
||||
tags_match=tags_match,
|
||||
tag_groups=tag_groups,
|
||||
|
||||
@@ -22,7 +22,7 @@ class GraphRetrievalTimings:
|
||||
pattern_count: int = 0 # Number of patterns executed
|
||||
fusion: float = 0.0 # Time for RRF fusion
|
||||
fetch: float = 0.0 # Time to fetch memory unit details
|
||||
seeds_time: float = 0.0 # Time to find semantic seeds (if fallback used)
|
||||
seeds_time: float = 0.0 # Time spent selecting semantic graph seeds
|
||||
result_count: int = 0 # Number of results returned
|
||||
# Detailed per-hop timing: list of {hop, exec_time, uncached, load_time, edges_loaded, total_time}
|
||||
hop_details: list[dict] = field(default_factory=list)
|
||||
|
||||
@@ -449,6 +449,7 @@ class SQLDialect(ABC):
|
||||
query_text: str,
|
||||
*,
|
||||
text_search_extension: str = "native",
|
||||
max_query_terms: int | None = None,
|
||||
) -> str:
|
||||
"""Prepare the text parameter value for BM25 search.
|
||||
|
||||
@@ -459,6 +460,8 @@ class SQLDialect(ABC):
|
||||
tokens: Tokenized query words.
|
||||
query_text: Original query text.
|
||||
text_search_extension: Full-text search backend variant.
|
||||
max_query_terms: Optional backend-specific token cap. 0 or None
|
||||
leaves query terms uncapped.
|
||||
|
||||
Returns:
|
||||
Prepared text string to bind as the BM25 text parameter.
|
||||
|
||||
@@ -303,6 +303,7 @@ class OracleDialect(SQLDialect):
|
||||
query_text: str,
|
||||
*,
|
||||
text_search_extension: str = "native",
|
||||
max_query_terms: int | None = None,
|
||||
) -> str:
|
||||
# Oracle Text: filter tokens with special chars, escape reserved words
|
||||
# with curly braces (e.g. "about" → "{about}"), and join with OR.
|
||||
|
||||
@@ -254,8 +254,11 @@ class PostgreSQLDialect(SQLDialect):
|
||||
query_text: str,
|
||||
*,
|
||||
text_search_extension: str = "native",
|
||||
max_query_terms: int | None = None,
|
||||
) -> str:
|
||||
if text_search_extension in ("vchord", "pg_textsearch", "pgroonga", "pg_search"):
|
||||
return query_text
|
||||
if max_query_terms is not None and max_query_terms > 0:
|
||||
tokens = tokens[:max_query_terms]
|
||||
# native tsvector: join tokens with OR operator
|
||||
return " | ".join(tokens)
|
||||
|
||||
@@ -128,6 +128,11 @@ def _extract_non_chinese_period(query: str, reference_date: datetime) -> DateRan
|
||||
"december|diciembre|dicembre|d[ée]cembre|dezember": 12,
|
||||
}
|
||||
for pattern, month_num in month_patterns.items():
|
||||
# Skip when a day number precedes the month ("13 июля 2026", "13 July 2026"):
|
||||
# that is an exact date, and collapsing it to the whole month loses precision.
|
||||
# dateparser resolves those correctly, so let them fall through to it.
|
||||
if re.search(rf"\b\d{{1,2}}\s+({pattern})\b", query, re.IGNORECASE):
|
||||
continue
|
||||
match = re.search(rf"\b({pattern})\s+(\d{{4}})\b", query, re.IGNORECASE)
|
||||
if match:
|
||||
year = int(match.group(2))
|
||||
|
||||
@@ -19,9 +19,12 @@ from decimal import Decimal
|
||||
from typing import Any
|
||||
from uuid import UUID
|
||||
|
||||
from ..causal_links import CAUSAL_LINK_TYPES
|
||||
from ..db_utils import acquire_with_retry
|
||||
from ..schema import fq_table
|
||||
from .schema import (
|
||||
CARRIED_HISTORY_TABLES,
|
||||
HISTORY_TABLES,
|
||||
SCHEMA_VERSION,
|
||||
TransferCausalRelation,
|
||||
TransferChunk,
|
||||
@@ -71,9 +74,7 @@ _BANK_ROW_TABLES = ("banks", "mental_models", "directives", "webhooks")
|
||||
# keep their (id, bank_id) across export/import, so their refresh history can be
|
||||
# re-attached. The surrogate ``id`` is dropped on dump so the target reassigns it
|
||||
# (see _dump_history_rows); restored after its parent table (mental_models).
|
||||
_CARRIED_HISTORY_TABLES = ("mental_model_history",)
|
||||
# Operational history — only carried with include_history=True.
|
||||
_HISTORY_TABLES = ("audit_log", "llm_requests")
|
||||
# Intentionally never exported.
|
||||
_SKIP_TABLES = frozenset(
|
||||
{
|
||||
@@ -125,10 +126,9 @@ class _LoadedExport:
|
||||
unit_index: dict[Any, _UnitLocation] = field(default_factory=dict)
|
||||
|
||||
|
||||
# Causal link types that retain persists between facts. Only these travel in the
|
||||
# archive; temporal/semantic/entity links are regenerated against the target bank.
|
||||
_CAUSAL_LINK_TYPES = ("caused_by", "causes", "enables", "prevents")
|
||||
|
||||
# Retain currently writes only ``caused_by``. The legacy types stay in archives
|
||||
# so importing a historical bank preserves its graph; temporal/semantic/entity
|
||||
# links are regenerated against the target bank.
|
||||
# Facts of these types are exported; observations are derived and excluded.
|
||||
_EXPORTED_FACT_TYPES = ("world", "experience")
|
||||
|
||||
@@ -287,11 +287,11 @@ async def export_bank(conn: Any, bank_id: str, *, include_history: bool = False)
|
||||
observations = await _load_observations(conn, bank_id, loaded.unit_index)
|
||||
|
||||
bank_rows = {table: await _dump_bank_rows(conn, table, bank_id) for table in _BANK_ROW_TABLES}
|
||||
for table in _CARRIED_HISTORY_TABLES:
|
||||
for table in CARRIED_HISTORY_TABLES:
|
||||
bank_rows[table] = await _dump_history_rows(conn, table, bank_id)
|
||||
history_rows: dict[str, list[dict]] = {}
|
||||
if include_history:
|
||||
history_rows = {table: await _dump_bank_rows(conn, table, bank_id) for table in _HISTORY_TABLES}
|
||||
history_rows = {table: await _dump_bank_rows(conn, table, bank_id) for table in HISTORY_TABLES}
|
||||
|
||||
archive = io.BytesIO()
|
||||
fact_total = 0
|
||||
@@ -543,7 +543,7 @@ async def _attach_causal_relations(conn: Any, loaded: _LoadedFacts) -> None:
|
||||
AND from_unit_id = ANY($2)
|
||||
AND to_unit_id = ANY($2)
|
||||
""",
|
||||
list(_CAUSAL_LINK_TYPES),
|
||||
list(CAUSAL_LINK_TYPES),
|
||||
list(loaded.unit_index.keys()),
|
||||
)
|
||||
for row in rows:
|
||||
|
||||
@@ -18,8 +18,9 @@ from dataclasses import dataclass, field
|
||||
from datetime import UTC, date, datetime
|
||||
from typing import Any, Literal
|
||||
|
||||
from ..causal_links import CANONICAL_CAUSAL_LINK_TYPE, LEGACY_CAUSAL_LINK_TYPES
|
||||
from ..db_utils import acquire_with_retry
|
||||
from ..retain import bank_utils, chunk_storage, embedding_processing, fact_storage, orchestrator
|
||||
from ..retain import bank_utils, chunk_storage, embedding_processing, fact_storage, link_utils, orchestrator
|
||||
from ..retain.types import (
|
||||
CausalRelation,
|
||||
ChunkMetadata,
|
||||
@@ -29,6 +30,8 @@ from ..retain.types import (
|
||||
)
|
||||
from ..schema import fq_table
|
||||
from .schema import (
|
||||
CARRIED_HISTORY_TABLES,
|
||||
HISTORY_TABLES,
|
||||
SCHEMA_VERSION,
|
||||
TransferDocument,
|
||||
TransferFact,
|
||||
@@ -221,8 +224,6 @@ _BANK_CHILD_TABLES = ("mental_models", "directives", "webhooks")
|
||||
# Child-history carried verbatim; restored after its parent (mental_models) so the
|
||||
# foreign key resolves. Surrogate ids were dropped on export (the target reassigns
|
||||
# them), so these restore via fresh IDENTITY values.
|
||||
_CARRIED_HISTORY_TABLES = ("mental_model_history",)
|
||||
_HISTORY_TABLES = ("audit_log", "llm_requests")
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -263,11 +264,11 @@ def parse_bank_archive(archive_bytes: bytes) -> ParsedBankArchive:
|
||||
f"Not a whole-bank archive (archive_type={manifest.archive_type!r}); use import_documents instead"
|
||||
)
|
||||
bank_rows: dict[str, list[dict]] = {}
|
||||
for table in ("banks", *_BANK_CHILD_TABLES, *_CARRIED_HISTORY_TABLES):
|
||||
for table in ("banks", *_BANK_CHILD_TABLES, *CARRIED_HISTORY_TABLES):
|
||||
fname = f"{table}.json"
|
||||
bank_rows[table] = json.loads(zf.read(fname)) if fname in names else []
|
||||
history_rows: dict[str, list[dict]] = {}
|
||||
for table in _HISTORY_TABLES:
|
||||
for table in HISTORY_TABLES:
|
||||
fname = f"history/{table}.json"
|
||||
if fname in names:
|
||||
history_rows[table] = json.loads(zf.read(fname))
|
||||
@@ -408,7 +409,7 @@ async def import_bank(
|
||||
result.directives_imported = await _restore_rows(conn, "directives", parsed.bank_rows.get("directives", []))
|
||||
result.webhooks_imported = await _restore_rows(conn, "webhooks", parsed.bank_rows.get("webhooks", []))
|
||||
if include_history:
|
||||
for table in _HISTORY_TABLES:
|
||||
for table in HISTORY_TABLES:
|
||||
result.history_rows_imported += await _restore_rows(conn, table, parsed.history_rows.get(table, []))
|
||||
|
||||
logger.info(
|
||||
@@ -474,12 +475,17 @@ async def _import_one_document(
|
||||
)
|
||||
|
||||
extracted_facts = [_to_extracted_fact(fact) for fact in document.facts]
|
||||
legacy_causal_relations = _legacy_causal_relations(document)
|
||||
|
||||
processed_facts: list[ProcessedFact] = []
|
||||
if extracted_facts:
|
||||
augmented = embedding_processing.augment_texts_with_dates(extracted_facts, format_date_fn)
|
||||
embeddings = await embedding_processing.generate_embeddings_batch(embeddings_model, augmented)
|
||||
processed_facts = [ProcessedFact.from_extracted_fact(ef, emb) for ef, emb in zip(extracted_facts, embeddings)]
|
||||
processed_facts = [
|
||||
pf
|
||||
for ef, emb in zip(extracted_facts, embeddings)
|
||||
if (pf := ProcessedFact.from_extracted_fact(ef, emb)) is not None
|
||||
]
|
||||
|
||||
contents = [RetainContent(content=document.original_text or "")]
|
||||
chunk_meta = [
|
||||
@@ -545,6 +551,18 @@ async def _import_one_document(
|
||||
ops=ops,
|
||||
)
|
||||
|
||||
# Retain writes only ``caused_by``. Restore legacy archive edges
|
||||
# separately so their distinct direction and semantics survive a
|
||||
# transfer without broadening the normal retain write contract.
|
||||
if result_unit_ids and legacy_causal_relations:
|
||||
await link_utils.restore_legacy_causal_links_batch(
|
||||
conn,
|
||||
bank_id,
|
||||
result_unit_ids[0],
|
||||
legacy_causal_relations,
|
||||
ops=ops,
|
||||
)
|
||||
|
||||
try:
|
||||
await entity_resolver.flush_pending_stats()
|
||||
except Exception:
|
||||
@@ -705,6 +723,7 @@ def _to_extracted_fact(fact: TransferFact) -> ExtractedFact:
|
||||
causal_relations=[
|
||||
CausalRelation(relation_type=rel.relation_type, target_fact_index=rel.target_fact_index)
|
||||
for rel in fact.causal_relations
|
||||
if rel.relation_type == CANONICAL_CAUSAL_LINK_TYPE
|
||||
],
|
||||
content_index=0,
|
||||
chunk_index=fact.chunk_index,
|
||||
@@ -714,3 +733,19 @@ def _to_extracted_fact(fact: TransferFact) -> ExtractedFact:
|
||||
tags=list(fact.tags),
|
||||
observation_scopes=fact.observation_scopes,
|
||||
)
|
||||
|
||||
|
||||
def _legacy_causal_relations(document: TransferDocument) -> list[list[CausalRelation]]:
|
||||
"""Return legacy archive edges for transfer-only restoration.
|
||||
|
||||
Invalid archive values are excluded. The write helper repeats the explicit
|
||||
compatibility allowlist as a persistence boundary.
|
||||
"""
|
||||
return [
|
||||
[
|
||||
CausalRelation(relation_type=relation.relation_type, target_fact_index=relation.target_fact_index)
|
||||
for relation in fact.causal_relations
|
||||
if relation.relation_type in LEGACY_CAUSAL_LINK_TYPES
|
||||
]
|
||||
for fact in document.facts
|
||||
]
|
||||
|
||||
@@ -22,6 +22,12 @@ from pydantic import BaseModel, Field
|
||||
# Bump when the archive layout changes in a backward-incompatible way.
|
||||
SCHEMA_VERSION = 1
|
||||
|
||||
# Whole-bank transfer table classifications shared by export and import.
|
||||
# Child history is always carried after its mental-model parent; operational
|
||||
# history is optional and included only when the caller requests it.
|
||||
CARRIED_HISTORY_TABLES = ("mental_model_history",)
|
||||
HISTORY_TABLES = ("audit_log", "llm_requests")
|
||||
|
||||
ObservationScopes = Literal["per_tag", "combined", "all_combinations", "shared"] | list[list[str]]
|
||||
|
||||
|
||||
|
||||
@@ -45,6 +45,7 @@ from hindsight_api.extensions.operation_validator import (
|
||||
# Consolidation operation
|
||||
ConsolidateContext,
|
||||
ConsolidateResult,
|
||||
CreateBankContext,
|
||||
# File Conversion
|
||||
FileConvertResult,
|
||||
# Mental Model operations
|
||||
@@ -105,6 +106,7 @@ __all__ = [
|
||||
"BankReadOperation",
|
||||
"BankWriteContext",
|
||||
"BankWriteOperation",
|
||||
"CreateBankContext",
|
||||
# Operation Validator - Consolidation
|
||||
"ConsolidateContext",
|
||||
"ConsolidateResult",
|
||||
|
||||
@@ -397,6 +397,14 @@ class BankWriteContext:
|
||||
request_context: "RequestContext"
|
||||
|
||||
|
||||
@dataclass
|
||||
class CreateBankContext:
|
||||
"""Context for validating creation of a new bank."""
|
||||
|
||||
bank_id: str
|
||||
request_context: "RequestContext"
|
||||
|
||||
|
||||
@dataclass
|
||||
class BankListContext:
|
||||
"""Context for filtering the bank list (post-query)."""
|
||||
@@ -881,6 +889,23 @@ class OperationValidatorExtension(Extension, ABC):
|
||||
"""
|
||||
return ValidationResult.accept()
|
||||
|
||||
async def validate_create_bank(self, ctx: CreateBankContext) -> ValidationResult:
|
||||
"""
|
||||
Validate creation of a new bank before the bank row is inserted.
|
||||
|
||||
Override to implement custom validation logic for operations that
|
||||
explicitly or implicitly create a bank.
|
||||
|
||||
Args:
|
||||
ctx: Context containing:
|
||||
- bank_id: Bank identifier
|
||||
- request_context: Request context with auth info
|
||||
|
||||
Returns:
|
||||
ValidationResult indicating whether the bank may be created.
|
||||
"""
|
||||
return ValidationResult.accept()
|
||||
|
||||
async def filter_bank_list(self, ctx: BankListContext) -> BankListResult:
|
||||
"""
|
||||
Filter the bank list after querying.
|
||||
|
||||
@@ -1225,7 +1225,9 @@ def _register_create_bank(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsCo
|
||||
"""
|
||||
try:
|
||||
request_context = _get_request_context(config)
|
||||
# get_bank_profile auto-creates bank if it doesn't exist
|
||||
# create_bank may auto-create the bank; validate that explicit
|
||||
# creation permission before reading the resulting profile.
|
||||
await memory._ensure_bank_exists(bank_id, request_context)
|
||||
profile = await memory.get_bank_profile(bank_id, request_context=request_context)
|
||||
|
||||
# Update name/mission if provided
|
||||
@@ -3145,7 +3147,10 @@ def _register_get_bank(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfi
|
||||
profile = await memory.get_bank_profile(
|
||||
target_bank,
|
||||
request_context=_get_request_context(config),
|
||||
create_if_missing=False,
|
||||
)
|
||||
if profile is None:
|
||||
return json.dumps({"error": f"Bank '{target_bank}' not found"})
|
||||
if "disposition" in profile and hasattr(profile["disposition"], "model_dump"):
|
||||
profile["disposition"] = profile["disposition"].model_dump()
|
||||
return json.dumps(profile, indent=2, default=str)
|
||||
@@ -3173,7 +3178,10 @@ def _register_get_bank(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfi
|
||||
profile = await memory.get_bank_profile(
|
||||
target_bank,
|
||||
request_context=_get_request_context(config),
|
||||
create_if_missing=False,
|
||||
)
|
||||
if profile is None:
|
||||
return {"error": f"Bank '{target_bank}' not found"}
|
||||
if "disposition" in profile and hasattr(profile["disposition"], "model_dump"):
|
||||
profile["disposition"] = profile["disposition"].model_dump()
|
||||
return profile
|
||||
|
||||
@@ -56,6 +56,11 @@ MIGRATION_LOCK_ID = 123456789
|
||||
_alembic_lock = threading.Lock()
|
||||
|
||||
|
||||
def _set_alembic_main_option(config: Config, name: str, value: str) -> None:
|
||||
"""Set an Alembic option without treating URL percent escapes as interpolation."""
|
||||
config.set_main_option(name, value.replace("%", "%%"))
|
||||
|
||||
|
||||
def _detect_vector_extension(conn, vector_extension: str = "pgvector") -> str:
|
||||
"""Validate configured vector extension and preserve Azure DiskANN detection."""
|
||||
return detect_vector_extension(conn, vector_extension)
|
||||
@@ -191,22 +196,22 @@ def _run_migrations_internal(database_url: str, script_location: str, schema: st
|
||||
alembic_cfg = Config()
|
||||
|
||||
# Set the script location (where alembic versions are stored)
|
||||
alembic_cfg.set_main_option("script_location", script_location)
|
||||
_set_alembic_main_option(alembic_cfg, "script_location", script_location)
|
||||
|
||||
# Set the database URL
|
||||
alembic_cfg.set_main_option("sqlalchemy.url", database_url)
|
||||
_set_alembic_main_option(alembic_cfg, "sqlalchemy.url", database_url)
|
||||
|
||||
# Configure logging (optional, but helps with debugging)
|
||||
# Uses Python's logging system instead of alembic.ini
|
||||
alembic_cfg.set_main_option("prepend_sys_path", ".")
|
||||
_set_alembic_main_option(alembic_cfg, "prepend_sys_path", ".")
|
||||
|
||||
# Set path_separator to avoid deprecation warning
|
||||
alembic_cfg.set_main_option("path_separator", "os")
|
||||
_set_alembic_main_option(alembic_cfg, "path_separator", "os")
|
||||
|
||||
# If targeting a specific schema, pass it to env.py via config
|
||||
# env.py will handle setting search_path and version_table_schema
|
||||
if schema:
|
||||
alembic_cfg.set_main_option("target_schema", schema)
|
||||
_set_alembic_main_option(alembic_cfg, "target_schema", schema)
|
||||
|
||||
# Run migrations under a process-level lock. Alembic uses module-level
|
||||
# global proxies that are not thread-safe, so concurrent command.upgrade()
|
||||
@@ -431,8 +436,8 @@ def check_migration_status(
|
||||
|
||||
# Create config programmatically
|
||||
alembic_cfg = Config()
|
||||
alembic_cfg.set_main_option("script_location", script_location)
|
||||
alembic_cfg.set_main_option("path_separator", "os")
|
||||
_set_alembic_main_option(alembic_cfg, "script_location", script_location)
|
||||
_set_alembic_main_option(alembic_cfg, "path_separator", "os")
|
||||
|
||||
script = ScriptDirectory.from_config(alembic_cfg)
|
||||
head_rev = script.get_current_head()
|
||||
|
||||
@@ -284,6 +284,8 @@ class MemoryLink(Base):
|
||||
entity = relationship("Entity", back_populates="memory_links")
|
||||
|
||||
__table_args__ = (
|
||||
# Retain writes ``caused_by`` only. Keep the historical causal values
|
||||
# valid so existing rows and transfer archives remain queryable.
|
||||
CheckConstraint(
|
||||
"link_type IN ('temporal', 'semantic', 'entity', 'causes', 'caused_by', 'enables', 'prevents')",
|
||||
name="memory_links_link_type_check",
|
||||
|
||||
@@ -2,6 +2,7 @@ from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -153,7 +154,22 @@ async def stop_embedded_postgres() -> None:
|
||||
await _default_instance.stop()
|
||||
|
||||
|
||||
def parse_pg0_url(db_url: str) -> tuple[bool, str | None, int | None]:
|
||||
@dataclass(frozen=True)
|
||||
class Pg0Url:
|
||||
"""Parsed representation of a ``pg0`` embedded-database URL.
|
||||
|
||||
``username``/``password`` are ``None`` when the URL omits credentials, in
|
||||
which case the pg0 defaults (``hindsight``/``hindsight``) apply.
|
||||
"""
|
||||
|
||||
is_pg0: bool
|
||||
instance_name: str | None = None
|
||||
port: int | None = None
|
||||
username: str | None = None
|
||||
password: str | None = None
|
||||
|
||||
|
||||
def parse_pg0_url(db_url: str) -> Pg0Url:
|
||||
"""
|
||||
Parse a database URL and check if it's a pg0:// embedded database URL.
|
||||
|
||||
@@ -161,29 +177,47 @@ def parse_pg0_url(db_url: str) -> tuple[bool, str | None, int | None]:
|
||||
- "pg0" -> default instance "hindsight"
|
||||
- "pg0://instance-name" -> named instance
|
||||
- "pg0://instance-name:port" -> named instance with explicit port
|
||||
- "pg0://user:pwd@instance-name:port" -> named instance with credentials
|
||||
(``user`` or ``user:pwd``; either half may be present)
|
||||
- Any other URL (e.g., postgresql://) -> not a pg0 URL
|
||||
|
||||
Args:
|
||||
db_url: The database URL to parse
|
||||
|
||||
Returns:
|
||||
Tuple of (is_pg0, instance_name, port)
|
||||
- is_pg0: True if this is a pg0 URL
|
||||
- instance_name: The instance name (or None if not pg0)
|
||||
- port: The explicit port (or None for auto-assign)
|
||||
A :class:`Pg0Url`. When ``is_pg0`` is False the remaining fields are None.
|
||||
"""
|
||||
if db_url == "pg0":
|
||||
return True, "hindsight", None
|
||||
return Pg0Url(is_pg0=True, instance_name="hindsight")
|
||||
|
||||
if db_url.startswith("pg0://"):
|
||||
url_part = db_url[6:] # Remove "pg0://"
|
||||
if ":" in url_part:
|
||||
instance_name, port_str = url_part.rsplit(":", 1)
|
||||
return True, instance_name or "hindsight", int(port_str)
|
||||
else:
|
||||
return True, url_part or "hindsight", None
|
||||
if not db_url.startswith("pg0://"):
|
||||
return Pg0Url(is_pg0=False)
|
||||
|
||||
return False, None, None
|
||||
url_part = db_url[6:] # Remove "pg0://"
|
||||
|
||||
# Split optional "user:pwd@" credentials from the "instance:port" host part.
|
||||
# rsplit on the last "@" so passwords may contain "@".
|
||||
username: str | None = None
|
||||
password: str | None = None
|
||||
if "@" in url_part:
|
||||
creds, url_part = url_part.rsplit("@", 1)
|
||||
user_part, sep, pwd_part = creds.partition(":")
|
||||
username = user_part or None
|
||||
password = pwd_part if sep else None
|
||||
|
||||
if ":" in url_part:
|
||||
instance_name, port_str = url_part.rsplit(":", 1)
|
||||
port: int | None = int(port_str)
|
||||
else:
|
||||
instance_name, port = url_part, None
|
||||
|
||||
return Pg0Url(
|
||||
is_pg0=True,
|
||||
instance_name=instance_name or "hindsight",
|
||||
port=port,
|
||||
username=username,
|
||||
password=password,
|
||||
)
|
||||
|
||||
|
||||
async def resolve_database_url(db_url: str) -> str:
|
||||
@@ -199,8 +233,13 @@ async def resolve_database_url(db_url: str) -> str:
|
||||
Returns:
|
||||
The resolved postgresql:// connection URL
|
||||
"""
|
||||
is_pg0, instance_name, port = parse_pg0_url(db_url)
|
||||
if is_pg0:
|
||||
pg0 = EmbeddedPostgres(name=instance_name, port=port)
|
||||
parsed = parse_pg0_url(db_url)
|
||||
if parsed.is_pg0:
|
||||
kwargs: dict[str, object] = {"name": parsed.instance_name, "port": parsed.port}
|
||||
if parsed.username is not None:
|
||||
kwargs["username"] = parsed.username
|
||||
if parsed.password is not None:
|
||||
kwargs["password"] = parsed.password
|
||||
pg0 = EmbeddedPostgres(**kwargs)
|
||||
return await pg0.ensure_running()
|
||||
return db_url
|
||||
|
||||
@@ -197,6 +197,11 @@ def main():
|
||||
shared_pool = max(0, config.worker_max_slots - sum(reservations.values()))
|
||||
print(f" Slot reservations: {reservations_str}")
|
||||
print(f" Shared pool: {shared_pool}")
|
||||
if config.operation_retention_days == 0:
|
||||
print(" Operation retention: disabled (terminal rows and payloads are kept)")
|
||||
else:
|
||||
print(f" Operation retention: {config.operation_retention_days} days (terminal rows, payloads, and metadata)")
|
||||
print(f" Operation cleanup batch: {config.operation_cleanup_batch_size} rows/schema/cycle")
|
||||
print(f" HTTP server: {args.http_host}:{args.http_port}")
|
||||
print()
|
||||
|
||||
|
||||
@@ -37,6 +37,19 @@ def _metric_operation_label(operation_type: str | None) -> str:
|
||||
return operation_type or "unknown"
|
||||
|
||||
|
||||
def _updated_row_count(result: Any) -> int:
|
||||
"""Extract a row count from backend execute() results."""
|
||||
if isinstance(result, int):
|
||||
return result
|
||||
if isinstance(result, str):
|
||||
try:
|
||||
return int(result.rsplit(" ", 1)[-1])
|
||||
except (TypeError, ValueError):
|
||||
return 0
|
||||
rowcount = getattr(result, "rowcount", None)
|
||||
return rowcount if isinstance(rowcount, int) else 0
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from hindsight_api.engine.db.base import DatabaseBackend, DatabaseConnection
|
||||
from hindsight_api.extensions.tenant import TenantExtension
|
||||
@@ -209,6 +222,8 @@ class WorkerPoller:
|
||||
# Rotation offset for per-tenant fair claiming. Advances past the last
|
||||
# schema we serviced so a busy tenant can't monopolize the poll order.
|
||||
self._next_schema_idx: int = 0
|
||||
# Retention cleanup runs outside the claim loop. Keep one task per
|
||||
# poller so maintenance cannot overlap with itself or block slot refill.
|
||||
|
||||
@staticmethod
|
||||
def _normalize_poll_schema(schema: str | None) -> str | None:
|
||||
@@ -504,17 +519,20 @@ class WorkerPoller:
|
||||
return result
|
||||
|
||||
async def _mark_completed(self, operation_id: str, schema: str | None):
|
||||
"""Mark a task as completed."""
|
||||
"""Mark a processing task as completed, then propagate to parent if needed."""
|
||||
table = fq_table("async_operations", schema)
|
||||
async with self._backend.acquire() as conn:
|
||||
await conn.execute(
|
||||
f"""
|
||||
UPDATE {table}
|
||||
SET status = 'completed', completed_at = now(), updated_at = now()
|
||||
WHERE operation_id = $1
|
||||
""",
|
||||
operation_id,
|
||||
)
|
||||
async with conn.transaction():
|
||||
result = await conn.execute(
|
||||
f"""
|
||||
UPDATE {table}
|
||||
SET status = 'completed', completed_at = now(), updated_at = now()
|
||||
WHERE operation_id = $1 AND status = 'processing'
|
||||
""",
|
||||
operation_id,
|
||||
)
|
||||
if _updated_row_count(result):
|
||||
await self._maybe_update_parent_operation(operation_id, schema, conn)
|
||||
|
||||
async def _mark_failed(self, operation_id: str, error_message: str, schema: str | None):
|
||||
"""Mark a task as failed with error message, then propagate to parent if applicable."""
|
||||
@@ -749,6 +767,7 @@ class WorkerPoller:
|
||||
task.task_dict["_schema"] = task.schema
|
||||
await self._executor(task.task_dict)
|
||||
logger.debug(f"Task {task.operation_id} execution finished")
|
||||
await self._mark_completed(task.operation_id, task.schema)
|
||||
terminal_success = True
|
||||
except DeferOperation as e:
|
||||
# Deferral is not a terminal outcome — do not record a completion.
|
||||
@@ -938,6 +957,7 @@ class WorkerPoller:
|
||||
for task in tasks:
|
||||
await self.execute_task(task)
|
||||
|
||||
if tasks:
|
||||
# Continue immediately to claim more tasks (if slots available)
|
||||
continue
|
||||
|
||||
|
||||
@@ -51,7 +51,7 @@ dependencies = [
|
||||
"anthropic>=0.40.0",
|
||||
"typer>=0.9.0",
|
||||
"cohere>=5.0.0",
|
||||
"litellm>=1.83.14", # 1.82.7/1.82.8 had a supply chain compromise (yanked); 1.83.0+ also fixes GHSA-jjhc-v7c2-5hh6 / GHSA-53mr-6c8q-9789 / GHSA-pq44-5pcq-4r5g / GHSA-8cjq-wjmh-q42r
|
||||
"litellm>=1.84.0", # 1.82.7/1.82.8 had a supply chain compromise (yanked); 1.83.0+ also fixes GHSA-jjhc-v7c2-5hh6 / GHSA-53mr-6c8q-9789 / GHSA-pq44-5pcq-4r5g / GHSA-8cjq-wjmh-q42r; 1.84.0 fixes GHSA-4xpc-pv4p-pm3w
|
||||
"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)
|
||||
"winloop>=0.1.0; sys_platform == 'win32'",
|
||||
@@ -81,10 +81,11 @@ dependencies = [
|
||||
local-ml = [
|
||||
# Local ML models for embeddings/reranking
|
||||
"sentence-transformers>=3.3.0",
|
||||
"transformers>=4.53.0", # Security fixes for ReDoS vulnerabilities
|
||||
# transformers (incl. latest 5.x) hard-requires tokenizers<=0.23.0 via a
|
||||
# runtime check; without this cap an in-place upgrade can pull tokenizers
|
||||
# 0.23.1 and break local embeddings/reranker startup. See issue #2055.
|
||||
"transformers>=5.5.0", # ReDoS fixes; 5.5.0 clears GHSA-fgcw-684q-jj6r (LightGlue RCE)
|
||||
# transformers enforces tokenizers<=0.23.0 with a runtime check, but has
|
||||
# shipped metadata declaring a wider range than it actually enforces. Keep
|
||||
# this cap: without it an in-place upgrade can pull tokenizers 0.23.1 and
|
||||
# break local embeddings/reranker startup. See issue #2055.
|
||||
"tokenizers>=0.22.0,<=0.23.0",
|
||||
"torch>=2.6.0", # CVE fix for remote code execution
|
||||
"einops>=0.8.2",
|
||||
@@ -104,7 +105,7 @@ local-llm = [
|
||||
local-onnx = [
|
||||
# In-process ONNX Runtime embeddings without an Ollama/TEI sidecar
|
||||
"onnxruntime>=1.17.0",
|
||||
"transformers>=4.53.0",
|
||||
"transformers>=5.5.0", # 5.5.0 clears GHSA-fgcw-684q-jj6r (LightGlue RCE)
|
||||
"tokenizers>=0.22.0,<=0.23.0", # See issue #2055 (transformers caps tokenizers<=0.23.0)
|
||||
"huggingface-hub>=0.20.0",
|
||||
"numpy>=1.26.0",
|
||||
|
||||
@@ -24,6 +24,33 @@ from dotenv import load_dotenv
|
||||
# per worker process. Guarded so slim/no-torch environments still collect.
|
||||
try:
|
||||
import torch # noqa: F401 # eager one-time init; see comment above
|
||||
|
||||
# Same class of problem, different torch module. transformers' lazy loader
|
||||
# imports `torch._inductor.test_operators` while resolving classes such as
|
||||
# AutoModelForSequenceClassification / GenerationMixin (exercised by the
|
||||
# cross-encoder / reranker tests). That module registers an `_inductor_test`
|
||||
# TORCH_LIBRARY namespace at module-body level, and under pytest-xdist its
|
||||
# body can execute twice, raising "Only a single TORCH_LIBRARY can be used
|
||||
# to register the namespace _inductor_test". The failure surfaces on
|
||||
# whichever shard runs the reranker tests, masked by transformers as a
|
||||
# misleading "sentence-transformers is required for LocalSTEmbeddings"
|
||||
# ImportError. Seed it once here so the later lazy import is a sys.modules
|
||||
# cache hit and the body never re-executes.
|
||||
import torch._inductor.test_operators # noqa: F401 # see comment above
|
||||
|
||||
# Seed the rest of the native embedding/reranker stack the same way, and for
|
||||
# the same reason. transformers and safetensors/tokenizers ship PyO3/Rust
|
||||
# and C extensions whose module bodies are not safe to execute twice
|
||||
# (safetensors raises "PyO3 modules ... may only be initialized once per
|
||||
# interpreter process"). When these are first imported lazily from inside a
|
||||
# fixture's event loop / sentence-transformers' thread pools, or re-executed
|
||||
# by transformers' lazy-loader retry path, the second init aborts and — like
|
||||
# the torch cases above — is re-raised as a misleading
|
||||
# "sentence-transformers is required" ImportError on the reranker shard.
|
||||
# Importing the whole chain here (single-threaded, at collection time) puts
|
||||
# every submodule in sys.modules so later imports are cache hits.
|
||||
import transformers # noqa: F401 # seeds safetensors/tokenizers once
|
||||
import sentence_transformers # noqa: F401
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
@@ -136,7 +163,7 @@ def pg0_db_url(db_url, tmp_path_factory, worker_id):
|
||||
from hindsight_api.pg0 import parse_pg0_url as _parse_pg0_url
|
||||
|
||||
# Determine pg0 instance name/port from db_url (if it's a pg0:// URL) or use defaults
|
||||
if db_url and not _parse_pg0_url(db_url)[0]:
|
||||
if db_url and not _parse_pg0_url(db_url).is_pg0:
|
||||
# Plain postgresql:// URL - use it directly but still run migrations
|
||||
from hindsight_api.migrations import run_migrations
|
||||
|
||||
@@ -144,9 +171,9 @@ def pg0_db_url(db_url, tmp_path_factory, worker_id):
|
||||
return db_url
|
||||
|
||||
if db_url:
|
||||
_, pg0_name, pg0_port = _parse_pg0_url(db_url)
|
||||
pg0_instance_name = pg0_name or DEFAULT_PG0_INSTANCE_NAME
|
||||
pg0_instance_port = pg0_port or DEFAULT_PG0_PORT
|
||||
_parsed = _parse_pg0_url(db_url)
|
||||
pg0_instance_name = _parsed.instance_name or DEFAULT_PG0_INSTANCE_NAME
|
||||
pg0_instance_port = _parsed.port or DEFAULT_PG0_PORT
|
||||
else:
|
||||
pg0_instance_name = DEFAULT_PG0_INSTANCE_NAME
|
||||
pg0_instance_port = DEFAULT_PG0_PORT
|
||||
|
||||
@@ -0,0 +1,273 @@
|
||||
"""Anthropic Message Batches support for the provider batch interface.
|
||||
|
||||
The engine's batch path (retain fact extraction, gated on
|
||||
``retain_batch_enabled``) speaks the OpenAI batch wire shape: JSONL entries
|
||||
with ``custom_id``/``method``/``url``/``body`` going in, and
|
||||
``response.body.choices[0].message.content`` (+ OpenAI-keyed ``usage``) coming
|
||||
out. ``AnthropicLLM`` translates both directions onto the Message Batches API,
|
||||
which bills all token usage at 50% of standard price.
|
||||
|
||||
Translation rules mirror the provider's synchronous ``call()`` path:
|
||||
- system messages fold into the ``system`` param;
|
||||
- ``max_completion_tokens`` becomes ``max_tokens`` (default 4096);
|
||||
- ``temperature`` is dropped (the sync path never sends it either — current
|
||||
Claude models reject non-default sampling params);
|
||||
- ``response_format`` with ``strict=True`` becomes a single forced tool_use
|
||||
tool (native constrained decoding, issue #1002); non-strict injects the
|
||||
schema into the system prompt and expects JSON text back.
|
||||
"""
|
||||
|
||||
import json
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
pytestmark = pytest.mark.asyncio
|
||||
|
||||
|
||||
def _make_provider():
|
||||
with patch("anthropic.AsyncAnthropic") as mock_client_cls:
|
||||
mock_client_cls.return_value = MagicMock()
|
||||
from hindsight_api.engine.providers.anthropic_llm import AnthropicLLM
|
||||
|
||||
provider = AnthropicLLM(
|
||||
provider="anthropic",
|
||||
api_key="fake-key",
|
||||
base_url="",
|
||||
model="claude-sonnet-5",
|
||||
)
|
||||
provider._client = MagicMock()
|
||||
return provider
|
||||
|
||||
|
||||
_SCHEMA = {
|
||||
"type": "object",
|
||||
"properties": {"facts": {"type": "array", "items": {"type": "string"}}},
|
||||
"required": ["facts"],
|
||||
}
|
||||
|
||||
|
||||
def _openai_request(custom_id: str, *, strict: bool = True, temperature: float | None = 0.1) -> dict:
|
||||
body = {
|
||||
"model": "claude-sonnet-5",
|
||||
"messages": [
|
||||
{"role": "system", "content": "Extract facts."},
|
||||
{"role": "user", "content": f"Text for {custom_id}"},
|
||||
],
|
||||
"max_completion_tokens": 2000,
|
||||
"response_format": {
|
||||
"type": "json_schema",
|
||||
"json_schema": {"name": "facts", "schema": _SCHEMA, "strict": strict},
|
||||
},
|
||||
}
|
||||
if temperature is not None:
|
||||
body["temperature"] = temperature
|
||||
return {"custom_id": custom_id, "method": "POST", "url": "/v1/chat/completions", "body": body}
|
||||
|
||||
|
||||
def _batch(status: str = "in_progress", **counts) -> SimpleNamespace:
|
||||
defaults = {"processing": 0, "succeeded": 0, "errored": 0, "canceled": 0, "expired": 0}
|
||||
defaults.update(counts)
|
||||
return SimpleNamespace(
|
||||
id="msgbatch_test1",
|
||||
processing_status=status,
|
||||
created_at="2026-07-08T00:00:00Z",
|
||||
ended_at="2026-07-08T00:30:00Z" if status == "ended" else None,
|
||||
request_counts=SimpleNamespace(**defaults),
|
||||
)
|
||||
|
||||
|
||||
class _AsyncIter:
|
||||
def __init__(self, items):
|
||||
self._items = list(items)
|
||||
|
||||
def __aiter__(self):
|
||||
self._iter = iter(self._items)
|
||||
return self
|
||||
|
||||
async def __anext__(self):
|
||||
try:
|
||||
return next(self._iter)
|
||||
except StopIteration:
|
||||
raise StopAsyncIteration from None
|
||||
|
||||
|
||||
def _succeeded_entry(custom_id: str, tool_input: dict) -> SimpleNamespace:
|
||||
block = SimpleNamespace(type="tool_use", name="structured_response", input=tool_input, text=None)
|
||||
message = SimpleNamespace(
|
||||
content=[block],
|
||||
usage=SimpleNamespace(input_tokens=100, output_tokens=40, cache_read_input_tokens=0),
|
||||
stop_reason="tool_use",
|
||||
)
|
||||
return SimpleNamespace(custom_id=custom_id, result=SimpleNamespace(type="succeeded", message=message))
|
||||
|
||||
|
||||
def _errored_entry(custom_id: str) -> SimpleNamespace:
|
||||
error = SimpleNamespace(type="invalid_request", message="bad request")
|
||||
return SimpleNamespace(custom_id=custom_id, result=SimpleNamespace(type="errored", error=error))
|
||||
|
||||
|
||||
async def test_supports_batch_api():
|
||||
provider = _make_provider()
|
||||
assert await provider.supports_batch_api() is True
|
||||
|
||||
|
||||
async def test_submit_batch_translates_openai_requests():
|
||||
provider = _make_provider()
|
||||
provider._client.messages.batches.create = AsyncMock(return_value=_batch("in_progress", processing=2))
|
||||
|
||||
requests = [_openai_request("chunk_0"), _openai_request("chunk_1")]
|
||||
metadata = await provider.submit_batch(requests)
|
||||
|
||||
provider._client.messages.batches.create.assert_awaited_once()
|
||||
submitted = provider._client.messages.batches.create.await_args.kwargs["requests"]
|
||||
assert [r["custom_id"] for r in submitted] == ["chunk_0", "chunk_1"]
|
||||
|
||||
params = submitted[0]["params"]
|
||||
assert params["model"] == "claude-sonnet-5"
|
||||
# System message folded into the system param (as the cached block list
|
||||
# the sync call() path sends), not left in messages.
|
||||
assert "Extract facts." in params["system"][0]["text"]
|
||||
assert all(m["role"] != "system" for m in params["messages"])
|
||||
assert params["messages"] == [{"role": "user", "content": "Text for chunk_0"}]
|
||||
assert params["max_tokens"] == 2000
|
||||
# temperature is dropped, mirroring the sync call() path.
|
||||
assert "temperature" not in params
|
||||
# strict=True → forced tool_use (native constrained decoding).
|
||||
assert params["tools"][0]["input_schema"] == _SCHEMA
|
||||
assert params["tool_choice"] == {"type": "tool", "name": "structured_response"}
|
||||
|
||||
assert metadata["batch_id"] == "msgbatch_test1"
|
||||
assert metadata["status"] == "in_progress"
|
||||
assert metadata["request_count"] == 2
|
||||
|
||||
|
||||
async def test_submit_batch_non_strict_schema_injects_into_system():
|
||||
provider = _make_provider()
|
||||
provider._client.messages.batches.create = AsyncMock(return_value=_batch("in_progress", processing=1))
|
||||
|
||||
await provider.submit_batch([_openai_request("chunk_0", strict=False)])
|
||||
|
||||
params = provider._client.messages.batches.create.await_args.kwargs["requests"][0]["params"]
|
||||
assert "tools" not in params
|
||||
assert "tool_choice" not in params
|
||||
# Schema is injected into the system prompt for JSON-text output —
|
||||
# inside the cached block, so the injection is part of the cached prefix.
|
||||
assert "facts" in params["system"][0]["text"]
|
||||
assert "valid JSON" in params["system"][0]["text"]
|
||||
|
||||
|
||||
async def test_submit_batch_system_carries_cache_control_marker():
|
||||
"""Batch items share their system prompt, so it gets the cache marker.
|
||||
|
||||
Mirrors the sync ``call()`` one-shot rule: system is the sole cache
|
||||
breakpoint. Within a Message Batch every request carries the same fact-
|
||||
extraction system prompt, so the first request's cache write serves the
|
||||
rest as best-effort reads (and stacks with the 50% batch discount).
|
||||
"""
|
||||
provider = _make_provider()
|
||||
provider._client.messages.batches.create = AsyncMock(return_value=_batch("in_progress", processing=1))
|
||||
|
||||
await provider.submit_batch([_openai_request("chunk_0")])
|
||||
|
||||
params = provider._client.messages.batches.create.await_args.kwargs["requests"][0]["params"]
|
||||
assert params["system"] == [{"type": "text", "text": "Extract facts.", "cache_control": {"type": "ephemeral"}}]
|
||||
# One-shot items: no end-marker on messages (that breakpoint only pays
|
||||
# off on the sync tool loop, where the next iteration reads it back).
|
||||
assert "cache_control" not in json.dumps(params["messages"])
|
||||
|
||||
|
||||
async def test_submit_batch_without_system_message_sends_no_system_param():
|
||||
provider = _make_provider()
|
||||
provider._client.messages.batches.create = AsyncMock(return_value=_batch("in_progress", processing=1))
|
||||
|
||||
body = {
|
||||
"model": "claude-sonnet-5",
|
||||
"messages": [{"role": "user", "content": "no system here"}],
|
||||
"max_completion_tokens": 1000,
|
||||
}
|
||||
request = {"custom_id": "chunk_0", "method": "POST", "url": "/v1/chat/completions", "body": body}
|
||||
await provider.submit_batch([request])
|
||||
|
||||
params = provider._client.messages.batches.create.await_args.kwargs["requests"][0]["params"]
|
||||
assert "system" not in params
|
||||
assert "cache_control" not in json.dumps(params["messages"])
|
||||
|
||||
|
||||
async def test_get_batch_status_in_progress():
|
||||
provider = _make_provider()
|
||||
provider._client.messages.batches.retrieve = AsyncMock(
|
||||
return_value=_batch("in_progress", processing=3, succeeded=1)
|
||||
)
|
||||
|
||||
status = await provider.get_batch_status("msgbatch_test1")
|
||||
|
||||
assert status["batch_id"] == "msgbatch_test1"
|
||||
assert status["status"] == "in_progress"
|
||||
assert status["request_counts"]["total"] == 4
|
||||
assert status["request_counts"]["completed"] == 1
|
||||
|
||||
|
||||
async def test_get_batch_status_ended_maps_to_completed():
|
||||
"""The engine's poll loop breaks on the OpenAI-vocabulary status 'completed'."""
|
||||
provider = _make_provider()
|
||||
provider._client.messages.batches.retrieve = AsyncMock(return_value=_batch("ended", succeeded=3, errored=1))
|
||||
|
||||
status = await provider.get_batch_status("msgbatch_test1")
|
||||
|
||||
assert status["status"] == "completed"
|
||||
assert status["request_counts"]["total"] == 4
|
||||
assert status["request_counts"]["completed"] == 4
|
||||
assert status["request_counts"]["failed"] == 1
|
||||
assert status["completed_at"] == "2026-07-08T00:30:00Z"
|
||||
|
||||
|
||||
async def test_retrieve_batch_results_translates_to_openai_shape():
|
||||
provider = _make_provider()
|
||||
provider._client.messages.batches.retrieve = AsyncMock(return_value=_batch("ended", succeeded=1, errored=1))
|
||||
entries = [
|
||||
_succeeded_entry("chunk_0", {"facts": ["Alice is an engineer."]}),
|
||||
_errored_entry("chunk_1"),
|
||||
]
|
||||
provider._client.messages.batches.results = AsyncMock(return_value=_AsyncIter(entries))
|
||||
|
||||
results = await provider.retrieve_batch_results("msgbatch_test1")
|
||||
|
||||
by_id = {r["custom_id"]: r for r in results}
|
||||
ok = by_id["chunk_0"]
|
||||
body = ok["response"]["body"]
|
||||
# The engine reads choices[0].message.content and json.loads() it.
|
||||
assert json.loads(body["choices"][0]["message"]["content"]) == {"facts": ["Alice is an engineer."]}
|
||||
# Usage arrives under the OpenAI key names the engine sums.
|
||||
assert body["usage"] == {"prompt_tokens": 100, "completion_tokens": 40, "total_tokens": 140}
|
||||
|
||||
failed = by_id["chunk_1"]
|
||||
assert failed["error"]
|
||||
assert "response" not in failed
|
||||
|
||||
|
||||
async def test_retrieve_batch_results_text_content_passthrough():
|
||||
"""Non-strict requests come back as text blocks; concatenate them as content."""
|
||||
provider = _make_provider()
|
||||
provider._client.messages.batches.retrieve = AsyncMock(return_value=_batch("ended", succeeded=1))
|
||||
text_block = SimpleNamespace(type="text", text='{"facts": []}')
|
||||
message = SimpleNamespace(
|
||||
content=[text_block],
|
||||
usage=SimpleNamespace(input_tokens=10, output_tokens=5, cache_read_input_tokens=0),
|
||||
stop_reason="end_turn",
|
||||
)
|
||||
entry = SimpleNamespace(custom_id="chunk_0", result=SimpleNamespace(type="succeeded", message=message))
|
||||
provider._client.messages.batches.results = AsyncMock(return_value=_AsyncIter([entry]))
|
||||
|
||||
results = await provider.retrieve_batch_results("msgbatch_test1")
|
||||
|
||||
assert results[0]["response"]["body"]["choices"][0]["message"]["content"] == '{"facts": []}'
|
||||
|
||||
|
||||
async def test_retrieve_batch_results_raises_when_not_ended():
|
||||
provider = _make_provider()
|
||||
provider._client.messages.batches.retrieve = AsyncMock(return_value=_batch("in_progress", processing=2))
|
||||
|
||||
with pytest.raises(ValueError, match="not completed"):
|
||||
await provider.retrieve_batch_results("msgbatch_test1")
|
||||
@@ -0,0 +1,181 @@
|
||||
"""Anthropic prompt caching via inline cache_control markers.
|
||||
|
||||
``LLMInterface.get_or_create_cached_prefix`` documents Anthropic as an
|
||||
"inline-marker provider": rather than returning an explicit cache handle, the
|
||||
provider marks the reusable prefix inside ``call`` / ``call_with_tools`` with
|
||||
``cache_control`` breakpoints. Cache reads bill at ~10% of the base input
|
||||
price; a marker below the model's minimum cacheable prefix is silently
|
||||
ignored by the API (no premium), so marking is safe unconditionally.
|
||||
|
||||
Two breakpoints (of the 4 allowed):
|
||||
- the system prompt, in both entry points — it is stable per scope (fact
|
||||
extraction reuses it across every chunk; reflect/consolidation put their
|
||||
stable instructions there), so tools+system cache across calls;
|
||||
- the last message content block, in ``call_with_tools`` only — the reflect
|
||||
agent loop resends the whole growing conversation each iteration, so each
|
||||
request's end-marker becomes the next iteration's cache read point.
|
||||
"""
|
||||
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from pydantic import BaseModel
|
||||
|
||||
pytestmark = pytest.mark.asyncio
|
||||
|
||||
EPHEMERAL = {"type": "ephemeral"}
|
||||
|
||||
|
||||
def _make_provider():
|
||||
with patch("anthropic.AsyncAnthropic") as mock_client_cls:
|
||||
mock_client_cls.return_value = MagicMock()
|
||||
from hindsight_api.engine.providers.anthropic_llm import AnthropicLLM
|
||||
|
||||
provider = AnthropicLLM(
|
||||
provider="anthropic",
|
||||
api_key="fake-key",
|
||||
base_url="",
|
||||
model="claude-sonnet-5",
|
||||
)
|
||||
provider._client = MagicMock()
|
||||
return provider
|
||||
|
||||
|
||||
def _text_response(text: str = "ok"):
|
||||
block = MagicMock()
|
||||
block.type = "text"
|
||||
block.text = text
|
||||
resp = MagicMock()
|
||||
resp.content = [block]
|
||||
resp.usage = MagicMock(input_tokens=10, output_tokens=2, cache_read_input_tokens=0)
|
||||
resp.stop_reason = "end_turn"
|
||||
return resp
|
||||
|
||||
|
||||
def _tool_response():
|
||||
resp = MagicMock()
|
||||
resp.content = []
|
||||
resp.usage = MagicMock(input_tokens=10, output_tokens=2, cache_read_input_tokens=0)
|
||||
resp.stop_reason = "end_turn"
|
||||
return resp
|
||||
|
||||
|
||||
class _Out(BaseModel):
|
||||
facts: list[str]
|
||||
|
||||
|
||||
async def test_call_marks_system_prompt_for_caching():
|
||||
provider = _make_provider()
|
||||
provider._client.messages.create = AsyncMock(return_value=_text_response())
|
||||
|
||||
with patch("hindsight_api.engine.providers.anthropic_llm.get_metrics_collector"):
|
||||
await provider.call(
|
||||
messages=[
|
||||
{"role": "system", "content": "Stable extraction instructions."},
|
||||
{"role": "user", "content": "Chunk text."},
|
||||
],
|
||||
scope="test",
|
||||
max_retries=0,
|
||||
)
|
||||
|
||||
params = provider._client.messages.create.await_args.kwargs
|
||||
assert params["system"] == [{"type": "text", "text": "Stable extraction instructions.", "cache_control": EPHEMERAL}]
|
||||
# User messages are untouched in call() — one-shot calls share no
|
||||
# conversation prefix with each other, only the system prompt.
|
||||
assert params["messages"] == [{"role": "user", "content": "Chunk text."}]
|
||||
|
||||
|
||||
async def test_call_non_strict_schema_lands_inside_cached_system_block():
|
||||
"""Schema injection happens before marking, so the marked block includes it."""
|
||||
provider = _make_provider()
|
||||
provider._client.messages.create = AsyncMock(return_value=_text_response('{"facts": []}'))
|
||||
|
||||
with patch("hindsight_api.engine.providers.anthropic_llm.get_metrics_collector"):
|
||||
await provider.call(
|
||||
messages=[
|
||||
{"role": "system", "content": "Extract."},
|
||||
{"role": "user", "content": "Text."},
|
||||
],
|
||||
response_format=_Out,
|
||||
scope="test",
|
||||
max_retries=0,
|
||||
)
|
||||
|
||||
params = provider._client.messages.create.await_args.kwargs
|
||||
assert len(params["system"]) == 1
|
||||
system_block = params["system"][0]
|
||||
assert system_block["cache_control"] == EPHEMERAL
|
||||
assert "Extract." in system_block["text"]
|
||||
assert "valid JSON" in system_block["text"]
|
||||
|
||||
|
||||
async def test_call_without_system_prompt_sends_no_system_param():
|
||||
provider = _make_provider()
|
||||
provider._client.messages.create = AsyncMock(return_value=_text_response())
|
||||
|
||||
with patch("hindsight_api.engine.providers.anthropic_llm.get_metrics_collector"):
|
||||
await provider.call(
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
scope="test",
|
||||
max_retries=0,
|
||||
)
|
||||
|
||||
assert "system" not in provider._client.messages.create.await_args.kwargs
|
||||
|
||||
|
||||
async def test_call_with_tools_marks_system_and_last_message():
|
||||
provider = _make_provider()
|
||||
provider._client.messages.create = AsyncMock(return_value=_tool_response())
|
||||
|
||||
with patch("hindsight_api.engine.providers.anthropic_llm.get_metrics_collector"):
|
||||
await provider.call_with_tools(
|
||||
messages=[
|
||||
{"role": "system", "content": "Reflect agent instructions."},
|
||||
{"role": "user", "content": "Question?"},
|
||||
{"role": "assistant", "content": "Working on it."},
|
||||
{"role": "user", "content": "Latest turn."},
|
||||
],
|
||||
tools=[{"function": {"name": "recall", "description": "d", "parameters": {"type": "object"}}}],
|
||||
max_retries=0,
|
||||
)
|
||||
|
||||
params = provider._client.messages.create.await_args.kwargs
|
||||
assert params["system"] == [{"type": "text", "text": "Reflect agent instructions.", "cache_control": EPHEMERAL}]
|
||||
|
||||
messages = params["messages"]
|
||||
# Earlier messages carry no markers — only the final block gets one, so
|
||||
# the next iteration of the agent loop reads the whole prefix from cache.
|
||||
assert messages[0] == {"role": "user", "content": "Question?"}
|
||||
assert messages[1] == {"role": "assistant", "content": "Working on it."}
|
||||
assert messages[2]["content"] == [{"type": "text", "text": "Latest turn.", "cache_control": EPHEMERAL}]
|
||||
|
||||
|
||||
async def test_call_with_tools_marks_last_block_of_tool_result_message():
|
||||
"""Tool-result turns arrive as block lists; the marker goes on the last block."""
|
||||
provider = _make_provider()
|
||||
provider._client.messages.create = AsyncMock(return_value=_tool_response())
|
||||
|
||||
with patch("hindsight_api.engine.providers.anthropic_llm.get_metrics_collector"):
|
||||
await provider.call_with_tools(
|
||||
messages=[
|
||||
{"role": "user", "content": "Question?"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"tool_calls": [
|
||||
{"id": "t1", "function": {"name": "recall", "arguments": "{}"}},
|
||||
{"id": "t2", "function": {"name": "recall", "arguments": "{}"}},
|
||||
],
|
||||
},
|
||||
{"role": "tool", "tool_call_id": "t1", "content": "result one"},
|
||||
{"role": "tool", "tool_call_id": "t2", "content": "result two"},
|
||||
],
|
||||
tools=[{"function": {"name": "recall", "description": "d", "parameters": {"type": "object"}}}],
|
||||
max_retries=0,
|
||||
)
|
||||
|
||||
messages = provider._client.messages.create.await_args.kwargs["messages"]
|
||||
last_blocks = messages[-1]["content"]
|
||||
assert last_blocks[-1]["type"] == "tool_result"
|
||||
assert last_blocks[-1]["cache_control"] == EPHEMERAL
|
||||
# The earlier tool-result message is unmarked.
|
||||
assert all("cache_control" not in block for block in messages[-2]["content"])
|
||||
@@ -107,5 +107,8 @@ async def test_non_strict_keeps_text_injection_fallback():
|
||||
)
|
||||
kwargs = provider._client.messages.create.call_args.kwargs
|
||||
assert "tools" not in kwargs # no forced tool when not strict
|
||||
assert "valid JSON matching this schema" in (kwargs.get("system") or "")
|
||||
# system is a cache_control-marked block list; the schema text-injection
|
||||
# lands inside the (single) block.
|
||||
system_text = "".join(block["text"] for block in (kwargs.get("system") or []))
|
||||
assert "valid JSON matching this schema" in system_text
|
||||
assert isinstance(result, _Decision)
|
||||
|
||||
@@ -2,6 +2,7 @@
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import uuid
|
||||
|
||||
import pytest
|
||||
@@ -515,6 +516,105 @@ async def test_retain_outcome_metadata_records_zero_counts(memory, request_conte
|
||||
assert "extraction_errors_sample" not in parent["result_metadata"]
|
||||
|
||||
|
||||
async def _seed_retain_op_with_errors(pool, bank_id: str, error_count: int) -> uuid.UUID:
|
||||
"""Insert a pending retain operation whose outcome metadata records extraction errors."""
|
||||
operation_id = uuid.uuid4()
|
||||
await pool.execute(
|
||||
"""
|
||||
INSERT INTO async_operations (operation_id, bank_id, operation_type, result_metadata, status)
|
||||
VALUES ($1, $2, $3, $4, $5)
|
||||
""",
|
||||
operation_id,
|
||||
bank_id,
|
||||
"retain",
|
||||
json.dumps(
|
||||
{
|
||||
"unit_ids_count": 3,
|
||||
"extraction_errors_count": error_count,
|
||||
"extraction_errors_sample": ["chunk 2 failed to parse"],
|
||||
}
|
||||
),
|
||||
"pending",
|
||||
)
|
||||
return operation_id
|
||||
|
||||
|
||||
async def _op_row(pool, operation_id: uuid.UUID):
|
||||
return await pool.fetchrow(
|
||||
"SELECT status, error_message FROM async_operations WHERE operation_id = $1",
|
||||
operation_id,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_completion_marks_failed_when_flag_on_and_errors_present(memory):
|
||||
"""With the escape hatch on, a retain that dropped facts ends 'failed' (issue #2700)."""
|
||||
from hindsight_api.config import ENV_FAIL_ON_EXTRACTION_ERRORS, clear_config_cache
|
||||
|
||||
bank_id = "test_fail_on_extraction_errors_on"
|
||||
pool = await memory._get_pool()
|
||||
await _ensure_bank(pool, bank_id)
|
||||
operation_id = await _seed_retain_op_with_errors(pool, bank_id, error_count=2)
|
||||
|
||||
os.environ[ENV_FAIL_ON_EXTRACTION_ERRORS] = "true"
|
||||
clear_config_cache()
|
||||
try:
|
||||
await memory._mark_operation_completed(str(operation_id))
|
||||
finally:
|
||||
del os.environ[ENV_FAIL_ON_EXTRACTION_ERRORS]
|
||||
clear_config_cache()
|
||||
|
||||
row = await _op_row(pool, operation_id)
|
||||
assert row["status"] == "failed"
|
||||
assert row["error_message"] is not None
|
||||
assert "2" in row["error_message"]
|
||||
assert "extraction error" in row["error_message"].lower()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_completion_stays_completed_when_flag_off(memory):
|
||||
"""Default behavior is preserved: extraction errors still complete the operation."""
|
||||
from hindsight_api.config import ENV_FAIL_ON_EXTRACTION_ERRORS, clear_config_cache
|
||||
|
||||
bank_id = "test_fail_on_extraction_errors_off"
|
||||
pool = await memory._get_pool()
|
||||
await _ensure_bank(pool, bank_id)
|
||||
operation_id = await _seed_retain_op_with_errors(pool, bank_id, error_count=2)
|
||||
|
||||
os.environ.pop(ENV_FAIL_ON_EXTRACTION_ERRORS, None)
|
||||
clear_config_cache()
|
||||
try:
|
||||
await memory._mark_operation_completed(str(operation_id))
|
||||
finally:
|
||||
clear_config_cache()
|
||||
|
||||
row = await _op_row(pool, operation_id)
|
||||
assert row["status"] == "completed"
|
||||
assert row["error_message"] is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_completion_completed_when_flag_on_but_no_errors(memory):
|
||||
"""The flag only fails operations that actually accumulated extraction errors."""
|
||||
from hindsight_api.config import ENV_FAIL_ON_EXTRACTION_ERRORS, clear_config_cache
|
||||
|
||||
bank_id = "test_fail_on_extraction_errors_none"
|
||||
pool = await memory._get_pool()
|
||||
await _ensure_bank(pool, bank_id)
|
||||
operation_id = await _seed_retain_op_with_errors(pool, bank_id, error_count=0)
|
||||
|
||||
os.environ[ENV_FAIL_ON_EXTRACTION_ERRORS] = "true"
|
||||
clear_config_cache()
|
||||
try:
|
||||
await memory._mark_operation_completed(str(operation_id))
|
||||
finally:
|
||||
del os.environ[ENV_FAIL_ON_EXTRACTION_ERRORS]
|
||||
clear_config_cache()
|
||||
|
||||
row = await _op_row(pool, operation_id)
|
||||
assert row["status"] == "completed"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retain_records_user_provided_document_ids(memory, request_context):
|
||||
"""User-supplied document_ids land in child op result_metadata.document_ids."""
|
||||
|
||||
@@ -17,14 +17,17 @@ import pytest
|
||||
from pydantic import BaseModel
|
||||
|
||||
from hindsight_api.engine.bank_attribution import apply_bank_attribution
|
||||
from hindsight_api.engine.cross_encoder import LiteLLMCrossEncoder
|
||||
from hindsight_api.engine.embeddings import OpenAIEmbeddings
|
||||
from hindsight_api.engine.memory_engine import (
|
||||
MemoryEngine,
|
||||
_bind_bank_id,
|
||||
_current_bank_id,
|
||||
get_current_bank_id,
|
||||
)
|
||||
from hindsight_api.engine.providers.openai_compatible_llm import OpenAICompatibleLLM
|
||||
from hindsight_api.engine.retain.embedding_utils import generate_embeddings_batch
|
||||
from hindsight_api.models import RequestContext
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
@@ -56,31 +59,9 @@ class TestBankContextVar:
|
||||
def test_default_is_none(self):
|
||||
assert get_current_bank_id() is None
|
||||
|
||||
def test_set_and_reset(self):
|
||||
token = _current_bank_id.set("user-42")
|
||||
try:
|
||||
assert get_current_bank_id() == "user-42"
|
||||
finally:
|
||||
_current_bank_id.reset(token)
|
||||
assert get_current_bank_id() is None
|
||||
|
||||
def test_reset_runs_even_on_exception(self):
|
||||
"""A finally-based reset must unwind the binding even when the body raises."""
|
||||
token = _current_bank_id.set("user-boom")
|
||||
try:
|
||||
with pytest.raises(ValueError):
|
||||
try:
|
||||
assert get_current_bank_id() == "user-boom"
|
||||
raise ValueError("boom")
|
||||
finally:
|
||||
_current_bank_id.reset(token)
|
||||
finally:
|
||||
pass
|
||||
assert get_current_bank_id() is None
|
||||
|
||||
|
||||
class TestBindBankIdDecorator:
|
||||
"""The engine binds the bank via @_bind_bank_id on recall/retain/batch/task methods."""
|
||||
"""The engine binds the bank via @_bind_bank_id on recall/retain/batch/reflect/task methods."""
|
||||
|
||||
async def test_binds_named_arg_positional_and_keyword(self):
|
||||
@_bind_bank_id()
|
||||
@@ -117,6 +98,21 @@ class TestBindBankIdDecorator:
|
||||
|
||||
assert await op(12345) is None
|
||||
|
||||
async def test_reflect_async_binds_and_resets_its_bank_argument(self):
|
||||
engine = object.__new__(MemoryEngine)
|
||||
engine._reflect_llm_config = None
|
||||
observed_bank_ids: list[str | None] = []
|
||||
|
||||
with patch(
|
||||
"hindsight_api.engine.memory_engine.sanitize_text",
|
||||
side_effect=lambda value: observed_bank_ids.append(get_current_bank_id()) or value,
|
||||
):
|
||||
with pytest.raises(ValueError, match="Memory LLM API key not set"):
|
||||
await engine.reflect_async("user-reflect", "question", request_context=RequestContext())
|
||||
|
||||
assert observed_bank_ids == ["user-reflect", "user-reflect"]
|
||||
assert get_current_bank_id() is None
|
||||
|
||||
|
||||
# ── LLM provider: user injection ──────────────────────────────────────────────
|
||||
|
||||
@@ -258,27 +254,22 @@ def test_embeddings_user_injected_when_flag_on_and_bank_set():
|
||||
assert captured[0]["user"] == "user-emb"
|
||||
|
||||
|
||||
def test_embeddings_user_not_injected_when_flag_off():
|
||||
_set_flag(False)
|
||||
emb = _openai_embeddings()
|
||||
captured: list[dict] = []
|
||||
emb._client = _fake_embed_client(captured)
|
||||
token = _current_bank_id.set("user-emb")
|
||||
try:
|
||||
emb.encode(["hello"])
|
||||
finally:
|
||||
_current_bank_id.reset(token)
|
||||
assert "user" not in captured[0]
|
||||
async def test_litellm_proxy_sends_bank_header():
|
||||
encoder = LiteLLMCrossEncoder(api_base="https://rerank.example", model="will-memory-rerank")
|
||||
response = SimpleNamespace(
|
||||
raise_for_status=lambda: None,
|
||||
json=lambda: {"results": [{"index": 0, "relevance_score": 0.91}]},
|
||||
)
|
||||
encoder._async_client = SimpleNamespace(post=AsyncMock(return_value=response))
|
||||
|
||||
with patch(
|
||||
"hindsight_api.engine.cross_encoder.reranker_bank_attribution_headers",
|
||||
return_value={"X-Hindsight-Bank-Id": "bank-litellm-proxy"},
|
||||
):
|
||||
scores = await encoder.predict([("query", "document")])
|
||||
|
||||
def test_embeddings_user_not_injected_when_bank_unset():
|
||||
_set_flag(True)
|
||||
emb = _openai_embeddings()
|
||||
captured: list[dict] = []
|
||||
emb._client = _fake_embed_client(captured)
|
||||
assert get_current_bank_id() is None
|
||||
emb.encode(["hello"])
|
||||
assert "user" not in captured[0]
|
||||
assert scores == [0.91]
|
||||
assert encoder._async_client.post.call_args.kwargs["headers"] == {"X-Hindsight-Bank-Id": "bank-litellm-proxy"}
|
||||
|
||||
|
||||
# ── Executor context propagation ──────────────────────────────────────────────
|
||||
|
||||
@@ -2,6 +2,7 @@
|
||||
Config wiring for per-bank attribution and the configurable OpenRouter rerank URL.
|
||||
|
||||
- HINDSIGHT_API_LLM_SEND_BANK_AS_USER (default off, opt-in bool)
|
||||
- HINDSIGHT_API_RERANKER_SEND_BANK_AS_HEADER (default off, opt-in bool)
|
||||
- HINDSIGHT_API_RERANKER_OPENROUTER_BASE_URL (default = previously hardcoded URL)
|
||||
|
||||
Deterministic, no network.
|
||||
@@ -12,7 +13,9 @@ from dataclasses import fields
|
||||
from unittest.mock import patch
|
||||
|
||||
from hindsight_api.config import DEFAULT_RERANKER_OPENROUTER_BASE_URL, HindsightConfig
|
||||
from hindsight_api.engine.bank_attribution import reranker_bank_attribution_headers
|
||||
from hindsight_api.engine.cross_encoder import create_cross_encoder_from_env
|
||||
from hindsight_api.engine.memory_engine import _current_bank_id
|
||||
|
||||
|
||||
def _restore_env(saved: dict[str, str | None]) -> None:
|
||||
@@ -77,14 +80,29 @@ class TestSendBankAsUserConfig:
|
||||
finally:
|
||||
_restore_env(saved)
|
||||
|
||||
def test_one_enables(self):
|
||||
def test_reranker_bank_header_default_false_and_true(self):
|
||||
from hindsight_api.config import clear_config_cache
|
||||
|
||||
saved = {"HINDSIGHT_API_LLM_SEND_BANK_AS_USER": os.environ.get("HINDSIGHT_API_LLM_SEND_BANK_AS_USER")}
|
||||
os.environ["HINDSIGHT_API_LLM_SEND_BANK_AS_USER"] = "1"
|
||||
key = "HINDSIGHT_API_RERANKER_SEND_BANK_AS_HEADER"
|
||||
saved = {key: os.environ.get(key)}
|
||||
os.environ.pop(key, None)
|
||||
clear_config_cache()
|
||||
try:
|
||||
assert HindsightConfig.from_env().llm_send_bank_as_user is True
|
||||
assert HindsightConfig.from_env().reranker_send_bank_as_header is False
|
||||
token = _current_bank_id.set("bank-disabled")
|
||||
try:
|
||||
assert reranker_bank_attribution_headers() == {}
|
||||
finally:
|
||||
_current_bank_id.reset(token)
|
||||
os.environ[key] = "true"
|
||||
clear_config_cache()
|
||||
assert HindsightConfig.from_env().reranker_send_bank_as_header is True
|
||||
assert reranker_bank_attribution_headers() == {}
|
||||
token = _current_bank_id.set("bank-configured")
|
||||
try:
|
||||
assert reranker_bank_attribution_headers() == {"X-Hindsight-Bank-Id": "bank-configured"}
|
||||
finally:
|
||||
_current_bank_id.reset(token)
|
||||
finally:
|
||||
_restore_env(saved)
|
||||
|
||||
|
||||
@@ -5,6 +5,7 @@ import pytest_asyncio
|
||||
import httpx
|
||||
from datetime import datetime
|
||||
from hindsight_api.api import create_app
|
||||
from hindsight_api.api.http import BankTemplateManifest, validate_bank_template
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
@@ -81,6 +82,17 @@ class TestImportValidation:
|
||||
assert set(data["mental_models_created"]) == {"test-model-one", "test-model-two"}
|
||||
assert set(data["directives_created"]) == {"Be concise", "Use examples"}
|
||||
|
||||
def test_verbatim_extraction_mode_is_valid(self):
|
||||
"""verbatim is a valid retain extraction mode in bank manifests."""
|
||||
manifest = BankTemplateManifest.model_validate(
|
||||
{
|
||||
"version": "1",
|
||||
"bank": {"retain_extraction_mode": "verbatim"},
|
||||
}
|
||||
)
|
||||
|
||||
assert validate_bank_template(manifest) == []
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_import_invalid_version(self, api_client, bank_id):
|
||||
"""Reject manifest with unsupported version."""
|
||||
|
||||
@@ -0,0 +1,76 @@
|
||||
from contextlib import asynccontextmanager
|
||||
|
||||
import pytest
|
||||
|
||||
from hindsight_api.engine.retain import bank_utils
|
||||
|
||||
|
||||
class _FailingIndexOps:
|
||||
async def create_bank_vector_indexes(self, *args, **kwargs) -> None:
|
||||
raise RuntimeError("simulated per-bank vector index DDL failure")
|
||||
|
||||
|
||||
class _FakeTransaction:
|
||||
def __init__(self, conn: "_FakeConnection") -> None:
|
||||
self._conn = conn
|
||||
|
||||
async def __aenter__(self) -> None:
|
||||
self._conn.in_transaction = True
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb) -> None:
|
||||
if exc_type is None:
|
||||
self._conn.committed_bank = self._conn.pending_bank
|
||||
self._conn.pending_bank = None
|
||||
self._conn.in_transaction = False
|
||||
|
||||
|
||||
class _FakeConnection:
|
||||
def __init__(self) -> None:
|
||||
self.committed_bank: str | None = None
|
||||
self.pending_bank: str | None = None
|
||||
self.in_transaction = False
|
||||
|
||||
def transaction(self) -> _FakeTransaction:
|
||||
return _FakeTransaction(self)
|
||||
|
||||
async def fetchrow(self, query: str, bank_id: str):
|
||||
visible_bank = self.pending_bank if self.in_transaction else self.committed_bank
|
||||
if visible_bank != bank_id:
|
||||
return None
|
||||
return {
|
||||
"name": bank_id,
|
||||
"disposition": bank_utils.DEFAULT_DISPOSITION,
|
||||
"mission": "",
|
||||
}
|
||||
|
||||
async def fetchval(self, query: str, bank_id: str, *args):
|
||||
if self.in_transaction:
|
||||
self.pending_bank = bank_id
|
||||
else:
|
||||
self.committed_bank = bank_id
|
||||
return bank_id
|
||||
|
||||
|
||||
class _FakePool:
|
||||
def __init__(self, conn: _FakeConnection) -> None:
|
||||
self.conn = conn
|
||||
self.ops = _FailingIndexOps()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_lazy_bank_create_rolls_back_on_vector_index_failure(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""A failed per-bank index DDL must not leave an orphaned bank row."""
|
||||
conn = _FakeConnection()
|
||||
pool = _FakePool(conn)
|
||||
|
||||
@asynccontextmanager
|
||||
async def acquire_without_transaction(*args, **kwargs):
|
||||
yield conn
|
||||
|
||||
monkeypatch.setattr(bank_utils, "acquire_with_retry", acquire_without_transaction)
|
||||
|
||||
with pytest.raises(RuntimeError, match="simulated per-bank vector index DDL failure"):
|
||||
await bank_utils.get_or_create_bank_profile(pool, "atomicity-test-bank")
|
||||
|
||||
profile = await bank_utils.get_bank_profile_if_exists(pool, "atomicity-test-bank")
|
||||
assert profile is None, "bank row should roll back when per-bank vector index creation fails"
|
||||
@@ -8,18 +8,15 @@ Tests cover:
|
||||
- Worker recovery on restart
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import uuid
|
||||
from datetime import datetime, timezone
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from hindsight_api import RequestContext
|
||||
from hindsight_api.config import HindsightConfig
|
||||
from hindsight_api.engine.llm_wrapper import create_llm_provider
|
||||
from hindsight_api.engine.retain.fact_extraction import (
|
||||
RetainContent,
|
||||
extract_facts_from_contents,
|
||||
@@ -202,6 +199,224 @@ async def test_batch_api_normal_flow(mock_llm_config, test_contents, hindsight_c
|
||||
pass
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_batch_api_accepts_top_level_fact_list(mock_llm_config, test_contents, hindsight_config):
|
||||
"""Batch extraction accepts a recoverable top-level facts array."""
|
||||
batch_id = "batch_top_level_facts"
|
||||
mock_llm_config._provider_impl.supports_batch_api = AsyncMock(return_value=True)
|
||||
mock_llm_config._provider_impl.submit_batch = AsyncMock(
|
||||
return_value={
|
||||
"batch_id": batch_id,
|
||||
"status": "validating",
|
||||
"request_counts": {"total": 1, "completed": 0, "failed": 0},
|
||||
}
|
||||
)
|
||||
mock_llm_config._provider_impl.get_batch_status = AsyncMock(
|
||||
return_value={"status": "completed", "request_counts": {"total": 1, "completed": 1, "failed": 0}}
|
||||
)
|
||||
mock_llm_config._provider_impl.retrieve_batch_results = AsyncMock(
|
||||
return_value=[
|
||||
{
|
||||
"custom_id": "chunk_0",
|
||||
"response": {
|
||||
"body": {
|
||||
"choices": [
|
||||
{
|
||||
"message": {
|
||||
"content": json.dumps(
|
||||
[
|
||||
{
|
||||
"what": "Alice is a senior software engineer at TechCorp",
|
||||
"when": "present",
|
||||
"where": "TechCorp",
|
||||
"who": "Alice",
|
||||
"why": "Professional background information",
|
||||
"fact_type": "world",
|
||||
"fact_kind": "conversation",
|
||||
}
|
||||
]
|
||||
)
|
||||
}
|
||||
}
|
||||
],
|
||||
"usage": {"prompt_tokens": 100, "completion_tokens": 50, "total_tokens": 150},
|
||||
}
|
||||
},
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
facts, chunks, usage = await extract_facts_from_contents_batch_api(
|
||||
contents=[test_contents[0]],
|
||||
llm_config=mock_llm_config,
|
||||
agent_name="test_agent",
|
||||
config=hindsight_config,
|
||||
pool=None,
|
||||
operation_id=None,
|
||||
schema=None,
|
||||
)
|
||||
|
||||
assert len(facts) == 1
|
||||
assert "Alice" in facts[0].fact_text
|
||||
assert len(chunks) == 1
|
||||
assert chunks[0].fact_count == 1
|
||||
assert usage.total_tokens == 150
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_batch_api_rejects_top_level_non_fact_list(mock_llm_config, test_contents, hindsight_config):
|
||||
"""Batch extraction records malformed top-level lists instead of crashing."""
|
||||
batch_id = "batch_malformed_list"
|
||||
mock_llm_config._provider_impl.supports_batch_api = AsyncMock(return_value=True)
|
||||
mock_llm_config._provider_impl.submit_batch = AsyncMock(return_value={"batch_id": batch_id})
|
||||
mock_llm_config._provider_impl.get_batch_status = AsyncMock(
|
||||
return_value={"status": "completed", "request_counts": {"total": 1, "completed": 1, "failed": 0}}
|
||||
)
|
||||
mock_llm_config._provider_impl.retrieve_batch_results = AsyncMock(
|
||||
return_value=[
|
||||
{
|
||||
"custom_id": "chunk_0",
|
||||
"response": {
|
||||
"body": {
|
||||
"choices": [{"message": {"content": json.dumps(["not a fact dict"])}}],
|
||||
"usage": {"prompt_tokens": 100, "completion_tokens": 50, "total_tokens": 150},
|
||||
}
|
||||
},
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
facts, chunks, usage = await extract_facts_from_contents_batch_api(
|
||||
contents=[test_contents[0]],
|
||||
llm_config=mock_llm_config,
|
||||
agent_name="test_agent",
|
||||
config=hindsight_config,
|
||||
pool=None,
|
||||
operation_id=None,
|
||||
schema=None,
|
||||
)
|
||||
|
||||
assert facts == []
|
||||
assert len(chunks) == 1
|
||||
assert chunks[0].fact_count == 0
|
||||
assert usage.total_tokens == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_batch_api_recovers_fenced_and_control_char_json(mock_llm_config, test_contents, hindsight_config):
|
||||
"""#2701: batch content that bare json.loads can't parse but parse_llm_json can
|
||||
(markdown code fences + an embedded raw control character, e.g. a transient
|
||||
Gemini quirk) must still yield facts instead of dropping the whole chunk."""
|
||||
batch_id = "batch_recoverable_json"
|
||||
|
||||
# Valid facts JSON, but wrapped in ```json fences AND containing a raw
|
||||
# control character (\x01) inside a string value. Bare json.loads fails on
|
||||
# both; parse_llm_json strips the fences and scrubs the control char.
|
||||
inner_json = json.dumps(
|
||||
{
|
||||
"facts": [
|
||||
{
|
||||
"what": "Alice is a senior software engineer at TechCorp",
|
||||
"when": "present",
|
||||
"where": "TechCorp",
|
||||
"who": "Alice",
|
||||
"why": "Professional background\x01information",
|
||||
"fact_type": "world",
|
||||
"fact_kind": "conversation",
|
||||
}
|
||||
]
|
||||
}
|
||||
)
|
||||
unparseable_content = f"```json\n{inner_json}\n```"
|
||||
|
||||
# Sanity: the raw content is NOT parseable by the bare parser.
|
||||
with pytest.raises(json.JSONDecodeError):
|
||||
json.loads(unparseable_content)
|
||||
|
||||
mock_llm_config._provider_impl.supports_batch_api = AsyncMock(return_value=True)
|
||||
mock_llm_config._provider_impl.submit_batch = AsyncMock(return_value={"batch_id": batch_id})
|
||||
mock_llm_config._provider_impl.get_batch_status = AsyncMock(
|
||||
return_value={"status": "completed", "request_counts": {"total": 1, "completed": 1, "failed": 0}}
|
||||
)
|
||||
mock_llm_config._provider_impl.retrieve_batch_results = AsyncMock(
|
||||
return_value=[
|
||||
{
|
||||
"custom_id": "chunk_0",
|
||||
"response": {
|
||||
"body": {
|
||||
"choices": [{"message": {"content": unparseable_content}}],
|
||||
"usage": {"prompt_tokens": 100, "completion_tokens": 50, "total_tokens": 150},
|
||||
}
|
||||
},
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
facts, chunks, usage = await extract_facts_from_contents_batch_api(
|
||||
contents=[test_contents[0]],
|
||||
llm_config=mock_llm_config,
|
||||
agent_name="test_agent",
|
||||
config=hindsight_config,
|
||||
pool=None,
|
||||
operation_id=None,
|
||||
schema=None,
|
||||
)
|
||||
|
||||
# The facts are recovered rather than lost.
|
||||
assert len(facts) == 1
|
||||
assert "Alice" in facts[0].fact_text
|
||||
assert len(chunks) == 1
|
||||
assert chunks[0].fact_count == 1
|
||||
assert usage.total_tokens == 150
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_batch_api_unparseable_json_still_records_error(mock_llm_config, test_contents, hindsight_config):
|
||||
"""#2701: genuinely unparseable content (not recoverable by parse_llm_json)
|
||||
must preserve the existing behavior — record the error, fact_count=0, no crash."""
|
||||
batch_id = "batch_unparseable_json"
|
||||
|
||||
# Not JSON at all, and not recoverable by fence-stripping or control-char scrubbing.
|
||||
unparseable_content = "this is not json {{{ ["
|
||||
|
||||
with pytest.raises(json.JSONDecodeError):
|
||||
json.loads(unparseable_content)
|
||||
|
||||
mock_llm_config._provider_impl.supports_batch_api = AsyncMock(return_value=True)
|
||||
mock_llm_config._provider_impl.submit_batch = AsyncMock(return_value={"batch_id": batch_id})
|
||||
mock_llm_config._provider_impl.get_batch_status = AsyncMock(
|
||||
return_value={"status": "completed", "request_counts": {"total": 1, "completed": 1, "failed": 0}}
|
||||
)
|
||||
mock_llm_config._provider_impl.retrieve_batch_results = AsyncMock(
|
||||
return_value=[
|
||||
{
|
||||
"custom_id": "chunk_0",
|
||||
"response": {
|
||||
"body": {
|
||||
"choices": [{"message": {"content": unparseable_content}}],
|
||||
"usage": {"prompt_tokens": 100, "completion_tokens": 50, "total_tokens": 150},
|
||||
}
|
||||
},
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
facts, chunks, usage = await extract_facts_from_contents_batch_api(
|
||||
contents=[test_contents[0]],
|
||||
llm_config=mock_llm_config,
|
||||
agent_name="test_agent",
|
||||
config=hindsight_config,
|
||||
pool=None,
|
||||
operation_id=None,
|
||||
schema=None,
|
||||
)
|
||||
|
||||
assert facts == []
|
||||
assert len(chunks) == 1
|
||||
assert chunks[0].fact_count == 0
|
||||
assert usage.total_tokens == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_batch_api_crash_recovery(mock_llm_config, test_contents, hindsight_config, memory, request_context):
|
||||
"""Test crash recovery: resume polling from existing batch_id."""
|
||||
|
||||
@@ -0,0 +1,34 @@
|
||||
"""Regression tests for the causal link taxonomy used by retain."""
|
||||
|
||||
import pytest
|
||||
from pydantic import ValidationError
|
||||
|
||||
from hindsight_api.engine.causal_links import (
|
||||
CANONICAL_CAUSAL_LINK_TYPE,
|
||||
CANONICAL_CAUSAL_LINK_TYPES,
|
||||
CAUSAL_LINK_TYPES,
|
||||
LEGACY_CAUSAL_LINK_TYPES,
|
||||
)
|
||||
from hindsight_api.engine.retain.fact_extraction import CausalRelation, FactCausalRelation
|
||||
|
||||
|
||||
@pytest.mark.parametrize("relation_type", ["causes", "enables", "prevents"])
|
||||
def test_retain_causal_models_reject_legacy_relation_types(relation_type: str) -> None:
|
||||
"""New retain output is canonical even though storage reads legacy links."""
|
||||
with pytest.raises(ValidationError):
|
||||
CausalRelation(target_fact_index=0, relation_type=relation_type)
|
||||
|
||||
with pytest.raises(ValidationError):
|
||||
FactCausalRelation(target_index=0, relation_type=relation_type)
|
||||
|
||||
|
||||
def test_retain_causal_models_accept_caused_by() -> None:
|
||||
"""The canonical causal relationship remains valid in both extraction schemas."""
|
||||
assert CausalRelation(target_fact_index=0, relation_type="caused_by").relation_type == "caused_by"
|
||||
assert FactCausalRelation(target_index=0, relation_type="caused_by").relation_type == "caused_by"
|
||||
|
||||
|
||||
def test_causal_link_taxonomy_keeps_canonical_and_legacy_types_separate() -> None:
|
||||
assert CANONICAL_CAUSAL_LINK_TYPES == {CANONICAL_CAUSAL_LINK_TYPE}
|
||||
assert LEGACY_CAUSAL_LINK_TYPES == {"causes", "enables", "prevents"}
|
||||
assert CAUSAL_LINK_TYPES == (CANONICAL_CAUSAL_LINK_TYPE, "causes", "enables", "prevents")
|
||||
@@ -63,9 +63,7 @@ class TestCausalRelationsValidation:
|
||||
assert rel.target_fact_index >= 0, (
|
||||
f"Fact {i} has negative causal relation index: {rel.target_fact_index}"
|
||||
)
|
||||
assert rel.relation_type in ["caused_by", "enabled_by", "prevented_by"], (
|
||||
f"Invalid relation_type: {rel.relation_type}"
|
||||
)
|
||||
assert rel.relation_type == "caused_by", f"Invalid relation_type: {rel.relation_type}"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_first_fact_has_no_causal_relations(self):
|
||||
@@ -196,10 +194,10 @@ class TestCausalRelationsValidation:
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_relation_types_are_backward_looking(self):
|
||||
async def test_relation_types_use_the_canonical_form(self):
|
||||
"""
|
||||
Test that all relation types describe how the current fact
|
||||
relates to a previous fact (caused_by, enabled_by, prevented_by).
|
||||
Test that all extracted relation types use the canonical backward-looking
|
||||
``caused_by`` form.
|
||||
"""
|
||||
text = """
|
||||
Alice learned Python programming.
|
||||
@@ -220,12 +218,9 @@ class TestCausalRelationsValidation:
|
||||
config=_get_raw_config(),
|
||||
)
|
||||
|
||||
# Verify relation types are all backward-looking
|
||||
valid_types = {"caused_by", "enabled_by", "prevented_by"}
|
||||
|
||||
for i, fact in enumerate(facts):
|
||||
if fact.causal_relations:
|
||||
for rel in fact.causal_relations:
|
||||
assert rel.relation_type in valid_types, (
|
||||
f"Invalid relation_type '{rel.relation_type}'. Must be one of: {valid_types}"
|
||||
assert rel.relation_type == "caused_by", (
|
||||
f"Invalid relation_type '{rel.relation_type}'. Must be 'caused_by'"
|
||||
)
|
||||
|
||||
@@ -85,11 +85,10 @@ After searching for weeks, I finally found a cheaper apartment in Brooklyn.
|
||||
f"Got {len(all_causal_relations)}: {all_causal_relations}"
|
||||
)
|
||||
|
||||
# Verify relation types are valid (passive only - facts reference PREVIOUS facts)
|
||||
valid_types = {"caused_by", "enabled_by", "prevented_by"}
|
||||
# Retain writes the single canonical passive form for previous facts.
|
||||
for rel in all_causal_relations:
|
||||
assert rel["relation_type"] in valid_types, (
|
||||
f"Invalid relation_type '{rel['relation_type']}'. Must be one of {valid_types}"
|
||||
assert rel["relation_type"] == "caused_by", (
|
||||
f"Invalid relation_type '{rel['relation_type']}'. Must be 'caused_by'"
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -165,10 +164,9 @@ Machine learning fascinated me so much that I changed my career to data science.
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bidirectional_causal_relationships(self):
|
||||
async def test_causal_relationships_use_backward_references(self):
|
||||
"""
|
||||
Test that bidirectional causal relationships (causes and caused_by)
|
||||
are handled correctly.
|
||||
Test that causal relationships are represented as backward references.
|
||||
"""
|
||||
text = """
|
||||
My promotion at work caused me to move to New York.
|
||||
|
||||
@@ -0,0 +1,50 @@
|
||||
"""Regression coverage for deterministic chunk deletion ordering."""
|
||||
|
||||
import pytest
|
||||
|
||||
from hindsight_api.engine.retain import chunk_storage
|
||||
|
||||
|
||||
class RecordingConn:
|
||||
def __init__(self) -> None:
|
||||
self.calls: list[tuple[str, tuple[object, ...]]] = []
|
||||
|
||||
async def execute(self, sql: str, *args: object) -> None:
|
||||
self.calls.append((sql, args))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_chunks_by_ids_predeletes_links_before_chunks():
|
||||
conn = RecordingConn()
|
||||
chunk_ids = ["chunk-b", "chunk-a"]
|
||||
|
||||
await chunk_storage.delete_chunks_by_ids(conn, chunk_ids)
|
||||
|
||||
assert len(conn.calls) == 2
|
||||
link_sql, link_args = conn.calls[0]
|
||||
chunk_sql, chunk_args = conn.calls[1]
|
||||
|
||||
assert link_args == (chunk_ids,)
|
||||
assert chunk_args == (chunk_ids,)
|
||||
|
||||
assert "DELETE FROM" in link_sql
|
||||
assert "memory_links" in link_sql
|
||||
assert "target_units AS MATERIALIZED" in link_sql
|
||||
assert "ordered_links AS MATERIALIZED" in link_sql
|
||||
assert "ORDER BY" in link_sql
|
||||
assert "FOR UPDATE OF ml" in link_sql
|
||||
|
||||
assert "DELETE FROM" in chunk_sql
|
||||
assert "chunks" in chunk_sql
|
||||
assert "ordered_chunks AS MATERIALIZED" in chunk_sql
|
||||
assert "ORDER BY chunk_id" in chunk_sql
|
||||
assert "FOR UPDATE" in chunk_sql
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_chunks_by_ids_noops_without_chunks():
|
||||
conn = RecordingConn()
|
||||
|
||||
await chunk_storage.delete_chunks_by_ids(conn, [])
|
||||
|
||||
assert conn.calls == []
|
||||
@@ -0,0 +1,199 @@
|
||||
"""Regression test for surfacing the CLI's real error text (issue #2702).
|
||||
|
||||
The Claude Code CLI can report a failure with ``is_error=True`` while
|
||||
``subtype`` still reads ``"success"``, putting the actual detail in
|
||||
``result`` — e.g. quota exhaustion:
|
||||
|
||||
{"type":"result","subtype":"success","is_error":true,
|
||||
"api_error_status":429,
|
||||
"result":"You've hit your weekly limit · resets Jul 18, 12pm (UTC)"}
|
||||
|
||||
The Agent SDK's fallback exception is built from ``errors`` (empty here)
|
||||
or ``subtype``, producing the misleading "Claude Code returned an error
|
||||
result: success". These tests assert that both provider call paths inspect
|
||||
the ResultMessage directly and raise with the CLI's actual error text.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
QUOTA_ERROR_TEXT = "You've hit your weekly limit · resets Jul 18, 12pm (UTC)"
|
||||
|
||||
|
||||
@dataclass
|
||||
class _FakeOptions:
|
||||
"""Stand-in for ClaudeAgentOptions; captures kwargs without importing SDK."""
|
||||
|
||||
system_prompt: str | None = None
|
||||
max_turns: int | None = None
|
||||
allowed_tools: list[str] = field(default_factory=list)
|
||||
tools: list[str] = field(default_factory=list)
|
||||
env: dict[str, str] = field(default_factory=dict)
|
||||
mcp_servers: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
|
||||
class _FakeAssistantMessage:
|
||||
def __init__(self, content: list[Any]) -> None:
|
||||
self.content = content
|
||||
|
||||
|
||||
class _FakeTextBlock:
|
||||
def __init__(self, text: str) -> None:
|
||||
self.text = text
|
||||
|
||||
|
||||
class _FakeResultMessage:
|
||||
def __init__(self, subtype: str, is_error: bool, result: str | None) -> None:
|
||||
self.subtype = subtype
|
||||
self.is_error = is_error
|
||||
self.result = result
|
||||
|
||||
|
||||
def _instantiate_provider():
|
||||
from hindsight_api.engine.providers.claude_code_llm import ClaudeCodeLLM
|
||||
|
||||
return ClaudeCodeLLM(
|
||||
provider="claude-code",
|
||||
api_key="",
|
||||
base_url="",
|
||||
model="claude-haiku-4-5",
|
||||
reasoning_effort="low",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_raises_with_result_text_on_error_result(monkeypatch):
|
||||
"""call() must surface ResultMessage.result, not the 'success' subtype."""
|
||||
import claude_agent_sdk
|
||||
|
||||
async def fake_query(prompt: str, options: _FakeOptions):
|
||||
yield _FakeResultMessage(subtype="success", is_error=True, result=QUOTA_ERROR_TEXT)
|
||||
|
||||
monkeypatch.setattr(claude_agent_sdk, "ClaudeAgentOptions", _FakeOptions)
|
||||
monkeypatch.setattr(claude_agent_sdk, "AssistantMessage", _FakeAssistantMessage)
|
||||
monkeypatch.setattr(claude_agent_sdk, "TextBlock", _FakeTextBlock)
|
||||
monkeypatch.setattr(claude_agent_sdk, "ResultMessage", _FakeResultMessage)
|
||||
monkeypatch.setattr(claude_agent_sdk, "query", fake_query)
|
||||
|
||||
provider = _instantiate_provider()
|
||||
with pytest.raises(RuntimeError) as excinfo:
|
||||
await provider.call(
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
max_retries=0,
|
||||
scope="test",
|
||||
)
|
||||
|
||||
assert QUOTA_ERROR_TEXT in str(excinfo.value)
|
||||
assert "error result: success" not in str(excinfo.value)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_falls_back_to_subtype_when_result_empty(monkeypatch):
|
||||
"""With no result text, the subtype is still better than nothing."""
|
||||
import claude_agent_sdk
|
||||
|
||||
async def fake_query(prompt: str, options: _FakeOptions):
|
||||
yield _FakeResultMessage(subtype="error_max_turns", is_error=True, result=None)
|
||||
|
||||
monkeypatch.setattr(claude_agent_sdk, "ClaudeAgentOptions", _FakeOptions)
|
||||
monkeypatch.setattr(claude_agent_sdk, "AssistantMessage", _FakeAssistantMessage)
|
||||
monkeypatch.setattr(claude_agent_sdk, "TextBlock", _FakeTextBlock)
|
||||
monkeypatch.setattr(claude_agent_sdk, "ResultMessage", _FakeResultMessage)
|
||||
monkeypatch.setattr(claude_agent_sdk, "query", fake_query)
|
||||
|
||||
provider = _instantiate_provider()
|
||||
with pytest.raises(RuntimeError, match="error_max_turns"):
|
||||
await provider.call(
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
max_retries=0,
|
||||
scope="test",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_ignores_non_error_result_message(monkeypatch):
|
||||
"""A normal is_error=False ResultMessage must not affect the response."""
|
||||
import claude_agent_sdk
|
||||
|
||||
async def fake_query(prompt: str, options: _FakeOptions):
|
||||
yield _FakeAssistantMessage(content=[_FakeTextBlock(text="ok")])
|
||||
yield _FakeResultMessage(subtype="success", is_error=False, result="ok")
|
||||
|
||||
monkeypatch.setattr(claude_agent_sdk, "ClaudeAgentOptions", _FakeOptions)
|
||||
monkeypatch.setattr(claude_agent_sdk, "AssistantMessage", _FakeAssistantMessage)
|
||||
monkeypatch.setattr(claude_agent_sdk, "TextBlock", _FakeTextBlock)
|
||||
monkeypatch.setattr(claude_agent_sdk, "ResultMessage", _FakeResultMessage)
|
||||
monkeypatch.setattr(claude_agent_sdk, "query", fake_query)
|
||||
|
||||
provider = _instantiate_provider()
|
||||
result = await provider.call(
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
max_retries=0,
|
||||
scope="test",
|
||||
)
|
||||
|
||||
assert result == "ok"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_with_tools_raises_with_result_text_on_error_result(monkeypatch):
|
||||
"""call_with_tools() must surface ResultMessage.result the same way."""
|
||||
import claude_agent_sdk
|
||||
|
||||
class _FakeClient:
|
||||
def __init__(self, options: _FakeOptions) -> None:
|
||||
self.options = options
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb):
|
||||
return False
|
||||
|
||||
async def query(self, prompt: str) -> None:
|
||||
return None
|
||||
|
||||
async def receive_response(self):
|
||||
yield _FakeResultMessage(subtype="success", is_error=True, result=QUOTA_ERROR_TEXT)
|
||||
|
||||
@dataclass
|
||||
class _FakeSdkMcpTool:
|
||||
name: str
|
||||
description: str
|
||||
input_schema: dict[str, Any]
|
||||
handler: Any
|
||||
|
||||
def fake_create_sdk_mcp_server(name: str, version: str, tools=None):
|
||||
return {"name": name, "version": version, "tools": tools}
|
||||
|
||||
monkeypatch.setattr(claude_agent_sdk, "ClaudeAgentOptions", _FakeOptions)
|
||||
monkeypatch.setattr(claude_agent_sdk, "AssistantMessage", _FakeAssistantMessage)
|
||||
monkeypatch.setattr(claude_agent_sdk, "TextBlock", _FakeTextBlock)
|
||||
monkeypatch.setattr(claude_agent_sdk, "ResultMessage", _FakeResultMessage)
|
||||
monkeypatch.setattr(claude_agent_sdk, "ToolUseBlock", type("ToolUseBlock", (), {}))
|
||||
monkeypatch.setattr(claude_agent_sdk, "ClaudeSDKClient", _FakeClient)
|
||||
monkeypatch.setattr(claude_agent_sdk, "SdkMcpTool", _FakeSdkMcpTool)
|
||||
monkeypatch.setattr(claude_agent_sdk, "create_sdk_mcp_server", fake_create_sdk_mcp_server)
|
||||
|
||||
provider = _instantiate_provider()
|
||||
with pytest.raises(RuntimeError) as excinfo:
|
||||
await provider.call_with_tools(
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
tools=[
|
||||
{
|
||||
"function": {
|
||||
"name": "noop",
|
||||
"description": "no-op",
|
||||
"parameters": {"type": "object", "properties": {}},
|
||||
}
|
||||
}
|
||||
],
|
||||
max_retries=0,
|
||||
scope="test",
|
||||
)
|
||||
|
||||
assert QUOTA_ERROR_TEXT in str(excinfo.value)
|
||||
@@ -24,10 +24,11 @@ from __future__ import annotations
|
||||
import asyncio
|
||||
import base64
|
||||
import json
|
||||
import os
|
||||
import stat
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
@@ -38,6 +39,7 @@ import pytest
|
||||
from hindsight_api.engine.providers.codex_llm import (
|
||||
_CODEX_CLIENT_ID,
|
||||
_CODEX_REFRESH_TOKEN_URL,
|
||||
CodexAuthManager,
|
||||
CodexLLM,
|
||||
CodexRefreshExpiredError,
|
||||
)
|
||||
@@ -396,6 +398,63 @@ async def test_concurrent_ensure_fresh_token_calls_produce_one_refresh(tmp_path:
|
||||
assert call_count == 1, f"expected 1 network refresh under contention, got {call_count}"
|
||||
|
||||
|
||||
def test_sibling_auth_manager_adopts_rotated_codex_credentials(tmp_path: Path):
|
||||
"""A stale sibling manager should adopt auth.json rotation before reusing the old RT."""
|
||||
expired = _make_jwt(int(time.time()) - 60)
|
||||
new_access = _make_jwt(int(time.time()) + 3600)
|
||||
auth_file = _make_codex_auth_file(tmp_path, expired, refresh_token="rt-old")
|
||||
|
||||
first = CodexAuthManager.from_file(auth_file)
|
||||
sibling = CodexAuthManager.from_file(auth_file)
|
||||
|
||||
refresh_resp = _refresh_response(200, {"access_token": new_access, "refresh_token": "rt-new"})
|
||||
with patch.object(first._http_client, "post", return_value=refresh_resp):
|
||||
first.refresh_tokens(reason="test")
|
||||
|
||||
def unexpected_post(*args, **kwargs):
|
||||
raise AssertionError("stale sibling should not call refresh endpoint with old refresh_token")
|
||||
|
||||
with patch.object(sibling._http_client, "post", new=unexpected_post):
|
||||
sibling.refresh_tokens(reason="test", force=True)
|
||||
|
||||
assert sibling.access_token == new_access
|
||||
assert sibling.refresh_token == "rt-new"
|
||||
|
||||
|
||||
def test_parallel_auth_managers_share_one_refresh_for_same_auth_file(tmp_path: Path):
|
||||
"""Separate managers in one process should single-flight per canonical auth path."""
|
||||
expired = _make_jwt(int(time.time()) - 60)
|
||||
new_access = _make_jwt(int(time.time()) + 3600)
|
||||
auth_file = _make_codex_auth_file(tmp_path, expired, refresh_token="rt-old")
|
||||
|
||||
managers = [CodexAuthManager.from_file(auth_file), CodexAuthManager.from_file(auth_file)]
|
||||
call_count = 0
|
||||
call_count_lock = threading.Lock()
|
||||
|
||||
def fake_post(*args, **kwargs):
|
||||
nonlocal call_count
|
||||
with call_count_lock:
|
||||
call_count += 1
|
||||
time.sleep(0.02)
|
||||
return _refresh_response(200, {"access_token": new_access, "refresh_token": "rt-new"})
|
||||
|
||||
patches = [patch.object(manager._http_client, "post", new=fake_post) for manager in managers]
|
||||
for patcher in patches:
|
||||
patcher.start()
|
||||
try:
|
||||
with ThreadPoolExecutor(max_workers=2) as executor:
|
||||
futures = [executor.submit(manager.refresh_tokens, "test") for manager in managers]
|
||||
for future in futures:
|
||||
future.result()
|
||||
finally:
|
||||
for patcher in patches:
|
||||
patcher.stop()
|
||||
|
||||
assert call_count == 1
|
||||
assert [manager.access_token for manager in managers] == [new_access, new_access]
|
||||
assert [manager.refresh_token for manager in managers] == ["rt-new", "rt-new"]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Reactive 401 retry on the request path
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -422,6 +481,7 @@ async def test_call_reactively_refreshes_on_401_and_retries(tmp_path: Path):
|
||||
refresh_resp = _refresh_response(200, {"access_token": new_access, "refresh_token": "rt-new"})
|
||||
|
||||
call_count = {"refresh": 0, "post": 0}
|
||||
sent_headers: list[httpx.Headers] = []
|
||||
|
||||
# Sync mock for the auth manager's HTTP client (used for token refresh).
|
||||
def fake_refresh_post(*args, **kwargs):
|
||||
@@ -431,6 +491,7 @@ async def test_call_reactively_refreshes_on_401_and_retries(tmp_path: Path):
|
||||
# Async mock for the LLM's HTTP client (used for backend calls).
|
||||
async def fake_backend_post(url, **kwargs):
|
||||
call_count["post"] += 1
|
||||
sent_headers.append(httpx.Headers(kwargs["headers"]))
|
||||
if call_count["post"] == 1:
|
||||
raise httpx.HTTPStatusError("401", request=MagicMock(), response=fail_response)
|
||||
return success_resp
|
||||
@@ -451,6 +512,10 @@ async def test_call_reactively_refreshes_on_401_and_retries(tmp_path: Path):
|
||||
assert call_count["refresh"] == 1
|
||||
assert call_count["post"] == 2 # one 401, one success after refresh
|
||||
assert llm.access_token == new_access
|
||||
assert sent_headers[0]["Authorization"] == f"Bearer {fresh}"
|
||||
assert sent_headers[1]["Authorization"] == f"Bearer {new_access}"
|
||||
for header_name in ("Content-Type", "OpenAI-Account-ID", "User-Agent", "Origin", "originator"):
|
||||
assert sent_headers[1][header_name] == sent_headers[0][header_name]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
||||
@@ -0,0 +1,63 @@
|
||||
"""Regression tests for Codex request identity headers."""
|
||||
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
from hindsight_api.engine.providers.codex_llm import CodexLLM
|
||||
|
||||
|
||||
def build_llm() -> CodexLLM:
|
||||
with (
|
||||
patch.object(CodexLLM, "_load_codex_auth", return_value=("token", "account")),
|
||||
patch.object(CodexLLM, "_load_codex_refresh_token", return_value=None),
|
||||
):
|
||||
return CodexLLM(
|
||||
provider="openai-codex",
|
||||
api_key="ignored",
|
||||
base_url="https://chatgpt.com/backend-api",
|
||||
model="gpt-5.6-luna",
|
||||
)
|
||||
|
||||
|
||||
def assert_codex_request_identity(headers: httpx.Headers) -> None:
|
||||
assert headers["originator"] == "codex_cli_rs"
|
||||
assert headers["User-Agent"] == "codex_cli_rs/0.0.0 (Hindsight)"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_sends_codex_request_identity() -> None:
|
||||
llm = build_llm()
|
||||
response = MagicMock()
|
||||
response.raise_for_status.return_value = None
|
||||
|
||||
with (
|
||||
patch.object(llm._client, "post", new_callable=AsyncMock) as mock_post,
|
||||
patch.object(llm, "_parse_sse_stream", new_callable=AsyncMock, return_value="ok"),
|
||||
):
|
||||
mock_post.return_value = response
|
||||
await llm.call(messages=[{"role": "user", "content": "hello"}], max_retries=0)
|
||||
|
||||
assert_codex_request_identity(mock_post.call_args.kwargs["headers"])
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_with_tools_sends_codex_request_identity() -> None:
|
||||
llm = build_llm()
|
||||
response = MagicMock()
|
||||
response.status_code = 200
|
||||
response.raise_for_status.return_value = None
|
||||
|
||||
with (
|
||||
patch.object(llm._client, "post", new_callable=AsyncMock) as mock_post,
|
||||
patch.object(llm, "_parse_sse_tool_stream", new_callable=AsyncMock, return_value=(None, [])),
|
||||
):
|
||||
mock_post.return_value = response
|
||||
await llm.call_with_tools(
|
||||
messages=[{"role": "user", "content": "hello"}],
|
||||
tools=[],
|
||||
max_retries=0,
|
||||
)
|
||||
|
||||
assert_codex_request_identity(mock_post.call_args.kwargs["headers"])
|
||||
@@ -0,0 +1,188 @@
|
||||
"""
|
||||
Regression tests for Codex structured output (issue #2504).
|
||||
|
||||
Before the fix, ``CodexLLM.call(strict_schema=True)`` was a dead no-op: structured
|
||||
output always went through prompt-injected schema + raw ``json.loads`` on the
|
||||
model's free-form text. Escape-heavy content (code, serial/CLI commands, Windows
|
||||
paths, regexes) makes weaker models emit invalid ``\\escape`` sequences, so every
|
||||
parse attempt fails and retain/consolidation burn all retries and fail.
|
||||
|
||||
The fix:
|
||||
- ``strict_schema=True`` routes structured output through a single forced function
|
||||
tool (constrained decoding into the response schema).
|
||||
- The non-strict fallback now repairs invalid ``\\escape`` sequences before giving up.
|
||||
"""
|
||||
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from pydantic import BaseModel
|
||||
|
||||
from hindsight_api.engine.providers.codex_llm import (
|
||||
CodexLLM,
|
||||
_repair_invalid_json_escapes,
|
||||
)
|
||||
from hindsight_api.engine.response_models import LLMToolCall
|
||||
|
||||
|
||||
class _Fact(BaseModel):
|
||||
fact: str
|
||||
|
||||
|
||||
def build_llm() -> CodexLLM:
|
||||
with patch.object(CodexLLM, "_load_codex_auth", return_value=("token", "account")):
|
||||
return CodexLLM(
|
||||
provider="openai-codex",
|
||||
api_key="ignored",
|
||||
base_url="https://chatgpt.com/backend-api",
|
||||
model="gpt-5.4-mini",
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _repair_invalid_json_escapes — pure unit tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_repair_fixes_invalid_escape_in_json():
|
||||
# `\d` and `\s` are not valid JSON escapes; raw json.loads fails.
|
||||
broken = r'{"fact": "regex \d+\s matches digits"}'
|
||||
import json
|
||||
|
||||
with pytest.raises(json.JSONDecodeError):
|
||||
json.loads(broken)
|
||||
repaired = _repair_invalid_json_escapes(broken)
|
||||
assert json.loads(repaired) == {"fact": r"regex \d+\s matches digits"}
|
||||
|
||||
|
||||
def test_repair_preserves_valid_escapes():
|
||||
import json
|
||||
|
||||
valid = r'{"fact": "line1\nline2\ttab \"quoted\" \\backslash é"}'
|
||||
# Already valid — repair must not corrupt it.
|
||||
assert json.loads(_repair_invalid_json_escapes(valid)) == json.loads(valid)
|
||||
|
||||
|
||||
def test_repair_handles_windows_paths():
|
||||
import json
|
||||
|
||||
# Uses path segments whose first char isn't a valid JSON escape letter
|
||||
# (b/f/n/r/t/u), where the repair is unambiguous.
|
||||
broken = r'{"path": "C:\Windows\System32\app.exe"}'
|
||||
assert json.loads(_repair_invalid_json_escapes(broken)) == {"path": r"C:\Windows\System32\app.exe"}
|
||||
|
||||
|
||||
def test_repair_handles_trailing_backslash():
|
||||
# A lone trailing backslash must be escaped, not dropped.
|
||||
assert _repair_invalid_json_escapes("abc\\") == "abc\\\\"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# strict_schema forced-tool path
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_strict_schema_uses_forced_function_tool():
|
||||
llm = build_llm()
|
||||
response = MagicMock()
|
||||
response.raise_for_status.return_value = None
|
||||
tool_call = LLMToolCall(id="call-1", name="structured_response", arguments={"fact": "the sky is blue"})
|
||||
|
||||
with patch.object(llm._client, "post", new_callable=AsyncMock) as mock_post:
|
||||
mock_post.return_value = response
|
||||
with patch.object(llm, "_parse_sse_tool_stream", new_callable=AsyncMock) as mock_parse:
|
||||
mock_parse.return_value = (None, [tool_call])
|
||||
result = await llm.call(
|
||||
messages=[{"role": "user", "content": "The sky is blue"}],
|
||||
response_format=_Fact,
|
||||
strict_schema=True,
|
||||
max_retries=0,
|
||||
)
|
||||
sent_payload = mock_post.call_args.kwargs["json"]
|
||||
sent_headers = mock_post.call_args.kwargs["headers"]
|
||||
|
||||
# Forced tool wired into the request payload.
|
||||
assert sent_payload["tool_choice"] == {"type": "function", "name": "structured_response"}
|
||||
assert len(sent_payload["tools"]) == 1
|
||||
assert sent_payload["tools"][0]["name"] == "structured_response"
|
||||
assert sent_payload["parallel_tool_calls"] is False
|
||||
assert sent_headers["originator"] == "codex_cli_rs"
|
||||
assert sent_headers["User-Agent"] == "codex_cli_rs/0.0.0 (Hindsight)"
|
||||
# No prompt-injected schema in the instructions.
|
||||
assert "You must respond with valid JSON" not in sent_payload["instructions"]
|
||||
|
||||
assert isinstance(result, _Fact)
|
||||
assert result.fact == "the sky is blue"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_strict_schema_skip_validation_returns_dict():
|
||||
llm = build_llm()
|
||||
response = MagicMock()
|
||||
response.raise_for_status.return_value = None
|
||||
tool_call = LLMToolCall(id="c", name="structured_response", arguments={"fact": "x"})
|
||||
|
||||
with patch.object(llm._client, "post", new_callable=AsyncMock) as mock_post:
|
||||
mock_post.return_value = response
|
||||
with patch.object(llm, "_parse_sse_tool_stream", new_callable=AsyncMock) as mock_parse:
|
||||
mock_parse.return_value = (None, [tool_call])
|
||||
result = await llm.call(
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
response_format=_Fact,
|
||||
strict_schema=True,
|
||||
skip_validation=True,
|
||||
max_retries=0,
|
||||
)
|
||||
|
||||
assert result == {"fact": "x"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_strict_schema_retries_when_forced_tool_missing():
|
||||
llm = build_llm()
|
||||
response = MagicMock()
|
||||
response.raise_for_status.return_value = None
|
||||
|
||||
with patch.object(llm._client, "post", new_callable=AsyncMock) as mock_post:
|
||||
mock_post.return_value = response
|
||||
# Model returns no tool call at all — should raise after retries exhausted.
|
||||
with patch.object(llm, "_parse_sse_tool_stream", new_callable=AsyncMock) as mock_parse:
|
||||
mock_parse.return_value = ("some prose", [])
|
||||
with pytest.raises(RuntimeError, match="structured_response"):
|
||||
await llm.call(
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
response_format=_Fact,
|
||||
strict_schema=True,
|
||||
max_retries=0,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Non-strict fallback: escape repair keeps the retry storm from happening
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_non_strict_repairs_invalid_escapes_without_retrying():
|
||||
llm = build_llm()
|
||||
response = MagicMock()
|
||||
response.raise_for_status.return_value = None
|
||||
# Escape-heavy content the model would emit as invalid JSON.
|
||||
escape_heavy = r'{"fact": "run rig-control \d serial \s command"}'
|
||||
|
||||
with patch.object(llm._client, "post", new_callable=AsyncMock) as mock_post:
|
||||
mock_post.return_value = response
|
||||
with patch.object(llm, "_parse_sse_stream", new_callable=AsyncMock) as mock_parse:
|
||||
mock_parse.return_value = escape_heavy
|
||||
result = await llm.call(
|
||||
messages=[{"role": "user", "content": "coding transcript"}],
|
||||
response_format=_Fact,
|
||||
strict_schema=False,
|
||||
max_retries=3,
|
||||
)
|
||||
|
||||
# Parsed on the first attempt (no retry storm): the SSE stream was read once.
|
||||
assert mock_post.await_count == 1
|
||||
assert isinstance(result, _Fact)
|
||||
assert result.fact == r"run rig-control \d serial \s command"
|
||||
@@ -174,7 +174,11 @@ class TestCohereCrossEncoder:
|
||||
("What is Python?", "Python is a British comedy group"),
|
||||
]
|
||||
|
||||
scores = await encoder.predict(pairs)
|
||||
with patch(
|
||||
"hindsight_api.engine.cross_encoder.reranker_bank_attribution_headers",
|
||||
return_value={"X-Hindsight-Bank-Id": "bank-cohere-http"},
|
||||
):
|
||||
scores = await encoder.predict(pairs)
|
||||
|
||||
assert len(scores) == 3
|
||||
assert scores == [0.9, 0.7, 0.5]
|
||||
@@ -184,6 +188,7 @@ class TestCohereCrossEncoder:
|
||||
call_args = encoder._http_client._async_client.post.call_args
|
||||
assert call_args[0][0] == "https://my-endpoint.inference.ai.azure.com/models/cohere-rerank-v3-english/invoke"
|
||||
assert call_args.kwargs["json"]["model"] == "cohere-rerank-v3-english"
|
||||
assert call_args.kwargs["headers"] == {"X-Hindsight-Bank-Id": "bank-cohere-http"}
|
||||
assert call_args.kwargs["json"]["query"] == "What is Python?"
|
||||
assert len(call_args.kwargs["json"]["documents"]) == 3
|
||||
assert call_args.kwargs["json"]["return_documents"] is False
|
||||
|
||||
@@ -122,6 +122,77 @@ def test_retain_structured_chunk_size_reads_from_env():
|
||||
assert config.retain_structured_chunk_size == 9000
|
||||
|
||||
|
||||
def test_fail_on_extraction_errors_defaults_to_false(monkeypatch):
|
||||
"""Silent-success behavior is preserved by default (issue #2700)."""
|
||||
from hindsight_api.config import ENV_FAIL_ON_EXTRACTION_ERRORS, HindsightConfig
|
||||
|
||||
monkeypatch.delenv(ENV_FAIL_ON_EXTRACTION_ERRORS, raising=False)
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_PROVIDER", "mock")
|
||||
|
||||
config = HindsightConfig.from_env()
|
||||
|
||||
assert config.fail_on_extraction_errors is False
|
||||
|
||||
|
||||
def test_fail_on_extraction_errors_reads_true_from_env(monkeypatch):
|
||||
"""The opt-in escape hatch parses truthy values from the environment."""
|
||||
from hindsight_api.config import ENV_FAIL_ON_EXTRACTION_ERRORS, HindsightConfig
|
||||
|
||||
monkeypatch.setenv(ENV_FAIL_ON_EXTRACTION_ERRORS, "true")
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_PROVIDER", "mock")
|
||||
|
||||
config = HindsightConfig.from_env()
|
||||
|
||||
assert config.fail_on_extraction_errors is True
|
||||
|
||||
|
||||
def test_llm_ollama_num_ctx_defaults_to_none(monkeypatch):
|
||||
"""Unset Ollama num_ctx override lets Ollama use its model/server default."""
|
||||
from hindsight_api.config import ENV_LLM_OLLAMA_NUM_CTX, HindsightConfig
|
||||
|
||||
monkeypatch.delenv(ENV_LLM_OLLAMA_NUM_CTX, raising=False)
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_PROVIDER", "mock")
|
||||
|
||||
config = HindsightConfig.from_env()
|
||||
|
||||
assert config.llm_ollama_num_ctx is None
|
||||
|
||||
|
||||
def test_llm_ollama_num_ctx_keeps_direct_construction_default():
|
||||
"""Direct HindsightConfig construction should not require the new field."""
|
||||
from dataclasses import fields
|
||||
|
||||
from hindsight_api.config import HindsightConfig
|
||||
|
||||
config_field = next(item for item in fields(HindsightConfig) if item.name == "llm_ollama_num_ctx")
|
||||
|
||||
assert config_field.default is None
|
||||
assert config_field.kw_only
|
||||
|
||||
|
||||
def test_llm_ollama_num_ctx_reads_positive_int(monkeypatch):
|
||||
"""The native Ollama context override is parsed as a positive integer."""
|
||||
from hindsight_api.config import ENV_LLM_OLLAMA_NUM_CTX, HindsightConfig
|
||||
|
||||
monkeypatch.setenv(ENV_LLM_OLLAMA_NUM_CTX, "65536")
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_PROVIDER", "mock")
|
||||
|
||||
config = HindsightConfig.from_env()
|
||||
|
||||
assert config.llm_ollama_num_ctx == 65536
|
||||
|
||||
|
||||
def test_llm_ollama_num_ctx_rejects_non_positive_values(monkeypatch):
|
||||
"""Zero would be accepted by neither Ollama nor downstream range logic."""
|
||||
from hindsight_api.config import ENV_LLM_OLLAMA_NUM_CTX, HindsightConfig
|
||||
|
||||
monkeypatch.setenv(ENV_LLM_OLLAMA_NUM_CTX, "0")
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_PROVIDER", "mock")
|
||||
|
||||
with pytest.raises(ValueError, match=ENV_LLM_OLLAMA_NUM_CTX):
|
||||
HindsightConfig.from_env()
|
||||
|
||||
|
||||
def test_retain_structured_chunk_size_can_be_less_than_chunk_size():
|
||||
"""Structured-chunk cap can be smaller than the retain chunk target."""
|
||||
from hindsight_api.config import HindsightConfig
|
||||
@@ -706,3 +777,61 @@ def test_gemini_service_tier_empty_env_is_unset(monkeypatch):
|
||||
|
||||
config = HindsightConfig.from_env()
|
||||
assert config.llm_gemini_service_tier is None
|
||||
|
||||
|
||||
def test_operation_retention_defaults(monkeypatch):
|
||||
from hindsight_api.config import (
|
||||
ENV_OPERATION_CLEANUP_BATCH_SIZE,
|
||||
ENV_OPERATION_RETENTION_DAYS,
|
||||
HindsightConfig,
|
||||
)
|
||||
|
||||
monkeypatch.delenv(ENV_OPERATION_RETENTION_DAYS, raising=False)
|
||||
monkeypatch.delenv(ENV_OPERATION_CLEANUP_BATCH_SIZE, raising=False)
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_PROVIDER", "mock")
|
||||
|
||||
config = HindsightConfig.from_env()
|
||||
|
||||
assert config.operation_retention_days == 0
|
||||
assert config.operation_cleanup_batch_size == 1000
|
||||
assert "operation_retention_days" in HindsightConfig.get_static_fields()
|
||||
assert "operation_cleanup_batch_size" in HindsightConfig.get_static_fields()
|
||||
|
||||
|
||||
def test_operation_retention_env_overrides(monkeypatch):
|
||||
from hindsight_api.config import (
|
||||
ENV_OPERATION_CLEANUP_BATCH_SIZE,
|
||||
ENV_OPERATION_RETENTION_DAYS,
|
||||
HindsightConfig,
|
||||
)
|
||||
|
||||
monkeypatch.setenv(ENV_OPERATION_RETENTION_DAYS, "0")
|
||||
monkeypatch.setenv(ENV_OPERATION_CLEANUP_BATCH_SIZE, "37")
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_PROVIDER", "mock")
|
||||
|
||||
config = HindsightConfig.from_env()
|
||||
|
||||
assert config.operation_retention_days == 0
|
||||
assert config.operation_cleanup_batch_size == 37
|
||||
|
||||
|
||||
@pytest.mark.parametrize("raw", ["-1", "not-an-int"])
|
||||
def test_operation_retention_rejects_invalid_values(monkeypatch, raw):
|
||||
from hindsight_api.config import ENV_OPERATION_RETENTION_DAYS, HindsightConfig
|
||||
|
||||
monkeypatch.setenv(ENV_OPERATION_RETENTION_DAYS, raw)
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_PROVIDER", "mock")
|
||||
|
||||
with pytest.raises(ValueError, match=ENV_OPERATION_RETENTION_DAYS):
|
||||
HindsightConfig.from_env()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("raw", ["0", "-1", "not-an-int"])
|
||||
def test_operation_cleanup_batch_size_requires_positive_integer(monkeypatch, raw):
|
||||
from hindsight_api.config import ENV_OPERATION_CLEANUP_BATCH_SIZE, HindsightConfig
|
||||
|
||||
monkeypatch.setenv(ENV_OPERATION_CLEANUP_BATCH_SIZE, raw)
|
||||
monkeypatch.setenv("HINDSIGHT_API_LLM_PROVIDER", "mock")
|
||||
|
||||
with pytest.raises(ValueError, match=ENV_OPERATION_CLEANUP_BATCH_SIZE):
|
||||
HindsightConfig.from_env()
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user