vellum-ai 0.12.7__py3-none-any.whl → 0.12.9__py3-none-any.whl

vellum/workflows/nodes/core/templating_node/tests/test_templating_node.py

@@ -2,6 +2,8 @@ import json
 
 from vellum.workflows.nodes.bases.base import BaseNode
 from vellum.workflows.nodes.core.templating_node.node import TemplatingNode
+from vellum.workflows.state import BaseState
+from vellum.workflows.types.core import Json
 
 
 def test_templating_node__dict_output():
@@ -22,6 +24,70 @@ def test_templating_node__dict_output():
     assert json.loads(outputs.result) == {"key": "value"}
 
 
+def test_templating_node__int_output():
+    # GIVEN a templating node that outputs an integer
+    class IntTemplateNode(TemplatingNode[BaseState, int]):
+        template = "{{ data }}"
+        inputs = {
+            "data": 42,
+        }
+
+    # WHEN the node is run
+    node = IntTemplateNode()
+    outputs = node.run()
+
+    # THEN the output is the expected integer
+    assert outputs.result == 42
+
+
+def test_templating_node__float_output():
+    # GIVEN a templating node that outputs a float
+    class FloatTemplateNode(TemplatingNode[BaseState, float]):
+        template = "{{ data }}"
+        inputs = {
+            "data": 42.5,
+        }
+
+    # WHEN the node is run
+    node = FloatTemplateNode()
+    outputs = node.run()
+
+    # THEN the output is the expected float
+    assert outputs.result == 42.5
+
+
+def test_templating_node__bool_output():
+    # GIVEN a templating node that outputs a bool
+    class BoolTemplateNode(TemplatingNode[BaseState, bool]):
+        template = "{{ data }}"
+        inputs = {
+            "data": True,
+        }
+
+    # WHEN the node is run
+    node = BoolTemplateNode()
+    outputs = node.run()
+
+    # THEN the output is the expected bool
+    assert outputs.result is True
+
+
+def test_templating_node__json_output():
+    # GIVEN a templating node that outputs JSON
+    class JSONTemplateNode(TemplatingNode[BaseState, Json]):
+        template = "{{ data }}"
+        inputs = {
+            "data": {"key": "value"},
+        }
+
+    # WHEN the node is run
+    node = JSONTemplateNode()
+    outputs = node.run()
+
+    # THEN the output is the expected JSON
+    assert outputs.result == {"key": "value"}
+
+
 def test_templating_node__execution_count_reference():
     # GIVEN a random node
     class OtherNode(BaseNode):

vellum/workflows/nodes/core/try_node/node.py

