ibm-watsonx-orchestrate 1.6.2__py3-none-any.whl → 1.6.4__py3-none-any.whl

ibm_watsonx_orchestrate/flow_builder/flows/flow.py

@@ -15,6 +15,7 @@ import logging
 import copy
 import uuid
 import pytz
+import os
 
 from typing_extensions import Self
 from pydantic import BaseModel, Field, SerializeAsAny
@@ -25,11 +26,12 @@ from ibm_watsonx_orchestrate.client.tools.tempus_client import TempusClient
 from ibm_watsonx_orchestrate.client.utils import instantiate_client
 from ..types import (
     EndNodeSpec, Expression, ForeachPolicy, ForeachSpec, LoopSpec, BranchNodeSpec, MatchPolicy, PromptLLMParameters, PromptNodeSpec, 
-    StartNodeSpec, ToolSpec, JsonSchemaObject, ToolRequestBody, ToolResponseBody, UserFieldKind, UserFieldOption, UserFlowSpec, UserNodeSpec, WaitPolicy
+    StartNodeSpec, ToolSpec, JsonSchemaObject, ToolRequestBody, ToolResponseBody, UserFieldKind, UserFieldOption, UserFlowSpec, UserNodeSpec, WaitPolicy,
+    DocProcSpec, TextExtractionResponse, File, DecisionsNodeSpec, DecisionsRule
 )
 from .constants import CURRENT_USER, START, END, ANY_USER
 from ..node import (
-    EndNode, Node, PromptNode, StartNode, UserNode, AgentNode, DataMap, ToolNode
+    EndNode, Node, PromptNode, StartNode, UserNode, AgentNode, DataMap, ToolNode, DocProcNode, DecisionsNode
 )
 from ..types import (
     AgentNodeSpec, extract_node_spec, FlowContext, FlowEventType, FlowEvent, FlowSpec,
@@ -114,7 +116,7 @@ class Flow(Node):
                 # pydantic suppport nested comparison by default
 
                 schema.title = title
-                
+
                 if schema == existing_schema:
                     return existing_schema
                 # we need to do a deep compare
@@ -354,6 +356,7 @@ class Flow(Node):
               name: str, 
               agent: str, 
               display_name: str|None=None,
+              title: str | None = None,
               message: str | None = "Follow the agent instructions.",
               description: str | None = None,
               input_schema: type[BaseModel]|None = None, 
@@ -371,6 +374,7 @@ class Flow(Node):
             display_name=display_name,
             description=description,
             agent=agent,
+            title=title,
             message=message,
             guidelines=guidelines,
             input_schema=_get_tool_request_body(input_schema_obj),
@@ -428,6 +432,91 @@ class Flow(Node):
         # add the node to the list of node
         node = self._add_node(node)
         return cast(PromptNode, node)
+    
+    def decisions(self, 
+            name: str, 
+            display_name: str|None=None,
+            rules: list[DecisionsRule] | None = None,
+            default_actions: dict[str, Any] = None,
+            locale: str | None = None,
+            description: str | None = None,
+            input_schema: type[BaseModel]|None = None, 
+            output_schema: type[BaseModel]|None=None,
+            input_map: DataMap = None) -> PromptNode:
+
+        if name is None:
+            raise ValueError("name must be provided.")
+        
+        if rules is None:
+            raise ValueError("rules must be specified.")
+
+         # create input spec
+        input_schema_obj = _get_json_schema_obj(parameter_name = "input", type_def = input_schema)
+        output_schema_obj = _get_json_schema_obj("output", output_schema)
+
+        # Create the tool spec
+        task_spec = DecisionsNodeSpec(
+            name=name,
+            display_name=display_name if display_name is not None else name,
+            description=description,
+            rules=rules,
+            default_actions=default_actions,
+            locale=locale,
+            input_schema=_get_tool_request_body(input_schema_obj),
+            output_schema=_get_tool_response_body(output_schema_obj),
+            output_schema_object = output_schema_obj
+        )
+
+        node = DecisionsNode(spec=task_spec)
+        # setup input map
+        if input_map:
+            node.input_map = self._get_data_map(input_map)
+        
+        # add the node to the list of node
+        node = self._add_node(node)
+        return cast(DecisionsNode, node)
+    
+    def docproc(self, 
+            name: str, 
+            task: str,
+            display_name: str|None=None,
+            description: str | None = None,
+            input_map: DataMap = None) -> DocProcNode:
+
+        if name is None :
+            raise ValueError("name must be provided.")
+        
+        if task is None:
+            raise ValueError("task must be provided.")
+        
+        output_schema_dict = { 
+            "text_extraction" : TextExtractionResponse
+        }
+         # create input spec
+        input_schema_obj = _get_json_schema_obj(parameter_name = "input", type_def = File)
+        output_schema_obj = _get_json_schema_obj("output", output_schema_dict[task])
+        if "$defs" in output_schema_obj.model_extra:
+            output_schema_obj.model_extra.pop("$defs")
+        # Create the docproc spec
+        task_spec = DocProcSpec(
+            name=name,
+            display_name=display_name if display_name is not None else name,
+            description=description,
+            input_schema=_get_tool_request_body(input_schema_obj),
+            output_schema=_get_tool_response_body(output_schema_obj),
+            output_schema_object = output_schema_obj,
+            task=task
+        )
+
+        node = DocProcNode(spec=task_spec)
+        # setup input map
+        if input_map:
+            node.input_map = self._get_data_map(input_map)
+        
+        # add the node to the list of node
+        node = self._add_node(node)
+        return cast(DocProcNode, node)
+    
 
     def node_exists(self, node: Union[str, Node]):
      
@@ -921,7 +1010,8 @@ class FlowFactory(BaseModel):
                     description: str|None=None,
                     initiators: Sequence[str]|None=None,
                     input_schema: type[BaseModel]|None=None,
-                    output_schema: type[BaseModel]|None=None) -> Flow:
+                    output_schema: type[BaseModel]|None=None,
+                    schedulable: bool=False) -> Flow:
         if isinstance(name, Callable):
             flow_spec = getattr(name, "__flow_spec__", None)
             if not flow_spec:
