ibm-watsonx-orchestrate 1.5.1__py3-none-any.whl → 1.6.0__py3-none-any.whl

ibm_watsonx_orchestrate/docker/default.env

@@ -22,7 +22,6 @@ ASSISTANT_EMBEDDINGS_API_BASE=https://us-south.ml.cloud.ibm.com
 #ASSISTANT_EMBEDDINGS_SPACE_ID=#Will default to WATSONX_SPACE_ID if specified
 ROUTING_LLM_API_BASE=https://us-south.ml.cloud.ibm.com
 ROUTING_LLM_MODEL_ID=watsonx/ibm/granite-8b-unified-api-model-v2
-AUTHORIZATION_URL=https://iam.test.cloud.ibm.com/identity/token
 #ROUTING_LLM_API_KEY=#Will default to WATSONX_APIKEY if specified
 #ROUTING_LLM_SPACE_ID=#Will default to WATSONX_SPACE_ID if specified
 ASSISTANT_INDEX_CHUNK_SIZE=1000
@@ -53,44 +52,59 @@ EVENT_BROKER_TTL="-1"
 # See get_default_registry_env_vars_by_dev_edition_source() in src/ibm_watsonx_orchestrate/cli/commands/server/server_command.py for more details.
 REGISTRY_URL=
 
-SERVER_TAG=05-06-2025-v1
+
+SERVER_TAG=27-06-2025
 SERVER_REGISTRY=
 
-WORKER_TAG=03-06-2025
+WORKER_TAG=27-06-2025
 WORKER_REGISTRY=
 
-AI_GATEWAY_TAG=03-06-2025
+AI_GATEWAY_TAG=23-06-2025
 AI_GATEWAY_REGISTRY=
 
+AGENT_GATEWAY_TAG=23-06-2025
+AGENT_GATEWAY_REGISTRY=
+
 DB_REGISTRY=
 # If you build multiarch set all three of these to the same, we have a pr against main
 # to not have this separation, but we can merge it later
-DBTAG=07-06-2025
-AMDDBTAG=07-06-2025
-ARM64DBTAG=07-06-2025
+DBTAG=24-06-2025-v1
+AMDDBTAG=24-06-2025-v1
+ARM64DBTAG=24-06-2025-v1
 
 UI_REGISTRY=
-UITAG=03-06-2025
+UITAG=27-06-2025
 
 CM_REGISTRY=
-CM_TAG=13-05-2025
+CM_TAG=27-06-2025
 
-TRM_TAG=11-06-2025-1-5-0
+TRM_TAG=26-06-2025
 TRM_REGISTRY=
 
-TR_TAG=11-06-2025-1-5-0
+TR_TAG=26-06-2025
 TR_REGISTRY=
 
-BUILDER_TAG=03-06-2025
+BUILDER_TAG=26-06-2025-a
 BUILDER_REGISTRY=
 
-FLOW_RUNTIME_TAG=13-05-2025
+FLOW_RUNTIME_TAG=23-06-2025
 FLOW_RUMTIME_REGISTRY=
 
+
+AGENT_ANALYTICS_TAG=27-06-2025-v2
+AGENT_ANALYTICS_REGISTRY=
+
+JAEGER_PROXY_TAG=27-06-2025
+JAEGER_PROXY_REGISTRY=
+
+SOCKET_HANDLER_TAG=29-05-2025
+SOCKET_HANDLER_REGISTRY=
+
+
 # END -- IMAGE REGISTRIES AND TAGS
 
 TAVILY_API_KEY=dummy_tavily_api_key
-PREFERRED_MODELS=meta-llama/llama-3-1-8b-instruct,meta-llama/llama-3-1-70b-instruct,meta-llama/llama-3-405b-instruct,meta-llama/llama-3-2-1b-instruct,meta-llama/llama-3-2-3b-instruct,meta-llama/llama-3-2-11b-vision-instruct,meta-llama/llama-3-2-90b-vision-instruct,mistralai/mistral-large,mistralai/mixtral-8x7b-instruct-v01,ibm/granite-3-2b-instruct,ibm/granite-3-8b-instruct
+PREFERRED_MODELS=meta-llama/llama-3-2-90b-vision-instruct,meta-llama/llama-3-405b-instruct
 INCOMPATIBLE_MODELS=flan,embedding,cross-encoder,tinytimemixers
 #WATSONX_APIKEY= #Must define in .env
 WATSONX_URL=https://us-south.ml.cloud.ibm.com
