nebu 0.1.45__py3-none-any.whl → 0.1.48__py3-none-any.whl

nebu/processors/consumer.py

@@ -1,12 +1,14 @@
 #!/usr/bin/env python3
+import importlib
 import json
 import os
 import socket
 import sys
 import time
 import traceback
-from datetime import datetime
-from typing import Dict, TypeVar
+import types  # Added for ModuleType
+from datetime import datetime, timezone
+from typing import Any, Callable, Dict, List, Optional, Tuple, TypeVar, cast
 
 import redis
 import socks
@@ -18,67 +20,168 @@ T = TypeVar("T")
 # Environment variable name used as a guard in the decorator
 _NEBU_INSIDE_CONSUMER_ENV_VAR = "_NEBU_INSIDE_CONSUMER_EXEC"
 
-# Get function and model source code and create them dynamically
-try:
-    function_source = os.environ.get("FUNCTION_SOURCE")
-    function_name = os.environ.get("FUNCTION_NAME")
-    stream_message_source = os.environ.get("STREAM_MESSAGE_SOURCE")
-    input_model_source = os.environ.get("INPUT_MODEL_SOURCE")
-    output_model_source = os.environ.get("OUTPUT_MODEL_SOURCE")
-    content_type_source = os.environ.get("CONTENT_TYPE_SOURCE")
-    is_stream_message = os.environ.get("IS_STREAM_MESSAGE") == "True"
-    param_type_name = os.environ.get("PARAM_TYPE_NAME")
-    return_type_name = os.environ.get("RETURN_TYPE_NAME")
-    param_type_str = os.environ.get("PARAM_TYPE_STR")
-    return_type_str = os.environ.get("RETURN_TYPE_STR")
-    content_type_name = os.environ.get("CONTENT_TYPE_NAME")
+# --- Global variables for dynamically loaded code ---
+target_function: Optional[Callable] = None
+init_function: Optional[Callable] = None
+imported_module: Optional[types.ModuleType] = None
+local_namespace: Dict[str, Any] = {}  # Namespace for included objects
+last_load_mtime: float = 0.0
+entrypoint_abs_path: Optional[str] = None
+
+
+# --- Function to Load/Reload User Code ---
+def load_or_reload_user_code(
+    module_path: str,
+    function_name: str,
+    entrypoint_abs_path: str,
+    init_func_name: Optional[str] = None,
+    included_object_sources: Optional[List[Tuple[str, List[str]]]] = None,
+) -> Tuple[
+    Optional[Callable],
+    Optional[Callable],
+    Optional[types.ModuleType],
+    Dict[str, Any],
+    float,
+]:
+    """Loads or reloads the user code module, executes includes, and returns functions/module."""
+    global _NEBU_INSIDE_CONSUMER_ENV_VAR  # Access the global guard var name
+
+    current_mtime = 0.0
+    loaded_target_func = None
+    loaded_init_func = None
+    loaded_module = None
+    exec_namespace: Dict[str, Any] = {}  # Use a local namespace for this load attempt
+
+    print(f"[Code Loader] Attempting to load/reload module: '{module_path}'")
+    os.environ[_NEBU_INSIDE_CONSUMER_ENV_VAR] = "1"  # Set guard *before* import/reload
+    print(f"[Code Loader] Set environment variable {_NEBU_INSIDE_CONSUMER_ENV_VAR}=1")
 
-    # Get source for the file containing the decorated function
-    decorated_func_file_source = os.environ.get("DECORATED_FUNC_FILE_SOURCE")
-    # Get sources for the directory containing the decorated function
-    # decorated_dir_sources_json = os.environ.get("DECORATED_DIR_SOURCES") # Removed
+    try:
+        current_mtime = os.path.getmtime(entrypoint_abs_path)
+
+        # Execute included object sources FIRST (if any)
+        if included_object_sources:
+            print("[Code Loader] Executing @include object sources...")
+            # Include necessary imports for the exec context
+            exec("from pydantic import BaseModel, Field", exec_namespace)
+            exec(
+                "from typing import Optional, List, Dict, Any, Generic, TypeVar",
+                exec_namespace,
+            )
+            exec("T_exec = TypeVar('T_exec')", exec_namespace)
+            exec("from nebu.processors.models import *", exec_namespace)
+            # ... add other common imports if needed by included objects ...
 
-    # Get init_func source if provided
-    init_func_source = os.environ.get("INIT_FUNC_SOURCE")
-    init_func_name = os.environ.get("INIT_FUNC_NAME")
+            for i, (obj_source, args_sources) in enumerate(included_object_sources):
+                try:
+                    exec(obj_source, exec_namespace)
+                    print(
+                        f"[Code Loader] Successfully executed included object {i} base source"
+                    )
+                    for j, arg_source in enumerate(args_sources):
+                        try:
+                            exec(arg_source, exec_namespace)
+                            print(
+                                f"[Code Loader] Successfully executed included object {i} arg {j} source"
+                            )
+                        except Exception as e_arg:
+                            print(
+                                f"Error executing included object {i} arg {j} source: {e_arg}"
+                            )
+                            traceback.print_exc()  # Log specific error but continue? Or fail reload?
+                except Exception as e_base:
+                    print(f"Error executing included object {i} base source: {e_base}")
+                    traceback.print_exc()  # Log specific error but continue? Or fail reload?
+            print("[Code Loader] Finished executing included object sources.")
+
+        # Check if module is already loaded and needs reload
+        if module_path in sys.modules:
+            print(
+                f"[Code Loader] Module '{module_path}' already imported. Reloading..."
+            )
+            # Pass the exec_namespace as globals? Usually reload works within its own context.
+            # If included objects *modify* the module's global scope upon exec,
+            # reload might not pick that up easily. Might need a fresh import instead.
+            # Let's try reload first.
+            loaded_module = importlib.reload(sys.modules[module_path])
+            print(f"[Code Loader] Successfully reloaded module: {module_path}")
+        else:
+            # Import the main module
+            loaded_module = importlib.import_module(module_path)
+            print(
+                f"[Code Loader] Successfully imported module for the first time: {module_path}"
+            )
 
