vellum-ai 0.14.16__py3-none-any.whl → 0.14.18__py3-none-any.whl

vellum/workflows/nodes/displayable/bases/inline_prompt_node/tests/test_inline_prompt_node.py

@@ -0,0 +1,298 @@
+import pytest
+import json
+from unittest import mock
+from uuid import uuid4
+from typing import Any, Iterator, List
+
+from vellum import (
+    ChatMessagePromptBlock,
+    JinjaPromptBlock,
+    PlainTextPromptBlock,
+    PromptBlock,
+    PromptParameters,
+    RichTextPromptBlock,
+    VariablePromptBlock,
+)
+from vellum.client.types.execute_prompt_event import ExecutePromptEvent
+from vellum.client.types.fulfilled_execute_prompt_event import FulfilledExecutePromptEvent
+from vellum.client.types.initiated_execute_prompt_event import InitiatedExecutePromptEvent
+from vellum.client.types.prompt_output import PromptOutput
+from vellum.client.types.prompt_request_string_input import PromptRequestStringInput
+from vellum.client.types.string_vellum_value import StringVellumValue
+from vellum.workflows.errors import WorkflowErrorCode
+from vellum.workflows.exceptions import NodeException
+from vellum.workflows.inputs import BaseInputs
+from vellum.workflows.nodes import InlinePromptNode
+from vellum.workflows.nodes.displayable.bases.inline_prompt_node import BaseInlinePromptNode
+from vellum.workflows.state import BaseState
+from vellum.workflows.state.base import StateMeta
+
+
+def test_validation_with_missing_variables():
+    """Test that validation correctly identifies missing variables."""
+    test_blocks: List[PromptBlock] = [
+        VariablePromptBlock(input_variable="required_var1"),
+        VariablePromptBlock(input_variable="required_var2"),
+        RichTextPromptBlock(
+            blocks=[
+                PlainTextPromptBlock(text="Some text"),
+                VariablePromptBlock(input_variable="required_var3"),
+            ],
+        ),
+        JinjaPromptBlock(template="Template without variables"),
+        ChatMessagePromptBlock(
+            chat_role="USER",
+            blocks=[
+                RichTextPromptBlock(
+                    blocks=[
+                        PlainTextPromptBlock(text="Nested text"),
+                        VariablePromptBlock(input_variable="required_var4"),
+                    ],
+                ),
+            ],
+        ),
+    ]
+
+    # GIVEN a BaseInlinePromptNode
+    class TestNode(BaseInlinePromptNode):
+        ml_model = "test-model"
+        blocks = test_blocks
+        prompt_inputs = {
+            "required_var1": "value1",
+            # required_var2 is missing
+            # required_var3 is missing
+            # required_var4 is missing
+        }
+
+    # WHEN the node is run
+    node = TestNode()
+    with pytest.raises(NodeException) as excinfo:
+        list(node.run())
+
+    # THEN the node raises the correct NodeException
+    assert excinfo.value.code == WorkflowErrorCode.INVALID_INPUTS
+    assert "required_var2" in str(excinfo.value)
+    assert "required_var3" in str(excinfo.value)
+    assert "required_var4" in str(excinfo.value)
+
+
+def test_validation_with_all_variables_provided(vellum_adhoc_prompt_client):
+    """Test that validation passes when all variables are provided."""
+    test_blocks: List[PromptBlock] = [
+        VariablePromptBlock(input_variable="required_var1"),
+        VariablePromptBlock(input_variable="required_var2"),
+        RichTextPromptBlock(
+            blocks=[
+                PlainTextPromptBlock(text="Some text"),
+                VariablePromptBlock(input_variable="required_var3"),
+            ],
+        ),
+        JinjaPromptBlock(template="Template without variables"),
+        ChatMessagePromptBlock(
+            chat_role="USER",
+            blocks=[
+                RichTextPromptBlock(
+                    blocks=[
+                        PlainTextPromptBlock(text="Nested text"),
+                        VariablePromptBlock(input_variable="required_var4"),
+                    ],
+                ),
+            ],
+        ),
+    ]
+
+    # GIVEN a BaseInlinePromptNode
+    class TestNode(BaseInlinePromptNode):
+        ml_model = "test-model"
+        blocks = test_blocks
+        prompt_inputs = {
+            "required_var1": "value1",
+            "required_var2": "value2",
+            "required_var3": "value3",
+            "required_var4": "value4",
+        }
+
+    expected_outputs: List[PromptOutput] = [
+        StringVellumValue(value="Test response"),
+    ]
+
+    def generate_prompt_events(*args: Any, **kwargs: Any) -> Iterator[ExecutePromptEvent]:
+        execution_id = str(uuid4())
+        events: List[ExecutePromptEvent] = [
+            InitiatedExecutePromptEvent(execution_id=execution_id),
+            FulfilledExecutePromptEvent(
+                execution_id=execution_id,
+                outputs=expected_outputs,
+            ),
+        ]
+        yield from events
+
+    vellum_adhoc_prompt_client.adhoc_execute_prompt_stream.side_effect = generate_prompt_events
+
+    # WHEN the node is run
+    node = TestNode()
+    list(node.run())
+
+    # THEN the prompt is executed with the correct inputs
+    mock_api = vellum_adhoc_prompt_client.adhoc_execute_prompt_stream
+    assert mock_api.call_count == 1
+    assert mock_api.call_args.kwargs["input_values"] == [
+        PromptRequestStringInput(key="required_var1", type="STRING", value="value1"),
+        PromptRequestStringInput(key="required_var2", type="STRING", value="value2"),
+        PromptRequestStringInput(key="required_var3", type="STRING", value="value3"),
+        PromptRequestStringInput(key="required_var4", type="STRING", value="value4"),
+    ]
+
+
+def test_validation_with_extra_variables(vellum_adhoc_prompt_client):
+    """Test that validation passes when extra variables are provided."""
+    test_blocks: List[PromptBlock] = [
+        VariablePromptBlock(input_variable="required_var"),
+    ]
+
+    # GIVEN a BaseInlinePromptNode
+    class TestNode(BaseInlinePromptNode):
+        ml_model = "test-model"
+        blocks = test_blocks
+        prompt_inputs = {
+            "required_var": "value",
+            "extra_var": "extra_value",  # This is not required
+        }
+
+    expected_outputs: List[PromptOutput] = [
+        StringVellumValue(value="Test response"),
+    ]
+
+    def generate_prompt_events(*args: Any, **kwargs: Any) -> Iterator[ExecutePromptEvent]:
+        execution_id = str(uuid4())
+        events: List[ExecutePromptEvent] = [
+            InitiatedExecutePromptEvent(execution_id=execution_id),
+            FulfilledExecutePromptEvent(
+                execution_id=execution_id,
+                outputs=expected_outputs,
+            ),
+        ]
+        yield from events
+
+    vellum_adhoc_prompt_client.adhoc_execute_prompt_stream.side_effect = generate_prompt_events
+
+    # WHEN the node is run
+    node = TestNode()
+    list(node.run())
+
+    # THEN the prompt is executed with the correct inputs
+    mock_api = vellum_adhoc_prompt_client.adhoc_execute_prompt_stream
+    assert mock_api.call_count == 1
+    assert mock_api.call_args.kwargs["input_values"] == [
+        PromptRequestStringInput(key="required_var", type="STRING", value="value"),
+        PromptRequestStringInput(key="extra_var", type="STRING", value="extra_value"),
+    ]
+
+
+def test_inline_prompt_node__json_output(vellum_adhoc_prompt_client):
+    """Confirm that InlinePromptNodes output the expected JSON when run."""
+
+    # GIVEN a node that subclasses InlinePromptNode
+    class Inputs(BaseInputs):
+        input: str
+
+    class State(BaseState):
+        pass
+
+    class MyInlinePromptNode(InlinePromptNode):
+        ml_model = "gpt-4o"
+        blocks = []
+        parameters = PromptParameters(
+            stop=[],
+            temperature=0.0,
+            max_tokens=4096,
+            top_p=1.0,
+            top_k=0,
+            frequency_penalty=0.0,
+            presence_penalty=0.0,
+            logit_bias=None,
+            custom_parameters={
+                "json_mode": False,
+                "json_schema": {
+                    "name": "get_result",
+                    "schema": {
+                        "type": "object",
+                        "required": ["result"],
+                        "properties": {"result": {"type": "string", "description": ""}},
+                    },
+                },
+            },
+        )
+
+    # AND a known JSON response from invoking an inline prompt
+    expected_json = {"result": "Hello, world!"}
+    expected_outputs: List[PromptOutput] = [
+        StringVellumValue(value=json.dumps(expected_json)),
+    ]
+
+    def generate_prompt_events(*args: Any, **kwargs: Any) -> Iterator[ExecutePromptEvent]:
+        execution_id = str(uuid4())
+        events: List[ExecutePromptEvent] = [
+            InitiatedExecutePromptEvent(execution_id=execution_id),
+            FulfilledExecutePromptEvent(
+                execution_id=execution_id,
+                outputs=expected_outputs,
+            ),
+        ]
+        yield from events
+
+    vellum_adhoc_prompt_client.adhoc_execute_prompt_stream.side_effect = generate_prompt_events
+
+    # WHEN the node is run
+    node = MyInlinePromptNode(
+        state=State(
+            meta=StateMeta(workflow_inputs=Inputs(input="Generate JSON.")),
+        )
+    )
+    outputs = [o for o in node.run()]
+
+    # THEN the node should have produced the outputs we expect
+    results_output = outputs[0]
+    assert results_output.name == "results"
+    assert results_output.value == expected_outputs
+
+    text_output = outputs[1]
+    assert text_output.name == "text"
+    assert text_output.value == '{"result": "Hello, world!"}'
+
+    json_output = outputs[2]
+    assert json_output.name == "json"
+    assert json_output.value == expected_json
+
+    # AND we should have made the expected call to Vellum search
+    vellum_adhoc_prompt_client.adhoc_execute_prompt_stream.assert_called_once_with(
+        blocks=[],
+        expand_meta=Ellipsis,
+        functions=None,
+        input_values=[],
+        input_variables=[],
+        ml_model="gpt-4o",
+        parameters=PromptParameters(
+            stop=[],
+            temperature=0.0,
+            max_tokens=4096,
+            top_p=1.0,
+            top_k=0,
+            frequency_penalty=0.0,
+            presence_penalty=0.0,
+            logit_bias=None,
+            custom_parameters={
+                "json_mode": False,
+                "json_schema": {
+                    "name": "get_result",
+                    "schema": {
+                        "type": "object",
+                        "required": ["result"],
+                        "properties": {"result": {"type": "string", "description": ""}},
+                    },
+                },
+            },
+        ),
+        request_options=mock.ANY,
+        settings=None,
+    )

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

