fix(polygraf): align request auth/payload with API format

Author: Ali-Khodabakhsh

nemoguardrails/library/polygraf/request.py

@@ -39,11 +39,20 @@ async def polygraf_request(
     Raises:
         ValueError: If the API call fails or the response cannot be parsed as JSON.
     """
-    payload = {"text": text}
+    # Polygraf request payload. Some deployments accept/require additional flags
+    # controlling PII/PID detection and aggregation.
+    payload = {
+        "text": text,
+        # NOTE: Kept as `detect_pid` to match the working Polygraf API format
+        # provided by users of this integration.
+        "detect_pid": True,
+        "pid_granularity": 3,
+        "aggregate_entities": True,
+    }
     headers: Dict[str, str] = {"Content-Type": "application/json"}
 
     if api_key:
-        headers["API_Key"] = f"Bearer {api_key}"
+        headers["Authorization"] = f"Bearer {api_key}"
 
     async with aiohttp.ClientSession() as session:
         async with session.post(
@@ -55,8 +64,20 @@ async def polygraf_request(
                 )
 
             try:
-                return await resp.json()
+                data = await resp.json()
             except aiohttp.ContentTypeError:
                 raise ValueError(
                     f"Failed to parse Polygraf response as JSON. Status: {resp.status}, Content: {await resp.text()}"
                 )
+
+            # Polygraf may return either a raw list of entities or a wrapper object.
+            if isinstance(data, list):
+                return data
+            if isinstance(data, dict):
+                entities = data.get("entities")
+                if isinstance(entities, list):
+                    return entities
+
+            raise ValueError(
+                "Invalid response from Polygraf service: expected a list of entities or an object with an 'entities' list."
+            )