pangea-sdk: add llm_input, llm_info, and log_fields to AI Guard

Author: kenany

CHANGELOG.md

@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## Unreleased
 
+### Added
+
+- AI Guard and Prompt Guard services.
+
 ### Changed
 
 - Clarified what `PangeaConfig.environment` affects.

examples/ai_guard/ai_guard_examples/guard.py

@@ -20,6 +20,6 @@
 # Structured input.
 structured_input = [{"role": "user", "content": "hello world"}]
 print("Guarding structured input:", structured_input)
-structured_response = ai_guard.guard_text(structured_input)
+structured_response = ai_guard.guard_text(messages=structured_input)
 assert structured_response.result
 print("Response:", structured_response.result.prompt_messages)

examples/ai_guard/poetry.lock

@@ -23,7 +23,7 @@ async def main() -> None:
     # Structured input.
     structured_input = [{"role": "user", "content": "hello world"}]
     print("Guarding structured input:", structured_input)
-    structured_response = await ai_guard.guard_text(structured_input)
+    structured_response = await ai_guard.guard_text(messages=structured_input)
     assert structured_response.result
     print("Response:", structured_response.result.prompt_messages)
 

examples/asyncio/ai_guard/ai_guard_examples/guard.py

@@ -7,7 +7,7 @@
 from pangea.asyncio.services.base import ServiceBaseAsync
 from pangea.config import PangeaConfig
 from pangea.response import PangeaResponse
-from pangea.services.ai_guard import TextGuardResult
+from pangea.services.ai_guard import LogFields, TextGuardResult
 
 _T = TypeVar("_T")
 
