ibm-watsonx-orchestrate 1.10.0b0__py3-none-any.whl → 1.11.0b0__py3-none-any.whl

ibm_watsonx_orchestrate/flow_builder/flows/flow.py

@@ -25,13 +25,13 @@ from ibm_watsonx_orchestrate.client.tools.tool_client import ToolClient
 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, 
+    DocProcKVPSchema, Assignment, Conditions, EndNodeSpec, Expression, ForeachPolicy, ForeachSpec, LoopSpec, BranchNodeSpec, MatchPolicy, NodeIdCondition, PlainTextReadingOrder, PromptExample, PromptLLMParameters, PromptNodeSpec, TimerNodeSpec,
     StartNodeSpec, ToolSpec, JsonSchemaObject, ToolRequestBody, ToolResponseBody, UserFieldKind, UserFieldOption, UserFlowSpec, UserNodeSpec, WaitPolicy,
-    DocProcSpec, TextExtractionResponse, DocProcInput, DecisionsNodeSpec, DecisionsRule, DocExtSpec, File
+    DocProcSpec, TextExtractionResponse, DocProcInput, DecisionsNodeSpec, DecisionsRule, DocExtSpec, File, DocumentClassificationResponse, DocClassifierSpec, DocumentProcessingCommonInput
 )
 from .constants import CURRENT_USER, START, END, ANY_USER
 from ..node import (
-    EndNode, Node, PromptNode, StartNode, UserNode, AgentNode, DataMap, ToolNode, DocProcNode, DecisionsNode, DocExtNode
+    EndNode, Node, PromptNode, StartNode, UserNode, AgentNode, DataMap, ToolNode, DocProcNode, DecisionsNode, DocExtNode, DocClassifierNode
 )
 from ..types import (
     AgentNodeSpec, extract_node_spec, FlowContext, FlowEventType, FlowEvent, FlowSpec,
@@ -64,7 +64,7 @@ class FlowEdge(BaseModel):
 
 class Flow(Node):
     '''Flow represents a flow that will be run by wxO Flow engine.'''
-    output_map: DataMap | None = None
+    output_map: dict[str, DataMap] | None = None
     nodes: dict[str, SerializeAsAny[Node]] = {}
     edges: List[FlowEdge] = []
     schemas: dict[str, JsonSchemaObject] = {}
@@ -401,6 +401,7 @@ class Flow(Node):
             display_name: str|None=None,
             system_prompt: str | list[str] | None = None,
             user_prompt: str | list[str] | None = None,
+            prompt_examples: list[PromptExample] | None = None,
             llm: str | None = None,
             llm_parameters: PromptLLMParameters | None = None,
             description: str | None = None,
@@ -422,6 +423,7 @@ class Flow(Node):
             description=description,
             system_prompt=system_prompt,
             user_prompt=user_prompt,
+            prompt_examples=prompt_examples,
             llm=llm,
             llm_parameters=llm_parameters,
             input_schema=_get_tool_request_body(input_schema_obj),
@@ -438,23 +440,95 @@ class Flow(Node):
         node = self._add_node(node)
         return cast(PromptNode, node)
     
+    def docclassfier(self, 
+            name: str, 
+            llm : str = "watsonx/meta-llama/llama-3-2-90b-vision-instruct",
+            version: str = "TIP",
+            display_name: str| None = None,
+            classes: type[BaseModel]| None = None, 
+            description: str | None = None,
+            min_confidence: float = 0.0,
+            input_map: DataMap = None) -> DocClassifierNode:
+        
+        if name is None :
+            raise ValueError("name must be provided.")
+        
+        doc_classifier_config = DocClassifierNode.generate_config(llm=llm, min_confidence=min_confidence,input_classes=classes)
+
+        input_schema_obj = _get_json_schema_obj(parameter_name = "input", type_def = DocumentProcessingCommonInput)
+        output_schema_obj = _get_json_schema_obj(parameter_name = "output", type_def = DocumentClassificationResponse)
+        
+        if "$defs" in output_schema_obj.model_extra:
+            output_schema_obj.model_extra.pop("$defs")
+        # Create the docclassifier spec
+        task_spec = DocClassifierSpec(
+            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,
+            config=doc_classifier_config,
+            version=version
+        )
+        node = DocClassifierNode(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(DocClassifierNode, node)
+        
+        
+    def timer(self,
+          name: str,
+          delay: int,
+          display_name: str | None = None,
+          description: str | None = None,
+          input_map: DataMap = None) -> Node:
+
+        if name is None:
+            raise ValueError("name must be provided.")
+        if delay < 0:
+            raise ValueError("delay must be non-negative.")
+
+        timer_spec = TimerNodeSpec(
+            name=name,
+            display_name=display_name if display_name is not None else name,
+            description=description,
+            delay=delay
+        )
+
+        node = Node(spec=timer_spec)
+
+        if input_map:
+            node.input_map = self._get_data_map(input_map)
+
+        node = self._add_node(node)
+        return node
+
+    
     def docext(self, 
             name: str, 
-            llm : str = "meta-llama/llama-3-2-11b-vision-instruct",
+            llm : str = "watsonx/meta-llama/llama-3-2-90b-vision-instruct",
             version: str = "TIP",
             display_name: str| None = None,
-            input_entities: type[BaseModel]| None = None, 
+            fields: type[BaseModel]| None = None, 
             description: str | None = None,
-            input_map: DataMap = None) -> tuple[DocExtNode, type[BaseModel]]:
+            input_map: DataMap = None,
+            enable_hw: bool = False) -> tuple[DocExtNode, type[BaseModel]]:
         
         if name is None :
             raise ValueError("name must be provided.")
 
-        doc_ext_config = DocExtNode.generate_config(llm=llm, input_entites=input_entities)
+        doc_ext_config = DocExtNode.generate_config(llm=llm, fields=fields)
 
-        DocExtFieldValue = DocExtNode.generate_docext_field_value_model(input_entities=input_entities)
+        DocExtFieldValue = DocExtNode.generate_docext_field_value_model(fields=fields)
         
-        input_schema_obj = _get_json_schema_obj(parameter_name = "input", type_def = File)
+        input_schema_obj = _get_json_schema_obj(parameter_name = "input", type_def = DocumentProcessingCommonInput)
         output_schema_obj = _get_json_schema_obj("output", DocExtFieldValue)
 
         if "$defs" in output_schema_obj.model_extra:
@@ -469,7 +543,8 @@ class Flow(Node):
             output_schema=_get_tool_response_body(output_schema_obj),
             output_schema_object = output_schema_obj,
             config=doc_ext_config,
-            version=version
+            version=version,
+            enable_hw=enable_hw
         )
         node = DocExtNode(spec=task_spec)
 
