Compare commits

..
Author SHA1 Message Date
Nicolò Boschi 50c38f8a53 Add Python client get_version helper
Adds HindsightClient.get_version()/aget_version() convenience wrappers for
the existing /version endpoint, re-exports VersionResponse for typed callers,
and tests both paths against a mocked MonitoringApi.

Python parity for #2252 (TypeScript getVersion). Fixes #2248.
2026-06-17 11:23:02 +02:00
626 changed files with 13408 additions and 25522 deletions
+1 -22
View File
@@ -2,7 +2,7 @@
# Copy this file to .env and fill in your values
# LLM Configuration (Required)
# Supported providers: openai, groq, ollama, gemini, anthropic, lmstudio, vertexai, minimax, deepseek, zai, atlas, volcano
# Supported providers: openai, groq, ollama, gemini, anthropic, lmstudio, vertexai, minimax, deepseek, zai, volcano
HINDSIGHT_API_LLM_PROVIDER=openai
HINDSIGHT_API_LLM_API_KEY=your-api-key-here
HINDSIGHT_API_LLM_MODEL=gpt-4o-mini
@@ -37,11 +37,6 @@ HINDSIGHT_API_LLM_BASE_URL=https://api.openai.com/v1
# HINDSIGHT_API_LLM_API_KEY=your-zai-api-key
# HINDSIGHT_API_LLM_MODEL=glm-4.5-flash # or glm-4.5-air for the paid tier
# Example: Atlas Cloud configuration (OpenAI-compatible, https://www.atlascloud.ai)
# HINDSIGHT_API_LLM_PROVIDER=atlas
# HINDSIGHT_API_LLM_API_KEY=your-atlascloud-api-key
# HINDSIGHT_API_LLM_MODEL=deepseek-ai/deepseek-v4-pro # reasoning model; also Qwen / GLM / Kimi / MiniMax, etc.
# Example: LM Studio local configuration (Qwen 2.5 32B recommended)
# HINDSIGHT_API_LLM_PROVIDER=lmstudio
# HINDSIGHT_API_LLM_API_KEY=lmstudio
@@ -92,18 +87,6 @@ HINDSIGHT_API_LOG_LEVEL=info
# korean_lindera/lindera(korean), ngram(min,max), edge_ngram(min,max)
# HINDSIGHT_API_TEXT_SEARCH_EXTENSION_PG_SEARCH_TOKENIZER=
# File Parser (Optional - uses markitdown by default)
# HINDSIGHT_API_FILE_PARSER=markitdown
# Enable image OCR for MarkItDown using an OpenAI-compatible OCR/vision endpoint.
# These OCR settings are independent from HINDSIGHT_API_LLM_* because MarkItDown
# uses the OpenAI SDK directly and requires Chat Completions image input support.
# When OCR is enabled, API_KEY, BASE_URL, and MODEL are required.
# HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_ENABLED=false
# HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_API_KEY=
# HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_BASE_URL=
# HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_MODEL=
# HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_PROMPT=
# 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
@@ -167,10 +150,6 @@ HINDSIGHT_API_LOG_LEVEL=info
# Custom service name and environment (optional, defaults: hindsight-api, development)
# HINDSIGHT_API_OTEL_SERVICE_NAME=hindsight-production
# HINDSIGHT_API_OTEL_DEPLOYMENT_ENVIRONMENT=production
#
# 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
# -----------------------------------------------------------------------------
# Control Plane (Optional)
+6
View File
@@ -0,0 +1,6 @@
version: 2
updates:
- package-ecosystem: "github-actions"
directory: "/"
schedule:
interval: "weekly"
+2 -2
View File
@@ -266,7 +266,7 @@ jobs:
strategy:
matrix:
include:
- os: ubuntu-22.04
- os: ubuntu-latest
target: x86_64-unknown-linux-gnu
artifact_name: hindsight
asset_name: hindsight-linux-amd64
@@ -278,7 +278,7 @@ jobs:
target: aarch64-apple-darwin
artifact_name: hindsight
asset_name: hindsight-darwin-arm64
- os: ubuntu-22.04-arm
- os: ubuntu-24.04-arm
target: aarch64-unknown-linux-gnu
artifact_name: hindsight
asset_name: hindsight-linux-arm64
-120
View File
@@ -45,20 +45,17 @@ jobs:
integrations-pydantic-ai: ${{ steps.filter.outputs.integrations-pydantic-ai }}
integrations-ag2: ${{ steps.filter.outputs.integrations-ag2 }}
integrations-autogen: ${{ steps.filter.outputs.integrations-autogen }}
integrations-aider: ${{ steps.filter.outputs.integrations-aider }}
integrations-langgraph: ${{ steps.filter.outputs.integrations-langgraph }}
integrations-llamaindex: ${{ steps.filter.outputs.integrations-llamaindex }}
integrations-paperclip: ${{ steps.filter.outputs.integrations-paperclip }}
integrations-opencode: ${{ steps.filter.outputs.integrations-opencode }}
integrations-cursor: ${{ steps.filter.outputs.integrations-cursor }}
integrations-zed: ${{ steps.filter.outputs.integrations-zed }}
integrations-n8n: ${{ steps.filter.outputs.integrations-n8n }}
integrations-zapier: ${{ steps.filter.outputs.integrations-zapier }}
integrations-cloudflare-oauth-proxy: ${{ steps.filter.outputs.integrations-cloudflare-oauth-proxy }}
integrations-superagent: ${{ steps.filter.outputs.integrations-superagent }}
integrations-lockfiles: ${{ steps.filter.outputs.integrations-lockfiles }}
integrations-openai-agents: ${{ steps.filter.outputs.integrations-openai-agents }}
integrations-openhands: ${{ steps.filter.outputs.integrations-openhands }}
integrations-pipecat: ${{ steps.filter.outputs.integrations-pipecat }}
integrations-agentcore: ${{ steps.filter.outputs.integrations-agentcore }}
integrations-smolagents: ${{ steps.filter.outputs.integrations-smolagents }}
@@ -157,8 +154,6 @@ jobs:
- 'hindsight-integrations/ag2/**'
integrations-autogen:
- 'hindsight-integrations/autogen/**'
integrations-aider:
- 'hindsight-integrations/aider/**'
integrations-langgraph:
- 'hindsight-integrations/langgraph/**'
integrations-llamaindex:
@@ -171,8 +166,6 @@ jobs:
- 'hindsight-integrations/opencode/**'
integrations-cursor:
- 'hindsight-integrations/cursor/**'
integrations-zed:
- 'hindsight-integrations/zed/**'
integrations-n8n:
- 'hindsight-integrations/n8n/**'
integrations-zapier:
@@ -187,8 +180,6 @@ jobs:
- 'scripts/check-integration-lockfiles.sh'
integrations-openai-agents:
- 'hindsight-integrations/openai-agents/**'
integrations-openhands:
- 'hindsight-integrations/openhands/**'
integrations-pipecat:
- 'hindsight-integrations/pipecat/**'
integrations-agentcore:
@@ -497,37 +488,6 @@ jobs:
working-directory: ./hindsight-integrations/cursor
run: python -m pytest tests/ -v
test-zed-integration:
needs: [detect-changes]
if: >-
github.event_name != 'pull_request_review' &&
(github.event_name == 'workflow_dispatch' ||
needs.detect-changes.outputs.integrations-zed == 'true' ||
needs.detect-changes.outputs.ci == 'true')
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v6
with:
ref: ${{ github.event.pull_request.head.sha || '' }}
- name: Set up Python
uses: actions/setup-python@v6
with:
python-version: '3.11'
- name: Install package and pytest
working-directory: ./hindsight-integrations/zed
# Installs the package (incl. the zstandard runtime dep) so the threads.db
# reader tests can decompress Zed's zstd blobs.
run: pip install -e . pytest
- name: Run tests
working-directory: ./hindsight-integrations/zed
# PR CI runs only the deterministic bucket; the real-LLM E2E bucket
# (requires_real_llm) needs a live Hindsight server and runs separately.
run: python -m pytest tests/ -v -m "not requires_real_llm"
test-omo-integration:
needs: [detect-changes]
if: >-
@@ -3061,45 +3021,6 @@ jobs:
working-directory: ./hindsight-integrations/ag2
run: uv run pytest tests -v
test-aider-integration:
needs: [detect-changes]
if: >-
(github.event_name == 'workflow_dispatch' ||
needs.detect-changes.outputs.integrations-aider == '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 aider integration
working-directory: ./hindsight-integrations/aider
run: uv build
- name: Install dependencies
working-directory: ./hindsight-integrations/aider
run: uv sync --frozen
- name: Run tests
working-directory: ./hindsight-integrations/aider
# PR CI runs only the deterministic bucket; the real-LLM E2E bucket
# (requires_real_llm) needs a live Hindsight server and runs separately.
run: uv run pytest tests -v -m "not requires_real_llm"
test-autogen-integration:
needs: [detect-changes]
if: >-
@@ -3748,45 +3669,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-openhands-integration:
needs: [detect-changes]
if: >-
(github.event_name == 'workflow_dispatch' ||
needs.detect-changes.outputs.integrations-openhands == '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 openhands integration
working-directory: ./hindsight-integrations/openhands
run: uv build
- name: Install dependencies
working-directory: ./hindsight-integrations/openhands
run: uv sync --frozen
- name: Run tests
working-directory: ./hindsight-integrations/openhands
# PR CI runs only the deterministic bucket; the real-LLM E2E bucket
# (requires_real_llm) needs a live Hindsight server and runs separately.
run: uv run pytest tests -v -m "not requires_real_llm"
test-claude-agent-sdk-integration:
needs: [detect-changes]
if: >-
@@ -4821,7 +4703,6 @@ jobs:
- test-openclaw-integration
- test-integration
- test-ag2-integration
- test-aider-integration
- test-autogen-integration
- test-continue-integration
- test-smolagents-integration
@@ -4836,7 +4717,6 @@ jobs:
- test-pydantic-ai-integration
- test-llamaindex-integration
- test-openai-agents-integration
- test-openhands-integration
- test-agentcore-integration
- test-haystack-integration
- test-pip-slim
+1 -37
View File
@@ -16,42 +16,6 @@
---
### Powered by Atlas Cloud (OpenAI-compatible)
<p align="center">
<a href="https://www.atlascloud.ai/?utm_source=github&utm_medium=link&utm_campaign=hindsight">
<img src="./hindsight-docs/static/img/atlas-cloud-logo.png" alt="Atlas Cloud" width="200">
</a>
</p>
> 🎁 **[Atlas Cloud](https://www.atlascloud.ai/?utm_source=github&utm_medium=link&utm_campaign=hindsight)** is a full-modal, OpenAI-compatible AI inference platform — plug it in as a drop-in LLM backend for Hindsight's fact extraction, reflection and consolidation, with one API for DeepSeek, Qwen, GLM, Kimi, MiniMax and more. No multi-vendor setup needed.
> Budget-friendly: [coding plan](https://www.atlascloud.ai/console/coding-plan)
```bash
export HINDSIGHT_API_LLM_PROVIDER=atlas
export HINDSIGHT_API_LLM_API_KEY=your-atlascloud-api-key # base_url defaults to https://api.atlascloud.ai/v1
export HINDSIGHT_API_LLM_MODEL=deepseek-ai/deepseek-v4-pro
```
`deepseek-ai/deepseek-v4-pro` is a reasoning model — give it enough `max_tokens` (>= 512).
<details>
<summary>All Atlas Cloud chat models (59)</summary>
- **Anthropic (Claude):** `anthropic/claude-haiku-4.5-20251001`, `anthropic/claude-opus-4.8`, `anthropic/claude-sonnet-4.6`
- **OpenAI (GPT):** `openai/gpt-5.4`, `openai/gpt-5.5`
- **Google (Gemini):** `google/gemini-3.1-flash-lite`, `google/gemini-3.1-pro-preview`, `google/gemini-3.5-flash`
- **Alibaba Qwen:** `qwen/qwen2.5-7b-instruct`, `Qwen/Qwen3-235B-A22B-Instruct-2507`, `qwen/qwen3-235b-a22b-thinking-2507`, `qwen/qwen3-30b-a3b`, `Qwen/Qwen3-30B-A3B-Instruct-2507`, `qwen/qwen3-30b-a3b-thinking-2507`, `qwen/qwen3-32b`, `qwen/qwen3-8b`, `Qwen/Qwen3-Coder`, `qwen/qwen3-coder-next`, `qwen/qwen3-max-2026-01-23`, `Qwen/Qwen3-Next-80B-A3B-Instruct`, `Qwen/Qwen3-Next-80B-A3B-Thinking`, `Qwen/Qwen3-VL-235B-A22B-Instruct`, `qwen/qwen3-vl-235b-a22b-thinking`, `qwen/qwen3-vl-30b-a3b-instruct`, `qwen/qwen3-vl-30b-a3b-thinking`, `qwen/qwen3-vl-8b-instruct`, `qwen/qwen3.5-122b-a10b`, `qwen/qwen3.5-27b`, `qwen/qwen3.5-35b-a3b`, `qwen/qwen3.5-397b-a17b`, `qwen/qwen3.6-35b-a3b`, `qwen/qwen3.6-plus`
- **DeepSeek:** `deepseek-ai/deepseek-ocr`, `deepseek-ai/deepseek-r1-0528`, `deepseek-ai/DeepSeek-V3-0324`, `deepseek-ai/DeepSeek-V3.1`, `deepseek-ai/DeepSeek-V3.1-Terminus`, `deepseek-ai/deepseek-v3.2`, `deepseek-ai/DeepSeek-V3.2-Exp`, `deepseek-ai/deepseek-v4-flash`, `deepseek-ai/deepseek-v4-pro`
- **Moonshot (Kimi):** `moonshotai/Kimi-K2-Instruct`, `moonshotai/Kimi-K2-Instruct-0905`, `moonshotai/Kimi-K2-Thinking`, `moonshotai/kimi-k2.5`, `moonshotai/kimi-k2.6`
- **Zhipu GLM:** `zai-org/GLM-4.6`, `zai-org/glm-4.7`, `zai-org/glm-5`, `zai-org/glm-5-turbo`, `zai-org/glm-5.1`, `zai-org/glm-5v-turbo`
- **MiniMax:** `MiniMaxAI/MiniMax-M2`, `minimaxai/minimax-m2.1`, `minimaxai/minimax-m2.5`, `minimaxai/minimax-m2.7`
- **xAI:** `xai/grok-4.3`
- **Kuaishou KAT:** `kwaipilot/kat-coder-pro-v2`
- **Other:** `owl`
</details>
## What is Hindsight?
Hindsight™ is an agent memory system built to create smarter agents that learn over time. Most agent memory systems focus on recalling conversation history. Hindsight is focused on making agents that learn, not just remember.
@@ -106,7 +70,7 @@ docker run -it --pull always --name hindsight --restart unless-stopped -p 8888:8
>API: http://localhost:8888
>UI: http://localhost:9999
You can modify the LLM provider by setting `HINDSIGHT_API_LLM_PROVIDER`. Valid options are `openai`, `anthropic`, `gemini`, `groq`, `ollama`, `lmstudio`, `minimax`, and `atlas` ([Atlas Cloud](https://www.atlascloud.ai/?utm_source=github&utm_medium=link&utm_campaign=hindsight)). The documentation provides more details on [supported models](https://hindsight.vectorize.io/developer/models).
You can modify the LLM provider by setting `HINDSIGHT_API_LLM_PROVIDER`. Valid options are `openai`, `anthropic`, `gemini`, `groq`, `ollama`, `lmstudio`, and `minimax`. The documentation provides more details on [supported models](https://hindsight.vectorize.io/developer/models).
+2 -2
View File
@@ -2,8 +2,8 @@ apiVersion: v2
name: hindsight
description: Hindsight helm chart
type: application
version: 0.8.3
appVersion: "0.8.3"
version: 0.8.2
appVersion: "0.8.2"
keywords:
- ai
- memory
+1 -1
View File
@@ -1,6 +1,6 @@
{
"name": "@vectorize-io/hindsight-all",
"version": "0.8.3",
"version": "0.8.2",
"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",
+2 -2
View File
@@ -4,12 +4,12 @@ build-backend = "setuptools.build_meta"
[project]
name = "hindsight-all-slim"
version = "0.8.3"
version = "0.8.2"
description = "Hindsight: Agent Memory That Works Like Human Memory - Slim All-in-One Bundle"
readme = "README.md"
requires-python = ">=3.11"
dependencies = [
"hindsight-api-slim==0.8.3",
"hindsight-api-slim==0.8.2",
"hindsight-client>=0.0.7",
"hindsight-embed>=0.1.0",
]
+3 -3
View File
@@ -4,12 +4,12 @@ build-backend = "hatchling.build"
[project]
name = "hindsight-all"
version = "0.8.3"
version = "0.8.2"
description = "Hindsight: Agent Memory That Works Like Human Memory - All-in-One Bundle"
readme = "README.md"
requires-python = ">=3.11"
dependencies = [
"hindsight-api-slim[all]==0.8.3",
"hindsight-api-slim[all]==0.8.2",
"hindsight-client>=0.0.7",
"hindsight-embed>=0.1.0",
]
@@ -21,7 +21,7 @@ hindsight-embed = { workspace = true }
[project.optional-dependencies]
local-llm = [
"hindsight-api-slim[local-llm]==0.8.3",
"hindsight-api-slim[local-llm]==0.8.2",
]
test = [
"pytest>=7.0.0",
+1 -1
View File
@@ -53,4 +53,4 @@ __all__ = [
"RemoteTEICrossEncoder",
"LLMConfig",
]
__version__ = "0.8.3"
__version__ = "0.8.2"
+1 -17
View File
@@ -256,7 +256,6 @@ async def _run_migration(
schema: str | None = None,
base_schema: str = DEFAULT_DATABASE_SCHEMA,
embedding_dimension: int | None = None,
ensure_extensions: bool = True,
) -> list[str]:
"""Resolve database URL and run migrations for one schema or all discovered schemas."""
from ..migrations import run_migrations_for_schemas
@@ -293,7 +292,7 @@ async def _run_migration(
vector_extension=config.vector_extension,
text_search_extension=config.text_search_extension,
pg_search_tokenizer=config.text_search_extension_pg_search_tokenizer,
ensure_extensions=ensure_extensions,
ensure_extensions=True,
)
return schemas
@@ -312,18 +311,6 @@ def run_db_migration(
"--embedding-dimension",
help="Expected embedding dimension to enforce after migrations. Omit to skip dimension sync.",
),
skip_extension_reconcile: bool = typer.Option(
False,
"--skip-extension-reconcile",
help=(
"Skip the post-migration vector / text-search index reconcile. This step only does "
"work when the configured backend (HINDSIGHT_API_VECTOR_EXTENSION / "
"HINDSIGHT_API_TEXT_SEARCH_EXTENSION) differs from a schema's existing indexes — a "
"rare, operator-driven change. Skipping it makes a no-change re-migration over many "
"tenant schemas much faster. Only use when you have NOT changed the backend; a "
"backend change still needs a normal run to reshape the indexes."
),
),
):
"""Run database migrations to the latest version."""
config = HindsightConfig.from_env()
@@ -337,8 +324,6 @@ def run_db_migration(
typer.echo(f"Running database migrations for schema: {schema}...")
else:
typer.echo("Running database migrations for base schema and all discovered tenant schemas...")
if skip_extension_reconcile:
typer.echo("Skipping post-migration extension reconcile (--skip-extension-reconcile).")
schemas = asyncio.run(
_run_migration(
@@ -346,7 +331,6 @@ def run_db_migration(
schema=schema,
base_schema=config.database_schema,
embedding_dimension=embedding_dimension,
ensure_extensions=not skip_extension_reconcile,
)
)
@@ -1,158 +0,0 @@
"""Make maintenance routines resilient to schemas that vanish mid-scan.
``public.banks_needing_consolidation()`` and
``public.schemas_with_expired_rows(...)`` snapshot the set of schemas owning a
target table from ``pg_class`` and then run a dynamic query against each schema
in turn. That is a time-of-check/time-of-use race: a schema (or its tables) can
be dropped — a tenant being deleted, or a tenant migration that recreates
tables — between the snapshot and the per-schema query, which then aborts the
whole routine with::
relation "<schema>.memory_units" does not exist
relation "<schema>.audit_log" does not exist
In the test suite this surfaces as cross-worker contamination: the multi-tenant
maintenance test creates and drops ~100 ``mt<hash>_NNN`` schemas while
``test_maintenance_routines`` (on another xdist worker, same DB) calls the
routines. In production the background maintenance loop hits the same race when
a tenant is removed or mid-migration.
Wrap each per-schema query in its own ``BEGIN ... EXCEPTION`` block so a schema
that disappears (``undefined_table`` / ``invalid_schema_name`` /
``undefined_column``) is skipped instead of aborting the scan. The routines stay
``CREATE OR REPLACE`` and PostgreSQL-only, and are (re)installed only on the run
that targets the shared ``public`` schema — same gating as the original
install (``e5f6a7b8c9d0``) and its repair (``b2d4f6a8c1e3``).
Revision ID: c7e9f1a3b5d2
Revises: e1f2a3b4c5d6
Create Date: 2026-06-19
"""
from collections.abc import Sequence
from alembic import context, op
from hindsight_api.alembic._dialect import run_for_dialect
revision: str = "c7e9f1a3b5d2"
down_revision: str | Sequence[str] | None = "e1f2a3b4c5d6"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
def _should_install_public_routines(target_schema: str | None) -> bool:
"""True for the run that must (re)create the shared ``public.*`` routines.
The routines physically live in ``public``, so they are installed exactly
once — on the base run (no ``target_schema``) or the run that explicitly
targets ``public``. Mirrors ``b2d4f6a8c1e3``.
"""
return not target_schema or target_schema == "public"
def _pg_upgrade() -> None:
if not _should_install_public_routines(context.config.get_main_option("target_schema")):
return
# Same body as b2d4f6a8c1e3, but each per-schema query runs in its own
# subtransaction so a schema dropped mid-scan is skipped, not fatal.
op.execute(
"""
CREATE OR REPLACE FUNCTION public.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(
"""
CREATE OR REPLACE FUNCTION public.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$;
"""
)
def _pg_downgrade() -> None:
# No-op: e5f6a7b8c9d0 owns these functions' lifecycle and drops them on its
# own downgrade. This migration only re-installs them (the resilient body is
# a strict superset of the previous behaviour), so there is nothing to undo
# without racing that migration's DROP.
pass
def upgrade() -> None:
run_for_dialect(pg=_pg_upgrade)
def downgrade() -> None:
run_for_dialect(pg=_pg_downgrade)
+23 -55
View File
@@ -158,12 +158,7 @@ from hindsight_api.engine.response_models import (
)
from hindsight_api.engine.search.tags import TagGroup, TagsMatch
from hindsight_api.extensions import HttpExtension, OperationValidationError, load_extension
from hindsight_api.metrics import (
create_metrics_collector,
get_metrics_collector,
initialize_metrics,
normalize_http_endpoint,
)
from hindsight_api.metrics import create_metrics_collector, get_metrics_collector, initialize_metrics
from hindsight_api.models import RequestContext
logger = logging.getLogger(__name__)
@@ -270,16 +265,6 @@ class RecallRequest(BaseModel):
default=None,
description="List of fact types to recall: 'world', 'experience', 'observation'. Defaults to world and experience if not specified.",
)
prefer_observations: bool = Field(
default=False,
description=(
"When recalling raw facts ('world'/'experience') together with 'observation', drop any raw "
"fact that an observation in the results was consolidated from, so the observation supersedes "
"it and you don't get duplicate content. The freed slots are backfilled with the next results, "
"keeping the result count at the requested budget. Disabled by default; set to true to enable. "
"No effect unless 'observation' and at least one raw type are both requested."
),
)
budget: Budget = Budget.MID
max_tokens: int = 4096
trace: bool = False
@@ -296,16 +281,12 @@ class RecallRequest(BaseModel):
)
tags: list[str] | None = Field(
default=None,
description="Filter memories by tags. If not specified, all memories are returned. "
"Omitting tags (or passing []) together with tags_match='exact' filters to "
"untagged/global observations only (the scope written by observation_scopes='shared').",
description="Filter memories by tags. If not specified, all memories are returned.",
)
tags_match: TagsMatch = Field(
default="any",
description="How to match tags: 'any' (OR, includes untagged), 'all' (AND, includes untagged), "
"'any_strict' (OR, excludes untagged), 'all_strict' (AND, excludes untagged), "
"'exact' (set-equality on the full scope, excludes untagged). With 'exact' and no tags "
"(or []), the empty global scope is selected and only untagged memories match.",
"'any_strict' (OR, excludes untagged), 'all_strict' (AND, excludes untagged).",
)
tag_groups: list[TagGroup] | None = Field(
default=None,
@@ -1461,13 +1442,6 @@ class DryRunExtractRequest(BaseModel):
entities_allow_free_form: bool | None = None
llm_output_language: str | None = None
@field_validator("content")
@classmethod
def validate_content(cls, v: str) -> str:
if not v.strip():
raise ValueError("content cannot be empty")
return v
class ListDocumentsResponse(BaseModel):
"""Response model for list documents endpoint."""
@@ -3263,9 +3237,15 @@ def create_app(
@app.middleware("http")
async def http_metrics_middleware(request, call_next):
"""Record HTTP request metrics."""
# Template id segments (bank ids, UUIDs, numeric ids) so the endpoint
# metric label stays bounded-cardinality.
path = normalize_http_endpoint(request.url.path)
# Normalize endpoint path to reduce cardinality
# Replace UUIDs and numeric IDs with placeholders
import re
path = request.url.path
# Replace UUIDs
path = re.sub(r"/[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}", "/{id}", path)
# Replace numeric IDs
path = re.sub(r"/\d+(?=/|$)", "/{id}", path)
status_code = [500] # Default to 500, will be updated
metrics_collector = get_metrics_collector()
@@ -3353,7 +3333,6 @@ def _register_routes(app: FastAPI):
async def _precheck_dep(
bank_id: str,
request: Request,
request_context: RequestContext = Depends(get_request_context),
) -> None:
validator = getattr(app.state.memory, "_operation_validator", None)
@@ -3362,20 +3341,10 @@ def _register_routes(app: FastAPI):
from hindsight_api.extensions import PrecheckContext
await app.state.memory._authenticate_tenant(request_context)
cl_header = request.headers.get("content-length")
content_length: int | None = None
if cl_header is not None:
try:
parsed = int(cl_header)
except ValueError:
parsed = -1
if parsed >= 0:
content_length = parsed
ctx = PrecheckContext(
operation=operation,
bank_id=bank_id,
request_context=request_context,
content_length=content_length,
)
result = await validator.precheck(ctx)
if not result.allowed:
@@ -3479,7 +3448,7 @@ def _register_routes(app: FastAPI):
async def api_graph(
bank_id: str,
type: str | None = None,
limit: int = Query(default=1000, ge=0),
limit: int = 1000,
q: str | None = None,
tags: list[str] | None = Query(None),
tags_match: str = "all_strict",
@@ -3527,8 +3496,8 @@ def _register_routes(app: FastAPI):
consolidation_state: str | None = None,
state: str | None = None,
document_id: str | None = None,
limit: int = Query(default=100, ge=0),
offset: int = Query(default=0, ge=0),
limit: int = 100,
offset: int = 0,
request_context: RequestContext = Depends(get_request_context),
):
"""
@@ -3840,7 +3809,6 @@ def _register_routes(app: FastAPI):
max_tokens=request.max_tokens,
enable_trace=request.trace,
fact_type=fact_types,
prefer_observations=request.prefer_observations,
question_date=question_date,
include_entities=include_entities,
max_entity_tokens=max_entity_tokens,
@@ -4263,8 +4231,8 @@ def _register_routes(app: FastAPI):
)
async def api_list_entities(
bank_id: str,
limit: int = Query(default=100, ge=0, description="Maximum number of entities to return"),
offset: int = Query(default=0, ge=0, description="Offset for pagination"),
limit: int = Query(default=100, description="Maximum number of entities to return"),
offset: int = Query(default=0, description="Offset for pagination"),
request_context: RequestContext = Depends(get_request_context),
):
"""List entities for a memory bank with pagination."""
@@ -4299,7 +4267,7 @@ def _register_routes(app: FastAPI):
)
async def api_entity_graph(
bank_id: str,
limit: int = Query(default=1000, ge=0, description="Maximum number of co-occurrence edges to return"),
limit: int = Query(default=1000, description="Maximum number of co-occurrence edges to return"),
min_count: int = Query(default=1, description="Minimum cooccurrence_count to include an edge"),
request_context: RequestContext = Depends(get_request_context),
):
@@ -4926,8 +4894,8 @@ def _register_routes(app: FastAPI):
tags_match: str = Query(
"any_strict", description="How to match tags: 'any', 'all', 'any_strict', 'all_strict'"
),
limit: int = Query(default=100, ge=0),
offset: int = Query(default=0, ge=0),
limit: int = 100,
offset: int = 0,
request_context: RequestContext = Depends(get_request_context),
):
"""
@@ -5111,8 +5079,8 @@ def _register_routes(app: FastAPI):
default="memories",
description="Where to read tags from: 'memories' (memory_units, default) or 'mental_models'.",
),
limit: int = Query(default=100, ge=0, description="Maximum number of tags to return"),
offset: int = Query(default=0, ge=0, description="Offset for pagination"),
limit: int = Query(default=100, description="Maximum number of tags to return"),
offset: int = Query(default=0, description="Offset for pagination"),
request_context: RequestContext = Depends(get_request_context),
):
"""
@@ -6781,7 +6749,7 @@ def _register_routes(app: FastAPI):
description="Upload files (PDF, DOCX, etc.), convert them to markdown, and retain as memories.\n\n"
"This endpoint handles file upload, conversion, and memory creation in a single operation.\n\n"
"**Features:**\n"
"- Supports PDF, DOCX, PPTX, XLSX, images (parser-dependent OCR), audio (with transcription)\n"
"- Supports PDF, DOCX, PPTX, XLSX, images (with OCR), audio (with transcription)\n"
"- Automatic file-to-markdown conversion using pluggable parsers\n"
"- Files stored in object storage (PostgreSQL by default, S3 for production)\n"
"- Each file becomes a separate document with optional metadata/tags\n"
+1 -20
View File
@@ -9,7 +9,7 @@ from fastmcp import FastMCP
from hindsight_api import MemoryEngine
from hindsight_api import __version__ as HINDSIGHT_VERSION
from hindsight_api.config import DEFAULT_MCP_RECALL_DESCRIPTION, DEFAULT_MCP_RETAIN_DESCRIPTION, _get_raw_config
from hindsight_api.config import _get_raw_config
from hindsight_api.engine.memory_engine import _current_schema
from hindsight_api.extensions import MCPExtension, load_extension
from hindsight_api.extensions.tenant import AuthenticationError
@@ -78,19 +78,6 @@ def get_current_mcp_authenticated() -> bool:
return _current_mcp_authenticated.get()
def _build_mcp_tool_descriptions(extra_instructions: str | None) -> tuple[str | None, str | None]:
"""Return custom retain/recall descriptions when server-level MCP instructions are set."""
if not isinstance(extra_instructions, str):
return None, None
extra_instructions = extra_instructions.strip()
if not extra_instructions:
return None, None
suffix = f"\n\nAdditional instructions: {extra_instructions}"
return DEFAULT_MCP_RETAIN_DESCRIPTION + suffix, DEFAULT_MCP_RECALL_DESCRIPTION + suffix
def create_mcp_server(memory: MemoryEngine, multi_bank: bool = True) -> FastMCP:
"""
Create and configure the Hindsight MCP server.
@@ -148,10 +135,6 @@ def create_mcp_server(memory: MemoryEngine, multi_bank: bool = True) -> FastMCP:
allowed = frozenset(global_config.mcp_enabled_tools)
base_tools = (base_tools if base_tools is not None else _ALL_TOOLS) & allowed
retain_description, recall_description = _build_mcp_tool_descriptions(
getattr(global_config, "mcp_instructions", None)
)
# Configure and register tools using shared module
config = MCPToolsConfig(
bank_id_resolver=get_current_bank_id,
@@ -161,8 +144,6 @@ def create_mcp_server(memory: MemoryEngine, multi_bank: bool = True) -> FastMCP:
mcp_authenticated_resolver=get_current_mcp_authenticated, # Propagate MCP pre-auth flag
include_bank_id_param=multi_bank,
tools=base_tools,
retain_description=retain_description,
recall_description=recall_description,
)
register_mcp_tools(mcp, memory, config)
-102
View File
@@ -142,7 +142,6 @@ ENV_LLM_REASONING_EFFORT = "HINDSIGHT_API_LLM_REASONING_EFFORT"
ENV_LLM_GROQ_SERVICE_TIER = "HINDSIGHT_API_LLM_GROQ_SERVICE_TIER"
ENV_LLM_OPENAI_SERVICE_TIER = "HINDSIGHT_API_LLM_OPENAI_SERVICE_TIER"
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"
ENV_LLM_STRICT_SCHEMA = "HINDSIGHT_API_LLM_STRICT_SCHEMA"
@@ -160,25 +159,11 @@ ENV_LLM_LITELLMROUTER_CONFIG = "HINDSIGHT_API_LLM_LITELLMROUTER_CONFIG"
DEFAULT_LLM_GROQ_SERVICE_TIER = "auto" # "on_demand", "flex", or "auto"
DEFAULT_LLM_OPENAI_SERVICE_TIER = None # None (default) or "flex" (50% cheaper)
DEFAULT_LLM_BEDROCK_SERVICE_TIER = None # None (default), "flex", "priority", or "reserved"
DEFAULT_LLM_GEMINI_SERVICE_TIER = None # None (default) or "flex" (50% cheaper best-effort tier)
DEFAULT_LLM_EXTRA_BODY = None # None = no extra body params; JSON dict merged into OpenAI extra_body
DEFAULT_LLM_DEFAULT_HEADERS = (
None # None = no extra headers; JSON dict passed as default_headers to provider SDK clients
)
def parse_gemini_service_tier(value: str | None) -> str | None:
"""Normalize and validate the Gemini service tier."""
tier = value or None
valid_tiers = (None, "flex")
if tier not in valid_tiers:
raise ValueError(
f"Invalid HINDSIGHT_API_LLM_GEMINI_SERVICE_TIER: "
f"{tier!r}. Must be one of: {', '.join(t for t in valid_tiers if t is not None)}."
)
return tier
# 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"
@@ -369,7 +354,6 @@ ENV_ACCESS_LOG = "HINDSIGHT_API_ACCESS_LOG"
ENV_MCP_ENABLED = "HINDSIGHT_API_MCP_ENABLED"
ENV_MCP_ENABLED_TOOLS = "HINDSIGHT_API_MCP_ENABLED_TOOLS"
ENV_MCP_STATELESS = "HINDSIGHT_API_MCP_STATELESS"
ENV_MCP_INSTRUCTIONS = "HINDSIGHT_API_MCP_INSTRUCTIONS"
ENV_ENABLE_BANK_CONFIG_API = "HINDSIGHT_API_ENABLE_BANK_CONFIG_API"
ENV_ENABLE_BANK_LLM_HEALTH = "HINDSIGHT_API_ENABLE_BANK_LLM_HEALTH"
ENV_ENABLE_DRY_RUN_EXTRACT = "HINDSIGHT_API_ENABLE_DRY_RUN_EXTRACT"
@@ -391,7 +375,6 @@ ENV_OTEL_EXPORTER_OTLP_HEADERS = "HINDSIGHT_API_OTEL_EXPORTER_OTLP_HEADERS"
ENV_OTEL_SERVICE_NAME = "HINDSIGHT_API_OTEL_SERVICE_NAME"
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"
# Vertex AI configuration
ENV_LLM_VERTEXAI_PROJECT_ID = "HINDSIGHT_API_LLM_VERTEXAI_PROJECT_ID"
@@ -441,11 +424,6 @@ ENV_FILE_STORAGE_AZURE_ACCOUNT_NAME = "HINDSIGHT_API_FILE_STORAGE_AZURE_ACCOUNT_
ENV_FILE_STORAGE_AZURE_ACCOUNT_KEY = "HINDSIGHT_API_FILE_STORAGE_AZURE_ACCOUNT_KEY"
ENV_FILE_PARSER = "HINDSIGHT_API_FILE_PARSER"
ENV_FILE_PARSER_ALLOWLIST = "HINDSIGHT_API_FILE_PARSER_ALLOWLIST"
ENV_FILE_PARSER_MARKITDOWN_OCR_ENABLED = "HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_ENABLED"
ENV_FILE_PARSER_MARKITDOWN_OCR_API_KEY = "HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_API_KEY"
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_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"
@@ -575,14 +553,6 @@ ENV_RECALL_MAX_CANDIDATES_PER_SOURCE = "HINDSIGHT_API_RECALL_MAX_CANDIDATES_PER_
# Empty disables the feature.
ENV_RECALL_STRATEGY_BOOSTS = "HINDSIGHT_API_RECALL_STRATEGY_BOOSTS"
# Recency decay used by recall reranking (engine/search/reranking.py). The decay
# function maps a memory's age onto a freshness signal that nudges its final
# ranking via a small multiplicative boost. "linear" (default) preserves the
# historical behaviour; "exponential" decays by half-life; "none" disables it.
ENV_RECENCY_DECAY_FUNCTION = "HINDSIGHT_API_RECENCY_DECAY_FUNCTION"
ENV_RECENCY_DECAY_LINEAR_WINDOW_DAYS = "HINDSIGHT_API_RECENCY_DECAY_LINEAR_WINDOW_DAYS"
ENV_RECENCY_DECAY_HALFLIFE_DAYS = "HINDSIGHT_API_RECENCY_DECAY_HALFLIFE_DAYS"
# Audit log settings
ENV_AUDIT_LOG_ENABLED = "HINDSIGHT_API_AUDIT_LOG_ENABLED"
ENV_AUDIT_LOG_ACTIONS = "HINDSIGHT_API_AUDIT_LOG_ACTIONS"
@@ -618,7 +588,6 @@ PROVIDER_DEFAULT_MODELS = {
"deepseek": "deepseek-v4-flash",
"zai": "glm-4.5-flash",
"opencode-go": "deepseek-v4-flash",
"atlas": "deepseek-ai/deepseek-v4-pro",
"ollama": "gemma3:12b",
"ollama-cloud": "gemma3:12b",
"llamacpp": "gemma-4-e2b-it",
@@ -722,14 +691,6 @@ DEFAULT_RECALL_MAX_CANDIDATES_PER_SOURCE = 0
# "graph:high,semantic:low"). Empty disables the feature. See
# ENV_RECALL_STRATEGY_BOOSTS for the full rationale.
DEFAULT_RECALL_STRATEGY_BOOSTS = ""
# Recency decay shape used by recall reranking. "linear" reproduces the
# historical straight-line decay; defaults below keep behaviour unchanged.
RECENCY_DECAY_FUNCTIONS = ("linear", "exponential", "none")
DEFAULT_RECENCY_DECAY_FUNCTION = "linear"
# Linear: days over which freshness decays from 1.0 to its 0.1 floor.
DEFAULT_RECENCY_DECAY_LINEAR_WINDOW_DAYS = 365.0
# Exponential: age (days) at which the recency signal is neutral (0.5).
DEFAULT_RECENCY_DECAY_HALFLIFE_DAYS = 90.0
# Retrieval arms that can be boosted; mirrors fusion.py source_names.
RECALL_STRATEGY_NAMES = ("semantic", "bm25", "graph", "temporal")
# User-facing priority levels. Kept in sync with recall_boost.BOOST_LEVELS by a
@@ -843,7 +804,6 @@ DEFAULT_ACCESS_LOG = False
DEFAULT_MCP_ENABLED = True
DEFAULT_MCP_ENABLED_TOOLS: list[str] | None = None # None = all tools enabled
DEFAULT_MCP_STATELESS = False # False = stateful (supports SSE/GET); True = stateless (POST-only)
DEFAULT_MCP_INSTRUCTIONS = None
DEFAULT_ENABLE_BANK_CONFIG_API = True
# Dry-run extraction is a preview tool that makes a real LLM call but stores nothing. Enabled by
# default; set HINDSIGHT_API_ENABLE_DRY_RUN_EXTRACT=false to remove the endpoint (e.g. to cap
@@ -887,10 +847,6 @@ DEFAULT_RETAIN_BATCH_POLL_INTERVAL_SECONDS = 60 # Batch API polling interval in
DEFAULT_FILE_STORAGE_TYPE = "native" # PostgreSQL BYTEA storage
DEFAULT_FILE_PARSER = "markitdown" # Default parser fallback chain (comma-separated, e.g. "iris,markitdown")
DEFAULT_FILE_PARSER_ALLOWLIST = None # Allowlist of parsers clients may request (None = all registered parsers)
DEFAULT_FILE_PARSER_MARKITDOWN_OCR_ENABLED = False
DEFAULT_FILE_PARSER_MARKITDOWN_OCR_PROMPT = """You are a precise OCR transcription engine.
Transcribe only the visible text in the image. Do not describe the image, summarize it, translate it, infer missing content, or add commentary. Preserve the original language, wording, numbers, punctuation, capitalization, and reading order. Reconstruct headings, lists, key-value fields, stamps, and tables as clean Markdown when the layout is clear. If text is unreadable or uncertain, write [unclear] for that span. Return only the extracted Markdown."""
DEFAULT_FILE_CONVERSION_MAX_BATCH_SIZE_MB = 100 # Max total batch size in MB (all files combined)
DEFAULT_FILE_CONVERSION_MAX_BATCH_SIZE = 10 # Max files per batch upload
DEFAULT_ENABLE_FILE_UPLOAD_API = True # Enable file upload endpoint
@@ -1009,7 +965,6 @@ DEFAULT_OTEL_TRACES_ENABLED = False # Disabled by default for backward compatib
DEFAULT_OTEL_SERVICE_NAME = "hindsight-api"
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
# Audit log defaults
DEFAULT_AUDIT_LOG_ENABLED = False # Disabled by default
@@ -1225,18 +1180,6 @@ def _validate_recall_budget_function(function: str) -> str:
return function_lower
def _validate_recency_decay_function(function: str) -> str:
"""Validate and normalize the recency decay function."""
function_lower = function.lower()
if function_lower not in RECENCY_DECAY_FUNCTIONS:
logger.warning(
f"Invalid recency decay function '{function}', must be one of {RECENCY_DECAY_FUNCTIONS}. "
f"Defaulting to '{DEFAULT_RECENCY_DECAY_FUNCTION}'."
)
return DEFAULT_RECENCY_DECAY_FUNCTION
return function_lower
def _parse_bank_priority(raw: str) -> dict[str, int]:
"""Parse ``bank-pattern:priority,...`` into ``{pattern: priority}``.
@@ -1355,7 +1298,6 @@ class HindsightConfig:
llm_groq_service_tier: str # Groq: "on_demand", "flex", or "auto"
llm_openai_service_tier: str | None # OpenAI: None (default) or "flex" (50% cheaper)
llm_bedrock_service_tier: str | None # Bedrock: None (default), "flex", "priority", or "reserved"
llm_gemini_service_tier: str | None # Gemini: None (default) or "flex" (50% cheaper)
llm_extra_body: (
dict | None
) # Extra body params merged into OpenAI-compatible API calls (e.g. {"chat_template_kwargs": {"enable_thinking": true}})
@@ -1493,9 +1435,6 @@ class HindsightConfig:
bm25_min_score: float
recall_max_candidates_per_source: int
recall_strategy_boosts: dict[str, str]
recency_decay_function: str
recency_decay_linear_window_days: float
recency_decay_halflife_days: float
reranker_cohere_api_key: str | None
reranker_cohere_model: str
reranker_cohere_base_url: str | None
@@ -1539,7 +1478,6 @@ class HindsightConfig:
mcp_enabled: bool
mcp_enabled_tools: list[str] | None # None = all tools; explicit list = allowlist
mcp_stateless: bool # True = stateless HTTP (POST-only); False = stateful (supports GET/SSE)
mcp_instructions: str | None # Additional instructions appended to retain/recall MCP tool descriptions
enable_bank_config_api: bool
enable_bank_llm_health: bool
enable_dry_run_extract: bool
@@ -1704,7 +1642,6 @@ class HindsightConfig:
otel_service_name: str
otel_deployment_environment: str
metrics_include_bank_id: bool
metrics_backlog_enabled: bool
# Audit log configuration (static - server-level only)
audit_log_enabled: bool # Master switch for audit logging
@@ -1739,11 +1676,6 @@ class HindsightConfig:
embeddings_zeroentropy_encoding_format: str = DEFAULT_EMBEDDINGS_ZEROENTROPY_ENCODING_FORMAT
embeddings_zeroentropy_batch_size: int = DEFAULT_EMBEDDINGS_ZEROENTROPY_BATCH_SIZE
embeddings_zeroentropy_latency: str | None = DEFAULT_EMBEDDINGS_ZEROENTROPY_LATENCY
file_parser_markitdown_ocr_enabled: bool = DEFAULT_FILE_PARSER_MARKITDOWN_OCR_ENABLED
file_parser_markitdown_ocr_api_key: str | None = None
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
# Class-level sets for configuration categorization
@@ -1784,8 +1716,6 @@ class HindsightConfig:
"file_storage_gcs_service_account_key",
"file_storage_azure_account_key",
# File parser credentials
"file_parser_markitdown_ocr_api_key",
"file_parser_markitdown_ocr_base_url",
"file_parser_iris_token",
"file_parser_llama_parse_api_key",
}
@@ -1949,9 +1879,6 @@ class HindsightConfig:
f"Note: 'standard' is not a valid Bedrock service tier -- use unset for default tier."
)
# Validate gemini_service_tier
self.llm_gemini_service_tier = parse_gemini_service_tier(self.llm_gemini_service_tier)
# When LLM provider is "none", force chunks-only mode and disable LLM-dependent features
if self.llm_provider == "none":
self.retain_extraction_mode = "chunks"
@@ -2069,11 +1996,6 @@ class HindsightConfig:
llm_groq_service_tier=os.getenv(ENV_LLM_GROQ_SERVICE_TIER, DEFAULT_LLM_GROQ_SERVICE_TIER),
llm_openai_service_tier=os.getenv(ENV_LLM_OPENAI_SERVICE_TIER, DEFAULT_LLM_OPENAI_SERVICE_TIER),
llm_bedrock_service_tier=os.getenv(ENV_LLM_BEDROCK_SERVICE_TIER) or None,
llm_gemini_service_tier=(
parse_gemini_service_tier(os.getenv(ENV_LLM_GEMINI_SERVICE_TIER) or DEFAULT_LLM_GEMINI_SERVICE_TIER)
if llm_provider.lower() == "gemini"
else None
),
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"),
@@ -2348,15 +2270,6 @@ class HindsightConfig:
recall_strategy_boosts=_parse_strategy_boosts(
os.getenv(ENV_RECALL_STRATEGY_BOOSTS, DEFAULT_RECALL_STRATEGY_BOOSTS)
),
recency_decay_function=_validate_recency_decay_function(
os.getenv(ENV_RECENCY_DECAY_FUNCTION, DEFAULT_RECENCY_DECAY_FUNCTION)
),
recency_decay_linear_window_days=float(
os.getenv(ENV_RECENCY_DECAY_LINEAR_WINDOW_DAYS, str(DEFAULT_RECENCY_DECAY_LINEAR_WINDOW_DAYS))
),
recency_decay_halflife_days=float(
os.getenv(ENV_RECENCY_DECAY_HALFLIFE_DAYS, str(DEFAULT_RECENCY_DECAY_HALFLIFE_DAYS))
),
# Cohere reranker (with backward-compatible fallback to shared API key)
reranker_cohere_api_key=os.getenv(ENV_RERANKER_COHERE_API_KEY) or os.getenv(ENV_COHERE_API_KEY),
reranker_cohere_model=os.getenv(ENV_RERANKER_COHERE_MODEL, DEFAULT_RERANKER_COHERE_MODEL),
@@ -2432,7 +2345,6 @@ class HindsightConfig:
if os.getenv(ENV_MCP_ENABLED_TOOLS)
else DEFAULT_MCP_ENABLED_TOOLS,
mcp_stateless=os.getenv(ENV_MCP_STATELESS, str(DEFAULT_MCP_STATELESS)).lower() == "true",
mcp_instructions=os.getenv(ENV_MCP_INSTRUCTIONS) or DEFAULT_MCP_INSTRUCTIONS,
enable_bank_llm_health=os.getenv(ENV_ENABLE_BANK_LLM_HEALTH, str(DEFAULT_ENABLE_BANK_LLM_HEALTH)).lower()
== "true",
enable_bank_config_api=os.getenv(ENV_ENABLE_BANK_CONFIG_API, str(DEFAULT_ENABLE_BANK_CONFIG_API)).lower()
@@ -2512,18 +2424,6 @@ class HindsightConfig:
file_parser_allowlist=_parse_str_list(os.getenv(ENV_FILE_PARSER_ALLOWLIST))
if os.getenv(ENV_FILE_PARSER_ALLOWLIST)
else None,
file_parser_markitdown_ocr_enabled=os.getenv(
ENV_FILE_PARSER_MARKITDOWN_OCR_ENABLED,
str(DEFAULT_FILE_PARSER_MARKITDOWN_OCR_ENABLED),
).lower()
in ("1", "true", "yes", "on"),
file_parser_markitdown_ocr_api_key=os.getenv(ENV_FILE_PARSER_MARKITDOWN_OCR_API_KEY) or None,
file_parser_markitdown_ocr_base_url=os.getenv(ENV_FILE_PARSER_MARKITDOWN_OCR_BASE_URL) or None,
file_parser_markitdown_ocr_model=os.getenv(ENV_FILE_PARSER_MARKITDOWN_OCR_MODEL) or None,
file_parser_markitdown_ocr_prompt=os.getenv(
ENV_FILE_PARSER_MARKITDOWN_OCR_PROMPT,
DEFAULT_FILE_PARSER_MARKITDOWN_OCR_PROMPT,
),
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,
@@ -2714,8 +2614,6 @@ class HindsightConfig:
otel_deployment_environment=os.getenv(ENV_OTEL_DEPLOYMENT_ENVIRONMENT, DEFAULT_OTEL_DEPLOYMENT_ENVIRONMENT),
metrics_include_bank_id=os.getenv(ENV_METRICS_INCLUDE_BANK_ID, str(DEFAULT_METRICS_INCLUDE_BANK_ID)).lower()
in ("true", "1", "yes"),
metrics_backlog_enabled=os.getenv(ENV_METRICS_BACKLOG_ENABLED, str(DEFAULT_METRICS_BACKLOG_ENABLED)).lower()
in ("true", "1", "yes"),
# 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=[
@@ -8,7 +8,6 @@ Config values are resolved on every request to ensure consistency across
multiple API servers.
"""
import asyncio
import json
import logging
from dataclasses import asdict, replace
@@ -162,83 +161,26 @@ class ConfigResolver:
resolved_config = await self.resolve_full_config(bank_id, context)
config_dict = asdict(resolved_config)
# SECURITY: drop static/infrastructure + credential fields, then permission-filter.
filtered = self._strip_static_and_credential_fields(config_dict)
return await self._apply_permission_filter(filtered, bank_id, context)
# SECURITY: Filter to only configurable fields (exclude static/infrastructure)
filtered = {k: v for k, v in config_dict.items() if k in self._configurable_fields}
def _strip_static_and_credential_fields(self, config_dict: dict[str, Any]) -> dict[str, Any]:
"""Keep only configurable, non-credential fields.
# SECURITY: Remove ALL credential fields (API keys, base URLs, etc.)
filtered = {k: v for k, v in filtered.items() if k not in self._credential_fields}
SECURITY: excludes static/infrastructure fields and ALL credential fields
(API keys, base URLs, etc.) so a resolved config is safe to return over the API.
"""
return {
k: v for k, v in config_dict.items() if k in self._configurable_fields and k not in self._credential_fields
}
async def _apply_permission_filter(
self, filtered: dict[str, Any], bank_id: str, context: RequestContext | None
) -> dict[str, Any]:
"""Further restrict already-stripped config to the tenant/bank permission allow-list.
On extension error, leaves ``filtered`` unchanged (parity with the historical
single-bank path: a permissions lookup failure must not leak or drop fields).
"""
if not (self.tenant_extension and context):
return filtered
try:
allowed_fields = await self.tenant_extension.get_allowed_config_fields(context, bank_id)
if allowed_fields is not None: # None means "allow all"
filtered = {k: v for k, v in filtered.items() if k in allowed_fields}
logger.debug(
f"Applied permission filter for bank {bank_id}: allowed={len(allowed_fields)} fields, "
f"returned={len(filtered)} fields"
)
except Exception as e:
logger.warning(f"Failed to load permissions for bank {bank_id}: {e}")
return filtered
async def get_bank_configs(
self, bank_ids: list[str], context: RequestContext | None = None
) -> dict[str, dict[str, Any]]:
"""Batch variant of :meth:`get_bank_config` for many banks.
Equivalent to calling ``get_bank_config`` per bank, but resolves the
global + tenant base once and loads every bank's ``banks.config`` JSONB
in a single query, instead of one config round-trip per bank. Used by
``list_banks`` to overlay disposition + mission without an N+1.
Returns a mapping of bank_id -> filtered configurable-field dict. A bank
with no config row still appears, mapped to the global+tenant base.
"""
if not bank_ids:
return {}
# Global + tenant base, resolved once (tenant override is per-request, not per-bank).
base_dict = asdict(self._global_config)
# PERMISSIONS: Further filter based on tenant/bank permissions
if self.tenant_extension and context:
try:
tenant_overrides = await self.tenant_extension.get_tenant_config(context)
if tenant_overrides:
normalized_tenant = normalize_config_dict(tenant_overrides)
base_dict.update({k: v for k, v in normalized_tenant.items() if k in self._configurable_fields})
allowed_fields = await self.tenant_extension.get_allowed_config_fields(context, bank_id)
if allowed_fields is not None: # None means "allow all"
filtered = {k: v for k, v in filtered.items() if k in allowed_fields}
logger.debug(
f"Applied permission filter for bank {bank_id}: allowed={len(allowed_fields)} fields, "
f"returned={len(filtered)} fields"
)
except Exception as e:
logger.warning(f"Failed to load tenant config for bulk resolve: {e}")
logger.warning(f"Failed to load permissions for bank {bank_id}: {e}")
# All bank overrides in one query, then merge + strip per bank.
bank_overrides = await self._load_bank_configs(bank_ids)
stripped = {
bank_id: self._strip_static_and_credential_fields({**base_dict, **bank_overrides.get(bank_id, {})})
for bank_id in bank_ids
}
# Permission filter is per-bank; resolve concurrently when an extension is present.
if not (self.tenant_extension and context):
return stripped
permission_filtered = await asyncio.gather(
*(self._apply_permission_filter(stripped[bank_id], bank_id, context) for bank_id in bank_ids)
)
return dict(zip(bank_ids, permission_filtered, strict=True))
return filtered
async def _load_bank_config(self, bank_id: str) -> dict[str, Any]:
"""
@@ -277,45 +219,6 @@ class ConfigResolver:
return {}
async def _load_bank_configs(self, bank_ids: list[str]) -> dict[str, dict[str, Any]]:
"""Bulk variant of :meth:`_load_bank_config`: load many banks' overrides in one query.
Returns a mapping of bank_id -> normalized active overrides. Banks with no row
(or an empty/all-tombstone config) are simply absent from the mapping.
"""
result: dict[str, dict[str, Any]] = {}
if not bank_ids:
return result
try:
async with self._backend.acquire() as conn:
rows = await conn.fetch(
f"""
SELECT bank_id, config FROM {fq_table("banks")} WHERE bank_id = ANY($1)
""",
bank_ids,
)
for row in rows:
config_data = row["config"]
if not config_data:
continue
# Handle case where JSONB is returned as JSON string
if isinstance(config_data, str):
config_data = json.loads(config_data)
# Normalize keys (handle both env var format and Python field format)
normalized = normalize_config_dict(config_data)
# Only active overrides for configurable fields. JSON null is a tombstone
# for "Server Default" in the bank-config UI and must not override defaults.
overrides = {
k: v for k, v in normalized.items() if k in self._configurable_fields and v is not None
}
if overrides:
result[row["bank_id"]] = overrides
except Exception as e:
logger.error(f"Failed to bulk-load bank configs: {e}")
return result
async def update_bank_config(
self, bank_id: str, updates: dict[str, Any], context: RequestContext | None = None
) -> None:
@@ -402,9 +305,6 @@ class ConfigResolver:
# Validate recall budget fields
_validate_recall_budget_updates(normalized_updates)
# Validate disposition trait fields (1-5 integer scale)
_validate_disposition_updates(normalized_updates)
chunking_fields_updated = (
"retain_chunk_size" in normalized_updates
or "retain_structured_chunk_size" in normalized_updates
@@ -519,31 +419,6 @@ def _validate_recall_budget_updates(updates: dict[str, Any]) -> None:
)
_DISPOSITION_KEYS = (
"disposition_skepticism",
"disposition_literalism",
"disposition_empathy",
)
def _validate_disposition_updates(updates: dict[str, Any]) -> None:
"""Validate disposition trait config updates. Raises ValueError on invalid input.
Each trait is an integer on a 1-5 scale (or None to clear the per-bank
override). The read overlay injects the stored value verbatim into a strict
``DispositionTraits(int, ge=1, le=5)``; an out-of-contract value (a float, a
0-1 scale, or an int outside 1-5) accepted here would later 500 the whole
bank list when any bank profile is serialized (issue #2348).
"""
for key in _DISPOSITION_KEYS:
if key in updates:
value = updates[key]
if value is None:
continue
if not isinstance(value, int) or isinstance(value, bool) or not (1 <= value <= 5):
raise ValueError(f"{key} must be an integer between 1 and 5, got {value!r}")
def apply_strategy(config: HindsightConfig, strategy_name: str) -> HindsightConfig:
"""
Apply a named retain strategy's overrides on top of a resolved config.
@@ -96,10 +96,7 @@ class BankStatsCache:
value = await loader()
except BaseException as exc:
async with self._lock:
# Invalidation may have detached this loader and allowed a new
# one to claim the key. Never remove that newer loader's slot.
if self._in_flight.get(key) is in_flight:
self._in_flight.pop(key, None)
self._in_flight.pop(key, None)
if not in_flight.done():
in_flight.set_exception(exc)
# Suppress "Future exception was never retrieved" when no other
@@ -109,12 +106,8 @@ class BankStatsCache:
raise
async with self._lock:
# Only the loader that still owns the key may populate the cache.
# An invalidated loader can finish for its original callers, but its
# pre-invalidation result must not overwrite a newer load.
if self._in_flight.get(key) is in_flight:
self._store_unlocked(key, value)
self._in_flight.pop(key, None)
self._store_unlocked(key, value)
self._in_flight.pop(key, None)
if not in_flight.done():
in_flight.set_result(value)
return value
@@ -122,13 +115,8 @@ class BankStatsCache:
async def invalidate(self, schema: str, bank_id: str) -> None:
"""Drop any cached stats for `(schema, bank_id)`."""
async with self._lock:
key = (schema, bank_id)
self._entries.pop(key, None)
# Detach rather than cancel: existing callers may finish with the
# snapshot they requested, while post-invalidation callers reload.
self._in_flight.pop(key, None)
self._entries.pop((schema, bank_id), None)
async def clear(self) -> None:
async with self._lock:
self._entries.clear()
self._in_flight.clear()
@@ -449,13 +449,6 @@ class _CreateAction(BaseModel):
def sanitize_text(cls, v: str) -> str:
return sanitize_llm_output(v) or ""
@field_validator("source_fact_ids", mode="before")
@classmethod
def ensure_list(cls, v: str | list[str]) -> list[str]:
if isinstance(v, str):
return [v]
return v
class _UpdateAction(BaseModel):
text: str
@@ -468,13 +461,6 @@ class _UpdateAction(BaseModel):
def sanitize_text(cls, v: str) -> str:
return sanitize_llm_output(v) or ""
@field_validator("source_fact_ids", mode="before")
@classmethod
def ensure_list(cls, v: str | list[str]) -> list[str]:
if isinstance(v, str):
return [v]
return v
class _DeleteAction(BaseModel):
observation_id: str # UUID of the observation to remove
@@ -654,7 +640,6 @@ class ConsolidationPerfLog:
self.start_time = time.time()
self.lines: list[str] = []
self.timings: dict[str, float] = {}
self.timing_counts: dict[str, int] = {}
self.llm_calls: int = 0
self.total_obs_in_context: int = 0
self.total_prompt_chars: int = 0
@@ -664,13 +649,11 @@ class ConsolidationPerfLog:
self.lines.append(message)
def record_timing(self, key: str, duration: float) -> None:
"""Record a timing measurement.
Tracks both total seconds and call count so the summary can
distinguish one slow call from many fast calls in aggregate.
"""
self.timings[key] = self.timings.get(key, 0.0) + duration
self.timing_counts[key] = self.timing_counts.get(key, 0) + 1
"""Record a timing measurement."""
if key in self.timings:
self.timings[key] += duration
else:
self.timings[key] = duration
def record_llm_call(self, obs_count: int, prompt_chars: int) -> None:
"""Record stats for a single LLM call."""
@@ -693,8 +676,6 @@ class ConsolidationPerfLog:
"""
for key, value in other.timings.items():
self.timings[key] = self.timings.get(key, 0.0) + value
for key, count in other.timing_counts.items():
self.timing_counts[key] = self.timing_counts.get(key, 0) + count
self.llm_calls += other.llm_calls
self.total_obs_in_context += other.total_obs_in_context
self.total_prompt_chars += other.total_prompt_chars
@@ -1295,22 +1276,16 @@ async def _run_consolidation_job(
f"{stats['skipped']} skipped)"
)
# Add timing breakdown. Each phase is recorded once per call, so the count
# disambiguates a single slow call from many fast calls — important for
# operators triaging "the recall phase took 15s" log lines, where the
# total is the sum of many serial sub-calls rather than one slow query.
def _fmt(key: str) -> str:
total = perf.timings[key]
count = perf.timing_counts.get(key, 0)
if count > 1:
avg_ms = total * 1000.0 / count
return f"{key}={total:.3f}s ({count} calls, avg={avg_ms:.0f}ms)"
return f"{key}={total:.3f}s"
# Add timing breakdown
timing_parts = []
for key in ("recall", "llm", "embedding", "db_write"):
if key in perf.timings:
timing_parts.append(_fmt(key))
if "recall" in perf.timings:
timing_parts.append(f"recall={perf.timings['recall']:.3f}s")
if "llm" in perf.timings:
timing_parts.append(f"llm={perf.timings['llm']:.3f}s")
if "embedding" in perf.timings:
timing_parts.append(f"embedding={perf.timings['embedding']:.3f}s")
if "db_write" in perf.timings:
timing_parts.append(f"db_write={perf.timings['db_write']:.3f}s")
if perf.llm_calls > 0:
timing_parts.append(f"avg_obs={perf.total_obs_in_context / perf.llm_calls:.1f}")
@@ -212,7 +212,7 @@ class LocalSTCrossEncoder(CrossEncoderModel):
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
# Check for GPU (CUDA) or Apple Silicon (MPS)
# 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
@@ -220,13 +220,10 @@ class LocalSTCrossEncoder(CrossEncoderModel):
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
device = None # Let sentence-transformers auto-detect GPU/MPS
except Exception as e:
logger.warning(f"Failed to detect GPU/MPS/XPU, falling back to CPU: {e}")
logger.warning(f"Failed to detect GPU/MPS, 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
@@ -256,21 +256,13 @@ class OracleOps(DataAccessOps):
# 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 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"""
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],
[(bank_id, uid) for uid in unit_ids],
)
async def claim_graph_maintenance_batch(
@@ -348,15 +348,6 @@ class PostgreSQLOps(DataAccessOps):
) -> None:
if not unit_ids:
return
# Sort to enforce a global lock-acquisition order on the
# (bank_id, unit_id) unique-key. Without this, two concurrent
# transactions inserting overlapping unit_id sets in different
# orders can deadlock on the ON CONFLICT row locks — Postgres
# acquires a short-lived lock per row being checked, and cycle
# detection then aborts one transaction. 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.execute(
f"""
INSERT INTO {table} (bank_id, unit_id)
@@ -364,7 +355,7 @@ class PostgreSQLOps(DataAccessOps):
ON CONFLICT (bank_id, unit_id) DO NOTHING
""",
bank_id,
sorted_unit_ids,
unit_ids,
)
async def claim_graph_maintenance_batch(
@@ -190,7 +190,7 @@ class LocalSTEmbeddings(Embeddings):
device = "cpu"
logger.info("Embeddings: forcing CPU mode")
else:
# Check for GPU (CUDA), Apple Silicon (MPS), or Intel XPU
# Check for GPU (CUDA) or Apple Silicon (MPS)
# 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
@@ -198,13 +198,10 @@ class LocalSTEmbeddings(Embeddings):
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
device = None # Let sentence-transformers auto-detect GPU/MPS
except Exception as e:
logger.warning(f"Failed to detect GPU/MPS/XPU, falling back to CPU: {e}")
logger.warning(f"Failed to detect GPU/MPS, falling back to CPU: {e}")
# Suppress verbose transformers warnings during model loading
# This suppresses the "UNEXPECTED" warnings from BertModel which are harmless
@@ -712,8 +709,7 @@ class OpenAIEmbeddings(Embeddings):
class CodexOAuthEmbeddings(OpenAIEmbeddings):
"""
OpenAI embeddings using the Codex/ChatGPT OAuth token from the Codex
``auth.json`` (``$CODEX_HOME/auth.json``, or ``~/.codex/auth.json`` when unset).
OpenAI embeddings using the Codex/ChatGPT OAuth token from ``~/.codex/auth.json``.
Codex OAuth is an LLM-provider auth path in Hindsight, but the same bearer token
can also authenticate against the standard OpenAI embeddings endpoint. This keeps
@@ -782,6 +782,236 @@ class EntityResolver:
return entity_ids
async def resolve_entity(
self,
bank_id: str,
entity_text: str,
context: str,
nearby_entities: list[dict],
unit_event_date,
) -> str:
"""
Resolve an entity to a canonical entity ID.
Args:
bank_id: bank ID (entities are scoped to agents)
entity_text: Entity text ("Alice", "Google", etc.)
context: Context where entity appears
nearby_entities: Other entities in the same unit
unit_event_date: When this unit was created
Returns:
Entity ID (creates new entity if needed)
"""
async with acquire_with_retry(self.pool) as conn:
# Find candidate entities with similar name
candidates = await conn.fetch(
f"""
SELECT id, canonical_name, metadata, last_seen
FROM {fq_table("entities")}
WHERE bank_id = $1
AND (
canonical_name ILIKE $2
OR canonical_name ILIKE $3
OR $2 ILIKE canonical_name || '%%'
)
ORDER BY mention_count DESC
""",
bank_id,
entity_text,
f"%{entity_text}%",
)
if not candidates:
# New entity - create it
return await self._create_entity(conn, bank_id, entity_text, unit_event_date)
# Score candidates based on:
# 1. Name similarity
# 2. Context overlap (TODO: could use embeddings)
# 3. Co-occurring entities
# 4. Temporal proximity
best_candidate = None
best_score = 0.0
nearby_entity_set = {e["text"].lower() for e in nearby_entities if e["text"] != entity_text}
for row in candidates:
candidate_id = row["id"]
canonical_name = row["canonical_name"]
last_seen = row["last_seen"]
score = 0.0
# 1. Name similarity (0-1)
name_similarity = SequenceMatcher(None, entity_text.lower(), canonical_name.lower()).ratio()
score += name_similarity * 0.5
# 2. Co-occurring entities (0-0.5)
# Get entities that co-occurred with this candidate before
# Use the materialized co-occurrence cache for fast lookup
co_entity_rows = await conn.fetch(
f"""
SELECT e.canonical_name, ec.cooccurrence_count
FROM {fq_table("entity_cooccurrences")} ec
JOIN {fq_table("entities")} e ON (
CASE
WHEN ec.entity_id_1 = $1 THEN ec.entity_id_2
WHEN ec.entity_id_2 = $1 THEN ec.entity_id_1
END = e.id
)
WHERE ec.entity_id_1 = $1 OR ec.entity_id_2 = $1
""",
candidate_id,
)
co_entities = {r["canonical_name"].lower() for r in co_entity_rows}
# Check overlap with nearby entities
overlap = len(nearby_entity_set & co_entities)
if nearby_entity_set:
co_entity_score = overlap / len(nearby_entity_set)
score += co_entity_score * 0.3
# 3. Temporal proximity (0-0.2)
if last_seen:
# Normalize both to UTC-aware to avoid naive/aware mismatch
# (Oracle returns naive datetimes from fromisoformat)
_evt = unit_event_date if unit_event_date.tzinfo else unit_event_date.replace(tzinfo=UTC)
_seen = last_seen if last_seen.tzinfo else last_seen.replace(tzinfo=UTC)
days_diff = abs((_evt - _seen).total_seconds() / 86400)
if days_diff < 7: # Within a week
temporal_score = max(0, 1.0 - (days_diff / 7))
score += temporal_score * 0.2
if score > best_score:
best_score = score
best_candidate = candidate_id
# Threshold for considering it the same entity
threshold = 0.6
if best_score > threshold:
# Update entity
await conn.execute(
f"""
UPDATE {fq_table("entities")}
SET mention_count = mention_count + 1,
last_seen = $1
WHERE id = $2
""",
unit_event_date,
best_candidate,
)
return best_candidate
else:
# Not confident - create new entity
return await self._create_entity(conn, bank_id, entity_text, unit_event_date)
async def _create_entity(
self,
conn,
bank_id: str,
entity_text: str,
event_date,
) -> str:
"""
Create a new entity or get existing one if it already exists.
Uses INSERT ... ON CONFLICT to handle race conditions where
two concurrent transactions try to create the same entity.
Args:
conn: Database connection
bank_id: bank ID
entity_text: Entity text
event_date: When first seen
Returns:
Entity ID
"""
entity_id = await conn.fetchval(
f"""
INSERT INTO {fq_table("entities")} (bank_id, canonical_name, first_seen, last_seen, mention_count)
VALUES ($1, $2, COALESCE($3, now()), COALESCE($4, now()), 1)
ON CONFLICT (bank_id, LOWER(canonical_name))
DO UPDATE SET
mention_count = {fq_table("entities")}.mention_count + 1,
last_seen = EXCLUDED.last_seen
RETURNING id
""",
bank_id,
entity_text,
event_date,
event_date,
)
return entity_id
async def link_unit_to_entity(self, unit_id: str, entity_id: str):
"""
Link a memory unit to an entity.
Also updates co-occurrence cache with other entities in the same unit.
Args:
unit_id: Memory unit ID
entity_id: Entity ID
"""
async with acquire_with_retry(self.pool) as conn:
# Insert unit-entity link
await conn.execute(
f"""
INSERT INTO {fq_table("unit_entities")} (unit_id, entity_id)
VALUES ($1, $2)
ON CONFLICT DO NOTHING
""",
unit_id,
entity_id,
)
# Update co-occurrence cache: find other entities in this unit
rows = await conn.fetch(
f"""
SELECT entity_id
FROM {fq_table("unit_entities")}
WHERE unit_id = $1 AND entity_id != $2
""",
unit_id,
entity_id,
)
other_entities = [row["entity_id"] for row in rows]
# Update co-occurrences for each pair
for other_entity_id in other_entities:
await self._update_cooccurrence(conn, entity_id, other_entity_id)
async def _update_cooccurrence(self, conn, entity_id_1: str, entity_id_2: str):
"""
Update the co-occurrence cache for two entities.
Uses CHECK constraint ordering (entity_id_1 < entity_id_2) to avoid duplicates.
Args:
conn: Database connection
entity_id_1: First entity ID
entity_id_2: Second entity ID
"""
# Ensure consistent ordering (smaller UUID first)
if entity_id_1 > entity_id_2:
entity_id_1, entity_id_2 = entity_id_2, entity_id_1
await conn.execute(
f"""
INSERT INTO {fq_table("entity_cooccurrences")} (entity_id_1, entity_id_2, cooccurrence_count, last_cooccurred)
VALUES ($1, $2, 1, NOW())
ON CONFLICT (entity_id_1, entity_id_2)
DO UPDATE SET
cooccurrence_count = {fq_table("entity_cooccurrences")}.cooccurrence_count + 1,
last_cooccurred = NOW()
""",
entity_id_1,
entity_id_2,
)
async def link_units_to_entities_batch(
self,
unit_entity_pairs: list[tuple[str, str]] | list[tuple[str, str, datetime | None]],
@@ -6,7 +6,6 @@ enabling support for multiple LLM backends (OpenAI, Anthropic, Gemini, Codex, et
"""
from abc import ABC, abstractmethod
from datetime import datetime
from typing import Any
from .response_models import LLMToolCallResult
@@ -253,11 +252,3 @@ class OutputTooLongError(Exception):
"""
pass
class ProviderRateLimitResetError(Exception):
"""Raised when an upstream provider says quota will reopen at a known time."""
def __init__(self, retry_at: datetime, message: str = "") -> None:
self.retry_at = retry_at
super().__init__(message)
@@ -10,6 +10,7 @@ import re
import time
import uuid
from contextlib import AsyncExitStack
from pathlib import Path
from typing import TYPE_CHECKING, Any
# Vertex AI imports (conditional - for LLMProvider to pass credentials to GeminiLLM)
@@ -252,7 +253,6 @@ def create_llm_provider(
gemini_safety_settings: list | None = None,
prompt_cache_enabled: bool = False,
litellmrouter_config: dict[str, Any] | None = None,
gemini_service_tier: str | None = None,
) -> Any: # Returns LLMInterface
"""
Factory function to create the appropriate LLM provider implementation.
@@ -266,7 +266,6 @@ def create_llm_provider(
groq_service_tier: Groq service tier (for Groq provider) - "on_demand", "flex", or "auto".
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).
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
@@ -297,12 +296,6 @@ def create_llm_provider(
)
provider_lower = provider.lower()
if provider_lower == "gemini":
from ..config import parse_gemini_service_tier
gemini_service_tier = parse_gemini_service_tier(gemini_service_tier)
else:
gemini_service_tier = None
if provider_lower == "openai-codex":
return CodexLLM(
@@ -351,7 +344,6 @@ def create_llm_provider(
vertexai_region=vertexai_region,
vertexai_credentials=vertexai_credentials,
gemini_safety_settings=gemini_safety_settings,
gemini_service_tier=gemini_service_tier,
prompt_cache_enabled=prompt_cache_enabled,
extra_body=extra_body,
)
@@ -466,7 +458,6 @@ def create_llm_provider(
"openrouter",
"zai",
"opencode-go",
"atlas",
):
return OpenAICompatibleLLM(
provider=provider,
@@ -505,7 +496,6 @@ class LLMProvider:
extra_body: dict[str, Any] | None = None,
default_headers: dict[str, str] | None = None,
litellmrouter_config: dict[str, Any] | None = None,
gemini_service_tier: str | None = None,
):
"""
Initialize LLM provider.
@@ -519,7 +509,6 @@ class LLMProvider:
groq_service_tier: Groq service tier ("on_demand", "flex", "auto") - from config.
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.
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).
@@ -543,7 +532,6 @@ class LLMProvider:
self.groq_service_tier = groq_service_tier
self.openai_service_tier = openai_service_tier
self.bedrock_service_tier = bedrock_service_tier
self.gemini_service_tier = gemini_service_tier
# 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
@@ -589,7 +577,6 @@ class LLMProvider:
"openrouter",
"zai",
"opencode-go",
"atlas",
"fireworks",
"nous",
]
@@ -616,8 +603,6 @@ class LLMProvider:
self.base_url = "https://api.z.ai/api/coding/paas/v4"
elif self.provider == "opencode-go":
self.base_url = "https://opencode.ai/zen/go/v1"
elif self.provider == "atlas":
self.base_url = "https://api.atlascloud.ai/v1"
elif self.provider == "nous":
self.base_url = "https://inference-api.nousresearch.com/v1"
@@ -675,22 +660,6 @@ class LLMProvider:
except Exception:
pass # Config may not be initialized in test environments
if self.provider == "gemini":
from ..config import parse_gemini_service_tier
self.gemini_service_tier = parse_gemini_service_tier(self.gemini_service_tier)
if self.provider == "gemini" and self.gemini_service_tier is None:
from ..config import _get_raw_config
try:
raw_config = _get_raw_config()
self.gemini_service_tier = raw_config.llm_gemini_service_tier
except Exception:
pass # Config may not be initialized in test environments
elif self.provider != "gemini":
self.gemini_service_tier = None
# Prompt-prefix caching is a provider-agnostic toggle (default on): resolve
# it from the static server config for every provider when the caller didn't
# pass an explicit override. Providers that don't support caching ignore the
@@ -729,7 +698,6 @@ class LLMProvider:
groq_service_tier=self.groq_service_tier,
openai_service_tier=self.openai_service_tier,
bedrock_service_tier=self.bedrock_service_tier,
gemini_service_tier=self.gemini_service_tier,
extra_body=self.extra_body,
default_headers=self.default_headers,
vertexai_project_id=vertexai_project_id,
@@ -1055,9 +1023,7 @@ class LLMProvider:
def _load_codex_auth(self) -> tuple[str, str]:
"""
Load OAuth credentials from the Codex ``auth.json``.
Honors ``CODEX_HOME`` (falling back to ``~/.codex``).
Load OAuth credentials from ~/.codex/auth.json.
Returns:
Tuple of (access_token, account_id).
@@ -1066,9 +1032,7 @@ class LLMProvider:
FileNotFoundError: If auth file doesn't exist.
ValueError: If auth file is invalid.
"""
from .providers.codex_auth import default_codex_auth_file
auth_file = default_codex_auth_file()
auth_file = Path.home() / ".codex" / "auth.json"
if not auth_file.exists():
raise FileNotFoundError(
@@ -1178,12 +1142,10 @@ class LLMProvider:
ENV_LLM_BEDROCK_SERVICE_TIER,
ENV_LLM_DEFAULT_HEADERS,
ENV_LLM_EXTRA_BODY,
ENV_LLM_GEMINI_SERVICE_TIER,
ENV_LLM_MODEL,
ENV_LLM_PROVIDER,
ENV_LLM_REASONING_EFFORT,
_get_default_model_for_provider,
parse_gemini_service_tier,
)
provider = os.getenv(ENV_LLM_PROVIDER, DEFAULT_LLM_PROVIDER)
@@ -1210,11 +1172,6 @@ class LLMProvider:
extra_body=extra_body,
default_headers=default_headers,
bedrock_service_tier=os.getenv(ENV_LLM_BEDROCK_SERVICE_TIER) or None,
gemini_service_tier=(
parse_gemini_service_tier(os.getenv(ENV_LLM_GEMINI_SERVICE_TIER))
if provider.lower() == "gemini"
else None
),
)
@@ -45,7 +45,6 @@ from .audit import AuditLogger, audit_context
from .bank_stats_cache import BankStatsCache
from .db import DatabaseBackend, create_database_backend
from .db_budget import budgeted_operation
from .llm_interface import ProviderRateLimitResetError
from .llm_trace import (
LLMRequestEntry,
LLMRequestListResponse,
@@ -745,37 +744,6 @@ def _resolve_refresh_tag_filtering(
return RefreshTagFiltering(tags=model_tags, tags_match=tags_match, tag_groups=None)
@dataclass
class ResolvedDispositionMission:
"""Disposition + mission after overlaying resolved bank config on the legacy columns."""
disposition: dict[str, int]
mission: str
def _overlay_bank_config_disposition_mission(
disposition: dict[str, int], mission: str, config_dict: dict[str, Any]
) -> ResolvedDispositionMission:
"""Overlay resolved bank config on top of the legacy banks.disposition /
banks.mission column values.
``reflect_mission`` and ``disposition_*`` in the resolved bank config take
precedence over the legacy DB columns. Shared by ``get_bank_profile`` and
``list_banks`` so the single-bank and list paths return identical
disposition + mission for the same bank.
"""
resolved_mission = config_dict.get("reflect_mission") or mission
cfg_skep = config_dict.get("disposition_skepticism")
cfg_lit = config_dict.get("disposition_literalism")
cfg_emp = config_dict.get("disposition_empathy")
resolved_disposition = {
"skepticism": cfg_skep if cfg_skep is not None else disposition["skepticism"],
"literalism": cfg_lit if cfg_lit is not None else disposition["literalism"],
"empathy": cfg_emp if cfg_emp is not None else disposition["empathy"],
}
return ResolvedDispositionMission(disposition=resolved_disposition, mission=resolved_mission)
class MemoryEngine(MemoryEngineInterface):
"""
Advanced memory system using temporal and semantic linking with PostgreSQL.
@@ -965,7 +933,6 @@ class MemoryEngine(MemoryEngineInterface):
default_headers=config.llm_default_headers,
litellmrouter_config=config.llm_litellmrouter_config,
bedrock_service_tier=config.llm_bedrock_service_tier,
gemini_service_tier=config.llm_gemini_service_tier,
)
# Store client and model for convenience (deprecated: use _llm_config.call() instead)
@@ -999,7 +966,6 @@ class MemoryEngine(MemoryEngineInterface):
default_headers=config.llm_default_headers,
litellmrouter_config=config.retain_llm_litellmrouter_config or config.llm_litellmrouter_config,
bedrock_service_tier=config.llm_bedrock_service_tier,
gemini_service_tier=config.llm_gemini_service_tier,
)
# Reflect LLM config - for think/observe operations (can use lighter models)
@@ -1028,7 +994,6 @@ class MemoryEngine(MemoryEngineInterface):
default_headers=config.llm_default_headers,
litellmrouter_config=config.reflect_llm_litellmrouter_config or config.llm_litellmrouter_config,
bedrock_service_tier=config.llm_bedrock_service_tier,
gemini_service_tier=config.llm_gemini_service_tier,
)
# Consolidation LLM config - for mental model consolidation (can use efficient models)
@@ -1057,7 +1022,6 @@ class MemoryEngine(MemoryEngineInterface):
default_headers=config.llm_default_headers,
litellmrouter_config=config.consolidation_llm_litellmrouter_config or config.llm_litellmrouter_config,
bedrock_service_tier=config.llm_bedrock_service_tier,
gemini_service_tier=config.llm_gemini_service_tier,
)
# Initialize cross-encoder reranker (cached for performance)
@@ -1788,9 +1752,6 @@ class MemoryEngine(MemoryEngineInterface):
audit_entry.response = {"status": "completed", "operation_id": operation_id}
except ProviderRateLimitResetError as e:
logger.warning(f"Task deferred until provider quota resets at {e.retry_at}: {e}")
raise DeferOperation(exec_date=e.retry_at, reason=str(e)) from e
except RetryTaskAt:
# Task-owned retry: let the poller handle scheduling
raise
@@ -2790,15 +2751,7 @@ class MemoryEngine(MemoryEngineInterface):
self._parser_registry = FileParserRegistry()
try:
self._parser_registry.register(
MarkitdownParser(
ocr_enabled=config.file_parser_markitdown_ocr_enabled,
ocr_api_key=config.file_parser_markitdown_ocr_api_key,
ocr_base_url=config.file_parser_markitdown_ocr_base_url,
ocr_model=config.file_parser_markitdown_ocr_model,
ocr_prompt=config.file_parser_markitdown_ocr_prompt,
)
)
self._parser_registry.register(MarkitdownParser())
logger.debug("Registered markitdown parser")
except ImportError:
logger.warning("markitdown not available - file parsing disabled")
@@ -3334,28 +3287,6 @@ class MemoryEngine(MemoryEngineInterface):
sub_doc_id = document_id or (sub_batch[0].get("document_id") if len(sub_batch) == 1 else None)
sub_offset = chunk_offsets.get(sub_doc_id, 0) if sub_doc_id else 0
# Count the chunks this sub-batch will produce BEFORE handing it
# to the orchestrator. retain_batch consumes (pops) each item's
# "content" while streaming, so reading it back after the call
# yields "" — and chunk_text("") returns [""] (count 1),
# advancing the per-document cursor by 1 regardless of the real
# chunk count. For slices that each span several chunks the next
# sub-batch then restarts ~1 slot in, colliding chunk_ids and
# overwriting earlier chunks (only ~1 new chunk survives per
# sub-batch). Capture it here while content is still present.
sub_chunk_count = 0
if sub_doc_id:
sub_chunk_count = sum(
len(
fact_extraction.chunk_text(
item.get("content", "") or "",
chunking_config.chunk_size,
structured_chunk_size=chunking_config.structured_chunk_size,
)
)
for item in sub_batch
)
sub_results, sub_usage, sub_processed = await self._retain_batch_async_internal(
bank_id=bank_id,
contents=sub_batch,
@@ -3375,10 +3306,20 @@ class MemoryEngine(MemoryEngineInterface):
)
# Advance the document's chunk_index cursor by the number of
# chunks this sub-batch produced (counted above, before the
# orchestrator consumed the content), so the next sub-batch
# sharing the document continues the sequence.
# chunks this sub-batch produced (computed with the same chunk
# size the orchestrator uses), so the next sub-batch sharing the
# document continues the sequence.
if sub_doc_id:
sub_chunk_count = sum(
len(
fact_extraction.chunk_text(
item.get("content", "") or "",
chunking_config.chunk_size,
structured_chunk_size=chunking_config.structured_chunk_size,
)
)
for item in sub_batch
)
# retain_batch only prepends the existing body on the global
# first sub-batch (is_first_batch == i == 1), so fold its chunk
# count in only there.
@@ -3442,8 +3383,6 @@ class MemoryEngine(MemoryEngineInterface):
llm_input_tokens=total_usage.input_tokens,
llm_output_tokens=total_usage.output_tokens,
llm_total_tokens=total_usage.total_tokens,
llm_cached_input_tokens=getattr(total_usage, "cached_tokens", 0) or 0,
llm_thoughts_tokens=getattr(total_usage, "thoughts_tokens", 0) or 0,
processed_content_tokens=total_processed_content_tokens,
)
try:
@@ -3845,10 +3784,6 @@ class MemoryEngine(MemoryEngineInterface):
max_tokens: int = 4096,
enable_trace: bool = False,
fact_type: list[str] | None = None,
# Opt-in (default False). Internal callers that recall raw facts on purpose —
# notably consolidation, which needs the raw facts it folds into observations —
# must leave this off so they aren't silently deduped away.
prefer_observations: bool = False,
question_date: datetime | None = None,
include_entities: bool = False,
max_entity_tokens: int = 500,
@@ -3881,10 +3816,6 @@ class MemoryEngine(MemoryEngineInterface):
bank_id: bank ID to recall for
query: Recall query
fact_type: List of fact types to recall (e.g., ['world', 'experience'])
prefer_observations: When True and both 'observation' and a raw type ('world'/'experience')
are requested, drop raw facts that a returned observation was consolidated from
(deduplication by provenance). Freed slots backfill, keeping the result count at
the budget. No-op unless both observation and raw types are requested.
budget: Budget level for graph traversal (low=100, mid=300, high=600 units)
max_tokens: Maximum tokens to return (counts only 'text' field, default 4096)
Results are returned until token budget is reached, stopping before
@@ -4022,7 +3953,6 @@ class MemoryEngine(MemoryEngineInterface):
max_chunk_tokens,
request_context,
semaphore_wait=semaphore_wait,
prefer_observations=prefer_observations,
tags=tags,
tags_match=tags_match,
tag_groups=tag_groups,
@@ -4159,7 +4089,6 @@ class MemoryEngine(MemoryEngineInterface):
max_chunk_tokens: int = 8192,
request_context: "RequestContext" = None,
semaphore_wait: float = 0.0,
prefer_observations: bool = False,
tags: list[str] | None = None,
tags_match: TagsMatch = "any",
tag_groups: list[TagGroup] | None = None,
@@ -4634,14 +4563,10 @@ class MemoryEngine(MemoryEngineInterface):
ce = reranker_instance.cross_encoder
# "rrf" mode is passthrough by construction; so is a configured "rrf" CE.
is_passthrough = (reranking == "rrf") or (ce is not None and ce.provider_name == "rrf")
scoring_config = get_config()
apply_combined_scoring(
scored_results,
now=_recall_scoring_now(question_date),
is_passthrough_reranker=is_passthrough,
recency_decay_function=scoring_config.recency_decay_function,
recency_decay_linear_window_days=scoring_config.recency_decay_linear_window_days,
recency_decay_halflife_days=scoring_config.recency_decay_halflife_days,
)
# Per-strategy additive boost: nudge candidates surfaced by a
# prioritised retrieval arm up the final ordering.
@@ -4676,48 +4601,6 @@ class MemoryEngine(MemoryEngineInterface):
if request_context is not None:
request_context.raise_if_cancelled()
# Step 4.8: prefer-observations dedup. When the caller asked for observations
# alongside raw facts, an observation supersedes the raw facts it was
# consolidated from: drop those raw facts so the same content isn't returned
# twice. Runs BEFORE the Step 5 truncation so the freed slots backfill with
# the next-best results, keeping the result count at the budget. No-op unless
# 'observation' and at least one raw type were both requested.
raw_types_requested = {"world", "experience"} & set(fact_type)
if prefer_observations and "observation" in fact_type and raw_types_requested:
# "The observation list" = observations within the window we would return.
# Only those can supersede a raw fact; a far-down observation should not
# suppress a top raw fact it merely happens to reference.
observation_ids = [
uuid.UUID(sr.id)
for sr in scored_results[: thinking_budget * 2]
if sr.retrieval.fact_type == "observation"
]
if observation_ids:
superseded_ids: set[str] = set()
async with acquire_with_retry(backend) as dedup_conn:
obs_rows = await dedup_conn.fetch(
f"""
SELECT source_memory_ids
FROM {fq_table("memory_units")}
WHERE id = ANY($1::uuid[]) AND fact_type = 'observation'
""",
observation_ids,
)
for obs_row in obs_rows:
for sid in obs_row["source_memory_ids"] or []:
superseded_ids.add(str(sid))
if superseded_ids:
before_count = len(scored_results)
scored_results = [
sr
for sr in scored_results
if not (sr.retrieval.fact_type in ("world", "experience") and sr.id in superseded_ids)
]
log_buffer.append(
f" [4.8] prefer_observations: dropped {before_count - len(scored_results)} "
f"raw fact(s) superseded by {len(observation_ids)} observation(s)"
)
# Step 5: Truncate to thinking_budget * 2 for token filtering
rerank_limit = thinking_budget * 2
top_scored = scored_results[:rerank_limit]
@@ -5392,17 +5275,6 @@ class MemoryEngine(MemoryEngineInterface):
"memory_units_deleted": units_count if deleted else 0,
}
# Drop any cached stats for this bank — deleting the document changed
# the document count and (via cascade) the memory-unit/link counts
# get_bank_stats reports, which the TTL would otherwise serve at
# pre-delete values for up to a minute (mirrors delete_bank). Best-effort:
# a cache-eviction failure must not fail an already-committed delete.
if deleted:
try:
await self._bank_stats_cache.invalidate(get_current_schema(), bank_id)
except Exception as e:
logger.warning(f"Failed to invalidate bank stats cache after document deletion for bank {bank_id}: {e}")
if invalidated_obs > 0:
config = await self._config_resolver.resolve_full_config(bank_id, request_context)
if config.enable_auto_consolidation:
@@ -5570,14 +5442,6 @@ class MemoryEngine(MemoryEngineInterface):
)
if invalidated_obs > 0:
# Observation units were deleted, changing the counts get_bank_stats
# reports — drop the cached stats so the TTL does not serve pre-update
# values for up to a minute (mirrors delete_bank). Best-effort: a
# cache-eviction failure must not fail an already-committed update.
try:
await self._bank_stats_cache.invalidate(get_current_schema(), bank_id)
except Exception as e:
logger.warning(f"Failed to invalidate bank stats cache after document update for bank {bank_id}: {e}")
config = await self._config_resolver.resolve_full_config(bank_id, request_context)
if config.enable_auto_consolidation:
try:
@@ -5669,19 +5533,6 @@ class MemoryEngine(MemoryEngineInterface):
else "Memory unit not found",
}
# Drop any cached stats for this bank — the deleted unit (and its
# cascaded links/entities) changed the counts get_bank_stats reports,
# which the TTL would otherwise serve at pre-delete values for up to a
# minute (mirrors delete_bank). Best-effort: a cache-eviction failure
# must not fail an already-committed delete.
if deleted and bank_id:
try:
await self._bank_stats_cache.invalidate(get_current_schema(), bank_id)
except Exception as e:
logger.warning(
f"Failed to invalidate bank stats cache after memory unit deletion for bank {bank_id}: {e}"
)
if bank_id_for_consolidation:
config = await self._config_resolver.resolve_full_config(bank_id_for_consolidation, request_context)
if config.enable_auto_consolidation:
@@ -5904,17 +5755,7 @@ class MemoryEngine(MemoryEngineInterface):
bank_id,
)
# Drop any cached stats for this bank — clearing observations changed
# the memory-unit/observation counts and the consolidation timestamps
# get_bank_stats reports, which the TTL would otherwise serve at stale
# values for up to a minute (mirrors delete_bank). Best-effort: a
# cache-eviction failure must not fail an already-committed clear.
try:
await self._bank_stats_cache.invalidate(get_current_schema(), bank_id)
except Exception as e:
logger.warning(f"Failed to invalidate bank stats cache after clearing observations for bank {bank_id}: {e}")
return {"deleted_count": count or 0}
return {"deleted_count": count or 0}
async def list_observation_scopes(
self,
@@ -8271,15 +8112,25 @@ class MemoryEngine(MemoryEngineInterface):
# reflect_mission and disposition in config take precedence over the legacy DB columns
config_dict = await self._config_resolver.get_bank_config(bank_id, request_context)
mission = config_dict.get("reflect_mission") or profile["mission"]
# Overlay disposition from config if explicitly set; fall back to DB values
db_disp = profile["disposition"]
db_disp_dict = db_disp.model_dump() if hasattr(db_disp, "model_dump") else dict(db_disp)
resolved = _overlay_bank_config_disposition_mission(db_disp_dict, profile["mission"], config_dict)
cfg_skep = config_dict.get("disposition_skepticism")
cfg_lit = config_dict.get("disposition_literalism")
cfg_emp = config_dict.get("disposition_empathy")
disposition = {
"skepticism": cfg_skep if cfg_skep is not None else db_disp_dict["skepticism"],
"literalism": cfg_lit if cfg_lit is not None else db_disp_dict["literalism"],
"empathy": cfg_emp if cfg_emp is not None else db_disp_dict["empathy"],
}
return {
"bank_id": bank_id,
"name": profile["name"],
"disposition": resolved.disposition,
"mission": resolved.mission,
"disposition": disposition,
"mission": mission,
}
async def _ensure_bank_exists(
@@ -8494,17 +8345,6 @@ class MemoryEngine(MemoryEngineInterface):
BankListContext(banks=banks, request_context=request_context)
)
banks = result.banks
# Overlay resolved bank config (reflect_mission + disposition_*) on top of the
# legacy banks.disposition / banks.mission columns, mirroring get_bank_profile so
# the list and get paths return identical disposition + mission for a bank.
# Resolve every bank's config in one batch (single config-column query + a single
# tenant-config resolve) rather than one round-trip per bank.
configs = await self._config_resolver.get_bank_configs([bank["bank_id"] for bank in banks], request_context)
for bank in banks:
resolved = _overlay_bank_config_disposition_mission(
bank["disposition"], bank["mission"], configs.get(bank["bank_id"], {})
)
bank["disposition"], bank["mission"] = resolved.disposition, resolved.mission
return banks
# ==================== Reflect Methods ====================
@@ -11855,8 +11695,6 @@ class MemoryEngine(MemoryEngineInterface):
**task_payload,
}
from hindsight_api.extensions.operation_validator import OperationValidationError
async with acquire_with_retry(backend) as conn:
async with conn.transaction():
if dedupe_by_bank:
@@ -11878,20 +11716,10 @@ class MemoryEngine(MemoryEngineInterface):
# serialize) but not with FOR KEY SHARE (so those inserts proceed).
# On Oracle this rewrites to FOR UPDATE, which there does not block
# indexed-FK child inserts.
#
# Use fetchval so we can also verify the bank actually exists.
# Without this check, callers that race against bank deletion
# or that derive bank IDs before creating the bank reach the
# INSERT below and get an asyncpg.ForeignKeyViolationError, which
# surfaces as an opaque 500 from the API. A clean
# OperationValidationError(404) is the right shape — the FastAPI
# handler already converts it via its existing except clause.
bank_exists = await conn.fetchval(
await conn.execute(
f"SELECT 1 FROM {fq_table('banks')} WHERE bank_id = $1 FOR NO KEY UPDATE",
bank_id,
)
if bank_exists is None:
raise OperationValidationError(f"Bank '{bank_id}' not found", status_code=404)
# Only check 'pending', not 'processing': a processing task uses a
# watermark from when it started, so memories added after that need
# a fresh run regardless.
@@ -11921,16 +11749,6 @@ class MemoryEngine(MemoryEngineInterface):
"operation_id": str(row["operation_id"]),
"deduplicated": True,
}
else:
# Scoped/non-dedupe submits skip the lock + dedup above.
# Still verify the bank exists so an FK violation can't
# escape as a 500.
bank_exists = await conn.fetchval(
f"SELECT 1 FROM {fq_table('banks')} WHERE bank_id = $1",
bank_id,
)
if bank_exists is None:
raise OperationValidationError(f"Bank '{bank_id}' not found", status_code=404)
await conn.execute(
f"""
@@ -3,116 +3,43 @@
import asyncio
import logging
import tempfile
from dataclasses import dataclass
from pathlib import Path
from hindsight_api.config import DEFAULT_FILE_PARSER_MARKITDOWN_OCR_PROMPT
from .base import FileParser
logger = logging.getLogger(__name__)
@dataclass(frozen=True)
class MarkitdownOcrOptions:
"""OpenAI-compatible OCR options passed through to MarkItDown."""
# Keep this typed as object so the OpenAI SDK import stays lazy for non-OCR users.
llm_client: object
llm_model: str
llm_prompt: str
class MarkitdownParser(FileParser):
"""
Markitdown file parser.
Uses Microsoft's markitdown library to convert various file formats
to markdown including PDF, Office docs, images with optional OCR,
audio, HTML.
to markdown including PDF, Office docs, images (via OCR), audio, HTML.
Supported formats:
- PDF (.pdf)
- Word (.docx, .doc)
- PowerPoint (.pptx, .ppt)
- Excel (.xlsx, .xls)
- Images (.jpg, .jpeg, .png) - optional OCR
- Images (.jpg, .jpeg, .png) - with OCR
- HTML (.html, .htm)
- Text (.txt, .md)
- Audio (.mp3, .wav) - with transcription
"""
def __init__(
self,
*,
ocr_enabled: bool = False,
ocr_api_key: str | None = None,
ocr_base_url: str | None = None,
ocr_model: str | None = None,
ocr_prompt: str | None = None,
):
def __init__(self):
"""Initialize markitdown parser."""
# Lazy import to avoid requiring markitdown for all users
try:
from markitdown import MarkItDown
self._markitdown = MarkItDown()
except ImportError as e:
raise ImportError(
"markitdown package is required for file parsing. Install with: pip install markitdown"
) from e
self._ocr_enabled = ocr_enabled
if ocr_enabled:
ocr_options = self._build_ocr_options(
api_key=ocr_api_key,
base_url=ocr_base_url,
model=ocr_model,
prompt=ocr_prompt,
)
self._markitdown = MarkItDown(
llm_client=ocr_options.llm_client,
llm_model=ocr_options.llm_model,
llm_prompt=ocr_options.llm_prompt,
)
else:
self._markitdown = MarkItDown()
def _build_ocr_options(
self,
*,
api_key: str | None,
base_url: str | None,
model: str | None,
prompt: str | None,
) -> MarkitdownOcrOptions:
"""Build MarkItDown options for OpenAI-compatible image OCR."""
if not model or not model.strip():
raise ValueError(
"Markitdown OCR is enabled but no model is configured. "
"Set HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_MODEL to an OpenAI-compatible OCR/vision model "
"with image-input support."
)
if not api_key:
raise ValueError(
"Markitdown OCR is enabled but no API key is configured. "
"Set HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_API_KEY."
)
if not base_url or not base_url.strip():
raise ValueError(
"Markitdown OCR is enabled but no base URL is configured. "
"Set HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_BASE_URL to an OpenAI-compatible OCR/vision endpoint."
)
try:
from openai import OpenAI
except ImportError as e:
raise RuntimeError("openai package is required when Markitdown OCR is enabled.") from e
return MarkitdownOcrOptions(
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,
)
async def convert(self, file_data: bytes, filename: str) -> str:
"""Parse file to markdown using markitdown."""
# markitdown is synchronous, so we run it in executor to avoid blocking
@@ -121,13 +48,6 @@ class MarkitdownParser(FileParser):
def _convert_sync(self, file_data: bytes, filename: str) -> str:
"""Synchronous parsing (runs in thread pool)."""
if self._is_image_file(filename) and not self._ocr_enabled:
raise RuntimeError(
"Image OCR is not enabled for the markitdown parser. "
"Set HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_ENABLED=true and configure an OpenAI-compatible "
"OCR/vision endpoint with image-input support, or choose an OCR-capable parser."
)
# Write to temp file (markitdown requires file path)
with tempfile.NamedTemporaryFile(suffix=Path(filename).suffix, delete=False) as tmp:
tmp.write(file_data)
@@ -153,11 +73,6 @@ class MarkitdownParser(FileParser):
except Exception:
pass
@staticmethod
def _is_image_file(filename: str) -> bool:
"""Return whether the file type needs OCR to extract useful text."""
return Path(filename).suffix.lower() in {".jpg", ".jpeg", ".png"}
def supports(self, filename: str, content_type: str | None = None) -> bool:
"""Check if markitdown supports this file type."""
# Supported extensions (from markitdown docs)
@@ -170,7 +85,7 @@ class MarkitdownParser(FileParser):
".ppt",
".xlsx",
".xls",
# Images (optional OCR)
# Images (with OCR)
".jpg",
".jpeg",
".png",
@@ -136,9 +136,7 @@ class AnthropicLLM(LLMInterface):
initial_backoff: Initial backoff time in seconds.
max_backoff: Maximum backoff time in seconds.
skip_validation: Return raw JSON without Pydantic validation.
strict_schema: Route structured output through a forced tool_use tool for
native constrained decoding (issue #1002). When False, falls back to
schema-in-prompt + JSON parse.
strict_schema: Use strict JSON schema enforcement (not supported by Anthropic).
return_usage: If True, return tuple (result, TokenUsage) instead of just result.
Returns:
@@ -169,21 +167,14 @@ class AnthropicLLM(LLMInterface):
else:
anthropic_messages.append({"role": role, "content": content})
# Structured output: prefer Anthropic-native constrained decoding via a single
# forced tool_use tool (strict_schema) over text-injecting the schema and
# parsing the reply. Native constrained decoding guarantees schema-valid JSON,
# eliminating the invalid-JSON retry storm (issue #1002). When strict_schema is
# off we keep the text-inject + json.loads fallback for backward compatibility.
schema = None
use_forced_tool = False
_tool_name = "structured_response"
# Add JSON schema instruction if response_format is provided
if response_format is not None and hasattr(response_format, "model_json_schema"):
schema = response_format.model_json_schema()
if strict_schema:
use_forced_tool = True
schema_msg = f"\n\nYou must respond with valid JSON matching this schema:\n{json.dumps(schema, indent=2, ensure_ascii=False)}"
if system_prompt:
system_prompt += schema_msg
else:
schema_msg = f"\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
system_prompt = schema_msg
# Prepare parameters
call_params: dict[str, Any] = {
@@ -195,14 +186,6 @@ class AnthropicLLM(LLMInterface):
if system_prompt:
call_params["system"] = system_prompt
if use_forced_tool:
# Single tool whose input_schema IS the response schema; force the model to
# emit it via tool_choice so the SDK does constrained decoding for us.
call_params["tools"] = [
{"name": _tool_name, "description": "Return the structured response.", "input_schema": schema}
]
call_params["tool_choice"] = {"type": "tool", "name": _tool_name}
if self._extra_body:
call_params["extra_body"] = self._extra_body
@@ -212,49 +195,32 @@ class AnthropicLLM(LLMInterface):
try:
response = await self._client.messages.create(**call_params)
if use_forced_tool:
# Forced tool_use → the validated args are already a dict; no parsing,
# no markdown-strip, no JSON-decode retry possible.
tool_input = None
for block in response.content:
if block.type == "tool_use" and block.name == _tool_name:
tool_input = block.input or {}
break
if tool_input is None:
# Model ignored the forced tool (rare, e.g. a gateway that drops
# tool_choice). Fall back to text parse so we don't hard-fail; the
# existing retry loop still covers genuine errors.
content = "".join(b.text for b in response.content if b.type == "text")
tool_input = json.loads(content)
content = json.dumps(tool_input)
result = tool_input if skip_validation else response_format.model_validate(tool_input)
else:
# Anthropic response content is a list of blocks
content = ""
for block in response.content:
if block.type == "text":
content += block.text
# Anthropic response content is a list of blocks
content = ""
for block in response.content:
if block.type == "text":
content += block.text
if response_format is not None:
# Models may wrap JSON in markdown code blocks
clean_content = content
if "```json" in content:
clean_content = content.split("```json")[1].split("```")[0].strip()
elif "```" in content:
clean_content = content.split("```")[1].split("```")[0].strip()
if response_format is not None:
# Models may wrap JSON in markdown code blocks
clean_content = content
if "```json" in content:
clean_content = content.split("```json")[1].split("```")[0].strip()
elif "```" in content:
clean_content = content.split("```")[1].split("```")[0].strip()
try:
json_data = json.loads(clean_content)
except json.JSONDecodeError:
# Fallback to parsing raw content if markdown stripping failed
json_data = json.loads(content)
try:
json_data = json.loads(clean_content)
except json.JSONDecodeError:
# Fallback to parsing raw content if markdown stripping failed
json_data = json.loads(content)
if skip_validation:
result = json_data
else:
result = response_format.model_validate(json_data)
if skip_validation:
result = json_data
else:
result = content
result = response_format.model_validate(json_data)
else:
result = content
# Record metrics and log slow calls
duration = time.time() - start_time
@@ -60,22 +60,6 @@ _CODEX_TERMINAL_REFRESH_ERROR_CODES = frozenset(
)
def default_codex_auth_file() -> Path:
"""Return the path to Codex's ``auth.json``.
Honors the ``CODEX_HOME`` environment variable — the same variable the
canonical ``@openai/codex`` CLI uses to relocate its config/credentials
directory — and falls back to ``~/.codex`` when it is unset or empty.
Resolved lazily on each call (rather than cached at import time) so that
the environment is read at the point of use.
"""
codex_home = os.environ.get("CODEX_HOME")
if codex_home:
return Path(codex_home) / "auth.json"
return Path.home() / ".codex" / "auth.json"
class CodexRefreshExpiredError(RuntimeError):
"""Raised when the Codex refresh_token itself is no longer valid.
@@ -102,7 +86,7 @@ class CodexAuthManager:
The OAuth refresh token. May be ``None`` when the auth file omits it;
the provider still works as a one-shot loader in that case.
auth_file:
Path to the Codex ``auth.json``. Used for re-reading the refresh token
Path to ``~/.codex/auth.json``. Used for re-reading the refresh token
on demand and for atomic persistence of rotated credentials.
"""
@@ -131,8 +115,7 @@ class CodexAuthManager:
Parameters
----------
auth_file:
Defaults to ``$CODEX_HOME/auth.json`` (or ``~/.codex/auth.json``
when ``CODEX_HOME`` is unset).
Defaults to ``~/.codex/auth.json``.
Raises
------
@@ -143,7 +126,7 @@ class CodexAuthManager:
``auth_mode``.
"""
if auth_file is None:
auth_file = default_codex_auth_file()
auth_file = Path.home() / ".codex" / "auth.json"
if not auth_file.exists():
raise FileNotFoundError(f"Codex auth file not found: {auth_file}. Run 'codex auth login' to authenticate.")
@@ -2,9 +2,8 @@
OpenAI Codex LLM provider using ChatGPT Plus/Pro OAuth authentication.
This provider enables using ChatGPT Plus/Pro subscriptions for API calls
without separate OpenAI Platform API credits. It uses OAuth tokens from the
Codex ``auth.json`` (``$CODEX_HOME/auth.json``, or ``~/.codex/auth.json`` when
``CODEX_HOME`` is unset) and communicates with the ChatGPT backend API.
without separate OpenAI Platform API credits. It uses OAuth tokens from
~/.codex/auth.json and communicates with the ChatGPT backend API.
Tokens are refreshed automatically: the provider decodes the access_token
JWT's ``exp`` claim and proactively refreshes via
@@ -36,7 +35,6 @@ from .codex_auth import (
_CODEX_TOKEN_REFRESH_SKEW_SECONDS,
CodexAuthManager,
CodexRefreshExpiredError,
default_codex_auth_file,
)
# Re-export for backward compatibility (tests import from this module).
@@ -57,15 +55,14 @@ class CodexLLM(LLMInterface):
"""
LLM provider using OpenAI Codex OAuth authentication.
Authenticates using ChatGPT Plus/Pro credentials stored in the Codex
``auth.json`` (honoring ``CODEX_HOME``, default ``~/.codex``) and makes API
calls to chatgpt.com/backend-api/codex/responses.
Authenticates using ChatGPT Plus/Pro credentials stored in ~/.codex/auth.json
and makes API calls to chatgpt.com/backend-api/codex/responses.
"""
def __init__(
self,
provider: str,
api_key: str, # Will be ignored, reads from the Codex auth.json (CODEX_HOME or ~/.codex)
api_key: str, # Will be ignored, reads from ~/.codex/auth.json
base_url: str,
model: str,
reasoning_effort: str = "low",
@@ -84,14 +81,12 @@ class CodexLLM(LLMInterface):
refresh_token = self._load_codex_refresh_token()
logger.info(f"Loaded Codex OAuth credentials for account: {account_id}")
except Exception as e:
auth_file = default_codex_auth_file()
raise RuntimeError(
f"Failed to load Codex OAuth credentials from {auth_file}: {e}\n\n"
f"Failed to load Codex OAuth credentials from ~/.codex/auth.json: {e}\n\n"
"To set up Codex authentication:\n"
"1. Install Codex CLI: npm install -g @openai/codex\n"
"2. Login: codex auth login\n"
f"3. Verify: ls {auth_file}\n\n"
"(Set CODEX_HOME to use a credentials directory other than ~/.codex.)\n\n"
"3. Verify: ls ~/.codex/auth.json\n\n"
"Or use a different provider (openai, anthropic, gemini) with API keys."
) from e
@@ -99,7 +94,7 @@ class CodexLLM(LLMInterface):
access_token=access_token,
account_id=account_id,
refresh_token=refresh_token,
auth_file=default_codex_auth_file(),
auth_file=Path.home() / ".codex" / "auth.json",
)
# Use ChatGPT backend API endpoint. Codex auth is tied to
@@ -161,7 +156,7 @@ class CodexLLM(LLMInterface):
def _load_codex_auth(self) -> tuple[str, str]:
"""
Load OAuth credentials from the Codex ``auth.json`` (CODEX_HOME or ~/.codex).
Load OAuth credentials from ~/.codex/auth.json.
Returns:
Tuple of (access_token, account_id).
@@ -170,7 +165,7 @@ class CodexLLM(LLMInterface):
FileNotFoundError: If auth file doesn't exist.
ValueError: If auth file is invalid.
"""
auth_file = default_codex_auth_file()
auth_file = Path.home() / ".codex" / "auth.json"
if not auth_file.exists():
raise FileNotFoundError(
@@ -202,7 +197,9 @@ class CodexLLM(LLMInterface):
pre- and post-``__init__`` because it does not depend on
``_auth_manager`` being constructed yet.
"""
auth_file = self._auth_manager._auth_file if hasattr(self, "_auth_manager") else default_codex_auth_file()
auth_file = (
self._auth_manager._auth_file if hasattr(self, "_auth_manager") else Path.home() / ".codex" / "auth.json"
)
return CodexAuthManager.load_refresh_token_from_file(auth_file)
@staticmethod
@@ -76,7 +76,6 @@ class GeminiLLM(LLMInterface):
# Safety settings: None means use Gemini's defaults
self._safety_settings: list | None = kwargs.get("gemini_safety_settings")
self._service_tier: str | None = kwargs.get("gemini_service_tier")
# User-configured extra params merged into the GenerateContentConfig of
# every call. Gemini's request body nests generation params, so we expose
@@ -107,16 +106,6 @@ class GeminiLLM(LLMInterface):
self._client = genai.Client(api_key=self.api_key)
logger.info(f"Gemini API: model={self.model}")
def _apply_service_tier(self, config_kwargs: dict[str, Any]) -> None:
if not self._service_tier:
return
http_options = dict(config_kwargs.get("http_options") or {})
extra_body = dict(http_options.get("extra_body") or {})
extra_body.setdefault("service_tier", self._service_tier)
http_options["extra_body"] = extra_body
config_kwargs["http_options"] = http_options
def _init_vertexai(self, **kwargs: Any) -> None:
"""Initialize Vertex AI client with project, region, and credentials."""
# Extract Vertex AI config from kwargs
@@ -258,13 +247,16 @@ class GeminiLLM(LLMInterface):
else:
gemini_contents.append(genai_types.Content(role="user", parts=[genai_types.Part(text=content)]))
def _system_instruction_with_schema() -> str:
# Add the JSON schema as a textual hint in the system_instruction (matching
# the normal uncached path). Structured output is still enforced via
# response_schema regardless; this is just guidance text.
if response_format is not None and hasattr(response_format, "model_json_schema"):
schema = response_format.model_json_schema()
schema_msg = (
f"\n\nYou must respond with valid JSON matching this schema:\n"
f"{json.dumps(schema, indent=2, ensure_ascii=False)}"
)
return (system_instruction + schema_msg) if system_instruction else schema_msg
schema_msg = f"\n\nYou must respond with valid JSON matching this schema:\n{json.dumps(schema, indent=2, ensure_ascii=False)}"
if system_instruction:
system_instruction += schema_msg
else:
system_instruction = schema_msg
# Apply safety settings: context var (per-request bank override) takes precedence over instance default
effective_safety_settings = _safety_settings_ctx.get()
@@ -281,18 +273,11 @@ class GeminiLLM(LLMInterface):
def _build_generation_config(use_cache: bool) -> "genai_types.GenerateContentConfig | None":
# Seed with user-configured extra params; explicit settings below win.
config_kwargs: dict[str, Any] = dict(self._extra_body)
self._apply_service_tier(config_kwargs)
if use_cache:
config_kwargs["cached_content"] = cached_prefix
elif (
use_schema_prompt_fallback
and response_format is not None
and hasattr(response_format, "model_json_schema")
):
config_kwargs["system_instruction"] = _system_instruction_with_schema()
elif system_instruction:
config_kwargs["system_instruction"] = system_instruction
if response_format is not None and not use_schema_prompt_fallback:
if response_format is not None:
config_kwargs["response_mime_type"] = "application/json"
config_kwargs["response_schema"] = response_format
if temperature is not None:
@@ -310,7 +295,6 @@ class GeminiLLM(LLMInterface):
return genai_types.GenerateContentConfig(**config_kwargs) if config_kwargs else None
cache_active = using_cache
use_schema_prompt_fallback = False
generation_config = _build_generation_config(cache_active)
last_exception = None
@@ -428,26 +412,12 @@ class GeminiLLM(LLMInterface):
output_tokens=output_tokens,
total_tokens=input_tokens + output_tokens,
cached_tokens=cached_tokens,
thoughts_tokens=thoughts_tokens,
)
return result, token_usage
return result
except json.JSONDecodeError as e:
last_exception = e
if (
attempt < max_retries
and response_format is not None
and hasattr(response_format, "model_json_schema")
and not cache_active
and not use_schema_prompt_fallback
):
logger.warning("Gemini returned invalid JSON, retrying with prompt-side schema guidance...")
cache_active = False
use_schema_prompt_fallback = True
generation_config = _build_generation_config(cache_active)
await asyncio.sleep(min(initial_backoff * (2**attempt), max_backoff))
continue
if attempt < max_retries:
logger.warning("Gemini returned invalid JSON, retrying...")
backoff = min(initial_backoff * (2**attempt), max_backoff)
@@ -634,7 +604,6 @@ class GeminiLLM(LLMInterface):
def _build_tools_config(use_cache: bool) -> "genai_types.GenerateContentConfig":
# Seed with user-configured extra params; explicit settings below win.
config_kwargs: dict[str, Any] = dict(self._extra_body)
self._apply_service_tier(config_kwargs)
if use_cache:
config_kwargs["cached_content"] = cached_prefix
else:
@@ -780,8 +749,6 @@ class GeminiLLM(LLMInterface):
finish_reason=finish_reason,
input_tokens=input_tokens,
output_tokens=output_tokens,
cached_tokens=cached_input_tokens,
thoughts_tokens=thoughts_tokens,
)
except genai_errors.APIError as e:
@@ -15,13 +15,9 @@ is handled automatically by LiteLLM.
import asyncio
import json
import logging
import os
import time
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 LLMInterface, OutputTooLongError
from hindsight_api.engine.response_models import LLMToolCall, LLMToolCallResult, TokenUsage
from hindsight_api.metrics import get_metrics_collector
@@ -51,15 +47,13 @@ class LiteLLMLLM(LLMInterface):
base_url: str,
model: str,
reasoning_effort: str = "low",
timeout: float | None = None,
timeout: float = 300.0,
extra_body: dict[str, Any] | None = None,
bedrock_service_tier: str | None = None,
**kwargs: Any,
):
super().__init__(provider, api_key, base_url, model, reasoning_effort, **kwargs)
# ``None`` falls back to HINDSIGHT_API_LLM_TIMEOUT, then DEFAULT_LLM_TIMEOUT — never None,
# so the hard ``asyncio.wait_for`` backstop in ``call`` is always bounded.
self.timeout = timeout if timeout is not None else float(os.getenv(ENV_LLM_TIMEOUT, str(DEFAULT_LLM_TIMEOUT)))
self.timeout = timeout
self._litellm: Any = None
# User-configured extra params merged as top-level kwargs into every
# completion call so LiteLLM normalizes them per-provider (e.g. maps
@@ -215,10 +209,7 @@ class LiteLLMLLM(LLMInterface):
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),
timeout=self.timeout,
)
response = await self._acompletion(**call_kwargs)
content = response.choices[0].message.content or ""
finish_reason = response.choices[0].finish_reason
@@ -313,25 +304,6 @@ class LiteLLMLLM(LLMInterface):
logger.error(f"LiteLLM returned invalid JSON after {max_retries + 1} attempts")
raise
except (TimeoutError, asyncio.TimeoutError, LiteLLMTimeout) as e:
# litellm/httpx don't always honor their own ``timeout=`` (e.g. a connection held
# open with no token progress), so ``wait_for`` is the hard cap that cancels a hung
# call regardless — otherwise one straggler pins a worker slot and stalls its gather.
last_exception = e
exc_name = type(e).__name__
if attempt < max_retries:
logger.warning(
f"LiteLLM call exceeded timeout={self.timeout}s ({exc_name}, scope={scope}), retrying..."
)
backoff = min(initial_backoff * (2**attempt), max_backoff)
await asyncio.sleep(backoff)
continue
logger.error(
f"LiteLLM call timed out after {self.timeout}s on {attempt + 1} attempts "
f"({exc_name}, scope={scope})"
)
raise
except Exception as e:
error_str = str(e).lower()
# Fast fail on auth errors
@@ -382,10 +354,7 @@ class LiteLLMLLM(LLMInterface):
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),
timeout=self.timeout,
)
response = await self._acompletion(**call_kwargs)
message = response.choices[0].message
content = message.content
@@ -455,23 +424,6 @@ class LiteLLMLLM(LLMInterface):
output_tokens=output_tokens,
)
except (TimeoutError, asyncio.TimeoutError, LiteLLMTimeout) as e:
# See ``call`` — hard cap so a hung completion cannot block
# forever and pin a worker slot / concurrency permit.
last_exception = e
exc_name = type(e).__name__
if attempt < max_retries:
logger.warning(
f"LiteLLM tool call exceeded timeout={self.timeout}s ({exc_name}, scope={scope}), retrying..."
)
await asyncio.sleep(min(initial_backoff * (2**attempt), max_backoff))
continue
logger.error(
f"LiteLLM tool call timed out after {self.timeout}s on {attempt + 1} attempts "
f"({exc_name}, scope={scope})"
)
raise
except Exception as e:
error_str = str(e).lower()
if "401" in error_str or "403" in error_str or "unauthorized" in error_str:
@@ -67,7 +67,7 @@ class LiteLLMRouterLLM(LiteLLMLLM):
model: str,
config: dict[str, Any],
reasoning_effort: str = "low",
timeout: float | None = None,
timeout: float = 300.0,
**kwargs: Any,
):
super().__init__(
@@ -26,8 +26,6 @@ import logging
import os
import re
import time
from datetime import UTC, datetime, timedelta
from email.utils import parsedate_to_datetime
from typing import Any
from urllib.parse import parse_qs, urlparse, urlunparse
@@ -36,7 +34,7 @@ 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 LLMInterface, OutputTooLongError, ProviderRateLimitResetError
from hindsight_api.engine.llm_interface import LLMInterface, OutputTooLongError
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
@@ -85,49 +83,6 @@ def _strip_code_fences(content: str) -> str:
return content
# Reasoning/thinking tags emitted by extended-thinking models. Some providers
# (e.g. MiniMax-M3) leak the chain-of-thought wrapped in these tags into the
# response body instead of a separate reasoning_content field. Each entry is
# (open_tag, close_tag); the open tag also matches when the close tag is missing
# (truncated output) so a dangling block is removed to end-of-string.
_REASONING_TAG_PAIRS: tuple[tuple[str, str], ...] = (
("<think>", "</think>"),
("<thinking>", "</thinking>"),
("<thought>", "</thought>"),
("<reasoning>", "</reasoning>"),
("|startthink|", "|endthink|"),
)
def _strip_reasoning_tags(text: str) -> str:
"""Strip extended-thinking/reasoning blocks from an LLM response.
Removes the full set of tag styles emitted by reasoning models:
``<think>``, ``<thinking>``, ``<thought>``, ``<reasoning>`` and the
``|startthink|...|endthink|`` markers. Both the structured (JSON) path and
the free-form path must call this — otherwise a non-structured response
(e.g. a mental-model markdown blob from MiniMax-M3) leaks the raw
``<think>...</think>`` verbatim into stored memories.
Handles two cases:
1. Closed blocks: ``<think>...</think>`` removed wherever they appear.
2. Unclosed blocks: a dangling ``<think>`` with no closing tag (model output
truncated mid-thought) is removed from the open tag to end-of-string.
Returns the input unchanged (modulo surrounding whitespace) when no tags are
present.
"""
if not text:
return text
for open_tag, close_tag in _REASONING_TAG_PAIRS:
open_re = re.escape(open_tag)
close_re = re.escape(close_tag)
# Closed blocks first, then any remaining unclosed (truncated) block.
text = re.sub(rf"{open_re}.*?{close_re}", "", text, flags=re.DOTALL)
text = re.sub(rf"{open_re}.*", "", text, flags=re.DOTALL)
return text.strip()
def _response_get(response: Any, key: str, default: Any = None) -> Any:
if isinstance(response, dict):
return response.get(key, default)
@@ -279,122 +234,6 @@ def _summarize_status_error(e: APIStatusError, body_max: int = 400) -> str:
return f"HTTP {e.status_code}: {body_str or '<no body>'}"
_RATE_LIMIT_RESET_AT_RE = re.compile(
r"\breset at\s+"
r"(?P<reset_at>\d{4}-\d{2}-\d{2}[ T]\d{2}:\d{2}:\d{2}(?:\s*(?:Z|[+-]\d{2}:?\d{2}))?)",
re.IGNORECASE,
)
_RATE_LIMIT_WINDOW_RE = re.compile(
r"\b(?:for|in)\s+(?P<amount>\d+)\s*(?P<unit>second|minute|hour|day)s?\b",
re.IGNORECASE,
)
def _status_error_body_text(e: APIStatusError) -> str:
body: Any = getattr(e, "body", None)
if body is None:
try:
body = e.response.text
except Exception:
body = None
if isinstance(body, (dict, list)):
try:
return json.dumps(body, default=str, ensure_ascii=False)
except Exception:
return str(body)
return str(body or "").strip()
def _parse_retry_after_header(value: str | None, now: datetime) -> datetime | None:
if not value:
return None
raw = value.strip()
try:
seconds = float(raw)
except ValueError:
seconds = -1.0
if seconds >= 0:
return now + timedelta(seconds=seconds)
try:
parsed = parsedate_to_datetime(raw)
except (TypeError, ValueError, IndexError, OverflowError):
return None
if parsed.tzinfo is None:
parsed = parsed.replace(tzinfo=UTC)
return parsed.astimezone(UTC)
def _parse_reset_at_datetime(value: str) -> datetime | None:
raw = value.strip().replace(" ", "T")
if raw.endswith("Z"):
raw = f"{raw[:-1]}+00:00"
elif re.search(r"[+-]\d{4}$", raw):
raw = f"{raw[:-2]}:{raw[-2:]}"
try:
parsed = datetime.fromisoformat(raw)
except ValueError:
return None
if parsed.tzinfo is None:
# Some providers (z.ai included) return a wall-clock reset timestamp
# without a zone. Interpret it in the host's local zone so logs, status
# pages, and the queued next_retry_at describe the same operator-facing
# clock instead of silently shifting by UTC offset.
parsed = parsed.astimezone()
return parsed.astimezone(UTC)
def _rate_limit_retry_at(e: APIStatusError) -> datetime | None:
now = datetime.now(UTC)
response = getattr(e, "response", None)
headers = getattr(response, "headers", None)
if headers is not None:
retry_at = _parse_retry_after_header(headers.get("retry-after") or headers.get("Retry-After"), now)
if retry_at is not None and retry_at > now:
return retry_at
body_text = _status_error_body_text(e)
reset_match = _RATE_LIMIT_RESET_AT_RE.search(body_text)
if reset_match:
retry_at = _parse_reset_at_datetime(reset_match.group("reset_at"))
if retry_at is not None and retry_at > now:
return retry_at
window_match = _RATE_LIMIT_WINDOW_RE.search(body_text)
if not window_match:
return None
amount = int(window_match.group("amount"))
unit = window_match.group("unit").lower()
if unit == "second":
seconds = amount
elif unit == "minute":
seconds = amount * 60
elif unit == "hour":
seconds = amount * 3600
else:
seconds = amount * 86400
return now + timedelta(seconds=seconds)
def _raise_provider_quota_defer(
e: APIStatusError, *, provider: str, model: str, scope: str, max_backoff: float
) -> None:
if e.status_code != 429:
return
retry_at = _rate_limit_retry_at(e)
if retry_at is None:
return
if (retry_at - datetime.now(UTC)).total_seconds() <= max_backoff:
return
summary = _summarize_status_error(e)
raise ProviderRateLimitResetError(
retry_at=retry_at,
message=(
f"Provider quota exhausted ({provider}/{model}, scope={scope}); retry at {retry_at.isoformat()}: {summary}"
),
) from e
class OpenAICompatibleLLM(LLMInterface):
"""
LLM provider for OpenAI-compatible APIs.
@@ -430,7 +269,7 @@ class OpenAICompatibleLLM(LLMInterface):
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 ("low", "medium", "high").
timeout: Request timeout in seconds (uses env var or 120s default).
timeout: Request timeout in seconds (uses env var or 300s default).
groq_service_tier: Groq service tier ("on_demand", "flex", "auto").
extra_body: Extra body params merged into every API call.
**kwargs: Additional provider-specific parameters.
@@ -451,7 +290,6 @@ class OpenAICompatibleLLM(LLMInterface):
"openrouter",
"zai",
"opencode-go",
"atlas",
"fireworks",
]
if self.provider not in valid_providers:
@@ -477,8 +315,6 @@ class OpenAICompatibleLLM(LLMInterface):
self.base_url = "https://api.z.ai/api/coding/paas/v4"
elif self.provider == "opencode-go":
self.base_url = "https://opencode.ai/zen/go/v1"
elif self.provider == "atlas":
self.base_url = "https://api.atlascloud.ai/v1"
elif self.provider == "fireworks":
# OpenAI-compatible inference host (online path). The batch API
# lives on a separate control-plane host — see FireworksLLM.
@@ -499,7 +335,6 @@ class OpenAICompatibleLLM(LLMInterface):
"openrouter",
"zai",
"opencode-go",
"atlas",
"ollama-cloud",
)
and not self.api_key
@@ -782,10 +617,15 @@ class OpenAICompatibleLLM(LLMInterface):
scope=scope,
)
# Strip reasoning model thinking tags (closed and unclosed).
# Strip reasoning model thinking tags
# Supports: <think>, <thinking>, <thought>, <reasoning>, |startthink|/|endthink|
original_len = len(content)
content = _strip_reasoning_tags(content)
content = re.sub(r"<think>.*?</think>", "", content, flags=re.DOTALL)
content = re.sub(r"<thinking>.*?</thinking>", "", content, flags=re.DOTALL)
content = re.sub(r"<thought>.*?</thought>", "", content, flags=re.DOTALL)
content = re.sub(r"<reasoning>.*?</reasoning>", "", content, flags=re.DOTALL)
content = re.sub(r"\|startthink\|.*?\|endthink\|", "", content, flags=re.DOTALL)
content = content.strip()
if len(content) < original_len:
logger.debug(f"Stripped {original_len - len(content)} chars of reasoning tokens")
@@ -834,13 +674,6 @@ class OpenAICompatibleLLM(LLMInterface):
scope=scope,
)
# Free-form (non-structured) output also leaks reasoning tags:
# reasoning models like MiniMax-M3 wrap their chain-of-thought
# in <think>...</think> in the response body. Without this strip
# a mental-model markdown blob is stored verbatim with the raw
# thinking tags. Mirrors the structured-output path above.
result = _strip_reasoning_tags(result)
# Record token usage metrics
duration = time.time() - start_time
usage = response.usage
@@ -928,10 +761,6 @@ class OpenAICompatibleLLM(LLMInterface):
logger.error(f"Auth error (HTTP {e.status_code}), not retrying: {str(e)}")
raise
_raise_provider_quota_defer(
e, provider=self.provider, model=self.model, scope=scope, max_backoff=max_backoff
)
# Handle tool_use_failed error - model outputted in tool call format
if e.status_code == 400 and response_format is not None:
try:
@@ -985,6 +814,7 @@ class OpenAICompatibleLLM(LLMInterface):
f"scope={scope}): {_summarize_status_error(e)}"
)
raise
except ProviderResponseError as e:
last_exception = e
if e.retryable and attempt < max_retries:
@@ -1217,10 +1047,6 @@ class OpenAICompatibleLLM(LLMInterface):
f"not retrying: {_summarize_status_error(e)}"
)
raise
_raise_provider_quota_defer(
e, provider=self.provider, model=self.model, scope=scope, max_backoff=max_backoff
)
last_exception = e
if attempt < max_retries:
logger.warning(
@@ -1234,6 +1060,7 @@ class OpenAICompatibleLLM(LLMInterface):
f"({self.provider}/{self.model}, scope={scope}): {_summarize_status_error(e)}"
)
raise
except Exception:
raise
@@ -15,7 +15,7 @@ import time
from typing import TYPE_CHECKING, Any, Awaitable, Callable
from ...config import get_config
from .models import DirectiveInfo, LLMCall, ReflectAgentResult, StructuredOutputResult, TokenUsageSummary, ToolCall
from .models import DirectiveInfo, LLMCall, ReflectAgentResult, TokenUsageSummary, ToolCall
from .prompts import (
_extract_directive_rules,
build_final_prompt,
@@ -141,7 +141,7 @@ async def _generate_structured_output(
response_schema: dict,
llm_config: "LLMProvider",
reflect_id: str,
) -> StructuredOutputResult:
) -> tuple[dict[str, Any] | None, int, int]:
"""Generate structured output from an answer using the provided JSON schema.
Args:
@@ -151,8 +151,8 @@ async def _generate_structured_output(
reflect_id: Reflect ID for logging
Returns:
A StructuredOutputResult carrying the structured output (None if
generation fails) and the call's token usage.
Tuple of (structured_output, input_tokens, output_tokens).
structured_output is None if generation fails.
"""
try:
from typing import Any as TypingAny
@@ -186,7 +186,7 @@ async def _generate_structured_output(
if not fields:
logger.warning(f"[REFLECT {reflect_id}] No fields found in response_schema, skipping structured output")
return StructuredOutputResult()
return None, 0, 0
DynamicModel = create_model("StructuredResponse", **fields)
@@ -259,17 +259,11 @@ OUTPUT:"""
logger.warning(f"[REFLECT {reflect_id}] Required field '{field_name}' is empty in structured output")
logger.info(f"[REFLECT {reflect_id}] Generated structured output with {len(structured_output)} fields")
return StructuredOutputResult(
structured_output=structured_output,
input_tokens=usage.input_tokens,
output_tokens=usage.output_tokens,
cached_tokens=usage.cached_tokens,
thoughts_tokens=usage.thoughts_tokens,
)
return structured_output, usage.input_tokens, usage.output_tokens
except Exception as e:
logger.warning(f"[REFLECT {reflect_id}] Failed to generate structured output: {e}")
return StructuredOutputResult()
return None, 0, 0
def _count_messages_tokens(messages: list[dict[str, Any]]) -> int:
@@ -441,14 +435,9 @@ async def run_reflect_agent(
llm_trace: list[dict[str, Any]] = []
context_history: list[dict[str, Any]] = [] # For final prompt fallback
# Token usage tracking - accumulate across all LLM calls.
# cached_tokens and thoughts_tokens are surfaced for cost attribution
# and prompt-cache tuning. Both are subsets of (or parallel to) the
# input/output counts and are NOT double-counted in total_tokens.
# Token usage tracking - accumulate across all LLM calls
total_input_tokens = 0
total_output_tokens = 0
total_cached_tokens = 0
total_thoughts_tokens = 0
# Track available IDs for validation (prevents hallucinated citations)
available_memory_ids: set[str] = set()
@@ -471,8 +460,6 @@ async def run_reflect_agent(
input_tokens=total_input_tokens,
output_tokens=total_output_tokens,
total_tokens=total_input_tokens + total_output_tokens,
cached_tokens=total_cached_tokens,
thoughts_tokens=total_thoughts_tokens,
)
def _log_completion(answer: str, iterations: int, forced: bool = False):
@@ -539,8 +526,6 @@ async def run_reflect_agent(
llm_duration = int((time.time() - llm_start) * 1000)
total_input_tokens += usage.input_tokens
total_output_tokens += usage.output_tokens
total_cached_tokens += getattr(usage, "cached_tokens", 0) or 0
total_thoughts_tokens += getattr(usage, "thoughts_tokens", 0) or 0
llm_trace.append(
{
"scope": "final",
@@ -554,12 +539,11 @@ async def run_reflect_agent(
# Generate structured output if schema provided
structured_output = None
if response_schema and answer:
struct = await _generate_structured_output(answer, response_schema, llm_config, reflect_id)
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
structured_output, struct_in, struct_out = await _generate_structured_output(
answer, response_schema, llm_config, reflect_id
)
total_input_tokens += struct_in
total_output_tokens += struct_out
_log_completion(answer, iteration + 1, forced=True)
return ReflectAgentResult(
@@ -604,8 +588,6 @@ async def run_reflect_agent(
llm_duration = int((time.time() - llm_start) * 1000)
total_input_tokens += usage.input_tokens
total_output_tokens += usage.output_tokens
total_cached_tokens += getattr(usage, "cached_tokens", 0) or 0
total_thoughts_tokens += getattr(usage, "thoughts_tokens", 0) or 0
llm_trace.append(
{
"scope": "final",
@@ -618,12 +600,11 @@ async def run_reflect_agent(
structured_output = None
if response_schema and answer:
struct = await _generate_structured_output(answer, response_schema, llm_config, reflect_id)
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
structured_output, struct_in, struct_out = await _generate_structured_output(
answer, response_schema, llm_config, reflect_id
)
total_input_tokens += struct_in
total_output_tokens += struct_out
_log_completion(answer, iteration + 1, forced=True)
return ReflectAgentResult(
@@ -680,8 +661,6 @@ async def run_reflect_agent(
consecutive_errors = 0
total_input_tokens += result.input_tokens
total_output_tokens += result.output_tokens
total_cached_tokens += getattr(result, "cached_tokens", 0) or 0
total_thoughts_tokens += getattr(result, "thoughts_tokens", 0) or 0
llm_trace.append(
{
"scope": f"agent_{iteration + 1}",
@@ -730,8 +709,6 @@ async def run_reflect_agent(
llm_duration = int((time.time() - llm_start) * 1000)
total_input_tokens += usage.input_tokens
total_output_tokens += usage.output_tokens
total_cached_tokens += getattr(usage, "cached_tokens", 0) or 0
total_thoughts_tokens += getattr(usage, "thoughts_tokens", 0) or 0
llm_trace.append(
{
"scope": "final",
@@ -745,12 +722,11 @@ async def run_reflect_agent(
# Generate structured output if schema provided
structured_output = None
if response_schema and answer:
struct = await _generate_structured_output(answer, response_schema, llm_config, reflect_id)
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
structured_output, struct_in, struct_out = await _generate_structured_output(
answer, response_schema, llm_config, reflect_id
)
total_input_tokens += struct_in
total_output_tokens += struct_out
_log_completion(answer, iteration + 1, forced=True)
return ReflectAgentResult(
@@ -807,8 +783,6 @@ async def run_reflect_agent(
)
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",
@@ -822,12 +796,11 @@ async def run_reflect_agent(
# Generate structured output if schema provided
structured_output = None
if response_schema and answer:
struct = await _generate_structured_output(answer, response_schema, llm_config, reflect_id)
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
structured_output, struct_in, struct_out = await _generate_structured_output(
answer, response_schema, llm_config, reflect_id
)
total_input_tokens += struct_in
total_output_tokens += struct_out
_log_completion(answer, iteration + 1)
return ReflectAgentResult(
@@ -862,8 +835,6 @@ async def run_reflect_agent(
llm_duration = int((time.time() - llm_start) * 1000)
total_input_tokens += usage.input_tokens
total_output_tokens += usage.output_tokens
total_cached_tokens += getattr(usage, "cached_tokens", 0) or 0
total_thoughts_tokens += getattr(usage, "thoughts_tokens", 0) or 0
llm_trace.append(
{
"scope": "final",
@@ -877,12 +848,11 @@ async def run_reflect_agent(
# Generate structured output if schema provided
structured_output = None
if response_schema and answer:
struct = await _generate_structured_output(answer, response_schema, llm_config, reflect_id)
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
structured_output, struct_in, struct_out = await _generate_structured_output(
answer, response_schema, llm_config, reflect_id
)
total_input_tokens += struct_in
total_output_tokens += struct_out
_log_completion(answer, iteration + 1, forced=True)
return ReflectAgentResult(
@@ -1177,15 +1147,14 @@ async def _process_done_tool(
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)
structured_output = struct.structured_output
structured_output, struct_in, struct_out = await _generate_structured_output(
answer, response_schema, llm_config, reflect_id
)
# Add structured output tokens to usage
final_usage = TokenUsageSummary(
input_tokens=usage.input_tokens + struct.input_tokens,
output_tokens=usage.output_tokens + struct.output_tokens,
total_tokens=usage.total_tokens + struct.input_tokens + struct.output_tokens,
cached_tokens=usage.cached_tokens + struct.cached_tokens,
thoughts_tokens=usage.thoughts_tokens + struct.thoughts_tokens,
input_tokens=usage.input_tokens + struct_in,
output_tokens=usage.output_tokens + struct_out,
total_tokens=usage.total_tokens + struct_in + struct_out,
)
log_completion(answer, iterations)
@@ -78,32 +78,9 @@ class DirectiveInfo(BaseModel):
class TokenUsageSummary(BaseModel):
"""Total token usage across all LLM calls."""
input_tokens: int = Field(default=0, description="Total input tokens used (includes any cached prefix tokens)")
output_tokens: int = Field(default=0, description="Total visible output tokens used (excludes reasoning/thoughts)")
total_tokens: int = Field(default=0, description="Total tokens (input + output, excludes thoughts)")
cached_tokens: int = Field(
default=0,
description="Cached/cache-read prompt tokens summed across calls. Subset of input_tokens.",
)
thoughts_tokens: int = Field(
default=0,
description=(
"Reasoning/thinking tokens summed across calls. Billed at the output rate by some providers "
"but not part of visible output."
),
)
class StructuredOutputResult(BaseModel):
"""Result of structured-output generation, including token usage for the call."""
structured_output: dict[str, Any] | None = Field(
default=None, description="Generated structured output, or None if generation failed"
)
input_tokens: int = Field(default=0, description="Input tokens used")
output_tokens: int = Field(default=0, description="Visible output tokens used")
cached_tokens: int = Field(default=0, description="Cached prefix tokens. Subset of input_tokens.")
thoughts_tokens: int = Field(default=0, description="Reasoning/thinking tokens, when reported by the provider")
input_tokens: int = Field(default=0, description="Total input tokens used")
output_tokens: int = Field(default=0, description="Total output tokens used")
total_tokens: int = Field(default=0, description="Total tokens (input + output)")
class ReflectAgentResult(BaseModel):
@@ -31,20 +31,8 @@ class LLMToolCallResult(BaseModel):
content: str | None = Field(default=None, description="Text content if any")
tool_calls: list[LLMToolCall] = Field(default_factory=list, description="Tool calls requested by the LLM")
finish_reason: str | None = Field(default=None, description="Reason the LLM stopped: 'stop', 'tool_calls', etc.")
input_tokens: int = Field(
default=0,
description="Input tokens used in this call (includes any cached prefix tokens reported by the provider)",
)
output_tokens: int = Field(
default=0, description="Visible output tokens used in this call (excludes reasoning/thoughts)"
)
cached_tokens: int = Field(
default=0, description="Cached prefix tokens, when reported by the provider. Subset of input_tokens."
)
thoughts_tokens: int = Field(
default=0,
description="Reasoning/thinking tokens. Billed at the output rate by some providers but not part of visible output.",
)
input_tokens: int = Field(default=0, description="Input tokens used in this call")
output_tokens: int = Field(default=0, description="Output tokens used in this call")
class ToolCallTrace(BaseModel):
@@ -103,18 +91,9 @@ class TokenUsage(BaseModel):
)
input_tokens: int = Field(default=0, description="Number of input/prompt tokens consumed")
output_tokens: int = Field(
default=0, description="Number of visible output/completion tokens generated (excludes reasoning/thoughts)"
)
total_tokens: int = Field(default=0, description="Total tokens (input + output, excludes thoughts)")
output_tokens: int = Field(default=0, description="Number of output/completion tokens generated")
total_tokens: int = Field(default=0, description="Total tokens (input + output)")
cached_tokens: int = Field(default=0, description="Cached/cache-read prompt tokens, when reported by the provider")
thoughts_tokens: int = Field(
default=0,
description=(
"Reasoning/thinking tokens generated by the model. Billed at the output rate by some providers "
"(e.g. Gemini 2.5+ family) but not surfaced in the visible response."
),
)
def __add__(self, other: "TokenUsage") -> "TokenUsage":
"""Allow aggregating token usage from multiple calls."""
@@ -123,7 +102,6 @@ class TokenUsage(BaseModel):
output_tokens=self.output_tokens + other.output_tokens,
total_tokens=self.total_tokens + other.total_tokens,
cached_tokens=self.cached_tokens + other.cached_tokens,
thoughts_tokens=self.thoughts_tokens + other.thoughts_tokens,
)
@@ -329,8 +307,7 @@ class ReflectResult(BaseModel):
],
"experience": [],
"opinion": [],
"observation": [],
"mental-models": [],
"mental_models": [],
"directives": [
{
"id": "directive-123",
@@ -347,7 +324,7 @@ class ReflectResult(BaseModel):
text: str = Field(description="The formulated answer text")
based_on: dict[str, Any] = Field(
description="Facts used to formulate the answer, organized by type (world, experience, observation, mental-models, directives)"
description="Facts used to formulate the answer, organized by type (world, experience, mental_models, directives)"
)
structured_output: dict[str, Any] | None = Field(
default=None,
@@ -14,7 +14,6 @@ from typing import Any, Literal, cast
from pydantic import BaseModel, ConfigDict, Field, create_model, field_validator
from ..llm_interface import ProviderRateLimitResetError
from ..llm_wrapper import LLMConfig, OutputTooLongError, sanitize_llm_output
from ..operation_metadata import RetainExtractionErrors
from ..response_models import TokenUsage
@@ -452,11 +451,6 @@ def chunk_text(text: str, max_chars: int, structured_chunk_size: int | None = No
``structured_chunk_size``. When unset, that limit defaults to ``max_chars``.
For plain text, uses sentence-aware splitting.
The result is idempotent: re-chunking any chunk this returns yields that chunk
unchanged. The streaming retain pipeline pre-chunks each document once and then
re-chunks every piece during extraction; if a piece re-split, its sub-chunks
would inherit one chunk_index and collide on ``chunk_id`` (issue #2301).
Args:
text: Input text to chunk (plain text, JSON conversation, or JSONL)
max_chars: Target maximum characters per chunk
@@ -475,23 +469,11 @@ def chunk_text(text: str, max_chars: int, structured_chunk_size: int | None = No
# Try to parse as JSON conversation array
try:
parsed = json.loads(text)
if isinstance(parsed, list) and all(isinstance(turn, dict) for turn in parsed):
# This looks like a conversation - chunk at turn boundaries
return _chunk_conversation(parsed, max_chars, structured_limit)
except (json.JSONDecodeError, ValueError):
parsed = None
if isinstance(parsed, list) and all(isinstance(turn, dict) for turn in parsed):
# This looks like a conversation - chunk at turn boundaries
return _chunk_conversation(parsed, max_chars, structured_limit)
if isinstance(parsed, dict):
# A single JSON object — e.g. one JSONL line handed back to the extractor
# after the producer already pre-chunked it. It is one structured unit:
# keep it whole up to the structured limit, else split it as text within
# the chunk budget. Without this, a lone object (one line, so _chunk_jsonl
# declines) would fall through to plain-text splitting and re-split a chunk
# the producer deliberately kept whole — breaking idempotency (issue #2301).
if len(text) <= structured_limit:
return [text]
return _split_oversized_unit(text, max_chars)
pass
# Try to parse as JSONL (newline-delimited JSON objects, e.g. session logs)
jsonl_chunks = _chunk_jsonl(text, max_chars, structured_limit)
@@ -533,12 +515,10 @@ def _chunk_conversation(turns: list[dict], max_chars: int, structured_limit: int
turn_size = turn_unit_size + 1 # +1 for comma
# A turn too large to keep whole even alone: flush, then split it as
# text. Fragment within min(structured_limit, max_chars) so no fragment
# exceeds the chunk budget — otherwise a downstream re-chunk would split
# it again and collide on chunk_id (issue #2301).
# text so no chunk runs far over budget (the extractor won't re-chunk).
if turn_unit_size > structured_limit:
_flush()
chunks.extend(_split_oversized_unit(turn_json, min(structured_limit, max_chars)))
chunks.extend(_split_oversized_unit(turn_json, structured_limit))
continue
# If adding this turn would exceed limit and we have turns, save current chunk
@@ -601,12 +581,10 @@ def _chunk_jsonl(text: str, max_chars: int, structured_limit: int) -> list[str]
line_size = len(line) + 1 # +1 for the joining newline
# A line too large to keep whole even alone: flush, then split it as
# text. Fragment within min(structured_limit, max_chars) so no fragment
# exceeds the chunk budget — otherwise a downstream re-chunk would split
# it again and collide on chunk_id (issue #2301).
# text so no chunk runs far over budget (the extractor won't re-chunk).
if line_unit_size > structured_limit:
_flush()
chunks.extend(_split_oversized_unit(line, min(structured_limit, max_chars)))
chunks.extend(_split_oversized_unit(line, structured_limit))
continue
# If adding this line would exceed the limit and we have lines, flush.
@@ -1814,21 +1792,10 @@ async def extract_facts_from_text(
total_usage = total_usage + chunk_usage
if failed_chunks:
failed_summary = ", ".join(f"chunk {idx}: {type(err).__name__}" for idx, err in failed_chunks[:5])
quota_errors = [err for _, err in failed_chunks if isinstance(err, ProviderRateLimitResetError)]
if quota_errors and len(quota_errors) == len(failed_chunks):
retry_at = max(err.retry_at for err in quota_errors)
raise ProviderRateLimitResetError(
retry_at=retry_at,
message=(
f"Fact extraction deferred by provider quota: {len(failed_chunks)}/{len(chunks)} chunks failed. "
f"First failures: {failed_summary}. Provider detail: {quota_errors[0]}"
),
) from quota_errors[0]
# Fail the entire retain — partial extraction is not acceptable.
# All successfully extracted facts are discarded because the transaction
# hasn't committed yet. The worker poller will retry the entire task.
failed_summary = ", ".join(f"chunk {idx}: {type(err).__name__}" for idx, err in failed_chunks[:5])
raise RuntimeError(
f"Fact extraction failed: {len(failed_chunks)}/{len(chunks)} chunks failed. "
f"First failures: {failed_summary}"
@@ -1615,19 +1615,8 @@ async def _streaming_retain_batch(
# Check if facts are already committed (recovery from previous crash).
# If so, skip extraction+writes and jump straight to final ANN pass.
# ---------------------------------------------------------------------------
# Only the call that starts a document at chunk 0 may take the whole-document
# skip. When an oversized single item is split into several sequential
# sub-batches that SHARE one document_id AND one operation_id (see
# _split_contents_into_sub_batches), the first sub-batch commits its chunks
# and stamps effective_doc_id into result_metadata.facts_committed_document_ids.
# Without the offset gate, every later sub-batch (chunk_index_offset > 0) would
# then see its own document already "committed" and skip extraction, dropping
# all chunks past the first slice. A non-zero offset inherently means this call
# continues a document another sub-batch already started, so it must always do
# its work — crash-safety for those chunks still comes from the per-chunk hash
# recovery (existing_chunk_hashes) below.
facts_already_committed = False
if operation_id and chunk_index_offset == 0:
if operation_id:
try:
async with acquire_with_retry(pool) as conn:
row = await conn.fetchrow(
@@ -251,10 +251,8 @@ class LinkExpansionRetriever(GraphRetriever):
result.activation = row["score"]
results.append(result)
# filter_results_by_tags is a no-op when no filter applies (tags falsy and not
# the exact-empty/global scope), so call it unconditionally — gating on `if tags:`
# would skip the untagged-only filter for tags=[] + tags_match="exact".
results = filter_results_by_tags(results, tags, match=tags_match)
if tags:
results = filter_results_by_tags(results, tags, match=tags_match)
if tag_groups:
results = filter_results_by_tag_groups(results, tag_groups)
@@ -16,44 +16,6 @@ _RECENCY_ALPHA: float = 0.2
_TEMPORAL_ALPHA: float = 0.2
_PROOF_COUNT_ALPHA: float = 0.1 # Conservative: max ±5% for evidence strength
# Recency decay: maps a memory's age (days) onto a freshness signal in [0, 1]
# where 0.5 is neutral (no boost). The signal is then folded into the
# multiplicative recency_boost via `1 + recency_alpha * (recency - 0.5)`.
#
# "linear" — straight line from 1.0 (today) to a floor of 0.1, reaching
# the floor at `linear_window_days`. The historical default.
# "exponential" — 0.5 ** (days_ago / halflife_days). The half-life is the age
# at which the signal is exactly neutral (0.5): younger
# memories are boosted, older ones penalised, with a smooth
# asymptote toward 0 (no hard cutoff).
# "none" — always neutral (0.5), disabling the recency boost entirely.
# The validated set of names lives in config.RECENCY_DECAY_FUNCTIONS.
_RECENCY_DECAY_FUNCTION: str = "linear"
_RECENCY_DECAY_LINEAR_WINDOW_DAYS: float = 365.0
_RECENCY_DECAY_HALFLIFE_DAYS: float = 90.0
def compute_recency_decay(
days_ago: float,
function: str = _RECENCY_DECAY_FUNCTION,
linear_window_days: float = _RECENCY_DECAY_LINEAR_WINDOW_DAYS,
halflife_days: float = _RECENCY_DECAY_HALFLIFE_DAYS,
) -> float:
"""Map a memory's age in days to a freshness signal in [0, 1] (neutral 0.5).
Future-dated memories (negative ``days_ago``) clamp to the maximum freshness
so they are never penalised. See ``RECENCY_DECAY_FUNCTIONS`` for the shapes.
"""
if function == "none":
return 0.5
if function == "exponential":
if halflife_days <= 0:
return 0.5
return min(1.0, 0.5 ** (days_ago / halflife_days))
# "linear" (default): straight decay to a 0.1 floor over the window.
window = linear_window_days if linear_window_days > 0 else _RECENCY_DECAY_LINEAR_WINDOW_DAYS
return max(0.1, min(1.0, 1.0 - (days_ago / window)))
def apply_combined_scoring(
scored_results: list[ScoredResult],
@@ -62,9 +24,6 @@ def apply_combined_scoring(
temporal_alpha: float = _TEMPORAL_ALPHA,
proof_count_alpha: float = _PROOF_COUNT_ALPHA,
is_passthrough_reranker: bool = False,
recency_decay_function: str = _RECENCY_DECAY_FUNCTION,
recency_decay_linear_window_days: float = _RECENCY_DECAY_LINEAR_WINDOW_DAYS,
recency_decay_halflife_days: float = _RECENCY_DECAY_HALFLIFE_DAYS,
) -> None:
"""Apply combined scoring to a list of ScoredResults in-place.
@@ -98,12 +57,6 @@ def apply_combined_scoring(
recency_alpha: Max relative recency adjustment (default 0.2 ±10%).
temporal_alpha: Max relative temporal adjustment (default 0.2 ±10%).
proof_count_alpha: Max relative proof count adjustment (default 0.1 ±5%).
recency_decay_function: Agefreshness curve "linear" (default),
"exponential", or "none". See compute_recency_decay.
recency_decay_linear_window_days: Days over which the linear curve
decays to its floor (default 365).
recency_decay_halflife_days: For the exponential curve, the age at which
the recency signal is neutral (0.5) (default 90).
"""
if now.tzinfo is None:
now = now.replace(tzinfo=UTC)
@@ -145,8 +98,7 @@ def apply_combined_scoring(
sr.cross_encoder_score_normalized = 1.0 - (0.9 * new_rank / denom)
for sr in scored_results:
# Recency: configurable decay (linear default; see compute_recency_decay)
# → [0.0, 1.0]; neutral 0.5 if no date.
# Recency: linear decay over 365 days → [0.1, 1.0]; neutral 0.5 if no date.
# Use the unit's effective time (occurred_start, then mentioned_at, then
# occurred_end) — the same COALESCE order as retrieval._coalesce_date — so a
# memory that carries only a mentioned_at / occurred_end (e.g. conversation
@@ -159,12 +111,7 @@ def apply_combined_scoring(
if occurred.tzinfo is None:
occurred = occurred.replace(tzinfo=UTC)
days_ago = (now - occurred).total_seconds() / 86400
sr.recency = compute_recency_decay(
days_ago,
recency_decay_function,
recency_decay_linear_window_days,
recency_decay_halflife_days,
)
sr.recency = max(0.1, min(1.0, 1.0 - (days_ago / 365)))
# Temporal proximity: meaningful only for temporal queries; neutral otherwise.
sr.temporal = sr.retrieval.temporal_proximity if sr.retrieval.temporal_proximity is not None else 0.5
@@ -14,12 +14,6 @@ AND matching (all/all_strict): Memory matches if ALL request tags are present in
EXACT matching: Memory matches only if its tag set EQUALS the request tag set (order-
independent). Used for observation "scope" filtering, where each observation lives
under exactly one scope (its full tag set) and "scope [a]" must not match "[a, b]".
An EMPTY request scope (no tags ``[]`` or ``None``) is the global/untagged scope and
matches only untagged memories the scope that ``observation_scopes="shared"``
consolidation writes to. This is the one mode where absent tags filter rather than
meaning "no filter"; all other modes treat empty/absent tags as "no filtering". This
mirrors the ``GET .../graph`` endpoint, where ``tags_match="exact"`` with no tags also
selects the global scope.
"""
from __future__ import annotations
@@ -88,16 +82,11 @@ def build_tags_where_clause(
>>> clause, params, next_offset = build_tags_where_clause(['user_a'], 3, 'mu.', 'any_strict')
>>> print(clause) # "AND mu.tags IS NOT NULL AND mu.tags != '{}' AND mu.tags && $3"
"""
column = f"{table_alias}tags" if table_alias else "tags"
if match == "exact" and not tags:
# Empty/absent scope = global/untagged: match only untagged rows. No bind param
# needed (callers gate the param on truthy `tags`, so none is appended).
return f"AND ({column} IS NULL OR {column} = '{{}}')", [], param_offset
if not tags:
return "", [], param_offset
column = f"{table_alias}tags" if table_alias else "tags"
if match == "exact":
# Set equality (order-independent): superset AND subset. Untagged rows
# (empty array) never satisfy `@>` of a non-empty scope, so they're excluded.
@@ -137,16 +126,11 @@ def build_tags_where_clause_simple(
Returns:
SQL clause string or empty string.
"""
column = f"{table_alias}tags" if table_alias else "tags"
if match == "exact" and not tags:
# Empty/absent scope = global/untagged: match only untagged rows. No bind param
# needed (callers gate the param on truthy `tags`, so none is appended).
return f"AND ({column} IS NULL OR {column} = '{{}}')"
if not tags:
return ""
column = f"{table_alias}tags" if table_alias else "tags"
if match == "exact":
# Set equality (order-independent): superset AND subset. Untagged rows
# (empty array) never satisfy `@>` of a non-empty scope, so they're excluded.
@@ -180,10 +164,6 @@ def filter_results_by_tags(
Returns:
Filtered list of results.
"""
if match == "exact" and not tags:
# Empty/absent scope = global/untagged: keep only untagged results.
return [r for r in results if not getattr(r, "tags", None)]
if not tags:
return results
@@ -287,9 +267,6 @@ def _build_group_clause(
if isinstance(group, TagGroupLeaf):
column = f"{table_alias}tags" if table_alias else "tags"
if group.match == "exact":
if len(group.tags) == 0:
# Empty scope = global/untagged: match only untagged rows (no bind param).
return f"({column} IS NULL OR {column} = '{{}}')", [], param_offset
clause = f"({column} @> ${param_offset} AND {column} <@ ${param_offset})"
return clause, [group.tags], param_offset + 1
operator, include_untagged = _parse_tags_match(group.match)
@@ -392,9 +369,6 @@ def _match_group(result: object, group: TagGroup) -> bool:
if isinstance(group, TagGroupLeaf):
result_tags = getattr(result, "tags", None)
is_untagged = result_tags is None or len(result_tags) == 0
if group.match == "exact" and len(group.tags) == 0:
# Empty scope = global/untagged: match only untagged results.
return is_untagged
_, include_untagged = _parse_tags_match(group.match)
is_any_match = group.match in ("any", "any_strict")
tags_set = set(group.tags)
@@ -97,10 +97,6 @@ class PrecheckContext:
- ``bank_id``: parsed from the URL path.
- ``request_context``: the authenticated :class:`RequestContext` (tenant
already resolved by the tenant extension).
- ``content_length``: value of the ``Content-Length`` request header as an
int, or ``None`` when the header is absent or unparseable (e.g. chunked
transfer encoding). Lets a precheck make size-aware decisions such as
an upper-bound cost estimate without reading or deserialising the body.
Implementations should keep precheck cheap and side-effect-free. The
full per-request validators (``validate_retain`` / ``validate_recall``
@@ -111,7 +107,6 @@ class PrecheckContext:
operation: str
bank_id: str
request_context: "RequestContext"
content_length: int | None = None
@dataclass
@@ -208,16 +203,6 @@ class RetainResult:
llm_input_tokens: int | None = None
llm_output_tokens: int | None = None
llm_total_tokens: int | None = None
# Diagnostic token splits surfaced for cost attribution and prompt-cache
# tuning. ``llm_cached_input_tokens`` is the subset of llm_input_tokens
# served from the provider's prompt cache (e.g. Gemini's
# cached_content_token_count). ``llm_thoughts_tokens`` is reasoning tokens
# that are billed at the output rate by some providers (Gemini 2.5+) but
# are not part of the visible response. Both default to None when the
# engine/provider didn't report them; downstream metering extensions
# should treat None as 0.
llm_cached_input_tokens: int | None = None
llm_thoughts_tokens: int | None = None
# Content tokens the retain pipeline actually processed, after
# chunk-level content-hash deduplication. Semantics:
# None — no dedup signal available (e.g. a first-time retain or a
+65 -141
View File
@@ -12,7 +12,6 @@ from datetime import datetime, timezone
from typing import Any, Callable
from fastmcp import FastMCP
from mcp.types import ToolAnnotations
from pydantic import TypeAdapter
from hindsight_api import MemoryEngine
@@ -200,47 +199,6 @@ def build_content_dict(
return content_dict, None
# MCP tool annotations. Hindsight is a closed memory store (no open-world / internet
# access), so openWorldHint=False throughout. readOnlyHint lets clients group and
# auto-approve safe reads; destructiveHint flags tools that delete or clear memory.
_READ_ONLY_TOOLS = {
"recall",
"reflect",
"list_banks",
"get_bank",
"get_bank_stats",
"list_mental_models",
"get_mental_model",
"list_directives",
"list_memories",
"get_memory",
"list_documents",
"get_document",
"list_operations",
"get_operation",
"list_tags",
}
_DESTRUCTIVE_TOOLS = {
"delete_bank",
"clear_memories",
"clear_mental_model",
"delete_mental_model",
"delete_directive",
"delete_document",
"invalidate_memory",
}
def _tool_annotations(name: str) -> ToolAnnotations:
if name in _READ_ONLY_TOOLS:
return ToolAnnotations(readOnlyHint=True, openWorldHint=False)
if name in _DESTRUCTIVE_TOOLS:
return ToolAnnotations(readOnlyHint=False, destructiveHint=True, openWorldHint=False)
# Everything else writes but does not destructively delete/clear memory
# (retain, create_*, update_*, refresh_mental_model, cancel_operation).
return ToolAnnotations(readOnlyHint=False, destructiveHint=False, openWorldHint=False)
def register_mcp_tools(
mcp: FastMCP,
memory: MemoryEngine,
@@ -594,7 +552,7 @@ def _register_retain(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig)
if config.include_bank_id_param:
@mcp.tool(description=description, annotations=_tool_annotations("retain"))
@mcp.tool(description=description)
async def retain(
content: str,
context: str = "general",
@@ -650,7 +608,7 @@ def _register_retain(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig)
else:
@mcp.tool(description=description, annotations=_tool_annotations("retain"))
@mcp.tool(description=description)
async def retain(
content: str,
context: str = "general",
@@ -708,7 +666,7 @@ def _register_sync_retain(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsCo
if config.include_bank_id_param:
@mcp.tool(annotations=_tool_annotations("sync_retain"))
@mcp.tool()
async def sync_retain(
content: str,
context: str = "general",
@@ -766,7 +724,7 @@ def _register_sync_retain(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsCo
else:
@mcp.tool(annotations=_tool_annotations("sync_retain"))
@mcp.tool()
async def sync_retain(
content: str,
context: str = "general",
@@ -827,13 +785,12 @@ def _register_recall(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig)
if config.include_bank_id_param:
@mcp.tool(description=description, annotations=_tool_annotations("recall"))
@mcp.tool(description=description)
async def recall(
query: str,
max_tokens: int = 4096,
budget: str = "high",
types: list[str] | None = None,
prefer_observations: bool = False,
tags: list[str] | None = None,
tags_match: str = "any",
tag_groups: list[dict] | None = None,
@@ -846,10 +803,6 @@ def _register_recall(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig)
max_tokens: Maximum tokens to return in results (default: 4096)
budget: Search budget - 'low', 'mid', or 'high' (default: 'high'). Higher budgets search more thoroughly.
types: Fact types to include (e.g., ['world', 'experience']). Default: all types.
prefer_observations: When recalling raw facts together with 'observation', drop any raw fact
that a returned observation was consolidated from, so the observation supersedes it (no
duplicate content). Disabled by default; set true to enable. No effect unless
'observation' and a raw type are both in types. Default: False.
tags: Optional tags to filter results by (e.g., ['project:alpha']). Mutually exclusive with tag_groups.
tags_match: How to match tags - 'any' (match any tag) or 'all' (match all tags). Default: 'any'
tag_groups: Compound tag filter using boolean groups (AND-ed together). Each group is a leaf
@@ -878,7 +831,6 @@ def _register_recall(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig)
"bank_id": target_bank,
"query": query,
"fact_type": fact_types,
"prefer_observations": prefer_observations,
"budget": budget_enum,
"max_tokens": max_tokens,
"request_context": _get_request_context(config),
@@ -905,13 +857,12 @@ def _register_recall(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig)
else:
@mcp.tool(description=description, annotations=_tool_annotations("recall"))
@mcp.tool(description=description)
async def recall(
query: str,
max_tokens: int = 4096,
budget: str = "high",
types: list[str] | None = None,
prefer_observations: bool = False,
tags: list[str] | None = None,
tags_match: str = "any",
tag_groups: list[dict] | None = None,
@@ -923,10 +874,6 @@ def _register_recall(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig)
max_tokens: Maximum tokens to return in results (default: 4096)
budget: Search budget - 'low', 'mid', or 'high' (default: 'high'). Higher budgets search more thoroughly.
types: Fact types to include (e.g., ['world', 'experience']). Default: all types.
prefer_observations: When recalling raw facts together with 'observation', drop any raw fact
that a returned observation was consolidated from, so the observation supersedes it (no
duplicate content). Disabled by default; set true to enable. No effect unless
'observation' and a raw type are both in types. Default: False.
tags: Optional tags to filter results by (e.g., ['project:alpha']). Mutually exclusive with tag_groups.
tags_match: How to match tags - 'any' (match any tag) or 'all' (match all tags). Default: 'any'
tag_groups: Compound tag filter using boolean groups (AND-ed together). Each group is a leaf
@@ -954,7 +901,6 @@ def _register_recall(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig)
"bank_id": target_bank,
"query": query,
"fact_type": fact_types,
"prefer_observations": prefer_observations,
"budget": budget_enum,
"max_tokens": max_tokens,
"request_context": _get_request_context(config),
@@ -985,7 +931,7 @@ def _register_reflect(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig
if config.include_bank_id_param:
@mcp.tool(annotations=_tool_annotations("reflect"))
@mcp.tool()
async def reflect(
query: str,
context: str | None = None,
@@ -995,7 +941,6 @@ def _register_reflect(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig
tags: list[str] | None = None,
tags_match: str = "any",
include_based_on: bool = False,
include_trace: bool = False,
bank_id: str | None = None,
) -> str:
"""
@@ -1026,7 +971,6 @@ def _register_reflect(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig
tags: Optional tags to filter memories by (e.g., ['project:alpha'])
tags_match: How to match tags - 'any' (match any tag) or 'all' (match all tags). Default: 'any'
include_based_on: Include source facts used for synthesis. Defaults to false because broad reflections can exceed MCP client result limits.
include_trace: Include the reflection's internal trace fields (tool_trace/llm_trace and directives_applied). Defaults to false because the trace can be tens of KB and overflow MCP client context; enable only for debugging.
bank_id: Optional bank to reflect in (defaults to session bank). Use for cross-bank operations.
"""
try:
@@ -1056,15 +1000,6 @@ def _register_reflect(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig
result_data = json.loads(reflect_result.model_dump_json(indent=2))
if not include_based_on:
result_data.pop("based_on", None)
if not include_trace:
# The agentic reflect loop's trace fields can be tens of KB (full
# mental-model text) and silently overflow MCP client context; the
# REST API omits them by default too. directives_applied is built by
# the engine "for the trace" and carries full directive content, so it
# belongs with tool_trace/llm_trace here. Opt in via include_trace.
result_data.pop("tool_trace", None)
result_data.pop("llm_trace", None)
result_data.pop("directives_applied", None)
if response_schema is not None and hasattr(reflect_result, "structured_output"):
result_data["structured_output"] = reflect_result.structured_output
return json.dumps(result_data, indent=2)
@@ -1077,7 +1012,7 @@ def _register_reflect(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig
else:
@mcp.tool(annotations=_tool_annotations("reflect"))
@mcp.tool()
async def reflect(
query: str,
context: str | None = None,
@@ -1087,7 +1022,6 @@ def _register_reflect(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig
tags: list[str] | None = None,
tags_match: str = "any",
include_based_on: bool = False,
include_trace: bool = False,
) -> dict:
"""
Generate thoughtful analysis by synthesizing stored memories with the bank's personality.
@@ -1117,7 +1051,6 @@ def _register_reflect(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig
tags: Optional tags to filter memories by (e.g., ['project:alpha'])
tags_match: How to match tags - 'any' (match any tag) or 'all' (match all tags). Default: 'any'
include_based_on: Include source facts used for synthesis. Defaults to false because broad reflections can exceed MCP client result limits.
include_trace: Include the reflection's internal trace fields (tool_trace/llm_trace and directives_applied). Defaults to false because the trace can be tens of KB and overflow MCP client context; enable only for debugging.
"""
try:
target_bank = config.bank_id_resolver()
@@ -1146,15 +1079,6 @@ def _register_reflect(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig
result_data = reflect_result.model_dump()
if not include_based_on:
result_data.pop("based_on", None)
if not include_trace:
# The agentic reflect loop's trace fields can be tens of KB (full
# mental-model text) and silently overflow MCP client context; the
# REST API omits them by default too. directives_applied is built by
# the engine "for the trace" and carries full directive content, so it
# belongs with tool_trace/llm_trace here. Opt in via include_trace.
result_data.pop("tool_trace", None)
result_data.pop("llm_trace", None)
result_data.pop("directives_applied", None)
if response_schema is not None and hasattr(reflect_result, "structured_output"):
result_data["structured_output"] = reflect_result.structured_output
return result_data
@@ -1169,7 +1093,7 @@ def _register_reflect(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig
def _register_list_banks(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig) -> None:
"""Register the list_banks tool."""
@mcp.tool(annotations=_tool_annotations("list_banks"))
@mcp.tool()
async def list_banks() -> str:
"""
List all available memory banks.
@@ -1194,7 +1118,7 @@ def _register_list_banks(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsCon
def _register_create_bank(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig) -> None:
"""Register the create_bank tool."""
@mcp.tool(annotations=_tool_annotations("create_bank"))
@mcp.tool()
async def create_bank(bank_id: str, name: str | None = None, mission: str | None = None) -> str:
"""
Create a new memory bank or get an existing one.
@@ -1258,7 +1182,7 @@ def _register_list_mental_models(mcp: FastMCP, memory: MemoryEngine, config: MCP
if config.include_bank_id_param:
@mcp.tool(annotations=_tool_annotations("list_mental_models"))
@mcp.tool()
async def list_mental_models(
tags: list[str] | None = None,
detail: str = "full",
@@ -1297,7 +1221,7 @@ def _register_list_mental_models(mcp: FastMCP, memory: MemoryEngine, config: MCP
else:
@mcp.tool(annotations=_tool_annotations("list_mental_models"))
@mcp.tool()
async def list_mental_models(
tags: list[str] | None = None,
detail: str = "full",
@@ -1338,7 +1262,7 @@ def _register_get_mental_model(mcp: FastMCP, memory: MemoryEngine, config: MCPTo
if config.include_bank_id_param:
@mcp.tool(annotations=_tool_annotations("get_mental_model"))
@mcp.tool()
async def get_mental_model(
mental_model_id: str,
detail: str = "full",
@@ -1378,7 +1302,7 @@ def _register_get_mental_model(mcp: FastMCP, memory: MemoryEngine, config: MCPTo
else:
@mcp.tool(annotations=_tool_annotations("get_mental_model"))
@mcp.tool()
async def get_mental_model(
mental_model_id: str,
detail: str = "full",
@@ -1420,7 +1344,7 @@ def _register_create_mental_model(mcp: FastMCP, memory: MemoryEngine, config: MC
if config.include_bank_id_param:
@mcp.tool(annotations=_tool_annotations("create_mental_model"))
@mcp.tool()
async def create_mental_model(
name: str,
source_query: str,
@@ -1504,7 +1428,7 @@ def _register_create_mental_model(mcp: FastMCP, memory: MemoryEngine, config: MC
else:
@mcp.tool(annotations=_tool_annotations("create_mental_model"))
@mcp.tool()
async def create_mental_model(
name: str,
source_query: str,
@@ -1586,7 +1510,7 @@ def _register_update_mental_model(mcp: FastMCP, memory: MemoryEngine, config: MC
if config.include_bank_id_param:
@mcp.tool(annotations=_tool_annotations("update_mental_model"))
@mcp.tool()
async def update_mental_model(
mental_model_id: str,
name: str | None = None,
@@ -1647,7 +1571,7 @@ def _register_update_mental_model(mcp: FastMCP, memory: MemoryEngine, config: MC
else:
@mcp.tool(annotations=_tool_annotations("update_mental_model"))
@mcp.tool()
async def update_mental_model(
mental_model_id: str,
name: str | None = None,
@@ -1710,7 +1634,7 @@ def _register_delete_mental_model(mcp: FastMCP, memory: MemoryEngine, config: MC
if config.include_bank_id_param:
@mcp.tool(annotations=_tool_annotations("delete_mental_model"))
@mcp.tool()
async def delete_mental_model(
mental_model_id: str,
bank_id: str | None = None,
@@ -1746,7 +1670,7 @@ def _register_delete_mental_model(mcp: FastMCP, memory: MemoryEngine, config: MC
else:
@mcp.tool(annotations=_tool_annotations("delete_mental_model"))
@mcp.tool()
async def delete_mental_model(
mental_model_id: str,
) -> dict:
@@ -1784,7 +1708,7 @@ def _register_refresh_mental_model(mcp: FastMCP, memory: MemoryEngine, config: M
if config.include_bank_id_param:
@mcp.tool(annotations=_tool_annotations("refresh_mental_model"))
@mcp.tool()
async def refresh_mental_model(
mental_model_id: str,
bank_id: str | None = None,
@@ -1828,7 +1752,7 @@ def _register_refresh_mental_model(mcp: FastMCP, memory: MemoryEngine, config: M
else:
@mcp.tool(annotations=_tool_annotations("refresh_mental_model"))
@mcp.tool()
async def refresh_mental_model(
mental_model_id: str,
) -> dict:
@@ -1872,7 +1796,7 @@ def _register_clear_mental_model(mcp: FastMCP, memory: MemoryEngine, config: MCP
if config.include_bank_id_param:
@mcp.tool(annotations=_tool_annotations("clear_mental_model"))
@mcp.tool()
async def clear_mental_model(
mental_model_id: str,
bank_id: str | None = None,
@@ -1918,7 +1842,7 @@ def _register_clear_mental_model(mcp: FastMCP, memory: MemoryEngine, config: MCP
else:
@mcp.tool(annotations=_tool_annotations("clear_mental_model"))
@mcp.tool()
async def clear_mental_model(
mental_model_id: str,
) -> dict:
@@ -1969,7 +1893,7 @@ def _register_list_directives(mcp: FastMCP, memory: MemoryEngine, config: MCPToo
if config.include_bank_id_param:
@mcp.tool(annotations=_tool_annotations("list_directives"))
@mcp.tool()
async def list_directives(
tags: list[str] | None = None,
active_only: bool = True,
@@ -2007,7 +1931,7 @@ def _register_list_directives(mcp: FastMCP, memory: MemoryEngine, config: MCPToo
else:
@mcp.tool(annotations=_tool_annotations("list_directives"))
@mcp.tool()
async def list_directives(
tags: list[str] | None = None,
active_only: bool = True,
@@ -2047,7 +1971,7 @@ def _register_create_directive(mcp: FastMCP, memory: MemoryEngine, config: MCPTo
if config.include_bank_id_param:
@mcp.tool(annotations=_tool_annotations("create_directive"))
@mcp.tool()
async def create_directive(
name: str,
content: str,
@@ -2093,7 +2017,7 @@ def _register_create_directive(mcp: FastMCP, memory: MemoryEngine, config: MCPTo
else:
@mcp.tool(annotations=_tool_annotations("create_directive"))
@mcp.tool()
async def create_directive(
name: str,
content: str,
@@ -2141,7 +2065,7 @@ def _register_delete_directive(mcp: FastMCP, memory: MemoryEngine, config: MCPTo
if config.include_bank_id_param:
@mcp.tool(annotations=_tool_annotations("delete_directive"))
@mcp.tool()
async def delete_directive(
directive_id: str,
bank_id: str | None = None,
@@ -2177,7 +2101,7 @@ def _register_delete_directive(mcp: FastMCP, memory: MemoryEngine, config: MCPTo
else:
@mcp.tool(annotations=_tool_annotations("delete_directive"))
@mcp.tool()
async def delete_directive(
directive_id: str,
) -> dict:
@@ -2220,7 +2144,7 @@ def _register_list_memories(mcp: FastMCP, memory: MemoryEngine, config: MCPTools
if config.include_bank_id_param:
@mcp.tool(annotations=_tool_annotations("list_memories"))
@mcp.tool()
async def list_memories(
type: str | None = None,
q: str | None = None,
@@ -2235,7 +2159,7 @@ def _register_list_memories(mcp: FastMCP, memory: MemoryEngine, config: MCPTools
browse/search without relevance ranking.
Args:
type: Filter by fact type: 'world', 'experience', or 'observation'
type: Filter by fact type: 'world', 'experience', or 'opinion'
q: Optional text search query to filter memories
limit: Maximum number of results (default: 100)
offset: Pagination offset (default: 0)
@@ -2264,7 +2188,7 @@ def _register_list_memories(mcp: FastMCP, memory: MemoryEngine, config: MCPTools
else:
@mcp.tool(annotations=_tool_annotations("list_memories"))
@mcp.tool()
async def list_memories(
type: str | None = None,
q: str | None = None,
@@ -2278,7 +2202,7 @@ def _register_list_memories(mcp: FastMCP, memory: MemoryEngine, config: MCPTools
browse/search without relevance ranking.
Args:
type: Filter by fact type: 'world', 'experience', or 'observation'
type: Filter by fact type: 'world', 'experience', or 'opinion'
q: Optional text search query to filter memories
limit: Maximum number of results (default: 100)
offset: Pagination offset (default: 0)
@@ -2310,7 +2234,7 @@ def _register_get_memory(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsCon
if config.include_bank_id_param:
@mcp.tool(annotations=_tool_annotations("get_memory"))
@mcp.tool()
async def get_memory(
memory_id: str,
bank_id: str | None = None,
@@ -2346,7 +2270,7 @@ def _register_get_memory(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsCon
else:
@mcp.tool(annotations=_tool_annotations("get_memory"))
@mcp.tool()
async def get_memory(
memory_id: str,
) -> dict:
@@ -2397,7 +2321,7 @@ def _register_update_memory(mcp: FastMCP, memory: MemoryEngine, config: MCPTools
if config.include_bank_id_param:
@mcp.tool(description=_EDIT_DOC, annotations=_tool_annotations("update_memory"))
@mcp.tool(description=_EDIT_DOC)
async def update_memory(
memory_id: str,
text: str | None = None,
@@ -2443,7 +2367,7 @@ def _register_update_memory(mcp: FastMCP, memory: MemoryEngine, config: MCPTools
else:
@mcp.tool(description=_EDIT_DOC, annotations=_tool_annotations("update_memory"))
@mcp.tool(description=_EDIT_DOC)
async def update_memory(
memory_id: str,
text: str | None = None,
@@ -2502,7 +2426,7 @@ def _register_invalidate_memory(mcp: FastMCP, memory: MemoryEngine, config: MCPT
if config.include_bank_id_param:
@mcp.tool(description=_INVALIDATE_DOC, annotations=_tool_annotations("invalidate_memory"))
@mcp.tool(description=_INVALIDATE_DOC)
async def invalidate_memory(
memory_id: str,
reason: str | None = None,
@@ -2542,7 +2466,7 @@ def _register_invalidate_memory(mcp: FastMCP, memory: MemoryEngine, config: MCPT
else:
@mcp.tool(description=_INVALIDATE_DOC, annotations=_tool_annotations("invalidate_memory"))
@mcp.tool(description=_INVALIDATE_DOC)
async def invalidate_memory(
memory_id: str,
reason: str | None = None,
@@ -2589,7 +2513,7 @@ def _register_list_documents(mcp: FastMCP, memory: MemoryEngine, config: MCPTool
if config.include_bank_id_param:
@mcp.tool(annotations=_tool_annotations("list_documents"))
@mcp.tool()
async def list_documents(
q: str | None = None,
limit: int = 100,
@@ -2627,7 +2551,7 @@ def _register_list_documents(mcp: FastMCP, memory: MemoryEngine, config: MCPTool
else:
@mcp.tool(annotations=_tool_annotations("list_documents"))
@mcp.tool()
async def list_documents(
q: str | None = None,
limit: int = 100,
@@ -2667,7 +2591,7 @@ def _register_get_document(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsC
if config.include_bank_id_param:
@mcp.tool(annotations=_tool_annotations("get_document"))
@mcp.tool()
async def get_document(
document_id: str,
bank_id: str | None = None,
@@ -2703,7 +2627,7 @@ def _register_get_document(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsC
else:
@mcp.tool(annotations=_tool_annotations("get_document"))
@mcp.tool()
async def get_document(
document_id: str,
) -> dict:
@@ -2741,7 +2665,7 @@ def _register_delete_document(mcp: FastMCP, memory: MemoryEngine, config: MCPToo
if config.include_bank_id_param:
@mcp.tool(annotations=_tool_annotations("delete_document"))
@mcp.tool()
async def delete_document(
document_id: str,
bank_id: str | None = None,
@@ -2775,7 +2699,7 @@ def _register_delete_document(mcp: FastMCP, memory: MemoryEngine, config: MCPToo
else:
@mcp.tool(annotations=_tool_annotations("delete_document"))
@mcp.tool()
async def delete_document(
document_id: str,
) -> dict:
@@ -2816,7 +2740,7 @@ def _register_list_operations(mcp: FastMCP, memory: MemoryEngine, config: MCPToo
if config.include_bank_id_param:
@mcp.tool(annotations=_tool_annotations("list_operations"))
@mcp.tool()
async def list_operations(
status: str | None = None,
limit: int = 20,
@@ -2853,7 +2777,7 @@ def _register_list_operations(mcp: FastMCP, memory: MemoryEngine, config: MCPToo
else:
@mcp.tool(annotations=_tool_annotations("list_operations"))
@mcp.tool()
async def list_operations(
status: str | None = None,
limit: int = 20,
@@ -2892,7 +2816,7 @@ def _register_get_operation(mcp: FastMCP, memory: MemoryEngine, config: MCPTools
if config.include_bank_id_param:
@mcp.tool(annotations=_tool_annotations("get_operation"))
@mcp.tool()
async def get_operation(
operation_id: str,
bank_id: str | None = None,
@@ -2926,7 +2850,7 @@ def _register_get_operation(mcp: FastMCP, memory: MemoryEngine, config: MCPTools
else:
@mcp.tool(annotations=_tool_annotations("get_operation"))
@mcp.tool()
async def get_operation(
operation_id: str,
) -> dict:
@@ -2962,7 +2886,7 @@ def _register_cancel_operation(mcp: FastMCP, memory: MemoryEngine, config: MCPTo
if config.include_bank_id_param:
@mcp.tool(annotations=_tool_annotations("cancel_operation"))
@mcp.tool()
async def cancel_operation(
operation_id: str,
bank_id: str | None = None,
@@ -2994,7 +2918,7 @@ def _register_cancel_operation(mcp: FastMCP, memory: MemoryEngine, config: MCPTo
else:
@mcp.tool(annotations=_tool_annotations("cancel_operation"))
@mcp.tool()
async def cancel_operation(
operation_id: str,
) -> dict:
@@ -3033,7 +2957,7 @@ def _register_list_tags(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConf
if config.include_bank_id_param:
@mcp.tool(annotations=_tool_annotations("list_tags"))
@mcp.tool()
async def list_tags(
q: str | None = None,
limit: int = 100,
@@ -3070,7 +2994,7 @@ def _register_list_tags(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConf
else:
@mcp.tool(annotations=_tool_annotations("list_tags"))
@mcp.tool()
async def list_tags(
q: str | None = None,
limit: int = 100,
@@ -3109,7 +3033,7 @@ def _register_get_bank(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfi
if config.include_bank_id_param:
@mcp.tool(annotations=_tool_annotations("get_bank"))
@mcp.tool()
async def get_bank(
bank_id: str | None = None,
) -> str:
@@ -3142,7 +3066,7 @@ def _register_get_bank(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfi
else:
@mcp.tool(annotations=_tool_annotations("get_bank"))
@mcp.tool()
async def get_bank() -> dict:
"""
Get the profile of this memory bank.
@@ -3172,7 +3096,7 @@ def _register_get_bank(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfi
def _register_get_bank_stats(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig) -> None:
"""Register the get_bank_stats tool (multi-bank only)."""
@mcp.tool(annotations=_tool_annotations("get_bank_stats"))
@mcp.tool()
async def get_bank_stats(
bank_id: str | None = None,
) -> str:
@@ -3245,7 +3169,7 @@ def _register_update_bank(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsCo
if config.include_bank_id_param:
@mcp.tool(annotations=_tool_annotations("update_bank"))
@mcp.tool()
async def update_bank(
name: str | None = None,
mission: str | None = None,
@@ -3306,7 +3230,7 @@ def _register_update_bank(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsCo
else:
@mcp.tool(annotations=_tool_annotations("update_bank"))
@mcp.tool()
async def update_bank(
name: str | None = None,
mission: str | None = None,
@@ -3369,7 +3293,7 @@ def _register_delete_bank(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsCo
if config.include_bank_id_param:
@mcp.tool(annotations=_tool_annotations("delete_bank"))
@mcp.tool()
async def delete_bank(
bank_id: str | None = None,
) -> str:
@@ -3401,7 +3325,7 @@ def _register_delete_bank(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsCo
else:
@mcp.tool(annotations=_tool_annotations("delete_bank"))
@mcp.tool()
async def delete_bank() -> dict:
"""
Delete this memory bank and all its data.
@@ -3432,7 +3356,7 @@ def _register_clear_memories(mcp: FastMCP, memory: MemoryEngine, config: MCPTool
if config.include_bank_id_param:
@mcp.tool(annotations=_tool_annotations("clear_memories"))
@mcp.tool()
async def clear_memories(
type: str | None = None,
bank_id: str | None = None,
@@ -3443,7 +3367,7 @@ def _register_clear_memories(mcp: FastMCP, memory: MemoryEngine, config: MCPTool
Optionally filter by fact type to only clear specific kinds of memories.
Args:
type: Optional fact type filter: 'world', 'experience', or 'observation'. If not specified, clears all.
type: Optional fact type filter: 'world', 'experience', or 'opinion'. If not specified, clears all.
bank_id: Optional bank (defaults to session bank). Use for cross-bank operations.
"""
try:
@@ -3467,7 +3391,7 @@ def _register_clear_memories(mcp: FastMCP, memory: MemoryEngine, config: MCPTool
else:
@mcp.tool(annotations=_tool_annotations("clear_memories"))
@mcp.tool()
async def clear_memories(
type: str | None = None,
) -> dict:
@@ -3477,7 +3401,7 @@ def _register_clear_memories(mcp: FastMCP, memory: MemoryEngine, config: MCPTool
Optionally filter by fact type to only clear specific kinds of memories.
Args:
type: Optional fact type filter: 'world', 'experience', or 'observation'. If not specified, clears all.
type: Optional fact type filter: 'world', 'experience', or 'opinion'. If not specified, clears all.
"""
try:
target_bank = config.bank_id_resolver()
+19 -312
View File
@@ -11,17 +11,15 @@ This module provides metrics for:
- Database connection pool metrics
"""
import asyncio
import importlib
import logging
import os
import re
_resource_mod = importlib.import_module("resource") if importlib.util.find_spec("resource") else None
import threading
import time
from contextlib import contextmanager
from typing import TYPE_CHECKING, Callable, NamedTuple
from typing import TYPE_CHECKING, Callable
from opentelemetry import metrics
from opentelemetry.exporter.prometheus import PrometheusMetricReader
@@ -77,28 +75,6 @@ LLM_DURATION_BUCKETS = (0.1, 0.25, 0.5, 1.0, 2.0, 3.0, 5.0, 10.0, 15.0, 30.0, 60
# HTTP request duration buckets (millisecond-level for fast endpoints)
HTTP_DURATION_BUCKETS = (0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0, 30.0)
# How often the backlog / queue-depth gauge caches are refreshed (seconds).
# The counts are aggregate COUNT queries, so a background task refreshes a
# cache and the observable gauges read from it — keeping the /metrics scrape
# path synchronous (the same reason the db-pool gauges read cached state).
BACKLOG_METRICS_REFRESH_SECONDS = 30
class _AsyncOpKey(NamedTuple):
"""Cache / label key for the async-operation queue gauge."""
tenant: str
operation_type: str
status: str
bank_id: str | None
class _BacklogKey(NamedTuple):
"""Cache / label key for the consolidation backlog and failed gauges."""
tenant: str
bank_id: str | None
def get_token_bucket(token_count: int) -> str:
"""
@@ -137,27 +113,6 @@ def get_token_bucket(token_count: int) -> str:
return "50k+"
# Template unbounded id segments before a path is used as the low-cardinality
# "endpoint" metric label. A raw per-bank path segment (e.g. user-123) would
# otherwise create one never-evicted OTel series per bank.
_METRIC_BANK_SEGMENT_RE = re.compile(r"(/banks/)[^/]+")
_METRIC_UUID_RE = re.compile(r"/[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}")
_METRIC_NUMERIC_ID_RE = re.compile(r"/\d+(?=/|$)")
def normalize_http_endpoint(path: str) -> str:
"""Template high-cardinality id segments in an HTTP path for safe metric labeling.
Collapses the "/banks/<id>" segment (any bank id, including non-numeric ones like
"user-123"), UUIDs, and numeric ids to placeholders so the "endpoint" metric label
has bounded cardinality. Analogous to get_token_bucket for token counts.
"""
path = _METRIC_BANK_SEGMENT_RE.sub(r"\g<1>{bank_id}", path)
path = _METRIC_UUID_RE.sub("/{id}", path)
path = _METRIC_NUMERIC_ID_RE.sub("/{id}", path)
return path
logger = logging.getLogger(__name__)
# Global meter instance
@@ -246,19 +201,6 @@ class MetricsCollectorBase:
"""Context manager to record operation duration and status."""
raise NotImplementedError
def record_operation_result(
self,
operation: str,
bank_id: str,
success: bool,
duration: float,
source: str = "api",
budget: str | None = None,
max_tokens: int | None = None,
):
"""Record a single completed operation with an explicit success label."""
raise NotImplementedError
def record_llm_call(
self,
provider: str,
@@ -312,19 +254,6 @@ class NoOpMetricsCollector(MetricsCollectorBase):
"""No-op context manager."""
yield
def record_operation_result(
self,
operation: str,
bank_id: str,
success: bool,
duration: float,
source: str = "api",
budget: str | None = None,
max_tokens: int | None = None,
):
"""No-op operation result recording."""
pass
def record_llm_call(
self,
provider: str,
@@ -432,13 +361,6 @@ class MetricsCollector(MetricsCollectorBase):
# DB pool metrics holder (set via set_db_pool)
self._db_pool: "asyncpg.Pool | None" = None
# Backlog / queue-depth gauge caches, refreshed by a background task
# (see _setup_backlog_metrics) so the scrape path stays synchronous.
self._async_ops_counts: dict[_AsyncOpKey, int] = {}
self._consolidation_backlog: dict[_BacklogKey, int] = {}
self._consolidation_failed: dict[_BacklogKey, int] = {}
self._backlog_task: "asyncio.Task | None" = None
@contextmanager
def record_operation(
self,
@@ -464,6 +386,18 @@ class MetricsCollector(MetricsCollectorBase):
max_tokens: Optional max tokens for the operation
"""
start_time = time.time()
attributes = {
"operation": operation,
"source": source,
"tenant": _get_tenant(),
}
if self._include_bank_id:
attributes["bank_id"] = bank_id
if budget:
attributes["budget"] = budget
if max_tokens:
attributes["max_tokens"] = str(max_tokens)
success = True
cancelled = False
try:
@@ -482,51 +416,14 @@ class MetricsCollector(MetricsCollectorBase):
raise
finally:
if not cancelled:
self.record_operation_result(
operation,
bank_id,
success=success,
duration=time.time() - start_time,
source=source,
budget=budget,
max_tokens=max_tokens,
)
duration = time.time() - start_time
attributes["success"] = str(success).lower()
def record_operation_result(
self,
operation: str,
bank_id: str,
success: bool,
duration: float,
source: str = "api",
budget: str | None = None,
max_tokens: int | None = None,
):
"""Record a single completed operation (duration + count) with a success label.
# Record duration
self.operation_duration.record(duration, attributes)
Direct (non-context-manager) recording for code paths that need explicit
success control rather than the exception-based ``record_operation`` e.g.
the async worker, where deferrals/retries are not terminal outcomes and must
not be counted as completions.
"""
attributes = {
"operation": operation,
"source": source,
"tenant": _get_tenant(),
}
if self._include_bank_id:
attributes["bank_id"] = bank_id
if budget:
attributes["budget"] = budget
if max_tokens:
attributes["max_tokens"] = str(max_tokens)
attributes["success"] = str(success).lower()
# Record duration
self.operation_duration.record(duration, attributes)
# Record operation count
self.operation_total.add(1, attributes)
# Record operation count
self.operation_total.add(1, attributes)
def record_llm_call(
self,
@@ -731,10 +628,6 @@ class MetricsCollector(MetricsCollectorBase):
"""
self._db_pool = pool
self._setup_db_pool_metrics()
from .config import get_config
if get_config().metrics_backlog_enabled:
self._setup_backlog_metrics()
def _setup_db_pool_metrics(self):
"""Set up observable gauges for database pool metrics."""
@@ -800,192 +693,6 @@ class MetricsCollector(MetricsCollectorBase):
unit="{connections}",
)
def _setup_backlog_metrics(self):
"""Observable gauges for the async-operation queue and the
consolidation backlog.
These mirror fields the bank-stats endpoint already computes
(``operations_by_status``, ``pending_consolidation``,
``failed_consolidation``) but expose them as scrapable gauges, so
queue depth and backlog can be trended and alerted on instead of only
polled per-bank over HTTP. The two motivating questions both come for
free here: "is the worker keeping up?" (async-op queue) and "is the
knowledge base caught up?" (consolidation backlog) — including the
``processing`` state, which is the only signal that surfaces a hung
operation stuck holding a worker slot.
Counts are aggregate ``COUNT`` queries, so a background task refreshes
a cache every ``BACKLOG_METRICS_REFRESH_SECONDS`` and these callbacks
read it keeping the scrape path synchronous, the same approach as
the db-pool gauges above.
"""
if self._backlog_task is not None:
return # already started for this collector
def get_async_operations(_options):
for key, value in list(self._async_ops_counts.items()):
attrs = {"tenant": key.tenant, "operation_type": key.operation_type, "status": key.status}
if key.bank_id is not None:
attrs["bank_id"] = key.bank_id
yield metrics.Observation(value, attrs)
def get_consolidation_backlog(_options):
for key, value in list(self._consolidation_backlog.items()):
attrs = {"tenant": key.tenant}
if key.bank_id is not None:
attrs["bank_id"] = key.bank_id
yield metrics.Observation(value, attrs)
def get_consolidation_failed(_options):
for key, value in list(self._consolidation_failed.items()):
attrs = {"tenant": key.tenant}
if key.bank_id is not None:
attrs["bank_id"] = key.bank_id
yield metrics.Observation(value, attrs)
self.meter.create_observable_gauge(
name="hindsight.async_operations",
callbacks=[get_async_operations],
description="Async operations in a non-terminal state, by operation_type and status "
"(pending=queued backlog, processing=in-flight, failed=stranded)",
unit="{operations}",
)
self.meter.create_observable_gauge(
name="hindsight.consolidation.backlog",
callbacks=[get_consolidation_backlog],
description="Source memories (experience/world) not yet consolidated into observations",
unit="{memories}",
)
self.meter.create_observable_gauge(
name="hindsight.consolidation.failed",
callbacks=[get_consolidation_failed],
description="Source memories whose consolidation permanently failed "
"(recoverable via the consolidation recovery endpoint)",
unit="{memories}",
)
# Drive the caches from a background task on the running loop.
# set_db_pool runs during async startup, so a loop is normally present;
# if not, the gauges simply stay empty rather than crashing collection.
try:
loop = asyncio.get_running_loop()
except RuntimeError:
logger.warning("No running event loop; backlog metrics disabled")
return
# Process-lifetime task: there is no collector teardown hook to cancel it
# on, so it's torn down with the event loop at process shutdown. If a
# shutdown path is ever added, cancel self._backlog_task there.
self._backlog_task = loop.create_task(self._backlog_refresh_loop())
async def _backlog_refresh_loop(self):
"""Periodically refresh the backlog / queue-depth caches."""
while True:
try:
await self._refresh_backlog()
except Exception:
logger.debug("Backlog metrics refresh failed", exc_info=True)
await asyncio.sleep(BACKLOG_METRICS_REFRESH_SECONDS)
async def _refresh_backlog(self):
"""Recount the async-operation queue and consolidation backlog across
every provisioned Hindsight schema.
Per-bank labels are gated behind ``metrics_include_bank_id`` (off by
default) to keep cardinality bounded; when off, counts are aggregated
per tenant/schema. All SQL here is PostgreSQL-specific (``FILTER``,
``information_schema``), which is consistent with this collector
already being bound to an asyncpg pool.
"""
if self._db_pool is None:
return
async_ops: dict[_AsyncOpKey, int] = {}
backlog: dict[_BacklogKey, int] = {}
failed: dict[_BacklogKey, int] = {}
per_bank = self._include_bank_id
bank_sel = "bank_id, " if per_bank else ""
bank_grp = " GROUP BY bank_id" if per_bank else ""
async with self._db_pool.acquire() as conn:
# memory_units is the central per-tenant table; its presence marks a
# provisioned Hindsight schema.
schema_rows = await conn.fetch(
"SELECT table_schema FROM information_schema.tables WHERE table_name = 'memory_units'"
)
for schema_row in schema_rows:
schema = schema_row["table_schema"]
# Worker queue depth — mirrors operations_by_status, split by
# operation_type. Terminal states (completed/cancelled) are
# excluded on purpose: a gauge of finished work grows without
# bound and says nothing about current load.
# Index: idx_async_operations_status.
ops_grp = "operation_type, status" + (", bank_id" if per_bank else "")
try:
rows = await conn.fetch(
f"SELECT operation_type, status, {bank_sel}COUNT(*) AS count "
f'FROM "{schema}".async_operations '
"WHERE status IN ('pending', 'processing', 'failed') "
f"GROUP BY {ops_grp}"
)
for row in rows:
bank = row["bank_id"] if per_bank else None
key = _AsyncOpKey(schema, row["operation_type"] or "unknown", row["status"], bank)
async_ops[key] = async_ops.get(key, 0) + int(row["count"])
except Exception:
logger.debug("Async-ops queue query failed for schema %s", schema, exc_info=True)
# Consolidation backlog + stranded counts. Two separate COUNT(*)
# queries rather than one with two FILTERs — each WHERE matches a
# partial-index predicate exactly:
# idx_memory_units_unconsolidated WHERE consolidated_at IS NULL ...
# idx_memory_units_consolidation_failed WHERE consolidation_failed_at IS NOT NULL ...
# GROUP BY bank_id still composes — bank_id is each index's lead column.
#
# The backlog count runs with seqscan disabled in a scoped
# transaction. The partial index matches its predicate, but
# `consolidated_at IS NULL` is true for a large fraction of the
# table (every observation has a null consolidated_at), so the
# planner misjudges selectivity and otherwise seq-scans the whole
# (largest) table on every refresh — verified on a 114k-row table
# via EXPLAIN: seq scan ~92 ms vs index scan ~0.1 ms. SET LOCAL
# forces the index path and resets at transaction end. The failed
# count below needs no such nudge: `consolidation_failed_at IS NOT
# NULL` is rare, so its index is chosen on cost.
try:
async with conn.transaction():
await conn.execute("SET LOCAL enable_seqscan = off")
rows = await conn.fetch(
f"SELECT {bank_sel}COUNT(*) AS count "
f'FROM "{schema}".memory_units '
"WHERE consolidated_at IS NULL AND fact_type IN ('experience', 'world')"
f"{bank_grp}"
)
for row in rows:
bank = row["bank_id"] if per_bank else None
key = _BacklogKey(schema, bank)
backlog[key] = backlog.get(key, 0) + int(row["count"])
except Exception:
logger.debug("Consolidation backlog query failed for schema %s", schema, exc_info=True)
try:
rows = await conn.fetch(
f"SELECT {bank_sel}COUNT(*) AS count "
f'FROM "{schema}".memory_units '
"WHERE consolidation_failed_at IS NOT NULL AND fact_type IN ('experience', 'world')"
f"{bank_grp}"
)
for row in rows:
bank = row["bank_id"] if per_bank else None
key = _BacklogKey(schema, bank)
failed[key] = failed.get(key, 0) + int(row["count"])
except Exception:
logger.debug("Consolidation failed query failed for schema %s", schema, exc_info=True)
self._async_ops_counts = async_ops
self._consolidation_backlog = backlog
self._consolidation_failed = failed
# Global metrics collector instance (defaults to no-op)
_metrics_collector: MetricsCollectorBase = NoOpMetricsCollector()
+18 -14
View File
@@ -686,20 +686,24 @@ def ensure_vector_extension(
if not current_index_info:
if table_name == "memory_units" and uses_per_bank_vector_indexes(target_ext):
# Per-bank backends never use a GLOBAL memory_units vector index.
# Every vector search is bank + fact_type scoped and served by the
# per-(bank, fact_type) partial indexes created at bank-creation time
# (bank_utils.create_bank_vector_indexes); the planner never picks a
# global index when bank_id is in the WHERE clause, which is exactly
# why migration d5e6f7a8b9c0 drops it for these backends. So don't
# create one here either — not even on an empty schema with no per-bank
# indexes yet (those are built when the first bank is created). Verified
# via EXPLAIN: the query uses idx_mu_emb_* whether or not the global
# index exists, so creating it is dead weight.
logger.debug(
f"Per-bank vector backend ({target_ext}); skipping global {index_name} creation on {table_name}"
)
continue
# Check whether per-bank partial vector indexes already cover this table
# (created by the bank_utils lifecycle — no global index needed in that case)
per_bank_index_count = conn.execute(
text("""
SELECT COUNT(*)
FROM pg_indexes
WHERE schemaname = :schema
AND tablename = :table_name
AND indexname LIKE 'idx_mu_emb_%'
"""),
{"schema": schema_name, "table_name": table_name},
).scalar()
if per_bank_index_count and per_bank_index_count > 0:
logger.debug(
f"No global embedding index on {table_name}, but {per_bank_index_count} "
f"per-bank partial vector indexes exist — skipping global index creation"
)
continue
logger.warning(f"No embedding index found for {table_name}, will create it if safe")
mismatched_tables.append((table_name, index_name, None, row_count))
continue
@@ -20,23 +20,9 @@ from dataclasses import dataclass
from typing import TYPE_CHECKING, Any
from ..engine.schema import fq_table_explicit as fq_table
from ..metrics import get_metrics_collector
from .exceptions import DeferOperation, RetryTaskAt
from .stage import StageHolder, bind_holder
# Map DB operation_type -> metric `operation` label, collapsing the retain
# variants onto "retain" so async worker completions land on the same
# operation="retain" series the synchronous API path emits. Unknown types
# pass through unchanged.
_RETAIN_OP_TYPES = {"retain", "batch_retain", "file_convert_retain"}
def _metric_operation_label(operation_type: str | None) -> str:
if operation_type in _RETAIN_OP_TYPES:
return "retain"
return operation_type or "unknown"
if TYPE_CHECKING:
from hindsight_api.engine.db.base import DatabaseBackend, DatabaseConnection
from hindsight_api.extensions.tenant import TenantExtension
@@ -715,24 +701,6 @@ class WorkerPoller:
"""
task_type = task.task_dict.get("type", "unknown")
bank_id = task.task_dict.get("bank_id", "unknown")
# Operation metric (source="worker"): record on terminal outcomes only, so
# async worker throughput and latency (retain, consolidation and the other
# worker task types) are visible in Prometheus. Prefer the DB-authoritative
# operation_type.
#
# success semantics are deliberately narrow: success=false means the task
# raised out to the poller (an unexpected error, or retry-exhausted). It does
# NOT capture deterministic failures that the executor handles itself and
# returns from normally (file_convert_retain, non-retryable errors via
# memory_engine.execute_task) — those record success=true here. Treat this as
# a completion-throughput signal, not a failure-rate one: for authoritative
# failure visibility use the hindsight_async_operations{status="failed"} gauge,
# which reads each operation's final DB status.
op_label = _metric_operation_label(task.task_dict.get("operation_type") or task_type)
op_start = time.time()
metrics = get_metrics_collector()
# None = not a terminal outcome (deferred/retried) → no metric.
terminal_success: bool | None = None
# Bind the stage holder in this task's own contextvar scope so engine
# code running under us can update it via stage.set_stage(). If holder
@@ -749,28 +717,14 @@ class WorkerPoller:
task.task_dict["_schema"] = task.schema
await self._executor(task.task_dict)
logger.debug(f"Task {task.operation_id} execution finished")
terminal_success = True
except DeferOperation as e:
# Deferral is not a terminal outcome — do not record a completion.
await self._defer_operation(task.operation_id, e.exec_date, e.reason, task.schema)
except RetryTaskAt as e:
# Retry is not a terminal outcome — do not record a completion.
await self._schedule_retry(task.operation_id, e.retry_at, str(e), task.schema)
except Exception as e:
logger.error(f"Task {task.operation_id} failed: {e}")
traceback.print_exc()
await self._mark_failed(task.operation_id, str(e), task.schema)
terminal_success = False
# Record the metric outside the executor's exception scope so a metrics
# reporting failure can never be mistaken for a task failure and flip terminal state.
if terminal_success is not None:
try:
metrics.record_operation_result(
op_label, bank_id, success=terminal_success, duration=time.time() - op_start, source="worker"
)
except Exception:
logger.warning(f"Failed to record worker operation metric for {task.operation_id}", exc_info=True)
async def recover_own_tasks(self) -> int:
"""
+3 -3
View File
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
[project]
name = "hindsight-api-slim"
version = "0.8.3"
version = "0.8.2"
description = "Hindsight: Agent Memory That Works Like Human Memory"
readme = "README.md"
requires-python = ">=3.11"
@@ -60,10 +60,10 @@ dependencies = [
"pyasn1>=0.6.3", # DoS vulnerability fix
"urllib3>=2.7.0", # Decompression-bomb safeguards bypass + sensitive header forwarding fixes
"langchain-core>=1.2.22", # Path traversal in legacy load_prompt functions fix
"langsmith>=0.8.18", # GHSA-f4xh-w4cj-qxq8: arbitrary server-side file read in TracingMiddleware fix (supersedes >=0.6.3 SSRF tracing-header-injection floor)
"langsmith>=0.6.3", # SSRF via tracing header injection fix
"protobuf>=6.33.5", # JSON recursion depth bypass fix
"pillow>=12.1.1", # Out-of-bounds write in PSD image loading fix
"cryptography>=48.0.1", # GHSA-537c-gmf6-5ccf: bundled-OpenSSL OOB read fix needs >=48.0.1. Prior <47 cap (47.0.0 SIGILL on ARM64 Docker/Podman, pyca/cryptography#14733) lifted — 47/48/49 verified importing + RSA sign/verify cleanly on linux/arm64 (Docker on Apple Silicon) and native arm64 macOS; upstream issue closed unconfirmed.
"cryptography>=46.0.6,<47", # Incomplete DNS name constraint enforcement fix; cap <47 — 47.0.0 SIGILLs on some ARM64 Linux VMs (Docker/Podman on Apple Silicon), pyca/cryptography#14733
"filelock>=3.20.1", # TOCTOU race condition fix
"authlib>=1.6.9", # Account takeover/JWS header injection vulnerability fix
"pyjwt>=2.12.0", # Accepts unknown crit header extensions fix
@@ -547,36 +547,3 @@ async def test_run_migration_with_schema_only_runs_requested_schema(monkeypatch)
assert calls["run_migrations"] == [("resolved::postgresql://test", "tenant_demo")]
assert calls["ensure_vector_extension"] == [("resolved::postgresql://test", "pgvector", "tenant_demo")]
assert calls["ensure_text_search_extension"] == [("resolved::postgresql://test", "native", "", "tenant_demo")]
@pytest.mark.parametrize(
("ensure_extensions", "expected"),
[(True, True), (False, False)],
)
@pytest.mark.asyncio
async def test_run_migration_threads_ensure_extensions_flag(monkeypatch, ensure_extensions, expected):
"""The --skip-extension-reconcile flag (ensure_extensions=False) must reach run_migrations_for_schemas.
The post-migration vector/text-search reconcile only does work on a backend change, so operators
can skip it on a no-change re-migration over many tenant schemas. Verify the flag is threaded through
rather than silently dropped.
"""
monkeypatch.setenv("HINDSIGHT_API_DATABASE_URL", "postgresql://test")
captured: dict = {}
async def fake_resolve_database_url(db_url: str) -> str:
return f"resolved::{db_url}"
def fake_run_migrations_for_schemas(database_url, schemas, **kwargs):
captured["ensure_extensions"] = kwargs.get("ensure_extensions")
monkeypatch.setattr(admin_cli, "load_extension", lambda *args, **kwargs: None)
monkeypatch.setattr(admin_cli, "resolve_database_url", fake_resolve_database_url)
from hindsight_api import migrations as migrations_module
monkeypatch.setattr(migrations_module, "run_migrations_for_schemas", fake_run_migrations_for_schemas)
await admin_cli._run_migration("postgresql://test", schema="tenant_demo", ensure_extensions=ensure_extensions)
assert captured["ensure_extensions"] is expected
@@ -1,111 +0,0 @@
"""Regression tests for issue #1002 — Anthropic structured output via forced tool_use.
When strict_schema=True, AnthropicLLM.call() must request the schema through a single
forced tool_use tool (tool_choice={"type":"tool",...}) and read the validated args from
the tool_use block, NOT inject the schema as text and json.loads() the reply (which caused
a ~1:1 invalid-JSON retry storm / OOM in production).
"""
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from pydantic import BaseModel
class _Decision(BaseModel):
action: str
reason: str
def _make_anthropic_provider():
with patch("anthropic.AsyncAnthropic") as mock_client_cls:
mock_client_cls.return_value = MagicMock()
from hindsight_api.engine.providers.anthropic_llm import AnthropicLLM
provider = AnthropicLLM(
provider="anthropic",
api_key="fake-key",
base_url="",
model="claude-sonnet-4-20250514",
)
provider._client = MagicMock()
return provider
def _tool_use_response(args: dict):
block = MagicMock()
block.type = "tool_use"
block.name = "structured_response"
block.input = args
resp = MagicMock()
resp.content = [block]
resp.usage = MagicMock(input_tokens=5, output_tokens=2, cache_read_input_tokens=0)
resp.stop_reason = "tool_use"
return resp
@pytest.mark.asyncio
async def test_strict_schema_uses_forced_tool_choice():
"""strict_schema=True ⇒ a single tool is defined and tool_choice forces it (no schema text-injection)."""
provider = _make_anthropic_provider()
provider._client.messages.create = AsyncMock(return_value=_tool_use_response({"action": "skip", "reason": "dup"}))
with patch("hindsight_api.engine.providers.anthropic_llm.get_metrics_collector"):
result = await provider.call(
messages=[{"role": "user", "content": "decide"}],
response_format=_Decision,
strict_schema=True,
scope="test",
max_retries=0,
)
kwargs = provider._client.messages.create.call_args.kwargs
# forced tool_use requested
assert "tools" in kwargs and len(kwargs["tools"]) == 1
assert kwargs["tool_choice"] == {"type": "tool", "name": "structured_response"}
# schema NOT injected as text into the system prompt
assert "valid JSON matching this schema" not in (kwargs.get("system") or "")
# validated model returned straight from tool_use.input
assert isinstance(result, _Decision)
assert result.action == "skip"
@pytest.mark.asyncio
async def test_strict_schema_tool_use_never_hits_json_retry_loop():
"""A tool_use response is structurally valid → no second messages.create call (no retry storm)."""
provider = _make_anthropic_provider()
create = AsyncMock(return_value=_tool_use_response({"action": "keep", "reason": "novel"}))
provider._client.messages.create = create
with patch("hindsight_api.engine.providers.anthropic_llm.get_metrics_collector"):
await provider.call(
messages=[{"role": "user", "content": "x"}],
response_format=_Decision,
strict_schema=True,
scope="test",
max_retries=10, # would allow 11 attempts on the old text-parse path
)
assert create.await_count == 1 # exactly one call — the bug was N retries on malformed text
@pytest.mark.asyncio
async def test_non_strict_keeps_text_injection_fallback():
"""strict_schema=False (default) preserves the legacy schema-in-prompt behavior."""
provider = _make_anthropic_provider()
block = MagicMock()
block.type = "text"
block.text = '{"action":"skip","reason":"d"}'
resp = MagicMock()
resp.content = [block]
resp.usage = MagicMock(input_tokens=5, output_tokens=2, cache_read_input_tokens=0)
resp.stop_reason = "end_turn"
provider._client.messages.create = AsyncMock(return_value=resp)
with patch("hindsight_api.engine.providers.anthropic_llm.get_metrics_collector"):
result = await provider.call(
messages=[{"role": "user", "content": "decide"}],
response_format=_Decision,
strict_schema=False,
scope="test",
max_retries=0,
)
kwargs = provider._client.messages.create.call_args.kwargs
assert "tools" not in kwargs # no forced tool when not strict
assert "valid JSON matching this schema" in (kwargs.get("system") or "")
assert isinstance(result, _Decision)
@@ -1,90 +0,0 @@
"""Regression test: submitting an async op for a bank that doesn't exist must
raise a clean validation error, not a raw asyncpg `ForeignKeyViolationError`.
`_submit_async_operation` inserts into `async_operations`, which has an FK to
`banks.bank_id`. If a caller submits for a missing bank (typo, race against a
deletion, integration that derives bank IDs before the bank is created), the
INSERT raises `asyncpg.exceptions.ForeignKeyViolationError`. The FastAPI
endpoint's broad `except Exception` then surfaces it as a 500 — but this is
a client error, not a server error, and should be a 404.
This test exercises the call directly via `MemoryEngine.submit_async_*` so
the failure mode is observable without spinning up the HTTP layer.
"""
import uuid
import pytest
from hindsight_api.extensions.operation_validator import OperationValidationError
pytestmark = pytest.mark.xdist_group("async_submit_bank_not_found_tests")
@pytest.fixture
def no_inline_execution(memory):
"""Prevent SyncTaskBackend from running the submitted op inline so we
only test the submit-path failure, not downstream execution."""
async def _noop(_payload):
return None
original = memory._task_backend.submit_task
memory._task_backend.submit_task = _noop
yield
memory._task_backend.submit_task = original
@pytest.mark.asyncio
async def test_consolidation_submit_on_missing_bank_raises_validation_error(
memory, request_context, no_inline_execution
):
"""A `/consolidate` submit against a bank that doesn't exist must raise
OperationValidationError(404), not a raw asyncpg FK violation that bubbles
out as a 500 from the API."""
missing_bank = f"does-not-exist-{uuid.uuid4().hex[:8]}"
with pytest.raises(OperationValidationError) as exc_info:
await memory.submit_async_consolidation(
bank_id=missing_bank,
request_context=request_context,
)
assert exc_info.value.status_code == 404
assert missing_bank in exc_info.value.reason
@pytest.mark.asyncio
async def test_scoped_consolidation_submit_on_missing_bank_raises_validation_error(
memory, request_context, no_inline_execution
):
"""Scoped consolidates (with `observation_scopes`) take the
`dedupe_by_bank=False` branch, which historically skipped the bank lock
entirely and went straight to the FK-violating INSERT. Same 404 contract."""
missing_bank = f"does-not-exist-{uuid.uuid4().hex[:8]}"
with pytest.raises(OperationValidationError) as exc_info:
await memory.submit_async_consolidation(
bank_id=missing_bank,
request_context=request_context,
observation_scopes=[{"tag": "anything"}],
)
assert exc_info.value.status_code == 404
assert missing_bank in exc_info.value.reason
@pytest.mark.asyncio
async def test_graph_maintenance_on_missing_bank_short_circuits(memory, request_context, no_inline_execution):
"""`submit_async_graph_maintenance` has its own short-circuit that checks
the per-bank queue before calling `_submit_async_operation`. A missing
bank means an empty queue, so it returns `no_work=True` without reaching
the FK-violating INSERT. This test pins that behaviour."""
missing_bank = f"does-not-exist-{uuid.uuid4().hex[:8]}"
result = await memory.submit_async_graph_maintenance(
bank_id=missing_bank,
request_context=request_context,
)
assert result == {"operation_id": None, "no_work": True}
@@ -1,244 +0,0 @@
"""
Tests for the async-operation queue and consolidation backlog gauges
(``_setup_backlog_metrics`` / ``_refresh_backlog`` in metrics.py).
These gauges expose, as scrapable time-series, the same counts the bank-stats
endpoint already returns per bank (``operations_by_status``,
``pending_consolidation``, ``failed_consolidation``):
- ``hindsight_async_operations{operation_type,status}`` worker queue depth
(pending=backlog, processing=in-flight, failed=stranded)
- ``hindsight_consolidation_backlog`` source memories not yet consolidated
- ``hindsight_consolidation_failed`` source memories permanently failed
"""
from unittest.mock import MagicMock, patch
import pytest
from hindsight_api.metrics import MetricsCollector, _AsyncOpKey, _BacklogKey
class _FakeTxn:
async def __aenter__(self):
return None
async def __aexit__(self, *exc):
return False
class _FakeConn:
"""asyncpg-like connection whose fetch() is dispatched by SQL substring."""
def __init__(self, fetch_fn):
self._fetch_fn = fetch_fn
self.executed = []
async def fetch(self, sql, *args):
return self._fetch_fn(sql, *args)
async def execute(self, sql, *args):
self.executed.append(sql)
def transaction(self):
return _FakeTxn()
class _FakeAcquire:
def __init__(self, conn):
self._conn = conn
async def __aenter__(self):
return self._conn
async def __aexit__(self, *exc):
return False
class _FakePool:
def __init__(self, fetch_fn):
self._conn = _FakeConn(fetch_fn)
def acquire(self):
return _FakeAcquire(self._conn)
def _collector(include_bank_id=False):
mock_config = MagicMock()
mock_config.metrics_include_bank_id = include_bank_id
with (
patch("hindsight_api.metrics.get_meter", return_value=MagicMock()),
patch("hindsight_api.config.get_config", return_value=mock_config),
):
return MetricsCollector()
def _set_db_pool_with_backlog_enabled(collector, pool):
"""Call set_db_pool with the backlog flag forced on (it's off by default)."""
mock_config = MagicMock()
mock_config.metrics_backlog_enabled = True
with patch("hindsight_api.config.get_config", return_value=mock_config):
collector.set_db_pool(pool)
def _rows_for(sql):
"""Canned results, keyed off distinctive substrings of each query."""
if "information_schema.tables" in sql:
return [{"table_schema": "public"}]
if "async_operations" in sql:
return [
{"operation_type": "retain", "status": "pending", "count": 5},
{"operation_type": "consolidation", "status": "pending", "count": 12},
{"operation_type": "consolidation", "status": "processing", "count": 1},
{"operation_type": "consolidation", "status": "failed", "count": 2},
]
if "memory_units" in sql and "consolidated_at IS NULL" in sql:
return [{"count": 42}]
if "memory_units" in sql and "consolidation_failed_at IS NOT NULL" in sql:
return [{"count": 3}]
return []
@pytest.mark.asyncio
async def test_refresh_backlog_aggregates_queue_and_consolidation():
collector = _collector(include_bank_id=False)
collector._db_pool = _FakePool(lambda sql, *a: _rows_for(sql))
await collector._refresh_backlog()
# Worker queue depth keyed by (schema, operation_type, status, bank=None)
assert collector._async_ops_counts[("public", "retain", "pending", None)] == 5
assert collector._async_ops_counts[("public", "consolidation", "pending", None)] == 12
assert collector._async_ops_counts[("public", "consolidation", "processing", None)] == 1
assert collector._async_ops_counts[("public", "consolidation", "failed", None)] == 2
# Consolidation backlog (source memories), keyed by (schema, bank=None)
assert collector._consolidation_backlog[("public", None)] == 42
assert collector._consolidation_failed[("public", None)] == 3
@pytest.mark.asyncio
async def test_refresh_backlog_uses_index_matched_predicates_not_filter_scan():
"""Backlog/failed must be two separate COUNT(*) queries whose WHERE matches
a partial-index predicate exactly (no FILTER over a full-table scan), and
the queue query must exclude terminal statuses."""
captured = []
collector = _collector()
collector._db_pool = _FakePool(lambda sql, *a: (captured.append(sql), _rows_for(sql))[1])
await collector._refresh_backlog()
mem_queries = [s for s in captured if "memory_units" in s and "COUNT(*)" in s]
assert len(mem_queries) == 2 # split, not a single two-FILTER aggregate
assert all("FILTER" not in s for s in mem_queries)
assert any("consolidated_at IS NULL AND fact_type IN ('experience', 'world')" in s for s in mem_queries)
assert any("consolidation_failed_at IS NOT NULL AND fact_type IN ('experience', 'world')" in s for s in mem_queries)
ops_sql = next(s for s in captured if "async_operations" in s and "GROUP BY" in s)
assert "status IN ('pending', 'processing', 'failed')" in ops_sql
assert "completed" not in ops_sql and "cancelled" not in ops_sql
@pytest.mark.asyncio
async def test_backlog_count_runs_with_seqscan_disabled():
"""`consolidated_at IS NULL` is true for a large fraction of the table, so
the planner misjudges selectivity and won't use the partial index without a
nudge the backlog count must issue SET LOCAL enable_seqscan=off."""
collector = _collector()
pool = _FakePool(lambda sql, *a: _rows_for(sql))
collector._db_pool = pool
await collector._refresh_backlog()
assert any("enable_seqscan" in s.lower() and "off" in s.lower() for s in pool._conn.executed)
# the result is still correct under the nudge
assert collector._consolidation_backlog[("public", None)] == 42
@pytest.mark.asyncio
async def test_refresh_backlog_per_bank_labels_and_group_by_when_enabled():
"""With metrics_include_bank_id on, bank_id enters the cache key and the
SQL switches to GROUP BY bank_id."""
captured = []
def fetch(sql, *a):
captured.append(sql)
if "information_schema.tables" in sql:
return [{"table_schema": "public"}]
if "async_operations" in sql:
return [{"operation_type": "retain", "status": "pending", "bank_id": "bankA", "count": 4}]
if "memory_units" in sql and "consolidated_at IS NULL" in sql:
return [{"bank_id": "bankA", "count": 11}]
if "memory_units" in sql and "consolidation_failed_at IS NOT NULL" in sql:
return [{"bank_id": "bankA", "count": 2}]
return []
collector = _collector(include_bank_id=True)
collector._db_pool = _FakePool(fetch)
await collector._refresh_backlog()
assert collector._async_ops_counts[("public", "retain", "pending", "bankA")] == 4
assert collector._consolidation_backlog[("public", "bankA")] == 11
assert collector._consolidation_failed[("public", "bankA")] == 2
# bank_id must be grouped in every per-bank count query
assert all("GROUP BY bank_id" in s for s in captured if "memory_units" in s and "COUNT(*)" in s)
def test_gauges_register_and_emit_cached_values_without_bank_id():
collector = _collector(include_bank_id=False)
# Sync call: no running loop, so gauges register but no background task spawns.
_set_db_pool_with_backlog_enabled(collector, MagicMock())
gauges = {
c.kwargs["name"]: c.kwargs["callbacks"][0]
for c in collector.meter.create_observable_gauge.call_args_list
if "callbacks" in c.kwargs
}
assert "hindsight.async_operations" in gauges
assert "hindsight.consolidation.backlog" in gauges
assert "hindsight.consolidation.failed" in gauges
collector._async_ops_counts = {
_AsyncOpKey("public", "retain", "pending", None): 7,
_AsyncOpKey("public", "consolidation", "processing", None): 1,
}
collector._consolidation_backlog = {_BacklogKey("public", None): 9}
obs = list(gauges["hindsight.async_operations"](None))
by_label = {(o.attributes["operation_type"], o.attributes["status"]): o.value for o in obs}
assert by_label[("retain", "pending")] == 7
assert by_label[("consolidation", "processing")] == 1
assert all("bank_id" not in o.attributes for o in obs) # cardinality guard
backlog_obs = list(gauges["hindsight.consolidation.backlog"](None))
assert backlog_obs[0].value == 9
assert backlog_obs[0].attributes["tenant"] == "public"
def test_gauge_emits_bank_id_attribute_when_present():
collector = _collector(include_bank_id=True)
_set_db_pool_with_backlog_enabled(collector, MagicMock())
gauges = {
c.kwargs["name"]: c.kwargs["callbacks"][0]
for c in collector.meter.create_observable_gauge.call_args_list
if "callbacks" in c.kwargs
}
collector._consolidation_backlog = {_BacklogKey("public", "bankA"): 4}
obs = list(gauges["hindsight.consolidation.backlog"](None))
assert obs[0].value == 4
assert obs[0].attributes["bank_id"] == "bankA"
def test_backlog_gauges_not_registered_when_flag_disabled():
"""Backlog metrics are off by default: set_db_pool must not register the
gauges unless metrics_backlog_enabled is set."""
collector = _collector()
mock_config = MagicMock()
mock_config.metrics_backlog_enabled = False
with patch("hindsight_api.config.get_config", return_value=mock_config):
collector.set_db_pool(MagicMock())
names = [
c.kwargs.get("name") for c in collector.meter.create_observable_gauge.call_args_list if "callbacks" in c.kwargs
]
assert "hindsight.async_operations" not in names
assert "hindsight.consolidation.backlog" not in names
assert "hindsight.consolidation.failed" not in names
assert collector._backlog_task is None
@@ -186,43 +186,6 @@ async def test_invalidate_drops_entry() -> None:
assert calls[0] == 2
@pytest.mark.asyncio
async def test_invalidate_detaches_in_flight_loader() -> None:
cache = BankStatsCache(ttl_seconds=60, max_entries=100)
stale_started = asyncio.Event()
release_stale = asyncio.Event()
fresh_started = asyncio.Event()
async def stale_loader() -> dict[str, Any]:
stale_started.set()
await release_stale.wait()
return {"v": "stale"}
async def fresh_loader() -> dict[str, Any]:
fresh_started.set()
return {"v": "fresh"}
stale_task = asyncio.create_task(cache.get_or_load("schema", "bank", stale_loader))
await stale_started.wait()
await cache.invalidate("schema", "bank")
# A request after invalidation must start a new load instead of joining the
# pre-invalidation query, which may contain data from before a bank write.
fresh_result = await asyncio.wait_for(cache.get_or_load("schema", "bank", fresh_loader), timeout=1)
assert fresh_started.is_set()
assert fresh_result == {"v": "fresh"}
release_stale.set()
assert await stale_task == {"v": "stale"}
# The stale loader completed last, but must not overwrite the fresh value.
async def should_not_run() -> dict[str, Any]:
raise AssertionError("fresh value was not cached")
cached = await cache.get_or_load("schema", "bank", should_not_run)
assert cached == {"v": "fresh"}
@pytest.mark.asyncio
async def test_clear_drops_all_entries() -> None:
cache = BankStatsCache(ttl_seconds=60, max_entries=100)
@@ -236,27 +199,3 @@ async def test_clear_drops_all_entries() -> None:
await cache.get_or_load("s", "a", loader)
await cache.get_or_load("s", "b", loader)
assert calls[0] == 4
@pytest.mark.asyncio
async def test_clear_detaches_in_flight_loaders() -> None:
cache = BankStatsCache(ttl_seconds=60, max_entries=100)
stale_started = asyncio.Event()
release_stale = asyncio.Event()
async def stale_loader() -> dict[str, Any]:
stale_started.set()
await release_stale.wait()
return {"v": "stale"}
async def fresh_loader() -> dict[str, Any]:
return {"v": "fresh"}
stale_task = asyncio.create_task(cache.get_or_load("schema", "bank", stale_loader))
await stale_started.wait()
await cache.clear()
assert await cache.get_or_load("schema", "bank", fresh_loader) == {"v": "fresh"}
release_stale.set()
assert await stale_task == {"v": "stale"}
assert await cache.get_or_load("schema", "bank", fresh_loader) == {"v": "fresh"}
@@ -1,164 +0,0 @@
"""Regression tests: get_bank_stats cache must be invalidated by mutations.
`get_bank_stats` is served from a short-TTL per-process cache (`BankStatsCache`).
`delete_bank` already invalidates that cache after it mutates counts, but the
other operations that change the same counts `delete_memory_unit`,
`delete_document`, `clear_observations`, and `update_document` (when a tag
change deletes observations) did not, so a client polling stats right after a
deletion would see pre-mutation counts until the TTL expired (up to a minute).
Each test pins a long TTL on the engine's stats cache so that, *without* the
invalidation fix, the second `get_bank_stats` call would be served the stale
cached value and the assertion would fail.
"""
import uuid
import pytest
from hindsight_api import RequestContext
from hindsight_api.engine.bank_stats_cache import BankStatsCache
from hindsight_api.engine.memory_engine import MemoryEngine
# A TTL long enough that, absent invalidation, the warmed cache would still be
# served on the post-mutation read within the same test.
_PINNED_TTL_SECONDS = 300.0
async def _insert_memory(conn, bank_id: str, text: str, fact_type: str = "experience") -> uuid.UUID:
"""Insert a memory unit directly, bypassing the LLM retain pipeline."""
mem_id = uuid.uuid4()
await conn.execute(
"""
INSERT INTO memory_units (id, bank_id, text, fact_type, event_date, created_at, updated_at, consolidated_at)
VALUES ($1, $2, $3, $4, NOW(), NOW(), NOW(), NOW())
""",
mem_id,
bank_id,
text,
fact_type,
)
return mem_id
async def _insert_observation(conn, bank_id: str, text: str, source_memory_ids: list[uuid.UUID]) -> uuid.UUID:
"""Insert an observation unit directly."""
obs_id = uuid.uuid4()
await conn.execute(
"""
INSERT INTO memory_units (
id, bank_id, text, fact_type, event_date, source_memory_ids, proof_count, created_at, updated_at
) VALUES ($1, $2, $3, 'observation', NOW(), $4, $5, NOW(), NOW())
""",
obs_id,
bank_id,
text,
source_memory_ids,
len(source_memory_ids),
)
return obs_id
async def _insert_document(conn, bank_id: str, doc_id: str) -> None:
await conn.execute(
"""
INSERT INTO documents (id, bank_id, original_text, content_hash)
VALUES ($1, $2, $3, $4)
""",
doc_id,
bank_id,
f"text-for-{doc_id}",
doc_id,
)
async def _attach_unit_to_doc(conn, unit_id: uuid.UUID, doc_id: str) -> None:
await conn.execute("UPDATE memory_units SET document_id = $1 WHERE id = $2", doc_id, unit_id)
async def _ensure_bank(memory: MemoryEngine, bank_id: str, request_context: RequestContext) -> None:
await memory.get_bank_profile(bank_id=bank_id, request_context=request_context)
def _pin_cache(memory: MemoryEngine) -> None:
"""Replace the stats cache with one that has a deterministic long TTL."""
memory._bank_stats_cache = BankStatsCache(ttl_seconds=_PINNED_TTL_SECONDS, max_entries=128)
class TestBankStatsCacheInvalidation:
@pytest.mark.asyncio
async def test_delete_memory_unit_invalidates_stats_cache(
self, memory: MemoryEngine, request_context: RequestContext
):
bank_id = f"test-stats-cache-delunit-{uuid.uuid4().hex[:8]}"
await _ensure_bank(memory, bank_id, request_context)
pool = await memory._get_pool()
async with pool.acquire() as conn:
m1 = await _insert_memory(conn, bank_id, "Alice loves hiking.")
await _insert_memory(conn, bank_id, "Bob enjoys cycling.")
_pin_cache(memory)
try:
before = await memory.get_bank_stats(bank_id, request_context=request_context)
assert before["node_counts"].get("experience") == 2
await memory.delete_memory_unit(str(m1), request_context=request_context)
after = await memory.get_bank_stats(bank_id, request_context=request_context)
# Without invalidation the long-TTL cache would still report 2.
assert after["node_counts"].get("experience") == 1
finally:
await memory.delete_bank(bank_id, request_context=request_context)
@pytest.mark.asyncio
async def test_delete_document_invalidates_stats_cache(self, memory: MemoryEngine, request_context: RequestContext):
bank_id = f"test-stats-cache-deldoc-{uuid.uuid4().hex[:8]}"
await _ensure_bank(memory, bank_id, request_context)
document_id = f"doc-{uuid.uuid4().hex[:8]}"
pool = await memory._get_pool()
async with pool.acquire() as conn:
await _insert_document(conn, bank_id, document_id)
unit_id = await _insert_memory(conn, bank_id, "Alice works at Acme.")
await _attach_unit_to_doc(conn, unit_id, document_id)
_pin_cache(memory)
try:
before = await memory.get_bank_stats(bank_id, request_context=request_context)
assert before["total_documents"] == 1
assert before["node_counts"].get("experience") == 1
await memory.delete_document(document_id, bank_id, request_context=request_context)
after = await memory.get_bank_stats(bank_id, request_context=request_context)
# Without invalidation the long-TTL cache would still report 1 document.
assert after["total_documents"] == 0
assert after["node_counts"].get("experience", 0) == 0
finally:
await memory.delete_bank(bank_id, request_context=request_context)
@pytest.mark.asyncio
async def test_clear_observations_invalidates_stats_cache(
self, memory: MemoryEngine, request_context: RequestContext
):
bank_id = f"test-stats-cache-clearobs-{uuid.uuid4().hex[:8]}"
await _ensure_bank(memory, bank_id, request_context)
pool = await memory._get_pool()
async with pool.acquire() as conn:
m1 = await _insert_memory(conn, bank_id, "Alice loves hiking.")
await _insert_observation(conn, bank_id, "Alice enjoys hiking regularly.", [m1])
_pin_cache(memory)
try:
before = await memory.get_bank_stats(bank_id, request_context=request_context)
assert before["total_observations"] == 1
await memory.clear_observations(bank_id, request_context=request_context)
after = await memory.get_bank_stats(bank_id, request_context=request_context)
# Without invalidation the long-TTL cache would still report 1 observation.
assert after["total_observations"] == 0
finally:
await memory.delete_bank(bank_id, request_context=request_context)
-73
View File
@@ -322,76 +322,3 @@ def test_plain_text_lines_not_treated_as_jsonl():
# Sanity: these are not JSON objects (so the JSONL path correctly declined).
with pytest.raises(json.JSONDecodeError):
json.loads(chunks[0])
# ---------------------------------------------------------------------------
# Idempotency — re-chunking a produced chunk must be a no-op (issue #2301)
# ---------------------------------------------------------------------------
#
# The streaming retain pipeline pre-chunks each document once (producer) and then
# re-chunks every piece during extraction (consumer), stamping all sub-chunks of
# one piece with that piece's single chunk_index. If a piece re-split, its
# sub-chunks would derive the same chunk_id = {bank}_{doc}_{index} and the
# ON CONFLICT upsert would fail with CardinalityViolationError. This can only
# happen when structured_chunk_size > max_chars (a chunk legitimately exceeds the
# re-chunk budget); the defaults (structured == max_chars) never trip it.
def _assert_idempotent(text: str, *, max_chars: int, structured_chunk_size: int) -> list[str]:
chunks = chunk_text(text, max_chars=max_chars, structured_chunk_size=structured_chunk_size)
for chunk in chunks:
rechunked = chunk_text(chunk, max_chars=max_chars, structured_chunk_size=structured_chunk_size)
assert rechunked == [chunk], (
f"re-chunking a produced chunk split it again ({len(chunk)} chars -> "
f"{len(rechunked)} pieces) — not idempotent (issue #2301)"
)
return chunks
def test_conversation_turn_over_chunk_size_is_rechunk_stable():
"""A conversation turn larger than max_chars but kept whole by the larger
structured cap must survive a re-chunk unchanged (issue #2301)."""
content = json.dumps([{"role": "assistant", "content": "x" * 6000}])
_assert_idempotent(content, max_chars=3000, structured_chunk_size=5000)
def test_jsonl_line_over_chunk_size_is_rechunk_stable():
"""A single oversized JSONL line, kept whole within the structured cap, must
not be re-split when handed back through chunk_text (issue #2301)."""
text = "\n".join([json.dumps({"event": "x" * 3800}), json.dumps({"event": "small"})])
_assert_idempotent(text, max_chars=3000, structured_chunk_size=4500)
def test_oversized_unit_fragments_stay_within_chunk_budget():
"""A unit past even the structured cap is fragmented as text; no fragment may
exceed max_chars, so a re-chunk leaves the fragments intact (issue #2301)."""
text = "\n".join([json.dumps({"event": "z" * 9000}), json.dumps({"e": "s"})])
chunks = _assert_idempotent(text, max_chars=3000, structured_chunk_size=4500)
assert all(len(c) <= 3000 for c in chunks)
def test_single_json_object_kept_whole_within_structured_cap():
"""A lone JSON object over max_chars but within the structured cap is returned
whole rather than plain-text-split (the basis of re-chunk stability)."""
obj = json.dumps({"role": "assistant", "content": "x" * 4000})
assert chunk_text(obj, max_chars=3000, structured_chunk_size=5000) == [obj]
def test_rechunk_preserves_one_chunk_id_per_pre_chunk():
"""End-to-end of the producer/consumer chunk_id derivation: each pre-chunk
(one global index) must re-chunk to exactly one piece, so the derived
chunk_ids stay unique within an upsert batch (issue #2301)."""
content = json.dumps([{"role": "assistant", "content": "x" * 6000}])
pre_chunks = chunk_text(content, max_chars=3000, structured_chunk_size=5000)
chunk_ids = []
for global_idx, pre in enumerate(pre_chunks):
for _ in chunk_text(pre, max_chars=3000, structured_chunk_size=5000):
chunk_ids.append(f"bank_doc_{global_idx}")
assert len(chunk_ids) == len(set(chunk_ids)), f"duplicate chunk_ids in one batch: {chunk_ids}"
@@ -1,109 +0,0 @@
"""Tests for ``CODEX_HOME`` resolution of the Codex ``auth.json`` location.
Codex stores its OAuth credentials under a configurable home directory. The
canonical ``@openai/codex`` CLI honors the ``CODEX_HOME`` environment variable
and falls back to ``~/.codex``. Hindsight's Codex auth/LLM/embeddings paths
must resolve the same way so that a user who relocates ``CODEX_HOME`` is still
authenticated.
"""
import json
from pathlib import Path
from hindsight_api.engine.providers.codex_auth import (
CodexAuthManager,
default_codex_auth_file,
)
from hindsight_api.engine.providers.codex_llm import CodexLLM
def _write_auth(auth_dir: Path, access_token: str = "at-test") -> Path:
"""Write a minimal chatgpt-mode auth.json under ``auth_dir``."""
auth_dir.mkdir(parents=True, exist_ok=True)
auth_file = auth_dir / "auth.json"
auth_file.write_text(
json.dumps(
{
"auth_mode": "chatgpt",
"tokens": {
"access_token": access_token,
"refresh_token": "rt-test",
"account_id": "acct-test",
},
}
)
)
return auth_file
# ---------------------------------------------------------------------------
# default_codex_auth_file()
# ---------------------------------------------------------------------------
def test_default_auth_file_falls_back_to_home_codex_when_unset(tmp_path, monkeypatch):
monkeypatch.delenv("CODEX_HOME", raising=False)
monkeypatch.setattr(Path, "home", classmethod(lambda cls: tmp_path))
assert default_codex_auth_file() == tmp_path / ".codex" / "auth.json"
def test_default_auth_file_honors_codex_home_when_set(tmp_path, monkeypatch):
codex_home = tmp_path / "custom-codex"
monkeypatch.setenv("CODEX_HOME", str(codex_home))
assert default_codex_auth_file() == codex_home / "auth.json"
def test_default_auth_file_empty_codex_home_falls_back(tmp_path, monkeypatch):
"""An empty ``CODEX_HOME`` is treated as unset (matches shell semantics)."""
monkeypatch.setenv("CODEX_HOME", "")
monkeypatch.setattr(Path, "home", classmethod(lambda cls: tmp_path))
assert default_codex_auth_file() == tmp_path / ".codex" / "auth.json"
def test_default_auth_file_resolved_lazily(tmp_path, monkeypatch):
"""The env var is read on each call, not cached at import time."""
monkeypatch.setenv("CODEX_HOME", str(tmp_path / "a"))
assert default_codex_auth_file() == tmp_path / "a" / "auth.json"
monkeypatch.setenv("CODEX_HOME", str(tmp_path / "b"))
assert default_codex_auth_file() == tmp_path / "b" / "auth.json"
# ---------------------------------------------------------------------------
# CodexAuthManager.from_file() — honors CODEX_HOME by default
# ---------------------------------------------------------------------------
def test_auth_manager_from_file_uses_codex_home(tmp_path, monkeypatch):
codex_home = tmp_path / "custom-codex"
_write_auth(codex_home, access_token="at-from-codex-home")
monkeypatch.setenv("CODEX_HOME", str(codex_home))
mgr = CodexAuthManager.from_file()
assert mgr.access_token == "at-from-codex-home"
assert mgr._auth_file == codex_home / "auth.json"
# ---------------------------------------------------------------------------
# CodexLLM — loads credentials from CODEX_HOME
# ---------------------------------------------------------------------------
def test_codex_llm_loads_from_codex_home(tmp_path, monkeypatch):
codex_home = tmp_path / "custom-codex"
_write_auth(codex_home, access_token="at-llm")
monkeypatch.setenv("CODEX_HOME", str(codex_home))
llm = CodexLLM(
provider="codex",
api_key="ignored",
base_url="",
model="gpt-5-codex",
)
assert llm.access_token == "at-llm"
assert llm._auth_file == codex_home / "auth.json"
@@ -8,12 +8,9 @@ relevance score, independent of the cross-encoder model's score calibration.
from datetime import datetime, timedelta, timezone
from hindsight_api.engine.search.reranking import (
_RECENCY_ALPHA,
_TEMPORAL_ALPHA,
apply_combined_scoring,
compute_recency_decay,
)
import pytest
from hindsight_api.engine.search.reranking import apply_combined_scoring, _RECENCY_ALPHA, _TEMPORAL_ALPHA
from hindsight_api.engine.search.types import MergedCandidate, RetrievalResult, ScoredResult
UTC = timezone.utc
@@ -203,52 +200,3 @@ class TestBoostFormula:
def test_empty_list_is_noop(self):
apply_combined_scoring([], now=NOW) # must not raise
class TestRecencyDecayFunction:
"""The configurable age→freshness curve (compute_recency_decay)."""
def test_linear_is_default_and_unchanged(self):
"""Default function reproduces the historical linear decay over 365 days."""
assert compute_recency_decay(0) == 1.0
assert abs(compute_recency_decay(182.5) - 0.5) < 1e-6 # neutral at half the window
assert compute_recency_decay(400) == 0.1 # floored past the window
def test_linear_window_is_configurable(self):
"""A custom window moves the neutral crossing; 730d window → neutral at 365d."""
assert abs(compute_recency_decay(365, "linear", linear_window_days=730) - 0.5) < 1e-6
def test_exponential_neutral_at_halflife(self):
"""Exponential decay is exactly neutral (0.5) at the configured half-life."""
assert compute_recency_decay(0, "exponential", halflife_days=90) == 1.0
assert abs(compute_recency_decay(90, "exponential", halflife_days=90) - 0.5) < 1e-9
assert abs(compute_recency_decay(180, "exponential", halflife_days=90) - 0.25) < 1e-9
def test_exponential_penalises_old_less_harshly_than_linear(self):
"""A 1-year-old memory keeps more freshness under a 90d-halflife exponential
than under the linear floor the curve never hard-cuts to 0.1."""
lin = compute_recency_decay(365, "linear")
exp = compute_recency_decay(365, "exponential", halflife_days=180)
assert exp > lin
def test_none_is_always_neutral(self):
"""'none' disables the recency signal — always neutral, no boost."""
assert compute_recency_decay(0, "none") == 0.5
assert compute_recency_decay(10_000, "none") == 0.5
def test_future_dates_clamp_to_max(self):
"""Negative ages (future-dated memories) never exceed full freshness."""
assert compute_recency_decay(-100, "linear") == 1.0
assert compute_recency_decay(-100, "exponential", halflife_days=90) == 1.0
def test_nonpositive_halflife_falls_back_to_neutral(self):
"""A misconfigured (<=0) half-life degrades to neutral rather than dividing by zero."""
assert compute_recency_decay(30, "exponential", halflife_days=0) == 0.5
def test_function_threads_through_apply_combined_scoring(self):
"""The decay function chosen at the call site is what scores sr.recency."""
old = NOW - timedelta(days=180)
sr = _make_result(ce_norm=0.5, occurred_start=old)
apply_combined_scoring([sr], now=NOW, recency_decay_function="none")
assert sr.recency == 0.5
assert abs(sr.weight - 0.5) < 1e-9 # neutral → no recency boost
@@ -25,7 +25,6 @@ def setup_test_env():
"HINDSIGHT_API_LLM_MODEL",
"HINDSIGHT_API_LLM_REASONING_EFFORT",
"HINDSIGHT_API_LLM_BEDROCK_SERVICE_TIER",
"HINDSIGHT_API_LLM_GEMINI_SERVICE_TIER",
"HINDSIGHT_API_SEMANTIC_MIN_SIMILARITY",
"HINDSIGHT_API_DATABASE_URL",
"HINDSIGHT_API_MIGRATION_DATABASE_URL",
@@ -453,53 +452,6 @@ def test_llm_output_language_empty_string_is_unset(monkeypatch):
assert config.llm_output_language is None
def test_markitdown_ocr_defaults_disabled(monkeypatch):
from hindsight_api.config import HindsightConfig
monkeypatch.setenv("HINDSIGHT_API_LLM_PROVIDER", "mock")
config = HindsightConfig.from_env()
assert config.file_parser_markitdown_ocr_enabled is False
def test_markitdown_ocr_does_not_fall_back_to_main_llm_config(monkeypatch):
from hindsight_api.config import DEFAULT_FILE_PARSER_MARKITDOWN_OCR_PROMPT, HindsightConfig
monkeypatch.setenv("HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_ENABLED", "true")
monkeypatch.setenv("HINDSIGHT_API_LLM_PROVIDER", "anthropic")
monkeypatch.setenv("HINDSIGHT_API_LLM_API_KEY", "main-key")
monkeypatch.setenv("HINDSIGHT_API_LLM_BASE_URL", "https://main.example/v1")
monkeypatch.setenv("HINDSIGHT_API_LLM_MODEL", "main-vision-model")
config = HindsightConfig.from_env()
assert config.file_parser_markitdown_ocr_enabled is True
assert config.file_parser_markitdown_ocr_api_key is None
assert config.file_parser_markitdown_ocr_base_url is None
assert config.file_parser_markitdown_ocr_model is None
assert config.file_parser_markitdown_ocr_prompt == DEFAULT_FILE_PARSER_MARKITDOWN_OCR_PROMPT
def test_markitdown_ocr_uses_explicit_config(monkeypatch):
from hindsight_api.config import HindsightConfig
monkeypatch.setenv("HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_ENABLED", "true")
monkeypatch.setenv("HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_API_KEY", "parser-key")
monkeypatch.setenv("HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_BASE_URL", "https://parser.example/v1")
monkeypatch.setenv("HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_MODEL", "parser-vision-model")
monkeypatch.setenv("HINDSIGHT_API_FILE_PARSER_MARKITDOWN_OCR_PROMPT", "Extract this document exactly.")
monkeypatch.setenv("HINDSIGHT_API_LLM_PROVIDER", "mock")
monkeypatch.setenv("HINDSIGHT_API_LLM_API_KEY", "main-key")
monkeypatch.setenv("HINDSIGHT_API_LLM_BASE_URL", "https://main.example/v1")
monkeypatch.setenv("HINDSIGHT_API_LLM_MODEL", "main-vision-model")
config = HindsightConfig.from_env()
assert config.file_parser_markitdown_ocr_enabled is True
assert config.file_parser_markitdown_ocr_api_key == "parser-key"
assert config.file_parser_markitdown_ocr_base_url == "https://parser.example/v1"
assert config.file_parser_markitdown_ocr_model == "parser-vision-model"
assert config.file_parser_markitdown_ocr_prompt == "Extract this document exactly."
def test_llm_reasoning_effort_defaults_to_low(monkeypatch):
from hindsight_api.config import HindsightConfig
@@ -628,81 +580,3 @@ def test_bedrock_service_tier_rejects_invalid_value(monkeypatch):
assert "HINDSIGHT_API_LLM_BEDROCK_SERVICE_TIER" in error_message
assert "standard" in error_message
assert "'standard' is not a valid Bedrock service tier" in error_message
# ---------------------------------------------------------------------------
# Gemini service tier (HINDSIGHT_API_LLM_GEMINI_SERVICE_TIER)
# ---------------------------------------------------------------------------
def test_gemini_service_tier_defaults_to_none(monkeypatch):
"""Gemini service tier defaults to None (standard tier) when unset."""
from hindsight_api.config import HindsightConfig
monkeypatch.delenv("HINDSIGHT_API_LLM_GEMINI_SERVICE_TIER", raising=False)
monkeypatch.setenv("HINDSIGHT_API_LLM_PROVIDER", "mock")
config = HindsightConfig.from_env()
assert config.llm_gemini_service_tier is None
def test_gemini_service_tier_flex(monkeypatch):
"""Flex tier is accepted for Gemini."""
from hindsight_api.config import HindsightConfig
monkeypatch.setenv("HINDSIGHT_API_LLM_GEMINI_SERVICE_TIER", "flex")
monkeypatch.setenv("HINDSIGHT_API_LLM_PROVIDER", "gemini")
monkeypatch.setenv("HINDSIGHT_API_LLM_API_KEY", "fake-key")
config = HindsightConfig.from_env()
assert config.llm_gemini_service_tier == "flex"
def test_gemini_service_tier_accepts_mixed_case_provider(monkeypatch):
"""Gemini tier parsing follows provider's case-insensitive handling."""
from hindsight_api.config import HindsightConfig
monkeypatch.setenv("HINDSIGHT_API_LLM_GEMINI_SERVICE_TIER", "flex")
monkeypatch.setenv("HINDSIGHT_API_LLM_PROVIDER", "Gemini")
monkeypatch.setenv("HINDSIGHT_API_LLM_API_KEY", "fake-key")
config = HindsightConfig.from_env()
assert config.llm_gemini_service_tier == "flex"
def test_gemini_service_tier_rejects_invalid_value(monkeypatch):
"""Unknown Gemini service tiers are rejected early."""
from hindsight_api.config import HindsightConfig
monkeypatch.setenv("HINDSIGHT_API_LLM_GEMINI_SERVICE_TIER", "standard")
monkeypatch.setenv("HINDSIGHT_API_LLM_PROVIDER", "gemini")
monkeypatch.setenv("HINDSIGHT_API_LLM_API_KEY", "fake-key")
with pytest.raises(ValueError) as exc_info:
HindsightConfig.from_env()
error_message = str(exc_info.value)
assert "HINDSIGHT_API_LLM_GEMINI_SERVICE_TIER" in error_message
assert "standard" in error_message
def test_gemini_service_tier_ignored_for_non_gemini_provider(monkeypatch):
"""Invalid Gemini-only tiers do not break unrelated providers."""
from hindsight_api.config import HindsightConfig
monkeypatch.setenv("HINDSIGHT_API_LLM_PROVIDER", "mock")
monkeypatch.setenv("HINDSIGHT_API_LLM_GEMINI_SERVICE_TIER", "standard")
config = HindsightConfig.from_env()
assert config.llm_gemini_service_tier is None
def test_gemini_service_tier_empty_env_is_unset(monkeypatch):
"""Empty env values are treated as unset for templated deployments."""
from hindsight_api.config import HindsightConfig
monkeypatch.setenv("HINDSIGHT_API_LLM_GEMINI_SERVICE_TIER", "")
monkeypatch.setenv("HINDSIGHT_API_LLM_PROVIDER", "mock")
config = HindsightConfig.from_env()
assert config.llm_gemini_service_tier is None
@@ -1,64 +0,0 @@
"""Tests for write-side validation of bank disposition config overrides.
Disposition traits (skepticism / literalism / empathy) are integers on a 1-5
scale. The ``PATCH /v1/{tenant}/banks/{id}/config`` write path must reject
out-of-contract values (floats, 0-1 scales, ints outside 1-5) at write time;
otherwise a single malformed bank 500s the entire bank list because the read
overlay injects the stored value verbatim into a strict
``DispositionTraits(int, ge=1, le=5)``. See issue #2348.
"""
import pytest
from hindsight_api.config_resolver import _validate_disposition_updates
_DISPOSITION_FIELD_NAMES = (
"disposition_skepticism",
"disposition_literalism",
"disposition_empathy",
)
class TestValidateDispositionUpdates:
def test_no_op_passes(self):
_validate_disposition_updates({})
_validate_disposition_updates({"unrelated_field": 123})
def test_valid_in_range_integers_pass(self):
for key in _DISPOSITION_FIELD_NAMES:
for value in (1, 2, 3, 4, 5):
_validate_disposition_updates({key: value})
def test_none_clears_override(self):
# None is the "unset this per-bank override" sentinel (field is int | None).
for key in _DISPOSITION_FIELD_NAMES:
_validate_disposition_updates({key: None})
def test_out_of_range_integer_raises(self):
for key in _DISPOSITION_FIELD_NAMES:
with pytest.raises(ValueError, match=key):
_validate_disposition_updates({key: 0})
with pytest.raises(ValueError, match=key):
_validate_disposition_updates({key: 6})
with pytest.raises(ValueError, match=key):
_validate_disposition_updates({key: -1})
def test_float_raises(self):
# The reported v0.8.3 case: a 0-1 scale used by mistake.
for key in _DISPOSITION_FIELD_NAMES:
with pytest.raises(ValueError, match=key):
_validate_disposition_updates({key: 0.7})
with pytest.raises(ValueError, match=key):
_validate_disposition_updates({key: 3.0}) # float, even if in 1-5 range
def test_bool_raises(self):
# bool is an int subclass and would sneak past a naive isinstance(int) check.
for key in _DISPOSITION_FIELD_NAMES:
with pytest.raises(ValueError, match=key):
_validate_disposition_updates({key: True})
def test_string_raises(self):
for key in _DISPOSITION_FIELD_NAMES:
with pytest.raises(ValueError, match=key):
_validate_disposition_updates({key: "3"})
@@ -124,10 +124,6 @@ def test_openai_codex_provider_uses_codex_oauth_token_and_configured_batch_size(
)
monkeypatch.setenv("HOME", str(tmp_path))
# Codex auth resolves via CODEX_HOME first (falling back to ~/.codex), so a
# CODEX_HOME leaking in from the runner's environment would point auth.json
# away from the tmp_path fixture. Pin resolution to the patched HOME.
monkeypatch.delenv("CODEX_HOME", raising=False)
os.environ["HINDSIGHT_API_LLM_PROVIDER"] = "mock"
os.environ["HINDSIGHT_API_EMBEDDINGS_PROVIDER"] = "openai-codex"
os.environ["HINDSIGHT_API_EMBEDDINGS_OPENAI_MODEL"] = "text-embedding-3-small"
@@ -1,115 +0,0 @@
"""Regression test: `enqueue_graph_maintenance` must insert unit_ids in a
deterministic sorted order so concurrent transactions can't deadlock on the
graph_maintenance_queue unique-key check.
Symptom (production): under load, concurrent `PATCH /memories/{id}` requests
on the same bank generate overlapping `victim_ids` sets (the surviving units
whose outgoing links pointed at the updated unit). Each transaction inserts
those victims into `graph_maintenance_queue` with
`ON CONFLICT (bank_id, unit_id) DO NOTHING`. The conflict check takes a
short-lived row-level lock per (bank_id, unit_id) being inserted, and when
two transactions insert overlapping sets in different orders Postgres
detects a deadlock and aborts one of them surfacing as
`asyncpg.exceptions.DeadlockDetectedError` from the API, which becomes a 500.
The fix sorts the input list inside both `ops_postgresql` and `ops_oracle`
before passing it to the INSERT, so every transaction acquires the per-row
locks in the same global (sorted-UUID) order. With a total order over the
lock set, deadlock is mathematically impossible Postgres still serializes
the conflicting inserts but they queue cleanly instead of cycling.
This test pins that post-condition by capturing the array passed to the
underlying `conn.execute` (PG path) / `conn.executemany` (Oracle path) and
asserting it's sorted.
"""
from __future__ import annotations
import uuid
from unittest.mock import AsyncMock
import pytest
from hindsight_api.engine.db.ops_oracle import OracleOps
from hindsight_api.engine.db.ops_postgresql import PostgreSQLOps
def _shuffled_uuids(n: int) -> list[uuid.UUID]:
"""Generate n UUIDs in a deliberately non-monotonic order. Hex literals
avoid `uuid.uuid4()` because uuid4 is random and we want determinism."""
raw = [
"ffffffff-ffff-4fff-8fff-ffffffffffff",
"00000000-0000-4000-8000-000000000001",
"88888888-8888-4888-8888-888888888888",
"11111111-1111-4111-8111-111111111111",
"ccccccccc-cccc-4ccc-8ccc-cccccccccccc"[:36],
"44444444-4444-4444-8444-444444444444",
]
return [uuid.UUID(s) for s in raw[:n]]
@pytest.mark.asyncio
async def test_pg_enqueue_graph_maintenance_inserts_in_sorted_order():
"""The PostgreSQL ops impl must pass the unit_ids to the INSERT in
sorted order, regardless of how the caller ordered them."""
ops = PostgreSQLOps()
conn = AsyncMock()
unit_ids = _shuffled_uuids(6)
assert unit_ids != sorted(unit_ids), "test inputs must be unsorted"
await ops.enqueue_graph_maintenance(
conn=conn,
table="graph_maintenance_queue",
bank_id="test-bank",
unit_ids=unit_ids,
)
assert conn.execute.await_count == 1
_sql, bank_id_arg, ids_arg = conn.execute.await_args.args
assert bank_id_arg == "test-bank"
assert ids_arg == sorted(unit_ids), f"expected sorted unit_ids for deadlock-free concurrent inserts, got {ids_arg}"
@pytest.mark.asyncio
async def test_oracle_enqueue_graph_maintenance_inserts_in_sorted_order():
"""The Oracle ops impl applies the same sort. `executemany` receives a
list of (bank_id, unit_id) tuples; the unit_id projection must be
sorted."""
ops = OracleOps()
conn = AsyncMock()
unit_ids = _shuffled_uuids(6)
assert unit_ids != sorted(unit_ids), "test inputs must be unsorted"
await ops.enqueue_graph_maintenance(
conn=conn,
table="graph_maintenance_queue",
bank_id="test-bank",
unit_ids=unit_ids,
)
assert conn.executemany.await_count == 1
_sql, rows = conn.executemany.await_args.args
assert [r[0] for r in rows] == ["test-bank"] * len(unit_ids)
assert [r[1] for r in rows] == sorted(unit_ids), (
f"expected sorted unit_ids for deadlock-free concurrent inserts, got {[r[1] for r in rows]}"
)
@pytest.mark.asyncio
async def test_pg_empty_unit_ids_short_circuits():
"""Empty input must remain a no-op — the early return predates this fix
and must continue to skip the INSERT entirely."""
ops = PostgreSQLOps()
conn = AsyncMock()
await ops.enqueue_graph_maintenance(conn, "graph_maintenance_queue", "b", [])
conn.execute.assert_not_awaited()
@pytest.mark.asyncio
async def test_oracle_empty_unit_ids_short_circuits():
ops = OracleOps()
conn = AsyncMock()
await ops.enqueue_graph_maintenance(conn, "graph_maintenance_queue", "b", [])
conn.executemany.assert_not_awaited()
@@ -1,70 +0,0 @@
"""ensure_vector_extension must not create the (unused) global memory_units index.
For per-bank backends (pgvector / pgvectorscale / vchord) every vector search is
bank + fact_type scoped and served by the per-(bank, fact_type) partial indexes
created at bank-creation time. The global `idx_memory_units_embedding` is never
chosen by the planner (migration d5e6f7a8b9c0 drops it for exactly this reason),
so the post-migration reconcile must not recreate it on a fresh schema.
"""
import asyncio
import pytest
from sqlalchemy import create_engine, text
from hindsight_api._vector_index import uses_per_bank_vector_indexes
from hindsight_api.config import HindsightConfig
from hindsight_api.migrations import ensure_vector_extension, run_migrations
@pytest.fixture(scope="module")
def vec_db_url():
"""A dedicated pg0 instance so the test owns its schema/index state."""
from hindsight_api.pg0 import EmbeddedPostgres
pg0 = EmbeddedPostgres(name="hindsight-vecidx-test", port=5570)
loop = asyncio.new_event_loop()
try:
return loop.run_until_complete(pg0.ensure_running())
finally:
loop.close()
def test_per_bank_backend_does_not_create_global_memory_units_index(vec_db_url):
config = HindsightConfig.from_env()
vec = config.vector_extension
if not uses_per_bank_vector_indexes(vec):
pytest.skip(f"backend {vec!r} uses a global vector index by design (no per-bank indexes)")
schema = "vecidx_fresh"
engine = create_engine(vec_db_url)
try:
with engine.connect() as conn:
conn.execute(text(f'DROP SCHEMA IF EXISTS "{schema}" CASCADE'))
conn.commit()
finally:
engine.dispose()
run_migrations(vec_db_url, schema=schema)
# Fresh, empty schema (no banks yet) → the reconcile must be a no-op for the
# global index, not recreate it.
ensure_vector_extension(vec_db_url, vector_extension=vec, schema=schema)
engine = create_engine(vec_db_url)
try:
with engine.connect() as conn:
global_index_count = conn.execute(
text(
"SELECT COUNT(*) FROM pg_indexes "
"WHERE schemaname = :schema AND tablename = 'memory_units' "
"AND indexname = 'idx_memory_units_embedding'"
),
{"schema": schema},
).scalar()
conn.execute(text(f'DROP SCHEMA IF EXISTS "{schema}" CASCADE'))
conn.commit()
finally:
engine.dispose()
assert global_index_count == 0
+1 -95
View File
@@ -911,7 +911,7 @@ class TestPrecheckHttpWiring:
def _build_app(validator):
"""Mirror the precheck wiring from ``hindsight_api.api.http`` in a
standalone FastAPI app."""
from fastapi import Depends, FastAPI, HTTPException, Request
from fastapi import Depends, FastAPI, HTTPException
from pydantic import BaseModel, model_validator
from hindsight_api.extensions import PrecheckContext
@@ -952,23 +952,12 @@ class TestPrecheckHttpWiring:
def _precheck_for(operation: str):
async def _dep(
bank_id: str,
request: Request,
request_context: RequestContext = Depends(_request_context),
) -> None:
cl_header = request.headers.get("content-length")
content_length: int | None = None
if cl_header is not None:
try:
parsed = int(cl_header)
except ValueError:
parsed = -1
if parsed >= 0:
content_length = parsed
ctx = PrecheckContext(
operation=operation,
bank_id=bank_id,
request_context=request_context,
content_length=content_length,
)
result = await validator.precheck(ctx)
if not result.allowed:
@@ -1092,86 +1081,3 @@ class TestPrecheckHttpWiring:
resp = client.get("/v1/default/banks/precheck-bank/memories/list")
assert resp.status_code == 200
assert len(validator.precheck_calls) == 0
def test_precheck_context_carries_content_length(self):
"""Content-Length header is exposed to the precheck so a validator
can make size-aware decisions (e.g. upper-bound cost estimate)
before the body is deserialised."""
validator = RecordingPrecheckValidator(reject=False)
app, _ = self._build_app(validator)
client = TestClient(app)
# Body must contain at least 500 'x' bytes; check the surfaced
# Content-Length is within a tight band around that floor (allows
# for JSON envelope + httpx's serialisation choices without
# depending on exact byte counts).
payload = {"items": [{"content": "x" * 500}]}
resp = client.post(
"/v1/default/banks/precheck-bank/memories",
json=payload,
)
assert resp.status_code == 200
assert len(validator.precheck_calls) == 1
ctx = validator.precheck_calls[0]
assert ctx.content_length is not None
assert 500 <= ctx.content_length <= 600
def test_precheck_context_content_length_zero_is_not_none(self):
"""An empty POST body has Content-Length: 0. That should surface
as the int 0, not None None means 'unknown', 0 means 'known to
be empty'."""
validator = RecordingPrecheckValidator(reject=False)
app, _ = self._build_app(validator)
client = TestClient(app)
# Empty body fails Pydantic parse (422), but precheck runs first
# and records the Content-Length.
client.post(
"/v1/default/banks/precheck-bank/memories",
content=b"",
headers={"content-type": "application/json"},
)
assert len(validator.precheck_calls) >= 1
ctx = validator.precheck_calls[-1]
assert ctx.content_length == 0
@pytest.mark.asyncio
async def test_precheck_context_content_length_none_when_header_missing(self):
"""When the Content-Length header isn't set (e.g. chunked transfer
encoding) the validator sees None, not a crash and not a default 0."""
from starlette.requests import Request as _StarletteRequest
from hindsight_api.extensions import PrecheckContext
from hindsight_api.models import RequestContext
validator = RecordingPrecheckValidator(reject=False)
# Replicate the wiring's parse step inline so the test exercises
# the same code-path semantics introduced in
# ``hindsight_api.api.http._precheck_dep``.
scope = {
"type": "http",
"method": "POST",
"path": "/v1/default/banks/bank-x/memories",
"headers": [], # no content-length
"query_string": b"",
}
req = _StarletteRequest(scope)
cl_header = req.headers.get("content-length")
content_length: int | None = None
if cl_header is not None:
try:
parsed = int(cl_header)
except ValueError:
parsed = -1
if parsed >= 0:
content_length = parsed
ctx = PrecheckContext(
operation="retain",
bank_id="bank-x",
request_context=RequestContext(),
content_length=content_length,
)
await validator.precheck(ctx)
assert validator.precheck_calls[-1].content_length is None
@@ -77,27 +77,6 @@ async def test_dry_run_extracts_without_persisting(api_client, memory):
assert after["total"] == before["total"]
@pytest.mark.asyncio
async def test_dry_run_rejects_empty_content(api_client, memory):
"""Empty/whitespace-only content is rejected by request validation (422) before the
billable LLM extraction call runs matching retain (RetainItem.content) and recall
(RecallRequest.query), which already reject empty input."""
bank_id = f"dryrun-{uuid.uuid4().hex[:8]}"
await memory.get_bank_profile(bank_id=bank_id, request_context=RequestContext())
before = await memory.list_memory_units(bank_id=bank_id, request_context=RequestContext())
for content in ("", " ", "\n\t "):
resp = await api_client.post(
f"/v1/default/banks/{bank_id}/memories/dry-run-extract",
json={"content": content},
)
assert resp.status_code == 422, resp.text
# Rejected before extraction: nothing was persisted.
after = await memory.list_memory_units(bank_id=bank_id, request_context=RequestContext())
assert after["total"] == before["total"]
@pytest.mark.asyncio
async def test_dry_run_disabled_returns_404(api_client, memory):
"""With HINDSIGHT_API_ENABLE_DRY_RUN_EXTRACT=false the endpoint is removed (returns 404)."""
@@ -346,149 +346,6 @@ async def test_markitdown_converter():
assert "test document" in result.lower() or "multiple lines" in result.lower()
def test_markitdown_converter_does_not_enable_ocr_by_default(monkeypatch):
"""Markitdown should keep its local/default behavior unless OCR is explicitly enabled."""
import markitdown
from hindsight_api.engine.parsers import MarkitdownParser
calls = []
class FakeMarkItDown:
def __init__(self, **kwargs):
calls.append(kwargs)
monkeypatch.setattr(markitdown, "MarkItDown", FakeMarkItDown)
MarkitdownParser()
assert calls == [{}]
@pytest.mark.asyncio
async def test_markitdown_image_without_ocr_has_actionable_error(monkeypatch):
"""Image uploads should explain that MarkItDown OCR is disabled instead of surfacing a low-level error."""
import markitdown
from hindsight_api.engine.parsers import MarkitdownParser
class FakeMarkItDown:
def __init__(self, **kwargs):
pass
def convert(self, path):
raise AssertionError("MarkItDown should not be called when image OCR is disabled")
monkeypatch.setattr(markitdown, "MarkItDown", FakeMarkItDown)
parser = MarkitdownParser()
with pytest.raises(RuntimeError, match="Image OCR is not enabled for the markitdown parser"):
await parser.convert(b"\x89PNG\r\n\x1a\n", "screenshot.png")
def test_markitdown_converter_can_enable_ocr(monkeypatch):
"""When enabled, Markitdown receives an OpenAI-compatible client, model, and OCR prompt."""
import markitdown
import openai
from hindsight_api.config import DEFAULT_FILE_PARSER_MARKITDOWN_OCR_PROMPT
from hindsight_api.engine.parsers import MarkitdownParser
markitdown_calls = []
openai_calls = []
class FakeMarkItDown:
def __init__(self, **kwargs):
markitdown_calls.append(kwargs)
class FakeOpenAI:
def __init__(self, **kwargs):
openai_calls.append(kwargs)
monkeypatch.delenv("OPENAI_API_KEY", raising=False)
monkeypatch.setattr(markitdown, "MarkItDown", FakeMarkItDown)
monkeypatch.setattr(openai, "OpenAI", FakeOpenAI)
MarkitdownParser(
ocr_enabled=True,
ocr_api_key="parser-key",
ocr_base_url="https://vision.example/v1",
ocr_model="vision-model",
)
assert openai_calls == [
{
"api_key": "parser-key",
"base_url": "https://vision.example/v1",
}
]
assert markitdown_calls[0]["llm_client"].__class__ is FakeOpenAI
assert markitdown_calls[0]["llm_model"] == "vision-model"
assert markitdown_calls[0]["llm_prompt"] == DEFAULT_FILE_PARSER_MARKITDOWN_OCR_PROMPT
def test_markitdown_converter_requires_model_when_ocr_enabled(monkeypatch):
"""OCR should fail fast when enabled without a model."""
import markitdown
from hindsight_api.engine.parsers import MarkitdownParser
class FakeMarkItDown:
def __init__(self, **kwargs):
pass
monkeypatch.setattr(markitdown, "MarkItDown", FakeMarkItDown)
with pytest.raises(ValueError, match="no model"):
MarkitdownParser(ocr_enabled=True, ocr_api_key="parser-key")
def test_markitdown_converter_requires_base_url_when_ocr_enabled(monkeypatch):
"""OCR should fail fast when enabled without a dedicated OpenAI-compatible endpoint."""
import markitdown
from hindsight_api.engine.parsers import MarkitdownParser
class FakeMarkItDown:
def __init__(self, **kwargs):
pass
monkeypatch.setattr(markitdown, "MarkItDown", FakeMarkItDown)
with pytest.raises(ValueError, match="no base URL"):
MarkitdownParser(ocr_enabled=True, ocr_api_key="parser-key", ocr_model="vision-model")
def test_markitdown_converter_reports_missing_openai_when_ocr_enabled(monkeypatch):
"""Missing OpenAI SDK should not be reported as missing MarkItDown."""
import builtins
import markitdown
from hindsight_api.engine.parsers import MarkitdownParser
real_import = builtins.__import__
class FakeMarkItDown:
def __init__(self, **kwargs):
pass
def fake_import(name, globals=None, locals=None, fromlist=(), level=0):
if name == "openai":
raise ImportError("no openai")
return real_import(name, globals, locals, fromlist, level)
monkeypatch.setattr(markitdown, "MarkItDown", FakeMarkItDown)
monkeypatch.setattr(builtins, "__import__", fake_import)
with pytest.raises(RuntimeError, match="openai package is required"):
MarkitdownParser(
ocr_enabled=True,
ocr_api_key="parser-key",
ocr_base_url="https://vision.example/v1",
ocr_model="vision-model",
)
@pytest.mark.asyncio
async def test_converter_registry():
"""Test file parser registry."""
@@ -12,8 +12,6 @@ import subprocess
import tempfile
import time
import uuid
from collections.abc import Iterator
from contextlib import contextmanager
import httpx
import pytest
@@ -23,7 +21,6 @@ logger = logging.getLogger(__name__)
try:
from testcontainers.core.container import DockerContainer
from testcontainers.core.docker_client import DockerClient as _DockerClient
_has_testcontainers = True
except ImportError:
@@ -41,8 +38,6 @@ SEAWEEDFS_S3_PORT = 8333
TEST_BUCKET = "hindsight-test"
ACCESS_KEY = "test_access_key"
SECRET_KEY = "test_secret_key"
_PORT_MAPPING_RETRY_TIMEOUT_SECONDS = 10.0
_PORT_MAPPING_RETRY_INTERVAL_SECONDS = 0.1
# SeaweedFS S3 IAM config granting full access to our test credentials
_S3_CONFIG = {
@@ -69,33 +64,6 @@ def _docker_available() -> bool:
return False
if _has_testcontainers:
@contextmanager
def _retry_testcontainers_port_mapping() -> Iterator[None]:
original_port = _DockerClient.port
def port_with_retry(self: _DockerClient, container_id: str, port: int) -> str:
deadline = time.monotonic() + _PORT_MAPPING_RETRY_TIMEOUT_SECONDS
while True:
try:
return original_port(self, container_id, port)
except ConnectionError:
# Docker Desktop can report a container as running before its
# published port appears in NetworkSettings.Ports. This affects
# both Ryuk's 8080 lookup inside testcontainers and the
# SeaweedFS S3 port lookup below.
if time.monotonic() >= deadline:
raise
time.sleep(_PORT_MAPPING_RETRY_INTERVAL_SECONDS)
_DockerClient.port = port_with_retry
try:
yield
finally:
_DockerClient.port = original_port
def _wait_for_seaweedfs(endpoint: str, timeout: int = 30) -> None:
"""Poll SeaweedFS S3 endpoint until ready."""
deadline = time.time() + timeout
@@ -133,11 +101,11 @@ def seaweedfs_container():
.with_command(f"server -s3 -s3.port={SEAWEEDFS_S3_PORT} -s3.config=/etc/seaweedfs/s3.json -ip.bind=0.0.0.0")
)
container.start()
try:
with _retry_testcontainers_port_mapping():
container.start()
host = container.get_container_host_ip()
port = container.get_exposed_port(SEAWEEDFS_S3_PORT)
host = container.get_container_host_ip()
port = container.get_exposed_port(SEAWEEDFS_S3_PORT)
endpoint = f"http://{host}:{port}"
_wait_for_seaweedfs(endpoint, timeout=240)
@@ -1,103 +0,0 @@
"""Plumbing tests for the Gemini service tier flag."""
from unittest.mock import MagicMock, patch
import pytest
from hindsight_api.engine.llm_wrapper import LLMConfig
def test_llm_config_threads_gemini_service_tier_to_provider_impl():
"""End-to-end: LLMConfig -> create_llm_provider -> GeminiLLM carries the tier."""
pytest.importorskip("google.genai")
with patch("google.genai.Client", return_value=MagicMock()):
llm = LLMConfig(
provider="gemini",
api_key="fake-key",
base_url="",
model="gemini-2.5-flash",
gemini_service_tier="flex",
)
assert llm._provider_impl._service_tier == "flex"
def test_llm_provider_from_env_validates_gemini_service_tier(monkeypatch):
"""Direct env construction rejects the same invalid tiers as HindsightConfig."""
from hindsight_api.config import clear_config_cache
from hindsight_api.engine.llm_wrapper import LLMProvider
monkeypatch.setenv("HINDSIGHT_API_LLM_PROVIDER", "gemini")
monkeypatch.setenv("HINDSIGHT_API_LLM_API_KEY", "fake-key")
monkeypatch.setenv("HINDSIGHT_API_LLM_GEMINI_SERVICE_TIER", "standard")
clear_config_cache()
with pytest.raises(ValueError, match="HINDSIGHT_API_LLM_GEMINI_SERVICE_TIER"):
LLMProvider.from_env()
clear_config_cache()
def test_llm_provider_from_env_ignores_gemini_tier_for_non_gemini(monkeypatch):
"""Invalid Gemini-only tier env values do not break other providers."""
from hindsight_api.config import clear_config_cache
from hindsight_api.engine.llm_wrapper import LLMProvider
monkeypatch.setenv("HINDSIGHT_API_LLM_PROVIDER", "mock")
monkeypatch.setenv("HINDSIGHT_API_LLM_GEMINI_SERVICE_TIER", "standard")
clear_config_cache()
provider = LLMProvider.from_env()
assert provider.gemini_service_tier is None
clear_config_cache()
def test_llm_provider_from_env_keeps_lightweight_loader(monkeypatch):
"""Reading the Gemini tier must not construct the full application config."""
from hindsight_api.config import clear_config_cache
from hindsight_api.engine.llm_wrapper import LLMProvider
monkeypatch.setenv("HINDSIGHT_API_LLM_PROVIDER", "gemini")
monkeypatch.setenv("HINDSIGHT_API_LLM_API_KEY", "fake-key")
monkeypatch.setenv("HINDSIGHT_API_LLM_GEMINI_SERVICE_TIER", "flex")
monkeypatch.setenv("HINDSIGHT_API_RETAIN_MAX_COMPLETION_TOKENS", "1000")
monkeypatch.setenv("HINDSIGHT_API_RETAIN_CHUNK_SIZE", "2000")
clear_config_cache()
with patch("google.genai.Client", return_value=MagicMock()):
provider = LLMProvider.from_env()
assert provider.gemini_service_tier == "flex"
clear_config_cache()
def test_llm_provider_constructor_validates_gemini_service_tier():
"""Direct Gemini construction rejects invalid tiers before API calls."""
from hindsight_api.engine.llm_wrapper import LLMProvider
with pytest.raises(ValueError, match="HINDSIGHT_API_LLM_GEMINI_SERVICE_TIER"):
LLMProvider(
provider="gemini",
api_key="fake-key",
base_url="",
model="gemini-2.5-flash",
gemini_service_tier="standard",
)
def test_vertexai_ignores_gemini_service_tier():
"""The Gemini-only tier flag is not forwarded to Vertex AI providers."""
from hindsight_api.engine.llm_wrapper import create_llm_provider
with patch("hindsight_api.engine.providers.GeminiLLM") as mock_gemini:
create_llm_provider(
provider="vertexai",
api_key="",
base_url="",
model="gemini-2.5-flash",
reasoning_effort="low",
gemini_service_tier="flex",
)
assert mock_gemini.call_args.kwargs["gemini_service_tier"] is None
@@ -1,119 +0,0 @@
"""Reproduces the concurrent-insert deadlock on ``graph_maintenance_queue``
that PR #2353 targets, and demonstrates that a shared insertion order cures it.
A deadlock is a *database*-level phenomenon, so unlike the PR's own tests (which
only assert that the Python list handed to ``conn.execute`` is sorted) these run
against the real Postgres test DB and drive two genuinely-concurrent
transactions, forcing the exact interleaving that produces a lock cycle.
Modelling note
--------------
Production enqueues a whole victim set in ONE statement::
INSERT INTO graph_maintenance_queue (bank_id, unit_id)
SELECT $1, v FROM unnest($2::uuid[]) ON CONFLICT (bank_id, unit_id) DO NOTHING
That single statement still takes the per-row unique-key locks one row at a time,
in the order ``unnest`` yields we just can't pause *inside* a single statement.
So each worker here issues the rows one at a time with a barrier between them.
That makes the otherwise-racy interleaving deterministic while exercising the
identical lock: ``ON CONFLICT`` on the ``(bank_id, unit_id)`` primary key.
"""
from __future__ import annotations
import asyncio
import uuid
import pytest
from asyncpg.exceptions import DeadlockDetectedError
from hindsight_api.engine.memory_engine import MemoryEngine
# Two keys with an unambiguous sort order (low < high as UUIDs / as text).
K_LOW = uuid.UUID("00000000-0000-4000-8000-000000000001")
K_HIGH = uuid.UUID("ffffffff-ffff-4fff-8fff-ffffffffffff")
async def _insert_one(conn, bank_id: str, unit_id: uuid.UUID) -> None:
"""One row of the production INSERT ... ON CONFLICT DO NOTHING."""
await conn.execute(
"""
INSERT INTO graph_maintenance_queue (bank_id, unit_id)
VALUES ($1, $2)
ON CONFLICT (bank_id, unit_id) DO NOTHING
""",
bank_id,
unit_id,
)
@pytest.mark.asyncio
async def test_unordered_concurrent_enqueue_deadlocks(memory: MemoryEngine):
"""Two transactions inserting the same two keys in OPPOSITE orders deadlock.
This is the pre-fix reality: ``enqueue_relink_victims`` feeds whatever order
``SELECT DISTINCT`` returns, so two overlapping victim sets can acquire the
unique-key locks in opposite orders and cycle. Postgres aborts one with
``DeadlockDetectedError``, which the API surfaces as a 500.
"""
pool = await memory._get_pool()
bank_id = f"dl-bug-{uuid.uuid4().hex[:8]}"
# Both transactions hold their first lock before either takes its second,
# so the cross-wait (and thus the cycle) is guaranteed rather than racy.
barrier = asyncio.Barrier(2)
async def worker(order: list[uuid.UUID]) -> None:
async with pool.acquire() as conn:
async with conn.transaction():
await _insert_one(conn, bank_id, order[0])
await barrier.wait()
await _insert_one(conn, bank_id, order[1])
results = await asyncio.wait_for(
asyncio.gather(
worker([K_LOW, K_HIGH]),
worker([K_HIGH, K_LOW]),
return_exceptions=True,
),
timeout=30,
)
deadlocks = [r for r in results if isinstance(r, DeadlockDetectedError)]
assert deadlocks, f"expected one transaction aborted with DeadlockDetectedError, got {results!r}"
@pytest.mark.asyncio
async def test_ordered_concurrent_enqueue_does_not_deadlock(memory: MemoryEngine):
"""With both transactions inserting in the SAME (sorted) order — exactly what
PR #2353's ``sorted(unit_ids)`` guarantees per call — there is no cycle. The
second transaction simply waits on the first shared key and proceeds once the
first commits; both victim sets land in the queue.
"""
pool = await memory._get_pool()
bank_id = f"dl-fix-{uuid.uuid4().hex[:8]}"
order = sorted([K_LOW, K_HIGH]) # identical order for both workers
async def worker() -> None:
async with pool.acquire() as conn:
async with conn.transaction():
for uid in order:
await _insert_one(conn, bank_id, uid)
# Sorted order cannot cycle; the timeout only guards against an unexpected hang.
results = await asyncio.wait_for(
asyncio.gather(worker(), worker(), return_exceptions=True),
timeout=30,
)
errors = [r for r in results if isinstance(r, BaseException)]
assert not errors, f"sorted concurrent inserts must not deadlock, got {results!r}"
async with pool.acquire() as conn:
rows = await conn.fetch(
"SELECT unit_id FROM graph_maintenance_queue WHERE bank_id = $1 ORDER BY unit_id",
bank_id,
)
assert [r["unit_id"] for r in rows] == order
@@ -5,7 +5,6 @@ import pytest
from datetime import datetime, timezone, timedelta
from unittest.mock import AsyncMock, MagicMock
from hindsight_api.config import clear_config_cache
from hindsight_api.engine.retain.link_utils import (
_normalize_datetime,
_cap_links_per_unit,
@@ -409,19 +408,6 @@ class TestComputeSemanticLinksAnnPgBouncerSafety:
following the CREATE TEMP TABLE.
"""
@pytest.fixture(autouse=True)
def _reset_config_cache(self):
# Tests below monkeypatch HINDSIGHT_API_VECTOR_EXTENSION. The ANN code
# path reads it through the process-global config cache, and monkeypatch
# reverts only the env var — not the cache. Left uncleared, a leaked
# "vchord" makes every later bank-creating test on the same xdist worker
# emit `USING vchordrq` against the pgvector-only test DB and fail with
# `access method "vchordrq" does not exist`. Clear before and after so
# the cache is rebuilt from the current env for each test.
clear_config_cache()
yield
clear_config_cache()
@pytest.fixture
def mock_conn(self):
"""An asyncpg-like connection mock with an async `transaction()`
@@ -1,71 +0,0 @@
"""Regression test: list_banks must apply the same disposition + mission
config overlay that get_bank_profile applies.
Bug (reproduced live against 0.8.1): for a bank whose disposition and
mission were evolved/overridden via bank *config* (the banks.config JSONB:
reflect_mission, disposition_skepticism/literalism/empathy), the single-bank
get path returns the real values while the list path returns the stale legacy
DB-column defaults ({skepticism:3, literalism:3, empathy:3} and "").
Root cause: MemoryEngine.get_bank_profile overlays the resolved bank config
on top of the legacy banks.disposition/banks.mission columns, but
MemoryEngine.list_banks returned bank_utils.list_banks rows straight from
those columns with no overlay. The two endpoints disagreed for the same bank.
This test sets disposition + mission through the config path (so the legacy
columns keep their defaults) and asserts list_banks agrees with
get_bank_profile for that bank.
Runs via: uv run pytest tests/test_list_banks_config_overlay.py -v
"""
from __future__ import annotations
import pytest
from hindsight_api.models import RequestContext
@pytest.mark.asyncio
async def test_list_banks_overlays_config_disposition_and_mission(memory):
bank_id = "list_banks_config_overlay_bank"
request_context = RequestContext(api_key=None, api_key_id=None, tenant_id=None, internal=False)
# Values that differ from the 3/3/3 defaults on every trait, and a
# clearly non-empty mission, so a stale-default regression is unmissable.
overrides = {
"reflect_mission": "I am the shared long-term memory for this regression test.",
"disposition_skepticism": 4,
"disposition_literalism": 5,
"disposition_empathy": 2,
}
try:
# Create the bank. Its legacy banks.disposition/banks.mission columns
# keep their defaults (3/3/3 and "") — the real values live in config.
await memory.get_bank_profile(bank_id, request_context=request_context)
# Set disposition + mission via the *config* path (banks.config JSONB),
# exactly the path that triggered the live bug.
await memory._config_resolver.update_bank_config(bank_id, overrides, request_context)
# Source of truth: the single-bank get path already overlays config.
profile = await memory.get_bank_profile(bank_id, request_context=request_context)
assert profile["mission"] == overrides["reflect_mission"]
assert profile["disposition"] == {"skepticism": 4, "literalism": 5, "empathy": 2}
# The list path must agree with the get path for this bank.
banks = await memory.list_banks(request_context=request_context)
entry = next((b for b in banks if b["bank_id"] == bank_id), None)
assert entry is not None, f"bank {bank_id!r} not present in list_banks output"
assert entry["mission"] == profile["mission"], (
"list_banks returned a different mission than get_bank_profile: "
f"list={entry['mission']!r} get={profile['mission']!r}"
)
assert entry["disposition"] == profile["disposition"], (
"list_banks returned a different disposition than get_bank_profile: "
f"list={entry['disposition']!r} get={profile['disposition']!r}"
)
finally:
await memory.delete_bank(bank_id, request_context=request_context)
@@ -1,80 +0,0 @@
"""Regression: user-facing GET list endpoints must reject negative limit/offset
with a clean 422 at the FastAPI boundary instead of letting the value reach
Postgres (``LIMIT/OFFSET must not be negative``) and surfacing as an opaque 500
that also leaks the raw Postgres error string.
This makes pagination validation consistent with the sibling list endpoints in
the same router (document-chunks / directives / async-ops / audit) that already
declare ``Query(..., ge=...)``. The engine emits ``LIMIT $n OFFSET $n`` with no
``max(0, ...)`` clamp, so the guard has to live at the request boundary.
"""
import uuid
import httpx
import pytest
import pytest_asyncio
from hindsight_api import RequestContext
from hindsight_api.api import create_app
@pytest_asyncio.fixture
async def api_client(memory):
app = create_app(memory, initialize_memory=False)
transport = httpx.ASGITransport(app=app)
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
yield client
def _url(bank_id: str, suffix: str) -> str:
return f"/v1/default/banks/{bank_id}/{suffix}"
# Endpoints that accept a ``limit`` query param.
LIMIT_ENDPOINTS = [
"graph",
"memories/list",
"entities",
"entities/graph",
"documents",
"tags",
]
# Subset that also accept an ``offset`` query param.
OFFSET_ENDPOINTS = [
"memories/list",
"entities",
"documents",
"tags",
]
@pytest.mark.asyncio
@pytest.mark.parametrize("suffix", LIMIT_ENDPOINTS)
async def test_negative_limit_returns_422_not_500(api_client, suffix):
bank_id = f"pag-{uuid.uuid4().hex[:8]}"
resp = await api_client.get(_url(bank_id, suffix), params={"limit": -1})
# FastAPI validation runs before the handler / DB, so a bad pagination input
# is a clean 422 — never a 500 leaking the raw Postgres error.
assert resp.status_code == 422, resp.text
@pytest.mark.asyncio
@pytest.mark.parametrize("suffix", OFFSET_ENDPOINTS)
async def test_negative_offset_returns_422_not_500(api_client, suffix):
bank_id = f"pag-{uuid.uuid4().hex[:8]}"
resp = await api_client.get(_url(bank_id, suffix), params={"offset": -1})
assert resp.status_code == 422, resp.text
@pytest.mark.asyncio
@pytest.mark.parametrize("limit", [0, 1, 100])
async def test_valid_limit_is_accepted_including_zero(api_client, memory, limit):
# Positive control: ge=0 rejects only NEGATIVE limits. A non-negative limit
# — including limit=0 (a valid empty page, LIMIT 0) — must still be accepted,
# so this fix does not change behavior for any previously-valid input.
bank_id = f"pag-{uuid.uuid4().hex[:8]}"
await memory.get_bank_profile(bank_id=bank_id, request_context=RequestContext())
resp = await api_client.get(_url(bank_id, "memories/list"), params={"limit": limit, "offset": 0})
assert resp.status_code == 200, resp.text
@@ -1,82 +0,0 @@
"""
Regression test for the hard timeout on the LiteLLM provider.
A completion that never returns a connection held open with no token
progress, or one straggler inside a concurrent ``asyncio.gather`` fan-out
must not block forever. ``call`` / ``call_with_tools`` wrap the request in
``asyncio.wait_for`` so it is cancelled after ``timeout`` seconds and surfaced
as a retryable ``TimeoutError`` instead of pinning a worker slot and a
concurrency permit indefinitely.
"""
import asyncio
import time
import pytest
from hindsight_api.config import DEFAULT_LLM_TIMEOUT, ENV_LLM_TIMEOUT
from hindsight_api.engine.providers.litellm_llm import LiteLLMLLM
def _make_provider(timeout: float | None) -> LiteLLMLLM:
return LiteLLMLLM(
provider="litellm",
api_key="unused",
base_url="http://localhost:0/v1",
model="litellm_proxy/test-model",
timeout=timeout,
)
async def test_call_cancels_hung_completion(monkeypatch):
"""A hung ``_acompletion`` is cancelled per attempt and raises TimeoutError."""
provider = _make_provider(timeout=0.1)
calls = 0
async def _hang(**kwargs):
nonlocal calls
calls += 1
await asyncio.Event().wait() # never resolves
monkeypatch.setattr(provider, "_acompletion", _hang)
started = time.monotonic()
with pytest.raises((TimeoutError, asyncio.TimeoutError)):
await provider.call(
messages=[{"role": "user", "content": "hi"}],
max_retries=1,
initial_backoff=0.01,
max_backoff=0.01,
)
elapsed = time.monotonic() - started
# max_retries=1 -> attempts 0 and 1, each bounded by the timeout.
assert calls == 2
# Bounded by ~2 * timeout + backoff — nowhere near hanging forever.
assert elapsed < 2.0
async def test_call_with_tools_cancels_hung_completion(monkeypatch):
provider = _make_provider(timeout=0.1)
async def _hang(**kwargs):
await asyncio.Event().wait()
monkeypatch.setattr(provider, "_acompletion", _hang)
with pytest.raises((TimeoutError, asyncio.TimeoutError)):
await provider.call_with_tools(
messages=[{"role": "user", "content": "hi"}],
tools=[],
max_retries=0,
initial_backoff=0.01,
max_backoff=0.01,
)
async def test_unset_timeout_falls_back_to_default(monkeypatch):
"""``None`` must resolve to a finite default — never ``None``, which would
make ``asyncio.wait_for`` wait forever and reintroduce the hang."""
monkeypatch.delenv(ENV_LLM_TIMEOUT, raising=False)
provider = _make_provider(timeout=None)
assert provider.timeout == DEFAULT_LLM_TIMEOUT
+1 -170
View File
@@ -120,7 +120,7 @@ async def test_anthropic_no_extra_body_omits_key():
# ─── Gemini ───────────────────────────────────────────────────────────────────
def _make_gemini_provider(extra_body=None, gemini_service_tier=None):
def _make_gemini_provider(extra_body=None):
pytest.importorskip("google.genai")
with patch("google.genai.Client") as mock_client_cls:
mock_client_cls.return_value = MagicMock()
@@ -132,7 +132,6 @@ def _make_gemini_provider(extra_body=None, gemini_service_tier=None):
base_url="",
model="gemini-2.5-flash",
extra_body=extra_body,
gemini_service_tier=gemini_service_tier,
)
provider._client = MagicMock()
return provider
@@ -177,174 +176,6 @@ async def test_gemini_explicit_temperature_overrides_extra_body():
assert config_arg.temperature == 0.9
@pytest.mark.asyncio
async def test_gemini_service_tier_applies_to_http_options_extra_body():
"""The native Gemini service tier flag reaches GenerateContentConfig."""
provider = _make_gemini_provider(gemini_service_tier="flex")
provider._client.aio.models.generate_content = AsyncMock(return_value=_fake_gemini_response())
await provider.call(messages=[{"role": "user", "content": "hi"}], scope="test")
config_arg = provider._client.aio.models.generate_content.call_args.kwargs.get("config")
assert config_arg.http_options.extra_body["service_tier"] == "flex"
@pytest.mark.asyncio
async def test_gemini_extra_body_service_tier_takes_precedence():
"""The explicit extra_body escape hatch wins over the native flag."""
provider = _make_gemini_provider(
extra_body={"http_options": {"extra_body": {"service_tier": "standard"}}},
gemini_service_tier="flex",
)
provider._client.aio.models.generate_content = AsyncMock(return_value=_fake_gemini_response())
await provider.call(messages=[{"role": "user", "content": "hi"}], scope="test")
config_arg = provider._client.aio.models.generate_content.call_args.kwargs.get("config")
assert config_arg.http_options.extra_body["service_tier"] == "standard"
assert provider._extra_body["http_options"]["extra_body"]["service_tier"] == "standard"
@pytest.mark.asyncio
async def test_gemini_structured_call_uses_native_schema_without_prompt_duplicate():
"""Structured Gemini calls send schema through response_schema only."""
from pydantic import BaseModel
class StructuredAnswer(BaseModel):
answer: str
provider = _make_gemini_provider()
response = _fake_gemini_response()
response.text = '{"answer": "ok"}'
provider._client.aio.models.generate_content = AsyncMock(return_value=response)
result = await provider.call(
messages=[
{"role": "system", "content": "Return concise JSON."},
{"role": "user", "content": "hello"},
],
response_format=StructuredAnswer,
scope="test",
)
config_arg = provider._client.aio.models.generate_content.call_args.kwargs.get("config")
assert result.answer == "ok"
assert config_arg.response_mime_type == "application/json"
assert config_arg.response_schema is StructuredAnswer
assert config_arg.system_instruction == "Return concise JSON."
assert "valid JSON matching this schema" not in config_arg.system_instruction
@pytest.mark.asyncio
async def test_gemini_cached_structured_call_keeps_native_schema():
"""Cached Gemini calls still send response_schema per request."""
from pydantic import BaseModel
class StructuredAnswer(BaseModel):
answer: str
provider = _make_gemini_provider()
response = _fake_gemini_response()
response.text = '{"answer": "ok"}'
provider._client.aio.models.generate_content = AsyncMock(return_value=response)
result = await provider.call(
messages=[
{"role": "system", "content": "Return concise JSON."},
{"role": "user", "content": "hello"},
],
response_format=StructuredAnswer,
cached_prefix="cachedContents/test",
scope="test",
)
config_arg = provider._client.aio.models.generate_content.call_args.kwargs.get("config")
assert result.answer == "ok"
assert config_arg.cached_content == "cachedContents/test"
assert config_arg.system_instruction is None
assert config_arg.response_mime_type == "application/json"
assert config_arg.response_schema is StructuredAnswer
@pytest.mark.asyncio
async def test_gemini_structured_parse_failure_falls_back_to_prompt_schema():
"""Malformed native-schema output gets one prompt-schema compatibility retry."""
from pydantic import BaseModel
class StructuredAnswer(BaseModel):
answer: str
provider = _make_gemini_provider()
invalid = _fake_gemini_response()
invalid.text = "not json"
valid = _fake_gemini_response()
valid.text = '{"answer": "ok"}'
provider._client.aio.models.generate_content = AsyncMock(side_effect=[invalid, valid])
result = await provider.call(
messages=[
{"role": "system", "content": "Return concise JSON."},
{"role": "user", "content": "hello"},
],
response_format=StructuredAnswer,
scope="test",
max_retries=1,
initial_backoff=0,
max_backoff=0,
)
first_config = provider._client.aio.models.generate_content.call_args_list[0].kwargs["config"]
fallback_config = provider._client.aio.models.generate_content.call_args_list[1].kwargs["config"]
assert result.answer == "ok"
assert first_config.response_schema is StructuredAnswer
assert first_config.system_instruction == "Return concise JSON."
assert fallback_config.response_schema is None
assert fallback_config.response_mime_type is None
assert fallback_config.system_instruction.startswith("Return concise JSON.")
assert "valid JSON matching this schema" in fallback_config.system_instruction
assert '"answer"' in fallback_config.system_instruction
@pytest.mark.asyncio
async def test_gemini_cached_parse_retry_keeps_cached_native_schema():
"""Cached structured retries keep cache context instead of switching prompts."""
from pydantic import BaseModel
class StructuredAnswer(BaseModel):
answer: str
provider = _make_gemini_provider()
invalid = _fake_gemini_response()
invalid.text = "not json"
valid = _fake_gemini_response()
valid.text = '{"answer": "ok"}'
provider._client.aio.models.generate_content = AsyncMock(side_effect=[invalid, valid])
result = await provider.call(
messages=[
{"role": "system", "content": "Return concise JSON."},
{"role": "user", "content": "hello"},
],
response_format=StructuredAnswer,
cached_prefix="cachedContents/test",
scope="test",
max_retries=1,
initial_backoff=0,
max_backoff=0,
)
first_config = provider._client.aio.models.generate_content.call_args_list[0].kwargs["config"]
retry_config = provider._client.aio.models.generate_content.call_args_list[1].kwargs["config"]
assert result.answer == "ok"
assert first_config.cached_content == "cachedContents/test"
assert retry_config.cached_content == "cachedContents/test"
assert retry_config.response_schema is StructuredAnswer
assert retry_config.response_mime_type == "application/json"
assert retry_config.system_instruction is None
# ─── LiteLLM ──────────────────────────────────────────────────────────────────
@@ -39,20 +39,11 @@ async def hundred_tenant_schemas(memory: MemoryEngine):
"""Create N_TENANTS isolated schemas cloning the loop's tables; drop them after."""
prefix = f"mt{uuid.uuid4().hex[:8]}"
schemas = [f"{prefix}_{i:03d}" for i in range(N_TENANTS)]
# Create all schemas + their tables in ONE transaction so the schemas become
# visible to other connections only once fully built. Without this, each DDL
# autocommits, leaving a window where a schema exists with only some of its
# tables. The global maintenance routines (schemas_with_expired_rows /
# banks_needing_consolidation) discover schemas by table presence and are run
# concurrently by test_maintenance_routines on another xdist worker against
# the shared test DB; they would query a not-yet-created table in a half-built
# schema and fail with `relation "<schema>.<table>" does not exist`.
async with memory._pool.acquire() as conn:
async with conn.transaction():
for s in schemas:
await conn.execute(f'CREATE SCHEMA "{s}"')
for table in _CLONED_TABLES:
await conn.execute(f'CREATE TABLE "{s}".{table} (LIKE public.{table} INCLUDING DEFAULTS)')
for s in schemas:
await conn.execute(f'CREATE SCHEMA "{s}"')
for table in _CLONED_TABLES:
await conn.execute(f'CREATE TABLE "{s}".{table} (LIKE public.{table} INCLUDING DEFAULTS)')
try:
yield prefix, schemas
finally:
@@ -134,33 +134,6 @@ async def test_banks_needing_consolidation_includes_in_flight_after_completion(m
assert bank in {r["bank_id"] for r in rows}
@pytest.mark.asyncio
async def test_banks_needing_consolidation_skips_schema_with_vanished_table(memory: MemoryEngine):
"""A schema discovered via its ``memory_units`` table but missing the
``banks`` table the routine joins must be skipped, not abort the scan.
This reproduces the time-of-check/time-of-use race deterministically: the
routine snapshots schemas owning ``memory_units`` from ``pg_class`` and then
joins each schema's ``banks`` table. A tenant being dropped or migrated (and,
in the test suite, the concurrent multi-tenant maintenance test) can leave a
schema whose ``banks`` table is gone. Before the fix the dynamic query raised
``undefined_table`` and aborted the whole routine (migration c7e9f1a3b5d2)."""
schema = f"mtvanish{uuid.uuid4().hex[:8]}"
try:
async with memory._pool.acquire() as conn:
await conn.execute(f'CREATE SCHEMA "{schema}"')
# Discovered by the FOR loop (has memory_units) but the JOIN target
# `banks` is absent — exactly a half-built / vanishing schema.
await conn.execute(f'CREATE TABLE "{schema}".memory_units (LIKE public.memory_units INCLUDING DEFAULTS)')
# Must not raise; the bad schema is simply skipped.
rows = await conn.fetch("SELECT schema_name, bank_id FROM public.banks_needing_consolidation()")
assert schema not in {r["schema_name"] for r in rows}
finally:
async with memory._pool.acquire() as conn:
await conn.execute(f'DROP SCHEMA IF EXISTS "{schema}" CASCADE')
@pytest.mark.asyncio
async def test_schemas_with_expired_rows(memory: MemoryEngine):
"""Returns schemas holding a row older than p_days; respects the p_days<=0 guard."""
+8 -22
View File
@@ -423,28 +423,6 @@ def test_global_mcp_enabled_tools_intersects_with_single_bank_mode(mock_memory):
assert "list_banks" not in tools # single-bank mode excludes it regardless
def test_mcp_instructions_append_to_retain_and_recall_descriptions(mock_memory):
"""HINDSIGHT_API_MCP_INSTRUCTIONS customizes retain/recall tool descriptions."""
from unittest.mock import MagicMock, patch
from hindsight_api.api.mcp import create_mcp_server
custom_instructions = "Also store every action you take."
mock_cfg = MagicMock()
mock_cfg.mcp_enabled_tools = ["retain", "recall", "reflect"]
mock_cfg.mcp_instructions = custom_instructions
with patch("hindsight_api.api.mcp._get_raw_config", return_value=mock_cfg):
mcp_server = create_mcp_server(mock_memory, multi_bank=True)
tools = _tools(mcp_server)
expected_suffix = f"Additional instructions: {custom_instructions}"
assert expected_suffix in tools["retain"].description
assert expected_suffix in tools["recall"].description
assert expected_suffix not in tools["reflect"].description
@pytest.mark.asyncio
async def test_routing_logic_from_url_path():
"""Test that routing correctly selects server based on URL structure.
@@ -452,6 +430,14 @@ async def test_routing_logic_from_url_path():
Simulates the path parsing logic from MCPMiddleware.__call__ after the
prefix has been stripped. Any first path segment is treated as a bank_id.
"""
from hindsight_api.api.mcp import MCPMiddleware
# Mock memory
mock_memory = MagicMock()
# Create middleware
middleware = MCPMiddleware(None, mock_memory)
# Simulate different URL patterns and verify routing
# Path is what remains after stripping the /mcp prefix
test_cases = [
-104
View File
@@ -1911,107 +1911,3 @@ class TestBankToolFiltering:
# Filter bypassed — config resolver was never consulted, all tools visible
assert "recall" in visible
mock_memory_with_resolver._config_resolver.get_bank_config.assert_not_called()
@pytest.mark.asyncio
class TestToolAnnotations:
"""Every MCP tool must carry read-only / destructive hints (openWorldHint=False)."""
async def test_read_only_tool(self, mock_memory):
ann = _tools(_make_mcp_server(mock_memory, {"recall"}))["recall"].annotations
assert ann is not None
assert ann.readOnlyHint is True
assert ann.openWorldHint is False
async def test_reflect_is_read_only(self, mock_memory):
# reflect synthesizes an answer and persists nothing (memory_engine.reflect_async),
# so it carries readOnlyHint=True like recall.
ann = _tools(_make_mcp_server(mock_memory, {"reflect"}))["reflect"].annotations
assert ann is not None
assert ann.readOnlyHint is True
assert ann.openWorldHint is False
async def test_destructive_tool(self, mock_memory):
ann = _tools(_make_mcp_server(mock_memory, {"delete_bank"}))["delete_bank"].annotations
assert ann is not None
assert ann.readOnlyHint is False
assert ann.destructiveHint is True
async def test_write_tool_is_not_destructive(self, mock_memory):
ann = _tools(_make_mcp_server(mock_memory, {"retain"}))["retain"].annotations
assert ann is not None
assert ann.readOnlyHint is False
assert ann.destructiveHint is False
async def test_annotations_apply_in_single_bank_mode(self, mock_memory):
ann = _tools(_make_mcp_server(mock_memory, {"recall"}, include_bank_id=False))["recall"].annotations
assert ann is not None
assert ann.readOnlyHint is True
def _reflect_mcp_with_trace(include_bank_id_param: bool):
"""An MCP server whose reflect returns a result carrying tool_trace/llm_trace."""
from fastmcp import FastMCP
# Mirrors ReflectResult: the agentic loop's trace fields are large and present.
reflect_payload = {
"text": "answer",
"based_on": {"world": []},
"tool_trace": [{"tool": "recall", "output": "x" * 1000}],
"llm_trace": [{"model": "test", "output": "y" * 1000}],
"directives_applied": [{"id": "d1", "name": "Tone", "content": "z" * 1000}],
}
memory = MagicMock()
memory.reflect_async = AsyncMock(
return_value=MagicMock(
model_dump_json=lambda indent=None: json.dumps(reflect_payload),
model_dump=lambda: dict(reflect_payload),
structured_output=None,
)
)
mcp = FastMCP("test")
config = MCPToolsConfig(
bank_id_resolver=lambda: "test-bank",
include_bank_id_param=include_bank_id_param,
tools={"reflect"},
)
register_mcp_tools(mcp, memory, config)
return mcp
def _reflect_result_data(result) -> dict:
"""The multi-bank reflect returns a JSON string; single-bank returns a dict."""
return json.loads(result) if isinstance(result, str) else result
@pytest.mark.asyncio
class TestReflectTraceOmission:
"""reflect must not leak the agentic tool_trace/llm_trace/directives_applied into MCP responses by default."""
@pytest.mark.parametrize("multi_bank", [True, False])
async def test_trace_omitted_by_default(self, multi_bank):
mcp = _reflect_mcp_with_trace(multi_bank)
data = _reflect_result_data(await _tools(mcp)["reflect"].fn(query="q"))
assert data["text"] == "answer"
assert "tool_trace" not in data
assert "llm_trace" not in data
# directives_applied is built "for the trace" and carries full directive content,
# so it must be omitted by default like the other trace fields.
assert "directives_applied" not in data
@pytest.mark.parametrize("multi_bank", [True, False])
async def test_trace_included_when_requested(self, multi_bank):
mcp = _reflect_mcp_with_trace(multi_bank)
data = _reflect_result_data(await _tools(mcp)["reflect"].fn(query="q", include_trace=True))
assert "tool_trace" in data
assert "llm_trace" in data
assert "directives_applied" in data
@pytest.mark.parametrize("multi_bank", [True, False])
async def test_based_on_flag_is_independent_of_trace(self, multi_bank):
# include_based_on keeps based_on but must not pull the trace back in.
mcp = _reflect_mcp_with_trace(multi_bank)
data = _reflect_result_data(await _tools(mcp)["reflect"].fn(query="q", include_based_on=True))
assert "based_on" in data
assert "tool_trace" not in data
assert "directives_applied" not in data
@@ -334,9 +334,6 @@ def _make_minimal_engine():
mock_embeddings = MagicMock()
mock_embeddings.dimension = 384
from hindsight_api.config import clear_config_cache
from hindsight_api.engine.memory_engine import MemoryEngine
with patch.dict(
os.environ,
{
@@ -346,17 +343,11 @@ def _make_minimal_engine():
},
clear=False,
):
clear_config_cache()
engine = MemoryEngine(db_url="postgresql://localhost/hindsight_test", embeddings=mock_embeddings)
from hindsight_api.config import clear_config_cache
from hindsight_api.engine.memory_engine import MemoryEngine
# Constructing the engine above repopulated the process-global config cache
# from the patched env (provider="none" forces retain_extraction_mode="chunks").
# Now that the patched env is gone, drop that cache so the leaked "none"/chunks
# config does not bleed into other tests on this xdist worker — their retains
# would silently skip entity extraction (0 unit_entities) and fail unrelated
# assertions. The next get_config() rebuilds from the real env.
clear_config_cache()
return engine
clear_config_cache()
return MemoryEngine(db_url="postgresql://localhost/hindsight_test", embeddings=mock_embeddings)
def test_engine_memory_defense_shares_ext_ctx() -> None:
-30
View File
@@ -11,7 +11,6 @@ from hindsight_api.metrics import (
get_token_bucket,
create_metrics_collector,
initialize_metrics,
normalize_http_endpoint,
)
@@ -222,17 +221,6 @@ class TestMetricsCollector:
assert reflect_attrs["operation"] == "reflect"
assert reflect_attrs["source"] == "api"
def test_record_operation_result_records_with_explicit_success(self, collector):
"""Direct recording path used by the worker (source=worker, explicit success)."""
collector.record_operation_result("retain", bank_id="test_bank", success=False, duration=1.5, source="worker")
duration, attributes = collector.operation_duration.record.call_args[0]
assert duration == 1.5
assert attributes["operation"] == "retain"
assert attributes["source"] == "worker"
assert attributes["success"] == "false"
collector.operation_total.add.assert_called_once_with(1, attributes)
def test_record_operation_includes_bank_id_when_enabled(self):
"""Test that bank_id is included in attributes when metrics_include_bank_id is enabled."""
mock_config = MagicMock()
@@ -336,24 +324,6 @@ class TestGetTokenBucket:
assert get_token_bucket(1000000) == "50k+"
class TestNormalizeHttpEndpoint:
"""Tests for normalize_http_endpoint (low-cardinality HTTP metric labels)."""
def test_templates_high_cardinality_segments(self):
"""Bank ids (incl. non-numeric), UUIDs, and numeric ids collapse to placeholders."""
cases = [
("/v1/default/banks/user-1680/memories/recall", "/v1/default/banks/{bank_id}/memories/recall"),
("/v1/default/banks/tenant-acme/memories", "/v1/default/banks/{bank_id}/memories"),
("/v1/default/banks/user-1680", "/v1/default/banks/{bank_id}"),
("/v1/default/banks/3f8c1e2a-1111-2222-3333-444455556666/config", "/v1/default/banks/{bank_id}/config"),
("/v1/default/banks/42/config", "/v1/default/banks/{bank_id}/config"),
("/v1/default/banks", "/v1/default/banks"),
("/health", "/health"),
]
for raw, expected in cases:
assert normalize_http_endpoint(raw) == expected, raw
class TestLLMMetrics:
"""Tests for LLM-specific metrics recording."""
@@ -11,7 +11,7 @@ from pathlib import Path
import pytest
from alembic import command
from alembic.config import Config
from sqlalchemy import create_engine, text
from sqlalchemy import create_engine, text, inspect
# ---------------------------------------------------------------------------
# Helpers
@@ -33,17 +33,8 @@ def _upgrade(db_url: str, revision: str) -> None:
command.upgrade(_alembic_cfg(db_url), revision)
def _reset_public_schema(db_url: str) -> None:
engine = create_engine(db_url, isolation_level="AUTOCOMMIT")
try:
with engine.connect() as conn:
# This test rewinds/replays migration history against a persistent
# pg0 instance. Rebuild only its dedicated public schema so a
# previous run cannot leave alembic_version ahead of the real DDL.
conn.execute(text("DROP SCHEMA IF EXISTS public CASCADE"))
conn.execute(text("CREATE SCHEMA public"))
finally:
engine.dispose()
def _downgrade(db_url: str, revision: str) -> None:
command.downgrade(_alembic_cfg(db_url), revision)
# ---------------------------------------------------------------------------
@@ -52,16 +43,16 @@ def _reset_public_schema(db_url: str) -> None:
@pytest.fixture(scope="module")
def pre_backsweep_db_url() -> str:
def pre_backsweep_db_url():
"""
Spin up a dedicated pg0 instance and ensure schema is at the revision
just before the backsweep so each test can seed orphan data and then
apply the backsweep itself.
Because pg0 data directories persist across test runs, the DB may already
have schema from a previous test run. Reset this test's dedicated schema
first, then migrate to the real pre-backsweep revision instead of stamping
a head schema backward.
Because pg0 data directories persist across test runs, the DB may
already be at head. We upgrade to head first (to ensure all tables
exist), then stamp the revision back to pre-backsweep so Alembic
treats the backsweep as not-yet-applied.
"""
from hindsight_api.pg0 import EmbeddedPostgres
@@ -72,8 +63,10 @@ def pre_backsweep_db_url() -> str:
finally:
loop.close()
_reset_public_schema(url)
_upgrade(url, "f6g7h8i9j0k1")
# Ensure all tables exist (upgrade to head), then stamp back to
# pre-backsweep so the backsweep migration will actually run.
_upgrade(url, "heads")
command.stamp(_alembic_cfg(url), "f6g7h8i9j0k1")
return url
@@ -82,7 +75,7 @@ def pre_backsweep_db_url() -> str:
# ---------------------------------------------------------------------------
def test_backsweep_removes_orphans_and_preserves_legit_rows(pre_backsweep_db_url: str) -> None:
def test_backsweep_removes_orphans_and_preserves_legit_rows(pre_backsweep_db_url):
"""
Seed four kinds of rows then apply the backsweep migration and verify:
@@ -121,20 +114,15 @@ def test_backsweep_removes_orphans_and_preserves_legit_rows(pre_backsweep_db_url
conn.execute(text("INSERT INTO banks (bank_id) VALUES (:b)"), {"b": alive_bank})
# --- seed memory_units ---
def insert_mu(
uid: uuid.UUID,
bank: str,
fact_type: str,
sources: list[uuid.UUID] | None = None,
) -> None:
def insert_mu(uid, bank, fact_type, sources=None):
src_arr = "{" + ",".join(str(s) for s in (sources or [])) + "}"
conn.execute(
text(
"""
INSERT INTO memory_units
(id, bank_id, text, event_date, fact_type, source_memory_ids)
(id, bank_id, text, fact_type, source_memory_ids)
VALUES
(:id, :bank, :text, now(), :ft, CAST(:src AS uuid[]))
(:id, :bank, :text, :ft, CAST(:src AS uuid[]))
"""
),
{"id": uid, "bank": bank, "text": "test", "ft": fact_type, "src": src_arr},
@@ -162,7 +150,7 @@ def test_backsweep_removes_orphans_and_preserves_legit_rows(pre_backsweep_db_url
# --- verify ---
with engine.connect() as conn:
def exists(uid: uuid.UUID) -> bool:
def exists(uid):
return conn.execute(text("SELECT 1 FROM memory_units WHERE id = :id"), {"id": uid}).fetchone() is not None
# Must be gone
@@ -1,110 +0,0 @@
"""Provider quota reset windows defer worker retries instead of failing retains."""
from datetime import UTC, datetime, timedelta
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from openai import APIStatusError
from hindsight_api.engine.llm_interface import ProviderRateLimitResetError
from hindsight_api.engine.providers.openai_compatible_llm import OpenAICompatibleLLM
def _llm() -> OpenAICompatibleLLM:
return OpenAICompatibleLLM(
provider="zai",
model="glm-5-turbo",
api_key="test",
base_url="https://example.com/v1",
)
def _usage_limit_error(reset_at: str) -> APIStatusError:
body = {
"code": "1308",
"message": f"Usage limit reached for 5 hour. Your limit will reset at {reset_at}",
}
response = MagicMock()
response.status_code = 429
response.text = '{"code": "1308", "message": "usage limit"}'
response.headers = {}
return APIStatusError("rate limited", response=response, body=body)
def _short_retry_after_error() -> APIStatusError:
response = MagicMock()
response.status_code = 429
response.text = '{"code": "rate_limit", "message": "retry shortly"}'
response.headers = {"retry-after": "1"}
return APIStatusError("rate limited", response=response, body={"message": "retry shortly"})
@pytest.mark.asyncio
async def test_usage_limit_429_with_reset_defers_without_inner_retry() -> None:
llm = _llm()
reset_at = (datetime.now(UTC) + timedelta(hours=5)).replace(microsecond=0)
create = AsyncMock(side_effect=_usage_limit_error(reset_at.isoformat().replace("+00:00", "Z")))
llm._client = SimpleNamespace(chat=SimpleNamespace(completions=SimpleNamespace(create=create)))
with patch(
"hindsight_api.engine.providers.openai_compatible_llm.asyncio.sleep",
new_callable=AsyncMock,
) as sleep:
with pytest.raises(ProviderRateLimitResetError) as exc_info:
await llm.call(
messages=[{"role": "user", "content": "x"}],
scope="retain_extract_facts",
max_retries=2,
)
assert create.await_count == 1
sleep.assert_not_awaited()
assert abs((exc_info.value.retry_at - reset_at).total_seconds()) < 1
assert "Provider quota exhausted" in str(exc_info.value)
@pytest.mark.asyncio
async def test_short_retry_after_429_uses_normal_retry_loop() -> None:
llm = _llm()
create = AsyncMock(side_effect=_short_retry_after_error())
llm._client = SimpleNamespace(chat=SimpleNamespace(completions=SimpleNamespace(create=create)))
with patch(
"hindsight_api.engine.providers.openai_compatible_llm.asyncio.sleep",
new_callable=AsyncMock,
) as sleep:
with pytest.raises(APIStatusError):
await llm.call(
messages=[{"role": "user", "content": "x"}],
scope="retain_extract_facts",
max_retries=2,
max_backoff=60,
)
assert create.await_count == 3
assert sleep.await_count == 2
@pytest.mark.asyncio
async def test_extract_facts_from_text_preserves_provider_quota_reset(monkeypatch) -> None:
from hindsight_api.engine.retain import fact_extraction
retry_at = (datetime.now(UTC) + timedelta(hours=2)).replace(microsecond=0)
async def quota_limited_chunk(**_: object) -> None:
raise ProviderRateLimitResetError(retry_at=retry_at, message="quota resets later")
monkeypatch.setattr(fact_extraction, "_extract_facts_with_auto_split", quota_limited_chunk)
with pytest.raises(ProviderRateLimitResetError) as exc_info:
await fact_extraction.extract_facts_from_text(
text="Alice moved to Berlin.",
event_date=None,
llm_config=object(),
agent_name="TestAgent",
config=SimpleNamespace(retain_chunk_size=1000, retain_structured_chunk_size=None),
)
assert exc_info.value.retry_at == retry_at
assert "Fact extraction deferred by provider quota" in str(exc_info.value)
@@ -1,184 +0,0 @@
"""Tests for the recall `prefer_observations` deduplication flag.
When the caller recalls raw facts ('world'/'experience') together with
'observation' and sets prefer_observations=True, any raw fact that a returned
observation was consolidated from (tracked via memory_units.source_memory_ids)
is dropped so the observation supersedes it no duplicate content.
Dedup is provenance-based, not semantic: a raw fact that is semantically
similar to an observation but NOT listed in its source_memory_ids must survive.
No LLM required inserts memory_units directly via SQL with real embeddings.
"""
import uuid
import pytest
import pytest_asyncio
from hindsight_api import MemoryEngine, RequestContext
from hindsight_api.engine.retain import embedding_utils
RC = RequestContext(tenant_id="default")
QUERY = "Alice mountain hiking"
# Two raw facts the observation is consolidated from (must be dropped when the
# flag is on), one raw fact that is semantically similar but NOT a source (must
# survive), and the observation itself.
SRC1_TEXT = "Alice loves hiking in the mountains"
SRC2_TEXT = "Alice hikes the Alps every summer"
NON_SRC_TEXT = "Alice enjoys exploring mountain hiking trails"
OBS_TEXT = "Alice is an avid mountain hiker"
async def _insert_unit(
conn,
*,
unit_id: str,
text: str,
bank_id: str,
embedding_str: str,
fact_type: str = "world",
source_memory_ids: list[uuid.UUID] | None = None,
) -> None:
await conn.execute(
"""
INSERT INTO memory_units (id, bank_id, text, fact_type, embedding, source_memory_ids)
VALUES ($1, $2, $3, $4, $5::vector, $6::uuid[])
""",
unit_id,
bank_id,
text,
fact_type,
embedding_str,
source_memory_ids,
)
def _to_str(emb: list[float]) -> str:
return "[" + ",".join(str(v) for v in emb) + "]"
def _result_ids(result) -> set[str]:
return {str(r.id) for r in result.results}
@pytest_asyncio.fixture
async def seeded_obs_memory(memory_no_llm_verify: MemoryEngine):
"""Seed two source facts, one non-source fact, and an observation over the two sources."""
engine = memory_no_llm_verify
bank_id = f"test-prefer-obs-{uuid.uuid4().hex[:8]}"
await engine.get_bank_profile(bank_id, request_context=RC)
src1_id = str(uuid.uuid4())
src2_id = str(uuid.uuid4())
non_src_id = str(uuid.uuid4())
obs_id = str(uuid.uuid4())
embeddings = await embedding_utils.generate_embeddings_batch(
engine.embeddings,
[SRC1_TEXT, SRC2_TEXT, NON_SRC_TEXT, OBS_TEXT],
)
pool = await engine._get_pool()
async with pool.acquire() as conn:
await _insert_unit(conn, unit_id=src1_id, text=SRC1_TEXT, bank_id=bank_id, embedding_str=_to_str(embeddings[0]))
await _insert_unit(conn, unit_id=src2_id, text=SRC2_TEXT, bank_id=bank_id, embedding_str=_to_str(embeddings[1]))
await _insert_unit(
conn, unit_id=non_src_id, text=NON_SRC_TEXT, bank_id=bank_id, embedding_str=_to_str(embeddings[2])
)
await _insert_unit(
conn,
unit_id=obs_id,
text=OBS_TEXT,
bank_id=bank_id,
embedding_str=_to_str(embeddings[3]),
fact_type="observation",
source_memory_ids=[uuid.UUID(src1_id), uuid.UUID(src2_id)],
)
ids = {"src1": src1_id, "src2": src2_id, "non_src": non_src_id, "obs": obs_id}
yield engine, bank_id, ids
await engine.delete_bank(bank_id, request_context=RC)
class TestPreferObservations:
async def test_disabled_returns_sources_and_observation(self, seeded_obs_memory):
"""Without the flag, the source facts AND the observation are all returned."""
engine, bank_id, ids = seeded_obs_memory
result = await engine.recall_async(
bank_id=bank_id,
query=QUERY,
request_context=RC,
fact_type=["world", "experience", "observation"],
prefer_observations=False,
max_tokens=10000,
)
found = _result_ids(result)
assert ids["src1"] in found
assert ids["src2"] in found
assert ids["obs"] in found
async def test_enabled_drops_source_facts_keeps_observation(self, seeded_obs_memory):
"""With the flag, the observation supersedes the facts it was consolidated from."""
engine, bank_id, ids = seeded_obs_memory
result = await engine.recall_async(
bank_id=bank_id,
query=QUERY,
request_context=RC,
fact_type=["world", "experience", "observation"],
prefer_observations=True,
max_tokens=10000,
)
found = _result_ids(result)
assert ids["obs"] in found, "the observation must remain"
assert ids["src1"] not in found, "source fact 1 is superseded by the observation"
assert ids["src2"] not in found, "source fact 2 is superseded by the observation"
async def test_enabled_keeps_non_source_fact(self, seeded_obs_memory):
"""Dedup is provenance-based: a similar fact NOT in source_memory_ids survives."""
engine, bank_id, ids = seeded_obs_memory
result = await engine.recall_async(
bank_id=bank_id,
query=QUERY,
request_context=RC,
fact_type=["world", "experience", "observation"],
prefer_observations=True,
max_tokens=10000,
)
found = _result_ids(result)
assert ids["non_src"] in found, "a non-source fact must not be dropped, even if semantically similar"
async def test_noop_without_observation_type(self, seeded_obs_memory):
"""The flag is a no-op when 'observation' is not among the requested types."""
engine, bank_id, ids = seeded_obs_memory
result = await engine.recall_async(
bank_id=bank_id,
query=QUERY,
request_context=RC,
fact_type=["world", "experience"],
prefer_observations=True,
max_tokens=10000,
)
found = _result_ids(result)
assert ids["src1"] in found
assert ids["src2"] in found
def test_flag_is_opt_in_by_default():
"""prefer_observations is opt-in: off at the API surface and the engine method.
The engine default in particular must stay False so internal callers notably
consolidation, which needs the raw facts it folds into observations are never
silently deduped.
"""
import inspect
from hindsight_api.api.http import RecallRequest
from hindsight_api.engine.memory_engine import MemoryEngine
assert RecallRequest(query="anything").prefer_observations is False
engine_default = inspect.signature(MemoryEngine.recall_async).parameters["prefer_observations"].default
assert engine_default is False
@@ -100,11 +100,7 @@ async def test_recall_async_passes_question_date_to_combined_scoring(monkeypatch
)
def apply_combined_scoring(
scored_results: list[ScoredResult],
*,
now: datetime,
is_passthrough_reranker: bool,
**_kwargs: object,
scored_results: list[ScoredResult], *, now: datetime, is_passthrough_reranker: bool
) -> None:
nonlocal captured_now
assert is_passthrough_reranker is False
@@ -1,78 +0,0 @@
"""Tests for _strip_reasoning_tags helper in OpenAI-compatible LLM provider."""
from hindsight_api.engine.providers.openai_compatible_llm import _strip_reasoning_tags
class TestStripReasoningTags:
"""Test reasoning/thinking tag stripping from LLM responses."""
def test_plain_text_unchanged(self):
"""Text without reasoning tags passes through (modulo edge whitespace)."""
content = "User prefers functional programming patterns."
assert _strip_reasoning_tags(content) == content
def test_empty_string(self):
"""Empty string passes through."""
assert _strip_reasoning_tags("") == ""
def test_closed_think_stripped(self):
"""A closed <think>...</think> block is removed."""
content = "<think>let me reason</think>The answer is 42."
assert _strip_reasoning_tags(content) == "The answer is 42."
def test_closed_thinking_stripped(self):
assert _strip_reasoning_tags("<thinking>reasoning</thinking>Result") == "Result"
def test_closed_thought_stripped(self):
assert _strip_reasoning_tags("<thought>hmm</thought>Result") == "Result"
def test_closed_reasoning_stripped(self):
assert _strip_reasoning_tags("<reasoning>step by step</reasoning>Result") == "Result"
def test_startthink_endthink_stripped(self):
"""The |startthink|...|endthink| marker style is removed."""
content = "|startthink|internal monologue|endthink|Final output"
assert _strip_reasoning_tags(content) == "Final output"
def test_multiline_think_stripped(self):
"""DOTALL: a multi-line thinking block is fully removed."""
content = "<think>\nline one\nline two\n</think>\nThe real content."
assert _strip_reasoning_tags(content) == "The real content."
def test_unclosed_think_stripped_to_end(self):
"""An unclosed <think> (truncated output) is removed to end-of-string."""
content = "Partial answer.\n<think>I started thinking but got cut off"
assert _strip_reasoning_tags(content) == "Partial answer."
def test_unclosed_thinking_stripped_to_end(self):
content = "result text\n<thinking>dangling reasoning with no close"
assert _strip_reasoning_tags(content) == "result text"
def test_only_unclosed_think_becomes_empty(self):
"""Content that is entirely an unclosed thinking block collapses to empty."""
content = "<think>everything is reasoning and it never closed"
assert _strip_reasoning_tags(content) == ""
def test_multiple_blocks_stripped(self):
"""Multiple closed blocks are all removed."""
content = "<think>a</think>Hello <think>b</think>World"
assert _strip_reasoning_tags(content) == "Hello World"
def test_mental_model_markdown_contamination(self):
"""Real-world MiniMax-M3 free-form leak: <think> wrapping a markdown mental model."""
content = (
"<think>\n"
"The user keeps asking about FP. I should consolidate this.\n"
"</think>\n"
"# Mental Model: Coding Preferences\n\n"
"The user prefers functional programming patterns and immutable data."
)
result = _strip_reasoning_tags(content)
assert "<think>" not in result
assert "</think>" not in result
assert result.startswith("# Mental Model: Coding Preferences")
def test_unclosed_think_after_json_payload(self):
"""Truncated <think> trailing valid JSON is stripped (closing tag absent)."""
content = '{"facts": [{"what": "test"}]}\n<think>oops truncated'
assert _strip_reasoning_tags(content) == '{"facts": [{"what": "test"}]}'
@@ -1,81 +0,0 @@
"""Regression test for https://github.com/vectorize-io/hindsight/issues/2301
Raising ``HINDSIGHT_API_RETAIN_STRUCTURED_CHUNK_SIZE`` above
``HINDSIGHT_API_RETAIN_CHUNK_SIZE`` and retaining a JSONL/conversation document
whose line/turn overflows the chunk size used to crash with::
asyncpg.exceptions.CardinalityViolationError:
ON CONFLICT DO UPDATE command cannot affect row a second time
The streaming retain pipeline pre-chunks each document once (one ``chunk_index``
per piece) and then re-chunks every piece during extraction. With the structured
cap above the chunk size, a pre-chunk could legitimately exceed the re-chunk
budget, so it re-split into several sub-chunks that all inherited the one
``chunk_index`` colliding on ``chunk_id = {bank}_{doc}_{index}`` in a single
upsert batch. ``chunk_text`` is now idempotent, so the re-chunk is a no-op.
"""
from datetime import datetime, timezone
import pytest
from hindsight_api.config import clear_config_cache
def _ts() -> float:
return datetime.now(timezone.utc).timestamp()
@pytest.fixture(autouse=True)
def _structured_chunk_env(monkeypatch):
# Structured cap ABOVE the chunk size — the configuration that triggers #2301.
# Default-scale sizes (matching the issue) so the document yields only a
# handful of chunks: a smaller chunk size explodes the embedding/link work and
# destabilises the shared session fixture under xdist.
monkeypatch.setenv("HINDSIGHT_API_RETAIN_CHUNK_SIZE", "3000")
monkeypatch.setenv("HINDSIGHT_API_RETAIN_STRUCTURED_CHUNK_SIZE", "4500")
monkeypatch.setenv("HINDSIGHT_API_ENABLE_AUTO_CONSOLIDATION", "false")
monkeypatch.setenv("HINDSIGHT_API_ENABLE_OBSERVATIONS", "false")
clear_config_cache()
yield
clear_config_cache()
@pytest.mark.asyncio
async def test_jsonl_line_over_chunk_size_retains_without_collision(memory, request_context):
"""A JSONL document with a line longer than the chunk size retains cleanly
when the structured cap is raised above it (issue #2301)."""
import json
bank_id = f"test_2301_{_ts()}"
document_id = "doc-2301"
try:
body = "\n".join(
[
json.dumps({"role": "user", "content": "short opening line"}),
# Between chunk size (3000) and structured cap (4500): kept whole by
# the producer, would re-split on re-chunk without the fix.
json.dumps({"role": "assistant", "content": "k" * 3800}),
# Past even the structured cap: fragmented as text.
json.dumps({"role": "assistant", "content": "m" * 9000}),
json.dumps({"role": "user", "content": "short closing line"}),
]
)
# The bug raised CardinalityViolationError here.
await memory.retain_async(
bank_id=bank_id,
content=body,
context="jsonl with oversized line",
document_id=document_id,
request_context=request_context,
)
chunks = await memory.list_document_chunks(bank_id, document_id, limit=10000, request_context=request_context)
indices = sorted(c["chunk_index"] for c in chunks["items"])
# No collisions: chunk_index values are unique.
assert len(indices) == len(set(indices)), f"duplicate chunk_index values: {indices}"
assert indices == list(range(len(indices))), f"chunk_index sequence not contiguous: {indices}"
finally:
await memory.delete_bank(bank_id, request_context=request_context)
@@ -1,168 +0,0 @@
"""Regression: sub-batch slices that each span MULTIPLE extraction chunks must
keep full chunk coverage on BOTH the sync (inline) and async (submitted) retain
paths.
Two distinct bugs hid behind the same symptom ingesting a large plain-text
document dropped most of its body (and any fact past the first slice). Both only
trigger when an oversized single item is split into sequential sub-batches whose
*slices each re-chunk into several extraction chunks* (the default config: batch
tokens 10k ~30k-char slices, re-chunked at 3k ~10 chunks/slice):
1. chunk_index offset (sync + async). retain_batch_async advanced the per-document
chunk_index cursor by re-chunking ``item["content"]`` AFTER the orchestrator
had consumed (popped) it ``chunk_text("")`` returns ``[""]`` (count 1), so
the cursor moved by 1 per sub-batch instead of by the real chunk count. Later
slices restarted ~1 slot in, colliding ``chunk_id = {bank}_{doc}_{index}`` and
overwriting earlier chunks via upsert.
2. whole-document recovery skip (async only). All sub-batches of one submitted
operation share one ``operation_id``; the first slice stamps the document into
``result_metadata.facts_committed_document_ids``. The crash-recovery fast-path
then saw every later slice's document already "committed" and skipped
extraction entirely, so only the first slice survived.
The existing #1888 coverage tests use ``RETAIN_BATCH_TOKENS=100`` (a ~300-char
budget, under the chunk size) so every slice collapses to ONE chunk which masks
both bugs (offset-by-1 happens to equal the real count, and a 1-chunk doc isn't
re-sliced). These tests size the body so each slice fans out to ~6 chunks, with
globally-unique tokens so no chunk-hash dedup hides a dropped slice, and assert
full coverage + contiguous indices + a needle planted in a late slice.
"""
from datetime import datetime, timezone
import pytest
from hindsight_api.config import clear_config_cache
# The async test submits via submit_async_retain, which inserts parent/child rows
# into async_operations. test_worker.py drives its own WorkerPoller.claim_batch()
# against the same pool, so running the two files on different xdist workers lets
# them steal each other's pending rows. Share the "worker_tests" group so they
# serialize on the same xdist process (matches test_async_batch_retain.py).
pytestmark = pytest.mark.xdist_group("worker_tests")
# Planted in a late paragraph so it lands in a late sub-batch slice — the first
# thing either bug drops (mirrors the field-reported "165 commits" fact that
# vanished on the async path). A single no-space token so it can't straddle a
# chunk boundary (a multi-word phrase can split across two chunks at this test's
# small 500-char chunk size and read as "dropped" when it wasn't).
NEEDLE = "NEEDLE_165_COMMITS_MERGED_INTO_THE_MAIN_BRANCH"
def _ts() -> float:
return datetime.now(timezone.utc).timestamp()
@pytest.fixture(autouse=True)
def _multichunk_split_env(monkeypatch):
# Small extraction chunks (500 chars) with a batch-token budget whose char
# budget (700 * 3 = 2100) spans several chunks, so each oversized sub-batch
# slice fans out to ~6 extraction chunks. Skip consolidation/observations to
# keep the test fast and deterministic.
monkeypatch.setenv("HINDSIGHT_API_RETAIN_CHUNK_SIZE", "500")
monkeypatch.setenv("HINDSIGHT_API_RETAIN_BATCH_TOKENS", "700")
monkeypatch.setenv("HINDSIGHT_API_ENABLE_AUTO_CONSOLIDATION", "false")
monkeypatch.setenv("HINDSIGHT_API_ENABLE_OBSERVATIONS", "false")
clear_config_cache()
yield
clear_config_cache()
def _make_body(paragraphs: int = 24, needle_at: int = 20) -> str:
"""Plain-text transcript whose every token is unique across the whole body,
so no two extraction chunks can hash-collide (a real content-hash collision
would legitimately dedup and mask a dropped slice). The needle sits in a late
paragraph."""
lines = []
for i in range(paragraphs):
toks = " ".join(f"w{i:03d}t{j:03d}" for j in range(60))
if i == needle_at:
lines.append(f"[Turn {i}] Assistant: {NEEDLE} fact {toks}")
else:
lines.append(f"[Turn {i}] Assistant: progress {i}: {toks}")
return "\n\n".join(lines)
async def _chunk_coverage(memory, bank_id, document_id, request_context):
doc = await memory.get_document(document_id, bank_id, request_context=request_context)
assert doc is not None
original_len = len(doc["original_text"])
chunks = await memory.list_document_chunks(bank_id, document_id, limit=10000, request_context=request_context)
items = chunks["items"]
sum_chunk_text = sum(len(c["chunk_text"]) for c in items)
indices = sorted(c["chunk_index"] for c in items)
needle_present = any(NEEDLE in c["chunk_text"] for c in items)
return original_len, sum_chunk_text, indices, needle_present
def _assert_full_coverage(label, original_len, sum_chunk_text, indices, needle_present):
# Sanity: the body must actually fan out to many chunks across several
# multi-chunk slices, or the test wouldn't exercise the bug at all.
assert len(indices) >= 16, f"{label}: only {len(indices)} chunks — body too small to exercise multi-chunk slices"
assert sum_chunk_text >= original_len * 0.9, (
f"{label}: chunks cover only {sum_chunk_text}/{original_len} chars "
f"(~{100 * sum_chunk_text // original_len}%) — a sub-batch slice was overwritten or skipped"
)
assert indices == list(range(len(indices))), (
f"{label}: chunk_index sequence is not contiguous: {indices} — sub-batch slices collided on chunk_id"
)
assert needle_present, f"{label}: the late-slice needle fact was dropped (offset collision or recovery skip)"
@pytest.mark.asyncio
async def test_sync_inline_multichunk_subbatch_coverage(memory, request_context):
"""Sync inline path (retain_batch_async): an oversized doc whose slices each
span several extraction chunks must keep full coverage (offset bug)."""
bank_id = f"test_multichunk_sync_{_ts()}"
document_id = "doc-multichunk-sync"
try:
body = _make_body()
await memory.retain_batch_async(
bank_id=bank_id,
contents=[{"content": body, "context": "big doc", "document_id": document_id}],
request_context=request_context,
)
cov = await _chunk_coverage(memory, bank_id, document_id, request_context)
_assert_full_coverage("sync", *cov)
finally:
await memory.delete_bank(bank_id, request_context=request_context)
@pytest.mark.asyncio
@pytest.mark.timeout(600)
async def test_async_submit_multichunk_subbatch_coverage(memory, request_context):
"""Async submit path (submit_async_retain → child op → worker): the same
oversized doc must keep full coverage too. Exercises both the offset bug and
the shared-operation_id whole-document recovery skip."""
import asyncio
bank_id = f"test_multichunk_async_{_ts()}"
document_id = "doc-multichunk-async"
try:
body = _make_body()
result = await memory.submit_async_retain(
bank_id=bank_id,
contents=[{"content": body, "context": "big doc", "document_id": document_id}],
request_context=request_context,
)
operation_id = result["operation_id"]
# SyncTaskBackend (test backend) drains children inline; wait for the
# parent to reach a terminal state before reading chunks.
status = None
for _ in range(600):
status = await memory.get_operation_status(
bank_id=bank_id, operation_id=operation_id, request_context=request_context
)
if status["status"] in ("completed", "failed"):
break
await asyncio.sleep(0.1)
assert status is not None and status["status"] == "completed", (
f"async retain did not complete: {status['status'] if status else 'no status'}"
)
cov = await _chunk_coverage(memory, bank_id, document_id, request_context)
_assert_full_coverage("async", *cov)
finally:
await memory.delete_bank(bank_id, request_context=request_context)
@@ -23,7 +23,6 @@ from hindsight_api.engine.search.tags import (
TagGroupNot,
TagGroupOr,
build_tag_groups_where_clause,
build_tags_where_clause,
build_tags_where_clause_simple,
filter_results_by_tag_groups,
filter_results_by_tags,
@@ -137,46 +136,6 @@ class TestTagsWhereClauseBuilder:
assert "@>" in result
assert "<@" in result
# ---- Test "exact" mode with the empty scope ([]) = untagged/global only ----
def test_tags_match_exact_empty_list_matches_untagged_only(self):
"""match='exact' with [] filters to untagged rows only (no bind param)."""
result = build_tags_where_clause_simple([], 5, match="exact")
assert "IS NULL" in result
assert "= '{}'" in result
# Untagged-only is param-free: callers append no tags param for an empty list.
assert "$5" not in result
# Must not use set-equality operators (which would need a bound scope).
assert "@>" not in result
assert "<@" not in result
def test_tags_match_exact_empty_list_with_table_alias(self):
"""Empty-scope exact clause respects the table alias."""
result = build_tags_where_clause_simple([], 5, table_alias="mu.", match="exact")
assert "mu.tags IS NULL" in result
assert "mu.tags = '{}'" in result
def test_tags_match_exact_none_matches_untagged_only(self):
"""match='exact' with None (no tags) selects the global scope, like the graph endpoint."""
result = build_tags_where_clause_simple(None, 5, match="exact")
assert "IS NULL" in result
assert "= '{}'" in result
assert "$5" not in result
def test_tags_match_any_empty_list_still_no_filter(self):
"""Empty list only filters under 'exact'; other modes treat [] as no filter."""
assert build_tags_where_clause_simple([], 5, match="any") == ""
assert build_tags_where_clause_simple([], 5, match="any_strict") == ""
@pytest.mark.parametrize("tags", [None, []])
def test_tags_where_clause_exact_empty_scope_keeps_param_offset(self, tags):
"""The parameterized builder must not consume a bind index for the empty scope,
so following clauses stay aligned with their params."""
clause, params, next_offset = build_tags_where_clause(tags, param_offset=4, match="exact")
assert clause == "AND (tags IS NULL OR tags = '{}')"
assert params == []
assert next_offset == 4
# ---- Test table alias with all modes ----
def test_tags_match_any_with_table_alias(self):
@@ -296,20 +255,6 @@ class TestFilterResultsByTags:
assert len(filtered) == 1
assert filtered[0].tags == ["a"]
def test_exact_mode_empty_scope_matches_untagged_only(self):
"""'exact' mode with [] should keep only untagged results (NULL or empty)."""
results = [MockResult(["a"]), MockResult(["a", "b"]), MockResult(None), MockResult([])]
filtered = filter_results_by_tags(results, [], match="exact")
assert len(filtered) == 2
assert all(not r.tags for r in filtered)
def test_exact_mode_none_matches_untagged_only(self):
"""'exact' mode with None (no tags) selects the global scope (untagged only)."""
results = [MockResult(["a"]), MockResult(None), MockResult([])]
filtered = filter_results_by_tags(results, None, match="exact")
assert len(filtered) == 2
assert all(not r.tags for r in filtered)
def test_all_mode_includes_untagged(self):
"""'all' mode should include untagged results."""
results = [MockResult(["a", "b"]), MockResult(None), MockResult([])]
@@ -557,16 +502,6 @@ class TestBuildTagGroupsWhereClause:
assert len(params) == 2
assert next_offset == 3
def test_exact_leaf_empty_scope_matches_untagged_only(self):
"""An exact leaf with [] becomes an untagged-only clause with no bind param."""
groups = [TagGroupLeaf(tags=[], match="exact")]
clause, params, next_offset = build_tag_groups_where_clause(groups, 5)
assert "IS NULL" in clause
assert "= '{}'" in clause
assert "$5" not in clause # param-free
assert params == []
assert next_offset == 5 # offset unchanged — no param consumed
# ============================================================================
# Unit Tests for filter_results_by_tag_groups (Python-side)
@@ -596,14 +531,6 @@ class TestFilterResultsByTagGroups:
assert len(filtered) == 1
assert filtered[0].tags == ["step:5"]
def test_exact_leaf_empty_scope_matches_untagged_only(self):
"""An exact leaf with [] keeps only untagged results (matches SQL builder)."""
groups = [TagGroupLeaf(tags=[], match="exact")]
results = [MockResult(["a"]), MockResult(["a", "b"]), MockResult(None), MockResult([])]
filtered = filter_results_by_tag_groups(results, groups)
assert len(filtered) == 2
assert all(not r.tags for r in filtered)
def test_single_leaf_all_strict_matches_superset(self):
"""Single all_strict leaf matches results that contain all tags."""
groups = [TagGroupLeaf(tags=["user:alice", "step:5"], match="all_strict")]
@@ -977,37 +904,6 @@ async def test_recall_with_empty_tags_returns_all(api_client, test_bank_id):
assert any("Rachel" in t for t in texts), "Should find Rachel"
@pytest.mark.asyncio
async def test_recall_empty_tags_exact_returns_untagged_only(api_client, test_bank_id):
"""tags=[] with tags_match='exact' returns only untagged/global memories."""
# One untagged (global) memory and one tagged memory.
response = await api_client.post(
f"/v1/default/banks/{test_bank_id}/memories",
json={
"items": [
{"content": "Sam studies astronomy."}, # no tags -> global scope
{"content": "Tina studies geology.", "tags": ["user_tina"]},
]
},
)
assert response.status_code == 200
# exact match on the empty scope -> only the untagged memory.
response = await api_client.post(
f"/v1/default/banks/{test_bank_id}/memories/recall",
json={"query": "Who studies what?", "budget": "low", "tags": [], "tags_match": "exact"},
)
assert response.status_code == 200
results = response.json()["results"]
texts = [r["text"] for r in results]
assert any("Sam" in t for t in texts), "Should find the untagged memory"
assert not any("Tina" in t for t in texts), "Should NOT find the tagged memory"
# Every returned memory must be untagged.
for r in results:
assert not r.get("tags"), f"Expected untagged result, got tags={r.get('tags')}"
@pytest.mark.asyncio
async def test_multi_user_agent_visibility(api_client):
"""
@@ -1,162 +0,0 @@
"""Tests for cached / thoughts token propagation through TokenUsage,
LLMToolCallResult, TokenUsageSummary, and RetainResult.
The Gemini 2.5+ family (and any future provider with prompt caching +
reasoning tokens) reports four distinct token counts on every response:
prompt, candidates (visible output), cached_content, and thoughts. The
last two are billed separately by the provider but were previously not
threaded through to downstream return contexts, so application-layer
metering had no way to attribute prompt-cache hit rate or reasoning cost
per operation.
These tests pin the propagation: when a provider populates cached or
thoughts on the way out, every accumulator and aggregate type carries
the value through unchanged.
"""
from __future__ import annotations
import pytest
from hindsight_api.engine.reflect.agent import _generate_structured_output
from hindsight_api.engine.reflect.models import StructuredOutputResult, TokenUsageSummary
from hindsight_api.engine.response_models import LLMToolCallResult, TokenUsage
from hindsight_api.extensions.operation_validator import RetainResult
def test_token_usage_carries_cached_and_thoughts():
"""TokenUsage defaults both new fields to 0 and accepts non-zero values."""
u = TokenUsage(input_tokens=1500, output_tokens=500, total_tokens=2000)
assert u.cached_tokens == 0
assert u.thoughts_tokens == 0
u = TokenUsage(
input_tokens=1500,
output_tokens=500,
total_tokens=2000,
cached_tokens=200,
thoughts_tokens=80,
)
assert u.cached_tokens == 200
assert u.thoughts_tokens == 80
def test_token_usage_aggregates_thoughts_tokens():
"""TokenUsage.__add__ sums thoughts_tokens alongside the existing fields.
Multi-iteration agentic loops accumulate per-call usage via ``+``. If
thoughts_tokens isn't summed, the per-op total undercounts reasoning
spend by a factor of N (the number of LLM sub-calls).
"""
a = TokenUsage(input_tokens=10, output_tokens=5, total_tokens=15, cached_tokens=2, thoughts_tokens=7)
b = TokenUsage(input_tokens=20, output_tokens=8, total_tokens=28, cached_tokens=3, thoughts_tokens=11)
c = a + b
assert c.input_tokens == 30
assert c.output_tokens == 13
assert c.total_tokens == 43
assert c.cached_tokens == 5
assert c.thoughts_tokens == 18
def test_llm_tool_call_result_carries_cached_and_thoughts():
"""call_with_tools returns LLMToolCallResult — both new fields default to 0
and accept non-zero values from the provider."""
r = LLMToolCallResult(content="ok", input_tokens=1234, output_tokens=56)
assert r.cached_tokens == 0
assert r.thoughts_tokens == 0
r = LLMToolCallResult(
content="ok",
input_tokens=1234,
output_tokens=56,
cached_tokens=200,
thoughts_tokens=78,
)
assert r.cached_tokens == 200
assert r.thoughts_tokens == 78
def test_token_usage_summary_carries_cached_and_thoughts():
"""TokenUsageSummary is what reflect agent returns to its caller — needs
to propagate the aggregate so per-op cost attribution works."""
s = TokenUsageSummary(
input_tokens=10000,
output_tokens=200,
total_tokens=10200,
cached_tokens=3000,
thoughts_tokens=150,
)
assert s.cached_tokens == 3000
assert s.thoughts_tokens == 150
def test_token_usage_summary_defaults_cached_and_thoughts_to_zero():
"""Defaults preserve backward compatibility for callers built before the
fields existed."""
s = TokenUsageSummary(input_tokens=100, output_tokens=50, total_tokens=150)
assert s.cached_tokens == 0
assert s.thoughts_tokens == 0
def test_retain_result_carries_cached_input_and_thoughts():
"""RetainResult is the contract between the engine and any metering
extension. The two new fields are optional (None) so older extensions
that don't read them are unaffected; engines that DO populate them get
end-to-end attribution into the metering hook."""
class _Ctx:
pass
r = RetainResult(
bank_id="b",
contents=[],
request_context=_Ctx(),
document_id=None,
fact_type_override=None,
unit_ids=[],
llm_input_tokens=1000,
llm_output_tokens=50,
llm_total_tokens=1050,
llm_cached_input_tokens=300,
llm_thoughts_tokens=25,
)
assert r.llm_cached_input_tokens == 300
assert r.llm_thoughts_tokens == 25
# Defaults stay None for engines that don't surface the data, so
# downstream extensions can use ``or 0`` without breaking on a
# core-only build.
r2 = RetainResult(
bank_id="b",
contents=[],
request_context=_Ctx(),
document_id=None,
fact_type_override=None,
unit_ids=[],
)
assert r2.llm_cached_input_tokens is None
assert r2.llm_thoughts_tokens is None
@pytest.mark.asyncio
async def test_generate_structured_output_returns_dataclass_on_no_fields():
"""_generate_structured_output returns a StructuredOutputResult, not a tuple.
Regression guard: the function and all six call sites must agree on a single
return type. A previous tuple-based contract drifted out of sync (the failure
branch returned 3 values while callers unpacked 5), which would crash reflect
with a ValueError on any structured-output failure. An empty schema exercises
the no-LLM-call branch deterministically.
"""
result = await _generate_structured_output(
answer="anything",
response_schema={},
llm_config=None,
reflect_id="test",
)
assert isinstance(result, StructuredOutputResult)
assert result.structured_output is None
assert result.input_tokens == 0
assert result.output_tokens == 0
assert result.cached_tokens == 0
assert result.thoughts_tokens == 0
+10 -141
View File
@@ -13,7 +13,6 @@ Tests cover:
import asyncio
import json
import uuid
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
import pytest_asyncio
@@ -71,99 +70,6 @@ async def clean_operations(pool):
)
def test_metric_operation_label_normalises_retain_variants():
"""Worker completion metrics collapse retain variants onto operation="retain"
so they share the API path's series; other types pass through unchanged."""
from hindsight_api.worker.poller import _metric_operation_label
assert _metric_operation_label("retain") == "retain"
assert _metric_operation_label("batch_retain") == "retain"
assert _metric_operation_label("file_convert_retain") == "retain"
assert _metric_operation_label("consolidation") == "consolidation"
assert _metric_operation_label("reflect") == "reflect"
assert _metric_operation_label(None) == "unknown"
class TestWorkerOperationMetrics:
"""_execute_task_inner emits operation metrics on terminal outcomes only (no DB)."""
def _make_poller(self, executor):
from hindsight_api.worker import WorkerPoller
poller = WorkerPoller(backend=MagicMock(), worker_id="w-test", executor=executor)
# Stub terminal-state handlers so _execute_task_inner never touches the DB.
poller._mark_failed = AsyncMock()
poller._defer_operation = AsyncMock()
poller._schedule_retry = AsyncMock()
return poller
async def _run(self, executor, task_type="batch_retain"):
from hindsight_api.worker.poller import ClaimedTask
poller = self._make_poller(executor)
task = ClaimedTask(
operation_id=str(uuid.uuid4()),
task_dict={"type": task_type, "operation_type": task_type, "bank_id": "bank-1"},
schema=None,
)
collector = MagicMock()
with patch("hindsight_api.worker.poller.get_metrics_collector", return_value=collector):
await poller._execute_task_inner(task)
return collector
@pytest.mark.asyncio
async def test_executor_returning_normally_records_success(self):
"""Success is inferred from the executor returning without raising to the
poller. This deliberately includes deterministic failures that
memory_engine.execute_task handles itself and returns from normally
(file_convert_retain, non-retryable errors) at the poller boundary they
are indistinguishable from a clean completion, so they also record
success=true. The worker counter is therefore a completion-throughput
signal; authoritative failure visibility comes from the
hindsight_async_operations{status="failed"} gauge, which reads each
operation's final DB status.
"""
collector = await self._run(AsyncMock()) # executor returns normally
collector.record_operation_result.assert_called_once()
call = collector.record_operation_result.call_args
assert call.args[0] == "retain" # batch_retain normalised
assert call.kwargs["success"] is True
assert call.kwargs["source"] == "worker"
@pytest.mark.asyncio
async def test_failure_records_failure(self):
async def boom(_):
raise RuntimeError("kaboom")
collector = await self._run(boom)
collector.record_operation_result.assert_called_once()
assert collector.record_operation_result.call_args.kwargs["success"] is False
@pytest.mark.asyncio
async def test_deferral_not_counted(self):
from datetime import datetime, timezone
from hindsight_api.worker.exceptions import DeferOperation
async def defer(_):
raise DeferOperation(exec_date=datetime.now(timezone.utc), reason="later")
collector = await self._run(defer)
collector.record_operation_result.assert_not_called()
@pytest.mark.asyncio
async def test_retry_not_counted(self):
from datetime import datetime, timezone
from hindsight_api.worker.exceptions import RetryTaskAt
async def retry(_):
raise RetryTaskAt(retry_at=datetime.now(timezone.utc), message="transient")
collector = await self._run(retry)
collector.record_operation_result.assert_not_called()
def test_all_operation_types_have_slot_reservation_config():
"""Every operation_type used in memory_engine must be listed in
WORKER_SLOT_RESERVATION_TYPES so it can be reserved via env var.
@@ -926,34 +832,6 @@ class TestWorkerPoller:
# Defensive: confirm it wasn't a RetryTaskAt masquerading as Defer.
assert not isinstance(exc_info.value, RetryTaskAt)
@pytest.mark.asyncio
async def test_memory_engine_provider_quota_reset_becomes_defer_operation(self, memory, monkeypatch):
"""Provider quota windows should park worker tasks until the reset time."""
from datetime import UTC, datetime, timedelta
from hindsight_api.engine.llm_interface import ProviderRateLimitResetError
from hindsight_api.worker.exceptions import DeferOperation, RetryTaskAt
retry_at = (datetime.now(UTC) + timedelta(hours=5)).replace(microsecond=0)
async def quota_limited_retain(_task_dict: object) -> None:
raise ProviderRateLimitResetError(retry_at=retry_at, message="quota resets later")
monkeypatch.setattr(memory, "_handle_batch_retain", quota_limited_retain)
with pytest.raises(DeferOperation) as exc_info:
await memory.execute_task(
{
"type": "batch_retain",
"bank_id": "test-provider-quota-defer",
"contents": [{"content": "x"}],
}
)
assert exc_info.value.exec_date == retry_at
assert exc_info.value.reason == "quota resets later"
assert not isinstance(exc_info.value, RetryTaskAt)
@pytest.mark.asyncio
async def test_claim_batch_skips_consolidation_when_same_bank_processing(self, pool, backend, clean_operations):
"""Test that pending consolidation is skipped if same bank has one processing."""
@@ -1069,11 +947,9 @@ class TestWorkerPoller:
claimed = await poller.claim_batch()
# Should claim the retain task (non-consolidation tasks are unaffected).
# Filter to our bank — parallel tests may contribute other claims.
my_claims = [c for c in claimed if c.task_dict.get("bank_id") == bank_id]
assert len(my_claims) == 1, f"Expected 1 claim for our bank, got {len(my_claims)}"
assert my_claims[0].operation_id == str(retain_op_id)
# Should claim the retain task (non-consolidation tasks are unaffected)
assert len(claimed) == 1
assert claimed[0].operation_id == str(retain_op_id)
class TestWorkerRecovery:
@@ -1710,9 +1586,7 @@ class TestDynamicTenantDiscovery:
# First claim_batch should call list_tenants
claimed1 = await poller.claim_batch()
assert mock_extension.list_tenants_calls == 1
# Filter to our bank — parallel tests may contribute other claims.
my_claims1 = [c for c in claimed1 if c.task_dict.get("bank_id") == bank_id]
assert len(my_claims1) == 2, f"Expected 2 claims for our bank, got {len(my_claims1)}"
assert len(claimed1) == 2
# Add more tasks
for i in range(2):
@@ -1731,8 +1605,7 @@ class TestDynamicTenantDiscovery:
# Second claim_batch should call list_tenants again
claimed2 = await poller.claim_batch()
assert mock_extension.list_tenants_calls == 2
my_claims2 = [c for c in claimed2 if c.task_dict.get("bank_id") == bank_id]
assert len(my_claims2) == 2, f"Expected 2 claims for our bank, got {len(my_claims2)}"
assert len(claimed2) == 2
@pytest.mark.asyncio
async def test_poller_picks_up_new_tenants_without_restart(self, pool, backend, clean_operations):
@@ -1777,12 +1650,10 @@ class TestDynamicTenantDiscovery:
tenant_extension=dynamic_extension,
)
# First poll - only public schema. Filter to our bank — parallel tests
# may contribute other claims.
# First poll - only public schema
claimed1 = await poller.claim_batch()
my_claims1 = [c for c in claimed1 if c.task_dict.get("bank_id") == bank_id]
assert len(my_claims1) == 1, f"Expected 1 claim for our bank, got {len(my_claims1)}"
assert my_claims1[0].schema is None # public is represented as None
assert len(claimed1) == 1
assert claimed1[0].schema is None # public is represented as None
assert dynamic_extension.list_tenants_calls == 1
# Simulate tenant list changing (but we won't add a non-existent schema)
@@ -1804,14 +1675,12 @@ class TestDynamicTenantDiscovery:
# Second poll - list_tenants should be called again
claimed2 = await poller.claim_batch()
my_claims2 = [c for c in claimed2 if c.task_dict.get("bank_id") == bank_id]
assert len(my_claims2) == 1, f"Expected 1 claim for our bank, got {len(my_claims2)}"
assert len(claimed2) == 1
assert dynamic_extension.list_tenants_calls == 2 # Called again on second poll
# Third poll with no tasks - still calls list_tenants
claimed3 = await poller.claim_batch()
my_claims3 = [c for c in claimed3 if c.task_dict.get("bank_id") == bank_id]
assert len(my_claims3) == 0, f"Expected 0 claims for our bank, got {len(my_claims3)}"
assert len(claimed3) == 0
assert dynamic_extension.list_tenants_calls == 3 # Called again even with no tasks
@pytest.mark.asyncio
+2 -2
View File
@@ -4,12 +4,12 @@ build-backend = "setuptools.build_meta"
[project]
name = "hindsight-api"
version = "0.8.3"
version = "0.8.2"
description = "Hindsight: Agent Memory That Works Like Human Memory"
readme = "README.md"
requires-python = ">=3.11"
dependencies = [
"hindsight-api-slim[all]==0.8.3",
"hindsight-api-slim[all]==0.8.2",
]
[tool.uv.sources]
+1 -1
View File
@@ -1,6 +1,6 @@
[package]
name = "hindsight-cli"
version = "0.8.3"
version = "0.8.2"
edition = "2021"
authors = ["Hindsight Team"]
description = "A beautiful CLI for Hindsight - semantic memory system"
+7 -29
View File
@@ -425,8 +425,8 @@ impl ApiClient {
.client
.list_documents(
agent_id,
limit.map(|l| l as u64),
offset.map(|o| o as u64),
limit.map(|l| l as i64),
offset.map(|o| o as i64),
q,
None,
None,
@@ -524,8 +524,8 @@ impl ApiClient {
bank_id,
None, // consolidation_state
None, // document_id
limit.map(|l| l as u64),
offset.map(|o| o as u64),
limit,
offset,
q,
None, // state
type_filter,
@@ -546,12 +546,7 @@ impl ApiClient {
self.runtime.block_on(async {
let response = self
.client
.list_entities(
bank_id,
limit.map(|l| l as u64),
offset.map(|o| o as u64),
None,
)
.list_entities(bank_id, limit, offset, None)
.await?;
Ok(response.into_inner())
})
@@ -669,17 +664,7 @@ impl ApiClient {
self.runtime.block_on(async {
let response = self
.client
.get_graph(
bank_id,
None,
None,
limit.map(|l| l as u64),
None,
None,
None,
type_filter,
None,
)
.get_graph(bank_id, None, None, limit, None, None, None, type_filter, None)
.await?;
Ok(response.into_inner())
})
@@ -741,14 +726,7 @@ impl ApiClient {
self.runtime.block_on(async {
let response = self
.client
.list_tags(
bank_id,
limit.map(|l| l as u64),
offset.map(|o| o as u64),
q,
None,
None,
)
.list_tags(bank_id, limit, offset, q, None, None)
.await?;
Ok(response.into_inner())
})
-1
View File
@@ -340,7 +340,6 @@ impl App {
max_tokens: query_max_tokens,
trace: false,
query_timestamp: None,
prefer_observations: false,
include: None,
tags: None,
tags_match: TagsMatch::Any,
-2
View File
@@ -277,7 +277,6 @@ pub fn recall(
tags: Vec<String>,
tags_match: Option<String>,
query_timestamp: Option<String>,
prefer_observations: bool,
verbose: bool,
output_format: OutputFormat,
) -> Result<()> {
@@ -311,7 +310,6 @@ pub fn recall(
max_tokens,
trace,
query_timestamp,
prefer_observations,
include,
tags: if tags.is_empty() { None } else { Some(tags) },
tags_match: parse_tags_match(&tags_match),
-6
View File
@@ -532,10 +532,6 @@ enum MemoryCommands {
/// Reference timestamp for recall (ISO 8601, e.g. 2023-05-30T23:40:00)
#[arg(long)]
query_timestamp: Option<String>,
/// Prefer observations: drop raw facts a returned observation was consolidated from (no effect unless observation + a raw type are both recalled)
#[arg(long)]
prefer_observations: bool,
},
/// Generate answers using bank identity (reflect/reasoning)
@@ -1405,7 +1401,6 @@ fn run() -> Result<()> {
tags,
tags_match,
query_timestamp,
prefer_observations,
} => commands::memory::recall(
&client,
&bank_id,
@@ -1419,7 +1414,6 @@ fn run() -> Result<()> {
tags,
tags_match,
query_timestamp,
prefer_observations,
verbose,
output_format,
),
+5 -36
View File
@@ -7,7 +7,7 @@ info:
name: Apache 2.0
url: https://www.apache.org/licenses/LICENSE-2.0.html
title: Hindsight HTTP API
version: 0.8.3
version: 0.8.2
servers:
- url: /
paths:
@@ -80,7 +80,6 @@ paths:
required: false
schema:
default: 1000
minimum: 0
title: Limit
type: integer
style: form
@@ -213,7 +212,6 @@ paths:
required: false
schema:
default: 100
minimum: 0
title: Limit
type: integer
style: form
@@ -223,7 +221,6 @@ paths:
required: false
schema:
default: 0
minimum: 0
title: Offset
type: integer
style: form
@@ -735,7 +732,6 @@ paths:
schema:
default: 100
description: Maximum number of entities to return
minimum: 0
title: Limit
type: integer
style: form
@@ -747,7 +743,6 @@ paths:
schema:
default: 0
description: Offset for pagination
minimum: 0
title: Offset
type: integer
style: form
@@ -797,7 +792,6 @@ paths:
schema:
default: 1000
description: Maximum number of co-occurrence edges to return
minimum: 0
title: Limit
type: integer
style: form
@@ -1696,7 +1690,6 @@ paths:
required: false
schema:
default: 100
minimum: 0
title: Limit
type: integer
style: form
@@ -1706,7 +1699,6 @@ paths:
required: false
schema:
default: 0
minimum: 0
title: Offset
type: integer
style: form
@@ -2054,7 +2046,6 @@ paths:
schema:
default: 100
description: Maximum number of tags to return
minimum: 0
title: Limit
type: integer
style: form
@@ -2066,7 +2057,6 @@ paths:
schema:
default: 0
description: Offset for pagination
minimum: 0
title: Offset
type: integer
style: form
@@ -3551,7 +3541,7 @@ paths:
This endpoint handles file upload, conversion, and memory creation in a single operation.
**Features:**
- Supports PDF, DOCX, PPTX, XLSX, images (parser-dependent OCR), audio (with transcription)
- Supports PDF, DOCX, PPTX, XLSX, images (with OCR), audio (with transcription)
- Automatic file-to-markdown conversion using pluggable parsers
- Files stored in object storage (PostgreSQL by default, S3 for production)
- Each file becomes a separate document with optional metadata/tags
@@ -7123,17 +7113,6 @@ components:
type: string
nullable: true
type: array
prefer_observations:
default: false
description: "When recalling raw facts ('world'/'experience') together with\
\ 'observation', drop any raw fact that an observation in the results\
\ was consolidated from, so the observation supersedes it and you don't\
\ get duplicate content. The freed slots are backfilled with the next\
\ results, keeping the result count at the requested budget. Disabled\
\ by default; set to true to enable. No effect unless 'observation' and\
\ at least one raw type are both requested."
title: Prefer Observations
type: boolean
budget:
$ref: '#/components/schemas/Budget'
max_tokens:
@@ -7158,9 +7137,7 @@ components:
default: any
description: "How to match tags: 'any' (OR, includes untagged), 'all' (AND,\
\ includes untagged), 'any_strict' (OR, excludes untagged), 'all_strict'\
\ (AND, excludes untagged), 'exact' (set-equality on the full scope, excludes\
\ untagged). With 'exact' and no tags (or []), the empty global scope\
\ is selected and only untagged memories match."
\ (AND, excludes untagged)."
enum:
- any
- all
@@ -7866,13 +7843,12 @@ components:
type: integer
output_tokens:
default: 0
description: Number of visible output/completion tokens generated (excludes
reasoning/thoughts)
description: Number of output/completion tokens generated
title: Output Tokens
type: integer
total_tokens:
default: 0
description: "Total tokens (input + output, excludes thoughts)"
description: Total tokens (input + output)
title: Total Tokens
type: integer
cached_tokens:
@@ -7880,13 +7856,6 @@ components:
description: "Cached/cache-read prompt tokens, when reported by the provider"
title: Cached Tokens
type: integer
thoughts_tokens:
default: 0
description: Reasoning/thinking tokens generated by the model. Billed at
the output rate by some providers (e.g. Gemini 2.5+ family) but not surfaced
in the visible response.
title: Thoughts Tokens
type: integer
title: TokenUsage
ToolCallsIncludeOptions:
description: Options for including tool calls in reflect results.
+1 -1
View File
@@ -3,7 +3,7 @@ Hindsight HTTP API
HTTP API for Hindsight
API version: 0.8.3
API version: 0.8.2
*/
// Code generated by OpenAPI Generator (https://openapi-generator.tech); DO NOT EDIT.
+1 -1
View File
@@ -3,7 +3,7 @@ Hindsight HTTP API
HTTP API for Hindsight
API version: 0.8.3
API version: 0.8.2
*/
// Code generated by OpenAPI Generator (https://openapi-generator.tech); DO NOT EDIT.
+1 -1
View File
@@ -3,7 +3,7 @@ Hindsight HTTP API
HTTP API for Hindsight
API version: 0.8.3
API version: 0.8.2
*/
// Code generated by OpenAPI Generator (https://openapi-generator.tech); DO NOT EDIT.
+1 -1
View File
@@ -3,7 +3,7 @@ Hindsight HTTP API
HTTP API for Hindsight
API version: 0.8.3
API version: 0.8.2
*/
// Code generated by OpenAPI Generator (https://openapi-generator.tech); DO NOT EDIT.

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