[wip] Custom session handlers

python/model_hosting_container_standards/common/fastapi/utils.py

@@ -1,6 +1,8 @@
+import json
 from typing import Any, Dict, Optional, Union
 
-from fastapi import Request
+from fastapi import Request, Response
+from fastapi.responses import JSONResponse
 from pydantic import BaseModel
 
 
@@ -33,3 +35,29 @@ def serialize_request(
         "query_params": raw_request.query_params,
         "path_params": raw_request.path_params,
     }
+
+
+def serialize_response(response: Union[Response, JSONResponse]):
+    """Create a structured data dictionary for JMESPath transformations.
+
+    Extracts and organizes response data into a standardized format that can be used
+    with JMESPath expressions to transform and extract specific data elements.
+
+    :param Union[Response, JSONResponse] response: Response body data - can be:
+        - FastAPI Response object
+        - JSONResponse object
+    :return Dict[str, Any]: Structured data with body, headers, status_code, and media_type
+    """
+    # Process response body based on type
+    body = response.body.decode(response.charset)
+    try:
+        body = json.loads(body)
+    except json.JSONDecodeError:
+        # If body is not JSON, keep it as a string
+        # logger.warning(f"Response body is not JSON, keeping as string: {e}")
+        pass
+
+    return {
+        "body": body,
+        "headers": response.headers,
+    }

python/model_hosting_container_standards/common/transforms/base_api_transform.py

@@ -7,7 +7,7 @@
 from pydantic import BaseModel, Field
 
 from ...logging_config import logger
-from ..fastapi.utils import serialize_request
+from ..fastapi.utils import serialize_request, serialize_response
 from .utils import _compile_jmespath_expressions
 
 
@@ -103,6 +103,19 @@ async def transform_request(self, raw_request: Request):
         """
         raise NotImplementedError()
 
+    def _transform_response(self, response: Response):
+        """Transform the response based on the request processing results.
+
+        Subclasses must implement this method to handle request parsing, validation,
+        and transformation according to their specific operation requirements.
+
+        :param Response response: The response to transform
+        :param transform_request_output: Output from the request transformation
+        :raises NotImplementedError: Must be implemented by subclasses
+        """
+        response_data = serialize_response(response)
+        return self._transform(response_data, self._response_shape)
+
     def _transform_request(
         self, request: Optional[BaseModel], raw_request: Request
     ) -> Dict[str, Any]:

python/model_hosting_container_standards/sagemaker/__init__.py

@@ -21,7 +21,11 @@
 from .lora.models import AppendOperation
 from .sagemaker_loader import SageMakerFunctionLoader
 from .sagemaker_router import create_sagemaker_router
-from .sessions import create_session_transform_decorator
+from .sessions import (
+    create_session_transform_decorator,
+    register_engine_session_handler,
+)
+from .sessions.models import SageMakerSessionHeader
 
 # SageMaker decorator instances - created using utility functions
 
@@ -118,17 +122,47 @@ def inject_adapter_id(
     )
 
 
-def stateful_session_manager():
+def stateful_session_manager(request_session_id_path: Optional[str] = None):
     """Create a decorator for session-based sticky routing.
 
     This decorator enables stateful session management without JMESPath transformations.
     Pass empty dicts to enable transform infrastructure (for intercept functionality)
     without requiring JMESPath expressions.
 
+    Args:
+        request_session_id_path: JMESPath target path where session ID should be
+                                 injected INTO the request body from the session header
+
     Returns:
         A decorator that can be applied to route handlers to enable session management
     """
-    return create_session_transform_decorator()(request_shape={}, response_shape={})
+    request_shape = {}
+    if request_session_id_path:
+        request_shape[request_session_id_path] = (
+            f'headers."{SageMakerSessionHeader.SESSION_ID}"'
+        )
+    return create_session_transform_decorator()(
+        request_shape=request_shape, response_shape={}
+    )
+
+
+def register_create_session_handler(
+    request_shape, response_session_id_path: str, content_path: Optional[str] = None
+):
+    return register_engine_session_handler(
+        "create_session",
+        request_shape=request_shape,
+        response_session_id_path=response_session_id_path,
+        content_path=content_path or "`successfully created session.`",
+    )
+
+
+def register_close_session_handler(request_shape, content_path: Optional[str] = None):
+    return register_engine_session_handler(
+        "close_session",
+        request_shape=request_shape,
+        content_path=content_path or "`successfully closed session.`",
+    )
 
 
 def bootstrap(app: FastAPI) -> FastAPI:

python/model_hosting_container_standards/sagemaker/sagemaker_router.py

@@ -41,6 +41,11 @@ def get_sagemaker_route_config(handler_type: str) -> Optional[RouteConfig]:
             summary="Model inference endpoint",
         )
 
+    if handler_type in ["create_session", "close_session"]:
+        # It's a request transformer, not a standalone API endpoint
+        # It modifies requests in-flight but doesn't expose its own route
+        return None
+
     # Delegate to LoRA route resolver for LoRA-specific handlers
     return get_lora_route_config(handler_type)