chalkpy 2.166.2__tar.gz → 2.167.1__tar.gz

{chalkpy-2.166.2 → chalkpy-2.167.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: chalkpy
-Version: 2.166.2
+Version: 2.167.1
 Summary: Python SDK for Chalk
 Author: Chalk AI, Inc.
 Project-URL: Homepage, https://chalk.ai

{chalkpy-2.166.2 → chalkpy-2.167.1}/chalk/_gen/chalk/dataframe/v1/dataframe_pb2.py

@@ -19,7 +19,7 @@ from chalk._gen.chalk.graph.v2 import sources_pb2 as chalk_dot_graph_dot_v2_dot_
 
 
 DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(
-    b'\n"chalk/dataframe/v1/dataframe.proto\x12\x12\x63halk.dataframe.v1\x1a\x1a\x63halk/arrow/v1/arrow.proto\x1a$chalk/expression/v1/expression.proto\x1a\x1c\x63halk/graph/v2/sources.proto"\x80\x06\n\x10\x44\x61taFrameOperand\x12#\n\x0cvalue_string\x18\x01 \x01(\tH\x00R\x0bvalueString\x12\x1d\n\tvalue_int\x18\x02 \x01(\x03H\x00R\x08valueInt\x12\x1f\n\nvalue_bool\x18\x03 \x01(\x08H\x00R\tvalueBool\x12;\n\nvalue_none\x18\x04 \x01(\x0b\x32\x1a.chalk.dataframe.v1.PyNoneH\x00R\tvalueNone\x12;\n\nvalue_list\x18\x05 \x01(\x0b\x32\x1a.chalk.dataframe.v1.PyListH\x00R\tvalueList\x12;\n\nvalue_dict\x18\x06 \x01(\x0b\x32\x1a.chalk.dataframe.v1.PyDictH\x00R\tvalueDict\x12X\n\x15value_dataframe_index\x18\x07 \x01(\x0b\x32".chalk.dataframe.v1.DataFrameIndexH\x00R\x13valueDataframeIndex\x12;\n\x0c\x61rrow_schema\x18\x08 \x01(\x0b\x32\x16.chalk.arrow.v1.SchemaH\x00R\x0b\x61rrowSchema\x12\x44\n\x0b\x61rrow_table\x18\t \x01(\x0b\x32!.chalk.arrow.v1.TableParquetBytesH\x00R\narrowTable\x12O\n\x0funderscore_expr\x18\n \x01(\x0b\x32$.chalk.expression.v1.LogicalExprNodeH\x00R\x0eunderscoreExpr\x12K\n\rlibchalk_expr\x18\x0b \x01(\x0b\x32$.chalk.expression.v1.LogicalExprNodeH\x00R\x0clibchalkExpr\x12J\n\x0b\x64\x61ta_source\x18\x0c \x01(\x0b\x32\'.chalk.dataframe.v1.DataFrameDataSourceH\x00R\ndataSourceB\t\n\x07operand"^\n\x13\x44\x61taFrameDataSource\x12G\n\x0f\x64\x61tabase_source\x18\x01 \x01(\x0b\x32\x1e.chalk.graph.v2.DatabaseSourceR\x0e\x64\x61tabaseSource">\n\x0e\x44\x61taFrameIndex\x12,\n\x12\x64\x61taframe_op_index\x18\x01 \x01(\x05R\x10\x64\x61taframeOpIndex"\x08\n\x06PyNone"M\n\x06PyList\x12\x43\n\nlist_items\x18\x01 \x03(\x0b\x32$.chalk.dataframe.v1.DataFrameOperandR\tlistItems"\x97\x01\n\x0bPyDictEntry\x12\x41\n\tentry_key\x18\x01 \x01(\x0b\x32$.chalk.dataframe.v1.DataFrameOperandR\x08\x65ntryKey\x12\x45\n\x0b\x65ntry_value\x18\x02 \x01(\x0b\x32$.chalk.dataframe.v1.DataFrameOperandR\nentryValue"L\n\x06PyDict\x12\x42\n\x0c\x64ict_entries\x18\x01 \x03(\x0b\x32\x1f.chalk.dataframe.v1.PyDictEntryR\x0b\x64ictEntries"\xfc\x01\n\x14\x44\x61taFrameConstructor\x12J\n\x0cself_operand\x18\x01 \x01(\x0b\x32".chalk.dataframe.v1.DataFrameIndexH\x00R\x0bselfOperand\x88\x01\x01\x12#\n\rfunction_name\x18\x02 \x01(\tR\x0c\x66unctionName\x12.\n\x04\x61rgs\x18\x03 \x01(\x0b\x32\x1a.chalk.dataframe.v1.PyListR\x04\x61rgs\x12\x32\n\x06kwargs\x18\x04 \x01(\x0b\x32\x1a.chalk.dataframe.v1.PyDictR\x06kwargsB\x0f\n\r_self_operand"]\n\rDataFramePlan\x12L\n\x0c\x63onstructors\x18\x01 \x03(\x0b\x32(.chalk.dataframe.v1.DataFrameConstructorR\x0c\x63onstructorsB\x92\x01\n\x16\x63om.chalk.dataframe.v1B\x0e\x44\x61taframeProtoP\x01\xa2\x02\x03\x43\x44X\xaa\x02\x12\x43halk.Dataframe.V1\xca\x02\x12\x43halk\\Dataframe\\V1\xe2\x02\x1e\x43halk\\Dataframe\\V1\\GPBMetadata\xea\x02\x14\x43halk::Dataframe::V1b\x06proto3'
+    b'\n"chalk/dataframe/v1/dataframe.proto\x12\x12\x63halk.dataframe.v1\x1a\x1a\x63halk/arrow/v1/arrow.proto\x1a$chalk/expression/v1/expression.proto\x1a\x1c\x63halk/graph/v2/sources.proto"\xa5\x06\n\x10\x44\x61taFrameOperand\x12#\n\x0cvalue_string\x18\x01 \x01(\tH\x00R\x0bvalueString\x12\x1d\n\tvalue_int\x18\x02 \x01(\x03H\x00R\x08valueInt\x12\x1f\n\nvalue_bool\x18\x03 \x01(\x08H\x00R\tvalueBool\x12;\n\nvalue_none\x18\x04 \x01(\x0b\x32\x1a.chalk.dataframe.v1.PyNoneH\x00R\tvalueNone\x12;\n\nvalue_list\x18\x05 \x01(\x0b\x32\x1a.chalk.dataframe.v1.PyListH\x00R\tvalueList\x12;\n\nvalue_dict\x18\x06 \x01(\x0b\x32\x1a.chalk.dataframe.v1.PyDictH\x00R\tvalueDict\x12X\n\x15value_dataframe_index\x18\x07 \x01(\x0b\x32".chalk.dataframe.v1.DataFrameIndexH\x00R\x13valueDataframeIndex\x12;\n\x0c\x61rrow_schema\x18\x08 \x01(\x0b\x32\x16.chalk.arrow.v1.SchemaH\x00R\x0b\x61rrowSchema\x12\x44\n\x0b\x61rrow_table\x18\t \x01(\x0b\x32!.chalk.arrow.v1.TableParquetBytesH\x00R\narrowTable\x12O\n\x0funderscore_expr\x18\n \x01(\x0b\x32$.chalk.expression.v1.LogicalExprNodeH\x00R\x0eunderscoreExpr\x12K\n\rlibchalk_expr\x18\x0b \x01(\x0b\x32$.chalk.expression.v1.LogicalExprNodeH\x00R\x0clibchalkExpr\x12J\n\x0b\x64\x61ta_source\x18\x0c \x01(\x0b\x32\'.chalk.dataframe.v1.DataFrameDataSourceH\x00R\ndataSource\x12#\n\x0cvalue_double\x18\r \x01(\x01H\x00R\x0bvalueDoubleB\t\n\x07operand"^\n\x13\x44\x61taFrameDataSource\x12G\n\x0f\x64\x61tabase_source\x18\x01 \x01(\x0b\x32\x1e.chalk.graph.v2.DatabaseSourceR\x0e\x64\x61tabaseSource">\n\x0e\x44\x61taFrameIndex\x12,\n\x12\x64\x61taframe_op_index\x18\x01 \x01(\x05R\x10\x64\x61taframeOpIndex"\x08\n\x06PyNone"M\n\x06PyList\x12\x43\n\nlist_items\x18\x01 \x03(\x0b\x32$.chalk.dataframe.v1.DataFrameOperandR\tlistItems"\x97\x01\n\x0bPyDictEntry\x12\x41\n\tentry_key\x18\x01 \x01(\x0b\x32$.chalk.dataframe.v1.DataFrameOperandR\x08\x65ntryKey\x12\x45\n\x0b\x65ntry_value\x18\x02 \x01(\x0b\x32$.chalk.dataframe.v1.DataFrameOperandR\nentryValue"L\n\x06PyDict\x12\x42\n\x0c\x64ict_entries\x18\x01 \x03(\x0b\x32\x1f.chalk.dataframe.v1.PyDictEntryR\x0b\x64ictEntries"\xfc\x01\n\x14\x44\x61taFrameConstructor\x12J\n\x0cself_operand\x18\x01 \x01(\x0b\x32".chalk.dataframe.v1.DataFrameIndexH\x00R\x0bselfOperand\x88\x01\x01\x12#\n\rfunction_name\x18\x02 \x01(\tR\x0c\x66unctionName\x12.\n\x04\x61rgs\x18\x03 \x01(\x0b\x32\x1a.chalk.dataframe.v1.PyListR\x04\x61rgs\x12\x32\n\x06kwargs\x18\x04 \x01(\x0b\x32\x1a.chalk.dataframe.v1.PyDictR\x06kwargsB\x0f\n\r_self_operand"]\n\rDataFramePlan\x12L\n\x0c\x63onstructors\x18\x01 \x03(\x0b\x32(.chalk.dataframe.v1.DataFrameConstructorR\x0c\x63onstructorsB\x92\x01\n\x16\x63om.chalk.dataframe.v1B\x0e\x44\x61taframeProtoP\x01\xa2\x02\x03\x43\x44X\xaa\x02\x12\x43halk.Dataframe.V1\xca\x02\x12\x43halk\\Dataframe\\V1\xe2\x02\x1e\x43halk\\Dataframe\\V1\\GPBMetadata\xea\x02\x14\x43halk::Dataframe::V1b\x06proto3'
 )
 
 _globals = globals()
@@ -31,21 +31,21 @@ if _descriptor._USE_C_DESCRIPTORS == False:
         "DESCRIPTOR"
     ]._serialized_options = b"\n\026com.chalk.dataframe.v1B\016DataframeProtoP\001\242\002\003CDX\252\002\022Chalk.Dataframe.V1\312\002\022Chalk\\Dataframe\\V1\342\002\036Chalk\\Dataframe\\V1\\GPBMetadata\352\002\024Chalk::Dataframe::V1"
     _globals["_DATAFRAMEOPERAND"]._serialized_start = 155
