feat: extensions

docker/standalone/Dockerfile

@@ -2,16 +2,19 @@
 # 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)
+#   PRELOAD_ML_MODELS=true/false - Pre-download ML models during build (default: 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
 
 ARG INCLUDE_API=true
 ARG INCLUDE_CP=true
+ARG PRELOAD_ML_MODELS=true
 
 # =============================================================================
 # Stage: API Builder
@@ -159,14 +162,17 @@ 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)
+ARG PRELOAD_ML_MODELS
+RUN if [ "$PRELOAD_ML_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')"; \
+    else echo "Skipping ML model preload"; fi
 
 EXPOSE 8888
 
@@ -277,14 +283,17 @@ 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)
+ARG PRELOAD_ML_MODELS
+RUN if [ "$PRELOAD_ML_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')"; \
+    else echo "Skipping ML model preload"; fi
 
 EXPOSE 8888 9999
 

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/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/http.py

@@ -12,7 +12,7 @@ from contextlib import asynccontextmanager
 from datetime import datetime
 from typing import Any
 
-from fastapi import FastAPI, HTTPException, Query
+from fastapi import Depends, FastAPI, Header, HTTPException, Query
 
 
 def _parse_metadata(metadata: Any) -> dict[str, Any]:
@@ -33,9 +33,11 @@ from pydantic import BaseModel, ConfigDict, Field
 
 from hindsight_api import MemoryEngine
 from hindsight_api.engine.db_utils import acquire_with_retry
-from hindsight_api.engine.memory_engine import Budget
+from hindsight_api.engine.memory_engine import Budget, fq_table
 from hindsight_api.engine.response_models import VALID_RECALL_FACT_TYPES
+from hindsight_api.extensions import HttpExtension, load_extension
 from hindsight_api.metrics import create_metrics_collector, get_metrics_collector, initialize_metrics
+from hindsight_api.models import RequestContext
 
 logger = logging.getLogger(__name__)
 
@@ -337,7 +339,7 @@ class RetainResponse(BaseModel):
     success: bool
     bank_id: str
     items_count: int
-    async_: bool = Field(
+    is_async: bool = Field(
         alias="async", serialization_alias="async", description="Whether the operation was processed asynchronously"
     )
 
@@ -706,7 +708,11 @@ class DeleteResponse(BaseModel):
     deleted_count: int | None = None
 
 
-def create_app(memory: MemoryEngine, initialize_memory: bool = True) -> FastAPI:
+def create_app(
+    memory: MemoryEngine,
+    initialize_memory: bool = True,
+    http_extension: HttpExtension | None = None,
+) -> FastAPI:
     """
     Create and configure the FastAPI application.
 
@@ -714,6 +720,8 @@ def create_app(memory: MemoryEngine, initialize_memory: bool = True) -> FastAPI:
         memory: MemoryEngine instance (already initialized with required parameters).
                 Migrations are controlled by the MemoryEngine's run_migrations parameter.
         initialize_memory: Whether to initialize memory system on startup (default: True)
+        http_extension: Optional HTTP extension to mount custom endpoints under /extension/.
+                       If None, attempts to load from HINDSIGHT_API_HTTP_EXTENSION env var.
 
     Returns:
         Configured FastAPI application
@@ -723,6 +731,11 @@ def create_app(memory: MemoryEngine, initialize_memory: bool = True) -> FastAPI:
         In that case, you should call memory.initialize() manually before starting the server
         and memory.close() when shutting down.
     """
+    # Load HTTP extension from environment if not provided
+    if http_extension is None:
+        http_extension = load_extension("HTTP", HttpExtension)
+        if http_extension:
+            logging.info(f"Loaded HTTP extension: {http_extension.__class__.__name__}")
 
     @asynccontextmanager
     async def lifespan(app: FastAPI):
@@ -746,8 +759,18 @@ def create_app(memory: MemoryEngine, initialize_memory: bool = True) -> FastAPI:
             await memory.initialize()
             logging.info("Memory system initialized")
 
+        # Call HTTP extension startup hook
+        if http_extension:
+            await http_extension.on_startup()
+            logging.info("HTTP extension started")
+
         yield
 
+        # Call HTTP extension shutdown hook
+        if http_extension:
+            await http_extension.on_shutdown()
+            logging.info("HTTP extension stopped")
+
         # Shutdown: Cleanup memory system
         await memory.close()
         logging.info("Memory system closed")
@@ -775,12 +798,36 @@ def create_app(memory: MemoryEngine, initialize_memory: bool = True) -> FastAPI:
     # Register all routes
     _register_routes(app)
 
+    # Mount HTTP extension router if available
+    if http_extension:
+        extension_router = http_extension.get_router(memory)
+        app.include_router(extension_router, prefix="/ext", tags=["Extension"])
+        logging.info("HTTP extension router mounted at /ext/")
+
     return app
 
 
 def _register_routes(app: FastAPI):
     """Register all API routes on the given app instance."""
 
+    def get_request_context(authorization: str | None = Header(default=None)) -> RequestContext:
+        """
+        Extract request context from Authorization header.
+
+        Supports:
+        - Bearer token: "Bearer <api_key>"
+        - Direct API key: "<api_key>"
+
+        Returns RequestContext with extracted API key (may be None if no auth header).
+        """
+        api_key = None
+        if authorization:
+            if authorization.lower().startswith("bearer "):
+                api_key = authorization[7:].strip()
+            else:
+                api_key = authorization.strip()
+        return RequestContext(api_key=api_key)
+
     @app.get(
         "/health",
         summary="Health check endpoint",
@@ -821,10 +868,12 @@ def _register_routes(app: FastAPI):
         operation_id="get_graph",
         tags=["Memory"],
     )
-    async def api_graph(bank_id: str, type: str | None = None):
+    async def api_graph(
+        bank_id: str, type: str | None = None, request_context: RequestContext = Depends(get_request_context)
+    ):
         """Get graph data from database, filtered by bank_id and optionally by type."""
         try:
-            data = await app.state.memory.get_graph_data(bank_id, type)
+            data = await app.state.memory.get_graph_data(bank_id, type, request_context=request_context)
             return data
         except Exception as e:
             import traceback
@@ -841,7 +890,14 @@ def _register_routes(app: FastAPI):
         operation_id="list_memories",
         tags=["Memory"],
     )
-    async def api_list(bank_id: str, type: str | None = None, q: str | None = None, limit: int = 100, offset: int = 0):
+    async def api_list(
+        bank_id: str,
+        type: str | None = None,
+        q: str | None = None,
+        limit: int = 100,
+        offset: int = 0,
+        request_context: RequestContext = Depends(get_request_context),
+    ):
         """
         List memory units for table view with optional full-text search.
 
@@ -857,7 +913,12 @@ def _register_routes(app: FastAPI):
         """
         try:
             data = await app.state.memory.list_memory_units(
-                bank_id=bank_id, fact_type=type, search_query=q, limit=limit, offset=offset
+                bank_id=bank_id,
+                fact_type=type,
+                search_query=q,
+                limit=limit,
+                offset=offset,
+                request_context=request_context,
             )
             return data
         except Exception as e:
@@ -880,7 +941,9 @@ def _register_routes(app: FastAPI):
         operation_id="recall_memories",
         tags=["Memory"],
     )
-    async def api_recall(bank_id: str, request: RecallRequest):
+    async def api_recall(
+        bank_id: str, request: RecallRequest, request_context: RequestContext = Depends(get_request_context)
+    ):
         """Run a recall and return results with trace."""
         metrics = get_metrics_collector()
 
@@ -923,6 +986,7 @@ def _register_routes(app: FastAPI):
                     max_entity_tokens=max_entity_tokens,
                     include_chunks=include_chunks,
                     max_chunk_tokens=max_chunk_tokens,
+                    request_context=request_context,
                 )
 
             # Convert core MemoryFact objects to API RecallResult objects (excluding internal metrics)
@@ -995,14 +1059,20 @@ def _register_routes(app: FastAPI):
         operation_id="reflect",
         tags=["Memory"],
     )
-    async def api_reflect(bank_id: str, request: ReflectRequest):
+    async def api_reflect(
+        bank_id: str, request: ReflectRequest, request_context: RequestContext = Depends(get_request_context)
+    ):
         metrics = get_metrics_collector()
 
         try:
             # Use the memory system's reflect_async method (record metrics)
             with metrics.record_operation("reflect", bank_id=bank_id, budget=request.budget.value):
                 core_result = await app.state.memory.reflect_async(
-                    bank_id=bank_id, query=request.query, budget=request.budget, context=request.context
+                    bank_id=bank_id,
+                    query=request.query,
+                    budget=request.budget,
+                    context=request.context,
+                    request_context=request_context,
                 )
 
             # Convert core MemoryFact objects to API ReflectFact objects if facts are requested
@@ -1041,10 +1111,10 @@ def _register_routes(app: FastAPI):
         operation_id="list_banks",
         tags=["Banks"],
     )
-    async def api_list_banks():
+    async def api_list_banks(request_context: RequestContext = Depends(get_request_context)):
         """Get list of all banks with their profiles."""
         try:
-            banks = await app.state.memory.list_banks()
+            banks = await app.state.memory.list_banks(request_context=request_context)
             return BankListResponse(banks=banks)
         except Exception as e:
             import traceback
@@ -1067,9 +1137,9 @@ def _register_routes(app: FastAPI):
             async with acquire_with_retry(pool) as conn:
                 # Get node counts by fact_type
                 node_stats = await conn.fetch(
-                    """
+                    f"""
                     SELECT fact_type, COUNT(*) as count
-                    FROM memory_units
+                    FROM {fq_table("memory_units")}
                     WHERE bank_id = $1
                     GROUP BY fact_type
                     """,
@@ -1078,10 +1148,10 @@ def _register_routes(app: FastAPI):
 
                 # Get link counts by link_type
                 link_stats = await conn.fetch(
-                    """
+                    f"""
                     SELECT ml.link_type, COUNT(*) as count
-                    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
                     GROUP BY ml.link_type
                     """,
@@ -1090,10 +1160,10 @@ def _register_routes(app: FastAPI):
 
                 # Get link counts by fact_type (from nodes)
                 link_fact_type_stats = await conn.fetch(
-                    """
+                    f"""
                     SELECT mu.fact_type, COUNT(*) as count
-                    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
                     GROUP BY mu.fact_type
                     """,
@@ -1102,10 +1172,10 @@ def _register_routes(app: FastAPI):
 
                 # Get link counts by fact_type AND link_type
                 link_breakdown_stats = await conn.fetch(
-                    """
+                    f"""
                     SELECT mu.fact_type, ml.link_type, COUNT(*) as count
-                    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
                     GROUP BY mu.fact_type, ml.link_type
                     """,
@@ -1114,9 +1184,9 @@ def _register_routes(app: FastAPI):
 
                 # Get pending and failed operations counts
                 ops_stats = await conn.fetch(
-                    """
+                    f"""
                     SELECT status, COUNT(*) as count
-                    FROM async_operations
+                    FROM {fq_table("async_operations")}
                     WHERE bank_id = $1
                     GROUP BY status
                     """,
@@ -1128,9 +1198,9 @@ def _register_routes(app: FastAPI):
 
                 # Get document count
                 doc_count_result = await conn.fetchrow(
-                    """
+                    f"""
                     SELECT COUNT(*) as count
-                    FROM documents
+                    FROM {fq_table("documents")}
                     WHERE bank_id = $1
                     """,
                     bank_id,
@@ -1184,11 +1254,13 @@ def _register_routes(app: FastAPI):
         tags=["Entities"],
     )
     async def api_list_entities(
-        bank_id: str, limit: int = Query(default=100, description="Maximum number of entities to return")
+        bank_id: str,
+        limit: int = Query(default=100, description="Maximum number of entities to return"),
+        request_context: RequestContext = Depends(get_request_context),
     ):
         """List entities for a memory bank."""
         try:
-            entities = await app.state.memory.list_entities(bank_id, limit=limit)
+            entities = await app.state.memory.list_entities(bank_id, limit=limit, request_context=request_context)
             return EntityListResponse(items=[EntityListItem(**e) for e in entities])
         except Exception as e:
             import traceback
@@ -1205,37 +1277,26 @@ def _register_routes(app: FastAPI):
         operation_id="get_entity",
         tags=["Entities"],
     )
-    async def api_get_entity(bank_id: str, entity_id: str):
+    async def api_get_entity(
+        bank_id: str, entity_id: str, request_context: RequestContext = Depends(get_request_context)
+    ):
         """Get entity details with observations."""
         try:
-            # First get the entity metadata
-            pool = await app.state.memory._get_pool()
-            async with acquire_with_retry(pool) as conn:
-                entity_row = await conn.fetchrow(
-                    """
-                    SELECT id, canonical_name, mention_count, first_seen, last_seen, metadata
-                    FROM entities
-                    WHERE bank_id = $1 AND id = $2
-                    """,
-                    bank_id,
-                    uuid.UUID(entity_id),
-                )
+            entity = await app.state.memory.get_entity(bank_id, entity_id, request_context=request_context)
 
-            if not entity_row:
+            if entity is None:
                 raise HTTPException(status_code=404, detail=f"Entity {entity_id} not found")
 
-            # Get observations for the entity
-            observations = await app.state.memory.get_entity_observations(bank_id, entity_id, limit=20)
-
             return EntityDetailResponse(
-                id=str(entity_row["id"]),
-                canonical_name=entity_row["canonical_name"],
-                mention_count=entity_row["mention_count"],
-                first_seen=entity_row["first_seen"].isoformat() if entity_row["first_seen"] else None,
-                last_seen=entity_row["last_seen"].isoformat() if entity_row["last_seen"] else None,
-                metadata=_parse_metadata(entity_row["metadata"]),
+                id=entity["id"],
+                canonical_name=entity["canonical_name"],
+                mention_count=entity["mention_count"],
+                first_seen=entity["first_seen"],
+                last_seen=entity["last_seen"],
+                metadata=_parse_metadata(entity["metadata"]),
                 observations=[
-                    EntityObservationResponse(text=obs.text, mentioned_at=obs.mentioned_at) for obs in observations
+                    EntityObservationResponse(text=obs.text, mentioned_at=obs.mentioned_at)
+                    for obs in entity["observations"]
                 ],
             )
         except HTTPException:
@@ -1255,42 +1316,40 @@ def _register_routes(app: FastAPI):
         operation_id="regenerate_entity_observations",
         tags=["Entities"],
     )
-    async def api_regenerate_entity_observations(bank_id: str, entity_id: str):
+    async def api_regenerate_entity_observations(
+        bank_id: str,
+        entity_id: str,
+        request_context: RequestContext = Depends(get_request_context),
+    ):
         """Regenerate observations for an entity."""
         try:
-            # First get the entity metadata
-            pool = await app.state.memory._get_pool()
-            async with acquire_with_retry(pool) as conn:
-                entity_row = await conn.fetchrow(
-                    """
-                    SELECT id, canonical_name, mention_count, first_seen, last_seen, metadata
-                    FROM entities
-                    WHERE bank_id = $1 AND id = $2
-                    """,
-                    bank_id,
-                    uuid.UUID(entity_id),
-                )
+            # Get the entity to verify it exists and get canonical_name
+            entity = await app.state.memory.get_entity(bank_id, entity_id, request_context=request_context)
 
-            if not entity_row:
+            if entity is None:
                 raise HTTPException(status_code=404, detail=f"Entity {entity_id} not found")
 
             # Regenerate observations
             await app.state.memory.regenerate_entity_observations(
-                bank_id=bank_id, entity_id=entity_id, entity_name=entity_row["canonical_name"]
+                bank_id=bank_id,
+                entity_id=entity_id,
+                entity_name=entity["canonical_name"],
+                request_context=request_context,
             )
 
-            # Get updated observations
-            observations = await app.state.memory.get_entity_observations(bank_id, entity_id, limit=20)
+            # Get updated entity with new observations
+            entity = await app.state.memory.get_entity(bank_id, entity_id, request_context=request_context)
 
             return EntityDetailResponse(
-                id=str(entity_row["id"]),
-                canonical_name=entity_row["canonical_name"],
-                mention_count=entity_row["mention_count"],
-                first_seen=entity_row["first_seen"].isoformat() if entity_row["first_seen"] else None,
-                last_seen=entity_row["last_seen"].isoformat() if entity_row["last_seen"] else None,
-                metadata=_parse_metadata(entity_row["metadata"]),
+                id=entity["id"],
+                canonical_name=entity["canonical_name"],
+                mention_count=entity["mention_count"],
+                first_seen=entity["first_seen"],
+                last_seen=entity["last_seen"],
+                metadata=_parse_metadata(entity["metadata"]),
                 observations=[
-                    EntityObservationResponse(text=obs.text, mentioned_at=obs.mentioned_at) for obs in observations
+                    EntityObservationResponse(text=obs.text, mentioned_at=obs.mentioned_at)
+                    for obs in entity["observations"]
                 ],
             )
         except HTTPException:
@@ -1310,7 +1369,13 @@ def _register_routes(app: FastAPI):
         operation_id="list_documents",
         tags=["Documents"],
     )
-    async def api_list_documents(bank_id: str, q: str | None = None, limit: int = 100, offset: int = 0):
+    async def api_list_documents(
+        bank_id: str,
+        q: str | None = None,
+        limit: int = 100,
+        offset: int = 0,
+        request_context: RequestContext = Depends(get_request_context),
+    ):
         """
         List documents for a memory bank with optional search.
 
@@ -1321,7 +1386,9 @@ def _register_routes(app: FastAPI):
             offset: Offset for pagination (default: 0)
         """
         try:
