vellum-ai 0.14.63__py3-none-any.whl → 1.13.5__py3-none-any.whl

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

@@ -1,18 +1,33 @@
 from abc import abstractmethod
-from typing import ClassVar, Generator, Generic, Iterator, List, Optional, Union
+from itertools import chain
+from typing import ClassVar, Generator, Generic, Iterator, List, Optional, Union, cast
 
 from vellum import AdHocExecutePromptEvent, ExecutePromptEvent, PromptOutput
+from vellum.client.core import RequestOptions
 from vellum.client.core.api_error import ApiError
-from vellum.core import RequestOptions
+from vellum.workflows.descriptors.base import BaseDescriptor
 from vellum.workflows.errors.types import WorkflowErrorCode, vellum_error_to_workflow_error
+from vellum.workflows.events.node import NodeExecutionStreamingEvent
 from vellum.workflows.exceptions import NodeException
+from vellum.workflows.expressions.coalesce_expression import CoalesceExpression
 from vellum.workflows.nodes.bases import BaseNode
 from vellum.workflows.outputs.base import BaseOutput, BaseOutputs
+from vellum.workflows.references.output import OutputReference
 from vellum.workflows.types.core import EntityInputsInterface, MergeBehavior
 from vellum.workflows.types.generics import StateType
 
 
-class BasePromptNode(BaseNode, Generic[StateType]):
+def _contains_reference_to_output(reference: BaseDescriptor, target_reference: OutputReference) -> bool:
+    if reference == target_reference:
+        return True
+    if isinstance(reference, CoalesceExpression):
+        return _contains_reference_to_output(reference.lhs, target_reference) or _contains_reference_to_output(
+            reference.rhs, target_reference
+        )
+    return False
+
+
+class BasePromptNode(BaseNode[StateType], Generic[StateType]):
     # Inputs that are passed to the Prompt
     prompt_inputs: ClassVar[Optional[EntityInputsInterface]] = None
 
@@ -46,25 +61,41 @@ class BasePromptNode(BaseNode, Generic[StateType]):
         except ApiError as e:
             self._handle_api_error(e)
 
-        # We don't use the INITIATED event anyway, so we can just skip it
-        # and use the exception handling to catch other api level errors
         try:
-            next(prompt_event_stream)
+            first_event = next(prompt_event_stream)
         except ApiError as e:
             self._handle_api_error(e)
+        else:
+            if first_event.state == "REJECTED":
+                workflow_error = vellum_error_to_workflow_error(first_event.error)
+                raise NodeException.of(workflow_error)
+            if first_event.state != "INITIATED":
+                prompt_event_stream = cast(
+                    Union[Iterator[AdHocExecutePromptEvent], Iterator[ExecutePromptEvent]],
+                    chain([first_event], prompt_event_stream),
+                )
 
         outputs: Optional[List[PromptOutput]] = None
+        exception: Optional[NodeException] = None
         for event in prompt_event_stream:
+            if exception:
+                continue
+
             if event.state == "INITIATED":
                 continue
             elif event.state == "STREAMING":
                 yield BaseOutput(name="results", delta=event.output.value)
+                if event.output.type == "STRING":
+                    yield BaseOutput(name="text", delta=event.output.value)
             elif event.state == "FULFILLED":
                 outputs = event.outputs
                 yield BaseOutput(name="results", value=event.outputs)
             elif event.state == "REJECTED":
                 workflow_error = vellum_error_to_workflow_error(event.error)
-                raise NodeException.of(workflow_error)
+                exception = NodeException.of(workflow_error)
+
+        if exception:
+            raise exception
 
         return outputs
 
@@ -73,15 +104,62 @@ class BasePromptNode(BaseNode, Generic[StateType]):
             raise NodeException(
                 message=e.body.get("detail", "Provider credentials is missing or unavailable"),
                 code=WorkflowErrorCode.PROVIDER_CREDENTIALS_UNAVAILABLE,
+                raw_data=e.body,
             )
 
         elif e.status_code and e.status_code >= 400 and e.status_code < 500 and isinstance(e.body, dict):
             raise NodeException(
                 message=e.body.get("detail", "Failed to execute Prompt"),
                 code=WorkflowErrorCode.INVALID_INPUTS,
+                raw_data=e.body,
             ) from e
 
         raise NodeException(
             message="Failed to execute Prompt",
             code=WorkflowErrorCode.INTERNAL_ERROR,
+            raw_data=e.body,
         ) from e