-    _globals["_DATAFRAMEOPERAND"]._serialized_end = 923
-    _globals["_DATAFRAMEDATASOURCE"]._serialized_start = 925
-    _globals["_DATAFRAMEDATASOURCE"]._serialized_end = 1019
-    _globals["_DATAFRAMEINDEX"]._serialized_start = 1021
-    _globals["_DATAFRAMEINDEX"]._serialized_end = 1083
-    _globals["_PYNONE"]._serialized_start = 1085
-    _globals["_PYNONE"]._serialized_end = 1093
-    _globals["_PYLIST"]._serialized_start = 1095
-    _globals["_PYLIST"]._serialized_end = 1172
-    _globals["_PYDICTENTRY"]._serialized_start = 1175
-    _globals["_PYDICTENTRY"]._serialized_end = 1326
-    _globals["_PYDICT"]._serialized_start = 1328
-    _globals["_PYDICT"]._serialized_end = 1404
-    _globals["_DATAFRAMECONSTRUCTOR"]._serialized_start = 1407
-    _globals["_DATAFRAMECONSTRUCTOR"]._serialized_end = 1659
-    _globals["_DATAFRAMEPLAN"]._serialized_start = 1661
-    _globals["_DATAFRAMEPLAN"]._serialized_end = 1754
+    _globals["_DATAFRAMEOPERAND"]._serialized_end = 960
+    _globals["_DATAFRAMEDATASOURCE"]._serialized_start = 962
+    _globals["_DATAFRAMEDATASOURCE"]._serialized_end = 1056
+    _globals["_DATAFRAMEINDEX"]._serialized_start = 1058
+    _globals["_DATAFRAMEINDEX"]._serialized_end = 1120
+    _globals["_PYNONE"]._serialized_start = 1122
+    _globals["_PYNONE"]._serialized_end = 1130
+    _globals["_PYLIST"]._serialized_start = 1132
+    _globals["_PYLIST"]._serialized_end = 1209
+    _globals["_PYDICTENTRY"]._serialized_start = 1212
+    _globals["_PYDICTENTRY"]._serialized_end = 1363
+    _globals["_PYDICT"]._serialized_start = 1365
+    _globals["_PYDICT"]._serialized_end = 1441
+    _globals["_DATAFRAMECONSTRUCTOR"]._serialized_start = 1444
+    _globals["_DATAFRAMECONSTRUCTOR"]._serialized_end = 1696
+    _globals["_DATAFRAMEPLAN"]._serialized_start = 1698
+    _globals["_DATAFRAMEPLAN"]._serialized_end = 1791
 # @@protoc_insertion_point(module_scope)

