fix: ui shows only 1000 memories

* feat: Record LLM token metrics via Prometheus

Wire up the existing token metrics infrastructure to actually record
token usage from LLM calls. The MetricsCollector already had
record_tokens() method and Prometheus counters (hindsight.tokens.input,
hindsight.tokens.output), but they were never being populated.

Changes:
- Import get_metrics_collector in llm_wrapper.py
- Call record_tokens() after successful LLM calls for:
  - OpenAI/Groq (using response.usage.prompt_tokens, completion_tokens)
  - Anthropic (using response.usage.input_tokens, output_tokens)
  - Gemini (using response.usage_metadata.prompt_token_count, candidates_token_count)
- Add test file to verify token metrics are recorded

Note: Ollama's native API doesn't return token usage, so metrics
are not recorded for that provider.

The token metrics will now be available via /metrics endpoint:
- hindsight_tokens_input_total
- hindsight_tokens_output_total

* feat: add per-request token usage tracking to retain and reflect endpoints

- Add TokenUsage model with input_tokens, output_tokens, total_tokens
- Return usage metrics in retain response (sync operations only)
- Return usage metrics in reflect response
- Update Python, TypeScript, and Rust clients
- Add API documentation for usage fields
- Add changelog entry

CLAUDE.md

@@ -84,29 +84,15 @@ PostgreSQL with pgvector. Schema managed via Alembic migrations in `hindsight-ap
 
 Key tables: `banks`, `memory_units`, `documents`, `entities`, `entity_links`
 
-### Database Backups (IMPORTANT)
-**Before any operation that may affect the database, run a backup:**
-```bash
-docker exec hindsight /backups/backup.sh
-```
-
-Operations requiring backup:
-- Running database migrations
-- Modifying Alembic migration files
-- Rebuilding Docker images
-- Resetting or recreating containers
-- Any schema changes
-- Bulk data operations
-
-Backups are stored in `~/hindsight-backups/` on the host.
-
-To restore:
-```bash
-docker exec -it hindsight /backups/restore.sh <backup-file.sql.gz>
-```
-
 ## Key Conventions
 
+### Code Quality
+**Always run the lint script after making Python or TypeScript/Node changes:**
+```bash
+./scripts/hooks/lint.sh
+```
+This runs the same checks as the pre-commit hook (Ruff for Python, ESLint/Prettier for TypeScript).
+
 ### Memory Banks
 - Each bank is isolated (no cross-bank data access)
 - Banks have dispositions (skepticism, literalism, empathy traits 1-5) affecting reflect
@@ -127,6 +113,29 @@ docker exec -it hindsight /backups/restore.sh <backup-file.sql.gz>
 - Next.js App Router for control plane
 - Tailwind CSS with shadcn/ui components
 
+### Adding New API Configuration Flags
+
+When adding a new environment variable configuration:
+
+1. **config.py** (`hindsight-api/hindsight_api/config.py`):
+   - Add `ENV_*` constant for the environment variable name
+   - Add `DEFAULT_*` constant for the default value
+   - Add field to `HindsightConfig` dataclass
+   - Add initialization in `from_env()` method
+
+2. **main.py** (`hindsight-api/hindsight_api/main.py`):
+   - Add field to the manual `HindsightConfig()` constructor call (search for "CLI override")
+
+3. **Use the config** in code:
+   ```python
+   from ...config import get_config
+   config = get_config()
+   value = config.your_new_field
+   ```
+
+4. **Documentation** (`hindsight-docs/docs/developer/configuration.md`):
+   - Add to appropriate section table with Variable, Description, Default
+
 ## Environment Setup
 
 ```bash

helm/hindsight/templates/_helpers.tpl

@@ -110,3 +110,14 @@ API URL for control plane
 {{- define "hindsight.apiUrl" -}}
 {{- printf "http://%s-api:%d" (include "hindsight.fullname" .) (.Values.api.service.port | int) }}
 {{- end }}
+
+{{/*
+Get the name of the secret to use
+*/}}
+{{- define "hindsight.secretName" -}}
+{{- if .Values.existingSecret }}
+{{- .Values.existingSecret }}
+{{- else }}
+{{- printf "%s-secret" (include "hindsight.fullname" .) }}
+{{- end }}
+{{- end }}

helm/hindsight/templates/api-deployment.yaml

@@ -15,7 +15,9 @@ spec:
   template:
     metadata:
       annotations:
+        {{- if not .Values.existingSecret }}
         checksum/secret: {{ include (print $.Template.BasePath "/secret.yaml") . | sha256sum }}
+        {{- end }}
         {{- with .Values.podAnnotations }}
         {{- toYaml . | nindent 8 }}
         {{- end }}
@@ -37,27 +39,36 @@ spec:
         - name: http
           containerPort: {{ .Values.api.service.targetPort }}
           protocol: TCP
+        {{- if .Values.existingSecret }}
+        envFrom:
+        - secretRef:
+            name: {{ .Values.existingSecret }}
+        {{- end }}
         env:
-        - name: HINDSIGHT_API_DATABASE_URL
-          value: {{ include "hindsight.databaseUrl" . | quote }}
+        {{- /* POSTGRES_PASSWORD must be defined before DATABASE_URL for $(VAR) interpolation */}}
         {{- if not .Values.postgresql.enabled }}
         - name: POSTGRES_PASSWORD
           valueFrom:
             secretKeyRef:
-              name: {{ include "hindsight.fullname" . }}-secret
+              name: {{ include "hindsight.secretName" . }}
               key: postgres-password
         {{- end }}
+        - name: HINDSIGHT_API_DATABASE_URL
+          value: {{ include "hindsight.databaseUrl" . | quote }}
         {{- range $key, $value := .Values.api.env }}
         - name: {{ $key }}
           value: {{ $value | quote }}
         {{- end }}
+        {{- /* Only use api.secrets when not using existingSecret (for chart-managed secrets) */}}
+        {{- if not .Values.existingSecret }}
         {{- range $key, $value := .Values.api.secrets }}
         - name: {{ $key }}
           valueFrom:
             secretKeyRef:
-              name: {{ include "hindsight.fullname" $ }}-secret
+              name: {{ include "hindsight.secretName" $ }}
               key: {{ $key }}
         {{- end }}
+        {{- end }}
         livenessProbe:
           {{- toYaml .Values.api.livenessProbe | nindent 10 }}
         readinessProbe:

helm/hindsight/templates/controlplane-deployment.yaml

@@ -15,7 +15,9 @@ spec:
   template:
     metadata:
       annotations:
+        {{- if not .Values.existingSecret }}
         checksum/secret: {{ include (print $.Template.BasePath "/secret.yaml") . | sha256sum }}
+        {{- end }}
         {{- with .Values.podAnnotations }}
         {{- toYaml . | nindent 8 }}
         {{- end }}
@@ -37,6 +39,11 @@ spec:
         - name: http
           containerPort: {{ .Values.controlPlane.service.targetPort }}
           protocol: TCP
+        {{- if .Values.existingSecret }}
+        envFrom:
+        - secretRef:
+            name: {{ .Values.existingSecret }}
+        {{- end }}
         env:
         - name: HINDSIGHT_CP_DATAPLANE_API_URL
           value: {{ include "hindsight.apiUrl" . | quote }}
@@ -44,13 +51,16 @@ spec:
         - name: {{ $key }}
           value: {{ $value | quote }}
         {{- end }}
+        {{- /* Only use controlPlane.secrets when not using existingSecret (for chart-managed secrets) */}}
+        {{- if not .Values.existingSecret }}
         {{- range $key, $value := .Values.controlPlane.secrets }}
         - name: {{ $key }}
           valueFrom:
             secretKeyRef:
-              name: {{ include "hindsight.fullname" $ }}-secret
+              name: {{ include "hindsight.secretName" $ }}
               key: {{ $key }}
         {{- end }}
+        {{- end }}
         livenessProbe:
           {{- toYaml .Values.controlPlane.livenessProbe | nindent 10 }}
         readinessProbe:

helm/hindsight/templates/secret.yaml

@@ -1,7 +1,8 @@
+{{- if not .Values.existingSecret }}
 apiVersion: v1
 kind: Secret
 metadata:
-  name: {{ include "hindsight.fullname" . }}-secret
+  name: {{ include "hindsight.secretName" . }}
   labels:
     {{- include "hindsight.labels" . | nindent 4 }}
 type: Opaque
@@ -15,3 +16,4 @@ data:
   {{- if and (not .Values.postgresql.enabled) .Values.postgresql.external.password }}
   postgres-password: {{ .Values.postgresql.external.password | b64enc | quote }}
   {{- end }}
+{{- end }}

helm/hindsight/values.yaml

@@ -3,6 +3,15 @@
 # Chart version - use this to set a consistent image tag across all components
 version: "0.1.1"
 
+# Use an existing secret instead of creating one from values
+# When set, all keys from this secret are injected as environment variables via envFrom
+# Required keys:
+#   - postgres-password: PostgreSQL password (when postgresql.enabled=false)
+# Optional keys (any key becomes an env var):
+#   - HINDSIGHT_API_LLM_API_KEY: API key for LLM provider
+#   - Any other env vars you want to inject
+# existingSecret: "my-hindsight-secret"
+
 # Global settings
 replicaCount: 1
 

hindsight-api/hindsight_api/admin/__init__.py

@@ -0,0 +1 @@
+# Admin CLI for Hindsight

hindsight-api/hindsight_api/admin/cli.py

@@ -0,0 +1,252 @@
+"""
+Hindsight Admin CLI - backup and restore operations.
+"""
+
+import asyncio
+import io
+import json
+import logging
+import zipfile
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any
+
+import asyncpg
+import typer
+
+from ..config import HindsightConfig
+from ..pg0 import parse_pg0_url, resolve_database_url
+
+
+def _fq_table(table: str, schema: str) -> str:
+    """Get fully-qualified table name with schema prefix."""
+    return f"{schema}.{table}"
+
+
+# Setup logging
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(message)s",
+)
+logger = logging.getLogger(__name__)
+
+app = typer.Typer(name="hindsight-admin", help="Hindsight administrative commands")
+
+# Tables to backup/restore in dependency order
+# Import must happen in this order due to foreign key constraints
+BACKUP_TABLES = [
+    "banks",
+    "documents",
+    "entities",
+    "chunks",
+    "memory_units",
+    "unit_entities",
+    "entity_cooccurrences",
+    "memory_links",
+]
+
+MANIFEST_VERSION = "1"
+
+
+async def _backup(database_url: str, output_path: Path, schema: str = "public") -> dict[str, Any]:
+    """Backup all tables to a zip file using binary COPY protocol."""
+    conn = await asyncpg.connect(database_url)
+    try:
+        tables: dict[str, Any] = {}
+        manifest: dict[str, Any] = {
+            "version": MANIFEST_VERSION,
+            "created_at": datetime.now(timezone.utc).isoformat(),
+            "schema": schema,
+            "tables": tables,
+        }
+
+        # Use a transaction with REPEATABLE READ isolation to get a consistent
+        # snapshot across all tables. This prevents race conditions where
+        # entity_cooccurrences could reference entities created after the
+        # entities table was backed up.
+        async with conn.transaction(isolation="repeatable_read"):
+            with zipfile.ZipFile(output_path, "w", zipfile.ZIP_DEFLATED) as zf:
+                for i, table in enumerate(BACKUP_TABLES, 1):
+                    typer.echo(f"  [{i}/{len(BACKUP_TABLES)}] Backing up {table}...", nl=False)
+
+                    buffer = io.BytesIO()
+
+                    # Use binary COPY for exact type preservation
+                    # asyncpg requires schema_name as separate parameter
+                    await conn.copy_from_table(table, schema_name=schema, output=buffer, format="binary")
+
+                    data = buffer.getvalue()
+                    zf.writestr(f"{table}.bin", data)
+
+                    # Get row count for manifest
+                    qualified_table = _fq_table(table, schema)
+                    row_count = await conn.fetchval(f"SELECT COUNT(*) FROM {qualified_table}")
+                    tables[table] = {
+                        "rows": row_count,
+                        "size_bytes": len(data),
+                    }
+
+                    typer.echo(f" {row_count} rows")
+
+                zf.writestr("manifest.json", json.dumps(manifest, indent=2))
+
+        return manifest
+    finally:
+        await conn.close()
+
+
+async def _restore(database_url: str, input_path: Path, schema: str = "public") -> dict[str, Any]:
+    """Restore all tables from a zip file using binary COPY protocol."""
+    conn = await asyncpg.connect(database_url)
+    try:
+        with zipfile.ZipFile(input_path, "r") as zf:
+            # Read and validate manifest
+            manifest: dict[str, Any] = json.loads(zf.read("manifest.json"))
+            if manifest.get("version") != MANIFEST_VERSION:
+                raise ValueError(f"Unsupported backup version: {manifest.get('version')}")
+
+            # Use a transaction for atomic restore - either all tables are
+            # restored or none are, preventing partial/inconsistent state.
+            async with conn.transaction():
+                typer.echo("  Clearing existing data...")
+                # Truncate tables in reverse order (respects FK constraints)
+                for table in reversed(BACKUP_TABLES):
+                    qualified_table = _fq_table(table, schema)
+                    await conn.execute(f"TRUNCATE TABLE {qualified_table} CASCADE")
+
+                # Restore tables in forward order
+                for i, table in enumerate(BACKUP_TABLES, 1):
+                    filename = f"{table}.bin"
+                    if filename not in zf.namelist():
+                        typer.echo(f"  [{i}/{len(BACKUP_TABLES)}] {table}: skipped (not in backup)")
+                        continue
+
+                    expected_rows = manifest["tables"].get(table, {}).get("rows", "?")
+                    typer.echo(f"  [{i}/{len(BACKUP_TABLES)}] Restoring {table}... {expected_rows} rows")
+
+                    data = zf.read(filename)
+                    buffer = io.BytesIO(data)
+                    # asyncpg requires schema_name as separate parameter
+                    await conn.copy_to_table(table, schema_name=schema, source=buffer, format="binary")
+
+                # Refresh materialized view
+                typer.echo("  Refreshing materialized views...")
+                await conn.execute(f"REFRESH MATERIALIZED VIEW {_fq_table('memory_units_bm25', schema)}")
+
+        return manifest
+    finally:
+        await conn.close()
+
+
+async def _run_backup(db_url: str, output: Path, schema: str = "public") -> dict[str, Any]:
+    """Resolve database URL and run backup."""
+    is_pg0, instance_name, _ = parse_pg0_url(db_url)
+    if is_pg0:
+        typer.echo(f"Starting embedded PostgreSQL (instance: {instance_name})...")
+    resolved_url = await resolve_database_url(db_url)
+    return await _backup(resolved_url, output, schema)
+
+
+async def _run_restore(db_url: str, input_file: Path, schema: str = "public") -> dict[str, Any]:
+    """Resolve database URL and run restore."""
+    is_pg0, instance_name, _ = parse_pg0_url(db_url)
+    if is_pg0:
+        typer.echo(f"Starting embedded PostgreSQL (instance: {instance_name})...")
+    resolved_url = await resolve_database_url(db_url)
+    return await _restore(resolved_url, input_file, schema)
+
+
+@app.command()
+def backup(
+    output: Path = typer.Argument(..., help="Output file path (.zip)"),
+    schema: str = typer.Option("public", "--schema", "-s", help="Database schema to backup"),
+):
+    """Backup the Hindsight database to a zip file."""
+    config = HindsightConfig.from_env()
+
+    if not config.database_url:
+        typer.echo("Error: Database URL not configured.", err=True)
+        typer.echo("Set HINDSIGHT_API_DATABASE_URL environment variable.", err=True)
+        raise typer.Exit(1)
+
+    if output.suffix != ".zip":
+        output = output.with_suffix(".zip")
+
+    typer.echo(f"Backing up database (schema: {schema}) to {output}...")
+
+    manifest = asyncio.run(_run_backup(config.database_url, output, schema))
+
+    total_rows = sum(t["rows"] for t in manifest["tables"].values())
+    typer.echo(f"Backed up {total_rows} rows across {len(BACKUP_TABLES)} tables")
+    typer.echo(f"Backup saved to {output}")
+
+
+@app.command()
+def restore(
+    input_file: Path = typer.Argument(..., help="Input backup file (.zip)"),
+    schema: str = typer.Option("public", "--schema", "-s", help="Database schema to restore to"),
+    yes: bool = typer.Option(False, "--yes", "-y", help="Skip confirmation prompt"),
+):
+    """Restore the database from a backup file. WARNING: This deletes all existing data."""
+    config = HindsightConfig.from_env()
+
+    if not config.database_url:
+        typer.echo("Error: Database URL not configured.", err=True)
+        typer.echo("Set HINDSIGHT_API_DATABASE_URL environment variable.", err=True)
+        raise typer.Exit(1)
+
+    if not input_file.exists():
+        typer.echo(f"Error: File not found: {input_file}", err=True)
+        raise typer.Exit(1)
+
+    if not yes:
+        typer.confirm(
+            "This will DELETE all existing data and replace it with the backup. Continue?",
+            abort=True,
+        )
+
+    typer.echo(f"Restoring database (schema: {schema}) from {input_file}...")
+
+    manifest = asyncio.run(_run_restore(config.database_url, input_file, schema))
+
+    total_rows = sum(t["rows"] for t in manifest["tables"].values())
+    typer.echo(f"Restored {total_rows} rows across {len(BACKUP_TABLES)} tables")
+    typer.echo("Restore complete")
+
+
+async def _run_migration(db_url: str, schema: str = "public") -> None:
+    """Resolve database URL and run migrations."""
+    from ..migrations import run_migrations
+
+    is_pg0, instance_name, _ = parse_pg0_url(db_url)
+    if is_pg0:
+        typer.echo(f"Starting embedded PostgreSQL (instance: {instance_name})...")
+    resolved_url = await resolve_database_url(db_url)
+    run_migrations(resolved_url, schema=schema)
+
+
+@app.command(name="run-db-migration")
+def run_db_migration(
+    schema: str = typer.Option("public", "--schema", "-s", help="Database schema to run migrations on"),
+):
+    """Run database migrations to the latest version."""
+    config = HindsightConfig.from_env()
+
+    if not config.database_url:
+        typer.echo("Error: Database URL not configured.", err=True)
+        typer.echo("Set HINDSIGHT_API_DATABASE_URL environment variable.", err=True)
+        raise typer.Exit(1)
+
+    typer.echo(f"Running database migrations (schema: {schema})...")
+
+    asyncio.run(_run_migration(config.database_url, schema))
+
+    typer.echo("Database migrations completed successfully")
+
+
+def main():
+    app()
+
+
+if __name__ == "__main__":
+    main()

