vellum-ai 0.14.51__tar.gz → 0.14.53__tar.gz

{vellum_ai-0.14.51 → vellum_ai-0.14.53}/PKG-INFO

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

{vellum_ai-0.14.51 → vellum_ai-0.14.53}/ee/vellum_cli/__init__.py

@@ -259,6 +259,11 @@ Helpful for running and debugging workflows locally.""",
     help="""Directory to pull the workflow into. If not specified, \
 the workflow will be pulled into the current working directory.""",
 )
+@click.option(
+    "--workspace",
+    type=str,
+    help="The specific Workspace config to use when pulling",
+)
 def workflows_pull(
     module: Optional[str],
     include_json: Optional[bool],
@@ -268,6 +273,7 @@ def workflows_pull(
     strict: Optional[bool],
     include_sandbox: Optional[bool],
     target_directory: Optional[str],
+    workspace: Optional[str],
 ) -> None:
     """
     Pull Workflows from Vellum. If a module is provided, only the Workflow for that module will be pulled.
@@ -283,6 +289,7 @@ def workflows_pull(
         strict=strict,
         include_sandbox=include_sandbox,
         target_directory=target_directory,
+        workspace=workspace,
     )
 
 

{vellum_ai-0.14.51 → vellum_ai-0.14.53}/ee/vellum_cli/pull.py

@@ -13,7 +13,13 @@ from vellum.client.core.api_error import ApiError
 from vellum.client.core.pydantic_utilities import UniversalBaseModel
 from vellum.utils.uuid import is_valid_uuid
 from vellum.workflows.vellum_client import create_vellum_client
-from vellum_cli.config import VellumCliConfig, WorkflowConfig, WorkflowDeploymentConfig, load_vellum_cli_config
+from vellum_cli.config import (
+    DEFAULT_WORKSPACE_CONFIG,
+    VellumCliConfig,
+    WorkflowConfig,
+    WorkflowDeploymentConfig,
+    load_vellum_cli_config,
+)
 from vellum_cli.logger import load_cli_logger
 
 ERROR_LOG_FILE_NAME = "error.log"
@@ -43,6 +49,7 @@ def _resolve_workflow_config(
     module: Optional[str] = None,
     workflow_sandbox_id: Optional[str] = None,
     workflow_deployment: Optional[str] = None,
+    workspace: Optional[str] = None,
 ) -> WorkflowConfigResolutionResult:
     if workflow_sandbox_id and workflow_deployment:
         raise ValueError("Cannot specify both workflow_sandbox_id and workflow_deployment")