@@ -119,7 +133,10 @@ TOOLS_RUNTIME_MANAGER_BASE_URL="http://tools-runtime-manager:8080"
 CONNECTION_SERVICE_BASE_URL="http://wxo-server-connection-manager:3001"
 AI_GATEWAY_BASE_URL="http://ai-gateway:8787/v1"
 AI_GATEWAY_ENABLED=True
-
+AGENT_GATEWAY_URI="http://wxo-agent-gateway:8989"
+DEFAULT_TENANT_ID=10000000-0000-0000-0000-000000000000
+ES_USERNAME=elastic
+ES_PASSWORD=changeme
 #To Prevent warnings
 VECTOR_STORE_PROVIDER=
 MILVUS_URI=
@@ -135,3 +152,11 @@ DB_CONN_LIFE=
 DB_MAX_IDLE_CONN=
 DB_MAX_CONN=
 SERVER_HOST=
+
+# Use your machine's local IP address for external async tool communication.
+CALLBACK_HOST_URL=
+
+AGENTOPS_API_KEY_AUTH_ENABLED=true
+AGENTOPS_API_KEY=qwertyuiop
+
+RUNTIME_MANAGER_API_KEY=example

ibm_watsonx_orchestrate/docker/proxy-config-single.yaml

@@ -0,0 +1,12 @@
+# Default tenant to use when tenant_id is not found
+default_tenant: "default"
+
+# Tenant configurations
+tenants:
+  # Default tenant on ES instance 1
+  default:
+    store_type: elasticsearch
+    hostname: "http://elasticsearch:9200"
+    username: "elastic"
+    password: "changeme"
+    index_prefix: "default"

ibm_watsonx_orchestrate/{experimental/flow_builder → flow_builder}/flows/__init__.py

@@ -1,8 +1,8 @@
 from .constants import START, END, RESERVED
 from ..types import FlowContext, TaskData, TaskEventType
-from ..node import UserNode, AgentNode, StartNode, EndNode, PromptNode
+from ..node import UserNode, AgentNode, StartNode, EndNode, PromptNode, ToolNode
 from .flow import Flow, CompiledFlow, FlowRun, FlowEvent, FlowEventType, FlowFactory, MatchPolicy, WaitPolicy, ForeachPolicy, Branch, Foreach, Loop
-from .decorators import user, flow_spec, flow
+from .decorators import flow
 from ..data_map import Assignment, DataMap
 
 
