vellum-ai 0.14.63__py3-none-any.whl → 1.13.5__py3-none-any.whl

vellum/workflows/utils/tests/test_functions.py

@@ -1,14 +1,26 @@
+import pytest
 from dataclasses import dataclass
-from typing import Dict, List, Optional, Union
+from enum import Enum
+from unittest.mock import Mock
+from typing import Annotated, Dict, List, Literal, Optional, Tuple, Union
 
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 
 from vellum.client.types.function_definition import FunctionDefinition
+from vellum.client.types.string_vellum_value import StringVellumValue
+from vellum.client.types.vellum_variable import VellumVariable
 from vellum.workflows import BaseWorkflow
 from vellum.workflows.inputs.base import BaseInputs
 from vellum.workflows.nodes.bases.base import BaseNode
 from vellum.workflows.state.base import BaseState
-from vellum.workflows.utils.functions import compile_function_definition, compile_workflow_function_definition
+from vellum.workflows.types.definition import DeploymentDefinition
+from vellum.workflows.utils.functions import (
+    compile_function_definition,
+    compile_inline_workflow_function_definition,
+    compile_workflow_deployment_function_definition,
+    tool,
+    use_tool_inputs,
+)
 
 
 def test_compile_function_definition__just_name():
@@ -177,14 +189,15 @@ def test_compile_function_definition__dataclasses():
     compiled_function = compile_function_definition(my_function)
 
     # THEN it should return the compiled function definition