hindsight-api/hindsight_api/api/http.py

@@ -36,7 +36,7 @@ from pydantic import BaseModel, ConfigDict, Field, field_validator
 from hindsight_api import MemoryEngine
 from hindsight_api.engine.db_utils import acquire_with_retry
 from hindsight_api.engine.memory_engine import Budget, fq_table
-from hindsight_api.engine.response_models import VALID_RECALL_FACT_TYPES
+from hindsight_api.engine.response_models import VALID_RECALL_FACT_TYPES, TokenUsage
 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.models import RequestContext
@@ -364,7 +364,15 @@ class RetainResponse(BaseModel):
 
     model_config = ConfigDict(
         populate_by_name=True,
-        json_schema_extra={"example": {"success": True, "bank_id": "user123", "items_count": 2, "async": False}},
+        json_schema_extra={
+            "example": {
+                "success": True,
+                "bank_id": "user123",
+                "items_count": 2,
+                "async": False,
+                "usage": {"input_tokens": 500, "output_tokens": 100, "total_tokens": 600},
+            }
+        },
     )
 
     success: bool
@@ -373,6 +381,10 @@ class RetainResponse(BaseModel):
     is_async: bool = Field(
         alias="async", serialization_alias="async", description="Whether the operation was processed asynchronously"
     )
+    usage: TokenUsage | None = Field(
+        default=None,
+        description="Token usage metrics for LLM calls during fact extraction (only present for synchronous operations)",
+    )
 
 
 class FactsIncludeOptions(BaseModel):
@@ -472,6 +484,7 @@ class ReflectResponse(BaseModel):
                     "summary": "AI is transformative",
                     "key_points": ["Used in healthcare", "Discussed recently"],
                 },
+                "usage": {"input_tokens": 1500, "output_tokens": 500, "total_tokens": 2000},
             }
         }
     )
@@ -482,6 +495,10 @@ class ReflectResponse(BaseModel):
         default=None,
         description="Structured output parsed according to the request's response_schema. Only present when response_schema was provided in the request.",
     )
+    usage: TokenUsage | None = Field(
+        default=None,
+        description="Token usage metrics for LLM calls during reflection.",
+    )
 
 
 class BanksResponse(BaseModel):
@@ -630,6 +647,7 @@ class GraphDataResponse(BaseModel):
                     }
                 ],
                 "total_units": 2,
+                "limit": 1000,
             }
         }
     )
@@ -638,6 +656,7 @@ class GraphDataResponse(BaseModel):
     edges: list[dict[str, Any]]
     table_rows: list[dict[str, Any]]
     total_units: int
+    limit: int
 
 
 class ListMemoryUnitsResponse(BaseModel):
@@ -1049,16 +1068,19 @@ def _register_routes(app: FastAPI):
         "/v1/default/banks/{bank_id}/graph",
         response_model=GraphDataResponse,
         summary="Get memory graph data",
-        description="Retrieve graph data for visualization, optionally filtered by type (world/experience/opinion). Limited to 1000 most recent items.",
+        description="Retrieve graph data for visualization, optionally filtered by type (world/experience/opinion).",
         operation_id="get_graph",
         tags=["Memory"],
     )
     async def api_graph(
-        bank_id: str, type: str | None = None, request_context: RequestContext = Depends(get_request_context)
+        bank_id: str,
+        type: str | None = None,
+        limit: int = 1000,
+        request_context: RequestContext = Depends(get_request_context),
     ):
         """Get graph data from database, filtered by bank_id and optionally by type."""
         try:
-            data = await app.state.memory.get_graph_data(bank_id, type, request_context=request_context)
+            data = await app.state.memory.get_graph_data(bank_id, type, limit=limit, request_context=request_context)
             return data
         except (AuthenticationError, HTTPException):
             raise
@@ -1290,6 +1312,7 @@ def _register_routes(app: FastAPI):
                 text=core_result.text,
                 based_on=based_on_facts,
                 structured_output=core_result.structured_output,
+                usage=core_result.usage,
             )
 
         except OperationValidationError as e:
@@ -2016,12 +2039,12 @@ def _register_routes(app: FastAPI):
             else:
                 # Synchronous processing: wait for completion (record metrics)
                 with metrics.record_operation("retain", bank_id=bank_id):
-                    result = await app.state.memory.retain_batch_async(
-                        bank_id=bank_id, contents=contents, request_context=request_context
+                    result, usage = await app.state.memory.retain_batch_async(
+                        bank_id=bank_id, contents=contents, request_context=request_context, return_usage=True
                     )
 
                 return RetainResponse.model_validate(
-                    {"success": True, "bank_id": bank_id, "items_count": len(contents), "async": False}
+                    {"success": True, "bank_id": bank_id, "items_count": len(contents), "async": False, "usage": usage}
                 )
         except OperationValidationError as e:
             raise HTTPException(status_code=e.status_code, detail=e.reason)

hindsight-api/hindsight_api/config.py

@@ -8,6 +8,11 @@ import logging
 import os
 from dataclasses import dataclass
 
+from dotenv import find_dotenv, load_dotenv
+
+# Load .env file, searching current and parent directories (overrides existing env vars)
+load_dotenv(find_dotenv(usecwd=True), override=True)
+
 logger = logging.getLogger(__name__)
 
 # Environment variable names
@@ -42,10 +47,16 @@ ENV_MCP_INSTRUCTIONS = "HINDSIGHT_API_MCP_INSTRUCTIONS"
 ENV_OBSERVATION_MIN_FACTS = "HINDSIGHT_API_OBSERVATION_MIN_FACTS"
 ENV_OBSERVATION_TOP_ENTITIES = "HINDSIGHT_API_OBSERVATION_TOP_ENTITIES"
 
+# Retain settings
+ENV_RETAIN_MAX_COMPLETION_TOKENS = "HINDSIGHT_API_RETAIN_MAX_COMPLETION_TOKENS"
+
 # Optimization flags
 ENV_SKIP_LLM_VERIFICATION = "HINDSIGHT_API_SKIP_LLM_VERIFICATION"
 ENV_LAZY_RERANKER = "HINDSIGHT_API_LAZY_RERANKER"
 
+# Database migrations
+ENV_RUN_MIGRATIONS_ON_STARTUP = "HINDSIGHT_API_RUN_MIGRATIONS_ON_STARTUP"
+
 # Default values
 DEFAULT_DATABASE_URL = "pg0"
 DEFAULT_LLM_PROVIDER = "openai"
@@ -72,6 +83,12 @@ DEFAULT_MCP_LOCAL_BANK_ID = "mcp"
 DEFAULT_OBSERVATION_MIN_FACTS = 5  # Min facts required to generate entity observations
 DEFAULT_OBSERVATION_TOP_ENTITIES = 5  # Max entities to process per retain batch
 
+# Retain settings
+DEFAULT_RETAIN_MAX_COMPLETION_TOKENS = 64000  # Max tokens for fact extraction LLM call
+
+# Database migrations
+DEFAULT_RUN_MIGRATIONS_ON_STARTUP = True
+
 # Default MCP tool descriptions (can be customized via env vars)
 DEFAULT_MCP_RETAIN_DESCRIPTION = """Store important information to long-term memory.
 
@@ -134,10 +151,16 @@ class HindsightConfig:
     observation_min_facts: int
     observation_top_entities: int
 
+    # Retain settings
+    retain_max_completion_tokens: int
+
     # Optimization flags
     skip_llm_verification: bool
     lazy_reranker: bool
 
+    # Database migrations
+    run_migrations_on_startup: bool
+
     @classmethod
     def from_env(cls) -> "HindsightConfig":
         """Create configuration from environment variables."""
@@ -174,6 +197,12 @@ class HindsightConfig:
             observation_top_entities=int(
                 os.getenv(ENV_OBSERVATION_TOP_ENTITIES, str(DEFAULT_OBSERVATION_TOP_ENTITIES))
             ),
+            # Retain settings
+            retain_max_completion_tokens=int(
+                os.getenv(ENV_RETAIN_MAX_COMPLETION_TOKENS, str(DEFAULT_RETAIN_MAX_COMPLETION_TOKENS))
+            ),
+            # Database migrations
+            run_migrations_on_startup=os.getenv(ENV_RUN_MIGRATIONS_ON_STARTUP, "true").lower() == "true",
         )
 
     def get_llm_base_url(self) -> str:
@@ -220,6 +249,19 @@ class HindsightConfig:
         logger.info(f"Graph retriever: {self.graph_retriever}")
 
 
+# Cached config instance
+_config_cache: HindsightConfig | None = None
+
+
 def get_config() -> HindsightConfig:
-    """Get the current configuration from environment variables."""
-    return HindsightConfig.from_env()
+    """Get the cached configuration, loading from environment on first call."""
+    global _config_cache
+    if _config_cache is None:
+        _config_cache = HindsightConfig.from_env()
+    return _config_cache
+
+
+def clear_config_cache() -> None:
+    """Clear the config cache. Useful for testing or reloading config."""
+    global _config_cache
+    _config_cache = None

hindsight-api/hindsight_api/engine/interface.py

@@ -289,6 +289,7 @@ class MemoryEngineInterface(ABC):
         bank_id: str,
         *,
         fact_type: str | None = None,
+        limit: int = 1000,
         request_context: "RequestContext",
     ) -> dict[str, Any]:
         """
@@ -297,10 +298,11 @@ class MemoryEngineInterface(ABC):
         Args:
             bank_id: The memory bank ID.
             fact_type: Filter by fact type.
+            limit: Maximum number of items to return (default: 1000).
             request_context: Request context for authentication.
 
         Returns:
-            Dict with nodes, edges, table_rows, total_units.
+            Dict with nodes, edges, table_rows, total_units, limit.
         """
         ...
 

hindsight-api/hindsight_api/engine/llm_wrapper.py

@@ -23,6 +23,8 @@ from ..config import (
     ENV_LLM_MAX_CONCURRENT,
     ENV_LLM_TIMEOUT,
 )
+from ..metrics import get_metrics_collector
+from .response_models import TokenUsage
 
 # Seed applied to every Groq request for deterministic behavior.
 DEFAULT_LLM_SEED = 4242
@@ -174,6 +176,7 @@ class LLMProvider:
         max_backoff: float = 60.0,
         skip_validation: bool = False,
         strict_schema: bool = False,
+        return_usage: bool = False,
     ) -> Any:
         """
         Make an LLM API call with retry logic.
@@ -189,9 +192,11 @@ class LLMProvider:
             max_backoff: Maximum backoff time in seconds.
             skip_validation: Return raw JSON without Pydantic validation.
             strict_schema: Use strict JSON schema enforcement (OpenAI only). Guarantees all required fields.
+            return_usage: If True, return tuple (result, TokenUsage) instead of just result.
 
         Returns:
-            Parsed response if response_format is provided, otherwise text content.
+            If return_usage=False: Parsed response if response_format is provided, otherwise text content.
+            If return_usage=True: Tuple of (result, TokenUsage) with token counts from the LLM call.
 
         Raises:
             OutputTooLongError: If output exceeds token limits.
@@ -203,7 +208,14 @@ class LLMProvider:
             # Handle Gemini provider separately
             if self.provider == "gemini":
                 return await self._call_gemini(
-                    messages, response_format, max_retries, initial_backoff, max_backoff, skip_validation, start_time
+                    messages,
+                    response_format,
+                    max_retries,
+                    initial_backoff,
+                    max_backoff,
+                    skip_validation,
+                    start_time,
+                    return_usage,
                 )
 
             # Handle Anthropic provider separately
@@ -217,6 +229,7 @@ class LLMProvider:
                     max_backoff,
                     skip_validation,
                     start_time,
+                    return_usage,
                 )
 
             # Handle Ollama with native API for structured output (better schema enforcement)
@@ -231,6 +244,7 @@ class LLMProvider:
                     max_backoff,
                     skip_validation,
                     start_time,
+                    return_usage,
                 )
 
             call_params = {
@@ -379,21 +393,41 @@ class LLMProvider:
                         response = await self._client.chat.completions.create(**call_params)
                         result = response.choices[0].message.content
 
-                    # Log slow calls
+                    # Record token usage metrics
                     duration = time.time() - start_time
                     usage = response.usage
-                    if duration > 10.0:
-                        ratio = max(1, usage.completion_tokens) / usage.prompt_tokens
+                    input_tokens = usage.prompt_tokens or 0 if usage else 0
+                    output_tokens = usage.completion_tokens or 0 if usage else 0
+                    total_tokens = usage.total_tokens or 0 if usage else 0
+
+                    if usage:
+                        get_metrics_collector().record_tokens(
+                            operation=scope,
+                            bank_id="llm",
+                            input_tokens=input_tokens,
+                            output_tokens=output_tokens,
+                        )
+
+                    # Log slow calls
+                    if duration > 10.0 and usage:
+                        ratio = max(1, output_tokens) / max(1, input_tokens)
                         cached_tokens = 0
                         if hasattr(usage, "prompt_tokens_details") and usage.prompt_tokens_details:
                             cached_tokens = getattr(usage.prompt_tokens_details, "cached_tokens", 0) or 0
                         cache_info = f", cached_tokens={cached_tokens}" if cached_tokens > 0 else ""
                         logger.info(
                             f"slow llm call: model={self.provider}/{self.model}, "
-                            f"input_tokens={usage.prompt_tokens}, output_tokens={usage.completion_tokens}, "
-                            f"total_tokens={usage.total_tokens}{cache_info}, time={duration:.3f}s, ratio out/in={ratio:.2f}"
+                            f"input_tokens={input_tokens}, output_tokens={output_tokens}, "
+                            f"total_tokens={total_tokens}{cache_info}, time={duration:.3f}s, ratio out/in={ratio:.2f}"
                         )
 
+                    if return_usage:
+                        token_usage = TokenUsage(
+                            input_tokens=input_tokens,
+                            output_tokens=output_tokens,
+                            total_tokens=total_tokens,
+                        )
+                        return result, token_usage
                     return result
 
                 except LengthFinishReasonError as e:
@@ -452,6 +486,7 @@ class LLMProvider:
         max_backoff: float,
         skip_validation: bool,
         start_time: float,
+        return_usage: bool = False,
     ) -> Any:
         """Handle Anthropic-specific API calls."""
         from anthropic import APIConnectionError, APIStatusError, RateLimitError
@@ -524,17 +559,35 @@ class LLMProvider:
                 else:
                     result = content
 
-                # Log slow calls
+                # Record token usage metrics
                 duration = time.time() - start_time
-                if duration > 10.0:
-                    input_tokens = response.usage.input_tokens
-                    output_tokens = response.usage.output_tokens
+                input_tokens = response.usage.input_tokens or 0 if response.usage else 0
+                output_tokens = response.usage.output_tokens or 0 if response.usage else 0
+                total_tokens = input_tokens + output_tokens
+
+                if response.usage:
+                    get_metrics_collector().record_tokens(
+                        operation="memory",
+                        bank_id="llm",
+                        input_tokens=input_tokens,
+                        output_tokens=output_tokens,
+                    )
+
+                # Log slow calls
+                if duration > 10.0 and response.usage:
                     logger.info(
                         f"slow llm call: model={self.provider}/{self.model}, "
                         f"input_tokens={input_tokens}, output_tokens={output_tokens}, "
                         f"time={duration:.3f}s"
                     )
 
+                if return_usage:
+                    token_usage = TokenUsage(
+                        input_tokens=input_tokens,
+                        output_tokens=output_tokens,
+                        total_tokens=total_tokens,
+                    )
+                    return result, token_usage
                 return result
 
             except json.JSONDecodeError as e:
@@ -589,6 +642,7 @@ class LLMProvider:
         max_backoff: float,
         skip_validation: bool,
         start_time: float,