@@ -20,6 +20,7 @@ __all__ = [
     "StartNode",
     "EndNode",
     "PromptNode",
+    "ToolNode",
     "Assignment",
     "DataMap",
 

ibm_watsonx_orchestrate/flow_builder/flows/decorators.py

@@ -0,0 +1,77 @@
+"""
+A set of decorators to help define different Flow constructs.
+"""
+
+import asyncio
+from functools import wraps
+import logging
+import inspect
+from typing import Callable, Optional, Sequence
+from pydantic import BaseModel
+from ..types import extract_node_spec, UserNodeSpec, FlowSpec
+
+from .flow import FlowFactory, Flow
+
+logger = logging.getLogger(__name__)
+
+
+class FlowWrapper:
+    def __init__(self, func, a_model):
+        self.func = func
+        self.a_model = a_model
+        wraps(func)(self)  # Preserve metadata
+
+    def __call__(self, *args, **kwargs):
+        result = self.func(self.a_model)
+        if not isinstance(result, Flow):
+            raise ValueError("Return value must be of type Flow")
+        return result
+    
+def flow(*args, 
+         name: Optional[str]=None, 
+         display_name: Optional[str]=None,
+         description: str|None=None,
+         input_schema: type[BaseModel] | None = None,
+         output_schema: type[BaseModel] | None = None,
+         initiators: Sequence[str] = ()):
+    """Decorator to mark a function as a flow model builder."""
+
+    def decorator(func: Callable):
+        """
+        Decorator that takes a function as an argument and returns a wrapper function.
+        The wrapper function takes a single argument of type Flow and calls the original function with the created flow as an argument.
+        """
+
+        sig = inspect.signature(func)
+        if len(sig.parameters) != 1:
+            raise ValueError("Only one argument is allowed")
+        param = list(sig.parameters.values())[0]
+        if param.annotation != Flow:
+            raise ValueError("Argument must be of type Flow")
+        if sig.return_annotation != Flow:
+            raise ValueError("Return value must be of type Flow")
+        
+        node_spec = extract_node_spec(func, name, description)
+        a_model = FlowFactory.create_flow(
+                             name = node_spec.name,
+                             display_name = display_name,
+                             description = node_spec.description,
+                             input_schema = input_schema,
+                             output_schema = output_schema,
+                             initiators = initiators)
+
+        # logger.info("Creating flow model: %s", a_model.spec.name)
+
+        # @wraps(func)
+        # def wrapper(*args, **kwargs):
+        #     result = func(a_model)
+        #     if not isinstance(result, Flow):
+        #         raise ValueError("Return value must be of type Flow")
+        #     return result
+
+        return FlowWrapper(func, a_model)
+
+    if len(args) == 1 and callable(args[0]):
+        return decorator(args[0])
+    else:
+        return decorator

ibm_watsonx_orchestrate/{experimental/flow_builder → flow_builder}/flows/events.py

@@ -47,8 +47,13 @@ class StreamConsumer:
 def deserialize_flow_event(byte_data: bytes) -> FlowEvent:
     """Deserialize byte data into a FlowEvent object."""
     # Decode the byte data
-    decoded_data = byte_data[b'message'].decode('utf-8')
+    decoded_data = byte_data[b'data'].decode('utf-8') if b'data' in byte_data else None
+    if not decoded_data:
+        decoded_data = byte_data[b'message'].decode('utf-8') if b'message' in byte_data else None
     
+    if not decoded_data:
+        raise ValueError("No data in received event.")
+
     # Parse the JSON string into a dictionary
     parsed_data = json.loads(decoded_data)
     

ibm_watsonx_orchestrate/{experimental/flow_builder → flow_builder}/flows/flow.py

@@ -20,6 +20,7 @@ from typing_extensions import Self
 from pydantic import BaseModel, Field, SerializeAsAny
 import yaml
 from ibm_watsonx_orchestrate.agent_builder.tools.python_tool import PythonTool
+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 (
@@ -36,7 +37,7 @@ from ..types import (
 )
 
 from ..data_map import DataMap
-from ..utils import _get_json_schema_obj, get_valid_name, import_flow_model
+from ..utils import _get_json_schema_obj, get_valid_name, import_flow_model, _get_tool_request_body, _get_tool_response_body
 
 from .events import StreamConsumer
 
@@ -70,6 +71,7 @@ class Flow(Node):
     metadata: dict[str, str] = {}
     parent: Any = None
     _sequence_id: int = 0 # internal-id
+    _tool_client: ToolClient = None
 
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
@@ -77,6 +79,9 @@ class Flow(Node):
         # extract data schemas
         self._refactor_node_to_schemaref(self)
 
+        # get Tool Client
+        self._tool_client = instantiate_client(ToolClient)
+
     def _find_topmost_flow(self) -> Self:
         if self.parent:
             return self.parent._find_topmost_flow()
@@ -104,7 +109,27 @@ class Flow(Node):
         # if there is already a schema with the same name, return it
         if title:
             if title in top_flow.schemas:
-                return top_flow.schemas[title]
+                existing_schema = top_flow.schemas[title]
+                # we need a deep compare if the incoming schema and existing_schema is the same
+                # pydantic suppport nested comparison by default
+
+                schema.title = title
+                
+                if schema == existing_schema:
+                    return existing_schema
+                # we need to do a deep compare
+                incoming_model = schema.model_dump(exclude_none=True, exclude_unset=True)
+                existing_model = existing_schema.model_dump(exclude_none=True, exclude_unset=True)
+
+                # log the model
+                # logger.info(f"incoming_model: {incoming_model}")
+                # logger.info(f"existing_model: {existing_model}")
+
+                if incoming_model == existing_model:
+                    return existing_schema
+                
+                # else we need a new name, and create a new schema
+                title = title + "_" + str(self._next_sequence_id())
 
         # otherwise, create a deep copy of the schema, add it to the dictionary and return it
         if schema:
@@ -122,7 +147,7 @@ class Flow(Node):
                 elif schema.aliasName:
                     title = get_valid_name(schema.aliasName)
                 else:
-                    title = "bo_" + uuid.uuid4().hex
+                    title = "bo_" + str(self._next_sequence_id())
             
             if new_schema.type == "object":
             # iterate the properties and add schema recursively
@@ -193,7 +218,7 @@ class Flow(Node):
                             required= spec.output_schema.required)
                     spec.output_schema = self._add_schema_ref(json_obj, f"{spec.name}_output")
                 elif spec.output_schema.type == "array":
-                    if spec.output_schema.items.type == "object":
+                    if hasattr(spec.output_schema, "items") and hasattr(spec.output_schema.items, "type") and spec.output_schema.items.type == "object":
                         schema_ref = self._add_schema_ref(spec.output_schema.items)
                         spec.output_schema.items = JsonSchemaObjectRef(ref=f"{schema_ref.ref}")
 
@@ -265,23 +290,30 @@ class Flow(Node):
         
         if isinstance(tool, str):        
             name = name if name is not None and name != "" else tool
-            input_schema_obj = _get_json_schema_obj(parameter_name = "input", type_def = input_schema)
-            output_schema_obj = _get_json_schema_obj("output", output_schema)
+
+            if input_schema is None and output_schema is None:
+                # try to retrieve the schema from server
+                tool_specs: List[dict] = self._tool_client.get_draft_by_name(name)
+                if (tool_specs is None) or (len(tool_specs) == 0):
+                    raise ValueError(f"tool '{name}' not found")
+
+                # use the first spec
+                tool_spec: ToolSpec = ToolSpec.model_validate(tool_specs[0])
+                # just pick the first one that is found
+                if hasattr(tool_spec, "input_schema"):
+                    input_schema_obj = _get_json_schema_obj("input", tool_spec.input_schema, True)
+                if hasattr(tool_spec, "output_schema"):
+                    output_schema_obj = _get_json_schema_obj("output", tool_spec.output_schema)
+            else: 
+                input_schema_obj = _get_json_schema_obj("input", input_schema)
+                output_schema_obj = _get_json_schema_obj("output", output_schema)
 
             toolnode_spec = ToolNodeSpec(type = "tool",
                                      name = name,
                                      display_name = display_name,
                                      description = description,
-                                     input_schema=ToolRequestBody(
-                                        type=input_schema_obj.type,
-                                        properties=input_schema_obj.properties,
-                                        required=input_schema_obj.required,
-                                    ) if input_schema is not None else None,
-                                    output_schema=ToolResponseBody(
-                                        type=output_schema_obj.type,
-                                        properties=output_schema_obj.properties,
-                                        required=output_schema_obj.required
-                                    ) if output_schema is not None else None,
+                                     input_schema= _get_tool_request_body(input_schema_obj),
+                                     output_schema= _get_tool_response_body(output_schema_obj),
                                      output_schema_object = output_schema_obj,
                                      tool = tool)
 
@@ -341,16 +373,8 @@ class Flow(Node):
             agent=agent,
             message=message,
             guidelines=guidelines,
-            input_schema=ToolRequestBody(
-                type=input_schema_obj.type,
-                properties=input_schema_obj.properties,
-                required=input_schema_obj.required,
-            ),
-            output_schema=ToolResponseBody(
-                type=output_schema_obj.type,
-                properties=output_schema_obj.properties,
-                required=output_schema_obj.required
-            ),
+            input_schema=_get_tool_request_body(input_schema_obj),
+            output_schema=_get_tool_response_body(output_schema_obj),
             output_schema_object = output_schema_obj
         )
 