-    # Check for generic type arguments
-    input_model_args = []
-    output_model_args = []
-    content_type_args = []
+        # Get the target function from the loaded/reloaded module
+        loaded_target_func = getattr(loaded_module, function_name)
+        print(
+            f"[Code Loader] Successfully loaded function '{function_name}' from module '{module_path}'"
+        )
 
-    # Get input model arg sources
-    i = 0
-    while True:
-        arg_source = os.environ.get(f"INPUT_MODEL_ARG_{i}_SOURCE")
-        if arg_source:
-            input_model_args.append(arg_source)
-            i += 1
-        else:
-            break
+        # Get the init function if specified
+        if init_func_name:
+            loaded_init_func = getattr(loaded_module, init_func_name)
+            print(
+                f"[Code Loader] Successfully loaded init function '{init_func_name}' from module '{module_path}'"
+            )
+            # Execute init_func
+            print(f"[Code Loader] Executing init_func: {init_func_name}...")
+            loaded_init_func()  # Call the function
+            print(f"[Code Loader] Successfully executed init_func: {init_func_name}")
+
+        print("[Code Loader] Code load/reload successful.")
+        return (
+            loaded_target_func,
+            loaded_init_func,
+            loaded_module,
+            exec_namespace,
+            current_mtime,
+        )
 
-    # Get output model arg sources
-    i = 0
-    while True:
-        arg_source = os.environ.get(f"OUTPUT_MODEL_ARG_{i}_SOURCE")
-        if arg_source:
-            output_model_args.append(arg_source)
-            i += 1
-        else:
-            break
+    except FileNotFoundError:
+        print(
+            f"[Code Loader] Error: Entrypoint file not found at '{entrypoint_abs_path}'. Cannot load/reload."
+        )
+        return None, None, None, {}, 0.0  # Indicate failure
+    except ImportError as e:
+        print(f"[Code Loader] Error importing/reloading module '{module_path}': {e}")
+        traceback.print_exc()
+        return None, None, None, {}, 0.0  # Indicate failure
+    except AttributeError as e:
+        print(
+            f"[Code Loader] Error accessing function '{function_name}' or '{init_func_name}' in module '{module_path}': {e}"
+        )
+        traceback.print_exc()
+        return None, None, None, {}, 0.0  # Indicate failure
+    except Exception as e:
+        print(f"[Code Loader] Unexpected error during code load/reload: {e}")
+        traceback.print_exc()
+        return None, None, None, {}, 0.0  # Indicate failure
+    finally:
+        # Unset the guard environment variable
+        os.environ.pop(_NEBU_INSIDE_CONSUMER_ENV_VAR, None)
+        print(
+            f"[Code Loader] Unset environment variable {_NEBU_INSIDE_CONSUMER_ENV_VAR}"
+        )
 
-    # Get content type arg sources
-    i = 0
-    while True:
-        arg_source = os.environ.get(f"CONTENT_TYPE_ARG_{i}_SOURCE")
-        if arg_source:
-            content_type_args.append(arg_source)
-            i += 1
-        else:
-            break
 
-    # Get included object sources
-    included_object_sources = []
+# --- Get Environment Variables ---
+try:
+    # Core function info
+    _function_name = os.environ.get("FUNCTION_NAME")
+    _entrypoint_rel_path = os.environ.get("NEBU_ENTRYPOINT_MODULE_PATH")
+
+    # Type info
+    is_stream_message = os.environ.get("IS_STREAM_MESSAGE") == "True"
+    param_type_str = os.environ.get("PARAM_TYPE_STR")
+    return_type_str = os.environ.get("RETURN_TYPE_STR")
+    content_type_name = os.environ.get("CONTENT_TYPE_NAME")
+
+    # Init func info
+    _init_func_name = os.environ.get("INIT_FUNC_NAME")
+
+    # Included object sources
+    _included_object_sources = []
     i = 0
     while True:
         obj_source = os.environ.get(f"INCLUDED_OBJECT_{i}_SOURCE")
@@ -92,213 +195,87 @@ try:
                     j += 1
                 else:
                     break
-            included_object_sources.append((obj_source, args))
+            _included_object_sources.append((obj_source, args))
             i += 1
         else:
             break
 
-    if not function_source or not function_name:
-        print("FUNCTION_SOURCE or FUNCTION_NAME environment variables not set")
+    if not _function_name or not _entrypoint_rel_path:
+        print(
+            "FATAL: FUNCTION_NAME or NEBU_ENTRYPOINT_MODULE_PATH environment variables not set"
+        )
         sys.exit(1)
 