{chalkpy-2.166.2 → chalkpy-2.167.1}/chalk/_gen/chalk/dataframe/v1/dataframe_pb2.pyi

@@ -28,6 +28,7 @@ class DataFrameOperand(_message.Message):
         "underscore_expr",
         "libchalk_expr",
         "data_source",
+        "value_double",
     )
     VALUE_STRING_FIELD_NUMBER: _ClassVar[int]
     VALUE_INT_FIELD_NUMBER: _ClassVar[int]
@@ -41,6 +42,7 @@ class DataFrameOperand(_message.Message):
     UNDERSCORE_EXPR_FIELD_NUMBER: _ClassVar[int]
     LIBCHALK_EXPR_FIELD_NUMBER: _ClassVar[int]
     DATA_SOURCE_FIELD_NUMBER: _ClassVar[int]
+    VALUE_DOUBLE_FIELD_NUMBER: _ClassVar[int]
     value_string: str
     value_int: int
     value_bool: bool
@@ -53,6 +55,7 @@ class DataFrameOperand(_message.Message):
     underscore_expr: _expression_pb2.LogicalExprNode
     libchalk_expr: _expression_pb2.LogicalExprNode
     data_source: DataFrameDataSource
+    value_double: float
     def __init__(
         self,
         value_string: _Optional[str] = ...,
@@ -67,6 +70,7 @@ class DataFrameOperand(_message.Message):
         underscore_expr: _Optional[_Union[_expression_pb2.LogicalExprNode, _Mapping]] = ...,
         libchalk_expr: _Optional[_Union[_expression_pb2.LogicalExprNode, _Mapping]] = ...,
         data_source: _Optional[_Union[DataFrameDataSource, _Mapping]] = ...,
+        value_double: _Optional[float] = ...,
     ) -> None: ...
 
 class DataFrameDataSource(_message.Message):