@@ -73,9 +73,7 @@ class TryNode(BaseNode[StateType], Generic[StateType], metaclass=_TryNodeMeta):
     def run(self) -> Iterator[BaseOutput]:
         subworkflow = self.subworkflow(
             parent_state=self.state,
-            context=WorkflowContext(
-                _vellum_client=self._context._vellum_client,
-            ),
+            context=WorkflowContext(vellum_client=self._context.vellum_client),
         )
         subworkflow_stream = subworkflow.stream(
             event_filter=all_workflow_event_filter,

vellum/workflows/nodes/core/try_node/tests/test_node.py

@@ -103,7 +103,7 @@ def test_try_node__use_parent_execution_context():
     # WHEN the node is run with a custom vellum client
     node = TestNode(
         context=WorkflowContext(
-            _vellum_client=Vellum(api_key="test-key"),
+            vellum_client=Vellum(api_key="test-key"),
         )
     )
     outputs = list(node.run())

vellum/workflows/nodes/displayable/bases/api_node/node.py

@@ -26,11 +26,11 @@ class BaseAPINode(BaseNode, Generic[StateType]):
     url: str
     method: APIRequestMethod
     data: Optional[str] = None
-    json: Optional["Json"] = None
+    json: Optional[Json] = None
     headers: Optional[Dict[str, Union[str, VellumSecret]]] = None
 
     class Outputs(BaseOutputs):
-        json: Optional["Json"]
+        json: Optional[Json]
         headers: Dict[str, str]
         status_code: int
         text: str

vellum/workflows/nodes/displayable/bases/search_node.py

@@ -12,6 +12,7 @@ from vellum import (
     SearchWeightsRequest,
 )
 from vellum.core import ApiError, RequestOptions
+from vellum.utils.uuid import is_valid_uuid
 from vellum.workflows.errors import WorkflowErrorCode
 from vellum.workflows.exceptions import NodeException
 from vellum.workflows.nodes.bases import BaseNode
@@ -73,11 +74,13 @@ class BaseSearchNode(BaseNode[StateType], Generic[StateType]):
         results: List[SearchResult]
 
     def _perform_search(self) -> SearchResponse:
+        index_is_uuid = True if isinstance(self.document_index, UUID) else is_valid_uuid(self.document_index)
+
         try:
             return self._context.vellum_client.search(
                 query=self.query,
-                index_id=str(self.document_index) if isinstance(self.document_index, UUID) else None,
-                index_name=self.document_index if isinstance(self.document_index, str) else None,
+                index_id=str(self.document_index) if index_is_uuid else None,
+                index_name=str(self.document_index) if not index_is_uuid else None,
                 options=self.options,
             )
         except NotFoundError:

vellum/workflows/nodes/displayable/subworkflow_deployment_node/node.py

@@ -17,12 +17,13 @@ from vellum.workflows.context import get_parent_context
 from vellum.workflows.errors import WorkflowErrorCode
 from vellum.workflows.errors.types import workflow_event_error_to_workflow_error
 from vellum.workflows.exceptions import NodeException
-from vellum.workflows.nodes.bases.base_subworkflow_node.node import BaseSubworkflowNode
+from vellum.workflows.nodes.bases.base import BaseNode
 from vellum.workflows.outputs.base import BaseOutput
+from vellum.workflows.types.core import EntityInputsInterface
 from vellum.workflows.types.generics import StateType
 
 
-class SubworkflowDeploymentNode(BaseSubworkflowNode[StateType], Generic[StateType]):
+class SubworkflowDeploymentNode(BaseNode[StateType], Generic[StateType]):
     """
     Used to execute a Workflow Deployment.
 
@@ -38,6 +39,7 @@ class SubworkflowDeploymentNode(BaseSubworkflowNode[StateType], Generic[StateTyp
 
     # Either the Workflow Deployment's UUID or its name.
     deployment: ClassVar[Union[UUID, str]]
+    subworkflow_inputs: ClassVar[EntityInputsInterface] = {}
 
     release_tag: str = LATEST_RELEASE_TAG
     external_id: Optional[str] = OMIT

vellum/workflows/nodes/experimental/README.md

@@ -0,0 +1,6 @@
+# 🧪 Experimental
+
+This section is a proofing ground for new ideas and concepts. It's not meant to be used in production and is
+subject to breaking changes at any time.
+
+If a concept within is sufficiently interesting and validated, then it may be introduced as a first-class node.

vellum/workflows/nodes/experimental/openai_chat_completion_node/__init__.py

@@ -0,0 +1,5 @@
+from .node import OpenAIChatCompletionNode
+
+__all__ = [
+    "OpenAIChatCompletionNode",
+]

vellum/workflows/nodes/experimental/openai_chat_completion_node/node.py

@@ -0,0 +1,260 @@
+import json
+import logging
+import os
+from uuid import uuid4
+from typing import Any, Iterable, Iterator, List, Literal, Union, cast
+
+from openai import OpenAI
+from openai.types.chat import (
+    ChatCompletionAssistantMessageParam,
+    ChatCompletionContentPartImageParam,
+    ChatCompletionContentPartInputAudioParam,
+    ChatCompletionContentPartParam,
+    ChatCompletionContentPartRefusalParam,
+    ChatCompletionContentPartTextParam,
+    ChatCompletionMessageParam,
+    ChatCompletionSystemMessageParam,
+    ChatCompletionUserMessageParam,
+)
+from openai.types.chat.chat_completion_chunk import Choice
+
+from vellum import (
+    AdHocExecutePromptEvent,
+    FulfilledAdHocExecutePromptEvent,
+    InitiatedAdHocExecutePromptEvent,
+    RejectedAdHocExecutePromptEvent,
+    StreamingAdHocExecutePromptEvent,
+    StringVellumValue,
+    VellumAudio,
+    VellumError,
+)
+from vellum.prompts.blocks.compilation import compile_prompt_blocks
+from vellum.prompts.blocks.types import CompiledChatMessagePromptBlock
+from vellum.workflows.errors import WorkflowErrorCode
+from vellum.workflows.exceptions import NodeException
+from vellum.workflows.nodes import InlinePromptNode
+from vellum.workflows.types.generics import StateType
+
+logger = logging.getLogger(__name__)
+
+
+class OpenAIChatCompletionNode(InlinePromptNode[StateType]):
+    """
+    Used to execute a Prompt using the OpenAI API.
+    """
+
+    # Override
+    def _get_prompt_event_stream(self) -> Iterator[AdHocExecutePromptEvent]:
+        client = self._get_client()
+
+        execution_id = str(uuid4())
+
+        yield InitiatedAdHocExecutePromptEvent(
+            execution_id=execution_id,
+        )
+
+        try:
+            stream = client.chat.completions.create(
+                messages=self._get_messages(),
+                model=self.ml_model,
+                # TODO: Add support for additional parameters
+                stream=True,
+            )
+        except Exception as exc:
+            yield RejectedAdHocExecutePromptEvent(
+                error=VellumError(
+                    code=WorkflowErrorCode.PROVIDER_ERROR,
+                    message=exc.args[0],
+                ),
+                execution_id=execution_id,
+            )
+            return
+
+        combined_delta_content = ""
+        for chunk in stream:
+            choices: List[Choice] = chunk.choices
+            if len(choices) != 1:
+                yield RejectedAdHocExecutePromptEvent(
+                    error=VellumError(
+                        code=WorkflowErrorCode.PROVIDER_ERROR,
+                        message="Expected one choice per chunk, but found more than one.",
+                    ),
+                    execution_id=execution_id,
+                )
+                return
+
+            choice = choices[0]
+            delta = choice.delta
+
+            if delta.tool_calls:
+                # TODO: Add support for tool calls
+                raise NotImplementedError("This node hasn't been extended to support tool calling yet.")
+
+            if delta.content:
+                combined_delta_content += delta.content
+
+            StreamingAdHocExecutePromptEvent(
+                output=StringVellumValue(value=delta.content),
+                # TODO: Add support for multiple outputs
+                output_index=1,
+                execution_id=execution_id,
+            )
+
+        yield FulfilledAdHocExecutePromptEvent(
+            # TODO: Add support for multiple outputs
+            outputs=[
+                StringVellumValue(value=combined_delta_content),
+            ],
+            execution_id=execution_id,
+        )
+
+    def _get_client(self) -> OpenAI:
+        """Used to retrieve an API client for interacting with the OpenAI API.
+
+        Note: This method can be overridden if you'd like to use your own API client that conforms to the same
+            interfaces as that of OpenAI.
+        """
+
+        openai_api_key = os.environ.get("OPENAI_API_KEY")
+
+        if not openai_api_key:
+            raise NodeException(
+                code=WorkflowErrorCode.INTERNAL_ERROR,
+                message="Unable to determine an OpenAI API key.",
+            )
+
+        client = OpenAI(api_key=openai_api_key)
+        return client
+
+    def _get_messages(self) -> Iterable[ChatCompletionMessageParam]:
+        input_variables, input_values = self._compile_prompt_inputs()
+
+        compiled_blocks = compile_prompt_blocks(
+            blocks=self.blocks, inputs=input_values, input_variables=input_variables
+        )
+
+        chat_message_blocks: list[CompiledChatMessagePromptBlock] = [
+            block for block in compiled_blocks if block.block_type == "CHAT_MESSAGE"
+        ]
+        messages = [self._create_message(block) for block in chat_message_blocks]
+
+        return messages
+
+    @classmethod
+    def _create_message(cls, chat_message_block: CompiledChatMessagePromptBlock) -> ChatCompletionMessageParam:
+        name = chat_message_block.source
+        content = cls._create_message_content(chat_message_block)
+
+        if chat_message_block.role == "SYSTEM":
+            relevant_system_content = [
+                cast(ChatCompletionContentPartTextParam, c) for c in content if c["type"] == "text"
+            ]
+            system_message: ChatCompletionSystemMessageParam = {
+                "content": relevant_system_content,
+                "role": "system",
+            }
+            if name:
+                system_message["name"] = name
+
+            return system_message
+        elif chat_message_block.role == "USER":
+            user_message: ChatCompletionUserMessageParam = {
+                "content": content,
+                "role": "user",
+            }
+            if name:
+                user_message["name"] = name
+
+            return user_message
+        elif chat_message_block.role == "ASSISTANT":
+            relevant_assistant_content = [
+                cast(Union[ChatCompletionContentPartTextParam, ChatCompletionContentPartRefusalParam], c)
+                for c in content
+                if c["type"] in ["text", "refusal"]
+            ]
+            assistant_message: ChatCompletionAssistantMessageParam = {
+                "content": relevant_assistant_content,
+                "role": "assistant",
+            }
+            if name:
+                assistant_message["name"] = name
+
+            return assistant_message
+        else:
+            logger.error(f"Unexpected role: {chat_message_block.role}")
+            raise NodeException(
+                code=WorkflowErrorCode.INTERNAL_ERROR, message="Unexpected role found when compiling prompt blocks"
+            )
+
+    @classmethod
+    def _create_message_content(
+        cls,
+        chat_message_block: CompiledChatMessagePromptBlock,
+    ) -> List[ChatCompletionContentPartParam]:
+        content: List[ChatCompletionContentPartParam] = []
+        for block in chat_message_block.blocks:
+            if block.content.type == "STRING":
+                string_value = cast(str, block.content.value)
+                string_content_item: ChatCompletionContentPartTextParam = {"type": "text", "text": string_value}
+                content.append(string_content_item)
+            elif block.content.type == "JSON":
+                json_value = cast(Any, block.content.value)
+                json_content_item: ChatCompletionContentPartTextParam = {"type": "text", "text": json.dumps(json_value)}
+                content.append(json_content_item)
+            elif block.content.type == "IMAGE":
+                image_value = cast(VellumAudio, block.content.value)
+                image_content_item: ChatCompletionContentPartImageParam = {
+                    "type": "image_url",
+                    "image_url": {"url": image_value.src},
+                }
+                if image_value.metadata and image_value.metadata.get("detail"):
+                    detail = image_value.metadata["detail"]
+
+                    if detail not in ["auto", "low", "high"]:
+                        raise NodeException(
+                            code=WorkflowErrorCode.INTERNAL_ERROR,
+                            message="Image detail must be one of 'auto', 'low', or 'high.",
+                        )
+
+                    image_content_item["image_url"]["detail"] = cast(Literal["auto", "low", "high"], detail)
+
+                content.append(image_content_item)
+            elif block.content.type == "AUDIO":
+                audio_value = cast(VellumAudio, block.content.value)
+                audio_value_src_parts = audio_value.src.split(",")
+                if len(audio_value_src_parts) != 2:
+                    raise NodeException(
+                        code=WorkflowErrorCode.INTERNAL_ERROR, message="Audio data is not properly encoded."
+                    )
+                _, cleaned_audio_value = audio_value_src_parts
+                if not audio_value.metadata:
+                    raise NodeException(
+                        code=WorkflowErrorCode.INTERNAL_ERROR, message="Audio metadata is required for audio input."
+                    )
+                audio_format = audio_value.metadata.get("format")
+                if not audio_format:
+                    raise NodeException(
+                        code=WorkflowErrorCode.INTERNAL_ERROR, message="Audio format is required for audio input."
+                    )
+                if audio_format not in {"wav", "mp3"}:
+                    raise NodeException(
+                        code=WorkflowErrorCode.INTERNAL_ERROR,
+                        message="Audio format must be one of 'wav' or 'mp3'.",
+                    )
+
+                audio_content_item: ChatCompletionContentPartInputAudioParam = {
+                    "type": "input_audio",
+                    "input_audio": {
+                        "data": cleaned_audio_value,
+                        "format": cast(Literal["wav", "mp3"], audio_format),
+                    },
+                }
+
+                content.append(audio_content_item)
+            else:
+                raise NodeException(
+                    code=WorkflowErrorCode.INTERNAL_ERROR,
+                    message=f"Failed to parse chat message block {block.content.type}",
+                )
+
+        return content

vellum/workflows/sandbox.py

@@ -1,4 +1,4 @@
-from typing import Generic, Sequence, Type
+from typing import Generic, Sequence
 
 import dotenv
 
@@ -9,8 +9,8 @@ from vellum.workflows.types.generics import WorkflowType
 from vellum.workflows.workflows.event_filters import root_workflow_event_filter
 
 
-class SandboxRunner(Generic[WorkflowType]):
-    def __init__(self, workflow: Type[WorkflowType], inputs: Sequence[BaseInputs]):
+class WorkflowSandboxRunner(Generic[WorkflowType]):
+    def __init__(self, workflow: WorkflowType, inputs: Sequence[BaseInputs]):
         if not inputs:
             raise ValueError("Inputs are required to have at least one defined inputs")
 
@@ -30,8 +30,7 @@ class SandboxRunner(Generic[WorkflowType]):
 
         selected_inputs = self._inputs[index]
 
-        workflow = self._workflow()
-        events = workflow.stream(
+        events = self._workflow.stream(
             inputs=selected_inputs,
             event_filter=root_workflow_event_filter,
         )
@@ -49,3 +48,6 @@ class SandboxRunner(Generic[WorkflowType]):
                 self._logger.info(f"Just started Node: {event.node_definition.__name__}")
             elif event.name == "node.execution.fulfilled":
                 self._logger.info(f"Just finished Node: {event.node_definition.__name__}")
+            elif event.name == "node.execution.rejected":
+                self._logger.debug(f"Error: {event.error}")
+                self._logger.error(f"Failed to run Node: {event.node_definition.__name__}")

vellum/workflows/state/context.py

@@ -13,11 +13,12 @@ if TYPE_CHECKING:
 class WorkflowContext:
     def __init__(
         self,
-        _vellum_client: Optional[Vellum] = None,
-        _parent_context: Optional[ParentContext] = None,
+        *,
+        vellum_client: Optional[Vellum] = None,
+        parent_context: Optional[ParentContext] = None,
     ):
-        self._vellum_client = _vellum_client
-        self._parent_context = _parent_context
+        self._vellum_client = vellum_client
+        self._parent_context = parent_context
         self._event_queue: Optional[Queue["WorkflowEvent"]] = None
 
     @cached_property

vellum/workflows/tests/test_sandbox.py

@@ -3,7 +3,7 @@ from typing import List
 
 from vellum.workflows.inputs.base import BaseInputs
 from vellum.workflows.nodes.bases.base import BaseNode
-from vellum.workflows.sandbox import SandboxRunner
+from vellum.workflows.sandbox import WorkflowSandboxRunner
 from vellum.workflows.state.base import BaseState
 from vellum.workflows.workflows.base import BaseWorkflow
 
@@ -49,7 +49,7 @@ def test_sandbox_runner__happy_path(mock_logger, run_kwargs, expected_last_log):
     ]
 
     # WHEN we run the sandbox
-    runner = SandboxRunner(workflow=Workflow, inputs=inputs)
+    runner = WorkflowSandboxRunner(workflow=Workflow(), inputs=inputs)
     runner.run(**run_kwargs)
 
     # THEN we see the logs

vellum/workflows/utils/tests/test_vellum_variables.py

@@ -2,6 +2,7 @@ import pytest
 from typing import List, Optional
 
 from vellum import ChatMessage, SearchResult
+from vellum.workflows.types.core import Json
 from vellum.workflows.utils.vellum_variables import primitive_type_to_vellum_variable_type
 
 
@@ -18,6 +19,8 @@ from vellum.workflows.utils.vellum_variables import primitive_type_to_vellum_var
         (Optional[List[ChatMessage]], "CHAT_HISTORY"),
         (List[SearchResult], "SEARCH_RESULTS"),
         (Optional[List[SearchResult]], "SEARCH_RESULTS"),
+        (Json, "JSON"),
+        (Optional[Json], "JSON"),
     ],
 )
 def test_primitive_type_to_vellum_variable_type(type_, expected):

vellum/workflows/utils/vellum_variables.py

@@ -35,14 +35,15 @@ def primitive_type_to_vellum_variable_type(type_: Union[Type, BaseDescriptor]) -
         if len(types) != 1:
             # Check explicitly for our internal JSON type.
             # Matches the type found at vellum.workflows.utils.vellum_variables.Json
-            if types == [
+            actual_types_with_explicit_ref = [
                 bool,
                 int,
                 float,
                 str,
-                typing.List[typing.ForwardRef("Json")],  # type: ignore [misc]
-                typing.Dict[str, typing.ForwardRef("Json")],  # type: ignore [misc]
-            ]:
+                typing.List[Json],
+                typing.Dict[str, Json],
+            ]
+            if types == actual_types_with_explicit_ref:
                 return "JSON"
             raise ValueError(f"Expected Descriptor to only have one type, got {types}")
 

vellum/workflows/workflows/base.py

@@ -121,10 +121,11 @@ class BaseWorkflow(Generic[WorkflowInputsType, StateType], metaclass=_BaseWorkfl
 
     def __init__(
         self,
+        *,
+        context: Optional[WorkflowContext] = None,
         parent_state: Optional[BaseState] = None,
         emitters: Optional[List[BaseWorkflowEmitter]] = None,
         resolvers: Optional[List[BaseWorkflowResolver]] = None,
-        context: Optional[WorkflowContext] = None,
     ):
         self._parent_state = parent_state
         self.emitters = emitters or (self.emitters if hasattr(self, "emitters") else [])
@@ -188,6 +189,7 @@ class BaseWorkflow(Generic[WorkflowInputsType, StateType], metaclass=_BaseWorkfl
     def run(
         self,
         inputs: Optional[WorkflowInputsType] = None,
+        *,
         state: Optional[StateType] = None,
         entrypoint_nodes: Optional[RunFromNodeArg] = None,
         external_inputs: Optional[ExternalInputsArg] = None,
@@ -281,8 +283,9 @@ class BaseWorkflow(Generic[WorkflowInputsType, StateType], metaclass=_BaseWorkfl
 
     def stream(
         self,
-        event_filter: Optional[Callable[[Type["BaseWorkflow"], WorkflowEvent], bool]] = None,
         inputs: Optional[WorkflowInputsType] = None,
+        *,
+        event_filter: Optional[Callable[[Type["BaseWorkflow"], WorkflowEvent], bool]] = None,
         state: Optional[StateType] = None,
         entrypoint_nodes: Optional[RunFromNodeArg] = None,
         external_inputs: Optional[ExternalInputsArg] = None,

{vellum_ai-0.12.7.dist-info → vellum_ai-0.12.9.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vellum-ai
-Version: 0.12.7
+Version: 0.12.9
 Summary: 
 License: MIT
 Requires-Python: >=3.9,<4.0
@@ -25,6 +25,7 @@ Requires-Dist: cdktf (>=0.20.5,<0.21.0)
 Requires-Dist: click (==8.1.7)
 Requires-Dist: docker (==7.1.0)
 Requires-Dist: httpx (>=0.21.2)
+Requires-Dist: openai (>=1.0.0)
 Requires-Dist: orderly-set (>=5.2.2,<6.0.0)
 Requires-Dist: publication (==0.0.3)
 Requires-Dist: pydantic (>=1.9.2)