@@ -528,9 +603,12 @@ class Flow(Node):
     def docproc(self, 
             name: str, 
             task: str,
+            plain_text_reading_order : PlainTextReadingOrder = PlainTextReadingOrder.block_structure,
             display_name: str|None=None,
             description: str | None = None,
-            input_map: DataMap = None) -> DocProcNode:
+            input_map: DataMap = None,
+            kvp_schemas: list[DocProcKVPSchema] = None,
+            enable_hw: bool = False) -> DocProcNode:
 
         if name is None :
             raise ValueError("name must be provided.")
@@ -552,7 +630,10 @@ class Flow(Node):
             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
+            task=task,
+            plain_text_reading_order=plain_text_reading_order,
+            enable_hw=enable_hw,
+            kvp_schemas=kvp_schemas
         )
 
         node = DocProcNode(spec=task_spec)
@@ -642,7 +723,7 @@ class Flow(Node):
         '''Create a single node flow with an automatic START and END node.'''
         return self.sequence(START, node, END)
 
-    def branch(self, evaluator: Union[Callable, Expression]) -> "Branch":
+    def branch(self, evaluator: Union[Callable, Expression, Conditions]) -> 'Branch':
         '''Create a BRANCH node'''
         e = evaluator
         if isinstance(evaluator, Callable):
@@ -656,11 +737,19 @@ class Flow(Node):
             # e = new_script_spec
         elif isinstance(evaluator, str):
             e = Expression(expression=evaluator)
+        elif isinstance(evaluator, list):
+            e = Conditions(conditions=evaluator)
 
         spec = BranchNodeSpec(name = "branch_" + str(self._next_sequence_id()), evaluator=e)
         branch_node = Branch(spec = spec, containing_flow=self)
         return cast(Branch, self._add_node(branch_node))
     
+    def conditions(self) -> 'Branch':
+        '''Create a Branch node with empty Conditions evaluator (if-else)'''
+        spec = BranchNodeSpec(name = "branch_" + str(self._next_sequence_id()), evaluator=Conditions(conditions=[]))
+        branch_conditions_node = Branch(spec = spec, containing_flow=self)
+        return cast(Branch, self._add_node(branch_conditions_node))
+    
     def wait_for(self, *args) -> "Wait":
         '''Wait for all incoming nodes to complete.'''
         raise ValueError("Not implemented yet.")
