vellum-ai 0.14.45__py3-none-any.whl → 0.14.47__py3-none-any.whl

vellum/workflows/nodes/displayable/code_execution_node/tests/test_code_execution_node.py

@@ -409,6 +409,56 @@ def main(word: str) -> int:
     assert exc_info.value.message == "Expected an output of type 'int', but received 'str'"
 
 
+def test_run_node__run_inline__incorrect_output_type_list():
+    """Confirm that CodeExecutionNodes raise an error if the output type is incorrect during inline execution."""
+
+    # GIVEN a node that subclasses CodeExecutionNode that returns a list but is defined to return an int
+    class ExampleCodeExecutionNode(CodeExecutionNode[BaseState, int]):
+        code = """\
+def main(output: list[str]) -> int:
+    return output.value
+"""
+        runtime = "PYTHON_3_11_6"
+
+        code_inputs = {
+            "output": ["hello", "world"],
+        }
+
+    # WHEN we run the node
+    node = ExampleCodeExecutionNode()
+    with pytest.raises(NodeException) as exc_info:
+        node.run()
+
+    # THEN the node should have produced the exception we expected
+    assert exc_info.value.message == "Expected an output of type 'int', but received 'list'"
+
+
+def test_run_node__run_inline__incorrect_output_type_dict():
+    """Confirm that CodeExecutionNodes raise an error if the output type is incorrect during inline execution."""
+
+    # GIVEN a node that subclasses CodeExecutionNode that returns a dict but is defined to return an int
+    class ExampleCodeExecutionNode(CodeExecutionNode[BaseState, int]):
+        code = """\
+def main(output: dict[str, str]) -> int:
+    return output.value
+"""
+        runtime = "PYTHON_3_11_6"
+
+        code_inputs = {
+            "output": {
+                "foo": "bar",
+            }
+        }
+
+    # WHEN we run the node
+    node = ExampleCodeExecutionNode()
+    with pytest.raises(NodeException) as exc_info:
+        node.run()
+
+    # THEN the node should have produced the exception we expected
+    assert exc_info.value.message == "Expected an output of type 'int', but received 'dict'"
+
+
 def test_run_node__run_inline__valid_dict_to_pydantic():
     """Confirm that CodeExecutionNodes can convert a dict to a Pydantic model during inline execution."""
 
@@ -821,3 +871,65 @@ def main(arg1: list) -> str:
 
     # AND the result should be the correct output
     assert outputs == {"result": "bar", "log": ""}
+
+
+def test_run_node__string_value_wrapper__get_attr():
+    # GIVEN a node that accesses the 'value' property of a string input
+    class ExampleCodeExecutionNode(CodeExecutionNode[BaseState, str]):
+        code = """\
+def main(text: str) -> str:
+    return text.value
+"""
+        code_inputs = {
+            "text": "hello",
+        }
+        runtime = "PYTHON_3_11_6"
+
+    # WHEN we run the node
+    node = ExampleCodeExecutionNode()
+    outputs = node.run()
+
+    # THEN the node should successfully access the string value through the .value property
+    assert outputs == {"result": "hello", "log": ""}
+
+
+def test_run_node__string_value_wrapper__get_item():
+    # GIVEN a node that accesses the 'value' property of a string input
+    class ExampleCodeExecutionNode(CodeExecutionNode[BaseState, str]):
+        code = """\
+def main(text: str) -> str:
+    return text["value"]
+"""
+        code_inputs = {
+            "text": "hello",
+        }
+        runtime = "PYTHON_3_11_6"
+
+    # WHEN we run the node
+    node = ExampleCodeExecutionNode()
+    outputs = node.run()
+
+    # THEN the node should successfully access the string value through the .value property
+    assert outputs == {"result": "hello", "log": ""}
+
+
+def test_run_node__string_value_wrapper__list_of_dicts():
+    # GIVEN a node that accesses the 'value' property of a string input
+    class ExampleCodeExecutionNode(CodeExecutionNode[BaseState, Any]):
+        code = """\
+def main(output: list[str]) -> list[str]:
+    results = []
+    for item in output:
+        results.append(item['value'])
+
+    return results
+"""
+        code_inputs = {"output": ['{"foo": "bar"}', '{"foo2": "bar2"}']}
+        runtime = "PYTHON_3_11_6"
+
+    # WHEN we run the node
+    node = ExampleCodeExecutionNode()
+    outputs = node.run()
+
+    # THEN the node should successfully access the string value
+    assert outputs == {"result": ['{"foo": "bar"}', '{"foo2": "bar2"}'], "log": ""}

vellum/workflows/nodes/displayable/code_execution_node/utils.py

@@ -8,6 +8,7 @@ from vellum.workflows.errors.types import WorkflowErrorCode
 from vellum.workflows.exceptions import NodeException
 from vellum.workflows.nodes.utils import cast_to_output_type
 from vellum.workflows.state.context import WorkflowContext