@@ -1,6 +1,7 @@
 import json
-from typing import Iterator
+from typing import Any, Dict, Iterator, Type, Union
 
+from vellum.workflows.constants import undefined
 from vellum.workflows.errors import WorkflowErrorCode
 from vellum.workflows.exceptions import NodeException
 from vellum.workflows.nodes.displayable.bases import BaseInlinePromptNode as BaseInlinePromptNode
@@ -30,9 +31,11 @@ class InlinePromptNode(BaseInlinePromptNode[StateType]):
         The outputs of the InlinePromptNode.
 
         text: str - The result of the Prompt Execution
+        json: Optional[Dict[Any, Any]] - The result of the Prompt Execution in JSON format
         """
 
         text: str
+        json: Union[Dict[Any, Any], Type[undefined]] = undefined
 
     def run(self) -> Iterator[BaseOutput]:
         outputs = yield from self._process_prompt_event_stream()
@@ -43,14 +46,31 @@ class InlinePromptNode(BaseInlinePromptNode[StateType]):
             )
 
         string_outputs = []
+        json_output = None
+
+        should_parse_json = False
+        if hasattr(self, "parameters"):
+            custom_params = self.parameters.custom_parameters
+            if custom_params and isinstance(custom_params, dict):
+                json_schema = custom_params.get("json_schema", {})
+                if (isinstance(json_schema, dict) and "schema" in json_schema) or custom_params.get("json_mode", {}):
+                    should_parse_json = True
+
         for output in outputs:
             if output.value is None:
                 continue
 
             if output.type == "STRING":
                 string_outputs.append(output.value)
+                if should_parse_json:
+                    try:
+                        parsed_json = json.loads(output.value)
+                        json_output = parsed_json
+                    except (json.JSONDecodeError, TypeError):
+                        pass
             elif output.type == "JSON":
                 string_outputs.append(json.dumps(output.value, indent=4))
+                json_output = output.value
             elif output.type == "FUNCTION_CALL":
                 string_outputs.append(output.value.model_dump_json(indent=4))
             else:
@@ -58,3 +78,6 @@ class InlinePromptNode(BaseInlinePromptNode[StateType]):
 
         value = "\n".join(string_outputs)
         yield BaseOutput(name="text", value=value)
+
+        if json_output:
+            yield BaseOutput(name="json", value=json_output)

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

@@ -27,6 +27,7 @@ from vellum import (
     StringVellumValue,
     VellumAudio,
     VellumError,
+    VellumImage,
 )
 from vellum.prompts.blocks.compilation import compile_prompt_blocks
 from vellum.prompts.blocks.types import CompiledChatMessagePromptBlock
@@ -202,7 +203,7 @@ class OpenAIChatCompletionNode(InlinePromptNode[StateType]):
                 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_value = cast(VellumImage, block.content.value)
                 image_content_item: ChatCompletionContentPartImageParam = {
                     "type": "image_url",
                     "image_url": {"url": image_value.src},
@@ -251,6 +252,11 @@ class OpenAIChatCompletionNode(InlinePromptNode[StateType]):
                 }
 
                 content.append(audio_content_item)
+            elif block.content.type == "DOCUMENT":
+                raise NodeException(
+                    code=WorkflowErrorCode.PROVIDER_ERROR,
+                    message="Document chat message content type is not currently supported",
+                )
             else:
                 raise NodeException(
                     code=WorkflowErrorCode.INTERNAL_ERROR,

vellum/workflows/runner/runner.py

@@ -43,7 +43,7 @@ from vellum.workflows.events.workflow import (
     WorkflowExecutionSnapshottedEvent,
     WorkflowExecutionStreamingBody,
 )
-from vellum.workflows.exceptions import NodeException
+from vellum.workflows.exceptions import NodeException, WorkflowInitializationException
 from vellum.workflows.nodes.bases import BaseNode
 from vellum.workflows.nodes.bases.base import NodeRunResponse
 from vellum.workflows.nodes.mocks import MockNodeExecutionArg
@@ -332,6 +332,18 @@ class WorkflowRunner(Generic[StateType]):
                     parent=parent_context,
                 )
             )
+        except WorkflowInitializationException as e:
+            self._workflow_event_inner_queue.put(
+                NodeExecutionRejectedEvent(
+                    trace_id=node.state.meta.trace_id,
+                    span_id=span_id,
+                    body=NodeExecutionRejectedBody(
+                        node_definition=node.__class__,
+                        error=e.error,
+                    ),
+                    parent=parent_context,
+                )
+            )
         except Exception as e:
             logger.exception(f"An unexpected error occurred while running node {node.__class__.__name__}")
 
@@ -563,6 +575,9 @@ class WorkflowRunner(Generic[StateType]):
             except NodeException as e:
                 self._workflow_event_outer_queue.put(self._reject_workflow_event(e.error))
                 return
+            except WorkflowInitializationException as e:
+                self._workflow_event_outer_queue.put(self._reject_workflow_event(e.error))
+                return
             except Exception:
                 err_message = f"An unexpected error occurred while initializing node {node_cls.__name__}"
                 logger.exception(err_message)

vellum/workflows/utils/tests/test_vellum_variables.py

@@ -1,7 +1,7 @@
 import pytest
 from typing import List, Optional
 
-from vellum import ChatMessage, SearchResult
+from vellum import ChatMessage, SearchResult, VellumAudio, VellumDocument, VellumImage
 from vellum.workflows.types.core import Json
 from vellum.workflows.utils.vellum_variables import primitive_type_to_vellum_variable_type
 
@@ -21,6 +21,12 @@ from vellum.workflows.utils.vellum_variables import primitive_type_to_vellum_var
         (Optional[List[SearchResult]], "SEARCH_RESULTS"),
         (Json, "JSON"),
         (Optional[Json], "JSON"),
+        (VellumDocument, "DOCUMENT"),
+        (Optional[VellumDocument], "DOCUMENT"),
+        (VellumAudio, "AUDIO"),
+        (Optional[VellumAudio], "AUDIO"),
+        (VellumImage, "IMAGE"),
+        (Optional[VellumImage], "IMAGE"),
     ],
 )
 def test_primitive_type_to_vellum_variable_type(type_, expected):

vellum/workflows/utils/vellum_variables.py

@@ -10,6 +10,8 @@ from vellum import (
     SearchResultRequest,
     VellumAudio,
     VellumAudioRequest,
+    VellumDocument,
+    VellumDocumentRequest,
     VellumError,
     VellumErrorRequest,
     VellumImage,
@@ -62,6 +64,8 @@ def primitive_type_to_vellum_variable_type(type_: Union[Type, BaseDescriptor]) -
         return "IMAGE"
     elif _is_type_optionally_in(type_, (VellumAudio, VellumAudioRequest)):
         return "AUDIO"
+    elif _is_type_optionally_in(type_, (VellumDocument, VellumDocumentRequest)):
+        return "DOCUMENT"
     elif _is_type_optionally_in(type_, (VellumError, VellumErrorRequest)):
         return "ERROR"
     elif _is_type_optionally_in(type_, (List[ChatMessage], List[ChatMessageRequest])):

{vellum_ai-0.14.16.dist-info → vellum_ai-0.14.18.dist-info}/METADATA

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