-    # Create a local namespace for executing the function
-    local_namespace = {}
-
-    # Include pydantic BaseModel and typing tools for type annotations
-    exec("from pydantic import BaseModel, Field", local_namespace)
-    exec(
-        "from typing import Optional, List, Dict, Any, Generic, TypeVar",
-        local_namespace,
-    )
-    exec("T = TypeVar('T')", local_namespace)
-    exec("from nebu.processors.models import *", local_namespace)
-    exec("from nebu.processors.processor import *", local_namespace)
-    # Add import for the processor decorator itself
-    exec("from nebu.processors.decorate import processor", local_namespace)
-    # Add import for chatx openai types
-    exec("from nebu.chatx.openai import *", local_namespace)
-
-    # Set the guard environment variable before executing any source code
-    os.environ[_NEBU_INSIDE_CONSUMER_ENV_VAR] = "1"
-    print(f"[Consumer] Set environment variable {_NEBU_INSIDE_CONSUMER_ENV_VAR}=1")
-
-    try:
-        # Execute the source file of the decorated function FIRST
-        if decorated_func_file_source:
-            print("[Consumer] Executing decorated function's file source...")
-            try:
-                exec(decorated_func_file_source, local_namespace)
-                print(
-                    "[Consumer] Successfully executed decorated function's file source."
-                )
-            except Exception as e:
-                print(f"Error executing decorated function's file source: {e}")
-                traceback.print_exc()  # Warn and continue
-        else:
-            print(
-                "[Consumer] No decorated function's file source found in environment."
-            )
-
-        # Execute the sources from the decorated function's directory
-        # if decorated_dir_sources_json:
-        #     print("[Consumer] Executing decorated function's directory sources...")
-        #     try:
-        #         dir_sources = json.loads(decorated_dir_sources_json)
-        #         # Sort by relative path for some predictability (e.g., __init__.py first)
-        #         for rel_path, source_code in sorted(dir_sources.items()):
-        #             print(f"[Consumer] Executing source from: {rel_path}...")
-        #             try:
-        #                 exec(source_code, local_namespace)
-        #                 print(f"[Consumer] Successfully executed source from: {rel_path}")
-        #             except Exception as e:
-        #                 print(f"Error executing source from {rel_path}: {e}")
-        #                 traceback.print_exc() # Warn and continue
-        #     except json.JSONDecodeError as e:
-        #         print(f"Error decoding DECORATED_DIR_SOURCES JSON: {e}")
-        #         traceback.print_exc()
-        #     except Exception as e:
-        #         print(f"Unexpected error processing directory sources: {e}")
-        #         traceback.print_exc()
-        # else:
-        #     print("[Consumer] No decorated function's directory sources found in environment.")
-
-        # Execute included object sources NEXT, as they might define types needed by others
-        print("[Consumer] Executing included object sources...")
-        for i, (obj_source, args_sources) in enumerate(included_object_sources):
-            try:
-                exec(obj_source, local_namespace)
+    # Calculate absolute path for modification time checking
+    # Assuming CWD or PYTHONPATH allows finding the relative path
+    # This might need adjustment based on deployment specifics
+    entrypoint_abs_path = os.path.abspath(_entrypoint_rel_path)
+    if not os.path.exists(entrypoint_abs_path):
+        # Try constructing path based on PYTHONPATH if direct abspath fails
+        python_path = os.environ.get("PYTHONPATH", "").split(os.pathsep)
+        found_path = False
+        for p_path in python_path:
+            potential_path = os.path.abspath(os.path.join(p_path, _entrypoint_rel_path))
+            if os.path.exists(potential_path):
+                entrypoint_abs_path = potential_path
+                found_path = True
                 print(
-                    f"[Consumer] Successfully executed included object {i} base source"
+                    f"[Consumer] Found entrypoint absolute path via PYTHONPATH: {entrypoint_abs_path}"
                 )
-                for j, arg_source in enumerate(args_sources):
-                    try:
-                        exec(arg_source, local_namespace)
-                        print(
-                            f"[Consumer] Successfully executed included object {i} arg {j} source"
-                        )
-                    except Exception as e:
-                        print(
-                            f"Error executing included object {i} arg {j} source: {e}"
-                        )
-                        traceback.print_exc()
-            except Exception as e:
-                print(f"Error executing included object {i} base source: {e}")
-                traceback.print_exc()
-        print("[Consumer] Finished executing included object sources.")
-
-        # First try to import the module to get any needed dependencies
-        # This is a fallback in case the module is available
-        module_name = os.environ.get("MODULE_NAME")
-        try:
-            if module_name:
-                exec(f"import {module_name}", local_namespace)
-                print(f"Successfully imported module {module_name}")
-        except Exception as e:
-            print(f"Warning: Could not import module {module_name}: {e}")
+                break
+        if not found_path:
             print(
-                "This is expected if running in a Jupyter notebook. Will use dynamic execution."
+                f"FATAL: Could not find entrypoint file via relative path '{_entrypoint_rel_path}' or in PYTHONPATH."
             )
