vellum-ai 0.14.69__py3-none-any.whl → 0.14.71__py3-none-any.whl

vellum_ee/workflows/display/nodes/vellum/tests/test_search_node.py

@@ -0,0 +1,104 @@
+from vellum.workflows import BaseWorkflow
+from vellum.workflows.inputs import BaseInputs
+from vellum.workflows.nodes.displayable import SearchNode
+from vellum.workflows.nodes.displayable.bases.types import (
+    MetadataLogicalCondition,
+    MetadataLogicalConditionGroup,
+    SearchFilters,
+)
+from vellum.workflows.state.base import BaseState
+from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
+
+
+def test_search_filters_with_input_reference():
+    """Test that SearchFilters with MetadataLogicalCondition using input references can be serialized"""
+
+    # GIVEN a search node with a metadata filter that uses an input reference
+    class TestInputs(BaseInputs):
+        file_id: str
+
+    class MySearchNode(SearchNode):
+        query = "my query"
+        document_index = "document_index"
+        filters = SearchFilters(
+            external_ids=None,
+            metadata=MetadataLogicalConditionGroup(
+                combinator="AND",
+                negated=False,
+                conditions=[MetadataLogicalCondition(lhs_variable="ID", operator="=", rhs_variable=TestInputs.file_id)],
+            ),
+        )
+
+    # AND a workflow with the Search Node
+    class Workflow(BaseWorkflow[TestInputs, BaseState]):
+        graph = MySearchNode
+
+    # WHEN the workflow is serialized
+    workflow_display = get_workflow_display(workflow_class=Workflow)
+    serialized_workflow: dict = workflow_display.serialize()
+
+    # THEN the node should properly serialize the filter reference
+    serialized_search_node = next(
+        node for node in serialized_workflow["workflow_raw_data"]["nodes"] if node["id"] == str(MySearchNode.__id__)
+    )
+
+    serialized_metadata_filter = next(
+        inp for inp in serialized_search_node["inputs"] if inp["key"] == "metadata_filters"
+    )
+
+    assert serialized_metadata_filter == {
+        "id": "4a9f96aa-ba3b-4c4e-9ce4-370fe64f717f",
+        "key": "metadata_filters",
+        "value": {
+            "combinator": "OR",
+            "rules": [
+                {
+                    "data": {
+                        "type": "JSON",
+                        "value": {
+                            "combinator": "AND",
+                            "conditions": [
+                                {
+                                    "lhs_variable_id": "9aedaffa-c2a4-4c37-9969-184e1ff43ded",
+                                    "operator": "=",
+                                    "rhs_variable_id": "c2151ef1-ad98-4940-b0e9-28dabe47a951",
+                                    "type": "LOGICAL_CONDITION",
+                                }
+                            ],
+                            "negated": False,
+                            "type": "LOGICAL_CONDITION_GROUP",
+                        },
+                    },
+                    "type": "CONSTANT_VALUE",
+                }
+            ],
+        },
+    }
+
+    # AND the LHS filter references should be present as node inputs
+    serialized_lhs_input = next(
+        inp for inp in serialized_search_node["inputs"] if inp["id"] == "9aedaffa-c2a4-4c37-9969-184e1ff43ded"
+    )
+    assert serialized_lhs_input["value"] == {
+        "combinator": "OR",
+        "rules": [
+            {
+                "data": {"type": "STRING", "value": "ID"},
+                "type": "CONSTANT_VALUE",
+            }
+        ],
+    }
+
+    # AND the RHS filter references should be present as node inputs
+    serialized_rhs_input = next(
+        inp for inp in serialized_search_node["inputs"] if inp["id"] == "c2151ef1-ad98-4940-b0e9-28dabe47a951"
+    )
+    assert serialized_rhs_input["value"] == {
+        "combinator": "OR",
+        "rules": [
+            {
+                "data": {"input_variable_id": "e2f4fff9-1277-47cb-8988-12f8ada450ba"},
+                "type": "INPUT_VARIABLE",
+            }
+        ],
+    }

vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_attributes_serialization.py

@@ -7,6 +7,7 @@ from vellum.client.types.chat_message import ChatMessage
 from vellum.workflows.inputs.base import BaseInputs
 from vellum.workflows.nodes.bases.base import BaseNode
 from vellum.workflows.references.constant import ConstantValueReference
+from vellum.workflows.references.environment_variable import EnvironmentVariableReference
 from vellum.workflows.references.lazy import LazyReference
 from vellum.workflows.references.vellum_secret import VellumSecretReference
 from vellum.workflows.state.base import BaseState
@@ -475,6 +476,21 @@ def test_serialize_node__node_execution(serialize_node):
     )
 
 
+def test_serialize_node__environment_variable(serialize_node):
+    class EnvironmentVariableGenericNode(BaseNode):
+        attr = EnvironmentVariableReference(name="API_KEY")
+
+    serialized_node = serialize_node(EnvironmentVariableGenericNode)
+
+    expected_value = {
+        "type": "ENVIRONMENT_VARIABLE",
+        "environment_variable": "API_KEY",
+    }
+
+    actual_value = serialized_node["attributes"][0]["value"]
+    assert actual_value == expected_value
+
+
 def test_serialize_node__coalesce(serialize_node):
     class CoalesceNodeA(BaseNode):
         class Outputs(BaseNode.Outputs):

vellum_ee/workflows/display/tests/workflow_serialization/test_basic_inline_prompt_node_serialization.py

@@ -174,6 +174,34 @@ def test_serialize_workflow():
                     "name": "ml_model",
                     "value": {"type": "CONSTANT_VALUE", "value": {"type": "STRING", "value": "gpt-4o"}},
                 },
+                {
+                    "id": "25f935f3-363f-4ead-a5a0-db234ca67e1e",
+                    "name": "blocks",
+                    "value": {
+                        "type": "CONSTANT_VALUE",
+                        "value": {
+                            "type": "JSON",
+                            "value": [
+                                {
+                                    "block_type": "CHAT_MESSAGE",
+                                    "state": None,
+                                    "cache_config": None,
+                                    "chat_role": "SYSTEM",
+                                    "chat_source": None,
+                                    "chat_message_unterminated": None,
+                                    "blocks": [
+                                        {
+                                            "block_type": "JINJA",
+                                            "state": None,
+                                            "cache_config": None,
+                                            "template": "What's your favorite {{noun}}?",
+                                        }
+                                    ],
+                                }
+                            ],
+                        },
+                    },
+                },
                 {
                     "id": "ffabe7d2-8ab6-4201-9d41-c4d7be1386e1",
                     "name": "prompt_inputs",
@@ -355,3 +383,57 @@ def test_serialize_workflow_with_descriptor_functions():
         "node_output_id": "470fadb9-b8b5-477e-a502-5209d398bcf9",
         "type": "NODE_OUTPUT",
     }