+
+    def __directly_emit_workflow_output__(
+        self,
+        event: NodeExecutionStreamingEvent,
+        workflow_output_descriptor: OutputReference,
+    ) -> bool:
+        # Check if workflow output directly references this node's text output
+        text_output = getattr(event.node_definition.Outputs, "text", None)
+
+        if (
+            text_output is not None
+            and event.output.name == "text"
+            and isinstance(workflow_output_descriptor.instance, BaseDescriptor)
+            and _contains_reference_to_output(workflow_output_descriptor.instance, text_output)
+        ):
+            return True
+
+        if event.output.name != "results":
+            return False
+
+        if not isinstance(event.output.delta, str) and not event.output.is_initiated:
+            return False
+
+        # Check if workflow output references this node's text output through a FinalOutputNode
+        target_nodes = [e.to_node for port in self.Ports for e in port.edges if e.to_node.__simulates_workflow_output__]
+        target_node_output = next(
+            (
+                o
+                for target_node in target_nodes
+                for o in target_node.Outputs
+                if o == workflow_output_descriptor.instance
+            ),
+            None,
+        )
+        if not target_node_output:
+            return False
+
+        if not isinstance(target_node_output.instance, BaseDescriptor):
+            return False
+
+        if text_output is None:
+            return False
+
+        return _contains_reference_to_output(target_node_output.instance, text_output)

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

@@ -1,12 +1,30 @@
+from itertools import chain
 import json
 from uuid import uuid4
