planar 0.5.0__py3-none-any.whl

planar/workflows/step_metadata.py

@@ -0,0 +1,357 @@
+"""
+Utility functions for gathering rich metadata for workflow steps.
+
+This module provides helper functions for getting step-specific
+metadata for the various step types in Planar workflows.
+"""
+
+from typing import Annotated, Dict, List, Literal, Optional, Union
+from uuid import UUID
+
+from pydantic import BaseModel, Field
+from sqlmodel import select
+
+from planar.ai.models import ToolCall
+from planar.ai.utils import AgentSerializeable, get_agent_serializable
+from planar.human import HumanTask, get_human_tasks
+from planar.logging import get_logger
+from planar.object_registry import ObjectRegistry
+from planar.rules.models import (
+    RuleSerializeable,
+)
+from planar.rules.rule_configuration import rule_configuration
+from planar.session import get_session
+from planar.workflows.models import StepType, WorkflowStep
+
+logger = get_logger(__name__)
+
+
+class HumanTaskMetadata(BaseModel):
+    """Metadata wrapper for human task steps."""
+
+    step_type: Literal[StepType.HUMAN_IN_THE_LOOP] = StepType.HUMAN_IN_THE_LOOP
+    human_task: HumanTask
+
+
+class AgentMetadata(BaseModel):
+    """Metadata wrapper for agent steps."""
+
+    step_type: Literal[StepType.AGENT] = StepType.AGENT
+    agent: AgentSerializeable
+
+
+class RuleMetadata(BaseModel):
+    """Metadata wrapper for rule steps."""
+
+    step_type: Literal[StepType.RULE] = StepType.RULE
+    rule: RuleSerializeable
+
+
+class ToolCallMetadata(BaseModel):
+    """Metadata wrapper for tool call steps."""
+
+    step_type: Literal[StepType.TOOL_CALL] = StepType.TOOL_CALL
+    tool_call: ToolCall
+
+
+StepMetadata = Annotated[
+    Union[
+        HumanTaskMetadata,
+        AgentMetadata,
+        RuleMetadata,
+        ToolCallMetadata,
+    ],
+    Field(discriminator="step_type"),
+]
+
+
+def extract_simple_name(fully_qualified_name: str) -> str:
+    """
+    Extract the last part of a fully qualified name.
+
+    For example: 'module.directory.fn_name' becomes 'fn_name'.
+    If there are no periods in the string, returns the original string.
+
+    Args:
+        fully_qualified_name: A possibly fully qualified name with dot separators
+
+    Returns:
+        The last part of the name
+    """
+    return (
+        fully_qualified_name.split(".")[-1]
+        if "." in fully_qualified_name
+        else fully_qualified_name
+    )
+
+
+async def get_human_step_metadata(
+    workflow_id: UUID, step_id: int
+) -> Optional[HumanTaskMetadata]:
+    """
+    Get metadata for a human-in-the-loop step.
+
+    Args:
+        workflow_id: The ID of the workflow
+        step_id: The ID of the step
+
+    Returns:
+        A HumanTaskMetadata object, or None if no task is found
+    """
+    # Get human tasks for this workflow
+    tasks = await get_human_tasks(workflow_id=workflow_id)
+
+    # Tasks are ordered by creation time, newest first
+    # For now, we'll assume the most recent task is associated with this step
+    # In the future, we may want to add a step_id field to HumanTask for direct correlation
+    if not tasks:
+        logger.debug("no human tasks found", workflow_id=workflow_id, step_id=step_id)
+        return None
+
+    task = tasks[0]
+    logger.debug(
+        "found human task",
+        task_id=task.id,
+        workflow_id=workflow_id,
+        step_id=step_id,
+    )
+    return HumanTaskMetadata(human_task=task)
+
+
+async def get_agent_step_metadata(
+    workflow_id: UUID, step_id: int, registry: ObjectRegistry
+) -> Optional[AgentMetadata]:
+    """
+    Get metadata for an agent step.
+
+    Args:
+        workflow_id: The ID of the workflow
+        step_id: The ID of the step
+        registry: ObjectRegistry instance for looking up agents
+
+    Returns:
+        An AgentMetadata object, or None if no metadata is found
+    """
+    session = get_session()
+
+    logger.debug(
+        "getting agent step metadata", workflow_id=workflow_id, step_id=step_id
+    )
+    async with session.begin_read():
+        step = (
+            await session.exec(
+                select(WorkflowStep).where(
+                    (WorkflowStep.workflow_id == workflow_id)
+                    & (WorkflowStep.step_id == step_id)
+                )
+            )
+        ).first()
+
+    if not step or not step.display_name:
+        logger.debug(
+            "agent step or display_name not found",
+            workflow_id=workflow_id,
+            step_id=step_id,
+        )
+        return None
+
+    agent_name = step.display_name
+    logger.debug("agent name from step display_name", agent_name=agent_name)
+
+    agent_serializable = await get_agent_serializable(
+        agent_name=agent_name, registry=registry
+    )
+
+    if not agent_serializable:
+        logger.debug("agent serializable not found", agent_name=agent_name)
+        return None
+
+    logger.info("agent metadata retrieved", agent_name=agent_name)
+    return AgentMetadata(agent=agent_serializable)
+
+
+async def get_rule_step_metadata(
+    workflow_id: UUID, step_id: int, registry: ObjectRegistry
+) -> Optional[RuleMetadata]:
+    """
+    Get metadata for a rule step.
+
+    Args:
+        workflow_id: The ID of the workflow
+        step_id: The ID of the step
+
+    Returns:
+        A RuleMetadata object, or None if no metadata is found
+    """
+    session = get_session()
+
+    logger.debug("getting rule step metadata", workflow_id=workflow_id, step_id=step_id)
+    async with session.begin_read():
+        step = (
+            await session.exec(
+                select(WorkflowStep).where(
+                    (WorkflowStep.workflow_id == workflow_id)
+                    & (WorkflowStep.step_id == step_id)
+                )
+            )
+        ).first()
+
+    if not step:
+        logger.debug("rule step not found", workflow_id=workflow_id, step_id=step_id)
+        return None
+
+    rule_name = extract_simple_name(step.function_name)
+    logger.debug("rule name extracted", rule_name=rule_name)
+
+    rule = next((r for r in registry.get_rules() if r.name == rule_name), None)
+
+    if not rule:
+        logger.debug("rule not found in registry", rule_name=rule_name)
+        return None
+
+    configs = await rule_configuration.read_configs_with_default(
+        rule_name, rule.to_config()
+    )
+    logger.debug(
+        "retrieved configs for rule",
+        count=len(configs),
+        rule_name=rule_name,
+    )
+
+    rule_serializable = RuleSerializeable(
+        input_schema=rule.input.model_json_schema(),
+        output_schema=rule.output.model_json_schema(),
+        name=rule_name,
+        description=step.display_name or rule_name,
+        configs=configs,
+    )
+    logger.info("rule metadata retrieved", rule_name=rule_name)
+    return RuleMetadata(rule=rule_serializable)
+
+
+async def get_tool_call_step_metadata(
+    workflow_id: UUID, step_id: int
+) -> Optional[ToolCallMetadata]:
+    """
+    Get metadata for a tool call step.
+
+    Args:
+        workflow_id: The ID of the workflow
+        step_id: The ID of the step
+
+    Returns:
+        A ToolCallMetadata object, or None if no metadata is found
+    """
+    session = get_session()
+
+    logger.debug(
+        "getting tool call step metadata", workflow_id=workflow_id, step_id=step_id
+    )
+    async with session.begin_read():
+        step = (
+            await session.exec(
+                select(WorkflowStep).where(
+                    (WorkflowStep.workflow_id == workflow_id)
+                    & (WorkflowStep.step_id == step_id)
+                )
+            )
+        ).first()
+
+        if not step:
+            logger.debug(
+                "tool call step not found", workflow_id=workflow_id, step_id=step_id
+            )
+            return None
+        tool_name = extract_simple_name(step.function_name)
+        logger.debug("tool name extracted", tool_name=tool_name)
+
+        # Get the parent step if available
+        parent_step = None
+        if step.parent_step_id:
+            logger.debug(
+                "fetching parent step for tool call step",
+                parent_step_id=step.parent_step_id,
+                step_id=step_id,
+            )
+            parent_step = (
+                await session.exec(
+                    select(WorkflowStep).where(
+                        (WorkflowStep.workflow_id == workflow_id)
+                        & (WorkflowStep.step_id == step.parent_step_id)
+                    )
+                )
+            ).first()
+            logger.debug("parent step found", found=parent_step is not None)
+
+    tool_call_id = None
+    if parent_step and parent_step.result and isinstance(parent_step.result, dict):
+        if "tool_calls" in parent_step.result:
+            for tc in parent_step.result["tool_calls"]:
+                if tc.get("name") == tool_name:
+                    tool_call_id = tc.get("id")
+                    logger.debug(
+                        "found tool_call_id for tool in parent step result",
+                        tool_call_id=tool_call_id,
+                        tool_name=tool_name,
+                    )
+                    break
+
+    if not tool_call_id:
+        logger.debug("could not determine tool_call_id for tool", tool_name=tool_name)
+
+    tool_call_obj = ToolCall(
+        id=tool_call_id, name=tool_name, arguments=step.kwargs or {}
+    )
+    logger.info("tool call metadata retrieved", tool_name=tool_name)
+    return ToolCallMetadata(tool_call=tool_call_obj)
+
+
+async def get_steps_metadata(
+    steps: List[WorkflowStep], registry: ObjectRegistry
+) -> Dict[int, StepMetadata]:
+    """
+    Get metadata for multiple steps efficiently.
+
+    Args:
+        steps: A list of workflow steps
+        registry: Optional ObjectRegistry instance for looking up agents
+
+    Returns:
+        A dictionary mapping step_id to strongly-typed StepMetadata objects
+    """
+    result = {}
+
+    human_steps = [s for s in steps if s.step_type == StepType.HUMAN_IN_THE_LOOP]
+    agent_steps = [s for s in steps if s.step_type == StepType.AGENT]
+    rule_steps = [s for s in steps if s.step_type == StepType.RULE]
+    tool_call_steps = [s for s in steps if s.step_type == StepType.TOOL_CALL]
+
+    if human_steps:
+        for step in human_steps:
+            metadata = await get_human_step_metadata(step.workflow_id, step.step_id)
+            if metadata:
+                result[step.step_id] = metadata
+
+    if agent_steps:
+        for step in agent_steps:
+            metadata = await get_agent_step_metadata(
+                step.workflow_id, step.step_id, registry
+            )
+            if metadata:
+                result[step.step_id] = metadata
+
+    if rule_steps:
+        for step in rule_steps:
+            metadata = await get_rule_step_metadata(
+                step.workflow_id, step.step_id, registry
+            )
+            if metadata:
+                result[step.step_id] = metadata
+
+    if tool_call_steps:
+        for step in tool_call_steps:
+            metadata = await get_tool_call_step_metadata(step.workflow_id, step.step_id)
+            if metadata:
+                result[step.step_id] = metadata
+
+    return result