@@ -53,6 +60,7 @@ def _resolve_workflow_config(
             workflow_config = WorkflowConfig(
                 workflow_sandbox_id=workflow_sandbox_id,
                 module=module,
+                workspace=workspace or DEFAULT_WORKSPACE_CONFIG.name,
             )
             config.workflows.append(workflow_config)
             return WorkflowConfigResolutionResult(
@@ -132,8 +140,9 @@ def pull_command(
     strict: Optional[bool] = None,
     include_sandbox: Optional[bool] = None,
     target_directory: Optional[str] = None,
+    workspace: Optional[str] = None,
 ) -> None:
-    load_dotenv()
+    load_dotenv(dotenv_path=os.path.join(os.getcwd(), ".env"))
     logger = load_cli_logger()
     config = load_vellum_cli_config()
 
@@ -142,6 +151,7 @@ def pull_command(
         module=module,
         workflow_sandbox_id=workflow_sandbox_id,
         workflow_deployment=workflow_deployment,
+        workspace=workspace,
     )
 
     workflow_config = workflow_config_result.workflow_config
@@ -157,7 +167,20 @@ def pull_command(
     else:
         logger.info(f"Pulling workflow from {pk}...")
 
-    client = create_vellum_client()
+    resolved_workspace = workspace or workflow_config.workspace or DEFAULT_WORKSPACE_CONFIG.name
+    workspace_config = (
+        next((w for w in config.workspaces if w.name == resolved_workspace), DEFAULT_WORKSPACE_CONFIG)
+        if workspace
+        else DEFAULT_WORKSPACE_CONFIG
+    )
+    api_key = os.getenv(workspace_config.api_key)
+    if not api_key:
+        raise ValueError(f"No API key value found in environment for workspace '{workspace_config.name}'.")
+
+    client = create_vellum_client(
+        api_key=api_key,
+        api_url=workspace_config.api_url,
+    )
     query_parameters = {}
 
     if include_json:

{vellum_ai-0.14.51 → vellum_ai-0.14.53}/ee/vellum_cli/push.py

@@ -45,8 +45,11 @@ def push_command(
         else config.workflows
     )
 
-    if len(workflow_configs) > 1 and workspace:
-        workflow_configs = [w for w in workflow_configs if w.workspace == workspace]
+    if len(workflow_configs) > 1:
+        if workspace:
+            workflow_configs = [w for w in workflow_configs if w.workspace == workspace]
+        else:
+            workflow_configs = [w for w in workflow_configs if w.workspace == DEFAULT_WORKSPACE_CONFIG.name]
 
     if len(workflow_configs) == 0:
         if module and module_exists(module):

{vellum_ai-0.14.51 → vellum_ai-0.14.53}/ee/vellum_cli/tests/test_pull.py

@@ -3,6 +3,7 @@ import io
 import json
 import os
 import tempfile
+from unittest import mock
 from uuid import uuid4
 import zipfile
 
@@ -1284,3 +1285,74 @@ def test_pull__workflow_deployment_with_name_and_id(vellum_client):
         assert lock_data["workflows"][0]["deployments"][0]["label"] == deployment_label
 
     os.chdir(current_dir)
+
+
+def test_pull__workspace_option__uses_different_api_key(mock_module, vellum_client_class):
+    # GIVEN a module and workflow_sandbox_id
+    temp_dir = mock_module.temp_dir
+    module = mock_module.module
+    set_pyproject_toml = mock_module.set_pyproject_toml
+    workflow_sandbox_id = str(uuid4())
+
+    # AND a different workspace is set in the pyproject.toml
+    set_pyproject_toml(
+        {
+            "workflows": [],
+            "workspaces": [
+                {
+                    "name": "my_other_workspace",
+                    "api_key": "MY_OTHER_VELLUM_API_KEY",
+                }
+            ],
+        }
+    )
+
+    # AND the .env file has the other api key stored
+    with open(os.path.join(temp_dir, ".env"), "w") as f:
+        f.write(
+            """
+VELLUM_API_KEY=abcdef123456
+MY_OTHER_VELLUM_API_KEY=aaabbbcccddd
+"""
+        )
+
+    # AND the workflow pull API call returns a zip file
+    vellum_client_class.return_value.workflows.pull.return_value = iter(
+        [_zip_file_map({"workflow.py": "print('hello')"})]
+    )
+
+    # WHEN calling `vellum pull` with --workspace
+    runner = CliRunner()
+    result = runner.invoke(
+        cli_main,
+        [
+            "workflows",
+            "pull",
+            module,
+            "--workflow-sandbox-id",
+            workflow_sandbox_id,
+            "--workspace",
+            "my_other_workspace",
+        ],
+    )
+
+    # THEN it should succeed
+    assert result.exit_code == 0, result.output
+
+    # AND we should have called the vellum client with the correct api key
+    vellum_client_class.assert_called_once_with(
+        api_key="aaabbbcccddd",
+        environment=mock.ANY,
+    )
+
+    # AND the vellum lock file should have been updated with the correct workspace
+    with open(os.path.join(temp_dir, "vellum.lock.json")) as f:
+        lock_file_content = json.load(f)
+        assert lock_file_content["workflows"][0]["workspace"] == "my_other_workspace"
+        assert lock_file_content["workflows"][0]["workflow_sandbox_id"] == workflow_sandbox_id
+
+    # AND the workflow.py file is written as expected
+    workflow_py = os.path.join(temp_dir, *module.split("."), "workflow.py")
+    assert os.path.exists(workflow_py)
+    with open(workflow_py) as f:
+        assert f.read() == "print('hello')"

{vellum_ai-0.14.51 → vellum_ai-0.14.53}/ee/vellum_cli/tests/test_push.py

@@ -870,3 +870,45 @@ def test_push__create_new_config_for_existing_module(mock_module, vellum_client)
         new_config = new_configs[0]
         assert new_config["workflow_sandbox_id"] == new_workflow_sandbox_id
         assert new_config["workspace"] == "default"
+
+
+def test_push__use_default_workspace_if_not_specified__multiple_workflows_configured(mock_module, vellum_client):
+    # GIVEN a config with a workspace configured
+    temp_dir = mock_module.temp_dir
+    module = mock_module.module
+    workflow_sandbox_id = str(uuid4())
+    mock_module.set_pyproject_toml(
+        {
+            "workspaces": [
+                {"name": "my_other_workspace"},
+            ],
+            "workflows": [
+                {"module": module, "workflow_sandbox_id": workflow_sandbox_id, "workspace": "default"},
+                {"module": module, "workflow_sandbox_id": str(uuid4()), "workspace": "my_other_workspace"},
+            ],
+        }
+    )
+
+    # AND a workflow exists in the module successfully
+    _ensure_workflow_py(temp_dir, module)
+
+    # AND the push API call returns successfully
+    vellum_client.workflows.push.return_value = WorkflowPushResponse(
+        workflow_sandbox_id=workflow_sandbox_id,
+    )
+
+    # WHEN calling `vellum push` with a module without a workspace specified
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["workflows", "push", module])
+
+    # THEN it should succeed
+    assert result.exit_code == 0, result.output
+
+    # AND check that lockfile should maintain that this workflow is using the default workspace
+    with open(os.path.join(temp_dir, "vellum.lock.json")) as f:
+        lock_file_content = json.load(f)
+        configs = [w for w in lock_file_content["workflows"] if w["module"] == module]
+        assert len(configs) == 2
+        config = configs[0]
+        assert config["workflow_sandbox_id"] == workflow_sandbox_id
+        assert config["workspace"] == "default"

{vellum_ai-0.14.51 → vellum_ai-0.14.53}/ee/vellum_ee/workflows/display/nodes/vellum/tests/test_prompt_node.py

@@ -198,7 +198,7 @@ def test_serialize_node__prompt_inputs__state_reference():
                 {
                     "key": "foo",
                     "value": {
-                        "type": "STATE_VALUE",
+                        "type": "WORKFLOW_STATE",
                         "state_variable_id": "45649791-c642-4405-aff9-a1fafd780ea1",
                     },
                 },

{vellum_ai-0.14.51 → vellum_ai-0.14.53}/ee/vellum_ee/workflows/display/utils/expressions.py

@@ -207,7 +207,7 @@ def serialize_value(display_context: "WorkflowDisplayContext", value: Any) -> Js
     if isinstance(value, StateValueReference):
         state_value_display = display_context.global_state_value_displays[value]
         return {
-            "type": "STATE_VALUE",
+            "type": "WORKFLOW_STATE",
             "state_variable_id": str(state_value_display.id),
         }
 

{vellum_ai-0.14.51 → vellum_ai-0.14.53}/pyproject.toml

@@ -3,7 +3,7 @@ name = "vellum-ai"
 
 [tool.poetry]
 name = "vellum-ai"
-version = "0.14.51"
+version = "0.14.53"
 description = ""
 readme = "README.md"
 authors = []

{vellum_ai-0.14.51 → vellum_ai-0.14.53}/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.14.51",
+            "X-Fern-SDK-Version": "0.14.53",
         }
         headers["X-API-KEY"] = self.api_key
         return headers

{vellum_ai-0.14.51 → vellum_ai-0.14.53}/src/vellum/workflows/nodes/core/map_node/tests/test_node.py

@@ -194,7 +194,7 @@ def test_map_node_parallel_execution_with_workflow():
             thread_ids[self.item] = current_thread_id
 
             # Simulate work
-            time.sleep(0.01)
+            time.sleep(0.1)
 
             end = time.time()
             end_str = datetime.datetime.fromtimestamp(end).strftime("%Y-%m-%d %H:%M:%S.%f")

{vellum_ai-0.14.51 → vellum_ai-0.14.53}/src/vellum/workflows/nodes/displayable/bases/inline_prompt_node/node.py

@@ -18,6 +18,7 @@ from vellum import (
 )
 from vellum.client import ApiError, RequestOptions
 from vellum.client.types.chat_message_request import ChatMessageRequest
+from vellum.client.types.prompt_exec_config import PromptExecConfig
 from vellum.client.types.prompt_settings import PromptSettings
 from vellum.client.types.rich_text_child_block import RichTextChildBlock
 from vellum.workflows.context import get_execution_context
@@ -134,6 +135,23 @@ class BaseInlinePromptNode(BasePromptNode[StateType], Generic[StateType]):
             )
 
     def _process_prompt_event_stream(self) -> Generator[BaseOutput, None, Optional[List[PromptOutput]]]:
