hindsight-api 0.4.6__py3-none-any.whl → 0.4.8__py3-none-any.whl

hindsight_api/engine/memory_engine.py

@@ -303,8 +303,10 @@ class MemoryEngine(MemoryEngineInterface):
         db_url = db_url or config.database_url
         memory_llm_provider = memory_llm_provider or config.llm_provider
         memory_llm_api_key = memory_llm_api_key or config.llm_api_key
-        # Ollama and mock don't require an API key
-        if not memory_llm_api_key and memory_llm_provider not in ("ollama", "mock"):
+        # Ollama, openai-codex, claude-code, and mock don't require an API key
+        # openai-codex uses OAuth tokens from ~/.codex/auth.json
+        # claude-code uses OAuth tokens from macOS Keychain
+        if not memory_llm_api_key and memory_llm_provider not in ("ollama", "openai-codex", "claude-code", "mock"):
             raise ValueError("LLM API key is required. Set HINDSIGHT_API_LLM_API_KEY environment variable.")
         memory_llm_model = memory_llm_model or config.llm_model
         memory_llm_base_url = memory_llm_base_url or config.get_llm_base_url() or None
@@ -457,7 +459,11 @@ class MemoryEngine(MemoryEngineInterface):
         # Store operation validator extension (optional)
         self._operation_validator = operation_validator
 
-        # Store tenant extension (optional)
+        # Store tenant extension (always set, use default if none provided)
+        if tenant_extension is None:
+            from ..extensions.builtin.tenant import DefaultTenantExtension
+
+            tenant_extension = DefaultTenantExtension(config={})
         self._tenant_extension = tenant_extension
 
     async def _validate_operation(self, validation_coro) -> None:
@@ -495,22 +501,18 @@ class MemoryEngine(MemoryEngineInterface):
         Raises:
             AuthenticationError: If authentication fails or request_context is missing when required.
         """
-        if self._tenant_extension is None:
-            _current_schema.set("public")
-            return "public"
-
         from hindsight_api.extensions import AuthenticationError
 
         if request_context is None:
-            raise AuthenticationError("RequestContext is required when tenant extension is configured")
+            raise AuthenticationError("RequestContext is required")
 
         # For internal/background operations (e.g., worker tasks), skip extension authentication.
         # The task was already authenticated at submission time, and execute_task sets _current_schema
-        # from the task's _schema field. For public schema tasks, _current_schema keeps its default "public".
+        # from the task's _schema field.
         if request_context.internal:
             return _current_schema.get()
 
-        # Let AuthenticationError propagate - HTTP layer will convert to 401
+        # Authenticate through tenant extension (always set, may be default no-auth extension)
         tenant_context = await self._tenant_extension.authenticate(request_context)
 
         _current_schema.set(tenant_context.schema_name)
@@ -536,10 +538,15 @@ class MemoryEngine(MemoryEngineInterface):
             f"[BATCH_RETAIN_TASK] Starting background batch retain for bank_id={bank_id}, {len(contents)} items"
         )
 
-        # Use internal request context for background tasks (skips tenant auth when schema is pre-set)
+        # Restore tenant_id/api_key_id from task payload so downstream operations
+        # (e.g., consolidation and mental model refreshes) can attribute usage.
         from hindsight_api.models import RequestContext
 
-        internal_context = RequestContext(internal=True)
+        internal_context = RequestContext(
+            internal=True,
+            tenant_id=task_dict.get("_tenant_id"),
+            api_key_id=task_dict.get("_api_key_id"),
+        )
         await self.retain_batch_async(bank_id=bank_id, contents=contents, request_context=internal_context)
 
         logger.info(f"[BATCH_RETAIN_TASK] Completed background batch retain for bank_id={bank_id}")
@@ -565,7 +572,13 @@ class MemoryEngine(MemoryEngineInterface):
 
         from .consolidation import run_consolidation_job
 
-        internal_context = RequestContext(internal=True)
+        # Restore tenant_id/api_key_id from task payload so downstream operations
+        # (e.g., mental model refreshes) can attribute usage to the correct org.
+        internal_context = RequestContext(
+            internal=True,
+            tenant_id=task_dict.get("_tenant_id"),
+            api_key_id=task_dict.get("_api_key_id"),
+        )
         result = await run_consolidation_job(
             memory_engine=self,
             bank_id=bank_id,
@@ -597,7 +610,13 @@ class MemoryEngine(MemoryEngineInterface):
 
         from hindsight_api.models import RequestContext
 
-        internal_context = RequestContext(internal=True)
+        # Restore tenant_id/api_key_id from task payload so extensions can
+        # attribute the mental_model_refresh operation to the correct org.
+        internal_context = RequestContext(
+            internal=True,
+            tenant_id=task_dict.get("_tenant_id"),
+            api_key_id=task_dict.get("_api_key_id"),
+        )
 
         # Get the current mental model to get source_query
         mental_model = await self.get_mental_model(bank_id, mental_model_id, request_context=internal_context)
@@ -641,6 +660,42 @@ class MemoryEngine(MemoryEngineInterface):
             request_context=internal_context,
         )
 
+        # Call post-operation hook if validator is configured
+        if self._operation_validator:
+            from hindsight_api.extensions.operation_validator import MentalModelRefreshResult
+
+            # Count facts and mental models from based_on
+            facts_used = 0
+            mental_models_used = 0
+            if reflect_result.based_on:
+                for fact_type, facts in reflect_result.based_on.items():
+                    if facts:
+                        if fact_type == "mental_models":
+                            mental_models_used += len(facts)
+                        else:
+                            facts_used += len(facts)
+
+            # Estimate tokens
+            query_tokens = len(source_query) // 4 if source_query else 0
+            output_tokens = len(generated_content) // 4 if generated_content else 0
+            context_tokens = 0  # refresh doesn't use additional context
+
+            result_ctx = MentalModelRefreshResult(
+                bank_id=bank_id,
+                mental_model_id=mental_model_id,
+                request_context=internal_context,
+                query_tokens=query_tokens,
+                output_tokens=output_tokens,
+                context_tokens=context_tokens,
+                facts_used=facts_used,
+                mental_models_used=mental_models_used,
+                success=True,
+            )
+            try:
+                await self._operation_validator.on_mental_model_refresh_complete(result_ctx)
+            except Exception as hook_err:
+                logger.warning(f"Post-mental-model-refresh hook error (non-fatal): {hook_err}")
+
         logger.info(f"[REFRESH_MENTAL_MODEL_TASK] Completed for bank_id={bank_id}, mental_model_id={mental_model_id}")
 
     async def execute_task(self, task_dict: dict[str, Any]):
@@ -884,30 +939,34 @@ class MemoryEngine(MemoryEngineInterface):
 
             if not self.db_url:
                 raise ValueError("Database URL is required for migrations")
-            logger.info("Running database migrations...")
-            # Use configured database schema for migrations (defaults to "public")
-            run_migrations(self.db_url, schema=get_config().database_schema)
 
-            # Migrate all existing tenant schemas (if multi-tenant)
-            if self._tenant_extension is not None:
-                try:
-                    tenants = await self._tenant_extension.list_tenants()
-                    if tenants:
-                        logger.info(f"Running migrations on {len(tenants)} tenant schemas...")
-                        for tenant in tenants:
-                            schema = tenant.schema
-                            if schema and schema != "public":
-                                try:
-                                    run_migrations(self.db_url, schema=schema)
-                                except Exception as e:
-                                    logger.warning(f"Failed to migrate tenant schema {schema}: {e}")
-                        logger.info("Tenant schema migrations completed")
-                except Exception as e:
-                    logger.warning(f"Failed to run tenant schema migrations: {e}")
-
-            # Ensure embedding column dimension matches the model's dimension
-            # This is done after migrations and after embeddings.initialize()
-            ensure_embedding_dimension(self.db_url, self.embeddings.dimension, schema=get_config().database_schema)
+            # Migrate all schemas from the tenant extension
+            # The tenant extension is the single source of truth for which schemas exist
+            logger.info("Running database migrations...")
+            try:
+                tenants = await self._tenant_extension.list_tenants()
+                if tenants:
+                    logger.info(f"Running migrations on {len(tenants)} schema(s)...")
+                    for tenant in tenants:
+                        schema = tenant.schema
+                        if schema:
+                            try:
+                                run_migrations(self.db_url, schema=schema)
+                            except Exception as e:
+                                logger.warning(f"Failed to migrate schema {schema}: {e}")
+                    logger.info("Schema migrations completed")
+
+                    # Ensure embedding column dimension matches the model's dimension
+                    # This is done after migrations and after embeddings.initialize()
+                    for tenant in tenants:
+                        schema = tenant.schema
+                        if schema:
+                            try:
+                                ensure_embedding_dimension(self.db_url, self.embeddings.dimension, schema=schema)
+                            except Exception as e:
+                                logger.warning(f"Failed to ensure embedding dimension for schema {schema}: {e}")
+            except Exception as e:
+                logger.warning(f"Failed to run schema migrations: {e}")
 
         logger.info(f"Connecting to PostgreSQL at {self.db_url}")
 
@@ -5416,6 +5475,13 @@ class MemoryEngine(MemoryEngineInterface):
         task_payload: dict[str, Any] = {"contents": contents}
         if document_tags:
             task_payload["document_tags"] = document_tags
+        # Pass tenant_id and api_key_id through task payload so the worker
+        # can propagate request context to downstream operations (e.g.,
+        # consolidation and mental model refreshes triggered after retain).
+        if request_context.tenant_id:
+            task_payload["_tenant_id"] = request_context.tenant_id
+        if request_context.api_key_id:
+            task_payload["_api_key_id"] = request_context.api_key_id
 
         result = await self._submit_async_operation(
             bank_id=bank_id,
@@ -5448,11 +5514,21 @@ class MemoryEngine(MemoryEngineInterface):
             Dict with operation_id
         """
         await self._authenticate_tenant(request_context)