@@ -942,7 +1032,8 @@ class FlowFactory(BaseModel):
             initiators=initiators,
             input_schema=_get_tool_request_body(input_schema_obj),
             output_schema=_get_tool_response_body(output_schema_obj),
-            output_schema_object = output_schema_obj
+            output_schema_object = output_schema_obj,
+            schedulable=schedulable,
         )
 
         return Flow(spec = flow_spec)
@@ -1228,10 +1319,12 @@ class UserFlow(Flow):
               kind: UserFieldKind = UserFieldKind.Text,
               display_name: str | None = None,
               description: str | None = None,
-              owners: list[str] = [],
               default: Any | None = None,
-              text: str = None,
+              text: str = None, # The text used to ask question to the user, e.g. 'what is your name?'
               option: UserFieldOption | None = None,
+              is_list: bool = False,
+              min: Any | None = None,
+              max: Any | None = None,
               input_map: DataMap = None,
               custom: dict[str, Any] = {}) -> UserNode:
         '''create a node in the flow'''
@@ -1246,20 +1339,42 @@ class UserFlow(Flow):
         schema_obj.properties = {}
         schema_obj.properties[name] = UserFieldKind.convert_kind_to_schema_property(kind, name, description, default, option, custom)
 
-        return self.user(name, 
-                         display_name=display_name,
-                         description=description,
-                         owners=owners,
-                         text=text,
-                         output_schema=schema_obj,
-                         input_map=input_map)
+        task_spec = UserNodeSpec(
+            name=name,
+            display_name=display_name,
+            description=description,
+            owners=[CURRENT_USER],
+            input_schema=_get_tool_request_body(schema_obj),
+            output_schema=_get_tool_response_body(schema_obj),
+            text=text,
+            output_schema_object = schema_obj
+        )
+
+        node = UserNode(spec = task_spec)
+        node.field(name = name,
+                   kind = kind,
+                   display_name = display_name,
+                   description = description,
+                   default = default,
+                   text = text,
+                   option = option,
+                   is_list = is_list,
+                   min = min,
+                   max = max,
+                   custom = custom)
+
+        # setup input map
+        if input_map:
+            node.input_map = self._get_data_map(input_map)
+
+        node = self._add_node(node)
+        return cast(UserNode, node)
 
     def user(
         self,
         name: str | None = None,
         display_name: str | None = None,
         description: str | None = None,
-        owners: list[str] = [],
         text: str | None = None,
         output_schema: type[BaseModel] | JsonSchemaObject| None = None,
         input_map: DataMap = None,
@@ -1273,16 +1388,12 @@ class UserFlow(Flow):
         # input and output is always the same in an user node
         output_schema_obj = output_schema_obj
 
-        # identify owner
-        if not owners:
-            owners = [ANY_USER]
-
         # Create the tool spec
         task_spec = UserNodeSpec(
             name=name,
             display_name=display_name,
             description=description,
-            owners=owners,
+            owners=[CURRENT_USER],
             input_schema=_get_tool_request_body(output_schema_obj),
             output_schema=_get_tool_response_body(output_schema_obj),
             text=text,

ibm_watsonx_orchestrate/flow_builder/node.py

@@ -5,7 +5,7 @@ import uuid
 import yaml
 from pydantic import BaseModel, Field, SerializeAsAny
 
-from .types import EndNodeSpec, NodeSpec, AgentNodeSpec, PromptNodeSpec, StartNodeSpec, ToolNodeSpec, UserFieldKind, UserFieldOption, UserNodeSpec
+from .types import EndNodeSpec, NodeSpec, AgentNodeSpec, PromptNodeSpec, StartNodeSpec, ToolNodeSpec, UserFieldKind, UserFieldOption, UserNodeSpec, DocProcSpec, DecisionsNodeSpec
 from .data_map import DataMap
 
 class Node(BaseModel):
@@ -78,6 +78,8 @@ class UserNode(Node):
               description: str | None = None,
               default: Any | None = None,
               option: UserFieldOption | None = None,
+              min: Any | None = None,
+              max: Any | None = None,
               is_list: bool = False,
               custom: dict[str, Any] | None = None,
               widget: str | None = None):
@@ -88,6 +90,8 @@ class UserNode(Node):
                               description=description,
                               default=default,
                               option=option,
+                              min=min,
+                              max=max,
                               is_list=is_list,
                               custom=custom,
                               widget=widget)