-            data = await app.state.memory.list_documents(bank_id=bank_id, search_query=q, limit=limit, offset=offset)
+            data = await app.state.memory.list_documents(
+                bank_id=bank_id, search_query=q, limit=limit, offset=offset, request_context=request_context
+            )
             return data
         except Exception as e:
             import traceback
@@ -1338,7 +1405,9 @@ def _register_routes(app: FastAPI):
         operation_id="get_document",
         tags=["Documents"],
     )
-    async def api_get_document(bank_id: str, document_id: str):
+    async def api_get_document(
+        bank_id: str, document_id: str, request_context: RequestContext = Depends(get_request_context)
+    ):
         """
         Get a specific document with its original text.
 
@@ -1347,7 +1416,7 @@ def _register_routes(app: FastAPI):
             document_id: Document ID (from path)
         """
         try:
-            document = await app.state.memory.get_document(document_id, bank_id)
+            document = await app.state.memory.get_document(document_id, bank_id, request_context=request_context)
             if not document:
                 raise HTTPException(status_code=404, detail="Document not found")
             return document
@@ -1368,7 +1437,7 @@ def _register_routes(app: FastAPI):
         operation_id="get_chunk",
         tags=["Documents"],
     )
-    async def api_get_chunk(chunk_id: str):
+    async def api_get_chunk(chunk_id: str, request_context: RequestContext = Depends(get_request_context)):
         """
         Get a specific chunk with its text.
 
@@ -1376,7 +1445,7 @@ def _register_routes(app: FastAPI):
             chunk_id: Chunk ID (from path, format: bank_id_document_id_chunk_index)
         """
         try:
-            chunk = await app.state.memory.get_chunk(chunk_id)
+            chunk = await app.state.memory.get_chunk(chunk_id, request_context=request_context)
             if not chunk:
                 raise HTTPException(status_code=404, detail="Chunk not found")
             return chunk
@@ -1401,7 +1470,9 @@ def _register_routes(app: FastAPI):
         operation_id="delete_document",
         tags=["Documents"],
     )
-    async def api_delete_document(bank_id: str, document_id: str):
+    async def api_delete_document(
+        bank_id: str, document_id: str, request_context: RequestContext = Depends(get_request_context)
+    ):
         """
         Delete a document and all its associated memory units and links.
 
@@ -1410,7 +1481,7 @@ def _register_routes(app: FastAPI):
             document_id: Document ID to delete (from path)
         """
         try:
-            result = await app.state.memory.delete_document(document_id, bank_id)
+            result = await app.state.memory.delete_document(document_id, bank_id, request_context=request_context)
 
             if result["document_deleted"] == 0:
                 raise HTTPException(status_code=404, detail="Document not found")
@@ -1437,45 +1508,14 @@ def _register_routes(app: FastAPI):
         operation_id="list_operations",
         tags=["Operations"],
     )
-    async def api_list_operations(bank_id: str):
+    async def api_list_operations(bank_id: str, request_context: RequestContext = Depends(get_request_context)):
         """List all async operations (pending and failed) for a memory bank."""
         try:
-            pool = await app.state.memory._get_pool()
-            async with acquire_with_retry(pool) as conn:
-                operations = await conn.fetch(
-                    """
-                    SELECT operation_id, bank_id, operation_type, created_at, status, error_message, result_metadata
-                    FROM async_operations
-                    WHERE bank_id = $1
-                    ORDER BY created_at DESC
-                    """,
-                    bank_id,
-                )
-
-                def parse_metadata(metadata):
-                    """Parse result_metadata which may be a string or dict."""
-                    if metadata is None:
-                        return {}
-                    if isinstance(metadata, str):
-                        return json.loads(metadata)
-                    return metadata
-
-                return {
-                    "bank_id": bank_id,
-                    "operations": [
-                        {
-                            "id": str(row["operation_id"]),
-                            "task_type": row["operation_type"],
-                            "items_count": parse_metadata(row["result_metadata"]).get("items_count", 0),
-                            "document_id": parse_metadata(row["result_metadata"]).get("document_id"),
-                            "created_at": row["created_at"].isoformat(),
-                            "status": row["status"],
-                            "error_message": row["error_message"],
-                        }
-                        for row in operations
-                    ],
-                }
-
+            operations = await app.state.memory.list_operations(bank_id, request_context=request_context)
+            return {
+                "bank_id": bank_id,
+                "operations": operations,
+            }
         except Exception as e:
             import traceback
 
@@ -1490,39 +1530,21 @@ def _register_routes(app: FastAPI):
         operation_id="cancel_operation",
         tags=["Operations"],
     )
-    async def api_cancel_operation(bank_id: str, operation_id: str):
+    async def api_cancel_operation(
+        bank_id: str, operation_id: str, request_context: RequestContext = Depends(get_request_context)
+    ):
         """Cancel a pending async operation."""
         try:
             # Validate UUID format
             try:
-                op_uuid = uuid.UUID(operation_id)
+                uuid.UUID(operation_id)
             except ValueError:
                 raise HTTPException(status_code=400, detail=f"Invalid operation_id format: {operation_id}")
 
-            pool = await app.state.memory._get_pool()
-            async with acquire_with_retry(pool) as conn:
-                # Check if operation exists and belongs to this memory bank
-                result = await conn.fetchrow(
-                    "SELECT bank_id FROM async_operations WHERE operation_id = $1 AND bank_id = $2", op_uuid, bank_id
-                )
-
-                if not result:
-                    raise HTTPException(
-                        status_code=404, detail=f"Operation {operation_id} not found for memory bank {bank_id}"
-                    )
-
-                # Delete the operation
-                await conn.execute("DELETE FROM async_operations WHERE operation_id = $1", op_uuid)
-
-                return {
-                    "success": True,
-                    "message": f"Operation {operation_id} cancelled",
-                    "operation_id": operation_id,
-                    "bank_id": bank_id,
-                }
-
-        except HTTPException:
-            raise
+            result = await app.state.memory.cancel_operation(bank_id, operation_id, request_context=request_context)
+            return result
+        except ValueError as e:
+            raise HTTPException(status_code=404, detail=str(e))
         except Exception as e:
             import traceback
 
@@ -1538,10 +1560,10 @@ def _register_routes(app: FastAPI):
         operation_id="get_bank_profile",
         tags=["Banks"],
     )
-    async def api_get_bank_profile(bank_id: str):
+    async def api_get_bank_profile(bank_id: str, request_context: RequestContext = Depends(get_request_context)):
         """Get memory bank profile (disposition + background)."""
         try:
-            profile = await app.state.memory.get_bank_profile(bank_id)
+            profile = await app.state.memory.get_bank_profile(bank_id, request_context=request_context)
             # Convert DispositionTraits object to dict for Pydantic
             disposition_dict = (
                 profile["disposition"].model_dump()
@@ -1569,14 +1591,18 @@ def _register_routes(app: FastAPI):
         operation_id="update_bank_disposition",
         tags=["Banks"],
     )
-    async def api_update_bank_disposition(bank_id: str, request: UpdateDispositionRequest):
+    async def api_update_bank_disposition(
+        bank_id: str, request: UpdateDispositionRequest, request_context: RequestContext = Depends(get_request_context)
+    ):
         """Update bank disposition traits."""
         try:
             # Update disposition
-            await app.state.memory.update_bank_disposition(bank_id, request.disposition.model_dump())
+            await app.state.memory.update_bank_disposition(
+                bank_id, request.disposition.model_dump(), request_context=request_context
+            )
 
             # Get updated profile
-            profile = await app.state.memory.get_bank_profile(bank_id)
+            profile = await app.state.memory.get_bank_profile(bank_id, request_context=request_context)
             disposition_dict = (
                 profile["disposition"].model_dump()
                 if hasattr(profile["disposition"], "model_dump")
@@ -1603,11 +1629,13 @@ def _register_routes(app: FastAPI):
         operation_id="add_bank_background",
         tags=["Banks"],
     )
-    async def api_add_bank_background(bank_id: str, request: AddBackgroundRequest):
+    async def api_add_bank_background(
+        bank_id: str, request: AddBackgroundRequest, request_context: RequestContext = Depends(get_request_context)
+    ):
         """Add or merge bank background information. Optionally infer disposition traits."""
         try:
             result = await app.state.memory.merge_bank_background(
-                bank_id, request.content, update_disposition=request.update_disposition
+                bank_id, request.content, update_disposition=request.update_disposition, request_context=request_context
             )
 
             response = BackgroundResponse(background=result["background"])
@@ -1630,51 +1658,31 @@ def _register_routes(app: FastAPI):
         operation_id="create_or_update_bank",
         tags=["Banks"],
     )
-    async def api_create_or_update_bank(bank_id: str, request: CreateBankRequest):
+    async def api_create_or_update_bank(
+        bank_id: str, request: CreateBankRequest, request_context: RequestContext = Depends(get_request_context)
+    ):
         """Create or update an agent with disposition and background."""
         try:
-            # Get existing profile or create with defaults
-            profile = await app.state.memory.get_bank_profile(bank_id)
+            # Ensure bank exists by getting profile (auto-creates with defaults)
+            await app.state.memory.get_bank_profile(bank_id, request_context=request_context)
 
-            # Update name if provided
-            if request.name is not None:
-                pool = await app.state.memory._get_pool()
-                async with acquire_with_retry(pool) as conn:
-                    await conn.execute(
-                        """
-                        UPDATE banks
-                        SET name = $2,
-                            updated_at = NOW()
-                        WHERE bank_id = $1
-                        """,
-                        bank_id,
-                        request.name,
-                    )
-                profile["name"] = request.name
+            # Update name and/or background if provided
+            if request.name is not None or request.background is not None:
+                await app.state.memory.update_bank(
+                    bank_id,
+                    name=request.name,
+                    background=request.background,
+                    request_context=request_context,
+                )
 
             # Update disposition if provided
             if request.disposition is not None:
-                await app.state.memory.update_bank_disposition(bank_id, request.disposition.model_dump())
-                profile["disposition"] = request.disposition.model_dump()
-
-            # Update background if provided (replace, not merge)
-            if request.background is not None:
-                pool = await app.state.memory._get_pool()
-                async with acquire_with_retry(pool) as conn:
-                    await conn.execute(
-                        """
-                        UPDATE banks
-                        SET background = $2,
-                            updated_at = NOW()
-                        WHERE bank_id = $1
-                        """,
-                        bank_id,
-                        request.background,
-                    )
-                profile["background"] = request.background
+                await app.state.memory.update_bank_disposition(
+                    bank_id, request.disposition.model_dump(), request_context=request_context
+                )
 
             # Get final profile
-            final_profile = await app.state.memory.get_bank_profile(bank_id)
+            final_profile = await app.state.memory.get_bank_profile(bank_id, request_context=request_context)
             disposition_dict = (
                 final_profile["disposition"].model_dump()
                 if hasattr(final_profile["disposition"], "model_dump")
@@ -1702,10 +1710,10 @@ def _register_routes(app: FastAPI):
         operation_id="delete_bank",
         tags=["Banks"],
     )
-    async def api_delete_bank(bank_id: str):
+    async def api_delete_bank(bank_id: str, request_context: RequestContext = Depends(get_request_context)):
         """Delete an entire memory bank and all its data."""
         try:
-            result = await app.state.memory.delete_bank(bank_id)
+            result = await app.state.memory.delete_bank(bank_id, request_context=request_context)
             return DeleteResponse(
                 success=True,
                 message=f"Bank '{bank_id}' and all associated data deleted successfully",
@@ -1745,7 +1753,9 @@ def _register_routes(app: FastAPI):
         operation_id="retain_memories",
         tags=["Memory"],
     )
-    async def api_retain(bank_id: str, request: RetainRequest):
+    async def api_retain(
+        bank_id: str, request: RetainRequest, request_context: RequestContext = Depends(get_request_context)
+    ):
         """Retain memories with optional async processing."""
         metrics = get_metrics_collector()
 
@@ -1766,43 +1776,25 @@ def _register_routes(app: FastAPI):
 
             if request.async_:
                 # Async processing: queue task and return immediately
-                operation_id = uuid.uuid4()
-
-                # Insert operation record into database
-                pool = await app.state.memory._get_pool()
-                async with acquire_with_retry(pool) as conn:
-                    await conn.execute(
-                        """
-                        INSERT INTO async_operations (operation_id, bank_id, operation_type, result_metadata)
-                        VALUES ($1, $2, $3, $4)
-                        """,
-                        operation_id,
-                        bank_id,
-                        "retain",
-                        json.dumps({"items_count": len(contents)}),
-                    )
-
-                # Submit task to background queue
-                await app.state.memory._task_backend.submit_task(
+                result = await app.state.memory.submit_async_retain(bank_id, contents, request_context=request_context)
+                return RetainResponse.model_validate(
                     {
-                        "type": "batch_retain",
-                        "operation_id": str(operation_id),
+                        "success": True,
                         "bank_id": bank_id,
-                        "contents": contents,
+                        "items_count": result["items_count"],
+                        "async": True,
                     }
                 )
-
-                logging.info(
-                    f"Retain task queued for bank_id={bank_id}, {len(contents)} items, operation_id={operation_id}"
-                )
-
-                return RetainResponse(success=True, bank_id=bank_id, items_count=len(contents), async_=True)
             else:
                 # Synchronous processing: wait for completion (record metrics)
                 with metrics.record_operation("retain", bank_id=bank_id):
-                    result = await app.state.memory.retain_batch_async(bank_id=bank_id, contents=contents)
+                    result = await app.state.memory.retain_batch_async(
+                        bank_id=bank_id, contents=contents, request_context=request_context
+                    )
 
-                return RetainResponse(success=True, bank_id=bank_id, items_count=len(contents), async_=False)
+                return RetainResponse.model_validate(
+                    {"success": True, "bank_id": bank_id, "items_count": len(contents), "async": False}
+                )
         except Exception as e:
             import traceback
 
@@ -1832,10 +1824,11 @@ def _register_routes(app: FastAPI):
     async def api_clear_bank_memories(
         bank_id: str,
         type: str | None = Query(None, description="Optional fact type filter (world, experience, opinion)"),
+        request_context: RequestContext = Depends(get_request_context),
     ):
         """Clear memories for a memory bank, optionally filtered by type."""
         try:
-            await app.state.memory.delete_bank(bank_id, fact_type=type)
+            await app.state.memory.delete_bank(bank_id, fact_type=type, request_context=request_context)
 
             return DeleteResponse(success=True)
         except Exception as e:

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()
@@ -67,7 +68,11 @@ def create_mcp_server(memory: MemoryEngine) -> FastMCP:
         """
         try:
             bank_id = get_current_bank_id()
-            await memory.retain_batch_async(bank_id=bank_id, contents=[{"content": content, "context": context}])
+            if bank_id is None:
+                return "Error: No bank_id configured"
+            await memory.retain_batch_async(
+                bank_id=bank_id, contents=[{"content": content, "context": context}], request_context=RequestContext()
+            )
             return "Memory stored successfully"
         except Exception as e:
             logger.error(f"Error storing memory: {e}", exc_info=True)
@@ -90,10 +95,16 @@ def create_mcp_server(memory: MemoryEngine) -> FastMCP:
         """
         try:
             bank_id = get_current_bank_id()
+            if bank_id 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
+                bank_id=bank_id,
+                query=query,
+                fact_type=list(VALID_RECALL_FACT_TYPES),
+                budget=Budget.LOW,
+                request_context=RequestContext(),
             )
 
             results = [
@@ -102,7 +113,7 @@ def create_mcp_server(memory: MemoryEngine) -> FastMCP:
                     "text": fact.text,
                     "type": fact.fact_type,
                     "context": fact.context,
-                    "event_date": fact.event_date,
+                    "occurred_start": fact.occurred_start,
                 }
                 for fact in search_result.results[:max_results]
             ]

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/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,592 @@
+"""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,
+        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.
+            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,
+        request_context: "RequestContext",
+    ) -> dict[str, Any]:
+        """
+        Get graph data for visualization.
+
+        Args:
+            bank_id: The memory bank ID.
+            fact_type: Filter by fact type.
+            request_context: Request context for authentication.
+
+        Returns:
+            Dict with nodes, edges, table_rows, total_units.
+        """
+        ...
+
+    # =========================================================================
+    # 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

@@ -96,7 +96,7 @@ class LLMProvider:
             client_kwargs = {"api_key": self.api_key, "max_retries": 0}
             if self.base_url:
                 client_kwargs["base_url"] = self.base_url
-            self._client = AsyncOpenAI(**client_kwargs)
+            self._client = AsyncOpenAI(**client_kwargs)  # type: ignore[invalid-argument-type] - dict kwargs
             self._gemini_client = None
 
     async def verify_connection(self) -> None:
@@ -467,6 +467,8 @@ class 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")
 
@@ -477,6 +479,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"))
 
@@ -487,6 +493,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/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/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,7 @@ import time
 import uuid
 from datetime import UTC, datetime
 
+from ..memory_engine import fq_table
 from ..search import observation_utils
 from . import embedding_utils
 from .types import EntityLink
@@ -75,8 +76,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 +87,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 +155,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 +194,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 +218,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 +241,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
@@ -29,7 +28,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,7 +42,7 @@ 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,

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

@@ -7,9 +7,33 @@ 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)
+    """
+
+    content: str  # Required
+    context: str
+    event_date: datetime
+    metadata: dict[str, str]
+    document_id: str
+
+
+def _now_utc() -> datetime:
+    """Factory function for default event_date."""
+    return datetime.now(UTC)
+
+
 @dataclass
 class RetainContent:
     """
@@ -20,16 +44,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)
-
 
 @dataclass
 class ChunkMetadata:

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/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)

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/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,110 @@
+"""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
+
+        run_migrations(self._database_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,325 @@
+"""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):
+        self.reason = reason
+        super().__init__(f"Operation validation failed: {reason}")
+
+
+@dataclass
+class ValidationResult:
+    """Result of an operation validation."""
+
+    allowed: bool
+    reason: str | None = None
+
+    @classmethod
+    def accept(cls) -> "ValidationResult":
+        """Create an accepted validation result."""
+        return cls(allowed=True)
+
+    @classmethod
+    def reject(cls, reason: str) -> "ValidationResult":
+        """Create a rejected validation result with a reason."""
+        return cls(allowed=False, reason=reason)
+
+
+# =============================================================================
+# 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

@@ -185,7 +185,7 @@ def main():
         mcp_enabled=config.mcp_enabled,
     )
 
-    uvicorn.run(**uvicorn_config)
+    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/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:

hindsight-api/hindsight_api/models.py

@@ -2,9 +2,24 @@
 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
+
+
 from pgvector.sqlalchemy import Vector
 from sqlalchemy import (
     CheckConstraint,

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:

hindsight-api/pyproject.toml

@@ -92,6 +92,7 @@ dev = [
     "python-dotenv>=1.2.1",
     "filelock>=3.0.0",
     "ruff>=0.8.0",
+    "ty>=0.0.1",
 ]
 
 [tool.ruff]
@@ -121,3 +122,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_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_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_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

@@ -897,7 +897,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 +905,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 +924,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 +956,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 +967,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 +975,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 +985,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 +993,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 +1008,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 +1016,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 +1033,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 +1041,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_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,7 @@ async def test_regenerate_entity_observations(memory):
 
 
 @pytest.mark.asyncio
-async def test_search_with_include_entities(memory):
+async def test_search_with_include_entities(memory, request_context):
     """
     Test that search with include_entities=True returns entity observations.
 
@@ -196,7 +200,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 +214,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 +269,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 +281,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 +309,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 +332,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 +344,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 +383,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 +419,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 +476,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_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,4 @@ 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)

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 ===")