+        return_usage: bool = False,
     ) -> Any:
         """
         Call Ollama using native API with JSON schema enforcement.
@@ -663,11 +717,35 @@ class LLMProvider:
                         else:
                             raise
 
+                    # Extract token usage from Ollama response
+                    # Ollama returns prompt_eval_count (input) and eval_count (output)
+                    input_tokens = result.get("prompt_eval_count", 0) or 0
+                    output_tokens = result.get("eval_count", 0) or 0
+                    total_tokens = input_tokens + output_tokens
+
+                    # Record to metrics
+                    if input_tokens > 0 or output_tokens > 0:
+                        get_metrics_collector().record_tokens(
+                            operation="memory",
+                            bank_id="llm",
+                            input_tokens=input_tokens,
+                            output_tokens=output_tokens,
+                        )
+
                     # Validate against Pydantic model or return raw JSON
                     if skip_validation:
-                        return json_data
+                        validated_result = json_data
                     else:
-                        return response_format.model_validate(json_data)
+                        validated_result = response_format.model_validate(json_data)
+
+                    if return_usage:
+                        token_usage = TokenUsage(
+                            input_tokens=input_tokens,
+                            output_tokens=output_tokens,
+                            total_tokens=total_tokens,
+                        )
+                        return validated_result, token_usage
+                    return validated_result
 
                 except httpx.HTTPStatusError as e:
                     last_exception = e
@@ -710,6 +788,7 @@ class LLMProvider:
         max_backoff: float,
         skip_validation: bool,
         start_time: float,
+        return_usage: bool = False,
     ) -> Any:
         """Handle Gemini-specific API calls."""
         # Convert OpenAI-style messages to Gemini format
@@ -786,16 +865,36 @@ class LLMProvider:
                 else:
                     result = content
 
-                # Log slow calls
+                # Record token usage metrics
                 duration = time.time() - start_time
-                if duration > 10.0 and hasattr(response, "usage_metadata") and response.usage_metadata:
+                input_tokens = 0
+                output_tokens = 0
+                if hasattr(response, "usage_metadata") and response.usage_metadata:
                     usage = response.usage_metadata
-                    logger.info(
-                        f"slow llm call: model={self.provider}/{self.model}, "
-                        f"input_tokens={usage.prompt_token_count}, output_tokens={usage.candidates_token_count}, "
-                        f"time={duration:.3f}s"
+                    input_tokens = usage.prompt_token_count or 0
+                    output_tokens = usage.candidates_token_count or 0
+                    get_metrics_collector().record_tokens(
+                        operation="memory",
+                        bank_id="llm",
+                        input_tokens=input_tokens,
+                        output_tokens=output_tokens,
                     )
 
+                    # Log slow calls
+                    if duration > 10.0:
+                        logger.info(
+                            f"slow llm call: model={self.provider}/{self.model}, "
+                            f"input_tokens={input_tokens}, output_tokens={output_tokens}, "
+                            f"time={duration:.3f}s"
+                        )
+
+                if return_usage:
+                    token_usage = TokenUsage(
+                        input_tokens=input_tokens,
+                        output_tokens=output_tokens,
+                        total_tokens=input_tokens + output_tokens,
+                    )
+                    return result, token_usage
                 return result
 
             except json.JSONDecodeError as e:

hindsight-api/hindsight_api/engine/memory_engine.py

@@ -132,11 +132,18 @@ if TYPE_CHECKING:
 
 from enum import Enum
 
-from ..pg0 import EmbeddedPostgres
+from ..pg0 import EmbeddedPostgres, parse_pg0_url
 from .entity_resolver import EntityResolver
 from .llm_wrapper import LLMConfig
 from .query_analyzer import QueryAnalyzer
-from .response_models import VALID_RECALL_FACT_TYPES, EntityObservation, EntityState, MemoryFact, ReflectResult
+from .response_models import (
+    VALID_RECALL_FACT_TYPES,
+    EntityObservation,
+    EntityState,
+    MemoryFact,
+    ReflectResult,
+    TokenUsage,
+)
 from .response_models import RecallResult as RecallResultModel
 from .retain import bank_utils, embedding_utils
 from .retain.types import RetainContentDict
@@ -259,31 +266,14 @@ class MemoryEngine(MemoryEngineInterface):
         memory_llm_base_url = memory_llm_base_url or config.get_llm_base_url() or None
         # Track pg0 instance (if used)
         self._pg0: EmbeddedPostgres | None = None
-        self._pg0_instance_name: str | None = None
 
         # Initialize PostgreSQL connection URL
         # The actual URL will be set during initialize() after starting the server
         # Supports: "pg0" (default instance), "pg0://instance-name" (named instance), or regular postgresql:// URL
-        if db_url == "pg0":
-            self._use_pg0 = True
-            self._pg0_instance_name = "hindsight"
-            self._pg0_port = None  # Use default port
-            self.db_url = None
-        elif db_url.startswith("pg0://"):
-            self._use_pg0 = True
-            # Parse instance name and optional port: pg0://instance-name or pg0://instance-name:port
-            url_part = db_url[6:]  # Remove "pg0://"
-            if ":" in url_part:
-                self._pg0_instance_name, port_str = url_part.rsplit(":", 1)
-                self._pg0_port = int(port_str)
-            else:
-                self._pg0_instance_name = url_part or "hindsight"
-                self._pg0_port = None  # Use default port
+        self._use_pg0, self._pg0_instance_name, self._pg0_port = parse_pg0_url(db_url)
+        if self._use_pg0:
             self.db_url = None
         else:
-            self._use_pg0 = False
-            self._pg0_instance_name = None
-            self._pg0_port = None
             self.db_url = db_url
 
         # Set default base URL if not provided
@@ -971,7 +961,8 @@ class MemoryEngine(MemoryEngineInterface):
         document_id: str | None = None,
         fact_type_override: str | None = None,
         confidence_score: float | None = None,
-    ) -> list[list[str]]:
+        return_usage: bool = False,
+    ):
         """
         Store multiple content items as memory units in ONE batch operation.
 
@@ -992,9 +983,11 @@ class MemoryEngine(MemoryEngineInterface):
                         Applies the same document_id to ALL content items that don't specify their own.
             fact_type_override: Override fact type for all facts ('world', 'experience', 'opinion')
             confidence_score: Confidence score for opinions (0.0 to 1.0)
+            return_usage: If True, returns tuple of (unit_ids, TokenUsage). Default False for backward compatibility.
 
         Returns:
-            List of lists of unit IDs (one list per content item)
+            If return_usage=False: List of lists of unit IDs (one list per content item)
+            If return_usage=True: Tuple of (unit_ids, TokenUsage)
 
         Example (new style - per-content document_id):
             unit_ids = await memory.retain_batch_async(
@@ -1021,6 +1014,8 @@ class MemoryEngine(MemoryEngineInterface):
         start_time = time.time()
 
         if not contents:
+            if return_usage:
+                return [], TokenUsage()
             return []
 
         # Authenticate tenant and set schema in context (for fq_table())
@@ -1050,6 +1045,7 @@ class MemoryEngine(MemoryEngineInterface):
         # Auto-chunk large batches by character count to avoid timeouts and memory issues
         # Calculate total character count
         total_chars = sum(len(item.get("content", "")) for item in contents)
+        total_usage = TokenUsage()
 
         CHARS_PER_BATCH = 600_000
 
@@ -1090,7 +1086,7 @@ class MemoryEngine(MemoryEngineInterface):
                     f"Processing sub-batch {i}/{len(sub_batches)}: {len(sub_batch)} items, {sub_batch_chars:,} chars"
                 )
 
-                sub_results = await self._retain_batch_async_internal(
+                sub_results, sub_usage = await self._retain_batch_async_internal(
                     bank_id=bank_id,
                     contents=sub_batch,
                     document_id=document_id,
@@ -1099,6 +1095,7 @@ class MemoryEngine(MemoryEngineInterface):
                     confidence_score=confidence_score,
                 )
                 all_results.extend(sub_results)
+                total_usage = total_usage + sub_usage
 
             total_time = time.time() - start_time
             logger.info(
@@ -1107,7 +1104,7 @@ class MemoryEngine(MemoryEngineInterface):
             result = all_results
         else:
             # Small batch - use internal method directly
-            result = await self._retain_batch_async_internal(
+            result, total_usage = await self._retain_batch_async_internal(
                 bank_id=bank_id,
                 contents=contents,
                 document_id=document_id,
@@ -1136,6 +1133,8 @@ class MemoryEngine(MemoryEngineInterface):
             except Exception as e:
                 logger.warning(f"Post-retain hook error (non-fatal): {e}")
 
+        if return_usage:
+            return result, total_usage
         return result
 
     async def _retain_batch_async_internal(
@@ -1146,7 +1145,7 @@ class MemoryEngine(MemoryEngineInterface):
         is_first_batch: bool = True,
         fact_type_override: str | None = None,
         confidence_score: float | None = None,
-    ) -> list[list[str]]:
+    ) -> tuple[list[list[str]], "TokenUsage"]:
         """
         Internal method for batch processing without chunking logic.
 
@@ -1162,6 +1161,9 @@ class MemoryEngine(MemoryEngineInterface):
             is_first_batch: Whether this is the first batch (for chunked operations, only delete on first batch)
             fact_type_override: Override fact type for all facts
             confidence_score: Confidence score for opinions
+
+        Returns:
+            Tuple of (unit ID lists, token usage for fact extraction)
         """
         # Backpressure: limit concurrent retains to prevent database contention
         async with self._put_semaphore:
@@ -2262,6 +2264,7 @@ class MemoryEngine(MemoryEngineInterface):
         bank_id: str | None = None,
         fact_type: str | None = None,
         *,
+        limit: int = 1000,
         request_context: "RequestContext",
     ):
         """
@@ -2270,10 +2273,11 @@ class MemoryEngine(MemoryEngineInterface):
         Args:
             bank_id: Filter by bank ID
             fact_type: Filter by fact type (world, experience, opinion)
+            limit: Maximum number of items to return (default: 1000)
             request_context: Request context for authentication.
 
         Returns:
-            Dict with nodes, edges, and table_rows
+            Dict with nodes, edges, table_rows, total_units, and limit
         """
         await self._authenticate_tenant(request_context)
         pool = await self._get_pool()
@@ -2295,15 +2299,29 @@ class MemoryEngine(MemoryEngineInterface):
 
             where_clause = "WHERE " + " AND ".join(query_conditions) if query_conditions else ""
 
+            # Get total count first
+            total_count_result = await conn.fetchrow(
+                f"""
+                SELECT COUNT(*) as total
+                FROM {fq_table("memory_units")}
+                {where_clause}
+            """,
+                *query_params,
+            )
+            total_count = total_count_result["total"] if total_count_result else 0
+
+            # Get units with limit
+            param_count += 1
             units = await conn.fetch(
                 f"""
                 SELECT id, text, event_date, context, occurred_start, occurred_end, mentioned_at, document_id, chunk_id, fact_type
                 FROM {fq_table("memory_units")}
                 {where_clause}
                 ORDER BY mentioned_at DESC NULLS LAST, event_date DESC
-                LIMIT 1000
+                LIMIT ${param_count}
             """,
                 *query_params,
+                limit,
             )
 
             # Get links, filtering to only include links between units of the selected agent
@@ -2440,7 +2458,7 @@ class MemoryEngine(MemoryEngineInterface):
                 }
             )
 