planar/workflows/step_testing_utils.py

@@ -0,0 +1,86 @@
+from sqlmodel import col, select
+
+from planar.session import get_session
+from planar.workflows.models import WorkflowStep
+
+
+async def get_step_parent(step: WorkflowStep) -> WorkflowStep | None:
+    """Get the parent step of the given step.
+
+    Args:
+        step: The step to get the parent of
+
+    Returns:
+        The parent step, or None if the step has no parent
+    """
+    if step.parent_step_id is None:
+        return None
+
+    session = get_session()
+    return (
+        await session.exec(
+            select(WorkflowStep)
+            .where(col(WorkflowStep.workflow_id) == step.workflow_id)
+            .where(col(WorkflowStep.step_id) == step.parent_step_id)
+        )
+    ).first()
+
+
+async def get_step_children(step: WorkflowStep) -> list[WorkflowStep]:
+    """Get all direct child steps of the given step.
+
+    Args:
+        step: The step to get the children of
+
+    Returns:
+        A list of child steps
+    """
+    session = get_session()
+    result = await session.exec(
+        select(WorkflowStep)
+        .where(col(WorkflowStep.workflow_id) == step.workflow_id)
+        .where(col(WorkflowStep.parent_step_id) == step.step_id)
+    )
+    return list(result.all())
+
+
+async def get_step_descendants(step: WorkflowStep) -> list[WorkflowStep]:
+    """Get all descendant steps (children, grandchildren, etc.) of the given step.
+
+    Args:
+        step: The step to get the descendants of
+
+    Returns:
+        A list of all descendant steps
+    """
+    descendants = await get_step_children(step)
+    result_descendants = descendants.copy()
+
+    # For each child, recursively get their descendants
+    for child in descendants:
+        child_descendants = await get_step_descendants(child)
+        result_descendants.extend(child_descendants)
+
+    return result_descendants
+
+
+async def get_step_ancestors(step: WorkflowStep) -> list[WorkflowStep]:
+    """Get all ancestor steps (parent, grandparent, etc.) of the given step.
+
+    Args:
+        step: The step to get the ancestors of
+
+    Returns:
+        A list of all ancestor steps, ordered from parent to oldest ancestor
+    """
+    ancestors = []
+    current = step
+
+    while current.parent_step_id is not None:
+        parent = await get_step_parent(current)
+        if parent is None:
+            break
+        ancestors.append(parent)
+        current = parent
+
+    return ancestors

