pangea-sdk: make AI Guard messages parameter stricter

Author: kenany

CHANGELOG.md

@@ -7,6 +7,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## Unreleased
 
+### Changed
+
+- AI Guard: `messages` parameter is no longer a generic. A new `Message` model
+  has been introduced, and `messages` is now a `Sequence[Message]`.
+
 ## 6.1.1 - 2025-05-12
 
 ### Fixed

examples/ai_guard/ai_guard_examples/guard_text.py

@@ -3,7 +3,7 @@
 import os
 
 from pangea import PangeaConfig
-from pangea.services.ai_guard import AIGuard
+from pangea.services.ai_guard import AIGuard, Message
 
 token = os.getenv("PANGEA_AI_GUARD_TOKEN", "")
 domain = os.getenv("PANGEA_DOMAIN", "aws.us.pangea.cloud")
@@ -18,7 +18,7 @@
 print("Response:", text_response.result.prompt_text)
 
 # Structured input.
-structured_input = [{"role": "user", "content": "hello world"}]
+structured_input = [Message(role="user", content="hello world")]
 print("Guarding structured input:", structured_input)
 structured_response = ai_guard.guard_text(messages=structured_input)
 assert structured_response.result

examples/asyncio/ai_guard/async_ai_guard_examples/guard_text.py

@@ -5,6 +5,7 @@
 
 from pangea import PangeaConfig
 from pangea.asyncio.services import AIGuardAsync
+from pangea.services.ai_guard import Message
 
 token = os.getenv("PANGEA_AI_GUARD_TOKEN", "")
 domain = os.getenv("PANGEA_DOMAIN", "aws.us.pangea.cloud")
@@ -21,7 +22,7 @@ async def main() -> None:
     print("Response:", text_response.result.prompt_text)
 
     # Structured input.
-    structured_input = [{"role": "user", "content": "hello world"}]
+    structured_input = [Message(role="user", content="hello world")]
     print("Guarding structured input:", structured_input)
     structured_response = await ai_guard.guard_text(messages=structured_input)
     assert structured_response.result

packages/pangea-sdk/pangea/asyncio/services/ai_guard.py

@@ -1,13 +1,14 @@
 from __future__ import annotations
 
+from collections.abc import Sequence
 from typing import overload
 
 from typing_extensions import TypeVar
 
 from pangea.asyncio.services.base import ServiceBaseAsync
 from pangea.config import PangeaConfig
 from pangea.response import PangeaResponse
-from pangea.services.ai_guard import LogFields, Overrides, TextGuardResult
+from pangea.services.ai_guard import LogFields, Message, Overrides, TextGuardResult
 
 _T = TypeVar("_T")
 
@@ -60,7 +61,7 @@ async def guard_text(
         debug: bool | None = None,
         overrides: Overrides | None = None,
         log_fields: LogFields | None = None,
-    ) -> PangeaResponse[TextGuardResult[None]]:
+    ) -> PangeaResponse[TextGuardResult]:
         """
         Text Guard for scanning LLM inputs and outputs
 
@@ -88,12 +89,12 @@ async def guard_text(
     async def guard_text(
         self,
         *,
-        messages: _T,
+        messages: Sequence[Message],
         recipe: str | None = None,
         debug: bool | None = None,
         overrides: Overrides | None = None,
         log_fields: LogFields | None = None,
-    ) -> PangeaResponse[TextGuardResult[_T]]:
+    ) -> PangeaResponse[TextGuardResult]:
         """
         Text Guard for scanning LLM inputs and outputs
 