@@ -105,7 +109,20 @@ class PromptNode(Node):
 
     def get_spec(self) -> PromptNodeSpec:
         return cast(PromptNodeSpec, self.spec)
+    
+class DocProcNode(Node):
+    def __repr__(self):
+        return f"DocProcNode(name='{self.spec.name}', description='{self.spec.description}')"
 
+    def get_spec(self) -> DocProcSpec:
+        return cast(DocProcSpec, self.spec)
+class DecisionsNode(Node):
+    def __repr__(self):
+        return f"DecisionsNode(name='{self.spec.name}', description='{self.spec.description}')"
+
+    def get_spec(self) -> DecisionsNodeSpec:
+        return cast(DecisionsNodeSpec, self.spec)
+    
 class NodeInstance(BaseModel):
     node: Node
     id: str # unique id of this task instance

ibm_watsonx_orchestrate/flow_builder/types.py

@@ -1,5 +1,7 @@
 from dataclasses import dataclass
-from enum import Enum
+from enum import Enum, StrEnum, auto
+from datetime import date
+import numbers
 import inspect
 import logging
 from typing import (
@@ -7,7 +9,6 @@ from typing import (
 )
 
 import docstring_parser
-from munch import Munch
 from pydantic import BaseModel, Field
 
 from langchain_core.tools.base import create_schema_from_function
@@ -116,7 +117,7 @@ def _to_json_from_output_schema(schema: Union[ToolResponseBody, SchemaRef]) -> d
     return model_spec
 
 class NodeSpec(BaseModel):
-    kind: Literal["node", "tool", "user", "agent", "flow", "start", "decisions", "prompt", "branch", "wait", "foreach", "loop", "userflow", "end"] = "node"
+    kind: Literal["node", "tool", "user", "agent", "flow", "start", "decisions", "prompt", "branch", "wait", "foreach", "loop", "userflow", "end", "docproc" ] = "node"
     name: str
     display_name: str | None = None
     description: str | None = None
@@ -162,6 +163,24 @@ class NodeSpec(BaseModel):
 
         return model_spec
 
+class DocProcTask(StrEnum):
+    '''
+    Possible names for the Document processing task parameter
+    '''
+    text_extraction = auto()
+
+class DocProcSpec(NodeSpec):
+    task: DocProcTask = Field(description='The document processing operation name', default=DocProcTask.text_extraction)
+    
+    def __init__(self, **data):
+        super().__init__(**data)
+        self.kind = "docproc"
+
+    def to_json(self) -> dict[str, Any]:
+        model_spec = super().to_json()
+        model_spec["task"] = self.task
+        return model_spec
+    
 class StartNodeSpec(NodeSpec):
     def __init__(self, **data):
         super().__init__(**data)
@@ -296,6 +315,8 @@ class UserField(BaseModel):
     description: str | None = None
     default: Any | None = None
     option: UserFieldOption | None = None
+    min: Any | None = None,
+    max: Any | None = None,
     is_list: bool = False
     custom: dict[str, Any] | None = None
     widget: str | None = None
@@ -314,6 +335,10 @@ class UserField(BaseModel):
             model_spec["description"] = self.description
         if self.default:
             model_spec["default"] = self.default
+        if self.min:
+            model_spec["min"] = self.min
+        if self.max:
+            model_spec["min"] = self.max
         if self.is_list:
             model_spec["is_list"] = self.is_list
         if self.option:
@@ -356,7 +381,10 @@ class UserNodeSpec(NodeSpec):
               display_name: str | None = None, 
               description: str | None = None, 
               default: Any | None = None, 
-              option: list[str] | None = None, is_list: bool = False,
+              option: list[str] | None = None, 
+              min: Any | None = None,
+              max: Any | None = None,
+              is_list: bool = False,
               custom: dict[str, Any] | None = None,
               widget: str | None = None):
         userfield = UserField(name=name, 
@@ -366,6 +394,8 @@ class UserNodeSpec(NodeSpec):
                               description=description, 
                               default=default, 
                               option=option, 
+                              min=min,
+                              max=max,
                               is_list=is_list,
                               custom=custom,
                               widget=widget)