+    ref_name = f"{__name__}.test_compile_function_definition__dataclasses.<locals>.MyDataClass"
     assert compiled_function == FunctionDefinition(
         name="my_function",
         parameters={
             "type": "object",
-            "properties": {"c": {"$ref": "#/$defs/MyDataClass"}},
+            "properties": {"c": {"$ref": f"#/$defs/{ref_name}"}},
             "required": ["c"],
             "$defs": {
-                "MyDataClass": {
+                ref_name: {
                     "type": "object",
                     "properties": {"a": {"type": "integer"}, "b": {"type": "string"}},
                     "required": ["a", "b"],
@@ -197,7 +210,7 @@ def test_compile_function_definition__dataclasses():
 def test_compile_function_definition__pydantic():
     # GIVEN a function with a pydantic model
     class MyPydanticModel(BaseModel):
-        a: int
+        a: int = Field(description="The first number")
         b: str
 
     def my_function(c: MyPydanticModel):
@@ -207,16 +220,20 @@ def test_compile_function_definition__pydantic():
     compiled_function = compile_function_definition(my_function)
 
     # THEN it should return the compiled function definition
+    ref_name = f"{__name__}.test_compile_function_definition__pydantic.<locals>.MyPydanticModel"
     assert compiled_function == FunctionDefinition(
         name="my_function",
         parameters={
             "type": "object",
-            "properties": {"c": {"$ref": "#/$defs/MyPydanticModel"}},
+            "properties": {"c": {"$ref": f"#/$defs/{ref_name}"}},
             "required": ["c"],
             "$defs": {
-                "MyPydanticModel": {
+                ref_name: {
                     "type": "object",
-                    "properties": {"a": {"type": "integer"}, "b": {"type": "string"}},
+                    "properties": {
+                        "a": {"type": "integer", "description": "The first number"},
+                        "b": {"type": "string"},
+                    },
                     "required": ["a", "b"],
                 }
             },
@@ -238,14 +255,15 @@ def test_compile_function_definition__default_dataclass():
     compiled_function = compile_function_definition(my_function)
 
     # THEN it should return the compiled function definition
+    ref_name = f"{__name__}.test_compile_function_definition__default_dataclass.<locals>.MyDataClass"
     assert compiled_function == FunctionDefinition(
         name="my_function",
         parameters={
             "type": "object",
-            "properties": {"c": {"$ref": "#/$defs/MyDataClass", "default": {"a": 1, "b": "hello"}}},
+            "properties": {"c": {"$ref": f"#/$defs/{ref_name}", "default": {"a": 1, "b": "hello"}}},
             "required": [],
             "$defs": {
-                "MyDataClass": {
+                ref_name: {
                     "type": "object",
                     "properties": {"a": {"type": "integer"}, "b": {"type": "string"}},
                     "required": ["a", "b"],
@@ -268,14 +286,15 @@ def test_compile_function_definition__default_pydantic():
     compiled_function = compile_function_definition(my_function)
 
     # THEN it should return the compiled function definition
+    ref_name = f"{__name__}.test_compile_function_definition__default_pydantic.<locals>.MyPydanticModel"
     assert compiled_function == FunctionDefinition(
         name="my_function",
         parameters={
             "type": "object",
-            "properties": {"c": {"$ref": "#/$defs/MyPydanticModel", "default": {"a": 1, "b": "hello"}}},
+            "properties": {"c": {"$ref": f"#/$defs/{ref_name}", "default": {"a": 1, "b": "hello"}}},
             "required": [],
             "$defs": {
-                "MyPydanticModel": {
+                ref_name: {
                     "type": "object",
                     "properties": {"a": {"type": "integer"}, "b": {"type": "string"}},
                     "required": ["a", "b"],
@@ -299,7 +318,7 @@ def test_compile_function_definition__lambda():
     )
 
 
-def test_compile_workflow_function_definition():
+def test_compile_inline_workflow_function_definition():
     class MyNode(BaseNode):
         pass
 
@@ -307,7 +326,7 @@ def test_compile_workflow_function_definition():
         graph = MyNode
 
     # WHEN compiling the function
-    compiled_function = compile_workflow_function_definition(MyWorkflow)
+    compiled_function = compile_inline_workflow_function_definition(MyWorkflow)
 
     # THEN it should return the compiled function definition
     assert compiled_function == FunctionDefinition(
@@ -316,7 +335,7 @@ def test_compile_workflow_function_definition():
     )
 
 
-def test_compile_workflow_function_definition__docstring():
+def test_compile_inline_workflow_function_definition__docstring():
     class MyNode(BaseNode):
         pass
 
@@ -328,7 +347,7 @@ def test_compile_workflow_function_definition__docstring():
         graph = MyNode
 
     # WHEN compiling the function
-    compiled_function = compile_workflow_function_definition(MyWorkflow)
+    compiled_function = compile_inline_workflow_function_definition(MyWorkflow)
 
     # THEN it should return the compiled function definition
     assert compiled_function == FunctionDefinition(
@@ -338,7 +357,7 @@ def test_compile_workflow_function_definition__docstring():
     )
 
 
-def test_compile_workflow_function_definition__all_args():
+def test_compile_inline_workflow_function_definition__all_args():
     class MyInputs(BaseInputs):
         a: str
         b: int
@@ -354,7 +373,7 @@ def test_compile_workflow_function_definition__all_args():
         graph = MyNode
 
     # WHEN compiling the workflow
-    compiled_function = compile_workflow_function_definition(MyWorkflow)
+    compiled_function = compile_inline_workflow_function_definition(MyWorkflow)
 
     # THEN it should return the compiled function definition
     assert compiled_function == FunctionDefinition(
@@ -374,7 +393,7 @@ def test_compile_workflow_function_definition__all_args():
     )
 
 
-def test_compile_workflow_function_definition__unions():
+def test_compile_inline_workflow_function_definition__unions():
     # GIVEN a workflow with a union
     class MyInputs(BaseInputs):
         a: Union[str, int]
@@ -386,7 +405,7 @@ def test_compile_workflow_function_definition__unions():
         graph = MyNode
 
     # WHEN compiling the workflow
-    compiled_function = compile_workflow_function_definition(MyWorkflow)
+    compiled_function = compile_inline_workflow_function_definition(MyWorkflow)
 
     # THEN it should return the compiled function definition
     assert compiled_function == FunctionDefinition(
@@ -399,7 +418,7 @@ def test_compile_workflow_function_definition__unions():
     )
 
 
-def test_compile_workflow_function_definition__optionals():
+def test_compile_inline_workflow_function_definition__optionals():
     class MyInputs(BaseInputs):
         a: str
         b: Optional[str]
@@ -414,7 +433,7 @@ def test_compile_workflow_function_definition__optionals():
         graph = MyNode
 
     # WHEN compiling the workflow
-    compiled_function = compile_workflow_function_definition(MyWorkflow)
+    compiled_function = compile_inline_workflow_function_definition(MyWorkflow)
 
     # THEN it should return the compiled function definition
     assert compiled_function == FunctionDefinition(
@@ -431,3 +450,441 @@ def test_compile_workflow_function_definition__optionals():
             "required": ["a", "b", "c"],
         },
     )
+
+
+def test_compile_workflow_deployment_function_definition__just_name():
+    # GIVEN a mock Vellum client and deployment
+    mock_client = Mock()
+    mock_release = Mock()
+    mock_release.deployment.name = "my_deployment"
+    mock_release.workflow_version.input_variables = []
+    mock_release.description = "This is a test deployment"
+    mock_client.workflow_deployments.retrieve_workflow_deployment_release.return_value = mock_release
+
+    deployment_definition = DeploymentDefinition(deployment="my_deployment", release_tag="LATEST")
+
+    # WHEN compiling the deployment workflow function
+    compiled_function = compile_workflow_deployment_function_definition(deployment_definition, mock_client)
+
+    # THEN it should return the compiled function definition (same structure as function test)
+    assert compiled_function == FunctionDefinition(
+        name="my_deployment",
+        description="This is a test deployment",
+        parameters={"type": "object", "properties": {}, "required": []},
+    )
+
+
+def test_compile_workflow_deployment_function_definition__all_args():
+    # GIVEN a mock Vellum client and deployment
+    mock_client = Mock()
+    mock_release = Mock()
+
+    mock_inputs = []
+    for key, vellum_type in [
+        ("a", "STRING"),
+        ("b", "NUMBER"),
+        ("c", "JSON"),
+        ("d", "CHAT_HISTORY"),
+        ("e", "SEARCH_RESULTS"),
+        ("f", "ERROR"),
+        ("g", "ARRAY"),
+        ("h", "FUNCTION_CALL"),
+        ("i", "IMAGE"),
+        ("j", "AUDIO"),
+        ("k", "DOCUMENT"),
+        ("l", "NULL"),
+    ]:
+        mock_input = VellumVariable(
+            id=f"input_{key}",
+            key=key,
+            type=vellum_type,
+            required=True,
+            default=None,
+        )
+        mock_inputs.append(mock_input)
+
+    mock_release.workflow_version.input_variables = mock_inputs
+    mock_release.deployment.name = "my_deployment"
+    mock_release.description = "This is a test deployment"
+    mock_client.workflow_deployments.retrieve_workflow_deployment_release.return_value = mock_release
+
+    deployment_definition = DeploymentDefinition(deployment="my_deployment", release_tag="latest")
+
+    # WHEN compiling the deployment workflow function
+    compiled_function = compile_workflow_deployment_function_definition(deployment_definition, mock_client)
+
+    # THEN it should return the compiled function definition
+    assert compiled_function == FunctionDefinition(
+        name="my_deployment",
+        description="This is a test deployment",
+        parameters={
+            "type": "object",
+            "properties": {
+                "a": {"type": "string"},
+                "b": {"type": "number"},
+                "c": {"type": "object"},
+                "d": {"type": "array"},
+                "e": {"type": "array"},
+                "f": {"type": "object"},
+                "g": {"type": "array"},
+                "h": {"type": "object"},
+                "i": {"type": "object"},
+                "j": {"type": "object"},
+                "k": {"type": "object"},
+                "l": {"type": "null"},
+            },
+            "required": ["a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l"],
+        },
+    )
+
+
+def test_compile_workflow_deployment_function_definition__defaults():
+    # GIVEN a mock Vellum client and deployment
+    mock_client = Mock()
+    mock_release = Mock()
+
+    mock_inputs = []
+
+    mock_input_no_default = VellumVariable(
+        id="no_default",
+        key="no_default",
+        type="STRING",
+        required=True,
+        default=None,
+    )
+    mock_inputs.append(mock_input_no_default)
+
+    mock_input_null_default = VellumVariable(
+        id="null_default",
+        key="null_default",
+        type="STRING",
+        required=False,
+        default=StringVellumValue(value=None),
+    )
+    mock_inputs.append(mock_input_null_default)
+
+    mock_input_actual_default = VellumVariable(
+        id="actual_default",
+        key="actual_default",
+        type="STRING",
+        required=False,
+        default=StringVellumValue(value="hello world"),
+    )
+    mock_inputs.append(mock_input_actual_default)
+
+    mock_release.workflow_version.input_variables = mock_inputs
+    mock_release.deployment.name = "my_deployment"
+    mock_release.description = "This is a test deployment"
+    mock_client.workflow_deployments.retrieve_workflow_deployment_release.return_value = mock_release
+
+    deployment_definition = DeploymentDefinition(deployment="my_deployment", release_tag="LATEST")
+
+    # WHEN compiling the deployment workflow function
+    compiled_function = compile_workflow_deployment_function_definition(deployment_definition, mock_client)
+
+    # THEN it should return the compiled function definition with proper default handling
+    assert compiled_function == FunctionDefinition(
+        name="my_deployment",
+        description="This is a test deployment",
+        parameters={
+            "type": "object",
+            "properties": {
+                "no_default": {"type": "string"},
+                "null_default": {"type": "string", "default": None},
+                "actual_default": {"type": "string", "default": "hello world"},
+            },
+            "required": ["no_default"],
+        },
+    )
+
+
+@pytest.mark.parametrize(
+    "annotation,expected_schema",
+    [
+        (Literal["a", "b"], {"type": "string", "enum": ["a", "b"]}),
+        (Literal["a", 1], {"enum": ["a", 1]}),
+    ],
+)
+def test_compile_function_definition__literal(annotation, expected_schema):
+    def my_function(a: annotation):  # type: ignore
+        pass
+
+    compiled_function = compile_function_definition(my_function)
+    assert isinstance(compiled_function.parameters, dict)
+    assert compiled_function.parameters["properties"]["a"] == expected_schema
+
+
+def test_compile_function_definition__literal_type_not_in_map():
+    class MyEnum(Enum):
+        FOO = "foo"
+        BAR = "bar"
+
+    def my_function(a: Literal[MyEnum.FOO, MyEnum.BAR]):
+        pass
+
+    compiled_function = compile_function_definition(my_function)
+    assert isinstance(compiled_function.parameters, dict)
+    assert compiled_function.parameters["properties"]["a"] == {"enum": [MyEnum.FOO, MyEnum.BAR]}
+
+
+def test_compile_function_definition__annotated_descriptions():
+    # GIVEN a function with annotated parameters that include descriptions
+    def my_function(
+        bar: Annotated[str, "My bar parameter"],
+        other: Annotated[int, "My other parameter"],
+        regular_param: str,
+        optional_param: Annotated[bool, "Optional boolean parameter"] = True,
+    ):
+        """Test function with annotated parameters."""
+        pass
+
+    # WHEN compiling the function
+    compiled_function = compile_function_definition(my_function)
+
+    # THEN it should return the compiled function definition with descriptions
+    assert compiled_function == FunctionDefinition(
+        name="my_function",
+        description="Test function with annotated parameters.",
+        parameters={
+            "type": "object",
+            "properties": {
+                "bar": {"type": "string", "description": "My bar parameter"},
+                "other": {"type": "integer", "description": "My other parameter"},
+                "regular_param": {"type": "string"},
+                "optional_param": {"type": "boolean", "description": "Optional boolean parameter", "default": True},
+            },
+            "required": ["bar", "other", "regular_param"],
+        },
+    )
+
+
+def test_compile_function_definition__annotated_without_description():
+    # GIVEN a function with annotated parameters but no description metadata
+    def my_function(param: Annotated[str, None]):
+        pass
+
+    # WHEN compiling the function
+    compiled_function = compile_function_definition(my_function)
+
+    # THEN it should return the compiled function definition without description
+    assert compiled_function == FunctionDefinition(
+        name="my_function",
+        parameters={
+            "type": "object",
+            "properties": {
+                "param": {"type": "string"},
+            },
+            "required": ["param"],
+        },
+    )
+
+
+def test_compile_function_definition__annotated_complex_types():
+    # GIVEN a function with annotated types
+    def my_function(
+        location: Annotated[Literal["New York", "Portland"], "The location you found"],
+        items: Annotated[List[str], "List of string items"],
+        config: Annotated[Dict[str, int], "Configuration mapping"],
+    ):
+        pass
+
+    # WHEN compiling the function
+    compiled_function = compile_function_definition(my_function)
+
+    # THEN it should return the compiled function definition with descriptions for complex types
+    assert compiled_function == FunctionDefinition(
+        name="my_function",
+        parameters={
+            "type": "object",
+            "properties": {
+                "location": {
+                    "type": "string",
+                    "enum": ["New York", "Portland"],
+                    "description": "The location you found",
+                },
+                "items": {"type": "array", "items": {"type": "string"}, "description": "List of string items"},
+                "config": {
+                    "type": "object",
+                    "additionalProperties": {"type": "integer"},
+                    "description": "Configuration mapping",
+                },
+            },
+            "required": ["location", "items", "config"],
+        },
+    )
+
+
+@pytest.mark.parametrize(
+    "annotation,expected_schema",
+    [
+        (
+            Tuple[int, ...],
+            {"type": "array", "items": {"type": "integer"}},
+        ),
+        (
+            Tuple[int, str],
+            {
+                "type": "array",
+                "prefixItems": [{"type": "integer"}, {"type": "string"}],
+                "minItems": 2,
+                "maxItems": 2,
+            },
+        ),
+    ],
+)
+def test_compile_function_definition__tuples(annotation, expected_schema):
+    def my_function(a: annotation):  # type: ignore
+        pass
+
+    compiled_function = compile_function_definition(my_function)
+    assert isinstance(compiled_function.parameters, dict)
+    assert compiled_function.parameters["properties"]["a"] == expected_schema
+
+
+def test_compile_function_definition__string_annotations_with_future_imports():
+    """Test that string annotations work with __future__ import annotations."""
+    # This simulates what happens when using `from __future__ import annotations`
+    # where type annotations become string literals at runtime
+
+    def my_function_with_string_annotations(
+        a: "str",
+        b: "int",
+        c: "float",
+        d: "bool",
+        e: "list",
+        f: "dict",
+        g: "None",
+    ):
+        """Function with string type annotations."""
+        pass
+
+    # WHEN compiling the function
+    compiled_function = compile_function_definition(my_function_with_string_annotations)
+
+    # THEN it should return the compiled function definition with proper types
+    assert compiled_function == FunctionDefinition(
+        name="my_function_with_string_annotations",
+        description="Function with string type annotations.",
+        parameters={
+            "type": "object",
+            "properties": {
+                "a": {"type": "string"},
+                "b": {"type": "integer"},
+                "c": {"type": "number"},
+                "d": {"type": "boolean"},
+                "e": {"type": "array"},
+                "f": {"type": "object"},
+                "g": {"type": "null"},
+            },
+            "required": ["a", "b", "c", "d", "e", "f", "g"],
+        },
+    )
+
+
+def test_use_tool_inputs__inline_vs_decorator():
+    """
+    Tests that inline use_tool_inputs(...)(func) behaves the same as @use_tool_inputs(...) decorator.
+    """
+
+    # GIVEN a function with some parameters
+    def my_function(a: str, b: int, c: float) -> str:
+        """A test function."""
+        return f"{a}-{b}-{c}"
+
+    # WHEN using use_tool_inputs as a decorator
+    @use_tool_inputs(a="fixed_a", b=42)
+    def decorated_function(a: str, b: int, c: float) -> str:
+        """A test function."""
+        return f"{a}-{b}-{c}"
+
+    # AND using use_tool_inputs inline
+    inline_function = use_tool_inputs(a="fixed_a", b=42)(my_function)
+
+    # THEN both should have the same __vellum_inputs__ attribute
+    assert hasattr(decorated_function, "__vellum_inputs__")
+    assert hasattr(inline_function, "__vellum_inputs__")
+
+    # AND the inputs should be identical
+    assert decorated_function.__vellum_inputs__ == inline_function.__vellum_inputs__
+    assert decorated_function.__vellum_inputs__ == {"a": "fixed_a", "b": 42}
+
+
+def test_tool__inline_vs_decorator():
+    """
+    Tests that inline tool(inputs={...})(func) behaves the same as @tool(inputs={...}) decorator.
+    """
+
+    # GIVEN a function with some parameters
+    def my_function(a: str, b: int, c: float) -> str:
+        """A test function."""
+        return f"{a}-{b}-{c}"
+
+    # WHEN using tool as a decorator
+    @tool(inputs={"a": "fixed_a", "b": 42})
+    def decorated_function(a: str, b: int, c: float) -> str:
+        """A test function."""
+        return f"{a}-{b}-{c}"
+
+    # AND using tool inline
+    inline_function = tool(inputs={"a": "fixed_a", "b": 42})(my_function)
+
+    # THEN both should have the same __vellum_inputs__ attribute
+    assert hasattr(decorated_function, "__vellum_inputs__")
+    assert hasattr(inline_function, "__vellum_inputs__")
+
+    # AND the inputs should be identical
+    assert decorated_function.__vellum_inputs__ == inline_function.__vellum_inputs__
+    assert decorated_function.__vellum_inputs__ == {"a": "fixed_a", "b": 42}
+
+
+def test_tool__backward_compatibility_with_use_tool_inputs():
+    """
+    Tests that tool(inputs={...}) and use_tool_inputs(**inputs) produce the same __vellum_inputs__ attribute.
+    """
+
+    # GIVEN a function with some parameters
+    def my_function(a: str, b: int) -> str:
+        """A test function."""
+        return f"{a}-{b}"
+
+    # WHEN using tool as a decorator with inputs dict
+    @tool(inputs={"a": "value_a"})
+    def tool_decorated(a: str, b: int) -> str:
+        """A test function."""
+        return f"{a}-{b}"
+
+    # AND using use_tool_inputs as a decorator with kwargs
+    @use_tool_inputs(a="value_a")
+    def use_tool_inputs_decorated(a: str, b: int) -> str:
+        """A test function."""
+        return f"{a}-{b}"
+
+    # THEN both should have identical __vellum_inputs__ attributes
+    assert getattr(tool_decorated, "__vellum_inputs__") == getattr(use_tool_inputs_decorated, "__vellum_inputs__")
+    assert getattr(tool_decorated, "__vellum_inputs__") == {"a": "value_a"}
+
+
+def test_tool_examples_included_in_schema():
+    @tool(
+        examples=[
+            {"location": "San Francisco"},
+            {"location": "New York", "units": "celsius"},
+        ]
+    )
+    def get_current_weather(location: str, units: str = "fahrenheit") -> str:
+        return "sunny"
+
+    compiled = compile_function_definition(get_current_weather)
+    assert isinstance(compiled.parameters, dict)
+    assert compiled.parameters == {
+        "type": "object",
+        "properties": {
+            "location": {"type": "string"},
+            "units": {"type": "string", "default": "fahrenheit"},
+        },
+        "required": ["location"],
+        "examples": [
+            {"location": "San Francisco"},
+            {"location": "New York", "units": "celsius"},
+        ],
+    }

vellum/workflows/utils/tests/test_names.py

@@ -8,6 +8,15 @@ from vellum.workflows.utils.names import pascal_to_title_case
     [
         ("MyPascalCaseString", "My Pascal Case String"),
         ("AnotherPascalCaseString", "Another Pascal Case String"),
+        ("FetchDeploymentScheme", "Fetch Deployment Scheme"),
+        ("CheckDeploymentExists", "Check Deployment Exists"),
+        ("APINode", "API Node"),
+        ("HTTPSConnection", "HTTPS Connection"),
+        ("XMLParser", "XML Parser"),
+        ("SimpleWord", "Simple Word"),
+        ("A", "A"),
+        ("AB", "AB"),
+        ("ABCDef", "ABC Def"),
     ],
 )
 def test_pascal_to_title_case(input_str, expected):