chalkpy 2.90.1__py3-none-any.whl → 2.95.3__py3-none-any.whl

chalk/features/inference.py

@@ -12,12 +12,11 @@ from chalk.utils.collections import ensure_tuple
 
 
 def build_inference_function(
-    model_version: ModelVersion, pkey: Feature, output_feature: Optional[Feature] = None, text_model: bool = False
+    model_version: ModelVersion, pkey: Feature, output_features: Optional[Feature | list[Feature]] = None
 ) -> Callable[[DataFrame], DataFrame]:
     """Build the core inference function that takes a DataFrame and returns predictions.
 
-    This is the function that gets called by the inference resolver. It excludes the primary key
-    from the input, makes predictions, and optionally adds them back to the DataFrame.
+    Uses ModelInference.prepare_input() and extract_output() for model-specific logic.
 
     Parameters
     ----------
@@ -25,40 +24,49 @@ def build_inference_function(
         The model version to use for prediction
     pkey
         The primary key feature to exclude from predictions
-    output_feature
-        Optional output feature to add predictions to the DataFrame. If None, just returns raw predictions.
+    output_features
+        Optional output feature(s) to add predictions to the DataFrame.
+        Can be a single Feature or a list of Features for multi-output models.
 
     Returns
     -------
     Callable[[DataFrame], DataFrame]
         Function that takes a DataFrame and returns predictions
     """
+    # For all other models, use the ModelInference prepare_input/extract_output methods
     pkey_string = str(pkey)
-    if text_model:
-        # For text models, we expect a single column of string data
-        def fn(inp: DataFrame):
-            texts = inp[[c for c in inp.columns if c != pkey_string]].to_pyarrow().column(0).to_numpy()
-            result = model_version.predict(texts)
-            if output_feature is not None:
-                return inp[pkey_string].with_columns({output_feature: result})
-            return result
-
-        return fn
 
     def fn(inp: DataFrame):
-        # Convert features to PyArrow table, excluding primary key
-        arr = inp[[c for c in inp.columns if c != pkey_string]].to_pyarrow().__array__()
+        # Get features (excluding primary key) as PyArrow table
+        feature_table = inp[[c for c in inp.columns if c != pkey_string]].to_pyarrow()
+
+        # Use model-specific input preparation (default: __array__(), ONNX: struct array)
+        model_input = model_version.predictor.prepare_input(feature_table)
+
+        # Run prediction
+        result = model_version.predict(model_input)
+
+        if output_features is not None:
+            # Normalize to list for uniform processing
+            features_list = output_features if isinstance(output_features, list) else [output_features]
+
+            # Extract output for each feature and build columns dict
+            columns_dict = {}
+            for output_feature in features_list:
+                # Use model-specific output extraction (default: identity, ONNX: extract field)
+                output_feature_name = str(output_feature).split(".")[-1]
+                result_data = model_version.predictor.extract_output(result, output_feature_name)
+                columns_dict[output_feature] = result_data
+
+            return inp[pkey_string].with_columns(columns_dict)
 
-        result = model_version.predict(arr)
-        if output_feature is not None:
-            return inp[pkey_string].with_columns({output_feature: result})
         return result
 
     return fn
 
 
 def generate_inference_resolver(
-    inputs: list[Underscore], model_version: ModelVersion, resource_hint: Optional[ResourceHint] = None
+    inputs: list[Underscore] | Underscore, model_version: ModelVersion, resource_hint: Optional[ResourceHint] = None
 ) -> Feature:
     output_feature = Feature()
     previous_hook = output_feature.hook