@@ -402,6 +432,8 @@ class UserNodeSpec(NodeSpec):
                                              default=prop_schema.default,
                                              option=self.setup_field_options(prop_schema.title, prop_schema.enum),
                                              is_list=prop_schema.type == "array",
+                                             min=prop_schema.minimum,
+                                             max=prop_schema.maximum,
                                              custom=prop_schema.model_extra))
 
     def setup_field_options(self, name: str, enums: List[str]) -> UserFieldOption:
@@ -415,6 +447,7 @@ class UserNodeSpec(NodeSpec):
 
 class AgentNodeSpec(ToolNodeSpec):
     message: str | None = Field(default=None, description="The instructions for the task.")
+    title: str | None = Field(default=None, description="The title of the message.")
     guidelines: str | None = Field(default=None, description="The guidelines for the task.")
     agent: str
 
@@ -430,6 +463,8 @@ class AgentNodeSpec(ToolNodeSpec):
             model_spec["guidelines"] = self.guidelines
         if self.agent:
             model_spec["agent"] = self.agent
+        if self.title:
+            model_spec["title"] = self.title
         return model_spec
 
 class PromptLLMParameters(BaseModel):
@@ -479,6 +514,7 @@ class PromptNodeSpec(NodeSpec):
             model_spec["llm_parameters"] = self.llm_parameters.to_json()
 
         return model_spec
+    
 
 class Expression(BaseModel):
     '''An expression could return a boolean or a value'''
@@ -553,10 +589,9 @@ class WaitNodeSpec(FlowControlNodeSpec):
         return my_dict
 
 class FlowSpec(NodeSpec):
- 
-
     # who can initiate the flow
     initiators: Sequence[str] = [ANY_USER]
+    schedulable: bool = False
 
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
@@ -566,6 +601,8 @@ class FlowSpec(NodeSpec):
         model_spec = super().to_json()
         if self.initiators:
             model_spec["initiators"] = self.initiators
+        
+        model_spec["schedulable"] = self.schedulable
 
         return model_spec
 
@@ -631,11 +668,11 @@ class TaskData(NamedTuple):
 
 class TaskEventType(Enum):
  
-    ON_TASK_WAIT = "on_task_wait" # the task is waiting for inputs before proceeding
-    ON_TASK_START = "on_task_start"
-    ON_TASK_END = "on_task_end"
-    ON_TASK_STREAM = "on_task_stream"
-    ON_TASK_ERROR = "on_task_error"
+    ON_TASK_WAIT = "task:on_task_wait" # the task is waiting for inputs before proceeding
+    ON_TASK_START = "task:on_task_start"
+    ON_TASK_END = "task:on_task_end"
+    ON_TASK_STREAM = "task:on_task_stream"
+    ON_TASK_ERROR = "task:on_task_error"
 
 class FlowData(BaseModel):
     '''This class represents the data that is passed between tasks in a flow.'''