+            # Attempting abspath anyway for the error message in load function
+            entrypoint_abs_path = os.path.abspath(_entrypoint_rel_path)
+
+    # Convert entrypoint file path to module path
+    _module_path = _entrypoint_rel_path.replace(os.sep, ".")
+    if _module_path.endswith(".py"):
+        _module_path = _module_path[:-3]
+    if _module_path.endswith(".__init__"):
+        _module_path = _module_path[: -len(".__init__")]
+    elif _module_path == "__init__":
+        print(
+            f"FATAL: Entrypoint '{_entrypoint_rel_path}' resolves to ambiguous top-level __init__. Please use a named file or package."
+        )
+        sys.exit(1)
+    if not _module_path:
+        print(
+            f"FATAL: Could not derive a valid module path from entrypoint '{_entrypoint_rel_path}'"
+        )
+        sys.exit(1)
 
-        # Define the models
-        # First define stream message class if needed
-        if stream_message_source:
-            try:
-                exec(stream_message_source, local_namespace)
-                print("Successfully defined Message class")
-            except Exception as e:
-                print(f"Error defining Message: {e}")
-                traceback.print_exc()
-
-        # Define content type if available
-        if content_type_source:
-            try:
-                exec(content_type_source, local_namespace)
-                print(f"Successfully defined content type {content_type_name}")
-
-                # Define any content type args
-                for arg_source in content_type_args:
-                    try:
-                        exec(arg_source, local_namespace)
-                        print("Successfully defined content type argument")
-                    except Exception as e:
-                        print(f"Error defining content type argument: {e}")
-                        traceback.print_exc()
-            except Exception as e:
-                print(f"Error defining content type: {e}")
-                traceback.print_exc()
-
-        # Define input model if different from stream message
-        if input_model_source and (
-            not is_stream_message or input_model_source != stream_message_source
-        ):
-            try:
-                exec(input_model_source, local_namespace)
-                print(f"Successfully defined input model {param_type_str}")
-
-                # Define any input model args
-                for arg_source in input_model_args:
-                    try:
-                        exec(arg_source, local_namespace)
-                        print("Successfully defined input model argument")
-                    except Exception as e:
-                        print(f"Error defining input model argument: {e}")
-                        traceback.print_exc()
-            except Exception as e:
-                print(f"Error defining input model: {e}")
-                traceback.print_exc()
-
-        # Define output model
-        if output_model_source:
-            try:
-                exec(output_model_source, local_namespace)
-                print(f"Successfully defined output model {return_type_str}")
+    print(
+        f"[Consumer] Initializing. Entrypoint: '{_entrypoint_rel_path}', Module: '{_module_path}', Function: '{_function_name}', Init: '{_init_func_name}'"
+    )
 
-                # Define any output model args
-                for arg_source in output_model_args:
-                    try:
-                        exec(arg_source, local_namespace)
-                        print("Successfully defined output model argument")
-                    except Exception as e:
-                        print(f"Error defining output model argument: {e}")
-                        traceback.print_exc()
-            except Exception as e:
-                print(f"Error defining output model: {e}")
-                traceback.print_exc()
+    # --- Initial Load of User Code ---
+    (
+        target_function,
+        init_function,
+        imported_module,
+        local_namespace,
+        last_load_mtime,
+    ) = load_or_reload_user_code(
+        _module_path,
+        _function_name,
+        entrypoint_abs_path,
+        _init_func_name,
+        _included_object_sources,
+    )
 
-        # Finally, execute the function code
-        try:
-            exec(function_source, local_namespace)
-            target_function = local_namespace[function_name]
-            print(f"Successfully loaded function {function_name}")
-        except Exception as e:
-            print(f"Error creating function from source: {e}")
-            traceback.print_exc()
-            sys.exit(1)
+    if target_function is None or imported_module is None:
+        print("FATAL: Initial load of user code failed. Exiting.")
+        sys.exit(1)
+    print(
+        f"[Consumer] Initial code load successful. Last modified time: {last_load_mtime}"
+    )
 
-        # Execute init_func if provided
-        if init_func_source and init_func_name:
-            print(f"Executing init_func: {init_func_name}...")
-            try:
-                exec(init_func_source, local_namespace)
-                init_function = local_namespace[init_func_name]
-                print(
-                    f"[Consumer] Environment before calling init_func {init_func_name}: {os.environ}"
-                )
-                init_function()  # Call the function
-                print(f"Successfully executed init_func: {init_func_name}")
-            except Exception as e:
-                print(f"Error executing init_func '{init_func_name}': {e}")
-                traceback.print_exc()
-                # Decide if failure is critical. For now, let's exit.
-                print("Exiting due to init_func failure.")
-                sys.exit(1)
-    finally:
-        # Unset the guard environment variable after all execs are done
-        os.environ.pop(_NEBU_INSIDE_CONSUMER_ENV_VAR, None)
-        print(f"[Consumer] Unset environment variable {_NEBU_INSIDE_CONSUMER_ENV_VAR}")
 
 except Exception as e:
-    print(f"Error setting up function: {e}")
+    print(f"FATAL: Error during initial environment setup or code load: {e}")
     traceback.print_exc()
     sys.exit(1)
 
@@ -349,46 +326,66 @@ except ResponseError as e:
 
 # Function to process messages
 def process_message(message_id: str, message_data: Dict[str, str]) -> None:
