vellum-ai 0.13.24__tar.gz → 0.13.25__tar.gz

{vellum_ai-0.13.24 → vellum_ai-0.13.25}/PKG-INFO

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

{vellum_ai-0.13.24 → vellum_ai-0.13.25}/ee/vellum_ee/workflows/display/nodes/base_node_display.py

@@ -343,6 +343,27 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
             return self.serialize_value(display_context, value._value)
 
         if isinstance(value, LazyReference):
+            if isinstance(value._get, str):
+                reference_parts = value._get.split(".")
+                if len(reference_parts) < 3:
+                    raise Exception(
+                        f"Failed to parse lazy reference: {value._get}. Only Node Output references are supported."
+                    )
+
+                output_name = reference_parts[-1]
+                nested_class_name = reference_parts[-2]
+                if nested_class_name != "Outputs":
+                    raise Exception(
+                        f"Failed to parse lazy reference: {value._get}. Outputs are the only node reference supported."
+                    )
+
+                node_class_name = ".".join(reference_parts[:-2])
+                for node in display_context.node_displays.keys():
+                    if node.__name__ == node_class_name:
+                        return self.serialize_value(display_context, getattr(node.Outputs, output_name))
+
+                raise NotImplementedError(f"Failed to find a LazyReference for: {value._get}")
+
             return self.serialize_value(display_context, value._get())
 
         if isinstance(value, WorkflowInputReference):

{vellum_ai-0.13.24 → vellum_ai-0.13.25}/ee/vellum_ee/workflows/display/nodes/vellum/inline_prompt_node.py

@@ -1,9 +1,10 @@
 from uuid import UUID
-from typing import ClassVar, Dict, Generic, List, Optional, Tuple, Type, TypeVar, Union
+from typing import Callable, ClassVar, Dict, Generic, List, Optional, Tuple, Type, TypeVar, Union
 
-from vellum import PromptBlock, RichTextChildBlock, VellumVariable
+from vellum import FunctionDefinition, PromptBlock, RichTextChildBlock, VellumVariable
 from vellum.workflows.nodes import InlinePromptNode
 from vellum.workflows.types.core import JsonObject
+from vellum.workflows.utils.functions import compile_function_definition
 from vellum.workflows.utils.uuids import uuid4_from_hash
 from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
 from vellum_ee.workflows.display.nodes.utils import raise_if_descriptor