+        try:
+            # Compile dict blocks into PromptBlocks
+            exec_config = PromptExecConfig.model_validate(
+                {
+                    "ml_model": "",
+                    "input_variables": [],
+                    "parameters": {},
+                    "blocks": self.blocks,
+                }
+            )
+            self.blocks = exec_config.blocks  # type: ignore
+        except Exception:
+            raise NodeException(
+                message="Failed to compile blocks",
+                code=WorkflowErrorCode.INVALID_INPUTS,
+            )
+
         self._validate()
         try:
             prompt_event_stream = self._get_prompt_event_stream()

{vellum_ai-0.14.51 → vellum_ai-0.14.53}/src/vellum/workflows/nodes/displayable/bases/inline_prompt_node/tests/test_inline_prompt_node.py

@@ -473,3 +473,149 @@ def test_inline_prompt_node__json_output_with_streaming_disabled(vellum_adhoc_pr
         request_options=mock.ANY,
         settings=PromptSettings(stream_enabled=False),
     )
+
+
+def test_inline_prompt_node__dict_blocks(vellum_adhoc_prompt_client):
+    # GIVEN a node that has dict blocks
+    class MyInlinePromptNode(InlinePromptNode):
+        ml_model = "gpt-4o"
+        blocks = [
+            {  # type: ignore
+                "state": None,
+                "blocks": [
+                    {
+                        "state": None,
+                        "blocks": [
+                            {
+                                "text": "You are a weather expert",
+                                "state": None,
+                                "block_type": "PLAIN_TEXT",
+                                "cache_config": None,
+                            }
+                        ],
+                        "block_type": "RICH_TEXT",
+                        "cache_config": None,
+                    }
+                ],
+                "chat_role": "SYSTEM",
+                "block_type": "CHAT_MESSAGE",
+                "chat_source": None,
+                "cache_config": None,
+                "chat_message_unterminated": None,
+            },
+            {  # type: ignore
+                "state": None,
+                "blocks": [
+                    {
+                        "state": None,
+                        "blocks": [
+                            {
+                                "state": None,
+                                "block_type": "VARIABLE",
+                                "cache_config": None,
+                                "input_variable": "question",
+                            }
+                        ],
+                        "block_type": "RICH_TEXT",
+                        "cache_config": None,
+                    }
+                ],
+                "chat_role": "USER",
+                "block_type": "CHAT_MESSAGE",
+                "chat_source": None,
+                "cache_config": None,
+                "chat_message_unterminated": None,
+            },
+            VariablePromptBlock(block_type="VARIABLE", state=None, cache_config=None, input_variable="chat_history"),
+        ]
+        prompt_inputs = {
+            "question": "What is the weather in Tokyo?",
+            "chat_history": "You are a weather expert",
+        }
+        settings = PromptSettings(stream_enabled=False)
+
+    # 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_event(*args: Any, **kwargs: Any) -> AdHocExecutePromptEvent:
+        execution_id = str(uuid4())
+        return FulfilledAdHocExecutePromptEvent(
+            execution_id=execution_id,
+            outputs=expected_outputs,
+        )
+
+    vellum_adhoc_prompt_client.adhoc_execute_prompt.side_effect = generate_prompt_event
+
+    # WHEN the node is run
+    node = MyInlinePromptNode()
+    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!"}'
+
+
+def test_inline_prompt_node__dict_blocks_error(vellum_adhoc_prompt_client):
+    # GIVEN a node that has an error (wrong block type)
+    class MyInlinePromptNode(InlinePromptNode):
+        ml_model = "gpt-4o"
+        blocks = [
+            {  # type: ignore
+                "state": None,
+                "blocks": [
+                    {
+                        "state": None,
+                        "blocks": [
+                            {
+                                "text": "You are a weather expert",
+                                "state": None,
+                                "block_type": "PLAIN_TEXT",
+                                "cache_config": None,
+                            }
+                        ],
+                        "block_type": "WRONG_BLOCK_TYPE",
+                        "cache_config": None,
+                    }
+                ],
+                "chat_role": "SYSTEM",
+                "block_type": "CHAT_MESSAGE",
+                "chat_source": None,
+                "cache_config": None,
+                "chat_message_unterminated": None,
+            },
+        ]
+        prompt_inputs = {
+            "question": "What is the weather in Tokyo?",
+        }
+        settings = PromptSettings(stream_enabled=False)
+
+    # 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_event(*args: Any, **kwargs: Any) -> AdHocExecutePromptEvent:
+        execution_id = str(uuid4())
+        return FulfilledAdHocExecutePromptEvent(
+            execution_id=execution_id,
+            outputs=expected_outputs,
+        )
+
+    vellum_adhoc_prompt_client.adhoc_execute_prompt.side_effect = generate_prompt_event
+
+    node = MyInlinePromptNode()
+    with pytest.raises(NodeException) as excinfo:
+        list(node.run())
+
+    # THEN the node should raise the correct NodeException
+    assert excinfo.value.code == WorkflowErrorCode.INVALID_INPUTS
+    assert "Failed to compile blocks" == str(excinfo.value)