-from typing import Callable, ClassVar, Generator, Generic, Iterator, List, Optional, Set, Tuple, Union
+from typing import (
+    TYPE_CHECKING,
+    Callable,
+    ClassVar,
+    Generator,
+    Generic,
+    Iterator,
+    List,
+    Optional,
+    Set,
+    Tuple,
+    Type,
+    Union,
+)
+
+import httpx
+import jsonschema
 
 from vellum import (
     AdHocExecutePromptEvent,
     AdHocExpandMeta,
     ChatMessage,
     FunctionDefinition,
+    InitiatedAdHocExecutePromptEvent,
     PromptBlock,
     PromptOutput,
     PromptParameters,
@@ -17,21 +35,106 @@ from vellum import (
     VellumVariable,
 )
 from vellum.client import ApiError, RequestOptions
+from vellum.client.types import (
+    PromptRequestAudioInput,
+    PromptRequestDocumentInput,
+    PromptRequestImageInput,
+    PromptRequestVideoInput,
+    VellumAudio,
+    VellumAudioRequest,
+    VellumDocument,
+    VellumDocumentRequest,
+    VellumImage,
+    VellumImageRequest,
+    VellumVideo,
+    VellumVideoRequest,
+)
 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.prompts.constants import DEFAULT_PROMPT_PARAMETERS
 from vellum.workflows.context import get_execution_context
 from vellum.workflows.errors import WorkflowErrorCode
 from vellum.workflows.errors.types import vellum_error_to_workflow_error
 from vellum.workflows.events.types import default_serializer
 from vellum.workflows.exceptions import NodeException
 from vellum.workflows.nodes.displayable.bases.base_prompt_node import BasePromptNode
-from vellum.workflows.nodes.displayable.bases.inline_prompt_node.constants import DEFAULT_PROMPT_PARAMETERS
+from vellum.workflows.nodes.displayable.bases.utils import process_additional_prompt_outputs
 from vellum.workflows.outputs import BaseOutput
 from vellum.workflows.types import MergeBehavior
+from vellum.workflows.types.definition import (
+    ComposioToolDefinition,
+    DeploymentDefinition,
+    MCPServer,
+    MCPToolDefinition,
+    VellumIntegrationToolDefinition,
+)
 from vellum.workflows.types.generics import StateType, is_workflow_class
-from vellum.workflows.utils.functions import compile_function_definition, compile_workflow_function_definition
+from vellum.workflows.utils.functions import (
+    compile_composio_tool_definition,
+    compile_function_definition,
+    compile_inline_workflow_function_definition,
+    compile_mcp_tool_definition,
+    compile_vellum_integration_tool_definition,
+    compile_workflow_deployment_function_definition,
+    get_mcp_tool_name,
+)
+from vellum.workflows.utils.pydantic_schema import normalize_json
+
+if TYPE_CHECKING:
+    from vellum.workflows.workflows.base import BaseWorkflow
+
+
+def _get_json_schema_to_validate(parameters_ref: object) -> Optional[dict]:
+    """
+    Extracts the JSON schema to validate from a parameters reference.
+
+    Also normalizes Pydantic models to JSON schema dicts so they can be validated.
+
+    Args:
+        parameters_ref: The parameters reference (NodeReference wrapping PromptParameters)
+
+    Returns:
+        The JSON schema dict to validate, or None if no schema found.
+    """
+    parameters_instance = getattr(parameters_ref, "instance", None)
+    if not parameters_instance:
+        return None
+
+    custom_params = getattr(parameters_instance, "custom_parameters", None)
+    if not isinstance(custom_params, dict):
+        return None
+
+    json_schema = custom_params.get("json_schema")
+    if json_schema is None:
+        return None
+
+    # Normalize Pydantic models to JSON schema dicts before validation
+    # This handles cases like {"name": "...", "schema": SomePydanticModel}
+    json_schema = normalize_json(json_schema)
+
+    # After normalization, we expect a dict; anything else we ignore
+    if not isinstance(json_schema, dict):
+        return None
+
+    return json_schema
+
+
+def _validate_json_schema_structure(schema: dict) -> None:
+    """
+    Validates the structure of a JSON schema using the jsonschema library.
+
+    This uses the JSON Schema meta-schema to validate that the provided schema
+    is a valid JSON Schema according to the specification.
+
+    Args:
+        schema: The JSON schema dictionary to validate
+
+    Raises:
+        jsonschema.exceptions.SchemaError: If the schema structure is invalid
+    """
+    jsonschema.Draft7Validator.check_schema(schema)
 
 
 class BaseInlinePromptNode(BasePromptNode[StateType], Generic[StateType]):
@@ -53,10 +156,22 @@ class BaseInlinePromptNode(BasePromptNode[StateType], Generic[StateType]):
     blocks: ClassVar[List[PromptBlock]]
 
     # The functions/tools that a Prompt has access to
-    functions: Optional[List[Union[FunctionDefinition, Callable]]] = None
+    functions: Optional[
+        List[
+            Union[
+                FunctionDefinition,
+                Callable,
+                DeploymentDefinition,
+                Type["BaseWorkflow"],
+                VellumIntegrationToolDefinition,
+                MCPServer,
+                MCPToolDefinition,
+            ]
+        ]
+    ] = None
 
     parameters: PromptParameters = DEFAULT_PROMPT_PARAMETERS
-    expand_meta: Optional[AdHocExpandMeta] = None
+    expand_meta: Optional[AdHocExpandMeta] = AdHocExpandMeta(finish_reason=True)
 
     settings: Optional[PromptSettings] = None
 
@@ -93,6 +208,9 @@ class BaseInlinePromptNode(BasePromptNode[StateType], Generic[StateType]):
         execution_context = get_execution_context()
         request_options = self.request_options or RequestOptions()
 
+        processed_parameters = self.process_parameters(self.parameters)
+        processed_blocks = self.process_blocks(self.blocks)
+
         request_options["additional_body_parameters"] = {
             "execution_context": execution_context.model_dump(mode="json"),
             **request_options.get("additional_body_parameters", {}),
@@ -105,10 +223,46 @@ class BaseInlinePromptNode(BasePromptNode[StateType], Generic[StateType]):
             for function in self.functions:
                 if isinstance(function, FunctionDefinition):
                     normalized_functions.append(function)
+                elif isinstance(function, DeploymentDefinition):
+                    normalized_functions.append(
+                        compile_workflow_deployment_function_definition(
+                            function,
+                            vellum_client=self._context.vellum_client,
+                        )
+                    )
                 elif is_workflow_class(function):
-                    normalized_functions.append(compile_workflow_function_definition(function))
-                else:
+                    normalized_functions.append(compile_inline_workflow_function_definition(function))
+                elif callable(function):
                     normalized_functions.append(compile_function_definition(function))
+                elif isinstance(function, ComposioToolDefinition):
+                    normalized_functions.append(compile_composio_tool_definition(function))
+                elif isinstance(function, VellumIntegrationToolDefinition):
+                    normalized_functions.append(
+                        compile_vellum_integration_tool_definition(function, self._context.vellum_client)
+                    )
+                elif isinstance(function, MCPToolDefinition):
+                    normalized_functions.append(
+                        FunctionDefinition(
+                            name=get_mcp_tool_name(function),
+                            description=function.description,
+                            parameters=function.parameters,
+                        )
+                    )
+                elif isinstance(function, MCPServer):
+                    tool_definitions = compile_mcp_tool_definition(function)
+                    for tool_def in tool_definitions:
+                        normalized_functions.append(
+                            FunctionDefinition(
+                                name=get_mcp_tool_name(tool_def),
+                                description=tool_def.description,
+                                parameters=tool_def.parameters,
+                            )
+                        )
+                else:
+                    raise NodeException(
+                        message=f"`{function}` is not a valid function definition",
+                        code=WorkflowErrorCode.INVALID_INPUTS,
+                    )
 
         if self.settings and not self.settings.stream_enabled:
             # This endpoint is returning a single event, so we need to wrap it in a generator
@@ -117,21 +271,22 @@ class BaseInlinePromptNode(BasePromptNode[StateType], Generic[StateType]):
                 ml_model=self.ml_model,
                 input_values=input_values,
                 input_variables=input_variables,
-                parameters=self.parameters,
-                blocks=self.blocks,
+                parameters=processed_parameters,
+                blocks=processed_blocks,
                 settings=self.settings,
                 functions=normalized_functions,
                 expand_meta=self.expand_meta,
                 request_options=request_options,
             )
-            return iter([response])
+            initiated_event = InitiatedAdHocExecutePromptEvent(execution_id=response.execution_id)
+            return iter([initiated_event, response])
         else:
             return self._context.vellum_client.ad_hoc.adhoc_execute_prompt_stream(
                 ml_model=self.ml_model,
                 input_values=input_values,
                 input_variables=input_variables,
-                parameters=self.parameters,
-                blocks=self.blocks,
+                parameters=processed_parameters,
+                blocks=processed_blocks,
                 settings=self.settings,
                 functions=normalized_functions,
                 expand_meta=self.expand_meta,
@@ -161,14 +316,27 @@ class BaseInlinePromptNode(BasePromptNode[StateType], Generic[StateType]):
             prompt_event_stream = self._get_prompt_event_stream()
         except ApiError as e:
             self._handle_api_error(e)
+        except httpx.TransportError:
+            raise NodeException(
+                message="Failed to connect to Vellum server",
+                code=WorkflowErrorCode.INTERNAL_ERROR,
+            )
 
-        if not self.settings or (self.settings and self.settings.stream_enabled):
-            # We don't use the INITIATED event anyway, so we can just skip it
-            # and use the exception handling to catch other api level errors
-            try:
-                next(prompt_event_stream)
-            except ApiError as e:
-                self._handle_api_error(e)
+        try:
+            first_event = next(prompt_event_stream)
+        except ApiError as e:
+            self._handle_api_error(e)
+        except httpx.TransportError:
+            raise NodeException(
+                message="Failed to connect to Vellum server",
+                code=WorkflowErrorCode.INTERNAL_ERROR,
+            )
+        else:
+            if first_event.state == "REJECTED":
+                workflow_error = vellum_error_to_workflow_error(first_event.error)
+                raise NodeException.of(workflow_error)
+            if first_event.state != "INITIATED":
+                prompt_event_stream = chain([first_event], prompt_event_stream)
 
         outputs: Optional[List[PromptOutput]] = None
         for event in prompt_event_stream:
@@ -176,7 +344,20 @@ class BaseInlinePromptNode(BasePromptNode[StateType], Generic[StateType]):
                 continue
             elif event.state == "STREAMING":
                 yield BaseOutput(name="results", delta=event.output.value)
+                if event.output.type == "STRING":
+                    yield BaseOutput(name="text", delta=event.output.value)
             elif event.state == "FULFILLED":
+                if event.meta and event.meta.finish_reason == "LENGTH":
+                    text_value, json_value = process_additional_prompt_outputs(event.outputs)
+                    if text_value == "":
+                        raise NodeException(
+                            message=(
+                                "Maximum tokens reached before model could output any content. "
+                                "Consider increasing the max_tokens Prompt Parameter."
+                            ),
+                            code=WorkflowErrorCode.INVALID_OUTPUTS,
+                        )
+
                 outputs = event.outputs
                 yield BaseOutput(name="results", value=event.outputs)
             elif event.state == "REJECTED":
@@ -240,6 +421,97 @@ class BaseInlinePromptNode(BasePromptNode[StateType], Generic[StateType]):
                         value=chat_history,
                     )
                 )
