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

pangea/services/ai_guard.py

@@ -1,13 +1,11 @@
 from __future__ import annotations
 
-from collections.abc import Mapping, Sequence
-from typing import Annotated, Any, Generic, Literal, Optional, Union, overload
-
-from pydantic import BaseModel, ConfigDict, Field, RootModel
-from typing_extensions import TypeAlias, TypedDict, TypeVar
+from collections.abc import Sequence
+from typing import Any, Generic, Literal, Optional, overload
 
+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 +20,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,12 +280,9 @@ class CodeDetectionResult(APIResponseModel):
     """The action taken by this Detector"""
 
 
-_T = TypeVar("_T")
-
-
-class TextGuardDetector(APIResponseModel, Generic[_T]):
+class TextGuardDetector(APIResponseModel, Generic[T]):
     detected: Optional[bool] = None
-    data: Optional[_T] = None
+    data: Optional[T] = None
 
 
 class TextGuardDetectors(APIResponseModel):
@@ -302,6 +302,11 @@ class TextGuardDetectors(APIResponseModel):
     topic: Optional[TextGuardDetector[TopicDetectionResult]] = None
 
 
+class PromptMessage(APIResponseModel):
+    role: str
+    content: str
+
+
 class TextGuardResult(PangeaResponseResult):
     detectors: TextGuardDetectors
     """Result of the recipe analyzing and input prompt."""
@@ -318,7 +323,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,539 +336,6 @@ 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 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 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")
-
-    enabled: Optional[bool] = None
-
-
-class Code(BaseModel):
-    model_config = ConfigDict(extra="forbid")
-
-    enabled: Optional[bool] = None
-
-
-class ConnectionsConfig(BaseModel):
-    model_config = ConfigDict(extra="forbid")
-
-    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
-
-
-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)] = "*"
-
-
-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
-    """
-    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).
-    """
-    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.
-
-    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
-    """
-    If `true`, applies redaction or transformation when the detected value is
-    determined to be malicious by intel analysis
-    """
-
-
-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
-
-
-class RedactConnectorSettings(BaseModel):
-    fpe_tweak_vault_secret_id: Optional[str] = None
-
-
-class ConnectorSettings(BaseModel):
-    model_config = ConfigDict(extra="forbid")
-
-    redact: Optional[RedactConnectorSettings] = None
-
-
-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.
-    """
-
-    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.
-    """
-
-
-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.
-    """
-
-
-class ServiceConfigsPage(PangeaResponseResult):
-    count: Optional[int] = None
-    """The total number of service configs matched by the list request."""
-    last: Optional[str] = None
-    """
-    Used to fetch the next page of the current listing when provided in a
-    repeated request's last parameter.
-    """
-    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"""
-
-
-ExtraInfo: TypeAlias = Union[ExtraInfoTyped, dict[str, object]]
-
-
 class AIGuard(ServiceBase):
     """AI Guard service client.
 
@@ -938,11 +410,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 +438,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 +449,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 +480,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 +490,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,156 +507,64 @@ class AIGuard(ServiceBase):
             },
         )
 
-    def guard(
-        self,
-        input: Mapping[str, Any],
-        *,
-        recipe: str | None = None,
-        debug: bool | None = None,
-        overrides: Overrides | None = None,
-        app_id: str | None = None,
-        actor_id: str | None = None,
-        llm_provider: str | None = None,
-        model: str | None = None,
-        model_version: str | None = None,
-        request_token_count: int | None = None,
-        response_token_count: int | None = None,
-        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,
-        extra_info: ExtraInfo | None = None,
-        count_tokens: bool | None = None,
-    ) -> PangeaResponse[GuardResult]:
-        """
-        Guard LLM input and output
+        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]
 
-        Analyze and redact content to avoid manipulation of the model, addition
-        of malicious content, and other undesirable data transfers.
+        return response
 
-        OperationId: ai_guard_post_v1beta_guard
 
-        Args:
-            input: 'messages' (required) contains Prompt content and role array
-                in JSON format. The `content` is the multimodal text or image
-                input that will be analyzed. Additional properties such as
-                'tools' may be provided for analysis.
-            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
-            app_name: Name of source application.
-            llm_provider: Underlying LLM.  Example: 'OpenAI'.
-            model: Model used to perform the event. Example: 'gpt'.
-            model_version: Model version used to perform the event. Example: '3.5'.
-            request_token_count: Number of tokens in the request.
-            response_token_count: Number of tokens in the response.
-            source_ip: IP address of user or app or agent.
-            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.
-            extra_info: (AIDR) Logging schema.
-            count_tokens: Provide input and output token count.
-        """
-        return self.request.post(
-            "v1beta/guard",
-            GuardResult,
-            data={
-                "input": input,
-                "recipe": recipe,
-                "debug": debug,
-                "overrides": overrides,
-                "app_id": app_id,
-                "actor_id": actor_id,
-                "llm_provider": llm_provider,
-                "model": model,
-                "model_version": model_version,
-                "request_token_count": request_token_count,
-                "response_token_count": response_token_count,
-                "source_ip": source_ip,
-                "source_location": source_location,
-                "tenant_id": tenant_id,
-                "event_type": event_type,
-                "sensor_instance_id": sensor_instance_id,
-                "extra_info": extra_info,
-                "count_tokens": count_tokens,
-            },
-        )
+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 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)
+    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 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,
-        )
+    if len(messages) == 0:
+        return [], []
 
-    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,
-        )
+    system_messages = [msg for msg in messages if msg.role == "system"]
+    system_indices = [i for i, msg in enumerate(messages) if msg.role == "system"]
 
-    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 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]
 
-    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,
-        )
+    # 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)
+    ]