Skip to content

Commit b5cef85

Browse files
Add validation to raise error when ChatMessageTrigger and entrypoint are both present
- Add ChatMessageTrigger to WorkflowTriggerType enum (CHAT_MESSAGE) - Add validation in serialize() to raise WorkflowValidationError when both ChatMessageTrigger and entrypoint nodes are present (mutually exclusive) - Add ChatMessageTrigger handling in _serialize_workflow_trigger() - Add tests for the validation Co-Authored-By: vincent@vellum.ai <0426vincent@gmail.com>
1 parent d2f6997 commit b5cef85

File tree

2 files changed

+101
-0
lines changed

2 files changed

+101
-0
lines changed
Lines changed: 88 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,88 @@
1+
import pytest
2+
from typing import List
3+
4+
from pydantic import Field
5+
6+
from vellum.client.types import ChatMessage
7+
from vellum.workflows import BaseWorkflow
8+
from vellum.workflows.graph import Graph
9+
from vellum.workflows.inputs.base import BaseInputs
10+
from vellum.workflows.nodes.bases.base import BaseNode
11+
from vellum.workflows.references import LazyReference
12+
from vellum.workflows.state.base import BaseState
13+
from vellum.workflows.triggers.chat_message import ChatMessageTrigger
14+
from vellum_ee.workflows.display.utils.exceptions import WorkflowValidationError
15+
from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
16+
17+
18+
class ChatState(BaseState):
19+
chat_history: List[ChatMessage] = Field(default_factory=list)
20+
21+
22+
class SimpleNode(BaseNode):
23+
class Outputs(BaseNode.Outputs):
24+
response: str = "Hello!"
25+
26+
27+
class AnotherNode(BaseNode):
28+
class Outputs(BaseNode.Outputs):
29+
result: str = "Another!"
30+
31+
32+
def test_chat_message_trigger_with_entrypoint_raises_error():
33+
"""
34+
Tests that serializing a workflow with both ChatMessageTrigger and an entrypoint raises an error.
35+
"""
36+
37+
# GIVEN a workflow with ChatMessageTrigger that also has a separate entrypoint (multiple subgraphs)
38+
class SimpleChatTrigger(ChatMessageTrigger):
39+
class Config(ChatMessageTrigger.Config):
40+
output = LazyReference(lambda: WorkflowWithEntrypoint.Outputs.response) # type: ignore[has-type]
41+
42+
class WorkflowWithEntrypoint(BaseWorkflow[BaseInputs, ChatState]):
43+
graph = {
44+
SimpleChatTrigger >> SimpleNode,
45+
Graph.from_node(AnotherNode),
46+
}
47+
48+
class Outputs(BaseWorkflow.Outputs):
49+
response = SimpleNode.Outputs.response
50+
51+
# WHEN we try to serialize the workflow
52+
# THEN it should raise a WorkflowValidationError
53+
with pytest.raises(WorkflowValidationError) as exc_info:
54+
get_workflow_display(workflow_class=WorkflowWithEntrypoint).serialize()
55+
56+
# AND the error message should indicate the mutual exclusivity
57+
assert "ChatMessageTrigger and entrypoint nodes are mutually exclusive" in str(exc_info.value)
58+
59+
60+
def test_chat_message_trigger_without_entrypoint_does_not_raise_mutual_exclusivity_error():
61+
"""
62+
Tests that serializing a workflow with only ChatMessageTrigger (no entrypoint) does not raise
63+
the mutual exclusivity error.
64+
"""
65+
66+
# GIVEN a workflow with only ChatMessageTrigger (no entrypoint)
67+
class SimpleChatTrigger(ChatMessageTrigger):
68+
class Config(ChatMessageTrigger.Config):
69+
output = LazyReference(lambda: ChatOnlyWorkflow.Outputs.response) # type: ignore[has-type]
70+
71+
class ChatOnlyWorkflow(BaseWorkflow[BaseInputs, ChatState]):
72+
graph = SimpleChatTrigger >> SimpleNode
73+
74+
class Outputs(BaseWorkflow.Outputs):
75+
response = SimpleNode.Outputs.response
76+
77+
# WHEN we try to serialize the workflow
78+
# THEN it should NOT raise a WorkflowValidationError about mutual exclusivity
79+
# (it may raise other errors due to ChatMessageTrigger attribute serialization limitations)
80+
try:
81+
get_workflow_display(workflow_class=ChatOnlyWorkflow).serialize()
82+
except WorkflowValidationError as e:
83+
# If we get a WorkflowValidationError, it should NOT be about mutual exclusivity
84+
assert "ChatMessageTrigger and entrypoint nodes are mutually exclusive" not in str(
85+
e
86+
), f"Should not raise mutual exclusivity error for ChatMessageTrigger-only workflow, got: {e}"
87+
except ValueError:
88+
pass

ee/vellum_ee/workflows/display/workflows/base_workflow_display.py

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -274,6 +274,19 @@ def serialize(self) -> JsonObject:
274274
has_triggers = len(trigger_edges) > 0
275275
needs_entrypoint_node = has_manual_trigger or not has_triggers or len(non_trigger_entrypoint_nodes) > 0
276276

277+
# Check for ChatMessageTrigger - it's mutually exclusive with having an entrypoint node
278+
chat_message_trigger_edges = [
279+
edge for edge in trigger_edges if issubclass(edge.trigger_class, ChatMessageTrigger)
280+
]
281+
has_chat_message_trigger = len(chat_message_trigger_edges) > 0
282+
283+
if has_chat_message_trigger and needs_entrypoint_node:
284+
raise WorkflowValidationError(
285+
message="ChatMessageTrigger and entrypoint nodes are mutually exclusive. "
286+
"A workflow cannot have both a ChatMessageTrigger and an entrypoint node.",
287+
workflow_class_name=self._workflow.__name__,
288+
)
289+
277290
entrypoint_node_id: Optional[UUID] = None
278291
entrypoint_node_source_handle_id: Optional[UUID] = None
279292
entrypoint_node_display = self.display_context.workflow_display.entrypoint_node_display

0 commit comments

Comments
 (0)