{vellum_ai-0.14.51 → vellum_ai-0.14.53}/src/vellum/workflows/nodes/displayable/search_node/node.py

@@ -1,6 +1,8 @@
+import json
 from typing import ClassVar
 
 from vellum.workflows.nodes.displayable.bases import BaseSearchNode as BaseSearchNode
+from vellum.workflows.state.encoder import DefaultStateEncoder
 from vellum.workflows.types import MergeBehavior
 from vellum.workflows.types.generics import StateType
 
@@ -33,6 +35,9 @@ class SearchNode(BaseSearchNode[StateType]):
         text: str
 
     def run(self) -> Outputs:
+        if not isinstance(self.query, str):
+            self.query = json.dumps(self.query, cls=DefaultStateEncoder)
+
         results = self._perform_search().results
         text = self.chunk_separator.join([r.text for r in results])
         return self.Outputs(results=results, text=text)

{vellum_ai-0.14.51 → vellum_ai-0.14.53}/src/vellum/workflows/nodes/displayable/search_node/tests/test_node.py

@@ -1,4 +1,7 @@
+import json
+
 from vellum import SearchResponse, SearchResult, SearchResultDocument
+from vellum.client.types.chat_message import ChatMessage
 from vellum.client.types.json_vellum_value_request import JsonVellumValueRequest
 from vellum.client.types.search_filters_request import SearchFiltersRequest
 from vellum.client.types.search_request_options_request import SearchRequestOptionsRequest
