pangea-sdk 6.5.0b1__py3-none-any.whl → 6.7.0__py3-none-any.whl

pangea/services/ai_guard.py

@@ -1,13 +1,13 @@
 from __future__ import annotations
 
 from collections.abc import Mapping, Sequence
-from typing import Annotated, Any, Generic, Literal, Optional, Union, overload
+from typing import Annotated, Any, Generic, Literal, Optional, overload
 
 from pydantic import BaseModel, ConfigDict, Field, RootModel
-from typing_extensions import TypeAlias, TypedDict, TypeVar
 
+from pangea._typing import T
 from pangea.config import PangeaConfig
-from pangea.response import APIRequestModel, APIResponseModel, PangeaDateTime, PangeaResponse, PangeaResponseResult
+from pangea.response import APIRequestModel, APIResponseModel, PangeaResponse, PangeaResponseResult
 from pangea.services.base import ServiceBase
 
 # This is named "prompt injection" in the API spec even though it is also used
@@ -22,10 +22,15 @@ PiiEntityAction = Literal["disabled", "report", "block", "mask", "partial_maskin
 
 
 class Message(APIRequestModel):
-    role: str
+    role: Optional[str] = None
     content: str
 
 
+class McpToolsMessage(APIRequestModel):
+    role: Literal["tools"]
+    content: list[dict[str, Any]]
+
+
 class CodeDetectionOverride(APIRequestModel):
     disabled: Optional[bool] = None
     action: Optional[Literal["report", "block"]] = None
@@ -277,29 +282,31 @@ class CodeDetectionResult(APIResponseModel):
     """The action taken by this Detector"""
 
 
-_T = TypeVar("_T")
-
-
-class TextGuardDetector(APIResponseModel, Generic[_T]):
+class GuardDetector(APIResponseModel, Generic[T]):
     detected: Optional[bool] = None
-    data: Optional[_T] = None
+    data: Optional[T] = None
 
 
 class TextGuardDetectors(APIResponseModel):
-    code_detection: Optional[TextGuardDetector[CodeDetectionResult]] = None
-    competitors: Optional[TextGuardDetector[object]] = None
-    custom_entity: Optional[TextGuardDetector[object]] = None
-    gibberish: Optional[TextGuardDetector[object]] = None
-    hardening: Optional[TextGuardDetector[object]] = None
-    language_detection: Optional[TextGuardDetector[LanguageDetectionResult]] = None
-    malicious_entity: Optional[TextGuardDetector[MaliciousEntityResult]] = None
-    pii_entity: Optional[TextGuardDetector[PiiEntityResult]] = None
-    profanity_and_toxicity: Optional[TextGuardDetector[object]] = None
-    prompt_injection: Optional[TextGuardDetector[PromptInjectionResult]] = None
-    secrets_detection: Optional[TextGuardDetector[SecretsEntityResult]] = None
-    selfharm: Optional[TextGuardDetector[object]] = None
-    sentiment: Optional[TextGuardDetector[object]] = None
-    topic: Optional[TextGuardDetector[TopicDetectionResult]] = None
+    code_detection: Optional[GuardDetector[CodeDetectionResult]] = None
+    competitors: Optional[GuardDetector[object]] = None
+    custom_entity: Optional[GuardDetector[object]] = None
+    gibberish: Optional[GuardDetector[object]] = None
+    hardening: Optional[GuardDetector[object]] = None
+    language_detection: Optional[GuardDetector[LanguageDetectionResult]] = None
+    malicious_entity: Optional[GuardDetector[MaliciousEntityResult]] = None
+    pii_entity: Optional[GuardDetector[PiiEntityResult]] = None
+    profanity_and_toxicity: Optional[GuardDetector[object]] = None
+    prompt_injection: Optional[GuardDetector[PromptInjectionResult]] = None
+    secrets_detection: Optional[GuardDetector[SecretsEntityResult]] = None
+    selfharm: Optional[GuardDetector[object]] = None
+    sentiment: Optional[GuardDetector[object]] = None
+    topic: Optional[GuardDetector[TopicDetectionResult]] = None
+
+
+class PromptMessage(APIResponseModel):
+    role: str
+    content: str
 
 
 class TextGuardResult(PangeaResponseResult):
@@ -318,7 +325,7 @@ class TextGuardResult(PangeaResponseResult):
     unredact.
     """
 
-    prompt_messages: Optional[object] = None
+    prompt_messages: Optional[list[PromptMessage]] = None
     """Updated structured prompt, if applicable."""
 
     prompt_text: Optional[str] = None
@@ -331,537 +338,119 @@ class TextGuardResult(PangeaResponseResult):
     """Whether or not the original input was transformed."""
 
 
-class GuardDetectors(APIResponseModel):
-    code: Optional[object] = None
-    competitors: Optional[object] = None
-    confidential_and_pii_entity: Optional[object] = None
-    custom_entity: Optional[object] = None
-    language: Optional[object] = None
-    malicious_entity: Optional[object] = None
-    malicious_prompt: Optional[object] = None
-    prompt_hardening: Optional[object] = None
-    secret_and_key_entity: Optional[object] = None
-    topic: Optional[object] = None
-
-
-class GuardResult(PangeaResponseResult):
-    detectors: GuardDetectors
-    """Result of the recipe analyzing and input prompt."""
-
-    access_rules: Optional[object] = None
-    """Result of the recipe evaluating configured rules"""
-
-    blocked: Optional[bool] = None
-    """Whether or not the prompt triggered a block detection."""
-
-    fpe_context: Optional[str] = None
-    """
-    If an FPE redaction method returned results, this will be the context passed
-    to unredact.
-    """
-
-    input_token_count: Optional[float] = None
-    """Number of tokens counted in the input"""
-
-    output: Optional[object] = None
-    """Updated structured prompt."""
-
-    output_token_count: Optional[float] = None
-    """Number of tokens counted in the output"""
-
-    recipe: Optional[str] = None
-    """The Recipe that was used."""
-
-    transformed: Optional[bool] = None
-    """Whether or not the original input was transformed."""
-
-
-class Areas(BaseModel):
-    model_config = ConfigDict(extra="forbid")
-
-    text_guard: bool
-
-
-class AuditDataActivityConfig(BaseModel):
-    model_config = ConfigDict(extra="forbid")
-
-    enabled: bool
-    audit_service_config_id: str
-    areas: Areas
-
-
-class PromptGuard(BaseModel):
-    model_config = ConfigDict(extra="forbid")
-
-    enabled: Optional[bool] = None
-    config_id: Optional[str] = None
-    confidence_threshold: Optional[float] = None
-
+class Tool(RootModel[str]):
+    root: Annotated[str, Field(min_length=1)]
+    """Tool name"""
 
-class IpIntel(BaseModel):
-    model_config = ConfigDict(extra="forbid")
 
-    enabled: Optional[bool] = None
-    config_id: Optional[str] = None
-    reputation_provider: Optional[str] = None
-    risk_threshold: Optional[float] = None
+class McpTool(APIRequestModel):
+    server_name: Annotated[str, Field(min_length=1)]
+    """MCP server name"""
 
+    tools: Annotated[list[Tool], Field(min_length=1)]
 
-class UserIntel(BaseModel):
-    model_config = ConfigDict(extra="forbid")
 
-    enabled: Optional[bool] = None
-    config_id: Optional[str] = None
-    breach_provider: Optional[str] = None
-
-
-class UrlIntel(BaseModel):
-    model_config = ConfigDict(extra="forbid")
-
-    enabled: Optional[bool] = None
-    config_id: Optional[str] = None
-    reputation_provider: Optional[str] = None
-    risk_threshold: Optional[float] = None
-
-
-class DomainIntel(BaseModel):
-    model_config = ConfigDict(extra="forbid")
-
-    enabled: Optional[bool] = None
-    config_id: Optional[str] = None
-    reputation_provider: Optional[str] = None
-    risk_threshold: Optional[float] = None
-
-
-class FileScan(BaseModel):
-    model_config = ConfigDict(extra="forbid")
-
-    enabled: Optional[bool] = None
-    config_id: Optional[str] = None
-    scan_provider: Optional[str] = None
-    risk_threshold: Optional[float] = None
-
-
-class Redact(BaseModel):
-    model_config = ConfigDict(extra="forbid")
-
-    enabled: Optional[bool] = None
-    config_id: Optional[str] = None
-
-
-class Vault(BaseModel):
-    model_config = ConfigDict(extra="forbid")
-
-    config_id: Optional[str] = None
-
-
-class Lingua(BaseModel):
-    model_config = ConfigDict(extra="forbid")
+class ExtraInfo(BaseModel):
+    """(AIDR) Logging schema."""
 
-    enabled: Optional[bool] = None
+    # Additional properties are allowed here.
+    model_config = ConfigDict(extra="allow")
 
+    app_name: Optional[str] = None
+    """Name of source application/agent."""
 
-class Code(BaseModel):
-    model_config = ConfigDict(extra="forbid")
+    app_group: Optional[str] = None
+    """The group of source application/agent."""
 
-    enabled: Optional[bool] = None
+    app_version: Optional[str] = None
+    """Version of the source application/agent."""
 
+    actor_name: Optional[str] = None
+    """Name of subject actor/service account."""
 
-class ConnectionsConfig(BaseModel):
-    model_config = ConfigDict(extra="forbid")
+    actor_group: Optional[str] = None
+    """The group of subject actor."""
 
-    prompt_guard: Optional[PromptGuard] = None
-    ip_intel: Optional[IpIntel] = None
-    user_intel: Optional[UserIntel] = None
-    url_intel: Optional[UrlIntel] = None
-    domain_intel: Optional[DomainIntel] = None
-    file_scan: Optional[FileScan] = None
-    redact: Optional[Redact] = None
-    vault: Optional[Vault] = None
-    lingua: Optional[Lingua] = None
-    code: Optional[Code] = None
+    source_region: Optional[str] = None
+    """Geographic region or data center."""
 
+    sub_tenant: Optional[str] = None
+    """Sub tenant of the user or organization"""
+    mcp_tools: Optional[Sequence[McpTool]] = None
 
-class PartialMasking(BaseModel):
-    masking_type: Optional[Literal["unmask", "mask"]] = "unmask"
-    unmasked_from_left: Annotated[Optional[int], Field(ge=0)] = None
-    unmasked_from_right: Annotated[Optional[int], Field(ge=0)] = None
-    masked_from_left: Annotated[Optional[int], Field(ge=0)] = None
-    masked_from_right: Annotated[Optional[int], Field(ge=0)] = None
-    chars_to_ignore: Optional[list[CharsToIgnoreItem]] = None
-    masking_char: Annotated[Optional[str], Field(max_length=1, min_length=1)] = "*"
+    """Each item groups tools for a given MCP server."""
 
 
-class RuleRedactionConfig1(APIResponseModel):
-    redaction_type: Literal[
-        "mask",
-        "partial_masking",
-        "replacement",
-        "hash",
-        "detect_only",
-        "fpe",
-        "mask",
-        "detect_only",
-    ]
-    """Redaction method to apply for this rule"""
-    redaction_value: Optional[str] = None
-    partial_masking: Optional[PartialMasking] = None
-    hash: Optional[Hash] = None
-    fpe_alphabet: Optional[
-        Literal[
-            "numeric",
-            "alphalower",
-            "alphaupper",
-            "alpha",
-            "alphanumericlower",
-            "alphanumericupper",
-            "alphanumeric",
-        ]
-    ] = None
-
-
-class PartialMasking1(BaseModel):
-    masking_type: Optional[Literal["unmask", "mask"]] = "unmask"
-    unmasked_from_left: Annotated[Optional[int], Field(ge=0)] = None
-    unmasked_from_right: Annotated[Optional[int], Field(ge=0)] = None
-    masked_from_left: Annotated[Optional[int], Field(ge=0)] = None
-    masked_from_right: Annotated[Optional[int], Field(ge=0)] = None
-    chars_to_ignore: Optional[list[CharsToIgnoreItem]] = None
-    masking_char: Annotated[Optional[str], Field(max_length=1, min_length=1)] = "*"
-
-
-class RuleRedactionConfig2(BaseModel):
-    model_config = ConfigDict(extra="forbid")
-
-    redaction_type: Literal["replacement"]
-    redaction_value: str
-    partial_masking: Optional[PartialMasking1] = None
-    hash: Optional[Hash] = None
-    fpe_alphabet: Optional[
-        Literal[
-            "numeric",
-            "alphalower",
-            "alphaupper",
-            "alpha",
-            "alphanumericlower",
-            "alphanumericupper",
-            "alphanumeric",
-        ]
-    ] = None
-
-
-class PartialMasking2(BaseModel):
-    masking_type: Optional[Literal["unmask", "mask"]] = "unmask"
-    unmasked_from_left: Annotated[Optional[int], Field(ge=0)] = None
-    unmasked_from_right: Annotated[Optional[int], Field(ge=0)] = None
-    masked_from_left: Annotated[Optional[int], Field(ge=0)] = None
-    masked_from_right: Annotated[Optional[int], Field(ge=0)] = None
-    chars_to_ignore: Optional[list[CharsToIgnoreItem]] = None
-    masking_char: Annotated[Optional[str], Field(max_length=1, min_length=1)] = "*"
-
-
-class RuleRedactionConfig3(BaseModel):
-    model_config = ConfigDict(extra="forbid")
-
-    redaction_type: Literal["partial_masking"]
-    redaction_value: str
-    partial_masking: PartialMasking2
-    hash: Optional[Hash] = None
-    fpe_alphabet: Optional[
-        Literal[
-            "numeric",
-            "alphalower",
-            "alphaupper",
-            "alpha",
-            "alphanumericlower",
-            "alphanumericupper",
-            "alphanumeric",
-        ]
-    ] = None
-
-
-class PartialMasking3(BaseModel):
-    masking_type: Optional[Literal["unmask", "mask"]] = "unmask"
-    unmasked_from_left: Annotated[Optional[int], Field(ge=0)] = None
-    unmasked_from_right: Annotated[Optional[int], Field(ge=0)] = None
-    masked_from_left: Annotated[Optional[int], Field(ge=0)] = None
-    masked_from_right: Annotated[Optional[int], Field(ge=0)] = None
-    chars_to_ignore: Optional[list[CharsToIgnoreItem]] = None
-    masking_char: Annotated[Optional[str], Field(max_length=1, min_length=1)] = "*"
-
-
-class RuleRedactionConfig4(BaseModel):
-    model_config = ConfigDict(extra="forbid")
-
-    redaction_type: Literal["hash"]
-    redaction_value: str
-    partial_masking: PartialMasking3
-    hash: Optional[Hash] = None
-    fpe_alphabet: Optional[
-        Literal[
-            "numeric",
-            "alphalower",
-            "alphaupper",
-            "alpha",
-            "alphanumericlower",
-            "alphanumericupper",
-            "alphanumeric",
-        ]
-    ] = None
-
-
-class CharsToIgnoreItem(RootModel[str]):
-    root: Annotated[str, Field(max_length=1, min_length=1)]
-
-
-class PartialMasking4(BaseModel):
-    masking_type: Optional[Literal["unmask", "mask"]] = "unmask"
-    unmasked_from_left: Annotated[Optional[int], Field(ge=0)] = None
-    unmasked_from_right: Annotated[Optional[int], Field(ge=0)] = None
-    masked_from_left: Annotated[Optional[int], Field(ge=0)] = None
-    masked_from_right: Annotated[Optional[int], Field(ge=0)] = None
-    chars_to_ignore: Optional[list[CharsToIgnoreItem]] = None
-    masking_char: Annotated[Optional[str], Field(max_length=1, min_length=1)] = "*"
-
-
-class Hash(BaseModel):
-    hash_type: Literal["md5", "sha256"]
-    """The type of hashing algorithm"""
-
-
-class RuleRedactionConfig5(BaseModel):
-    model_config = ConfigDict(extra="forbid")
-
-    redaction_type: Literal["fpe"]
-    redaction_value: str
-    partial_masking: PartialMasking4
-    hash: Optional[Hash] = None
-    fpe_alphabet: Optional[
-        Literal[
-            "numeric",
-            "alphalower",
-            "alphaupper",
-            "alpha",
-            "alphanumericlower",
-            "alphanumericupper",
-            "alphanumeric",
-        ]
-    ] = None
-
-
-class Rule(APIResponseModel):
-    redact_rule_id: str
+class AccessRuleResult(APIResponseModel):
     """
-    Identifier of the redaction rule to apply. This should match a rule defined
-    in the [Redact service](https://pangea.cloud/docs/redact/using-redact/using-redact).
+    Details about the evaluation of a single rule, including whether it matched,
+    the action to take, the rule name, and optional debugging information.
     """
-    redaction: Union[
-        RuleRedactionConfig1,
-        RuleRedactionConfig2,
-        RuleRedactionConfig3,
-        RuleRedactionConfig4,
-        RuleRedactionConfig5,
-    ]
-    """
-    Configuration for the redaction method applied to detected values.
 
-    Each rule supports one redaction type, such as masking, replacement,
-    hashing, Format-Preserving Encryption (FPE), or detection-only mode.
-    Additional parameters may be required depending on the selected redaction
-    type.
+    matched: bool
+    """Whether this rule's logic evaluated to true for the input."""
 
-    For more details, see the [AI Guard Recipe Actions](https://pangea.cloud/docs/ai-guard/recipes#actions)
-    documentation.
-    """
-    block: Optional[bool] = None
-    """
-    If `true`, indicates that further processing should be stopped when this
-    rule is triggered
-    """
-    disabled: Optional[bool] = None
-    """
-    If `true`, disables this specific rule even if the detector is enabled
-    """
-    reputation_check: Optional[bool] = None
-    """
-    If `true`, performs a reputation check using the configured intel provider.
-    Applies to the Malicious Entity detector when using IP, URL, or Domain Intel
-    services.
-    """
-    transform_if_malicious: Optional[bool] = None
+    action: str
     """
-    If `true`, applies redaction or transformation when the detected value is
-    determined to be malicious by intel analysis
+    The action resulting from the rule evaluation. One of 'allowed', 'blocked',
+    or 'reported'.
     """
 
+    name: str
+    """A human-readable name for the rule."""
 
-class Settings(BaseModel):
-    rules: Optional[list[Rule]] = None
-
-
-class DetectorSetting(BaseModel):
-    model_config = ConfigDict(extra="forbid")
-
-    detector_name: str
-    state: Literal["disabled", "enabled"]
-    settings: Settings
+    logic: Optional[dict[str, Any]] = None
+    """The JSON logic expression evaluated for this rule."""
 
+    attributes: Optional[dict[str, Any]] = None
+    """The input attribute values that were available during rule evaluation."""
 
-class RedactConnectorSettings(BaseModel):
-    fpe_tweak_vault_secret_id: Optional[str] = None
 
+class GuardDetectors(APIResponseModel):
+    """Result of the recipe analyzing and input prompt."""
 
-class ConnectorSettings(BaseModel):
-    model_config = ConfigDict(extra="forbid")
+    code: Optional[GuardDetector[CodeDetectionResult]] = None
+    competitors: Optional[GuardDetector[object]] = None
+    confidential_and_pii_entity: Optional[GuardDetector[PiiEntityResult]] = None
+    custom_entity: Optional[GuardDetector[object]] = None
+    language: Optional[GuardDetector[LanguageDetectionResult]] = None
+    malicious_entity: Optional[GuardDetector[MaliciousEntityResult]] = None
+    malicious_prompt: Optional[GuardDetector[PromptInjectionResult]] = None
+    prompt_hardening: Optional[GuardDetector[object]] = None
+    secret_and_key_entity: Optional[GuardDetector[SecretsEntityResult]] = None
+    topic: Optional[GuardDetector[TopicDetectionResult]] = None
 
-    redact: Optional[RedactConnectorSettings] = None
 
+class GuardResult(PangeaResponseResult):
+    output: Optional[dict[str, Any]] = None
+    """Updated structured prompt."""
 
-class AccessRuleSettings(APIResponseModel):
-    """
-    Configuration for an individual access rule used in an AI Guard recipe. Each
-    rule defines its matching logic and the action to apply when the logic
-    evaluates to true.
-    """
+    blocked: Optional[bool] = None
+    """Whether or not the prompt triggered a block detection."""
 
-    rule_key: Annotated[str, Field(pattern="^([a-zA-Z0-9_][a-zA-Z0-9/|_]*)$")]
-    """
-    Unique identifier for this rule. Should be user-readable and consistent
-    across recipe updates.
-    """
-    name: str
-    """Display label for the rule shown in user interfaces."""
-    state: Literal["block", "report"]
-    """
-    Action to apply if the rule matches. Use 'block' to stop further processing
-    or 'report' to simply log the match.
-    """
+    transformed: Optional[bool] = None
+    """Whether or not the original input was transformed."""
 
+    recipe: Optional[str] = None
+    """The Recipe that was used."""
 
-class RecipeConfig(APIResponseModel):
-    name: str
-    """Human-readable name of the recipe"""
-    description: str
-    """Detailed description of the recipe's purpose or use case"""
-    version: Optional[str] = "v1"
-    """Optional version identifier for the recipe. Can be used to track changes."""
-    detectors: Optional[list[DetectorSetting]] = None
-    """Setting for Detectors"""
-    access_rules: Optional[list[AccessRuleSettings]] = None
-    """Configuration for access rules used in an AI Guard recipe."""
-    connector_settings: Optional[ConnectorSettings] = None
-
-
-class ServiceConfig(PangeaResponseResult):
-    id: Optional[str] = None
-    """ID of an AI Guard service configuration"""
-    name: Optional[str] = None
-    """Human-readable name of the AI Guard service configuration"""
-    audit_data_activity: Optional[AuditDataActivityConfig] = None
-    connections: Optional[ConnectionsConfig] = None
-    recipes: Optional[dict[str, RecipeConfig]] = None
-
-
-class ServiceConfigFilter(BaseModel):
-    model_config = ConfigDict(extra="forbid")
-
-    id: Optional[str] = None
-    """
-    Only records where id equals this value.
-    """
-    id__contains: Optional[list[str]] = None
-    """
-    Only records where id includes each substring.
-    """
-    id__in: Optional[list[str]] = None
-    """
-    Only records where id equals one of the provided substrings.
-    """
-    created_at: Optional[PangeaDateTime] = None
-    """
-    Only records where created_at equals this value.
-    """
-    created_at__gt: Optional[PangeaDateTime] = None
-    """
-    Only records where created_at is greater than this value.
-    """
-    created_at__gte: Optional[PangeaDateTime] = None
-    """
-    Only records where created_at is greater than or equal to this value.
-    """
-    created_at__lt: Optional[PangeaDateTime] = None
-    """
-    Only records where created_at is less than this value.
-    """
-    created_at__lte: Optional[PangeaDateTime] = None
-    """
-    Only records where created_at is less than or equal to this value.
-    """
-    updated_at: Optional[PangeaDateTime] = None
-    """
-    Only records where updated_at equals this value.
-    """
-    updated_at__gt: Optional[PangeaDateTime] = None
-    """
-    Only records where updated_at is greater than this value.
-    """
-    updated_at__gte: Optional[PangeaDateTime] = None
-    """
-    Only records where updated_at is greater than or equal to this value.
-    """
-    updated_at__lt: Optional[PangeaDateTime] = None
-    """
-    Only records where updated_at is less than this value.
-    """
-    updated_at__lte: Optional[PangeaDateTime] = None
-    """
-    Only records where updated_at is less than or equal to this value.
-    """
+    detectors: GuardDetectors
+    """Result of the recipe analyzing and input prompt."""
 
+    access_rules: Optional[dict[str, AccessRuleResult]] = None
+    """Result of the recipe evaluating configured rules"""
 
-class ServiceConfigsPage(PangeaResponseResult):
-    count: Optional[int] = None
-    """The total number of service configs matched by the list request."""
-    last: Optional[str] = None
+    fpe_context: Optional[str] = None
     """
-    Used to fetch the next page of the current listing when provided in a
-    repeated request's last parameter.
+    If an FPE redaction method returned results, this will be the context passed
+    to unredact.
     """
-    items: Optional[list[ServiceConfig]] = None
-
-
-class ExtraInfoTyped(TypedDict, total=False):
-    """(AIDR) Logging schema."""
-
-    app_name: str
-    """Name of source application."""
-
-    app_group: str
-    """The group of source application."""
-
-    app_version: str
-    """Version of the source application."""
-
-    actor_name: str
-    """Name of subject actor."""
-
-    actor_group: str
-    """The group of subject actor."""
-
-    source_region: str
-    """Geographic region or data center."""
-
-    data_sensitivity: str
-    """Sensitivity level of data involved"""
-
-    customer_tier: str
-    """Tier of the user or organization"""
-
-    use_case: str
-    """Business-specific use case"""
 
+    input_token_count: Optional[float] = None
+    """Number of tokens counted in the input"""
 
-ExtraInfo: TypeAlias = Union[ExtraInfoTyped, dict[str, object]]
+    output_token_count: Optional[float] = None
+    """Number of tokens counted in the output"""
 
 
 class AIGuard(ServiceBase):
@@ -938,11 +527,12 @@ class AIGuard(ServiceBase):
     def guard_text(
         self,
         *,
-        messages: Sequence[Message],
+        messages: Sequence[Message | McpToolsMessage],
         recipe: str | None = None,
         debug: bool | None = None,
         overrides: Overrides | None = None,
         log_fields: LogFields | None = None,
+        only_relevant_content: bool = False,
     ) -> PangeaResponse[TextGuardResult]:
         """
         Guard LLM input and output text
@@ -965,6 +555,8 @@ class AIGuard(ServiceBase):
             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.
+            only_relevant_content: Whether or not to only send relevant content
+                to AI Guard.
 
         Examples:
             response = ai_guard.guard_text(messages=[Message(role="user", content="hello world")])
@@ -974,11 +566,12 @@ class AIGuard(ServiceBase):
         self,
         text: str | None = None,
         *,
-        messages: Sequence[Message] | None = None,
+        messages: Sequence[Message | McpToolsMessage] | None = None,
         debug: bool | None = None,
         log_fields: LogFields | None = None,
         overrides: Overrides | None = None,
         recipe: str | None = None,
+        only_relevant_content: bool = False,
     ) -> PangeaResponse[TextGuardResult]:
         """
         Guard LLM input and output text
@@ -1004,6 +597,8 @@ class AIGuard(ServiceBase):
             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.
+            only_relevant_content: Whether or not to only send relevant content
+                to AI Guard.
 
         Examples:
             response = ai_guard.guard_text("text")
@@ -1012,7 +607,11 @@ class AIGuard(ServiceBase):
         if text is not None and messages is not None:
             raise ValueError("Exactly one of `text` or `messages` must be given")
 
-        return self.request.post(
+        if only_relevant_content and messages is not None:
+            original_messages = messages
+            messages, original_indices = get_relevant_content(messages)
+
+        response = self.request.post(
             "v1/text/guard",
             TextGuardResult,
             data={
@@ -1025,6 +624,13 @@ class AIGuard(ServiceBase):
             },
         )
 
+        if only_relevant_content and response.result and response.result.prompt_messages:
+            response.result.prompt_messages = patch_messages(
+                original_messages, original_indices, response.result.prompt_messages
+            )  # type: ignore[assignment]
+
+        return response
+
     def guard(
         self,
         input: Mapping[str, Any],
@@ -1042,8 +648,8 @@ class AIGuard(ServiceBase):
         source_ip: str | None = None,
         source_location: str | None = None,
         tenant_id: str | None = None,
-        event_type: Literal["input", "output"] | None = None,
-        sensor_instance_id: str | None = None,
+        event_type: Literal["input", "output", "tool_input", "tool_output", "tool_listing"] | None = None,
+        collector_instance_id: str | None = None,
         extra_info: ExtraInfo | None = None,
         count_tokens: bool | None = None,
     ) -> PangeaResponse[GuardResult]:
@@ -1053,7 +659,7 @@ class AIGuard(ServiceBase):
         Analyze and redact content to avoid manipulation of the model, addition
         of malicious content, and other undesirable data transfers.
 
-        OperationId: ai_guard_post_v1beta_guard
+        OperationId: ai_guard_post_v1_guard
 
         Args:
             input: 'messages' (required) contains Prompt content and role array
@@ -1072,12 +678,12 @@ class AIGuard(ServiceBase):
             source_location: Location of user or app or agent.
             tenant_id: For gateway-like integrations with multi-tenant support.
             event_type: (AIDR) Event Type.
-            sensor_instance_id: (AIDR) sensor instance id.
+            collector_instance_id: (AIDR) collector instance id.
             extra_info: (AIDR) Logging schema.
             count_tokens: Provide input and output token count.
         """
         return self.request.post(
-            "v1beta/guard",
+            "v1/guard",
             GuardResult,
             data={
                 "input": input,
@@ -1095,86 +701,63 @@ class AIGuard(ServiceBase):
                 "source_location": source_location,
                 "tenant_id": tenant_id,
                 "event_type": event_type,
-                "sensor_instance_id": sensor_instance_id,
+                "collector_instance_id": collector_instance_id,
                 "extra_info": extra_info,
                 "count_tokens": count_tokens,
             },
         )
 
-    def get_service_config(self, id: str) -> PangeaResponse[ServiceConfig]:
-        """
-        OperationId: ai_guard_post_v1beta_config
-        """
-        return self.request.post("v1beta/config", data={"id": id}, result_class=ServiceConfig)
 
-    def create_service_config(
-        self,
-        name: str,
-        *,
-        id: str | None = None,
-        audit_data_activity: AuditDataActivityConfig | None = None,
-        connections: ConnectionsConfig | None = None,
-        recipes: Mapping[str, RecipeConfig] | None = None,
-    ) -> PangeaResponse[ServiceConfig]:
-        """
-        OperationId: ai_guard_post_v1beta_config_create
-        """
-        return self.request.post(
-            "v1beta/config/create",
-            data={
-                "name": name,
-                "id": id,
-                "audit_data_activity": audit_data_activity,
-                "connections": connections,
-                "recipes": recipes,
-            },
-            result_class=ServiceConfig,
-        )
+def get_relevant_content(
+    messages: Sequence[Message | McpToolsMessage],
+) -> tuple[list[Message | McpToolsMessage], list[int]]:
+    """
+    Returns relevant messages and their indices in the original list.
 
-    def update_service_config(
-        self,
-        id: str,
-        name: str,
-        *,
-        audit_data_activity: AuditDataActivityConfig | None = None,
-        connections: ConnectionsConfig | None = None,
-        recipes: Mapping[str, RecipeConfig] | None = None,
-    ) -> PangeaResponse[ServiceConfig]:
-        """
-        OperationId: ai_guard_post_v1beta_config_update
-        """
-        return self.request.post(
-            "v1beta/config/update",
-            data={
-                "id": id,
-                "name": name,
-                "audit_data_activity": audit_data_activity,
-                "connections": connections,
-                "recipes": recipes,
-            },
-            result_class=ServiceConfig,
-        )
+    1, If last message is "assistant", then the relevant messages are all system
+      messages that come before it, plus that last assistant message.
+    2. Else, find the last "assistant" message. Then the relevant messages are
+      all system messages that come before it, and all messages that come after
+      it.
+    """
 
-    def delete_service_config(self, id: str) -> PangeaResponse[ServiceConfig]:
-        """
-        OperationId: ai_guard_post_v1beta_config_delete
-        """
-        return self.request.post("v1beta/config/delete", data={"id": id}, result_class=ServiceConfig)
+    if len(messages) == 0:
+        return [], []
 
-    def list_service_configs(
-        self,
-        *,
-        filter: ServiceConfigFilter | None = None,
-        last: str | None = None,
-        order: Literal["asc", "desc"] | None = None,
-        order_by: Literal["id", "created_at", "updated_at"] | None = None,
-        size: int | None = None,
-    ) -> PangeaResponse[ServiceConfigsPage]:
-        """
-        OperationId: ai_guard_post_v1beta_config_list
-        """
-        return self.request.post(
-            "v1beta/config/list",
-            data={"filter": filter, "last": last, "order": order, "order_by": order_by, "size": size},
-            result_class=ServiceConfigsPage,
-        )
+    system_messages = [msg for msg in messages if msg.role == "system"]
+    system_indices = [i for i, msg in enumerate(messages) if msg.role == "system"]
+
+    # If the last message is assistant, then return all system messages and that
+    # assistant message.
+    if messages[-1].role == "assistant":
+        return system_messages + [messages[-1]], system_indices + [len(messages) - 1]
+
+    # Otherwise, work backwards until we find the last assistant message, then
+    # return all messages after that.
+    last_assistant_index = -1
+    for i in range(len(messages) - 1, -1, -1):
+        if messages[i].role == "assistant":
+            last_assistant_index = i
+            break
+
+    relevant_messages = []
+    indices = []
+    for i, msg in enumerate(messages):
+        if msg.role == "system" or i > last_assistant_index:
+            relevant_messages.append(msg)
+            indices.append(i)
+
+    return relevant_messages, indices
+
+
+def patch_messages(
+    original: Sequence[Message | McpToolsMessage],
+    original_indices: list[int],
+    transformed: Sequence[PromptMessage],
+) -> list[Message | McpToolsMessage | PromptMessage]:
+    if len(original) == len(transformed):
+        return list(transformed)
+
+    return [
+        transformed[original_indices.index(i)] if i in original_indices else orig for i, orig in enumerate(original)
+    ]