vellum-ai 1.3.1__py3-none-any.whl → 1.3.3__py3-none-any.whl

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

@@ -47,17 +47,20 @@ def test_serialize_workflow():
 
     # AND its output variables should be what we expect
     output_variables = serialized_workflow["output_variables"]
-    assert len(output_variables) == 1
+    assert len(output_variables) == 2
     assert not DeepDiff(
-        [{"id": "15a0ab89-8ed4-43b9-afa2-3c0b29d4dc3e", "key": "results", "type": "JSON"}],
+        [
+            {"id": "15a0ab89-8ed4-43b9-afa2-3c0b29d4dc3e", "key": "results", "type": "JSON"},
+            {"id": "0ef1608e-1737-41cc-9b90-a8e124138f70", "key": "json", "type": "JSON"},
+        ],
         output_variables,
         ignore_order=True,
     )
 
     # AND its raw data should be what we expect
     workflow_raw_data = serialized_workflow["workflow_raw_data"]
-    assert len(workflow_raw_data["edges"]) == 2
-    assert len(workflow_raw_data["nodes"]) == 3
+    assert len(workflow_raw_data["edges"]) == 3
+    assert len(workflow_raw_data["nodes"]) == 4
 
     # AND each node should be serialized correctly
     entrypoint_node = workflow_raw_data["nodes"][0]
@@ -305,7 +308,7 @@ def test_serialize_workflow():
                     },
                 }
             ],
-            "display_data": {"position": {"x": 400.0, "y": -50.0}},
+            "display_data": {"position": {"x": 400.0, "y": 75.0}},
             "base": {
                 "name": "FinalOutputNode",
                 "module": ["vellum", "workflows", "nodes", "displayable", "final_output_node", "node"],
@@ -336,6 +339,14 @@ def test_serialize_workflow():
                 "target_handle_id": "46c99277-2b4b-477d-851c-ea497aef6b16",
                 "type": "DEFAULT",
             },
+            {
+                "id": "0b1a2960-4cd5-4045-844f-42b6c87487aa",
+                "source_node_id": "8450dd06-975a-41a4-a564-808ee8808fe6",
+                "source_handle_id": "d4a097ab-e22d-42f1-b6bc-2ed96856377a",
+                "target_node_id": "1f4e3b7b-6af1-42c8-ab33-05b0f01e2b62",
+                "target_handle_id": "7d94907f-c840-4ced-b813-ee3b17f2a8a9",
+                "type": "DEFAULT",
+            },
         ],
         serialized_edges,
         ignore_order=True,

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

@@ -0,0 +1,81 @@
+from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
+
+from tests.workflows.web_search.workflow import WebSearchWorkflow
+
+
+def test_serialize_web_search_workflow():
+    # GIVEN a WebSearchWorkflow with a node that has various input types
+    # - query from workflow input
+    # - api_key from VellumSecretReference
+    # - num_results as a constant integer
+    # - location as a constant string
+
+    # WHEN we serialize the workflow through the display system
+    workflow_display = get_workflow_display(workflow_class=WebSearchWorkflow)
+    serialized_workflow: dict = workflow_display.serialize()
+
+    # THEN we should get a properly structured serialized workflow
+    assert serialized_workflow.keys() == {
+        "workflow_raw_data",
+        "input_variables",
+        "state_variables",
+        "output_variables",
+    }
+
+    # AND it should have input variables
+    input_variables = serialized_workflow["input_variables"]
+    assert len(input_variables) == 1
+    assert input_variables[0]["key"] == "search_query"
+    assert input_variables[0]["type"] == "STRING"
+
+    # AND it should have output variables
+    output_variables = serialized_workflow["output_variables"]
+    assert len(output_variables) == 2
+    output_keys = {var["key"] for var in output_variables}
+    assert output_keys == {"search_results", "result_urls"}
+
+    # AND the web search node should be serialized as a GENERIC type node
+    workflow_raw_data = serialized_workflow["workflow_raw_data"]
+    nodes = workflow_raw_data["nodes"]
+
+    web_search_node = next(
+        (
+            node
+            for node in nodes
+            if node and node.get("type") == "GENERIC" and node.get("base", {}).get("name") == "WebSearchNode"
+        ),
+        None,
+    )
+
+    assert web_search_node is not None, "WebSearchNode should be serialized as GENERIC type"
+
+    # AND it should have the correct base module reference
+    expected_base = {
+        "name": "WebSearchNode",
+        "module": ["vellum", "workflows", "nodes", "displayable", "web_search_node", "node"],
+    }
+    assert (
+        web_search_node["base"] == expected_base
+    ), f"Base mismatch. Expected: {expected_base}, Found: {web_search_node['base']}"
+
+    # AND it should have all four serializable attributes defined in our display class
+    web_search_attributes = web_search_node["attributes"]
+    attribute_names = {attr["name"] for attr in web_search_attributes}
+
+    # AND the attributes should match exactly what we defined in __serializable_inputs__
+    assert attribute_names == {"query", "api_key", "num_results", "location"}
+
+    # AND it should have all three expected outputs from WebSearchNode
+    assert {output["name"] for output in web_search_node["outputs"]} == {"text", "urls", "results"}
+
+    # AND each attribute should have the correct value type based on its source
+    attr_types_by_name = {attr["name"]: attr["value"]["type"] for attr in web_search_attributes}
+
+    # AND query should reference the workflow input
+    assert attr_types_by_name["query"] == "WORKFLOW_INPUT"
+    # AND api_key should reference a Vellum secret
+    assert attr_types_by_name["api_key"] == "VELLUM_SECRET"
+    # AND num_results should be a constant value
+    assert attr_types_by_name["num_results"] == "CONSTANT_VALUE"
+    # AND location should be a constant value
+    assert attr_types_by_name["location"] == "CONSTANT_VALUE"