Inline Prompt Node

1
from vellum import (
2
    ChatMessagePromptBlock,
3
    JinjaPromptBlock,
4
    PlainTextPromptBlock,
5
    PromptParameters,
6
    PromptSettings,
7
    RichTextPromptBlock,
8
    VariablePromptBlock,
9
)
10
from vellum.workflows.nodes.displayable import InlinePromptNode
11
from vellum.workflows.ports import Port
12
from vellum.workflows.references import LazyReference
13
14
# nodes/router_prompt.py
15
class RouterPrompt(InlinePromptNode):
16
    """
17
    This prompt is used to route to the appropriate handler based on the type of document being parsed.
18
    """
19
    ml_model = "gpt-5"
20
    blocks = [
21
        ChatMessagePromptBlock(
22
            chat_role="SYSTEM",
23
            blocks=[
24
                RichTextPromptBlock(
25
                    blocks=[
26
                        # Prefer RichTextPromptBlock for a nicer UI editing experience
27
                        PlainTextPromptBlock(text="Answer the user's question: "),
28
                        VariablePromptBlock(input_variable="question"),
29
                    ]
30
                ),
31
                JinjaPromptBlock(
32
                    template="Using a templating block to write Jinja templates inline: {{ query | upper | truncate(3) }}"
33
                ),
34
            ],
35
        ),
36
        # Use VariablePromptBlock at the top level to include chat history in context, for any chatbot / chat agent use-cases
37
        VariablePromptBlock(input_variable="chat_history"),
38
    ]
39
    prompt_inputs = {
40
        "document_text": Inputs.document_text,  # Reference workflow input
41
        "chat_history": Inputs.chat_history,  # List[ChatMessage]
42
    }
43
    custom_parameters={
44
        # prefer json_schema over json_mode if strict types are required
45
        # json_mode is a more flexible way to produce valid JSON through schemas defined in the prompt itself
46
        "json_mode": True,
47
        "json_schema": {
48
            "strict": True,
49
            "name": "schema",
50
            "schema": {
51
                "type": "object",
52
                "properties": {
53
                    "classification": {
54
                        "type": "string",
55
                        "description": "What type of document to classify as",
56
                        "enum": [
57
                            "policy",
58
                            "certificate_of_insurance",
59
                        ],
60
                    },
61
                },
62
                "required": [
63
                    "classification",
64
                ],
65
            },
66
        },
67
    },
68
69
    class Ports(InlinePromptNode.Ports):
70
        group_1_if_port = Port.on_if(LazyReference(lambda: RouterPrompt.Outputs.json)["classification"].equals("policy"))
71
        group_1_else_port = Port.on_else()
72
73
class FinalOutputNode(BaseOutputs):
74
    classification: str
75
76
77
# workflow.py
78
class Workflow(BaseWorkflow[Inputs, BaseState]):
79
    graph = GetDocument >> {
80
        RouterPrompt.Ports.group_1_if_port >> MyPolicyParserWorkflow,
81
        RouterPrompt.Ports.group_1_else_port >> MyCOIParserWorkflow,
82
    } >> FinalOutputNode
83
84
    class Outputs(BaseWorkflow.Outputs):
85
        final_output = FinalOutput.Outputs.value