@@ -73,21 +81,15 @@ def generate_inference_resolver(
 
         def resolver_factory():
             # Use the extracted build_inference_function
-
             cleaned_inputs = []
-            for i in inputs:
+            inputs_list = inputs if isinstance(inputs, list) else [inputs]
+            for i in inputs_list:
                 try:
                     cleaned_inputs.append(Feature.from_root_fqn(output_feature.namespace + str(i)[1:]))
                 except Exception as e:
                     raise ValueError(f"Could not find feature for input {i}: {e}")
 
-            try:
-                text_model = len(cleaned_inputs) == 1 and cleaned_inputs[0].typ.parsed_annotation is str
-            except Exception:
-                # Fallback in case of any issues determining the type
-                text_model = False
-
-            fn = build_inference_function(model_version, pkey, output_feature, text_model=text_model)
+            fn = build_inference_function(model_version, pkey, output_feature)
 
             identifier = model_version.identifier or ""
             model_reference = MODEL_REFERENCE_REGISTRY.get((model_version.name, identifier), None)
@@ -111,7 +113,7 @@ def generate_inference_resolver(
                 when=None,
                 tags=None,
                 owner=None,
-                resource_hint=resource_hint,
+                resource_hint=resource_hint or model_version.resource_hint,
                 data_sources=None,
                 is_sql_file_resolver=False,
                 source_line=None,

chalk/features/resolver.py

@@ -70,6 +70,7 @@ from pydantic import BaseModel
 
 from chalk._lsp._class_finder import get_function_caller_info
 from chalk._lsp.error_builder import FunctionCallErrorBuilder, ResolverErrorBuilder, get_resolver_error_builder
+from chalk.df.LazyFramePlaceholder import LazyFramePlaceholder
 from chalk.features._encoding.protobuf import (
     convert_proto_message_type_to_pyarrow_type,
     serialize_message_file_descriptor,
@@ -114,6 +115,7 @@ if TYPE_CHECKING:
 
     from chalk.features import Underscore
     from chalk.features.underscore import UnderscoreAttr, UnderscoreCall, UnderscoreCast, UnderscoreFunction
+    from chalk.ml.model_version import ModelVersion
     from chalk.sql import BaseSQLSourceProtocol, SQLSourceGroup
     from chalk.sql._internal.sql_settings import SQLResolverSettings
     from chalk.sql._internal.sql_source import BaseSQLSource
@@ -555,28 +557,18 @@ class ResolverRegistry:
         short_name = resolver.name
         if short_name in self._short_name_to_resolver:
             if not override and not notebook.is_notebook():
-                if resolver.fqn == self._short_name_to_resolver[short_name]:
-                    # Same resolver was redefined
-                    resolver.lsp_builder.add_diagnostic(
-                        message=f"Duplicate resolver '{resolver.fqn}'. Multiple resolvers cannot have the same name.",
-                        code="71",
-                        label="duplicate name",
-                        range=resolver.lsp_builder.function_name(),
-                        raise_error=ValueError,
-                    )
-                else:
-                    # Same short name was reused
-                    resolver.lsp_builder.add_diagnostic(
-                        message=(
-                            f"Another resolver with the same function name '{resolver.name}' in module "
-                            f"'{self._short_name_to_resolver[short_name].__module__}' exists. "
-                            f"Resolver function names must be unique. Please rename this resolver in module '{resolver.__module__}'."
-                        ),
-                        label="duplicate resolver shortname",
-                        code="71",
-                        range=resolver.lsp_builder.function_name(),
-                        raise_error=None,
-                    )
+                # Same short name was reused
+                resolver.lsp_builder.add_diagnostic(
+                    message=(
+                        f"Another resolver with the same function name '{resolver.name}' in module "
+                        f"'{self._short_name_to_resolver[short_name].__module__}' exists. "
+                        f"Resolver function names must be unique. Please rename this resolver in module '{resolver.__module__}'."
+                    ),
+                    label="duplicate resolver shortname",
+                    code="71",
+                    range=resolver.lsp_builder.function_name(),
+                    raise_error=None,
+                )
                 return
             existing_resolver = self._short_name_to_resolver[short_name]
             # Need to remove the resolver from the typed registry
@@ -649,6 +641,7 @@ class Resolver(ResolverProtocol[P, T], abc.ABC):
         output_row_order: Literal["one-to-one"] | None = None,
         venv: str | None = None,
         name: None = None,  # deprecated
+        postprocessing: Underscore | None = None,
     ):
         self._function_definition = ... if function_definition is None else function_definition
         self._function_captured_globals = ... if function_captured_globals is None else function_captured_globals
@@ -692,6 +685,7 @@ class Resolver(ResolverProtocol[P, T], abc.ABC):
         self._data_lineage = data_lineage
         self._sql_settings = sql_settings
         self.output_row_order = output_row_order
+        self.postprocessing = postprocessing
         super().__init__()
 
     @property
@@ -849,14 +843,20 @@ class Resolver(ResolverProtocol[P, T], abc.ABC):
                 annotation = None
 
             if annotation is not None:
-                if not isinstance(val, DataFrame):
+                if self.static and type(val).__name__ == "DataFrame" and type(val).__module__ == "chalkdf.dataframe":
+                    # No need to wrap this class in DataFrame.
+                    pass
+                elif self.static and isinstance(val, LazyFramePlaceholder):
+                    # No need to wrap this class in DataFrame.
+                    pass
+                elif not isinstance(val, DataFrame):
                     val = DataFrame(val)
 
                 if time_is_frozen():
                     frozen_filter = Filter(lhs=CHALK_TS_FEATURE, operation="<=", rhs=TimeDelta(hours_ago=0))
                     annotation.filters = (frozen_filter, *annotation.filters)
 
-                if annotation.filters and len(annotation.filters) > 0:
+                if annotation.filters and len(annotation.filters) > 0 and not isinstance(val, LazyFramePlaceholder):
                     try:
                         val = val[annotation.filters]
                         val._materialize()  # pyright: ignore[reportPrivateUsage]
@@ -2829,6 +2829,8 @@ class StreamResolver(Resolver[P, T]):
         sql_settings: SQLResolverSettings | None,
         feature_expressions: dict[Feature, Underscore] | None,
         message_producer_parsed: StreamResolverMessageProducerParsed | None,
+        skip_online: bool = False,
+        skip_offline: bool = False,
     ):
         super().__init__(
             function_definition=function_definition,
@@ -2897,6 +2899,8 @@ class StreamResolver(Resolver[P, T]):
 
         self.feature_expressions: dict[Feature, Underscore] | None = feature_expressions
         self.message_producer_parsed: StreamResolverMessageProducerParsed | None = message_producer_parsed
+        self.skip_online = skip_online
+        self.skip_offline = skip_offline
 
     @property
     def output_features(self) -> Sequence[Feature]:
@@ -3808,6 +3812,8 @@ def make_stream_resolver(
     owner: Optional[str] = None,
     doc: str | None = None,
     sink: Sink | None = None,
+    skip_online: bool = False,
+    skip_offline: bool = False,
 ) -> StreamResolver:
     """Constructs a streaming resolver that, instead of a Python function,
     defines its output features as column projections on an input message.
@@ -3840,6 +3846,14 @@ def make_stream_resolver(
     sink
         An optional message producer configuration that specifies where to send messages.
         Read more at https://docs.chalk.ai/api-docs#Sink
+    skip_online
+        If True, skip online persistence (no writes to Redis/DynamoDB/etc).
+        Results will still be processed but not stored in online stores.
+        Note: Only applies to native streaming. Default: False
+    skip_offline
+        If True, skip offline persistence (no result bus publishing for offline storage).
+        Results will still be processed but not stored in offline stores (S3/BigQuery/etc).
+        Note: Only applies to native streaming. Default: False
 
     Returns
     -------
@@ -4017,6 +4031,8 @@ def make_stream_resolver(
         sql_settings=None,
         feature_expressions={unwrap_feature(x): u for x, u in output_features.items()},
         message_producer_parsed=message_producer_parsed,
+        skip_online=skip_online,
+        skip_offline=skip_offline,
     )
     resolver.add_to_registry(override=False)
     return resolver
@@ -4307,6 +4323,9 @@ def validate_field_chain(
 
     # Base case: if parent is UnderscoreRoot (_), validate field against current_type
     if isinstance(underscore_attr._chalk__parent, UnderscoreRoot):
+        if underscore_attr._chalk__attr == "chalk_now":
+            return datetime
+
         # Check if current_type allows field access
         if current_type in (str, bytes):
             error_builder.add_diagnostic(
@@ -4476,3 +4495,184 @@ def validate_message_attributes(
         message_type = message_type.__args__[0]  # pyright: ignore[reportAttributeAccessIssue]
     for expression in expressions:
         validate_underscore_expression(expression, message_type, error_builder, name)
+
+
+def make_model_resolver(
+    name: str,
+    model: "ModelVersion",
+    inputs: Dict[Feature, str] | List[Feature],
+    output: Feature | List[Feature] | Dict[Feature, str],
+    feature_class: Optional[type[Features]] = None,
+    resource_group: Optional[str] = None,
+    resource_hint: Optional[ResourceHint] = None,
+) -> OnlineResolver:
+    """
+    Create an online resolver that runs inference on a model.
+
+    This function provides an imperative API for creating model inference resolvers,
+    as an alternative to using F.inference in feature definitions. It uses the same
+    underlying implementation as F.inference but allows you to create resolvers
+    programmatically.
+
+    Parameters
+    ----------
+    name
+        The name of the resolver
+    model
+        A ModelVersion reference to a deployed model
+    inputs
+        Either a dict mapping Feature objects to model input names (strings), or a list of
+        Feature objects. If a dict, the values represent the model's expected input names
+        (for future use). If a list, the features will be passed as a single DataFrame to
+        the model.
+    output
+        The output feature(s) that will contain the predictions.
+        Can be a single Feature, a list of Features, or a dict mapping Feature objects to
+        model output names (strings) for future use with multi-output models.
+    feature_class
+        Optional feature class to use. If not provided, will be inferred from the inputs.
+    resource_group
+        Optional resource group for the resolver
+    resource_hint
+        Optional resource hint for execution (e.g., CPU/GPU preferences)
+
+    Returns
+    -------
+    OnlineResolver
+        The created resolver
+
+    Examples
+    --------
+    >>> from chalk.features import features, feature
+    >>> from chalk.features.resolver import make_model_resolver
+    >>> from chalk.ml import ModelVersion
+    >>>
+    >>> @features
+    ... class User:
+    ...     id: str = feature(primary=True)
+    ...     age: float
+    ...     income: float
+    ...     risk_score: float
+    ...     credit_score: float
+    >>>
+    >>> # Create a model version reference
+    >>> model = ModelVersion(
+    ...     name="risk_model",
+    ...     version=1,
+    ...     model_type="sklearn",
+    ...     model_encoding="pickle",
+    ...     filename="model.pkl"
+    ... )
+    >>>
+    >>> # Create resolver with single output
+    >>> resolver = make_model_resolver(
+    ...     name="risk_model",
+    ...     model=model,
+    ...     inputs=[User.age, User.income],
+    ...     output=User.risk_score,
+    ... )
+    >>>
+    >>> # Create resolver with multiple outputs (list)
+    >>> resolver = make_model_resolver(
+    ...     name="multi_output_model",
+    ...     model=model,
+    ...     inputs=[User.age, User.income],
+    ...     output=[User.risk_score, User.credit_score],
+    ... )
+    >>>
+    >>> # Create resolver with named inputs and outputs (dict)
+    >>> resolver = make_model_resolver(
+    ...     name="named_model",
+    ...     model=model,
+    ...     inputs={User.age: "age_input", User.income: "income_input"},
+    ...     output={User.risk_score: "risk_output", User.credit_score: "credit_output"},
+    ... )
+    """
+    from chalk.features.inference import build_inference_function
+
+    if isinstance(inputs, dict):
+        input_features_raw = list(inputs.keys())
+    else:
+        input_features_raw = inputs
+
+    input_features = [unwrap_feature(f) for f in input_features_raw]
+
+    if isinstance(output, dict):
+        output_features = [unwrap_feature(f) for f in output.keys()]
+    elif isinstance(output, list):
+        output_features = [unwrap_feature(f) for f in output]
+    else:
+        output_features = [unwrap_feature(output)]
+
+    # If feature_class is not provided, try to infer it from the first input feature
+    if feature_class is None:
+        if not input_features:
+            raise ValueError("Cannot infer feature class: no input features provided and feature_class not specified")
+
+        first_input = input_features[0]
+
+        if hasattr(first_input, "features_cls") and first_input.features_cls is not None:
+            feature_class = first_input.features_cls
+        else:
+            raise ValueError(
+                "Cannot infer feature class from inputs. Please provide feature_class parameter explicitly."
+            )
+
+    pkey = feature_class.__chalk_primary__
+    if pkey is None:
+        raise ValueError(f"Feature class {feature_class} does not have a primary key defined")
+
+    first_output = output_features[0]
+
+    output_namespace = (
+        first_output.namespace
+        if hasattr(first_output, "namespace") and first_output.namespace
+        else feature_class.__name__.lower()
+    )
+
+    # Use the same underlying inference function as F.inference
+    # Pass list of outputs if multiple, single if only one
+    output_for_inference = output_features if len(output_features) > 1 else output_features[0]
+    inference_fn = build_inference_function(model, pkey, output_for_inference)
+
+    if len(output_features) == 1:
+        output_names = output_features[0].name
+    else:
+        output_names = "_".join(f.name for f in output_features)
+
+    resolver = OnlineResolver(
+        function_definition="",
+        filename="",
+        fqn=f"{name}__{output_namespace}_{output_names}",
+        doc=None,
+        inputs=[DataFrame[[pkey, *ensure_tuple(input_features)]]],
+        state=None,
+        output=Features[DataFrame[tuple([*output_features, pkey])]],  # type: ignore[misc]
+        fn=inference_fn,
+        environment=None,
+        machine_type=None,
+        default_args=[None],
+        timeout=None,
+        cron=None,
+        when=None,
+        tags=None,
+        owner=None,
+        resource_hint=resource_hint or model.resource_hint,
+        data_sources=None,
+        is_sql_file_resolver=False,
+        source_line=None,
+        lsp_builder=get_resolver_error_builder(inference_fn),
+        parse=None,
+        static=False,
+        total=False,
+        autogenerated=False,
+        unique_on=None,
+        partitioned_by=None,
+        data_lineage=None,
+        sql_settings=None,
+    )
+
+    # Register the resolver
+    RESOLVER_REGISTRY.add_to_registry(resolver, override=False)
+
+    return resolver

chalk/functions/__init__.py

@@ -1635,6 +1635,69 @@ def sagemaker_predict(
     )
 
 
+def openai_complete(
+    api_key: Underscore | str,
+    prompt: Underscore | str,
+    model: Underscore | str,
+    max_tokens: Underscore | int,
+    temperature: Underscore | float,
+):
+    """
+    Makes a completion request to OpenAI's chat API and returns the response.
+
+    This is a blocking expression that calls OpenAI's API during feature computation.
+    The response includes the completion text along with token usage statistics.
+
+    Parameters
+    ----------
+    api_key
+        The OpenAI API key to use for authentication.
+    prompt
+        The prompt text to send to the model.
+    model
+        The OpenAI model to use (e.g., "gpt-4", "gpt-3.5-turbo").
+    max_tokens
+        The maximum number of tokens to generate in the completion.
+    temperature
+        The sampling temperature to use, between 0 and 2. Higher values make
+        output more random, lower values make it more deterministic.
+
+    Returns
+    -------
+    A struct containing:
+        - completion: The generated text response
+        - prompt_tokens: Number of tokens in the prompt
+        - completion_tokens: Number of tokens in the completion
+        - total_tokens: Total tokens used (prompt + completion)
+        - model: The model used for the completion
+        - finish_reason: Why the completion stopped (e.g., "stop", "length")
+
+    Examples
+    --------
+    >>> import chalk.functions as F
+    >>> from chalk.features import _, features
+    >>> @features
+    ... class Document:
+    ...    id: str
+    ...    content: str
+    ...    summary: str = F.openai_complete(
+    ...        api_key="sk-...",
+    ...        prompt=_.content,
+    ...        model="gpt-4",
+    ...        max_tokens=100,
+    ...        temperature=0.7,
+    ...    ).completion
+    """
+    return UnderscoreFunction(
+        "openai_complete",
+        api_key,
+        prompt,
+        model,
+        max_tokens,
+        temperature,
+    )
+
+
 def json_value(expr: Underscore, path: Union[str, Underscore]):
     """
     Extract structured data from a JSON string feature using a JSONPath expression.
@@ -3884,7 +3947,6 @@ def array_filter(
     ...    id: str
     ...    recent_activities: list[float]
     ...    average_activity: float
-    ...
     ...    recent_high_value_activities: list[float] = F.array_filter(
     ...        _.recent_activities,
     ...        lambda amount: amount > _.average_activity,
@@ -3990,7 +4052,6 @@ def array_reduce(
     ... class User:
     ...    id: str
     ...    scores: list[int]
-    ...
     ...    total_score: int = F.array_reduce(
     ...        arr=_.scores,
     ...        initial_value=0,
@@ -4671,7 +4732,7 @@ def nth_bucket_end(value: Underscore, bucket_duration: str, n: int, initial_buck
 
 
 def inference(
-    model: ModelVersion, inputs: list[Underscore | Any], resource_hint: ResourceHint | None = None
+    model: ModelVersion, inputs: list[Underscore | Any] | Underscore, resource_hint: ResourceHint | None = None
 ) -> Underscore | Feature:
     """
     Run inference on a deployed ML model.
@@ -5917,6 +5978,7 @@ __all__ = (
     "normal_cdf",
     "nth_bucket_end",
     "nth_bucket_start",
+    "openai_complete",
     "parse_datetime",
     "partial_ratio",
     "pi",

chalk/gitignore/gitignore_parser.py

@@ -114,6 +114,10 @@ def _rule_from_pattern(pattern: str, base_path: Optional[Path] = None, source: O
     regex = _fnmatch_pathname_to_regex(pattern, directory_only)
     if anchored:
         regex = f"^{regex}"
+    else:
+        # For non-anchored patterns, match at path component boundaries
+        # (start of string or after a path separator)
+        regex = f"(^|/){regex}"
     regex = f"(?ms){regex}"
     return IgnoreRule(
         pattern=orig_pattern,
@@ -215,6 +219,6 @@ def _fnmatch_pathname_to_regex(pattern: str, directory_only: bool):
     if directory_only:
         res.append(r"/.*$")
     else:
-        res.append("(/.*|[^/]*)$")
+        res.append("(/.*)?$")
 
     return "".join(res)