Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add ability to choose LLM prompt template for value generation and flow generation #903

Open
wants to merge 3 commits into
base: develop
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 18 additions & 9 deletions nemoguardrails/actions/v2_x/generation.py
Original file line number Diff line number Diff line change
Expand Up @@ -763,16 +763,21 @@ async def generate_value(
if "GenerateValueAction" not in result.text:
examples += f"{result.text}\n\n"

llm_call_info_var.set(
LLMCallInfo(task=Task.GENERATE_VALUE_FROM_INSTRUCTION.value)
out_variables: dict[str, Any] = {}
rendered_instructions = self.llm_task_manager._render_string(
instructions,
out_variables=out_variables,
)

task = out_variables.get("template", Task.GENERATE_VALUE_FROM_INSTRUCTION)
llm_call_info_var.set(LLMCallInfo(task=task))

prompt = self.llm_task_manager.render_task_prompt(
task=Task.GENERATE_VALUE_FROM_INSTRUCTION,
task=task,
events=events,
context={
"examples": examples,
"instructions": instructions,
"instructions": rendered_instructions,
"var_name": var_name if var_name else "result",
"context": state.context,
},
Expand Down Expand Up @@ -869,16 +874,20 @@ async def generate_flow(
render_context["tool_names"] = ", ".join(tool_names)

# TODO: add the context of the flow
out_variables = {}
flow_nld = self.llm_task_manager._render_string(
textwrap.dedent(docstring), context=render_context, events=events
textwrap.dedent(docstring),
context=render_context,
events=events,
out_variables=out_variables,
)

llm_call_info_var.set(
LLMCallInfo(task=Task.GENERATE_FLOW_CONTINUATION_FROM_NLD.value)
)
task = out_variables.get("template", Task.GENERATE_FLOW_CONTINUATION_FROM_NLD)

llm_call_info_var.set(LLMCallInfo(task=task))

prompt = self.llm_task_manager.render_task_prompt(
task=Task.GENERATE_FLOW_CONTINUATION_FROM_NLD,
task=task,
events=events,
context={
"flow_nld": flow_nld,
Expand Down
12 changes: 11 additions & 1 deletion nemoguardrails/llm/taskmanager.py
Original file line number Diff line number Diff line change
Expand Up @@ -109,12 +109,14 @@ def _render_string(
template_str: str,
context: Optional[dict] = None,
events: Optional[List[dict]] = None,
out_variables: Optional[dict] = None,
) -> str:
"""Render a template using the provided context information.

:param template_str: The template to render.
:param context: The context for rendering the prompt.
:param events: The history of events so far.
:param out_variables: If not None the dict will be populated with variables set in the template
:return: The rendered template.
:rtype: str.
"""
Expand Down Expand Up @@ -152,7 +154,15 @@ def _render_string(

render_context[variable] = value

return template.render(render_context)
rendered = template.render(render_context)

if out_variables is not None:
mod = template.module
template_vars = {
n: getattr(mod, n) for n in dir(mod) if not n.startswith("_")
}
out_variables.update(template_vars)
return rendered

def _render_messages(
self,
Expand Down
6 changes: 4 additions & 2 deletions tests/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@
CallbackManagerForLLMRun,
)
from langchain_core.language_models.llms import LLM
from pydantic import Field

from nemoguardrails import LLMRails, RailsConfig
from nemoguardrails.colang import parse_colang_file
Expand All @@ -40,6 +41,7 @@ class FakeLLM(LLM):
"""Fake LLM wrapper for testing purposes."""

responses: List
prompt_history: List[str] = Field(default_factory=list, exclude=True)
i: int = 0
streaming: bool = False

Expand All @@ -60,7 +62,7 @@ def _call(
f"No responses available for query number {self.i + 1} in FakeLLM. "
"Most likely, too many LLM calls are made or additional responses need to be provided."
)

self.prompt_history.append(prompt)
response = self.responses[self.i]
self.i += 1
return response
Expand All @@ -77,7 +79,7 @@ async def _acall(
f"No responses available for query number {self.i + 1} in FakeLLM. "
"Most likely, too many LLM calls are made or additional responses need to be provided."
)

self.prompt_history.append(prompt)
response = self.responses[self.i]
self.i += 1

Expand Down
122 changes: 122 additions & 0 deletions tests/v2_x/test_llm_template_selection.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,122 @@
# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import logging

from rich.logging import RichHandler

from nemoguardrails import RailsConfig
from tests.utils import TestChat

FORMAT = "%(message)s"
logging.basicConfig(
level=logging.DEBUG,
format=FORMAT,
datefmt="[%X,%f]",
handlers=[RichHandler(markup=True)],
)

config_1 = """
colang_version: "2.x"

models:
- type: main
engine: openai
model: gpt-3.5-turbo-instruct

prompts:
- task: generate_antonym
models:
- openai/gpt-3.5-turbo
- openai/gpt-4
messages:
- type: user
content: |-
Generate the antonym of the bot expression below. Use the syntax: bot say "<antonym goes here>".
- type: user
content: |-
YOUR TASK:
{{ flow_nld }}

- task: repeat
models:
- openai/gpt-3.5-turbo
- openai/gpt-4
messages:
- type: system
content: |
Your are a value generation bot that needs to generate a value for the ${{ var_name }} variable based on instructions form the user.
Be very precised and always pick the most suitable variable type (e.g. double quotes for strings). Only generated the value and do not provide any additional response.
- type: user
content: |
{{ instructions }} three times
Assign the generated value to:
${{ var_name }} =

"""


def test_template_choice_in_value_generation():
"""Test template selection in value generation"""
config = RailsConfig.from_content(
colang_content="""
flow main
match UtteranceUserActionFinished(final_transcript="hi")
$test = ..."a random bird name{{% set template = 'repeat' %}}"
await UtteranceBotAction(script=$test)
""",
yaml_content=config_1,
)

chat = TestChat(
config,
llm_completions=["'parrot, raven, peacock'"],
)

expected_prompt = "System: Your are a value generation bot that needs to generate a value for the $test variable based on instructions form the user.\nBe very precised and always pick the most suitable variable type (e.g. double quotes for strings). Only generated the value and do not provide any additional response.\nHuman: a random bird name three times\nAssign the generated value to:\n$test ="

chat >> "hi"
chat << "parrot, raven, peacock"
assert chat.llm.prompt_history[0] == expected_prompt


def test_template_choice_in_flow_generation():
"""Test template selection in flow generation"""
config = RailsConfig.from_content(
colang_content="""
import core
flow generate antonym
\"\"\"
{% set template = "generate_antonym" %}
bot say "lucky"
\"\"\"
...
flow main
match UtteranceUserActionFinished(final_transcript="hi")
generate antonym
""",
yaml_content=config_1,
)

chat = TestChat(
config,
llm_completions=["bot say 'unfortunate'"],
)

expected_prompt = 'Human: Generate the antonym of the bot expression below. Use the syntax: bot say "<antonym goes here>".\nHuman: YOUR TASK:\n\n\nbot say "lucky"'

chat >> "hi"
chat << "unfortunate"
assert chat.llm.prompt_history[0] == expected_prompt
Loading