hindsight-cli/src/api.rs

@@ -64,13 +64,24 @@ pub struct ApiClient {
 }
 
 impl ApiClient {
-    pub fn new(base_url: String) -> Result<Self> {
+    pub fn new(base_url: String, api_key: Option<String>) -> Result<Self> {
         let runtime = std::sync::Arc::new(tokio::runtime::Runtime::new()?);
 
-        // Create HTTP client with 2-minute timeout
-        let http_client = reqwest::Client::builder()
-            .timeout(std::time::Duration::from_secs(120))
-            .build()?;
+        // Create HTTP client with 2-minute timeout and optional auth header
+        let mut client_builder = reqwest::Client::builder()
+            .timeout(std::time::Duration::from_secs(120));
+
+        if let Some(key) = api_key {
+            let mut headers = reqwest::header::HeaderMap::new();
+            let auth_value = format!("Bearer {}", key);
+            headers.insert(
+                reqwest::header::AUTHORIZATION,
+                reqwest::header::HeaderValue::from_str(&auth_value)?,
+            );
+            client_builder = client_builder.default_headers(headers);
+        }
+
+        let http_client = client_builder.build()?;
 
         let client = AsyncClient::new_with_client(&base_url, http_client);
         Ok(ApiClient { client, runtime })

hindsight-cli/src/config.rs

@@ -10,6 +10,7 @@ const CONFIG_DIR_NAME: &str = ".hindsight";
 
 pub struct Config {
     pub api_url: String,
+    pub api_key: Option<String>,
     pub source: ConfigSource,
 }
 
@@ -32,22 +33,27 @@ impl std::fmt::Display for ConfigSource {
 
 impl Config {
     /// Load configuration with the following priority:
-    /// 1. Environment variable (HINDSIGHT_API_URL) - highest priority, for overrides
+    /// 1. Environment variable (HINDSIGHT_API_URL, HINDSIGHT_API_KEY) - highest priority, for overrides
     /// 2. Local config file (~/.hindsight/config.toml)
     /// 3. Default (http://localhost:8888)
     pub fn load() -> Result<Self> {
+        // Load API key from environment (highest priority)
+        let env_api_key = env::var("HINDSIGHT_API_KEY").ok();
+
         // 1. Environment variable takes highest priority (for overrides)
         if let Ok(api_url) = env::var("HINDSIGHT_API_URL") {
-            return Self::validate_and_create(api_url, ConfigSource::Environment);
+            return Self::validate_and_create(api_url, env_api_key, ConfigSource::Environment);
         }
 
         // 2. Try local config file
-        if let Some(api_url) = Self::load_from_file()? {
-            return Self::validate_and_create(api_url, ConfigSource::LocalFile);
+        if let Some((api_url, file_api_key)) = Self::load_from_file()? {
+            // Environment api_key takes precedence over file api_key
+            let api_key = env_api_key.or(file_api_key);
+            return Self::validate_and_create(api_url, api_key, ConfigSource::LocalFile);
         }
 
         // 3. Fall back to default
-        Self::validate_and_create(DEFAULT_API_URL.to_string(), ConfigSource::Default)
+        Self::validate_and_create(DEFAULT_API_URL.to_string(), env_api_key, ConfigSource::Default)
     }
 
     /// Legacy method for backwards compatibility
@@ -55,14 +61,14 @@ impl Config {
         Self::load()
     }
 
-    fn validate_and_create(api_url: String, source: ConfigSource) -> Result<Self> {
+    fn validate_and_create(api_url: String, api_key: Option<String>, source: ConfigSource) -> Result<Self> {
         if !api_url.starts_with("http://") && !api_url.starts_with("https://") {
             anyhow::bail!(
                 "Invalid API URL: {}. Must start with http:// or https://",
                 api_url
             );
         }
-        Ok(Config { api_url, source })
+        Ok(Config { api_url, api_key, source })
     }
 
     fn config_dir() -> Option<PathBuf> {
@@ -73,7 +79,7 @@ impl Config {
         Self::config_dir().map(|dir| dir.join(CONFIG_FILE_NAME))
     }
 
-    fn load_from_file() -> Result<Option<String>> {
+    fn load_from_file() -> Result<Option<(String, Option<String>)>> {
         let config_path = match Self::config_file_path() {
             Some(path) => path,
             None => return Ok(None),
@@ -86,23 +92,40 @@ impl Config {
         let content = fs::read_to_string(&config_path)
             .with_context(|| format!("Failed to read config file: {}", config_path.display()))?;
 
-        // Simple TOML parsing for api_url
+        let mut api_url: Option<String> = None;
+        let mut api_key: Option<String> = None;
+
+        // Simple TOML parsing for api_url and api_key
         for line in content.lines() {
             let line = line.trim();
             if line.starts_with("api_url") {
                 if let Some(value) = line.split('=').nth(1) {
                     let value = value.trim().trim_matches('"').trim_matches('\'');
                     if !value.is_empty() {
-                        return Ok(Some(value.to_string()));
+                        api_url = Some(value.to_string());
+                    }
+                }
+            } else if line.starts_with("api_key") {
+                if let Some(value) = line.split('=').nth(1) {
+                    let value = value.trim().trim_matches('"').trim_matches('\'');
+                    if !value.is_empty() {
+                        api_key = Some(value.to_string());
                     }
                 }
             }
         }
 
-        Ok(None)
+        match api_url {
+            Some(url) => Ok(Some((url, api_key))),
+            None => Ok(None),
+        }
     }
 
     pub fn save_api_url(api_url: &str) -> Result<PathBuf> {
+        Self::save_config(api_url, None)
+    }
+
+    pub fn save_config(api_url: &str, api_key: Option<&str>) -> Result<PathBuf> {
         let config_dir = Self::config_dir()
             .ok_or_else(|| anyhow::anyhow!("Could not determine home directory"))?;
 
@@ -113,7 +136,10 @@ impl Config {
         }
 
         let config_path = config_dir.join(CONFIG_FILE_NAME);
-        let content = format!("api_url = \"{}\"\n", api_url);
+        let mut content = format!("api_url = \"{}\"\n", api_url);
+        if let Some(key) = api_key {
+            content.push_str(&format!("api_key = \"{}\"\n", key));
+        }
 
         fs::write(&config_path, content)
             .with_context(|| format!("Failed to write config file: {}", config_path.display()))?;

hindsight-cli/src/main.rs

@@ -94,12 +94,15 @@ enum Commands {
     /// Launch the web-based control plane UI
     Ui,
 
-    /// Configure the CLI (API URL, etc.)
-    #[command(after_help = "Configuration priority:\n  1. Environment variable (HINDSIGHT_API_URL) - highest priority\n  2. Config file (~/.hindsight/config)\n  3. Default (http://localhost:8888)")]
+    /// Configure the CLI (API URL, API key, etc.)
+    #[command(after_help = "Configuration priority:\n  1. Environment variables (HINDSIGHT_API_URL, HINDSIGHT_API_KEY) - highest priority\n  2. Config file (~/.hindsight/config)\n  3. Default (http://localhost:8888)")]
     Configure {
         /// API URL to connect to (interactive prompt if not provided)
         #[arg(long)]
         api_url: Option<String>,
+        /// API key for authentication (sent as Bearer token)
+        #[arg(long)]
+        api_key: Option<String>,
     },
 }
 
@@ -372,8 +375,8 @@ fn run() -> Result<()> {
     let verbose = cli.verbose;
 
     // Handle configure command before loading full config (it doesn't need API client)
-    if let Commands::Configure { api_url } = cli.command {
-        return handle_configure(api_url, output_format);
+    if let Commands::Configure { api_url, api_key } = cli.command {
+        return handle_configure(api_url, api_key, output_format);
     }
 
     // Handle ui command - needs config but not API client
@@ -389,9 +392,10 @@ fn run() -> Result<()> {
     });
 
     let api_url = config.api_url().to_string();
+    let api_key = config.api_key.clone();
 
     // Create API client
-    let client = ApiClient::new(api_url.clone()).unwrap_or_else(|e| {
+    let client = ApiClient::new(api_url.clone(), api_key).unwrap_or_else(|e| {
         errors::handle_api_error(e, &api_url);
     });
 
@@ -476,7 +480,7 @@ fn run() -> Result<()> {
     Ok(())
 }
 
-fn handle_configure(api_url: Option<String>, output_format: OutputFormat) -> Result<()> {
+fn handle_configure(api_url: Option<String>, api_key: Option<String>, output_format: OutputFormat) -> Result<()> {
     // Load current config to show current state
     let current_config = Config::load().ok();
 
@@ -487,6 +491,15 @@ fn handle_configure(api_url: Option<String>, output_format: OutputFormat) -> Res
         // Show current configuration
         if let Some(ref config) = current_config {
             println!("  Current API URL: {}", config.api_url);
+            if let Some(ref key) = config.api_key {
+                // Mask the API key for display
+                let masked = if key.len() > 8 {
+                    format!("{}...{}", &key[..4], &key[key.len()-4..])
+                } else {
+                    "****".to_string()
+                };
+                println!("  Current API Key: {}", masked);
+            }
             println!("  Source: {}", config.source);
             println!();
         }
@@ -511,18 +524,30 @@ fn handle_configure(api_url: Option<String>, output_format: OutputFormat) -> Res
         return Ok(());
     }
 
+    // Use provided api_key, or keep existing one if not provided
+    let new_api_key = api_key.or_else(|| current_config.as_ref().and_then(|c| c.api_key.clone()));
+
     // Save to config file
-    let config_path = Config::save_api_url(&new_api_url)?;
+    let config_path = Config::save_config(&new_api_url, new_api_key.as_deref())?;
 
     if output_format == OutputFormat::Pretty {
         ui::print_success(&format!("Configuration saved to {}", config_path.display()));
         println!();
         println!("  API URL: {}", new_api_url);
+        if let Some(ref key) = new_api_key {
+            let masked = if key.len() > 8 {
+                format!("{}...{}", &key[..4], &key[key.len()-4..])
+            } else {
+                "****".to_string()
+            };
+            println!("  API Key: {}", masked);
+        }
         println!();
-        println!("Note: Environment variable HINDSIGHT_API_URL will override this setting.");
+        println!("Note: Environment variables HINDSIGHT_API_URL and HINDSIGHT_API_KEY will override these settings.");
     } else {
         let result = serde_json::json!({
             "api_url": new_api_url,
+            "api_key_set": new_api_key.is_some(),
             "config_path": config_path.display().to_string(),
         });
         output::print_output(&result, output_format)?;

hindsight-cli/src/utils.rs

@@ -5,7 +5,7 @@ use crate::output::OutputFormat;
 
 /// Get API client from config
 pub fn get_client(config: &Config) -> Result<ApiClient> {
-    ApiClient::new(config.api_url.clone())
+    ApiClient::new(config.api_url.clone(), config.api_key.clone())
         .context("Failed to create API client")
 }
 

hindsight-clients/python/hindsight_client/hindsight_client.py

@@ -44,8 +44,12 @@ class Hindsight:
         ```python
         from hindsight_client import Hindsight
 
+        # Without authentication
         client = Hindsight(base_url="http://localhost:8888")
 
+        # With API key authentication
+        client = Hindsight(base_url="http://localhost:8888", api_key="your-api-key")
+
         # Store a memory
         client.retain(bank_id="alice", content="Alice loves AI")
 
@@ -59,15 +63,16 @@ class Hindsight:
         ```
     """
 
-    def __init__(self, base_url: str, timeout: float = 30.0):
+    def __init__(self, base_url: str, api_key: Optional[str] = None, timeout: float = 30.0):
         """
         Initialize the Hindsight client.
 
         Args:
             base_url: The base URL of the Hindsight API server
+            api_key: Optional API key for authentication (sent as Bearer token)
             timeout: Request timeout in seconds (default: 30.0)
         """
-        config = hindsight_client_api.Configuration(host=base_url)
+        config = hindsight_client_api.Configuration(host=base_url, access_token=api_key)
         self._api_client = hindsight_client_api.ApiClient(config)
         self._api = default_api.DefaultApi(self._api_client)
 

hindsight-clients/typescript/src/index.ts

@@ -5,8 +5,15 @@
  * ```typescript
  * import { HindsightClient } from '@vectorize-io/hindsight-client';
  *
+ * // Without authentication
  * const client = new HindsightClient({ baseUrl: 'http://localhost:8888' });
  *
+ * // With API key authentication
+ * const client = new HindsightClient({
+ *   baseUrl: 'http://localhost:8888',
+ *   apiKey: 'your-api-key'
+ * });
+ *
  * // Retain a memory
  * await client.retain('alice', 'Alice loves AI');
  *
@@ -37,6 +44,10 @@ import type {
 
 export interface HindsightClientOptions {
     baseUrl: string;
+    /**
+     * Optional API key for authentication (sent as Bearer token in Authorization header)
+     */
+    apiKey?: string;
 }
 
 export interface MemoryItemInput {
@@ -54,6 +65,9 @@ export class HindsightClient {
         this.client = createClient(
             createConfig({
                 baseUrl: options.baseUrl,
+                headers: options.apiKey
+                    ? { Authorization: `Bearer ${options.apiKey}` }
+                    : undefined,
             })
         );
     }

hindsight-dev/benchmarks/locomo/locomo_benchmark.py

@@ -3,28 +3,28 @@ LoComo-specific benchmark implementations.
 
 Provides dataset, answer generator, and evaluator for the LoComo benchmark.
 """
-import sys
-from pathlib import Path
 
-from benchmarks.common.benchmark_runner import BenchmarkRunner
-
-import json
-from datetime import datetime, timezone
-from typing import List, Dict, Any, Tuple, Optional
 import asyncio
-import pydantic
-from openai import AsyncOpenAI
+import json
 import os
+import sys
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple
 
-from benchmarks.common.benchmark_runner import BenchmarkDataset, LLMAnswerGenerator, LLMAnswerEvaluator
+import pydantic
 from hindsight_api.engine.llm_wrapper import LLMConfig
+from openai import AsyncOpenAI
+
+from benchmarks.common.benchmark_runner import BenchmarkDataset, BenchmarkRunner, LLMAnswerEvaluator, LLMAnswerGenerator
+
 
 class LoComoDataset(BenchmarkDataset):
     """LoComo dataset implementation."""
 
     def load(self, path: Path, max_items: Optional[int] = None) -> List[Dict[str, Any]]:
         """Load LoComo dataset from JSON file."""
-        with open(path, 'r') as f:
+        with open(path, "r") as f:
             dataset = json.load(f)
 
         if max_items:
@@ -34,7 +34,7 @@ class LoComoDataset(BenchmarkDataset):
 
     def get_item_id(self, item: Dict) -> str:
         """Get sample ID from LoComo item."""
-        return item['sample_id']
+        return item["sample_id"]
 
     def prepare_sessions_for_ingestion(self, item: Dict) -> List[Dict[str, Any]]:
         """
@@ -45,12 +45,12 @@ class LoComoDataset(BenchmarkDataset):
         Returns:
             List of session dicts, each containing 'content', 'context', 'event_date', 'document_id'
         """
-        conv = item['conversation']
-        speaker_a = conv['speaker_a']
-        speaker_b = conv['speaker_b']
+        conv = item["conversation"]
+        speaker_a = conv["speaker_a"]
+        speaker_b = conv["speaker_b"]
 
         # Get all session keys sorted
-        session_keys = sorted([k for k in conv.keys() if k.startswith('session_') and not k.endswith('_date_time')])
+        session_keys = sorted([k for k in conv.keys() if k.startswith("session_") and not k.endswith("_date_time")])
 
         session_items = []
 
@@ -65,12 +65,14 @@ class LoComoDataset(BenchmarkDataset):
             session_date = self._parse_date(conv.get(date_key))
             session_content = json.dumps(session_data)
             document_id = f"{item['sample_id']}_{session_key}"
-            session_items.append({
-                "content": session_content,
-                "context": f"Conversation between {speaker_a} and {speaker_b} ({session_key} of {item['sample_id']})",
-                "event_date": session_date,
-                "document_id": document_id
-            })
+            session_items.append(
+                {
+                    "content": session_content,
+                    "context": f"Conversation between {speaker_a} and {speaker_b} ({session_key} of {item['sample_id']})",
+                    "event_date": session_date,
+                    "document_id": document_id,
+                }
+            )
 
         return session_items
 
@@ -81,7 +83,7 @@ class LoComoDataset(BenchmarkDataset):
         Returns:
             List of QA dicts with 'question', 'answer', 'category'
         """
-        return item['qa']
+        return item["qa"]
 
     def _parse_date(self, date_string: str) -> datetime:
         """Parse LoComo date format to datetime."""
@@ -95,6 +97,7 @@ class LoComoDataset(BenchmarkDataset):
 
 class QuestionAnswer(pydantic.BaseModel):
     """Answer format for LoComo questions."""
+
     answer: str
     reasoning: str
 
@@ -113,7 +116,7 @@ class LoComoAnswerGenerator(LLMAnswerGenerator):
         question: str,
         recall_result: Dict[str, Any],
         question_date: Optional[datetime] = None,
-        question_type: Optional[str] = None
+        question_type: Optional[str] = None,
     ) -> Tuple[str, str, Optional[List[Dict[str, Any]]]]:
         """
         Generate answer from retrieved memories using Groq.
@@ -141,7 +144,7 @@ class LoComoAnswerGenerator(LLMAnswerGenerator):
                 messages=[
                     {
                         "role": "system",
-                        "content": "You are a helpful expert assistant answering questions from lme_experiment users based on the provided context."
+                        "content": "You are a helpful expert assistant answering questions from lme_experiment users based on the provided context.",
                     },
                     {
                         "role": "user",
@@ -165,11 +168,11 @@ Context:
 Question: {question}
 Answer:
 
-"""
-                    }
+""",
+                    },
                 ],
                 response_format=QuestionAnswer,
-                scope="memory"
+                scope="memory",
             )
             return answer_obj.answer, answer_obj.reasoning, None
         except Exception as e:
@@ -183,7 +186,7 @@ class LoComoThinkAnswerGenerator(LLMAnswerGenerator):
     so it doesn't need external search to be performed by the benchmark runner.
     """
 
-    def __init__(self, memory: 'MemoryEngine', agent_id: str, thinking_budget: int = 500):
+    def __init__(self, memory: "MemoryEngine", agent_id: str, thinking_budget: int = 500):
         """Initialize with memory instance and agent_id.
 
         Args:
@@ -204,7 +207,7 @@ class LoComoThinkAnswerGenerator(LLMAnswerGenerator):
         question: str,
         recall_result: Dict[str, Any],
         question_date: Optional[datetime] = None,
-        question_type: Optional[str] = None
+        question_type: Optional[str] = None,
     ) -> Tuple[str, str, Optional[List[Dict[str, Any]]]]:
         """
         Generate answer using the integrated think API.
@@ -235,9 +238,9 @@ class LoComoThinkAnswerGenerator(LLMAnswerGenerator):
 
             # Extract memories from based_on
             based_on = result.based_on
-            world_facts = based_on.get('world', [])
-            agent_facts = based_on.get('agent', [])
-            opinion_facts = based_on.get('opinion', [])
+            world_facts = based_on.get("world", [])
+            agent_facts = based_on.get("agent", [])
+            opinion_facts = based_on.get("opinion", [])
 
             # Combine all facts into retrieved_memories
             retrieved_memories = []
@@ -271,7 +274,7 @@ async def run_benchmark(
     api_url: str = None,
     max_concurrent_questions_override: int = None,
     only_failed: bool = False,
-    only_invalid: bool = False
+    only_invalid: bool = False,
 ):
     """
     Run the LoComo benchmark.
@@ -287,6 +290,7 @@ async def run_benchmark(
         only_invalid: If True, only run conversations that have invalid questions (is_invalid=True)
     """
     from rich.console import Console
+
     console = Console()
 
     # Load previous results if filtering for failed/invalid conversations
@@ -294,35 +298,41 @@ async def run_benchmark(
     invalid_conversation_ids = set()
     if only_failed or only_invalid:
         suffix = "_think" if use_think else ""
-        results_filename = f'benchmark_results{suffix}.json'
-        results_path = Path(__file__).parent / 'results' / results_filename
+        results_filename = f"benchmark_results{suffix}.json"
+        results_path = Path(__file__).parent / "results" / results_filename
 
         if not results_path.exists():
-            console.print(f"[red]Error: Cannot use --only-failed or --only-invalid without existing results file[/red]")
+            console.print("[red]Error: Cannot use --only-failed or --only-invalid without existing results file[/red]")
             console.print(f"[yellow]Results file not found: {results_path}[/yellow]")
             return
 
-        with open(results_path, 'r') as f:
+        with open(results_path, "r") as f:
             previous_results = json.load(f)
 
         # Extract conversation IDs that have failed or invalid questions
-        for item_result in previous_results.get('item_results', []):
-            item_id = item_result['item_id']
-            for detail in item_result['metrics'].get('detailed_results', []):
-                if only_failed and detail.get('is_correct') == False and not detail.get('is_invalid', False):
+        for item_result in previous_results.get("item_results", []):
+            item_id = item_result["item_id"]
+            for detail in item_result["metrics"].get("detailed_results", []):
+                if only_failed and detail.get("is_correct") == False and not detail.get("is_invalid", False):
                     failed_conversation_ids.add(item_id)
-                if only_invalid and detail.get('is_invalid', False):
+                if only_invalid and detail.get("is_invalid", False):
                     invalid_conversation_ids.add(item_id)
 
         if only_failed:
-            console.print(f"[cyan]Filtering to {len(failed_conversation_ids)} conversations with failed questions (is_correct=False)[/cyan]")
+            console.print(
+                f"[cyan]Filtering to {len(failed_conversation_ids)} conversations with failed questions (is_correct=False)[/cyan]"
+            )
         if only_invalid:
-            console.print(f"[cyan]Filtering to {len(invalid_conversation_ids)} conversations with invalid questions (is_invalid=True)[/cyan]")
+            console.print(
+                f"[cyan]Filtering to {len(invalid_conversation_ids)} conversations with invalid questions (is_invalid=True)[/cyan]"
+            )
 
         target_ids = failed_conversation_ids if only_failed else invalid_conversation_ids
         if not target_ids:
             filter_type = "failed" if only_failed else "invalid"
-            console.print(f"[yellow]No conversations with {filter_type} questions found in previous results. Nothing to run.[/yellow]")
+            console.print(
+                f"[yellow]No conversations with {filter_type} questions found in previous results. Nothing to run.[/yellow]"
+            )
             return
 
     # Initialize components
@@ -331,18 +341,16 @@ async def run_benchmark(
     # Use remote API client if api_url is provided, otherwise use local memory
     if api_url:
         from benchmarks.common.benchmark_runner import HindsightClientAdapter
+
         memory = HindsightClientAdapter(base_url=api_url)
         await memory.initialize()
     else:
         from benchmarks.common.benchmark_runner import create_memory_engine
+
         memory = await create_memory_engine()
 
     if use_think:
-        answer_generator = LoComoThinkAnswerGenerator(
-            memory=memory,
-            agent_id="locomo",
-            thinking_budget=500
-        )
+        answer_generator = LoComoThinkAnswerGenerator(memory=memory, agent_id="locomo", thinking_budget=500)
         max_concurrent_questions = max_concurrent_questions_override or 4
         eval_semaphore_size = 4
     else:
@@ -356,14 +364,11 @@ async def run_benchmark(
 
     # Create benchmark runner
     runner = BenchmarkRunner(
-        dataset=dataset,
-        answer_generator=answer_generator,
-        answer_evaluator=answer_evaluator,
-        memory=memory
+        dataset=dataset, answer_generator=answer_generator, answer_evaluator=answer_evaluator, memory=memory
     )
 
     # Filter dataset if using --only-failed or --only-invalid
-    dataset_path = Path(__file__).parent / 'datasets' / 'locomo10.json'
+    dataset_path = Path(__file__).parent / "datasets" / "locomo10.json"
 
     if only_failed or only_invalid:
         # Load and filter dataset
@@ -374,14 +379,16 @@ async def run_benchmark(
 
         # Temporarily replace dataset's load method
         original_load = dataset.load
+
         def filtered_load(path: Path, max_items: Optional[int] = None):
             return filtered_items[:max_items] if max_items else filtered_items
+
         dataset.load = filtered_load
 
     # Determine output filename based on mode
     suffix = "_think" if use_think else ""
-    results_filename = f'benchmark_results{suffix}.json'
-    output_path = Path(__file__).parent / 'results' / results_filename
+    results_filename = f"benchmark_results{suffix}.json"
+    output_path = Path(__file__).parent / "results" / results_filename
 
     # Create results directory if it doesn't exist
     output_path.parent.mkdir(parents=True, exist_ok=True)
@@ -406,7 +413,7 @@ async def run_benchmark(
         clear_agent_per_item=True,  # Use unique agent ID per conversation
         max_concurrent_items=3,  # Process up to 3 conversations in parallel
         output_path=output_path,  # Save results incrementally
-        merge_with_existing=merge_with_existing
+        merge_with_existing=merge_with_existing,
     )
 
     # Display results (final save already happened incrementally)
@@ -430,14 +437,10 @@ def generate_markdown_table(results: dict, use_think: bool = False):
     4 = Open-domain
     """
     from rich.console import Console
+
     console = Console()
 
-    category_names = {
-        '1': 'Multi-hop',
-        '2': 'Single-hop',
-        '3': 'Temporal',
-        '4': 'Open-domain'
-    }
+    category_names = {"1": "Multi-hop", "2": "Single-hop", "3": "Temporal", "4": "Open-domain"}
 
     # Build markdown content
     lines = []
@@ -446,33 +449,41 @@ def generate_markdown_table(results: dict, use_think: bool = False):
     lines.append("")
 
     # Add model configuration
-    if 'model_config' in results:
-        config = results['model_config']
+    if "model_config" in results:
+        config = results["model_config"]
         lines.append("## Model Configuration")
         lines.append("")
         lines.append(f"- **Hindsight**: {config['hindsight']['provider']}/{config['hindsight']['model']}")
-        lines.append(f"- **Answer Generation**: {config['answer_generation']['provider']}/{config['answer_generation']['model']}")
+        lines.append(
+            f"- **Answer Generation**: {config['answer_generation']['provider']}/{config['answer_generation']['model']}"
+        )
         lines.append(f"- **LLM Judge**: {config['judge']['provider']}/{config['judge']['model']}")
         lines.append("")
 
-    lines.append(f"**Overall Accuracy**: {results['overall_accuracy']:.2f}% ({results['total_correct']}/{results['total_questions']})")
+    lines.append(
+        f"**Overall Accuracy**: {results['overall_accuracy']:.2f}% ({results['total_correct']}/{results['total_questions']})"
+    )
     lines.append("")
-    lines.append("| Sample ID | Sessions | Questions | Correct | Accuracy | Multi-hop | Single-hop | Temporal | Open-domain |")
-    lines.append("|-----------|----------|-----------|---------|----------|-----------|------------|----------|-------------|")
+    lines.append(
+        "| Sample ID | Sessions | Questions | Correct | Accuracy | Multi-hop | Single-hop | Temporal | Open-domain |"
+    )
+    lines.append(
+        "|-----------|----------|-----------|---------|----------|-----------|------------|----------|-------------|"
+    )
 
-    for item_result in results['item_results']:
-        item_id = item_result['item_id']
-        num_sessions = item_result['num_sessions']
-        metrics = item_result['metrics']
+    for item_result in results["item_results"]:
+        item_id = item_result["item_id"]
+        num_sessions = item_result["num_sessions"]
+        metrics = item_result["metrics"]
 
         # Calculate category accuracies
-        cat_stats = metrics.get('category_stats', {})
+        cat_stats = metrics.get("category_stats", {})
         cat_accuracies = {}
 
-        for cat_id in ['1', '2', '3', '4']:
+        for cat_id in ["1", "2", "3", "4"]:
             if cat_id in cat_stats:
                 stats = cat_stats[cat_id]
-                acc = (stats['correct'] / stats['total'] * 100) if stats['total'] > 0 else 0
+                acc = (stats["correct"] / stats["total"] * 100) if stats["total"] > 0 else 0
                 cat_accuracies[cat_id] = f"{acc:.1f}% ({stats['correct']}/{stats['total']})"
             else:
                 cat_accuracies[cat_id] = "N/A"
@@ -485,28 +496,48 @@ def generate_markdown_table(results: dict, use_think: bool = False):
 
     # Write to file with suffix
     suffix = "_think" if use_think else ""
-    output_file = Path(__file__).parent / 'results' / f'results_table{suffix}.md'
+    output_file = Path(__file__).parent / "results" / f"results_table{suffix}.md"
     output_file.parent.mkdir(parents=True, exist_ok=True)
-    output_file.write_text('\n'.join(lines))
+    output_file.write_text("\n".join(lines))
     console.print(f"\n[green]✓[/green] Results table saved to {output_file}")
 
 
 if __name__ == "__main__":
-    import logging
     import argparse
+    import logging
 
-    logging.basicConfig(level=logging.INFO, format='%(asctime)s %(levelname)s %(message)s')
+    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
 
-    parser = argparse.ArgumentParser(description='Run LoComo benchmark')
-    parser.add_argument('--max-conversations', type=int, default=None, help='Maximum conversations to evaluate')
-    parser.add_argument('--max-questions', type=int, default=None, help='Maximum questions per conversation')
-    parser.add_argument('--skip-ingestion', action='store_true', help='Skip ingestion and use existing data')
-    parser.add_argument('--use-think', action='store_true', help='Use think API instead of search + LLM')
-    parser.add_argument('--conversation', type=str, default=None, help='Run only specific conversation (e.g., "conv-26")')
-    parser.add_argument('--api-url', type=str, default=None, help='Hindsight API URL (default: use local memory, example: http://localhost:8888)')
-    parser.add_argument('--max-concurrent-questions', type=int, default=None, help='Max concurrent questions per conversation (default: 4 for think, 10 for search)')
-    parser.add_argument('--only-failed', action='store_true', help='Only run conversations that have failed questions (is_correct=False). Requires existing results file.')
-    parser.add_argument('--only-invalid', action='store_true', help='Only run conversations that have invalid questions (is_invalid=True). Requires existing results file.')
+    parser = argparse.ArgumentParser(description="Run LoComo benchmark")
+    parser.add_argument("--max-conversations", type=int, default=None, help="Maximum conversations to evaluate")
+    parser.add_argument("--max-questions", type=int, default=None, help="Maximum questions per conversation")
+    parser.add_argument("--skip-ingestion", action="store_true", help="Skip ingestion and use existing data")
+    parser.add_argument("--use-think", action="store_true", help="Use think API instead of search + LLM")
+    parser.add_argument(
+        "--conversation", type=str, default=None, help='Run only specific conversation (e.g., "conv-26")'
+    )
+    parser.add_argument(
+        "--api-url",
+        type=str,
+        default=None,
+        help="Hindsight API URL (default: use local memory, example: http://localhost:8888)",
+    )
+    parser.add_argument(
+        "--max-concurrent-questions",
+        type=int,
+        default=None,
+        help="Max concurrent questions per conversation (default: 4 for think, 10 for search)",
+    )
+    parser.add_argument(
+        "--only-failed",
+        action="store_true",
+        help="Only run conversations that have failed questions (is_correct=False). Requires existing results file.",
+    )
+    parser.add_argument(
+        "--only-invalid",
+        action="store_true",
+        help="Only run conversations that have invalid questions (is_invalid=True). Requires existing results file.",
+    )
 
     args = parser.parse_args()
 
@@ -514,14 +545,16 @@ if __name__ == "__main__":
     if args.only_failed and args.only_invalid:
         parser.error("Cannot use both --only-failed and --only-invalid at the same time")
 
-    results = asyncio.run(run_benchmark(
-        max_conversations=args.max_conversations,
-        max_questions_per_conv=args.max_questions,
-        skip_ingestion=args.skip_ingestion,
-        use_think=args.use_think,
-        conversation=args.conversation,
-        api_url=args.api_url,
-        max_concurrent_questions_override=args.max_concurrent_questions,
-        only_failed=args.only_failed,
-        only_invalid=args.only_invalid
-    ))
+    results = asyncio.run(
+        run_benchmark(
+            max_conversations=args.max_conversations,
+            max_questions_per_conv=args.max_questions,
+            skip_ingestion=args.skip_ingestion,
+            use_think=args.use_think,
+            conversation=args.conversation,
+            api_url=args.api_url,
+            max_concurrent_questions_override=args.max_concurrent_questions,
+            only_failed=args.only_failed,
+            only_invalid=args.only_invalid,
+        )
+    )

hindsight-dev/benchmarks/longmemeval/longmemeval_benchmark.py

@@ -3,21 +3,20 @@ LongMemEval-specific benchmark implementations.
 
 Provides dataset, answer generator, and evaluator for the LongMemEval benchmark.
 """
-import sys
-from pathlib import Path
 
-from benchmarks.common.benchmark_runner import BenchmarkRunner
-
-import json
-from datetime import datetime, timezone
-from typing import List, Dict, Any, Tuple, Optional
 import asyncio
-import pydantic
-from openai import AsyncOpenAI
+import json
 import os
+import sys
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple
 
-from benchmarks.common.benchmark_runner import BenchmarkDataset, LLMAnswerGenerator, LLMAnswerEvaluator
+import pydantic
 from hindsight_api.engine.llm_wrapper import LLMConfig
+from openai import AsyncOpenAI
+
+from benchmarks.common.benchmark_runner import BenchmarkDataset, BenchmarkRunner, LLMAnswerEvaluator, LLMAnswerGenerator
 
 
 class LongMemEvalDataset(BenchmarkDataset):
@@ -25,7 +24,7 @@ class LongMemEvalDataset(BenchmarkDataset):
 
     def load(self, path: Path, max_items: Optional[int] = None) -> List[Dict[str, Any]]:
         """Load LongMemEval dataset from JSON file."""
-        with open(path, 'r') as f:
+        with open(path, "r") as f:
             dataset = json.load(f)
 
         if max_items:
@@ -67,7 +66,7 @@ class LongMemEvalDataset(BenchmarkDataset):
             for turn in session_turns:
                 if isinstance(turn, dict):
                     # Create a copy without has_answer
-                    cleaned_turn = {k: v for k, v in turn.items() if k != 'has_answer'}
+                    cleaned_turn = {k: v for k, v in turn.items() if k != "has_answer"}
                     cleaned_turns.append(cleaned_turn)
                 else:
                     cleaned_turns.append(turn)
@@ -75,12 +74,14 @@ class LongMemEvalDataset(BenchmarkDataset):
             session_content = json.dumps(cleaned_turns)
             question_id = item.get("question_id", "unknown")
             document_id = f"{question_id}_{session_id}"
-            batch_contents.append({
-                "content": session_content,
-                "context": f"Session {document_id} - you are the assistant in this conversation - happened on {session_date.strftime('%Y-%m-%d %H:%M:%S')} UTC.",
-                "event_date": session_date,
-                "document_id": document_id
-            })
+            batch_contents.append(
+                {
+                    "content": session_content,
+                    "context": f"Session {document_id} - you are the assistant in this conversation - happened on {session_date.strftime('%Y-%m-%d %H:%M:%S')} UTC.",
+                    "event_date": session_date,
+                    "document_id": document_id,
+                }
+            )
 
         return batch_contents
 
@@ -95,22 +96,24 @@ class LongMemEvalDataset(BenchmarkDataset):
         """
         # Parse question_date if available
         question_date = None
-        if 'question_date' in item:
-            question_date = self._parse_date(item['question_date'])
+        if "question_date" in item:
+            question_date = self._parse_date(item["question_date"])
 
-        return [{
-            'question': item.get("question", ""),
-            'answer': item.get("answer", ""),
-            'category': item.get("question_type", "unknown"),
-            'question_date': question_date
-        }]
+        return [
+            {
+                "question": item.get("question", ""),
+                "answer": item.get("answer", ""),
+                "category": item.get("question_type", "unknown"),
+                "question_date": question_date,
+            }
+        ]
 
     def _parse_date(self, date_str: str) -> datetime:
         """Parse date string to datetime object."""
         try:
             # LongMemEval format: "2023/05/20 (Sat) 02:21"
             # Try to parse the main part before the day name
-            date_str_cleaned = date_str.split('(')[0].strip() if '(' in date_str else date_str
+            date_str_cleaned = date_str.split("(")[0].strip() if "(" in date_str else date_str
 
             # Try multiple formats
             for fmt in ["%Y/%m/%d %H:%M", "%Y-%m-%d %H:%M:%S", "%Y-%m-%d", "%Y/%m/%d"]:
@@ -121,7 +124,7 @@ class LongMemEvalDataset(BenchmarkDataset):
                     continue
 
             # Fallback: try ISO format
-            return datetime.fromisoformat(date_str.replace('Z', '+00:00'))
+            return datetime.fromisoformat(date_str.replace("Z", "+00:00"))
         except Exception:
             raise ValueError(f"Failed to parse date string: {date_str}")
 
@@ -130,6 +133,7 @@ class QuestionAnswer(pydantic.BaseModel):
     answer: str
     reasoning: Optional[str] = None
 
+
 class LongMemEvalAnswerGenerator(LLMAnswerGenerator):
     """LongMemEval-specific answer generator using configurable LLM provider."""
 
@@ -202,10 +206,7 @@ class LongMemEvalAnswerGenerator(LLMAnswerGenerator):
                 chunk_text = chunk_info.get("chunk_text", "")
 
             # Build the formatted fact entry
-            entry_parts = [
-                f"Fact {i} ({fact_type}): {fact_text}",
-                f"When: {when_str}"
-            ]
+            entry_parts = [f"Fact {i} ({fact_type}): {fact_text}", f"When: {when_str}"]
 
             # Add context field if present
             context = fact.get("context")
@@ -217,7 +218,7 @@ class LongMemEvalAnswerGenerator(LLMAnswerGenerator):
                 # Truncate very long chunks
                 if len(chunk_text) > 1000:
                     chunk_text = chunk_text[:1000] + "..."
-                entry_parts.append(f"Source chunk:\n  \"{chunk_text}\"")
+                entry_parts.append(f'Source chunk:\n  "{chunk_text}"')
 
             formatted_parts.append("\n".join(entry_parts))
 
@@ -326,43 +327,43 @@ The context contains memory facts extracted from previous conversations, each wi
             return ""
 
     async def generate_answer(
-                self,
-                question: str,
-                recall_result: Dict[str, Any],
-                question_date: Optional[datetime] = None,
-                question_type: Optional[str] = None
-        ) -> Tuple[str, str, Optional[List[Dict[str, Any]]]]:
-            """
-            Generate answer from retrieved memories using Groq.
+        self,
+        question: str,
+        recall_result: Dict[str, Any],
+        question_date: Optional[datetime] = None,
+        question_type: Optional[str] = None,
+    ) -> Tuple[str, str, Optional[List[Dict[str, Any]]]]:
+        """
+        Generate answer from retrieved memories using Groq.
 
-            Args:
-                question: The question text
-                recall_result: Full RecallResult dict containing results, entities, chunks, and trace
-                question_date: Date when the question was asked (for temporal context)
-                question_type: Question category (e.g., 'single-session-user', 'multi-session-assistant')
+        Args:
+            question: The question text
+            recall_result: Full RecallResult dict containing results, entities, chunks, and trace
+            question_date: Date when the question was asked (for temporal context)
+            question_type: Question category (e.g., 'single-session-user', 'multi-session-assistant')
 
-            Returns:
-                Tuple of (answer, reasoning, None)
-                - None indicates to use the memories from recall_result
-            """
-            # Format context based on selected mode
-            if self.context_format == "structured":
-                context = self._format_context_structured(recall_result)
-            else:
-                context = self._format_context_json(recall_result)
+        Returns:
+            Tuple of (answer, reasoning, None)
+            - None indicates to use the memories from recall_result
+        """
+        # Format context based on selected mode
+        if self.context_format == "structured":
+            context = self._format_context_structured(recall_result)
+        else:
+            context = self._format_context_json(recall_result)
 
-            context_instructions = self._get_context_instructions()
+        context_instructions = self._get_context_instructions()
 
-            # Format question date if provided
-            formatted_question_date = question_date.strftime('%Y-%m-%d %H:%M:%S UTC') if question_date else "Not specified"
+        # Format question date if provided
+        formatted_question_date = question_date.strftime("%Y-%m-%d %H:%M:%S UTC") if question_date else "Not specified"
 
-            # Use LLM to generate answer
-            try:
-                answer_obj = await self.llm_config.call(
-                    messages=[
-                        {
-                            "role": "user",
-                            "content": f"""You are a helpful assistant that must answer user questions based on the previous conversations.
+        # Use LLM to generate answer
+        try:
+            answer_obj = await self.llm_config.call(
+                messages=[
+                    {
+                        "role": "user",
+                        "content": f"""You are a helpful assistant that must answer user questions based on the previous conversations.
 
 {context_instructions}**Answer Guidelines:**
 1. Start by scanning retrieved context to understand the facts and events that happened and the timeline.
@@ -393,20 +394,20 @@ Retrieved Context:
 
 
 Answer:
-"""
-                        }
-                    ],
-                    response_format=QuestionAnswer,
-                    scope="memory",
-                    max_completion_tokens=32768,
-                )
-                reasoning_text = answer_obj.reasoning or ""
-                if reasoning_text:
-                    reasoning_text = reasoning_text + " "
-                reasoning_text += f"(question date: {formatted_question_date})"
-                return answer_obj.answer, reasoning_text, None
-            except Exception as e:
-                return f"Error generating answer: {str(e)}", "Error occurred during answer generation.", None
+""",
+                    }
+                ],
+                response_format=QuestionAnswer,
+                scope="memory",
+                max_completion_tokens=32768,
+            )
+            reasoning_text = answer_obj.reasoning or ""
+            if reasoning_text:
+                reasoning_text = reasoning_text + " "
+            reasoning_text += f"(question date: {formatted_question_date})"
+            return answer_obj.answer, reasoning_text, None
+        except Exception as e:
+            return f"Error generating answer: {str(e)}", "Error occurred during answer generation.", None
 
 
 async def run_benchmark(
@@ -425,7 +426,7 @@ async def run_benchmark(
     max_concurrent_items: int = 1,
     results_filename: str = "benchmark_results.json",
     context_format: str = "json",
-    source_results: str = None
+    source_results: str = None,
 ):
     """
     Run the LongMemEval benchmark.
@@ -449,13 +450,16 @@ async def run_benchmark(
         source_results: Source results file to read failed/invalid questions from (for --only-failed/--only-invalid). Defaults to benchmark_results.json.
     """
     from rich.console import Console
+
     console = Console()
 
     # Validate mutually exclusive arguments
     # --max-instances-per-category can't be combined with --max-instances or --category
     # But --category CAN be combined with --max-instances (to limit questions within a category)
     if max_instances_per_category is not None and (max_instances is not None or category is not None):
-        console.print("[red]Error: --max-questions-per-category cannot be combined with --max-instances or --category[/red]")
+        console.print(
+            "[red]Error: --max-questions-per-category cannot be combined with --max-instances or --category[/red]"
+        )
         return
 
     # Validate --only-ingested can't be combined with other dataset filters
@@ -480,8 +484,10 @@ async def run_benchmark(
     dataset_path = Path(__file__).parent / "datasets" / "longmemeval_s_cleaned.json"
     if not dataset_path.exists():
         if not download_dataset(dataset_path):
-            console.print(f"[red]Failed to download dataset. Please download manually:[/red]")
-            console.print("[yellow]curl -L 'https://huggingface.co/datasets/xiaowu0162/longmemeval-cleaned/resolve/main/longmemeval_s_cleaned.json' -o benchmarks/longmemeval/datasets/longmemeval_s_cleaned.json[/yellow]")
+            console.print("[red]Failed to download dataset. Please download manually:[/red]")
+            console.print(
+                "[yellow]curl -L 'https://huggingface.co/datasets/xiaowu0162/longmemeval-cleaned/resolve/main/longmemeval_s_cleaned.json' -o benchmarks/longmemeval/datasets/longmemeval_s_cleaned.json[/yellow]"
+            )
             return
 
     # Initialize components
@@ -499,9 +505,10 @@ async def run_benchmark(
 
         # Group by category and take max_instances_per_category from each
         from collections import defaultdict
+
         category_items = defaultdict(list)
         for item in original_dataset_items:
-            cat = item.get('question_type', 'unknown')
+            cat = item.get("question_type", "unknown")
             category_items[cat].append(item)
 
         # Take up to max_instances_per_category from each category
@@ -518,30 +525,34 @@ async def run_benchmark(
     invalid_question_ids = set()
     if only_failed or only_invalid:
         # Use source_results if specified, otherwise default to benchmark_results.json
-        source_file = source_results if source_results else 'benchmark_results.json'
-        results_path = Path(__file__).parent / 'results' / source_file
+        source_file = source_results if source_results else "benchmark_results.json"
+        results_path = Path(__file__).parent / "results" / source_file
         if not results_path.exists():
-            console.print(f"[red]Error: Cannot use --only-failed or --only-invalid without existing results file[/red]")
+            console.print("[red]Error: Cannot use --only-failed or --only-invalid without existing results file[/red]")
             console.print(f"[yellow]Results file not found: {results_path}[/yellow]")
             return
 
         console.print(f"[cyan]Reading failed/invalid questions from: {source_file}[/cyan]")
-        with open(results_path, 'r') as f:
+        with open(results_path, "r") as f:
             previous_results = json.load(f)
 
         # Extract question IDs that failed or are invalid
-        for item_result in previous_results.get('item_results', []):
-            item_id = item_result['item_id']
-            for detail in item_result['metrics'].get('detailed_results', []):
-                if only_failed and detail.get('is_correct') == False and not detail.get('is_invalid', False):
+        for item_result in previous_results.get("item_results", []):
+            item_id = item_result["item_id"]
+            for detail in item_result["metrics"].get("detailed_results", []):
+                if only_failed and detail.get("is_correct") == False and not detail.get("is_invalid", False):
                     failed_question_ids.add(item_id)
-                if only_invalid and detail.get('is_invalid', False):
+                if only_invalid and detail.get("is_invalid", False):
                     invalid_question_ids.add(item_id)
 
         if only_failed:
-            console.print(f"[cyan]Filtering to {len(failed_question_ids)} questions that failed (is_correct=False)[/cyan]")
+            console.print(
+                f"[cyan]Filtering to {len(failed_question_ids)} questions that failed (is_correct=False)[/cyan]"
+            )
         if only_invalid:
-            console.print(f"[cyan]Filtering to {len(invalid_question_ids)} questions that were invalid (is_invalid=True)[/cyan]")
+            console.print(
+                f"[cyan]Filtering to {len(invalid_question_ids)} questions that were invalid (is_invalid=True)[/cyan]"
+            )
 
     # Filter dataset by category if specified
     if category:
@@ -550,11 +561,11 @@ async def run_benchmark(
             # Load full dataset without max_instances limit for filtering
             original_dataset_items = dataset.load(dataset_path, max_items=None)
 
-        filtered_items = [item for item in original_dataset_items if item.get('question_type') == category]
+        filtered_items = [item for item in original_dataset_items if item.get("question_type") == category]
 
         if not filtered_items:
             console.print(f"[yellow]No questions found for category '{category}'. Available categories:[/yellow]")
-            available_categories = set(item.get('question_type', 'unknown') for item in original_dataset_items)
+            available_categories = set(item.get("question_type", "unknown") for item in original_dataset_items)
             for cat in sorted(available_categories):
                 console.print(f"  - {cat}")
             return
@@ -562,7 +573,9 @@ async def run_benchmark(
         total_found = len(filtered_items)
         will_run = min(total_found, max_instances) if max_instances else total_found
         if max_instances and total_found > max_instances:
-            console.print(f"[green]Found {total_found} questions for category '{category}' (will run {will_run} due to --max-instances)[/green]")
+            console.print(
+                f"[green]Found {total_found} questions for category '{category}' (will run {will_run} due to --max-instances)[/green]"
+            )
         else:
             console.print(f"[green]Found {total_found} questions for category '{category}'[/green]")
 
@@ -588,7 +601,9 @@ async def run_benchmark(
         total_found = len(filtered_items)
         will_run = min(total_found, max_instances) if max_instances else total_found
         if max_instances and total_found > max_instances:
-            console.print(f"[green]Found {total_found} {filter_type} items to re-evaluate (will run {will_run} due to --max-instances)[/green]")
+            console.print(
+                f"[green]Found {total_found} {filter_type} items to re-evaluate (will run {will_run} due to --max-instances)[/green]"
+            )
         else:
             console.print(f"[green]Found {total_found} {filter_type} items to re-evaluate[/green]")
 
@@ -600,6 +615,7 @@ async def run_benchmark(
 
     # Create local memory engine
     from benchmarks.common.benchmark_runner import create_memory_engine
+
     memory = await create_memory_engine()
 
     # Filter by only_ingested: only run items whose memory bank already exists
@@ -623,10 +639,9 @@ async def run_benchmark(
             # Check if bank has any memory units
             async with pool.acquire() as conn:
                 result = await conn.fetchrow(
-                    "SELECT COUNT(*) as count FROM memory_units WHERE bank_id = $1 LIMIT 1",
-                    agent_id
+                    "SELECT COUNT(*) as count FROM memory_units WHERE bank_id = $1 LIMIT 1", agent_id
                 )
-                if result['count'] > 0:
+                if result["count"] > 0:
                     ingested_items.append(item)
 
         filtered_items = ingested_items
@@ -638,34 +653,43 @@ async def run_benchmark(
 
     # Create benchmark runner
     runner = BenchmarkRunner(
-        dataset=dataset,
-        answer_generator=answer_generator,
-        answer_evaluator=answer_evaluator,
-        memory=memory
+        dataset=dataset, answer_generator=answer_generator, answer_evaluator=answer_evaluator, memory=memory
     )
 
     # If filtering by category, failed, invalid, only_ingested, or max_instances_per_category, we need to use a custom dataset that only returns those items
     # We'll temporarily replace the dataset's load method
     if filtered_items is not None:
         original_load = dataset.load
+
         def filtered_load(path: Path, max_items: Optional[int] = None):
             return filtered_items[:max_items] if max_items else filtered_items
+
         dataset.load = filtered_load
 
     # Run benchmark
     # Single-phase approach: each question gets its own isolated agent_id
     # This ensures each question only has access to its own context
-    output_path = Path(__file__).parent / 'results' / results_filename
+    output_path = Path(__file__).parent / "results" / results_filename
 
     # Create results directory if it doesn't exist
     output_path.parent.mkdir(parents=True, exist_ok=True)
 
-    merge_with_existing = (filln or question_id is not None or only_failed or only_invalid or only_ingested or category is not None or max_instances_per_category is not None)
+    merge_with_existing = (
+        filln
+        or question_id is not None
+        or only_failed
+        or only_invalid
+        or only_ingested
+        or category is not None
+        or max_instances_per_category is not None
+    )
 
     results = await runner.run(
         dataset_path=dataset_path,
         agent_id="longmemeval",  # Will be suffixed with question_id per item
-        max_items=max_instances if not max_instances_per_category else None,  # Don't apply max_items when using per-category limit
+        max_items=max_instances
+        if not max_instances_per_category
+        else None,  # Don't apply max_items when using per-category limit
         max_questions_per_item=max_questions_per_instance,
         thinking_budget=thinking_budget,
         max_tokens=max_tokens,
@@ -678,7 +702,7 @@ async def run_benchmark(
         specific_item=question_id,  # Optional filter for specific question ID
         max_concurrent_items=max_concurrent_items,  # Parallel instance processing
         output_path=output_path,  # Save results incrementally
-        merge_with_existing=merge_with_existing  # Merge when using --fill, --category, --only-failed, --only-invalid flags or specific question
+        merge_with_existing=merge_with_existing,  # Merge when using --fill, --category, --only-failed, --only-invalid flags or specific question
     )
 
     # Display results (final save already happened incrementally)
@@ -702,12 +726,14 @@ def download_dataset(dataset_path: Path) -> bool:
         True if successful, False otherwise
     """
     import subprocess
+
     from rich.console import Console
+
     console = Console()
 
     url = "https://huggingface.co/datasets/xiaowu0162/longmemeval-cleaned/resolve/main/longmemeval_s_cleaned.json"
 
-    console.print(f"[yellow]Dataset not found. Downloading from HuggingFace...[/yellow]")
+    console.print("[yellow]Dataset not found. Downloading from HuggingFace...[/yellow]")
     console.print(f"[dim]URL: {url}[/dim]")
     console.print(f"[dim]Destination: {dataset_path}[/dim]")
 
@@ -720,18 +746,18 @@ def download_dataset(dataset_path: Path) -> bool:
             ["curl", "-L", "-o", str(dataset_path), url],
             capture_output=True,
             text=True,
-            timeout=300  # 5 minute timeout
+            timeout=300,  # 5 minute timeout
         )
 
         if result.returncode == 0 and dataset_path.exists():
-            console.print(f"[green]✓ Dataset downloaded successfully[/green]")
+            console.print("[green]✓ Dataset downloaded successfully[/green]")
             return True
         else:
             console.print(f"[red]✗ Download failed: {result.stderr}[/red]")
             return False
 
     except subprocess.TimeoutExpired:
-        console.print(f"[red]✗ Download timed out after 5 minutes[/red]")
+        console.print("[red]✗ Download timed out after 5 minutes[/red]")
         return False
     except Exception as e:
         console.print(f"[red]✗ Download error: {e}[/red]")
@@ -740,22 +766,23 @@ def download_dataset(dataset_path: Path) -> bool:
 
 def generate_type_report(results: dict):
     """Generate a detailed report by question type."""
-    from rich.table import Table
     from rich.console import Console
+    from rich.table import Table
+
     console = Console()
 
     # Aggregate stats by question type
     type_stats = {}
 
-    for item_result in results['item_results']:
-        metrics = item_result['metrics']
-        by_category = metrics.get('category_stats', {})
+    for item_result in results["item_results"]:
+        metrics = item_result["metrics"]
+        by_category = metrics.get("category_stats", {})
 
         for qtype, stats in by_category.items():
             if qtype not in type_stats:
-                type_stats[qtype] = {'total': 0, 'correct': 0}
-            type_stats[qtype]['total'] += stats['total']
-            type_stats[qtype]['correct'] += stats['correct']
+                type_stats[qtype] = {"total": 0, "correct": 0}
+            type_stats[qtype]["total"] += stats["total"]
+            type_stats[qtype]["correct"] += stats["correct"]
 
     # Display table
     table = Table(title="Performance by Question Type")
@@ -765,13 +792,8 @@ def generate_type_report(results: dict):
     table.add_column("Accuracy", justify="right", style="magenta")
 
     for qtype, stats in sorted(type_stats.items()):
-        acc = (stats['correct'] / stats['total'] * 100) if stats['total'] > 0 else 0
-        table.add_row(
-            qtype,
-            str(stats['total']),
-            str(stats['correct']),
-            f"{acc:.1f}%"
-        )
+        acc = (stats["correct"] / stats["total"] * 100) if stats["total"] > 0 else 0
+        table.add_row(qtype, str(stats["total"]), str(stats["correct"]), f"{acc:.1f}%")
 
     console.print("\n")
     console.print(table)
@@ -780,21 +802,22 @@ def generate_type_report(results: dict):
 def generate_markdown_table(results: dict, json_output_path: Path):
     """Generate a markdown results table with model configuration."""
     from rich.console import Console
+
     console = Console()
 
     # Aggregate stats by question type
     type_stats = {}
 
-    for item_result in results['item_results']:
-        metrics = item_result['metrics']
-        by_category = metrics.get('category_stats', {})
+    for item_result in results["item_results"]:
+        metrics = item_result["metrics"]
+        by_category = metrics.get("category_stats", {})
 
         for qtype, stats in by_category.items():
             if qtype not in type_stats:
-                type_stats[qtype] = {'total': 0, 'correct': 0, 'invalid': 0}
-            type_stats[qtype]['total'] += stats['total']
-            type_stats[qtype]['correct'] += stats['correct']
-            type_stats[qtype]['invalid'] += stats.get('invalid', 0)
+                type_stats[qtype] = {"total": 0, "correct": 0, "invalid": 0}
+            type_stats[qtype]["total"] += stats["total"]
+            type_stats[qtype]["correct"] += stats["correct"]
+            type_stats[qtype]["invalid"] += stats.get("invalid", 0)
 
     # Build markdown content
     lines = []
@@ -802,16 +825,20 @@ def generate_markdown_table(results: dict, json_output_path: Path):
     lines.append("")
 
     # Add model configuration
-    if 'model_config' in results:
-        config = results['model_config']
+    if "model_config" in results:
+        config = results["model_config"]
         lines.append("## Model Configuration")
         lines.append("")
         lines.append(f"- **Hindsight**: {config['hindsight']['provider']}/{config['hindsight']['model']}")
-        lines.append(f"- **Answer Generation**: {config['answer_generation']['provider']}/{config['answer_generation']['model']}")
+        lines.append(
+            f"- **Answer Generation**: {config['answer_generation']['provider']}/{config['answer_generation']['model']}"
+        )
         lines.append(f"- **LLM Judge**: {config['judge']['provider']}/{config['judge']['model']}")
         lines.append("")
 
-    lines.append(f"**Overall Accuracy**: {results['overall_accuracy']:.2f}% ({results['total_correct']}/{results['total_questions']})")
+    lines.append(
+        f"**Overall Accuracy**: {results['overall_accuracy']:.2f}% ({results['total_correct']}/{results['total_questions']})"
+    )
     lines.append("")
 
     # Results by question type
@@ -822,34 +849,36 @@ def generate_markdown_table(results: dict, json_output_path: Path):
 
     for qtype in sorted(type_stats.keys()):
         stats = type_stats[qtype]
-        valid_total = stats['total'] - stats['invalid']
-        acc = (stats['correct'] / valid_total * 100) if valid_total > 0 else 0
-        invalid_str = str(stats['invalid']) if stats['invalid'] > 0 else "-"
+        valid_total = stats["total"] - stats["invalid"]
+        acc = (stats["correct"] / valid_total * 100) if valid_total > 0 else 0
+        invalid_str = str(stats["invalid"]) if stats["invalid"] > 0 else "-"
         lines.append(f"| {qtype} | {stats['total']} | {stats['correct']} | {invalid_str} | {acc:.1f}% |")
 
     # Add overall row
-    total_invalid = results.get('total_invalid', 0)
+    total_invalid = results.get("total_invalid", 0)
     invalid_str = str(total_invalid) if total_invalid > 0 else "-"
-    lines.append(f"| **OVERALL** | **{results['total_questions']}** | **{results['total_correct']}** | **{invalid_str}** | **{results['overall_accuracy']:.1f}%** |")
+    lines.append(
+        f"| **OVERALL** | **{results['total_questions']}** | **{results['total_correct']}** | **{invalid_str}** | **{results['overall_accuracy']:.1f}%** |"
+    )
 
     # Write to file (same directory as JSON, but .md extension)
-    md_output_path = json_output_path.with_suffix('.md')
-    md_output_path.write_text('\n'.join(lines))
+    md_output_path = json_output_path.with_suffix(".md")
+    md_output_path.write_text("\n".join(lines))
     console.print(f"\n[green]✓[/green] Results table saved to {md_output_path}")
 
 
 if __name__ == "__main__":
-    import logging
     import argparse
+    import logging
 
-    logging.basicConfig(level=logging.INFO, format='%(asctime)s %(levelname)s %(message)s')
+    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
 
     parser = argparse.ArgumentParser(description="Run LongMemEval benchmark")
     parser.add_argument(
         "--max-instances",
         type=int,
         default=None,
-        help="Limit TOTAL number of questions to evaluate (default: all 500). For per-category limits, use --max-questions-per-category instead."
+        help="Limit TOTAL number of questions to evaluate (default: all 500). For per-category limits, use --max-questions-per-category instead.",
     )
     parser.add_argument(
         "--max-instances-per-category",
@@ -857,87 +886,72 @@ if __name__ == "__main__":
         type=int,
         default=None,
         dest="max_instances_per_category",
-        help="Limit number of questions per category (e.g., 20 = 20 questions from each of the 6 categories = 120 total). Cannot be combined with --max-instances or --category."
+        help="Limit number of questions per category (e.g., 20 = 20 questions from each of the 6 categories = 120 total). Cannot be combined with --max-instances or --category.",
     )
     parser.add_argument(
-        "--max-questions",
-        type=int,
-        default=None,
-        help="Limit number of questions per instance (for quick testing)"
+        "--max-questions", type=int, default=None, help="Limit number of questions per instance (for quick testing)"
     )
     parser.add_argument(
-        "--thinking-budget",
-        type=int,
-        default=500,
-        help="Thinking budget for spreading activation search"
-    )
-    parser.add_argument(
-        "--max-tokens",
-        type=int,
-        default=8192,
-        help="Maximum tokens to retrieve from memories"
-    )
-    parser.add_argument(
-        "--skip-ingestion",
-        action="store_true",
-        help="Skip ingestion and use existing data"
+        "--thinking-budget", type=int, default=500, help="Thinking budget for spreading activation search"
     )
+    parser.add_argument("--max-tokens", type=int, default=8192, help="Maximum tokens to retrieve from memories")
+    parser.add_argument("--skip-ingestion", action="store_true", help="Skip ingestion and use existing data")
     parser.add_argument(
         "--fill",
         action="store_true",
-        help="Only process questions not already in results file (for resuming interrupted runs)"
+        help="Only process questions not already in results file (for resuming interrupted runs)",
     )
     parser.add_argument(
         "--question-id",
         type=str,
         default=None,
-        help="Filter to specific question ID (e.g., 'e47becba'). Useful with --skip-ingestion to test a single question."
+        help="Filter to specific question ID (e.g., 'e47becba'). Useful with --skip-ingestion to test a single question.",
     )
     parser.add_argument(
         "--only-failed",
         action="store_true",
-        help="Only run questions that were previously marked as incorrect (is_correct=False). Requires existing results file."
+        help="Only run questions that were previously marked as incorrect (is_correct=False). Requires existing results file.",
     )
     parser.add_argument(
         "--only-invalid",
         action="store_true",
-        help="Only run questions that were previously marked as invalid (is_invalid=True). Requires existing results file."
+        help="Only run questions that were previously marked as invalid (is_invalid=True). Requires existing results file.",
     )
     parser.add_argument(
         "--only-ingested",
         action="store_true",
-        help="Only run questions whose memory bank already exists (has been ingested). Automatically skips ingestion. Cannot be combined with --only-failed, --only-invalid, --category, --question-id, or --max-instances-per-category."
+        help="Only run questions whose memory bank already exists (has been ingested). Automatically skips ingestion. Cannot be combined with --only-failed, --only-invalid, --category, --question-id, or --max-instances-per-category.",
     )
     parser.add_argument(
         "--category",
         type=str,
         default=None,
-        help="Filter questions by category/question_type. Available categories: 'single-session-user', 'multi-session', 'single-session-preference', 'temporal-reasoning', 'knowledge-update', 'single-session-assistant'. Can be combined with --max-instances to limit questions within the category."
+        help="Filter questions by category/question_type. Available categories: 'single-session-user', 'multi-session', 'single-session-preference', 'temporal-reasoning', 'knowledge-update', 'single-session-assistant'. Can be combined with --max-instances to limit questions within the category.",
     )
     parser.add_argument(
         "--parallel",
         type=int,
         default=1,
-        help="Number of instances to process in parallel (default: 1 for sequential). Higher values speed up evaluation but use more memory."
+        help="Number of instances to process in parallel (default: 1 for sequential). Higher values speed up evaluation but use more memory.",
     )
     parser.add_argument(
         "--results-filename",
         type=str,
         default="benchmark_results.json",
-        help="Filename for results output (default: benchmark_results.json). Saved in results/ directory."
+        help="Filename for results output (default: benchmark_results.json). Saved in results/ directory.",
     )
     parser.add_argument(
         "--context-format",
         type=str,
         choices=["json", "structured"],
         default="json",
-        help="How to format context for answer generation. 'json' (raw JSON dump, original behavior) or 'structured' (human-readable format with facts grouped with source chunks). Default: json."
+        help="How to format context for answer generation. 'json' (raw JSON dump, original behavior) or 'structured' (human-readable format with facts grouped with source chunks). Default: json.",
     )
     parser.add_argument(
         "--source-results",
         type=str,
         default=None,
-        help="Source results file to read failed/invalid questions from (for --only-failed/--only-invalid). Defaults to benchmark_results.json if not specified."
+        help="Source results file to read failed/invalid questions from (for --only-failed/--only-invalid). Defaults to benchmark_results.json if not specified.",
     )
 
     args = parser.parse_args()
@@ -951,21 +965,23 @@ if __name__ == "__main__":
     if args.max_instances_per_category is not None and (args.max_instances is not None or args.category is not None):
         parser.error("--max-questions-per-category cannot be combined with --max-instances or --category")
 
-    results = asyncio.run(run_benchmark(
-        max_instances=args.max_instances,
-        max_instances_per_category=args.max_instances_per_category,
-        max_questions_per_instance=args.max_questions,
-        thinking_budget=args.thinking_budget,
-        max_tokens=args.max_tokens,
-        skip_ingestion=args.skip_ingestion,
-        filln=args.fill,
-        question_id=args.question_id,
-        only_failed=args.only_failed,
-        only_invalid=args.only_invalid,
-        only_ingested=args.only_ingested,
-        category=args.category,
-        max_concurrent_items=args.parallel,
-        results_filename=args.results_filename,
-        context_format=args.context_format,
-        source_results=args.source_results
-    ))
+    results = asyncio.run(
+        run_benchmark(
+            max_instances=args.max_instances,
+            max_instances_per_category=args.max_instances_per_category,
+            max_questions_per_instance=args.max_questions,
+            thinking_budget=args.thinking_budget,
+            max_tokens=args.max_tokens,
+            skip_ingestion=args.skip_ingestion,
+            filln=args.fill,
+            question_id=args.question_id,
+            only_failed=args.only_failed,
+            only_invalid=args.only_invalid,
+            only_ingested=args.only_ingested,
+            category=args.category,
+            max_concurrent_items=args.parallel,
+            results_filename=args.results_filename,
+            context_format=args.context_format,
+            source_results=args.source_results,
+        )
+    )

hindsight-dev/hindsight_dev/generate_changelog.py

@@ -5,6 +5,7 @@ Generate changelog entry for a new release.
 This script fetches the commit diff between releases, uses an LLM to summarize,
 and prepends the entry to the changelog page.
 """
+
 import argparse
 import json
 import os
@@ -29,6 +30,7 @@ CHANGELOG_PATH = REPO_PATH / "hindsight-docs" / "docs" / "changelog" / "index.md
 
 class ChangelogEntry(BaseModel):
     """A single changelog entry."""
+
     category: str  # "feature", "improvement", "bugfix", "breaking", "other"
     summary: str  # Brief description of the change
     commit_id: str  # Short commit hash
@@ -36,12 +38,14 @@ class ChangelogEntry(BaseModel):
 
 class ChangelogResponse(BaseModel):
     """Structured response from LLM."""
+
     entries: list[ChangelogEntry]
 
 
 @dataclass
 class Commit:
     """Parsed commit from git log."""
+
     hash: str
     message: str
 
@@ -151,10 +155,7 @@ def analyze_commits_with_llm(
     file_diff: str,
 ) -> list[ChangelogEntry]:
     """Use LLM to analyze commits and return structured changelog entries."""
-    commits_json = json.dumps(
-        [{"commit_id": c.hash, "message": c.message} for c in commits],
-        indent=2
-    )
+    commits_json = json.dumps([{"commit_id": c.hash, "message": c.message} for c in commits], indent=2)
 
     prompt = f"""Analyze the following git commits for release {version} of Hindsight (an AI memory system).
 
@@ -252,8 +253,8 @@ For full release details, see [GitHub Releases](https://github.com/vectorize-io/
 
     match = re.search(r"^## ", content, re.MULTILINE)
     if match:
-        header = content[:match.start()].rstrip() + "\n\n"
-        releases = content[match.start():]
+        header = content[: match.start()].rstrip() + "\n\n"
+        releases = content[match.start() :]
     else:
         header = content.rstrip() + "\n\n"
         releases = ""
@@ -283,7 +284,7 @@ def generate_changelog_entry(
     tag = version if version.startswith("v") else f"v{version}"
     display_version = version.lstrip("v")
 
-    console.print(f"[blue]Fetching tags from repository...[/blue]")
+    console.print("[blue]Fetching tags from repository...[/blue]")
     existing_tags = get_git_tags()
 
     if tag not in existing_tags and display_version not in existing_tags:
@@ -300,7 +301,7 @@ def generate_changelog_entry(
     else:
         console.print("[yellow]No previous version found, will include all commits[/yellow]")
 
-    console.print(f"[blue]Getting commits...[/blue]")
+    console.print("[blue]Getting commits...[/blue]")
     commits = get_commits(previous_tag, actual_tag)
     file_diff = get_detailed_diff(previous_tag, actual_tag)
 

hindsight-dev/hindsight_dev/generate_openapi.py

@@ -4,13 +4,15 @@ Generate OpenAPI specification from FastAPI app.
 
 This script imports the FastAPI app and exports its OpenAPI schema to a JSON file.
 """
+
 import json
-import sys
 import os
+import sys
 from pathlib import Path
 
-from hindsight_api.api import create_app
 from hindsight_api import MemoryEngine
+from hindsight_api.api import create_app
+
 
 def generate_openapi_spec(output_path: str = None):
     """Generate OpenAPI spec and save to file."""
@@ -34,7 +36,7 @@ def generate_openapi_spec(output_path: str = None):
 
     # Write to file
     output_file = Path(output_path)
-    with open(output_file, 'w') as f:
+    with open(output_file, "w") as f:
         json.dump(openapi_schema, f, indent=2)
 
     print(f"✓ OpenAPI specification generated: {output_file.absolute()}")
@@ -44,14 +46,15 @@ def generate_openapi_spec(output_path: str = None):
 
     # List endpoints
     print("\n  Endpoints:")
-    for path, methods in openapi_schema['paths'].items():
+    for path, methods in openapi_schema["paths"].items():
         for method in methods.keys():
-            if method.upper() in ['GET', 'POST', 'PUT', 'DELETE', 'PATCH']:
+            if method.upper() in ["GET", "POST", "PUT", "DELETE", "PATCH"]:
                 endpoint_info = methods[method]
-                summary = endpoint_info.get('summary', 'No summary')
-                tags = ', '.join(endpoint_info.get('tags', ['untagged']))
+                summary = endpoint_info.get("summary", "No summary")
+                tags = ", ".join(endpoint_info.get("tags", ["untagged"]))
                 print(f"    {method.upper():6} {path:30} [{tags}] - {summary}")
 
+
 if __name__ == "__main__":
     output = sys.argv[1] if len(sys.argv) > 1 else "openapi.json"
     generate_openapi_spec(output)

hindsight-dev/hindsight_dev/sync_cookbook.py

@@ -212,9 +212,7 @@ This recipe is available as an interactive Jupyter notebook.
         # Insert callout after first heading
         first_heading_match = re.search(r"^(#\s+.+\n)", md_content, re.MULTILINE)
         if first_heading_match:
-            idx = md_content.index(first_heading_match.group(0)) + len(
-                first_heading_match.group(0)
-            )
+            idx = md_content.index(first_heading_match.group(0)) + len(first_heading_match.group(0))
             final_content = md_content[:idx] + "\n" + callout + "\n" + md_content[idx:]
         else:
             final_content = callout + "\n" + md_content
@@ -247,9 +245,7 @@ def process_applications(cookbook_dir: Path, apps_dir: Path) -> list[dict]:
             continue
 
         slug = entry.name
-        title = extract_title_from_readme(readme_path) or " ".join(
-            word.capitalize() for word in slug.split("-")
-        )
+        title = extract_title_from_readme(readme_path) or " ".join(word.capitalize() for word in slug.split("-"))
 
         print(f"  Processing app: {entry.name} → {slug}.md")
 
@@ -275,12 +271,8 @@ This is a complete, runnable application demonstrating Hindsight integration.
         # Insert callout after first heading
         first_heading_match = re.search(r"^(#\s+.+\n)", readme_content, re.MULTILINE)
         if first_heading_match:
-            idx = readme_content.index(first_heading_match.group(0)) + len(
-                first_heading_match.group(0)
-            )
-            final_content = (
-                readme_content[:idx] + "\n" + callout + "\n" + readme_content[idx:]
-            )
+            idx = readme_content.index(first_heading_match.group(0)) + len(first_heading_match.group(0))
+            final_content = readme_content[:idx] + "\n" + callout + "\n" + readme_content[idx:]
         else:
             final_content = callout + "\n" + readme_content
 
@@ -407,26 +399,20 @@ def clean_description(desc: str) -> str:
     return desc
 
 
-def update_cookbook_index(
-    recipes: list[dict], apps: list[dict], docs_dir: Path
-):
+def update_cookbook_index(recipes: list[dict], apps: list[dict], docs_dir: Path):
     """Update cookbook/index.mdx with recipe and app carousels."""
     # Build recipe items for the carousel
     recipe_items = []
     for r in recipes:
         title = r["title"].replace('"', '\\"')
-        recipe_items.append(
-            f'    {{ title: "{title}", href: "/cookbook/recipes/{r["slug"]}" }}'
-        )
+        recipe_items.append(f'    {{ title: "{title}", href: "/cookbook/recipes/{r["slug"]}" }}')
     recipes_json = ",\n".join(recipe_items)
 
     # Build app items for the carousel
     app_items = []
     for a in apps:
         title = a["title"].replace('"', '\\"')
-        app_items.append(
-            f'    {{ title: "{title}", href: "/cookbook/applications/{a["slug"]}" }}'
-        )
+        app_items.append(f'    {{ title: "{title}", href: "/cookbook/applications/{a["slug"]}" }}')
     apps_json = ",\n".join(app_items)
 
     content = f"""---

hindsight-dev/pyproject.toml

@@ -27,3 +27,58 @@ generate-openapi = "hindsight_dev.generate_openapi:generate_openapi_spec"
 generate-changelog = "hindsight_dev.generate_changelog:main"
 sync-cookbook = "hindsight_dev.sync_cookbook:main"
 generate-llms-full = "hindsight_dev.generate_llms_full:main"
+
+[dependency-groups]
+dev = [
+    "ruff>=0.8.0",
+    "ty>=0.0.1",
+]
+
+[tool.ruff]
+line-length = 120
+target-version = "py311"
+
+[tool.ruff.lint]
+select = [
+    "E",      # pycodestyle errors
+    "W",      # pycodestyle warnings
+    "F",      # Pyflakes
+    "I",      # isort
+]
+ignore = [
+    "E501",   # line too long (handled by formatter)
+    "E402",   # module import not at top of file
+    "E712",   # avoid equality comparisons to False (intentional for optional bools)
+    "F401",   # unused import (too noisy during development)
+    "F403",   # star imports (fasthtml uses this pattern)
+    "F405",   # may be undefined from star imports (fasthtml uses this pattern)
+    "F841",   # unused variable (too noisy during development)
+    "F811",   # redefined while unused
+    "F821",   # undefined name (forward references in type hints)
+    "W291",   # trailing whitespace (in multiline strings)
+    "W293",   # blank line contains whitespace (in multiline strings)
+]
+
+[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.rules]
+# Disable noisy rules
+invalid-argument-type = "ignore"    # Too many false positives
+invalid-return-type = "ignore"      # Often intentional
+# missing-argument and unknown-argument are enabled (default)
+invalid-parameter-default = "ignore" # Optional params with None default
+possibly-missing-attribute = "ignore" # Common with Optional types
+unsupported-operator = "ignore"     # Pandas DataFrame operations
+unresolved-reference = "ignore"     # Forward references
+unresolved-import = "ignore"        # Dynamic imports
+invalid-assignment = "ignore"       # Intentional monkey-patching
+no-matching-overload = "ignore"     # Complex generic issues

hindsight-docs/docusaurus.config.ts

@@ -116,7 +116,18 @@ const config: Config = {
     ],
   ],
 
-  themes: ['@docusaurus/theme-mermaid'],
+  themes: [
+    '@docusaurus/theme-mermaid',
+    [
+      '@easyops-cn/docusaurus-search-local',
+      {
+        hashed: true,
+        docsRouteBasePath: '/',
+        indexBlog: false,
+        highlightSearchTermsOnTargetPage: false,
+      },
+    ],
+  ],
 
   themeConfig: {
     ...(ANNOUNCEMENT_BAR && {

hindsight-docs/package.json

@@ -19,6 +19,7 @@
     "@docusaurus/preset-classic": "3.9.2",
     "@docusaurus/theme-common": "^3.9.2",
     "@docusaurus/theme-mermaid": "^3.9.2",
+    "@easyops-cn/docusaurus-search-local": "^0.52.2",
     "@mdx-js/react": "^3.0.0",
     "clsx": "^2.0.0",
     "prism-react-renderer": "^2.3.0",

hindsight-docs/src/css/custom.css

@@ -890,3 +890,309 @@ div[class*="announcementBar"] a:hover {
   background: rgba(255, 255, 255, 0.3);
 }
 
+/* ===== Local Search Plugin Styling ===== */
+
+/* Search bar in navbar */
+.navbar__search {
+  margin-left: auto;
+}
+
+
+/* Hide keyboard shortcut hints */
+[class*="searchHint"],
+[class*="searchHintContainer"] {
+  display: none !important;
+}
+
+/* Minimal search button */
+[class*="searchBarContainer"] button,
+[class*="searchBar_"] {
+  font-family: 'Space Grotesk', -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif !important;
+  font-size: 0.75rem !important;
+  font-weight: 500 !important;
+  background: transparent !important;
+  border: none !important;
+  border-bottom: 1px solid transparent !important;
+  border-image: linear-gradient(90deg, #0074d9, #009296) 1 !important;
+  border-radius: 0 !important;
+  padding: 0.25rem 0 !important;
+  color: var(--ifm-color-emphasis-600) !important;
+  transition: all 0.2s ease !important;
+  height: auto !important;
+  min-height: unset !important;
+}
+
+[class*="searchBarContainer"] button:hover,
+[class*="searchBar_"]:hover {
+  color: var(--ifm-font-color-base) !important;
+}
+
+[data-theme='dark'] [class*="searchBarContainer"] button,
+[data-theme='dark'] [class*="searchBar_"] {
+  color: var(--ifm-color-emphasis-500) !important;
+}
+
+[data-theme='dark'] [class*="searchBarContainer"] button:hover,
+[data-theme='dark'] [class*="searchBar_"]:hover {
+  color: #e2e8f0 !important;
+}
+
+/* Search input field - remove ALL borders */
+[class*="searchInput"],
+[class*="searchQueryInput"],
+[class*="searchInput"] *,
+[class*="searchQueryInput"] *,
+[class*="searchBarContainer"] input,
+[class*="searchBarContainer"] input * {
+  font-family: 'Space Grotesk', -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif !important;
+  font-size: 0.875rem !important;
+  font-weight: 500 !important;
+  background: transparent !important;
+  border: 0 !important;
+  border-top: 0 !important;
+  border-right: 0 !important;
+  border-bottom: 0 !important;
+  border-left: 0 !important;
+  border-radius: 0 !important;
+  color: var(--ifm-font-color-base) !important;
+  padding: 0.5rem 0 !important;
+  box-shadow: none !important;
+  outline: none !important;
+}
+
+/* Add single underline via pseudo or direct */
+[class*="searchQueryInput"],
+[class*="searchBarContainer"] input[type="search"] {
+  border-bottom: 1px solid var(--ifm-toc-border-color) !important;
+}
+
+[data-theme='dark'] [class*="searchQueryInput"],
+[data-theme='dark'] [class*="searchBarContainer"] input[type="search"] {
+  border-bottom: 1px solid #27272a !important;
+}
+
+[class*="searchQueryInput"]:focus,
+[class*="searchBarContainer"] input[type="search"]:focus {
+  border-bottom: 1px solid #0074d9 !important;
+}
+
+[class*="searchInput"]::placeholder,
+[class*="searchQueryInput"]::placeholder,
+[class*="searchBarContainer"] input::placeholder {
+  font-family: 'Space Grotesk', -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif !important;
+  font-weight: 500 !important;
+  color: var(--ifm-color-emphasis-400) !important;
+}
+
+/* Minimal dropdown */
+[class*="dropdownMenu"],
+[class*="suggestionsContainer"],
+[class*="searchResultsContainer"],
+[class*="suggestions"],
+ul[class*="suggestion"] {
+  font-family: 'Inter', -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif !important;
+  background: #ffffff !important;
+  background-color: #ffffff !important;
+  border: none !important;
+  border-radius: 0.25rem !important;
+  box-shadow: 0 4px 20px -4px rgba(0, 0, 0, 0.15) !important;
+  padding: 0.25rem !important;
+}
+
+[data-theme='dark'] [class*="dropdownMenu"],
+[data-theme='dark'] [class*="suggestionsContainer"],
+[data-theme='dark'] [class*="searchResultsContainer"],
+[data-theme='dark'] [class*="suggestions"],
+[data-theme='dark'] ul[class*="suggestion"] {
+  background: #0f0f11 !important;
+  background-color: #0f0f11 !important;
+  box-shadow: 0 4px 20px -4px rgba(0, 0, 0, 0.5) !important;
+}
+
+/* Minimal result items */
+[class*="searchResultItem"],
+[class*="suggestion"] {
+  font-family: 'Inter', -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif !important;
+  background: transparent !important;
+  border-radius: 0.25rem !important;
+  padding: 0.5rem 0.75rem !important;
+  margin: 0 !important;
+  transition: background-color 0.1s ease !important;
+  cursor: pointer !important;
+  border-left: 2px solid transparent !important;
+}
+
+/* No background hover - keep transparent */
+[class*="searchResultItem"]:hover,
+[class*="searchResultItem"][class*="cursor"],
+[class*="suggestion"]:hover,
+[class*="suggestion"][class*="cursor"] {
+  background: transparent !important;
+  border-left-color: transparent !important;
+}
+
+/* Icons - default color */
+[class*="hitIcon"],
+[class*="hitTree"] {
+  color: #71717a !important;
+  transition: all 0.15s ease !important;
+}
+
+[class*="hitIcon"] svg path,
+[class*="hitIcon"] svg,
+[class*="hitTree"] svg path,
+[class*="hitTree"] svg {
+  stroke: #71717a !important;
+  color: #71717a !important;
+  transition: all 0.15s ease !important;
+}
+
+/* Icons - hover gradient effect */
+[class*="suggestion"]:hover [class*="hitIcon"] svg path,
+[class*="suggestion"]:hover [class*="hitIcon"] svg,
+[class*="suggestion"]:hover [class*="hitTree"] svg path,
+[class*="suggestion"]:hover [class*="hitTree"] svg {
+  stroke: #0074d9 !important;
+  color: #0074d9 !important;
+}
+
+[data-theme='dark'] [class*="hitIcon"] svg path,
+[data-theme='dark'] [class*="hitIcon"] svg,
+[data-theme='dark'] [class*="hitTree"] svg path,
+[data-theme='dark'] [class*="hitTree"] svg {
+  stroke: #71717a !important;
+  color: #71717a !important;
+}
+
+[data-theme='dark'] [class*="suggestion"]:hover [class*="hitIcon"] svg path,
+[data-theme='dark'] [class*="suggestion"]:hover [class*="hitIcon"] svg,
+[data-theme='dark'] [class*="suggestion"]:hover [class*="hitTree"] svg path,
+[data-theme='dark'] [class*="suggestion"]:hover [class*="hitTree"] svg {
+  stroke: #3396e8 !important;
+  color: #3396e8 !important;
+}
+
+/* === PAGE-LEVEL RESULTS (doc icon, no tree) === */
+/* Title is the page name - GRADIENT */
+[class*="suggestion"]:not(:has([class*="hitTree"])) [class*="hitTitle"] {
+  font-family: 'Space Grotesk', -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif !important;
+  font-weight: 600 !important;
+  font-size: 0.9375rem !important;
+  background: linear-gradient(90deg, #0074d9, #009296) !important;
+  -webkit-background-clip: text !important;
+  background-clip: text !important;
+  -webkit-text-fill-color: transparent !important;
+  margin: 0 !important;
+}
+
+/* === SECTION-LEVEL RESULTS (hashtag icon, has tree) === */
+/* Title is the section name - regular gray text */
+[class*="suggestion"]:has([class*="hitTree"]) [class*="hitTitle"] {
+  font-family: 'Inter', -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif !important;
+  font-weight: 500 !important;
+  font-size: 0.8125rem !important;
+  background: none !important;
+  -webkit-background-clip: unset !important;
+  background-clip: unset !important;
+  color: #3f3f46 !important;
+  -webkit-text-fill-color: #3f3f46 !important;
+  margin: 0 !important;
+}
+
+[data-theme='dark'] [class*="suggestion"]:has([class*="hitTree"]) [class*="hitTitle"] {
+  color: #d4d4d8 !important;
+  -webkit-text-fill-color: #d4d4d8 !important;
+}
+
+/* Path is the page name - GRADIENT */
+[class*="suggestion"]:has([class*="hitTree"]) [class*="hitPath"] {
+  font-family: 'Space Grotesk', -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif !important;
+  font-weight: 600 !important;
+  font-size: 0.75rem !important;
+  background: linear-gradient(90deg, #0074d9, #009296) !important;
+  -webkit-background-clip: text !important;
+  background-clip: text !important;
+  -webkit-text-fill-color: transparent !important;
+  text-transform: none !important;
+  letter-spacing: 0 !important;
+  margin-top: 0.125rem !important;
+  display: block !important;
+}
+
+/* === FALLBACK for browsers without :has() === */
+[class*="hitPath"] {
+  font-family: 'Inter', -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif !important;
+  font-size: 0.625rem !important;
+  font-weight: 400 !important;
+  color: #71717a !important;
+  -webkit-text-fill-color: #71717a !important;
+  text-transform: none !important;
+  letter-spacing: 0 !important;
+}
+
+[class*="hitTitle"] {
+  font-family: 'Space Grotesk', -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif !important;
+  font-weight: 600 !important;
+  font-size: 0.9375rem !important;
+  background: linear-gradient(90deg, #0074d9, #009296) !important;
+  -webkit-background-clip: text !important;
+  background-clip: text !important;
+  -webkit-text-fill-color: transparent !important;
+  margin: 0 !important;
+}
+
+/* Highlight matched text - must override color rules */
+[class*="searchResultItem"] mark,
+[class*="suggestion"] mark,
+.search-result-match mark,
+[class*="searchResultItem"]:hover mark,
+[class*="suggestion"]:hover mark,
+[class*="searchResultItem"][class*="cursor"] mark,
+[class*="suggestion"][class*="cursor"] mark {
+  background: transparent !important;
+  color: #0074d9 !important;
+  padding: 0 !important;
+  border-radius: 0 !important;
+  font-weight: 600 !important;
+  text-decoration: none !important;
+}
+
+[data-theme='dark'] [class*="searchResultItem"] mark,
+[data-theme='dark'] [class*="suggestion"] mark,
+[data-theme='dark'] .search-result-match mark,
+[data-theme='dark'] [class*="searchResultItem"]:hover mark,
+[data-theme='dark'] [class*="suggestion"]:hover mark,
+[data-theme='dark'] [class*="searchResultItem"][class*="cursor"] mark,
+[data-theme='dark'] [class*="suggestion"][class*="cursor"] mark {
+  color: #66b3f0 !important;
+}
+
+/* No results */
+[class*="noResults"],
+[class*="searchNoResult"] {
+  font-family: 'Inter', -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif !important;
+  font-size: 0.75rem !important;
+  color: var(--ifm-color-emphasis-500) !important;
+  padding: 1rem !important;
+  text-align: center !important;
+}
+
+/* Loading */
+[class*="searchIndexLoading"],
+[class*="loadingRing"] {
+  font-family: 'Inter', -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif !important;
+  font-size: 0.75rem !important;
+  color: var(--ifm-color-emphasis-400) !important;
+}
+
+/* Hide footer */
+[class*="hitFooter"],
+[class*="searchFooter"] {
+  display: none !important;
+}
+
+/* Hide action icons */
+[class*="hitAction"] {
+  display: none !important;
+}
+

scripts/hooks/lint-node.sh

@@ -1,22 +0,0 @@
-#!/bin/bash
-# Lint Node/TypeScript code with ESLint and Prettier
-
-REPO_ROOT="$(git rev-parse --show-toplevel)"
-cd "$REPO_ROOT/hindsight-control-plane"
-
-echo "  Linting Node/TS with ESLint and Prettier..."
-
-# Capture output and only show on failure
-OUTPUT=$(npx eslint --fix "src/**/*.{ts,tsx}" 2>&1)
-if [ $? -ne 0 ]; then
-    echo "$OUTPUT"
-    exit 1
-fi
-
-OUTPUT=$(npx prettier --write "src/**/*.{ts,tsx}" 2>&1)
-if [ $? -ne 0 ]; then
-    echo "$OUTPUT"
-    exit 1
-fi
-
-echo "  Node lint complete"

scripts/hooks/lint-python.sh

@@ -1,22 +0,0 @@
-#!/bin/bash
-# Lint Python code with Ruff
-
-REPO_ROOT="$(git rev-parse --show-toplevel)"
-cd "$REPO_ROOT/hindsight-api"
-
-echo "  Linting Python with Ruff..."
-
-# Capture output and only show on failure
-OUTPUT=$(uv run ruff check --fix . 2>&1)
-if [ $? -ne 0 ]; then
-    echo "$OUTPUT"
-    exit 1
-fi
-
-OUTPUT=$(uv run ruff format . 2>&1)
-if [ $? -ne 0 ]; then
-    echo "$OUTPUT"
-    exit 1
-fi
-
-echo "  Python lint complete"

scripts/hooks/lint.sh

@@ -0,0 +1,76 @@
+#!/bin/bash
+# Parallel linting for all code (Node, Python)
+# Runs all linting tasks concurrently for faster execution
+
+set -e
+
+REPO_ROOT="$(git rev-parse --show-toplevel)"
+TEMP_DIR=$(mktemp -d)
+trap "rm -rf $TEMP_DIR" EXIT
+
+# Track all background jobs
+declare -a PIDS
+declare -a NAMES
+
+run_task() {
+    local name="$1"
+    local dir="$2"
+    shift 2
+    local cmd="$@"
+
+    (
+        cd "$dir"
+        if OUTPUT=$($cmd 2>&1); then
+            echo "OK" > "$TEMP_DIR/$name.status"
+        else
+            echo "FAIL" > "$TEMP_DIR/$name.status"
+            echo "$OUTPUT" > "$TEMP_DIR/$name.output"
+        fi
+    ) &
+    PIDS+=($!)
+    NAMES+=("$name")
+}
+
+echo "  Running lints in parallel..."
+
+# Node/TypeScript tasks
+run_task "eslint" "$REPO_ROOT/hindsight-control-plane" "npx eslint --fix src/**/*.{ts,tsx}"
+run_task "prettier" "$REPO_ROOT/hindsight-control-plane" "npx prettier --write src/**/*.{ts,tsx}"
+
+# Python hindsight-api tasks
+run_task "ruff-api-check" "$REPO_ROOT/hindsight-api" "uv run ruff check --fix ."
+run_task "ruff-api-format" "$REPO_ROOT/hindsight-api" "uv run ruff format ."
+run_task "ty-api" "$REPO_ROOT/hindsight-api" "uv run ty check hindsight_api"
+
+# Python hindsight-dev tasks
+run_task "ruff-dev-check" "$REPO_ROOT/hindsight-dev" "uv run ruff check --fix ."
+run_task "ruff-dev-format" "$REPO_ROOT/hindsight-dev" "uv run ruff format ."
+run_task "ty-dev" "$REPO_ROOT/hindsight-dev" "uv run ty check hindsight_dev benchmarks"
+
+# Wait for all tasks to complete
+for pid in "${PIDS[@]}"; do
+    wait "$pid" 2>/dev/null || true
+done
+
+# Check results
+FAILED=0
+for name in "${NAMES[@]}"; do
+    if [ -f "$TEMP_DIR/$name.status" ]; then
+        STATUS=$(cat "$TEMP_DIR/$name.status")
+        if [ "$STATUS" = "FAIL" ]; then
+            echo ""
+            echo "  ❌ $name failed:"
+            cat "$TEMP_DIR/$name.output"
+            FAILED=1
+        fi
+    else
+        echo "  ❌ $name: no status (crashed?)"
+        FAILED=1
+    fi
+done
+
+if [ $FAILED -eq 1 ]; then
+    exit 1
+fi
+
+echo "  All lints passed ✓"

uv.lock

@@ -1216,6 +1216,7 @@ dev = [
     { name = "pytest-xdist" },
     { name = "python-dotenv" },
     { name = "ruff" },
+    { name = "ty" },
 ]
 
 [package.metadata]
@@ -1265,6 +1266,7 @@ dev = [
     { name = "pytest-xdist", specifier = ">=3.8.0" },
     { name = "python-dotenv", specifier = ">=1.2.1" },
     { name = "ruff", specifier = ">=0.8.0" },
+    { name = "ty", specifier = ">=0.0.1" },
 ]
 
 [[package]]
@@ -1314,6 +1316,12 @@ dependencies = [
     { name = "streamlit" },
 ]
 
+[package.dev-dependencies]
+dev = [
+    { name = "ruff" },
+    { name = "ty" },
+]
+
 [package.metadata]
 requires-dist = [
     { name = "hindsight-api", editable = "hindsight-api" },
@@ -1324,6 +1332,12 @@ requires-dist = [
     { name = "streamlit", specifier = ">=1.51.0" },
 ]
 
+[package.metadata.requires-dev]
+dev = [
+    { name = "ruff", specifier = ">=0.8.0" },
+    { name = "ty", specifier = ">=0.0.1" },
+]
+
 [[package]]
 name = "httpcore"
 version = "1.0.9"
@@ -4266,6 +4280,31 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/78/eb/65f5ba83c2a123f6498a3097746607e5b2f16add29e36765305e4ac7fdd8/triton-3.1.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:c8182f42fd8080a7d39d666814fa36c5e30cc00ea7eeeb1a2983dbb4c99a0fdc", size = 209551444 },
 ]
 
+[[package]]
+name = "ty"
+version = "0.0.5"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/9e/db/6299d478000f4f1c6f9bf2af749359381610ffc4cbe6713b66e436ecf6e7/ty-0.0.5.tar.gz", hash = "sha256:983da6330773ff71e2b249810a19c689f9a0372f6e21bbf7cde37839d05b4346", size = 4806218 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/7c/98/c1f61ba378b4191e641bb36c07b7fcc70ff844d61be7a4bf2fea7472b4a9/ty-0.0.5-py3-none-linux_armv6l.whl", hash = "sha256:1594cd9bb68015eb2f5a3c68a040860f3c9306dc6667d7a0e5f4df9967b460e2", size = 9785554 },
+    { url = "https://files.pythonhosted.org/packages/ab/f9/b37b77c03396bd779c1397dae4279b7ad79315e005b3412feed8812a4256/ty-0.0.5-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:7c0140ba980233d28699d9ddfe8f43d0b3535d6a3bbff9935df625a78332a3cf", size = 9603995 },
+    { url = "https://files.pythonhosted.org/packages/7d/70/4e75c11903b0e986c0203040472627cb61d6a709e1797fb08cdf9d565743/ty-0.0.5-py3-none-macosx_11_0_arm64.whl", hash = "sha256:15de414712cde92048ae4b1a77c4dc22920bd23653fe42acaf73028bad88f6b9", size = 9145815 },
+    { url = "https://files.pythonhosted.org/packages/89/05/93983dfcf871a41dfe58e5511d28e6aa332a1f826cc67333f77ae41a2f8a/ty-0.0.5-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:438aa51ad6c5fae64191f8d58876266e26f9250cf09f6624b6af47a22fa88618", size = 9619849 },
+    { url = "https://files.pythonhosted.org/packages/82/b6/896ab3aad59f846823f202e94be6016fb3f72434d999d2ae9bd0f28b3af9/ty-0.0.5-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:1b3d373fd96af1564380caf153600481c676f5002ee76ba8a7c3508cdff82ee0", size = 9606611 },
+    { url = "https://files.pythonhosted.org/packages/ca/ae/098e33fc92330285ed843e2750127e896140c4ebd2d73df7732ea496f588/ty-0.0.5-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:8453692503212ad316cf8b99efbe85a91e5f63769c43be5345e435a1b16cba5a", size = 10029523 },
+    { url = "https://files.pythonhosted.org/packages/04/5a/f4b4c33758b9295e9aca0de9645deca0f4addd21d38847228723a6e780fc/ty-0.0.5-py3-none-manylinux_2_17_ppc64.manylinux2014_ppc64.whl", hash = "sha256:2e4c454139473abbd529767b0df7a795ed828f780aef8d0d4b144558c0dc4446", size = 10870892 },
+    { url = "https://files.pythonhosted.org/packages/c3/c5/4e3e7e88389365aa1e631c99378711cf0c9d35a67478cb4720584314cf44/ty-0.0.5-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:426d4f3b82475b1ec75f3cc9ee5a667c8a4ae8441a09fcd8e823a53b706d00c7", size = 10599291 },
+    { url = "https://files.pythonhosted.org/packages/c1/5d/138f859ea87bd95e17b9818e386ae25a910e46521c41d516bf230ed83ffc/ty-0.0.5-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:5710817b67c6b2e4c0224e4f319b7decdff550886e9020f6d46aa1ce8f89a609", size = 10413515 },
+    { url = "https://files.pythonhosted.org/packages/27/21/1cbcd0d3b1182172f099e88218137943e0970603492fb10c7c9342369d9a/ty-0.0.5-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:e23c55ef08882c7c5ced1ccb90b4eeefa97f690aea254f58ac0987896c590f76", size = 10144992 },
+    { url = "https://files.pythonhosted.org/packages/ad/30/fdac06a5470c09ad2659a0806497b71f338b395d59e92611f71b623d05a0/ty-0.0.5-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:b9e4c1a28a23b14cf8f4f793f4da396939f16c30bfa7323477c8cc234e352ac4", size = 9606408 },
+    { url = "https://files.pythonhosted.org/packages/09/93/e99dcd7f53295192d03efd9cbcec089a916f49cad4935c0160ea9adbd53d/ty-0.0.5-py3-none-musllinux_1_2_armv7l.whl", hash = "sha256:4e9ebb61529b9745af662e37c37a01ad743cdd2c95f0d1421705672874d806cd", size = 9630040 },
+    { url = "https://files.pythonhosted.org/packages/d7/f8/6d1e87186e4c35eb64f28000c1df8fd5f73167ce126c5e3dd21fd1204a23/ty-0.0.5-py3-none-musllinux_1_2_i686.whl", hash = "sha256:5eb191a8e332f50f56dfe45391bdd7d43dd4ef6e60884710fd7ce84c5d8c1eb5", size = 9754016 },
+    { url = "https://files.pythonhosted.org/packages/28/e6/20f989342cb3115852dda404f1d89a10a3ce93f14f42b23f095a3d1a00c9/ty-0.0.5-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:92ed7451a1e82ee134a2c24ca43b74dd31e946dff2b08e5c34473e6b051de542", size = 10252877 },
+    { url = "https://files.pythonhosted.org/packages/57/9d/fc66fa557443233dfad9ae197ff3deb70ae0efcfb71d11b30ef62f5cdcc3/ty-0.0.5-py3-none-win32.whl", hash = "sha256:71f6707e4c1c010c158029a688a498220f28bb22fdb6707e5c20e09f11a5e4f2", size = 9212640 },
+    { url = "https://files.pythonhosted.org/packages/68/b6/05c35f6dea29122e54af0e9f8dfedd0a100c721affc8cc801ebe2bc2ed13/ty-0.0.5-py3-none-win_amd64.whl", hash = "sha256:2b8b754a0d7191e94acdf0c322747fec34371a4d0669f5b4e89549aef28814ae", size = 10034701 },
+    { url = "https://files.pythonhosted.org/packages/df/ca/4201ed5cb2af73912663d0c6ded927c28c28b3c921c9348aa8d2cfef4853/ty-0.0.5-py3-none-win_arm64.whl", hash = "sha256:83bea5a5296caac20d52b790ded2b830a7ff91c4ed9f36730fe1f393ceed6654", size = 9566474 },
+]
+
 [[package]]
 name = "typer"
 version = "0.20.0"