@@ -391,16 +415,8 @@ class Flow(Node):
             user_prompt=user_prompt,
             llm=llm,
             llm_parameters=llm_parameters,
-            input_schema=ToolRequestBody(
-                type=input_schema_obj.type,
-                properties=input_schema_obj.properties,
-                required=input_schema_obj.required,
-            ),
-            output_schema=ToolResponseBody(
-                type=output_schema_obj.type,
-                properties=output_schema_obj.properties,
-                required=output_schema_obj.required
-            ),
+            input_schema=_get_tool_request_body(input_schema_obj),
+            output_schema=_get_tool_response_body(output_schema_obj),
             output_schema_object = output_schema_obj
         )
 
@@ -543,21 +559,13 @@ class Flow(Node):
                 },
                 required = ["items"])
 
+        new_foreach_item_schema = self._add_schema(foreach_item_schema)
         spec = ForeachSpec(name = "foreach_" + str(self._next_sequence_id()),
-                           input_schema=ToolRequestBody(
-                                type=input_schema_obj.type,
-                                properties=input_schema_obj.properties,
-                                required=input_schema_obj.required,
-                            ),
-                            output_schema=ToolResponseBody(
-                                type=output_schema_obj.type,
-                                properties=output_schema_obj.properties,
-                                required=output_schema_obj.required
-                            ) if output_schema_obj is not None else None,
-                           item_schema = foreach_item_schema)
+                           input_schema=_get_tool_request_body(input_schema_obj),
+                           output_schema=_get_tool_response_body(output_schema_obj),
+                           item_schema = new_foreach_item_schema)
         foreach_obj = Foreach(spec = spec, parent = self)
         foreach_node = self._add_node(foreach_obj)
