Compare commits

..
Author SHA1 Message Date
Nicolò Boschi 0022d427d3 feat(integrations): add hindsight-opencode-coding plugin
Reflect-only long-term memory for coding agents in OpenCode, with a git+chat
backfill and (opt-in) live session write-back.

- reflect + INJECT: on a task, reflect() the symptom and push the root-cause
  answer into the system prompt (no tools/recall).
- backfill: every commit (full message + full diff, commit timestamp + git
  metadata) under a 'git' retain strategy; each chat as a JSON user/assistant
  transcript with custom extraction (<=2 coherent facts) under a 'chat' strategy;
  observations on; optional codebase knowledge pages.
- live write-back (opt-in HINDSIGHT_RETAIN_SESSIONS): every N turns upsert the
  tool-filtered transcript under a stable conversation:<sessionID> document_id.
2026-07-02 13:55:40 +02:00
Nicolò Boschi 1f9bad0858 feat(knowledge-base): default pages to living-document trigger + 4096 tokens
Client-created pages had no server curation applying a trigger, so they fell back
to the plain mental-model default (no refresh, full mode, all fact types). Make a
knowledge page a living document by default: when the client omits `trigger`, use
observation-only + delta + exclude_mental_models + refresh_after_consolidation;
when it omits `max_tokens`, default to 4096 (vs the mental-model 2048). Clients
can still override either.
2026-07-02 10:30:40 +02:00
Nicolò Boschi 138bf02f29 refactor(knowledge-base): drop server-side curation + folder missions
The knowledge base is now purely client-managed (CRUD over folders/pages); the
server no longer auto-curates. Removes the folder curator entirely and the
folder `mission` concept, and leads the sidebar with Knowledge Base.

- Remove engine/knowledge_curator.py, the curate_folder task (handler + dispatch
  + submit_async_curate_folder / _bank_folders), the post-consolidation curation
  hook, and the folder-create / mission-update curation triggers.
- Remove folder `mission` and `last_curated_at` (columns + engine + API + UI);
  keep `managed` as a client-set flag. Migration a5b6 now adds `managed` only;
  the last_curated_at migration is dropped and the unique-index migration
  repointed. Single alembic head preserved.
- API: KnowledgeNode/CreateFolderRequest/UpdateNodeRequest lose `mission`;
  PATCH node handles name/parent_id only.
- Control plane: sidebar leads with Knowledge Base (before Memories); remove the
  mission field, edit-mission dialog, and mission display from the KB view.
- Delete the curator tests; regenerate OpenAPI + SDK clients.
2026-07-02 10:30:39 +02:00
Nicolò Boschi df178aae8a refactor(hindsight-fs): mirror the knowledge-base tree, not mental models
Re-point hindsight-fs at the knowledge base so it projects a bank's folder/page
hierarchy as nested directories + .md files, instead of a flat list of mental
models.

- client: fetch GET /knowledge-base/tree + /export (two calls, any bank size)
  and join by page id; replaces the paginated mental-models list.
- format: planMirror() walks the tree into folder dirs + page files at nested
  paths (slug per segment, collision-safe); pages render the page's OKF doc.
- sync: create folder dirs, write pages at nested paths, prune removed pages and
  emptied folders; state keyed by relative path + tracked dirs.
- config/cli: drop the mental-model `detail` flag; `list` prints folders+pages;
  help/README updated. Tests rewritten for the tree/export model.

Verified live against a bank's knowledge base: the `people` folder mirrors to
people/anna.md + people/marco.md with OKF frontmatter.
2026-07-02 10:30:39 +02:00
Nicolò Boschi a43026b8f4 feat(hindsight-fs): mirror a bank's mental models as a live local folder
Add @vectorize-io/hindsight-fs, a CLI under hindsight-tools/ that mirrors a
Hindsight bank's mental models as real markdown files (YAML frontmatter + body)
in a local directory, refreshed from the API on an interval. Once mounted,
ordinary shell tools (ls, cat, grep, find, ...) work against current memory.

- Pull-based sync engine: full list each tick, write changed/new/tampered
  files, skip unchanged (content-hashed), prune deleted models. Atomic writes;
  a transient API error never wipes the mirror.
- One-way mirror enforced two ways: files are read-only (0444) so agent edits
  fail with EACCES, plus a tamper-revert backstop that compares on-disk bytes
  and overwrites drift on the next pass. --writable opts out.
- Commands: mount/start/stop/restart/sync/status/list/logs/unmount. Background
  daemon via detached process + pidfile; per-mount config is remembered.
- status doubles as a healthcheck: --json report and a non-zero exit when the
  mount is dead/failed/stale (--stale-after overrides the threshold).
- Tests: unit (sync engine, frontmatter, health) + e2e that spawns the real
  CLI against a mock API and exercises real bash commands. 26 tests.
2026-07-02 10:30:39 +02:00
Nicolò Boschi 5c425e276e feat(knowledge-base): self-curating knowledge base (OKF pages + folder missions)
Server-side knowledge base: a hierarchy of folders and pages over mental
models, projected to the Open Knowledge Format, with a mission-driven curator
that maintains pages automatically after each consolidation.

- knowledge_pages table (PG + Oracle): parent_id tree, kind folder/page,
  mission, managed, last_curated_at; partial unique index on (folder, name)
  for concurrency-safe dedup; added to BACKUP_TABLES.
- api/okf.py: OKF serializer (frontmatter + body, index/log, constellation graph).
- engine/knowledge_curator.py: folder curator (LLM op plan + safe apply); reads
  new memories since last curation (delta, not recall); ops create/merge/delete
  page + spawn sub-folder (bounded depth<=3, <=8). Runs as an async curate_folder
  task on folder/mission create and after consolidation. Curator pages use an
  observation-only delta trigger with exclude_mental_models.