-    # Initialize variables that need to be accessible in the except block
+    # Access the globally managed user code elements
+    global target_function, imported_module, local_namespace
+
+    # Check if target_function is loaded (might be None if reload failed)
+    if target_function is None or imported_module is None:
+        print(
+            f"Error processing message {message_id}: User code (target_function or module) is not loaded. Skipping."
+        )
+        # Decide how to handle this - skip and ack? Send error?
+        # Sending error for now, but not acking yet.
+        # This requires the main loop to handle potential ack failure later if needed.
+        _send_error_response(
+            message_id,
+            "User code is not loaded (likely due to a failed reload)",
+            traceback.format_exc(),
+            None,
+            None,
+        )  # Pass None for user_id if unavailable here
+        return  # Skip processing
+
     return_stream = None
     user_id = None
     try:
-        # Extract the JSON string payload from the 'data' field
         payload_str = message_data.get("data")
-
-        # decode_responses=True should mean payload_str is a string if found.
-        if not payload_str:
+        if not payload_str or not isinstance(payload_str, str):
             raise ValueError(
-                f"Missing or invalid 'data' field (expected string): {message_data}"
+                f"Missing or invalid 'data' field (expected non-empty string): {message_data}"
             )
-
-        # Parse the JSON string into a dictionary
         try:
             raw_payload = json.loads(payload_str)
         except json.JSONDecodeError as json_err:
             raise ValueError(f"Failed to parse JSON payload: {json_err}") from json_err
-
-        # Validate that raw_payload is a dictionary as expected
         if not isinstance(raw_payload, dict):
             raise TypeError(
                 f"Expected parsed payload to be a dictionary, but got {type(raw_payload)}"
             )
 
-        print(f"Raw payload: {raw_payload}")
+        # print(f"Raw payload: {raw_payload}") # Reduce verbosity
 
-        # Extract fields from the parsed payload
-        # These fields are extracted for completeness and potential future use
-        kind = raw_payload.get("kind", "")  # kind
-        msg_id = raw_payload.get("id", "")  # msg_id
+        kind = raw_payload.get("kind", "")
+        msg_id = raw_payload.get("id", "")
         content_raw = raw_payload.get("content", {})
-        created_at = raw_payload.get("created_at", 0)  # created_at
+        created_at_str = raw_payload.get("created_at")  # Get as string or None
+        # Attempt to parse created_at, fallback to now()
+        try:
+            created_at = (
+                datetime.fromisoformat(created_at_str)
+                if created_at_str
+                else datetime.now(timezone.utc)
+            )
+        except ValueError:
+            created_at = datetime.now(timezone.utc)
+
         return_stream = raw_payload.get("return_stream")
         user_id = raw_payload.get("user_id")
-        orgs = raw_payload.get("organizations")  # organizations
-        handle = raw_payload.get("handle")  # handle
-        adapter = raw_payload.get("adapter")  # adapter
+        orgs = raw_payload.get("organizations")
+        handle = raw_payload.get("handle")
+        adapter = raw_payload.get("adapter")
 
