refactor(library): migrate HTTP request helpers

nemoguardrails/library/activefence/actions.py

@@ -17,10 +17,9 @@
 import os
 from typing import Literal, Optional
 
-import aiohttp
-
 from nemoguardrails.actions import action
 from nemoguardrails.actions.rail_outcome import RailOutcome
+from nemoguardrails.http import HTTPClient, http_call
 from nemoguardrails.utils import new_uuid
 
 log = logging.getLogger(__name__)
@@ -103,6 +102,7 @@ def _activefence_outcome(
 async def call_activefence_api(
     text: Optional[str] = None,
     threshold_mode: Literal["simple", "detailed"] = "simple",
+    http_client: Optional[HTTPClient] = None,
     **kwargs,
 ) -> RailOutcome:
     api_key = os.environ.get("ACTIVEFENCE_API_KEY")
@@ -117,25 +117,25 @@ async def call_activefence_api(
         "content_id": "ng-" + new_uuid(),
     }
 
-    async with aiohttp.ClientSession() as session:
-        async with session.post(
-            url=url,
-            headers=headers,
-            json=data,
-        ) as response:
-            if response.status != 200:
-                raise ValueError(
-                    f"ActiveFence call failed with status code {response.status}.\nDetails: {await response.text()}"
-                )
-            response_json = await response.json()
-            log.info(json.dumps(response_json, indent=True))
-            violations = response_json["violations"]
-
-            violations_dict = {}
-            max_risk_score = 0.0
-            for violation in violations:
-                if violation["risk_score"] > max_risk_score:
-                    max_risk_score = violation["risk_score"]
-                violations_dict[violation["violation_type"]] = violation["risk_score"]
-
-            return _activefence_outcome(max_risk_score, violations_dict, threshold_mode)
+    response = await http_call(
+        http_client,
+        "POST",
+        url,
+        headers=headers,
+        json=data,
+        raise_for_status=False,
+    )
+    if response.status_code != 200:
+        raise ValueError(f"ActiveFence call failed with status code {response.status_code}.\nDetails: {response.text}")
+    response_json = response.json()
+    log.info(json.dumps(response_json, indent=True))
+    violations = response_json["violations"]
+
+    violations_dict = {}
+    max_risk_score = 0.0
+    for violation in violations:
+        if violation["risk_score"] > max_risk_score:
+            max_risk_score = violation["risk_score"]
+        violations_dict[violation["violation_type"]] = violation["risk_score"]
+
+    return _activefence_outcome(max_risk_score, violations_dict, threshold_mode)

nemoguardrails/library/factchecking/align_score/actions.py

@@ -19,6 +19,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.factchecking.align_score.request import alignscore_request
 from nemoguardrails.library.self_check.facts.actions import _fact_check_outcome, self_check_facts
 from nemoguardrails.llm.taskmanager import LLMTaskManager
@@ -33,6 +34,7 @@ async def alignscore_check_facts(
     context: Optional[dict] = None,
     llm: Optional[LLMModel] = None,
     config: Optional[RailsConfig] = None,
+    http_client: Optional[HTTPClient] = None,
     **kwargs,
 ) -> RailOutcome:
     """Checks the facts for the bot response using an information alignment score."""
@@ -45,7 +47,8 @@ async def alignscore_check_facts(
     evidence = context.get("relevant_chunks", [])
     response = context.get("bot_message")
 
-    alignscore = await alignscore_request(alignscore_api_url, evidence, response)
+    request_kwargs = {"http_client": http_client} if http_client is not None else {}
+    alignscore = await alignscore_request(alignscore_api_url, evidence, response, **request_kwargs)
     if alignscore is None:
         log.warning("AlignScore endpoint not set up properly. Falling back to the ask_llm approach for fact-checking.")
         if fallback_to_self_check:

nemoguardrails/library/factchecking/align_score/request.py

@@ -16,9 +16,8 @@
 import logging
 from typing import Optional
 
-import aiohttp
-
 from nemoguardrails.actions import action
+from nemoguardrails.http import HTTPClient, http_call
 
 log = logging.getLogger(__name__)
 
@@ -28,24 +27,30 @@ async def alignscore_request(
     api_url: str = "http://localhost:5000/alignscore_large",
     evidence: Optional[list] = None,
     response: Optional[str] = None,
+    http_client: Optional[HTTPClient] = None,
 ):
     """Checks the facts for the bot response by making a request to the AlignScore API."""
     if not evidence:
         return 1.0
 
     payload = {"evidence": evidence, "claim": response}
 
-    async with aiohttp.ClientSession() as session:
-        async with session.post(api_url, json=payload) as resp:
-            if resp.status != 200:
-                log.error(f"AlignScore API request failed with status {resp.status}")
-                return None
-
-            result = await resp.json()
-
-            log.info(f"AlignScore was {result}.")
-            try:
-                result = result["alignscore"]
-            except Exception:
-                result = None
-            return result
+    http_response = await http_call(
+        http_client,
+        "POST",
+        api_url,
+        json=payload,
+        raise_for_status=False,
+    )
+    if http_response.status_code != 200:
+        log.error(f"AlignScore API request failed with status {http_response.status_code}")
+        return None
+
+    result = http_response.json()
+
+    log.info(f"AlignScore was {result}.")
+    try:
+        result = result["alignscore"]
+    except Exception:
+        result = None
+    return result

nemoguardrails/library/gliner/actions.py

@@ -22,6 +22,7 @@
 from nemoguardrails import RailsConfig
 from nemoguardrails.actions import action
 from nemoguardrails.actions.rail_outcome import RailOutcome, TransformTarget
+from nemoguardrails.http import HTTPClient
 from nemoguardrails.library.gliner.request import gliner_request
 from nemoguardrails.rails.llm.config import GLiNERDetection
 
@@ -98,6 +99,7 @@ async def gliner_detect_pii(
     source: str,
     text: str,
     config: RailsConfig,
+    http_client: HTTPClient | None = None,
     **kwargs,
 ) -> RailOutcome:
     """Checks whether the provided text contains any PII using GLiNER.
@@ -128,6 +130,7 @@ async def gliner_detect_pii(
 
     api_key = _resolve_api_key(gliner_config)
 
+    request_kwargs = {"http_client": http_client} if http_client is not None else {}
     gliner_response = await gliner_request(
         text=text,
         server_endpoint=server_endpoint,
@@ -138,6 +141,7 @@ async def gliner_detect_pii(
         flat_ner=gliner_config.flat_ner,
         api_key=api_key,
         model=gliner_config.model,
+        **request_kwargs,
     )
 
     try:
@@ -148,7 +152,12 @@ async def gliner_detect_pii(
 
 
 @action(is_system_action=False)
-async def gliner_mask_pii(source: str, text: str, config: RailsConfig) -> RailOutcome:
+async def gliner_mask_pii(
+    source: str,
+    text: str,
+    config: RailsConfig,
+    http_client: HTTPClient | None = None,
+) -> RailOutcome:
     """Masks any detected PII in the provided text using GLiNER.
 
     Args:
@@ -177,6 +186,7 @@ async def gliner_mask_pii(source: str, text: str, config: RailsConfig) -> RailOu
 
     api_key = _resolve_api_key(gliner_config)
 
+    request_kwargs = {"http_client": http_client} if http_client is not None else {}
     gliner_response = await gliner_request(
         text=text,
         server_endpoint=server_endpoint,
@@ -187,6 +197,7 @@ async def gliner_mask_pii(source: str, text: str, config: RailsConfig) -> RailOu
         flat_ner=gliner_config.flat_ner,
         api_key=api_key,
         model=gliner_config.model,
+        **request_kwargs,
     )
 
     if not gliner_response or not isinstance(gliner_response, dict):

nemoguardrails/library/gliner/request.py

@@ -19,8 +19,7 @@
 import logging
 from typing import Any, Dict, List, Optional
 
-import aiohttp
-
+from nemoguardrails.http import HTTPClient, HTTPResponseDecodeError, http_call
 from nemoguardrails.library.gliner.models import GLiNERRequest
 
 log = logging.getLogger(__name__)
@@ -36,6 +35,7 @@ async def gliner_request(
     flat_ner: Optional[bool] = None,
     api_key: Optional[str] = None,
     model: str = "nvidia/gliner-pii",
+    http_client: Optional[HTTPClient] = None,
 ) -> Dict[str, Any]:
     """Send a PII detection request to the GLiNER API.
 
@@ -107,44 +107,47 @@ async def gliner_request(
     if request.labels:
         payload["labels"] = request.labels
 
-    async with aiohttp.ClientSession() as session:
-        async with session.post(server_endpoint, json=payload, headers=headers) as resp:
-            if resp.status != 200:
-                raise ValueError(f"GLiNER call failed with status code {resp.status}.\nDetails: {await resp.text()}")
-
-            try:
-                raw = await resp.json()
-            except aiohttp.ContentTypeError as err:
-                raise ValueError(
-                    f"Failed to parse GLiNER response as JSON. Status: {resp.status}, Content: {await resp.text()}"
-                ) from err
-
-            if not use_chat_completions:
-                return raw
-
-            # Unwrap chat completions envelope and normalize entity field names.
-            # NIM uses: text, label, start, end
-            # Custom server uses: value, suggested_label, start_position, end_position
-            try:
-                nim_data = json.loads(raw["choices"][0]["message"]["content"])
-            except (KeyError, IndexError, json.JSONDecodeError, TypeError) as e:
-                raise ValueError(f"Failed to parse NIM response content: {e}") from e
-
-            if not isinstance(nim_data, dict):
-                raise ValueError(f"Expected NIM response content to be a JSON object, got {type(nim_data).__name__}")
-
-            normalized_entities = [
-                {
-                    "value": e.get("text", e.get("value", "")),
-                    "suggested_label": e.get("label", e.get("suggested_label", "")),
-                    "start_position": e.get("start", e.get("start_position", 0)),
-                    "end_position": e.get("end", e.get("end_position", 0)),
-                    "score": e.get("score", 0.0),
-                }
-                for e in (nim_data.get("entities") or [])
-            ]
-            return {
-                "entities": normalized_entities,
-                "total_entities": nim_data.get("total_entities", len(normalized_entities)),
-                "tagged_text": nim_data.get("tagged_text", ""),
-            }
+    response = await http_call(
+        http_client,
+        "POST",
+        server_endpoint,
+        json=payload,
+        headers=headers,
+        raise_for_status=False,
+    )
+    if response.status_code != 200:
+        raise ValueError(f"GLiNER call failed with status code {response.status_code}.\nDetails: {response.text}")
+
+    try:
+        raw = response.json()
+    except HTTPResponseDecodeError as error:
+        raise ValueError(
+            f"Failed to parse GLiNER response as JSON. Status: {response.status_code}, Content: {response.text}"
+        ) from error
+
+    if not use_chat_completions:
+        return raw
+
+    try:
+        nim_data = json.loads(raw["choices"][0]["message"]["content"])
+    except (KeyError, IndexError, json.JSONDecodeError, TypeError) as e:
+        raise ValueError(f"Failed to parse NIM response content: {e}") from e
+
+    if not isinstance(nim_data, dict):
+        raise ValueError(f"Expected NIM response content to be a JSON object, got {type(nim_data).__name__}")
+
+    normalized_entities = [
+        {
+            "value": e.get("text", e.get("value", "")),
+            "suggested_label": e.get("label", e.get("suggested_label", "")),
+            "start_position": e.get("start", e.get("start_position", 0)),
+            "end_position": e.get("end", e.get("end_position", 0)),
+            "score": e.get("score", 0.0),
+        }
+        for e in (nim_data.get("entities") or [])
+    ]
+    return {
+        "entities": normalized_entities,
+        "total_entities": nim_data.get("total_entities", len(normalized_entities)),
+        "tagged_text": nim_data.get("tagged_text", ""),
+    }

nemoguardrails/library/jailbreak_detection/actions.py

@@ -35,6 +35,7 @@
 from nemoguardrails.actions import action
 from nemoguardrails.actions.rail_outcome import RailOutcome
 from nemoguardrails.context import llm_call_info_var
+from nemoguardrails.http import HTTPClient
 from nemoguardrails.library.jailbreak_detection.request import (
     jailbreak_detection_heuristics_request,
     jailbreak_detection_model_request,
@@ -57,6 +58,7 @@
 async def jailbreak_detection_heuristics(
     llm_task_manager: LLMTaskManager,
     context: Optional[dict] = None,
+    http_client: Optional[HTTPClient] = None,
     **kwargs,
 ) -> RailOutcome:
     """Checks the user's prompt to determine if it is attempt to jailbreak the model."""
@@ -80,7 +82,14 @@ async def jailbreak_detection_heuristics(
         jailbreak = any([lp_check["jailbreak"], ps_ppl_check["jailbreak"]])
         return RailOutcome.block() if jailbreak else RailOutcome.allow()
 
-    jailbreak = await jailbreak_detection_heuristics_request(prompt, jailbreak_api_url, lp_threshold, ps_ppl_threshold)
+    request_kwargs = {"http_client": http_client} if http_client is not None else {}
+    jailbreak = await jailbreak_detection_heuristics_request(
+        prompt,
+        jailbreak_api_url,
+        lp_threshold,
+        ps_ppl_threshold,
+        **request_kwargs,
+    )
     if jailbreak is None:
         log.warning("Jailbreak endpoint not set up properly.")
         # If no result, assume not a jailbreak
@@ -94,6 +103,7 @@ async def jailbreak_detection_model(
     llm_task_manager: LLMTaskManager,
     context: Optional[dict] = None,
     model_caches: Optional[Dict[str, CacheInterface]] = None,
+    http_client: Optional[HTTPClient] = None,
 ) -> RailOutcome:
     """Uses a trained classifier to determine if a user input is a jailbreak attempt"""
     prompt: str = ""
@@ -151,15 +161,21 @@ async def jailbreak_detection_model(
             )
             jailbreak_result = False
     else:
+        request_kwargs = {"http_client": http_client} if http_client is not None else {}
         if nim_base_url:
             jailbreak = await jailbreak_nim_request(
                 prompt=prompt,
                 nim_url=nim_base_url,
                 nim_auth_token=nim_auth_token,
                 nim_classification_path=nim_classification_path,
+                **request_kwargs,
             )
         elif jailbreak_api_url:
-            jailbreak = await jailbreak_detection_model_request(prompt=prompt, api_url=jailbreak_api_url)
+            jailbreak = await jailbreak_detection_model_request(
+                prompt=prompt,
+                api_url=jailbreak_api_url,
+                **request_kwargs,
+            )
 
         if jailbreak is None:
             log.warning("Jailbreak endpoint not set up properly.")