Improve prefect task logic

Author: jlowin

src/controlflow/core/flow.py

@@ -45,7 +45,7 @@ def add_task(self, task: "Task"):
         self._tasks[task.id] = task
 
     @contextmanager
-    def _context(self, create_prefect_flow_context: bool = True):
+    def create_context(self, create_prefect_flow_context: bool = True):
         if create_prefect_flow_context:
             prefect_ctx = prefect_flow_context(name=self.name)
         else:
@@ -55,7 +55,7 @@ def _context(self, create_prefect_flow_context: bool = True):
 
     def __enter__(self):
         # use stack so we can enter the context multiple times
-        cm = self._context()
+        cm = self.create_context()
         self._cm_stack.append(cm)
         return cm.__enter__()
 

src/controlflow/core/task.py

@@ -140,7 +140,7 @@ def __init__(
         super().__init__(**kwargs)
 
         self._prefect_task = PrefectTrackingTask(
-            name=f"Working on {self.friendly_name()}",
+            name=f"Working on {self.friendly_name()}...",
             description=self.instructions,
             tags=[self.__class__.__name__],
         )
@@ -307,6 +307,7 @@ async def run_once_async(self, agent: "Agent" = None, flow: "Flow" = None):
         controller = controlflow.Controller(tasks=[self], agents=agent, flow=flow)
         await controller.run_once_async()
 
+    @prefect.task(task_run_name=lambda _, args: f"Run {args['self'].friendly_name()}")
     def _run(
         self,
         raise_on_error: bool = True,
@@ -319,6 +320,8 @@ def _run(
         """
         from controlflow.core.flow import Flow, get_flow
 
+        self._prefect_task.is_started = True
+
         if max_iterations == NOTSET:
             max_iterations = controlflow.settings.max_task_iterations
         if max_iterations is None:
@@ -362,21 +365,17 @@ def run(
         If max_iterations is provided, the task will run at most that many times before raising an error.
         """
 
-        @prefect.task(task_run_name=f"Run Task {self.id}")
-        def _run():
-            gen = self._run(
-                raise_on_error=raise_on_error,
-                max_iterations=max_iterations,
-                flow=flow,
-                run_async=False,
-            )
-            while True:
-                try:
-                    next(gen)
-                except StopIteration as e:
-                    return e.value
-
-        return _run()
+        gen = self._run(
+            raise_on_error=raise_on_error,
+            max_iterations=max_iterations,
+            flow=flow,
+            run_async=False,
+        )
+        while True:
+            try:
+                next(gen)
+            except StopIteration as e:
+                return e.value
 
     async def run_async(
         self,
@@ -390,31 +389,27 @@ async def run_async(
         If max_iterations is provided, the task will run at most that many times before raising an error.
         """
 
-        @prefect.task(task_run_name=f"Run Task {self.id}")
-        async def _run():
-            gen = self._run(
-                raise_on_error=raise_on_error,
-                max_iterations=max_iterations,
-                flow=flow,
-                run_async=True,
-            )
-            while True:
-                try:
-                    await next(gen)
-                except StopIteration as e:
-                    return e.value
-
-        return await _run()
+        gen = self._run(
+            raise_on_error=raise_on_error,
+            max_iterations=max_iterations,
+            flow=flow,
+            run_async=True,
+        )
+        while True:
+            try:
+                await next(gen)
+            except StopIteration as e:
+                return e.value
 
     @contextmanager
-    def _context(self):
+    def create_context(self):
         stack = ctx.get("tasks", [])
         with ctx(tasks=stack + [self]):
             yield self
 
     def __enter__(self):
         # use stack so we can enter the context multiple times
-        self._cm_stack.append(self._context())
+        self._cm_stack.append(self.create_context())
         return self._cm_stack[-1].__enter__()
 
     def __exit__(self, *exc_info):

src/controlflow/decorators.py

@@ -101,7 +101,7 @@ def wrapper(
             **flow_kwargs,
         )
 
-        with flow_obj._context(create_prefect_flow_context=False):
+        with flow_obj.create_context(create_prefect_flow_context=False):
             with controlflow.instructions(instructions):
                 result = fn(*args, **kwargs)
 

src/controlflow/utilities/prefect.py

@@ -328,6 +328,30 @@ def skip(self):
 
 
 def prefect_task_context(**kwargs):
+    """
+    Creates a Prefect task that starts when the context is entered and ends when
+    it closes. This is useful for creating a Prefect task that is not tied to a
+    specific function but governs a block of code. Note that some features, like
+    retries and caching, will not work.
+    """
+    supported_kwargs = {
+        "name",
+        "description",
+        "task_run_name",
+        "tags",
+        "version",
+        "timeout_seconds",
+        "log_prints",
+        "on_completion",
+        "on_failure",
+    }
+    unsupported_kwargs = set(kwargs.keys()) - set(supported_kwargs)
+    if unsupported_kwargs:
+        raise ValueError(
+            f"Unsupported keyword arguments for a task context provided: "
+            f"{unsupported_kwargs}. Consider using a @task-decorated function instead."
+        )
+
     @contextmanager
     @prefect.task(**kwargs)
     def task_context():
@@ -337,6 +361,31 @@ def task_context():
 
 
 def prefect_flow_context(**kwargs):
+    """
+    Creates a Prefect flow that starts when the context is entered and ends when
+    it closes. This is useful for creating a Prefect flow that is not tied to a
+    specific function but governs a block of code. Note that some features, like
+    retries and caching, will not work.
+    """
+
+    supported_kwargs = {
+        "name",
+        "description",
+        "flow_run_name",
+        "tags",
+        "version",
+        "timeout_seconds",
+        "log_prints",
+        "on_completion",
+        "on_failure",
+    }
+    unsupported_kwargs = set(kwargs.keys()) - set(supported_kwargs)
+    if unsupported_kwargs:
+        raise ValueError(
+            f"Unsupported keyword arguments for a flow context provided: "
+            f"{unsupported_kwargs}. Consider using a @flow-decorated function instead."
+        )
+
     @contextmanager
     @prefect.flow(**kwargs)
     def flow_context():