1
# nodes/extract_user_message.py
2
from vellum.workflows.nodes import TemplatingNode
3
4
class ExtractUserMessage(TemplatingNode):
5
  # Here, we reference the chat_history input that we've connected to this node.
6
  template = """\
7
    {{ chat_history[-1]["text"] }}\
8
  """
9
10
  # Here, we define the inputs to _this_ node.
11
  inputs = {
12
    "chat_history": Inputs.chat_history,
13
  }

1
# nodes/search_node.py
2
from vellum.workflows.nodes import BaseSearchNode
3
4
class SearchNode(BaseSearchNode):
5
  document_index = "vellum-trust-center-policies"
6
  query = ExtractUserMessage.Outputs.result

1
# nodes/format_search_results_node.py
2
from vellum.workflows.nodes import TemplatingNode
3
4
class FormatSearchResultsNode(TemplatingNode):
5
  template = """\
6
    {% for result in results -%}
7
    Policy: {{ result.document.label }}
8
    ------
9
    {{ result.text }}
10
    {% if not loop.last %}
11
    #####
12
    {% endif -%}
13
    {% endfor %}\
14
  """
15
16
  inputs = {
17
    "results": SearchNode.Outputs.results,
18
  }

1
# nodes/prompt_node.py
2
from vellum.workflows.nodes import InlinePromptNode
3
from vellum import (
4
  ChatMessagePromptBlock,
5
  ChatMessageRequest,
6
  JinjaPromptBlock,
7
)
8
9
class PromptNode(InlinePromptNode):
10
  ml_model = "gpt-4o"
11
  prompt_inputs = {
12
    "question": ExtractUserMessage.Outputs.result,
13
    "context": FormatSearchResultsNode.Outputs.result,
14
  }
15
  blocks = [
16
      ChatMessagePromptBlock(
17
          chat_role="SYSTEM",
18
          blocks=[
19
              JinjaPromptBlock(
20
                  block_type="JINJA",
21
                  template="""\
22
                      Answer user question based on the context provided below, if you don't know the answer say "Sorry I don't know"
23
24
                      **Context**
25
                      ``
26
                      {{ context }}
27
                      ``
28
29
                      Limit your answer to 250 words and provide a citation at the end of your answer\
30
                  """,
31
              ),
32
          ],
33
      ),
34
      ChatMessagePromptBlock(
35
          chat_role="USER", 
36
          blocks=[
37
              JinjaPromptBlock(
38
                  block_type="JINJA",
39
                  template="""\
40
                      {{ question }}\
41
                  """,
42
              ),
43
          ],
44
      ),
45
  ]

1
# workflow.py
2
from vellum.workflows import BaseWorkflow
3
from vellum.workflows.state import BaseState
4
5
class BasicRAGWorkflow(BaseWorkflow[Inputs, BaseState]):
6
  graph = ExtractUserMessage >> SearchNode >> FormatSearchResultsNode >> PromptNode
7
8
class Outputs(BaseWorkflow.Outputs):
9
  result = PromptNode.Outputs.text

1
## From any file / function from which you want to reference the Workflow
2
workflow = BasicRAGWorkflow()

1
## From any file / function from which you want to run the Workflow
2
terminal_event = workflow.run(
3
  inputs=Inputs(
4
    chat_history=[
5
        ChatMessageRequest(
6
            role="USER",
7
            text="How often is employee training??",
8
        )
9
    ]
10
  )
11
)
12
## Output:
13
print(terminal_event.outputs.result)
14
15
"""
16
Employee training, as outlined in the Information Security Policy
17
  occurs on an annual basis. All new hires are required to complete
18
  information security awareness training as part of their new employee
19
  onboarding process and then annually thereafter. This ongoing training
20
  includes security and privacy requirements, the correct use of 
21
  information assets and facilities, and, consistent with assigned roles 
22
  and responsibilities, incident response and contingency training. 
23
  Additionally, individuals responsible for supporting or writing code for 
24
  internet-facing applications or internal applications that handle customer 
25
  information must complete annual security training specific to secure coding 
26
  practices, which includes OWASP secure development principles and OWASP top 10 
27
  vulnerability awareness for the most recent year available.
28
29
Citation: Policy Information Security Policy - v1.pdf & Policy Software Development Life Cycle Policy - v1.pdf\
30
"""