diff --git a/hindsight-api/hindsight_api/api/http.py b/hindsight-api/hindsight_api/api/http.py index 9acdf0a0..e128de00 100644 --- a/hindsight-api/hindsight_api/api/http.py +++ b/hindsight-api/hindsight_api/api/http.py @@ -2290,6 +2290,23 @@ def _register_routes(app: FastAPI): ): """Get a mental model by ID.""" try: + # Pre-operation validation hook + validator = app.state.memory._operation_validator + if validator: + from hindsight_api.extensions.operation_validator import MentalModelGetContext + + ctx = MentalModelGetContext( + bank_id=bank_id, + mental_model_id=mental_model_id, + request_context=request_context, + ) + validation = await validator.validate_mental_model_get(ctx) + if not validation.allowed: + raise OperationValidationError( + validation.reason or "Operation not allowed", + status_code=validation.status_code, + ) + mental_model = await app.state.memory.get_mental_model( bank_id=bank_id, mental_model_id=mental_model_id, @@ -2297,9 +2314,31 @@ def _register_routes(app: FastAPI): ) if mental_model is None: raise HTTPException(status_code=404, detail=f"Mental model '{mental_model_id}' not found") + + # Post-operation hook + if validator: + from hindsight_api.extensions.operation_validator import MentalModelGetResult + + content = mental_model.get("content", "") + output_tokens = len(content) // 4 if content else 0 + + result_ctx = MentalModelGetResult( + bank_id=bank_id, + mental_model_id=mental_model_id, + request_context=request_context, + output_tokens=output_tokens, + success=True, + ) + try: + await validator.on_mental_model_get_complete(result_ctx) + except Exception as hook_err: + logger.warning(f"Post-mental-model-get hook error (non-fatal): {hook_err}") + return MentalModelResponse(**mental_model) except (AuthenticationError, HTTPException): raise + except OperationValidationError as e: + raise HTTPException(status_code=e.status_code, detail=e.reason) except Exception as e: import traceback diff --git a/hindsight-api/hindsight_api/engine/memory_engine.py b/hindsight-api/hindsight_api/engine/memory_engine.py index 65572e2e..6d697dfb 100644 --- a/hindsight-api/hindsight_api/engine/memory_engine.py +++ b/hindsight-api/hindsight_api/engine/memory_engine.py @@ -597,7 +597,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 +647,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]): @@ -5482,13 +5524,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, ) diff --git a/hindsight-api/hindsight_api/extensions/__init__.py b/hindsight-api/hindsight_api/extensions/__init__.py index 592d0f8d..5a7dabeb 100644 --- a/hindsight-api/hindsight_api/extensions/__init__.py +++ b/hindsight-api/hindsight_api/extensions/__init__.py @@ -24,6 +24,10 @@ from hindsight_api.extensions.operation_validator import ( # Consolidation operation ConsolidateContext, ConsolidateResult, + # Mental Model operations + MentalModelGetContext, + MentalModelGetResult, + MentalModelRefreshResult, # Core operations OperationValidationError, OperationValidatorExtension, @@ -65,6 +69,10 @@ __all__ = [ # Operation Validator - Consolidation "ConsolidateContext", "ConsolidateResult", + # Operation Validator - Mental Model + "MentalModelGetContext", + "MentalModelGetResult", + "MentalModelRefreshResult", # Tenant/Auth "ApiKeyTenantExtension", "AuthenticationError", diff --git a/hindsight-api/hindsight_api/extensions/operation_validator.py b/hindsight-api/hindsight_api/extensions/operation_validator.py index de4fb814..0b778ca7 100644 --- a/hindsight-api/hindsight_api/extensions/operation_validator.py +++ b/hindsight-api/hindsight_api/extensions/operation_validator.py @@ -196,6 +196,48 @@ class ConsolidateResult: error: str | None = None +# ============================================================================= +# Mental Model Contexts +# ============================================================================= + + +@dataclass +class MentalModelGetContext: + """Context for a mental model GET operation validation (pre-operation).""" + + bank_id: str + mental_model_id: str + request_context: "RequestContext" + + +@dataclass +class MentalModelGetResult: + """Result context for post-mental-model-GET hook.""" + + bank_id: str + mental_model_id: str + request_context: "RequestContext" + output_tokens: int # tokens in the returned content + success: bool = True + error: str | None = None + + +@dataclass +class MentalModelRefreshResult: + """Result context for post-mental-model-refresh hook.""" + + bank_id: str + mental_model_id: str + request_context: "RequestContext" + query_tokens: int # tokens in source_query + output_tokens: int # tokens in generated content + context_tokens: int # tokens in context (if any) + facts_used: int # facts referenced in based_on + mental_models_used: int # mental models referenced in based_on + success: bool = True + error: str | None = None + + class OperationValidatorExtension(Extension, ABC): """ Validates and hooks into retain/recall/reflect/consolidate operations. @@ -402,3 +444,64 @@ class OperationValidatorExtension(Extension, ABC): - error: Error message (if failed) """ pass + + # ========================================================================= + # Mental Model - Pre-operation validation hook (optional - override to implement) + # ========================================================================= + + async def validate_mental_model_get(self, ctx: MentalModelGetContext) -> ValidationResult: + """ + Validate a mental model GET operation before execution. + + Override to implement custom validation logic for mental model retrieval. + + Args: + ctx: Context containing: + - bank_id: Bank identifier + - mental_model_id: Mental model identifier + - request_context: Request context with auth info + + Returns: + ValidationResult indicating whether the operation is allowed. + """ + return ValidationResult.accept() + + # ========================================================================= + # Mental Model - Post-operation hooks (optional - override to implement) + # ========================================================================= + + async def on_mental_model_get_complete(self, result: MentalModelGetResult) -> None: + """ + Called after a mental model GET operation completes (success or failure). + + Override to implement post-operation logic such as tracking or audit logging. + + Args: + result: Result context containing: + - bank_id: Bank identifier + - mental_model_id: Mental model identifier + - output_tokens: Token count of the returned content + - success: Whether the operation succeeded + - error: Error message (if failed) + """ + pass + + async def on_mental_model_refresh_complete(self, result: MentalModelRefreshResult) -> None: + """ + Called after a mental model refresh operation completes (success or failure). + + Override to implement post-operation logic such as tracking or audit logging. + + Args: + result: Result context containing: + - bank_id: Bank identifier + - mental_model_id: Mental model identifier + - query_tokens: Tokens in source_query + - output_tokens: Tokens in generated content + - context_tokens: Tokens in context + - facts_used: Number of facts referenced + - mental_models_used: Number of mental models referenced + - success: Whether the operation succeeded + - error: Error message (if failed) + """ + pass diff --git a/hindsight-api/tests/test_mental_model_hooks.py b/hindsight-api/tests/test_mental_model_hooks.py new file mode 100644 index 00000000..c88655c5 --- /dev/null +++ b/hindsight-api/tests/test_mental_model_hooks.py @@ -0,0 +1,206 @@ +"""Unit tests for mental model operation validator hooks. + +Tests that the operation validator hooks are called correctly for +mental model GET and refresh operations. +""" + +import pytest + +from hindsight_api.extensions.operation_validator import ( + MentalModelGetContext, + MentalModelGetResult, + MentalModelRefreshResult, + OperationValidatorExtension, + ValidationResult, +) + + +class TestMentalModelGetContextDataclass: + """Tests for MentalModelGetContext dataclass.""" + + def test_create_context(self): + """Test creating a MentalModelGetContext.""" + from unittest.mock import MagicMock + + request_context = MagicMock() + ctx = MentalModelGetContext( + bank_id="bank-1", + mental_model_id="mm-1", + request_context=request_context, + ) + + assert ctx.bank_id == "bank-1" + assert ctx.mental_model_id == "mm-1" + assert ctx.request_context is request_context + + +class TestMentalModelGetResultDataclass: + """Tests for MentalModelGetResult dataclass.""" + + def test_create_result_success(self): + """Test creating a successful MentalModelGetResult.""" + from unittest.mock import MagicMock + + request_context = MagicMock() + result = MentalModelGetResult( + bank_id="bank-1", + mental_model_id="mm-1", + request_context=request_context, + output_tokens=250, + ) + + assert result.bank_id == "bank-1" + assert result.mental_model_id == "mm-1" + assert result.output_tokens == 250 + assert result.success is True + assert result.error is None + + def test_create_result_failure(self): + """Test creating a failed MentalModelGetResult.""" + from unittest.mock import MagicMock + + result = MentalModelGetResult( + bank_id="bank-1", + mental_model_id="mm-1", + request_context=MagicMock(), + output_tokens=0, + success=False, + error="Not found", + ) + + assert result.success is False + assert result.error == "Not found" + + +class TestMentalModelRefreshResultDataclass: + """Tests for MentalModelRefreshResult dataclass.""" + + def test_create_result_with_all_fields(self): + """Test creating a MentalModelRefreshResult with all fields.""" + from unittest.mock import MagicMock + + result = MentalModelRefreshResult( + bank_id="bank-1", + mental_model_id="mm-1", + request_context=MagicMock(), + query_tokens=50, + output_tokens=500, + context_tokens=0, + facts_used=10, + mental_models_used=2, + ) + + assert result.query_tokens == 50 + assert result.output_tokens == 500 + assert result.context_tokens == 0 + assert result.facts_used == 10 + assert result.mental_models_used == 2 + assert result.success is True + assert result.error is None + + def test_create_result_failure(self): + """Test creating a failed MentalModelRefreshResult.""" + from unittest.mock import MagicMock + + result = MentalModelRefreshResult( + bank_id="bank-1", + mental_model_id="mm-1", + request_context=MagicMock(), + query_tokens=50, + output_tokens=0, + context_tokens=0, + facts_used=0, + mental_models_used=0, + success=False, + error="Reflect failed", + ) + + assert result.success is False + assert result.error == "Reflect failed" + + +class TestDefaultHookBehavior: + """Tests for default (no-op) behavior of mental model hooks on OperationValidatorExtension.""" + + @pytest.fixture + def validator(self): + """Create a concrete subclass for testing default behavior.""" + from unittest.mock import MagicMock + + # Create a concrete subclass that implements the abstract methods + class TestValidator(OperationValidatorExtension): + async def validate_retain(self, ctx): + return ValidationResult.accept() + + async def validate_recall(self, ctx): + return ValidationResult.accept() + + async def validate_reflect(self, ctx): + return ValidationResult.accept() + + return TestValidator(config={}) + + @pytest.mark.asyncio + async def test_validate_mental_model_get_default_accepts(self, validator): + """Test that default validate_mental_model_get accepts.""" + from unittest.mock import MagicMock + + ctx = MentalModelGetContext( + bank_id="bank-1", + mental_model_id="mm-1", + request_context=MagicMock(), + ) + + result = await validator.validate_mental_model_get(ctx) + + assert result.allowed is True + + @pytest.mark.asyncio + async def test_on_mental_model_get_complete_default_noop(self, validator): + """Test that default on_mental_model_get_complete is a no-op.""" + from unittest.mock import MagicMock + + result = MentalModelGetResult( + bank_id="bank-1", + mental_model_id="mm-1", + request_context=MagicMock(), + output_tokens=100, + ) + + # Should not raise + await validator.on_mental_model_get_complete(result) + + @pytest.mark.asyncio + async def test_on_mental_model_refresh_complete_default_noop(self, validator): + """Test that default on_mental_model_refresh_complete is a no-op.""" + from unittest.mock import MagicMock + + result = MentalModelRefreshResult( + bank_id="bank-1", + mental_model_id="mm-1", + request_context=MagicMock(), + query_tokens=50, + output_tokens=500, + context_tokens=0, + facts_used=5, + mental_models_used=1, + ) + + # Should not raise + await validator.on_mental_model_refresh_complete(result) + + +class TestExportsAvailable: + """Test that mental model hooks are properly exported.""" + + def test_imports_from_extensions_package(self): + """Test that all mental model types can be imported from hindsight_api.extensions.""" + from hindsight_api.extensions import ( + MentalModelGetContext, + MentalModelGetResult, + MentalModelRefreshResult, + ) + + assert MentalModelGetContext is not None + assert MentalModelGetResult is not None + assert MentalModelRefreshResult is not None