-        self._add_schema(foreach_item_schema)
 
         return cast(Flow, foreach_node)
 
@@ -580,16 +588,8 @@ class Flow(Node):
 
         loop_spec = LoopSpec(name = "loop_" + str(self._next_sequence_id()), 
                              evaluator = e, 
-                             input_schema=ToolRequestBody(
-                                type=input_schema_obj.type,
-                                properties=input_schema_obj.properties,
-                                required=input_schema_obj.required,
-                             ) if input_schema_obj is not None else None,
-                             output_schema=ToolResponseBody(
-                                type=output_schema_obj.type,
-                                properties=output_schema_obj.properties,
-                                required=output_schema_obj.required
-                             ) if output_schema_obj is not None else None)
+                             input_schema=_get_tool_request_body(input_schema_obj),
+                             output_schema=_get_tool_response_body(output_schema_obj))
         while_loop = Loop(spec = loop_spec, parent = self)
         while_node = self._add_node(while_loop)
         return cast(Loop, while_node)
@@ -599,23 +599,15 @@ class Flow(Node):
                  input_schema: type[BaseModel] |None=None,
                  output_schema: type[BaseModel] |None=None) -> "UserFlow": # return a UserFlow object
 
-        logger.warning("userflow is NOT working yet.")
+        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)
 
         spec = UserFlowSpec(name = "userflow_" + str(self._next_sequence_id()),
-                           input_schema=ToolRequestBody(
-                                type=input_schema_obj.type,
-                                properties=input_schema_obj.properties,
-                                required=input_schema_obj.required,
-                            ) if input_schema_obj is not None else None,
-                            output_schema=ToolResponseBody(
-                                type=output_schema_obj.type,
-                                properties=output_schema_obj.properties,
-                                required=output_schema_obj.required
-                            ) if output_schema_obj is not None else None,
-                           owners = owners)
+                            input_schema=_get_tool_request_body(input_schema_obj),
+                            output_schema=_get_tool_response_body(output_schema_obj),
+                            owners = owners)
         userflow_obj = UserFlow(spec = spec, parent = self)
         userflow_node = self._add_node(userflow_obj)
 
@@ -679,8 +671,9 @@ class Flow(Node):
         
         # Deploy flow to the engine
         model = self.to_json()
-        await import_flow_model(model)
+        tool_id = await import_flow_model(model)
 
+        compiled_flow.flow_id = tool_id
         compiled_flow.deployed = True
 
         return compiled_flow
@@ -746,6 +739,7 @@ class FlowRun(BaseModel):
     name: str | None = None
     id: str = None
     flow: Flow
+    deployed_flow_id: str
     status: FlowRunStatus = FlowRunStatus.NOT_STARTED
     output: Any = None
     error: Any = None
@@ -765,10 +759,12 @@ class FlowRun(BaseModel):
 
         # Start the flow
         client:TempusClient = instantiate_client(client=TempusClient)