chalkpy-2.167.1/chalk/_version.py

@@ -0,0 +1 @@
+__version__ = "2.167.1"

{chalkpy-2.166.2 → chalkpy-2.167.1}/chalk/df/LazyFramePlaceholder.py

@@ -2175,6 +2175,10 @@ def _convert_to_dataframe_proto(
             return dataframe_pb2.DataFrameOperand(
                 value_bool=value,
             )
+        if isinstance(value, float):
+            return dataframe_pb2.DataFrameOperand(
+                value_double=value,
+            )
         if isinstance(value, int):
             return dataframe_pb2.DataFrameOperand(
                 value_int=value,
@@ -2324,6 +2328,8 @@ def _convert_from_dataframe_proto(
             return value.value_string
         if value.HasField("value_int"):
             return value.value_int
+        if value.HasField("value_double"):
+            return value.value_double
         if value.HasField("value_bool"):
             return value.value_bool
         if value.HasField("value_none"):

chalkpy-2.167.1/chalk/workflows/__init__.py

@@ -0,0 +1,36 @@
+from chalk.workflows._children import current_workflow_id, parent_workflow_id, start_child
+from chalk.workflows._context import ChildWorkflowHandle, TaskHandle
+from chalk.workflows._definitions import (
+    TASK_REGISTRY,
+    WORKFLOW_REGISTRY,
+    TaskDefinition,
+    WorkflowDefinition,
+    task,
+    workflow,
+)
+from chalk.workflows._messages import (
+    WorkflowMessage,
+    WorkflowMessageTimeoutError,
+    local_messages,
+    send_message,
+    wait_for_message,
+)
+
+__all__ = (
+    "ChildWorkflowHandle",
+    "TASK_REGISTRY",
+    "WORKFLOW_REGISTRY",
+    "TaskDefinition",
+    "TaskHandle",
+    "WorkflowDefinition",
+    "WorkflowMessage",
+    "WorkflowMessageTimeoutError",
+    "current_workflow_id",
+    "local_messages",
+    "parent_workflow_id",
+    "send_message",
+    "start_child",
+    "task",
+    "wait_for_message",
+    "workflow",
+)

chalkpy-2.167.1/chalk/workflows/_children.py

@@ -0,0 +1,108 @@
+from __future__ import annotations
+
+import inspect
+from typing import Any, Mapping, TypeVar, overload
+
+from chalk.workflows._context import ChildWorkflowHandle, current_orchestrator
+from chalk.workflows._definitions import WorkflowDefinition
+
+R = TypeVar("R")
+
+
+def current_workflow_id() -> str:
+    """Return the id of the running workflow, from inside a `@workflow`.
+
+    It is the id `ChalkClient.trigger_workflow` returned (or the child's id, in a workflow
+    started with `start_child`), which other workflows can pass to `send_message`. A workflow
+    run directly in this process gets a generated `local-...` id.
+    """
+    orchestrator = current_orchestrator()
+    if orchestrator is None:
+        raise RuntimeError("current_workflow_id can only be called inside a @workflow")
+    return orchestrator.current_workflow_id()
+
+
+def parent_workflow_id() -> str | None:
+    """Return the id of the workflow that started this one with `start_child`, from inside a
+    `@workflow`, or `None` if it was not started as a child.
+
+    A child can use it to message its parent without waiting to hear from it first:
+
+    >>> from chalk.workflows import current_workflow_id, parent_workflow_id, send_message, workflow
+    >>> @workflow
+    ... async def sub_agent() -> None:
+    ...     parent_id = parent_workflow_id()
+    ...     if parent_id is not None:
+    ...         await send_message(parent_id, "ready", {"agent": current_workflow_id()})
+    """
+    orchestrator = current_orchestrator()
+    if orchestrator is None:
+        raise RuntimeError("parent_workflow_id can only be called inside a @workflow")
+    return orchestrator.parent_workflow_id()
+
+
+@overload
+async def start_child(
+    workflow: WorkflowDefinition[..., R], input: Mapping[str, Any] | None = None, *, workflow_id: str | None = None
+) -> ChildWorkflowHandle[R]: ...
+
+
+@overload
+async def start_child(
+    workflow: str, input: Mapping[str, Any] | None = None, *, workflow_id: str | None = None
+) -> ChildWorkflowHandle[Any]: ...
+
+
+async def start_child(
+    workflow: WorkflowDefinition[..., Any] | str,
+    input: Mapping[str, Any] | None = None,
+    *,
+    workflow_id: str | None = None,
+) -> ChildWorkflowHandle[Any]:
+    """Start another workflow as a child of the calling `@workflow`, and return a handle to it.
+
+    The child runs as its own workflow, on the same workers as the parent, with its own
+    history; await the handle for its result and use `handle.send_message(...)` to message
+    it. The workflow lineage view shows it as a child of the calling workflow.
+
+    A child's lifetime is bounded by its parent's: if the parent completes, fails or is
+    cancelled while the child is still running, the child is cancelled. It receives the
+    cancellation like any cancelled workflow, so its `finally` blocks still run (for example
+    to release a sandbox). Await every child before returning from the parent.
+
+    Like a workflow started with `ChalkClient.trigger_workflow`, the child is not retried as
+    a whole; retries are configured per `@task`.
+
+    Parameters
+    ----------
+    workflow
+        The child `@workflow`, or the name of one served by this workflow's workers.
+    input
+        The child's keyword arguments, by name, as for `ChalkClient.trigger_workflow`. Must
+        be JSON-serializable. `None` starts a child that takes no arguments.
+    workflow_id
+        The child's workflow id. Defaults to one derived from the parent's id and the child's
+        name. Starting a child with the id of a workflow that is still running fails.
+
+    Examples
+    --------
+    >>> from chalk.workflows import start_child, workflow
+    >>> @workflow
+    ... async def score_segment(segment: str) -> float:
+    ...     return 0.5
+    >>> @workflow
+    ... async def score_all(segments: list[str]) -> list[float]:
+    ...     children = [await start_child(score_segment, {"segment": s}) for s in segments]
+    ...     return [await child for child in children]
+    """
+    orchestrator = current_orchestrator()
+    if orchestrator is None:
+        raise RuntimeError("start_child can only be called inside a @workflow")
+    kwargs = dict(input) if input is not None else {}
+    if isinstance(workflow, WorkflowDefinition):
+        # Fail in the parent, where the mistake is, rather than in a child that cannot start.
+        try:
+            inspect.signature(workflow.fn).bind(**kwargs)
+        except TypeError as e:
+            raise TypeError(f"Invalid input for child workflow {workflow.name!r}: {e}") from None
+    return await orchestrator.start_child(workflow, kwargs, workflow_id)

chalkpy-2.167.1/chalk/workflows/_context.py

@@ -0,0 +1,241 @@
+from __future__ import annotations
+
+import asyncio
+import contextvars
+import uuid
+from datetime import timedelta
+from typing import TYPE_CHECKING, Any, Awaitable, Callable, Generator, Generic, Mapping, Protocol, Sequence, TypeVar
+
+if TYPE_CHECKING:
+    from chalk.workflows._definitions import TaskDefinition, WorkflowDefinition
+    from chalk.workflows._messages import WorkflowMessage
+
+R = TypeVar("R")
+
+
+class TaskHandle(Generic[R]):
+    """A handle to a task submitted with `task.submit(...)` inside a workflow.
+
+    Await the handle to obtain the task's result:
+
+    >>> handle = my_task.submit(1)  # doctest: +SKIP
+    >>> result = await handle  # doctest: +SKIP
+    """
+
+    def __init__(self, awaitable: Awaitable[R]):
+        super().__init__()
+        self._awaitable = awaitable
+
+    def __await__(self) -> Generator[Any, None, R]:
+        return self._awaitable.__await__()
+
+
+class ChildWorkflowHandle(Generic[R]):
+    """A handle to a child workflow started with `chalk.workflows.start_child(...)`.
+
+    Await the handle to obtain the child's result, and use `send_message` to message it:
+
+    >>> child = await start_child(review_account, {"account_id": "acct-1"})  # doctest: +SKIP
+    >>> await child.send_message("review", {"approved": True})  # doctest: +SKIP
+    >>> result = await child  # doctest: +SKIP
+    """
+
+    def __init__(
+        self,
+        workflow_id: str,
+        result: Awaitable[R],
+        send_message: Callable[[str, Any], Awaitable[None]],
+    ):
+        super().__init__()
+        self.workflow_id = workflow_id
+        self._result = result
+        self._send_message = send_message
+
+    def __await__(self) -> Generator[Any, None, R]:
+        return self._result.__await__()
+
+    async def send_message(self, name: str, payload: Any) -> None:
+        """Send the child a message, which it receives with `wait_for_message(name)`.
+        Equivalent to `chalk.workflows.send_message(handle.workflow_id, name, payload)`."""
+        await self._send_message(name, payload)
+
+
+class WorkflowOrchestrator(Protocol):
+    """Backend that executes task invocations made inside a running workflow.
+
+    Implementations: `LocalOrchestrator` (in-process asyncio) and the Temporal-backed
+    orchestrator in `chalk.workflows._temporal`. User code never sees this type.
+    """
+
+    def run_task(
+        self,
+        task: TaskDefinition[Any, Any],
+        args: Sequence[Any],
+        kwargs: Mapping[str, Any],
+    ) -> Awaitable[Any]: ...
+
+    def submit_task(
+        self,
+        task: TaskDefinition[Any, Any],
+        args: Sequence[Any],
+        kwargs: Mapping[str, Any],
+    ) -> TaskHandle[Any]: ...
+
+    def wait_for_message(self, name: str, timeout: timedelta | None) -> Awaitable[WorkflowMessage]: ...
+
+    def start_child(
+        self,
+        workflow: WorkflowDefinition[Any, Any] | str,
+        kwargs: Mapping[str, Any],
+        workflow_id: str | None,
+    ) -> Awaitable[ChildWorkflowHandle[Any]]: ...
+
+    def send_message(self, workflow_id: str, name: str, payload: Any) -> Awaitable[None]: ...
+
+    def current_workflow_id(self) -> str: ...
+
+    def parent_workflow_id(self) -> str | None: ...
+
+
+# Messages queued by `chalk.workflows.local_messages` for workflows run by a LocalOrchestrator.
+local_messages_var: contextvars.ContextVar[list[WorkflowMessage] | None] = contextvars.ContextVar(
+    "chalk_workflow_local_messages", default=None
+)
+
+current_orchestrator_var: contextvars.ContextVar[WorkflowOrchestrator | None] = contextvars.ContextVar(
+    "chalk_workflow_orchestrator", default=None
+)
+
+
+def current_orchestrator() -> WorkflowOrchestrator | None:
+    return current_orchestrator_var.get()
+
+
+def local_workflow_id(name: str) -> str:
+    return f"local-{name}-{uuid.uuid4().hex[:12]}"
+
+
+class LocalOrchestrator:
+    """Executes tasks in-process. Used when a workflow function is invoked directly
+    (local runs and unit tests) rather than under a durable execution backend."""
+
+    def __init__(self, workflow_id: str, parent_workflow_id: str | None = None):
+        super().__init__()
+        self.workflow_id = workflow_id
+        self._parent_workflow_id = parent_workflow_id
+        # This invocation's running children, by id; a child leaves when it finishes.
+        self._children: dict[str, asyncio.Future[Any]] = {}
+
+    def current_workflow_id(self) -> str:
+        return self.workflow_id
+
+    def parent_workflow_id(self) -> str | None:
+        return self._parent_workflow_id
+
+    async def _execute(
+        self,
+        task: TaskDefinition[Any, Any],
+        args: Sequence[Any],
+        kwargs: Mapping[str, Any],
+    ) -> Any:
+        attempts = task.retries + 1
+        last_error: BaseException | None = None
+        for attempt in range(attempts):
+            try:
+                if task.is_async:
+                    return await task.fn(*args, **kwargs)
+                # Sync tasks run in a thread so local workflows can still make progress
+                # on other submitted tasks, mirroring how sync tasks execute in a
+                # worker thread pool under a durable backend.
+                return await asyncio.to_thread(task.fn, *args, **kwargs)
+            except Exception as e:  # noqa: PERF203
+                last_error = e
+                if attempt + 1 < attempts and task.retry_delay is not None:
+                    await asyncio.sleep(task.retry_delay.total_seconds())
+        assert last_error is not None
+        raise last_error
+
+    def run_task(
+        self,
+        task: TaskDefinition[Any, Any],
+        args: Sequence[Any],
+        kwargs: Mapping[str, Any],
+    ) -> Awaitable[Any]:
+        return self._execute(task, args, kwargs)
+
+    def submit_task(
+        self,
+        task: TaskDefinition[Any, Any],
+        args: Sequence[Any],
+        kwargs: Mapping[str, Any],
+    ) -> TaskHandle[Any]:
+        return TaskHandle(asyncio.ensure_future(self._execute(task, args, kwargs)))
+
+    async def wait_for_message(self, name: str, timeout: timedelta | None) -> WorkflowMessage:
+        # A direct local run has no address anyone could send to, so it only sees messages
+        # queued in advance; waiting would hang forever.
+        queued = local_messages_var.get()
+        if queued is not None:
+            for i, message in enumerate(queued):
+                if message.name == name:
+                    return queued.pop(i)
+        raise RuntimeError(
+            f"No local message named {name!r}. A workflow run directly in this process only receives "
+            + "messages queued with chalk.workflows.local_messages(...); run it on the workflow orchestrator "
+            + "(ChalkClient.trigger_workflow or run_workflow) to receive messages sent to it."
+        )
+
+    async def start_child(
+        self,
+        workflow: WorkflowDefinition[Any, Any] | str,
+        kwargs: Mapping[str, Any],
+        workflow_id: str | None,
+    ) -> ChildWorkflowHandle[Any]:
+        if isinstance(workflow, str):
+            from chalk.workflows._definitions import WORKFLOW_REGISTRY
+
+            definition = WORKFLOW_REGISTRY.get(workflow)
+            if definition is None:
+                raise ValueError(f"No workflow named {workflow!r} is defined in this process")
+        else:
+            definition = workflow
+
+        child_id = workflow_id if workflow_id is not None else local_workflow_id(definition.name)
+        # The orchestrator refuses to start a workflow whose id is still running.
+        if child_id in self._children:
+            raise RuntimeError(f"Cannot start child workflow {child_id!r}: a child with that id is still running")
+
+        async def run_child() -> Any:
+            # The child runs in its own task (and so its own context): messages queued with
+            # local_messages(...) are the parent's, not the child's.
+            local_messages_var.set(None)
+            orchestrator = LocalOrchestrator(child_id, parent_workflow_id=self.workflow_id)
+            current_orchestrator_var.set(orchestrator)
+            try:
+                return await definition.fn(**kwargs)
+            finally:
+                await orchestrator.close()
+
+        child = asyncio.ensure_future(run_child())
+        self._children[child_id] = child
+        child.add_done_callback(lambda _: self._children.pop(child_id, None))
+        return ChildWorkflowHandle(
+            child_id,
+            child,
+            lambda name, payload: self.send_message(child_id, name, payload),
+        )
+
+    async def close(self) -> None:
+        """Cancel the children still running when this invocation exits, and wait for their
+        `finally` blocks, as the orchestrator does for a child whose parent closes first."""
+        children = list(self._children.values())
+        for child in children:
+            child.cancel()
+        await asyncio.gather(*children, return_exceptions=True)
+
+    async def send_message(self, workflow_id: str, name: str, payload: Any) -> None:
+        raise RuntimeError(
+            f"Cannot send message {name!r} to workflow {workflow_id!r}: a workflow run directly in this process "
+            + "has no orchestrator to deliver messages between workflows. Run it on the workflow orchestrator "
+            + "(ChalkClient.trigger_workflow or run_workflow) to message other workflows."
+        )

{chalkpy-2.166.2 → chalkpy-2.167.1}/chalk/workflows/_definitions.py

@@ -8,7 +8,13 @@ from typing import Any, Callable, Coroutine, Generic, Mapping, Sequence, TypeVar
 from typing_extensions import ParamSpec
 
 from chalk.utils.duration import Duration, parse_chalk_duration
-from chalk.workflows._context import LocalOrchestrator, TaskHandle, current_orchestrator, current_orchestrator_var
+from chalk.workflows._context import (
+    LocalOrchestrator,
+    TaskHandle,
+    current_orchestrator,
+    current_orchestrator_var,
+    local_workflow_id,
+)
 
 P = ParamSpec("P")
 R = TypeVar("R")
@@ -118,11 +124,13 @@ class WorkflowDefinition(Generic[P, R]):
             # Calling one workflow from another: run it inline in the caller's
             # orchestration context, so its tasks remain durable task invocations.
             return await self.fn(*args, **kwargs)
-        token: contextvars.Token[Any] = current_orchestrator_var.set(LocalOrchestrator())
+        orchestrator = LocalOrchestrator(local_workflow_id(self.name))
+        token: contextvars.Token[Any] = current_orchestrator_var.set(orchestrator)
         try:
             return await self.fn(*args, **kwargs)
         finally:
             current_orchestrator_var.reset(token)
+            await orchestrator.close()
 
     def __repr__(self) -> str:
         return f"chalk.workflows.workflow(name={self.name!r})"

{chalkpy-2.166.2 → chalkpy-2.167.1}/chalk/workflows/_messages.py

@@ -16,10 +16,15 @@ WORKFLOW_MESSAGE_SIGNAL = "chalk.workflow_message"
 class WorkflowMessage:
     """A message delivered to a running workflow, e.g. a reviewer's decision.
 
-    `sender_kind` and `sender_id` identify who sent it, as authenticated by the Chalk API
-    server (for example `"user"` and a user id, or `"service_token"` and a token id). They
-    cannot be set by the sender, so a workflow can use them to decide whose decisions to
-    accept.
+    `sender_kind` and `sender_id` identify who sent it. For a message sent through the Chalk
+    API server (`ChalkClient.send_workflow_message`), they are authenticated by the server
+    (for example `"user"` and a user id, or `"service_token"` and a token id) and cannot be
+    set by the sender, so a workflow can use them to decide whose decisions to accept.
+
+    A message sent by another workflow with `chalk.workflows.send_message` has `sender_kind`
+    `"workflow"` and the sending workflow's id as `sender_id`. Those are written by the
+    sending workflow itself and do not pass through the API server, so they are not
+    authenticated: do not use them to authorize a decision.
     """
 
     name: str
@@ -79,6 +84,46 @@ async def wait_for_message(name: str, *, timeout: timedelta | None = None) -> Wo
     return await orchestrator.wait_for_message(name, timeout)
 
 
+async def send_message(workflow_id: str, name: str, payload: Any) -> None:
+    """Send a message from inside a `@workflow` to another running workflow.
+
+    The receiving workflow gets it from `wait_for_message(name)`, with `sender_kind`
+    `"workflow"` and this workflow's id as `sender_id`. Use it, for example, for a child
+    started with `start_child` to report back to its parent (see `parent_workflow_id`), or to
+    reply to the `sender_id` of a message from another workflow.
+
+    The message is delivered by the workflow orchestrator directly, not through the Chalk API
+    server, and is recorded in this workflow's history, so the workflow lineage view shows it
+    as a message edge from this workflow to `workflow_id`. Because it does not pass through
+    the API server, the sender is self-declared rather than authenticated as it is for
+    `ChalkClient.send_workflow_message`; see `WorkflowMessage`.
+
+    Parameters
+    ----------
+    workflow_id
+        The id of the running workflow to message.
+    name
+        The message name the receiver waits for.
+    payload
+        The message body. Must be JSON-serializable.
+
+    Examples
+    --------
+    >>> from chalk.workflows import send_message, wait_for_message, workflow
+    >>> @workflow
+    ... async def echo() -> None:
+    ...     ping = await wait_for_message("ping")
+    ...     await send_message(ping.sender_id, "pong", ping.payload)
+    """
+    orchestrator = current_orchestrator()
+    if orchestrator is None:
+        raise RuntimeError(
+            "chalk.workflows.send_message can only be called inside a @workflow; "
+            + "use ChalkClient.send_workflow_message to message a workflow from outside one"
+        )
+    await orchestrator.send_message(workflow_id, name, payload)
+
+
 @contextlib.contextmanager
 def local_messages(messages: Sequence[tuple[str, Any]]) -> Iterator[None]:
     """Queue messages for workflows invoked directly in this process, such as in unit tests.