@@ -675,6 +764,77 @@ class Flow(Node):
             
         # return cast(Wait, self.node(wait_node))
             
+    def map_flow_output_with_variable(self, target_output_variable: str, variable: str, default_value: str = None) -> Self:
+        if self.output_map and "spec" in self.output_map:
+            maps = self.output_map["spec"].maps or []
+        else:
+            maps = []
+        
+        curr_map_metadata = {
+            "assignmentType": "variable"
+        }
+
+        target_variable = "flow.output." + target_output_variable
+        value_expression = "flow." + variable
+
+        if default_value:
+            maps.append(Assignment(target_variable=target_variable, value_expression=value_expression, default_value=default_value, metadata=curr_map_metadata))
+        else:
+            maps.append(Assignment(target_variable=target_variable, value_expression=value_expression, metadata=curr_map_metadata))
+
+        flow_output_map_spec = DataMap(maps=maps)
+
+        if self.output_map and "spec" in self.output_map:
+            self.output_map["spec"] = flow_output_map_spec
+        else:
+            self.output_map = {"spec": flow_output_map_spec}
+        return self
+    
+    def map_output(self, output_variable: str, expression: str, default_value: str = None) -> Self:
+        if self.output_map and "spec" in self.output_map:
+            maps = self.output_map["spec"].maps or []
+        else:
+            maps = []
+        
+        curr_map_metadata = {
+            "assignmentType": "pyExpression"
+        }
+
+        target_variable = "flow.output." + output_variable
+        value_expression = expression
+
+        if default_value:
+            maps.append(Assignment(target_variable=target_variable, value_expression=value_expression, default_value=default_value, metadata=curr_map_metadata))
+        else:
+            maps.append(Assignment(target_variable=target_variable, value_expression=value_expression, metadata=curr_map_metadata))
+
+        flow_output_map_spec = DataMap(maps=maps)
+
+        if self.output_map and "spec" in self.output_map:
+            self.output_map["spec"] = flow_output_map_spec
+        else:
+            self.output_map = {"spec": flow_output_map_spec}
+        return self
+    
+    def map_flow_output_with_none(self, target_output_variable: str) -> Self:
+        if self.output_map and "spec" in self.output_map:
+            maps = self.output_map["spec"].maps or []
+        else:
+            maps = []
+        
+
+        target_variable = "flow.output." + target_output_variable
+
+        maps.append(Assignment(target_variable=target_variable, value_expression=None))
+
+        flow_output_map_spec = DataMap(maps=maps)
+
+        if self.output_map and "spec" in self.output_map:
+            self.output_map["spec"] = flow_output_map_spec
+        else:
+            self.output_map = {"spec": flow_output_map_spec}
+        return self
+
 
     def foreach(self, item_schema: type[BaseModel],
                 input_schema: type[BaseModel] |None=None,
@@ -735,8 +895,6 @@ class Flow(Node):
                  input_schema: type[BaseModel] |None=None,
                  output_schema: type[BaseModel] |None=None) -> "UserFlow": # return a UserFlow object
 
-        raise ValueError("userflow is NOT supported yet and it's interface will change.")
-
         output_schema_obj = _get_json_schema_obj("output", output_schema)
         input_schema_obj = _get_json_schema_obj("input", input_schema)
 
@@ -838,6 +996,11 @@ class Flow(Node):
         for key, value in self.metadata.items():
             metadata_dict[key] = value
         flow_dict["metadata"] = metadata_dict
+
+        if self.output_map and "spec" in self.output_map:
+            flow_dict["output_map"] = {
+                "spec": self.output_map["spec"].to_json()
+            }
         return flow_dict
 
     def _get_node_id(self, node: Union[str, Node]) -> str:
@@ -1147,6 +1310,27 @@ class Branch(FlowControl):
             raise ValueError("Cannot have custom label __default__. Use default() instead.")
 
         return self._add_case(label, node)