@@ -115,19 +116,19 @@ async def guard_text(
             log_field: Additional fields to include in activity log
 
         Examples:
-            response = await ai_guard.guard_text(messages=[{"role": "user", "content": "hello world"}])
+            response = await ai_guard.guard_text(messages=[Message(role="user", content="hello world")])
         """
 
-    async def guard_text(  # type: ignore[misc]
+    async def guard_text(
         self,
         text: str | None = None,
         *,
-        messages: _T | None = None,
+        messages: Sequence[Message] | None = None,
         recipe: str | None = None,
         debug: bool | None = None,
         overrides: Overrides | None = None,
         log_fields: LogFields | None = None,
-    ) -> PangeaResponse[TextGuardResult[None]]:
+    ) -> PangeaResponse[TextGuardResult]:
         """
         Text Guard for scanning LLM inputs and outputs
 

packages/pangea-sdk/pangea/services/ai_guard.py

@@ -1,7 +1,9 @@
 from __future__ import annotations
 
+from collections.abc import Sequence
 from typing import Generic, Literal, Optional, overload
 
+from pydantic import BaseModel, ConfigDict
 from typing_extensions import TypeVar
 
 from pangea.config import PangeaConfig
@@ -19,6 +21,13 @@
 PiiEntityAction = Literal["disabled", "report", "block", "mask", "partial_masking", "replacement", "hash", "fpe"]
 
 
+class Message(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    role: str
+    content: str
+
+
 class CodeDetectionOverride(APIRequestModel):
     disabled: Optional[bool] = None
     action: Optional[Literal["report", "block"]] = None
@@ -276,14 +285,14 @@ class TextGuardDetectors(APIResponseModel):
     code_detection: Optional[TextGuardDetector[CodeDetectionResult]] = None
 
 
-class TextGuardResult(PangeaResponseResult, Generic[_T]):
+class TextGuardResult(PangeaResponseResult):
     detectors: TextGuardDetectors
     """Result of the recipe analyzing and input prompt."""
 
     prompt_text: Optional[str] = None
     """Updated prompt text, if applicable."""
 
-    prompt_messages: Optional[_T] = None
+    prompt_messages: Optional[object] = None
     """Updated structured prompt, if applicable."""
 
     blocked: bool
@@ -347,7 +356,7 @@ def guard_text(
         debug: bool | None = None,
         overrides: Overrides | None = None,
         log_fields: LogFields | None = None,
-    ) -> PangeaResponse[TextGuardResult[None]]:
+    ) -> PangeaResponse[TextGuardResult]:
         """
         Text Guard for scanning LLM inputs and outputs
 
@@ -375,12 +384,12 @@ def guard_text(
     def guard_text(
         self,
         *,
-        messages: _T,
+        messages: Sequence[Message],
         recipe: str | None = None,
         debug: bool | None = None,
         overrides: Overrides | None = None,
         log_fields: LogFields | None = None,
-    ) -> PangeaResponse[TextGuardResult[_T]]:
+    ) -> PangeaResponse[TextGuardResult]:
         """
         Text Guard for scanning LLM inputs and outputs
 
@@ -402,19 +411,19 @@ def guard_text(
             log_field: Additional fields to include in activity log
 
         Examples:
-            response = ai_guard.guard_text(messages=[{"role": "user", "content": "hello world"}])
+            response = ai_guard.guard_text(messages=[Message(role="user", content="hello world")])
         """
 
-    def guard_text(  # type: ignore[misc]
+    def guard_text(
         self,
         text: str | None = None,
         *,
-        messages: _T | None = None,
+        messages: Sequence[Message] | None = None,
         recipe: str | None = None,
         debug: bool | None = None,
         overrides: Overrides | None = None,
         log_fields: LogFields | None = None,
-    ) -> PangeaResponse[TextGuardResult[None]]:
+    ) -> PangeaResponse[TextGuardResult]:
         """
         Text Guard for scanning LLM inputs and outputs
 

packages/pangea-sdk/tests/integration/asyncio/test_ai_guard.py

@@ -4,7 +4,7 @@
 
 from pangea import PangeaConfig
 from pangea.asyncio.services import AIGuardAsync
-from pangea.services.ai_guard import LogFields
+from pangea.services.ai_guard import LogFields, Message
 from pangea.tools import TestEnvironment, get_test_domain, get_test_token, logger_set_pangea_config
 from tests.test_tools import load_test_environment
 
@@ -42,7 +42,7 @@ async def test_text_guard(self) -> None:
 
     async def test_text_guard_messages(self) -> None:
         response = await self.client.guard_text(
-            messages=[{"role": "user", "content": "hello world"}], log_fields=LogFields(source="Acme Wizard")
+            messages=[Message(role="user", content="hello world")], log_fields=LogFields(source="Acme Wizard")
         )
         assert response.status == "Success"
         assert response.result

packages/pangea-sdk/tests/integration/test_ai_guard.py

@@ -4,7 +4,7 @@
 
 from pangea import PangeaConfig
 from pangea.services import AIGuard
-from pangea.services.ai_guard import LogFields
+from pangea.services.ai_guard import LogFields, Message
 from pangea.tools import TestEnvironment, get_test_domain, get_test_token, logger_set_pangea_config
 from tests.test_tools import load_test_environment
 
@@ -35,7 +35,7 @@ def test_text_guard(self) -> None:
 
     def test_text_guard_messages(self) -> None:
         response = self.client.guard_text(
-            messages=[{"role": "user", "content": "hello world"}], log_fields=LogFields(source="Acme Wizard")
+            messages=[Message(role="user", content="hello world")], log_fields=LogFields(source="Acme Wizard")
         )
         assert response.status == "Success"
         assert response.result

packages/pangea-sdk/tests/integration2/test_ai_guard.py

@@ -8,7 +8,7 @@
 from pangea import PangeaConfig
 from pangea.asyncio.services.ai_guard import AIGuardAsync
 from pangea.services import AIGuard
-from pangea.services.ai_guard import LogFields, TextGuardResult
+from pangea.services.ai_guard import LogFields, Message, TextGuardResult
 
 from ..utils import assert_matches_type
 
@@ -34,7 +34,7 @@ def test_text_guard(self, client: AIGuard) -> None:
 
     def test_text_guard_messages(self, client: AIGuard) -> None:
         response = client.guard_text(
-            messages=[{"role": "user", "content": "hello world"}],
+            messages=[Message(role="user", content="hello world")],
             debug=False,
             log_fields=LogFields(source="Acme Wizard"),
         )
@@ -52,7 +52,7 @@ async def test_text_guard(self, async_client: AIGuardAsync) -> None:
 
     async def test_text_guard_messages(self, async_client: AIGuardAsync) -> None:
         response = await async_client.guard_text(
-            messages=[{"role": "user", "content": "hello world"}],
+            messages=[Message(role="user", content="hello world")],
             debug=False,
             log_fields=LogFields(source="Acme Wizard"),
         )