@@ -54,10 +54,12 @@ def __init__(
     @overload
     async def guard_text(
         self,
-        text_or_messages: str,
+        text: str,
         *,
-        recipe: str = "pangea_prompt_guard",
-        debug: bool = False,
+        recipe: str | None = None,
+        debug: bool | None = None,
+        llm_info: str | None = None,
+        log_fields: LogFields | None = None,
     ) -> PangeaResponse[TextGuardResult[None]]:
         """
         Text Guard for scanning LLM inputs and outputs
@@ -76,6 +78,8 @@ async def guard_text(
                 are to be applied to the text, such as defang malicious URLs.
             debug: Setting this value to true will provide a detailed analysis
                 of the text data
+            llm_info: Short string hint for the LLM Provider information
+            log_field: Additional fields to include in activity log
 
         Examples:
             response = await ai_guard.guard_text("text")
@@ -84,10 +88,12 @@ async def guard_text(
     @overload
     async def guard_text(
         self,
-        text_or_messages: _T,
         *,
-        recipe: str = "pangea_prompt_guard",
-        debug: bool = False,
+        messages: _T,
+        recipe: str | None = None,
+        debug: bool | None = None,
+        llm_info: str | None = None,
+        log_fields: LogFields | None = None,
     ) -> PangeaResponse[TextGuardResult[_T]]:
         """
         Text Guard for scanning LLM inputs and outputs
@@ -98,27 +104,31 @@ async def guard_text(
         OperationId: ai_guard_post_v1_text_guard
 
         Args:
-            text_or_messages: Structured data to be scanned by AI Guard for PII,
-                sensitive data, malicious content, and other data types defined
-                by the configuration. Supports processing up to 10KB of text.
+            messages: Structured messages data to be scanned by AI Guard for
+                PII, sensitive data, malicious content, and other data types
+                defined by the configuration. Supports processing up to 10KB of
+                JSON text
             recipe: Recipe key of a configuration of data types and settings
                 defined in the Pangea User Console. It specifies the rules that
                 are to be applied to the text, such as defang malicious URLs.
             debug: Setting this value to true will provide a detailed analysis
                 of the text data
+            llm_info: Short string hint for the LLM Provider information
+            log_field: Additional fields to include in activity log
 
         Examples:
-            response = await ai_guard.guard_text([
-                {"role": "user", "content": "hello world"}
-            ])
+            response = await ai_guard.guard_text(messages=[{"role": "user", "content": "hello world"}])
         """
 
+    @overload
     async def guard_text(
         self,
-        text_or_messages: str | _T,
         *,
-        recipe: str = "pangea_prompt_guard",
-        debug: bool = False,
+        llm_input: _T,
+        recipe: str | None = None,
+        debug: bool | None = None,
+        llm_info: str | None = None,
+        log_fields: LogFields | None = None,
     ) -> PangeaResponse[TextGuardResult[_T]]:
         """
         Text Guard for scanning LLM inputs and outputs
@@ -129,25 +139,83 @@ async def guard_text(
         OperationId: ai_guard_post_v1_text_guard
 
         Args:
-            text_or_messages: Text or structured data to be scanned by AI Guard
-                for PII, sensitive data, malicious content, and other data types
-                defined by the configuration. Supports processing up to 10KB of text.
+            llm_input: Structured full llm payload data to be scanned by AI
+                Guard for PII, sensitive data, malicious content, and other data
+                types defined by the configuration. Supports processing up to
+                10KB of JSON text
             recipe: Recipe key of a configuration of data types and settings
                 defined in the Pangea User Console. It specifies the rules that
                 are to be applied to the text, such as defang malicious URLs.
             debug: Setting this value to true will provide a detailed analysis
                 of the text data
+            llm_info: Short string hint for the LLM Provider information
+            log_field: Additional fields to include in activity log
+
+        Examples:
+            response = await ai_guard.guard_text(
+                llm_input={"model": "gpt-4o", "messages": [{"role": "user", "content": "hello world"}]}
+            )
+        """
+
+    async def guard_text(  # type: ignore[misc]
+        self,
+        text: str | None = None,
+        *,
+        messages: _T | None = None,
+        llm_input: _T | None = None,
+        recipe: str | None = None,
+        debug: bool | None = None,
+        llm_info: str | None = None,
+        log_fields: LogFields | None = None,
+    ) -> PangeaResponse[TextGuardResult[None]]:
+        """
+        Text Guard for scanning LLM inputs and outputs
+
+        Analyze and redact text to avoid manipulation of the model, addition of
+        malicious content, and other undesirable data transfers.
+
+        OperationId: ai_guard_post_v1_text_guard
+
+        Args:
+            text: Text to be scanned by AI Guard for PII, sensitive data,
+                malicious content, and other data types defined by the
+                configuration. Supports processing up to 10KB of text.
+            messages: Structured messages data to be scanned by AI Guard for
+                PII, sensitive data, malicious content, and other data types
+                defined by the configuration. Supports processing up to 10KB of
+                JSON text
+            llm_input: Structured full llm payload data to be scanned by AI
+                Guard for PII, sensitive data, malicious content, and other data
+                types defined by the configuration. Supports processing up to
+                10KB of JSON text
+            recipe: Recipe key of a configuration of data types and settings
+                defined in the Pangea User Console. It specifies the rules that
+                are to be applied to the text, such as defang malicious URLs.
+            debug: Setting this value to true will provide a detailed analysis
+                of the text data
+            llm_info: Short string hint for the LLM Provider information
+            log_field: Additional fields to include in activity log
 
         Examples:
             response = await ai_guard.guard_text("text")
         """
 
+        if not any((text, messages, llm_input)):
+            raise ValueError("Exactly one of `text`, `messages`, or `llm_input` must be given")
+
+        if sum((text is not None, messages is not None, llm_input is not None)) > 1:
+            raise ValueError("Only one of `text`, `messages`, or `llm_input` can be given at once")
+
         return await self.request.post(
             "v1/text/guard",
             TextGuardResult,
             data={
-                "text" if isinstance(text_or_messages, str) else "messages": text_or_messages,
+                "text": text,
+                "messages": messages,
+                "llm_input": llm_input,
                 "recipe": recipe,
                 "debug": debug,
+                "llm_info": llm_info,
+                "log_fields": log_fields,
             },
         )