-        ack = client.arun_flow(self.flow.spec.name,input_data)
+        logger.info(f"Launching flow instance...")
+        ack = client.arun_flow(self.deployed_flow_id,input_data)
         self.id=ack["instance_id"]
         self.name = f"{self.flow.spec.name}:{self.id}"
         self.status = FlowRunStatus.IN_PROGRESS
+        logger.info(f"Flow instance `{self.name}` started.")
 
         # Listen for events
         consumer = StreamConsumer(self.id)
@@ -781,6 +777,11 @@ class FlowRun(BaseModel):
             
             self._update_status(event)
 
+            if event.kind == FlowEventType.ON_FLOW_END:
+                logger.info(f"Flow instance `{self.name}` completed.")
+            elif event.kind == FlowEventType.ON_FLOW_ERROR:
+                logger.error(f"Flow instance `{self.name}` failed with error: {event.error}")
+
             yield event
     
     def _update_status(self, event:FlowEvent):
@@ -794,7 +795,7 @@ class FlowRun(BaseModel):
 
 
         if self.debug:
-            logger.debug(f"Flow instance `{self.name}` status change: `{self.status}`")
+            logger.debug(f"Flow instance `{self.name}` status change: `{self.status}`.  \nEvent: {event}")
 
     async def _arun(self, input_data: dict=None, **kwargs):
         
@@ -822,7 +823,7 @@ class FlowRun(BaseModel):
     def _on_flow_end(self, event:FlowEvent):
 
         self.status = FlowRunStatus.COMPLETED
-        self.output = event.context.data["output"] if "output" in event.context.data else None
+        self.output = event.context.data.output
 
         if self.debug:
             logger.debug(f"Flow run `{self.name}`: on_complete handler called. Output: {self.output}")
@@ -846,6 +847,7 @@ class FlowRun(BaseModel):
 class CompiledFlow(BaseModel):
     '''A compiled version of the flow'''
     flow: Flow
+    flow_id: str | None = None
     deployed: bool = False
     
     async def invoke(self, input_data:dict=None, on_flow_end_handler: Callable=None, on_flow_error_handler: Callable=None, debug:bool=False, **kwargs) -> FlowRun:
@@ -868,7 +870,7 @@ class CompiledFlow(BaseModel):
         if self.deployed is False:
             raise ValueError("Flow has not been deployed yet. Please deploy the flow before invoking it by using the Flow.compile_deploy() function.")
 
-        flow_run = FlowRun(flow=self.flow, on_flow_end_handler=on_flow_end_handler, on_flow_error_handler=on_flow_error_handler, debug=debug, **kwargs)
+        flow_run = FlowRun(flow=self.flow,  deployed_flow_id=self.flow_id, on_flow_end_handler=on_flow_end_handler, on_flow_error_handler=on_flow_error_handler, debug=debug, **kwargs)
         asyncio.create_task(flow_run._arun(input_data=input_data, **kwargs))
         return flow_run
     
@@ -890,7 +892,7 @@ class CompiledFlow(BaseModel):
         if self.deployed is False:
             raise ValueError("Flow has not been deployed yet. Please deploy the flow before invoking it by using the Flow.compile_deploy() function.")
         
-        flow_run = FlowRun(flow=self.flow, debug=debug)
+        flow_run = FlowRun(flow=self.flow, deployed_flow_id=self.flow_id, debug=debug)
         async for event in flow_run._arun_events(input_data=input_data, filters=filters):
             yield (event, flow_run)
     
@@ -938,16 +940,8 @@ class FlowFactory(BaseModel):
             display_name=display_name,
             description=description,
             initiators=initiators,
-            input_schema=ToolRequestBody(
-                type=input_schema_obj.type,
-                properties=input_schema_obj.properties,
-                required=input_schema_obj.required,
-            ) if input_schema_obj else None,
-            output_schema=ToolResponseBody(
-                type=output_schema_obj.type,
-                properties=output_schema_obj.properties,
-                required=output_schema_obj.required
-            ) if output_schema_obj else None,
+            input_schema=_get_tool_request_body(input_schema_obj),
+            output_schema=_get_tool_response_body(output_schema_obj),
             output_schema_object = output_schema_obj
         )
 
