mirror of
https://github.com/hoshikawa2/nemo_guardrails_oci_generative_ai.git
synced 2026-07-10 08:24:20 +00:00
113 lines
3.3 KiB
Python
113 lines
3.3 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import List, Optional, TypedDict
|
|
|
|
from langchain_core.language_models import BaseLLM
|
|
from nemoguardrails import RailsConfig
|
|
from nemoguardrails.actions.actions import ActionResult, action
|
|
from nemoguardrails.actions.llm.utils import llm_call
|
|
from nemoguardrails.context import llm_call_info_var
|
|
from nemoguardrails.llm.taskmanager import LLMTaskManager
|
|
from nemoguardrails.llm.types import Task
|
|
from nemoguardrails.logging.explain import LLMCallInfo
|
|
from nemoguardrails.utils import new_event_dict
|
|
|
|
|
|
class KeywordDetectionResult(TypedDict):
|
|
is_match: bool
|
|
text: str
|
|
detections: List[str]
|
|
|
|
|
|
@action(is_system_action=True)
|
|
async def self_check_input(
|
|
llm_task_manager: LLMTaskManager,
|
|
context: Optional[dict] = None,
|
|
llm: Optional[BaseLLM] = None,
|
|
config: Optional[RailsConfig] = None,
|
|
**kwargs,
|
|
):
|
|
"""Run the self check prompt without forcing max_tokens.
|
|
|
|
Some OpenAI-compatible backends used for moderation return an empty message
|
|
when max_tokens is constrained too aggressively, which NeMo then treats as unsafe.
|
|
"""
|
|
|
|
user_input = (context or {}).get("user_message")
|
|
task = Task.SELF_CHECK_INPUT
|
|
|
|
if not user_input:
|
|
return True
|
|
|
|
prompt = llm_task_manager.render_task_prompt(
|
|
task=task,
|
|
context={"user_input": user_input},
|
|
)
|
|
stop = llm_task_manager.get_stop_tokens(task=task)
|
|
|
|
llm_call_info_var.set(LLMCallInfo(task=task.value))
|
|
|
|
response = await llm_call(
|
|
llm,
|
|
prompt,
|
|
stop=stop,
|
|
llm_params={
|
|
"temperature": config.lowest_temperature if config else 0,
|
|
},
|
|
)
|
|
|
|
if llm_task_manager.has_output_parser(task):
|
|
result = llm_task_manager.parse_task_output(task, output=response)
|
|
else:
|
|
result = llm_task_manager.parse_task_output(
|
|
task,
|
|
output=response,
|
|
forced_output_parser="is_content_safe",
|
|
)
|
|
|
|
is_safe = result[0]
|
|
|
|
if not is_safe:
|
|
return ActionResult(
|
|
return_value=False,
|
|
events=[new_event_dict("mask_prev_user_message", intent="unanswerable message")],
|
|
)
|
|
|
|
return True
|
|
|
|
|
|
@action(is_system_action=True)
|
|
async def detect_keywords(
|
|
source: str,
|
|
text: str,
|
|
config: RailsConfig,
|
|
) -> KeywordDetectionResult:
|
|
if source not in ("input", "output", "retrieval"):
|
|
raise ValueError("source must be one of 'input', 'output', or 'retrieval'")
|
|
|
|
keyword_config = getattr(config.rails.config, "keyword_detection", None)
|
|
if keyword_config is None:
|
|
return KeywordDetectionResult(is_match=False, text=text, detections=[])
|
|
|
|
options = getattr(keyword_config, source, None)
|
|
if options is None or not text:
|
|
return KeywordDetectionResult(is_match=False, text=text, detections=[])
|
|
|
|
keywords = getattr(options, "keywords", None) or []
|
|
if not keywords:
|
|
return KeywordDetectionResult(is_match=False, text=text, detections=[])
|
|
|
|
haystack = text.lower() if getattr(options, "case_insensitive", True) else text
|
|
|
|
matched = []
|
|
for keyword in keywords:
|
|
needle = keyword.lower() if getattr(options, "case_insensitive", True) else keyword
|
|
if needle and needle in haystack:
|
|
matched.append(keyword)
|
|
|
|
return KeywordDetectionResult(
|
|
is_match=bool(matched),
|
|
text=text,
|
|
detections=matched,
|
|
)
|