Compare commits
60
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
4e37381eba | ||
|
|
ee81c65e4b | ||
|
|
ccd3eb24c9 | ||
|
|
51cb32896f | ||
|
|
af42382983 | ||
|
|
955b0c523c | ||
|
|
80281e2543 | ||
|
|
bb4dd4f393 | ||
|
|
ba5ddd59af | ||
|
|
aab7032071 | ||
|
|
adb6dcd683 | ||
|
|
ae93c93182 | ||
|
|
1c0c53c062 | ||
|
|
731add1fdf | ||
|
|
a7f82453a3 | ||
|
|
2bd6be8d85 | ||
|
|
e1014cc790 | ||
|
|
da2125cf13 | ||
|
|
f4bac2d41d | ||
|
|
39abf0ad3f | ||
|
|
8426b0c359 | ||
|
|
f4a0a31f70 | ||
|
|
65862c4fef | ||
|
|
ef548833fd | ||
|
|
55f70e1d27 | ||
|
|
b8cfddd7b6 | ||
|
|
52cb9a2bae | ||
|
|
539101af38 | ||
|
|
efa37cb15f | ||
|
|
faaa97d4a0 | ||
|
|
70804fe2c6 | ||
|
|
ae2532b165 | ||
|
|
81865bf873 | ||
|
|
9e47759347 | ||
|
|
d8665d7ab0 | ||
|
|
9bde15331e | ||
|
|
2fb2de1aa8 | ||
|
|
d68bd07423 | ||
|
|
44972d3215 | ||
|
|
aa308ad201 | ||
|
|
cb73790c27 | ||
|
|
b1fe23fbe4 | ||
|
|
ce81217381 | ||
|
|
94619ce52b | ||
|
|
27aa6bbf46 | ||
|
|
acf4d5c860 | ||
|
|
551932991d | ||
|
|
5ee53c512f | ||
|
|
4efa204727 | ||
|
|
c3bb647640 | ||
|
|
12851bc7ee | ||
|
|
a5f4d30ea6 | ||
|
|
0c9bc765ce | ||
|
|
cd34efa596 | ||
|
|
2b521c3a09 | ||
|
|
9681d96195 | ||
|
|
a32ecfeb33 | ||
|
|
bf73a1dfbe | ||
|
|
ca2ce5c16d | ||
|
|
7b17da7a0c |
@@ -87,6 +87,18 @@ 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
|
||||
@@ -150,6 +162,10 @@ 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)
|
||||
|
||||
@@ -1,6 +0,0 @@
|
||||
version: 2
|
||||
updates:
|
||||
- package-ecosystem: "github-actions"
|
||||
directory: "/"
|
||||
schedule:
|
||||
interval: "weekly"
|
||||
@@ -45,17 +45,20 @@ 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 }}
|
||||
@@ -154,6 +157,8 @@ jobs:
|
||||
- 'hindsight-integrations/ag2/**'
|
||||
integrations-autogen:
|
||||
- 'hindsight-integrations/autogen/**'
|
||||
integrations-aider:
|
||||
- 'hindsight-integrations/aider/**'
|
||||
integrations-langgraph:
|
||||
- 'hindsight-integrations/langgraph/**'
|
||||
integrations-llamaindex:
|
||||
@@ -166,6 +171,8 @@ jobs:
|
||||
- 'hindsight-integrations/opencode/**'
|
||||
integrations-cursor:
|
||||
- 'hindsight-integrations/cursor/**'
|
||||
integrations-zed:
|
||||
- 'hindsight-integrations/zed/**'
|
||||
integrations-n8n:
|
||||
- 'hindsight-integrations/n8n/**'
|
||||
integrations-zapier:
|
||||
@@ -180,6 +187,8 @@ 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:
|
||||
@@ -488,6 +497,37 @@ 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: >-
|
||||
@@ -3021,6 +3061,45 @@ 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: >-
|
||||
@@ -3669,6 +3748,45 @@ 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: >-
|
||||
@@ -4703,6 +4821,7 @@ jobs:
|
||||
- test-openclaw-integration
|
||||
- test-integration
|
||||
- test-ag2-integration
|
||||
- test-aider-integration
|
||||
- test-autogen-integration
|
||||
- test-continue-integration
|
||||
- test-smolagents-integration
|
||||
@@ -4717,6 +4836,7 @@ jobs:
|
||||
- test-pydantic-ai-integration
|
||||
- test-llamaindex-integration
|
||||
- test-openai-agents-integration
|
||||
- test-openhands-integration
|
||||
- test-agentcore-integration
|
||||
- test-haystack-integration
|
||||
- test-pip-slim
|
||||
|
||||
@@ -2,8 +2,8 @@ apiVersion: v2
|
||||
name: hindsight
|
||||
description: Hindsight helm chart
|
||||
type: application
|
||||
version: 0.8.2
|
||||
appVersion: "0.8.2"
|
||||
version: 0.8.3
|
||||
appVersion: "0.8.3"
|
||||
keywords:
|
||||
- ai
|
||||
- memory
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@vectorize-io/hindsight-all",
|
||||
"version": "0.8.2",
|
||||
"version": "0.8.3",
|
||||
"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",
|
||||
|
||||
@@ -4,12 +4,12 @@ build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "hindsight-all-slim"
|
||||
version = "0.8.2"
|
||||
version = "0.8.3"
|
||||
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.2",
|
||||
"hindsight-api-slim==0.8.3",
|
||||
"hindsight-client>=0.0.7",
|
||||
"hindsight-embed>=0.1.0",
|
||||
]
|
||||
|
||||
@@ -4,12 +4,12 @@ build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "hindsight-all"
|
||||
version = "0.8.2"
|
||||
version = "0.8.3"
|
||||
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.2",
|
||||
"hindsight-api-slim[all]==0.8.3",
|
||||
"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.2",
|
||||
"hindsight-api-slim[local-llm]==0.8.3",
|
||||
]
|
||||
test = [
|
||||
"pytest>=7.0.0",
|
||||
|
||||
@@ -53,4 +53,4 @@ __all__ = [
|
||||
"RemoteTEICrossEncoder",
|
||||
"LLMConfig",
|
||||
]
|
||||
__version__ = "0.8.2"
|
||||
__version__ = "0.8.3"
|
||||
|
||||
@@ -256,6 +256,7 @@ 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
|
||||
@@ -292,7 +293,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=True,
|
||||
ensure_extensions=ensure_extensions,
|
||||
)
|
||||
|
||||
return schemas
|
||||
@@ -311,6 +312,18 @@ 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()
|
||||
@@ -324,6 +337,8 @@ 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(
|
||||
@@ -331,6 +346,7 @@ def run_db_migration(
|
||||
schema=schema,
|
||||
base_schema=config.database_schema,
|
||||
embedding_dimension=embedding_dimension,
|
||||
ensure_extensions=not skip_extension_reconcile,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
+158
@@ -0,0 +1,158 @@
|
||||
"""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)
|
||||
@@ -158,7 +158,12 @@ 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
|
||||
from hindsight_api.metrics import (
|
||||
create_metrics_collector,
|
||||
get_metrics_collector,
|
||||
initialize_metrics,
|
||||
normalize_http_endpoint,
|
||||
)
|
||||
from hindsight_api.models import RequestContext
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -1442,6 +1447,13 @@ 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."""
|
||||
@@ -3237,15 +3249,9 @@ def create_app(
|
||||
@app.middleware("http")
|
||||
async def http_metrics_middleware(request, call_next):
|
||||
"""Record HTTP request metrics."""
|
||||
# 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)
|
||||
# Template id segments (bank ids, UUIDs, numeric ids) so the endpoint
|
||||
# metric label stays bounded-cardinality.
|
||||
path = normalize_http_endpoint(request.url.path)
|
||||
|
||||
status_code = [500] # Default to 500, will be updated
|
||||
metrics_collector = get_metrics_collector()
|
||||
@@ -3333,6 +3339,7 @@ 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)
|
||||
@@ -3341,10 +3348,20 @@ 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:
|
||||
@@ -6749,7 +6766,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 (with OCR), audio (with transcription)\n"
|
||||
"- Supports PDF, DOCX, PPTX, XLSX, images (parser-dependent 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"
|
||||
|
||||
@@ -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 _get_raw_config
|
||||
from hindsight_api.config import DEFAULT_MCP_RECALL_DESCRIPTION, DEFAULT_MCP_RETAIN_DESCRIPTION, _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,6 +78,19 @@ 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.
|
||||
@@ -135,6 +148,10 @@ 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,
|
||||
@@ -144,6 +161,8 @@ 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)
|
||||
|
||||
@@ -142,6 +142,7 @@ 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"
|
||||
@@ -159,11 +160,25 @@ 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"
|
||||
@@ -354,6 +369,7 @@ 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"
|
||||
@@ -375,6 +391,7 @@ 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"
|
||||
@@ -424,6 +441,11 @@ 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"
|
||||
@@ -804,6 +826,7 @@ 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
|
||||
@@ -847,6 +870,10 @@ 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
|
||||
@@ -965,6 +992,7 @@ 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
|
||||
@@ -1298,6 +1326,7 @@ 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}})
|
||||
@@ -1478,6 +1507,7 @@ 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
|
||||
@@ -1642,6 +1672,7 @@ 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
|
||||
@@ -1676,6 +1707,11 @@ 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
|
||||
|
||||
@@ -1716,6 +1752,8 @@ 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",
|
||||
}
|
||||
@@ -1879,6 +1917,9 @@ 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"
|
||||
@@ -1996,6 +2037,11 @@ 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"),
|
||||
@@ -2345,6 +2391,7 @@ 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()
|
||||
@@ -2424,6 +2471,18 @@ 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,
|
||||
@@ -2614,6 +2673,8 @@ 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,6 +8,7 @@ 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
|
||||
@@ -161,26 +162,83 @@ class ConfigResolver:
|
||||
resolved_config = await self.resolve_full_config(bank_id, context)
|
||||
config_dict = asdict(resolved_config)
|
||||
|
||||
# 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}
|
||||
# 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: 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}
|
||||
def _strip_static_and_credential_fields(self, config_dict: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Keep only configurable, non-credential fields.
|
||||
|
||||
# PERMISSIONS: Further filter based on tenant/bank permissions
|
||||
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)
|
||||
if self.tenant_extension and context:
|
||||
try:
|
||||
allowed_fields = await self.tenant_extension.get_allowed_config_fields(context, bank_id)
|
||||
if allowed_fields is not None: # None means "allow all"
|
||||
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"
|
||||
)
|
||||
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})
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to load permissions for bank {bank_id}: {e}")
|
||||
logger.warning(f"Failed to load tenant config for bulk resolve: {e}")
|
||||
|
||||
return filtered
|
||||
# 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))
|
||||
|
||||
async def _load_bank_config(self, bank_id: str) -> dict[str, Any]:
|
||||
"""
|
||||
@@ -219,6 +277,45 @@ 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:
|
||||
|
||||
@@ -96,7 +96,10 @@ class BankStatsCache:
|
||||
value = await loader()
|
||||
except BaseException as exc:
|
||||
async with self._lock:
|
||||
self._in_flight.pop(key, None)
|
||||
# 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)
|
||||
if not in_flight.done():
|
||||
in_flight.set_exception(exc)
|
||||
# Suppress "Future exception was never retrieved" when no other
|
||||
@@ -106,8 +109,12 @@ class BankStatsCache:
|
||||
raise
|
||||
|
||||
async with self._lock:
|
||||
self._store_unlocked(key, value)
|
||||
self._in_flight.pop(key, None)
|
||||
# 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)
|
||||
if not in_flight.done():
|
||||
in_flight.set_result(value)
|
||||
return value
|
||||
@@ -115,8 +122,13 @@ class BankStatsCache:
|
||||
async def invalidate(self, schema: str, bank_id: str) -> None:
|
||||
"""Drop any cached stats for `(schema, bank_id)`."""
|
||||
async with self._lock:
|
||||
self._entries.pop((schema, bank_id), None)
|
||||
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)
|
||||
|
||||
async def clear(self) -> None:
|
||||
async with self._lock:
|
||||
self._entries.clear()
|
||||
self._in_flight.clear()
|
||||
|
||||
@@ -449,6 +449,13 @@ 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
|
||||
@@ -461,6 +468,13 @@ 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
|
||||
@@ -640,6 +654,7 @@ 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
|
||||
@@ -649,11 +664,13 @@ class ConsolidationPerfLog:
|
||||
self.lines.append(message)
|
||||
|
||||
def record_timing(self, key: str, duration: float) -> None:
|
||||
"""Record a timing measurement."""
|
||||
if key in self.timings:
|
||||
self.timings[key] += duration
|
||||
else:
|
||||
self.timings[key] = duration
|
||||
"""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
|
||||
|
||||
def record_llm_call(self, obs_count: int, prompt_chars: int) -> None:
|
||||
"""Record stats for a single LLM call."""
|
||||
@@ -676,6 +693,8 @@ 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
|
||||
@@ -1276,16 +1295,22 @@ async def _run_consolidation_job(
|
||||
f"{stats['skipped']} skipped)"
|
||||
)
|
||||
|
||||
# Add timing breakdown
|
||||
# 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"
|
||||
|
||||
timing_parts = []
|
||||
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")
|
||||
for key in ("recall", "llm", "embedding", "db_write"):
|
||||
if key in perf.timings:
|
||||
timing_parts.append(_fmt(key))
|
||||
|
||||
if perf.llm_calls > 0:
|
||||
timing_parts.append(f"avg_obs={perf.total_obs_in_context / perf.llm_calls:.1f}")
|
||||
|
||||
@@ -190,7 +190,7 @@ class LocalSTEmbeddings(Embeddings):
|
||||
device = "cpu"
|
||||
logger.info("Embeddings: forcing CPU mode")
|
||||
else:
|
||||
# Check for GPU (CUDA) or Apple Silicon (MPS)
|
||||
# Check for GPU (CUDA), Apple Silicon (MPS), or Intel XPU
|
||||
# Wrap in try-except to gracefully handle any device detection issues
|
||||
# (e.g., in CI environments or when PyTorch is built without GPU support)
|
||||
device = "cpu" # Default to CPU
|
||||
@@ -198,10 +198,13 @@ 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
|
||||
device = None # Let sentence-transformers auto-detect GPU/MPS/XPU
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to detect GPU/MPS, falling back to CPU: {e}")
|
||||
logger.warning(f"Failed to detect GPU/MPS/XPU, falling back to CPU: {e}")
|
||||
|
||||
# Suppress verbose transformers warnings during model loading
|
||||
# This suppresses the "UNEXPECTED" warnings from BertModel which are harmless
|
||||
@@ -709,7 +712,8 @@ class OpenAIEmbeddings(Embeddings):
|
||||
|
||||
class CodexOAuthEmbeddings(OpenAIEmbeddings):
|
||||
"""
|
||||
OpenAI embeddings using the Codex/ChatGPT OAuth token from ``~/.codex/auth.json``.
|
||||
OpenAI embeddings using the Codex/ChatGPT OAuth token from the Codex
|
||||
``auth.json`` (``$CODEX_HOME/auth.json``, or ``~/.codex/auth.json`` when unset).
|
||||
|
||||
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
|
||||
|
||||
@@ -6,6 +6,7 @@ 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
|
||||
@@ -252,3 +253,11 @@ 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,7 +10,6 @@ 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)
|
||||
@@ -253,6 +252,7 @@ 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,6 +266,7 @@ 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
|
||||
@@ -296,6 +297,12 @@ 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(
|
||||
@@ -344,6 +351,7 @@ 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,
|
||||
)
|
||||
@@ -496,6 +504,7 @@ 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.
|
||||
@@ -509,6 +518,7 @@ 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).
|
||||
@@ -532,6 +542,7 @@ 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
|
||||
@@ -660,6 +671,22 @@ 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
|
||||
@@ -698,6 +725,7 @@ 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,
|
||||
@@ -1023,7 +1051,9 @@ class LLMProvider:
|
||||
|
||||
def _load_codex_auth(self) -> tuple[str, str]:
|
||||
"""
|
||||
Load OAuth credentials from ~/.codex/auth.json.
|
||||
Load OAuth credentials from the Codex ``auth.json``.
|
||||
|
||||
Honors ``CODEX_HOME`` (falling back to ``~/.codex``).
|
||||
|
||||
Returns:
|
||||
Tuple of (access_token, account_id).
|
||||
@@ -1032,7 +1062,9 @@ class LLMProvider:
|
||||
FileNotFoundError: If auth file doesn't exist.
|
||||
ValueError: If auth file is invalid.
|
||||
"""
|
||||
auth_file = Path.home() / ".codex" / "auth.json"
|
||||
from .providers.codex_auth import default_codex_auth_file
|
||||
|
||||
auth_file = default_codex_auth_file()
|
||||
|
||||
if not auth_file.exists():
|
||||
raise FileNotFoundError(
|
||||
@@ -1142,10 +1174,12 @@ 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)
|
||||
@@ -1172,6 +1206,11 @@ 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,6 +45,7 @@ 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,
|
||||
@@ -744,6 +745,37 @@ 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.
|
||||
@@ -933,6 +965,7 @@ 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)
|
||||
@@ -966,6 +999,7 @@ 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)
|
||||
@@ -994,6 +1028,7 @@ 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)
|
||||
@@ -1022,6 +1057,7 @@ 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)
|
||||
@@ -1752,6 +1788,9 @@ 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
|
||||
@@ -2751,7 +2790,15 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
|
||||
self._parser_registry = FileParserRegistry()
|
||||
try:
|
||||
self._parser_registry.register(MarkitdownParser())
|
||||
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,
|
||||
)
|
||||
)
|
||||
logger.debug("Registered markitdown parser")
|
||||
except ImportError:
|
||||
logger.warning("markitdown not available - file parsing disabled")
|
||||
@@ -3287,6 +3334,28 @@ 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,
|
||||
@@ -3306,20 +3375,10 @@ class MemoryEngine(MemoryEngineInterface):
|
||||
)
|
||||
|
||||
# Advance the document's chunk_index cursor by the number of
|
||||
# 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.
|
||||
# chunks this sub-batch produced (counted above, before the
|
||||
# orchestrator consumed the content), 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.
|
||||
@@ -8112,25 +8171,15 @@ 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)
|
||||
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"],
|
||||
}
|
||||
resolved = _overlay_bank_config_disposition_mission(db_disp_dict, profile["mission"], config_dict)
|
||||
|
||||
return {
|
||||
"bank_id": bank_id,
|
||||
"name": profile["name"],
|
||||
"disposition": disposition,
|
||||
"mission": mission,
|
||||
"disposition": resolved.disposition,
|
||||
"mission": resolved.mission,
|
||||
}
|
||||
|
||||
async def _ensure_bank_exists(
|
||||
@@ -8345,6 +8394,17 @@ 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 ====================
|
||||
|
||||
@@ -3,43 +3,116 @@
|
||||
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 (via OCR), audio, HTML.
|
||||
to markdown including PDF, Office docs, images with optional OCR,
|
||||
audio, HTML.
|
||||
|
||||
Supported formats:
|
||||
- PDF (.pdf)
|
||||
- Word (.docx, .doc)
|
||||
- PowerPoint (.pptx, .ppt)
|
||||
- Excel (.xlsx, .xls)
|
||||
- Images (.jpg, .jpeg, .png) - with OCR
|
||||
- Images (.jpg, .jpeg, .png) - optional OCR
|
||||
- HTML (.html, .htm)
|
||||
- Text (.txt, .md)
|
||||
- Audio (.mp3, .wav) - with transcription
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
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,
|
||||
):
|
||||
"""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
|
||||
@@ -48,6 +121,13 @@ 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)
|
||||
@@ -73,6 +153,11 @@ 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)
|
||||
@@ -85,7 +170,7 @@ class MarkitdownParser(FileParser):
|
||||
".ppt",
|
||||
".xlsx",
|
||||
".xls",
|
||||
# Images (with OCR)
|
||||
# Images (optional OCR)
|
||||
".jpg",
|
||||
".jpeg",
|
||||
".png",
|
||||
|
||||
@@ -60,6 +60,22 @@ _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.
|
||||
|
||||
@@ -86,7 +102,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 ``~/.codex/auth.json``. Used for re-reading the refresh token
|
||||
Path to the Codex ``auth.json``. Used for re-reading the refresh token
|
||||
on demand and for atomic persistence of rotated credentials.
|
||||
"""
|
||||
|
||||
@@ -115,7 +131,8 @@ class CodexAuthManager:
|
||||
Parameters
|
||||
----------
|
||||
auth_file:
|
||||
Defaults to ``~/.codex/auth.json``.
|
||||
Defaults to ``$CODEX_HOME/auth.json`` (or ``~/.codex/auth.json``
|
||||
when ``CODEX_HOME`` is unset).
|
||||
|
||||
Raises
|
||||
------
|
||||
@@ -126,7 +143,7 @@ class CodexAuthManager:
|
||||
``auth_mode``.
|
||||
"""
|
||||
if auth_file is None:
|
||||
auth_file = Path.home() / ".codex" / "auth.json"
|
||||
auth_file = default_codex_auth_file()
|
||||
|
||||
if not auth_file.exists():
|
||||
raise FileNotFoundError(f"Codex auth file not found: {auth_file}. Run 'codex auth login' to authenticate.")
|
||||
|
||||
@@ -2,8 +2,9 @@
|
||||
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
|
||||
~/.codex/auth.json and communicates with the ChatGPT backend API.
|
||||
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.
|
||||
|
||||
Tokens are refreshed automatically: the provider decodes the access_token
|
||||
JWT's ``exp`` claim and proactively refreshes via
|
||||
@@ -35,6 +36,7 @@ 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).
|
||||
@@ -55,14 +57,15 @@ class CodexLLM(LLMInterface):
|
||||
"""
|
||||
LLM provider using OpenAI Codex OAuth authentication.
|
||||
|
||||
Authenticates using ChatGPT Plus/Pro credentials stored in ~/.codex/auth.json
|
||||
and makes API calls to chatgpt.com/backend-api/codex/responses.
|
||||
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.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
provider: str,
|
||||
api_key: str, # Will be ignored, reads from ~/.codex/auth.json
|
||||
api_key: str, # Will be ignored, reads from the Codex auth.json (CODEX_HOME or ~/.codex)
|
||||
base_url: str,
|
||||
model: str,
|
||||
reasoning_effort: str = "low",
|
||||
@@ -81,12 +84,14 @@ 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 ~/.codex/auth.json: {e}\n\n"
|
||||
f"Failed to load Codex OAuth credentials from {auth_file}: {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"
|
||||
"3. Verify: ls ~/.codex/auth.json\n\n"
|
||||
f"3. Verify: ls {auth_file}\n\n"
|
||||
"(Set CODEX_HOME to use a credentials directory other than ~/.codex.)\n\n"
|
||||
"Or use a different provider (openai, anthropic, gemini) with API keys."
|
||||
) from e
|
||||
|
||||
@@ -94,7 +99,7 @@ class CodexLLM(LLMInterface):
|
||||
access_token=access_token,
|
||||
account_id=account_id,
|
||||
refresh_token=refresh_token,
|
||||
auth_file=Path.home() / ".codex" / "auth.json",
|
||||
auth_file=default_codex_auth_file(),
|
||||
)
|
||||
|
||||
# Use ChatGPT backend API endpoint. Codex auth is tied to
|
||||
@@ -156,7 +161,7 @@ class CodexLLM(LLMInterface):
|
||||
|
||||
def _load_codex_auth(self) -> tuple[str, str]:
|
||||
"""
|
||||
Load OAuth credentials from ~/.codex/auth.json.
|
||||
Load OAuth credentials from the Codex ``auth.json`` (CODEX_HOME or ~/.codex).
|
||||
|
||||
Returns:
|
||||
Tuple of (access_token, account_id).
|
||||
@@ -165,7 +170,7 @@ class CodexLLM(LLMInterface):
|
||||
FileNotFoundError: If auth file doesn't exist.
|
||||
ValueError: If auth file is invalid.
|
||||
"""
|
||||
auth_file = Path.home() / ".codex" / "auth.json"
|
||||
auth_file = default_codex_auth_file()
|
||||
|
||||
if not auth_file.exists():
|
||||
raise FileNotFoundError(
|
||||
@@ -197,9 +202,7 @@ 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 Path.home() / ".codex" / "auth.json"
|
||||
)
|
||||
auth_file = self._auth_manager._auth_file if hasattr(self, "_auth_manager") else default_codex_auth_file()
|
||||
return CodexAuthManager.load_refresh_token_from_file(auth_file)
|
||||
|
||||
@staticmethod
|
||||
|
||||
@@ -76,6 +76,7 @@ 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
|
||||
@@ -106,6 +107,16 @@ 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
|
||||
@@ -273,6 +284,7 @@ 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 system_instruction:
|
||||
@@ -604,6 +616,7 @@ 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:
|
||||
|
||||
@@ -15,9 +15,13 @@ 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
|
||||
@@ -47,13 +51,15 @@ class LiteLLMLLM(LLMInterface):
|
||||
base_url: str,
|
||||
model: str,
|
||||
reasoning_effort: str = "low",
|
||||
timeout: float = 300.0,
|
||||
timeout: float | None = None,
|
||||
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)
|
||||
self.timeout = timeout
|
||||
# ``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._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
|
||||
@@ -209,7 +215,10 @@ class LiteLLMLLM(LLMInterface):
|
||||
if attempt > 0:
|
||||
set_stage(f"llm.{self._stage_label}.{scope}.attempt={attempt + 1}/{max_retries + 1}")
|
||||
try:
|
||||
response = await self._acompletion(**call_kwargs)
|
||||
response = await asyncio.wait_for(
|
||||
self._acompletion(**call_kwargs),
|
||||
timeout=self.timeout,
|
||||
)
|
||||
|
||||
content = response.choices[0].message.content or ""
|
||||
finish_reason = response.choices[0].finish_reason
|
||||
@@ -304,6 +313,25 @@ 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
|
||||
@@ -354,7 +382,10 @@ class LiteLLMLLM(LLMInterface):
|
||||
if attempt > 0:
|
||||
set_stage(f"llm.{self._stage_label}.tools.attempt={attempt + 1}/{max_retries + 1}")
|
||||
try:
|
||||
response = await self._acompletion(**call_kwargs)
|
||||
response = await asyncio.wait_for(
|
||||
self._acompletion(**call_kwargs),
|
||||
timeout=self.timeout,
|
||||
)
|
||||
|
||||
message = response.choices[0].message
|
||||
content = message.content
|
||||
@@ -424,6 +455,23 @@ 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 = 300.0,
|
||||
timeout: float | None = None,
|
||||
**kwargs: Any,
|
||||
):
|
||||
super().__init__(
|
||||
|
||||
@@ -26,6 +26,8 @@ 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
|
||||
|
||||
@@ -34,7 +36,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
|
||||
from hindsight_api.engine.llm_interface import LLMInterface, OutputTooLongError, ProviderRateLimitResetError
|
||||
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
|
||||
@@ -83,6 +85,49 @@ 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)
|
||||
@@ -234,6 +279,122 @@ 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.
|
||||
@@ -269,7 +430,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 300s default).
|
||||
timeout: Request timeout in seconds (uses env var or 120s default).
|
||||
groq_service_tier: Groq service tier ("on_demand", "flex", "auto").
|
||||
extra_body: Extra body params merged into every API call.
|
||||
**kwargs: Additional provider-specific parameters.
|
||||
@@ -617,15 +778,10 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
scope=scope,
|
||||
)
|
||||
|
||||
# Strip reasoning model thinking tags
|
||||
# Strip reasoning model thinking tags (closed and unclosed).
|
||||
# Supports: <think>, <thinking>, <thought>, <reasoning>, |startthink|/|endthink|
|
||||
original_len = len(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()
|
||||
content = _strip_reasoning_tags(content)
|
||||
if len(content) < original_len:
|
||||
logger.debug(f"Stripped {original_len - len(content)} chars of reasoning tokens")
|
||||
|
||||
@@ -674,6 +830,13 @@ 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
|
||||
@@ -761,6 +924,10 @@ 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:
|
||||
@@ -814,7 +981,6 @@ 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:
|
||||
@@ -1047,6 +1213,10 @@ 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(
|
||||
@@ -1060,7 +1230,6 @@ class OpenAICompatibleLLM(LLMInterface):
|
||||
f"({self.provider}/{self.model}, scope={scope}): {_summarize_status_error(e)}"
|
||||
)
|
||||
raise
|
||||
|
||||
except Exception:
|
||||
raise
|
||||
|
||||
|
||||
@@ -14,6 +14,7 @@ from typing import Any, Literal, cast
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field, create_model, field_validator
|
||||
|
||||
from ..llm_interface import ProviderRateLimitResetError
|
||||
from ..llm_wrapper import LLMConfig, OutputTooLongError, sanitize_llm_output
|
||||
from ..operation_metadata import RetainExtractionErrors
|
||||
from ..response_models import TokenUsage
|
||||
@@ -451,6 +452,11 @@ 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
|
||||
@@ -469,11 +475,23 @@ 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):
|
||||
pass
|
||||
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)
|
||||
|
||||
# Try to parse as JSONL (newline-delimited JSON objects, e.g. session logs)
|
||||
jsonl_chunks = _chunk_jsonl(text, max_chars, structured_limit)
|
||||
@@ -515,10 +533,12 @@ 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 so no chunk runs far over budget (the extractor won't re-chunk).
|
||||
# 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).
|
||||
if turn_unit_size > structured_limit:
|
||||
_flush()
|
||||
chunks.extend(_split_oversized_unit(turn_json, structured_limit))
|
||||
chunks.extend(_split_oversized_unit(turn_json, min(structured_limit, max_chars)))
|
||||
continue
|
||||
|
||||
# If adding this turn would exceed limit and we have turns, save current chunk
|
||||
@@ -581,10 +601,12 @@ 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 so no chunk runs far over budget (the extractor won't re-chunk).
|
||||
# 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).
|
||||
if line_unit_size > structured_limit:
|
||||
_flush()
|
||||
chunks.extend(_split_oversized_unit(line, structured_limit))
|
||||
chunks.extend(_split_oversized_unit(line, min(structured_limit, max_chars)))
|
||||
continue
|
||||
|
||||
# If adding this line would exceed the limit and we have lines, flush.
|
||||
@@ -1792,10 +1814,21 @@ 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,8 +1615,19 @@ 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:
|
||||
if operation_id and chunk_index_offset == 0:
|
||||
try:
|
||||
async with acquire_with_retry(pool) as conn:
|
||||
row = await conn.fetchrow(
|
||||
|
||||
@@ -97,6 +97,10 @@ 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``
|
||||
@@ -107,6 +111,7 @@ class PrecheckContext:
|
||||
operation: str
|
||||
bank_id: str
|
||||
request_context: "RequestContext"
|
||||
content_length: int | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
|
||||
@@ -12,6 +12,7 @@ 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
|
||||
@@ -199,6 +200,47 @@ 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,
|
||||
@@ -552,7 +594,7 @@ def _register_retain(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig)
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool(description=description)
|
||||
@mcp.tool(description=description, annotations=_tool_annotations("retain"))
|
||||
async def retain(
|
||||
content: str,
|
||||
context: str = "general",
|
||||
@@ -608,7 +650,7 @@ def _register_retain(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig)
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool(description=description)
|
||||
@mcp.tool(description=description, annotations=_tool_annotations("retain"))
|
||||
async def retain(
|
||||
content: str,
|
||||
context: str = "general",
|
||||
@@ -666,7 +708,7 @@ def _register_sync_retain(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsCo
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("sync_retain"))
|
||||
async def sync_retain(
|
||||
content: str,
|
||||
context: str = "general",
|
||||
@@ -724,7 +766,7 @@ def _register_sync_retain(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsCo
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("sync_retain"))
|
||||
async def sync_retain(
|
||||
content: str,
|
||||
context: str = "general",
|
||||
@@ -785,7 +827,7 @@ def _register_recall(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig)
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool(description=description)
|
||||
@mcp.tool(description=description, annotations=_tool_annotations("recall"))
|
||||
async def recall(
|
||||
query: str,
|
||||
max_tokens: int = 4096,
|
||||
@@ -857,7 +899,7 @@ def _register_recall(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig)
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool(description=description)
|
||||
@mcp.tool(description=description, annotations=_tool_annotations("recall"))
|
||||
async def recall(
|
||||
query: str,
|
||||
max_tokens: int = 4096,
|
||||
@@ -931,7 +973,7 @@ def _register_reflect(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("reflect"))
|
||||
async def reflect(
|
||||
query: str,
|
||||
context: str | None = None,
|
||||
@@ -941,6 +983,7 @@ 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:
|
||||
"""
|
||||
@@ -971,6 +1014,7 @@ 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 tool_trace/llm_trace. 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:
|
||||
@@ -1000,6 +1044,12 @@ 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 tool_trace/llm_trace can be tens of KB
|
||||
# (full mental-model text) and silently overflow MCP client context;
|
||||
# the REST API omits it by default too. Opt in via include_trace.
|
||||
result_data.pop("tool_trace", None)
|
||||
result_data.pop("llm_trace", 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)
|
||||
@@ -1012,7 +1062,7 @@ def _register_reflect(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfig
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("reflect"))
|
||||
async def reflect(
|
||||
query: str,
|
||||
context: str | None = None,
|
||||
@@ -1022,6 +1072,7 @@ 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.
|
||||
@@ -1051,6 +1102,7 @@ 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 tool_trace/llm_trace. 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()
|
||||
@@ -1079,6 +1131,12 @@ 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 tool_trace/llm_trace can be tens of KB
|
||||
# (full mental-model text) and silently overflow MCP client context;
|
||||
# the REST API omits it by default too. Opt in via include_trace.
|
||||
result_data.pop("tool_trace", None)
|
||||
result_data.pop("llm_trace", None)
|
||||
if response_schema is not None and hasattr(reflect_result, "structured_output"):
|
||||
result_data["structured_output"] = reflect_result.structured_output
|
||||
return result_data
|
||||
@@ -1093,7 +1151,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()
|
||||
@mcp.tool(annotations=_tool_annotations("list_banks"))
|
||||
async def list_banks() -> str:
|
||||
"""
|
||||
List all available memory banks.
|
||||
@@ -1118,7 +1176,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()
|
||||
@mcp.tool(annotations=_tool_annotations("create_bank"))
|
||||
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.
|
||||
@@ -1182,7 +1240,7 @@ def _register_list_mental_models(mcp: FastMCP, memory: MemoryEngine, config: MCP
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("list_mental_models"))
|
||||
async def list_mental_models(
|
||||
tags: list[str] | None = None,
|
||||
detail: str = "full",
|
||||
@@ -1221,7 +1279,7 @@ def _register_list_mental_models(mcp: FastMCP, memory: MemoryEngine, config: MCP
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("list_mental_models"))
|
||||
async def list_mental_models(
|
||||
tags: list[str] | None = None,
|
||||
detail: str = "full",
|
||||
@@ -1262,7 +1320,7 @@ def _register_get_mental_model(mcp: FastMCP, memory: MemoryEngine, config: MCPTo
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("get_mental_model"))
|
||||
async def get_mental_model(
|
||||
mental_model_id: str,
|
||||
detail: str = "full",
|
||||
@@ -1302,7 +1360,7 @@ def _register_get_mental_model(mcp: FastMCP, memory: MemoryEngine, config: MCPTo
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("get_mental_model"))
|
||||
async def get_mental_model(
|
||||
mental_model_id: str,
|
||||
detail: str = "full",
|
||||
@@ -1344,7 +1402,7 @@ def _register_create_mental_model(mcp: FastMCP, memory: MemoryEngine, config: MC
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("create_mental_model"))
|
||||
async def create_mental_model(
|
||||
name: str,
|
||||
source_query: str,
|
||||
@@ -1428,7 +1486,7 @@ def _register_create_mental_model(mcp: FastMCP, memory: MemoryEngine, config: MC
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("create_mental_model"))
|
||||
async def create_mental_model(
|
||||
name: str,
|
||||
source_query: str,
|
||||
@@ -1510,7 +1568,7 @@ def _register_update_mental_model(mcp: FastMCP, memory: MemoryEngine, config: MC
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("update_mental_model"))
|
||||
async def update_mental_model(
|
||||
mental_model_id: str,
|
||||
name: str | None = None,
|
||||
@@ -1571,7 +1629,7 @@ def _register_update_mental_model(mcp: FastMCP, memory: MemoryEngine, config: MC
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("update_mental_model"))
|
||||
async def update_mental_model(
|
||||
mental_model_id: str,
|
||||
name: str | None = None,
|
||||
@@ -1634,7 +1692,7 @@ def _register_delete_mental_model(mcp: FastMCP, memory: MemoryEngine, config: MC
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("delete_mental_model"))
|
||||
async def delete_mental_model(
|
||||
mental_model_id: str,
|
||||
bank_id: str | None = None,
|
||||
@@ -1670,7 +1728,7 @@ def _register_delete_mental_model(mcp: FastMCP, memory: MemoryEngine, config: MC
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("delete_mental_model"))
|
||||
async def delete_mental_model(
|
||||
mental_model_id: str,
|
||||
) -> dict:
|
||||
@@ -1708,7 +1766,7 @@ def _register_refresh_mental_model(mcp: FastMCP, memory: MemoryEngine, config: M
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("refresh_mental_model"))
|
||||
async def refresh_mental_model(
|
||||
mental_model_id: str,
|
||||
bank_id: str | None = None,
|
||||
@@ -1752,7 +1810,7 @@ def _register_refresh_mental_model(mcp: FastMCP, memory: MemoryEngine, config: M
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("refresh_mental_model"))
|
||||
async def refresh_mental_model(
|
||||
mental_model_id: str,
|
||||
) -> dict:
|
||||
@@ -1796,7 +1854,7 @@ def _register_clear_mental_model(mcp: FastMCP, memory: MemoryEngine, config: MCP
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("clear_mental_model"))
|
||||
async def clear_mental_model(
|
||||
mental_model_id: str,
|
||||
bank_id: str | None = None,
|
||||
@@ -1842,7 +1900,7 @@ def _register_clear_mental_model(mcp: FastMCP, memory: MemoryEngine, config: MCP
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("clear_mental_model"))
|
||||
async def clear_mental_model(
|
||||
mental_model_id: str,
|
||||
) -> dict:
|
||||
@@ -1893,7 +1951,7 @@ def _register_list_directives(mcp: FastMCP, memory: MemoryEngine, config: MCPToo
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("list_directives"))
|
||||
async def list_directives(
|
||||
tags: list[str] | None = None,
|
||||
active_only: bool = True,
|
||||
@@ -1931,7 +1989,7 @@ def _register_list_directives(mcp: FastMCP, memory: MemoryEngine, config: MCPToo
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("list_directives"))
|
||||
async def list_directives(
|
||||
tags: list[str] | None = None,
|
||||
active_only: bool = True,
|
||||
@@ -1971,7 +2029,7 @@ def _register_create_directive(mcp: FastMCP, memory: MemoryEngine, config: MCPTo
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("create_directive"))
|
||||
async def create_directive(
|
||||
name: str,
|
||||
content: str,
|
||||
@@ -2017,7 +2075,7 @@ def _register_create_directive(mcp: FastMCP, memory: MemoryEngine, config: MCPTo
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("create_directive"))
|
||||
async def create_directive(
|
||||
name: str,
|
||||
content: str,
|
||||
@@ -2065,7 +2123,7 @@ def _register_delete_directive(mcp: FastMCP, memory: MemoryEngine, config: MCPTo
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("delete_directive"))
|
||||
async def delete_directive(
|
||||
directive_id: str,
|
||||
bank_id: str | None = None,
|
||||
@@ -2101,7 +2159,7 @@ def _register_delete_directive(mcp: FastMCP, memory: MemoryEngine, config: MCPTo
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("delete_directive"))
|
||||
async def delete_directive(
|
||||
directive_id: str,
|
||||
) -> dict:
|
||||
@@ -2144,7 +2202,7 @@ def _register_list_memories(mcp: FastMCP, memory: MemoryEngine, config: MCPTools
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("list_memories"))
|
||||
async def list_memories(
|
||||
type: str | None = None,
|
||||
q: str | None = None,
|
||||
@@ -2159,7 +2217,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 'opinion'
|
||||
type: Filter by fact type: 'world', 'experience', or 'observation'
|
||||
q: Optional text search query to filter memories
|
||||
limit: Maximum number of results (default: 100)
|
||||
offset: Pagination offset (default: 0)
|
||||
@@ -2188,7 +2246,7 @@ def _register_list_memories(mcp: FastMCP, memory: MemoryEngine, config: MCPTools
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("list_memories"))
|
||||
async def list_memories(
|
||||
type: str | None = None,
|
||||
q: str | None = None,
|
||||
@@ -2202,7 +2260,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 'opinion'
|
||||
type: Filter by fact type: 'world', 'experience', or 'observation'
|
||||
q: Optional text search query to filter memories
|
||||
limit: Maximum number of results (default: 100)
|
||||
offset: Pagination offset (default: 0)
|
||||
@@ -2234,7 +2292,7 @@ def _register_get_memory(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsCon
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("get_memory"))
|
||||
async def get_memory(
|
||||
memory_id: str,
|
||||
bank_id: str | None = None,
|
||||
@@ -2270,7 +2328,7 @@ def _register_get_memory(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsCon
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("get_memory"))
|
||||
async def get_memory(
|
||||
memory_id: str,
|
||||
) -> dict:
|
||||
@@ -2321,7 +2379,7 @@ def _register_update_memory(mcp: FastMCP, memory: MemoryEngine, config: MCPTools
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool(description=_EDIT_DOC)
|
||||
@mcp.tool(description=_EDIT_DOC, annotations=_tool_annotations("update_memory"))
|
||||
async def update_memory(
|
||||
memory_id: str,
|
||||
text: str | None = None,
|
||||
@@ -2367,7 +2425,7 @@ def _register_update_memory(mcp: FastMCP, memory: MemoryEngine, config: MCPTools
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool(description=_EDIT_DOC)
|
||||
@mcp.tool(description=_EDIT_DOC, annotations=_tool_annotations("update_memory"))
|
||||
async def update_memory(
|
||||
memory_id: str,
|
||||
text: str | None = None,
|
||||
@@ -2426,7 +2484,7 @@ def _register_invalidate_memory(mcp: FastMCP, memory: MemoryEngine, config: MCPT
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool(description=_INVALIDATE_DOC)
|
||||
@mcp.tool(description=_INVALIDATE_DOC, annotations=_tool_annotations("invalidate_memory"))
|
||||
async def invalidate_memory(
|
||||
memory_id: str,
|
||||
reason: str | None = None,
|
||||
@@ -2466,7 +2524,7 @@ def _register_invalidate_memory(mcp: FastMCP, memory: MemoryEngine, config: MCPT
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool(description=_INVALIDATE_DOC)
|
||||
@mcp.tool(description=_INVALIDATE_DOC, annotations=_tool_annotations("invalidate_memory"))
|
||||
async def invalidate_memory(
|
||||
memory_id: str,
|
||||
reason: str | None = None,
|
||||
@@ -2513,7 +2571,7 @@ def _register_list_documents(mcp: FastMCP, memory: MemoryEngine, config: MCPTool
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("list_documents"))
|
||||
async def list_documents(
|
||||
q: str | None = None,
|
||||
limit: int = 100,
|
||||
@@ -2551,7 +2609,7 @@ def _register_list_documents(mcp: FastMCP, memory: MemoryEngine, config: MCPTool
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("list_documents"))
|
||||
async def list_documents(
|
||||
q: str | None = None,
|
||||
limit: int = 100,
|
||||
@@ -2591,7 +2649,7 @@ def _register_get_document(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsC
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("get_document"))
|
||||
async def get_document(
|
||||
document_id: str,
|
||||
bank_id: str | None = None,
|
||||
@@ -2627,7 +2685,7 @@ def _register_get_document(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsC
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("get_document"))
|
||||
async def get_document(
|
||||
document_id: str,
|
||||
) -> dict:
|
||||
@@ -2665,7 +2723,7 @@ def _register_delete_document(mcp: FastMCP, memory: MemoryEngine, config: MCPToo
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("delete_document"))
|
||||
async def delete_document(
|
||||
document_id: str,
|
||||
bank_id: str | None = None,
|
||||
@@ -2699,7 +2757,7 @@ def _register_delete_document(mcp: FastMCP, memory: MemoryEngine, config: MCPToo
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("delete_document"))
|
||||
async def delete_document(
|
||||
document_id: str,
|
||||
) -> dict:
|
||||
@@ -2740,7 +2798,7 @@ def _register_list_operations(mcp: FastMCP, memory: MemoryEngine, config: MCPToo
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("list_operations"))
|
||||
async def list_operations(
|
||||
status: str | None = None,
|
||||
limit: int = 20,
|
||||
@@ -2777,7 +2835,7 @@ def _register_list_operations(mcp: FastMCP, memory: MemoryEngine, config: MCPToo
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("list_operations"))
|
||||
async def list_operations(
|
||||
status: str | None = None,
|
||||
limit: int = 20,
|
||||
@@ -2816,7 +2874,7 @@ def _register_get_operation(mcp: FastMCP, memory: MemoryEngine, config: MCPTools
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("get_operation"))
|
||||
async def get_operation(
|
||||
operation_id: str,
|
||||
bank_id: str | None = None,
|
||||
@@ -2850,7 +2908,7 @@ def _register_get_operation(mcp: FastMCP, memory: MemoryEngine, config: MCPTools
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("get_operation"))
|
||||
async def get_operation(
|
||||
operation_id: str,
|
||||
) -> dict:
|
||||
@@ -2886,7 +2944,7 @@ def _register_cancel_operation(mcp: FastMCP, memory: MemoryEngine, config: MCPTo
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("cancel_operation"))
|
||||
async def cancel_operation(
|
||||
operation_id: str,
|
||||
bank_id: str | None = None,
|
||||
@@ -2918,7 +2976,7 @@ def _register_cancel_operation(mcp: FastMCP, memory: MemoryEngine, config: MCPTo
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("cancel_operation"))
|
||||
async def cancel_operation(
|
||||
operation_id: str,
|
||||
) -> dict:
|
||||
@@ -2957,7 +3015,7 @@ def _register_list_tags(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConf
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("list_tags"))
|
||||
async def list_tags(
|
||||
q: str | None = None,
|
||||
limit: int = 100,
|
||||
@@ -2994,7 +3052,7 @@ def _register_list_tags(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConf
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("list_tags"))
|
||||
async def list_tags(
|
||||
q: str | None = None,
|
||||
limit: int = 100,
|
||||
@@ -3033,7 +3091,7 @@ def _register_get_bank(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfi
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("get_bank"))
|
||||
async def get_bank(
|
||||
bank_id: str | None = None,
|
||||
) -> str:
|
||||
@@ -3066,7 +3124,7 @@ def _register_get_bank(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsConfi
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("get_bank"))
|
||||
async def get_bank() -> dict:
|
||||
"""
|
||||
Get the profile of this memory bank.
|
||||
@@ -3096,7 +3154,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()
|
||||
@mcp.tool(annotations=_tool_annotations("get_bank_stats"))
|
||||
async def get_bank_stats(
|
||||
bank_id: str | None = None,
|
||||
) -> str:
|
||||
@@ -3169,7 +3227,7 @@ def _register_update_bank(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsCo
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("update_bank"))
|
||||
async def update_bank(
|
||||
name: str | None = None,
|
||||
mission: str | None = None,
|
||||
@@ -3230,7 +3288,7 @@ def _register_update_bank(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsCo
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("update_bank"))
|
||||
async def update_bank(
|
||||
name: str | None = None,
|
||||
mission: str | None = None,
|
||||
@@ -3293,7 +3351,7 @@ def _register_delete_bank(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsCo
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("delete_bank"))
|
||||
async def delete_bank(
|
||||
bank_id: str | None = None,
|
||||
) -> str:
|
||||
@@ -3325,7 +3383,7 @@ def _register_delete_bank(mcp: FastMCP, memory: MemoryEngine, config: MCPToolsCo
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("delete_bank"))
|
||||
async def delete_bank() -> dict:
|
||||
"""
|
||||
Delete this memory bank and all its data.
|
||||
@@ -3356,7 +3414,7 @@ def _register_clear_memories(mcp: FastMCP, memory: MemoryEngine, config: MCPTool
|
||||
|
||||
if config.include_bank_id_param:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("clear_memories"))
|
||||
async def clear_memories(
|
||||
type: str | None = None,
|
||||
bank_id: str | None = None,
|
||||
@@ -3367,7 +3425,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 'opinion'. If not specified, clears all.
|
||||
type: Optional fact type filter: 'world', 'experience', or 'observation'. If not specified, clears all.
|
||||
bank_id: Optional bank (defaults to session bank). Use for cross-bank operations.
|
||||
"""
|
||||
try:
|
||||
@@ -3391,7 +3449,7 @@ def _register_clear_memories(mcp: FastMCP, memory: MemoryEngine, config: MCPTool
|
||||
|
||||
else:
|
||||
|
||||
@mcp.tool()
|
||||
@mcp.tool(annotations=_tool_annotations("clear_memories"))
|
||||
async def clear_memories(
|
||||
type: str | None = None,
|
||||
) -> dict:
|
||||
@@ -3401,7 +3459,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 'opinion'. If not specified, clears all.
|
||||
type: Optional fact type filter: 'world', 'experience', or 'observation'. If not specified, clears all.
|
||||
"""
|
||||
try:
|
||||
target_bank = config.bank_id_resolver()
|
||||
|
||||
@@ -11,15 +11,17 @@ 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
|
||||
from typing import TYPE_CHECKING, Callable, NamedTuple
|
||||
|
||||
from opentelemetry import metrics
|
||||
from opentelemetry.exporter.prometheus import PrometheusMetricReader
|
||||
@@ -75,6 +77,28 @@ 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:
|
||||
"""
|
||||
@@ -113,6 +137,27 @@ 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
|
||||
@@ -201,6 +246,19 @@ 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,
|
||||
@@ -254,6 +312,19 @@ 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,
|
||||
@@ -361,6 +432,13 @@ 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,
|
||||
@@ -386,18 +464,6 @@ 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:
|
||||
@@ -416,14 +482,51 @@ class MetricsCollector(MetricsCollectorBase):
|
||||
raise
|
||||
finally:
|
||||
if not cancelled:
|
||||
duration = time.time() - start_time
|
||||
attributes["success"] = str(success).lower()
|
||||
self.record_operation_result(
|
||||
operation,
|
||||
bank_id,
|
||||
success=success,
|
||||
duration=time.time() - start_time,
|
||||
source=source,
|
||||
budget=budget,
|
||||
max_tokens=max_tokens,
|
||||
)
|
||||
|
||||
# Record duration
|
||||
self.operation_duration.record(duration, attributes)
|
||||
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 operation count
|
||||
self.operation_total.add(1, 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)
|
||||
|
||||
def record_llm_call(
|
||||
self,
|
||||
@@ -628,6 +731,10 @@ 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."""
|
||||
@@ -693,6 +800,192 @@ 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()
|
||||
|
||||
@@ -686,24 +686,20 @@ def ensure_vector_extension(
|
||||
|
||||
if not current_index_info:
|
||||
if table_name == "memory_units" and uses_per_bank_vector_indexes(target_ext):
|
||||
# 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
|
||||
# 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
|
||||
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,9 +20,23 @@ 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
|
||||
@@ -701,6 +715,24 @@ 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
|
||||
@@ -717,14 +749,28 @@ 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:
|
||||
"""
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "hindsight-api-slim"
|
||||
version = "0.8.2"
|
||||
version = "0.8.3"
|
||||
description = "Hindsight: Agent Memory That Works Like Human Memory"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.11"
|
||||
@@ -63,7 +63,7 @@ dependencies = [
|
||||
"langsmith>=0.6.3", # SSRF via tracing header injection fix
|
||||
"protobuf>=6.33.5", # JSON recursion depth bypass fix
|
||||
"pillow>=12.1.1", # Out-of-bounds write in PSD image loading fix
|
||||
"cryptography>=46.0.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
|
||||
"cryptography>=48.0.1", # GHSA-537c-gmf6-5ccf: bundled-OpenSSL OOB read fix needs >=48.0.1. Prior <47 cap (47.0.0 SIGILL on ARM64 Docker/Podman, pyca/cryptography#14733) lifted — 47/48/49 verified importing + RSA sign/verify cleanly on linux/arm64 (Docker on Apple Silicon) and native arm64 macOS; upstream issue closed unconfirmed.
|
||||
"filelock>=3.20.1", # TOCTOU race condition fix
|
||||
"authlib>=1.6.9", # Account takeover/JWS header injection vulnerability fix
|
||||
"pyjwt>=2.12.0", # Accepts unknown crit header extensions fix
|
||||
|
||||
@@ -547,3 +547,36 @@ 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
|
||||
|
||||
@@ -0,0 +1,244 @@
|
||||
"""
|
||||
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,6 +186,43 @@ 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)
|
||||
@@ -199,3 +236,27 @@ 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"}
|
||||
|
||||
@@ -322,3 +322,76 @@ 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}"
|
||||
|
||||
@@ -0,0 +1,109 @@
|
||||
"""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"
|
||||
@@ -25,6 +25,7 @@ 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",
|
||||
@@ -452,6 +453,53 @@ 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
|
||||
|
||||
@@ -580,3 +628,81 @@ 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
|
||||
|
||||
@@ -124,6 +124,10 @@ 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"
|
||||
|
||||
@@ -0,0 +1,70 @@
|
||||
"""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
|
||||
@@ -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
|
||||
from fastapi import Depends, FastAPI, HTTPException, Request
|
||||
from pydantic import BaseModel, model_validator
|
||||
|
||||
from hindsight_api.extensions import PrecheckContext
|
||||
@@ -952,12 +952,23 @@ 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:
|
||||
@@ -1081,3 +1092,86 @@ 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,6 +77,27 @@ 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,6 +346,149 @@ 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,6 +12,8 @@ import subprocess
|
||||
import tempfile
|
||||
import time
|
||||
import uuid
|
||||
from collections.abc import Iterator
|
||||
from contextlib import contextmanager
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
@@ -21,6 +23,7 @@ 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:
|
||||
@@ -38,6 +41,8 @@ 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 = {
|
||||
@@ -64,6 +69,33 @@ 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
|
||||
@@ -101,11 +133,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:
|
||||
host = container.get_container_host_ip()
|
||||
port = container.get_exposed_port(SEAWEEDFS_S3_PORT)
|
||||
with _retry_testcontainers_port_mapping():
|
||||
container.start()
|
||||
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)
|
||||
|
||||
@@ -0,0 +1,103 @@
|
||||
"""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
|
||||
@@ -5,6 +5,7 @@ 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,
|
||||
@@ -408,6 +409,19 @@ 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()`
|
||||
|
||||
@@ -0,0 +1,71 @@
|
||||
"""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)
|
||||
@@ -0,0 +1,82 @@
|
||||
"""
|
||||
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
|
||||
@@ -120,7 +120,7 @@ async def test_anthropic_no_extra_body_omits_key():
|
||||
# ─── Gemini ───────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def _make_gemini_provider(extra_body=None):
|
||||
def _make_gemini_provider(extra_body=None, gemini_service_tier=None):
|
||||
pytest.importorskip("google.genai")
|
||||
with patch("google.genai.Client") as mock_client_cls:
|
||||
mock_client_cls.return_value = MagicMock()
|
||||
@@ -132,6 +132,7 @@ def _make_gemini_provider(extra_body=None):
|
||||
base_url="",
|
||||
model="gemini-2.5-flash",
|
||||
extra_body=extra_body,
|
||||
gemini_service_tier=gemini_service_tier,
|
||||
)
|
||||
provider._client = MagicMock()
|
||||
return provider
|
||||
@@ -176,6 +177,34 @@ 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"
|
||||
|
||||
|
||||
# ─── LiteLLM ──────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
|
||||
@@ -39,11 +39,20 @@ 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:
|
||||
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)')
|
||||
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)')
|
||||
try:
|
||||
yield prefix, schemas
|
||||
finally:
|
||||
|
||||
@@ -134,6 +134,33 @@ 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."""
|
||||
|
||||
@@ -423,6 +423,28 @@ 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.
|
||||
@@ -430,14 +452,6 @@ 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 = [
|
||||
|
||||
@@ -1911,3 +1911,101 @@ 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}],
|
||||
}
|
||||
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 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
|
||||
|
||||
@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
|
||||
|
||||
@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
|
||||
|
||||
@@ -334,6 +334,9 @@ 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,
|
||||
{
|
||||
@@ -343,11 +346,17 @@ def _make_minimal_engine():
|
||||
},
|
||||
clear=False,
|
||||
):
|
||||
from hindsight_api.config import clear_config_cache
|
||||
from hindsight_api.engine.memory_engine import MemoryEngine
|
||||
|
||||
clear_config_cache()
|
||||
return MemoryEngine(db_url="postgresql://localhost/hindsight_test", embeddings=mock_embeddings)
|
||||
engine = MemoryEngine(db_url="postgresql://localhost/hindsight_test", embeddings=mock_embeddings)
|
||||
|
||||
# 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
|
||||
|
||||
|
||||
def test_engine_memory_defense_shares_ext_ctx() -> None:
|
||||
|
||||
@@ -11,6 +11,7 @@ from hindsight_api.metrics import (
|
||||
get_token_bucket,
|
||||
create_metrics_collector,
|
||||
initialize_metrics,
|
||||
normalize_http_endpoint,
|
||||
)
|
||||
|
||||
|
||||
@@ -221,6 +222,17 @@ 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()
|
||||
@@ -324,6 +336,24 @@ 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, inspect
|
||||
from sqlalchemy import create_engine, text
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
@@ -33,8 +33,17 @@ def _upgrade(db_url: str, revision: str) -> None:
|
||||
command.upgrade(_alembic_cfg(db_url), revision)
|
||||
|
||||
|
||||
def _downgrade(db_url: str, revision: str) -> None:
|
||||
command.downgrade(_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()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -43,16 +52,16 @@ def _downgrade(db_url: str, revision: str) -> None:
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def pre_backsweep_db_url():
|
||||
def pre_backsweep_db_url() -> str:
|
||||
"""
|
||||
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 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.
|
||||
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.
|
||||
"""
|
||||
from hindsight_api.pg0 import EmbeddedPostgres
|
||||
|
||||
@@ -63,10 +72,8 @@ def pre_backsweep_db_url():
|
||||
finally:
|
||||
loop.close()
|
||||
|
||||
# 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")
|
||||
_reset_public_schema(url)
|
||||
_upgrade(url, "f6g7h8i9j0k1")
|
||||
return url
|
||||
|
||||
|
||||
@@ -75,7 +82,7 @@ def pre_backsweep_db_url():
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_backsweep_removes_orphans_and_preserves_legit_rows(pre_backsweep_db_url):
|
||||
def test_backsweep_removes_orphans_and_preserves_legit_rows(pre_backsweep_db_url: str) -> None:
|
||||
"""
|
||||
Seed four kinds of rows then apply the backsweep migration and verify:
|
||||
|
||||
@@ -114,15 +121,20 @@ 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, bank, fact_type, sources=None):
|
||||
def insert_mu(
|
||||
uid: uuid.UUID,
|
||||
bank: str,
|
||||
fact_type: str,
|
||||
sources: list[uuid.UUID] | None = None,
|
||||
) -> None:
|
||||
src_arr = "{" + ",".join(str(s) for s in (sources or [])) + "}"
|
||||
conn.execute(
|
||||
text(
|
||||
"""
|
||||
INSERT INTO memory_units
|
||||
(id, bank_id, text, fact_type, source_memory_ids)
|
||||
(id, bank_id, text, event_date, fact_type, source_memory_ids)
|
||||
VALUES
|
||||
(:id, :bank, :text, :ft, CAST(:src AS uuid[]))
|
||||
(:id, :bank, :text, now(), :ft, CAST(:src AS uuid[]))
|
||||
"""
|
||||
),
|
||||
{"id": uid, "bank": bank, "text": "test", "ft": fact_type, "src": src_arr},
|
||||
@@ -150,7 +162,7 @@ def test_backsweep_removes_orphans_and_preserves_legit_rows(pre_backsweep_db_url
|
||||
# --- verify ---
|
||||
with engine.connect() as conn:
|
||||
|
||||
def exists(uid):
|
||||
def exists(uid: uuid.UUID) -> bool:
|
||||
return conn.execute(text("SELECT 1 FROM memory_units WHERE id = :id"), {"id": uid}).fetchone() is not None
|
||||
|
||||
# Must be gone
|
||||
|
||||
@@ -0,0 +1,110 @@
|
||||
"""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)
|
||||
@@ -0,0 +1,78 @@
|
||||
"""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"}]}'
|
||||
@@ -0,0 +1,81 @@
|
||||
"""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)
|
||||
@@ -0,0 +1,168 @@
|
||||
"""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)
|
||||
@@ -13,6 +13,7 @@ Tests cover:
|
||||
import asyncio
|
||||
import json
|
||||
import uuid
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
@@ -70,6 +71,99 @@ 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.
|
||||
@@ -832,6 +926,34 @@ 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."""
|
||||
@@ -947,9 +1069,11 @@ class TestWorkerPoller:
|
||||
|
||||
claimed = await poller.claim_batch()
|
||||
|
||||
# Should claim the retain task (non-consolidation tasks are unaffected)
|
||||
assert len(claimed) == 1
|
||||
assert claimed[0].operation_id == str(retain_op_id)
|
||||
# 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)
|
||||
|
||||
|
||||
class TestWorkerRecovery:
|
||||
@@ -1586,7 +1710,9 @@ class TestDynamicTenantDiscovery:
|
||||
# First claim_batch should call list_tenants
|
||||
claimed1 = await poller.claim_batch()
|
||||
assert mock_extension.list_tenants_calls == 1
|
||||
assert len(claimed1) == 2
|
||||
# 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)}"
|
||||
|
||||
# Add more tasks
|
||||
for i in range(2):
|
||||
@@ -1605,7 +1731,8 @@ class TestDynamicTenantDiscovery:
|
||||
# Second claim_batch should call list_tenants again
|
||||
claimed2 = await poller.claim_batch()
|
||||
assert mock_extension.list_tenants_calls == 2
|
||||
assert len(claimed2) == 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)}"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_poller_picks_up_new_tenants_without_restart(self, pool, backend, clean_operations):
|
||||
@@ -1650,10 +1777,12 @@ class TestDynamicTenantDiscovery:
|
||||
tenant_extension=dynamic_extension,
|
||||
)
|
||||
|
||||
# First poll - only public schema
|
||||
# First poll - only public schema. Filter to our bank — parallel tests
|
||||
# may contribute other claims.
|
||||
claimed1 = await poller.claim_batch()
|
||||
assert len(claimed1) == 1
|
||||
assert claimed1[0].schema is None # public is represented as None
|
||||
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 dynamic_extension.list_tenants_calls == 1
|
||||
|
||||
# Simulate tenant list changing (but we won't add a non-existent schema)
|
||||
@@ -1675,12 +1804,14 @@ class TestDynamicTenantDiscovery:
|
||||
|
||||
# Second poll - list_tenants should be called again
|
||||
claimed2 = await poller.claim_batch()
|
||||
assert len(claimed2) == 1
|
||||
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 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()
|
||||
assert len(claimed3) == 0
|
||||
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 dynamic_extension.list_tenants_calls == 3 # Called again even with no tasks
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
||||
@@ -4,12 +4,12 @@ build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "hindsight-api"
|
||||
version = "0.8.2"
|
||||
version = "0.8.3"
|
||||
description = "Hindsight: Agent Memory That Works Like Human Memory"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.11"
|
||||
dependencies = [
|
||||
"hindsight-api-slim[all]==0.8.2",
|
||||
"hindsight-api-slim[all]==0.8.3",
|
||||
]
|
||||
|
||||
[tool.uv.sources]
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
[package]
|
||||
name = "hindsight-cli"
|
||||
version = "0.8.2"
|
||||
version = "0.8.3"
|
||||
edition = "2021"
|
||||
authors = ["Hindsight Team"]
|
||||
description = "A beautiful CLI for Hindsight - semantic memory system"
|
||||
|
||||
@@ -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.2
|
||||
version: 0.8.3
|
||||
servers:
|
||||
- url: /
|
||||
paths:
|
||||
@@ -3541,7 +3541,7 @@ paths:
|
||||
This endpoint handles file upload, conversion, and memory creation in a single operation.
|
||||
|
||||
**Features:**
|
||||
- Supports PDF, DOCX, PPTX, XLSX, images (with OCR), audio (with transcription)
|
||||
- Supports PDF, DOCX, PPTX, XLSX, images (parser-dependent 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
|
||||
|
||||
@@ -3,7 +3,7 @@ Hindsight HTTP API
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
API version: 0.8.2
|
||||
API version: 0.8.3
|
||||
*/
|
||||
|
||||
// Code generated by OpenAPI Generator (https://openapi-generator.tech); DO NOT EDIT.
|
||||
|
||||
@@ -3,7 +3,7 @@ Hindsight HTTP API
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
API version: 0.8.2
|
||||
API version: 0.8.3
|
||||
*/
|
||||
|
||||
// Code generated by OpenAPI Generator (https://openapi-generator.tech); DO NOT EDIT.
|
||||
|
||||
@@ -3,7 +3,7 @@ Hindsight HTTP API
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
API version: 0.8.2
|
||||
API version: 0.8.3
|
||||
*/
|
||||
|
||||
// Code generated by OpenAPI Generator (https://openapi-generator.tech); DO NOT EDIT.
|
||||
|
||||
@@ -3,7 +3,7 @@ Hindsight HTTP API
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
API version: 0.8.2
|
||||
API version: 0.8.3
|
||||
*/
|
||||
|
||||
// Code generated by OpenAPI Generator (https://openapi-generator.tech); DO NOT EDIT.
|
||||
|
||||
@@ -3,7 +3,7 @@ Hindsight HTTP API
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
API version: 0.8.2
|
||||
API version: 0.8.3
|
||||
*/
|
||||
|
||||
// Code generated by OpenAPI Generator (https://openapi-generator.tech); DO NOT EDIT.
|
||||
|
||||
@@ -3,7 +3,7 @@ Hindsight HTTP API
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
API version: 0.8.2
|
||||
API version: 0.8.3
|
||||
*/
|
||||
|
||||
// Code generated by OpenAPI Generator (https://openapi-generator.tech); DO NOT EDIT.
|
||||
|
||||
@@ -3,7 +3,7 @@ Hindsight HTTP API
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
API version: 0.8.2
|
||||
API version: 0.8.3
|
||||
*/
|
||||
|
||||
// Code generated by OpenAPI Generator (https://openapi-generator.tech); DO NOT EDIT.
|
||||
|
||||
@@ -3,7 +3,7 @@ Hindsight HTTP API
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
API version: 0.8.2
|
||||
API version: 0.8.3
|
||||
*/
|
||||
|
||||
// Code generated by OpenAPI Generator (https://openapi-generator.tech); DO NOT EDIT.
|
||||
@@ -62,7 +62,7 @@ Upload files (PDF, DOCX, etc.), convert them to markdown, and retain as memories
|
||||
This endpoint handles file upload, conversion, and memory creation in a single operation.
|
||||
|
||||
**Features:**
|
||||
- Supports PDF, DOCX, PPTX, XLSX, images (with OCR), audio (with transcription)
|
||||
- Supports PDF, DOCX, PPTX, XLSX, images (parser-dependent 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
|
||||
|
||||
@@ -3,7 +3,7 @@ Hindsight HTTP API
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
API version: 0.8.2
|
||||
API version: 0.8.3
|
||||
*/
|
||||
|
||||
// Code generated by OpenAPI Generator (https://openapi-generator.tech); DO NOT EDIT.
|
||||
|
||||
@@ -3,7 +3,7 @@ Hindsight HTTP API
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
API version: 0.8.2
|
||||
API version: 0.8.3
|
||||
*/
|
||||
|
||||
// Code generated by OpenAPI Generator (https://openapi-generator.tech); DO NOT EDIT.
|
||||
|
||||
@@ -3,7 +3,7 @@ Hindsight HTTP API
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
API version: 0.8.2
|
||||
API version: 0.8.3
|
||||
*/
|
||||
|
||||
// Code generated by OpenAPI Generator (https://openapi-generator.tech); DO NOT EDIT.
|
||||
|
||||
@@ -3,7 +3,7 @@ Hindsight HTTP API
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
API version: 0.8.2
|
||||
API version: 0.8.3
|
||||
*/
|
||||
|
||||
// Code generated by OpenAPI Generator (https://openapi-generator.tech); DO NOT EDIT.
|
||||
|
||||
@@ -3,7 +3,7 @@ Hindsight HTTP API
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
API version: 0.8.2
|
||||
API version: 0.8.3
|
||||
*/
|
||||
|
||||
// Code generated by OpenAPI Generator (https://openapi-generator.tech); DO NOT EDIT.
|
||||
|
||||
@@ -3,7 +3,7 @@ Hindsight HTTP API
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
API version: 0.8.2
|
||||
API version: 0.8.3
|
||||
*/
|
||||
|
||||
// Code generated by OpenAPI Generator (https://openapi-generator.tech); DO NOT EDIT.
|
||||
|
||||
@@ -3,7 +3,7 @@ Hindsight HTTP API
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
API version: 0.8.2
|
||||
API version: 0.8.3
|
||||
*/
|
||||
|
||||
// Code generated by OpenAPI Generator (https://openapi-generator.tech); DO NOT EDIT.
|
||||
@@ -41,7 +41,7 @@ var (
|
||||
queryDescape = strings.NewReplacer( "%5B", "[", "%5D", "]" )
|
||||
)
|
||||
|
||||
// APIClient manages communication with the Hindsight HTTP API API v0.8.2
|
||||
// APIClient manages communication with the Hindsight HTTP API API v0.8.3
|
||||
// In most cases there should be only one, shared, APIClient.
|
||||
type APIClient struct {
|
||||
cfg *Configuration
|
||||
|
||||
@@ -3,7 +3,7 @@ Hindsight HTTP API
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
API version: 0.8.2
|
||||
API version: 0.8.3
|
||||
*/
|
||||
|
||||
// Code generated by OpenAPI Generator (https://openapi-generator.tech); DO NOT EDIT.
|
||||
|
||||
@@ -3,7 +3,7 @@ Hindsight HTTP API
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
API version: 0.8.2
|
||||
API version: 0.8.3
|
||||
*/
|
||||
|
||||
// Code generated by OpenAPI Generator (https://openapi-generator.tech); DO NOT EDIT.
|
||||
|
||||
@@ -3,7 +3,7 @@ Hindsight HTTP API
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
API version: 0.8.2
|
||||
API version: 0.8.3
|
||||
*/
|
||||
|
||||
// Code generated by OpenAPI Generator (https://openapi-generator.tech); DO NOT EDIT.
|
||||
|
||||
@@ -3,7 +3,7 @@ Hindsight HTTP API
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
API version: 0.8.2
|
||||
API version: 0.8.3
|
||||
*/
|
||||
|
||||
// Code generated by OpenAPI Generator (https://openapi-generator.tech); DO NOT EDIT.
|
||||
|
||||
@@ -3,7 +3,7 @@ Hindsight HTTP API
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
API version: 0.8.2
|
||||
API version: 0.8.3
|
||||
*/
|
||||
|
||||
// Code generated by OpenAPI Generator (https://openapi-generator.tech); DO NOT EDIT.
|
||||
|
||||
@@ -3,7 +3,7 @@ Hindsight HTTP API
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
API version: 0.8.2
|
||||
API version: 0.8.3
|
||||
*/
|
||||
|
||||
// Code generated by OpenAPI Generator (https://openapi-generator.tech); DO NOT EDIT.
|
||||
|
||||
@@ -3,7 +3,7 @@ Hindsight HTTP API
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
API version: 0.8.2
|
||||
API version: 0.8.3
|
||||
*/
|
||||
|
||||
// Code generated by OpenAPI Generator (https://openapi-generator.tech); DO NOT EDIT.
|
||||
|
||||
@@ -3,7 +3,7 @@ Hindsight HTTP API
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
API version: 0.8.2
|
||||
API version: 0.8.3
|
||||
*/
|
||||
|
||||
// Code generated by OpenAPI Generator (https://openapi-generator.tech); DO NOT EDIT.
|
||||
|
||||
@@ -3,7 +3,7 @@ Hindsight HTTP API
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
API version: 0.8.2
|
||||
API version: 0.8.3
|
||||
*/
|
||||
|
||||
// Code generated by OpenAPI Generator (https://openapi-generator.tech); DO NOT EDIT.
|
||||
|
||||
@@ -3,7 +3,7 @@ Hindsight HTTP API
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
API version: 0.8.2
|
||||
API version: 0.8.3
|
||||
*/
|
||||
|
||||
// Code generated by OpenAPI Generator (https://openapi-generator.tech); DO NOT EDIT.
|
||||
|
||||
@@ -3,7 +3,7 @@ Hindsight HTTP API
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
API version: 0.8.2
|
||||
API version: 0.8.3
|
||||
*/
|
||||
|
||||
// Code generated by OpenAPI Generator (https://openapi-generator.tech); DO NOT EDIT.
|
||||
|
||||
@@ -3,7 +3,7 @@ Hindsight HTTP API
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
API version: 0.8.2
|
||||
API version: 0.8.3
|
||||
*/
|
||||
|
||||
// Code generated by OpenAPI Generator (https://openapi-generator.tech); DO NOT EDIT.
|
||||
|
||||
@@ -3,7 +3,7 @@ Hindsight HTTP API
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
API version: 0.8.2
|
||||
API version: 0.8.3
|
||||
*/
|
||||
|
||||
// Code generated by OpenAPI Generator (https://openapi-generator.tech); DO NOT EDIT.
|
||||
|
||||
@@ -3,7 +3,7 @@ Hindsight HTTP API
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
API version: 0.8.2
|
||||
API version: 0.8.3
|
||||
*/
|
||||
|
||||
// Code generated by OpenAPI Generator (https://openapi-generator.tech); DO NOT EDIT.
|
||||
|
||||
@@ -3,7 +3,7 @@ Hindsight HTTP API
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
API version: 0.8.2
|
||||
API version: 0.8.3
|
||||
*/
|
||||
|
||||
// Code generated by OpenAPI Generator (https://openapi-generator.tech); DO NOT EDIT.
|
||||
|
||||
@@ -3,7 +3,7 @@ Hindsight HTTP API
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
API version: 0.8.2
|
||||
API version: 0.8.3
|
||||
*/
|
||||
|
||||
// Code generated by OpenAPI Generator (https://openapi-generator.tech); DO NOT EDIT.
|
||||
|
||||
@@ -3,7 +3,7 @@ Hindsight HTTP API
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
API version: 0.8.2
|
||||
API version: 0.8.3
|
||||
*/
|
||||
|
||||
// Code generated by OpenAPI Generator (https://openapi-generator.tech); DO NOT EDIT.
|
||||
|
||||
@@ -3,7 +3,7 @@ Hindsight HTTP API
|
||||
|
||||
HTTP API for Hindsight
|
||||
|
||||
API version: 0.8.2
|
||||
API version: 0.8.3
|
||||
*/
|
||||
|
||||
// Code generated by OpenAPI Generator (https://openapi-generator.tech); DO NOT EDIT.
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user