@@ -667,9 +704,9 @@ class FlowContext(BaseModel):
     
 class FlowEventType(Enum):
  
-    ON_FLOW_START = "on_flow_start"
-    ON_FLOW_END = "on_flow_end"
-    ON_FLOW_ERROR = "on_flow_error"
+    ON_FLOW_START = "flow:on_flow_start"
+    ON_FLOW_END = "flow:on_flow_end"
+    ON_FLOW_ERROR = "flow:on_flow_error"
 
 
 @dataclass
@@ -691,9 +728,227 @@ class Assignment(BaseModel):
             e.g. "node.input.name" or "=f'{node.output.name}_{node.output.id}'"
 
     '''
-    target: str
-    source: str
+    target_variable: str
+    value_expression: str | None = None
+    has_no_value: bool = False
+    default_value: Any | None = None
+    metadata: dict = Field(default_factory=dict[str, Any])
+
+class LanguageCode(StrEnum):
+    '''
+    The ISO-639 language codes understood by Document Processing functions.
+    A special 'en_hw' code is used to enable an English handwritten model.
+    '''
+    en = auto()
+    fr = auto()
+    en_hw = auto()
+
+class File(BaseModel):
+    '''
+    This class represents the input of a Document processing task. 
+
+    Attributes:
+        document_ref (bytes|str): This is either a URL to the location of the document bytes or an ID that we use to resolve the location of the document
+        language (LanguageCode): Optional language code used when processing the input document
+    '''
+    # This is declared as bytes but the runtime will understand if a URL is send in as input.
+    # We need to use bytes here for Chat-with-doc to recognize the input as a File.
+    document_ref: bytes | str = Field(
+        description="Either an ID or a URL identifying the document to be used.", 
+        title='Document reference', 
+        default=None, 
+        json_schema_extra={"format": "binary"})
+    language: Optional[LanguageCode] = Field(
+        description='Optional language code of the document, defaults to "en"', 
+        title='Document language code', 
+        default=LanguageCode.en)
+
+class TextExtraction(BaseModel):
+    '''
+    This class represents the output generated by a "text_extraction" document processing (docproc) operation.
+    Attributes:
+        text (str): the text extracted from the input document.
+    '''
+    text: str = Field(description='The text extracted from the input document', title='Text extraction')
+
+class TextExtractionResponse(BaseModel):
+    '''
+    The text extraction operation response.
+    Attributes:
+        output (TextExtraction): a wrapper for the text extraction response
+    '''
+    output: TextExtraction = Field(description='The text extraction response')
+
+
+class DecisionsCondition(BaseModel):
+    _condition: str | None = None
+
+    def greater_than(self, value: Union[numbers.Number, date, str]) -> Self:
+        self._check_type_is_number_or_date_or_str(value)
+        self._condition = f"> {self._format_value(value)}"
+        return self
+
+    def greater_than_or_equal(self, value: Union[numbers.Number, date, str]) -> Self:
+        self._check_type_is_number_or_date_or_str(value)
+        self._condition = f">= {self._format_value(value)}"
+        return self
+
+    def less_than(self, value: Union[numbers.Number, date, str]) -> Self:
+        self._check_type_is_number_or_date_or_str(value)
+        self._condition = f"< {self._format_value(value)}"
+        return self
+
+    def less_than_or_equal(self, value: Union[numbers.Number, date, str]) -> Self:
+        self._check_type_is_number_or_date_or_str(value)
+        self._condition = f"<= {self._format_value(value)}"
+        return self
+
+    def equal(self, value: Union[numbers.Number, date, str]) -> Self:
+        self._check_type_is_number_or_date_or_str(value)
+        self._condition = f"== {self._format_value(value)}"
+        return self
+
+    def not_equal(self, value: Union[numbers.Number, date, str]) -> Self:
+        self._check_type_is_number_or_date_or_str(value)
+        self._condition = f"== {self._format_value(value)}"
+        return self
+    
+    def contains(self, value: str) -> Self:
+        self._check_type_is_str(value)
+        self._condition = f"contains {self._format_value(value)}"
+        return self
+
+    def not_contains(self, value: str) -> Self:
+        self._check_type_is_str(value)
+        self._condition = f"doesNotContain {self._format_value(value)}"
+        return self
+
+    def is_in(self, value: str) -> Self:
+        self._check_type_is_str(value)
+        self._condition = f"in {self._format_value(value)}"
+        return self
+
+    def is_not_in(self, value: str) -> Self:
+        self._check_type_is_str(value)
+        self._condition = f"notIn {self._format_value(value)}"
+        return self
+
+    def startswith(self, value: str) -> Self:
+        self._check_type_is_str(value)
+        self._condition = f"startsWith {self._format_value(value)}"
+        return self
+
+    def endswith(self, value: str) -> Self:
+        self._check_type_is_str(value)
+        self._condition = f"endsWith {self._format_value(value)}"
+        return self
+
+
+    def in_range(self, startValue: Union[numbers.Number, date], endValue: Union[numbers.Number, date], 
+                 startsInclusive: bool = False, endsInclusive: bool = False) -> Self:
+        self._check_type_is_number_or_date_or_str(startValue)
+        self._check_type_is_number_or_date_or_str(endValue)
+        if type(startValue) is not type(endValue):
+            raise TypeError("startValue and endValue must be of the same type")
+        start_op = "[" if startsInclusive else "("    # [ is inclusive, ( is exclusive
+        end_op =  "]" if endsInclusive else ")" 
+        self._condition = f"{start_op}{self._format_value(startValue)}:{self._format_value(endValue)}{end_op}"
+        return self
+
+    def _check_type_is_number_or_date(self, value: Union[numbers.Number, date]):
+        if not isinstance(value, (numbers.Number, date)):
+            raise TypeError("Value must be a number or a date")
+
+    def _check_type_is_number_or_date_or_str(self, value: Union[numbers.Number, date, str]):
+        if not isinstance(value, (numbers.Number, date, str)):
+            raise TypeError("Value must be a number or a date or a string")
+        
+    def _check_type_is_str(self, value: str):
+        if not isinstance(value, str):
+            raise TypeError("Value must be a string")
+    
+    @staticmethod
+    def _format_value(value: Union[numbers.Number, date, str]):
+        if isinstance(value, numbers.Number):
+            return f"{value}"
+        if isinstance(value, date):
+            return f"\"{value.strftime('%B %d, %Y')}\""
+        return f"\"{value}\""
+    
+    def condition(self):
+        return self._condition
+
+
+
+class DecisionsRule(BaseModel):
+    '''
+    A set of decisions rules.
+    '''
+    _conditions: dict[str, str]
+    _actions: dict[str, Union[numbers.Number, str]]
+
+    def __init__(self, **data):
+        super().__init__(**data)
+        self._conditions = {}
+        self._actions = {}
+
+    def condition(self, key: str, cond: DecisionsCondition) -> Self:
+        self._conditions[key] = cond.condition()
+        return self
     
+    def action(self, key: str, value: Union[numbers.Number, date, str]) -> Self:
+        if isinstance(value, date):
+            self._actions[key] = value.strftime("%B %d, %Y")
+            return self
+        self._actions[key] = value
+        return self
+
+    def to_json(self) -> dict[str, Any]:
+        '''
+        Serialize the rules into JSON object
+        '''
+        model_spec = {}
+        if self._conditions:
+            model_spec["conditions"] = self._conditions
+        if self._actions:
+            model_spec["actions"] = self._actions
+        return model_spec
+
+
+class DecisionsNodeSpec(NodeSpec):
+    '''
+    Node specification for Decision Table
+    '''
+    locale: str | None = None
+    rules: list[DecisionsRule]
+    default_actions: dict[str, Union[int, float, complex, str]] | None
+
+    def __init__(self, **data):
+        super().__init__(**data)
+        self.kind = "decisions"
+
+    def default_action(self, key: str, value: Union[int, float, complex, date, str]) -> Self:
+        '''
+        create a new default action
+        '''
+        if isinstance(value, date):
+            self.default_actions[key] = value.strftime("%B %d, %Y")
+            return self
+        self.default_actions[key] = value
+        return self
+
+    def to_json(self) -> dict[str, Any]:
+        model_spec = super().to_json()
+        if self.locale:
+            model_spec["locale"] = self.locale
+        if self.rules:
+            model_spec["rules"] = [rule.to_json() for rule in self.rules]
+        if self.default_actions:
+            model_spec["default_actions"] = self.default_actions
+
+        return model_spec
+
+
 def extract_node_spec(
         fn: Callable | PythonTool,
         name: Optional[str] = None,