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

chalk/client/client.py

@@ -32,6 +32,7 @@ from chalk.client.models import (
     GetIncrementalProgressResponse,
     GetRegisteredModelResponse,
     GetRegisteredModelVersionResponse,
+    ManualTriggerScheduledQueryResponse,
     OfflineQueryInputUri,
     OnlineQuery,
     OnlineQueryContext,
@@ -40,6 +41,7 @@ from chalk.client.models import (
     RegisterModelVersionResponse,
     ResolverRunResponse,
     ResourceRequests,
+    ScheduledQueryRun,
     StreamResolverTestResponse,
     WhoAmIResponse,
 )
@@ -924,6 +926,8 @@ class ChalkClient:
         max_retries: int | None = None,
         query_name: str | None = None,
         query_name_version: str | None = None,
+        *,
+        input_sql: str | None = None,
     ) -> Dataset:
         """Compute feature values from the offline store or by running offline/online resolvers.
         See `Dataset` for more information.
@@ -944,6 +948,9 @@ class ChalkClient:
             times, the list must match the length of the `input` lists. Each element of input_time corresponds with the
             feature values at the same index of the `input` lists.
             See https://docs.chalk.ai/docs/temporal-consistency for more information.
+        input_sql
+            An alternative to `input`: a ChalkSQL query that returns values
+            to use as inputs.
         output
             The features that you'd like to sample, if they exist.
             If an output feature was never computed for a sample (row) in
@@ -1080,6 +1087,82 @@ class ChalkClient:
         """
         ...
 
+    def run_scheduled_query(
+        self,
+        name: str,
+        planner_options: Optional[Mapping[str, Any]],
+        incremental_resolvers: Optional[Sequence[str]],
+        max_samples: Optional[int],
+        env_overrides: Optional[Mapping[str, str]],
+    ) -> ManualTriggerScheduledQueryResponse:
+        """
+        Manually trigger a scheduled query request.
+
+        Parameters
+        ----------
+        name
+            The name of the scheduled query to be triggered.
+        incremental_resolvers
+            If set to None, Chalk will incrementalize resolvers in the query's root namespaces.
+            If set to a list of resolvers, this set will be used for incrementalization.
+            Incremental resolvers must return a feature time in its output, and must return a `DataFrame`.
+            Most commonly, this will be the name of a SQL file resolver. Chalk will ingest all new data
+            from these resolvers and propagate changes to values in the root namespace.
+        max_samples
+            The maximum number of samples to compute.
+        env_overrides:
+            A dictionary of environment values to override during this specific triggered query.
+
+        Other Parameters
+        ----------------
+        planner_options
+            A dictionary of options to pass to the planner.
+            These are typically provided by Chalk Support for specific use cases.
+
+        Returns
+        -------
+        ManualTriggerScheduledQueryResponse
+            A response message containing metadata around the triggered run.
+
+        Examples
+        --------
+        >>> from chalk.client.client_grpc import ChalkGRPCClient
+        >>> ChalkGRPCClient().run_scheduled_query(
+        ...     name="my_scheduled_query",
+        ... )
+        """
+        ...
+
+    def get_scheduled_query_run_history(
+        self,
+        name: str,
+        limit: int = 10,
+    ) -> List[ScheduledQueryRun]:
+        """
+        Get the run history for a scheduled query.
+
+        Parameters
+        ----------
+        name
+            The name of the scheduled query.
+        limit
+            The maximum number of runs to return. Defaults to 10.
+
+        Returns
+        -------
+        list[ScheduledQueryRun]
+            A response message containing the list of scheduled query runs.
+
+        Examples
+        --------
+        >>> from chalk.client import ChalkClient
+        >>> ChalkClient().get_scheduled_query_run_history(
+        ...     name="my_scheduled_query",
+        ...     limit=20,
+        ... )
+        """
+        ...
+
     def prompt_evaluation(
         self,
         prompts: list[Prompt | str],
@@ -2074,21 +2157,20 @@ class ChalkClient:
         name: str,
         version: Optional[int] = None,
     ) -> Union[GetRegisteredModelResponse, GetRegisteredModelVersionResponse]:
-        """
-        Retrieve a registered model from the Chalk model registry.
+        """Retrieve a registered model from the Chalk model registry.
 
         Parameters
         ----------
