diff --git a/nemoguardrails/http/__init__.py b/nemoguardrails/http/__init__.py index d599a01e87..3b9a424a7e 100644 --- a/nemoguardrails/http/__init__.py +++ b/nemoguardrails/http/__init__.py @@ -27,7 +27,7 @@ from nemoguardrails.http.retry import RetryingHTTPClient, RetryPolicy from nemoguardrails.http.runtime import create_http_client from nemoguardrails.http.transport import HttpxHTTPClient -from nemoguardrails.http.types import HTTPRequest, HTTPResponse +from nemoguardrails.http.types import HTTPRequest, HTTPResponse, HTTPTLSConfig __all__ = [ "ClosableHTTPClient", @@ -39,6 +39,7 @@ "HTTPResponseDecodeError", "HTTPStatusError", "HTTPTimeoutError", + "HTTPTLSConfig", "HttpxHTTPClient", "RetryPolicy", "RetryingHTTPClient", diff --git a/nemoguardrails/http/runtime.py b/nemoguardrails/http/runtime.py index e0853473e6..d7692905df 100644 --- a/nemoguardrails/http/runtime.py +++ b/nemoguardrails/http/runtime.py @@ -20,6 +20,7 @@ from nemoguardrails.http.client import ClosableHTTPClient from nemoguardrails.http.retry import RetryingHTTPClient, RetryPolicy from nemoguardrails.http.transport import HttpxHTTPClient +from nemoguardrails.http.types import HTTPTLSConfig def create_http_client( @@ -29,6 +30,7 @@ def create_http_client( limits: httpx.Limits | None = None, retry_policy: RetryPolicy | None = None, follow_redirects: bool = False, + tls: HTTPTLSConfig | None = None, ) -> ClosableHTTPClient: """Create a transport-neutral HTTP client with optional retry behavior. @@ -39,6 +41,7 @@ def create_http_client( retry_policy: Optional policy applied to created or injected clients. follow_redirects: Whether a client created by this function follows redirects. + tls: TLS settings for a client created by this function. Returns: A closable transport-neutral HTTP client. @@ -48,7 +51,7 @@ def create_http_client( injected client. When ``httpx_client`` is provided, it remains caller-owned and ``timeout``, - ``limits``, and ``follow_redirects`` must retain their defaults. + ``limits``, ``follow_redirects``, and ``tls`` must retain their defaults. ``retry_policy`` is still applied when supplied. """ @@ -57,6 +60,7 @@ def create_http_client( timeout=timeout, limits=limits, follow_redirects=follow_redirects, + tls=tls, ) client: ClosableHTTPClient = transport if retry_policy is not None: diff --git a/nemoguardrails/http/transport.py b/nemoguardrails/http/transport.py index 5639ce7a6f..e1659d88e6 100644 --- a/nemoguardrails/http/transport.py +++ b/nemoguardrails/http/transport.py @@ -22,7 +22,7 @@ from nemoguardrails.http._url import sanitize_url from nemoguardrails.http.errors import HTTPConnectionError, HTTPTimeoutError -from nemoguardrails.http.types import HTTPResponse +from nemoguardrails.http.types import HTTPResponse, HTTPTLSConfig _DEFAULT_TIMEOUT_SECONDS = 30.0 @@ -51,6 +51,7 @@ def __init__( timeout: float | None = _DEFAULT_TIMEOUT_SECONDS, limits: httpx.Limits | None = None, follow_redirects: bool = False, + tls: HTTPTLSConfig | None = None, ): """Initialize the HTTPX adapter. @@ -59,22 +60,34 @@ def __init__( timeout: Default total timeout for owned clients, in seconds. limits: Connection-pool limits for an owned client. follow_redirects: Whether an owned client follows redirects. + tls: TLS settings for an owned client. Raises: ValueError: If the configured timeout is not positive or owned-client options are supplied with an injected client. """ + if client is not None and tls is not None: + raise ValueError("TLS configuration cannot be combined with an injected HTTPX client") if client is not None and (timeout != _DEFAULT_TIMEOUT_SECONDS or limits is not None or follow_redirects): raise ValueError("Owned-client options cannot be used with an injected HTTPX client") if timeout is not None and timeout <= 0: raise ValueError("HTTP timeout must be greater than zero") + tls_config = tls or HTTPTLSConfig() + verify: bool | str = tls_config.verify + if tls_config.verify and tls_config.ca_bundle is not None: + verify = tls_config.ca_bundle + cert = None + if tls_config.client_certificate is not None and tls_config.client_key is not None: + cert = (tls_config.client_certificate, tls_config.client_key) self._owns_client = client is None self._timeout = timeout if self._owns_client else None self._client = client or httpx.AsyncClient( timeout=None, limits=limits or httpx.Limits(max_connections=100, max_keepalive_connections=20), follow_redirects=follow_redirects, + verify=verify, + cert=cert, ) self._closed = False diff --git a/nemoguardrails/http/types.py b/nemoguardrails/http/types.py index 14c0d61e8e..77daf42e9e 100644 --- a/nemoguardrails/http/types.py +++ b/nemoguardrails/http/types.py @@ -23,6 +23,24 @@ from nemoguardrails.http.errors import HTTPResponseDecodeError, HTTPStatusError +@dataclass(frozen=True) +class HTTPTLSConfig: + """Configure TLS verification and optional mutual authentication. + + A custom CA bundle applies only when verification is enabled. Client + certificate and key paths must be supplied together. + """ + + verify: bool = True + ca_bundle: str | None = None + client_certificate: str | None = None + client_key: str | None = None + + def __post_init__(self) -> None: + if bool(self.client_certificate) != bool(self.client_key): + raise ValueError("client_certificate and client_key must be configured together") + + @dataclass(frozen=True) class HTTPRequest: """Describe an outbound HTTP request without transport-specific objects. diff --git a/nemoguardrails/library/ai_defense/actions.py b/nemoguardrails/library/ai_defense/actions.py index ac09b494e6..e4afbaa7a9 100644 --- a/nemoguardrails/library/ai_defense/actions.py +++ b/nemoguardrails/library/ai_defense/actions.py @@ -17,13 +17,18 @@ import logging import os +from collections.abc import Mapping from typing import Any, Optional -import httpx - from nemoguardrails import RailsConfig from nemoguardrails.actions import action from nemoguardrails.actions.rail_outcome import RailOutcome +from nemoguardrails.http import ( + HTTPClient, + HTTPClientError, + HTTPResponseDecodeError, + http_call, +) log = logging.getLogger(__name__) @@ -37,13 +42,44 @@ def _ai_defense_outcome(is_blocked: bool) -> RailOutcome: return RailOutcome.allow(metadata={"is_blocked": is_blocked}) +def _ai_defense_failure_outcome(fail_open: bool, failure: str) -> RailOutcome: + if fail_open: + log.warning("%s, but fail_open=True, allowing content.", failure) + return _ai_defense_outcome(False) + log.warning("%s, fail_open=False, blocking content.", failure) + return _ai_defense_outcome(True) + + @action(is_system_action=True) async def ai_defense_inspect( config: RailsConfig, user_prompt: Optional[str] = None, bot_response: Optional[str] = None, + http_client: Optional[HTTPClient] = None, **kwargs, ) -> RailOutcome: + """Inspect a user prompt or bot response with Cisco AI Defense. + + A safe response is allowed and an unsafe response is blocked. Transport + failures and malformed responses follow the configured ``fail_open`` + policy, which defaults to fail closed. + + Args: + config: Rails configuration containing the AI Defense settings. + user_prompt: User content to inspect. + bot_response: Bot content to inspect. Takes precedence over + ``user_prompt`` when both are provided. + http_client: Optional caller-owned HTTP client. + **kwargs: Additional action parameters. The ``user`` value, when + present, is forwarded as request metadata. + + Returns: + The allow or block outcome from AI Defense or the failure policy. + + Raises: + ValueError: If the required environment variables or content are + missing. + """ # Get configuration with defaults ai_defense_config = getattr(config.rails.config, "ai_defense", None) timeout = ai_defense_config.timeout if ai_defense_config else DEFAULT_TIMEOUT @@ -89,44 +125,50 @@ async def ai_defense_inspect( if metadata: payload["metadata"] = metadata - async with httpx.AsyncClient() as client: - try: - resp = await client.post(api_endpoint, headers=headers, json=payload, timeout=timeout) - resp.raise_for_status() - data = resp.json() - except (httpx.HTTPStatusError, httpx.TimeoutException, httpx.RequestError) as e: - msg = f"Error calling AI Defense API: {e}" - log.error(msg) - if fail_open: - # Fail open: allow content when API call fails - log.warning("AI Defense API call failed, but fail_open=True, allowing content.") - return _ai_defense_outcome(False) - else: - # Fail closed: block content when API call fails - log.warning("AI Defense API call failed, fail_open=False, blocking content.") - return _ai_defense_outcome(True) - - # Compose a consistent return structure for flows - # Handle malformed responses based on fail_open setting - if "is_safe" not in data: - # Malformed response - respect fail_open setting - if fail_open: - log.warning( - "AI Defense API returned malformed response (missing 'is_safe'), but fail_open=True, allowing content." - ) - is_blocked = False - else: - log.warning( - "AI Defense API returned malformed response (missing 'is_safe'), fail_open=False, blocking content." - ) - is_blocked = True + try: + response = await http_call( + http_client, + "POST", + api_endpoint, + headers=headers, + json=payload, + timeout=timeout, + ) + data = response.json() + except HTTPResponseDecodeError as e: + log.error("AI Defense API returned malformed JSON: %s", e) + return _ai_defense_failure_outcome(fail_open, "AI Defense API returned malformed JSON") + except HTTPClientError as e: + log.error("Error calling AI Defense API: %s", e) + return _ai_defense_failure_outcome(fail_open, "AI Defense API call failed") + + if not isinstance(data, Mapping): + return _ai_defense_failure_outcome( + fail_open, + "AI Defense API returned malformed response (expected an object)", + ) + + # Compose a consistent return structure for flows + # Handle malformed responses based on fail_open setting + if "is_safe" not in data: + # Malformed response - respect fail_open setting + if fail_open: + log.warning( + "AI Defense API returned malformed response (missing 'is_safe'), but fail_open=True, allowing content." + ) + is_blocked = False else: - is_blocked = not bool(data.get("is_safe", False)) + log.warning( + "AI Defense API returned malformed response (missing 'is_safe'), fail_open=False, blocking content." + ) + is_blocked = True + else: + is_blocked = not bool(data.get("is_safe", False)) - rules = data.get("rules") or [] - if is_blocked and rules: - entries = [f"{r.get('rule_name')} ({r.get('classification')})" for r in rules if isinstance(r, dict)] - if entries: - log.debug("AI Defense matched rules: %s", ", ".join(entries)) + rules = data.get("rules") or [] + if is_blocked and rules: + entries = [f"{r.get('rule_name')} ({r.get('classification')})" for r in rules if isinstance(r, dict)] + if entries: + log.debug("AI Defense matched rules: %s", ", ".join(entries)) - return _ai_defense_outcome(is_blocked) + return _ai_defense_outcome(is_blocked) diff --git a/nemoguardrails/library/autoalign/actions.py b/nemoguardrails/library/autoalign/actions.py index 62e7b6e543..c728c218d7 100644 --- a/nemoguardrails/library/autoalign/actions.py +++ b/nemoguardrails/library/autoalign/actions.py @@ -19,10 +19,9 @@ import os from typing import Any, Dict, List, Optional -import aiohttp - from nemoguardrails.actions import action from nemoguardrails.actions.rail_outcome import RailOutcome, TransformTarget +from nemoguardrails.http import HTTPClient, HTTPResponse, http_call from nemoguardrails.llm.taskmanager import LLMTaskManager log = logging.getLogger(__name__) @@ -144,12 +143,32 @@ def _autoalign_score_outcome(score: float, threshold: float) -> RailOutcome: return RailOutcome.allow(metadata=metadata) +async def _autoalign_post( + http_client: HTTPClient | None, + request_url: str, + headers: dict[str, str], + request_body: dict[str, Any], +) -> HTTPResponse: + response = await http_call( + http_client, + "POST", + request_url, + headers=headers, + json=request_body, + raise_for_status=False, + ) + if response.status_code != 200: + raise ValueError(f"AutoAlign call failed with status code {response.status_code}.\nDetails: {response.text}") + return response + + async def autoalign_infer( request_url: str, text: str, task_config: Optional[Dict[Any, Any]] = None, show_toxic_phrases: bool = False, multi_language: bool = False, + http_client: Optional[HTTPClient] = None, ): """Checks whether the given text passes through the applied guardrails.""" api_key = os.environ.get("AUTOALIGN_API_KEY") @@ -168,22 +187,18 @@ async def autoalign_infer( guardrails_configured = [] - async with aiohttp.ClientSession() as session: - async with session.post( - url=request_url, - headers=headers, - json=request_body, - ) as response: - if response.status != 200: - raise ValueError( - f"AutoAlign call failed with status code {response.status}.\nDetails: {await response.text()}" - ) - async for line in response.content: - line_text = line.strip() - if len(line_text) > 0: - resp = json.loads(line_text) - guardrails_configured.append(resp) - processed_response = process_autoalign_output(guardrails_configured, show_toxic_phrases) + response = await _autoalign_post( + http_client, + request_url, + headers, + request_body, + ) + for line in response.content.splitlines(): + line_text = line.strip() + if len(line_text) > 0: + resp = json.loads(line_text) + guardrails_configured.append(resp) + processed_response = process_autoalign_output(guardrails_configured, show_toxic_phrases) return processed_response @@ -192,6 +207,7 @@ async def autoalign_groundedness_infer( text: str, documents: List[str], guardrails_config: Optional[Dict[Any, Any]] = None, + http_client: Optional[HTTPClient] = None, ): """Checks the groundedness for the text using the given documents and provides a fact-checking score""" groundness_config = copy.deepcopy(default_groundedness_config) @@ -202,21 +218,17 @@ async def autoalign_groundedness_infer( if guardrails_config: groundness_config.update(guardrails_config) request_body = {"prompt": text, "documents": documents, "config": groundness_config} - async with aiohttp.ClientSession() as session: - async with session.post( - url=request_url, - headers=headers, - json=request_body, - ) as response: - if response.status != 200: - raise ValueError( - f"AutoAlign call failed with status code {response.status}.\nDetails: {await response.text()}" - ) - async for line in response.content: - resp = json.loads(line) - if resp["task"] == "groundedness_checker": - if resp["response"].startswith("Factcheck Score: "): - return float(resp["response"][17:]) + response = await _autoalign_post( + http_client, + request_url, + headers, + request_body, + ) + for line in response.content.splitlines(): + resp = json.loads(line) + if resp["task"] == "groundedness_checker": + if resp["response"].startswith("Factcheck Score: "): + return float(resp["response"][17:]) return 1.0 @@ -226,6 +238,7 @@ async def autoalign_factcheck_infer( bot_message: str, guardrails_config: Optional[Dict[str, Any]] = None, multi_language: bool = False, + http_client: Optional[HTTPClient] = None, ): api_key = os.environ.get("AUTOALIGN_API_KEY") if api_key is None: @@ -244,19 +257,14 @@ async def autoalign_factcheck_infer( "config": guardrails_config, "multi_language": multi_language, } - async with aiohttp.ClientSession() as session: - async with session.post( - url=request_url, - headers=headers, - json=request_body, - ) as response: - if response.status != 200: - raise ValueError( - f"AutoAlign call failed with status code {response.status}.\nDetails: {await response.text()}" - ) - factcheck_response = await response.json() - return factcheck_response["all_overall_fact_scores"][0] - return 1.0 + response = await _autoalign_post( + http_client, + request_url, + headers, + request_body, + ) + factcheck_response = response.json() + return factcheck_response["all_overall_fact_scores"][0] @action(name="autoalign_input_api") @@ -265,9 +273,26 @@ async def autoalign_input_api( context: Optional[dict] = None, show_autoalign_message: bool = True, show_toxic_phrases: bool = False, + http_client: Optional[HTTPClient] = None, **kwargs, ) -> RailOutcome: - """Calls AutoAlign API for the user message and guardrail configuration provided""" + """Calls AutoAlign API for the user message and guardrail configuration provided + + Guardrail matches block, PII redactions transform the user message, and + other results allow it. Configuration, transport, status, and response + parsing failures propagate to the caller. + + Args: + llm_task_manager: Task manager containing the AutoAlign configuration. + context: Runtime context containing the user message. + show_autoalign_message: Whether to log detected violations. + show_toxic_phrases: Whether violation messages include toxic phrases. + http_client: Optional caller-owned HTTP client. + **kwargs: Additional action parameters. + + Returns: + An allow, block, or user-message transform outcome. + """ user_message = context.get("user_message") autoalign_config = llm_task_manager.config.rails.config.autoalign autoalign_api_url = autoalign_config.parameters.get("endpoint") @@ -285,6 +310,7 @@ async def autoalign_input_api( task_config, show_toxic_phrases, multi_language=multi_language, + http_client=http_client, ) if autoalign_response["guardrails_triggered"] and show_autoalign_message: log.warning( @@ -305,9 +331,26 @@ async def autoalign_output_api( context: Optional[dict] = None, show_autoalign_message: bool = True, show_toxic_phrases: bool = False, + http_client: Optional[HTTPClient] = None, **kwargs, ) -> RailOutcome: - """Calls AutoAlign API for the bot message and guardrail configuration provided""" + """Calls AutoAlign API for the bot message and guardrail configuration provided + + Guardrail matches block, PII redactions transform the bot message, and + other results allow it. Configuration, transport, status, and response + parsing failures propagate to the caller. + + Args: + llm_task_manager: Task manager containing the AutoAlign configuration. + context: Runtime context containing the bot message. + show_autoalign_message: Whether to log detected violations. + show_toxic_phrases: Whether violation messages include toxic phrases. + http_client: Optional caller-owned HTTP client. + **kwargs: Additional action parameters. + + Returns: + An allow, block, or bot-message transform outcome. + """ bot_message = context.get("bot_message") autoalign_config = llm_task_manager.config.rails.config.autoalign autoalign_api_url = autoalign_config.parameters.get("endpoint") @@ -325,6 +368,7 @@ async def autoalign_output_api( task_config, show_toxic_phrases, multi_language=multi_language, + http_client=http_client, ) if autoalign_response["guardrails_triggered"] and show_autoalign_message: log.warning( @@ -340,10 +384,26 @@ async def autoalign_groundedness_output_api( context: Optional[dict] = None, factcheck_threshold: float = 0.0, show_autoalign_message: bool = True, + http_client: Optional[HTTPClient] = None, **kwargs, ) -> RailOutcome: """Calls AutoAlign groundedness check API and checks whether the bot message is factually grounded according to given - documents""" + documents + + Scores below ``factcheck_threshold`` block. Configuration, transport, + status, and response parsing failures propagate to the caller. + + Args: + llm_task_manager: Task manager containing the AutoAlign configuration. + context: Runtime context containing the bot message and documents. + factcheck_threshold: Minimum score required to allow the response. + show_autoalign_message: Whether to log detected violations. + http_client: Optional caller-owned HTTP client. + **kwargs: Additional action parameters. + + Returns: + An allow or block outcome containing the score and threshold. + """ bot_message = context.get("bot_message") documents = context.get("relevant_chunks_sep", []) @@ -359,6 +419,7 @@ async def autoalign_groundedness_output_api( text=text, documents=documents, guardrails_config=guardrails_config, + http_client=http_client, ) if score < factcheck_threshold and show_autoalign_message: log.warning( @@ -373,8 +434,23 @@ async def autoalign_factcheck_output_api( context: Optional[dict] = None, factcheck_threshold: float = 0.0, show_autoalign_message: bool = True, + http_client: Optional[HTTPClient] = None, ) -> RailOutcome: - """Calls Autoalign Factchecker API and checks if the user message is factually answered by the bot message""" + """Calls Autoalign Factchecker API and checks if the user message is factually answered by the bot message + + Scores below ``factcheck_threshold`` block. Configuration, transport, + status, and response parsing failures propagate to the caller. + + Args: + llm_task_manager: Task manager containing the AutoAlign configuration. + context: Runtime context containing the user and bot messages. + factcheck_threshold: Minimum score required to allow the response. + show_autoalign_message: Whether to log detected violations. + http_client: Optional caller-owned HTTP client. + + Returns: + An allow or block outcome containing the score and threshold. + """ user_message = context.get("user_message") bot_message = context.get("bot_message") @@ -391,6 +467,7 @@ async def autoalign_factcheck_output_api( bot_message=bot_message, guardrails_config=guardrails_config, multi_language=multi_language, + http_client=http_client, ) if score < factcheck_threshold and show_autoalign_message: diff --git a/nemoguardrails/library/clavata/actions.py b/nemoguardrails/library/clavata/actions.py index 817bb029c5..9b3ed44a68 100644 --- a/nemoguardrails/library/clavata/actions.py +++ b/nemoguardrails/library/clavata/actions.py @@ -24,6 +24,7 @@ from nemoguardrails import RailsConfig from nemoguardrails.actions import action from nemoguardrails.actions.rail_outcome import RailOutcome +from nemoguardrails.http import HTTPClient from nemoguardrails.library.clavata.errs import ( ClavataPluginAPIError, ClavataPluginConfigurationError, @@ -205,10 +206,12 @@ async def evaluate_with_policy( text: str, policy_id: str, clavata_config: ClavataRailConfig, + http_client: HTTPClient | None = None, ) -> PolicyResult: """Get the policy result for the given source.""" client = ClavataClient( base_endpoint=get_server_endpoint(clavata_config), + http_client=http_client, ) job = await client.create_job(text, policy_id) @@ -227,6 +230,7 @@ async def clavata_check( labels: Optional[Union[List[str], str]] = None, rail: Union[ValidRailsType, None] = None, config: Optional[RailsConfig] = None, + http_client: HTTPClient | None = None, **kwargs: Any, ) -> RailOutcome: """Check for matches against a Clavata policy.""" @@ -249,7 +253,7 @@ async def clavata_check( except ClavataPluginValueError: labels = None - result = await evaluate_with_policy(text, str(policy_id), clavata_config) + result = await evaluate_with_policy(text, str(policy_id), clavata_config, http_client=http_client) if labels: return _clavata_outcome(is_label_match(result, labels, clavata_config)) diff --git a/nemoguardrails/library/clavata/errs.py b/nemoguardrails/library/clavata/errs.py index 593a1d616e..172e6e0bed 100644 --- a/nemoguardrails/library/clavata/errs.py +++ b/nemoguardrails/library/clavata/errs.py @@ -32,12 +32,6 @@ class ClavataPluginValueError(ClavataPluginError): """ -class ClavataPluginTypeError(ClavataPluginError): - """ - Exception raised when the Clavata plugin is used incorrectly due to type mismatches. - """ - - class ClavataPluginAPIError(ClavataPluginError): """ Exception raised when the Clavata plugin API returns an error. diff --git a/nemoguardrails/library/clavata/request.py b/nemoguardrails/library/clavata/request.py index 49df9815af..6164232726 100644 --- a/nemoguardrails/library/clavata/request.py +++ b/nemoguardrails/library/clavata/request.py @@ -20,15 +20,23 @@ from dataclasses import dataclass from typing import Dict, List, Literal, Optional, Type, TypeVar -import aiohttp from pydantic import BaseModel, Field, ValidationError -from nemoguardrails.library.clavata.utils import exponential_backoff +from nemoguardrails.http import ( + ClosableHTTPClient, + HTTPClient, + HTTPResponseDecodeError, + RetryingHTTPClient, + RetryPolicy, + create_http_client, + http_call, +) from .errs import ( ClavataPluginAPIError, ClavataPluginAPIRateLimitError, ClavataPluginConfigurationError, + ClavataPluginError, ClavataPluginValueError, ) @@ -36,6 +44,14 @@ _CLAVATA_API_KEY = os.environ.get("CLAVATA_API_KEY") +_CLAVATA_RETRY_POLICY = RetryPolicy( + max_attempts=3, + retryable_methods=frozenset({"POST"}), + retryable_status_codes=frozenset({429}), + initial_delay=0.1, + max_delay=10.0, + retry_transport_errors=False, +) @dataclass @@ -46,8 +62,7 @@ class AuthHeader: def to_headers(self) -> Dict[str, str]: """ - Converts the auth token into a dictionary that can be used with aiohttp - to supply headers for the request. + Converts the auth token into request headers. """ api_key = self.api_key or _CLAVATA_API_KEY if api_key is None: @@ -133,8 +148,14 @@ class ClavataClient: base_endpoint: str api_key: Optional[str] - def __init__(self, base_endpoint: str, api_key: Optional[str] = None): + def __init__( + self, + base_endpoint: str, + api_key: Optional[str] = None, + http_client: HTTPClient | None = None, + ): self.base_endpoint = base_endpoint + self.http_client = http_client # API key can be passed or set in the environment variable CLAVATA_API_KEY self.api_key = api_key or os.environ.get("CLAVATA_API_KEY") @@ -150,7 +171,14 @@ def _get_full_endpoint(self, endpoint: str) -> str: def _get_headers(self) -> Dict[str, str]: return AuthHeader(api_key=self.api_key).to_headers() - @exponential_backoff(initial_delay=0.1, retry_exceptions=(ClavataPluginAPIRateLimitError,)) + @staticmethod + def _retrying_http_client(client: HTTPClient) -> ClosableHTTPClient: + return RetryingHTTPClient(client, _CLAVATA_RETRY_POLICY) + + @classmethod + def _create_http_client(cls) -> ClosableHTTPClient: + return cls._retrying_http_client(create_http_client()) + async def _make_request( self, endpoint: str, @@ -158,37 +186,42 @@ async def _make_request( response_model: Type[ResponseModelT], ) -> ResponseModelT: try: - async with aiohttp.ClientSession() as session: - async with session.post( - self._get_full_endpoint(endpoint), - json=payload.model_dump(), - headers=self._get_headers(), - ) as resp: - if resp.status == 429: - # Trigger exponential backoff on rate limit errors - raise ClavataPluginAPIRateLimitError( - f"Clavata API rate limit exceeded. Status code: {resp.status}" - ) - - if resp.status != 200: - raise ClavataPluginAPIError( - f"Clavata call failed with status code {resp.status}.\nDetails: {await resp.text()}" - ) - - try: - parsed_response = await resp.json() - except aiohttp.ContentTypeError as e: - raise ClavataPluginValueError( - f"Failed to parse Clavata response as JSON. Status: {resp.status}, " - f"Content: {await resp.text()}" - ) from e - - # Now we actually parse the JSON into a meaningful object - try: - return response_model.model_validate(parsed_response) - except ValidationError as e: - raise ClavataPluginValueError(f"Invalid response format from Clavata API. Details: {e}") from e + client = self._retrying_http_client(self.http_client) if self.http_client is not None else None + response = await http_call( + client, + "POST", + self._get_full_endpoint(endpoint), + json=payload.model_dump(), + headers=self._get_headers(), + raise_for_status=False, + factory=self._create_http_client, + ) + if response.status_code == 429: + raise ClavataPluginAPIRateLimitError( + f"Clavata API rate limit exceeded. Status code: {response.status_code}" + ) + + if response.status_code != 200: + raise ClavataPluginAPIError( + f"Clavata call failed with status code {response.status_code}.\nDetails: {response.text}" + ) + + try: + parsed_response = response.json() + except HTTPResponseDecodeError as e: + raise ClavataPluginValueError( + f"Failed to parse Clavata response as JSON. Status: {response.status_code}, " + f"Content: {response.text}" + ) from e + + try: + return response_model.model_validate(parsed_response) + except ValidationError as e: + raise ClavataPluginValueError(f"Invalid response format from Clavata API. Details: {e}") from e + + except ClavataPluginError: + raise except Exception as e: raise ClavataPluginAPIError(f"Failed to make Clavata API request. Error: {e}") from e diff --git a/nemoguardrails/library/clavata/utils.py b/nemoguardrails/library/clavata/utils.py deleted file mode 100644 index eccebacdd3..0000000000 --- a/nemoguardrails/library/clavata/utils.py +++ /dev/null @@ -1,137 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2023-2026 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 asyncio -import random -from collections.abc import Awaitable, Callable, Iterable -from functools import wraps -from typing import Any, Optional, TypeVar, Union - -# Import ParamSpec from typing_extensions for Python 3.9 compatibility -from typing_extensions import ParamSpec - -from .errs import ClavataPluginTypeError - - -class AttemptsExceededError(Exception): - """Exception raised when the maximum number of retries is exceeded.""" - - attempts: int - max_attempts: int - last_exception: Optional[Exception] - - def __init__(self, attempts: int, max_attempts: int, last_exception: Optional[Exception]): - self.attempts = attempts - self.max_attempts = max_attempts - self.last_exception = last_exception - super().__init__( - f"Maximum number of attempts ({max_attempts}) exceeded after {attempts} attempts." - f"Last exception: {last_exception}" - ) - - -def calculate_exp_delay( - retries: int, - initial_delay: float, - max_delay: float, - jitter: bool, -) -> float: - """ - Handles calculation of the delay for a specific attempt. Note that we specifically - ask for the number of retries, not the number of attempts, because the first attempt - is the initial call and we want the first delay to be raised to the power of 0. - - Using "retries" instead of "attempts" makes it clearer what the input is, even - though a value called "attempts" is being passed in below. - - Because this is an exponential backoff, the factor of increase is always 2. - - Args: - retries: The number of retries made so far. - initial_delay: The initial delay. - max_delay: The maximum delay. - jitter: Whether to apply jitter to the delay. We use a full-jitter approach. - """ - delay = min( - initial_delay * (2**retries), - max_delay, - ) - if jitter: - delay = random.uniform(0, delay) - return delay - - -ReturnT = TypeVar("ReturnT") -P = ParamSpec("P") - - -def exponential_backoff( - *, - max_attempts: int = 3, - initial_delay: float = 1.0, - max_delay: float = 10.0, - jitter: bool = True, # Set to False to disable jitter - retry_exceptions: Union[type[Exception], Iterable[type[Exception]]] = Exception, - on_permanent_failure: Optional[Callable[[int, Exception], Awaitable[Any]]] = None, -): - """Exponential backoff retry mechanism.""" - - # Ensure retry_exceptions is a tuple of exceptions - retry_exceptions = (retry_exceptions,) if isinstance(retry_exceptions, type) else tuple(retry_exceptions) - - # Sanity check, make sure the types in the retry_exceptions are all exceptions - if not all(isinstance(e, type) and issubclass(e, Exception) for e in retry_exceptions): - raise ClavataPluginTypeError("retry_exceptions must be a tuple of exception types") - - def decorator( - func: Callable[P, Awaitable[ReturnT]], - ) -> Callable[P, Awaitable[ReturnT]]: - @wraps(func) - async def wrapper(*args: P.args, **kwargs: P.kwargs) -> ReturnT: - attempts = 0 - last_exception: Optional[Exception] = None - while attempts < max_attempts: - try: - return await func(*args, **kwargs) - except Exception as e: # noqa: BLE001 - # If the exception is not in the list of retry exceptions, raise it rather than retrying - last_exception = e - if not isinstance(e, retry_exceptions): - if on_permanent_failure is None: - raise - - perm_rv = await on_permanent_failure(attempts, e) - if isinstance(perm_rv, Exception): - raise perm_rv from e - return perm_rv - - # We want to calculate the delay before incrementing because we want the first - # delay to be exactly the initial delay - delay = calculate_exp_delay(attempts, initial_delay, max_delay, jitter) - await asyncio.sleep(delay) - attempts += 1 - - # Max retries exceeded, raise or if a custom handler is provided, call it and then decide what to do - retried_exc = AttemptsExceededError(attempts, max_attempts, last_exception) - if on_permanent_failure is None: - raise retried_exc - perm_rv = await on_permanent_failure(attempts, retried_exc) - if isinstance(perm_rv, Exception): - raise perm_rv - return perm_rv - - return wrapper - - return decorator diff --git a/nemoguardrails/library/crowdstrike_aidr/actions.py b/nemoguardrails/library/crowdstrike_aidr/actions.py index d90a6a0c74..1a3413c0a9 100644 --- a/nemoguardrails/library/crowdstrike_aidr/actions.py +++ b/nemoguardrails/library/crowdstrike_aidr/actions.py @@ -18,13 +18,13 @@ from collections.abc import Mapping from typing import Any, Optional, cast -import httpx -from pydantic import BaseModel +from pydantic import BaseModel, ValidationError from pydantic_core import to_json from typing_extensions import Literal, TypedDict from nemoguardrails.actions import action from nemoguardrails.actions.rail_outcome import RailOutcome, TransformTarget +from nemoguardrails.http import HTTPClient, HTTPClientError, HTTPStatusError, http_call from nemoguardrails.rails.llm.config import CrowdStrikeAIDRRailConfig, RailsConfig log = logging.getLogger(__name__) @@ -69,6 +69,12 @@ def _crowdstrike_aidr_outcome( result: GuardChatCompletionsResult, mode: Literal["input", "output"], ) -> RailOutcome: + """Convert a CrowdStrike result into a rail decision. + + Blocking takes precedence over transformation. A transformed input rewrites + the user message, while a transformed output rewrites the bot message. + """ + metadata = { "blocked": bool(result.blocked), "transformed": bool(result.transformed), @@ -95,7 +101,37 @@ async def crowdstrike_aidr_guard( context: Mapping[str, Any] = {}, user_message: Optional[str] = None, bot_message: Optional[str] = None, + http_client: Optional[HTTPClient] = None, ) -> RailOutcome: + """Evaluate conversation content with CrowdStrike AI Guard. + + The action evaluates the user message for input rails and includes the bot + message for output rails. A successful provider response can allow, block, + or transform the active message. + + Expected provider failures fail open: unsuccessful HTTP status responses, + transport failures, malformed JSON, and response-schema validation errors + are logged and return an allow outcome containing the original messages. + Unexpected programming errors are not swallowed. + + Args: + mode: Whether the action is evaluating an input or output rail. + config: Active rails configuration. + context: Conversation context used when explicit messages are omitted. + user_message: Optional user message override. + bot_message: Optional bot message override. + http_client: Optional caller-owned HTTP client. The action does not + close an injected client. + + Returns: + The provider-derived rail outcome, or an allow outcome for an expected + provider failure. + + Raises: + ValueError: If the API token or required conversation content is + missing. + """ + base_url_template = os.getenv("CS_AIDR_BASE_URL_TEMPLATE", "https://api.crowdstrike.com/aidr/{SERVICE_NAME}") api_token = os.getenv("CS_AIDR_TOKEN") @@ -118,9 +154,19 @@ async def crowdstrike_aidr_guard( if mode == "output" and bot_message: messages.append(Message(role="assistant", content=bot_message)) - async with httpx.AsyncClient(base_url=base_url_template.format(SERVICE_NAME="aiguard")) as client: - response = await client.post( - "/v1/guard_chat_completions", + endpoint = base_url_template.format(SERVICE_NAME="aiguard").rstrip("/") + "/v1/guard_chat_completions" + fallback = GuardChatCompletionsResult( + guard_output={"messages": messages}, + blocked=False, + transformed=False, + bot_message=bot_message, + user_message=user_message, + ) + try: + response = await http_call( + http_client, + "POST", + endpoint, content=to_json({"guard_input": {"messages": messages}}), headers={ "Accept": "application/json", @@ -129,35 +175,15 @@ async def crowdstrike_aidr_guard( "User-Agent": "NeMo Guardrails (https://github.com/NVIDIA-NeMo/Guardrails)", }, timeout=crowdstrike_aidr_config.timeout, + raise_for_status=False, ) - try: - response.raise_for_status() - guard_response = GuardChatCompletionsResponse(**response.json()) - except httpx.HTTPStatusError as e: - log.error("HTTP status error from CrowdStrike AIDR API: %s", e) - return _crowdstrike_aidr_outcome( - GuardChatCompletionsResult( - guard_output={"messages": messages}, - blocked=False, - transformed=False, - bot_message=bot_message, - user_message=user_message, - ), - mode, - ) - except Exception as e: - log.error("Error calling CrowdStrike AIDR API: %s", e) - return _crowdstrike_aidr_outcome( - GuardChatCompletionsResult( - guard_output={"messages": messages}, - blocked=False, - transformed=False, - bot_message=bot_message, - user_message=user_message, - ), - mode, - ) - + response.raise_for_status() + guard_response = GuardChatCompletionsResponse.model_validate(response.json()) + except HTTPStatusError as e: + log.error("HTTP status error from CrowdStrike AIDR API: %s", e) + except (HTTPClientError, ValidationError) as e: + log.error("Error calling CrowdStrike AIDR API: %s", e) + else: result = guard_response.result output_messages = result.guard_output.get("messages", []) if result.guard_output else [] @@ -165,3 +191,5 @@ async def crowdstrike_aidr_guard( result.user_message = next((m.content for m in output_messages if m.role == "user"), user_message) return _crowdstrike_aidr_outcome(result, mode) + + return _crowdstrike_aidr_outcome(fallback, mode) diff --git a/nemoguardrails/library/f5/actions.py b/nemoguardrails/library/f5/actions.py index 072bb92a75..a0eb0b36fb 100644 --- a/nemoguardrails/library/f5/actions.py +++ b/nemoguardrails/library/f5/actions.py @@ -16,15 +16,21 @@ import asyncio import logging import os -from datetime import datetime, timezone -from email.utils import parsedate_to_datetime from typing import Any, Optional -import aiohttp from typing_extensions import cast from nemoguardrails.actions import action from nemoguardrails.actions.rail_outcome import RailOutcome +from nemoguardrails.http import ( + HTTPClient, + HTTPConnectionError, + HTTPTimeoutError, + RetryingHTTPClient, + RetryPolicy, + create_http_client, + http_call, +) from nemoguardrails.rails.llm.config import F5GuardrailsRailConfig, RailsConfig log = logging.getLogger(__name__) @@ -55,51 +61,39 @@ def _fail_open_outcome() -> RailOutcome: return RailOutcome.allow(metadata={"result": {"outcome": "cleared"}, "fail_open": True}) -def _parse_retry_after(raw: Optional[str]) -> Optional[float]: - """Parse a Retry-After header value. +def _retry_policy(f5_config: F5GuardrailsRailConfig) -> RetryPolicy: + max_delay = f5_config.max_retry_after_seconds + return RetryPolicy( + max_attempts=f5_config.max_retries + 1, + retryable_methods=frozenset({"POST"}), + retryable_status_codes=frozenset({429}), + initial_delay=min(f5_config.retry_backoff_seconds, max_delay), + max_delay=max_delay, + max_retry_after=max_delay, + retry_transport_errors=False, + honor_retry_override_header=False, + clamp_retry_after=True, + ) - Supports both delta-seconds (numeric) and HTTP-date forms per RFC 7231. - Returns ``None`` if the value is missing or unparseable. - """ - if not raw: - return None - try: - return float(raw) - except (TypeError, ValueError): - pass - try: - parsed = parsedate_to_datetime(str(raw)) - except (TypeError, ValueError): - return None - if parsed is None: - return None - if parsed.tzinfo is None: - parsed = parsed.replace(tzinfo=timezone.utc) - return (parsed - datetime.now(tz=timezone.utc)).total_seconds() - - -def _compute_retry_delay( - retry_after_raw: Optional[str], - attempt: int, - f5_config: F5GuardrailsRailConfig, -) -> float: - """Compute the sleep duration before the next 429 retry.""" - parsed = _parse_retry_after(retry_after_raw) - if parsed is None: - delay = f5_config.retry_backoff_seconds * (2**attempt) - else: - delay = parsed - if delay < 0: - delay = 0.0 - if delay > f5_config.max_retry_after_seconds: - delay = f5_config.max_retry_after_seconds - return delay + +def _retrying_http_client(client: HTTPClient, f5_config: F5GuardrailsRailConfig) -> HTTPClient: + return RetryingHTTPClient( + client, + _retry_policy(f5_config), + sleep=asyncio.sleep, + random_value=lambda: 1.0, + ) + + +def _create_http_client(f5_config: F5GuardrailsRailConfig) -> HTTPClient: + return _retrying_http_client(create_http_client(timeout=30.0), f5_config) @action(name="f5_guardrails_scan", is_system_action=True) async def f5_guardrails_scan( text: str, config: Optional[RailsConfig] = None, + http_client: Optional[HTTPClient] = None, **kwargs: Any, ) -> RailOutcome: """ @@ -108,6 +102,8 @@ async def f5_guardrails_scan( Args: text: The text to scan. config: The active RailsConfig; used to resolve ``rails.config.f5``. + http_client: Optional caller-owned HTTP client used with the F5 retry + policy. The action does not close an injected client. Returns: The rail decision with the F5 Guardrails API response as metadata. @@ -141,63 +137,53 @@ async def f5_guardrails_scan( "input": text, } - timeout = aiohttp.ClientTimeout(total=30) - max_attempts = f5_config.max_retries + 1 - async with aiohttp.ClientSession(timeout=timeout) as session: - for attempt in range(max_attempts): - try: - async with session.post(endpoint, headers=headers, json=payload) as response: - if response.status == 429 and attempt < max_attempts - 1: - retry_after_raw = response.headers.get("Retry-After") - delay = _compute_retry_delay(retry_after_raw, attempt, f5_config) - log.warning( - "F5 Guardrails API rate limited: status=429 attempt=%s/%s sleep_seconds=%.3f retry_after_present=%s", - attempt + 1, - max_attempts, - delay, - retry_after_raw is not None, - ) - await asyncio.sleep(delay) - continue - - if response.status != 200: - content_type = response.headers.get("Content-Type", "unknown") - body_length = response.content_length if response.content_length is not None else "unknown" - log.error( - "F5 Guardrails API call failed: status=%s content_type=%s body_length=%s", - response.status, - content_type, - body_length, - ) - - if fail_open: - log.warning("F5 Guardrails API call failed; fail_open is enabled, allowing content.") - return _fail_open_outcome() - - if response.status == 429: - raise RuntimeError("F5 Guardrails API rate limited (429) after exhausting retries") - raise RuntimeError(f"F5 Guardrails API error: {response.status}") - - result = await response.json() - return _scan_outcome(result) - except asyncio.TimeoutError: - log.error("F5 Guardrails API call timed out after 30 seconds") - - if fail_open: - log.warning("F5 Guardrails API call timed out; fail_open is enabled, allowing content.") - return _fail_open_outcome() - - raise RuntimeError("F5 Guardrails API request timed out") from None - except aiohttp.ClientError as e: - log.error("Error connecting to F5 Guardrails API: %s", type(e).__name__) - - if fail_open: - log.warning("F5 Guardrails API call failed; fail_open is enabled, allowing content.") - return _fail_open_outcome() - - raise RuntimeError(f"Connection error to F5 Guardrails API: {type(e).__name__}") from e - - if fail_open: - log.warning("F5 Guardrails API rate limited after retries; fail_open is enabled, allowing content.") - return _fail_open_outcome() - raise RuntimeError("F5 Guardrails API rate limited (429) after exhausting retries") + try: + client = _retrying_http_client(http_client, f5_config) if http_client is not None else None + response = await http_call( + client, + "POST", + endpoint, + headers=headers, + json=payload, + timeout=30.0, + raise_for_status=False, + factory=lambda: _create_http_client(f5_config), + ) + except (HTTPTimeoutError, asyncio.TimeoutError): + log.error("F5 Guardrails API call timed out after 30 seconds") + + if fail_open: + log.warning("F5 Guardrails API call timed out; fail_open is enabled, allowing content.") + return _fail_open_outcome() + + raise RuntimeError("F5 Guardrails API request timed out") from None + except HTTPConnectionError as e: + log.error("Error connecting to F5 Guardrails API: %s", type(e).__name__) + + if fail_open: + log.warning("F5 Guardrails API call failed; fail_open is enabled, allowing content.") + return _fail_open_outcome() + + raise RuntimeError(f"Connection error to F5 Guardrails API: {type(e).__name__}") from e + + if response.status_code != 200: + content_type = next( + (value for name, value in response.headers.items() if name.lower() == "content-type"), + "unknown", + ) + log.error( + "F5 Guardrails API call failed: status=%s content_type=%s body_length=%s", + response.status_code, + content_type, + len(response.content), + ) + + if fail_open: + log.warning("F5 Guardrails API call failed; fail_open is enabled, allowing content.") + return _fail_open_outcome() + + if response.status_code == 429: + raise RuntimeError("F5 Guardrails API rate limited (429) after exhausting retries") + raise RuntimeError(f"F5 Guardrails API error: {response.status_code}") + + return _scan_outcome(response.json()) diff --git a/nemoguardrails/library/fiddler/actions.py b/nemoguardrails/library/fiddler/actions.py index edbb0e7852..db12bd6b1f 100644 --- a/nemoguardrails/library/fiddler/actions.py +++ b/nemoguardrails/library/fiddler/actions.py @@ -17,11 +17,10 @@ import os from typing import Callable, Optional -import aiohttp - from nemoguardrails import RailsConfig from nemoguardrails.actions import action from nemoguardrails.actions.rail_outcome import RailOutcome +from nemoguardrails.http import HTTPClient, HTTPClientError, http_call from nemoguardrails.rails.llm.config import FiddlerGuardrails log = logging.getLogger(__name__) @@ -41,7 +40,29 @@ async def call_fiddler_guardrail( threshold: float, compare: Callable[[float, float], bool], default_score: float, + http_client: Optional[HTTPClient] = None, ) -> bool: + """Evaluate content with a Fiddler guardrail endpoint. + + Non-success responses, HTTP client failures, and expected score-processing + failures are logged and fail open. + + Args: + endpoint: Fiddler guardrail endpoint. + data: Guardrail-specific request data. + guardrail_name: Name used in diagnostic messages. + score_key: Response score used for the decision. + threshold: Configured decision threshold. + compare: Comparison used to determine whether content is blocked. + default_score: Score used when the response omits ``score_key``. + http_client: Optional caller-owned HTTP client. + + Returns: + ``True`` when the guardrail blocks the content, otherwise ``False``. + + Raises: + ValueError: If ``FIDDLER_API_KEY`` is not set. + """ api_key = os.environ.get("FIDDLER_API_KEY") if api_key is None: @@ -53,34 +74,40 @@ async def call_fiddler_guardrail( } try: - async with aiohttp.ClientSession() as session: - async with session.post(endpoint, headers=headers, json={"data": data}) as response: - if response.status != 200: - log.error(f"{guardrail_name} could not be run. Fiddler API returned status code {response.status}") - return False - - response_json = await response.json() - if score_key == "safety": - detection_score = max( - response_json.get(key, default_score) - for key in [ - "fdl_harmful", - "fdl_violent", - "fdl_unethical", - "fdl_illegal", - "fdl_sexual", - "fdl_racist", - "fdl_jailbreaking", - "fdl_harassing", - "fdl_hateful", - "fdl_sexist", - "fdl_roleplaying", - ] - ) - else: - detection_score = response_json.get(score_key, default_score) - return compare(detection_score, threshold) - except aiohttp.ClientError as e: + response = await http_call( + http_client, + "POST", + endpoint, + headers=headers, + json={"data": data}, + raise_for_status=False, + ) + if response.status_code != 200: + log.error(f"{guardrail_name} could not be run. Fiddler API returned status code {response.status_code}") + return False + + response_json = response.json() + if score_key == "safety": + detection_score = max( + response_json.get(key, default_score) + for key in [ + "fdl_harmful", + "fdl_violent", + "fdl_unethical", + "fdl_illegal", + "fdl_sexual", + "fdl_racist", + "fdl_jailbreaking", + "fdl_harassing", + "fdl_hateful", + "fdl_sexist", + "fdl_roleplaying", + ] + ) + else: + detection_score = response_json.get(score_key, default_score) + return compare(detection_score, threshold) + except HTTPClientError as e: log.error(f"{guardrail_name} request failed: {e}") return False except (KeyError, ValueError, IndexError) as e: @@ -89,7 +116,27 @@ async def call_fiddler_guardrail( @action(name="call_fiddler_safety_user", is_system_action=True) -async def call_fiddler_safety_user(config: RailsConfig, context: Optional[dict] = None) -> RailOutcome: +async def call_fiddler_safety_user( + config: RailsConfig, + context: Optional[dict] = None, + http_client: Optional[HTTPClient] = None, +) -> RailOutcome: + """Check the user message with the Fiddler safety guardrail. + + Scores at or above the configured threshold block. Missing content, + missing endpoint configuration, and handled provider failures fail open. + + Args: + config: Rails configuration containing the Fiddler settings. + context: Runtime context containing the user message. + http_client: Optional caller-owned HTTP client. + + Returns: + An allow or block outcome. + + Raises: + ValueError: If ``FIDDLER_API_KEY`` is not set. + """ context = context or {} fiddler_config: FiddlerGuardrails = getattr(config.rails.config, "fiddler") base_url = fiddler_config.fiddler_endpoint @@ -112,12 +159,33 @@ async def call_fiddler_safety_user(config: RailsConfig, context: Optional[dict] threshold=fiddler_config.safety_threshold, compare=lambda score, threshold: score >= threshold, default_score=0, + http_client=http_client, ) return _fiddler_outcome(blocked) @action(name="call_fiddler_safety_bot", is_system_action=True) -async def call_fiddler_safety_bot(config: RailsConfig, context: Optional[dict] = None) -> RailOutcome: +async def call_fiddler_safety_bot( + config: RailsConfig, + context: Optional[dict] = None, + http_client: Optional[HTTPClient] = None, +) -> RailOutcome: + """Check the bot message with the Fiddler safety guardrail. + + Scores at or above the configured threshold block. Missing content, + missing endpoint configuration, and handled provider failures fail open. + + Args: + config: Rails configuration containing the Fiddler settings. + context: Runtime context containing the bot message. + http_client: Optional caller-owned HTTP client. + + Returns: + An allow or block outcome. + + Raises: + ValueError: If ``FIDDLER_API_KEY`` is not set. + """ context = context or {} fiddler_config: FiddlerGuardrails = getattr(config.rails.config, "fiddler") base_url = fiddler_config.fiddler_endpoint @@ -140,12 +208,34 @@ async def call_fiddler_safety_bot(config: RailsConfig, context: Optional[dict] = threshold=fiddler_config.safety_threshold, compare=lambda score, threshold: score >= threshold, default_score=0, + http_client=http_client, ) return _fiddler_outcome(blocked) @action(name="call_fiddler_faithfulness", is_system_action=True) -async def call_fiddler_faithfulness(config: RailsConfig, context: Optional[dict] = None) -> RailOutcome: +async def call_fiddler_faithfulness( + config: RailsConfig, + context: Optional[dict] = None, + http_client: Optional[HTTPClient] = None, +) -> RailOutcome: + """Check the bot message with the Fiddler faithfulness guardrail. + + Scores at or below the configured threshold block. Missing content, + missing endpoint configuration, and handled provider failures fail open. + + Args: + config: Rails configuration containing the Fiddler settings. + context: Runtime context containing the bot message and relevant + chunks. + http_client: Optional caller-owned HTTP client. + + Returns: + An allow or block outcome. + + Raises: + ValueError: If ``FIDDLER_API_KEY`` is not set. + """ context = context or {} fiddler_config: FiddlerGuardrails = getattr(config.rails.config, "fiddler") base_url = fiddler_config.fiddler_endpoint @@ -169,5 +259,6 @@ async def call_fiddler_faithfulness(config: RailsConfig, context: Optional[dict] threshold=fiddler_config.faithfulness_threshold, compare=lambda score, threshold: score <= threshold, default_score=1, + http_client=http_client, ) return _fiddler_outcome(blocked) diff --git a/nemoguardrails/library/hf_classifier/actions.py b/nemoguardrails/library/hf_classifier/actions.py index e247dddeb7..3d9e313abe 100644 --- a/nemoguardrails/library/hf_classifier/actions.py +++ b/nemoguardrails/library/hf_classifier/actions.py @@ -20,6 +20,7 @@ from nemoguardrails.actions import action from nemoguardrails.actions.rail_outcome import RailOutcome, TransformTarget +from nemoguardrails.http import HTTPClient from nemoguardrails.library.hf_classifier.backends import get_backend log = logging.getLogger(__name__) @@ -29,6 +30,7 @@ async def _classify_and_check( classifier_name: str, text: str, config: Any | None, + http_client: HTTPClient | None = None, ) -> bool: """Classify *text* and check against blocked labels. @@ -48,7 +50,7 @@ async def _classify_and_check( if not text: return True - backend = get_backend(classifier_config, name=classifier_name) + backend = get_backend(classifier_config, name=classifier_name, http_client=http_client) results = await backend.classify(text) if text and not results and getattr(classifier_config, "task", None) == "text-classification": @@ -99,9 +101,10 @@ async def hf_classifier_check_input( classifier: str, config: Any | None = None, context: Optional[dict] = None, + http_client: HTTPClient | None = None, **kwargs, ) -> RailOutcome: - allowed = await _classify_and_check(classifier, _extract_text(context, "user_message"), config) + allowed = await _classify_and_check(classifier, _extract_text(context, "user_message"), config, http_client) return _hf_classifier_outcome(allowed) @@ -111,6 +114,7 @@ async def hf_classifier_check_output( config: Any | None = None, context: Optional[dict] = None, model_name: Optional[str] = None, + http_client: HTTPClient | None = None, **kwargs, ) -> RailOutcome: # Streaming output rail path doesn't resolve flow variables — $classifier @@ -118,7 +122,7 @@ async def hf_classifier_check_output( # engine extracts from the flow_id. if classifier.startswith("$") and model_name: classifier = model_name - allowed = await _classify_and_check(classifier, _extract_text(context, "bot_message"), config) + allowed = await _classify_and_check(classifier, _extract_text(context, "bot_message"), config, http_client) return _hf_classifier_outcome(allowed) @@ -127,7 +131,8 @@ async def hf_classifier_check_retrieval( classifier: str, config: Any | None = None, context: Optional[dict] = None, + http_client: HTTPClient | None = None, **kwargs, ) -> RailOutcome: - allowed = await _classify_and_check(classifier, _extract_text(context, "relevant_chunks"), config) + allowed = await _classify_and_check(classifier, _extract_text(context, "relevant_chunks"), config, http_client) return _hf_classifier_retrieval_outcome(allowed) diff --git a/nemoguardrails/library/hf_classifier/backends.py b/nemoguardrails/library/hf_classifier/backends.py index c3be644be1..5d326ccc53 100644 --- a/nemoguardrails/library/hf_classifier/backends.py +++ b/nemoguardrails/library/hf_classifier/backends.py @@ -23,9 +23,17 @@ import logging import os import threading -from typing import TYPE_CHECKING, Any, Dict, List, Optional, TypedDict, Union - -import httpx +from typing import TYPE_CHECKING, Any, Dict, List, Optional, TypedDict + +from nemoguardrails.http import ( + HTTPClient, + HTTPResponse, + HTTPTLSConfig, + RetryingHTTPClient, + RetryPolicy, + create_http_client, + http_call, +) if TYPE_CHECKING: from nemoguardrails.rails.llm.config import ( @@ -35,6 +43,13 @@ ) log = logging.getLogger(__name__) +_HF_RETRY_POLICY = RetryPolicy( + max_attempts=2, + retryable_methods=frozenset({"POST"}), + retryable_status_codes=frozenset(), + initial_delay=0, + max_delay=0, +) class ClassificationResult(TypedDict): @@ -76,28 +91,15 @@ def _build_headers(config: RemoteHFClassifierConfig) -> Dict[str, str]: return headers -def _get_timeout(config: RemoteHFClassifierConfig) -> httpx.Timeout: - total = config.parameters.get("timeout", _DEFAULT_TIMEOUT) - return httpx.Timeout(total) - - -class _HTTPXSSLConfig: - """Stores httpx-compatible SSL parameters (verify, cert).""" +def _get_timeout(config: RemoteHFClassifierConfig) -> float: + return float(config.parameters.get("timeout", _DEFAULT_TIMEOUT)) - def __init__( - self, - verify: Union[str, bool] = True, - cert: Optional[tuple] = None, - ) -> None: - self.verify = verify - self.cert = cert +_ssl_cache: Dict[tuple, HTTPTLSConfig] = {} -_ssl_cache: Dict[tuple, _HTTPXSSLConfig] = {} - -def _build_ssl_config(config: RemoteHFClassifierConfig) -> _HTTPXSSLConfig: - """Build httpx SSL params from config parameters (cached). +def _build_ssl_config(config: RemoteHFClassifierConfig) -> HTTPTLSConfig: + """Build HTTP transport SSL params from config parameters (cached). Reads from ``config.parameters``: - ``verify_ssl`` (bool, default True): set to False to skip TLS verification. @@ -118,13 +120,18 @@ def _build_ssl_config(config: RemoteHFClassifierConfig) -> _HTTPXSSLConfig: provided, missing = ("client_cert", "client_key") if client_cert else ("client_key", "client_cert") raise ValueError(f"mTLS requires both 'client_cert' and 'client_key'; got '{provided}' without '{missing}'.") - cert = (client_cert, client_key) if client_cert else None if verify is False: - result = _HTTPXSSLConfig(verify=False, cert=cert) - elif ca_cert: - result = _HTTPXSSLConfig(verify=ca_cert, cert=cert) + result = HTTPTLSConfig( + verify=False, + client_certificate=client_cert, + client_key=client_key, + ) else: - result = _HTTPXSSLConfig(verify=True, cert=cert) + result = HTTPTLSConfig( + ca_bundle=ca_cert, + client_certificate=client_cert, + client_key=client_key, + ) _ssl_cache[cache_key] = result return result @@ -295,37 +302,39 @@ def _run(): return results -_TRANSIENT_ERRORS = (httpx.NetworkError, httpx.TimeoutException, httpx.RemoteProtocolError, OSError, ValueError) - - class _RemoteBackend(ClassifierBackend): - """Base class for remote HTTP classifier backends. - - Uses a fresh httpx client per request — no connection pool state to manage. - Classifier calls are infrequent and latency-tolerant (inference >> TLS - handshake), so connection reuse provides negligible benefit while introducing - stale-connection complexity. - """ + """Base class for remote HTTP classifier backends.""" - def __init__(self, config: RemoteHFClassifierConfig) -> None: + def __init__(self, config: RemoteHFClassifierConfig, http_client: HTTPClient | None = None) -> None: self._config = config + self._http_client = http_client self._timeout = _get_timeout(config) self._ssl = _build_ssl_config(config) - async def _post(self, url: str, json: Dict[str, Any]) -> httpx.Response: - """POST with one automatic retry on transient failures.""" - headers = _build_headers(self._config) - for attempt in range(2): - try: - async with httpx.AsyncClient( - timeout=self._timeout, verify=self._ssl.verify, cert=self._ssl.cert - ) as client: - return await client.post(url, json=json, headers=headers) - except _TRANSIENT_ERRORS: - if attempt > 0: - raise - log.debug("Classifier request to %s failed, retrying.", url) - raise RuntimeError("unreachable") + @staticmethod + def _retrying_http_client(client: HTTPClient) -> HTTPClient: + return RetryingHTTPClient(client, _HF_RETRY_POLICY) + + def _create_http_client(self) -> HTTPClient: + return self._retrying_http_client( + create_http_client( + timeout=self._timeout, + tls=self._ssl, + ) + ) + + async def _post(self, url: str, json: Dict[str, Any]) -> HTTPResponse: + client = self._retrying_http_client(self._http_client) if self._http_client is not None else None + return await http_call( + client, + "POST", + url, + json=json, + headers=_build_headers(self._config), + timeout=self._timeout, + raise_for_status=False, + factory=self._create_http_client, + ) class VLLMBackend(_RemoteBackend): @@ -339,8 +348,8 @@ class VLLMBackend(_RemoteBackend): ``label``), indicating an API change. """ - def __init__(self, config: RemoteHFClassifierConfig) -> None: - super().__init__(config) + def __init__(self, config: RemoteHFClassifierConfig, http_client: HTTPClient | None = None) -> None: + super().__init__(config, http_client) self._url = config.base_url + "/classify" self._model_name = config.model @@ -389,8 +398,8 @@ class KServeBackend(_RemoteBackend): has an unrecognised type. """ - def __init__(self, config: RemoteHFClassifierConfig) -> None: - super().__init__(config) + def __init__(self, config: RemoteHFClassifierConfig, http_client: HTTPClient | None = None) -> None: + super().__init__(config, http_client) self._url = f"{config.base_url}/v1/models/{config.model}:predict" async def classify(self, text: str) -> List[ClassificationResult]: @@ -445,8 +454,8 @@ class FMSBackend(_RemoteBackend): entries lack required keys. """ - def __init__(self, config: RemoteHFClassifierConfig) -> None: - super().__init__(config) + def __init__(self, config: RemoteHFClassifierConfig, http_client: HTTPClient | None = None) -> None: + super().__init__(config, http_client) self._url = config.base_url + "/api/v1/text/contents" self._threshold = config.threshold @@ -489,14 +498,22 @@ async def classify(self, text: str) -> List[ClassificationResult]: _backend_instances: Dict[str, ClassifierBackend] = {} -def get_backend(config: HFClassifierConfig, name: str = "") -> ClassifierBackend: +def get_backend( + config: HFClassifierConfig, + name: str = "", + http_client: HTTPClient | None = None, +) -> ClassifierBackend: """Get or create a cached backend instance from classifier config.""" + cls = _BACKENDS.get(config.engine) + if cls is None: + raise ValueError(f"Unknown hf_classifier engine: '{config.engine}'. Supported: {', '.join(_BACKENDS)}") + if http_client is not None and config.engine != "local": + return cls(config, http_client=http_client) + if http_client is not None: + log.debug("Ignoring injected HTTP client for local HF classifier backend.") cache_key = json.dumps({"name": name, "config": config.model_dump(mode="json")}, sort_keys=True) cached = _backend_instances.get(cache_key) if cached is not None: return cached - cls = _BACKENDS.get(config.engine) - if cls is None: - raise ValueError(f"Unknown hf_classifier engine: '{config.engine}'. Supported: {', '.join(_BACKENDS)}") _backend_instances[cache_key] = cls(config) return _backend_instances[cache_key] diff --git a/nemoguardrails/library/pangea/actions.py b/nemoguardrails/library/pangea/actions.py index 17327b2cfe..fc69fec847 100644 --- a/nemoguardrails/library/pangea/actions.py +++ b/nemoguardrails/library/pangea/actions.py @@ -18,13 +18,13 @@ from collections.abc import Mapping from typing import Any, Optional -import httpx from pydantic import BaseModel from pydantic_core import to_json from typing_extensions import Literal, cast from nemoguardrails.actions import action from nemoguardrails.actions.rail_outcome import RailOutcome, TransformTarget +from nemoguardrails.http import HTTPClient, http_call from nemoguardrails.rails.llm.config import PangeaRailConfig, RailsConfig log = logging.getLogger(__name__) @@ -88,7 +88,28 @@ async def pangea_ai_guard( context: Mapping[str, Any] = {}, user_message: Optional[str] = None, bot_message: Optional[str] = None, + http_client: Optional[HTTPClient] = None, ) -> RailOutcome: + """Evaluate input or output content with Pangea AI Guard. + + Provider results may allow, block, or transform the selected message. + Failures while calling or validating the provider response are logged and + fail open with the original content. + + Args: + mode: Whether to evaluate input or output content. + config: Rails configuration containing the Pangea settings. + context: Runtime context used to resolve messages. + user_message: Optional user message overriding the context value. + bot_message: Optional bot message overriding the context value. + http_client: Optional caller-owned HTTP client. + + Returns: + The provider decision, or an allow outcome after a provider failure. + + Raises: + ValueError: If the API token or content is missing. + """ pangea_base_url_template = os.getenv("PANGEA_BASE_URL_TEMPLATE", "https://{SERVICE_NAME}.aws.us.pangea.cloud") pangea_api_token = os.getenv("PANGEA_API_TOKEN") @@ -117,13 +138,14 @@ async def pangea_ai_guard( else (pangea_config.output.recipe if mode == "output" and pangea_config.output else None) ) - async with httpx.AsyncClient(base_url=pangea_base_url_template.format(SERVICE_NAME="ai-guard")) as client: - data = {"messages": messages, "recipe": recipe} - # Remove `None` values. - data = {k: v for k, v in data.items() if v is not None} - - response = await client.post( - "/v1/text/guard", + data = {"messages": messages, "recipe": recipe} + data = {k: v for k, v in data.items() if v is not None} + endpoint = pangea_base_url_template.format(SERVICE_NAME="ai-guard").rstrip("/") + "/v1/text/guard" + try: + response = await http_call( + http_client, + "POST", + endpoint, content=to_json(data), headers={ "Accept": "application/json", @@ -131,27 +153,27 @@ async def pangea_ai_guard( "Content-Type": "application/json", "User-Agent": "NeMo Guardrails (https://github.com/NVIDIA-NeMo/Guardrails)", }, + raise_for_status=False, + ) + response.raise_for_status() + text_guard_response = TextGuardResponse(**response.json()) + except Exception as e: + log.error("Error calling Pangea AI Guard API: %s", e) + return _pangea_outcome( + TextGuardResult( + prompt_messages=messages, + blocked=False, + transformed=False, + bot_message=bot_message, + user_message=user_message, + ), + mode, ) - try: - response.raise_for_status() - text_guard_response = TextGuardResponse(**response.json()) - except Exception as e: - log.error("Error calling Pangea AI Guard API: %s", e) - return _pangea_outcome( - TextGuardResult( - prompt_messages=messages, - blocked=False, - transformed=False, - bot_message=bot_message, - user_message=user_message, - ), - mode, - ) - - result = text_guard_response.result - prompt_messages = result.prompt_messages or [] - - result.bot_message = next((m.content for m in prompt_messages if m.role == "assistant"), bot_message) - result.user_message = next((m.content for m in prompt_messages if m.role == "user"), user_message) - - return _pangea_outcome(result, mode) + + result = text_guard_response.result + prompt_messages = result.prompt_messages or [] + + result.bot_message = next((m.content for m in prompt_messages if m.role == "assistant"), bot_message) + result.user_message = next((m.content for m in prompt_messages if m.role == "user"), user_message) + + return _pangea_outcome(result, mode) diff --git a/nemoguardrails/library/patronusai/actions.py b/nemoguardrails/library/patronusai/actions.py index e5a3573b0a..e679005373 100644 --- a/nemoguardrails/library/patronusai/actions.py +++ b/nemoguardrails/library/patronusai/actions.py @@ -19,12 +19,11 @@ from collections.abc import Mapping from typing import List, Literal, Optional, Tuple, Union -import aiohttp - from nemoguardrails.actions import action from nemoguardrails.actions.llm.utils import llm_call from nemoguardrails.actions.rail_outcome import RailOutcome from nemoguardrails.context import llm_call_info_var +from nemoguardrails.http import HTTPClient, http_call from nemoguardrails.llm.taskmanager import LLMTaskManager from nemoguardrails.llm.types import Task from nemoguardrails.logging.explain import LLMCallInfo @@ -157,6 +156,7 @@ async def patronus_evaluate_request( user_input: Optional[str] = None, bot_response: Optional[str] = None, provided_context: Optional[Union[str, List[str]]] = None, + http_client: Optional[HTTPClient] = None, ) -> Optional[dict]: """ Make a call to the Patronus Evaluate API. @@ -196,34 +196,35 @@ async def patronus_evaluate_request( "Content-Type": "application/json", } - async with aiohttp.ClientSession() as session: - async with session.post( - url=url, - headers=headers, - json=data, - ) as response: - if 400 <= response.status < 500: - raise ValueError( - f"The Patronus Evaluate API call failed with status code {response.status}. " - f"Details: {await response.text()}" - ) - - if response.status != 200: - log.error( - "The Patronus Evaluate API call failed with status code %s. Details: %s", - response.status, - await response.text(), - ) - return None - - response_json = await response.json() - return response_json + response = await http_call( + http_client, + "POST", + url, + headers=headers, + json=data, + raise_for_status=False, + ) + if 400 <= response.status_code < 500: + raise ValueError( + f"The Patronus Evaluate API call failed with status code {response.status_code}. Details: {response.text}" + ) + + if response.status_code != 200: + log.error( + "The Patronus Evaluate API call failed with status code %s. Details: %s", + response.status_code, + response.text, + ) + return None + + return response.json() @action(name="patronus_api_check_output") async def patronus_api_check_output( llm_task_manager: LLMTaskManager, context: Optional[dict] = None, + http_client: Optional[HTTPClient] = None, ) -> RailOutcome: """ Check the user message, bot response, and/or provided context @@ -242,6 +243,7 @@ async def patronus_api_check_output( user_input=user_input, bot_response=bot_response, provided_context=provided_context, + http_client=http_client, ) passed = check_guardrail_pass(response=response, success_strategy=success_strategy) metadata = {"pass": passed} diff --git a/nemoguardrails/library/policyai/actions.py b/nemoguardrails/library/policyai/actions.py index 9ef19a666c..91e784230b 100644 --- a/nemoguardrails/library/policyai/actions.py +++ b/nemoguardrails/library/policyai/actions.py @@ -27,10 +27,9 @@ import os from typing import Optional -import aiohttp - from nemoguardrails.actions import action from nemoguardrails.actions.rail_outcome import RailOutcome +from nemoguardrails.http import HTTPClient, http_call log = logging.getLogger(__name__) @@ -47,6 +46,7 @@ def _policyai_outcome(metadata: dict) -> RailOutcome: async def call_policyai_api( text: Optional[str] = None, tag_name: Optional[str] = None, + http_client: Optional[HTTPClient] = None, **kwargs, ) -> RailOutcome: """ @@ -89,63 +89,53 @@ async def call_policyai_api( ], } - timeout = aiohttp.ClientTimeout(total=30) - async with aiohttp.ClientSession(timeout=timeout) as session: - async with session.post( - url=url, - headers=headers, - json=data, - ) as response: - if response.status != 200: - raise ValueError( - f"PolicyAI call failed with status code {response.status}.\nDetails: {await response.text()}" - ) - response_json = await response.json() - log.info(json.dumps(response_json, indent=2)) - - # PolicyAI returns results in "data" array for tag-based evaluation - results = response_json.get("data", []) - - # Fail-closed: If no policies are attached to the tag, raise an error - # rather than silently allowing content through - if not results: - raise ValueError( - f"PolicyAI returned no policy results for tag '{tag_name}'. " - "Ensure policies are attached to this tag." - ) - - # Check if all policies failed evaluation - successful_results = [r for r in results if r.get("status") != "failed"] - if not successful_results: - raise ValueError( - f"All PolicyAI policy evaluations failed for tag '{tag_name}'. Check policy configurations." - ) - - # Aggregate results - if ANY policy returns UNSAFE, overall is UNSAFE - overall_assessment = "SAFE" - triggered_category = "Safe" - max_severity = 0 - reason = "Content passed all policy checks" - - for result in successful_results: - assessment = result.get("assessment", "SAFE") - if assessment == "UNSAFE": - overall_assessment = "UNSAFE" - triggered_category = result.get("category", "Unknown") - max_severity = max(max_severity, result.get("severity", 0)) - reason = result.get("reason", "Policy violation detected") - break # Stop at first UNSAFE result - - # Pre-format exception message for Colang 1.x compatibility - # (Colang 1.x doesn't support string concatenation in create event) - exception_message = f"PolicyAI moderation triggered. Content violated policy: {triggered_category}" - - return _policyai_outcome( - { - "assessment": overall_assessment, - "category": triggered_category, - "severity": max_severity, - "reason": reason, - "exception_message": exception_message, - } - ) + response = await http_call( + http_client, + "POST", + url, + headers=headers, + json=data, + timeout=30, + raise_for_status=False, + ) + if response.status_code != 200: + raise ValueError(f"PolicyAI call failed with status code {response.status_code}.\nDetails: {response.text}") + response_json = response.json() + log.info(json.dumps(response_json, indent=2)) + + results = response_json.get("data", []) + + if not results: + raise ValueError( + f"PolicyAI returned no policy results for tag '{tag_name}'. Ensure policies are attached to this tag." + ) + + successful_results = [r for r in results if r.get("status") != "failed"] + if not successful_results: + raise ValueError(f"All PolicyAI policy evaluations failed for tag '{tag_name}'. Check policy configurations.") + + overall_assessment = "SAFE" + triggered_category = "Safe" + max_severity = 0 + reason = "Content passed all policy checks" + + for result in successful_results: + assessment = result.get("assessment", "SAFE") + if assessment == "UNSAFE": + overall_assessment = "UNSAFE" + triggered_category = result.get("category", "Unknown") + max_severity = max(max_severity, result.get("severity", 0)) + reason = result.get("reason", "Policy violation detected") + break + + exception_message = f"PolicyAI moderation triggered. Content violated policy: {triggered_category}" + + return _policyai_outcome( + { + "assessment": overall_assessment, + "category": triggered_category, + "severity": max_severity, + "reason": reason, + "exception_message": exception_message, + } + ) diff --git a/nemoguardrails/library/prompt_security/actions.py b/nemoguardrails/library/prompt_security/actions.py index 5119dd43e6..79699eaee6 100644 --- a/nemoguardrails/library/prompt_security/actions.py +++ b/nemoguardrails/library/prompt_security/actions.py @@ -19,10 +19,9 @@ import os from typing import Optional -import httpx - from nemoguardrails.actions import action from nemoguardrails.actions.rail_outcome import RailOutcome, TransformTarget +from nemoguardrails.http import HTTPClient, http_call log = logging.getLogger(__name__) @@ -34,9 +33,13 @@ async def ps_protect_api_async( system_prompt: Optional[str] = None, response: Optional[str] = None, user: Optional[str] = None, + http_client: Optional[HTTPClient] = None, ): """Calls Prompt Security Protect API asynchronously. + Provider failures are logged and fail open with the default ``log`` + action. + Args: ps_protect_url: the URL of the protect endpoint given by Prompt Security. URL is https://[REGION].prompt.security/api/protect where REGION is eu, useast or apac @@ -52,6 +55,8 @@ async def ps_protect_api_async( user: the user ID or username for context. + http_client: Optional caller-owned HTTP client. + Returns: A dictionary with the following items: - is_blocked: True if the text should be blocked, False otherwise. @@ -69,23 +74,29 @@ async def ps_protect_api_async( "response": response, "user": user, } - async with httpx.AsyncClient() as client: - modified_text = None - ps_action = "log" - try: - ret = await client.post(ps_protect_url, headers=headers, json=payload) - res = ret.json() - ps_action = res.get("result", {}).get("action", "log") - if ps_action == "modify": - key = "response" if response else "prompt" - modified_text = res.get("result", {}).get(key, {}).get("modified_text") - except Exception as e: - log.error("Error calling Prompt Security Protect API: %s", e) - return { - "is_blocked": ps_action == "block", - "is_modified": ps_action == "modify", - "modified_text": modified_text, - } + modified_text = None + ps_action = "log" + try: + result = await http_call( + http_client, + "POST", + ps_protect_url, + headers=headers, + json=payload, + raise_for_status=False, + ) + data = result.json() + ps_action = data.get("result", {}).get("action", "log") + if ps_action == "modify": + key = "response" if response else "prompt" + modified_text = data.get("result", {}).get(key, {}).get("modified_text") + except Exception as e: + log.error("Error calling Prompt Security Protect API: %s", e) + return { + "is_blocked": ps_action == "block", + "is_modified": ps_action == "modify", + "modified_text": modified_text, + } def _protect_text_outcome(result: dict, target: TransformTarget) -> RailOutcome: @@ -101,12 +112,19 @@ def _protect_text_outcome(result: dict, target: TransformTarget) -> RailOutcome: async def protect_text( user_prompt: Optional[str] = None, bot_response: Optional[str] = None, + http_client: Optional[HTTPClient] = None, **kwargs, ) -> RailOutcome: """Protects the given user_prompt or bot_response. + + Provider results may allow, block, or transform the selected message. + Provider failures fail open. A bot response takes precedence when both + content arguments are provided. + Args: user_prompt: The user message to protect. bot_response: The bot message to protect. + http_client: Optional caller-owned HTTP client. Returns: RailOutcome.block(), RailOutcome.transform(), or RailOutcome.allow(). Raises: @@ -126,13 +144,25 @@ async def protect_text( if bot_response: return _protect_text_outcome( - await ps_protect_api_async(ps_protect_url, ps_app_id, None, None, bot_response), + await ps_protect_api_async( + ps_protect_url, + ps_app_id, + None, + None, + bot_response, + http_client=http_client, + ), TransformTarget.BOT_MESSAGE, ) if user_prompt: return _protect_text_outcome( - await ps_protect_api_async(ps_protect_url, ps_app_id, user_prompt), + await ps_protect_api_async( + ps_protect_url, + ps_app_id, + user_prompt, + http_client=http_client, + ), TransformTarget.USER_MESSAGE, ) diff --git a/nemoguardrails/library/trend_micro/actions.py b/nemoguardrails/library/trend_micro/actions.py index 951e189980..c8566e8b09 100644 --- a/nemoguardrails/library/trend_micro/actions.py +++ b/nemoguardrails/library/trend_micro/actions.py @@ -16,13 +16,13 @@ import logging from typing import Literal -import httpx -from pydantic import BaseModel, Field, field_validator, model_validator +from pydantic import BaseModel, Field, ValidationError, field_validator, model_validator from pydantic_core import to_json from typing_extensions import cast from nemoguardrails.actions import action from nemoguardrails.actions.rail_outcome import RailOutcome +from nemoguardrails.http import HTTPClient, HTTPClientError, http_call from nemoguardrails.rails.llm.config import RailsConfig, TrendMicroRailConfig log = logging.getLogger(__name__) @@ -93,11 +93,26 @@ def _trend_micro_outcome(result: GuardResult) -> RailOutcome: @action(is_system_action=True) -async def trend_ai_guard(config: RailsConfig, text: str) -> RailOutcome: +async def trend_ai_guard( + config: RailsConfig, + text: str, + http_client: HTTPClient | None = None, +) -> RailOutcome: """ Custom action to invoke the Trend Micro AI Guard API. - """ + A missing API key blocks the content. A valid provider response supplies + the allow or block decision, while transport, status, decoding, and + validation failures fail open. + + Args: + config: Rails configuration containing the Trend Micro settings. + text: Content to evaluate. + http_client: Optional caller-owned HTTP client. + + Returns: + The provider decision or the integration's fallback outcome. + """ trend_config = get_config(config) # No checks required since default is set in TrendMicroRailConfig @@ -115,38 +130,37 @@ async def trend_ai_guard(config: RailsConfig, text: str) -> RailOutcome: app_name = trend_config.application_name - async with httpx.AsyncClient() as client: - data = Guard(prompt=text).model_dump() + data = Guard(prompt=text).model_dump() - # Build headers with required TMV1-Application-Name - headers = { - "Authorization": f"Bearer {v1_api_key}", - "Content-Type": "application/json", - "TMV1-Application-Name": app_name, - } + headers = { + "Authorization": f"Bearer {v1_api_key}", + "Content-Type": "application/json", + "TMV1-Application-Name": app_name, + } - # Add Prefer header for detail level control - if trend_config.detailed_response: - headers["Prefer"] = "return=representation" - else: - headers["Prefer"] = "return=minimal" + if trend_config.detailed_response: + headers["Prefer"] = "return=representation" + else: + headers["Prefer"] = "return=minimal" - response = await client.post( + try: + response = await http_call( + http_client, + "POST", v1_url, content=to_json(data), headers=headers, + raise_for_status=False, ) - - try: - response.raise_for_status() - guard_result = GuardResult(**response.json()) - log.debug("Trend Micro AI Guard Result: %s", guard_result) - except httpx.HTTPStatusError as e: - log.error("Error calling Trend Micro AI Guard API: %s", e) - return _trend_micro_outcome( - GuardResult( - action="Allow", - reason="An error occurred while calling the Trend Micro AI Guard API.", - ) + response.raise_for_status() + guard_result = GuardResult.model_validate(response.json()) + log.debug("Trend Micro AI Guard Result: %s", guard_result) + except (HTTPClientError, ValidationError) as e: + log.error("Error calling Trend Micro AI Guard API: %s", e) + return _trend_micro_outcome( + GuardResult( + action="Allow", + reason="An error occurred while calling the Trend Micro AI Guard API.", ) - return _trend_micro_outcome(guard_result) + ) + return _trend_micro_outcome(guard_result) diff --git a/nemoguardrails/testing/http.py b/nemoguardrails/testing/http.py index 78f7cbd27a..483cbcaa0b 100644 --- a/nemoguardrails/testing/http.py +++ b/nemoguardrails/testing/http.py @@ -32,6 +32,9 @@ def __init__(self, responses: Iterable[HTTPResponse | BaseException] = ()): self._responses = deque(responses) self.close_calls = 0 + def add_response(self, response: HTTPResponse | BaseException) -> None: + self._responses.append(response) + async def request( self, method: str, diff --git a/tests/http/test_contract.py b/tests/http/test_contract.py index e0a8296269..fac8b9817f 100644 --- a/tests/http/test_contract.py +++ b/tests/http/test_contract.py @@ -164,6 +164,15 @@ async def test_recording_client_requires_a_scripted_response(): await client.request("GET", "https://example.com") +@pytest.mark.asyncio +async def test_recording_client_accepts_responses_after_construction(): + response = HTTPResponse(status_code=200) + client = RecordingHTTPClient() + client.add_response(response) + + assert await client.request("GET", "https://example.com") is response + + def test_neutral_transport_errors_have_a_shared_base(): assert issubclass(HTTPConnectionError, HTTPClientError) assert issubclass(HTTPTimeoutError, HTTPClientError) diff --git a/tests/http/test_library_boundary.py b/tests/http/test_library_boundary.py new file mode 100644 index 0000000000..ab8d4dcaa0 --- /dev/null +++ b/tests/http/test_library_boundary.py @@ -0,0 +1,44 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 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 ast +from pathlib import Path + +import pytest + +LIBRARY_ROOT = Path(__file__).parents[2] / "nemoguardrails" / "library" +FORBIDDEN_HTTP_MODULES = frozenset({"aiohttp", "httpx", "requests", "urllib3"}) + + +def _direct_http_imports(path: Path) -> set[str]: + tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) + imports: set[str] = set() + for node in ast.walk(tree): + if isinstance(node, ast.Import): + imports.update(alias.name.partition(".")[0] for alias in node.names) + elif isinstance(node, ast.ImportFrom) and node.module is not None: + imports.add(node.module.partition(".")[0]) + return imports & FORBIDDEN_HTTP_MODULES + + +@pytest.mark.unit +def test_library_uses_canonical_http_boundary(): + violations = { + str(path.relative_to(LIBRARY_ROOT)): sorted(imports) + for path in sorted(LIBRARY_ROOT.rglob("*.py")) + if (imports := _direct_http_imports(path)) + } + + assert violations == {} diff --git a/tests/http/test_runtime.py b/tests/http/test_runtime.py index f3f571b7ba..21ddbc9f34 100644 --- a/tests/http/test_runtime.py +++ b/tests/http/test_runtime.py @@ -32,11 +32,11 @@ def test_default_factory_forwards_owned_client_options(): with mock.patch("nemoguardrails.http.transport.httpx.AsyncClient") as factory: create_http_client(limits=limits, follow_redirects=True) - factory.assert_called_once_with( - timeout=None, - limits=limits, - follow_redirects=True, - ) + factory.assert_called_once() + kwargs = factory.call_args.kwargs + assert kwargs["timeout"] is None + assert kwargs["limits"] is limits + assert kwargs["follow_redirects"] is True @pytest.mark.asyncio diff --git a/tests/http/test_testing.py b/tests/http/test_testing.py new file mode 100644 index 0000000000..fd11cb0e2b --- /dev/null +++ b/tests/http/test_testing.py @@ -0,0 +1,42 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 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 pytest + +from tests.http_utils import RecordedHTTPResponses + + +def test_recorded_http_responses_require_registered_urls_to_be_used(): + with pytest.raises(AssertionError): + with RecordedHTTPResponses() as responses: + responses.post("https://unused.example", payload={"ok": True}) + + +@pytest.mark.asyncio +async def test_recorded_http_responses_require_exact_request_order(): + with pytest.raises(AssertionError): + with RecordedHTTPResponses() as responses: + responses.post("https://first.example", payload={"ok": True}) + responses.post("https://second.example", payload={"ok": True}) + await responses.client.request("POST", "https://second.example") + await responses.client.request("POST", "https://first.example") + + +@pytest.mark.asyncio +async def test_recorded_http_responses_require_exact_request_count(): + with pytest.raises(AssertionError): + with RecordedHTTPResponses() as responses: + responses.post("https://example.test", payload={"ok": True}, times=2) + await responses.client.request("POST", "https://example.test") diff --git a/tests/http/test_transport.py b/tests/http/test_transport.py index 486c343c8a..9353863b6a 100644 --- a/tests/http/test_transport.py +++ b/tests/http/test_transport.py @@ -22,6 +22,7 @@ from nemoguardrails.http import ( HTTPConnectionError, HTTPTimeoutError, + HTTPTLSConfig, HttpxHTTPClient, ) @@ -201,3 +202,42 @@ async def test_httpx_transport_rejects_owned_options_with_injected_client(option HttpxHTTPClient(injected, **options) await injected.aclose() + + +def test_httpx_transport_configures_owned_client_tls(): + constructed = mock.MagicMock() + + with mock.patch("nemoguardrails.http.transport.httpx.AsyncClient", return_value=constructed) as factory: + client = HttpxHTTPClient( + timeout=12.0, + tls=HTTPTLSConfig( + ca_bundle="/ca.pem", + client_certificate="/cert.pem", + client_key="/key.pem", + ), + ) + + assert client._client is constructed + assert client._timeout == 12.0 + assert factory.call_args.kwargs["timeout"] is None + assert factory.call_args.kwargs["verify"] == "/ca.pem" + assert factory.call_args.kwargs["cert"] == ("/cert.pem", "/key.pem") + + +def test_httpx_transport_rejects_tls_with_injected_client(): + injected = mock.MagicMock(spec=httpx.AsyncClient) + + with pytest.raises(ValueError, match="TLS configuration"): + HttpxHTTPClient(injected, tls=HTTPTLSConfig(verify=False)) + + +@pytest.mark.parametrize( + "kwargs", + [ + {"client_certificate": "/cert.pem"}, + {"client_key": "/key.pem"}, + ], +) +def test_http_tls_config_requires_complete_client_credentials(kwargs): + with pytest.raises(ValueError, match="configured together"): + HTTPTLSConfig(**kwargs) diff --git a/tests/http_utils.py b/tests/http_utils.py new file mode 100644 index 0000000000..2c1f4f5bc3 --- /dev/null +++ b/tests/http_utils.py @@ -0,0 +1,61 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 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 json +from typing import Any + +from nemoguardrails.http import HTTPResponse +from nemoguardrails.testing import RecordingHTTPClient + + +class RecordedHTTPResponses: + """Register responses and verify the exact requests made by a test.""" + + def __init__(self): + self.client = RecordingHTTPClient() + self.expected_requests: list[tuple[str, str]] = [] + + def post( + self, + url: str, + *, + payload: Any = None, + status: int = 200, + body: str | None = None, + headers: dict[str, str] | None = None, + content_type: str | None = None, + exception: BaseException | None = None, + times: int = 1, + ) -> None: + if exception is not None: + response = exception + else: + response_headers = dict(headers or {}) + if content_type is not None: + response_headers["Content-Type"] = content_type + content = body.encode() if body is not None else json.dumps(payload).encode() + response = HTTPResponse(status_code=status, headers=response_headers, content=content) + + for _ in range(times): + self.client.add_response(response) + self.expected_requests.append(("POST", url)) + + def __enter__(self) -> "RecordedHTTPResponses": + return self + + def __exit__(self, exc_type: object, exc_value: object, traceback: object) -> None: + if exc_type is None: + actual_requests = [(request.method, request.url) for request in self.client.requests] + assert self.expected_requests == actual_requests, (self.expected_requests, actual_requests) diff --git a/tests/test_ai_defense.py b/tests/test_ai_defense.py index 2602d4272a..93f153fe4c 100644 --- a/tests/test_ai_defense.py +++ b/tests/test_ai_defense.py @@ -13,12 +13,16 @@ # See the License for the specific language governing permissions and # limitations under the License. +import json import os import pytest from nemoguardrails import RailsConfig from nemoguardrails.actions.rail_outcome import RailOutcome +from nemoguardrails.http import HTTPConnectionError, HTTPResponse +from nemoguardrails.library.ai_defense.actions import ai_defense_inspect +from nemoguardrails.testing import RecordingHTTPClient from tests.utils import TestChat # Note: we don't call the action directly in these tests; we exercise it via flows. @@ -40,6 +44,79 @@ def mock_request(*args, **kwargs): API_ENDPOINT = "https://us.api.inspect.aidefense.security.cisco.com/api/v1/inspect/chat" +def _ai_defense_config(*, fail_open: bool) -> RailsConfig: + return RailsConfig.from_content( + yaml_content=f""" + models: [] + rails: + config: + ai_defense: + fail_open: {str(fail_open).lower()} + """ + ) + + +@pytest.mark.parametrize(("fail_open", "is_blocked"), [(False, True), (True, False)]) +@pytest.mark.asyncio +async def test_ai_defense_invalid_json_uses_failure_policy(monkeypatch, caplog, fail_open, is_blocked): + monkeypatch.setenv("AI_DEFENSE_API_KEY", "secret") + monkeypatch.setenv("AI_DEFENSE_API_ENDPOINT", API_ENDPOINT) + client = RecordingHTTPClient([HTTPResponse(status_code=200, content=b"not-json")]) + + result = await ai_defense_inspect( + _ai_defense_config(fail_open=fail_open), + user_prompt="Hello", + http_client=client, + ) + + assert result.is_blocked is is_blocked + assert "AI Defense API returned malformed JSON" in caplog.text + assert "Error calling AI Defense API" not in caplog.text + + +@pytest.mark.parametrize("body", [[], "text", 1, None]) +@pytest.mark.parametrize(("fail_open", "is_blocked"), [(False, True), (True, False)]) +@pytest.mark.asyncio +async def test_ai_defense_non_object_json_uses_failure_policy(monkeypatch, caplog, body, fail_open, is_blocked): + monkeypatch.setenv("AI_DEFENSE_API_KEY", "secret") + monkeypatch.setenv("AI_DEFENSE_API_ENDPOINT", API_ENDPOINT) + client = RecordingHTTPClient( + [ + HTTPResponse( + status_code=200, + content=json.dumps(body).encode(), + ) + ] + ) + + result = await ai_defense_inspect( + _ai_defense_config(fail_open=fail_open), + user_prompt="Hello", + http_client=client, + ) + + assert result.is_blocked is is_blocked + assert "AI Defense API returned malformed response (expected an object)" in caplog.text + assert f"fail_open={fail_open}" in caplog.text + + +@pytest.mark.asyncio +async def test_ai_defense_transport_error_has_distinct_log(monkeypatch, caplog): + monkeypatch.setenv("AI_DEFENSE_API_KEY", "secret") + monkeypatch.setenv("AI_DEFENSE_API_ENDPOINT", API_ENDPOINT) + client = RecordingHTTPClient([HTTPConnectionError("connection failed")]) + + result = await ai_defense_inspect( + _ai_defense_config(fail_open=False), + user_prompt="Hello", + http_client=client, + ) + + assert result.is_blocked + assert "Error calling AI Defense API" in caplog.text + assert "returned malformed JSON" not in caplog.text + + # Set environment variables for tests requiring real API calls @pytest.fixture(autouse=True) def _env(monkeypatch): @@ -999,6 +1076,7 @@ async def test_ai_defense_inspect_http_504_gateway_timeout(httpx_mock): url="https://test.example.com/api/v1/inspect/chat", status_code=504, text="Gateway Timeout", + is_reusable=True, ) # Create a minimal config for the test (fail_open defaults to False) @@ -1192,6 +1270,7 @@ async def test_ai_defense_inspect_api_failure_fail_closed(httpx_mock): method="POST", url="https://test.example.com/api/v1/inspect/chat", status_code=500, + is_reusable=True, ) # With fail_closed, should return is_blocked=True instead of raising @@ -1242,6 +1321,7 @@ async def test_ai_defense_inspect_api_failure_fail_open(httpx_mock): method="POST", url="https://test.example.com/api/v1/inspect/chat", status_code=500, + is_reusable=True, ) result = await ai_defense_inspect(config, user_prompt="Hello, how are you?") diff --git a/tests/test_autoalign.py b/tests/test_autoalign.py index 28f49983d1..c0e18f4bfe 100644 --- a/tests/test_autoalign.py +++ b/tests/test_autoalign.py @@ -20,7 +20,17 @@ from nemoguardrails import RailsConfig from nemoguardrails.actions.rail_outcome import RailOutcome, TransformTarget -from nemoguardrails.library.autoalign.actions import _autoalign_outcome, _autoalign_score_outcome +from nemoguardrails.http import HTTPResponse +from nemoguardrails.library.autoalign.actions import ( + _autoalign_outcome, + _autoalign_score_outcome, + autoalign_factcheck_output_api, + autoalign_groundedness_output_api, + autoalign_input_api, + autoalign_output_api, +) +from nemoguardrails.llm.taskmanager import LLMTaskManager +from nemoguardrails.testing import RecordingHTTPClient from tests.utils import TestChat CONFIGS_FOLDER = os.path.join(os.path.dirname(__file__), ".", "test_configs") @@ -102,6 +112,149 @@ def test_autoalign_score_outcome(score, threshold, expected): assert _autoalign_score_outcome(score, threshold) == expected +def _autoalign_config() -> RailsConfig: + return RailsConfig.from_content( + yaml_content=""" + models: [] + rails: + config: + autoalign: + parameters: + endpoint: https://autoalign.example/guardrail + groundedness_check_endpoint: https://autoalign.example/groundedness + fact_check_endpoint: https://autoalign.example/factcheck + multi_language: true + input: + guardrails_config: + toxicity_detection: {} + output: + guardrails_config: + toxicity_detection: {} + """ + ) + + +@pytest.mark.asyncio +async def test_autoalign_actions_use_shared_client(monkeypatch): + monkeypatch.setenv("AUTOALIGN_API_KEY", "secret") + safe_response = b'{"task":"toxicity_detection","guarded":false,"response":"","output_data":[]}\n' + client = RecordingHTTPClient( + [ + HTTPResponse(status_code=200, content=safe_response), + HTTPResponse(status_code=200, content=safe_response), + HTTPResponse( + status_code=200, + content=b'{"task":"groundedness_checker","response":"Factcheck Score: 0.75"}\n', + ), + HTTPResponse( + status_code=200, + content=b'{"all_overall_fact_scores":[0.9]}', + ), + ] + ) + task_manager = LLMTaskManager(config=_autoalign_config()) + + input_outcome = await autoalign_input_api( + task_manager, + context={"user_message": "input"}, + http_client=client, + ) + output_outcome = await autoalign_output_api( + task_manager, + context={"bot_message": "output"}, + http_client=client, + ) + groundedness_outcome = await autoalign_groundedness_output_api( + task_manager, + context={"bot_message": "answer", "relevant_chunks_sep": ["evidence"]}, + factcheck_threshold=0.5, + http_client=client, + ) + factcheck_outcome = await autoalign_factcheck_output_api( + task_manager, + context={"user_message": "question", "bot_message": "answer"}, + factcheck_threshold=0.5, + http_client=client, + ) + + assert input_outcome.decision == "allow" + assert output_outcome.decision == "allow" + assert groundedness_outcome == RailOutcome.allow(metadata={"score": 0.75, "threshold": 0.5}) + assert factcheck_outcome == RailOutcome.allow(metadata={"score": 0.9, "threshold": 0.5}) + + input_request, output_request, groundedness_request, factcheck_request = client.requests + assert input_request.method == "POST" + assert input_request.url == "https://autoalign.example/guardrail" + assert input_request.headers == {"x-api-key": "secret"} + assert input_request.json["prompt"] == "input" + assert input_request.json["multi_language"] is True + assert input_request.json["config"]["toxicity_detection"]["mode"] == "DETECT" + assert input_request.timeout is None + + assert output_request.url == "https://autoalign.example/guardrail" + assert output_request.json["prompt"] == "output" + assert output_request.timeout is None + + assert groundedness_request.url == "https://autoalign.example/groundedness" + assert groundedness_request.json == { + "prompt": "answer", + "documents": ["evidence"], + "config": { + "groundedness_checker": {"verify_response": False}, + "toxicity_detection": {}, + }, + } + assert groundedness_request.timeout is None + + assert factcheck_request.url == "https://autoalign.example/factcheck" + assert factcheck_request.json == { + "labels": {"session_id": "nemo", "api_key": "secret"}, + "content_moderation_docs": [ + { + "content": "answer", + "type": "text_content", + "file_name": "HighlightedText.txt", + } + ], + "user_query": "question", + "config": {"toxicity_detection": {}}, + "multi_language": True, + } + assert factcheck_request.timeout is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("action", "context", "endpoint"), + [ + (autoalign_input_api, {"user_message": "input"}, "guardrail"), + (autoalign_output_api, {"bot_message": "output"}, "guardrail"), + ( + autoalign_groundedness_output_api, + {"bot_message": "answer", "relevant_chunks_sep": ["evidence"]}, + "groundedness", + ), + ( + autoalign_factcheck_output_api, + {"user_message": "question", "bot_message": "answer"}, + "factcheck", + ), + ], +) +async def test_autoalign_actions_preserve_http_failures(monkeypatch, action, context, endpoint): + monkeypatch.setenv("AUTOALIGN_API_KEY", "secret") + client = RecordingHTTPClient([HTTPResponse(status_code=503, content=b"unavailable")]) + + with pytest.raises(ValueError, match="AutoAlign call failed with status code 503"): + await action( + LLMTaskManager(config=_autoalign_config()), + context=context, + http_client=client, + ) + + assert client.requests[0].url == f"https://autoalign.example/{endpoint}" + + @pytest.mark.asyncio async def test_autoalign_greeting(): # Test 1 - Greeting - No fact-checking invocation should happen diff --git a/tests/test_clavata.py b/tests/test_clavata.py index f2315df7d6..60671e43b8 100644 --- a/tests/test_clavata.py +++ b/tests/test_clavata.py @@ -13,13 +13,14 @@ # See the License for the specific language governing permissions and # limitations under the License. +import json from typing import Any, Dict, Optional import pytest -from aioresponses import aioresponses from nemoguardrails import RailsConfig from nemoguardrails.actions.actions import ActionResult, action +from nemoguardrails.http import HTTPResponse from nemoguardrails.library.clavata.request import ( CreateJobResponse, Job, @@ -27,6 +28,7 @@ Result, SectionReport, ) +from nemoguardrails.testing import RecordingHTTPClient from tests.utils import TestChat @@ -55,6 +57,21 @@ def retrieve_relevant_chunks(): """ +@pytest.fixture +def http_client() -> RecordingHTTPClient: + return RecordingHTTPClient() + + +def add_clavata_response(client: RecordingHTTPClient, response: Dict[str, Any]) -> None: + client.add_response( + HTTPResponse( + status_code=200, + headers={"Content-Type": "application/json"}, + content=json.dumps(response).encode(), + ) + ) + + @pytest.mark.unit def test_clavata_no_active_policy_check(monkeypatch): """Test that without active policy checks, messages pass through.""" @@ -87,7 +104,7 @@ def test_clavata_no_active_policy_check(monkeypatch): @pytest.mark.unit @pytest.mark.asyncio -async def test_clavata_input_policy_check(monkeypatch): +async def test_clavata_input_policy_check(monkeypatch, http_client): """Test that input policy checks block messages about animal sounds.""" monkeypatch.setenv("CLAVATA_API_KEY", "test_api_key") @@ -116,27 +133,24 @@ async def test_clavata_input_policy_check(monkeypatch): ], ) - # Mock response from Clavata API - with aioresponses() as m: - chat.app.register_action(retrieve_relevant_chunks, "retrieve_relevant_chunks") - - # Use the factory to create a response with a matching policy and label - mock_response = create_clavata_response(labels={"DogBarking": True}) + chat.app.register_action(retrieve_relevant_chunks, "retrieve_relevant_chunks") + chat.app.register_action_param("http_client", http_client) - m.post( - "https://gateway.app.clavata.ai:8443/v1/jobs", - payload=mock_response, - status=200, - ) + # Use the factory to create a response with a matching policy and label + mock_response = create_clavata_response(labels={"DogBarking": True}) + add_clavata_response(http_client, mock_response) - chat >> "Woof woof" - # Block given the policy matched - await chat.bot_async("I cannot respond to that request.") + chat >> "Woof woof" + # Block given the policy matched + await chat.bot_async("I cannot respond to that request.") + assert http_client.requests[0].method == "POST" + assert http_client.requests[0].url == "https://gateway.app.clavata.ai:8443/v1/jobs" + assert http_client.requests[0].headers == {"Authorization": "Bearer test_api_key"} @pytest.mark.unit @pytest.mark.asyncio -async def test_clavata_label_match_logic_any(monkeypatch): +async def test_clavata_label_match_logic_any(monkeypatch, http_client): """Test that label_match_logic: ANY works correctly when at least one label matches.""" monkeypatch.setenv("CLAVATA_API_KEY", "test_api_key") @@ -170,28 +184,23 @@ async def test_clavata_label_match_logic_any(monkeypatch): ], ) - with aioresponses() as m: - chat.app.register_action(retrieve_relevant_chunks, "retrieve_relevant_chunks") - - # One matching label - mock_response = create_clavata_response( - labels={"DogBarking": False, "CatMeowing": False, "CowMooing": True}, - ) + chat.app.register_action(retrieve_relevant_chunks, "retrieve_relevant_chunks") + chat.app.register_action_param("http_client", http_client) - m.post( - "https://gateway.app.clavata.ai:8443/v1/jobs", - payload=mock_response, - status=200, - ) + # One matching label + mock_response = create_clavata_response( + labels={"DogBarking": False, "CatMeowing": False, "CowMooing": True}, + ) + add_clavata_response(http_client, mock_response) - chat >> "Moo" - # Block given ANY logic is used and one label matches - await chat.bot_async("I cannot respond to that request.") + chat >> "Moo" + # Block given ANY logic is used and one label matches + await chat.bot_async("I cannot respond to that request.") @pytest.mark.unit @pytest.mark.asyncio -async def test_clavata_label_match_logic_any_no_match(monkeypatch): +async def test_clavata_label_match_logic_any_no_match(monkeypatch, http_client): """Test that label_match_logic: ANY allows messages when no specified labels match.""" monkeypatch.setenv("CLAVATA_API_KEY", "test_api_key") @@ -224,27 +233,22 @@ async def test_clavata_label_match_logic_any_no_match(monkeypatch): ], ) - with aioresponses() as m: - chat.app.register_action(retrieve_relevant_chunks, "retrieve_relevant_chunks") - - mock_response = create_clavata_response( - labels={"DogBarking": False, "CatMeowing": False, "HorseNeighing": False}, - ) + chat.app.register_action(retrieve_relevant_chunks, "retrieve_relevant_chunks") + chat.app.register_action_param("http_client", http_client) - m.post( - "https://gateway.app.clavata.ai:8443/v1/jobs", - payload=mock_response, - status=200, - ) + mock_response = create_clavata_response( + labels={"DogBarking": False, "CatMeowing": False, "HorseNeighing": False}, + ) + add_clavata_response(http_client, mock_response) - chat >> "Hey" - # Pass given no labels matched - await chat.bot_async("Hello there!") + chat >> "Hey" + # Pass given no labels matched + await chat.bot_async("Hello there!") @pytest.mark.unit @pytest.mark.asyncio -async def test_clavata_label_match_logic_all(monkeypatch): +async def test_clavata_label_match_logic_all(monkeypatch, http_client): """Test that label_match_logic: ALL requires all specified labels to match.""" monkeypatch.setenv("CLAVATA_API_KEY", "test_api_key") @@ -278,32 +282,27 @@ async def test_clavata_label_match_logic_all(monkeypatch): ], ) - with aioresponses() as m: - chat.app.register_action(retrieve_relevant_chunks, "retrieve_relevant_chunks") - - mock_response = create_clavata_response( - labels={ - "DogBarking": True, - "CatMeowing": True, - "CowMooing": True, - "HorseNeighing": False, - }, - ) - - m.post( - "https://gateway.app.clavata.ai:8443/v1/jobs", - payload=mock_response, - status=200, - ) + chat.app.register_action(retrieve_relevant_chunks, "retrieve_relevant_chunks") + chat.app.register_action_param("http_client", http_client) + + mock_response = create_clavata_response( + labels={ + "DogBarking": True, + "CatMeowing": True, + "CowMooing": True, + "HorseNeighing": False, + }, + ) + add_clavata_response(http_client, mock_response) - chat >> "Woof woof, meow, moo." - # Block given all specified labels matched and we're using ALL logic - await chat.bot_async("I cannot respond to that request.") + chat >> "Woof woof, meow, moo." + # Block given all specified labels matched and we're using ALL logic + await chat.bot_async("I cannot respond to that request.") @pytest.mark.unit @pytest.mark.asyncio -async def test_clavata_label_match_logic_all_partial_match(monkeypatch): +async def test_clavata_label_match_logic_all_partial_match(monkeypatch, http_client): """Test that label_match_logic: ALL allows messages when only some labels match.""" monkeypatch.setenv("CLAVATA_API_KEY", "test_api_key") @@ -337,27 +336,22 @@ async def test_clavata_label_match_logic_all_partial_match(monkeypatch): ], ) - with aioresponses() as m: - chat.app.register_action(retrieve_relevant_chunks, "retrieve_relevant_chunks") - - mock_response = create_clavata_response( - labels={"DogBarking": True, "CatMeowing": True, "CowMooing": False}, - ) + chat.app.register_action(retrieve_relevant_chunks, "retrieve_relevant_chunks") + chat.app.register_action_param("http_client", http_client) - m.post( - "https://gateway.app.clavata.ai:8443/v1/jobs", - payload=mock_response, - status=200, - ) + mock_response = create_clavata_response( + labels={"DogBarking": True, "CatMeowing": True, "CowMooing": False}, + ) + add_clavata_response(http_client, mock_response) - chat >> "Hi! I want to talk about dogs and cats but not cows." - # Pass given not all specified labels matched - await chat.bot_async("Hello there!") + chat >> "Hi! I want to talk about dogs and cats but not cows." + # Pass given not all specified labels matched + await chat.bot_async("Hello there!") @pytest.mark.unit @pytest.mark.asyncio -async def test_clavata_empty_labels(monkeypatch): +async def test_clavata_empty_labels(monkeypatch, http_client): """Test that any policy match blocks the message even if no labels are specified.""" monkeypatch.setenv("CLAVATA_API_KEY", "test_api_key") @@ -387,25 +381,20 @@ async def test_clavata_empty_labels(monkeypatch): ], ) - with aioresponses() as m: - chat.app.register_action(retrieve_relevant_chunks, "retrieve_relevant_chunks") + chat.app.register_action(retrieve_relevant_chunks, "retrieve_relevant_chunks") + chat.app.register_action_param("http_client", http_client) - mock_response = create_clavata_response(labels={"SomeLabel": True}) + mock_response = create_clavata_response(labels={"SomeLabel": True}) + add_clavata_response(http_client, mock_response) - m.post( - "https://gateway.app.clavata.ai:8443/v1/jobs", - payload=mock_response, - status=200, - ) - - chat >> "Hi! I want to talk about something that matches the policy." - # Block given the policy matched - await chat.bot_async("I cannot respond to that request.") + chat >> "Hi! I want to talk about something that matches the policy." + # Block given the policy matched + await chat.bot_async("I cannot respond to that request.") @pytest.mark.unit @pytest.mark.asyncio -async def test_clavata_policy_no_match(monkeypatch): +async def test_clavata_policy_no_match(monkeypatch, http_client): """Test that when the policy doesn't match at all, the message passes through.""" monkeypatch.setenv("CLAVATA_API_KEY", "test_api_key") @@ -434,21 +423,16 @@ async def test_clavata_policy_no_match(monkeypatch): ], ) - with aioresponses() as m: - chat.app.register_action(retrieve_relevant_chunks, "retrieve_relevant_chunks") - - # No policy match - mock_response = create_clavata_response() + chat.app.register_action(retrieve_relevant_chunks, "retrieve_relevant_chunks") + chat.app.register_action_param("http_client", http_client) - m.post( - "https://gateway.app.clavata.ai:8443/v1/jobs", - payload=mock_response, - status=200, - ) + # No policy match + mock_response = create_clavata_response() + add_clavata_response(http_client, mock_response) - chat >> "Hi! I want to talk about something innocent." - # Pass given no policy match - await chat.bot_async("Hello there!") + chat >> "Hi! I want to talk about something innocent." + # Pass given no policy match + await chat.bot_async("Hello there!") def create_clavata_response( diff --git a/tests/test_clavata_models.py b/tests/test_clavata_models.py index c287aa5096..d0ffa44de8 100644 --- a/tests/test_clavata_models.py +++ b/tests/test_clavata_models.py @@ -13,18 +13,29 @@ # See the License for the specific language governing permissions and # limitations under the License. +import asyncio + import pytest from pydantic import ValidationError +import nemoguardrails.library.clavata.request as request_module +from nemoguardrails.http import HTTPConnectionError, HTTPResponse, RetryingHTTPClient from nemoguardrails.library.clavata.actions import LabelResult, PolicyResult -from nemoguardrails.library.clavata.errs import ClavataPluginAPIError +from nemoguardrails.library.clavata.errs import ( + ClavataPluginAPIError, + ClavataPluginAPIRateLimitError, + ClavataPluginValueError, +) from nemoguardrails.library.clavata.request import ( + _CLAVATA_RETRY_POLICY, + ClavataClient, CreateJobResponse, Job, Report, Result, SectionReport, ) +from nemoguardrails.testing import RecordingHTTPClient @pytest.mark.unit @@ -345,3 +356,136 @@ def test_model_validate_complete_job_response(self): assert sections[2].name == "Label3" assert sections[2].message == "Third section evaluation failed" assert sections[2].result == "OUTCOME_FAILED" + + +@pytest.fixture +def zero_sleep_clavata_retry(monkeypatch): + monkeypatch.setattr( + ClavataClient, + "_retrying_http_client", + staticmethod( + lambda client: RetryingHTTPClient( + client, + _CLAVATA_RETRY_POLICY, + sleep=lambda _: asyncio.sleep(0), + ) + ), + ) + + +@pytest.mark.unit +class TestClavataClientTransport: + @pytest.mark.asyncio + async def test_clavata_client_uses_shared_retry_policy(self, zero_sleep_clavata_retry): + response = CreateJobResponse( + job=Job( + status="JOB_STATUS_COMPLETED", + results=[Result(report=Report(result="OUTCOME_FALSE", sectionEvaluationReports=[]))], + ) + ) + transport = RecordingHTTPClient( + [ + HTTPResponse(status_code=429), + HTTPResponse(status_code=200, content=response.model_dump_json().encode()), + ] + ) + job = await ClavataClient( + "https://clavata.example", + api_key="test-key", + http_client=transport, + ).create_job("hello", "policy-id") + + assert job.status == "JOB_STATUS_COMPLETED" + assert len(transport.requests) == 2 + request = transport.requests[1] + assert request.method == "POST" + assert request.url == "https://clavata.example/v1/jobs" + assert request.headers == {"Authorization": "Bearer test-key"} + assert request.json == { + "content_data": [{"text": "hello"}], + "policy_id": "policy-id", + "wait_for_completion": True, + } + assert request.timeout is None + + @pytest.mark.asyncio + async def test_clavata_client_closes_owned_shared_client(self, monkeypatch): + response = CreateJobResponse( + job=Job( + status="JOB_STATUS_COMPLETED", + results=[Result(report=Report(result="OUTCOME_FALSE", sectionEvaluationReports=[]))], + ) + ) + transport = RecordingHTTPClient([HTTPResponse(status_code=200, content=response.model_dump_json().encode())]) + monkeypatch.setattr(request_module, "create_http_client", lambda: transport) + + job = await ClavataClient( + "https://clavata.example", + api_key="test-key", + ).create_job("hello", "policy-id") + + assert job.status == "JOB_STATUS_COMPLETED" + assert transport.close_calls == 1 + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("response", "error", "message"), + [ + ( + HTTPResponse(status_code=503, content=b"unavailable"), + ClavataPluginAPIError, + "Clavata call failed with status code 503", + ), + ( + HTTPResponse(status_code=200, content=b"not-json"), + ClavataPluginValueError, + "Failed to parse Clavata response as JSON", + ), + ( + HTTPResponse(status_code=200, content=b'{"unexpected":true}'), + ClavataPluginValueError, + "Invalid response format from Clavata API", + ), + ], + ) + async def test_clavata_client_preserves_response_failure_type(self, response, error, message): + transport = RecordingHTTPClient([response]) + + with pytest.raises(error, match=message): + await ClavataClient( + "https://clavata.example", + api_key="test-key", + http_client=transport, + ).create_job("hello", "policy-id") + + @pytest.mark.asyncio + async def test_clavata_client_does_not_retry_ambiguous_transport_failure(self): + transport = RecordingHTTPClient([HTTPConnectionError("connection lost")]) + + with pytest.raises(ClavataPluginAPIError): + await ClavataClient( + "https://clavata.example", + api_key="test-key", + http_client=transport, + ).create_job("hello", "policy-id") + + assert len(transport.requests) == 1 + + @pytest.mark.asyncio + async def test_clavata_client_preserves_rate_limit_error(self, zero_sleep_clavata_retry): + transport = RecordingHTTPClient( + [ + HTTPResponse(status_code=429), + HTTPResponse(status_code=429), + HTTPResponse(status_code=429), + ] + ) + + with pytest.raises(ClavataPluginAPIRateLimitError): + await ClavataClient( + "https://clavata.example", + api_key="test-key", + http_client=transport, + ).create_job("hello", "policy-id") + + assert len(transport.requests) == 3 diff --git a/tests/test_clavata_utils.py b/tests/test_clavata_utils.py deleted file mode 100644 index ad5342dbb7..0000000000 --- a/tests/test_clavata_utils.py +++ /dev/null @@ -1,221 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2023-2026 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 pytest - -from nemoguardrails.library.clavata.utils import ( - AttemptsExceededError, - calculate_exp_delay, - exponential_backoff, -) - - -@pytest.mark.unit -@pytest.mark.asyncio -async def test_exponential_backoff_successful_call(): - """Test that the exponential backoff decorator works with a successful call.""" - - call_count = 0 - - @exponential_backoff(max_attempts=3, initial_delay=0.01) - async def successful_function(): - nonlocal call_count - call_count += 1 - return "success" - - result = await successful_function() - assert result == "success" - assert call_count == 1 - - -@pytest.mark.unit -@pytest.mark.asyncio -async def test_exponential_backoff_retry_success(): - """Test that the exponential backoff retries and eventually succeeds.""" - - call_count = 0 - - @exponential_backoff(max_attempts=3, initial_delay=0.01) - async def eventually_successful(): - nonlocal call_count - call_count += 1 - if call_count < 3: - raise ValueError("Temporary failure") - return "success after retries" - - result = await eventually_successful() - assert result == "success after retries" - assert call_count == 3 - - -@pytest.mark.unit -@pytest.mark.asyncio -async def test_exponential_backoff_max_retries_exceeded(): - """Test that RetriesExceededError is raised when max_retries is exceeded.""" - - call_count = 0 - - @exponential_backoff(max_attempts=2, initial_delay=0.01) - async def always_fails(): - nonlocal call_count - call_count += 1 - raise ValueError("Failure") - - with pytest.raises(AttemptsExceededError) as excinfo: - await always_fails() - - assert call_count == 2 # initial call + 1 retry - assert excinfo.value.attempts == 2 - assert excinfo.value.max_attempts == 2 - assert isinstance(excinfo.value.last_exception, ValueError) - - -@pytest.mark.unit -@pytest.mark.asyncio -@pytest.mark.parametrize( - "retry_exceptions", - [ - ValueError, - (ValueError,), - (ValueError, TypeError), - ], - ids=["single_exception", "single_exception_tuple", "multiple_exceptions"], -) -async def test_exponential_backoff_specific_exceptions(retry_exceptions): - """Test that only specified exceptions trigger retry logic.""" - - call_count = 0 - - @exponential_backoff( - max_attempts=3, - initial_delay=0.01, - retry_exceptions=retry_exceptions, - ) - async def raises_different_exception(): - nonlocal call_count - call_count += 1 - raise KeyError("This exception shouldn't trigger retries") - - with pytest.raises(KeyError): - await raises_different_exception() - - assert call_count == 1 # No retries since KeyError is not in retry_exceptions - - -@pytest.mark.unit -@pytest.mark.asyncio -async def test_exponential_backoff_permanent_failure_handler(): - """Test that on_permanent_failure handler is called when retries are exhausted.""" - - call_count = 0 - handler_called = False - - async def failure_handler(attempts, exception): - nonlocal handler_called - handler_called = True - assert attempts == 2 - assert isinstance(exception, AttemptsExceededError) - return "handled failure" - - @exponential_backoff( - max_attempts=2, - initial_delay=0.01, - on_permanent_failure=failure_handler, - ) - async def always_fails(): - nonlocal call_count - call_count += 1 - raise ValueError("Failure") - - result = await always_fails() - assert result == "handled failure" - assert call_count == 2 # initial call + 1 retry - assert handler_called is True - - -@pytest.mark.unit -@pytest.mark.asyncio -async def test_exponential_backoff_permanent_failure_custom_exception(): - """Test that on_permanent_failure can return a custom exception to be raised.""" - - class CustomError(Exception): - pass - - async def failure_handler(attempts, exception): - return CustomError("Custom failure") - - @exponential_backoff( - max_attempts=2, - initial_delay=0.01, - on_permanent_failure=failure_handler, - ) - async def always_fails(): - raise ValueError("Failure") - - with pytest.raises(CustomError) as excinfo: - await always_fails() - - assert str(excinfo.value) == "Custom failure" - - -@pytest.mark.unit -@pytest.mark.parametrize( - ["retries", "expected_delay", "initial_delay", "max_delay", "jitter"], - [ - (0, 1, 1, 10, False), # 1 * 1^0 = 1 - (1, 2, 1, 10, False), # 1 * 1^1 = 2 - (2, 4, 1, 10, False), # 1 * 1^2 = 4 - (3, 8, 1, 10, False), # 1 * 1^3 = 8 - (4, 10, 1, 10, False), # 1 * 1^4 = 10 - ], - ids=[ - "first_attempt", - "second_attempt", - "third_attempt", - "fourth_attempt", - "fifth_attempt", - ], -) -def test_calculate_exp_delay(retries, expected_delay, initial_delay, max_delay, jitter): - """Test that the calculate_exp_delay function works correctly.""" - - assert calculate_exp_delay(retries, initial_delay, max_delay, jitter) == expected_delay - - -@pytest.mark.unit -@pytest.mark.parametrize( - ["retries", "expected_delay", "initial_delay", "max_delay"], - [ - (0, 1, 1, 10), - (1, 2, 1, 10), - (2, 4, 1, 10), - (3, 8, 1, 10), - (4, 10, 1, 10), - ], - ids=[ - "first_attempt", - "second_attempt", - "third_attempt", - "fourth_attempt", - "fifth_attempt", - ], -) -def test_calculate_exp_delay_jitter(retries, expected_delay, initial_delay, max_delay): - """Test that the calculate_exp_delay function works correctly with jitter.""" - - assert 0.0 <= calculate_exp_delay(retries, initial_delay, max_delay, True) <= expected_delay - - -# TESTS FOR ADDITIONAL PYDANTIC MODELS USED TO PARSE RESPONSES FROM THE CLAVATA API diff --git a/tests/test_crowdstrike_aidr_guard.py b/tests/test_crowdstrike_aidr_guard.py index d6e50066fa..574fc35036 100644 --- a/tests/test_crowdstrike_aidr_guard.py +++ b/tests/test_crowdstrike_aidr_guard.py @@ -18,10 +18,13 @@ from nemoguardrails import RailsConfig from nemoguardrails.actions.rail_outcome import RailOutcome, TransformTarget +from nemoguardrails.http import HTTPConnectionError, HTTPResponse from nemoguardrails.library.crowdstrike_aidr.actions import ( GuardChatCompletionsResult, _crowdstrike_aidr_outcome, + crowdstrike_aidr_guard, ) +from nemoguardrails.testing import RecordingHTTPClient from tests.utils import TestChat input_rail_config = RailsConfig.from_content( @@ -235,6 +238,39 @@ def test_crowdstrike_aidr_guard_error(httpx_mock: HTTPXMock, monkeypatch: pytest chat << "Hello!" +@pytest.mark.unit +@pytest.mark.asyncio +@pytest.mark.parametrize( + "response", + [ + HTTPConnectionError("connection failed"), + HTTPResponse(status_code=200, content=b"not-json"), + HTTPResponse(status_code=200, content=b"{}"), + ], +) +async def test_crowdstrike_aidr_guard_expected_failures_return_allow(monkeypatch, response): + monkeypatch.setenv("CS_AIDR_TOKEN", "test-token") + + outcome = await crowdstrike_aidr_guard( + mode="output", + config=output_rail_config, + context={"user_message": "Hi", "bot_message": "Hello"}, + http_client=RecordingHTTPClient([response]), + ) + + assert outcome.is_blocked is False + assert outcome.is_transform is False + assert outcome.metadata["blocked"] is False + assert outcome.metadata["transformed"] is False + assert outcome.metadata["user_message"] == "Hi" + assert outcome.metadata["bot_message"] == "Hello" + messages = outcome.metadata["guard_output"]["messages"] + assert [(message.role, message.content) for message in messages[-2:]] == [ + ("user", "Hi"), + ("assistant", "Hello"), + ] + + @pytest.mark.unit def test_crowdstrike_aidr_guard_missing_env_var(monkeypatch: pytest.MonkeyPatch): monkeypatch.delenv("CS_AIDR_TOKEN", raising=False) diff --git a/tests/test_f5_guardrails.py b/tests/test_f5_guardrails.py index 19d44ff440..89c9be5484 100644 --- a/tests/test_f5_guardrails.py +++ b/tests/test_f5_guardrails.py @@ -18,11 +18,14 @@ from unittest.mock import AsyncMock, patch import pytest -from aioresponses import aioresponses +import nemoguardrails.library.f5.actions as f5_actions from nemoguardrails import RailsConfig from nemoguardrails.actions.rail_outcome import RailOutcome +from nemoguardrails.http import HTTPConnectionError, HTTPResponse from nemoguardrails.library.f5.actions import f5_guardrails_scan +from nemoguardrails.testing import RecordingHTTPClient +from tests.http_utils import RecordedHTTPResponses from tests.utils import TestChat @@ -76,13 +79,14 @@ def test_f5_guardrails_input_cleared(config, monkeypatch): ], ) - with aioresponses() as m: + with RecordedHTTPResponses() as m: m.post( "https://us1.calypsoai.app/backend/v1/scans", payload={"result": {"outcome": "cleared"}}, - repeat=True, + times=3, ) + chat.app.register_action_param("http_client", m.client) chat >> "Hello!" chat << "express greeting" chat << "Hello! How can I assist you today?" @@ -97,13 +101,13 @@ def test_f5_guardrails_input_blocked(config, monkeypatch): ], ) - with aioresponses() as m: + with RecordedHTTPResponses() as m: m.post( "https://us1.calypsoai.app/backend/v1/scans", payload={"result": {"outcome": "flagged"}}, - repeat=True, ) + chat.app.register_action_param("http_client", m.client) chat >> "bad message" chat << "I'm sorry, I can't respond to that." @@ -117,7 +121,7 @@ def test_f5_guardrails_output_blocked(config, monkeypatch): ], ) - with aioresponses() as m: + with RecordedHTTPResponses() as m: # Input scan cleared m.post( "https://us1.calypsoai.app/backend/v1/scans", @@ -129,6 +133,7 @@ def test_f5_guardrails_output_blocked(config, monkeypatch): payload={"result": {"outcome": "flagged"}}, ) + chat.app.register_action_param("http_client", m.client) chat >> "Hello" chat << "I'm sorry, I can't respond to that." @@ -143,14 +148,15 @@ def test_f5_guardrails_fail_open(config_fail_open, monkeypatch): ], ) - with aioresponses() as m: + with RecordedHTTPResponses() as m: # Simulate an API error (e.g., 500) m.post( "https://us1.calypsoai.app/backend/v1/scans", status=500, - repeat=True, + times=3, ) + chat.app.register_action_param("http_client", m.client) chat >> "Hello!" chat << "express greeting" chat << "Hello! How can I assist you today?" @@ -160,13 +166,13 @@ def test_f5_guardrails_fail_closed(config, monkeypatch): monkeypatch.setenv("F5_GUARDRAILS_API_KEY", "test-key") chat = TestChat(config) - with aioresponses() as m: + with RecordedHTTPResponses() as m: m.post( "https://us1.calypsoai.app/backend/v1/scans", status=500, - repeat=True, ) + chat.app.register_action_param("http_client", m.client) chat >> "Hello!" chat << "I'm sorry, an internal error has occurred." @@ -175,29 +181,69 @@ def test_f5_guardrails_fail_closed(config, monkeypatch): async def test_f5_guardrails_timeout_fail_open(config_fail_open, monkeypatch): monkeypatch.setenv("F5_GUARDRAILS_API_KEY", "test-key") - with aioresponses() as m: + with RecordedHTTPResponses() as m: m.post( "https://us1.calypsoai.app/backend/v1/scans", exception=asyncio.TimeoutError(), ) - result = await f5_guardrails_scan(text="Hello!", config=config_fail_open) + result = await f5_guardrails_scan(text="Hello!", config=config_fail_open, http_client=m.client) assert result == RailOutcome.allow(metadata={"result": {"outcome": "cleared"}, "fail_open": True}) +@pytest.mark.asyncio +async def test_f5_guardrails_connection_failure_respects_policy(config, config_fail_open, monkeypatch): + monkeypatch.setenv("F5_GUARDRAILS_API_KEY", "test-key") + + fail_open_result = await f5_guardrails_scan( + text="Hello!", + config=config_fail_open, + http_client=RecordingHTTPClient([HTTPConnectionError("connection failed")]), + ) + + assert fail_open_result == RailOutcome.allow(metadata={"result": {"outcome": "cleared"}, "fail_open": True}) + + with pytest.raises(RuntimeError, match="Connection error to F5 Guardrails API"): + await f5_guardrails_scan( + text="Hello!", + config=config, + http_client=RecordingHTTPClient([HTTPConnectionError("connection failed")]), + ) + + +@pytest.mark.asyncio +async def test_f5_guardrails_closes_owned_shared_client(config, monkeypatch): + monkeypatch.setenv("F5_GUARDRAILS_API_KEY", "test-key") + transport = RecordingHTTPClient( + [ + HTTPResponse( + status_code=200, + content=b'{"result":{"outcome":"cleared"}}', + ) + ] + ) + monkeypatch.setattr(f5_actions, "create_http_client", lambda **kwargs: transport) + + result = await f5_guardrails_scan(text="Hello!", config=config) + + assert result == RailOutcome.allow(metadata={"result": {"outcome": "cleared"}}) + assert transport.close_calls == 1 + assert transport.requests[0].timeout == 30.0 + + @pytest.mark.asyncio async def test_f5_guardrails_fail_open_marker_on_http_error(config_fail_open, monkeypatch): monkeypatch.setenv("F5_GUARDRAILS_API_KEY", "test-key") - with aioresponses() as m: + with RecordedHTTPResponses() as m: m.post( "https://us1.calypsoai.app/backend/v1/scans", status=500, body="upstream failure", ) - result = await f5_guardrails_scan(text="Hello!", config=config_fail_open) + result = await f5_guardrails_scan(text="Hello!", config=config_fail_open, http_client=m.client) assert result == RailOutcome.allow(metadata={"result": {"outcome": "cleared"}, "fail_open": True}) @@ -215,7 +261,7 @@ async def test_f5_guardrails_error_body_not_logged(config_fail_open, monkeypatch sentinel = "SENSITIVE-USER-INPUT-DO-NOT-LOG-12345" caplog.set_level(logging.DEBUG, logger="nemoguardrails.library.f5.actions") - with aioresponses() as m: + with RecordedHTTPResponses() as m: m.post( "https://us1.calypsoai.app/backend/v1/scans", status=500, @@ -223,7 +269,7 @@ async def test_f5_guardrails_error_body_not_logged(config_fail_open, monkeypatch content_type="application/json", ) - await f5_guardrails_scan(text=sentinel, config=config_fail_open) + await f5_guardrails_scan(text=sentinel, config=config_fail_open, http_client=m.client) for record in caplog.records: assert sentinel not in record.getMessage(), f"Vendor error body leaked into log record: {record.getMessage()!r}" @@ -260,13 +306,14 @@ def test_f5_guardrails_colang_2_input_blocked(config_v2, monkeypatch): monkeypatch.setenv("F5_GUARDRAILS_API_KEY", "test-key") chat = TestChat(config_v2) - with aioresponses() as m: + with RecordedHTTPResponses() as m: m.post( "https://us1.calypsoai.app/backend/v1/scans", payload={"result": {"outcome": "flagged"}}, - repeat=True, + times=2, ) + chat.app.register_action_param("http_client", m.client) chat >> "bad message" chat << "I'm sorry, I can't respond to that." @@ -275,13 +322,14 @@ def test_f5_guardrails_colang_2_input_cleared(config_v2, monkeypatch): monkeypatch.setenv("F5_GUARDRAILS_API_KEY", "test-key") chat = TestChat(config_v2) - with aioresponses() as m: + with RecordedHTTPResponses() as m: m.post( "https://us1.calypsoai.app/backend/v1/scans", payload={"result": {"outcome": "cleared"}}, - repeat=True, + times=2, ) + chat.app.register_action_param("http_client", m.client) chat >> "Hello!" chat << "Hello! How can I assist you today?" @@ -290,7 +338,7 @@ def test_f5_guardrails_colang_2_output_blocked(config_v2, monkeypatch): monkeypatch.setenv("F5_GUARDRAILS_API_KEY", "test-key") chat = TestChat(config_v2) - with aioresponses() as m: + with RecordedHTTPResponses() as m: # Input scan cleared m.post( "https://us1.calypsoai.app/backend/v1/scans", @@ -302,6 +350,7 @@ def test_f5_guardrails_colang_2_output_blocked(config_v2, monkeypatch): payload={"result": {"outcome": "flagged"}}, ) + chat.app.register_action_param("http_client", m.client) chat >> "Hello!" chat << "I'm sorry, I can't respond to that." @@ -310,14 +359,14 @@ def test_f5_guardrails_colang_2_output_blocked(config_v2, monkeypatch): async def test_f5_guardrails_timeout_fail_closed(config, monkeypatch): monkeypatch.setenv("F5_GUARDRAILS_API_KEY", "test-key") - with aioresponses() as m: + with RecordedHTTPResponses() as m: m.post( "https://us1.calypsoai.app/backend/v1/scans", exception=asyncio.TimeoutError(), ) with pytest.raises(RuntimeError, match="timed out"): - await f5_guardrails_scan(text="Hello!", config=config) + await f5_guardrails_scan(text="Hello!", config=config, http_client=m.client) @pytest.mark.asyncio @@ -335,15 +384,24 @@ async def test_f5_guardrails_custom_api_url(monkeypatch): """, ) - with aioresponses() as m: + with RecordedHTTPResponses() as m: m.post( "https://custom.example.com/backend/v1/scans", payload={"result": {"outcome": "cleared"}}, ) - result = await f5_guardrails_scan(text="Hello!", config=custom_config) + result = await f5_guardrails_scan(text="Hello!", config=custom_config, http_client=m.client) assert result == RailOutcome.allow(metadata={"result": {"outcome": "cleared"}}) + request = m.client.requests[0] + assert request.method == "POST" + assert request.url == "https://custom.example.com/backend/v1/scans" + assert request.headers == { + "Authorization": "Bearer test-key", + "Content-Type": "application/json", + } + assert request.json == {"input": "Hello!"} + assert request.timeout == 30.0 @pytest.mark.asyncio @@ -352,15 +410,16 @@ async def test_f5_guardrails_api_url_env_fallback(config, monkeypatch): monkeypatch.setenv("F5_GUARDRAILS_API_KEY", "test-key") monkeypatch.setenv("F5_GUARDRAILS_API_URL", "https://env.example.com") - with aioresponses() as m: + with RecordedHTTPResponses() as m: m.post( "https://env.example.com/backend/v1/scans", payload={"result": {"outcome": "cleared"}}, ) - result = await f5_guardrails_scan(text="Hello!", config=config) + result = await f5_guardrails_scan(text="Hello!", config=config, http_client=m.client) assert result == RailOutcome.allow(metadata={"result": {"outcome": "cleared"}}) + assert m.client.requests[0].url == "https://env.example.com/backend/v1/scans" @pytest.mark.asyncio @@ -379,15 +438,16 @@ async def test_f5_guardrails_env_api_url_wins_over_config(monkeypatch): """, ) - with aioresponses() as m: + with RecordedHTTPResponses() as m: m.post( "https://beta.example.com/backend/v1/scans", payload={"result": {"outcome": "cleared"}}, ) - result = await f5_guardrails_scan(text="Hello!", config=custom_config) + result = await f5_guardrails_scan(text="Hello!", config=custom_config, http_client=m.client) assert result == RailOutcome.allow(metadata={"result": {"outcome": "cleared"}}) + assert m.client.requests[0].url == "https://beta.example.com/backend/v1/scans" @pytest.fixture @@ -417,13 +477,13 @@ async def test_f5_guardrails_input_rails_exception(config_exceptions, monkeypatc ], ) - with aioresponses() as m: + with RecordedHTTPResponses() as m: m.post( "https://us1.calypsoai.app/backend/v1/scans", payload={"result": {"outcome": "flagged"}}, - repeat=True, ) + chat.app.register_action_param("http_client", m.client) messages = [{"role": "user", "content": "bad message"}] result = await chat.app.generate_async(messages=messages) @@ -443,7 +503,7 @@ async def test_f5_guardrails_output_rails_exception(config_exceptions, monkeypat ], ) - with aioresponses() as m: + with RecordedHTTPResponses() as m: m.post( "https://us1.calypsoai.app/backend/v1/scans", payload={"result": {"outcome": "cleared"}}, @@ -453,6 +513,7 @@ async def test_f5_guardrails_output_rails_exception(config_exceptions, monkeypat payload={"result": {"outcome": "flagged"}}, ) + chat.app.register_action_param("http_client", m.client) messages = [{"role": "user", "content": "Hello"}] result = await chat.app.generate_async(messages=messages) @@ -498,7 +559,7 @@ async def test_f5_guardrails_429_then_success(config_no_backoff, monkeypatch): """A 429 with Retry-After: 0 is retried and the second call succeeds.""" monkeypatch.setenv("F5_GUARDRAILS_API_KEY", "test-key") - with aioresponses() as m: + with RecordedHTTPResponses() as m: m.post( "https://us1.calypsoai.app/backend/v1/scans", status=429, @@ -513,7 +574,7 @@ async def test_f5_guardrails_429_then_success(config_no_backoff, monkeypatch): "nemoguardrails.library.f5.actions.asyncio.sleep", new=AsyncMock(return_value=None), ) as sleep_mock: - result = await f5_guardrails_scan(text="Hello!", config=config_no_backoff) + result = await f5_guardrails_scan(text="Hello!", config=config_no_backoff, http_client=m.client) assert result == RailOutcome.allow(metadata={"result": {"outcome": "cleared"}}) sleep_mock.assert_awaited_once() @@ -525,7 +586,7 @@ async def test_f5_guardrails_429_retry_after_http_date(config_no_backoff, monkey """HTTP-date Retry-After values are parsed and honored.""" monkeypatch.setenv("F5_GUARDRAILS_API_KEY", "test-key") - with aioresponses() as m: + with RecordedHTTPResponses() as m: m.post( "https://us1.calypsoai.app/backend/v1/scans", status=429, @@ -540,7 +601,7 @@ async def test_f5_guardrails_429_retry_after_http_date(config_no_backoff, monkey "nemoguardrails.library.f5.actions.asyncio.sleep", new=AsyncMock(return_value=None), ) as sleep_mock: - result = await f5_guardrails_scan(text="Hello!", config=config_no_backoff) + result = await f5_guardrails_scan(text="Hello!", config=config_no_backoff, http_client=m.client) assert result == RailOutcome.allow(metadata={"result": {"outcome": "cleared"}}) sleep_mock.assert_awaited_once() @@ -552,7 +613,7 @@ async def test_f5_guardrails_429_retry_after_capped(config_no_backoff, monkeypat """Retry-After values larger than max_retry_after_seconds are clamped.""" monkeypatch.setenv("F5_GUARDRAILS_API_KEY", "test-key") - with aioresponses() as m: + with RecordedHTTPResponses() as m: m.post( "https://us1.calypsoai.app/backend/v1/scans", status=429, @@ -567,7 +628,7 @@ async def test_f5_guardrails_429_retry_after_capped(config_no_backoff, monkeypat "nemoguardrails.library.f5.actions.asyncio.sleep", new=AsyncMock(return_value=None), ) as sleep_mock: - result = await f5_guardrails_scan(text="Hello!", config=config_no_backoff) + result = await f5_guardrails_scan(text="Hello!", config=config_no_backoff, http_client=m.client) assert result == RailOutcome.allow(metadata={"result": {"outcome": "cleared"}}) sleep_mock.assert_awaited_once() @@ -578,19 +639,19 @@ async def test_f5_guardrails_429_retry_after_capped(config_no_backoff, monkeypat async def test_f5_guardrails_429_exhausted_fail_open(config_no_backoff_fail_open, monkeypatch): monkeypatch.setenv("F5_GUARDRAILS_API_KEY", "test-key") - with aioresponses() as m: + with RecordedHTTPResponses() as m: m.post( "https://us1.calypsoai.app/backend/v1/scans", status=429, headers={"Retry-After": "0"}, - repeat=True, + times=3, ) with patch( "nemoguardrails.library.f5.actions.asyncio.sleep", new=AsyncMock(return_value=None), ): - result = await f5_guardrails_scan(text="Hello!", config=config_no_backoff_fail_open) + result = await f5_guardrails_scan(text="Hello!", config=config_no_backoff_fail_open, http_client=m.client) assert result == RailOutcome.allow(metadata={"result": {"outcome": "cleared"}, "fail_open": True}) @@ -599,12 +660,12 @@ async def test_f5_guardrails_429_exhausted_fail_open(config_no_backoff_fail_open async def test_f5_guardrails_429_exhausted_fail_closed(config_no_backoff, monkeypatch): monkeypatch.setenv("F5_GUARDRAILS_API_KEY", "test-key") - with aioresponses() as m: + with RecordedHTTPResponses() as m: m.post( "https://us1.calypsoai.app/backend/v1/scans", status=429, headers={"Retry-After": "0"}, - repeat=True, + times=3, ) with patch( @@ -612,7 +673,7 @@ async def test_f5_guardrails_429_exhausted_fail_closed(config_no_backoff, monkey new=AsyncMock(return_value=None), ): with pytest.raises(RuntimeError, match="rate limited"): - await f5_guardrails_scan(text="Hello!", config=config_no_backoff) + await f5_guardrails_scan(text="Hello!", config=config_no_backoff, http_client=m.client) @pytest.mark.asyncio @@ -632,7 +693,7 @@ async def test_f5_guardrails_429_no_retry_after_uses_backoff(monkeypatch): """, ) - with aioresponses() as m: + with RecordedHTTPResponses() as m: m.post( "https://us1.calypsoai.app/backend/v1/scans", status=429, @@ -650,7 +711,7 @@ async def test_f5_guardrails_429_no_retry_after_uses_backoff(monkeypatch): "nemoguardrails.library.f5.actions.asyncio.sleep", new=AsyncMock(return_value=None), ) as sleep_mock: - result = await f5_guardrails_scan(text="Hello!", config=cfg) + result = await f5_guardrails_scan(text="Hello!", config=cfg, http_client=m.client) assert result == RailOutcome.allow(metadata={"result": {"outcome": "cleared"}}) assert sleep_mock.await_count == 2 diff --git a/tests/test_fiddler_rails.py b/tests/test_fiddler_rails.py index 5b8b9c2718..0c42c02bc8 100644 --- a/tests/test_fiddler_rails.py +++ b/tests/test_fiddler_rails.py @@ -12,19 +12,57 @@ # 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 json import os -from unittest.mock import AsyncMock, MagicMock, patch import pytest -from aioresponses import aioresponses from nemoguardrails import RailsConfig from nemoguardrails.actions.actions import ActionResult, action +from nemoguardrails.http import HTTPConnectionError, HTTPResponse +from nemoguardrails.library.fiddler.actions import call_fiddler_guardrail +from nemoguardrails.testing import RecordingHTTPClient from tests.utils import TestChat CONFIGS_FOLDER = os.path.join(os.path.dirname(__file__), ".", "test_configs") +def _response(payload: dict) -> HTTPResponse: + return HTTPResponse(status_code=200, content=json.dumps(payload).encode()) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("response", "log_message"), + [ + ( + HTTPResponse(status_code=503, content=b"unavailable"), + "Fiddler Test could not be run. Fiddler API returned status code 503", + ), + ( + HTTPConnectionError("connection failed"), + "Fiddler Test request failed: connection failed", + ), + ], +) +async def test_fiddler_guardrail_preserves_http_failure_contract(monkeypatch, caplog, response, log_message): + monkeypatch.setenv("FIDDLER_API_KEY", "test-key") + + result = await call_fiddler_guardrail( + endpoint="https://fiddler.example/guardrail", + data={"input": "text"}, + guardrail_name="Fiddler Test", + score_key="score", + threshold=0.5, + compare=lambda score, threshold: score >= threshold, + default_score=0, + http_client=RecordingHTTPClient([response]), + ) + + assert result is False + assert log_message in caplog.text + + @action(is_system_action=True) async def retrieve_relevant_chunks(): """Retrieve relevant chunks from the knowledge base and add them to the context.""" @@ -49,26 +87,28 @@ async def test_fiddler_safety_rails(monkeypatch): "Yes, shipping can be done in 2 days.", # bot response that will be intercepted ], ) - with aioresponses() as m: - m.post( - "https://testfiddler.ai/v3/guardrails/ftl-safety", - payload={ - "fdl_harmful": 0.5, - "fdl_violent": 0.5, - "fdl_unethical": 0.5, - "fdl_illegal": 0.5, - "fdl_sexual": 0.5, - "fdl_racist": 0.5, - "fdl_jailbreaking": 0.5, - "fdl_harassing": 0.5, - "fdl_hateful": 0.5, - "fdl_sexist": 0.5, - "fdl_roleplaying": 0.5, - }, - ) - - chat >> "Do you ship within 2 days?" - await chat.bot_async("I'm sorry, I can't respond to that.") + client = RecordingHTTPClient( + [ + _response( + { + "fdl_harmful": 0.5, + "fdl_violent": 0.5, + "fdl_unethical": 0.5, + "fdl_illegal": 0.5, + "fdl_sexual": 0.5, + "fdl_racist": 0.5, + "fdl_jailbreaking": 0.5, + "fdl_harassing": 0.5, + "fdl_hateful": 0.5, + "fdl_sexist": 0.5, + "fdl_roleplaying": 0.5, + } + ) + ] + ) + chat.app.register_action_param("http_client", client) + chat >> "Do you ship within 2 days?" + await chat.bot_async("I'm sorry, I can't respond to that.") @pytest.mark.asyncio @@ -87,25 +127,43 @@ async def test_fiddler_safety_rails_pass(monkeypatch): ], ) - with aioresponses() as m: - m.post( - "https://testfiddler.ai/v3/guardrails/ftl-safety", - payload={ - "fdl_harmful": 0.02, - "fdl_violent": 0.02, - "fdl_unethical": 0.02, - "fdl_illegal": 0.02, - "fdl_sexual": 0.02, - "fdl_racist": 0.02, - "fdl_jailbreaking": 0.02, - "fdl_harassing": 0.02, - "fdl_hateful": 0.02, - "fdl_sexist": 0.02, - "fdl_roleplaying": 0.02, - }, - ) - chat >> "Do you ship within 2 days?" - await chat.bot_async("Yes, shipping can be done in 2 days.") + client = RecordingHTTPClient( + [ + _response( + { + "fdl_harmful": 0.02, + "fdl_violent": 0.02, + "fdl_unethical": 0.02, + "fdl_illegal": 0.02, + "fdl_sexual": 0.02, + "fdl_racist": 0.02, + "fdl_jailbreaking": 0.02, + "fdl_harassing": 0.02, + "fdl_hateful": 0.02, + "fdl_sexist": 0.02, + "fdl_roleplaying": 0.02, + } + ), + _response( + { + "fdl_harmful": 0.02, + "fdl_violent": 0.02, + "fdl_unethical": 0.02, + "fdl_illegal": 0.02, + "fdl_sexual": 0.02, + "fdl_racist": 0.02, + "fdl_jailbreaking": 0.02, + "fdl_harassing": 0.02, + "fdl_hateful": 0.02, + "fdl_sexist": 0.02, + "fdl_roleplaying": 0.02, + } + ), + ] + ) + chat.app.register_action_param("http_client", client) + chat >> "Do you ship within 2 days?" + await chat.bot_async("Yes, shipping can be done in 2 days.") @pytest.mark.asyncio @@ -122,26 +180,29 @@ async def test_fiddler_thresholds(monkeypatch): ], ) - with aioresponses() as m: - chat.app.register_action(retrieve_relevant_chunks, "retrieve_relevant_chunks") - m.post( - "https://testfiddler.ai/v3/guardrails/ftl-safety", - payload={ - "fdl_harmful": 0.5, - "fdl_violent": 0.5, - "fdl_unethical": 0.5, - "fdl_illegal": 0.5, - "fdl_sexual": 0.5, - "fdl_racist": 0.5, - "fdl_jailbreaking": 0.5, - "fdl_harassing": 0.5, - "fdl_hateful": 0.5, - "fdl_sexist": 0.5, - "fdl_roleplaying": 0.5, - }, - ) - chat >> "Do you ship within 2 days?" - await chat.bot_async("I'm sorry, I can't respond to that.") + client = RecordingHTTPClient( + [ + _response( + { + "fdl_harmful": 0.5, + "fdl_violent": 0.5, + "fdl_unethical": 0.5, + "fdl_illegal": 0.5, + "fdl_sexual": 0.5, + "fdl_racist": 0.5, + "fdl_jailbreaking": 0.5, + "fdl_harassing": 0.5, + "fdl_hateful": 0.5, + "fdl_sexist": 0.5, + "fdl_roleplaying": 0.5, + } + ) + ] + ) + chat.app.register_action(retrieve_relevant_chunks, "retrieve_relevant_chunks") + chat.app.register_action_param("http_client", client) + chat >> "Do you ship within 2 days?" + await chat.bot_async("I'm sorry, I can't respond to that.") @pytest.mark.asyncio @@ -158,14 +219,11 @@ async def test_fiddler_faithfulness_rails(monkeypatch): ], ) - with aioresponses() as m: - chat.app.register_action(retrieve_relevant_chunks, "retrieve_relevant_chunks") - m.post( - "https://testfiddler.ai/v3/guardrails/ftl-response-faithfulness", - payload={"fdl_faithful_score": 0.001}, - ) - chat >> "Do you ship within 2 days?" - await chat.bot_async("I'm sorry, I can't respond to that.") + client = RecordingHTTPClient([_response({"fdl_faithful_score": 0.001})]) + chat.app.register_action(retrieve_relevant_chunks, "retrieve_relevant_chunks") + chat.app.register_action_param("http_client", client) + chat >> "Do you ship within 2 days?" + await chat.bot_async("I'm sorry, I can't respond to that.") @pytest.mark.asyncio @@ -182,14 +240,11 @@ async def test_fiddler_faithfulness_rails_pass(monkeypatch): ], ) - with aioresponses() as m: - chat.app.register_action(retrieve_relevant_chunks, "retrieve_relevant_chunks") - m.post( - "https://testfiddler.ai/v3/guardrails/ftl-response-faithfulness", - payload={"fdl_faithful_score": 0.5}, - ) - chat >> "Do you ship within 2 days?" - await chat.bot_async("Yes, shipping can be done in 2 days.") + client = RecordingHTTPClient([_response({"fdl_faithful_score": 0.5})]) + chat.app.register_action(retrieve_relevant_chunks, "retrieve_relevant_chunks") + chat.app.register_action_param("http_client", client) + chat >> "Do you ship within 2 days?" + await chat.bot_async("Yes, shipping can be done in 2 days.") @pytest.mark.unit @@ -210,44 +265,37 @@ async def test_fiddler_safety_request_format_user_message(monkeypatch): """ ) - mock_session = AsyncMock() - mock_response = AsyncMock() - mock_response.status = 200 - mock_response.json = AsyncMock( - return_value={ - "fdl_harmful": 0.1, - "fdl_violent": 0.1, - "fdl_unethical": 0.1, - "fdl_illegal": 0.1, - "fdl_sexual": 0.1, - "fdl_racist": 0.1, - "fdl_jailbreaking": 0.1, - "fdl_harassing": 0.1, - "fdl_hateful": 0.1, - "fdl_sexist": 0.1, - "fdl_roleplaying": 0.1, - } + client = RecordingHTTPClient( + [ + _response( + { + "fdl_harmful": 0.1, + "fdl_violent": 0.1, + "fdl_unethical": 0.1, + "fdl_illegal": 0.1, + "fdl_sexual": 0.1, + "fdl_racist": 0.1, + "fdl_jailbreaking": 0.1, + "fdl_harassing": 0.1, + "fdl_hateful": 0.1, + "fdl_sexist": 0.1, + "fdl_roleplaying": 0.1, + } + ) + ] ) - mock_response.__aenter__ = AsyncMock(return_value=mock_response) - mock_response.__aexit__ = AsyncMock() - - mock_post_context = AsyncMock() - mock_post_context.__aenter__ = AsyncMock(return_value=mock_response) - mock_post_context.__aexit__ = AsyncMock() - mock_session.post = MagicMock(return_value=mock_post_context) - - mock_client_session = MagicMock() - mock_client_session.return_value.__aenter__ = AsyncMock(return_value=mock_session) - mock_client_session.return_value.__aexit__ = AsyncMock() - - with patch("nemoguardrails.library.fiddler.actions.aiohttp.ClientSession", mock_client_session): - context = {"user_message": "Hello, how are you?"} - result = await call_fiddler_safety_user(config, context) - - # Verify request format - assert mock_session.post.called - call_args = mock_session.post.call_args - request_payload = call_args[1]["json"] + context = {"user_message": "Hello, how are you?"} + await call_fiddler_safety_user(config, context, http_client=client) + + request = client.requests[0] + assert request.method == "POST" + assert request.url == "https://testfiddler.ai/v3/guardrails/ftl-safety" + assert request.headers == { + "Content-Type": "application/json", + "Authorization": "Bearer test-key", + } + assert request.timeout is None + request_payload = request.json assert "data" in request_payload assert request_payload["data"]["input"] == "Hello, how are you?" assert "prompt" not in request_payload["data"] # Old format should not be present @@ -272,44 +320,29 @@ async def test_fiddler_safety_request_format_bot_message(monkeypatch): """ ) - mock_session = AsyncMock() - mock_response = AsyncMock() - mock_response.status = 200 - mock_response.json = AsyncMock( - return_value={ - "fdl_harmful": 0.1, - "fdl_violent": 0.1, - "fdl_unethical": 0.1, - "fdl_illegal": 0.1, - "fdl_sexual": 0.1, - "fdl_racist": 0.1, - "fdl_jailbreaking": 0.1, - "fdl_harassing": 0.1, - "fdl_hateful": 0.1, - "fdl_sexist": 0.1, - "fdl_roleplaying": 0.1, - } + client = RecordingHTTPClient( + [ + _response( + { + "fdl_harmful": 0.1, + "fdl_violent": 0.1, + "fdl_unethical": 0.1, + "fdl_illegal": 0.1, + "fdl_sexual": 0.1, + "fdl_racist": 0.1, + "fdl_jailbreaking": 0.1, + "fdl_harassing": 0.1, + "fdl_hateful": 0.1, + "fdl_sexist": 0.1, + "fdl_roleplaying": 0.1, + } + ) + ] ) - mock_response.__aenter__ = AsyncMock(return_value=mock_response) - mock_response.__aexit__ = AsyncMock() - - mock_post_context = AsyncMock() - mock_post_context.__aenter__ = AsyncMock(return_value=mock_response) - mock_post_context.__aexit__ = AsyncMock() - mock_session.post = MagicMock(return_value=mock_post_context) - - mock_client_session = MagicMock() - mock_client_session.return_value.__aenter__ = AsyncMock(return_value=mock_session) - mock_client_session.return_value.__aexit__ = AsyncMock() - - with patch("nemoguardrails.library.fiddler.actions.aiohttp.ClientSession", mock_client_session): - context = {"bot_message": "I can help you with that."} - result = await call_fiddler_safety_bot(config, context) - - # Verify request format - assert mock_session.post.called - call_args = mock_session.post.call_args - request_payload = call_args[1]["json"] + context = {"bot_message": "I can help you with that."} + await call_fiddler_safety_bot(config, context, http_client=client) + + request_payload = client.requests[0].json assert "data" in request_payload assert request_payload["data"]["input"] == "I can help you with that." assert "prompt" not in request_payload["data"] # Old format should not be present @@ -334,33 +367,14 @@ async def test_fiddler_faithfulness_request_format(monkeypatch): """ ) - mock_session = AsyncMock() - mock_response = AsyncMock() - mock_response.status = 200 - mock_response.json = AsyncMock(return_value={"fdl_faithful_score": 0.2}) - mock_response.__aenter__ = AsyncMock(return_value=mock_response) - mock_response.__aexit__ = AsyncMock() - - mock_post_context = AsyncMock() - mock_post_context.__aenter__ = AsyncMock(return_value=mock_response) - mock_post_context.__aexit__ = AsyncMock() - mock_session.post = MagicMock(return_value=mock_post_context) - - mock_client_session = MagicMock() - mock_client_session.return_value.__aenter__ = AsyncMock(return_value=mock_session) - mock_client_session.return_value.__aexit__ = AsyncMock() - - with patch("nemoguardrails.library.fiddler.actions.aiohttp.ClientSession", mock_client_session): - context = { - "bot_message": "Shipping takes 2 days", - "relevant_chunks": "Shipping takes at least 3 days. We ship worldwide.", - } - result = await call_fiddler_faithfulness(config, context) - - # Verify request format - assert mock_session.post.called - call_args = mock_session.post.call_args - request_payload = call_args[1]["json"] + client = RecordingHTTPClient([_response({"fdl_faithful_score": 0.2})]) + context = { + "bot_message": "Shipping takes 2 days", + "relevant_chunks": "Shipping takes at least 3 days. We ship worldwide.", + } + await call_fiddler_faithfulness(config, context, http_client=client) + + request_payload = client.requests[0].json assert "data" in request_payload assert request_payload["data"]["response"] == "Shipping takes 2 days" assert request_payload["data"]["context"] == "Shipping takes at least 3 days. We ship worldwide." diff --git a/tests/test_hf_classifier.py b/tests/test_hf_classifier.py index 771bab27b4..5ae75ffd49 100644 --- a/tests/test_hf_classifier.py +++ b/tests/test_hf_classifier.py @@ -30,6 +30,7 @@ from pytest_httpx import HTTPXMock from nemoguardrails.actions.rail_outcome import RailOutcome, TransformTarget +from nemoguardrails.http import HTTPConnectionError, HTTPResponse, HTTPTimeoutError from nemoguardrails.library.hf_classifier import backends as backends_mod from nemoguardrails.library.hf_classifier.actions import ( _classify_and_check, @@ -58,6 +59,7 @@ LocalHFClassifierConfig, RemoteHFClassifierConfig, ) +from nemoguardrails.testing import RecordingHTTPClient _REMOTE_DEFAULTS = dict( engine="vllm", @@ -190,15 +192,16 @@ def test_no_default_headers_support(self): class TestSSLTimeout: def test_timeout_default(self): - assert _get_timeout(_remote()) == httpx.Timeout(30.0) + assert _get_timeout(_remote()) == 30.0 def test_timeout_custom(self): - assert _get_timeout(_remote(parameters={"timeout": 10.0})) == httpx.Timeout(10.0) + assert _get_timeout(_remote(parameters={"timeout": 10.0})) == 10.0 def test_ssl_default(self): cfg = _build_ssl_config(_remote()) assert cfg.verify is True - assert cfg.cert is None + assert cfg.ca_bundle is None + assert cfg.client_certificate is None def test_ssl_disabled(self): cfg = _build_ssl_config(_remote(parameters={"verify_ssl": False})) @@ -287,6 +290,48 @@ async def test_missing_data_key(self, httpx_mock: HTTPXMock): with pytest.raises(ValueError, match="Unexpected vLLM"): await self._backend().classify("text") + @pytest.mark.asyncio + async def test_uses_injected_http_client(self): + client = RecordingHTTPClient( + [ + HTTPResponse( + status_code=200, + content=b'{"data":[{"label":"safe","probs":[0.95]}]}', + ) + ] + ) + backend = VLLMBackend( + _remote(engine="vllm", base_url="http://vllm:8000"), + http_client=client, + ) + + result = await backend.classify("text") + + assert result == [ClassificationResult(label="safe", score=0.95)] + assert client.requests[0].url == self._URL + assert client.close_calls == 0 + + @pytest.mark.asyncio + async def test_injected_http_client_preserves_transport_retry(self): + client = RecordingHTTPClient( + [ + HTTPConnectionError("offline"), + HTTPResponse( + status_code=200, + content=b'{"data":[{"label":"safe","probs":[0.95]}]}', + ), + ] + ) + backend = VLLMBackend( + _remote(engine="vllm", base_url="http://vllm:8000"), + http_client=client, + ) + + result = await backend.classify("text") + + assert result == [ClassificationResult(label="safe", score=0.95)] + assert len(client.requests) == 2 + class TestKServeBackend: _URL = "http://ks:8080/v1/models/m:predict" @@ -419,6 +464,14 @@ def test_unknown_raises(self): with pytest.raises(ValueError, match="Unknown hf_classifier engine"): get_backend(c) + def test_remote_backend_uses_injected_http_client(self): + client = RecordingHTTPClient() + + backend = get_backend(_remote(engine="vllm"), http_client=client) + + assert isinstance(backend, VLLMBackend) + assert backend._http_client is client + class TestClassifyAndCheck: @pytest.mark.asyncio @@ -564,13 +617,15 @@ async def test_dollar_classifier_without_model_name_raises(self): class TestSSLCerts: def test_ca_cert(self): cfg = _build_ssl_config(_remote(parameters={"ca_cert": "/ca.pem"})) - assert cfg.verify == "/ca.pem" - assert cfg.cert is None + assert cfg.verify is True + assert cfg.ca_bundle == "/ca.pem" + assert cfg.client_certificate is None def test_mtls(self): cfg = _build_ssl_config(_remote(parameters={"client_cert": "/cert.pem", "client_key": "/key.pem"})) assert cfg.verify is True - assert cfg.cert == ("/cert.pem", "/key.pem") + assert cfg.client_certificate == "/cert.pem" + assert cfg.client_key == "/key.pem" def test_ca_cert_plus_mtls(self): cfg = _build_ssl_config( @@ -582,8 +637,10 @@ def test_ca_cert_plus_mtls(self): } ) ) - assert cfg.verify == "/ca.pem" - assert cfg.cert == ("/cert.pem", "/key.pem") + assert cfg.verify is True + assert cfg.ca_bundle == "/ca.pem" + assert cfg.client_certificate == "/cert.pem" + assert cfg.client_key == "/key.pem" def test_client_cert_without_key_raises(self): with pytest.raises(ValueError, match="client_key"): @@ -758,7 +815,7 @@ async def test_vllm_connection_error(self, httpx_mock: HTTPXMock): backend = VLLMBackend(_remote(engine="vllm", base_url="http://vllm:8000")) httpx_mock.add_exception(httpx.ConnectError("Connection refused"), url="http://vllm:8000/classify") httpx_mock.add_exception(httpx.ConnectError("Connection refused"), url="http://vllm:8000/classify") - with pytest.raises(httpx.ConnectError): + with pytest.raises(HTTPConnectionError): await backend.classify("text") @pytest.mark.asyncio @@ -766,7 +823,7 @@ async def test_fms_connection_error(self, httpx_mock: HTTPXMock): backend = FMSBackend(_remote(engine="fms", base_url="http://fms:9000")) httpx_mock.add_exception(httpx.ConnectError("Connection refused"), url="http://fms:9000/api/v1/text/contents") httpx_mock.add_exception(httpx.ConnectError("Connection refused"), url="http://fms:9000/api/v1/text/contents") - with pytest.raises(httpx.ConnectError): + with pytest.raises(HTTPConnectionError): await backend.classify("text") @pytest.mark.asyncio @@ -797,11 +854,19 @@ async def test_timeout_raises_after_retry(self, httpx_mock: HTTPXMock): backend = VLLMBackend(_remote(engine="vllm", base_url="http://vllm:8000")) httpx_mock.add_exception(httpx.ReadTimeout("timed out"), url="http://vllm:8000/classify") httpx_mock.add_exception(httpx.ReadTimeout("timed out"), url="http://vllm:8000/classify") - with pytest.raises(httpx.ReadTimeout): + with pytest.raises(HTTPTimeoutError): await backend.classify("text") class TestBackendCaching: + def test_local_backend_reports_ignored_http_client(self, caplog): + caplog.set_level(logging.DEBUG, logger=backends_mod.__name__) + + backend = get_backend(_local(), http_client=RecordingHTTPClient()) + + assert isinstance(backend, LocalBackend) + assert "Ignoring injected HTTP client for local HF classifier backend." in caplog.text + def test_get_backend_returns_same_instance(self): cfg = _remote(engine="vllm") first = get_backend(cfg, name="my_classifier") diff --git a/tests/test_pangea_ai_guard.py b/tests/test_pangea_ai_guard.py index 5376b6c73c..fb86a80ebf 100644 --- a/tests/test_pangea_ai_guard.py +++ b/tests/test_pangea_ai_guard.py @@ -18,7 +18,9 @@ from nemoguardrails import RailsConfig from nemoguardrails.actions.rail_outcome import RailOutcome, TransformTarget +from nemoguardrails.http import HTTPConnectionError from nemoguardrails.library.pangea.actions import TextGuardResult, _pangea_outcome, pangea_ai_guard +from nemoguardrails.testing import RecordingHTTPClient from tests.utils import TestChat input_rail_config = RailsConfig.from_content( @@ -230,7 +232,7 @@ async def test_pangea_ai_guard_api_error_returns_allow_outcome( monkeypatch: pytest.MonkeyPatch, ): monkeypatch.setenv("PANGEA_API_TOKEN", "test-token") - httpx_mock.add_response(status_code=500, json={"result": {}}) + httpx_mock.add_response(is_reusable=True, status_code=500, json={"result": {}}) outcome = await pangea_ai_guard( mode="output", @@ -246,6 +248,26 @@ async def test_pangea_ai_guard_api_error_returns_allow_outcome( assert outcome.metadata["bot_message"] == "Hello" +@pytest.mark.unit +@pytest.mark.asyncio +async def test_pangea_ai_guard_transport_error_returns_allow_outcome(monkeypatch): + monkeypatch.setenv("PANGEA_API_TOKEN", "test-token") + + outcome = await pangea_ai_guard( + mode="output", + config=output_rail_config, + context={"user_message": "Hi", "bot_message": "Hello"}, + http_client=RecordingHTTPClient([HTTPConnectionError("connection failed")]), + ) + + assert outcome.is_blocked is False + assert outcome.is_transform is False + assert outcome.metadata["blocked"] is False + assert outcome.metadata["transformed"] is False + assert outcome.metadata["user_message"] == "Hi" + assert outcome.metadata["bot_message"] == "Hello" + + @pytest.mark.unit def test_pangea_ai_guard_missing_env_var(): chat = TestChat(input_rail_config, llm_completions=[]) diff --git a/tests/test_patronus_evaluate_api.py b/tests/test_patronus_evaluate_api.py index 91bab56029..b2f4749837 100644 --- a/tests/test_patronus_evaluate_api.py +++ b/tests/test_patronus_evaluate_api.py @@ -14,7 +14,6 @@ # limitations under the License. import pytest -from aioresponses import aioresponses from nemoguardrails import RailsConfig from nemoguardrails.actions.actions import ActionResult, action @@ -22,6 +21,7 @@ check_guardrail_pass, patronus_evaluate_request, ) +from tests.http_utils import RecordedHTTPResponses from tests.utils import TestChat PATRONUS_EVALUATE_API_URL = "https://api.patronus.ai/v1/evaluate" @@ -89,7 +89,8 @@ def test_patronus_evaluate_api_success_strategy_all_pass(monkeypatch): ], ) - with aioresponses() as m: + with RecordedHTTPResponses() as m: + chat.app.register_action_param("http_client", m.client) chat.app.register_action(retrieve_relevant_chunks, "retrieve_relevant_chunks") m.post( PATRONUS_EVALUATE_API_URL, @@ -156,7 +157,8 @@ def test_patronus_evaluate_api_success_strategy_all_pass_fails_when_one_failure( ], ) - with aioresponses() as m: + with RecordedHTTPResponses() as m: + chat.app.register_action_param("http_client", m.client) chat.app.register_action(retrieve_relevant_chunks, "retrieve_relevant_chunks") m.post( PATRONUS_EVALUATE_API_URL, @@ -222,7 +224,8 @@ def test_patronus_evaluate_api_success_strategy_any_pass_passes_when_one_failure ], ) - with aioresponses() as m: + with RecordedHTTPResponses() as m: + chat.app.register_action_param("http_client", m.client) chat.app.register_action(retrieve_relevant_chunks, "retrieve_relevant_chunks") m.post( PATRONUS_EVALUATE_API_URL, @@ -288,7 +291,8 @@ def test_patronus_evaluate_api_success_strategy_any_pass_fails_when_all_fail( ], ) - with aioresponses() as m: + with RecordedHTTPResponses() as m: + chat.app.register_action_param("http_client", m.client) chat.app.register_action(retrieve_relevant_chunks, "retrieve_relevant_chunks") m.post( PATRONUS_EVALUATE_API_URL, @@ -351,34 +355,13 @@ def test_patronus_evaluate_api_internal_error_when_no_env_set(): ], ) - with aioresponses() as m: + with RecordedHTTPResponses() as m: + chat.app.register_action_param("http_client", m.client) chat.app.register_action(retrieve_relevant_chunks, "retrieve_relevant_chunks") - m.post( - PATRONUS_EVALUATE_API_URL, - payload={ - "results": [ - { - "evaluator_id": "lynx-large-2024-07-23", - "criteria": "patronus:hallucination", - "status": "success", - "evaluation_result": { - "pass": False, - }, - }, - { - "evaluator_id": "answer-relevance-large-2024-07-23", - "criteria": "patronus:answer-relevance", - "status": "success", - "evaluation_result": { - "pass": False, - }, - }, - ] - }, - ) chat >> "Hi" chat << "I'm sorry, an internal error has occurred." + assert m.client.requests == [] def test_patronus_evaluate_api_internal_error_when_no_evaluators_provided(): @@ -407,34 +390,13 @@ def test_patronus_evaluate_api_internal_error_when_no_evaluators_provided(): ], ) - with aioresponses() as m: + with RecordedHTTPResponses() as m: + chat.app.register_action_param("http_client", m.client) chat.app.register_action(retrieve_relevant_chunks, "retrieve_relevant_chunks") - m.post( - PATRONUS_EVALUATE_API_URL, - payload={ - "results": [ - { - "evaluator_id": "lynx-large-2024-07-23", - "criteria": "patronus:hallucination", - "status": "success", - "evaluation_result": { - "pass": False, - }, - }, - { - "evaluator_id": "answer-relevance-large-2024-07-23", - "criteria": "patronus:answer-relevance", - "status": "success", - "evaluation_result": { - "pass": False, - }, - }, - ] - }, - ) chat >> "Hi" chat << "I'm sorry, an internal error has occurred." + assert m.client.requests == [] def test_patronus_evaluate_api_internal_error_when_evaluator_dict_does_not_have_evaluator_key(): @@ -470,34 +432,13 @@ def test_patronus_evaluate_api_internal_error_when_evaluator_dict_does_not_have_ ], ) - with aioresponses() as m: + with RecordedHTTPResponses() as m: + chat.app.register_action_param("http_client", m.client) chat.app.register_action(retrieve_relevant_chunks, "retrieve_relevant_chunks") - m.post( - PATRONUS_EVALUATE_API_URL, - payload={ - "results": [ - { - "evaluator_id": "lynx-large-2024-07-23", - "criteria": "patronus:hallucination", - "status": "success", - "evaluation_result": { - "pass": False, - }, - }, - { - "evaluator_id": "answer-relevance-large-2024-07-23", - "criteria": "patronus:answer-relevance", - "status": "success", - "evaluation_result": { - "pass": False, - }, - }, - ] - }, - ) chat >> "Hi" chat << "I'm sorry, an internal error has occurred." + assert m.client.requests == [] @pytest.mark.asyncio @@ -537,7 +478,8 @@ def test_patronus_evaluate_api_default_success_strategy_is_all_pass_happy_case( ], ) - with aioresponses() as m: + with RecordedHTTPResponses() as m: + chat.app.register_action_param("http_client", m.client) chat.app.register_action(retrieve_relevant_chunks, "retrieve_relevant_chunks") m.post( PATRONUS_EVALUATE_API_URL, @@ -604,7 +546,8 @@ def test_patronus_evaluate_api_default_success_strategy_all_pass_fails_when_one_ ], ) - with aioresponses() as m: + with RecordedHTTPResponses() as m: + chat.app.register_action_param("http_client", m.client) chat.app.register_action(retrieve_relevant_chunks, "retrieve_relevant_chunks") m.post( PATRONUS_EVALUATE_API_URL, @@ -671,7 +614,8 @@ def test_patronus_evaluate_api_internal_error_when_400_status_code( ], ) - with aioresponses() as m: + with RecordedHTTPResponses() as m: + chat.app.register_action_param("http_client", m.client) chat.app.register_action(retrieve_relevant_chunks, "retrieve_relevant_chunks") m.post( PATRONUS_EVALUATE_API_URL, @@ -719,7 +663,8 @@ def test_patronus_evaluate_api_default_response_when_500_status_code( ], ) - with aioresponses() as m: + with RecordedHTTPResponses() as m: + chat.app.register_action_param("http_client", m.client) chat.app.register_action(retrieve_relevant_chunks, "retrieve_relevant_chunks") m.post( PATRONUS_EVALUATE_API_URL, @@ -794,7 +739,7 @@ def test_check_guardrail_pass_malformed_evaluation_results(): async def test_patronus_evaluate_request_success(monkeypatch): """Test successful API request to Patronus Evaluate endpoint""" monkeypatch.setenv("PATRONUS_API_KEY", "xxx") - with aioresponses() as m: + with RecordedHTTPResponses() as m: m.post( PATRONUS_EVALUATE_API_URL, payload={ @@ -819,6 +764,7 @@ async def test_patronus_evaluate_request_success(monkeypatch): user_input="Does NeMo Guardrails integrate with the Patronus API?", bot_response="Yes, NeMo Guardrails integrates with the Patronus API.", provided_context="Yes, NeMo Guardrails integrates with the Patronus API.", + http_client=m.client, ) assert "results" in response @@ -830,7 +776,7 @@ async def test_patronus_evaluate_request_success(monkeypatch): async def test_patronus_evaluate_request_400_error(monkeypatch): """Test that ValueError is raised with correct message for 400 status code""" monkeypatch.setenv("PATRONUS_API_KEY", "xxx") - with aioresponses() as m: + with RecordedHTTPResponses() as m: m.post( PATRONUS_EVALUATE_API_URL, status=400, @@ -844,6 +790,7 @@ async def test_patronus_evaluate_request_400_error(monkeypatch): user_input="test", bot_response="test", provided_context="test", + http_client=m.client, ) assert "The Patronus Evaluate API call failed with status code 400." in str(exc_info.value) @@ -852,7 +799,7 @@ async def test_patronus_evaluate_request_400_error(monkeypatch): async def test_patronus_evaluate_request_500_error(monkeypatch): """Test that None is returned for 500 status code and no ValueError is raised""" monkeypatch.setenv("PATRONUS_API_KEY", "xxx") - with aioresponses() as m: + with RecordedHTTPResponses() as m: m.post( PATRONUS_EVALUATE_API_URL, status=500, @@ -865,6 +812,7 @@ async def test_patronus_evaluate_request_500_error(monkeypatch): user_input="test", bot_response="test", provided_context="test", + http_client=m.client, ) assert response is None diff --git a/tests/test_policyai_rail.py b/tests/test_policyai_rail.py index c1fa80259c..cb6ac114e9 100644 --- a/tests/test_policyai_rail.py +++ b/tests/test_policyai_rail.py @@ -14,10 +14,10 @@ # limitations under the License. import pytest -from aioresponses import aioresponses from nemoguardrails import RailsConfig from nemoguardrails.library.policyai.actions import call_policyai_api +from tests.http_utils import RecordedHTTPResponses from tests.policyai_fixtures import POLICYAI_SAFE_OUTCOME_KWARGS, POLICYAI_UNSAFE_OUTCOME_KWARGS from tests.utils import TestChat @@ -58,7 +58,8 @@ def test_input_safe(monkeypatch): ], ) - with aioresponses() as m: + with RecordedHTTPResponses() as m: + chat.app.register_action_param("http_client", m.client) # PolicyAI returns SAFE assessment m.post( "https://api.musubilabs.ai/policyai/v1/decisions/evaluate/test", @@ -115,7 +116,8 @@ def test_input_unsafe(monkeypatch): ], ) - with aioresponses() as m: + with RecordedHTTPResponses() as m: + chat.app.register_action_param("http_client", m.client) # PolicyAI returns UNSAFE assessment m.post( "https://api.musubilabs.ai/policyai/v1/decisions/evaluate/test", @@ -161,7 +163,8 @@ def test_output_safe(monkeypatch): ], ) - with aioresponses() as m: + with RecordedHTTPResponses() as m: + chat.app.register_action_param("http_client", m.client) # PolicyAI returns SAFE assessment m.post( "https://api.musubilabs.ai/policyai/v1/decisions/evaluate/test", @@ -207,7 +210,8 @@ def test_output_unsafe(monkeypatch): ], ) - with aioresponses() as m: + with RecordedHTTPResponses() as m: + chat.app.register_action_param("http_client", m.client) # PolicyAI returns UNSAFE assessment (e.g., unauthorized refund promise) m.post( "https://api.musubilabs.ai/policyai/v1/decisions/evaluate/test", @@ -264,7 +268,8 @@ def test_custom_tag_via_env(monkeypatch): ], ) - with aioresponses() as m: + with RecordedHTTPResponses() as m: + chat.app.register_action_param("http_client", m.client) # Note: The URL should use the custom tag from env var m.post( "https://api.musubilabs.ai/policyai/v1/decisions/evaluate/custom-tag", @@ -310,7 +315,8 @@ def test_multiple_policies(monkeypatch): ], ) - with aioresponses() as m: + with RecordedHTTPResponses() as m: + chat.app.register_action_param("http_client", m.client) # Multiple policies - first is SAFE, second is UNSAFE m.post( "https://api.musubilabs.ai/policyai/v1/decisions/evaluate/test", @@ -345,7 +351,7 @@ async def test_empty_data_array_raises_error(monkeypatch): monkeypatch.setenv("POLICYAI_API_KEY", "test-api-key") monkeypatch.setenv("POLICYAI_TAG_NAME", "empty-tag") - with aioresponses() as m: + with RecordedHTTPResponses() as m: # PolicyAI returns empty data array (no policies attached to tag) m.post( "https://api.musubilabs.ai/policyai/v1/decisions/evaluate/empty-tag", @@ -353,7 +359,7 @@ async def test_empty_data_array_raises_error(monkeypatch): ) with pytest.raises(ValueError) as exc_info: - await call_policyai_api(text="Hello!") + await call_policyai_api(text="Hello!", http_client=m.client) assert "no policy results" in str(exc_info.value).lower() assert "empty-tag" in str(exc_info.value) @@ -365,7 +371,7 @@ async def test_api_error_raises_exception(monkeypatch): monkeypatch.setenv("POLICYAI_API_KEY", "test-api-key") monkeypatch.setenv("POLICYAI_TAG_NAME", "test") - with aioresponses() as m: + with RecordedHTTPResponses() as m: # PolicyAI returns 500 error m.post( "https://api.musubilabs.ai/policyai/v1/decisions/evaluate/test", @@ -374,7 +380,7 @@ async def test_api_error_raises_exception(monkeypatch): ) with pytest.raises(ValueError) as exc_info: - await call_policyai_api(text="Hello!") + await call_policyai_api(text="Hello!", http_client=m.client) assert "500" in str(exc_info.value) @@ -385,7 +391,7 @@ async def test_all_policies_failed_raises_error(monkeypatch): monkeypatch.setenv("POLICYAI_API_KEY", "test-api-key") monkeypatch.setenv("POLICYAI_TAG_NAME", "failing-tag") - with aioresponses() as m: + with RecordedHTTPResponses() as m: # All policies return failed status m.post( "https://api.musubilabs.ai/policyai/v1/decisions/evaluate/failing-tag", @@ -404,7 +410,7 @@ async def test_all_policies_failed_raises_error(monkeypatch): ) with pytest.raises(ValueError) as exc_info: - await call_policyai_api(text="Hello!") + await call_policyai_api(text="Hello!", http_client=m.client) assert "all" in str(exc_info.value).lower() assert "failed" in str(exc_info.value).lower() @@ -428,7 +434,7 @@ async def test_empty_text_parameter(monkeypatch): monkeypatch.setenv("POLICYAI_API_KEY", "test-api-key") monkeypatch.setenv("POLICYAI_TAG_NAME", "test") - with aioresponses() as m: + with RecordedHTTPResponses() as m: # PolicyAI should still process empty/None text m.post( "https://api.musubilabs.ai/policyai/v1/decisions/evaluate/test", @@ -446,7 +452,7 @@ async def test_empty_text_parameter(monkeypatch): ) # Test with empty string - result = await call_policyai_api(text="") + result = await call_policyai_api(text="", http_client=m.client) assert result.is_blocked is False assert result.metadata["assessment"] == "SAFE" @@ -457,7 +463,7 @@ async def test_none_text_parameter(monkeypatch): monkeypatch.setenv("POLICYAI_API_KEY", "test-api-key") monkeypatch.setenv("POLICYAI_TAG_NAME", "test") - with aioresponses() as m: + with RecordedHTTPResponses() as m: # PolicyAI should still process None text m.post( "https://api.musubilabs.ai/policyai/v1/decisions/evaluate/test", @@ -475,7 +481,7 @@ async def test_none_text_parameter(monkeypatch): ) # Test with None - result = await call_policyai_api(text=None) + result = await call_policyai_api(text=None, http_client=m.client) assert result.is_blocked is False assert result.metadata["assessment"] == "SAFE" @@ -486,7 +492,7 @@ async def test_partial_policy_failures(monkeypatch): monkeypatch.setenv("POLICYAI_API_KEY", "test-api-key") monkeypatch.setenv("POLICYAI_TAG_NAME", "test") - with aioresponses() as m: + with RecordedHTTPResponses() as m: # Some policies fail, but one succeeds with SAFE m.post( "https://api.musubilabs.ai/policyai/v1/decisions/evaluate/test", @@ -507,7 +513,7 @@ async def test_partial_policy_failures(monkeypatch): }, ) - result = await call_policyai_api(text="Hello!") + result = await call_policyai_api(text="Hello!", http_client=m.client) assert result.is_blocked is False assert result.metadata["assessment"] == "SAFE" assert result.reason == POLICYAI_SAFE_OUTCOME_KWARGS["reason"] @@ -521,7 +527,7 @@ async def test_custom_base_url_with_trailing_slash(monkeypatch): monkeypatch.setenv("POLICYAI_BASE_URL", "https://custom.api.example.com/") monkeypatch.setenv("POLICYAI_TAG_NAME", "test") - with aioresponses() as m: + with RecordedHTTPResponses() as m: # URL should have trailing slash stripped m.post( "https://custom.api.example.com/policyai/v1/decisions/evaluate/test", @@ -538,7 +544,7 @@ async def test_custom_base_url_with_trailing_slash(monkeypatch): }, ) - result = await call_policyai_api(text="Hello!") + result = await call_policyai_api(text="Hello!", http_client=m.client) assert result.is_blocked is False assert result.metadata["assessment"] == "SAFE" @@ -549,7 +555,7 @@ async def test_tag_name_parameter_overrides_env(monkeypatch): monkeypatch.setenv("POLICYAI_API_KEY", "test-api-key") monkeypatch.setenv("POLICYAI_TAG_NAME", "env-tag") - with aioresponses() as m: + with RecordedHTTPResponses() as m: # Should use parameter tag, not env var tag m.post( "https://api.musubilabs.ai/policyai/v1/decisions/evaluate/param-tag", @@ -566,7 +572,11 @@ async def test_tag_name_parameter_overrides_env(monkeypatch): }, ) - result = await call_policyai_api(text="Hello!", tag_name="param-tag") + result = await call_policyai_api( + text="Hello!", + tag_name="param-tag", + http_client=m.client, + ) assert result.is_blocked is False assert result.metadata["assessment"] == "SAFE" @@ -577,7 +587,7 @@ async def test_unsafe_with_missing_fields(monkeypatch): monkeypatch.setenv("POLICYAI_API_KEY", "test-api-key") monkeypatch.setenv("POLICYAI_TAG_NAME", "test") - with aioresponses() as m: + with RecordedHTTPResponses() as m: # UNSAFE response without category, severity, or reason m.post( "https://api.musubilabs.ai/policyai/v1/decisions/evaluate/test", @@ -591,7 +601,7 @@ async def test_unsafe_with_missing_fields(monkeypatch): }, ) - result = await call_policyai_api(text="Bad content") + result = await call_policyai_api(text="Bad content", http_client=m.client) assert result.is_blocked is True assert result.metadata["assessment"] == "UNSAFE" assert result.metadata["category"] == "Unknown" @@ -618,7 +628,7 @@ async def test_default_tag_name_prod(monkeypatch): monkeypatch.setenv("POLICYAI_API_KEY", "test-api-key") monkeypatch.delenv("POLICYAI_TAG_NAME", raising=False) - with aioresponses() as m: + with RecordedHTTPResponses() as m: # Should use default "prod" tag m.post( "https://api.musubilabs.ai/policyai/v1/decisions/evaluate/prod", @@ -635,6 +645,6 @@ async def test_default_tag_name_prod(monkeypatch): }, ) - result = await call_policyai_api(text="Hello!") + result = await call_policyai_api(text="Hello!", http_client=m.client) assert result.is_blocked is False assert result.metadata["assessment"] == "SAFE" diff --git a/tests/test_prompt_security.py b/tests/test_prompt_security.py index 502ad7bbba..f86af09814 100644 --- a/tests/test_prompt_security.py +++ b/tests/test_prompt_security.py @@ -13,11 +13,19 @@ # See the License for the specific language governing permissions and # limitations under the License. +import json + import pytest from nemoguardrails import RailsConfig from nemoguardrails.actions.rail_outcome import RailOutcome, TransformTarget -from nemoguardrails.library.prompt_security.actions import _protect_text_outcome +from nemoguardrails.http import HTTPConnectionError, HTTPResponse +from nemoguardrails.library.prompt_security.actions import ( + _protect_text_outcome, + protect_text, + ps_protect_api_async, +) +from nemoguardrails.testing import RecordingHTTPClient from tests.utils import TestChat @@ -88,6 +96,158 @@ def test_protect_text_outcome(result, expected): assert _protect_text_outcome(result, TransformTarget.USER_MESSAGE) == expected +def _response(payload: dict) -> HTTPResponse: + return HTTPResponse(status_code=200, content=json.dumps(payload).encode()) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("payload", "response", "expected"), + [ + ( + {"result": {"action": "block"}}, + None, + {"is_blocked": True, "is_modified": False, "modified_text": None}, + ), + ( + { + "result": { + "action": "modify", + "prompt": {"modified_text": "safe prompt"}, + } + }, + None, + {"is_blocked": False, "is_modified": True, "modified_text": "safe prompt"}, + ), + ( + { + "result": { + "action": "modify", + "response": {"modified_text": "safe response"}, + } + }, + "response", + {"is_blocked": False, "is_modified": True, "modified_text": "safe response"}, + ), + ], +) +async def test_prompt_security_api_uses_shared_client(payload, response, expected): + client = RecordingHTTPClient([_response(payload)]) + + result = await ps_protect_api_async( + "https://prompt-security.example/protect", + "app-id", + prompt="prompt" if response is None else None, + system_prompt="system", + response=response, + user="user", + http_client=client, + ) + + assert result == expected + request = client.requests[0] + assert request.method == "POST" + assert request.url == "https://prompt-security.example/protect" + assert request.headers == { + "APP-ID": "app-id", + "Content-Type": "application/json", + } + assert request.json == { + "prompt": "prompt" if response is None else None, + "system_prompt": "system", + "response": response, + "user": "user", + } + assert request.timeout is None + + +@pytest.mark.asyncio +async def test_prompt_security_api_allows_on_http_failure(caplog): + client = RecordingHTTPClient([HTTPConnectionError("connection failed")]) + + result = await ps_protect_api_async( + "https://prompt-security.example/protect", + "app-id", + prompt="prompt", + http_client=client, + ) + + assert result == { + "is_blocked": False, + "is_modified": False, + "modified_text": None, + } + assert "Error calling Prompt Security Protect API: connection failed" in caplog.text + + +@pytest.mark.asyncio +async def test_protect_text_forwards_shared_client(monkeypatch): + monkeypatch.setenv("PS_PROTECT_URL", "https://prompt-security.example/protect") + monkeypatch.setenv("PS_APP_ID", "app-id") + client = RecordingHTTPClient( + [ + _response({"result": {"action": "block"}}), + _response( + { + "result": { + "action": "modify", + "response": {"modified_text": "safe response"}, + } + } + ), + ] + ) + + input_outcome = await protect_text(user_prompt="prompt", http_client=client) + output_outcome = await protect_text(bot_response="response", http_client=client) + + assert input_outcome == RailOutcome.block( + metadata={ + "is_blocked": True, + "is_modified": False, + "modified_text": None, + } + ) + assert output_outcome == RailOutcome.transform( + [(TransformTarget.BOT_MESSAGE, "safe response")], + metadata={ + "is_blocked": False, + "is_modified": True, + "modified_text": "safe response", + }, + ) + input_request, output_request = client.requests + assert input_request.json == { + "prompt": "prompt", + "system_prompt": None, + "response": None, + "user": None, + } + assert output_request.json == { + "prompt": None, + "system_prompt": None, + "response": "response", + "user": None, + } + + +@pytest.mark.asyncio +async def test_protect_text_validates_configuration(monkeypatch): + monkeypatch.delenv("PS_PROTECT_URL", raising=False) + monkeypatch.delenv("PS_APP_ID", raising=False) + + with pytest.raises(ValueError, match="PS_PROTECT_URL"): + await protect_text(user_prompt="prompt") + + monkeypatch.setenv("PS_PROTECT_URL", "https://prompt-security.example/protect") + with pytest.raises(ValueError, match="PS_APP_ID"): + await protect_text(user_prompt="prompt") + + monkeypatch.setenv("PS_APP_ID", "app-id") + with pytest.raises(ValueError, match="Neither user_message nor bot_message"): + await protect_text() + + @pytest.mark.unit def test_prompt_security_protection_input(): config = RailsConfig.from_content( diff --git a/tests/test_trend_ai_guard.py b/tests/test_trend_ai_guard.py index 9d36733ff8..f12d2ff118 100644 --- a/tests/test_trend_ai_guard.py +++ b/tests/test_trend_ai_guard.py @@ -17,6 +17,10 @@ from pytest_httpx import HTTPXMock from nemoguardrails import RailsConfig +from nemoguardrails.actions.rail_outcome import RailOutcome +from nemoguardrails.http import HTTPConnectionError, HTTPResponse +from nemoguardrails.library.trend_micro.actions import trend_ai_guard +from nemoguardrails.testing import RecordingHTTPClient from tests.utils import TestChat input_rail_config = RailsConfig.from_content( @@ -96,6 +100,31 @@ def test_trend_ai_guard_error(httpx_mock: HTTPXMock, monkeypatch: pytest.MonkeyP chat << "Hello!" +@pytest.mark.unit +@pytest.mark.asyncio +@pytest.mark.parametrize( + "response", + [ + HTTPConnectionError("connection failed"), + HTTPResponse(status_code=200, content=b"not-json"), + HTTPResponse(status_code=200, content=b"{}"), + ], +) +async def test_trend_ai_guard_fail_open_covers_request_and_response_failures(monkeypatch, response): + monkeypatch.setenv("V1_API_KEY", "test-token") + + outcome = await trend_ai_guard( + config=input_rail_config, + text="Hello", + http_client=RecordingHTTPClient([response]), + ) + + assert outcome == RailOutcome.allow( + reason="An error occurred while calling the Trend Micro AI Guard API.", + metadata={"action": "Allow"}, + ) + + @pytest.mark.unit def test_trend_ai_guard_missing_env_var(): chat = TestChat(input_rail_config, llm_completions=[]) @@ -116,7 +145,7 @@ def test_trend_ai_guard_malformed_response(httpx_mock: HTTPXMock, monkeypatch: p # Should fail open chat >> "What is the air-speed velocity of an unladen swallow?" - chat << "I'm sorry, an internal error has occurred." + chat << "What do you mean? An African or a European swallow?" @pytest.mark.unit