+
+        # Pass tenant_id and api_key_id through task payload so the worker
+        # can provide request context to extension hooks (e.g., usage metering
+        # for mental model refreshes triggered by consolidation).
+        task_payload: dict[str, Any] = {}
+        if request_context.tenant_id:
+            task_payload["_tenant_id"] = request_context.tenant_id
+        if request_context.api_key_id:
+            task_payload["_api_key_id"] = request_context.api_key_id
+
         return await self._submit_async_operation(
             bank_id=bank_id,
             operation_type="consolidation",
             task_type="consolidation",
-            task_payload={},
+            task_payload=task_payload,
             dedupe_by_bank=True,
         )
 
@@ -5482,13 +5558,21 @@ class MemoryEngine(MemoryEngineInterface):
         if not mental_model:
             raise ValueError(f"Mental model {mental_model_id} not found in bank {bank_id}")
 
+        # Pass tenant_id and api_key_id through task payload so the worker
+        # can provide request context to extension hooks.
+        task_payload: dict[str, Any] = {
+            "mental_model_id": mental_model_id,
+        }
+        if request_context.tenant_id:
+            task_payload["_tenant_id"] = request_context.tenant_id
+        if request_context.api_key_id:
+            task_payload["_api_key_id"] = request_context.api_key_id
+
         return await self._submit_async_operation(
             bank_id=bank_id,
             operation_type="refresh_mental_model",
             task_type="refresh_mental_model",
-            task_payload={
-                "mental_model_id": mental_model_id,
-            },
+            task_payload=task_payload,
             result_metadata={"mental_model_id": mental_model_id, "name": mental_model["name"]},
             dedupe_by_bank=False,
         )

hindsight_api/engine/providers/__init__.py

@@ -0,0 +1,14 @@
+"""
+LLM provider implementations.
+
+This package contains concrete implementations of the LLMInterface for various providers.
+"""
+
+from .anthropic_llm import AnthropicLLM
+from .claude_code_llm import ClaudeCodeLLM
+from .codex_llm import CodexLLM
+from .gemini_llm import GeminiLLM
+from .mock_llm import MockLLM
+from .openai_compatible_llm import OpenAICompatibleLLM
+
+__all__ = ["AnthropicLLM", "ClaudeCodeLLM", "CodexLLM", "GeminiLLM", "MockLLM", "OpenAICompatibleLLM"]