-        name : str
-           Name of the model to retrieve
-        version : int, optional
-           Specific version number to retrieve. If not provided, returns
-           information about all versions of the model
+        name
+            Name of the model to retrieve.
+        version
+            Specific version number to retrieve. If not provided, returns
+            information about all versions of the model.
 
         Returns
         -------
-        GetRegisteredModelResponse or GetRegisteredModelVersionResponse
-           Model information including metadata, versions, and configuration details
+        Union[GetRegisteredModelResponse, GetRegisteredModelVersionResponse]
+            Model information including metadata, versions, and configuration details.
 
         Examples
         --------
@@ -2165,49 +2247,49 @@ class ChalkClient:
         source_config: Optional[SourceConfig] = None,
         dependencies: Optional[List[str]] = None,
     ) -> RegisterModelVersionResponse:
-        """
-        Register a model in the Chalk model registry.
+        """Register a model in the Chalk model registry.
 
         Parameters
         ----------
-        name : str
-           Unique name for the model
-        aliases : list of str, optional
-           List of version aliases (e.g., ["v1.0", "latest"])
-        model : object, optional
-           Python model object (for object-based registration)
-        model_paths : list of str, optional
-           Paths to model files (for file-based registration)
-        additional_files : List[str], optional
-           Additional files needed for inference (tokenizers, configs, etc.)
-        model_type : ModelType, optional
-           Type of model framework
-        model_encoding : ModelEncoding, optional
-           Serialization format
-        input_schema : dict, list, or Any
-           Definition of the input schema. Can be:
-           - dict: Dictionary mapping column names to dtypes for tabular data
-           - list: List of (shape, dtype) tuples for tensor data
-        output_schema : dict, list, or Any
-           Definition of the output schema. Can be:
-           - dict: Dictionary mapping column names to dtypes for tabular data
-           - list: List of (shape, dtype) tuples for tensor data
-        metadata : dict, optional
-           Additional metadata dictionary containing framework info,
-           training details, performance metrics, etc.
-        input_features : FeatureReference, str, optional
+        name
+            Unique name for the model.
+        aliases
+            List of version aliases (e.g., `["v1.0", "latest"]`).
+        model
+            Python model object (for object-based registration).
+        model_paths
+            Paths to model files (for file-based registration).
+        additional_files
+            Additional files needed for inference (tokenizers, configs, etc.)
+        model_type
+            Type of model framework
+        model_encoding
+            Serialization format
+        input_schema
+            Definition of the input schema. Can be:
+            - `dict`: Dictionary mapping column names to dtypes for tabular data
+            - `list`: List of `(shape, dtype)` tuples for tensor data
+        output_schema
+            Definition of the output schema. Can be:
+            - `dict`: Dictionary mapping column names to dtypes for tabular data
+            - `list`: List of `(shape, dtype)` tuples for tensor data
+        metadata
+            Additional metadata dictionary containing framework info,
+            training details, performance metrics, etc.
+        input_features
             The features to be used as inputs to the model.
             For example, `[User.message]`. Features can also be expressed as snakecased strings,
-            e.g. `["user.message"]`
-        output_features : FeatureReference, str, optional
+            e.g. `["user.message"]`.
+        output_features
             The features to be used as outputs to the model.
             For example, `[User.is_spam]`. Features can also be expressed as snakecased strings,
-            e.g. `["user.is_spam"]`
-        source_config : LocalSourceConfig, S3SourceConfig, HFSourceConfig, optional
+            e.g. `["user.is_spam"]`.
+        source_config
             Config to pass credentials to access files from a remote source.
-        dependencies : List[str], optional
+        dependencies
             List of package dependencies needed to run this model.
-            e.g. ["torch==2.7.1", "numpy==1.26.4"]
+            e.g. `["torch==2.7.1", "numpy==1.26.4"]`.
+
         Returns
         -------
         ModelVersion
@@ -2250,6 +2332,7 @@ class ChalkClient:
         name: str,
         model_artifact_id: Optional[str] = None,
         run_id: Optional[str] = None,
+        run_name: Optional[str] = None,
         criterion: Optional[ModelRunCriterion] = None,
         aliases: Optional[List[str]] = None,
     ) -> RegisterModelVersionResponse:
