nebu 0.1.14__py3-none-any.whl → 0.1.17__py3-none-any.whl

nebu/containers/decorator.py

@@ -1,9 +1,9 @@
 import base64
-import inspect
 import pickle
 import time
 from typing import Any, Callable, List, Optional
 
+import dill  # Import dill
 import requests
 
 from nebu.containers.container import Container
@@ -47,8 +47,16 @@ def container(
                 )
                 time.sleep(1)
 
-            # Get function source code
-            func_code = inspect.getsource(func)
+            # Get function source code using dill
+            try:
+                func_code = dill.source.getsource(func)
+            except (OSError, TypeError) as e:
+                raise RuntimeError(
+                    f"Failed to retrieve source code for function '{func.__name__}'. "
+                    "This can happen with functions defined dynamically or interactively "
+                    "(e.g., in a Jupyter notebook or REPL). Ensure the function is defined "
+                    f"in a standard Python module if possible. Original error: {e}"
+                )
 
             # Serialize arguments using pickle for complex objects
             serialized_args = base64.b64encode(pickle.dumps(args)).decode("utf-8")

nebu/processors/consumer.py

@@ -1,6 +1,7 @@
 #!/usr/bin/env python3
 import json
 import os
+import socket
 import sys
 import time
 import traceback
@@ -8,6 +9,7 @@ from datetime import datetime
 from typing import Dict, TypeVar
 
 import redis
+import socks
 
 # Define TypeVar for generic models
 T = TypeVar("T")
@@ -97,6 +99,26 @@ try:
     exec("from nebu.processors.models import *", local_namespace)
     exec("from nebu.processors.processor import *", local_namespace)
 
+    # Execute included object sources FIRST, 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)
+            print(f"[Consumer] Successfully executed included object {i} base source")
+            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")
@@ -176,22 +198,6 @@ try:
             print(f"Error defining output model: {e}")
             traceback.print_exc()
 
-    # Execute included object sources
-    for i, (obj_source, args_sources) in enumerate(included_object_sources):
-        try:
-            exec(obj_source, local_namespace)
-            print(f"Successfully executed included object {i} base source")
-            for j, arg_source in enumerate(args_sources):
-                try:
-                    exec(arg_source, local_namespace)
-                    print(f"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()
-
     # Finally, execute the function code
     try:
         exec(function_source, local_namespace)
@@ -216,13 +222,24 @@ if not all([REDIS_URL, REDIS_CONSUMER_GROUP, REDIS_STREAM]):
     print("Missing required Redis environment variables")
     sys.exit(1)
 
+# Configure SOCKS proxy before connecting to Redis
+# Use the proxy settings provided by tailscaled
+socks.set_default_proxy(socks.SOCKS5, "localhost", 1055)
+socket.socket = socks.socksocket
+print("Configured SOCKS5 proxy for socket connections via localhost:1055")
+
 # Connect to Redis
 try:
-    r = redis.from_url(REDIS_URL)
+    # Parse the Redis URL to handle potential credentials or specific DBs if needed
+    # Although from_url should work now with the patched socket
+    r = redis.from_url(
+        REDIS_URL, decode_responses=True
+    )  # Added decode_responses for convenience
+    r.ping()  # Test connection
     redis_info = REDIS_URL.split("@")[-1] if "@" in REDIS_URL else REDIS_URL
-    print(f"Connected to Redis at {redis_info}")
+    print(f"Connected to Redis via SOCKS proxy at {redis_info}")
 except Exception as e:
-    print(f"Failed to connect to Redis: {e}")
+    print(f"Failed to connect to Redis via SOCKS proxy: {e}")
     traceback.print_exc()
     sys.exit(1)
 

nebu/processors/decorate.py

@@ -1,7 +1,20 @@
+import ast  # For parsing notebook code
 import inspect
+import re  # Import re for fallback check
 import textwrap
-from typing import Any, Callable, Dict, List, Optional, TypeVar, get_type_hints
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    List,
+    Optional,
+    TypeVar,
+    get_args,
+    get_origin,
+    get_type_hints,
+)
 
