Skip to content

Commit ef295b4

Browse files
RobGeadaRobGeada
authored
Merge pull request #1 from Pouyanpi/pr-1874-first-class-self-check-tasks
fix(self-check): preserve custom task routing
2 parents fe720c6 + 2895b99 commit ef295b4

10 files changed

Lines changed: 677 additions & 190 deletions

File tree

‎nemoguardrails/library/self_check/input_check/actions.py‎

Lines changed: 29 additions & 60 deletions
Original file line numberDiff line numberDiff line change
@@ -14,49 +14,35 @@
1414
# limitations under the License.
1515

1616
import logging
17-
from typing import Dict, Optional
17+
from typing import Dict, List, Optional
1818

1919
from nemoguardrails import RailsConfig
2020
from nemoguardrails.actions.actions import ActionResult, action
21-
from nemoguardrails.actions.llm.utils import llm_call, warn_if_truncated
22-
from nemoguardrails.context import llm_call_info_var
21+
from nemoguardrails.library.self_check.utils import (
22+
SELF_CHECK_INPUT_DEFAULT_TASK,
23+
SELF_CHECK_INPUT_FLOW,
24+
SELF_CHECK_INPUT_TASK_PARAM,
25+
resolve_self_check_task,
26+
run_self_check_task,
27+
)
2328
from nemoguardrails.llm.taskmanager import LLMTaskManager
24-
from nemoguardrails.logging.explain import LLMCallInfo
2529
from nemoguardrails.types import LLMModel
2630
from nemoguardrails.utils import new_event_dict
2731

2832
log = logging.getLogger(__name__)
2933

30-
DEFAULT_TASK = "self_check_input"
31-
32-
33-
def _get_llm(
34-
llms: Dict[str, LLMModel], task: str, default_task: str, default_llm: Optional[LLMModel]
35-
) -> Optional[LLMModel]:
36-
if task in llms:
37-
return llms[task]
38-
elif default_task in llms:
39-
log.debug(f"No model found with type={task}, falling back to default {default_task} model")
40-
return llms[default_task]
41-
elif default_llm is not None:
42-
log.debug(f"No model found with type={task} or type={default_task}, falling back to main model")
43-
return default_llm
44-
else:
45-
error_msg = (
46-
f"No matching model for task={task} found. "
47-
f"Please configure a model with type={task}, type={default_task}, or type=main"
48-
)
49-
raise ValueError(error_msg)
34+
DEFAULT_TASK = SELF_CHECK_INPUT_DEFAULT_TASK
5035

5136

5237
@action(is_system_action=True)
5338
async def self_check_input(
5439
llms: Dict[str, LLMModel],
5540
llm_task_manager: LLMTaskManager,
5641
context: Optional[dict] = None,
42+
events: Optional[List[dict]] = None,
5743
llm: Optional[LLMModel] = None,
5844
config: Optional[RailsConfig] = None,
59-
task: str = DEFAULT_TASK,
45+
task: Optional[str] = None,
6046
**kwargs,
6147
):
6248
"""Checks the input from the user.
@@ -68,55 +54,38 @@ async def self_check_input(
6854
True if the input should be allowed, False otherwise.
6955
"""
7056

71-
_MAX_TOKENS = 1024
57+
context = context or {}
7258
user_input = context.get("user_message")
7359

74-
# guard against an unset $task variable
75-
if task.startswith("$") and task.endswith("_task"):
76-
task = DEFAULT_TASK
77-
78-
llm = _get_llm(llms, task, default_task=DEFAULT_TASK, default_llm=llm)
60+
task = resolve_self_check_task(
61+
task,
62+
context,
63+
events,
64+
triggered_rail_key="triggered_input_rail",
65+
start_rail_event_type="StartInputRail",
66+
flow_id=SELF_CHECK_INPUT_FLOW,
67+
task_param=SELF_CHECK_INPUT_TASK_PARAM,
68+
default_task=DEFAULT_TASK,
69+
)
7970