@@ -13,6 +16,7 @@ from vellum.workflows.nodes.displayable.bases.types import (
     SearchFilters,
 )
 from vellum.workflows.nodes.displayable.search_node.node import SearchNode
+from vellum.workflows.state.base import BaseState
 
 
 def test_run_workflow__happy_path(vellum_client):
@@ -172,3 +176,42 @@ def test_run_workflow__happy_path__options_attribute(vellum_client):
             ),
         ),
     )
+
+
+def test_run_workflow__chat_history_as_query(vellum_client):
+    """
+    Confirm that we can successfully invoke a Search node with a chat history as the query param,
+    backwards compatible with original workflows
+    """
+
+    # GIVEN a state definition with a chat history
+    class MyState(BaseState):
+        chat_history: list[ChatMessage]
+
+    # AND a Search Node that uses the chat history as the query param
+    class MySearchNode(SearchNode[MyState]):
+        query = MyState.chat_history  # type: ignore[assignment]
+        document_index = "document_index"
+        limit = 1
+
+    # AND a Search request that will return a 200 ok resposne
+    search_response = SearchResponse(
+        results=[
+            SearchResult(
+                text="Search query", score="0.0", keywords=["keywords"], document=SearchResultDocument(label="label")
+            )
+        ]
+    )
+
+    vellum_client.search.return_value = search_response
+
+    # WHEN we run the workflow
+    outputs = MySearchNode(state=MyState(chat_history=[ChatMessage(role="USER", text="Hello, world!")])).run()
+
+    # THEN the workflow should have completed successfully
+    assert outputs.text == "Search query"
+
+    # AND the options should be as expected
+    assert json.loads(vellum_client.search.call_args.kwargs["query"]) == [
+        {"role": "USER", "text": "Hello, world!", "source": None, "content": None}
+    ]