chore: add flag to not include ml libs in docker image

* fix: duplicated causal relationships and token optimization

* doc

* doc

.env.example

@@ -2,11 +2,23 @@
 # Copy this file to .env and fill in your values
 
 # LLM Configuration (Required)
+# Supported providers: openai, groq, ollama, gemini, anthropic, lmstudio
 HINDSIGHT_API_LLM_PROVIDER=openai
 HINDSIGHT_API_LLM_API_KEY=your-api-key-here
 HINDSIGHT_API_LLM_MODEL=o3-mini
 HINDSIGHT_API_LLM_BASE_URL=https://api.openai.com/v1
 
+# Example: Anthropic Claude configuration
+# HINDSIGHT_API_LLM_PROVIDER=anthropic
+# HINDSIGHT_API_LLM_API_KEY=your-anthropic-api-key
+# HINDSIGHT_API_LLM_MODEL=claude-sonnet-4-20250514
+
+# Example: LM Studio local configuration (Qwen 2.5 32B recommended)
+# HINDSIGHT_API_LLM_PROVIDER=lmstudio
+# HINDSIGHT_API_LLM_API_KEY=lmstudio
+# HINDSIGHT_API_LLM_BASE_URL=http://localhost:1234/v1
+# HINDSIGHT_API_LLM_MODEL=qwen2.5-32b-instruct
+
 # API Configuration (Optional)
 HINDSIGHT_API_HOST=0.0.0.0
 HINDSIGHT_API_PORT=8888

.github/ISSUE_TEMPLATE/bug_report.yml

@@ -0,0 +1,71 @@
+name: Bug Report
+description: Report a bug or unexpected behavior
+labels: ["bug", "triage"]
+body:
+  - type: markdown
+    attributes:
+      value: |
+        Thanks for taking the time to report a bug! Please fill out the sections below.
+
+  - type: textarea
+    id: description
+    attributes:
+      label: Bug Description
+      description: A clear and concise description of the bug
+      placeholder: What happened?
+    validations:
+      required: true
+
+  - type: textarea
+    id: reproduction
+    attributes:
+      label: Steps to Reproduce
+      description: Steps to reproduce the behavior
+      placeholder: |
+        1. Configure '...'
+        2. Call '...'
+        3. See error
+    validations:
+      required: true
+
+  - type: textarea
+    id: expected
+    attributes:
+      label: Expected Behavior
+      description: What did you expect to happen?
+    validations:
+      required: true
+
+  - type: textarea
+    id: actual
+    attributes:
+      label: Actual Behavior
+      description: What actually happened?
+    validations:
+      required: true
+
+  - type: input
+    id: version
+    attributes:
+      label: Version
+      description: What version are you using?
+      placeholder: e.g., 0.1.0 or commit hash
+    validations:
+      required: false
+
+  - type: dropdown
+    id: llm-provider
+    attributes:
+      label: LLM Provider
+      description: Which LLM provider are you using?
+      options:
+        - OpenAI
+        - Anthropic
+        - Gemini
+        - Groq
+        - Ollama
+        - LM Studio
+        - Other
+    validations:
+      required: false
+

.github/ISSUE_TEMPLATE/config.yml

@@ -0,0 +1,8 @@
+blank_issues_enabled: false
+contact_links:
+  - name: Questions & Help
+    url: https://github.com/vectorize-io/hindsight/discussions/categories/q-a
+    about: Please ask questions and get help in Discussions instead of opening an issue.
+  - name: Ideas & Feedback
+    url: https://github.com/vectorize-io/hindsight/discussions/categories/ideas
+    about: Share ideas or give feedback in Discussions.

.github/ISSUE_TEMPLATE/feature_request.yml

@@ -0,0 +1,82 @@
+name: Feature Request
+description: Suggest a new feature or enhancement
+labels: ["enhancement", "triage"]
+body:
+  - type: markdown
+    attributes:
+      value: |
+        Thanks for suggesting a feature! Please describe what you'd like to see added.
+
+  - type: textarea
+    id: use-case
+    attributes:
+      label: Use Case
+      description: Describe your specific use case. What are you building? What's your goal?
+      placeholder: |
+        I'm building an AI agent that needs to...
+        My application handles...
+    validations:
+      required: true
+
+  - type: textarea
+    id: problem
+    attributes:
+      label: Problem Statement
+      description: What problem are you facing? What's missing or difficult today?
+      placeholder: Currently I have to... which causes...
+    validations:
+      required: true
+
+  - type: textarea
+    id: benefit
+    attributes:
+      label: How This Feature Would Help
+      description: Explain how this feature would improve your workflow or solve your problem
+      placeholder: With this feature, I would be able to...
+    validations:
+      required: true
+
+  - type: textarea
+    id: solution
+    attributes:
+      label: Proposed Solution
+      description: Describe your ideal solution (optional - we may have ideas too!)
+      placeholder: It would be great if Hindsight could...
+    validations:
+      required: false
+
+  - type: textarea
+    id: alternatives
+    attributes:
+      label: Alternatives Considered
+      description: Have you considered any alternative solutions or workarounds?
+    validations:
+      required: false
+
+  - type: dropdown
+    id: priority
+    attributes:
+      label: Priority
+      description: How important is this feature to you?
+      options:
+        - Nice to have
+        - Important - affects my workflow
+        - Critical - blocking my use case
+    validations:
+      required: true
+
+  - type: textarea
+    id: additional
+    attributes:
+      label: Additional Context
+      description: Any other context, mockups, or examples?
+    validations:
+      required: false
+
+  - type: checkboxes
+    id: checklist
+    attributes:
+      label: Checklist
+      options:
+        - label: I would be willing to contribute this feature
+          required: false

.github/workflows/release.yml

@@ -42,6 +42,10 @@ jobs:
       working-directory: ./hindsight-integrations/litellm
       run: uv build --out-dir dist
 
+    - name: Build hindsight-embed
+      working-directory: ./hindsight-embed
+      run: uv build --out-dir dist
+
     # Publish in order (client and api first, then hindsight-all which depends on them)
     - name: Publish hindsight-client to PyPI
       uses: pypa/gh-action-pypi-publish@release/v1
@@ -67,6 +71,12 @@ jobs:
         packages-dir: ./hindsight-integrations/litellm/dist
         skip-existing: true
 
+    - name: Publish hindsight-embed to PyPI
+      uses: pypa/gh-action-pypi-publish@release/v1
+      with:
+        packages-dir: ./hindsight-embed/dist
+        skip-existing: true
+
     # Upload artifacts for GitHub release
     - name: Upload artifacts
       uses: actions/upload-artifact@v4
@@ -77,6 +87,7 @@ jobs:
           hindsight-api/dist/*
           hindsight/dist/*
           hindsight-integrations/litellm/dist/*
+          hindsight-embed/dist/*
         retention-days: 1
 
   release-typescript-client:
@@ -416,6 +427,7 @@ jobs:
         cp artifacts/python-packages/hindsight-api/dist/* release-assets/ || true
         cp artifacts/python-packages/hindsight/dist/* release-assets/ || true
         cp artifacts/python-packages/hindsight-integrations/litellm/dist/* release-assets/ || true
+        cp artifacts/python-packages/hindsight-embed/dist/* release-assets/ || true
         # TypeScript client
         cp artifacts/typescript-client/*.tgz release-assets/ || true
         # Control Plane

.github/workflows/test.yml

@@ -20,6 +20,8 @@ jobs:
             path: hindsight-api
           - name: hindsight-client
             path: hindsight-clients/python
+          - name: hindsight-embed
+            path: hindsight-embed
 
     steps:
     - uses: actions/checkout@v4
@@ -173,6 +175,90 @@ jobs:
       working-directory: hindsight-cli
       run: cargo build --release
 
+    - name: Upload CLI artifact
+      uses: actions/upload-artifact@v4
+      with:
+        name: hindsight-cli
+        path: hindsight-cli/target/release/hindsight
+        retention-days: 1
+
+  test-rust-cli:
+    runs-on: ubuntu-latest
+    needs: build-rust-cli
+    env:
+      HINDSIGHT_API_LLM_PROVIDER: groq
+      HINDSIGHT_API_LLM_API_KEY: ${{ secrets.GROQ_API_KEY }}
+      HINDSIGHT_API_LLM_MODEL: openai/gpt-oss-20b
+      HINDSIGHT_API_URL: http://localhost:8888
+      GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+      UV_INDEX: pytorch=https://download.pytorch.org/whl/cpu
+
+    steps:
+    - uses: actions/checkout@v4
+
+    - name: Download CLI artifact
+      uses: actions/download-artifact@v4
+      with:
+        name: hindsight-cli
+        path: /tmp/cli
+
+    - name: Make CLI executable
+      run: chmod +x /tmp/cli/hindsight
+
+    - name: Install uv
+      uses: astral-sh/setup-uv@v5
+      with:
+        enable-cache: true
+        prune-cache: false
+
+    - name: Set up Python
+      uses: actions/setup-python@v5
+      with:
+        python-version-file: ".python-version"
+
+    - name: Build API
+      working-directory: ./hindsight-api
+      run: uv build
+
+    - name: Install API dependencies
+      working-directory: ./hindsight-api
+      run: uv sync --no-install-project --index-strategy unsafe-best-match
+
+    - name: Create .env file
+      run: |
+        cat > .env << EOF
+        HINDSIGHT_API_LLM_PROVIDER=${{ env.HINDSIGHT_API_LLM_PROVIDER }}
+        HINDSIGHT_API_LLM_API_KEY=${{ env.HINDSIGHT_API_LLM_API_KEY }}
+        HINDSIGHT_API_LLM_MODEL=${{ env.HINDSIGHT_API_LLM_MODEL }}
+        EOF
+
+    - name: Start API server
+      run: |
+        ./scripts/dev/start-api.sh > /tmp/api-server.log 2>&1 &
+        echo "Waiting for API server to be ready..."
+        for i in {1..60}; do
+          if curl -sf http://localhost:8888/health > /dev/null 2>&1; then
+            echo "API server is ready after ${i}s"
+            break
+          fi
+          if [ $i -eq 60 ]; then
+            echo "API server failed to start after 60s"
+            cat /tmp/api-server.log
+            exit 1
+          fi
+          sleep 1
+        done
+
+    - name: Run CLI smoke test
+      run: |
+        HINDSIGHT_CLI=/tmp/cli/hindsight ./hindsight-cli/smoke-test.sh
+
+    - name: Show API server logs
+      if: always()
+      run: |
+        echo "=== API Server Logs ==="
+        cat /tmp/api-server.log || echo "No API server log found"
+
   lint-helm-chart:
     runs-on: ubuntu-latest
 
@@ -239,6 +325,8 @@ jobs:
       GROQ_API_KEY: ${{ secrets.GROQ_API_KEY }}
       GEMINI_API_KEY: ${{ secrets.GEMINI_API_KEY }}
       OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
+      COHERE_API_KEY: ${{ secrets.COHERE_API_KEY }}
+      HINDSIGHT_API_EMBEDDINGS_OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
       HINDSIGHT_API_LLM_MODEL: openai/gpt-oss-20b
       GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
       # Prefer CPU-only PyTorch in CI (but keep PyPI for everything else)
@@ -527,6 +615,97 @@ jobs:
         echo "=== API Server Logs ==="
         cat /tmp/api-server.log || echo "No API server log found"
 
+  test-integration:
+    runs-on: ubuntu-latest
+    env:
+      HINDSIGHT_API_LLM_PROVIDER: groq
+      HINDSIGHT_API_LLM_API_KEY: ${{ secrets.GROQ_API_KEY }}
+      HINDSIGHT_API_LLM_MODEL: openai/gpt-oss-20b
+      HINDSIGHT_API_URL: http://localhost:8888
+      GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+      UV_INDEX: pytorch=https://download.pytorch.org/whl/cpu
+
+    steps:
+    - uses: actions/checkout@v4
+
+    - name: Install uv
+      uses: astral-sh/setup-uv@v5
+      with:
+        enable-cache: true
+        prune-cache: false
+
+    - name: Set up Python
+      uses: actions/setup-python@v5
+      with:
+        python-version-file: ".python-version"
+
+    - name: Build API
+      working-directory: ./hindsight-api
+      run: uv build
+
+    - name: Install API dependencies
+      working-directory: ./hindsight-api
+      run: uv sync --no-install-project --index-strategy unsafe-best-match
+
+    - name: Install integration test dependencies
+      working-directory: ./hindsight-integration-tests
+      run: uv sync
+
+    - name: Cache HuggingFace models
+      uses: actions/cache@v4
+      with:
+        path: ~/.cache/huggingface
+        key: ${{ runner.os }}-huggingface-${{ hashFiles('hindsight-api/pyproject.toml') }}
+        restore-keys: |
+          ${{ runner.os }}-huggingface-
+
+    - name: Pre-download models
+      working-directory: ./hindsight-api
+      run: |
+        uv run python -c "
+        from sentence_transformers import SentenceTransformer, CrossEncoder
+        print('Downloading embedding model...')
+        SentenceTransformer('BAAI/bge-small-en-v1.5')
+        print('Downloading cross-encoder model...')
+        CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2')
+        print('Models downloaded successfully')
+        "
+
+    - name: Create .env file
+      run: |
+        cat > .env << EOF
+        HINDSIGHT_API_LLM_PROVIDER=${{ env.HINDSIGHT_API_LLM_PROVIDER }}
+        HINDSIGHT_API_LLM_API_KEY=${{ env.HINDSIGHT_API_LLM_API_KEY }}
+        HINDSIGHT_API_LLM_MODEL=${{ env.HINDSIGHT_API_LLM_MODEL }}
+        EOF
+
+    - name: Start API server
+      run: |
+        ./scripts/dev/start-api.sh > /tmp/api-server.log 2>&1 &
+        echo "Waiting for API server to be ready..."
+        for i in {1..60}; do
+          if curl -sf http://localhost:8888/health > /dev/null 2>&1; then
+            echo "API server is ready after ${i}s"
+            break
+          fi
+          if [ $i -eq 60 ]; then
+            echo "API server failed to start after 60s"
+            cat /tmp/api-server.log
+            exit 1
+          fi
+          sleep 1
+        done
+
+    - name: Run integration tests
+      working-directory: ./hindsight-integration-tests
+      run: uv run pytest tests/ -v
+
+    - name: Show API server logs
+      if: always()
+      run: |
+        echo "=== API Server Logs ==="
+        cat /tmp/api-server.log || echo "No API server log found"
+
   test-litellm-integration:
     runs-on: ubuntu-latest
 
@@ -556,8 +735,49 @@ jobs:
       working-directory: ./hindsight-integrations/litellm
       run: uv run pytest tests -v
 
+  test-embed:
+    runs-on: ubuntu-latest
+    env:
+      HINDSIGHT_EMBED_LLM_PROVIDER: groq
+      HINDSIGHT_EMBED_LLM_API_KEY: ${{ secrets.GROQ_API_KEY }}
+      HINDSIGHT_EMBED_LLM_MODEL: openai/gpt-oss-20b
+      # Prefer CPU-only PyTorch in CI
+      UV_INDEX: pytorch=https://download.pytorch.org/whl/cpu
+
+    steps:
+    - uses: actions/checkout@v4
+
+    - name: Install uv
+      uses: astral-sh/setup-uv@v5
+      with:
+        enable-cache: true
+        prune-cache: false
+
+    - name: Set up Python
+      uses: actions/setup-python@v5
+      with:
+        python-version-file: ".python-version"
+
+    - name: Install dependencies
+      working-directory: ./hindsight-embed
+      run: uv sync --index-strategy unsafe-best-match
+
+    - name: Cache HuggingFace models
+      uses: actions/cache@v4
+      with:
+        path: ~/.cache/huggingface
+        key: ${{ runner.os }}-huggingface-embed-${{ hashFiles('hindsight-embed/pyproject.toml') }}
+        restore-keys: |
+          ${{ runner.os }}-huggingface-embed-
+          ${{ runner.os }}-huggingface-
+
+    - name: Run smoke test
+      working-directory: ./hindsight-embed
+      run: ./test.sh
+
   test-doc-examples:
     runs-on: ubuntu-latest
+    needs: build-rust-cli
     env:
       HINDSIGHT_API_LLM_PROVIDER: groq
       HINDSIGHT_API_LLM_API_KEY: ${{ secrets.GROQ_API_KEY }}
@@ -569,6 +789,15 @@ jobs:
     steps:
     - uses: actions/checkout@v4
 
+    - name: Download CLI artifact
+      uses: actions/download-artifact@v4
+      with:
+        name: hindsight-cli
+        path: /usr/local/bin
+
+    - name: Make CLI executable
+      run: chmod +x /usr/local/bin/hindsight
+
     - name: Install uv
       uses: astral-sh/setup-uv@v5
       with:
@@ -642,8 +871,90 @@ jobs:
           node "$f"
         done
 
+    - name: Configure CLI
+      run: hindsight configure --api-url http://localhost:8888
+
+    - name: Run CLI doc examples
+      run: |
+        for f in hindsight-docs/examples/api/*.sh; do
+          echo "Running $f..."
+          bash "$f"
+        done
+
     - name: Show API server logs
       if: always()
       run: |
         echo "=== API Server Logs ==="
-        cat /tmp/api-server.log || echo "No API server log found"
+        cat /tmp/api-server.log || echo "No API server log found"
+
+  verify-generated-files:
+    runs-on: ubuntu-latest
+    env:
+      UV_INDEX: pytorch=https://download.pytorch.org/whl/cpu
+
+    steps:
+    - uses: actions/checkout@v4
+
+    - name: Install uv
+      uses: astral-sh/setup-uv@v5
+      with:
+        enable-cache: true
+
+    - name: Set up Python
+      uses: actions/setup-python@v5
+      with:
+        python-version-file: ".python-version"
+
+    - name: Set up Node.js
+      uses: actions/setup-node@v4
+      with:
+        node-version: '20'
+        cache: 'npm'
+        cache-dependency-path: package-lock.json
+
+    - name: Install Rust
+      uses: dtolnay/rust-toolchain@stable
+
+    - name: Cache cargo
+      uses: actions/cache@v4
+      with:
+        path: |
+          ~/.cargo/registry
+          ~/.cargo/git
+        key: ${{ runner.os }}-cargo-gen-${{ hashFiles('**/Cargo.lock') }}
+
+    - name: Install Node dependencies
+      run: npm ci
+
+    - name: Install Python dependencies
+      run: |
+        cd hindsight-dev && uv sync --index-strategy unsafe-best-match
+        cd ../hindsight-api && uv sync --index-strategy unsafe-best-match
+        cd ../hindsight-embed && uv sync --index-strategy unsafe-best-match
+
+    - name: Run generate-openapi
+      run: ./scripts/generate-openapi.sh
+
+    - name: Run generate-clients
+      run: ./scripts/generate-clients.sh
+
+    - name: Run lint
+      run: ./scripts/hooks/lint.sh
+
+    - name: Verify no uncommitted changes
+      run: |
+        if [ -n "$(git status --porcelain)" ]; then
+          echo "❌ Error: Generated files are out of sync with committed files."
+          echo ""
+          echo "The following files have changed after running generation scripts:"
+          git status --porcelain
+          echo ""
+          echo "Please run the following commands locally and commit the changes:"
+          echo "  ./scripts/generate-openapi.sh"
+          echo "  ./scripts/generate-clients.sh"
+          echo "  ./scripts/hooks/lint.sh"
+          echo ""
+          git diff --stat
+          exit 1
+        fi
+        echo "✓ All generated files are up to date"

.gitignore

@@ -5,15 +5,18 @@ build/
 dist/
 wheels/
 *.egg-info
-
+.mcp.json
+.osgrep
 # Virtual environments
 .venv
 
 # Node
 node_modules/
 
-# Environment variables
+# Environment variables and local config
 .env
+docker-compose.yml
+docker-compose.override.yml
 
 # IDE
 .idea/
@@ -39,4 +42,8 @@ hindsight-docs/static/llms-full.txt
 hindsight-dev/benchmarks/locomo/results/
 hindsight-dev/benchmarks/longmemeval/results/
 hindsight-cli/target
-hindsight-clients/rust/target
+hindsight-clients/rust/target
+.claude
+whats-next.md
+TASK.md
+CHANGELOG.md

AGENTS.md

@@ -14,6 +14,7 @@ This document captures architectural decisions and coding conventions for the Hi
 hindsight/              # Python package for embedded usage
 hindsight-api/          # FastAPI server (core memory engine)
 hindsight-cli/          # Rust CLI client
+hindsight-embed/        # Embedded CLI (no server needed)
 hindsight-control-plane/ # Next.js admin UI
 hindsight-docs/         # Docusaurus documentation site
 hindsight-dev/          # Development tools and benchmarks
@@ -148,4 +149,5 @@ Note: The maintained wrapper `hindsight_client.py` and `README.md` are preserved
 
 # Branding
 ## Colors
-- Primary: gradient from #0074d9 to #009296  
+- Primary: gradient from #0074d9 to #009296
+

CLAUDE.md

@@ -0,0 +1,186 @@
+# CLAUDE.md
+
+This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.
+
+## Project Overview
+
+Hindsight is an agent memory system that provides long-term memory for AI agents using biomimetic data structures. Memories are organized as:
+- **World facts**: General knowledge ("The sky is blue")
+- **Experience facts**: Personal experiences ("I visited Paris in 2023")
+- **Opinion facts**: Beliefs with confidence scores ("Paris is beautiful" - 0.9 confidence)
+- **Observations**: Complex mental models derived from reflection
+
+## Development Commands
+
+### API Server (Python/FastAPI)
+```bash
+# Start API server (loads .env automatically)
+./scripts/dev/start-api.sh
+
+# Run all tests (parallelized with pytest-xdist)
+cd hindsight-api && uv run pytest tests/
+
+# Run specific test file
+cd hindsight-api && uv run pytest tests/test_http_api_integration.py -v
+
+# Run single test function
+cd hindsight-api && uv run pytest tests/test_retain.py::test_retain_simple -v
+
+# Lint and format
+cd hindsight-api && uv run ruff check .
+cd hindsight-api && uv run ruff format .
+
+# Type checking (uses ty - extremely fast type checker from Astral)
+cd hindsight-api && uv run ty check hindsight_api/
+```
+
+### Control Plane (Next.js)
+```bash
+./scripts/dev/start-control-plane.sh
+# Or manually:
+cd hindsight-control-plane && npm run dev
+```
+
+### Documentation Site (Docusaurus)
+```bash
+./scripts/dev/start-docs.sh
+```
+
+### Generating Clients/OpenAPI
+```bash
+# Regenerate OpenAPI spec after API changes (REQUIRED after changing endpoints)
+./scripts/generate-openapi.sh
+
+# Regenerate all client SDKs (Python, TypeScript, Rust)
+./scripts/generate-clients.sh
+```
+
+### Benchmarks
+```bash
+./scripts/benchmarks/run-longmemeval.sh
+./scripts/benchmarks/run-locomo.sh
+./scripts/benchmarks/start-visualizer.sh  # View results at localhost:8001
+```
+
+## Architecture
+
+### Monorepo Structure
+- **hindsight-api/**: Core FastAPI server with memory engine (Python, uv)
+- **hindsight/**: Embedded Python bundle (hindsight-all package)
+- **hindsight-control-plane/**: Admin UI (Next.js, npm)
+- **hindsight-cli/**: CLI tool (Rust, cargo, uses progenitor for API client)
+- **hindsight-clients/**: Generated SDK clients (Python, TypeScript, Rust)
+- **hindsight-docs/**: Docusaurus documentation site
+- **hindsight-integrations/**: Framework integrations (LiteLLM, OpenAI)
+- **hindsight-dev/**: Development tools and benchmarks
+
+### Core Engine (hindsight-api/hindsight_api/engine/)
+- `memory_engine.py`: Main orchestrator (~170KB) for retain/recall/reflect operations
+- `llm_wrapper.py`: LLM abstraction supporting OpenAI, Anthropic, Gemini, Groq, Ollama, LM Studio
+- `embeddings.py`: Embedding generation (local sentence-transformers or TEI)
+- `cross_encoder.py`: Reranking (local or TEI)
+- `entity_resolver.py`: Entity extraction and normalization
+- `query_analyzer.py`: Query intent analysis
+
+**retain/**: Memory ingestion pipeline
+- `orchestrator.py`: Coordinates the retain flow
+- `fact_extraction.py`: LLM-based fact extraction from content
+- `link_utils.py`: Entity link creation and management
+
+**search/**: Multi-strategy retrieval
+- `retrieval.py`: Main retrieval orchestrator
+- `graph_retrieval.py`: Entity/relationship graph traversal
+- `mpfp_retrieval.py`: Multi-Path Fact Propagation retrieval
+- `fusion.py`: Reciprocal rank fusion for combining results
+- `reranking.py`: Cross-encoder reranking
+
+### API Layer (hindsight-api/hindsight_api/api/)
+- `http.py`: FastAPI HTTP routers (~80KB) for all REST endpoints
+- `mcp.py`: Model Context Protocol server implementation
+
+Main operations:
+- **Retain**: Store memories, extracts facts/entities/relationships
+- **Recall**: Retrieve memories via 4 parallel strategies (semantic, BM25, graph, temporal) + reranking
+- **Reflect**: Deep analysis forming new opinions/observations (disposition-aware)
+
+### Database
+PostgreSQL with pgvector. Schema managed via Alembic migrations in `hindsight-api/hindsight_api/alembic/`. Migrations run automatically on API startup.
+
+Key tables: `banks`, `memory_units`, `documents`, `entities`, `entity_links`
+
+## 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 an isolated memory store (like a "brain" for one user/agent)
+- Banks have dispositions (skepticism, literalism, empathy traits 1-5) affecting reflect
+- Banks can have background context
+- Bank isolation is strict - no cross-bank data leakage
+
+### API Design
+- All endpoints operate on a single bank per request
+- Multi-bank queries are client responsibility to orchestrate
+- Disposition traits only affect reflect, not recall
+
+### Python Style
+- Python 3.11+, type hints required
+- Async throughout (asyncpg, async FastAPI)
+- Pydantic models for request/response
+- Ruff for linting (line-length 120)
+- No Python files at project root - maintain clean directory structure
+
+### TypeScript Style
+- 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
+cp .env.example .env
+# Edit .env with LLM API key
+
+# Python deps
+uv sync --directory hindsight-api/
+
+# Node deps (uses npm workspaces)
+npm install
+```
+
+Required env vars:
+- `HINDSIGHT_API_LLM_PROVIDER`: openai, anthropic, gemini, groq, ollama, lmstudio
+- `HINDSIGHT_API_LLM_API_KEY`: Your API key
+- `HINDSIGHT_API_LLM_MODEL`: Model name (e.g., o3-mini, claude-sonnet-4-20250514)
+
+Optional (uses local models by default):
+- `HINDSIGHT_API_EMBEDDINGS_PROVIDER`: local (default) or tei
+- `HINDSIGHT_API_RERANKER_PROVIDER`: local (default) or tei
+- `HINDSIGHT_API_DATABASE_URL`: External PostgreSQL (uses embedded pg0 by default)

CONTRIBUTING.md

@@ -51,7 +51,36 @@ cd hindsight-api
 uv run pytest tests/
 ```
 
-### Code style
+### Code Style
+
+We use [Ruff](https://docs.astral.sh/ruff/) for Python linting and formatting, and ESLint/Prettier for TypeScript.
+
+#### Setting up git hooks (recommended)
+
+Set up git hooks to automatically lint and format code before each commit:
+
+```bash
+./scripts/setup-hooks.sh
+```
+
+This configures git to use the hooks in `.githooks/`, which run all scripts in `scripts/hooks/` on commit. The lint hook runs in parallel:
+- **Python**: `ruff check --fix`, `ruff format`, `ty check`
+- **TypeScript**: `eslint --fix`, `prettier`
+
+#### Manual linting and formatting
+
+```bash
+# Run all lints (same as pre-commit)
+./scripts/hooks/lint.sh
+
+# Or run individually for Python:
+cd hindsight-api
+uv run ruff check --fix .   # Lint and auto-fix
+uv run ruff format .        # Format code
+uv run ty check hindsight_api  # Type check
+```
+
+#### Style guidelines
 
 - Use Python type hints
 - Follow existing code patterns

README.md

@@ -81,6 +81,8 @@ docker run --rm -it --pull always -p 8888:8888 -p 9999:9999 \
   ghcr.io/vectorize-io/hindsight:latest
 ```
 
+You can modify the LLM provider by setting `HINDSIGHT_API_LLM_PROVIDER`. Valid options are `openai`, `anthropic`, `gemini`, `groq`, `ollama`, and `lmstudio`. The documentation provides more details on [supported models](https://hindsight.vectorize.io/developer/models).
+
 API: http://localhost:8888  
 UI: http://localhost:9999
 

docker/standalone/Dockerfile

@@ -2,16 +2,24 @@
 # Supports building API-only, Control Plane-only, or both
 #
 # Build args:
-#   INCLUDE_API=true/false  - Include API (default: true)
-#   INCLUDE_CP=true/false   - Include Control Plane (default: true)
+#   INCLUDE_API=true/false         - Include API (default: true)
+#   INCLUDE_CP=true/false          - Include Control Plane (default: true)
+#   INCLUDE_LOCAL_MODELS=true/false - Include local ML models for embeddings/reranking (default: true)
+#                                     Set to false when using external providers (TEI, OpenAI, Cohere)
+#   PRELOAD_ML_MODELS=true/false   - Pre-download ML models during build (default: true)
+#                                     Only effective when INCLUDE_LOCAL_MODELS=true
 #
 # Examples:
-#   docker build -t hindsight .                                    # Both (standalone)
-#   docker build -t hindsight-api --build-arg INCLUDE_CP=false .   # API only
-#   docker build -t hindsight-cp --build-arg INCLUDE_API=false .   # Control Plane only
+#   docker build -t hindsight .                                          # Both (standalone)
+#   docker build -t hindsight-api --build-arg INCLUDE_CP=false .         # API only
+#   docker build -t hindsight-cp --build-arg INCLUDE_API=false .         # Control Plane only
+#   docker build -t hindsight --build-arg PRELOAD_ML_MODELS=false .      # Skip ML model preload
+#   docker build -t hindsight --build-arg INCLUDE_LOCAL_MODELS=false .   # Skip local ML deps (for external providers)
 
 ARG INCLUDE_API=true
 ARG INCLUDE_CP=true
+ARG PRELOAD_ML_MODELS=true
+ARG INCLUDE_LOCAL_MODELS=true
 
 # =============================================================================
 # Stage: API Builder
@@ -19,6 +27,7 @@ ARG INCLUDE_CP=true
 FROM python:3.11-slim AS api-builder
 
 ARG INCLUDE_API
+ARG INCLUDE_LOCAL_MODELS
 RUN if [ "$INCLUDE_API" != "true" ]; then echo "Skipping API build" && exit 0; fi
 
 WORKDIR /app
@@ -37,6 +46,15 @@ COPY hindsight-api/README.md ./api/
 
 WORKDIR /app/api
 
+# Remove local ML model dependencies if INCLUDE_LOCAL_MODELS=false
+# This creates a smaller image when using external providers (TEI, OpenAI, Cohere)
+RUN if [ "$INCLUDE_LOCAL_MODELS" != "true" ]; then \
+        echo "Removing local-models dependencies (sentence-transformers, torch, transformers)..." && \
+        sed -i '/"sentence-transformers/d' pyproject.toml && \
+        sed -i '/"transformers/d' pyproject.toml && \
+        sed -i '/"torch/d' pyproject.toml; \
+    fi
+
 # Sync dependencies (will create lock file if needed)
 RUN uv sync
 
@@ -122,7 +140,6 @@ FROM python:3.11-slim AS api-only
 
 WORKDIR /app
 
-# Install pg0 dependencies (procps provides 'kill' command needed by pg0)
 # Note: libicu version varies by Debian version - try common versions in order
 RUN apt-get update && apt-get install -y \
     curl \
@@ -135,7 +152,6 @@ RUN apt-get update && apt-get install -y \
     && rm -rf /var/lib/apt/lists/* \
     && pip install --no-cache-dir uv
 
-# Create non-root user (PostgreSQL cannot run as root)
 RUN useradd -m -s /bin/bash hindsight
 
 # Copy API with virtual environment from builder
@@ -145,28 +161,26 @@ COPY --from=api-builder /app/api /app/api
 COPY docker/standalone/start-all.sh /app/start-all.sh
 RUN chmod +x /app/start-all.sh
 
-# Create data directory for pg0 and set ownership
-RUN mkdir -p /app/data && chown -R hindsight:hindsight /app
+RUN chown -R hindsight:hindsight /app
 
-# Switch to non-root user
 USER hindsight
 
-# Set PATH for hindsight user
 ENV PATH="/app/api/.venv/bin:${PATH}"
 
-# Pre-cache PostgreSQL binaries by starting/stopping pg0-embedded
-ENV PG0_HOME=/home/hindsight/.pg0-cache
-
-ENV PG0_HOME=/home/hindsight/.pg0
-
-# Pre-download ML models to avoid runtime download
-RUN /app/api/.venv/bin/python -c "\
+# Pre-download ML models to avoid runtime download (conditional)
+# Only runs if both PRELOAD_ML_MODELS=true AND INCLUDE_LOCAL_MODELS=true
+ARG PRELOAD_ML_MODELS
+ARG INCLUDE_LOCAL_MODELS
+RUN if [ "$PRELOAD_ML_MODELS" = "true" ] && [ "$INCLUDE_LOCAL_MODELS" = "true" ]; then \
+    /app/api/.venv/bin/python -c "\
 from sentence_transformers import SentenceTransformer, CrossEncoder; \
 print('Downloading embedding model...'); \
 SentenceTransformer('BAAI/bge-small-en-v1.5'); \
 print('Downloading cross-encoder model...'); \
 CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2'); \
-print('Models cached successfully')"
+print('Models cached successfully')"; \
+    elif [ "$INCLUDE_LOCAL_MODELS" != "true" ]; then echo "Skipping ML model preload (local-models not included)"; \
+    else echo "Skipping ML model preload"; fi
 
 EXPOSE 8888
 
@@ -220,7 +234,7 @@ FROM python:3.11-slim AS standalone
 
 WORKDIR /app
 
-# Install Node.js, curl, uv, and pg0 dependencies (procps provides 'kill' command needed by pg0)
+# Install Node.js, curl, uv, and system dependencies
 # Note: libicu version varies by Debian version - try common versions in order
 RUN apt-get update && apt-get install -y \
     curl \
@@ -235,7 +249,6 @@ RUN apt-get update && apt-get install -y \
     && rm -rf /var/lib/apt/lists/* \
     && pip install --no-cache-dir uv
 
-# Create non-root user (PostgreSQL cannot run as root)
 RUN useradd -m -s /bin/bash hindsight
 
 # Copy API with virtual environment from builder
@@ -256,35 +269,26 @@ WORKDIR /app
 COPY docker/standalone/start-all.sh /app/start-all.sh
 RUN chmod +x /app/start-all.sh
 
-# Create data directory for pg0 and set ownership
-RUN mkdir -p /app/data && chown -R hindsight:hindsight /app
+RUN chown -R hindsight:hindsight /app
 
-# Switch to non-root user
 USER hindsight
 
-# Set PATH for hindsight user
 ENV PATH="/app/api/.venv/bin:${PATH}"
 
-# Pre-cache PostgreSQL binaries by starting/stopping pg0-embedded
-ENV PG0_HOME=/home/hindsight/.pg0-cache
-RUN /app/api/.venv/bin/python -c "\
-from pg0 import Pg0; \
-print('Pre-caching PostgreSQL binaries...'); \
-pg = Pg0(name='hindsight', port=5555, username='hindsight', password='hindsight', database='hindsight'); \
-pg.start(); \
-pg.stop(); \
-print('PostgreSQL pre-cached to PG0_HOME')" || echo "Pre-download skipped"
-
-ENV PG0_HOME=/home/hindsight/.pg0
-
-# Pre-download ML models to avoid runtime download
-RUN /app/api/.venv/bin/python -c "\
+# Pre-download ML models to avoid runtime download (conditional)
+# Only runs if both PRELOAD_ML_MODELS=true AND INCLUDE_LOCAL_MODELS=true
+ARG PRELOAD_ML_MODELS
+ARG INCLUDE_LOCAL_MODELS
+RUN if [ "$PRELOAD_ML_MODELS" = "true" ] && [ "$INCLUDE_LOCAL_MODELS" = "true" ]; then \
+    /app/api/.venv/bin/python -c "\
 from sentence_transformers import SentenceTransformer, CrossEncoder; \
 print('Downloading embedding model...'); \
 SentenceTransformer('BAAI/bge-small-en-v1.5'); \
 print('Downloading cross-encoder model...'); \
 CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2'); \
-print('Models cached successfully')"
+print('Models cached successfully')"; \
+    elif [ "$INCLUDE_LOCAL_MODELS" != "true" ]; then echo "Skipping ML model preload (local-models not included)"; \
+    else echo "Skipping ML model preload"; fi
 
 EXPOSE 8888 9999
 

docker/standalone/start-all.sh

@@ -5,16 +5,70 @@ set -e
 ENABLE_API="${HINDSIGHT_ENABLE_API:-true}"
 ENABLE_CP="${HINDSIGHT_ENABLE_CP:-true}"
 
-# Copy pre-cached PostgreSQL data if runtime directory is empty (first run with volume)
-if [ "$ENABLE_API" = "true" ]; then
-    PG0_CACHE="/home/hindsight/.pg0-cache"
-    PG0_HOME="/home/hindsight/.pg0"
-    if [ -d "$PG0_CACHE" ] && [ "$(ls -A $PG0_CACHE 2>/dev/null)" ]; then
-        if [ ! "$(ls -A $PG0_HOME 2>/dev/null)" ]; then
-            echo "📦 Copying pre-cached PostgreSQL data..."
-            cp -r "$PG0_CACHE"/* "$PG0_HOME"/ 2>/dev/null || true
-        fi
+# =============================================================================
+# Dependency waiting (opt-in via HINDSIGHT_WAIT_FOR_DEPS=true)
+#
+# Problem: When running with LM Studio, the LLM may take time to load models.
+# If Hindsight starts before LM Studio is ready, it fails on LLM verification.
+# This wait loop ensures dependencies are ready before starting.
+# =============================================================================
+if [ "${HINDSIGHT_WAIT_FOR_DEPS:-false}" = "true" ]; then
+    LLM_BASE_URL="${HINDSIGHT_API_LLM_BASE_URL:-http://host.docker.internal:1234/v1}"
+    MAX_RETRIES="${HINDSIGHT_RETRY_MAX:-0}"  # 0 = infinite
+    RETRY_INTERVAL="${HINDSIGHT_RETRY_INTERVAL:-10}"
+
+    # Check if external database is configured (skip check for embedded pg0)
+    SKIP_DB_CHECK=false
+    if [ -z "${HINDSIGHT_API_DATABASE_URL}" ]; then
+        SKIP_DB_CHECK=true
+    else
+        DB_CHECK_HOST=$(echo "$HINDSIGHT_API_DATABASE_URL" | sed -E 's|.*@([^:/]+):([0-9]+)/.*|\1 \2|')
     fi
+
+    check_db() {
+        if $SKIP_DB_CHECK; then
+            return 0
+        fi
+        if command -v pg_isready &> /dev/null; then
+            pg_isready -h $(echo $DB_CHECK_HOST | cut -d' ' -f1) -p $(echo $DB_CHECK_HOST | cut -d' ' -f2) &>/dev/null
+        else
+            python3 -c "import socket; s=socket.socket(); s.settimeout(5); exit(0 if s.connect_ex(('$(echo $DB_CHECK_HOST | cut -d' ' -f1)', $(echo $DB_CHECK_HOST | cut -d' ' -f2))) == 0 else 1)" 2>/dev/null
+        fi
+    }
+
+    check_llm() {
+        curl -sf "${LLM_BASE_URL}/models" --connect-timeout 5 &>/dev/null
+    }
+
+    echo "⏳ Waiting for dependencies to be ready..."
+    attempt=1
+
+    while true; do
+        db_ok=false
+        llm_ok=false
+
+        if check_db; then
+            db_ok=true
+        fi
+
+        if check_llm; then
+            llm_ok=true
+        fi
+
+        if $db_ok && $llm_ok; then
+            echo "✅ Dependencies ready!"
+            break
+        fi
+
+        if [ "$MAX_RETRIES" -ne 0 ] && [ "$attempt" -ge "$MAX_RETRIES" ]; then
+            echo "❌ Max retries ($MAX_RETRIES) reached. Dependencies not available."
+            exit 1
+        fi
+
+        echo "   Attempt $attempt: DB=$( $db_ok && echo 'ok' || echo 'waiting' ), LLM=$( $llm_ok && echo 'ok' || echo 'waiting' )"
+        sleep "$RETRY_INTERVAL"
+        ((attempt++))
+    done
 fi
 
 # Track PIDs for wait

helm/hindsight/Chart.yaml

@@ -2,8 +2,8 @@ apiVersion: v2
 name: hindsight
 description: Hindsight helm chart
 type: application
-version: 0.1.11
-appVersion: "0.1.11"
+version: 0.2.1
+appVersion: "0.2.1"
 keywords:
   - ai
   - memory

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/README.md

@@ -80,7 +80,7 @@ Configure via environment variables:
 | Variable | Description | Default |
 |----------|-------------|---------|
 | `HINDSIGHT_API_DATABASE_URL` | PostgreSQL connection string | `pg0` (embedded) |
-| `HINDSIGHT_API_LLM_PROVIDER` | `openai`, `groq`, `gemini`, `ollama` | `openai` |
+| `HINDSIGHT_API_LLM_PROVIDER` | `openai`, `anthropic`, `gemini`, `groq`, `ollama`, `lmstudio` | `openai` |
 | `HINDSIGHT_API_LLM_API_KEY` | API key for LLM provider | - |
 | `HINDSIGHT_API_LLM_MODEL` | Model name | `gpt-4o-mini` |
 | `HINDSIGHT_API_HOST` | Server bind address | `0.0.0.0` |

hindsight-api/hindsight_api/__init__.py

@@ -21,9 +21,11 @@ from .engine.search.trace import (
     WeightComponents,
 )
 from .engine.search.tracer import SearchTracer
+from .models import RequestContext
 
 __all__ = [
     "MemoryEngine",
+    "RequestContext",
     "HindsightConfig",
     "get_config",
     "SearchTrace",

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/alembic/env.py

@@ -109,6 +109,9 @@ def run_migrations_online() -> None:
 
     get_database_url()  # Process and set the database URL in config
 
+    # Check if we're targeting a specific schema (for multi-tenant isolation)
+    target_schema = config.get_main_option("target_schema")
+
     connectable = engine_from_config(
         config.get_section(config.config_ini_section, {}),
         prefix="sqlalchemy.",
@@ -121,14 +124,34 @@ def run_migrations_online() -> None:
     def set_read_write_mode(dbapi_connection, connection_record):
         cursor = dbapi_connection.cursor()
         cursor.execute("SET SESSION CHARACTERISTICS AS TRANSACTION READ WRITE")
+        # If targeting a specific schema, set search_path
+        # Include public in search_path for access to shared extensions (pgvector)
+        if target_schema:
+            cursor.execute(f'CREATE SCHEMA IF NOT EXISTS "{target_schema}"')
+            cursor.execute(f'SET search_path TO "{target_schema}", public')
         cursor.close()
 
     with connectable.connect() as connection:
         # Also explicitly set read-write mode on this connection
         connection.execute(text("SET SESSION CHARACTERISTICS AS TRANSACTION READ WRITE"))
+
+        # If targeting a specific schema, set search_path
+        # Include public in search_path for access to shared extensions (pgvector)
+        if target_schema:
+            connection.execute(text(f'CREATE SCHEMA IF NOT EXISTS "{target_schema}"'))
+            connection.execute(text(f'SET search_path TO "{target_schema}", public'))
+
         connection.commit()  # Commit the SET command
 
-        context.configure(connection=connection, target_metadata=target_metadata)
+        # Configure context with version_table_schema if using a specific schema
+        context_opts = {
+            "connection": connection,
+            "target_metadata": target_metadata,
+        }
+        if target_schema:
+            context_opts["version_table_schema"] = target_schema
+
+        context.configure(**context_opts)
 
         with context.begin_transaction():
             context.run_migrations()

hindsight-api/hindsight_api/alembic/versions/d9f6a3b4c5e2_rename_bank_to_interactions.py

@@ -6,7 +6,7 @@ Create Date: 2024-12-04 15:00:00.000000
 
 """
 
-from alembic import op
+from alembic import context, op
 
 # revision identifiers, used by Alembic.
 revision = "d9f6a3b4c5e2"
@@ -15,14 +15,22 @@ branch_labels = None
 depends_on = None
 
 
+def _get_schema_prefix() -> str:
+    """Get schema prefix for table names (e.g., 'tenant_x.' or '' for public)."""
+    schema = context.config.get_main_option("target_schema")
+    return f'"{schema}".' if schema else ""
+
+
 def upgrade():
+    schema = _get_schema_prefix()
+
     # Drop old check constraint FIRST (before updating data)
     op.drop_constraint("memory_units_fact_type_check", "memory_units", type_="check")
 
     # Update existing 'bank' values to 'experience'
-    op.execute("UPDATE memory_units SET fact_type = 'experience' WHERE fact_type = 'bank'")
+    op.execute(f"UPDATE {schema}memory_units SET fact_type = 'experience' WHERE fact_type = 'bank'")
     # Also update any 'interactions' values (in case of partial migration)
-    op.execute("UPDATE memory_units SET fact_type = 'experience' WHERE fact_type = 'interactions'")
+    op.execute(f"UPDATE {schema}memory_units SET fact_type = 'experience' WHERE fact_type = 'interactions'")
 
     # Create new check constraint with 'experience' instead of 'bank'
     op.create_check_constraint(
@@ -31,11 +39,13 @@ def upgrade():
 
 
 def downgrade():
+    schema = _get_schema_prefix()
+
     # Drop new check constraint FIRST
     op.drop_constraint("memory_units_fact_type_check", "memory_units", type_="check")
 
     # Update 'experience' back to 'bank'
-    op.execute("UPDATE memory_units SET fact_type = 'bank' WHERE fact_type = 'experience'")
+    op.execute(f"UPDATE {schema}memory_units SET fact_type = 'bank' WHERE fact_type = 'experience'")
 
     # Recreate old check constraint
     op.create_check_constraint(

hindsight-api/hindsight_api/alembic/versions/e0a1b2c3d4e5_disposition_to_3_traits.py

@@ -12,7 +12,7 @@ system (skepticism, literalism, empathy with 1-5 integer values).
 from collections.abc import Sequence
 
 import sqlalchemy as sa
-from alembic import op
+from alembic import context, op
 
 # revision identifiers, used by Alembic.
 revision: str = "e0a1b2c3d4e5"
@@ -21,9 +21,36 @@ branch_labels: str | Sequence[str] | None = None
 depends_on: str | Sequence[str] | None = None
 
 
+def _get_schema_prefix() -> str:
+    """Get schema prefix for table names (e.g., 'tenant_x.' or '' for public)."""
+    schema = context.config.get_main_option("target_schema")
+    return f'"{schema}".' if schema else ""
+
+
+def _get_target_schema() -> str:
+    """Get the target schema name (tenant schema or 'public')."""
+    schema = context.config.get_main_option("target_schema")
+    return schema if schema else "public"
+
+
 def upgrade() -> None:
     """Convert Big Five disposition to 3-trait disposition."""
     conn = op.get_bind()
+    schema = _get_schema_prefix()
+    target_schema = _get_target_schema()
+
+    # Check if disposition column exists (should have been created by previous migration)
+    result = conn.execute(
+        sa.text("""
+        SELECT column_name
+        FROM information_schema.columns
+        WHERE table_schema = :schema AND table_name = 'banks' AND column_name = 'disposition'
+    """),
+        {"schema": target_schema},
+    )
+    if not result.fetchone():
+        # Column doesn't exist yet (shouldn't happen but be safe)
+        return
 
     # Update all existing banks to use the new disposition format
     # Convert from old format to new format with reasonable mappings:
@@ -32,18 +59,18 @@ def upgrade() -> None:
     # - empathy: derived from agreeableness + inverse of neuroticism
     # Default all to 3 (neutral) for simplicity
     conn.execute(
-        sa.text("""
-        UPDATE banks
-        SET disposition = '{"skepticism": 3, "literalism": 3, "empathy": 3}'::jsonb
+        sa.text(f"""
+        UPDATE {schema}banks
+        SET disposition = '{{"skepticism": 3, "literalism": 3, "empathy": 3}}'::jsonb
         WHERE disposition IS NOT NULL
     """)
     )
 
     # Update the default for new banks
     conn.execute(
-        sa.text("""
-        ALTER TABLE banks
-        ALTER COLUMN disposition SET DEFAULT '{"skepticism": 3, "literalism": 3, "empathy": 3}'::jsonb
+        sa.text(f"""
+        ALTER TABLE {schema}banks
+        ALTER COLUMN disposition SET DEFAULT '{{"skepticism": 3, "literalism": 3, "empathy": 3}}'::jsonb
     """)
     )
 
@@ -51,20 +78,34 @@ def upgrade() -> None:
 def downgrade() -> None:
     """Convert back to Big Five disposition."""
     conn = op.get_bind()
+    schema = _get_schema_prefix()
+    target_schema = _get_target_schema()
+
+    # Check if disposition column exists
+    result = conn.execute(
+        sa.text("""
+        SELECT column_name
+        FROM information_schema.columns
+        WHERE table_schema = :schema AND table_name = 'banks' AND column_name = 'disposition'
+    """),
+        {"schema": target_schema},
+    )
+    if not result.fetchone():
+        return
 
     # Revert to Big Five format with default values
     conn.execute(
-        sa.text("""
-        UPDATE banks
-        SET disposition = '{"openness": 0.5, "conscientiousness": 0.5, "extraversion": 0.5, "agreeableness": 0.5, "neuroticism": 0.5, "bias_strength": 0.5}'::jsonb
+        sa.text(f"""
+        UPDATE {schema}banks
+        SET disposition = '{{"openness": 0.5, "conscientiousness": 0.5, "extraversion": 0.5, "agreeableness": 0.5, "neuroticism": 0.5, "bias_strength": 0.5}}'::jsonb
         WHERE disposition IS NOT NULL
     """)
     )
 
     # Update the default for new banks
     conn.execute(
-        sa.text("""
-        ALTER TABLE banks
-        ALTER COLUMN disposition SET DEFAULT '{"openness": 0.5, "conscientiousness": 0.5, "extraversion": 0.5, "agreeableness": 0.5, "neuroticism": 0.5, "bias_strength": 0.5}'::jsonb
+        sa.text(f"""
+        ALTER TABLE {schema}banks
+        ALTER COLUMN disposition SET DEFAULT '{{"openness": 0.5, "conscientiousness": 0.5, "extraversion": 0.5, "agreeableness": 0.5, "neuroticism": 0.5, "bias_strength": 0.5}}'::jsonb
     """)
     )

hindsight-api/hindsight_api/alembic/versions/rename_personality_to_disposition.py

@@ -9,7 +9,7 @@ Create Date: 2024-12-04
 from collections.abc import Sequence
 
 import sqlalchemy as sa
-from alembic import op
+from alembic import context, op
 from sqlalchemy.dialects import postgresql
 
 # revision identifiers, used by Alembic.
@@ -19,17 +19,25 @@ branch_labels: str | Sequence[str] | None = None
 depends_on: str | Sequence[str] | None = None
 
 
+def _get_target_schema() -> str:
+    """Get the target schema name (tenant schema or 'public')."""
+    schema = context.config.get_main_option("target_schema")
+    return schema if schema else "public"
+
+
 def upgrade() -> None:
     """Rename personality column to disposition in banks table (if it exists)."""
     conn = op.get_bind()
+    target_schema = _get_target_schema()
 
     # Check if 'personality' column exists (old database)
     result = conn.execute(
         sa.text("""
         SELECT column_name
         FROM information_schema.columns
-        WHERE table_name = 'banks' AND column_name = 'personality'
-    """)
+        WHERE table_schema = :schema AND table_name = 'banks' AND column_name = 'personality'
+    """),
+        {"schema": target_schema},
     )
     has_personality = result.fetchone() is not None
 
@@ -38,8 +46,9 @@ def upgrade() -> None:
         sa.text("""
         SELECT column_name
         FROM information_schema.columns
-        WHERE table_name = 'banks' AND column_name = 'disposition'
-    """)
+        WHERE table_schema = :schema AND table_name = 'banks' AND column_name = 'disposition'
+    """),
+        {"schema": target_schema},
     )
     has_disposition = result.fetchone() is not None
 
@@ -63,12 +72,14 @@ def upgrade() -> None:
 def downgrade() -> None:
     """Revert disposition column back to personality."""
     conn = op.get_bind()
+    target_schema = _get_target_schema()
     result = conn.execute(
         sa.text("""
         SELECT column_name
         FROM information_schema.columns
-        WHERE table_name = 'banks' AND column_name = 'disposition'
-    """)
+        WHERE table_schema = :schema AND table_name = 'banks' AND column_name = 'disposition'
+    """),
+        {"schema": target_schema},
     )
     if result.fetchone():
         op.alter_column("banks", "disposition", new_column_name="personality")

hindsight-api/hindsight_api/api/__init__.py

@@ -5,6 +5,7 @@ Provides both HTTP REST API and MCP (Model Context Protocol) server.
 """
 
 import logging
+from contextlib import asynccontextmanager
 from typing import Optional
 
 from fastapi import FastAPI
@@ -45,6 +46,18 @@ def create_app(
         # Both HTTP and MCP
         app = create_app(memory, mcp_api_enabled=True)
     """
+    mcp_app = None
+
+    # Create MCP app first if enabled (we need its lifespan for chaining)
+    if mcp_api_enabled:
+        try:
+            from .mcp import create_mcp_app
+
+            mcp_app = create_mcp_app(memory=memory)
+        except ImportError as e:
+            logger.error(f"MCP server requested but dependencies not available: {e}")
+            logger.error("Install with: pip install hindsight-api[mcp]")
+            raise
 
     # Import and create HTTP API if enabled
     if http_api_enabled:
@@ -57,20 +70,31 @@ def create_app(
         app = FastAPI(title="Hindsight API", version="0.0.7")
         logger.info("HTTP REST API disabled")
 
-    # Mount MCP server if enabled
-    if mcp_api_enabled:
-        try:
-            from .mcp import create_mcp_app
+    # Mount MCP server and chain its lifespan if enabled
+    if mcp_app is not None:
+        # Get the MCP app's underlying Starlette app for lifespan access
+        mcp_starlette_app = mcp_app.mcp_app
 
-            # Create MCP app with dynamic bank_id support
-            # Supports: /mcp/{bank_id}/sse (bank-specific SSE endpoint)
-            mcp_app = create_mcp_app(memory=memory)
-            app.mount(mcp_mount_path, mcp_app)
-            logger.info(f"MCP server enabled at {mcp_mount_path}/{{bank_id}}/sse")
-        except ImportError as e:
-            logger.error(f"MCP server requested but dependencies not available: {e}")
-            logger.error("Install with: pip install hindsight-api[mcp]")
-            raise
+        # Store the original lifespan
+        original_lifespan = app.router.lifespan_context
+
+        @asynccontextmanager
+        async def chained_lifespan(app_instance: FastAPI):
+            """Chain the MCP lifespan with the main app lifespan."""
+            # Start MCP lifespan first
+            async with mcp_starlette_app.router.lifespan_context(mcp_starlette_app):
+                logger.info("MCP lifespan started")
+                # Then start the original app lifespan
+                async with original_lifespan(app_instance):
+                    yield
+            logger.info("MCP lifespan stopped")
+
+        # Replace the app's lifespan with the chained version
+        app.router.lifespan_context = chained_lifespan
+
+        # Mount the MCP middleware
+        app.mount(mcp_mount_path, mcp_app)
+        logger.info(f"MCP server enabled at {mcp_mount_path}/")
 
     return app
 

hindsight-api/hindsight_api/api/mcp.py

@@ -9,6 +9,7 @@ from fastmcp import FastMCP
 
 from hindsight_api import MemoryEngine
 from hindsight_api.engine.response_models import VALID_RECALL_FACT_TYPES
+from hindsight_api.models import RequestContext
 
 # Configure logging from HINDSIGHT_API_LOG_LEVEL environment variable
 _log_level_str = os.environ.get("HINDSIGHT_API_LOG_LEVEL", "info").lower()
@@ -26,12 +27,15 @@ logging.basicConfig(
 )
 logger = logging.getLogger(__name__)
 
-# Context variable to hold the current bank_id from the URL path
+# Default bank_id from environment variable
+DEFAULT_BANK_ID = os.environ.get("HINDSIGHT_MCP_BANK_ID", "default")
+
+# Context variable to hold the current bank_id
 _current_bank_id: ContextVar[str | None] = ContextVar("current_bank_id", default=None)
 
 
 def get_current_bank_id() -> str | None:
-    """Get the current bank_id from context (set from URL path)."""
+    """Get the current bank_id from context."""
     return _current_bank_id.get()
 
 
@@ -43,12 +47,18 @@ def create_mcp_server(memory: MemoryEngine) -> FastMCP:
         memory: MemoryEngine instance (required)
 
     Returns:
-        Configured FastMCP server instance
+        Configured FastMCP server instance with stateless_http enabled
     """
-    mcp = FastMCP("hindsight-mcp-server")
+    # Use stateless_http=True for Claude Code compatibility
+    mcp = FastMCP("hindsight-mcp-server", stateless_http=True)
 
     @mcp.tool()
-    async def retain(content: str, context: str = "general") -> str:
+    async def retain(
+        content: str,
+        context: str = "general",
+        async_processing: bool = True,
+        bank_id: str | None = None,
+    ) -> str:
         """
         Store important information to long-term memory.
 
@@ -64,17 +74,34 @@ def create_mcp_server(memory: MemoryEngine) -> FastMCP:
         Args:
             content: The fact/memory to store (be specific and include relevant details)
             context: Category for the memory (e.g., 'preferences', 'work', 'hobbies', 'family'). Default: 'general'
+            async_processing: If True, queue for background processing and return immediately. If False, wait for completion. Default: True
+            bank_id: Optional bank to store in (defaults to session bank). Use for cross-bank operations.
         """
         try:
-            bank_id = get_current_bank_id()
-            await memory.retain_batch_async(bank_id=bank_id, contents=[{"content": content, "context": context}])
-            return "Memory stored successfully"
+            target_bank = bank_id or get_current_bank_id()
+            if target_bank is None:
+                return "Error: No bank_id configured"
+            contents = [{"content": content, "context": context}]
+            if async_processing:
+                # Queue for background processing and return immediately
+                result = await memory.submit_async_retain(
+                    bank_id=target_bank, contents=contents, request_context=RequestContext()
+                )
+                return f"Memory queued for background processing (operation_id: {result.get('operation_id', 'N/A')})"
+            else:
+                # Wait for completion
+                await memory.retain_batch_async(
+                    bank_id=target_bank,
+                    contents=contents,
+                    request_context=RequestContext(),
+                )
+                return f"Memory stored successfully in bank '{target_bank}'"
         except Exception as e:
             logger.error(f"Error storing memory: {e}", exc_info=True)
             return f"Error: {str(e)}"
 
     @mcp.tool()
-    async def recall(query: str, max_results: int = 10) -> str:
+    async def recall(query: str, max_tokens: int = 4096, bank_id: str | None = None) -> str:
         """
         Search memories to provide personalized, context-aware responses.
 
@@ -86,43 +113,165 @@ def create_mcp_server(memory: MemoryEngine) -> FastMCP:
 
         Args:
             query: Natural language search query (e.g., "user's food preferences", "what projects is user working on")
-            max_results: Maximum number of results to return (default: 10)
+            max_tokens: Maximum tokens in the response (default: 4096)
+            bank_id: Optional bank to search in (defaults to session bank). Use for cross-bank operations.
         """
         try:
-            bank_id = get_current_bank_id()
+            target_bank = bank_id or get_current_bank_id()
+            if target_bank is None:
+                return "Error: No bank_id configured"
             from hindsight_api.engine.memory_engine import Budget
 
-            search_result = await memory.recall_async(
-                bank_id=bank_id, query=query, fact_type=list(VALID_RECALL_FACT_TYPES), budget=Budget.LOW
+            recall_result = await memory.recall_async(
+                bank_id=target_bank,
+                query=query,
+                fact_type=list(VALID_RECALL_FACT_TYPES),
+                budget=Budget.HIGH,
+                max_tokens=max_tokens,
+                request_context=RequestContext(),
             )
 
-            results = [
-                {
-                    "id": fact.id,
-                    "text": fact.text,
-                    "type": fact.fact_type,
-                    "context": fact.context,
-                    "event_date": fact.event_date,
-                }
-                for fact in search_result.results[:max_results]
-            ]
-
-            return json.dumps({"results": results}, indent=2)
+            # Use model's JSON serialization
+            return recall_result.model_dump_json(indent=2)
         except Exception as e:
             logger.error(f"Error searching: {e}", exc_info=True)
-            return json.dumps({"error": str(e), "results": []})
+            return f'{{"error": "{e}", "results": []}}'
+
+    @mcp.tool()
+    async def reflect(query: str, context: str | None = None, budget: str = "low", bank_id: str | None = None) -> str:
+        """
+        Generate thoughtful analysis by synthesizing stored memories with the bank's personality.
+
+        WHEN TO USE THIS TOOL:
+        Use reflect when you need reasoned analysis, not just fact retrieval. This tool
+        thinks through the question using everything the bank knows and its personality traits.
+
+        EXAMPLES OF GOOD QUERIES:
+        - "What patterns have emerged in how I approach debugging?"
+        - "Based on my past decisions, what architectural style do I prefer?"
+        - "What might be the best approach for this problem given what you know about me?"
+        - "How should I prioritize these tasks based on my goals?"
+
+        HOW IT DIFFERS FROM RECALL:
+        - recall: Returns raw facts matching your search (fast lookup)
+        - reflect: Reasons across memories to form a synthesized answer (deeper analysis)
+
+        Use recall for "what did I say about X?" and reflect for "what should I do about X?"
+
+        Args:
+            query: The question or topic to reflect on
+            context: Optional context about why this reflection is needed
+            budget: Search budget - 'low', 'mid', or 'high' (default: 'low')
+            bank_id: Optional bank to reflect in (defaults to session bank). Use for cross-bank operations.
+        """
+        try:
+            target_bank = bank_id or get_current_bank_id()
+            if target_bank is None:
+                return "Error: No bank_id configured"
+            from hindsight_api.engine.memory_engine import Budget
+
+            # Map string budget to enum
+            budget_map = {"low": Budget.LOW, "mid": Budget.MID, "high": Budget.HIGH}
+            budget_enum = budget_map.get(budget.lower(), Budget.LOW)
+
+            reflect_result = await memory.reflect_async(
+                bank_id=target_bank,
+                query=query,
+                budget=budget_enum,
+                context=context,
+                request_context=RequestContext(),
+            )
+
+            return reflect_result.model_dump_json(indent=2)
+        except Exception as e:
+            logger.error(f"Error reflecting: {e}", exc_info=True)
+            return f'{{"error": "{e}", "text": ""}}'
+
+    @mcp.tool()
+    async def list_banks() -> str:
+        """
+        List all available memory banks.
+
+        Use this tool to discover what memory banks exist in the system.
+        Each bank is an isolated memory store (like a separate "brain").
+
+        Returns:
+            JSON list of banks with their IDs, names, dispositions, and backgrounds.
+        """
+        try:
+            banks = await memory.list_banks(request_context=RequestContext())
+            return json.dumps({"banks": banks}, indent=2)
+        except Exception as e:
+            logger.error(f"Error listing banks: {e}", exc_info=True)
+            return f'{{"error": "{e}", "banks": []}}'
+
+    @mcp.tool()
+    async def create_bank(bank_id: str, name: str | None = None, background: str | None = None) -> str:
+        """
+        Create a new memory bank or get an existing one.
+
+        Memory banks are isolated stores - each one is like a separate "brain" for a user/agent.
+        Banks are auto-created with default settings if they don't exist.
+
+        Args:
+            bank_id: Unique identifier for the bank (e.g., 'user-123', 'agent-alpha')
+            name: Optional human-friendly name for the bank
+            background: Optional background context about the bank's owner/purpose
+        """
+        try:
+            # get_bank_profile auto-creates bank if it doesn't exist
+            profile = await memory.get_bank_profile(bank_id, request_context=RequestContext())
+
+            # Update name/background if provided
+            if name is not None or background is not None:
+                await memory.update_bank(
+                    bank_id,
+                    name=name,
+                    background=background,
+                    request_context=RequestContext(),
+                )
+                # Fetch updated profile
+                profile = await memory.get_bank_profile(bank_id, request_context=RequestContext())
+
+            # Serialize disposition if it's a Pydantic model
+            if "disposition" in profile and hasattr(profile["disposition"], "model_dump"):
+                profile["disposition"] = profile["disposition"].model_dump()
+            return json.dumps(profile, indent=2)
+        except Exception as e:
+            logger.error(f"Error creating bank: {e}", exc_info=True)
+            return f'{{"error": "{e}"}}'
 
     return mcp
 
 
 class MCPMiddleware:
-    """ASGI middleware that extracts bank_id from path and sets context."""
+    """ASGI middleware that extracts bank_id from header or path and sets context.
+
+    Bank ID can be provided via:
+    1. X-Bank-Id header (recommended for Claude Code)
+    2. URL path: /mcp/{bank_id}/
+    3. Environment variable HINDSIGHT_MCP_BANK_ID (fallback default)
+
+    For Claude Code, configure with:
+        claude mcp add --transport http hindsight http://localhost:8888/mcp \\
+            --header "X-Bank-Id: my-bank"
+    """
 
     def __init__(self, app, memory: MemoryEngine):
         self.app = app
         self.memory = memory
         self.mcp_server = create_mcp_server(memory)
-        self.mcp_app = self.mcp_server.http_app()
+        self.mcp_app = self.mcp_server.http_app(path="/")
+        # Expose the lifespan for the parent app to chain
+        self.lifespan = self.mcp_app.lifespan_handler if hasattr(self.mcp_app, "lifespan_handler") else None
+
+    def _get_header(self, scope: dict, name: str) -> str | None:
+        """Extract a header value from ASGI scope."""
+        name_lower = name.lower().encode()
+        for header_name, header_value in scope.get("headers", []):
+            if header_name.lower() == name_lower:
+                return header_value.decode()
+        return None
 
     async def __call__(self, scope, receive, send):
         if scope["type"] != "http":
@@ -139,32 +288,39 @@ class MCPMiddleware:
         # Also handle case where mount path wasn't stripped (e.g., /mcp/...)
         if path.startswith("/mcp/"):
             path = path[4:]  # Remove /mcp prefix
+        elif path == "/mcp":
+            path = "/"
 
-        # Extract bank_id from path: /{bank_id}/ or /{bank_id}
-        # http_app expects requests at /
-        if not path.startswith("/") or len(path) <= 1:
-            # No bank_id in path - return error
-            await self._send_error(send, 400, "bank_id required in path: /mcp/{bank_id}/")
-            return
+        # Try to get bank_id from header first (for Claude Code compatibility)
+        bank_id = self._get_header(scope, "X-Bank-Id")
 
-        # Extract bank_id from first path segment
-        parts = path[1:].split("/", 1)
-        if not parts[0]:
-            await self._send_error(send, 400, "bank_id required in path: /mcp/{bank_id}/")
-            return
+        # MCP endpoint paths that should not be treated as bank_ids
+        MCP_ENDPOINTS = {"sse", "messages"}
 
-        bank_id = parts[0]
-        new_path = "/" + parts[1] if len(parts) > 1 else "/"
+        # If no header, try to extract from path: /{bank_id}/...
+        new_path = path
+        if not bank_id and path.startswith("/") and len(path) > 1:
+            parts = path[1:].split("/", 1)
+            # Don't treat MCP endpoints as bank_ids
+            if parts[0] and parts[0] not in MCP_ENDPOINTS:
+                # First segment looks like a bank_id
+                bank_id = parts[0]
+                new_path = "/" + parts[1] if len(parts) > 1 else "/"
+
+        # Fall back to default bank_id
+        if not bank_id:
+            bank_id = DEFAULT_BANK_ID
+            logger.debug(f"Using default bank_id: {bank_id}")
 
         # Set bank_id context
         token = _current_bank_id.set(bank_id)
         try:
             new_scope = scope.copy()
             new_scope["path"] = new_path
+            # Clear root_path since we're passing directly to the app
+            new_scope["root_path"] = ""
 
-            # Wrap send to rewrite the SSE endpoint URL to include bank_id
-            # The SSE app sends "event: endpoint\ndata: /messages\n" but we need
-            # the client to POST to /{bank_id}/messages instead
+            # Wrap send to rewrite the SSE endpoint URL to include bank_id if using path-based routing
             async def send_wrapper(message):
                 if message["type"] == "http.response.body":
                     body = message.get("body", b"")
@@ -200,9 +356,10 @@ def create_mcp_app(memory: MemoryEngine):
     """
     Create an ASGI app that handles MCP requests.
 
-    URL pattern: /mcp/{bank_id}/
-
-    The bank_id is extracted from the URL path and made available to tools.
+    Bank ID can be provided via:
+    1. X-Bank-Id header: claude mcp add --transport http hindsight http://localhost:8888/mcp --header "X-Bank-Id: my-bank"
+    2. URL path: /mcp/{bank_id}/
+    3. Environment variable HINDSIGHT_MCP_BANK_ID (fallback, default: "default")
 
     Args:
         memory: MemoryEngine instance

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
@@ -16,10 +21,30 @@ ENV_LLM_PROVIDER = "HINDSIGHT_API_LLM_PROVIDER"
 ENV_LLM_API_KEY = "HINDSIGHT_API_LLM_API_KEY"
 ENV_LLM_MODEL = "HINDSIGHT_API_LLM_MODEL"
 ENV_LLM_BASE_URL = "HINDSIGHT_API_LLM_BASE_URL"
+ENV_LLM_MAX_CONCURRENT = "HINDSIGHT_API_LLM_MAX_CONCURRENT"
+ENV_LLM_TIMEOUT = "HINDSIGHT_API_LLM_TIMEOUT"
+ENV_LLM_GROQ_SERVICE_TIER = "HINDSIGHT_API_LLM_GROQ_SERVICE_TIER"
+
+# Per-operation LLM configuration (optional, falls back to global LLM config)
+ENV_RETAIN_LLM_PROVIDER = "HINDSIGHT_API_RETAIN_LLM_PROVIDER"
+ENV_RETAIN_LLM_API_KEY = "HINDSIGHT_API_RETAIN_LLM_API_KEY"
+ENV_RETAIN_LLM_MODEL = "HINDSIGHT_API_RETAIN_LLM_MODEL"
+ENV_RETAIN_LLM_BASE_URL = "HINDSIGHT_API_RETAIN_LLM_BASE_URL"
+
+ENV_REFLECT_LLM_PROVIDER = "HINDSIGHT_API_REFLECT_LLM_PROVIDER"
+ENV_REFLECT_LLM_API_KEY = "HINDSIGHT_API_REFLECT_LLM_API_KEY"
+ENV_REFLECT_LLM_MODEL = "HINDSIGHT_API_REFLECT_LLM_MODEL"
+ENV_REFLECT_LLM_BASE_URL = "HINDSIGHT_API_REFLECT_LLM_BASE_URL"
 
 ENV_EMBEDDINGS_PROVIDER = "HINDSIGHT_API_EMBEDDINGS_PROVIDER"
 ENV_EMBEDDINGS_LOCAL_MODEL = "HINDSIGHT_API_EMBEDDINGS_LOCAL_MODEL"
 ENV_EMBEDDINGS_TEI_URL = "HINDSIGHT_API_EMBEDDINGS_TEI_URL"
+ENV_EMBEDDINGS_OPENAI_API_KEY = "HINDSIGHT_API_EMBEDDINGS_OPENAI_API_KEY"
+ENV_EMBEDDINGS_OPENAI_MODEL = "HINDSIGHT_API_EMBEDDINGS_OPENAI_MODEL"
+
+ENV_COHERE_API_KEY = "HINDSIGHT_API_COHERE_API_KEY"
+ENV_EMBEDDINGS_COHERE_MODEL = "HINDSIGHT_API_EMBEDDINGS_COHERE_MODEL"
+ENV_RERANKER_COHERE_MODEL = "HINDSIGHT_API_RERANKER_COHERE_MODEL"
 
 ENV_RERANKER_PROVIDER = "HINDSIGHT_API_RERANKER_PROVIDER"
 ENV_RERANKER_LOCAL_MODEL = "HINDSIGHT_API_RERANKER_LOCAL_MODEL"
@@ -33,17 +58,49 @@ ENV_GRAPH_RETRIEVER = "HINDSIGHT_API_GRAPH_RETRIEVER"
 ENV_MCP_LOCAL_BANK_ID = "HINDSIGHT_API_MCP_LOCAL_BANK_ID"
 ENV_MCP_INSTRUCTIONS = "HINDSIGHT_API_MCP_INSTRUCTIONS"
 
+# Observation thresholds
+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"
+ENV_RETAIN_CHUNK_SIZE = "HINDSIGHT_API_RETAIN_CHUNK_SIZE"
+
+# 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"
+
+# Database connection pool
+ENV_DB_POOL_MIN_SIZE = "HINDSIGHT_API_DB_POOL_MIN_SIZE"
+ENV_DB_POOL_MAX_SIZE = "HINDSIGHT_API_DB_POOL_MAX_SIZE"
+ENV_DB_COMMAND_TIMEOUT = "HINDSIGHT_API_DB_COMMAND_TIMEOUT"
+ENV_DB_ACQUIRE_TIMEOUT = "HINDSIGHT_API_DB_ACQUIRE_TIMEOUT"
+
+# Background task processing
+ENV_TASK_BATCH_SIZE = "HINDSIGHT_API_TASK_BATCH_SIZE"
+ENV_TASK_BATCH_INTERVAL = "HINDSIGHT_API_TASK_BATCH_INTERVAL"
+
 # Default values
 DEFAULT_DATABASE_URL = "pg0"
 DEFAULT_LLM_PROVIDER = "openai"
 DEFAULT_LLM_MODEL = "gpt-5-mini"
+DEFAULT_LLM_MAX_CONCURRENT = 32
+DEFAULT_LLM_TIMEOUT = 120.0  # seconds
 
 DEFAULT_EMBEDDINGS_PROVIDER = "local"
 DEFAULT_EMBEDDINGS_LOCAL_MODEL = "BAAI/bge-small-en-v1.5"
+DEFAULT_EMBEDDINGS_OPENAI_MODEL = "text-embedding-3-small"
+DEFAULT_EMBEDDING_DIMENSION = 384
 
 DEFAULT_RERANKER_PROVIDER = "local"
 DEFAULT_RERANKER_LOCAL_MODEL = "cross-encoder/ms-marco-MiniLM-L-6-v2"
 
+DEFAULT_EMBEDDINGS_COHERE_MODEL = "embed-english-v3.0"
+DEFAULT_RERANKER_COHERE_MODEL = "rerank-english-v3.0"
+
 DEFAULT_HOST = "0.0.0.0"
 DEFAULT_PORT = 8888
 DEFAULT_LOG_LEVEL = "info"
@@ -51,6 +108,27 @@ DEFAULT_MCP_ENABLED = True
 DEFAULT_GRAPH_RETRIEVER = "bfs"  # Options: "bfs", "mpfp"
 DEFAULT_MCP_LOCAL_BANK_ID = "mcp"
 
+# Observation thresholds
+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
+DEFAULT_RETAIN_CHUNK_SIZE = 3000  # Max chars per chunk for fact extraction
+
+# Database migrations
+DEFAULT_RUN_MIGRATIONS_ON_STARTUP = True
+
+# Database connection pool
+DEFAULT_DB_POOL_MIN_SIZE = 5
+DEFAULT_DB_POOL_MAX_SIZE = 100
+DEFAULT_DB_COMMAND_TIMEOUT = 60  # seconds
+DEFAULT_DB_ACQUIRE_TIMEOUT = 30  # seconds
+
+# Background task processing
+DEFAULT_TASK_BATCH_SIZE = 10
+DEFAULT_TASK_BATCH_INTERVAL = 1.0  # seconds
+
 # Default MCP tool descriptions (can be customized via env vars)
 DEFAULT_MCP_RETAIN_DESCRIPTION = """Store important information to long-term memory.
 
@@ -71,8 +149,8 @@ Use this tool PROACTIVELY to:
 - Remember user's goals and context
 - Personalize responses based on past interactions"""
 
-# Required embedding dimension for database schema
-EMBEDDING_DIMENSION = 384
+# Default embedding dimension (used by initial migration, adjusted at runtime)
+EMBEDDING_DIMENSION = DEFAULT_EMBEDDING_DIMENSION
 
 
 @dataclass
@@ -82,11 +160,24 @@ class HindsightConfig:
     # Database
     database_url: str
 
-    # LLM
+    # LLM (default, used as fallback for per-operation config)
     llm_provider: str
     llm_api_key: str | None
     llm_model: str
     llm_base_url: str | None
+    llm_max_concurrent: int
+    llm_timeout: float
+
+    # Per-operation LLM configuration (None = use default LLM config)
+    retain_llm_provider: str | None
+    retain_llm_api_key: str | None
+    retain_llm_model: str | None
+    retain_llm_base_url: str | None
+
+    reflect_llm_provider: str | None
+    reflect_llm_api_key: str | None
+    reflect_llm_model: str | None
+    reflect_llm_base_url: str | None
 
     # Embeddings
     embeddings_provider: str
@@ -107,6 +198,31 @@ class HindsightConfig:
     # Recall
     graph_retriever: str
 
+    # Observation thresholds
+    observation_min_facts: int
+    observation_top_entities: int
+
+    # Retain settings
+    retain_max_completion_tokens: int
+    retain_chunk_size: int
+
+    # Optimization flags
+    skip_llm_verification: bool
+    lazy_reranker: bool
+
+    # Database migrations
+    run_migrations_on_startup: bool
+
+    # Database connection pool
+    db_pool_min_size: int
+    db_pool_max_size: int
+    db_command_timeout: int
+    db_acquire_timeout: int
+
+    # Background task processing
+    task_batch_size: int
+    task_batch_interval: float
+
     @classmethod
     def from_env(cls) -> "HindsightConfig":
         """Create configuration from environment variables."""
@@ -118,6 +234,17 @@ class HindsightConfig:
             llm_api_key=os.getenv(ENV_LLM_API_KEY),
             llm_model=os.getenv(ENV_LLM_MODEL, DEFAULT_LLM_MODEL),
             llm_base_url=os.getenv(ENV_LLM_BASE_URL) or None,
+            llm_max_concurrent=int(os.getenv(ENV_LLM_MAX_CONCURRENT, str(DEFAULT_LLM_MAX_CONCURRENT))),
+            llm_timeout=float(os.getenv(ENV_LLM_TIMEOUT, str(DEFAULT_LLM_TIMEOUT))),
+            # Per-operation LLM config (None = use default)
+            retain_llm_provider=os.getenv(ENV_RETAIN_LLM_PROVIDER) or None,
+            retain_llm_api_key=os.getenv(ENV_RETAIN_LLM_API_KEY) or None,
+            retain_llm_model=os.getenv(ENV_RETAIN_LLM_MODEL) or None,
+            retain_llm_base_url=os.getenv(ENV_RETAIN_LLM_BASE_URL) or None,
+            reflect_llm_provider=os.getenv(ENV_REFLECT_LLM_PROVIDER) or None,
+            reflect_llm_api_key=os.getenv(ENV_REFLECT_LLM_API_KEY) or None,
+            reflect_llm_model=os.getenv(ENV_REFLECT_LLM_MODEL) or None,
+            reflect_llm_base_url=os.getenv(ENV_REFLECT_LLM_BASE_URL) or None,
             # Embeddings
             embeddings_provider=os.getenv(ENV_EMBEDDINGS_PROVIDER, DEFAULT_EMBEDDINGS_PROVIDER),
             embeddings_local_model=os.getenv(ENV_EMBEDDINGS_LOCAL_MODEL, DEFAULT_EMBEDDINGS_LOCAL_MODEL),
@@ -133,6 +260,29 @@ class HindsightConfig:
             mcp_enabled=os.getenv(ENV_MCP_ENABLED, str(DEFAULT_MCP_ENABLED)).lower() == "true",
             # Recall
             graph_retriever=os.getenv(ENV_GRAPH_RETRIEVER, DEFAULT_GRAPH_RETRIEVER),
+            # Optimization flags
+            skip_llm_verification=os.getenv(ENV_SKIP_LLM_VERIFICATION, "false").lower() == "true",
+            lazy_reranker=os.getenv(ENV_LAZY_RERANKER, "false").lower() == "true",
+            # Observation thresholds
+            observation_min_facts=int(os.getenv(ENV_OBSERVATION_MIN_FACTS, str(DEFAULT_OBSERVATION_MIN_FACTS))),
+            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))
+            ),
+            retain_chunk_size=int(os.getenv(ENV_RETAIN_CHUNK_SIZE, str(DEFAULT_RETAIN_CHUNK_SIZE))),
+            # Database migrations
+            run_migrations_on_startup=os.getenv(ENV_RUN_MIGRATIONS_ON_STARTUP, "true").lower() == "true",
+            # Database connection pool
+            db_pool_min_size=int(os.getenv(ENV_DB_POOL_MIN_SIZE, str(DEFAULT_DB_POOL_MIN_SIZE))),
+            db_pool_max_size=int(os.getenv(ENV_DB_POOL_MAX_SIZE, str(DEFAULT_DB_POOL_MAX_SIZE))),
+            db_command_timeout=int(os.getenv(ENV_DB_COMMAND_TIMEOUT, str(DEFAULT_DB_COMMAND_TIMEOUT))),
+            db_acquire_timeout=int(os.getenv(ENV_DB_ACQUIRE_TIMEOUT, str(DEFAULT_DB_ACQUIRE_TIMEOUT))),
+            # Background task processing
+            task_batch_size=int(os.getenv(ENV_TASK_BATCH_SIZE, str(DEFAULT_TASK_BATCH_SIZE))),
+            task_batch_interval=float(os.getenv(ENV_TASK_BATCH_INTERVAL, str(DEFAULT_TASK_BATCH_INTERVAL))),
         )
 
     def get_llm_base_url(self) -> str:
@@ -145,6 +295,8 @@ class HindsightConfig:
             return "https://api.groq.com/openai/v1"
         elif provider == "ollama":
             return "http://localhost:11434/v1"
+        elif provider == "lmstudio":
+            return "http://localhost:1234/v1"
         else:
             return ""
 
@@ -172,11 +324,32 @@ class HindsightConfig:
         """Log the current configuration (without sensitive values)."""
         logger.info(f"Database: {self.database_url}")
         logger.info(f"LLM: provider={self.llm_provider}, model={self.llm_model}")
+        if self.retain_llm_provider or self.retain_llm_model:
+            retain_provider = self.retain_llm_provider or self.llm_provider
+            retain_model = self.retain_llm_model or self.llm_model
+            logger.info(f"LLM (retain): provider={retain_provider}, model={retain_model}")
+        if self.reflect_llm_provider or self.reflect_llm_model:
+            reflect_provider = self.reflect_llm_provider or self.llm_provider
+            reflect_model = self.reflect_llm_model or self.llm_model
+            logger.info(f"LLM (reflect): provider={reflect_provider}, model={reflect_model}")
         logger.info(f"Embeddings: provider={self.embeddings_provider}")
         logger.info(f"Reranker: provider={self.reranker_provider}")
         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/daemon.py

@@ -0,0 +1,204 @@
+"""
+Daemon mode support for Hindsight API.
+
+Provides idle timeout and lockfile management for running as a background daemon.
+"""
+
+import asyncio
+import fcntl
+import logging
+import os
+import sys
+import time
+from pathlib import Path
+
+logger = logging.getLogger(__name__)
+
+# Default daemon configuration
+DEFAULT_DAEMON_PORT = 8889
+DEFAULT_IDLE_TIMEOUT = 0  # 0 = no auto-exit (hindsight-embed passes its own timeout)
+LOCKFILE_PATH = Path.home() / ".hindsight" / "daemon.lock"
+DAEMON_LOG_PATH = Path.home() / ".hindsight" / "daemon.log"
+
+
+class IdleTimeoutMiddleware:
+    """ASGI middleware that tracks activity and exits after idle timeout."""
+
+    def __init__(self, app, idle_timeout: int = DEFAULT_IDLE_TIMEOUT):
+        self.app = app
+        self.idle_timeout = idle_timeout
+        self.last_activity = time.time()
+        self._checker_task = None
+
+    async def __call__(self, scope, receive, send):
+        # Update activity timestamp on each request
+        self.last_activity = time.time()
+        await self.app(scope, receive, send)
+
+    def start_idle_checker(self):
+        """Start the background task that checks for idle timeout."""
+        self._checker_task = asyncio.create_task(self._check_idle())
+
+    async def _check_idle(self):
+        """Background task that exits the process after idle timeout."""
+        # If idle_timeout is 0, don't auto-exit
+        if self.idle_timeout <= 0:
+            return
+
+        while True:
+            await asyncio.sleep(30)  # Check every 30 seconds
+            idle_time = time.time() - self.last_activity
+            if idle_time > self.idle_timeout:
+                logger.info(f"Idle timeout reached ({self.idle_timeout}s), shutting down daemon")
+                # Give a moment for any in-flight requests
+                await asyncio.sleep(1)
+                os._exit(0)
+
+
+class DaemonLock:
+    """
+    File-based lock to prevent multiple daemon instances.
+
+    Uses fcntl.flock for atomic locking on Unix systems.
+    """
+
+    def __init__(self, lockfile: Path = LOCKFILE_PATH):
+        self.lockfile = lockfile
+        self._fd = None
+
+    def acquire(self) -> bool:
+        """
+        Try to acquire the daemon lock.
+
+        Returns True if lock acquired, False if another daemon is running.
+        """
+        self.lockfile.parent.mkdir(parents=True, exist_ok=True)
+
+        try:
+            self._fd = open(self.lockfile, "w")
+            fcntl.flock(self._fd.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
+            # Write PID for debugging
+            self._fd.write(str(os.getpid()))
+            self._fd.flush()
+            return True
+        except (IOError, OSError):
+            # Lock is held by another process
+            if self._fd:
+                self._fd.close()
+                self._fd = None
+            return False
+
+    def release(self):
+        """Release the daemon lock."""
+        if self._fd:
+            try:
+                fcntl.flock(self._fd.fileno(), fcntl.LOCK_UN)
+                self._fd.close()
+            except Exception:
+                pass
+            finally:
+                self._fd = None
+                # Remove lockfile
+                try:
+                    self.lockfile.unlink()
+                except Exception:
+                    pass
+
+    def is_locked(self) -> bool:
+        """Check if the lock is held by another process."""
+        if not self.lockfile.exists():
+            return False
+
+        try:
+            fd = open(self.lockfile, "r")
+            fcntl.flock(fd.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
+            # We got the lock, so no one else has it
+            fcntl.flock(fd.fileno(), fcntl.LOCK_UN)
+            fd.close()
+            return False
+        except (IOError, OSError):
+            return True
+
+    def get_pid(self) -> int | None:
+        """Get the PID of the daemon holding the lock."""
+        if not self.lockfile.exists():
+            return None
+        try:
+            with open(self.lockfile, "r") as f:
+                return int(f.read().strip())
+        except (ValueError, IOError):
+            return None
+
+
+def daemonize():
+    """
+    Fork the current process into a background daemon.
+
+    Uses double-fork technique to properly detach from terminal.
+    """
+    # First fork
+    pid = os.fork()
+    if pid > 0:
+        # Parent exits
+        sys.exit(0)
+
+    # Create new session
+    os.setsid()
+
+    # Second fork to prevent zombie processes
+    pid = os.fork()
+    if pid > 0:
+        sys.exit(0)
+
+    # Redirect standard file descriptors to log file
+    DAEMON_LOG_PATH.parent.mkdir(parents=True, exist_ok=True)
+
+    sys.stdout.flush()
+    sys.stderr.flush()
+
+    # Redirect stdin to /dev/null
+    with open("/dev/null", "r") as devnull:
+        os.dup2(devnull.fileno(), sys.stdin.fileno())
+
+    # Redirect stdout/stderr to log file
+    log_fd = open(DAEMON_LOG_PATH, "a")
+    os.dup2(log_fd.fileno(), sys.stdout.fileno())
+    os.dup2(log_fd.fileno(), sys.stderr.fileno())
+
+
+def check_daemon_running(port: int = DEFAULT_DAEMON_PORT) -> bool:
+    """Check if a daemon is running and responsive on the given port."""
+    import socket
+
+    try:
+        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        sock.settimeout(1)
+        result = sock.connect_ex(("127.0.0.1", port))
+        sock.close()
+        return result == 0
+    except Exception:
+        return False
+
+
+def stop_daemon(port: int = DEFAULT_DAEMON_PORT) -> bool:
+    """Stop a running daemon by sending SIGTERM to the process."""
+    lock = DaemonLock()
+    pid = lock.get_pid()
+
+    if pid is None:
+        return False
+
+    try:
+        import signal
+
+        os.kill(pid, signal.SIGTERM)
+        # Wait for process to exit
+        for _ in range(50):  # Wait up to 5 seconds
+            time.sleep(0.1)
+            try:
+                os.kill(pid, 0)  # Check if process exists
+            except OSError:
+                return True  # Process exited
+        return False
+    except OSError:
+        return False

hindsight-api/hindsight_api/engine/__init__.py

@@ -11,7 +11,13 @@ from .cross_encoder import CrossEncoderModel, LocalSTCrossEncoder, RemoteTEICros
 from .db_utils import acquire_with_retry
 from .embeddings import Embeddings, LocalSTEmbeddings, RemoteTEIEmbeddings
 from .llm_wrapper import LLMConfig
-from .memory_engine import MemoryEngine
+from .memory_engine import (
+    MemoryEngine,
+    UnqualifiedTableError,
+    fq_table,
+    get_current_schema,
+    validate_sql_schema,
+)
 from .response_models import MemoryFact, RecallResult, ReflectResult
 from .search.trace import (
     EntryPoint,
@@ -49,4 +55,9 @@ __all__ = [
     "RecallResult",
     "ReflectResult",
     "MemoryFact",
+    # Schema safety utilities
+    "fq_table",
+    "get_current_schema",
+    "validate_sql_schema",
+    "UnqualifiedTableError",
 ]

hindsight-api/hindsight_api/engine/cross_encoder.py

@@ -13,8 +13,11 @@ from abc import ABC, abstractmethod
 import httpx
 
 from ..config import (
+    DEFAULT_RERANKER_COHERE_MODEL,
     DEFAULT_RERANKER_LOCAL_MODEL,
     DEFAULT_RERANKER_PROVIDER,
+    ENV_COHERE_API_KEY,
+    ENV_RERANKER_COHERE_MODEL,
     ENV_RERANKER_LOCAL_MODEL,
     ENV_RERANKER_PROVIDER,
     ENV_RERANKER_TEI_URL,
@@ -278,6 +281,96 @@ class RemoteTEICrossEncoder(CrossEncoderModel):
         return all_scores
 
 
+class CohereCrossEncoder(CrossEncoderModel):
+    """
+    Cohere cross-encoder implementation using the Cohere Rerank API.
+
+    Supports rerank-english-v3.0 and rerank-multilingual-v3.0 models.
+    """
+
+    def __init__(
+        self,
+        api_key: str,
+        model: str = DEFAULT_RERANKER_COHERE_MODEL,
+        timeout: float = 60.0,
+    ):
+        """
+        Initialize Cohere cross-encoder client.
+
+        Args:
+            api_key: Cohere API key
+            model: Cohere rerank model name (default: rerank-english-v3.0)
+            timeout: Request timeout in seconds (default: 60.0)
+        """
+        self.api_key = api_key
+        self.model = model
+        self.timeout = timeout
+        self._client = None
+
+    @property
+    def provider_name(self) -> str:
+        return "cohere"
+
+    async def initialize(self) -> None:
+        """Initialize the Cohere client."""
+        if self._client is not None:
+            return
+
+        try:
+            import cohere
+        except ImportError:
+            raise ImportError("cohere is required for CohereCrossEncoder. Install it with: pip install cohere")
+
+        logger.info(f"Reranker: initializing Cohere provider with model {self.model}")
+        self._client = cohere.Client(api_key=self.api_key, timeout=self.timeout)
+        logger.info("Reranker: Cohere provider initialized")
+
+    def predict(self, pairs: list[tuple[str, str]]) -> list[float]:
+        """
+        Score query-document pairs using the Cohere Rerank API.
+
+        Args:
+            pairs: List of (query, document) tuples to score
+
+        Returns:
+            List of relevance scores
+        """
+        if self._client is None:
+            raise RuntimeError("Reranker not initialized. Call initialize() first.")
+
+        if not pairs:
+            return []
+
+        # Group pairs by query for efficient batching
+        # Cohere rerank expects one query with multiple documents
+        query_groups: dict[str, list[tuple[int, str]]] = {}
+        for idx, (query, text) in enumerate(pairs):
+            if query not in query_groups:
+                query_groups[query] = []
+            query_groups[query].append((idx, text))
+
+        all_scores = [0.0] * len(pairs)
+
+        for query, indexed_texts in query_groups.items():
+            texts = [text for _, text in indexed_texts]
+            indices = [idx for idx, _ in indexed_texts]
+
+            response = self._client.rerank(
+                query=query,
+                documents=texts,
+                model=self.model,
+                return_documents=False,
+            )
+
+            # Map scores back to original positions
+            for result in response.results:
+                original_idx = result.index
+                score = result.relevance_score
+                all_scores[indices[original_idx]] = score
+
+        return all_scores
+
+
 def create_cross_encoder_from_env() -> CrossEncoderModel:
     """
     Create a CrossEncoderModel instance based on environment variables.
@@ -298,5 +391,11 @@ def create_cross_encoder_from_env() -> CrossEncoderModel:
         model = os.environ.get(ENV_RERANKER_LOCAL_MODEL)
         model_name = model or DEFAULT_RERANKER_LOCAL_MODEL
         return LocalSTCrossEncoder(model_name=model_name)
+    elif provider == "cohere":
+        api_key = os.environ.get(ENV_COHERE_API_KEY)
+        if not api_key:
+            raise ValueError(f"{ENV_COHERE_API_KEY} is required when {ENV_RERANKER_PROVIDER} is 'cohere'")
+        model = os.environ.get(ENV_RERANKER_COHERE_MODEL, DEFAULT_RERANKER_COHERE_MODEL)
+        return CohereCrossEncoder(api_key=api_key, model=model)
     else:
-        raise ValueError(f"Unknown reranker provider: {provider}. Supported: 'local', 'tei'")
+        raise ValueError(f"Unknown reranker provider: {provider}. Supported: 'local', 'tei', 'cohere'")

hindsight-api/hindsight_api/engine/embeddings.py

@@ -3,8 +3,8 @@ Embeddings abstraction for the memory system.
 
 Provides an interface for generating embeddings with different backends.
 
-IMPORTANT: All embeddings must produce 384-dimensional vectors to match
-the database schema (pgvector column defined as vector(384)).
+The embedding dimension is auto-detected from the model at initialization.
+The database schema is automatically adjusted to match the model's dimension.
 
 Configuration via environment variables - see hindsight_api.config for all env var names.
 """
@@ -16,12 +16,18 @@ from abc import ABC, abstractmethod
 import httpx
 
 from ..config import (
+    DEFAULT_EMBEDDINGS_COHERE_MODEL,
     DEFAULT_EMBEDDINGS_LOCAL_MODEL,
+    DEFAULT_EMBEDDINGS_OPENAI_MODEL,
     DEFAULT_EMBEDDINGS_PROVIDER,
-    EMBEDDING_DIMENSION,
+    ENV_COHERE_API_KEY,
+    ENV_EMBEDDINGS_COHERE_MODEL,
     ENV_EMBEDDINGS_LOCAL_MODEL,
+    ENV_EMBEDDINGS_OPENAI_API_KEY,
+    ENV_EMBEDDINGS_OPENAI_MODEL,
     ENV_EMBEDDINGS_PROVIDER,
     ENV_EMBEDDINGS_TEI_URL,
+    ENV_LLM_API_KEY,
 )
 
 logger = logging.getLogger(__name__)
@@ -31,8 +37,8 @@ class Embeddings(ABC):
     """
     Abstract base class for embedding generation.
 
-    All implementations MUST generate 384-dimensional embeddings to match
-    the database schema.
+    The embedding dimension is determined by the model and detected at initialization.
+    The database schema is automatically adjusted to match the model's dimension.
     """
 
     @property
@@ -41,6 +47,12 @@ class Embeddings(ABC):
         """Return a human-readable name for this provider (e.g., 'local', 'tei')."""
         pass
 
+    @property
+    @abstractmethod
+    def dimension(self) -> int:
+        """Return the embedding dimension produced by this model."""
+        pass
+
     @abstractmethod
     async def initialize(self) -> None:
         """
@@ -54,13 +66,13 @@ class Embeddings(ABC):
     @abstractmethod
     def encode(self, texts: list[str]) -> list[list[float]]:
         """
-        Generate 384-dimensional embeddings for a list of texts.
+        Generate embeddings for a list of texts.
 
         Args:
             texts: List of text strings to encode
 
         Returns:
-            List of 384-dimensional embedding vectors (each is a list of floats)
+            List of embedding vectors (each is a list of floats)
         """
         pass
 
@@ -70,9 +82,7 @@ class LocalSTEmbeddings(Embeddings):
     Local embeddings implementation using SentenceTransformers.
 
     Call initialize() during startup to load the model and avoid cold starts.
-
-    Default model is BAAI/bge-small-en-v1.5 which produces 384-dimensional
-    embeddings matching the database schema.
+    The embedding dimension is auto-detected from the model.
     """
 
     def __init__(self, model_name: str | None = None):
@@ -81,16 +91,22 @@ class LocalSTEmbeddings(Embeddings):
 
         Args:
             model_name: Name of the SentenceTransformer model to use.
-                       Must produce 384-dimensional embeddings.
                        Default: BAAI/bge-small-en-v1.5
         """
         self.model_name = model_name or DEFAULT_EMBEDDINGS_LOCAL_MODEL
         self._model = None
+        self._dimension: int | None = None
 
     @property
     def provider_name(self) -> str:
         return "local"
 
+    @property
+    def dimension(self) -> int:
+        if self._dimension is None:
+            raise RuntimeError("Embeddings not initialized. Call initialize() first.")
+        return self._dimension
+
     async def initialize(self) -> None:
         """Load the embedding model."""
         if self._model is not None:
@@ -112,26 +128,18 @@ class LocalSTEmbeddings(Embeddings):
             model_kwargs={"low_cpu_mem_usage": False, "device_map": None},
         )
 
-        # Validate dimension matches database schema
-        model_dim = self._model.get_sentence_embedding_dimension()
-        if model_dim != EMBEDDING_DIMENSION:
-            raise ValueError(
-                f"Model {self.model_name} produces {model_dim}-dimensional embeddings, "
-                f"but database schema requires {EMBEDDING_DIMENSION} dimensions. "
-                f"Use a model that produces {EMBEDDING_DIMENSION}-dimensional embeddings."
-            )
-
-        logger.info(f"Embeddings: local provider initialized (dim: {model_dim})")
+        self._dimension = self._model.get_sentence_embedding_dimension()
+        logger.info(f"Embeddings: local provider initialized (dim: {self._dimension})")
 
     def encode(self, texts: list[str]) -> list[list[float]]:
         """
-        Generate 384-dimensional embeddings for a list of texts.
+        Generate embeddings for a list of texts.
 
         Args:
             texts: List of text strings to encode
 
         Returns:
-            List of 384-dimensional embedding vectors
+            List of embedding vectors
         """
         if self._model is None:
             raise RuntimeError("Embeddings not initialized. Call initialize() first.")
@@ -146,7 +154,7 @@ class RemoteTEIEmbeddings(Embeddings):
     TEI provides a high-performance inference server for embedding models.
     See: https://github.com/huggingface/text-embeddings-inference
 
-    The server should be running a model that produces 384-dimensional embeddings.
+    The embedding dimension is auto-detected from the server at initialization.
     """
 
     def __init__(
@@ -174,11 +182,18 @@ class RemoteTEIEmbeddings(Embeddings):
         self.retry_delay = retry_delay
         self._client: httpx.Client | None = None
         self._model_id: str | None = None
+        self._dimension: int | None = None
 
     @property
     def provider_name(self) -> str:
         return "tei"
 
+    @property
+    def dimension(self) -> int:
+        if self._dimension is None:
+            raise RuntimeError("Embeddings not initialized. Call initialize() first.")
+        return self._dimension
+
     def _request_with_retry(self, method: str, url: str, **kwargs) -> httpx.Response:
         """Make an HTTP request with automatic retries on transient errors."""
         import time
@@ -229,7 +244,24 @@ class RemoteTEIEmbeddings(Embeddings):
             response = self._request_with_retry("GET", f"{self.base_url}/info")
             info = response.json()
             self._model_id = info.get("model_id", "unknown")
-            logger.info(f"Embeddings: TEI provider initialized (model: {self._model_id})")
+
+            # Get dimension from server info or by doing a test embedding
+            if "max_input_length" in info and "model_dtype" in info:
+                # Try to get dimension from info endpoint (some TEI versions expose it)
+                # If not available, do a test embedding
+                pass
+
+            # Do a test embedding to detect dimension
+            test_response = self._request_with_retry(
+                "POST",
+                f"{self.base_url}/embed",
+                json={"inputs": ["test"]},
+            )
+            test_embeddings = test_response.json()
+            if test_embeddings and len(test_embeddings) > 0:
+                self._dimension = len(test_embeddings[0])
+
+            logger.info(f"Embeddings: TEI provider initialized (model: {self._model_id}, dim: {self._dimension})")
         except httpx.HTTPError as e:
             raise RuntimeError(f"Failed to connect to TEI server at {self.base_url}: {e}")
 
@@ -269,6 +301,234 @@ class RemoteTEIEmbeddings(Embeddings):
         return all_embeddings
 
 
+class OpenAIEmbeddings(Embeddings):
+    """
+    OpenAI embeddings implementation using the OpenAI API.
+
+    Supports text-embedding-3-small (1536 dims), text-embedding-3-large (3072 dims),
+    and text-embedding-ada-002 (1536 dims, legacy).
+
+    The embedding dimension is auto-detected from the model at initialization.
+    """
+
+    # Known dimensions for OpenAI embedding models
+    MODEL_DIMENSIONS = {
+        "text-embedding-3-small": 1536,
+        "text-embedding-3-large": 3072,
+        "text-embedding-ada-002": 1536,
+    }
+
+    def __init__(
+        self,
+        api_key: str,
+        model: str = DEFAULT_EMBEDDINGS_OPENAI_MODEL,
+        batch_size: int = 100,
+        max_retries: int = 3,
+    ):
+        """
+        Initialize OpenAI embeddings client.
+
+        Args:
+            api_key: OpenAI API key
+            model: OpenAI embedding model name (default: text-embedding-3-small)
+            batch_size: Maximum batch size for embedding requests (default: 100)
+            max_retries: Maximum number of retries for failed requests (default: 3)
+        """
+        self.api_key = api_key
+        self.model = model
+        self.batch_size = batch_size
+        self.max_retries = max_retries
+        self._client = None
+        self._dimension: int | None = None
+
+    @property
+    def provider_name(self) -> str:
+        return "openai"
+
+    @property
+    def dimension(self) -> int:
+        if self._dimension is None:
+            raise RuntimeError("Embeddings not initialized. Call initialize() first.")
+        return self._dimension
+
+    async def initialize(self) -> None:
+        """Initialize the OpenAI client and detect dimension."""
+        if self._client is not None:
+            return
+
+        try:
+            from openai import OpenAI
+        except ImportError:
+            raise ImportError("openai is required for OpenAIEmbeddings. Install it with: pip install openai")
+
+        logger.info(f"Embeddings: initializing OpenAI provider with model {self.model}")
+        self._client = OpenAI(api_key=self.api_key, max_retries=self.max_retries)
+
+        # Try to get dimension from known models, otherwise do a test embedding
+        if self.model in self.MODEL_DIMENSIONS:
+            self._dimension = self.MODEL_DIMENSIONS[self.model]
+        else:
+            # Do a test embedding to detect dimension
+            response = self._client.embeddings.create(
+                model=self.model,
+                input=["test"],
+            )
+            if response.data:
+                self._dimension = len(response.data[0].embedding)
+
+        logger.info(f"Embeddings: OpenAI provider initialized (model: {self.model}, dim: {self._dimension})")
+
+    def encode(self, texts: list[str]) -> list[list[float]]:
+        """
+        Generate embeddings using the OpenAI API.
+
+        Args:
+            texts: List of text strings to encode
+
+        Returns:
+            List of embedding vectors
+        """
+        if self._client is None:
+            raise RuntimeError("Embeddings not initialized. Call initialize() first.")
+
+        if not texts:
+            return []
+
+        all_embeddings = []
+
+        # Process in batches
+        for i in range(0, len(texts), self.batch_size):
+            batch = texts[i : i + self.batch_size]
+
+            response = self._client.embeddings.create(
+                model=self.model,
+                input=batch,
+            )
+
+            # Sort by index to ensure correct order
+            batch_embeddings = sorted(response.data, key=lambda x: x.index)
+            all_embeddings.extend([e.embedding for e in batch_embeddings])
+
+        return all_embeddings
+
+
+class CohereEmbeddings(Embeddings):
+    """
+    Cohere embeddings implementation using the Cohere API.
+
+    Supports embed-english-v3.0 (1024 dims) and embed-multilingual-v3.0 (1024 dims).
+
+    The embedding dimension is auto-detected from the model at initialization.
+    """
+
+    # Known dimensions for Cohere embedding models
+    MODEL_DIMENSIONS = {
+        "embed-english-v3.0": 1024,
+        "embed-multilingual-v3.0": 1024,
+        "embed-english-light-v3.0": 384,
+        "embed-multilingual-light-v3.0": 384,
+        "embed-english-v2.0": 4096,
+        "embed-multilingual-v2.0": 768,
+    }
+
+    def __init__(
+        self,
+        api_key: str,
+        model: str = DEFAULT_EMBEDDINGS_COHERE_MODEL,
+        batch_size: int = 96,
+        timeout: float = 60.0,
+        input_type: str = "search_document",
+    ):
+        """
+        Initialize Cohere embeddings client.
+
+        Args:
+            api_key: Cohere API key
+            model: Cohere embedding model name (default: embed-english-v3.0)
+            batch_size: Maximum batch size for embedding requests (default: 96, Cohere's limit)
+            timeout: Request timeout in seconds (default: 60.0)
+            input_type: Input type for embeddings (default: search_document).
+                       Options: search_document, search_query, classification, clustering
+        """
+        self.api_key = api_key
+        self.model = model
+        self.batch_size = batch_size
+        self.timeout = timeout
+        self.input_type = input_type
+        self._client = None
+        self._dimension: int | None = None
+
+    @property
+    def provider_name(self) -> str:
+        return "cohere"
+
+    @property
+    def dimension(self) -> int:
+        if self._dimension is None:
+            raise RuntimeError("Embeddings not initialized. Call initialize() first.")
+        return self._dimension
+
+    async def initialize(self) -> None:
+        """Initialize the Cohere client and detect dimension."""
+        if self._client is not None:
+            return
+
+        try:
+            import cohere
+        except ImportError:
+            raise ImportError("cohere is required for CohereEmbeddings. Install it with: pip install cohere")
+
+        logger.info(f"Embeddings: initializing Cohere provider with model {self.model}")
+        self._client = cohere.Client(api_key=self.api_key, timeout=self.timeout)
+
+        # Try to get dimension from known models, otherwise do a test embedding
+        if self.model in self.MODEL_DIMENSIONS:
+            self._dimension = self.MODEL_DIMENSIONS[self.model]
+        else:
+            # Do a test embedding to detect dimension
+            response = self._client.embed(
+                texts=["test"],
+                model=self.model,
+                input_type=self.input_type,
+            )
+            if response.embeddings:
+                self._dimension = len(response.embeddings[0])
+
+        logger.info(f"Embeddings: Cohere provider initialized (model: {self.model}, dim: {self._dimension})")
+
+    def encode(self, texts: list[str]) -> list[list[float]]:
+        """
+        Generate embeddings using the Cohere API.
+
+        Args:
+            texts: List of text strings to encode
+
+        Returns:
+            List of embedding vectors
+        """
+        if self._client is None:
+            raise RuntimeError("Embeddings not initialized. Call initialize() first.")
+
+        if not texts:
+            return []
+
+        all_embeddings = []
+
+        # Process in batches
+        for i in range(0, len(texts), self.batch_size):
+            batch = texts[i : i + self.batch_size]
+
+            response = self._client.embed(
+                texts=batch,
+                model=self.model,
+                input_type=self.input_type,
+            )
+
+            all_embeddings.extend(response.embeddings)
+
+        return all_embeddings
+
+
 def create_embeddings_from_env() -> Embeddings:
     """
     Create an Embeddings instance based on environment variables.
@@ -289,5 +549,21 @@ def create_embeddings_from_env() -> Embeddings:
         model = os.environ.get(ENV_EMBEDDINGS_LOCAL_MODEL)
         model_name = model or DEFAULT_EMBEDDINGS_LOCAL_MODEL
         return LocalSTEmbeddings(model_name=model_name)
+    elif provider == "openai":
+        # Use dedicated embeddings API key, or fall back to LLM API key
+        api_key = os.environ.get(ENV_EMBEDDINGS_OPENAI_API_KEY) or os.environ.get(ENV_LLM_API_KEY)
+        if not api_key:
+            raise ValueError(
+                f"{ENV_EMBEDDINGS_OPENAI_API_KEY} or {ENV_LLM_API_KEY} is required "
+                f"when {ENV_EMBEDDINGS_PROVIDER} is 'openai'"
+            )
+        model = os.environ.get(ENV_EMBEDDINGS_OPENAI_MODEL, DEFAULT_EMBEDDINGS_OPENAI_MODEL)
+        return OpenAIEmbeddings(api_key=api_key, model=model)
+    elif provider == "cohere":
+        api_key = os.environ.get(ENV_COHERE_API_KEY)
+        if not api_key:
+            raise ValueError(f"{ENV_COHERE_API_KEY} is required when {ENV_EMBEDDINGS_PROVIDER} is 'cohere'")
+        model = os.environ.get(ENV_EMBEDDINGS_COHERE_MODEL, DEFAULT_EMBEDDINGS_COHERE_MODEL)
+        return CohereEmbeddings(api_key=api_key, model=model)
     else:
-        raise ValueError(f"Unknown embeddings provider: {provider}. Supported: 'local', 'tei'")
+        raise ValueError(f"Unknown embeddings provider: {provider}. Supported: 'local', 'tei', 'openai', 'cohere'")

hindsight-api/hindsight_api/engine/entity_resolver.py

@@ -11,6 +11,7 @@ from difflib import SequenceMatcher
 import asyncpg
 
 from .db_utils import acquire_with_retry
+from .memory_engine import fq_table
 
 # Load spaCy model (singleton)
 _nlp = None
@@ -68,9 +69,9 @@ class EntityResolver:
     ) -> list[str]:
         # Query ALL candidates for this bank
         all_entities = await conn.fetch(
-            """
+            f"""
             SELECT canonical_name, id, metadata, last_seen, mention_count
-            FROM entities
+            FROM {fq_table("entities")}
             WHERE bank_id = $1
             """,
             bank_id,
@@ -82,11 +83,11 @@ class EntityResolver:
         # Query ALL co-occurrences for this bank's entities in one query
         # This builds a map of entity_id -> set of co-occurring entity names
         all_cooccurrences = await conn.fetch(
-            """
+            f"""
             SELECT ec.entity_id_1, ec.entity_id_2, ec.cooccurrence_count
-            FROM entity_cooccurrences ec
-            WHERE ec.entity_id_1 IN (SELECT id FROM entities WHERE bank_id = $1)
-               OR ec.entity_id_2 IN (SELECT id FROM entities WHERE bank_id = $1)
+            FROM {fq_table("entity_cooccurrences")} ec
+            WHERE ec.entity_id_1 IN (SELECT id FROM {fq_table("entities")} WHERE bank_id = $1)
+               OR ec.entity_id_2 IN (SELECT id FROM {fq_table("entities")} WHERE bank_id = $1)
             """,
             bank_id,
         )
@@ -195,8 +196,8 @@ class EntityResolver:
         # Batch update existing entities
         if entities_to_update:
             await conn.executemany(
-                """
-                UPDATE entities SET
+                f"""
+                UPDATE {fq_table("entities")} SET
                     mention_count = mention_count + 1,
                     last_seen = $2
                 WHERE id = $1::uuid
@@ -232,13 +233,13 @@ class EntityResolver:
             # Batch INSERT ... ON CONFLICT with RETURNING
             # This is much faster than individual inserts
             rows = await conn.fetch(
-                """
-                INSERT INTO entities (bank_id, canonical_name, first_seen, last_seen, mention_count)
+                f"""
+                INSERT INTO {fq_table("entities")} (bank_id, canonical_name, first_seen, last_seen, mention_count)
                 SELECT $1, name, event_date, event_date, 1
                 FROM unnest($2::text[], $3::timestamptz[]) AS t(name, event_date)
                 ON CONFLICT (bank_id, LOWER(canonical_name))
                 DO UPDATE SET
-                    mention_count = entities.mention_count + 1,
+                    mention_count = {fq_table("entities")}.mention_count + 1,
                     last_seen = EXCLUDED.last_seen
                 RETURNING id
                 """,
@@ -279,9 +280,9 @@ class EntityResolver:
         async with acquire_with_retry(self.pool) as conn:
             # Find candidate entities with similar name
             candidates = await conn.fetch(
-                """
+                f"""
                 SELECT id, canonical_name, metadata, last_seen
-                FROM entities
+                FROM {fq_table("entities")}
                 WHERE bank_id = $1
                   AND (
                     canonical_name ILIKE $2
@@ -326,10 +327,10 @@ class EntityResolver:
                 # Get entities that co-occurred with this candidate before
                 # Use the materialized co-occurrence cache for fast lookup
                 co_entity_rows = await conn.fetch(
-                    """
+                    f"""
                     SELECT e.canonical_name, ec.cooccurrence_count
-                    FROM entity_cooccurrences ec
-                    JOIN entities e ON (
+                    FROM {fq_table("entity_cooccurrences")} ec
+                    JOIN {fq_table("entities")} e ON (
                         CASE
                             WHEN ec.entity_id_1 = $1 THEN ec.entity_id_2
                             WHEN ec.entity_id_2 = $1 THEN ec.entity_id_1
@@ -365,8 +366,8 @@ class EntityResolver:
             if best_score > threshold:
                 # Update entity
                 await conn.execute(
-                    """
-                    UPDATE entities
+                    f"""
+                    UPDATE {fq_table("entities")}
                     SET mention_count = mention_count + 1,
                         last_seen = $1
                     WHERE id = $2
@@ -402,12 +403,12 @@ class EntityResolver:
             Entity ID
         """
         entity_id = await conn.fetchval(
-            """
-            INSERT INTO entities (bank_id, canonical_name, first_seen, last_seen, mention_count)
+            f"""
+            INSERT INTO {fq_table("entities")} (bank_id, canonical_name, first_seen, last_seen, mention_count)
             VALUES ($1, $2, $3, $4, 1)
             ON CONFLICT (bank_id, LOWER(canonical_name))
             DO UPDATE SET
-                mention_count = entities.mention_count + 1,
+                mention_count = {fq_table("entities")}.mention_count + 1,
                 last_seen = EXCLUDED.last_seen
             RETURNING id
             """,
@@ -430,8 +431,8 @@ class EntityResolver:
         async with acquire_with_retry(self.pool) as conn:
             # Insert unit-entity link
             await conn.execute(
-                """
-                INSERT INTO unit_entities (unit_id, entity_id)
+                f"""
+                INSERT INTO {fq_table("unit_entities")} (unit_id, entity_id)
                 VALUES ($1, $2)
                 ON CONFLICT DO NOTHING
                 """,
@@ -441,9 +442,9 @@ class EntityResolver:
 
             # Update co-occurrence cache: find other entities in this unit
             rows = await conn.fetch(
-                """
+                f"""
                 SELECT entity_id
-                FROM unit_entities
+                FROM {fq_table("unit_entities")}
                 WHERE unit_id = $1 AND entity_id != $2
                 """,
                 unit_id,
@@ -472,12 +473,12 @@ class EntityResolver:
             entity_id_1, entity_id_2 = entity_id_2, entity_id_1
 
         await conn.execute(
-            """
-            INSERT INTO entity_cooccurrences (entity_id_1, entity_id_2, cooccurrence_count, last_cooccurred)
+            f"""
+            INSERT INTO {fq_table("entity_cooccurrences")} (entity_id_1, entity_id_2, cooccurrence_count, last_cooccurred)
             VALUES ($1, $2, 1, NOW())
             ON CONFLICT (entity_id_1, entity_id_2)
             DO UPDATE SET
-                cooccurrence_count = entity_cooccurrences.cooccurrence_count + 1,
+                cooccurrence_count = {fq_table("entity_cooccurrences")}.cooccurrence_count + 1,
                 last_cooccurred = NOW()
             """,
             entity_id_1,
@@ -506,8 +507,8 @@ class EntityResolver:
     async def _link_units_to_entities_batch_impl(self, conn, unit_entity_pairs: list[tuple[str, str]]):
         # Batch insert all unit-entity links
         await conn.executemany(
-            """
-            INSERT INTO unit_entities (unit_id, entity_id)
+            f"""
+            INSERT INTO {fq_table("unit_entities")} (unit_id, entity_id)
             VALUES ($1, $2)
             ON CONFLICT DO NOTHING
             """,
@@ -541,12 +542,12 @@ class EntityResolver:
         if cooccurrence_pairs:
             now = datetime.now(UTC)
             await conn.executemany(
-                """
-                INSERT INTO entity_cooccurrences (entity_id_1, entity_id_2, cooccurrence_count, last_cooccurred)
+                f"""
+                INSERT INTO {fq_table("entity_cooccurrences")} (entity_id_1, entity_id_2, cooccurrence_count, last_cooccurred)
                 VALUES ($1, $2, $3, $4)
                 ON CONFLICT (entity_id_1, entity_id_2)
                 DO UPDATE SET
-                    cooccurrence_count = entity_cooccurrences.cooccurrence_count + 1,
+                    cooccurrence_count = {fq_table("entity_cooccurrences")}.cooccurrence_count + 1,
                     last_cooccurred = EXCLUDED.last_cooccurred
                 """,
                 [(e1, e2, 1, now) for e1, e2 in cooccurrence_pairs],
@@ -565,9 +566,9 @@ class EntityResolver:
         """
         async with acquire_with_retry(self.pool) as conn:
             rows = await conn.fetch(
-                """
+                f"""
                 SELECT unit_id
-                FROM unit_entities
+                FROM {fq_table("unit_entities")}
                 WHERE entity_id = $1
                 ORDER BY unit_id
                 LIMIT $2
@@ -594,8 +595,8 @@ class EntityResolver:
         """
         async with acquire_with_retry(self.pool) as conn:
             row = await conn.fetchrow(
-                """
-                SELECT id FROM entities
+                f"""
+                SELECT id FROM {fq_table("entities")}
                 WHERE bank_id = $1
                   AND canonical_name ILIKE $2
                 ORDER BY mention_count DESC

hindsight-api/hindsight_api/engine/interface.py

@@ -0,0 +1,598 @@
+"""Abstract interface for MemoryEngine public methods.
+
+This module defines the public API that HTTP endpoints and extensions should use
+to interact with the memory system. All methods require a RequestContext for
+authentication when a TenantExtension is configured.
+"""
+
+from abc import ABC, abstractmethod
+from datetime import datetime
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from hindsight_api.engine.memory_engine import Budget
+    from hindsight_api.engine.response_models import RecallResult, ReflectResult
+    from hindsight_api.models import RequestContext
+
+
+class MemoryEngineInterface(ABC):
+    """
+    Abstract interface for the Memory Engine.
+
+    This defines the public API that should be used by HTTP endpoints and extensions.
+    All methods require a RequestContext for authentication.
+    """
+
+    # =========================================================================
+    # Health & Status
+    # =========================================================================
+
+    @abstractmethod
+    async def health_check(self) -> dict:
+        """
+        Check the health of the memory system.
+
+        Returns:
+            Dict with 'status' key ('healthy' or 'unhealthy') and additional info.
+        """
+        ...
+
+    # =========================================================================
+    # Core Memory Operations
+    # =========================================================================
+
+    @abstractmethod
+    async def retain_batch_async(
+        self,
+        bank_id: str,
+        contents: list[dict[str, Any]],
+        *,
+        request_context: "RequestContext",
+    ) -> dict[str, Any]:
+        """
+        Retain a batch of memory items.
+
+        Args:
+            bank_id: The memory bank ID.
+            contents: List of content dicts with 'content', optional 'event_date',
+                     'context', 'metadata', 'document_id'.
+            request_context: Request context for authentication.
+
+        Returns:
+            Dict with processing results.
+        """
+        ...
+
+    @abstractmethod
+    async def recall_async(
+        self,
+        bank_id: str,
+        query: str,
+        *,
+        budget: "Budget | None" = None,
+        max_tokens: int = 4096,
+        enable_trace: bool = False,
+        fact_type: list[str] | None = None,
+        question_date: datetime | None = None,
+        include_entities: bool = False,
+        max_entity_tokens: int = 500,
+        include_chunks: bool = False,
+        max_chunk_tokens: int = 8192,
+        request_context: "RequestContext",
+    ) -> "RecallResult":
+        """
+        Recall memories relevant to a query.
+
+        Args:
+            bank_id: The memory bank ID.
+            query: The search query.
+            budget: Search budget (LOW, MID, HIGH).
+            max_tokens: Maximum tokens in response.
+            enable_trace: Include trace information.
+            fact_type: Filter by fact types.
+            question_date: Context date for temporal relevance.
+            include_entities: Include entity observations.
+            max_entity_tokens: Max tokens for entity observations.
+            include_chunks: Include raw chunks.
+            max_chunk_tokens: Max tokens for chunks.
+            request_context: Request context for authentication.
+
+        Returns:
+            RecallResult with matching memories.
+        """
+        ...
+
+    @abstractmethod
+    async def reflect_async(
+        self,
+        bank_id: str,
+        query: str,
+        *,
+        budget: "Budget | None" = None,
+        context: str | None = None,
+        max_tokens: int = 4096,
+        response_schema: dict | None = None,
+        request_context: "RequestContext",
+    ) -> "ReflectResult":
+        """
+        Reflect on a query and generate a thoughtful response.
+
+        Args:
+            bank_id: The memory bank ID.
+            query: The question to reflect on.
+            budget: Search budget for retrieving context.
+            context: Additional context for the reflection.
+            max_tokens: Maximum tokens for the response.
+            response_schema: Optional JSON Schema for structured output.
+            request_context: Request context for authentication.
+
+        Returns:
+            ReflectResult with generated response and supporting facts.
+        """
+        ...
+
+    # =========================================================================
+    # Bank Management
+    # =========================================================================
+
+    @abstractmethod
+    async def list_banks(
+        self,
+        *,
+        request_context: "RequestContext",
+    ) -> list[dict[str, Any]]:
+        """
+        List all memory banks.
+
+        Args:
+            request_context: Request context for authentication.
+
+        Returns:
+            List of bank info dicts.
+        """
+        ...
+
+    @abstractmethod
+    async def get_bank_profile(
+        self,
+        bank_id: str,
+        *,
+        request_context: "RequestContext",
+    ) -> dict[str, Any]:
+        """
+        Get bank profile including disposition and background.
+
+        Args:
+            bank_id: The memory bank ID.
+            request_context: Request context for authentication.
+
+        Returns:
+            Bank profile dict.
+        """
+        ...
+
+    @abstractmethod
+    async def update_bank_disposition(
+        self,
+        bank_id: str,
+        disposition: dict[str, int],
+        *,
+        request_context: "RequestContext",
+    ) -> None:
+        """
+        Update bank disposition traits.
+
+        Args:
+            bank_id: The memory bank ID.
+            disposition: Dict with trait values.
+            request_context: Request context for authentication.
+        """
+        ...
+
+    @abstractmethod
+    async def merge_bank_background(
+        self,
+        bank_id: str,
+        new_info: str,
+        *,
+        update_disposition: bool = True,
+        request_context: "RequestContext",
+    ) -> dict[str, Any]:
+        """
+        Merge new background information into bank profile.
+
+        Args:
+            bank_id: The memory bank ID.
+            new_info: New background information to merge.
+            update_disposition: Whether to infer disposition from background.
+            request_context: Request context for authentication.
+
+        Returns:
+            Updated background info.
+        """
+        ...
+
+    @abstractmethod
+    async def delete_bank(
+        self,
+        bank_id: str,
+        *,
+        fact_type: str | None = None,
+        request_context: "RequestContext",
+    ) -> dict[str, int]:
+        """
+        Delete a bank or its memories.
+
+        Args:
+            bank_id: The memory bank ID.
+            fact_type: If specified, only delete memories of this type.
+            request_context: Request context for authentication.
+
+        Returns:
+            Dict with deletion counts.
+        """
+        ...
+
+    # =========================================================================
+    # Memory Units
+    # =========================================================================
+
+    @abstractmethod
+    async def list_memory_units(
+        self,
+        bank_id: str,
+        *,
+        fact_type: str | None = None,
+        search_query: str | None = None,
+        limit: int = 100,
+        offset: int = 0,
+        request_context: "RequestContext",
+    ) -> dict[str, Any]:
+        """
+        List memory units with pagination.
+
+        Args:
+            bank_id: The memory bank ID.
+            fact_type: Filter by fact type.
+            search_query: Full-text search query.
+            limit: Maximum results.
+            offset: Pagination offset.
+            request_context: Request context for authentication.
+
+        Returns:
+            Dict with 'items', 'total', 'limit', 'offset'.
+        """
+        ...
+
+    @abstractmethod
+    async def delete_memory_unit(
+        self,
+        unit_id: str,
+        *,
+        request_context: "RequestContext",
+    ) -> dict[str, Any]:
+        """
+        Delete a specific memory unit.
+
+        Args:
+            unit_id: The memory unit ID.
+            request_context: Request context for authentication.
+
+        Returns:
+            Deletion result.
+        """
+        ...
+
+    @abstractmethod
+    async def get_graph_data(
+        self,
+        bank_id: str,
+        *,
+        fact_type: str | None = None,
+        limit: int = 1000,
+        request_context: "RequestContext",
+    ) -> dict[str, Any]:
+        """
+        Get graph data for visualization.
+
+        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, limit.
+        """
+        ...
+
+    # =========================================================================
+    # Documents
+    # =========================================================================
+
+    @abstractmethod
+    async def list_documents(
+        self,
+        bank_id: str,
+        *,
+        search_query: str | None = None,
+        limit: int = 100,
+        offset: int = 0,
+        request_context: "RequestContext",
+    ) -> dict[str, Any]:
+        """
+        List documents with pagination.
+
+        Args:
+            bank_id: The memory bank ID.
+            search_query: Search query.
+            limit: Maximum results.
+            offset: Pagination offset.
+            request_context: Request context for authentication.
+
+        Returns:
+            Dict with 'items', 'total', 'limit', 'offset'.
+        """
+        ...
+
+    @abstractmethod
+    async def get_document(
+        self,
+        document_id: str,
+        bank_id: str,
+        *,
+        request_context: "RequestContext",
+    ) -> dict[str, Any] | None:
+        """
+        Get a specific document.
+
+        Args:
+            document_id: The document ID.
+            bank_id: The memory bank ID.
+            request_context: Request context for authentication.
+
+        Returns:
+            Document dict or None if not found.
+        """
+        ...
+
+    @abstractmethod
+    async def delete_document(
+        self,
+        document_id: str,
+        bank_id: str,
+        *,
+        request_context: "RequestContext",
+    ) -> dict[str, int]:
+        """
+        Delete a document and its memory units.
+
+        Args:
+            document_id: The document ID.
+            bank_id: The memory bank ID.
+            request_context: Request context for authentication.
+
+        Returns:
+            Dict with deletion counts.
+        """
+        ...
+
+    @abstractmethod
+    async def get_chunk(
+        self,
+        chunk_id: str,
+        *,
+        request_context: "RequestContext",
+    ) -> dict[str, Any] | None:
+        """
+        Get a specific chunk.
+
+        Args:
+            chunk_id: The chunk ID.
+            request_context: Request context for authentication.
+
+        Returns:
+            Chunk dict or None if not found.
+        """
+        ...
+
+    # =========================================================================
+    # Entities
+    # =========================================================================
+
+    @abstractmethod
+    async def list_entities(
+        self,
+        bank_id: str,
+        *,
+        limit: int = 100,
+        request_context: "RequestContext",
+    ) -> list[dict[str, Any]]:
+        """
+        List entities for a bank.
+
+        Args:
+            bank_id: The memory bank ID.
+            limit: Maximum results.
+            request_context: Request context for authentication.
+
+        Returns:
+            List of entity dicts.
+        """
+        ...
+
+    @abstractmethod
+    async def get_entity_observations(
+        self,
+        bank_id: str,
+        entity_id: str,
+        *,
+        limit: int = 10,
+        request_context: "RequestContext",
+    ) -> list[Any]:
+        """
+        Get observations for an entity.
+
+        Args:
+            bank_id: The memory bank ID.
+            entity_id: The entity ID.
+            limit: Maximum observations.
+            request_context: Request context for authentication.
+
+        Returns:
+            List of EntityObservation objects.
+        """
+        ...
+
+    @abstractmethod
+    async def regenerate_entity_observations(
+        self,
+        bank_id: str,
+        entity_id: str,
+        entity_name: str,
+        *,
+        request_context: "RequestContext",
+    ) -> None:
+        """
+        Regenerate observations for an entity.
+
+        Args:
+            bank_id: The memory bank ID.
+            entity_id: The entity ID.
+            entity_name: The entity's canonical name.
+            request_context: Request context for authentication.
+        """
+        ...
+
+    # =========================================================================
+    # Statistics & Operations
+    # =========================================================================
+
+    @abstractmethod
+    async def get_bank_stats(
+        self,
+        bank_id: str,
+        *,
+        request_context: "RequestContext",
+    ) -> dict[str, Any]:
+        """
+        Get statistics about memory nodes and links for a bank.
+
+        Args:
+            bank_id: The memory bank ID.
+            request_context: Request context for authentication.
+
+        Returns:
+            Dict with node_counts, link_counts, link_counts_by_fact_type,
+            link_breakdown, and operations stats.
+        """
+        ...
+
+    @abstractmethod
+    async def get_entity(
+        self,
+        bank_id: str,
+        entity_id: str,
+        *,
+        request_context: "RequestContext",
+    ) -> dict[str, Any] | None:
+        """
+        Get entity details including metadata and observations.
+
+        Args:
+            bank_id: The memory bank ID.
+            entity_id: The entity ID.
+            request_context: Request context for authentication.
+
+        Returns:
+            Entity dict with id, canonical_name, mention_count, first_seen,
+            last_seen, metadata, and observations. None if not found.
+        """
+        ...
+
+    @abstractmethod
+    async def list_operations(
+        self,
+        bank_id: str,
+        *,
+        request_context: "RequestContext",
+    ) -> list[dict[str, Any]]:
+        """
+        List async operations for a bank.
+
+        Args:
+            bank_id: The memory bank ID.
+            request_context: Request context for authentication.
+
+        Returns:
+            List of operation dicts with id, task_type, status, etc.
+        """
+        ...
+
+    @abstractmethod
+    async def cancel_operation(
+        self,
+        bank_id: str,
+        operation_id: str,
+        *,
+        request_context: "RequestContext",
+    ) -> dict[str, Any]:
+        """
+        Cancel a pending async operation.
+
+        Args:
+            bank_id: The memory bank ID.
+            operation_id: The operation ID to cancel.
+            request_context: Request context for authentication.
+
+        Returns:
+            Dict with success status and message.
+
+        Raises:
+            ValueError: If operation not found.
+        """
+        ...
+
+    @abstractmethod
+    async def update_bank(
+        self,
+        bank_id: str,
+        *,
+        name: str | None = None,
+        background: str | None = None,
+        request_context: "RequestContext",
+    ) -> dict[str, Any]:
+        """
+        Update bank name and/or background.
+
+        Args:
+            bank_id: The memory bank ID.
+            name: New bank name (optional).
+            background: New background text (optional, replaces existing).
+            request_context: Request context for authentication.
+
+        Returns:
+            Updated bank profile dict.
+        """
+        ...
+
+    @abstractmethod
+    async def submit_async_retain(
+        self,
+        bank_id: str,
+        contents: list[dict[str, Any]],
+        *,
+        request_context: "RequestContext",
+    ) -> dict[str, Any]:
+        """
+        Submit a batch retain operation to run asynchronously.
+
+        Args:
+            bank_id: The memory bank ID.
+            contents: List of content dicts to retain.
+            request_context: Request context for authentication.
+
+        Returns:
+            Dict with operation_id and items_count.
+        """
+        ...

hindsight-api/hindsight_api/engine/llm_wrapper.py

@@ -3,16 +3,29 @@ LLM wrapper for unified configuration across providers.
 """
 
 import asyncio
+import json
 import logging
 import os
+import re
 import time
 from typing import Any
 
+import httpx
 from google import genai
 from google.genai import errors as genai_errors
 from google.genai import types as genai_types
 from openai import APIConnectionError, APIStatusError, AsyncOpenAI, LengthFinishReasonError
 
+from ..config import (
+    DEFAULT_LLM_MAX_CONCURRENT,
+    DEFAULT_LLM_TIMEOUT,
+    ENV_LLM_GROQ_SERVICE_TIER,
+    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
 
@@ -22,7 +35,9 @@ logger = logging.getLogger(__name__)
 logging.getLogger("httpx").setLevel(logging.WARNING)
 
 # Global semaphore to limit concurrent LLM requests across all instances
-_global_llm_semaphore = asyncio.Semaphore(32)
+# Set HINDSIGHT_API_LLM_MAX_CONCURRENT=1 for local LLMs (LM Studio, Ollama)
+_llm_max_concurrent = int(os.getenv(ENV_LLM_MAX_CONCURRENT, str(DEFAULT_LLM_MAX_CONCURRENT)))
+_global_llm_semaphore = asyncio.Semaphore(_llm_max_concurrent)
 
 
 class OutputTooLongError(Exception):
@@ -51,53 +66,87 @@ class LLMProvider:
         base_url: str,
         model: str,
         reasoning_effort: str = "low",
+        groq_service_tier: str | None = None,
     ):
         """
         Initialize LLM provider.
 
         Args:
-            provider: Provider name ("openai", "groq", "ollama", "gemini").
+            provider: Provider name ("openai", "groq", "ollama", "gemini", "anthropic", "lmstudio").
             api_key: API key.
             base_url: Base URL for the API.
             model: Model name.
             reasoning_effort: Reasoning effort level for supported providers.
+            groq_service_tier: Groq service tier ("on_demand", "flex", "auto"). Default: None (uses Groq's default).
         """
         self.provider = provider.lower()
         self.api_key = api_key
         self.base_url = base_url
         self.model = model
         self.reasoning_effort = reasoning_effort
+        # Default to 'auto' for best performance, users can override to 'on_demand' for free tier
+        self.groq_service_tier = groq_service_tier or os.getenv(ENV_LLM_GROQ_SERVICE_TIER, "auto")
 
         # Validate provider
-        valid_providers = ["openai", "groq", "ollama", "gemini"]
+        valid_providers = ["openai", "groq", "ollama", "gemini", "anthropic", "lmstudio", "mock"]
         if self.provider not in valid_providers:
             raise ValueError(f"Invalid LLM provider: {self.provider}. Must be one of: {', '.join(valid_providers)}")
 
+        # Mock provider tracking (for testing)
+        self._mock_calls: list[dict] = []
+        self._mock_response: Any = None
+
         # Set default base URLs
         if not self.base_url:
             if self.provider == "groq":
                 self.base_url = "https://api.groq.com/openai/v1"
             elif self.provider == "ollama":
                 self.base_url = "http://localhost:11434/v1"
+            elif self.provider == "lmstudio":
+                self.base_url = "http://localhost:1234/v1"
 
-        # Validate API key (not needed for ollama)
-        if self.provider != "ollama" and not self.api_key:
+        # Validate API key (not needed for ollama, lmstudio, or mock)
+        if self.provider not in ("ollama", "lmstudio", "mock") and not self.api_key:
             raise ValueError(f"API key not found for {self.provider}")
 
+        # Get timeout config (set HINDSIGHT_API_LLM_TIMEOUT for local LLMs that need longer timeouts)
+        self.timeout = float(os.getenv(ENV_LLM_TIMEOUT, str(DEFAULT_LLM_TIMEOUT)))
+
         # Create client based on provider
-        if self.provider == "gemini":
+        self._client = None
+        self._gemini_client = None
+        self._anthropic_client = None
+
+        if self.provider == "mock":
+            # Mock provider - no client needed
+            pass
+        elif self.provider == "gemini":
             self._gemini_client = genai.Client(api_key=self.api_key)
-            self._client = None
-        elif self.provider == "ollama":
-            self._client = AsyncOpenAI(api_key="ollama", base_url=self.base_url, max_retries=0)
-            self._gemini_client = None
+        elif self.provider == "anthropic":
+            from anthropic import AsyncAnthropic
+
+            # Only pass base_url if it's set (Anthropic uses default URL otherwise)
+            anthropic_kwargs = {"api_key": self.api_key}
+            if self.base_url:
+                anthropic_kwargs["base_url"] = self.base_url
+            if self.timeout:
+                anthropic_kwargs["timeout"] = self.timeout
+            self._anthropic_client = AsyncAnthropic(**anthropic_kwargs)
+        elif self.provider in ("ollama", "lmstudio"):
+            # Use dummy key if not provided for local
+            api_key = self.api_key or "local"
+            client_kwargs = {"api_key": api_key, "base_url": self.base_url, "max_retries": 0}
+            if self.timeout:
+                client_kwargs["timeout"] = self.timeout
+            self._client = AsyncOpenAI(**client_kwargs)
         else:
             # Only pass base_url if it's set (OpenAI uses default URL otherwise)
             client_kwargs = {"api_key": self.api_key, "max_retries": 0}
             if self.base_url:
                 client_kwargs["base_url"] = self.base_url
+            if self.timeout:
+                client_kwargs["timeout"] = self.timeout
             self._client = AsyncOpenAI(**client_kwargs)
-            self._gemini_client = None
 
     async def verify_connection(self) -> None:
         """
@@ -133,6 +182,8 @@ class LLMProvider:
         initial_backoff: float = 1.0,
         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.
@@ -147,9 +198,12 @@ class LLMProvider:
             initial_backoff: Initial backoff time in seconds.
             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.
@@ -157,12 +211,56 @@ class LLMProvider:
         """
         async with _global_llm_semaphore:
             start_time = time.time()
-            import json
+
+            # Handle Mock provider (for testing)
+            if self.provider == "mock":
+                return await self._call_mock(
+                    messages,
+                    response_format,
+                    scope,
+                    return_usage,
+                )
 
             # 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
+            if self.provider == "anthropic":
+                return await self._call_anthropic(
+                    messages,
+                    response_format,
+                    max_completion_tokens,
+                    max_retries,
+                    initial_backoff,
+                    max_backoff,
+                    skip_validation,
+                    start_time,
+                    return_usage,
+                )
+
+            # Handle Ollama with native API for structured output (better schema enforcement)
+            if self.provider == "ollama" and response_format is not None:
+                return await self._call_ollama_native(
+                    messages,
+                    response_format,
+                    max_completion_tokens,
+                    temperature,
+                    max_retries,
+                    initial_backoff,
+                    max_backoff,
+                    skip_validation,
+                    start_time,
+                    return_usage,
                 )
 
             call_params = {
@@ -200,34 +298,108 @@ class LLMProvider:
             # Provider-specific parameters
             if self.provider == "groq":
                 call_params["seed"] = DEFAULT_LLM_SEED
-                extra_body = {"service_tier": "auto"}
-                # Only add reasoning parameters for reasoning models
+                extra_body: dict[str, Any] = {}
+                # Add service_tier if configured (requires paid plan for flex/auto)
+                if self.groq_service_tier:
+                    extra_body["service_tier"] = self.groq_service_tier
+                # Add reasoning parameters for reasoning models
                 if is_reasoning_model:
                     extra_body["include_reasoning"] = False
-                call_params["extra_body"] = extra_body
+                if extra_body:
+                    call_params["extra_body"] = extra_body
 
             last_exception = None
 
             for attempt in range(max_retries + 1):
                 try:
                     if response_format is not None:
-                        # Add schema to system message for JSON mode
+                        schema = None
                         if hasattr(response_format, "model_json_schema"):
                             schema = response_format.model_json_schema()
-                            schema_msg = f"\n\nYou must respond with valid JSON matching this schema:\n{json.dumps(schema, indent=2)}"
 
-                            if call_params["messages"] and call_params["messages"][0].get("role") == "system":
-                                call_params["messages"][0]["content"] += schema_msg
-                            elif call_params["messages"]:
-                                call_params["messages"][0]["content"] = (
-                                    schema_msg + "\n\n" + call_params["messages"][0]["content"]
-                                )
+                        if strict_schema and schema is not None:
+                            # Use OpenAI's strict JSON schema enforcement
+                            # This guarantees all required fields are returned
+                            call_params["response_format"] = {
+                                "type": "json_schema",
+                                "json_schema": {
+                                    "name": "response",
+                                    "strict": True,
+                                    "schema": schema,
+                                },
+                            }
+                        else:
+                            # Soft enforcement: add schema to prompt and use json_object mode
+                            if schema is not None:
+                                schema_msg = f"\n\nYou must respond with valid JSON matching this schema:\n{json.dumps(schema, indent=2)}"
 
-                        call_params["response_format"] = {"type": "json_object"}
+                                if call_params["messages"] and call_params["messages"][0].get("role") == "system":
+                                    call_params["messages"][0]["content"] += schema_msg
+                                elif call_params["messages"]:
+                                    call_params["messages"][0]["content"] = (
+                                        schema_msg + "\n\n" + call_params["messages"][0]["content"]
+                                    )
+                            if self.provider not in ("lmstudio", "ollama"):
+                                # LM Studio and Ollama don't support json_object response format reliably
+                                # We rely on the schema in the system message instead
+                                call_params["response_format"] = {"type": "json_object"}
+
+                        logger.debug(f"Sending request to {self.provider}/{self.model} (timeout={self.timeout})")
                         response = await self._client.chat.completions.create(**call_params)
+                        logger.debug(f"Received response from {self.provider}/{self.model}")
 
                         content = response.choices[0].message.content
-                        json_data = json.loads(content)
+
+                        # Strip reasoning model thinking tags
+                        # Supports: <think>, <thinking>, <reasoning>, |startthink|/|endthink|
+                        # for reasoning models that embed thinking in their output (e.g., Qwen3, DeepSeek)
+                        if content:
+                            original_len = len(content)
+                            content = re.sub(r"<think>.*?</think>", "", content, flags=re.DOTALL)
+                            content = re.sub(r"<thinking>.*?</thinking>", "", content, flags=re.DOTALL)
+                            content = re.sub(r"<reasoning>.*?</reasoning>", "", content, flags=re.DOTALL)
+                            content = re.sub(r"\|startthink\|.*?\|endthink\|", "", content, flags=re.DOTALL)
+                            content = content.strip()
+                            if len(content) < original_len:
+                                logger.debug(f"Stripped {original_len - len(content)} chars of reasoning tokens")
+
+                        # For local models, they may wrap JSON in markdown code blocks
+                        if self.provider in ("lmstudio", "ollama"):
+                            clean_content = content
+                            if "```json" in content:
+                                clean_content = content.split("```json")[1].split("```")[0].strip()
+                            elif "```" in content:
+                                clean_content = content.split("```")[1].split("```")[0].strip()
+                            try:
+                                json_data = json.loads(clean_content)
+                            except json.JSONDecodeError:
+                                # Fallback to parsing raw content
+                                json_data = json.loads(content)
+                        else:
+                            # Log raw LLM response for debugging JSON parse issues
+                            try:
+                                json_data = json.loads(content)
+                            except json.JSONDecodeError as json_err:
+                                # Truncate content for logging (first 500 and last 200 chars)
+                                content_preview = content[:500] if content else "<empty>"
+                                if content and len(content) > 700:
+                                    content_preview = f"{content[:500]}...TRUNCATED...{content[-200:]}"
+                                logger.warning(
+                                    f"JSON parse error from LLM response (attempt {attempt + 1}/{max_retries + 1}): {json_err}\n"
+                                    f"  Model: {self.provider}/{self.model}\n"
+                                    f"  Content length: {len(content) if content else 0} chars\n"
+                                    f"  Content preview: {content_preview!r}\n"
+                                    f"  Finish reason: {response.choices[0].finish_reason if response.choices else 'unknown'}"
+                                )
+                                # Retry on JSON parse errors - LLM may return valid JSON on next attempt
+                                if attempt < max_retries:
+                                    backoff = min(initial_backoff * (2**attempt), max_backoff)
+                                    await asyncio.sleep(backoff)
+                                    last_exception = json_err
+                                    continue
+                                else:
+                                    logger.error(f"JSON parse error after {max_retries + 1} attempts, giving up")
+                                    raise
 
                         if skip_validation:
                             result = json_data
@@ -237,21 +409,45 @@ 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
+
+                    # Record LLM metrics
+                    metrics = get_metrics_collector()
+                    metrics.record_llm_call(
+                        provider=self.provider,
+                        model=self.model,
+                        scope=scope,
+                        duration=duration,
+                        input_tokens=input_tokens,
+                        output_tokens=output_tokens,
+                        success=True,
+                    )
+
+                    # 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:
@@ -300,6 +496,317 @@ class LLMProvider:
                 raise last_exception
             raise RuntimeError("LLM call failed after all retries with no exception captured")
 
+    async def _call_anthropic(
+        self,
+        messages: list[dict[str, str]],
+        response_format: Any | None,
+        max_completion_tokens: int | None,
+        max_retries: int,
+        initial_backoff: float,
+        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
+
+        # Convert OpenAI-style messages to Anthropic format
+        system_prompt = None
+        anthropic_messages = []
+
+        for msg in messages:
+            role = msg.get("role", "user")
+            content = msg.get("content", "")
+
+            if role == "system":
+                if system_prompt:
+                    system_prompt += "\n\n" + content
+                else:
+                    system_prompt = content
+            else:
+                anthropic_messages.append({"role": role, "content": content})
+
+        # Add JSON schema instruction if response_format is provided
+        if response_format is not None and hasattr(response_format, "model_json_schema"):
+            schema = response_format.model_json_schema()
+            schema_msg = f"\n\nYou must respond with valid JSON matching this schema:\n{json.dumps(schema, indent=2)}"
+            if system_prompt:
+                system_prompt += schema_msg
+            else:
+                system_prompt = schema_msg
+
+        # Prepare parameters
+        call_params = {
+            "model": self.model,
+            "messages": anthropic_messages,
+            "max_tokens": max_completion_tokens if max_completion_tokens is not None else 4096,
+        }
+
+        if system_prompt:
+            call_params["system"] = system_prompt
+
+        last_exception = None
+
+        for attempt in range(max_retries + 1):
+            try:
+                response = await self._anthropic_client.messages.create(**call_params)
+
+                # Anthropic response content is a list of blocks
+                content = ""
+                for block in response.content:
+                    if block.type == "text":
+                        content += block.text
+
+                if response_format is not None:
+                    # Models may wrap JSON in markdown code blocks
+                    clean_content = content
+                    if "```json" in content:
+                        clean_content = content.split("```json")[1].split("```")[0].strip()
+                    elif "```" in content:
+                        clean_content = content.split("```")[1].split("```")[0].strip()
+
+                    try:
+                        json_data = json.loads(clean_content)
+                    except json.JSONDecodeError:
+                        # Fallback to parsing raw content if markdown stripping failed
+                        json_data = json.loads(content)
+
+                    if skip_validation:
+                        result = json_data
+                    else:
+                        result = response_format.model_validate(json_data)
+                else:
+                    result = content
+
+                # Record metrics and log slow calls
+                duration = time.time() - start_time
+                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
+
+                # Record LLM metrics
+                metrics = get_metrics_collector()
+                metrics.record_llm_call(
+                    provider=self.provider,
+                    model=self.model,
+                    scope="memory",
+                    duration=duration,
+                    input_tokens=input_tokens,
+                    output_tokens=output_tokens,
+                    success=True,
+                )
+
+                # Log slow calls
+                if duration > 10.0:
+                    logger.info(
+                        f"slow llm call: scope=memory, 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:
+                last_exception = e
+                if attempt < max_retries:
+                    logger.warning("Anthropic returned invalid JSON, retrying...")
+                    backoff = min(initial_backoff * (2**attempt), max_backoff)
+                    await asyncio.sleep(backoff)
+                    continue
+                else:
+                    logger.error(f"Anthropic returned invalid JSON after {max_retries + 1} attempts")
+                    raise
+
+            except (APIConnectionError, RateLimitError, APIStatusError) as e:
+                # Fast fail on 401/403
+                if isinstance(e, APIStatusError) and e.status_code in (401, 403):
+                    logger.error(f"Anthropic auth error (HTTP {e.status_code}), not retrying: {str(e)}")
+                    raise
+
+                last_exception = e
+                if attempt < max_retries:
+                    # Check if it's a rate limit or server error
+                    should_retry = isinstance(e, (APIConnectionError, RateLimitError)) or (
+                        isinstance(e, APIStatusError) and e.status_code >= 500
+                    )
+
+                    if should_retry:
+                        backoff = min(initial_backoff * (2**attempt), max_backoff)
+                        jitter = backoff * 0.2 * (2 * (time.time() % 1) - 1)
+                        await asyncio.sleep(backoff + jitter)
+                        continue
+
+                logger.error(f"Anthropic API error after {max_retries + 1} attempts: {str(e)}")
+                raise
+
+            except Exception as e:
+                logger.error(f"Unexpected error during Anthropic call: {type(e).__name__}: {str(e)}")
+                raise
+
+        if last_exception:
+            raise last_exception
+        raise RuntimeError("Anthropic call failed after all retries")
+
+    async def _call_ollama_native(
+        self,
+        messages: list[dict[str, str]],
+        response_format: Any,
+        max_completion_tokens: int | None,
+        temperature: float | None,
+        max_retries: int,
+        initial_backoff: float,
+        max_backoff: float,
+        skip_validation: bool,
+        start_time: float,
+        return_usage: bool = False,
+    ) -> Any:
+        """
+        Call Ollama using native API with JSON schema enforcement.
+
+        Ollama's native API supports passing a full JSON schema in the 'format' parameter,
+        which provides better structured output control than the OpenAI-compatible API.
+        """
+        # Get the JSON schema from the Pydantic model
+        schema = response_format.model_json_schema() if hasattr(response_format, "model_json_schema") else None
+
+        # Build the base URL for Ollama's native API
+        # Default OpenAI-compatible URL is http://localhost:11434/v1
+        # Native API is at http://localhost:11434/api/chat
+        base_url = self.base_url or "http://localhost:11434/v1"
+        if base_url.endswith("/v1"):
+            native_url = base_url[:-3] + "/api/chat"
+        else:
+            native_url = base_url.rstrip("/") + "/api/chat"
+
+        # Build request payload
+        payload = {
+            "model": self.model,
+            "messages": messages,
+            "stream": False,
+        }
+
+        # Add schema as format parameter for structured output
+        if schema:
+            payload["format"] = schema
+
+        # Add optional parameters with optimized defaults for Ollama
+        # Benchmarking shows num_ctx=16384 + num_batch=512 is optimal
+        options = {
+            "num_ctx": 16384,  # 16k context window for larger prompts
+            "num_batch": 512,  # Optimal batch size for prompt processing
+        }
+        if max_completion_tokens:
+            options["num_predict"] = max_completion_tokens
+        if temperature is not None:
+            options["temperature"] = temperature
+        payload["options"] = options
+
+        last_exception = None
+
+        async with httpx.AsyncClient(timeout=300.0) as client:
+            for attempt in range(max_retries + 1):
+                try:
+                    response = await client.post(native_url, json=payload)
+                    response.raise_for_status()
+
+                    result = response.json()
+                    content = result.get("message", {}).get("content", "")
+
+                    # Parse JSON response
+                    try:
+                        json_data = json.loads(content)
+                    except json.JSONDecodeError as json_err:
+                        content_preview = content[:500] if content else "<empty>"
+                        if content and len(content) > 700:
+                            content_preview = f"{content[:500]}...TRUNCATED...{content[-200:]}"
+                        logger.warning(
+                            f"Ollama JSON parse error (attempt {attempt + 1}/{max_retries + 1}): {json_err}\n"
+                            f"  Model: ollama/{self.model}\n"
+                            f"  Content length: {len(content) if content else 0} chars\n"
+                            f"  Content preview: {content_preview!r}"
+                        )
+                        if attempt < max_retries:
+                            backoff = min(initial_backoff * (2**attempt), max_backoff)
+                            await asyncio.sleep(backoff)
+                            last_exception = json_err
+                            continue
+                        else:
+                            raise
+
+                    # Extract token usage from Ollama response
+                    # Ollama returns prompt_eval_count (input) and eval_count (output)
+                    duration = time.time() - start_time
+                    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 LLM metrics
+                    metrics = get_metrics_collector()
+                    metrics.record_llm_call(
+                        provider=self.provider,
+                        model=self.model,
+                        scope="memory",
+                        duration=duration,
+                        input_tokens=input_tokens,
+                        output_tokens=output_tokens,
+                        success=True,
+                    )
+
+                    # Validate against Pydantic model or return raw JSON
+                    if skip_validation:
+                        validated_result = json_data
+                    else:
+                        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
+                    if attempt < max_retries:
+                        logger.warning(
+                            f"Ollama HTTP error (attempt {attempt + 1}/{max_retries + 1}): {e.response.status_code}"
+                        )
+                        backoff = min(initial_backoff * (2**attempt), max_backoff)
+                        await asyncio.sleep(backoff)
+                        continue
+                    else:
+                        logger.error(f"Ollama HTTP error after {max_retries + 1} attempts: {e}")
+                        raise
+
+                except httpx.RequestError as e:
+                    last_exception = e
+                    if attempt < max_retries:
+                        logger.warning(f"Ollama connection error (attempt {attempt + 1}/{max_retries + 1}): {e}")
+                        backoff = min(initial_backoff * (2**attempt), max_backoff)
+                        await asyncio.sleep(backoff)
+                        continue
+                    else:
+                        logger.error(f"Ollama connection error after {max_retries + 1} attempts: {e}")
+                        raise
+
+                except Exception as e:
+                    logger.error(f"Unexpected error during Ollama call: {type(e).__name__}: {e}")
+                    raise
+
+        if last_exception:
+            raise last_exception
+        raise RuntimeError("Ollama call failed after all retries")
+
     async def _call_gemini(
         self,
         messages: list[dict[str, str]],
@@ -309,10 +816,9 @@ class LLMProvider:
         max_backoff: float,
         skip_validation: bool,
         start_time: float,
+        return_usage: bool = False,
     ) -> Any:
         """Handle Gemini-specific API calls."""
-        import json
-
         # Convert OpenAI-style messages to Gemini format
         system_instruction = None
         gemini_contents = []
@@ -387,16 +893,42 @@ class LLMProvider:
                 else:
                     result = content
 
-                # Log slow calls
+                # Record metrics and log slow calls
                 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
+                    input_tokens = usage.prompt_token_count or 0
+                    output_tokens = usage.candidates_token_count or 0
+
+                # Record LLM metrics
+                metrics = get_metrics_collector()
+                metrics.record_llm_call(
+                    provider=self.provider,
+                    model=self.model,
+                    scope="memory",
+                    duration=duration,
+                    input_tokens=input_tokens,
+                    output_tokens=output_tokens,
+                    success=True,
+                )
+
+                # Log slow calls
+                if duration > 10.0 and input_tokens > 0:
                     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"slow llm call: scope=memory, 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:
@@ -438,11 +970,68 @@ class LLMProvider:
             raise last_exception
         raise RuntimeError("Gemini call failed after all retries")
 
+    async def _call_mock(
+        self,
+        messages: list[dict[str, str]],
+        response_format: Any | None,
+        scope: str,
+        return_usage: bool,
+    ) -> Any:
+        """
+        Handle mock provider calls for testing.
+
+        Records the call and returns a configurable mock response.
+        """
+        # Record the call for test verification
+        call_record = {
+            "provider": self.provider,
+            "model": self.model,
+            "messages": messages,
+            "response_format": response_format.__name__
+            if response_format and hasattr(response_format, "__name__")
+            else str(response_format),
+            "scope": scope,
+        }
+        self._mock_calls.append(call_record)
+        logger.debug(f"Mock LLM call recorded: scope={scope}, model={self.model}")
+
+        # Return mock response
+        if self._mock_response is not None:
+            result = self._mock_response
+        elif response_format is not None:
+            # Try to create a minimal valid instance of the response format
+            try:
+                # For Pydantic models, try to create with minimal valid data
+                result = {"mock": True}
+            except Exception:
+                result = {"mock": True}
+        else:
+            result = "mock response"
+
+        if return_usage:
+            token_usage = TokenUsage(input_tokens=10, output_tokens=5, total_tokens=15)
+            return result, token_usage
+        return result
+
+    def set_mock_response(self, response: Any) -> None:
+        """Set the response to return from mock calls."""
+        self._mock_response = response
+
+    def get_mock_calls(self) -> list[dict]:
+        """Get the list of recorded mock calls."""
+        return self._mock_calls
+
+    def clear_mock_calls(self) -> None:
+        """Clear the recorded mock calls."""
+        self._mock_calls = []
+
     @classmethod
     def for_memory(cls) -> "LLMProvider":
         """Create provider for memory operations from environment variables."""
         provider = os.getenv("HINDSIGHT_API_LLM_PROVIDER", "groq")
         api_key = os.getenv("HINDSIGHT_API_LLM_API_KEY")
+        if not api_key:
+            raise ValueError("HINDSIGHT_API_LLM_API_KEY environment variable is required")
         base_url = os.getenv("HINDSIGHT_API_LLM_BASE_URL", "")
         model = os.getenv("HINDSIGHT_API_LLM_MODEL", "openai/gpt-oss-120b")
 
@@ -453,6 +1042,10 @@ class LLMProvider:
         """Create provider for answer generation. Falls back to memory config if not set."""
         provider = os.getenv("HINDSIGHT_API_ANSWER_LLM_PROVIDER", os.getenv("HINDSIGHT_API_LLM_PROVIDER", "groq"))
         api_key = os.getenv("HINDSIGHT_API_ANSWER_LLM_API_KEY", os.getenv("HINDSIGHT_API_LLM_API_KEY"))
+        if not api_key:
+            raise ValueError(
+                "HINDSIGHT_API_LLM_API_KEY or HINDSIGHT_API_ANSWER_LLM_API_KEY environment variable is required"
+            )
         base_url = os.getenv("HINDSIGHT_API_ANSWER_LLM_BASE_URL", os.getenv("HINDSIGHT_API_LLM_BASE_URL", ""))
         model = os.getenv("HINDSIGHT_API_ANSWER_LLM_MODEL", os.getenv("HINDSIGHT_API_LLM_MODEL", "openai/gpt-oss-120b"))
 
@@ -463,6 +1056,10 @@ class LLMProvider:
         """Create provider for judge/evaluator operations. Falls back to memory config if not set."""
         provider = os.getenv("HINDSIGHT_API_JUDGE_LLM_PROVIDER", os.getenv("HINDSIGHT_API_LLM_PROVIDER", "groq"))
         api_key = os.getenv("HINDSIGHT_API_JUDGE_LLM_API_KEY", os.getenv("HINDSIGHT_API_LLM_API_KEY"))
+        if not api_key:
+            raise ValueError(
+                "HINDSIGHT_API_LLM_API_KEY or HINDSIGHT_API_JUDGE_LLM_API_KEY environment variable is required"
+            )
         base_url = os.getenv("HINDSIGHT_API_JUDGE_LLM_BASE_URL", os.getenv("HINDSIGHT_API_LLM_BASE_URL", ""))
         model = os.getenv("HINDSIGHT_API_JUDGE_LLM_MODEL", os.getenv("HINDSIGHT_API_LLM_MODEL", "openai/gpt-oss-120b"))
 

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.
@@ -123,7 +154,8 @@ class ReflectResult(BaseModel):
     Result from a reflect operation.
 
     Contains the formulated answer, the facts it was based on (organized by type),
-    and any new opinions that were formed during the reflection process.
+    any new opinions that were formed during the reflection process, and optionally
+    structured output if a response schema was provided.
     """
 
     model_config = ConfigDict(
@@ -145,6 +177,8 @@ class ReflectResult(BaseModel):
                     "opinion": [],
                 },
                 "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},
             }
         }
     )
@@ -154,6 +188,14 @@ class ReflectResult(BaseModel):
         description="Facts used to formulate the answer, organized by type (world, experience, opinion)"
     )
     new_opinions: list[str] = Field(default_factory=list, description="List of newly formed opinions during reflection")
+    structured_output: dict[str, Any] | None = Field(
+        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/bank_utils.py

@@ -10,6 +10,7 @@ from typing import TypedDict
 from pydantic import BaseModel, Field
 
 from ..db_utils import acquire_with_retry
+from ..memory_engine import fq_table
 from ..response_models import DispositionTraits
 
 logger = logging.getLogger(__name__)
@@ -51,9 +52,9 @@ async def get_bank_profile(pool, bank_id: str) -> BankProfile:
     async with acquire_with_retry(pool) as conn:
         # Try to get existing bank
         row = await conn.fetchrow(
-            """
+            f"""
             SELECT name, disposition, background
-            FROM banks WHERE bank_id = $1
+            FROM {fq_table("banks")} WHERE bank_id = $1
             """,
             bank_id,
         )
@@ -70,8 +71,8 @@ async def get_bank_profile(pool, bank_id: str) -> BankProfile:
 
         # Bank doesn't exist, create with defaults
         await conn.execute(
-            """
-            INSERT INTO banks (bank_id, name, disposition, background)
+            f"""
+            INSERT INTO {fq_table("banks")} (bank_id, name, disposition, background)
             VALUES ($1, $2, $3::jsonb, $4)
             ON CONFLICT (bank_id) DO NOTHING
             """,
@@ -98,8 +99,8 @@ async def update_bank_disposition(pool, bank_id: str, disposition: dict[str, int
 
     async with acquire_with_retry(pool) as conn:
         await conn.execute(
-            """
-            UPDATE banks
+            f"""
+            UPDATE {fq_table("banks")}
             SET disposition = $2::jsonb,
                 updated_at = NOW()
             WHERE bank_id = $1
@@ -140,8 +141,8 @@ async def merge_bank_background(pool, llm_config, bank_id: str, new_info: str, u
         if inferred_disposition:
             # Update both background and disposition
             await conn.execute(
-                """
-                UPDATE banks
+                f"""
+                UPDATE {fq_table("banks")}
                 SET background = $2,
                     disposition = $3::jsonb,
                     updated_at = NOW()
@@ -154,8 +155,8 @@ async def merge_bank_background(pool, llm_config, bank_id: str, new_info: str, u
         else:
             # Update only background
             await conn.execute(
-                """
-                UPDATE banks
+                f"""
+                UPDATE {fq_table("banks")}
                 SET background = $2,
                     updated_at = NOW()
                 WHERE bank_id = $1
@@ -361,9 +362,9 @@ async def list_banks(pool) -> list:
     """
     async with acquire_with_retry(pool) as conn:
         rows = await conn.fetch(
-            """
+            f"""
             SELECT bank_id, name, disposition, background, created_at, updated_at
-            FROM banks
+            FROM {fq_table("banks")}
             ORDER BY updated_at DESC
             """
         )

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

@@ -6,6 +6,7 @@ Handles storage of document chunks in the database.
 
 import logging
 
+from ..memory_engine import fq_table
 from .types import ChunkMetadata
 
 logger = logging.getLogger(__name__)
@@ -42,8 +43,8 @@ async def store_chunks_batch(conn, bank_id: str, document_id: str, chunks: list[
 
     # Batch insert all chunks
     await conn.execute(
-        """
-        INSERT INTO chunks (chunk_id, document_id, bank_id, chunk_text, chunk_index)
+        f"""
+        INSERT INTO {fq_table("chunks")} (chunk_id, document_id, bank_id, chunk_text, chunk_index)
         SELECT * FROM unnest($1::text[], $2::text[], $3::text[], $4::text[], $5::integer[])
         """,
         chunk_ids,

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

@@ -13,16 +13,23 @@ logger = logging.getLogger(__name__)
 
 
 async def process_entities_batch(
-    entity_resolver, conn, bank_id: str, unit_ids: list[str], facts: list[ProcessedFact], log_buffer: list[str] = None
+    entity_resolver,
+    conn,
+    bank_id: str,
+    unit_ids: list[str],
+    facts: list[ProcessedFact],
+    log_buffer: list[str] = None,
+    user_entities_per_content: dict[int, list[dict]] = None,
 ) -> list[EntityLink]:
     """
     Process entities for all facts and create entity links.
 
     This function:
     1. Extracts entity mentions from fact texts
-    2. Resolves entity names to canonical entities
-    3. Creates entity records in the database
-    4. Returns entity links ready for insertion
+    2. Merges user-provided entities with LLM-extracted entities
+    3. Resolves entity names to canonical entities
+    4. Creates entity records in the database
+    5. Returns entity links ready for insertion
 
     Args:
         entity_resolver: EntityResolver instance for entity resolution
@@ -31,6 +38,7 @@ async def process_entities_batch(
         unit_ids: List of unit IDs (same length as facts)
         facts: List of ProcessedFact objects
         log_buffer: Optional buffer for detailed logging
+        user_entities_per_content: Dict mapping content_index to list of user-provided entities
 
     Returns:
         List of EntityLink objects for batch insertion
@@ -41,14 +49,35 @@ async def process_entities_batch(
     if len(unit_ids) != len(facts):
         raise ValueError(f"Mismatch between unit_ids ({len(unit_ids)}) and facts ({len(facts)})")
 
+    user_entities_per_content = user_entities_per_content or {}
+
     # Extract data for link_utils function
     fact_texts = [fact.fact_text for fact in facts]
     # Use occurred_start if available, otherwise use mentioned_at for entity timestamps
     fact_dates = [fact.occurred_start if fact.occurred_start is not None else fact.mentioned_at for fact in facts]
-    # Convert EntityRef objects to dict format expected by link_utils
-    entities_per_fact = [
-        [{"text": entity.name, "type": "CONCEPT"} for entity in (fact.entities or [])] for fact in facts
-    ]
+
+    # Convert EntityRef objects to dict format and merge with user-provided entities
+    entities_per_fact = []
+    for fact in facts:
+        # Start with LLM-extracted entities
+        llm_entities = [{"text": entity.name, "type": "CONCEPT"} for entity in (fact.entities or [])]
+
+        # Get user entities for this content (use content_index from fact)
+        user_entities = user_entities_per_content.get(fact.content_index, [])
+
+        # Merge with case-insensitive deduplication
+        seen_texts = {e["text"].lower() for e in llm_entities}
+        for user_entity in user_entities:
+            if user_entity["text"].lower() not in seen_texts:
+                llm_entities.append(
+                    {
+                        "text": user_entity["text"],
+                        "type": user_entity.get("type", "CONCEPT"),
+                    }
+                )
+                seen_texts.add(user_entity["text"].lower())
+
+        entities_per_fact.append(llm_entities)
 
     # Use existing link_utils function for entity processing
     entity_links = await link_utils.extract_entities_batch_optimized(

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

@@ -14,7 +14,47 @@ 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:
+    """
+    Infer a temporal date from fact text when LLM didn't provide occurred_start.
+
+    This is a fallback for when the LLM fails to extract temporal information
+    from relative time expressions like "last night", "yesterday", etc.
+    """
+    import re
+
+    fact_lower = fact_text.lower()
+
+    # Map relative time expressions to day offsets
+    temporal_patterns = {
+        r"\blast night\b": -1,
+        r"\byesterday\b": -1,
+        r"\btoday\b": 0,
+        r"\bthis morning\b": 0,
+        r"\bthis afternoon\b": 0,
+        r"\bthis evening\b": 0,
+        r"\btonigh?t\b": 0,
+        r"\btomorrow\b": 1,
+        r"\blast week\b": -7,
+        r"\bthis week\b": 0,
+        r"\bnext week\b": 7,
+        r"\blast month\b": -30,
+        r"\bthis month\b": 0,
+        r"\bnext month\b": 30,
+    }
+
+    for pattern, offset_days in temporal_patterns.items():
+        if re.search(pattern, fact_lower):
+            target_date = event_date + timedelta(days=offset_days)
+            return target_date.replace(hour=0, minute=0, second=0, microsecond=0).isoformat()
+
+    # If no relative time expression found, return None
+    return None
 
 
 def _sanitize_text(text: str) -> str:
@@ -71,22 +111,44 @@ class Fact(BaseModel):
 
 
 class CausalRelation(BaseModel):
-    """Causal relationship between facts."""
+    """Causal relationship from this fact to a previous fact (stored format)."""
 
-    target_fact_index: int = Field(
-        description="Index of the related fact in the facts array (0-based). "
-        "This creates a directed causal link to another fact in the extraction."
-    )
-    relation_type: Literal["causes", "caused_by", "enables", "prevents"] = Field(
-        description="Type of causal relationship: "
-        "'causes' = this fact directly causes the target fact, "
-        "'caused_by' = this fact was caused by the target fact, "
-        "'enables' = this fact enables/allows the target fact, "
-        "'prevents' = this fact prevents/blocks the target fact"
+    target_fact_index: int = Field(description="Index of the related fact in the facts array (0-based).")
+    relation_type: Literal["caused_by", "enabled_by", "prevented_by"] = Field(
+        description="How this fact relates to the target: "
+        "'caused_by' = this fact was caused by the target, "
+        "'enabled_by' = this fact was enabled by the target, "
+        "'prevented_by' = this fact was prevented by the target"
     )
     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",
+        description="Strength of relationship (0.0 to 1.0)",
+        ge=0.0,
+        le=1.0,
+        default=1.0,
+    )
+
+
+class FactCausalRelation(BaseModel):
+    """
+    Causal relationship from this fact to a PREVIOUS fact (embedded in each fact).
+
+    Uses index-based references but ONLY allows referencing facts that appear
+    BEFORE this fact in the list. This prevents hallucination of invalid indices.
+    """
+
+    target_index: int = Field(
+        description="Index of the PREVIOUS fact this relates to (0-based). "
+        "MUST be less than this fact's position in the list. "
+        "Example: if this is fact #5, target_index can only be 0, 1, 2, 3, or 4."
+    )
+    relation_type: Literal["caused_by", "enabled_by", "prevented_by"] = Field(
+        description="How this fact relates to the target fact: "
+        "'caused_by' = this fact was caused by the target fact, "
+        "'enabled_by' = this fact was enabled by the target fact, "
+        "'prevented_by' = this fact was blocked/prevented by the target fact"
+    )
+    strength: float = Field(
+        description="Strength of relationship (0.0 to 1.0). 1.0 = strong, 0.5 = moderate",
         ge=0.0,
         le=1.0,
         default=1.0,
@@ -176,8 +238,12 @@ class ExtractedFact(BaseModel):
         default=None,
         description="Named entities, objects, AND abstract concepts from the fact. Include: people names, organizations, places, significant objects (e.g., 'coffee maker', 'car'), AND abstract concepts/themes (e.g., 'friendship', 'career growth', 'loss', 'celebration'). Extract anything that could help link related facts together.",
     )
-    causal_relations: list[CausalRelation] | None = Field(
-        default=None, description="Causal links to other facts. Can be null."
+
+    # Causal relations to PREVIOUS facts only (prevents hallucination of invalid indices)
+    causal_relations: list[FactCausalRelation] | None = Field(
+        default=None,
+        description="Causal links to PREVIOUS facts only. target_index MUST be less than this fact's position. "
+        "Example: fact #3 can only reference facts 0, 1, or 2. Max 2 relations per fact.",
     )
 
     @field_validator("entities", mode="before")
@@ -188,14 +254,6 @@ class ExtractedFact(BaseModel):
             return []
         return v
 
-    @field_validator("causal_relations", mode="before")
-    @classmethod
-    def ensure_causal_relations_list(cls, v):
-        """Ensure causal_relations is always a list (convert None to empty list)."""
-        if v is None:
-            return []
-        return v
-
     def build_fact_text(self) -> str:
         """Combine all dimensions into a single comprehensive fact string."""
         parts = [self.what]
@@ -215,7 +273,7 @@ class ExtractedFact(BaseModel):
 
 
 class FactExtractionResponse(BaseModel):
-    """Response containing all extracted facts."""
+    """Response containing all extracted facts (causal relations are embedded in each fact)."""
 
     facts: list[ExtractedFact] = Field(description="List of extracted factual statements")
 
@@ -318,7 +376,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).
 
@@ -339,6 +397,9 @@ async def _extract_facts_from_chunk(
 
     prompt = f"""Extract facts from text into structured format with FOUR required dimensions - BE EXTREMELY DETAILED.
 
+LANGUAGE REQUIREMENT: Detect the language of the input text. All extracted facts, entity names, descriptions,
+and other output MUST be in the SAME language as the input. Do not translate to English if the input is in another language.
+
 {fact_types_instruction}
 
 
@@ -534,7 +595,52 @@ 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 (EMBEDDED IN EACH FACT - REFERENCE PREVIOUS FACTS ONLY)
+══════════════════════════════════════════════════════════════════════════
+
+Each fact can have a `causal_relations` array that links to PREVIOUS facts only.
+⚠️ CRITICAL: target_index MUST be less than this fact's position in the list!
+
+If you're writing fact #5, you can only reference facts 0, 1, 2, 3, or 4.
+This ensures all references are valid.
+
+Relationship types (all describe how THIS fact relates to the target):
+- "caused_by": This fact was caused by the target fact
+- "enabled_by": This fact was enabled/allowed by the target fact
+- "prevented_by": This fact was blocked/prevented by the target fact
+
+Max 2 causal relations per fact. Only add if there's a clear causal link.
+
+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."
+
+Output facts:
+```json
+{{
+  "facts": [
+    {{
+      "what": "User lost their job in January due to company layoffs",
+      ...other fields...
+      "causal_relations": null  // First fact - nothing to reference
+    }},
+    {{
+      "what": "User couldn't pay rent because of job loss",
+      ...other fields...
+      "causal_relations": [{{"target_index": 0, "relation_type": "caused_by", "strength": 1.0}}]
+    }},
+    {{
+      "what": "User moved to a cheaper apartment",
+      ...other fields...
+      "causal_relations": [{{"target_index": 1, "relation_type": "caused_by", "strength": 0.9}}]
+    }}
+  ]
+}}
+```
+
+This creates: Job loss (0) ← Can't pay rent (1) ← Moved apartment (2)"""
 
     import logging
 
@@ -545,6 +651,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)
@@ -563,16 +670,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 = []
@@ -590,9 +700,10 @@ 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", [])
+
             if not raw_facts:
                 logger.debug(
                     f"LLM response missing 'facts' field or returned empty list. "
@@ -676,13 +787,18 @@ Text:
                 if fact_kind == "event":
                     occurred_start = get_value("occurred_start")
                     occurred_end = get_value("occurred_end")
-                    if occurred_start:
+
+                    # If LLM didn't set temporal fields, try to extract them from the fact text
+                    if not occurred_start:
+                        fact_data["occurred_start"] = _infer_temporal_date(combined_text, event_date)
+                    else:
                         fact_data["occurred_start"] = occurred_start
-                        # For point events: if occurred_end not set, default to occurred_start
-                        if occurred_end:
-                            fact_data["occurred_end"] = occurred_end
-                        else:
-                            fact_data["occurred_end"] = occurred_start
+
+                    # For point events: if occurred_end not set, default to occurred_start
+                    if occurred_end:
+                        fact_data["occurred_end"] = occurred_end
+                    elif fact_data.get("occurred_start"):
+                        fact_data["occurred_end"] = fact_data["occurred_start"]
 
                 # Add entities if present (validate as Entity objects)
                 # LLM sometimes returns strings instead of {"text": "..."} format
@@ -702,19 +818,41 @@ 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:
-                        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
+                # Add per-fact causal relations (new schema: target_index must be < current fact index)
+                validated_relations = []
+                causal_relations_raw = get_value("causal_relations")
+                if causal_relations_raw:
+                    for rel in causal_relations_raw:
+                        if not isinstance(rel, dict):
+                            continue
+                        # New schema uses target_index
+                        target_idx = rel.get("target_index")
+                        relation_type = rel.get("relation_type")
+                        strength = rel.get("strength", 1.0)
+
+                        if target_idx is None or relation_type is None:
+                            continue
+
+                        # Validate: target_index must be < current fact index
+                        if target_idx < 0 or target_idx >= i:
+                            logger.debug(
+                                f"Invalid target_index {target_idx} for fact {i} (must be 0 to {i - 1}). Skipping."
+                            )
+                            continue
+
+                        try:
+                            validated_relations.append(
+                                CausalRelation(
+                                    target_fact_index=target_idx,
+                                    relation_type=relation_type,
+                                    strength=strength,
+                                )
+                            )
+                        except Exception as e:
+                            logger.debug(f"Invalid causal relation {rel}: {e}")
+
+                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()
@@ -735,7 +873,7 @@ Text:
                 )
                 continue
 
-            return chunk_facts
+            return chunk_facts, usage
 
         except BadRequestError as e:
             last_error = e
@@ -762,7 +900,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.
 
@@ -780,7 +918,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
 
@@ -859,12 +997,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(
@@ -874,7 +1014,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.
 
@@ -893,11 +1033,13 @@ 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)
+    config = get_config()
+    chunks = chunk_text(text, max_chars=config.retain_chunk_size)
     tasks = [
         _extract_facts_with_auto_split(
             chunk=chunk,
@@ -914,10 +1056,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
 
 
 # ============================================================================
@@ -938,7 +1082,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.
 
@@ -955,10 +1099,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 = []
@@ -981,11 +1125,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
@@ -1039,7 +1187,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/fact_storage.py

@@ -7,6 +7,7 @@ Handles insertion of facts into the database.
 import json
 import logging
 
+from ..memory_engine import fq_table
 from .types import ProcessedFact
 
 logger = logging.getLogger(__name__)
@@ -67,8 +68,8 @@ async def insert_facts_batch(
 
     # Batch insert all facts
     results = await conn.fetch(
-        """
-        INSERT INTO memory_units (bank_id, text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
+        f"""
+        INSERT INTO {fq_table("memory_units")} (bank_id, text, embedding, event_date, occurred_start, occurred_end, mentioned_at,
                                  context, fact_type, confidence_score, access_count, metadata, chunk_id, document_id)
         SELECT $1, * FROM unnest(
             $2::text[], $3::vector[], $4::timestamptz[], $5::timestamptz[], $6::timestamptz[], $7::timestamptz[],
@@ -107,8 +108,8 @@ async def ensure_bank_exists(conn, bank_id: str) -> None:
         bank_id: Bank identifier
     """
     await conn.execute(
-        """
-        INSERT INTO banks (bank_id, disposition, background)
+        f"""
+        INSERT INTO {fq_table("banks")} (bank_id, disposition, background)
         VALUES ($1, $2::jsonb, $3)
         ON CONFLICT (bank_id) DO UPDATE
         SET updated_at = NOW()
@@ -141,12 +142,14 @@ async def handle_document_tracking(
     # Always delete old document first if it exists (cascades to units and links)
     # Only delete on the first batch to avoid deleting data we just inserted
     if is_first_batch:
-        await conn.fetchval("DELETE FROM documents WHERE id = $1 AND bank_id = $2 RETURNING id", document_id, bank_id)
+        await conn.fetchval(
+            f"DELETE FROM {fq_table('documents')} WHERE id = $1 AND bank_id = $2 RETURNING id", document_id, bank_id
+        )
 
     # Insert document (or update if exists from concurrent operations)
     await conn.execute(
-        """
-        INSERT INTO documents (id, bank_id, original_text, content_hash, metadata, retain_params)
+        f"""
+        INSERT INTO {fq_table("documents")} (id, bank_id, original_text, content_hash, metadata, retain_params)
         VALUES ($1, $2, $3, $4, $5, $6)
         ON CONFLICT (id, bank_id) DO UPDATE
         SET original_text = EXCLUDED.original_text,

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

@@ -7,6 +7,7 @@ import time
 from datetime import UTC, datetime, timedelta
 from uuid import UUID
 
+from ..memory_engine import fq_table
 from .types import EntityLink
 
 logger = logging.getLogger(__name__)
@@ -290,9 +291,9 @@ async def extract_entities_batch_optimized(
 
             entity_id_list = [uuid.UUID(eid) if isinstance(eid, str) else eid for eid in all_entity_ids]
             rows = await conn.fetch(
-                """
+                f"""
                 SELECT entity_id, unit_id
-                FROM unit_entities
+                FROM {fq_table("unit_entities")}
                 WHERE entity_id = ANY($1::uuid[])
                 """,
                 entity_id_list,
@@ -413,9 +414,9 @@ async def create_temporal_links_batch_per_fact(
         # Get the event_date for each new unit
         fetch_dates_start = time_mod.time()
         rows = await conn.fetch(
-            """
+            f"""
             SELECT id, event_date
-            FROM memory_units
+            FROM {fq_table("memory_units")}
             WHERE id::text = ANY($1)
             """,
             unit_ids,
@@ -432,9 +433,9 @@ async def create_temporal_links_batch_per_fact(
 
         fetch_neighbors_start = time_mod.time()
         all_candidates = await conn.fetch(
-            """
+            f"""
             SELECT id, event_date
-            FROM memory_units
+            FROM {fq_table("memory_units")}
             WHERE bank_id = $1
               AND event_date BETWEEN $2 AND $3
               AND id::text != ALL($4)
@@ -479,8 +480,8 @@ async def create_temporal_links_batch_per_fact(
         if links:
             insert_start = time_mod.time()
             await conn.executemany(
-                """
-                INSERT INTO memory_links (from_unit_id, to_unit_id, link_type, weight, entity_id)
+                f"""
+                INSERT INTO {fq_table("memory_links")} (from_unit_id, to_unit_id, link_type, weight, entity_id)
                 VALUES ($1, $2, $3, $4, $5)
                 ON CONFLICT (from_unit_id, to_unit_id, link_type, COALESCE(entity_id, '00000000-0000-0000-0000-000000000000'::uuid)) DO NOTHING
                 """,
@@ -535,9 +536,9 @@ async def create_semantic_links_batch(
         # Fetch ALL existing units with embeddings in ONE query
         fetch_start = time_mod.time()
         all_existing = await conn.fetch(
-            """
+            f"""
             SELECT id, embedding
-            FROM memory_units
+            FROM {fq_table("memory_units")}
             WHERE bank_id = $1
               AND embedding IS NOT NULL
               AND id::text != ALL($2)
@@ -644,8 +645,8 @@ async def create_semantic_links_batch(
         if all_links:
             insert_start = time_mod.time()
             await conn.executemany(
-                """
-                INSERT INTO memory_links (from_unit_id, to_unit_id, link_type, weight, entity_id)
+                f"""
+                INSERT INTO {fq_table("memory_links")} (from_unit_id, to_unit_id, link_type, weight, entity_id)
                 VALUES ($1, $2, $3, $4, $5)
                 ON CONFLICT (from_unit_id, to_unit_id, link_type, COALESCE(entity_id, '00000000-0000-0000-0000-000000000000'::uuid)) DO NOTHING
                 """,
@@ -721,8 +722,8 @@ async def insert_entity_links_batch(conn, links: list[EntityLink], chunk_size: i
 
     # Insert from temp table with ON CONFLICT (single query for all rows)
     insert_start = time_mod.time()
-    await conn.execute("""
-        INSERT INTO memory_links (from_unit_id, to_unit_id, link_type, weight, entity_id)
+    await conn.execute(f"""
+        INSERT INTO {fq_table("memory_links")} (from_unit_id, to_unit_id, link_type, weight, entity_id)
         SELECT from_unit_id, to_unit_id, link_type, weight, entity_id
         FROM _temp_entity_links
         ON CONFLICT (from_unit_id, to_unit_id, link_type, COALESCE(entity_id, '00000000-0000-0000-0000-000000000000'::uuid)) DO NOTHING
@@ -808,8 +809,8 @@ async def create_causal_links_batch(
             insert_start = time_mod.time()
             try:
                 await conn.executemany(
-                    """
-                    INSERT INTO memory_links (from_unit_id, to_unit_id, link_type, weight, entity_id)
+                    f"""
+                    INSERT INTO {fq_table("memory_links")} (from_unit_id, to_unit_id, link_type, weight, entity_id)
                     VALUES ($1, $2, $3, $4, $5)
                     ON CONFLICT (from_unit_id, to_unit_id, link_type, COALESCE(entity_id, '00000000-0000-0000-0000-000000000000'::uuid)) DO NOTHING
                     """,

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

@@ -9,6 +9,8 @@ import time
 import uuid
 from datetime import UTC, datetime
 
+from ...config import get_config
+from ..memory_engine import fq_table
 from ..search import observation_utils
 from . import embedding_utils
 from .types import EntityLink
@@ -48,8 +50,9 @@ async def regenerate_observations_batch(
         entity_links: Entity links from this batch
         log_buffer: Optional log buffer for timing
     """
-    TOP_N_ENTITIES = 5
-    MIN_FACTS_THRESHOLD = 5
+    config = get_config()
+    TOP_N_ENTITIES = config.observation_top_entities
+    MIN_FACTS_THRESHOLD = config.observation_min_facts
 
     if not entity_links:
         return
@@ -75,8 +78,8 @@ async def regenerate_observations_batch(
 
     # Batch query for entity names
     entity_rows = await conn.fetch(
-        """
-        SELECT id, canonical_name FROM entities
+        f"""
+        SELECT id, canonical_name FROM {fq_table("entities")}
         WHERE id = ANY($1) AND bank_id = $2
         """,
         entity_uuids,
@@ -86,10 +89,10 @@ async def regenerate_observations_batch(
 
     # Batch query for fact counts
     fact_counts = await conn.fetch(
-        """
+        f"""
         SELECT ue.entity_id, COUNT(*) as cnt
-        FROM unit_entities ue
-        JOIN memory_units mu ON ue.unit_id = mu.id
+        FROM {fq_table("unit_entities")} ue
+        JOIN {fq_table("memory_units")} mu ON ue.unit_id = mu.id
         WHERE ue.entity_id = ANY($1) AND mu.bank_id = $2
         GROUP BY ue.entity_id
         """,
@@ -154,10 +157,10 @@ async def _regenerate_entity_observations(
 
     # Get all facts mentioning this entity (exclude observations themselves)
     rows = await conn.fetch(
-        """
+        f"""
         SELECT mu.id, mu.text, mu.context, mu.occurred_start, mu.fact_type
-        FROM memory_units mu
-        JOIN unit_entities ue ON mu.id = ue.unit_id
+        FROM {fq_table("memory_units")} mu
+        JOIN {fq_table("unit_entities")} ue ON mu.id = ue.unit_id
         WHERE mu.bank_id = $1
           AND ue.entity_id = $2
           AND mu.fact_type IN ('world', 'experience')
@@ -193,12 +196,12 @@ async def _regenerate_entity_observations(
 
     # Delete old observations for this entity
     await conn.execute(
-        """
-        DELETE FROM memory_units
+        f"""
+        DELETE FROM {fq_table("memory_units")}
         WHERE id IN (
             SELECT mu.id
-            FROM memory_units mu
-            JOIN unit_entities ue ON mu.id = ue.unit_id
+            FROM {fq_table("memory_units")} mu
+            JOIN {fq_table("unit_entities")} ue ON mu.id = ue.unit_id
             WHERE mu.bank_id = $1
               AND mu.fact_type = 'observation'
               AND ue.entity_id = $2
@@ -217,8 +220,8 @@ async def _regenerate_entity_observations(
 
     for obs_text, embedding in zip(observations, embeddings):
         result = await conn.fetchrow(
-            """
-            INSERT INTO memory_units (
+            f"""
+            INSERT INTO {fq_table("memory_units")} (
                 bank_id, text, embedding, context, event_date,
                 occurred_start, occurred_end, mentioned_at,
                 fact_type, access_count
@@ -240,8 +243,8 @@ async def _regenerate_entity_observations(
 
         # Link observation to entity
         await conn.execute(
-            """
-            INSERT INTO unit_entities (unit_id, entity_id)
+            f"""
+            INSERT INTO {fq_table("unit_entities")} (unit_id, entity_id)
             VALUES ($1, $2)
             """,
             uuid.UUID(obs_id),

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

@@ -8,7 +8,6 @@ import logging
 import time
 import uuid
 from datetime import UTC, datetime
-from typing import Any
 
 from ..db_utils import acquire_with_retry
 from . import bank_utils
@@ -19,6 +18,7 @@ def utcnow():
     return datetime.now(UTC)
 
 
+from ..response_models import TokenUsage
 from . import (
     chunk_storage,
     deduplication,
@@ -29,7 +29,7 @@ from . import (
     link_creation,
     observation_regeneration,
 )
-from .types import ExtractedFact, ProcessedFact, RetainContent
+from .types import ExtractedFact, ProcessedFact, RetainContent, RetainContentDict
 
 logger = logging.getLogger(__name__)
 
@@ -43,12 +43,12 @@ async def retain_batch(
     format_date_fn,
     duplicate_checker_fn,
     bank_id: str,
-    contents_dicts: list[dict[str, Any]],
+    contents_dicts: list[RetainContentDict],
     document_id: str | None = None,
     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.
 
@@ -68,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)
@@ -92,6 +92,7 @@ async def retain_batch(
             context=item.get("context", ""),
             event_date=item.get("event_date") or utcnow(),
             metadata=item.get("metadata", {}),
+            entities=item.get("entities", []),
         )
         contents.append(content)
 
@@ -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(
@@ -107,11 +108,64 @@ async def retain_batch(
     )
 
     if not extracted_facts:
+        # Still need to create document if document_id was provided
+        async with acquire_with_retry(pool) as conn:
+            async with conn.transaction():
+                await fact_storage.ensure_bank_exists(conn, bank_id)
+
+                # Handle document tracking even with no facts
+                if document_id:
+                    combined_content = "\n".join([c.get("content", "") for c in contents_dicts])
+                    retain_params = {}
+                    if contents_dicts:
+                        first_item = contents_dicts[0]
+                        if first_item.get("context"):
+                            retain_params["context"] = first_item["context"]
+                        if first_item.get("event_date"):
+                            retain_params["event_date"] = (
+                                first_item["event_date"].isoformat()
+                                if hasattr(first_item["event_date"], "isoformat")
+                                else str(first_item["event_date"])
+                            )
+                        if first_item.get("metadata"):
+                            retain_params["metadata"] = first_item["metadata"]
+                    await fact_storage.handle_document_tracking(
+                        conn, bank_id, document_id, combined_content, is_first_batch, retain_params
+                    )
+                else:
+                    # Check for per-item document_ids
+                    from collections import defaultdict
+
+                    contents_by_doc = defaultdict(list)
+                    for idx, content_dict in enumerate(contents_dicts):
+                        doc_id = content_dict.get("document_id")
+                        if doc_id:
+                            contents_by_doc[doc_id].append((idx, content_dict))
+
+                    for doc_id, doc_contents in contents_by_doc.items():
+                        combined_content = "\n".join([c.get("content", "") for _, c in doc_contents])
+                        retain_params = {}
+                        if doc_contents:
+                            first_item = doc_contents[0][1]
+                            if first_item.get("context"):
+                                retain_params["context"] = first_item["context"]
+                            if first_item.get("event_date"):
+                                retain_params["event_date"] = (
+                                    first_item["event_date"].isoformat()
+                                    if hasattr(first_item["event_date"], "isoformat")
+                                    else str(first_item["event_date"])
+                                )
+                            if first_item.get("metadata"):
+                                retain_params["metadata"] = first_item["metadata"]
+                        await fact_storage.handle_document_tracking(
+                            conn, bank_id, doc_id, combined_content, is_first_batch, retain_params
+                        )
+
         total_time = time.time() - start_time
         logger.info(
-            f"RETAIN_BATCH COMPLETE: 0 facts extracted from {len(contents)} contents in {total_time:.3f}s (nothing to store)"
+            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:
@@ -291,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()
@@ -300,8 +354,18 @@ async def retain_batch(
 
             # Process entities
             step_start = time.time()
+            # Build map of content_index -> user entities for merging
+            user_entities_per_content = {
+                idx: content.entities for idx, content in enumerate(contents) if content.entities
+            }
             entity_links = await entity_processing.process_entities_batch(
-                entity_resolver, conn, bank_id, unit_ids, non_duplicate_facts, log_buffer
+                entity_resolver,
+                conn,
+                bank_id,
+                unit_ids,
+                non_duplicate_facts,
+                log_buffer,
+                user_entities_per_content=user_entities_per_content,
             )
             log_buffer.append(f"[6] Process entities: {len(entity_links)} links in {time.time() - step_start:.3f}s")
 
@@ -352,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/retain/types.py

@@ -7,9 +7,35 @@ from content input to fact storage.
 
 from dataclasses import dataclass, field
 from datetime import UTC, datetime
+from typing import TypedDict
 from uuid import UUID
 
 
+class RetainContentDict(TypedDict, total=False):
+    """Type definition for content items in retain_batch_async.
+
+    Fields:
+        content: Text content to store (required)
+        context: Context about the content (optional)
+        event_date: When the content occurred (optional, defaults to now)
+        metadata: Custom key-value metadata (optional)
+        document_id: Document ID for this content item (optional)
+        entities: User-provided entities to merge with extracted entities (optional)
+    """
+
+    content: str  # Required
+    context: str
+    event_date: datetime
+    metadata: dict[str, str]
+    document_id: str
+    entities: list[dict[str, str]]  # [{"text": "...", "type": "..."}]
+
+
+def _now_utc() -> datetime:
+    """Factory function for default event_date."""
+    return datetime.now(UTC)
+
+
 @dataclass
 class RetainContent:
     """
@@ -20,15 +46,9 @@ class RetainContent:
 
     content: str
     context: str = ""
-    event_date: datetime | None = None
+    event_date: datetime = field(default_factory=_now_utc)
     metadata: dict[str, str] = field(default_factory=dict)
-
-    def __post_init__(self):
-        """Ensure event_date is set."""
-        if self.event_date is None:
-            from datetime import datetime
-
-            self.event_date = datetime.now(UTC)
+    entities: list[dict[str, str]] = field(default_factory=list)  # User-provided entities
 
 
 @dataclass
@@ -135,6 +155,9 @@ class ProcessedFact:
     # DB fields (set after insertion)
     unit_id: UUID | None = None
 
+    # Track which content this fact came from (for user entity merging)
+    content_index: int = 0
+
     @property
     def is_duplicate(self) -> bool:
         """Check if this fact was marked as a duplicate."""
@@ -177,6 +200,7 @@ class ProcessedFact:
             entities=entities,
             causal_relations=extracted_fact.causal_relations,
             chunk_id=chunk_id,
+            content_index=extracted_fact.content_index,
         )
 
 

hindsight-api/hindsight_api/engine/search/graph_retrieval.py

@@ -10,6 +10,7 @@ import logging
 from abc import ABC, abstractmethod
 
 from ..db_utils import acquire_with_retry
+from ..memory_engine import fq_table
 from .types import RetrievalResult
 
 logger = logging.getLogger(__name__)
@@ -139,11 +140,11 @@ class BFSGraphRetriever(GraphRetriever):
 
         # Step 1: Find entry points
         entry_points = await conn.fetch(
-            """
+            f"""
             SELECT id, text, context, event_date, occurred_start, occurred_end,
                    mentioned_at, access_count, embedding, fact_type, document_id, chunk_id,
                    1 - (embedding <=> $1::vector) AS similarity
-            FROM memory_units
+            FROM {fq_table("memory_units")}
             WHERE bank_id = $2
               AND embedding IS NOT NULL
               AND fact_type = $3
@@ -188,13 +189,13 @@ class BFSGraphRetriever(GraphRetriever):
             if batch_nodes and budget_remaining > 0:
                 max_neighbors = len(batch_nodes) * 20
                 neighbors = await conn.fetch(
-                    """
+                    f"""
                     SELECT mu.id, mu.text, mu.context, mu.occurred_start, mu.occurred_end,
                            mu.mentioned_at, mu.access_count, mu.embedding, mu.fact_type,
                            mu.document_id, mu.chunk_id,
                            ml.weight, ml.link_type, ml.from_unit_id
-                    FROM memory_links ml
-                    JOIN memory_units mu ON ml.to_unit_id = mu.id
+                    FROM {fq_table("memory_links")} ml
+                    JOIN {fq_table("memory_units")} mu ON ml.to_unit_id = mu.id
                     WHERE ml.from_unit_id = ANY($1::uuid[])
                       AND ml.weight >= $2
                       AND mu.fact_type = $3

hindsight-api/hindsight_api/engine/search/mpfp_retrieval.py

@@ -20,6 +20,7 @@ from collections import defaultdict
 from dataclasses import dataclass, field
 
 from ..db_utils import acquire_with_retry
+from ..memory_engine import fq_table
 from .graph_retrieval import GraphRetriever
 from .types import RetrievalResult
 
@@ -217,10 +218,10 @@ async def load_typed_adjacency(pool, bank_id: str) -> TypedAdjacency:
     """
     async with acquire_with_retry(pool) as conn:
         rows = await conn.fetch(
-            """
+            f"""
             SELECT ml.from_unit_id, ml.to_unit_id, ml.link_type, ml.weight
-            FROM memory_links ml
-            JOIN memory_units mu ON ml.from_unit_id = mu.id
+            FROM {fq_table("memory_links")} ml
+            JOIN {fq_table("memory_units")} mu ON ml.from_unit_id = mu.id
             WHERE mu.bank_id = $1
               AND ml.weight >= 0.1
             ORDER BY ml.from_unit_id, ml.weight DESC
@@ -252,10 +253,10 @@ async def fetch_memory_units_by_ids(
 
     async with acquire_with_retry(pool) as conn:
         rows = await conn.fetch(
-            """
+            f"""
             SELECT id, text, context, event_date, occurred_start, occurred_end,
                    mentioned_at, access_count, embedding, fact_type, document_id, chunk_id
-            FROM memory_units
+            FROM {fq_table("memory_units")}
             WHERE id = ANY($1::uuid[])
               AND fact_type = $2
             """,
@@ -418,9 +419,9 @@ class MPFPGraphRetriever(GraphRetriever):
         """Fallback: find semantic seeds via embedding search."""
         async with acquire_with_retry(pool) as conn:
             rows = await conn.fetch(
-                """
+                f"""
                 SELECT id, 1 - (embedding <=> $1::vector) AS similarity
-                FROM memory_units
+                FROM {fq_table("memory_units")}
                 WHERE bank_id = $2
                   AND embedding IS NOT NULL
                   AND fact_type = $3

hindsight-api/hindsight_api/engine/search/reranking.py

@@ -26,6 +26,23 @@ class CrossEncoderReranker:
 
             cross_encoder = create_cross_encoder_from_env()
         self.cross_encoder = cross_encoder
+        self._initialized = False
+
+    async def ensure_initialized(self):
+        """Ensure the cross-encoder model is initialized (for lazy initialization)."""
+        if self._initialized:
+            return
+
+        import asyncio
+
+        cross_encoder = self.cross_encoder
+        # For local providers, run in thread pool to avoid blocking event loop
+        if cross_encoder.provider_name == "local":
+            loop = asyncio.get_event_loop()
+            await loop.run_in_executor(None, lambda: asyncio.run(cross_encoder.initialize()))
+        else:
+            await cross_encoder.initialize()
+        self._initialized = True
 
     def rerank(self, query: str, candidates: list[MergedCandidate]) -> list[ScoredResult]:
         """

hindsight-api/hindsight_api/engine/search/retrieval.py

@@ -16,6 +16,7 @@ from typing import Optional
 
 from ...config import get_config
 from ..db_utils import acquire_with_retry
+from ..memory_engine import fq_table
 from .graph_retrieval import BFSGraphRetriever, GraphRetriever
 from .mpfp_retrieval import MPFPGraphRetriever
 from .types import RetrievalResult
@@ -80,10 +81,10 @@ async def retrieve_semantic(
         List of RetrievalResult objects
     """
     results = await conn.fetch(
-        """
+        f"""
         SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, access_count, embedding, fact_type, document_id, chunk_id,
                1 - (embedding <=> $1::vector) AS similarity
-        FROM memory_units
+        FROM {fq_table("memory_units")}
         WHERE bank_id = $2
           AND embedding IS NOT NULL
           AND fact_type = $3
@@ -131,10 +132,10 @@ async def retrieve_bm25(conn, query_text: str, bank_id: str, fact_type: str, lim
     query_tsquery = " | ".join(tokens)
 
     results = await conn.fetch(
-        """
+        f"""
         SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, access_count, embedding, fact_type, document_id, chunk_id,
                ts_rank_cd(search_vector, to_tsquery('english', $1)) AS bm25_score
-        FROM memory_units
+        FROM {fq_table("memory_units")}
         WHERE bank_id = $2
           AND fact_type = $3
           AND search_vector @@ to_tsquery('english', $1)
@@ -188,10 +189,10 @@ async def retrieve_temporal(
         end_date = end_date.replace(tzinfo=UTC)
 
     entry_points = await conn.fetch(
-        """
+        f"""
         SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at, access_count, embedding, fact_type, document_id, chunk_id,
                1 - (embedding <=> $1::vector) AS similarity
-        FROM memory_units
+        FROM {fq_table("memory_units")}
         WHERE bank_id = $2
           AND fact_type = $3
           AND embedding IS NOT NULL
@@ -272,12 +273,12 @@ async def retrieve_temporal(
         # Get neighbors via temporal and causal links
         if budget_remaining > 0:
             neighbors = await conn.fetch(
-                """
+                f"""
                 SELECT mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start, mu.occurred_end, mu.mentioned_at, mu.access_count, mu.embedding, mu.fact_type, mu.document_id, mu.chunk_id,
                        ml.weight, ml.link_type,
                        1 - (mu.embedding <=> $1::vector) AS similarity
-                FROM memory_links ml
-                JOIN memory_units mu ON ml.to_unit_id = mu.id
+                FROM {fq_table("memory_links")} ml
+                JOIN {fq_table("memory_units")} mu ON ml.to_unit_id = mu.id
                 WHERE ml.from_unit_id = $2
                   AND ml.link_type IN ('temporal', 'causes', 'caused_by', 'enables', 'prevents')
                   AND ml.weight >= 0.1
@@ -546,11 +547,11 @@ async def _get_temporal_entry_points(
         end_date = end_date.replace(tzinfo=UTC)
 
     rows = await conn.fetch(
-        """
+        f"""
         SELECT id, text, context, event_date, occurred_start, occurred_end, mentioned_at,
                access_count, embedding, fact_type, document_id, chunk_id,
                1 - (embedding <=> $1::vector) AS similarity
-        FROM memory_units
+        FROM {fq_table("memory_units")}
         WHERE bank_id = $2
           AND fact_type = $3
           AND embedding IS NOT NULL

hindsight-api/hindsight_api/engine/search/think_utils.py

@@ -101,7 +101,7 @@ def build_think_prompt(
     name: str,
     disposition: DispositionTraits,
     background: str,
-    context: str = None,
+    context: str | None = None,
 ) -> str:
     """Build the think prompt for the LLM."""
     disposition_desc = build_disposition_description(disposition)
@@ -172,7 +172,7 @@ def get_system_message(disposition: DispositionTraits) -> str:
         " ".join(instructions) if instructions else "Balance your disposition traits when interpreting information."
     )
 
-    return f"You are a person with your own thoughts, experiences, opinions, and disposition. When asked a question, respond naturally like a human would - share what you genuinely think based on what you know and have experienced. {disposition_instruction} Be direct, express your views confidently, and use 'I think', 'I believe', 'in my view', etc. Respond in plain text without markdown formatting."
+    return f"You are a person with your own thoughts, experiences, opinions, and disposition. When asked a question, respond naturally like a human would - share what you genuinely think based on what you know and have experienced. {disposition_instruction} Be direct, express your views confidently, and use 'I think', 'I believe', 'in my view', etc. Respond in plain text without markdown formatting. IMPORTANT: Detect the language of the question and respond in the SAME language. Do not translate to English if the question is in another language."
 
 
 async def extract_opinions_from_text(llm_config, text: str, query: str) -> list[Opinion]:

hindsight-api/hindsight_api/engine/search/tracer.py

@@ -115,7 +115,7 @@ class SearchTracer:
         node_id: str,
         text: str,
         context: str,
-        event_date: datetime,
+        event_date: datetime | None,
         access_count: int,
         is_entry_point: bool,
         parent_node_id: str | None,

hindsight-api/hindsight_api/engine/task_backend.py

@@ -89,6 +89,38 @@ class TaskBackend(ABC):
             traceback.print_exc()
 
 
+class SyncTaskBackend(TaskBackend):
+    """
+    Synchronous task backend that executes tasks immediately.
+
+    This is useful for embedded/CLI usage where we don't want background
+    workers that prevent clean exit. Tasks are executed inline rather than
+    being queued.
+    """
+
+    async def initialize(self):
+        """No-op for sync backend."""
+        self._initialized = True
+        logger.debug("SyncTaskBackend initialized")
+
+    async def submit_task(self, task_dict: dict[str, Any]):
+        """
+        Execute the task immediately (synchronously).
+
+        Args:
+            task_dict: Task dictionary to execute
+        """
+        if not self._initialized:
+            await self.initialize()
+
+        await self._execute_task(task_dict)
+
+    async def shutdown(self):
+        """No-op for sync backend."""
+        self._initialized = False
+        logger.debug("SyncTaskBackend shutdown")
+
+
 class AsyncIOQueueBackend(TaskBackend):
     """
     Task backend implementation using asyncio queues.

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/extensions/__init__.py

@@ -0,0 +1,66 @@
+"""
+Hindsight Extensions System.
+
+Extensions allow customizing and extending Hindsight behavior without modifying core code.
+Extensions are loaded via environment variables pointing to implementation classes.
+
+Example:
+    HINDSIGHT_API_OPERATION_VALIDATOR_EXTENSION=mypackage.validators:MyValidator
+    HINDSIGHT_API_OPERATION_VALIDATOR_MAX_RETRIES=3
+
+    HINDSIGHT_API_HTTP_EXTENSION=mypackage.http:MyHttpExtension
+    HINDSIGHT_API_HTTP_SOME_CONFIG=value
+
+Extensions receive an ExtensionContext that provides a controlled API for interacting
+with the system (e.g., running migrations for tenant schemas).
+"""
+
+from hindsight_api.extensions.base import Extension
+from hindsight_api.extensions.builtin import ApiKeyTenantExtension
+from hindsight_api.extensions.context import DefaultExtensionContext, ExtensionContext
+from hindsight_api.extensions.http import HttpExtension
+from hindsight_api.extensions.loader import load_extension
+from hindsight_api.extensions.operation_validator import (
+    OperationValidationError,
+    OperationValidatorExtension,
+    RecallContext,
+    RecallResult,
+    ReflectContext,
+    ReflectResultContext,
+    RetainContext,
+    RetainResult,
+    ValidationResult,
+)
+from hindsight_api.extensions.tenant import (
+    AuthenticationError,
+    TenantContext,
+    TenantExtension,
+)
+from hindsight_api.models import RequestContext
+
+__all__ = [
+    # Base
+    "Extension",
+    "load_extension",
+    # Context
+    "ExtensionContext",
+    "DefaultExtensionContext",
+    # HTTP Extension
+    "HttpExtension",
+    # Operation Validator
+    "OperationValidationError",
+    "OperationValidatorExtension",
+    "RecallContext",
+    "RecallResult",
+    "ReflectContext",
+    "ReflectResultContext",
+    "RetainContext",
+    "RetainResult",
+    "ValidationResult",
+    # Tenant/Auth
+    "ApiKeyTenantExtension",
+    "AuthenticationError",
+    "RequestContext",
+    "TenantContext",
+    "TenantExtension",
+]

hindsight-api/hindsight_api/extensions/base.py

@@ -0,0 +1,81 @@
+"""Base Extension class for all Hindsight extensions."""
+
+from abc import ABC
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from hindsight_api.extensions.context import ExtensionContext
+
+
+class Extension(ABC):
+    """
+    Base class for all Hindsight extensions.
+
+    Extensions are loaded via environment variables and receive configuration
+    from prefixed environment variables.
+
+    Example:
+        HINDSIGHT_API_MY_EXTENSION=mypackage.ext:MyExtension
+        HINDSIGHT_API_MY_SOME_CONFIG=value
+
+        The extension receives: {"some_config": "value"}
+
+    Extensions also receive an ExtensionContext that provides a controlled API
+    for interacting with the system (e.g., running migrations for tenant schemas).
+    """
+
+    def __init__(self, config: dict[str, str]):
+        """
+        Initialize the extension with configuration.
+
+        Args:
+            config: Dictionary of configuration values from environment variables.
+                    Keys are lowercased with the prefix stripped.
+        """
+        self.config = config
+        self._context: "ExtensionContext | None" = None
+
+    def set_context(self, context: "ExtensionContext") -> None:
+        """
+        Set the extension context.
+
+        Called by the extension loader after instantiation.
+        Extensions should not call this directly.
+
+        Args:
+            context: The ExtensionContext providing system APIs.
+        """
+        self._context = context
+
+    @property
+    def context(self) -> "ExtensionContext":
+        """
+        Get the extension context.
+
+        Returns:
+            The ExtensionContext providing system APIs.
+
+        Raises:
+            RuntimeError: If context has not been set yet.
+        """
+        if self._context is None:
+            raise RuntimeError(
+                "Extension context not set. Context is available after the extension is loaded by the system."
+            )
+        return self._context
+
+    async def on_startup(self) -> None:
+        """
+        Called when the application starts.
+
+        Override to perform initialization tasks like connecting to external services.
+        """
+        pass
+
+    async def on_shutdown(self) -> None:
+        """
+        Called when the application shuts down.
+
+        Override to perform cleanup tasks like closing connections.
+        """
+        pass

hindsight-api/hindsight_api/extensions/builtin/__init__.py

@@ -0,0 +1,18 @@
+"""
+Built-in extension implementations.
+
+These are ready-to-use implementations of the extension interfaces.
+They can be used directly or serve as examples for custom implementations.
+
+Available built-in extensions:
+    - ApiKeyTenantExtension: Simple API key validation with public schema
+
+Example usage:
+    HINDSIGHT_API_TENANT_EXTENSION=hindsight_api.extensions.builtin.tenant:ApiKeyTenantExtension
+"""
+
+from hindsight_api.extensions.builtin.tenant import ApiKeyTenantExtension
+
+__all__ = [
+    "ApiKeyTenantExtension",
+]

hindsight-api/hindsight_api/extensions/builtin/tenant.py

@@ -0,0 +1,33 @@
+"""Built-in tenant extension implementations."""
+
+from hindsight_api.extensions.tenant import AuthenticationError, TenantContext, TenantExtension
+from hindsight_api.models import RequestContext
+
+
+class ApiKeyTenantExtension(TenantExtension):
+    """
+    Built-in tenant extension that validates API key against an environment variable.
+
+    This is a simple implementation that:
+    1. Validates the API key matches HINDSIGHT_API_TENANT_API_KEY
+    2. Returns 'public' as the schema for all authenticated requests
+
+    Configuration:
+        HINDSIGHT_API_TENANT_EXTENSION=hindsight_api.extensions.builtin.tenant:ApiKeyTenantExtension
+        HINDSIGHT_API_TENANT_API_KEY=your-secret-key
+
+    For multi-tenant setups with separate schemas per tenant, implement a custom
+    TenantExtension that looks up the schema based on the API key or token claims.
+    """
+
+    def __init__(self, config: dict[str, str]):
+        super().__init__(config)
+        self.expected_api_key = config.get("api_key")
+        if not self.expected_api_key:
+            raise ValueError("HINDSIGHT_API_TENANT_API_KEY is required when using ApiKeyTenantExtension")
+
+    async def authenticate(self, context: RequestContext) -> TenantContext:
+        """Validate API key and return public schema context."""
+        if context.api_key != self.expected_api_key:
+            raise AuthenticationError("Invalid API key")
+        return TenantContext(schema_name="public")

hindsight-api/hindsight_api/extensions/context.py

@@ -0,0 +1,117 @@
+"""Extension context providing a controlled API for extensions to interact with the system."""
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from hindsight_api.engine.interface import MemoryEngineInterface
+
+
+class ExtensionContext(ABC):
+    """
+    Abstract context providing a controlled API for extensions.
+
+    Extensions receive this context instead of direct access to internal
+    components like MemoryEngine or database connections. This provides:
+    - A stable API that won't break when internals change
+    - Security by limiting what extensions can access
+    - Clear documentation of what extensions can do
+
+    Built-in implementation:
+        hindsight_api.extensions.builtin.context.DefaultExtensionContext
+
+    Example usage in an extension:
+        class MyTenantExtension(TenantExtension):
+            async def on_startup(self) -> None:
+                # Run migrations for a new tenant schema
+                await self.context.run_migration("tenant_acme")
+
+        class MyHttpExtension(HttpExtension):
+            def get_router(self, memory):
+                # Use memory engine for custom endpoints
+                engine = self.context.get_memory_engine()
+                ...
+    """
+
+    @abstractmethod
+    async def run_migration(self, schema: str) -> None:
+        """
+        Run database migrations for a specific schema.
+
+        This creates the schema if it doesn't exist and runs all pending
+        migrations. Uses advisory locks to coordinate between distributed workers.
+
+        Args:
+            schema: PostgreSQL schema name (e.g., "tenant_acme").
+                    The schema will be created if it doesn't exist.
+
+        Raises:
+            RuntimeError: If migrations fail to complete.
+
+        Example:
+            # Provision a new tenant schema
+            await context.run_migration("tenant_acme")
+        """
+        ...
+
+    @abstractmethod
+    def get_memory_engine(self) -> "MemoryEngineInterface":
+        """
+        Get the memory engine interface.
+
+        Returns the MemoryEngineInterface for performing memory operations
+        like retain, recall, reflect, and entity/document management.
+
+        Returns:
+            MemoryEngineInterface instance.
+
+        Example:
+            engine = context.get_memory_engine()
+            result = await engine.recall_async(bank_id, query)
+        """
+        ...
+
+
+class DefaultExtensionContext(ExtensionContext):
+    """
+    Default implementation of ExtensionContext.
+
+    Uses the system's database URL and migration infrastructure.
+    """
+
+    def __init__(
+        self,
+        database_url: str,
+        memory_engine: "MemoryEngineInterface | None" = None,
+    ):
+        """
+        Initialize the context.
+
+        Args:
+            database_url: SQLAlchemy database URL for migrations.
+            memory_engine: Optional MemoryEngine instance for memory operations.
+        """
+        self._database_url = database_url
+        self._memory_engine = memory_engine
+
+    async def run_migration(self, schema: str) -> None:
+        """Run migrations for a specific schema."""
+        from hindsight_api.migrations import run_migrations
+
+        # Prefer getting URL from memory engine (handles pg0 case where URL is set after init)
+        db_url = self._database_url
+        if self._memory_engine is not None:
+            engine_url = getattr(self._memory_engine, "db_url", None)
+            if engine_url:
+                db_url = engine_url
+
+        run_migrations(db_url, schema=schema)
+
+    def get_memory_engine(self) -> "MemoryEngineInterface":
+        """Get the memory engine interface."""
+        if self._memory_engine is None:
+            raise RuntimeError(
+                "Memory engine not configured in ExtensionContext. "
+                "Ensure the context was created with a memory_engine parameter."
+            )
+        return self._memory_engine

hindsight-api/hindsight_api/extensions/http.py

@@ -0,0 +1,89 @@
+"""
+HTTP Extension for adding custom endpoints to the Hindsight API.
+
+This extension allows adding custom HTTP endpoints under the /ext/ path prefix.
+The extension provides a FastAPI router that is mounted on the main application.
+"""
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING
+
+from fastapi import APIRouter
+
+from hindsight_api.extensions.base import Extension
+
+if TYPE_CHECKING:
+    from hindsight_api import MemoryEngine
+
+
+class HttpExtension(Extension, ABC):
+    """
+    Base class for HTTP extensions that add custom API endpoints.
+
+    HTTP extensions provide a FastAPI router that gets mounted under /ext/.
+    The extension has full control over the routes, request/response models, and handlers.
+
+    Example:
+        ```python
+        from fastapi import APIRouter
+        from hindsight_api.extensions import HttpExtension
+
+        class MyHttpExtension(HttpExtension):
+            def get_router(self, memory: MemoryEngine) -> APIRouter:
+                router = APIRouter()
+
+                @router.get("/hello")
+                async def hello():
+                    return {"message": "Hello from extension!"}
+
+                @router.post("/custom/{bank_id}/action")
+                async def custom_action(bank_id: str):
+                    # Access memory engine for database operations
+                    pool = await memory._get_pool()
+                    # ... custom logic
+                    return {"status": "ok"}
+
+                return router
+        ```
+
+    The routes will be available at:
+        - GET /ext/hello
+        - POST /ext/custom/{bank_id}/action
+
+    Configuration via environment variables:
+        HINDSIGHT_API_HTTP_EXTENSION=mypackage.ext:MyHttpExtension
+        HINDSIGHT_API_HTTP_SOME_CONFIG=value
+
+    The extension receives config: {"some_config": "value"}
+    """
+
+    @abstractmethod
+    def get_router(self, memory: "MemoryEngine") -> APIRouter:
+        """
+        Return a FastAPI router with custom endpoints.
+
+        The router will be mounted at /ext/ on the main application.
+        All routes defined in the router will be prefixed with /ext/.
+
+        Args:
+            memory: The MemoryEngine instance for database access and core operations.
+                   Use this to access the connection pool, run queries, or call
+                   memory operations like retain, recall, etc.
+
+        Returns:
+            A FastAPI APIRouter with the custom endpoints defined.
+
+        Example:
+            ```python
+            def get_router(self, memory: MemoryEngine) -> APIRouter:
+                router = APIRouter(tags=["My Extension"])
+
+                @router.get("/status")
+                async def status():
+                    health = await memory.health_check()
+                    return {"extension": "healthy", "memory": health}
+
+                return router
+            ```
+        """
+        pass

hindsight-api/hindsight_api/extensions/loader.py

@@ -0,0 +1,125 @@
+"""Extension loader utilities."""
+
+import importlib
+import logging
+import os
+from typing import TYPE_CHECKING, TypeVar
+
+from hindsight_api.extensions.base import Extension
+
+if TYPE_CHECKING:
+    from hindsight_api.extensions.context import ExtensionContext
+
+logger = logging.getLogger(__name__)
+
+T = TypeVar("T", bound=Extension)
+
+
+class ExtensionLoadError(Exception):
+    """Raised when an extension fails to load."""
+
+    pass
+
+
+def load_extension(
+    prefix: str,
+    base_class: type[T],
+    env_prefix: str = "HINDSIGHT_API",
+    context: "ExtensionContext | None" = None,
+) -> T | None:
+    """
+    Load an extension from environment variable configuration.
+
+    The extension class is specified via {env_prefix}_{prefix}_EXTENSION environment
+    variable in the format "module.path:ClassName".
+
+    Configuration for the extension is collected from all environment variables
+    matching {env_prefix}_{prefix}_* (excluding the EXTENSION variable itself).
+
+    Args:
+        prefix: The extension prefix (e.g., "OPERATION_VALIDATOR").
+        base_class: The base class that the extension must inherit from.
+        env_prefix: The environment variable prefix (default: "HINDSIGHT_API").
+        context: Optional ExtensionContext to provide system APIs to the extension.
+
+    Returns:
+        An instance of the extension, or None if not configured.
+
+    Raises:
+        ExtensionLoadError: If the extension fails to load or validate.
+
+    Example:
+        HINDSIGHT_API_OPERATION_VALIDATOR_EXTENSION=mypackage.validators:MyValidator
+        HINDSIGHT_API_OPERATION_VALIDATOR_MAX_REQUESTS=100
+
+        ext = load_extension("OPERATION_VALIDATOR", OperationValidatorExtension)
+        # ext.config == {"max_requests": "100"}
+    """
+    env_var = f"{env_prefix}_{prefix}_EXTENSION"
+    ext_path = os.getenv(env_var)
+
+    if not ext_path:
+        logger.debug(f"No extension configured for {env_var}")
+        return None
+
+    logger.info(f"Loading extension from {env_var}={ext_path}")
+
+    # Parse "module.path:ClassName"
+    if ":" not in ext_path:
+        raise ExtensionLoadError(f"Invalid extension path '{ext_path}'. Expected format: 'module.path:ClassName'")
+
+    module_path, class_name = ext_path.rsplit(":", 1)
+
+    # Import the module
+    try:
+        module = importlib.import_module(module_path)
+    except ImportError as e:
+        raise ExtensionLoadError(f"Failed to import extension module '{module_path}': {e}") from e
+
+    # Get the class
+    try:
+        ext_class = getattr(module, class_name)
+    except AttributeError as e:
+        raise ExtensionLoadError(f"Extension class '{class_name}' not found in module '{module_path}'") from e
+
+    # Validate inheritance
+    if not isinstance(ext_class, type) or not issubclass(ext_class, base_class):
+        raise ExtensionLoadError(f"Extension class '{ext_class.__name__}' must inherit from '{base_class.__name__}'")
+
+    # Collect configuration from environment variables
+    config = _collect_config(env_prefix, prefix)
+
+    logger.info(f"Loaded extension {ext_class.__name__} with config keys: {list(config.keys())}")
+
+    # Instantiate the extension
+    try:
+        extension = ext_class(config)
+    except Exception as e:
+        raise ExtensionLoadError(f"Failed to instantiate extension '{ext_class.__name__}': {e}") from e
+
+    # Set the context if provided
+    if context is not None:
+        extension.set_context(context)
+        logger.debug(f"Set context on extension {ext_class.__name__}")
+
+    return extension
+
+
+def _collect_config(env_prefix: str, prefix: str) -> dict[str, str]:
+    """
+    Collect configuration from environment variables.
+
+    Collects all variables matching {env_prefix}_{prefix}_* except for
+    {env_prefix}_{prefix}_EXTENSION, strips the prefix, and lowercases keys.
+    """
+    config = {}
+    full_prefix = f"{env_prefix}_{prefix}_"
+    extension_var = f"{full_prefix}EXTENSION"
+
+    for key, value in os.environ.items():
+        if key.startswith(full_prefix) and key != extension_var:
+            # Strip prefix and lowercase the key
+            config_key = key[len(full_prefix) :].lower()
+            config[config_key] = value
+
+    return config

hindsight-api/hindsight_api/extensions/operation_validator.py

@@ -0,0 +1,327 @@
+"""Operation Validator Extension for validating retain/recall/reflect operations."""
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from datetime import datetime
+from typing import TYPE_CHECKING, Any
+
+from hindsight_api.extensions.base import Extension
+
+if TYPE_CHECKING:
+    from hindsight_api.engine.memory_engine import Budget
+    from hindsight_api.engine.response_models import RecallResult as RecallResultModel
+    from hindsight_api.engine.response_models import ReflectResult
+    from hindsight_api.models import RequestContext
+
+
+class OperationValidationError(Exception):
+    """Raised when an operation fails validation."""
+
+    def __init__(self, reason: str, status_code: int = 403):
+        self.reason = reason
+        self.status_code = status_code
+        super().__init__(f"Operation validation failed: {reason}")
+
+
+@dataclass
+class ValidationResult:
+    """Result of an operation validation."""
+
+    allowed: bool
+    reason: str | None = None
+    status_code: int = 403  # Default to Forbidden
+
+    @classmethod
+    def accept(cls) -> "ValidationResult":
+        """Create an accepted validation result."""
+        return cls(allowed=True)
+
+    @classmethod
+    def reject(cls, reason: str, status_code: int = 403) -> "ValidationResult":
+        """Create a rejected validation result with a reason and HTTP status code."""
+        return cls(allowed=False, reason=reason, status_code=status_code)
+
+
+# =============================================================================
+# Pre-operation Contexts (all user-provided parameters)
+# =============================================================================
+
+
+@dataclass
+class RetainContext:
+    """Context for a retain operation validation (pre-operation).
+
+    Contains ALL user-provided parameters for the retain operation.
+    """
+
+    bank_id: str
+    contents: list[dict]  # List of {content, context, event_date, document_id}
+    request_context: "RequestContext"
+    document_id: str | None = None
+    fact_type_override: str | None = None
+    confidence_score: float | None = None
+
+
+@dataclass
+class RecallContext:
+    """Context for a recall operation validation (pre-operation).
+
+    Contains ALL user-provided parameters for the recall operation.
+    """
+
+    bank_id: str
+    query: str
+    request_context: "RequestContext"
+    budget: "Budget | None" = None
+    max_tokens: int = 4096
+    enable_trace: bool = False
+    fact_types: list[str] = field(default_factory=list)
+    question_date: datetime | None = None
+    include_entities: bool = False
+    max_entity_tokens: int = 500
+    include_chunks: bool = False
+    max_chunk_tokens: int = 8192
+
+
+@dataclass
+class ReflectContext:
+    """Context for a reflect operation validation (pre-operation).
+
+    Contains ALL user-provided parameters for the reflect operation.
+    """
+
+    bank_id: str
+    query: str
+    request_context: "RequestContext"
+    budget: "Budget | None" = None
+    context: str | None = None
+
+
+# =============================================================================
+# Post-operation Contexts (includes results)
+# =============================================================================
+
+
+@dataclass
+class RetainResult:
+    """Result context for post-retain hook.
+
+    Contains the operation parameters and the result.
+    """
+
+    bank_id: str
+    contents: list[dict]
+    request_context: "RequestContext"
+    document_id: str | None
+    fact_type_override: str | None
+    confidence_score: float | None
+    # Result
+    unit_ids: list[list[str]]  # List of unit IDs per content item
+    success: bool = True
+    error: str | None = None
+
+
+@dataclass
+class RecallResult:
+    """Result context for post-recall hook.
+
+    Contains the operation parameters and the result.
+    """
+
+    bank_id: str
+    query: str
+    request_context: "RequestContext"
+    budget: "Budget | None"
+    max_tokens: int
+    enable_trace: bool
+    fact_types: list[str]
+    question_date: datetime | None
+    include_entities: bool
+    max_entity_tokens: int
+    include_chunks: bool
+    max_chunk_tokens: int
+    # Result
+    result: "RecallResultModel | None" = None
+    success: bool = True
+    error: str | None = None
+
+
+@dataclass
+class ReflectResultContext:
+    """Result context for post-reflect hook.
+
+    Contains the operation parameters and the result.
+    """
+
+    bank_id: str
+    query: str
+    request_context: "RequestContext"
+    budget: "Budget | None"
+    context: str | None
+    # Result
+    result: "ReflectResult | None" = None
+    success: bool = True
+    error: str | None = None
+
+
+class OperationValidatorExtension(Extension, ABC):
+    """
+    Validates and hooks into retain/recall/reflect operations.
+
+    This extension allows implementing custom logic such as:
+    - Rate limiting (pre-operation)
+    - Quota enforcement (pre-operation)
+    - Permission checks (pre-operation)
+    - Content filtering (pre-operation)
+    - Usage tracking (post-operation)
+    - Audit logging (post-operation)
+    - Metrics collection (post-operation)
+
+    Enable via environment variable:
+        HINDSIGHT_API_OPERATION_VALIDATOR_EXTENSION=mypackage.validators:MyValidator
+
+    Configuration is passed from prefixed environment variables:
+        HINDSIGHT_API_OPERATION_VALIDATOR_MAX_REQUESTS=100
+        -> config = {"max_requests": "100"}
+
+    Hook execution order:
+        1. validate_retain/validate_recall/validate_reflect (pre-operation)
+        2. [operation executes]
+        3. on_retain_complete/on_recall_complete/on_reflect_complete (post-operation)
+    """
+
+    # =========================================================================
+    # Pre-operation validation hooks (abstract - must be implemented)
+    # =========================================================================
+
+    @abstractmethod
+    async def validate_retain(self, ctx: RetainContext) -> ValidationResult:
+        """
+        Validate a retain operation before execution.
+
+        Called before the retain operation is processed. Return ValidationResult.reject()
+        to prevent the operation from executing.
+
+        Args:
+            ctx: Context containing all user-provided parameters:
+                - bank_id: Bank identifier
+                - contents: List of content dicts
+                - request_context: Request context with auth info
+                - document_id: Optional document ID
+                - fact_type_override: Optional fact type override
+                - confidence_score: Optional confidence score
+
+        Returns:
+            ValidationResult indicating whether the operation is allowed.
+        """
+        ...
+
+    @abstractmethod
+    async def validate_recall(self, ctx: RecallContext) -> ValidationResult:
+        """
+        Validate a recall operation before execution.
+
+        Called before the recall operation is processed. Return ValidationResult.reject()
+        to prevent the operation from executing.
+
+        Args:
+            ctx: Context containing all user-provided parameters:
+                - bank_id: Bank identifier
+                - query: Search query
+                - request_context: Request context with auth info
+                - budget: Budget level
+                - max_tokens: Maximum tokens to return
+                - enable_trace: Whether to include trace info
+                - fact_types: List of fact types to search
+                - question_date: Optional date context for query
+                - include_entities: Whether to include entity data
+                - max_entity_tokens: Max tokens for entities
+                - include_chunks: Whether to include chunks
+                - max_chunk_tokens: Max tokens for chunks
+
+        Returns:
+            ValidationResult indicating whether the operation is allowed.
+        """
+        ...
+
+    @abstractmethod
+    async def validate_reflect(self, ctx: ReflectContext) -> ValidationResult:
+        """
+        Validate a reflect operation before execution.
+
+        Called before the reflect operation is processed. Return ValidationResult.reject()
+        to prevent the operation from executing.
+
+        Args:
+            ctx: Context containing all user-provided parameters:
+                - bank_id: Bank identifier
+                - query: Question to answer
+                - request_context: Request context with auth info
+                - budget: Budget level
+                - context: Optional additional context
+
+        Returns:
+            ValidationResult indicating whether the operation is allowed.
+        """
+        ...
+
+    # =========================================================================
+    # Post-operation hooks (optional - override to implement)
+    # =========================================================================
+
+    async def on_retain_complete(self, result: RetainResult) -> None:
+        """
+        Called after a retain operation completes (success or failure).
+
+        Override this method to implement post-operation logic such as:
+        - Usage tracking
+        - Audit logging
+        - Metrics collection
+        - Notifications
+
+        Args:
+            result: Result context containing:
+                - All original operation parameters
+                - unit_ids: List of created unit IDs (if success)
+                - success: Whether the operation succeeded
+                - error: Error message (if failed)
+        """
+        pass
+
+    async def on_recall_complete(self, result: RecallResult) -> None:
+        """
+        Called after a recall operation completes (success or failure).
+
+        Override this method to implement post-operation logic such as:
+        - Usage tracking
+        - Audit logging
+        - Metrics collection
+        - Query analytics
+
+        Args:
+            result: Result context containing:
+                - All original operation parameters
+                - result: RecallResultModel (if success)
+                - success: Whether the operation succeeded
+                - error: Error message (if failed)
+        """
+        pass
+
+    async def on_reflect_complete(self, result: ReflectResultContext) -> None:
+        """
+        Called after a reflect operation completes (success or failure).
+
+        Override this method to implement post-operation logic such as:
+        - Usage tracking
+        - Audit logging
+        - Metrics collection
+        - Response analytics
+
+        Args:
+            result: Result context containing:
+                - All original operation parameters
+                - result: ReflectResult (if success)
+                - success: Whether the operation succeeded
+                - error: Error message (if failed)
+        """
+        pass

hindsight-api/hindsight_api/extensions/tenant.py

@@ -0,0 +1,63 @@
+"""Tenant Extension for multi-tenancy and API key authentication."""
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+
+from hindsight_api.extensions.base import Extension
+from hindsight_api.models import RequestContext
+
+
+class AuthenticationError(Exception):
+    """Raised when authentication fails."""
+
+    def __init__(self, reason: str):
+        self.reason = reason
+        super().__init__(f"Authentication failed: {reason}")
+
+
+@dataclass
+class TenantContext:
+    """
+    Tenant context returned by authentication.
+
+    Contains the PostgreSQL schema name for tenant isolation.
+    All database queries will use fully-qualified table names
+    with this schema (e.g., schema_name.memory_units).
+    """
+
+    schema_name: str
+
+
+class TenantExtension(Extension, ABC):
+    """
+    Extension for multi-tenancy and API key authentication.
+
+    This extension validates incoming requests and returns the tenant context
+    including the PostgreSQL schema to use for database operations.
+
+    Built-in implementation:
+        hindsight_api.extensions.builtin.tenant.ApiKeyTenantExtension
+
+    Enable via environment variable:
+        HINDSIGHT_API_TENANT_EXTENSION=hindsight_api.extensions.builtin.tenant:ApiKeyTenantExtension
+        HINDSIGHT_API_TENANT_API_KEY=your-secret-key
+
+    The returned schema_name is used for fully-qualified table names in queries,
+    enabling tenant isolation at the database level.
+    """
+
+    @abstractmethod
+    async def authenticate(self, context: RequestContext) -> TenantContext:
+        """
+        Authenticate the action context and return tenant context.
+
+        Args:
+            context: The action context containing API key and other auth data.
+
+        Returns:
+            TenantContext with the schema_name for database operations.
+
+        Raises:
+            AuthenticationError: If authentication fails.
+        """
+        ...

hindsight-api/hindsight_api/main.py

@@ -4,6 +4,9 @@ Command-line interface for Hindsight API.
 Run the server with:
     hindsight-api
 
+Run as background daemon:
+    hindsight-api --daemon
+
 Stop with Ctrl+C.
 """
 
@@ -21,9 +24,14 @@ from . import MemoryEngine
 from .api import create_app
 from .banner import print_banner
 from .config import HindsightConfig, get_config
-
-print()
-print_banner()
+from .daemon import (
+    DEFAULT_DAEMON_PORT,
+    DEFAULT_IDLE_TIMEOUT,
+    DaemonLock,
+    IdleTimeoutMiddleware,
+    daemonize,
+)
+from .extensions import DefaultExtensionContext, OperationValidatorExtension, TenantExtension, load_extension
 
 # Filter deprecation warnings from third-party libraries
 warnings.filterwarnings("ignore", message="websockets.legacy is deprecated")
@@ -106,8 +114,52 @@ def main():
     parser.add_argument("--ssl-keyfile", default=None, help="SSL key file")
     parser.add_argument("--ssl-certfile", default=None, help="SSL certificate file")
 
+    # Daemon mode options
+    parser.add_argument(
+        "--daemon",
+        action="store_true",
+        help=f"Run as background daemon (uses port {DEFAULT_DAEMON_PORT}, auto-exits after idle)",
+    )
+    parser.add_argument(
+        "--idle-timeout",
+        type=int,
+        default=DEFAULT_IDLE_TIMEOUT,
+        help=f"Idle timeout in seconds before auto-exit in daemon mode (default: {DEFAULT_IDLE_TIMEOUT})",
+    )
+
     args = parser.parse_args()
 
+    # Daemon mode handling
+    if args.daemon:
+        # Use fixed daemon port
+        args.port = DEFAULT_DAEMON_PORT
+        args.host = "127.0.0.1"  # Only bind to localhost for security
+
+        # Check if another daemon is already running
+        daemon_lock = DaemonLock()
+        if not daemon_lock.acquire():
+            print(f"Daemon already running (PID: {daemon_lock.get_pid()})", file=sys.stderr)
+            sys.exit(1)
+
+        # Fork into background
+        daemonize()
+
+        # Re-acquire lock in child process
+        daemon_lock = DaemonLock()
+        if not daemon_lock.acquire():
+            sys.exit(1)
+
+        # Register cleanup to release lock
+        def release_lock():
+            daemon_lock.release()
+
+        atexit.register(release_lock)
+
+    # Print banner (not in daemon mode)
+    if not args.daemon:
+        print()
+        print_banner()
+
     # Configure Python logging based on log level
     # Update config with CLI override if provided
     if args.log_level != config.log_level:
@@ -117,6 +169,16 @@ def main():
             llm_api_key=config.llm_api_key,
             llm_model=config.llm_model,
             llm_base_url=config.llm_base_url,
+            llm_max_concurrent=config.llm_max_concurrent,
+            llm_timeout=config.llm_timeout,
+            retain_llm_provider=config.retain_llm_provider,
+            retain_llm_api_key=config.retain_llm_api_key,
+            retain_llm_model=config.retain_llm_model,
+            retain_llm_base_url=config.retain_llm_base_url,
+            reflect_llm_provider=config.reflect_llm_provider,
+            reflect_llm_api_key=config.reflect_llm_api_key,
+            reflect_llm_model=config.reflect_llm_model,
+            reflect_llm_base_url=config.reflect_llm_base_url,
             embeddings_provider=config.embeddings_provider,
             embeddings_local_model=config.embeddings_local_model,
             embeddings_tei_url=config.embeddings_tei_url,
@@ -128,17 +190,58 @@ def main():
             log_level=args.log_level,
             mcp_enabled=config.mcp_enabled,
             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,
+            retain_chunk_size=config.retain_chunk_size,
+            skip_llm_verification=config.skip_llm_verification,
+            lazy_reranker=config.lazy_reranker,
+            run_migrations_on_startup=config.run_migrations_on_startup,
+            db_pool_min_size=config.db_pool_min_size,
+            db_pool_max_size=config.db_pool_max_size,
+            db_command_timeout=config.db_command_timeout,
+            db_acquire_timeout=config.db_acquire_timeout,
+            task_batch_size=config.task_batch_size,
+            task_batch_interval=config.task_batch_interval,
         )
     config.configure_logging()
-    config.log_config()
+    if not args.daemon:
+        config.log_config()
 
     # Register cleanup handlers
     atexit.register(_cleanup)
     signal.signal(signal.SIGINT, _signal_handler)
     signal.signal(signal.SIGTERM, _signal_handler)
 
+    # Load operation validator extension if configured
+    operation_validator = load_extension("OPERATION_VALIDATOR", OperationValidatorExtension)
+    if operation_validator:
+        import logging
+
+        logging.info(f"Loaded operation validator: {operation_validator.__class__.__name__}")
+
+    # Load tenant extension if configured
+    tenant_extension = load_extension("TENANT", TenantExtension)
+    if tenant_extension:
+        import logging
+
+        logging.info(f"Loaded tenant extension: {tenant_extension.__class__.__name__}")
+
     # Create MemoryEngine (reads configuration from environment)
-    _memory = MemoryEngine()
+    _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:
+        extension_context = DefaultExtensionContext(
+            database_url=config.database_url,
+            memory_engine=_memory,
+        )
+        tenant_extension.set_context(extension_context)
+        logging.info("Extension context set on tenant extension")
 
     # Create FastAPI app
     app = create_app(
@@ -149,6 +252,12 @@ def main():
         initialize_memory=True,
     )
 
+    # Wrap with idle timeout middleware in daemon mode
+    idle_middleware = None
+    if args.daemon:
+        idle_middleware = IdleTimeoutMiddleware(app, idle_timeout=args.idle_timeout)
+        app = idle_middleware
+
     # Prepare uvicorn config
     uvicorn_config = {
         "app": app,
@@ -172,20 +281,40 @@ def main():
     if args.ssl_certfile:
         uvicorn_config["ssl_certfile"] = args.ssl_certfile
 
-    from .banner import print_startup_info
+    # Print startup info (not in daemon mode)
+    if not args.daemon:
+        from .banner import print_startup_info
 
-    print_startup_info(
-        host=args.host,
-        port=args.port,
-        database_url=config.database_url,
-        llm_provider=config.llm_provider,
-        llm_model=config.llm_model,
-        embeddings_provider=config.embeddings_provider,
-        reranker_provider=config.reranker_provider,
-        mcp_enabled=config.mcp_enabled,
-    )
+        print_startup_info(
+            host=args.host,
+            port=args.port,
+            database_url=config.database_url,
+            llm_provider=config.llm_provider,
+            llm_model=config.llm_model,
+            embeddings_provider=config.embeddings_provider,
+            reranker_provider=config.reranker_provider,
+            mcp_enabled=config.mcp_enabled,
+        )
 
-    uvicorn.run(**uvicorn_config)
+    # Start idle checker in daemon mode
+    if idle_middleware is not None:
+        # Start the idle checker in a background thread with its own event loop
+        import threading
+
+        def run_idle_checker():
+            import time
+
+            time.sleep(2)  # Wait for uvicorn to start
+            try:
+                loop = asyncio.new_event_loop()
+                asyncio.set_event_loop(loop)
+                loop.run_until_complete(idle_middleware._check_idle())
+            except Exception:
+                pass
+
+        threading.Thread(target=run_idle_checker, daemon=True).start()
+
+    uvicorn.run(**uvicorn_config)  # type: ignore[invalid-argument-type] - dict kwargs
 
 
 if __name__ == "__main__":

hindsight-api/hindsight_api/mcp_local.py

@@ -87,6 +87,7 @@ def create_local_mcp_server(bank_id: str, memory=None) -> FastMCP:
     from hindsight_api import MemoryEngine
     from hindsight_api.engine.memory_engine import Budget
     from hindsight_api.engine.response_models import VALID_RECALL_FACT_TYPES
+    from hindsight_api.models import RequestContext
 
     # Create memory engine with pg0 embedded database if not provided
     if memory is None:
@@ -115,7 +116,11 @@ def create_local_mcp_server(bank_id: str, memory=None) -> FastMCP:
 
         async def _retain():
             try:
-                await memory.retain_batch_async(bank_id=bank_id, contents=[{"content": content, "context": context}])
+                await memory.retain_batch_async(
+                    bank_id=bank_id,
+                    contents=[{"content": content, "context": context}],
+                    request_context=RequestContext(),
+                )
             except Exception as e:
                 logger.error(f"Error storing memory: {e}", exc_info=True)
 
@@ -142,6 +147,7 @@ def create_local_mcp_server(bank_id: str, memory=None) -> FastMCP:
                 fact_type=list(VALID_RECALL_FACT_TYPES),
                 budget=budget_enum,
                 max_tokens=max_tokens,
+                request_context=RequestContext(),
             )
 
             return search_result.model_dump()

hindsight-api/hindsight_api/metrics.py

@@ -5,6 +5,7 @@ This module provides metrics for:
 - Operation latency (retain, recall, reflect) with percentiles
 - Token usage (input/output) per operation
 - Per-bank granularity via labels
+- LLM call latency and token usage with scope dimension
 """
 
 import logging
@@ -14,8 +15,54 @@ from contextlib import contextmanager
 from opentelemetry import metrics
 from opentelemetry.exporter.prometheus import PrometheusMetricReader
 from opentelemetry.sdk.metrics import MeterProvider
+from opentelemetry.sdk.metrics.view import ExplicitBucketHistogramAggregation, View
 from opentelemetry.sdk.resources import Resource
 
+# Custom bucket boundaries for operation duration (in seconds)
+# Fine granularity in 0-30s range where most operations complete
+DURATION_BUCKETS = (0.1, 0.25, 0.5, 0.75, 1.0, 2.0, 3.0, 5.0, 7.5, 10.0, 15.0, 20.0, 30.0, 60.0, 120.0)
+
+# LLM duration buckets (finer granularity for faster LLM calls)
+LLM_DURATION_BUCKETS = (0.1, 0.25, 0.5, 1.0, 2.0, 3.0, 5.0, 10.0, 15.0, 30.0, 60.0, 120.0)
+
+
+def get_token_bucket(token_count: int) -> str:
+    """
+    Convert a token count to a bucket label for use as a dimension.
+
+    This allows analyzing token usage patterns without high-cardinality issues.
+
+    Buckets:
+    - "0-100": Very small requests/responses
+    - "100-500": Small requests/responses
+    - "500-1k": Medium requests/responses
+    - "1k-5k": Large requests/responses
+    - "5k-10k": Very large requests/responses
+    - "10k-50k": Huge requests/responses
+    - "50k+": Extremely large requests/responses
+
+    Args:
+        token_count: Number of tokens
+
+    Returns:
+        Bucket label string
+    """
+    if token_count < 100:
+        return "0-100"
+    elif token_count < 500:
+        return "100-500"
+    elif token_count < 1000:
+        return "500-1k"
+    elif token_count < 5000:
+        return "1k-5k"
+    elif token_count < 10000:
+        return "5k-10k"
+    elif token_count < 50000:
+        return "10k-50k"
+    else:
+        return "50k+"
+
+
 logger = logging.getLogger(__name__)
 
 # Global meter instance
@@ -48,8 +95,22 @@ def initialize_metrics(service_name: str = "hindsight-api", service_version: str
     # Create Prometheus metric reader
     prometheus_reader = PrometheusMetricReader()
 
-    # Create meter provider with Prometheus exporter
-    provider = MeterProvider(resource=resource, metric_readers=[prometheus_reader])
+    # Create view with custom bucket boundaries for duration histogram
+    duration_view = View(
+        instrument_name="hindsight.operation.duration",
+        aggregation=ExplicitBucketHistogramAggregation(boundaries=DURATION_BUCKETS),
+    )
+
+    # Create view with custom bucket boundaries for LLM duration histogram
+    llm_duration_view = View(
+        instrument_name="hindsight.llm.duration",
+        aggregation=ExplicitBucketHistogramAggregation(boundaries=LLM_DURATION_BUCKETS),
+    )
+
+    # Create meter provider with Prometheus exporter and custom views
+    provider = MeterProvider(
+        resource=resource, metric_readers=[prometheus_reader], views=[duration_view, llm_duration_view]
+    )
 
     # Set the global meter provider
     metrics.set_meter_provider(provider)
@@ -71,20 +132,39 @@ class MetricsCollectorBase:
     """Base class for metrics collectors."""
 
     @contextmanager
-    def record_operation(self, operation: str, bank_id: str, budget: str | None = None, max_tokens: int | None = None):
-        """Context manager to record operation duration and status."""
-        raise NotImplementedError
-
-    def record_tokens(
+    def record_operation(
         self,
         operation: str,
         bank_id: str,
-        input_tokens: int = 0,
-        output_tokens: int = 0,
+        source: str = "api",
         budget: str | None = None,
         max_tokens: int | None = None,
     ):
-        """Record token usage for an operation."""
+        """Context manager to record operation duration and status."""
+        raise NotImplementedError
+
+    def record_llm_call(
+        self,
+        provider: str,
+        model: str,
+        scope: str,
+        duration: float,
+        input_tokens: int = 0,
+        output_tokens: int = 0,
+        success: bool = True,
+    ):
+        """
+        Record metrics for an LLM call.
+
+        Args:
+            provider: LLM provider name (openai, anthropic, gemini, groq, ollama, lmstudio)
+            model: Model name
+            scope: Scope identifier (e.g., "memory", "reflect", "entity_observation")
+            duration: Call duration in seconds
+            input_tokens: Number of input/prompt tokens
+            output_tokens: Number of output/completion tokens
+            success: Whether the call was successful
+        """
         raise NotImplementedError
 
 
@@ -92,20 +172,28 @@ class NoOpMetricsCollector(MetricsCollectorBase):
     """No-op metrics collector that does nothing. Used when metrics are disabled."""
 
     @contextmanager
-    def record_operation(self, operation: str, bank_id: str, budget: str | None = None, max_tokens: int | None = None):
-        """No-op context manager."""
-        yield
-
-    def record_tokens(
+    def record_operation(
         self,
         operation: str,
         bank_id: str,
-        input_tokens: int = 0,
-        output_tokens: int = 0,
+        source: str = "api",
         budget: str | None = None,
         max_tokens: int | None = None,
     ):
-        """No-op token recording."""
+        """No-op context manager."""
+        yield
+
+    def record_llm_call(
+        self,
+        provider: str,
+        model: str,
+        scope: str,
+        duration: float,
+        input_tokens: int = 0,
+        output_tokens: int = 0,
+        success: bool = True,
+    ):
+        """No-op LLM call recording."""
         pass
 
 
@@ -125,33 +213,52 @@ class MetricsCollector(MetricsCollectorBase):
             name="hindsight.operation.duration", description="Duration of Hindsight operations in seconds", unit="s"
         )
 
-        # Token usage counters
-        self.tokens_input = self.meter.create_counter(
-            name="hindsight.tokens.input", description="Number of input tokens consumed", unit="tokens"
-        )
-
-        self.tokens_output = self.meter.create_counter(
-            name="hindsight.tokens.output", description="Number of output tokens generated", unit="tokens"
-        )
-
         # Operation counter (success/failure)
         self.operation_total = self.meter.create_counter(
             name="hindsight.operation.total", description="Total number of operations executed", unit="operations"
         )
 
+        # LLM call latency histogram (in seconds)
+        # Records duration of LLM API calls with provider, model, and scope dimensions
+        self.llm_duration = self.meter.create_histogram(
+            name="hindsight.llm.duration", description="Duration of LLM API calls in seconds", unit="s"
+        )
+
+        # LLM token usage counters with bucket labels
+        self.llm_tokens_input = self.meter.create_counter(
+            name="hindsight.llm.tokens.input", description="Number of input tokens for LLM calls", unit="tokens"
+        )
+
+        self.llm_tokens_output = self.meter.create_counter(
+            name="hindsight.llm.tokens.output", description="Number of output tokens from LLM calls", unit="tokens"
+        )
+
+        # LLM call counter (success/failure)
+        self.llm_calls_total = self.meter.create_counter(
+            name="hindsight.llm.calls.total", description="Total number of LLM API calls", unit="calls"
+        )
+
     @contextmanager
-    def record_operation(self, operation: str, bank_id: str, budget: str | None = None, max_tokens: int | None = None):
+    def record_operation(
+        self,
+        operation: str,
+        bank_id: str,
+        source: str = "api",
+        budget: str | None = None,
+        max_tokens: int | None = None,
+    ):
         """
         Context manager to record operation duration and status.
 
         Usage:
-            with metrics.record_operation("recall", bank_id="user123", budget="mid", max_tokens=4096):
+            with metrics.record_operation("recall", bank_id="user123", source="api", budget="mid", max_tokens=4096):
                 # ... perform operation
                 pass
 
         Args:
-            operation: Operation name (retain, recall, reflect)
+            operation: Operation name (retain, recall, reflect, entity_observation)
             bank_id: Memory bank ID
+            source: Source of the operation (api, reflect, internal)
             budget: Optional budget level (low, mid, high)
             max_tokens: Optional max tokens for the operation
         """
@@ -159,6 +266,7 @@ class MetricsCollector(MetricsCollectorBase):
         attributes = {
             "operation": operation,
             "bank_id": bank_id,
+            "source": source,
         }
         if budget:
             attributes["budget"] = budget
@@ -181,40 +289,56 @@ class MetricsCollector(MetricsCollectorBase):
             # Record operation count
             self.operation_total.add(1, attributes)
 
-    def record_tokens(
+    def record_llm_call(
         self,
-        operation: str,
-        bank_id: str,
+        provider: str,
+        model: str,
+        scope: str,
+        duration: float,
         input_tokens: int = 0,
         output_tokens: int = 0,
-        budget: str | None = None,
-        max_tokens: int | None = None,
+        success: bool = True,
     ):
         """
-        Record token usage for an operation.
+        Record metrics for an LLM call.
 
         Args:
-            operation: Operation name (retain, recall, reflect)
-            bank_id: Memory bank ID
-            input_tokens: Number of input tokens
-            output_tokens: Number of output tokens
-            budget: Optional budget level
-            max_tokens: Optional max tokens for the operation
+            provider: LLM provider name (openai, anthropic, gemini, groq, ollama, lmstudio)
+            model: Model name
+            scope: Scope identifier (e.g., "memory", "reflect", "entity_observation")
+            duration: Call duration in seconds
+            input_tokens: Number of input/prompt tokens
+            output_tokens: Number of output/completion tokens
+            success: Whether the call was successful
         """
-        attributes = {
-            "operation": operation,
-            "bank_id": bank_id,
+        # Base attributes for all metrics
+        base_attributes = {
+            "provider": provider,
+            "model": model,
+            "scope": scope,
+            "success": str(success).lower(),
         }
-        if budget:
-            attributes["budget"] = budget
-        if max_tokens:
-            attributes["max_tokens"] = str(max_tokens)
 
+        # Record duration
+        self.llm_duration.record(duration, base_attributes)
+
+        # Record call count
+        self.llm_calls_total.add(1, base_attributes)
+
+        # Record tokens with bucket labels for cardinality control
         if input_tokens > 0:
-            self.tokens_input.add(input_tokens, attributes)
+            input_attributes = {
+                **base_attributes,
+                "token_bucket": get_token_bucket(input_tokens),
+            }
+            self.llm_tokens_input.add(input_tokens, input_attributes)
 
         if output_tokens > 0:
-            self.tokens_output.add(output_tokens, attributes)
+            output_attributes = {
+                **base_attributes,
+                "token_bucket": get_token_bucket(output_tokens),
+            }
+            self.llm_tokens_output.add(output_tokens, output_attributes)
 
 
 # Global metrics collector instance (defaults to no-op)

hindsight-api/hindsight_api/migrations.py

@@ -6,12 +6,16 @@ on application startup. It is designed to be safe for concurrent
 execution using PostgreSQL advisory locks to coordinate between
 distributed workers.
 
+Supports multi-tenant schema isolation: migrations can target a specific
+PostgreSQL schema, allowing each tenant to have isolated tables.
+
 Important: All migrations must be backward-compatible to allow
 safe rolling deployments.
 
 No alembic.ini required - all configuration is done programmatically.
 """
 
+import hashlib
 import logging
 import os
 from pathlib import Path
@@ -26,11 +30,29 @@ logger = logging.getLogger(__name__)
 MIGRATION_LOCK_ID = 123456789
 
 
-def _run_migrations_internal(database_url: str, script_location: str) -> None:
+def _get_schema_lock_id(schema: str) -> int:
+    """
+    Generate a unique advisory lock ID for a schema.
+
+    Uses hash of schema name to create a deterministic lock ID.
+    """
+    # Use hash to create a unique lock ID per schema
+    # Keep within PostgreSQL's bigint range
+    hash_bytes = hashlib.sha256(schema.encode()).digest()[:8]
+    return int.from_bytes(hash_bytes, byteorder="big") % (2**31)
+
+
+def _run_migrations_internal(database_url: str, script_location: str, schema: str | None = None) -> None:
     """
     Internal function to run migrations without locking.
+
+    Args:
+        database_url: SQLAlchemy database URL
+        script_location: Path to alembic scripts
+        schema: Target schema (None for default/public)
     """
-    logger.info("Running database migrations to head...")
+    schema_name = schema or "public"
+    logger.info(f"Running database migrations to head for schema '{schema_name}'...")
     logger.info(f"Database URL: {database_url}")
     logger.info(f"Script location: {script_location}")
 
@@ -50,13 +72,22 @@ def _run_migrations_internal(database_url: str, script_location: str) -> None:
     # Set path_separator to avoid deprecation warning
     alembic_cfg.set_main_option("path_separator", "os")
 
-    # Run migrations to head (latest version)
+    # If targeting a specific schema, pass it to env.py via config
+    # env.py will handle setting search_path and version_table_schema
+    if schema:
+        alembic_cfg.set_main_option("target_schema", schema)
+
+    # Run migrations
     command.upgrade(alembic_cfg, "head")
 
-    logger.info("Database migrations completed successfully")
+    logger.info(f"Database migrations completed successfully for schema '{schema_name}'")
 
 
-def run_migrations(database_url: str, script_location: str | None = None) -> None:
+def run_migrations(
+    database_url: str,
+    script_location: str | None = None,
+    schema: str | None = None,
+) -> None:
     """
     Run database migrations to the latest version using programmatic Alembic configuration.
 
@@ -65,19 +96,28 @@ def run_migrations(database_url: str, script_location: str | None = None) -> Non
     - Other workers wait for the lock, then verify migrations are complete
     - If schema is already up-to-date, this is a fast no-op
 
+    Supports multi-tenant schema isolation: when a schema is specified, migrations
+    run in that schema instead of public. This allows tenant extensions to provision
+    new tenant schemas with their own isolated tables.
+
     Args:
         database_url: SQLAlchemy database URL (e.g., "postgresql://user:pass@host/db")
         script_location: Path to alembic migrations directory (e.g., "/path/to/alembic").
                         If None, defaults to hindsight-api/alembic directory.
+        schema: Target PostgreSQL schema name. If None, uses default (public).
+                When specified, creates the schema if needed and runs migrations there.
 
     Raises:
         RuntimeError: If migrations fail to complete
         FileNotFoundError: If script_location doesn't exist
 
     Example:
-        # Using default location (hindsight_api package)
+        # Using default location and public schema
         run_migrations("postgresql://user:pass@host/db")
 
+        # Run migrations for a specific tenant schema
+        run_migrations("postgresql://user:pass@host/db", schema="tenant_acme")
+
         # Using custom location (when importing from another project)
         run_migrations(
             "postgresql://user:pass@host/db",
@@ -99,21 +139,25 @@ def run_migrations(database_url: str, script_location: str | None = None) -> Non
                 f"Alembic script location not found at {script_location}. Database migrations cannot be run."
             )
 
+        # Use schema-specific lock ID for multi-tenant isolation
+        lock_id = _get_schema_lock_id(schema) if schema else MIGRATION_LOCK_ID
+        schema_name = schema or "public"
+
         # Use PostgreSQL advisory lock to coordinate between distributed workers
         engine = create_engine(database_url)
         with engine.connect() as conn:
             # pg_advisory_lock blocks until the lock is acquired
             # The lock is automatically released when the connection closes
-            logger.debug(f"Acquiring migration advisory lock (id={MIGRATION_LOCK_ID})...")
-            conn.execute(text(f"SELECT pg_advisory_lock({MIGRATION_LOCK_ID})"))
+            logger.debug(f"Acquiring migration advisory lock for schema '{schema_name}' (id={lock_id})...")
+            conn.execute(text(f"SELECT pg_advisory_lock({lock_id})"))
             logger.debug("Migration advisory lock acquired")
 
             try:
                 # Run migrations while holding the lock
-                _run_migrations_internal(database_url, script_location)
+                _run_migrations_internal(database_url, script_location, schema=schema)
             finally:
                 # Explicitly release the lock (also released on connection close)
-                conn.execute(text(f"SELECT pg_advisory_unlock({MIGRATION_LOCK_ID})"))
+                conn.execute(text(f"SELECT pg_advisory_unlock({lock_id})"))
                 logger.debug("Migration advisory lock released")
 
     except FileNotFoundError:
@@ -185,3 +229,131 @@ def check_migration_status(
     except Exception as e:
         logger.warning(f"Unable to check migration status: {e}")
         return None, None
+
+
+def ensure_embedding_dimension(
+    database_url: str,
+    required_dimension: int,
+    schema: str | None = None,
+) -> None:
+    """
+    Ensure the embedding column dimension matches the model's dimension.
+
+    This function checks the current vector column dimension in the database
+    and adjusts it if necessary:
+    - If dimensions match: no action needed
+    - If dimensions differ and table is empty: ALTER COLUMN to new dimension
+    - If dimensions differ and table has data: raise error with migration guidance
+
+    Args:
+        database_url: SQLAlchemy database URL
+        required_dimension: The embedding dimension required by the model
+        schema: Target PostgreSQL schema name (None for public)
+
+    Raises:
+        RuntimeError: If dimension mismatch with existing data
+    """
+    schema_name = schema or "public"
+
+    engine = create_engine(database_url)
+    with engine.connect() as conn:
+        # Check if memory_units table exists
+        table_exists = conn.execute(
+            text("""
+                SELECT EXISTS (
+                    SELECT 1 FROM information_schema.tables
+                    WHERE table_schema = :schema AND table_name = 'memory_units'
+                )
+            """),
+            {"schema": schema_name},
+        ).scalar()
+
+        if not table_exists:
+            logger.debug(f"memory_units table does not exist in schema '{schema_name}', skipping dimension check")
+            return
+
+        # Get current column dimension from pg_attribute
+        # pgvector stores dimension in atttypmod
+        current_dim = conn.execute(
+            text("""
+                SELECT atttypmod
+                FROM pg_attribute a
+                JOIN pg_class c ON a.attrelid = c.oid
+                JOIN pg_namespace n ON c.relnamespace = n.oid
+                WHERE n.nspname = :schema
+                  AND c.relname = 'memory_units'
+                  AND a.attname = 'embedding'
+            """),
+            {"schema": schema_name},
+        ).scalar()
+
+        if current_dim is None:
+            logger.warning("Could not determine current embedding dimension, skipping check")
+            return
+
+        # pgvector stores dimension directly in atttypmod (no offset like other types)
+        current_dimension = current_dim
+
+        if current_dimension == required_dimension:
+            logger.debug(f"Embedding dimension OK: {current_dimension}")
+            return
+
+        logger.info(
+            f"Embedding dimension mismatch: database has {current_dimension}, model requires {required_dimension}"
+        )
+
+        # Check if table has data
+        row_count = conn.execute(
+            text(f"SELECT COUNT(*) FROM {schema_name}.memory_units WHERE embedding IS NOT NULL")
+        ).scalar()
+
+        if row_count > 0:
+            raise RuntimeError(
+                f"Cannot change embedding dimension from {current_dimension} to {required_dimension}: "
+                f"memory_units table contains {row_count} rows with embeddings. "
+                f"To change dimensions, you must either:\n"
+                f"  1. Re-embed all data: DELETE FROM {schema_name}.memory_units; then restart\n"
+                f"  2. Use a model with {current_dimension}-dimensional embeddings"
+            )
+
+        # Table is empty, safe to alter column
+        logger.info(f"Altering embedding column dimension from {current_dimension} to {required_dimension}")
+
+        # Drop the HNSW index on embedding column if it exists
+        # Only drop indexes that use 'hnsw' and reference the 'embedding' column
+        conn.execute(
+            text(f"""
+                DO $$
+                DECLARE idx_name TEXT;
+                BEGIN
+                    FOR idx_name IN
+                        SELECT indexname FROM pg_indexes
+                        WHERE schemaname = '{schema_name}'
+                          AND tablename = 'memory_units'
+                          AND indexdef LIKE '%hnsw%'
+                          AND indexdef LIKE '%embedding%'
+                    LOOP
+                        EXECUTE 'DROP INDEX IF EXISTS {schema_name}.' || idx_name;
+                    END LOOP;
+                END $$;
+            """)
+        )
+
+        # Alter the column type
+        conn.execute(
+            text(f"ALTER TABLE {schema_name}.memory_units ALTER COLUMN embedding TYPE vector({required_dimension})")
+        )
+        conn.commit()
+
+        # Recreate the HNSW index
+        conn.execute(
+            text(f"""
+                CREATE INDEX IF NOT EXISTS idx_memory_units_embedding_hnsw
+                ON {schema_name}.memory_units
+                USING hnsw (embedding vector_cosine_ops)
+                WITH (m = 16, ef_construction = 64)
+            """)
+        )
+        conn.commit()
+
+        logger.info(f"Successfully changed embedding dimension to {required_dimension}")

hindsight-api/hindsight_api/models.py

@@ -2,9 +2,27 @@
 SQLAlchemy models for the memory system.
 """
 
+from dataclasses import dataclass
 from datetime import datetime
 from uuid import UUID as PyUUID
 
+
+@dataclass
+class RequestContext:
+    """
+    Context for request authentication and authorization.
+
+    This dataclass carries authentication data from HTTP requests to the
+    memory engine operations. It can be extended to include additional
+    context like headers, tokens, user info, etc.
+    """
+
+    api_key: str | None = None
+    api_key_id: str | None = None  # UUID of the API key used for authentication
+    tenant_id: str | None = None  # Tenant identifier (set by extension after auth)
+    internal: bool = False  # True for background/internal operations (not user-visible)
+
+
 from pgvector.sqlalchemy import Vector
 from sqlalchemy import (
     CheckConstraint,
@@ -23,6 +41,8 @@ from sqlalchemy.dialects.postgresql import JSONB, TIMESTAMP, UUID
 from sqlalchemy.ext.asyncio import AsyncAttrs
 from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, relationship
 
+from .config import EMBEDDING_DIMENSION
+
 
 class Base(AsyncAttrs, DeclarativeBase):
     """Base class for all models."""
@@ -63,7 +83,7 @@ class MemoryUnit(Base):
     bank_id: Mapped[str] = mapped_column(Text, nullable=False)
     document_id: Mapped[str | None] = mapped_column(Text)
     text: Mapped[str] = mapped_column(Text, nullable=False)
-    embedding = mapped_column(Vector(384))  # pgvector type
+    embedding = mapped_column(Vector(EMBEDDING_DIMENSION))  # pgvector type
     context: Mapped[str | None] = mapped_column(Text)
     event_date: Mapped[datetime] = mapped_column(
         TIMESTAMP(timezone=True), nullable=False

hindsight-api/hindsight_api/pg0.py

@@ -40,7 +40,7 @@ class EmbeddedPostgres:
             # Only set port if explicitly specified
             if self.port is not None:
                 kwargs["port"] = self.port
-            self._pg0 = Pg0(**kwargs)
+            self._pg0 = Pg0(**kwargs)  # type: ignore[invalid-argument-type] - dict kwargs
         return self._pg0
 
     async def start(self, max_retries: int = 5, retry_delay: float = 4.0) -> str:
@@ -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

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "hindsight-api"
-version = "0.1.11"
+version = "0.2.1"
 description = "Hindsight: Agent Memory That Works Like Human Memory"
 readme = "README.md"
 requires-python = ">=3.11"
@@ -14,7 +14,6 @@ dependencies = [
     "openai>=1.0.0",
     "pydantic>=2.0.0",
     "rich>=13.0.0",
-    "sentence-transformers>=3.0.0,<3.3.0",
     "langchain-text-splitters>=0.3.0",
     "fastapi[standard]>=0.120.3",
     "uvicorn>=0.38.0",
@@ -24,8 +23,6 @@ dependencies = [
     "pgvector>=0.4.1",
     "greenlet>=3.2.4",
     "psycopg2-binary>=2.9.11",
-    "transformers>=4.30.0,<4.46.0",
-    "torch>=2.0.0",
     "tiktoken>=0.12.0",
     "httpx>=0.27.0",
     "fastmcp>=2.3.0",
@@ -37,6 +34,13 @@ dependencies = [
     "opentelemetry-exporter-prometheus>=0.41b0",
     "dateparser>=1.2.2",
     "google-genai>=1.0.0",
+    "anthropic>=0.40.0",
+    "typer>=0.9.0",
+    "cohere>=5.0.0",
+    # Local ML models for embeddings/reranking - can be excluded in Docker with INCLUDE_LOCAL_MODELS=false
+    "sentence-transformers>=3.0.0,<3.3.0",
+    "transformers>=4.30.0,<4.46.0",
+    "torch>=2.0.0",
 ]
 
 [project.optional-dependencies]
@@ -51,6 +55,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"]
@@ -74,7 +79,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
@@ -92,6 +97,7 @@ dev = [
     "python-dotenv>=1.2.1",
     "filelock>=3.0.0",
     "ruff>=0.8.0",
+    "ty>=0.0.1",
 ]
 
 [tool.ruff]
@@ -121,3 +127,28 @@ ignore = [
 [tool.ruff.format]
 quote-style = "double"
 indent-style = "space"
+
+[tool.ty]
+# Type checking configuration
+# ty is an extremely fast Python type checker from Astral (same team as ruff/uv)
+
+[tool.ty.environment]
+python-version = "3.11"
+
+[tool.ty.src]
+exclude = [
+    "tests/",
+    "hindsight_api/alembic/",
+]
+
+[tool.ty.rules]
+# Disable noisy rules while keeping important ones
+invalid-argument-type = "ignore"    # False positives with **kwargs patterns
+invalid-return-type = "ignore"      # Often intentional in async code
+invalid-parameter-default = "ignore" # Optional params with None default
+possibly-missing-attribute = "ignore" # Common with Optional types
+invalid-raise = "ignore"            # False positives with exception tracking
+call-non-callable = "ignore"        # False positives with Optional types
+invalid-key = "ignore"              # Pydantic ConfigDict not understood
+invalid-method-override = "ignore"  # Intentional signature differences
+unresolved-reference = "ignore"     # Forward references not always resolved

hindsight-api/tests/conftest.py

@@ -8,7 +8,7 @@ import os
 import filelock
 from pathlib import Path
 from dotenv import load_dotenv
-from hindsight_api import MemoryEngine, LLMConfig, LocalSTEmbeddings
+from hindsight_api import MemoryEngine, LLMConfig, LocalSTEmbeddings, RequestContext
 
 from hindsight_api.engine.cross_encoder import LocalSTCrossEncoder
 from hindsight_api.engine.query_analyzer import DateparserQueryAnalyzer
@@ -99,6 +99,12 @@ def pg0_db_url(db_url, tmp_path_factory, worker_id):
     return url
 
 
+@pytest.fixture(scope="function")
+def request_context():
+    """Provide a default RequestContext for tests."""
+    return RequestContext()
+
+
 @pytest.fixture(scope="session")
 def llm_config():
     """

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_agents_api.py

@@ -3,7 +3,7 @@ Tests for agent management API (profile, disposition, background).
 """
 import pytest
 import uuid
-from hindsight_api import MemoryEngine
+from hindsight_api import MemoryEngine, RequestContext
 from hindsight_api.api import CreateBankRequest, DispositionTraits
 from hindsight_api.engine.memory_engine import Budget
 
@@ -17,11 +17,11 @@ class TestAgentProfile:
     """Tests for agent profile management."""
 
     @pytest.mark.asyncio
-    async def test_get_agent_profile_creates_default(self, memory: MemoryEngine):
+    async def test_get_agent_profile_creates_default(self, memory: MemoryEngine, request_context):
         """Test that getting a profile for a new agent creates default disposition."""
         bank_id = unique_agent_id("test_profile_default")
 
-        profile = await memory.get_bank_profile(bank_id)
+        profile = await memory.get_bank_profile(bank_id, request_context=request_context)
 
         assert profile is not None
         assert "disposition" in profile
@@ -35,11 +35,11 @@ class TestAgentProfile:
         assert profile["background"] == ""
 
     @pytest.mark.asyncio
-    async def test_update_agent_disposition(self, memory: MemoryEngine):
+    async def test_update_agent_disposition(self, memory: MemoryEngine, request_context):
         """Test updating agent disposition traits."""
         bank_id = unique_agent_id("test_profile_update")
 
-        profile = await memory.get_bank_profile(bank_id)
+        profile = await memory.get_bank_profile(bank_id, request_context=request_context)
         assert profile["disposition"].skepticism == 3
 
         new_disposition = {
@@ -47,26 +47,26 @@ class TestAgentProfile:
             "literalism": 4,
             "empathy": 2,
         }
-        await memory.update_bank_disposition(bank_id, new_disposition)
+        await memory.update_bank_disposition(bank_id, new_disposition, request_context=request_context)
 
-        updated_profile = await memory.get_bank_profile(bank_id)
+        updated_profile = await memory.get_bank_profile(bank_id, request_context=request_context)
         disposition = updated_profile["disposition"]
         assert disposition.skepticism == new_disposition["skepticism"]
         assert disposition.literalism == new_disposition["literalism"]
         assert disposition.empathy == new_disposition["empathy"]
 
     @pytest.mark.asyncio
-    async def test_list_agents(self, memory: MemoryEngine):
+    async def test_list_agents(self, memory: MemoryEngine, request_context):
         """Test listing all agents."""
         agent_id_1 = unique_agent_id("test_list")
         agent_id_2 = unique_agent_id("test_list")
         agent_id_3 = unique_agent_id("test_list")
 
-        await memory.get_bank_profile(agent_id_1)
-        await memory.get_bank_profile(agent_id_2)
-        await memory.get_bank_profile(agent_id_3)
+        await memory.get_bank_profile(agent_id_1, request_context=request_context)
+        await memory.get_bank_profile(agent_id_2, request_context=request_context)
+        await memory.get_bank_profile(agent_id_3, request_context=request_context)
 
-        agents = await memory.list_banks()
+        agents = await memory.list_banks(request_context=request_context)
 
         agent_ids = [a["bank_id"] for a in agents]
         assert agent_id_1 in agent_ids
@@ -85,46 +85,50 @@ class TestAgentBackground:
     """Tests for agent background management."""
 
     @pytest.mark.asyncio
-    async def test_merge_agent_background(self, memory: MemoryEngine):
+    async def test_merge_agent_background(self, memory: MemoryEngine, request_context):
         """Test merging agent background information."""
         bank_id = unique_agent_id("test_profile_merge")
 
-        profile = await memory.get_bank_profile(bank_id)
+        profile = await memory.get_bank_profile(bank_id, request_context=request_context)
         assert profile["background"] == ""
 
         result1 = await memory.merge_bank_background(
             bank_id,
             "I was born in Texas",
-            update_disposition=False
+            update_disposition=False,
+            request_context=request_context,
         )
         assert "Texas" in result1["background"]
 
         result2 = await memory.merge_bank_background(
             bank_id,
             "I have 10 years of startup experience",
-            update_disposition=False
+            update_disposition=False,
+            request_context=request_context,
         )
         assert "Texas" in result2["background"] or "startup" in result2["background"]
 
-        final_profile = await memory.get_bank_profile(bank_id)
+        final_profile = await memory.get_bank_profile(bank_id, request_context=request_context)
         assert final_profile["background"] != ""
 
     @pytest.mark.asyncio
-    async def test_merge_background_handles_conflicts(self, memory: MemoryEngine):
+    async def test_merge_background_handles_conflicts(self, memory: MemoryEngine, request_context):
         """Test that merging background handles conflicts (new overwrites old)."""
         bank_id = unique_agent_id("test_profile_conflict")
 
         result1 = await memory.merge_bank_background(
             bank_id,
             "I was born in Colorado",
-            update_disposition=False
+            update_disposition=False,
+            request_context=request_context,
         )
         assert "Colorado" in result1["background"]
 
         result2 = await memory.merge_bank_background(
             bank_id,
             "You were born in Texas",
-            update_disposition=False
+            update_disposition=False,
+            request_context=request_context,
         )
         assert "Texas" in result2["background"]
 
@@ -133,7 +137,7 @@ class TestAgentEndpoint:
     """Tests for agent PUT endpoint logic."""
 
     @pytest.mark.asyncio
-    async def test_put_agent_create(self, memory: MemoryEngine):
+    async def test_put_agent_create(self, memory: MemoryEngine, request_context):
         """Test creating an agent via PUT endpoint."""
         bank_id = unique_agent_id("test_put_create")
 
@@ -146,12 +150,13 @@ class TestAgentEndpoint:
             background="I am a creative software engineer"
         )
 
-        profile = await memory.get_bank_profile(bank_id)
+        profile = await memory.get_bank_profile(bank_id, request_context=request_context)
 
         if request.disposition is not None:
             await memory.update_bank_disposition(
                 bank_id,
-                request.disposition.model_dump()
+                request.disposition.model_dump(),
+                request_context=request_context,
             )
 
         if request.background is not None:
@@ -168,14 +173,14 @@ class TestAgentEndpoint:
                     request.background
                 )
 
-        final_profile = await memory.get_bank_profile(bank_id)
+        final_profile = await memory.get_bank_profile(bank_id, request_context=request_context)
 
         assert final_profile["disposition"].skepticism == 4
         assert final_profile["disposition"].literalism == 5
         assert final_profile["background"] == "I am a creative software engineer"
 
     @pytest.mark.asyncio
-    async def test_put_agent_partial_update(self, memory: MemoryEngine):
+    async def test_put_agent_partial_update(self, memory: MemoryEngine, request_context):
         """Test updating only background."""
         bank_id = unique_agent_id("test_put_partial")
 
@@ -183,7 +188,7 @@ class TestAgentEndpoint:
             background="I am a data scientist"
         )
 
-        profile = await memory.get_bank_profile(bank_id)
+        profile = await memory.get_bank_profile(bank_id, request_context=request_context)
 
         if request.background is not None:
             pool = await memory._get_pool()
@@ -199,7 +204,7 @@ class TestAgentEndpoint:
                     request.background
                 )
 
-        final_profile = await memory.get_bank_profile(bank_id)
+        final_profile = await memory.get_bank_profile(bank_id, request_context=request_context)
 
         assert final_profile["disposition"].skepticism == 3  # Default
         assert final_profile["background"] == "I am a data scientist"
@@ -209,7 +214,7 @@ class TestAgentDispositionIntegration:
     """Tests for disposition integration with other features."""
 
     @pytest.mark.asyncio
-    async def test_think_uses_disposition(self, memory: MemoryEngine):
+    async def test_think_uses_disposition(self, memory: MemoryEngine, request_context):
         """Test that THINK operation uses agent disposition."""
         bank_id = unique_agent_id("test_think")
 
@@ -218,12 +223,13 @@ class TestAgentDispositionIntegration:
             "literalism": 4,  # High literalism
             "empathy": 2,     # Low empathy
         }
-        await memory.update_bank_disposition(bank_id, disposition)
+        await memory.update_bank_disposition(bank_id, disposition, request_context=request_context)
 
         await memory.merge_bank_background(
             bank_id,
             "I am a creative artist who values innovation over tradition",
-            update_disposition=False
+            update_disposition=False,
+            request_context=request_context,
         )
 
         await memory.retain_batch_async(
@@ -232,13 +238,14 @@ class TestAgentDispositionIntegration:
                 {"content": "Traditional painting techniques have been used for centuries"},
                 {"content": "Modern digital art is changing the art world"}
             ],
-            document_id="art_facts"
+            request_context=request_context,
         )
 
         result = await memory.reflect_async(
             bank_id=bank_id,
             query="What do you think about traditional vs modern art?",
-            budget=Budget.LOW
+            budget=Budget.LOW,
+            request_context=request_context,
         )
 
         assert result.text is not None

hindsight-api/tests/test_batch_chunking.py

@@ -6,7 +6,7 @@ import os
 
 
 @pytest.mark.asyncio
-async def test_large_batch_auto_chunks(memory):
+async def test_large_batch_auto_chunks(memory, request_context):
     bank_id = "test_chunking_agent"
     # Create a large batch that should trigger chunking
     # Each item is ~2000 chars, so 30 items = 60k chars (exceeds 50k threshold)
@@ -24,7 +24,8 @@ async def test_large_batch_auto_chunks(memory):
     # Ingest the large batch (should auto-chunk)
     result = await memory.retain_batch_async(
         bank_id=bank_id,
-        contents=contents
+        contents=contents,
+        request_context=request_context,
     )
 
     # Verify we got results back
@@ -33,7 +34,7 @@ async def test_large_batch_auto_chunks(memory):
 
 
 @pytest.mark.asyncio
-async def test_small_batch_no_chunking(memory):
+async def test_small_batch_no_chunking(memory, request_context):
     bank_id = "test_no_chunking_agent"
 
     # Create a small batch that should NOT trigger chunking
@@ -50,7 +51,8 @@ async def test_small_batch_no_chunking(memory):
     # Ingest the small batch (should NOT auto-chunk)
     result = await memory.retain_batch_async(
         bank_id=bank_id,
-        contents=contents
+        contents=contents,
+        request_context=request_context,
     )
 
     # Verify we got results back

hindsight-api/tests/test_causal_relations.py

@@ -0,0 +1,223 @@
+"""
+Test suite for causal relations extraction and validation.
+
+Tests that:
+1. Causal relations only reference previous facts (target_index < current fact index)
+2. Invalid causal relation indices are rejected
+3. The new per-fact causal relations schema works correctly
+"""
+
+from datetime import datetime
+
+import pytest
+
+from hindsight_api import LLMConfig
+from hindsight_api.engine.retain.fact_extraction import extract_facts_from_text
+
+
+class TestCausalRelationsValidation:
+    """Tests for causal relations index validation."""
+
+    @pytest.mark.asyncio
+    async def test_causal_relations_only_reference_previous_facts(self):
+        """
+        Test that causal relations can only reference facts that appear before them.
+
+        This test verifies the new schema that prevents hallucination of invalid
+        fact indices by constraining target_index to be less than the current fact's index.
+        """
+        # Text with clear causal chain
+        text = """
+        I lost my job in January due to company layoffs.
+        Because I lost my job, I couldn't pay my rent.
+        Since I couldn't afford rent, I had to move to a cheaper apartment.
+        After moving, I started looking for a new job.
+        """
+
+        context = "Personal life update"
+        llm_config = LLMConfig.for_memory()
+        event_date = datetime(2024, 3, 15)
+
+        facts, _, usage = await extract_facts_from_text(
+            text=text,
+            event_date=event_date,
+            context=context,
+            llm_config=llm_config,
+            agent_name="TestUser",
+        )
+
+        assert len(facts) > 0, "Should extract at least one fact"
+
+        # Verify all causal relations reference valid previous facts
+        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 causal relation to fact {rel.target_fact_index}, "
+                        f"but target_index must be < current index ({i})"
+                    )
+                    assert rel.target_fact_index >= 0, (
+                        f"Fact {i} has negative causal relation index: {rel.target_fact_index}"
+                    )
+                    assert rel.relation_type in ["caused_by", "enabled_by", "prevented_by"], (
+                        f"Invalid relation_type: {rel.relation_type}"
+                    )
+
+    @pytest.mark.asyncio
+    async def test_first_fact_has_no_causal_relations(self):
+        """
+        Test that the first fact (index 0) cannot have causal relations.
+
+        Since causal relations can only reference previous facts,
+        and there are no facts before index 0, the first fact should
+        have no causal relations.
+        """
+        text = """
+        The user started a new machine learning project.
+        The project requires learning TensorFlow.
+        Learning TensorFlow is challenging but rewarding.
+        """
+
+        context = "Project update"
+        llm_config = LLMConfig.for_memory()
+        event_date = datetime(2024, 6, 1)
+
+        facts, _, _ = await extract_facts_from_text(
+            text=text,
+            event_date=event_date,
+            context=context,
+            llm_config=llm_config,
+            agent_name="TestUser",
+        )
+
+        assert len(facts) > 0, "Should extract at least one fact"
+
+        # First fact should have no causal relations (nothing to reference)
+        if facts[0].causal_relations:
+            # If there are causal relations on the first fact, they should be empty
+            # or the validation should have filtered them out
+            for rel in facts[0].causal_relations:
+                # This should never happen due to validation
+                assert False, (
+                    f"First fact should not have causal relations, "
+                    f"but found: target_index={rel.target_fact_index}"
+                )
+
+    @pytest.mark.asyncio
+    async def test_causal_chain_extraction(self):
+        """
+        Test that a clear causal chain is extracted with valid relations.
+        """
+        text = """
+        Emily got promoted to senior engineer last month.
+        Because of her promotion, she received a significant salary increase.
+        With the extra money, she decided to buy a new car.
+        """
+
+        context = "Personal achievement story"
+        llm_config = LLMConfig.for_memory()
+        event_date = datetime(2024, 7, 15)
+
+        facts, _, _ = await extract_facts_from_text(
+            text=text,
+            event_date=event_date,
+            context=context,
+            llm_config=llm_config,
+            agent_name="TestUser",
+        )
+
+        assert len(facts) > 0, "Should extract facts about the causal chain"
+
+        # Collect all causal relations
+        all_relations = []
+        for i, fact in enumerate(facts):
+            if fact.causal_relations:
+                for rel in fact.causal_relations:
+                    all_relations.append({
+                        "from_fact": i,
+                        "to_fact": rel.target_fact_index,
+                        "type": rel.relation_type,
+                    })
+
+        # If causal relations were extracted, verify they form a valid chain
+        if all_relations:
+            for rel in all_relations:
+                assert rel["to_fact"] < rel["from_fact"], (
+                    f"Causal relation from fact {rel['from_fact']} to fact {rel['to_fact']} "
+                    f"is invalid (target must be < source)"
+                )
+
+    @pytest.mark.asyncio
+    async def test_token_efficiency_with_causal_relations(self):
+        """
+        Test that causal relations don't cause excessive output tokens.
+
+        This test verifies that the new schema (per-fact causal relations
+        with index constraints) doesn't waste tokens on invalid relations.
+        """
+        text = """
+        The company announced budget cuts in Q1.
+        Due to the budget cuts, the marketing team was reduced.
+        The reduced team meant fewer campaigns could be run.
+        With fewer campaigns, lead generation dropped.
+        Lower leads resulted in decreased sales.
+        """
+
+        context = "Business impact analysis"
+        llm_config = LLMConfig.for_memory()
+        event_date = datetime(2024, 4, 1)
+
+        facts, _, usage = await extract_facts_from_text(
+            text=text,
+            event_date=event_date,
+            context=context,
+            llm_config=llm_config,
+            agent_name="TestUser",
+        )
+
+        assert len(facts) > 0, "Should extract facts"
+
+        # Calculate output/input ratio
+        if usage.input_tokens > 0:
+            ratio = usage.output_tokens / usage.input_tokens
+            # The ratio should be reasonable (< 5x) with the new schema
+            # Previously it could be 7-10x due to hallucinated indices
+            assert ratio < 6, (
+                f"Output/input token ratio {ratio:.2f}x is too high. "
+                f"Input: {usage.input_tokens}, Output: {usage.output_tokens}"
+            )
+
+    @pytest.mark.asyncio
+    async def test_relation_types_are_backward_looking(self):
+        """
+        Test that all relation types describe how the current fact
+        relates to a previous fact (caused_by, enabled_by, prevented_by).
+        """
+        text = """
+        Alice learned Python programming.
+        Because she knew Python, she got a job as a data scientist.
+        Her data science skills enabled her to lead the analytics team.
+        """
+
+        context = "Career progression"
+        llm_config = LLMConfig.for_memory()
+        event_date = datetime(2024, 5, 1)
+
+        facts, _, _ = await extract_facts_from_text(
+            text=text,
+            event_date=event_date,
+            context=context,
+            llm_config=llm_config,
+            agent_name="TestUser",
+        )
+
+        # Verify relation types are all backward-looking
+        valid_types = {"caused_by", "enabled_by", "prevented_by"}
+
+        for i, fact in enumerate(facts):
+            if fact.causal_relations:
+                for rel in fact.causal_relations:
+                    assert rel.relation_type in valid_types, (
+                        f"Invalid relation_type '{rel.relation_type}'. "
+                        f"Must be one of: {valid_types}"
+                    )

hindsight-api/tests/test_causal_relationships.py

@@ -0,0 +1,202 @@
+"""
+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
+        # New constraint: target_index must be < from_fact_index (can only reference PREVIOUS facts)
+        num_facts = len(facts)
+        invalid_relations = []
+        for rel in all_causal_relations:
+            # Must be non-negative and less than the current fact's index
+            if rel["to_fact_index"] < 0 or rel["to_fact_index"] >= rel["from_fact_index"]:
+                invalid_relations.append(rel)
+
+        assert len(invalid_relations) == 0, (
+            f"Found {len(invalid_relations)} causal relations with invalid indices! "
+            f"Each target_fact_index must be < from_fact_index (can only reference previous facts). "
+            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 (passive only - facts reference PREVIOUS facts)
+        valid_types = {"caused_by", "enabled_by", "prevented_by"}
+        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 (must reference PREVIOUS facts only)
+        for i, fact in enumerate(facts):
+            if fact.causal_relations:
+                for rel in fact.causal_relations:
+                    assert 0 <= rel.target_fact_index < i, (
+                        f"Fact {i} has causal relation to invalid index {rel.target_fact_index}. "
+                        f"Must reference previous facts only (valid range: 0 to {i - 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! 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"
+        )
+
+        # Validate all indices (must reference PREVIOUS facts only)
+        for i, fact in enumerate(facts):
+            if fact.causal_relations:
+                for rel in fact.causal_relations:
+                    assert 0 <= rel.target_fact_index < i, (
+                        f"Invalid target_fact_index {rel.target_fact_index} in fact {i}. "
+                        f"Must reference previous facts only (valid range: 0 to {i - 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_combined_scoring.py

@@ -10,6 +10,7 @@ import pytest
 from datetime import datetime, timezone
 from hindsight_api.engine.search.types import RetrievalResult, MergedCandidate, ScoredResult
 from hindsight_api.engine.memory_engine import Budget
+from hindsight_api import RequestContext
 
 
 class TestRRFNormalization:
@@ -125,7 +126,7 @@ class TestCombinedScoringFormula:
 
 
 @pytest.mark.asyncio
-async def test_trace_has_normalized_rrf(memory):
+async def test_trace_has_normalized_rrf(memory, request_context):
     """Integration test: verify trace contains normalized RRF values, not raw."""
     bank_id = f"test_scoring_{datetime.now(timezone.utc).timestamp()}"
 
@@ -135,21 +136,25 @@ async def test_trace_has_normalized_rrf(memory):
             bank_id=bank_id,
             content="Python is a programming language created by Guido van Rossum",
             context="tech facts",
+            request_context=request_context,
         )
         await memory.retain_async(
             bank_id=bank_id,
             content="JavaScript was created by Brendan Eich at Netscape",
             context="tech facts",
+            request_context=request_context,
         )
         await memory.retain_async(
             bank_id=bank_id,
             content="The Eiffel Tower is located in Paris, France",
             context="geography facts",
+            request_context=request_context,
         )
         await memory.retain_async(
             bank_id=bank_id,
             content="Mount Everest is the tallest mountain on Earth",
             context="geography facts",
+            request_context=request_context,
         )
 
         # Search with tracing
@@ -160,6 +165,7 @@ async def test_trace_has_normalized_rrf(memory):
             budget=Budget.LOW,
             max_tokens=1024,
             enable_trace=True,
+            request_context=request_context,
         )
 
         assert result.trace is not None, "Trace should be present"
@@ -210,11 +216,11 @@ async def test_trace_has_normalized_rrf(memory):
             print(f"  - First result score components: {sc}")
 
     finally:
-        await memory.delete_bank(bank_id)
+        await memory.delete_bank(bank_id, request_context=request_context)
 
 
 @pytest.mark.asyncio
-async def test_rrf_normalized_not_raw_in_trace(memory):
+async def test_rrf_normalized_not_raw_in_trace(memory, request_context):
     """Verify that raw RRF scores (0.04-0.06 range) don't appear as normalized values."""
     bank_id = f"test_rrf_raw_{datetime.now(timezone.utc).timestamp()}"
 
@@ -225,6 +231,7 @@ async def test_rrf_normalized_not_raw_in_trace(memory):
                 bank_id=bank_id,
                 content=f"Test fact number {i} about various topics",
                 context="test context",
+                request_context=request_context,
             )
 
         result = await memory.recall_async(
@@ -234,6 +241,7 @@ async def test_rrf_normalized_not_raw_in_trace(memory):
             budget=Budget.LOW,
             max_tokens=512,
             enable_trace=True,
+            request_context=request_context,
         )
 
         trace = result.trace
@@ -268,11 +276,11 @@ async def test_rrf_normalized_not_raw_in_trace(memory):
         print("\n✓ RRF raw vs normalized test passed!")
 
     finally:
-        await memory.delete_bank(bank_id)
+        await memory.delete_bank(bank_id, request_context=request_context)
 
 
 @pytest.mark.asyncio
-async def test_combined_score_matches_components(memory):
+async def test_combined_score_matches_components(memory, request_context):
     """Verify the final score actually equals the weighted sum of components."""
     bank_id = f"test_combined_{datetime.now(timezone.utc).timestamp()}"
 
@@ -281,11 +289,13 @@ async def test_combined_score_matches_components(memory):
             bank_id=bank_id,
             content="The quick brown fox jumps over the lazy dog",
             context="test",
+            request_context=request_context,
         )
         await memory.retain_async(
             bank_id=bank_id,
             content="A quick test of the emergency broadcast system",
             context="test",
+            request_context=request_context,
         )
 
         result = await memory.recall_async(
@@ -295,6 +305,7 @@ async def test_combined_score_matches_components(memory):
             budget=Budget.LOW,
             max_tokens=512,
             enable_trace=True,
+            request_context=request_context,
         )
 
         trace = result.trace
@@ -320,4 +331,4 @@ async def test_combined_score_matches_components(memory):
         print("\n✓ Combined score verification test passed!")
 
     finally:
-        await memory.delete_bank(bank_id)
+        await memory.delete_bank(bank_id, request_context=request_context)

hindsight-api/tests/test_custom_embedding_dimension.py

@@ -0,0 +1,602 @@
+"""
+Tests for custom embedding dimensions and automatic dimension detection.
+
+Uses isolated PostgreSQL schemas to avoid affecting other tests.
+Includes tests for:
+- Automatic embedding dimension detection and database schema adjustment
+- OpenAI embeddings provider with 1536 dimensions
+"""
+
+import asyncio
+import os
+import pytest
+from datetime import datetime
+from sqlalchemy import create_engine, text
+
+from hindsight_api import MemoryEngine, RequestContext
+from hindsight_api.engine.embeddings import LocalSTEmbeddings, OpenAIEmbeddings, CohereEmbeddings
+from hindsight_api.engine.cross_encoder import LocalSTCrossEncoder, CohereCrossEncoder
+from hindsight_api.engine.query_analyzer import DateparserQueryAnalyzer
+from hindsight_api.extensions import TenantExtension, TenantContext
+from hindsight_api.migrations import run_migrations, ensure_embedding_dimension
+
+
+# =============================================================================
+# Shared Utilities
+# =============================================================================
+
+
+class SchemaTenantExtension(TenantExtension):
+    """Tenant extension that routes all requests to a specific schema (for testing)."""
+
+    def __init__(self, schema_name: str):
+        self.schema_name = schema_name
+
+    async def authenticate(self, request_context: RequestContext) -> TenantContext:
+        return TenantContext(schema_name=self.schema_name)
+
+
+def get_test_schema(prefix: str, worker_id: str) -> str:
+    """Get unique schema name per xdist worker."""
+    if worker_id == "master" or not worker_id:
+        return prefix
+    return f"{prefix}_{worker_id}"
+
+
+def create_isolated_schema(db_url: str, schema_name: str, dimension: int | None = None):
+    """Create an isolated schema with migrations and optional dimension adjustment."""
+    engine = create_engine(db_url)
+
+    # Create schema (drop first if exists from previous failed run)
+    with engine.connect() as conn:
+        conn.execute(text(f"DROP SCHEMA IF EXISTS {schema_name} CASCADE"))
+        conn.execute(text(f"CREATE SCHEMA {schema_name}"))
+        conn.commit()
+
+    # Run migrations in the isolated schema
+    run_migrations(db_url, schema=schema_name)
+
+    # Adjust embedding dimension if specified
+    if dimension is not None:
+        ensure_embedding_dimension(db_url, dimension, schema=schema_name)
+
+
+def drop_schema(db_url: str, schema_name: str):
+    """Drop an isolated schema."""
+    engine = create_engine(db_url)
+    with engine.connect() as conn:
+        conn.execute(text(f"DROP SCHEMA IF EXISTS {schema_name} CASCADE"))
+        conn.commit()
+
+
+def get_column_dimension(db_url: str, schema: str = "public") -> int | None:
+    """Get the current embedding column dimension from the database."""
+    engine = create_engine(db_url)
+    with engine.connect() as conn:
+        result = conn.execute(
+            text("""
+                SELECT atttypmod
+                FROM pg_attribute a
+                JOIN pg_class c ON a.attrelid = c.oid
+                JOIN pg_namespace n ON c.relnamespace = n.oid
+                WHERE n.nspname = :schema
+                  AND c.relname = 'memory_units'
+                  AND a.attname = 'embedding'
+            """),
+            {"schema": schema},
+        ).scalar()
+        return result
+
+
+def get_row_count(db_url: str, schema: str = "public") -> int:
+    """Get the number of rows with embeddings in memory_units."""
+    engine = create_engine(db_url)
+    with engine.connect() as conn:
+        return conn.execute(
+            text(f"SELECT COUNT(*) FROM {schema}.memory_units WHERE embedding IS NOT NULL")
+        ).scalar()
+
+
+def insert_test_embedding(db_url: str, schema: str, dimension: int):
+    """Insert a test row with a dummy embedding."""
+    engine = create_engine(db_url)
+    embedding = [0.1] * dimension
+    embedding_str = "[" + ",".join(str(x) for x in embedding) + "]"
+
+    with engine.connect() as conn:
+        conn.execute(
+            text(f"""
+                INSERT INTO {schema}.memory_units (bank_id, text, embedding, event_date, fact_type)
+                VALUES ('test-bank', 'test text', '{embedding_str}'::vector, NOW(), 'world')
+            """)
+        )
+        conn.commit()
+
+
+def clear_embeddings(db_url: str, schema: str):
+    """Clear all rows from memory_units."""
+    engine = create_engine(db_url)
+    with engine.connect() as conn:
+        conn.execute(text(f"DELETE FROM {schema}.memory_units"))
+        conn.commit()
+
+
+# =============================================================================
+# Embedding Dimension Tests (Local Embeddings)
+# =============================================================================
+
+
+@pytest.fixture(scope="class")
+def dimension_test_schema(pg0_db_url, worker_id):
+    """Create an isolated schema for dimension tests."""
+    schema_name = get_test_schema("test_embed_dim", worker_id)
+    create_isolated_schema(pg0_db_url, schema_name)
+    yield pg0_db_url, schema_name
+    drop_schema(pg0_db_url, schema_name)
+
+
+class TestEmbeddingDimension:
+    """Tests for embedding dimension detection and adjustment."""
+
+    def test_dimension_matches_no_change(self, dimension_test_schema):
+        """When dimension matches, no changes should be made."""
+        db_url, schema = dimension_test_schema
+
+        # Get initial dimension (should be 384 from migration)
+        initial_dim = get_column_dimension(db_url, schema)
+        assert initial_dim == 384, f"Expected 384, got {initial_dim}"
+
+        # Call ensure_embedding_dimension with matching dimension
+        ensure_embedding_dimension(db_url, 384, schema=schema)
+
+        # Dimension should still be 384
+        assert get_column_dimension(db_url, schema) == 384
+
+    def test_dimension_change_empty_table(self, dimension_test_schema):
+        """When table is empty, dimension can be changed."""
+        db_url, schema = dimension_test_schema
+
+        # Ensure table is empty
+        clear_embeddings(db_url, schema)
+        assert get_row_count(db_url, schema) == 0
+
+        # Change dimension to 768
+        ensure_embedding_dimension(db_url, 768, schema=schema)
+
+        # Verify dimension changed
+        new_dim = get_column_dimension(db_url, schema)
+        assert new_dim == 768, f"Expected 768, got {new_dim}"
+
+        # Change back to 384 for other tests
+        ensure_embedding_dimension(db_url, 384, schema=schema)
+        assert get_column_dimension(db_url, schema) == 384
+
+    def test_dimension_change_blocked_with_data(self, dimension_test_schema):
+        """When table has data, dimension change should be blocked."""
+        db_url, schema = dimension_test_schema
+
+        # Ensure table is empty first
+        clear_embeddings(db_url, schema)
+
+        # Insert a test row with 384-dim embedding
+        insert_test_embedding(db_url, schema, 384)
+        assert get_row_count(db_url, schema) == 1
+
+        # Try to change dimension - should raise error
+        with pytest.raises(RuntimeError) as exc_info:
+            ensure_embedding_dimension(db_url, 768, schema=schema)
+
+        assert "Cannot change embedding dimension" in str(exc_info.value)
+        assert "1 rows with embeddings" in str(exc_info.value)
+
+        # Dimension should be unchanged
+        assert get_column_dimension(db_url, schema) == 384
+
+        # Cleanup
+        clear_embeddings(db_url, schema)
+
+    def test_local_embeddings_dimension_detection(self, embeddings):
+        """Test that LocalSTEmbeddings correctly detects dimension."""
+        # Initialize embeddings if not already done
+        loop = asyncio.new_event_loop()
+        try:
+            loop.run_until_complete(embeddings.initialize())
+        finally:
+            loop.close()
+
+        # bge-small-en-v1.5 produces 384-dim embeddings
+        assert embeddings.dimension == 384
+
+        # Verify by generating an actual embedding
+        result = embeddings.encode(["test"])
+        assert len(result) == 1
+        assert len(result[0]) == 384
+
+
+# =============================================================================
+# OpenAI Embeddings Tests
+# =============================================================================
+
+
+def has_openai_api_key() -> bool:
+    """Check if OpenAI API key is available."""
+    return bool(os.environ.get("HINDSIGHT_API_EMBEDDINGS_OPENAI_API_KEY"))
+
+
+def get_openai_api_key() -> str:
+    """Get OpenAI API key from environment."""
+    return os.environ.get("HINDSIGHT_API_EMBEDDINGS_OPENAI_API_KEY", "")
+
+
+@pytest.fixture(scope="module")
+def openai_embeddings():
+    """Create OpenAI embeddings instance."""
+    if not has_openai_api_key():
+        pytest.skip("OpenAI API key not available (set HINDSIGHT_API_EMBEDDINGS_OPENAI_API_KEY)")
+
+    embeddings = OpenAIEmbeddings(
+        api_key=get_openai_api_key(),
+        model="text-embedding-3-small",
+    )
+    loop = asyncio.new_event_loop()
+    try:
+        loop.run_until_complete(embeddings.initialize())
+    finally:
+        loop.close()
+    return embeddings
+
+
+@pytest.fixture(scope="module")
+def openai_test_schema(pg0_db_url, worker_id, openai_embeddings):
+    """Create an isolated schema for OpenAI embedding tests."""
+    schema_name = get_test_schema("test_openai_embed", worker_id)
+    create_isolated_schema(pg0_db_url, schema_name, dimension=openai_embeddings.dimension)
+    yield pg0_db_url, schema_name
+    drop_schema(pg0_db_url, schema_name)
+
+
+@pytest.fixture
+def cross_encoder():
+    """Provide a cross encoder for tests."""
+    return LocalSTCrossEncoder()
+
+
+@pytest.fixture
+def query_analyzer():
+    """Provide a query analyzer for tests."""
+    return DateparserQueryAnalyzer()
+
+
+@pytest.fixture
+def test_bank_id():
+    """Provide a unique bank ID for this test run."""
+    return f"openai_test_{datetime.now().timestamp()}"
+
+
+@pytest.fixture
+def request_context():
+    """Provide a default RequestContext for tests."""
+    return RequestContext()
+
+
+class TestOpenAIEmbeddings:
+    """Tests for OpenAI embeddings provider."""
+
+    def test_openai_embeddings_initialization(self, openai_embeddings):
+        """Test that OpenAI embeddings initializes correctly."""
+        assert openai_embeddings.dimension == 1536
+        assert openai_embeddings.provider_name == "openai"
+
+    def test_openai_embeddings_encode(self, openai_embeddings):
+        """Test that OpenAI embeddings can encode text."""
+        texts = ["Hello, world!", "This is a test."]
+        embeddings = openai_embeddings.encode(texts)
+
+        assert len(embeddings) == 2
+        assert len(embeddings[0]) == 1536
+        assert len(embeddings[1]) == 1536
+        assert all(isinstance(x, float) for x in embeddings[0])
+
+    @pytest.mark.asyncio
+    async def test_openai_embeddings_retain_recall(
+        self,
+        openai_test_schema,
+        openai_embeddings,
+        cross_encoder,
+        query_analyzer,
+        test_bank_id,
+        request_context,
+    ):
+        """Test retain and recall operations with OpenAI embeddings."""
+        db_url, schema_name = openai_test_schema
+
+        memory = MemoryEngine(
+            db_url=db_url,
+            memory_llm_provider=os.getenv("HINDSIGHT_API_LLM_PROVIDER", "groq"),
+            memory_llm_api_key=os.getenv("HINDSIGHT_API_LLM_API_KEY"),
+            memory_llm_model=os.getenv("HINDSIGHT_API_LLM_MODEL", "openai/gpt-oss-120b"),
+            memory_llm_base_url=os.getenv("HINDSIGHT_API_LLM_BASE_URL") or None,
+            embeddings=openai_embeddings,
+            cross_encoder=cross_encoder,
+            query_analyzer=query_analyzer,
+            pool_min_size=1,
+            pool_max_size=3,
+            run_migrations=False,
+            tenant_extension=SchemaTenantExtension(schema_name),
+        )
+
+        try:
+            await memory.initialize()
+
+            # Store some memories
+            await memory.retain_async(
+                bank_id=test_bank_id,
+                content="Alice works as a software engineer at Google.",
+                context="career discussion",
+                request_context=request_context,
+            )
+
+            await memory.retain_async(
+                bank_id=test_bank_id,
+                content="Bob is a data scientist specializing in machine learning.",
+                context="team introductions",
+                request_context=request_context,
+            )
+
+            # Recall memories
+            result = await memory.recall_async(
+                bank_id=test_bank_id,
+                query="Who works in technology?",
+                request_context=request_context,
+            )
+
+            assert result is not None
+            assert len(result.results) > 0
+
+            memory_texts = [m.text for m in result.results]
+            assert any(
+                "Alice" in text or "Bob" in text or "software" in text or "data scientist" in text
+                for text in memory_texts
+            ), f"Expected to find relevant memories, got: {memory_texts}"
+
+        finally:
+            try:
+                if memory._pool and not memory._pool._closing:
+                    await memory.close()
+            except Exception:
+                pass
+
+    @pytest.mark.asyncio
+    async def test_openai_embeddings_batch_retain(
+        self,
+        openai_test_schema,
+        openai_embeddings,
+        cross_encoder,
+        query_analyzer,
+        test_bank_id,
+        request_context,
+    ):
+        """Test batch retain with OpenAI embeddings."""
+        db_url, schema_name = openai_test_schema
+
+        memory = MemoryEngine(
+            db_url=db_url,
+            memory_llm_provider=os.getenv("HINDSIGHT_API_LLM_PROVIDER", "groq"),
+            memory_llm_api_key=os.getenv("HINDSIGHT_API_LLM_API_KEY"),
+            memory_llm_model=os.getenv("HINDSIGHT_API_LLM_MODEL", "openai/gpt-oss-120b"),
+            memory_llm_base_url=os.getenv("HINDSIGHT_API_LLM_BASE_URL") or None,
+            embeddings=openai_embeddings,
+            cross_encoder=cross_encoder,
+            query_analyzer=query_analyzer,
+            pool_min_size=1,
+            pool_max_size=3,
+            run_migrations=False,
+            tenant_extension=SchemaTenantExtension(schema_name),
+        )
+
+        try:
+            await memory.initialize()
+
+            contents = [
+                {"content": "Python is my favorite programming language.", "context": "preferences"},
+                {"content": "I prefer dark mode for all my applications.", "context": "preferences"},
+                {"content": "Coffee is essential for morning productivity.", "context": "habits"},
+            ]
+
+            result = await memory.retain_batch_async(
+                bank_id=test_bank_id,
+                contents=contents,
+                request_context=request_context,
+            )
+
+            assert len(result) == 3
+
+            recall_result = await memory.recall_async(
+                bank_id=test_bank_id,
+                query="What are my preferences?",
+                request_context=request_context,
+            )
+
+            assert recall_result is not None
+            assert len(recall_result.results) > 0
+
+        finally:
+            try:
+                if memory._pool and not memory._pool._closing:
+                    await memory.close()
+            except Exception:
+                pass
+
+
+# =============================================================================
+# Cohere Embeddings Tests
+# =============================================================================
+
+
+def has_cohere_api_key() -> bool:
+    """Check if Cohere API key is available."""
+    return bool(os.environ.get("COHERE_API_KEY"))
+
+
+def get_cohere_api_key() -> str:
+    """Get Cohere API key from environment."""
+    return os.environ.get("COHERE_API_KEY", "")
+
+
+@pytest.fixture(scope="module")
+def cohere_embeddings():
+    """Create Cohere embeddings instance."""
+    if not has_cohere_api_key():
+        pytest.skip("Cohere API key not available (set COHERE_API_KEY)")
+
+    embeddings = CohereEmbeddings(
+        api_key=get_cohere_api_key(),
+        model="embed-english-v3.0",
+    )
+    loop = asyncio.new_event_loop()
+    try:
+        loop.run_until_complete(embeddings.initialize())
+    finally:
+        loop.close()
+    return embeddings
+
+
+@pytest.fixture(scope="module")
+def cohere_cross_encoder():
+    """Create Cohere cross-encoder instance."""
+    if not has_cohere_api_key():
+        pytest.skip("Cohere API key not available (set COHERE_API_KEY)")
+
+    cross_encoder = CohereCrossEncoder(
+        api_key=get_cohere_api_key(),
+        model="rerank-english-v3.0",
+    )
+    loop = asyncio.new_event_loop()
+    try:
+        loop.run_until_complete(cross_encoder.initialize())
+    finally:
+        loop.close()
+    return cross_encoder
+
+
+@pytest.fixture(scope="module")
+def cohere_test_schema(pg0_db_url, worker_id, cohere_embeddings):
+    """Create an isolated schema for Cohere embedding tests."""
+    schema_name = get_test_schema("test_cohere_embed", worker_id)
+    create_isolated_schema(pg0_db_url, schema_name, dimension=cohere_embeddings.dimension)
+    yield pg0_db_url, schema_name
+    drop_schema(pg0_db_url, schema_name)
+
+
+class TestCohereEmbeddings:
+    """Tests for Cohere embeddings provider."""
+
+    def test_cohere_embeddings_initialization(self, cohere_embeddings):
+        """Test that Cohere embeddings initializes correctly."""
+        assert cohere_embeddings.dimension == 1024
+        assert cohere_embeddings.provider_name == "cohere"
+
+    def test_cohere_embeddings_encode(self, cohere_embeddings):
+        """Test that Cohere embeddings can encode text."""
+        texts = ["Hello, world!", "This is a test."]
+        embeddings = cohere_embeddings.encode(texts)
+
+        assert len(embeddings) == 2
+        assert len(embeddings[0]) == 1024
+        assert len(embeddings[1]) == 1024
+        assert all(isinstance(x, float) for x in embeddings[0])
+
+
+class TestCohereCrossEncoder:
+    """Tests for Cohere cross-encoder/reranker."""
+
+    def test_cohere_cross_encoder_initialization(self, cohere_cross_encoder):
+        """Test that Cohere cross-encoder initializes correctly."""
+        assert cohere_cross_encoder.provider_name == "cohere"
+
+    def test_cohere_cross_encoder_predict(self, cohere_cross_encoder):
+        """Test that Cohere cross-encoder can score pairs."""
+        pairs = [
+            ("What is the capital of France?", "Paris is the capital of France."),
+            ("What is the capital of France?", "The Eiffel Tower is in Paris."),
+            ("What is the capital of France?", "Python is a programming language."),
+        ]
+        scores = cohere_cross_encoder.predict(pairs)
+
+        assert len(scores) == 3
+        assert all(isinstance(s, float) for s in scores)
+        # The first result should be most relevant
+        assert scores[0] > scores[2], "Direct answer should score higher than unrelated text"
+
+
+class TestCohereIntegration:
+    """Integration tests for Cohere embeddings with memory engine."""
+
+    @pytest.mark.asyncio
+    async def test_cohere_embeddings_retain_recall(
+        self,
+        cohere_test_schema,
+        cohere_embeddings,
+        cohere_cross_encoder,
+        query_analyzer,
+        request_context,
+    ):
+        """Test retain and recall operations with Cohere embeddings."""
+        db_url, schema_name = cohere_test_schema
+        test_bank_id = f"cohere_test_{datetime.now().timestamp()}"
+
+        memory = MemoryEngine(
+            db_url=db_url,
+            memory_llm_provider=os.getenv("HINDSIGHT_API_LLM_PROVIDER", "groq"),
+            memory_llm_api_key=os.getenv("HINDSIGHT_API_LLM_API_KEY"),
+            memory_llm_model=os.getenv("HINDSIGHT_API_LLM_MODEL", "openai/gpt-oss-120b"),
+            memory_llm_base_url=os.getenv("HINDSIGHT_API_LLM_BASE_URL") or None,
+            embeddings=cohere_embeddings,
+            cross_encoder=cohere_cross_encoder,
+            query_analyzer=query_analyzer,
+            pool_min_size=1,
+            pool_max_size=3,
+            run_migrations=False,
+            tenant_extension=SchemaTenantExtension(schema_name),
+        )
+
+        try:
+            await memory.initialize()
+
+            # Store some memories
+            await memory.retain_async(
+                bank_id=test_bank_id,
+                content="Alice works as a software engineer at Google.",
+                context="career discussion",
+                request_context=request_context,
+            )
+
+            await memory.retain_async(
+                bank_id=test_bank_id,
+                content="Bob is a data scientist specializing in machine learning.",
+                context="team introductions",
+                request_context=request_context,
+            )
+
+            # Recall memories
+            result = await memory.recall_async(
+                bank_id=test_bank_id,
+                query="Who works in technology?",
+                request_context=request_context,
+            )
+
+            assert result is not None
+            assert len(result.results) > 0
+
+            memory_texts = [m.text for m in result.results]
+            assert any(
+                "Alice" in text or "Bob" in text or "software" in text or "data scientist" in text
+                for text in memory_texts
+            ), f"Expected to find relevant memories, got: {memory_texts}"
+
+        finally:
+            try:
+                if memory._pool and not memory._pool._closing:
+                    await memory.close()
+            except Exception:
+                pass

hindsight-api/tests/test_document_tracking.py

@@ -4,10 +4,11 @@ Tests for document tracking and upsert functionality.
 import logging
 import pytest
 from datetime import datetime, timezone
+from hindsight_api import RequestContext
 
 
 @pytest.mark.asyncio
-async def test_document_creation_and_retrieval(memory):
+async def test_document_creation_and_retrieval(memory, request_context):
     """Test that documents are created and can be retrieved."""
     bank_id = f"test_doc_{datetime.now(timezone.utc).timestamp()}"
 
@@ -19,11 +20,12 @@ async def test_document_creation_and_retrieval(memory):
             bank_id=bank_id,
             content="Alice works at Google. Bob works at Microsoft.",
             context="Team meeting",
-            document_id=document_id
+            document_id=document_id,
+            request_context=request_context,
         )
 
         # Retrieve document
-        doc = await memory.get_document(document_id, bank_id)
+        doc = await memory.get_document(document_id, bank_id, request_context=request_context)
 
         assert doc is not None
         assert doc["id"] == document_id
@@ -32,11 +34,11 @@ async def test_document_creation_and_retrieval(memory):
         assert doc["memory_unit_count"] > 0
 
     finally:
-        await memory.delete_bank(bank_id)
+        await memory.delete_bank(bank_id, request_context=request_context)
 
 
 @pytest.mark.asyncio
-async def test_document_upsert(memory):
+async def test_document_upsert(memory, request_context):
     """Test that providing the same document_id automatically upserts (deletes old units and creates new ones)."""
     bank_id = f"test_upsert_{datetime.now(timezone.utc).timestamp()}"
 
@@ -48,11 +50,12 @@ async def test_document_upsert(memory):
             bank_id=bank_id,
             content="Alice works at Google.",
             context="Initial",
-            document_id=document_id
+            document_id=document_id,
+            request_context=request_context,
         )
 
         # Get document stats
-        doc_v1 = await memory.get_document(document_id, bank_id)
+        doc_v1 = await memory.get_document(document_id, bank_id, request_context=request_context)
         count_v1 = doc_v1["memory_unit_count"]
 
         # Update with different content (automatic upsert when same document_id is provided)
@@ -60,11 +63,12 @@ async def test_document_upsert(memory):
             bank_id=bank_id,
             content="Alice works at Microsoft. Bob works at Apple.",
             context="Updated",
-            document_id=document_id
+            document_id=document_id,
+            request_context=request_context,
         )
 
         # Get updated document stats
-        doc_v2 = await memory.get_document(document_id, bank_id)
+        doc_v2 = await memory.get_document(document_id, bank_id, request_context=request_context)
         count_v2 = doc_v2["memory_unit_count"]
 
         # Verify old units were replaced
@@ -75,11 +79,11 @@ async def test_document_upsert(memory):
         assert set(units_v1).isdisjoint(set(units_v2))
 
     finally:
-        await memory.delete_bank(bank_id)
+        await memory.delete_bank(bank_id, request_context=request_context)
 
 
 @pytest.mark.asyncio
-async def test_document_deletion(memory):
+async def test_document_deletion(memory, request_context):
     """Test that deleting a document cascades to memory units."""
     bank_id = f"test_delete_{datetime.now(timezone.utc).timestamp()}"
 
@@ -91,29 +95,30 @@ async def test_document_deletion(memory):
             bank_id=bank_id,
             content="Alice works at Google.",
             context="Test",
-            document_id=document_id
+            document_id=document_id,
+            request_context=request_context,
         )
 
         # Verify it exists
-        doc = await memory.get_document(document_id, bank_id)
+        doc = await memory.get_document(document_id, bank_id, request_context=request_context)
         assert doc is not None
         assert doc["memory_unit_count"] > 0
 
         # Delete document
-        result = await memory.delete_document(document_id, bank_id)
+        result = await memory.delete_document(document_id, bank_id, request_context=request_context)
         assert result["document_deleted"] == 1
         assert result["memory_units_deleted"] > 0
 
         # Verify it's gone
-        doc_after = await memory.get_document(document_id, bank_id)
+        doc_after = await memory.get_document(document_id, bank_id, request_context=request_context)
         assert doc_after is None
 
     finally:
-        await memory.delete_bank(bank_id)
+        await memory.delete_bank(bank_id, request_context=request_context)
 
 
 @pytest.mark.asyncio
-async def test_memory_without_document(memory):
+async def test_memory_without_document(memory, request_context):
     """Test that memories can still be created without document tracking."""
     bank_id = f"test_no_doc_{datetime.now(timezone.utc).timestamp()}"
 
@@ -122,10 +127,11 @@ async def test_memory_without_document(memory):
         units = await memory.retain_async(
             bank_id=bank_id,
             content="Alice works at Google.",
-            context="Test"
+            context="Test",
+            request_context=request_context,
         )
 
         assert len(units) > 0
 
     finally:
-        await memory.delete_bank(bank_id)
+        await memory.delete_bank(bank_id, request_context=request_context)

hindsight-api/tests/test_extensions.py

@@ -0,0 +1,796 @@
+"""Tests for the Hindsight extensions system."""
+
+from collections import defaultdict
+
+import pytest
+from fastapi import APIRouter
+from fastapi.testclient import TestClient
+
+from hindsight_api.extensions import (
+    ApiKeyTenantExtension,
+    AuthenticationError,
+    Extension,
+    HttpExtension,
+    OperationValidationError,
+    OperationValidatorExtension,
+    RecallContext,
+    RecallResult,
+    ReflectContext,
+    ReflectResultContext,
+    RequestContext,
+    RetainContext,
+    RetainResult,
+    TenantContext,
+    TenantExtension,
+    ValidationResult,
+    load_extension,
+)
+
+
+class TestExtensionLoader:
+    """Tests for extension loading and lifecycle."""
+
+    def test_load_extension_with_config(self, monkeypatch):
+        """Extension receives config from prefixed env vars and supports lifecycle."""
+        monkeypatch.setenv(
+            "HINDSIGHT_API_TEST_EXTENSION",
+            "tests.test_extensions:LifecycleTestExtension",
+        )
+        monkeypatch.setenv("HINDSIGHT_API_TEST_API_URL", "https://example.com")
+        monkeypatch.setenv("HINDSIGHT_API_TEST_MAX_RETRIES", "5")
+
+        ext = load_extension("TEST", Extension)
+
+        assert ext is not None
+        assert ext.config["api_url"] == "https://example.com"
+        assert ext.config["max_retries"] == "5"
+
+    @pytest.mark.asyncio
+    async def test_extension_lifecycle(self, monkeypatch):
+        """Extension on_startup and on_shutdown are called."""
+        monkeypatch.setenv(
+            "HINDSIGHT_API_TEST_EXTENSION",
+            "tests.test_extensions:LifecycleTestExtension",
+        )
+
+        ext = load_extension("TEST", Extension)
+
+        assert not ext.started
+        assert not ext.stopped
+
+        await ext.on_startup()
+        assert ext.started
+
+        await ext.on_shutdown()
+        assert ext.stopped
+
+
+class LifecycleTestExtension(Extension):
+    """Test extension for config and lifecycle tests."""
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.started = False
+        self.stopped = False
+
+    async def on_startup(self):
+        self.started = True
+
+    async def on_shutdown(self):
+        self.stopped = True
+
+
+class RateLimitingValidator(OperationValidatorExtension):
+    """
+    Mock validator that blocks after N attempts per bank_id.
+
+    Used for testing the extension integration with MemoryEngine.
+    """
+
+    def __init__(self, config: dict):
+        super().__init__(config)
+        self.max_attempts = int(config.get("max_attempts", "2"))
+        self.retain_counts: dict[str, int] = defaultdict(int)
+        self.recall_counts: dict[str, int] = defaultdict(int)
+        self.reflect_counts: dict[str, int] = defaultdict(int)
+
+    async def validate_retain(self, ctx: RetainContext) -> ValidationResult:
+        self.retain_counts[ctx.bank_id] += 1
+        if self.retain_counts[ctx.bank_id] > self.max_attempts:
+            return ValidationResult.reject(
+                f"Retain limit exceeded for bank {ctx.bank_id}"
+            )
+        return ValidationResult.accept()
+
+    async def validate_recall(self, ctx: RecallContext) -> ValidationResult:
+        self.recall_counts[ctx.bank_id] += 1
+        if self.recall_counts[ctx.bank_id] > self.max_attempts:
+            return ValidationResult.reject(
+                f"Recall limit exceeded for bank {ctx.bank_id}"
+            )
+        return ValidationResult.accept()
+
+    async def validate_reflect(self, ctx: ReflectContext) -> ValidationResult:
+        self.reflect_counts[ctx.bank_id] += 1
+        if self.reflect_counts[ctx.bank_id] > self.max_attempts:
+            return ValidationResult.reject(
+                f"Reflect limit exceeded for bank {ctx.bank_id}"
+            )
+        return ValidationResult.accept()
+
+
+class TrackingValidator(OperationValidatorExtension):
+    """
+    Mock validator that tracks all pre and post hook calls with full parameters.
+
+    Used for testing that hooks receive all user-provided parameters.
+    """
+
+    def __init__(self, config: dict):
+        super().__init__(config)
+        # Pre-hook tracking
+        self.pre_retain_calls: list[RetainContext] = []
+        self.pre_recall_calls: list[RecallContext] = []
+        self.pre_reflect_calls: list[ReflectContext] = []
+        # Post-hook tracking
+        self.post_retain_calls: list[RetainResult] = []
+        self.post_recall_calls: list[RecallResult] = []
+        self.post_reflect_calls: list[ReflectResultContext] = []
+
+    async def validate_retain(self, ctx: RetainContext) -> ValidationResult:
+        self.pre_retain_calls.append(ctx)
+        return ValidationResult.accept()
+
+    async def validate_recall(self, ctx: RecallContext) -> ValidationResult:
+        self.pre_recall_calls.append(ctx)
+        return ValidationResult.accept()
+
+    async def validate_reflect(self, ctx: ReflectContext) -> ValidationResult:
+        self.pre_reflect_calls.append(ctx)
+        return ValidationResult.accept()
+
+    async def on_retain_complete(self, result: RetainResult) -> None:
+        self.post_retain_calls.append(result)
+
+    async def on_recall_complete(self, result: RecallResult) -> None:
+        self.post_recall_calls.append(result)
+
+    async def on_reflect_complete(self, result: ReflectResultContext) -> None:
+        self.post_reflect_calls.append(result)
+
+
+class TestMemoryEngineValidation:
+    """Tests for validation integration with MemoryEngine.
+
+    The OperationValidatorExtension is integrated at the MemoryEngine level,
+    so all interfaces (HTTP API, MCP, SDK) get the same validation behavior.
+
+    For retain, the batch is validated as a whole (all or nothing) using
+    retain_batch_async which is the public method used by the HTTP API.
+    """
+
+    @pytest.mark.asyncio
+    async def test_retain_batch_validation(self, memory_with_validator):
+        """Retain batch is validated as a whole - accepts or rejects entire batch."""
+        memory = memory_with_validator
+        bank_id = "test-retain-batch"
+        ctx = RequestContext()
+
+        # First batch should succeed
+        await memory.retain_batch_async(
+            bank_id=bank_id,
+            contents=[
+                {"content": "First item"},
+                {"content": "Second item"},
+            ],
+            request_context=ctx,
+        )
+
+        # Second batch should succeed (2nd attempt)
+        await memory.retain_batch_async(
+            bank_id=bank_id,
+            contents=[{"content": "Third item"}],
+            request_context=ctx,
+        )
+
+        # Third batch should be blocked entirely (exceeds limit)
+        with pytest.raises(OperationValidationError) as exc_info:
+            await memory.retain_batch_async(
+                bank_id=bank_id,
+                contents=[
+                    {"content": "Should not be stored"},
+                    {"content": "Neither should this"},
+                ],
+                request_context=ctx,
+            )
+
+        assert "limit exceeded" in str(exc_info.value).lower()
+
+    @pytest.mark.asyncio
+    async def test_recall_validation(self, memory_with_validator):
+        """Recall is validated before execution."""
+        memory = memory_with_validator
+        bank_id = "test-recall-validation"
+        ctx = RequestContext()
+
+        # First recall should pass validation
+        await memory.recall_async(bank_id, "test query", fact_type=["world"], request_context=ctx)
+
+        # Second recall should pass validation
+        await memory.recall_async(bank_id, "another query", fact_type=["world"], request_context=ctx)
+
+        # Third recall should be blocked by validator
+        with pytest.raises(OperationValidationError) as exc_info:
+            await memory.recall_async(bank_id, "blocked query", fact_type=["world"], request_context=ctx)
+
+        assert "limit exceeded" in str(exc_info.value).lower()
+
+    @pytest.mark.asyncio
+    async def test_reflect_validation(self, memory_with_validator):
+        """Reflect is validated before execution."""
+        memory = memory_with_validator
+        bank_id = "test-reflect-validation"
+        ctx = RequestContext()
+
+        # First reflect should pass validation (may fail internally but validation passes)
+        try:
+            await memory.reflect_async(bank_id, "test question", request_context=ctx)
+        except OperationValidationError:
+            raise  # Re-raise validation errors
+        except Exception:
+            pass  # Other errors are fine (e.g., no data)
+
+        # Second reflect should pass validation
+        try:
+            await memory.reflect_async(bank_id, "another question", request_context=ctx)
+        except OperationValidationError:
+            raise
+        except Exception:
+            pass
+
+        # Third reflect should be blocked by validator
+        with pytest.raises(OperationValidationError) as exc_info:
+            await memory.reflect_async(bank_id, "blocked question", request_context=ctx)
+
+        assert "limit exceeded" in str(exc_info.value).lower()
+
+
+@pytest.fixture
+def memory_with_validator(memory):
+    """Memory engine with a rate-limiting validator (max 2 attempts per bank)."""
+    validator = RateLimitingValidator({"max_attempts": "2"})
+    memory._operation_validator = validator
+    return memory
+
+
+@pytest.fixture
+def memory_with_tracking_validator(memory):
+    """Memory engine with a tracking validator that records all hook calls."""
+    validator = TrackingValidator({})
+    memory._operation_validator = validator
+    return memory, validator
+
+
+class TestOperationHooksParameters:
+    """Tests for pre and post operation hooks receiving all user-provided parameters."""
+
+    @pytest.mark.asyncio
+    async def test_retain_pre_hook_receives_all_parameters(self, memory_with_tracking_validator):
+        """Pre-retain hook receives all user-provided parameters."""
+        memory, validator = memory_with_tracking_validator
+        bank_id = "test-retain-params"
+        ctx = RequestContext(api_key="test-key")
+        contents = [{"content": "Test content", "context": "test context"}]
+        document_id = "doc-123"
+
+        await memory.retain_batch_async(
+            bank_id=bank_id,
+            contents=contents,
+            document_id=document_id,
+            fact_type_override="world",
+            confidence_score=0.9,
+            request_context=ctx,
+        )
+
+        assert len(validator.pre_retain_calls) == 1
+        pre_ctx = validator.pre_retain_calls[0]
+
+        # Verify all parameters are present
+        assert pre_ctx.bank_id == bank_id
+        # Note: contents is copied before document_id is applied to individual items
+        assert len(pre_ctx.contents) == len(contents)
+        assert pre_ctx.contents[0]["content"] == contents[0]["content"]
+        assert pre_ctx.document_id == document_id
+        assert pre_ctx.fact_type_override == "world"
+        assert pre_ctx.confidence_score == 0.9
+        assert pre_ctx.request_context == ctx
+
+    @pytest.mark.asyncio
+    async def test_retain_post_hook_receives_all_parameters_and_result(self, memory_with_tracking_validator):
+        """Post-retain hook receives all parameters plus the result."""
+        memory, validator = memory_with_tracking_validator
+        bank_id = "test-retain-post"
+        ctx = RequestContext(api_key="test-key")
+        contents = [{"content": "Test content for post hook"}]
+        document_id = "doc-456"
+
+        result = await memory.retain_batch_async(
+            bank_id=bank_id,
+            contents=contents,
+            document_id=document_id,
+            fact_type_override="experience",
+            confidence_score=0.8,
+            request_context=ctx,
+        )
+
+        assert len(validator.post_retain_calls) == 1
+        post_result = validator.post_retain_calls[0]
+
+        # Verify all parameters are present
+        assert post_result.bank_id == bank_id
+        assert post_result.document_id == document_id
+        assert post_result.fact_type_override == "experience"
+        assert post_result.confidence_score == 0.8
+        assert post_result.request_context == ctx
+
+        # Verify result data
+        assert post_result.success is True
+        assert post_result.error is None
+        assert post_result.unit_ids == result  # Should match the return value
+
+    @pytest.mark.asyncio
+    async def test_recall_pre_hook_receives_all_parameters(self, memory_with_tracking_validator):
+        """Pre-recall hook receives all user-provided parameters."""
+        from datetime import datetime, timezone
+        from hindsight_api.engine.memory_engine import Budget
+
+        memory, validator = memory_with_tracking_validator
+        bank_id = "test-recall-params"
+        ctx = RequestContext(api_key="test-key")
+        query = "test query"
+        question_date = datetime(2024, 1, 15, tzinfo=timezone.utc)
+
+        await memory.recall_async(
+            bank_id=bank_id,
+            query=query,
+            budget=Budget.HIGH,
+            max_tokens=2048,
+            enable_trace=True,
+            fact_type=["world", "experience"],
+            question_date=question_date,
+            include_entities=True,
+            max_entity_tokens=300,
+            include_chunks=True,
+            max_chunk_tokens=4096,
+            request_context=ctx,
+        )
+
+        assert len(validator.pre_recall_calls) == 1
+        pre_ctx = validator.pre_recall_calls[0]
+
+        # Verify all parameters are present
+        assert pre_ctx.bank_id == bank_id
+        assert pre_ctx.query == query
+        assert pre_ctx.budget == Budget.HIGH
+        assert pre_ctx.max_tokens == 2048
+        assert pre_ctx.enable_trace is True
+        assert pre_ctx.fact_types == ["world", "experience"]
+        assert pre_ctx.question_date == question_date
+        assert pre_ctx.include_entities is True
+        assert pre_ctx.max_entity_tokens == 300
+        assert pre_ctx.include_chunks is True
+        assert pre_ctx.max_chunk_tokens == 4096
+        assert pre_ctx.request_context == ctx
+
+    @pytest.mark.asyncio
+    async def test_recall_post_hook_receives_all_parameters_and_result(self, memory_with_tracking_validator):
+        """Post-recall hook receives all parameters plus the result."""
+        from hindsight_api.engine.memory_engine import Budget
+
+        memory, validator = memory_with_tracking_validator
+        bank_id = "test-recall-post"
+        ctx = RequestContext(api_key="test-key")
+
+        result = await memory.recall_async(
+            bank_id=bank_id,
+            query="test query for post",
+            budget=Budget.LOW,
+            max_tokens=1024,
+            fact_type=["world"],
+            request_context=ctx,
+        )
+
+        assert len(validator.post_recall_calls) == 1
+        post_result = validator.post_recall_calls[0]
+
+        # Verify all parameters are present
+        assert post_result.bank_id == bank_id
+        assert post_result.query == "test query for post"
+        assert post_result.budget == Budget.LOW
+        assert post_result.max_tokens == 1024
+        assert post_result.fact_types == ["world"]
+        assert post_result.request_context == ctx
+
+        # Verify result data
+        assert post_result.success is True
+        assert post_result.error is None
+        assert post_result.result == result  # Should match the return value
+
+    @pytest.mark.asyncio
+    async def test_reflect_pre_hook_receives_all_parameters(self, memory_with_tracking_validator):
+        """Pre-reflect hook receives all user-provided parameters."""
+        from hindsight_api.engine.memory_engine import Budget
+
+        memory, validator = memory_with_tracking_validator
+        bank_id = "test-reflect-params"
+        ctx = RequestContext(api_key="test-key")
+
+        try:
+            await memory.reflect_async(
+                bank_id=bank_id,
+                query="test question",
+                budget=Budget.MID,
+                context="additional context",
+                request_context=ctx,
+            )
+        except Exception:
+            pass  # May fail if no data, but pre-hook should still be called
+
+        assert len(validator.pre_reflect_calls) == 1
+        pre_ctx = validator.pre_reflect_calls[0]
+
+        # Verify all parameters are present
+        assert pre_ctx.bank_id == bank_id
+        assert pre_ctx.query == "test question"
+        assert pre_ctx.budget == Budget.MID
+        assert pre_ctx.context == "additional context"
+        assert pre_ctx.request_context == ctx
+
+    @pytest.mark.asyncio
+    async def test_reflect_post_hook_receives_all_parameters_and_result(self, memory_with_tracking_validator):
+        """Post-reflect hook receives all parameters plus the result on success."""
+        from hindsight_api.engine.memory_engine import Budget
+
+        memory, validator = memory_with_tracking_validator
+        bank_id = "test-reflect-post"
+        ctx = RequestContext(api_key="test-key")
+
+        # Store some content first so reflect has something to work with
+        await memory.retain_batch_async(
+            bank_id=bank_id,
+            contents=[{"content": "Alice is a software engineer at Google."}],
+            request_context=ctx,
+        )
+
+        result = await memory.reflect_async(
+            bank_id=bank_id,
+            query="What does Alice do?",
+            budget=Budget.LOW,
+            context="work context",
+            request_context=ctx,
+        )
+
+        assert len(validator.post_reflect_calls) == 1
+        post_result = validator.post_reflect_calls[0]
+
+        # Verify all parameters are present
+        assert post_result.bank_id == bank_id
+        assert post_result.query == "What does Alice do?"
+        assert post_result.budget == Budget.LOW
+        assert post_result.context == "work context"
+        assert post_result.request_context == ctx
+
+        # Verify result data
+        assert post_result.success is True
+        assert post_result.error is None
+        assert post_result.result == result  # Should match the return value
+        assert post_result.result.text is not None
+
+    @pytest.mark.asyncio
+    async def test_post_hooks_called_in_order_after_pre_hooks(self, memory_with_tracking_validator):
+        """Post hooks are called after pre hooks and after operation completes."""
+        memory, validator = memory_with_tracking_validator
+        bank_id = "test-hook-order"
+        ctx = RequestContext()
+
+        # Retain operation
+        await memory.retain_batch_async(
+            bank_id=bank_id,
+            contents=[{"content": "Test content"}],
+            request_context=ctx,
+        )
+
+        # Pre-hook should be called before post-hook
+        assert len(validator.pre_retain_calls) == 1
+        assert len(validator.post_retain_calls) == 1
+
+        # Recall operation
+        await memory.recall_async(
+            bank_id=bank_id,
+            query="test",
+            fact_type=["world"],
+            request_context=ctx,
+        )
+
+        assert len(validator.pre_recall_calls) == 1
+        assert len(validator.post_recall_calls) == 1
+
+
+class TestTenantExtension:
+    """Tests for TenantExtension and ApiKeyTenantExtension."""
+
+    @pytest.mark.asyncio
+    async def test_api_key_tenant_extension_valid_key(self):
+        """ApiKeyTenantExtension accepts valid API key."""
+        ext = ApiKeyTenantExtension({"api_key": "secret-key-123"})
+
+        result = await ext.authenticate(RequestContext(api_key="secret-key-123"))
+
+        assert result.schema_name == "public"
+
+    @pytest.mark.asyncio
+    async def test_api_key_tenant_extension_invalid_key(self):
+        """ApiKeyTenantExtension rejects invalid API key."""
+        ext = ApiKeyTenantExtension({"api_key": "secret-key-123"})
+
+        with pytest.raises(AuthenticationError) as exc_info:
+            await ext.authenticate(RequestContext(api_key="wrong-key"))
+
+        assert "Invalid API key" in str(exc_info.value)
+
+    @pytest.mark.asyncio
+    async def test_api_key_tenant_extension_missing_key(self):
+        """ApiKeyTenantExtension rejects missing API key."""
+        ext = ApiKeyTenantExtension({"api_key": "secret-key-123"})
+
+        with pytest.raises(AuthenticationError):
+            await ext.authenticate(RequestContext(api_key=None))
+
+    def test_api_key_tenant_extension_requires_config(self):
+        """ApiKeyTenantExtension requires api_key in config."""
+        with pytest.raises(ValueError) as exc_info:
+            ApiKeyTenantExtension({})
+
+        assert "HINDSIGHT_API_TENANT_API_KEY is required" in str(exc_info.value)
+
+
+class TestMemoryEngineTenantAuth:
+    """Tests for tenant authentication in MemoryEngine."""
+
+    @pytest.mark.asyncio
+    async def test_retain_requires_tenant_request_when_extension_configured(
+        self, memory_with_tenant
+    ):
+        """Retain fails without RequestContext when tenant extension is configured."""
+        memory = memory_with_tenant
+
+        with pytest.raises(AuthenticationError) as exc_info:
+            await memory.retain_batch_async(
+                bank_id="test-bank",
+                contents=[{"content": "test"}],
+                request_context=None,  # Missing!
+            )
+
+        assert "RequestContext is required" in str(exc_info.value)
+
+    @pytest.mark.asyncio
+    async def test_retain_succeeds_with_valid_tenant_request(self, memory_with_tenant):
+        """Retain succeeds with valid RequestContext."""
+        memory = memory_with_tenant
+
+        # Should not raise
+        await memory.retain_batch_async(
+            bank_id="test-bank-tenant",
+            contents=[{"content": "test content"}],
+            request_context=RequestContext(api_key="test-api-key"),
+        )
+
+    @pytest.mark.asyncio
+    async def test_retain_fails_with_invalid_api_key(self, memory_with_tenant):
+        """Retain fails with invalid API key."""
+        memory = memory_with_tenant
+
+        with pytest.raises(AuthenticationError) as exc_info:
+            await memory.retain_batch_async(
+                bank_id="test-bank",
+                contents=[{"content": "test"}],
+                request_context=RequestContext(api_key="wrong-key"),
+            )
+
+        assert "Invalid API key" in str(exc_info.value)
+
+    @pytest.mark.asyncio
+    async def test_recall_requires_tenant_request_when_extension_configured(
+        self, memory_with_tenant
+    ):
+        """Recall fails without RequestContext when tenant extension is configured."""
+        memory = memory_with_tenant
+
+        with pytest.raises(AuthenticationError):
+            await memory.recall_async(
+                bank_id="test-bank",
+                query="test query",
+                fact_type=["world"],
+                request_context=None,
+            )
+
+    @pytest.mark.asyncio
+    async def test_no_tenant_request_needed_without_extension(self, memory):
+        """Operations work with empty RequestContext when no tenant extension configured."""
+        # Should not raise - no tenant extension configured, just pass empty RequestContext
+        await memory.retain_batch_async(
+            bank_id="test-bank-no-tenant",
+            contents=[{"content": "test content"}],
+            request_context=RequestContext(),
+        )
+
+
+@pytest.fixture
+def memory_with_tenant(memory):
+    """Memory engine with a tenant extension (API key auth)."""
+    tenant_ext = ApiKeyTenantExtension({"api_key": "test-api-key"})
+    memory._tenant_extension = tenant_ext
+    return memory
+
+
+class SampleHttpExtension(HttpExtension):
+    """Sample HTTP extension for testing that provides custom endpoints."""
+
+    def __init__(self, config: dict):
+        super().__init__(config)
+        self.started = False
+        self.stopped = False
+        self.request_count = 0
+
+    async def on_startup(self):
+        self.started = True
+
+    async def on_shutdown(self):
+        self.stopped = True
+
+    def get_router(self, memory) -> APIRouter:
+        router = APIRouter()
+
+        @router.get("/hello")
+        async def hello():
+            self.request_count += 1
+            return {"message": "Hello from extension!"}
+
+        @router.get("/config")
+        async def get_config():
+            return {"config": self.config}
+
+        @router.get("/health-check")
+        async def extension_health():
+            health = await memory.health_check()
+            return {"extension": "healthy", "memory": health}
+
+        @router.post("/echo")
+        async def echo(data: dict):
+            return {"echoed": data}
+
+        return router
+
+
+class TestHttpExtensionIntegration:
+    """Tests for HTTP extension integration."""
+
+    def test_load_http_extension(self, monkeypatch):
+        """HttpExtension can be loaded from environment variable."""
+        monkeypatch.setenv(
+            "HINDSIGHT_API_HTTP_EXTENSION",
+            "tests.test_extensions:SampleHttpExtension",
+        )
+        monkeypatch.setenv("HINDSIGHT_API_HTTP_CUSTOM_PARAM", "custom_value")
+
+        ext = load_extension("HTTP", HttpExtension)
+
+        assert ext is not None
+        assert isinstance(ext, SampleHttpExtension)
+        assert ext.config["custom_param"] == "custom_value"
+
+    def test_http_extension_router_mounted_at_ext(self, memory):
+        """HTTP extension router is mounted at /ext/."""
+        from hindsight_api.api.http import create_app
+
+        ext = SampleHttpExtension({"test_key": "test_value"})
+        app = create_app(memory, initialize_memory=False, http_extension=ext)
+
+        client = TestClient(app)
+
+        # Extension endpoint should be accessible at /ext/
+        response = client.get("/ext/hello")
+        assert response.status_code == 200
+        assert response.json() == {"message": "Hello from extension!"}
+
+        # Should track request count
+        assert ext.request_count == 1
+
+        # Old path should NOT work
+        response = client.get("/extension/hello")
+        assert response.status_code == 404
+
+    def test_http_extension_config_endpoint(self, memory):
+        """Extension can expose its config via custom endpoint."""
+        from hindsight_api.api.http import create_app
+
+        ext = SampleHttpExtension({"api_key": "secret", "limit": "100"})
+        app = create_app(memory, initialize_memory=False, http_extension=ext)
+
+        client = TestClient(app)
+
+        response = client.get("/ext/config")
+        assert response.status_code == 200
+        assert response.json()["config"]["api_key"] == "secret"
+        assert response.json()["config"]["limit"] == "100"
+
+    def test_http_extension_can_access_memory(self, memory):
+        """Extension endpoints can access memory engine."""
+        from hindsight_api.api.http import create_app
+
+        ext = SampleHttpExtension({})
+        app = create_app(memory, initialize_memory=False, http_extension=ext)
+
+        client = TestClient(app)
+
+        response = client.get("/ext/health-check")
+        assert response.status_code == 200
+        data = response.json()
+        assert data["extension"] == "healthy"
+        assert "memory" in data
+
+    def test_http_extension_post_endpoint(self, memory):
+        """Extension can handle POST requests with JSON body."""
+        from hindsight_api.api.http import create_app
+
+        ext = SampleHttpExtension({})
+        app = create_app(memory, initialize_memory=False, http_extension=ext)
+
+        client = TestClient(app)
+
+        response = client.post("/ext/echo", json={"key": "value", "number": 42})
+        assert response.status_code == 200
+        assert response.json() == {"echoed": {"key": "value", "number": 42}}
+
+    def test_http_extension_not_mounted_when_none(self, memory):
+        """No extension routes when http_extension is None."""
+        from hindsight_api.api.http import create_app
+
+        app = create_app(memory, initialize_memory=False, http_extension=None)
+
+        client = TestClient(app)
+
+        # Extension endpoint should not exist
+        response = client.get("/ext/hello")
+        assert response.status_code == 404
+
+    @pytest.mark.asyncio
+    async def test_http_extension_lifecycle(self):
+        """HTTP extension on_startup and on_shutdown are called."""
+        ext = SampleHttpExtension({})
+
+        assert not ext.started
+        assert not ext.stopped
+
+        await ext.on_startup()
+        assert ext.started
+
+        await ext.on_shutdown()
+        assert ext.stopped
+
+    def test_core_routes_still_work_with_extension(self, memory):
+        """Core API routes still work when extension is mounted."""
+        from hindsight_api.api.http import create_app
+
+        ext = SampleHttpExtension({})
+        app = create_app(memory, initialize_memory=False, http_extension=ext)
+
+        client = TestClient(app)
+
+        # Health endpoint should work
+        response = client.get("/health")
+        assert response.status_code in (200, 503)  # May be unhealthy if DB not connected
+
+        # Banks list endpoint should work
+        response = client.get("/v1/default/banks")
+        assert response.status_code in (200, 500)  # May fail if DB not ready

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,
@@ -402,6 +402,8 @@ I'm planning to visit Tokyo next month.
 
         Ideally: If conversation is on August 14, 2023 and text says "last night",
         the date field should be August 13. We accept 13 or 14 as LLM may vary.
+
+        Retries up to 3 times to account for LLM inconsistencies.
         """
         text = """
 Melanie: Hey Caroline! Last night was amazing! We celebrated my daughter's birthday
@@ -410,41 +412,69 @@ with a concert surrounded by music, joy and the warm summer breeze.
 
         context = "Conversation between Melanie and Caroline"
         llm_config = LLMConfig.for_memory()
-
         event_date = datetime(2023, 8, 14, 14, 24)
 
-        facts, _ = await extract_facts_from_text(
-            text=text,
-            event_date=event_date,
-            context=context,
-            llm_config=llm_config,
-            agent_name="Melanie"
-        )
+        last_error = None
+        max_retries = 3
 
-        assert len(facts) > 0, "Should extract at least one fact"
+        for attempt in range(max_retries):
+            try:
+                facts, _, _ = await extract_facts_from_text(
+                    text=text,
+                    event_date=event_date,
+                    context=context,
+                    llm_config=llm_config,
+                    agent_name="Melanie"
+                )
 
-        birthday_fact = None
-        for fact in facts:
-            if "birthday" in fact.fact.lower() or "concert" in fact.fact.lower():
-                birthday_fact = fact
-                break
+                assert len(facts) > 0, "Should extract at least one fact"
 
-        assert birthday_fact is not None, "Should extract fact about birthday celebration"
+                birthday_fact = None
+                for fact in facts:
+                    if "birthday" in fact.fact.lower() or "concert" in fact.fact.lower():
+                        birthday_fact = fact
+                        break
 
-        fact_date_str = birthday_fact.occurred_start
-        assert fact_date_str is not None, "occurred_start should not be None for temporal events"
+                assert birthday_fact is not None, "Should extract fact about birthday celebration"
 
-        if 'T' in fact_date_str:
-            fact_date = datetime.fromisoformat(fact_date_str.replace('Z', '+00:00'))
-        else:
-            fact_date = datetime.fromisoformat(fact_date_str)
+                fact_date_str = birthday_fact.occurred_start
+                assert fact_date_str is not None, "occurred_start should not be None for temporal events"
 
-        assert fact_date.year == 2023, "Year should be 2023"
-        assert fact_date.month == 8, "Month should be August"
-        # Accept day 13 (ideal: last night) or 14 (conversation date) as valid
-        assert fact_date.day in (13, 14), (
-            f"Day should be 13 or 14 (around Aug 14 event), but got {fact_date.day}."
-        )
+                if 'T' in fact_date_str:
+                    fact_date = datetime.fromisoformat(fact_date_str.replace('Z', '+00:00'))
+                else:
+                    fact_date = datetime.fromisoformat(fact_date_str)
+
+                assert fact_date.year == 2023, "Year should be 2023"
+                assert fact_date.month == 8, "Month should be August"
+                # Accept day 13 (ideal: last night) or 14 (conversation date) as valid
+                assert fact_date.day in (13, 14), (
+                    f"Day should be 13 or 14 (around Aug 14 event), but got {fact_date.day}."
+                )
+
+                # If we reach here, test passed
+                return
+
+            except AssertionError as e:
+                last_error = e
+                if attempt < max_retries - 1:
+                    print(f"Test attempt {attempt + 1} failed: {e}. Retrying...")
+                    continue
+                else:
+                    # Last attempt failed, re-raise the error
+                    raise e
+            except Exception as e:
+                last_error = e
+                if attempt < max_retries - 1:
+                    print(f"Test attempt {attempt + 1} failed with exception: {e}. Retrying...")
+                    continue
+                else:
+                    # Last attempt failed, re-raise the error
+                    raise e
+
+        # Should not reach here, but just in case
+        if last_error:
+            raise last_error
 
     @pytest.mark.asyncio
     async def test_date_field_calculation_yesterday(self):
@@ -458,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,
@@ -507,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,
@@ -537,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,
@@ -564,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,
@@ -615,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,
@@ -665,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,
@@ -728,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,
@@ -773,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,
@@ -808,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,
@@ -867,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,
@@ -897,7 +927,7 @@ class TestDispositionInference:
     """Tests for LLM-based disposition trait inference from background."""
 
     @pytest.mark.asyncio
-    async def test_background_merge_with_disposition_inference(self, memory):
+    async def test_background_merge_with_disposition_inference(self, memory, request_context):
         """Test that background merge infers disposition traits by default."""
         import uuid
         bank_id = f"test_infer_{uuid.uuid4().hex[:8]}"
@@ -905,7 +935,8 @@ class TestDispositionInference:
         result = await memory.merge_bank_background(
             bank_id,
             "I am a creative software engineer who loves innovation and trying new technologies",
-            update_disposition=True
+            update_disposition=True,
+            request_context=request_context,
         )
 
         assert "background" in result
@@ -923,30 +954,31 @@ class TestDispositionInference:
             assert 1 <= disposition[trait] <= 5
 
     @pytest.mark.asyncio
-    async def test_background_merge_without_disposition_inference(self, memory):
+    async def test_background_merge_without_disposition_inference(self, memory, request_context):
         """Test that background merge skips disposition inference when disabled."""
         import uuid
         bank_id = f"test_no_infer_{uuid.uuid4().hex[:8]}"
 
-        initial_profile = await memory.get_bank_profile(bank_id)
+        initial_profile = await memory.get_bank_profile(bank_id, request_context=request_context)
         initial_disposition = initial_profile["disposition"]
 
         result = await memory.merge_bank_background(
             bank_id,
             "I am a data scientist",
-            update_disposition=False
+            update_disposition=False,
+            request_context=request_context,
         )
 
         assert "background" in result
         assert "disposition" not in result
 
-        final_profile = await memory.get_bank_profile(bank_id)
+        final_profile = await memory.get_bank_profile(bank_id, request_context=request_context)
         final_disposition = final_profile["disposition"]
 
         assert initial_disposition == final_disposition
 
     @pytest.mark.asyncio
-    async def test_disposition_inference_for_lawyer(self, memory):
+    async def test_disposition_inference_for_lawyer(self, memory, request_context):
         """Test disposition inference for lawyer profile (high skepticism, high literalism)."""
         import uuid
         bank_id = f"test_lawyer_{uuid.uuid4().hex[:8]}"
@@ -954,7 +986,8 @@ class TestDispositionInference:
         result = await memory.merge_bank_background(
             bank_id,
             "I am a lawyer who focuses on contract details and never takes claims at face value",
-            update_disposition=True
+            update_disposition=True,
+            request_context=request_context,
         )
 
         disposition = result["disposition"]
@@ -964,7 +997,7 @@ class TestDispositionInference:
         assert disposition["literalism"] >= 3
 
     @pytest.mark.asyncio
-    async def test_disposition_inference_for_therapist(self, memory):
+    async def test_disposition_inference_for_therapist(self, memory, request_context):
         """Test disposition inference for therapist profile (high empathy)."""
         import uuid
         bank_id = f"test_therapist_{uuid.uuid4().hex[:8]}"
@@ -972,7 +1005,8 @@ class TestDispositionInference:
         result = await memory.merge_bank_background(
             bank_id,
             "I am a therapist who deeply understands and connects with people's emotional struggles",
-            update_disposition=True
+            update_disposition=True,
+            request_context=request_context,
         )
 
         disposition = result["disposition"]
@@ -981,7 +1015,7 @@ class TestDispositionInference:
         assert disposition["empathy"] >= 3
 
     @pytest.mark.asyncio
-    async def test_disposition_updates_in_database(self, memory):
+    async def test_disposition_updates_in_database(self, memory, request_context):
         """Test that inferred disposition is actually stored in database."""
         import uuid
         bank_id = f"test_db_update_{uuid.uuid4().hex[:8]}"
@@ -989,12 +1023,13 @@ class TestDispositionInference:
         result = await memory.merge_bank_background(
             bank_id,
             "I am an innovative designer",
-            update_disposition=True
+            update_disposition=True,
+            request_context=request_context,
         )
 
         inferred_disposition = result["disposition"]
 
-        profile = await memory.get_bank_profile(bank_id)
+        profile = await memory.get_bank_profile(bank_id, request_context=request_context)
         db_disposition = profile["disposition"]
 
         # Compare values (db_disposition is a Pydantic model)
@@ -1003,7 +1038,7 @@ class TestDispositionInference:
         assert db_disposition.empathy == inferred_disposition["empathy"]
 
     @pytest.mark.asyncio
-    async def test_multiple_background_merges_update_disposition(self, memory):
+    async def test_multiple_background_merges_update_disposition(self, memory, request_context):
         """Test that each background merge can update disposition."""
         import uuid
         bank_id = f"test_multi_merge_{uuid.uuid4().hex[:8]}"
@@ -1011,14 +1046,16 @@ class TestDispositionInference:
         result1 = await memory.merge_bank_background(
             bank_id,
             "I am a software engineer",
-            update_disposition=True
+            update_disposition=True,
+            request_context=request_context,
         )
         disposition1 = result1["disposition"]
 
         result2 = await memory.merge_bank_background(
             bank_id,
             "I love creative problem solving and innovation",
-            update_disposition=True
+            update_disposition=True,
+            request_context=request_context,
         )
         disposition2 = result2["disposition"]
 
@@ -1026,7 +1063,7 @@ class TestDispositionInference:
         assert "creative" in result2["background"].lower() or "innovation" in result2["background"].lower()
 
     @pytest.mark.asyncio
-    async def test_background_merge_conflict_resolution_with_disposition(self, memory):
+    async def test_background_merge_conflict_resolution_with_disposition(self, memory, request_context):
         """Test that conflicts are resolved and disposition reflects final background."""
         import uuid
         bank_id = f"test_conflict_{uuid.uuid4().hex[:8]}"
@@ -1034,13 +1071,15 @@ class TestDispositionInference:
         await memory.merge_bank_background(
             bank_id,
             "I was born in Colorado and prefer stability",
-            update_disposition=True
+            update_disposition=True,
+            request_context=request_context,
         )
 
         result = await memory.merge_bank_background(
             bank_id,
             "You were born in Texas and are very skeptical of people",
-            update_disposition=True
+            update_disposition=True,
+            request_context=request_context,
         )
 
         background = result["background"]

hindsight-api/tests/test_fact_ordering.py

@@ -7,24 +7,24 @@ distinguish between things said earlier vs later.
 """
 import pytest
 from datetime import datetime, timezone
-from hindsight_api import MemoryEngine
+from hindsight_api import MemoryEngine, RequestContext
 from hindsight_api.engine.memory_engine import Budget
 import os
 
 
 @pytest.mark.asyncio
-async def test_fact_ordering_within_conversation(memory):
+async def test_fact_ordering_within_conversation(memory, request_context):
     bank_id = "test_ordering_agent"
 
     # Get/create agent (auto-creates with defaults)
-    await memory.get_bank_profile(bank_id)
+    await memory.get_bank_profile(bank_id, request_context=request_context)
 
     # Update disposition to match Marcus
     await memory.update_bank_disposition(bank_id, {
         "skepticism": 3,
         "literalism": 3,
         "empathy": 3
-    })
+    }, request_context=request_context)
 
     # A conversation where Marcus changes his position
     conversation = """
@@ -43,7 +43,8 @@ Marcus: Yeah, I realized I was being too optimistic about their defense.
         content=conversation,
         context="podcast discussion about NFL game",
         event_date=base_event_date,
-        document_id="test_conv_1"
+        document_id="test_conv_1",
+        request_context=request_context,
     )
 
     # Search for all facts about Marcus's predictions
@@ -52,7 +53,8 @@ Marcus: Yeah, I realized I was being too optimistic about their defense.
         query="Marcus prediction Rams",
         fact_type=['opinion', 'experience', 'world'],
         budget=Budget.LOW,
-        max_tokens=8192
+        max_tokens=8192,
+        request_context=request_context,
     )
 
     print(f"\n=== Retrieved {len(results.results)} facts ===")
@@ -113,17 +115,17 @@ Marcus: Yeah, I realized I was being too optimistic about their defense.
         print(f"\n✅ Temporal ordering preserved: First prediction came before changed prediction")
 
     # Cleanup
-    await memory.delete_bank(bank_id)
+    await memory.delete_bank(bank_id, request_context=request_context)
 
     print(f"\n✅ Test passed: Fact ordering within conversation is preserved")
 
 
 @pytest.mark.asyncio
-async def test_multiple_documents_ordering(memory):
+async def test_multiple_documents_ordering(memory, request_context):
 
     bank_id = "test_multi_doc_agent"
 
-    await memory.get_bank_profile(bank_id)  # Auto-creates with defaults
+    await memory.get_bank_profile(bank_id, request_context=request_context)  # Auto-creates with defaults
 
     # Two separate conversations with same base time
     base_time = datetime(2024, 11, 14, 10, 0, 0, tzinfo=timezone.utc)
@@ -146,7 +148,8 @@ Alice: I reconsidered the team's experience level.
         contents=[
             {"content": conv1, "context": "project discussion 1", "event_date": base_time},
             {"content": conv2, "context": "project discussion 2", "event_date": base_time}
-        ]
+        ],
+        request_context=request_context,
     )
 
     # Search for Alice's preferences
@@ -155,7 +158,8 @@ Alice: I reconsidered the team's experience level.
         query="Alice preference React Vue",
         fact_type=['opinion', 'experience'],
         budget=Budget.LOW,
-        max_tokens=8192
+        max_tokens=8192,
+        request_context=request_context,
     )
 
     print(f"\n=== Retrieved {len(results.results)} agent facts ===")
@@ -175,6 +179,6 @@ Alice: I reconsidered the team's experience level.
         print(f"\n✅ Facts from {len(agent_facts)} statements have {len(unique_timestamps)} unique timestamps")
 
     # Cleanup
-    await memory.delete_bank(bank_id)
+    await memory.delete_bank(bank_id, request_context=request_context)
 
     print(f"\n✅ Test passed: Multiple documents maintain separate ordering")

hindsight-api/tests/test_http_api_integration.py

@@ -428,6 +428,66 @@ async def test_document_deletion(api_client):
     assert response.status_code == 404
 
 
+@pytest.mark.asyncio
+async def test_document_deletion_with_slashes_in_id(api_client):
+    """
+    Test document deletion when document_id contains forward slashes.
+
+    Regression test for https://github.com/vectorize-io/hindsight/issues/92
+
+    Document IDs with slashes (e.g., "folder/file.md") should work correctly
+    for all operations including creation, listing, retrieval, and deletion.
+    """
+    import urllib.parse
+
+    test_bank_id = f"doc_slash_test_{datetime.now().timestamp()}"
+    document_id_with_slash = "reports/quarterly/q1-2024.md"
+
+    try:
+        # 1. Create a document with slashes in its ID
+        response = await api_client.post(
+            f"/v1/default/banks/{test_bank_id}/memories",
+            json={
+                "items": [
+                    {
+                        "content": "The Q1 2024 report shows significant growth in user engagement.",
+                        "context": "quarterly report",
+                        "document_id": document_id_with_slash
+                    }
+                ]
+            }
+        )
+        assert response.status_code == 200, f"Failed to create document: {response.text}"
+
+        # 2. Verify document exists via list endpoint
+        response = await api_client.get(f"/v1/default/banks/{test_bank_id}/documents")
+        assert response.status_code == 200
+        documents = response.json()
+        doc_ids = [doc["id"] for doc in documents["items"]]
+        assert document_id_with_slash in doc_ids, f"Document should be in list: {doc_ids}"
+
+        # 3. Delete the document (slashes in document_id should work with :path converter)
+        encoded_doc_id = urllib.parse.quote(document_id_with_slash, safe="")
+        response = await api_client.delete(
+            f"/v1/default/banks/{test_bank_id}/documents/{encoded_doc_id}"
+        )
+        assert response.status_code == 200, (
+            f"Failed to delete document with slashes in ID. "
+            f"Status: {response.status_code}, Response: {response.text}"
+        )
+
+        # Verify document is deleted
+        response = await api_client.get(f"/v1/default/banks/{test_bank_id}/documents")
+        assert response.status_code == 200
+        documents = response.json()
+        doc_ids = [doc["id"] for doc in documents["items"]]
+        assert document_id_with_slash not in doc_ids, "Document should be deleted"
+
+    finally:
+        # Cleanup - delete the bank
+        await api_client.delete(f"/v1/default/banks/{test_bank_id}")
+
+
 @pytest.mark.asyncio
 async def test_async_retain(api_client):
     """Test asynchronous retain functionality.
@@ -608,3 +668,298 @@ async def test_async_retain_parallel(api_client):
         assert response.status_code == 200
         results = response.json()["results"]
         assert len(results) > 0, f"Should find memories for document {i}"
+
+
+@pytest.mark.asyncio
+async def test_reflect_structured_output(api_client):
+    """Test reflect endpoint with structured output via response_schema.
+
+    When response_schema is provided, the reflect endpoint should return
+    both the natural language text response and a structured_output field
+    containing the response parsed according to the provided JSON schema.
+    """
+    test_bank_id = f"reflect_structured_test_{datetime.now().timestamp()}"
+
+    # Store some memories to reflect on
+    response = await api_client.post(
+        f"/v1/default/banks/{test_bank_id}/memories",
+        json={
+            "items": [
+                {
+                    "content": "Alice is a senior machine learning engineer with 8 years of experience.",
+                    "context": "team member info"
+                },
+                {
+                    "content": "Bob is a junior data scientist who joined last month.",
+                    "context": "team member info"
+                },
+                {
+                    "content": "The team uses Python and TensorFlow for most projects.",
+                    "context": "tech stack"
+                }
+            ]
+        }
+    )
+    assert response.status_code == 200
+
+    # Define a JSON schema for structured output
+    response_schema = {
+        "type": "object",
+        "properties": {
+            "team_members": {
+                "type": "array",
+                "items": {
+                    "type": "object",
+                    "properties": {
+                        "name": {"type": "string"},
+                        "role": {"type": "string"},
+                        "experience_level": {"type": "string"}
+                    }
+                }
+            },
+            "technologies": {
+                "type": "array",
+                "items": {"type": "string"}
+            },
+            "summary": {"type": "string"}
+        },
+        "required": ["team_members", "summary"]
+    }
+
+    # Call reflect with response_schema
+    response = await api_client.post(
+        f"/v1/default/banks/{test_bank_id}/reflect",
+        json={
+            "query": "Give me an overview of the team and their tech stack",
+            "response_schema": response_schema
+        }
+    )
+    assert response.status_code == 200
+    result = response.json()
+
+    # Verify text field exists (empty when using structured output)
+    assert "text" in result
+    assert result["text"] == ""
+
+    # Verify structured output exists and has expected structure
+    assert "structured_output" in result
+    assert result["structured_output"] is not None
+
+    structured = result["structured_output"]
+    assert "team_members" in structured
+    assert "summary" in structured
+    assert isinstance(structured["team_members"], list)
+    assert isinstance(structured["summary"], str)
+
+    # Verify team members have the expected fields
+    if len(structured["team_members"]) > 0:
+        member = structured["team_members"][0]
+        assert "name" in member or "role" in member  # At least some fields should be present
+
+
+@pytest.mark.asyncio
+async def test_reflect_without_structured_output(api_client):
+    """Test that reflect works normally without response_schema.
+
+    When response_schema is not provided, the structured_output field
+    should be null/None in the response.
+    """
+    test_bank_id = f"reflect_no_structured_test_{datetime.now().timestamp()}"
+
+    # Store a memory
+    response = await api_client.post(
+        f"/v1/default/banks/{test_bank_id}/memories",
+        json={
+            "items": [
+                {
+                    "content": "The project deadline is next Friday.",
+                    "context": "project timeline"
+                }
+            ]
+        }
+    )
+    assert response.status_code == 200
+
+    # Call reflect without response_schema
+    response = await api_client.post(
+        f"/v1/default/banks/{test_bank_id}/reflect",
+        json={
+            "query": "When is the project deadline?"
+        }
+    )
+    assert response.status_code == 200
+    result = response.json()
+
+    # Verify response has text but structured_output is null
+    assert "text" in result
+    assert len(result["text"]) > 0
+    assert result.get("structured_output") is None
+
+
+@pytest.mark.asyncio
+async def test_reflect_with_max_tokens(api_client):
+    """Test reflect endpoint with custom max_tokens parameter.
+
+    The max_tokens parameter controls the maximum tokens for the LLM response.
+    """
+    test_bank_id = f"reflect_max_tokens_test_{datetime.now().timestamp()}"
+
+    # Store a memory
+    response = await api_client.post(
+        f"/v1/default/banks/{test_bank_id}/memories",
+        json={
+            "items": [
+                {
+                    "content": "Python is a popular programming language for data science and machine learning.",
+                    "context": "tech"
+                }
+            ]
+        }
+    )
+    assert response.status_code == 200
+
+    # Call reflect with custom max_tokens
+    response = await api_client.post(
+        f"/v1/default/banks/{test_bank_id}/reflect",
+        json={
+            "query": "What is Python used for?",
+            "max_tokens": 500
+        }
+    )
+    assert response.status_code == 200
+    result = response.json()
+
+    # 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_provider.py

@@ -26,6 +26,9 @@ MODEL_MATRIX = [
     ("gemini", "gemini-2.5-flash"),
     ("gemini", "gemini-2.5-flash-lite"),
     ("gemini", "gemini-3-pro-preview"),
+    # Ollama models (local)
+    ("ollama", "gemma3:12b"),
+    ("ollama", "gemma3:1b"),
 ]
 
 
@@ -48,12 +51,18 @@ async def test_llm_provider_memory_operations(provider: str, model: str):
     All models must pass this test.
     """
     api_key = get_api_key_for_provider(provider)
-    if not api_key:
+
+    # Skip Ollama tests in CI (no models available)
+    if provider == "ollama" and os.getenv("CI"):
+        pytest.skip(f"Skipping {provider}/{model}: Ollama not available in CI")
+
+    # Other providers need an API key
+    if provider != "ollama" and not api_key:
         pytest.skip(f"Skipping {provider}/{model}: no API key available")
 
     llm = LLMProvider(
         provider=provider,
-        api_key=api_key,
+        api_key=api_key or "",
         base_url="",
         model=model,
     )

hindsight-api/tests/test_llm_token_metrics.py

@@ -0,0 +1,245 @@
+"""
+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,
+)
+
+
+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_llm_metrics_recorded_for_groq():
+    """
+    Test that LLM 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_llm_call 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_llm_call was called - this is the main test
+        assert mock_collector.record_llm_call.called, "record_llm_call should have been called"
+
+        # Get the call arguments
+        call_kwargs = mock_collector.record_llm_call.call_args.kwargs
+
+        # Verify the call had correct structure
+        assert call_kwargs["provider"] == "groq", f"Expected provider='groq', got {call_kwargs}"
+        assert call_kwargs["model"] == "openai/gpt-oss-20b", f"Expected model='openai/gpt-oss-20b', got {call_kwargs}"
+        assert call_kwargs["scope"] == "test_metrics", f"Expected scope='test_metrics', got {call_kwargs}"
+        assert call_kwargs["duration"] > 0, f"Expected duration > 0, got {call_kwargs['duration']}"
+        assert call_kwargs["input_tokens"] > 0, f"Expected input_tokens > 0, got {call_kwargs['input_tokens']}"
+        assert call_kwargs["output_tokens"] >= 0, f"Expected output_tokens >= 0, got {call_kwargs['output_tokens']}"
+        assert call_kwargs["success"] is True, f"Expected success=True, got {call_kwargs['success']}"
+
+        print(f"\nLLM metrics recorded:")
+        print(f"  provider: {call_kwargs['provider']}")
+        print(f"  model: {call_kwargs['model']}")
+        print(f"  scope: {call_kwargs['scope']}")
+        print(f"  duration: {call_kwargs['duration']:.3f}s")
+        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_llm_metrics_recorded_for_structured_output():
+    """
+    Test that LLM 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_llm_call was called
+        assert mock_collector.record_llm_call.called, "record_llm_call should have been called"
+
+        call_kwargs = mock_collector.record_llm_call.call_args.kwargs
+        assert call_kwargs["input_tokens"] > 0
+        assert call_kwargs["output_tokens"] > 0
+
+        print(f"\nStructured output LLM 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-api/tests/test_mcp_api_integration.py

@@ -1,177 +0,0 @@
-"""
-Integration test for the MCP (Model Context Protocol) server.
-
-Tests MCP endpoints by starting a FastAPI server with MCP enabled and using the MCP client.
-
-Note: MCP server is integrated with the web server. These tests require HINDSIGHT_API_MCP_ENABLED=true.
-"""
-import asyncio
-import pytest
-import pytest_asyncio
-import httpx
-from mcp import ClientSession
-from mcp.client.sse import sse_client
-from hindsight_api.api import create_app
-
-
-@pytest_asyncio.fixture
-async def mcp_server(memory):
-    """Start the FastAPI app with MCP enabled and return the SSE URL."""
-    # Memory is already initialized by the conftest fixture (with migrations)
-    app = create_app(
-        memory,
-        initialize_memory=False,
-        mcp_api_enabled=True
-    )
-
-    # Use httpx to create a test server
-    transport = httpx.ASGITransport(app=app)
-    async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
-        # The MCP SSE endpoint is at /mcp/sse
-        # We need to yield the base URL for sse_client to connect
-        # However, sse_client expects a real URL, not a test client
-        # So we'll start a real server on a random port
-        pass
-
-    # For now, skip these tests as they require a real server
-    # The sse_client doesn't work with ASGI test transport
-    pytest.skip("MCP tests require a real running server. Run: HINDSIGHT_API_MCP_ENABLED=true uvicorn hindsight_api.api:app")
-
-
-@pytest.mark.asyncio
-async def test_mcp_server_tools_via_sse(mcp_server):
-    """Test MCP server tools via SSE transport using proper MCP client."""
-    sse_url = mcp_server
-
-    async with sse_client(sse_url) as (read, write):
-        async with ClientSession(read, write) as session:
-            await session.initialize()
-
-            # Test 1: List tools
-            tools_list = await session.list_tools()
-            print(f"Tools: {tools_list}")
-            tool_names = [t.name for t in tools_list.tools]
-            assert "hindsight_search" in tool_names
-            assert "hindsight_put" in tool_names
-
-            # Test 2: Call hindsight_put
-            put_result = await session.call_tool(
-                "hindsight_put",
-                arguments={
-                    "content": "User loves Python programming",
-                    "context": "programming_preferences",
-                    "explanation": "Storing user's programming language preference"
-                }
-            )
-            print(f"Put result: {put_result}")
-            assert put_result is not None
-
-            # Wait a bit for indexing
-            await asyncio.sleep(1)
-
-            # Test 3: Call hindsight_search
-            search_result = await session.call_tool(
-                "hindsight_search",
-                arguments={
-                    "query": "What programming languages does the user like?",
-                    "max_tokens": 4096,
-                    "explanation": "Searching for programming preferences"
-                }
-            )
-            print(f"Search result: {search_result}")
-            assert search_result is not None
-
-
-@pytest.mark.asyncio
-async def test_multiple_concurrent_requests(mcp_server):
-    """Test multiple concurrent requests from a single session."""
-    sse_url = mcp_server
-
-    async with sse_client(sse_url) as (read, write):
-        async with ClientSession(read, write) as session:
-            await session.initialize()
-
-            # Fire off 10 concurrent search requests from same session
-            async def make_search(idx):
-                try:
-                    result = await session.call_tool(
-                        "hindsight_search",
-                        arguments={
-                            "query": f"test query {idx}",
-                            "explanation": f"Concurrent test {idx}"
-                        }
-                    )
-                    return idx, "success", result
-                except Exception as e:
-                    return idx, "error", str(e)
-
-            tasks = [make_search(i) for i in range(10)]
-            results = await asyncio.gather(*tasks, return_exceptions=True)
-
-            # Check results
-            successes = 0
-            failures = 0
-
-            for result in results:
-                if isinstance(result, Exception):
-                    print(f"Request failed with exception: {result}")
-                    failures += 1
-                else:
-                    idx, status, data = result
-                    if status == "success":
-                        successes += 1
-                    else:
-                        print(f"Request {idx} failed: {data}")
-                        failures += 1
-
-            print(f"Successes: {successes}, Failures: {failures}")
-
-            # We expect all requests to succeed
-            assert successes >= 8, f"Too many failures: {failures}/10"
-
-
-@pytest.mark.asyncio
-async def test_race_condition_with_rapid_requests(mcp_server):
-    """Test rapid-fire requests with multiple sessions to trigger race condition."""
-    sse_url = mcp_server
-
-    async def rapid_session_search(idx):
-        """Create a new session and immediately make a request."""
-        try:
-            async with sse_client(sse_url) as (read, write):
-                async with ClientSession(read, write) as session:
-                    await session.initialize()
-
-                    # Make request immediately after initialization
-                    result = await session.call_tool(
-                        "hindsight_search",
-                        arguments={
-                            "query": f"rapid query {idx}",
-                            "max_tokens": 2048
-                        }
-                    )
-                    return idx, "success", result
-        except Exception as e:
-            return idx, "error", str(e)
-
-    # Fire 20 requests with minimal delay, each with its own session
-    tasks = [rapid_session_search(i) for i in range(20)]
-    results = await asyncio.gather(*tasks)
-
-    # Analyze results
-    errors = []
-    for idx, status, data in results:
-        if status == "error":
-            errors.append((idx, data))
-
-    if errors:
-        print(f"Found {len(errors)} errors:")
-        for idx, error_msg in errors:
-            print(f"  Request {idx}: {error_msg}")
-
-    # Most requests should succeed
-    assert len(errors) < 5, f"Too many errors: {len(errors)}/20"
-
-
-if __name__ == "__main__":
-    pytest.main([__file__, "-v", "-s"])

hindsight-api/tests/test_mcp_routing.py

@@ -9,6 +9,7 @@ def mock_memory():
     """Create a mock MemoryEngine."""
     memory = MagicMock()
     memory.retain_batch_async = AsyncMock()
+    memory.submit_async_retain = AsyncMock(return_value={"operation_id": "test-op-123"})
     memory.recall_async = AsyncMock(return_value=MagicMock(results=[]))
     return memory
 
@@ -44,11 +45,11 @@ async def test_mcp_tools_use_context_bank_id(mock_memory):
     assert "retain" in tools
     assert "recall" in tools
 
-    # Test retain with bank_id from context
+    # Test retain with bank_id from context (use async_processing=False for synchronous test)
     token = _current_bank_id.set("context-bank-id")
     try:
         retain_tool = tools["retain"]
-        result = await retain_tool.fn(content="test content", context="test_context")
+        result = await retain_tool.fn(content="test content", context="test_context", async_processing=False)
         assert "successfully" in result.lower()
 
         # Verify the memory was called with the context bank_id

hindsight-api/tests/test_metrics.py

@@ -0,0 +1,391 @@
+"""Tests for metrics instrumentation."""
+import pytest
+from unittest.mock import MagicMock, patch
+
+from hindsight_api.metrics import (
+    MetricsCollector,
+    MetricsCollectorBase,
+    NoOpMetricsCollector,
+    get_metrics_collector,
+    get_token_bucket,
+    create_metrics_collector,
+    initialize_metrics,
+)
+
+
+class TestNoOpMetricsCollector:
+    """Tests for the no-op metrics collector."""
+
+    def test_record_operation_is_noop(self):
+        """Test that record_operation does nothing."""
+        collector = NoOpMetricsCollector()
+
+        # Should not raise any exception
+        with collector.record_operation("recall", bank_id="test_bank", source="api"):
+            pass
+
+    def test_nested_contexts_work(self):
+        """Test that nested context managers work correctly."""
+        collector = NoOpMetricsCollector()
+
+        # Nested contexts should work without issues
+        with collector.record_operation("reflect", bank_id="test_bank", source="api"):
+            with collector.record_operation("recall", bank_id="test_bank", source="reflect"):
+                pass
+
+    def test_exception_propagates(self):
+        """Test that exceptions inside context are propagated."""
+        collector = NoOpMetricsCollector()
+
+        with pytest.raises(ValueError, match="test error"):
+            with collector.record_operation("recall", bank_id="test_bank"):
+                raise ValueError("test error")
+
+    def test_record_llm_call_is_noop(self):
+        """Test that record_llm_call does nothing."""
+        collector = NoOpMetricsCollector()
+
+        # Should not raise any exception
+        collector.record_llm_call(
+            provider="openai",
+            model="gpt-4",
+            scope="memory",
+            duration=1.5,
+            input_tokens=100,
+            output_tokens=50,
+            success=True,
+        )
+
+
+class TestMetricsCollector:
+    """Tests for the real metrics collector."""
+
+    @pytest.fixture
+    def mock_meter(self):
+        """Create a mock meter for testing."""
+        meter = MagicMock()
+        # Create separate mocks for each histogram (operation_duration, llm_duration)
+        histogram_mocks = [MagicMock(), MagicMock()]
+        meter.create_histogram.side_effect = histogram_mocks
+        # Create separate mocks for each counter
+        # (operation_total, llm_tokens_input, llm_tokens_output, llm_calls_total)
+        counter_mocks = [MagicMock() for _ in range(4)]
+        meter.create_counter.side_effect = counter_mocks
+        return meter
+
+    @pytest.fixture
+    def collector(self, mock_meter):
+        """Create a MetricsCollector with a mock meter."""
+        with patch("hindsight_api.metrics.get_meter", return_value=mock_meter):
+            return MetricsCollector()
+
+    def test_record_operation_records_duration(self, collector):
+        """Test that record_operation records duration."""
+        with collector.record_operation("recall", bank_id="test_bank", source="api"):
+            pass
+
+        # Histogram should have been called
+        collector.operation_duration.record.assert_called_once()
+        call_args = collector.operation_duration.record.call_args
+
+        # First arg is duration (should be > 0)
+        duration = call_args[0][0]
+        assert duration >= 0
+
+        # Second arg is attributes dict
+        attributes = call_args[0][1]
+        assert attributes["operation"] == "recall"
+        assert attributes["bank_id"] == "test_bank"
+        assert attributes["source"] == "api"
+        assert attributes["success"] == "true"
+
+    def test_record_operation_records_failure_on_exception(self, collector):
+        """Test that record_operation records failure when exception occurs."""
+        with pytest.raises(RuntimeError):
+            with collector.record_operation("retain", bank_id="test_bank", source="api"):
+                raise RuntimeError("Test error")
+
+        # Should have recorded with success=false
+        call_args = collector.operation_duration.record.call_args
+        attributes = call_args[0][1]
+        assert attributes["success"] == "false"
+
+    def test_record_operation_with_budget(self, collector):
+        """Test that budget is included in attributes when provided."""
+        with collector.record_operation("recall", bank_id="test_bank", source="api", budget="mid"):
+            pass
+
+        call_args = collector.operation_duration.record.call_args
+        attributes = call_args[0][1]
+        assert attributes["budget"] == "mid"
+
+    def test_record_operation_with_max_tokens(self, collector):
+        """Test that max_tokens is included in attributes when provided."""
+        with collector.record_operation("recall", bank_id="test_bank", source="api", max_tokens=4096):
+            pass
+
+        call_args = collector.operation_duration.record.call_args
+        attributes = call_args[0][1]
+        assert attributes["max_tokens"] == "4096"
+
+    def test_record_operation_source_values(self, collector):
+        """Test different source values: api, reflect, internal."""
+        sources = ["api", "reflect", "internal"]
+
+        for source in sources:
+            collector.operation_duration.record.reset_mock()
+
+            with collector.record_operation("recall", bank_id="test_bank", source=source):
+                pass
+
+            call_args = collector.operation_duration.record.call_args
+            attributes = call_args[0][1]
+            assert attributes["source"] == source
+
+    def test_nested_contexts_track_separately(self, collector):
+        """Test that nested operations are tracked separately with different sources."""
+        # Simulate reflect (api) calling recall (reflect)
+        with collector.record_operation("reflect", bank_id="test_bank", source="api"):
+            with collector.record_operation("recall", bank_id="test_bank", source="reflect"):
+                pass
+
+        # Should have 2 calls to record
+        assert collector.operation_duration.record.call_count == 2
+        assert collector.operation_total.add.call_count == 2
+
+        # Check the calls
+        calls = collector.operation_duration.record.call_args_list
+
+        # First call should be recall (inner context exits first)
+        recall_attrs = calls[0][0][1]
+        assert recall_attrs["operation"] == "recall"
+        assert recall_attrs["source"] == "reflect"
+
+        # Second call should be reflect (outer context exits last)
+        reflect_attrs = calls[1][0][1]
+        assert reflect_attrs["operation"] == "reflect"
+        assert reflect_attrs["source"] == "api"
+
+
+class TestGetMetricsCollector:
+    """Tests for the get_metrics_collector function."""
+
+    def test_returns_noop_by_default(self):
+        """Test that get_metrics_collector returns NoOpMetricsCollector by default."""
+        # Reset global state
+        import hindsight_api.metrics as metrics_module
+        original_collector = metrics_module._metrics_collector
+
+        try:
+            metrics_module._metrics_collector = NoOpMetricsCollector()
+            collector = get_metrics_collector()
+            assert isinstance(collector, NoOpMetricsCollector)
+        finally:
+            metrics_module._metrics_collector = original_collector
+
+
+class TestMetricsCollectorBase:
+    """Tests for the MetricsCollectorBase abstract class."""
+
+    def test_is_abstract(self):
+        """Test that MetricsCollectorBase methods are abstract."""
+        # Create a class that inherits but doesn't implement
+        class IncompleteCollector(MetricsCollectorBase):
+            pass
+
+        collector = IncompleteCollector()
+
+        # Abstract methods should raise NotImplementedError
+        with pytest.raises(NotImplementedError):
+            with collector.record_operation("test", "test"):
+                pass
+
+        with pytest.raises(NotImplementedError):
+            collector.record_llm_call("test", "test", "test", 1.0)
+
+
+class TestGetTokenBucket:
+    """Tests for the get_token_bucket function."""
+
+    def test_bucket_0_100(self):
+        """Test tokens < 100 return '0-100' bucket."""
+        assert get_token_bucket(0) == "0-100"
+        assert get_token_bucket(50) == "0-100"
+        assert get_token_bucket(99) == "0-100"
+
+    def test_bucket_100_500(self):
+        """Test tokens 100-499 return '100-500' bucket."""
+        assert get_token_bucket(100) == "100-500"
+        assert get_token_bucket(250) == "100-500"
+        assert get_token_bucket(499) == "100-500"
+
+    def test_bucket_500_1k(self):
+        """Test tokens 500-999 return '500-1k' bucket."""
+        assert get_token_bucket(500) == "500-1k"
+        assert get_token_bucket(750) == "500-1k"
+        assert get_token_bucket(999) == "500-1k"
+
+    def test_bucket_1k_5k(self):
+        """Test tokens 1000-4999 return '1k-5k' bucket."""
+        assert get_token_bucket(1000) == "1k-5k"
+        assert get_token_bucket(2500) == "1k-5k"
+        assert get_token_bucket(4999) == "1k-5k"
+
+    def test_bucket_5k_10k(self):
+        """Test tokens 5000-9999 return '5k-10k' bucket."""
+        assert get_token_bucket(5000) == "5k-10k"
+        assert get_token_bucket(7500) == "5k-10k"
+        assert get_token_bucket(9999) == "5k-10k"
+
+    def test_bucket_10k_50k(self):
+        """Test tokens 10000-49999 return '10k-50k' bucket."""
+        assert get_token_bucket(10000) == "10k-50k"
+        assert get_token_bucket(25000) == "10k-50k"
+        assert get_token_bucket(49999) == "10k-50k"
+
+    def test_bucket_50k_plus(self):
+        """Test tokens >= 50000 return '50k+' bucket."""
+        assert get_token_bucket(50000) == "50k+"
+        assert get_token_bucket(100000) == "50k+"
+        assert get_token_bucket(1000000) == "50k+"
+
+
+class TestLLMMetrics:
+    """Tests for LLM-specific metrics recording."""
+
+    @pytest.fixture
+    def mock_meter(self):
+        """Create a mock meter for testing."""
+        meter = MagicMock()
+        # Create separate mocks for each histogram (operation_duration, llm_duration)
+        histogram_mocks = [MagicMock(), MagicMock()]
+        meter.create_histogram.side_effect = histogram_mocks
+        # Create separate mocks for each counter
+        # (operation_total, llm_tokens_input, llm_tokens_output, llm_calls_total)
+        counter_mocks = [MagicMock() for _ in range(4)]
+        meter.create_counter.side_effect = counter_mocks
+        return meter
+
+    @pytest.fixture
+    def collector(self, mock_meter):
+        """Create a MetricsCollector with a mock meter."""
+        with patch("hindsight_api.metrics.get_meter", return_value=mock_meter):
+            return MetricsCollector()
+
+    def test_record_llm_call_records_duration(self, collector):
+        """Test that record_llm_call records duration."""
+        collector.record_llm_call(
+            provider="openai",
+            model="gpt-4",
+            scope="memory",
+            duration=1.5,
+            input_tokens=100,
+            output_tokens=50,
+            success=True,
+        )
+
+        # LLM duration histogram should be called
+        collector.llm_duration.record.assert_called_once()
+        call_args = collector.llm_duration.record.call_args
+
+        # First arg is duration
+        assert call_args[0][0] == 1.5
+
+        # Second arg is attributes dict
+        attributes = call_args[0][1]
+        assert attributes["provider"] == "openai"
+        assert attributes["model"] == "gpt-4"
+        assert attributes["scope"] == "memory"
+        assert attributes["success"] == "true"
+
+    def test_record_llm_call_records_failure(self, collector):
+        """Test that record_llm_call records failure status."""
+        collector.record_llm_call(
+            provider="anthropic",
+            model="claude-3",
+            scope="reflect",
+            duration=0.5,
+            success=False,
+        )
+
+        # Check success is false
+        call_args = collector.llm_duration.record.call_args
+        attributes = call_args[0][1]
+        assert attributes["success"] == "false"
+
+    def test_record_llm_call_records_tokens_with_buckets(self, collector):
+        """Test that record_llm_call records tokens with bucket labels."""
+        collector.record_llm_call(
+            provider="openai",
+            model="gpt-4",
+            scope="memory",
+            duration=1.0,
+            input_tokens=2500,  # Should be "1k-5k" bucket
+            output_tokens=150,  # Should be "100-500" bucket
+            success=True,
+        )
+
+        # Input tokens should be recorded with bucket
+        collector.llm_tokens_input.add.assert_called_once()
+        input_call = collector.llm_tokens_input.add.call_args
+        assert input_call[0][0] == 2500
+        assert input_call[0][1]["token_bucket"] == "1k-5k"
+
+        # Output tokens should be recorded with bucket
+        collector.llm_tokens_output.add.assert_called_once()
+        output_call = collector.llm_tokens_output.add.call_args
+        assert output_call[0][0] == 150
+        assert output_call[0][1]["token_bucket"] == "100-500"
+
+    def test_record_llm_call_skips_zero_tokens(self, collector):
+        """Test that zero token values don't record."""
+        collector.record_llm_call(
+            provider="openai",
+            model="gpt-4",
+            scope="memory",
+            duration=1.0,
+            input_tokens=0,
+            output_tokens=0,
+            success=True,
+        )
+
+        # Token counters should not be called
+        collector.llm_tokens_input.add.assert_not_called()
+        collector.llm_tokens_output.add.assert_not_called()
+
+    def test_record_llm_call_increments_call_counter(self, collector):
+        """Test that record_llm_call increments the call counter."""
+        collector.record_llm_call(
+            provider="gemini",
+            model="gemini-pro",
+            scope="entity_observation",
+            duration=2.0,
+            success=True,
+        )
+
+        # Call counter should be incremented
+        collector.llm_calls_total.add.assert_called_once()
+        call_args = collector.llm_calls_total.add.call_args
+        assert call_args[0][0] == 1
+        assert call_args[0][1]["provider"] == "gemini"
+        assert call_args[0][1]["model"] == "gemini-pro"
+        assert call_args[0][1]["scope"] == "entity_observation"
+
+    def test_record_llm_call_different_scopes(self, collector):
+        """Test recording LLM calls with different scopes."""
+        scopes = ["memory", "reflect", "entity_observation", "answer"]
+
+        for scope in scopes:
+            collector.llm_duration.record.reset_mock()
+
+            collector.record_llm_call(
+                provider="openai",
+                model="gpt-4",
+                scope=scope,
+                duration=1.0,
+                success=True,
+            )
+
+            call_args = collector.llm_duration.record.call_args
+            attributes = call_args[0][1]
+            assert attributes["scope"] == scope

hindsight-api/tests/test_multilingual.py

@@ -0,0 +1,286 @@
+"""
+Test multilingual support for retain and reflect operations.
+
+Tests that the system correctly handles non-English input and produces
+output in the same language as the input.
+"""
+
+import pytest
+import logging
+from datetime import datetime, timezone
+from hindsight_api.engine.memory_engine import Budget
+from hindsight_api import RequestContext
+
+logger = logging.getLogger(__name__)
+
+
+@pytest.mark.asyncio
+async def test_retain_chinese_content(memory, request_context):
+    """
+    Test that retain correctly extracts facts from Chinese content
+    and keeps the output in Chinese.
+
+    This test verifies:
+    1. Facts are extracted from Chinese text
+    2. The extracted facts contain Chinese characters
+    3. Entity names are preserved in Chinese
+    """
+    bank_id = f"test_chinese_retain_{datetime.now(timezone.utc).timestamp()}"
+
+    try:
+        # Chinese content about a person and their activities
+        chinese_content = """
+        张伟是一位资深软件工程师，在腾讯工作了五年。他专门研究分布式系统，
+        并领导了公司微服务架构的开发。他以编写干净、文档完善的代码而闻名。
+
+        李明上个月加入团队担任初级开发人员。他正在学习React和Node.js。
+        李明很有热情，在代码审查中提出很好的问题。他最近完成了他的第一个功能，
+        这是一个用户认证流程。
+
+        团队使用Kubernetes进行容器编排，并部署到阿里云。他们遵循敏捷方法论，
+        采用两周冲刺周期。合并前必须进行代码审查。
+        """
+
+        # Retain the Chinese content
+        unit_ids = await memory.retain_async(
+            bank_id=bank_id,
+            content=chinese_content,
+            context="团队概述",  # Chinese context
+            event_date=datetime(2024, 1, 15, tzinfo=timezone.utc),
+            request_context=request_context,
+        )
+
+        logger.info(f"Retained {len(unit_ids)} facts from Chinese content")
+        assert len(unit_ids) > 0, "Should have extracted and stored facts from Chinese content"
+
+        # Recall the facts with a Chinese query
+        result = await memory.recall_async(
+            bank_id=bank_id,
+            query="告诉我关于张伟的信息",  # "Tell me about Zhang Wei"
+            budget=Budget.MID,
+            max_tokens=1000,
+            fact_type=["world"],
+            request_context=request_context,
+        )
+
+        logger.info(f"Recalled {len(result.results)} facts")
+        assert len(result.results) > 0, "Should recall facts about Zhang Wei"
+
+        # Verify that the facts contain Chinese characters
+        # At least one fact should mention 张伟 (Zhang Wei) or related Chinese content
+        chinese_facts_found = 0
+        for fact in result.results:
+            logger.info(f"Fact: {fact.text[:100]}...")
+            # Check for common Chinese characters or the name
+            if any(
+                char in fact.text
+                for char in ["张", "伟", "腾讯", "软件", "工程师", "分布式", "系统", "代码"]
+            ):
+                chinese_facts_found += 1
+
+        logger.info(f"Found {chinese_facts_found} facts with Chinese content")
+        assert chinese_facts_found > 0, (
+            f"Expected facts to contain Chinese characters, but none found. "
+            f"Facts: {[f.text for f in result.results]}"
+        )
+
+        logger.info("Chinese retain test passed - facts preserved in Chinese")
+
+    finally:
+        await memory.delete_bank(bank_id, request_context=request_context)
+
+
+@pytest.mark.asyncio
+async def test_reflect_chinese_content(memory, request_context):
+    """
+    Test that reflect correctly generates responses in Chinese
+    when given Chinese facts and a Chinese query.
+
+    This test verifies:
+    1. Reflection produces a response in Chinese
+    2. The response references the Chinese facts
+    3. Opinions are formed and expressed in Chinese
+    """
+    bank_id = f"test_chinese_reflect_{datetime.now(timezone.utc).timestamp()}"
+
+    try:
+        # Store some Chinese facts to give context for opinion formation
+        await memory.retain_async(
+            bank_id=bank_id,
+            content="张伟是一位优秀的软件工程师，完成了五个重大项目。他总是按时交付，代码整洁有良好的文档。",
+            context="绩效评估",  # "Performance review"
+            event_date=datetime(2024, 1, 15, tzinfo=timezone.utc),
+            request_context=request_context,
+        )
+
+        await memory.retain_async(
+            bank_id=bank_id,
+            content="李明最近加入团队。他错过了第一个截止日期，代码有很多bug。",
+            context="绩效评估",
+            event_date=datetime(2024, 2, 1, tzinfo=timezone.utc),
+            request_context=request_context,
+        )
+
+        # Reflect with a Chinese query
+        query = "谁是更可靠的工程师？"  # "Who is a more reliable engineer?"
+        result = await memory.reflect_async(
+            bank_id=bank_id,
+            query=query,
+            budget=Budget.LOW,
+            request_context=request_context,
+        )
+
+        logger.info(f"Reflection answer: {result.text}")
+
+        # Verify we got an answer
+        assert result.text, "Reflection should return an answer"
+
+        # Check that the response contains Chinese characters
+        # The response should be in Chinese, not English
+        chinese_chars_found = sum(1 for char in result.text if "\u4e00" <= char <= "\u9fff")
+        total_chars = len(result.text.replace(" ", "").replace("\n", ""))
+
+        logger.info(f"Chinese characters: {chinese_chars_found}, Total characters: {total_chars}")
+
+        # At least 30% of characters should be Chinese (allowing for numbers, punctuation)
+        chinese_ratio = chinese_chars_found / max(total_chars, 1)
+        assert chinese_ratio > 0.3, (
+            f"Expected response to be in Chinese (>30% Chinese characters), "
+            f"but only {chinese_ratio:.1%} are Chinese. Response: {result.text}"
+        )
+
+        # Check that Chinese names are mentioned
+        assert "张伟" in result.text or "李明" in result.text, (
+            f"Expected response to mention Chinese names 张伟 or 李明. Response: {result.text}"
+        )
+
+        logger.info("Chinese reflect test passed - response generated in Chinese")
+
+    finally:
+        await memory.delete_bank(bank_id, request_context=request_context)
+
+
+@pytest.mark.asyncio
+async def test_retain_japanese_content(memory, request_context):
+    """
+    Test that retain correctly handles Japanese content.
+
+    This test verifies multilingual support extends beyond Chinese
+    to other non-Latin languages.
+    """
+    bank_id = f"test_japanese_retain_{datetime.now(timezone.utc).timestamp()}"
+
+    try:
+        # Japanese content about a developer
+        japanese_content = """
+        田中さんはソフトウェアエンジニアで、東京のスタートアップで働いています。
+        彼女はPythonとTypeScriptが得意で、毎日コードレビューをしています。
+        先週、新しいAPIを完成させました。
+        """
+
+        unit_ids = await memory.retain_async(
+            bank_id=bank_id,
+            content=japanese_content,
+            context="チームプロフィール",  # "Team profile"
+            event_date=datetime(2024, 1, 15, tzinfo=timezone.utc),
+            request_context=request_context,
+        )
+
+        logger.info(f"Retained {len(unit_ids)} facts from Japanese content")
+        assert len(unit_ids) > 0, "Should have extracted facts from Japanese content"
+
+        # Recall with Japanese query
+        result = await memory.recall_async(
+            bank_id=bank_id,
+            query="田中さんについて教えてください",  # "Tell me about Tanaka-san"
+            budget=Budget.MID,
+            max_tokens=1000,
+            fact_type=["world"],
+            request_context=request_context,
+        )
+
+        assert len(result.results) > 0, "Should recall facts about Tanaka"
+
+        # Check for Japanese content in facts
+        japanese_facts_found = 0
+        for fact in result.results:
+            logger.info(f"Fact: {fact.text[:100]}...")
+            # Check for Japanese characters (hiragana, katakana, or kanji)
+            if any(
+                ("\u3040" <= char <= "\u309f")  # Hiragana
+                or ("\u30a0" <= char <= "\u30ff")  # Katakana
+                or ("\u4e00" <= char <= "\u9fff")  # Kanji
+                for char in fact.text
+            ):
+                japanese_facts_found += 1
+
+        assert japanese_facts_found > 0, (
+            f"Expected facts to contain Japanese characters. "
+            f"Facts: {[f.text for f in result.results]}"
+        )
+
+        logger.info("Japanese retain test passed - facts preserved in Japanese")
+
+    finally:
+        await memory.delete_bank(bank_id, request_context=request_context)
+
+
+@pytest.mark.asyncio
+async def test_mixed_language_entities(memory, request_context):
+    """
+    Test that entity extraction works correctly with mixed language content.
+
+    Some entities (like company names) might be in English while the
+    description is in Chinese.
+    """
+    bank_id = f"test_mixed_lang_{datetime.now(timezone.utc).timestamp()}"
+
+    try:
+        # Mixed language content - Chinese with English company names
+        mixed_content = """
+        王芳在Google北京办公室工作，她是一名高级产品经理。
+        之前她在Microsoft和Amazon工作过。
+        她负责管理YouTube在中国市场的推广策略。
+        """
+
+        unit_ids = await memory.retain_async(
+            bank_id=bank_id,
+            content=mixed_content,
+            context="员工资料",
+            event_date=datetime(2024, 1, 15, tzinfo=timezone.utc),
+            request_context=request_context,
+        )
+
+        assert len(unit_ids) > 0, "Should extract facts from mixed language content"
+
+        # Recall and check entities
+        result = await memory.recall_async(
+            bank_id=bank_id,
+            query="王芳在哪里工作？",  # "Where does Wang Fang work?"
+            budget=Budget.MID,
+            max_tokens=1000,
+            fact_type=["world"],
+            include_entities=True,
+            request_context=request_context,
+        )
+
+        assert len(result.results) > 0, "Should recall facts about Wang Fang"
+
+        # Check that both Chinese and English entities are preserved
+        all_text = " ".join(f.text for f in result.results)
+        logger.info(f"Combined facts: {all_text}")
+
+        # Should contain Chinese name and/or English company names
+        has_chinese_name = "王芳" in all_text
+        has_english_company = any(
+            company in all_text for company in ["Google", "Microsoft", "Amazon", "YouTube"]
+        )
+
+        assert has_chinese_name or has_english_company, (
+            f"Expected mixed language entities. Facts: {all_text}"
+        )
+
+        logger.info("Mixed language entity test passed")
+
+    finally:
+        await memory.delete_bank(bank_id, request_context=request_context)

hindsight-api/tests/test_observations.py

@@ -3,11 +3,12 @@ Test observation generation and entity state functionality.
 """
 import pytest
 from hindsight_api.engine.memory_engine import Budget
+from hindsight_api import RequestContext
 from datetime import datetime, timezone
 
 
 @pytest.mark.asyncio
-async def test_observation_generation_on_put(memory):
+async def test_observation_generation_on_put(memory, request_context):
     """
     Test that observations are generated SYNCHRONOUSLY when new facts are added.
 
@@ -36,7 +37,8 @@ async def test_observation_generation_on_put(memory):
                 bank_id=bank_id,
                 content=content,
                 context="work info",
-                event_date=datetime(2024, 1, 15 + i, tzinfo=timezone.utc)
+                event_date=datetime(2024, 1, 15 + i, tzinfo=timezone.utc),
+                request_context=request_context,
             )
 
         # Observations are generated SYNCHRONOUSLY during retain,
@@ -75,7 +77,7 @@ async def test_observation_generation_on_put(memory):
         print(f"Entity: {entity_name} (id: {entity_id})")
 
         # Get observations for the entity - should be available immediately
-        observations = await memory.get_entity_observations(bank_id, entity_id, limit=10)
+        observations = await memory.get_entity_observations(bank_id, entity_id, limit=10, request_context=request_context)
 
         print(f"\n=== Observations for {entity_name} ===")
         print(f"Total observations: {len(observations)}")
@@ -102,7 +104,7 @@ async def test_observation_generation_on_put(memory):
 
 
 @pytest.mark.asyncio
-async def test_regenerate_entity_observations(memory):
+async def test_regenerate_entity_observations(memory, request_context):
     """
     Test explicit regeneration of observations for an entity.
     """
@@ -114,7 +116,8 @@ async def test_regenerate_entity_observations(memory):
             bank_id=bank_id,
             content="Sarah is a product manager who loves user research and data analysis.",
             context="work info",
-            event_date=datetime(2024, 1, 15, tzinfo=timezone.utc)
+            event_date=datetime(2024, 1, 15, tzinfo=timezone.utc),
+            request_context=request_context,
         )
 
         await memory.wait_for_background_tasks()
@@ -140,14 +143,15 @@ async def test_regenerate_entity_observations(memory):
             created_ids = await memory.regenerate_entity_observations(
                 bank_id=bank_id,
                 entity_id=entity_id,
-                entity_name=entity_name
+                entity_name=entity_name,
+                request_context=request_context,
             )
 
             print(f"\n=== Regenerated Observations ===")
             print(f"Created {len(created_ids)} observations for {entity_name}")
 
             # Get the observations
-            observations = await memory.get_entity_observations(bank_id, entity_id, limit=10)
+            observations = await memory.get_entity_observations(bank_id, entity_id, limit=10, request_context=request_context)
             for obs in observations:
                 print(f"  - {obs.text}")
 
@@ -170,7 +174,108 @@ async def test_regenerate_entity_observations(memory):
 
 
 @pytest.mark.asyncio
-async def test_search_with_include_entities(memory):
+async def test_manual_regenerate_with_few_facts(memory, request_context):
+    """
+    Test that manual regeneration works even with fewer than 5 facts.
+
+    This is important because:
+    - Automatic generation during retain requires MIN_FACTS_THRESHOLD (5)
+    - But manual regeneration via API should work with any number of facts
+    - The UI triggers manual regeneration, so it should work regardless of fact count
+    """
+    bank_id = f"test_manual_regen_{datetime.now(timezone.utc).timestamp()}"
+
+    try:
+        # Store only 2 facts - below the automatic threshold
+        await memory.retain_async(
+            bank_id=bank_id,
+            content="Alice works at Google as a senior software engineer.",
+            context="work info",
+            event_date=datetime(2024, 1, 15, tzinfo=timezone.utc),
+            request_context=request_context,
+        )
+        await memory.retain_async(
+            bank_id=bank_id,
+            content="Alice loves hiking and outdoor photography.",
+            context="hobbies",
+            event_date=datetime(2024, 1, 16, tzinfo=timezone.utc),
+            request_context=request_context,
+        )
+
+        # Find the Alice entity
+        pool = await memory._get_pool()
+        async with pool.acquire() as conn:
+            entity_row = await conn.fetchrow(
+                """
+                SELECT id, canonical_name
+                FROM entities
+                WHERE bank_id = $1 AND LOWER(canonical_name) LIKE '%alice%'
+                LIMIT 1
+                """,
+                bank_id
+            )
+
+        assert entity_row is not None, "Alice entity should have been extracted"
+
+        entity_id = str(entity_row['id'])
+        entity_name = entity_row['canonical_name']
+
+        # Check fact count - should be < 5
+        async with pool.acquire() as conn:
+            fact_count = await conn.fetchval(
+                "SELECT COUNT(*) FROM unit_entities WHERE entity_id = $1",
+                entity_row['id']
+            )
+
+        print(f"\n=== Manual Regeneration Test ===")
+        print(f"Entity: {entity_name} (id: {entity_id})")
+        print(f"Linked facts: {fact_count}")
+
+        # Verify we're testing with fewer than the automatic threshold
+        assert fact_count < 5, f"Test requires < 5 facts, but entity has {fact_count}"
+
+        # Before regeneration - should have no observations (auto threshold not met)
+        obs_before = await memory.get_entity_observations(bank_id, entity_id, limit=10, request_context=request_context)
+        print(f"Observations before manual regenerate: {len(obs_before)}")
+
+        # Manually regenerate observations - this should work regardless of fact count
+        created_ids = await memory.regenerate_entity_observations(
+            bank_id=bank_id,
+            entity_id=entity_id,
+            entity_name=entity_name,
+            request_context=request_context,
+        )
+
+        print(f"Observations created by manual regenerate: {len(created_ids)}")
+
+        # Get observations after regeneration
+        observations = await memory.get_entity_observations(bank_id, entity_id, limit=10, request_context=request_context)
+        print(f"Observations after manual regenerate: {len(observations)}")
+        for obs in observations:
+            print(f"  - {obs.text}")
+
+        # Manual regeneration should create observations even with < 5 facts
+        assert len(observations) > 0, \
+            f"Manual regeneration should create observations even with only {fact_count} facts. " \
+            f"The LLM should synthesize at least 1 observation from the available facts."
+
+        # Verify observations contain relevant content
+        obs_texts = " ".join([o.text.lower() for o in observations])
+        assert any(keyword in obs_texts for keyword in ["google", "engineer", "hiking", "photography", "alice"]), \
+            "Observations should contain relevant information about Alice"
+
+        print(f"✓ Manual regeneration works with {fact_count} facts (below automatic threshold of 5)")
+
+    finally:
+        # Cleanup
+        pool = await memory._get_pool()
+        async with pool.acquire() as conn:
+            await conn.execute("DELETE FROM memory_units WHERE bank_id = $1", bank_id)
+            await conn.execute("DELETE FROM entities WHERE bank_id = $1", bank_id)
+
+
+@pytest.mark.asyncio
+async def test_search_with_include_entities(memory, request_context):
     """
     Test that search with include_entities=True returns entity observations.
 
@@ -196,7 +301,8 @@ async def test_search_with_include_entities(memory):
                 bank_id=bank_id,
                 content=content,
                 context="work info",
-                event_date=datetime(2024, 1, 15 + i, tzinfo=timezone.utc)
+                event_date=datetime(2024, 1, 15 + i, tzinfo=timezone.utc),
+                request_context=request_context,
             )
 
         # Observations are generated synchronously during retain, no need to wait
@@ -209,7 +315,8 @@ async def test_search_with_include_entities(memory):
             budget=Budget.LOW,
             max_tokens=2000,
             include_entities=True,
-            max_entity_tokens=500
+            max_entity_tokens=500,
+            request_context=request_context,
         )
 
         print(f"\n=== Search Results ===")
@@ -263,7 +370,7 @@ async def test_search_with_include_entities(memory):
 
 
 @pytest.mark.asyncio
-async def test_get_entity_state(memory):
+async def test_get_entity_state(memory, request_context):
     """
     Test getting the full state of an entity.
     """
@@ -275,7 +382,8 @@ async def test_get_entity_state(memory):
             bank_id=bank_id,
             content="Bob is a frontend developer who specializes in React and TypeScript.",
             context="work info",
-            event_date=datetime(2024, 1, 15, tzinfo=timezone.utc)
+            event_date=datetime(2024, 1, 15, tzinfo=timezone.utc),
+            request_context=request_context,
         )
 
         await memory.wait_for_background_tasks()
@@ -302,7 +410,8 @@ async def test_get_entity_state(memory):
                 bank_id=bank_id,
                 entity_id=entity_id,
                 entity_name=entity_name,
-                limit=10
+                limit=10,
+                request_context=request_context,
             )
 
             print(f"\n=== Entity State for {entity_name} ===")
@@ -324,7 +433,7 @@ async def test_get_entity_state(memory):
 
 
 @pytest.mark.asyncio
-async def test_observation_fact_type_in_database(memory):
+async def test_observation_fact_type_in_database(memory, request_context):
     """
     Test that observations are stored with correct fact_type in database.
     """
@@ -336,7 +445,8 @@ async def test_observation_fact_type_in_database(memory):
             bank_id=bank_id,
             content="Charlie is a DevOps engineer who manages the Kubernetes infrastructure.",
             context="work info",
-            event_date=datetime(2024, 1, 15, tzinfo=timezone.utc)
+            event_date=datetime(2024, 1, 15, tzinfo=timezone.utc),
+            request_context=request_context,
         )
 
         await memory.wait_for_background_tasks()
@@ -374,7 +484,7 @@ async def test_observation_fact_type_in_database(memory):
 
 
 @pytest.mark.asyncio
-async def test_user_entity_prioritized_for_observations(memory):
+async def test_user_entity_prioritized_for_observations(memory, request_context):
     """
     Test that the 'user' entity gets observations even when many other entities exist.
 
@@ -410,7 +520,8 @@ async def test_user_entity_prioritized_for_observations(memory):
                 bank_id=bank_id,
                 content=content,
                 context="personal info",
-                event_date=datetime(2024, 1, 15 + i, tzinfo=timezone.utc)
+                event_date=datetime(2024, 1, 15 + i, tzinfo=timezone.utc),
+                request_context=request_context,
             )
 
         # Observations are generated synchronously during retain
@@ -466,7 +577,7 @@ async def test_user_entity_prioritized_for_observations(memory):
             f"User entity should have at least 5 facts, but has {user_fact_count}"
 
         # Get observations for user entity
-        observations = await memory.get_entity_observations(bank_id, user_entity_id, limit=10)
+        observations = await memory.get_entity_observations(bank_id, user_entity_id, limit=10, request_context=request_context)
 
         print(f"\n=== User Entity Observations ===")
         print(f"Total observations: {len(observations)}")

hindsight-api/tests/test_per_operation_llm_config.py

@@ -0,0 +1,277 @@
+"""
+Tests for per-operation LLM configuration.
+
+Verifies that retain and reflect operations use their respective LLM configs.
+"""
+
+import os
+
+import pytest
+
+
+@pytest.fixture(autouse=True)
+def setup_test_env():
+    """Set up environment for each test, restoring original values after."""
+    from hindsight_api.config import clear_config_cache
+
+    # Save original environment values
+    env_vars_to_set = {
+        "HINDSIGHT_API_SKIP_LLM_VERIFICATION": "true",
+        "HINDSIGHT_API_LAZY_RERANKER": "true",
+        "HINDSIGHT_API_LLM_PROVIDER": "mock",
+        "HINDSIGHT_API_LLM_MODEL": "default-model",
+        "HINDSIGHT_API_RETAIN_LLM_PROVIDER": "mock",
+        "HINDSIGHT_API_RETAIN_LLM_MODEL": "retain-model",
+        "HINDSIGHT_API_REFLECT_LLM_PROVIDER": "mock",
+        "HINDSIGHT_API_REFLECT_LLM_MODEL": "reflect-model",
+    }
+
+    # Save original values
+    original_values = {}
+    for key in env_vars_to_set:
+        original_values[key] = os.environ.get(key)
+
+    # Set test values
+    for key, value in env_vars_to_set.items():
+        os.environ[key] = value
+
+    clear_config_cache()
+
+    yield
+
+    # Restore original environment
+    for key, original_value in original_values.items():
+        if original_value is None:
+            os.environ.pop(key, None)
+        else:
+            os.environ[key] = original_value
+
+    clear_config_cache()
+
+
+class TestPerOperationLLMConfig:
+    """Test that per-operation LLM configs are correctly applied."""
+
+    def test_config_loads_per_operation_settings(self):
+        """Test that config correctly loads per-operation LLM settings."""
+        from hindsight_api.config import get_config
+
+        config = get_config()
+
+        # Default config
+        assert config.llm_provider == "mock"
+        assert config.llm_model == "default-model"
+
+        # Retain config
+        assert config.retain_llm_provider == "mock"
+        assert config.retain_llm_model == "retain-model"
+
+        # Reflect config
+        assert config.reflect_llm_provider == "mock"
+        assert config.reflect_llm_model == "reflect-model"
+
+    def test_memory_engine_creates_separate_llm_configs(self):
+        """Test that MemoryEngine creates separate LLM configs for each operation."""
+        from hindsight_api import MemoryEngine
+
+        engine = MemoryEngine(
+            skip_llm_verification=True,
+            lazy_reranker=True,
+        )
+
+        # Verify default config
+        assert engine._llm_config.provider == "mock"
+        assert engine._llm_config.model == "default-model"
+
+        # Verify retain config
+        assert engine._retain_llm_config.provider == "mock"
+        assert engine._retain_llm_config.model == "retain-model"
+
+        # Verify reflect config
+        assert engine._reflect_llm_config.provider == "mock"
+        assert engine._reflect_llm_config.model == "reflect-model"
+
+    def test_memory_engine_with_explicit_params(self):
+        """Test that explicit params override env config."""
+        from hindsight_api import MemoryEngine
+
+        engine = MemoryEngine(
+            memory_llm_provider="mock",
+            memory_llm_model="explicit-default",
+            retain_llm_provider="mock",
+            retain_llm_model="explicit-retain",
+            reflect_llm_provider="mock",
+            reflect_llm_model="explicit-reflect",
+            skip_llm_verification=True,
+            lazy_reranker=True,
+        )
+
+        assert engine._llm_config.model == "explicit-default"
+        assert engine._retain_llm_config.model == "explicit-retain"
+        assert engine._reflect_llm_config.model == "explicit-reflect"
+
+    def test_memory_engine_fallback_when_no_per_operation_config(self):
+        """Test that per-operation configs fall back to default when not set."""
+        from hindsight_api.config import clear_config_cache as clear_cache
+
+        # Temporarily clear per-operation env vars
+        retain_provider = os.environ.pop("HINDSIGHT_API_RETAIN_LLM_PROVIDER", None)
+        retain_model = os.environ.pop("HINDSIGHT_API_RETAIN_LLM_MODEL", None)
+        reflect_provider = os.environ.pop("HINDSIGHT_API_REFLECT_LLM_PROVIDER", None)
+        reflect_model = os.environ.pop("HINDSIGHT_API_REFLECT_LLM_MODEL", None)
+
+        try:
+            clear_cache()
+            from hindsight_api import MemoryEngine
+
+            engine = MemoryEngine(
+                skip_llm_verification=True,
+                lazy_reranker=True,
+            )
+
+            # All should fall back to default
+            assert engine._llm_config.model == "default-model"
+            assert engine._retain_llm_config.model == "default-model"
+            assert engine._reflect_llm_config.model == "default-model"
+        finally:
+            # Restore env vars
+            if retain_provider:
+                os.environ["HINDSIGHT_API_RETAIN_LLM_PROVIDER"] = retain_provider
+            if retain_model:
+                os.environ["HINDSIGHT_API_RETAIN_LLM_MODEL"] = retain_model
+            if reflect_provider:
+                os.environ["HINDSIGHT_API_REFLECT_LLM_PROVIDER"] = reflect_provider
+            if reflect_model:
+                os.environ["HINDSIGHT_API_REFLECT_LLM_MODEL"] = reflect_model
+            clear_cache()
+
+
+class TestMockLLMProvider:
+    """Test the mock LLM provider functionality."""
+
+    def test_mock_provider_records_calls(self):
+        """Test that mock provider records calls."""
+        from hindsight_api.engine.llm_wrapper import LLMProvider
+
+        provider = LLMProvider(
+            provider="mock",
+            api_key="",
+            base_url="",
+            model="test-model",
+        )
+
+        import asyncio
+
+        async def make_call():
+            return await provider.call(
+                messages=[{"role": "user", "content": "test"}],
+                scope="test_scope",
+            )
+
+        result = asyncio.get_event_loop().run_until_complete(make_call())
+
+        # Verify call was recorded
+        calls = provider.get_mock_calls()
+        assert len(calls) == 1
+        assert calls[0]["model"] == "test-model"
+        assert calls[0]["scope"] == "test_scope"
+        assert calls[0]["messages"] == [{"role": "user", "content": "test"}]
+
+    def test_mock_provider_returns_custom_response(self):
+        """Test that mock provider can return custom responses."""
+        from hindsight_api.engine.llm_wrapper import LLMProvider
+
+        provider = LLMProvider(
+            provider="mock",
+            api_key="",
+            base_url="",
+            model="test-model",
+        )
+
+        provider.set_mock_response({"custom": "response"})
+
+        import asyncio
+
+        async def make_call():
+            return await provider.call(
+                messages=[{"role": "user", "content": "test"}],
+            )
+
+        result = asyncio.get_event_loop().run_until_complete(make_call())
+        assert result == {"custom": "response"}
+
+    def test_mock_provider_returns_usage_when_requested(self):
+        """Test that mock provider returns token usage."""
+        from hindsight_api.engine.llm_wrapper import LLMProvider
+
+        provider = LLMProvider(
+            provider="mock",
+            api_key="",
+            base_url="",
+            model="test-model",
+        )
+
+        import asyncio
+
+        async def make_call():
+            return await provider.call(
+                messages=[{"role": "user", "content": "test"}],
+                return_usage=True,
+            )
+
+        result, usage = asyncio.get_event_loop().run_until_complete(make_call())
+        assert usage.input_tokens == 10
+        assert usage.output_tokens == 5
+        assert usage.total_tokens == 15
+
+
+class TestRetainUsesRetainLLMConfig:
+    """Test that retain operations use the retain LLM config."""
+
+    def test_retain_llm_config_is_passed_to_orchestrator(self):
+        """Verify retain operation is configured to use _retain_llm_config."""
+        from hindsight_api import MemoryEngine
+
+        engine = MemoryEngine(
+            memory_llm_provider="mock",
+            memory_llm_model="default-model",
+            retain_llm_provider="mock",
+            retain_llm_model="retain-specific-model",
+            reflect_llm_provider="mock",
+            reflect_llm_model="reflect-specific-model",
+            skip_llm_verification=True,
+            lazy_reranker=True,
+        )
+
+        # Verify the retain LLM config is set correctly
+        assert engine._retain_llm_config.model == "retain-specific-model"
+        assert engine._retain_llm_config.provider == "mock"
+
+        # Verify it's different from the reflect config
+        assert engine._retain_llm_config.model != engine._reflect_llm_config.model
+
+
+class TestReflectUsesReflectLLMConfig:
+    """Test that reflect operations use the reflect LLM config."""
+
+    def test_reflect_llm_config_is_set_correctly(self):
+        """Verify reflect/think operation is configured to use _reflect_llm_config."""
+        from hindsight_api import MemoryEngine
+
+        engine = MemoryEngine(
+            memory_llm_provider="mock",
+            memory_llm_model="default-model",
+            retain_llm_provider="mock",
+            retain_llm_model="retain-specific-model",
+            reflect_llm_provider="mock",
+            reflect_llm_model="reflect-specific-model",
+            skip_llm_verification=True,
+            lazy_reranker=True,
+        )
+
+        # Verify the reflect LLM config is set correctly
+        assert engine._reflect_llm_config.model == "reflect-specific-model"
+        assert engine._reflect_llm_config.provider == "mock"
+
+        # Verify it's different from the retain config
+        assert engine._reflect_llm_config.model != engine._retain_llm_config.model

hindsight-api/tests/test_retain.py

@@ -5,12 +5,13 @@ import pytest
 import logging
 from datetime import datetime, timezone, timedelta
 from hindsight_api.engine.memory_engine import Budget
+from hindsight_api import RequestContext
 
 logger = logging.getLogger(__name__)
 
 
 @pytest.mark.asyncio
-async def test_retain_with_chunks(memory):
+async def test_retain_with_chunks(memory, request_context):
     """
     Test that retain function:
     1. Stores facts with associated chunks
@@ -41,7 +42,8 @@ async def test_retain_with_chunks(memory):
             content=long_content,
             context="team overview",
             event_date=datetime(2024, 1, 15, tzinfo=timezone.utc),
-            document_id=document_id
+            document_id=document_id,
+            request_context=request_context,
         )
 
         print(f"\n=== Retained {len(unit_ids)} facts ===")
@@ -56,7 +58,8 @@ async def test_retain_with_chunks(memory):
             fact_type=["world"],  # Search for world facts
             include_entities=False,  # Disable entities for simpler test
             include_chunks=True,  # Enable chunks
-            max_chunk_tokens=8192
+            max_chunk_tokens=8192,
+            request_context=request_context,
         )
 
         print(f"\n=== Recall Results (with chunks) ===")
@@ -88,12 +91,12 @@ async def test_retain_with_chunks(memory):
 
     finally:
         # Cleanup - delete the test bank
-        await memory.delete_bank(bank_id)
+        await memory.delete_bank(bank_id, request_context=request_context)
         print(f"\n=== Cleaned up bank: {bank_id} ===")
 
 
 @pytest.mark.asyncio
-async def test_chunks_and_entities_follow_fact_order(memory):
+async def test_chunks_and_entities_follow_fact_order(memory, request_context):
     """
     Test that chunks and entities in recall results follow the same order as facts.
     This is critical because token limits may truncate later items.
@@ -130,7 +133,8 @@ async def test_chunks_and_entities_follow_fact_order(memory):
                 content=item["content"],
                 context=item["context"],
                 event_date=datetime(2024, 1, 15, tzinfo=timezone.utc),
-                document_id=item["document_id"]
+                document_id=item["document_id"],
+                request_context=request_context,
             )
 
         print("\n=== Stored 3 separate documents ===")
@@ -144,7 +148,8 @@ async def test_chunks_and_entities_follow_fact_order(memory):
             fact_type=["world"],
             include_entities=True,
             include_chunks=True,
-            max_chunk_tokens=8192
+            max_chunk_tokens=8192,
+            request_context=request_context,
         )
 
         print(f"\n=== Recall Results ===")
@@ -214,12 +219,12 @@ async def test_chunks_and_entities_follow_fact_order(memory):
 
     finally:
         # Cleanup
-        await memory.delete_bank(bank_id)
+        await memory.delete_bank(bank_id, request_context=request_context)
         print(f"\n=== Cleaned up bank: {bank_id} ===")
 
 
 @pytest.mark.asyncio
-async def test_event_date_storage(memory):
+async def test_event_date_storage(memory, request_context):
     """
     Test that event_date is correctly stored as occurred_start.
     Verifies that we can track when events actually happened vs when they were stored.
@@ -235,7 +240,8 @@ async def test_event_date_storage(memory):
             bank_id=bank_id,
             content="Alice completed the Q2 product launch on June 15th, 2023.",
             context="project history",
-            event_date=past_event_date
+            event_date=past_event_date,
+            request_context=request_context,
         )
 
         assert len(unit_ids) > 0, "Should have created at least one memory unit"
@@ -246,7 +252,8 @@ async def test_event_date_storage(memory):
             query="When did Alice complete the product launch?",
             budget=Budget.LOW,
             max_tokens=500,
-            fact_type=["world"]
+            fact_type=["world"],
+            request_context=request_context,
         )
 
         assert len(result.results) > 0, "Should recall the stored fact"
@@ -268,11 +275,11 @@ async def test_event_date_storage(memory):
         print(f"\n✓ Event date correctly stored: {occurred_dt}")
 
     finally:
-        await memory.delete_bank(bank_id)
+        await memory.delete_bank(bank_id, request_context=request_context)
 
 
 @pytest.mark.asyncio
-async def test_temporal_ordering(memory):
+async def test_temporal_ordering(memory, request_context):
     """
     Test that facts can be stored and retrieved with correct temporal ordering.
     Stores facts with different event_dates and verifies temporal relationships.
@@ -305,7 +312,8 @@ async def test_temporal_ordering(memory):
                 bank_id=bank_id,
                 content=event["content"],
                 context=event["context"],
-                event_date=event["event_date"]
+                event_date=event["event_date"],
+                request_context=request_context,
             )
 
         print("\n=== Stored 3 events with different temporal dates ===")
@@ -316,7 +324,8 @@ async def test_temporal_ordering(memory):
             query="Tell me about Alice's career progression",
             budget=Budget.MID,
             max_tokens=1000,
-            fact_type=["world"]
+            fact_type=["world"],
+            request_context=request_context,
         )
 
         assert len(result.results) >= 3, f"Should recall all 3 events, got {len(result.results)}"
@@ -345,11 +354,11 @@ async def test_temporal_ordering(memory):
         print(f"\n✓ Temporal ordering preserved: {min_date.date()} to {max_date.date()}")
 
     finally:
-        await memory.delete_bank(bank_id)
+        await memory.delete_bank(bank_id, request_context=request_context)
 
 
 @pytest.mark.asyncio
-async def test_mentioned_at_vs_occurred(memory):
+async def test_mentioned_at_vs_occurred(memory, request_context):
     """
     Test distinction between when fact occurred vs when it was mentioned.
 
@@ -369,7 +378,8 @@ async def test_mentioned_at_vs_occurred(memory):
             bank_id=bank_id,
             content="Alice graduated from MIT in March 2020.",
             context="education history",
-            event_date=conversation_date  # When this conversation happened
+            event_date=conversation_date,  # When this conversation happened
+            request_context=request_context,
         )
 
         assert len(unit_ids) > 0, "Should create memory unit"
@@ -380,7 +390,8 @@ async def test_mentioned_at_vs_occurred(memory):
             query="Where did Alice go to school?",
             budget=Budget.LOW,
             max_tokens=500,
-            fact_type=["world"]
+            fact_type=["world"],
+            request_context=request_context,
         )
 
         assert len(result.results) > 0, "Should recall the fact"
@@ -415,11 +426,11 @@ async def test_mentioned_at_vs_occurred(memory):
         print(f"✓ Test passed: Historical conversation correctly ingested with event_date=2020")
 
     finally:
-        await memory.delete_bank(bank_id)
+        await memory.delete_bank(bank_id, request_context=request_context)
 
 
 @pytest.mark.asyncio
-async def test_occurred_dates_not_defaulted(memory):
+async def test_occurred_dates_not_defaulted(memory, request_context):
     """
     Test that occurred_start and occurred_end are NOT defaulted to mentioned_at.
 
@@ -441,7 +452,8 @@ async def test_occurred_dates_not_defaulted(memory):
             bank_id=bank_id,
             content="Alice likes coffee. The weather is sunny today.",
             context="current observations",
-            event_date=event_date
+            event_date=event_date,
+            request_context=request_context,
         )
 
         assert len(unit_ids) > 0, "Should create memory unit"
@@ -452,7 +464,8 @@ async def test_occurred_dates_not_defaulted(memory):
             query="What does Alice like?",
             budget=Budget.LOW,
             max_tokens=500,
-            fact_type=["world", "opinion"]
+            fact_type=["world", "opinion"],
+            request_context=request_context,
         )
 
         assert len(result.results) > 0, "Should recall the fact"
@@ -504,11 +517,11 @@ async def test_occurred_dates_not_defaulted(memory):
         print(f"✓ Test passed: occurred dates are not incorrectly defaulted to mentioned_at")
 
     finally:
-        await memory.delete_bank(bank_id)
+        await memory.delete_bank(bank_id, request_context=request_context)
 
 
 @pytest.mark.asyncio
-async def test_mentioned_at_from_context_string(memory):
+async def test_mentioned_at_from_context_string(memory, request_context):
     """
     Test that mentioned_at is extracted from context string by LLM.
 
@@ -527,7 +540,8 @@ async def test_mentioned_at_from_context_string(memory):
             bank_id=bank_id,
             content="Alice mentioned she loves hiking in the mountains.",
             context=f"Session ABC123 - you are the assistant in this conversation - happened on {session_date.strftime('%Y-%m-%d %H:%M:%S')} UTC.",
-            event_date=None  # Not providing event_date - should default to now() if LLM doesn't extract
+            event_date=None,  # Not providing event_date - should default to now() if LLM doesn't extract
+            request_context=request_context,
         )
 
         assert len(unit_ids) > 0, "Should create memory unit"
@@ -538,7 +552,8 @@ async def test_mentioned_at_from_context_string(memory):
             query="What does Alice like?",
             budget=Budget.LOW,
             max_tokens=500,
-            fact_type=["world"]
+            fact_type=["world"],
+            request_context=request_context,
         )
 
         assert len(result.results) > 0, "Should recall the fact"
@@ -574,7 +589,7 @@ async def test_mentioned_at_from_context_string(memory):
         print(f"✓ mentioned_at is always set (never None)")
 
     finally:
-        await memory.delete_bank(bank_id)
+        await memory.delete_bank(bank_id, request_context=request_context)
 
 
 # ============================================================
@@ -582,7 +597,7 @@ async def test_mentioned_at_from_context_string(memory):
 # ============================================================
 
 @pytest.mark.asyncio
-async def test_context_preservation(memory):
+async def test_context_preservation(memory, request_context):
     """
     Test that context is preserved and retrievable.
     Context helps understand why/how memory was formed.
@@ -597,7 +612,8 @@ async def test_context_preservation(memory):
             bank_id=bank_id,
             content="The team decided to prioritize mobile development for next quarter.",
             context=specific_context,
-            event_date=datetime(2024, 1, 15, tzinfo=timezone.utc)
+            event_date=datetime(2024, 1, 15, tzinfo=timezone.utc),
+            request_context=request_context,
         )
 
         assert len(unit_ids) > 0, "Should create at least one memory unit"
@@ -608,7 +624,8 @@ async def test_context_preservation(memory):
             query="What did the team decide?",
             budget=Budget.LOW,
             max_tokens=500,
-            fact_type=["world"]
+            fact_type=["world"],
+            request_context=request_context,
         )
 
         assert len(result.results) > 0, "Should recall the stored fact"
@@ -620,11 +637,11 @@ async def test_context_preservation(memory):
         print(f"  Retrieved {len(result.results)} facts")
 
     finally:
-        await memory.delete_bank(bank_id)
+        await memory.delete_bank(bank_id, request_context=request_context)
 
 
 @pytest.mark.asyncio
-async def test_context_with_batch(memory):
+async def test_context_with_batch(memory, request_context):
     """
     Test that each item in a batch can have different contexts.
     """
@@ -650,7 +667,8 @@ async def test_context_with_batch(memory):
                     "context": "incident response",
                     "event_date": datetime(2024, 1, 12, tzinfo=timezone.utc)
                 }
-            ]
+            ],
+            request_context=request_context,
         )
 
         # Should have created facts from all items
@@ -661,7 +679,7 @@ async def test_context_with_batch(memory):
         print(f"  Created {total_units} total memory units")
 
     finally:
-        await memory.delete_bank(bank_id)
+        await memory.delete_bank(bank_id, request_context=request_context)
 
 
 # ============================================================
@@ -669,7 +687,7 @@ async def test_context_with_batch(memory):
 # ============================================================
 
 @pytest.mark.asyncio
-async def test_metadata_storage_and_retrieval(memory):
+async def test_metadata_storage_and_retrieval(memory, request_context):
     """
     Test that user-defined metadata is preserved.
     Metadata allows arbitrary key-value data to be stored with facts.
@@ -692,7 +710,8 @@ async def test_metadata_storage_and_retrieval(memory):
             bank_id=bank_id,
             content="The product launch is scheduled for March 1st.",
             context="planning meeting",
-            event_date=datetime(2024, 1, 15, tzinfo=timezone.utc)
+            event_date=datetime(2024, 1, 15, tzinfo=timezone.utc),
+            request_context=request_context,
         )
 
         assert len(unit_ids) > 0, "Should create memory units"
@@ -703,7 +722,8 @@ async def test_metadata_storage_and_retrieval(memory):
             query="When is the product launch?",
             budget=Budget.LOW,
             max_tokens=500,
-            fact_type=["world"]
+            fact_type=["world"],
+            request_context=request_context,
         )
 
         assert len(result.results) > 0, "Should recall stored facts"
@@ -712,7 +732,7 @@ async def test_metadata_storage_and_retrieval(memory):
         print(f"  (Note: Metadata support depends on API implementation)")
 
     finally:
-        await memory.delete_bank(bank_id)
+        await memory.delete_bank(bank_id, request_context=request_context)
 
 
 # ============================================================
@@ -720,7 +740,7 @@ async def test_metadata_storage_and_retrieval(memory):
 # ============================================================
 
 @pytest.mark.asyncio
-async def test_empty_batch(memory):
+async def test_empty_batch(memory, request_context):
     """
     Test that empty batch is handled gracefully without errors.
     """
@@ -730,7 +750,8 @@ async def test_empty_batch(memory):
         # Attempt to store empty batch
         unit_ids = await memory.retain_batch_async(
             bank_id=bank_id,
-            contents=[]
+            contents=[],
+            request_context=request_context,
         )
 
         # Should return empty list or handle gracefully
@@ -741,11 +762,11 @@ async def test_empty_batch(memory):
 
     finally:
         # Clean up (though nothing should be stored)
-        await memory.delete_bank(bank_id)
+        await memory.delete_bank(bank_id, request_context=request_context)
 
 
 @pytest.mark.asyncio
-async def test_single_item_batch(memory):
+async def test_single_item_batch(memory, request_context):
     """
     Test that batch with one item works correctly.
     """
@@ -761,7 +782,8 @@ async def test_single_item_batch(memory):
                     "context": "deployment log",
                     "event_date": datetime(2024, 1, 15, tzinfo=timezone.utc)
                 }
-            ]
+            ],
+            request_context=request_context,
         )
 
         assert len(unit_ids) == 1, "Should return one list of unit IDs"
@@ -770,11 +792,11 @@ async def test_single_item_batch(memory):
         print(f"✓ Single-item batch created {len(unit_ids[0])} units")
 
     finally:
-        await memory.delete_bank(bank_id)
+        await memory.delete_bank(bank_id, request_context=request_context)
 
 
 @pytest.mark.asyncio
-async def test_mixed_content_batch(memory):
+async def test_mixed_content_batch(memory, request_context):
     """
     Test batch with varying content sizes (short and long).
     """
@@ -798,7 +820,8 @@ async def test_mixed_content_batch(memory):
                 {"content": short_content, "context": "onboarding"},
                 {"content": long_content, "context": "performance review"},
                 {"content": "Charlie is on vacation this week.", "context": "team status"}
-            ]
+            ],
+            request_context=request_context,
         )
 
         # All items should be processed
@@ -813,11 +836,11 @@ async def test_mixed_content_batch(memory):
         print(f"  Long content: {long_units} units")
 
     finally:
-        await memory.delete_bank(bank_id)
+        await memory.delete_bank(bank_id, request_context=request_context)
 
 
 @pytest.mark.asyncio
-async def test_batch_with_missing_optional_fields(memory):
+async def test_batch_with_missing_optional_fields(memory, request_context):
     """
     Test that batch handles items with missing optional fields.
     """
@@ -842,7 +865,8 @@ async def test_batch_with_missing_optional_fields(memory):
                     "context": "code review",
                     # No event_date
                 }
-            ]
+            ],
+            request_context=request_context,
         )
 
         # All items should be processed successfully
@@ -852,7 +876,7 @@ async def test_batch_with_missing_optional_fields(memory):
         print(f"✓ Batch with mixed optional fields created {total_units} total units")
 
     finally:
-        await memory.delete_bank(bank_id)
+        await memory.delete_bank(bank_id, request_context=request_context)
 
 
 # ============================================================
@@ -860,7 +884,7 @@ async def test_batch_with_missing_optional_fields(memory):
 # ============================================================
 
 @pytest.mark.asyncio
-async def test_single_batch_multiple_documents(memory):
+async def test_single_batch_multiple_documents(memory, request_context):
     """
     Test storing multiple distinct documents in a single batch call.
     Each should be tracked separately.
@@ -876,21 +900,24 @@ async def test_single_batch_multiple_documents(memory):
             bank_id=bank_id,
             content="Alice's resume: 10 years Python experience, worked at Google.",
             context="resume review",
-            document_id="resume_alice"
+            document_id="resume_alice",
+            request_context=request_context,
         )
 
         doc2_units = await memory.retain_async(
             bank_id=bank_id,
             content="Bob's resume: 5 years JavaScript experience, worked at Meta.",
             context="resume review",
-            document_id="resume_bob"
+            document_id="resume_bob",
+            request_context=request_context,
         )
 
         doc3_units = await memory.retain_async(
             bank_id=bank_id,
             content="Charlie's resume: 8 years Go experience, worked at Amazon.",
             context="resume review",
-            document_id="resume_charlie"
+            document_id="resume_charlie",
+            request_context=request_context,
         )
 
         # All documents should be stored
@@ -907,17 +934,18 @@ async def test_single_batch_multiple_documents(memory):
             query="Who worked at Google?",
             budget=Budget.MID,
             max_tokens=1000,
-            fact_type=["world"]
+            fact_type=["world"],
+            request_context=request_context,
         )
 
         assert len(result.results) > 0, "Should find facts about Alice"
 
     finally:
-        await memory.delete_bank(bank_id)
+        await memory.delete_bank(bank_id, request_context=request_context)
 
 
 @pytest.mark.asyncio
-async def test_document_upsert_behavior(memory):
+async def test_document_upsert_behavior(memory, request_context):
     """
     Test that upserting a document replaces the old content.
     """
@@ -930,7 +958,8 @@ async def test_document_upsert_behavior(memory):
             bank_id=bank_id,
             content="Project is in planning phase. Alice is the lead.",
             context="status update v1",
-            document_id=document_id
+            document_id=document_id,
+            request_context=request_context,
         )
 
         assert len(v1_units) > 0, "Should create units for v1"
@@ -940,7 +969,8 @@ async def test_document_upsert_behavior(memory):
             bank_id=bank_id,
             content="Project is in development phase. Bob has joined as co-lead.",
             context="status update v2",
-            document_id=document_id
+            document_id=document_id,
+            request_context=request_context,
         )
 
         assert len(v2_units) > 0, "Should create units for v2"
@@ -951,7 +981,8 @@ async def test_document_upsert_behavior(memory):
             query="What is the project status?",
             budget=Budget.MID,
             max_tokens=1000,
-            fact_type=["world"]
+            fact_type=["world"],
+            request_context=request_context,
         )
 
         assert len(result.results) > 0, "Should recall facts"
@@ -959,7 +990,7 @@ async def test_document_upsert_behavior(memory):
         print(f"✓ Document upsert created v1: {len(v1_units)} units, v2: {len(v2_units)} units")
 
     finally:
-        await memory.delete_bank(bank_id)
+        await memory.delete_bank(bank_id, request_context=request_context)
 
 
 # ============================================================
@@ -967,7 +998,7 @@ async def test_document_upsert_behavior(memory):
 # ============================================================
 
 @pytest.mark.asyncio
-async def test_chunk_fact_mapping(memory):
+async def test_chunk_fact_mapping(memory, request_context):
     """
     Test that facts correctly reference their source chunks via chunk_id.
     """
@@ -990,7 +1021,8 @@ async def test_chunk_fact_mapping(memory):
             bank_id=bank_id,
             content=content,
             context="technical documentation",
-            document_id=document_id
+            document_id=document_id,
+            request_context=request_context,
         )
 
         assert len(unit_ids) > 0, "Should create memory units"
@@ -1003,7 +1035,8 @@ async def test_chunk_fact_mapping(memory):
             max_tokens=1000,
             fact_type=["world"],
             include_chunks=True,
-            max_chunk_tokens=8192
+            max_chunk_tokens=8192,
+            request_context=request_context,
         )
 
         assert len(result.results) > 0, "Should recall facts"
@@ -1026,11 +1059,11 @@ async def test_chunk_fact_mapping(memory):
             print(f"  Returned {len(result.chunks)} chunks matching fact references")
 
     finally:
-        await memory.delete_bank(bank_id)
+        await memory.delete_bank(bank_id, request_context=request_context)
 
 
 @pytest.mark.asyncio
-async def test_chunk_ordering_preservation(memory):
+async def test_chunk_ordering_preservation(memory, request_context):
     """
     Test that chunk_index reflects the correct order within a document.
     """
@@ -1070,7 +1103,8 @@ async def test_chunk_ordering_preservation(memory):
             bank_id=bank_id,
             content=content,
             context="multi-section document",
-            document_id=document_id
+            document_id=document_id,
+            request_context=request_context,
         )
 
         assert len(unit_ids) > 0, "Should create units"
@@ -1083,7 +1117,8 @@ async def test_chunk_ordering_preservation(memory):
             max_tokens=2000,
             fact_type=["world"],
             include_chunks=True,
-            max_chunk_tokens=8192
+            max_chunk_tokens=8192,
+            request_context=request_context,
         )
 
         if result.chunks:
@@ -1103,11 +1138,11 @@ async def test_chunk_ordering_preservation(memory):
             print("✓ Content stored (may have created single chunk or no chunks returned)")
 
     finally:
-        await memory.delete_bank(bank_id)
+        await memory.delete_bank(bank_id, request_context=request_context)
 
 
 @pytest.mark.asyncio
-async def test_chunks_truncation_behavior(memory):
+async def test_chunks_truncation_behavior(memory, request_context):
     """
     Test that when chunks exceed max_chunk_tokens, truncation is indicated.
     """
@@ -1165,7 +1200,8 @@ async def test_chunks_truncation_behavior(memory):
             bank_id=bank_id,
             content=large_content,
             context="large document test",
-            document_id=document_id
+            document_id=document_id,
+            request_context=request_context,
         )
 
         assert len(unit_ids) > 0, "Should create units"
@@ -1178,7 +1214,8 @@ async def test_chunks_truncation_behavior(memory):
             max_tokens=1000,
             fact_type=["world"],
             include_chunks=True,
-            max_chunk_tokens=500  # Small limit to test truncation
+            max_chunk_tokens=500,  # Small limit to test truncation
+            request_context=request_context,
         )
 
         if result.chunks:
@@ -1198,7 +1235,7 @@ async def test_chunks_truncation_behavior(memory):
             print("✓ No chunks returned (may be under token limit)")
 
     finally:
-        await memory.delete_bank(bank_id)
+        await memory.delete_bank(bank_id, request_context=request_context)
 
 
 # ============================================================
@@ -1206,7 +1243,7 @@ async def test_chunks_truncation_behavior(memory):
 # ============================================================
 
 @pytest.mark.asyncio
-async def test_temporal_links_creation(memory):
+async def test_temporal_links_creation(memory, request_context):
     """
     Test that temporal links are created between facts with nearby event dates.
 
@@ -1223,7 +1260,8 @@ async def test_temporal_links_creation(memory):
             bank_id=bank_id,
             content="Alice started working on the authentication module.",
             context="daily standup",
-            event_date=base_date
+            event_date=base_date,
+            request_context=request_context,
         )
 
         # Fact 2 at 2:00 PM same day (4 hours later)
@@ -1231,7 +1269,8 @@ async def test_temporal_links_creation(memory):
             bank_id=bank_id,
             content="Bob reviewed the API design document.",
             context="daily standup",
-            event_date=base_date.replace(hour=14)
+            event_date=base_date.replace(hour=14),
+            request_context=request_context,
         )
 
         # Fact 3 at 9:00 AM next day (23 hours later)
@@ -1239,7 +1278,8 @@ async def test_temporal_links_creation(memory):
             bank_id=bank_id,
             content="Charlie deployed the new database schema.",
             context="daily standup",
-            event_date=base_date.replace(day=16, hour=9)
+            event_date=base_date.replace(day=16, hour=9),
+            request_context=request_context,
         )
 
         assert len(unit_ids_1) > 0 and len(unit_ids_2) > 0 and len(unit_ids_3) > 0
@@ -1278,11 +1318,11 @@ async def test_temporal_links_creation(memory):
             logger.info("Temporal links created successfully with proper weights")
 
     finally:
-        await memory.delete_bank(bank_id)
+        await memory.delete_bank(bank_id, request_context=request_context)
 
 
 @pytest.mark.asyncio
-async def test_semantic_links_creation(memory):
+async def test_semantic_links_creation(memory, request_context):
     """
     Test that semantic links are created between facts with similar content.
 
@@ -1295,21 +1335,24 @@ async def test_semantic_links_creation(memory):
         unit_ids_1 = await memory.retain_async(
             bank_id=bank_id,
             content="Alice is an expert in Python programming and has built many web applications.",
-            context="team skills"
+            context="team skills",
+            request_context=request_context,
         )
 
         # Similar content - should create semantic link
         unit_ids_2 = await memory.retain_async(
             bank_id=bank_id,
             content="Bob is proficient in Python development and specializes in building APIs.",
-            context="team skills"
+            context="team skills",
+            request_context=request_context,
         )
 
         # Different content - less likely to create strong semantic link
         unit_ids_3 = await memory.retain_async(
             bank_id=bank_id,
             content="The quarterly sales meeting is scheduled for next Tuesday at 3 PM.",
-            context="calendar events"
+            context="calendar events",
+            request_context=request_context,
         )
 
         assert len(unit_ids_1) > 0 and len(unit_ids_2) > 0 and len(unit_ids_3) > 0
@@ -1349,11 +1392,11 @@ async def test_semantic_links_creation(memory):
             logger.info("Semantic links created successfully between similar content")
 
     finally:
-        await memory.delete_bank(bank_id)
+        await memory.delete_bank(bank_id, request_context=request_context)
 
 
 @pytest.mark.asyncio
-async def test_entity_links_creation(memory):
+async def test_entity_links_creation(memory, request_context):
     """
     Test that entity links are created between facts that mention the same entities.
 
@@ -1367,28 +1410,32 @@ async def test_entity_links_creation(memory):
         unit_ids_1 = await memory.retain_async(
             bank_id=bank_id,
             content="Alice joined Google as a software engineer in 2020.",
-            context="career history"
+            context="career history",
+            request_context=request_context,
         )
 
         # Mentions same entity (Alice) - should create entity link
         unit_ids_2 = await memory.retain_async(
             bank_id=bank_id,
             content="Alice led the development of the new authentication system.",
-            context="project updates"
+            context="project updates",
+            request_context=request_context,
         )
 
         # Mentions same entity (Google) - should create entity link
         unit_ids_3 = await memory.retain_async(
             bank_id=bank_id,
             content="Google announced new cloud services at their annual conference.",
-            context="tech news"
+            context="tech news",
+            request_context=request_context,
         )
 
         # Different entities - no entity link expected
         unit_ids_4 = await memory.retain_async(
             bank_id=bank_id,
             content="Bob works at Meta on machine learning infrastructure.",
-            context="career history"
+            context="career history",
+            request_context=request_context,
         )
 
         assert len(unit_ids_1) > 0 and len(unit_ids_2) > 0 and len(unit_ids_3) > 0 and len(unit_ids_4) > 0
@@ -1445,11 +1492,11 @@ async def test_entity_links_creation(memory):
             logger.info("Entity links are properly bidirectional")
 
     finally:
-        await memory.delete_bank(bank_id)
+        await memory.delete_bank(bank_id, request_context=request_context)
 
 
 @pytest.mark.asyncio
-async def test_causal_links_creation(memory):
+async def test_causal_links_creation(memory, request_context):
     """
     Test that causal links are created between facts with causal relationships.
 
@@ -1471,7 +1518,8 @@ async def test_causal_links_creation(memory):
         unit_ids = await memory.retain_async(
             bank_id=bank_id,
             content=content,
-            context="project timeline"
+            context="project timeline",
+            request_context=request_context,
         )
 
         assert len(unit_ids) > 0, "Should have created facts"
@@ -1517,11 +1565,11 @@ async def test_causal_links_creation(memory):
         logger.info("Test completed (causal link extraction is LLM-dependent)")
 
     finally:
-        await memory.delete_bank(bank_id)
+        await memory.delete_bank(bank_id, request_context=request_context)
 
 
 @pytest.mark.asyncio
-async def test_all_link_types_together(memory):
+async def test_all_link_types_together(memory, request_context):
     """
     Integration test: Verify all link types can be created in a single retain operation.
 
@@ -1539,7 +1587,8 @@ async def test_all_link_types_together(memory):
             bank_id=bank_id,
             content="Alice completed the Python backend service for the authentication system.",
             context="sprint review",
-            event_date=base_date
+            event_date=base_date,
+            request_context=request_context,
         )
 
         # Fact 2: Related to Alice, similar topic (Python), close in time
@@ -1547,7 +1596,8 @@ async def test_all_link_types_together(memory):
             bank_id=bank_id,
             content="Alice optimized the Python code and improved the authentication performance by 40%.",
             context="sprint review",
-            event_date=base_date.replace(hour=14)  # Same day, 4 hours later
+            event_date=base_date.replace(hour=14),  # Same day, 4 hours later
+            request_context=request_context,
         )
 
         # Fact 3: Related to Alice, different topic but same entity
@@ -1555,7 +1605,8 @@ async def test_all_link_types_together(memory):
             bank_id=bank_id,
             content="Alice presented the security architecture at the team meeting.",
             context="team meeting",
-            event_date=base_date.replace(day=16)  # Next day
+            event_date=base_date.replace(day=16),  # Next day
+            request_context=request_context,
         )
 
         assert len(unit_ids_1) > 0 and len(unit_ids_2) > 0 and len(unit_ids_3) > 0
@@ -1594,11 +1645,11 @@ async def test_all_link_types_together(memory):
             logger.info("All major link types (temporal, semantic, entity) are working correctly")
 
     finally:
-        await memory.delete_bank(bank_id)
+        await memory.delete_bank(bank_id, request_context=request_context)
 
 
 @pytest.mark.asyncio
-async def test_semantic_links_within_same_batch(memory):
+async def test_semantic_links_within_same_batch(memory, request_context):
     """
     Test that semantic links are created between facts retained in the SAME batch.
 
@@ -1617,7 +1668,8 @@ async def test_semantic_links_within_same_batch(memory):
 
         result = await memory.retain_batch_async(
             bank_id=bank_id,
-            contents=contents
+            contents=contents,
+            request_context=request_context,
         )
 
         # Flatten the list of lists
@@ -1652,11 +1704,11 @@ async def test_semantic_links_within_same_batch(memory):
                 logger.info(f"  Semantic link: {str(link['from_unit_id'])[:8]}... -> {str(link['to_unit_id'])[:8]}... (weight: {link['weight']:.3f})")
 
     finally:
-        await memory.delete_bank(bank_id)
+        await memory.delete_bank(bank_id, request_context=request_context)
 
 
 @pytest.mark.asyncio
-async def test_temporal_links_within_same_batch(memory):
+async def test_temporal_links_within_same_batch(memory, request_context):
     """
     Test that temporal links are created between facts retained in the SAME batch.
 
@@ -1689,7 +1741,8 @@ async def test_temporal_links_within_same_batch(memory):
 
         result = await memory.retain_batch_async(
             bank_id=bank_id,
-            contents=contents
+            contents=contents,
+            request_context=request_context,
         )
 
         # Flatten the list of lists
@@ -1724,4 +1777,72 @@ async def test_temporal_links_within_same_batch(memory):
                 logger.info(f"  Temporal link: {str(link['from_unit_id'])[:8]}... -> {str(link['to_unit_id'])[:8]}... (weight: {link['weight']:.3f})")
 
     finally:
-        await memory.delete_bank(bank_id)
+        await memory.delete_bank(bank_id, request_context=request_context)
+
+
+@pytest.mark.asyncio
+async def test_user_provided_entities(memory, request_context):
+    """
+    Test that user-provided entities are merged with auto-extracted entities.
+
+    This tests the feature added in PR #91 where users can provide entities
+    via the 'entities' field in the retain request. These should be combined
+    with LLM-extracted entities, with case-insensitive deduplication.
+    """
+    bank_id = f"test_user_entities_{datetime.now(timezone.utc).timestamp()}"
+
+    try:
+        # Store content with user-provided entities
+        # The content mentions "Alice" which LLM might extract,
+        # but we also provide "ProjectX" and "ACME Corp" which may not be in the text
+        contents = [
+            {
+                "content": "Alice completed the quarterly report.",
+                "context": "work update",
+                "entities": [
+                    {"text": "ProjectX", "type": "PROJECT"},
+                    {"text": "ACME Corp", "type": "ORG"},
+                    {"text": "Alice"},  # May also be extracted by LLM (dedup test)
+                ],
+            }
+        ]
+
+        result = await memory.retain_batch_async(
+            bank_id=bank_id,
+            contents=contents,
+            request_context=request_context,
+        )
+
+        # Flatten the list of lists
+        unit_ids = [uid for sublist in result for uid in sublist]
+        assert len(unit_ids) > 0, "Should have created at least one fact"
+
+        logger.info(f"Created {len(unit_ids)} facts with user-provided entities")
+
+        # Query entity links to verify user-provided entities were stored
+        async with memory._pool.acquire() as conn:
+            # Get all entities linked to our facts via the unit_entities junction table
+            entity_rows = await conn.fetch(
+                """
+                SELECT DISTINCT e.canonical_name
+                FROM entities e
+                JOIN unit_entities ue ON e.id = ue.entity_id
+                WHERE ue.unit_id::text = ANY($1)
+                """,
+                unit_ids
+            )
+
+            entity_names = {row['canonical_name'].lower() for row in entity_rows}
+            logger.info(f"Found entities linked to facts: {[row['canonical_name'] for row in entity_rows]}")
+
+            # Verify user-provided entities are present
+            assert "projectx" in entity_names, "User-provided entity 'ProjectX' should be linked"
+            assert "acme corp" in entity_names, "User-provided entity 'ACME Corp' should be linked"
+
+            # Alice should be present (either from LLM extraction or user-provided)
+            assert "alice" in entity_names, "Entity 'Alice' should be linked"
+
+            logger.info("✓ User-provided entities successfully merged with extracted entities")
+
+    finally:
+        await memory.delete_bank(bank_id, request_context=request_context)

hindsight-api/tests/test_schema_isolation.py

@@ -0,0 +1,406 @@
+"""
+Tests for multi-tenant schema isolation.
+
+Verifies that concurrent retain operations from different tenants
+are properly isolated in their respective PostgreSQL schemas.
+"""
+
+import asyncio
+import uuid
+
+import pytest
+import pytest_asyncio
+
+from hindsight_api.extensions import RequestContext, TenantContext, TenantExtension
+from hindsight_api.engine.memory_engine import _current_schema, fq_table
+from hindsight_api.migrations import run_migrations
+
+
+class MultiSchemaTestTenantExtension(TenantExtension):
+    """
+    Test tenant extension that maps API keys to schema names.
+
+    API keys are in format: "key-{schema_name}"
+    Provisions schemas on first access using run_migrations(schema=name).
+    """
+
+    def __init__(self, config: dict):
+        super().__init__(config)
+        self.db_url = config.get("db_url")
+        # Pre-configured valid schemas for test
+        self.valid_schemas = config.get("valid_schemas", set())
+        # Track provisioned schemas
+        self._provisioned: set[str] = set()
+
+    async def authenticate(self, context: RequestContext) -> TenantContext:
+        if not context.api_key:
+            from hindsight_api.extensions import AuthenticationError
+
+            raise AuthenticationError("API key required")
+
+        # Parse schema from API key (format: "key-{schema}")
+        if context.api_key.startswith("key-"):
+            schema = context.api_key[4:]  # Remove "key-" prefix
+            if schema in self.valid_schemas:
+                # Provision schema on first access
+                if schema not in self._provisioned and self.db_url:
+                    run_migrations(self.db_url, schema=schema)
+                    self._provisioned.add(schema)
+                return TenantContext(schema_name=schema)
+
+        from hindsight_api.extensions import AuthenticationError
+
+        raise AuthenticationError(f"Unknown API key: {context.api_key}")
+
+
+async def drop_schema(conn, schema_name: str) -> None:
+    """Drop a schema and all its contents."""
+    await conn.execute(f'DROP SCHEMA IF EXISTS "{schema_name}" CASCADE')
+
+
+async def count_memories_in_schema(conn, schema_name: str, bank_id: str) -> int:
+    """Count memory units in a specific schema for a bank."""
+    result = await conn.fetchval(
+        f'SELECT COUNT(*) FROM "{schema_name}".memory_units WHERE bank_id = $1',
+        bank_id,
+    )
+    return result or 0
+
+
+async def get_memory_texts_in_schema(conn, schema_name: str, bank_id: str) -> list[str]:
+    """Get all memory texts in a specific schema for a bank."""
+    rows = await conn.fetch(
+        f'SELECT text FROM "{schema_name}".memory_units WHERE bank_id = $1 ORDER BY text',
+        bank_id,
+    )
+    return [row["text"] for row in rows]
+
+
+class TestSchemaIsolation:
+    """Tests for multi-tenant schema isolation."""
+
+    @pytest.mark.asyncio
+    async def test_concurrent_inserts_isolated_by_schema(self, memory, pg0_db_url):
+        """
+        Multiple concurrent database operations from different tenants
+        should store data in their respective schemas without cross-contamination.
+
+        Uses run_migrations(schema=x) to provision schemas like a real extension.
+        """
+        import asyncpg
+
+        # Test schemas
+        schemas = ["tenant_alpha", "tenant_beta", "tenant_gamma"]
+        bank_id = f"test-isolation-{uuid.uuid4().hex[:8]}"
+
+        # Clean up any existing schemas
+        conn = await asyncpg.connect(pg0_db_url)
+        try:
+            for schema in schemas:
+                await drop_schema(conn, schema)
+        finally:
+            await conn.close()
+
+        # Configure tenant extension that provisions schemas via run_migrations
+        tenant_ext = MultiSchemaTestTenantExtension({
+            "db_url": pg0_db_url,
+            "valid_schemas": set(schemas),
+        })
+        memory._tenant_extension = tenant_ext
+
+        # Define concurrent insert tasks for each tenant
+        async def insert_for_tenant(schema_name: str, content_prefix: str):
+            """Insert memories for a specific tenant using schema context."""
+            # Authenticate to set the schema context
+            tenant_request = RequestContext(api_key=f"key-{schema_name}")
+            await memory._authenticate_tenant(tenant_request)
+
+            # Now fq_table will use the correct schema
+            pool = await memory._get_pool()
+            from hindsight_api.engine.db_utils import acquire_with_retry
+
+            async with acquire_with_retry(pool) as conn:
+                # Insert 3 memories for this tenant
+                for i in range(3):
+                    await conn.execute(
+                        f"""
+                        INSERT INTO {fq_table('memory_units')} (bank_id, text, event_date, fact_type)
+                        VALUES ($1, $2, now(), 'world')
+                        """,
+                        bank_id,
+                        f"MARKER_{content_prefix}_DOC{i}: Memory for {schema_name}",
+                    )
+
+        # Run concurrent inserts for all tenants
+        await asyncio.gather(
+            insert_for_tenant("tenant_alpha", "ALPHA"),
+            insert_for_tenant("tenant_beta", "BETA"),
+            insert_for_tenant("tenant_gamma", "GAMMA"),
+        )
+
+        # Verify isolation - each schema should only have its own data
+        conn = await asyncpg.connect(pg0_db_url)
+        try:
+            for schema in schemas:
+                texts = await get_memory_texts_in_schema(conn, schema, bank_id)
+                prefix = schema.replace("tenant_", "").upper()
+
+                # Should have exactly 3 memories
+                assert len(texts) == 3, f"Schema {schema} should have 3 memories, got {len(texts)}"
+
+                # All texts should contain the schema's marker
+                for text in texts:
+                    assert f"MARKER_{prefix}" in text, (
+                        f"Memory in {schema} missing its marker: {text}"
+                    )
+
+                # Should NOT contain other tenants' markers
+                other_prefixes = ["ALPHA", "BETA", "GAMMA"]
+                other_prefixes.remove(prefix)
+                for other in other_prefixes:
+                    for text in texts:
+                        assert f"MARKER_{other}" not in text, (
+                            f"Cross-contamination! Schema {schema} has {other}'s marker: {text}"
+                        )
+
+        finally:
+            # Cleanup
+            for schema in schemas:
+                await drop_schema(conn, schema)
+            await conn.close()
+
+        # Reset tenant extension
+        memory._tenant_extension = None
+        _current_schema.set("public")
+
+    @pytest.mark.asyncio
+    async def test_schema_context_isolation_in_concurrent_tasks(self, pg0_db_url):
+        """
+        Verify that _current_schema contextvar is properly isolated
+        between concurrent async tasks.
+        """
+        results = {}
+        errors = []
+
+        async def check_schema_context(schema_name: str, delay: float):
+            """Set schema context, wait, then verify it's still correct."""
+            try:
+                # Set the schema
+                _current_schema.set(schema_name)
+
+                # Small delay to allow interleaving
+                await asyncio.sleep(delay)
+
+                # Verify schema is still correct
+                current = _current_schema.get()
+                if current != schema_name:
+                    errors.append(f"Expected {schema_name}, got {current}")
+
+                # Verify fq_table uses correct schema
+                table = fq_table("memory_units")
+                expected = f"{schema_name}.memory_units"
+                if table != expected:
+                    errors.append(f"Expected {expected}, got {table}")
+
+                results[schema_name] = current
+
+            except Exception as e:
+                errors.append(f"Error in {schema_name}: {e}")
+
+        # Run many concurrent tasks with different schemas
+        tasks = []
+        for i in range(10):
+            for schema in ["schema_a", "schema_b", "schema_c"]:
+                # Vary delays to create interleaving
+                delay = 0.01 * (i % 3)
+                tasks.append(check_schema_context(f"{schema}_{i}", delay))
+
+        await asyncio.gather(*tasks)
+
+        # No errors should have occurred
+        assert not errors, f"Schema context isolation errors: {errors}"
+
+    @pytest.mark.asyncio
+    async def test_list_memories_respects_schema(self, memory, pg0_db_url):
+        """
+        list_memory_units should only return memories from the current schema.
+
+        Uses run_migrations(schema=x) to provision schemas.
+        """
+        import asyncpg
+
+        schemas = ["tenant_list_a", "tenant_list_b"]
+        bank_id = f"test-list-{uuid.uuid4().hex[:8]}"
+
+        # Clean up any existing schemas and provision via migrations
+        conn = await asyncpg.connect(pg0_db_url)
+        try:
+            for schema in schemas:
+                await drop_schema(conn, schema)
+        finally:
+            await conn.close()
+
+        # Provision schemas using run_migrations
+        for schema in schemas:
+            run_migrations(pg0_db_url, schema=schema)
+
+        # Insert test data directly into each schema
+        conn = await asyncpg.connect(pg0_db_url)
+        try:
+            for schema in schemas:
+                await conn.execute(
+                    f"""
+                    INSERT INTO "{schema}".memory_units (bank_id, text, event_date, fact_type)
+                    VALUES ($1, $2, now(), 'world')
+                    """,
+                    bank_id,
+                    f"Direct insert for {schema}",
+                )
+        finally:
+            await conn.close()
+
+        # Configure tenant extension
+        tenant_ext = MultiSchemaTestTenantExtension({
+            "db_url": pg0_db_url,
+            "valid_schemas": set(schemas),
+        })
+        memory._tenant_extension = tenant_ext
+
+        try:
+            # Query as tenant_list_a - should only see tenant_list_a's data
+            tenant_a_request = RequestContext(api_key="key-tenant_list_a")
+            await memory._authenticate_tenant(tenant_a_request)
+
+            result_a = await memory.list_memory_units(bank_id=bank_id, request_context=tenant_a_request)
+            texts_a = [item["text"] for item in result_a.get("items", [])]
+
+            assert len(texts_a) == 1, f"Expected 1 memory for tenant_list_a, got {len(texts_a)}"
+            assert "tenant_list_a" in texts_a[0], f"Wrong content: {texts_a[0]}"
+
+            # Query as tenant_list_b - should only see tenant_list_b's data
+            tenant_b_request = RequestContext(api_key="key-tenant_list_b")
+            await memory._authenticate_tenant(tenant_b_request)
+
+            result_b = await memory.list_memory_units(bank_id=bank_id, request_context=tenant_b_request)
+            texts_b = [item["text"] for item in result_b.get("items", [])]
+
+            assert len(texts_b) == 1, f"Expected 1 memory for tenant_list_b, got {len(texts_b)}"
+            assert "tenant_list_b" in texts_b[0], f"Wrong content: {texts_b[0]}"
+
+        finally:
+            # Cleanup
+            conn = await asyncpg.connect(pg0_db_url)
+            try:
+                for schema in schemas:
+                    await drop_schema(conn, schema)
+            finally:
+                await conn.close()
+
+            memory._tenant_extension = None
+            _current_schema.set("public")
+
+    @pytest.mark.asyncio
+    async def test_high_concurrency_schema_isolation(self, memory, pg0_db_url):
+        """
+        Stress test: Many concurrent operations across multiple schemas
+        should maintain perfect isolation.
+
+        Uses run_migrations(schema=x) to provision schemas like a real extension.
+        """
+        import asyncpg
+
+        # Create more schemas for stress test
+        num_schemas = 5
+        ops_per_schema = 10
+        schemas = [f"stress_tenant_{i}" for i in range(num_schemas)]
+        bank_id = f"test-stress-{uuid.uuid4().hex[:8]}"
+
+        # Clean up any existing schemas first
+        conn = await asyncpg.connect(pg0_db_url)
+        try:
+            for schema in schemas:
+                await drop_schema(conn, schema)
+        finally:
+            await conn.close()
+
+        # Provision schemas using run_migrations
+        for schema in schemas:
+            run_migrations(pg0_db_url, schema=schema)
+
+        # Configure tenant extension (schemas already provisioned)
+        tenant_ext = MultiSchemaTestTenantExtension({
+            "db_url": pg0_db_url,
+            "valid_schemas": set(schemas),
+        })
+        # Mark schemas as already provisioned so extension doesn't re-run migrations
+        tenant_ext._provisioned = set(schemas)
+        memory._tenant_extension = tenant_ext
+
+        errors = []
+
+        async def insert_one(schema: str, item_id: int):
+            """Single insert operation for tracking."""
+            try:
+                # Authenticate to set the schema context
+                tenant_request = RequestContext(api_key=f"key-{schema}")
+                await memory._authenticate_tenant(tenant_request)
+
+                # Insert using fq_table
+                pool = await memory._get_pool()
+                from hindsight_api.engine.db_utils import acquire_with_retry
+
+                async with acquire_with_retry(pool) as conn:
+                    await conn.execute(
+                        f"""
+                        INSERT INTO {fq_table('memory_units')} (bank_id, text, event_date, fact_type)
+                        VALUES ($1, $2, now(), 'world')
+                        """,
+                        bank_id,
+                        f"STRESS_MARKER_{schema}_ITEM{item_id}: Memory for {schema}",
+                    )
+            except Exception as e:
+                errors.append(f"Insert error for {schema}: {e}")
+
+        # Run many concurrent operations
+        tasks = []
+        for i in range(ops_per_schema):
+            for schema in schemas:
+                tasks.append(insert_one(schema, i))
+
+        await asyncio.gather(*tasks)
+
+        # Check for errors during insert
+        assert not errors, f"Errors during insert: {errors}"
+
+        # Verify no cross-contamination
+        conn = await asyncpg.connect(pg0_db_url)
+        try:
+            for schema in schemas:
+                texts = await get_memory_texts_in_schema(conn, schema, bank_id)
+
+                # Should have exactly ops_per_schema memories
+                assert len(texts) == ops_per_schema, (
+                    f"Schema {schema} should have {ops_per_schema} memories, got {len(texts)}"
+                )
+
+                # All memories should reference this schema only
+                for text in texts:
+                    # Check it contains our schema marker
+                    assert f"STRESS_MARKER_{schema}" in text, (
+                        f"Memory in {schema} doesn't contain schema marker: {text}"
+                    )
+
+                    # Check it doesn't contain other schema markers
+                    for other_schema in schemas:
+                        if other_schema != schema:
+                            assert f"STRESS_MARKER_{other_schema}" not in text, (
+                                f"Cross-contamination! {schema} has {other_schema}'s data: {text}"
+                            )
+        finally:
+            # Cleanup
+            for schema in schemas:
+                await drop_schema(conn, schema)
+            await conn.close()
+
+        memory._tenant_extension = None
+        _current_schema.set("public")

hindsight-api/tests/test_search_trace.py

@@ -3,12 +3,12 @@ Test search tracing functionality.
 """
 import pytest
 from hindsight_api.engine.memory_engine import Budget
-from hindsight_api import SearchTrace
+from hindsight_api import SearchTrace, RequestContext
 from datetime import datetime, timezone
 
 
 @pytest.mark.asyncio
-async def test_search_with_trace(memory):
+async def test_search_with_trace(memory, request_context):
     """Test that search with enable_trace=True returns a valid SearchTrace."""
     # Generate a unique agent ID for this test
     bank_id = f"test_trace_{datetime.now(timezone.utc).timestamp()}"
@@ -20,16 +20,19 @@ async def test_search_with_trace(memory):
             bank_id=bank_id,
             content="Alice works at Google in Mountain View",
             context="test context",
+            request_context=request_context,
         )
         await memory.retain_async(
             bank_id=bank_id,
             content="Bob also works at Google but in New York",
             context="test context",
+            request_context=request_context,
         )
         await memory.retain_async(
             bank_id=bank_id,
             content="Charlie founded a startup called TechCorp",
             context="test context",
+            request_context=request_context,
         )
 
         # Search with tracing enabled
@@ -40,6 +43,7 @@ async def test_search_with_trace(memory):
             budget=Budget.LOW, # 20,
             max_tokens=512,
             enable_trace=True,
+            request_context=request_context,
         )
 
         # Verify results
@@ -102,11 +106,11 @@ async def test_search_with_trace(memory):
 
     finally:
         # Cleanup
-        await memory.delete_bank(bank_id)
+        await memory.delete_bank(bank_id, request_context=request_context)
 
 
 @pytest.mark.asyncio
-async def test_search_without_trace(memory):
+async def test_search_without_trace(memory, request_context):
     """Test that search with enable_trace=False returns None for trace."""
     bank_id = f"test_no_trace_{datetime.now(timezone.utc).timestamp()}"
 
@@ -117,6 +121,7 @@ async def test_search_without_trace(memory):
             bank_id=bank_id,
             content="Test memory without trace",
             context="test",
+            request_context=request_context,
         )
 
         # Search without tracing
@@ -127,6 +132,7 @@ async def test_search_without_trace(memory):
             budget=Budget.LOW, # 10,
             max_tokens=512,
             enable_trace=False,
+            request_context=request_context,
         )
 
         # Verify trace is None
@@ -137,4 +143,4 @@ async def test_search_without_trace(memory):
 
     finally:
         # Cleanup
-        await memory.delete_bank(bank_id)
+        await memory.delete_bank(bank_id, request_context=request_context)

hindsight-api/tests/test_sql_schema_safety.py

@@ -0,0 +1,136 @@
+"""
+Safety tests to ensure all SQL queries use fully-qualified table names.
+
+This prevents cross-tenant data access by ensuring every table reference
+includes the schema prefix (e.g., public.memory_units instead of just memory_units).
+"""
+
+import re
+from pathlib import Path
+
+import pytest
+
+# All tables that MUST be schema-qualified in SQL queries
+TABLES = [
+    "memory_units",
+    "memory_links",
+    "unit_entities",
+    "entities",
+    "entity_cooccurrences",
+    "banks",
+    "documents",
+    "chunks",
+    "async_operations",
+]
+
+# Files to scan for SQL queries
+SCAN_PATHS = [
+    "hindsight_api/engine",
+    "hindsight_api/api",
+]
+
+# Files to exclude (e.g., migrations, tests)
+EXCLUDE_PATTERNS = [
+    "alembic",
+    "__pycache__",
+    "test_",
+]
+
+
+def get_python_files() -> list[Path]:
+    """Get all Python files to scan."""
+    root = Path(__file__).parent.parent
+    files = []
+    for scan_path in SCAN_PATHS:
+        path = root / scan_path
+        if path.exists():
+            for py_file in path.rglob("*.py"):
+                # Check exclusions
+                if any(excl in str(py_file) for excl in EXCLUDE_PATTERNS):
+                    continue
+                files.append(py_file)
+    return files
+
+
+def find_unqualified_table_refs(content: str, filename: str) -> list[tuple[int, str, str]]:
+    """
+    Find SQL statements with unqualified table references.
+
+    Returns list of (line_number, table_name, line_content).
+    """
+    violations = []
+
+    # Patterns that indicate SQL context
+    sql_keywords = r"(?:FROM|JOIN|INTO|UPDATE|DELETE\s+FROM)\s+"
+
+    # Additional SQL indicators to confirm this is actually SQL, not prose
+    sql_indicators = re.compile(
+        r"(SELECT|INSERT|DELETE|UPDATE|CREATE|ALTER|DROP|WHERE|SET|VALUES|"
+        r'f"""|f\'\'\'|""".*SELECT|\'\'\'.*SELECT)',
+        re.IGNORECASE,
+    )
+
+    lines = content.split("\n")
+    for line_num, line in enumerate(lines, 1):
+        # Skip comments and strings that are clearly not SQL
+        stripped = line.strip()
+        if stripped.startswith("#"):
+            continue
+
+        for table in TABLES:
+            # Pattern: SQL keyword followed by unqualified table name
+            # Should match: FROM memory_units, JOIN memory_units, INTO memory_units
+            # Should NOT match: FROM public.memory_units, FROM {schema}.memory_units
+            # Should NOT match: fq_table("memory_units")
+
+            # Check for unqualified table after SQL keyword
+            pattern = rf"{sql_keywords}{table}(?:\s|$|,|\))"
+
+            if re.search(pattern, line, re.IGNORECASE):
+                # Check if it's actually qualified (has schema prefix)
+                qualified_pattern = rf"\.\s*{table}(?:\s|$|,|\))"
+                fq_table_pattern = rf'fq_table\s*\(\s*["\']?{table}'
+
+                if not re.search(qualified_pattern, line) and not re.search(
+                    fq_table_pattern, line
+                ):
+                    # Additional check: line must have SQL indicators
+                    # This avoids false positives in docstrings like "split into chunks"
+                    if sql_indicators.search(line):
+                        violations.append((line_num, table, stripped))
+
+    return violations
+
+
+class TestSQLSchemaSafety:
+    """Ensure all SQL uses schema-qualified table names."""
+
+    def test_no_unqualified_table_references(self):
+        """All SQL queries must use fq_table() or schema.table format."""
+        all_violations = []
+
+        for py_file in get_python_files():
+            content = py_file.read_text()
+            violations = find_unqualified_table_refs(content, py_file.name)
+
+            for line_num, table, line in violations:
+                all_violations.append(
+                    f"{py_file.relative_to(py_file.parent.parent)}:{line_num} - "
+                    f"unqualified '{table}': {line[:80]}..."
+                )
+
+        if all_violations:
+            msg = (
+                f"Found {len(all_violations)} unqualified table references!\n"
+                "These could cause cross-tenant data access.\n"
+                "Use fq_table('table_name') for all table references.\n\n"
+                + "\n".join(all_violations[:20])  # Show first 20
+            )
+            if len(all_violations) > 20:
+                msg += f"\n... and {len(all_violations) - 20} more"
+            pytest.fail(msg)
+
+    def test_tables_list_is_complete(self):
+        """Verify we're checking for all tables (sanity check)."""
+        # This is a sanity check - if you add a new table, add it to TABLES
+        assert len(TABLES) >= 9, "Update TABLES list if you added new tables"

hindsight-api/tests/test_temporal_ranges.py

@@ -3,16 +3,17 @@ import asyncio
 from datetime import datetime, timezone, timedelta
 import pytest
 from hindsight_api.engine.memory_engine import Budget
+from hindsight_api import RequestContext
 
 
 @pytest.mark.asyncio
-async def test_temporal_ranges_are_written(memory):
+async def test_temporal_ranges_are_written(memory, request_context):
     """Test that occurred_start, occurred_end, and mentioned_at are actually written to database."""
     bank_id = "test_temporal_ranges"
 
     # Clean up any existing data
     try:
-        await memory.delete_bank(bank_id)
+        await memory.delete_bank(bank_id, request_context=request_context)
     except Exception:
         pass
 
@@ -23,7 +24,8 @@ async def test_temporal_ranges_are_written(memory):
     await memory.retain_async(
         bank_id=bank_id,
         content=text1,
-        event_date=conversation_date
+        event_date=conversation_date,
+        request_context=request_context,
     )
 
     # Test 2: Period event (month range)
@@ -32,7 +34,8 @@ async def test_temporal_ranges_are_written(memory):
     await memory.retain_async(
         bank_id=bank_id,
         content=text2,
-        event_date=conversation_date
+        event_date=conversation_date,
+        request_context=request_context,
     )
 
     # Give it a moment for async processing
@@ -114,7 +117,8 @@ async def test_temporal_ranges_are_written(memory):
         query="pottery workshop",
         fact_type=["world", "experience"],
         budget=Budget.LOW,
-        max_tokens=4096
+        max_tokens=4096,
+        request_context=request_context,
     )
 
     print(f"Found {len(search_result.results)} search results")
@@ -132,4 +136,4 @@ async def test_temporal_ranges_are_written(memory):
             print("⚠ Temporal fields not yet populated in search results (known issue)")
 
     # Clean up
-    await memory.delete_bank(bank_id)
+    await memory.delete_bank(bank_id, request_context=request_context)

hindsight-api/tests/test_think.py

@@ -4,10 +4,11 @@ Test think function for opinion generation and consistency.
 import pytest
 from datetime import datetime, timezone
 from hindsight_api.engine.memory_engine import Budget
+from hindsight_api import RequestContext
 
 
 @pytest.mark.asyncio
-async def test_think_opinion_consistency(memory):
+async def test_think_opinion_consistency(memory, request_context):
     """
     Test that think function:
     1. Generates an opinion
@@ -23,14 +24,16 @@ async def test_think_opinion_consistency(memory):
             bank_id=bank_id,
             content="Alice is a software engineer who has worked on 5 major projects. She always delivers on time and writes clean, well-documented code.",
             context="performance review",
-            event_date=datetime(2024, 1, 15, tzinfo=timezone.utc)
+            event_date=datetime(2024, 1, 15, tzinfo=timezone.utc),
+            request_context=request_context,
         )
 
         await memory.retain_async(
             bank_id=bank_id,
             content="Bob recently joined the team. He missed his first deadline and his code had many bugs.",
             context="performance review",
-            event_date=datetime(2024, 2, 1, tzinfo=timezone.utc)
+            event_date=datetime(2024, 2, 1, tzinfo=timezone.utc),
+            request_context=request_context,
         )
 
         # First think call - should generate opinions
@@ -39,6 +42,7 @@ async def test_think_opinion_consistency(memory):
             bank_id=bank_id,
             query=query,
             budget=Budget.LOW,
+            request_context=request_context,
         )
 
         print(f"\n=== First Think Call ===")
@@ -82,6 +86,7 @@ async def test_think_opinion_consistency(memory):
             bank_id=bank_id,
             query=query,
             budget=Budget.LOW,
+            request_context=request_context,
         )
 
         print(f"\n=== Second Think Call ===")
@@ -122,13 +127,13 @@ async def test_think_opinion_consistency(memory):
     finally:
         # Clean up agent data
         try:
-            await memory.delete_bank(bank_id)
+            await memory.delete_bank(bank_id, request_context=request_context)
         except Exception as e:
             print(f"Warning: Error during cleanup: {e}")
 
 
 @pytest.mark.asyncio
-async def test_think_without_prior_context(memory):
+async def test_think_without_prior_context(memory, request_context):
     """
     Test that think function handles queries when there's no relevant context.
     """
@@ -139,6 +144,7 @@ async def test_think_without_prior_context(memory):
         bank_id=bank_id,
         query="What is the capital of France?",
         budget=Budget.LOW,
+        request_context=request_context,
     )
 
     print(f"\n=== Think Without Context ===")