+from vellum.workflows.types.code_execution_node_wrappers import ListWrapper, clean_for_dict_wrapper
 from vellum.workflows.types.core import EntityInputsInterface
 
 
@@ -35,58 +36,6 @@ def read_file_from_path(
         return None
 
 
-class ListWrapper(list):
-    def __getitem__(self, key):
-        item = super().__getitem__(key)
-        if not isinstance(item, DictWrapper) and not isinstance(item, ListWrapper):
-            self.__setitem__(key, _clean_for_dict_wrapper(item))
-
-        return super().__getitem__(key)
-
-
-class DictWrapper(dict):
-    """
-    This wraps a dict object to make it behave basically the same as a standard javascript object
-    and enables us to use vellum types here without a shared library since we don't actually
-    typecheck things here.
-    """
-
-    def __getitem__(self, key):
-        return self.__getattr__(key)
-
-    def __getattr__(self, attr):
-        if attr not in self:
-            if attr == "value":
-                # In order to be backwards compatible with legacy Workflows, which wrapped
-                # several values as VellumValue objects, we use the "value" key to return itself
-                return self
-
-            raise AttributeError(f"dict has no key: '{attr}'")
-
-        item = super().__getitem__(attr)
-        if not isinstance(item, DictWrapper) and not isinstance(item, ListWrapper):
-            self.__setattr__(attr, _clean_for_dict_wrapper(item))
-
-        return super().__getitem__(attr)
-
-    def __setattr__(self, name, value):
-        self[name] = value
-
-
-def _clean_for_dict_wrapper(obj):
-    if isinstance(obj, dict):
-        wrapped = DictWrapper(obj)
-        for key in wrapped:
-            wrapped[key] = _clean_for_dict_wrapper(wrapped[key])
-
-        return wrapped
-
-    elif isinstance(obj, list):
-        return ListWrapper(map(lambda item: _clean_for_dict_wrapper(item), obj))
-
-    return obj
-
-
 def run_code_inline(
     code: str,
     inputs: EntityInputsInterface,
@@ -107,12 +56,12 @@ def run_code_inline(
                     (
                         item.model_dump()
                         if isinstance(item, BaseModel)
-                        else _clean_for_dict_wrapper(item) if isinstance(item, (dict, list)) else item
+                        else clean_for_dict_wrapper(item) if isinstance(item, (dict, list, str)) else item
                     )
                     for item in value
                 ]
             )
-        return _clean_for_dict_wrapper(value)
+        return clean_for_dict_wrapper(value)
 
     exec_globals = {
         "__arg__inputs": {name: wrap_value(value) for name, value in inputs.items()},

vellum/workflows/nodes/displayable/tests/test_text_prompt_deployment_node.py

@@ -9,7 +9,6 @@ from vellum import (
     PromptOutput,
     StringVellumValue,
 )
-from vellum.workflows.constants import OMIT
 from vellum.workflows.inputs import BaseInputs
 from vellum.workflows.nodes import PromptDeploymentNode
 from vellum.workflows.state import BaseState
@@ -66,14 +65,14 @@ def test_text_prompt_deployment_node__basic(vellum_client):
 
     # AND we should have made the expected call to stream the prompt execution
     vellum_client.execute_prompt_stream.assert_called_once_with(
-        expand_meta=OMIT,
-        expand_raw=OMIT,
-        external_id=OMIT,
+        expand_meta=None,
+        expand_raw=None,
+        external_id=None,
         inputs=[],
-        metadata=OMIT,
+        metadata=None,
         prompt_deployment_id=None,
         prompt_deployment_name="my-deployment",
-        raw_overrides=OMIT,
+        raw_overrides=None,
         release_tag="LATEST",
         request_options={
             "additional_body_parameters": {"execution_context": {"parent_context": None, "trace_id": mock.ANY}}

vellum/workflows/nodes/utils.py

@@ -1,4 +1,5 @@
 from functools import cache
+import inspect
 import json
 import sys
 from types import ModuleType
@@ -14,6 +15,7 @@ from vellum.workflows.nodes import BaseNode
 from vellum.workflows.nodes.bases.base_adornment_node import BaseAdornmentNode
 from vellum.workflows.ports.port import Port
 from vellum.workflows.state.base import BaseState
+from vellum.workflows.types.code_execution_node_wrappers import DictWrapper, ListWrapper, StringValueWrapper
 from vellum.workflows.types.core import Json
 from vellum.workflows.types.generics import NodeType
 
@@ -176,6 +178,12 @@ def parse_type_from_str(result_as_str: str, output_type: Any) -> Any:
 
 
 def _get_type_name(obj: Any) -> str:
+    if inspect.isclass(obj) and issubclass(obj, StringValueWrapper):
+        return "str"
+    if inspect.isclass(obj) and issubclass(obj, ListWrapper):
+        return "list"
+    if inspect.isclass(obj) and issubclass(obj, DictWrapper):
+        return "dict"
     if isinstance(obj, type):
         return obj.__name__
 

vellum/workflows/types/code_execution_node_wrappers.py

@@ -0,0 +1,69 @@
+class StringValueWrapper(str):
+    def __getitem__(self, key):
+        if key == "value":
+            return self
+        raise KeyError(key)
+
+    def __getattr__(self, attr):
+        if attr == "value":
+            return self
+        raise AttributeError(f"'str' object has no attribute '{attr}'")
+
+
+class ListWrapper(list):
+    def __getitem__(self, key):
+        item = super().__getitem__(key)
+        if not isinstance(item, DictWrapper) and not isinstance(item, ListWrapper):
+            self.__setitem__(key, clean_for_dict_wrapper(item))
+
+        return super().__getitem__(key)
+
+    def __getattr__(self, attr):
+        if attr == "value":
+            return self
+        raise AttributeError(f"'list' object has no attribute '{attr}'")
+
+
+class DictWrapper(dict):
+    """
+    This wraps a dict object to make it behave basically the same as a standard javascript object
+    and enables us to use vellum types here without a shared library since we don't actually
+    typecheck things here.
+    """
+
+    def __getitem__(self, key):
+        return self.__getattr__(key)
+
+    def __getattr__(self, attr):
+        if attr not in self:
+            if attr == "value":
+                # In order to be backwards compatible with legacy Workflows, which wrapped
+                # several values as VellumValue objects, we use the "value" key to return itself
+                return self
+
+            raise AttributeError(f"dict has no key: '{attr}'")
+
+        item = super().__getitem__(attr)
+        if not isinstance(item, DictWrapper) and not isinstance(item, ListWrapper):
+            self.__setattr__(attr, clean_for_dict_wrapper(item))
+
+        return super().__getitem__(attr)
+
+    def __setattr__(self, name, value):
+        self[name] = value
+
+
+def clean_for_dict_wrapper(obj):
+    if isinstance(obj, dict):
+        wrapped = DictWrapper(obj)
+        for key in wrapped:
+            wrapped[key] = clean_for_dict_wrapper(wrapped[key])
+
+        return wrapped
+
+    elif isinstance(obj, list):
+        return ListWrapper(map(lambda item: clean_for_dict_wrapper(item), obj))
+    elif isinstance(obj, str):
+        return StringValueWrapper(obj)
+
+    return obj

vellum/workflows/vellum_client.py

@@ -1,22 +1,34 @@
 import os
-from typing import Optional
+from typing import List, Optional
 
 from vellum import Vellum, VellumEnvironment
 
 
-def create_vellum_client(api_key: Optional[str] = None) -> Vellum:
+def create_vellum_client(api_key: Optional[str] = None, api_url: Optional[str] = None) -> Vellum:
     if api_key is None:
         api_key = os.getenv("VELLUM_API_KEY", default="")
 
     return Vellum(
         api_key=api_key,
-        environment=create_vellum_environment(),
+        environment=create_vellum_environment(api_url),
     )
 
 
-def create_vellum_environment() -> VellumEnvironment:
+def create_vellum_environment(api_url: Optional[str] = None) -> VellumEnvironment:
     return VellumEnvironment(
-        default=os.getenv("VELLUM_DEFAULT_API_URL", os.getenv("VELLUM_API_URL", "https://api.vellum.ai")),
-        documents=os.getenv("VELLUM_DOCUMENTS_API_URL", os.getenv("VELLUM_API_URL", "https://documents.vellum.ai")),
-        predict=os.getenv("VELLUM_PREDICT_API_URL", os.getenv("VELLUM_API_URL", "https://predict.vellum.ai")),
+        default=_resolve_env([api_url, "VELLUM_DEFAULT_API_URL", "VELLUM_API_URL"], "https://api.vellum.ai"),
+        documents=_resolve_env([api_url, "VELLUM_DOCUMENTS_API_URL", "VELLUM_API_URL"], "https://documents.vellum.ai"),
+        predict=_resolve_env([api_url, "VELLUM_PREDICT_API_URL", "VELLUM_API_URL"], "https://predict.vellum.ai"),
     )
+
+
+def _resolve_env(names: List[Optional[str]], default: str = "") -> str:
+    for name in names:
+        if not name:
+            continue
+
+        value = os.getenv(name)
+        if value:
+            return value
+
+    return default

{vellum_ai-0.14.45.dist-info → vellum_ai-0.14.47.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vellum-ai
-Version: 0.14.45
+Version: 0.14.47
 Summary: 
 License: MIT
 Requires-Python: >=3.9,<4.0