@@ -2264,6 +2347,8 @@ class ChalkClient:
             Artifact UUID to promote to a model version.
         run_id: str, optional
             run id that produce the artifact to promote.
+        run_name: str, optional
+            run name used in the checkpointer for artifact to promote.
         criterion: ModelRunCriterion, optional
             criterion on which to select the artifact from the training run.
             If none provided, the latest artifact in the run will be selected.

chalk/client/client_async.py

@@ -9,11 +9,13 @@ from uuid import UUID
 from chalk import DataFrame
 from chalk.client.models import (
     BulkOnlineQueryResponse,
+    BulkOnlineQueryResult,
     FeatureReference,
     FeatureStatisticsResponse,
     OfflineQueryInputUri,
     OnlineQuery,
     OnlineQueryContext,
+    PlanQueryResponse,
     ResourceRequests,
     UploadFeaturesResponse,
     WhoAmIResponse,
@@ -26,6 +28,7 @@ from chalk.prompts import Prompt
 if TYPE_CHECKING:
     import pandas as pd
     import polars as pl
+    import pyarrow as pa
 
     QueryInput = Mapping[FeatureReference, Any] | pd.DataFrame | pl.DataFrame | DataFrame | str
 
@@ -467,6 +470,147 @@ class AsyncChalkClient:
         """
         ...
 
+    async def plan_query(
+        self,
+        input: Sequence[FeatureReference],
+        output: Sequence[FeatureReference],
+        staleness: Mapping[FeatureReference, str] | None = None,
+        environment: EnvironmentId | None = None,
+        tags: list[str] | None = None,
+        preview_deployment_id: str | None = None,
+        branch: Union[BranchId, None] = ...,
+        query_name: str | None = None,
+        query_name_version: str | None = None,
+        meta: Mapping[str, str] | None = None,
+        store_plan_stages: bool = False,
+        explain: bool = False,
+        num_input_rows: Optional[int] = None,
+        headers: Mapping[str, str] | None = None,
+    ) -> PlanQueryResponse:
+        """Plan a query without executing it.
+
+        Parameters
+        ----------
+        input
+            The features for which there are known values, mapped to those values.
+            For example, `{User.id: 1234}`. Features can also be expressed as snakecased strings,
+            e.g. `{"user.id": 1234}`
+        output
+            Outputs are the features that you'd like to compute from the inputs.
+            For example, `[User.age, User.name, User.email]`.
+        staleness
+            Maximum staleness overrides for any output features or intermediate features.
+            See https://docs.chalk.ai/docs/query-caching for more information.
+        environment
+            The environment under which to run the resolvers.
+            API tokens can be scoped to an environment.
+            If no environment is specified in the query,
+            but the token supports only a single environment,
+            then that environment will be taken as the scope
+            for executing the request.
+        tags
+            The tags used to scope the resolvers.
+            See https://docs.chalk.ai/docs/resolver-tags for more information.
+        branch
+            If specified, Chalk will route your request to the relevant branch.
+        preview_deployment_id
+            If specified, Chalk will route your request to the relevant preview deployment.
+        query_name
+            The semantic name for the query you're making, for example, `"loan_application_model"`.
+            Typically, each query that you make from your application should have a name.
+            Chalk will present metrics and dashboard functionality grouped by 'query_name'.
+            If your query name matches a `NamedQuery`, the query will automatically pull outputs
+            and options specified in the matching `NamedQuery`.
+        query_name_version
+            If `query_name` is specified, this specifies the version of the named query you're making.
+            This is only useful if you want your query to use a `NamedQuery` with a specific name and a
+            specific version. If a `query_name` has not been supplied, then this parameter is ignored.
+        meta
+            Arbitrary `key:value` pairs to associate with a query.
+        store_plan_stages
+            If true, the plan will store the intermediate values at each stage in the plan
+        explain
+            If true, the plan will emit additional output to assist with debugging.
+        num_input_rows:
+            The number of input rows that this plan will be run with. If unknown, specify `None`.
+        headers
+            Additional headers to provide with the request
+
+        Returns
+        -------
+        PlanQueryResponse
+            The query plan, including the resolver execution order and the
+            resolver execution plan for each resolver.
+
+        Examples
+        --------
+        >>> from chalk.client import AsyncChalkClient
+        >>> result = await AsyncChalkClient().plan_query(
+        ...     input=[User.id],
+        ...     output=[User.fico_score],
+        ...     staleness={User.fico_score: "10m"},
+        ... )
+        >>> result.rendered_plan
+        >>> result.output_schema
+        """
+        ...
+
+    async def _run_serialized_query(
+        self,
+        serialized_plan_bytes: bytes,
+        input: Union[Mapping[FeatureReference, Sequence[Any]], pa.Table],
+        output: Sequence[FeatureReference] = (),
+        staleness: Mapping[FeatureReference, str] | None = None,
+        context: OnlineQueryContext | None = None,
+        query_name: str | None = None,
+        query_name_version: str | None = None,
+        correlation_id: str | None = None,
+        include_meta: bool = False,
+        explain: bool = False,
+        store_plan_stages: bool = False,
+        meta: Mapping[str, str] | None = None,
+        headers: Mapping[str, str] | None = None,
+    ) -> BulkOnlineQueryResult:
+        """Run a query using a pre-serialized plan.
+
+        This is a protected method for internal use and testing.
+
+        Parameters
+        ----------
+        serialized_plan_bytes
+            The serialized BatchPlan protobuf bytes
+        input
+            The input data, either as a mapping of features to values or as a PyArrow table
+        output
+            The output features to compute
+        staleness
+            Maximum staleness overrides for features
+        context
+            Query context including environment and tags
+        query_name
+            The name of the query
+        query_name_version
+            The version of the query
+        correlation_id
+            Correlation ID for logging
+        include_meta
+            Whether to include metadata in the response
+        explain
+            Whether to include explain output
+        store_plan_stages
+            Whether to store plan stages
+        meta
+            Customer metadata tags
+        headers
+            Additional headers to provide with the request
+
+        Returns
+        -------
+        OnlineQueryResult
+            The query result
+        """
+        ...
+
     async def offline_query(
         self,
         input: Union[QueryInput, OfflineQueryInputUri] | None = None,
@@ -502,6 +646,8 @@ class AsyncChalkClient:
         max_retries: int | None = None,
         query_name: str | None = None,
         query_name_version: str | None = None,
+        *,
+        input_sql: str | None = None,
     ) -> Dataset:
         """Compute feature values from the offline store or by running offline/online resolvers.
         See `Dataset` for more information.
@@ -516,6 +662,9 @@ class AsyncChalkClient:
             an observation in line with the timestamp in `input_times`.
         input_times
             A list of the times of the observations from `input`.
+        input_sql
+            An alternative to `input`: a ChalkSQL query that returns values
+            to use as inputs.
         output
             The features that you'd like to sample, if they exist.
             If an output feature was never computed for a sample (row) in

chalk/client/client_async_impl.py

@@ -18,8 +18,10 @@ from chalk.client.client_impl import ChalkAPIClientImpl, OnlineQueryResponseImpl
 from chalk.client.dataset import Dataset, DatasetImpl
 from chalk.client.models import (
     BulkOnlineQueryResponse,
+    BulkOnlineQueryResult,
     ChalkError,
     FeatureStatisticsResponse,
+    PlanQueryResponse,
     UploadFeaturesResponse,
     WhoAmIResponse,
 )
@@ -142,6 +144,26 @@ class AsyncChalkClientImpl(AsyncChalkClient):
             partial(self._client.query_bulk, *args, **kwargs),
         )
 
+    async def plan_query(
+        self,
+        *args: Any,
+        **kwargs: Any,
+    ) -> PlanQueryResponse:
+        return await self._current_loop().run_in_executor(
+            self._executor,
+            partial(self._client.plan_query, *args, **kwargs),
+        )
+
+    async def _run_serialized_query(
+        self,
+        *args: Any,
+        **kwargs: Any,
+    ) -> BulkOnlineQueryResult:
+        return await self._current_loop().run_in_executor(
+            self._executor,
+            partial(self._client._run_serialized_query, *args, **kwargs),  # pyright: ignore[reportPrivateUsage]
+        )
+
     async def offline_query(
         self,
         *args: Any,