Compare commits
6
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
0022d427d3 | ||
|
|
1f9bad0858 | ||
|
|
138bf02f29 | ||
|
|
df178aae8a | ||
|
|
a43026b8f4 | ||
|
|
5c425e276e |
@@ -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"
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
@@ -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.
|
||||
@@ -159,18 +154,6 @@ If any files in `hindsight-api-slim/hindsight_api/api/` were changed:
|
||||
- Were the client SDKs regenerated? (`./scripts/generate-clients.sh`)
|
||||
- Were the control plane proxy routes updated? (`hindsight-control-plane/src/app/api/`)
|
||||
|
||||
### 7a. Check TS/Python wrapper-client parity
|
||||
|
||||
Two of the generated SDKs ship a **hand-written, maintained convenience wrapper** on top of the auto-generated low-level client — and *only* these two:
|
||||
- **TypeScript**: `hindsight-clients/typescript/src/index.ts` (`HindsightClient`)
|
||||
- **Python**: `hindsight-clients/python/hindsight_client/hindsight_client.py` (`Hindsight`)
|
||||
|
||||
(The Rust/Go/etc. clients are generated-only — no wrapper to keep in sync.)
|
||||
|
||||
These wrappers are what most third-party consumers actually call, and they must expose the same surface. **If a change touches one wrapper's method — adds/removes a parameter, changes a default, forwards a new query/body field — the equivalent method in the *other* wrapper must get the same change in the same (or an immediately-following) PR.** A parameter that exists in the generated SDK but is dropped by one wrapper silently strips it for every consumer of that language (this is exactly what #2975 / #3042 fixed for `detail`/`tags_match`/`limit`/`offset` on `listMentalModels`/`getMentalModel`). **Should fix** — flag any wrapper method that gains capabilities in one language but not the other, and add a matching mapping regression test on both sides.
|
||||
|
||||
Note: the `client-coverage-check` CI tool only validates **request-body** fields, not GET **query** parameters — so query-param parity gaps are *not* caught automatically and must be checked by hand here.
|
||||
|
||||
### 7b. Check API-layer data-access boundary
|
||||
|
||||
For each changed handler in `hindsight-api-slim/hindsight_api/api/` (e.g. `http.py`, `mcp.py`):
|
||||
@@ -221,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:
|
||||
|
||||
@@ -21,25 +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
|
||||
|
||||
# Some backends, including Bedrock Converse, reject JSON Schema maxItems.
|
||||
# Disable it only for those backends; consolidation still enforces the cap.
|
||||
# HINDSIGHT_API_LLM_SUPPORTS_MAX_ITEMS=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
|
||||
@@ -78,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
|
||||
@@ -108,16 +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
|
||||
|
||||
# Wall-clock ceiling (seconds) for one retain task in the worker. A retain that
|
||||
# blocks indefinitely is cancelled and marked 'failed' — and so becomes
|
||||
# retryable — instead of holding its worker slot until the process restarts.
|
||||
# Set well above your slowest healthy retain; 0 disables. Default 3600.
|
||||
# HINDSIGHT_API_RETAIN_WALL_TIMEOUT=3600
|
||||
|
||||
# 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
|
||||
@@ -133,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)
|
||||
@@ -156,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
|
||||
@@ -173,19 +127,12 @@ HINDSIGHT_API_LOG_LEVEL=info
|
||||
# HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_BASE_URL=
|
||||
# HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_MODEL=
|
||||
# HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_PROMPT=
|
||||
# Optional JSON dict of custom headers for the OCR OpenAI client (e.g. proxies / request tracing).
|
||||
# HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_DEFAULT_HEADERS=
|
||||
|
||||
# Embeddings Configuration (Optional - uses local by default)
|
||||
# Provider: "local" (default), "onnx", "tei", "openai", "cohere", "google", "openrouter", "zeroentropy", "litellm", or "litellm-sdk"
|
||||
# HINDSIGHT_API_EMBEDDINGS_PROVIDER=local
|
||||
# For local provider:
|
||||
# HINDSIGHT_API_EMBEDDINGS_LOCAL_MODEL=BAAI/bge-small-en-v1.5
|
||||
# Force CPU if local embeddings hit MPS/XPC instability on macOS:
|
||||
# HINDSIGHT_API_EMBEDDINGS_LOCAL_FORCE_CPU=false
|
||||
# Opt in to the Apple Silicon MPS GPU (off by default: MPS leaks memory under
|
||||
# variable-length workloads). CUDA/XPU still auto-select regardless:
|
||||
# HINDSIGHT_API_EMBEDDINGS_LOCAL_ALLOW_MPS=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
|
||||
@@ -222,28 +169,11 @@ HINDSIGHT_API_LOG_LEVEL=info
|
||||
# DeepSeek note: DeepSeek is supported for LLM calls, but not for embeddings.
|
||||
# If using DeepSeek as LLM provider, keep embeddings on local/openai/cohere/google/etc.
|
||||
|
||||
# Embedding similarity thresholds. These defaults preserve the behavior calibrated
|
||||
# for BAAI/bge-small-en-v1.5. Recalibrate each threshold independently when changing
|
||||
# embedding models because cosine-similarity distributions are model-dependent.
|
||||
# HINDSIGHT_API_SEMANTIC_MIN_SIMILARITY=0.3
|
||||
# HINDSIGHT_API_GRAPH_SEED_MIN_SIMILARITY=0.3
|
||||
# HINDSIGHT_API_TEMPORAL_SEMANTIC_MIN_SIMILARITY=0.1
|
||||
# HINDSIGHT_API_SEMANTIC_LINK_MIN_SIMILARITY=0.7
|
||||
# HINDSIGHT_API_CONSOLIDATION_DEDUP_THRESHOLD=0.97
|
||||
|
||||
# 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
|
||||
# Opt in to the Apple Silicon MPS GPU (off by default: MPS leaks memory under
|
||||
# variable-length workloads). CUDA/XPU still auto-select regardless:
|
||||
# HINDSIGHT_API_RERANKER_LOCAL_ALLOW_MPS=false
|
||||
# For TEI provider:
|
||||
# HINDSIGHT_API_RERANKER_TEI_URL=http://localhost:8081
|
||||
|
||||
@@ -265,16 +195,6 @@ HINDSIGHT_API_LOG_LEVEL=info
|
||||
# Expose async-operation queue + consolidation-backlog gauges on /metrics.
|
||||
# Runs periodic per-schema COUNT queries on a background task (disabled by default).
|
||||
# HINDSIGHT_API_METRICS_BACKLOG_ENABLED=true
|
||||
#
|
||||
# Runtime-stall observability (enabled by default). When a liveness probe fails,
|
||||
# these tell you WHY: a blocked event loop vs DB connection-pool exhaustion.
|
||||
# The loop watchdog logs the offending stack when the loop is unresponsive; the
|
||||
# DB-pool acquire timing logs (and exposes hindsight.db.pool.waiting) when
|
||||
# callers queue for a connection. Both are cheap; tune or disable if needed.
|
||||
# HINDSIGHT_API_LOOP_WATCHDOG_ENABLED=false
|
||||
# HINDSIGHT_API_LOOP_WATCHDOG_STALL_THRESHOLD_MS=1000
|
||||
# HINDSIGHT_API_LOOP_WATCHDOG_POLL_INTERVAL_MS=250
|
||||
# HINDSIGHT_API_DB_ACQUIRE_WARN_THRESHOLD_MS=1000
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Control Plane (Optional)
|
||||
|
||||
@@ -25,20 +25,6 @@ jobs:
|
||||
with:
|
||||
python-version-file: ".python-version"
|
||||
|
||||
# Each package is built from its own directory, so stage the repository's
|
||||
# canonical license inside each isolated build context.
|
||||
- name: Stage Python package licenses
|
||||
run: |
|
||||
for package in \
|
||||
hindsight-clients/python \
|
||||
hindsight-api-slim \
|
||||
hindsight-api \
|
||||
hindsight-all \
|
||||
hindsight-all-slim \
|
||||
hindsight-embed; do
|
||||
cp LICENSE "$package/LICENSE"
|
||||
done
|
||||
|
||||
# Build all packages
|
||||
- name: Build hindsight-client
|
||||
working-directory: ./hindsight-clients/python
|
||||
@@ -64,24 +50,6 @@ jobs:
|
||||
working-directory: ./hindsight-embed
|
||||
run: uv build --out-dir dist
|
||||
|
||||
- name: Verify Python package licenses
|
||||
run: |
|
||||
for package in \
|
||||
hindsight-clients/python \
|
||||
hindsight-api-slim \
|
||||
hindsight-api \
|
||||
hindsight-all \
|
||||
hindsight-all-slim \
|
||||
hindsight-embed; do
|
||||
wheel=$(find "$package/dist" -maxdepth 1 -name '*.whl' -print -quit)
|
||||
sdist=$(find "$package/dist" -maxdepth 1 -name '*.tar.gz' -print -quit)
|
||||
|
||||
unzip -Z1 "$wheel" | grep -Eq '\.dist-info/licenses/LICENSE$'
|
||||
unzip -p "$wheel" '*/METADATA' | grep -Fxq 'License-Expression: MIT'
|
||||
unzip -p "$wheel" '*/METADATA' | grep -Fxq 'License-File: LICENSE'
|
||||
tar -tzf "$sdist" | grep -Eq '/LICENSE$'
|
||||
done
|
||||
|
||||
# Publish in order (client and api-slim first, then api/all wrappers which depend on them)
|
||||
- name: Publish hindsight-client to PyPI
|
||||
uses: pypa/gh-action-pypi-publish@release/v1
|
||||
|
||||
+17
-169
@@ -41,8 +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-copilot-cli: ${{ steps.filter.outputs.integrations-copilot-cli }}
|
||||
integrations-crewai: ${{ steps.filter.outputs.integrations-crewai }}
|
||||
integrations-litellm: ${{ steps.filter.outputs.integrations-litellm }}
|
||||
integrations-pydantic-ai: ${{ steps.filter.outputs.integrations-pydantic-ai }}
|
||||
@@ -154,8 +152,6 @@ jobs:
|
||||
- 'hindsight-integrations/continue/**'
|
||||
integrations-cursor-cli:
|
||||
- 'hindsight-integrations/cursor-cli/**'
|
||||
integrations-copilot-cli:
|
||||
- 'hindsight-integrations/copilot-cli/**'
|
||||
integrations-crewai:
|
||||
- 'hindsight-integrations/crewai/**'
|
||||
integrations-litellm:
|
||||
@@ -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:
|
||||
@@ -289,18 +283,6 @@ jobs:
|
||||
working-directory: ./hindsight-api-slim
|
||||
run: uv build
|
||||
|
||||
# `uv build` only packages the source; it does not prove the dependency set
|
||||
# resolves or that the code imports on this interpreter. Install into a fresh
|
||||
# env and run a byte-compile + import smoke test so the matrix actually
|
||||
# exercises each Python version (notably 3.14).
|
||||
- name: Install and smoke-test on Python ${{ matrix.python-version }}
|
||||
working-directory: ./hindsight-api-slim
|
||||
run: |
|
||||
uv venv --python ${{ matrix.python-version }} .venv-smoke
|
||||
VIRTUAL_ENV=.venv-smoke uv pip install .
|
||||
.venv-smoke/bin/python -m compileall -q hindsight_api
|
||||
.venv-smoke/bin/python -c "import hindsight_api, hindsight_api.main, hindsight_api.config; from hindsight_api.engine import memory_engine, llm_wrapper; print('import OK')"
|
||||
|
||||
build-typescript-client:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
@@ -538,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]
|
||||
@@ -715,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: >-
|
||||
@@ -1279,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
|
||||
|
||||
@@ -1886,27 +1836,6 @@ jobs:
|
||||
with:
|
||||
ref: ${{ github.event.pull_request.head.sha || '' }}
|
||||
|
||||
# The Oracle 23ai `free` service image is large; combined with the Python ML
|
||||
# deps (torch) it exhausts the runner's ~14 GB root disk, so uv fails to
|
||||
# extract a wheel with "No space left on device (os error 28)".
|
||||
# Only the cheap, high-yield reclaims are enabled — the Android/.NET/Haskell
|
||||
# dirs plus swap are a few `rm -rf`s worth ~16-21 GB, which is ample headroom:
|
||||
# - large-packages runs apt-get remove and costs minutes for little gain;
|
||||
# - tool-cache would delete the preinstalled Python that actions/setup-python
|
||||
# then has to re-download, making the job slower, not faster;
|
||||
# - docker-images is pointless here (the Oracle service container is already
|
||||
# running, so its image is in use and cannot be pruned anyway).
|
||||
- name: Free Disk Space
|
||||
uses: jlumbroso/free-disk-space@main
|
||||
with:
|
||||
tool-cache: false
|
||||
android: true
|
||||
dotnet: true
|
||||
haskell: true
|
||||
large-packages: false
|
||||
docker-images: false
|
||||
swap-storage: true
|
||||
|
||||
- name: Setup Oracle test user
|
||||
# The SYSTEM tablespace uses manual segment space management which
|
||||
# doesn't support VECTOR types. Create an ASSM tablespace and a
|
||||
@@ -2267,27 +2196,6 @@ jobs:
|
||||
with:
|
||||
ref: ${{ github.event.pull_request.head.sha || '' }}
|
||||
|
||||
# The Oracle 23ai `free` service image is large; combined with the Python ML
|
||||
# deps (torch) it exhausts the runner's ~14 GB root disk, so uv fails to
|
||||
# extract a wheel with "No space left on device (os error 28)".
|
||||
# Only the cheap, high-yield reclaims are enabled — the Android/.NET/Haskell
|
||||
# dirs plus swap are a few `rm -rf`s worth ~16-21 GB, which is ample headroom:
|
||||
# - large-packages runs apt-get remove and costs minutes for little gain;
|
||||
# - tool-cache would delete the preinstalled Python that actions/setup-python
|
||||
# then has to re-download, making the job slower, not faster;
|
||||
# - docker-images is pointless here (the Oracle service container is already
|
||||
# running, so its image is in use and cannot be pruned anyway).
|
||||
- name: Free Disk Space
|
||||
uses: jlumbroso/free-disk-space@main
|
||||
with:
|
||||
tool-cache: false
|
||||
android: true
|
||||
dotnet: true
|
||||
haskell: true
|
||||
large-packages: false
|
||||
docker-images: false
|
||||
swap-storage: true
|
||||
|
||||
- name: Setup Oracle test user
|
||||
run: |
|
||||
pip install oracledb
|
||||
@@ -2448,27 +2356,6 @@ jobs:
|
||||
with:
|
||||
ref: ${{ github.event.pull_request.head.sha || '' }}
|
||||
|
||||
# The Oracle 23ai `free` service image is large; combined with the Python ML
|
||||
# deps (torch) it exhausts the runner's ~14 GB root disk, so uv fails to
|
||||
# extract a wheel with "No space left on device (os error 28)".
|
||||
# Only the cheap, high-yield reclaims are enabled — the Android/.NET/Haskell
|
||||
# dirs plus swap are a few `rm -rf`s worth ~16-21 GB, which is ample headroom:
|
||||
# - large-packages runs apt-get remove and costs minutes for little gain;
|
||||
# - tool-cache would delete the preinstalled Python that actions/setup-python
|
||||
# then has to re-download, making the job slower, not faster;
|
||||
# - docker-images is pointless here (the Oracle service container is already
|
||||
# running, so its image is in use and cannot be pruned anyway).
|
||||
- name: Free Disk Space
|
||||
uses: jlumbroso/free-disk-space@main
|
||||
with:
|
||||
tool-cache: false
|
||||
android: true
|
||||
dotnet: true
|
||||
haskell: true
|
||||
large-packages: false
|
||||
docker-images: false
|
||||
swap-storage: true
|
||||
|
||||
- name: Setup Oracle test user
|
||||
run: |
|
||||
pip install oracledb
|
||||
@@ -3594,43 +3481,6 @@ jobs:
|
||||
# (requires_real_llm) needs a live Hindsight server and runs separately.
|
||||
run: uv run pytest tests -v -m "not requires_real_llm"
|
||||
|
||||
test-copilot-cli-integration:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
(github.event_name == 'workflow_dispatch' ||
|
||||
needs.detect-changes.outputs.integrations-copilot-cli == '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 copilot-cli integration
|
||||
working-directory: ./hindsight-integrations/copilot-cli
|
||||
run: uv build
|
||||
|
||||
- name: Install dependencies
|
||||
working-directory: ./hindsight-integrations/copilot-cli
|
||||
run: uv sync --frozen
|
||||
|
||||
- name: Run tests
|
||||
working-directory: ./hindsight-integrations/copilot-cli
|
||||
run: uv run pytest tests -v
|
||||
|
||||
test-crewai-integration:
|
||||
needs: [detect-changes]
|
||||
if: >-
|
||||
@@ -4921,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
|
||||
@@ -5060,7 +4909,6 @@ jobs:
|
||||
- test-github-copilot-integration
|
||||
- test-codex-integration
|
||||
- test-cursor-cli-integration
|
||||
- test-zcode-integration
|
||||
- build-ai-sdk-integration
|
||||
- test-ai-sdk-integration-deno
|
||||
- test-opencode-integration
|
||||
|
||||
@@ -5,13 +5,6 @@ build/
|
||||
dist/
|
||||
wheels/
|
||||
*.egg-info
|
||||
# Release builds stage the canonical root license in each package context.
|
||||
/hindsight-clients/python/LICENSE
|
||||
/hindsight-api-slim/LICENSE
|
||||
/hindsight-api/LICENSE
|
||||
/hindsight-all/LICENSE
|
||||
/hindsight-all-slim/LICENSE
|
||||
/hindsight-embed/LICENSE
|
||||
.mcp.json
|
||||
.playwright-mcp/
|
||||
.osgrep
|
||||
@@ -48,7 +41,6 @@ nltk_data/
|
||||
logs/
|
||||
|
||||
.DS_Store
|
||||
.sesskey
|
||||
|
||||
# Generated docs files
|
||||
hindsight-docs/static/llms-full.txt
|
||||
|
||||
@@ -286,35 +286,6 @@ When adding or modifying parameters in the dataplane API (hindsight-api), you mu
|
||||
- Update the client type definition in `lib/api.ts`
|
||||
- Update any UI components that need to use the new parameter
|
||||
|
||||
### Harness Attribution (which coding agent wrote a document)
|
||||
|
||||
`hindsight-integrations/hindsight-coding-agents/` stamps the coding agent on every
|
||||
document it retains, so the control plane can show its logo instead of another
|
||||
`key=value` chip:
|
||||
|
||||
- `metadata.harness = "<id>"` — the authoritative field
|
||||
- tag `harness:<id>` — the same value, so the documents list can filter on it
|
||||
|
||||
The ids are defined by that integration's HookSpecs
|
||||
(`src/harness/hook-lifecycle.ts`) plus the persistent-plugin entrypoints
|
||||
registered in `src/harness/registry.ts`, whose id is their
|
||||
`createPluginEntry(...)` argument — currently `antigravity-cli`, `claude-code`,
|
||||
`cline-cli`, `codex`, `copilot-cli`, `cursor-cli`, `devin-cli`, `grok-build`,
|
||||
`kilo`, `opencode`.
|
||||
|
||||
The control plane resolves the value in
|
||||
`hindsight-control-plane/src/lib/harness-logo.ts` (metadata wins over the tag) and
|
||||
renders it with `components/ui/harness-logo.tsx` in the documents table and the
|
||||
document detail dialog. **Adding a harness to the integration means adding it to
|
||||
that registry in the same change**: copy its icon from
|
||||
`hindsight-docs/static/img/icons/` (or take it from the agent's own brand assets
|
||||
when the docs site carries none) into
|
||||
`hindsight-control-plane/public/img/harness/` and add one entry. Don't register
|
||||
ids nothing writes — a test asserts the registry matches the emitted set, plus an
|
||||
explicit list of retired ids kept so already-retained documents keep their logo.
|
||||
An unregistered harness is not an error: it renders no logo and still shows as
|
||||
ordinary metadata.
|
||||
|
||||
### Adding New Integrations
|
||||
|
||||
Every new integration in `hindsight-integrations/` must satisfy all of the following before it can be merged:
|
||||
|
||||
Binary file not shown.
|
Before Width: | Height: | Size: 303 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 152 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 153 KiB |
@@ -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
|
||||
|
||||
|
||||
@@ -2,8 +2,8 @@ apiVersion: v2
|
||||
name: hindsight
|
||||
description: Hindsight helm chart
|
||||
type: application
|
||||
version: 0.8.6
|
||||
appVersion: "0.8.6"
|
||||
version: 0.8.4
|
||||
appVersion: "0.8.4"
|
||||
keywords:
|
||||
- ai
|
||||
- memory
|
||||
|
||||
@@ -60,13 +60,13 @@ spec:
|
||||
valueFrom:
|
||||
fieldRef:
|
||||
fieldPath: metadata.name
|
||||
{{- /* Explicitly set port to override K8s service discovery env var (HINDSIGHT_API_PORT) */}}
|
||||
- name: HINDSIGHT_API_PORT
|
||||
value: {{ .Values.worker.service.targetPort | quote }}
|
||||
{{- /* Inherit LLM config from api.env, then apply worker-specific env.
|
||||
Merge (worker.env wins) so a key set in both does not emit a
|
||||
duplicate env entry, which server-side apply rejects. */}}
|
||||
{{- range $key, $value := merge (deepCopy (.Values.worker.env | default dict)) (.Values.api.env | default dict) }}
|
||||
{{- /* Inherit LLM config from api.env */}}
|
||||
{{- range $key, $value := .Values.api.env }}
|
||||
- name: {{ $key }}
|
||||
value: {{ $value | quote }}
|
||||
{{- end }}
|
||||
{{- /* Worker-specific env vars */}}
|
||||
{{- range $key, $value := .Values.worker.env }}
|
||||
- name: {{ $key }}
|
||||
value: {{ $value | quote }}
|
||||
{{- end }}
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@vectorize-io/hindsight-all",
|
||||
"version": "0.8.6",
|
||||
"version": "0.8.4",
|
||||
"description": "Node.js programmatic lifecycle manager for Hindsight — embeds a local hindsight daemon in a Node application. Pair with @vectorize-io/hindsight-client for memory operations.",
|
||||
"main": "dist/index.js",
|
||||
"types": "dist/index.d.ts",
|
||||
|
||||
@@ -1,16 +1,15 @@
|
||||
[build-system]
|
||||
requires = ["setuptools>=77"]
|
||||
requires = ["setuptools>=61"]
|
||||
build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "hindsight-all-slim"
|
||||
version = "0.8.6"
|
||||
version = "0.8.4"
|
||||
description = "Hindsight: Agent Memory That Works Like Human Memory - Slim All-in-One Bundle"
|
||||
license = "MIT"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.11"
|
||||
dependencies = [
|
||||
"hindsight-api-slim==0.8.6",
|
||||
"hindsight-api-slim==0.8.4",
|
||||
"hindsight-client>=0.0.7",
|
||||
"hindsight-embed>=0.1.0",
|
||||
]
|
||||
|
||||
@@ -1,16 +1,15 @@
|
||||
[build-system]
|
||||
requires = ["hatchling>=1.27"]
|
||||
requires = ["hatchling"]
|
||||
build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "hindsight-all"
|
||||
version = "0.8.6"
|
||||
version = "0.8.4"
|
||||
description = "Hindsight: Agent Memory That Works Like Human Memory - All-in-One Bundle"
|
||||
license = "MIT"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.11"
|
||||
dependencies = [
|
||||
"hindsight-api-slim[all]==0.8.6",
|
||||
"hindsight-api-slim[all]==0.8.4",
|
||||
"hindsight-client>=0.0.7",
|
||||
"hindsight-embed>=0.1.0",
|
||||
]
|
||||
@@ -22,7 +21,7 @@ hindsight-embed = { workspace = true }
|
||||
|
||||
[project.optional-dependencies]
|
||||
local-llm = [
|
||||
"hindsight-api-slim[local-llm]==0.8.6",
|
||||
"hindsight-api-slim[local-llm]==0.8.4",
|
||||
]
|
||||
test = [
|
||||
"pytest>=7.0.0",
|
||||
|
||||
@@ -53,4 +53,4 @@ __all__ = [
|
||||
"RemoteTEICrossEncoder",
|
||||
"LLMConfig",
|
||||
]
|
||||
__version__ = "0.8.6"
|
||||
__version__ = "0.8.4"
|
||||
|
||||
@@ -9,7 +9,6 @@ import io
|
||||
import json
|
||||
import logging
|
||||
import zipfile
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
@@ -17,12 +16,10 @@ from typing import Any
|
||||
import asyncpg
|
||||
import typer
|
||||
|
||||
from ..config import DEFAULT_DATABASE_SCHEMA, HindsightConfig, load_dotenv_for_entrypoint
|
||||
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
|
||||
|
||||
@@ -69,88 +66,7 @@ BACKUP_TABLES = [
|
||||
"graph_maintenance_queue",
|
||||
]
|
||||
|
||||
MANIFEST_VERSION = "2"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class BackupColumn:
|
||||
"""A PostgreSQL column shape required to decode a binary COPY stream."""
|
||||
|
||||
name: str
|
||||
type_name: str
|
||||
|
||||
|
||||
async def _table_columns(conn: asyncpg.Connection, schema: str, table: str) -> list[BackupColumn]:
|
||||
rows = await conn.fetch(
|
||||
"""
|
||||
SELECT a.attname AS name, pg_catalog.format_type(a.atttypid, a.atttypmod) AS type_name
|
||||
FROM pg_catalog.pg_attribute AS a
|
||||
JOIN pg_catalog.pg_class AS c ON c.oid = a.attrelid
|
||||
JOIN pg_catalog.pg_namespace AS n ON n.oid = c.relnamespace
|
||||
WHERE n.nspname = $1 AND c.relname = $2 AND a.attnum > 0 AND NOT a.attisdropped
|
||||
AND a.attgenerated = ''
|
||||
ORDER BY a.attnum
|
||||
""",
|
||||
schema,
|
||||
table,
|
||||
)
|
||||
return [BackupColumn(name=row["name"], type_name=row["type_name"]) for row in rows]
|
||||
|
||||
|
||||
async def _validate_restore_schema(
|
||||
conn: asyncpg.Connection, manifest: dict[str, Any], schema: str
|
||||
) -> dict[str, list[str]]:
|
||||
"""Validate every COPY stream against the target before destructive work starts.
|
||||
|
||||
Type equality is an exact ``format_type`` string match. This is deliberately
|
||||
stricter than binary-COPY wire compatibility (e.g. ``varchar`` and ``text``
|
||||
share a binary format yet compare unequal here): we would rather fail a
|
||||
genuinely-restorable backup with a clear, actionable error than silently risk
|
||||
a subtle binary mismatch. Restores blocked this way can be recovered by
|
||||
aligning the target schema.
|
||||
"""
|
||||
restore_columns: dict[str, list[str]] = {}
|
||||
errors: list[str] = []
|
||||
for table, table_manifest in manifest["tables"].items():
|
||||
source_columns = [BackupColumn(**column) for column in table_manifest["columns"]]
|
||||
target_by_name = {column.name: column for column in await _table_columns(conn, schema, table)}
|
||||
missing = [column.name for column in source_columns if column.name not in target_by_name]
|
||||
mismatched = [
|
||||
f"{column.name} ({column.type_name} in backup, {target_by_name[column.name].type_name} in target)"
|
||||
for column in source_columns
|
||||
if column.name in target_by_name and target_by_name[column.name].type_name != column.type_name
|
||||
]
|
||||
if missing:
|
||||
errors.append(f"{table}: target is missing backup columns {', '.join(missing)}")
|
||||
if mismatched:
|
||||
errors.append(f"{table}: incompatible column types: {', '.join(mismatched)}")
|
||||
restore_columns[table] = [column.name for column in source_columns]
|
||||
|
||||
if errors:
|
||||
details = "; ".join(errors)
|
||||
raise ValueError(f"Backup schema is incompatible with target schema '{schema}': {details}")
|
||||
return restore_columns
|
||||
|
||||
|
||||
def _effective_backup_tables() -> list[str]:
|
||||
"""Core backup tables plus any bank-scoped tables a loaded extension declares.
|
||||
|
||||
``BACKUP_TABLES`` covers only the tables core owns. An extension that
|
||||
provisions its own bank-scoped tables (via ``TenantExtension``) declares
|
||||
them through ``extra_bank_tables()`` so they aren't dropped on restore.
|
||||
Extension tables are appended *after* the core set so restore's forward
|
||||
COPY inserts them after their FK parents (e.g. ``banks``) and the reversed
|
||||
TRUNCATE clears them before those parents.
|
||||
"""
|
||||
tables = list(BACKUP_TABLES)
|
||||
tenant_extension = load_extension("TENANT", TenantExtension)
|
||||
if tenant_extension is not None:
|
||||
seen = set(tables)
|
||||
for spec in tenant_extension.extra_bank_tables():
|
||||
if spec.include_in_backup and spec.name not in seen:
|
||||
tables.append(spec.name)
|
||||
seen.add(spec.name)
|
||||
return tables
|
||||
MANIFEST_VERSION = "1"
|
||||
|
||||
|
||||
async def _admin_connect(db_url: str) -> asyncpg.Connection:
|
||||
@@ -161,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))
|
||||
@@ -171,18 +86,8 @@ async def _admin_connect(db_url: str) -> asyncpg.Connection:
|
||||
return conn
|
||||
|
||||
|
||||
async def _backup(
|
||||
database_url: str,
|
||||
output_path: Path,
|
||||
schema: str = "public",
|
||||
backup_tables: list[str] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Backup all tables to a zip file using binary COPY protocol.
|
||||
|
||||
``backup_tables`` defaults to the core ``BACKUP_TABLES``; callers pass the
|
||||
extension-augmented list from ``_effective_backup_tables()``.
|
||||
"""
|
||||
backup_tables = backup_tables if backup_tables is not None else BACKUP_TABLES
|
||||
async def _backup(database_url: str, output_path: Path, schema: str = "public") -> dict[str, Any]:
|
||||
"""Backup all tables to a zip file using binary COPY protocol."""
|
||||
conn = await asyncpg.connect(database_url)
|
||||
try:
|
||||
tables: dict[str, Any] = {}
|
||||
@@ -199,24 +104,14 @@ async def _backup(
|
||||
# entities table was backed up.
|
||||
async with conn.transaction(isolation="repeatable_read"):
|
||||
with zipfile.ZipFile(output_path, "w", zipfile.ZIP_DEFLATED) as zf:
|
||||
for i, table in enumerate(backup_tables, 1):
|
||||
typer.echo(f" [{i}/{len(backup_tables)}] Backing up {table}...", nl=False)
|
||||
for i, table in enumerate(BACKUP_TABLES, 1):
|
||||
typer.echo(f" [{i}/{len(BACKUP_TABLES)}] Backing up {table}...", nl=False)
|
||||
|
||||
buffer = io.BytesIO()
|
||||
|
||||
columns = await _table_columns(conn, schema, table)
|
||||
|
||||
# Pin the ordered columns into both the stream and manifest.
|
||||
# PostgreSQL binary COPY does not encode column identities, so
|
||||
# restore must validate this shape before truncating any data.
|
||||
# Use binary COPY for exact type preservation
|
||||
# asyncpg requires schema_name as separate parameter
|
||||
await conn.copy_from_table(
|
||||
table,
|
||||
schema_name=schema,
|
||||
columns=[column.name for column in columns],
|
||||
output=buffer,
|
||||
format="binary",
|
||||
)
|
||||
await conn.copy_from_table(table, schema_name=schema, output=buffer, format="binary")
|
||||
|
||||
data = buffer.getvalue()
|
||||
zf.writestr(f"{table}.bin", data)
|
||||
@@ -227,7 +122,6 @@ async def _backup(
|
||||
tables[table] = {
|
||||
"rows": row_count,
|
||||
"size_bytes": len(data),
|
||||
"columns": [{"name": column.name, "type_name": column.type_name} for column in columns],
|
||||
}
|
||||
|
||||
typer.echo(f" {row_count} rows")
|
||||
@@ -239,20 +133,8 @@ async def _backup(
|
||||
await conn.close()
|
||||
|
||||
|
||||
async def _restore(
|
||||
database_url: str,
|
||||
input_path: Path,
|
||||
schema: str = "public",
|
||||
backup_tables: list[str] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Restore all tables from a zip file using binary COPY protocol.
|
||||
|
||||
``backup_tables`` defaults to the core ``BACKUP_TABLES``; callers pass the
|
||||
extension-augmented list from ``_effective_backup_tables()``. Tables named
|
||||
here but absent from the archive are truncated then skipped for restore, so
|
||||
a stale extension registration never leaves pre-restore rows behind.
|
||||
"""
|
||||
backup_tables = backup_tables if backup_tables is not None else BACKUP_TABLES
|
||||
async def _restore(database_url: str, input_path: Path, schema: str = "public") -> dict[str, Any]:
|
||||
"""Restore all tables from a zip file using binary COPY protocol."""
|
||||
conn = await asyncpg.connect(database_url)
|
||||
try:
|
||||
with zipfile.ZipFile(input_path, "r") as zf:
|
||||
@@ -261,40 +143,29 @@ async def _restore(
|
||||
if manifest.get("version") != MANIFEST_VERSION:
|
||||
raise ValueError(f"Unsupported backup version: {manifest.get('version')}")
|
||||
|
||||
# Complete the compatibility check before entering the transaction
|
||||
# that truncates tables. This turns historical schema drift into an
|
||||
# actionable error without risking the target's existing data.
|
||||
restore_columns = await _validate_restore_schema(conn, manifest, schema)
|
||||
|
||||
# Use a transaction for atomic restore - either all tables are
|
||||
# restored or none are, preventing partial/inconsistent state.
|
||||
async with conn.transaction():
|
||||
typer.echo(" Clearing existing data...")
|
||||
# Truncate tables in reverse order (respects FK constraints)
|
||||
for table in reversed(backup_tables):
|
||||
for table in reversed(BACKUP_TABLES):
|
||||
qualified_table = _fq_table(table, schema)
|
||||
await conn.execute(f"TRUNCATE TABLE {qualified_table} CASCADE")
|
||||
|
||||
# Restore tables in forward order
|
||||
for i, table in enumerate(backup_tables, 1):
|
||||
for i, table in enumerate(BACKUP_TABLES, 1):
|
||||
filename = f"{table}.bin"
|
||||
if filename not in zf.namelist():
|
||||
typer.echo(f" [{i}/{len(backup_tables)}] {table}: skipped (not in backup)")
|
||||
typer.echo(f" [{i}/{len(BACKUP_TABLES)}] {table}: skipped (not in backup)")
|
||||
continue
|
||||
|
||||
expected_rows = manifest["tables"].get(table, {}).get("rows", "?")
|
||||
typer.echo(f" [{i}/{len(backup_tables)}] Restoring {table}... {expected_rows} rows")
|
||||
typer.echo(f" [{i}/{len(BACKUP_TABLES)}] Restoring {table}... {expected_rows} rows")
|
||||
|
||||
data = zf.read(filename)
|
||||
buffer = io.BytesIO(data)
|
||||
# asyncpg requires schema_name as separate parameter
|
||||
await conn.copy_to_table(
|
||||
table,
|
||||
schema_name=schema,
|
||||
columns=restore_columns[table],
|
||||
source=buffer,
|
||||
format="binary",
|
||||
)
|
||||
await conn.copy_to_table(table, schema_name=schema, source=buffer, format="binary")
|
||||
|
||||
# Refresh materialized view
|
||||
typer.echo(" Refreshing materialized views...")
|
||||
@@ -307,22 +178,20 @@ async def _restore(
|
||||
|
||||
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)
|
||||
return await _backup(resolved_url, output, schema, backup_tables=_effective_backup_tables())
|
||||
return await _backup(resolved_url, output, schema)
|
||||
|
||||
|
||||
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)
|
||||
return await _restore(resolved_url, input_file, schema, backup_tables=_effective_backup_tables())
|
||||
return await _restore(resolved_url, input_file, schema)
|
||||
|
||||
|
||||
@app.command()
|
||||
@@ -346,7 +215,7 @@ def backup(
|
||||
manifest = asyncio.run(_run_backup(config.database_url, output, schema))
|
||||
|
||||
total_rows = sum(t["rows"] for t in manifest["tables"].values())
|
||||
typer.echo(f"Backed up {total_rows} rows across {len(manifest['tables'])} tables")
|
||||
typer.echo(f"Backed up {total_rows} rows across {len(BACKUP_TABLES)} tables")
|
||||
typer.echo(f"Backup saved to {output}")
|
||||
|
||||
|
||||
@@ -379,7 +248,7 @@ def restore(
|
||||
manifest = asyncio.run(_run_restore(config.database_url, input_file, schema))
|
||||
|
||||
total_rows = sum(t["rows"] for t in manifest["tables"].values())
|
||||
typer.echo(f"Restored {total_rows} rows across {len(manifest['tables'])} tables")
|
||||
typer.echo(f"Restored {total_rows} rows across {len(BACKUP_TABLES)} tables")
|
||||
typer.echo("Restore complete")
|
||||
|
||||
|
||||
@@ -393,17 +262,17 @@ 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)
|
||||
|
||||
config = HindsightConfig.from_env()
|
||||
tenant_extension = load_extension("TENANT", TenantExtension)
|
||||
if schema:
|
||||
schemas = [schema]
|
||||
else:
|
||||
tenant_extension = load_extension("TENANT", TenantExtension)
|
||||
|
||||
schemas = [base_schema or DEFAULT_DATABASE_SCHEMA]
|
||||
if tenant_extension:
|
||||
tenants = await tenant_extension.list_tenants()
|
||||
@@ -428,36 +297,9 @@ async def _run_migration(
|
||||
ensure_extensions=ensure_extensions,
|
||||
)
|
||||
|
||||
# After core migrations, provision any extension-owned bank-scoped tables
|
||||
# per schema so extension schema evolves on the same lifecycle as core
|
||||
# schema (rather than via a lazy first-request path).
|
||||
if tenant_extension is not None:
|
||||
await _provision_extra_bank_tables(resolved_url, schemas, tenant_extension)
|
||||
|
||||
return schemas
|
||||
|
||||
|
||||
async def _provision_extra_bank_tables(
|
||||
resolved_url: str, schemas: list[str], tenant_extension: TenantExtension
|
||||
) -> None:
|
||||
"""Run the tenant extension's table provisioner for each migrated schema.
|
||||
|
||||
Fires after core migrations complete so extension-owned bank tables are
|
||||
created/evolved on the same lifecycle as core schema. A failure aborts the
|
||||
migration command (and names the offending schema) rather than being
|
||||
swallowed — provisioning is idempotent, so the operator can fix and re-run.
|
||||
"""
|
||||
for schema in schemas:
|
||||
conn = await asyncpg.connect(resolved_url)
|
||||
try:
|
||||
await tenant_extension.provision_bank_tables(conn, schema)
|
||||
except Exception as e:
|
||||
typer.echo(f" Failed to provision extension tables for schema '{schema}': {e}", err=True)
|
||||
raise
|
||||
finally:
|
||||
await conn.close()
|
||||
|
||||
|
||||
@app.command(name="run-db-migration")
|
||||
def run_db_migration(
|
||||
schema: str | None = typer.Option(
|
||||
@@ -512,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)
|
||||
@@ -647,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()
|
||||
|
||||
@@ -766,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)
|
||||
@@ -826,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)
|
||||
@@ -892,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)
|
||||
@@ -961,7 +665,6 @@ def worker_status(
|
||||
|
||||
|
||||
def main():
|
||||
load_dotenv_for_entrypoint()
|
||||
app()
|
||||
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
|
||||
-67
@@ -1,67 +0,0 @@
|
||||
"""Drop observation_history's FK to memory_units.
|
||||
|
||||
The history table records one snapshot per observation change, keyed by
|
||||
``(bank_id, observation_id)``. Its foreign key to ``memory_units`` existed only to
|
||||
cascade-delete history when the observation row went away.
|
||||
|
||||
That assumes every observation *is* a ``memory_units`` row, which is true only
|
||||
while Postgres is the memories store. When another store owns the memories the
|
||||
observation lives there and Postgres holds no row for it, so every history insert
|
||||
raises a foreign-key violation — swallowed by the writer as "a race with parallel
|
||||
consolidation" and logged at warning level. The audit trail goes silently empty.
|
||||
|
||||
Dropping the constraint lets history be recorded wherever the observation is
|
||||
stored. The cleanup the cascade used to do is now explicit, in the paths that
|
||||
delete observations (``_execute_delete_action``, ``clear_observations``,
|
||||
``delete_bank``). Rows orphaned by a path that misses — a document delete
|
||||
cascading through ``memory_units``, for instance — are invisible to readers,
|
||||
which always filter by ``(bank_id, observation_id)``, and are reclaimed when the
|
||||
bank is deleted.
|
||||
|
||||
Oracle builds this schema through its own DDL runner and never had the
|
||||
constraint, so the Oracle slot is a deliberate no-op.
|
||||
|
||||
Revision ID: a1c9e7f3b2d8
|
||||
Revises: c7d1e9a4b3f2
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import op
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "a1c9e7f3b2d8"
|
||||
down_revision: str | Sequence[str] | None = "c7d1e9a4b3f2"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
_CONSTRAINT = "observation_history_observation_id_fkey"
|
||||
|
||||
|
||||
def _pg_upgrade() -> None:
|
||||
op.execute(f"ALTER TABLE observation_history DROP CONSTRAINT IF EXISTS {_CONSTRAINT}")
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
# Re-adding the FK requires every row to reference a live memory_unit, so
|
||||
# clear any history whose observation is not a Postgres row first — those are
|
||||
# exactly the rows this migration made possible.
|
||||
op.execute(
|
||||
"DELETE FROM observation_history h "
|
||||
"WHERE NOT EXISTS (SELECT 1 FROM memory_units m WHERE m.id = h.observation_id)"
|
||||
)
|
||||
op.execute(
|
||||
f"ALTER TABLE observation_history ADD CONSTRAINT {_CONSTRAINT} "
|
||||
"FOREIGN KEY (observation_id) REFERENCES memory_units(id) ON DELETE CASCADE"
|
||||
)
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
# Oracle never had the constraint (its schema is built by a separate DDL
|
||||
# runner), so only Postgres has anything to drop.
|
||||
run_for_dialect(pg=_pg_upgrade, oracle=None)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
run_for_dialect(pg=_pg_downgrade, oracle=None)
|
||||
+52
@@ -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)
|
||||
-82
@@ -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)
|
||||
+2
-18
@@ -7,12 +7,8 @@ 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.
|
||||
|
||||
``managed`` lets a client tag a node as system-owned vs. hand-authored; it
|
||||
carries no server-side behaviour. A partial unique index keeps page names unique
|
||||
within a folder (case-insensitive; root pages compared under an empty parent).
|
||||
|
||||
Revision ID: a9b8c7d6e5f4
|
||||
Revises: a1c9e7f3b2d8
|
||||
Revises: b57a7c9e0d13
|
||||
Create Date: 2026-06-25
|
||||
"""
|
||||
|
||||
@@ -23,7 +19,7 @@ from alembic import context, op
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "a9b8c7d6e5f4"
|
||||
down_revision: str | Sequence[str] | None = "a1c9e7f3b2d8"
|
||||
down_revision: str | Sequence[str] | None = "b57a7c9e0d13"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
@@ -51,7 +47,6 @@ def _pg_upgrade() -> None:
|
||||
name TEXT NOT NULL,
|
||||
mental_model_id VARCHAR(64),
|
||||
sort_order INTEGER NOT NULL DEFAULT 0,
|
||||
managed BOOLEAN NOT NULL DEFAULT false,
|
||||
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),
|
||||
@@ -68,25 +63,15 @@ def _pg_upgrade() -> None:
|
||||
op.execute(
|
||||
f"CREATE INDEX IF NOT EXISTS idx_kp_bank_parent ON {schema}knowledge_pages (bank_id, parent_id, sort_order)"
|
||||
)
|
||||
# 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")
|
||||
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:
|
||||
# No case-insensitive unique index on Oracle: `name` is a CLOB and cannot be
|
||||
# indexed with lower(); page-name uniqueness is enforced on PG only.
|
||||
op.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS knowledge_pages (
|
||||
@@ -97,7 +82,6 @@ def _oracle_upgrade() -> None:
|
||||
name CLOB NOT NULL,
|
||||
mental_model_id VARCHAR2(64),
|
||||
sort_order NUMBER DEFAULT 0 NOT NULL,
|
||||
managed NUMBER(1) 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),
|
||||
|
||||
-259
@@ -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)
|
||||
+71
@@ -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)
|
||||
-90
@@ -1,90 +0,0 @@
|
||||
"""Add ``causal_links`` to the curation archive (invalidated_memory_units).
|
||||
|
||||
Causal edges (``caused_by`` and the historical ``causes``/``enables``/
|
||||
``prevents``) are retain-time extraction output: unlike temporal and semantic
|
||||
links they cannot be recomputed from dates or embeddings, and graph maintenance
|
||||
never rebuilds them. Invalidation MOVES a fact out of ``memory_units``, so the
|
||||
``memory_links → memory_units`` FK cascade deletes every incident edge — and
|
||||
revert had no way to bring the causal ones back (#2864).
|
||||
|
||||
This column parks the descriptors of the causal edges incident to an archived
|
||||
fact — ``[{"from_unit_id", "to_unit_id", "link_type", "weight"}, ...]`` — so
|
||||
revert can rematerialize them. It is deliberately unindexed and lives only on
|
||||
the archive: live facts keep their causal edges in ``memory_links`` (curation
|
||||
edits no longer delete them), and the archive is small, cold, and only read by
|
||||
low-frequency curation operations.
|
||||
|
||||
Revision ID: c7d1e9a4b3f2
|
||||
Revises: d7b2f8a1c934
|
||||
Create Date: 2026-07-24
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
from alembic import context, op
|
||||
|
||||
from hindsight_api.alembic._dialect import run_for_dialect
|
||||
|
||||
revision: str = "c7d1e9a4b3f2"
|
||||
down_revision: str | Sequence[str] | None = "d7b2f8a1c934"
|
||||
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()
|
||||
# NOT NULL DEFAULT is metadata-only on PG 11+, so this is cheap even on a
|
||||
# large archive. Existing rows read as "no causal edges captured" — edges
|
||||
# lost before this migration cannot be reconstructed and are not guessed.
|
||||
op.execute(
|
||||
f"ALTER TABLE {schema}invalidated_memory_units "
|
||||
f"ADD COLUMN IF NOT EXISTS causal_links JSONB NOT NULL DEFAULT '[]'::jsonb"
|
||||
)
|
||||
|
||||
|
||||
def _pg_downgrade() -> None:
|
||||
schema = _pg_schema_prefix()
|
||||
op.execute(f"ALTER TABLE {schema}invalidated_memory_units DROP COLUMN IF EXISTS causal_links")
|
||||
|
||||
|
||||
def _oracle_upgrade() -> None:
|
||||
# Kept in sync with PG for schema parity (curation itself is PostgreSQL-only
|
||||
# today — it introspects pg_attribute to move rows between the two tables).
|
||||
# Swallow ORA-01430 (column already exists) so the migration is idempotent.
|
||||
op.execute(
|
||||
"""
|
||||
BEGIN
|
||||
EXECUTE IMMEDIATE 'ALTER TABLE invalidated_memory_units ADD (causal_links CLOB DEFAULT ''[]''
|
||||
CONSTRAINT imu_causal_links_json CHECK (causal_links IS JSON))';
|
||||
EXCEPTION WHEN OTHERS THEN
|
||||
IF SQLCODE != -1430 THEN RAISE; END IF;
|
||||
END;
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def _oracle_downgrade() -> None:
|
||||
# Swallow ORA-00904 (column does not exist).
|
||||
op.execute(
|
||||
"""
|
||||
BEGIN
|
||||
EXECUTE IMMEDIATE 'ALTER TABLE invalidated_memory_units DROP COLUMN causal_links';
|
||||
EXCEPTION WHEN OTHERS THEN
|
||||
IF SQLCODE != -904 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)
|
||||
-150
@@ -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)
|
||||
-96
@@ -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)
|
||||
File diff suppressed because it is too large
Load Diff
+113
-21
@@ -1,40 +1,76 @@
|
||||
"""Markdown rendering for knowledge pages.
|
||||
"""Open Knowledge Format (OKF) projection for knowledge pages.
|
||||
|
||||
Knowledge pages render as *read-only* markdown documents over the existing mental
|
||||
models: each mental model becomes a markdown body with a YAML frontmatter block
|
||||
(``type`` required; ``title``/``description``/``tags``/``timestamp`` optional).
|
||||
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 rendering unit-testable without a DB or
|
||||
LLM and lets the HTTP layer stay a thin wrapper.
|
||||
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 dataclasses import dataclass
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
# Every page carries exactly one ``type`` frontmatter field. We default to this
|
||||
# when a page does not declare one via a ``type:<x>`` tag.
|
||||
# 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 ``type`` through a tag of the form ``type:runbook``.
|
||||
# This keeps rendering schema-free (no new mental_models column): the type is
|
||||
# lifted from the existing tags array.
|
||||
# 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 ``type`` and the tags that remain after the type tag is split off."""
|
||||
"""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.
|
||||
|
||||
@@ -47,11 +83,11 @@ def _scalar(value: Any) -> str:
|
||||
|
||||
|
||||
def page_type(tags: list[str] | None) -> PageType:
|
||||
"""Split a ``type`` out of the tag list.
|
||||
"""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 leak into the page's displayed tags.
|
||||
Falls back to :data:`DEFAULT_PAGE_TYPE`.
|
||||
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] = []
|
||||
@@ -70,7 +106,7 @@ def _timestamp(mm: dict[str, Any]) -> str | None:
|
||||
|
||||
|
||||
def frontmatter(mm: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Build the ordered frontmatter mapping for a mental model.
|
||||
"""Build the ordered OKF frontmatter mapping for a mental model.
|
||||
|
||||
``None``/empty values are dropped by :func:`render_frontmatter`.
|
||||
"""
|
||||
@@ -103,23 +139,23 @@ def render_frontmatter(fm: dict[str, Any]) -> str:
|
||||
|
||||
|
||||
def render_document(mm: dict[str, Any]) -> str:
|
||||
"""Render a full markdown document: frontmatter block + markdown body."""
|
||||
"""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:
|
||||
"""Bundle filename for a page id."""
|
||||
"""OKF bundle filename for a page id."""
|
||||
return f"{page_id}.md"
|
||||
|
||||
|
||||
def log_filename(page_id: str) -> str:
|
||||
"""Reserved per-page history filename."""
|
||||
"""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 markdown navigation over the tree.
|
||||
"""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``.
|
||||
@@ -169,3 +205,59 @@ def render_log(mm: dict[str, Any], history: list[dict[str, Any]]) -> str:
|
||||
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)
|
||||
@@ -66,6 +66,11 @@ def color_end(text: str) -> str:
|
||||
return color(text, 1.0)
|
||||
|
||||
|
||||
def color_mid(text: str) -> str:
|
||||
"""Color text with gradient middle color."""
|
||||
return color(text, 0.5)
|
||||
|
||||
|
||||
def dim(text: str) -> str:
|
||||
"""Dim/gray text."""
|
||||
return f"\033[38;2;128;128;128m{text}\033[0m"
|
||||
|
||||
@@ -19,28 +19,12 @@ from ._pg_search import normalize_pg_search_tokenizer
|
||||
from ._vector_index import validate_extension
|
||||
from .utils import mask_network_location
|
||||
|
||||
# Load .env file, searching current and parent directories (overrides existing env vars)
|
||||
load_dotenv(find_dotenv(usecwd=True), override=True)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def load_dotenv_for_entrypoint() -> None:
|
||||
"""Load a discovered ``.env`` file for Hindsight's own entry points.
|
||||
|
||||
Importing ``hindsight_api`` (or anything that pulls it in) must NOT mutate
|
||||
the host application's ``os.environ``. See issue #2961: doing so at module
|
||||
scope let an upward ``.env`` walk from the process cwd silently overwrite an
|
||||
embedding application's own configuration.
|
||||
|
||||
This helper is therefore called explicitly from Hindsight's standalone entry
|
||||
points only — the API server (CLI and ``hindsight_api.server:app``), the
|
||||
worker, and the admin CLI. ``override=True`` is deliberate: it preserves the
|
||||
exact precedence those entry points have always had (a discovered ``.env``
|
||||
is authoritative over the ambient process environment). Because a library
|
||||
import never reaches this code path, that precedence no longer leaks into
|
||||
embedders.
|
||||
"""
|
||||
load_dotenv(find_dotenv(usecwd=True), override=True)
|
||||
|
||||
|
||||
class ConfigFieldAccessError(AttributeError):
|
||||
"""Raised when trying to access a bank-configurable field from global config."""
|
||||
|
||||
@@ -161,18 +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_SUPPORTS_MAX_ITEMS = "HINDSIGHT_API_LLM_SUPPORTS_MAX_ITEMS"
|
||||
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
|
||||
@@ -274,35 +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")
|
||||
|
||||
|
||||
def _parse_boolean_env(env_name: str, default: bool) -> bool:
|
||||
"""Parse a boolean environment variable, rejecting ambiguous values."""
|
||||
raw = os.getenv(env_name)
|
||||
if raw is None:
|
||||
return default
|
||||
|
||||
normalized = raw.strip().lower()
|
||||
if normalized in ("true", "1"):
|
||||
return True
|
||||
if normalized in ("false", "0"):
|
||||
return False
|
||||
raise ValueError(f"Invalid {env_name} value {raw!r}: expected true, false, 1, or 0")
|
||||
|
||||
|
||||
# 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"
|
||||
@@ -314,7 +259,6 @@ ENV_RETAIN_LLM_INITIAL_BACKOFF = "HINDSIGHT_API_RETAIN_LLM_INITIAL_BACKOFF"
|
||||
ENV_RETAIN_LLM_MAX_BACKOFF = "HINDSIGHT_API_RETAIN_LLM_MAX_BACKOFF"
|
||||
ENV_RETAIN_LLM_TIMEOUT = "HINDSIGHT_API_RETAIN_LLM_TIMEOUT"
|
||||
ENV_RETAIN_LLM_LITELLMROUTER_CONFIG = "HINDSIGHT_API_RETAIN_LLM_LITELLMROUTER_CONFIG"
|
||||
ENV_RETAIN_LLM_REASONING_EFFORT = "HINDSIGHT_API_RETAIN_LLM_REASONING_EFFORT"
|
||||
|
||||
# Fireworks AI batch inference. Fireworks' batch API is a proprietary
|
||||
# account-scoped dataset/job REST API on a control-plane host, distinct from the
|
||||
@@ -337,7 +281,6 @@ ENV_REFLECT_LLM_INITIAL_BACKOFF = "HINDSIGHT_API_REFLECT_LLM_INITIAL_BACKOFF"
|
||||
ENV_REFLECT_LLM_MAX_BACKOFF = "HINDSIGHT_API_REFLECT_LLM_MAX_BACKOFF"
|
||||
ENV_REFLECT_LLM_TIMEOUT = "HINDSIGHT_API_REFLECT_LLM_TIMEOUT"
|
||||
ENV_REFLECT_LLM_LITELLMROUTER_CONFIG = "HINDSIGHT_API_REFLECT_LLM_LITELLMROUTER_CONFIG"
|
||||
ENV_REFLECT_LLM_REASONING_EFFORT = "HINDSIGHT_API_REFLECT_LLM_REASONING_EFFORT"
|
||||
|
||||
ENV_CONSOLIDATION_LLM_PROVIDER = "HINDSIGHT_API_CONSOLIDATION_LLM_PROVIDER"
|
||||
ENV_CONSOLIDATION_LLM_API_KEY = "HINDSIGHT_API_CONSOLIDATION_LLM_API_KEY"
|
||||
@@ -349,12 +292,10 @@ ENV_CONSOLIDATION_LLM_INITIAL_BACKOFF = "HINDSIGHT_API_CONSOLIDATION_LLM_INITIAL
|
||||
ENV_CONSOLIDATION_LLM_MAX_BACKOFF = "HINDSIGHT_API_CONSOLIDATION_LLM_MAX_BACKOFF"
|
||||
ENV_CONSOLIDATION_LLM_TIMEOUT = "HINDSIGHT_API_CONSOLIDATION_LLM_TIMEOUT"
|
||||
ENV_CONSOLIDATION_LLM_LITELLMROUTER_CONFIG = "HINDSIGHT_API_CONSOLIDATION_LLM_LITELLMROUTER_CONFIG"
|
||||
ENV_CONSOLIDATION_LLM_REASONING_EFFORT = "HINDSIGHT_API_CONSOLIDATION_LLM_REASONING_EFFORT"
|
||||
|
||||
ENV_EMBEDDINGS_PROVIDER = "HINDSIGHT_API_EMBEDDINGS_PROVIDER"
|
||||
ENV_EMBEDDINGS_LOCAL_MODEL = "HINDSIGHT_API_EMBEDDINGS_LOCAL_MODEL"
|
||||
ENV_EMBEDDINGS_LOCAL_FORCE_CPU = "HINDSIGHT_API_EMBEDDINGS_LOCAL_FORCE_CPU"
|
||||
ENV_EMBEDDINGS_LOCAL_ALLOW_MPS = "HINDSIGHT_API_EMBEDDINGS_LOCAL_ALLOW_MPS"
|
||||
ENV_EMBEDDINGS_LOCAL_TRUST_REMOTE_CODE = "HINDSIGHT_API_EMBEDDINGS_LOCAL_TRUST_REMOTE_CODE"
|
||||
ENV_EMBEDDINGS_ONNX_MODEL_ID = "HINDSIGHT_API_EMBEDDINGS_ONNX_MODEL_ID"
|
||||
ENV_EMBEDDINGS_ONNX_MODEL_PATH = "HINDSIGHT_API_EMBEDDINGS_ONNX_MODEL_PATH"
|
||||
@@ -432,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"
|
||||
@@ -442,10 +382,8 @@ 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_ALLOW_MPS = "HINDSIGHT_API_RERANKER_LOCAL_ALLOW_MPS"
|
||||
ENV_RERANKER_LOCAL_MAX_CONCURRENT = "HINDSIGHT_API_RERANKER_LOCAL_MAX_CONCURRENT"
|
||||
ENV_RERANKER_LOCAL_TRUST_REMOTE_CODE = "HINDSIGHT_API_RERANKER_LOCAL_TRUST_REMOTE_CODE"
|
||||
ENV_RERANKER_LOCAL_FP16 = "HINDSIGHT_API_RERANKER_LOCAL_FP16"
|
||||
@@ -465,9 +403,6 @@ ENV_RERANKER_LITELLM_SDK_TIMEOUT = "HINDSIGHT_API_RERANKER_LITELLM_SDK_TIMEOUT"
|
||||
ENV_RERANKER_GOOGLE_TIMEOUT = "HINDSIGHT_API_RERANKER_GOOGLE_TIMEOUT"
|
||||
ENV_RERANKER_MAX_CANDIDATES = "HINDSIGHT_API_RERANKER_MAX_CANDIDATES"
|
||||
ENV_SEMANTIC_MIN_SIMILARITY = "HINDSIGHT_API_SEMANTIC_MIN_SIMILARITY"
|
||||
ENV_GRAPH_SEED_MIN_SIMILARITY = "HINDSIGHT_API_GRAPH_SEED_MIN_SIMILARITY"
|
||||
ENV_TEMPORAL_SEMANTIC_MIN_SIMILARITY = "HINDSIGHT_API_TEMPORAL_SEMANTIC_MIN_SIMILARITY"
|
||||
ENV_SEMANTIC_LINK_MIN_SIMILARITY = "HINDSIGHT_API_SEMANTIC_LINK_MIN_SIMILARITY"
|
||||
ENV_RERANKER_FLASHRANK_MODEL = "HINDSIGHT_API_RERANKER_FLASHRANK_MODEL"
|
||||
ENV_RERANKER_FLASHRANK_CACHE_DIR = "HINDSIGHT_API_RERANKER_FLASHRANK_CACHE_DIR"
|
||||
ENV_RERANKER_FLASHRANK_CPU_MEM_ARENA = "HINDSIGHT_API_RERANKER_FLASHRANK_CPU_MEM_ARENA"
|
||||
@@ -532,12 +467,6 @@ ENV_OTEL_DEPLOYMENT_ENVIRONMENT = "HINDSIGHT_API_OTEL_DEPLOYMENT_ENVIRONMENT"
|
||||
ENV_METRICS_INCLUDE_BANK_ID = "HINDSIGHT_API_METRICS_INCLUDE_BANK_ID"
|
||||
ENV_METRICS_BACKLOG_ENABLED = "HINDSIGHT_API_METRICS_BACKLOG_ENABLED"
|
||||
|
||||
# Runtime-stall observability (loop watchdog + DB pool acquire instrumentation)
|
||||
ENV_LOOP_WATCHDOG_ENABLED = "HINDSIGHT_API_LOOP_WATCHDOG_ENABLED"
|
||||
ENV_LOOP_WATCHDOG_STALL_THRESHOLD_MS = "HINDSIGHT_API_LOOP_WATCHDOG_STALL_THRESHOLD_MS"
|
||||
ENV_LOOP_WATCHDOG_POLL_INTERVAL_MS = "HINDSIGHT_API_LOOP_WATCHDOG_POLL_INTERVAL_MS"
|
||||
ENV_DB_ACQUIRE_WARN_THRESHOLD_MS = "HINDSIGHT_API_DB_ACQUIRE_WARN_THRESHOLD_MS"
|
||||
|
||||
# Vertex AI configuration
|
||||
ENV_LLM_VERTEXAI_PROJECT_ID = "HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID"
|
||||
ENV_LLM_VERTEXAI_REGION = "HINDSIGHT_API_LLM_VERTEXAI_REGION"
|
||||
@@ -556,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"
|
||||
@@ -596,7 +520,6 @@ ENV_FILE_PARSER_MARKITDOWN_OCR_API_KEY = "HINDSIGHT_API_FILE_PARSER_MARKITDOWN_O
|
||||
ENV_FILE_PARSER_MARKITDOWN_OCR_BASE_URL = "HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_BASE_URL"
|
||||
ENV_FILE_PARSER_MARKITDOWN_OCR_MODEL = "HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_MODEL"
|
||||
ENV_FILE_PARSER_MARKITDOWN_OCR_PROMPT = "HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_PROMPT"
|
||||
ENV_FILE_PARSER_MARKITDOWN_OCR_DEFAULT_HEADERS = "HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_DEFAULT_HEADERS"
|
||||
ENV_FILE_PARSER_IRIS_TOKEN = "HINDSIGHT_API_FILE_PARSER_IRIS_TOKEN"
|
||||
ENV_FILE_PARSER_IRIS_ORG_ID = "HINDSIGHT_API_FILE_PARSER_IRIS_ORG_ID"
|
||||
ENV_FILE_PARSER_LLAMA_PARSE_API_KEY = "HINDSIGHT_API_FILE_PARSER_LLAMA_PARSE_API_KEY"
|
||||
@@ -661,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
|
||||
@@ -676,71 +598,25 @@ 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 worker slot reservations: op_type -> default reserved count.
|
||||
# Each entry reserves a guaranteed *minimum* number of slots (a floor) for that
|
||||
# operation type within the global WORKER_MAX_SLOTS pool, so a saturated pool can't
|
||||
# starve it. Remaining capacity (WORKER_MAX_SLOTS - sum of reservations) is a shared
|
||||
# pool usable by any type, so a reservation does NOT cap the type — it may overflow
|
||||
# the shared pool. Adding a type here is the only change needed to make it reservable.
|
||||
#
|
||||
# op_type matches the value stored in async_operations.operation_type; the env var
|
||||
# names are derived from it (see _parse_worker_slot_reservations).
|
||||
WORKER_SLOT_TYPE_DEFAULTS: dict[str, int] = {
|
||||
"consolidation": 2,
|
||||
"retain": 0,
|
||||
"file_convert_retain": 0,
|
||||
"refresh_mental_model": 0,
|
||||
"graph_maintenance": 0,
|
||||
"import_documents": 0,
|
||||
# Per-operation-type slot reservations. Each entry maps an operation_type
|
||||
# (as stored in async_operations.operation_type) to its env var and default.
|
||||
# Adding a new operation type here is the ONLY change needed to make it
|
||||
# reservable via env var — config fields, from_env(), and the
|
||||
# worker_slot_reservations property all derive from this dict.
|
||||
WORKER_SLOT_RESERVATION_TYPES: dict[str, tuple[str, int]] = {
|
||||
"consolidation": ("HINDSIGHT_API_WORKER_CONSOLIDATION_MAX_SLOTS", 2),
|
||||
"retain": ("HINDSIGHT_API_WORKER_RETAIN_MAX_SLOTS", 0),
|
||||
"file_convert_retain": ("HINDSIGHT_API_WORKER_FILE_CONVERT_RETAIN_MAX_SLOTS", 0),
|
||||
"refresh_mental_model": ("HINDSIGHT_API_WORKER_REFRESH_MENTAL_MODEL_MAX_SLOTS", 0),
|
||||
"graph_maintenance": ("HINDSIGHT_API_WORKER_GRAPH_MAINTENANCE_MAX_SLOTS", 0),
|
||||
"import_documents": ("HINDSIGHT_API_WORKER_IMPORT_DOCUMENTS_MAX_SLOTS", 0),
|
||||
}
|
||||
|
||||
|
||||
def _parse_worker_slot_reservations() -> dict[str, int]:
|
||||
"""Parse per-type RESERVED_SLOTS (and the deprecated _MAX_SLOTS alias).
|
||||
|
||||
``HINDSIGHT_API_WORKER_<TYPE>_MAX_SLOTS`` is a deprecated alias for
|
||||
``..._RESERVED_SLOTS`` — despite its name it always set the reservation floor,
|
||||
never a ceiling. It still works but logs a warning. Returns op_type -> reserved
|
||||
floor for entries with a reservation > 0.
|
||||
"""
|
||||
reservations: dict[str, int] = {}
|
||||
for op_type, default in WORKER_SLOT_TYPE_DEFAULTS.items():
|
||||
reserved_env = f"HINDSIGHT_API_WORKER_{op_type.upper()}_RESERVED_SLOTS"
|
||||
legacy_env = f"HINDSIGHT_API_WORKER_{op_type.upper()}_MAX_SLOTS"
|
||||
raw_reserved = os.getenv(reserved_env)
|
||||
raw_legacy = os.getenv(legacy_env)
|
||||
if raw_reserved is not None and raw_legacy is not None:
|
||||
raise ValueError(
|
||||
f"Both {reserved_env} and the deprecated {legacy_env} are set; "
|
||||
f"they configure the same value. Keep only {reserved_env}."
|
||||
)
|
||||
if raw_legacy is not None:
|
||||
logger.warning(
|
||||
"%s is deprecated and will be removed in a future release. Despite its name it "
|
||||
"reserves a *minimum* (floor), not a maximum. Rename it to %s.",
|
||||
legacy_env,
|
||||
reserved_env,
|
||||
)
|
||||
reserved_source = raw_reserved if raw_reserved is not None else raw_legacy
|
||||
reserved = int(reserved_source) if reserved_source is not None else default
|
||||
if reserved < 0:
|
||||
raise ValueError(f"{reserved_env} must be >= 0, got {reserved}")
|
||||
if reserved > 0:
|
||||
reservations[op_type] = reserved
|
||||
return reservations
|
||||
|
||||
|
||||
ENV_WORKER_CONSOLIDATION_BANK_PRIORITY = "HINDSIGHT_API_WORKER_CONSOLIDATION_BANK_PRIORITY"
|
||||
ENV_RETAIN_MAX_CONCURRENT = "HINDSIGHT_API_RETAIN_MAX_CONCURRENT"
|
||||
ENV_RETAIN_WALL_TIMEOUT = "HINDSIGHT_API_RETAIN_WALL_TIMEOUT"
|
||||
|
||||
# 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"
|
||||
@@ -762,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.
|
||||
@@ -782,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"
|
||||
@@ -856,7 +725,6 @@ DEFAULT_LLAMACPP_EXTRA_ARGS = None # Space-separated extra CLI args for llama.c
|
||||
# (prose preambles, markdown fences, invalid JSON) — wedging retain/consolidation
|
||||
# on parse retries.
|
||||
DEFAULT_LLM_STRICT_SCHEMA = False
|
||||
DEFAULT_LLM_SUPPORTS_MAX_ITEMS = True
|
||||
|
||||
DEFAULT_LLM_MAX_CONCURRENT = 32
|
||||
DEFAULT_LLM_MAX_RETRIES = 3 # Max retry attempts for LLM API calls
|
||||
@@ -877,9 +745,6 @@ DEFAULT_LLM_GEMINI_SAFETY_SETTINGS = None # None = use Gemini default safety se
|
||||
DEFAULT_EMBEDDINGS_PROVIDER = "local"
|
||||
DEFAULT_EMBEDDINGS_LOCAL_MODEL = "BAAI/bge-small-en-v1.5"
|
||||
DEFAULT_EMBEDDINGS_LOCAL_FORCE_CPU = False # Force CPU mode for local embeddings
|
||||
# Apple Silicon MPS is opt-in: it leaks memory under variable-length workloads
|
||||
# (unbounded per-shape kernel/allocator cache). CUDA/XPU still auto-select.
|
||||
DEFAULT_EMBEDDINGS_LOCAL_ALLOW_MPS = False
|
||||
DEFAULT_EMBEDDINGS_LOCAL_TRUST_REMOTE_CODE = False # Security: disabled by default, required for some models
|
||||
DEFAULT_EMBEDDINGS_ONNX_MODEL_ID = "intfloat/multilingual-e5-small"
|
||||
DEFAULT_EMBEDDINGS_ONNX_FILE = "onnx/model.onnx"
|
||||
@@ -896,12 +761,8 @@ 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
|
||||
# Apple Silicon MPS is opt-in: it leaks memory under variable-length workloads
|
||||
# (unbounded per-shape kernel/allocator cache). CUDA/XPU still auto-select.
|
||||
DEFAULT_RERANKER_LOCAL_ALLOW_MPS = False
|
||||
DEFAULT_RERANKER_LOCAL_MAX_CONCURRENT = 4 # Limit concurrent CPU-bound reranking to prevent thrashing
|
||||
DEFAULT_RERANKER_LOCAL_TRUST_REMOTE_CODE = (
|
||||
False # Security: disabled by default, required for some models like jina-reranker-v2
|
||||
@@ -924,16 +785,10 @@ DEFAULT_RERANKER_LITELLM_SDK_TIMEOUT = 60.0
|
||||
DEFAULT_RERANKER_GOOGLE_TIMEOUT = 60.0
|
||||
DEFAULT_RERANKER_MAX_CANDIDATES = 300
|
||||
DEFAULT_SEMANTIC_MIN_SIMILARITY = 0.3
|
||||
DEFAULT_GRAPH_SEED_MIN_SIMILARITY = 0.3
|
||||
DEFAULT_TEMPORAL_SEMANTIC_MIN_SIMILARITY = 0.1
|
||||
DEFAULT_SEMANTIC_LINK_MIN_SIMILARITY = 0.7
|
||||
# Minimum BM25 score a row must exceed to enter fusion. 0.0 gates out
|
||||
# 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.
|
||||
@@ -1054,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"
|
||||
@@ -1108,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
|
||||
@@ -1190,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)
|
||||
@@ -1208,28 +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.
|
||||
# Wall-clock ceiling for one retain task in the worker (0 disables). This is a
|
||||
# deadlock/wedge backstop, not a latency target: a retain that blocks forever on
|
||||
# a lock, an LLM permit or a queue put would otherwise hold its worker slot until
|
||||
# the process restarts, and 'processing' is neither retryable nor cancellable
|
||||
# through the API. Set well above any healthy retain so it only ever fires on a
|
||||
# genuine wedge — the per-attempt LLM timeout and the retry budget already bound
|
||||
# the normal slow path.
|
||||
DEFAULT_RETAIN_WALL_TIMEOUT = 3600 # seconds (1 hour)
|
||||
|
||||
# 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)
|
||||
@@ -1265,26 +1089,11 @@ DEFAULT_OTEL_DEPLOYMENT_ENVIRONMENT = "development"
|
||||
DEFAULT_METRICS_INCLUDE_BANK_ID = False # Disabled by default to avoid high-cardinality OTel metric growth
|
||||
DEFAULT_METRICS_BACKLOG_ENABLED = False # Disabled by default: runs periodic per-schema COUNT queries
|
||||
|
||||
# Runtime-stall observability defaults. Both are cheap and on by default: the
|
||||
# watchdog is a single background thread pinging the loop; the DB-pool acquire
|
||||
# timing is a monotonic() delta per acquire. They turn a failing liveness probe
|
||||
# from "pod restarted, cause unknown" into a logged root cause (blocked loop vs
|
||||
# pool exhaustion).
|
||||
DEFAULT_LOOP_WATCHDOG_ENABLED = True
|
||||
DEFAULT_LOOP_WATCHDOG_STALL_THRESHOLD_MS = 1000 # log a stall once the loop is unresponsive this long
|
||||
DEFAULT_LOOP_WATCHDOG_POLL_INTERVAL_MS = 250 # how often the watchdog thread pings the loop
|
||||
DEFAULT_DB_ACQUIRE_WARN_THRESHOLD_MS = 1000 # log a warning when a pool acquire waits this long
|
||||
|
||||
# Audit log defaults
|
||||
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
|
||||
@@ -1400,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 == "":
|
||||
@@ -1420,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}")
|
||||
@@ -1795,24 +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
|
||||
llm_supports_max_items: bool = field(
|
||||
default=DEFAULT_LLM_SUPPORTS_MAX_ITEMS,
|
||||
kw_only=True,
|
||||
) # Whether structured-output schemas accept JSON Schema maxItems
|
||||
# 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
|
||||
@@ -1840,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
|
||||
@@ -1863,7 +1623,6 @@ class HindsightConfig:
|
||||
retain_llm_max_backoff: float | None
|
||||
retain_llm_timeout: float | None
|
||||
retain_llm_litellmrouter_config: dict | None
|
||||
retain_llm_reasoning_effort: str | None
|
||||
|
||||
# Fireworks AI batch inference (static, server-level)
|
||||
fireworks_account_id: str | None
|
||||
@@ -1880,7 +1639,6 @@ class HindsightConfig:
|
||||
reflect_llm_max_backoff: float | None
|
||||
reflect_llm_timeout: float | None
|
||||
reflect_llm_litellmrouter_config: dict | None
|
||||
reflect_llm_reasoning_effort: str | None
|
||||
|
||||
consolidation_llm_provider: str | None
|
||||
consolidation_llm_api_key: str | None
|
||||
@@ -1892,13 +1650,11 @@ class HindsightConfig:
|
||||
consolidation_llm_max_backoff: float | None
|
||||
consolidation_llm_timeout: float | None
|
||||
consolidation_llm_litellmrouter_config: dict | None
|
||||
consolidation_llm_reasoning_effort: str | None
|
||||
|
||||
# Embeddings
|
||||
embeddings_provider: str
|
||||
embeddings_local_model: str
|
||||
embeddings_local_force_cpu: bool
|
||||
embeddings_local_allow_mps: bool
|
||||
embeddings_local_trust_remote_code: bool
|
||||
embeddings_onnx_model_id: str
|
||||
embeddings_onnx_model_path: str | None
|
||||
@@ -1929,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
|
||||
@@ -1941,10 +1696,8 @@ class HindsightConfig:
|
||||
|
||||
# Reranker
|
||||
reranker_provider: str
|
||||
reranker_send_bank_as_header: bool
|
||||
reranker_local_model: str
|
||||
reranker_local_force_cpu: bool
|
||||
reranker_local_allow_mps: bool
|
||||
reranker_local_max_concurrent: int
|
||||
reranker_local_trust_remote_code: bool
|
||||
reranker_local_fp16: bool
|
||||
@@ -1956,9 +1709,6 @@ class HindsightConfig:
|
||||
reranker_tei_http_timeout: float
|
||||
reranker_max_candidates: int
|
||||
semantic_min_similarity: float
|
||||
graph_seed_min_similarity: float
|
||||
temporal_semantic_min_similarity: float
|
||||
semantic_link_min_similarity: float
|
||||
bm25_min_score: float
|
||||
recall_max_candidates_per_source: int
|
||||
recall_strategy_boosts: dict[str, str]
|
||||
@@ -2146,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)
|
||||
@@ -2159,16 +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
|
||||
retain_wall_timeout: 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
|
||||
@@ -2179,25 +1924,11 @@ class HindsightConfig:
|
||||
metrics_include_bank_id: bool
|
||||
metrics_backlog_enabled: bool
|
||||
|
||||
# Runtime-stall observability (static, server-level only)
|
||||
loop_watchdog_enabled: bool
|
||||
loop_watchdog_stall_threshold_ms: int
|
||||
loop_watchdog_poll_interval_ms: int
|
||||
db_acquire_warn_threshold_ms: int
|
||||
|
||||
# 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)
|
||||
@@ -2234,7 +1965,6 @@ class HindsightConfig:
|
||||
file_parser_markitdown_ocr_base_url: str | None = None
|
||||
file_parser_markitdown_ocr_model: str | None = None
|
||||
file_parser_markitdown_ocr_prompt: str = DEFAULT_FILE_PARSER_MARKITDOWN_OCR_PROMPT
|
||||
file_parser_markitdown_ocr_default_headers: dict | None = None
|
||||
|
||||
# Multi-LLM chains (static, server-level). Index 0 of each chain is the
|
||||
# corresponding unindexed/base LLM config above; these hold the extra indexed
|
||||
@@ -2249,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
|
||||
|
||||
@@ -2307,13 +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",
|
||||
# Persist raw source text (documents.original_text / chunks.chunk_text).
|
||||
# Per-bank so a data-minimizing bank can keep only derived facts while
|
||||
# others retain the raw source for expansion/re-extraction.
|
||||
"store_document_text",
|
||||
# Retention settings (behavioral)
|
||||
"retain_chunk_size",
|
||||
"retain_structured_chunk_size",
|
||||
@@ -2452,18 +2174,10 @@ class HindsightConfig:
|
||||
self.text_search_extension_pg_search_tokenizer
|
||||
)
|
||||
|
||||
for field_name in (
|
||||
"semantic_min_similarity",
|
||||
"graph_seed_min_similarity",
|
||||
"temporal_semantic_min_similarity",
|
||||
"semantic_link_min_similarity",
|
||||
):
|
||||
value = getattr(self, field_name)
|
||||
if not 0.0 <= value <= 1.0:
|
||||
raise ValueError(f"Invalid {field_name}: {value}. 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")
|
||||
if not 0.0 <= self.semantic_min_similarity <= 1.0:
|
||||
raise ValueError(
|
||||
f"Invalid semantic_min_similarity: {self.semantic_min_similarity}. Must be between 0.0 and 1.0"
|
||||
)
|
||||
|
||||
# Validate bedrock_service_tier
|
||||
valid_bedrock_tiers = (None, "flex", "priority", "reserved")
|
||||
@@ -2554,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."""
|
||||
@@ -2568,9 +2275,6 @@ class HindsightConfig:
|
||||
llm_provider = os.getenv(ENV_LLM_PROVIDER, DEFAULT_LLM_PROVIDER)
|
||||
llm_model = os.getenv(ENV_LLM_MODEL) or _get_default_model_for_provider(llm_provider)
|
||||
|
||||
# Parse per-type worker slot reservations (floors) once.
|
||||
worker_slot_reservations = _parse_worker_slot_reservations()
|
||||
|
||||
config = cls(
|
||||
# Database
|
||||
database_backend=os.getenv(ENV_DATABASE_BACKEND, DEFAULT_DATABASE_BACKEND).lower(),
|
||||
@@ -2613,19 +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_supports_max_items=_parse_boolean_env(
|
||||
ENV_LLM_SUPPORTS_MAX_ITEMS,
|
||||
DEFAULT_LLM_SUPPORTS_MAX_ITEMS,
|
||||
),
|
||||
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
|
||||
),
|
||||
@@ -2650,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))),
|
||||
@@ -2689,7 +2380,6 @@ class HindsightConfig:
|
||||
else None,
|
||||
retain_llm_timeout=float(os.getenv(ENV_RETAIN_LLM_TIMEOUT)) if os.getenv(ENV_RETAIN_LLM_TIMEOUT) else None,
|
||||
retain_llm_litellmrouter_config=_parse_llm_router_config(ENV_RETAIN_LLM_LITELLMROUTER_CONFIG),
|
||||
retain_llm_reasoning_effort=os.getenv(ENV_RETAIN_LLM_REASONING_EFFORT) or None,
|
||||
reflect_llm_provider=os.getenv(ENV_REFLECT_LLM_PROVIDER) or None,
|
||||
reflect_llm_api_key=os.getenv(ENV_REFLECT_LLM_API_KEY) or None,
|
||||
reflect_llm_model=os.getenv(ENV_REFLECT_LLM_MODEL)
|
||||
@@ -2715,7 +2405,6 @@ class HindsightConfig:
|
||||
if os.getenv(ENV_REFLECT_LLM_TIMEOUT)
|
||||
else None,
|
||||
reflect_llm_litellmrouter_config=_parse_llm_router_config(ENV_REFLECT_LLM_LITELLMROUTER_CONFIG),
|
||||
reflect_llm_reasoning_effort=os.getenv(ENV_REFLECT_LLM_REASONING_EFFORT) or None,
|
||||
consolidation_llm_provider=os.getenv(ENV_CONSOLIDATION_LLM_PROVIDER) or None,
|
||||
consolidation_llm_api_key=os.getenv(ENV_CONSOLIDATION_LLM_API_KEY) or None,
|
||||
consolidation_llm_model=os.getenv(ENV_CONSOLIDATION_LLM_MODEL)
|
||||
@@ -2741,7 +2430,6 @@ class HindsightConfig:
|
||||
if os.getenv(ENV_CONSOLIDATION_LLM_TIMEOUT)
|
||||
else None,
|
||||
consolidation_llm_litellmrouter_config=_parse_llm_router_config(ENV_CONSOLIDATION_LLM_LITELLMROUTER_CONFIG),
|
||||
consolidation_llm_reasoning_effort=os.getenv(ENV_CONSOLIDATION_LLM_REASONING_EFFORT) or None,
|
||||
# Multi-LLM chains (indexed members + routing strategy)
|
||||
llm_members=_parse_llm_members(""),
|
||||
llm_strategy=_parse_llm_strategy(os.getenv(ENV_LLM_STRATEGY)),
|
||||
@@ -2758,10 +2446,6 @@ class HindsightConfig:
|
||||
ENV_EMBEDDINGS_LOCAL_FORCE_CPU, str(DEFAULT_EMBEDDINGS_LOCAL_FORCE_CPU)
|
||||
).lower()
|
||||
in ("true", "1"),
|
||||
embeddings_local_allow_mps=os.getenv(
|
||||
ENV_EMBEDDINGS_LOCAL_ALLOW_MPS, str(DEFAULT_EMBEDDINGS_LOCAL_ALLOW_MPS)
|
||||
).lower()
|
||||
in ("true", "1"),
|
||||
embeddings_local_trust_remote_code=os.getenv(
|
||||
ENV_EMBEDDINGS_LOCAL_TRUST_REMOTE_CODE, str(DEFAULT_EMBEDDINGS_LOCAL_TRUST_REMOTE_CODE)
|
||||
).lower()
|
||||
@@ -2871,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),
|
||||
@@ -2895,20 +2576,11 @@ class HindsightConfig:
|
||||
or os.getenv(ENV_LLM_VERTEXAI_SERVICE_ACCOUNT_KEY),
|
||||
# Reranker
|
||||
reranker_provider=os.getenv(ENV_RERANKER_PROVIDER, DEFAULT_RERANKER_PROVIDER),
|
||||
reranker_send_bank_as_header=os.getenv(
|
||||
ENV_RERANKER_SEND_BANK_AS_HEADER,
|
||||
str(DEFAULT_RERANKER_SEND_BANK_AS_HEADER),
|
||||
).lower()
|
||||
in ("true", "1"),
|
||||
reranker_local_model=os.getenv(ENV_RERANKER_LOCAL_MODEL, DEFAULT_RERANKER_LOCAL_MODEL),
|
||||
reranker_local_force_cpu=os.getenv(
|
||||
ENV_RERANKER_LOCAL_FORCE_CPU, str(DEFAULT_RERANKER_LOCAL_FORCE_CPU)
|
||||
).lower()
|
||||
in ("true", "1"),
|
||||
reranker_local_allow_mps=os.getenv(
|
||||
ENV_RERANKER_LOCAL_ALLOW_MPS, str(DEFAULT_RERANKER_LOCAL_ALLOW_MPS)
|
||||
).lower()
|
||||
in ("true", "1"),
|
||||
reranker_local_max_concurrent=int(
|
||||
os.getenv(ENV_RERANKER_LOCAL_MAX_CONCURRENT, str(DEFAULT_RERANKER_LOCAL_MAX_CONCURRENT))
|
||||
),
|
||||
@@ -2935,21 +2607,7 @@ 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))),
|
||||
graph_seed_min_similarity=float(
|
||||
os.getenv(ENV_GRAPH_SEED_MIN_SIMILARITY, str(DEFAULT_GRAPH_SEED_MIN_SIMILARITY))
|
||||
),
|
||||
temporal_semantic_min_similarity=float(
|
||||
os.getenv(ENV_TEMPORAL_SEMANTIC_MIN_SIMILARITY, str(DEFAULT_TEMPORAL_SEMANTIC_MIN_SIMILARITY))
|
||||
),
|
||||
semantic_link_min_similarity=float(
|
||||
os.getenv(ENV_SEMANTIC_LINK_MIN_SIMILARITY, str(DEFAULT_SEMANTIC_LINK_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))
|
||||
),
|
||||
@@ -3131,9 +2789,6 @@ class HindsightConfig:
|
||||
ENV_FILE_PARSER_MARKITDOWN_OCR_PROMPT,
|
||||
DEFAULT_FILE_PARSER_MARKITDOWN_OCR_PROMPT,
|
||||
),
|
||||
file_parser_markitdown_ocr_default_headers=json.loads(
|
||||
os.getenv(ENV_FILE_PARSER_MARKITDOWN_OCR_DEFAULT_HEADERS, "null")
|
||||
),
|
||||
file_parser_iris_token=os.getenv(ENV_FILE_PARSER_IRIS_TOKEN) or None,
|
||||
file_parser_iris_org_id=os.getenv(ENV_FILE_PARSER_IRIS_ORG_ID) or None,
|
||||
file_parser_llama_parse_api_key=os.getenv(ENV_FILE_PARSER_LLAMA_PARSE_API_KEY) or None,
|
||||
@@ -3247,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",
|
||||
@@ -3265,28 +2916,17 @@ class HindsightConfig:
|
||||
),
|
||||
worker_http_port=int(os.getenv(ENV_WORKER_HTTP_PORT, str(DEFAULT_WORKER_HTTP_PORT))),
|
||||
worker_max_slots=int(os.getenv(ENV_WORKER_MAX_SLOTS, str(DEFAULT_WORKER_MAX_SLOTS))),
|
||||
worker_slot_reservations=worker_slot_reservations,
|
||||
worker_slot_reservations={
|
||||
op_type: int(os.getenv(env_var, str(default)))
|
||||
for op_type, (env_var, default) in WORKER_SLOT_RESERVATION_TYPES.items()
|
||||
if int(os.getenv(env_var, str(default))) > 0
|
||||
},
|
||||
worker_consolidation_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))),
|
||||
retain_wall_timeout=int(os.getenv(ENV_RETAIN_WALL_TIMEOUT, str(DEFAULT_RETAIN_WALL_TIMEOUT))),
|
||||
# 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))
|
||||
),
|
||||
@@ -3341,18 +2981,6 @@ class HindsightConfig:
|
||||
in ("true", "1", "yes"),
|
||||
metrics_backlog_enabled=os.getenv(ENV_METRICS_BACKLOG_ENABLED, str(DEFAULT_METRICS_BACKLOG_ENABLED)).lower()
|
||||
in ("true", "1", "yes"),
|
||||
# Runtime-stall observability (static, server-level only)
|
||||
loop_watchdog_enabled=os.getenv(ENV_LOOP_WATCHDOG_ENABLED, str(DEFAULT_LOOP_WATCHDOG_ENABLED)).lower()
|
||||
in ("true", "1", "yes"),
|
||||
loop_watchdog_stall_threshold_ms=int(
|
||||
os.getenv(ENV_LOOP_WATCHDOG_STALL_THRESHOLD_MS, str(DEFAULT_LOOP_WATCHDOG_STALL_THRESHOLD_MS))
|
||||
),
|
||||
loop_watchdog_poll_interval_ms=int(
|
||||
os.getenv(ENV_LOOP_WATCHDOG_POLL_INTERVAL_MS, str(DEFAULT_LOOP_WATCHDOG_POLL_INTERVAL_MS))
|
||||
),
|
||||
db_acquire_warn_threshold_ms=int(
|
||||
os.getenv(ENV_DB_ACQUIRE_WARN_THRESHOLD_MS, str(DEFAULT_DB_ACQUIRE_WARN_THRESHOLD_MS))
|
||||
),
|
||||
# Audit log configuration (static, server-level only)
|
||||
audit_log_enabled=os.getenv(ENV_AUDIT_LOG_ENABLED, str(DEFAULT_AUDIT_LOG_ENABLED)).lower() == "true",
|
||||
audit_log_actions=[
|
||||
@@ -3361,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=[
|
||||
|
||||
@@ -331,17 +331,11 @@ class ConfigResolver:
|
||||
logger.error(f"Failed to bulk-load bank configs: {e}")
|
||||
return result
|
||||
|
||||
async def validate_bank_config_updates(
|
||||
self,
|
||||
bank_id: str,
|
||||
updates: dict[str, Any],
|
||||
context: RequestContext | None = None,
|
||||
*,
|
||||
projected_bank_overrides: dict[str, Any] | None = None,
|
||||
check_permissions: bool = True,
|
||||
) -> dict[str, Any]:
|
||||
async def update_bank_config(
|
||||
self, bank_id: str, updates: dict[str, Any], context: RequestContext | None = None
|
||||
) -> None:
|
||||
"""
|
||||
Normalize and validate bank configuration overrides.
|
||||
Update bank configuration overrides (with permission checking).
|
||||
|
||||
Args:
|
||||
bank_id: Bank identifier
|
||||
@@ -350,16 +344,9 @@ class ConfigResolver:
|
||||
or Python field format (llm_provider).
|
||||
Only configurable fields are allowed.
|
||||
context: Request context for permission checking
|
||||
projected_bank_overrides: Bank overrides to use as the validation
|
||||
base instead of loading the current bank row.
|
||||
check_permissions: Whether client field permissions apply to these
|
||||
updates. Server-owned projected values set this to false.
|
||||
|
||||
Returns:
|
||||
Normalized updates ready to persist.
|
||||
|
||||
Raises:
|
||||
ValueError: If attempting to override invalid/disallowed fields.
|
||||
ValueError: If attempting to override invalid/disallowed fields
|
||||
"""
|
||||
# Normalize keys
|
||||
normalized_updates = normalize_config_dict(updates)
|
||||
@@ -391,7 +378,7 @@ class ConfigResolver:
|
||||
)
|
||||
|
||||
# PERMISSIONS: Check tenant/bank permissions
|
||||
if check_permissions and self.tenant_extension and context:
|
||||
if self.tenant_extension and context:
|
||||
try:
|
||||
allowed_fields = await self.tenant_extension.get_allowed_config_fields(context, bank_id)
|
||||
if allowed_fields is not None: # None means "allow all"
|
||||
@@ -440,11 +427,7 @@ class ConfigResolver:
|
||||
)
|
||||
if chunking_fields_updated:
|
||||
config_dict = await self._resolve_parent_config_dict(bank_id, context)
|
||||
active_bank_overrides = (
|
||||
await self._load_bank_config(bank_id)
|
||||
if projected_bank_overrides is None
|
||||
else dict(projected_bank_overrides)
|
||||
)
|
||||
active_bank_overrides = await self._load_bank_config(bank_id)
|
||||
for key, value in normalized_updates.items():
|
||||
if key not in self._configurable_fields:
|
||||
continue
|
||||
@@ -460,26 +443,17 @@ class ConfigResolver:
|
||||
)
|
||||
_validate_retain_strategy_chunking(base_config, base_config.retain_strategies)
|
||||
|
||||
return normalized_updates
|
||||
# Persist the override. Banks are created lazily (on first retain), so a
|
||||
# PATCH that precedes any ingestion would otherwise UPDATE zero rows and
|
||||
# silently no-op while returning 200. Ensure the bank row exists first
|
||||
# (this also creates its per-bank vector indexes), then merge defensively:
|
||||
# COALESCE guards against a NULL config column (NULL || jsonb is NULL),
|
||||
# which would drop the override even when a row is updated.
|
||||
from .engine.retain.fact_storage import ensure_bank_exists
|
||||
|
||||
async def update_bank_config(
|
||||
self, bank_id: str, updates: dict[str, Any], context: RequestContext | None = None
|
||||
) -> None:
|
||||
"""Validate and persist bank configuration overrides for an existing bank.
|
||||
|
||||
Bank creation belongs to ``MemoryEngine``; this raises ``ValueError`` if
|
||||
the bank does not exist rather than silently discarding the overrides.
|
||||
"""
|
||||
normalized_updates = await self.validate_bank_config_updates(bank_id, updates, context)
|
||||
await self._persist_bank_config(bank_id, normalized_updates)
|
||||
|
||||
async def _persist_bank_config(self, bank_id: str, normalized_updates: dict[str, Any]) -> None:
|
||||
"""Persist already-validated overrides without changing bank lifecycle state."""
|
||||
# Bank lifecycle belongs to MemoryEngine. Callers must create the row
|
||||
# before reaching this persistence step. COALESCE guards against a NULL
|
||||
# config column (NULL || jsonb is NULL), which would drop the override.
|
||||
async with self._backend.acquire() as conn:
|
||||
result = await conn.execute(
|
||||
await ensure_bank_exists(conn, bank_id, ops=self._backend.ops)
|
||||
await conn.execute(
|
||||
f"""
|
||||
UPDATE {fq_table("banks")}
|
||||
SET config = COALESCE(config, '{{}}'::jsonb) || $1::jsonb,
|
||||
@@ -490,14 +464,6 @@ class ConfigResolver:
|
||||
bank_id,
|
||||
)
|
||||
|
||||
# A missing bank row matches zero rows, which would otherwise persist
|
||||
# nothing while reporting success. Fail loudly instead: reaching here
|
||||
# without the row means a caller skipped the engine's provisioning step.
|
||||
# (The Oracle wrapper reshapes rowcount into the same "UPDATE <n>" form.)
|
||||
updated = int(result.split()[-1]) if isinstance(result, str) and result.startswith("UPDATE") else 0
|
||||
if updated == 0:
|
||||
raise ValueError(f"Cannot update config for bank '{bank_id}': the bank does not exist")
|
||||
|
||||
logger.info(f"Updated bank config for {bank_id}: {list(normalized_updates.keys())}")
|
||||
|
||||
async def reset_bank_config(self, bank_id: str) -> None:
|
||||
|
||||
@@ -14,7 +14,10 @@ import subprocess
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import IO
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from typing import IO
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -39,28 +42,37 @@ class IdleTimeoutMiddleware:
|
||||
self.app = app
|
||||
self.idle_timeout = idle_timeout
|
||||
self.last_activity = time.time()
|
||||
self._checker_task = None
|
||||
|
||||
async def __call__(self, scope, receive, send):
|
||||
# Update activity timestamp on each request
|
||||
self.last_activity = time.time()
|
||||
await self.app(scope, receive, send)
|
||||
|
||||
def start_idle_checker(self):
|
||||
"""Start the background task that checks for idle timeout."""
|
||||
self._checker_task = asyncio.create_task(self._check_idle())
|
||||
|
||||
async def _check_idle(self):
|
||||
"""Exit the daemon after the configured period without requests."""
|
||||
"""Background task that exits the process after idle timeout."""
|
||||
# If idle_timeout is 0, don't auto-exit
|
||||
if self.idle_timeout <= 0:
|
||||
return
|
||||
|
||||
while True:
|
||||
await asyncio.sleep(30)
|
||||
await asyncio.sleep(30) # Check every 30 seconds
|
||||
idle_time = time.time() - self.last_activity
|
||||
if idle_time > self.idle_timeout:
|
||||
logger.info(f"Idle timeout reached ({self.idle_timeout}s), shutting down daemon")
|
||||
# Give a moment for any in-flight requests
|
||||
await asyncio.sleep(1)
|
||||
# Send SIGTERM to ourselves to trigger graceful shutdown
|
||||
import signal
|
||||
|
||||
os.kill(os.getpid(), signal.SIGTERM)
|
||||
|
||||
|
||||
def _detach_popen_kwargs(log_handle: IO[bytes]) -> dict:
|
||||
def _detach_popen_kwargs(log_handle: "IO[bytes]") -> dict:
|
||||
"""Cross-platform kwargs to spawn a subprocess detached from the caller.
|
||||
|
||||
On POSIX, ``start_new_session=True`` calls ``setsid(2)`` so the child
|
||||
@@ -157,3 +169,17 @@ def daemonize():
|
||||
subprocess.Popen(cmd, env=env, **_detach_popen_kwargs(log_handle))
|
||||
|
||||
sys.exit(0)
|
||||
|
||||
|
||||
def check_daemon_running(port: int = DEFAULT_DAEMON_PORT) -> bool:
|
||||
"""Check if a daemon is running and responsive on the given port."""
|
||||
import socket
|
||||
|
||||
try:
|
||||
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
||||
sock.settimeout(1)
|
||||
result = sock.connect_ex(("127.0.0.1", port))
|
||||
sock.close()
|
||||
return result == 0
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
@@ -3,7 +3,7 @@ Memory Engine - Core implementation of the memory system.
|
||||
|
||||
This package contains all the implementation details of the memory engine:
|
||||
- MemoryEngine: Main class for memory operations
|
||||
- Utility modules: embedding_utils, link_utils, bank_utils
|
||||
- Utility modules: embedding_utils, link_utils, think_utils, bank_utils
|
||||
- Supporting modules: embeddings, cross_encoder, entity_resolver, etc.
|
||||
"""
|
||||
|
||||
|
||||
@@ -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,8 +19,6 @@ from typing import Any
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from ..engine.db_utils import acquire_with_retry
|
||||
from ..models import RequestContext
|
||||
from .schema import fq_table_explicit
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -121,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))
|
||||
@@ -189,12 +150,8 @@ class AuditLogger:
|
||||
logger.debug("Audit log skipped: pool not available")
|
||||
return
|
||||
try:
|
||||
# fq_table_explicit qualifies per dialect: "schema".audit_log on
|
||||
# PostgreSQL, bare audit_log on Oracle (where the schema is set at the
|
||||
# session level). A raw f"{schema}.audit_log" produced public.audit_log
|
||||
# on Oracle, where "public" is a reserved word — every write failed
|
||||
# with ORA-00903 even though the table exists.
|
||||
table = fq_table_explicit("audit_log", self._schema_getter())
|
||||
schema = self._schema_getter()
|
||||
table = f"{schema}.audit_log"
|
||||
async with acquire_with_retry(pool, max_retries=1) as conn:
|
||||
await conn.execute(
|
||||
f"""
|
||||
@@ -225,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.
|
||||
|
||||
@@ -234,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,70 +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.
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
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)
|
||||
|
||||
DEFAULT_CAUSAL_LINK_WEIGHT = 1.0
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class CausalLinkDescriptor:
|
||||
"""One causal edge, parked on the curation archive while an endpoint is invalidated.
|
||||
|
||||
Invalidation moves a fact out of ``memory_units``, so the FK cascade deletes
|
||||
its ``memory_links`` rows — and nothing could recreate a causal edge, which
|
||||
is extraction output rather than derived data. The descriptor is what the
|
||||
archive row stores so revert can rematerialize the edge (#2864).
|
||||
"""
|
||||
|
||||
from_unit_id: str
|
||||
to_unit_id: str
|
||||
link_type: str
|
||||
weight: float = DEFAULT_CAUSAL_LINK_WEIGHT
|
||||
|
||||
def as_json_dict(self) -> dict[str, Any]:
|
||||
"""Serializable form written to ``invalidated_memory_units.causal_links``.
|
||||
|
||||
The key names double as the column list of the ``jsonb_to_recordset``
|
||||
read in ``snapshot_causal_links`` — keep them in sync.
|
||||
"""
|
||||
return {
|
||||
"from_unit_id": self.from_unit_id,
|
||||
"to_unit_id": self.to_unit_id,
|
||||
"link_type": self.link_type,
|
||||
"weight": self.weight,
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_json_dict(cls, raw: Any) -> "CausalLinkDescriptor | None":
|
||||
"""Parse one stored descriptor, or None when it isn't a usable causal edge.
|
||||
|
||||
The archive column is plain JSON with no schema enforcement (a restore
|
||||
from an older backup, or a hand-edited row, can put anything there), and
|
||||
``memory_links`` has a ``link_type`` CHECK constraint — so an unusable
|
||||
entry is skipped rather than allowed to abort the whole revert.
|
||||
"""
|
||||
if not isinstance(raw, dict):
|
||||
return None
|
||||
from_unit_id = raw.get("from_unit_id")
|
||||
to_unit_id = raw.get("to_unit_id")
|
||||
link_type = raw.get("link_type")
|
||||
if not from_unit_id or not to_unit_id or link_type not in CAUSAL_LINK_TYPES:
|
||||
return None
|
||||
return cls(
|
||||
from_unit_id=str(from_unit_id),
|
||||
to_unit_id=str(to_unit_id),
|
||||
link_type=str(link_type),
|
||||
weight=float(raw.get("weight") or DEFAULT_CAUSAL_LINK_WEIGHT),
|
||||
)
|
||||
@@ -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:
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -37,36 +37,19 @@ _PROCESSING_RULES = """## PROCESSING RULES
|
||||
|
||||
9. KEEP DISTINCT TOPICS DISTINCT: do not merge observations about different people, entities, or unrelated topics. Merging is for the same canonical fact recurring — not for related-but-distinct claims."""
|
||||
|
||||
# Field-by-field definitions of the input shape used by the cached system
|
||||
# prefix. The call site runs .format(), so these strings must contain no braces.
|
||||
_FACT_FIELDS = """One per line, formatted as `[uuid] fact text (temporal fields)`:
|
||||
- `[uuid]`: the fact's identifier — copy it verbatim into `source_fact_ids`
|
||||
- `occurred_start` / `occurred_end`: when the described event happened. This can be long before the fact was stated — a fact recorded today may describe a 2019 event.
|
||||
- `mentioned_at`: when the source material that states this fact was written. This is the fact's recency: how up to date the statement is, NOT when it was added to memory. A fact taken from an old document keeps its old `mentioned_at` even if it was only just processed."""
|
||||
|
||||
_OBSERVATION_FIELDS = """- `id`: unique identifier — copy this exactly when issuing an UPDATE or DELETE
|
||||
- `text`: the observation content
|
||||
- `proof_count`: how many source facts this observation has already merged
|
||||
- `occurred_start` / `occurred_end`: the span of the events behind the observation — earliest start and latest end across its source facts
|
||||
- `mentioned_at`: the latest of the `mentioned_at` values of its source facts — the most recent point at which this observation was stated
|
||||
- `source_memories`: the supporting facts behind this observation. May be partial or absent for large observations — the count above remains the true total. Each entry carries the same `text` and temporal fields as a new fact, plus:
|
||||
- `context`: optional surrounding context for that fact"""
|
||||
|
||||
# Stable description of the input shape. For the cached split path this lives in
|
||||
# the system prefix (build_consolidation_system_prompt) so it is not re-sent on
|
||||
# every batch; the per-batch user message then carries only the actual data.
|
||||
_INPUT_FORMAT_NOTE = f"""## INPUT FORMAT
|
||||
_INPUT_FORMAT_NOTE = """## INPUT FORMAT
|
||||
|
||||
Each request provides new facts and existing observations. Every temporal field is optional and is omitted when unknown.
|
||||
|
||||
### New facts
|
||||
|
||||
{_FACT_FIELDS}
|
||||
|
||||
### Existing observations
|
||||
|
||||
A JSON array pooled from recalls across the new facts. Each entry has:
|
||||
{_OBSERVATION_FIELDS}"""
|
||||
Each request provides new facts and existing observations:
|
||||
- New facts: one per line, each prefixed with its `[uuid]`, followed by the fact text and optional temporal fields.
|
||||
- Existing observations: a JSON array pooled from recalls across the new facts. Each entry has:
|
||||
- `id`: unique identifier — copy this exactly when issuing an UPDATE or DELETE
|
||||
- `text`: the observation content
|
||||
- `proof_count`: number of supporting memories
|
||||
- `occurred_start` / `occurred_end`: temporal range of source facts
|
||||
- `source_memories`: array of supporting facts with their text and dates"""
|
||||
|
||||
# Per-batch data section for the cached split path — the stable format
|
||||
# explanation above is omitted here (it lives in the cached prefix); only the
|
||||
@@ -81,6 +64,24 @@ _SPLIT_INPUT_SECTION = """## INPUT
|
||||
|
||||
{observations_text}"""
|
||||
|
||||
# Data section — format placeholders {facts_text} and {observations_text} are substituted at call time
|
||||
_INPUT_SECTION = """## INPUT
|
||||
|
||||
### New facts
|
||||
|
||||
{facts_text}
|
||||
|
||||
### Existing observations
|
||||
|
||||
JSON array, pooled from recalls across all new facts above. Each entry has:
|
||||
- `id`: unique identifier — copy this exactly when issuing an UPDATE or DELETE
|
||||
- `text`: the observation content
|
||||
- `proof_count`: number of supporting memories
|
||||
- `occurred_start` / `occurred_end`: temporal range of source facts
|
||||
- `source_memories`: array of supporting facts with their text and dates
|
||||
|
||||
{observations_text}"""
|
||||
|
||||
_DECISION_GUIDE = """## DECISION GUIDE
|
||||
|
||||
- **Same canonical event, decision, claim, or facet as an existing observation → UPDATE** (use `observation_id` + new `source_fact_ids`).
|
||||
@@ -141,6 +142,39 @@ Expected output (UPDATE for the state change; CREATE for the unrelated work-hour
|
||||
- Return `{{"creates": [], "updates": [], "deletes": []}}` if nothing durable is found."""
|
||||
|
||||
|
||||
def build_batch_consolidation_prompt(
|
||||
observations_mission: str | None = None,
|
||||
observation_capacity_note: str | None = None,
|
||||
llm_output_language: str | None = None,
|
||||
) -> str:
|
||||
"""
|
||||
Build the consolidation prompt for batch mode (multiple facts per LLM call).
|
||||
|
||||
The mission defines *what* to track (customisable per bank) and takes
|
||||
priority over the built-in processing rules when the two conflict.
|
||||
Processing rules, decision guide, and output format are always present.
|
||||
When ``llm_output_language`` is set, observations are emitted in that
|
||||
language.
|
||||
"""
|
||||
mission = escape_for_prompt(observations_mission or _DEFAULT_MISSION)
|
||||
|
||||
capacity_section = ""
|
||||
if observation_capacity_note:
|
||||
capacity_section = f"\n\n## CAPACITY CONSTRAINT\n\n{escape_for_prompt(observation_capacity_note)}"
|
||||
|
||||
return (
|
||||
"You are a memory consolidation system. Synthesize new facts into "
|
||||
"observations, merging with existing observations when appropriate.\n\n"
|
||||
f"## MISSION\n\n{mission}\n\n"
|
||||
f"{_MISSION_PRIORITY_NOTE}"
|
||||
f"{capacity_section}\n\n"
|
||||
f"{_PROCESSING_RULES}\n\n"
|
||||
f"{_INPUT_SECTION}\n\n"
|
||||
f"{_DECISION_GUIDE}\n\n"
|
||||
f"{_OUTPUT_SECTION}" + output_language_directive(llm_output_language)
|
||||
)
|
||||
|
||||
|
||||
def build_consolidation_system_prompt(
|
||||
llm_output_language: str | None = None,
|
||||
) -> str:
|
||||
|
||||
@@ -12,7 +12,6 @@ import os
|
||||
import warnings
|
||||
from abc import ABC, abstractmethod
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
@@ -46,17 +45,47 @@ from ..config import (
|
||||
ENV_RERANKER_TEI_URL,
|
||||
ENV_RERANKER_ZEROENTROPY_API_KEY,
|
||||
)
|
||||
from .bank_attribution import reranker_bank_attribution_headers
|
||||
from .local_device import (
|
||||
release_local_inference_memory,
|
||||
resolve_model_device_type,
|
||||
select_local_device,
|
||||
)
|
||||
from .tei_retry import tei_retry_delay
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _resolve_malloc_trim():
|
||||
"""Return a callable that asks glibc to release freed heap pages to the OS.
|
||||
|
||||
Local CPU rerankers (FlashRank/ONNX, SentenceTransformers/torch) allocate
|
||||
large transient numpy/tensor buffers per call. On Linux glibc, those pages
|
||||
are freed at the Python level but kept by the allocator as a high-water
|
||||
mark — RSS grows monotonically across many recalls (see issue #1717).
|
||||
Calling `malloc_trim(0)` after each batch returns those pages to the OS.
|
||||
|
||||
Resolved once at import; returns a no-op on non-glibc platforms (macOS,
|
||||
musl, Windows) where the call is unavailable or unnecessary.
|
||||
"""
|
||||
import sys
|
||||
|
||||
if sys.platform != "linux":
|
||||
return lambda: None
|
||||
|
||||
import ctypes
|
||||
import ctypes.util
|
||||
|
||||
libc_path = ctypes.util.find_library("c")
|
||||
if libc_path is None:
|
||||
return lambda: None
|
||||
try:
|
||||
libc = ctypes.CDLL(libc_path)
|
||||
trim = libc.malloc_trim
|
||||
except (OSError, AttributeError):
|
||||
# Not glibc (musl has no malloc_trim) or libc lookup failed.
|
||||
return lambda: None
|
||||
trim.argtypes = [ctypes.c_size_t]
|
||||
trim.restype = ctypes.c_int
|
||||
return lambda: trim(0)
|
||||
|
||||
|
||||
_malloc_trim = _resolve_malloc_trim()
|
||||
|
||||
|
||||
class CrossEncoderModel(ABC):
|
||||
"""
|
||||
Abstract base class for cross-encoder reranking.
|
||||
@@ -121,7 +150,6 @@ class LocalSTCrossEncoder(CrossEncoderModel):
|
||||
fp16: bool = False,
|
||||
bucket_batching: bool = False,
|
||||
batch_size: int = DEFAULT_RERANKER_LOCAL_BATCH_SIZE,
|
||||
allow_mps: bool = False,
|
||||
):
|
||||
"""
|
||||
Initialize local SentenceTransformers cross-encoder.
|
||||
@@ -143,9 +171,6 @@ class LocalSTCrossEncoder(CrossEncoderModel):
|
||||
Default: False (opt-in via env var).
|
||||
batch_size: Batch size for predict() calls. Optimal values vary by
|
||||
hardware and model (MPS: 32, CUDA: 128+). Default: 32.
|
||||
allow_mps: Opt in to the Apple Silicon MPS GPU. Disabled by default
|
||||
because MPS leaks memory under variable-length workloads
|
||||
(see engine/local_device.py). Default: False
|
||||
"""
|
||||
self.model_name = model_name or DEFAULT_RERANKER_LOCAL_MODEL
|
||||
self.force_cpu = force_cpu
|
||||
@@ -153,9 +178,7 @@ class LocalSTCrossEncoder(CrossEncoderModel):
|
||||
self.fp16 = fp16
|
||||
self.bucket_batching = bucket_batching
|
||||
self.batch_size = batch_size
|
||||
self.allow_mps = allow_mps
|
||||
self._model = None
|
||||
self._device_type: str = "cpu"
|
||||
LocalSTCrossEncoder._max_concurrent = max_concurrent
|
||||
|
||||
@property
|
||||
@@ -177,13 +200,33 @@ class LocalSTCrossEncoder(CrossEncoderModel):
|
||||
|
||||
logger.info(f"Reranker: initializing local provider with model {self.model_name}")
|
||||
|
||||
# Determine device based on hardware availability. We always set
|
||||
# low_cpu_mem_usage=False to prevent lazy loading (meta tensors) which can
|
||||
# cause issues when accelerate is installed but no GPU is available.
|
||||
# Determine device based on hardware availability.
|
||||
# We always set low_cpu_mem_usage=False to prevent lazy loading (meta tensors)
|
||||
# which can cause issues when accelerate is installed but no GPU is available.
|
||||
# Note: We do NOT use device_map because CrossEncoder internally calls .to(device)
|
||||
# after loading, which conflicts with accelerate's device_map handling.
|
||||
# MPS is opt-in (allow_mps) — see engine/local_device.py for why.
|
||||
device = select_local_device(self.force_cpu, self.allow_mps)
|
||||
import torch
|
||||
|
||||
# Force CPU mode if configured (used in daemon mode to avoid MPS/XPC issues on macOS)
|
||||
if self.force_cpu:
|
||||
device = "cpu"
|
||||
logger.info("Reranker: forcing CPU mode (HINDSIGHT_API_RERANKER_LOCAL_FORCE_CPU=1)")
|
||||
else:
|
||||
# Check for GPU (CUDA), Apple Silicon (MPS), or Intel XPU
|
||||
# Wrap in try-except to gracefully handle any device detection issues
|
||||
# (e.g., in CI environments or when PyTorch is built without GPU support)
|
||||
device = "cpu" # Default to CPU
|
||||
try:
|
||||
has_gpu = torch.cuda.is_available() or (
|
||||
hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
|
||||
)
|
||||
# Intel Arc XPU support — torch.xpu is available when the XPU build is loaded
|
||||
if not has_gpu and hasattr(torch, "xpu"):
|
||||
has_gpu = torch.xpu.is_available()
|
||||
if has_gpu:
|
||||
device = None # Let sentence-transformers auto-detect GPU/MPS/XPU
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to detect GPU/MPS/XPU, falling back to CPU: {e}")
|
||||
|
||||
# Patch transformers 5.x compatibility for models using XLM-RoBERTa
|
||||
# (e.g., jina-reranker-v2-base-multilingual). transformers 5.x removed
|
||||
@@ -227,11 +270,9 @@ class LocalSTCrossEncoder(CrossEncoderModel):
|
||||
# Restore original logging level
|
||||
transformers_logger.setLevel(original_level)
|
||||
|
||||
self._device_type = resolve_model_device_type(self._model)
|
||||
|
||||
# FP16 inference: convert model weights to half precision.
|
||||
# Empirically validated: 27-36% faster on MPS, quality-identical (20/20 overlap).
|
||||
if self.fp16 and self._device_type != "cpu":
|
||||
if self.fp16 and device != "cpu":
|
||||
self._model.model.half()
|
||||
logger.info("Reranker: FP16 inference enabled")
|
||||
|
||||
@@ -274,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_local_inference_memory(self._device_type)
|
||||
_malloc_trim()
|
||||
|
||||
async def predict(self, pairs: list[tuple[str, str]]) -> list[float]:
|
||||
"""
|
||||
@@ -390,20 +431,14 @@ class RemoteTEICrossEncoder(CrossEncoderModel):
|
||||
await asyncio.sleep(delay)
|
||||
delay *= 2 # Exponential backoff
|
||||
except httpx.HTTPStatusError as e:
|
||||
# TEI uses 429 as normal overload backpressure. Retry it with
|
||||
# the same bounded budget as transient server errors.
|
||||
if (e.response.status_code == 429 or e.response.status_code >= 500) and attempt < self.max_retries:
|
||||
# Retry on 5xx server errors
|
||||
if e.response.status_code >= 500 and attempt < self.max_retries:
|
||||
last_error = e
|
||||
sleep_delay = tei_retry_delay(
|
||||
e.response,
|
||||
delay,
|
||||
request_timeout=self.timeout,
|
||||
)
|
||||
logger.warning(
|
||||
f"TEI transient error (attempt {attempt + 1}/{self.max_retries + 1}): {e}. "
|
||||
f"Retrying in {sleep_delay:.2f}s..."
|
||||
f"TEI server error (attempt {attempt + 1}/{self.max_retries + 1}): {e}. "
|
||||
f"Retrying in {delay}s..."
|
||||
)
|
||||
await asyncio.sleep(sleep_delay)
|
||||
await asyncio.sleep(delay)
|
||||
delay *= 2
|
||||
else:
|
||||
raise
|
||||
@@ -449,7 +484,6 @@ class RemoteTEICrossEncoder(CrossEncoderModel):
|
||||
semaphore,
|
||||
"POST",
|
||||
f"{self.base_url}/rerank",
|
||||
headers=reranker_bank_attribution_headers(),
|
||||
json={
|
||||
"query": query,
|
||||
"texts": texts,
|
||||
@@ -590,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()
|
||||
|
||||
@@ -889,7 +919,6 @@ class FlashRankCrossEncoder(CrossEncoderModel):
|
||||
self.max_length = max_length
|
||||
self.cpu_mem_arena = cpu_mem_arena
|
||||
self._ranker = None
|
||||
self._device_type: str = "cpu" # FlashRank runs on CPU via ONNX Runtime
|
||||
FlashRankCrossEncoder._max_concurrent = max_concurrent
|
||||
|
||||
@property
|
||||
@@ -961,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]]] = {}
|
||||
@@ -994,7 +1023,7 @@ class FlashRankCrossEncoder(CrossEncoderModel):
|
||||
|
||||
return all_scores
|
||||
finally:
|
||||
release_local_inference_memory(self._device_type)
|
||||
_malloc_trim()
|
||||
|
||||
async def predict(self, pairs: list[tuple[str, str]]) -> list[float]:
|
||||
"""
|
||||
@@ -1122,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,
|
||||
@@ -1241,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
|
||||
@@ -1254,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
|
||||
|
||||
@@ -1608,7 +1647,6 @@ def create_cross_encoder_from_env() -> CrossEncoderModel:
|
||||
fp16=config.reranker_local_fp16,
|
||||
bucket_batching=config.reranker_local_bucket_batching,
|
||||
batch_size=config.reranker_local_batch_size,
|
||||
allow_mps=config.reranker_local_allow_mps,
|
||||
)
|
||||
elif provider == "cohere":
|
||||
api_key = config.reranker_cohere_api_key
|
||||
|
||||
@@ -67,28 +67,16 @@ def create_database_backend(backend_type: str) -> DatabaseBackend:
|
||||
return _get_backend_class(backend_type)()
|
||||
|
||||
|
||||
_OPS_CACHE: dict[str, DataAccessOps] = {}
|
||||
|
||||
|
||||
def create_data_access_ops(backend_type: str) -> DataAccessOps:
|
||||
"""Factory: the DataAccessOps for a backend name.
|
||||
|
||||
Returns a per-dialect SINGLETON: ``DataAccessOps`` is stateless (it only builds and runs SQL),
|
||||
so one shared instance per dialect is correct — and it means the database backend and the
|
||||
memories store hold the *same* ops object, so a test that patches a method on it (e.g.
|
||||
``enqueue_graph_maintenance``) observes every caller regardless of which layer issued it.
|
||||
"""Factory: create a DataAccessOps by backend name.
|
||||
|
||||
Args:
|
||||
backend_type: One of "postgresql" or "oracle".
|
||||
|
||||
Returns:
|
||||
The shared DataAccessOps instance for that backend.
|
||||
A DataAccessOps instance.
|
||||
|
||||
Raises:
|
||||
ValueError: If backend_type is not recognized.
|
||||
"""
|
||||
ops = _OPS_CACHE.get(backend_type)
|
||||
if ops is None:
|
||||
ops = _get_ops_class(backend_type)()
|
||||
_OPS_CACHE[backend_type] = ops
|
||||
return ops
|
||||
return _get_ops_class(backend_type)()
|
||||
|
||||
@@ -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,
|
||||
@@ -282,29 +253,22 @@ class OracleOps(DataAccessOps):
|
||||
) -> None:
|
||||
if not unit_ids:
|
||||
return
|
||||
# Locking upsert (#3034), the Oracle analogue of the PG
|
||||
# ``ON CONFLICT DO UPDATE``. The old IGNORE_ROW_ON_DUPKEY_INDEX insert
|
||||
# skipped duplicates WITHOUT locking the existing row, so a mutation
|
||||
# re-enqueueing an already-queued unit could not block a worker from
|
||||
# concurrently claiming (deleting) that row and processing the unit's
|
||||
# pre-mutation state — the re-enqueue signal was silently lost. MERGE
|
||||
# WHEN MATCHED takes an exclusive row lock on the existing queue row
|
||||
# (the SET is a deliberate no-op that preserves enqueued_at); WHEN NOT
|
||||
# MATCHED inserts a fresh row. That serialises the mutation against the
|
||||
# worker's claim for the same (bank_id, unit_id).
|
||||
# Oracle doesn't support ON CONFLICT; rely on the PK and the
|
||||
# IGNORE_ROW_ON_DUPKEY_INDEX hint to skip duplicates server-side.
|
||||
# The hint name must match the PK constraint exactly.
|
||||
#
|
||||
# Sort to enforce a global (bank_id, unit_id) lock-acquisition order,
|
||||
# matching claim_graph_maintenance_batch's delete order, so overlapping
|
||||
# mutation/worker sets acquire the shared row locks ascending and cannot
|
||||
# cycle.
|
||||
# Sort to enforce a global lock-acquisition order on the
|
||||
# (bank_id, unit_id) PK. Without this, two concurrent
|
||||
# transactions inserting overlapping unit_id sets in different
|
||||
# orders can deadlock on the unique-check row locks. Sorting
|
||||
# gives every concurrent caller the same lock order, so
|
||||
# conflicting inserts queue cleanly instead of cycling.
|
||||
sorted_unit_ids = sorted(unit_ids)
|
||||
await conn.executemany(
|
||||
f"""
|
||||
MERGE INTO {table} q
|
||||
USING (SELECT $1 AS bank_id, $2 AS unit_id FROM dual) s
|
||||
ON (q.bank_id = s.bank_id AND q.unit_id = s.unit_id)
|
||||
WHEN MATCHED THEN UPDATE SET q.enqueued_at = q.enqueued_at
|
||||
WHEN NOT MATCHED THEN INSERT (bank_id, unit_id) VALUES (s.bank_id, s.unit_id)
|
||||
INSERT /*+ IGNORE_ROW_ON_DUPKEY_INDEX({table}, pk_graph_maintenance_queue) */
|
||||
INTO {table} (bank_id, unit_id)
|
||||
VALUES ($1, $2)
|
||||
""",
|
||||
[(bank_id, uid) for uid in sorted_unit_ids],
|
||||
)
|
||||
@@ -329,15 +293,7 @@ class OracleOps(DataAccessOps):
|
||||
bank_id,
|
||||
limit,
|
||||
)
|
||||
# Ordered locking (#3034): the per-row DELETE takes the queue rows'
|
||||
# exclusive locks in executemany array order. Sort the claimed keys by
|
||||
# unit_id so those locks are acquired in the same (bank_id, unit_id)
|
||||
# order the enqueue MERGE uses — overlapping mutation/worker sets then
|
||||
# lock the shared rows ascending and cannot cycle. (The batch is still
|
||||
# *chosen* oldest-first by enqueued_at above; only the lock/delete order
|
||||
# is normalised.) The Pass 1 retry wrap in run_graph_maintenance_job is
|
||||
# the ORA-00060 backstop for any residual interleaving.
|
||||
claimed = sorted(str(row["unit_id"]) for row in rows)
|
||||
claimed = [str(row["unit_id"]) for row in rows]
|
||||
if claimed:
|
||||
await conn.executemany(
|
||||
f"DELETE FROM {table} WHERE bank_id = $1 AND unit_id = $2",
|
||||
@@ -373,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}
|
||||
@@ -497,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
|
||||
@@ -517,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
|
||||
)"""
|
||||
@@ -881,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,
|
||||
@@ -253,20 +286,11 @@ class PostgreSQLOps(DataAccessOps):
|
||||
entity_names: list[str],
|
||||
entity_dates: list,
|
||||
) -> dict[str, str]:
|
||||
# ORDER BY LOWER(name) so every concurrent batch inserts in the same order
|
||||
# as the conflict target (bank_id, LOWER(canonical_name)). ON CONFLICT DO
|
||||
# NOTHING takes a ShareLock on the inserting transaction of any speculative
|
||||
# row it collides with, so two batches with overlapping names inserting in
|
||||
# different orders deadlock. The caller already sorts by Python's
|
||||
# ``str.lower()``, which agrees with the index for ASCII but not for every
|
||||
# locale (see the Turkish-İ note in entity_resolver) — ordering in SQL makes
|
||||
# the database's own collation the single arbiter for all writers.
|
||||
inserted_rows = await conn.fetch(
|
||||
f"""
|
||||
INSERT INTO {table} (bank_id, canonical_name, first_seen, last_seen, mention_count)
|
||||
SELECT $1, name, COALESCE(event_date, now()), COALESCE(event_date, now()), 0
|
||||
FROM unnest($2::text[], $3::timestamptz[]) AS t(name, event_date)
|
||||
ORDER BY LOWER(name)
|
||||
ON CONFLICT (bank_id, LOWER(canonical_name))
|
||||
DO NOTHING
|
||||
RETURNING id, LOWER(canonical_name) AS name_lower
|
||||
@@ -286,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
|
||||
@@ -298,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,
|
||||
@@ -369,24 +357,11 @@ class PostgreSQLOps(DataAccessOps):
|
||||
# concurrent caller the same lock order, so conflicting inserts
|
||||
# queue cleanly instead of cycling.
|
||||
sorted_unit_ids = sorted(unit_ids)
|
||||
# DO UPDATE (not DO NOTHING) on a duplicate enqueue — #3034. The SET is a
|
||||
# deliberate no-op that preserves enqueued_at; its only purpose is to take
|
||||
# the existing row's lock. DO NOTHING does NOT lock the conflicting row, so
|
||||
# a mutation that re-enqueues an already-queued unit could not block a
|
||||
# worker from concurrently claiming (deleting) that row and processing the
|
||||
# unit's pre-mutation state; the re-enqueue signal was then silently lost
|
||||
# and the unit's derived links stayed stale with an empty queue. Locking
|
||||
# the row serialises the mutation against the worker's claim for that
|
||||
# (bank_id, unit_id): the worker either waits for the committed post-mutation
|
||||
# state, or (if it claimed first) this INSERT lands a fresh row after the
|
||||
# worker's delete commits. Row locks are acquired in sorted unit_id order,
|
||||
# matching claim_graph_maintenance_batch, so the two never cycle.
|
||||
await conn.execute(
|
||||
f"""
|
||||
INSERT INTO {table} (bank_id, unit_id)
|
||||
SELECT $1, v FROM unnest($2::uuid[]) AS t(v)
|
||||
ON CONFLICT (bank_id, unit_id)
|
||||
DO UPDATE SET enqueued_at = {table}.enqueued_at
|
||||
ON CONFLICT (bank_id, unit_id) DO NOTHING
|
||||
""",
|
||||
bank_id,
|
||||
sorted_unit_ids,
|
||||
@@ -399,35 +374,16 @@ class PostgreSQLOps(DataAccessOps):
|
||||
bank_id: str,
|
||||
limit: int,
|
||||
) -> list[str]:
|
||||
# Ordered locking (#3034). Choose the oldest batch by enqueued_at, but
|
||||
# acquire the row locks in (bank_id, unit_id) order — the same order the
|
||||
# enqueue upsert takes them — so a foreground mutation re-enqueueing an
|
||||
# overlapping unit set can never cycle against a worker draining it. The
|
||||
# `chosen` CTE is MATERIALIZED so the enqueued_at pick is fenced from the
|
||||
# locking clause; `FOR UPDATE OF q ... ORDER BY q.unit_id` then puts
|
||||
# LockRows above the Sort, so locks are taken ascending by unit_id (same
|
||||
# idiom as prune_stale_cooccurrences' #2529 ordered lock). A concurrent
|
||||
# enqueue holding one of these rows blocks this claim until it commits, at
|
||||
# which point the worker deletes and processes the committed state.
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
WITH chosen AS MATERIALIZED (
|
||||
DELETE FROM {table}
|
||||
WHERE (bank_id, unit_id) IN (
|
||||
SELECT bank_id, unit_id FROM {table}
|
||||
WHERE bank_id = $1
|
||||
ORDER BY enqueued_at
|
||||
LIMIT $2
|
||||
),
|
||||
locked AS (
|
||||
SELECT q.bank_id, q.unit_id
|
||||
FROM {table} q
|
||||
JOIN chosen c ON c.bank_id = q.bank_id AND c.unit_id = q.unit_id
|
||||
ORDER BY q.unit_id
|
||||
FOR UPDATE OF q
|
||||
)
|
||||
DELETE FROM {table} q
|
||||
USING locked l
|
||||
WHERE q.bank_id = l.bank_id AND q.unit_id = l.unit_id
|
||||
RETURNING q.unit_id
|
||||
RETURNING unit_id
|
||||
""",
|
||||
bank_id,
|
||||
limit,
|
||||
@@ -469,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,
|
||||
)
|
||||
@@ -522,21 +449,11 @@ class PostgreSQLOps(DataAccessOps):
|
||||
mu_table: str,
|
||||
unit_ids: list[str],
|
||||
) -> list[ResultRow]:
|
||||
# Cast only canonical UUID text inputs, never the indexed column. The old
|
||||
# ``id::text`` predicate silently ignored malformed, uppercase, braced,
|
||||
# and unhyphenated inputs; filtering before the cast preserves that
|
||||
# behavior while allowing the primary-key index to serve the lookup.
|
||||
return await conn.fetch(
|
||||
f"""
|
||||
SELECT id, event_date, fact_type
|
||||
FROM {mu_table}
|
||||
WHERE id = ANY(
|
||||
ARRAY(
|
||||
SELECT input.unit_id::uuid
|
||||
FROM unnest($1::text[]) AS input(unit_id)
|
||||
WHERE input.unit_id ~ '^[0-9a-f]{{8}}-[0-9a-f]{{4}}-[0-9a-f]{{4}}-[0-9a-f]{{4}}-[0-9a-f]{{12}}$'
|
||||
)
|
||||
)
|
||||
WHERE id::text = ANY($1)
|
||||
""",
|
||||
unit_ids,
|
||||
)
|
||||
@@ -624,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
|
||||
@@ -851,16 +761,10 @@ class PostgreSQLOps(DataAccessOps):
|
||||
internal_id: str,
|
||||
fact_types: dict[str, str],
|
||||
) -> None:
|
||||
# CONCURRENTLY so the drop takes ShareUpdateExclusive, not ACCESS
|
||||
# EXCLUSIVE, on the shared memory_units table. A plain DROP INDEX blocks
|
||||
# (and deadlocks with) every other bank's concurrent reads/writes on the
|
||||
# table; CONCURRENTLY does not conflict with DML. The caller
|
||||
# (delete_bank) runs this on an autocommit connection after its delete
|
||||
# transaction has committed — CONCURRENTLY cannot run inside a tx.
|
||||
for ft, suffix in fact_types.items():
|
||||
uid = str(internal_id).replace("-", "")[:16]
|
||||
idx = f"idx_mu_emb_{suffix}_{uid}"
|
||||
await conn.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}.{idx}")
|
||||
await conn.execute(f"DROP INDEX IF EXISTS {schema}.{idx}")
|
||||
|
||||
def get_entity_resolution_strategy(self) -> str:
|
||||
return "trigram"
|
||||
@@ -990,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,
|
||||
|
||||
@@ -23,8 +23,6 @@ from collections.abc import AsyncIterator
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import Any, NamedTuple
|
||||
|
||||
from .pool_instrumentation import PoolStats, acquire_conn
|
||||
|
||||
|
||||
class _OracleJSONEncoder(json.JSONEncoder):
|
||||
"""JSON encoder that handles datetime and UUID objects."""
|
||||
@@ -80,9 +78,7 @@ _LIKE_ANY_RE = re.compile(r"(\w+)\s+LIKE\s+ANY\s*\(\s*:(\d+)\s*\)", re.IGNORECAS
|
||||
_NOT_LIKE_ALL_RE = re.compile(r"(\w+)\s+NOT\s+LIKE\s+ALL\s*\(\s*:(\d+)\s*\)", re.IGNORECASE)
|
||||
|
||||
_JSON_ARROW_TEXT_RE = re.compile(r'("?\w+"?)\s*->>\s*\'(\w+)\'') # handles both col and "col"
|
||||
# Reserved-word columns ("trigger") are already quoted by the time this runs, so the
|
||||
# column group must accept the quoted form too — same shape as the arrow regex above.
|
||||
_JSON_HAS_KEY_RE = re.compile(r"(\"?\w+\"?)\s*\?\s*'(\w+)'")
|
||||
_JSON_HAS_KEY_RE = re.compile(r"(\w+)\s*\?\s*'(\w+)'")
|
||||
_JSONB_CONTAINS_RE = re.compile(r"(\w+)\s*@>\s*:(\d+)")
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -150,7 +146,6 @@ _JSON_COL_NAMES = {
|
||||
"config",
|
||||
"observation_scopes",
|
||||
"source_memory_ids",
|
||||
"causal_links",
|
||||
"trigger",
|
||||
"http_config",
|
||||
"event_types",
|
||||
@@ -449,6 +444,9 @@ def _rewrite_pg_to_oracle(query: str) -> RewriteResult:
|
||||
if has_for_update:
|
||||
# FOR UPDATE path: use ROWNUM instead of FETCH FIRST.
|
||||
# Extract and remove LIMIT clause, inject ROWNUM into WHERE.
|
||||
def _limit_to_rownum(m):
|
||||
return "" # Remove the LIMIT clause; we'll add ROWNUM below
|
||||
|
||||
limit_val = None
|
||||
limit_match = re.search(r"\bLIMIT\s+(\d+|:\w+)\b", query, re.IGNORECASE)
|
||||
if limit_match:
|
||||
@@ -1244,11 +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
|
||||
self._acquire_warn_threshold_s: float = 1.0
|
||||
|
||||
async def initialize(
|
||||
self,
|
||||
@@ -1264,10 +1257,6 @@ class OracleBackend(DatabaseBackend):
|
||||
oracledb = _import_oracledb()
|
||||
self._oracledb = oracledb
|
||||
|
||||
from ...config import get_config
|
||||
|
||||
self._acquire_warn_threshold_s = get_config().db_acquire_warn_threshold_ms / 1000.0
|
||||
|
||||
# Parse URL-format DSN (oracle://user:pass@host:port/service)
|
||||
from urllib.parse import urlparse
|
||||
|
||||
@@ -1288,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.
|
||||
|
||||
@@ -1311,41 +1294,15 @@ 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()
|
||||
|
||||
def _pool_stats(self) -> PoolStats | None:
|
||||
"""Snapshot for slow-acquire logs, from oracledb pool attributes."""
|
||||
pool = self._pool
|
||||
if pool is None:
|
||||
return None
|
||||
try:
|
||||
busy = pool.busy
|
||||
return PoolStats(in_use=busy, max=pool.max, idle=pool.opened - busy)
|
||||
except Exception:
|
||||
return None
|
||||
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]:
|
||||
pool = self._ensure_pool()
|
||||
conn = await acquire_conn(
|
||||
pool.acquire(), pool_stats=self._pool_stats, warn_threshold_s=self._acquire_warn_threshold_s
|
||||
)
|
||||
conn = await pool.acquire()
|
||||
try:
|
||||
await self._set_session_schema(conn)
|
||||
yield OracleConnection(conn)
|
||||
@@ -1361,9 +1318,7 @@ class OracleBackend(DatabaseBackend):
|
||||
@asynccontextmanager
|
||||
async def transaction(self) -> AsyncIterator[OracleConnection]:
|
||||
pool = self._ensure_pool()
|
||||
conn = await acquire_conn(
|
||||
pool.acquire(), pool_stats=self._pool_stats, warn_threshold_s=self._acquire_warn_threshold_s
|
||||
)
|
||||
conn = await pool.acquire()
|
||||
try:
|
||||
await self._set_session_schema(conn)
|
||||
yield OracleConnection(conn)
|
||||
|
||||
@@ -1,137 +0,0 @@
|
||||
"""Instrumentation for database connection-pool acquisition.
|
||||
|
||||
asyncpg exposes pool *size* and *idle* counts, but not how many callers are
|
||||
currently **queued waiting** for a connection — and that queue depth is the
|
||||
signal that actually distinguishes a saturated pool from a healthy one. When the
|
||||
pool is exhausted, ``/health`` (which itself acquires a connection to run
|
||||
``SELECT 1``) blocks in ``pool.acquire()`` until a connection frees or the acquire
|
||||
times out, so a liveness probe can fail **with the event loop completely idle**.
|
||||
|
||||
This module tracks the process-wide count of in-flight acquisitions that have not
|
||||
yet obtained a connection, and times each acquire so a slow one logs with full
|
||||
pool stats. It is the DB-side counterpart to ``loop_watchdog`` (which covers loop
|
||||
stalls); together, a stuck ``/health`` can be attributed to either a blocked loop
|
||||
or pool exhaustion from the logs alone.
|
||||
|
||||
The counter is a plain int mutated only from the event-loop thread (asyncpg
|
||||
acquisitions are awaited on the loop), so no lock is needed.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import time
|
||||
from collections.abc import AsyncIterator, Callable
|
||||
from contextlib import asynccontextmanager
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
logger = logging.getLogger("hindsight.db.pool")
|
||||
|
||||
_waiting = 0 # callers currently blocked in pool.acquire(), process-wide
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class PoolStats:
|
||||
"""Point-in-time connection-pool utilization snapshot."""
|
||||
|
||||
in_use: int
|
||||
max: int
|
||||
idle: int
|
||||
|
||||
|
||||
def waiting_count() -> int:
|
||||
"""Number of callers currently blocked waiting to acquire a pooled connection."""
|
||||
return _waiting
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def instrument_acquire(
|
||||
acquire_cm: Any,
|
||||
*,
|
||||
pool_stats: Callable[[], PoolStats | None] | None = None,
|
||||
warn_threshold_s: float,
|
||||
) -> AsyncIterator[Any]:
|
||||
"""Wrap a pool's ``acquire()`` context manager with wait tracking + slow-acquire logging.
|
||||
|
||||
Args:
|
||||
acquire_cm: an async context manager yielding a connection (e.g. the object
|
||||
returned by ``asyncpg.Pool.acquire()``).
|
||||
pool_stats: optional zero-arg callable returning a ``PoolStats`` snapshot for
|
||||
the slow-acquire log line.
|
||||
warn_threshold_s: log a warning when the acquire itself takes at least this long.
|
||||
|
||||
Yields:
|
||||
The acquired connection.
|
||||
"""
|
||||
global _waiting
|
||||
_waiting += 1
|
||||
start = time.monotonic()
|
||||
acquired = False
|
||||
try:
|
||||
async with acquire_cm as conn:
|
||||
acquired = True
|
||||
_waiting -= 1
|
||||
_record_acquire_wait(time.monotonic() - start, pool_stats, warn_threshold_s)
|
||||
yield conn
|
||||
finally:
|
||||
# If __aenter__ raised (acquire timeout / cancellation), we never
|
||||
# decremented above — do it here so the waiter count can't leak.
|
||||
if not acquired:
|
||||
_waiting -= 1
|
||||
|
||||
|
||||
async def acquire_conn(
|
||||
acquire_awaitable: Any,
|
||||
*,
|
||||
pool_stats: Callable[[], PoolStats | None] | None = None,
|
||||
warn_threshold_s: float,
|
||||
) -> Any:
|
||||
"""Await a pool acquire that returns a connection, with wait tracking + slow log.
|
||||
|
||||
For pools whose acquire is ``conn = await pool.acquire()`` (oracledb) rather than
|
||||
an async context manager (asyncpg — use ``instrument_acquire`` for those). The
|
||||
caller is responsible for releasing the returned connection.
|
||||
"""
|
||||
global _waiting
|
||||
_waiting += 1
|
||||
start = time.monotonic()
|
||||
try:
|
||||
conn = await acquire_awaitable
|
||||
finally:
|
||||
_waiting -= 1
|
||||
_record_acquire_wait(time.monotonic() - start, pool_stats, warn_threshold_s)
|
||||
return conn
|
||||
|
||||
|
||||
def _record_acquire_wait(
|
||||
wait_s: float,
|
||||
pool_stats: Callable[[], PoolStats | None] | None,
|
||||
warn_threshold_s: float,
|
||||
) -> None:
|
||||
try:
|
||||
from ...metrics import get_metrics_collector
|
||||
|
||||
get_metrics_collector().record_db_acquire_wait(wait_s)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
if wait_s < warn_threshold_s:
|
||||
return
|
||||
|
||||
stats: PoolStats | None = None
|
||||
if pool_stats is not None:
|
||||
try:
|
||||
stats = pool_stats()
|
||||
except Exception:
|
||||
stats = None
|
||||
logger.warning(
|
||||
"slow DB pool acquire: waited %.3fs for a connection "
|
||||
"(in_use=%s max=%s idle=%s waiting=%s). The pool is likely saturated; "
|
||||
"/health can stall on connection acquisition while the event loop is free.",
|
||||
wait_s,
|
||||
stats.in_use if stats else None,
|
||||
stats.max if stats else None,
|
||||
stats.idle if stats else None,
|
||||
_waiting,
|
||||
)
|
||||
@@ -15,7 +15,6 @@ from typing import Any
|
||||
import asyncpg # noqa: F401
|
||||
|
||||
from .base import DatabaseBackend, DatabaseConnection
|
||||
from .pool_instrumentation import PoolStats, instrument_acquire
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -77,8 +76,6 @@ class PostgreSQLBackend(DatabaseBackend):
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._pool: asyncpg.Pool | None = None
|
||||
self._acquire_warn_threshold_s: float = 1.0
|
||||
self._acquire_timeout_s: float | None = None
|
||||
|
||||
async def initialize(
|
||||
self,
|
||||
@@ -91,16 +88,6 @@ class PostgreSQLBackend(DatabaseBackend):
|
||||
statement_cache_size: int = 0,
|
||||
init_callback: Any | None = None,
|
||||
) -> None:
|
||||
from ...config import get_config
|
||||
|
||||
self._acquire_warn_threshold_s = get_config().db_acquire_warn_threshold_ms / 1000.0
|
||||
# Kept for acquire() below: asyncpg's ``timeout`` create_pool kwarg is a
|
||||
# *connect* kwarg (how long establishing a new connection may take), and
|
||||
# ``Pool.acquire()`` defaults to waiting for a free connection forever.
|
||||
# Passing it here alone made HINDSIGHT_API_DB_ACQUIRE_TIMEOUT a no-op for
|
||||
# the wait it names: a pool-exhaustion stall never surfaced as an error,
|
||||
# it just hung (#3002). 0 restores the unbounded behaviour.
|
||||
self._acquire_timeout_s = acquire_timeout if acquire_timeout > 0 else None
|
||||
self._pool = await asyncpg.create_pool(
|
||||
dsn,
|
||||
min_size=min_size,
|
||||
@@ -108,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}, "
|
||||
@@ -121,45 +103,21 @@ 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
|
||||
|
||||
def _pool_stats(self) -> PoolStats | None:
|
||||
"""Snapshot for slow-acquire logs. in_use = live connections minus idle ones."""
|
||||
pool = self._pool
|
||||
if pool is None:
|
||||
return None
|
||||
idle = pool.get_idle_size()
|
||||
return PoolStats(in_use=pool.get_size() - idle, max=pool.get_max_size(), idle=idle)
|
||||
|
||||
@asynccontextmanager
|
||||
async def acquire(self) -> AsyncIterator[PostgresConnection]:
|
||||
pool = self._ensure_pool()
|
||||
async with instrument_acquire(
|
||||
pool.acquire(timeout=self._acquire_timeout_s),
|
||||
pool_stats=self._pool_stats,
|
||||
warn_threshold_s=self._acquire_warn_threshold_s,
|
||||
) as conn:
|
||||
async with pool.acquire() as conn:
|
||||
yield PostgresConnection(conn)
|
||||
|
||||
@asynccontextmanager
|
||||
async def transaction(self) -> AsyncIterator[PostgresConnection]:
|
||||
pool = self._ensure_pool()
|
||||
async with instrument_acquire(
|
||||
pool.acquire(timeout=self._acquire_timeout_s),
|
||||
pool_stats=self._pool_stats,
|
||||
warn_threshold_s=self._acquire_warn_threshold_s,
|
||||
) as conn:
|
||||
async with pool.acquire() as conn:
|
||||
async with conn.transaction():
|
||||
yield PostgresConnection(conn)
|
||||
|
||||
|
||||
@@ -164,6 +164,35 @@ class BudgetedOperation:
|
||||
"""
|
||||
return BudgetedPool(pool, self)
|
||||
|
||||
async def acquire_many(
|
||||
self,
|
||||
pool: Any,
|
||||
count: int,
|
||||
) -> AsyncIterator[list[Any]]:
|
||||
"""
|
||||
Acquire multiple connections within the budget.
|
||||
|
||||
Note: This acquires connections sequentially to respect the budget.
|
||||
For parallel acquisition, use multiple acquire() calls with asyncio.gather().
|
||||
This method is intended for use with raw asyncpg pools only, not DatabaseBackend.
|
||||
|
||||
Args:
|
||||
pool: asyncpg connection pool (raw pool only)
|
||||
count: Number of connections to acquire
|
||||
|
||||
Yields:
|
||||
List of database connections
|
||||
"""
|
||||
connections = []
|
||||
try:
|
||||
for _ in range(count):
|
||||
conn = await pool.acquire()
|
||||
connections.append(conn)
|
||||
yield connections
|
||||
finally:
|
||||
for conn in connections:
|
||||
await pool.release(conn)
|
||||
|
||||
|
||||
# Global default manager instance
|
||||
_default_manager: ConnectionBudgetManager | None = None
|
||||
|
||||
@@ -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..."
|
||||
|
||||
@@ -48,12 +48,6 @@ from ..config import (
|
||||
ENV_LLM_API_KEY,
|
||||
)
|
||||
from .bank_attribution import apply_bank_attribution
|
||||
from .local_device import (
|
||||
release_local_inference_memory,
|
||||
resolve_model_device_type,
|
||||
select_local_device,
|
||||
)
|
||||
from .tei_retry import tei_retry_delay
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -82,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.
|
||||
@@ -161,13 +136,7 @@ class LocalSTEmbeddings(Embeddings):
|
||||
The embedding dimension is auto-detected from the model.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model_name: str | None = None,
|
||||
force_cpu: bool = False,
|
||||
trust_remote_code: bool = False,
|
||||
allow_mps: bool = False,
|
||||
):
|
||||
def __init__(self, model_name: str | None = None, force_cpu: bool = False, trust_remote_code: bool = False):
|
||||
"""
|
||||
Initialize local SentenceTransformers embeddings.
|
||||
|
||||
@@ -179,17 +148,12 @@ class LocalSTEmbeddings(Embeddings):
|
||||
trust_remote_code: Allow loading models with custom code (security risk).
|
||||
Required for some models with custom architectures.
|
||||
Default: False (disabled for security)
|
||||
allow_mps: Opt in to the Apple Silicon MPS GPU. Disabled by default
|
||||
because MPS leaks memory under variable-length workloads
|
||||
(see engine/local_device.py). Default: False
|
||||
"""
|
||||
self.model_name = model_name or DEFAULT_EMBEDDINGS_LOCAL_MODEL
|
||||
self.force_cpu = force_cpu
|
||||
self.trust_remote_code = trust_remote_code
|
||||
self.allow_mps = allow_mps
|
||||
self._model = None
|
||||
self._dimension: int | None = None
|
||||
self._device_type: str = "cpu"
|
||||
|
||||
@property
|
||||
def provider_name(self) -> str:
|
||||
@@ -216,11 +180,31 @@ class LocalSTEmbeddings(Embeddings):
|
||||
|
||||
logger.info(f"Embeddings: initializing local provider with model {self.model_name}")
|
||||
|
||||
# Determine device based on hardware availability. We always set
|
||||
# low_cpu_mem_usage=False to prevent lazy loading (meta tensors) which can
|
||||
# cause issues when accelerate is installed but no GPU is available.
|
||||
# MPS is opt-in (allow_mps) — see engine/local_device.py for why.
|
||||
device = select_local_device(self.force_cpu, self.allow_mps)
|
||||
# Determine device based on hardware availability.
|
||||
# We always set low_cpu_mem_usage=False to prevent lazy loading (meta tensors)
|
||||
# which can cause issues when accelerate is installed but no GPU is available.
|
||||
import torch
|
||||
|
||||
# Force CPU mode if configured (used in daemon mode to avoid MPS/XPC issues on macOS)
|
||||
if self.force_cpu:
|
||||
device = "cpu"
|
||||
logger.info("Embeddings: forcing CPU mode")
|
||||
else:
|
||||
# Check for GPU (CUDA), Apple Silicon (MPS), or Intel XPU
|
||||
# Wrap in try-except to gracefully handle any device detection issues
|
||||
# (e.g., in CI environments or when PyTorch is built without GPU support)
|
||||
device = "cpu" # Default to CPU
|
||||
try:
|
||||
has_gpu = torch.cuda.is_available() or (
|
||||
hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
|
||||
)
|
||||
# Intel Arc XPU support — torch.xpu is available when the XPU build is loaded
|
||||
if not has_gpu and hasattr(torch, "xpu"):
|
||||
has_gpu = torch.xpu.is_available()
|
||||
if has_gpu:
|
||||
device = None # Let sentence-transformers auto-detect GPU/MPS/XPU
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to detect GPU/MPS/XPU, falling back to CPU: {e}")
|
||||
|
||||
# Suppress verbose transformers warnings during model loading
|
||||
# This suppresses the "UNEXPECTED" warnings from BertModel which are harmless
|
||||
@@ -247,8 +231,7 @@ class LocalSTEmbeddings(Embeddings):
|
||||
transformers_logger.setLevel(original_level)
|
||||
|
||||
self._dimension = self._model.get_sentence_embedding_dimension()
|
||||
self._device_type = resolve_model_device_type(self._model)
|
||||
logger.info(f"Embeddings: local provider initialized (dim: {self._dimension}, device: {self._device_type})")
|
||||
logger.info(f"Embeddings: local provider initialized (dim: {self._dimension})")
|
||||
|
||||
def encode(self, texts: list[str]) -> list[list[float]]:
|
||||
"""
|
||||
@@ -263,19 +246,8 @@ class LocalSTEmbeddings(Embeddings):
|
||||
if self._model is None:
|
||||
raise RuntimeError("Embeddings not initialized. Call initialize() first.")
|
||||
|
||||
try:
|
||||
embeddings = self._model.encode(texts, convert_to_numpy=True, show_progress_bar=False)
|
||||
return [emb.tolist() for emb in embeddings]
|
||||
finally:
|
||||
# Only reclaim the GPU allocator pool here, and only when actually on a
|
||||
# GPU (opt-in MPS/CUDA/XPU). encode() runs in tight retain loops, so a
|
||||
# gc.collect()/malloc_trim on every call is too costly on the CPU default
|
||||
# — and unnecessary: refcounting frees the small transient buffers
|
||||
# immediately and the allocator reuses them for the next batch. (The
|
||||
# reranker keeps its per-batch heap trim for the #1717 CPU case; it runs
|
||||
# on the lighter recall path.) See engine/local_device.py.
|
||||
if self._device_type != "cpu":
|
||||
release_local_inference_memory(self._device_type)
|
||||
embeddings = self._model.encode(texts, convert_to_numpy=True, show_progress_bar=False)
|
||||
return [emb.tolist() for emb in embeddings]
|
||||
|
||||
|
||||
class OnnxEmbeddings(Embeddings):
|
||||
@@ -515,20 +487,13 @@ class RemoteTEIEmbeddings(Embeddings):
|
||||
time.sleep(delay)
|
||||
delay *= 2 # Exponential backoff
|
||||
except httpx.HTTPStatusError as e:
|
||||
# TEI uses 429 as normal overload backpressure. Retry it with
|
||||
# the same bounded budget as transient server errors.
|
||||
if (e.response.status_code == 429 or e.response.status_code >= 500) and attempt < self.max_retries:
|
||||
# Retry on 5xx server errors
|
||||
if e.response.status_code >= 500 and attempt < self.max_retries:
|
||||
last_error = e
|
||||
sleep_delay = tei_retry_delay(
|
||||
e.response,
|
||||
delay,
|
||||
request_timeout=self.timeout,
|
||||
)
|
||||
logger.warning(
|
||||
f"TEI transient error (attempt {attempt + 1}/{self.max_retries + 1}): {e}. "
|
||||
f"Retrying in {sleep_delay:.2f}s..."
|
||||
f"TEI server error (attempt {attempt + 1}/{self.max_retries + 1}): {e}. Retrying in {delay}s..."
|
||||
)
|
||||
time.sleep(sleep_delay)
|
||||
time.sleep(delay)
|
||||
delay *= 2
|
||||
else:
|
||||
raise
|
||||
@@ -1237,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.
|
||||
@@ -1252,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
|
||||
@@ -1264,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
|
||||
|
||||
@@ -1341,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
|
||||
@@ -1653,7 +1585,6 @@ def create_embeddings_from_env() -> Embeddings:
|
||||
model_name=config.embeddings_local_model,
|
||||
force_cpu=config.embeddings_local_force_cpu,
|
||||
trust_remote_code=config.embeddings_local_trust_remote_code,
|
||||
allow_mps=config.embeddings_local_allow_mps,
|
||||
)
|
||||
elif provider == "onnx":
|
||||
return OnnxEmbeddings(
|
||||
@@ -1760,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,70 +771,21 @@ 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,
|
||||
unit_entity_pairs: list[tuple[str, str]] | list[tuple[str, str, datetime | None]],
|
||||
conn=None,
|
||||
bank_id: str | None = None,
|
||||
):
|
||||
"""
|
||||
Link multiple memory units to entities in batch (MUCH faster than sequential).
|
||||
@@ -892,32 +813,22 @@ class EntityResolver:
|
||||
|
||||
if conn is None:
|
||||
async with acquire_with_retry(self.pool) as conn:
|
||||
return await self._link_units_to_entities_batch_impl(conn, normalized, bank_id)
|
||||
return await self._link_units_to_entities_batch_impl(conn, normalized)
|
||||
else:
|
||||
return await self._link_units_to_entities_batch_impl(conn, normalized, bank_id)
|
||||
return await self._link_units_to_entities_batch_impl(conn, normalized)
|
||||
|
||||
async def _link_units_to_entities_batch_impl(
|
||||
self, conn, unit_entity_pairs: list[tuple[str, str, datetime | None]], bank_id: str | None = None
|
||||
):
|
||||
async def _link_units_to_entities_batch_impl(self, conn, unit_entity_pairs: list[tuple[str, str, datetime | None]]):
|
||||
# Sorted bulk insert to prevent deadlocks from inconsistent lock ordering
|
||||
# across concurrent transactions on the unit_entities unique index.
|
||||
sorted_pairs = sorted(unit_entity_pairs, key=lambda t: (t[0], t[1]))
|
||||
unit_ids = [p[0] for p in sorted_pairs]
|
||||
entity_ids = [p[1] for p in sorted_pairs]
|
||||
|
||||
# The unit→entity posting belongs to whoever stores the memory, so the
|
||||
# memories store records it. Co-occurrence below is separate and unaffected:
|
||||
# it references only `entities`, which stays in Postgres either way, and is
|
||||
# read by the entity-graph endpoint and by resolution's disambiguation signal.
|
||||
from .memories import get_memories
|
||||
|
||||
await get_memories().record_unit_entities(
|
||||
conn=conn,
|
||||
ops=self._ops,
|
||||
fq_table=fq_table,
|
||||
bank_id=bank_id,
|
||||
unit_ids=unit_ids,
|
||||
entity_ids=entity_ids,
|
||||
await self._ops.bulk_insert_unit_entities(
|
||||
conn,
|
||||
fq_table("unit_entities"),
|
||||
unit_ids,
|
||||
entity_ids,
|
||||
)
|
||||
|
||||
# Build maps keyed by unit_id:
|
||||
@@ -942,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
|
||||
@@ -959,3 +878,58 @@ class EntityResolver:
|
||||
_CooccurrencePair(entity_id_1=e1, entity_id_2=e2, event_date=ed)
|
||||
for (e1, e2), ed in cooccurrence_pairs.items()
|
||||
)
|
||||
|
||||
async def get_units_by_entity(self, entity_id: str, limit: int = 100) -> list[str]:
|
||||
"""
|
||||
Get all units that mention an entity.
|
||||
|
||||
Args:
|
||||
entity_id: Entity ID
|
||||
limit: Max results
|
||||
|
||||
Returns:
|
||||
List of unit IDs
|
||||
"""
|
||||
async with acquire_with_retry(self.pool) as conn:
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT unit_id
|
||||
FROM {fq_table("unit_entities")}
|
||||
WHERE entity_id = $1
|
||||
ORDER BY unit_id
|
||||
LIMIT $2
|
||||
""",
|
||||
entity_id,
|
||||
limit,
|
||||
)
|
||||
return [row["unit_id"] for row in rows]
|
||||
|
||||
async def get_entity_by_text(
|
||||
self,
|
||||
bank_id: str,
|
||||
entity_text: str,
|
||||
) -> str | None:
|
||||
"""
|
||||
Find an entity by text (for query resolution).
|
||||
|
||||
Args:
|
||||
bank_id: bank ID
|
||||
entity_text: Entity text to search for
|
||||
|
||||
Returns:
|
||||
Entity ID if found, None otherwise
|
||||
"""
|
||||
async with acquire_with_retry(self.pool) as conn:
|
||||
row = await conn.fetchrow(
|
||||
f"""
|
||||
SELECT id FROM {fq_table("entities")}
|
||||
WHERE bank_id = $1
|
||||
AND canonical_name ILIKE $2
|
||||
ORDER BY mention_count DESC
|
||||
LIMIT 1
|
||||
""",
|
||||
bank_id,
|
||||
entity_text,
|
||||
)
|
||||
|
||||
return row["id"] if row else None
|
||||
|
||||
@@ -5,25 +5,25 @@ Three reconciliation passes run together on every worker invocation:
|
||||
1. **Relink top-up.** Drain ``graph_maintenance_queue`` (units whose
|
||||
outgoing temporal/semantic links lost a neighbour to a delete). For
|
||||
each, count current outgoing links per type; if below cap, run the
|
||||
same probes retain uses and insert the missing links.
|
||||
same probes retain uses (:func:`fetch_temporal_neighbors`,
|
||||
:func:`compute_semantic_links_ann`) and insert the missing links.
|
||||
``bulk_insert_links`` has ``ON CONFLICT DO NOTHING`` on the uniqueness
|
||||
key, so we can re-probe freely and the DB de-dupes.
|
||||
|
||||
2. **Orphan entity prune.** Delete ``entities`` rows in the bank that no
|
||||
longer have any live memory references. FK ON DELETE CASCADE on
|
||||
longer have any ``unit_entities`` references. FK ON DELETE CASCADE on
|
||||
``entity_cooccurrences`` then removes any cooccurrence row pointing
|
||||
at the pruned entities.
|
||||
|
||||
3. **Stale cooccurrence prune.** Defensive sweep for cooccurrence rows
|
||||
where both endpoints still exist but no current memory references
|
||||
where both endpoints still exist but no current memory_unit references
|
||||
both of them — the cooccurrence was real at the time it was recorded,
|
||||
but every unit that witnessed it has since been deleted.
|
||||
|
||||
Each pass is work the *memories store* owns, because each is a query over
|
||||
`memory_links`, `unit_entities` and `entities` — the slice the store carves
|
||||
out. This module orchestrates them (drain the queue, wrap the sweep in a
|
||||
deadlock-retry) and asks the store to do the part that touches storage. A store
|
||||
whose links travel inside its memories has no `memory_links` to dangle and no
|
||||
join table to sweep, so its relink and cooccurrence passes are no-ops and the
|
||||
job simply prunes the orphan `entities` rows, which stay in Postgres regardless.
|
||||
All three passes run on every invocation. The queue is the only source
|
||||
of work for pass 1; passes 2 and 3 are bank-wide sweeps backed by indexes
|
||||
on ``entities(bank_id)`` and ``unit_entities(entity_id)``, so they're
|
||||
cheap when there's nothing to do.
|
||||
|
||||
The worker dedupes on bank: a second job for the same bank is dropped
|
||||
while one is pending. Once processing starts, a new job becomes the
|
||||
@@ -35,18 +35,18 @@ from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import time
|
||||
import uuid as uuid_module
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from ..models import RequestContext
|
||||
from .db.base import DatabaseConnection
|
||||
|
||||
# Re-exported for callers and tests that import the link caps from here; the caps
|
||||
# themselves live with the link builders the relink pass mirrors — the temporal one
|
||||
# with the retain-time builders, the semantic one with the store's relink pass — so
|
||||
# there is a single definition of each and the two cannot drift.
|
||||
from .memories.pg.graph import MAX_SEMANTIC_LINKS_PER_UNIT # noqa: F401
|
||||
from .retain.link_utils import MAX_TEMPORAL_LINKS_PER_UNIT # noqa: F401
|
||||
from .retain.link_utils import (
|
||||
MAX_TEMPORAL_LINKS_PER_UNIT,
|
||||
_bulk_insert_links,
|
||||
_normalize_datetime,
|
||||
compute_semantic_links_ann,
|
||||
)
|
||||
from .schema import fq_table
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -54,19 +54,17 @@ if TYPE_CHECKING:
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# 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
|
||||
# Mirrors the ``top_k`` default in ``compute_semantic_links_ann`` at retain
|
||||
# time. If you change one, change the other — otherwise victims would either
|
||||
# never reach the cap (probe returns less than the cap) or stay perpetually
|
||||
# under it (cap is higher than retain creates).
|
||||
MAX_SEMANTIC_LINKS_PER_UNIT = 50
|
||||
|
||||
|
||||
@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
|
||||
# Worker fetches this many rows per relink-loop iteration. Bounds
|
||||
# per-iteration probe/insert latency so a 10k-row backlog doesn't hold a
|
||||
# worker slot for minutes. Chosen so the typical iteration runs in well
|
||||
# under 1s.
|
||||
_DRAIN_BATCH_SIZE = 50
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -90,52 +88,67 @@ class JobResult:
|
||||
async def enqueue_relink_victims(
|
||||
conn: DatabaseConnection,
|
||||
bank_id: str,
|
||||
affected_unit_ids: list[str],
|
||||
include_affected_units: bool = False,
|
||||
deleted_unit_ids: list[str],
|
||||
ops: Any,
|
||||
) -> int:
|
||||
"""Enqueue surviving units whose outgoing temporal/semantic links pointed at
|
||||
``affected_unit_ids`` for later link top-up.
|
||||
``deleted_unit_ids`` for later link top-up.
|
||||
|
||||
Must run inside the same transaction that drops those links, *before* the
|
||||
delete (or cascade) fires — once the rows are gone, the join that finds the
|
||||
victims returns nothing.
|
||||
|
||||
``include_affected_units`` covers the case where the affected units are NOT
|
||||
being removed: an edit deletes every link incident to the edited unit but
|
||||
leaves it live, so the unit needs its own outgoing adjacency rebuilt too.
|
||||
Passing it for a unit that will be gone at commit is harmless but pointless
|
||||
— the drain skips queue rows with no live unit — so callers should only set
|
||||
it when the unit survives the transaction.
|
||||
|
||||
Delegated to the memories store: finding the victims is a `memory_links`
|
||||
query, and a store whose links are inline has none, so it returns 0 and the
|
||||
relink pass has nothing to do. The store resolves the dialect it needs from
|
||||
``conn``.
|
||||
Must run inside the same transaction that deletes the units, *before* the
|
||||
cascade fires — once the rows are gone, the join that finds the victims
|
||||
returns nothing.
|
||||
|
||||
Args:
|
||||
conn: Database connection inside the active transaction.
|
||||
bank_id: Bank owning the affected units.
|
||||
affected_unit_ids: Memory_unit IDs whose incident temporal/semantic
|
||||
links are about to be (or are being) removed.
|
||||
include_affected_units: Also enqueue ``affected_unit_ids`` themselves,
|
||||
for callers that leave them live.
|
||||
conn: Database connection inside the active delete transaction.
|
||||
bank_id: Bank owning the deleted units.
|
||||
deleted_unit_ids: Memory_unit IDs about to be (or being) deleted.
|
||||
ops: ``DataAccessOps`` instance, supplies the dialect-specific
|
||||
bulk-insert path.
|
||||
|
||||
Returns:
|
||||
Number of distinct victim units enqueued (0 for a store with no links).
|
||||
Number of distinct victim units enqueued (after dedup against rows
|
||||
already in the queue).
|
||||
"""
|
||||
if not affected_unit_ids:
|
||||
if not deleted_unit_ids:
|
||||
return 0
|
||||
|
||||
from .memories import get_memories
|
||||
deleted_uuids = [uuid_module.UUID(uid) if isinstance(uid, str) else uid for uid in deleted_unit_ids]
|
||||
deleted_str_set = {str(uid) for uid in deleted_uuids}
|
||||
|
||||
return await get_memories().enqueue_relink_victims(
|
||||
conn=conn,
|
||||
fq_table=fq_table,
|
||||
bank_id=bank_id,
|
||||
affected_unit_ids=affected_unit_ids,
|
||||
include_affected_units=include_affected_units,
|
||||
# Find units (other than the ones being deleted) that have an outgoing
|
||||
# temporal/semantic link pointing at a doomed unit. Entity links are
|
||||
# intentionally excluded — they're scheduled for removal and would only
|
||||
# add noise to the recompute job.
|
||||
victim_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT DISTINCT from_unit_id
|
||||
FROM {fq_table("memory_links")}
|
||||
WHERE to_unit_id = ANY($1::uuid[])
|
||||
AND bank_id = $2
|
||||
AND link_type IN ('temporal', 'semantic')
|
||||
""",
|
||||
deleted_uuids,
|
||||
bank_id,
|
||||
)
|
||||
|
||||
victim_ids = [row["from_unit_id"] for row in victim_rows if str(row["from_unit_id"]) not in deleted_str_set]
|
||||
|
||||
if not victim_ids:
|
||||
return 0
|
||||
|
||||
await ops.enqueue_graph_maintenance(
|
||||
conn,
|
||||
fq_table("graph_maintenance_queue"),
|
||||
bank_id,
|
||||
victim_ids,
|
||||
)
|
||||
|
||||
logger.debug(
|
||||
f"[GRAPH_MAINT] Enqueued {len(victim_ids)} relink victims in "
|
||||
f"bank={bank_id} (deleted {len(deleted_unit_ids)} units)"
|
||||
)
|
||||
return len(victim_ids)
|
||||
|
||||
|
||||
async def run_graph_maintenance_job(
|
||||
memory_engine: "MemoryEngine",
|
||||
@@ -150,69 +163,195 @@ async def run_graph_maintenance_job(
|
||||
Per-pass counters from :class:`JobResult`.
|
||||
"""
|
||||
del request_context # accepted for symmetry with other run_*_job helpers
|
||||
from ..config import get_config
|
||||
from .memories import get_memories
|
||||
|
||||
backend = await memory_engine._get_backend()
|
||||
store = get_memories()
|
||||
config = get_config()
|
||||
ops = backend.ops
|
||||
|
||||
result = JobResult()
|
||||
job_start = time.time()
|
||||
|
||||
# --- Pass 1: relink ---
|
||||
# The store owns the whole drain loop: it is a claim → top-up → commit over
|
||||
# its own link table, so how it batches and re-probes is its business — including
|
||||
# the #3034 serialisation (the claim takes queue rows FOR UPDATE in (bank_id,
|
||||
# unit_id) order against a concurrent re-enqueue), which lives in the store's
|
||||
# claim (`ops.claim_graph_maintenance_batch`). A store with no links returns an
|
||||
# empty dict and this is a no-op.
|
||||
relink = await store.relink_pass(backend=backend, fq_table=fq_table, bank_id=bank_id, config=config)
|
||||
result.relink_units_processed = relink.get("relink_units_processed", 0)
|
||||
result.relink_links_added = relink.get("relink_links_added", 0)
|
||||
# Per-iteration loop: claim → top up → commit. We rely on submit-time
|
||||
# dedup to keep at most one job per bank running, so no need for
|
||||
# SKIP LOCKED.
|
||||
iterations = 0
|
||||
while True:
|
||||
from .memory_engine import acquire_with_retry
|
||||
|
||||
async with acquire_with_retry(backend) as conn:
|
||||
async with conn.transaction():
|
||||
unit_ids = await ops.claim_graph_maintenance_batch(
|
||||
conn,
|
||||
fq_table("graph_maintenance_queue"),
|
||||
bank_id,
|
||||
_DRAIN_BATCH_SIZE,
|
||||
)
|
||||
if not unit_ids:
|
||||
break
|
||||
|
||||
result.relink_links_added += await _relink_batch(conn, bank_id, unit_ids, ops, backend)
|
||||
|
||||
result.relink_units_processed += len(unit_ids)
|
||||
iterations += 1
|
||||
|
||||
if iterations > 10000:
|
||||
# Defensive guard against runaway loops — at 50 units/iter that's
|
||||
# 500k targets, far beyond any realistic single-bank backlog.
|
||||
logger.error(
|
||||
f"[GRAPH_MAINT] bank={bank_id} hit iteration cap ({iterations}); aborting relink ({result.as_dict()})"
|
||||
)
|
||||
break
|
||||
|
||||
# --- Pass 2 & 3: entity / cooccurrence sweeps ---
|
||||
# Bank-wide single-statement deletes. Cheap when there's nothing to do.
|
||||
#
|
||||
# Unlike Pass 1's queue claim, these DELETEs aren't protected by any
|
||||
# consistent lock-ordering guarantee: the stale-cooccurrence prune 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.
|
||||
#
|
||||
# The prunes themselves are the store's: the orphan-`entities` sweep applies
|
||||
# to every store (that registry stays in Postgres), while the cooccurrence
|
||||
# sweep is a no-op for a store that never wrote `unit_entities`.
|
||||
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 store.prune_orphan_entities(conn=conn, fq_table=fq_table, bank_id=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 store.prune_stale_cooccurrences(conn=conn, fq_table=fq_table, bank_id=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(
|
||||
f"[GRAPH_MAINT] bank={bank_id} done: {result.as_dict()}, elapsed={elapsed:.2f}s, operation_id={operation_id}"
|
||||
)
|
||||
return result.as_dict()
|
||||
|
||||
|
||||
async def _relink_batch(
|
||||
conn: DatabaseConnection,
|
||||
bank_id: str,
|
||||
victim_ids: list[str],
|
||||
ops: Any,
|
||||
backend: Any,
|
||||
) -> int:
|
||||
"""Top up temporal/semantic links for a batch of victim units. Returns rows inserted."""
|
||||
# Load each victim's metadata. Victims whose units were deleted between
|
||||
# enqueue and now silently drop out — exactly the no-op behaviour we want
|
||||
# for stale queue rows.
|
||||
victim_uuids = [uuid_module.UUID(vid) for vid in victim_ids]
|
||||
victim_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT id::text AS id, event_date, fact_type, embedding::text AS embedding
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE id = ANY($1::uuid[])
|
||||
AND bank_id = $2
|
||||
AND fact_type IN ('experience', 'world')
|
||||
""",
|
||||
victim_uuids,
|
||||
bank_id,
|
||||
)
|
||||
|
||||
if not victim_rows:
|
||||
return 0
|
||||
|
||||
alive_uuids = [uuid_module.UUID(row["id"]) for row in victim_rows]
|
||||
|
||||
# Count current outgoing temporal/semantic links per victim so we only
|
||||
# probe for the ones genuinely below cap. Saves the bulk of the work when
|
||||
# most victims still have plenty of links.
|
||||
count_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT from_unit_id, link_type, COUNT(*) AS cnt
|
||||
FROM {fq_table("memory_links")}
|
||||
WHERE from_unit_id = ANY($1::uuid[])
|
||||
AND bank_id = $2
|
||||
AND link_type IN ('temporal', 'semantic')
|
||||
GROUP BY from_unit_id, link_type
|
||||
""",
|
||||
alive_uuids,
|
||||
bank_id,
|
||||
)
|
||||
counts: dict[tuple[str, str], int] = {}
|
||||
for row in count_rows:
|
||||
counts[(str(row["from_unit_id"]), row["link_type"])] = int(row["cnt"])
|
||||
|
||||
# --- Temporal top-up ---
|
||||
temporal_needs = [r for r in victim_rows if counts.get((r["id"], "temporal"), 0) < MAX_TEMPORAL_LINKS_PER_UNIT]
|
||||
new_links: list[tuple] = []
|
||||
|
||||
if temporal_needs:
|
||||
lateral_unit_ids = [uuid_module.UUID(r["id"]) for r in temporal_needs if r["event_date"] is not None]
|
||||
lateral_event_dates = [
|
||||
_normalize_datetime(r["event_date"]) for r in temporal_needs if r["event_date"] is not None
|
||||
]
|
||||
lateral_fact_types = [r["fact_type"] for r in temporal_needs if r["event_date"] is not None]
|
||||
|
||||
if lateral_unit_ids:
|
||||
rows = await ops.fetch_temporal_neighbors(
|
||||
conn,
|
||||
fq_table("memory_units"),
|
||||
bank_id,
|
||||
lateral_unit_ids,
|
||||
lateral_event_dates,
|
||||
lateral_fact_types,
|
||||
MAX_TEMPORAL_LINKS_PER_UNIT,
|
||||
)
|
||||
for row in rows:
|
||||
time_diff_h = float(row["time_diff_hours"])
|
||||
# Mirror the 24h window enforced at retain time. The bidirectional
|
||||
# index scan returns the K closest neighbours regardless of
|
||||
# window, so we filter here.
|
||||
if time_diff_h > 24:
|
||||
continue
|
||||
weight = max(0.3, 1.0 - (time_diff_h / 24))
|
||||
new_links.append((row["from_id"], str(row["id"]), "temporal", weight, None))
|
||||
|
||||
# --- Semantic top-up ---
|
||||
# ANN must run on its own connection: it opens a nested transaction with
|
||||
# SET LOCAL hnsw.ef_search + CREATE TEMP TABLE ON COMMIT DROP, and nesting
|
||||
# that inside our current write transaction would commit our writes early.
|
||||
semantic_needs = [
|
||||
r
|
||||
for r in victim_rows
|
||||
if counts.get((r["id"], "semantic"), 0) < MAX_SEMANTIC_LINKS_PER_UNIT and r["embedding"] is not None
|
||||
]
|
||||
if semantic_needs:
|
||||
from .memory_engine import acquire_with_retry
|
||||
|
||||
seed_ids = [r["id"] for r in semantic_needs]
|
||||
seed_embs = [r["embedding"] for r in semantic_needs]
|
||||
seed_ftypes = [r["fact_type"] for r in semantic_needs]
|
||||
async with acquire_with_retry(backend) as ann_conn:
|
||||
try:
|
||||
ann_links = await compute_semantic_links_ann(
|
||||
ann_conn,
|
||||
bank_id,
|
||||
seed_ids,
|
||||
seed_embs,
|
||||
fact_types=seed_ftypes,
|
||||
)
|
||||
# Strip self-links (rare but possible because the ANN probe
|
||||
# has no exclude list — see the comment in compute_semantic_links_ann).
|
||||
ann_links = [lnk for lnk in ann_links if lnk[0] != lnk[1]]
|
||||
new_links.extend(ann_links)
|
||||
except Exception as e:
|
||||
# ANN uses PG-specific HNSW syntax; on dialects/configs where
|
||||
# it isn't available we still want the temporal top-up to land.
|
||||
logger.warning(f"[GRAPH_MAINT] Semantic top-up failed for bank={bank_id}: {type(e).__name__}: {e}")
|
||||
|
||||
if not new_links:
|
||||
return 0
|
||||
|
||||
await _bulk_insert_links(
|
||||
conn,
|
||||
new_links,
|
||||
bank_id=bank_id,
|
||||
skip_exists_check=False,
|
||||
ops=ops,
|
||||
)
|
||||
return len(new_links)
|
||||
|
||||
@@ -6,7 +6,6 @@ authentication when a TenantExtension is configured.
|
||||
"""
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
@@ -14,26 +13,9 @@ if TYPE_CHECKING:
|
||||
from hindsight_api.engine.memory_engine import BankLlmHealthInfo, Budget
|
||||
from hindsight_api.engine.response_models import RecallResult, ReflectResult
|
||||
from hindsight_api.engine.search.tags import TagsMatch
|
||||
from hindsight_api.extensions import BankWriteOperation
|
||||
from hindsight_api.models import RequestContext
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class BankConfigState:
|
||||
"""Resolved bank configuration and its bank-level overrides."""
|
||||
|
||||
config: dict[str, Any]
|
||||
overrides: dict[str, Any]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class BankTemplateImportWrite:
|
||||
"""One bank-write decision reserved for a specific imported resource."""
|
||||
|
||||
operation: "BankWriteOperation"
|
||||
target: str | None = None
|
||||
|
||||
|
||||
class MemoryEngineInterface(ABC):
|
||||
"""
|
||||
Abstract interface for the Memory Engine.
|
||||
@@ -198,37 +180,6 @@ class MemoryEngineInterface(ABC):
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def get_bank_config(
|
||||
self,
|
||||
bank_id: str,
|
||||
*,
|
||||
request_context: "RequestContext",
|
||||
) -> BankConfigState:
|
||||
"""Return resolved configuration after authenticating and authorizing the read."""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def update_bank_config(
|
||||
self,
|
||||
bank_id: str,
|
||||
updates: dict[str, Any],
|
||||
*,
|
||||
request_context: "RequestContext",
|
||||
) -> BankConfigState:
|
||||
"""Create a bank if needed and persist validated configuration overrides."""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def reset_bank_config(
|
||||
self,
|
||||
bank_id: str,
|
||||
*,
|
||||
request_context: "RequestContext",
|
||||
) -> BankConfigState:
|
||||
"""Remove all bank configuration overrides after authorization."""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
async def update_bank_disposition(
|
||||
self,
|
||||
@@ -324,8 +275,6 @@ class MemoryEngineInterface(ABC):
|
||||
*,
|
||||
fact_type: str | None = None,
|
||||
search_query: str | None = None,
|
||||
entity_id: str | None = None,
|
||||
created_before: datetime | None = None,
|
||||
limit: int = 100,
|
||||
offset: int = 0,
|
||||
request_context: "RequestContext",
|
||||
@@ -337,8 +286,6 @@ class MemoryEngineInterface(ABC):
|
||||
bank_id: The memory bank ID.
|
||||
fact_type: Filter by fact type.
|
||||
search_query: Full-text search query.
|
||||
entity_id: Filter to memory units linked to this entity ID.
|
||||
created_before: Keep units with ``created_at`` before this instant.
|
||||
limit: Maximum results.
|
||||
offset: Pagination offset.
|
||||
request_context: Request context for authentication.
|
||||
@@ -618,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,
|
||||
@@ -649,8 +572,6 @@ class MemoryEngineInterface(ABC):
|
||||
*,
|
||||
name: str | None = None,
|
||||
mission: str | None = None,
|
||||
config_updates: dict[str, Any] | None = None,
|
||||
create_if_missing: bool = True,
|
||||
request_context: "RequestContext",
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
@@ -660,9 +581,6 @@ class MemoryEngineInterface(ABC):
|
||||
bank_id: The memory bank ID.
|
||||
name: New bank name (optional).
|
||||
mission: New mission text (optional, replaces existing).
|
||||
config_updates: Bank configuration overrides to apply with the profile update.
|
||||
create_if_missing: Create a missing bank when True; otherwise raise
|
||||
a 404 operation error.
|
||||
request_context: Request context for authentication.
|
||||
|
||||
Returns:
|
||||
|
||||
@@ -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]],
|
||||
|
||||
@@ -36,21 +36,6 @@ from .db_utils import acquire_with_retry
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _llm_requests_persistable() -> bool:
|
||||
"""Whether the ``llm_requests`` table exists on the active backend.
|
||||
|
||||
``llm_requests`` is PostgreSQL-only: its migration is ``run_for_dialect(pg=...)``
|
||||
with the Oracle slot intentionally absent, and MaintenanceLoop skips its
|
||||
retention sweep on Oracle for the same reason. On Oracle the table does not
|
||||
exist, so best-effort trace writes must be skipped rather than attempted —
|
||||
otherwise every LLM call fires an INSERT that fails with ORA-00903 and spams
|
||||
the error log. Mirrors the ``_is_oracle()`` gate in MaintenanceLoop.start.
|
||||
"""
|
||||
from .schema import _is_oracle
|
||||
|
||||
return not _is_oracle()
|
||||
|
||||
|
||||
# ── bank/operation attribution (carried across the async call chain) ──────────
|
||||
|
||||
|
||||
@@ -391,32 +376,10 @@ 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:
|
||||
return False
|
||||
if not _llm_requests_persistable():
|
||||
return False
|
||||
if self._allowed_scopes is not None:
|
||||
return scope in self._allowed_scopes
|
||||
return True
|
||||
@@ -510,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
|
||||
@@ -583,7 +546,7 @@ class LLMTraceRecorder:
|
||||
ids are snapshotted synchronously here because the caller may reset the
|
||||
context immediately after.
|
||||
"""
|
||||
if not self._enabled or not _llm_requests_persistable() or trace_ctx is None or not trace_ctx.trace_id:
|
||||
if not self._enabled or trace_ctx is None or not trace_ctx.trace_id:
|
||||
return
|
||||
created_ids = list(dict.fromkeys([*(created or []), *trace_ctx.created_memory_ids]))
|
||||
source_ids = list(dict.fromkeys([*(source or []), *trace_ctx.source_memory_ids]))
|
||||
@@ -605,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:
|
||||
@@ -888,13 +844,7 @@ class LLMProvider:
|
||||
from ..worker.stage import set_stage
|
||||
|
||||
structured = "+structured" if response_format is not None else ""
|
||||
# `.queued` until the concurrency permits are in hand — see the acquire
|
||||
# below. Without it, a call waiting on a saturated semaphore is
|
||||
# indistinguishable from one the provider is actively running, and the
|
||||
# label points at the provider (#3002: an operator lost an hour to
|
||||
# "llm.bedrock.*" for tasks that had never reached Bedrock).
|
||||
base_stage = f"llm.{self.provider}.{scope}{structured}"
|
||||
set_stage(f"{base_stage}.queued")
|
||||
set_stage(f"llm.{self.provider}.{scope}{structured}")
|
||||
|
||||
# Resolve the retry policy: explicit per-call arg wins, else the provider's
|
||||
# configured per-operation/global default, else this method's own fallback.
|
||||
@@ -911,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
|
||||
@@ -954,7 +900,6 @@ class LLMProvider:
|
||||
async with AsyncExitStack() as stack:
|
||||
for sem in _semaphores_for_scope(scope):
|
||||
await stack.enter_async_context(sem)
|
||||
set_stage(base_stage)
|
||||
|
||||
# cached_prefix is only set for providers that returned a handle
|
||||
# from get_or_create_cached_prefix() (e.g. Gemini); it's None for
|
||||
@@ -1020,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.
|
||||
@@ -1039,16 +983,14 @@ 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.
|
||||
"""
|
||||
from ..worker.stage import set_stage
|
||||
|
||||
# `.queued` until the permits are held — see the structured path above.
|
||||
base_stage = f"llm.{self.provider}.{scope}+tools"
|
||||
set_stage(f"{base_stage}.queued")
|
||||
set_stage(f"llm.{self.provider}.{scope}+tools")
|
||||
|
||||
# Resolve the retry policy: explicit per-call arg wins, else the provider's
|
||||
# configured per-operation/global default, else this method's own fallback.
|
||||
@@ -1090,17 +1032,11 @@ class LLMProvider:
|
||||
async with AsyncExitStack() as stack:
|
||||
for sem in _semaphores_for_scope(scope):
|
||||
await stack.enter_async_context(sem)
|
||||
set_stage(base_stage)
|
||||
|
||||
# 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(
|
||||
@@ -1324,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,
|
||||
@@ -1335,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,
|
||||
)
|
||||
|
||||
@@ -1380,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,
|
||||
|
||||
@@ -1,172 +0,0 @@
|
||||
"""Device selection and post-inference memory release for local (in-process)
|
||||
SentenceTransformer / CrossEncoder models.
|
||||
|
||||
Two concerns live here, both about keeping a local API instance's memory flat:
|
||||
|
||||
**1. Device selection — MPS is opt-in.**
|
||||
On Apple Silicon the PyTorch **MPS** (Metal) backend caches a distinct compiled
|
||||
kernel graph *and* allocator pool per unique input tensor shape, and never
|
||||
releases them. Under the variable-length, high-volume recall/rerank/embed traffic
|
||||
the engine generates (documents and candidate sets of every size), that per-shape
|
||||
cache grows without bound. A local instance was observed idling at ~20 GB — ~9.4 GB
|
||||
of Metal graphics memory plus ~8 GB of native heap, essentially all of it stale
|
||||
per-shape MPS cache. CPU inference has no per-shape cache: the same workload holds
|
||||
flat at a few hundred MB, with negligible latency cost for the small default
|
||||
models (and MPS actually *slows down* over time as it recompiles graphs for new
|
||||
shapes). So MPS is excluded from auto-detection and must be opted into explicitly;
|
||||
CUDA and Intel XPU still auto-select.
|
||||
|
||||
This is a confirmed, still-open PyTorch bug in the MPSGraph compilation cache
|
||||
(keyed on tensor shape, no eviction path). We are tracking it upstream:
|
||||
- https://github.com/pytorch/pytorch/issues/181213
|
||||
([MPS] unbounded RSS growth with varying-shape inference — our exact case)
|
||||
- https://github.com/pytorch/pytorch/issues/164299 (graphCache identified as
|
||||
the primary leak culprit)
|
||||
- https://github.com/pytorch/pytorch/issues/182815 (proposes, but has not yet
|
||||
shipped, a torch.mps.invalidate_graph_cache() API / PYTORCH_MPS_DISABLE_GRAPH_CACHE
|
||||
env var that would let us keep MPS)
|
||||
No released mitigation exists today: empty_cache(), synchronize(),
|
||||
PYTORCH_MPS_HIGH_WATERMARK_RATIO, and autorelease pools were all confirmed
|
||||
ineffective upstream. Revisit MPS-as-default once one of those knobs lands.
|
||||
|
||||
**2. Memory release after each batch.**
|
||||
Local CPU inference allocates large transient numpy/tensor buffers per call. The
|
||||
allocator keeps those freed pages as a high-water mark, so RSS grows monotonically
|
||||
across many calls (issue #1717). We return them to the OS after each batch —
|
||||
``malloc_trim`` on glibc/Linux, ``malloc_zone_pressure_relief`` on macOS (the
|
||||
original #1717 fix covered only Linux). When the model ran on a GPU we also empty
|
||||
that backend's allocator pool via ``torch.<backend>.empty_cache()``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import ctypes
|
||||
import ctypes.util
|
||||
import gc
|
||||
import logging
|
||||
import sys
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def select_local_device(force_cpu: bool, allow_mps: bool) -> str | None:
|
||||
"""Choose the device for a local SentenceTransformer / CrossEncoder.
|
||||
|
||||
Returns a value suitable to pass as the model's ``device`` argument:
|
||||
|
||||
- ``"cpu"`` — forced CPU, or the only accelerator is MPS and it is not allowed.
|
||||
- ``None`` — let sentence-transformers auto-detect (picks CUDA / XPU,
|
||||
handling multi-GPU correctly).
|
||||
- ``"mps"`` — Apple Silicon GPU, only when ``allow_mps`` is set.
|
||||
|
||||
MPS is never auto-selected because its per-shape cache leaks unbounded memory
|
||||
under the engine's variable-length workload (see the module docstring). Set the
|
||||
matching ``*_ALLOW_MPS`` config flag to opt back in.
|
||||
"""
|
||||
if force_cpu:
|
||||
return "cpu"
|
||||
|
||||
try:
|
||||
import torch
|
||||
|
||||
if torch.cuda.is_available():
|
||||
return None # auto-detect CUDA
|
||||
if hasattr(torch, "xpu") and torch.xpu.is_available():
|
||||
return None # auto-detect Intel XPU
|
||||
|
||||
mps_available = hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
|
||||
if mps_available:
|
||||
if allow_mps:
|
||||
return "mps"
|
||||
logger.info(
|
||||
"Local model: MPS (Apple Silicon GPU) is available but disabled by "
|
||||
"default because its per-shape cache leaks memory under variable-length "
|
||||
"workloads; running on CPU. Set the *_ALLOW_MPS flag to opt in."
|
||||
)
|
||||
return "cpu"
|
||||
return "cpu"
|
||||
except Exception as e: # pragma: no cover - defensive
|
||||
logger.warning("Local device detection failed, falling back to CPU: %s", e)
|
||||
return "cpu"
|
||||
|
||||
|
||||
def resolve_model_device_type(model: object) -> str:
|
||||
"""Best-effort device *type* ("cpu" / "cuda" / "mps" / "xpu") of a loaded model.
|
||||
|
||||
Used to decide which GPU allocator pool to empty after inference. Falls back to
|
||||
``"cpu"`` (the safe no-op choice for release) if the device can't be read.
|
||||
"""
|
||||
device = getattr(model, "device", None)
|
||||
if device is None:
|
||||
inner = getattr(model, "model", None) # CrossEncoder wraps the HF model
|
||||
device = getattr(inner, "device", None)
|
||||
try:
|
||||
return device.type if device is not None else "cpu"
|
||||
except Exception: # pragma: no cover - defensive
|
||||
return "cpu"
|
||||
|
||||
|
||||
def _resolve_heap_trim():
|
||||
"""Return a callable that asks the C allocator to release freed pages to the OS.
|
||||
|
||||
glibc (Linux) exposes ``malloc_trim``; macOS exposes
|
||||
``malloc_zone_pressure_relief``. Resolved once at import; returns a no-op on
|
||||
platforms where neither is available (musl, Windows).
|
||||
"""
|
||||
if sys.platform == "linux":
|
||||
libc_path = ctypes.util.find_library("c")
|
||||
if libc_path is None:
|
||||
return lambda: None
|
||||
try:
|
||||
libc = ctypes.CDLL(libc_path)
|
||||
trim = libc.malloc_trim
|
||||
except (OSError, AttributeError):
|
||||
# Not glibc (musl has no malloc_trim) or libc lookup failed.
|
||||
return lambda: None
|
||||
trim.argtypes = [ctypes.c_size_t]
|
||||
trim.restype = ctypes.c_int
|
||||
return lambda: trim(0)
|
||||
|
||||
if sys.platform == "darwin":
|
||||
try:
|
||||
libc = ctypes.CDLL("/usr/lib/libSystem.dylib")
|
||||
default_zone = libc.malloc_default_zone
|
||||
default_zone.restype = ctypes.c_void_p
|
||||
relief = libc.malloc_zone_pressure_relief
|
||||
relief.argtypes = [ctypes.c_void_p, ctypes.c_size_t]
|
||||
relief.restype = ctypes.c_size_t
|
||||
except (OSError, AttributeError):
|
||||
return lambda: None
|
||||
# pressure_relief(zone, goal=0) reclaims as much as possible.
|
||||
return lambda: relief(default_zone(), 0)
|
||||
|
||||
return lambda: None
|
||||
|
||||
|
||||
_heap_trim = _resolve_heap_trim()
|
||||
|
||||
|
||||
def _empty_gpu_cache(device_type: str | None) -> None:
|
||||
"""Empty the allocator pool of the GPU backend the model ran on, if any."""
|
||||
if not device_type or device_type == "cpu":
|
||||
return
|
||||
try:
|
||||
import torch
|
||||
|
||||
backend = getattr(torch, device_type, None) # torch.cuda / torch.mps / torch.xpu
|
||||
if backend is not None and hasattr(backend, "empty_cache"):
|
||||
backend.empty_cache()
|
||||
except Exception: # pragma: no cover - defensive
|
||||
pass
|
||||
|
||||
|
||||
def release_local_inference_memory(device_type: str | None = None) -> None:
|
||||
"""Release transient heap (and GPU allocator) memory after a local inference batch.
|
||||
|
||||
Frees Python objects, returns freed native pages to the OS, and empties the GPU
|
||||
allocator pool when the model ran on a GPU. Safe to call on every platform and
|
||||
device; the pieces that don't apply are cheap no-ops.
|
||||
"""
|
||||
gc.collect()
|
||||
_heap_trim()
|
||||
_empty_gpu_cache(device_type)
|
||||
@@ -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,18 +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
|
||||
# Cross-store txn recovery (only when the memories store keeps its rows outside SQL): a backstop
|
||||
# for a writer that crashed between its external writes and the decide. The happy path decides
|
||||
# inline after commit, so this rarely finds work; five minutes bounds how long a crashed txn stalls
|
||||
# its namespace's fold.
|
||||
_TXN_RECOVERY_INTERVAL_SECONDS = 300
|
||||
# A pending txn is left alone for this long from first sighting before the sweep aborts an
|
||||
# unwitnessed one — the writer may still be mid-flight (PendingTxn carries no timestamp).
|
||||
_TXN_RECOVERY_GRACE_SECONDS = 300
|
||||
|
||||
|
||||
class MaintenanceLoop:
|
||||
@@ -73,9 +60,6 @@ class MaintenanceLoop:
|
||||
self._stop = asyncio.Event()
|
||||
# Monotonic timestamps of the last run per job, keyed by job name.
|
||||
self._last_run: dict[str, float] = {}
|
||||
# Cross-store txn recovery: first-sighting time per pending txn_id, so an unwitnessed
|
||||
# txn gets a grace period before the sweep aborts it. Persists across ticks.
|
||||
self._txn_first_seen: dict[str, float] = {}
|
||||
|
||||
# ── lifecycle ──────────────────────────────────────────────────────────
|
||||
|
||||
@@ -113,32 +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
|
||||
or MaintenanceLoop._cross_store_recovery_enabled()
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _cross_store_recovery_enabled() -> bool:
|
||||
"""True when the memories store keeps memories outside SQL and therefore has
|
||||
cross-store write-group txns a crashed writer could leave undecided."""
|
||||
try:
|
||||
from .memories import get_memories
|
||||
|
||||
return not get_memories().writes_memory_rows_in_sql
|
||||
except Exception:
|
||||
return False
|
||||
return reconcile_on or audit_on or llm_on or mm_refresh_on
|
||||
|
||||
# ── loop ───────────────────────────────────────────────────────────────
|
||||
|
||||
@@ -172,10 +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))
|
||||
if self._cross_store_recovery_enabled() and self._is_due("txn_recovery", _TXN_RECOVERY_INTERVAL_SECONDS):
|
||||
await self._run_timed("cross-store txn recovery", self._run_txn_recovery())
|
||||
|
||||
async def _run_timed(self, name: str, coro: Coroutine[Any, Any, None]) -> None:
|
||||
"""Run a maintenance job and emit one timing line for it.
|
||||
@@ -194,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)
|
||||
@@ -208,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]
|
||||
@@ -223,109 +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")
|
||||
|
||||
# ── cross-store txn recovery ─────────────────────────────────────────────
|
||||
|
||||
async def _run_txn_recovery(self) -> None:
|
||||
"""Resolve write-group txns a crashed writer left undecided, for a store that keeps its
|
||||
rows outside SQL.
|
||||
|
||||
For each bank, the store lists its namespace's pending txns and decides each against the
|
||||
Postgres witness table (present ⇒ commit, absent past the grace ⇒ abort — never on
|
||||
assumption), then reaps expired witness rows. A no-op for the SQL stores. Best-effort: a
|
||||
failure here only delays a stalled fold until the next tick.
|
||||
"""
|
||||
from .memories import get_memories
|
||||
|
||||
store = get_memories()
|
||||
if store.writes_memory_rows_in_sql:
|
||||
return
|
||||
backend = self._engine._backend
|
||||
try:
|
||||
async with acquire_with_retry(backend, max_retries=1) as conn:
|
||||
bank_ids = [r[0] for r in await conn.fetch(f"SELECT bank_id FROM {fq_table('banks')}")]
|
||||
if not bank_ids:
|
||||
return
|
||||
decided = await store.recover_pending_txns(
|
||||
conn=conn,
|
||||
fq_table=fq_table,
|
||||
bank_ids=bank_ids,
|
||||
first_seen=self._txn_first_seen,
|
||||
now=time.monotonic(),
|
||||
grace_seconds=_TXN_RECOVERY_GRACE_SECONDS,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"Cross-store txn recovery failed: {e}")
|
||||
return
|
||||
if decided:
|
||||
logger.info(f"Cross-store txn recovery: decided {decided} undecided txn(s)")
|
||||
|
||||
# ── consolidation reconcile ──────────────────────────────────────────────
|
||||
|
||||
async def _run_reconcile(self) -> None:
|
||||
@@ -333,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
|
||||
@@ -394,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
|
||||
@@ -406,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}")
|
||||
|
||||
@@ -1,86 +0,0 @@
|
||||
"""The memories store: which one is installed, and how the engine reaches it.
|
||||
|
||||
Resolved through the ordinary extension loader — ``HINDSIGHT_API_MEMORIES_EXTENSION``
|
||||
names a ``module:Class``, and ``HINDSIGHT_API_MEMORIES_*`` becomes its config — so
|
||||
this behaves like every other extension point. Unset (the normal case) means
|
||||
:class:`~hindsight_api.engine.memories.postgres.PostgresMemories`: rows in
|
||||
`memory_units`, links in `memory_links` / `unit_entities`, retrieval as SQL.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
|
||||
from .base import (
|
||||
META_CHUNK_ID,
|
||||
CausalEdgeRecord,
|
||||
DeletePredicate,
|
||||
FactRecord,
|
||||
MemoriesExtension,
|
||||
MemoryPatch,
|
||||
ScanPage,
|
||||
StoredMemory,
|
||||
build_fact_records,
|
||||
build_text_signals,
|
||||
source_key,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_memories: MemoriesExtension | None = None
|
||||
|
||||
|
||||
def create_memories(context=None) -> MemoriesExtension:
|
||||
"""Build the configured memories store, or the Postgres default."""
|
||||
from ...extensions.loader import load_extension
|
||||
|
||||
loaded = load_extension("MEMORIES", MemoriesExtension, context=context)
|
||||
if loaded is not None:
|
||||
logger.info("[memories] store=%s (memory rows do not go to postgres)", loaded.name)
|
||||
return loaded
|
||||
|
||||
from .postgres import PostgresMemories
|
||||
|
||||
return PostgresMemories({})
|
||||
|
||||
|
||||
def get_memories() -> MemoriesExtension:
|
||||
"""The process-wide memories store, built on first use.
|
||||
|
||||
Retrieval and the retain pipeline reach it through call chains that do not
|
||||
carry the engine, so it is resolved here rather than threaded through every
|
||||
signature.
|
||||
"""
|
||||
global _memories
|
||||
if _memories is None:
|
||||
_memories = create_memories()
|
||||
return _memories
|
||||
|
||||
|
||||
def set_memories(memories: MemoriesExtension | None) -> None:
|
||||
"""Override the store (tests, and engine startup after initialize())."""
|
||||
global _memories
|
||||
_memories = memories
|
||||
# The graph arm's retriever is chosen from the store and then cached, so it
|
||||
# has to be re-resolved whenever the store changes.
|
||||
from ..search.retrieval import set_default_graph_retriever
|
||||
|
||||
set_default_graph_retriever(None)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"META_CHUNK_ID",
|
||||
"CausalEdgeRecord",
|
||||
"DeletePredicate",
|
||||
"FactRecord",
|
||||
"MemoriesExtension",
|
||||
"MemoryPatch",
|
||||
"ScanPage",
|
||||
"StoredMemory",
|
||||
"build_fact_records",
|
||||
"build_text_signals",
|
||||
"create_memories",
|
||||
"get_memories",
|
||||
"set_memories",
|
||||
"source_key",
|
||||
]
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,20 +0,0 @@
|
||||
"""The Postgres memories implementation, split by what calls it.
|
||||
|
||||
:class:`~hindsight_api.engine.memories.postgres.PostgresMemories` is a thin class
|
||||
over these modules; the queries live here, grouped by concern rather than piled
|
||||
behind one object:
|
||||
|
||||
* :mod:`counts` — the stats/admin aggregates (freshness, per-doc, timeseries, scopes)
|
||||
* :mod:`curation` — the memory/entity list and detail views
|
||||
* :mod:`graph` — the graph view, entity postings, and the maintenance passes
|
||||
* :mod:`reads` — addressed reads: get, scan, count, tags, consolidation state
|
||||
* :mod:`writes` — inserts, deletes, and observation invalidation
|
||||
|
||||
Every function here takes the live connection and Hindsight's ``fq_table``
|
||||
resolver rather than reaching for globals, so each is callable from a
|
||||
transaction the caller already owns.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
__all__ = ["counts", "curation", "graph", "reads", "writes"]
|
||||
@@ -1,151 +0,0 @@
|
||||
"""The count/aggregate surfaces: consolidation freshness, per-document counts,
|
||||
ingestion over time, observation scopes.
|
||||
|
||||
Each is one ``GROUP BY`` (or filtered ``COUNT``) over `memory_units`. They back
|
||||
the stats and admin views, not retrieval, so they are grouped here away from the
|
||||
addressed reads. The SQL is lifted verbatim from the engine methods that used to
|
||||
carry it; only the connection and ``fq_table`` resolver are now parameters.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable
|
||||
from datetime import datetime
|
||||
from typing import Any
|
||||
|
||||
|
||||
async def consolidation_freshness(*, conn, fq_table: Callable[[str], str], bank_id: str) -> dict[str, Any]:
|
||||
"""Last consolidation time plus the pending / failed fact counts, in one scan.
|
||||
|
||||
All three come from a single pass so keeping ``failed`` — part of the
|
||||
published contract — costs nothing over reflect()'s ``pending`` read.
|
||||
"""
|
||||
row = await conn.fetchrow(
|
||||
f"""
|
||||
SELECT
|
||||
MAX(consolidated_at) AS last_consolidated_at,
|
||||
COUNT(*) FILTER (WHERE consolidated_at IS NULL AND fact_type IN ('experience', 'world')) AS pending,
|
||||
COUNT(*) FILTER (WHERE consolidation_failed_at IS NOT NULL AND fact_type IN ('experience', 'world')) AS failed
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE bank_id = $1
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
if row is None:
|
||||
return {"last_consolidated_at": None, "pending": 0, "failed": 0}
|
||||
return {
|
||||
"last_consolidated_at": row["last_consolidated_at"],
|
||||
"pending": row["pending"] or 0,
|
||||
"failed": row["failed"] or 0,
|
||||
}
|
||||
|
||||
|
||||
async def link_counts(*, conn, fq_table: Callable[[str], str], bank_id: str) -> dict[str, int]:
|
||||
"""``{link_type: count}`` of live links in a bank.
|
||||
|
||||
Non-entity links (temporal / semantic / caused_by) are a single ``GROUP BY`` over
|
||||
``memory_links``. Entity links are no longer stored there — they are derived on demand
|
||||
from ``unit_entities``, replicating the historical writer cap of ``MAX_LINKS_PER_ENTITY``
|
||||
bidirectional edges per shared entity — so they are aggregated to one ``entity`` scalar.
|
||||
"""
|
||||
max_links_per_entity = 10
|
||||
non_entity_link_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT link_type, COUNT(*) as count
|
||||
FROM {fq_table("memory_links")}
|
||||
WHERE bank_id = $1
|
||||
GROUP BY link_type
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
entity_total_row = await conn.fetchrow(
|
||||
f"""
|
||||
WITH per_entity AS (
|
||||
SELECT ue.entity_id, COUNT(*) AS n
|
||||
FROM {fq_table("unit_entities")} ue
|
||||
JOIN {fq_table("memory_units")} mu ON mu.id = ue.unit_id
|
||||
WHERE mu.bank_id = $1
|
||||
GROUP BY ue.entity_id
|
||||
)
|
||||
SELECT COALESCE(SUM(LEAST(n - 1, $2)), 0)::bigint AS count
|
||||
FROM per_entity
|
||||
""",
|
||||
bank_id,
|
||||
max_links_per_entity,
|
||||
)
|
||||
entity_link_total = int(entity_total_row["count"] or 0) if entity_total_row else 0
|
||||
counts: dict[str, int] = {row["link_type"]: row["count"] for row in non_entity_link_rows}
|
||||
if entity_link_total > 0:
|
||||
counts["entity"] = entity_link_total
|
||||
return counts
|
||||
|
||||
|
||||
async def document_memory_counts(
|
||||
*, conn, fq_table: Callable[[str], str], bank_id: str, document_ids: list[str]
|
||||
) -> dict[str, int]:
|
||||
"""Live memory count per document id, for the ids given."""
|
||||
if not document_ids:
|
||||
return {}
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT document_id, COUNT(*) AS unit_count
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE bank_id = $1 AND document_id = ANY($2::text[])
|
||||
GROUP BY document_id
|
||||
""",
|
||||
bank_id,
|
||||
list(document_ids),
|
||||
)
|
||||
return {row["document_id"]: row["unit_count"] for row in rows}
|
||||
|
||||
|
||||
async def memories_timeseries(
|
||||
*, conn, fq_table: Callable[[str], str], bank_id: str, time_field: str, trunc: str, since: datetime
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Memories bucketed by ``time_field`` (truncated to ``trunc``) and fact_type.
|
||||
|
||||
``time_field`` is whitelisted by the caller before it reaches here — it is
|
||||
interpolated into SQL. Event-time fields fall back to ``created_at`` per row so
|
||||
rows without an event timestamp still appear.
|
||||
"""
|
||||
bucket_expr = time_field if time_field == "created_at" else f"COALESCE({time_field}, created_at)"
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT date_trunc('{trunc}', {bucket_expr} AT TIME ZONE 'UTC') AS bucket,
|
||||
fact_type, COUNT(*) AS count
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE bank_id = $1 AND {bucket_expr} >= $2
|
||||
GROUP BY bucket, fact_type
|
||||
ORDER BY bucket
|
||||
""",
|
||||
bank_id,
|
||||
since,
|
||||
)
|
||||
return [{"bucket": r["bucket"], "fact_type": r["fact_type"], "count": r["count"]} for r in rows]
|
||||
|
||||
|
||||
async def observation_scope_counts(*, conn, fq_table: Callable[[str], str], bank_id: str) -> list[dict[str, Any]]:
|
||||
"""Observations grouped by scope (their sorted tag set), most-populous first."""
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT scope, COUNT(*) AS count
|
||||
FROM (
|
||||
SELECT COALESCE(ARRAY(SELECT unnest(tags) ORDER BY 1), '{{}}'::text[]) AS scope
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE bank_id = $1 AND fact_type = 'observation'
|
||||
) s
|
||||
GROUP BY scope
|
||||
ORDER BY count DESC, scope
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
return [{"tags": list(r["scope"]), "count": r["count"]} for r in rows]
|
||||
|
||||
|
||||
__all__ = [
|
||||
"consolidation_freshness",
|
||||
"document_memory_counts",
|
||||
"link_counts",
|
||||
"memories_timeseries",
|
||||
"observation_scope_counts",
|
||||
]
|
||||
@@ -1,507 +0,0 @@
|
||||
"""Curation reads: the memory list, the memory detail view, and the entity list.
|
||||
|
||||
These back the curation UI — the table of memories a bank holds, the detail panel
|
||||
for one of them, and the entity roster beside it. They are paged and filtered
|
||||
rather than ranked: nothing here scores anything, and nothing walks the corpus.
|
||||
|
||||
Two things separate them from the addressed reads in :mod:`reads`. They render
|
||||
*view* dicts (ISO strings, joined entity names, a ``state`` discriminator) rather
|
||||
than :class:`~hindsight_api.engine.memories.base.StoredMemory`, because the HTTP
|
||||
layer serialises what comes back verbatim. And they read the archive as well as
|
||||
the live table: curation moves an invalidated fact to `invalidated_memory_units`,
|
||||
so "show me the invalidated ones" is a different table, not a different predicate.
|
||||
|
||||
Authentication, operation validation and audit stay with the engine methods that
|
||||
call these — only the queries and their row rendering live here.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from datetime import datetime
|
||||
from typing import Any
|
||||
|
||||
from ...search.tags import build_tags_where_clause
|
||||
|
||||
|
||||
def _entity_rows_for_units_sql(*, ops, fq_table, unit_ids_placeholder: int) -> str:
|
||||
"""SQL SELECT producing ``(unit_id, entity_id, canonical_name)`` rows for
|
||||
the given unit IDs.
|
||||
|
||||
Direct rows come from ``unit_entities``. Observations rarely carry
|
||||
direct rows there; their entity association lives transitively through
|
||||
their source memories (``source_memory_ids`` on PG, the
|
||||
``observation_sources`` junction on Oracle). When an observation has
|
||||
no direct entity rows the SELECT inherits its source memories'
|
||||
entities, so the result is the same set callers would get from
|
||||
``get_memory_unit``.
|
||||
|
||||
``unit_ids_placeholder`` is the 1-based parameter index that holds the
|
||||
``uuid[]`` of unit IDs. The placeholder is referenced twice — both
|
||||
sides of the UNION need it — so callers should not reuse the slot.
|
||||
"""
|
||||
ue = fq_table("unit_entities")
|
||||
ents = fq_table("entities")
|
||||
mu = fq_table("memory_units")
|
||||
p = unit_ids_placeholder
|
||||
|
||||
direct = (
|
||||
f"SELECT ue.unit_id, e.id AS entity_id, e.canonical_name "
|
||||
f"FROM {ue} ue "
|
||||
f"JOIN {ents} e ON e.id = ue.entity_id "
|
||||
f"WHERE ue.unit_id = ANY(${p}::uuid[])"
|
||||
)
|
||||
|
||||
if ops.uses_observation_sources_table:
|
||||
os_t = fq_table("observation_sources")
|
||||
inherited = (
|
||||
f"SELECT os.observation_id AS unit_id, e.id AS entity_id, e.canonical_name "
|
||||
f"FROM {os_t} os "
|
||||
f"JOIN {ue} src_ue ON src_ue.unit_id = os.source_id "
|
||||
f"JOIN {ents} e ON e.id = src_ue.entity_id "
|
||||
f"WHERE os.observation_id = ANY(${p}::uuid[]) "
|
||||
f"AND NOT EXISTS (SELECT 1 FROM {ue} d WHERE d.unit_id = os.observation_id)"
|
||||
)
|
||||
else:
|
||||
inherited = (
|
||||
f"SELECT obs.id AS unit_id, e.id AS entity_id, e.canonical_name "
|
||||
f"FROM {mu} obs "
|
||||
f"CROSS JOIN LATERAL unnest(obs.source_memory_ids) AS src_id "
|
||||
f"JOIN {ue} src_ue ON src_ue.unit_id = src_id "
|
||||
f"JOIN {ents} e ON e.id = src_ue.entity_id "
|
||||
f"WHERE obs.id = ANY(${p}::uuid[]) "
|
||||
f"AND obs.fact_type = 'observation' "
|
||||
f"AND obs.source_memory_ids IS NOT NULL "
|
||||
f"AND NOT EXISTS (SELECT 1 FROM {ue} d WHERE d.unit_id = obs.id)"
|
||||
)
|
||||
|
||||
return f"({direct}) UNION ({inherited})"
|
||||
|
||||
|
||||
async def list_memory_units(
|
||||
*,
|
||||
conn,
|
||||
ops,
|
||||
fq_table,
|
||||
bank_id: str,
|
||||
fact_type: str | None = None,
|
||||
search_query: str | None = None,
|
||||
consolidation_state: str | None = None,
|
||||
state: str | None = None,
|
||||
document_id: str | None = None,
|
||||
entity_id: str | None = None,
|
||||
tags: list[str] | None = None,
|
||||
tags_match: str = "any",
|
||||
created_before: datetime | None = None,
|
||||
limit: int = 100,
|
||||
offset: int = 0,
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
List memory units for table view with optional full-text search.
|
||||
|
||||
Args:
|
||||
conn: Open database connection (the caller owns the transaction).
|
||||
ops: Dialect ops. Unused by this query; part of the interface signature.
|
||||
fq_table: Table-name resolver.
|
||||
bank_id: Filter by bank ID
|
||||
fact_type: Filter by fact type (world, experience)
|
||||
search_query: Full-text search query (searches text and context fields)
|
||||
document_id: Optional filter to a single source document.
|
||||
tags: Optional list of tag names to filter by. When omitted, no tag
|
||||
filtering is applied (except tags_match='exact', which then selects
|
||||
the untagged/global scope).
|
||||
tags_match: How to combine tags (same modes as recall): 'any' (OR,
|
||||
default) or 'all' (AND) both also include untagged units;
|
||||
'any_strict'/'all_strict' exclude untagged units; 'exact' matches
|
||||
units whose tag set equals the given tags exactly.
|
||||
state: Optional curation-state filter ('valid' or 'invalidated').
|
||||
Invalidated facts live in a separate archive table; 'invalidated'
|
||||
reads that archive. Omitted/('valid') lists live facts.
|
||||
consolidation_state: Optional filter on consolidation state. One of
|
||||
'failed' (consolidation permanently failed and awaiting recovery),
|
||||
'pending' (not yet consolidated, no failure), or
|
||||
'done' (successfully consolidated). Only applies to source memory
|
||||
types (world/experience).
|
||||
limit: Maximum number of results to return
|
||||
offset: Offset for pagination
|
||||
|
||||
Returns:
|
||||
Dict with items (list of memory units) and total count
|
||||
"""
|
||||
if state is not None and state not in ("valid", "invalidated"):
|
||||
raise ValueError(f"Invalid state '{state}': expected 'valid' or 'invalidated'.")
|
||||
if entity_id is not None:
|
||||
import uuid as _uuid
|
||||
|
||||
try:
|
||||
_uuid.UUID(entity_id)
|
||||
except ValueError:
|
||||
raise ValueError(f"Invalid entity_id: '{entity_id}' is not a valid UUID") from None
|
||||
# Invalidated facts live in a separate archive table; pick the source
|
||||
# accordingly. Default (state is None) lists live facts.
|
||||
is_archived = state == "invalidated"
|
||||
source_table = fq_table("invalidated_memory_units") if is_archived else fq_table("memory_units")
|
||||
|
||||
# Build query conditions
|
||||
query_conditions = []
|
||||
query_params = []
|
||||
param_count = 0
|
||||
|
||||
if bank_id:
|
||||
param_count += 1
|
||||
query_conditions.append(f"bank_id = ${param_count}")
|
||||
query_params.append(bank_id)
|
||||
|
||||
if fact_type:
|
||||
param_count += 1
|
||||
query_conditions.append(f"fact_type = ${param_count}")
|
||||
query_params.append(fact_type)
|
||||
|
||||
if document_id:
|
||||
param_count += 1
|
||||
query_conditions.append(f"document_id = ${param_count}")
|
||||
query_params.append(document_id)
|
||||
|
||||
if entity_id:
|
||||
# Reverse lookup via the stored entity links. Entity links reference live memory units, so
|
||||
# this yields nothing against the invalidated archive (documented on the method).
|
||||
param_count += 1
|
||||
query_conditions.append(
|
||||
f"id IN (SELECT unit_id FROM {fq_table('unit_entities')} WHERE entity_id = ${param_count}::uuid)"
|
||||
)
|
||||
query_params.append(entity_id)
|
||||
|
||||
if search_query:
|
||||
# Full-text search on text and context fields using ILIKE
|
||||
param_count += 1
|
||||
query_conditions.append(f"(text ILIKE ${param_count} OR context ILIKE ${param_count})")
|
||||
query_params.append(f"%{search_query}%")
|
||||
|
||||
if consolidation_state:
|
||||
# Named apart from `state`, which the engine method used to shadow here;
|
||||
# `is_archived` was already resolved above, so behaviour is unchanged.
|
||||
wanted = consolidation_state.lower()
|
||||
if wanted == "failed":
|
||||
query_conditions.append("consolidation_failed_at IS NOT NULL AND fact_type IN ('experience', 'world')")
|
||||
elif wanted == "pending":
|
||||
query_conditions.append(
|
||||
"consolidated_at IS NULL AND consolidation_failed_at IS NULL AND fact_type IN ('experience', 'world')"
|
||||
)
|
||||
elif wanted == "done":
|
||||
query_conditions.append("consolidated_at IS NOT NULL AND fact_type IN ('experience', 'world')")
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Invalid consolidation_state '{consolidation_state}': expected 'failed', 'pending', or 'done'."
|
||||
)
|
||||
|
||||
if tags:
|
||||
tags_clause, tags_params, next_param = build_tags_where_clause(tags, param_count + 1, "", tags_match)
|
||||
if tags_clause:
|
||||
query_conditions.append(tags_clause.removeprefix("AND "))
|
||||
query_params.extend(tags_params)
|
||||
param_count = next_param - 1
|
||||
elif tags_match == "exact":
|
||||
# Exact match with no tags is the "global" scope: rows that carry no
|
||||
# tags at all. (Other match modes treat empty tags as "no filter".)
|
||||
query_conditions.append("(tags IS NULL OR tags = '{}')")
|
||||
|
||||
if created_before is not None:
|
||||
param_count += 1
|
||||
query_conditions.append(f"created_at < ${param_count}")
|
||||
query_params.append(created_before)
|
||||
|
||||
where_clause = "WHERE " + " AND ".join(query_conditions) if query_conditions else ""
|
||||
|
||||
# Get total count
|
||||
count_query = f"""
|
||||
SELECT COUNT(*) as total
|
||||
FROM {source_table}
|
||||
{where_clause}
|
||||
"""
|
||||
count_result = await conn.fetchrow(count_query, *query_params)
|
||||
total = count_result["total"]
|
||||
|
||||
# Get units with limit and offset
|
||||
param_count += 1
|
||||
limit_param = f"${param_count}"
|
||||
query_params.append(limit)
|
||||
|
||||
param_count += 1
|
||||
offset_param = f"${param_count}"
|
||||
query_params.append(offset)
|
||||
|
||||
# The archive carries invalidation bookkeeping; the live table doesn't.
|
||||
curation_cols = (
|
||||
"invalidation_reason, invalidated_at"
|
||||
if is_archived
|
||||
else "NULL::text AS invalidation_reason, NULL::timestamptz AS invalidated_at"
|
||||
)
|
||||
units = await conn.fetch(
|
||||
f"""
|
||||
SELECT id, text, event_date, context, fact_type, document_id,
|
||||
mentioned_at, occurred_start, occurred_end, chunk_id, proof_count,
|
||||
tags, metadata, consolidated_at, consolidation_failed_at, edited_at, {curation_cols}
|
||||
FROM {source_table}
|
||||
{where_clause}
|
||||
ORDER BY mentioned_at DESC NULLS LAST, created_at DESC
|
||||
LIMIT {limit_param} OFFSET {offset_param}
|
||||
""",
|
||||
*query_params,
|
||||
)
|
||||
|
||||
# Get entity information for these units
|
||||
if units:
|
||||
unit_ids = [row["id"] for row in units]
|
||||
unit_entities = await conn.fetch(
|
||||
f"""
|
||||
SELECT ue.unit_id, e.canonical_name
|
||||
FROM {fq_table("unit_entities")} ue
|
||||
JOIN {fq_table("entities")} e ON ue.entity_id = e.id
|
||||
WHERE ue.unit_id = ANY($1::uuid[])
|
||||
ORDER BY ue.unit_id
|
||||
""",
|
||||
unit_ids,
|
||||
)
|
||||
else:
|
||||
unit_entities = []
|
||||
|
||||
# Build entity mapping
|
||||
entity_map: dict[Any, list[str]] = {}
|
||||
for row in unit_entities:
|
||||
unit_id = row["unit_id"]
|
||||
entity_name = row["canonical_name"]
|
||||
if unit_id not in entity_map:
|
||||
entity_map[unit_id] = []
|
||||
entity_map[unit_id].append(entity_name)
|
||||
|
||||
# Build result items
|
||||
items = []
|
||||
for row in units:
|
||||
unit_id = row["id"]
|
||||
entities = entity_map.get(unit_id, [])
|
||||
|
||||
items.append(
|
||||
{
|
||||
"id": str(unit_id),
|
||||
"text": row["text"],
|
||||
"context": row["context"] if row["context"] else "",
|
||||
"date": row["event_date"].isoformat() if row["event_date"] else "",
|
||||
"fact_type": row["fact_type"],
|
||||
"document_id": row["document_id"],
|
||||
"mentioned_at": row["mentioned_at"].isoformat() if row["mentioned_at"] else None,
|
||||
"occurred_start": row["occurred_start"].isoformat() if row["occurred_start"] else None,
|
||||
"occurred_end": row["occurred_end"].isoformat() if row["occurred_end"] else None,
|
||||
"entities": ", ".join(entities) if entities else "",
|
||||
"chunk_id": row["chunk_id"] if row["chunk_id"] else None,
|
||||
"proof_count": row["proof_count"] if row["proof_count"] is not None else 1,
|
||||
"tags": list(row["tags"]) if row["tags"] else [],
|
||||
"metadata": conn.parse_json(row["metadata"]) if row["metadata"] is not None else {},
|
||||
"consolidated_at": row["consolidated_at"].isoformat() if row["consolidated_at"] else None,
|
||||
"consolidation_failed_at": (
|
||||
row["consolidation_failed_at"].isoformat() if row["consolidation_failed_at"] else None
|
||||
),
|
||||
"state": "invalidated" if is_archived else "valid",
|
||||
"invalidation_reason": row["invalidation_reason"],
|
||||
"invalidated_at": row["invalidated_at"].isoformat() if row["invalidated_at"] else None,
|
||||
"edited_at": row["edited_at"].isoformat() if row["edited_at"] else None,
|
||||
}
|
||||
)
|
||||
|
||||
return {"items": items, "total": total, "limit": limit, "offset": offset}
|
||||
|
||||
|
||||
async def get_memory_unit(*, conn, ops, fq_table, bank_id: str, unit_id: str) -> dict[str, Any] | None:
|
||||
"""
|
||||
Get a single memory unit by ID.
|
||||
|
||||
Args:
|
||||
conn: Open database connection (the caller owns the transaction).
|
||||
ops: Dialect ops, for the observation→source entity inheritance shape.
|
||||
fq_table: Table-name resolver.
|
||||
bank_id: Bank ID
|
||||
unit_id: Memory unit ID (the caller validates it is a UUID)
|
||||
|
||||
Returns:
|
||||
Dict with memory unit data or None if not found
|
||||
"""
|
||||
# Get the memory unit (include source_memory_ids for mental models).
|
||||
# Curation moves invalidated facts to invalidated_memory_units, so fall
|
||||
# back to the archive (with its invalidation bookkeeping) on a miss.
|
||||
select_cols = (
|
||||
"id, text, context, event_date, occurred_start, occurred_end, "
|
||||
"mentioned_at, fact_type, document_id, chunk_id, tags, metadata, source_memory_ids, "
|
||||
"observation_scopes, edited_at"
|
||||
)
|
||||
row = await conn.fetchrow(
|
||||
f"SELECT {select_cols}, NULL::text AS invalidation_reason, NULL::timestamptz AS invalidated_at "
|
||||
f"FROM {fq_table('memory_units')} WHERE id = $1 AND bank_id = $2",
|
||||
unit_id,
|
||||
bank_id,
|
||||
)
|
||||
unit_state = "valid"
|
||||
if not row:
|
||||
row = await conn.fetchrow(
|
||||
f"SELECT {select_cols}, invalidation_reason, invalidated_at "
|
||||
f"FROM {fq_table('invalidated_memory_units')} WHERE id = $1 AND bank_id = $2",
|
||||
unit_id,
|
||||
bank_id,
|
||||
)
|
||||
unit_state = "invalidated"
|
||||
|
||||
if not row:
|
||||
return None
|
||||
|
||||
# Get entity information. _entity_rows_for_units_sql handles the
|
||||
# observation→source_memory_ids inheritance fallback in SQL, so a
|
||||
# single query covers direct rows and inherited ones.
|
||||
entities_rows = await conn.fetch(
|
||||
_entity_rows_for_units_sql(ops=ops, fq_table=fq_table, unit_ids_placeholder=1),
|
||||
[row["id"]],
|
||||
)
|
||||
entities = [r["canonical_name"] for r in entities_rows]
|
||||
|
||||
result: dict[str, Any] = {
|
||||
"id": str(row["id"]),
|
||||
"text": row["text"],
|
||||
"context": row["context"] if row["context"] else "",
|
||||
"date": row["event_date"].isoformat() if row["event_date"] else "",
|
||||
"type": row["fact_type"],
|
||||
"mentioned_at": row["mentioned_at"].isoformat() if row["mentioned_at"] else None,
|
||||
"occurred_start": row["occurred_start"].isoformat() if row["occurred_start"] else None,
|
||||
"occurred_end": row["occurred_end"].isoformat() if row["occurred_end"] else None,
|
||||
"entities": entities,
|
||||
"document_id": row["document_id"] if row["document_id"] else None,
|
||||
"chunk_id": str(row["chunk_id"]) if row["chunk_id"] else None,
|
||||
"tags": row["tags"] if row["tags"] else [],
|
||||
"metadata": conn.parse_json(row["metadata"]) if row["metadata"] is not None else {},
|
||||
"observation_scopes": (
|
||||
conn.parse_json(row["observation_scopes"]) if row["observation_scopes"] is not None else None
|
||||
),
|
||||
"state": unit_state,
|
||||
"invalidation_reason": row["invalidation_reason"],
|
||||
"invalidated_at": row["invalidated_at"].isoformat() if row["invalidated_at"] else None,
|
||||
"edited_at": row["edited_at"].isoformat() if row["edited_at"] else None,
|
||||
}
|
||||
|
||||
# For observations, include source_memory_ids
|
||||
# history is deprecated here - use GET /memories/{id}/history instead
|
||||
if row["fact_type"] == "observation":
|
||||
result["history"] = []
|
||||
|
||||
if row["fact_type"] == "observation" and row["source_memory_ids"]:
|
||||
source_ids = row["source_memory_ids"]
|
||||
result["source_memory_ids"] = [str(sid) for sid in source_ids]
|
||||
|
||||
# Fetch source memories
|
||||
source_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT id, text, fact_type, context, occurred_start, mentioned_at
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE id = ANY($1::uuid[])
|
||||
ORDER BY mentioned_at DESC NULLS LAST
|
||||
""",
|
||||
source_ids,
|
||||
)
|
||||
result["source_memories"] = [
|
||||
{
|
||||
"id": str(r["id"]),
|
||||
"text": r["text"],
|
||||
"type": r["fact_type"],
|
||||
"context": r["context"],
|
||||
"occurred_start": r["occurred_start"].isoformat() if r["occurred_start"] else None,
|
||||
"mentioned_at": r["mentioned_at"].isoformat() if r["mentioned_at"] else None,
|
||||
}
|
||||
for r in source_rows
|
||||
]
|
||||
|
||||
return result
|
||||
|
||||
|
||||
async def list_entities(
|
||||
*,
|
||||
conn,
|
||||
fq_table,
|
||||
bank_id: str,
|
||||
search: str | None = None,
|
||||
limit: int = 100,
|
||||
offset: int = 0,
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
List all entities for a bank with pagination.
|
||||
|
||||
Args:
|
||||
conn: Open database connection (the caller owns the transaction).
|
||||
fq_table: Table-name resolver.
|
||||
bank_id: bank IDentifier
|
||||
search: Optional case-insensitive substring match on canonical_name.
|
||||
limit: Maximum number of entities to return
|
||||
offset: Offset for pagination
|
||||
|
||||
Returns:
|
||||
Dict with items, total, limit, offset
|
||||
"""
|
||||
conditions = ["bank_id = $1"]
|
||||
params: list[Any] = [bank_id]
|
||||
if search:
|
||||
# Substring match, same ILIKE shape entity lookup uses elsewhere. Applied
|
||||
# to the count too, so the UI pages over the filtered set.
|
||||
params.append(f"%{search}%")
|
||||
conditions.append(f"canonical_name ILIKE ${len(params)}")
|
||||
where_clause = " AND ".join(conditions)
|
||||
|
||||
# Get total count
|
||||
total_row = await conn.fetchrow(
|
||||
f"""
|
||||
SELECT COUNT(*) as total
|
||||
FROM {fq_table("entities")}
|
||||
WHERE {where_clause}
|
||||
""",
|
||||
*params,
|
||||
)
|
||||
total = total_row["total"] if total_row else 0
|
||||
|
||||
# Get paginated entities
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT id, canonical_name, mention_count, first_seen, last_seen, metadata
|
||||
FROM {fq_table("entities")}
|
||||
WHERE {where_clause}
|
||||
ORDER BY mention_count DESC, last_seen DESC, id ASC
|
||||
LIMIT ${len(params) + 1} OFFSET ${len(params) + 2}
|
||||
""",
|
||||
*params,
|
||||
limit,
|
||||
offset,
|
||||
)
|
||||
|
||||
entities = []
|
||||
for row in rows:
|
||||
# Handle metadata - may be dict, JSON string, or None
|
||||
metadata = row["metadata"]
|
||||
if metadata is None:
|
||||
metadata = {}
|
||||
elif isinstance(metadata, str):
|
||||
try:
|
||||
metadata = json.loads(metadata)
|
||||
except json.JSONDecodeError:
|
||||
metadata = {}
|
||||
|
||||
entities.append(
|
||||
{
|
||||
"id": str(row["id"]),
|
||||
"canonical_name": row["canonical_name"],
|
||||
"mention_count": row["mention_count"],
|
||||
"first_seen": row["first_seen"].isoformat() if row["first_seen"] else None,
|
||||
"last_seen": row["last_seen"].isoformat() if row["last_seen"] else None,
|
||||
"metadata": metadata,
|
||||
}
|
||||
)
|
||||
return {
|
||||
"items": entities,
|
||||
"total": total,
|
||||
"limit": limit,
|
||||
"offset": offset,
|
||||
}
|
||||
|
||||
|
||||
__all__ = ["get_memory_unit", "list_entities", "list_memory_units"]
|
||||
@@ -1,779 +0,0 @@
|
||||
"""Graph-shaped reads and the link-maintenance passes, in SQL.
|
||||
|
||||
Everything here is a query over the *joins* around `memory_units` rather than
|
||||
over the memories themselves: `unit_entities` (which entities a memory mentions)
|
||||
and `memory_links` (memory-to-memory temporal/semantic/causal edges).
|
||||
|
||||
Two groups of callers:
|
||||
|
||||
* **The graph view.** :func:`graph_units`, :func:`graph_entity_rows` and
|
||||
:func:`graph_direct_links` return raw rows; the engine still owns the
|
||||
filtering, the observation inheritance, the derived entity edges, the
|
||||
colouring and the response assembly. These functions answer only "which
|
||||
memories", "which entity postings" and "which stored edges".
|
||||
* **The graph-maintenance job.** :func:`enqueue_relink_victims` runs inside the
|
||||
delete transaction; :func:`relink_pass`, :func:`prune_orphan_entities` and
|
||||
:func:`prune_stale_cooccurrences` are the three reconciliation passes the job
|
||||
drives. The job keeps the orchestration (pass ordering, the deadlock retry
|
||||
around the sweeps, the timing log); each function here does the pass's work.
|
||||
|
||||
:func:`entity_memory_counts` and :func:`entities_for_units` are the two entity
|
||||
postings reads that are not part of the graph view but read the same join table.
|
||||
|
||||
A store whose links travel inside the memory has nothing to relink and no join
|
||||
table to sweep, which is why these are methods on the interface at all: it
|
||||
answers them with zeroes rather than with SQL.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import uuid as uuid_module
|
||||
from collections.abc import Callable
|
||||
from typing import Any
|
||||
|
||||
from ....config import get_config
|
||||
from ...db.base import DatabaseConnection
|
||||
from ...retain.link_utils import (
|
||||
MAX_TEMPORAL_LINKS_PER_UNIT,
|
||||
_bulk_insert_links,
|
||||
_normalize_datetime,
|
||||
compute_semantic_links_ann,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Mirrors the ``top_k`` default in ``compute_semantic_links_ann`` at retain
|
||||
# time. If you change one, change the other — otherwise victims would either
|
||||
# never reach the cap (probe returns less than the cap) or stay perpetually
|
||||
# under it (cap is higher than retain creates).
|
||||
MAX_SEMANTIC_LINKS_PER_UNIT = 50
|
||||
|
||||
# Worker fetches this many rows per relink-loop iteration. Bounds
|
||||
# per-iteration probe/insert latency so a 10k-row backlog doesn't hold a
|
||||
# worker slot for minutes. Chosen so the typical iteration runs in well
|
||||
# under 1s.
|
||||
_DRAIN_BATCH_SIZE = 50
|
||||
|
||||
# Defensive guard against runaway relink loops — at _DRAIN_BATCH_SIZE units per
|
||||
# iteration that's 500k targets, far beyond any realistic single-bank backlog.
|
||||
_RELINK_ITERATION_CAP = 10000
|
||||
|
||||
# Cap at 10k edges — the UI can't usefully render more, and uncapped queries
|
||||
# on highly-connected graphs (e.g. 1000 nodes with 500k+ edges) are too slow.
|
||||
_GRAPH_MAX_EDGES = 10000
|
||||
|
||||
# Columns the graph view renders: nodes take id/text/date/context/entities,
|
||||
# the table rows take the rest, and `source_memory_ids` is what lets the caller
|
||||
# inherit an observation's links and entities from the facts behind it.
|
||||
_GRAPH_UNIT_COLUMNS = (
|
||||
"id, text, event_date, context, occurred_start, occurred_end, mentioned_at, "
|
||||
"document_id, chunk_id, fact_type, tags, created_at, proof_count, source_memory_ids"
|
||||
)
|
||||
|
||||
|
||||
def _ops_for(conn: DatabaseConnection) -> Any:
|
||||
"""The ``DataAccessOps`` matching the connection's SQL dialect.
|
||||
|
||||
This is the SQL memories store, and SQL means Postgres *or* Oracle — the two
|
||||
speak different dialects (Oracle inherits entity links through the
|
||||
``observation_sources`` junction, Postgres through ``source_memory_ids``
|
||||
arrays), so the ops must follow the connection rather than assume Postgres.
|
||||
The ops go by ``conn.backend_type`` — the connection objects carry the dialect
|
||||
but not the backend's ``ops`` handle, so resolve through the per-dialect cache
|
||||
of ``create_data_access_ops`` (a dict lookup after the first call, and the same
|
||||
instance the backend holds). The default covers callers that hand in a bare
|
||||
asyncpg connection with no dialect to report.
|
||||
"""
|
||||
from ...db import create_data_access_ops
|
||||
|
||||
return create_data_access_ops(getattr(conn, "backend_type", "postgresql"))
|
||||
|
||||
|
||||
def _as_uuids(unit_ids: list) -> list:
|
||||
"""Coerce a mixed list of uuid strings / UUIDs to UUIDs for a ``uuid[]`` bind."""
|
||||
return [uuid_module.UUID(uid) if isinstance(uid, str) else uid for uid in unit_ids]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------- graph view
|
||||
|
||||
|
||||
def _observations_via_source_match(
|
||||
fq_table: Callable[[str], str],
|
||||
ops: Any,
|
||||
source_column: str,
|
||||
source_placeholder: int,
|
||||
bank_placeholder: int | None,
|
||||
) -> str:
|
||||
"""A predicate matching observations whose *sources* satisfy ``<col> = $n``.
|
||||
|
||||
Observations carry no `document_id` / `chunk_id` of their own; the link to a
|
||||
source row lives in `source_memory_ids` (native array) or the
|
||||
`observation_sources` junction, depending on the dialect.
|
||||
"""
|
||||
if ops.uses_observation_sources_table:
|
||||
bank_clause = f" AND src.bank_id = ${bank_placeholder}" if bank_placeholder else ""
|
||||
return (
|
||||
f"id IN (SELECT os.observation_id "
|
||||
f"FROM {fq_table('observation_sources')} os "
|
||||
f"JOIN {fq_table('memory_units')} src ON src.id = os.source_id "
|
||||
f"WHERE src.{source_column} = ${source_placeholder}{bank_clause})"
|
||||
)
|
||||
bank_clause = f" AND bank_id = ${bank_placeholder}" if bank_placeholder else ""
|
||||
return (
|
||||
f"source_memory_ids && (SELECT array_agg(id) "
|
||||
f"FROM {fq_table('memory_units')} "
|
||||
f"WHERE {source_column} = ${source_placeholder}{bank_clause})"
|
||||
)
|
||||
|
||||
|
||||
async def graph_units(
|
||||
*,
|
||||
conn: DatabaseConnection,
|
||||
fq_table: Callable[[str], str],
|
||||
bank_id: str | None = None,
|
||||
fact_type: str | None = None,
|
||||
search_query: str | None = None,
|
||||
document_id: str | None = None,
|
||||
chunk_id: str | None = None,
|
||||
tags: list[str] | None = None,
|
||||
tags_match: str = "all_strict",
|
||||
limit: int = 1000,
|
||||
) -> dict[str, Any]:
|
||||
"""Memory nodes for the graph view, plus the total matching count.
|
||||
|
||||
Returns ``{"units": [...], "total": int}``: ``units`` is the page (newest
|
||||
first, capped at ``limit``); ``total`` is how many match the filters, which
|
||||
the UI shows alongside the page. ``document_id`` / ``chunk_id`` also match an
|
||||
observation whose *sources* carry them, since observations have neither of
|
||||
their own.
|
||||
"""
|
||||
from ...search.tags import build_tags_where_clause_simple
|
||||
|
||||
ops = _ops_for(conn)
|
||||
conditions: list[str] = []
|
||||
params: list[Any] = []
|
||||
|
||||
bank_placeholder: int | None = None
|
||||
if bank_id:
|
||||
params.append(bank_id)
|
||||
bank_placeholder = len(params)
|
||||
conditions.append(f"bank_id = ${bank_placeholder}")
|
||||
|
||||
if fact_type:
|
||||
params.append(fact_type)
|
||||
conditions.append(f"fact_type = ${len(params)}")
|
||||
|
||||
if document_id:
|
||||
params.append(document_id)
|
||||
obs = _observations_via_source_match(fq_table, ops, "document_id", len(params), bank_placeholder)
|
||||
conditions.append(f"(document_id = ${len(params)} OR (fact_type = 'observation' AND {obs}))")
|
||||
|
||||
if chunk_id:
|
||||
params.append(chunk_id)
|
||||
obs = _observations_via_source_match(fq_table, ops, "chunk_id", len(params), bank_placeholder)
|
||||
conditions.append(f"(chunk_id = ${len(params)} OR (fact_type = 'observation' AND {obs}))")
|
||||
|
||||
if search_query:
|
||||
params.append(f"%{search_query}%")
|
||||
conditions.append(f"(text ILIKE ${len(params)} OR context ILIKE ${len(params)})")
|
||||
|
||||
if tags:
|
||||
tag_clause = build_tags_where_clause_simple(tags, len(params) + 1, match=tags_match)
|
||||
if tag_clause:
|
||||
conditions.append(tag_clause.removeprefix("AND "))
|
||||
params.append(tags)
|
||||
elif tags_match == "exact":
|
||||
# Exact match with no tags is the "global" scope: rows carrying no tags at
|
||||
# all. (Other modes treat empty tags as "no filter".)
|
||||
conditions.append("(tags IS NULL OR tags = '{}')")
|
||||
|
||||
where_clause = "WHERE " + " AND ".join(conditions) if conditions else ""
|
||||
|
||||
total_row = await conn.fetchrow(
|
||||
f"SELECT COUNT(*) AS total FROM {fq_table('memory_units')} {where_clause}",
|
||||
*params,
|
||||
)
|
||||
total = total_row["total"] if total_row else 0
|
||||
|
||||
params.append(limit)
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT {_GRAPH_UNIT_COLUMNS}
|
||||
FROM {fq_table("memory_units")}
|
||||
{where_clause}
|
||||
ORDER BY mentioned_at DESC NULLS LAST, event_date DESC
|
||||
LIMIT ${len(params)}
|
||||
""",
|
||||
*params,
|
||||
)
|
||||
return {"units": [dict(row) for row in rows], "total": total}
|
||||
|
||||
|
||||
async def graph_entity_rows(
|
||||
*,
|
||||
conn: DatabaseConnection,
|
||||
fq_table: Callable[[str], str],
|
||||
bank_id: str,
|
||||
unit_ids: list[str],
|
||||
) -> list[dict[str, Any]]:
|
||||
"""``(unit_id, entity_id, canonical_name)`` rows for the graph view's entity edges.
|
||||
|
||||
Direct `unit_entities` postings only. An observation's entities are inherited
|
||||
from its source memories by the caller, which is why the ids it passes here
|
||||
are the visible units *plus* their source memories.
|
||||
|
||||
Scoped by unit id rather than by bank: the ids already came from a
|
||||
bank-scoped :func:`graph_units`, and `unit_entities` carries no bank column.
|
||||
"""
|
||||
if not unit_ids:
|
||||
return []
|
||||
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT ue.unit_id, e.id AS entity_id, e.canonical_name
|
||||
FROM {fq_table("unit_entities")} ue
|
||||
JOIN {fq_table("entities")} e ON ue.entity_id = e.id
|
||||
WHERE ue.unit_id = ANY($1::uuid[])
|
||||
ORDER BY ue.unit_id
|
||||
""",
|
||||
_as_uuids(unit_ids),
|
||||
)
|
||||
return [dict(row) for row in rows]
|
||||
|
||||
|
||||
async def graph_direct_links(
|
||||
*,
|
||||
conn: DatabaseConnection,
|
||||
fq_table: Callable[[str], str],
|
||||
bank_id: str,
|
||||
unit_ids: list[str],
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Memory-to-memory edges with *both* endpoints in ``unit_ids``.
|
||||
|
||||
Entity edges are derived by the caller from `unit_entities` so we don't
|
||||
materialize them in `memory_links` anymore (dropped in migration
|
||||
e9b2c7d1f3a4) — no link_type filter is needed. ``entity_name`` is selected as
|
||||
NULL so the row shape matches the derived edges the caller mixes these with.
|
||||
|
||||
Pass the visible units *and* the source memories they inherit from: the
|
||||
caller copies a source memory's links onto the observations built on it.
|
||||
"""
|
||||
if not unit_ids:
|
||||
return []
|
||||
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT ml.from_unit_id,
|
||||
ml.to_unit_id,
|
||||
ml.link_type,
|
||||
ml.weight,
|
||||
NULL::text AS entity_name
|
||||
FROM {fq_table("memory_links")} ml
|
||||
WHERE ml.from_unit_id = ANY($1::uuid[])
|
||||
AND ml.to_unit_id = ANY($1::uuid[])
|
||||
ORDER BY ml.weight DESC NULLS LAST
|
||||
LIMIT $2
|
||||
""",
|
||||
_as_uuids(unit_ids),
|
||||
_GRAPH_MAX_EDGES,
|
||||
)
|
||||
return [dict(row) for row in rows]
|
||||
|
||||
|
||||
# ------------------------------------------------------------ entity postings
|
||||
|
||||
|
||||
async def entity_memory_counts(
|
||||
*,
|
||||
conn: DatabaseConnection,
|
||||
fq_table: Callable[[str], str],
|
||||
bank_id: str,
|
||||
entity_ids: list[str] | None = None,
|
||||
) -> dict[str, int]:
|
||||
"""Live memory count per entity id, for the entities in ``bank_id``.
|
||||
|
||||
The GROUP BY is what makes this an orphan test: an entity with no surviving
|
||||
`unit_entities` row produces no group, so it is simply absent from the
|
||||
result rather than present with a zero.
|
||||
|
||||
Scoped through ``memory_units.bank_id`` — `unit_entities` has no bank column,
|
||||
and joining is what keeps the count to *live* memories (deleted units take
|
||||
their postings with them via ON DELETE CASCADE).
|
||||
"""
|
||||
params: list[Any] = [bank_id]
|
||||
entity_filter = ""
|
||||
if entity_ids is not None:
|
||||
if not entity_ids:
|
||||
return {}
|
||||
params.append(_as_uuids(entity_ids))
|
||||
entity_filter = f"AND ue.entity_id = ANY(${len(params)}::uuid[])"
|
||||
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT ue.entity_id, COUNT(*) AS n
|
||||
FROM {fq_table("unit_entities")} ue
|
||||
JOIN {fq_table("memory_units")} mu ON mu.id = ue.unit_id
|
||||
WHERE mu.bank_id = $1
|
||||
{entity_filter}
|
||||
GROUP BY ue.entity_id
|
||||
""",
|
||||
*params,
|
||||
)
|
||||
return {str(row["entity_id"]): int(row["n"]) for row in rows}
|
||||
|
||||
|
||||
def _entity_rows_for_units_sql(
|
||||
fq_table: Callable[[str], str],
|
||||
ops: Any,
|
||||
unit_ids_placeholder: int,
|
||||
) -> str:
|
||||
"""SQL SELECT producing ``(unit_id, entity_id, canonical_name)`` rows for
|
||||
the given unit IDs.
|
||||
|
||||
Direct rows come from ``unit_entities``. Observations rarely carry
|
||||
direct rows there; their entity association lives transitively through
|
||||
their source memories (``source_memory_ids`` on PG, the
|
||||
``observation_sources`` junction on Oracle). When an observation has
|
||||
no direct entity rows the SELECT inherits its source memories'
|
||||
entities, so the result is the same set callers would get from
|
||||
``get_memory_unit``.
|
||||
|
||||
``unit_ids_placeholder`` is the 1-based parameter index that holds the
|
||||
``uuid[]`` of unit IDs. The placeholder is referenced twice — both
|
||||
sides of the UNION need it — so callers should not reuse the slot.
|
||||
"""
|
||||
ue = fq_table("unit_entities")
|
||||
ents = fq_table("entities")
|
||||
mu = fq_table("memory_units")
|
||||
p = unit_ids_placeholder
|
||||
|
||||
direct = (
|
||||
f"SELECT ue.unit_id, e.id AS entity_id, e.canonical_name "
|
||||
f"FROM {ue} ue "
|
||||
f"JOIN {ents} e ON e.id = ue.entity_id "
|
||||
f"WHERE ue.unit_id = ANY(${p}::uuid[])"
|
||||
)
|
||||
|
||||
if ops.uses_observation_sources_table:
|
||||
os_t = fq_table("observation_sources")
|
||||
inherited = (
|
||||
f"SELECT os.observation_id AS unit_id, e.id AS entity_id, e.canonical_name "
|
||||
f"FROM {os_t} os "
|
||||
f"JOIN {ue} src_ue ON src_ue.unit_id = os.source_id "
|
||||
f"JOIN {ents} e ON e.id = src_ue.entity_id "
|
||||
f"WHERE os.observation_id = ANY(${p}::uuid[]) "
|
||||
f"AND NOT EXISTS (SELECT 1 FROM {ue} d WHERE d.unit_id = os.observation_id)"
|
||||
)
|
||||
else:
|
||||
inherited = (
|
||||
f"SELECT obs.id AS unit_id, e.id AS entity_id, e.canonical_name "
|
||||
f"FROM {mu} obs "
|
||||
f"CROSS JOIN LATERAL unnest(obs.source_memory_ids) AS src_id "
|
||||
f"JOIN {ue} src_ue ON src_ue.unit_id = src_id "
|
||||
f"JOIN {ents} e ON e.id = src_ue.entity_id "
|
||||
f"WHERE obs.id = ANY(${p}::uuid[]) "
|
||||
f"AND obs.fact_type = 'observation' "
|
||||
f"AND obs.source_memory_ids IS NOT NULL "
|
||||
f"AND NOT EXISTS (SELECT 1 FROM {ue} d WHERE d.unit_id = obs.id)"
|
||||
)
|
||||
|
||||
return f"({direct}) UNION ({inherited})"
|
||||
|
||||
|
||||
async def entities_for_units(
|
||||
*,
|
||||
conn: DatabaseConnection,
|
||||
fq_table: Callable[[str], str],
|
||||
bank_id: str,
|
||||
unit_ids: list[str],
|
||||
) -> dict[str, list[str]]:
|
||||
"""The entity ids each unit carries, keyed by unit id.
|
||||
|
||||
Observations inherit their source memories' entities when they carry no
|
||||
direct postings of their own — see :func:`_entity_rows_for_units_sql`. Units
|
||||
with no entities are absent rather than mapped to an empty list.
|
||||
"""
|
||||
if not unit_ids:
|
||||
return {}
|
||||
|
||||
rows = await conn.fetch(
|
||||
_entity_rows_for_units_sql(fq_table, _ops_for(conn), unit_ids_placeholder=1),
|
||||
_as_uuids(unit_ids),
|
||||
)
|
||||
|
||||
# UNION already de-duplicates whole rows, but a unit can reach the same
|
||||
# entity through more than one source memory, so dedupe per unit while
|
||||
# preserving the order the rows arrived in.
|
||||
by_unit: dict[str, list[str]] = {}
|
||||
for row in rows:
|
||||
unit_key = str(row["unit_id"])
|
||||
entity_id = str(row["entity_id"])
|
||||
bucket = by_unit.setdefault(unit_key, [])
|
||||
if entity_id not in bucket:
|
||||
bucket.append(entity_id)
|
||||
return by_unit
|
||||
|
||||
|
||||
async def entity_map_for_units(
|
||||
*,
|
||||
conn: DatabaseConnection,
|
||||
fq_table: Callable[[str], str],
|
||||
bank_id: str,
|
||||
unit_ids: list[str],
|
||||
) -> dict[str, list[dict[str, str]]]:
|
||||
"""``{unit_id: [{entity_id, canonical_name}]}`` — the recall/curation shape.
|
||||
|
||||
The named twin of :func:`entities_for_units`: recall renders the entity name
|
||||
on each fact, so it needs the label, not just the id. Observation-via-source
|
||||
inheritance and the per-unit dedupe are identical.
|
||||
"""
|
||||
if not unit_ids:
|
||||
return {}
|
||||
|
||||
rows = await conn.fetch(
|
||||
_entity_rows_for_units_sql(fq_table, _ops_for(conn), unit_ids_placeholder=1),
|
||||
_as_uuids(unit_ids),
|
||||
)
|
||||
by_unit: dict[str, list[dict[str, str]]] = {}
|
||||
for row in rows:
|
||||
unit_key = str(row["unit_id"])
|
||||
entity_id = str(row["entity_id"])
|
||||
bucket = by_unit.setdefault(unit_key, [])
|
||||
if not any(existing["entity_id"] == entity_id for existing in bucket):
|
||||
bucket.append({"entity_id": entity_id, "canonical_name": row["canonical_name"]})
|
||||
return by_unit
|
||||
|
||||
|
||||
# --------------------------------------------------------------- maintenance
|
||||
|
||||
|
||||
async def enqueue_relink_victims(
|
||||
*,
|
||||
conn: DatabaseConnection,
|
||||
fq_table: Callable[[str], str],
|
||||
bank_id: str,
|
||||
affected_unit_ids: list,
|
||||
include_affected_units: bool = False,
|
||||
) -> int:
|
||||
"""Enqueue surviving units whose outgoing temporal/semantic links pointed at
|
||||
``affected_unit_ids`` for later link top-up.
|
||||
|
||||
Must run inside the same transaction that drops those links, *before* the
|
||||
delete (or cascade) fires — once the rows are gone, the join that finds the
|
||||
victims returns nothing.
|
||||
|
||||
Args:
|
||||
conn: Database connection inside the active transaction.
|
||||
fq_table: Schema-qualifying table-name resolver.
|
||||
bank_id: Bank owning the affected units.
|
||||
affected_unit_ids: Memory_unit IDs whose incident temporal/semantic links
|
||||
are about to be (or are being) removed.
|
||||
include_affected_units: Also enqueue ``affected_unit_ids`` themselves — for
|
||||
an edit that deletes a unit's links but leaves the unit live, so its own
|
||||
outgoing adjacency is rebuilt too. One combined insert keeps the queue's
|
||||
sorted lock ordering intact.
|
||||
|
||||
Returns:
|
||||
Number of distinct victim units enqueued (after dedup against rows
|
||||
already in the queue).
|
||||
"""
|
||||
if not affected_unit_ids:
|
||||
return 0
|
||||
|
||||
ops = _ops_for(conn)
|
||||
affected_uuids = _as_uuids(affected_unit_ids)
|
||||
affected_str_set = {str(uid) for uid in affected_uuids}
|
||||
|
||||
# Find units (other than the affected ones) that have an outgoing
|
||||
# temporal/semantic link pointing at an affected unit. Entity links are
|
||||
# intentionally excluded — they're scheduled for removal and would only
|
||||
# add noise to the recompute job.
|
||||
victim_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT DISTINCT from_unit_id
|
||||
FROM {fq_table("memory_links")}
|
||||
WHERE to_unit_id = ANY($1::uuid[])
|
||||
AND bank_id = $2
|
||||
AND link_type IN ('temporal', 'semantic')
|
||||
""",
|
||||
affected_uuids,
|
||||
bank_id,
|
||||
)
|
||||
|
||||
victim_ids = {row["from_unit_id"] for row in victim_rows if str(row["from_unit_id"]) not in affected_str_set}
|
||||
if include_affected_units:
|
||||
victim_ids.update(affected_uuids)
|
||||
|
||||
if not victim_ids:
|
||||
return 0
|
||||
|
||||
await ops.enqueue_graph_maintenance(
|
||||
conn,
|
||||
fq_table("graph_maintenance_queue"),
|
||||
bank_id,
|
||||
list(victim_ids),
|
||||
)
|
||||
|
||||
logger.debug(
|
||||
f"[GRAPH_MAINT] Enqueued {len(victim_ids)} relink victims in "
|
||||
f"bank={bank_id} ({len(affected_unit_ids)} units affected)"
|
||||
)
|
||||
return len(victim_ids)
|
||||
|
||||
|
||||
async def relink_pass(
|
||||
*,
|
||||
backend: Any,
|
||||
fq_table: Callable[[str], str],
|
||||
bank_id: str,
|
||||
config: Any,
|
||||
) -> dict:
|
||||
"""Drain ``graph_maintenance_queue`` for ``bank_id``, topping up lost links.
|
||||
|
||||
Per-iteration loop: claim → top up → commit. We rely on submit-time
|
||||
dedup to keep at most one job per bank running, so no need for
|
||||
SKIP LOCKED.
|
||||
|
||||
Takes ``backend`` rather than a connection because the loop spans several
|
||||
transactions — one per claimed batch, plus a separate connection for the ANN
|
||||
probe — so it has to acquire its own.
|
||||
|
||||
``config`` is the caller's resolved configuration. The Postgres pass takes
|
||||
its caps from retain's link_utils (so relink and retain agree on what "full"
|
||||
means) and never reads it; it is accepted so a store that *does* tune its
|
||||
relinking gets it.
|
||||
|
||||
Returns:
|
||||
``{"relink_units_processed": int, "relink_links_added": int}``.
|
||||
"""
|
||||
del config # accepted for symmetry with stores that tune their own relinking
|
||||
ops = backend.ops
|
||||
|
||||
units_processed = 0
|
||||
links_added = 0
|
||||
iterations = 0
|
||||
while True:
|
||||
from ...memory_engine import acquire_with_retry
|
||||
|
||||
async with acquire_with_retry(backend) as conn:
|
||||
async with conn.transaction():
|
||||
unit_ids = await ops.claim_graph_maintenance_batch(
|
||||
conn,
|
||||
fq_table("graph_maintenance_queue"),
|
||||
bank_id,
|
||||
_DRAIN_BATCH_SIZE,
|
||||
)
|
||||
if not unit_ids:
|
||||
break
|
||||
|
||||
links_added += await _relink_batch(conn, fq_table, bank_id, unit_ids, ops, backend)
|
||||
|
||||
units_processed += len(unit_ids)
|
||||
iterations += 1
|
||||
|
||||
if iterations > _RELINK_ITERATION_CAP:
|
||||
# Defensive guard against runaway loops — at 50 units/iter that's
|
||||
# 500k targets, far beyond any realistic single-bank backlog.
|
||||
logger.error(
|
||||
f"[GRAPH_MAINT] bank={bank_id} hit iteration cap ({iterations}); aborting relink "
|
||||
f"(units_processed={units_processed}, links_added={links_added})"
|
||||
)
|
||||
break
|
||||
|
||||
return {"relink_units_processed": units_processed, "relink_links_added": links_added}
|
||||
|
||||
|
||||
async def _relink_batch(
|
||||
conn: DatabaseConnection,
|
||||
fq_table: Callable[[str], str],
|
||||
bank_id: str,
|
||||
victim_ids: list[str],
|
||||
ops: Any,
|
||||
backend: Any,
|
||||
) -> int:
|
||||
"""Top up temporal/semantic links for a batch of victim units. Returns rows inserted."""
|
||||
# Load each victim's metadata. Victims whose units were deleted between
|
||||
# enqueue and now silently drop out — exactly the no-op behaviour we want
|
||||
# for stale queue rows.
|
||||
victim_uuids = [uuid_module.UUID(vid) for vid in victim_ids]
|
||||
victim_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT id::text AS id, event_date, fact_type, embedding::text AS embedding
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE id = ANY($1::uuid[])
|
||||
AND bank_id = $2
|
||||
AND fact_type IN ('experience', 'world')
|
||||
""",
|
||||
victim_uuids,
|
||||
bank_id,
|
||||
)
|
||||
|
||||
if not victim_rows:
|
||||
return 0
|
||||
|
||||
alive_uuids = [uuid_module.UUID(row["id"]) for row in victim_rows]
|
||||
|
||||
# Count current outgoing temporal/semantic links per victim so we only
|
||||
# probe for the ones genuinely below cap. Saves the bulk of the work when
|
||||
# most victims still have plenty of links.
|
||||
count_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT from_unit_id, link_type, COUNT(*) AS cnt
|
||||
FROM {fq_table("memory_links")}
|
||||
WHERE from_unit_id = ANY($1::uuid[])
|
||||
AND bank_id = $2
|
||||
AND link_type IN ('temporal', 'semantic')
|
||||
GROUP BY from_unit_id, link_type
|
||||
""",
|
||||
alive_uuids,
|
||||
bank_id,
|
||||
)
|
||||
counts: dict[tuple[str, str], int] = {}
|
||||
for row in count_rows:
|
||||
counts[(str(row["from_unit_id"]), row["link_type"])] = int(row["cnt"])
|
||||
|
||||
# --- Temporal top-up ---
|
||||
temporal_needs = [r for r in victim_rows if counts.get((r["id"], "temporal"), 0) < MAX_TEMPORAL_LINKS_PER_UNIT]
|
||||
new_links: list[tuple] = []
|
||||
|
||||
if temporal_needs:
|
||||
lateral_unit_ids = [uuid_module.UUID(r["id"]) for r in temporal_needs if r["event_date"] is not None]
|
||||
lateral_event_dates = [
|
||||
_normalize_datetime(r["event_date"]) for r in temporal_needs if r["event_date"] is not None
|
||||
]
|
||||
lateral_fact_types = [r["fact_type"] for r in temporal_needs if r["event_date"] is not None]
|
||||
|
||||
if lateral_unit_ids:
|
||||
rows = await ops.fetch_temporal_neighbors(
|
||||
conn,
|
||||
fq_table("memory_units"),
|
||||
bank_id,
|
||||
lateral_unit_ids,
|
||||
lateral_event_dates,
|
||||
lateral_fact_types,
|
||||
MAX_TEMPORAL_LINKS_PER_UNIT,
|
||||
)
|
||||
for row in rows:
|
||||
time_diff_h = float(row["time_diff_hours"])
|
||||
# Mirror the 24h window enforced at retain time. The bidirectional
|
||||
# index scan returns the K closest neighbours regardless of
|
||||
# window, so we filter here.
|
||||
if time_diff_h > 24:
|
||||
continue
|
||||
weight = max(0.3, 1.0 - (time_diff_h / 24))
|
||||
new_links.append((row["from_id"], str(row["id"]), "temporal", weight, None))
|
||||
|
||||
# --- Semantic top-up ---
|
||||
# ANN must run on its own connection: it opens a nested transaction with
|
||||
# SET LOCAL hnsw.ef_search + CREATE TEMP TABLE ON COMMIT DROP, and nesting
|
||||
# that inside our current write transaction would commit our writes early.
|
||||
semantic_needs = [
|
||||
r
|
||||
for r in victim_rows
|
||||
if counts.get((r["id"], "semantic"), 0) < MAX_SEMANTIC_LINKS_PER_UNIT and r["embedding"] is not None
|
||||
]
|
||||
if semantic_needs:
|
||||
from ...memory_engine import acquire_with_retry
|
||||
|
||||
seed_ids = [r["id"] for r in semantic_needs]
|
||||
seed_embs = [r["embedding"] for r in semantic_needs]
|
||||
seed_ftypes = [r["fact_type"] for r in semantic_needs]
|
||||
async with acquire_with_retry(backend) as ann_conn:
|
||||
try:
|
||||
ann_links = await compute_semantic_links_ann(
|
||||
ann_conn,
|
||||
bank_id,
|
||||
seed_ids,
|
||||
seed_embs,
|
||||
fact_types=seed_ftypes,
|
||||
threshold=get_config().semantic_link_min_similarity,
|
||||
)
|
||||
# Strip self-links (rare but possible because the ANN probe
|
||||
# has no exclude list — see the comment in compute_semantic_links_ann).
|
||||
ann_links = [lnk for lnk in ann_links if lnk[0] != lnk[1]]
|
||||
new_links.extend(ann_links)
|
||||
except Exception as e:
|
||||
# ANN uses PG-specific HNSW syntax; on dialects/configs where
|
||||
# it isn't available we still want the temporal top-up to land.
|
||||
logger.warning(f"[GRAPH_MAINT] Semantic top-up failed for bank={bank_id}: {type(e).__name__}: {e}")
|
||||
|
||||
if not new_links:
|
||||
return 0
|
||||
|
||||
await _bulk_insert_links(
|
||||
conn,
|
||||
new_links,
|
||||
bank_id=bank_id,
|
||||
skip_exists_check=False,
|
||||
ops=ops,
|
||||
)
|
||||
return len(new_links)
|
||||
|
||||
|
||||
async def prune_orphan_entities(
|
||||
*,
|
||||
conn: DatabaseConnection,
|
||||
fq_table: Callable[[str], str],
|
||||
bank_id: str,
|
||||
) -> int:
|
||||
"""Delete ``entities`` rows in the bank with no remaining ``unit_entities``
|
||||
references. Returns the number pruned.
|
||||
|
||||
FK ON DELETE CASCADE on ``entity_cooccurrences`` then removes any
|
||||
cooccurrence row pointing at the pruned entities — which is why this runs
|
||||
before :func:`prune_stale_cooccurrences` rather than after.
|
||||
|
||||
A bank-wide single-statement delete, cheap when there's nothing to do. It is
|
||||
idempotent (rerunning only deletes what is still orphaned), so the caller is
|
||||
free to retry the whole transaction on deadlock.
|
||||
"""
|
||||
ops = _ops_for(conn)
|
||||
return await ops.prune_orphan_entities(
|
||||
conn,
|
||||
fq_table("entities"),
|
||||
fq_table("unit_entities"),
|
||||
bank_id,
|
||||
)
|
||||
|
||||
|
||||
async def prune_stale_cooccurrences(
|
||||
*,
|
||||
conn: DatabaseConnection,
|
||||
fq_table: Callable[[str], str],
|
||||
bank_id: str,
|
||||
) -> int:
|
||||
"""Delete cooccurrence rows no current memory witnesses. Returns the count.
|
||||
|
||||
Defensive sweep for rows where both endpoints still exist but no current
|
||||
memory_unit references both of them — the cooccurrence was real at the time
|
||||
it was recorded, but every unit that witnessed it has since been deleted.
|
||||
:func:`prune_orphan_entities` cascades the *missing-entity* case via FK; this
|
||||
pass catches the *stale-count* case it cannot see.
|
||||
|
||||
Like the orphan prune, a bank-wide idempotent sweep backed by indexes, so
|
||||
it's cheap when there's nothing to do and safe for the caller to retry.
|
||||
"""
|
||||
ops = _ops_for(conn)
|
||||
return await ops.prune_stale_cooccurrences(
|
||||
conn,
|
||||
fq_table("entity_cooccurrences"),
|
||||
fq_table("unit_entities"),
|
||||
fq_table("entities"),
|
||||
bank_id,
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"MAX_SEMANTIC_LINKS_PER_UNIT",
|
||||
"enqueue_relink_victims",
|
||||
"entities_for_units",
|
||||
"entity_map_for_units",
|
||||
"entity_memory_counts",
|
||||
"graph_direct_links",
|
||||
"graph_entity_rows",
|
||||
"graph_units",
|
||||
"prune_orphan_entities",
|
||||
"prune_stale_cooccurrences",
|
||||
"relink_pass",
|
||||
]
|
||||
@@ -1,539 +0,0 @@
|
||||
"""Addressed reads over `memory_units`: get, scan, count, tags, consolidation state.
|
||||
|
||||
Not retrieval — nothing here ranks. These are the queries behind the curation
|
||||
detail view, export, the bank-stats panel and the consolidation queue, lifted out
|
||||
of the call sites that used to issue them inline (``memory_engine``,
|
||||
``transfer/export``, ``consolidation/consolidator``) so
|
||||
:class:`~hindsight_api.engine.memories.postgres.PostgresMemories` can delegate
|
||||
rather than embed SQL.
|
||||
|
||||
Every function takes the live connection and Hindsight's ``fq_table`` resolver, so
|
||||
each one runs inside whatever transaction the caller already holds; none of them
|
||||
acquires a connection of its own.
|
||||
|
||||
**Cursor semantics.** ``scan_memories``'s ``page_token`` is opaque to callers, and
|
||||
for Postgres it is simply a *numeric offset rendered as a decimal string* against
|
||||
the scan's fixed ``ORDER BY created_at, id``. An empty token means "start at the
|
||||
beginning", and an empty token comes back once the walk is exhausted (i.e. the
|
||||
final short page). An offset cursor is a position rather than a snapshot — exactly
|
||||
the guarantee :class:`~hindsight_api.engine.memories.base.ScanPage` documents:
|
||||
rows written or deleted mid-walk can shift later pages, so a scan is
|
||||
eventually-complete browsing rather than a consistent iterator. ``skip`` is applied
|
||||
*on top of* the decoded cursor, so a caller that pages with both should pass
|
||||
``skip`` only on the first call — the returned token already accounts for it.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import uuid
|
||||
from collections.abc import Callable
|
||||
from datetime import datetime
|
||||
from typing import Any
|
||||
|
||||
from ...search.tags import (
|
||||
build_tag_groups_where_clause,
|
||||
build_tags_where_clause,
|
||||
build_tags_where_clause_simple,
|
||||
)
|
||||
from ..base import ScanPage, StoredMemory
|
||||
|
||||
# The `memory_units` projection every read here shares. Superset of the by-id
|
||||
# SELECT the recall source-facts path used (text/fact_type/context/timestamps/
|
||||
# document_id/chunk_id/tags/metadata), plus the observation bookkeeping columns
|
||||
# `StoredMemory` carries: source_memory_ids and consolidated_at.
|
||||
_MEMORY_COLUMNS = """
|
||||
id, text, fact_type, context, document_id, chunk_id, tags, metadata,
|
||||
proof_count, event_date, occurred_start, occurred_end, mentioned_at,
|
||||
created_at, source_memory_ids, consolidated_at, observation_scopes
|
||||
"""
|
||||
|
||||
# The scan's order. Fixed (created_at, id) like the export loader's, because an
|
||||
# offset cursor is only meaningful against a total order.
|
||||
_SCAN_ORDER = "ORDER BY created_at, id"
|
||||
|
||||
|
||||
def _as_json(value: Any) -> Any:
|
||||
"""Coerce an asyncpg JSONB column (str or already-decoded) to a Python object.
|
||||
|
||||
Connections differ in whether a JSONB codec is registered, so the column
|
||||
arrives either as text or as the decoded object.
|
||||
"""
|
||||
if value is None:
|
||||
return None
|
||||
if isinstance(value, str):
|
||||
try:
|
||||
return json.loads(value)
|
||||
except json.JSONDecodeError:
|
||||
# A valid scalar such as `"combined"` arrives already decoded on
|
||||
# connections that do register a decoder.
|
||||
return value
|
||||
return value
|
||||
|
||||
|
||||
def _as_uuids(unit_ids: list[Any]) -> list[uuid.UUID]:
|
||||
"""Unit ids as UUIDs, dropping anything unparseable.
|
||||
|
||||
A malformed id is treated the same way a deleted one is — simply absent from
|
||||
the result — rather than failing the whole read.
|
||||
"""
|
||||
out: list[uuid.UUID] = []
|
||||
for unit_id in unit_ids or []:
|
||||
if isinstance(unit_id, uuid.UUID):
|
||||
out.append(unit_id)
|
||||
continue
|
||||
try:
|
||||
out.append(uuid.UUID(str(unit_id)))
|
||||
except (ValueError, AttributeError, TypeError):
|
||||
continue
|
||||
return out
|
||||
|
||||
|
||||
def _column(row: Any, name: str, default: Any = None) -> Any:
|
||||
"""One column of an asyncpg Record, tolerating a narrower projection."""
|
||||
try:
|
||||
return row[name]
|
||||
except (KeyError, IndexError):
|
||||
return default
|
||||
|
||||
|
||||
def _stored_from_row(row: Any) -> StoredMemory:
|
||||
"""Map a `memory_units` row onto :class:`StoredMemory`.
|
||||
|
||||
Shared by every read in this module so the row → dataclass mapping exists
|
||||
once. ``entity_ids`` stays empty: the unit→entity posting lives in
|
||||
`unit_entities` and is served by ``entities_for_units``, not by a join here.
|
||||
"""
|
||||
source_ids = _column(row, "source_memory_ids") or []
|
||||
return StoredMemory(
|
||||
unit_id=str(row["id"]),
|
||||
text=row["text"],
|
||||
fact_type=row["fact_type"],
|
||||
context=_column(row, "context"),
|
||||
document_id=_column(row, "document_id"),
|
||||
chunk_id=str(_column(row, "chunk_id")) if _column(row, "chunk_id") else None,
|
||||
tags=list(_column(row, "tags") or []),
|
||||
metadata=_as_json(_column(row, "metadata")),
|
||||
proof_count=_column(row, "proof_count") or 1,
|
||||
event_date=_column(row, "event_date"),
|
||||
occurred_start=_column(row, "occurred_start"),
|
||||
occurred_end=_column(row, "occurred_end"),
|
||||
mentioned_at=_column(row, "mentioned_at"),
|
||||
created_at=_column(row, "created_at"),
|
||||
source_memory_ids=[str(sid) for sid in source_ids],
|
||||
consolidated_at=_column(row, "consolidated_at"),
|
||||
# Consolidation routes a candidate by its scopes, so this has to survive
|
||||
# the trip through the store rather than being re-queried per memory.
|
||||
observation_scopes=_as_json(_column(row, "observation_scopes")),
|
||||
)
|
||||
|
||||
|
||||
def _decode_page_token(page_token: str) -> int:
|
||||
"""Decode the offset cursor. Empty, malformed or negative all mean "start"."""
|
||||
if not page_token:
|
||||
return 0
|
||||
try:
|
||||
offset = int(page_token)
|
||||
except (TypeError, ValueError):
|
||||
return 0
|
||||
return offset if offset > 0 else 0
|
||||
|
||||
|
||||
async def get_memories(
|
||||
*, conn, fq_table: Callable[[str], str], bank_id: str, unit_ids: list[str]
|
||||
) -> list[StoredMemory]:
|
||||
"""Fetch memories by id. Missing or deleted ids are simply absent."""
|
||||
ids = _as_uuids(unit_ids)
|
||||
if not ids:
|
||||
return []
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT {_MEMORY_COLUMNS}
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE bank_id = $1 AND id = ANY($2::uuid[])
|
||||
""",
|
||||
bank_id,
|
||||
ids,
|
||||
)
|
||||
return [_stored_from_row(row) for row in rows]
|
||||
|
||||
|
||||
async def _semantic_edges(
|
||||
*, conn, fq_table: Callable[[str], str], bank_id: str, unit_ids: list[uuid.UUID]
|
||||
) -> dict[str, list[tuple[str, float]]]:
|
||||
"""Derived kNN edges for ``unit_ids``, keyed by unit id.
|
||||
|
||||
Walked in both directions, like the graph arm's semantic expansion: a
|
||||
`memory_links` row is written once, so a unit's neighbourhood is the union of
|
||||
the edges leaving it and those arriving at it.
|
||||
"""
|
||||
if not unit_ids:
|
||||
return {}
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT from_unit_id AS unit_id, to_unit_id AS target_id, weight
|
||||
FROM {fq_table("memory_links")}
|
||||
WHERE bank_id = $1 AND link_type = 'semantic' AND from_unit_id = ANY($2::uuid[])
|
||||
UNION ALL
|
||||
SELECT to_unit_id AS unit_id, from_unit_id AS target_id, weight
|
||||
FROM {fq_table("memory_links")}
|
||||
WHERE bank_id = $1 AND link_type = 'semantic' AND to_unit_id = ANY($2::uuid[])
|
||||
""",
|
||||
bank_id,
|
||||
unit_ids,
|
||||
)
|
||||
edges: dict[str, list[tuple[str, float]]] = {}
|
||||
for row in rows:
|
||||
edges.setdefault(str(row["unit_id"]), []).append((str(row["target_id"]), float(row["weight"] or 0.0)))
|
||||
return edges
|
||||
|
||||
|
||||
async def scan_memories(
|
||||
*,
|
||||
conn,
|
||||
fq_table: Callable[[str], str],
|
||||
bank_id: str,
|
||||
fact_types: list[str] | None = None,
|
||||
limit: int = 100,
|
||||
page_token: str = "",
|
||||
tags: list[str] | None = None,
|
||||
tags_match: str = "any",
|
||||
tag_groups: list | None = None,
|
||||
document_id: str | None = None,
|
||||
metadata_equals: dict[str, str] | None = None,
|
||||
skip: int = 0,
|
||||
include_edges: bool = False,
|
||||
) -> ScanPage:
|
||||
"""Page through stored memories. A full walk — for browsing and export only.
|
||||
|
||||
See the module docstring for the ``page_token`` (offset) cursor semantics.
|
||||
"""
|
||||
if limit is None or limit <= 0:
|
||||
return ScanPage()
|
||||
|
||||
where: list[str] = ["bank_id = $1"]
|
||||
params: list[Any] = [bank_id]
|
||||
|
||||
if fact_types:
|
||||
params.append(list(fact_types))
|
||||
where.append(f"fact_type = ANY(${len(params)})")
|
||||
|
||||
if document_id is not None:
|
||||
# A real column here, which is why it is not folded into
|
||||
# `metadata_equals`: only a store without the column keeps it in the bag.
|
||||
params.append(document_id)
|
||||
where.append(f"document_id = ${len(params)}")
|
||||
|
||||
if metadata_equals:
|
||||
# str→str equality across every key, which is exactly JSONB containment.
|
||||
params.append(json.dumps(metadata_equals))
|
||||
where.append(f"metadata @> ${len(params)}::jsonb")
|
||||
|
||||
# The tags clause owns its own `AND` prefix and, per the helper's contract,
|
||||
# only consumes a bind param when `tags` is non-empty (match="exact" with no
|
||||
# tags is the untagged/global scope and needs none).
|
||||
tags_clause = build_tags_where_clause_simple(tags, len(params) + 1, match=tags_match)
|
||||
if tags:
|
||||
params.append(list(tags))
|
||||
|
||||
# Compound tag groups (AND/OR/NOT trees), AND-ed on. Also owns its `AND` prefix and appends
|
||||
# one bind param per leaf; empty/absent groups yield no clause and no params.
|
||||
groups_clause, group_params, _ = build_tag_groups_where_clause(tag_groups, param_offset=len(params) + 1)
|
||||
params.extend(group_params)
|
||||
|
||||
offset = _decode_page_token(page_token) + max(int(skip or 0), 0)
|
||||
params.append(limit)
|
||||
limit_idx = len(params)
|
||||
params.append(offset)
|
||||
offset_idx = len(params)
|
||||
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT {_MEMORY_COLUMNS}
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE {" AND ".join(where)} {tags_clause} {groups_clause}
|
||||
{_SCAN_ORDER}
|
||||
LIMIT ${limit_idx} OFFSET ${offset_idx}
|
||||
""",
|
||||
*params,
|
||||
)
|
||||
|
||||
memories = [_stored_from_row(row) for row in rows]
|
||||
if include_edges and memories:
|
||||
edges = await _semantic_edges(
|
||||
conn=conn, fq_table=fq_table, bank_id=bank_id, unit_ids=_as_uuids([m.unit_id for m in memories])
|
||||
)
|
||||
for memory in memories:
|
||||
memory.semantic_edges = edges.get(memory.unit_id, [])
|
||||
|
||||
# A short page means the walk is exhausted, so the cursor goes empty.
|
||||
next_token = str(offset + len(rows)) if len(rows) == limit else ""
|
||||
return ScanPage(memories=memories, next_page_token=next_token)
|
||||
|
||||
|
||||
async def count_memories(*, conn, fq_table: Callable[[str], str], bank_id: str) -> dict[str, int]:
|
||||
"""Live memory count per fact_type. The bank-stats node counts."""
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT fact_type, COUNT(*) as count
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE bank_id = $1
|
||||
GROUP BY fact_type
|
||||
""",
|
||||
bank_id,
|
||||
)
|
||||
return {row["fact_type"]: int(row["count"]) for row in rows}
|
||||
|
||||
|
||||
async def list_tags(
|
||||
*,
|
||||
conn,
|
||||
fq_table: Callable[[str], str],
|
||||
bank_id: str,
|
||||
pattern: str | None = None,
|
||||
limit: int = 100,
|
||||
offset: int = 0,
|
||||
) -> dict[str, Any]:
|
||||
"""One page of a bank's tag histogram: ``{"items": [{tag, count}], "total", "limit", "offset"}``.
|
||||
|
||||
``memory_units`` lives in SQL for this store, so the wildcard filter, the
|
||||
``count DESC, tag ASC`` ordering and the paging all run in SQL — the whole
|
||||
histogram never crosses the wire. The dialect fragments come from
|
||||
``build_tag_listing_parts`` (``unnest`` on Postgres, ``JSON_TABLE`` on Oracle):
|
||||
this module backs both dialects, so it must not inline either one's SQL.
|
||||
"""
|
||||
from ...db import create_data_access_ops
|
||||
|
||||
ops = create_data_access_ops(getattr(conn, "backend_type", "postgresql"))
|
||||
tag_parts = ops.build_tag_listing_parts(fq_table("memory_units"))
|
||||
tag_source = tag_parts.tag_source
|
||||
non_empty_check = tag_parts.non_empty_check
|
||||
tag_col = tag_parts.tag_col
|
||||
bank_prefix = tag_parts.bank_prefix
|
||||
|
||||
params: list[Any] = [bank_id]
|
||||
pattern_clause = ""
|
||||
if pattern:
|
||||
# '*' is the wildcard, matched case-insensitively — same anchored ILIKE semantics as before.
|
||||
params.append(pattern.replace("*", "%"))
|
||||
pattern_clause = f"AND {tag_col} ILIKE $2"
|
||||
|
||||
total_row = await conn.fetchrow(
|
||||
f"""
|
||||
SELECT COUNT(DISTINCT {tag_col}) as total
|
||||
FROM {tag_source}
|
||||
WHERE {bank_prefix}bank_id = $1 {non_empty_check}
|
||||
{pattern_clause}
|
||||
""",
|
||||
*params,
|
||||
)
|
||||
total = int(total_row["total"]) if total_row else 0
|
||||
|
||||
limit_param = len(params) + 1
|
||||
offset_param = len(params) + 2
|
||||
params.extend([limit, offset])
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT {tag_col} as tag, COUNT(*) as count
|
||||
FROM {tag_source}
|
||||
WHERE {bank_prefix}bank_id = $1 {non_empty_check}
|
||||
{pattern_clause}
|
||||
GROUP BY {tag_col}
|
||||
ORDER BY count DESC, {tag_col} ASC
|
||||
LIMIT ${limit_param} OFFSET ${offset_param}
|
||||
""",
|
||||
*params,
|
||||
)
|
||||
return {
|
||||
"items": [{"tag": row["tag"], "count": int(row["count"])} for row in rows],
|
||||
"total": total,
|
||||
"limit": limit,
|
||||
"offset": offset,
|
||||
}
|
||||
|
||||
|
||||
async def find_unconsolidated(
|
||||
*,
|
||||
conn,
|
||||
fq_table: Callable[[str], str],
|
||||
bank_id: str,
|
||||
fact_types: list[str],
|
||||
limit: int,
|
||||
scope_tags: list[str] | None = None,
|
||||
) -> list[StoredMemory]:
|
||||
"""Memories not yet folded into an observation, oldest first.
|
||||
|
||||
The consolidator's candidate query: never consolidated, never *failed* to
|
||||
consolidate (a memory the LLM could not handle must not be retried forever),
|
||||
ordered by ``created_at`` so the queue drains in arrival order. ``scope_tags``
|
||||
is the same ``tags @> scope`` containment the job's scope filter uses — the
|
||||
job ORs several scopes together; one scope is passed here.
|
||||
"""
|
||||
where = [
|
||||
"bank_id = $1",
|
||||
"consolidated_at IS NULL",
|
||||
"consolidation_failed_at IS NULL",
|
||||
]
|
||||
params: list[Any] = [bank_id]
|
||||
if fact_types:
|
||||
params.append(list(fact_types))
|
||||
where.append(f"fact_type = ANY(${len(params)})")
|
||||
if scope_tags:
|
||||
params.append(list(scope_tags))
|
||||
where.append(f"tags @> ${len(params)}::varchar[]")
|
||||
params.append(limit)
|
||||
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT {_MEMORY_COLUMNS}
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE {" AND ".join(where)}
|
||||
ORDER BY created_at ASC
|
||||
LIMIT ${len(params)}
|
||||
""",
|
||||
*params,
|
||||
)
|
||||
return [_stored_from_row(row) for row in rows]
|
||||
|
||||
|
||||
async def count_unconsolidated(
|
||||
*,
|
||||
conn,
|
||||
fq_table: Callable[[str], str],
|
||||
bank_id: str,
|
||||
fact_types: list[str],
|
||||
scopes: list[list[str] | None],
|
||||
limit: int,
|
||||
) -> int:
|
||||
"""Bounded ``COUNT(*)`` of unconsolidated candidates matching any scope — the cheap counterpart
|
||||
to :func:`find_unconsolidated` that never ships a row.
|
||||
|
||||
Same predicates as ``find_unconsolidated`` (never consolidated, never failed, matching
|
||||
fact_type), with the scopes OR'd as ``tags @> scope`` containment. ``id`` is the PK so each row
|
||||
counts once; the inner ``LIMIT`` floors the count at ``limit`` exactly as walking that many rows
|
||||
would, so a huge backlog stays a single index count instead of a 17-column fetch.
|
||||
"""
|
||||
where = ["bank_id = $1", "consolidated_at IS NULL", "consolidation_failed_at IS NULL"]
|
||||
params: list[Any] = [bank_id]
|
||||
if fact_types:
|
||||
params.append(list(fact_types))
|
||||
where.append(f"fact_type = ANY(${len(params)})")
|
||||
# An unscoped entry (None) matches every row, collapsing the OR to no tag filter at all.
|
||||
scope_clauses: list[str] = []
|
||||
unscoped = any(scope is None for scope in scopes)
|
||||
if not unscoped:
|
||||
for scope in scopes:
|
||||
params.append(list(scope or []))
|
||||
scope_clauses.append(f"tags @> ${len(params)}::varchar[]")
|
||||
if scope_clauses:
|
||||
where.append("(" + " OR ".join(scope_clauses) + ")")
|
||||
params.append(limit)
|
||||
row = await conn.fetchrow(
|
||||
f"""
|
||||
SELECT COUNT(*) AS c FROM (
|
||||
SELECT 1 FROM {fq_table("memory_units")}
|
||||
WHERE {" AND ".join(where)}
|
||||
LIMIT ${len(params)}
|
||||
) sub
|
||||
""",
|
||||
*params,
|
||||
)
|
||||
return int(row["c"]) if row else 0
|
||||
|
||||
|
||||
async def mark_consolidated(
|
||||
*,
|
||||
conn,
|
||||
fq_table: Callable[[str], str],
|
||||
bank_id: str,
|
||||
unit_ids: list[str],
|
||||
when: datetime | None,
|
||||
failed: bool = False,
|
||||
) -> None:
|
||||
"""Stamp (or clear, with ``when=None``) the consolidated marker on sources.
|
||||
|
||||
``failed`` writes ``consolidation_failed_at`` instead of ``consolidated_at``,
|
||||
which is what keeps a memory the LLM could not consolidate out of the queue.
|
||||
|
||||
``when=None`` clears the column rather than stamping it — that is how a source
|
||||
is requeued once the observation built on it is deleted. The clear keeps the
|
||||
``fact_type IN ('experience', 'world')`` guard the requeue sites carry:
|
||||
observations are never themselves consolidated, so nothing about them should
|
||||
be reset by a requeue.
|
||||
|
||||
``updated_at`` is deliberately left alone, matching the consolidator's own
|
||||
statements: consolidation bookkeeping is not an edit to the memory, and
|
||||
bumping it would make every consolidation pass look like a write to the
|
||||
staleness check below.
|
||||
"""
|
||||
ids = _as_uuids(unit_ids)
|
||||
if not ids:
|
||||
return
|
||||
column = "consolidation_failed_at" if failed else "consolidated_at"
|
||||
guard = "" if when is not None else " AND fact_type IN ('experience', 'world')"
|
||||
await conn.execute(
|
||||
f"""
|
||||
UPDATE {fq_table("memory_units")}
|
||||
SET {column} = $1
|
||||
WHERE bank_id = $2 AND id = ANY($3::uuid[]){guard}
|
||||
""",
|
||||
when,
|
||||
bank_id,
|
||||
ids,
|
||||
)
|
||||
|
||||
|
||||
async def any_memory_updated_since(
|
||||
*,
|
||||
conn,
|
||||
fq_table: Callable[[str], str],
|
||||
bank_id: str,
|
||||
since: datetime,
|
||||
fact_types: list[str] | None = None,
|
||||
tags: list[str] | None = None,
|
||||
tags_match: str = "any",
|
||||
tag_groups: list | None = None,
|
||||
) -> bool:
|
||||
"""Whether any memory in ``bank_id``'s scope was written after ``since``.
|
||||
|
||||
Backs the mental-model staleness check, so it is a bounded existence test —
|
||||
``LIMIT 1``, never a COUNT: the answer is "is there one", and the planner can
|
||||
stop at the first hit. The scope is the mental model's: its flat tags (or the
|
||||
compound ``tag_groups``) plus an optional ``fact_types`` restriction. This is
|
||||
where the staleness query's WHERE lives, so the same scope that gates a
|
||||
refresh decides whether one is due.
|
||||
"""
|
||||
params: list[Any] = [bank_id, since]
|
||||
where = ["bank_id = $1", "updated_at > $2"]
|
||||
|
||||
tag_clause, tag_params, next_param = build_tags_where_clause(tags, param_offset=len(params) + 1, match=tags_match)
|
||||
if tag_clause:
|
||||
where.append(tag_clause.removeprefix("AND "))
|
||||
params.extend(tag_params)
|
||||
|
||||
group_clause, group_params, _ = build_tag_groups_where_clause(tag_groups, param_offset=next_param)
|
||||
if group_clause:
|
||||
where.append(group_clause.removeprefix("AND "))
|
||||
params.extend(group_params)
|
||||
# Untagged, no tag_groups → no tag constraint, matching any memory in the bank.
|
||||
|
||||
if fact_types:
|
||||
params.append(list(fact_types))
|
||||
where.append(f"fact_type = ANY(${len(params)}::text[])")
|
||||
|
||||
row = await conn.fetchval(
|
||||
f"SELECT 1 FROM {fq_table('memory_units')} WHERE {' AND '.join(where)} LIMIT 1",
|
||||
*params,
|
||||
)
|
||||
return row is not None
|
||||
|
||||
|
||||
__all__ = [
|
||||
"any_memory_updated_since",
|
||||
"count_memories",
|
||||
"find_unconsolidated",
|
||||
"get_memories",
|
||||
"list_tags",
|
||||
"mark_consolidated",
|
||||
"scan_memories",
|
||||
]
|
||||
@@ -1,584 +0,0 @@
|
||||
"""Writes against `memory_units`: the fact insert, the deletes, and observation invalidation.
|
||||
|
||||
Everything here mutates the memories slice and nothing else. The document row,
|
||||
the chunks, the entity registry and the link tables stay with their own callers —
|
||||
what lands in this module is only the statements that touch `memory_units` (and,
|
||||
on backends that keep one, the `observation_sources` junction that hangs off it).
|
||||
|
||||
Each function takes the live connection and Hindsight's ``fq_table`` resolver, so
|
||||
it runs inside whatever transaction the caller already holds; ``ops`` is the
|
||||
dialect ops object, which is what lets the same code serve the PG (native array)
|
||||
and Oracle (junction table) shapes of the observation→source relation.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import uuid
|
||||
from collections.abc import Callable
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from ....config import get_config
|
||||
from ..base import StoredMemory
|
||||
|
||||
if TYPE_CHECKING: # pragma: no cover - typing only
|
||||
from ...retain.types import ProcessedFact
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
async def insert_facts(
|
||||
*,
|
||||
conn,
|
||||
ops,
|
||||
bank_id: str,
|
||||
facts: list[ProcessedFact],
|
||||
document_id: str | None = None,
|
||||
) -> list[str]:
|
||||
"""Insert facts into the database in batch.
|
||||
|
||||
Args:
|
||||
conn: Database connection
|
||||
bank_id: Bank identifier
|
||||
facts: List of ProcessedFact objects to insert
|
||||
document_id: Optional document ID to associate with facts
|
||||
|
||||
Returns:
|
||||
List of unit IDs (UUIDs as strings) for the inserted facts, in the same
|
||||
order as ``facts``.
|
||||
"""
|
||||
if not facts:
|
||||
return []
|
||||
|
||||
# Imported here: `retain` reaches back into the engine for `fq_table`, so a
|
||||
# module-level import would close the cycle once the engine imports this store.
|
||||
from ...retain.fact_extraction import _sanitize_text
|
||||
|
||||
# Prepare data for batch insert
|
||||
fact_texts = []
|
||||
embeddings = []
|
||||
event_dates = []
|
||||
occurred_starts = []
|
||||
occurred_ends = []
|
||||
mentioned_ats = []
|
||||
contexts = []
|
||||
fact_types = []
|
||||
metadata_jsons = []
|
||||
chunk_ids = []
|
||||
document_ids = []
|
||||
tags_list = []
|
||||
observation_scopes_list = []
|
||||
text_signals_list = []
|
||||
|
||||
for fact in facts:
|
||||
fact_texts.append(_sanitize_text(fact.fact_text))
|
||||
# Convert embedding to string for asyncpg vector type
|
||||
embeddings.append(str(fact.embedding))
|
||||
# event_date: Use occurred_start if available, otherwise use mentioned_at
|
||||
# This maintains backward compatibility while handling None occurred_start
|
||||
event_dates.append(fact.occurred_start if fact.occurred_start is not None else fact.mentioned_at)
|
||||
occurred_starts.append(fact.occurred_start)
|
||||
occurred_ends.append(fact.occurred_end)
|
||||
mentioned_ats.append(fact.mentioned_at)
|
||||
contexts.append(_sanitize_text(fact.context))
|
||||
fact_types.append(fact.fact_type)
|
||||
metadata_jsons.append(json.dumps(fact.metadata))
|
||||
chunk_ids.append(fact.chunk_id)
|
||||
# Use per-fact document_id if available, otherwise fallback to batch-level document_id
|
||||
document_ids.append(fact.document_id if fact.document_id else document_id)
|
||||
# Convert tags to JSON string for proper batch insertion (PostgreSQL unnest doesn't handle 2D arrays well)
|
||||
tags_list.append(json.dumps(fact.tags if fact.tags else []))
|
||||
# observation_scopes: stored as JSONB (string or 2D array), None if not provided
|
||||
observation_scopes_list.append(
|
||||
json.dumps(fact.observation_scopes) if fact.observation_scopes is not None else None
|
||||
)
|
||||
# Build text_signals: entity names + date tokens for enriched BM25 indexing
|
||||
signal_parts = []
|
||||
if fact.entities:
|
||||
signal_parts.extend(e.name for e in fact.entities)
|
||||
if fact.occurred_start:
|
||||
try:
|
||||
signal_parts.append(fact.occurred_start.strftime("%B %d %Y").lstrip("0").replace(" 0", " "))
|
||||
except (ValueError, AttributeError):
|
||||
pass
|
||||
if fact.occurred_end and fact.occurred_end != fact.occurred_start:
|
||||
try:
|
||||
signal_parts.append(fact.occurred_end.strftime("%B %d %Y").lstrip("0").replace(" 0", " "))
|
||||
except (ValueError, AttributeError):
|
||||
pass
|
||||
text_signals_list.append(" ".join(signal_parts) if signal_parts else None)
|
||||
|
||||
# Batch insert all facts — delegates to DataAccessOps which handles
|
||||
# unnest (PG) vs row-by-row (Oracle) transparently.
|
||||
config = get_config()
|
||||
|
||||
return await ops.insert_facts_batch(
|
||||
conn,
|
||||
bank_id,
|
||||
fact_texts,
|
||||
embeddings,
|
||||
event_dates,
|
||||
occurred_starts,
|
||||
occurred_ends,
|
||||
mentioned_ats,
|
||||
contexts,
|
||||
fact_types,
|
||||
metadata_jsons,
|
||||
chunk_ids,
|
||||
document_ids,
|
||||
tags_list,
|
||||
observation_scopes_list,
|
||||
text_signals_list,
|
||||
text_search_extension=config.text_search_extension,
|
||||
)
|
||||
|
||||
|
||||
async def delete_document(*, conn, fq_table: Callable[[str], str], bank_id: str, document_id: str) -> None:
|
||||
"""Delete every memory unit belonging to ``document_id``.
|
||||
|
||||
Explicitly delete memory_units by document_id BEFORE deleting the
|
||||
document row. The CASCADE from documents→chunks→memory_units only
|
||||
catches units that have a non-NULL chunk_id FK. Units with chunk_id=NULL
|
||||
(e.g. from partial writes or edge cases) would survive the cascade.
|
||||
This explicit delete ensures complete cleanup.
|
||||
|
||||
Called when a document is replaced, so it races the replacement's writes: it
|
||||
must remove only what was written *before* this call, never the facts
|
||||
arriving moments later — which the ``document_id``/``bank_id`` predicate
|
||||
gives for free inside the caller's transaction.
|
||||
"""
|
||||
await conn.execute(
|
||||
f"DELETE FROM {fq_table('memory_units')} WHERE document_id = $1 AND bank_id = $2",
|
||||
document_id,
|
||||
bank_id,
|
||||
)
|
||||
|
||||
|
||||
async def delete_observations(*, conn, fq_table: Callable[[str], str], bank_id: str) -> None:
|
||||
"""Delete all observations in a bank, leaving the facts behind them.
|
||||
|
||||
Only the observation rows: requeuing the surviving sources (clearing
|
||||
``consolidated_at``) and resetting the bank's consolidation timestamp belong
|
||||
to the caller, which owns the bank row.
|
||||
"""
|
||||
await conn.execute(
|
||||
f"DELETE FROM {fq_table('memory_units')} WHERE bank_id = $1 AND fact_type = 'observation'",
|
||||
bank_id,
|
||||
)
|
||||
|
||||
|
||||
async def observations_for_sources(
|
||||
*,
|
||||
conn,
|
||||
ops,
|
||||
fq_table: Callable[[str], str],
|
||||
bank_id: str,
|
||||
unit_ids: list[str | uuid.UUID],
|
||||
) -> list[StoredMemory]:
|
||||
"""Observations consolidated from any of ``unit_ids``.
|
||||
|
||||
Only ``unit_id`` and ``source_memory_ids`` are populated — the caller uses
|
||||
them to delete the observations and to work out which sources survive, and
|
||||
the rest of the row is about to be deleted anyway.
|
||||
"""
|
||||
if not unit_ids:
|
||||
return []
|
||||
|
||||
fact_uuids = [uuid.UUID(str(fid)) if not isinstance(fid, uuid.UUID) else fid for fid in unit_ids]
|
||||
|
||||
if ops is not None and not ops.uses_observation_sources_table:
|
||||
# PG: use native array overlap operator
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT id, source_memory_ids
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE bank_id = $1
|
||||
AND fact_type = 'observation'
|
||||
AND source_memory_ids && $2::uuid[]
|
||||
""",
|
||||
bank_id,
|
||||
fact_uuids,
|
||||
)
|
||||
else:
|
||||
# Oracle / default: use observation_sources junction table
|
||||
rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT mu.id, mu.source_memory_ids
|
||||
FROM {fq_table("memory_units")} mu
|
||||
WHERE mu.bank_id = $1
|
||||
AND mu.fact_type = 'observation'
|
||||
AND EXISTS (
|
||||
SELECT 1 FROM {fq_table("observation_sources")} os
|
||||
WHERE os.observation_id = mu.id
|
||||
AND os.source_id = ANY($2::uuid[])
|
||||
)
|
||||
""",
|
||||
bank_id,
|
||||
fact_uuids,
|
||||
)
|
||||
|
||||
return [
|
||||
StoredMemory(
|
||||
unit_id=str(row["id"]),
|
||||
text="",
|
||||
fact_type="observation",
|
||||
source_memory_ids=[str(src_id) for src_id in (row["source_memory_ids"] or [])],
|
||||
)
|
||||
for row in rows
|
||||
]
|
||||
|
||||
|
||||
async def delete_stale_observations(
|
||||
*,
|
||||
conn,
|
||||
ops,
|
||||
fq_table: Callable[[str], str],
|
||||
bank_id: str,
|
||||
fact_ids: list[str | uuid.UUID],
|
||||
) -> int:
|
||||
"""Delete observations whose source memories are about to be removed.
|
||||
|
||||
Mirrors the cleanup performed by ``MemoryEngine.delete_document`` so that
|
||||
every code path that removes ``memory_units`` also removes the
|
||||
observations derived from them. Without this, ingesting a fresh version
|
||||
of a document via the retain pipeline (which does a full-replace
|
||||
``DELETE FROM documents`` cascade) used to leave orphan observations
|
||||
pointing at memory IDs that no longer existed.
|
||||
|
||||
For each observation referencing any of ``fact_ids``:
|
||||
1. Delete the observation row (its text is stale once even one source
|
||||
memory disappears).
|
||||
2. Reset ``consolidated_at = NULL`` on the surviving source memories so
|
||||
they get re-consolidated under fresh observations on the next run.
|
||||
|
||||
Must be called within an active transaction, before the source memories
|
||||
are deleted.
|
||||
|
||||
Returns the number of observations deleted.
|
||||
"""
|
||||
if not fact_ids:
|
||||
return 0
|
||||
|
||||
fact_uuids = [uuid.UUID(str(fid)) if not isinstance(fid, uuid.UUID) else fid for fid in fact_ids]
|
||||
|
||||
affected_obs = await observations_for_sources(
|
||||
conn=conn, ops=ops, fq_table=fq_table, bank_id=bank_id, unit_ids=fact_uuids
|
||||
)
|
||||
if not affected_obs:
|
||||
return 0
|
||||
|
||||
deleted_set = {str(uid) for uid in fact_uuids}
|
||||
obs_ids = [uuid.UUID(obs.unit_id) for obs in affected_obs]
|
||||
seen_remaining: set[str] = set()
|
||||
remaining_source_ids: list[uuid.UUID] = []
|
||||
for obs in affected_obs:
|
||||
for src_str in obs.source_memory_ids:
|
||||
if src_str not in deleted_set and src_str not in seen_remaining:
|
||||
remaining_source_ids.append(uuid.UUID(src_str))
|
||||
seen_remaining.add(src_str)
|
||||
|
||||
await conn.execute(
|
||||
f"DELETE FROM {fq_table('memory_units')} WHERE id = ANY($1::uuid[])",
|
||||
obs_ids,
|
||||
)
|
||||
|
||||
# Their history is keyed by observation_id and no longer cascades from memory_units (that FK
|
||||
# was dropped so history can be recorded for observations kept outside SQL), so drop the
|
||||
# deleted observations' snapshots explicitly rather than leaving them to accumulate.
|
||||
await conn.execute(
|
||||
f"DELETE FROM {fq_table('observation_history')} WHERE bank_id = $1 AND observation_id = ANY($2::uuid[])",
|
||||
bank_id,
|
||||
obs_ids,
|
||||
)
|
||||
|
||||
if remaining_source_ids:
|
||||
await conn.execute(
|
||||
f"""
|
||||
UPDATE {fq_table("memory_units")}
|
||||
SET consolidated_at = NULL
|
||||
WHERE id = ANY($1::uuid[])
|
||||
AND fact_type IN ('experience', 'world')
|
||||
""",
|
||||
remaining_source_ids,
|
||||
)
|
||||
|
||||
logger.info(
|
||||
f"[OBSERVATIONS] Deleted {len(obs_ids)} observations, reset {len(remaining_source_ids)} "
|
||||
f"source memories for re-consolidation in bank {bank_id}"
|
||||
)
|
||||
return len(obs_ids)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------- curation archive
|
||||
#
|
||||
# Invalidation moves a rejected memory between two tables rather than flagging it,
|
||||
# so recall / consolidation / graph never carry a "valid?" predicate: live facts
|
||||
# live in `memory_units`, invalidated ones in `invalidated_memory_units`. The
|
||||
# archive is cold storage — no index, so it drops the `embedding` and
|
||||
# `search_vector` columns, which are recomputed on the way back.
|
||||
|
||||
# The two recall-surface columns the archive omits. Both follow server config
|
||||
# (embedding dimension, search backend), so keeping them out of the INSERT…SELECT
|
||||
# round-trip makes a model or text-backend switch structurally unable to trip a
|
||||
# type/dimension mismatch (#2209, #2503); each is recomputed on revert.
|
||||
_ARCHIVE_OMITTED = ('"embedding"', '"search_vector"')
|
||||
|
||||
|
||||
async def _memory_unit_columns(conn, fq_table: Callable[[str], str]) -> str:
|
||||
"""The quoted, ordinal column list of `memory_units`.
|
||||
|
||||
Read from the catalog rather than hardcoded so a schema migration cannot make
|
||||
the archive round-trip drift from the live table (the archive is created via
|
||||
``LIKE memory_units``, so the lists line up).
|
||||
"""
|
||||
rows = await conn.fetch(
|
||||
f"SELECT a.attname FROM pg_attribute a "
|
||||
f"WHERE a.attrelid = '{fq_table('memory_units')}'::regclass "
|
||||
f"AND a.attnum > 0 AND NOT a.attisdropped ORDER BY a.attnum"
|
||||
)
|
||||
return ", ".join(f'"{r["attname"]}"' for r in rows)
|
||||
|
||||
|
||||
async def _archive_columns(conn, fq_table: Callable[[str], str]) -> str:
|
||||
"""`_memory_unit_columns` minus the two the archive does not carry."""
|
||||
collist = await _memory_unit_columns(conn, fq_table)
|
||||
return ", ".join(c for c in (s.strip() for s in collist.split(",")) if c not in _ARCHIVE_OMITTED)
|
||||
|
||||
|
||||
_ARCHIVE_SELECT = (
|
||||
"id, text, fact_type, context, occurred_start, occurred_end, mentioned_at, "
|
||||
"document_id, chunk_id, tags, metadata, proof_count, event_date, created_at, "
|
||||
"consolidated_at, entity_ids"
|
||||
)
|
||||
|
||||
|
||||
def _archived_stored(row: Any) -> StoredMemory:
|
||||
"""Map an `invalidated_memory_units` row onto :class:`StoredMemory`."""
|
||||
return StoredMemory(
|
||||
unit_id=str(row["id"]),
|
||||
text=row["text"],
|
||||
fact_type=row["fact_type"],
|
||||
context=row["context"],
|
||||
document_id=row["document_id"],
|
||||
chunk_id=str(row["chunk_id"]) if row["chunk_id"] else None,
|
||||
tags=list(row["tags"] or []),
|
||||
metadata=row["metadata"] if isinstance(row["metadata"], dict) else None,
|
||||
proof_count=row["proof_count"] or 1,
|
||||
event_date=row["event_date"],
|
||||
occurred_start=row["occurred_start"],
|
||||
occurred_end=row["occurred_end"],
|
||||
mentioned_at=row["mentioned_at"],
|
||||
created_at=row["created_at"],
|
||||
consolidated_at=row["consolidated_at"],
|
||||
entity_ids=[str(e) for e in (row["entity_ids"] or [])],
|
||||
)
|
||||
|
||||
|
||||
async def get_archived_memory(*, conn, fq_table, bank_id: str, unit_id: str) -> StoredMemory | None:
|
||||
row = await conn.fetchrow(
|
||||
f"SELECT {_ARCHIVE_SELECT} FROM {fq_table('invalidated_memory_units')} WHERE id = $1 AND bank_id = $2",
|
||||
str(unit_id),
|
||||
bank_id,
|
||||
)
|
||||
return _archived_stored(row) if row else None
|
||||
|
||||
|
||||
async def invalidate_memory(*, conn, fq_table, bank_id: str, unit_id: str, reason: str | None) -> bool:
|
||||
mu = fq_table("memory_units")
|
||||
arch = fq_table("invalidated_memory_units")
|
||||
ue = fq_table("unit_entities")
|
||||
arch_cols = await _archive_columns(conn, fq_table)
|
||||
|
||||
# Snapshot the entity ids before the delete cascade takes `unit_entities`, so
|
||||
# revert can restore the postings the move is about to drop.
|
||||
entity_ids = [
|
||||
r["entity_id"] for r in await conn.fetch(f"SELECT entity_id FROM {ue} WHERE unit_id = $1", str(unit_id))
|
||||
]
|
||||
# Causal edges are retain-time extraction output the FK cascade would destroy for good —
|
||||
# unlike temporal/semantic links they can't be recomputed, so snapshot their descriptors onto
|
||||
# the archive row and revert rematerializes them (#2864).
|
||||
from ...retain.link_utils import snapshot_causal_links
|
||||
|
||||
causal_links = await snapshot_causal_links(conn, bank_id, str(unit_id))
|
||||
inserted = await conn.fetchval(
|
||||
f"INSERT INTO {arch} ({arch_cols}, invalidation_reason, invalidated_at, entity_ids, causal_links) "
|
||||
f"SELECT {arch_cols}, $2, now(), $3::uuid[], $5::jsonb FROM {mu} WHERE id = $1 AND bank_id = $4 "
|
||||
f"RETURNING id",
|
||||
str(unit_id),
|
||||
reason,
|
||||
entity_ids,
|
||||
bank_id,
|
||||
json.dumps([descriptor.as_json_dict() for descriptor in causal_links]),
|
||||
)
|
||||
if inserted is None:
|
||||
return False
|
||||
# The cascade prunes `unit_entities` and `memory_links` with the row.
|
||||
await conn.execute(f"DELETE FROM {mu} WHERE id = $1 AND bank_id = $2", str(unit_id), bank_id)
|
||||
return True
|
||||
|
||||
|
||||
async def set_invalidation_reason(*, conn, fq_table, bank_id: str, unit_id: str, reason: str | None) -> None:
|
||||
await conn.execute(
|
||||
f"UPDATE {fq_table('invalidated_memory_units')} SET invalidation_reason = $3 WHERE id = $1 AND bank_id = $2",
|
||||
str(unit_id),
|
||||
bank_id,
|
||||
reason,
|
||||
)
|
||||
|
||||
|
||||
async def restore_memory(*, conn, fq_table, bank_id: str, unit_id: str) -> StoredMemory | None:
|
||||
mu = fq_table("memory_units")
|
||||
arch = fq_table("invalidated_memory_units")
|
||||
ue = fq_table("unit_entities")
|
||||
ent = fq_table("entities")
|
||||
arch_cols = await _archive_columns(conn, fq_table)
|
||||
|
||||
arch_row = await conn.fetchrow(
|
||||
f"SELECT {_ARCHIVE_SELECT} FROM {arch} WHERE id = $1 AND bank_id = $2", str(unit_id), bank_id
|
||||
)
|
||||
if arch_row is None:
|
||||
return None
|
||||
|
||||
# Move the row back. The archive omits embedding/search_vector, so both default
|
||||
# to NULL here; search_vector is rebuilt now, the embedding by the caller.
|
||||
await conn.execute(
|
||||
f"INSERT INTO {mu} ({arch_cols}) SELECT {arch_cols} FROM {arch} WHERE id = $1 AND bank_id = $2",
|
||||
str(unit_id),
|
||||
bank_id,
|
||||
)
|
||||
# Rebuild search_vector with the *current* backend, so a backend change while
|
||||
# the fact sat archived cannot leave a stale/wrong-type vector (#2503). None
|
||||
# means the backend indexes base columns directly and leaves it empty.
|
||||
from ...db.ops_postgresql import pg_search_vector_expr
|
||||
|
||||
sv_expr = pg_search_vector_expr(get_config())
|
||||
if sv_expr is not None:
|
||||
await conn.execute(
|
||||
f"UPDATE {mu} SET search_vector = {sv_expr} WHERE id = $1 AND bank_id = $2", str(unit_id), bank_id
|
||||
)
|
||||
# Re-consolidate from scratch; links are rebuilt by graph maintenance.
|
||||
await conn.execute(
|
||||
f"UPDATE {mu} SET consolidated_at = NULL, consolidation_failed_at = NULL, updated_at = now() "
|
||||
f"WHERE id = $1 AND bank_id = $2",
|
||||
str(unit_id),
|
||||
bank_id,
|
||||
)
|
||||
# Restore the entity postings for entities that still exist — some may have
|
||||
# been swept as orphans while the memory was archived.
|
||||
if arch_row["entity_ids"]:
|
||||
await conn.execute(
|
||||
f"INSERT INTO {ue} (unit_id, entity_id) "
|
||||
f"SELECT $1, eid FROM unnest($2::uuid[]) AS eid "
|
||||
f"WHERE EXISTS (SELECT 1 FROM {ent} e WHERE e.id = eid AND e.bank_id = $3) "
|
||||
f"ON CONFLICT DO NOTHING",
|
||||
str(unit_id),
|
||||
arch_row["entity_ids"],
|
||||
bank_id,
|
||||
)
|
||||
# Rematerialize the causal edges parked at invalidation (#2864). Edges whose peer is still
|
||||
# archived or permanently deleted are skipped — the peer keeps its own copy and recreates the
|
||||
# edge when it reverts, so the restore is order-independent and idempotent.
|
||||
from ...retain.link_utils import rematerialize_causal_links
|
||||
from .graph import _ops_for
|
||||
|
||||
causal_json = await conn.fetchval(
|
||||
f"SELECT causal_links FROM {arch} WHERE id = $1 AND bank_id = $2", str(unit_id), bank_id
|
||||
)
|
||||
if causal_json:
|
||||
await rematerialize_causal_links(conn, bank_id, conn.parse_json(causal_json) or [], ops=_ops_for(conn))
|
||||
# Invalidation cascaded away this unit's derived outgoing links; queue it so graph maintenance
|
||||
# rebuilds them (the drain only touches queued units — it never scans for missing adjacency).
|
||||
await _ops_for(conn).enqueue_graph_maintenance(
|
||||
conn, fq_table("graph_maintenance_queue"), bank_id, [uuid.UUID(str(unit_id))]
|
||||
)
|
||||
await conn.execute(f"DELETE FROM {arch} WHERE id = $1 AND bank_id = $2", str(unit_id), bank_id)
|
||||
return _archived_stored(arch_row)
|
||||
|
||||
|
||||
async def set_memory_embedding(*, conn, fq_table, bank_id: str, unit_id: str, embedding) -> None:
|
||||
await conn.execute(
|
||||
f"UPDATE {fq_table('memory_units')} SET embedding = $3::vector WHERE id = $1 AND bank_id = $2",
|
||||
str(unit_id),
|
||||
bank_id,
|
||||
embedding,
|
||||
)
|
||||
|
||||
|
||||
async def clear_unit_entities(*, conn, fq_table, bank_id: str, unit_id: str) -> None:
|
||||
await conn.execute(f"DELETE FROM {fq_table('unit_entities')} WHERE unit_id = $1", str(unit_id))
|
||||
|
||||
|
||||
async def apply_edit(
|
||||
*,
|
||||
conn,
|
||||
fq_table,
|
||||
bank_id: str,
|
||||
unit_id: str,
|
||||
text: str,
|
||||
context: str | None,
|
||||
fact_type: str,
|
||||
occurred_start,
|
||||
occurred_end,
|
||||
event_date,
|
||||
mentioned_at,
|
||||
entity_ids: list[str] | None,
|
||||
) -> None:
|
||||
# `entity_ids` and `mentioned_at` are unused here: the entity postings are
|
||||
# re-linked into `unit_entities` by the caller, and an edit does not move the
|
||||
# mention time. Both are on the signature for a store that carries entities on
|
||||
# the memory and rebuilds it wholesale.
|
||||
from ...causal_links import CAUSAL_LINK_TYPES
|
||||
from ...db.ops_postgresql import pg_search_vector_expr
|
||||
|
||||
mu = fq_table("memory_units")
|
||||
ml = fq_table("memory_links")
|
||||
# The caller enqueues the relink victims (and the edited unit itself, via
|
||||
# ``include_affected_units``) before invoking this — one combined queue insert keeps the
|
||||
# graph-maintenance queue's lock ordering intact.
|
||||
# Keep the stored text-search vector in sync with the edited text/context.
|
||||
# Reference the bind parameters, not the columns: PostgreSQL evaluates the
|
||||
# UPDATE's RHS before the sibling SET assignments land, so a column reference
|
||||
# would see the pre-edit values.
|
||||
sv_expr = pg_search_vector_expr(get_config(), text_col="$3", context_col="$4")
|
||||
sv_clause = f", search_vector = {sv_expr}" if sv_expr else ""
|
||||
await conn.execute(
|
||||
f"""
|
||||
UPDATE {mu}
|
||||
SET text = $3, context = $4, fact_type = $5, occurred_start = $6, occurred_end = $7,
|
||||
event_date = $8, consolidated_at = NULL, consolidation_failed_at = NULL,
|
||||
edited_at = now(), updated_at = now(){sv_clause}
|
||||
WHERE id = $1 AND bank_id = $2
|
||||
""",
|
||||
str(unit_id),
|
||||
bank_id,
|
||||
text,
|
||||
context,
|
||||
fact_type,
|
||||
occurred_start,
|
||||
occurred_end,
|
||||
event_date,
|
||||
)
|
||||
# Drop only the DERIVED links — graph maintenance recomputes temporal/semantic. Causal edges
|
||||
# are retain-time extraction output that nothing recreates, so an edit preserves them (#2864).
|
||||
await conn.execute(
|
||||
f"DELETE FROM {ml} WHERE (from_unit_id = $1 OR to_unit_id = $1) AND NOT (link_type = ANY($2::text[]))",
|
||||
str(unit_id),
|
||||
list(CAUSAL_LINK_TYPES),
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"apply_edit",
|
||||
"clear_unit_entities",
|
||||
"delete_document",
|
||||
"delete_observations",
|
||||
"delete_stale_observations",
|
||||
"get_archived_memory",
|
||||
"insert_facts",
|
||||
"invalidate_memory",
|
||||
"observations_for_sources",
|
||||
"restore_memory",
|
||||
"set_invalidation_reason",
|
||||
"set_memory_embedding",
|
||||
]
|
||||
@@ -1,509 +0,0 @@
|
||||
"""The default memories store: Postgres holds the memories and the links.
|
||||
|
||||
This is the behaviour Hindsight has always had, stated as an implementation of
|
||||
:class:`~hindsight_api.engine.memories.base.MemoriesExtension` rather than as the
|
||||
absence of one. Rows go in `memory_units`, the joins around it are `memory_links`
|
||||
and `unit_entities`, and every read is SQL — writing a row *is* indexing it, so
|
||||
:meth:`index_facts` has nothing left to do.
|
||||
|
||||
The class is deliberately thin. Each method delegates to a plain function in
|
||||
:mod:`hindsight_api.engine.memories.pg`, split by what calls it — curation,
|
||||
graph, reads, writes — so a change to one area is a change to one file, and the
|
||||
SQL is grouped by concern rather than piled behind a class. The two retrieval
|
||||
arms delegate further out still, to the query functions that already own them in
|
||||
:mod:`hindsight_api.engine.search.retrieval`.
|
||||
|
||||
Keeping this as an explicit store (rather than an ``if store is None`` branch at
|
||||
each call site) means the default path is the one the whole test suite exercises,
|
||||
and a second implementation cannot change it by accident.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime
|
||||
from typing import Any
|
||||
|
||||
from .base import DeletePredicate, MemoriesExtension, MemoryPatch, ScanPage, StoredMemory
|
||||
from .pg import counts, curation, graph, reads, writes
|
||||
|
||||
|
||||
class PostgresMemories(MemoriesExtension):
|
||||
"""Memories in `memory_units`, links in `memory_links` / `unit_entities`."""
|
||||
|
||||
name = "postgres"
|
||||
|
||||
# ------------------------------------------------------------------ writes
|
||||
|
||||
async def insert_facts(
|
||||
self,
|
||||
*,
|
||||
conn,
|
||||
ops,
|
||||
bank_id: str,
|
||||
facts: list,
|
||||
document_id: str | None = None,
|
||||
defer_index: bool = False,
|
||||
txn=None,
|
||||
) -> list[str]:
|
||||
# `txn` is ignored: Postgres memories live in the caller's own transaction, so the
|
||||
# write is already atomic with it — there is no separate store to hold invisible.
|
||||
# `defer_index` is meaningless here: the INSERT that returns the ids is
|
||||
# also what indexes the facts, so there is nothing to defer.
|
||||
return await writes.insert_facts(conn=conn, ops=ops, bank_id=bank_id, facts=facts, document_id=document_id)
|
||||
|
||||
async def delete_facts(self, bank_id: str, unit_ids: list[str], *, txn=None) -> None:
|
||||
"""No-op: the caller's `memory_units` DELETE (or its FK cascade) removed them."""
|
||||
|
||||
async def delete_where(self, bank_id: str, predicate: DeletePredicate, txn=None) -> int:
|
||||
"""No-op: predicate deletes are issued as SQL by the caller that owns the transaction."""
|
||||
return 0
|
||||
|
||||
async def delete_document(self, *, conn, fq_table, bank_id: str, document_id: str, txn=None) -> None:
|
||||
# `txn` ignored: Postgres memories are covered by the caller's own transaction.
|
||||
await writes.delete_document(conn=conn, fq_table=fq_table, bank_id=bank_id, document_id=document_id)
|
||||
|
||||
async def drop_bank_storage(self, bank_id: str) -> None:
|
||||
"""No-op: deleting the bank cascades to its memories."""
|
||||
|
||||
async def delete_observations(self, *, conn, fq_table, bank_id: str, txn=None) -> None:
|
||||
await writes.delete_observations(conn=conn, fq_table=fq_table, bank_id=bank_id)
|
||||
|
||||
async def update_memories(self, bank_id: str, patches: list[MemoryPatch], txn=None) -> None:
|
||||
"""No-op: the caller's UPDATE already wrote the row it holds open."""
|
||||
|
||||
# ------------------------------------------------------------------ recall arms
|
||||
|
||||
async def search(
|
||||
self,
|
||||
*,
|
||||
conn,
|
||||
bank_id: str,
|
||||
fact_types: list[str],
|
||||
query_embedding: str,
|
||||
query_text: str,
|
||||
limit: int,
|
||||
tags: list[str] | None = None,
|
||||
tags_match: str = "any",
|
||||
tag_groups: list | None = None,
|
||||
created_after: datetime | None = None,
|
||||
created_before: datetime | None = None,
|
||||
min_semantic: float | None = None,
|
||||
min_keyword: float | None = None,
|
||||
graph_seed_min_similarity: float | None = None,
|
||||
) -> "dict[str, SemanticBm25Result]":
|
||||
# Imported here: retrieval imports this package, so a module-level import
|
||||
# would close the cycle.
|
||||
from ..search.retrieval import retrieve_semantic_bm25_combined_sql
|
||||
|
||||
return await retrieve_semantic_bm25_combined_sql(
|
||||
conn,
|
||||
query_embedding,
|
||||
query_text,
|
||||
bank_id,
|
||||
fact_types,
|
||||
limit,
|
||||
tags=tags,
|
||||
tags_match=tags_match,
|
||||
tag_groups=tag_groups,
|
||||
created_after=created_after,
|
||||
created_before=created_before,
|
||||
min_semantic=min_semantic,
|
||||
min_keyword=min_keyword,
|
||||
graph_seed_min_similarity=graph_seed_min_similarity,
|
||||
)
|
||||
|
||||
async def temporal_search(
|
||||
self,
|
||||
*,
|
||||
conn,
|
||||
bank_id: str,
|
||||
fact_types: list[str],
|
||||
query_embedding: str,
|
||||
start_date: datetime,
|
||||
end_date: datetime,
|
||||
limit: int,
|
||||
semantic_threshold: float = 0.1,
|
||||
tags: list[str] | None = None,
|
||||
tags_match: str = "any",
|
||||
tag_groups: list | None = None,
|
||||
created_after: datetime | None = None,
|
||||
created_before: datetime | None = None,
|
||||
) -> dict[str, list]:
|
||||
from ..search.retrieval import retrieve_temporal_combined_sql
|
||||
|
||||
return await retrieve_temporal_combined_sql(
|
||||
conn,
|
||||
query_embedding,
|
||||
bank_id,
|
||||
fact_types,
|
||||
start_date,
|
||||
end_date,
|
||||
limit,
|
||||
semantic_threshold=semantic_threshold,
|
||||
tags=tags,
|
||||
tags_match=tags_match,
|
||||
tag_groups=tag_groups,
|
||||
created_after=created_after,
|
||||
created_before=created_before,
|
||||
)
|
||||
|
||||
# ------------------------------------------------------------------ addressed reads
|
||||
|
||||
async def get_memories(self, *, conn, fq_table, bank_id: str, unit_ids: list[str]) -> list[StoredMemory]:
|
||||
return await reads.get_memories(conn=conn, fq_table=fq_table, bank_id=bank_id, unit_ids=unit_ids)
|
||||
|
||||
async def scan_memories(
|
||||
self,
|
||||
*,
|
||||
conn,
|
||||
fq_table,
|
||||
bank_id: str,
|
||||
fact_types: list[str] | None = None,
|
||||
limit: int = 100,
|
||||
page_token: str = "",
|
||||
tags: list[str] | None = None,
|
||||
tags_match: str = "any",
|
||||
tag_groups: list | None = None,
|
||||
document_id: str | None = None,
|
||||
metadata_equals: dict[str, str] | None = None,
|
||||
skip: int = 0,
|
||||
include_edges: bool = False,
|
||||
) -> ScanPage:
|
||||
return await reads.scan_memories(
|
||||
conn=conn,
|
||||
fq_table=fq_table,
|
||||
bank_id=bank_id,
|
||||
fact_types=fact_types,
|
||||
limit=limit,
|
||||
page_token=page_token,
|
||||
tags=tags,
|
||||
tags_match=tags_match,
|
||||
tag_groups=tag_groups,
|
||||
document_id=document_id,
|
||||
metadata_equals=metadata_equals,
|
||||
skip=skip,
|
||||
include_edges=include_edges,
|
||||
)
|
||||
|
||||
async def count_memories(self, *, conn, fq_table, bank_id: str) -> dict[str, int]:
|
||||
return await reads.count_memories(conn=conn, fq_table=fq_table, bank_id=bank_id)
|
||||
|
||||
async def list_tags(
|
||||
self,
|
||||
*,
|
||||
conn,
|
||||
fq_table,
|
||||
bank_id: str,
|
||||
pattern: str | None = None,
|
||||
limit: int = 100,
|
||||
offset: int = 0,
|
||||
) -> dict[str, Any]:
|
||||
return await reads.list_tags(
|
||||
conn=conn, fq_table=fq_table, bank_id=bank_id, pattern=pattern, limit=limit, offset=offset
|
||||
)
|
||||
|
||||
async def find_unconsolidated(
|
||||
self,
|
||||
*,
|
||||
conn,
|
||||
fq_table,
|
||||
bank_id: str,
|
||||
fact_types: list[str],
|
||||
limit: int,
|
||||
scope_tags: list[str] | None = None,
|
||||
) -> list[StoredMemory]:
|
||||
return await reads.find_unconsolidated(
|
||||
conn=conn,
|
||||
fq_table=fq_table,
|
||||
bank_id=bank_id,
|
||||
fact_types=fact_types,
|
||||
limit=limit,
|
||||
scope_tags=scope_tags,
|
||||
)
|
||||
|
||||
async def count_unconsolidated(
|
||||
self,
|
||||
*,
|
||||
conn,
|
||||
fq_table,
|
||||
bank_id: str,
|
||||
fact_types: list[str],
|
||||
scopes: list[list[str] | None],
|
||||
limit: int,
|
||||
) -> int:
|
||||
return await reads.count_unconsolidated(
|
||||
conn=conn, fq_table=fq_table, bank_id=bank_id, fact_types=fact_types, scopes=scopes, limit=limit
|
||||
)
|
||||
|
||||
async def mark_consolidated(
|
||||
self,
|
||||
*,
|
||||
conn,
|
||||
fq_table,
|
||||
bank_id: str,
|
||||
unit_ids: list[str],
|
||||
when: datetime | None,
|
||||
failed: bool = False,
|
||||
txn=None,
|
||||
) -> None:
|
||||
await reads.mark_consolidated(
|
||||
conn=conn, fq_table=fq_table, bank_id=bank_id, unit_ids=unit_ids, when=when, failed=failed
|
||||
)
|
||||
|
||||
async def any_memory_updated_since(
|
||||
self,
|
||||
*,
|
||||
conn,
|
||||
fq_table,
|
||||
bank_id: str,
|
||||
since: datetime,
|
||||
fact_types: list[str] | None = None,
|
||||
tags: list[str] | None = None,
|
||||
tags_match: str = "any",
|
||||
tag_groups: list | None = None,
|
||||
) -> bool:
|
||||
return await reads.any_memory_updated_since(
|
||||
conn=conn,
|
||||
fq_table=fq_table,
|
||||
bank_id=bank_id,
|
||||
since=since,
|
||||
fact_types=fact_types,
|
||||
tags=tags,
|
||||
tags_match=tags_match,
|
||||
tag_groups=tag_groups,
|
||||
)
|
||||
|
||||
# -- count surfaces --
|
||||
|
||||
async def consolidation_freshness(self, *, conn, fq_table, bank_id: str) -> dict[str, Any]:
|
||||
return await counts.consolidation_freshness(conn=conn, fq_table=fq_table, bank_id=bank_id)
|
||||
|
||||
async def document_memory_counts(self, *, conn, fq_table, bank_id: str, document_ids: list[str]) -> dict[str, int]:
|
||||
return await counts.document_memory_counts(
|
||||
conn=conn, fq_table=fq_table, bank_id=bank_id, document_ids=document_ids
|
||||
)
|
||||
|
||||
async def link_counts(self, *, conn, fq_table, bank_id: str) -> dict[str, int]:
|
||||
return await counts.link_counts(conn=conn, fq_table=fq_table, bank_id=bank_id)
|
||||
|
||||
async def memories_timeseries(
|
||||
self, *, conn, fq_table, bank_id: str, time_field: str, trunc: str, since: datetime
|
||||
) -> list[dict[str, Any]]:
|
||||
return await counts.memories_timeseries(
|
||||
conn=conn, fq_table=fq_table, bank_id=bank_id, time_field=time_field, trunc=trunc, since=since
|
||||
)
|
||||
|
||||
async def observation_scope_counts(self, *, conn, fq_table, bank_id: str) -> list[dict[str, Any]]:
|
||||
return await counts.observation_scope_counts(conn=conn, fq_table=fq_table, bank_id=bank_id)
|
||||
|
||||
# ------------------------------------------------------------------ observations
|
||||
|
||||
async def upsert_observation(self, *, conn, bank_id: str, record, txn=None) -> None:
|
||||
"""No-op: the observation was written as a `memory_units` row by the caller."""
|
||||
|
||||
async def observations_for_sources(
|
||||
self, *, conn, ops, fq_table, bank_id: str, unit_ids: list[str]
|
||||
) -> list[StoredMemory]:
|
||||
return await writes.observations_for_sources(
|
||||
conn=conn, ops=ops, fq_table=fq_table, bank_id=bank_id, unit_ids=unit_ids
|
||||
)
|
||||
|
||||
async def delete_stale_observations(self, *, conn, ops, fq_table, bank_id: str, fact_ids: list) -> int:
|
||||
return await writes.delete_stale_observations(
|
||||
conn=conn, ops=ops, fq_table=fq_table, bank_id=bank_id, fact_ids=fact_ids
|
||||
)
|
||||
|
||||
# ------------------------------------------------------------------ curation reads
|
||||
|
||||
async def list_memory_units(
|
||||
self,
|
||||
*,
|
||||
conn,
|
||||
ops,
|
||||
fq_table,
|
||||
bank_id: str,
|
||||
fact_type: str | None = None,
|
||||
search_query: str | None = None,
|
||||
consolidation_state: str | None = None,
|
||||
state: str | None = None,
|
||||
document_id: str | None = None,
|
||||
entity_id: str | None = None,
|
||||
tags: list[str] | None = None,
|
||||
tags_match: str = "any",
|
||||
created_before: datetime | None = None,
|
||||
limit: int = 100,
|
||||
offset: int = 0,
|
||||
) -> dict[str, Any]:
|
||||
return await curation.list_memory_units(
|
||||
conn=conn,
|
||||
ops=ops,
|
||||
fq_table=fq_table,
|
||||
bank_id=bank_id,
|
||||
fact_type=fact_type,
|
||||
search_query=search_query,
|
||||
consolidation_state=consolidation_state,
|
||||
state=state,
|
||||
document_id=document_id,
|
||||
entity_id=entity_id,
|
||||
tags=tags,
|
||||
tags_match=tags_match,
|
||||
created_before=created_before,
|
||||
limit=limit,
|
||||
offset=offset,
|
||||
)
|
||||
|
||||
async def get_memory_unit(self, *, conn, ops, fq_table, bank_id: str, unit_id: str) -> dict[str, Any] | None:
|
||||
return await curation.get_memory_unit(conn=conn, ops=ops, fq_table=fq_table, bank_id=bank_id, unit_id=unit_id)
|
||||
|
||||
# -- curation archive --
|
||||
|
||||
async def get_archived_memory(self, *, conn, fq_table, bank_id: str, unit_id: str) -> StoredMemory | None:
|
||||
return await writes.get_archived_memory(conn=conn, fq_table=fq_table, bank_id=bank_id, unit_id=unit_id)
|
||||
|
||||
async def invalidate_memory(
|
||||
self, *, conn, fq_table, bank_id: str, unit_id: str, reason: str | None, txn=None
|
||||
) -> bool:
|
||||
return await writes.invalidate_memory(
|
||||
conn=conn, fq_table=fq_table, bank_id=bank_id, unit_id=unit_id, reason=reason
|
||||
)
|
||||
|
||||
async def set_invalidation_reason(self, *, conn, fq_table, bank_id: str, unit_id: str, reason: str | None) -> None:
|
||||
await writes.set_invalidation_reason(
|
||||
conn=conn, fq_table=fq_table, bank_id=bank_id, unit_id=unit_id, reason=reason
|
||||
)
|
||||
|
||||
async def restore_memory(self, *, conn, fq_table, bank_id: str, unit_id: str, txn=None) -> StoredMemory | None:
|
||||
return await writes.restore_memory(conn=conn, fq_table=fq_table, bank_id=bank_id, unit_id=unit_id)
|
||||
|
||||
async def set_memory_embedding(self, *, conn, fq_table, bank_id: str, unit_id: str, embedding, txn=None) -> None:
|
||||
await writes.set_memory_embedding(
|
||||
conn=conn, fq_table=fq_table, bank_id=bank_id, unit_id=unit_id, embedding=embedding
|
||||
)
|
||||
|
||||
async def clear_unit_entities(self, *, conn, fq_table, bank_id: str, unit_id: str) -> None:
|
||||
await writes.clear_unit_entities(conn=conn, fq_table=fq_table, bank_id=bank_id, unit_id=unit_id)
|
||||
|
||||
async def apply_edit(
|
||||
self,
|
||||
*,
|
||||
conn,
|
||||
fq_table,
|
||||
bank_id: str,
|
||||
unit_id: str,
|
||||
text: str,
|
||||
context: str | None,
|
||||
fact_type: str,
|
||||
occurred_start,
|
||||
occurred_end,
|
||||
event_date,
|
||||
mentioned_at,
|
||||
entity_ids: list[str] | None,
|
||||
txn=None,
|
||||
) -> None:
|
||||
await writes.apply_edit(
|
||||
conn=conn,
|
||||
fq_table=fq_table,
|
||||
bank_id=bank_id,
|
||||
unit_id=unit_id,
|
||||
text=text,
|
||||
context=context,
|
||||
fact_type=fact_type,
|
||||
occurred_start=occurred_start,
|
||||
occurred_end=occurred_end,
|
||||
event_date=event_date,
|
||||
mentioned_at=mentioned_at,
|
||||
entity_ids=entity_ids,
|
||||
)
|
||||
|
||||
async def list_entities(
|
||||
self,
|
||||
*,
|
||||
conn,
|
||||
fq_table,
|
||||
bank_id: str,
|
||||
search: str | None = None,
|
||||
limit: int = 100,
|
||||
offset: int = 0,
|
||||
) -> dict[str, Any]:
|
||||
return await curation.list_entities(
|
||||
conn=conn, fq_table=fq_table, bank_id=bank_id, search=search, limit=limit, offset=offset
|
||||
)
|
||||
|
||||
# ------------------------------------------------------------------ graph
|
||||
|
||||
async def graph_units(
|
||||
self,
|
||||
*,
|
||||
conn,
|
||||
fq_table,
|
||||
bank_id: str,
|
||||
fact_type: str | None = None,
|
||||
search_query: str | None = None,
|
||||
document_id: str | None = None,
|
||||
chunk_id: str | None = None,
|
||||
tags: list[str] | None = None,
|
||||
tags_match: str = "all_strict",
|
||||
limit: int = 1000,
|
||||
) -> dict[str, Any]:
|
||||
return await graph.graph_units(
|
||||
conn=conn,
|
||||
fq_table=fq_table,
|
||||
bank_id=bank_id,
|
||||
fact_type=fact_type,
|
||||
search_query=search_query,
|
||||
document_id=document_id,
|
||||
chunk_id=chunk_id,
|
||||
tags=tags,
|
||||
tags_match=tags_match,
|
||||
limit=limit,
|
||||
)
|
||||
|
||||
async def graph_entity_rows(self, *, conn, fq_table, bank_id: str, unit_ids: list[str]) -> list[dict[str, Any]]:
|
||||
return await graph.graph_entity_rows(conn=conn, fq_table=fq_table, bank_id=bank_id, unit_ids=unit_ids)
|
||||
|
||||
async def graph_direct_links(self, *, conn, fq_table, bank_id: str, unit_ids: list[str]) -> list[dict[str, Any]]:
|
||||
return await graph.graph_direct_links(conn=conn, fq_table=fq_table, bank_id=bank_id, unit_ids=unit_ids)
|
||||
|
||||
async def entity_memory_counts(
|
||||
self, *, conn, fq_table, bank_id: str, entity_ids: list[str] | None = None
|
||||
) -> dict[str, int]:
|
||||
return await graph.entity_memory_counts(conn=conn, fq_table=fq_table, bank_id=bank_id, entity_ids=entity_ids)
|
||||
|
||||
async def entities_for_units(self, *, conn, fq_table, bank_id: str, unit_ids: list[str]) -> dict[str, list[str]]:
|
||||
return await graph.entities_for_units(conn=conn, fq_table=fq_table, bank_id=bank_id, unit_ids=unit_ids)
|
||||
|
||||
async def entity_map_for_units(
|
||||
self, *, conn, fq_table, bank_id: str, unit_ids: list[str]
|
||||
) -> dict[str, list[dict[str, str]]]:
|
||||
return await graph.entity_map_for_units(conn=conn, fq_table=fq_table, bank_id=bank_id, unit_ids=unit_ids)
|
||||
|
||||
# ------------------------------------------------------------------ maintenance
|
||||
|
||||
async def record_unit_entities(
|
||||
self, *, conn, ops, fq_table, bank_id: str | None = None, unit_ids: list[Any], entity_ids: list[Any]
|
||||
) -> None:
|
||||
# The join is keyed by global unit id, so bank_id is not needed here.
|
||||
await ops.bulk_insert_unit_entities(conn, fq_table("unit_entities"), unit_ids, entity_ids)
|
||||
|
||||
async def enqueue_relink_victims(
|
||||
self, *, conn, fq_table, bank_id: str, affected_unit_ids: list, include_affected_units: bool = False
|
||||
) -> int:
|
||||
return await graph.enqueue_relink_victims(
|
||||
conn=conn,
|
||||
fq_table=fq_table,
|
||||
bank_id=bank_id,
|
||||
affected_unit_ids=affected_unit_ids,
|
||||
include_affected_units=include_affected_units,
|
||||
)
|
||||
|
||||
async def relink_pass(self, *, backend, fq_table, bank_id: str, config) -> dict:
|
||||
return await graph.relink_pass(backend=backend, fq_table=fq_table, bank_id=bank_id, config=config)
|
||||
|
||||
async def prune_orphan_entities(self, *, conn, fq_table, bank_id: str) -> int:
|
||||
return await graph.prune_orphan_entities(conn=conn, fq_table=fq_table, bank_id=bank_id)
|
||||
|
||||
async def prune_stale_cooccurrences(self, *, conn, fq_table, bank_id: str) -> int:
|
||||
return await graph.prune_stale_cooccurrences(conn=conn, fq_table=fq_table, bank_id=bank_id)
|
||||
|
||||
|
||||
__all__ = ["PostgresMemories"]
|
||||
File diff suppressed because it is too large
Load Diff
@@ -19,18 +19,10 @@ class BatchRetainParentMetadata:
|
||||
total_tokens: int
|
||||
num_sub_batches: int
|
||||
is_parent: bool = True
|
||||
# Set only when the whole batch targets a single document, so the operations
|
||||
# list surfaces which document an in-flight retain is (re)writing. The
|
||||
# documents UI cross-checks this to badge rows as "updating". Multi-document
|
||||
# batches leave it None and are matched per single-document child instead.
|
||||
document_id: str | None = None
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
"""Convert to dict for JSON serialization, omitting document_id when unset."""
|
||||
data = asdict(self)
|
||||
if data.get("document_id") is None:
|
||||
data.pop("document_id", None)
|
||||
return data
|
||||
"""Convert to dict for JSON serialization."""
|
||||
return asdict(self)
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -41,15 +33,10 @@ class BatchRetainChildMetadata:
|
||||
parent_operation_id: str
|
||||
sub_batch_index: int
|
||||
total_sub_batches: int
|
||||
# Set only when this child processes a single document (see the parent's note).
|
||||
document_id: str | None = None
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
"""Convert to dict for JSON serialization, omitting document_id when unset."""
|
||||
data = asdict(self)
|
||||
if data.get("document_id") is None:
|
||||
data.pop("document_id", None)
|
||||
return data
|
||||
"""Convert to dict for JSON serialization."""
|
||||
return asdict(self)
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -155,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)
|
||||
|
||||
@@ -72,7 +72,6 @@ class MarkitdownParser(FileParser):
|
||||
ocr_base_url: str | None = None,
|
||||
ocr_model: str | None = None,
|
||||
ocr_prompt: str | None = None,
|
||||
ocr_default_headers: dict[str, str] | None = None,
|
||||
):
|
||||
"""Initialize markitdown parser."""
|
||||
# Lazy import to avoid requiring markitdown for all users
|
||||
@@ -90,7 +89,6 @@ class MarkitdownParser(FileParser):
|
||||
base_url=ocr_base_url,
|
||||
model=ocr_model,
|
||||
prompt=ocr_prompt,
|
||||
default_headers=ocr_default_headers,
|
||||
)
|
||||
self._markitdown = MarkItDown(
|
||||
llm_client=ocr_options.llm_client,
|
||||
@@ -107,7 +105,6 @@ class MarkitdownParser(FileParser):
|
||||
base_url: str | None,
|
||||
model: str | None,
|
||||
prompt: str | None,
|
||||
default_headers: dict[str, str] | None,
|
||||
) -> MarkitdownOcrOptions:
|
||||
"""Build MarkItDown options for OpenAI-compatible image OCR."""
|
||||
if not model or not model.strip():
|
||||
@@ -132,15 +129,8 @@ class MarkitdownParser(FileParser):
|
||||
except ImportError as e:
|
||||
raise RuntimeError("openai package is required when Markitdown OCR is enabled.") from e
|
||||
|
||||
client_kwargs: dict[str, object] = {
|
||||
"api_key": api_key,
|
||||
"base_url": base_url.strip(),
|
||||
}
|
||||
if default_headers:
|
||||
client_kwargs["default_headers"] = default_headers
|
||||
|
||||
return MarkitdownOcrOptions(
|
||||
llm_client=OpenAI(**client_kwargs),
|
||||
llm_client=OpenAI(api_key=api_key, base_url=base_url.strip()),
|
||||
llm_model=model.strip(),
|
||||
llm_prompt=prompt or DEFAULT_FILE_PARSER_MARKITDOWN_OCR_PROMPT,
|
||||
)
|
||||
@@ -196,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,27 +209,10 @@ 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
|
||||
# Pin the configured model (issue #2881). Without this the spawned CLI
|
||||
# runs its own default model — an Opus-class model on Pro/Max OAuth —
|
||||
# regardless of HINDSIGHT_API_*_LLM_MODEL, while metrics/logs still print
|
||||
# self.model, so the mismatch is invisible. The isolated CLAUDE_CONFIG_DIR
|
||||
# (fresh temp dir) means a host settings.json can't reach the CLI either,
|
||||
# so passing it through here is the only channel.
|
||||
model=self.model or None,
|
||||
env=_get_isolated_claude_env(),
|
||||
)
|
||||
|
||||
@@ -260,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
|
||||
@@ -324,7 +287,7 @@ class ClaudeCodeLLM(LLMInterface):
|
||||
|
||||
# Record trace span
|
||||
try:
|
||||
from hindsight_api.tracing import _serialize_for_span, get_span_recorder
|
||||
from hindsight_api.tracing import get_span_recorder
|
||||
|
||||
span_recorder = get_span_recorder()
|
||||
span_recorder.record_llm_call(
|
||||
@@ -332,17 +295,15 @@ class ClaudeCodeLLM(LLMInterface):
|
||||
model=self.model,
|
||||
scope=scope,
|
||||
messages=messages,
|
||||
response_content=_serialize_for_span(result),
|
||||
response_content=result if isinstance(result, str) else result.model_dump_json(),
|
||||
input_tokens=estimated_input,
|
||||
output_tokens=estimated_output,
|
||||
duration=duration,
|
||||
finish_reason=None,
|
||||
error=None,
|
||||
)
|
||||
except Exception as span_error:
|
||||
# Tracing must remain best-effort, but expose instrumentation
|
||||
# bugs that would otherwise silently erase spans (#3025).
|
||||
logger.debug("Claude Code span recording failed: %s", span_error, exc_info=True)
|
||||
except Exception:
|
||||
pass # logging failure must never affect the operation
|
||||
|
||||
# Log slow calls
|
||||
if duration > 10.0:
|
||||
@@ -401,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.
|
||||
@@ -419,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
|
||||
@@ -432,7 +393,6 @@ class ClaudeCodeLLM(LLMInterface):
|
||||
AssistantMessage,
|
||||
ClaudeAgentOptions,
|
||||
ClaudeSDKClient,
|
||||
ResultMessage,
|
||||
SdkMcpTool,
|
||||
TextBlock,
|
||||
ToolUseBlock,
|
||||
@@ -513,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 = {}
|
||||
@@ -541,33 +504,16 @@ class ClaudeCodeLLM(LLMInterface):
|
||||
# else: tool_choice == "auto" or unspecified - use default behavior (no changes needed)
|
||||
|
||||
# Configure SDK options with MCP server
|
||||
#
|
||||
# tools=[] disables built-in CLI tools (Read, Write, Bash, ToolSearch, etc.)
|
||||
# Without this, Claude Code CLI defers MCP tools when too many built-in tools
|
||||
# are loaded, forcing Claude to use ToolSearch first — which wastes the turn
|
||||
# are loaded, forcing Claude to use ToolSearch first — which wastes the max_turns
|
||||
# budget and prevents direct MCP tool calls.
|
||||
#
|
||||
# max_turns=1 is critical (issue #2966). call_with_tools() is one *round* of
|
||||
# an agentic loop the caller drives: the model proposes tool calls, we return
|
||||
# them, and the orchestrator (reflect/agent.py) executes the REAL tools and
|
||||
# feeds the results back on the next call. The SDK, however, runs its own
|
||||
# in-process loop: it invokes our SDK MCP handlers — which are deliberate
|
||||
# placeholders returning "[Tool <name> called successfully]" (no real data) —
|
||||
# and lets the model react. With max_turns >= 2 the model calls recall, sees
|
||||
# the empty placeholder, re-queries with reworded searches, exhausts the turn
|
||||
# budget, and the run ends in error_max_turns with its tool calls discarded —
|
||||
# exactly the "0 tool calls / no information" failure in #2966. Capping at a
|
||||
# single turn stops the SDK from acting on the placeholder results: the model
|
||||
# emits its first tool call (or a text answer) and we return that to the caller
|
||||
# unchanged, matching how every other provider's call_with_tools() behaves.
|
||||
options = ClaudeAgentOptions(
|
||||
system_prompt=system_prompt if system_prompt else None,
|
||||
tools=[], # Disable built-in tools so MCP tools load eagerly
|
||||
max_turns=1, # One round: propose tool calls (or answer); caller drives the loop
|
||||
max_turns=2, # Allow tool call + tool result round-trip
|
||||
mcp_servers=mcp_servers_config,
|
||||
allowed_tools=allowed_tool_names,
|
||||
# Pin the configured model (issue #2881) — see the call() options block.
|
||||
model=self.model or None,
|
||||
env=_get_isolated_claude_env(),
|
||||
)
|
||||
|
||||
@@ -604,21 +550,6 @@ class ClaudeCodeLLM(LLMInterface):
|
||||
arguments=block.input,
|
||||
)
|
||||
)
|
||||
if tool_calls:
|
||||
# This round proposed tool call(s). Stop consuming the
|
||||
# stream so the SDK does not run another turn against our
|
||||
# placeholder handlers (issue #2966) — the caller executes
|
||||
# the real tools and calls us again with the results.
|
||||
break
|
||||
elif isinstance(message, ResultMessage) and message.is_error:
|
||||
# With max_turns=1 the CLI reports error_max_turns whenever
|
||||
# the model spent its single turn issuing a tool call (there
|
||||
# was no follow-up turn to emit final text). That is expected
|
||||
# here and not a failure: we already captured the tool call
|
||||
# above and break before reaching this branch. Only a genuine
|
||||
# error with nothing to return should surface (issue #2702).
|
||||
if not tool_calls:
|
||||
raise RuntimeError(_result_error_detail(message))
|
||||
|
||||
# Record metrics
|
||||
duration = time.time() - start_time
|
||||
|
||||
@@ -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):
|
||||
"""
|
||||
@@ -172,8 +117,8 @@ class CodexLLM(LLMInterface):
|
||||
if self.model.startswith("openai/"):
|
||||
self.model = self.model[len("openai/") :]
|
||||
|
||||
# Reasoning summary controls presentation separately from the backend's
|
||||
# reasoning effort, which is sent unchanged in each request payload.
|
||||
# Map reasoning effort to Codex reasoning summary format
|
||||
# Codex supports: "auto", "concise", "detailed"
|
||||
self.reasoning_summary = self._map_reasoning_effort(reasoning_effort)
|
||||
|
||||
# HTTP client for SSE streaming
|
||||
@@ -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
|
||||
@@ -448,28 +385,20 @@ class CodexLLM(LLMInterface):
|
||||
"tools": [],
|
||||
"tool_choice": "auto",
|
||||
"parallel_tool_calls": True,
|
||||
"reasoning": {"effort": self.reasoning_effort, "summary": reasoning_summary},
|
||||
"reasoning": {"summary": reasoning_summary},
|
||||
"store": False, # Codex uses stateless mode
|
||||
"stream": True, # SSE streaming
|
||||
"include": ["reasoning.encrypted_content"],
|
||||
"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
|
||||
@@ -573,7 +467,7 @@ class CodexLLM(LLMInterface):
|
||||
|
||||
# Record trace span
|
||||
try:
|
||||
from hindsight_api.tracing import _serialize_for_span, get_span_recorder
|
||||
from hindsight_api.tracing import get_span_recorder
|
||||
|
||||
# Estimate tokens for tracing
|
||||
estimated_input = sum(len(m.get("content", "")) for m in messages) // 4
|
||||
@@ -584,17 +478,15 @@ class CodexLLM(LLMInterface):
|
||||
model=self.model,
|
||||
scope=scope,
|
||||
messages=messages,
|
||||
response_content=_serialize_for_span(result),
|
||||
response_content=result if isinstance(result, str) else result.model_dump_json(),
|
||||
input_tokens=estimated_input,
|
||||
output_tokens=estimated_output,
|
||||
duration=duration,
|
||||
finish_reason=None,
|
||||
error=None,
|
||||
)
|
||||
except Exception as span_error:
|
||||
# Tracing must remain best-effort, but expose instrumentation
|
||||
# bugs that would otherwise silently erase spans (#3025).
|
||||
logger.debug("Codex span recording failed: %s", span_error, exc_info=True)
|
||||
except Exception:
|
||||
pass # logging failure must never affect the operation
|
||||
|
||||
if return_usage:
|
||||
# Codex doesn't provide token counts, estimate based on content
|
||||
@@ -650,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)
|
||||
|
||||
@@ -748,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.
|
||||
@@ -765,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.
|
||||
@@ -827,20 +716,22 @@ 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": {"effort": self.reasoning_effort, "summary": reasoning_summary},
|
||||
"reasoning": {"summary": reasoning_summary},
|
||||
"store": False,
|
||||
"stream": True,
|
||||
"include": ["reasoning.encrypted_content"],
|
||||
"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"
|
||||
|
||||
@@ -937,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
|
||||
|
||||
@@ -983,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.
|
||||
@@ -424,7 +329,8 @@ class GeminiLLM(LLMInterface):
|
||||
last_exception = None
|
||||
|
||||
for attempt in range(max_retries + 1):
|
||||
set_stage(f"llm.gemini.{scope}.attempt={attempt + 1}/{max_retries + 1}")
|
||||
if attempt > 0:
|
||||
set_stage(f"llm.gemini.{scope}.attempt={attempt + 1}/{max_retries + 1}")
|
||||
try:
|
||||
response = await asyncio.wait_for(
|
||||
self._client.aio.models.generate_content(
|
||||
@@ -573,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
|
||||
@@ -630,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.
|
||||
@@ -646,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.
|
||||
@@ -669,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()
|
||||
@@ -737,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")
|
||||
)
|
||||
@@ -768,16 +698,13 @@ class GeminiLLM(LLMInterface):
|
||||
|
||||
last_exception = None
|
||||
for attempt in range(max_retries + 1):
|
||||
set_stage(f"llm.gemini.tools.attempt={attempt + 1}/{max_retries + 1}")
|
||||
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
|
||||
@@ -880,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
|
||||
@@ -967,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
|
||||
@@ -1164,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]] = []
|
||||
@@ -1193,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": {
|
||||
@@ -255,7 +246,8 @@ class LiteLLMLLM(LLMInterface):
|
||||
last_exception = None
|
||||
|
||||
for attempt in range(max_retries + 1):
|
||||
set_stage(f"llm.{self._stage_label}.{scope}.attempt={attempt + 1}/{max_retries + 1}")
|
||||
if attempt > 0:
|
||||
set_stage(f"llm.{self._stage_label}.{scope}.attempt={attempt + 1}/{max_retries + 1}")
|
||||
try:
|
||||
response = await asyncio.wait_for(
|
||||
self._acompletion(**call_kwargs),
|
||||
@@ -285,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
|
||||
@@ -396,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
|
||||
@@ -429,24 +408,18 @@ 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):
|
||||
set_stage(f"llm.{self._stage_label}.tools.attempt={attempt + 1}/{max_retries + 1}")
|
||||
if attempt > 0:
|
||||
set_stage(f"llm.{self._stage_label}.tools.attempt={attempt + 1}/{max_retries + 1}")
|
||||
try:
|
||||
response = await asyncio.wait_for(
|
||||
self._acompletion(**call_kwargs),
|
||||
@@ -551,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, LLMToolChoiceMode
|
||||
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.
|
||||
@@ -266,7 +266,7 @@ class MockLLM(LLMInterface):
|
||||
else:
|
||||
result = LLMToolCallResult(content="mock response", finish_reason="stop")
|
||||
else:
|
||||
result = self._compliant_tool_call(tools, tool_choice, messages)
|
||||
result = LLMToolCallResult(content="mock response", finish_reason="stop")
|
||||
|
||||
# Set mock token usage on result if not already set
|
||||
if result.input_tokens == 0:
|
||||
@@ -297,61 +297,6 @@ class MockLLM(LLMInterface):
|
||||
|
||||
return result
|
||||
|
||||
@staticmethod
|
||||
def _compliant_tool_call(
|
||||
tools: list[dict[str, Any]],
|
||||
tool_choice: LLMToolChoice,
|
||||
messages: list[dict[str, Any]],
|
||||
) -> LLMToolCallResult:
|
||||
"""Default tool response: simulate a compliant tool-calling model.
|
||||
|
||||
Real providers drive the reflect loop entirely through tool calls -- they
|
||||
honor a forced tool choice, then finish via ``done`` -- and the reflect
|
||||
agent now rejects a turn that yields no tool call at all (a transport that
|
||||
can't tool-call raises ReflectToolCallError). So the mock must behave like a
|
||||
working provider here rather than returning bare "mock response" prose,
|
||||
which used to be salvaged as the answer. Only this default path is affected;
|
||||
tests that script turns via ``_response_callback`` / ``_mock_response`` are not.
|
||||
"""
|
||||
tool_names = {t.get("function", {}).get("name") for t in tools}
|
||||
|
||||
def _mock_query() -> str:
|
||||
for message in reversed(messages):
|
||||
content = message.get("content")
|
||||
if message.get("role") == "user" and isinstance(content, str) and content.strip():
|
||||
return content[:200]
|
||||
return "mock query"
|
||||
|
||||
# Honor a forced retrieval tool so the loop actually runs recall/search and
|
||||
# gathers evidence (populates based_on for tests that assert on it).
|
||||
if tool_choice.mode is LLMToolChoiceMode.NAMED and tool_choice.function_name in {
|
||||
"search_mental_models",
|
||||
"search_observations",
|
||||
"recall",
|
||||
}:
|
||||
return LLMToolCallResult(
|
||||
tool_calls=[
|
||||
LLMToolCall(
|
||||
id="mock_forced",
|
||||
name=tool_choice.function_name,
|
||||
arguments={"reason": "mock", "query": _mock_query()},
|
||||
)
|
||||
],
|
||||
finish_reason="tool_calls",
|
||||
)
|
||||
|
||||
# Auto turn: finish via the done tool, mirroring a model that has gathered
|
||||
# enough. The reflect evidence guardrail handles the empty-bank case (no
|
||||
# evidence -> forced text synthesis on the final iteration).
|
||||
if "done" in tool_names:
|
||||
return LLMToolCallResult(
|
||||
tool_calls=[LLMToolCall(id="mock_done", name="done", arguments={"answer": "mock response"})],
|
||||
finish_reason="tool_calls",
|
||||
)
|
||||
|
||||
# No done tool offered (non-reflect tool call): fall back to plain text.
|
||||
return LLMToolCallResult(content="mock response", finish_reason="stop")
|
||||
|
||||
@staticmethod
|
||||
def _build_mock_facts(messages: list[dict]) -> dict:
|
||||
"""Build a canned fact extraction response from the user message text.
|
||||
|
||||
@@ -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,57 +47,17 @@ 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"})
|
||||
|
||||
# Local providers whose OpenAI-compatible surface always lives under a `/v1`
|
||||
# path (LM Studio: http://localhost:1234/v1, Ollama: http://localhost:11434/v1).
|
||||
# For these we know the exact endpoint shape, so a bare host base URL can be
|
||||
# normalized safely. Cloud/proxy endpoints are left untouched — their path is
|
||||
# provider-specific and must be supplied verbatim.
|
||||
_V1_PATH_LOCAL_PROVIDERS = frozenset({"lmstudio", "ollama"})
|
||||
|
||||
|
||||
def _ensure_v1_base_url(base_url: str) -> str:
|
||||
"""Append the OpenAI-compatible ``/v1`` prefix to a bare local base URL.
|
||||
|
||||
LM Studio's server UI advertises its address as ``http://localhost:1234``,
|
||||
so users commonly set ``HINDSIGHT_API_LLM_BASE_URL`` to that bare host. The
|
||||
OpenAI SDK then POSTs to ``<host>/chat/completions`` and LM Studio rejects it
|
||||
with ``Unexpected endpoint or method`` — its OpenAI-compatible routes live
|
||||
under ``/v1``. Only a base URL with no meaningful path (bare host or a lone
|
||||
trailing slash) is rewritten; anything with an explicit path (e.g. a reverse
|
||||
proxy mount or an already-correct ``/v1``) is returned unchanged. See #2922.
|
||||
"""
|
||||
parsed = urlparse(base_url)
|
||||
if parsed.path.strip("/"):
|
||||
return base_url
|
||||
return urlunparse(parsed._replace(path="/v1"))
|
||||
|
||||
|
||||
class ProviderResponseError(RuntimeError):
|
||||
"""Raised when a provider returns a success response without usable content."""
|
||||
@@ -116,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
|
||||
@@ -529,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,
|
||||
):
|
||||
"""
|
||||
@@ -541,15 +445,10 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
api_key: API key (optional for ollama/lmstudio).
|
||||
base_url: Base URL for the API (uses defaults for groq/ollama/lmstudio if empty).
|
||||
model: Model name.
|
||||
reasoning_effort: Reasoning effort level for supported models
|
||||
("none", "low", "medium", "high"). "none" is required when calling
|
||||
function tools on some reasoning models, which reject every other
|
||||
value — including omitting the parameter entirely.
|
||||
reasoning_effort: Reasoning effort level for supported models ("low", "medium", "high").
|
||||
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)
|
||||
@@ -604,11 +503,6 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
# lives on a separate control-plane host — see FireworksLLM.
|
||||
self.base_url = "https://api.fireworks.ai/inference/v1"
|
||||
|
||||
# Normalize bare local base URLs (e.g. a user pasting the address shown
|
||||
# in the LM Studio UI) so the OpenAI SDK targets the `/v1` routes. See #2922.
|
||||
if self.provider in _V1_PATH_LOCAL_PROVIDERS and self.base_url:
|
||||
self.base_url = _ensure_v1_base_url(self.base_url)
|
||||
|
||||
# For ollama/lmstudio, use dummy key if not provided
|
||||
if self.provider in ("ollama", "lmstudio") and not self.api_key:
|
||||
self.api_key = "local"
|
||||
@@ -635,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 {}
|
||||
|
||||
@@ -666,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:
|
||||
"""
|
||||
@@ -689,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,
|
||||
@@ -751,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]],
|
||||
@@ -843,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
|
||||
@@ -859,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
|
||||
@@ -902,7 +789,8 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
# Surface attempt count in worker stage so JSON-schema retry loops
|
||||
# are visible from logs (small models on strict structured output
|
||||
# often loop here). Cheap no-op outside worker context.
|
||||
set_stage(f"llm.{self.provider}.{scope}.attempt={attempt + 1}/{max_retries + 1}")
|
||||
if attempt > 0:
|
||||
set_stage(f"llm.{self.provider}.{scope}.attempt={attempt + 1}/{max_retries + 1}")
|
||||
try:
|
||||
if response_format is not None:
|
||||
response = await self._client.chat.completions.create(**call_params)
|
||||
@@ -998,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,
|
||||
@@ -1010,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
|
||||
@@ -1080,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
|
||||
)
|
||||
@@ -1173,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.
|
||||
@@ -1187,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.
|
||||
@@ -1258,16 +1147,8 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
temperature = max(0.01, min(temperature, 1.0))
|
||||
call_params["temperature"] = temperature
|
||||
|
||||
# Set reasoning_effort for reasoning models, matching call(). Omitting it
|
||||
# here is not a neutral default: OpenAI rejects function tools on a
|
||||
# reasoning model unless reasoning_effort is present and set to "none",
|
||||
# so leaving it out fails exactly like sending an unsupported value.
|
||||
if self._supports_reasoning_model():
|
||||
call_params["reasoning_effort"] = self.reasoning_effort
|
||||
|
||||
# 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:
|
||||
@@ -1278,7 +1159,8 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
last_exception = None
|
||||
|
||||
for attempt in range(max_retries + 1):
|
||||
set_stage(f"llm.{self.provider}.tools.attempt={attempt + 1}/{max_retries + 1}")
|
||||
if attempt > 0:
|
||||
set_stage(f"llm.{self.provider}.tools.attempt={attempt + 1}/{max_retries + 1}")
|
||||
try:
|
||||
response = await self._client.chat.completions.create(**call_params)
|
||||
|
||||
@@ -1314,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,
|
||||
@@ -1325,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
|
||||
@@ -1388,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
|
||||
)
|
||||
@@ -1463,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:
|
||||
@@ -1482,7 +1355,8 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
|
||||
async with httpx.AsyncClient(timeout=300.0) as client:
|
||||
for attempt in range(max_retries + 1):
|
||||
set_stage(f"llm.ollama_native.{scope}.attempt={attempt + 1}/{max_retries + 1}")
|
||||
if attempt > 0:
|
||||
set_stage(f"llm.ollama_native.{scope}.attempt={attempt + 1}/{max_retries + 1}")
|
||||
try:
|
||||
response = await client.post(native_url, json=payload, headers=headers)
|
||||
response.raise_for_status()
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -7,13 +7,12 @@ The reflect agent uses an iterative loop with tools to:
|
||||
3. Expand memories (get chunk/document context)
|
||||
"""
|
||||
|
||||
from .agent import ReflectAgentResult, ReflectToolCallError, run_reflect_agent
|
||||
from .agent import ReflectAgentResult, run_reflect_agent
|
||||
from .models import ReflectAction, ReflectActionBatch
|
||||
|
||||
__all__ = [
|
||||
"run_reflect_agent",
|
||||
"ReflectAgentResult",
|
||||
"ReflectToolCallError",
|
||||
"ReflectAction",
|
||||
"ReflectActionBatch",
|
||||
]
|
||||
|
||||
@@ -10,11 +10,11 @@ Uses hierarchical retrieval:
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import re
|
||||
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,
|
||||
@@ -49,25 +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."
|
||||
|
||||
|
||||
class ReflectToolCallError(RuntimeError):
|
||||
"""The model never produced a tool call reflect could understand.
|
||||
|
||||
Reflect is driven entirely by structured tool calls (``recall``, ``expand``,
|
||||
``done`` ...). Some provider transports do not actually support function
|
||||
calling and silently drop the tool definitions from the request (e.g. litellm's
|
||||
Vertex AI gpt-oss MaaS path strips ``tools``/``tool_choice`` when the model is
|
||||
flagged as not supporting them). The model then answers in free text that may
|
||||
mimic a ``done`` payload. Rather than salvage that untooled text -- and risk
|
||||
surfacing raw tool-call JSON as the answer -- we fail loudly so the caller can
|
||||
switch to a tool-calling-capable model/transport.
|
||||
"""
|
||||
|
||||
|
||||
def _normalize_tool_name(name: str) -> str:
|
||||
"""Normalize tool name from various LLM output formats.
|
||||
@@ -101,12 +82,148 @@ def _is_done_tool(name: str) -> bool:
|
||||
return _normalize_tool_name(name) == "done"
|
||||
|
||||
|
||||
# Pattern to match done() call as text - handles done({...}) with nested JSON
|
||||
_DONE_CALL_PATTERN = re.compile(r"done\s*\(\s*\{.*$", re.DOTALL)
|
||||
|
||||
# Patterns for leaked structured output in the answer field
|
||||
_LEAKED_JSON_SUFFIX = re.compile(
|
||||
r'\s*```(?:json)?\s*\{[^}]*(?:"(?:observation_ids|memory_ids|mental_model_ids)"|\})\s*```\s*$',
|
||||
re.DOTALL | re.IGNORECASE,
|
||||
)
|
||||
_TRAILING_IDS_PATTERN = re.compile(
|
||||
r"\s*(?:observation_ids|memory_ids|mental_model_ids)\s*[=:]\s*\[.*?\]\s*$", re.DOTALL | re.IGNORECASE
|
||||
)
|
||||
_JSON_CODE_FENCE_PATTERN = re.compile(r"^\s*```(?:json)?\s*(\{.*\})\s*```\s*$", re.DOTALL | re.IGNORECASE)
|
||||
|
||||
_DONE_ARGUMENT_KEYS = frozenset(
|
||||
{
|
||||
"answer",
|
||||
"directive_compliance",
|
||||
"memory_ids",
|
||||
"mental_model_ids",
|
||||
"observation_ids",
|
||||
"model_ids",
|
||||
}
|
||||
)
|
||||
_DONE_ARGUMENT_MARKER_KEYS = _DONE_ARGUMENT_KEYS - {"answer"}
|
||||
_LEAKED_JSON_ID_KEYS = frozenset({"memory_ids", "mental_model_ids", "observation_ids", "model_ids"})
|
||||
|
||||
|
||||
def _unwrap_leaked_done_arguments(text: str) -> str | None:
|
||||
"""Return the answer when a done tool call was rendered as JSON text.
|
||||
|
||||
Some providers leak the done tool's argument object instead of surfacing it
|
||||
as a native tool call, e.g. {"answer": "...", "memory_ids": [...]}. Only
|
||||
unwrap objects that match the done argument shape so normal JSON answers
|
||||
stay intact.
|
||||
"""
|
||||
candidate = text.strip()
|
||||
if not candidate:
|
||||
return None
|
||||
|
||||
fenced = _JSON_CODE_FENCE_PATTERN.match(candidate)
|
||||
if fenced:
|
||||
candidate = fenced.group(1).strip()
|
||||
|
||||
try:
|
||||
payload = json.loads(candidate)
|
||||
except json.JSONDecodeError:
|
||||
return None
|
||||
|
||||
if not isinstance(payload, dict):
|
||||
return None
|
||||
answer = payload.get("answer")
|
||||
if not isinstance(answer, str) or not answer.strip():
|
||||
return None
|
||||
|
||||
keys = set(payload)
|
||||
if not keys.intersection(_DONE_ARGUMENT_MARKER_KEYS):
|
||||
return None
|
||||
if not keys.issubset(_DONE_ARGUMENT_KEYS):
|
||||
return None
|
||||
|
||||
for key in ("memory_ids", "mental_model_ids", "observation_ids", "model_ids"):
|
||||
value = payload.get(key)
|
||||
if value is not None and not isinstance(value, list):
|
||||
return None
|
||||
|
||||
return answer.strip()
|
||||
|
||||
|
||||
def _strip_trailing_id_json_object(text: str) -> str:
|
||||
stripped = text.rstrip()
|
||||
if not stripped.endswith("}"):
|
||||
return text.strip()
|
||||
|
||||
start = stripped.rfind("{")
|
||||
if start < 0:
|
||||
return text.strip()
|
||||
|
||||
try:
|
||||
payload = json.loads(stripped[start:])
|
||||
except json.JSONDecodeError:
|
||||
return text.strip()
|
||||
|
||||
if not isinstance(payload, dict) or not payload:
|
||||
return text.strip()
|
||||
keys = set(payload)
|
||||
if not keys.issubset(_LEAKED_JSON_ID_KEYS):
|
||||
return text.strip()
|
||||
|
||||
return stripped[:start].strip()
|
||||
|
||||
|
||||
def _clean_answer_text(text: str) -> str:
|
||||
"""Clean up answer text by removing any done() tool call syntax.
|
||||
|
||||
Some LLMs output the done() call as text instead of a proper tool call.
|
||||
This strips out patterns like: done({"answer": "...", ...})
|
||||
"""
|
||||
unwrapped = _unwrap_leaked_done_arguments(text)
|
||||
if unwrapped is not None:
|
||||
return unwrapped
|
||||
|
||||
# Remove done() call pattern from the end of the text
|
||||
cleaned = _DONE_CALL_PATTERN.sub("", text).strip()
|
||||
return cleaned if cleaned else text
|
||||
|
||||
|
||||
def _clean_done_answer(text: str) -> str:
|
||||
"""Clean up the answer field from a done() tool call.
|
||||
|
||||
Some LLMs leak structured output patterns into the answer text, such as:
|
||||
- JSON code blocks with observation_ids/memory_ids at the end
|
||||
- Raw JSON objects with these fields
|
||||
- Plain text like "observation_ids: [...]"
|
||||
|
||||
This cleans those patterns while preserving the actual answer content.
|
||||
"""
|
||||
if not text:
|
||||
return text
|
||||
|
||||
unwrapped = _unwrap_leaked_done_arguments(text)
|
||||
if unwrapped is not None:
|
||||
return unwrapped
|
||||
|
||||
cleaned = text
|
||||
|
||||
# Remove leaked JSON in code blocks at the end
|
||||
cleaned = _LEAKED_JSON_SUFFIX.sub("", cleaned).strip()
|
||||
|
||||
# Remove leaked raw JSON objects at the end
|
||||
cleaned = _strip_trailing_id_json_object(cleaned)
|
||||
|
||||
# Remove trailing ID patterns
|
||||
cleaned = _TRAILING_IDS_PATTERN.sub("", cleaned).strip()
|
||||
|
||||
return cleaned if cleaned else text
|
||||
|
||||
|
||||
async def _generate_structured_output(
|
||||
answer: str,
|
||||
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.
|
||||
|
||||
@@ -115,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
|
||||
@@ -209,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,
|
||||
@@ -302,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,
|
||||
@@ -420,13 +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,
|
||||
store_document_text: bool = True,
|
||||
*,
|
||||
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.
|
||||
@@ -454,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
|
||||
@@ -464,8 +472,8 @@ async def _run_reflect_agent_inner(
|
||||
|
||||
# Get tools for this agent (with directive compliance field if directives exist).
|
||||
# The expand tool only reads back raw source text (chunks/documents), so it is
|
||||
# useless and excluded when document text storage is disabled (per bank).
|
||||
include_expand = store_document_text
|
||||
# useless and excluded when document text storage is disabled.
|
||||
include_expand = get_config().store_document_text
|
||||
tools = get_reflect_tools(
|
||||
directive_rules=directive_rules,
|
||||
include_mental_models=has_mental_models,
|
||||
@@ -490,75 +498,30 @@ 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
|
||||
# Whether the model has ever produced a tool call reflect could understand.
|
||||
# Stays False when a transport silently strips tool support (the model then
|
||||
# only ever returns free text) -- that case fails via ReflectToolCallError.
|
||||
saw_tool_call = False
|
||||
tool_trace: list[ToolCall] = []
|
||||
tool_trace_summary: list[dict[str, Any]] = []
|
||||
llm_trace: list[dict[str, Any]] = []
|
||||
@@ -672,12 +635,12 @@ async def _run_reflect_agent_inner(
|
||||
"output_tokens": usage.output_tokens,
|
||||
}
|
||||
)
|
||||
answer = response.strip()
|
||||
answer = _clean_answer_text(response.strip())
|
||||
|
||||
# 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
|
||||
@@ -737,11 +700,11 @@ async def _run_reflect_agent_inner(
|
||||
"output_tokens": usage.output_tokens,
|
||||
}
|
||||
)
|
||||
answer = response.strip()
|
||||
answer = _clean_answer_text(response.strip())
|
||||
|
||||
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
|
||||
@@ -776,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,
|
||||
@@ -812,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
|
||||
@@ -880,12 +826,12 @@ async def _run_reflect_agent_inner(
|
||||
"output_tokens": usage.output_tokens,
|
||||
}
|
||||
)
|
||||
answer = response.strip()
|
||||
answer = _clean_answer_text(response.strip())
|
||||
|
||||
# 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
|
||||
@@ -904,29 +850,83 @@ async def _run_reflect_agent_inner(
|
||||
directives_applied=directives_applied,
|
||||
)
|
||||
|
||||
# No tool calls this turn.
|
||||
# No tool calls - LLM wants to respond with text
|
||||
if not result.tool_calls:
|
||||
# Reflect is driven by structured tool calls. A turn with no tool call
|
||||
# means one of two things:
|
||||
# * the model already gathered evidence via earlier tool calls and is
|
||||
# now stopping -- fine, synthesize a clean final answer below;
|
||||
# * the transport can't produce tool calls at all, so it only ever
|
||||
# returns free text (e.g. litellm strips tools on the Vertex gpt-oss
|
||||
# MaaS path). In that case ``saw_tool_call`` is still False.
|
||||
# We no longer salvage that free text as the answer -- it can be a raw
|
||||
# done()-payload with sibling id fields leaking into user-visible text.
|
||||
# Fail loudly instead so the caller picks a tool-calling-capable model.
|
||||
if not saw_tool_call:
|
||||
snippet = (result.content or "").strip()
|
||||
if len(snippet) > 500:
|
||||
snippet = snippet[:500] + "..."
|
||||
detail = f" Response: {snippet!r}" if snippet else " The model returned no content."
|
||||
raise ReflectToolCallError(
|
||||
f"Reflect requires a tool-calling model, but {llm_config.provider}/{llm_config.model} "
|
||||
f"produced no usable tool call (the transport may not support function calling)." + detail
|
||||
# When directives are present but no evidence has been gathered,
|
||||
# the LLM tends to echo directive content verbatim as its answer.
|
||||
# Fall through to the final-prompt path which doesn't include
|
||||
# directives and handles "no data" gracefully.
|
||||
has_gathered_evidence = (
|
||||
bool(available_memory_ids) or bool(available_mental_model_ids) or bool(available_observation_ids)
|
||||
)
|
||||
directive_leak_risk = directives and not has_gathered_evidence
|
||||
if result.content and not directive_leak_risk:
|
||||
answer = _clean_answer_text(result.content.strip())
|
||||
|
||||
# The call_with_tools call above is intentionally uncapped so the
|
||||
# LLM has headroom to emit tool-call JSON plus any intermediate
|
||||
# reasoning. But when the LLM short-circuits and returns text
|
||||
# directly, that text becomes the user-visible final answer and
|
||||
# must respect max_tokens like the forced-final paths do. If it
|
||||
# overshoots, run one extra capped call to rewrite it within
|
||||
# the cap.
|
||||
if max_tokens is not None and 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,
|
||||
)
|
||||
total_input_tokens += rewrite_usage.input_tokens
|
||||
total_output_tokens += rewrite_usage.output_tokens
|
||||
total_cached_tokens += getattr(rewrite_usage, "cached_tokens", 0) or 0
|
||||
total_thoughts_tokens += getattr(rewrite_usage, "thoughts_tokens", 0) or 0
|
||||
llm_trace.append(
|
||||
{
|
||||
"scope": "final_rewrite",
|
||||
"duration_ms": int((time.time() - rewrite_start) * 1000),
|
||||
"input_tokens": rewrite_usage.input_tokens,
|
||||
"output_tokens": rewrite_usage.output_tokens,
|
||||
}
|
||||
)
|
||||
answer = _clean_answer_text(rewritten.strip())
|
||||
|
||||
# Generate structured output if schema provided
|
||||
structured_output = None
|
||||
if response_schema and answer:
|
||||
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
|
||||
total_cached_tokens += struct.cached_tokens
|
||||
total_thoughts_tokens += struct.thoughts_tokens
|
||||
|
||||
_log_completion(answer, iteration + 1)
|
||||
return ReflectAgentResult(
|
||||
text=answer,
|
||||
structured_output=structured_output,
|
||||
iterations=iteration + 1,
|
||||
tools_called=total_tools_called,
|
||||
tool_trace=tool_trace,
|
||||
llm_trace=_get_llm_trace(),
|
||||
usage=_get_usage(),
|
||||
directives_applied=directives_applied,
|
||||
)
|
||||
# Model tool-called earlier and is now stopping: fall through to a clean
|
||||
# forced final synthesis (tools disabled, prose expected).
|
||||
# Empty response, force final
|
||||
prompt = build_final_prompt(
|
||||
query, context_history, bank_profile, context, max_context_tokens=max_context_tokens
|
||||
)
|
||||
@@ -958,12 +958,12 @@ async def _run_reflect_agent_inner(
|
||||
"output_tokens": usage.output_tokens,
|
||||
}
|
||||
)
|
||||
answer = response.strip()
|
||||
answer = _clean_answer_text(response.strip())
|
||||
|
||||
# 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
|
||||
@@ -982,11 +982,6 @@ async def _run_reflect_agent_inner(
|
||||
directives_applied=directives_applied,
|
||||
)
|
||||
|
||||
# The model produced at least one tool call reflect could parse: it can
|
||||
# drive the loop, so a later text-only turn is a legitimate stop, not a
|
||||
# broken transport.
|
||||
saw_tool_call = True
|
||||
|
||||
# Check for done tool call (handle various LLM output formats)
|
||||
done_call = next((tc for tc in result.tool_calls if _is_done_tool(tc.name)), None)
|
||||
if done_call:
|
||||
@@ -1006,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."
|
||||
@@ -1039,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)
|
||||
@@ -1071,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."
|
||||
@@ -1082,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(
|
||||
@@ -1174,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),
|
||||
}
|
||||
)
|
||||
@@ -1257,56 +1244,15 @@ 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
|
||||
|
||||
# ``done`` is a structured tool call: trust its ``answer`` field verbatim.
|
||||
# Sibling id fields (memory_ids, ...) live in their own arguments and are
|
||||
# validated separately below -- they can't bleed into a parsed answer string.
|
||||
answer = args.get("answer", "").strip()
|
||||
# Extract and clean the answer - some LLMs leak structured output into the answer text
|
||||
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 = 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]
|
||||
@@ -1315,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)
|
||||
@@ -1439,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":
|
||||
@@ -1481,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", [])
|
||||
|
||||
@@ -421,6 +421,76 @@ def build_system_prompt_for_tools(
|
||||
return "\n".join(parts)
|
||||
|
||||
|
||||
def build_agent_prompt(
|
||||
query: str,
|
||||
context_history: list[dict],
|
||||
bank_profile: dict,
|
||||
additional_context: str | None = None,
|
||||
) -> str:
|
||||
"""Build the user prompt for the reflect agent."""
|
||||
parts = []
|
||||
|
||||
# Bank identity
|
||||
name = bank_profile.get("name", "Assistant")
|
||||
mission = bank_profile.get("mission", "")
|
||||
|
||||
parts.append(f"## Memory Bank Context\nName: {name}")
|
||||
if mission:
|
||||
parts.append(f"Mission: {mission}")
|
||||
|
||||
# Disposition traits if present
|
||||
disposition = bank_profile.get("disposition", {})
|
||||
if disposition:
|
||||
traits = []
|
||||
if "skepticism" in disposition:
|
||||
traits.append(f"skepticism={disposition['skepticism']}")
|
||||
if "literalism" in disposition:
|
||||
traits.append(f"literalism={disposition['literalism']}")
|
||||
if "empathy" in disposition:
|
||||
traits.append(f"empathy={disposition['empathy']}")
|
||||
if traits:
|
||||
parts.append(f"Disposition: {', '.join(traits)}")
|
||||
|
||||
# Additional context from caller
|
||||
if additional_context:
|
||||
parts.append(f"\n## Additional Context\n{additional_context}")
|
||||
|
||||
# Tool call history
|
||||
if context_history:
|
||||
parts.append("\n## Tool Results (synthesize and reason from this data)")
|
||||
for i, entry in enumerate(context_history, 1):
|
||||
tool = entry["tool"]
|
||||
output = entry["output"]
|
||||
# Format as proper JSON for LLM readability
|
||||
try:
|
||||
output_str = json.dumps(output, indent=2, default=str, ensure_ascii=False)
|
||||
except (TypeError, ValueError):
|
||||
output_str = str(output)
|
||||
parts.append(f"\n### Call {i}: {tool}\n```json\n{output_str}\n```")
|
||||
|
||||
# The question
|
||||
parts.append(f"\n## Question\n{query}")
|
||||
|
||||
# Instructions
|
||||
if context_history:
|
||||
parts.append(
|
||||
"\n## Instructions\n"
|
||||
"Based on the tool results above, either call more tools or provide your final answer. "
|
||||
"Synthesize and reason from the data - make reasonable inferences when helpful. "
|
||||
"If you have related information, use it to give the best possible answer."
|
||||
)
|
||||
else:
|
||||
parts.append(
|
||||
"\n## Instructions\n"
|
||||
"Start by searching for relevant information using the hierarchical retrieval strategy:\n"
|
||||
"1. Try search_mental_models() first for curated summaries\n"
|
||||
"2. Try search_observations() for consolidated knowledge\n"
|
||||
"3. Use recall() for specific details or to verify stale data"
|
||||
)
|
||||
|
||||
return "\n".join(parts)
|
||||
|
||||
|
||||
def build_final_prompt(
|
||||
query: str,
|
||||
context_history: list[dict],
|
||||
@@ -820,3 +890,38 @@ ABSOLUTE RULES:
|
||||
OUTPUT FORMAT:
|
||||
- Output ONLY the updated markdown document. No preamble, no explanation, no diff markers, no commentary.
|
||||
- Do not wrap the output in code fences unless the CURRENT DOCUMENT itself was entirely a code fence."""
|
||||
|
||||
|
||||
def build_delta_prompt(
|
||||
*,
|
||||
current_content: str,
|
||||
candidate_content: str,
|
||||
supporting_facts: list[dict[str, Any]],
|
||||
source_query: str,
|
||||
) -> str:
|
||||
"""Build the user prompt for a delta-mode mental model refresh.
|
||||
|
||||
Args:
|
||||
current_content: The existing mental model content (to preserve as much as possible).
|
||||
candidate_content: Fresh synthesis from the reflect agent reflecting new reality.
|
||||
supporting_facts: Flat list of fact dicts (id, text, type) supporting the candidate.
|
||||
source_query: The mental model's source query, for topical framing.
|
||||
"""
|
||||
fact_lines: list[str] = []
|
||||
for f in supporting_facts:
|
||||
fid = f.get("id", "")
|
||||
text = (f.get("text") or "").strip().replace("\n", " ")
|
||||
ftype = f.get("type", "")
|
||||
fact_lines.append(f"- [{ftype}:{fid}] {text}")
|
||||
facts_block = "\n".join(fact_lines) if fact_lines else "(no supporting facts retrieved)"
|
||||
|
||||
return (
|
||||
f"## Topic\n{source_query}\n\n"
|
||||
f"## CURRENT DOCUMENT\n```markdown\n{current_content}\n```\n\n"
|
||||
f"## CANDIDATE UPDATE\n```markdown\n{candidate_content}\n```\n\n"
|
||||
f"## SUPPORTING FACTS\n{facts_block}\n\n"
|
||||
"## Task\n"
|
||||
"Produce the updated mental model document by applying the minimum necessary changes "
|
||||
"to CURRENT DOCUMENT so that it reflects CANDIDATE UPDATE and SUPPORTING FACTS. "
|
||||
"Preserve unchanged content byte-for-byte. Output only the final markdown."
|
||||
)
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -95,10 +95,8 @@ async def tool_search_mental_models(
|
||||
params: list[Any] = [bank_id, str(query_embedding), max_results]
|
||||
next_param = 4
|
||||
|
||||
# Exact matching treats absent or empty tags as the global scope. Do not
|
||||
# skip the filter, or mental models would see every scope while the other
|
||||
# reflect retrieval tools correctly see only untagged data.
|
||||
if tags or tags_match == "exact":
|
||||
# Use the centralized tag filtering logic
|
||||
if tags:
|
||||
tag_clause, tag_params, next_param = build_tags_where_clause(tags, param_offset=next_param, match=tags_match)
|
||||
filters += f" {tag_clause}"
|
||||
params.extend(tag_params)
|
||||
@@ -330,52 +328,28 @@ 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. A store that keeps memories outside SQL answers by id
|
||||
# through the store; normalize its records to the same UUID-keyed dict shape the SQL rows
|
||||
# have so the result-building below stays store-agnostic.
|
||||
from ..memories import get_memories
|
||||
|
||||
_store = get_memories()
|
||||
if _store.writes_memory_rows_in_sql:
|
||||
memories = await conn.fetch(
|
||||
f"""
|
||||
SELECT id, text, chunk_id, document_id, fact_type, context
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE id = ANY($1) AND bank_id = $2
|
||||
""",
|
||||
valid_uuids,
|
||||
bank_id,
|
||||
)
|
||||
else:
|
||||
stored = await _store.get_memories(
|
||||
conn=conn, fq_table=fq_table, bank_id=bank_id, unit_ids=[str(u) for u in valid_uuids]
|
||||
)
|
||||
memories = [
|
||||
{
|
||||
"id": uuid.UUID(s.unit_id),
|
||||
"text": s.text,
|
||||
"chunk_id": s.chunk_id,
|
||||
"document_id": s.document_id,
|
||||
"fact_type": s.fact_type,
|
||||
"context": s.context,
|
||||
}
|
||||
for s in stored
|
||||
]
|
||||
# Batch fetch all memory units
|
||||
memories = await conn.fetch(
|
||||
f"""
|
||||
SELECT id, text, chunk_id, document_id, fact_type, context
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE id = ANY($1) AND bank_id = $2
|
||||
""",
|
||||
valid_uuids,
|
||||
bank_id,
|
||||
)
|
||||
memory_map = {row["id"]: row for row in memories}
|
||||
|
||||
# Collect chunk_ids and document_ids for batch fetching
|
||||
@@ -421,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(
|
||||
|
||||
@@ -6,14 +6,13 @@ import json
|
||||
import logging
|
||||
import uuid
|
||||
from dataclasses import dataclass
|
||||
from datetime import UTC, datetime
|
||||
from typing import TypedDict
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from ..._vector_index import index_using_clause, uses_per_bank_vector_indexes
|
||||
from ...config import get_config
|
||||
from ..db_utils import acquire_with_retry, retry_with_backoff
|
||||
from ..db_utils import acquire_with_retry
|
||||
from ..memory_engine import fq_table, get_current_schema
|
||||
from ..response_models import DispositionTraits
|
||||
|
||||
@@ -189,19 +188,8 @@ async def get_or_create_bank_profile(pool, bank_id: str) -> BankProfileResult:
|
||||
or rolls back atomically with the caller's write), use
|
||||
``get_or_create_bank_profile_on_conn`` instead.
|
||||
"""
|
||||
|
||||
# A fresh bank builds its per-(bank, fact_type) partial vector indexes with
|
||||
# a plain CREATE INDEX (it must — this runs inside the bank-create tx, and
|
||||
# CONCURRENTLY cannot). That CREATE takes a ShareLock on the shared
|
||||
# memory_units table, which can deadlock with concurrent writers. The build
|
||||
# is idempotent (INSERT ... ON CONFLICT + CREATE INDEX IF NOT EXISTS), so a
|
||||
# transient deadlock (40P01 / ORA-00060) is safe to retry as a whole tx.
|
||||
async def _create() -> BankProfileResult:
|
||||
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 retry_with_backoff(_create)
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
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:
|
||||
@@ -413,33 +401,15 @@ Merged mission:"""
|
||||
return {"mission": merged}
|
||||
|
||||
|
||||
# Sort floor for banks that have never been written to and carry no created_at.
|
||||
_UNIX_EPOCH = datetime(1970, 1, 1, tzinfo=UTC)
|
||||
|
||||
|
||||
def _as_utc(ts: datetime | None) -> datetime | None:
|
||||
"""Normalize a DB timestamp to an aware UTC datetime so values stay comparable."""
|
||||
if ts is None:
|
||||
return None
|
||||
return ts if ts.tzinfo is not None else ts.replace(tzinfo=UTC)
|
||||
|
||||
|
||||
async def list_banks(pool) -> list:
|
||||
"""
|
||||
List all banks in the system with summary stats.
|
||||
|
||||
``last_document_at`` is document *ingestion* time (when a document first
|
||||
landed), while ``last_write_at`` is the last time anything was written to
|
||||
the bank — a document re-retained/appended to, or a fact stored. Appending
|
||||
to a long-lived document does not move ``last_document_at``, which is why
|
||||
the two differ and why UIs showing "last write" must use ``last_write_at``.
|
||||
|
||||
Args:
|
||||
pool: Database connection pool
|
||||
|
||||
Returns:
|
||||
List of dicts with bank info and stats (fact_count, last_document_at, last_write_at),
|
||||
most recently written bank first.
|
||||
List of dicts with bank info and stats (document_count, fact_count, last_event_at)
|
||||
"""
|
||||
banks_table = fq_table("banks")
|
||||
docs_table = fq_table("documents")
|
||||
@@ -452,72 +422,41 @@ async def list_banks(pool) -> list:
|
||||
b.bank_id, b.name, b.disposition, b.mission,
|
||||
b.created_at, b.updated_at,
|
||||
COALESCE(m.fact_count, 0) AS fact_count,
|
||||
d.last_document_at,
|
||||
d.last_document_write_at,
|
||||
m.last_fact_at
|
||||
d.last_document_at
|
||||
FROM {banks_table} b
|
||||
LEFT JOIN (
|
||||
SELECT bank_id,
|
||||
MAX(created_at) AS last_document_at,
|
||||
MAX(updated_at) AS last_document_write_at
|
||||
SELECT bank_id, MAX(created_at) AS last_document_at
|
||||
FROM {docs_table}
|
||||
GROUP BY bank_id
|
||||
) d ON d.bank_id = b.bank_id
|
||||
LEFT JOIN (
|
||||
SELECT bank_id,
|
||||
COUNT(*) AS fact_count,
|
||||
MAX(created_at) AS last_fact_at
|
||||
SELECT bank_id, COUNT(*) AS fact_count
|
||||
FROM {mu_table}
|
||||
GROUP BY bank_id
|
||||
) m ON m.bank_id = b.bank_id
|
||||
ORDER BY b.bank_id
|
||||
ORDER BY d.last_document_at DESC NULLS LAST, b.updated_at DESC
|
||||
"""
|
||||
)
|
||||
|
||||
result = []
|
||||
# Banks are ordered by last write in Python rather than SQL: GREATEST() has
|
||||
# different NULL semantics on PostgreSQL vs Oracle, and the bank list is small.
|
||||
sort_keys: dict[str, datetime] = {}
|
||||
# A store that keeps memories outside SQL leaves the memory_units join empty, so its
|
||||
# per-bank fact_count comes from the store instead (one live count per bank).
|
||||
from ..memories import get_memories
|
||||
|
||||
_store = get_memories()
|
||||
|
||||
for row in rows:
|
||||
disposition_data = row["disposition"]
|
||||
if isinstance(disposition_data, str):
|
||||
disposition_data = json.loads(disposition_data)
|
||||
|
||||
last_doc = _as_utc(row["last_document_at"])
|
||||
created_at = _as_utc(row["created_at"])
|
||||
updated_at = _as_utc(row["updated_at"])
|
||||
# Last write = newest of "a document was (re-)retained" and "a fact was stored".
|
||||
# Appending to an existing document only bumps documents.updated_at, and facts
|
||||
# written outside a retain (consolidation, curation, import) only bump memory_units.
|
||||
write_times = [t for t in (_as_utc(row["last_document_write_at"]), _as_utc(row["last_fact_at"])) if t]
|
||||
last_write = max(write_times) if write_times else None
|
||||
last_doc = row["last_document_at"]
|
||||
|
||||
fact_count = row["fact_count"]
|
||||
if not _store.writes_memory_rows_in_sql:
|
||||
fact_count = sum(
|
||||
(await _store.count_memories(conn=conn, fq_table=fq_table, bank_id=row["bank_id"])).values()
|
||||
)
|
||||
|
||||
sort_keys[row["bank_id"]] = last_write or created_at or _UNIX_EPOCH
|
||||
result.append(
|
||||
{
|
||||
"bank_id": row["bank_id"],
|
||||
"name": row["name"],
|
||||
"disposition": disposition_data,
|
||||
"mission": row["mission"] or "",
|
||||
"created_at": created_at.isoformat() if created_at else None,
|
||||
"updated_at": updated_at.isoformat() if updated_at else None,
|
||||
"fact_count": fact_count,
|
||||
"created_at": row["created_at"].isoformat() if row["created_at"] else None,
|
||||
"updated_at": row["updated_at"].isoformat() if row["updated_at"] else None,
|
||||
"fact_count": row["fact_count"],
|
||||
"last_document_at": last_doc.isoformat() if last_doc else None,
|
||||
"last_write_at": last_write.isoformat() if last_write else None,
|
||||
}
|
||||
)
|
||||
|
||||
result.sort(key=lambda bank: sort_keys[bank["bank_id"]], reverse=True)
|
||||
return result
|
||||
|
||||
@@ -8,7 +8,7 @@ import hashlib
|
||||
import logging
|
||||
from dataclasses import dataclass
|
||||
|
||||
from ...config import _get_raw_config
|
||||
from ...config import get_config
|
||||
from ..memory_engine import fq_table
|
||||
from .types import ChunkMetadata
|
||||
|
||||
@@ -55,88 +55,23 @@ async def load_existing_chunks(conn, bank_id: str, document_id: str) -> list[Exi
|
||||
]
|
||||
|
||||
|
||||
async def delete_chunks_by_ids(conn, chunk_ids: list[str], bank_id: str | None = None, txn=None) -> None:
|
||||
async def delete_chunks_by_ids(conn, chunk_ids: list[str]) -> None:
|
||||
"""
|
||||
Delete specific chunks by their IDs.
|
||||
|
||||
This cascades to memory_units (via FK with CASCADE delete)
|
||||
and their links.
|
||||
|
||||
``txn`` carries a cross-store write-group handle when this delete is part of a re-ingest:
|
||||
the store's tombstones must ride the same txn as the replacement writes so they commit
|
||||
(become visible) together — otherwise an aborted re-ingest could drop the old memories
|
||||
without landing the new ones.
|
||||
"""
|
||||
if not chunk_ids:
|
||||
return
|
||||
|
||||
# The chunks->memory_units FK cascade below does not reach a store that keeps memories
|
||||
# outside SQL (its memory_units is empty), so drop the memories carrying each deleted
|
||||
# chunk_id through the store — otherwise a delta re-ingest leaves the old ones as duplicates.
|
||||
from ..memories import META_CHUNK_ID, DeletePredicate, get_memories
|
||||
|
||||
_store = get_memories()
|
||||
if bank_id and not _store.writes_memory_rows_in_sql:
|
||||
for _cid in chunk_ids:
|
||||
await _store.delete_where(bank_id, DeletePredicate(metadata_equals={META_CHUNK_ID: _cid}), txn=txn)
|
||||
|
||||
# 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,
|
||||
)
|
||||
|
||||
|
||||
async def store_chunks_batch(
|
||||
conn,
|
||||
bank_id: str,
|
||||
document_id: str,
|
||||
chunks: list[ChunkMetadata],
|
||||
ops=None,
|
||||
store_document_text: bool | None = None,
|
||||
conn, bank_id: str, document_id: str, chunks: list[ChunkMetadata], ops=None
|
||||
) -> dict[int, str]:
|
||||
"""
|
||||
Store document chunks in the database.
|
||||
@@ -147,9 +82,6 @@ async def store_chunks_batch(
|
||||
document_id: Document identifier
|
||||
chunks: List of ChunkMetadata objects
|
||||
ops: DataAccessOps instance (from backend.ops)
|
||||
store_document_text: Whether to persist raw chunk text. When ``None``,
|
||||
falls back to the server-level default; callers on the retain path
|
||||
pass the per-bank resolved value.
|
||||
|
||||
Returns:
|
||||
Dictionary mapping global chunk index to chunk_id
|
||||
@@ -160,16 +92,7 @@ async def store_chunks_batch(
|
||||
# When document text storage is disabled, persist empty chunk_text (the
|
||||
# column is NOT NULL) while still computing content_hash from the real text
|
||||
# so delta-retain dedup is unaffected.
|
||||
# Fallback to the raw global default (not get_config(), which guards
|
||||
# bank-configurable fields); the retain path always passes the resolved value.
|
||||
store_text = store_document_text if store_document_text is not None else _get_raw_config().store_document_text
|
||||
# A store that owns a dedicated document store keeps the chunk TEXT there, so the
|
||||
# SQL chunks row carries only its metadata (chunk_id, index, content_hash) with empty text —
|
||||
# same shape as store_document_text=False, and idempotency is unaffected (content_hash stays).
|
||||
from ..memories import get_memories
|
||||
|
||||
if get_memories().owns_document_store:
|
||||
store_text = False
|
||||
store_text = get_config().store_document_text
|
||||
|
||||
# Prepare chunk data for batch insert
|
||||
chunk_ids = []
|
||||
@@ -201,3 +124,21 @@ async def store_chunks_batch(
|
||||
)
|
||||
|
||||
return chunk_id_map
|
||||
|
||||
|
||||
def map_facts_to_chunks(facts_chunk_indices: list[int], chunk_id_map: dict[int, str]) -> list[str | None]:
|
||||
"""
|
||||
Map fact chunk indices to chunk IDs.
|
||||
|
||||
Args:
|
||||
facts_chunk_indices: List of chunk indices for each fact
|
||||
chunk_id_map: Dictionary mapping chunk index to chunk_id
|
||||
|
||||
Returns:
|
||||
List of chunk_ids (same length as facts_chunk_indices)
|
||||
"""
|
||||
chunk_ids = []
|
||||
for chunk_idx in facts_chunk_indices:
|
||||
chunk_id = chunk_id_map.get(chunk_idx)
|
||||
chunk_ids.append(chunk_id)
|
||||
return chunk_ids
|
||||
|
||||
@@ -33,6 +33,36 @@ def _validate_embedding_vector(vector: list[float], *, index: int, expected_dime
|
||||
return vector
|
||||
|
||||
|
||||
def generate_embedding(
|
||||
embeddings_backend: EmbeddingsBackend, text: str, input_type: EmbeddingInputType = "document"
|
||||
) -> list[float]:
|
||||
"""
|
||||
Generate embedding for text using the provided embeddings backend.
|
||||
|
||||
Args:
|
||||
embeddings_backend: Embeddings instance to use for encoding
|
||||
text: Text to embed
|
||||
input_type: Whether text is retained document text or recall/search query text.
|
||||
|
||||
Returns:
|
||||
Embedding vector (dimension depends on embeddings backend)
|
||||
"""
|
||||
try:
|
||||
embeddings = _encode_with_input_type(embeddings_backend, [text], input_type)
|
||||
except Exception as e:
|
||||
raise Exception(f"Failed to generate embedding: {str(e)}")
|
||||
|
||||
if len(embeddings) != 1:
|
||||
raise RuntimeError(
|
||||
f"Embeddings backend returned {len(embeddings)} vectors for 1 input text; expected exact 1:1 alignment"
|
||||
)
|
||||
return _validate_embedding_vector(
|
||||
embeddings[0],
|
||||
index=0,
|
||||
expected_dimension=embeddings_backend.dimension,
|
||||
)
|
||||
|
||||
|
||||
def _encode_with_input_type(
|
||||
embeddings_backend: EmbeddingsBackend, texts: list[str], input_type: EmbeddingInputType
|
||||
) -> list[list[float]]:
|
||||
|
||||
@@ -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,27 @@ 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."""
|
||||
parts = [self.what]
|
||||
|
||||
# Add 'who' if not N/A
|
||||
if self.who and self.who.upper() != "N/A":
|
||||
parts.append(f"Involving: {self.who}")
|
||||
|
||||
# Add 'why' if not N/A
|
||||
if self.why and self.why.upper() != "N/A":
|
||||
parts.append(self.why)
|
||||
|
||||
if len(parts) == 1:
|
||||
return parts[0]
|
||||
|
||||
return " | ".join(parts)
|
||||
|
||||
|
||||
class FactExtractionResponse(BaseModel):
|
||||
@@ -236,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."""
|
||||
|
||||
@@ -352,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(
|
||||
@@ -366,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):
|
||||
@@ -399,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):
|
||||
@@ -439,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):
|
||||
@@ -730,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}"""
|
||||
|
||||
@@ -1152,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)
|
||||
@@ -1277,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,
|
||||
@@ -1371,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
|
||||
@@ -1395,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,
|
||||
@@ -1414,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:
|
||||
@@ -1427,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", [])
|
||||
|
||||
@@ -1467,8 +1392,6 @@ async def _extract_facts_from_chunk(
|
||||
# Fallback to old format if new fields not present
|
||||
if not what:
|
||||
what = get_value("factual_core")
|
||||
if not what:
|
||||
what = get_value("text")
|
||||
if not what:
|
||||
# In verbatim mode, 'what' is intentionally absent — text is backfilled from chunk
|
||||
if extraction_mode != "verbatim":
|
||||
@@ -1528,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)
|
||||
@@ -1540,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:
|
||||
@@ -1565,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")
|
||||
@@ -1578,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
|
||||
@@ -1636,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
|
||||
|
||||
@@ -1677,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(
|
||||
@@ -1729,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"
|
||||
@@ -2186,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)
|
||||
@@ -2201,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 = []
|
||||
@@ -2231,8 +2161,6 @@ async def extract_facts_from_contents_batch_api(
|
||||
what = get_value("what")
|
||||
if not what:
|
||||
what = get_value("factual_core")
|
||||
if not what:
|
||||
what = get_value("text")
|
||||
if not what:
|
||||
continue
|
||||
|
||||
@@ -2282,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)
|
||||
@@ -2294,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:
|
||||
@@ -2319,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 = []
|
||||
@@ -2412,18 +2349,15 @@ async def extract_facts_from_contents_batch_api(
|
||||
|
||||
for chunk_meta, chunk_facts in facts_by_chunk:
|
||||
content = contents[chunk_meta.content_index]
|
||||
extraction_group_start_idx = global_fact_idx
|
||||
|
||||
for fact_from_llm in chunk_facts:
|
||||
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 [], extraction_group_start_idx, len(chunk_facts)
|
||||
),
|
||||
causal_relations=_convert_causal_relations(fact_from_llm.causal_relations or [], global_fact_idx),
|
||||
content_index=chunk_meta.content_index,
|
||||
chunk_index=chunk_meta.chunk_index,
|
||||
context=content.context,
|
||||
@@ -2548,7 +2482,8 @@ async def extract_facts_from_contents(
|
||||
# Step 1: Create parallel fact extraction tasks
|
||||
fact_extraction_tasks = []
|
||||
for item in contents:
|
||||
# Call extract_facts_from_text directly (defined earlier in this file).
|
||||
# Call extract_facts_from_text directly (defined earlier in this file)
|
||||
# to avoid circular import with utils.extract_facts
|
||||
task = extract_facts_from_text(
|
||||
text=item.content,
|
||||
event_date=item.event_date,
|
||||
@@ -2606,37 +2541,40 @@ async def extract_facts_from_contents(
|
||||
fact_idx_in_content = 0
|
||||
for chunk_idx_in_content, (chunk_text, chunk_fact_count) in enumerate(chunks_from_llm):
|
||||
chunk_global_idx = chunk_start_idx + chunk_idx_in_content
|
||||
extraction_group_start_idx = global_fact_idx
|
||||
chunk_facts = facts_from_llm[fact_idx_in_content : fact_idx_in_content + chunk_fact_count]
|
||||
|
||||
for fact_from_llm in chunk_facts:
|
||||
# Convert Fact model from LLM to ExtractedFactType dataclass
|
||||
# mentioned_at is always the event_date (when the conversation/document occurred)
|
||||
extracted_fact = ExtractedFactType(
|
||||
fact_text=fact_from_llm.fact,
|
||||
fact_type=fact_from_llm.fact_type,
|
||||
entities=list(fact_from_llm.entities or []),
|
||||
# occurred_start/end: from LLM only, leave None if not provided
|
||||
occurred_start=_parse_datetime(fact_from_llm.occurred_start)
|
||||
if fact_from_llm.occurred_start
|
||||
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 [], extraction_group_start_idx, len(chunk_facts)
|
||||
),
|
||||
content_index=content_index,
|
||||
chunk_index=chunk_global_idx,
|
||||
context=content.context,
|
||||
# mentioned_at: always the event_date (when the conversation/document occurred)
|
||||
mentioned_at=content.event_date,
|
||||
metadata=content.metadata,
|
||||
tags=content.tags,
|
||||
observation_scopes=content.observation_scopes,
|
||||
)
|
||||
for _ in range(chunk_fact_count):
|
||||
if fact_idx_in_content < len(facts_from_llm):
|
||||
fact_from_llm = facts_from_llm[fact_idx_in_content]
|
||||
|
||||
extracted_facts.append(extracted_fact)
|
||||
global_fact_idx += 1
|
||||
fact_idx_in_content += 1
|
||||
# Convert Fact model from LLM to ExtractedFactType dataclass
|
||||
# mentioned_at is always the event_date (when the conversation/document occurred)
|
||||
extracted_fact = ExtractedFactType(
|
||||
fact_text=fact_from_llm.fact,
|
||||
fact_type=fact_from_llm.fact_type,
|
||||
entities=[e.text for e in (fact_from_llm.entities or [])],
|
||||
# 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
|
||||
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
|
||||
),
|
||||
content_index=content_index,
|
||||
chunk_index=chunk_global_idx,
|
||||
context=content.context,
|
||||
# mentioned_at: always the event_date (when the conversation/document occurred)
|
||||
mentioned_at=content.event_date,
|
||||
metadata=content.metadata,
|
||||
tags=content.tags,
|
||||
observation_scopes=content.observation_scopes,
|
||||
)
|
||||
|
||||
extracted_facts.append(extracted_fact)
|
||||
global_fact_idx += 1
|
||||
fact_idx_in_content += 1
|
||||
|
||||
# Step 4: For verbatim mode, collapse to one fact per chunk with original text
|
||||
if config.retain_extraction_mode == "verbatim":
|
||||
@@ -2688,9 +2626,7 @@ def _parse_datetime(date_str: str):
|
||||
return None
|
||||
|
||||
|
||||
def _convert_causal_relations(
|
||||
relations_from_llm, extraction_group_start_idx: int, extraction_group_size: int
|
||||
) -> list[CausalRelationType]:
|
||||
def _convert_causal_relations(relations_from_llm, fact_start_idx: int) -> list[CausalRelationType]:
|
||||
"""
|
||||
Convert causal relations from LLM format to ExtractedFact format.
|
||||
|
||||
@@ -2698,16 +2634,9 @@ def _convert_causal_relations(
|
||||
"""
|
||||
causal_relations = []
|
||||
for rel in relations_from_llm:
|
||||
target_fact_index = rel.target_fact_index
|
||||
if (
|
||||
not isinstance(target_fact_index, int)
|
||||
or isinstance(target_fact_index, bool)
|
||||
or not 0 <= target_fact_index < extraction_group_size
|
||||
):
|
||||
continue
|
||||
causal_relation = CausalRelationType(
|
||||
relation_type=rel.relation_type,
|
||||
target_fact_index=extraction_group_start_idx + target_fact_index,
|
||||
target_fact_index=fact_start_idx + rel.target_fact_index,
|
||||
)
|
||||
causal_relations.append(causal_relation)
|
||||
return causal_relations
|
||||
|
||||
@@ -9,7 +9,7 @@ import logging
|
||||
import uuid
|
||||
from datetime import datetime
|
||||
|
||||
from ...config import _get_raw_config
|
||||
from ...config import get_config
|
||||
from ..memory_engine import fq_table
|
||||
from .bank_utils import DEFAULT_DISPOSITION, create_bank_vector_indexes
|
||||
from .fact_extraction import _sanitize_text
|
||||
@@ -17,11 +17,6 @@ from .types import ProcessedFact
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
#: Page size for walking a replaced document's outgoing memories. Large enough
|
||||
#: that one page covers any ordinary document, small enough that a pathological
|
||||
#: one does not arrive as a single result set.
|
||||
_OUTGOING_PAGE = 500
|
||||
|
||||
|
||||
async def get_document_content(
|
||||
conn,
|
||||
@@ -40,48 +35,17 @@ async def get_document_content(
|
||||
return row
|
||||
|
||||
|
||||
async def count_document_memory_units(
|
||||
conn,
|
||||
bank_id: str,
|
||||
document_id: str,
|
||||
) -> int:
|
||||
"""Count the memory units a document currently owns.
|
||||
|
||||
This is the number reported as ``memory_unit_count`` by the Documents API and
|
||||
by the ``retain.completed`` webhook. Zero means the document is stored but
|
||||
unreachable through recall/reflect — only memory units carry embeddings, so a
|
||||
document without them cannot be retrieved until it is reprocessed (#3040).
|
||||
"""
|
||||
count = await conn.fetchval(
|
||||
f"SELECT COUNT(*) FROM {fq_table('memory_units')} WHERE bank_id = $1 AND document_id = $2",
|
||||
bank_id,
|
||||
document_id,
|
||||
)
|
||||
return int(count or 0)
|
||||
|
||||
|
||||
async def insert_facts_batch(
|
||||
conn,
|
||||
bank_id: str,
|
||||
facts: list[ProcessedFact],
|
||||
document_id: str | None = None,
|
||||
ops=None,
|
||||
defer_index: bool = False,
|
||||
txn=None,
|
||||
conn, bank_id: str, facts: list[ProcessedFact], document_id: str | None = None, ops=None
|
||||
) -> list[str]:
|
||||
"""
|
||||
Store facts and return their unit ids, in order.
|
||||
Insert facts into the database in batch.
|
||||
|
||||
Args:
|
||||
conn: Database connection
|
||||
bank_id: Bank identifier
|
||||
facts: List of ProcessedFact objects to insert
|
||||
document_id: Optional document ID to associate with facts
|
||||
defer_index: Ask for ids without the write. The retain orchestrator needs
|
||||
this because it can only supply entity ids and causal edges after
|
||||
Phase-1 placeholders have been remapped onto real unit ids; it then
|
||||
calls `index_facts` with the complete picture. The Postgres store,
|
||||
whose write *is* the insert that mints the ids, ignores it.
|
||||
|
||||
Returns:
|
||||
List of unit IDs (UUIDs as strings) for the inserted facts
|
||||
@@ -89,37 +53,85 @@ async def insert_facts_batch(
|
||||
if not facts:
|
||||
return []
|
||||
|
||||
from ..memories import get_memories
|
||||
# Prepare data for batch insert
|
||||
fact_texts = []
|
||||
embeddings = []
|
||||
event_dates = []
|
||||
occurred_starts = []
|
||||
occurred_ends = []
|
||||
mentioned_ats = []
|
||||
contexts = []
|
||||
fact_types = []
|
||||
metadata_jsons = []
|
||||
chunk_ids = []
|
||||
document_ids = []
|
||||
tags_list = []
|
||||
observation_scopes_list = []
|
||||
text_signals_list = []
|
||||
|
||||
return await get_memories().insert_facts(
|
||||
conn=conn,
|
||||
ops=ops,
|
||||
bank_id=bank_id,
|
||||
facts=facts,
|
||||
document_id=document_id,
|
||||
defer_index=defer_index,
|
||||
txn=txn,
|
||||
for fact in facts:
|
||||
fact_texts.append(_sanitize_text(fact.fact_text))
|
||||
# Convert embedding to string for asyncpg vector type
|
||||
embeddings.append(str(fact.embedding))
|
||||
# event_date: Use occurred_start if available, otherwise use mentioned_at
|
||||
# This maintains backward compatibility while handling None occurred_start
|
||||
event_dates.append(fact.occurred_start if fact.occurred_start is not None else fact.mentioned_at)
|
||||
occurred_starts.append(fact.occurred_start)
|
||||
occurred_ends.append(fact.occurred_end)
|
||||
mentioned_ats.append(fact.mentioned_at)
|
||||
contexts.append(_sanitize_text(fact.context))
|
||||
fact_types.append(fact.fact_type)
|
||||
metadata_jsons.append(json.dumps(fact.metadata))
|
||||
chunk_ids.append(fact.chunk_id)
|
||||
# Use per-fact document_id if available, otherwise fallback to batch-level document_id
|
||||
document_ids.append(fact.document_id if fact.document_id else document_id)
|
||||
# Convert tags to JSON string for proper batch insertion (PostgreSQL unnest doesn't handle 2D arrays well)
|
||||
tags_list.append(json.dumps(fact.tags if fact.tags else []))
|
||||
# observation_scopes: stored as JSONB (string or 2D array), None if not provided
|
||||
observation_scopes_list.append(
|
||||
json.dumps(fact.observation_scopes) if fact.observation_scopes is not None else None
|
||||
)
|
||||
# Build text_signals: entity names + date tokens for enriched BM25 indexing
|
||||
signal_parts = []
|
||||
if fact.entities:
|
||||
signal_parts.extend(e.name for e in fact.entities)
|
||||
if fact.occurred_start:
|
||||
try:
|
||||
signal_parts.append(fact.occurred_start.strftime("%B %d %Y").lstrip("0").replace(" 0", " "))
|
||||
except (ValueError, AttributeError):
|
||||
pass
|
||||
if fact.occurred_end and fact.occurred_end != fact.occurred_start:
|
||||
try:
|
||||
signal_parts.append(fact.occurred_end.strftime("%B %d %Y").lstrip("0").replace(" 0", " "))
|
||||
except (ValueError, AttributeError):
|
||||
pass
|
||||
text_signals_list.append(" ".join(signal_parts) if signal_parts else None)
|
||||
|
||||
# Batch insert all facts — delegates to DataAccessOps which handles
|
||||
# unnest (PG) vs row-by-row (Oracle) transparently.
|
||||
config = get_config()
|
||||
|
||||
return await ops.insert_facts_batch(
|
||||
conn,
|
||||
bank_id,
|
||||
fact_texts,
|
||||
embeddings,
|
||||
event_dates,
|
||||
occurred_starts,
|
||||
occurred_ends,
|
||||
mentioned_ats,
|
||||
contexts,
|
||||
fact_types,
|
||||
metadata_jsons,
|
||||
chunk_ids,
|
||||
document_ids,
|
||||
tags_list,
|
||||
observation_scopes_list,
|
||||
text_signals_list,
|
||||
text_search_extension=config.text_search_extension,
|
||||
)
|
||||
|
||||
|
||||
async def index_facts(
|
||||
bank_id: str,
|
||||
unit_ids: list[str],
|
||||
facts: list[ProcessedFact],
|
||||
document_id: str | None = None,
|
||||
unit_entity_ids: dict[str, list[str]] | None = None,
|
||||
) -> None:
|
||||
"""Complete a deferred `insert_facts_batch`, now that the edges are known.
|
||||
|
||||
``unit_entity_ids`` is the unit→entity posting and each fact's causal
|
||||
relations are its edges; both travel with the memory for a store that owns
|
||||
them. A no-op for the Postgres store, which wrote all of it already.
|
||||
"""
|
||||
from ..memories import get_memories
|
||||
|
||||
await get_memories().index_facts(bank_id, unit_ids, facts, document_id, unit_entity_ids)
|
||||
|
||||
|
||||
async def ensure_bank_exists(conn, bank_id: str, ops=None) -> None:
|
||||
"""
|
||||
Ensure bank exists in the database.
|
||||
@@ -160,37 +172,95 @@ async def delete_stale_observations_for_memories(
|
||||
"""Delete observations whose source memories are about to be removed.
|
||||
|
||||
Mirrors the cleanup performed by ``MemoryEngine.delete_document`` so that
|
||||
every code path that removes memories also removes the observations derived
|
||||
from them. Without this, ingesting a fresh version of a document via the
|
||||
retain pipeline (which does a full-replace ``DELETE FROM documents``
|
||||
cascade) used to leave orphan observations pointing at memory IDs that no
|
||||
longer existed.
|
||||
every code path that removes ``memory_units`` also removes the
|
||||
observations derived from them. Without this, ingesting a fresh version
|
||||
of a document via the retain pipeline (which does a full-replace
|
||||
``DELETE FROM documents`` cascade) used to leave orphan observations
|
||||
pointing at memory IDs that no longer existed.
|
||||
|
||||
For each observation referencing any of ``fact_ids``:
|
||||
1. Delete the observation (its text is stale once even one source memory
|
||||
disappears).
|
||||
2. Reset the consolidated marker on the surviving source memories so they
|
||||
get re-consolidated under fresh observations on the next run.
|
||||
1. Delete the observation row (its text is stale once even one source
|
||||
memory disappears).
|
||||
2. Reset ``consolidated_at = NULL`` on the surviving source memories so
|
||||
they get re-consolidated under fresh observations on the next run.
|
||||
|
||||
Must be called within an active transaction, before the source memories are
|
||||
deleted.
|
||||
Must be called within an active transaction, before the source memories
|
||||
are deleted.
|
||||
|
||||
Returns:
|
||||
Number of observations deleted.
|
||||
Returns the number of observations deleted.
|
||||
"""
|
||||
if not fact_ids:
|
||||
return 0
|
||||
|
||||
from ..memories import get_memories
|
||||
fact_uuids = [uuid.UUID(str(fid)) if not isinstance(fid, uuid.UUID) else fid for fid in fact_ids]
|
||||
|
||||
return await get_memories().delete_stale_observations(
|
||||
conn=conn,
|
||||
ops=ops,
|
||||
fq_table=fq_table,
|
||||
bank_id=bank_id,
|
||||
fact_ids=fact_ids,
|
||||
if ops is not None and not ops.uses_observation_sources_table:
|
||||
# PG: use native array overlap operator
|
||||
affected_obs = await conn.fetch(
|
||||
f"""
|
||||
SELECT id, source_memory_ids
|
||||
FROM {fq_table("memory_units")}
|
||||
WHERE bank_id = $1
|
||||
AND fact_type = 'observation'
|
||||
AND source_memory_ids && $2::uuid[]
|
||||
""",
|
||||
bank_id,
|
||||
fact_uuids,
|
||||
)
|
||||
else:
|
||||
# Oracle / default: use observation_sources junction table
|
||||
affected_obs = await conn.fetch(
|
||||
f"""
|
||||
SELECT mu.id, mu.source_memory_ids
|
||||
FROM {fq_table("memory_units")} mu
|
||||
WHERE mu.bank_id = $1
|
||||
AND mu.fact_type = 'observation'
|
||||
AND EXISTS (
|
||||
SELECT 1 FROM {fq_table("observation_sources")} os
|
||||
WHERE os.observation_id = mu.id
|
||||
AND os.source_id = ANY($2::uuid[])
|
||||
)
|
||||
""",
|
||||
bank_id,
|
||||
fact_uuids,
|
||||
)
|
||||
|
||||
if not affected_obs:
|
||||
return 0
|
||||
|
||||
deleted_set = {str(uid) for uid in fact_uuids}
|
||||
obs_ids = [obs["id"] for obs in affected_obs]
|
||||
seen_remaining: set[str] = set()
|
||||
remaining_source_ids: list[uuid.UUID] = []
|
||||
for obs in affected_obs:
|
||||
for src_id in obs["source_memory_ids"] or []:
|
||||
src_str = str(src_id)
|
||||
if src_str not in deleted_set and src_str not in seen_remaining:
|
||||
remaining_source_ids.append(src_id)
|
||||
seen_remaining.add(src_str)
|
||||
|
||||
await conn.execute(
|
||||
f"DELETE FROM {fq_table('memory_units')} WHERE id = ANY($1::uuid[])",
|
||||
obs_ids,
|
||||
)
|
||||
|
||||
if remaining_source_ids:
|
||||
await conn.execute(
|
||||
f"""
|
||||
UPDATE {fq_table("memory_units")}
|
||||
SET consolidated_at = NULL
|
||||
WHERE id = ANY($1::uuid[])
|
||||
AND fact_type IN ('experience', 'world')
|
||||
""",
|
||||
remaining_source_ids,
|
||||
)
|
||||
|
||||
logger.info(
|
||||
f"[OBSERVATIONS] Deleted {len(obs_ids)} observations, reset {len(remaining_source_ids)} "
|
||||
f"source memories for re-consolidation in bank {bank_id}"
|
||||
)
|
||||
return len(obs_ids)
|
||||
|
||||
|
||||
async def handle_document_tracking(
|
||||
conn,
|
||||
@@ -201,8 +271,6 @@ async def handle_document_tracking(
|
||||
retain_params: dict | None = None,
|
||||
document_tags: list[str] | None = None,
|
||||
ops=None,
|
||||
store_document_text: bool | None = None,
|
||||
txn=None,
|
||||
) -> None:
|
||||
"""
|
||||
Handle document tracking in the database (full-replace mode).
|
||||
@@ -239,29 +307,14 @@ async def handle_document_tracking(
|
||||
# frozen). Same cleanup the explicit ``delete_document`` API performs.
|
||||
preserved_created_at = None
|
||||
if is_first_batch:
|
||||
from ..memories import get_memories
|
||||
|
||||
store = get_memories()
|
||||
# Which memories the outgoing version left behind. Asked of the store
|
||||
# rather than queried here, because it is the store that knows where they
|
||||
# are. Paged to exhaustion: every one of them is about to be deleted, and
|
||||
# a document whose facts overflow one page must not keep half of them.
|
||||
existing_unit_ids: list[str] = []
|
||||
page_token = ""
|
||||
while True:
|
||||
page = await store.scan_memories(
|
||||
conn=conn,
|
||||
fq_table=fq_table,
|
||||
bank_id=bank_id,
|
||||
fact_types=["experience", "world"],
|
||||
document_id=document_id,
|
||||
limit=_OUTGOING_PAGE,
|
||||
page_token=page_token,
|
||||
)
|
||||
existing_unit_ids.extend(m.unit_id for m in page.memories)
|
||||
page_token = page.next_page_token
|
||||
if not page_token:
|
||||
break
|
||||
existing_unit_rows = await conn.fetch(
|
||||
f"""
|
||||
SELECT id FROM {fq_table("memory_units")}
|
||||
WHERE document_id = $1 AND fact_type IN ('experience', 'world')
|
||||
""",
|
||||
document_id,
|
||||
)
|
||||
existing_unit_ids = [row["id"] for row in existing_unit_rows]
|
||||
if existing_unit_ids:
|
||||
invalidated = await delete_stale_observations_for_memories(conn, bank_id, existing_unit_ids, ops=ops)
|
||||
if invalidated:
|
||||
@@ -277,14 +330,17 @@ async def handle_document_tracking(
|
||||
if ops is not None:
|
||||
from ..graph_maintenance import enqueue_relink_victims
|
||||
|
||||
await enqueue_relink_victims(conn, bank_id, [str(uid) for uid in existing_unit_ids])
|
||||
|
||||
await enqueue_relink_victims(conn, bank_id, [str(uid) for uid in existing_unit_ids], ops=ops)
|
||||
# Explicitly delete memory_units by document_id BEFORE deleting the
|
||||
# document row. The CASCADE from documents→chunks→memory_units only
|
||||
# catches units that have a non-NULL chunk_id FK. Units with chunk_id=NULL
|
||||
# (e.g. from partial writes or edge cases) would survive the cascade.
|
||||
# This explicit delete ensures complete cleanup.
|
||||
await store.delete_document(conn=conn, fq_table=fq_table, bank_id=bank_id, document_id=document_id, txn=txn)
|
||||
await conn.execute(
|
||||
f"DELETE FROM {fq_table('memory_units')} WHERE document_id = $1 AND bank_id = $2",
|
||||
document_id,
|
||||
bank_id,
|
||||
)
|
||||
# Capture created_at before deletion so re-ingestion preserves it.
|
||||
preserved_created_at = await conn.fetchval(
|
||||
f"DELETE FROM {fq_table('documents')} WHERE id = $1 AND bank_id = $2 RETURNING created_at",
|
||||
@@ -302,7 +358,6 @@ async def handle_document_tracking(
|
||||
retain_params,
|
||||
document_tags,
|
||||
preserved_created_at=preserved_created_at,
|
||||
store_document_text=store_document_text,
|
||||
)
|
||||
|
||||
|
||||
@@ -313,7 +368,6 @@ async def upsert_document_metadata(
|
||||
combined_content: str,
|
||||
retain_params: dict | None = None,
|
||||
document_tags: list[str] | None = None,
|
||||
store_document_text: bool | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
Update document metadata without deleting existing facts/chunks.
|
||||
@@ -326,16 +380,7 @@ async def upsert_document_metadata(
|
||||
combined_content = _sanitize_text(combined_content) or ""
|
||||
content_hash = hashlib.sha256(combined_content.encode()).hexdigest()
|
||||
|
||||
await _upsert_document_row(
|
||||
conn,
|
||||
bank_id,
|
||||
document_id,
|
||||
combined_content,
|
||||
content_hash,
|
||||
retain_params,
|
||||
document_tags,
|
||||
store_document_text=store_document_text,
|
||||
)
|
||||
await _upsert_document_row(conn, bank_id, document_id, combined_content, content_hash, retain_params, document_tags)
|
||||
|
||||
|
||||
async def _upsert_document_row(
|
||||
@@ -347,7 +392,6 @@ async def _upsert_document_row(
|
||||
retain_params: dict | None = None,
|
||||
document_tags: list[str] | None = None,
|
||||
preserved_created_at: datetime | None = None,
|
||||
store_document_text: bool | None = None,
|
||||
) -> None:
|
||||
"""Insert or update a document row.
|
||||
|
||||
@@ -359,20 +403,8 @@ async def _upsert_document_row(
|
||||
When ``store_document_text`` is disabled, the raw source text
|
||||
is dropped and ``original_text`` is stored as NULL. The ``content_hash`` is
|
||||
still computed from the real content so delta-retain dedup is unaffected.
|
||||
``store_document_text`` defaults to the server-level config when ``None``;
|
||||
the retain path passes the per-bank resolved value.
|
||||
"""
|
||||
# Fallback to the raw global default (not get_config(), which guards
|
||||
# bank-configurable fields); the retain path always passes the resolved value.
|
||||
store_text = store_document_text if store_document_text is not None else _get_raw_config().store_document_text
|
||||
original_text = combined_content if store_text else None
|
||||
# A store that owns a dedicated document store keeps the extracted text there, so the
|
||||
# SQL documents row holds only its metadata (id, content_hash, tags) with original_text NULL —
|
||||
# the bulky body is written to the store up front (orchestrator._store_document_bodies).
|
||||
from ..memories import get_memories
|
||||
|
||||
if get_memories().owns_document_store:
|
||||
original_text = None
|
||||
original_text = combined_content if get_config().store_document_text else None
|
||||
await conn.execute(
|
||||
f"""
|
||||
INSERT INTO {fq_table("documents")} (id, bank_id, original_text, content_hash, retain_params, tags, created_at, updated_at)
|
||||
@@ -408,20 +440,6 @@ async def update_memory_units_tags(
|
||||
Returns:
|
||||
Number of memory units updated.
|
||||
"""
|
||||
from ..memories import MemoryPatch, get_memories
|
||||
|
||||
store = get_memories()
|
||||
if not store.writes_memory_rows_in_sql:
|
||||
# A store that keeps memories outside SQL: page the document's memories and patch each
|
||||
# one's tags through the store — the UPDATE below is a no-op on its empty memory_units.
|
||||
page = await store.scan_memories(
|
||||
conn=conn, fq_table=fq_table, bank_id=bank_id, document_id=document_id, limit=1_000_000
|
||||
)
|
||||
patches = [MemoryPatch(unit_id=m.unit_id, tags=list(tags or [])) for m in page.memories]
|
||||
if patches:
|
||||
await store.update_memories(bank_id, patches)
|
||||
return len(patches)
|
||||
|
||||
result = await conn.execute(
|
||||
f"""
|
||||
UPDATE {fq_table("memory_units")}
|
||||
|
||||
@@ -37,7 +37,6 @@ async def create_semantic_links_batch(
|
||||
bank_id: str,
|
||||
unit_ids: list[str],
|
||||
embeddings: list[list[float]],
|
||||
threshold: float,
|
||||
pre_computed_ann_links: list[tuple] | None = None,
|
||||
ops=None,
|
||||
) -> int:
|
||||
@@ -53,7 +52,6 @@ async def create_semantic_links_batch(
|
||||
bank_id: Bank identifier
|
||||
unit_ids: List of unit IDs to create links for
|
||||
embeddings: List of embedding vectors (same length as unit_ids)
|
||||
threshold: Minimum cosine similarity for semantic links
|
||||
pre_computed_ann_links: Pre-computed ANN results from Phase 1
|
||||
|
||||
Returns:
|
||||
@@ -66,14 +64,7 @@ async def create_semantic_links_batch(
|
||||
raise ValueError(f"Mismatch between unit_ids ({len(unit_ids)}) and embeddings ({len(embeddings)})")
|
||||
|
||||
return await link_utils.create_semantic_links_batch(
|
||||
conn,
|
||||
bank_id,
|
||||
unit_ids,
|
||||
embeddings,
|
||||
threshold=threshold,
|
||||
log_buffer=[],
|
||||
pre_computed_ann_links=pre_computed_ann_links,
|
||||
ops=ops,
|
||||
conn, bank_id, unit_ids, embeddings, log_buffer=[], pre_computed_ann_links=pre_computed_ann_links, ops=ops
|
||||
)
|
||||
|
||||
|
||||
@@ -83,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
|
||||
@@ -101,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)
|
||||
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user