+    
+    def condition(self, to_node: Node, expression: str="", default: bool=False) -> Self:
+        '''
+        Add a condition to this branch node. 
+
+        Parameters:
+        expression (str): The expression of this condition.
+        to_node (Node): The node to go to when expression is evaluated to true.
+        default (bool): The condition is the default (else) case.
+        '''
+
+        node_id = self.containing_flow._get_node_id(to_node)
+        if default:
+            condition = NodeIdCondition(node_id=node_id, default=default)
+        else:
+            condition = NodeIdCondition(expression=expression, node_id=node_id, default=default)
+
+        self.spec.evaluator.conditions.append(condition)
+        self.containing_flow.edge(self, to_node)
+
+        return self
 
     def default(self, node: Node) -> Self:
         '''
@@ -1366,13 +1550,14 @@ class UserFlow(Flow):
               kind: UserFieldKind = UserFieldKind.Text,
               display_name: str | None = None,
               description: str | None = None,
+              direction: str | None = None,
               default: Any | None = 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,
+              input_map: DataMap | None= None,
               custom: dict[str, Any] = {}) -> UserNode:
         '''create a node in the flow'''
         # create a json schema object based on the single field
@@ -1404,6 +1589,8 @@ class UserFlow(Flow):
                    description = description,
                    default = default,
                    text = text,
+                   direction = direction,
+                   input_map = input_map,
                    option = option,
                    is_list = is_list,
                    min = min,

ibm_watsonx_orchestrate/flow_builder/node.py

@@ -6,14 +6,14 @@ import yaml
 from pydantic import BaseModel, Field, SerializeAsAny, create_model
 from enum import Enum
 
-from .types import EndNodeSpec, NodeSpec, AgentNodeSpec, PromptNodeSpec, StartNodeSpec, ToolNodeSpec, UserFieldKind, UserFieldOption, UserNodeSpec, DocProcSpec, \
-                    DocExtSpec, DocExtConfig, LanguageCode, DecisionsNodeSpec
+from .types import Assignment, DocExtConfigField, EndNodeSpec, NodeSpec, AgentNodeSpec, PromptNodeSpec, TimerNodeSpec, StartNodeSpec, ToolNodeSpec, UserFieldKind, UserFieldOption, UserNodeSpec, DocProcSpec, \
+                    DocExtSpec, DocExtConfig, DocClassifierSpec, DecisionsNodeSpec, DocClassifierConfig
 
 from .data_map import DataMap
 
 class Node(BaseModel):
     spec: SerializeAsAny[NodeSpec]
-    input_map: DataMap | None = None
+    input_map: dict[str, DataMap] | None = None
 
     def __call__(self, **kwargs):
         pass
@@ -40,10 +40,77 @@ class Node(BaseModel):
     def to_json(self) -> dict[str, Any]:
         model_spec = {}
         model_spec["spec"] = self.spec.to_json()
-        if self.input_map is not None:
-            model_spec['input_map'] = self.input_map.to_json()
+        if self.input_map is not None and "spec" in self.input_map:
+            model_spec['input_map'] = {
+                "spec": self.input_map["spec"].to_json()
+            }
 
         return model_spec
+    
+    def map_node_input_with_variable(self, target_input_variable: str, variable: str, default_value: str = None) -> None:
+        if self.input_map and "spec" in self.input_map:
+            maps = self.input_map["spec"].maps or []
+        else:
+            maps = []
+        
+        curr_map_metadata = {
+            "assignmentType": "variable"
+        }
+
+        target_variable = "self.input." + target_input_variable
+        value_expression = "flow." + variable
+
+        if default_value:
+            maps.append(Assignment(target_variable=target_variable, value_expression=value_expression, default_value=default_value, metadata=curr_map_metadata))
+        else:
+            maps.append(Assignment(target_variable=target_variable, value_expression=value_expression, metadata=curr_map_metadata))
+
+        node_input_map_spec = DataMap(maps=maps)
+        if self.input_map and "spec" in self.input_map:
+            self.input_map["spec"] = node_input_map_spec
+        else:
+            self.input_map = {"spec": node_input_map_spec}
+
+    def map_input(self, input_variable: str, expression: str, default_value: str = None) -> None:
+        if self.input_map and "spec" in self.input_map:
+            maps = self.input_map["spec"].maps or []
+        else:
+            maps = []
+        
+        curr_map_metadata = {
+            "assignmentType": "pyExpression"
+        }
+
+        target_variable = "self.input." + input_variable
+        value_expression = expression
+
+        if default_value:
+            maps.append(Assignment(target_variable=target_variable, value_expression=value_expression, default_value=default_value, metadata=curr_map_metadata))
+        else:
+            maps.append(Assignment(target_variable=target_variable, value_expression=value_expression, metadata=curr_map_metadata))
+
+        node_input_map_spec = DataMap(maps=maps)
+        if self.input_map and "spec" in self.input_map:
+            self.input_map["spec"] = node_input_map_spec
+        else:
+            self.input_map = {"spec": node_input_map_spec}
+
+    def map_node_input_with_none(self, target_input_variable: str) -> None:
+        if self.input_map and "spec" in self.input_map:
+            maps = self.input_map["spec"].maps or []
+        else:
+            maps = []
+        
+
+        target_variable = "self.input." + target_input_variable
+
+        maps.append(Assignment(target_variable=target_variable, value_expression=None))
+
+        node_input_map_spec = DataMap(maps=maps)
+        if self.input_map and "spec" in self.input_map:
+            self.input_map["spec"] = node_input_map_spec
+        else:
+            self.input_map = {"spec": node_input_map_spec}
 
 class StartNode(Node):
     def __repr__(self):
@@ -83,6 +150,8 @@ class UserNode(Node):
               option: UserFieldOption | None = None,
               min: Any | None = None,
               max: Any | None = None,
+              direction: str | None = None,
+              input_map: DataMap | None = None,
               is_list: bool = False,
               custom: dict[str, Any] | None = None,
               widget: str | None = None):
@@ -97,7 +166,9 @@ class UserNode(Node):
                               max=max,
                               is_list=is_list,
                               custom=custom,
-                              widget=widget)
+                              widget=widget,
+                              direction=direction,
+                              input_map=input_map)
 
 class AgentNode(Node):
     def __repr__(self):
@@ -120,6 +191,24 @@ class DocProcNode(Node):
     def get_spec(self) -> DocProcSpec:
         return cast(DocProcSpec, self.spec)
     
+class DocClassifierNode(Node):
+    def __repr__(self):
+        return f"DocClassifierNode(name='{self.spec.name}', description='{self.spec.description}')"
+
+    def get_spec(self) -> DocClassifierSpec:
+        return cast(DocClassifierSpec, self.spec)
+
+    @staticmethod
+    def generate_config(llm: str, input_classes: type[BaseModel], min_confidence: float) -> DocClassifierConfig:
+        return DocClassifierConfig(llm=llm, classes=input_classes.__dict__.values(), min_confidence=min_confidence)
+    
+class TimerNode(Node):
+    def __repr__(self):
+        return f"TimerNode(name='{self.spec.name}', description='{self.spec.description}')"
+
+    def get_spec(self) -> TimerNodeSpec:
+        return cast(TimerNodeSpec, self.spec)
+    
 class DocExtNode(Node):
     def __repr__(self):
         return f"DocExtNode(name='{self.spec.name}', description='{self.spec.description}')"
@@ -128,23 +217,29 @@ class DocExtNode(Node):
         return cast(DocExtSpec, self.spec)
     
     @staticmethod
-    def generate_config(llm: str, input_entites: type[BaseModel]) -> DocExtConfig:
-        entities = input_entites.__dict__.values()
-        return DocExtConfig(llm=llm, entities=entities)
+    def generate_config(llm: str, fields: type[BaseModel]) -> DocExtConfig:
+        return DocExtConfig(llm=llm, fields=fields.__dict__.values())
     
     @staticmethod
-    def generate_docext_field_value_model(input_entities: type[BaseModel]) -> type[BaseModel]:
+    def generate_docext_field_value_model(fields: type[BaseModel]) -> type[BaseModel]:
         create_field_value_description = lambda field_name: "Extracted value for " + field_name
+        field_definitions = {}
+
+        for name, value in fields.model_dump().items():
+            field_type = str  
+            field_kwargs = {
+                "title": value['name'],
+                "description": create_field_value_description(value['name']),
+                "type": value["type"] if value["type"] != "date" else "string"
+            }
+
+            # Add json_schema_extra if type is 'date'
+            if value["type"] == "date":
+                field_kwargs["json_schema_extra"] = {"format": "date"}
+
+            field_definitions[name] = (field_type, Field(**field_kwargs))
 
-        DocExtFieldValue = create_model(
-                "DocExtFieldValue",
-                    **{
-                        name: (str, Field(
-                            title=value['name'], 
-                            description=create_field_value_description(value['name']),
-                            )
-                        )
-                        for name, value in input_entities.model_dump().items()})
+        DocExtFieldValue = create_model("DocExtFieldValue", **field_definitions)
         return DocExtFieldValue
     
 class DecisionsNode(Node):