-        # --- Health Check Logic based on kind ---
+        # --- Health Check Logic (Keep as is) ---
         if kind == "HealthCheck":
             print(f"Received HealthCheck message {message_id}")
             health_response = {
@@ -413,109 +410,180 @@ def process_message(message_id: str, message_data: Dict[str, str]) -> None:
             return  # Exit early for health checks
         # --- End Health Check Logic ---
 
-        # Parse the content field if it's a string
+        # Parse content if it's a string (e.g., double-encoded JSON)
         if isinstance(content_raw, str):
             try:
                 content = json.loads(content_raw)
             except json.JSONDecodeError:
-                content = content_raw
+                content = content_raw  # Keep as string if not valid JSON
         else:
             content = content_raw
 
-        print(f"Content: {content}")
+        # print(f"Content: {content}") # Reduce verbosity
 
-        # For StreamMessage, construct the proper input object
-        if is_stream_message and "Message" in local_namespace:
-            # If we have a content type, try to construct it
-            if content_type_name and content_type_name in local_namespace:
-                # Try to create the content type model first
-                try:
-                    content_model = local_namespace[content_type_name].model_validate(
-                        content
-                    )
-                    print(f"Content model: {content_model}")
-                    input_obj = local_namespace["Message"](
+        # --- Construct Input Object using Imported Types ---
+        input_obj: Any = None
+        input_type_class = None
+
+        try:
+            # Try to get the actual model classes (they should be available via import)
+            # Need to handle potential NameErrors if imports failed silently
+            # Note: This assumes models are defined in the imported module scope
+            # Or imported by the imported module.
+            from nebu.processors.models import Message  # Import needed message class
+
+            if is_stream_message:
+                message_class = Message  # Use imported class
+                content_model_class = None
+                if content_type_name:
+                    try:
+                        # Assume content_type_name refers to a class available in the global scope
+                        # (either from imported module or included objects)
+                        # Use the globally managed imported_module and local_namespace
+                        content_model_class = getattr(
+                            imported_module, content_type_name, None
+                        )
+                        if content_model_class is None:
+                            # Check in local_namespace from included objects as fallback
+                            content_model_class = local_namespace.get(content_type_name)
+                        if content_model_class is None:
+                            print(
+                                f"Warning: Content type class '{content_type_name}' not found in imported module or includes."
+                            )
+                        else:
+                            print(f"Found content model class: {content_model_class}")
+                    except AttributeError:
+                        print(
+                            f"Warning: Content type class '{content_type_name}' not found in imported module."
+                        )
+                    except Exception as e:
+                        print(
+                            f"Warning: Error resolving content type class '{content_type_name}': {e}"
+                        )
+
+                if content_model_class:
+                    try:
+                        content_model = content_model_class.model_validate(content)
+                        print(f"Validated content model: {content_model}")
+                        input_obj = message_class(
+                            kind=kind,
+                            id=msg_id,
+                            content=content_model,
+                            created_at=int(created_at.timestamp()),
+                            return_stream=return_stream,
+                            user_id=user_id,
+                            orgs=orgs,
+                            handle=handle,
+                            adapter=adapter,
+                        )
+                    except Exception as e:
+                        print(
+                            f"Error validating/creating content model '{content_type_name}': {e}. Falling back."
+                        )
+                        # Fallback to raw content in Message
+                        input_obj = message_class(
+                            kind=kind,
+                            id=msg_id,
+                            content=cast(Any, content),
+                            created_at=int(created_at.timestamp()),
+                            return_stream=return_stream,
+                            user_id=user_id,
+                            orgs=orgs,
+                            handle=handle,
+                            adapter=adapter,
+                        )
+                else:
+                    # No content type name or class found, use raw content
+                    input_obj = message_class(
                         kind=kind,
                         id=msg_id,
-                        content=content_model,
-                        created_at=created_at,
+                        content=cast(Any, content),
+                        created_at=int(created_at.timestamp()),
                         return_stream=return_stream,
                         user_id=user_id,
                         orgs=orgs,
                         handle=handle,
                         adapter=adapter,
                     )
+            else:  # Not a stream message, use the function's parameter type
+                param_type_name = (
+                    param_type_str  # Assume param_type_str holds the class name
+                )
+                # Attempt to resolve the parameter type class
+                try:
+                    # Use the globally managed imported_module and local_namespace
+                    input_type_class = (
+                        getattr(imported_module, param_type_name, None)
+                        if param_type_name
+                        else None
+                    )
+                    if input_type_class is None and param_type_name:
+                        input_type_class = local_namespace.get(param_type_name)
+                    if input_type_class is None:
+                        if param_type_name:  # Only warn if a name was expected
+                            print(
+                                f"Warning: Input type class '{param_type_name}' not found. Passing raw content."
+                            )
+                        input_obj = content
+                    else:
+                        print(f"Found input model class: {input_type_class}")
+                        input_obj = input_type_class.model_validate(content)
+                        print(f"Validated input model: {input_obj}")
+                except AttributeError:
+                    print(
+                        f"Warning: Input type class '{param_type_name}' not found in imported module."
+                    )
+                    input_obj = content
                 except Exception as e:
-                    print(f"Error creating content type model: {e}")
-                    # Fallback to using raw content
-                    input_obj = local_namespace["Message"](
-                        kind=kind,
-                        id=msg_id,
-                        content=content,
-                        created_at=created_at,
-                        return_stream=return_stream,
-                        user_id=user_id,
-                        orgs=orgs,
-                        handle=handle,
-                        adapter=adapter,
+                    print(
+                        f"Error resolving/validating input type '{param_type_name}': {e}. Passing raw content."
                     )
-            else:
-                # Just use the raw content
-                print("Using raw content")
-                input_obj = local_namespace["Message"](
-                    kind=kind,
-                    id=msg_id,
-                    content=content,
-                    created_at=created_at,
-                    return_stream=return_stream,
-                    user_id=user_id,
-                    orgs=orgs,
-                    handle=handle,
-                    adapter=adapter,
-                )
-        else:
-            # Otherwise use the param type directly
-            try:
-                if param_type_str in local_namespace:
-                    print(f"Validating content against {param_type_str}")
-                    input_obj = local_namespace[param_type_str].model_validate(content)
-                else:
-                    # If we can't find the exact type, just pass the content directly
                     input_obj = content
-            except Exception as e:
-                print(f"Error creating input model: {e}, using raw content")
-                raise e
 
-        print(f"Input object: {input_obj}")
+        except NameError as e:
+            print(
+                f"Error: Required class (e.g., Message or parameter type) not found. Import failed? {e}"
+            )
+            # Can't proceed without types, re-raise or handle error response
+            raise RuntimeError(f"Required class not found: {e}") from e
+        except Exception as e:
+            print(f"Error constructing input object: {e}")
+            raise  # Re-raise unexpected errors during input construction
+
+        # print(f"Input object: {input_obj}") # Reduce verbosity
 
         # Execute the function
+        print(f"Executing function...")
         result = target_function(input_obj)
-        print(f"Result: {result}")
-        # If the result is a Pydantic model, convert to dict
-        if hasattr(result, "model_dump"):
-            result = result.model_dump()
+        print(f"Result: {result}")  # Reduce verbosity
+
+        # Convert result to dict if it's a Pydantic model
+        if hasattr(result, "model_dump"):  # Use model_dump for Pydantic v2+
+            result_content = result.model_dump(mode="json")  # Serialize properly
+        elif hasattr(result, "dict"):  # Fallback for older Pydantic
+            result_content = result.dict()
+        else:
+            result_content = result  # Assume JSON-serializable
 
         # Prepare the response
         response = {
             "kind": "StreamResponseMessage",
             "id": message_id,
-            "content": result,
+            "content": result_content,
             "status": "success",
             "created_at": datetime.now().isoformat(),
-            "user_id": user_id,
+            "user_id": user_id,  # Pass user_id back
         }
 
-        print(f"Response: {response}")
+        # print(f"Response: {response}") # Reduce verbosity
 
         # Send the result to the return stream
         if return_stream:
-            # Assert type again closer to usage for type checker clarity
             assert isinstance(return_stream, str)
             r.xadd(return_stream, {"data": json.dumps(response)})
             print(f"Processed message {message_id}, result sent to {return_stream}")
 
         # Acknowledge the message
-        # Assert types again closer to usage for type checker clarity
         assert isinstance(REDIS_STREAM, str)
         assert isinstance(REDIS_CONSUMER_GROUP, str)
         r.xack(REDIS_STREAM, REDIS_CONSUMER_GROUP, message_id)
@@ -523,77 +591,201 @@ def process_message(message_id: str, message_data: Dict[str, str]) -> None:
     except Exception as e:
         print(f"Error processing message {message_id}: {e}")
         traceback.print_exc()
+        _send_error_response(
+            message_id, str(e), traceback.format_exc(), return_stream, user_id
+        )
 
-        # Prepare the error response
-        error_response = {
-            "kind": "StreamResponseMessage",
-            "id": message_id,
-            "content": {
-                "error": str(e),
-                "traceback": traceback.format_exc(),
-            },
-            "status": "error",
-            "created_at": datetime.now().isoformat(),
-            "user_id": user_id,
-        }
-
-        # Send the error to the return stream
-        if return_stream:
-            # Assert type again closer to usage for type checker clarity
-            assert isinstance(return_stream, str)
-            r.xadd(return_stream, {"data": json.dumps(error_response)})
-        else:
-            # Assert type again closer to usage for type checker clarity
+        # Acknowledge the message even if processing failed
+        try:
             assert isinstance(REDIS_STREAM, str)
-            r.xadd(f"{REDIS_STREAM}.errors", {"data": json.dumps(error_response)})
+            assert isinstance(REDIS_CONSUMER_GROUP, str)
+            r.xack(REDIS_STREAM, REDIS_CONSUMER_GROUP, message_id)
+            print(f"Acknowledged failed message {message_id}")
+        except Exception as e_ack:
+            print(
+                f"CRITICAL: Failed to acknowledge failed message {message_id}: {e_ack}"
+            )
 
-        # Still acknowledge the message so we don't reprocess it
-        # Assert types again closer to usage for type checker clarity
-        assert isinstance(REDIS_STREAM, str)
-        assert isinstance(REDIS_CONSUMER_GROUP, str)
-        r.xack(REDIS_STREAM, REDIS_CONSUMER_GROUP, message_id)
+
+# --- Helper to Send Error Response ---
+def _send_error_response(
+    message_id: str,
+    error_msg: str,
+    tb: str,
+    return_stream: Optional[str],
+    user_id: Optional[str],
+):
+    """Sends a standardized error response to Redis."""
+    global r, REDIS_STREAM  # Access global Redis connection and stream name
+
+    error_response = {
+        "kind": "StreamResponseMessage",
+        "id": message_id,
+        "content": {
+            "error": error_msg,
+            "traceback": tb,
+        },
+        "status": "error",
+        "created_at": datetime.now(timezone.utc).isoformat(),  # Use UTC
+        "user_id": user_id,
+    }
+
+    error_destination = f"{REDIS_STREAM}.errors"  # Default error stream
+    if return_stream:  # Prefer return_stream if available
+        error_destination = return_stream
+
+    try:
+        assert isinstance(error_destination, str)
+        r.xadd(error_destination, {"data": json.dumps(error_response)})
+        print(f"Sent error response for message {message_id} to {error_destination}")
+    except Exception as e_redis:
+        print(
+            f"CRITICAL: Failed to send error response for {message_id} to Redis: {e_redis}"
+        )
+        traceback.print_exc()
 
 
 # Main loop
 print(f"Starting consumer for stream {REDIS_STREAM} in group {REDIS_CONSUMER_GROUP}")
-consumer_name = f"consumer-{os.getpid()}"
+consumer_name = f"consumer-{os.getpid()}-{socket.gethostname()}"  # More unique name
 
-while True:
-    try:
-        # Assert types just before use in the loop
-        assert isinstance(REDIS_STREAM, str)
-        assert isinstance(REDIS_CONSUMER_GROUP, str)
+try:
+    while True:
+        try:
+            # --- Check for Code Updates ---
+            if entrypoint_abs_path:  # Should always be set after init
+                try:
+                    current_mtime = os.path.getmtime(entrypoint_abs_path)
+                    if current_mtime > last_load_mtime:
+                        print(
+                            f"[Consumer] Detected change in entrypoint file: {entrypoint_abs_path}. Reloading code..."
+                        )
+                        (
+                            reloaded_target_func,
+                            reloaded_init_func,
+                            reloaded_module,
+                            reloaded_namespace,
+                            new_mtime,
+                        ) = load_or_reload_user_code(
+                            _module_path,
+                            _function_name,
+                            entrypoint_abs_path,
+                            _init_func_name,
+                            _included_object_sources,
+                        )
 
-        # Read from stream with blocking
-        streams = {REDIS_STREAM: ">"}  # '>' means read only new messages
-        # The type checker still struggles here, but the runtime types are asserted.
-        print("reading from stream...")
-        messages = r.xreadgroup(  # type: ignore[arg-type]
-            REDIS_CONSUMER_GROUP, consumer_name, streams, count=1, block=5000
-        )
+                        if (
+                            reloaded_target_func is not None
+                            and reloaded_module is not None
+                        ):
+                            print(
+                                "[Consumer] Code reload successful. Updating functions."
+                            )
+                            target_function = reloaded_target_func
+                            init_function = reloaded_init_func  # Update init ref too, though it's already run
+                            imported_module = reloaded_module
+                            local_namespace = (
+                                reloaded_namespace  # Update namespace from includes
+                            )
+                            last_load_mtime = new_mtime
+                        else:
+                            print(
+                                "[Consumer] Code reload failed. Continuing with previously loaded code."
+                            )
+                            # Optionally: Send an alert/log prominently that reload failed
+
+                except FileNotFoundError:
+                    print(
+                        f"[Consumer] Error: Entrypoint file '{entrypoint_abs_path}' not found during check. Cannot reload."
+                    )
+                    # Mark as non-runnable? Or just log?
+                    target_function = None  # Stop processing until file reappears?
+                    imported_module = None
+                    last_load_mtime = 0  # Reset mtime to force check next time
+                except Exception as e_reload_check:
+                    print(f"[Consumer] Error checking/reloading code: {e_reload_check}")
+                    traceback.print_exc()
+            else:
+                print(
+                    "[Consumer] Warning: Entrypoint absolute path not set, cannot check for code updates."
+                )
 
-        if not messages:
-            # No messages received, continue waiting
-            continue
+            # --- Read from Redis Stream ---
+            if target_function is None:
+                # If code failed to load initially or during reload, wait before retrying
+                print(
+                    "[Consumer] Target function not loaded, waiting 5s before checking again..."
+                )
+                time.sleep(5)
+                continue  # Skip reading from Redis until code is loaded
 
-        # Assert that messages is a list (expected synchronous return type)
-        assert isinstance(
-            messages, list
-        ), f"Expected list from xreadgroup, got {type(messages)}"
-        assert len(messages) > 0  # Ensure the list is not empty before indexing
+            assert isinstance(REDIS_STREAM, str)
+            assert isinstance(REDIS_CONSUMER_GROUP, str)
 
-        stream_name, stream_messages = messages[0]
+            # Simplified xreadgroup call - redis-py handles encoding now
+            # With decode_responses=True, redis-py expects str types here
+            streams_arg = {REDIS_STREAM: ">"}
+            messages = r.xreadgroup(
+                REDIS_CONSUMER_GROUP,
+                consumer_name,
+                streams_arg,
+                count=1,
+                block=5000,  # Use milliseconds for block
+            )
 
-        for message_id, message_data in stream_messages:
-            print(f"Processing message {message_id}")
-            print(f"Message data: {message_data}")
-            process_message(message_id, message_data)
+            if not messages:
+                # print("[Consumer] No new messages.") # Reduce verbosity
+                continue
 
-    except ConnectionError as e:
-        print(f"Redis connection error: {e}")
-        time.sleep(5)  # Wait before retrying
+            # Assert messages is not None to help type checker (already implied by `if not messages`)
+            assert messages is not None
 
-    except Exception as e:
-        print(f"Unexpected error: {e}")
-        traceback.print_exc()
-        time.sleep(1)  # Brief pause before continuing
+            # Cast messages to expected type to satisfy type checker
+            typed_messages = cast(
+                List[Tuple[str, List[Tuple[str, Dict[str, str]]]]], messages
+            )
+            stream_name_str, stream_messages = typed_messages[0]
+
+            # for msg_id_bytes, msg_data_bytes_dict in stream_messages: # Original structure
+            for (
+                message_id_str,
+                message_data_str_dict,
+            ) in stream_messages:  # Structure with decode_responses=True
+                # message_id_str = msg_id_bytes.decode('utf-8') # No longer needed
+                # Decode keys/values in the message data dict
+                # message_data_str_dict = { k.decode('utf-8'): v.decode('utf-8')
+                #                          for k, v in msg_data_bytes_dict.items() } # No longer needed
+                # print(f"Processing message {message_id_str}") # Reduce verbosity
+                # print(f"Message data: {message_data_str_dict}") # Reduce verbosity
+                process_message(message_id_str, message_data_str_dict)
+
+        except ConnectionError as e:
+            print(f"Redis connection error: {e}. Reconnecting in 5s...")
+            time.sleep(5)
+            # Attempt to reconnect explicitly
+            try:
+                print("Attempting Redis reconnection...")
+                # Close existing potentially broken connection? `r.close()` if available
+                r = redis.from_url(REDIS_URL, decode_responses=True)
+                r.ping()
+                print("Reconnected to Redis.")
+            except Exception as recon_e:
+                print(f"Failed to reconnect to Redis: {recon_e}")
+                # Keep waiting
+
+        except ResponseError as e:
+            print(f"Redis command error: {e}")
+            # Should we exit or retry?
+            if "NOGROUP" in str(e):
+                print("Consumer group seems to have disappeared. Exiting.")
+                sys.exit(1)
+            time.sleep(1)
+
+        except Exception as e:
+            print(f"Unexpected error in main loop: {e}")
+            traceback.print_exc()
+            time.sleep(1)
+
+finally:
+    print("Consumer loop exited.")
+    # Any other cleanup needed?