-        return {"nodes": nodes, "edges": edges, "table_rows": table_rows, "total_units": len(units)}
+        return {"nodes": nodes, "edges": edges, "table_rows": table_rows, "total_units": total_count, "limit": limit}
 
     async def list_memory_units(
         self,
@@ -3209,7 +3227,7 @@ Guidelines:
             response_format = JsonSchemaWrapper(response_schema)
 
         llm_start = time.time()
-        result = await self._llm_config.call(
+        llm_result, usage = await self._llm_config.call(
             messages=messages,
             scope="memory_reflect",
             max_completion_tokens=max_tokens,
@@ -3218,17 +3236,18 @@ Guidelines:
             # Don't enforce strict_schema - not all providers support it and may retry forever
             # Soft enforcement (schema in prompt + json_object mode) is sufficient
             strict_schema=False,
+            return_usage=True,
         )
         llm_time = time.time() - llm_start
 
         # Handle response based on whether structured output was requested
         if response_schema is not None:
-            structured_output = result
+            structured_output = llm_result
             answer_text = ""  # Empty for backward compatibility
             log_buffer.append(f"[REFLECT {reflect_id}] Structured output generated")
         else:
             structured_output = None
-            answer_text = result.strip()
+            answer_text = llm_result.strip()
 
         # Submit form_opinion task for background processing
         # Pass tenant_id from request context for internal authentication in background task
@@ -3254,6 +3273,7 @@ Guidelines:
             based_on={"world": world_results, "experience": agent_results, "opinion": opinion_results},
             new_opinions=[],  # Opinions are being extracted asynchronously
             structured_output=structured_output,
+            usage=usage,
         )
 
         # Call post-operation hook if validator is configured

hindsight-api/hindsight_api/engine/response_models.py

@@ -14,6 +14,37 @@ from pydantic import BaseModel, ConfigDict, Field
 VALID_RECALL_FACT_TYPES = frozenset(["world", "experience", "opinion"])
 
 
+class TokenUsage(BaseModel):
+    """
+    Token usage metrics for LLM calls.
+
+    Tracks input/output tokens for a single request to enable
+    per-request cost tracking and monitoring.
+    """
+
+    model_config = ConfigDict(
+        json_schema_extra={
+            "example": {
+                "input_tokens": 1500,
+                "output_tokens": 500,
+                "total_tokens": 2000,
+            }
+        }
+    )
+
+    input_tokens: int = Field(default=0, description="Number of input/prompt tokens consumed")
+    output_tokens: int = Field(default=0, description="Number of output/completion tokens generated")
+    total_tokens: int = Field(default=0, description="Total tokens (input + output)")
+
+    def __add__(self, other: "TokenUsage") -> "TokenUsage":
+        """Allow aggregating token usage from multiple calls."""
+        return TokenUsage(
+            input_tokens=self.input_tokens + other.input_tokens,
+            output_tokens=self.output_tokens + other.output_tokens,
+            total_tokens=self.total_tokens + other.total_tokens,
+        )
+
+
 class DispositionTraits(BaseModel):
     """
     Disposition traits for a memory bank.
@@ -147,6 +178,7 @@ class ReflectResult(BaseModel):
                 },
                 "new_opinions": ["Machine learning has great potential in healthcare"],
                 "structured_output": {"summary": "ML in healthcare", "confidence": 0.9},
+                "usage": {"input_tokens": 1500, "output_tokens": 500, "total_tokens": 2000},
             }
         }
     )
@@ -160,6 +192,10 @@ class ReflectResult(BaseModel):
         default=None,
         description="Structured output parsed according to the provided response schema. Only present when response_schema was provided.",
     )
+    usage: TokenUsage | None = Field(
+        default=None,
+        description="Token usage metrics for the LLM calls made during this reflect operation.",
+    )
 
 
 class Opinion(BaseModel):

hindsight-api/hindsight_api/engine/retain/fact_extraction.py

@@ -14,7 +14,9 @@ from typing import Literal
 
 from pydantic import BaseModel, ConfigDict, Field, field_validator
 
+from ...config import get_config
 from ..llm_wrapper import LLMConfig, OutputTooLongError
+from ..response_models import TokenUsage
 
 
 def _infer_temporal_date(fact_text: str, event_date: datetime) -> str | None:
@@ -109,7 +111,7 @@ class Fact(BaseModel):
 
 
 class CausalRelation(BaseModel):
-    """Causal relationship between facts."""
+    """Causal relationship between facts (legacy - embedded in each fact)."""
 
     target_fact_index: int = Field(
         description="Index of the related fact in the facts array (0-based). "
@@ -131,6 +133,36 @@ class CausalRelation(BaseModel):
     )
 
 
+class TopLevelCausalRelation(BaseModel):
+    """
+    Causal relationship between two facts (top-level schema).
+
+    This is the preferred format - defined AFTER all facts are extracted,
+    allowing the LLM to see the full list of facts before specifying relationships.
+    """
+
+    from_fact_index: int = Field(
+        description="Index of the source fact (0-based). The fact that causes/enables/prevents."
+    )
+    to_fact_index: int = Field(
+        description="Index of the target fact (0-based). The fact that is caused/enabled/prevented."
+    )
+    relation_type: Literal["causes", "caused_by", "enables", "prevents"] = Field(
+        description="Type of causal relationship: "
+        "'causes' = source fact directly causes the target fact, "
+        "'caused_by' = source fact was caused by the target fact, "
+        "'enables' = source fact enables/allows the target fact, "
+        "'prevents' = source fact prevents/blocks the target fact"
+    )
+    strength: float = Field(
+        description="Strength of causal relationship (0.0 to 1.0). "
+        "1.0 = direct/strong causation, 0.5 = moderate, 0.3 = weak/indirect",
+        ge=0.0,
+        le=1.0,
+        default=1.0,
+    )
+
+
 class ExtractedFact(BaseModel):
     """A single extracted fact with 5 required dimensions for comprehensive capture."""
 
@@ -253,9 +285,15 @@ class ExtractedFact(BaseModel):
 
 
 class FactExtractionResponse(BaseModel):
-    """Response containing all extracted facts."""
+    """Response containing all extracted facts and their causal relationships."""
 
     facts: list[ExtractedFact] = Field(description="List of extracted factual statements")
+    causal_relationships: list[TopLevelCausalRelation] | None = Field(
+        default=None,
+        description="Causal relationships between facts. Define these AFTER listing all facts. "
+        "Each relationship specifies from_fact_index -> to_fact_index with a relation type. "
+        "Indices must be valid (0 to N-1 where N is the number of facts).",
+    )
 
 
 def chunk_text(text: str, max_chars: int) -> list[str]:
@@ -356,7 +394,7 @@ async def _extract_facts_from_chunk(
     llm_config: "LLMConfig",
     agent_name: str = None,
     extract_opinions: bool = False,
-) -> list[dict[str, str]]:
+) -> tuple[list[dict[str, str]], TokenUsage]:
     """
     Extract facts from a single chunk (internal helper for parallel processing).
 
@@ -572,7 +610,53 @@ WHAT TO EXTRACT vs SKIP
 ══════════════════════════════════════════════════════════════════════════
 
 ✅ EXTRACT: User preferences (ALWAYS as separate facts!), feelings, plans, events, relationships, achievements
-❌ SKIP: Greetings, filler ("thanks", "cool"), purely structural statements"""
+❌ SKIP: Greetings, filler ("thanks", "cool"), purely structural statements
+
+══════════════════════════════════════════════════════════════════════════
+CAUSAL RELATIONSHIPS (CRITICAL - DEFINE AFTER ALL FACTS)
+══════════════════════════════════════════════════════════════════════════
+
+⚠️ IMPORTANT: Causal relationships are defined at the TOP LEVEL, AFTER listing all facts!
+
+The `causal_relationships` array goes at the root of your response (NOT inside each fact).
+This allows you to see all facts first before defining how they relate.
+
+Format:
+```json
+{{
+  "facts": [...all your extracted facts...],
+  "causal_relationships": [
+    {{"from_fact_index": 0, "to_fact_index": 1, "relation_type": "causes", "strength": 0.9}},
+    {{"from_fact_index": 1, "to_fact_index": 2, "relation_type": "enables", "strength": 0.7}}
+  ]
+}}
+```
+
+Relationship types:
+- "causes": Fact A directly causes Fact B (A → B)
+- "caused_by": Fact A was caused by Fact B (A ← B)
+- "enables": Fact A enables/allows Fact B to happen
+- "prevents": Fact A prevents/blocks Fact B from happening
+
+⚠️ INDEX VALIDATION: If you extract N facts (indices 0 to N-1), both from_fact_index and to_fact_index MUST be in range [0, N-1].
+
+Example (Event Date: March 15, 2024):
+Input: "I lost my job in January. Because of that, I couldn't pay rent. So I had to move to a cheaper apartment."
+
+Facts extracted:
+- Fact 0: "User lost their job in January due to layoffs"
+- Fact 1: "User couldn't pay rent because of job loss"
+- Fact 2: "User moved to a cheaper apartment"
+
+Causal relationships (at root level):
+```json
+"causal_relationships": [
+  {{"from_fact_index": 0, "to_fact_index": 1, "relation_type": "causes", "strength": 1.0}},
+  {{"from_fact_index": 1, "to_fact_index": 2, "relation_type": "causes", "strength": 0.9}}
+]
+```
+
+This creates a chain: Job loss (0) → Can't pay rent (1) → Moved to cheaper apartment (2)"""
 
     import logging
 
@@ -583,6 +667,7 @@ WHAT TO EXTRACT vs SKIP
     # Retry logic for JSON validation errors
     max_retries = 2
     last_error = None
+    config = get_config()
 
     # Sanitize input text to prevent Unicode encoding errors (e.g., unpaired surrogates)
     sanitized_chunk = _sanitize_text(chunk)
@@ -601,16 +686,19 @@ Context: {sanitized_context}
 Text:
 {sanitized_chunk}"""
 
+    usage = TokenUsage()  # Track cumulative usage across retries
     for attempt in range(max_retries):
         try:
-            extraction_response_json = await llm_config.call(
+            extraction_response_json, call_usage = await llm_config.call(
                 messages=[{"role": "system", "content": prompt}, {"role": "user", "content": user_message}],
                 response_format=FactExtractionResponse,
                 scope="memory_extract_facts",
                 temperature=0.1,
-                max_completion_tokens=65000,
+                max_completion_tokens=config.retain_max_completion_tokens,
                 skip_validation=True,  # Get raw JSON, we'll validate leniently
+                return_usage=True,
             )
+            usage = usage + call_usage  # Aggregate usage across retries
 
             # Lenient parsing of facts from raw JSON
             chunk_facts = []
@@ -628,9 +716,12 @@ Text:
                         f"LLM returned non-dict JSON after {max_retries} attempts: {type(extraction_response_json).__name__}. "
                         f"Raw: {str(extraction_response_json)[:500]}"
                     )
-                    return []
+                    return [], usage
 
             raw_facts = extraction_response_json.get("facts", [])
+            # Get top-level causal relationships (new schema)
+            top_level_causal_relations = extraction_response_json.get("causal_relationships", [])
+
             if not raw_facts:
                 logger.debug(
                     f"LLM response missing 'facts' field or returned empty list. "
@@ -641,6 +732,47 @@ Text:
                     f"text: {chunk}"
                 )
 
+            # Build a map from fact index to causal relations (from top-level field)
+            # This converts from_fact_index -> [{target_fact_index, relation_type, strength}]
+            causal_relations_by_fact: dict[int, list[dict]] = {}
+            if top_level_causal_relations:
+                num_facts = len(raw_facts)
+                for rel in top_level_causal_relations:
+                    if not isinstance(rel, dict):
+                        continue
+                    from_idx = rel.get("from_fact_index")
+                    to_idx = rel.get("to_fact_index")
+                    relation_type = rel.get("relation_type")
+                    strength = rel.get("strength", 1.0)
+
+                    # Validate indices
+                    if from_idx is None or to_idx is None or relation_type is None:
+                        logger.warning(f"Skipping malformed top-level causal relation: {rel}")
+                        continue
+                    if from_idx < 0 or from_idx >= num_facts:
+                        logger.warning(
+                            f"Invalid from_fact_index {from_idx} in top-level causal relation "
+                            f"(valid range: 0-{num_facts - 1}). Skipping."
+                        )
+                        continue
+                    if to_idx < 0 or to_idx >= num_facts:
+                        logger.warning(
+                            f"Invalid to_fact_index {to_idx} in top-level causal relation "
+                            f"(valid range: 0-{num_facts - 1}). Skipping."
+                        )
+                        continue
+
+                    # Add to the map for the from_fact_index
+                    if from_idx not in causal_relations_by_fact:
+                        causal_relations_by_fact[from_idx] = []
+                    causal_relations_by_fact[from_idx].append(
+                        {
+                            "target_fact_index": to_idx,
+                            "relation_type": relation_type,
+                            "strength": strength,
+                        }
+                    )
+
             for i, llm_fact in enumerate(raw_facts):
                 # Skip non-dict entries but track them for retry
                 if not isinstance(llm_fact, dict):
@@ -745,19 +877,40 @@ Text:
                     if validated_entities:
                         fact_data["entities"] = validated_entities
 
-                # Add causal relations if present (validate as CausalRelation objects)
-                # Filter out invalid relations (missing required fields)
-                causal_relations = get_value("causal_relations")
-                if causal_relations:
-                    validated_relations = []
-                    for rel in causal_relations:
+                # Add causal relations from both sources:
+                # 1. Top-level causal_relationships (preferred, new schema)
+                # 2. Per-fact causal_relations (legacy, for backward compatibility)
+                validated_relations = []
+
+                # First, add relations from top-level (already validated above)
+                if i in causal_relations_by_fact:
+                    for rel in causal_relations_by_fact[i]:
+                        try:
+                            validated_relations.append(CausalRelation.model_validate(rel))
+                        except Exception as e:
+                            logger.warning(f"Invalid top-level causal relation for fact {i}: {rel}: {e}")
+
+                # Then, add any legacy per-fact relations (with index validation)
+                legacy_causal_relations = get_value("causal_relations")
+                if legacy_causal_relations:
+                    num_facts = len(raw_facts)
+                    for rel in legacy_causal_relations:
                         if isinstance(rel, dict) and "target_fact_index" in rel and "relation_type" in rel:
-                            try:
-                                validated_relations.append(CausalRelation.model_validate(rel))
-                            except Exception as e:
-                                logger.warning(f"Invalid causal relation {rel}: {e}")
-                    if validated_relations:
-                        fact_data["causal_relations"] = validated_relations
+                            target_idx = rel.get("target_fact_index")
+                            # Validate target index for legacy format too
+                            if target_idx is not None and 0 <= target_idx < num_facts:
+                                try:
+                                    validated_relations.append(CausalRelation.model_validate(rel))
+                                except Exception as e:
+                                    logger.warning(f"Invalid causal relation {rel}: {e}")
+                            else:
+                                logger.warning(
+                                    f"Invalid target_fact_index {target_idx} in per-fact causal relation "
+                                    f"from fact {i} (valid range: 0-{num_facts - 1}). Skipping."
+                                )
+
+                if validated_relations:
+                    fact_data["causal_relations"] = validated_relations
 
                 # Always set mentioned_at to the event_date (when the conversation/document occurred)
                 fact_data["mentioned_at"] = event_date.isoformat()
@@ -778,7 +931,7 @@ Text:
                 )
                 continue
 
-            return chunk_facts
+            return chunk_facts, usage
 
         except BadRequestError as e:
             last_error = e
@@ -805,7 +958,7 @@ async def _extract_facts_with_auto_split(
     llm_config: LLMConfig,
     agent_name: str = None,
     extract_opinions: bool = False,
-) -> list[dict[str, str]]:
+) -> tuple[list[dict[str, str]], TokenUsage]:
     """
     Extract facts from a chunk with automatic splitting if output exceeds token limits.
 
@@ -823,7 +976,7 @@ async def _extract_facts_with_auto_split(
         extract_opinions: If True, extract ONLY opinions. If False, extract world and agent facts (no opinions)
 
     Returns:
-        List of fact dictionaries extracted from the chunk (possibly from sub-chunks)
+        Tuple of (facts list, token usage) extracted from the chunk (possibly from sub-chunks)
     """
     import logging
 
@@ -902,12 +1055,14 @@ async def _extract_facts_with_auto_split(
 
         # Combine results from both halves
         all_facts = []
-        for sub_result in sub_results:
-            all_facts.extend(sub_result)
+        total_usage = TokenUsage()
+        for sub_facts, sub_usage in sub_results:
+            all_facts.extend(sub_facts)
+            total_usage = total_usage + sub_usage
 
         logger.info(f"Successfully extracted {len(all_facts)} facts from split chunk {chunk_index + 1}")
 
-        return all_facts
+        return all_facts, total_usage
 
 
 async def extract_facts_from_text(
@@ -917,7 +1072,7 @@ async def extract_facts_from_text(
     agent_name: str,
     context: str = "",
     extract_opinions: bool = False,
-) -> tuple[list[Fact], list[tuple[str, int]]]:
+) -> tuple[list[Fact], list[tuple[str, int]], TokenUsage]:
     """
     Extract semantic facts from conversational or narrative text using LLM.
 
@@ -936,9 +1091,10 @@ async def extract_facts_from_text(
         extract_opinions: If True, extract ONLY opinions. If False, extract world and bank facts (no opinions)
 
     Returns:
-        Tuple of (facts, chunks) where:
+        Tuple of (facts, chunks, usage) where:
         - facts: List of Fact model instances
         - chunks: List of tuples (chunk_text, fact_count) for each chunk
+        - usage: Aggregated token usage across all LLM calls
     """
     chunks = chunk_text(text, max_chars=3000)
     tasks = [
@@ -957,10 +1113,12 @@ async def extract_facts_from_text(
     chunk_results = await asyncio.gather(*tasks)
     all_facts = []
     chunk_metadata = []  # [(chunk_text, fact_count), ...]
-    for chunk, chunk_facts in zip(chunks, chunk_results):
+    total_usage = TokenUsage()
+    for chunk, (chunk_facts, chunk_usage) in zip(chunks, chunk_results):
         all_facts.extend(chunk_facts)
         chunk_metadata.append((chunk, len(chunk_facts)))
-    return all_facts, chunk_metadata
+        total_usage = total_usage + chunk_usage
+    return all_facts, chunk_metadata, total_usage
 
 
 # ============================================================================
@@ -981,7 +1139,7 @@ SECONDS_PER_FACT = 10
 
 async def extract_facts_from_contents(
     contents: list[RetainContent], llm_config, agent_name: str, extract_opinions: bool = False
-) -> tuple[list[ExtractedFactType], list[ChunkMetadata]]:
+) -> tuple[list[ExtractedFactType], list[ChunkMetadata], TokenUsage]:
     """
     Extract facts from multiple content items in parallel.
 
@@ -998,10 +1156,10 @@ async def extract_facts_from_contents(
         extract_opinions: If True, extract only opinions; otherwise world/bank facts
 
     Returns:
-        Tuple of (extracted_facts, chunks_metadata)
+        Tuple of (extracted_facts, chunks_metadata, usage)
     """
     if not contents:
-        return [], []
+        return [], [], TokenUsage()
 
     # Step 1: Create parallel fact extraction tasks
     fact_extraction_tasks = []
@@ -1024,11 +1182,15 @@ async def extract_facts_from_contents(
     # Step 3: Flatten and convert to typed objects
     extracted_facts: list[ExtractedFactType] = []
     chunks_metadata: list[ChunkMetadata] = []
+    total_usage = TokenUsage()
 
     global_chunk_idx = 0
     global_fact_idx = 0
 
-    for content_index, (content, (facts_from_llm, chunks_from_llm)) in enumerate(zip(contents, all_fact_results)):
+    for content_index, (content, (facts_from_llm, chunks_from_llm, content_usage)) in enumerate(
+        zip(contents, all_fact_results)
+    ):
+        total_usage = total_usage + content_usage
         chunk_start_idx = global_chunk_idx
 
         # Convert chunk tuples to ChunkMetadata objects
@@ -1082,7 +1244,7 @@ async def extract_facts_from_contents(
     # Step 4: Add time offsets to preserve ordering within each content
     _add_temporal_offsets(extracted_facts, contents)
 
-    return extracted_facts, chunks_metadata
+    return extracted_facts, chunks_metadata, total_usage
 
 
 def _parse_datetime(date_str: str):

hindsight-api/hindsight_api/engine/retain/orchestrator.py

@@ -18,6 +18,7 @@ def utcnow():
     return datetime.now(UTC)
 
 
+from ..response_models import TokenUsage
 from . import (
     chunk_storage,
     deduplication,
@@ -47,7 +48,7 @@ async def retain_batch(
     is_first_batch: bool = True,
     fact_type_override: str | None = None,
     confidence_score: float | None = None,
-) -> list[list[str]]:
+) -> tuple[list[list[str]], TokenUsage]:
     """
     Process a batch of content through the retain pipeline.
 
@@ -67,7 +68,7 @@ async def retain_batch(
         confidence_score: Confidence score for opinions
 
     Returns:
-        List of unit ID lists (one list per content item)
+        Tuple of (unit ID lists, token usage for fact extraction)
     """
     start_time = time.time()
     total_chars = sum(len(item.get("content", "")) for item in contents_dicts)
@@ -99,7 +100,7 @@ async def retain_batch(
     step_start = time.time()
     extract_opinions = fact_type_override == "opinion"
 
-    extracted_facts, chunks = await fact_extraction.extract_facts_from_contents(
+    extracted_facts, chunks, usage = await fact_extraction.extract_facts_from_contents(
         contents, llm_config, agent_name, extract_opinions
     )
     log_buffer.append(
@@ -164,7 +165,7 @@ async def retain_batch(
         logger.info(
             f"RETAIN_BATCH COMPLETE: 0 facts extracted from {len(contents)} contents in {total_time:.3f}s (document tracked, no facts)"
         )
-        return [[] for _ in contents]
+        return [[] for _ in contents], usage
 
     # Apply fact_type_override if provided
     if fact_type_override:
@@ -344,7 +345,7 @@ async def retain_batch(
             non_duplicate_facts = deduplication.filter_duplicates(processed_facts, is_duplicate_flags)
 
             if not non_duplicate_facts:
-                return [[] for _ in contents]
+                return [[] for _ in contents], usage
 
             # Insert facts (document_id is now stored per-fact)
             step_start = time.time()
@@ -415,7 +416,7 @@ async def retain_batch(
 
         logger.info("\n" + "\n".join(log_buffer) + "\n")
 
-        return result_unit_ids
+        return result_unit_ids, usage
 
 
 def _map_results_to_contents(

hindsight-api/hindsight_api/engine/utils.py

@@ -49,7 +49,7 @@ async def extract_facts(
     if not text or not text.strip():
         return [], []
 
-    facts, chunks = await extract_facts_from_text(
+    facts, chunks, _ = await extract_facts_from_text(
         text,
         event_date,
         context=context,

hindsight-api/hindsight_api/main.py

@@ -184,8 +184,10 @@ def main():
             graph_retriever=config.graph_retriever,
             observation_min_facts=config.observation_min_facts,
             observation_top_entities=config.observation_top_entities,
+            retain_max_completion_tokens=config.retain_max_completion_tokens,
             skip_llm_verification=config.skip_llm_verification,
             lazy_reranker=config.lazy_reranker,
+            run_migrations_on_startup=config.run_migrations_on_startup,
         )
     config.configure_logging()
     if not args.daemon:
@@ -211,7 +213,11 @@ def main():
         logging.info(f"Loaded tenant extension: {tenant_extension.__class__.__name__}")
 
     # Create MemoryEngine (reads configuration from environment)
-    _memory = MemoryEngine(operation_validator=operation_validator, tenant_extension=tenant_extension)
+    _memory = MemoryEngine(
+        operation_validator=operation_validator,
+        tenant_extension=tenant_extension,
+        run_migrations=config.run_migrations_on_startup,
+    )
 
     # Set extension context on tenant extension (needed for schema provisioning)
     if tenant_extension:

hindsight-api/hindsight_api/pg0.py

@@ -132,3 +132,56 @@ async def stop_embedded_postgres() -> None:
     global _default_instance
     if _default_instance:
         await _default_instance.stop()
+
+
+def parse_pg0_url(db_url: str) -> tuple[bool, str | None, int | None]:
+    """
+    Parse a database URL and check if it's a pg0:// embedded database URL.
+
+    Supports:
+    - "pg0" -> default instance "hindsight"
+    - "pg0://instance-name" -> named instance
+    - "pg0://instance-name:port" -> named instance with explicit port
+    - Any other URL (e.g., postgresql://) -> not a pg0 URL
+
+    Args:
+        db_url: The database URL to parse
+
+    Returns:
+        Tuple of (is_pg0, instance_name, port)
+        - is_pg0: True if this is a pg0 URL
+        - instance_name: The instance name (or None if not pg0)
+        - port: The explicit port (or None for auto-assign)
+    """
+    if db_url == "pg0":
+        return True, "hindsight", None
+
+    if db_url.startswith("pg0://"):
+        url_part = db_url[6:]  # Remove "pg0://"
+        if ":" in url_part:
+            instance_name, port_str = url_part.rsplit(":", 1)
+            return True, instance_name or "hindsight", int(port_str)
+        else:
+            return True, url_part or "hindsight", None
+
+    return False, None, None
+
+
+async def resolve_database_url(db_url: str) -> str:
+    """
+    Resolve a database URL, handling pg0:// embedded database URLs.
+
+    If the URL is a pg0:// URL, starts the embedded PostgreSQL and returns
+    the actual postgresql:// connection URL. Otherwise, returns the URL unchanged.
+
+    Args:
+        db_url: Database URL (pg0://, pg0, or postgresql://)
+
+    Returns:
+        The resolved postgresql:// connection URL
+    """
+    is_pg0, instance_name, port = parse_pg0_url(db_url)
+    if is_pg0:
+        pg0 = EmbeddedPostgres(name=instance_name, port=port)
+        return await pg0.ensure_running()
+    return db_url

hindsight-api/pyproject.toml

@@ -38,6 +38,7 @@ dependencies = [
     "dateparser>=1.2.2",
     "google-genai>=1.0.0",
     "anthropic>=0.40.0",
+    "typer>=0.9.0",
 ]
 
 [project.optional-dependencies]
@@ -52,6 +53,7 @@ test = [
 [project.scripts]
 hindsight-api = "hindsight_api.main:main"
 hindsight-local-mcp = "hindsight_api.mcp_local:main"
+hindsight-admin = "hindsight_api.admin.cli:main"
 
 [tool.hatch.build.targets.wheel]
 packages = ["hindsight_api"]
@@ -75,7 +77,7 @@ log_cli = true
 log_cli_level = "INFO"
 log_cli_format = "%(asctime)s - %(levelname)s - %(name)s - %(message)s"
 log_cli_date_format = "%Y-%m-%d %H:%M:%S"
-addopts = "--timeout 120 -n 8 --durations=10 -v"
+addopts = "--timeout 120 -n 8 --dist loadgroup --durations=10 -v"
 asyncio_mode = "auto"
 asyncio_default_fixture_loop_scope = "function"
 log_auto_indent = true

hindsight-api/tests/test_admin_backup_restore.py

@@ -0,0 +1,292 @@
+"""
+Tests for admin backup and restore functionality.
+
+These tests use an isolated schema to avoid interfering with other tests.
+The backup/restore operations truncate tables, which would cause deadlocks
+and race conditions if run against the shared public schema.
+"""
+
+import tempfile
+import uuid
+import zipfile
+from pathlib import Path
+
+import asyncpg
+import pytest
+import pytest_asyncio
+
+from hindsight_api.admin.cli import _backup, _restore, BACKUP_TABLES
+from hindsight_api.migrations import run_migrations
+
+
+# Run these tests sequentially since they do full DB backup/restore
+pytestmark = pytest.mark.xdist_group(name="backup_restore")
+
+
+@pytest_asyncio.fixture(scope="function")
+async def backup_test_schema(pg0_db_url, embeddings):
+    """Create an isolated schema for backup/restore tests.
+
+    Uses a unique schema name per test invocation to avoid conflicts with
+    parallel test runs or leftover state from interrupted runs.
+
+    Returns a tuple of (db_url, schema_name, fq_helper, embeddings).
+    """
+    # Initialize embeddings if not already done
+    await embeddings.initialize()
+
+    # Use unique schema name to avoid conflicts
+    schema_name = f"backup_test_{uuid.uuid4().hex[:8]}"
+
+    def _fq(table: str) -> str:
+        """Get fully-qualified table name in test schema."""
+        return f"{schema_name}.{table}"
+
+    conn = await asyncpg.connect(pg0_db_url)
+    try:
+        await conn.execute(f"CREATE SCHEMA {schema_name}")
+    finally:
+        await conn.close()
+
+    # Run migrations on the isolated schema
+    run_migrations(pg0_db_url, schema=schema_name)
+
+    yield pg0_db_url, schema_name, _fq, embeddings
+
+    # Cleanup after test
+    conn = await asyncpg.connect(pg0_db_url)
+    try:
+        await conn.execute(f"DROP SCHEMA IF EXISTS {schema_name} CASCADE")
+    finally:
+        await conn.close()
+
+
+@pytest.mark.asyncio
+async def test_backup_restore_roundtrip(backup_test_schema):
+    """Test that backup and restore preserves all data correctly."""
+    db_url, schema_name, _fq, embeddings = backup_test_schema
+    bank_id = f"test-backup-{uuid.uuid4().hex[:8]}"
+    conn = await asyncpg.connect(db_url)
+
+    try:
+        # Create a bank
+        await conn.execute(
+            f"INSERT INTO {_fq('banks')} (bank_id) VALUES ($1) ON CONFLICT DO NOTHING",
+            bank_id,
+        )
+
+        # Create some test memory units with embeddings
+        # Convert embedding list to pgvector format string
+        embedding_list = embeddings.encode(["Test content about Alice"])[0]
+        embedding_str = "[" + ",".join(str(x) for x in embedding_list) + "]"
+        for text in [
+            "Alice is a software engineer who loves Python.",
+            "Bob works with Alice on the backend team.",
+            "The team uses PostgreSQL for their database.",
+        ]:
+            await conn.execute(
+                f"""INSERT INTO {_fq('memory_units')}
+                    (bank_id, text, fact_type, embedding, event_date)
+                    VALUES ($1, $2, 'world', $3::vector, NOW())""",
+                bank_id,
+                text,
+                embedding_str,
+            )
+
+        # Get counts before backup
+        counts_before = {}
+        for table in BACKUP_TABLES:
+            counts_before[table] = await conn.fetchval(f"SELECT COUNT(*) FROM {_fq(table)}")
+
+        # Verify we have data
+        assert counts_before["banks"] > 0
+        assert counts_before["memory_units"] > 0
+
+    finally:
+        await conn.close()
+
+    # Backup to a temp file
+    with tempfile.NamedTemporaryFile(suffix=".zip", delete=False) as f:
+        backup_path = Path(f.name)
+
+    try:
+        manifest = await _backup(db_url, backup_path, schema=schema_name)
+
+        # Verify backup file exists and is valid
+        assert backup_path.exists()
+        assert backup_path.stat().st_size > 0
+
+        # Verify manifest
+        assert manifest["version"] == "1"
+        assert "created_at" in manifest
+        for table in BACKUP_TABLES:
+            assert table in manifest["tables"]
+            assert manifest["tables"][table]["rows"] == counts_before[table]
+
+        # Verify zip contents
+        with zipfile.ZipFile(backup_path, "r") as zf:
+            assert "manifest.json" in zf.namelist()
+            for table in BACKUP_TABLES:
+                assert f"{table}.bin" in zf.namelist()
+
+        # Clear all data
+        conn = await asyncpg.connect(db_url)
+        try:
+            for table in reversed(BACKUP_TABLES):
+                await conn.execute(f"TRUNCATE TABLE {_fq(table)} CASCADE")
+
+            # Verify data is gone
+            for table in BACKUP_TABLES:
+                count = await conn.fetchval(f"SELECT COUNT(*) FROM {_fq(table)}")
+                assert count == 0, f"Table {table} should be empty after truncate"
+        finally:
+            await conn.close()
+
+        # Restore from backup
+        await _restore(db_url, backup_path, schema=schema_name)
+
+        # Verify counts match original
+        conn = await asyncpg.connect(db_url)
+        try:
+            for table in BACKUP_TABLES:
+                count = await conn.fetchval(f"SELECT COUNT(*) FROM {_fq(table)}")
+                assert count == counts_before[table], f"Table {table} count mismatch after restore"
+
+            # Verify data content is preserved
+            texts = await conn.fetch(
+                f"SELECT text FROM {_fq('memory_units')} WHERE bank_id = $1",
+                bank_id,
+            )
+            text_content = " ".join(r["text"] for r in texts)
+            assert "Alice" in text_content or "software" in text_content
+        finally:
+            await conn.close()
+
+    finally:
+        # Cleanup
+        if backup_path.exists():
+            backup_path.unlink()
+
+
+@pytest.mark.asyncio
+async def test_backup_restore_preserves_all_column_types(backup_test_schema):
+    """Test that all column types are preserved: vectors, UUIDs, timestamps, JSONB."""
+    db_url, schema_name, _fq, embeddings = backup_test_schema
+    bank_id = f"test-types-{uuid.uuid4().hex[:8]}"
+    conn = await asyncpg.connect(db_url)
+
+    try:
+        # Create a bank
+        await conn.execute(
+            f"INSERT INTO {_fq('banks')} (bank_id) VALUES ($1) ON CONFLICT DO NOTHING",
+            bank_id,
+        )
+
+        # Create a memory unit with all column types
+        # Convert embedding list to pgvector format string
+        embedding_list = embeddings.encode(["John Smith engineer"])[0]
+        embedding_str = "[" + ",".join(str(x) for x in embedding_list) + "]"
+        await conn.execute(
+            f"""INSERT INTO {_fq('memory_units')}
+                (bank_id, text, fact_type, embedding, event_date, metadata)
+                VALUES ($1, $2, 'world', $3::vector, NOW(), $4)""",
+            bank_id,
+            "John Smith is a senior engineer at Acme Corp since 2020.",
+            embedding_str,
+            '{"key": "value"}',
+        )
+
+        # Create an entity
+        await conn.execute(
+            f"""INSERT INTO {_fq('entities')}
+                (bank_id, canonical_name, metadata)
+                VALUES ($1, $2, $3)""",
+            bank_id,
+            "John Smith",
+            '{"role": "engineer"}',
+        )
+
+        # Get original data
+        original_unit = await conn.fetchrow(
+            f"""SELECT id, embedding, event_date, created_at, metadata, text
+               FROM {_fq('memory_units')} WHERE bank_id = $1 LIMIT 1""",
+            bank_id,
+        )
+        original_entity = await conn.fetchrow(
+            f"""SELECT id, first_seen, last_seen, metadata, canonical_name
+               FROM {_fq('entities')} WHERE bank_id = $1 LIMIT 1""",
+            bank_id,
+        )
+        original_bank = await conn.fetchrow(
+            f"SELECT bank_id, created_at, updated_at FROM {_fq('banks')} WHERE bank_id = $1",
+            bank_id,
+        )
+    finally:
+        await conn.close()
+
+    assert original_unit is not None, "Should have created memory units"
+    assert original_unit["embedding"] is not None, "Should have embedding"
+    assert original_unit["id"] is not None, "Should have UUID"
+    assert original_entity is not None, "Should have created entities"
+
+    with tempfile.NamedTemporaryFile(suffix=".zip", delete=False) as f:
+        backup_path = Path(f.name)
+
+    try:
+        await _backup(db_url, backup_path, schema=schema_name)
+
+        # Clear all data
+        conn = await asyncpg.connect(db_url)
+        try:
+            for table in reversed(BACKUP_TABLES):
+                await conn.execute(f"TRUNCATE TABLE {_fq(table)} CASCADE")
+        finally:
+            await conn.close()
+
+        await _restore(db_url, backup_path, schema=schema_name)
+
+        # Verify all column types are preserved exactly
+        conn = await asyncpg.connect(db_url)
+        try:
+            restored_unit = await conn.fetchrow(
+                f"""SELECT id, embedding, event_date, created_at, metadata, text
+                   FROM {_fq('memory_units')} WHERE bank_id = $1 LIMIT 1""",
+                bank_id,
+            )
+            restored_entity = await conn.fetchrow(
+                f"""SELECT id, first_seen, last_seen, metadata, canonical_name
+                   FROM {_fq('entities')} WHERE bank_id = $1 LIMIT 1""",
+                bank_id,
+            )
+            restored_bank = await conn.fetchrow(
+                f"SELECT bank_id, created_at, updated_at FROM {_fq('banks')} WHERE bank_id = $1",
+                bank_id,
+            )
+        finally:
+            await conn.close()
+
+        # Verify memory_units
+        assert restored_unit is not None, "Should have restored memory unit"
+        assert restored_unit["id"] == original_unit["id"], "UUID should match exactly"
+        assert restored_unit["text"] == original_unit["text"], "Text should match"
+        assert list(restored_unit["embedding"]) == list(original_unit["embedding"]), "Vector embedding should match exactly"
+        assert restored_unit["event_date"] == original_unit["event_date"], "Timestamp should match exactly"
+        assert restored_unit["created_at"] == original_unit["created_at"], "Created timestamp should match"
+        assert restored_unit["metadata"] == original_unit["metadata"], "JSONB metadata should match"
+
+        # Verify entities
+        assert restored_entity is not None, "Should have restored entity"
+        assert restored_entity["id"] == original_entity["id"], "Entity UUID should match"
+        assert restored_entity["canonical_name"] == original_entity["canonical_name"], "Entity name should match"
+        assert restored_entity["first_seen"] == original_entity["first_seen"], "Entity first_seen should match"
+        assert restored_entity["last_seen"] == original_entity["last_seen"], "Entity last_seen should match"
+        assert restored_entity["metadata"] == original_entity["metadata"], "Entity metadata should match"
+
+        # Verify banks
+        assert restored_bank is not None, "Should have restored bank"
+        assert restored_bank["bank_id"] == original_bank["bank_id"], "Bank ID should match"
+        assert restored_bank["created_at"] == original_bank["created_at"], "Bank created_at should match"
+
+    finally:
+        if backup_path.exists():
+            backup_path.unlink()

hindsight-api/tests/test_causal_relationships.py

@@ -0,0 +1,222 @@
+"""
+Test suite for causal relationship extraction.
+
+Tests that the fact extraction system correctly identifies and validates
+causal relationships between facts, with valid indices.
+"""
+
+from datetime import datetime
+
+import pytest
+
+from hindsight_api import LLMConfig
+from hindsight_api.engine.retain.fact_extraction import extract_facts_from_text
+
+
+class TestCausalRelationships:
+    """Tests for causal relationship extraction and validation."""
+
+    @pytest.mark.asyncio
+    async def test_causal_chain_extraction(self):
+        """
+        Test that a clear causal chain is extracted with valid relationships.
+
+        Story: Lost job -> couldn't pay rent -> had to move -> found new apartment
+
+        This is a 4-fact causal chain where each fact causes the next.
+        The extracted causal relations should have valid indices (0-3).
+        """
+        text = """
+I lost my job at the tech company in January because of layoffs.
+Because I lost my job, I couldn't pay my rent anymore.
+Since I couldn't afford rent, I had to move out of my apartment.
+After searching for weeks, I finally found a cheaper apartment in Brooklyn.
+"""
+
+        context = "Personal story about housing change"
+        llm_config = LLMConfig.for_memory()
+
+        facts, _, _ = await extract_facts_from_text(
+            text=text,
+            event_date=datetime(2024, 3, 15),
+            context=context,
+            llm_config=llm_config,
+            agent_name="TestUser"
+        )
+
+        assert len(facts) >= 3, f"Should extract at least 3 facts from the causal chain. Got {len(facts)}"
+
+        # Collect all causal relations from all facts
+        all_causal_relations = []
+        for i, fact in enumerate(facts):
+            if fact.causal_relations:
+                for rel in fact.causal_relations:
+                    all_causal_relations.append({
+                        "from_fact_index": i,
+                        "to_fact_index": rel.target_fact_index,
+                        "relation_type": rel.relation_type,
+                        "strength": rel.strength,
+                        "from_fact_text": fact.fact[:50],
+                    })
+
+        # Verify that ALL causal relation indices are valid
+        num_facts = len(facts)
+        invalid_relations = []
+        for rel in all_causal_relations:
+            if rel["to_fact_index"] < 0 or rel["to_fact_index"] >= num_facts:
+                invalid_relations.append(rel)
+
+        assert len(invalid_relations) == 0, (
+            f"Found {len(invalid_relations)} causal relations with invalid indices! "
+            f"Valid range is 0-{num_facts - 1}. "
+            f"Invalid relations: {invalid_relations}"
+        )
+
+        # Should have at least some causal relations extracted
+        assert len(all_causal_relations) >= 2, (
+            f"Should extract at least 2 causal relationships from this clear chain. "
+            f"Got {len(all_causal_relations)}: {all_causal_relations}"
+        )
+
+        # Verify relation types are valid
+        valid_types = {"causes", "caused_by", "enables", "prevents"}
+        for rel in all_causal_relations:
+            assert rel["relation_type"] in valid_types, (
+                f"Invalid relation_type '{rel['relation_type']}'. Must be one of {valid_types}"
+            )
+
+    @pytest.mark.asyncio
+    async def test_complex_causal_web(self):
+        """
+        Test a more complex scenario with multiple interconnected causes.
+
+        This tests the LLM's ability to identify multiple causal links and
+        ensure all referenced indices exist.
+        """
+        text = """
+The heavy rain caused flooding in the basement.
+The flooding damaged the electrical system.
+Because of the electrical damage, we had to call an electrician.
+The electrician found that the wiring was old and needed replacement.
+We decided to renovate the entire basement while fixing the wiring.
+The renovation took three months and cost $15,000.
+"""
+
+        context = "Home repair story"
+        llm_config = LLMConfig.for_memory()
+
+        facts, _, _ = await extract_facts_from_text(
+            text=text,
+            event_date=datetime(2024, 6, 1),
+            context=context,
+            llm_config=llm_config,
+            agent_name="TestUser"
+        )
+
+        assert len(facts) >= 4, f"Should extract at least 4 facts. Got {len(facts)}"
+
+        # Validate all causal relation indices
+        num_facts = len(facts)
+        for i, fact in enumerate(facts):
+            if fact.causal_relations:
+                for rel in fact.causal_relations:
+                    assert 0 <= rel.target_fact_index < num_facts, (
+                        f"Fact {i} has causal relation to invalid index {rel.target_fact_index}. "
+                        f"Valid range is 0-{num_facts - 1}. "
+                        f"Fact text: {fact.fact[:80]}..."
+                    )
+
+    @pytest.mark.asyncio
+    async def test_no_self_referencing_causal_relations(self):
+        """
+        Test that facts don't have causal relations pointing to themselves.
+        """
+        text = """
+I started learning Python because I wanted to automate my work tasks.
+Learning Python led me to discover machine learning.
+Machine learning fascinated me so much that I changed my career to data science.
+"""
+
+        context = "Career change story"
+        llm_config = LLMConfig.for_memory()
+
+        facts, _, _ = await extract_facts_from_text(
+            text=text,
+            event_date=datetime(2024, 1, 1),
+            context=context,
+            llm_config=llm_config,
+            agent_name="TestUser"
+        )
+
+        # Check no fact references itself
+        for i, fact in enumerate(facts):
+            if fact.causal_relations:
+                for rel in fact.causal_relations:
+                    assert rel.target_fact_index != i, (
+                        f"Fact {i} has a self-referencing causal relation! "
+                        f"Fact text: {fact.fact}"
+                    )
+
+    @pytest.mark.asyncio
+    async def test_bidirectional_causal_relationships(self):
+        """
+        Test that bidirectional causal relationships (causes and caused_by)
+        are handled correctly.
+        """
+        text = """
+My promotion at work caused me to move to New York.
+Moving to New York was caused by my promotion at work.
+The new role enabled me to lead a team of engineers.
+"""
+
+        context = "Work promotion story"
+        llm_config = LLMConfig.for_memory()
+
+        facts, _, _ = await extract_facts_from_text(
+            text=text,
+            event_date=datetime(2024, 2, 15),
+            context=context,
+            llm_config=llm_config,
+            agent_name="TestUser"
+        )
+
+        num_facts = len(facts)
+
+        # Validate all indices
+        for i, fact in enumerate(facts):
+            if fact.causal_relations:
+                for rel in fact.causal_relations:
+                    assert 0 <= rel.target_fact_index < num_facts, (
+                        f"Invalid target_fact_index {rel.target_fact_index} in fact {i}. "
+                        f"Valid range: 0-{num_facts - 1}"
+                    )
+
+    @pytest.mark.asyncio
+    async def test_causal_relation_strength_values(self):
+        """
+        Test that causal relation strength values are within valid range [0.0, 1.0].
+        """
+        text = """
+The stock market crash directly caused the company to lay off employees.
+The layoffs indirectly led to reduced consumer spending in the area.
+Reduced spending somewhat affected local businesses.
+"""
+
+        context = "Economic impact story"
+        llm_config = LLMConfig.for_memory()
+
+        facts, _, _ = await extract_facts_from_text(
+            text=text,
+            event_date=datetime(2024, 4, 1),
+            context=context,
+            llm_config=llm_config,
+            agent_name="TestUser"
+        )
+
+        for i, fact in enumerate(facts):
+            if fact.causal_relations:
+                for rel in fact.causal_relations:
+                    assert 0.0 <= rel.strength <= 1.0, (
+                        f"Causal relation strength {rel.strength} is outside valid range [0.0, 1.0]. "
+                        f"Fact {i}: {fact.fact[:50]}..."
+                    )

hindsight-api/tests/test_fact_extraction_quality.py

@@ -43,7 +43,7 @@ Marcus felt anxious about the upcoming interview.
         context = "Personal journal entry"
         llm_config = LLMConfig.for_memory()
 
-        facts, _ = await extract_facts_from_text(
+        facts, _, _ = await extract_facts_from_text(
             text=text,
             event_date=datetime(2024, 11, 13),
             context=context,
@@ -75,7 +75,7 @@ The music was so loud I could barely hear myself think.
         context = "Personal experience"
         llm_config = LLMConfig.for_memory()
 
-        facts, _ = await extract_facts_from_text(
+        facts, _, _ = await extract_facts_from_text(
             text=text,
             event_date=datetime(2024, 11, 13),
             context=context,
@@ -108,7 +108,7 @@ Maybe we should reconsider the timeline.
         context = "Team discussion"
         llm_config = LLMConfig.for_memory()
 
-        facts, _ = await extract_facts_from_text(
+        facts, _, _ = await extract_facts_from_text(
             text=text,
             event_date=datetime(2024, 11, 13),
             context=context,
@@ -141,7 +141,7 @@ I'm unable to attend the conference due to scheduling conflicts.
         context = "Personal profile discussion"
         llm_config = LLMConfig.for_memory()
 
-        facts, _ = await extract_facts_from_text(
+        facts, _, _ = await extract_facts_from_text(
             text=text,
             event_date=datetime(2024, 11, 13),
             context=context,
@@ -173,7 +173,7 @@ Unlike last year, we're ahead of schedule.
         context = "Project review"
         llm_config = LLMConfig.for_memory()
 
-        facts, _ = await extract_facts_from_text(
+        facts, _, _ = await extract_facts_from_text(
             text=text,
             event_date=datetime(2024, 11, 13),
             context=context,
@@ -206,7 +206,7 @@ She's enthusiastic about the opportunity.
         context = "Team meeting"
         llm_config = LLMConfig.for_memory()
 
-        facts, _ = await extract_facts_from_text(
+        facts, _, _ = await extract_facts_from_text(
             text=text,
             event_date=datetime(2024, 11, 13),
             context=context,
@@ -239,7 +239,7 @@ I'm planning to switch careers because I'm not fulfilled in my current role.
         context = "Personal goals discussion"
         llm_config = LLMConfig.for_memory()
 
-        facts, _ = await extract_facts_from_text(
+        facts, _, _ = await extract_facts_from_text(
             text=text,
             event_date=datetime(2024, 11, 13),
             context=context,
@@ -276,7 +276,7 @@ Family is the most important thing to her.
         context = "Personal values discussion"
         llm_config = LLMConfig.for_memory()
 
-        facts, _ = await extract_facts_from_text(
+        facts, _, _ = await extract_facts_from_text(
             text=text,
             event_date=datetime(2024, 11, 13),
             context=context,
@@ -310,7 +310,7 @@ I prefer presenting in person rather than virtually because I can read the room
 
         event_date = datetime(2024, 11, 13)
 
-        facts, _ = await extract_facts_from_text(
+        facts, _, _ = await extract_facts_from_text(
             text=text,
             event_date=event_date,
             context=context,
@@ -366,7 +366,7 @@ I'm planning to visit Tokyo next month.
 
         event_date = datetime(2024, 11, 13)
 
-        facts, _ = await extract_facts_from_text(
+        facts, _, _ = await extract_facts_from_text(
             text=text,
             event_date=event_date,
             context=context,
@@ -419,7 +419,7 @@ with a concert surrounded by music, joy and the warm summer breeze.
 
         for attempt in range(max_retries):
             try:
-                facts, _ = await extract_facts_from_text(
+                facts, _, _ = await extract_facts_from_text(
                     text=text,
                     event_date=event_date,
                     context=context,
@@ -488,7 +488,7 @@ Yesterday I went for a morning jog for the first time in a nearby park.
 
         event_date = datetime(2024, 11, 13)
 
-        facts, _ = await extract_facts_from_text(
+        facts, _, _ = await extract_facts_from_text(
             text=text,
             event_date=event_date,
             context=context,
@@ -537,7 +537,7 @@ Yesterday I went for a morning jog for the first time in a nearby park.
         This morning I had coffee with Alice.
         """
 
-        facts, _ = await extract_facts_from_text(
+        facts, _, _ = await extract_facts_from_text(
             text=text,
             event_date=reference_date,
             llm_config=llm_config,
@@ -567,7 +567,7 @@ Yesterday I went for a morning jog for the first time in a nearby park.
 
         text = "Alice works at Google. She loves Python programming."
 
-        facts, _ = await extract_facts_from_text(
+        facts, _, _ = await extract_facts_from_text(
             text=text,
             event_date=reference_date,
             llm_config=llm_config,
@@ -594,7 +594,7 @@ Yesterday I went for a morning jog for the first time in a nearby park.
         Bob will start his vacation on April 1st.
         """
 
-        facts, _ = await extract_facts_from_text(
+        facts, _, _ = await extract_facts_from_text(
             text=text,
             event_date=reference_date,
             llm_config=llm_config,
@@ -645,7 +645,7 @@ great time! Every time I see it, I can't help but smile.
 
         event_date = datetime(2023, 2, 23)
 
-        facts, _ = await extract_facts_from_text(
+        facts, _, _ = await extract_facts_from_text(
             text=text,
             event_date=event_date,
             context=context,
@@ -695,7 +695,7 @@ I've learned so much from it.
         context = "Personal update"
         llm_config = LLMConfig.for_memory()
 
-        facts, _ = await extract_facts_from_text(
+        facts, _, _ = await extract_facts_from_text(
             text=text,
             event_date=datetime(2024, 11, 13),
             context=context,
@@ -758,7 +758,7 @@ Jamie: Congratulations! I'd love to read it.
 
         llm_config = LLMConfig.for_memory()
 
-        facts, _ = await extract_facts_from_text(
+        facts, _, _ = await extract_facts_from_text(
             text=transcript,
             event_date=datetime(2024, 11, 13),
             llm_config=llm_config,
@@ -803,7 +803,7 @@ We presented our findings to the team yesterday.
 
         llm_config = LLMConfig.for_memory()
 
-        facts, _ = await extract_facts_from_text(
+        facts, _, _ = await extract_facts_from_text(
             text=text,
             event_date=datetime(2024, 11, 13),
             llm_config=llm_config,
@@ -838,7 +838,7 @@ Jamie: [teasing] We'll see who's right, my Niners pick is solid.
 
         llm_config = LLMConfig.for_memory()
 
-        facts, _ = await extract_facts_from_text(
+        facts, _, _ = await extract_facts_from_text(
             text=transcript,
             event_date=datetime(2024, 11, 14),
             context=context,
@@ -897,7 +897,7 @@ so the algorithm learns to box out. See you next week!
 
         llm_config = LLMConfig.for_memory()
 
-        facts, _ = await extract_facts_from_text(
+        facts, _, _ = await extract_facts_from_text(
             text=transcript,
             event_date=datetime(2024, 11, 13),
             llm_config=llm_config,

hindsight-api/tests/test_http_api_integration.py

@@ -832,3 +832,134 @@ async def test_reflect_with_max_tokens(api_client):
     # Verify response has text
     assert "text" in result
     assert len(result["text"]) > 0
+
+
+@pytest.mark.asyncio
+async def test_reflect_returns_token_usage(api_client):
+    """Test that reflect endpoint returns token usage metrics.
+
+    The usage field should contain input_tokens, output_tokens, and total_tokens
+    from the LLM call made during reflection.
+    """
+    test_bank_id = f"reflect_usage_test_{datetime.now().timestamp()}"
+
+    # Store a memory to reflect on
+    response = await api_client.post(
+        f"/v1/default/banks/{test_bank_id}/memories",
+        json={
+            "items": [
+                {
+                    "content": "The capital of France is Paris.",
+                    "context": "geography"
+                }
+            ]
+        }
+    )
+    assert response.status_code == 200
+
+    # Call reflect
+    response = await api_client.post(
+        f"/v1/default/banks/{test_bank_id}/reflect",
+        json={
+            "query": "What is the capital of France?"
+        }
+    )
+    assert response.status_code == 200
+    result = response.json()
+
+    # Verify response has text
+    assert "text" in result
+    assert len(result["text"]) > 0
+
+    # Verify usage field exists and has expected structure
+    assert "usage" in result, "Response should include 'usage' field"
+    usage = result["usage"]
+    assert usage is not None, "Usage should not be None for reflect"
+    assert "input_tokens" in usage, "Usage should have 'input_tokens'"
+    assert "output_tokens" in usage, "Usage should have 'output_tokens'"
+    assert "total_tokens" in usage, "Usage should have 'total_tokens'"
+
+    # Verify token counts are valid
+    assert usage["input_tokens"] > 0, f"Expected input_tokens > 0, got {usage['input_tokens']}"
+    assert usage["output_tokens"] >= 0, f"Expected output_tokens >= 0, got {usage['output_tokens']}"
+    assert usage["total_tokens"] == usage["input_tokens"] + usage["output_tokens"]
+
+    print(f"Reflect token usage: input={usage['input_tokens']}, output={usage['output_tokens']}, total={usage['total_tokens']}")
+
+
+@pytest.mark.asyncio
+async def test_retain_returns_token_usage(api_client):
+    """Test that retain endpoint returns token usage metrics for synchronous operations.
+
+    The usage field should contain input_tokens, output_tokens, and total_tokens
+    from the LLM calls made during fact extraction.
+    """
+    test_bank_id = f"retain_usage_test_{datetime.now().timestamp()}"
+
+    # Store memory synchronously (async=false is default)
+    response = await api_client.post(
+        f"/v1/default/banks/{test_bank_id}/memories",
+        json={
+            "items": [
+                {
+                    "content": "Alice is a software engineer at TechCorp. She specializes in machine learning.",
+                    "context": "team introduction"
+                }
+            ]
+        }
+    )
+    assert response.status_code == 200
+    result = response.json()
+
+    # Verify basic response
+    assert result["success"] is True
+    assert result["items_count"] == 1
+    assert result["async"] is False
+
+    # Verify usage field exists and has expected structure
+    assert "usage" in result, "Response should include 'usage' field"
+    usage = result["usage"]
+    assert usage is not None, "Usage should not be None for synchronous retain"
+    assert "input_tokens" in usage, "Usage should have 'input_tokens'"
+    assert "output_tokens" in usage, "Usage should have 'output_tokens'"
+    assert "total_tokens" in usage, "Usage should have 'total_tokens'"
+
+    # Verify token counts are valid
+    assert usage["input_tokens"] > 0, f"Expected input_tokens > 0, got {usage['input_tokens']}"
+    assert usage["output_tokens"] >= 0, f"Expected output_tokens >= 0, got {usage['output_tokens']}"
+    assert usage["total_tokens"] == usage["input_tokens"] + usage["output_tokens"]
+
+    print(f"Retain token usage: input={usage['input_tokens']}, output={usage['output_tokens']}, total={usage['total_tokens']}")
+
+
+@pytest.mark.asyncio
+async def test_retain_async_no_usage(api_client):
+    """Test that async retain does not return usage (as it's processed in background).
+
+    When async=true, the usage field should be None since the actual
+    fact extraction happens asynchronously.
+    """
+    test_bank_id = f"retain_async_no_usage_test_{datetime.now().timestamp()}"
+
+    # Store memory asynchronously
+    response = await api_client.post(
+        f"/v1/default/banks/{test_bank_id}/memories",
+        json={
+            "async": True,
+            "items": [
+                {
+                    "content": "Bob is a data scientist.",
+                    "context": "team introduction"
+                }
+            ]
+        }
+    )
+    assert response.status_code == 200
+    result = response.json()
+
+    # Verify async response
+    assert result["success"] is True
+    assert result["async"] is True
+
+    # Usage should be None for async operations
+    assert result.get("usage") is None, "Async retain should not include usage"

hindsight-api/tests/test_llm_token_metrics.py

@@ -0,0 +1,242 @@
+"""
+Test that LLM calls record token metrics via the metrics collector.
+"""
+import os
+from unittest.mock import MagicMock, patch
+import pytest
+from hindsight_api.engine.llm_wrapper import LLMProvider
+from hindsight_api.metrics import (
+    MetricsCollector,
+    NoOpMetricsCollector,
+    get_metrics_collector,
+    initialize_metrics,
+    create_metrics_collector,
+)
+
+
+def get_groq_api_key() -> str | None:
+    """Get Groq API key from environment."""
+    return os.getenv("GROQ_API_KEY")
+
+
+@pytest.mark.asyncio
+async def test_token_metrics_recorded_for_groq():
+    """
+    Test that token metrics are recorded when making LLM calls via Groq.
+    Uses openai/gpt-oss-20b as recommended by Hindsight.
+    """
+    api_key = get_groq_api_key()
+    if not api_key:
+        pytest.skip("Skipping: GROQ_API_KEY not set")
+
+    # Create a mock metrics collector to track record_tokens calls
+    mock_collector = MagicMock(spec=MetricsCollector)
+
+    with patch("hindsight_api.engine.llm_wrapper.get_metrics_collector", return_value=mock_collector):
+        llm = LLMProvider(
+            provider="groq",
+            api_key=api_key,
+            base_url="",
+            model="openai/gpt-oss-20b",
+        )
+
+        # Make an LLM call with clear instruction
+        response = await llm.call(
+            messages=[
+                {"role": "system", "content": "You are a helpful assistant. Always respond."},
+                {"role": "user", "content": "What is 2+2? Reply with just the number."}
+            ],
+            max_completion_tokens=50,
+            scope="test_metrics",
+        )
+
+        # Verify record_tokens was called - this is the main test
+        assert mock_collector.record_tokens.called, "record_tokens should have been called"
+
+        # Get the call arguments
+        call_kwargs = mock_collector.record_tokens.call_args.kwargs
+
+        # Verify the call had correct structure
+        assert call_kwargs["operation"] == "test_metrics", f"Expected operation='test_metrics', got {call_kwargs}"
+        assert call_kwargs["bank_id"] == "llm", f"Expected bank_id='llm', got {call_kwargs}"
+        assert call_kwargs["input_tokens"] > 0, f"Expected input_tokens > 0, got {call_kwargs['input_tokens']}"
+        # Output tokens may be 0 for some edge cases, but input should always be > 0
+        assert call_kwargs["output_tokens"] >= 0, f"Expected output_tokens >= 0, got {call_kwargs['output_tokens']}"
+
+        print(f"\nToken metrics recorded:")
+        print(f"  operation: {call_kwargs['operation']}")
+        print(f"  input_tokens: {call_kwargs['input_tokens']}")
+        print(f"  output_tokens: {call_kwargs['output_tokens']}")
+        print(f"  response: {response}")
+
+
+@pytest.mark.asyncio
+async def test_token_metrics_recorded_for_structured_output():
+    """
+    Test that token metrics are recorded for structured output (JSON) calls.
+    """
+    api_key = get_groq_api_key()
+    if not api_key:
+        pytest.skip("Skipping: GROQ_API_KEY not set")
+
+    from pydantic import BaseModel
+
+    class SimpleResponse(BaseModel):
+        greeting: str
+        language: str
+
+    mock_collector = MagicMock(spec=MetricsCollector)
+
+    with patch("hindsight_api.engine.llm_wrapper.get_metrics_collector", return_value=mock_collector):
+        llm = LLMProvider(
+            provider="groq",
+            api_key=api_key,
+            base_url="",
+            model="openai/gpt-oss-20b",
+        )
+
+        # Make a structured output call
+        response = await llm.call(
+            messages=[{"role": "user", "content": "Say hello in French. Return greeting and language."}],
+            response_format=SimpleResponse,
+            max_completion_tokens=100,
+            scope="structured_output_test",
+        )
+
+        # Verify structured response
+        assert isinstance(response, SimpleResponse)
+        assert response.greeting is not None
+        assert response.language is not None
+
+        # Verify record_tokens was called
+        assert mock_collector.record_tokens.called, "record_tokens should have been called"
+
+        call_kwargs = mock_collector.record_tokens.call_args.kwargs
+        assert call_kwargs["input_tokens"] > 0
+        assert call_kwargs["output_tokens"] > 0
+
+        print(f"\nStructured output token metrics:")
+        print(f"  greeting: {response.greeting}")
+        print(f"  language: {response.language}")
+        print(f"  input_tokens: {call_kwargs['input_tokens']}")
+        print(f"  output_tokens: {call_kwargs['output_tokens']}")
+
+
+@pytest.mark.asyncio
+async def test_noop_collector_when_metrics_disabled():
+    """
+    Test that NoOpMetricsCollector is returned when metrics are not initialized.
+    This verifies the fallback behavior doesn't break LLM calls.
+    """
+    api_key = get_groq_api_key()
+    if not api_key:
+        pytest.skip("Skipping: GROQ_API_KEY not set")
+
+    # Without initializing metrics, get_metrics_collector returns NoOpMetricsCollector
+    collector = get_metrics_collector()
+    assert isinstance(collector, NoOpMetricsCollector), "Should return NoOpMetricsCollector when not initialized"
+
+    # Make an LLM call - should work fine with NoOp collector
+    llm = LLMProvider(
+        provider="groq",
+        api_key=api_key,
+        base_url="",
+        model="openai/gpt-oss-20b",
+    )
+
+    response = await llm.call(
+        messages=[{"role": "user", "content": "Say 'test' in one word."}],
+        max_completion_tokens=50,
+    )
+
+    assert response is not None
+    print(f"\nLLM call succeeded with NoOpMetricsCollector: {response}")
+
+
+@pytest.mark.asyncio
+async def test_return_usage_returns_tuple():
+    """
+    Test that return_usage=True returns (result, TokenUsage) tuple.
+    """
+    from hindsight_api.engine.response_models import TokenUsage
+
+    api_key = get_groq_api_key()
+    if not api_key:
+        pytest.skip("Skipping: GROQ_API_KEY not set")
+
+    llm = LLMProvider(
+        provider="groq",
+        api_key=api_key,
+        base_url="",
+        model="openai/gpt-oss-20b",
+    )
+
+    # Call with return_usage=True
+    result, usage = await llm.call(
+        messages=[
+            {"role": "system", "content": "You are a helpful assistant."},
+            {"role": "user", "content": "What is 2+2? Reply with just the number."}
+        ],
+        max_completion_tokens=50,
+        return_usage=True,
+    )
+
+    # Verify result is the response text
+    assert result is not None
+    assert isinstance(result, str)
+
+    # Verify usage is TokenUsage model with valid counts
+    assert isinstance(usage, TokenUsage)
+    assert usage.input_tokens > 0, f"Expected input_tokens > 0, got {usage.input_tokens}"
+    assert usage.output_tokens >= 0, f"Expected output_tokens >= 0, got {usage.output_tokens}"
+    assert usage.total_tokens == usage.input_tokens + usage.output_tokens
+
+    print(f"\nreturn_usage=True test:")
+    print(f"  result: {result}")
+    print(f"  usage: {usage}")
+
+
+@pytest.mark.asyncio
+async def test_return_usage_with_structured_output():
+    """
+    Test that return_usage=True works with structured output (JSON).
+    """
+    from pydantic import BaseModel
+    from hindsight_api.engine.response_models import TokenUsage
+
+    api_key = get_groq_api_key()
+    if not api_key:
+        pytest.skip("Skipping: GROQ_API_KEY not set")
+
+    class MathAnswer(BaseModel):
+        answer: int
+        explanation: str
+
+    llm = LLMProvider(
+        provider="groq",
+        api_key=api_key,
+        base_url="",
+        model="openai/gpt-oss-20b",
+    )
+
+    # Call with return_usage=True and structured output
+    result, usage = await llm.call(
+        messages=[{"role": "user", "content": "What is 5+3? Return the answer and a brief explanation."}],
+        response_format=MathAnswer,
+        max_completion_tokens=100,
+        return_usage=True,
+    )
+
+    # Verify result is the parsed response
+    assert isinstance(result, MathAnswer)
+    assert result.answer == 8
+    assert result.explanation is not None
+
+    # Verify usage is TokenUsage model
+    assert isinstance(usage, TokenUsage)
+    assert usage.input_tokens > 0
+    assert usage.output_tokens > 0
+
+    print(f"\nStructured output with return_usage=True:")
+    print(f"  result: {result}")
+    print(f"  usage: {usage}")

hindsight-clients/python/.openapi-generator/FILES

@@ -51,6 +51,7 @@ hindsight_client_api/models/reflect_request.py
 hindsight_client_api/models/reflect_response.py
 hindsight_client_api/models/retain_request.py
 hindsight_client_api/models/retain_response.py
+hindsight_client_api/models/token_usage.py
 hindsight_client_api/models/update_disposition_request.py
 hindsight_client_api/models/validation_error.py
 hindsight_client_api/models/validation_error_loc_inner.py

hindsight-clients/python/hindsight_client_api/__init__.py

@@ -76,6 +76,7 @@ from hindsight_client_api.models.reflect_request import ReflectRequest
 from hindsight_client_api.models.reflect_response import ReflectResponse
 from hindsight_client_api.models.retain_request import RetainRequest
 from hindsight_client_api.models.retain_response import RetainResponse
+from hindsight_client_api.models.token_usage import TokenUsage
 from hindsight_client_api.models.update_disposition_request import UpdateDispositionRequest
 from hindsight_client_api.models.validation_error import ValidationError
 from hindsight_client_api.models.validation_error_loc_inner import ValidationErrorLocInner

hindsight-clients/python/hindsight_client_api/api/memory_api.py

@@ -347,6 +347,7 @@ class MemoryApi:
         self,
         bank_id: StrictStr,
         type: Optional[StrictStr] = None,
+        limit: Optional[StrictInt] = None,
         authorization: Optional[StrictStr] = None,
         _request_timeout: Union[
             None,
@@ -363,12 +364,14 @@ class MemoryApi:
     ) -> GraphDataResponse:
         """Get memory graph data
 
-        Retrieve graph data for visualization, optionally filtered by type (world/experience/opinion). Limited to 1000 most recent items.
+        Retrieve graph data for visualization, optionally filtered by type (world/experience/opinion).
 
         :param bank_id: (required)
         :type bank_id: str
         :param type:
         :type type: str
+        :param limit:
+        :type limit: int
         :param authorization:
         :type authorization: str
         :param _request_timeout: timeout setting for this request. If one
@@ -396,6 +399,7 @@ class MemoryApi:
         _param = self._get_graph_serialize(
             bank_id=bank_id,
             type=type,
+            limit=limit,
             authorization=authorization,
             _request_auth=_request_auth,
             _content_type=_content_type,
@@ -423,6 +427,7 @@ class MemoryApi:
         self,
         bank_id: StrictStr,
         type: Optional[StrictStr] = None,
+        limit: Optional[StrictInt] = None,
         authorization: Optional[StrictStr] = None,
         _request_timeout: Union[
             None,
@@ -439,12 +444,14 @@ class MemoryApi:
     ) -> ApiResponse[GraphDataResponse]:
         """Get memory graph data
 
-        Retrieve graph data for visualization, optionally filtered by type (world/experience/opinion). Limited to 1000 most recent items.
+        Retrieve graph data for visualization, optionally filtered by type (world/experience/opinion).
 
         :param bank_id: (required)
         :type bank_id: str
         :param type:
         :type type: str
+        :param limit:
+        :type limit: int
         :param authorization:
         :type authorization: str
         :param _request_timeout: timeout setting for this request. If one
@@ -472,6 +479,7 @@ class MemoryApi:
         _param = self._get_graph_serialize(
             bank_id=bank_id,
             type=type,
+            limit=limit,
             authorization=authorization,
             _request_auth=_request_auth,
             _content_type=_content_type,
@@ -499,6 +507,7 @@ class MemoryApi:
         self,
         bank_id: StrictStr,
         type: Optional[StrictStr] = None,
+        limit: Optional[StrictInt] = None,
         authorization: Optional[StrictStr] = None,
         _request_timeout: Union[
             None,
@@ -515,12 +524,14 @@ class MemoryApi:
     ) -> RESTResponseType:
         """Get memory graph data
 
-        Retrieve graph data for visualization, optionally filtered by type (world/experience/opinion). Limited to 1000 most recent items.
+        Retrieve graph data for visualization, optionally filtered by type (world/experience/opinion).
 
         :param bank_id: (required)
         :type bank_id: str
         :param type:
         :type type: str
+        :param limit:
+        :type limit: int
         :param authorization:
         :type authorization: str
         :param _request_timeout: timeout setting for this request. If one
@@ -548,6 +559,7 @@ class MemoryApi:
         _param = self._get_graph_serialize(
             bank_id=bank_id,
             type=type,
+            limit=limit,
             authorization=authorization,
             _request_auth=_request_auth,
             _content_type=_content_type,
@@ -570,6 +582,7 @@ class MemoryApi:
         self,
         bank_id,
         type,
+        limit,
         authorization,
         _request_auth,
         _content_type,
@@ -599,6 +612,10 @@ class MemoryApi:
             
             _query_params.append(('type', type))
             
+        if limit is not None:
+            
+            _query_params.append(('limit', limit))
+            
         # process the header parameters
         if authorization is not None:
             _header_params['authorization'] = authorization

hindsight-clients/python/hindsight_client_api/models/__init__.py

@@ -54,6 +54,7 @@ from hindsight_client_api.models.reflect_request import ReflectRequest
 from hindsight_client_api.models.reflect_response import ReflectResponse
 from hindsight_client_api.models.retain_request import RetainRequest
 from hindsight_client_api.models.retain_response import RetainResponse
+from hindsight_client_api.models.token_usage import TokenUsage
 from hindsight_client_api.models.update_disposition_request import UpdateDispositionRequest
 from hindsight_client_api.models.validation_error import ValidationError
 from hindsight_client_api.models.validation_error_loc_inner import ValidationErrorLocInner

hindsight-clients/python/hindsight_client_api/models/graph_data_response.py

@@ -30,7 +30,8 @@ class GraphDataResponse(BaseModel):
     edges: List[Dict[str, Any]]
     table_rows: List[Dict[str, Any]]
     total_units: StrictInt
-    __properties: ClassVar[List[str]] = ["nodes", "edges", "table_rows", "total_units"]
+    limit: StrictInt
+    __properties: ClassVar[List[str]] = ["nodes", "edges", "table_rows", "total_units", "limit"]
 
     model_config = ConfigDict(
         populate_by_name=True,
@@ -86,7 +87,8 @@ class GraphDataResponse(BaseModel):
             "nodes": obj.get("nodes"),
             "edges": obj.get("edges"),
             "table_rows": obj.get("table_rows"),
-            "total_units": obj.get("total_units")
+            "total_units": obj.get("total_units"),
+            "limit": obj.get("limit")
         })
         return _obj
 

hindsight-clients/python/hindsight_client_api/models/reflect_response.py

@@ -20,6 +20,7 @@ import json
 from pydantic import BaseModel, ConfigDict, StrictStr
 from typing import Any, ClassVar, Dict, List, Optional
 from hindsight_client_api.models.reflect_fact import ReflectFact
+from hindsight_client_api.models.token_usage import TokenUsage
 from typing import Optional, Set
 from typing_extensions import Self
 
@@ -30,7 +31,8 @@ class ReflectResponse(BaseModel):
     text: StrictStr
     based_on: Optional[List[ReflectFact]] = None
     structured_output: Optional[Dict[str, Any]] = None
-    __properties: ClassVar[List[str]] = ["text", "based_on", "structured_output"]
+    usage: Optional[TokenUsage] = None
+    __properties: ClassVar[List[str]] = ["text", "based_on", "structured_output", "usage"]
 
     model_config = ConfigDict(
         populate_by_name=True,
@@ -78,11 +80,19 @@ class ReflectResponse(BaseModel):
                 if _item_based_on:
                     _items.append(_item_based_on.to_dict())
             _dict['based_on'] = _items
+        # override the default output from pydantic by calling `to_dict()` of usage
+        if self.usage:
+            _dict['usage'] = self.usage.to_dict()
         # set to None if structured_output (nullable) is None
         # and model_fields_set contains the field
         if self.structured_output is None and "structured_output" in self.model_fields_set:
             _dict['structured_output'] = None
 
+        # set to None if usage (nullable) is None
+        # and model_fields_set contains the field
+        if self.usage is None and "usage" in self.model_fields_set:
+            _dict['usage'] = None
+
         return _dict
 
     @classmethod
@@ -97,7 +107,8 @@ class ReflectResponse(BaseModel):
         _obj = cls.model_validate({
             "text": obj.get("text"),
             "based_on": [ReflectFact.from_dict(_item) for _item in obj["based_on"]] if obj.get("based_on") is not None else None,
-            "structured_output": obj.get("structured_output")
+            "structured_output": obj.get("structured_output"),
+            "usage": TokenUsage.from_dict(obj["usage"]) if obj.get("usage") is not None else None
         })
         return _obj
 

hindsight-clients/python/hindsight_client_api/models/retain_response.py

@@ -18,7 +18,8 @@ import re  # noqa: F401
 import json
 
 from pydantic import BaseModel, ConfigDict, Field, StrictBool, StrictInt, StrictStr
-from typing import Any, ClassVar, Dict, List
+from typing import Any, ClassVar, Dict, List, Optional
+from hindsight_client_api.models.token_usage import TokenUsage
 from typing import Optional, Set
 from typing_extensions import Self
 
@@ -30,7 +31,8 @@ class RetainResponse(BaseModel):
     bank_id: StrictStr
     items_count: StrictInt
     var_async: StrictBool = Field(description="Whether the operation was processed asynchronously", alias="async")
-    __properties: ClassVar[List[str]] = ["success", "bank_id", "items_count", "async"]
+    usage: Optional[TokenUsage] = None
+    __properties: ClassVar[List[str]] = ["success", "bank_id", "items_count", "async", "usage"]
 
     model_config = ConfigDict(
         populate_by_name=True,
@@ -71,6 +73,14 @@ class RetainResponse(BaseModel):
             exclude=excluded_fields,
             exclude_none=True,
         )
+        # override the default output from pydantic by calling `to_dict()` of usage
+        if self.usage:
+            _dict['usage'] = self.usage.to_dict()
+        # set to None if usage (nullable) is None
+        # and model_fields_set contains the field
+        if self.usage is None and "usage" in self.model_fields_set:
+            _dict['usage'] = None
+
         return _dict
 
     @classmethod
@@ -86,7 +96,8 @@ class RetainResponse(BaseModel):
             "success": obj.get("success"),
             "bank_id": obj.get("bank_id"),
             "items_count": obj.get("items_count"),
-            "async": obj.get("async")
+            "async": obj.get("async"),
+            "usage": TokenUsage.from_dict(obj["usage"]) if obj.get("usage") is not None else None
         })
         return _obj
 

hindsight-clients/python/hindsight_client_api/models/token_usage.py

@@ -0,0 +1,91 @@
+# coding: utf-8
+
+"""
+    Hindsight HTTP API
+
+    HTTP API for Hindsight
+
+    The version of the OpenAPI document: 0.1.0
+    Generated by OpenAPI Generator (https://openapi-generator.tech)
+
+    Do not edit the class manually.
+"""  # noqa: E501
+
+
+from __future__ import annotations
+import pprint
+import re  # noqa: F401
+import json
+
+from pydantic import BaseModel, ConfigDict, Field, StrictInt
+from typing import Any, ClassVar, Dict, List, Optional
+from typing import Optional, Set
+from typing_extensions import Self
+
+class TokenUsage(BaseModel):
+    """
+    Token usage metrics for LLM calls.  Tracks input/output tokens for a single request to enable per-request cost tracking and monitoring.
+    """ # noqa: E501
+    input_tokens: Optional[StrictInt] = Field(default=0, description="Number of input/prompt tokens consumed")
+    output_tokens: Optional[StrictInt] = Field(default=0, description="Number of output/completion tokens generated")
+    total_tokens: Optional[StrictInt] = Field(default=0, description="Total tokens (input + output)")
+    __properties: ClassVar[List[str]] = ["input_tokens", "output_tokens", "total_tokens"]
+
+    model_config = ConfigDict(
+        populate_by_name=True,
+        validate_assignment=True,
+        protected_namespaces=(),
+    )
+
+
+    def to_str(self) -> str:
+        """Returns the string representation of the model using alias"""
+        return pprint.pformat(self.model_dump(by_alias=True))
+
+    def to_json(self) -> str:
+        """Returns the JSON representation of the model using alias"""
+        # TODO: pydantic v2: use .model_dump_json(by_alias=True, exclude_unset=True) instead
+        return json.dumps(self.to_dict())
+
+    @classmethod
+    def from_json(cls, json_str: str) -> Optional[Self]:
+        """Create an instance of TokenUsage from a JSON string"""
+        return cls.from_dict(json.loads(json_str))
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Return the dictionary representation of the model using alias.
+
+        This has the following differences from calling pydantic's
+        `self.model_dump(by_alias=True)`:
+
+        * `None` is only added to the output dict for nullable fields that
+          were set at model initialization. Other fields with value `None`
+          are ignored.
+        """
+        excluded_fields: Set[str] = set([
+        ])
+
+        _dict = self.model_dump(
+            by_alias=True,
+            exclude=excluded_fields,
+            exclude_none=True,
+        )
+        return _dict
+
+    @classmethod
+    def from_dict(cls, obj: Optional[Dict[str, Any]]) -> Optional[Self]:
+        """Create an instance of TokenUsage from a dict"""
+        if obj is None:
+            return None
+
+        if not isinstance(obj, dict):
+            return cls.model_validate(obj)
+
+        _obj = cls.model_validate({
+            "input_tokens": obj.get("input_tokens") if obj.get("input_tokens") is not None else 0,
+            "output_tokens": obj.get("output_tokens") if obj.get("output_tokens") is not None else 0,
+            "total_tokens": obj.get("total_tokens") if obj.get("total_tokens") is not None else 0
+        })
+        return _obj
+
+

hindsight-clients/typescript/.prettierrc

@@ -0,0 +1,8 @@
+{
+  "tabWidth": 2,
+  "useTabs": false,
+  "semi": true,
+  "singleQuote": false,
+  "trailingComma": "all",
+  "printWidth": 80
+}

hindsight-clients/typescript/generated/sdk.gen.ts

@@ -123,7 +123,7 @@ export const metricsEndpointMetricsGet = <ThrowOnError extends boolean = false>(
 /**
  * Get memory graph data
  *
- * Retrieve graph data for visualization, optionally filtered by type (world/experience/opinion). Limited to 1000 most recent items.
+ * Retrieve graph data for visualization, optionally filtered by type (world/experience/opinion).
  */
 export const getGraph = <ThrowOnError extends boolean = false>(
   options: Options<GetGraphData, ThrowOnError>,

hindsight-clients/typescript/generated/types.gen.ts

@@ -570,6 +570,10 @@ export type GraphDataResponse = {
    * Total Units
    */
   total_units: number;
+  /**
+   * Limit
+   */
+  limit: number;
 };
 
 /**
@@ -966,6 +970,10 @@ export type ReflectResponse = {
   structured_output?: {
     [key: string]: unknown;
   } | null;
+  /**
+   * Token usage metrics for LLM calls during reflection.
+   */
+  usage?: TokenUsage | null;
 };
 
 /**
@@ -1010,6 +1018,39 @@ export type RetainResponse = {
    * Whether the operation was processed asynchronously
    */
   async: boolean;
+  /**
+   * Token usage metrics for LLM calls during fact extraction (only present for synchronous operations)
+   */
+  usage?: TokenUsage | null;
+};
+
+/**
+ * TokenUsage
+ *
+ * Token usage metrics for LLM calls.
+ *
+ * Tracks input/output tokens for a single request to enable
+ * per-request cost tracking and monitoring.
+ */
+export type TokenUsage = {
+  /**
+   * Input Tokens
+   *
+   * Number of input/prompt tokens consumed
+   */
+  input_tokens?: number;
+  /**
+   * Output Tokens
+   *
+   * Number of output/completion tokens generated
+   */
+  output_tokens?: number;
+  /**
+   * Total Tokens
+   *
+   * Total tokens (input + output)
+   */
+  total_tokens?: number;
 };
 
 /**
@@ -1086,6 +1127,10 @@ export type GetGraphData = {
      * Type
      */
     type?: string | null;
+    /**
+     * Limit
+     */
+    limit?: number;
   };
   url: "/v1/default/banks/{bank_id}/graph";
 };

hindsight-clients/typescript/package.json

@@ -26,7 +26,7 @@
     "directory": "hindsight-clients/typescript"
   },
   "devDependencies": {
-    "@hey-api/openapi-ts": "^0.88.0",
+    "@hey-api/openapi-ts": "0.88.0",
     "@types/jest": "^29.0.0",
     "@types/node": "^20.0.0",
     "jest": "^29.0.0",

hindsight-control-plane/src/app/api/graph/route.ts

@@ -12,12 +12,15 @@ export async function GET(request: NextRequest) {
 
     // Get optional query parameters
     const type = searchParams.get("type") || searchParams.get("fact_type") || undefined;
+    const limitParam = searchParams.get("limit");
+    const limit = limitParam ? parseInt(limitParam, 10) : undefined;
 
     const response = await sdk.getGraph({
       client: lowLevelClient,
       path: { bank_id: bankId },
       query: {
         type: type,
+        limit: limit,
       },
     });
 

hindsight-control-plane/src/components/data-view.tsx

@@ -52,6 +52,9 @@ export function DataView({ factType }: DataViewProps) {
   const [selectedTableMemory, setSelectedTableMemory] = useState<any>(null);
   const itemsPerPage = 100;
 
+  // Fetch limit state - how many memories to load from the API
+  const [fetchLimit, setFetchLimit] = useState(1000);
+
   // Graph controls state
   const [showLabels, setShowLabels] = useState(true);
   const [maxNodes, setMaxNodes] = useState<number | undefined>(undefined);
@@ -93,7 +96,7 @@ export function DataView({ factType }: DataViewProps) {
     }
   };
 
-  const loadData = async () => {
+  const loadData = async (limit?: number) => {
     if (!currentBank) return;
 
     setLoading(true);
@@ -101,6 +104,7 @@ export function DataView({ factType }: DataViewProps) {
       const graphData: any = await client.getGraph({
         bank_id: currentBank,
         type: factType,
+        limit: limit ?? fetchLimit,
       });
       setData(graphData);
     } catch (error) {
@@ -265,9 +269,25 @@ export function DataView({ factType }: DataViewProps) {
 
           <div className="flex items-center justify-between mb-6">
             <div className="text-sm text-muted-foreground">
-              {searchQuery
-                ? `${filteredTableRows.length} of ${data.total_units} memories`
-                : `${data.total_units} total memories`}
+              {searchQuery ? (
+                `${filteredTableRows.length} of ${data.table_rows?.length ?? 0} loaded memories`
+              ) : data.table_rows?.length < data.total_units ? (
+                <span>
+                  Showing {data.table_rows?.length ?? 0} of {data.total_units} total memories
+                  <button
+                    onClick={() => {
+                      const newLimit = Math.min(data.total_units, fetchLimit + 1000);
+                      setFetchLimit(newLimit);
+                      loadData(newLimit);
+                    }}
+                    className="ml-2 text-primary hover:underline"
+                  >
+                    Load more
+                  </button>
+                </span>
+              ) : (
+                `${data.total_units} total memories`
+              )}
             </div>
             <div className="flex items-center gap-2 bg-muted rounded-lg p-1">
               <button

hindsight-control-plane/src/lib/api.ts

@@ -109,10 +109,11 @@ export class ControlPlaneClient {
   /**
    * Get graph data
    */
-  async getGraph(params: { bank_id: string; type?: string }) {
+  async getGraph(params: { bank_id: string; type?: string; limit?: number }) {
     const queryParams = new URLSearchParams();
     queryParams.append("bank_id", params.bank_id);
     if (params.type) queryParams.append("type", params.type);
+    if (params.limit) queryParams.append("limit", params.limit.toString());
     return this.fetchApi(`/api/graph?${queryParams}`);
   }
 

hindsight-docs/docs/changelog/index.md

@@ -8,6 +8,12 @@ This changelog highlights user-facing changes only. Internal maintenance, CI/CD,
 
 For full release details, see [GitHub Releases](https://github.com/vectorize-io/hindsight/releases).
 
+## [Unreleased]
+
+**Features**
+
+- Add per-request token usage tracking to retain and reflect endpoints for cost monitoring and billing integration.
+
 ## [0.2.0](https://github.com/vectorize-io/hindsight/releases/tag/v0.2.0)
 
 **Features**

hindsight-docs/docs/developer/admin-cli.md

@@ -0,0 +1,145 @@
+# Admin CLI
+
+The `hindsight-admin` CLI provides administrative commands for managing your Hindsight deployment, including database migrations, backup, and restore operations.
+
+## Installation
+
+The admin CLI is included with the `hindsight-api` package:
+
+```bash
+pip install hindsight-api
+# or
+uv add hindsight-api
+```
+
+## Commands
+
+### run-db-migration
+
+Run database migrations to the latest version. This is useful when you want to run migrations separately from API startup (e.g., in CI/CD pipelines or before deploying a new version).
+
+```bash
+hindsight-admin run-db-migration [OPTIONS]
+```
+
+**Options:**
+
+| Option | Description | Default |
+|--------|-------------|---------|
+| `--schema`, `-s` | Database schema to run migrations on | `public` |
+
+**Examples:**
+
+```bash
+# Run migrations on the default public schema
+hindsight-admin run-db-migration
+
+# Run migrations on a specific tenant schema
+hindsight-admin run-db-migration --schema tenant_acme
+```
+
+:::tip Disabling Auto-Migrations
+To disable automatic migrations on API startup, set `HINDSIGHT_API_RUN_MIGRATIONS_ON_STARTUP=false`. This is useful when you want to run migrations as a separate step in your deployment pipeline.
+:::
+
+---
+
+### backup
+
+Create a backup of all Hindsight data to a zip file.
+
+```bash
+hindsight-admin backup OUTPUT [OPTIONS]
+```
+
+**Arguments:**
+
+| Argument | Description |
+|----------|-------------|
+| `OUTPUT` | Output file path (will add `.zip` extension if not present) |
+
+**Options:**
+
+| Option | Description | Default |
+|--------|-------------|---------|
+| `--schema`, `-s` | Database schema to backup | `public` |
+
+**Examples:**
+
+```bash
+# Backup to a file
+hindsight-admin backup /backups/hindsight-2024-01-15.zip
+
+# Backup a specific tenant schema
+hindsight-admin backup /backups/tenant-acme.zip --schema tenant_acme
+```
+
+The backup includes:
+- Memory banks and their configuration
+- Documents and chunks
+- Entities and their relationships
+- Memory units (facts, experiences, opinions, observations)
+- Entity cooccurrences and memory links
+
+:::note Consistency
+Backups are created within a database transaction with `REPEATABLE READ` isolation, ensuring a consistent snapshot across all tables.
+:::
+
+---
+
+### restore
+
+Restore data from a backup file. **Warning: This deletes all existing data in the target schema.**
+
+```bash
+hindsight-admin restore INPUT [OPTIONS]
+```
+
+**Arguments:**
+
+| Argument | Description |
+|----------|-------------|
+| `INPUT` | Input backup file (.zip) |
+
+**Options:**
+
+| Option | Description | Default |
+|--------|-------------|---------|
+| `--schema`, `-s` | Database schema to restore to | `public` |
+| `--yes`, `-y` | Skip confirmation prompt | `false` |
+
+**Examples:**
+
+```bash
+# Restore with confirmation prompt
+hindsight-admin restore /backups/hindsight-2024-01-15.zip
+
+# Restore without confirmation (for scripts)
+hindsight-admin restore /backups/hindsight-2024-01-15.zip --yes
+
+# Restore to a specific tenant schema
+hindsight-admin restore /backups/tenant-acme.zip --schema tenant_acme --yes
+```
+
+:::warning Data Loss
+Restore will **delete all existing data** in the target schema before importing the backup. Always verify you have a recent backup before performing a restore.
+:::
+
+---
+
+## Environment Variables
+
+The admin CLI uses the same environment variables as the API service. The most important one is:
+
+| Variable | Description | Default |
+|----------|-------------|---------|
+| `HINDSIGHT_API_DATABASE_URL` | PostgreSQL connection string | `pg0` (embedded) |
+
+**Example:**
+
+```bash
+# Use a specific database
+export HINDSIGHT_API_DATABASE_URL=postgresql://user:pass@localhost:5432/hindsight
+hindsight-admin backup /backups/mybackup.zip
+```
+

hindsight-docs/docs/developer/api/reflect.mdx

@@ -55,6 +55,20 @@ Make sure you've completed the [Quick Start](./quickstart) to install the client
 | `max_tokens` | int | 4096 | Maximum tokens for the response |
 | `response_schema` | object | None | JSON Schema for [structured output](#structured-output) |
 
+### Response Fields
+
+| Field | Type | Description |
+|-------|------|-------------|
+| `text` | string | The generated answer text |
+| `based_on` | array | Facts used to generate the response |
+| `structured_output` | object | Parsed structured output (when `response_schema` provided) |
+| `usage` | TokenUsage | Token usage metrics for the LLM call |
+
+The `usage` field contains:
+- `input_tokens`: Number of input/prompt tokens consumed
+- `output_tokens`: Number of output/completion tokens generated
+- `total_tokens`: Sum of input and output tokens
+
 <Tabs>
 <TabItem value="python" label="Python">
 <CodeSnippet code={reflectPy} section="reflect-with-params" language="python" />

hindsight-docs/docs/developer/api/retain.mdx

@@ -66,6 +66,25 @@ Always provide context and event dates for optimal memory extraction:
 
 The `timestamp` defaults to the current time if not specified. Providing explicit timestamps enables temporal queries like "What happened last spring?"
 
+### Response Fields
+
+The retain response includes:
+
+| Field | Type | Description |
+|-------|------|-------------|
+| `success` | bool | Whether the operation succeeded |
+| `bank_id` | string | The memory bank ID |
+| `items_count` | int | Number of items processed |
+| `async` | bool | Whether processed asynchronously |
+| `usage` | TokenUsage | Token usage metrics for LLM calls (synchronous only) |
+
+The `usage` field contains token metrics for cost tracking:
+- `input_tokens`: Tokens consumed by prompts
+- `output_tokens`: Tokens generated by the LLM
+- `total_tokens`: Sum of input and output tokens
+
+Note: `usage` is only present for synchronous operations. Async operations (`async: true`) do not return usage metrics.
+
 ## Batch Ingestion
 
 Store multiple items in a single request. **Batch ingestion is the recommended approach** as it significantly improves performance by reducing network overhead and allowing Hindsight to optimize the memory extraction process across related content.

hindsight-docs/docs/developer/configuration.md

@@ -20,9 +20,18 @@ The API service handles all memory operations (retain, recall, reflect).
 | Variable | Description | Default |
 |----------|-------------|---------|
 | `HINDSIGHT_API_DATABASE_URL` | PostgreSQL connection string | `pg0` (embedded) |
+| `HINDSIGHT_API_RUN_MIGRATIONS_ON_STARTUP` | Run database migrations on API startup | `true` |
 
 If not provided, the server uses embedded `pg0` — convenient for development but not recommended for production.
 
+To run migrations manually (e.g., before starting the API), use the admin CLI:
+
+```bash
+hindsight-admin run-db-migration
+# Or for a specific schema:
+hindsight-admin run-db-migration --schema tenant_acme
+```
+
 ### LLM Provider
 
 | Variable | Description | Default |
@@ -183,6 +192,14 @@ Controls when the system generates entity observations (summaries about entities
 | `HINDSIGHT_API_OBSERVATION_MIN_FACTS` | Minimum facts about an entity before generating observations | `5` |
 | `HINDSIGHT_API_OBSERVATION_TOP_ENTITIES` | Max entities to process per retain batch | `5` |
 
+### Retain
+
+Controls the retain (memory ingestion) pipeline.
+
+| Variable | Description | Default |
+|----------|-------------|---------|
+| `HINDSIGHT_API_RETAIN_MAX_COMPLETION_TOKENS` | Max completion tokens for fact extraction LLM calls | `64000` |
+
 ### Local MCP Server
 
 Configuration for the local MCP server (`hindsight-local-mcp` command).

hindsight-docs/sidebars.ts

@@ -111,6 +111,11 @@ const sidebars: SidebarsConfig = {
           id: 'developer/configuration',
           label: 'Configuration',
         },
+        {
+          type: 'doc',
+          id: 'developer/admin-cli',
+          label: 'Admin CLI',
+        },
         {
           type: 'doc',
           id: 'developer/extensions',

hindsight-docs/static/openapi.json

@@ -59,7 +59,7 @@
           "Memory"
         ],
         "summary": "Get memory graph data",
-        "description": "Retrieve graph data for visualization, optionally filtered by type (world/experience/opinion). Limited to 1000 most recent items.",
+        "description": "Retrieve graph data for visualization, optionally filtered by type (world/experience/opinion).",
         "operationId": "get_graph",
         "parameters": [
           {
@@ -87,6 +87,16 @@
               "title": "Type"
             }
           },
+          {
+            "name": "limit",
+            "in": "query",
+            "required": false,
+            "schema": {
+              "type": "integer",
+              "default": 1000,
+              "title": "Limit"
+            }
+          },
           {
             "name": "authorization",
             "in": "header",
@@ -2504,6 +2514,10 @@
           "total_units": {
             "type": "integer",
             "title": "Total Units"
+          },
+          "limit": {
+            "type": "integer",
+            "title": "Limit"
           }
         },
         "type": "object",
@@ -2511,7 +2525,8 @@
           "nodes",
           "edges",
           "table_rows",
-          "total_units"
+          "total_units",
+          "limit"
         ],
         "title": "GraphDataResponse",
         "description": "Response model for graph data endpoint.",
@@ -2524,6 +2539,7 @@
               "weight": 0.8
             }
           ],
+          "limit": 1000,
           "nodes": [
             {
               "id": "1",
@@ -3404,6 +3420,17 @@
             ],
             "title": "Structured Output",
             "description": "Structured output parsed according to the request's response_schema. Only present when response_schema was provided in the request."
+          },
+          "usage": {
+            "anyOf": [
+              {
+                "$ref": "#/components/schemas/TokenUsage"
+              },
+              {
+                "type": "null"
+              }
+            ],
+            "description": "Token usage metrics for LLM calls during reflection."
           }
         },
         "type": "object",
@@ -3432,7 +3459,12 @@
             ],
             "summary": "AI is transformative"
           },
-          "text": "Based on my understanding, AI is a transformative technology..."
+          "text": "Based on my understanding, AI is a transformative technology...",
+          "usage": {
+            "input_tokens": 1500,
+            "output_tokens": 500,
+            "total_tokens": 2000
+          }
         }
       },
       "RetainRequest": {
@@ -3491,6 +3523,17 @@
             "type": "boolean",
             "title": "Async",
             "description": "Whether the operation was processed asynchronously"
+          },
+          "usage": {
+            "anyOf": [
+              {
+                "$ref": "#/components/schemas/TokenUsage"
+              },
+              {
+                "type": "null"
+              }
+            ],
+            "description": "Token usage metrics for LLM calls during fact extraction (only present for synchronous operations)"
           }
         },
         "type": "object",
@@ -3506,7 +3549,42 @@
           "async": false,
           "bank_id": "user123",
           "items_count": 2,
-          "success": true
+          "success": true,
+          "usage": {
+            "input_tokens": 500,
+            "output_tokens": 100,
+            "total_tokens": 600
+          }
+        }
+      },
+      "TokenUsage": {
+        "properties": {
+          "input_tokens": {
+            "type": "integer",
+            "title": "Input Tokens",
+            "description": "Number of input/prompt tokens consumed",
+            "default": 0
+          },
+          "output_tokens": {
+            "type": "integer",
+            "title": "Output Tokens",
+            "description": "Number of output/completion tokens generated",
+            "default": 0
+          },
+          "total_tokens": {
+            "type": "integer",
+            "title": "Total Tokens",
+            "description": "Total tokens (input + output)",
+            "default": 0
+          }
+        },
+        "type": "object",
+        "title": "TokenUsage",
+        "description": "Token usage metrics for LLM calls.\n\nTracks input/output tokens for a single request to enable\nper-request cost tracking and monitoring.",
+        "example": {
+          "input_tokens": 1500,
+          "output_tokens": 500,
+          "total_tokens": 2000
         }
       },
       "UpdateDispositionRequest": {

package.json

@@ -8,5 +8,8 @@
   ],
   "scripts": {
     "prepare": "./scripts/setup-hooks.sh"
+  },
+  "overrides": {
+    "qs": "^6.14.1"
   }
 }

scripts/generate-clients.sh

@@ -51,16 +51,16 @@ echo "=================================================="
 
 RUST_CLIENT_DIR="$CLIENTS_DIR/rust"
 
-# Clean old generated files
+# Clean old generated files (keep Cargo.lock for reproducible builds)
 echo "Cleaning old Rust generated code..."
 rm -rf "$RUST_CLIENT_DIR/target"
-rm -f "$RUST_CLIENT_DIR/Cargo.lock"
 
 # Trigger regeneration by building
+# Use --locked to ensure reproducible builds from committed Cargo.lock
 echo "Regenerating Rust client (via build.rs)..."
 cd "$RUST_CLIENT_DIR"
 cargo clean
-cargo build --release
+cargo build --release --locked
 
 echo "✓ Rust client generated at $RUST_CLIENT_DIR"
 echo ""
@@ -324,9 +324,11 @@ rm -rf "$TYPESCRIPT_CLIENT_DIR/services"
 rm -f "$TYPESCRIPT_CLIENT_DIR/index.ts"
 
 # Generate new client using @hey-api/openapi-ts
+# Use npm run generate to use the locally installed version (pinned in package.json)
+# instead of npx --yes which would fetch the latest version
 echo "Generating from $OPENAPI_SPEC..."
 cd "$TYPESCRIPT_CLIENT_DIR"
-npx --yes @hey-api/openapi-ts
+npm run generate
 
 echo "✓ TypeScript client generated at $TYPESCRIPT_CLIENT_DIR"
 echo ""

uv.lock

@@ -1215,6 +1215,7 @@ dependencies = [
     { name = "tiktoken" },
     { name = "torch" },
     { name = "transformers" },
+    { name = "typer" },
     { name = "uvicorn" },
     { name = "wsproto" },
 ]
@@ -1274,6 +1275,7 @@ requires-dist = [
     { name = "tiktoken", specifier = ">=0.12.0" },
     { name = "torch", specifier = ">=2.0.0" },
     { name = "transformers", specifier = ">=4.30.0,<4.46.0" },
+    { name = "typer", specifier = ">=0.9.0" },
     { name = "uvicorn", specifier = ">=0.38.0" },
     { name = "wsproto", specifier = ">=1.0.0" },
 ]