planar/workflows/sub_workflow_runner.py

@@ -0,0 +1,191 @@
+# This module contains a code for running workflows and steps as background
+# tasks.
+# This functionality is very complex and not well tested, so it is going to be
+# disabled for now. In the future if there's demand for it, we can enable it
+# again.
+# For reference, this is the commit where it was enabled:
+# 8c6b7decbdfc01072b5a0be66597690a487455d5
+from asyncio import Future, create_task
+from dataclasses import dataclass
+from typing import Any, Coroutine, Mapping, Sequence, cast
+from uuid import UUID
+
+from planar.session import get_engine, session_context
+from planar.task_local import TaskLocal
+from planar.utils import P, R, T, U
+from planar.workflows.context import (
+    ExecutionContext,
+    get_context,
+    in_context,
+    set_context,
+)
+from planar.workflows.wrappers import StepWrapper, WorkflowWrapper
+
+
+@dataclass(kw_only=True)
+class SubWorkflowCall:
+    wrapper: WorkflowWrapper
+    args: Sequence[Any]
+    kwargs: Mapping[str, Any]
+    started: Future[UUID]
+
+
+# When a workflow is running, it will may call steps and other workflows. This
+# class encapsulates the logic for handling the various ways in which this can
+# happen.
+
+# The simplest case is when steps/workflows are called using `await` directly
+# in the current task (when "in_context").
+
+
+# For workflows this means we'll start the workflow as a step (every workflow has an
+# auto generated "start step"), which makes the operation durable. Additionally
+# we call the "wait_for_completion" helper, which waits for the workflow to
+# complete (up to a timeout) and returns the result.
+#
+# For steps it is even simpler after all, calling steps is the most common thing that
+# can be done in a workflow.
+#
+# Things start to get more complicated when workflows or steps are called using
+# `asyncio.create_task`, which means it will run in a separate task without
+# blocking the current one.
+#
+# A lot of the logic in SubWorkflowRunner is about dealing with multiple
+# background workflows/steps starting at the same time. To make these
+# operations durable, we must do some magic to force the concurrent workflows
+# to start in the order they were called.
+#
+# The initial work is done by the "run" method, which cannot be an async
+# function since it must be able to access the current workflow context. If we
+# used an async function, it would only start executing when the new task
+# started. OTOH we can only know if a step/workflow was called in a separate task
+# in an async function (the "context forwarder"). So this is what we do:
+#
+#   - in the "run" method, we collect the all information about the calls in a list,
+#     (pending_sub_workflow_calls) withou actually calling anything, and then we call
+#     the appropriate context forwarder for steps/worfklows.
+#   - in the workflow context forwarder, if we were called in a separate task we
+#     invoke the "start_sub_workflow" function, which will search for the current
+#     workflow start function in the pending list and invoke the workflow's "waiter",
+#     an async function which will wait for a signal that the workflow has started
+#     (signal passed via a future).
+#   - after the last workflow is discovered, we invoke a "starter" task that
+#     will start all workflows in the correct order, signaling futures associated
+#     with the workflow
+#
+# The above magic is what ensures that child concurrent workflows will start in a
+# predictable order, and will be durable against restarts.
+#
+# For steps, the only thing that makes sense is to run it in a separate
+# workflow, or else there would be no ordering of completion, and that's why
+# every step has an "auto workflow", thus entering in the same path as
+# concurrent subworkflows described above.
+class SubWorkflowRunner:
+    def __init__(self):
+        self._pending_sub_workflow_calls: list[SubWorkflowCall] = []
+
+    def run(
+        self,
+        wrapper: WorkflowWrapper[P, T, U, R] | StepWrapper[P, T, U, R],
+        *args: P.args,
+        **kwargs: P.kwargs,
+    ) -> Coroutine[T, U, R]:
+        async def starter(ctx: ExecutionContext):
+            # Main starter task that will call the start steps in the order
+            # they were called
+            set_context(ctx)
+            sub_workflow_calls = self._pending_sub_workflow_calls[:]
+            self._pending_sub_workflow_calls.clear()
+            async with session_context(get_engine()):
+                for sub_workflow_call in sub_workflow_calls:
+                    args = sub_workflow_call.args
+                    kwargs = sub_workflow_call.kwargs
+                    start_step = sub_workflow_call.wrapper.start_step
+                    workflow_id = await start_step(*args, **kwargs)
+                    sub_workflow_call.started.set_result(workflow_id)
+
+        async def waiter(future: Future[UUID], wf_wrapper: WorkflowWrapper[P, T, U, R]):
+            # only call wait_for_completion after receiving the workflow_id
+            # from the future
+            workflow_id = await future
+            completion_coro = wf_wrapper.wait_for_completion(workflow_id)
+            return await cast(Coroutine[T, U, R], completion_coro)
+
+        def start_sub_workflow(wf_wrapper: WorkflowWrapper[P, T, U, R]):
+            ctx = get_context()
+
+            # find the matching call instance
+            future = None
+            index = 0
+            for index, sub_workflow_call in enumerate(self._pending_sub_workflow_calls):
+                if sub_workflow_call.wrapper == wf_wrapper:
+                    future = sub_workflow_call.started
+                    break
+
+            if index == len(self._pending_sub_workflow_calls) - 1:
+                # Last subworkflow. Run the starter task that will start all
+                # workflows in the correct order
+                create_task(starter(ctx))
+
+            assert future
+            return waiter(future, wf_wrapper)
+
+        async def workflow_context_forwarder(
+            parent_execution_context: ExecutionContext,
+            wf_wrapper: WorkflowWrapper[P, T, U, R],
+        ):
+            if not in_context():
+                # invoke "start_sub_workflow" after forwarding the parent context
+                set_context(parent_execution_context)
+                return await cast(Coroutine[T, U, R], start_sub_workflow(wf_wrapper))
+
+            # Simple case, no need to start any starter task.
+            # Clear pending calls and invoke the start step directly.
+            assert len(self._pending_sub_workflow_calls) == 1
+            self._pending_sub_workflow_calls.clear()
+            workflow_id = await wf_wrapper.start_step(*args, **kwargs)
+            # Wait for completion
+            completion_coro = wf_wrapper.wait_for_completion(workflow_id)
+            return await cast(Coroutine[T, U, R], completion_coro)
+
+        async def step_context_forwarder(
+            parent_execution_context: ExecutionContext,
+            step_wrapper: StepWrapper,
+        ):
+            if not in_context():
+                # Invoke the workflow context forwarder, passing in the auto workflow
+                return await workflow_context_forwarder(
+                    parent_execution_context, step_wrapper.auto_workflow
+                )
+            # Simple case. Clear the pending calls and invoke the step
+            # wrapper directly.
+            assert len(self._pending_sub_workflow_calls) == 1
+            self._pending_sub_workflow_calls.clear()
+            return await step_wrapper.wrapper(*args, **kwargs)
+
+        # add the workflow wrapper, along with args and a future to the pending list
+        self._pending_sub_workflow_calls.append(
+            SubWorkflowCall(
+                wrapper=wrapper.auto_workflow
+                if isinstance(wrapper, StepWrapper)
+                else wrapper,
+                started=Future(),
+                args=args,
+                kwargs=kwargs,
+            )
+        )
+
+        # invoke the correct context forwarder
+        if isinstance(wrapper, StepWrapper):
+            return step_context_forwarder(get_context(), wrapper)
+        else:
+            return workflow_context_forwarder(get_context(), wrapper)
+
+
+data: TaskLocal[SubWorkflowRunner] = TaskLocal()
+
+
+def get_sub_workflow_runner() -> SubWorkflowRunner:
+    if not data.is_set():
+        data.set(SubWorkflowRunner())
+    return data.get()