+            elif isinstance(input_value, (VellumAudio, VellumAudioRequest)):
+                input_variables.append(
+                    VellumVariable(
+                        id=str(uuid4()),
+                        key=input_name,
+                        type="AUDIO",
+                    )
+                )
+                input_values.append(
+                    PromptRequestAudioInput(
+                        key=input_name,
+                        value=(
+                            input_value
+                            if isinstance(input_value, VellumAudio)
+                            else VellumAudio.model_validate(input_value.model_dump())
+                        ),
+                    )
+                )
+            elif isinstance(input_value, (VellumVideo, VellumVideoRequest)):
+                input_variables.append(
+                    VellumVariable(
+                        id=str(uuid4()),
+                        key=input_name,
+                        type="VIDEO",
+                    )
+                )
+                input_values.append(
+                    PromptRequestVideoInput(
+                        key=input_name,
+                        value=(
+                            input_value
+                            if isinstance(input_value, VellumVideo)
+                            else VellumVideo.model_validate(input_value.model_dump())
+                        ),
+                    )
+                )
+            elif isinstance(input_value, (VellumImage, VellumImageRequest)):
+                input_variables.append(
+                    VellumVariable(
+                        id=str(uuid4()),
+                        key=input_name,
+                        type="IMAGE",
+                    )
+                )
+                input_values.append(
+                    PromptRequestImageInput(
+                        key=input_name,
+                        value=(
+                            input_value
+                            if isinstance(input_value, VellumImage)
+                            else VellumImage.model_validate(input_value.model_dump())
+                        ),
+                    )
+                )
+            elif isinstance(input_value, (VellumDocument, VellumDocumentRequest)):
+                input_variables.append(
+                    VellumVariable(
+                        id=str(uuid4()),
+                        key=input_name,
+                        type="DOCUMENT",
+                    )
+                )
+                input_values.append(
+                    PromptRequestDocumentInput(
+                        key=input_name,
+                        value=(
+                            input_value
+                            if isinstance(input_value, VellumDocument)
+                            else VellumDocument.model_validate(input_value.model_dump())
+                        ),
+                    )
+                )
+            elif (
+                isinstance(input_value, dict)
+                and "src" in input_value
+                and isinstance(input_value.get("src"), str)
+                and input_value["src"].endswith(".pdf")
+            ):
+                input_variables.append(
+                    VellumVariable(
+                        id=str(uuid4()),
+                        key=input_name,
+                        type="DOCUMENT",
+                    )
+                )
+                input_values.append(
+                    PromptRequestDocumentInput(
+                        key=input_name,
+                        value=VellumDocument.model_validate(input_value),
+                    )
+                )
             else:
                 try:
                     input_value = default_serializer(input_value)
@@ -267,3 +539,36 @@ class BaseInlinePromptNode(BasePromptNode[StateType], Generic[StateType]):
                 )
 
         return input_variables, input_values
+
+    def process_parameters(self, parameters: PromptParameters) -> PromptParameters:
+        """
+        Process parameters to recursively convert any Pydantic models to JSON schema dictionaries.
+        """
+        if not parameters.custom_parameters:
+            return parameters
+
+        processed_custom_params = normalize_json(parameters.custom_parameters)
+
+        return parameters.model_copy(update={"custom_parameters": processed_custom_params})
+
+    def process_blocks(self, blocks: List[PromptBlock]) -> List[PromptBlock]:
+        """
+        Override this method to process the blocks before they are executed.
+        """
+        return blocks
+
+    @classmethod
+    def __validate__(cls) -> None:
+        """
+        Validates the node configuration, including JSON schema structure in parameters.
+
+        Raises:
+            jsonschema.exceptions.SchemaError: If the JSON schema structure is invalid
+        """
+        parameters_ref = getattr(cls, "parameters", None)
+        if parameters_ref is None:
+            return
+
+        schema = _get_json_schema_to_validate(parameters_ref)
+        if schema is not None:
+            _validate_json_schema_structure(schema)