- MemoryEngine: folder/page CRUD, tree, curate, async submit + worker handler.
- /v1/default/banks/{bank}/knowledge-base/* endpoints.
- Control plane: knowledge-base tree view + constellation toggle, missions,
  OKF page panel + bundle export; proxies, client, sidebar, i18n.
- Tests: okf unit, knowledge-base HTTP, curator apply + dedup guard, hs_llm_core e2e.
- Regenerated OpenAPI + SDK clients + docs-skill.
2026-07-02 10:30:39 +02:00
750 changed files with 27527 additions and 60052 deletions
+1 -6
View File
@@ -1,7 +1,7 @@
{
"$schema": "https://anthropic.com/claude-code/marketplace.schema.json",
"name": "hindsight",
"version": "0.7.5",
"version": "0.7.2",
"description": "Official Hindsight integrations for Claude Code",
"owner": {
"name": "vectorize-io"
@@ -11,11 +11,6 @@
"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"
}
]
}
-13
View File
@@ -78,11 +78,6 @@ 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.
@@ -209,14 +204,6 @@ 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:
-43
View File
@@ -21,21 +21,6 @@ HINDSIGHT_API_LLM_BASE_URL=https://api.openai.com/v1
# HINDSIGHT_API_LLM_TEMPERATURE_REFLECT=0.9
# HINDSIGHT_API_LLM_TEMPERATURE_CONSOLIDATION=0.0
# Grammar-enforce structured output (json_schema strict) instead of the soft
# schema-in-prompt path. Helps weaker self-hosted models that emit prose preambles
# or invalid JSON. The global override below applies to every operation;
# per-operation overrides take precedence, in both directions -- set one to false
# to opt that operation out while the global flag is on.
# HINDSIGHT_API_LLM_STRICT_SCHEMA=false
# HINDSIGHT_API_LLM_STRICT_SCHEMA_RETAIN=true
# HINDSIGHT_API_LLM_STRICT_SCHEMA_REFLECT=true
# HINDSIGHT_API_LLM_STRICT_SCHEMA_CONSOLIDATION=true
# Diagnostic: on any LLM 4xx, log the exact assembled request ([LLM_4XX_DUMP]) --
# serialized request config (message bodies stripped) + capped per-message previews.
# For debugging otherwise-unreproducible rejected calls. Off by default.
# HINDSIGHT_API_LLM_DEBUG_DUMP_4XX=false
# Example: Anthropic Claude configuration
# HINDSIGHT_API_LLM_PROVIDER=anthropic
# HINDSIGHT_API_LLM_API_KEY=your-anthropic-api-key
@@ -74,15 +59,6 @@ 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
@@ -104,10 +80,6 @@ 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
@@ -123,10 +95,7 @@ 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)
@@ -146,11 +115,6 @@ 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
@@ -169,8 +133,6 @@ 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
@@ -210,13 +172,8 @@ 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
+17 -89
View File
@@ -32,7 +32,6 @@ jobs:
integration-tests: ${{ steps.filter.outputs.integration-tests }}
integrations-openclaw: ${{ steps.filter.outputs.integrations-openclaw }}
integrations-ai-sdk: ${{ steps.filter.outputs.integrations-ai-sdk }}
integrations-agentos: ${{ steps.filter.outputs.integrations-agentos }}
integrations-agent-framework: ${{ steps.filter.outputs.integrations-agent-framework }}
integrations-composio: ${{ steps.filter.outputs.integrations-composio }}
integrations-chat: ${{ steps.filter.outputs.integrations-chat }}
@@ -42,7 +41,6 @@ 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 }}
@@ -136,8 +134,6 @@ jobs:
- 'hindsight-integrations/openclaw/**'
integrations-ai-sdk:
- 'hindsight-integrations/ai-sdk/**'
integrations-agentos:
- 'hindsight-integrations/agentos/**'
integrations-agent-framework:
- 'hindsight-integrations/agent-framework/**'
integrations-composio:
@@ -184,8 +180,6 @@ jobs:
- 'hindsight-integrations/cursor/**'
integrations-zed:
- 'hindsight-integrations/zed/**'
integrations-zcode:
- 'hindsight-integrations/zcode/**'
integrations-n8n:
- 'hindsight-integrations/n8n/**'
integrations-zapier:
@@ -526,17 +520,22 @@ jobs:
with:
ref: ${{ github.event.pull_request.head.sha || '' }}
- name: Set up Node.js
uses: actions/setup-node@v6
- name: Set up Python
uses: actions/setup-python@v6
with:
node-version: '22'
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
- name: Run tests
working-directory: ./hindsight-integrations/zed
# 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
# 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"
test-omo-integration:
needs: [detect-changes]
@@ -703,43 +702,6 @@ 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: >-
@@ -803,37 +765,6 @@ jobs:
working-directory: ./hindsight-integrations/ai-sdk
run: npm run test:deno
build-agentos-integration:
needs: [detect-changes]
if: >-
(github.event_name == 'workflow_dispatch' ||
needs.detect-changes.outputs.integrations-agentos == '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: Set up Node.js
uses: actions/setup-node@v6
with:
node-version: '22'
- name: Install dependencies
working-directory: ./hindsight-integrations/agentos
run: npm ci
- name: Run tests
working-directory: ./hindsight-integrations/agentos
run: npm test
- name: Build
working-directory: ./hindsight-integrations/agentos
run: npm run build
test-opencode-integration:
needs: [detect-changes]
if: >-
@@ -1298,10 +1229,10 @@ jobs:
build-docs:
needs: [detect-changes]
# 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.
if: >-
(github.event_name == 'workflow_dispatch' ||
needs.detect-changes.outputs.docs == 'true' ||
needs.detect-changes.outputs.ci == 'true')
runs-on: ubuntu-latest
timeout-minutes: 30
@@ -4840,8 +4771,7 @@ jobs:
cd ../hindsight-embed && uv sync --frozen --index-strategy unsafe-best-match
- name: Run generate-openapi
working-directory: hindsight-dev
run: uv run generate-openapi
run: ./scripts/generate-openapi.sh
- name: Run generate-bank-template-schema
run: ./scripts/generate-bank-template-schema.sh
@@ -4979,10 +4909,8 @@ jobs:
- test-github-copilot-integration
- test-codex-integration
- test-cursor-cli-integration
- test-zcode-integration
- build-ai-sdk-integration
- test-ai-sdk-integration-deno
- build-agentos-integration
- test-opencode-integration
- test-eve-integration
- test-omo-integration
-1
View File
@@ -41,7 +41,6 @@ nltk_data/
logs/
.DS_Store
.sesskey
# Generated docs files
hindsight-docs/static/llms-full.txt
+1
View File
@@ -0,0 +1 @@
fcac2839-1db5-432f-91e1-c5dac07d7290
+9 -30
View File
@@ -41,43 +41,28 @@ RUN apt-get update && apt-get install -y \
&& rm -rf /var/lib/apt/lists/* \
&& pip install --no-cache-dir uv
# Copy the workspace lock and member metadata before source code so dependency
# installation stays cacheable while matching the versions tested in CI.
COPY pyproject.toml uv.lock ./
COPY hindsight-all/pyproject.toml ./hindsight-all/
COPY hindsight-api/pyproject.toml ./hindsight-api/
# Copy dependency files and README (required by pyproject.toml)
COPY hindsight-api-slim/pyproject.toml ./api/
COPY hindsight-api-slim/README.md ./api/
COPY hindsight-all-slim/pyproject.toml ./hindsight-all-slim/
COPY hindsight-dev/pyproject.toml ./hindsight-dev/
COPY hindsight-clients/python/pyproject.toml ./hindsight-clients/python/
COPY hindsight-embed/pyproject.toml ./hindsight-embed/
RUN ln -s api hindsight-api-slim
WORKDIR /app/api
# Sync dependencies using appropriate extras based on INCLUDE_LOCAL_MODELS
# local-ml: torch, sentence-transformers, transformers, einops, flashrank, mlx (optional)
# embedded-db: pg0-embedded (always included for embedded PostgreSQL support)
# ONNX Runtime embeddings are intentionally not bundled into the official
# standalone image; install the local-onnx extra in custom images when needed.
ENV UV_PROJECT_ENVIRONMENT=/app/api/.venv
RUN if [ "$INCLUDE_LOCAL_MODELS" = "true" ]; then \
uv sync --locked --package hindsight-api-slim --no-install-package hindsight-api-slim --extra local-ml --extra embedded-db; \
uv sync --extra local-ml --extra embedded-db; \
else \
uv sync --locked --package hindsight-api-slim --no-install-package hindsight-api-slim --extra embedded-db; \
uv sync --extra embedded-db; \
fi
# Copy source code (alembic migrations are inside hindsight_api/)
WORKDIR /app/api
COPY hindsight-api-slim/hindsight_api ./hindsight_api
# Install the local package from the same validated lock after source is present.
WORKDIR /app
RUN if [ "$INCLUDE_LOCAL_MODELS" = "true" ]; then \
uv sync --locked --package hindsight-api-slim --extra local-ml --extra embedded-db; \
else \
uv sync --locked --package hindsight-api-slim --extra embedded-db; \
fi \
&& uv pip check --python /app/api/.venv/bin/python
# Install the local package (uv sync only installed dependencies, not the package itself)
RUN uv pip install -e .
# =============================================================================
# Stage: SDK Builder (needed for Control Plane)
@@ -160,8 +145,6 @@ FROM python:3.11-slim AS api-only
WORKDIR /app
# Note: libicu version varies by Debian version - try common versions in order
# Runtime images use uv directly; remove pip build tooling after installation so
# vulnerable setuptools-vendored packages and wheel are not shipped in production.
RUN apt-get update && apt-get install -y \
curl \
procps \
@@ -171,8 +154,7 @@ RUN apt-get update && apt-get install -y \
libossp-uuid16 \
&& (apt-get install -y libicu72 2>/dev/null || apt-get install -y libicu74 2>/dev/null || apt-get install -y libicu76 2>/dev/null || true) \
&& rm -rf /var/lib/apt/lists/* \
&& pip install --no-cache-dir uv \
&& pip uninstall --yes setuptools wheel
&& pip install --no-cache-dir uv
RUN useradd -m -s /bin/bash hindsight
@@ -310,8 +292,6 @@ WORKDIR /app
# Install Node.js, curl, uv, and system dependencies
# Note: libicu version varies by Debian version - try common versions in order
# Runtime images use uv directly; remove pip build tooling after installation so
# vulnerable setuptools-vendored packages and wheel are not shipped in production.
RUN apt-get update && apt-get install -y \
curl \
procps \
@@ -323,8 +303,7 @@ RUN apt-get update && apt-get install -y \
&& curl -fsSL https://deb.nodesource.com/setup_20.x | bash - \
&& apt-get install -y nodejs \
&& rm -rf /var/lib/apt/lists/* \
&& pip install --no-cache-dir uv \
&& pip uninstall --yes setuptools wheel
&& pip install --no-cache-dir uv
RUN useradd -m -s /bin/bash hindsight
+9 -152
View File
@@ -18,10 +18,8 @@ import typer
from ..config import DEFAULT_DATABASE_SCHEMA, HindsightConfig
from ..engine.memory_engine import _current_schema
from ..engine.retain.bank_utils import _vector_index_clause
from ..engine.schema import fq_table_explicit as _fq_table
from ..engine.transfer import export_bank
from ..engine.vector_index_health import SchemaVectorIndexResult, repair_vector_indexes
from ..extensions import TenantExtension, load_extension
from ..pg0 import parse_pg0_url, resolve_database_url
@@ -58,6 +56,7 @@ BACKUP_TABLES = [
"observation_history",
"mental_models",
"mental_model_history",
"knowledge_pages",
"directives",
"async_operations",
"webhooks",
@@ -78,8 +77,7 @@ 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).
"""
_pg0 = parse_pg0_url(db_url)
is_pg0, instance_name = _pg0.is_pg0, _pg0.instance_name
is_pg0, instance_name, _ = parse_pg0_url(db_url)
if is_pg0:
typer.echo(f"Starting embedded PostgreSQL (instance: {instance_name})...")
conn = await asyncpg.connect(await resolve_database_url(db_url))
@@ -180,8 +178,7 @@ 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."""
_pg0 = parse_pg0_url(db_url)
is_pg0, instance_name = _pg0.is_pg0, _pg0.instance_name
is_pg0, instance_name, _ = parse_pg0_url(db_url)
if is_pg0:
typer.echo(f"Starting embedded PostgreSQL (instance: {instance_name})...")
resolved_url = await resolve_database_url(db_url)
@@ -190,8 +187,7 @@ 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."""
_pg0 = parse_pg0_url(db_url)
is_pg0, instance_name = _pg0.is_pg0, _pg0.instance_name
is_pg0, instance_name, _ = parse_pg0_url(db_url)
if is_pg0:
typer.echo(f"Starting embedded PostgreSQL (instance: {instance_name})...")
resolved_url = await resolve_database_url(db_url)
@@ -266,8 +262,7 @@ async def _run_migration(
"""Resolve database URL and run migrations for one schema or all discovered schemas."""
from ..migrations import run_migrations_for_schemas
_pg0 = parse_pg0_url(db_url)
is_pg0, instance_name = _pg0.is_pg0, _pg0.instance_name
is_pg0, instance_name, _ = parse_pg0_url(db_url)
if is_pg0:
typer.echo(f"Starting embedded PostgreSQL (instance: {instance_name})...")
resolved_url = await resolve_database_url(db_url)
@@ -359,134 +354,6 @@ def run_db_migration(
typer.echo(f"Database migrations completed successfully for {len(schemas)} schema(s)")
async def _resolve_schemas(base_schema: str | None) -> list[str]:
"""Base schema plus every discovered tenant schema, de-duplicated in order."""
schemas = [base_schema or DEFAULT_DATABASE_SCHEMA]
tenant_extension = load_extension("TENANT", TenantExtension)
if tenant_extension:
tenants = await tenant_extension.list_tenants()
schemas.extend(tenant.schema for tenant in tenants if tenant.schema)
return list(dict.fromkeys(schemas))
async def _run_repair_bank(
db_url: str,
*,
base_schema: str,
schema: str | None,
bank_id: str | None,
dry_run: bool,
) -> list[SchemaVectorIndexResult]:
"""Reconcile per-(bank, fact_type) vector index coverage over a raw connection.
A single autocommit connection is used because ``CREATE INDEX CONCURRENTLY``
(used by ``repair_vector_indexes``) cannot run inside a transaction block.
"""
schemas = [schema] if schema else await _resolve_schemas(base_schema)
index_clause = _vector_index_clause()
# Guarded by the command, but assert so this helper is never called for a
# backend without per-bank indexes.
assert index_clause is not None
conn = await _admin_connect(db_url)
try:
results = await repair_vector_indexes(conn, schemas, index_clause, dry_run=dry_run, bank_id=bank_id)
for result in results:
typer.echo(
f" schema '{result.schema}': {result.banks_scanned} bank(s) scanned, "
f"{result.already_present} present, {result.created} created, "
f"{result.skipped} to-create (dry-run), {result.failed} failed"
)
return results
finally:
await conn.close()
@app.command(name="repair-bank")
def repair_bank(
bank_id: str | None = typer.Option(
None,
"--bank",
"-b",
help="Bank id to repair. Mutually exclusive with --all.",
),
all_banks: bool = typer.Option(
False,
"--all",
help="Repair every bank in the base schema and all discovered tenant schemas.",
),
schema: str | None = typer.Option(
None,
"--schema",
"-s",
help="Limit to a single schema. Defaults to the base schema plus discovered tenant schemas.",
),
dry_run: bool = typer.Option(
False,
"--dry-run",
help="Report what would be repaired without creating or dropping any index.",
),
):
"""Verify and repair a bank's per-(bank, fact_type) vector index coverage.
Per-bank partial vector indexes are created when a bank is first created
(instant on an empty bank). Banks that arrive populated — via logical
restore, a cross-version upgrade, or a vector-extension switch — never hit
that path, so their recall silently falls back to a global index +
post-filter (slower, under-returning). This command detects missing OR
invalid coverage (an INVALID leftover or an index whose access method
drifted counts as missing) and rebuilds it with CREATE INDEX CONCURRENTLY,
so it never blocks the live fleet. Idempotent and safe to re-run — the
escape hatch after a restore, upgrade, or backend switch.
"""
if bool(bank_id) == all_banks:
typer.echo("Error: pass exactly one of --bank <id> or --all.", err=True)
raise typer.Exit(2)
config = HindsightConfig.from_env()
if not config.database_url:
typer.echo("Error: Database URL not configured.", err=True)
typer.echo("Set HINDSIGHT_API_DATABASE_URL environment variable.", err=True)
raise typer.Exit(1)
# Backend guard: backends with a single global vector index (AlloyDB ScaNN,
# Oracle) have no per-bank indexes to repair.
if _vector_index_clause() is None:
typer.echo("Configured vector backend does not use per-bank vector indexes — nothing to repair.")
return
target = f"bank '{bank_id}'" if bank_id else "all banks"
scope = f"schema '{schema}'" if schema else "base schema and all discovered tenant schemas"
typer.echo(f"Repairing per-bank vector indexes for {target} across {scope}...")
if dry_run:
typer.echo("Dry run: no indexes will be created or dropped.")
results = asyncio.run(
_run_repair_bank(
config.database_url,
base_schema=config.database_schema,
schema=schema,
bank_id=bank_id,
dry_run=dry_run,
)
)
total_banks = sum(r.banks_scanned for r in results)
total_present = sum(r.already_present for r in results)
total_created = sum(r.created for r in results)
total_skipped = sum(r.skipped for r in results)
total_failed = sum(r.failed for r in results)
typer.echo(
f"Done: {len(results)} schema(s), {total_banks} bank(s) scanned, "
f"{total_present} already present, {total_created} created, "
f"{total_skipped} to-create (dry-run), {total_failed} failed"
)
if total_failed:
failed_names = [name for r in results for name in r.failed_indexes]
typer.echo(f"Failed indexes (dropped, retry with a re-run): {', '.join(failed_names)}", err=True)
raise typer.Exit(1)
async def _run_export_bank(db_url: str, bank_id: str, output: Path, schema: str, include_history: bool) -> int:
"""Export a whole bank to a ZIP archive."""
conn = await _admin_connect(db_url)
@@ -494,14 +361,7 @@ async def _run_export_bank(db_url: str, bank_id: str, output: Path, schema: str,
# export_bank resolves table names via fq_table (the _current_schema
# contextvar); set it so the raw connection targets the right schema.
_current_schema.set(schema)
# _admin_connect registers JSON codecs, so row dumps already contain
# decoded Python values (including JSON scalar strings).
data = await export_bank(
conn,
bank_id,
include_history=include_history,
bank_rows_json_encoding="decoded",
)
data = await export_bank(conn, bank_id, include_history=include_history)
finally:
await conn.close()
@@ -613,8 +473,7 @@ 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."""
_pg0 = parse_pg0_url(db_url)
is_pg0, instance_name = _pg0.is_pg0, _pg0.instance_name
is_pg0, instance_name, _ = parse_pg0_url(db_url)
if is_pg0:
typer.echo(f"Starting embedded PostgreSQL (instance: {instance_name})...")
resolved_url = await resolve_database_url(db_url)
@@ -673,8 +532,7 @@ 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."""
_pg0 = parse_pg0_url(db_url)
is_pg0, instance_name = _pg0.is_pg0, _pg0.instance_name
is_pg0, instance_name, _ = parse_pg0_url(db_url)
if is_pg0:
typer.echo(f"Starting embedded PostgreSQL (instance: {instance_name})...")
resolved_url = await resolve_database_url(db_url)
@@ -739,8 +597,7 @@ 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."""
_pg0 = parse_pg0_url(db_url)
is_pg0, instance_name = _pg0.is_pg0, _pg0.instance_name
is_pg0, instance_name, _ = parse_pg0_url(db_url)
if is_pg0:
typer.echo(f"Starting embedded PostgreSQL (instance: {instance_name})...")
resolved_url = await resolve_database_url(db_url)
@@ -96,8 +96,7 @@ def get_database_url() -> str:
# for the sync engine used during migrations.
database_url = to_libpq_url(database_url)
# Alembic stores options through ConfigParser, where '%' is interpolation.
config.set_main_option("sqlalchemy.url", database_url.replace("%", "%%"))
config.set_main_option("sqlalchemy.url", database_url)
return database_url
@@ -0,0 +1,52 @@
"""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)
@@ -1,82 +0,0 @@
"""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)
@@ -0,0 +1,110 @@
"""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)
@@ -1,259 +0,0 @@
"""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)
@@ -0,0 +1,71 @@
"""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)
@@ -1,150 +0,0 @@
"""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)
@@ -1,96 +0,0 @@
"""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)
+492 -112
View File
@@ -18,6 +18,7 @@ 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 (
@@ -52,7 +53,6 @@ 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', 'custom', 'verbatim', or 'chunks'.",
description="Fact extraction mode: 'concise' (default), 'verbose', or 'custom'.",
)
retain_custom_instructions: str | None = Field(
default=None,
@@ -1434,7 +1434,6 @@ class ListMemoryUnitsResponse(BaseModel):
"date": "2024-01-15T10:30:00Z",
"type": "world",
"entities": "Alice (PERSON), Google (ORGANIZATION)",
"metadata": {"source": "slack", "channel": "engineering"},
}
],
"total": 150,
@@ -1668,8 +1667,8 @@ class UpdateMemoryRequest(BaseModel):
@model_validator(mode="after")
def _require_an_edit(self) -> "UpdateMemoryRequest":
has_value_edit = any(
v is not None
if all(
v is None
for v in (
self.text,
self.context,
@@ -1679,9 +1678,7 @@ 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'.")
@@ -2114,6 +2111,150 @@ 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."""
@@ -2207,8 +2348,7 @@ 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', 'custom', 'verbatim', or 'chunks'",
default=None, description="Fact extraction mode: 'concise' (default), 'verbose', or 'custom'"
)
retain_custom_instructions: str | None = Field(
default=None, description="Custom extraction prompt (when mode='custom')"
@@ -2434,10 +2574,10 @@ def validate_bank_template(manifest: "BankTemplateManifest") -> list[str]:
if manifest.bank:
bank = manifest.bank
if bank.retain_extraction_mode is not None:
if bank.retain_extraction_mode not in RETAIN_EXTRACTION_MODES:
valid_modes = ("concise", "verbose", "custom", "chunks")
if bank.retain_extraction_mode not in valid_modes:
errors.append(
"bank.retain_extraction_mode: "
f"must be one of {RETAIN_EXTRACTION_MODES}, got '{bank.retain_extraction_mode}'"
f"bank.retain_extraction_mode: must be one of {valid_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'")
@@ -2730,24 +2870,6 @@ class RetryOperationResponse(BaseModel):
operation_id: str
class DeleteOperationResponse(BaseModel):
"""Response model for delete operation endpoint."""
model_config = ConfigDict(
json_schema_extra={
"example": {
"success": True,
"message": "Operation 550e8400-e29b-41d4-a716-446655440000 deleted",
"operation_id": "550e8400-e29b-41d4-a716-446655440000",
}
}
)
success: bool
message: str
operation_id: str
class ChildOperationStatus(BaseModel):
"""Status of a child operation (for batch operations)."""
@@ -2839,7 +2961,7 @@ class FeaturesInfo(BaseModel):
file_upload_api: bool = Field(description="Whether file upload/conversion API is enabled")
document_export_api: bool = Field(description="Whether the document export endpoint is enabled")
document_import_api: bool = Field(description="Whether the document import endpoint is enabled")
audit_log: bool = Field(description="Whether audit logging is enabled by default (overridable per bank)")
audit_log: bool = Field(description="Whether audit logging is enabled")
llm_trace: bool = Field(description="Whether per-bank LLM request tracing is enabled")
store_document_text: bool = Field(
description="Whether raw source text is persisted. When false, document/chunk source text is not stored."
@@ -3009,15 +3131,10 @@ def _make_audited_http(audit_logger_getter: Callable[[], AuditLogger | None]):
@wraps(func)
async def wrapper(*args, **kwargs):
al = audit_logger_getter()
# Cheap bank-independent pre-filter first, then the per-bank
# decision (audit_log_enabled is overridable per bank).
if al is None or not al.action_allowed(action):
if al is None or not al.is_enabled(action):
return await func(*args, **kwargs)
bank_id = kwargs.get("bank_id")
if not await al.should_log(action, bank_id, kwargs.get("request_context")):
return await func(*args, **kwargs)
started_at = _dt.now(_tz.utc)
req_data = None
@@ -3514,8 +3631,8 @@ def _register_routes(app: FastAPI):
Returns version info and feature flags that can be used by clients
to determine which capabilities are available.
Note: the observations and audit_log flags show the global default.
Individual banks may override these via bank-specific configuration.
Note: observations flag shows the global default. Individual banks
may override this setting via bank-specific configuration.
"""
from hindsight_api import __version__
from hindsight_api.config import _get_raw_config
@@ -3611,8 +3728,6 @@ def _register_routes(app: FastAPI):
consolidation_state: str | None = None,
state: str | None = None,
document_id: str | None = None,
tags: list[str] | None = Query(default=None),
tags_match: TagsMatch = Query(default="any"),
limit: int = Query(default=100, ge=0),
offset: int = Query(default=0, ge=0),
request_context: RequestContext = Depends(get_request_context),
@@ -3629,10 +3744,6 @@ def _register_routes(app: FastAPI):
q: Search query for full-text search (searches text and context)
consolidation_state: Filter by consolidation state for source memories
(world/experience). One of 'failed', 'pending', or 'done'.
tags: Optional list of tag names to filter by.
tags_match: How to combine tags: 'any' (OR, default) or 'all' (AND) both
also include untagged memories; 'any_strict'/'all_strict' exclude
untagged; 'exact' matches the tag set exactly.
limit: Maximum number of results (default: 100)
offset: Offset for pagination (default: 0)
"""
@@ -3644,8 +3755,6 @@ def _register_routes(app: FastAPI):
consolidation_state=consolidation_state,
state=state,
document_id=document_id,
tags=tags,
tags_match=tags_match,
limit=limit,
offset=offset,
request_context=request_context,
@@ -3778,7 +3887,6 @@ 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,
@@ -3787,23 +3895,13 @@ 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=occurred_start,
occurred_end=occurred_end,
occurred_start=request.occurred_start,
occurred_end=request.occurred_end,
new_fact_type=request.fact_type,
entities=request.entities,
state=request.state,
@@ -4828,6 +4926,333 @@ 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
# =========================================================================
@@ -5454,9 +5879,8 @@ 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', 'processing', 'completed', "
"'failed', or 'cancelled'. Completed operations remain queryable with their payload for the configured "
"retention window and are pruned afterward.",
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.",
operation_id="get_operation_status",
tags=["Operations"],
)
@@ -5561,42 +5985,6 @@ def _register_routes(app: FastAPI):
logger.error(f"Error in POST /v1/default/banks/{bank_id}/operations/{operation_id}/retry: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
@app.delete(
"/v1/default/banks/{bank_id}/operations/{operation_id}/delete",
response_model=DeleteOperationResponse,
summary="Delete a terminal async operation",
description="Permanently remove a failed, cancelled, or completed async operation record",
operation_id="delete_operation",
tags=["Operations"],
)
@audited("delete_operation", request_param=None)
async def api_delete_operation(
bank_id: str, operation_id: str, request_context: RequestContext = Depends(get_request_context)
):
"""Delete a terminal async operation record."""
try:
try:
uuid.UUID(operation_id)
except ValueError:
raise HTTPException(status_code=400, detail=f"Invalid operation_id format: {operation_id}")
result = await app.state.memory.delete_operation(bank_id, operation_id, request_context=request_context)
return DeleteOperationResponse(**result)
except ValueError as e:
raise HTTPException(status_code=404, 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 DELETE /v1/default/banks/{bank_id}/operations/{operation_id}/delete: {error_detail}"
)
raise HTTPException(status_code=500, detail=str(e))
@app.get(
"/v1/default/banks/{bank_id}/profile",
response_model=BankProfileResponse,
@@ -5732,12 +6120,8 @@ def _register_routes(app: FastAPI):
):
"""Create or update an agent with disposition and mission."""
try:
# 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,
)
# Ensure bank exists by getting profile (auto-creates with defaults)
await app.state.memory.get_bank_profile(bank_id, request_context=request_context)
# Update name if provided (stored in DB for display only, deprecated)
if request.name is not None:
@@ -5929,12 +6313,8 @@ def _register_routes(app: FastAPI):
dry_run=True,
)
# 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,
)
# Ensure bank exists (auto-creates with defaults if needed)
await app.state.memory.get_bank_profile(bank_id, request_context=request_context)
return await apply_bank_template_manifest(
memory=app.state.memory,
+263
View File
@@ -0,0 +1,263 @@
"""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)
+2 -194
View File
@@ -145,17 +145,8 @@ ENV_LLM_BEDROCK_SERVICE_TIER = "HINDSIGHT_API_LLM_BEDROCK_SERVICE_TIER"
ENV_LLM_GEMINI_SERVICE_TIER = "HINDSIGHT_API_LLM_GEMINI_SERVICE_TIER"
ENV_LLM_EXTRA_BODY = "HINDSIGHT_API_LLM_EXTRA_BODY"
ENV_LLM_DEFAULT_HEADERS = "HINDSIGHT_API_LLM_DEFAULT_HEADERS"
# Grammar-enforced structured output. The global flag applies to every internal
# LLM call; the per-operation variants override it for a single operation, so an
# operator can enable strict schema where it fixes malformed/truncated JSON
# without paying the retry cost on operations whose model can't satisfy it.
# Resolution per operation: per-operation env -> global env -> built-in default.
ENV_LLM_STRICT_SCHEMA = "HINDSIGHT_API_LLM_STRICT_SCHEMA"
ENV_LLM_STRICT_SCHEMA_RETAIN = "HINDSIGHT_API_LLM_STRICT_SCHEMA_RETAIN"
ENV_LLM_STRICT_SCHEMA_REFLECT = "HINDSIGHT_API_LLM_STRICT_SCHEMA_REFLECT"
ENV_LLM_STRICT_SCHEMA_CONSOLIDATION = "HINDSIGHT_API_LLM_STRICT_SCHEMA_CONSOLIDATION"
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
@@ -257,21 +248,6 @@ def _resolve_operation_temperature(operation_env: str, default: float) -> float
return _parse_temperature(raw)
def _resolve_operation_strict_schema(operation_env: str) -> bool:
"""Resolve a per-operation strict-schema flag: per-op env -> global env -> default.
Resolved to a concrete bool here rather than left as None, so the call site
passes an explicit value and a per-operation "false" can override a global
"true" (the wrapper honours an explicit False -- see LLMConfig.call).
"""
raw = os.getenv(operation_env)
if raw is None:
raw = os.getenv(ENV_LLM_STRICT_SCHEMA)
if raw is None:
return DEFAULT_LLM_STRICT_SCHEMA
return raw.strip().lower() in ("true", "1")
# Per-operation LLM configuration (optional, falls back to global LLM config)
ENV_RETAIN_LLM_PROVIDER = "HINDSIGHT_API_RETAIN_LLM_PROVIDER"
ENV_RETAIN_LLM_API_KEY = "HINDSIGHT_API_RETAIN_LLM_API_KEY"
@@ -397,7 +373,6 @@ 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"
@@ -407,7 +382,6 @@ 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"
@@ -511,11 +485,6 @@ ENV_LLM_GEMINI_SAFETY_SETTINGS = "HINDSIGHT_API_LLM_GEMINI_SAFETY_SETTINGS"
# banks, and creation soft-fails to an uncached call, so it never breaks a request.
ENV_LLM_PROMPT_CACHE_ENABLED = "HINDSIGHT_API_LLM_PROMPT_CACHE_ENABLED"
# Opt-in diagnostic: when truthy, log the exact request behind any LLM 4xx (the
# serialized request config with message bodies stripped + length-capped per-message
# previews). Off by default; server-level only. See engine/providers/llm_debug.py.
ENV_LLM_DEBUG_DUMP_4XX = "HINDSIGHT_API_LLM_DEBUG_DUMP_4XX"
# Retain settings
ENV_RETAIN_MAX_COMPLETION_TOKENS = "HINDSIGHT_API_RETAIN_MAX_COMPLETION_TOKENS"
ENV_RETAIN_CHUNK_SIZE = "HINDSIGHT_API_RETAIN_CHUNK_SIZE"
@@ -615,7 +584,6 @@ 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
@@ -630,8 +598,6 @@ 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.
@@ -651,7 +617,6 @@ 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"
@@ -673,7 +638,6 @@ 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.
@@ -693,16 +657,10 @@ ENV_RECENCY_DECAY_LINEAR_WINDOW_DAYS = "HINDSIGHT_API_RECENCY_DECAY_LINEAR_WINDO
ENV_RECENCY_DECAY_HALFLIFE_DAYS = "HINDSIGHT_API_RECENCY_DECAY_HALFLIFE_DAYS"
# Audit log settings
# AUDIT_LOG_ENABLED is the deployment-wide default and is overridable per bank
# (and per tenant) through the bank config API, so auditing can be turned on for
# individual banks without enabling it everywhere.
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"
@@ -803,7 +761,6 @@ 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
@@ -832,9 +789,6 @@ 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.
@@ -955,10 +909,6 @@ 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"
@@ -1009,7 +959,6 @@ DEFAULT_RETAIN_ENTITY_LOOKUP = "trigram" # "full" or "trigram"
DEFAULT_RETAIN_ENTITY_RESOLUTION_BATCH_SIZE = 100 # Unique entity names per pg_trgm candidate lookup query
DEFAULT_RETAIN_BATCH_ENABLED = False # Use LLM Batch API for fact extraction (only when async=True)
DEFAULT_LLM_PROMPT_CACHE_ENABLED = True # Reuse the fixed system prefix via provider prompt caching
DEFAULT_LLM_DEBUG_DUMP_4XX = False # Log the exact request behind any LLM 4xx (diagnostic, off by default)
DEFAULT_RETAIN_BATCH_POLL_INTERVAL_SECONDS = 60 # Batch API polling interval in seconds
# File storage defaults
@@ -1091,14 +1040,6 @@ 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)
@@ -1109,20 +1050,10 @@ 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)
@@ -1163,11 +1094,6 @@ 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
@@ -1283,19 +1209,6 @@ 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 == "":
@@ -1303,25 +1216,6 @@ 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}")
@@ -1678,20 +1572,11 @@ class HindsightConfig:
dict | None
) # Custom headers passed as default_headers to provider SDK clients (e.g. {"X-Component-Id": "hindsight"} for proxies / request tracing)
llm_strict_schema: bool # Grammar-enforce structured output via the provider's strongest schema mode (see DEFAULT_LLM_STRICT_SCHEMA)
# Per-operation strict-schema overrides. Resolved from the per-operation env
# var, falling back to llm_strict_schema's global env var. See
# ENV_LLM_STRICT_SCHEMA and _resolve_operation_strict_schema.
llm_strict_schema_retain: bool
llm_strict_schema_reflect: bool
llm_strict_schema_consolidation: bool
# Tags outbound OpenAI-compatible LLM + embedding calls with `user=<bank_id>` for
# per-bank cost attribution. Downstream cost gateways (OpenRouter usage accounting,
# 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
@@ -1719,10 +1604,6 @@ class HindsightConfig:
# CachedContent prefix for its system prompt + response schema.
llm_prompt_cache_enabled: bool
# Opt-in diagnostic: log the exact request behind any LLM 4xx. Off by default;
# server-level only (not per-bank overridable). See engine/providers/llm_debug.py.
llm_debug_dump_4xx: bool
# Built-in llama.cpp configuration (for provider=llamacpp)
llamacpp_model_path: str | None # Path to GGUF file (None = auto-download default)
llamacpp_gpu_layers: int # -1 = all layers on GPU, 0 = CPU only
@@ -1804,7 +1685,6 @@ 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
@@ -1816,7 +1696,6 @@ 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
@@ -2017,7 +1896,6 @@ 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)
@@ -2030,15 +1908,12 @@ 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
@@ -2049,19 +1924,11 @@ class HindsightConfig:
metrics_include_bank_id: bool
metrics_backlog_enabled: bool
# Audit log configuration
# audit_log_enabled is hierarchical (env -> tenant -> bank): a deployment can
# audit some banks and not others. The actions allowlist and retention window
# stay static (server-level): retention is a global sweep with no bank scope.
audit_log_enabled: bool # Whether audit logging is on (overridable per bank)
# Audit log configuration (static - server-level only)
audit_log_enabled: bool # Master switch for audit logging
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)
@@ -2112,7 +1979,6 @@ 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
@@ -2170,9 +2036,6 @@ class HindsightConfig:
_CONFIGURABLE_FIELDS = {
# MCP tool access control
"mcp_enabled_tools",
# Audit logging on/off, per bank. The actions allowlist and retention
# window remain server-level and are deliberately not configurable.
"audit_log_enabled",
# Retention settings (behavioral)
"retain_chunk_size",
"retain_structured_chunk_size",
@@ -2316,9 +2179,6 @@ 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:
@@ -2408,13 +2268,6 @@ 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."""
@@ -2464,15 +2317,8 @@ class HindsightConfig:
llm_extra_body=json.loads(os.getenv(ENV_LLM_EXTRA_BODY, "null")),
llm_default_headers=json.loads(os.getenv(ENV_LLM_DEFAULT_HEADERS, "null")),
llm_strict_schema=os.getenv(ENV_LLM_STRICT_SCHEMA, str(DEFAULT_LLM_STRICT_SCHEMA)).lower() in ("true", "1"),
llm_strict_schema_retain=_resolve_operation_strict_schema(ENV_LLM_STRICT_SCHEMA_RETAIN),
llm_strict_schema_reflect=_resolve_operation_strict_schema(ENV_LLM_STRICT_SCHEMA_REFLECT),
llm_strict_schema_consolidation=_resolve_operation_strict_schema(ENV_LLM_STRICT_SCHEMA_CONSOLIDATION),
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
),
@@ -2497,8 +2343,6 @@ class HindsightConfig:
ENV_LLM_PROMPT_CACHE_ENABLED, str(DEFAULT_LLM_PROMPT_CACHE_ENABLED)
).lower()
in ("1", "true", "yes", "on"),
llm_debug_dump_4xx=os.getenv(ENV_LLM_DEBUG_DUMP_4XX, str(DEFAULT_LLM_DEBUG_DUMP_4XX)).lower()
in ("1", "true", "yes", "on"),
# Built-in llama.cpp configuration
llamacpp_model_path=os.getenv(ENV_LLAMACPP_MODEL_PATH) or None,
llamacpp_gpu_layers=int(os.getenv(ENV_LLAMACPP_GPU_LAYERS, str(DEFAULT_LLAMACPP_GPU_LAYERS))),
@@ -2711,9 +2555,6 @@ 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),
@@ -2735,11 +2576,6 @@ 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)
@@ -2772,11 +2608,6 @@ 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))
),
@@ -3071,10 +2902,6 @@ 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",
@@ -3097,23 +2924,9 @@ 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))
),
@@ -3176,11 +2989,6 @@ 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=[
@@ -10,7 +10,7 @@ import asyncio
import json
import logging
import uuid
from collections.abc import Awaitable, Callable
from collections.abc import Callable
from contextlib import asynccontextmanager
from dataclasses import dataclass, field
from datetime import datetime, timezone
@@ -19,7 +19,6 @@ from typing import Any
from pydantic import BaseModel, Field
from ..engine.db_utils import acquire_with_retry
from ..models import RequestContext
logger = logging.getLogger(__name__)
@@ -120,60 +119,23 @@ class AuditLogger:
schema_getter: Callable[[], str],
enabled: bool,
allowed_actions: list[str],
bank_enabled_resolver: Callable[[str, RequestContext | None], Awaitable[bool]] | None = None,
) -> None:
self._pool_getter = pool_getter
self._schema_getter = schema_getter
self._enabled = enabled
self._allowed_actions: frozenset[str] | None = frozenset(allowed_actions) if allowed_actions else None
# Resolves the hierarchical ``audit_log_enabled`` for one bank
# (env -> tenant -> bank). None means "no per-bank resolution wired",
# in which case the global value alone decides.
self._bank_enabled_resolver = bank_enabled_resolver
def action_allowed(self, action: str) -> bool:
"""Global action-allowlist check. Cheap, synchronous, bank-independent.
The allowlist is deployment-wide, so this is a valid pre-filter to skip
work for actions that can never be audited. It deliberately does NOT
consult the enabled flag: that is per-bank overridable, so a bank may
turn auditing ON even when the deployment default is off.
"""
if self._allowed_actions is None:
return True
return action in self._allowed_actions
async def should_log(self, action: str, bank_id: str | None, context: RequestContext | None = None) -> bool:
"""Full audit decision: action allowlist AND the bank's resolved switch.
``audit_log_enabled`` is hierarchical (env -> tenant -> bank), so the
effective value depends on which bank the action targets. Falls back to
the global value when there is no bank in scope or no resolver wired.
"""
if not self.action_allowed(action):
def is_enabled(self, action: str) -> bool:
"""Check if audit logging is enabled for this action."""
if not self._enabled:
return False
if bank_id is None or self._bank_enabled_resolver is None:
return self._enabled
try:
return await self._bank_enabled_resolver(bank_id, context)
except Exception as e:
# Never let a config-resolution failure break the request. Fall back
# to the deployment default: a transient DB blip must not silently
# create an audit gap for a bank meant to be audited. The tradeoff is
# the opt-out direction — a bank that overrode to false under a
# default-on deployment will be audited during the outage. We accept
# that: a few extra audit rows during a DB blip is the safer failure
# than dropping records that compliance may require.
logger.warning(f"Audit config resolution failed for bank={bank_id}: {e}; using global default")
return self._enabled
if self._allowed_actions is not None:
return action in self._allowed_actions
return True
def log_fire_and_forget(self, entry: AuditEntry) -> None:
"""Schedule an audit write as a background task.
Assumes the caller already made the audit decision via ``should_log``;
only the bank-independent allowlist is re-checked here.
"""
if not self.action_allowed(entry.action):
"""Schedule an audit write as a background task."""
if not self.is_enabled(entry.action):
return
try:
asyncio.create_task(self._safe_log(entry))
@@ -220,7 +182,6 @@ async def audit_context(
bank_id: str | None = None,
request: dict[str, Any] | None = None,
metadata: dict[str, Any] | None = None,
context: RequestContext | None = None,
):
"""Async context manager that times the operation and writes audit on exit.
@@ -229,7 +190,7 @@ async def audit_context(
result = await do_work()
entry.response = result_dict
"""
if audit_logger is None or not await audit_logger.should_log(action, bank_id, context):
if audit_logger is None or not audit_logger.is_enabled(action):
entry = AuditEntry(action=action, transport=transport, bank_id=bank_id)
yield entry
return
@@ -13,8 +13,6 @@ 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.
@@ -34,14 +32,3 @@ 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 {}
@@ -1,13 +0,0 @@
"""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,11 +109,6 @@ 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
@@ -131,21 +126,11 @@ 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 | None:
def add_years(base_date: datetime, years: int) -> datetime:
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
@@ -453,11 +438,6 @@ 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:
@@ -470,7 +450,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 | None:
def relative_offset_datetime(amount: int, unit: str, direction: int) -> datetime:
if unit in ("", ""):
return reference_date + timedelta(days=direction * amount)
if unit in ("", "星期", "礼拜"):
@@ -479,15 +459,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 | NoTemporalConstraintSentinel:
def point_constraint_at_offset(amount: int, unit: str, direction: int) -> DateRange:
d = relative_offset_datetime(amount, unit, direction)
return safe_constraint(d, d)
return constraint(d, d)
def window_to_reference(amount: int, unit: str) -> DateRange | NoTemporalConstraintSentinel:
return safe_constraint(relative_offset_datetime(amount, unit, -1), reference_date)
def window_to_reference(amount: int, unit: str) -> DateRange:
return constraint(relative_offset_datetime(amount, unit, -1), reference_date)
def window_from_reference(amount: int, unit: str) -> DateRange | NoTemporalConstraintSentinel:
return safe_constraint(reference_date, relative_offset_datetime(amount, unit, 1))
def window_from_reference(amount: int, unit: str) -> DateRange:
return constraint(reference_date, relative_offset_datetime(amount, unit, 1))
# Chinese rule guide
#
@@ -801,8 +781,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 = add_days(base, fixed_day_offset(relative_year_fixed_day_since_match.group(2)))
return safe_since_constraint(d)
d = base + timedelta(days=fixed_day_offset(relative_year_fixed_day_since_match.group(2)))
return since_constraint(d)
fixed_day_since_match = chinese_search(
rf"(大大后天|大后天|后天|明天|明日|今天|今日|本日|当日|当天|昨天|昨日|大大前天|大前天|前天)"
@@ -819,7 +799,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 safe_since_constraint(relative_offset_datetime(amount, unit, -1))
return since_constraint(relative_offset_datetime(amount, unit, -1))
weekend_since_match = chinese_search(
rf"(?<![上下大小每个各隔])"
@@ -919,8 +899,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 = add_days(base, daypart_day_offset(relative_year_daypart_since_match.group(2)))
return safe_since_constraint(d)
d = base + timedelta(days=daypart_day_offset(relative_year_daypart_since_match.group(2)))
return since_constraint(d)
daypart_since_match = chinese_search(
rf"(昨晚|昨夜|前晚|前夜|今晚|今早|今晨|明早|明晚|明夜){chinese_since_suffix_pattern}"
@@ -935,17 +915,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 = add_days(base, daypart_day_offset(relative_year_daypart_match.group(2)))
return safe_constraint(d, d)
d = base + timedelta(days=daypart_day_offset(relative_year_daypart_match.group(2)))
return constraint(d, d)
# Day-part abbreviations still resolve only to date granularity.
if chinese_search(r"昨晚|昨夜"):
d = add_days(reference_date, daypart_day_offset("昨晚"))
return safe_constraint(d, d)
d = reference_date + timedelta(days=daypart_day_offset("昨晚"))
return constraint(d, d)
if chinese_search(r"前晚|前夜"):
d = add_days(reference_date, daypart_day_offset("前晚"))
return safe_constraint(d, d)
d = reference_date + timedelta(days=daypart_day_offset("前晚"))
return constraint(d, d)
if chinese_search(r"今晚|今早|今晨"):
return constraint(reference_date, reference_date)
@@ -961,8 +941,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 = add_days(base, fixed_day_offset(relative_year_fixed_day_match.group(2)))
return safe_constraint(d, d)
d = base + timedelta(days=fixed_day_offset(relative_year_fixed_day_match.group(2)))
return constraint(d, d)
if chinese_search(r"昨天|昨日"):
d = reference_date - timedelta(days=1)
@@ -1105,7 +1085,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 safe_constraint(
return constraint(
relative_offset_datetime(start_amount, unit, 1),
relative_offset_datetime(end_amount, unit, 1),
)
@@ -1113,7 +1093,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 safe_constraint(relative_offset_datetime(2, unit, 1), relative_offset_datetime(5, unit, 1))
return 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}"
@@ -1133,7 +1113,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 safe_constraint(
return constraint(
relative_offset_datetime(second_amount, unit, -1),
relative_offset_datetime(first_amount, unit, -1),
)
@@ -1164,7 +1144,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 safe_constraint(add_years(reference_date, -3), add_years(reference_date, -1))
return constraint(add_years(reference_date, -3), add_years(reference_date, -1))
rolling_this_adjacent_match = chinese_search(
r"这(一两|[两二]三|三两|三四|四五|五六|六七|七八|八九|九十)个?(天|日|周|星期|礼拜|月|年)"
@@ -1174,7 +1154,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 safe_constraint(relative_offset_datetime(end_amount, unit, -1), reference_date)
return 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:
@@ -1211,7 +1191,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 safe_constraint(relative_offset_datetime(end_amount, unit, -1), reference_date)
return constraint(relative_offset_datetime(end_amount, unit, -1), reference_date)
rolling_past_few_match = chinese_search(r"(过去|近|最近)几个?(天|日|周|星期|礼拜|月|年)")
if rolling_past_few_match:
@@ -28,7 +28,6 @@ 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
@@ -103,29 +102,6 @@ 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:
@@ -133,20 +109,9 @@ stored, and it is highly similar to an EXISTING one:
[NEW] {new}
[EXISTING] {existing}
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."""
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"."""
def _dedup_active(config: Any) -> bool:
@@ -224,13 +189,10 @@ async def _dedup_adjudicate(
if best_id is None:
return _DedupOutcome(best_id=None, merged_text="", should_merge=False)
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",
strict_schema=get_config().llm_strict_schema_consolidation,
)
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",
)
if decision.action != "merge":
return _DedupOutcome(best_id=best_id, merged_text="", should_merge=False)
@@ -1804,22 +1766,15 @@ async def _append_observation_history(
history from growing without bound.
"""
obs_uuid = uuid.UUID(observation_id)
try:
await conn.execute(
f"""
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)),
)
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
obs_uuid,
bank_id,
json.dumps(asdict(snapshot)),
)
if max_entries and max_entries > 0:
await conn.execute(
f"""
@@ -2299,11 +2254,6 @@ async def _consolidate_batch_with_llm(
],
"response_format": response_model,
"scope": "consolidation",
# Resolved per operation (HINDSIGHT_API_LLM_STRICT_SCHEMA_CONSOLIDATION, falling
# back to the global flag) so an operator can grammar-enforce consolidation's
# structured output -- which narrows the raw-JSON failure mode behind #2668 --
# without forcing strict schema on operations whose model can't satisfy it.
"strict_schema": config.llm_strict_schema_consolidation,
}
# Only request an explicit output budget when configured. Left unset by default the key is
# omitted, so each provider keeps its implicit default (backwards compatible). Operators on
@@ -7,13 +7,11 @@ 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
@@ -47,7 +45,6 @@ from ..config import (
ENV_RERANKER_TEI_URL,
ENV_RERANKER_ZEROENTROPY_API_KEY,
)
from .bank_attribution import reranker_bank_attribution_headers
logger = logging.getLogger(__name__)
@@ -89,12 +86,6 @@ 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.
@@ -324,7 +315,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:
_release_rerank_heap()
_malloc_trim()
async def predict(self, pairs: list[tuple[str, str]]) -> list[float]:
"""
@@ -493,7 +484,6 @@ class RemoteTEICrossEncoder(CrossEncoderModel):
semaphore,
"POST",
f"{self.base_url}/rerank",
headers=reranker_bank_attribution_headers(),
json={
"query": query,
"texts": texts,
@@ -634,11 +624,7 @@ class _CohereCompatibleRerankClient:
if self.include_top_n:
body["top_n"] = len(texts)
response = await self._async_client.post(
self.rerank_url,
headers=reranker_bank_attribution_headers(),
json=body,
)
response = await self._async_client.post(self.rerank_url, json=body)
response.raise_for_status()
result = response.json()
@@ -1004,11 +990,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]]] = {}
@@ -1037,7 +1023,7 @@ class FlashRankCrossEncoder(CrossEncoderModel):
return all_scores
finally:
_release_rerank_heap()
_malloc_trim()
async def predict(self, pairs: list[tuple[str, str]]) -> list[float]:
"""
@@ -1165,7 +1151,6 @@ 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,
@@ -1284,11 +1269,10 @@ class LiteLLMSDKCrossEncoder(CrossEncoderModel):
indices = [idx for idx, _ in indexed_texts]
# Build kwargs for rerank call
rerank_kwargs: dict[str, Any] = {
rerank_kwargs = {
"model": self.model,
"query": query,
"documents": texts,
"headers": reranker_bank_attribution_headers(),
}
if self.api_key:
rerank_kwargs["api_key"] = self.api_key
@@ -1297,9 +1281,21 @@ class LiteLLMSDKCrossEncoder(CrossEncoderModel):
response = await self._litellm.arerank(**rerank_kwargs)
for result in response.results:
original_idx = result["index"]
all_scores[indices[original_idx]] = result["relevance_score"]
# 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)}")
return all_scores
@@ -307,17 +307,6 @@ 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,7 +18,6 @@ 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
@@ -173,25 +172,6 @@ class DataAccessOps(ABC):
"""
...
@abstractmethod
async def bulk_reassert_entities(
self,
conn: DatabaseConnection,
table: str,
bank_id: str,
entity_ids: list[str],
canonical_names: list[str],
) -> None:
"""Lock resolved parents and re-create any pruned since Phase-1 resolution.
Closes the retain Phase-1/prune race (#2662): existing rows are locked
(PG ``FOR KEY SHARE`` / Oracle ``FOR UPDATE``) so a concurrent
``prune_orphan_entities`` blocks until the caller's transaction commits,
while rows already deleted are re-inserted idempotently. ``entity_ids``
must be sorted by the caller for a stable lock order.
"""
...
@abstractmethod
async def bulk_insert_unit_entities(
self,
@@ -504,23 +484,6 @@ 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,8 +13,6 @@ 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."""
@@ -218,7 +216,7 @@ class OracleOps(DataAccessOps):
for orig_name in missing_names:
row = await conn.fetchrow(
f"""
SELECT id, canonical_name, LOWER(canonical_name) AS name_lower
SELECT id, LOWER(canonical_name) AS name_lower
FROM {table}
WHERE bank_id = $1 AND LOWER(canonical_name) = LOWER($2)
""",
@@ -226,37 +224,10 @@ class OracleOps(DataAccessOps):
orig_name,
)
if row:
# Wrap in a dict-like to include input_name for downstream compat
results.append(row)
return results
async def bulk_reassert_entities(
self,
conn: DatabaseConnection,
table: str,
bank_id: str,
entity_ids: list[str],
canonical_names: list[str],
) -> None:
# Oracle has no FOR KEY SHARE; FOR UPDATE is the row-lock equivalent that
# blocks a concurrent prune DELETE until this transaction commits. Lock
# each surviving parent in the caller's stable id order (pruned ids are
# simply absent here), then re-insert any that vanished. The translation
# layer rewrites ON CONFLICT DO NOTHING to strip-and-catch ORA-00001, so
# a name recreated under a new id is suppressed rather than raising.
for entity_id in entity_ids:
await conn.fetchrow(
f"SELECT id FROM {table} WHERE id = $1 FOR UPDATE",
entity_id,
)
await conn.executemany(
f"""
INSERT INTO {table} (id, bank_id, canonical_name)
VALUES ($1, $2, $3)
ON CONFLICT DO NOTHING
""",
[(entity_id, bank_id, canonical_name) for entity_id, canonical_name in zip(entity_ids, canonical_names)],
)
async def bulk_insert_unit_entities(
self,
conn: DatabaseConnection,
@@ -358,12 +329,6 @@ 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}
@@ -482,14 +447,6 @@ 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
@@ -502,6 +459,7 @@ 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
)"""
@@ -866,157 +824,6 @@ 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,40 +4,11 @@ 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."""
@@ -122,39 +93,101 @@ class PostgreSQLOps(DataAccessOps):
config = get_config()
table = self._get_mu_table()
# 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
"""
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
"""
results = await conn.fetch(
query,
@@ -277,7 +310,7 @@ class PostgreSQLOps(DataAccessOps):
) -> list[ResultRow]:
return await conn.fetch(
f"""
SELECT e.id, e.canonical_name, LOWER(e.canonical_name) AS name_lower, inputs.input_name
SELECT e.id, LOWER(e.canonical_name) AS name_lower, inputs.input_name
FROM {table} e
JOIN (
SELECT LOWER(n) AS input_name_lower, n AS input_name
@@ -289,42 +322,6 @@ class PostgreSQLOps(DataAccessOps):
missing_names,
)
async def bulk_reassert_entities(
self,
conn: DatabaseConnection,
table: str,
bank_id: str,
entity_ids: list[str],
canonical_names: list[str],
) -> None:
# One statement, one round-trip (same shape as bulk_insert_links):
# * the CTE takes FOR KEY SHARE on every parent that still exists,
# held to COMMIT, so a concurrent prune_orphan_entities DELETE blocks
# until the caller's unit_entities insert has committed;
# * the INSERT re-creates only the parents that were already pruned
# (NOT IN locked), carrying the canonical_name resolved in Phase 1.
# ON CONFLICT DO NOTHING (no target) keeps the rare case where another
# worker recreated the name under a new id from raising — that row stays
# absent and its unit link is the sole casualty, never the whole batch.
await conn.execute(
f"""
WITH locked AS (
SELECT id FROM {table}
WHERE id = ANY($2::uuid[])
ORDER BY id
FOR KEY SHARE
)
INSERT INTO {table} (id, bank_id, canonical_name)
SELECT t.entity_id, $1, t.canonical_name
FROM unnest($2::uuid[], $3::text[]) AS t(entity_id, canonical_name)
WHERE t.entity_id NOT IN (SELECT id FROM locked)
ON CONFLICT DO NOTHING
""",
bank_id,
entity_ids,
canonical_names,
)
async def bulk_insert_unit_entities(
self,
conn: DatabaseConnection,
@@ -428,48 +425,19 @@ 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 victims v
WHERE c.entity_id_1 = v.entity_id_1
AND c.entity_id_2 = v.entity_id_2
USING {entities_table} e
WHERE e.id = c.entity_id_1
AND e.bank_id = $1
AND NOT EXISTS (
SELECT 1
FROM {ue_table} u1
JOIN {ue_table} u2 ON u1.unit_id = u2.unit_id
WHERE u1.entity_id = c.entity_id_1
AND u2.entity_id = c.entity_id_2
)
""",
bank_id,
)
@@ -573,18 +541,11 @@ 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
@@ -933,93 +894,6 @@ 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,10 +1242,6 @@ 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,
@@ -1281,17 +1277,11 @@ class OracleBackend(DatabaseBackend):
logger.info(f"Oracle pool created (min={min_size}, max={max_size})")
async def shutdown(self) -> 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)
if self._pool is not None:
await self._pool.close(force=True)
self._pool = None
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.
@@ -1304,23 +1294,10 @@ class OracleBackend(DatabaseBackend):
from ..memory_engine import get_current_schema
schema = get_current_schema()
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()
if schema and schema != "public":
cursor = conn.cursor()
await cursor.execute(f'ALTER SESSION SET CURRENT_SCHEMA = "{schema}"')
await cursor.close()
@asynccontextmanager
async def acquire(self) -> AsyncIterator[OracleConnection]:
@@ -95,12 +95,7 @@ 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}, "
@@ -108,19 +103,11 @@ class PostgreSQLBackend(DatabaseBackend):
)
async def shutdown(self) -> 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()
if self._pool is not None:
await self._pool.close()
self._pool = None
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,7 +4,6 @@ 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
@@ -17,20 +16,6 @@ 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(
{
@@ -93,7 +78,7 @@ async def retry_with_backoff(
raise
last_exception = e
if attempt < max_retries:
delay = _backoff_delay(attempt, base_delay, max_delay)
delay = min(base_delay * (2**attempt), max_delay)
if type(e).__name__ == "DeadlockDetectedError" or _is_oracle_deadlock(e):
logger.warning(
"Deadlock detected during parallel document processing — "
@@ -151,7 +136,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 = _backoff_delay(attempt, DEFAULT_BASE_DELAY, DEFAULT_MAX_DELAY)
delay = min(DEFAULT_BASE_DELAY * (2**attempt), DEFAULT_MAX_DELAY)
logger.warning(
f"Database acquire failed (attempt {attempt + 1}/{max_retries + 1}): {e}. "
f"Retrying in {delay:.1f}s..."
@@ -76,25 +76,6 @@ 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.
@@ -1221,7 +1202,6 @@ 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.
@@ -1236,10 +1216,6 @@ 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
@@ -1248,7 +1224,6 @@ 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
@@ -1325,33 +1300,6 @@ 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
@@ -1743,7 +1691,6 @@ 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,11 +9,10 @@ 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
from typing import Any, Final, cast
from typing import Any, Final
from .db_utils import acquire_with_retry
from .memory_engine import fq_table
@@ -26,7 +25,6 @@ from .retain.entity_labels import (
from .retain.entity_labels import (
parse_entity_labels as _parse_entity_labels,
)
from .retain.types import ResolvedEntity
logger = logging.getLogger(__name__)
@@ -77,22 +75,6 @@ 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)."""
@@ -248,7 +230,7 @@ class EntityResolver:
unit_event_date,
conn=None,
entity_labels: list | None = None,
) -> list[ResolvedEntity]:
) -> list[str]:
"""
Resolve multiple entities in batch (MUCH faster than sequential).
@@ -263,8 +245,7 @@ class EntityResolver:
conn: Optional connection to use (if None, acquires from pool)
Returns:
Resolved entity identities (id + stored canonical name) in the same
order as input.
List of entity IDs in same order as input
"""
if not entities_data:
return []
@@ -290,7 +271,7 @@ class EntityResolver:
unit_event_date,
taxonomy_lookup: set[str] | None = None,
labels_cfg=None,
) -> list[ResolvedEntity]:
) -> list[str]:
if self.entity_lookup == "trigram":
# Route to backend-specific fuzzy strategy.
# Non-PG backends (Oracle) use UTL_MATCH instead of pg_trgm.
@@ -330,7 +311,7 @@ class EntityResolver:
unit_event_date,
taxonomy_lookup: set[str] | None = None,
labels_cfg=None,
) -> list[ResolvedEntity]:
) -> list[str]:
"""Original strategy: load all bank entities then match in Python."""
# Query ALL candidates for this bank
all_entities = await conn.fetch(
@@ -414,7 +395,7 @@ class EntityResolver:
unit_event_date,
taxonomy_lookup: set[str] | None = None,
labels_cfg=None,
) -> list[ResolvedEntity]:
) -> list[str]:
"""
Trigram strategy: fetch only similar candidates per entity name using pg_trgm.
@@ -518,7 +499,7 @@ class EntityResolver:
unit_event_date: datetime | None,
taxonomy_lookup: set[str] | None = None,
labels_cfg=None,
) -> list[ResolvedEntity]:
) -> list[str]:
"""
Oracle strategy: fetch similar candidates using UTL_MATCH.JARO_WINKLER_SIMILARITY.
@@ -626,14 +607,11 @@ class EntityResolver:
cooccurrence_map: dict[str, set[str]],
taxonomy_lookup: set[str] | None = None,
labels_cfg=None,
) -> list[ResolvedEntity]:
) -> list[str]:
"""Shared scoring + upsert logic used by both lookup strategies."""
# Resolve each entity using pre-fetched candidates. A slot stays None
# only if find-or-create fails to produce a row for a mention (a DB
# inconsistency); it surfaces as a clear error at the reassert boundary
# rather than a silent NOT NULL violation deeper in Phase 2.
resolved: list[ResolvedEntity | None] = [None] * len(entities_data)
# Resolve each entity using pre-fetched candidates
entity_ids = [None] * len(entities_data)
entities_to_update: list[_EntityStat] = []
entities_to_create: list[_EntityToCreate] = []
@@ -660,23 +638,21 @@ class EntityResolver:
if is_label:
# Exact case-insensitive match only for label entities
exact_match: ResolvedEntity | None = None
exact_match = None
entity_text_lower = entity_text.lower()
for candidate_id, canonical_name, metadata, last_seen, mention_count in candidates:
if canonical_name.lower() == entity_text_lower:
exact_match = ResolvedEntity(entity_id=candidate_id, canonical_name=canonical_name)
exact_match = candidate_id
break
if exact_match:
resolved[idx] = exact_match
entities_to_update.append(
_EntityStat(entity_id=exact_match.entity_id, event_date=entity_event_date)
)
entity_ids[idx] = exact_match
entities_to_update.append(_EntityStat(entity_id=exact_match, event_date=entity_event_date))
else:
entities_to_create.append(_EntityToCreate(idx=idx, name=entity_text, event_date=entity_event_date))
continue
# Score candidates
best_candidate: ResolvedEntity | None = None
best_candidate = None
best_score = 0.0
nearby_entity_set = {e["text"].lower() for e in nearby_entities if e["text"] != entity_text}
@@ -709,14 +685,14 @@ class EntityResolver:
if score > best_score:
best_score = score
best_candidate = ResolvedEntity(entity_id=candidate_id, canonical_name=canonical_name)
best_candidate = candidate_id
# Apply unified threshold
threshold = 0.6
if best_score > threshold and best_candidate is not None:
resolved[idx] = best_candidate
entities_to_update.append(_EntityStat(entity_id=best_candidate.entity_id, event_date=entity_event_date))
if best_score > threshold:
entity_ids[idx] = best_candidate
entities_to_update.append(_EntityStat(entity_id=best_candidate, event_date=entity_event_date))
else:
entities_to_create.append(
_EntityToCreate(idx=idx, name=entity_data["text"], event_date=entity_event_date)
@@ -749,9 +725,6 @@ class EntityResolver:
sorted_groups = sorted(groups.items())
entity_names = [g.name for _, g in sorted_groups]
entity_dates = [g.event_date for _, g in sorted_groups]
# Stored canonical name per lowercase key, so a resurrected parent
# keeps the name it was created/matched with rather than a fallback.
canonical_by_name = {name_lower: g.name for name_lower, g in sorted_groups}
# INSERT ... ON CONFLICT DO NOTHING — no row lock on already-existing entities.
# mention_count starts at 0 here; flush_pending_stats() is the sole source of
@@ -786,14 +759,11 @@ class EntityResolver:
)
for row in existing_rows:
id_by_name[row["name_lower"]] = row["id"]
canonical_by_name[row["name_lower"]] = row["canonical_name"]
# Also index by Python's lower() of the original input name so the
# assignment loop (which uses Python-lowercased keys) finds it even
# when Python and the database produce different lowercase strings.
if "input_name" in row:
input_name_lower = row["input_name"].lower()
id_by_name[input_name_lower] = row["id"]
canonical_by_name[input_name_lower] = row["canonical_name"]
id_by_name[row["input_name"].lower()] = row["id"]
# Assign entity IDs back and queue one stat per original mention so that
# flush_pending_stats() increments mention_count by the true mention count,
@@ -801,64 +771,16 @@ class EntityResolver:
for name_lower, g in sorted_groups:
entity_id = id_by_name.get(name_lower)
if entity_id:
canonical_name = canonical_by_name.get(name_lower, g.name)
for original_idx in g.indices:
resolved[original_idx] = ResolvedEntity(entity_id=entity_id, canonical_name=canonical_name)
pending.append(_EntityStat(entity_id=str(entity_id), event_date=g.event_date))
entity_ids[original_idx] = entity_id
pending.append(_EntityStat(entity_id=entity_id, event_date=g.event_date))
# Accumulate into the resolver's pending list; the orchestrator flushes
# these with await entity_resolver.flush_pending_stats() after the txn.
key = self._task_key()
self._pending_stats.setdefault(key, []).extend(pending)
missing = [i for i, entity in enumerate(resolved) if entity is None]
if missing:
raise RuntimeError(
f"Entity resolution produced no row for {len(missing)} mention(s) "
f"(indices {missing[:5]}); refusing to link units to a missing parent."
)
return cast(list[ResolvedEntity], resolved)
async def reassert_entities_batch(
self,
bank_id: str,
resolved_entities: list[ResolvedEntity],
conn,
) -> None:
"""Lock (and, if pruned, re-create) resolved parents before linking units.
Phase-1 resolution and the Phase-2 ``unit_entities`` insert run on
different transactions. In the gap, ``prune_orphan_entities`` can delete
a just-resolved parent — it legitimately has no ``unit_entities`` row
yet — and the Phase-2 FK insert then fails, dropping the whole batch as
non-retryable (silent memory loss, #2662).
Called on the Phase-2 connection immediately before
``link_units_to_entities_batch``, this locks the parents that still
exist (so the pruner blocks until we commit) and re-inserts any that
already vanished, in one round-trip. An entity referenced by a live unit
is by definition not an orphan, so resurrecting it is correct.
"""
# Deduplicate by id and lock in a stable order so concurrent reasserts
# acquire row locks consistently (same convention as bulk_insert_links).
seen: set[str] = set()
unique: list[ResolvedEntity] = []
for entity in sorted(resolved_entities, key=lambda e: e.entity_id):
if entity.entity_id in seen:
continue
seen.add(entity.entity_id)
unique.append(entity)
if not unique:
return
await self._ops.bulk_reassert_entities(
conn,
fq_table("entities"),
bank_id,
[entity.entity_id for entity in unique],
[entity.canonical_name for entity in unique],
)
return entity_ids
async def link_units_to_entities_batch(
self,
@@ -931,12 +853,20 @@ 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 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)
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)
# Accumulate co-occurrence pairs for post-transaction flush.
# The actual INSERT/UPDATE is deferred to flush_pending_stats() to avoid
@@ -66,20 +66,6 @@ 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:
@@ -217,51 +203,27 @@ 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 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
async with acquire_with_retry(backend) as conn:
async with conn.transaction():
result.orphan_entities_pruned = await ops.prune_orphan_entities(
conn,
fq_table("entities"),
fq_table("unit_entities"),
bank_id,
)
# The orphan prune above cascades cooccurrences via FK. The
# explicit cooccurrence pass below catches the *stale-count*
# case: both entities still exist but no current unit witnesses
# them together.
result.stale_cooccurrences_pruned = await ops.prune_stale_cooccurrences(
conn,
fq_table("entity_cooccurrences"),
fq_table("unit_entities"),
fq_table("entities"),
bank_id,
)
elapsed = time.time() - job_start
logger.info(
@@ -565,30 +565,6 @@ class MemoryEngineInterface(ABC):
"""
...
@abstractmethod
async def delete_operation(
self,
bank_id: str,
operation_id: str,
*,
request_context: "RequestContext",
) -> dict[str, Any]:
"""
Delete a terminal async operation record.
Args:
bank_id: The memory bank ID.
operation_id: The operation ID to delete.
request_context: Request context for authentication.
Returns:
Dict with success status and message.
Raises:
ValueError: If operation not found.
"""
...
@abstractmethod
async def update_bank(
self,
@@ -6,53 +6,12 @@ enabling support for multiple LLM backends (OpenAI, Anthropic, Gemini, Codex, et
"""
from abc import ABC, abstractmethod
from dataclasses import dataclass
from datetime import datetime
from enum import StrEnum
from typing import Any, Self
from typing import Any
from .response_models import LLMToolCallResult
class LLMToolChoiceMode(StrEnum):
"""Canonical tool-selection modes shared by every LLM provider."""
AUTO = "auto"
NONE = "none"
REQUIRED = "required"
NAMED = "named"
@dataclass(frozen=True, slots=True)
class LLMToolChoice:
"""Typed internal tool selection serialized only at provider boundaries."""
mode: LLMToolChoiceMode
function_name: str | None = None
def __post_init__(self) -> None:
if self.mode is LLMToolChoiceMode.NAMED:
if self.function_name is None or not self.function_name or self.function_name != self.function_name.strip():
raise ValueError("Named tool choice requires a non-empty canonical function name")
elif self.function_name is not None:
raise ValueError(f"Tool choice mode {self.mode.value!r} cannot include a function name")
@classmethod
def named(cls, function_name: str) -> Self:
return cls(mode=LLMToolChoiceMode.NAMED, function_name=function_name)
@property
def selected_function_name(self) -> str:
if self.function_name is None:
raise ValueError("Tool choice does not select a named function")
return self.function_name
LLM_TOOL_CHOICE_AUTO = LLMToolChoice(mode=LLMToolChoiceMode.AUTO)
LLM_TOOL_CHOICE_NONE = LLMToolChoice(mode=LLMToolChoiceMode.NONE)
LLM_TOOL_CHOICE_REQUIRED = LLMToolChoice(mode=LLMToolChoiceMode.REQUIRED)
class LLMInterface(ABC):
"""
Abstract interface for LLM providers.
@@ -155,9 +114,8 @@ class LLMInterface(ABC):
max_retries: int = 5,
initial_backoff: float = 1.0,
max_backoff: float = 30.0,
tool_choice: LLMToolChoice = LLM_TOOL_CHOICE_AUTO,
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.
@@ -171,7 +129,7 @@ class LLMInterface(ABC):
max_retries: Maximum retry attempts.
initial_backoff: Initial backoff time in seconds.
max_backoff: Maximum backoff time in seconds.
tool_choice: Canonical tool-selection policy.
tool_choice: How to choose tools - "auto", "none", "required", or specific function.
Returns:
LLMToolCallResult with content and/or tool_calls.
@@ -227,45 +185,6 @@ 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,26 +376,6 @@ 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:
@@ -493,7 +473,7 @@ class LLMTraceRecorder:
async def _safe_write(self, record: LLMRequestRecord) -> None:
"""Write a trace row. Errors are logged, never raised."""
pool = self._writable()
pool = self._pool_getter()
if pool is None:
logger.debug("LLM trace skipped: pool not available")
return
@@ -588,9 +568,8 @@ 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._writable()
pool = self._pool_getter()
if pool is None:
logger.debug("LLM trace memory_id attach skipped: pool not available")
return
try:
schema = self._schema_getter()
@@ -12,8 +12,6 @@ import uuid
from contextlib import AsyncExitStack
from typing import TYPE_CHECKING, Any
from json_repair import repair_json
# Vertex AI imports (conditional - for LLMProvider to pass credentials to GeminiLLM)
try:
from google.oauth2 import service_account
@@ -29,11 +27,13 @@ from ..config import (
ENV_REFLECT_LLM_MAX_CONCURRENT,
ENV_RETAIN_LLM_MAX_CONCURRENT,
)
from .llm_interface import LLM_TOOL_CHOICE_AUTO, LLMToolChoice, LLMToolChoiceMode
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
@@ -113,7 +113,7 @@ def _request_params(
temperature: float | None = None,
scope: str | None = None,
response_format: Any | None = None,
tool_choice: LLMToolChoice | None = None,
tool_choice: str | dict[str, Any] | None = None,
) -> dict[str, Any] | None:
"""Build the requested-params bag for tracing — only values the caller set.
@@ -128,8 +128,8 @@ def _request_params(
params["temperature"] = temperature
if response_format is not None:
params["response_schema"] = getattr(response_format, "__name__", None) or "structured"
if tool_choice is not None and tool_choice.mode is not LLMToolChoiceMode.AUTO:
params["tool_choice"] = tool_choice.function_name or tool_choice.mode.value
if tool_choice is not None and tool_choice != "auto":
params["tool_choice"] = tool_choice if isinstance(tool_choice, str) else "named"
return params or None
@@ -184,14 +184,6 @@ def parse_llm_json(raw: str) -> Any:
1. Markdown code fences (```json ... ```) strip them before parsing.
2. Embedded control characters (\\x00-\\x1f, \\x7f) replace with space
and retry if the initial parse fails.
3. Structural malformation (trailing commas, unterminated strings, single
quotes, invalid ``\\escape`` sequences) repaired as a last resort via
``json_repair`` (#2547/#2544).
The repair pass is purely *structural*: it fixes JSON that ``json.loads``
cannot parse at all. It deliberately does NOT touch content semantics
degenerate-but-valid JSON (repetition loops or leaked scaffolding inside
string values) parses fine here and is out of scope for this helper.
Args:
raw: Raw text returned by the LLM.
@@ -200,8 +192,7 @@ def parse_llm_json(raw: str) -> Any:
Parsed Python object (dict, list, etc.).
Raises:
json.JSONDecodeError: If the text cannot be parsed even after cleanup
and structural repair (e.g. repair yields an empty result).
json.JSONDecodeError: If the text cannot be parsed even after cleanup.
"""
text = raw.strip()
@@ -218,19 +209,7 @@ def parse_llm_json(raw: str) -> Any:
# Some models (e.g. Gemini) embed raw control characters inside JSON
# string values. Replacing them with a space usually produces valid JSON.
cleaned = re.sub(r"[\x00-\x1f\x7f]", " ", text)
try:
return json.loads(cleaned)
except json.JSONDecodeError:
# Last resort: structural repair of malformed JSON. ``repair_json`` never
# raises — unrecoverable input yields an empty result ("" / {} / []). Keep
# failing loudly in that case rather than let an empty object masquerade
# as a successful parse: callers (retry ladders, the #1833 fail-loud path)
# rely on JSONDecodeError to retry or surface the failure.
repaired = repair_json(cleaned, return_objects=True)
if not repaired:
raise
return repaired
_PROVIDERS_WITHOUT_API_KEY = frozenset(
@@ -256,17 +235,6 @@ 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,
@@ -286,7 +254,6 @@ 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.
@@ -301,8 +268,6 @@ 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
@@ -326,8 +291,6 @@ 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,
@@ -531,7 +494,6 @@ 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,
)
@@ -569,7 +531,6 @@ 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.
@@ -584,8 +545,6 @@ 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).
@@ -639,7 +598,6 @@ 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
@@ -784,7 +742,6 @@ 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,
)
@@ -846,7 +803,7 @@ class LLMProvider:
initial_backoff: float | None = None,
max_backoff: float | None = None,
skip_validation: bool = False,
strict_schema: bool | None = None,
strict_schema: bool = False,
return_usage: bool = False,
cached_prefix: str | None = None,
) -> Any:
@@ -867,10 +824,9 @@ class LLMProvider:
configured default (``llm_max_backoff``), else 60.0.
skip_validation: Return raw JSON without Pydantic validation.
strict_schema: Per-call override requesting grammar-enforced (json_schema strict)
structured output instead of the soft json_object path. None (the default)
inherits the server-level HINDSIGHT_API_LLM_STRICT_SCHEMA flag; an explicit
True or False wins over it, so a caller can force strict output on -- or off --
for its own scope. Providers without a strict mode ignore it.
structured output instead of the soft json_object path. The server-level
HINDSIGHT_API_LLM_STRICT_SCHEMA flag is OR-ed in here so it applies to every call;
providers without a strict mode ignore it.
return_usage: If True, return tuple (result, TokenUsage) instead of just result.
Returns:
@@ -905,18 +861,14 @@ class LLMProvider:
)
# Resolve strict-schema once, here, rather than in each provider: the
# per-call argument, falling back to the server-level
# HINDSIGHT_API_LLM_STRICT_SCHEMA flag when the caller expressed no
# preference. Providers with a json_schema response_format (OpenAI-compatible,
# per-call argument OR the server-level HINDSIGHT_API_LLM_STRICT_SCHEMA
# flag. Providers with a json_schema response_format (OpenAI-compatible,
# LiteLLM) then grammar-enforce structured output instead of the fragile
# soft json_object path; Gemini already enforces its native response_schema,
# and providers without a strict mode simply ignore the flag.
from ..config import get_config
# An explicit per-call value wins in BOTH directions -- `or` would have made a
# per-call False indistinguishable from "unset", silently ignoring any caller
# that opts out while the global flag is on.
strict_schema = strict_schema if strict_schema is not None else get_config().llm_strict_schema
strict_schema = strict_schema or get_config().llm_strict_schema
# LLM call observability flows through the OTel GenAI recorder
# (tracing.get_span_recorder().record_llm_call). Provider implementations
@@ -1013,9 +965,8 @@ class LLMProvider:
max_retries: int | None = None,
initial_backoff: float | None = None,
max_backoff: float | None = None,
tool_choice: LLMToolChoice = LLM_TOOL_CHOICE_AUTO,
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.
@@ -1032,7 +983,7 @@ class LLMProvider:
configured default (``llm_initial_backoff``), else 1.0.
max_backoff: Maximum backoff time in seconds. ``None`` uses the provider's
configured default (``llm_max_backoff``), else 30.0.
tool_choice: Canonical tool-selection policy.
tool_choice: How to choose tools - "auto", "none", "required", or {"type": "function", "function": {"name": "..."}}
Returns:
LLMToolCallResult with content and/or tool_calls.
@@ -1083,14 +1034,9 @@ 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() / 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 {}
)
# 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 {}
try:
# Delegate to provider implementation
result = await self._provider_impl.call_with_tools(
@@ -1314,7 +1260,6 @@ 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,
@@ -1325,7 +1270,6 @@ 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,
)
@@ -1370,7 +1314,6 @@ 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,10 +21,9 @@ 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 (``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.
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.
"""
from __future__ import annotations
@@ -33,13 +32,13 @@ import asyncio
import logging
import time
from collections.abc import Coroutine
from datetime import datetime, timedelta, timezone
from datetime import datetime, 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_routine, fq_table, fq_table_explicit
from .schema import _is_oracle, fq_table
if TYPE_CHECKING:
from .memory_engine import MemoryEngine
@@ -50,10 +49,6 @@ 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:
@@ -102,14 +97,10 @@ class MaintenanceLoop:
def _any_job_enabled() -> bool:
cfg = get_config()
reconcile_on = cfg.consolidation_reconcile_interval_seconds > 0
# Not gated on audit_log_enabled: that is per-bank overridable, so rows
# can exist even when the deployment default is off. Retention is driven
# purely by the (server-level) window.
audit_on = cfg.audit_log_retention_days > 0
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
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
return reconcile_on or audit_on or llm_on or mm_refresh_on
# ── loop ───────────────────────────────────────────────────────────────
@@ -143,8 +134,6 @@ 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,10 +152,7 @@ class MaintenanceLoop:
async def _run_retention(self, cfg: HindsightConfig) -> None:
# Retention days are static server-level config, so one global cutoff
# applies to every tenant schema (the routine sweeps them all).
# Not gated on audit_log_enabled: it is per-bank overridable, so a bank
# may be writing audit rows while the deployment default is off. Gating
# the purge on the global flag would let those rows accumulate forever.
if cfg.audit_log_retention_days > 0:
if cfg.audit_log_enabled and cfg.audit_log_retention_days > 0:
await self._purge_expired("audit_log", "started_at", cfg.audit_log_retention_days)
if cfg.llm_trace_enabled and cfg.llm_trace_retention_days > 0:
await self._purge_expired("llm_requests", "started_at", cfg.llm_trace_retention_days)
@@ -177,7 +163,7 @@ class MaintenanceLoop:
try:
async with acquire_with_retry(backend, max_retries=1) as conn:
rows = await conn.fetch(
f"SELECT * FROM {fq_routine('schemas_with_expired_rows')}($1, $2, $3)", table, ts_col, days
"SELECT * FROM public.schemas_with_expired_rows($1, $2, $3)", table, ts_col, days
)
for row in rows:
schema = row[0]
@@ -192,73 +178,6 @@ 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:
@@ -266,9 +185,7 @@ class MaintenanceLoop:
engine = self._engine
try:
async with acquire_with_retry(engine._backend, max_retries=1) as conn:
rows = await conn.fetch(
f"SELECT schema_name, bank_id FROM {fq_routine('banks_needing_consolidation')}()"
)
rows = await conn.fetch("SELECT schema_name, bank_id FROM public.banks_needing_consolidation()")
except Exception as e:
logger.warning(f"Consolidation reconcile discovery failed: {e}")
return
@@ -327,7 +244,7 @@ class MaintenanceLoop:
Discovery (the set of cron-scheduled models, minus any with an in-flight
refresh) is one cross-tenant round-trip via
``mental_models_with_cron()``. Cron *due-ness* is evaluated here in
``public.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
@@ -339,7 +256,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 "
f"FROM {fq_routine('mental_models_with_cron')}()"
"FROM public.mental_models_with_cron()"
)
except Exception as e:
logger.warning(f"Scheduled mental model refresh discovery failed: {e}")
File diff suppressed because it is too large Load Diff
@@ -142,21 +142,3 @@ 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,11 +186,7 @@ class MarkitdownParser(FileParser):
if Path(filename).suffix.lower() not in _TEXT_EXTENSIONS:
return None
try:
# 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")
file_data.decode("utf-8")
except UnicodeDecodeError:
return None
from markitdown import StreamInfo
@@ -14,9 +14,8 @@ import logging
import time
from typing import Any
from hindsight_api.engine.llm_interface import LLM_TOOL_CHOICE_AUTO, LLMInterface, LLMToolChoice
from hindsight_api.engine.llm_interface import LLMInterface
from hindsight_api.engine.llm_trace import LLMResponseUsage, stash_response_usage
from hindsight_api.engine.providers.llm_debug import dump_request_on_4xx
from hindsight_api.engine.response_models import LLMToolCall, LLMToolCallResult, TokenUsage
from hindsight_api.metrics import get_metrics_collector
@@ -35,43 +34,6 @@ 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.
@@ -244,9 +206,7 @@ class AnthropicLLM(LLMInterface):
}
if 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)
call_params["system"] = system_prompt
if use_forced_tool:
# Single tool whose input_schema IS the response schema; force the model to
@@ -386,9 +346,6 @@ class AnthropicLLM(LLMInterface):
logger.error(f"Anthropic auth error (HTTP {e.status_code}), not retrying: {str(e)}")
raise
# Diagnostic dump (opt-in) of the exact request behind any 4xx.
dump_request_on_4xx(scope=scope, provider=self.provider, model=self.model, err=e, request=call_params)
last_exception = e
if attempt < max_retries:
# Check if it's a rate limit or server error
@@ -423,7 +380,7 @@ class AnthropicLLM(LLMInterface):
max_retries: int = 5,
initial_backoff: float = 1.0,
max_backoff: float = 30.0,
tool_choice: LLMToolChoice = LLM_TOOL_CHOICE_AUTO,
tool_choice: str | dict[str, Any] = "auto",
) -> LLMToolCallResult:
"""
Make an LLM API call with tool/function calling support.
@@ -493,11 +450,6 @@ 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,
@@ -505,7 +457,7 @@ class AnthropicLLM(LLMInterface):
"max_tokens": max_completion_tokens or 4096,
}
if system_prompt:
call_params["system"] = _cached_system_blocks(system_prompt)
call_params["system"] = system_prompt
if self._extra_body:
call_params["extra_body"] = self._extra_body
@@ -581,8 +533,6 @@ class AnthropicLLM(LLMInterface):
except (APIConnectionError, APIStatusError) as e:
if isinstance(e, APIStatusError) and e.status_code in (401, 403):
raise
# Diagnostic dump (opt-in) of the exact request behind any 4xx.
dump_request_on_4xx(scope=scope, provider=self.provider, model=self.model, err=e, request=call_params)
last_exception = e
if attempt < max_retries:
await asyncio.sleep(min(initial_backoff * (2**attempt), max_backoff))
@@ -593,217 +543,6 @@ 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:
@@ -15,7 +15,7 @@ from typing import Any
from pydantic import ValidationError
from hindsight_api.engine.llm_interface import LLM_TOOL_CHOICE_AUTO, LLMInterface, LLMToolChoice, LLMToolChoiceMode
from hindsight_api.engine.llm_interface import LLMInterface
from hindsight_api.engine.llm_trace import LLMResponseUsage, stash_response_usage
from hindsight_api.engine.response_models import LLMToolCall, LLMToolCallResult, TokenUsage
from hindsight_api.metrics import get_metrics_collector
@@ -49,20 +49,6 @@ 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.
@@ -190,7 +176,6 @@ class ClaudeCodeLLM(LLMInterface):
from claude_agent_sdk import ( # type: ignore[unresolved-import]
AssistantMessage,
ClaudeAgentOptions,
ResultMessage,
TextBlock,
query,
)
@@ -224,19 +209,9 @@ class ClaudeCodeLLM(LLMInterface):
user_content += schema_instruction
# Configure SDK options
#
# tools=[] is required here for the same reason call_with_tools() below
# already sets it: with `tools` left at its default (None -> full
# "claude_code" built-in preset), allowed_tools=[] alone does not stop
# the CLI from loading the full built-in toolset and deferring into
# ToolSearch before answering, which burns the single max_turns=1
# budget on a tool-deferral step instead of a text response. Without
# this, single-turn calls intermittently fail with "Reached maximum
# number of turns (1)" even though the prompt itself needs no tools.
options = ClaudeAgentOptions(
system_prompt=system_prompt if system_prompt else None,
max_turns=1, # Single-turn for API-style interactions
tools=[], # Disable built-in tools so nothing forces a ToolSearch deferral
allowed_tools=[], # Disable tools for standard LLM calls
env=_get_isolated_claude_env(),
)
@@ -253,11 +228,6 @@ 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
@@ -392,7 +362,7 @@ class ClaudeCodeLLM(LLMInterface):
max_retries: int = 5,
initial_backoff: float = 1.0,
max_backoff: float = 30.0,
tool_choice: LLMToolChoice = LLM_TOOL_CHOICE_AUTO,
tool_choice: str | dict[str, Any] = "auto",
) -> LLMToolCallResult:
"""
Make an LLM API call with tool/function calling support using Claude Agent SDK.
@@ -410,7 +380,7 @@ class ClaudeCodeLLM(LLMInterface):
max_retries: Maximum retry attempts.
initial_backoff: Initial backoff time in seconds.
max_backoff: Maximum backoff time in seconds.
tool_choice: Canonical tool-selection policy.
tool_choice: How to choose tools - "auto", "none", "required", or specific function dict.
- "auto": Model decides whether to call tools (default)
- "required": Model must call at least one tool
- "none": Model must not call any tools
@@ -423,7 +393,6 @@ class ClaudeCodeLLM(LLMInterface):
AssistantMessage,
ClaudeAgentOptions,
ClaudeSDKClient,
ResultMessage,
SdkMcpTool,
TextBlock,
ToolUseBlock,
@@ -504,27 +473,30 @@ class ClaudeCodeLLM(LLMInterface):
mcp_servers_config = {"hindsight_tools": mcp_server} if sdk_tools else {}
# Process tool_choice
if tool_choice.mode is LLMToolChoiceMode.NAMED:
if isinstance(tool_choice, dict) and tool_choice.get("type") == "function":
# Force a specific tool: filter allowed_tools to only that tool and add instruction
forced_name = tool_choice.selected_function_name
forced_tool_mcp_name = f"mcp__hindsight_tools__{forced_name}"
if forced_tool_mcp_name in allowed_tool_names:
allowed_tool_names = [forced_tool_mcp_name]
force_instruction = (
f"\n\nIMPORTANT: You MUST call the '{forced_name}' tool. Do not respond with text only."
)
system_prompt += force_instruction
logger.debug(f"Claude Code: Forcing tool call to '{forced_name}'")
else:
logger.warning(f"Claude Code: Forced tool '{forced_name}' not found in available tools")
elif tool_choice.mode is LLMToolChoiceMode.REQUIRED:
forced_name = tool_choice.get("function", {}).get("name")
if forced_name:
# Filter to only the forced tool (with MCP prefix)
forced_tool_mcp_name = f"mcp__hindsight_tools__{forced_name}"
if forced_tool_mcp_name in allowed_tool_names:
allowed_tool_names = [forced_tool_mcp_name]
# Add strong instruction to system prompt
force_instruction = (
f"\n\nIMPORTANT: You MUST call the '{forced_name}' tool. Do not respond with text only."
)
system_prompt += force_instruction
logger.debug(f"Claude Code: Forcing tool call to '{forced_name}'")
else:
logger.warning(f"Claude Code: Forced tool '{forced_name}' not found in available tools")
elif tool_choice == "required":
# Must call at least one tool
tool_instruction = (
"\n\nIMPORTANT: You MUST call at least one of the available tools. Do not respond with text only."
)
system_prompt += tool_instruction
logger.debug("Claude Code: Tool call required")
elif tool_choice.mode is LLMToolChoiceMode.NONE:
elif tool_choice == "none":
# No tools should be called - disable all tools
allowed_tool_names = []
mcp_servers_config = {}
@@ -560,9 +532,6 @@ 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,7 +19,6 @@ from __future__ import annotations
import base64
import binascii
import contextlib
import json
import logging
import os
@@ -32,11 +31,6 @@ from typing import Any
import httpx
try:
import fcntl
except ImportError: # pragma: no cover - Windows
fcntl = None # type: ignore[assignment]
logger = logging.getLogger(__name__)
@@ -64,9 +58,6 @@ _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:
@@ -85,44 +76,6 @@ 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.
@@ -239,34 +192,6 @@ 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.
@@ -300,11 +225,6 @@ 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
# ------------------------------------------------------------------
@@ -419,93 +339,78 @@ class CodexAuthManager:
if not self._token_is_stale():
return
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
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."
)
if not self.refresh_token:
raise RuntimeError(
"Codex access_token is expired but no refresh_token is available. "
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}). "
"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."
)
log_reason = f" ({reason})" if reason else ""
logger.info(f"Refreshing Codex OAuth access_token{log_reason}")
if response.status_code >= 400:
raise RuntimeError(f"Codex OAuth refresh failed with HTTP {response.status_code}")
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
try:
body = response.json()
except json.JSONDecodeError as e:
raise RuntimeError(f"Codex OAuth refresh returned non-JSON body: {e}") from e
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_access = body.get("access_token")
if not new_access:
raise RuntimeError("Codex OAuth refresh returned no access_token")
if response.status_code >= 400:
raise RuntimeError(f"Codex OAuth refresh failed with HTTP {response.status_code}")
new_refresh = body.get("refresh_token") or self.refresh_token
new_id_token = body.get("id_token")
try:
body = response.json()
except json.JSONDecodeError as e:
raise RuntimeError(f"Codex OAuth refresh returned non-JSON body: {e}") from e
# 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
new_access = body.get("access_token")
if not new_access:
raise RuntimeError("Codex OAuth refresh returned no access_token")
persisted: dict[str, Any] = {
"access_token": new_access,
"refresh_token": new_refresh,
}
if new_id_token:
persisted["id_token"] = new_id_token
new_refresh = body.get("refresh_token") or self.refresh_token
new_id_token = body.get("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."
)
# 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")
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.
@@ -25,11 +25,9 @@ from typing import Any
import httpx
from hindsight_api.engine.llm_interface import LLM_TOOL_CHOICE_AUTO, LLMInterface, LLMToolChoice, LLMToolChoiceMode
from hindsight_api.engine.llm_interface import LLMInterface
from hindsight_api.engine.llm_trace import LLMResponseUsage, stash_response_usage
from hindsight_api.engine.providers.llm_debug import dump_request_on_4xx
from hindsight_api.engine.response_models import LLMToolCall, LLMToolCallResult, TokenUsage
from hindsight_api.engine.structured_output import strict_json_schema
from hindsight_api.metrics import get_metrics_collector
from .codex_auth import (
@@ -55,59 +53,6 @@ __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):
"""
@@ -195,18 +140,6 @@ 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
@@ -343,6 +276,32 @@ class CodexLLM(LLMInterface):
}
return mapping.get(effort.lower(), "auto")
def _normalize_tool_choice(self, tool_choice: str | dict[str, Any]) -> str | dict[str, Any]:
"""Normalize forced function tool choice for the Codex Responses API.
Older agent paths may still pass OpenAI chat-completions style named
tool choice payloads such as:
{"type": "function", "function": {"name": "recall"}}
Codex Responses expects the named function at the top level instead:
{"type": "function", "name": "recall"}
"""
if not isinstance(tool_choice, dict):
return tool_choice
if str(tool_choice.get("type") or "").strip() != "function":
return tool_choice
function_payload = tool_choice.get("function")
if isinstance(function_payload, dict):
function_name = str(function_payload.get("name") or "").strip()
if function_name:
return {"type": "function", "name": function_name}
function_name = str(tool_choice.get("name") or "").strip()
if function_name:
return {"type": "function", "name": function_name}
return tool_choice
async def verify_connection(self) -> None:
"""Verify Codex connection by making a simple test call."""
try:
@@ -377,18 +336,7 @@ class CodexLLM(LLMInterface):
strict_schema: bool = False,
return_usage: bool = False,
) -> Any:
"""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.
"""
"""Make API call to Codex backend with SSE streaming."""
start_time = time.time()
# Proactively refresh the OAuth access_token if it's near expiry.
@@ -413,22 +361,11 @@ class CodexLLM(LLMInterface):
else:
user_messages.append(msg)
# 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
# Add JSON schema instruction if response_format is provided
if response_format is not None and hasattr(response_format, "model_json_schema"):
schema = strict_json_schema(response_format) if strict_schema else response_format.model_json_schema()
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
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
# gpt-5.2-codex only supports "detailed" reasoning summary
reasoning_summary = "detailed" if "5.2" in self.model else self.reasoning_summary
@@ -455,21 +392,13 @@ class CodexLLM(LLMInterface):
"prompt_cache_key": str(uuid.uuid4()),
}
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()
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",
}
url = f"{self.base_url}/codex/responses"
@@ -483,15 +412,8 @@ class CodexLLM(LLMInterface):
response = await self._client.post(url, json=payload, headers=headers, timeout=120.0)
response.raise_for_status()
# 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)
# Parse SSE stream
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
@@ -504,28 +426,7 @@ class CodexLLM(LLMInterface):
)
# Handle structured output
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:
if response_format is not None:
# Models may wrap JSON in markdown
clean_content = content
if "```json" in content:
@@ -536,20 +437,13 @@ class CodexLLM(LLMInterface):
try:
json_data = json.loads(clean_content)
except json.JSONDecodeError as e:
# 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
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
@@ -648,9 +542,6 @@ class CodexLLM(LLMInterface):
"Run 'codex auth login' to re-authenticate."
) from e
# Diagnostic dump (opt-in) of the exact request behind any 4xx.
dump_request_on_4xx(scope=scope, provider=self.provider, model=self.model, err=e, request=payload)
# Log the actual error message from the API
error_detail = e.response.text[:500] if hasattr(e.response, "text") else str(e)
@@ -746,7 +637,7 @@ class CodexLLM(LLMInterface):
max_retries: int = 5,
initial_backoff: float = 1.0,
max_backoff: float = 30.0,
tool_choice: LLMToolChoice = LLM_TOOL_CHOICE_AUTO,
tool_choice: str | dict[str, Any] = "auto",
) -> LLMToolCallResult:
"""
Make API call with tool calling support.
@@ -763,7 +654,7 @@ class CodexLLM(LLMInterface):
max_retries: Maximum retry attempts.
initial_backoff: Initial backoff time in seconds.
max_backoff: Maximum backoff time in seconds.
tool_choice: Canonical tool-selection policy.
tool_choice: How to choose tools - "auto", "none", "required", or a specific function.
Returns:
LLMToolCallResult with content and/or tool_calls.
@@ -825,11 +716,7 @@ class CodexLLM(LLMInterface):
"instructions": system_instruction,
"input": user_messages,
"tools": codex_tools,
"tool_choice": (
{"type": "function", "name": tool_choice.selected_function_name}
if tool_choice.mode is LLMToolChoiceMode.NAMED
else tool_choice.mode.value
),
"tool_choice": self._normalize_tool_choice(tool_choice),
"parallel_tool_calls": True,
"reasoning": {"summary": reasoning_summary},
"store": False,
@@ -838,7 +725,13 @@ class CodexLLM(LLMInterface):
"prompt_cache_key": str(uuid.uuid4()),
}
headers = self._build_request_headers()
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",
}
url = f"{self.base_url}/codex/responses"
@@ -935,8 +828,6 @@ class CodexLLM(LLMInterface):
)
except Exception as e:
# Diagnostic dump (opt-in) of the exact request behind any 4xx.
dump_request_on_4xx(scope=scope, provider=self.provider, model=self.model, err=e, request=payload)
logger.error(f"Codex tool call error: {e}")
raise
@@ -981,13 +872,8 @@ class CodexLLM(LLMInterface):
try:
arguments = json.loads(arguments_str)
except json.JSONDecodeError:
# 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 = {}
logger.warning(f"Failed to parse tool arguments: {arguments_str}")
arguments = {}
tool_calls.append(
LLMToolCall(
@@ -56,14 +56,6 @@ _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:
@@ -100,10 +92,6 @@ 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(
@@ -241,99 +229,18 @@ 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
@@ -347,10 +254,8 @@ 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"{ttl_seconds if ttl_seconds is not None else self._ttl_seconds}s",
"ttl": f"{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,17 +13,15 @@ import json
import logging
import time
from contextvars import ContextVar
from dataclasses import dataclass
from typing import Any
from google import genai
from google.genai import errors as genai_errors
from google.genai import types as genai_types
from hindsight_api.engine.llm_interface import LLM_TOOL_CHOICE_AUTO, LLMInterface, LLMToolChoice, LLMToolChoiceMode
from hindsight_api.engine.llm_interface import LLMInterface
from hindsight_api.engine.llm_trace import LLMResponseUsage, stash_response_usage
from hindsight_api.engine.llm_wrapper import parse_llm_json
from hindsight_api.engine.providers.llm_debug import dump_request_on_4xx
from hindsight_api.engine.response_models import LLMToolCall, LLMToolCallResult, TokenUsage
from hindsight_api.metrics import get_metrics_collector
from hindsight_api.worker.stage import set_stage
@@ -65,99 +63,6 @@ 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] = []
pending_tool_names_by_call_id: dict[str, str] = {}
i = 0
while i < len(msg_list):
msg = msg_list[i]
role = msg.get("role", "user")
content = msg.get("content", "")
if role != "tool" and pending_tool_names_by_call_id:
missing_ids = ", ".join(sorted(pending_tool_names_by_call_id))
raise ValueError(f"Gemini assistant tool calls require results before the next message: {missing_ids}")
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", "")
tool_call_id = tool_msg["tool_call_id"]
tool_name = pending_tool_names_by_call_id.pop(tool_call_id, None)
if tool_name is None:
raise ValueError(f"Gemini tool result references unknown tool_call_id {tool_call_id!r}")
parts.append(
genai_types.Part(
function_response=genai_types.FunctionResponse(
name=tool_name,
response={"result": tool_content},
)
)
)
i += 1
if pending_tool_names_by_call_id:
missing_ids = ", ".join(sorted(pending_tool_names_by_call_id))
raise ValueError(f"Gemini assistant tool calls are missing results: {missing_ids}")
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:
tool_call_id = tc["id"]
fn = tc["function"]
fn_name = fn["name"]
if tool_call_id in pending_tool_names_by_call_id:
raise ValueError(
f"Gemini assistant tool call id {tool_call_id!r} must be unique within its turn"
)
pending_tool_names_by_call_id[tool_call_id] = fn_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
if pending_tool_names_by_call_id:
missing_ids = ", ".join(sorted(pending_tool_names_by_call_id))
raise ValueError(f"Gemini assistant tool calls are missing results: {missing_ids}")
return _GeminiConversation(system_instruction=system_instruction, contents=gemini_contents)
class GeminiLLM(LLMInterface):
"""
LLM provider for Google Gemini and Vertex AI.
@@ -574,17 +479,6 @@ class GeminiLLM(LLMInterface):
logger.error(f"Gemini auth error (HTTP {e.code}), not retrying: {str(e)}")
raise
# Diagnostic dump (opt-in) of the exact request behind any 4xx, captured
# before the cache-drop retry below rebuilds the config so we see what failed.
dump_request_on_4xx(
scope=scope,
provider=self.provider,
model=self.model,
err=e,
request=generation_config,
messages=gemini_contents,
)
# Cached-request safety net: a stale/invalid/expired CachedContent
# (or an incompatibility like cache + tool_config) surfaces as a 400.
# Retrying the same cached request can't recover, so on the first
@@ -631,9 +525,8 @@ class GeminiLLM(LLMInterface):
max_retries: int = 5,
initial_backoff: float = 1.0,
max_backoff: float = 30.0,
tool_choice: LLMToolChoice = LLM_TOOL_CHOICE_AUTO,
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.
@@ -647,22 +540,15 @@ class GeminiLLM(LLMInterface):
max_retries: Maximum retry attempts.
initial_backoff: Initial backoff time in seconds.
max_backoff: Maximum backoff time in seconds.
tool_choice: Canonical tool-selection policy.
tool_choice: How to choose tools (Gemini uses "auto" only).
cached_prefix: Optional CachedContent resource name (from
``GeminiCacheManager.get_or_create`` or ``create_incremental``).
When set, the system_instruction and tool definitions are assumed
``GeminiCacheManager.get_or_create`` with ``tools=...``). 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.
@@ -670,45 +556,86 @@ class GeminiLLM(LLMInterface):
start_time = time.time()
using_cache = cached_prefix is not None
# 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.
# 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``.
gemini_tools = []
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"),
)
]
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`` 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
# 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", "")
# 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
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},
)
)
)
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
# Apply safety settings: context var (per-request bank override) takes precedence over instance default
effective_safety_settings = _safety_settings_ctx.get()
@@ -738,20 +665,22 @@ class GeminiLLM(LLMInterface):
config_kwargs["max_output_tokens"] = max_completion_tokens
# Map OpenAI-style tool_choice to Gemini FunctionCallingConfig
if tool_choice.mode is LLMToolChoiceMode.REQUIRED:
if tool_choice == "required":
config_kwargs["tool_config"] = genai_types.ToolConfig(
function_calling_config=genai_types.FunctionCallingConfig(
mode="ANY",
)
)
elif tool_choice.mode is LLMToolChoiceMode.NAMED:
config_kwargs["tool_config"] = genai_types.ToolConfig(
function_calling_config=genai_types.FunctionCallingConfig(
mode="ANY",
allowed_function_names=[tool_choice.selected_function_name],
elif isinstance(tool_choice, dict) and tool_choice.get("type") == "function":
fn_name = tool_choice.get("function", {}).get("name")
if fn_name:
config_kwargs["tool_config"] = genai_types.ToolConfig(
function_calling_config=genai_types.FunctionCallingConfig(
mode="ANY",
allowed_function_names=[fn_name],
)
)
)
elif tool_choice.mode is LLMToolChoiceMode.NONE:
elif tool_choice == "none":
config_kwargs["tool_config"] = genai_types.ToolConfig(
function_calling_config=genai_types.FunctionCallingConfig(mode="NONE")
)
@@ -772,14 +701,10 @@ 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=active_contents,
contents=gemini_contents,
config=config,
),
timeout=90.0, # Safety net for network hangs; valid slow responses are <90s
@@ -882,17 +807,6 @@ class GeminiLLM(LLMInterface):
logger.error(f"Gemini auth error (HTTP {e.code}), not retrying: {str(e)}")
raise
# Diagnostic dump (opt-in) of the exact request behind any 4xx, captured
# before the cache-drop retry below rebuilds the config so we see what failed.
dump_request_on_4xx(
scope=scope,
provider=self.provider,
model=self.model,
err=e,
request=config,
messages=active_contents,
)
# Cached-request safety net (see ``call``): a stale/invalid cache or
# a cache+tool_config conflict surfaces as a 400. Drop the cache,
# invalidate it for later operations, and retry THIS call inline
@@ -969,56 +883,6 @@ 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
@@ -1166,7 +1030,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`` whenever a schema is present.
grammar-enforces via ``responseJsonSchema`` when ``strict`` is set.
"""
system_texts: list[str] = []
contents: list[dict[str, Any]] = []
@@ -1195,13 +1059,8 @@ class GeminiLLM(LLMInterface):
system_texts.append(
"You must respond with valid JSON matching this schema:\n" + json.dumps(schema, ensure_ascii=False)
)
# #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
if json_schema.get("strict"):
generation_config["responseJsonSchema"] = schema
request: dict[str, Any] = {"contents": contents}
if system_texts:
@@ -22,18 +22,9 @@ from typing import Any
from litellm.exceptions import Timeout as LiteLLMTimeout
from hindsight_api.config import DEFAULT_LLM_TIMEOUT, ENV_LLM_TIMEOUT
from hindsight_api.engine.llm_interface import (
LLM_TOOL_CHOICE_AUTO,
LLMInterface,
LLMToolChoice,
LLMToolChoiceMode,
OutputTooLongError,
)
from hindsight_api.engine.llm_interface import LLMInterface, OutputTooLongError
from hindsight_api.engine.llm_trace import LLMResponseUsage, stash_response_usage
from hindsight_api.engine.llm_wrapper import parse_llm_json
from hindsight_api.engine.providers.llm_debug import dump_request_on_4xx
from hindsight_api.engine.response_models import LLMToolCall, LLMToolCallResult, TokenUsage
from hindsight_api.engine.structured_output import strict_json_schema
from hindsight_api.metrics import get_metrics_collector
from hindsight_api.worker.stage import set_stage
@@ -242,7 +233,7 @@ class LiteLLMLLM(LLMInterface):
# Add JSON schema response format if provided
if response_format is not None and hasattr(response_format, "model_json_schema"):
schema = strict_json_schema(response_format) if strict_schema else response_format.model_json_schema()
schema = response_format.model_json_schema()
call_kwargs["response_format"] = {
"type": "json_schema",
"json_schema": {
@@ -286,17 +277,7 @@ class LiteLLMLLM(LLMInterface):
try:
json_data = json.loads(clean_content)
except json.JSONDecodeError:
try:
json_data = json.loads(content)
except json.JSONDecodeError:
if attempt < max_retries:
# Prefer a clean re-roll first — a fresh generation
# usually beats repairing a malformed one.
raise
# Retry budget spent: structural repair as a last
# resort (#2547/#2544). Raises again if unrecoverable,
# which the outer handler surfaces loudly.
json_data = parse_llm_json(content)
json_data = json.loads(content)
if skip_validation:
result = json_data
@@ -397,9 +378,6 @@ class LiteLLMLLM(LLMInterface):
logger.error(f"LiteLLM auth error, not retrying: {e}")
raise
# Diagnostic dump (opt-in) of the exact request behind any 4xx.
dump_request_on_4xx(scope=scope, provider=self.provider, model=self.model, err=e, request=call_kwargs)
last_exception = e
if attempt < max_retries:
# Retry on rate limits, connection errors, server errors
@@ -430,20 +408,13 @@ class LiteLLMLLM(LLMInterface):
max_retries: int = 5,
initial_backoff: float = 1.0,
max_backoff: float = 30.0,
tool_choice: LLMToolChoice = LLM_TOOL_CHOICE_AUTO,
tool_choice: str | dict[str, Any] = "auto",
) -> LLMToolCallResult:
start_time = time.time()
call_kwargs = self._build_common_kwargs(messages, max_completion_tokens, temperature)
call_kwargs["tools"] = tools
call_kwargs["tool_choice"] = (
{
"type": "function",
"function": {"name": tool_choice.selected_function_name},
}
if tool_choice.mode is LLMToolChoiceMode.NAMED
else tool_choice.mode.value
)
call_kwargs["tool_choice"] = tool_choice
last_exception = None
for attempt in range(max_retries + 1):
@@ -553,9 +524,6 @@ class LiteLLMLLM(LLMInterface):
if "401" in error_str or "403" in error_str or "unauthorized" in error_str:
raise
# Diagnostic dump (opt-in) of the exact request behind any 4xx.
dump_request_on_4xx(scope=scope, provider=self.provider, model=self.model, err=e, request=call_kwargs)
last_exception = e
if attempt < max_retries:
is_retryable = any(
@@ -22,7 +22,7 @@ import time
from pathlib import Path
from typing import Any
from hindsight_api.engine.llm_interface import LLM_TOOL_CHOICE_AUTO, LLMInterface, LLMToolChoice
from hindsight_api.engine.llm_interface import LLMInterface
from hindsight_api.engine.response_models import LLMToolCallResult
logger = logging.getLogger(__name__)
@@ -394,7 +394,7 @@ class LlamaCppLLM(LLMInterface):
max_retries: int = 5,
initial_backoff: float = 1.0,
max_backoff: float = 30.0,
tool_choice: LLMToolChoice = LLM_TOOL_CHOICE_AUTO,
tool_choice: str | dict[str, Any] = "auto",
) -> LLMToolCallResult:
"""Delegate tool calls to the OpenAI-compatible API."""
await self._ensure_initialized()
@@ -1,168 +0,0 @@
"""Opt-in diagnostic: dump the exact request behind an LLM 4xx rejection.
Some ``400 INVALID_ARGUMENT`` / ``400 Bad Request`` rejections of structured-output
calls are not reproducible by reconstructing the request after the fact the failing
factor lives in the request as it was actually assembled at runtime. Reconstructed
replays of the same inputs return ``200``, so the only reliable way to see what the
model rejected is to capture the real request at the moment it fails.
This helper is provider-agnostic. Every provider's error handler calls
``dump_request_on_4xx`` with whatever it assembled a Pydantic config
(google-genai ``GenerateContentConfig``), a kwargs dict (OpenAI / Anthropic /
LiteLLM ``**call_params``), etc. plus the raised error. The helper self-gates:
it is a no-op unless the ``llm_debug_dump_4xx`` config flag
(``HINDSIGHT_API_LLM_DEBUG_DUMP_4XX``) is enabled AND the error carries a 4xx
status, so callers can drop one unconditional call into each ``except`` block.
Safety / scope:
- Off by default the config flag is unset in normal operation.
- The serialized config omits message bodies (the ``messages``/``contents``/``input``
keys are stripped); message previews are length-capped, so an enabled dump can't
flood logs or spill large bodies.
- Never raises diagnostics must not break the request path (falls back to ``repr``).
"""
from __future__ import annotations
import json
import logging
from dataclasses import dataclass
from typing import Any
logger = logging.getLogger(__name__)
# Top-level request keys whose values are message bodies. Stripped from the config
# view so the dump never spills large user content — previews are logged separately.
_CONTENT_KEYS = ("messages", "contents", "input")
_PREVIEW_CHARS = 1500
_CONFIG_REPR_CAP = 8000
_ERR_CAP = 200
def _enabled() -> bool:
from hindsight_api.config import get_config
return bool(get_config().llm_debug_dump_4xx)
def status_code_of(err: Any) -> int | None:
"""Best-effort HTTP status of a provider error, across SDK error shapes.
OpenAI/Anthropic expose ``status_code``; google-genai uses ``code``; some wrap the
status on a ``response``. Returns None when no integer status is discoverable.
"""
for attr in ("status_code", "code", "http_status"):
value = getattr(err, attr, None)
if isinstance(value, int):
return value
response = getattr(err, "response", None)
if response is not None:
value = getattr(response, "status_code", None)
if isinstance(value, int):
return value
return None
def _serialize_config(request: Any) -> str:
"""Render the request config to a string without message bodies, never raising."""
try:
if request is None:
return "null"
# Pydantic models (google-genai GenerateContentConfig, SDK params objects).
dump = getattr(request, "model_dump_json", None)
if callable(dump):
return dump(exclude_none=True)
if isinstance(request, dict):
view = {k: v for k, v in request.items() if k not in _CONTENT_KEYS}
return json.dumps(view, ensure_ascii=False, default=str)
return repr(request)[:_CONFIG_REPR_CAP]
except Exception:
return repr(request)[:_CONFIG_REPR_CAP]
@dataclass
class _MessagePreview:
"""A message rendered for the dump: role + extracted text (not yet length-capped)."""
role: str
text: str
def _message_preview(msg: Any) -> _MessagePreview:
"""Extract role + text from a message across dict and provider-object shapes."""
# OpenAI / Anthropic dict: {"role": ..., "content": str | list[block]}
if isinstance(msg, dict):
role = str(msg.get("role", "?"))
content = msg.get("content")
if isinstance(content, str):
return _MessagePreview(role, content)
if isinstance(content, list):
text = ""
for block in content:
if isinstance(block, dict):
text += block.get("text") or ""
else:
text += getattr(block, "text", "") or ""
return _MessagePreview(role, text)
return _MessagePreview(role, "" if content is None else str(content))
# google-genai Content: role + parts[].text
role = str(getattr(msg, "role", "?"))
text = ""
for part in getattr(msg, "parts", None) or []:
text += getattr(part, "text", None) or ""
if not text:
text = getattr(msg, "content", "") or ""
return _MessagePreview(role, text)
def _resolve_messages(request: Any, messages: Any) -> Any:
"""Where per-message previews come from: explicit ``messages``, else inside ``request``."""
if messages is not None:
return messages
if isinstance(request, dict):
for key in _CONTENT_KEYS:
if key in request:
return request[key]
return []
def dump_request_on_4xx(
*,
scope: str,
provider: str,
model: str,
err: Any,
request: Any = None,
messages: Any = None,
) -> None:
"""Log the exact request behind an LLM 4xx when the diagnostic is enabled.
No-op unless ``HINDSIGHT_API_LLM_DEBUG_DUMP_4XX`` is truthy and ``err`` carries a
4xx status. ``request`` is whatever the provider assembled (a Pydantic config, a
kwargs dict, ...); ``messages`` overrides where the per-message previews come from
(defaults to the message list found inside ``request``).
"""
if not _enabled():
return
code = status_code_of(err)
if code is None or not (400 <= code < 500):
return
try:
cfg_repr = _serialize_config(request)
summary = []
for msg in _resolve_messages(request, messages) or []:
m = _message_preview(msg)
summary.append({"role": m.role, "chars": len(m.text), "preview": m.text[:_PREVIEW_CHARS]})
logger.error(
"[LLM_4XX_DUMP] provider=%s model=%s scope=%s code=%s err=%s config=%s contents=%s",
provider,
model,
scope,
code,
str(err)[:_ERR_CAP],
cfg_repr,
json.dumps(summary, ensure_ascii=False),
)
except Exception as dump_exc: # never let diagnostics break the request path
logger.warning("[LLM_4XX_DUMP] failed to serialize rejected request: %s", dump_exc)
@@ -9,7 +9,7 @@ import logging
from collections.abc import Callable
from typing import Any
from ..llm_interface import LLM_TOOL_CHOICE_AUTO, LLMInterface, LLMToolChoice
from ..llm_interface import LLMInterface
from ..response_models import LLMToolCall, LLMToolCallResult, TokenUsage
logger = logging.getLogger(__name__)
@@ -200,7 +200,7 @@ class MockLLM(LLMInterface):
max_retries: int = 5,
initial_backoff: float = 1.0,
max_backoff: float = 30.0,
tool_choice: LLMToolChoice = LLM_TOOL_CHOICE_AUTO,
tool_choice: str | dict[str, Any] = "auto",
) -> LLMToolCallResult:
"""
Make a mock LLM API call with tool/function calling support.
@@ -10,7 +10,7 @@ it raises a clear error instead of a confusing connection failure.
import logging
from typing import Any
from ..llm_interface import LLM_TOOL_CHOICE_AUTO, LLMInterface, LLMToolChoice
from ..llm_interface import LLMInterface
from ..response_models import LLMToolCallResult
logger = logging.getLogger(__name__)
@@ -65,7 +65,7 @@ class NoneLLM(LLMInterface):
max_retries: int = 5,
initial_backoff: float = 1.0,
max_backoff: float = 30.0,
tool_choice: LLMToolChoice = LLM_TOOL_CHOICE_AUTO,
tool_choice: str | dict[str, Any] = "auto",
) -> LLMToolCallResult:
"""Raise LLMNotAvailableError — no LLM is configured."""
raise LLMNotAvailableError(
@@ -36,18 +36,9 @@ from openai import APIConnectionError, APIStatusError, AsyncOpenAI, LengthFinish
from hindsight_api.config import DEFAULT_LLM_TIMEOUT, ENV_LLM_TIMEOUT
from hindsight_api.engine.bank_attribution import apply_bank_attribution
from hindsight_api.engine.llm_interface import (
LLM_TOOL_CHOICE_AUTO,
LLMInterface,
LLMToolChoice,
LLMToolChoiceMode,
OutputTooLongError,
ProviderRateLimitResetError,
)
from hindsight_api.engine.llm_interface import LLMInterface, OutputTooLongError, ProviderRateLimitResetError
from hindsight_api.engine.llm_trace import LLMResponseUsage, stash_response_usage
from hindsight_api.engine.providers.llm_debug import dump_request_on_4xx
from hindsight_api.engine.response_models import LLMToolCall, LLMToolCallResult, TokenUsage
from hindsight_api.engine.structured_output import strict_json_schema
from hindsight_api.metrics import get_metrics_collector
from hindsight_api.worker.stage import set_stage
@@ -56,31 +47,15 @@ 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
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"
# Self-hosted OpenAI-compatible servers 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). 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).
# --enable-auto-tool-choice). llama-server (the "llamacpp" provider) honors
# "required" correctly and is intentionally excluded (#1179).
_TOOL_CHOICE_REQUIRED_UNSUPPORTED_PROVIDERS = frozenset({"lmstudio", "ollama"})
@@ -92,68 +67,23 @@ 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. 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.
response format is requested. This strips the fences while preserving
the JSON content inside. Returns the original content unchanged if
no fences are detected.
"""
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
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
# Reasoning/thinking tags emitted by extended-thinking models. Some providers
@@ -505,8 +435,6 @@ 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,
):
"""
@@ -521,8 +449,6 @@ 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)
@@ -603,7 +529,6 @@ 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 {}
@@ -634,17 +559,17 @@ class OpenAICompatibleLLM(LLMInterface):
def _drops_tool_choice_required(self) -> bool:
"""Whether this endpoint silently ignores ``tool_choice="required"``.
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.
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.
"""
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
if self.provider in _TOOL_CHOICE_REQUIRED_UNSUPPORTED_PROVIDERS:
return True
return self.provider == "openai" and bool(self.base_url)
async def verify_connection(self) -> None:
"""
@@ -657,7 +582,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=self._verification_max_completion_tokens(),
max_completion_tokens=100,
max_retries=2,
initial_backoff=0.5,
max_backoff=2.0,
@@ -719,11 +644,6 @@ 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]],
@@ -811,7 +731,6 @@ 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
@@ -827,7 +746,7 @@ class OpenAICompatibleLLM(LLMInterface):
if response_format is not None:
schema = None
if hasattr(response_format, "model_json_schema"):
schema = strict_json_schema(response_format) if strict_schema else response_format.model_json_schema()
schema = response_format.model_json_schema()
if strict_schema and schema is not None:
# Use OpenAI's strict JSON schema enforcement
@@ -967,9 +886,7 @@ class OpenAICompatibleLLM(LLMInterface):
output_tokens = max(0, output_tokens - thoughts_tokens)
total_tokens = max(0, total_tokens - thoughts_tokens)
# 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.
# Record LLM metrics
metrics = get_metrics_collector()
metrics.record_llm_call(
provider=self.provider,
@@ -979,8 +896,6 @@ 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
@@ -1049,9 +964,6 @@ class OpenAICompatibleLLM(LLMInterface):
logger.error(f"Auth error (HTTP {e.status_code}), not retrying: {str(e)}")
raise
# Diagnostic dump (opt-in) of the exact request behind any 4xx.
dump_request_on_4xx(scope=scope, provider=self.provider, model=self.model, err=e, request=call_params)
_raise_provider_quota_defer(
e, provider=self.provider, model=self.model, scope=scope, max_backoff=max_backoff
)
@@ -1142,7 +1054,7 @@ class OpenAICompatibleLLM(LLMInterface):
max_retries: int = 5,
initial_backoff: float = 1.0,
max_backoff: float = 30.0,
tool_choice: LLMToolChoice = LLM_TOOL_CHOICE_AUTO,
tool_choice: str | dict[str, Any] = "auto",
) -> LLMToolCallResult:
"""
Make an LLM API call with tool/function calling support.
@@ -1156,43 +1068,51 @@ class OpenAICompatibleLLM(LLMInterface):
max_retries: Maximum retry attempts.
initial_backoff: Initial backoff time in seconds.
max_backoff: Maximum backoff time in seconds.
tool_choice: Canonical tool-selection policy.
tool_choice: How to choose tools - "auto", "none", "required", or specific function.
Returns:
LLMToolCallResult with content and/or tool_calls.
"""
start_time = time.time()
request_tool_choice: str | None
if tool_choice.mode is LLMToolChoiceMode.NAMED:
forced_name = tool_choice.selected_function_name
filtered = [tool for tool in tools if tool.get("function", {}).get("name") == forced_name]
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 = LLMToolChoiceMode.REQUIRED.value
elif tool_choice.mode is LLMToolChoiceMode.AUTO:
request_tool_choice = None
else:
request_tool_choice = tool_choice.mode.value
request_tool_choice: str | dict[str, Any] | None = tool_choice
# Normalize named tool_choice dicts to "required" + filter tools.
# Some providers (e.g. LM Studio, Ollama) reject the OpenAI named format
# {"type": "function", "function": {"name": "..."}}. The semantics are
# identical to tool_choice="required" with the tools list restricted to
# just the requested tool, so we apply that transformation where supported.
if isinstance(request_tool_choice, dict) and request_tool_choice.get("type") == "function":
forced_name = request_tool_choice.get("function", {}).get("name")
if forced_name:
filtered = [t for t in tools if t.get("function", {}).get("name") == forced_name]
if filtered:
tools = filtered
request_tool_choice = "required"
# DeepSeek accepts tool calls but rejects explicit required/named
# tool_choice values. The tools list has already been narrowed for
# forced calls, so omitting tool_choice preserves the practical behavior.
if "deepseek" in self.model.lower() and tool_choice.mode is not LLMToolChoiceMode.AUTO:
if "deepseek" in self.model.lower() and request_tool_choice != "auto":
request_tool_choice = None
# LM Studio and Ollama silently drop tool_choice="required", returning an
# empty tool_calls array instead of forcing a call (#1563/#1179).
# "auto" is the OpenAI API default — omitting tool_choice is semantically
# identical. Some providers (e.g. DeepSeek's reasoner pathway, which
# deepseek-v4-flash falls into when thinking mode is enabled) reject the
# parameter outright, returning HTTP 400 even for value "auto". Sending it
# only when the caller asks for a non-default behaviour avoids those 400s
# without changing semantics for compliant providers.
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).
# 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. Generic OpenAI-compatible endpoints retain the canonical
# ``required`` contract regardless of whether they use a custom base URL.
if request_tool_choice == LLMToolChoiceMode.REQUIRED.value and self._drops_tool_choice_required():
# auto. The real OpenAI API honors "required" and is left untouched.
if request_tool_choice == "required" and self._drops_tool_choice_required():
request_tool_choice = None
# DeepSeek tool-call replies can carry provider-specific reasoning_content.
@@ -1229,7 +1149,6 @@ 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:
@@ -1277,8 +1196,6 @@ 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,
@@ -1288,8 +1205,6 @@ 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
@@ -1351,10 +1266,6 @@ class OpenAICompatibleLLM(LLMInterface):
f"not retrying: {_summarize_status_error(e)}"
)
raise
# Diagnostic dump (opt-in) of the exact request behind any 4xx.
dump_request_on_4xx(scope=scope, provider=self.provider, model=self.model, err=e, request=call_params)
_raise_provider_quota_defer(
e, provider=self.provider, model=self.model, scope=scope, max_backoff=max_backoff
)
@@ -1426,10 +1337,9 @@ 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,7 +6,6 @@ structured information like temporal constraints.
"""
import logging
import re
from abc import ABC, abstractmethod
from datetime import datetime, timedelta
@@ -20,103 +19,6 @@ 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):
"""
@@ -262,23 +164,20 @@ class DateparserQueryAnalyzer(QueryAnalyzer):
if not results:
return QueryAnalysis(temporal_constraint=None)
# 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)
# 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)
for text, date in results
if not is_embedded_cjk_dateparser_match(query, text)
if (text.lower() not in false_positives or len(text) > 3)
and not is_embedded_cjk_dateparser_match(query, text)
]
scored_results = [entry for entry in scored_results if entry[0] > 0]
if not scored_results:
if not valid_results:
return QueryAnalysis(temporal_constraint=None)
# Highest signal score wins; ties broken by the longest matched span.
_, _, parsed_date = max(scored_results, key=lambda entry: (entry[0], entry[1]))
# Use the first valid date found
_, parsed_date = valid_results[0]
# Create constraint for single day
start_date = parsed_date.replace(hour=0, minute=0, second=0, microsecond=0)
@@ -15,7 +15,6 @@ import time
from typing import TYPE_CHECKING, Any, Awaitable, Callable
from ...config import get_config
from ..llm_interface import LLM_TOOL_CHOICE_AUTO, LLMToolChoice
from .models import DirectiveInfo, LLMCall, ReflectAgentResult, StructuredOutputResult, TokenUsageSummary, ToolCall
from .prompts import (
_extract_directive_rules,
@@ -50,11 +49,6 @@ 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.
@@ -230,7 +224,6 @@ 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.
@@ -239,10 +232,6 @@ 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
@@ -333,8 +322,6 @@ OUTPUT:"""
],
response_format=DynamicModel,
scope="reflect_structured",
strict_schema=get_config().llm_strict_schema_reflect,
max_completion_tokens=max_tokens,
max_retries=1,
initial_backoff=0.25,
max_backoff=1.0,
@@ -426,104 +413,7 @@ 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,
@@ -544,12 +434,6 @@ async def _run_reflect_agent_inner(
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.
@@ -577,6 +461,7 @@ async def _run_reflect_agent_inner(
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
@@ -613,68 +498,27 @@ async def _run_reflect_agent_inner(
{"role": "user", "content": query},
]
# 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
# 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():
try:
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
)
cached_prefix_name = await provider_impl.get_or_create_cached_prefix(
system_instruction=system_prompt,
tools=tools,
)
except Exception:
return
pending_cache_boundary = upto
pending_cache_task = task
cache_tasks.append(task)
# Caching is a soft optimisation; never let a cache-side
# error block a reflect.
cached_prefix_name = None
# Tracking
total_tools_called = 0
@@ -796,7 +640,7 @@ async def _run_reflect_agent_inner(
# 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, max_tokens)
struct = await _generate_structured_output(answer, response_schema, llm_config, reflect_id)
structured_output = struct.structured_output
total_input_tokens += struct.input_tokens
total_output_tokens += struct.output_tokens
@@ -860,7 +704,7 @@ async def _run_reflect_agent_inner(
structured_output = None
if response_schema and answer:
struct = await _generate_structured_output(answer, response_schema, llm_config, reflect_id, max_tokens)
struct = await _generate_structured_output(answer, response_schema, llm_config, reflect_id)
structured_output = struct.structured_output
total_input_tokens += struct.input_tokens
total_output_tokens += struct.output_tokens
@@ -895,35 +739,12 @@ async def _run_reflect_agent_inner(
if stop_forcing_from_iteration is not None and iteration >= stop_forcing_from_iteration:
# A fresh mental model already short-circuited the forced path.
iter_tool_choice = LLM_TOOL_CHOICE_AUTO
iter_tool_choice: str | dict = "auto"
elif iteration < len(forced_sequence):
iter_tool_choice = LLMToolChoice.named(forced_sequence[iteration])
iter_tool_choice = {"type": "function", "function": {"name": forced_sequence[iteration]}}
else:
iter_tool_choice = LLM_TOOL_CHOICE_AUTO
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 is LLM_TOOL_CHOICE_AUTO:
await _resolve_pending_cache()
call_msg_count = len(messages)
try:
ct_kwargs: dict[str, Any] = dict(
messages=messages,
@@ -931,9 +752,15 @@ async def _run_reflect_agent_inner(
scope="reflect_tool_call",
tool_choice=iter_tool_choice,
)
if incremental_caching and iter_tool_choice is LLM_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
# 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
result = await llm_config.call_with_tools(**ct_kwargs)
llm_duration = int((time.time() - llm_start) * 1000)
consecutive_errors = 0
@@ -1004,7 +831,7 @@ async def _run_reflect_agent_inner(
# 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, max_tokens)
struct = await _generate_structured_output(answer, response_schema, llm_config, reflect_id)
structured_output = struct.structured_output
total_input_tokens += struct.input_tokens
total_output_tokens += struct.output_tokens
@@ -1081,9 +908,7 @@ async def _run_reflect_agent_inner(
# 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, max_tokens
)
struct = await _generate_structured_output(answer, response_schema, llm_config, reflect_id)
structured_output = struct.structured_output
total_input_tokens += struct.input_tokens
total_output_tokens += struct.output_tokens
@@ -1138,7 +963,7 @@ async def _run_reflect_agent_inner(
# 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, max_tokens)
struct = await _generate_structured_output(answer, response_schema, llm_config, reflect_id)
structured_output = struct.structured_output
total_input_tokens += struct.input_tokens
total_output_tokens += struct.output_tokens
@@ -1176,6 +1001,7 @@ async def _run_reflect_agent_inner(
{
"role": "tool",
"tool_call_id": done_call.id,
"name": done_call.name, # Required by Gemini
"content": json.dumps(
{
"error": "You must search for information first. Use search_mental_models(), search_observations(), or recall() before providing your final answer."
@@ -1209,7 +1035,6 @@ async def _run_reflect_agent_inner(
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)
@@ -1241,6 +1066,7 @@ async def _run_reflect_agent_inner(
{
"role": "tool",
"tool_call_id": tc.id,
"name": tc.name,
"content": json.dumps(
{
"error": f"Tool '{_normalize_tool_name(tc.name)}' is not available. Use only the tools provided to you."
@@ -1252,16 +1078,6 @@ async def _run_reflect_agent_inner(
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(
@@ -1344,6 +1160,7 @@ async def _run_reflect_agent_inner(
{
"role": "tool",
"tool_call_id": tc.id,
"name": tc.name, # Required by Gemini
"content": json.dumps(output, default=str, ensure_ascii=False),
}
)
@@ -1427,7 +1244,6 @@ 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
@@ -1436,46 +1252,7 @@ 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_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,
)
)
answer = "No answer provided."
# 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]
@@ -1484,16 +1261,17 @@ 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, max_tokens)
struct = await _generate_structured_output(answer, response_schema, llm_config, reflect_id)
structured_output = struct.structured_output
# Add structured output tokens to usage
final_usage = TokenUsageSummary(
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,
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,
)
log_completion(answer, iterations)
@@ -1608,35 +1386,22 @@ async def _execute_tool(
query = args.get("query")
if not query:
return {"error": "search_mental_models requires a query parameter"}
max_results, error = _parse_tool_int_arg_or_error(args, "max_results", default=5)
if error:
return {"error": error}
max_results = int(args.get("max_results") or 5)
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, error = _parse_tool_int_arg_or_error(args, "max_tokens", default=5000, minimum=1000)
if error:
return {"error": error}
max_tokens = max(int(args.get("max_tokens") or 5000), 1000) # Default 5000, min 1000
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, 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}
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
return await recall_fn(query, max_tokens, max_chunk_tokens)
elif tool_name == "expand":
@@ -1650,63 +1415,23 @@ 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_preview = _summarize_tool_query(args)
max_results = _summarize_tool_int_arg(args, "max_results", default=5)
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)
return f"(query={query_preview}, max_results={max_results})"
elif tool_name == "search_observations":
query_preview = _summarize_tool_query(args)
max_tokens = _summarize_tool_int_arg(args, "max_tokens", default=5000, minimum=1000)
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)
return f"(query={query_preview}, max_tokens={max_tokens})"
elif tool_name == "recall":
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)
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)
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,17 +177,20 @@ _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*$")
_SEPARATOR_RX = re.compile(r"\s*([-*_])\1{2,}\s*")
def _strip_separators(lines: list[str]) -> list[str]:
"""Drop horizontal-rule lines (`---`, `***`) used as section separators.
Our renderer never emits these, but LLM output frequently includes them
between sections; treating them as blank lines avoids parsing them as
paragraphs.
"""
return ["" if re.fullmatch(r"\s*([-*_])\1{2,}\s*", line) else line for line in lines]
def _split_blocks(lines: list[str]) -> list[list[str]]:
"""Group consecutive non-blank lines into block chunks.
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.
"""
"""Group consecutive non-blank lines into block chunks."""
chunks: list[list[str]] = []
current: list[str] = []
in_fence = False
@@ -199,7 +202,7 @@ def _split_blocks(lines: list[str]) -> list[list[str]]:
if in_fence:
current.append(line)
continue
if line.strip() == "" or _SEPARATOR_RX.fullmatch(line):
if line.strip() == "":
if current:
chunks.append(current)
current = []
@@ -247,7 +250,8 @@ def parse_markdown(markdown: str) -> StructuredDocument:
so we never silently drop user content. Section IDs are unique slugs of
their headings.
"""
lines = (markdown or "").splitlines()
raw_lines = (markdown or "").splitlines()
lines = _strip_separators(raw_lines)
sections: list[Section] = []
used_ids: set[str] = set()
@@ -328,21 +328,18 @@ async def tool_expand(
if not memory_ids:
return {"error": "memory_ids is required and must not be empty"}
# 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] = {}
# Validate and convert UUIDs
valid_uuids: list[uuid.UUID] = []
errors: dict[str, str] = {}
for mid in memory_ids:
try:
uuid_by_id[mid] = uuid.UUID(mid)
valid_uuids.append(uuid.UUID(mid))
except ValueError:
errors[mid] = f"Invalid memory_id format: {mid}"
if not uuid_by_id:
if not valid_uuids:
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"""
@@ -398,12 +395,12 @@ async def tool_expand(
# Build results
results: list[dict[str, Any]] = []
for mid in memory_ids:
for mid, mem_uuid in zip(memory_ids, valid_uuids):
if mid in errors:
results.append({"memory_id": mid, "error": errors[mid]})
continue
memory = memory_map.get(uuid_by_id[mid])
memory = memory_map.get(mem_uuid)
if not memory:
results.append({"memory_id": mid, "error": f"Memory not found: {mid}"})
continue
@@ -255,20 +255,13 @@ 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).
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"}.
"""
"""Parse metadata from JSON string if needed (asyncpg may return JSONB as str)."""
if v is None:
return None
if isinstance(v, str):
import json
v = json.loads(v)
if isinstance(v, dict):
return {str(k): str(val) for k, val in v.items()}
return json.loads(v)
return v
chunk_id: str | None = Field(
@@ -189,8 +189,7 @@ 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:
async with conn.transaction():
return await get_or_create_bank_profile_on_conn(conn, bank_id, ops=pool.ops)
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,53 +64,8 @@ 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"""
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
""",
f"DELETE FROM {fq_table('chunks')} WHERE chunk_id = ANY($1::text[])",
chunk_ids,
)
@@ -7,7 +7,7 @@ Handles entity extraction and resolution for stored facts.
import logging
from . import link_utils
from .types import EntityResolutionResult, ProcessedFact
from .types import ProcessedFact
logger = logging.getLogger(__name__)
@@ -58,7 +58,7 @@ async def resolve_entities(
log_buffer: list[str] = None,
user_entities_per_content: dict[int, list[dict]] = None,
entity_labels: list | None = None,
) -> EntityResolutionResult:
) -> tuple[list[str], list[tuple], dict[str, list[str]]]:
"""
Phase 1: Resolve entity names to canonical IDs (read-heavy).
@@ -76,10 +76,10 @@ async def resolve_entities(
entity_labels: Optional entity label taxonomy
Returns:
EntityResolutionResult with the resolved identities and unit mappings.
Tuple of (resolved_entity_ids, entity_to_unit, unit_to_entity_ids).
"""
if not unit_ids or not facts:
return EntityResolutionResult(resolved_entities=[], entity_to_unit=[], unit_to_entity_ids={})
return [], [], {}
if len(unit_ids) != len(facts):
raise ValueError(f"Mismatch between unit_ids ({len(unit_ids)}) and facts ({len(facts)})")
@@ -15,10 +15,9 @@ 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, parse_llm_json, sanitize_llm_output
from ..llm_wrapper import LLMConfig, OutputTooLongError, sanitize_llm_output
from ..operation_metadata import RetainExtractionErrors
from ..response_models import TokenUsage
from ..structured_output import strict_json_schema
from .entity_labels import (
EntityLabelsConfig,
MapField,
@@ -33,7 +32,7 @@ def _extract_map_entities(
entity_obj: dict,
fields: dict[str, MapField],
prefix: str,
validated_entities: list[str],
validated_entities: "list[Entity]",
existing_texts_lower: set[str],
) -> None:
"""Recursively extract key:field:value entity strings from a map entity dict."""
@@ -60,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(label_str)
validated_entities.append(Entity(text=label_str))
existing_texts_lower.add(label_str.lower())
else:
# text or value — single string
@@ -68,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(label_str)
validated_entities.append(Entity(text=label_str))
existing_texts_lower.add(label_str.lower())
@@ -115,33 +114,12 @@ def _sanitize_text(text: str | None) -> str | None:
return sanitize_llm_output(text)
def _coerce_entity_strings(v: Any) -> Any:
"""
Normalize the LLM's `entities` field to a plain list of strings.
class Entity(BaseModel):
"""An entity extracted from text."""
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
text: str = Field(
description="The specific, named entity as it appears in the fact. Must be a proper noun or specific identifier."
)
class Fact(BaseModel):
@@ -166,7 +144,7 @@ class Fact(BaseModel):
)
# Optional structured data
entities: list[str] | None = None
entities: list[Entity] | None = None
causal_relations: list["CausalRelation"] | None = None
@@ -217,9 +195,7 @@ 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[str] = Field(
default_factory=list, description='People, places, concepts - plain strings, e.g. ["Alice", "Kubernetes"]'
)
entities: list[Entity] | None = Field(default=None, description="People, places, concepts")
causal_relations: list[FactCausalRelation] | None = Field(
default=None, description="Links to previous facts (target_index < this fact's index)"
)
@@ -227,7 +203,10 @@ class ExtractedFact(BaseModel):
@field_validator("entities", mode="before")
@classmethod
def ensure_entities_list(cls, v):
return _coerce_entity_strings(v)
"""Ensure entities is always a list (convert None to empty list)."""
if v is None:
return []
return v
def build_fact_text(self) -> str:
"""Combine all dimensions into a single comprehensive fact string."""
@@ -253,55 +232,6 @@ 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."""
@@ -369,9 +299,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[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.",
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.",
)
causal_relations: list[FactCausalRelation] | None = Field(
@@ -383,7 +313,9 @@ class ExtractedFactVerbose(BaseModel):
@field_validator("entities", mode="before")
@classmethod
def ensure_entities_list(cls, v):
return _coerce_entity_strings(v)
if v is None:
return []
return v
class FactExtractionResponseVerbose(BaseModel):
@@ -416,15 +348,17 @@ 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[str] = Field(
default_factory=list,
description='Named entities, objects, and concepts from the fact, as plain strings (e.g. ["Alice", "Kubernetes"]).',
entities: list[Entity] | None = Field(
default=None,
description="Named entities, objects, and concepts from the fact.",
)
@field_validator("entities", mode="before")
@classmethod
def ensure_entities_list(cls, v):
return _coerce_entity_strings(v)
if v is None:
return []
return v
class FactExtractionResponseNoCausal(BaseModel):
@@ -456,14 +390,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[str] = Field(
default_factory=list, description='People, places, concepts - plain strings, e.g. ["Alice", "Kubernetes"]'
)
entities: list[Entity] | None = Field(default=None, description="People, places, concepts")
@field_validator("entities", mode="before")
@classmethod
def ensure_entities_list(cls, v):
return _coerce_entity_strings(v)
if v is None:
return []
return v
class VerbatimFactExtractionResponse(BaseModel):
@@ -747,11 +681,6 @@ 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}"""
@@ -1169,8 +1098,8 @@ def _build_extraction_prompt_and_schema(config) -> tuple[str, type]:
}
if not free_form_entities:
dynamic_fields["entities"] = (
list[str],
Field(default_factory=list, description="Leave empty — labels-only mode"),
list[Entity] | None,
Field(default=None, description="Leave empty — labels-only mode"),
)
# Inherit parent's required fields and add 'labels' so it appears in the JSON schema
# required array (the base class json_schema_extra overrides required entirely)
@@ -1294,32 +1223,19 @@ def _build_request_body(llm_config, config, prompt: str, user_message: str, resp
request_body["service_tier"] = llm_config._provider_impl.openai_service_tier
# Add response_format (JSON schema). The batch path builds the request body
# directly instead of going through LLMProvider.call(), so resolve the
# strict-schema flag here too: strict=True grammar-enforces the output on capable
# backends rather than relying on the model to emit clean JSON. Reads the
# retain-scoped field, which already folds in the global HINDSIGHT_API_LLM_STRICT_SCHEMA
# fallback, so the batch and streaming paths can't disagree.
# directly instead of going through LLMProvider.call(), so honour
# HINDSIGHT_API_LLM_STRICT_SCHEMA here too: strict=True grammar-enforces the
# output on capable backends rather than relying on the model to emit clean JSON.
if hasattr(response_schema, "model_json_schema"):
schema = (
strict_json_schema(response_schema) if config.llm_strict_schema else response_schema.model_json_schema()
)
schema = response_schema.model_json_schema()
request_body["response_format"] = {
"type": "json_schema",
"json_schema": {"name": "facts", "schema": schema, "strict": config.llm_strict_schema_retain},
"json_schema": {"name": "facts", "schema": schema, "strict": config.llm_strict_schema},
}
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,
@@ -1388,15 +1304,10 @@ async def _extract_facts_from_chunk(
llm_max_retries = (
config.retain_llm_max_retries if config.retain_llm_max_retries is not None else config.llm_max_retries
)
# OUTER content-validation attempts (re-prompts on malformed JSON). Follows the
# same `N + 1` convention as the providers' transport-retry loops — N retries after
# the initial request — so a zero budget still performs one request (#2731). The raw
# budget is forwarded unchanged to llm_config.call(), which owns transport retries.
outer_attempts = llm_max_retries + 1
last_error: Exception | None = None
usage = TokenUsage() # Track cumulative usage across retries
for attempt in range(outer_attempts):
for attempt in range(llm_max_retries):
try:
initial_backoff = (
config.retain_llm_initial_backoff
@@ -1412,7 +1323,6 @@ async def _extract_facts_from_chunk(
response_format=response_schema,
scope="retain_extract_facts",
temperature=config.llm_temperature_retain,
strict_schema=config.llm_strict_schema_retain,
max_completion_tokens=config.retain_max_completion_tokens,
max_retries=llm_max_retries,
initial_backoff=initial_backoff,
@@ -1431,11 +1341,10 @@ async def _extract_facts_from_chunk(
has_malformed_facts = False
# Handle malformed LLM responses
coerced_response_json = _coerce_fact_response(extraction_response_json)
if coerced_response_json is None:
if attempt < outer_attempts - 1:
if not isinstance(extraction_response_json, dict):
if attempt < llm_max_retries - 1:
logger.warning(
f"LLM returned non-dict JSON on attempt {attempt + 1}/{outer_attempts}: {type(extraction_response_json).__name__}. Retrying..."
f"LLM returned non-dict JSON on attempt {attempt + 1}/{llm_max_retries}: {type(extraction_response_json).__name__}. Retrying..."
)
continue
else:
@@ -1444,10 +1353,9 @@ async def _extract_facts_from_chunk(
# worker's retry machinery and ultimately fails loudly — never
# silently commit the document with 0 facts. See issue #1833.
raise RuntimeError(
f"Fact extraction failed: LLM returned non-dict JSON after {outer_attempts} attempts "
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", [])
@@ -1543,9 +1451,21 @@ async def _extract_facts_from_chunk(
elif fact_data.get("occurred_start"):
fact_data["occurred_end"] = fact_data["occurred_start"]
# 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"))
# 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}")
# Post-process label entities from structured labels object
entity_labels_raw = getattr(config, "entity_labels", None)
@@ -1555,7 +1475,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.lower() for e in validated_entities}
existing_texts_lower = {e.text.lower() for e in validated_entities}
for group in labels_cfg.attributes:
value = labels_data.get(group.key)
if not value:
@@ -1580,12 +1500,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(label_str)
validated_entities.append(Entity(text=label_str))
existing_texts_lower.add(label_str.lower())
elif (
label_str.lower() in labels_lookup and label_str.lower() not in existing_texts_lower
):
validated_entities.append(label_str)
validated_entities.append(Entity(text=label_str))
existing_texts_lower.add(label_str.lower())
else:
logger.warning(f"Label '{label_str}' not in valid label values, skipping")
@@ -1593,7 +1513,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, labels_cfg, labels_lookup)
e for e in validated_entities if is_label_entity(e.text, labels_cfg, labels_lookup)
]
elif not free_form_entities:
# No labels but free_form disabled: clear all entities
@@ -1651,9 +1571,9 @@ async def _extract_facts_from_chunk(
continue
# If we got malformed facts and haven't exhausted retries, try again
if has_malformed_facts and len(chunk_facts) < len(raw_facts) * 0.8 and attempt < outer_attempts - 1:
if has_malformed_facts and len(chunk_facts) < len(raw_facts) * 0.8 and attempt < llm_max_retries - 1:
logger.warning(
f"Got {len(raw_facts) - len(chunk_facts)} malformed facts out of {len(raw_facts)} on attempt {attempt + 1}/{outer_attempts}. Retrying..."
f"Got {len(raw_facts) - len(chunk_facts)} malformed facts out of {len(raw_facts)} on attempt {attempt + 1}/{llm_max_retries}. Retrying..."
)
continue
@@ -1692,7 +1612,7 @@ async def _extract_facts_from_chunk(
# If we exhausted all retries, raise the last error or a descriptive fallback
if last_error is not None:
raise last_error
raise RuntimeError(f"Fact extraction failed after {outer_attempts} attempts: LLM did not return valid JSON")
raise RuntimeError(f"Fact extraction failed after {llm_max_retries} attempts: LLM did not return valid JSON")
async def _extract_facts_with_auto_split(
@@ -1744,22 +1664,33 @@ async def _extract_facts_with_auto_split(
metadata=metadata,
)
except OutputTooLongError:
# Output exceeded token limits - split the chunk and retry. Conversation
# chunks are JSON arrays, so preserve array/turn boundaries when possible.
# Output exceeded token limits - split the chunk in half and retry
logger.warning(
f"Output too long for chunk {chunk_index + 1}/{total_chunks} "
f"({len(chunk)} chars). Splitting and retrying..."
f"({len(chunk)} chars). Splitting in half and retrying..."
)
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()
# Split at the midpoint, preferring sentence boundaries
mid_point = len(chunk) // 2
first_half, second_half = split_chunks
# 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()
logger.info(
f"Split chunk {chunk_index + 1} into two sub-chunks: {len(first_half)} chars and {len(second_half)} chars"
@@ -2201,10 +2132,7 @@ async def extract_facts_from_contents_batch_api(
content_str = message.get("content", "{}")
try:
# #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)
extraction_response_json = json.loads(content_str)
except json.JSONDecodeError as e:
message = f"{custom_id}: failed to parse JSON: {e}"
logger.error(message)
@@ -2216,19 +2144,6 @@ 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 = []
@@ -2295,9 +2210,18 @@ async def extract_facts_from_contents_batch_api(
elif fact_data.get("occurred_start"):
fact_data["occurred_end"] = fact_data["occurred_start"]
# 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"))
# 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
# Post-process label entities from structured labels object
entity_labels_raw = getattr(config, "entity_labels", None)
@@ -2307,7 +2231,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.lower() for e in validated_entities}
existing_texts_lower = {e.text.lower() for e in validated_entities}
for group in labels_cfg_batch.attributes:
value = labels_data.get(group.key)
if not value:
@@ -2332,18 +2256,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(label_str)
validated_entities.append(Entity(text=label_str))
existing_texts_lower.add(label_str.lower())
elif (
label_str.lower() in labels_lookup_batch
and label_str.lower() not in existing_texts_lower
):
validated_entities.append(label_str)
validated_entities.append(Entity(text=label_str))
existing_texts_lower.add(label_str.lower())
if not free_form_entities_batch:
validated_entities = [
e for e in validated_entities if is_label_entity(e, labels_cfg_batch, labels_lookup_batch)
e for e in validated_entities if is_label_entity(e.text, labels_cfg_batch, labels_lookup_batch)
]
elif not free_form_entities_batch:
validated_entities = []
@@ -2430,7 +2354,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=list(fact_from_llm.entities or []),
entities=[e.text for e in (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),
@@ -2627,7 +2551,7 @@ async def extract_facts_from_contents(
extracted_fact = ExtractedFactType(
fact_text=fact_from_llm.fact,
fact_type=fact_from_llm.fact_type,
entities=list(fact_from_llm.entities or []),
entities=[e.text for e in (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,9 +74,7 @@ async def create_causal_links_batch(
"""
Create causal links between facts.
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.
Links facts that have causal relationships (causes, enables, prevents).
Args:
conn: Database connection
@@ -92,7 +90,22 @@ 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)})")
causal_relations_per_fact = [fact.causal_relations or [] for fact in 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([])
link_count = await link_utils.create_causal_links_batch(conn, bank_id, unit_ids, causal_relations_per_fact, ops=ops)
@@ -7,11 +7,7 @@ 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, EntityResolutionResult
logger = logging.getLogger(__name__)
@@ -304,7 +300,7 @@ async def resolve_entities_only(
llm_entities: list[list[dict]],
log_buffer: list[str] = None,
entity_labels: list | None = None,
) -> EntityResolutionResult:
) -> tuple[list[str], list[tuple], dict[str, list[str]]]:
"""
Phase 1 of entity processing: resolve entity names to canonical IDs.
@@ -325,10 +321,10 @@ async def resolve_entities_only(
entity_labels: Optional entity label taxonomy
Returns:
EntityResolutionResult carrying the resolved entity identities (id +
stored canonical name, in flattened order), the flat-index unit map,
and the unit entity-id map used to remap placeholder unit IDs in
Phase 2.
Tuple of (resolved_entity_ids, entity_to_unit, unit_to_entity_ids) where:
- resolved_entity_ids: list of entity IDs in same order as flattened entities
- entity_to_unit: maps flat index to (unit_id, local_index, fact_date)
- unit_to_entity_ids: maps unit_id to list of resolved entity IDs
"""
all_entities_flat, _all_entities, entity_to_unit = _prepare_entities_for_resolution(
unit_ids, sentences, fact_dates, llm_entities, log_buffer
@@ -336,10 +332,10 @@ async def resolve_entities_only(
if not all_entities_flat:
_log(log_buffer, " [6.2] Entity resolution (batched): 0 entities", level="debug")
return EntityResolutionResult(resolved_entities=[], entity_to_unit=[], unit_to_entity_ids={})
return [], [], {}
step_start = time.time()
resolved_entities = await entity_resolver.resolve_entities_batch(
resolved_entity_ids = await entity_resolver.resolve_entities_batch(
bank_id=bank_id,
entities_data=all_entities_flat,
context=context,
@@ -358,7 +354,7 @@ async def resolve_entities_only(
for idx, (unit_id, _local_idx, _fact_date) in enumerate(entity_to_unit):
if unit_id not in unit_to_entity_ids:
unit_to_entity_ids[unit_id] = []
unit_to_entity_ids[unit_id].append(resolved_entities[idx].entity_id)
unit_to_entity_ids[unit_id].append(resolved_entity_ids[idx])
_log(
log_buffer,
@@ -366,11 +362,7 @@ async def resolve_entities_only(
level="debug",
)
return EntityResolutionResult(
resolved_entities=resolved_entities,
entity_to_unit=entity_to_unit,
unit_to_entity_ids=unit_to_entity_ids,
)
return resolved_entity_ids, entity_to_unit, unit_to_entity_ids
async def create_temporal_links_batch_per_fact(
@@ -779,61 +771,28 @@ async def create_semantic_links_batch(
async def create_causal_links_batch(
conn: DatabaseConnection,
conn,
bank_id: str,
unit_ids: list[str],
causal_relations_per_fact: list[list[CausalRelation]],
ops: DataAccessOps | None = None,
causal_relations_per_fact: list[list[dict]],
ops=None,
) -> int:
"""Create canonical causal links for the retain pipeline.
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,
)
Create causal links between facts based on LLM-extracted causal relationships.
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.
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"
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
@@ -850,13 +809,15 @@ async def _write_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"]
if relation_type not in allowed_relation_types:
# Validate relation_type - only "caused_by" is supported (DB constraint)
valid_types = {"caused_by"}
if relation_type not in valid_types:
logger.error(
f"Invalid relation_type '{relation_type}' (type: {type(relation_type).__name__}) "
f"from fact {fact_idx}. Must be one of: {allowed_relation_types}. "
f"from fact {fact_idx}. Must be one of: {valid_types}. "
f"Relation data: {relation}"
)
continue
@@ -143,7 +143,7 @@ async def _fire_memory_defense_webhook(
logger.warning("memory_defense webhook delivery failed", exc_info=True)
async def _audit_memory_defense(
def _audit_memory_defense(
audit_logger: Any,
*,
bank_id: str,
@@ -152,15 +152,11 @@ async def _audit_memory_defense(
) -> None:
"""Write a fire-and-forget ``memory_defense`` audit entry for a non-allow decision.
No-op when auditing is off for this bank. ``audit_log_enabled`` is per-bank
overridable, so the decision must be awaited here rather than relying on the
logger's synchronous allowlist check alone.
No-op when audit logging is disabled (the logger gates on its own config).
The action taken (redact/block) and what matched live in the entry metadata.
"""
if audit_logger is None:
return
if not await audit_logger.should_log("memory_defense", bank_id):
return
from ..audit import AuditEntry
entry = AuditEntry(
@@ -250,12 +246,10 @@ from . import (
link_creation,
)
from .types import (
CausalRelation,
ChunkMetadata,
ExtractedFact,
EntityResolutionResult,
Phase1Result,
ProcessedFact,
ResolvedEntity,
RetainContent,
RetainContentDict,
)
@@ -266,15 +260,6 @@ RetainOutboxCallback = Callable[[asyncpg.Connection], Awaitable[None]]
RetainOutboxCallbackFactory = Callable[[list[RetainContentDict]], RetainOutboxCallback | None]
@dataclass
class _ProcessedFactBatch:
"""Aligned survivors from converting extracted facts for storage."""
extracted_facts: list[ExtractedFact]
processed_facts: list[ProcessedFact]
retained_index_by_original: list[int | None]
def _resolve_narrator(profile_name: str, bank_id: str) -> str | None:
"""Resolve the narrator (memory owner) used to prime fact extraction.
@@ -359,7 +344,7 @@ async def _pre_resolve_phase1(
embeddings = [fact.embedding for fact in processed_facts]
async with acquire_with_retry(pool) as resolve_conn:
entity_resolution = await entity_processing.resolve_entities(
resolved_entity_ids, entity_to_unit, unit_to_entity_ids = await entity_processing.resolve_entities(
entity_resolver,
resolve_conn,
bank_id,
@@ -381,7 +366,11 @@ async def _pre_resolve_phase1(
)
return Phase1Result(
entities=entity_resolution,
entities=EntityResolutionResult(
resolved_entity_ids=resolved_entity_ids,
entity_to_unit=entity_to_unit,
unit_to_entity_ids=unit_to_entity_ids,
),
semantic_ann_links=semantic_ann_links,
)
@@ -432,7 +421,7 @@ async def _insert_facts_and_links(
processed_facts: list[ProcessedFact],
config,
log_buffer: list[str],
resolved_entities: list[ResolvedEntity],
resolved_entity_ids: list[str],
entity_to_unit: list[tuple],
unit_to_entity_ids: dict[str, list[str]],
semantic_ann_links: list[tuple],
@@ -459,7 +448,6 @@ async def _insert_facts_and_links(
# Entity resolution was done in Phase 1 (separate connection).
# Remap placeholder IDs to actual unit IDs.
step_start = time.time()
resolved_entity_ids = [entity.entity_id for entity in resolved_entities]
remapped_entity_to_unit, _remapped_unit_to_entity_ids, remapped_semantic = _remap_phase1_results(
resolved_entity_ids, entity_to_unit, unit_to_entity_ids, semantic_ann_links or [], unit_ids
)
@@ -472,10 +460,6 @@ async def _insert_facts_and_links(
(unit_id, resolved_entity_ids[idx], fact_date)
for idx, (unit_id, _local_idx, fact_date) in enumerate(remapped_entity_to_unit)
]
# Lock/re-create the resolved parents on THIS transaction before linking,
# closing the window where prune_orphan_entities could have deleted one
# between Phase-1 resolution and this insert (#2662).
await entity_resolver.reassert_entities_batch(bank_id, resolved_entities, conn=conn)
await entity_resolver.link_units_to_entities_batch(unit_entity_pairs, conn=conn)
log_buffer.append(f" Insert unit_entities: {len(unit_entity_pairs)} pairs in {time.time() - step_start:.3f}s")
@@ -565,90 +549,9 @@ 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")
fact_batch = _process_extracted_facts(extracted_facts, embeddings)
processed_facts = [ProcessedFact.from_extracted_fact(ef, emb) for ef, emb in zip(extracted_facts, embeddings)]
return fact_batch.extracted_facts, fact_batch.processed_facts, chunks, usage
def _remap_causal_relations(
relations_per_fact: list[list[CausalRelation]],
retained_index_by_original: list[int | None],
) -> list[list[CausalRelation]]:
"""Remap a causal relation matrix after facts have been filtered.
Both the source row and each ``target_fact_index`` use fact ordinals. A
rejected source disappears with its row; a relation to a rejected target
must disappear rather than silently pointing at the next surviving fact.
"""
remapped = [[] for retained_index in retained_index_by_original if retained_index is not None]
for original_source, retained_source in enumerate(retained_index_by_original):
if retained_source is None:
continue
for relation in relations_per_fact[original_source]:
original_target = relation.target_fact_index
retained_target = (
retained_index_by_original[original_target]
if 0 <= original_target < len(retained_index_by_original)
else None
)
if retained_target is None:
continue
remapped[retained_source].append(
CausalRelation(
relation_type=relation.relation_type,
target_fact_index=retained_target,
)
)
return remapped
def _process_extracted_facts(
extracted_facts: list[ExtractedFact],
embeddings: list[list[float]],
) -> _ProcessedFactBatch:
"""Process facts while preserving their positional relationships.
``ProcessedFact.from_extracted_fact`` may reject a degenerate fact. Keep
the surviving extracted and processed facts in lockstep, and translate
causal ordinals from the original extraction into that retained sequence.
The returned index table is also used by transfer import for archive-only
links and observation source references.
"""
if len(extracted_facts) != len(embeddings):
raise ValueError(
f"Extracted facts/embeddings length mismatch: {len(extracted_facts)} facts, {len(embeddings)} embeddings"
)
retained_extracted: list[ExtractedFact] = []
processed_facts: list[ProcessedFact] = []
retained_index_by_original: list[int | None] = [None] * len(extracted_facts)
for original_index, (extracted_fact, embedding) in enumerate(zip(extracted_facts, embeddings, strict=True)):
processed_fact = ProcessedFact.from_extracted_fact(extracted_fact, embedding)
if processed_fact is None:
continue
retained_index_by_original[original_index] = len(processed_facts)
retained_extracted.append(extracted_fact)
processed_facts.append(processed_fact)
remapped_relations = _remap_causal_relations(
[fact.causal_relations for fact in extracted_facts],
retained_index_by_original,
)
for extracted_fact, processed_fact, causal_relations in zip(
retained_extracted,
processed_facts,
remapped_relations,
strict=True,
):
extracted_fact.causal_relations = causal_relations
processed_fact.causal_relations = causal_relations
return _ProcessedFactBatch(
extracted_facts=retained_extracted,
processed_facts=processed_facts,
retained_index_by_original=retained_index_by_original,
)
return extracted_facts, processed_facts, chunks, usage
async def retain_batch(
@@ -829,7 +732,7 @@ async def retain_batch(
document_id=_item_doc_id,
decision=_decision,
)
await _audit_memory_defense(
_audit_memory_defense(
audit_logger,
bank_id=bank_id,
document_id=_item_doc_id,
@@ -928,12 +831,6 @@ 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]
@@ -955,12 +852,6 @@ 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):
@@ -1313,14 +1204,6 @@ async def _streaming_retain_batch(
retain_params, merged_tags = _build_retain_params(contents_dicts, document_tags)
# Track whether document tracking has been done (by the first batch)
doc_tracking_done = [False]
# Track whether the transactional-outbox callback has already fired inside a
# batch write TXN. The in-TXN fire only runs on a final facts-bearing batch
# (is_last=True); two success paths never reach it — a committed-chunk count
# that lands exactly on a chunk_batch_size boundary (the sentinel drains an
# empty batch), and a final batch that extracts zero facts (it returns before
# the insert). A post-loop fallback fires the callback in those cases, so this
# flag exists to guarantee the callback fires exactly once.
outbox_fired = [False]
# ---------------------------------------------------------------------------
# Producer-consumer pipeline: LLM extraction runs concurrently with DB writes
@@ -1493,26 +1376,12 @@ async def _streaming_retain_batch(
# from a single oversized item sharing one document_id — without it
# each sub-batch restarts at 0 and their chunk_ids collide (#1888).
doc_chunk_index = global_idx + chunk_index_offset
fact_index_offset = len(batch_processed)
for fact, processed_fact in zip(extracted, processed, strict=True):
for fact in extracted:
fact.content_index = content_idx_in_batch
if fact.chunk_index is not None:
fact.chunk_index = doc_chunk_index
processed_fact.content_index = content_idx_in_batch
# Each producer call extracts one chunk, so its causal ordinals
# start at zero. Translate them into the combined consumer-batch
# sequence before link creation; otherwise later chunks can point
# at equally numbered facts from the first completed chunk.
causal_relations = [
CausalRelation(
relation_type=relation.relation_type,
target_fact_index=relation.target_fact_index + fact_index_offset,
)
for relation in processed_fact.causal_relations
]
fact.causal_relations = causal_relations
processed_fact.causal_relations = causal_relations
for pf in processed:
pf.content_index = content_idx_in_batch
for cm in chunk_meta:
cm.chunk_index = doc_chunk_index
@@ -1525,12 +1394,7 @@ async def _streaming_retain_batch(
nonlocal total_usage
total_usage = total_usage + batch_usage
# ``batch_extracted`` contains only survivors after the degenerate-text
# guard. Chunk metadata still records whether extraction originally
# produced facts, so an all-rejected batch follows the normal write path
# and preserves chunk/outbox behavior from before filtering was added.
had_extracted_facts = bool(batch_extracted) or any(chunk.fact_count for chunk in batch_chunk_meta)
if not had_extracted_facts:
if not batch_extracted:
# Even with 0 facts, the first batch must still run document tracking
# (cascade-delete + insert doc row) to establish ownership and prevent
# concurrent requests from interleaving. Later batches can safely skip.
@@ -1726,7 +1590,7 @@ async def _streaming_retain_batch(
batch_processed,
config,
log_buffer,
resolved_entities=phase1.entities.resolved_entities,
resolved_entity_ids=phase1.entities.resolved_entity_ids,
entity_to_unit=phase1.entities.entity_to_unit,
unit_to_entity_ids=phase1.entities.unit_to_entity_ids,
semantic_ann_links=[],
@@ -1737,12 +1601,6 @@ async def _streaming_retain_batch(
logger.info(f"[streaming] Phase 2 (write txn): {time.time() - p2_start:.3f}s")
# The write TXN above committed the transactional-outbox row in the
# same transaction as this batch's facts. Record it so the post-loop
# fallback doesn't queue a duplicate delivery.
if is_last and outbox_callback is not None:
outbox_fired[0] = True
# Best-effort: flush entity_cooccurrences and other deferred stats.
try:
await entity_resolver.flush_pending_stats()
@@ -1876,21 +1734,6 @@ async def _streaming_retain_batch(
combined_content = ""
log_buffer.append(f"[streaming] Document {effective_doc_id} tracked (no facts extracted)")
# Transactional-outbox fallback. The in-TXN fire only runs on a final
# facts-bearing batch (is_last=True). When the committed-chunk count lands
# exactly on a chunk_batch_size boundary the sentinel drains an empty batch
# and never marks one last; when the final batch extracts zero facts it
# returns before the insert; and when every chunk is skipped as already
# committed no batch runs at all. In each of those the retain still
# succeeded, so the retain.completed delivery must be queued — exactly once,
# in its own transaction (there is no batch TXN left to attach it to). Skip
# it on a concurrent takeover: an aborted request must not emit completion.
if outbox_callback is not None and not outbox_fired[0] and not pipeline_aborted[0]:
async with acquire_with_retry(pool) as conn:
async with conn.transaction():
await outbox_callback(conn)
outbox_fired[0] = True
# Mark facts as committed in operation metadata (crash recovery checkpoint)
if operation_id and all_unit_ids:
try:
@@ -2343,7 +2186,7 @@ async def _try_delta_retain(
processed_facts,
config,
log_buffer,
resolved_entities=phase1.entities.resolved_entities,
resolved_entity_ids=phase1.entities.resolved_entity_ids,
entity_to_unit=phase1.entities.entity_to_unit,
unit_to_entity_ids=phase1.entities.unit_to_entity_ids,
semantic_ann_links=phase1.semantic_ann_links,
@@ -5,14 +5,11 @@ 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.
@@ -99,12 +96,10 @@ class CausalRelation:
"""
Causal relationship between facts.
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.
Represents how one fact was caused by another.
"""
relation_type: str # ``caused_by`` for retain; legacy types for transfer restore
relation_type: str # "caused_by"
target_fact_index: int # Index of the target fact in the batch
@@ -190,47 +185,10 @@ 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 | None":
) -> "ProcessedFact":
"""
Create ProcessedFact from ExtractedFact.
@@ -240,17 +198,8 @@ class ProcessedFact:
chunk_id: Optional chunk ID
Returns:
ProcessedFact ready for storage, or None if the fact text is degenerate
(zero information content punctuation-only, empty, etc.)
ProcessedFact ready for storage
"""
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
@@ -260,7 +209,7 @@ class ProcessedFact:
entities = [EntityRef(name=name) for name in extracted_fact.entities]
return ProcessedFact(
fact_text=fact_text,
fact_text=extracted_fact.fact_text,
fact_type=extracted_fact.fact_type,
embedding=embedding,
occurred_start=occurred_start,
@@ -277,43 +226,19 @@ class ProcessedFact:
)
@dataclass
class ResolvedEntity:
"""Identity of a resolved entity carried across the retain phase boundary.
``canonical_name`` is the value stored on the entity row (NOT the raw input
mention), captured during Phase-1 resolution. It is threaded to Phase 2 so a
parent pruned between phases can be re-created with its real name the row
is gone by then, so the name is otherwise unrecoverable (#2662).
"""
entity_id: str
canonical_name: str
def __post_init__(self) -> None:
# Callers pass UUID objects or strings; normalize once so downstream
# comparisons, set membership, and SQL binds all see a plain str.
self.entity_id = str(self.entity_id)
@dataclass
class EntityResolutionResult:
"""
Result of Phase 1 entity resolution.
Contains resolved entity identities and the mapping data needed to remap
Contains resolved entity IDs and the mapping data needed to remap
placeholder unit IDs to real IDs after fact insertion in Phase 2.
"""
resolved_entities: list[ResolvedEntity]
resolved_entity_ids: list[str]
entity_to_unit: list[tuple]
unit_to_entity_ids: dict[str, list[str]]
@property
def resolved_entity_ids(self) -> list[str]:
"""Entity IDs in flattened resolution order (used by link remapping)."""
return [entity.entity_id for entity in self.resolved_entities]
@dataclass
class Phase1Result:
@@ -27,23 +27,6 @@ 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,6 +40,8 @@ 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)
@@ -57,6 +59,8 @@ 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.
Selects bounded semantic seeds, then expands through three parallel,
first-class signals stored in memory_links:
Expands from semantic/temporal seeds 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,6 +127,8 @@ 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",
@@ -144,6 +146,8 @@ 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
@@ -154,28 +158,32 @@ class LinkExpansionRetriever(GraphRetriever):
timings = GraphRetrievalTimings(fact_type=fact_type)
async with acquire_with_retry(pool) as conn:
# 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})"
)
# 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)
if not all_seeds:
return [], timings
@@ -235,16 +243,12 @@ 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 fact_id in sorted_ids:
row = row_map[fact_id]
for row in rows:
result = RetrievalResult.from_db_row(dict(row))
# ``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]
result.activation = row["score"]
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 DEFAULT_BM25_MAX_QUERY_TERMS, get_config
from ...config import get_config
from ..db_utils import acquire_with_retry
from ..memory_engine import fq_table
from ..sql import create_sql_dialect
@@ -222,12 +222,7 @@ 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,
max_query_terms=getattr(config, "bm25_max_query_terms", DEFAULT_BM25_MAX_QUERY_TERMS),
)
bm25_text_param: str = dialect.prepare_bm25_text(tokens, query_text, text_search_extension=text_ext)
for i, ft in enumerate(fact_types):
arms.append(
dialect.build_bm25_arm(
@@ -621,7 +616,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, mu.proof_count,
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,
l.weight, l.link_type,
1 - (mu.embedding <=> $1::vector) AS similarity
FROM unnest($2::uuid[]) AS src(from_unit_id)
@@ -798,7 +793,7 @@ async def retrieve_all_fact_types_parallel(
tc_start,
tc_end,
budget=thinking_budget,
semantic_threshold=0.1,
semantic_threshold=min_semantic if min_semantic is not None else 0.1,
tags=tags,
tags_match=tags_match,
tag_groups=tag_groups,
@@ -822,6 +817,8 @@ 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 spent selecting semantic graph seeds
seeds_time: float = 0.0 # Time to find semantic seeds (if fallback used)
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)
@@ -300,6 +300,19 @@ class SQLDialect(ABC):
"""FOR UPDATE SKIP LOCKED clause (same on both PG and Oracle)."""
...
@abstractmethod
def advisory_lock(self, id_param: str) -> str:
"""Advisory lock expression.
Args:
id_param: Parameter placeholder for the lock ID.
Returns:
PG: "pg_try_advisory_lock($1)"
Oracle: "SELECT ... FOR UPDATE NOWAIT" equivalent.
"""
...
# -- UUID generation -------------------------------------------------
@abstractmethod
@@ -436,7 +449,6 @@ 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.
@@ -447,8 +459,6 @@ 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.
@@ -203,6 +203,10 @@ class OracleDialect(SQLDialect):
def for_update_skip_locked(self) -> str:
return "FOR UPDATE SKIP LOCKED"
def advisory_lock(self, id_param: str) -> str:
# Oracle doesn't have advisory locks. Use SELECT FOR UPDATE NOWAIT on a lock row.
return "SELECT 1 FROM dual FOR UPDATE NOWAIT"
# -- UUID generation -------------------------------------------------
def generate_uuid(self) -> str:
@@ -299,7 +303,6 @@ 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.
@@ -118,6 +118,9 @@ class PostgreSQLDialect(SQLDialect):
def for_update_skip_locked(self) -> str:
return "FOR UPDATE SKIP LOCKED"
def advisory_lock(self, id_param: str) -> str:
return f"pg_try_advisory_lock({id_param})"
# -- UUID generation -------------------------------------------------
def generate_uuid(self) -> str:
@@ -251,11 +254,8 @@ 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)
@@ -1,30 +0,0 @@
"""Canonical JSON Schema serialization for OpenAI strict output."""
from typing import Any
from pydantic import BaseModel
from pydantic.json_schema import GenerateJsonSchema, JsonSchemaValue
from pydantic_core import core_schema
class OpenAIStrictSchemaGenerator(GenerateJsonSchema):
"""Emit the strict JSON Schema subset required by OpenAI-compatible APIs."""
def model_schema(self, schema: core_schema.ModelSchema) -> JsonSchemaValue:
json_schema = super().model_schema(schema)
properties = json_schema.get("properties")
if type(properties) is dict:
json_schema["required"] = list(properties)
json_schema["additionalProperties"] = False
return json_schema
def default_schema(self, schema: core_schema.WithDefaultSchema) -> JsonSchemaValue:
json_schema = super().default_schema(schema)
if json_schema.get("default", object()) is None:
json_schema.pop("default")
return json_schema
def strict_json_schema(response_format: type[BaseModel]) -> dict[str, Any]:
"""Serialize a typed response model directly into OpenAI's strict subset."""
return response_format.model_json_schema(schema_generator=OpenAIStrictSchemaGenerator)
@@ -128,11 +128,6 @@ 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,14 +19,10 @@ 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,
BankRowsJSONEncoding,
TransferCausalRelation,
TransferChunk,
TransferDocument,
@@ -75,7 +71,9 @@ _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(
{
@@ -127,9 +125,10 @@ class _LoadedExport:
unit_index: dict[Any, _UnitLocation] = field(default_factory=dict)
# 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.
# 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")
# Facts of these types are exported; observations are derived and excluded.
_EXPORTED_FACT_TYPES = ("world", "experience")
@@ -139,12 +138,7 @@ def _as_jsonb(value: Any) -> Any:
if value is None:
return None
if isinstance(value, str):
try:
return json.loads(value)
except json.JSONDecodeError:
# Admin connections register a JSONB decoder, so a valid scalar such
# as `"combined"` arrives here as the already-decoded `combined`.
return value
return json.loads(value)
return value
@@ -272,13 +266,7 @@ async def _dump_history_rows(conn: Any, table: str, bank_id: str) -> list[dict]:
return [{k: v for k, v in dict(row).items() if k not in _DERIVED_COLUMNS and k != "id"} for row in rows]
async def export_bank(
conn: Any,
bank_id: str,
*,
include_history: bool = False,
bank_rows_json_encoding: BankRowsJSONEncoding = "serialized",
) -> bytes:
async def export_bank(conn: Any, bank_id: str, *, include_history: bool = False) -> bytes:
"""Export an entire bank into a portable ZIP archive (no embeddings).
Produces a superset of the documents archive: the logical
@@ -299,11 +287,11 @@ async def export_bank(
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
@@ -333,7 +321,6 @@ async def export_bank(
directive_count=len(bank_rows.get("directives", [])),
webhook_count=len(bank_rows.get("webhooks", [])),
includes_history=include_history,
bank_rows_json_encoding=bank_rows_json_encoding,
)
zf.writestr("manifest.json", manifest.model_dump_json(indent=2))
@@ -556,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,9 +18,8 @@ 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, link_utils, orchestrator
from ..retain import bank_utils, chunk_storage, embedding_processing, fact_storage, orchestrator
from ..retain.types import (
CausalRelation,
ChunkMetadata,
@@ -30,10 +29,7 @@ from ..retain.types import (
)
from ..schema import fq_table
from .schema import (
CARRIED_HISTORY_TABLES,
HISTORY_TABLES,
SCHEMA_VERSION,
BankRowsJSONEncoding,
TransferDocument,
TransferFact,
TransferManifest,
@@ -87,14 +83,6 @@ class _ObservationOutcome:
skipped: int = 0
@dataclass
class _ImportedFactBatch:
"""Inserted fact IDs paired with their ordinals in the source archive."""
unit_ids: list[str]
original_ordinals: list[int]
@dataclass
class ParsedArchive:
"""A transfer archive after parsing/validation."""
@@ -177,7 +165,7 @@ async def import_documents(
if target_id != document.id:
result.remapped_document_ids[document.id] = target_id
imported_facts = await _import_one_document(
unit_ids = await _import_one_document(
backend=backend,
embeddings_model=embeddings_model,
entity_resolver=entity_resolver,
@@ -190,16 +178,16 @@ async def import_documents(
outbox_callback_factory=outbox_callback_factory,
)
result.documents_imported += 1
result.facts_imported += len(imported_facts.unit_ids)
result.facts_imported += len(unit_ids)
result.imported_documents.append(
ImportedDocument(
document_id=target_id,
unit_ids=imported_facts.unit_ids,
unit_ids=unit_ids,
content=document.original_text or "",
tags=list(document.tags),
)
)
for ordinal, unit_id in zip(imported_facts.original_ordinals, imported_facts.unit_ids, strict=True):
for ordinal, unit_id in enumerate(unit_ids):
ref_map[(document.id, ordinal)] = unit_id
if parsed.observations:
@@ -233,6 +221,8 @@ _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
@@ -273,29 +263,18 @@ 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))
return ParsedBankArchive(manifest=manifest, bank_rows=bank_rows, history_rows=history_rows)
def _resolve_bank_rows_json_encoding(manifest: TransferManifest) -> BankRowsJSONEncoding:
"""Resolve row JSON provenance, including the released v1 archive contract."""
return manifest.bank_rows_json_encoding or "decoded"
async def _restore_rows(
conn: Any,
table: str,
rows: list[dict],
*,
bank_rows_json_encoding: BankRowsJSONEncoding = "decoded",
) -> int:
async def _restore_rows(conn: Any, table: str, rows: list[dict]) -> int:
"""Insert verbatim rows into a bank-scoped table, coercing JSON-encoded values
back to the column's type (timestamps, uuids, jsonb). ``ON CONFLICT DO NOTHING``
keeps an import idempotent and safe to re-run against a partially-filled target."""
@@ -322,12 +301,9 @@ async def _restore_rows(
value = row[col]
if data_type in ("jsonb", "json"):
# asyncpg has no JSON codec on these raw connections; pass JSON
# text and cast. Provenance is required because a decoded JSON
# scalar containing JSON text is indistinguishable from a raw
# serialized object after the outer archive JSON is parsed.
if value is not None and (bank_rows_json_encoding == "decoded" or not isinstance(value, str)):
value = json.dumps(value)
values.append(value)
# text and cast. Values may already be str (no codec on export) or
# a Python object (codec on export) — normalize to text either way.
values.append(value if isinstance(value, str) or value is None else json.dumps(value))
placeholders.append(f"${position}::jsonb")
continue
if value is not None and isinstance(value, str):
@@ -376,7 +352,6 @@ async def import_bank(
if ops is None:
ops = backend.ops
parsed = parse_bank_archive(archive_bytes)
bank_rows_json_encoding = _resolve_bank_rows_json_encoding(parsed.manifest)
source_bank_id = parsed.manifest.source_bank_id
bank_id = target_bank_id or source_bank_id
@@ -399,22 +374,10 @@ async def import_bank(
f"(it is not a merge). Delete the bank first, or pass a different target bank id."
)
# Bank row first — children (documents, mental_models, …) FK to it.
await _restore_rows(
conn,
"banks",
parsed.bank_rows.get("banks", []),
bank_rows_json_encoding=bank_rows_json_encoding,
)
# The restored banks row bypasses the fresh-INSERT gate that normally
# creates per-bank vector indexes, so create them explicitly here while
# the bank is still empty (facts are imported below, so the build is
# instant). get_or_create_bank_profile would NOT do this: the row now
# exists, so it takes the SELECT branch and skips index creation —
# leaving the restored bank falling back to the global index +
# post-filter (slower, under-returning recall). See #2645.
internal_id = await conn.fetchval(f"SELECT internal_id FROM {fq_table('banks')} WHERE bank_id = $1", bank_id)
if internal_id is not None:
await bank_utils.create_bank_vector_indexes(conn, bank_id, str(internal_id), ops=ops)
await _restore_rows(conn, "banks", parsed.bank_rows.get("banks", []))
# Ensure the bank's per-bank vector indexes exist (no-op for global-index
# extensions); idempotent and keeps the restored banks row (ON CONFLICT DO NOTHING).
await bank_utils.get_or_create_bank_profile(backend, bank_id)
doc_result = await import_documents(
backend=backend,
@@ -436,38 +399,17 @@ async def import_bank(
)
async with acquire_with_retry(backend) as conn:
result.mental_models_imported = await _restore_rows(
conn,
"mental_models",
parsed.bank_rows.get("mental_models", []),
bank_rows_json_encoding=bank_rows_json_encoding,
conn, "mental_models", parsed.bank_rows.get("mental_models", [])
)
# Restored after mental_models so the (mental_model_id, bank_id) FK resolves.
result.mental_model_history_imported = await _restore_rows(
conn,
"mental_model_history",
parsed.bank_rows.get("mental_model_history", []),
bank_rows_json_encoding=bank_rows_json_encoding,
)
result.directives_imported = await _restore_rows(
conn,
"directives",
parsed.bank_rows.get("directives", []),
bank_rows_json_encoding=bank_rows_json_encoding,
)
result.webhooks_imported = await _restore_rows(
conn,
"webhooks",
parsed.bank_rows.get("webhooks", []),
bank_rows_json_encoding=bank_rows_json_encoding,
conn, "mental_model_history", parsed.bank_rows.get("mental_model_history", [])
)
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:
result.history_rows_imported += await _restore_rows(
conn,
table,
parsed.history_rows.get(table, []),
bank_rows_json_encoding=bank_rows_json_encoding,
)
for table in _HISTORY_TABLES:
result.history_rows_imported += await _restore_rows(conn, table, parsed.history_rows.get(table, []))
logger.info(
"[transfer] Imported bank %s: %d doc(s), %d fact(s), %d observation(s), "
@@ -519,8 +461,8 @@ async def _import_one_document(
target_id: str,
ops: Any,
outbox_callback_factory: Any = None,
) -> _ImportedFactBatch:
"""Re-embed and insert a document; map original fact ordinals to new unit ids."""
) -> list[str]:
"""Re-embed and insert a single document; returns the new unit ids in fact order."""
log_buffer: list[str] = []
# Fire the same retain.completed webhook retain emits, transactionally inside
@@ -532,21 +474,12 @@ 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] = []
retained_index_by_original: list[int | None] = []
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)
fact_batch = orchestrator._process_extracted_facts(extracted_facts, embeddings)
extracted_facts = fact_batch.extracted_facts
processed_facts = fact_batch.processed_facts
retained_index_by_original = fact_batch.retained_index_by_original
legacy_causal_relations = orchestrator._remap_causal_relations(
legacy_causal_relations,
retained_index_by_original,
)
processed_facts = [ProcessedFact.from_extracted_fact(ef, emb) for ef, emb in zip(extracted_facts, embeddings)]
contents = [RetainContent(content=document.original_text or "")]
chunk_meta = [
@@ -582,15 +515,6 @@ async def _import_one_document(
document.tags,
ops=ops,
)
if document.created_at is not None:
# Transfer archives carry source provenance. Apply it here,
# without changing normal retain/upsert timestamp semantics.
await conn.execute(
f"UPDATE {fq_table('documents')} SET created_at = $1 WHERE id = $2 AND bank_id = $3",
document.created_at,
target_id,
bank_id,
)
chunk_id_map: dict[int, str] = {}
if chunk_meta:
@@ -612,7 +536,7 @@ async def _import_one_document(
processed_facts,
config,
log_buffer,
resolved_entities=phase1.entities.resolved_entities,
resolved_entity_ids=phase1.entities.resolved_entity_ids,
entity_to_unit=phase1.entities.entity_to_unit,
unit_to_entity_ids=phase1.entities.unit_to_entity_ids,
semantic_ann_links=phase1.semantic_ann_links,
@@ -621,34 +545,14 @@ 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:
logger.warning("[transfer] Entity stats flush failed for document %s", target_id, exc_info=True)
logger.debug("[transfer] Imported document %s:\n%s", target_id, "\n".join(log_buffer))
# Single content item -> result_unit_ids[0] follows the retained fact order.
retained_unit_ids = list(result_unit_ids[0]) if result_unit_ids else []
return _ImportedFactBatch(
unit_ids=retained_unit_ids,
original_ordinals=[
original_index
for original_index, retained_index in enumerate(retained_index_by_original)
if retained_index is not None
],
)
# Single content item -> result_unit_ids[0] holds the new unit ids in fact order.
return list(result_unit_ids[0]) if result_unit_ids else []
async def _import_observations(
@@ -718,19 +622,11 @@ async def _import_observations(
all_source_ids: set[uuid.UUID] = set()
for (obs, sources), obs_unit_id in zip(resolved, obs_unit_ids):
observation_uuid = uuid.UUID(obs_unit_id)
if obs.event_date is not None:
# insert_facts_batch derives event_date for normal writes;
# transfer restores the source value carried by the archive.
await conn.execute(
f"UPDATE {fq_table('memory_units')} SET event_date = $1 WHERE id = $2 AND bank_id = $3",
obs.event_date,
observation_uuid,
bank_id,
)
source_uuids = [uuid.UUID(s) for s in sources]
all_source_ids.update(source_uuids)
await _link_observation_sources(conn, ops, bank_id, observation_uuid, source_uuids, obs.proof_count)
await _link_observation_sources(
conn, ops, bank_id, uuid.UUID(obs_unit_id), source_uuids, obs.proof_count
)
# Mark source facts consolidated so the target consolidator skips them.
if all_source_ids:
@@ -809,7 +705,6 @@ 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,
@@ -819,19 +714,3 @@ 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,14 +22,7 @@ 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]]
BankRowsJSONEncoding = Literal["decoded", "serialized"]
class TransferCausalRelation(BaseModel):
@@ -143,7 +136,3 @@ class TransferManifest(BaseModel):
webhook_count: int = 0
# True when --include-history carried audit_log / llm_requests.
includes_history: bool = False
# How JSON/JSONB values in bank/history row files were represented by the
# producing connection. Absent on legacy v1 archives; import treats those as
# decoded because the released producer was the codec-enabled admin CLI.
bank_rows_json_encoding: BankRowsJSONEncoding | None = None
@@ -1,216 +0,0 @@
"""Per-bank vector index coverage checks and repair.
Per-(bank, fact_type) partial vector indexes are created only when a bank is
first created (instant on an empty bank). A bank that becomes *populated*
outside that fresh-INSERT path via a logical restore, a cross-version upgrade,
or a vector-extension switch (e.g. ScaNNpgvector) never gets them, so its
bank-scoped recall silently falls back to the global index + post-filter, which
is both slower and under-returns results. See issue #2645.
This module is the shared engine for detecting and repairing that gap. It is
driven by the ``repair-bank`` admin command; the build always uses
``CREATE INDEX CONCURRENTLY`` on a raw autocommit connection so it never takes
``ACCESS EXCLUSIVE`` on the shared ``memory_units`` table.
"""
from __future__ import annotations
import logging
from dataclasses import dataclass, field
from typing import Any
from .retain.bank_utils import _BANK_INDEX_FACT_TYPES, _bank_index_name
logger = logging.getLogger(__name__)
# Postgres renders the partial predicate of an indexdef with parenthesized
# comparison operands and an explicit ::text cast, e.g.
# `... WHERE ((fact_type = 'world'::text) AND (bank_id = 'b1'::text))`.
# fact_type is emitted first (it is written first in the CREATE INDEX). Match
# that exact rendering so a mere name collision never counts as healthy.
_BANK_INDEX_PARTIAL_SUFFIX = " WHERE ((fact_type = "
# Access methods that legitimately back a per-(bank, fact_type) partial index.
# An index whose access method drifted after a backend switch does not match,
# so the health check treats it as unhealthy (rebuild).
_SUPPORTED_INDEX_AM: tuple[str, ...] = (
"btree",
"gin",
"gist",
"hnsw",
"ivfflat",
"diskann",
"vchordrq",
)
@dataclass
class SchemaVectorIndexResult:
"""Per-schema outcome of a vector-index repair pass."""
schema: str
banks_scanned: int = 0
already_present: int = 0
created: int = 0
skipped: int = 0 # would-create, reported under --dry-run
failed: int = 0
failed_indexes: list[str] = field(default_factory=list)
def _quote_identifier(value: str) -> str:
return '"' + value.replace('"', '""') + '"'
async def _index_health(conn: Any, schema: str, index_names: list[str]) -> dict[str, bool]:
"""Return valid-and-usable state for each requested index in one query.
Health requires the index to be valid AND ready, defined over the expected
``memory_units`` table, to use a supported access method, and to carry our
partial predicate. A name-only match is *not* enough: an INVALID leftover
(from an interrupted concurrent build) or an index whose access method
drifted after a backend switch must count as unhealthy so it is rebuilt
``pg_indexes``/``IF NOT EXISTS`` alone would silently treat those as present.
"""
if not index_names:
return {}
rows = await conn.fetch(
"""
SELECT c.relname AS index_name,
(i.indisvalid AND i.indisready
AND t.relname = 'memory_units'
AND am.amname = ANY($3::text[])
AND pg_get_indexdef(i.indexrelid) LIKE $4
) AS healthy
FROM pg_class c
JOIN pg_namespace n ON n.oid = c.relnamespace
JOIN pg_index i ON i.indexrelid = c.oid
JOIN pg_class t ON t.oid = i.indrelid
JOIN pg_am am ON am.oid = c.relam
WHERE n.nspname = $1 AND c.relname = ANY($2::text[])
""",
schema,
index_names,
list(_SUPPORTED_INDEX_AM),
"%" + _BANK_INDEX_PARTIAL_SUFFIX + "%",
)
return {row["index_name"]: bool(row["healthy"]) for row in rows}
async def _repair_schema(
conn: Any,
schema: str,
index_clause: str,
*,
dry_run: bool,
bank_id: str | None,
) -> SchemaVectorIndexResult:
result = SchemaVectorIndexResult(schema=schema)
qschema = _quote_identifier(schema)
if bank_id is not None:
banks = await conn.fetch(
f"SELECT bank_id, internal_id FROM {qschema}.banks WHERE bank_id = $1", # noqa: S608 — schema is a quoted identifier
bank_id,
)
else:
banks = await conn.fetch(f"SELECT bank_id, internal_id FROM {qschema}.banks ORDER BY bank_id") # noqa: S608
result.banks_scanned = len(banks)
# Resolve expected index names for every bank, then check them all in one
# catalog query rather than one round-trip per index.
expected_by_bank: list[tuple[str, dict[str, str]]] = []
all_index_names: list[str] = []
for bank in banks:
expected = {ft: _bank_index_name(ft, str(bank["internal_id"])) for ft in _BANK_INDEX_FACT_TYPES}
expected_by_bank.append((bank["bank_id"], expected))
all_index_names.extend(expected.values())
health = await _index_health(conn, schema, all_index_names)
for bid, expected in expected_by_bank:
# Render the bank_id literal server-side so escaping does not depend on
# standard_conforming_strings (the predicate is inlined into the DDL).
bank_id_literal = await conn.fetchval("SELECT quote_literal($1::text)", bid)
for ft in _BANK_INDEX_FACT_TYPES:
index_name = expected[ft]
healthy = health.get(index_name)
if healthy is True:
result.already_present += 1
continue
if dry_run:
result.skipped += 1
continue
qindex = _quote_identifier(index_name)
qualified = f"{qschema}.{qindex}"
try:
# An unhealthy-but-present index (INVALID leftover, wrong access
# method) must be dropped first — IF NOT EXISTS cannot repair it.
if healthy is False:
await conn.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {qualified}")
await conn.execute(
f"CREATE INDEX CONCURRENTLY IF NOT EXISTS {qindex} "
f"ON {qschema}.memory_units {index_clause} "
f"WHERE fact_type = '{ft}' AND bank_id = {bank_id_literal}"
)
result.created += 1
except Exception as exc: # noqa: BLE001 — one failed index must not abort the rest
result.failed += 1
result.failed_indexes.append(qualified)
logger.warning(
"Failed to repair vector index %s (bank=%s, fact_type=%s): %s"
"dropping the invalid leftover so a re-run can retry.",
qualified,
bid,
ft,
exc,
)
# A failed concurrent build leaves an INVALID index behind that
# would shadow the good one; drop it so a re-run retries cleanly.
try:
await conn.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {qualified}")
except Exception as cleanup_exc: # noqa: BLE001
logger.warning("Cleanup DROP INDEX for %s also failed: %s", qualified, cleanup_exc)
return result
async def _safe_repair_schema(
conn: Any,
schema: str,
index_clause: str,
*,
dry_run: bool,
bank_id: str | None,
) -> SchemaVectorIndexResult:
try:
return await _repair_schema(conn, schema, index_clause, dry_run=dry_run, bank_id=bank_id)
except Exception as exc: # noqa: BLE001 — one bad schema must not abort the whole sweep
logger.warning("Vector index repair aborted for schema %s: %s", schema, exc)
return SchemaVectorIndexResult(schema=schema, failed=1, failed_indexes=[f"{schema}.<schema-error>"])
async def repair_vector_indexes(
conn: Any,
schemas: list[str],
index_clause: str,
*,
dry_run: bool = False,
bank_id: str | None = None,
) -> list[SchemaVectorIndexResult]:
"""Rebuild missing or invalid per-bank vector indexes across ``schemas``.
``conn`` must be a raw autocommit PostgreSQL connection: ``CREATE INDEX
CONCURRENTLY`` cannot run inside a transaction block. When ``bank_id`` is
given, only that bank is reconciled (in each schema); otherwise every bank
is scanned.
Concurrency is handled by idempotency, not a lock (project rule: no advisory
locks they are unreliable behind connection poolers). Every build is
``CREATE INDEX CONCURRENTLY IF NOT EXISTS`` guarded by a valid/ready health
check, so a second concurrent run is a no-op on already-built indexes; if two
runs race the *same* missing index, Postgres rejects one build and the
per-index handler drops the leftover so a re-run converges cleanly.
"""
return [
await _safe_repair_schema(conn, schema, index_clause, dry_run=dry_run, bank_id=bank_id) for schema in schemas
]
@@ -45,7 +45,6 @@ from hindsight_api.extensions.operation_validator import (
# Consolidation operation
ConsolidateContext,
ConsolidateResult,
CreateBankContext,
# File Conversion
FileConvertResult,
# Mental Model operations
@@ -106,7 +105,6 @@ __all__ = [
"BankReadOperation",
"BankWriteContext",
"BankWriteOperation",
"CreateBankContext",
# Operation Validator - Consolidation
"ConsolidateContext",
"ConsolidateResult",
@@ -359,7 +359,6 @@ class BankWriteOperation(StrEnum):
DELETE_DIRECTIVE = "delete_directive"
DELETE_DOCUMENT = "delete_document"
DELETE_MENTAL_MODEL = "delete_mental_model"
DELETE_OPERATION = "delete_operation"
DELETE_WEBHOOK = "delete_webhook"
MERGE_BANK_MISSION = "merge_bank_mission"
REPROCESS_DOCUMENT = "reprocess_document"
@@ -398,14 +397,6 @@ 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)."""
@@ -890,23 +881,6 @@ 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.
+10 -68
View File
@@ -9,7 +9,7 @@ import json
import logging
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import Any, Callable, get_args
from typing import Any, Callable
from fastmcp import FastMCP
from mcp.types import ToolAnnotations
@@ -23,7 +23,7 @@ from hindsight_api.config import (
from hindsight_api.engine.audit import AuditEntry, AuditLogger
from hindsight_api.engine.memory_engine import Budget
from hindsight_api.engine.response_models import VALID_RECALL_FACT_TYPES, MinScores
from hindsight_api.engine.search.tags import TagGroup, TagsMatch
from hindsight_api.engine.search.tags import TagGroup
from hindsight_api.extensions import OperationValidationError
from hindsight_api.models import RequestContext
@@ -512,8 +512,7 @@ def _apply_audit_logging(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsCon
"""Create an audited wrapper for a tool's run method."""
async def _audited_run(arguments, _name=tool_name, _orig=original_run):
# Cheap bank-independent pre-filter before resolving bank_id.
if not audit_logger.action_allowed(_name):
if not audit_logger.is_enabled(_name):
return await _orig(arguments)
bank_id = None
@@ -522,10 +521,6 @@ def _apply_audit_logging(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsCon
elif hasattr(arguments, "get"):
bank_id = arguments.get("bank_id")
# Per-bank decision, resolved after bank_id is known.
if not await audit_logger.should_log(_name, bank_id):
return await _orig(arguments)
entry = AuditEntry(
action=_name,
transport="mcp",
@@ -563,8 +558,7 @@ def _apply_audit_logging(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsCon
if original_call_tool:
async def _audited_call_tool(name, arguments=None, **kwargs):
# Cheap bank-independent pre-filter before resolving bank_id.
if name not in _AUDITABLE_MCP_TOOLS or not audit_logger.action_allowed(name):
if name not in _AUDITABLE_MCP_TOOLS or not audit_logger.is_enabled(name):
return await original_call_tool(name, arguments, **kwargs)
bank_id = None
@@ -573,10 +567,6 @@ def _apply_audit_logging(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsCon
config.bank_id_resolver() if config.bank_id_resolver else None
)
# Per-bank decision, resolved after bank_id is known.
if not await audit_logger.should_log(name, bank_id):
return await original_call_tool(name, arguments, **kwargs)
entry = AuditEntry(
action=name,
transport="mcp",
@@ -1235,9 +1225,7 @@ def _register_create_bank(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsCo
"""
try:
request_context = _get_request_context(config)
# 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)
# get_bank_profile auto-creates bank if it doesn't exist
profile = await memory.get_bank_profile(bank_id, request_context=request_context)
# Update name/mission if provided
@@ -1264,10 +1252,7 @@ def _register_create_bank(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsCo
def _validate_mental_model_inputs(
name: str | None = None,
source_query: str | None = None,
max_tokens: int | None = None,
tags_match: str | None = None,
name: str | None = None, source_query: str | None = None, max_tokens: int | None = None
) -> str | None:
"""Validate mental model inputs, returning an error message or None if valid."""
if name is not None and not name.strip():
@@ -1276,9 +1261,6 @@ def _validate_mental_model_inputs(
return "source_query cannot be empty"
if max_tokens is not None and (max_tokens < 256 or max_tokens > 8192):
return f"max_tokens must be between 256 and 8192, got {max_tokens}"
if tags_match is not None and tags_match not in get_args(TagsMatch):
valid = ", ".join(get_args(TagsMatch))
return f"tags_match must be one of {valid}, got {tags_match!r}"
return None
@@ -1460,7 +1442,6 @@ def _register_create_mental_model(mcp: FastMCP, memory: MemoryEngine, config: MC
source_query: str,
mental_model_id: str | None = None,
tags: list[str] | None = None,
tags_match: str | None = None,
max_tokens: int = 2048,
trigger_refresh_after_consolidation: bool = False,
bank_id: str | None = None,
@@ -1482,12 +1463,6 @@ def _register_create_mental_model(mcp: FastMCP, memory: MemoryEngine, config: MC
source_query: The query to run through reflect to generate content
mental_model_id: Optional custom ID (alphanumeric lowercase with hyphens). Auto-generated if not provided.
tags: Optional tags for scoped visibility filtering
tags_match: How this model's tags are matched against memories when the content
is (re)generated. One of 'any' (match any tag, like recall/reflect), 'all'
(match all tags), 'any_strict', 'all_strict', or 'exact'. If omitted, a tagged
model defaults to 'all_strict' a memory must carry EVERY one of the model's
tags to be included, which silently filters out memories that only carry a
subset. Pass 'any' when your memories use narrow single-topic tags.
max_tokens: Maximum tokens for generated content (256-8192, default: 2048)
trigger_refresh_after_consolidation: If True, automatically refresh this model after memory consolidation. Default: False
bank_id: Optional bank (defaults to session bank). Use for cross-bank operations.
@@ -1498,15 +1473,13 @@ def _register_create_mental_model(mcp: FastMCP, memory: MemoryEngine, config: MC
return '{"error": "No bank_id configured"}'
validation_error = _validate_mental_model_inputs(
name=name, source_query=source_query, max_tokens=max_tokens, tags_match=tags_match
name=name, source_query=source_query, max_tokens=max_tokens
)
if validation_error:
return json.dumps({"error": validation_error})
request_context = _get_request_context(config)
trigger: dict[str, Any] = {"refresh_after_consolidation": trigger_refresh_after_consolidation}
if tags_match is not None:
trigger["tags_match"] = tags_match
trigger = {"refresh_after_consolidation": trigger_refresh_after_consolidation}
# Create with placeholder content
model = await memory.create_mental_model(
@@ -1553,7 +1526,6 @@ def _register_create_mental_model(mcp: FastMCP, memory: MemoryEngine, config: MC
source_query: str,
mental_model_id: str | None = None,
tags: list[str] | None = None,
tags_match: str | None = None,
max_tokens: int = 2048,
trigger_refresh_after_consolidation: bool = False,
) -> dict:
@@ -1574,12 +1546,6 @@ def _register_create_mental_model(mcp: FastMCP, memory: MemoryEngine, config: MC
source_query: The query to run through reflect to generate content
mental_model_id: Optional custom ID (alphanumeric lowercase with hyphens). Auto-generated if not provided.
tags: Optional tags for scoped visibility filtering
tags_match: How this model's tags are matched against memories when the content
is (re)generated. One of 'any' (match any tag, like recall/reflect), 'all'
(match all tags), 'any_strict', 'all_strict', or 'exact'. If omitted, a tagged
model defaults to 'all_strict' a memory must carry EVERY one of the model's
tags to be included, which silently filters out memories that only carry a
subset. Pass 'any' when your memories use narrow single-topic tags.
max_tokens: Maximum tokens for generated content (256-8192, default: 2048)
trigger_refresh_after_consolidation: If True, automatically refresh this model after memory consolidation. Default: False
"""
@@ -1589,15 +1555,13 @@ def _register_create_mental_model(mcp: FastMCP, memory: MemoryEngine, config: MC
return {"error": "No bank_id configured"}
validation_error = _validate_mental_model_inputs(
name=name, source_query=source_query, max_tokens=max_tokens, tags_match=tags_match
name=name, source_query=source_query, max_tokens=max_tokens
)
if validation_error:
return {"error": validation_error}
request_context = _get_request_context(config)
trigger: dict[str, Any] = {"refresh_after_consolidation": trigger_refresh_after_consolidation}
if tags_match is not None:
trigger["tags_match"] = tags_match
trigger = {"refresh_after_consolidation": trigger_refresh_after_consolidation}
model = await memory.create_mental_model(
bank_id=target_bank,
@@ -2279,8 +2243,6 @@ def _register_list_memories(mcp: FastMCP, memory: MemoryEngine, config: MCPTools
limit: int = 100,
offset: int = 0,
bank_id: str | None = None,
tags: list[str] | None = None,
tags_match: TagsMatch = "any",
) -> str:
"""
Browse stored memories with optional filtering.
@@ -2294,10 +2256,6 @@ def _register_list_memories(mcp: FastMCP, memory: MemoryEngine, config: MCPTools
limit: Maximum number of results (default: 100)
offset: Pagination offset (default: 0)
bank_id: Optional bank (defaults to session bank). Use for cross-bank operations.
tags: Optional list of tag names to filter by.
tags_match: How to combine tags: 'any' (OR, default) or 'all' (AND)
both also include untagged memories; 'any_strict'/'all_strict'
exclude untagged; 'exact' matches the tag set exactly.
"""
try:
target_bank = bank_id or config.bank_id_resolver()
@@ -2310,8 +2268,6 @@ def _register_list_memories(mcp: FastMCP, memory: MemoryEngine, config: MCPTools
search_query=q,
limit=limit,
offset=offset,
tags=tags,
tags_match=tags_match,
request_context=_get_request_context(config),
)
return json.dumps(result, indent=2, default=str)
@@ -2330,8 +2286,6 @@ def _register_list_memories(mcp: FastMCP, memory: MemoryEngine, config: MCPTools
q: str | None = None,
limit: int = 100,
offset: int = 0,
tags: list[str] | None = None,
tags_match: TagsMatch = "any",
) -> dict:
"""
Browse stored memories with optional filtering.
@@ -2344,10 +2298,6 @@ def _register_list_memories(mcp: FastMCP, memory: MemoryEngine, config: MCPTools
q: Optional text search query to filter memories
limit: Maximum number of results (default: 100)
offset: Pagination offset (default: 0)
tags: Optional list of tag names to filter by.
tags_match: How to combine tags: 'any' (OR, default) or 'all' (AND)
both also include untagged memories; 'any_strict'/'all_strict'
exclude untagged; 'exact' matches the tag set exactly.
"""
try:
target_bank = config.bank_id_resolver()
@@ -2360,8 +2310,6 @@ def _register_list_memories(mcp: FastMCP, memory: MemoryEngine, config: MCPTools
search_query=q,
limit=limit,
offset=offset,
tags=tags,
tags_match=tags_match,
request_context=_get_request_context(config),
)
return result
@@ -3197,10 +3145,7 @@ 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)
@@ -3228,10 +3173,7 @@ 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
+7 -12
View File
@@ -56,11 +56,6 @@ 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)
@@ -196,22 +191,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)
_set_alembic_main_option(alembic_cfg, "script_location", script_location)
alembic_cfg.set_main_option("script_location", script_location)
# Set the database URL
_set_alembic_main_option(alembic_cfg, "sqlalchemy.url", database_url)
alembic_cfg.set_main_option("sqlalchemy.url", database_url)
# Configure logging (optional, but helps with debugging)
# Uses Python's logging system instead of alembic.ini
_set_alembic_main_option(alembic_cfg, "prepend_sys_path", ".")
alembic_cfg.set_main_option("prepend_sys_path", ".")
# Set path_separator to avoid deprecation warning
_set_alembic_main_option(alembic_cfg, "path_separator", "os")
alembic_cfg.set_main_option("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:
_set_alembic_main_option(alembic_cfg, "target_schema", schema)
alembic_cfg.set_main_option("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()
@@ -436,8 +431,8 @@ def check_migration_status(
# Create config programmatically
alembic_cfg = Config()
_set_alembic_main_option(alembic_cfg, "script_location", script_location)
_set_alembic_main_option(alembic_cfg, "path_separator", "os")
alembic_cfg.set_main_option("script_location", script_location)
alembic_cfg.set_main_option("path_separator", "os")
script = ScriptDirectory.from_config(alembic_cfg)
head_rev = script.get_current_head()
@@ -284,8 +284,6 @@ 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",
+17 -56
View File
@@ -2,7 +2,6 @@ from __future__ import annotations
import asyncio
import logging
from dataclasses import dataclass
from typing import TYPE_CHECKING
if TYPE_CHECKING:
@@ -154,22 +153,7 @@ async def stop_embedded_postgres() -> None:
await _default_instance.stop()
@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:
def parse_pg0_url(db_url: str) -> tuple[bool, str | None, int | None]:
"""
Parse a database URL and check if it's a pg0:// embedded database URL.
@@ -177,47 +161,29 @@ def parse_pg0_url(db_url: str) -> Pg0Url:
- "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:
A :class:`Pg0Url`. When ``is_pg0`` is False the remaining fields are None.
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)
"""
if db_url == "pg0":
return Pg0Url(is_pg0=True, instance_name="hindsight")
return True, "hindsight", None
if not db_url.startswith("pg0://"):
return Pg0Url(is_pg0=False)
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
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,
)
return False, None, None
async def resolve_database_url(db_url: str) -> str:
@@ -233,13 +199,8 @@ async def resolve_database_url(db_url: str) -> str:
Returns:
The resolved postgresql:// connection URL
"""
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)
is_pg0, instance_name, port = parse_pg0_url(db_url)
if is_pg0:
pg0 = EmbeddedPostgres(name=instance_name, port=port)
return await pg0.ensure_running()
return db_url
@@ -197,11 +197,6 @@ 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()
@@ -269,7 +264,6 @@ def main():
max_slots=config.worker_max_slots,
slot_reservations=config.worker_slot_reservations,
consolidation_bank_priority=config.worker_consolidation_bank_priority or None,
max_retries=config.worker_max_retries,
)
# Create the HTTP app for metrics/health
@@ -37,19 +37,6 @@ 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
@@ -161,7 +148,6 @@ class WorkerPoller:
max_slots: int = 10,
slot_reservations: dict[str, int] | None = None,
consolidation_bank_priority: dict[str, int] | None = None,
max_retries: int = 3,
):
"""
Initialize the worker poller.
@@ -184,8 +170,6 @@ class WorkerPoller:
Patterns support ``*`` as wildcard. A bare ``*`` key is the catch-all default.
When set, consolidation tasks are claimed in priority tiers rather than
pure created_at order. None or empty dict preserves current behavior.
max_retries: Maximum retry attempts before a task is marked failed.
Must be >= 0. Default 3 (matches DEFAULT_WORKER_MAX_RETRIES).
"""
self._backend = backend
self._worker_id = worker_id
@@ -207,7 +191,6 @@ class WorkerPoller:
self._consolidation_bank_priority: dict[str, int] | None = (
consolidation_bank_priority if consolidation_bank_priority else None
)
self._max_retries = max(0, max_retries) # Never negative
# Cache of which optional PG routines are installed on the server
# (probed once, memoised for the life of the poller).
from ..engine.db.optional_routines import OptionalRoutines
@@ -226,8 +209,6 @@ 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:
@@ -523,20 +504,17 @@ class WorkerPoller:
return result
async def _mark_completed(self, operation_id: str, schema: str | None):
"""Mark a processing task as completed, then propagate to parent if needed."""
"""Mark a task as completed."""
table = fq_table("async_operations", schema)
async with self._backend.acquire() as conn:
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)
await conn.execute(
f"""
UPDATE {table}
SET status = 'completed', completed_at = now(), updated_at = now()
WHERE operation_id = $1
""",
operation_id,
)
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."""
@@ -771,7 +749,6 @@ 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.
@@ -801,18 +778,14 @@ class WorkerPoller:
This handles the case where a worker crashes while processing tasks.
On startup, we reset any tasks stuck in 'processing' for this worker_id
back to 'pending' so they can be picked up again provided their
``retry_count`` has not reached ``max_retries``. Tasks at/over the
limit are moved to 'failed' with an explanatory error message,
breaking the infinite loop where a task that kills the worker
(OOM, infinite loop) is re-claimed forever.
back to 'pending' so they can be picked up again.
Also recovers batch API operations that were in-flight.
If tenant_extension is configured, recovers across all tenant schemas.
Returns:
Number of tasks recovered (reset to pending, not including failed)
Number of tasks recovered
"""
schemas = await self._get_schemas()
total_count = 0
@@ -825,71 +798,20 @@ class WorkerPoller:
batch_count = await self._recover_batch_operations(schema)
total_count += batch_count
# Then reset normal worker tasks. Crash-interrupted tasks count
# toward the retry budget (worker_max_retries / HINDSIGHT_API_WORKER_MAX_RETRIES).
# Without this, a task that kills the worker (OOM, infinite loop)
# is reset forever: claim → grind → crash → recover → re-claim…
# Two separate UPDATEs so their row counts are meaningful:
# 1. Tasks under limit → increment retry_count, reset to pending
# 2. Tasks at/over limit → move to failed with a clear reason
max_retries = self._max_retries
# Then reset normal worker tasks
async with self._backend.acquire() as conn:
# Tasks under the limit: increment retry_count and reset to pending
result = await conn.execute(
f"""
UPDATE {table}
SET status = 'pending', worker_id = NULL, claimed_at = NULL,
retry_count = COALESCE(retry_count, 0) + 1, updated_at = now()
WHERE status = 'processing' AND worker_id = $1
AND result_metadata->>'batch_id' IS NULL
AND COALESCE(retry_count, 0) < $2
SET status = 'pending', worker_id = NULL, claimed_at = NULL, updated_at = now()
WHERE status = 'processing' AND worker_id = $1 AND result_metadata->>'batch_id' IS NULL
""",
self._worker_id,
max_retries,
)
# Tasks that exceeded the limit: move to failed. RETURNING
# gives us the ids so their parent aggregators can be rolled
# up below — a batch_retain child sub-batch carries
# parent_operation_id (not batch_id) in its metadata, so it IS
# eligible to be failed here, and without propagating that
# terminal state the parent is stranded in 'processing' forever
# (the same crash loop this method fixes, one level up).
failed_rows = await conn.fetch(
f"""
UPDATE {table}
SET status = 'failed', worker_id = NULL, claimed_at = NULL,
error_message = 'exceeded max recovery attempts after crash (retry_count >= {max_retries})',
completed_at = now(), updated_at = now()
WHERE status = 'processing' AND worker_id = $1
AND result_metadata->>'batch_id' IS NULL
AND COALESCE(retry_count, 0) >= $2
RETURNING operation_id
""",
self._worker_id,
max_retries,
)
# Roll each failed child up to its parent aggregator, one
# transaction per child so a single problematic parent can't undo
# the others — mirrors the per-task transaction the in-process
# _mark_failed path uses. The failing UPDATE above already committed,
# so the children stay failed regardless; _maybe_update_parent_operation
# no-ops for tasks without a parent_operation_id.
for failed_row in failed_rows:
async with self._backend.acquire() as conn:
async with conn.transaction():
await self._maybe_update_parent_operation(str(failed_row["operation_id"]), schema, conn)
pending_count = int(result.split()[-1]) if result else 0
failed_count = len(failed_rows)
total_count += pending_count
if failed_count > 0:
schema_display = f'"{schema}"' if schema else str(schema)
logger.warning(
f"Worker {self._worker_id} moved {failed_count} tasks to 'failed' "
f"(exceeded {max_retries} recovery attempts in schema "
f"{schema_display})"
)
# Parse "UPDATE N" to get count
count = int(result.split()[-1]) if result else 0
total_count += count
except Exception as e:
# Format schema for logging: custom schemas in quotes, None as-is
schema_display = f'"{schema}"' if schema else str(schema)
@@ -1016,7 +938,6 @@ class WorkerPoller:
for task in tasks:
await self.execute_task(task)
if tasks:
# Continue immediately to claim more tasks (if slots available)
continue
+7 -9
View File
@@ -51,7 +51,7 @@ dependencies = [
"anthropic>=0.40.0",
"typer>=0.9.0",
"cohere>=5.0.0",
"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
"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
"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'",
@@ -62,7 +62,7 @@ dependencies = [
"langchain-core>=1.2.22", # Path traversal in legacy load_prompt functions fix
"langsmith>=0.8.18", # GHSA-f4xh-w4cj-qxq8: arbitrary server-side file read in TracingMiddleware fix (supersedes >=0.6.3 SSRF tracing-header-injection floor)
"protobuf>=6.33.5", # JSON recursion depth bypass fix
"pillow>=12.3.0", # Multiple HIGH image parsing vulnerabilities fixed in 12.3.0
"pillow>=12.1.1", # Out-of-bounds write in PSD image loading fix
"cryptography>=48.0.1", # GHSA-537c-gmf6-5ccf: bundled-OpenSSL OOB read fix needs >=48.0.1. Prior <47 cap (47.0.0 SIGILL on ARM64 Docker/Podman, pyca/cryptography#14733) lifted — 47/48/49 verified importing + RSA sign/verify cleanly on linux/arm64 (Docker on Apple Silicon) and native arm64 macOS; upstream issue closed unconfirmed.
"filelock>=3.20.1", # TOCTOU race condition fix
"authlib>=1.6.9", # Account takeover/JWS header injection vulnerability fix
@@ -75,18 +75,16 @@ dependencies = [
"claude-agent-sdk>=0.2.82",
"boto3>=1.42.74",
"croniter>=2.0.0", # Cron parsing for scheduled mental model refresh
"json-repair>=0.30.0", # Structural repair of malformed LLM JSON (last-resort parse fallback)
]
[project.optional-dependencies]
local-ml = [
# Local ML models for embeddings/reranking
"sentence-transformers>=3.3.0",
"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.
"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.
"tokenizers>=0.22.0,<=0.23.0",
"torch>=2.6.0", # CVE fix for remote code execution
"einops>=0.8.2",
@@ -106,7 +104,7 @@ local-llm = [
local-onnx = [
# In-process ONNX Runtime embeddings without an Ollama/TEI sidecar
"onnxruntime>=1.17.0",
"transformers>=5.5.0", # 5.5.0 clears GHSA-fgcw-684q-jj6r (LightGlue RCE)
"transformers>=4.53.0",
"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",
+4 -46
View File
@@ -24,33 +24,6 @@ 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
@@ -163,7 +136,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).is_pg0:
if db_url and not _parse_pg0_url(db_url)[0]:
# Plain postgresql:// URL - use it directly but still run migrations
from hindsight_api.migrations import run_migrations
@@ -171,9 +144,9 @@ def pg0_db_url(db_url, tmp_path_factory, worker_id):
return db_url
if db_url:
_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
_, 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
else:
pg0_instance_name = DEFAULT_PG0_INSTANCE_NAME
pg0_instance_port = DEFAULT_PG0_PORT
@@ -624,18 +597,3 @@ async def api_client(memory):
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
yield client
def enable_audit_default(memory, enabled: bool) -> None:
"""Set the deployment-wide default for the hierarchical ``audit_log_enabled``.
``audit_log_enabled`` resolves through env -> tenant -> bank, and the
ConfigResolver snapshots the global layer at construction time. Tests that
want "auditing on by default" therefore have to update that snapshot;
flipping ``AuditLogger._enabled`` alone only covers actions with no bank in
scope. Per-bank overrides are set with ``resolver.update_bank_config``.
"""
from dataclasses import replace
resolver = memory._config_resolver
resolver._global_config = replace(resolver._global_config, audit_log_enabled=enabled)
@@ -1,49 +0,0 @@
"""Tests for the admin CLI whole-bank transfer boundary."""
from pathlib import Path
from unittest.mock import AsyncMock
import pytest
from hindsight_api.admin import cli
class _FakeConnection:
def __init__(self) -> None:
self.closed = False
async def close(self) -> None:
self.closed = True
@pytest.mark.asyncio
async def test_run_export_bank_declares_decoded_json_rows(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
"""The codec-enabled admin producer must identify its rows as decoded."""
connection = _FakeConnection()
export_bank = AsyncMock(return_value=b"archive")
async def fake_admin_connect(db_url: str) -> _FakeConnection:
assert db_url == "postgresql://example"
return connection
monkeypatch.setattr(cli, "_admin_connect", fake_admin_connect)
monkeypatch.setattr(cli, "export_bank", export_bank)
output = tmp_path / "bank.zip"
size = await cli._run_export_bank(
"postgresql://example",
"source-bank",
output,
"tenant_schema",
include_history=True,
)
export_bank.assert_awaited_once_with(
connection,
"source-bank",
include_history=True,
bank_rows_json_encoding="decoded",
)
assert output.read_bytes() == b"archive"
assert size == len(b"archive")
assert connection.closed is True
@@ -1,273 +0,0 @@
"""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")
@@ -1,181 +0,0 @@
"""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,8 +107,5 @@ 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
# 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 "valid JSON matching this schema" in (kwargs.get("system") or "")
assert isinstance(result, _Decision)
@@ -2,7 +2,6 @@
import asyncio
import json
import os
import uuid
import pytest
@@ -516,259 +515,6 @@ async def test_retain_outcome_metadata_records_zero_counts(memory, request_conte
assert "extraction_errors_sample" not in parent["result_metadata"]
@pytest.mark.asyncio
async def test_all_degenerate_facts_still_persist_document_chunks(memory, request_context, monkeypatch):
"""Filtering every extracted fact must not turn an extracted chunk into the zero-extraction fast path."""
from hindsight_api.engine.response_models import TokenUsage
from hindsight_api.engine.retain import fact_extraction
from hindsight_api.engine.retain.types import ChunkMetadata, ExtractedFact
async def degenerate_extract_facts_from_contents(contents, *_args, **_kwargs):
return (
[ExtractedFact(fact_text="...", fact_type="world", content_index=0, chunk_index=0)],
[ChunkMetadata(chunk_text=contents[0].content, fact_count=1, content_index=0, chunk_index=0)],
TokenUsage(),
)
monkeypatch.setattr(fact_extraction, "extract_facts_from_contents", degenerate_extract_facts_from_contents)
bank_id = f"test_all_degenerate_{uuid.uuid4().hex[:8]}"
document_id = "all-degenerate-document"
try:
await memory.retain_async(
bank_id=bank_id,
content="A source chunk whose only extracted fact is rejected.",
document_id=document_id,
request_context=request_context,
)
chunks = await memory.list_document_chunks(bank_id, document_id, limit=10, request_context=request_context)
units = await memory.list_memory_units(bank_id, request_context=request_context)
assert [chunk["chunk_index"] for chunk in chunks["items"]] == [0]
assert units["total"] == 0
finally:
await memory.delete_bank(bank_id, request_context=request_context)
@pytest.mark.asyncio
async def test_streaming_offsets_chunk_local_causal_fact_indices(memory, request_context, monkeypatch):
"""Causal targets from independently extracted chunks must stay within their source chunk."""
from hindsight_api.engine.response_models import TokenUsage
from hindsight_api.engine.retain import fact_extraction
from hindsight_api.engine.retain.types import CausalRelation, ChunkMetadata, ExtractedFact
chunks = ["first-streaming-chunk", "second-streaming-chunk"]
monkeypatch.setattr(fact_extraction, "chunk_text", lambda *_args, **_kwargs: chunks)
async def extract_chunk_facts(contents, *_args, **_kwargs):
chunk_text = contents[0].content
return (
[
ExtractedFact(fact_text=f"{chunk_text} cause", fact_type="world", chunk_index=0),
ExtractedFact(
fact_text=f"{chunk_text} effect",
fact_type="world",
chunk_index=0,
causal_relations=[CausalRelation(relation_type="caused_by", target_fact_index=0)],
),
],
[ChunkMetadata(chunk_text=chunk_text, fact_count=2, content_index=0, chunk_index=0)],
TokenUsage(),
)
monkeypatch.setattr(fact_extraction, "extract_facts_from_contents", extract_chunk_facts)
bank_id = f"test_streaming_causal_{uuid.uuid4().hex[:8]}"
try:
await memory.retain_async(
bank_id=bank_id,
content="Content is replaced by the deterministic chunk_text stub.",
document_id="streaming-causal-document",
request_context=request_context,
)
pool = await memory._get_pool()
rows = await pool.fetch(
"""
SELECT source.text AS source_text, target.text AS target_text
FROM memory_links links
JOIN memory_units source ON source.id = links.from_unit_id
JOIN memory_units target ON target.id = links.to_unit_id
WHERE links.bank_id = $1 AND links.link_type = 'caused_by'
""",
bank_id,
)
assert {(row["source_text"], row["target_text"]) for row in rows} == {
(f"{chunk} effect", f"{chunk} cause") for chunk in chunks
}
finally:
await memory.delete_bank(bank_id, request_context=request_context)
@pytest.mark.asyncio
async def test_degenerate_fact_preserves_later_chunk_provenance(memory, request_context, monkeypatch):
"""A rejected degenerate fact must not shift chunk provenance onto a later chunk's survivor.
Regression for #2794: filtering only ``processed_facts`` left ``extracted_facts``
full-length, so the consumer's ``zip(batch_extracted, batch_processed)`` paired
every survivor after a rejected fact with the wrong extracted fact assigning it
the wrong chunk_id.
Each of the two chunks emits [real, degenerate]. After the first (real, degenerate)
pair the zip is off-by-one for the rest of the batch, so *both* surviving real facts
collapse onto whichever chunk sorted first regardless of the nondeterministic
producer completion order. The fix filters both lists in lockstep per chunk, so each
real fact keeps its own chunk_index.
"""
from hindsight_api.engine.response_models import TokenUsage
from hindsight_api.engine.retain import fact_extraction
from hindsight_api.engine.retain.types import ChunkMetadata, ExtractedFact
chunks = ["chunk-zero-source", "chunk-one-source"]
monkeypatch.setattr(fact_extraction, "chunk_text", lambda *_args, **_kwargs: chunks)
real_fact_by_chunk = {chunks[0]: "chunk zero real fact", chunks[1]: "chunk one real fact"}
async def extract_chunk_facts(contents, *_args, **_kwargs):
chunk_text = contents[0].content
facts = [
ExtractedFact(fact_text=real_fact_by_chunk[chunk_text], fact_type="world", chunk_index=0),
ExtractedFact(fact_text="...", fact_type="world", chunk_index=0),
]
return (
facts,
[ChunkMetadata(chunk_text=chunk_text, fact_count=len(facts), content_index=0, chunk_index=0)],
TokenUsage(),
)
monkeypatch.setattr(fact_extraction, "extract_facts_from_contents", extract_chunk_facts)
bank_id = f"test_degen_provenance_{uuid.uuid4().hex[:8]}"
try:
await memory.retain_async(
bank_id=bank_id,
content="Content is replaced by the deterministic chunk_text stub.",
document_id="degen-provenance-document",
request_context=request_context,
)
pool = await memory._get_pool()
rows = await pool.fetch(
"""
SELECT units.text AS fact_text, chunks.chunk_index AS chunk_index
FROM memory_units units
JOIN chunks ON chunks.chunk_id = units.chunk_id
WHERE units.bank_id = $1
""",
bank_id,
)
assert {(row["fact_text"], row["chunk_index"]) for row in rows} == {
("chunk zero real fact", 0),
("chunk one real fact", 1),
}
finally:
await memory.delete_bank(bank_id, request_context=request_context)
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."""
+2 -8
View File
@@ -13,19 +13,14 @@ import pytest_asyncio
from hindsight_api.api import create_app
from hindsight_api.config import get_config
from tests.conftest import enable_audit_default
@pytest_asyncio.fixture
async def audit_api_client(memory):
"""Create a test client with audit logging enabled deployment-wide."""
# audit_log_enabled is hierarchical, so the effective value for a
# bank-scoped action comes from the ConfigResolver, not the logger's own
# flag. Set both so this fixture models "enabled by default, no per-bank
# override" — the logger flag covers actions with no bank in scope.
"""Create a test client with audit logging enabled."""
# Enable audit logging on the memory engine's audit logger
memory._audit_logger._enabled = True
memory._audit_logger._allowed_actions = None # All actions
enable_audit_default(memory, True)
app = create_app(memory, initialize_memory=False)
transport = httpx.ASGITransport(app=app)
@@ -402,7 +397,6 @@ async def test_audit_log_action_allowlist(memory):
"""Test that only allowed actions are audited when allowlist is set."""
memory._audit_logger._enabled = True
memory._audit_logger._allowed_actions = frozenset({"recall"}) # Only audit recall
enable_audit_default(memory, True)
app = create_app(memory, initialize_memory=False)
transport = httpx.ASGITransport(app=app)
@@ -1,220 +0,0 @@
"""Tests for the per-bank ``audit_log_enabled`` override.
``audit_log_enabled`` resolves through env -> tenant -> bank, so a deployment
can audit some banks and not others. These tests cover both directions of the
override (a bank opting IN when the default is off, and a bank opting OUT when
the default is on), plus the fallbacks that decide behaviour when there is no
bank in scope or config resolution fails.
"""
import asyncio
from datetime import datetime
import httpx
import pytest
import pytest_asyncio
from hindsight_api.api import create_app
from hindsight_api.engine.audit import AuditLogger
from tests.conftest import enable_audit_default
# Audit writes are fire-and-forget; give the background task room to land.
_AUDIT_SETTLE_SECONDS = 1.0
@pytest_asyncio.fixture
async def client(memory):
"""HTTP client whose audit allowlist is open (all actions auditable)."""
memory._audit_logger._allowed_actions = None
app = create_app(memory, initialize_memory=False)
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url="http://test") as c:
yield c
async def _recall(client, bank_id: str) -> None:
r = await client.post(
f"/v1/default/banks/{bank_id}/memories/recall",
json={"query": "per-bank audit test"},
)
assert r.status_code == 200, r.text
await asyncio.sleep(_AUDIT_SETTLE_SECONDS)
async def _audited_actions(client, bank_id: str) -> list[str]:
r = await client.get(f"/v1/default/banks/{bank_id}/audit-logs")
assert r.status_code == 200, r.text
return [e["action"] for e in r.json()["items"]]
def _bank(prefix: str) -> str:
return f"{prefix}_{datetime.now().timestamp()}"
# ── the override, both directions ──────────────────────────────────────────
@pytest.mark.asyncio
async def test_bank_opts_in_when_default_off(client, memory):
"""A bank override of True audits even though the deployment default is off."""
enable_audit_default(memory, False)
memory._audit_logger._enabled = False
bank_id = _bank("audit_optin")
await client.put(f"/v1/default/banks/{bank_id}", json={})
await memory._config_resolver.update_bank_config(bank_id, {"audit_log_enabled": True})
await _recall(client, bank_id)
assert "recall" in await _audited_actions(client, bank_id)
@pytest.mark.asyncio
async def test_bank_opts_out_when_default_on(client, memory):
"""A bank override of False suppresses auditing that is otherwise on."""
enable_audit_default(memory, True)
memory._audit_logger._enabled = True
bank_id = _bank("audit_optout")
await client.put(f"/v1/default/banks/{bank_id}", json={})
await memory._config_resolver.update_bank_config(bank_id, {"audit_log_enabled": False})
await _recall(client, bank_id)
# The create_bank above was audited (it ran before the override existed),
# so assert on the action the override was meant to suppress.
assert "recall" not in await _audited_actions(client, bank_id)
@pytest.mark.asyncio
async def test_banks_are_independent(client, memory):
"""One bank's override does not leak into another bank."""
enable_audit_default(memory, False)
memory._audit_logger._enabled = False
audited, quiet = _bank("audit_on"), _bank("audit_off")
for b in (audited, quiet):
await client.put(f"/v1/default/banks/{b}", json={})
await memory._config_resolver.update_bank_config(audited, {"audit_log_enabled": True})
await _recall(client, audited)
await _recall(client, quiet)
assert "recall" in await _audited_actions(client, audited)
assert await _audited_actions(client, quiet) == []
@pytest.mark.asyncio
async def test_no_override_uses_deployment_default(client, memory):
"""A bank with no explicit override follows the deployment default."""
enable_audit_default(memory, True)
memory._audit_logger._enabled = True
bank_id = _bank("audit_inherit")
await client.put(f"/v1/default/banks/{bank_id}", json={})
await _recall(client, bank_id)
assert "recall" in await _audited_actions(client, bank_id)
# ── AuditLogger decision logic (no DB) ─────────────────────────────────────
def _logger(*, enabled: bool, resolver=None, actions: list[str] | None = None) -> AuditLogger:
return AuditLogger(
pool_getter=lambda: None,
schema_getter=lambda: "public",
enabled=enabled,
allowed_actions=actions or [],
bank_enabled_resolver=resolver,
)
@pytest.mark.asyncio
async def test_allowlist_short_circuits_before_resolution():
"""A disallowed action never triggers a per-bank lookup."""
calls: list[str] = []
async def resolver(bank_id, context=None):
calls.append(bank_id)
return True
log = _logger(enabled=True, resolver=resolver, actions=["recall"])
assert await log.should_log("reflect", "b1") is False
assert calls == [], "allowlist must short-circuit before resolving bank config"
@pytest.mark.asyncio
async def test_no_bank_in_scope_uses_global_default():
"""With no bank_id there is nothing to resolve, so the global value decides."""
async def resolver(bank_id, context=None):
raise AssertionError("must not resolve without a bank_id")
assert await _logger(enabled=True, resolver=resolver).should_log("recall", None) is True
assert await _logger(enabled=False, resolver=resolver).should_log("recall", None) is False
@pytest.mark.asyncio
async def test_resolution_failure_falls_back_to_global_default():
"""A resolver error must not drop audit rows for an audited deployment.
Fails OPEN (to the deployment default) rather than closed: a transient DB
blip should not silently create an audit gap for a bank meant to be audited.
"""
async def broken(bank_id, context=None):
raise RuntimeError("config backend down")
assert await _logger(enabled=True, resolver=broken).should_log("recall", "b1") is True
assert await _logger(enabled=False, resolver=broken).should_log("recall", "b1") is False
@pytest.mark.asyncio
async def test_unwired_resolver_uses_global_default():
"""Without a resolver the logger behaves exactly as before this feature."""
assert await _logger(enabled=True).should_log("recall", "b1") is True
assert await _logger(enabled=False).should_log("recall", "b1") is False
@pytest.mark.asyncio
async def test_gating_ignores_config_permission_filter(memory, request_context):
"""A tenant permission filter must not suppress a bank's audit override.
get_allowed_config_fields controls which fields a user may *modify* (and it
filters the API-facing config read). Audit gating is an internal decision
and must see the bank's true stored value: a deployment that makes
audit_log_enabled read-only for some users must still audit banks that
opted in. Regression guard for the resolve_full_config vs get_bank_config
distinction in MemoryEngine._resolve_bank_audit_enabled.
"""
from hindsight_api.config_resolver import ConfigResolver
from hindsight_api.extensions.tenant import Tenant, TenantContext, TenantExtension
bank_id = "audit-perm-filter"
class RestrictiveExtension(TenantExtension):
def __init__(self):
pass # skip Extension.__init__(config); nothing here needs config
async def authenticate(self, context):
return TenantContext(schema_name="public")
async def list_tenants(self):
return [Tenant(schema="public")]
async def get_allowed_config_fields(self, context, bank_id):
# audit_log_enabled deliberately absent: read-only for this user.
return {"retain_chunk_size"}
await memory.get_bank_profile(bank_id, request_context=request_context)
# Store the opt-in with an allow-all resolver (the write is a separate
# concern from gating), then swap in the restrictive extension.
await memory._config_resolver.update_bank_config(bank_id, {"audit_log_enabled": True}, request_context)
restrictive = ConfigResolver(backend=memory._backend, tenant_extension=RestrictiveExtension())
memory._config_resolver = restrictive
# The API-facing read is filtered (proving the filter is actually active)...
api_view = await restrictive.get_bank_config(bank_id, request_context)
assert "audit_log_enabled" not in api_view
# ...but gating still sees the real value and audits the bank.
assert await memory._resolve_bank_audit_enabled(bank_id, request_context) is True
@@ -17,17 +17,14 @@ 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)
@@ -59,9 +56,31 @@ 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/reflect/task methods."""
"""The engine binds the bank via @_bind_bank_id on recall/retain/batch/task methods."""
async def test_binds_named_arg_positional_and_keyword(self):
@_bind_bank_id()
@@ -98,41 +117,6 @@ class TestBindBankIdDecorator:
assert await op(12345) is None
async def test_engine_provider_paths_bind_and_reset_their_bank_arguments(self):
engine = object.__new__(MemoryEngine)
engine._reflect_llm_config = None
engine._operation_validator = 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
with (
patch.object(
engine,
"_authenticate_tenant",
AsyncMock(side_effect=lambda _context: observed_bank_ids.append(get_current_bank_id())),
),
patch.object(engine, "_get_backend", AsyncMock(side_effect=RuntimeError("stop after authentication"))),
):
with pytest.raises(RuntimeError, match="stop after authentication"):
await engine.update_memory_unit(
"user-update",
"54a647e5-0a22-4e5d-8504-b8bfca2a6142",
text="corrected",
request_context=RequestContext(),
)
assert observed_bank_ids[-1] == "user-update"
assert get_current_bank_id() is None
# ── LLM provider: user injection ──────────────────────────────────────────────
@@ -274,22 +258,27 @@ def test_embeddings_user_injected_when_flag_on_and_bank_set():
assert captured[0]["user"] == "user-emb"
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))
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]
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")])
assert scores == [0.91]
assert encoder._async_client.post.call_args.kwargs["headers"] == {"X-Hindsight-Bank-Id": "bank-litellm-proxy"}
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]
# ── Executor context propagation ──────────────────────────────────────────────

Some files were not shown because too many files have changed in this diff Show More