+
+
+def test_serialize_workflow_with_descriptor_blocks():
+    """Test that serialization handles BaseDescriptor instances in blocks list."""
+
+    class TestInputs(BaseInputs):
+        noun: str
+
+    class UpstreamNode(BaseNode):
+        class Outputs(BaseNode.Outputs):
+            results: list
+
+        def run(self) -> Outputs:
+            return self.Outputs(results=["test"])
+
+    class TestInlinePromptNodeWithDescriptorBlocks(InlinePromptNode):
+        ml_model = "gpt-4o"
+        blocks = [UpstreamNode.Outputs.results[0]]  # type: ignore
+        prompt_inputs = {"noun": TestInputs.noun}
+
+    class TestWorkflow(BaseWorkflow[TestInputs, BaseState]):
+        graph = UpstreamNode >> TestInlinePromptNodeWithDescriptorBlocks
+
+    workflow_display = get_workflow_display(workflow_class=TestWorkflow)
+    serialized: dict = workflow_display.serialize()
+
+    prompt_nodes = [node for node in serialized["workflow_raw_data"]["nodes"] if node["type"] == "PROMPT"]
+    prompt_node = prompt_nodes[0]
+
+    blocks = prompt_node["data"]["exec_config"]["prompt_template_block_data"]["blocks"]
+    descriptor_blocks = [block for block in blocks if not isinstance(block, dict) or not block.get("block_type")]
+    assert len(descriptor_blocks) == 0, "BaseDescriptor blocks should not appear in serialized blocks"
+
+    blocks_attr = next((attr for attr in prompt_node["attributes"] if attr["name"] == "blocks"), None)
+    assert blocks_attr is not None, "blocks attribute should be present when blocks contain BaseDescriptor"
+    assert blocks_attr["value"]["type"] == "ARRAY_REFERENCE", "blocks attribute should be serialized as ARRAY_REFERENCE"
+    assert blocks_attr["value"]["items"] == [
+        {
+            "type": "BINARY_EXPRESSION",
+            "lhs": {
+                "type": "NODE_OUTPUT",
+                "node_id": str(UpstreamNode.__id__),
+                "node_output_id": str(UpstreamNode.__output_ids__["results"]),
+            },
+            "operator": "accessField",
+            "rhs": {
+                "type": "CONSTANT_VALUE",
+                "value": {
+                    "type": "NUMBER",
+                    "value": 0.0,
+                },
+            },
+        }
+    ]

vellum_ee/workflows/display/tests/workflow_serialization/test_basic_inline_subworkflow_serialization.py

@@ -301,7 +301,15 @@ def test_serialize_workflow():
                         },
                     ],
                 },
-                "input_variables": [{"id": "704c4640-bfda-44f0-8da3-e9cfc4f21cf2", "key": "metro", "type": "STRING"}],
+                "input_variables": [
+                    {
+                        "id": "704c4640-bfda-44f0-8da3-e9cfc4f21cf2",
+                        "key": "metro",
+                        "type": "STRING",
+                        "required": True,
+                        "default": None,
+                    }
+                ],
                 "output_variables": [
                     {"id": "2fc57139-7420-49e5-96a6-dcbb3ff5d622", "key": "temperature", "type": "NUMBER"},
                     {"id": "fad5dd9f-3328-4e70-ad55-65a5325a4a82", "key": "reasoning", "type": "STRING"},

vellum_ee/workflows/display/tests/workflow_serialization/test_basic_search_node_serialization.py

@@ -167,7 +167,7 @@ def test_serialize_workflow():
                     "rules": [
                         {
                             "type": "INPUT_VARIABLE",
-                            "data": {"input_variable_id": "8ffffeb2-79b3-4105-acc7-78b0267da955"},
+                            "data": {"input_variable_id": "b118247f-96dd-4b3e-8289-9f277483c520"},
                         }
                     ],
                     "combinator": "OR",
@@ -180,7 +180,7 @@ def test_serialize_workflow():
                     "rules": [
                         {
                             "type": "INPUT_VARIABLE",
-                            "data": {"input_variable_id": "8ffffeb2-79b3-4105-acc7-78b0267da955"},
+                            "data": {"input_variable_id": "aae2c10a-88b7-40bd-87a2-5e1e60c1e906"},
                         }
                     ],
                     "combinator": "OR",
@@ -193,7 +193,7 @@ def test_serialize_workflow():
                     "rules": [
                         {
                             "type": "INPUT_VARIABLE",
-                            "data": {"input_variable_id": "f5eee974-b0c3-4775-bc8a-679a9e99d7ba"},
+                            "data": {"input_variable_id": "c9611a62-d1f5-4b41-bf9c-1aa3355760b4"},
                         }
                     ],
                     "combinator": "OR",
@@ -206,7 +206,7 @@ def test_serialize_workflow():
                     "rules": [
                         {
                             "type": "INPUT_VARIABLE",
-                            "data": {"input_variable_id": "f5eee974-b0c3-4775-bc8a-679a9e99d7ba"},
+                            "data": {"input_variable_id": "f374640e-a5c0-470e-ac71-c36c2b198c00"},
                         }
                     ],
                     "combinator": "OR",