@@ -32,6 +33,17 @@ class BaseInlinePromptNodeDisplay(BaseNodeVellumDisplay[_InlinePromptNodeType],
         _, output_display = display_context.global_node_output_displays[node.Outputs.text]
         _, array_display = display_context.global_node_output_displays[node.Outputs.results]
         node_blocks = raise_if_descriptor(node.blocks)
+        function_definitions = raise_if_descriptor(node.functions)
+
+        blocks: list = [
+            self._generate_prompt_block(block, input_variable_id_by_name, [i]) for i, block in enumerate(node_blocks)
+        ]
+        functions = (
+            [self._generate_function_tools(function) for function in function_definitions]
+            if function_definitions
+            else []
+        )
+        blocks.extend(functions)
 
         return {
             "id": str(node_id),
@@ -50,10 +62,7 @@ class BaseInlinePromptNodeDisplay(BaseNodeVellumDisplay[_InlinePromptNodeType],
                     "input_variables": [prompt_input.dict() for prompt_input in prompt_inputs],
                     "prompt_template_block_data": {
                         "version": 1,
-                        "blocks": [
-                            self._generate_prompt_block(block, input_variable_id_by_name, [i])
-                            for i, block in enumerate(node_blocks)
-                        ],
+                        "blocks": blocks,
                     },
                 },
                 "ml_model_name": raise_if_descriptor(node.ml_model),
@@ -91,6 +100,21 @@ class BaseInlinePromptNodeDisplay(BaseNodeVellumDisplay[_InlinePromptNodeType],
 
         return node_inputs, prompt_inputs
 
+    def _generate_function_tools(self, function: Union[FunctionDefinition, Callable]) -> JsonObject:
+        normalized_functions = (
+            function if isinstance(function, FunctionDefinition) else compile_function_definition(function)
+        )
+        return {
+            "block_type": "FUNCTION_DEFINITION",
+            "properties": {
+                "function_name": normalized_functions.name,
+                "function_description": normalized_functions.description,
+                "function_parameters": normalized_functions.parameters,
+                "function_forced": normalized_functions.forced,
+                "function_strict": normalized_functions.strict,
+            },
+        }
+
     def _generate_prompt_block(
         self,
         prompt_block: Union[PromptBlock, RichTextChildBlock],

{vellum_ai-0.13.24 → vellum_ai-0.13.25}/ee/vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_attributes_serialization.py

@@ -7,9 +7,12 @@ from vellum.workflows.nodes.bases.base import BaseNode
 from vellum.workflows.references.constant import ConstantValueReference
 from vellum.workflows.references.lazy import LazyReference
 from vellum.workflows.references.vellum_secret import VellumSecretReference
+from vellum.workflows.workflows.base import BaseWorkflow
 from vellum_ee.workflows.display.base import WorkflowInputsDisplay
 from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
 from vellum_ee.workflows.display.nodes.types import NodeOutputDisplay
+from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
+from vellum_ee.workflows.display.workflows.vellum_workflow_display import VellumWorkflowDisplay
 
 
 class Inputs(BaseInputs):
@@ -151,6 +154,43 @@ def test_serialize_node__lazy_reference(serialize_node):
     )
 
 
+def test_serialize_node__lazy_reference_with_string():
+    # GIVEN two nodes with one lazily referencing the other
+    class LazyReferenceGenericNode(BaseNode):
+        attr = LazyReference[str]("OtherNode.Outputs.result")
+
+    class OtherNode(BaseNode):
+        class Outputs(BaseNode.Outputs):
+            result: str
+
+    # AND a workflow with both nodes
+    class Workflow(BaseWorkflow):
+        graph = LazyReferenceGenericNode >> OtherNode
+
+    # WHEN the workflow is serialized
+    workflow_display = get_workflow_display(base_display_class=VellumWorkflowDisplay, workflow_class=Workflow)
+    serialized_workflow: dict = workflow_display.serialize()
+
+    # THEN the node should properly serialize the attribute reference
+    lazy_reference_node = next(
+        node
+        for node in serialized_workflow["workflow_raw_data"]["nodes"]
+        if node["id"] == str(LazyReferenceGenericNode.__id__)
+    )
+
+    assert lazy_reference_node["attributes"] == [
+        {
+            "id": "98833d71-42a8-47e9-81c4-6a35646e3d3c",
+            "name": "attr",
+            "value": {
+                "type": "NODE_OUTPUT",
+                "node_id": str(OtherNode.__id__),
+                "node_output_id": "7a3406a1-6f11-4568-8aa0-e5dba6534dc2",
+            },
+        }
+    ]
+
+
 class WorkflowInputGenericNode(BaseNode):
     attr: str = Inputs.input
 

vellum_ai-0.13.25/ee/vellum_ee/workflows/display/tests/workflow_serialization/test_basic_inline_prompt_node_serialization.py

@@ -0,0 +1,18 @@
+from vellum_ee.workflows.display.workflows import VellumWorkflowDisplay
+from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
+
+from tests.workflows.basic_inline_prompt_node_with_functions.workflow import BasicInlinePromptWithFunctionsWorkflow
+
+
+def test_serialize_workflow():
+    # WHEN we serialize it
+    workflow_display = get_workflow_display(
+        base_display_class=VellumWorkflowDisplay, workflow_class=BasicInlinePromptWithFunctionsWorkflow
+    )
+    serialized_workflow: dict = workflow_display.serialize()
+    assert (
+        serialized_workflow["workflow_raw_data"]["nodes"][-2]["data"]["exec_config"]["prompt_template_block_data"][
+            "blocks"
+        ][-1]["block_type"]
+        == "FUNCTION_DEFINITION"
+    )

{vellum_ai-0.13.24 → vellum_ai-0.13.25}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "vellum-ai"
-version = "0.13.24"
+version = "0.13.25"
 description = ""
 readme = "README.md"
 authors = []

{vellum_ai-0.13.24 → vellum_ai-0.13.25}/src/vellum/client/core/client_wrapper.py

@@ -18,7 +18,7 @@ class BaseClientWrapper:
         headers: typing.Dict[str, str] = {
             "X-Fern-Language": "Python",
             "X-Fern-SDK-Name": "vellum-ai",
-            "X-Fern-SDK-Version": "0.13.24",
+            "X-Fern-SDK-Version": "0.13.25",
         }
         headers["X_API_KEY"] = self.api_key
         return headers

{vellum_ai-0.13.24 → vellum_ai-0.13.25}/src/vellum/workflows/references/lazy.py

@@ -1,7 +1,8 @@
 import ast
 import inspect
-from typing import TYPE_CHECKING, Callable, Generic, TypeVar, get_args
+from typing import TYPE_CHECKING, Callable, Generic, TypeVar, Union, get_args
 
+from vellum.workflows.constants import UNDEF
 from vellum.workflows.descriptors.base import BaseDescriptor
 
 if TYPE_CHECKING:
@@ -13,7 +14,7 @@ _T = TypeVar("_T")
 class LazyReference(BaseDescriptor[_T], Generic[_T]):
     def __init__(
         self,
-        get: Callable[[], BaseDescriptor[_T]],
+        get: Union[Callable[[], BaseDescriptor[_T]], str],
     ) -> None:
         self._get = get
         # TODO: figure out this some times returns empty
@@ -25,6 +26,19 @@ class LazyReference(BaseDescriptor[_T], Generic[_T]):
     def resolve(self, state: "BaseState") -> _T:
         from vellum.workflows.descriptors.utils import resolve_value
 
+        if isinstance(self._get, str):
+            # The full solution will involve creating a nodes registry on `WorkflowContext`. I want to update
+            # how WorkflowContext works so that we could just access it directly instead of it needing to be
+            # passed in, similar to get_workflow_context(). Because we don't want this to slow down p1 issues
+            # that we are debugging with existing workflows, using the following workaround for now.
+            for output_reference, value in state.meta.node_outputs.items():
+                if str(output_reference) == self._get:
+                    return value
+
+            # Fix typing surrounding the return value of node outputs/output descriptors
+            # https://app.shortcut.com/vellum/story/4783
+            return UNDEF  # type: ignore[return-value]
+
         return resolve_value(self._get(), state)
 
     def _get_name(self) -> str:
@@ -33,6 +47,9 @@ class LazyReference(BaseDescriptor[_T], Generic[_T]):
         setting that as the descriptor's name. Names are only used for debugging, so
         we could flesh out edge cases over time.
         """
+        if isinstance(self._get, str):
+            return self._get
+
         source = inspect.getsource(self._get).strip()
         try:
             parsed = ast.parse(source)