@@ -974,6 +968,9 @@ class Branch(FlowControl):
         Returns:
         Self: The current node.
         '''
+        if kind == MatchPolicy.ANY_MATCH:
+            raise ValueError("Branch with policy ANY_MATCH is not supported yet.")
+        
         self.spec.match_policy = kind
         return self
 
@@ -1286,12 +1283,8 @@ class UserFlow(Flow):
             display_name=display_name,
             description=description,
             owners=owners,
-            input_schema=None,
-            output_schema=ToolResponseBody(
-                type=output_schema_obj.type,
-                properties=output_schema_obj.properties,
-                required=output_schema_obj.required
-            ),
+            input_schema=_get_tool_request_body(output_schema_obj),
+            output_schema=_get_tool_response_body(output_schema_obj),
             text=text,
             output_schema_object = output_schema_obj
         )

ibm_watsonx_orchestrate/{experimental/flow_builder → flow_builder}/types.py

@@ -14,7 +14,7 @@ from langchain_core.tools.base import create_schema_from_function
 from langchain_core.utils.json_schema import dereference_refs
 
 from ibm_watsonx_orchestrate.agent_builder.tools import PythonTool
-from ibm_watsonx_orchestrate.experimental.flow_builder.flows.constants import ANY_USER
+from ibm_watsonx_orchestrate.flow_builder.flows.constants import ANY_USER
 from ibm_watsonx_orchestrate.agent_builder.tools.types import (
     ToolSpec, ToolRequestBody, ToolResponseBody, JsonSchemaObject
 )
@@ -29,9 +29,9 @@ class SchemaRef(BaseModel):
  
     ref: str = Field(description="The id of the schema to be used.", serialization_alias="$ref")
 
-def _assign_attribute(obj, attr_name, schema):
+def _assign_attribute(model_spec, attr_name, schema):
     if hasattr(schema, attr_name) and (getattr(schema, attr_name) is not None):
-        obj[attr_name] = getattr(schema, attr_name)
+        model_spec[attr_name] = getattr(schema, attr_name)
 
 def _to_json_from_json_schema(schema: JsonSchemaObject) -> dict[str, Any]:
     model_spec = {}
@@ -62,12 +62,14 @@ def _to_json_from_json_schema(schema: JsonSchemaObject) -> dict[str, Any]:
         model_spec["anyOf"] = [_to_json_from_json_schema(schema) for schema in schema.anyOf]
 
     _assign_attribute(model_spec, "in_field", schema)
+    _assign_attribute(model_spec, "in", schema)
     _assign_attribute(model_spec, "aliasName", schema)
 
     if hasattr(schema, 'model_extra') and schema.model_extra:
         # for each extra fiels, add it to the model spec
         for key, value in schema.model_extra.items():
-            model_spec[key] = value
+            if value is not None:
+                model_spec[key] = value
 
     if isinstance(schema, JsonSchemaObjectRef):
         model_spec["$ref"] = schema.ref
@@ -635,6 +637,11 @@ class TaskEventType(Enum):
     ON_TASK_STREAM = "on_task_stream"
     ON_TASK_ERROR = "on_task_error"
 
+class FlowData(BaseModel):
+    '''This class represents the data that is passed between tasks in a flow.'''
+    input: dict[str, Any] = Field(default_factory=dict)
+    output: dict[str, Any] = Field(default_factory=dict)
+
 class FlowContext(BaseModel):
  
     name: str | None = None # name of the process or task
@@ -642,10 +649,12 @@ class FlowContext(BaseModel):
     flow_id: str | None = None # id of the flow, this is at the flow definition level
     instance_id: str | None = None
     thread_id: str | None = None
+    correlation_id: str | None = None
+    tenant_id: str | None = None
     parent_context: Any | None = None
     child_context: List["FlowContext"] | None = None
     metadata: dict = Field(default_factory=dict[str, Any])
-    data: dict = Field(default_factory=dict[str, Any])
+    data: Optional[FlowData] = None
 
     def get(self, key: str) -> Any:
      
@@ -655,7 +664,7 @@ class FlowContext(BaseModel):
         if self.parent_context:
             pc = cast(FlowContext, self.parent_conetxt)
             return pc.get(key)
-
+    
 class FlowEventType(Enum):
  
     ON_FLOW_START = "on_flow_start"