8071
if user_input:
81-
prompt = llm_task_manager.render_task_prompt(
72+
is_safe, response = await run_self_check_task(
8273
task=task,
83-
context={
74+
prompt_context={
8475
"user_input": user_input,
8576
},
77+
llms=llms,
78+
default_task=DEFAULT_TASK,
79+
main_llm=llm,
80+
llm_task_manager=llm_task_manager,
81+
lowest_temperature=config.lowest_temperature,
8682
)
87-
stop = llm_task_manager.get_stop_tokens(task=task)
88-
max_tokens = llm_task_manager.get_max_tokens(task=task)
89-
max_tokens = max_tokens or _MAX_TOKENS
90-
91-
# Initialize the LLMCallInfo object
92-
llm_call_info_var.set(LLMCallInfo(task=task))
93-
94-
llm_response = await llm_call(
95-
llm,
96-
prompt,
97-
stop=stop,
98-
llm_params={
99-
"temperature": config.lowest_temperature,
100-
"max_tokens": max_tokens,
101-
},
102-
)
103-
warn_if_truncated(llm_response, task)
104-
response = llm_response.content
10583

10684
if task == DEFAULT_TASK:
10785
log.info(f"Input self-checking result is: `{response}`.")
10886
else:
10987
log.info(f"Input self-checking result for task={task} is: `{response}`.")
11088

111-
# for sake of backward compatibility
112-
# if the output_parser is not registered we will use the default one
113-
if llm_task_manager.has_output_parser(task):
114-
result = llm_task_manager.parse_task_output(task, output=response)
115-
else:
116-
result = llm_task_manager.parse_task_output(task, output=response, forced_output_parser="is_content_safe")
117-
118-
is_safe = result[0]
119-
12089
if not is_safe:
12190
return ActionResult(
12291
return_value=False,

‎nemoguardrails/library/self_check/input_check/flows.co‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
1-
flow self check input $input_task="self_check_input"
2-
$allowed = await SelfCheckInputAction(task=$input_task)
1+
flow self check input $task="self_check_input"
2+
$allowed = await SelfCheckInputAction(task=$task)
33

44
if not $allowed
55
if $system.config.enable_rails_exceptions

‎nemoguardrails/library/self_check/input_check/flows.v1.co‎

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2,10 +2,7 @@ define bot refuse to respond
22
"I'm sorry, I can't respond to that."
33

44
define flow self check input
5-
if not $input_task
6-
$input_task = "self_check_input"
7-
8-
$allowed = execute self_check_input(task=$input_task)
5+
$allowed = execute self_check_input
96

107
if not $allowed
118
if $config.enable_rails_exceptions

‎nemoguardrails/library/self_check/output_check/actions.py‎

Lines changed: 29 additions & 60 deletions
Original file line numberDiff line numberDiff line change
@@ -14,48 +14,34 @@
1414
# limitations under the License.
1515

1616
import logging
17-
from typing import Dict, Optional
17+
from typing import Dict, List, Optional
1818

1919
from nemoguardrails import RailsConfig
2020
from nemoguardrails.actions import action
21-
from nemoguardrails.actions.llm.utils import llm_call, warn_if_truncated
22-
from nemoguardrails.context import llm_call_info_var
21+
from nemoguardrails.library.self_check.utils import (
22+
SELF_CHECK_OUTPUT_DEFAULT_TASK,
23+
SELF_CHECK_OUTPUT_FLOW,
24+
SELF_CHECK_OUTPUT_TASK_PARAM,
25+
resolve_self_check_task,
26+
run_self_check_task,
27+
)
2328
from nemoguardrails.llm.taskmanager import LLMTaskManager
24-
from nemoguardrails.logging.explain import LLMCallInfo
2529
from nemoguardrails.types import LLMModel
2630

2731
log = logging.getLogger(__name__)
2832

29-
DEFAULT_TASK = "self_check_output"
30-
31-
32-
def _get_llm(
33-
llms: Dict[str, LLMModel], task: str, default_task: str, default_llm: Optional[LLMModel]
34-
) -> Optional[LLMModel]:
35-
if task in llms:
36-
return llms[task]
37-
elif default_task in llms:
38-
log.debug(f"No model found with type={task}, falling back to default {default_task} model")
39-
return llms[default_task]
40-
elif default_llm is not None:
41-
log.debug(f"No model found with type={task} or type={default_task}, falling back to main model")
42-
return default_llm
43-
else:
44-
error_msg = (
45-
f"No matching model for task={task} found. "
46-
f"Please configure a model with type={task}, type={default_task}, or type=main"
47-
)
48-
raise ValueError(error_msg)
33+
DEFAULT_TASK = SELF_CHECK_OUTPUT_DEFAULT_TASK
4934

5035

5136
@action(is_system_action=True, output_mapping=lambda value: not value)
5237
async def self_check_output(
5338
llms: Dict[str, LLMModel],
5439
llm_task_manager: LLMTaskManager,
5540
context: Optional[dict] = None,
41+
events: Optional[List[dict]] = None,
5642
llm: Optional[LLMModel] = None,
5743
config: Optional[RailsConfig] = None,
58-
task: str = DEFAULT_TASK,
44+
task: Optional[str] = None,
5945
**kwargs,
6046
):
6147
"""Checks if the output from the bot.
@@ -70,57 +56,40 @@ async def self_check_output(
7056
True if the output should be allowed, False otherwise.
7157
"""
7258

73-
_MAX_TOKENS = 1024
59+
context = context or {}
7460
bot_response = context.get("bot_message")
7561
user_input = context.get("user_message")
7662
bot_thinking = context.get("bot_thinking")
7763

78-
# guard against an unset $task variable
79-
if task.startswith("$") and task.endswith("_task"):
80-
task = DEFAULT_TASK
81-
82-
llm = _get_llm(llms, task, default_task=DEFAULT_TASK, default_llm=llm)
64+
task = resolve_self_check_task(
65+
task,
66+
context,
67+
events,
68+
triggered_rail_key="triggered_output_rail",
69+
start_rail_event_type="StartOutputRail",
70+
flow_id=SELF_CHECK_OUTPUT_FLOW,
71+
task_param=SELF_CHECK_OUTPUT_TASK_PARAM,
72+
default_task=DEFAULT_TASK,
73+
)
8374

8475
if bot_response:
85-
prompt = llm_task_manager.render_task_prompt(
76+
is_safe, response = await run_self_check_task(
8677
task=task,
87-
context={
78+
prompt_context={
8879
"user_input": user_input,
8980
"bot_response": bot_response,
9081
"bot_thinking": bot_thinking,
9182
},
83+
llms=llms,
84+
default_task=DEFAULT_TASK,
85+
main_llm=llm,
86+
llm_task_manager=llm_task_manager,
87+
lowest_temperature=config.lowest_temperature,
9288
)
93-
stop = llm_task_manager.get_stop_tokens(task=task)
94-
max_tokens = llm_task_manager.get_max_tokens(task=task)
95-
max_tokens = max_tokens or _MAX_TOKENS
96-
97-
# Initialize the LLMCallInfo object
98-
llm_call_info_var.set(LLMCallInfo(task=task))
99-
100-
llm_response = await llm_call(
101-
llm,
102-
prompt,
103-
stop=stop,
104-
llm_params={
105-
"temperature": config.lowest_temperature,
106-
"max_tokens": max_tokens,
107-
},
108-
)
109-
warn_if_truncated(llm_response, task)
110-
response = llm_response.content
11189

11290
if task == DEFAULT_TASK:
11391
log.info(f"Output self-checking result is: `{response}`.")
11492
else:
11593
log.info(f"Output self-checking result for task={task} is: `{response}`.")
11694

117-
# for sake of backward compatibility
118-
# if the output_parser is not registered we will use the default one
119-
if llm_task_manager.has_output_parser(task):
120-
result = llm_task_manager.parse_task_output(task, output=response)
121-
else:
122-
result = llm_task_manager.parse_task_output(task, output=response, forced_output_parser="is_content_safe")
123-
124-
is_safe = result[0]
125-
12695
return is_safe

‎nemoguardrails/library/self_check/output_check/flows.co‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
1-
flow self check output $output_task="self_check_output"
2-
$allowed = await SelfCheckOutputAction(task=$output_task)
1+
flow self check output $task="self_check_output"
2+
$allowed = await SelfCheckOutputAction(task=$task)
33

44
if not $allowed
55
if $system.config.enable_rails_exceptions

‎nemoguardrails/library/self_check/output_check/flows.v1.co‎

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2,10 +2,7 @@ define bot refuse to respond
22
"I'm sorry, I can't respond to that."
33

44
define flow self check output
5-
if not $output_task
6-
$output_task = "self_check_output"
7-
8-
$allowed = execute self_check_output(task=$output_task)
5+
$allowed = execute self_check_output
96

107
if not $allowed
118
if $config.enable_rails_exceptions

0 commit comments

Comments
 (0)