+import dill  # Add dill import
 from pydantic import BaseModel
 
 from nebu.containers.models import (
@@ -24,35 +37,300 @@ from .default import DEFAULT_MAX_REPLICAS, DEFAULT_MIN_REPLICAS, DEFAULT_SCALE
 T = TypeVar("T", bound=BaseModel)
 R = TypeVar("R", bound=BaseModel)
 
+# Attribute name for explicitly stored source
+_NEBU_EXPLICIT_SOURCE_ATTR = "_nebu_explicit_source"
+
+# --- Jupyter Helper Functions ---
+
+
+def is_jupyter_notebook():
+    """
+    Determine if the current code is running inside a Jupyter notebook.
+    Returns bool: True if running inside a Jupyter notebook, False otherwise.
+    """
+    print("[DEBUG Helper] Checking if running in Jupyter...")
+    try:
+        import IPython
+
+        ip = IPython.get_ipython()
+        if ip is None:
+            print("[DEBUG Helper] is_jupyter_notebook: No IPython instance found.")
+            return False
+        class_name = str(ip.__class__)
+        print(f"[DEBUG Helper] is_jupyter_notebook: IPython class name: {class_name}")
+        if "ZMQInteractiveShell" in class_name:
+            print(
+                "[DEBUG Helper] is_jupyter_notebook: Jupyter detected (ZMQInteractiveShell)."
+            )
+            return True
+        print(
+            "[DEBUG Helper] is_jupyter_notebook: Not Jupyter (IPython instance found, but not ZMQInteractiveShell)."
+        )
+        return False
+    except Exception as e:
+        print(f"[DEBUG Helper] is_jupyter_notebook: Exception occurred: {e}")
+        return False
+
+
+def get_notebook_executed_code():
+    """
+    Returns all executed code from the current notebook session.
+    Returns str or None: All executed code as a string, or None if not possible.
+    """
+    print("[DEBUG Helper] Attempting to get notebook execution history...")
+    try:
+        import IPython
+
+        ip = IPython.get_ipython()
+        if ip is None or not hasattr(ip, "history_manager"):
+            print(
+                "[DEBUG Helper] get_notebook_executed_code: No IPython instance or history_manager."
+            )
+            return None
+        history_manager = ip.history_manager
+        # Limiting history range for debugging? Maybe get_tail(N)? For now, get all.
+        # history = history_manager.get_range(start=1) # type: ignore
+        history = list(history_manager.get_range(start=1))  # type: ignore # Convert to list to get length
+        print(
+            f"[DEBUG Helper] get_notebook_executed_code: Retrieved {len(history)} history entries."
+        )
+        source_code = ""
+        separator = "\n#<NEBU_CELL_SEP>#\n"
+        for _, _, content in history:  # Use _ for unused session, lineno
+            if isinstance(content, str) and content.strip():
+                source_code += content + separator
+        print(
+            f"[DEBUG Helper] get_notebook_executed_code: Total history source length: {len(source_code)}"
+        )
+        return source_code
+    except Exception as e:
+        print(f"[DEBUG Helper] get_notebook_executed_code: Error getting history: {e}")
+        return None
+
+
+def extract_definition_source_from_string(
+    source_string: str, def_name: str, def_type: type = ast.FunctionDef
+) -> Optional[str]:
+    """
+    Attempts to extract the source code of a function or class from a larger string
+    (like notebook history). Finds the *last* complete definition.
+    Uses AST parsing for robustness.
+    def_type can be ast.FunctionDef or ast.ClassDef.
+    """
+    print(
+        f"[DEBUG Helper] Extracting '{def_name}' ({def_type.__name__}) from history string (len: {len(source_string)})..."
+    )
+    if not source_string or not def_name:
+        print("[DEBUG Helper] extract: Empty source string or def_name.")
+        return None
+
+    cells = source_string.split("#<NEBU_CELL_SEP>#")
+    print(f"[DEBUG Helper] extract: Split history into {len(cells)} potential cells.")
+    last_found_source = None
+
+    for i, cell in enumerate(reversed(cells)):
+        cell_num = len(cells) - 1 - i
+        cell = cell.strip()
+        if not cell:
+            continue
+        # print(f"[DEBUG Helper] extract: Analyzing cell #{cell_num}...") # Can be very verbose
+        try:
+            tree = ast.parse(cell)
+            found_in_cell = False
+            for node in ast.walk(tree):
+                if (
+                    isinstance(node, def_type)
+                    and hasattr(node, "name")
+                    and node.name == def_name
+                ):
+                    print(
+                        f"[DEBUG Helper] extract: Found node for '{def_name}' in cell #{cell_num}."
+                    )
+                    try:
+                        # Use ast.get_source_segment for accurate extraction (Python 3.8+)
+                        func_source = ast.get_source_segment(cell, node)
+                        if func_source:
+                            print(
+                                f"[DEBUG Helper] extract: Successfully extracted source using get_source_segment for '{def_name}'."
+                            )
+                            last_found_source = func_source
+                            found_in_cell = True
+                            break  # Stop searching this cell
+                    except AttributeError:  # Fallback for Python < 3.8
+                        print(
+                            f"[DEBUG Helper] extract: get_source_segment failed (likely Py < 3.8), using fallback for '{def_name}'."
+                        )
+                        start_lineno = getattr(node, "lineno", 1) - 1
+                        end_lineno = getattr(node, "end_lineno", start_lineno + 1)
+
+                        if hasattr(node, "decorator_list") and node.decorator_list:
+                            first_decorator_start_line = (
+                                getattr(
+                                    node.decorator_list[0], "lineno", start_lineno + 1
+                                )
+                                - 1
+                            )  # type: ignore
+                            start_lineno = min(start_lineno, first_decorator_start_line)
+
+                        lines = cell.splitlines()
+                        if 0 <= start_lineno < len(lines) and end_lineno <= len(lines):
+                            extracted_lines = lines[start_lineno:end_lineno]
+                            if extracted_lines and (
+                                extracted_lines[0].strip().startswith("@")
+                                or extracted_lines[0]
+                                .strip()
+                                .startswith(("def ", "class "))
+                            ):
+                                last_found_source = "\n".join(extracted_lines)
+                                print(
+                                    f"[DEBUG Helper] extract: Extracted source via fallback for '{def_name}'."
+                                )
+                                found_in_cell = True
+                                break
+                        else:
+                            print(
+                                f"[DEBUG Helper] extract: Warning: Line numbers out of bounds for {def_name} in cell (fallback)."
+                            )
+
+            if found_in_cell:
+                print(
+                    f"[DEBUG Helper] extract: Found and returning source for '{def_name}' from cell #{cell_num}."
+                )
+                return last_found_source  # Found last definition, return immediately
+
+        except (SyntaxError, ValueError) as e:
+            # print(f"[DEBUG Helper] extract: Skipping cell #{cell_num} due to parse error: {e}") # Can be verbose
+            continue
+        except Exception as e:
+            print(
+                f"[DEBUG Helper] extract: Warning: AST processing error for {def_name} in cell #{cell_num}: {e}"
+            )
+            continue
+
+    if not last_found_source:
+        print(
+            f"[DEBUG Helper] extract: Definition '{def_name}' of type {def_type.__name__} not found in history search."
+        )
+    return last_found_source
+
+
+# --- End Jupyter Helper Functions ---
+
+
+def include(obj: Any) -> Any:
+    """
+    Decorator to explicitly capture the source code of a function or class,
+    intended for use in environments where inspect/dill might fail (e.g., Jupyter).
+    """
+    try:
+        source = dill.source.getsource(obj)
+        dedented_source = textwrap.dedent(source)
+        setattr(obj, _NEBU_EXPLICIT_SOURCE_ATTR, dedented_source)
+        print(
+            f"[DEBUG @include] Successfully captured source for: {getattr(obj, '__name__', str(obj))}"
+        )
+    except Exception as e:
+        # Don't fail the definition, just warn
+        print(
+            f"Warning: @include could not capture source for {getattr(obj, '__name__', str(obj))}: {e}. Automatic source retrieval will be attempted later."
+        )
+    return obj
+
+
+def get_model_source(
+    model_class: Any, notebook_code: Optional[str] = None
+) -> Optional[str]:
+    """
+    Get the source code of a model class.
+    Checks explicit source, then notebook history (if provided), then dill.
+    """
+    model_name_str = getattr(model_class, "__name__", str(model_class))
+    print(f"[DEBUG get_model_source] Getting source for: {model_name_str}")
+    # 1. Check explicit source
+    explicit_source = getattr(model_class, _NEBU_EXPLICIT_SOURCE_ATTR, None)
+    if explicit_source:
+        print(
+            f"[DEBUG get_model_source] Using explicit source (@include) for: {model_name_str}"
+        )
+        return explicit_source
+
+    # 2. Check notebook history
+    if notebook_code and hasattr(model_class, "__name__"):
+        print(
+            f"[DEBUG get_model_source] Attempting notebook history extraction for: {model_class.__name__}"
+        )
+        extracted_source = extract_definition_source_from_string(
+            notebook_code, model_class.__name__, ast.ClassDef
+        )
+        if extracted_source:
+            print(
+                f"[DEBUG get_model_source] Using notebook history source for: {model_class.__name__}"
+            )
+            return extracted_source
+        else:
+            print(
+                f"[DEBUG get_model_source] Notebook history extraction failed for: {model_class.__name__}. Proceeding to dill."
+            )
 
-def get_model_source(model_class: Any) -> Optional[str]:
-    """Get the source code of a model class."""
+    # 3. Fallback to dill
     try:
-        source = inspect.getsource(model_class)
+        print(
+            f"[DEBUG get_model_source] Attempting dill fallback for: {model_name_str}"
+        )
+        source = dill.source.getsource(model_class)
+        print(f"[DEBUG get_model_source] Using dill source for: {model_name_str}")
         return textwrap.dedent(source)
-    except (IOError, TypeError):
+    except (IOError, TypeError, OSError) as e:
+        print(
+            f"[DEBUG get_model_source] Failed dill fallback for: {model_name_str}: {e}"
+        )
         return None
 
 
-def get_type_source(type_obj: Any) -> Optional[Any]:
+# Reintroduce get_type_source to handle generics
+def get_type_source(
+    type_obj: Any, notebook_code: Optional[str] = None
+) -> Optional[Any]:
     """Get the source code for a type, including generic parameters."""
-    # If it's a class, get its source
-    if isinstance(type_obj, type):
-        return get_model_source(type_obj)
+    type_obj_str = str(type_obj)
+    print(f"[DEBUG get_type_source] Getting source for type: {type_obj_str}")
+    origin = get_origin(type_obj)
+    args = get_args(type_obj)
 
-    # If it's a GenericAlias (like V1StreamMessage[SomeType])
-    if hasattr(type_obj, "__origin__") and hasattr(type_obj, "__args__"):
-        origin_source = get_model_source(type_obj.__origin__)
+    if origin is not None:
+        # Use updated get_model_source for origin
+        print(
+            f"[DEBUG get_type_source] Detected generic type. Origin: {origin}, Args: {args}"
+        )
+        origin_source = get_model_source(origin, notebook_code)
         args_sources = []
 
-        # Get sources for all type arguments
-        for arg in type_obj.__args__:
-            arg_source = get_type_source(arg)
+        # Recursively get sources for all type arguments
+        for arg in args:
+            print(
+                f"[DEBUG get_type_source] Recursively getting source for generic arg #{arg}"
+            )
+            arg_source = get_type_source(arg, notebook_code)
             if arg_source:
                 args_sources.append(arg_source)
 
-        return origin_source, args_sources
+        # Return tuple only if origin source or some arg sources were found
+        if origin_source or args_sources:
+            print(
+                f"[DEBUG get_type_source] Returning tuple source for generic: {type_obj_str}"
+            )
+            return (origin_source, args_sources)
+
+    # Fallback if not a class or recognizable generic alias
+    # Try get_model_source as a last resort for unknown types
+    fallback_source = get_model_source(type_obj, notebook_code)
+    if fallback_source:
+        print(
+            f"[DEBUG get_type_source] Using fallback get_model_source for: {type_obj_str}"
+        )
+        return fallback_source
 
+    print(f"[DEBUG get_type_source] Failed to get source for: {type_obj_str}")
     return None
 
 
@@ -75,297 +353,435 @@ def processor(
     no_delete: bool = False,
     include: Optional[List[Any]] = None,
 ):
-    """
-    Decorator that converts a function into a Processor.
-
-    Args:
-        image: The container image to use for the processor
-        setup_script: Optional setup script to run before starting the processor
-        scale: Optional scaling configuration
-        min_replicas: Minimum number of replicas to maintain
-        max_replicas: Maximum number of replicas to scale to
-        platform: Optional compute platform to run on
-        accelerators: Optional list of accelerator types
-        namespace: Optional namespace for the processor
-        labels: Optional labels to apply to the processor
-        env: Optional environment variables
-        volumes: Optional volume mounts
-        resources: Optional resource requirements
-        meters: Optional metering configuration
-        authz: Optional authorization configuration
-        python_cmd: Optional python command to use
-        no_delete: Whether to prevent deleting the processor on updates
-        include: Optional list of Python objects whose source code should be included
-    """
+    def decorator(
+        func: Callable[[Any], Any],
+    ) -> Processor:
+        # Moved init print here
+        print(
+            f"[DEBUG Decorator Init] @processor decorating function '{func.__name__}'"
+        )
+        all_env = env or []
+        processor_name = func.__name__
 
-    def decorator(func: Callable[[T], R]) -> Processor:
-        # Validate that the function takes a single parameter that is a BaseModel
+        # --- Determine Environment and Get Notebook Code ---
+        print("[DEBUG Decorator] Determining execution environment...")
+        in_jupyter = is_jupyter_notebook()
+        notebook_code = None
+        if in_jupyter:
+            print("[DEBUG Decorator] Jupyter environment detected.")
+            notebook_code = get_notebook_executed_code()
+            if not notebook_code:
+                print(
+                    "[DEBUG Decorator] Warning: Failed to get Jupyter execution history. Will attempt dill."
+                )
+            else:
+                print(
+                    f"[DEBUG Decorator] Retrieved notebook history (length: {len(notebook_code)})."
+                )
+        else:
+            print("[DEBUG Decorator] Non-Jupyter environment detected.")
+        # --- End Environment Determination ---
+
+        # --- Process Manually Included Objects ---
+        included_sources: Dict[Any, Any] = {}
+        if include:
+            print(f"[DEBUG Decorator] Processing manually included objects: {include}")
+            for i, obj in enumerate(include):
+                obj_name_str = getattr(obj, "__name__", str(obj))
+                print(
+                    f"[DEBUG Decorator] Getting source for manually included object: {obj_name_str}"
+                )
+                obj_source = get_model_source(obj, notebook_code)
+                if obj_source:
+                    included_sources[obj] = obj_source
+                    env_key = f"INCLUDED_OBJECT_{i}_SOURCE"
+                    all_env.append(V1EnvVar(key=env_key, value=obj_source))
+                    print(
+                        f"[DEBUG Decorator] Added source to env for included obj: {obj_name_str}"
+                    )
+                else:
+                    print(
+                        f"Warning: Could not retrieve source for manually included object: {obj_name_str}"
+                    )
+            print(
+                f"[DEBUG Decorator] Finished processing included objects. Sources found: {len(included_sources)}"
+            )
+        else:
+            print("[DEBUG Decorator] No manually included objects specified.")
+        # --- End Manually Included Objects ---
+
+        # --- Validate Function Signature and Types ---
+        print(
+            f"[DEBUG Decorator] Validating signature and type hints for {processor_name}..."
+        )
         sig = inspect.signature(func)
         params = list(sig.parameters.values())
-
         if len(params) != 1:
-            raise TypeError(f"Function {func.__name__} must take exactly one parameter")
+            raise TypeError(f"{processor_name} must take one parameter")
+
+        try:
+            type_hints = get_type_hints(func, globalns=func.__globals__, localns=None)
+            print(f"[DEBUG Decorator] Raw type hints: {type_hints}")
+        except Exception as e:
+            print(f"[DEBUG Decorator] Error getting type hints: {e}")
+            raise TypeError(
+                f"Could not evaluate type hints for {processor_name}: {e}"
+            ) from e
 
-        # Check parameter type
-        type_hints = get_type_hints(func)
         param_name = params[0].name
         if param_name not in type_hints:
             raise TypeError(
-                f"Parameter {param_name} in function {func.__name__} must have a type annotation"
+                f"{processor_name} parameter '{param_name}' must have type hint"
             )
-
         param_type = type_hints[param_name]
+        param_type_str_repr = str(param_type)  # Use string for regex
+        print(
+            f"[DEBUG Decorator] Parameter '{param_name}' type hint: {param_type_str_repr}"
+        )
 
-        # Check if input type is V1StreamMessage or a subclass
+        if "return" not in type_hints:
+            raise TypeError(f"{processor_name} must have return type hint")
+        return_type = type_hints["return"]
+        print(f"[DEBUG Decorator] Return type hint: {return_type}")
+
+        # --- Determine Input Type (StreamMessage, ContentType) ---
+        print(
+            f"[DEBUG Decorator] Determining input type structure for param type hint: {param_type_str_repr}"
+        )
+        origin = get_origin(param_type)
+        args = get_args(param_type)
+        print(f"[DEBUG Decorator] get_origin result: {origin}, get_args result: {args}")
         is_stream_message = False
         content_type = None
 
-        # Handle generic V1StreamMessage
-        if (
-            hasattr(param_type, "__origin__")
-            and param_type.__origin__ == V1StreamMessage
+        # Check 1: Standard introspection
+        if origin is V1StreamMessage or (
+            origin is not None
+            and origin.__name__ == V1StreamMessage.__name__
+            and origin.__module__ == V1StreamMessage.__module__
         ):
+            print(
+                "[DEBUG Decorator] Input type identified as V1StreamMessage via get_origin."
+            )
             is_stream_message = True
-            # Extract the content type from V1StreamMessage[ContentType]
-            if hasattr(param_type, "__args__") and param_type.__args__:
-                content_type = param_type.__args__[0]
-        # Handle direct V1StreamMessage
-        elif param_type is V1StreamMessage:
+            if args:
+                content_type = args[0]
+                print(
+                    f"[DEBUG Decorator] Content type extracted via get_args: {content_type}"
+                )
+            else:
+                print(
+                    "[DEBUG Decorator] V1StreamMessage detected, but no generic arguments found via get_args."
+                )
+        # Check 2: Direct type check
+        elif isinstance(param_type, type) and param_type is V1StreamMessage:
+            print(
+                "[DEBUG Decorator] Input type identified as direct V1StreamMessage type."
+            )
             is_stream_message = True
-
-        # Ensure the parameter is a BaseModel
-        actual_type = (
-            param_type.__origin__ if hasattr(param_type, "__origin__") else param_type  # type: ignore
-        )
-        if not issubclass(actual_type, BaseModel):
-            raise TypeError(
-                f"Parameter {param_name} in function {func.__name__} must be a BaseModel"
+        # Check 3: Regex fallback on string representation
+        elif origin is None:
+            print(
+                f"[DEBUG Decorator] get_origin failed. Attempting regex fallback on type string: '{param_type_str_repr}'"
             )
-
-        # Check return type
-        if "return" not in type_hints:
-            raise TypeError(
-                f"Function {func.__name__} must have a return type annotation"
+            # Use param_type_str_repr in match calls
+            generic_match = re.match(
+                r"^<class '(?:[a-zA-Z0-9_.]+\.)?V1StreamMessage\[(.+?)\]'>$",
+                param_type_str_repr,
+            )
+            if generic_match:
+                print(
+                    "[DEBUG Decorator] Regex matched generic V1StreamMessage pattern!"
+                )
+                is_stream_message = True
+                content_type_name_str = generic_match.group(1).strip()
+                print(
+                    f"[DEBUG Decorator] Captured content type name via regex: '{content_type_name_str}'"
+                )
+                try:
+                    resolved_type = eval(content_type_name_str, func.__globals__)
+                    content_type = resolved_type
+                    print(
+                        f"[DEBUG Decorator] Successfully resolved content type name '{content_type_name_str}' to type: {content_type}"
+                    )
+                except NameError:
+                    print(
+                        f"[DEBUG Decorator] Warning: Regex found content type name '{content_type_name_str}', but it's not defined in function's globals. Consumer might fail."
+                    )
+                    content_type = None
+                except Exception as e:
+                    print(
+                        f"[DEBUG Decorator] Warning: Error evaluating content type name '{content_type_name_str}': {e}"
+                    )
+                    content_type = None
+            else:
+                # Use param_type_str_repr in match calls
+                simple_match = re.match(
+                    r"^<class '(?:[a-zA-Z0-9_.]+\.)?V1StreamMessage'>$",
+                    param_type_str_repr,
+                )
+                if simple_match:
+                    print(
+                        "[DEBUG Decorator] Regex identified direct V1StreamMessage (no generic) from string."
+                    )
+                    is_stream_message = True
+                else:
+                    print(
+                        f"[DEBUG Decorator] Regex did not match V1StreamMessage pattern for string '{param_type_str_repr}'. Assuming not StreamMessage."
+                    )
+        else:
+            print(
+                f"[DEBUG Decorator] Input parameter '{param_name}' type ({param_type_str_repr}) identified as non-StreamMessage type (origin was not None and not V1StreamMessage)."
             )
 
-        return_type = type_hints["return"]
-        actual_return_type = (
-            return_type.__origin__
-            if hasattr(return_type, "__origin__")
-            else return_type
+        print(
+            f"[DEBUG Decorator] Final Input Type Determination: is_stream_message={is_stream_message}, content_type={content_type}"
         )
-        if not issubclass(actual_return_type, BaseModel):
-            raise TypeError(
-                f"Return value of function {func.__name__} must be a BaseModel"
+        # --- End Input Type Determination ---
+
+        # --- Validate Types are BaseModel ---
+        print(
+            "[DEBUG Decorator] Validating parameter and return types are BaseModel subclasses..."
+        )
+
+        def check_basemodel(type_to_check, desc):
+            print(
+                f"[DEBUG Decorator] check_basemodel: Checking {desc} - Type: {type_to_check}"
             )
+            if not type_to_check:
+                print(
+                    f"[DEBUG Decorator] check_basemodel: Skipping check for {desc} (type is None/empty)."
+                )
+                return
+            actual_type = get_origin(type_to_check) or type_to_check
+            print(
+                f"[DEBUG Decorator] check_basemodel: Actual type for {desc}: {actual_type}"
+            )
+            if isinstance(actual_type, type) and not issubclass(actual_type, BaseModel):
+                print(
+                    f"[DEBUG Decorator] check_basemodel: Error - {desc} effective type ({actual_type.__name__}) is not a BaseModel subclass."
+                )
+                raise TypeError(
+                    f"{desc} effective type ({actual_type.__name__}) must be BaseModel subclass"
+                )
+            elif not isinstance(actual_type, type):
+                print(
+                    f"[DEBUG Decorator] check_basemodel: Warning - {desc} effective type '{actual_type}' is not a class. Cannot verify BaseModel subclass."
+                )
+            else:
+                print(
+                    f"[DEBUG Decorator] check_basemodel: OK - {desc} effective type ({actual_type.__name__}) is a BaseModel subclass."
+                )
 
-        # Get function name to use as processor name
-        processor_name = func.__name__
+        effective_param_type = (
+            content_type
+            if is_stream_message and content_type
+            else param_type
+            if not is_stream_message
+            else None
+        )
+        check_basemodel(effective_param_type, f"Parameter '{param_name}'")
+        check_basemodel(return_type, "Return value")
+        print("[DEBUG Decorator] Type validation complete.")
+        # --- End Type Validation ---
 
-        # Prepare environment variables
-        all_env = env or []
+        # --- Get Function Source ---
+        print(
+            f"[DEBUG Decorator] Getting source code for function '{processor_name}'..."
+        )
+        function_source = None
+        explicit_source = getattr(func, _NEBU_EXPLICIT_SOURCE_ATTR, None)
 
-        # Get the source code of the function
-        try:
-            raw_function_source = inspect.getsource(func)
+        if explicit_source:
             print(
-                f"[DEBUG Decorator] Raw source for {func.__name__}:\n{raw_function_source}"
+                f"[DEBUG Decorator] Using explicit source (@include) for function {processor_name}"
             )
-
-            # Clean up the indentation
-            dedented_function_source = textwrap.dedent(raw_function_source)
+            function_source = explicit_source
+        elif in_jupyter and notebook_code:
             print(
-                f"[DEBUG Decorator] Dedented source for {func.__name__}:\n{dedented_function_source}"
+                f"[DEBUG Decorator] Attempting notebook history extraction for function '{processor_name}'..."
             )
-
-            # Find the start of the function definition ('def')
-            # Skip lines starting with '@' or empty lines until 'def' is found
-            lines = dedented_function_source.splitlines()
-            func_def_index = -1
-            for i, line in enumerate(lines):
-                stripped_line = line.strip()
-                if stripped_line.startswith("def "):
-                    func_def_index = i
-                    break
-                # Simply continue if it's not the def line.
-                # This skips decorators and their arguments, regardless of multi-line formatting.
-                continue
-
-            if func_def_index != -1:
-                # Keep lines from the 'def' line onwards
-                function_source = "\n".join(
-                    lines[func_def_index:]
-                )  # Use \n for env var
+            function_source = extract_definition_source_from_string(
+                notebook_code, processor_name, ast.FunctionDef
+            )
+            if function_source:
+                print(
+                    f"[DEBUG Decorator] Found function '{processor_name}' source in notebook history."
+                )
             else:
-                # If 'def' wasn't found (shouldn't happen with valid function source)
-                raise ValueError(
-                    f"Could not find function definition 'def' in source for {func.__name__}"
+                print(
+                    f"[DEBUG Decorator] Failed to find function '{processor_name}' in notebook history, falling back to dill."
                 )
-
+        if function_source is None:
             print(
-                f"[DEBUG Decorator] Processed function source for {func.__name__}:\n{function_source}"
+                f"[DEBUG Decorator] Using dill fallback for function '{processor_name}'..."
             )
+            try:
+                raw_function_source = dill.source.getsource(func)
+                function_source = textwrap.dedent(raw_function_source)
+                print(
+                    f"[DEBUG Decorator] Successfully got source via dill for '{processor_name}'."
+                )
+            except (IOError, TypeError, OSError) as e:
+                print(
+                    f"[DEBUG Decorator] Dill fallback failed for '{processor_name}': {e}"
+                )
+                if not (in_jupyter and notebook_code):
+                    raise ValueError(
+                        f"Could not retrieve source for '{processor_name}' using dill: {e}"
+                    ) from e
 
-        except (IOError, TypeError) as e:
-            print(f"[DEBUG Decorator] Error getting source for {func.__name__}: {e}")
+        if function_source is None:  # Final check after all attempts
             raise ValueError(
-                f"Could not retrieve source code for function {func.__name__}: {e}"
-            ) from e
+                f"Failed to obtain source code for function '{processor_name}' using any method."
+            )
 
-        # Get source code for the models
+        print(f"[DEBUG Decorator] Final function source obtained for '{processor_name}' (len: {len(function_source)}). Source starts:\n-------\
+{function_source[:250]}...\n-------")
+        # --- End Function Source ---
+
+        # --- Get Model Sources ---
+        print("[DEBUG Decorator] Getting model sources...")
         input_model_source = None
         output_model_source = None
         content_type_source = None
+        print("[DEBUG Decorator] Getting base V1StreamMessage source...")
+        stream_message_source = get_type_source(V1StreamMessage, notebook_code)
 
-        # Get the V1StreamMessage class source
-        stream_message_source = get_model_source(V1StreamMessage)
-
-        # Get input model source
         if is_stream_message:
-            input_model_source = stream_message_source
+            print(
+                f"[DEBUG Decorator] Input is StreamMessage. Content type: {content_type}"
+            )
             if content_type:
-                content_type_source = get_type_source(content_type)
-        else:
-            input_model_source = get_type_source(param_type)
-
-        # Get output model source
-        output_model_source = get_type_source(return_type)
-
-        # Add function source code to environment variables
-        print(
-            f"[DEBUG Decorator] Setting FUNCTION_SOURCE: {function_source[:100]}..."
-        )  # Print first 100 chars
-        all_env.append(V1EnvVar(key="FUNCTION_SOURCE", value=function_source))
-        print(f"[DEBUG Decorator] Setting FUNCTION_NAME: {func.__name__}")
-        all_env.append(V1EnvVar(key="FUNCTION_NAME", value=func.__name__))
-
-        # Add model source codes
-        if input_model_source:
-            if isinstance(input_model_source, tuple):
-                all_env.append(
-                    V1EnvVar(key="INPUT_MODEL_SOURCE", value=input_model_source[0])
+                print(
+                    f"[DEBUG Decorator] Getting source for content_type: {content_type}"
                 )
-                # Add generic args sources
-                for i, arg_source in enumerate(input_model_source[1]):
-                    all_env.append(
-                        V1EnvVar(key=f"INPUT_MODEL_ARG_{i}_SOURCE", value=arg_source)
+                content_type_source = get_type_source(content_type, notebook_code)
+                if content_type_source is None:
+                    print(
+                        f"Warning: Failed to get source for content_type: {content_type}"
                     )
-            else:
-                all_env.append(
-                    V1EnvVar(key="INPUT_MODEL_SOURCE", value=input_model_source)
+        else:  # Not a stream message
+            print(
+                f"[DEBUG Decorator] Input is not StreamMessage. Getting source for param_type: {param_type}"
+            )
+            input_model_source = get_type_source(param_type, notebook_code)
+            if input_model_source is None:
+                print(
+                    f"Warning: Failed to get source for input param_type: {param_type}"
                 )
 
-        if output_model_source:
-            if isinstance(output_model_source, tuple):
-                all_env.append(
-                    V1EnvVar(key="OUTPUT_MODEL_SOURCE", value=output_model_source[0])
-                )
-                # Add generic args sources
-                for i, arg_source in enumerate(output_model_source[1]):
-                    all_env.append(
-                        V1EnvVar(key=f"OUTPUT_MODEL_ARG_{i}_SOURCE", value=arg_source)
-                    )
-            else:
-                all_env.append(
-                    V1EnvVar(key="OUTPUT_MODEL_SOURCE", value=output_model_source)
-                )
+        print(f"[DEBUG Decorator] Getting source for return_type: {return_type}")
+        output_model_source = get_type_source(return_type, notebook_code)
+        if output_model_source is None:
+            print(f"Warning: Failed to get source for return_type: {return_type}")
 
-        if stream_message_source:
-            all_env.append(
-                V1EnvVar(key="STREAM_MESSAGE_SOURCE", value=stream_message_source)
-            )
+        print(
+            f"[DEBUG Decorator] Source Result - Content Type: {'Found' if content_type_source else 'Not Found or N/A'}"
+        )
+        print(
+            f"[DEBUG Decorator] Source Result - Input Model (non-stream): {'Found' if input_model_source else 'Not Found or N/A'}"
+        )
+        print(
+            f"[DEBUG Decorator] Source Result - Output Model: {'Found' if output_model_source else 'Not Found'}"
+        )
+        print(
+            f"[DEBUG Decorator] Source Result - Base StreamMessage: {'Found' if stream_message_source else 'Not Found'}"
+        )
+        # --- End Model Sources ---
 
-        if content_type_source:
-            if isinstance(content_type_source, tuple):
-                all_env.append(
-                    V1EnvVar(key="CONTENT_TYPE_SOURCE", value=content_type_source[0])
-                )
-                # Add generic args sources for content type
-                for i, arg_source in enumerate(content_type_source[1]):
-                    all_env.append(
-                        V1EnvVar(key=f"CONTENT_TYPE_ARG_{i}_SOURCE", value=arg_source)
-                    )
-            else:
-                all_env.append(
-                    V1EnvVar(key="CONTENT_TYPE_SOURCE", value=content_type_source)
+        # --- Populate Environment Variables ---
+        print("[DEBUG Decorator] Populating environment variables...")
+        all_env.append(V1EnvVar(key="FUNCTION_SOURCE", value=function_source))
+        all_env.append(V1EnvVar(key="FUNCTION_NAME", value=processor_name))
+
+        def add_source_to_env(key_base: str, source: Any):
+            print(f"[DEBUG Decorator] add_source_to_env: Processing key '{key_base}'")
+            if not source:
+                print(
+                    f"[DEBUG Decorator] add_source_to_env: No source for '{key_base}', skipping."
                 )
+                return
 
-        # Add included object sources
-        if include:
-            for i, obj in enumerate(include):
-                obj_source = get_type_source(
-                    obj
-                )  # Reuse existing function for source retrieval
-                if obj_source:
-                    if isinstance(obj_source, tuple):
-                        # Handle complex types (like generics) if needed, similar to models
+            if isinstance(source, tuple):
+                origin_src, arg_srcs = source
+                print(
+                    f"[DEBUG Decorator] add_source_to_env: '{key_base}' is tuple source. Origin found: {bool(origin_src)}, Num args: {len(arg_srcs)}"
+                )
+                if origin_src and isinstance(origin_src, str):
+                    all_env.append(V1EnvVar(key=f"{key_base}_SOURCE", value=origin_src))
+                    print(f"[DEBUG Decorator] Added env var {key_base}_SOURCE (origin)")
+                for i, arg_src in enumerate(arg_srcs):
+                    if isinstance(arg_src, str):
                         all_env.append(
-                            V1EnvVar(
-                                key=f"INCLUDED_OBJECT_{i}_SOURCE", value=obj_source[0]
-                            )
+                            V1EnvVar(key=f"{key_base}_ARG_{i}_SOURCE", value=arg_src)
+                        )
+                        print(
+                            f"[DEBUG Decorator] Added env var {key_base}_ARG_{i}_SOURCE"
                         )
-                        for j, arg_source in enumerate(obj_source[1]):
+                    elif isinstance(arg_src, tuple):
+                        arg_origin_src, _ = arg_src
+                        if arg_origin_src and isinstance(arg_origin_src, str):
                             all_env.append(
                                 V1EnvVar(
-                                    key=f"INCLUDED_OBJECT_{i}_ARG_{j}_SOURCE",
-                                    value=arg_source,
+                                    key=f"{key_base}_ARG_{i}_SOURCE",
+                                    value=arg_origin_src,
                                 )
                             )
-                    else:
-                        all_env.append(
-                            V1EnvVar(
-                                key=f"INCLUDED_OBJECT_{i}_SOURCE", value=obj_source
+                            print(
+                                f"[DEBUG Decorator] Added env var {key_base}_ARG_{i}_SOURCE (nested origin)"
                             )
+                        else:
+                            print(
+                                f"[DEBUG Decorator] Skipping complex/non-string nested arg origin for {key_base}_ARG_{i}"
+                            )
+                    else:
+                        print(
+                            f"[DEBUG Decorator] Skipping complex/non-string arg source for {key_base}_ARG_{i}"
                         )
-                else:
-                    # Optionally raise an error or log a warning if source can't be found
-                    print(
-                        f"Warning: Could not retrieve source for included object: {obj}"
-                    )
+            elif isinstance(source, str):
+                all_env.append(V1EnvVar(key=f"{key_base}_SOURCE", value=source))
+                print(f"[DEBUG Decorator] Added env var {key_base}_SOURCE (string)")
+            else:
+                print(
+                    f"[DEBUG Decorator] Warning: Unknown source type for {key_base}: {type(source)}. Skipping."
+                )
 
-        # Add parameter and return type info for runtime validation
-        all_env.append(
-            V1EnvVar(
-                key="PARAM_TYPE_NAME",
-                value=param_type.__name__
-                if hasattr(param_type, "__name__")
-                else str(param_type),
-            )
-        )
-        all_env.append(
-            V1EnvVar(
-                key="RETURN_TYPE_NAME",
-                value=return_type.__name__
-                if hasattr(return_type, "__name__")
-                else str(return_type),
-            )
-        )
-        all_env.append(V1EnvVar(key="IS_STREAM_MESSAGE", value=str(is_stream_message)))
+        add_source_to_env("INPUT_MODEL", input_model_source)
+        add_source_to_env("OUTPUT_MODEL", output_model_source)
+        add_source_to_env("CONTENT_TYPE", content_type_source)
+        add_source_to_env("STREAM_MESSAGE", stream_message_source)
 
-        if content_type:
+        print("[DEBUG Decorator] Adding type info env vars...")
+        all_env.append(V1EnvVar(key="PARAM_TYPE_STR", value=param_type_str_repr))
+        all_env.append(V1EnvVar(key="RETURN_TYPE_STR", value=str(return_type)))
+        all_env.append(V1EnvVar(key="IS_STREAM_MESSAGE", value=str(is_stream_message)))
+        if content_type and hasattr(content_type, "__name__"):
             all_env.append(
-                V1EnvVar(
-                    key="CONTENT_TYPE_NAME",
-                    value=content_type.__name__
-                    if hasattr(content_type, "__name__")
-                    else str(content_type),
-                )
+                V1EnvVar(key="CONTENT_TYPE_NAME", value=content_type.__name__)
             )
-
-        # We still add the module for reference, but we won't rely on importing it
         all_env.append(V1EnvVar(key="MODULE_NAME", value=func.__module__))
+        print("[DEBUG Decorator] Finished populating environment variables.")
+        # --- End Environment Variables ---
 
-        # Prepare metadata
+        # --- Final Setup ---
+        print("[DEBUG Decorator] Preparing final Processor object...")
         metadata = V1ResourceMetaRequest(
             name=processor_name, namespace=namespace, labels=labels
         )
-
-        # Create the command to run the consumer directly
         consumer_command = f"{python_cmd} -m nebu.processors.consumer"
+        setup_commands = [
+            f"{python_cmd} -m pip install dill pydantic redis nebu",  # Base deps
+        ]
+        if setup_script:
+            print("[DEBUG Decorator] Adding setup script to command.")
+            setup_commands.append(f"\n{setup_script}\n")
+        setup_commands.append(consumer_command)
+        final_command = "\n".join(setup_commands)
+        print(
+            f"[DEBUG Decorator] Final container command:\n-------\n{final_command}\n-------"
+        )
 
-        final_command = f"{python_cmd} -m pip install redis nebu\n\n{setup_script}\n\n{consumer_command}"
-
-        # Create the V1ContainerRequest
         container_request = V1ContainerRequest(
             image=image,
             command=final_command,
@@ -379,23 +795,29 @@ def processor(
             platform=platform,
             metadata=metadata,
         )
-        print("container_request", container_request)
+        print("[DEBUG Decorator] Final Container Request Env Vars (Summary):")
+        for env_var in all_env:
+            if "SOURCE" in env_var.key:
+                print(f"[DEBUG Decorator]  {env_var.key}: <source code present>")
+            else:
+                print(f"[DEBUG Decorator]  {env_var.key}: {env_var.value}")
 
-        # Create the processor instance
         processor_instance = Processor(
             name=processor_name,
             stream=processor_name,
             namespace=namespace,
             labels=labels,
             container=container_request,
-            schema_=None,  # TODO
+            schema_=None,
             common_schema=None,
             min_replicas=min_replicas,
             max_replicas=max_replicas,
             scale_config=scale,
             no_delete=no_delete,
         )
-
+        print(
+            f"[DEBUG Decorator] Processor instance '{processor_name}' created successfully."
+        )
         return processor_instance
 
     return decorator

{nebu-0.1.14.dist-info → nebu-0.1.17.dist-info}/METADATA

@@ -1,13 +1,16 @@
 Metadata-Version: 2.4
 Name: nebu
-Version: 0.1.14
+Version: 0.1.17
 Summary: A globally distributed container runtime
 Requires-Python: >=3.10.14
 Description-Content-Type: text/markdown
 License-File: LICENSE
+Requires-Dist: dill>=0.3.8
 Requires-Dist: openai>=1.68.2
 Requires-Dist: pydantic>=2.10.6
+Requires-Dist: pysocks>=1.7.1
 Requires-Dist: pyyaml>=6.0.2
+Requires-Dist: redis[socks]>=5.0
 Requires-Dist: requests>=2.32.3
 Dynamic: license-file
 

{nebu-0.1.14.dist-info → nebu-0.1.17.dist-info}/RECORD

@@ -3,18 +3,18 @@ nebu/auth.py,sha256=rApCd-7_c3GpIb7gjCB79rR7SOcmkG7MmaTE6zMbvr0,1125
 nebu/config.py,sha256=XBY7uKgcJX9d1HGxqqpx87o_9DuF3maUlUnKkcpUrKU,4565
 nebu/meta.py,sha256=CzFHMND9seuewzq9zNNx9WTr6JvrCBExe7BLqDSr7lM,745
 nebu/containers/container.py,sha256=yb7KaPTVXnEEAlrpdlUi4HNqF6P7z9bmwAILGlq6iqU,13502
-nebu/containers/decorator.py,sha256=qiM7hbHne9MhSp1gDgX5z5bimsXr_YPjTIZoe09dwr4,2741
+nebu/containers/decorator.py,sha256=uFtzlAXRHYZECJ-NPusY7oN9GXvdHrHDd_JNrIGr8aQ,3244
 nebu/containers/models.py,sha256=0j6NGy4yto-enRDh_4JH_ZTbHrLdSpuMOqNQPnIrwC4,6815
 nebu/containers/server.py,sha256=yFa2Y9PzBn59E1HftKiv0iapPonli2rbGAiU6r-wwe0,2513
-nebu/processors/consumer.py,sha256=OrhiG6qW4CcKXaNcTHV4AqXZ5M01JRFXbZhmNdU5Nsg,15232
-nebu/processors/decorate.py,sha256=coIHJ6p1GhGqgoDFhNSrWRLAdpDooenr-L6JwGghRck,14953
+nebu/processors/consumer.py,sha256=0HJxRLoeRdN4xY6bjIxqr5bD5JpFSyKb5s-eS5oTy9s,16063
+nebu/processors/decorate.py,sha256=BlpRF9u-1q9IS-UDn0XFzoYIBqEU4Xkjkoge394exJs,34786
 nebu/processors/default.py,sha256=W4slJenG59rvyTlJ7gRp58eFfXcNOTT2Hfi6zzJAobI,365
 nebu/processors/models.py,sha256=GvnI8UJrQSjHo2snP07cPfisCH90cEGTY-PZV5_AtXI,3654
 nebu/processors/processor.py,sha256=oy2YdI-cy6qQWxrZhpZahJV46oWZlu_Im-jm811R_oo,9667
 nebu/redis/models.py,sha256=coPovAcVXnOU1Xh_fpJL4PO3QctgK9nBe5QYoqEcnxg,1230
 nebu/services/service.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-nebu-0.1.14.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-nebu-0.1.14.dist-info/METADATA,sha256=Lo_6Ww5-zJ8MCZ7yPqCYrK1OtKOF9KpmR-ytN2FG24o,1588
-nebu-0.1.14.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-nebu-0.1.14.dist-info/top_level.txt,sha256=uLIbEKJeGSHWOAJN5S0i5XBGwybALlF9bYoB1UhdEgQ,5
-nebu-0.1.14.dist-info/RECORD,,
+nebu-0.1.17.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+nebu-0.1.17.dist-info/METADATA,sha256=m9oiMDrMt5x1x7wa51dsxwgzj-tUcX_hZrQrGN8NtyE,1678
+nebu-0.1.17.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
+nebu-0.1.17.dist-info/top_level.txt,sha256=uLIbEKJeGSHWOAJN5S0i5XBGwybALlF9bYoB1UhdEgQ,5
+nebu-0.1.17.dist-info/RECORD,,